From 3bf01693a795f09e22960155c92329c86bf6f144 Mon Sep 17 00:00:00 2001 From: Carlos Trujillo Date: Sat, 19 Sep 2026 14:40:20 +0300 Subject: [PATCH 1/2] Fix SDPA output dtype following config.floatX instead of the input dtype The softmax scale and the causal mask in _sdpa_graph used config.floatX to type their values. When config.floatX is float64 (the default), a float32 (q @ k^T) gets promoted to float64 through the scale multiplication, and the causal mask keeps it there, so the attention output is float64 regardless of the input dtype. Use q.dtype instead, which tracks the caller's precision. The value delta is <2e-7 (float32 scale vs float64 scale on a typical d_k=64), well within existing test tolerances (atol=1e-5). --- pytensor_ml/layers/attention.py | 9 +++--- tests/test_attention.py | 51 +++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 4 deletions(-) diff --git a/pytensor_ml/layers/attention.py b/pytensor_ml/layers/attention.py index 1ef71cc..8a82ca9 100644 --- a/pytensor_ml/layers/attention.py +++ b/pytensor_ml/layers/attention.py @@ -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 @@ -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]) + compute_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(compute_dtype)) else: - scale_t = pt.as_tensor(scale, dtype=config.floatX) + scale_t = pt.as_tensor(scale, dtype=compute_dtype) scores = (q @ k.swapaxes(-1, -2)) * scale_t @@ -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(compute_dtype) scores = scores + causal if mask is not None: diff --git a/tests/test_attention.py b/tests/test_attention.py index 6fdd1f9..b5cc69b 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -351,3 +351,54 @@ def test_attention_vectorizes_over_independent_inputs(rng): def test_multihead_attention_validates_head_counts(kwargs, message): with pytest.raises(ValueError, match=message): MultiheadAttention("mha", **kwargs) + + +@pytest.mark.parametrize("is_causal", [False, True], ids=["full", "causal"]) +def test_sdpa_dtype_follows_input(is_causal, rng): + """The output dtype of scaled_dot_product_attention must match the input dtype, regardless of + ``config.floatX``. Regression test for pymc-devs/pytensor-ml#164.""" + b, n_head, seq, dim = 2, 3, 5, 4 + q = pt.tensor("q", shape=(b, n_head, seq, dim), dtype="float32") + k = pt.tensor("k", shape=(b, n_head, seq, dim), dtype="float32") + v = pt.tensor("v", shape=(b, n_head, seq, dim), dtype="float32") + out = scaled_dot_product_attention(q, k, v, is_causal=is_causal) + + q_np = rng.normal(size=q.type.shape).astype("float32") + k_np = rng.normal(size=k.type.shape).astype("float32") + v_np = rng.normal(size=v.type.shape).astype("float32") + result = out.eval({q: q_np, k: k_np, v: v_np}) + assert result.dtype == "float32" + + +def test_sdpa_dtype_with_explicit_scale(rng): + """An explicit float ``scale`` must not upcast a float32 graph to float64. + Regression test for pymc-devs/pytensor-ml#164.""" + b, n_head, seq, dim = 2, 3, 5, 4 + q = pt.tensor("q", shape=(b, n_head, seq, dim), dtype="float32") + k = pt.tensor("k", shape=(b, n_head, seq, dim), dtype="float32") + v = pt.tensor("v", shape=(b, n_head, seq, dim), dtype="float32") + out = scaled_dot_product_attention(q, k, v, scale=0.5) + + q_np = rng.normal(size=q.type.shape).astype("float32") + k_np = rng.normal(size=k.type.shape).astype("float32") + v_np = rng.normal(size=v.type.shape).astype("float32") + result = out.eval({q: q_np, k: k_np, v: v_np}) + assert result.dtype == "float32" + + +def test_sdpa_dtype_with_same_dtype_mask(rng): + """A float32 additive mask combined with a causal mask must keep float32 output. + Regression test for pymc-devs/pytensor-ml#164.""" + b, n_head, seq, dim = 2, 3, 5, 4 + q = pt.tensor("q", shape=(b, n_head, seq, dim), dtype="float32") + k = pt.tensor("k", shape=(b, n_head, seq, dim), dtype="float32") + v = pt.tensor("v", shape=(b, n_head, seq, dim), dtype="float32") + mask = pt.tensor("mask", shape=(b, n_head, seq, seq), dtype="float32") + out = scaled_dot_product_attention(q, k, v, mask=mask, is_causal=True) + + q_np = rng.normal(size=q.type.shape).astype("float32") + k_np = rng.normal(size=k.type.shape).astype("float32") + v_np = rng.normal(size=v.type.shape).astype("float32") + mask_np = rng.normal(size=mask.type.shape).astype("float32") + result = out.eval({q: q_np, k: k_np, v: v_np, mask: mask_np}) + assert result.dtype == "float32" From 7e8f24d8a6eaebf8456dd86a4a9cdbe8980be7b6 Mon Sep 17 00:00:00 2001 From: Carlos Trujillo Date: Sat, 19 Sep 2026 19:14:52 +0300 Subject: [PATCH 2/2] Address SDPA dtype review feedback --- pytensor_ml/layers/attention.py | 8 ++-- tests/test_attention.py | 73 ++++++--------------------------- 2 files changed, 17 insertions(+), 64 deletions(-) diff --git a/pytensor_ml/layers/attention.py b/pytensor_ml/layers/attention.py index 8a82ca9..ec2f093 100644 --- a/pytensor_ml/layers/attention.py +++ b/pytensor_ml/layers/attention.py @@ -40,12 +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]) - compute_dtype = q.dtype + input_dtype = q.dtype if scale is None: - scale_t = 1.0 / pt.sqrt(q.shape[-1].astype(compute_dtype)) + scale_t = 1.0 / pt.sqrt(q.shape[-1].astype(input_dtype)) else: - scale_t = pt.as_tensor(scale, dtype=compute_dtype) + scale_t = pt.as_tensor(scale, dtype=input_dtype) scores = (q @ k.swapaxes(-1, -2)) * scale_t @@ -55,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(compute_dtype) + causal = pt.where(k_idx <= q_idx + (sk - sq), 0.0, -np.inf).astype(input_dtype) scores = scores + causal if mask is not None: diff --git a/tests/test_attention.py b/tests/test_attention.py index b5cc69b..2d97091 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -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") + 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) @@ -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) @@ -351,54 +355,3 @@ def test_attention_vectorizes_over_independent_inputs(rng): def test_multihead_attention_validates_head_counts(kwargs, message): with pytest.raises(ValueError, match=message): MultiheadAttention("mha", **kwargs) - - -@pytest.mark.parametrize("is_causal", [False, True], ids=["full", "causal"]) -def test_sdpa_dtype_follows_input(is_causal, rng): - """The output dtype of scaled_dot_product_attention must match the input dtype, regardless of - ``config.floatX``. Regression test for pymc-devs/pytensor-ml#164.""" - b, n_head, seq, dim = 2, 3, 5, 4 - q = pt.tensor("q", shape=(b, n_head, seq, dim), dtype="float32") - k = pt.tensor("k", shape=(b, n_head, seq, dim), dtype="float32") - v = pt.tensor("v", shape=(b, n_head, seq, dim), dtype="float32") - out = scaled_dot_product_attention(q, k, v, is_causal=is_causal) - - q_np = rng.normal(size=q.type.shape).astype("float32") - k_np = rng.normal(size=k.type.shape).astype("float32") - v_np = rng.normal(size=v.type.shape).astype("float32") - result = out.eval({q: q_np, k: k_np, v: v_np}) - assert result.dtype == "float32" - - -def test_sdpa_dtype_with_explicit_scale(rng): - """An explicit float ``scale`` must not upcast a float32 graph to float64. - Regression test for pymc-devs/pytensor-ml#164.""" - b, n_head, seq, dim = 2, 3, 5, 4 - q = pt.tensor("q", shape=(b, n_head, seq, dim), dtype="float32") - k = pt.tensor("k", shape=(b, n_head, seq, dim), dtype="float32") - v = pt.tensor("v", shape=(b, n_head, seq, dim), dtype="float32") - out = scaled_dot_product_attention(q, k, v, scale=0.5) - - q_np = rng.normal(size=q.type.shape).astype("float32") - k_np = rng.normal(size=k.type.shape).astype("float32") - v_np = rng.normal(size=v.type.shape).astype("float32") - result = out.eval({q: q_np, k: k_np, v: v_np}) - assert result.dtype == "float32" - - -def test_sdpa_dtype_with_same_dtype_mask(rng): - """A float32 additive mask combined with a causal mask must keep float32 output. - Regression test for pymc-devs/pytensor-ml#164.""" - b, n_head, seq, dim = 2, 3, 5, 4 - q = pt.tensor("q", shape=(b, n_head, seq, dim), dtype="float32") - k = pt.tensor("k", shape=(b, n_head, seq, dim), dtype="float32") - v = pt.tensor("v", shape=(b, n_head, seq, dim), dtype="float32") - mask = pt.tensor("mask", shape=(b, n_head, seq, seq), dtype="float32") - out = scaled_dot_product_attention(q, k, v, mask=mask, is_causal=True) - - q_np = rng.normal(size=q.type.shape).astype("float32") - k_np = rng.normal(size=k.type.shape).astype("float32") - v_np = rng.normal(size=v.type.shape).astype("float32") - mask_np = rng.normal(size=mask.type.shape).astype("float32") - result = out.eval({q: q_np, k: k_np, v: v_np, mask: mask_np}) - assert result.dtype == "float32"