Repository navigation
[Perf][DSv4.1] MegaAttn: wo_a writes wo_b's MXFP8 input - #60118
Open
Juntian777 wants to merge 9 commits into
Open
Juntian777 wants to merge 9 commits into
Juntian777 wants to merge 9 commits into
Conversation
MegaAttn's wo_a einsum wrote a BF16 z that wo_b then quantized to MXFP8. From 1024 tokens, the einsum now writes z as MXFP8 with DeepGEMM's packed UE8M0 scales, a small Triton kernel converts the scales to FlashInfer's F8_128x4 layout, and wo_b takes the QuantizedActivation directly. The bytes wo_b sees are unchanged. - Gated on the per-32 einsum recipe, o_lora_rank % 128 == 0, an MXFP8 wo_b kernel that takes F8_128x4 scales, and no GEMM-RS (the fused GEMM + reduce-scatter takes a BF16 z only). - Below 1024 tokens (decode, small prefills) nothing changes: the swizzle launch costs about what the removed quantization saves there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
Juntian777
requested review from
AndreasKaratzas,
WoosukKwon,
mgoin,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
October 5, 2026 20:54
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
mergify
Bot
added
deepseek
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
6 tasks done
Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
DeepGEMM's FP8 einsum can now write its UE8M0 scales in the F8_128x4 layout that wo_b's FlashInfer MXFP8 GEMM reads (vllm-project/DeepGEMM#40). The Triton swizzle and the 1024-token threshold go away, so the fused path also runs in decode. Bump the DeepGEMM pin to that change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
Juntian777
requested review from
Harry-Chen and
LucasWilkinson
as code owners
October 16, 2026 09:31
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
vllm-project/DeepGEMM#40 now selects the F8_128x4 scale layout with `use_128x4_sfd_layout=True` instead of inferring it from a 3D scale tensor. Pass the knob and bump the DeepGEMM pin to that commit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
vllm-project/DeepGEMM#40 merged into the fork's dev branch as 727f847, whose tree is identical to the previously pinned PR head 4ffc70d. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
zyongye
approved these changes
Oct 17, 2026
Take main's removal of tools/install_deepgemm.sh (vllm-project#61265); the DeepGEMM pin now lives only in cmake/external_projects/deepgemm.cmake. Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #95555 for commit |
Contributor
Author
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #95555. |
Member
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #95555. |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #95650 for commit |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #95681 for commit |
This branch has not been deployed
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
2 participants
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.
Overview
MegaAttn's (
FLASHMLA_MEGA_ATTN_DSV41)wo_aeinsum now writeszdirectly as MXFP8, together with the F8_128x4 UE8M0 scales thatwo_b's FlashInfer MXFP8 GEMM reads.wo_btherefore no longer quantizes a BF16z, at any token count. It uses vllm-project/DeepGEMM#40 (merged), which adds the 128x4 scale layout to DeepGEMM's FP8 einsum epilogue. This PR bumps the DeepGEMM pin to it.Claims
_o_projdrops the standalonemxfp8_quantizekernel. On GB200 it is 0.4-2.0 us per layer faster in decode (1-256 tokens) and 2.8-141 us faster in prefill (512-32768 tokens). The one regression is +0.6 us at 384 tokens (see Details).wo_breceives byte-identical FP8 data and scales, so model outputs do not change.flash_mla_mega_attn.pychanges (+33/-5 lines vs main). Shared o_proj code and other backends are untouched.Validation
_o_projalone (wo_aeinsum, thenwo_bCuTe-DSL MXFP8 GEMM) at DeepSeek-V4.1-Flash TP1 shapes: 8 groups x 4096 -> 1024,wo_b8192 -> 5120. Setup: 1 x GB200, CUPTI, full CUDA graph, cold L2, median of 3 rounds, both arms in one process. Software: torch 2.13.0+cu132, FlashInfer 0.7.0.post1, DeepGEMM a7a9277 + DeepGEMM#40. Outputs are bitwise identical at all 20 sizes measured.The previous revision of this PR wrote MN-major scales, swizzled them with a separate Triton kernel (1.6-5.2 us per call), and was gated to 1024+ tokens. The E2E prefill (DEP4) and accuracy results below were measured on that previous revision and were not re-run. This revision only removes the swizzle kernel from the same prefill path, and it hands
wo_bthe same bytes. Decode E2E has not been measured.Previous revision: E2E prefill (DEP4, GB300) and accuracy
DeepSeek-V4.1-Flash, DEP4 (DP4 x TP1 + EP) on one GB300 node. Default config: SWA bounded replay on, MegaAttn, DeepGEMM mega MoE, fp8 KV cache,
max_num_batched_tokens=16384, prefix caching off. OfflineLLMper DP rank, one random-token prompt per rank,max_tokens=1. TTFT is measured fromscheduled_tstofirst_token_tson the slowest rank; median of 3 repeats after a warmup.GPQAEval, 198 problems, 0-shot chat, thinking on, temperature 1.0, top_p 0.95Tests:
tests/test_einsum.pypass, including the new bitwise check of the 128x4 scales against the MN-major path.tests/kernels/test_dsv41_mega_attn_layouts.py::test_quantized_o_proj_matches_quantized_bf16now covers 1, 129 and 8193 tokens (decode, partial and multiple 128-row tiles), including CUDA graph capture and replay. It checks that_o_proj's FP8 data and valid-row scales are byte-identical towo_bquantizing the BF16z. All 5 tests in the file pass on 1 x GB200 (torch 2.14.0+cu132, DeepGEMM with [DO NOT MERGE] Prefix sharing (bug fixed) #40):pytest tests/kernels/test_dsv41_mega_attn_layouts.py -v.pre-commit(ruff, format, typos, shellcheck, mypy) passes.Details
wo_bneeded a separate swizzle kernel. Below about 1024 tokens, that launch cost about as much as the quantization it saved. Write FP8 einsum output SFs in the 128x4 tiled layout DeepGEMM#40 changes only the scale address in the einsum epilogue, so no extra kernel is left.o_lora_rank % 128 == 0, andwo_bis an MXFP8 kernel that accepts F8_128x4 scales (can_fuse_query_quant). It also requires no GEMM-RS, because the fused GEMM + reduce-scatter takes a BF16zonly.wo_boutput rows that are never stored, so the test compares valid rows only.727f847, the merge commit of DeepGEMM#40 on the fork'sdevbranch. The einsum selects the 128x4 scales withuse_128x4_sfd_layout=True, which an older externally installeddeep_gemm(vLLM prefers site-packages) does not accept.mainto resolve conflicts.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code