[roborluo]

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 (vllm-project#57448), building on vllm-project#58635 and vllm-project#58586. Related but not overlapping: vllm-project#61868, vllm-project#60355, vllm-project#62174, vllm-project#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 vllm-project#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 (vllm-project#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 vllm-project/vllm#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.

• [Perf][DSv4.1] DeepSeek-V4.1-Flash specialized decode on GB300 vllm-project/vllm#61868 adds an opt-in decode executor (--additional-config '{"decode_backend": "latency_specialized"}', GB300 + TP4 + DSpark only) whose own fused all-reduce + mHC boundary takes the deferred finalizes of decode steps. This PR changes the default
  path on every eligible deployment and step, prefill included, and does the finalize and the shared add in one kernel inside the MoE layer, without fusing an all-reduce. The two merge cleanly in vllm/models/deepseek_v41/nvidia/model.py (git merge-file). [Perf][DSv4.1] DeepSeek-V4.1-Flash specialized decode on GB300 vllm-project/vllm#61868's MoE step saves and restores the deferral fields of FusedMoEConfig around its own calls and calls ffn(...) for layers it cannot defer, which keeps working with this change.
• [Perf][DSV4.1] Enable the fused all-reduce + mHC boundary at TP2 vllm-project/vllm#60355 lets the fused all-reduce + mHC run at TP2; it changes which calls fold into that boundary, and this PR handles the rest either way.
• [Perf] Fuse GLM-5.x MoE finalize into the next RMSNorm's all-reduce vllm-project/vllm#62174 (GLM-5.x) and [Model][Perf] Fuse Beam's TP all-reduce with its sandwich norms and MoE finalize vllm-project/vllm#62366 (Beam) fuse their models' MoE finalize into an all-reduce + norm; different models and a different fusion.

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.