Build attention's scale and causal mask at the queries' dtype

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,11 @@ 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])
 
+    # Built at the queries' dtype, so float32 attention stays float32 whatever floatX is
     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(q.dtype))
     else:
-        scale_t = pt.as_tensor(scale, dtype=config.floatX)
+        scale_t = pt.as_tensor(scale, dtype=q.dtype)
 
     scores = (q @ k.swapaxes(-1, -2)) * scale_t
 
@@ -54,7 +54,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(q.dtype)
         scores = scores + causal
 
     if mask is not None:

tests/test_attention.py

@@ -72,6 +72,19 @@ def test_sdpa_causal_query_shorter_than_its_keys_sees_the_whole_prefix(rng):
     np.testing.assert_allclose(out, sdpa_np(q, k, v, is_causal=True), atol=1e-5)
 
 
+@pytest.mark.parametrize("scale", [None, 0.5], ids=["default_scale", "custom_scale"])
+@pytest.mark.parametrize("is_causal", [False, True], ids=["full", "causal"])
+def test_sdpa_keeps_single_precision_under_a_double_precision_floatx(is_causal, scale):
+    """The scale and the causal mask follow the queries' dtype. Built at floatX's float64, they would
+    promote a float32 attention, and everything after it, to float64."""
+    q = pt.tensor("q", shape=(1, 2, 3, 4), dtype="float32")
+
+    with pytensor.config.change_flags(floatX="float64"):
+        out = scaled_dot_product_attention(q, q, q, is_causal=is_causal, scale=scale)
+
+    assert out.dtype == "float32"
+
+
 def test_sdpa_grouped_query(rng):
     b, n_head, n_kv_head, seq, dim = 2, 6, 2, 4, 5
     q = rng.normal(size=(b, n_head, seq, dim)).astype(floatX)