Skip to content

Add MXFP8 E4M3 flash attention for CDNA4 - #36

Draft
ysa2215 wants to merge 14 commits into
mainfrom
yue/mxfp8-attention
Draft

Add MXFP8 E4M3 flash attention for CDNA4#36
ysa2215 wants to merge 14 commits into
mainfrom
yue/mxfp8-attention

Conversation

@ysa2215

@ysa2215 ysa2215 commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

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_e4m3 branch to the attention dispatch layer.

Both forward and backward use tl.dot_scaled directly, 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.py and test_mxfp8_attention_reference.py: 75 passed
  • Backward kernel vs golden reference:
    • dQ/dK: 53-60 dB
    • dV: 59-63 dB
    • cosine similarity approximately 1.0
  • Public autograd end-to-end vs bf16 SDPA:
    • dQ/dK: approximately 23.3 dB, cosine similarity 0.9976
    • dV: approximately 26 dB, cosine similarity 0.9988

The 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:

  • CDNA4 / gfx950 only
  • layout == "bhsd" only
  • causal attention requires seqlen_q == seqlen_k
  • backward does not support alibi or dropout

All 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

  • The causal backward index range intentionally relies on square causal shapes, which are enforced at entry. This avoids extra dK/dV iterations for unsupported non-square causal shapes. Full details are in MXFP8_ATTENTION_PLAN.md section 9.2.
  • This PR also fixes a dO stride bug in the dK/dV group loop under GQA with head_dim_qk != head_dim_v. A crossed-shape test was added to cover this case.
  • Cleanup-only items that do not affect behavior are left out to keep this PR focused.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant