Add MXFP8 E4M3 flash attention for CDNA4 - #36
Draft
ysa2215 wants to merge 14 commits into
Draft
Conversation
Remove a stale Triton JIT global dependency in backward quantization and add a user-facing autograd test so the MXFP8 attention wrapper is covered end to end.
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.
Summary
Adds MXFP8 E4M3 flash attention for AMD CDNA4 / gfx950, including forward and supported backward paths for dQ, dK, and dV. This also adds the
mxfp8_e4m3branch to the attention dispatch layer.Both forward and backward use
tl.dot_scaleddirectly, with no CDNA3 dequant fallback. Backward reuses the E4M3 tensors and 2D-block scales saved by forward instead of requantizing Q/K/V, so the two directions operate on the same quantized inputs.Validation
Tested on gfx950 / MI350, single GPU.
test_mxfp8_attention.pyandtest_mxfp8_attention_reference.py: 75 passedThe kernel-vs-reference accuracy confirms the backward port and scale broadcasting. The lower end-to-end accuracy matches the pure-PyTorch MXFP8 reference, so the remaining gap is from MXFP8 quantization rather than the Triton kernel.
The forward reference and both backward references are pure PyTorch, which keeps kernel porting errors separable from quantization error.
Scope
The supported scope is enforced by entry assertions:
layout == "bhsd"onlyseqlen_q == seqlen_kAll operands are E4M3. E5M2 is not parameterized in this PR because switching formats would require kernel changes, and the measured backward accuracy does not justify adding that complexity.
Reviewer Notes
MXFP8_ATTENTION_PLAN.mdsection 9.2.head_dim_qk != head_dim_v. A crossed-shape test was added to cover this case.