Skip to content

[Perf][Kernel][DSv4.1] Fuse MoE top-k finalize with shared-expert add - #62635

Closed
roborluo wants to merge 2 commits into
vllm-project:mainfrom
roborluo:dsv41-finalize-shared-add
Closed

roborluo wants to merge 2 commits into
vllm-project:mainfrom
roborluo:dsv41-finalize-shared-add

Conversation

@roborluo

@roborluo roborluo commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Overview

Makes DeepSeek-V4.1's MoE top-k finalize + shared-expert add 1.3-2.3x faster by fusing them into one Triton kernel, which makes the TP2 model step 2.1% faster in decode and 2.5% faster in a 4096-token prefill, with unchanged accuracy. On TP-only
deployments with the modular TRTLLM-Gen MXFP4 experts, every MoE call that the next layer's fused all-reduce + mHC does not take now runs its deferred top-k finalize and the shared-expert add as one kernel, instead of the experts' finalize followed by
a separate add.
DeepSeek V4.1 (#57448), building on #58635 and #58586. Related but not overlapping: #61868, #60355, #62174, #62366 (see Details).

Claims

  • 1.3-2.3x faster finalize + shared add: 1.3-1.4x at 1-64 tokens, 2.0x at 1024 and 2.3x at 8192 (DeepSeek V4.1's hidden 5120, top-6; table below). One kernel replaces two, and the routed [num_tokens, hidden] bf16 intermediate is no longer written
    and read back. This covers every deferrable call on deployments without the fused all-reduce + mHC, and elsewhere the last local layer, layers before an engram layer and calls above MHC_OVERLAP_MAX_TOKENS (16) tokens, such as prefill.
  • 2.1-2.5% faster TP2 model step: in vllm bench forward, decode (8 requests at 8192 context) is 2.1% faster and a 4096-token prefill 2.5% faster.
  • Unchanged results: the fusion itself changes no bits: given the same GEMM2 output, the kernel is bit-identical to the TRTLLM-Gen finalize followed by the runner's shared_output + routed_output (tested on the raw bits). GSM8K: 0.894 on the base
    commit, 0.911 with this PR (one run each).
  • No new flag or config. Other experts, deployments, platforms and the DSpark drafter's layers keep the runner's path.

Validation

Tests

Command Result
pytest -v -s tests/kernels/test_dsv41_moe_finalize.py (dispatch tests run on CPU) 15 passed
pytest -v -s "tests/kernels/moe/test_ocp_mx_moe.py::test_trtllm_mxfp4_deferred_finalize" 5 passed (4 before this PR)
pytest -v -s tests/models/test_deepseek_v4_mega_moe.py (builds DeepseekV4MoE) 24 passed
  • test_finalize_shared_add_matches_finalize_then_add: the kernel against moe_unpermute plus the shared add at atol=rtol=0 on the int16 view, for 1 and 1000 tokens, DeepSeek V4.1's shape (hidden 5120, top-6) and a hidden size ending in a partial
    block (2880, top-4), with one route to another rank's expert.
  • test_ineligible_moe_keeps_runner_finalize, test_deferred_calls_fold_into_next_layer_up_to_its_limit, test_forward_finalizes_deferred_calls_with_shared_add: dispatch on a real FusedMoEConfig (CPU).
  • test_trtllm_mxfp4_deferred_finalize[False-1000-None] (new case): the existing zero-tolerance check of moe_unpermute against the experts' own finalize, now at a prefill-sized call, since calls above 16 tokens now defer too.

Kernel microbenchmark (hidden 5120, top-6; python bench_dsv41_finalize_shared_add.py below: FlashInfer's timing, which falls back to CUDA events without cupti-python, median; two runs agree within 6%). moe_unpermute runs the TRTLLM-Gen
finalize's algorithm and matches it bit for bit, so it stands in for the finalize that runs inside FlashInfer; the script asserts both arms are bit-identical before timing.

num_tokens Speedup (moe_unpermute + add / fused)
1 1.33x
4 1.36x
64 1.38x
256 1.67x
1024 1.98x
4096 2.25x
8192 2.31x
bench_dsv41_finalize_shared_add.py
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""DSv4.1 fused top-k finalize + shared add vs. moe_unpermute, then add.

moe_unpermute runs the TRTLLM-Gen finalize's algorithm and matches it bit for
bit (check_deferred_moe_finalize), so it stands in for the finalize that the
experts run inside FlashInfer. Timed with FlashInfer CUPTI.
"""

import statistics

import torch
from flashinfer.testing import bench_gpu_time_with_cupti

from vllm.model_executor.layers.fused_moe.moe_output import (
    MoEOutput,
    UnfinalizedMoEOutput,
)
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import moe_unpermute
from vllm.models.deepseek_v41.nvidia.model import _finalize_shared_add

HIDDEN, TOP_K = 5120, 6
WARMUP = 25


def bench_us(fn) -> float:
    for _ in range(WARMUP):
        fn()
    torch.accelerator.synchronize()
    return statistics.median(bench_gpu_time_with_cupti(fn)) * 1e3


def bench_case(num_tokens: int) -> tuple[float, float, int]:
    routes = num_tokens * TOP_K
    gemm2 = torch.randn(routes, HIDDEN, dtype=torch.bfloat16)
    idx = torch.randperm(routes).int().view(num_tokens, TOP_K)
    weights = torch.rand(num_tokens, TOP_K).softmax(dim=-1)
    shared = torch.randn(num_tokens, HIDDEN, dtype=torch.bfloat16)
    routed = torch.empty_like(shared)
    num_rows = torch.full((1,), routes, dtype=torch.int64)
    deferred = MoEOutput(UnfinalizedMoEOutput(gemm2, weights, idx), shared)

    def base():
        moe_unpermute(routed, gemm2, weights, idx, expert_first_token_offset=num_rows)
        return shared + routed

    def fused():
        return _finalize_shared_add(deferred)

    assert torch.equal(base().view(torch.int16), fused().view(torch.int16))
    # Each top-k row, weight and index once, the shared output in, the sum out.
    nbytes = routes * (HIDDEN * 2 + 8) + 2 * num_tokens * HIDDEN * 2
    return bench_us(base), bench_us(fused), nbytes


def main() -> None:
    torch.set_default_device("cuda")
    torch.manual_seed(0)
    print(f"{'tokens':>6} {'base_us':>8} {'fused_us':>8} {'fused_GBps':>10}")
    for num_tokens in (1, 4, 16, 64, 256, 1024, 4096, 8192):
        base_us, fused_us, nbytes = bench_case(num_tokens)
        gbps = nbytes / fused_us / 1e3
        print(f"{num_tokens:>6} {base_us:>8.2f} {fused_us:>8.2f} {gbps:>10.1f}")


if __name__ == "__main__":
    main()

End to end, DeepSeek-V4.1-Flash (MXFP4 experts) at TP2, base commit aa3564ce vs. this PR with identical flags. The in-tree DeepSeek V4.1 GSM8K config serves the NVFP4 checkpoint, whose experts do not take this path, so it is run here against the
MXFP4 checkpoint. The accuracy server runs with VLLM_GPU_SYNC_CHECK=error, to check that the new path never syncs the host with the device, and with --jit-monitor-mode warn, whose log must not list _finalize_shared_add_kernel among the kernels
compiled after warmup. (error mode is not usable here: on main, the DSA indexer's _block_scores_kernel already compiles at the first request, which stops the server.)

VLLM_GPU_SYNC_CHECK=error vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --tokenizer-mode deepseek_v41 --tensor-parallel-size 2 \
  --engram-config '{"cpu_offload":true}' --max-model-len 8192 \
  --max-num-batched-tokens 4096 --jit-monitor-mode warn --port 8000
python tests/evals/gsm8k/gsm8k_eval.py --port 8000 --num-questions 1319 --num-shots 5

# Per-step latency and per-kernel deltas: --profile-dir prof_base on the base
# commit, prof_pr on this PR.
torchrun --nproc-per-node=2 -m vllm.entrypoints.cli.main bench forward \
  --model deepseek-ai/DeepSeek-V4.1-Flash --tokenizer-mode deepseek_v41 -tp 2 \
  --engram-config '{"cpu_offload":true}' --max-model-len 8192 \
  --max-num-batched-tokens 4096 --case decode=8*1@8192 --case prefill=4096 \
  --profile-dir prof_pr
python -m vllm.benchmarks.forward.profile prof_base/prefill.rank0.kernels.json \
  prof_pr/prefill.rank0.kernels.json
  • Model step time (bench forward, median): decode 8*1@8192 -2.1%, prefill 4096 -2.5%.
  • Where it comes from (per-kernel profile, share of the step): in the prefill step, the experts' finalize kernel (3.3%) and the elementwise shared adds (0.7%) are replaced by _finalize_shared_add_kernel (1.6%); in the decode step, the finalize
    (1.7%) and the adds (0.9%) by the fused kernel (1.2%). The GEMM2 kernel switches to its GEMM2-only variant at the same cost.
  • Accuracy and checks: GSM8K 5-shot, all 1319 questions: 0.894 on the base commit, 0.911 with this PR (one run each). VLLM_GPU_SYNC_CHECK=error detected no host sync in either run, and the JIT monitor listed only the DSA indexer's
    _block_scores_kernel (already on main) as compiled after warmup in both runs, never _finalize_shared_add_kernel.

Details

Motivation. Since #58635 the TRTLLM-Gen MXFP4 modular experts can stop after GEMM2 and return an UnfinalizedMoEOutput. On DeepSeek V4.1 the only consumer is the fused all-reduce + mHC boundary (#58586), for calls of up to MHC_OVERLAP_MAX_TOKENS
tokens on deployments where fuse_mhc_all_reduce holds. Every other call runs the experts' finalize kernel and then the runner's elementwise shared_output + routed_output.

Design.

  • _finalize_shared_add_kernel: a (num_tokens, cdiv(hidden, 1024)) grid; each program reduces one token's top-k GEMM2 rows over 1024 columns, adds the shared output and stores bf16.
  • The model calls DeepseekV4MoE.defer_finalize(next_layer_max_tokens) on every local layer, passing MHC_OVERLAP_MAX_TOKENS when the model fuses the mHC all-reduce and the next local layer has no engram, else 0 (the last local layer, a layer before
    an engram layer, or a model without that fusion). An eligible layer then defers every call that fits one experts launch.
  • The decoder layer keeps handing the calls the next layer takes to forward_unfinalized, now via finalizes_in_next_layer. DeepseekV4MoE.forward finalizes every other deferred call with the kernel and all-reduces when reduce_results is set, which
    is exactly when the runner's late all-reduce would run on these deployments.
  • The decision depends only on the static token count, so each CUDA graph captures one path. Nothing on the new path reads GPU tensors on the host.
  • The consumer stays in the model: the shared MoERunner is unchanged, and the kernel depends only on MoEOutput / UnfinalizedMoEOutput, so it can move next to them if another model with shared experts wants the same consumer. No existing kernel
    fuses the finalize with the shared add outside an all-reduce (moe_unpermute and the experts' finalize stop at the top-k sum; the LoRA experts' _finalize_lora_kernel adds the W2 LoRA delta to the unrounded sum instead).

Tradeoffs.

  • The kernel keeps the runner path's two roundings (finalize to bf16, then the bf16 add), so for a given GEMM2 output, moving a call to it changes no bits. The fused all-reduce + mHC finalize, which rounds once, is untouched.
  • With deferral no longer capped at 16 tokens, the deferred FlashInfer autotune pass in kernel_warmup coincides with the max_num_batched_tokens pass whenever that fits one experts launch: deployments that already deferred run one warmup pass fewer,
    and the remaining pass tunes the GEMM2-only variant that eligible calls now use. Bit-exactness holds given the same GEMM2 output, and that variant is tuned on its own, which is why end-to-end accuracy is checked with GSM8K rather than token by token.
  • The kernel's constexprs depend only on the model (hidden size, top-k), so the startup profile run compiles it whenever max_num_batched_tokens fits one experts launch; the JIT monitor log above checks that it never compiles later.

Limitations. TP1 keeps the runner path (use_deferred_moe_finalize needs TP > 1), as do calls split across experts launches, the DSpark drafter's MoE layers (they never request deferral) and every other experts class, including the NVFP4 TRTLLM-Gen
experts.

Compatibility.

  • Backends: only TrtLlmMxfp4ExpertsModular (an identity check, so the LoRA experts class is excluded) with no routed scale left in the runner and with shared experts; native MegaMoE and every other experts class are unchanged.
  • Architectures: these experts run on SM100-family GPUs only; other platforms and the AMD model (a separate file) are unchanged.
  • Parallelism: TP-only (no DP, EP, PCP or sequence parallelism) and no hidden-dim padding, via use_deferred_moe_finalize. The split keys on fuse_mhc_all_reduce, so it holds for whatever TP sizes supports_mhc_all_reduce admits (TP4 on main; [Perf][DSV4.1] Enable the fused all-reduce + mHC boundary at TP2 #60355
    proposes TP2). With pipeline parallelism, each stage's last local layer finalizes in the MoE layer.
  • CUDA graphs and torch.compile: static dispatch per captured size; the DeepSeek V4.1 NVIDIA model is not torch.compiled, and deferred calls already bypass the runner's custom op through forward_unfinalized, as the fused all-reduce + mHC path does on
    main.
  • Batch invariance: under VLLM_BATCH_INVARIANT=1 the TRTLLM-Gen MXFP4 experts are not selected, so this path is not taken; the kernel itself reduces each token in a fixed order without atomics.

CI coverage. tests/kernels/test_dsv41_moe_finalize.py runs in ":nvidia: (B200) Miscellaneous Kernels". The new test_ocp_mx_moe.py case runs in ":nvidia: (B200) Kernels" shard 0, which this diff does not trigger (the L4 MoE job skips SM100-only
tests), so it is run manually above. The optional ":nvidia: (B200) LM Eval Spec Decode" job serves the NVFP4 checkpoint at TP4, which covers the unchanged runner path only; no CI job serves DeepSeek V4.1 with MXFP4 experts at TP > 1, so the end-to-end
runs above are manual.

Not a duplicate. Searches of open PRs and PRs since 2026-08-15 for "finalize shared expert", "defer finalize" and "MoE finalize shared add" found no PR that fuses the deferred TRTLLM-Gen top-k finalize with the shared-expert add on DeepSeek V4.1's
default path.

AI assistance. This change, its tests and this description were prepared with AI assistance (Claude Code). The submitter reviewed every changed line and ran the commands in Validation.


Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (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.

roborluo and others added 2 commits October 18, 2026 22:28
On TP-only deployments with the modular TRTLLM-Gen MXFP4 experts, a MoE
call that the next layer's fused all-reduce + mHC does not take runs the
experts' top-k finalize and then a separate shared_output +
routed_output: two passes over the hidden states. These experts can
already stop after GEMM2 (vllm-project#58635), so the MoE layer now finalizes such
calls itself, in one Triton kernel that also adds the shared-expert
output.

- _finalize_shared_add_kernel reduces each token's top-k rows of the
  permuted GEMM2 output with their routing weights, adds the
  shared-expert output and writes the bf16 sum.
- The model calls DeepseekV4MoE.defer_finalize(next_layer_max_tokens)
  for every local layer. An eligible layer defers every call that fits
  one experts launch and records how many tokens the next layer's fused
  all-reduce + mHC takes: MHC_OVERLAP_MAX_TOKENS when the model fuses
  the mHC all-reduce and the successor is a local layer without engram,
  0 otherwise.
- The decoder layer still hands the calls the next layer takes to
  forward_unfinalized. DeepseekV4MoE.forward finalizes every other
  deferred call with the kernel, then all-reduces when reduce_results
  is set, as the runner does.

Dispatch is automatic, with no new flag. A call takes the kernel when
the experts are TrtLlmMxfp4ExpertsModular (not native MegaMoE), the
runner applies no routed scale to their output and the layer has shared
experts; the deployment can defer (use_deferred_moe_finalize: TP > 1
without DP, EP, PCP or sequence parallelism, no hidden-dim padding);
the call is non-empty and fits one experts launch; and the next layer's
fused all-reduce + mHC does not take it: the model does not fuse that
all-reduce, the layer is the last local one or precedes an engram
layer, or the call has more than MHC_OVERLAP_MAX_TOKENS tokens. The
choice depends only on the token count, so each CUDA graph captures one
path and the host never waits on the device. Other experts, deployments
and the DSpark drafter's layers keep the runner's path, and the calls
handed to the next layer unfinalized are unchanged. The MoE before the
decoder replay layers may now defer, so their input comment now says
that input is never an unfinalized MoE output.

Exactness: per element, the kernel does the finalize's arithmetic in
the same order, an FP32 FMA over the top-k slots in slot order from
+0.0 that skips routes to other ranks, rounded to bf16, then the
runner's FP32 shared + routed add, rounded to bf16. Given the same
GEMM2 output the result is bit-identical (NaN payloads aside).

With deferral no longer capped at MHC_OVERLAP_MAX_TOKENS, the deferred
FlashInfer autotune pass in kernel_warmup coincides with the
max_num_batched_tokens pass whenever that fits one experts launch, and
tunes the GEMM2-only variant. The kernel's constexprs depend only on
the model, so the startup profile run compiles it whenever
max_num_batched_tokens fits one experts launch.

Tests: the kernel is checked bit for bit against moe_unpermute plus
the shared add, including a route to another rank; dispatch tests on a
real FusedMoEConfig cover eligibility, the next layer's token limit and
forward's finalize + all-reduce; and test_trtllm_mxfp4_deferred_finalize
gains a 1000-token case, holding moe_unpermute to the experts' own
finalize at zero tolerance for prefill-sized calls.

Signed-off-by: bofengl <bofengl@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models labels Oct 20, 2026
@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

@LopezCastroRoberto LopezCastroRoberto left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since MHC_OVERLAP_MAX_TOKENS = 16, decode batches <=16 already fold the finalize into the fused all-reduce on all but the last/pre-engram layers, so I’d expect ~no change there. If there is a a gain at batch <= 16, it’s either noise or something I’ve missed.

Could you share the actual decode step latency (in ms, before and after) that the −2.1% comes from, and at which batch size it was measured? I generally recommend to repeat e2e runs 3-4 times.

@roborluo

Copy link
Copy Markdown
Contributor Author

Thanks for the careful read — the ≤16-token fold you describe is real, but it doesn't cover the configuration the −2.1% came from.

The decode number was measured at TP2 (the bench forward command in the description, -tp 2). On main the fold requires fuse_mhc_all_reduce, and supports_mhc_all_reduce admits only tensor_parallel_size == 4 (plus MNNVL Lamport multicast on the custom all-reduce, no EP, hidden 5120, hc_mult 4) — vllm/models/deepseek_v41/nvidia/ops/mhc.py. At TP2 that gate is false, defer_finalize() is never called on any layer, and no call at any batch size defers: every MoE layer runs the experts' finalize kernel plus the separate shared-expert add. Those are the two kernels this PR fuses, so at TP2 the fused kernel is on the path in every decode step, batch 8 included — MHC_OVERLAP_MAX_TOKENS never engages there.

Where the fold does exist (TP4 + MNNVL), this PR leaves it alone: a deferred call still routes through forward_unfinalized into the next layer's fused all-reduce + mHC, and the new kernel runs only where main also finalizes separately — the last local layer, pre-engram layers, and steps above 16 tokens. So for TP4 decode batches ≤16 I agree with you: nearly no change — the fold itself is untouched, and only the layers main also finalizes separately (the last local layer and pre-engram layers) switch to the fused kernel.

On the measurement: agreed that percentages from single runs aren't enough. I'll rerun the decode case (decode=8*1@8192, i.e. 8 requests × 1 decode token at 8192 context — the decode step is 8 tokens) 3–4 times per arm, base commit vs this PR, and post absolute median step times in ms, plus — on a TP4 setup where the fused all-reduce + mHC actually engages (I'll check its init in the log, since without MNNVL multicast TP4 behaves like TP2 here) — a ≤16-token pair (expected ~flat) and a larger-batch point as cross-checks.

@roborluo

Copy link
Copy Markdown
Contributor Author

Absolute numbers, 4 repeats per arm — and you were right to push back on the result, though for a different reason than the ≤16 fold.

Setup. B300 (2 GPUs, TP2), vLLM at commit 00b6c4c; arm A stock, arm B = this PR's runtime hunks applied on top (they apply cleanly). vllm bench forward with the command from the description, arms alternating back to back on the same node, 4 runs each. Values are the per-run median model step time; run-to-run spread was ≤0.1% on decode and ≤0.3% on prefill.

Case base (ms) PR (ms) Δ
decode 1*1@8192 (1-token step) 4.742 4.715 −0.6%
decode 8*1@8192 (8-token step) 7.349 7.429 +1.1%
decode 16*1@8192 9.615 9.764 +1.5%
decode 32*1@8192 13.704 13.871 +1.2%
prefill 4096 62.19 68.42 +10.0%

So on this machine at today's main, the PR is a net regression, and the −2.1%/−2.5% from the description (measured on a different SM10x system) does not reproduce.

Where it goes. A profiled pair (--profile-dir) shows the fused kernel itself does what it claims — _finalize_shared_add_kernel 1.22 ms vs the stock finalizeKernelVecLoad 1.51 ms over the 40 prefill MoE calls (decode: 0.15 vs 0.18 ms) — but the deferred raw-row expert call sits outside the FlashInfer autotuner's tuned buckets: the PR arm logs 28 No tuned config covers … falling back … tactic=-1 warnings vs 12 on base, with a new (M, 0), (0,), (M, 6), (M, 6), (M, 5120), (M, 160) shape family, and the expert GEMMs dispatch to different, slower variants — one grouped GEMM alone goes 6.5 → 14.5 ms in the profiled prefill. The handoff costs ~8 ms where the fusion saves ~0.3 ms.

Next. I'm converting this to a draft until it's reworked so the fused finalize consumes the experts' output on their already-tuned path (or the deferred-call shapes are covered by the tuning buckets), with the A/B above re-run as the acceptance gate. Thanks for insisting on the absolute numbers — the percentages hid exactly this.

@roborluo
roborluo marked this pull request as draft October 28, 2026 19:43
@roborluo roborluo closed this Nov 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants