fix: preserve MXFP4 attention normalization with PNQ - #40
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes an attention normalization regression in the MXFP4 FlashAttention forward path when probability tiles are quantized (PNQ), ensuring the online softmax denominator uses the same quantized probability mass consumed by the P @ V update. It also adds a CPU reference/invariant test and an optional GPU cross-validation harness to validate PNQ behavior against a pre-PNQ baseline.
Changes:
- Quantize each softmax tile before both the
P @ Vupdate and the online row-sum update; derivel_ijfrom the dequantized packed tile. - Add a CPU reference PNQ recurrence plus an invariant test for the
V=1normalization property. - Add an optional script-mode GPU A/B cross-validation flow that can run kernels from a baseline Git worktree.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| tests/unittest/mxfp4/test_mxfp_pnq.py | Adds CPU PNQ reference + invariant test and a script-only cross-validation harness. |
| alto/kernels/fp4/mxfp4/triton_flash_attention_mxfp4.py | Updates the forward inner loop to compute the online row-sum from the dequantized packed probability tile. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+60
to
+65
| _install_quantization_test_stubs() | ||
|
|
||
| try: | ||
| from .utils import convert_from_mxfp4_pytorch, convert_to_mxfp4_pytorch | ||
| except ImportError: | ||
| from utils import convert_from_mxfp4_pytorch, convert_to_mxfp4_pytorch |
| output_pnq = mxfp4_pnq_reference_bhsd(q, k, v, causal=True, pnq=True) | ||
| output_without_pnq = mxfp4_pnq_reference_bhsd(q, k, v, causal=True, pnq=False) | ||
|
|
||
| assert torch.allclose(output_pnq, torch.ones_like(output_pnq), atol=0.0, rtol=0.0) |
Comment on lines
+352
to
356
| # TODO(fix): p_fp4 is now packed above, so the mask below never reaches | ||
| # the PV dot -- enabling dropout would silently become a no-op. Keeping | ||
| # both dropout and PNQ requires masking the packed tile instead. | ||
| if ENABLE_DROPOUT: | ||
| philox_offset = batch_philox_offset + start_m * BLOCK_M * actual_seqlen_k + start_n - BLOCK_N |
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
P @ Vupdate and online row-sum update.Validation
Normalization invariant:
V=1For constant
V=1, correctly normalized attention must produce1.norm_rationorm_ratioWith PNQ, MAE and relative-L2 are zero in all three cases. Without PNQ, the output loses up to approximately 2.4% of its amplitude systematically.
Numerical benefit: nonzero-mean
Vbiased V = randn + 2.0exposes the effect of missing probability mass on coherent outputs.pytest -q -p no:cacheprovider tests/unittest/mxfp4/test_mxfp_pnq.py