Repository navigation
fix(flex): avoid decoding kernel selection for short GQA - #25
Merged
Merged
Conversation
Default the shared FlexAttention wrapper to BACKEND=TRITON to avoid PyTorch 2.13 AUTO lowering short GQA rows to an invalid decoding kernel. Preserve explicit backend and legacy force selectors without mutating caller options. Cover option handling in CPU CI and compare short-row outputs and Q/K/V gradients with SDPA under wrapped and nested compilation.
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.
Closes #24
Outcome
On the pinned PyTorch 2.13 runtime, short GQA training rows can enter FlexAttention's automatic
decoding path and fail with
NoValidChoicesError. Default the shared wrapper toBACKEND=TRITON, which uses the standard attention kernel.Preserve explicit
BACKENDand legacyFORCE_USE_FLEX_ATTENTIONchoices without mutatingcaller options. Add CPU policy checks to CI and GPU output/gradient regressions for the
previously failing shape, including calls nested inside
torch.compile.Verification
AUTOfailure with Q[1,32,112,128], K/V[1,8,262,128]and a DFlash blockmask; this requires no model checkpoint. Also reproduced and fixed the initial patch's
legacy-selector conflict.
and Q/K/V-gradient relative L2 errors are below 1% against SDPA's math backend; existing
attention forward/backward comparisons also pass.
ruffcheck, maintained documentation links, and CI YAML/test selection checks pass.verified Linux copy. Other GPU architectures, Torch versions and throughput are unverified.
Compatibility and provenance
the caller's control; attention masks, model configuration and training objectives retain
their existing contracts.
6c042a87140a84d13839e341ece2c5c3ada918bc; source headers identify the adaptation and preservethe upstream MIT notice. New CPU tests carry ASpec attribution. Existing license texts remain
applicable.