Repository navigation
Fix CI: gpu-metal crashes on master since #30003 - #30208
Closed
ServeurpersoCom wants to merge 1 commit into
Closed
ServeurpersoCom wants to merge 1 commit into
ServeurpersoCom wants to merge 1 commit into
Conversation
Member
|
We can add instantiations for these head sizes. |
Contributor
Author
|
That's what I was thinking; I'm checking to see if any of our models use it, if not, no big deal, it'll just be done ahead of time! |
Contributor
Author
|
Superseded by #30209, which adds the 128/96 kernels instead of rejecting the pair, and accepts every pair Metal compiles |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
#30003 adds flash attention tests where the K and V heads have different sizes: K = 128, V = 96. Metal said it could run them, but it only has kernels when K and V are the same size, plus two special pairs, 320/256 and 576/512. So it went looking for a 128/96 kernel that doesn't exist, and test-backend-ops crashed.
Metal now only accepts the head size pairs it really has a kernel for. Anything else falls back to another backend instead of crashing. This also replaces two older partial checks with a single complete one.
Beyond the CI, this was a real Metal bug: any graph using flash attention with one of these head pairs crashed the process, it now falls back to the CPU for that op. -> Native Metal kernels for more pairs could come in a follow-up if a model needs them.
Additional information
Tested on an M5 Max: test-backend-ops -o FLASH_ATTN_EXT passes 4675/4675 on Metal with no crash, the 128/96 cases now show as not supported, and the 320/256 and 576/512 pairs still run on Metal.
Follow-up #30003
Requirements