What happens
On Apple Silicon, a model that calls scaled_dot_product_attention with is_causal=True and different query and key lengths gets wrong numbers back from the MLX backend. Nothing is raised, so the model simply produces bad output.
Measured against eager PyTorch on the same inputs, with q of shape [1, 4, 1, 64] and k/v of [1, 4, 16, 64]:
is_causal, query length 1, key length 16 -> max error 2.7 to 4.1 (varies with input)
is_causal, query length 16, key length 16 -> max error 4.8e-07
The equal-length case is fine. Only the unequal one is wrong.
Why
The two libraries anchor a causal mask at opposite corners. Torch puts it at the top left, so with one query and sixteen keys the single query attends to the first key only. MLX puts it at the bottom right, so the query attends to all sixteen. Both are self-consistent, and they disagree whenever the lengths differ.
Confirmed with torch alone, comparing is_causal=True against masks built both ways:
q len 1, k len 4: vs top-left 0.0, vs bottom-right 1.49
q len 1, k len 16: vs top-left 0.0, vs bottom-right 1.41
q len 2, k len 4: vs top-left 0.0, vs bottom-right 5.25
q len 4, k len 4: vs top-left 0.0, vs bottom-right 0.0
So torch means top-left, and that is what a model written against torch expects.
Why it matters
This is the decode step of any model that generates one token at a time: one new query against a key cache holding everything so far. Any such model on this backend is affected.
Two things that make it awkward to fix
Worth writing down, because the obvious fixes both reach beyond this bug.
Building the top-left mask explicitly and keeping the fused kernel is correct with respect to torch, and it changes what the speech example in this repository computes. That example passes is_causal=True against a cache window it has grown itself, and it wants attention over the whole window, which is the bottom-right reading. So the example and torch disagree about what it asked for, and the backend has been quietly siding with the example. Fixing the backend alone makes that example produce an empty transcript.
Leaving the call to be decomposed instead avoids the mask question, and hits a separate problem in the decomposed path:
[broadcast_shapes] Shapes (1,6,64,256) and (1,6,64,257) cannot be broadcast.
An off-by-one on the dynamic cache length, at the first decode step. That looks like its own bug, independent of the mask convention.
So a full fix is probably three things: the mask in the backend, an explicit mask in the example instead of is_causal, and the broadcast off-by-one in the decomposed path.
What happens
On Apple Silicon, a model that calls
scaled_dot_product_attentionwithis_causal=Trueand different query and key lengths gets wrong numbers back from the MLX backend. Nothing is raised, so the model simply produces bad output.Measured against eager PyTorch on the same inputs, with
qof shape[1, 4, 1, 64]andk/vof[1, 4, 16, 64]:The equal-length case is fine. Only the unequal one is wrong.
Why
The two libraries anchor a causal mask at opposite corners. Torch puts it at the top left, so with one query and sixteen keys the single query attends to the first key only. MLX puts it at the bottom right, so the query attends to all sixteen. Both are self-consistent, and they disagree whenever the lengths differ.
Confirmed with torch alone, comparing
is_causal=Trueagainst masks built both ways:So torch means top-left, and that is what a model written against torch expects.
Why it matters
This is the decode step of any model that generates one token at a time: one new query against a key cache holding everything so far. Any such model on this backend is affected.
Two things that make it awkward to fix
Worth writing down, because the obvious fixes both reach beyond this bug.
Building the top-left mask explicitly and keeping the fused kernel is correct with respect to torch, and it changes what the speech example in this repository computes. That example passes
is_causal=Trueagainst a cache window it has grown itself, and it wants attention over the whole window, which is the bottom-right reading. So the example and torch disagree about what it asked for, and the backend has been quietly siding with the example. Fixing the backend alone makes that example produce an empty transcript.Leaving the call to be decomposed instead avoids the mask question, and hits a separate problem in the decomposed path:
An off-by-one on the dynamic cache length, at the first decode step. That looks like its own bug, independent of the mask convention.
So a full fix is probably three things: the mask in the backend, an explicit mask in the example instead of
is_causal, and the broadcast off-by-one in the decomposed path.