[K3 Perf] Fuse MXFP4 top-k finalization into latent-tail, ~5% E2E latency reduction - #53152
Conversation
Signed-off-by: yewentao256 <zhyanwentao@126.com>
![claude[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
|
/ci run |
|
✅ Triggered Buildkite CI #84873 for commit |
yewentao256
changed the title
mgoin
changed the title
yewentao256
changed the title
mgoin
left a comment
mgoin
left a comment
There was a problem hiding this comment.
The optimization looks nice, but could we clean up the output contract before merging?
trtllm_fp4_block_scale_moe now uses both output (the finalized destination) and result (a mode dependent list), which makes the data flow difficult to follow. Please rename these clearly, immediately decompose the deferred outputs, and ideally centralize the FlashInfer return conversion in a shared helper.
Also, UnfinalizedMoEOutput now passes through APIs still typed as tensor-only. Those annotations/guards should be updated, with targeted finalized-vs-deferred parity tests added. The current structure feels too fragile to future changes.
zyongye
left a comment
zyongye
left a comment
There was a problem hiding this comment.
I will merge this first and doing cleaning up and some perf tuning in a separate PR. Thanks for the effort.
|
I'm so excited! |
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
|
/ci run |
|
✅ Triggered Buildkite CI #85008 for commit |
|
Note: current ci failure buildkite/ci/pr/nvidia-h200-quantized-models is not related, happening in other PRs as well |
|
/ci run |
|
✅ Triggered Buildkite CI #85035 for commit |
|
Accuracy validation of this change on GB300:
Kimi-K3 MXFP4, GB300 4×GPU/node, AI assistance was used to produce these runs. |
Extend the deferred MoE finalize protocol (#53152) to the modular prepare/finalize path so the TRTLLM FP8 block-scale and NVFP4 routed kernels can stop after GEMM2, and add a FlashInfer fused finalize + all-reduce + RMSNorm consumer for MiniMax M3. Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Two places went their own way instead of using what #53152 established: - The TRTLLM FP8 block-scale and NVFP4 routed kernels hand-built an UnfinalizedMoEOutput from the raw FlashInfer return. Route both through convert_flashinfer_moe_output, as every other TRTLLM expert does. It validates the deferred layout, and it absorbs FlashInfer's coming switch of the finalized return from a Tensor to List[Tensor]. - The consumer's workspace capacity travelled on a new MoEOutput field. The protocol already has a channel for it -- the consumer declares FusedMoEConfig.defer_moe_finalize_max_num_tokens at build time and the producer gates on should_defer_moe_finalize -- so use that and leave MoEOutput alone. This also moves the workspace build out of the forward-time custom op, where it ran a collective without the vLLM config, into the consuming layer's constructor: an unsupported (tp_size, hidden, top_k, dtype) is now a build-time fallback to finalizing in the MoE kernel rather than a forward-time assert. Signed-off-by: Yongye Zhu <yongye@inferact.ai>
…ency reduction (vllm-project#53152) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
…llm-project#53152 (vllm-project#53310) Signed-off-by: yewentao256 <zhyanwentao@126.com>
…ency reduction (vllm-project#53152) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: mikeshawcode <michaelwshaw2@gmail.com>
…llm-project#53152 (vllm-project#53310) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: mikeshawcode <michaelwshaw2@gmail.com>
Extend the deferred MoE finalize protocol (vllm-project#53152) to the modular prepare/finalize path so the TRTLLM FP8 block-scale and NVFP4 routed kernels can stop after GEMM2, and add a FlashInfer fused finalize + all-reduce + RMSNorm consumer for MiniMax M3. Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Two places went their own way instead of using what vllm-project#53152 established: - The TRTLLM FP8 block-scale and NVFP4 routed kernels hand-built an UnfinalizedMoEOutput from the raw FlashInfer return. Route both through convert_flashinfer_moe_output, as every other TRTLLM expert does. It validates the deferred layout, and it absorbs FlashInfer's coming switch of the finalized return from a Tensor to List[Tensor]. - The consumer's workspace capacity travelled on a new MoEOutput field. The protocol already has a channel for it -- the consumer declares FusedMoEConfig.defer_moe_finalize_max_num_tokens at build time and the producer gates on should_defer_moe_finalize -- so use that and leave MoEOutput alone. This also moves the workspace build out of the forward-time custom op, where it ran a collective without the vLLM config, into the consuming layer's constructor: an unsupported (tp_size, hidden, top_k, dtype) is now a build-time fallback to finalizing in the MoE kernel rather than a forward-time assert. Signed-off-by: Yongye Zhu <yongye@inferact.ai>
4 participants
Purpose
Part of #50587
Fuse MXFP4 top-k finalization into the Kimi K3 latent-tail kernel.
Before
MXFP4 MoE kernel │ ├─ GEMM2 │ └─finalize kernel ├─ unpermute ├─ times router weight ├─ top-k reduction └─ Write [M, 3584] tensor │ ▼ latent tail reads the tensor └─ AllReduce + RMSNorm + Up Projection + shared expertNow
MXFP4 MoE kernel（do_finalize=False） │ └─ return： ├─ GEMM2 output ├─ router weights └─ permutation map │ latent tail reads the tensor ├─ times router weight and topk ├─ AllReduce + RMSNorm + Up Projection + shared expertThis removes one kernel launch and avoids writing and rereading the finalized intermediate tensor.
Test
Acc
Perf
TTFT remains unchanged
output_vllm_kimik3_8k1k_0817_no_spec.json
output_vllm_kimik3_8k1k.json