JAX aligns the causal mask top-left, so generation with a KV cache attends to the wrong prefix

scaled_dot_product_attention(is_causal=True) builds its triangle bottom-right — k_idx <= q_idx + (sk - sq) — so a short query block sees the whole prefix, which is what a KV cache needs. The JAX dispatch passes is_causal straight through to jax.nn.dot_product_attention, which aligns top-left, so query row 0 sees only key 0. Every other backend takes the reference path and gets the documented alignment. Equal query and key lengths agree to 7e-8, so training and every test pass; the divergence appears only when q_len != kv_len, which is only ever autoregressive decoding. Supplying your own mask alongside is_causal=True does not fix it — jax composes the two rather than letting the bias win.

import numpy as np
import pytensor
import pytensor.tensor as pt

from pytensor_ml.layers.attention import scaled_dot_product_attention

q = pt.tensor("q", shape=(1, 1, 2, 4))
k = pt.tensor("k", shape=(1, 1, 4, 4))
v = pt.tensor("v", shape=(1, 1, 4, 4))
out = scaled_dot_product_attention(q, k, v, is_causal=True)

# all scores equal, so each output row is the uniform average of what it may attend to
qn, kn = np.zeros((1, 1, 2, 4)), np.zeros((1, 1, 4, 4))
vn = np.eye(4).reshape(1, 1, 4, 4)

for mode in ["NUMBA", "JAX"]:
    print(mode, pytensor.function([q, k, v], out, mode=mode)(qn, kn, vn).squeeze())

# NUMBA [[0.333 0.333 0.333 0.   ]   <- rows see prefix {0,1,2} and {0,1,2,3}
#        [0.25  0.25  0.25  0.25 ]]
# JAX   [[1.   0.   0.   0.  ]       <- rows see {0} and {0,1}
#        [0.5  0.5  0.  0.  ]]

The MLX dispatch already builds the bottom-right triangle itself rather than delegating; the JAX one should do the same when q_len != kv_len.

Start at the JAX attention dispatch and compare it with the MLX dispatch, which already constructs the bottom-right causal triangle. Run the provided NUMBA versus JAX reproduction with unequal query and key lengths; done means JAX matches the reference bottom-right alignment for KV-cache decoding while equal-length behavior remains unchanged.