Run attention on MLX wherever the fused kernel can compute it - #22419
Conversation
shoumikhin
added
release notes: apple
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22419
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 9eed27d with merge base 834a4fb ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
|
I'm not sure this is the best way to support rank 3 attention. The SDPA handler could instead emit: ExpandDimsNode (only if rank 3) → SdpaNode → SqueezeNode could (only if rank 3). The ExpandDimsNode/SqueezeNodecould are no-data copy on MLX. |
metascroy
left a comment
metascroy
left a comment
There was a problem hiding this comment.
Rather than skip, change emission logic on SDPA node to still hit fused path.
99bff01 to
8c4f8e7
Compare
shoumikhin
changed the title
|
Thanks, you were right. Switched to the emission approach: Two things I ran into while doing it, both worth a look: The added dimension has to go at the front, not the middle. For a rank-3 input the first dimension is already the head one, so It also cannot be limited to rank 3. Rank 2 is accepted by PyTorch and fails the same way, so it takes two added dimensions. Rank 5 needs the opposite treatment and is left to decompose: folding its leading dimensions changes which of them the kernel reads as heads, and with unequal batch sizes it silently pairs the wrong ones. While checking shapes I found the kernel is claiming three other forms it cannot compute. The one that matters: a causal mask where the query and key lengths differ returns wrong values with no error, because Torch anchors the mask top-left and MLX bottom-right. That is the decode step of any model generating one token at a time, it is happening today, and it reaches the speech example in this repo. Measured 2.94 against eager. That check is in this change rather than a follow-up because adapting rank 3 makes a rank-3 decode call reachable for the first time. Without it that call returns 3.12 instead of 0.0, so shape adaptation on its own would turn a loud failure into a silent one. Happy to split it out if you would rather review them separately. |
|
Holding this one: CI found a real problem with it, and it is mine, not a flake.
When one graph has two attention calls and the backend takes one but not the other, the operator is still kept whole for both, because that decision is made per operator rather than per call. The declined one is then neither adapted nor decomposed, and export stops. So declining a call is not free the way I assumed. Reproduced locally with a graph holding one square call and one non-square causal call, which is the shape a speech model has. I would rather get this right than narrow the checks until CI goes quiet, so I am reworking it. The shape adaptation itself is unaffected and still measures correct; the question is only what to do with a call the kernel cannot compute when a sibling call can. |
8c4f8e7 to
484d97d
Compare
|
Fixed, and the speech test passes now. Two things changed since the version that broke it. The declined causal call is no longer declined. Where the query and key lengths differ, the top-left mask is built explicitly with the lengths read from the graph, so the call keeps the fused kernel and follows a cache that grows at runtime. That removes the mixing rather than working around it, and it measures faster than either alternative on a decode-shaped call. The remaining declines needed a partitioner fix. Preservation from decomposition is per operator, not per call, so one claimed call was keeping the operator whole for a declined one in the same graph, and that call was then neither lowered nor decomposed. That is what the speech test hit. An operator is now given back whenever any of its calls is unsupported. Worth knowing: that second part is not about attention at all. Two calls to On the causal convention, since it changes what users get: the mismatch is silent on main, not loud. A decode-shaped call returns values 2.94 away from eager with no error raised, and the example in this repo takes that path on every step past the first. It is 0.0 here. |
484d97d to
b09152e
Compare
|
Reworked. The shape adaptation is what you asked for: Two notes on what else is in here, both of which I would rather flag than have you find. Declining a call is not free, and that is why my first attempt broke the speech test. Preservation from decomposition is per operator, not per call, so one claimed call was keeping the operator whole for a declined one in the same graph, and that call then neither lowered nor decomposed. There is a partitioner change for it. Worth knowing that this is not an attention problem: two calls to The causal guard changes what the speech example computes, so it deserves your eye. At a decode step that example passes |
b09152e to
f59ef08
Compare
|
Third push, and this one should be clean. The shape adaptation is what you asked for and it is unchanged from the last round: What changed is that the causal check is out of this change entirely. I had it declining So this now leaves only two shapes to decompose, rank 5 and up and unequal batch at rank 4, both of which fail at execute today. To check the speech model was really unaffected rather than assume it, I exported it with and without this change and compared: identical serialized node counts, and an identical greedy token sequence over twelve decode steps. One part is worth your eye since it is outside attention. Declining any call needed a partitioner change, because preservation from decomposition is per operator rather than per call, so a claimed call was keeping the operator whole for a declined one in the same graph. Two calls to |
f59ef08 to
474e78e
Compare
|
Thanks, this was a genuinely good catch list. Three of these were real bugs of mine and I have fixed them. One I am pushing back on with a measurement. The rest are done. The grouped key unwrap at rank 3. You are right, and this was the worst thing in the change. The unwrap keys on dimension 1, which is the head dimension at rank 4 and the key sequence below it, so admitting rank 3 let it absorb a repeat carrying real keys. I reproduced it: the repeat node disappears, and the answer is off by 5.1 with a causal mask while matching exactly without one, which is precisely why it hides. Both unwraps now run only at rank 4, and there is a test that keeps the repeat at rank 3 and one that still absorbs it at rank 4. Non-square causal at the newly lifted ranks. Also right, and measured the same way: rank 3 off by 2.6, rank 2 off by 2.2, both silent. Those two ranks are now declined when the lengths differ. I left rank 4 alone deliberately, since it already reaches the kernel today and changing it turns out to need either a change to what the speech example computes or a fix to an off-by-one in the decomposed path. Filed separately as #22426 with the numbers. The head count check raising. You classed this minor and rejected it, but I think it deserved to stay: a zero key head count aborts the entire export with a divide by zero rather than declining one node. Fixed, with a test. Giving the whole operator back, and the per-node filter. You are right that the docstring stated something untrue about the framework, and right that other backends return a real filter. I have corrected the docstring. But I am keeping the coarse behaviour, because I ran the filter: on an attention block that reshapes its output, with a declined call in the same graph, it fails to lower at all with The rest. Both comment reasons were wrong and are rewritten: folding lands in the batch slot rather than the head slot, and the kernel rejects a mismatched batch rather than indexing it. The rank-2 test now asserts the squeeze count, the rejection tests now assert the work stayed on this backend so they cannot pass when nothing is delegated, and the precondition helper reads its operands once. The description had the counts wrong, and it now says fifteen tests with eleven failing at base, and lists the rank-4 calls the head checks refuse. |
474e78e to
a033ec7
Compare
Attention on rank-3 tensors exports without complaint and then fails when you run
the model:
MLX execute failed: [scaled_dot_product_attention] input with shape
(2,16,64) expected to be rank 4
PyTorch accepts rank 2, 3, 4 and 5 here. The fused kernel takes rank 4 only, and has
other requirements besides. The handler matched the operator by name and checked none
of them, so it claimed calls the kernel cannot run.
Rank 2 and rank 3 gain the missing leading dimensions, run on the fused kernel, and the
result is squeezed back. The dimensions go at the front, because for a rank-3 input the
first dimension is already the head one, and inserting in the middle moves it into the
batch slot, which misaligns masks and grouped heads.
Admitting a lower rank also reaches code that was only ever given rank 4. The grouped
key and value unwrap looks for a repeat on dimension 1, which is the head dimension at
rank 4 and the key sequence below it, so at rank 3 it would absorb a repeat that carries
real keys. Without a mask the result still matches, because a duplicated key and its
value give the same weighted sum, so this hides. With a causal mask it is wrong by whole
units. Both unwraps now run only at rank 4.
Three shapes are left for the lowering step to decompose into primitives, which still
run on this backend:
Rank 5 and above, because folding the leading dimensions pairs the wrong operands as
soon as one of them broadcasts a batch the others do not.
Unequal batch sizes at rank 4, because the kernel requires them to match and rejects the
call, so it fails at execute today.
Causal attention at rank 2 or rank 3 whose query and key lengths differ. Torch anchors a
causal mask at the top left and MLX at the bottom right, so the two disagree there and
the disagreement is silent. Rank 4 already reaches the kernel today and keeps its current
behaviour, since changing it is a separate question.
Declining a call turns out not to be free, which is the other half of this change.
Preservation from decomposition is requested per operator, so one claimed call would keep
the operator whole for a declined one in the same graph, leaving that call neither lowered
nor decomposed and stopping export. Give the whole operator back when any of its calls is
unsupported. The framework does offer a per-node filter that would keep the other calls
fused, and it is deliberately not used: it puts the program on a path that fails on an
ordinary attention block that reshapes its output. The docstring says so.
That half is not specific to attention. Two calls to torch.roll in one graph, one
supported and one not, already fail at execute today for the same reason.
Test Plan:
Fifteen partitioner tests. Eleven fail before this change and all fifteen pass after. They
assert on the serialized nodes, so they check which path a call took rather than only that
it answered, and the rejection cases also assert the work stayed on this backend, so they
cannot pass when nothing is delegated at all. Added a rank-3 case to the operator suite
too, which runs the compiled runtime.
Ran every combination on an Apple Silicon Mac against eager: plain, causal, explicit
scale, float mask, boolean mask, grouped-query attention, batch size 1, float16 and
bfloat16, at ranks 2 through 5. Rank 3 matches to 3.6e-07 and keeps the fused kernel in
all seven of its variants, rank 2 to 3.6e-07, rank 5 decomposes and matches to 3.6e-07
where it previously failed, and unequal batch goes from failing to 2.4e-07.
The grouped key case was measured both ways: at rank 4 the repeat is absorbed and the
result matches, and at rank 3 with a causal mask absorbing it differs from what the model
asked for by 5.1, so the guard is what keeps that correct.
A zero key head count used to divide by zero and abort the whole export; it now declines
that node.
A graph holding one supported and one unsupported call exports and runs. Two rolls in one
graph go from failing at execute to matching eager exactly.
Exported the speech example and compared it against the same export without this change:
identical serialized node counts and an identical greedy token sequence over twelve decode
steps, so that model is unaffected.
Ran the neighbouring backend test files, 71 tests, all passing.
a033ec7 to
9eed27d
Compare
The problem
Attention on rank-3 tensors exports without complaint and then fails when you run the model:
PyTorch accepts rank 2, 3, 4 and 5 here. The fused kernel takes rank 4 only, and has other requirements besides. The handler matched the operator by name and checked none of them, so it claimed calls the kernel cannot run.
Adapting the shapes that can be adapted
ExpandDimsNode→SdpaNode→SqueezeNode, as suggested in review, so rank 3 keeps the fused kernel. Rank 2 needs it too and takes two added dimensions.The dimensions go at the front. For a rank-3 input the first dimension is already the head one, so inserting in the middle moves it into the batch slot, which misaligns masks and grouped heads.
Admitting a lower rank also reaches code that was only ever given rank 4. The grouped key and value unwrap looks for a repeat on dimension 1, which is the head dimension at rank 4 and the key sequence below it, so at rank 3 it would absorb a repeat carrying real keys. Without a mask the result still matches, because a duplicated key and its value give the same weighted sum, which is what makes this easy to miss. With a causal mask it is off by 5.1. Both unwraps now run only at rank 4.
Three shapes left to decompose
Decomposed calls still run on this backend, which the tests assert rather than assume.
The causal one needs a word. Torch anchors a causal mask at the top left and MLX at the bottom right, so the two disagree whenever the lengths differ, silently. Rank 4 already reaches the kernel today and keeps its current behaviour: fixing that needs either a change to what the speech example computes or a fix to an off-by-one in the decomposed path, so it is filed separately as #22426. What this change does is avoid opening two more ranks onto the same trap.
Declining a call is not free
Preservation from decomposition is requested per operator, so one claimed call keeps the operator whole for a declined one in the same graph, and that call is then neither lowered nor decomposed:
Give the whole operator back when any of its calls is unsupported.
The framework does offer a per-node filter that would keep the other calls fused. It is deliberately not used, and the docstring says why: on an attention block that reshapes its output, with a declined call in the same graph, that path fails to lower at all with
Cannot view a tensor with shape (1,16,4,16) and strides (1024,16,256,1). It trades the fusion cost for a hard failure on a very ordinary shape. The cost of the coarser choice is real and written down: one declined call unfuses the operator's other calls in that graph.This half is not specific to attention. Two calls to
torch.rollin one graph, one supported and one not, already fail at execute today for the same reason.Test plan
Fifteen partitioner tests. Eleven fail before this change, all fifteen pass after. They assert on serialized nodes, so they check which path a call took rather than only that it answered, and the rejection cases also assert the work stayed on this backend, so they cannot pass when nothing is delegated at all. A rank-3 case was added to the operator suite too, which runs the compiled runtime.
Ran every combination on an Apple Silicon Mac against eager: plain, causal, explicit scale, float mask, boolean mask, grouped-query attention, batch size 1, float16 and bfloat16, ranks 2 through 5.
Exported the speech example and compared it against the same export without this change: identical serialized node counts and an identical greedy token sequence over twelve decode steps, so that model is unaffected.
Ran the neighbouring backend test files, 71 tests, all passing.