diff --git a/pytensor_ml/layers/attention.py b/pytensor_ml/layers/attention.py index 1ef71cc..cef1e87 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,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: diff --git a/tests/test_attention.py b/tests/test_attention.py index 205de59..41cae65 100644 --- a/tests/test_attention.py +++ b/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)