From 4f9f45e4850e7e3a24382cd5cfc2697e7626d82e Mon Sep 17 00:00:00 2001 From: bofengl Date: Sun, 18 Oct 2026 01:57:17 -0700 Subject: [PATCH] [Perf][Kernel][DSv4.1] Fuse MoE top-k finalize with shared-expert add 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 (#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 Co-Authored-By: Claude Opus 5.5 --- tests/kernels/moe/test_ocp_mx_moe.py | 9 +- tests/kernels/test_dsv41_moe_finalize.py | 240 ++++++++++++++++++ .../deepseek_v41/decoder_replay_layers.py | 2 +- vllm/models/deepseek_v41/nvidia/model.py | 165 ++++++++++-- 4 files changed, 387 insertions(+), 29 deletions(-) create mode 100644 tests/kernels/test_dsv41_moe_finalize.py diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 968781a188de..4dbf74a19b9d 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -666,7 +666,14 @@ def test_trtllm_gen_mxfp4_fused_moe( @pytest.mark.parametrize( "monolithic,num_tokens,chunk_size", - [(False, 1, None), (False, 16, None), (False, 16, 5), (True, 16, None)], + [ + (False, 1, None), + (False, 16, None), + # Prefill-sized: DeepSeek V4.1 defers every call that fits one launch. + (False, 1000, None), + (False, 16, 5), + (True, 16, None), + ], ) @pytest.mark.skipif( not TRTLLM_GEN_MXFP4_AVAILABLE, diff --git a/tests/kernels/test_dsv41_moe_finalize.py b/tests/kernels/test_dsv41_moe_finalize.py new file mode 100644 index 000000000000..479ec558034b --- /dev/null +++ b/tests/kernels/test_dsv41_moe_finalize.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4.1's fused MoE top-k finalize + shared-expert add. + +The fused kernel replaces the TRT-LLM finalize and the shared-expert add that +follows it, so it must reproduce them bit for bit. ``moe_unpermute`` is that +finalize: ``check_deferred_moe_finalize`` holds it to the TRTLLM-Gen experts' +own finalized output. The dispatch tests need no GPU. +""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( + TrtLlmMxfp4ExpertsModular, + TrtLlmMxfp4ExpertsMonolithic, +) +from vllm.model_executor.layers.fused_moe.moe_output import ( + MoEOutput, + UnfinalizedMoEOutput, +) +from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( + moe_permute_unpermute_supported, + moe_unpermute, +) +from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MoE as DeepseekV4MoEBase +from vllm.models.deepseek_v41.nvidia.model import ( + DeepseekV4MoE, + _finalize_shared_add, +) +from vllm.models.deepseek_v41.nvidia.ops.mhc import MHC_OVERLAP_MAX_TOKENS +from vllm.platforms import current_platform + +MODEL_MODULE = "vllm.models.deepseek_v41.nvidia.model" +# Stands in for the experts' single-launch size, which caps every deferral. +MAX_LAUNCH_TOKENS = 1000 + + +@pytest.mark.parametrize("num_tokens", [1, 1000]) +# DeepSeek V4.1's shape, and a hidden size that ends in a partial block. +@pytest.mark.parametrize("hidden_size,top_k", [(5120, 6), (2880, 4)]) +def test_finalize_shared_add_matches_finalize_then_add( + num_tokens: int, hidden_size: int, top_k: int +): + if not current_platform.is_cuda() or not moe_permute_unpermute_supported(): + pytest.skip("moe_unpermute requires CUDA") + torch.manual_seed(0) + num_routes = num_tokens * top_k + # Rows past the routed ones stand in for the kernel's padding. + gemm2_permuted = torch.randn( + num_routes + 8, hidden_size, device="cuda", dtype=torch.bfloat16 + ) + permuted_idx = torch.randperm(num_routes, device="cuda").to(torch.int32) + permuted_idx = permuted_idx.view(num_tokens, top_k) + expert_weights = torch.rand(num_tokens, top_k, device="cuda").softmax(dim=-1) + shared_output = torch.randn( + num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16 + ) + # One route goes to another rank's expert (-1), which the finalize skips. + # moe_unpermute's input keeps a row for that route with zero weight + # instead, which leaves its sum unchanged. + unfinalized_idx = permuted_idx.clone() + unfinalized_idx[0, -1] = -1 + reference_weights = expert_weights.clone() + reference_weights[0, -1] = 0 + + routed_output = torch.empty_like(shared_output) + moe_unpermute( + routed_output, + gemm2_permuted, + reference_weights, + permuted_idx, + # The kernel reads its valid-row count through this pointer either way. + expert_first_token_offset=gemm2_permuted.new_full( + (1,), gemm2_permuted.shape[0], dtype=torch.int64 + ), + ) + expected = shared_output + routed_output + + actual = _finalize_shared_add( + MoEOutput( + routed=UnfinalizedMoEOutput( + gemm2_permuted, expert_weights, unfinalized_idx + ), + shared_output=shared_output, + ) + ) + torch.testing.assert_close( + actual.view(torch.int16), expected.view(torch.int16), atol=0, rtol=0 + ) + + +def _make_moe( + monkeypatch: pytest.MonkeyPatch, + *, + tp_size: int = 2, + use_native_mega_moe: bool = False, + has_shared_experts: bool = True, + experts_cls: type = TrtLlmMxfp4ExpertsModular, + routed_scaling_factor: float = 1.0, + reduce_results: bool = True, +) -> DeepseekV4MoE: + """A DeepSeek V4.1 MoE layer holding only what its deferral reads.""" + # Skip building the gate and the experts; the fields below stand in. + monkeypatch.setattr( + DeepseekV4MoEBase, + "__init__", + lambda self, *args, **kwargs: nn.Module.__init__(self), + ) + hf_config = SimpleNamespace( + n_routed_experts=8, num_experts_per_tok=2, num_hidden_layers=1 + ) + moe = DeepseekV4MoE( + SimpleNamespace(model_config=SimpleNamespace(hf_config=hf_config)), + prefix="model.layers.0.ffn", + ) + moe_config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=128, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + activation=MoEActivation.SILU, + device="cpu", + routing_method=RoutingMethodType.Renormalize, + moe_parallel_config=replace( + FusedMoEParallelConfig.make_no_parallel(), tp_size=tp_size + ), + in_dtype=torch.bfloat16, + ) + # The experts lower the deferral cap to their single-launch size. + moe_config.limit_deferred_moe_finalize(MAX_LAUNCH_TOKENS) + moe.use_native_mega_moe = use_native_mega_moe + moe.shared_experts = nn.Identity() if has_shared_experts else None + moe.reduce_results = reduce_results + moe.experts = SimpleNamespace( + moe_config=moe_config, + routed_experts=SimpleNamespace( + quant_method=SimpleNamespace(experts_cls=experts_cls) + ), + routed_scaling_factor=routed_scaling_factor, + ) + return moe + + +@pytest.mark.parametrize( + "case", + [ + "native_mega_moe", + "no_shared_experts", + "monolithic_experts", + "routed_scaling_factor", + "tp1", + ], +) +def test_ineligible_moe_keeps_runner_finalize(monkeypatch, case): + moe = _make_moe( + monkeypatch, + tp_size=1 if case == "tp1" else 2, + use_native_mega_moe=case == "native_mega_moe", + has_shared_experts=case != "no_shared_experts", + experts_cls=( + TrtLlmMxfp4ExpertsMonolithic + if case == "monolithic_experts" + else TrtLlmMxfp4ExpertsModular + ), + routed_scaling_factor=2.5 if case == "routed_scaling_factor" else 1.0, + ) + moe.defer_finalize(MHC_OVERLAP_MAX_TOKENS) + assert not moe.defers_finalize(1) + assert not moe.finalizes_in_next_layer(1) + + +# 0: the last local layer, a layer before an engram layer, or a model that does +# not fuse the mHC all-reduce. +@pytest.mark.parametrize("next_layer_max_tokens", [0, MHC_OVERLAP_MAX_TOKENS]) +def test_deferred_calls_fold_into_next_layer_up_to_its_limit( + monkeypatch, next_layer_max_tokens +): + moe = _make_moe(monkeypatch) + moe.defer_finalize(next_layer_max_tokens) + for num_tokens in ( + 0, + 1, + MHC_OVERLAP_MAX_TOKENS, + MHC_OVERLAP_MAX_TOKENS + 1, + MAX_LAUNCH_TOKENS, + MAX_LAUNCH_TOKENS + 1, + ): + defers = 0 < num_tokens <= MAX_LAUNCH_TOKENS + assert moe.defers_finalize(num_tokens) == defers + assert moe.finalizes_in_next_layer(num_tokens) == ( + defers and num_tokens <= next_layer_max_tokens + ) + + +@pytest.mark.parametrize("reduce_results", [False, True]) +@pytest.mark.parametrize( + "num_tokens", [MHC_OVERLAP_MAX_TOKENS + 1, MAX_LAUNCH_TOKENS + 1] +) +def test_forward_finalizes_deferred_calls_with_shared_add( + monkeypatch, reduce_results, num_tokens +): + moe = _make_moe(monkeypatch, reduce_results=reduce_results) + moe.defer_finalize(MHC_OVERLAP_MAX_TOKENS) + calls: list[str] = [] + + def record(name: str): + def call(*args): + calls.append(name) + return name + + return call + + monkeypatch.setattr(DeepseekV4MoEBase, "forward", record("runner")) + monkeypatch.setattr(moe, "forward_unfinalized", record("forward_unfinalized")) + monkeypatch.setattr( + f"{MODEL_MODULE}._finalize_shared_add", record("finalize_shared_add") + ) + monkeypatch.setattr( + f"{MODEL_MODULE}.tensor_model_parallel_all_reduce", record("all_reduce") + ) + moe(torch.empty(num_tokens, 128)) + if num_tokens <= MAX_LAUNCH_TOKENS: + expected = ["forward_unfinalized", "finalize_shared_add"] + expected += ["all_reduce"] if reduce_results else [] + else: + expected = ["runner"] + assert calls == expected diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py index 28aa5e399f63..87aa35958221 100644 --- a/vllm/models/deepseek_v41/decoder_replay_layers.py +++ b/vllm/models/deepseek_v41/decoder_replay_layers.py @@ -111,7 +111,7 @@ def _run( replay_batch = self.replay_batch if replay_batch is None: return self.run_layers(hidden_states, *states) - # Replay batch steps hold more than `window` tokens: no deferred MoE finalize. + # Replay batch steps hold more than `window` tokens: no unfinalized MoE output. assert isinstance(hidden_states, torch.Tensor) rows = replay_batch.rows inputs = (hidden_states, *states) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 84b3413eb66e..5f4472ace77c 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -3,7 +3,7 @@ import typing from collections.abc import Callable, Iterable from functools import partial -from itertools import islice +from itertools import islice, zip_longest import regex as re import torch @@ -95,6 +95,7 @@ from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.ubatching import dbo_current_ubatch_id @@ -123,6 +124,76 @@ logger = init_logger(__name__) +@triton.jit +def _finalize_shared_add_kernel( + perm_ptr, + w_ptr, + idx_ptr, + shared_ptr, + out_ptr, + H: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_H: tl.constexpr, +): + t = tl.program_id(0).to(tl.int64) + hb = tl.program_id(1) + cols = hb * BLOCK_H + tl.arange(0, BLOCK_H) + cmask = cols < H + acc = tl.zeros([BLOCK_H], dtype=tl.float32) + for k in tl.static_range(TOPK): + p = tl.load(idx_ptr + t * TOPK + k) + w = tl.load(w_ptr + t * TOPK + k).to(tl.float32) + # A negative row is an expert of another rank. + valid = p >= 0 + x = tl.load( + perm_ptr + p.to(tl.int64) * H + cols, mask=cmask & valid, other=0.0 + ).to(tl.float32) + acc = tl.where(valid, tl.fma(x, w, acc), acc) + fin = acc.to(tl.bfloat16) + s = tl.load(shared_ptr + t * H + cols, mask=cmask, other=0.0) + y = (s.to(tl.float32) + fin.to(tl.float32)).to(tl.bfloat16) + tl.store(out_ptr + t * H + cols, y, mask=cmask) + + +def _finalize_shared_add(output: MoEOutput) -> torch.Tensor: + """Top-k finalize of a deferred MoE output fused with the shared-expert add. + + Takes ``gemm2_permuted`` [num_permuted_rows, hidden] bf16, the float routing + weights and int32 permute map [num_tokens, top_k] (a negative index skips + the route) and ``shared_output`` [num_tokens, hidden] bf16, all contiguous, + with no routed scale left to apply. Returns [num_tokens, hidden] bf16 computed + like the TRTLLM-Gen finalize followed by the runner's ``shared + routed``: + per token, an fp32 FMA over its top-k rows in slot order, rounded to bf16, + then an fp32 add of the shared output, rounded to bf16. Given the same + GEMM2 output it matches those two kernels bit for bit (NaN payloads aside). + """ + routed, shared = output.routed, output.shared_output + assert isinstance(routed, UnfinalizedMoEOutput) and shared is not None + assert output.routed_scaling_factor == 1.0 + perm = routed.gemm2_permuted + weights = routed.expert_weights + idx = routed.expanded_idx_to_permuted_idx + num_tokens, hidden_size = shared.shape + assert perm.dtype == shared.dtype == torch.bfloat16 + assert perm.shape[1] == hidden_size + assert all(t.is_contiguous() for t in (perm, weights, idx, shared)) + out = torch.empty_like(shared) + block_h = 1024 + grid = (num_tokens, triton.cdiv(hidden_size, block_h)) + _finalize_shared_add_kernel[grid]( + perm, + weights, + idx, + shared, + out, + H=hidden_size, + TOPK=idx.shape[1], + BLOCK_H=block_h, + num_warps=4, + ) + return out + + class DeepseekV4MoE(DeepseekV4MoEBase): def __init__( self, @@ -151,35 +222,72 @@ def __init__( reduce_results=reduce_results, image_sentinel_lo=IMAGE_SENTINEL_BASE_ID, ) + self._defers_finalize = False + self._next_layer_finalize_max_tokens = 0 - def defer_finalize(self) -> None: - """Leave the routed top-k reduction to the next fused all-reduce + mHC. + def defer_finalize(self, next_layer_max_tokens: int) -> None: + """Stop the routed experts after GEMM2 and finalize outside them. - Only the modular TRTLLM MXFP4 experts can stop after GEMM2; the fused - kernel takes as many tokens as the mHC overlap path. + Only for the modular TRTLLM-Gen MXFP4 experts, whose finalize + ``_finalize_shared_add`` reproduces, on a TP-only deployment. Deferred + calls of up to ``next_layer_max_tokens`` tokens are finalized by the next + layer's fused all-reduce + mHC; ``forward`` finalizes the others. """ - moe_config = self.experts.moe_config + if self.use_native_mega_moe or self.shared_experts is None: + return experts_cls = getattr( self.experts.routed_experts.quant_method, "experts_cls", None ) if ( - experts_cls is TrtLlmMxfp4ExpertsModular - and self.experts.routed_scaling_factor == 1.0 + experts_cls is not TrtLlmMxfp4ExpertsModular + or self.experts.routed_scaling_factor != 1.0 ): - moe_config.defer_moe_finalize(MHC_OVERLAP_MAX_TOKENS) - if moe_config.use_deferred_moe_finalize: - logger.info_once( - "DSV4.1 mHC: MoE top-k finalize fused into the all-reduce for " - "up to %d tokens.", - MHC_OVERLAP_MAX_TOKENS, - ) + return + moe_config = self.experts.moe_config + moe_config.defer_moe_finalize() + self._defers_finalize = moe_config.use_deferred_moe_finalize + if not self._defers_finalize: + return + self._next_layer_finalize_max_tokens = next_layer_max_tokens + logger.info_once("DSV4.1: MoE top-k finalize fused with the shared-expert add.") + if next_layer_max_tokens > 0: + logger.info_once( + "DSV4.1 mHC: MoE top-k finalize fused into the all-reduce for " + "up to %d tokens.", + next_layer_max_tokens, + ) def defers_finalize(self, num_tokens: int) -> bool: + """Whether the routed experts stop after GEMM2 for this call.""" return ( - not self.use_native_mega_moe + self._defers_finalize and self.experts.moe_config.should_defer_moe_finalize(num_tokens) ) + def finalizes_in_next_layer(self, num_tokens: int) -> bool: + """Whether the next layer's fused all-reduce + mHC finalizes this call.""" + return ( + num_tokens <= self._next_layer_finalize_max_tokens + and self.defers_finalize(num_tokens) + ) + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None = None, + mega_gate_metadata: MegaGateRoutingMetadata | None = None, + ) -> torch.Tensor: + if not self.defers_finalize(hidden_states.shape[0]): + return super().forward(hidden_states, input_ids, mega_gate_metadata) + # Where the experts defer, the runner all-reduces exactly when + # reduce_results is set. + output = _finalize_shared_add( + self.forward_unfinalized(hidden_states, input_ids) + ) + if self.reduce_results: + output = tensor_model_parallel_all_reduce(output) + return output + def forward_unfinalized( self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None ) -> MoEOutput: @@ -643,7 +751,7 @@ def forward( stream=mhc_stream, reduce_results=self.fuse_mhc_all_reduce, ) - if self.fuse_mhc_all_reduce and self.ffn.defers_finalize(x.shape[0]): + if self.fuse_mhc_all_reduce and self.ffn.finalizes_in_next_layer(x.shape[0]): # The next layer's first boundary finalizes it with the all-reduce. x = self.ffn.forward_unfinalized(x, input_ids) else: @@ -829,16 +937,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ], vllm_config.scheduler_config.max_num_batched_tokens, ) - if self.fuse_mhc_all_reduce: - # A MoE's top-k finalize folds into the next layer's first mHC - # boundary, which the last local layer lacks and an engram layer - # replaces with its own all-reduce. - local_layers = list(islice(self.layers, self.start_layer, self.end_layer)) - for layer, successor in zip(local_layers, local_layers[1:]): - assert isinstance(layer, DeepseekV4DecoderLayer) - assert isinstance(successor, DeepseekV4DecoderLayer) - if successor.engram is None: - layer.ffn.defer_finalize() + # A MoE's top-k finalize folds into the next layer's first mHC + # boundary, which the last local layer lacks and an engram layer + # replaces with its own all-reduce. The MoE layer finalizes every other + # deferred call itself, together with the shared-expert add. + local_layers = list(islice(self.layers, self.start_layer, self.end_layer)) + for layer, successor in zip_longest(local_layers, local_layers[1:]): + assert isinstance(layer, DeepseekV4DecoderLayer) + folds = ( + self.fuse_mhc_all_reduce + and isinstance(successor, DeepseekV4DecoderLayer) + and successor.engram is None + ) + layer.ffn.defer_finalize(MHC_OVERLAP_MAX_TOKENS if folds else 0) # Decoder-side SWA bounded replay: in eager prefill steps the layers past # the last KV source run on each request's trailing window only