Skip to content

Fix SDPA output dtype following config.floatX instead of the input dtype #165

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cetagostini wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix SDPA output dtype following config.floatX instead of the input dtype #165

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3bf0169
Fix SDPA output dtype following config.floatX instead of the input dtype
cetagostini Sep 19, 2026
7e8f24d
Address SDPA dtype review feedback
cetagostini Sep 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions pytensor_ml/layers/attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import numpy as np
import pytensor.tensor as pt

from pytensor import config
from pytensor.tensor.variable import TensorVariable

from pytensor_ml.base import Layer, UnaryLayerOp, _resolve_layer_name
Expand Down Expand Up @@ -41,10 +40,12 @@ def _sdpa_graph(
k = _repeat_kv(k, q.type.shape[-3], k.type.shape[-3])
v = _repeat_kv(v, q.type.shape[-3], v.type.shape[-3])

input_dtype = q.dtype

if scale is None:
scale_t = 1.0 / pt.sqrt(q.shape[-1].astype(config.floatX))
scale_t = 1.0 / pt.sqrt(q.shape[-1].astype(input_dtype))

@jessegrabowski jessegrabowski Sep 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Type the 1.0 at input_dtype too, with pt.reciprocal(pt.sqrt(q.shape[-1].astype(input_dtype))). PyTensor types the bare literal at float32, so a float16 q gets a float32 scale here while the scale= branch below keeps float16.

else:
scale_t = pt.as_tensor(scale, dtype=config.floatX)
scale_t = pt.as_tensor(scale, dtype=input_dtype)

scores = (q @ k.swapaxes(-1, -2)) * scale_t

Expand All @@ -54,7 +55,7 @@ def _sdpa_graph(
# case) sees the whole prefix. Reduces to a plain lower triangle when sq == sk.
q_idx = pt.arange(sq)[:, None]
k_idx = pt.arange(sk)[None, :]
causal = pt.where(k_idx <= q_idx + (sk - sq), 0.0, -np.inf).astype(config.floatX)
causal = pt.where(k_idx <= q_idx + (sk - sq), 0.0, -np.inf).astype(input_dtype)
scores = scores + causal

if mask is not None:
Expand Down
22 changes: 13 additions & 9 deletions tests/test_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,13 +50,15 @@ def sdpa_np(q, k, v, *, is_causal=False, mask=None, scale=None):
@pytest.mark.parametrize("is_causal", [False, True], ids=["full", "causal"])
def test_sdpa_matches_reference(is_causal, scale, rng):
b, n_head, seq, qk_dim, v_dim = 2, 3, 5, 4, 6
q = rng.normal(size=(b, n_head, seq, qk_dim)).astype(floatX)
k = rng.normal(size=(b, n_head, seq, qk_dim)).astype(floatX)
v = rng.normal(size=(b, n_head, seq, v_dim)).astype(floatX)
q = rng.normal(size=(b, n_head, seq, qk_dim)).astype("float32")

@jessegrabowski jessegrabowski Sep 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as my note from last week about folding these into the existing tests: parametrize this one over dtype and key atol off it, rather than adding cases. float16 is the width the fix misses, and the hardcoded float32 is why nothing here catches it.

k = rng.normal(size=(b, n_head, seq, qk_dim)).astype("float32")
v = rng.normal(size=(b, n_head, seq, v_dim)).astype("float32")

out = scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=scale).eval()
with pytensor.config.change_flags(floatX="float64"):
out = scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=scale).eval()

assert out.shape == (b, n_head, seq, v_dim)
assert out.dtype == q.dtype
np.testing.assert_allclose(out, sdpa_np(q, k, v, is_causal=is_causal, scale=scale), atol=1e-5)


Expand All @@ -74,13 +76,15 @@ def test_sdpa_grouped_query(rng):
@pytest.mark.parametrize("is_causal", [False, True], ids=["mask_only", "mask_and_causal"])
def test_sdpa_additive_mask(is_causal, rng):
b, n_head, seq, dim = 2, 3, 5, 4
q = rng.normal(size=(b, n_head, seq, dim)).astype(floatX)
k = rng.normal(size=(b, n_head, seq, dim)).astype(floatX)
v = rng.normal(size=(b, n_head, seq, dim)).astype(floatX)
mask = np.where(rng.normal(size=(b, n_head, seq, seq)) > 0, 0.0, -1e9).astype(floatX)
q = rng.normal(size=(b, n_head, seq, dim)).astype("float32")
k = rng.normal(size=(b, n_head, seq, dim)).astype("float32")
v = rng.normal(size=(b, n_head, seq, dim)).astype("float32")
mask = np.where(rng.normal(size=(b, n_head, seq, seq)) > 0, 0.0, -1e9).astype("float32")

out = scaled_dot_product_attention(q, k, v, mask=mask, is_causal=is_causal).eval()
with pytensor.config.change_flags(floatX="float64"):
out = scaled_dot_product_attention(q, k, v, mask=mask, is_causal=is_causal).eval()

assert out.dtype == q.dtype
np.testing.assert_allclose(out, sdpa_np(q, k, v, mask=mask, is_causal=is_causal), atol=1e-5)


Expand Down
Loading