Align the JAX causal mask bottom-right when queries and keys differ

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

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

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)