From 00e2f2a3880ec3ea845a634dcbedc7854b2ea48d Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Tue, 6 Oct 2026 12:25:15 -0500 Subject: [PATCH 1/2] Pin the reference causal mask for a query shorter than its keys --- tests/test_attention.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_attention.py b/tests/test_attention.py index 6fdd1f9..205de59 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -60,6 +60,18 @@ def test_sdpa_matches_reference(is_causal, scale, rng): np.testing.assert_allclose(out, sdpa_np(q, k, v, is_causal=is_causal, scale=scale), atol=1e-5) +def test_sdpa_causal_query_shorter_than_its_keys_sees_the_whole_prefix(rng): + """Decoding from a key-value cache: the triangle aligns bottom-right, so each of the two queries + attends to every key up to its own position at the end of the five.""" + q = rng.normal(size=(2, 3, 2, 4)).astype(floatX) + k = rng.normal(size=(2, 3, 5, 4)).astype(floatX) + v = rng.normal(size=(2, 3, 5, 6)).astype(floatX) + + out = scaled_dot_product_attention(q, k, v, is_causal=True).eval() + + np.testing.assert_allclose(out, sdpa_np(q, k, v, is_causal=True), atol=1e-5) + + 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) From bb6a8cda6316581f25f4a71579c5b596d52fd71c Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Tue, 6 Oct 2026 12:25:15 -0500 Subject: [PATCH 2/2] Align the JAX causal mask bottom-right when queries and keys differ --- pytensor_ml/dispatch/jax/attention.py | 17 ++++++++++++++--- tests/dispatch/jax/test_attention.py | 15 +++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/pytensor_ml/dispatch/jax/attention.py b/pytensor_ml/dispatch/jax/attention.py index 7b68715..9def94b 100644 --- a/pytensor_ml/dispatch/jax/attention.py +++ b/pytensor_ml/dispatch/jax/attention.py @@ -13,13 +13,24 @@ def jax_funcify_AttentionLayer(op, node=None, **kwargs): scale = op.scale def attention(q, k, v, mask=None): + # jax's is_causal aligns the triangle top-left, so a query block shorter than its keys, as when + # decoding from a cache, would see only the first keys. The layer aligns it bottom-right, so + # every query sees the whole prefix before it; build that triangle here when the lengths differ. + q_len, kv_len = q.shape[-2], k.shape[-2] + causal, causal_mask = is_causal, None + if is_causal and q_len != kv_len: + rows = jnp.arange(q_len)[:, None] + cols = jnp.arange(kv_len)[None, :] + causal, causal_mask = False, (cols <= rows + (kv_len - q_len))[None, None] # jax expects (batch, seq, head, dim); our convention is (batch, head, seq, dim). The additive - # mask is (batch, head, seq, seq) in both, so only q/k/v are transposed. is_causal composes with - # bias, so a combined causal + padding mask needs no manual construction here. + # mask is (batch, head, seq, seq) in both, so only q/k/v are transposed. jax combines the + # additive bias with whichever triangle applies, its own or the boolean one built above. q = jnp.swapaxes(q, -3, -2) k = jnp.swapaxes(k, -3, -2) v = jnp.swapaxes(v, -3, -2) - out = jax.nn.dot_product_attention(q, k, v, bias=mask, scale=scale, is_causal=is_causal) + out = jax.nn.dot_product_attention( + q, k, v, bias=mask, mask=causal_mask, scale=scale, is_causal=causal + ) return jnp.swapaxes(out, -3, -2) return attention diff --git a/tests/dispatch/jax/test_attention.py b/tests/dispatch/jax/test_attention.py index e6dab66..e2999dd 100644 --- a/tests/dispatch/jax/test_attention.py +++ b/tests/dispatch/jax/test_attention.py @@ -55,6 +55,21 @@ def test_additive_mask(is_causal, rng): compare_jax_and_py([q_sym, k_sym, v_sym], out, [q, k, v], assert_fn=assert_close) +@pytest.mark.parametrize("with_mask", [False, True], ids=["causal_only", "mask_and_causal"]) +def test_a_query_shorter_than_its_keys_sees_the_whole_prefix(with_mask, rng): + """Decoding from a key-value cache: two queries over four keys. The triangle aligns bottom-right, + so the first query sees keys 0 to 2 and the second all four.""" + q = rng.normal(size=(2, 3, 2, 4)).astype(floatX) + k, v = (rng.normal(size=(2, 3, 4, 4)).astype(floatX) for _ in range(2)) + mask = np.where(rng.normal(size=(2, 3, 2, 4)) > 0, 0.0, -np.inf).astype(floatX) + mask[:, :, :, 0] = 0.0 + q_sym, k_sym, v_sym = symbolic_like(q, k, v) + out = scaled_dot_product_attention( + q_sym, k_sym, v_sym, mask=pt.as_tensor(mask) if with_mask else None, is_causal=True + ) + compare_jax_and_py([q_sym, k_sym, v_sym], out, [q, k, v], assert_fn=assert_close) + + def test_dispatch_auto_registers(): """Loading JAX's dispatch must auto-register ours -- no `import pytensor_ml.dispatch.jax` needed.""" from pytensor.link.jax.dispatch import jax_funcify