feat: MLA prefill NVFP4 fused output - #8
Open
carlyou wants to merge 4 commits into
Open
Conversation
Wire FA4's fused NVFP4 output into the MLA prefill path: when the o_proj input quant is kNvfp4Dynamic, forward_mha passes the packed-e2m1 output buffer, the global scale, and the 128x4-swizzled e4m3 block-scale buffer straight to flash-attention, eliminating the standalone scaled_fp4_quant kernel. FA writes the swizzled scale-factor layout consumed by the cutlass/flashinfer NVFP4 GEMMs directly. Fusion is gated to pure-prefill batches: the swizzled buffer's 128-row tiles cannot be split between FA (prefill tokens) and the separate decode-token post-quant. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com>
carlyou
commented
Jul 7, 2026
Comment on lines
+2337
to
+2344
| ( | ||
| # NVFP4: packed e2m1 codes, two per byte. | ||
| output.view(torch.float4_e2m1fn_x2).view( | ||
| -1, self.num_heads, self.v_head_dim // 2 | ||
| ) | ||
| if is_nvfp4 | ||
| else output.view(-1, self.num_heads, self.v_head_dim) | ||
| ) |
Owner
Author
There was a problem hiding this comment.
define a variable instead putting this inline.
Owner
Author
There was a problem hiding this comment.
Hoisted into a named fused_out (and fused_out_scales alongside it) in 7ae76c6.
carlyou
commented
Jul 7, 2026
carlyou
commented
Jul 7, 2026
Comment on lines
731
to
736
| output_scales=( | ||
| mha_output_scales[num_mqa_tokens:num_actual_toks] | ||
| mha_output_scales | ||
| if quant_key == kNvfp4Dynamic | ||
| else mha_output_scales[num_mqa_tokens:num_actual_toks] | ||
| if mha_output_scales is not None | ||
| else None |
Owner
Author
There was a problem hiding this comment.
define a variable above for better readibility
Owner
Author
There was a problem hiding this comment.
Hoisted the slicing into the mha_output_scales assignment in 7ae76c6. mypy also caught that the hoisted version would have sliced None for static FP8 - guarded.
Name the fused out/scales views instead of inlining them in the forward_mha call, and derive is_nvfp4 from _detect_output_quant_key rather than dtype-sniffing the scales buffer. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com>
carlyou
commented
Jul 7, 2026
| and attn_metadata.prefill.chunked_context is None | ||
| and self.impl.dcp_world_size <= 1 | ||
| and (not is_pergroup or pergroup_layout_ok) | ||
| and nvfp4_ok |
Owner
Author
There was a problem hiding this comment.
put quant_key != kNvfp4Dynamic or num_mqa_tokens == 0 inline here.
Owner
Author
There was a problem hiding this comment.
Inlined in 79f55c4 (comment moved onto the conjunct).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com>
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
1 participant
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.
Wire FA4's fused NVFP4 output into the MLA prefill path: when the o_proj input quant is kNvfp4Dynamic, forward_mha passes the packed-e2m1 output buffer, the global scale, and the 128x4-swizzled e4m3 block-scale buffer straight to flash-attention, eliminating the standalone scaled_fp4_quant kernel. FA writes the swizzled scale-factor layout consumed by the cutlass/flashinfer NVFP4 GEMMs directly.
Fusion is gated to pure-prefill batches: the swizzled buffer's 128-row tiles cannot be split between FA (prefill tokens) and the separate decode-token post-quant.
Purpose
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.