diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 94741c3e29be..a19cbaf7d767 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 443bf5975bb9..bf61b2980e3a 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: @@ -832,16 +940,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