From b5f08f636d26ea1535ee340974fc7ce4f02aca20 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Thu, 24 Sep 2026 22:49:38 +0000 Subject: [PATCH 1/8] [MoE] Defer the TRTLLM-Gen top-k finalize on the modular path The monolithic TRTLLM-Gen experts can already stop after GEMM2 and hand an UnfinalizedMoEOutput to a consumer that fuses the top-k reduction. Extend that to the modular path for the MXFP4, NVFP4 and FP8 experts: they defer when the layer asks for it and the tokens fit in one kernel launch, since each launch permutes into its own buffer. The modular kernel passes the result through when its prepare/finalize does nothing but the top-k reduction, as MoEPrepareAndFinalizeNoDPEPModular opts in to, and FusedMoEKernel.apply returns it the way apply_monolithic does. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yongye Zhu --- tests/kernels/moe/test_flashinfer.py | 91 +++++++++++- tests/kernels/moe/test_ocp_mx_moe.py | 130 +++++++++++++++++- tests/kernels/moe/test_trtllm_nvfp4_moe.py | 130 ++++++++++++++---- tests/kernels/moe/utils.py | 35 +++++ .../fused_moe/experts/trtllm_fp8_moe.py | 23 +++- .../fused_moe/experts/trtllm_mxfp4_moe.py | 33 +++-- .../fused_moe/experts/trtllm_nvfp4_moe.py | 30 +++- .../layers/fused_moe/fused_moe_method_base.py | 4 +- .../layers/fused_moe/modular_kernel.py | 51 +++++-- .../layers/fused_moe/moe_output.py | 11 +- .../fused_moe/prepare_finalize/no_dp_ep.py | 5 + .../layers/fused_moe/routed_experts.py | 4 +- .../compressed_tensors_moe_w4a4_mxfp4.py | 3 +- .../compressed_tensors_moe_w4a4_nvfp4.py | 2 +- .../compressed_tensors_moe_w8a8_fp8.py | 3 +- .../compressed_tensors_moe_w8a8_mxfp8.py | 3 +- .../model_executor/layers/quantization/fp8.py | 3 +- .../quantization/inc/schemes/inc_mxfp4_moe.py | 3 +- .../quantization/inc/schemes/inc_mxfp8_moe.py | 3 +- .../layers/quantization/modelopt.py | 6 +- .../layers/quantization/mxfp4.py | 4 +- .../layers/quantization/online/moe_base.py | 2 +- .../layers/quantization/quark/quark_moe.py | 6 +- 23 files changed, 496 insertions(+), 89 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 6329780f9dc9..2429f3e5d063 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -1,14 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing -from dataclasses import dataclass +from dataclasses import dataclass, replace +from types import SimpleNamespace from unittest.mock import MagicMock import pytest import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from tests.kernels.moe.utils import make_dummy_moe_config +from tests.kernels.moe.utils import ( + check_deferred_moe_finalize, + make_dummy_moe_config, + make_test_weights, +) from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -32,6 +37,11 @@ TrtLlmFp8ExpertsMonolithic, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts +from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + convert_to_fp8_moe_kernel_format, +) +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, swap_w13_to_w31, @@ -767,3 +777,80 @@ def _supports_quant_scheme( assert supported == expected, reason if not expected: assert "SwiGLU" in reason + + +@pytest.mark.parametrize("m", [1, 16]) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="Requires TRTLLM-Gen FP8 MoE (SM100)", +) +def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): + """TRTLLM-Gen block-FP8 modular experts can leave the top-k finalize open.""" + e, topk, n, k = 32, 4, 1024, 1024 + block_shape = [128, 128] + set_random_seed(7) + with set_current_vllm_config(vllm_config): + (_, w1, w1_scale, _), (_, w2, w2_scale, _) = make_test_weights( + e, n, k, quant_dtype=torch.float8_e4m3fn, block_shape=block_shape + ) + w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format( + Fp8MoeBackend.FLASHINFER_TRTLLM, + SimpleNamespace( + weight_block_size=block_shape, + moe_config=SimpleNamespace( + is_act_and_mul=True, intermediate_size_per_partition=n + ), + activation=MoEActivation.SILU, + ), + w1, + w2, + w1_scale, + w2_scale, + w13_input_scale=None, + w2_input_scale=None, + ) + quant_config = fp8_w8a8_moe_quant_config( + w1_scale=w1_scale, w2_scale=w2_scale, block_shape=block_shape + ) + # One rank of a TP group: the only topology that defers. + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=2 * n, + num_local_experts=e, + num_logical_experts=e, + activation=MoEActivation.SILU, + device="cuda", + moe_parallel_config=replace( + FusedMoEParallelConfig.make_no_parallel(), tp_size=2 + ), + in_dtype=torch.bfloat16, + routing_method=RoutingMethodType.TopK, + max_num_tokens=next_power_of_2(m), + ) + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, quant_config=quant_config, allow_new_interface=True + ), + TrtLlmFp8ExpertsModular(moe_config=moe_config, quant_config=quant_config), + ) + + a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 + score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + check_deferred_moe_finalize( + moe_config, + lambda: kernel.apply( + hidden_states=a, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + ), + topk_weights, + ) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 1cd204d4b00b..ed06cb1ce9e1 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -7,7 +7,7 @@ import pytest import torch -from tests.kernels.moe.utils import check_accuracy +from tests.kernels.moe.utils import check_accuracy, check_deferred_moe_finalize from vllm._aiter_ops import ( is_aiter_found, is_aiter_found_and_supported, @@ -663,7 +663,133 @@ def test_trtllm_gen_mxfp4_fused_moe( check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8) -@pytest.mark.parametrize("topk", [1, 4]) +@pytest.mark.parametrize("num_tokens,chunk_size", [(1, None), (16, None), (16, 5)]) +@pytest.mark.skipif( + not TRTLLM_GEN_MXFP4_AVAILABLE, + reason="nvidia gpu and compute capability sm100 is required for this test", +) +def test_trtllm_mxfp4_modular_deferred_finalize( + num_tokens: int, + chunk_size: int | None, + monkeypatch: pytest.MonkeyPatch, + workspace_init, +): + """TRTLLM-Gen MXFP4 modular experts can leave the top-k finalize open.""" + 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, + ) + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + convert_weight_to_mxfp4_moe_kernel_format, + make_mxfp4_moe_kernel, + make_mxfp4_moe_quant_config, + ) + + num_experts, topk, hidden_size, intermediate_size = 32, 4, 1024, 512 + device = torch.device("cuda:0") + torch.manual_seed(0) + if chunk_size is not None: + monkeypatch.setattr( + TrtLlmMxfp4ExpertsModular, + "_max_supported_tokens", + lambda self, top_k, global_num_experts: chunk_size, + ) + + def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + w_q, w_scale = fp4_quantize( + w, + torch.tensor(1.0, device=device), + 32, + sf_use_ue8m0=True, + is_sf_swizzled_layout=False, + ) + return ( + w_q.reshape(*w.shape[:-1], -1), + w_scale.view(torch.uint8).reshape(*w.shape[:-1], -1), + ) + + w13, w13_scale = quantize( + torch.randn( + num_experts, + 2 * intermediate_size, + hidden_size, + device=device, + dtype=torch.bfloat16, + ) + ) + w2, w2_scale = quantize( + torch.randn( + num_experts, + hidden_size, + intermediate_size, + device=device, + dtype=torch.bfloat16, + ) + ) + backend = Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8 + w13, w2, w13_scale, w2_scale, _, _ = convert_weight_to_mxfp4_moe_kernel_format( + backend, + torch.nn.Module(), + w13, + w2, + w13_scale, + w2_scale, + _cache_permute_indices={}, + ) + # One rank of a TP group: the only topology that defers. + moe_config = FusedMoEConfig( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=hidden_size, + intermediate_size=2 * intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + activation=MoEActivation.SILU, + device=device, + routing_method=RoutingMethodType.TopK, + moe_parallel_config=replace( + FusedMoEParallelConfig.make_no_parallel(), tp_size=2 + ), + in_dtype=torch.bfloat16, + ) + kernel = make_mxfp4_moe_kernel( + make_mxfp4_moe_quant_config(backend, w1_scale=w13_scale, w2_scale=w2_scale), + moe_config, + TrtLlmMxfp4ExpertsModular, + backend, + ) + + hidden_states = torch.randn( + num_tokens, hidden_size, device=device, dtype=torch.bfloat16 + ) + topk_ids = torch.rand(num_tokens, num_experts, device=device).argsort(dim=-1) + topk_ids = topk_ids[:, :topk].to(torch.int32) + topk_weights = torch.rand(num_tokens, topk, device=device).softmax(dim=-1) + + check_deferred_moe_finalize( + moe_config, + lambda: kernel.apply( + hidden_states=hidden_states, + w1=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ), + topk_weights, + chunked=chunk_size is not None, + ) + + @pytest.mark.parametrize("num_experts", [32]) @pytest.mark.parametrize("num_tokens", [1, 128]) @pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)]) diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index bc0711200b2f..40b6b074f111 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -8,11 +8,13 @@ intermediate_size 704) that exercises the non-256-aligned padding path. """ +from dataclasses import replace + import pytest import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from tests.kernels.moe.utils import make_test_quant_config +from tests.kernels.moe.utils import check_deferred_moe_finalize, make_test_quant_config from tests.kernels.quantization.nvfp4_utils import ( FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX, @@ -31,6 +33,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, + FusedMoEQuantConfig, RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( @@ -92,6 +95,36 @@ def __init__(self, *, compile_native: bool = True) -> None: SITU = op_registry[_SITU_OP_NAME] +def _make_trtllm_fp4_moe_kernel( + moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig +) -> mk.FusedMoEKernel: + experts = TrtLlmNvFp4ExpertsModular( + moe_config=moe_config, quant_config=quant_config + ) + # Mimic the production weight-loader path so per-expert tensors that + # are normally precomputed in process_weights_after_loading (g1_scale_c + # and the rescaled gemm1_clamp_limit) get materialized. The test's + # synthetic quant_config has g1_alphas/g2_alphas already at their + # post-fusion values, so we set w13_weight_scale_2 to alias g1_alphas + # (same tensor) and use input_scale=1 to make the in-place + # weight_scale_2 *= input_scale step a no-op. + fake_layer = torch.nn.Module() + fake_layer.w13_weight_scale_2 = quant_config.g1_alphas + fake_layer.w2_weight_scale_2 = quant_config.g2_alphas + fake_layer.w13_input_scale = torch.ones_like(quant_config.g1_alphas) + fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas) + experts.process_weights_after_loading(fake_layer) + return mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + experts, + ) + + ACTIVATION_CASES = [ pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"), pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"), @@ -191,34 +224,11 @@ def test_trtllm_fp4_moe_no_graph( ), ) - trtllm_inner = TrtLlmNvFp4ExpertsModular( - moe_config=moe_config, quant_config=quant_config - ) - # Mimic the production weight-loader path so per-expert tensors that - # are normally precomputed in process_weights_after_loading (g1_scale_c - # and the rescaled gemm1_clamp_limit) get materialized. The test's - # synthetic quant_config has g1_alphas/g2_alphas already at their - # post-fusion values, so we set w13_weight_scale_2 to alias g1_alphas - # (same tensor) and use input_scale=1 to make the in-place - # weight_scale_2 *= input_scale step a no-op. - fake_layer = torch.nn.Module() - fake_layer.w13_weight_scale_2 = quant_config.g1_alphas - fake_layer.w2_weight_scale_2 = quant_config.g2_alphas - fake_layer.w13_input_scale = torch.ones_like(quant_config.g1_alphas) - fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas) - trtllm_inner.process_weights_after_loading(fake_layer) + trtllm_experts = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) if activation == MoEActivation.SITU: - torch.testing.assert_close(trtllm_inner.g1_scale_c, quant_config.a2_gscale) - - trtllm_experts = mk.FusedMoEKernel( - maybe_make_prepare_finalize( - moe=moe_config, - quant_config=quant_config, - allow_new_interface=True, - use_monolithic=False, - ), - trtllm_inner, - ) + torch.testing.assert_close( + trtllm_experts.fused_experts.g1_scale_c, quant_config.a2_gscale + ) trtllm_output = trtllm_experts.apply( hidden_states=a, @@ -331,6 +341,70 @@ def test_trtllm_fp4_moe_reprocess_does_not_refold_swiglu_params(): torch.testing.assert_close(experts.gemm1_beta, beta * expected) +@pytest.mark.parametrize("m,chunk_size", [(1, None), (16, None), (16, 5)]) +@torch.inference_mode() +def test_trtllm_fp4_moe_deferred_finalize( + m: int, + chunk_size: int | None, + monkeypatch: pytest.MonkeyPatch, + workspace_init, +): + """TRTLLM-Gen NvFP4 modular experts can leave the top-k finalize open.""" + e, topk, n, k = 128, 8, 1024, 1024 + dtype = torch.bfloat16 + if chunk_size is not None: + monkeypatch.setattr( + TrtLlmNvFp4ExpertsModular, "_get_chunk_size", lambda self: chunk_size + ) + + set_random_seed(7) + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 + w1_q, w2_q, quant_config = make_test_quant_config( + e, n, k, in_dtype=dtype, quant_dtype="nvfp4", is_scale_swizzled=False + ) + score = torch.randn((m, e), device="cuda", dtype=dtype) + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + + # One rank of a TP group: the only topology that defers. + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=2 * n, + num_local_experts=e, + num_logical_experts=e, + activation=MoEActivation.SILU, + device="cuda", + moe_parallel_config=replace( + FusedMoEParallelConfig.make_no_parallel(), tp_size=2 + ), + in_dtype=dtype, + routing_method=RoutingMethodType.TopK, + max_num_tokens=next_power_of_2(m), + ) + kernel = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) + + check_deferred_moe_finalize( + moe_config, + lambda: kernel.apply( + hidden_states=a, + w1=w1_q, + w2=w2_q, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + ), + topk_weights, + chunked=chunk_size is not None, + ) + + if __name__ == "__main__": test_trtllm_fp4_moe_no_graph( 64, diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 36189e4dd7f1..e5bd5a66f195 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable import torch @@ -28,6 +29,7 @@ fused_experts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import ( BatchedPrepareAndFinalize, ) @@ -707,6 +709,39 @@ def check_accuracy(a, b, atol, rtol, percent): ) +def check_deferred_moe_finalize( + moe_config: FusedMoEConfig, + run: Callable[[], torch.Tensor | UnfinalizedMoEOutput], + topk_weights: torch.Tensor, + chunked: bool = False, +) -> None: + """Check that a deferred MoE run leaves open exactly the finalize it skips. + + ``run`` executes the MoE, first with ``moe_config``'s deferral off and then + on. Reducing the deferred GEMM2 rows with the router's weights must give the + kernel's own finalized output. Each launch permutes into its own buffer, so + a run chunked into several launches finalizes instead. + """ + moe_config.defer_moe_finalize = False + finalized = run() + moe_config.defer_moe_finalize = True + deferred = run() + assert isinstance(finalized, torch.Tensor) + if chunked: + assert isinstance(deferred, torch.Tensor) + torch.testing.assert_close(deferred, finalized, atol=0, rtol=0) + return + + assert isinstance(deferred, UnfinalizedMoEOutput) + # Modular experts get the router's weights, and hand them back as-is. + torch.testing.assert_close(deferred.expert_weights, topk_weights, atol=0, rtol=0) + permuted_idx = deferred.expanded_idx_to_permuted_idx + assert permuted_idx.shape == topk_weights.shape and (permuted_idx >= 0).all() + rows = deferred.gemm2_permuted[permuted_idx.long()].float() + reference = (rows * topk_weights[..., None]).sum(dim=1).to(finalized.dtype) + torch.testing.assert_close(reference, finalized) + + def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8): """Weight/scale layouts for mxfp4 MoE, as (layout, opts) pairs. diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index e860597bcba3..fb52445d8417 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -12,6 +12,10 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.moe_output import ( + UnfinalizedMoEOutput, + convert_flashinfer_moe_output, +) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) @@ -273,7 +277,7 @@ def apply( workspace2: torch.Tensor, expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - ): + ) -> UnfinalizedMoEOutput | None: import flashinfer from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout @@ -297,7 +301,9 @@ def apply( weight_layout = WeightLayout.BlockMajorK hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale) - flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe( + num_tokens = hidden_states.shape[0] + defer = self.moe_config.should_defer_moe_finalize(num_tokens) + flashinfer_output = flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe( topk_ids=(topk_ids, topk_weights), routing_bias=None, hidden_states=hidden_states, @@ -321,9 +327,20 @@ def apply( use_shuffled_weight=use_shuffled_weight, weight_layout=weight_layout, fp8_quantization_type=fp8_quant_type, - output=output, + do_finalize=not defer, + output=None if defer else output, tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) + if not defer: + return None + routed_output = convert_flashinfer_moe_output( + flashinfer_output, + do_finalize=False, + num_tokens=num_tokens, + top_k=topk_ids.size(1), + ) + assert isinstance(routed_output, UnfinalizedMoEOutput) + return routed_output class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolithic): diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 1c20bc1cbcdf..aae8ecff5283 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -343,7 +343,7 @@ def _max_supported_tokens(self, top_k: int, global_num_experts: int) -> int: def _invoke_kernel( self, - output: torch.Tensor, + output: torch.Tensor | None, x_quant: torch.Tensor, x_scale: torch.Tensor | None, topk_ids: torch.Tensor, @@ -355,10 +355,11 @@ def _invoke_kernel( local_num_experts: int, local_expert_offset: int, topk: int, - ) -> None: + ) -> UnfinalizedMoEOutput | None: + """Finalize into ``output``, or stop after GEMM2 when it is None.""" from flashinfer import trtllm_fp4_block_scale_routed_moe - trtllm_fp4_block_scale_routed_moe( + flashinfer_output = trtllm_fp4_block_scale_routed_moe( topk_ids=(topk_ids, topk_weights), routing_bias=None, hidden_states=x_quant, @@ -386,12 +387,22 @@ def _invoke_kernel( # Modular kernel receives pre-routed tokens, so routing is already # done. Use Renormalize as a safe default the TRTLLM kernel supports. routing_method_type=RoutingMethodType.Renormalize, - do_finalize=True, + do_finalize=output is not None, enable_pdl=True, activation_type=self._flashinfer_activation_type(activation), output=output, tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) + if output is not None: + return None + routed_output = convert_flashinfer_moe_output( + flashinfer_output, + do_finalize=False, + num_tokens=x_quant.shape[0], + top_k=topk, + ) + assert isinstance(routed_output, UnfinalizedMoEOutput) + return routed_output def apply( self, @@ -410,7 +421,7 @@ def apply( workspace2: torch.Tensor, expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - ): + ) -> UnfinalizedMoEOutput | None: topk_ids = topk_ids.to(dtype=torch.int32) topk = topk_ids.size(-1) @@ -428,13 +439,17 @@ def apply( assert self.w1_scale is not None assert self.w2_scale is not None - # Chunk tokens so the batched-GEMM grid stays within CUDA limits. + # Chunk tokens so the batched-GEMM grid stays within CUDA limits. Each + # launch permutes into its own buffer, so only a run that fits in one + # launch can leave the top-k reduction to a deferring consumer. M = x_quant.size(0) chunk_size = self._max_supported_tokens(topk, global_num_experts) + defer = chunk_size >= M and self.moe_config.should_defer_moe_finalize(M) + unfinalized: UnfinalizedMoEOutput | None = None for start in range(0, M, chunk_size): end = min(start + chunk_size, M) - self._invoke_kernel( - output[start:end], + unfinalized = self._invoke_kernel( + None if defer else output[start:end], x_quant[start:end], None if x_scale is None else x_scale[start:end], topk_ids[start:end], @@ -448,4 +463,4 @@ def apply( topk, ) - return output + return unfinalized diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 646b4c8ad6b5..1423fe444286 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -333,7 +333,7 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: def _invoke_kernel( self, - output: torch.Tensor, + output: torch.Tensor | None, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, @@ -342,7 +342,8 @@ def _invoke_kernel( activation: MoEActivation, global_num_experts: int, a1q_scale: torch.Tensor | None, - ): + ) -> UnfinalizedMoEOutput | None: + """Finalize into ``output``, or stop after GEMM2 when it is None.""" import flashinfer assert self.quant_config.w1_scale is not None @@ -359,7 +360,7 @@ def _invoke_kernel( output1_scale_gate_scalar = self.quant_config.g1_alphas # Invoke kernel. - flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( + flashinfer_output = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( topk_ids=(topk_ids, topk_weights), routing_bias=None, hidden_states=hidden_states, @@ -387,7 +388,7 @@ def _invoke_kernel( local_num_experts=self.local_num_experts, routed_scaling_factor=None, routing_method_type=1, # not used - do_finalize=True, + do_finalize=output is not None, activation_type=activation_to_flashinfer_int(activation), per_token_scale=per_token_scale, output=output, @@ -395,6 +396,16 @@ def _invoke_kernel( fi_moe_largest_bucket(self.moe_config), self._get_chunk_size() ), ) + if output is not None: + return None + routed_output = convert_flashinfer_moe_output( + flashinfer_output, + do_finalize=False, + num_tokens=hidden_states.shape[0], + top_k=self.topk, + ) + assert isinstance(routed_output, UnfinalizedMoEOutput) + return routed_output def apply( self, @@ -413,7 +424,7 @@ def apply( workspace2: torch.Tensor, expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - ): + ) -> UnfinalizedMoEOutput | None: assert self._supports_activation(activation) # Per-token defers input quant to _invoke_kernel, so a1q_scale is None. assert a1q_scale is not None or self.per_token_activation @@ -425,8 +436,12 @@ def apply( chunk_size = self._get_chunk_size() if chunk_size >= M: - self._invoke_kernel( - output, + # Each launch permutes into its own buffer, so only a run that fits + # in one launch can leave the top-k reduction to a deferring + # consumer. + defer = self.moe_config.should_defer_moe_finalize(M) + return self._invoke_kernel( + None if defer else output, hidden_states, w1, w2, @@ -450,6 +465,7 @@ def apply( global_num_experts, None if a1q_scale is None else a1q_scale[start:end], ) + return None class TrtLlmNvFp4ExpertsMonolithic( diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py index c64aeadcc7a5..d83ed2c507ef 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py @@ -148,7 +148,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: "SharedExperts | None", shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: """Apply the MoE operation using modular kernels. Args: @@ -160,7 +160,7 @@ def apply( shared_experts: The shared experts module, if any Returns: - Output tensor from routed experts. + Finalized routed states or a deferred-finalize output. """ raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index fba16f67b0f6..f563c819e4b3 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -238,6 +238,15 @@ def supports_async(self) -> bool: """ return False + def supports_deferred_moe_finalize(self) -> bool: + """Whether ``finalize`` can be skipped for a deferring consumer. + + An implementation opts in only if everything it does in ``finalize`` + -- the top-k reduction, and any combine or reduce-scatter -- is work + the consumer takes over. + """ + return False + # TODO: pass FusedMoEParallelConfig in as ctor parameter? class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): @@ -404,9 +413,6 @@ class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize): described above for the monolithic case. """ - def supports_deferred_moe_finalize(self) -> bool: - return False - @abstractmethod def prepare( self, @@ -895,10 +901,14 @@ def apply( workspace2: torch.Tensor, expert_tokens_meta: ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - ) -> None: + ) -> UnfinalizedMoEOutput | None: """This function computes the intermediate result of a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2. + Writes into `output` and returns None, unless the implementation stopped + after GEMM2 and left the top-k reduction to a fused consumer, in which + case `output` is untouched and the unfinalized result is returned. + Args: output: (torch.Tensor): The unweighted, unreduced output tensor. hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE @@ -1270,7 +1280,7 @@ def _fused_experts( apply_router_weight_on_input: bool, expert_tokens_meta: ExpertTokensMetadata | None, output_alias: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: _, M_full, N, K, top_k = self.fused_experts.moe_problem_size( a1q, w1, w2, topk_ids ) @@ -1318,7 +1328,7 @@ def _fused_experts( elif use_output_alias: fused_out = output_alias - self.fused_experts.apply( + unfinalized = self.fused_experts.apply( output=fused_out, hidden_states=a1q, w1=w1, @@ -1336,6 +1346,11 @@ def _fused_experts( apply_router_weight_on_input=apply_router_weight_on_input, ) + # Experts that stopped after GEMM2 wrote nothing into `fused_out`; the + # top-k reduction they left open belongs to whoever takes this. + if isinstance(unfinalized, UnfinalizedMoEOutput): + return unfinalized + return fused_out def _finalize( @@ -1426,7 +1441,7 @@ def apply( apply_router_weight_on_input: bool = False, shared_experts: SharedExperts | None = None, shared_experts_input: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: """This function computes a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2, and top-k gating mechanism. @@ -1453,7 +1468,9 @@ def apply( hidden_states before latent projection. Returns: - torch.Tensor: The output tensor after applying the MoE layer. + torch.Tensor: The output tensor after applying the MoE layer, or + the unfinalized output when the experts left the top-k reduction to + a fused consumer. """ output = torch.empty_like(hidden_states) @@ -1498,6 +1515,17 @@ def apply( if lora_ctx is not None: lora_ctx.original_hidden_states = None + if isinstance(fused_out, UnfinalizedMoEOutput): + # Nothing below can run on an unfinalized output: finalize is the + # local top-k reduction (and, for DP/EP, the combine) that the + # consumer takes over. + if not self.prepare_finalize.supports_deferred_moe_finalize(): + raise RuntimeError( + f"{type(self.prepare_finalize).__name__} cannot pass through " + "a deferred MoE output." + ) + return fused_out + return self._finalize( output, fused_out, @@ -1653,10 +1681,7 @@ def output_is_reduced(self) -> bool: return self.prepare_finalize.output_is_reduced() def supports_deferred_moe_finalize(self) -> bool: - return ( - isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) - and self.prepare_finalize.supports_deferred_moe_finalize() - ) + return self.prepare_finalize.supports_deferred_moe_finalize() def apply_monolithic( self, @@ -1703,7 +1728,7 @@ def apply( apply_router_weight_on_input: bool, shared_experts: SharedExperts | None = None, shared_experts_input: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert isinstance(self.impl, FusedMoEKernelModularImpl) return self.impl.apply( hidden_states=hidden_states, diff --git a/vllm/model_executor/layers/fused_moe/moe_output.py b/vllm/model_executor/layers/fused_moe/moe_output.py index 04081cf6d17b..8f9662b442ce 100644 --- a/vllm/model_executor/layers/fused_moe/moe_output.py +++ b/vllm/model_executor/layers/fused_moe/moe_output.py @@ -9,7 +9,7 @@ @dataclass class UnfinalizedMoEOutput: - """Unfinalized output of a monolithic MoE kernel. + """Unfinalized output of a MoE kernel. Kernels that can stop after GEMM2 (the TRTLLM-Gen ``do_finalize=False`` path) hand back their permuted, unweighted output plus the routing weights and the @@ -24,10 +24,11 @@ class UnfinalizedMoEOutput: # [num_permuted_rows, hidden_dim] permuted, unweighted GEMM2 output. gemm2_permuted: torch.Tensor - # [num_tokens, top_k] routing weights, in the activation dtype -- consumers - # index them as such, and a buffer typed wider than its contents reads as - # garbage. These already carry routed_scaling_factor for routing methods that - # fold it in. + # [num_tokens, top_k] routing weights, in the activation dtype when the + # kernel routed itself and in the router's dtype for modular experts -- + # consumers index them as such, and a buffer typed wider than its contents + # reads as garbage. These already carry routed_scaling_factor for routing + # methods that fold it in. expert_weights: torch.Tensor # [num_tokens, top_k] int32 permute map; -1 = expert not local to this rank. expanded_idx_to_permuted_idx: torch.Tensor diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py index d13bc20673ec..660807fea3e4 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py @@ -38,6 +38,11 @@ def _quantize_input( class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular): + def supports_deferred_moe_finalize(self) -> bool: + # finalize() here is the top-k weight-and-reduce and nothing else, + # which is exactly the work a deferring consumer takes over. + return True + @property def activation_format(self) -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index e3a287ddc2da..15f6c5ed44a4 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -1224,7 +1224,7 @@ def forward_modular( topk_ids: torch.Tensor, shared_experts: "SharedExperts | None" = None, shared_experts_input: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: """Execute routed experts using the quantization method's apply function. This is called by the runner after router selection (for modular kernels) @@ -1239,7 +1239,7 @@ def forward_modular( shared_experts_input: Input for shared experts (if any) Returns: - Output tensor from routed experts. + Finalized routed states or a deferred-finalize output. """ assert not self.quant_method.is_monolithic diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index f7d12c863593..643fd75e820a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsMxFp4, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( B12X_BACKENDS, Mxfp4MoeBackend, @@ -227,7 +228,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.moe_kernel is not None return self.moe_kernel.apply( x, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index 54f5aa86b34c..51b7d0460a97 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -303,7 +303,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.moe_kernel is not None return self.moe_kernel.apply( x, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py index 262cd307a849..82f3ada44666 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py @@ -18,6 +18,7 @@ FusedMoEConfig, FusedMoEQuantConfig, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, @@ -375,7 +376,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py index ee3b321363b2..b5326e1875e5 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py @@ -12,6 +12,7 @@ FusedMoEConfig, FusedMoEQuantConfig, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, @@ -188,7 +189,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 860c41e50199..37bb9c684014 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -24,6 +24,7 @@ UnquantizedFusedMoEMethod, ) from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, @@ -827,7 +828,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py index b1a8a294f1be..53e8ca2b0163 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( B12X_BACKENDS, Mxfp4MoeBackend, @@ -228,7 +229,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.moe_kernel is not None return self.moe_kernel.apply( x, diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_moe.py index 2c9a06b798bb..ed7bd9fed000 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_moe.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_moe.py @@ -12,6 +12,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, @@ -179,7 +180,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.moe_kernel is not None return self.moe_kernel.apply( x, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 88c63d97b8f2..90ac2b8c7951 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -690,7 +690,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( @@ -1102,7 +1102,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( @@ -1471,7 +1471,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 33f6afac80a9..dcba98c7e468 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -433,7 +433,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( @@ -860,7 +860,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index 866c88b3d649..414c4aba5b08 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -166,7 +166,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 3a9f6d63662b..2e7fc12484d3 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -976,7 +976,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, @@ -1935,7 +1935,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, @@ -2182,7 +2182,7 @@ def apply( topk_ids: torch.Tensor, shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.moe_kernel is not None return self.moe_kernel.apply( x, From 8b2ae3845fc6cfd07b15f1e60d1309b5e2797275 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Thu, 24 Sep 2026 23:37:32 +0000 Subject: [PATCH 2/8] [MoE] Make deferred MoE finalize a per-layer setting fixed at model init Deferral was a mutable FusedMoEConfig flag plus a token cap, which the experts re-evaluated on every call, so a consumer had to be ready for either output form. A model now marks a layer once, while it is built, with MoERunner.enable_deferred_moe_finalize; that layer returns an UnfinalizedMoEOutput from every call, through forward_unfinalized. The layer checks the experts' capability and the topology up front, and the kernel binds the flag to its experts on first use. Runs the experts would have to chunk raise instead of finalizing. Where a fused consumer can't take the output, finalize_moe_output runs the TRT-LLM finalize as a standalone kernel (vLLM's moe_unpermute), bit-identical to the experts' own finalize. Kimi-K3 moves onto this: its tail takes the deferred output up to its token cap and finalizes it standalone above. The warmup's bounded autotune passes for deferred layers go away, since the full-size pass now runs their kernels. Also restores the topk parametrization of test_flashinfer_cutlass_mxfp4_fused_moe, dropped by the previous commit. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yongye Zhu --- tests/kernels/moe/test_flashinfer.py | 16 ++-- tests/kernels/moe/test_ocp_mx_moe.py | 45 +++++++--- tests/kernels/moe/test_trtllm_nvfp4_moe.py | 15 ++-- tests/kernels/moe/utils.py | 42 +++++----- .../test_flashinfer_autotune_warmup.py | 46 +---------- tests/models/kimi_k3/test_latent_moe_tail.py | 12 +-- .../model_executor/layers/fused_moe/config.py | 37 --------- .../fused_moe/experts/trtllm_bf16_moe.py | 8 +- .../fused_moe/experts/trtllm_fp8_moe.py | 6 +- .../fused_moe/experts/trtllm_mxfp4_moe.py | 19 +++-- .../fused_moe/experts/trtllm_nvfp4_moe.py | 20 +++-- .../layers/fused_moe/modular_kernel.py | 25 +++++- .../layers/fused_moe/moe_output.py | 31 ++++++- .../layers/fused_moe/runner/moe_runner.py | 82 ++++++++++++++++++- vllm/model_executor/warmup/kernel_warmup.py | 23 +----- .../kimi_k3/nvidia/latent_moe_runner.py | 82 +++++++++---------- 16 files changed, 277 insertions(+), 232 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 2429f3e5d063..d26360737b07 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing -from dataclasses import dataclass, replace +from dataclasses import dataclass from types import SimpleNamespace from unittest.mock import MagicMock @@ -785,7 +785,8 @@ def _supports_quant_scheme( reason="Requires TRTLLM-Gen FP8 MoE (SM100)", ) def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): - """TRTLLM-Gen block-FP8 modular experts can leave the top-k finalize open.""" + """TRTLLM-Gen block-FP8 modular experts can leave the top-k finalize to the + caller.""" e, topk, n, k = 32, 4, 1024, 1024 block_shape = [128, 128] set_random_seed(7) @@ -812,19 +813,16 @@ def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): quant_config = fp8_w8a8_moe_quant_config( w1_scale=w1_scale, w2_scale=w2_scale, block_shape=block_shape ) - # One rank of a TP group: the only topology that defers. moe_config = FusedMoEConfig( num_experts=e, experts_per_token=topk, hidden_dim=k, - intermediate_size=2 * n, + intermediate_size=n, num_local_experts=e, num_logical_experts=e, activation=MoEActivation.SILU, device="cuda", - moe_parallel_config=replace( - FusedMoEParallelConfig.make_no_parallel(), tp_size=2 - ), + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), in_dtype=torch.bfloat16, routing_method=RoutingMethodType.TopK, max_num_tokens=next_power_of_2(m), @@ -840,7 +838,7 @@ def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) check_deferred_moe_finalize( - moe_config, + kernel, lambda: kernel.apply( hidden_states=a, w1=w1, @@ -852,5 +850,5 @@ def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): expert_map=None, apply_router_weight_on_input=False, ), - topk_weights, + router_weights=topk_weights, ) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index ed06cb1ce9e1..bddefe4e51b6 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -663,18 +663,22 @@ def test_trtllm_gen_mxfp4_fused_moe( check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8) -@pytest.mark.parametrize("num_tokens,chunk_size", [(1, None), (16, None), (16, 5)]) +@pytest.mark.parametrize( + "monolithic,num_tokens,chunk_size", + [(False, 1, None), (False, 16, None), (False, 16, 5), (True, 16, None)], +) @pytest.mark.skipif( not TRTLLM_GEN_MXFP4_AVAILABLE, reason="nvidia gpu and compute capability sm100 is required for this test", ) -def test_trtllm_mxfp4_modular_deferred_finalize( +def test_trtllm_mxfp4_deferred_finalize( + monolithic: bool, num_tokens: int, chunk_size: int | None, monkeypatch: pytest.MonkeyPatch, workspace_init, ): - """TRTLLM-Gen MXFP4 modular experts can leave the top-k finalize open.""" + """TRTLLM-Gen MXFP4 experts can leave the top-k finalize to the caller.""" from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -683,6 +687,7 @@ def test_trtllm_mxfp4_modular_deferred_finalize( ) from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( TrtLlmMxfp4ExpertsModular, + TrtLlmMxfp4ExpertsMonolithic, ) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( Mxfp4MoeBackend, @@ -742,38 +747,51 @@ def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: w2_scale, _cache_permute_indices={}, ) - # One rank of a TP group: the only topology that defers. moe_config = FusedMoEConfig( num_experts=num_experts, experts_per_token=topk, hidden_dim=hidden_size, - intermediate_size=2 * intermediate_size, + intermediate_size=intermediate_size, num_local_experts=num_experts, num_logical_experts=num_experts, activation=MoEActivation.SILU, device=device, - routing_method=RoutingMethodType.TopK, - moe_parallel_config=replace( - FusedMoEParallelConfig.make_no_parallel(), tp_size=2 - ), + routing_method=RoutingMethodType.Renormalize, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), in_dtype=torch.bfloat16, ) kernel = make_mxfp4_moe_kernel( make_mxfp4_moe_quant_config(backend, w1_scale=w13_scale, w2_scale=w2_scale), moe_config, - TrtLlmMxfp4ExpertsModular, + TrtLlmMxfp4ExpertsMonolithic if monolithic else TrtLlmMxfp4ExpertsModular, backend, ) hidden_states = torch.randn( num_tokens, hidden_size, device=device, dtype=torch.bfloat16 ) + if monolithic: + router_logits = torch.randn(num_tokens, num_experts, device=device) + check_deferred_moe_finalize( + kernel, + lambda: kernel.apply_monolithic( + hidden_states, + w13, + w2, + router_logits, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ), + ) + return + topk_ids = torch.rand(num_tokens, num_experts, device=device).argsort(dim=-1) topk_ids = topk_ids[:, :topk].to(torch.int32) topk_weights = torch.rand(num_tokens, topk, device=device).softmax(dim=-1) - check_deferred_moe_finalize( - moe_config, + kernel, lambda: kernel.apply( hidden_states=hidden_states, w1=w13, @@ -785,11 +803,12 @@ def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: expert_map=None, apply_router_weight_on_input=False, ), - topk_weights, + router_weights=topk_weights, chunked=chunk_size is not None, ) +@pytest.mark.parametrize("topk", [1, 4]) @pytest.mark.parametrize("num_experts", [32]) @pytest.mark.parametrize("num_tokens", [1, 128]) @pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)]) diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index 40b6b074f111..dc9c9e2ca54e 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -8,8 +8,6 @@ intermediate_size 704) that exercises the non-256-aligned padding path. """ -from dataclasses import replace - import pytest import torch @@ -349,7 +347,7 @@ def test_trtllm_fp4_moe_deferred_finalize( monkeypatch: pytest.MonkeyPatch, workspace_init, ): - """TRTLLM-Gen NvFP4 modular experts can leave the top-k finalize open.""" + """TRTLLM-Gen NvFP4 modular experts can leave the top-k finalize to the caller.""" e, topk, n, k = 128, 8, 1024, 1024 dtype = torch.bfloat16 if chunk_size is not None: @@ -368,19 +366,16 @@ def test_trtllm_fp4_moe_deferred_finalize( score = torch.randn((m, e), device="cuda", dtype=dtype) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) - # One rank of a TP group: the only topology that defers. moe_config = FusedMoEConfig( num_experts=e, experts_per_token=topk, hidden_dim=k, - intermediate_size=2 * n, + intermediate_size=n, num_local_experts=e, num_logical_experts=e, activation=MoEActivation.SILU, device="cuda", - moe_parallel_config=replace( - FusedMoEParallelConfig.make_no_parallel(), tp_size=2 - ), + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), in_dtype=dtype, routing_method=RoutingMethodType.TopK, max_num_tokens=next_power_of_2(m), @@ -388,7 +383,7 @@ def test_trtllm_fp4_moe_deferred_finalize( kernel = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) check_deferred_moe_finalize( - moe_config, + kernel, lambda: kernel.apply( hidden_states=a, w1=w1_q, @@ -400,7 +395,7 @@ def test_trtllm_fp4_moe_deferred_finalize( expert_map=None, apply_router_weight_on_input=False, ), - topk_weights, + router_weights=topk_weights, chunked=chunk_size is not None, ) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index e5bd5a66f195..9a3386b0a611 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable +import pytest import torch import vllm._custom_ops as ops @@ -29,7 +30,10 @@ fused_experts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel -from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.moe_output import ( + UnfinalizedMoEOutput, + finalize_moe_output, +) from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import ( BatchedPrepareAndFinalize, ) @@ -710,36 +714,34 @@ def check_accuracy(a, b, atol, rtol, percent): def check_deferred_moe_finalize( - moe_config: FusedMoEConfig, + kernel: FusedMoEKernel, run: Callable[[], torch.Tensor | UnfinalizedMoEOutput], - topk_weights: torch.Tensor, + router_weights: torch.Tensor | None = None, chunked: bool = False, ) -> None: - """Check that a deferred MoE run leaves open exactly the finalize it skips. + """Check a kernel that defers its finalize against the finalize it skips. - ``run`` executes the MoE, first with ``moe_config``'s deferral off and then - on. Reducing the deferred GEMM2 rows with the router's weights must give the - kernel's own finalized output. Each launch permutes into its own buffer, so - a run chunked into several launches finalizes instead. + ``run`` calls ``kernel`` on fixed inputs, first as built and then with the + finalize deferred. ``finalize_moe_output`` on the deferred output must give + the kernel's own finalized output bit for bit, and modular experts must hand + the router's weights back as-is. Each launch permutes into its own buffer, + so a deferring kernel refuses a run it would have to chunk. """ - moe_config.defer_moe_finalize = False finalized = run() - moe_config.defer_moe_finalize = True - deferred = run() assert isinstance(finalized, torch.Tensor) + kernel.enable_deferred_moe_finalize() if chunked: - assert isinstance(deferred, torch.Tensor) - torch.testing.assert_close(deferred, finalized, atol=0, rtol=0) + with pytest.raises(ValueError, match="one kernel launch"): + run() return + deferred = run() assert isinstance(deferred, UnfinalizedMoEOutput) - # Modular experts get the router's weights, and hand them back as-is. - torch.testing.assert_close(deferred.expert_weights, topk_weights, atol=0, rtol=0) - permuted_idx = deferred.expanded_idx_to_permuted_idx - assert permuted_idx.shape == topk_weights.shape and (permuted_idx >= 0).all() - rows = deferred.gemm2_permuted[permuted_idx.long()].float() - reference = (rows * topk_weights[..., None]).sum(dim=1).to(finalized.dtype) - torch.testing.assert_close(reference, finalized) + if router_weights is not None: + torch.testing.assert_close( + deferred.expert_weights, router_weights, atol=0, rtol=0 + ) + torch.testing.assert_close(finalize_moe_output(deferred), finalized, atol=0, rtol=0) def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8): diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index dd162629b8f4..b6f43db9d6a6 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -3,7 +3,6 @@ from inspect import signature from types import SimpleNamespace -from typing import Any from unittest.mock import Mock, call, patch import pytest @@ -16,22 +15,6 @@ pytestmark = pytest.mark.cpu_test -class _FakeMoERunner: - """Minimal MoERunner stand-in for token-count discovery tests.""" - - moe_config: Any - - -def _make_moe(*, max_deferred_tokens: int = 128, enabled: bool = True): - """Create a fake MoE layer with a deferred-finalize token limit.""" - moe = _FakeMoERunner() - moe.moe_config = SimpleNamespace( - use_deferred_moe_finalize=enabled, - defer_moe_finalize_max_num_tokens=max_deferred_tokens, - ) - return moe - - def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto"): """Create a runner carrying only the state used by FlashInfer warmup.""" return SimpleNamespace( @@ -46,35 +29,10 @@ def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto ) -def test_flashinfer_autotune_token_counts_include_deferred_moe_limits(): - runner = _make_runner( - [ - object(), - _make_moe(max_deferred_tokens=128), - _make_moe(max_deferred_tokens=128), - _make_moe(max_deferred_tokens=64, enabled=False), - _make_moe(max_deferred_tokens=-1), - ], - linear_backend="flashinfer_cutedsl", - ) - - with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): - token_counts = _flashinfer_autotune_token_counts(runner) - - assert token_counts == (8192, 32, 128) - - def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated(): - runner = _make_runner( - [_make_moe(max_deferred_tokens=4096)], - max_tokens=32, - linear_backend="flashinfer_cutedsl", - ) - - with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): - token_counts = _flashinfer_autotune_token_counts(runner) + runner = _make_runner([], max_tokens=32, linear_backend="flashinfer_cutedsl") - assert token_counts == (32,) + assert _flashinfer_autotune_token_counts(runner) == (32,) @pytest.mark.parametrize("skip_attn", [False, True]) diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py index 5fde9ff86d4a..cc30ce3193d6 100644 --- a/tests/models/kimi_k3/test_latent_moe_tail.py +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -43,12 +43,6 @@ class FakeMoEConfig: hidden_dim = LATENT_SIZE hidden_dim_unpadded = LATENT_SIZE experts_per_token = 16 - defer_moe_finalize = False - defer_moe_finalize_max_num_tokens = -1 - - @property - def use_deferred_moe_finalize(self) -> bool: - return self.defer_moe_finalize moe_config = FakeMoEConfig() quant_method = SimpleNamespace( @@ -65,6 +59,7 @@ def use_deferred_moe_finalize(self) -> bool: def fake_runner_init(runner, *args, **kwargs) -> None: runner.moe_config = moe_config + runner.router = object() runner.routed_experts = SimpleNamespace(quant_method=quant_method) runner._shared_experts = object() runner.routed_output_transform = transform @@ -97,10 +92,9 @@ def fake_tail_initialize(**kwargs): ) monkeypatch.setattr(KimiK3LatentMoETailOp, "initialize", fake_tail_initialize) - latent_moe_runner.LatentMoERunner() + runner = latent_moe_runner.LatentMoERunner() - assert moe_config.defer_moe_finalize - assert moe_config.defer_moe_finalize_max_num_tokens == 128 + assert runner.defer_moe_finalize assert initialized_with["experts_per_token"] == 16 diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 0c82d28096dc..7a06a523649d 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1290,15 +1290,6 @@ class FusedMoEConfig: # Only honored on the non-reduced (late-AR) TP path. Default False. skip_final_all_reduce: bool = False - # When True, experts that can stop after GEMM2 are allowed to hand back an - # UnfinalizedMoEOutput instead of finalized states, leaving the top-k - # reduction to fuse into the consumer. Set by layers that have such a - # consumer; read through `use_deferred_moe_finalize`, which applies the - # guards. Kernels without the capability ignore it. Default False. - defer_moe_finalize: bool = False - # Optional consumer capacity for deferred finalize. Negative means unbounded. - defer_moe_finalize_max_num_tokens: int = -1 - # SwiGLU clamp limit. When set, backends that do not implement the clamp # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle # cannot silently select one and drop the clamp. @@ -1418,34 +1409,6 @@ def ep_rank(self): def use_ep(self): return self.moe_parallel_config.use_ep - @property - def use_deferred_moe_finalize(self) -> bool: - """Whether experts may return an unfinalized output on this deployment. - - Evaluated on read rather than in ``__post_init__`` because - ``defer_moe_finalize`` is set after construction, like - ``skip_final_all_reduce``. - """ - # The consumer fuses a TP all-reduce. Other parallel modes require a - # combine or reduce-scatter after the experts and cannot defer it. - return ( - self.defer_moe_finalize - and self.tp_size > 1 - and self.dp_size == 1 - and self.ep_size == 1 - and self.pcp_size == 1 - and not self.is_sequence_parallel - ) - - def should_defer_moe_finalize(self, num_tokens: int) -> bool: - """Return whether this invocation may defer the top-k reduction.""" - max_num_tokens = self.defer_moe_finalize_max_num_tokens - return ( - self.use_deferred_moe_finalize - and num_tokens > 0 - and (max_num_tokens < 0 or num_tokens <= max_num_tokens) - ) - @property def use_deepep_ht_kernels(self): return self.moe_parallel_config.use_deepep_ht_kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 54227d3cd1eb..de7bff6454f7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -232,6 +232,10 @@ def apply( class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonolithic): """BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface.""" + @staticmethod + def supports_deferred_moe_finalize() -> bool: + return True + @staticmethod def _supports_parallel_config( moe_parallel_config: FusedMoEParallelConfig, @@ -278,9 +282,7 @@ def apply( assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] num_tokens = hidden_states.shape[0] - # The runner divides by the token count on the host, so an idle rank's - # dummy 0-token forward has to keep the finalized (empty) form. - defer = self.moe_config.should_defer_moe_finalize(num_tokens) + defer = self.defer_moe_finalize routing_replay_out = self._maybe_make_routing_replay_buffer( num_tokens=num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index fb52445d8417..41ee1ec1197b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -199,6 +199,10 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): """Fp8 TRTLLM-Gen MoE kernels. Supports modular interface.""" + @staticmethod + def supports_deferred_moe_finalize() -> bool: + return True + @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: return ( @@ -302,7 +306,7 @@ def apply( hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale) num_tokens = hidden_states.shape[0] - defer = self.moe_config.should_defer_moe_finalize(num_tokens) + defer = self.defer_moe_finalize flashinfer_output = flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe( topk_ids=(topk_ids, topk_weights), routing_bias=None, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index aae8ecff5283..79e17e6bf0b4 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -126,6 +126,10 @@ def _supports_current_device() -> bool: def _supports_no_act_and_mul() -> bool: return False + @staticmethod + def supports_deferred_moe_finalize() -> bool: + return True + @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, @@ -226,7 +230,7 @@ def apply( x_quant = hidden_states x_scale = None num_tokens = hidden_states.shape[0] - defer = self.moe_config.should_defer_moe_finalize(num_tokens) + defer = self.defer_moe_finalize finalized_output = None if not defer: finalized_output = torch.empty( @@ -439,17 +443,20 @@ def apply( assert self.w1_scale is not None assert self.w2_scale is not None - # Chunk tokens so the batched-GEMM grid stays within CUDA limits. Each - # launch permutes into its own buffer, so only a run that fits in one - # launch can leave the top-k reduction to a deferring consumer. + # Chunk tokens so the batched-GEMM grid stays within CUDA limits. M = x_quant.size(0) chunk_size = self._max_supported_tokens(topk, global_num_experts) - defer = chunk_size >= M and self.moe_config.should_defer_moe_finalize(M) + if self.defer_moe_finalize and chunk_size < M: + # Each launch permutes into its own buffer. + raise ValueError( + f"Deferring the MoE finalize needs one kernel launch, but {M} " + f"tokens take {-(-M // chunk_size)}." + ) unfinalized: UnfinalizedMoEOutput | None = None for start in range(0, M, chunk_size): end = min(start + chunk_size, M) unfinalized = self._invoke_kernel( - None if defer else output[start:end], + None if self.defer_moe_finalize else output[start:end], x_quant[start:end], None if x_scale is None else x_scale[start:end], topk_ids[start:end], diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 1423fe444286..274c5879a780 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -218,6 +218,10 @@ def _supports_no_act_and_mul() -> bool: """Supports non-gated MoE (i.e. Nemotron-Nano).""" return True + @staticmethod + def supports_deferred_moe_finalize() -> bool: + return True + @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, @@ -436,12 +440,8 @@ def apply( chunk_size = self._get_chunk_size() if chunk_size >= M: - # Each launch permutes into its own buffer, so only a run that fits - # in one launch can leave the top-k reduction to a deferring - # consumer. - defer = self.moe_config.should_defer_moe_finalize(M) return self._invoke_kernel( - None if defer else output, + None if self.defer_moe_finalize else output, hidden_states, w1, w2, @@ -451,6 +451,12 @@ def apply( global_num_experts, a1q_scale, ) + elif self.defer_moe_finalize: + # Each launch permutes into its own buffer. + raise ValueError( + f"Deferring the MoE finalize needs one kernel launch, but {M} " + f"tokens take {-(-M // chunk_size)}." + ) else: for start in range(0, M, chunk_size): end = min(start + chunk_size, M) @@ -551,9 +557,7 @@ def apply( output1_scale_gate_scalar = self.quant_config.g1_alphas num_tokens = hidden_states.shape[0] - # The runner divides by the token count on the host, so an idle rank's - # dummy 0-token forward has to keep the finalized (empty) form. - defer = self.moe_config.should_defer_moe_finalize(num_tokens) + defer = self.defer_moe_finalize routing_replay_out = self._maybe_make_routing_replay_buffer( num_tokens=num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index f563c819e4b3..8b9946c26032 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -461,6 +461,9 @@ class FusedMoEExperts(ABC): # expert_map. RoutedExperts.expert_map reads this flag to pick which to hand # the active experts kernel. consumes_expert_mask: bool = False + # Stop after GEMM2 on every call and return an UnfinalizedMoEOutput. Set + # once, through FusedMoEKernel.enable_deferred_moe_finalize. + defer_moe_finalize: bool = False def __init__( self, @@ -737,6 +740,13 @@ def supports_lora() -> bool: """ return False + @staticmethod + def supports_deferred_moe_finalize() -> bool: + """Whether this impl can stop after GEMM2 and leave the top-k + reduction to the caller, on every call it gets. + """ + return False + def supports_packed_ue8m0_act_scales(self) -> bool: """A flag indicating whether or not this class can process packed ue8m0 activation scales. @@ -1681,7 +1691,20 @@ def output_is_reduced(self) -> bool: return self.prepare_finalize.output_is_reduced() def supports_deferred_moe_finalize(self) -> bool: - return self.prepare_finalize.supports_deferred_moe_finalize() + return ( + self.prepare_finalize.supports_deferred_moe_finalize() + and self.fused_experts.supports_deferred_moe_finalize() + ) + + def enable_deferred_moe_finalize(self) -> None: + """Return an UnfinalizedMoEOutput from every call from now on.""" + if not self.supports_deferred_moe_finalize(): + raise ValueError( + f"{type(self.fused_experts).__name__} with " + f"{type(self.prepare_finalize).__name__} cannot defer the MoE " + "finalize." + ) + self.fused_experts.defer_moe_finalize = True def apply_monolithic( self, diff --git a/vllm/model_executor/layers/fused_moe/moe_output.py b/vllm/model_executor/layers/fused_moe/moe_output.py index 8f9662b442ce..971153761a98 100644 --- a/vllm/model_executor/layers/fused_moe/moe_output.py +++ b/vllm/model_executor/layers/fused_moe/moe_output.py @@ -124,6 +124,30 @@ def convert_flashinfer_moe_output( ) +def finalize_moe_output(routed: UnfinalizedMoEOutput) -> torch.Tensor: + """Run the top-k reduction a deferring MoE layer left open. + + For the calls a fused consumer can't take. This is the TRT-LLM finalize as + a standalone kernel, so it reduces the same way the experts would have. + A deferring layer has every expert local (``ep_size == 1``), so the permute + map has no -1 entries. + """ + from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( + moe_unpermute, + ) + + gemm2_permuted = routed.gemm2_permuted + num_tokens = routed.expanded_idx_to_permuted_idx.shape[0] + output = gemm2_permuted.new_empty(num_tokens, gemm2_permuted.shape[1]) + moe_unpermute( + output, + gemm2_permuted, + routed.expert_weights.float(), + routed.expanded_idx_to_permuted_idx, + ) + return output + + @dataclass class MoEOutput: """A MoE layer's output with its final reduction still open. @@ -136,10 +160,9 @@ class MoEOutput: unfinalized, the top-k reduction is open too and can fold into the same kernel. - Producers only leave the routed output unfinalized when a fused consumer can - actually take that form -- the token ceiling and topology support are theirs - to check -- so an ``UnfinalizedMoEOutput`` here means the fused path applies, - and a consumer need not re-derive that. + A layer that defers its finalize leaves the routed output unfinalized on + every call. Calls its fused consumer can't take go through + ``finalize_moe_output`` first. """ # Un-reduced routed output, either finalized or not. diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 08279d1c8938..0f86459f2760 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -28,7 +28,10 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) -from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.moe_output import ( + MoEOutput, + UnfinalizedMoEOutput, +) from vllm.model_executor.layers.fused_moe.routed_experts import ( RoutedExperts, ) @@ -247,6 +250,9 @@ class MoERunner(MoERunnerInterface): for different configurations (e.g., with/without shared experts, gates, etc.). """ + # Set once by the model, through enable_deferred_moe_finalize. + defer_moe_finalize: bool = False + def __init__( self, layer_name: str, @@ -423,6 +429,75 @@ def _fused_output_is_reduced(self) -> bool: and self._quant_method.moe_kernel.output_is_reduced() ) + def supports_deferred_moe_finalize(self) -> bool: + """Whether this layer can hand back its routed output unfinalized. + + That takes experts that can stop after GEMM2, and nothing between them + and the caller: no combine or reduce-scatter, no sequence parallelism, + no hidden-dim padding to strip and no zero expert to add. + """ + experts_cls = getattr(self._quant_method, "experts_cls", None) + moe_config = self.moe_config + return ( + experts_cls is not None + and experts_cls.supports_deferred_moe_finalize() + and moe_config.dp_size == 1 + and moe_config.ep_size == 1 + and moe_config.pcp_size == 1 + and not moe_config.is_sequence_parallel + and moe_config.hidden_dim == moe_config.hidden_dim_unpadded + and not isinstance(self.router, ZeroExpertRouter) + ) + + def enable_deferred_moe_finalize(self) -> None: + """Hand back the routed output unfinalized on every call. + + Called once while the model is built. The caller then runs this layer + through ``forward_unfinalized`` and owns everything after the experts. + """ + if not self.supports_deferred_moe_finalize(): + raise ValueError(f"MoE layer {self.layer_name} cannot defer its finalize.") + assert self._quant_method.moe_kernel is None, ( + "the MoE finalize is deferred while the model is built" + ) + self.defer_moe_finalize = True + + def forward_unfinalized( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> MoEOutput: + """``forward`` for a layer that defers its finalize. + + Runs outside the custom op, which returns tensors only. The top-k + reduction, the routed scale and output transform, the shared-expert add + and any all-reduce are the caller's; ``finalize_moe_output`` runs the + reduction when the caller's fused consumer can't. + """ + assert self.defer_moe_finalize + kernel = self._quant_method.moe_kernel + assert kernel is not None + if not kernel.fused_experts.defer_moe_finalize: + # The kernel is built after weight loading, long after the layer. + kernel.enable_deferred_moe_finalize() + if shared_experts_input is None: + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) + shared_output, routed = _unpack( + self._forward_impl( + hidden_states, router_logits, shared_experts_input, input_ids + ) + ) + assert isinstance(routed, UnfinalizedMoEOutput) + return MoEOutput( + routed=routed, + shared_output=shared_output, + routed_scaling_factor=self.routed_scaling_factor, + ) + def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, @@ -708,6 +783,11 @@ def forward( 1. pytorch cannot handle union types in custom op signatures so _moe_forward and _moe_forward_shared must be split. """ + if self.defer_moe_finalize: + raise RuntimeError( + f"MoE layer {self.layer_name} defers its finalize; call " + "forward_unfinalized." + ) # Apply transform for routed experts (e.g., latent projection for # latent MoE). When the caller pre-applies the routed input transform # outside the runner (e.g. to overlap it on a separate stream), it diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 33fe7002aa2f..745912f3861f 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -311,26 +311,6 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS = 32 -def _flashinfer_deferred_moe_token_counts( - runner: "GPUModelRunner", -) -> tuple[int, ...]: - """Return bounded token counts that exercise deferred MoE finalization.""" - from vllm.model_executor.layers.fused_moe import MoERunner - - max_tokens = runner.scheduler_config.max_num_batched_tokens - token_counts: list[int] = [] - for module in runner.get_model().modules(): - if not isinstance(module, MoERunner): - continue - - moe_config = module.moe_config - max_deferred_tokens = moe_config.defer_moe_finalize_max_num_tokens - if moe_config.use_deferred_moe_finalize and max_deferred_tokens > 0: - token_counts.append(min(max_tokens, max_deferred_tokens)) - - return tuple(dict.fromkeys(token_counts)) - - def _flashinfer_autotune_token_counts( runner: "GPUModelRunner", *, include_bf16: bool = True ) -> tuple[int, ...]: @@ -344,8 +324,7 @@ def _flashinfer_autotune_token_counts( and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS ): token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS) - token_counts.extend(_flashinfer_deferred_moe_token_counts(runner)) - return tuple(dict.fromkeys(token_counts)) + return tuple(token_counts) def _run_flashinfer_autotune_dummy_runs( diff --git a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py index 8eedb5ba89ed..68985bf6a871 100644 --- a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py +++ b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py @@ -12,8 +12,11 @@ tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner, _unpack +from vllm.model_executor.layers.fused_moe.moe_output import ( + UnfinalizedMoEOutput, + finalize_moe_output, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner from vllm.model_executor.layers.layernorm import RMSNorm from vllm.platforms import current_platform from vllm.utils.multi_stream_utils import maybe_execute_in_parallel @@ -115,10 +118,13 @@ def __init__( experts_cls = getattr(self._quant_method, "experts_cls", None) # The kernel instance is built after weight loading, while the tail # must register its CuTeDSL warmup units during runner construction. - self.moe_config.defer_moe_finalize = ( + # The tail reads the routing weights in the activation dtype, as the + # monolithic kernels return them. + if ( experts_cls is TrtLlmMxfp4ExpertsMonolithic or experts_cls is TrtLlmNvFp4ExpertsMonolithic - ) and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded + ) and self.supports_deferred_moe_finalize(): + self.enable_deferred_moe_finalize() from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import ( KimiK3LatentMoETailOp, ) @@ -127,29 +133,19 @@ def __init__( hidden_size=transform.up_proj.weight.shape[0], latent_size=norm.weight.shape[0], experts_per_token=( - self.moe_config.experts_per_token - if self.moe_config.use_deferred_moe_finalize - else 0 + self.moe_config.experts_per_token if self.defer_moe_finalize else 0 ), dtype=norm.weight.dtype, device=current_platform.current_device(), rms_eps=norm.variance_epsilon, ) self._k3_latent_moe_tail_op = op - self.moe_config.defer_moe_finalize_max_num_tokens = ( - op.contract.max_num_tokens - if self.moe_config.use_deferred_moe_finalize - else -1 - ) - if self.moe_config.use_deferred_moe_finalize: + if self.defer_moe_finalize: logger.info_once( - "K3 latent-MoE tail fusion with deferred top-k finalization " - "is enabled for up to %d tokens.", + "K3 latent-MoE tail fusion takes the deferred top-k " + "finalization for up to %d tokens.", op.contract.max_num_tokens, ) - else: - self.moe_config.defer_moe_finalize = False - self.moe_config.defer_moe_finalize_max_num_tokens = -1 def _get_zero_residual( self, @@ -337,37 +333,35 @@ def _fused_forward( ) ) - if ( - og_hidden_dim_pre_xform is None - and self.moe_config.should_defer_moe_finalize(hidden_states.shape[0]) - ): - unfinalized = self._forward_impl( + if self.defer_moe_finalize: + moe_output = self.forward_unfinalized( + hidden_states, router_logits, input_ids, shared_experts_input + ) + routed = moe_output.routed + shared_output = moe_output.shared_output + assert isinstance(routed, UnfinalizedMoEOutput) + assert shared_output is not None + num_tokens = hidden_states.shape[0] + if 0 < num_tokens <= self._k3_latent_moe_tail_op.contract.max_num_tokens: + result = self._small_batch_tail(routed, shared_output, None) + return self._maybe_add_zero_expert_output(result) + fused_output = finalize_moe_output(routed) + else: + result = self._forward_entry( hidden_states, router_logits, shared_experts_input, input_ids, + self._encode_layer_name(), + self.moe_config.hidden_dim_unpadded + if self._quant_method.has_unpadded_output + else 0, ) - shared_output, fused_output = _unpack(unfinalized) - assert shared_output is not None - assert isinstance(fused_output, UnfinalizedMoEOutput) - result = self._small_batch_tail(fused_output, shared_output, None) - return self._maybe_add_zero_expert_output(result) - - result = self._forward_entry( - hidden_states, - router_logits, - shared_experts_input, - input_ids, - self._encode_layer_name(), - self.moe_config.hidden_dim_unpadded - if self._quant_method.has_unpadded_output - else 0, - ) - - shared_output, fused_output = cast(tuple[torch.Tensor, torch.Tensor], result) - - if og_hidden_dim_pre_xform is not None: - fused_output = fused_output[..., :og_hidden_dim_pre_xform] + shared_output, fused_output = cast( + tuple[torch.Tensor, torch.Tensor], result + ) + if og_hidden_dim_pre_xform is not None: + fused_output = fused_output[..., :og_hidden_dim_pre_xform] tier = self._select_tail_tier(fused_output, shared_output) if tier is LatentTailTier.TAIL_FUSION: From a6d37d8e6d59c87cce21eb75a166d1a1d9e78c7d Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 25 Sep 2026 00:33:48 +0000 Subject: [PATCH 3/8] [MoE] Keep Kimi-K3 on its per-call deferral knob for now Restore FusedMoEConfig's defer_moe_finalize knob, its token cap and the warmup's bounded autotune passes, and leave Kimi-K3's runner as it is on main. The monolithic experts honor either the knob or the per-layer flag; the modular experts, which have no knob user, read only the flag. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yongye Zhu --- .../test_flashinfer_autotune_warmup.py | 46 ++++++++++- tests/models/kimi_k3/test_latent_moe_tail.py | 12 ++- .../model_executor/layers/fused_moe/config.py | 37 +++++++++ .../fused_moe/experts/trtllm_bf16_moe.py | 6 +- .../fused_moe/experts/trtllm_mxfp4_moe.py | 4 +- .../fused_moe/experts/trtllm_nvfp4_moe.py | 6 +- .../layers/fused_moe/modular_kernel.py | 4 +- vllm/model_executor/warmup/kernel_warmup.py | 23 +++++- .../kimi_k3/nvidia/latent_moe_runner.py | 82 ++++++++++--------- 9 files changed, 172 insertions(+), 48 deletions(-) diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index b6f43db9d6a6..dd162629b8f4 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -3,6 +3,7 @@ from inspect import signature from types import SimpleNamespace +from typing import Any from unittest.mock import Mock, call, patch import pytest @@ -15,6 +16,22 @@ pytestmark = pytest.mark.cpu_test +class _FakeMoERunner: + """Minimal MoERunner stand-in for token-count discovery tests.""" + + moe_config: Any + + +def _make_moe(*, max_deferred_tokens: int = 128, enabled: bool = True): + """Create a fake MoE layer with a deferred-finalize token limit.""" + moe = _FakeMoERunner() + moe.moe_config = SimpleNamespace( + use_deferred_moe_finalize=enabled, + defer_moe_finalize_max_num_tokens=max_deferred_tokens, + ) + return moe + + def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto"): """Create a runner carrying only the state used by FlashInfer warmup.""" return SimpleNamespace( @@ -29,10 +46,35 @@ def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto ) +def test_flashinfer_autotune_token_counts_include_deferred_moe_limits(): + runner = _make_runner( + [ + object(), + _make_moe(max_deferred_tokens=128), + _make_moe(max_deferred_tokens=128), + _make_moe(max_deferred_tokens=64, enabled=False), + _make_moe(max_deferred_tokens=-1), + ], + linear_backend="flashinfer_cutedsl", + ) + + with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): + token_counts = _flashinfer_autotune_token_counts(runner) + + assert token_counts == (8192, 32, 128) + + def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated(): - runner = _make_runner([], max_tokens=32, linear_backend="flashinfer_cutedsl") + runner = _make_runner( + [_make_moe(max_deferred_tokens=4096)], + max_tokens=32, + linear_backend="flashinfer_cutedsl", + ) + + with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): + token_counts = _flashinfer_autotune_token_counts(runner) - assert _flashinfer_autotune_token_counts(runner) == (32,) + assert token_counts == (32,) @pytest.mark.parametrize("skip_attn", [False, True]) diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py index cc30ce3193d6..5fde9ff86d4a 100644 --- a/tests/models/kimi_k3/test_latent_moe_tail.py +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -43,6 +43,12 @@ class FakeMoEConfig: hidden_dim = LATENT_SIZE hidden_dim_unpadded = LATENT_SIZE experts_per_token = 16 + defer_moe_finalize = False + defer_moe_finalize_max_num_tokens = -1 + + @property + def use_deferred_moe_finalize(self) -> bool: + return self.defer_moe_finalize moe_config = FakeMoEConfig() quant_method = SimpleNamespace( @@ -59,7 +65,6 @@ class FakeMoEConfig: def fake_runner_init(runner, *args, **kwargs) -> None: runner.moe_config = moe_config - runner.router = object() runner.routed_experts = SimpleNamespace(quant_method=quant_method) runner._shared_experts = object() runner.routed_output_transform = transform @@ -92,9 +97,10 @@ def fake_tail_initialize(**kwargs): ) monkeypatch.setattr(KimiK3LatentMoETailOp, "initialize", fake_tail_initialize) - runner = latent_moe_runner.LatentMoERunner() + latent_moe_runner.LatentMoERunner() - assert runner.defer_moe_finalize + assert moe_config.defer_moe_finalize + assert moe_config.defer_moe_finalize_max_num_tokens == 128 assert initialized_with["experts_per_token"] == 16 diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 7a06a523649d..0c82d28096dc 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1290,6 +1290,15 @@ class FusedMoEConfig: # Only honored on the non-reduced (late-AR) TP path. Default False. skip_final_all_reduce: bool = False + # When True, experts that can stop after GEMM2 are allowed to hand back an + # UnfinalizedMoEOutput instead of finalized states, leaving the top-k + # reduction to fuse into the consumer. Set by layers that have such a + # consumer; read through `use_deferred_moe_finalize`, which applies the + # guards. Kernels without the capability ignore it. Default False. + defer_moe_finalize: bool = False + # Optional consumer capacity for deferred finalize. Negative means unbounded. + defer_moe_finalize_max_num_tokens: int = -1 + # SwiGLU clamp limit. When set, backends that do not implement the clamp # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle # cannot silently select one and drop the clamp. @@ -1409,6 +1418,34 @@ def ep_rank(self): def use_ep(self): return self.moe_parallel_config.use_ep + @property + def use_deferred_moe_finalize(self) -> bool: + """Whether experts may return an unfinalized output on this deployment. + + Evaluated on read rather than in ``__post_init__`` because + ``defer_moe_finalize`` is set after construction, like + ``skip_final_all_reduce``. + """ + # The consumer fuses a TP all-reduce. Other parallel modes require a + # combine or reduce-scatter after the experts and cannot defer it. + return ( + self.defer_moe_finalize + and self.tp_size > 1 + and self.dp_size == 1 + and self.ep_size == 1 + and self.pcp_size == 1 + and not self.is_sequence_parallel + ) + + def should_defer_moe_finalize(self, num_tokens: int) -> bool: + """Return whether this invocation may defer the top-k reduction.""" + max_num_tokens = self.defer_moe_finalize_max_num_tokens + return ( + self.use_deferred_moe_finalize + and num_tokens > 0 + and (max_num_tokens < 0 or num_tokens <= max_num_tokens) + ) + @property def use_deepep_ht_kernels(self): return self.moe_parallel_config.use_deepep_ht_kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index de7bff6454f7..f5f5aba6a539 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -282,7 +282,11 @@ def apply( assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] num_tokens = hidden_states.shape[0] - defer = self.defer_moe_finalize + # The runner divides by the token count on the host, so an idle rank's + # dummy 0-token forward has to keep the finalized (empty) form. + defer = self.defer_moe_finalize or self.moe_config.should_defer_moe_finalize( + num_tokens + ) routing_replay_out = self._maybe_make_routing_replay_buffer( num_tokens=num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 79e17e6bf0b4..105809c3acaf 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -230,7 +230,9 @@ def apply( x_quant = hidden_states x_scale = None num_tokens = hidden_states.shape[0] - defer = self.defer_moe_finalize + defer = self.defer_moe_finalize or self.moe_config.should_defer_moe_finalize( + num_tokens + ) finalized_output = None if not defer: finalized_output = torch.empty( diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 274c5879a780..3d472f07b54e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -557,7 +557,11 @@ def apply( output1_scale_gate_scalar = self.quant_config.g1_alphas num_tokens = hidden_states.shape[0] - defer = self.defer_moe_finalize + # The runner divides by the token count on the host, so an idle rank's + # dummy 0-token forward has to keep the finalized (empty) form. + defer = self.defer_moe_finalize or self.moe_config.should_defer_moe_finalize( + num_tokens + ) routing_replay_out = self._maybe_make_routing_replay_buffer( num_tokens=num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 8b9946c26032..6ac1682bb3da 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -462,7 +462,9 @@ class FusedMoEExperts(ABC): # the active experts kernel. consumes_expert_mask: bool = False # Stop after GEMM2 on every call and return an UnfinalizedMoEOutput. Set - # once, through FusedMoEKernel.enable_deferred_moe_finalize. + # once, through FusedMoEKernel.enable_deferred_moe_finalize -- unlike + # moe_config.should_defer_moe_finalize, a per-call choice that only the + # monolithic experts read. defer_moe_finalize: bool = False def __init__( diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 745912f3861f..33fe7002aa2f 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -311,6 +311,26 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS = 32 +def _flashinfer_deferred_moe_token_counts( + runner: "GPUModelRunner", +) -> tuple[int, ...]: + """Return bounded token counts that exercise deferred MoE finalization.""" + from vllm.model_executor.layers.fused_moe import MoERunner + + max_tokens = runner.scheduler_config.max_num_batched_tokens + token_counts: list[int] = [] + for module in runner.get_model().modules(): + if not isinstance(module, MoERunner): + continue + + moe_config = module.moe_config + max_deferred_tokens = moe_config.defer_moe_finalize_max_num_tokens + if moe_config.use_deferred_moe_finalize and max_deferred_tokens > 0: + token_counts.append(min(max_tokens, max_deferred_tokens)) + + return tuple(dict.fromkeys(token_counts)) + + def _flashinfer_autotune_token_counts( runner: "GPUModelRunner", *, include_bf16: bool = True ) -> tuple[int, ...]: @@ -324,7 +344,8 @@ def _flashinfer_autotune_token_counts( and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS ): token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS) - return tuple(token_counts) + token_counts.extend(_flashinfer_deferred_moe_token_counts(runner)) + return tuple(dict.fromkeys(token_counts)) def _run_flashinfer_autotune_dummy_runs( diff --git a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py index 68985bf6a871..8eedb5ba89ed 100644 --- a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py +++ b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py @@ -12,11 +12,8 @@ tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.moe_output import ( - UnfinalizedMoEOutput, - finalize_moe_output, -) -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner, _unpack from vllm.model_executor.layers.layernorm import RMSNorm from vllm.platforms import current_platform from vllm.utils.multi_stream_utils import maybe_execute_in_parallel @@ -118,13 +115,10 @@ def __init__( experts_cls = getattr(self._quant_method, "experts_cls", None) # The kernel instance is built after weight loading, while the tail # must register its CuTeDSL warmup units during runner construction. - # The tail reads the routing weights in the activation dtype, as the - # monolithic kernels return them. - if ( + self.moe_config.defer_moe_finalize = ( experts_cls is TrtLlmMxfp4ExpertsMonolithic or experts_cls is TrtLlmNvFp4ExpertsMonolithic - ) and self.supports_deferred_moe_finalize(): - self.enable_deferred_moe_finalize() + ) and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import ( KimiK3LatentMoETailOp, ) @@ -133,19 +127,29 @@ def __init__( hidden_size=transform.up_proj.weight.shape[0], latent_size=norm.weight.shape[0], experts_per_token=( - self.moe_config.experts_per_token if self.defer_moe_finalize else 0 + self.moe_config.experts_per_token + if self.moe_config.use_deferred_moe_finalize + else 0 ), dtype=norm.weight.dtype, device=current_platform.current_device(), rms_eps=norm.variance_epsilon, ) self._k3_latent_moe_tail_op = op - if self.defer_moe_finalize: + self.moe_config.defer_moe_finalize_max_num_tokens = ( + op.contract.max_num_tokens + if self.moe_config.use_deferred_moe_finalize + else -1 + ) + if self.moe_config.use_deferred_moe_finalize: logger.info_once( - "K3 latent-MoE tail fusion takes the deferred top-k " - "finalization for up to %d tokens.", + "K3 latent-MoE tail fusion with deferred top-k finalization " + "is enabled for up to %d tokens.", op.contract.max_num_tokens, ) + else: + self.moe_config.defer_moe_finalize = False + self.moe_config.defer_moe_finalize_max_num_tokens = -1 def _get_zero_residual( self, @@ -333,35 +337,37 @@ def _fused_forward( ) ) - if self.defer_moe_finalize: - moe_output = self.forward_unfinalized( - hidden_states, router_logits, input_ids, shared_experts_input - ) - routed = moe_output.routed - shared_output = moe_output.shared_output - assert isinstance(routed, UnfinalizedMoEOutput) - assert shared_output is not None - num_tokens = hidden_states.shape[0] - if 0 < num_tokens <= self._k3_latent_moe_tail_op.contract.max_num_tokens: - result = self._small_batch_tail(routed, shared_output, None) - return self._maybe_add_zero_expert_output(result) - fused_output = finalize_moe_output(routed) - else: - result = self._forward_entry( + if ( + og_hidden_dim_pre_xform is None + and self.moe_config.should_defer_moe_finalize(hidden_states.shape[0]) + ): + unfinalized = self._forward_impl( hidden_states, router_logits, shared_experts_input, input_ids, - self._encode_layer_name(), - self.moe_config.hidden_dim_unpadded - if self._quant_method.has_unpadded_output - else 0, ) - shared_output, fused_output = cast( - tuple[torch.Tensor, torch.Tensor], result - ) - if og_hidden_dim_pre_xform is not None: - fused_output = fused_output[..., :og_hidden_dim_pre_xform] + shared_output, fused_output = _unpack(unfinalized) + assert shared_output is not None + assert isinstance(fused_output, UnfinalizedMoEOutput) + result = self._small_batch_tail(fused_output, shared_output, None) + return self._maybe_add_zero_expert_output(result) + + result = self._forward_entry( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + self._encode_layer_name(), + self.moe_config.hidden_dim_unpadded + if self._quant_method.has_unpadded_output + else 0, + ) + + shared_output, fused_output = cast(tuple[torch.Tensor, torch.Tensor], result) + + if og_hidden_dim_pre_xform is not None: + fused_output = fused_output[..., :og_hidden_dim_pre_xform] tier = self._select_tail_tier(fused_output, shared_output) if tier is LatentTailTier.TAIL_FUSION: From 3362604a3624abd2e06419d5300e0f58a828665b Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 25 Sep 2026 00:48:44 +0000 Subject: [PATCH 4/8] [MoE] Keep deferred MoE finalize in FusedMoEConfig Drop the per-layer flag on the runner and the experts, and read deferral from moe_config again, the way Kimi-K3 already does. A layer that defers on every call sets defer_moe_finalize once at init with no token cap; the modular experts raise instead of chunking such a run, and finalize_moe_output stays for the calls a fused consumer can't take. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yongye Zhu --- tests/kernels/moe/test_flashinfer.py | 11 ++- tests/kernels/moe/test_ocp_mx_moe.py | 11 ++- tests/kernels/moe/test_trtllm_nvfp4_moe.py | 11 ++- tests/kernels/moe/utils.py | 14 ++-- .../fused_moe/experts/trtllm_bf16_moe.py | 8 +- .../fused_moe/experts/trtllm_fp8_moe.py | 6 +- .../fused_moe/experts/trtllm_mxfp4_moe.py | 13 +-- .../fused_moe/experts/trtllm_nvfp4_moe.py | 13 +-- .../layers/fused_moe/modular_kernel.py | 27 +----- .../layers/fused_moe/moe_output.py | 7 +- .../layers/fused_moe/runner/moe_runner.py | 82 +------------------ 11 files changed, 45 insertions(+), 158 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index d26360737b07..fbc45d41a9bd 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import SimpleNamespace from unittest.mock import MagicMock @@ -813,16 +813,19 @@ def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): quant_config = fp8_w8a8_moe_quant_config( w1_scale=w1_scale, w2_scale=w2_scale, block_shape=block_shape ) + # One rank of a TP group, which deferral needs. moe_config = FusedMoEConfig( num_experts=e, experts_per_token=topk, hidden_dim=k, - intermediate_size=n, + intermediate_size=2 * n, num_local_experts=e, num_logical_experts=e, activation=MoEActivation.SILU, device="cuda", - moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + moe_parallel_config=replace( + FusedMoEParallelConfig.make_no_parallel(), tp_size=2 + ), in_dtype=torch.bfloat16, routing_method=RoutingMethodType.TopK, max_num_tokens=next_power_of_2(m), @@ -838,7 +841,7 @@ def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) check_deferred_moe_finalize( - kernel, + moe_config, lambda: kernel.apply( hidden_states=a, w1=w1, diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index bddefe4e51b6..e4b11c9548b2 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -747,17 +747,20 @@ def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: w2_scale, _cache_permute_indices={}, ) + # One rank of a TP group, which deferral needs. moe_config = FusedMoEConfig( num_experts=num_experts, experts_per_token=topk, hidden_dim=hidden_size, - intermediate_size=intermediate_size, + intermediate_size=2 * intermediate_size, num_local_experts=num_experts, num_logical_experts=num_experts, activation=MoEActivation.SILU, device=device, routing_method=RoutingMethodType.Renormalize, - moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + moe_parallel_config=replace( + FusedMoEParallelConfig.make_no_parallel(), tp_size=2 + ), in_dtype=torch.bfloat16, ) kernel = make_mxfp4_moe_kernel( @@ -773,7 +776,7 @@ def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: if monolithic: router_logits = torch.randn(num_tokens, num_experts, device=device) check_deferred_moe_finalize( - kernel, + moe_config, lambda: kernel.apply_monolithic( hidden_states, w13, @@ -791,7 +794,7 @@ def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: topk_ids = topk_ids[:, :topk].to(torch.int32) topk_weights = torch.rand(num_tokens, topk, device=device).softmax(dim=-1) check_deferred_moe_finalize( - kernel, + moe_config, lambda: kernel.apply( hidden_states=hidden_states, w1=w13, diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index dc9c9e2ca54e..292e4549bb7d 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -8,6 +8,8 @@ intermediate_size 704) that exercises the non-256-aligned padding path. """ +from dataclasses import replace + import pytest import torch @@ -366,16 +368,19 @@ def test_trtllm_fp4_moe_deferred_finalize( score = torch.randn((m, e), device="cuda", dtype=dtype) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + # One rank of a TP group, which deferral needs. moe_config = FusedMoEConfig( num_experts=e, experts_per_token=topk, hidden_dim=k, - intermediate_size=n, + intermediate_size=2 * n, num_local_experts=e, num_logical_experts=e, activation=MoEActivation.SILU, device="cuda", - moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + moe_parallel_config=replace( + FusedMoEParallelConfig.make_no_parallel(), tp_size=2 + ), in_dtype=dtype, routing_method=RoutingMethodType.TopK, max_num_tokens=next_power_of_2(m), @@ -383,7 +388,7 @@ def test_trtllm_fp4_moe_deferred_finalize( kernel = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) check_deferred_moe_finalize( - kernel, + moe_config, lambda: kernel.apply( hidden_states=a, w1=w1_q, diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 9a3386b0a611..4979161533a8 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -714,22 +714,22 @@ def check_accuracy(a, b, atol, rtol, percent): def check_deferred_moe_finalize( - kernel: FusedMoEKernel, + moe_config: FusedMoEConfig, run: Callable[[], torch.Tensor | UnfinalizedMoEOutput], router_weights: torch.Tensor | None = None, chunked: bool = False, ) -> None: """Check a kernel that defers its finalize against the finalize it skips. - ``run`` calls ``kernel`` on fixed inputs, first as built and then with the - finalize deferred. ``finalize_moe_output`` on the deferred output must give - the kernel's own finalized output bit for bit, and modular experts must hand - the router's weights back as-is. Each launch permutes into its own buffer, - so a deferring kernel refuses a run it would have to chunk. + ``run`` calls the kernel on fixed inputs, first as built and then with + ``moe_config`` deferring on every call. ``finalize_moe_output`` on the + deferred output must give the kernel's own finalized output bit for bit, and + modular experts must hand the router's weights back as-is. Each launch + permutes into its own buffer, so deferring refuses a run it would chunk. """ finalized = run() assert isinstance(finalized, torch.Tensor) - kernel.enable_deferred_moe_finalize() + moe_config.defer_moe_finalize = True if chunked: with pytest.raises(ValueError, match="one kernel launch"): run() diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index f5f5aba6a539..54227d3cd1eb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -232,10 +232,6 @@ def apply( class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonolithic): """BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface.""" - @staticmethod - def supports_deferred_moe_finalize() -> bool: - return True - @staticmethod def _supports_parallel_config( moe_parallel_config: FusedMoEParallelConfig, @@ -284,9 +280,7 @@ def apply( num_tokens = hidden_states.shape[0] # The runner divides by the token count on the host, so an idle rank's # dummy 0-token forward has to keep the finalized (empty) form. - defer = self.defer_moe_finalize or self.moe_config.should_defer_moe_finalize( - num_tokens - ) + defer = self.moe_config.should_defer_moe_finalize(num_tokens) routing_replay_out = self._maybe_make_routing_replay_buffer( num_tokens=num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 41ee1ec1197b..fb52445d8417 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -199,10 +199,6 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): """Fp8 TRTLLM-Gen MoE kernels. Supports modular interface.""" - @staticmethod - def supports_deferred_moe_finalize() -> bool: - return True - @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: return ( @@ -306,7 +302,7 @@ def apply( hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale) num_tokens = hidden_states.shape[0] - defer = self.defer_moe_finalize + defer = self.moe_config.should_defer_moe_finalize(num_tokens) flashinfer_output = flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe( topk_ids=(topk_ids, topk_weights), routing_bias=None, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 105809c3acaf..b0ddfa2b430d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -126,10 +126,6 @@ def _supports_current_device() -> bool: def _supports_no_act_and_mul() -> bool: return False - @staticmethod - def supports_deferred_moe_finalize() -> bool: - return True - @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, @@ -230,9 +226,7 @@ def apply( x_quant = hidden_states x_scale = None num_tokens = hidden_states.shape[0] - defer = self.defer_moe_finalize or self.moe_config.should_defer_moe_finalize( - num_tokens - ) + defer = self.moe_config.should_defer_moe_finalize(num_tokens) finalized_output = None if not defer: finalized_output = torch.empty( @@ -448,7 +442,8 @@ def apply( # Chunk tokens so the batched-GEMM grid stays within CUDA limits. M = x_quant.size(0) chunk_size = self._max_supported_tokens(topk, global_num_experts) - if self.defer_moe_finalize and chunk_size < M: + defer = self.moe_config.should_defer_moe_finalize(M) + if defer and chunk_size < M: # Each launch permutes into its own buffer. raise ValueError( f"Deferring the MoE finalize needs one kernel launch, but {M} " @@ -458,7 +453,7 @@ def apply( for start in range(0, M, chunk_size): end = min(start + chunk_size, M) unfinalized = self._invoke_kernel( - None if self.defer_moe_finalize else output[start:end], + None if defer else output[start:end], x_quant[start:end], None if x_scale is None else x_scale[start:end], topk_ids[start:end], diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 3d472f07b54e..10477515e343 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -218,10 +218,6 @@ def _supports_no_act_and_mul() -> bool: """Supports non-gated MoE (i.e. Nemotron-Nano).""" return True - @staticmethod - def supports_deferred_moe_finalize() -> bool: - return True - @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, @@ -438,10 +434,11 @@ def apply( M = hidden_states.shape[0] chunk_size = self._get_chunk_size() + defer = self.moe_config.should_defer_moe_finalize(M) if chunk_size >= M: return self._invoke_kernel( - None if self.defer_moe_finalize else output, + None if defer else output, hidden_states, w1, w2, @@ -451,7 +448,7 @@ def apply( global_num_experts, a1q_scale, ) - elif self.defer_moe_finalize: + elif defer: # Each launch permutes into its own buffer. raise ValueError( f"Deferring the MoE finalize needs one kernel launch, but {M} " @@ -559,9 +556,7 @@ def apply( num_tokens = hidden_states.shape[0] # The runner divides by the token count on the host, so an idle rank's # dummy 0-token forward has to keep the finalized (empty) form. - defer = self.defer_moe_finalize or self.moe_config.should_defer_moe_finalize( - num_tokens - ) + defer = self.moe_config.should_defer_moe_finalize(num_tokens) routing_replay_out = self._maybe_make_routing_replay_buffer( num_tokens=num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 6ac1682bb3da..f563c819e4b3 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -461,11 +461,6 @@ class FusedMoEExperts(ABC): # expert_map. RoutedExperts.expert_map reads this flag to pick which to hand # the active experts kernel. consumes_expert_mask: bool = False - # Stop after GEMM2 on every call and return an UnfinalizedMoEOutput. Set - # once, through FusedMoEKernel.enable_deferred_moe_finalize -- unlike - # moe_config.should_defer_moe_finalize, a per-call choice that only the - # monolithic experts read. - defer_moe_finalize: bool = False def __init__( self, @@ -742,13 +737,6 @@ def supports_lora() -> bool: """ return False - @staticmethod - def supports_deferred_moe_finalize() -> bool: - """Whether this impl can stop after GEMM2 and leave the top-k - reduction to the caller, on every call it gets. - """ - return False - def supports_packed_ue8m0_act_scales(self) -> bool: """A flag indicating whether or not this class can process packed ue8m0 activation scales. @@ -1693,20 +1681,7 @@ def output_is_reduced(self) -> bool: return self.prepare_finalize.output_is_reduced() def supports_deferred_moe_finalize(self) -> bool: - return ( - self.prepare_finalize.supports_deferred_moe_finalize() - and self.fused_experts.supports_deferred_moe_finalize() - ) - - def enable_deferred_moe_finalize(self) -> None: - """Return an UnfinalizedMoEOutput from every call from now on.""" - if not self.supports_deferred_moe_finalize(): - raise ValueError( - f"{type(self.fused_experts).__name__} with " - f"{type(self.prepare_finalize).__name__} cannot defer the MoE " - "finalize." - ) - self.fused_experts.defer_moe_finalize = True + return self.prepare_finalize.supports_deferred_moe_finalize() def apply_monolithic( self, diff --git a/vllm/model_executor/layers/fused_moe/moe_output.py b/vllm/model_executor/layers/fused_moe/moe_output.py index 971153761a98..814fea309fbb 100644 --- a/vllm/model_executor/layers/fused_moe/moe_output.py +++ b/vllm/model_executor/layers/fused_moe/moe_output.py @@ -160,9 +160,10 @@ class MoEOutput: unfinalized, the top-k reduction is open too and can fold into the same kernel. - A layer that defers its finalize leaves the routed output unfinalized on - every call. Calls its fused consumer can't take go through - ``finalize_moe_output`` first. + Producers only leave the routed output unfinalized when the layer asked for + it -- the token ceiling and topology support are its to check. Without a + ceiling that is every call, and ``finalize_moe_output`` covers the calls a + fused consumer can't take. """ # Un-reduced routed output, either finalized or not. diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 0f86459f2760..08279d1c8938 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -28,10 +28,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) -from vllm.model_executor.layers.fused_moe.moe_output import ( - MoEOutput, - UnfinalizedMoEOutput, -) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.routed_experts import ( RoutedExperts, ) @@ -250,9 +247,6 @@ class MoERunner(MoERunnerInterface): for different configurations (e.g., with/without shared experts, gates, etc.). """ - # Set once by the model, through enable_deferred_moe_finalize. - defer_moe_finalize: bool = False - def __init__( self, layer_name: str, @@ -429,75 +423,6 @@ def _fused_output_is_reduced(self) -> bool: and self._quant_method.moe_kernel.output_is_reduced() ) - def supports_deferred_moe_finalize(self) -> bool: - """Whether this layer can hand back its routed output unfinalized. - - That takes experts that can stop after GEMM2, and nothing between them - and the caller: no combine or reduce-scatter, no sequence parallelism, - no hidden-dim padding to strip and no zero expert to add. - """ - experts_cls = getattr(self._quant_method, "experts_cls", None) - moe_config = self.moe_config - return ( - experts_cls is not None - and experts_cls.supports_deferred_moe_finalize() - and moe_config.dp_size == 1 - and moe_config.ep_size == 1 - and moe_config.pcp_size == 1 - and not moe_config.is_sequence_parallel - and moe_config.hidden_dim == moe_config.hidden_dim_unpadded - and not isinstance(self.router, ZeroExpertRouter) - ) - - def enable_deferred_moe_finalize(self) -> None: - """Hand back the routed output unfinalized on every call. - - Called once while the model is built. The caller then runs this layer - through ``forward_unfinalized`` and owns everything after the experts. - """ - if not self.supports_deferred_moe_finalize(): - raise ValueError(f"MoE layer {self.layer_name} cannot defer its finalize.") - assert self._quant_method.moe_kernel is None, ( - "the MoE finalize is deferred while the model is built" - ) - self.defer_moe_finalize = True - - def forward_unfinalized( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - input_ids: torch.Tensor | None = None, - shared_experts_input: torch.Tensor | None = None, - ) -> MoEOutput: - """``forward`` for a layer that defers its finalize. - - Runs outside the custom op, which returns tensors only. The top-k - reduction, the routed scale and output transform, the shared-expert add - and any all-reduce are the caller's; ``finalize_moe_output`` runs the - reduction when the caller's fused consumer can't. - """ - assert self.defer_moe_finalize - kernel = self._quant_method.moe_kernel - assert kernel is not None - if not kernel.fused_experts.defer_moe_finalize: - # The kernel is built after weight loading, long after the layer. - kernel.enable_deferred_moe_finalize() - if shared_experts_input is None: - hidden_states, shared_experts_input = self.apply_routed_input_transform( - hidden_states - ) - shared_output, routed = _unpack( - self._forward_impl( - hidden_states, router_logits, shared_experts_input, input_ids - ) - ) - assert isinstance(routed, UnfinalizedMoEOutput) - return MoEOutput( - routed=routed, - shared_output=shared_output, - routed_scaling_factor=self.routed_scaling_factor, - ) - def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, @@ -783,11 +708,6 @@ def forward( 1. pytorch cannot handle union types in custom op signatures so _moe_forward and _moe_forward_shared must be split. """ - if self.defer_moe_finalize: - raise RuntimeError( - f"MoE layer {self.layer_name} defers its finalize; call " - "forward_unfinalized." - ) # Apply transform for routed experts (e.g., latent projection for # latent MoE). When the caller pre-applies the routed input transform # outside the runner (e.g. to overlap it on a separate stream), it From 85a765355c62837f7684f263bf79efa087727353 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 25 Sep 2026 00:59:55 +0000 Subject: [PATCH 5/8] [MoE] Finalize chunked calls of a deferring layer defer_moe_finalize is a request. The modular MXFP4 and NVFP4 experts lower defer_moe_finalize_max_num_tokens to what they run in one kernel launch, so a call they would chunk finalizes as before instead of raising, and should_defer_moe_finalize tells the model beforehand which form a call returns. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yongye Zhu --- tests/kernels/moe/utils.py | 22 ++++++++--------- .../model_executor/layers/fused_moe/config.py | 9 ++++++- .../fused_moe/experts/trtllm_mxfp4_moe.py | 24 ++++++++++++------- .../fused_moe/experts/trtllm_nvfp4_moe.py | 21 ++++++++++------ 4 files changed, 49 insertions(+), 27 deletions(-) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 4979161533a8..63cafe20ea93 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -import pytest import torch import vllm._custom_ops as ops @@ -722,26 +721,27 @@ def check_deferred_moe_finalize( """Check a kernel that defers its finalize against the finalize it skips. ``run`` calls the kernel on fixed inputs, first as built and then with - ``moe_config`` deferring on every call. ``finalize_moe_output`` on the - deferred output must give the kernel's own finalized output bit for bit, and - modular experts must hand the router's weights back as-is. Each launch - permutes into its own buffer, so deferring refuses a run it would chunk. + ``moe_config`` asking to defer, which ``should_defer_moe_finalize`` must + report truthfully. A deferred output reduced by ``finalize_moe_output`` must + give the kernel's own finalized output bit for bit, and modular experts hand + the router's weights back as-is. A call the experts split across kernel + launches finalizes instead. """ finalized = run() assert isinstance(finalized, torch.Tensor) moe_config.defer_moe_finalize = True + output = run() + assert moe_config.should_defer_moe_finalize(finalized.shape[0]) != chunked if chunked: - with pytest.raises(ValueError, match="one kernel launch"): - run() + torch.testing.assert_close(output, finalized, atol=0, rtol=0) return - deferred = run() - assert isinstance(deferred, UnfinalizedMoEOutput) + assert isinstance(output, UnfinalizedMoEOutput) if router_weights is not None: torch.testing.assert_close( - deferred.expert_weights, router_weights, atol=0, rtol=0 + output.expert_weights, router_weights, atol=0, rtol=0 ) - torch.testing.assert_close(finalize_moe_output(deferred), finalized, atol=0, rtol=0) + torch.testing.assert_close(finalize_moe_output(output), finalized, atol=0, rtol=0) def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8): diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 0c82d28096dc..3017513c4966 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1296,7 +1296,8 @@ class FusedMoEConfig: # consumer; read through `use_deferred_moe_finalize`, which applies the # guards. Kernels without the capability ignore it. Default False. defer_moe_finalize: bool = False - # Optional consumer capacity for deferred finalize. Negative means unbounded. + # Optional consumer capacity for deferred finalize, lowered by experts that + # would split a larger call across kernel launches. Negative means unbounded. defer_moe_finalize_max_num_tokens: int = -1 # SwiGLU clamp limit. When set, backends that do not implement the clamp @@ -1446,6 +1447,12 @@ def should_defer_moe_finalize(self, num_tokens: int) -> bool: and (max_num_tokens < 0 or num_tokens <= max_num_tokens) ) + def limit_deferred_moe_finalize(self, max_num_tokens: int) -> None: + """Finalize calls above ``max_num_tokens`` even when deferring.""" + current = self.defer_moe_finalize_max_num_tokens + if current < 0 or max_num_tokens < current: + self.defer_moe_finalize_max_num_tokens = max_num_tokens + @property def use_deepep_ht_kernels(self): return self.moe_parallel_config.use_deepep_ht_kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index b0ddfa2b430d..e1557a8f0d55 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -289,6 +289,18 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula Moved from trtllm_moe.py. """ + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + **kwargs, + ): + super().__init__(moe_config, quant_config, **kwargs) + # Each launch permutes into its own buffer, so a chunked call finalizes. + moe_config.limit_deferred_moe_finalize( + self._max_supported_tokens(self.topk, moe_config.num_experts) + ) + @staticmethod def _supports_parallel_config( moe_parallel_config: FusedMoEParallelConfig, @@ -439,16 +451,12 @@ def apply( assert self.w1_scale is not None assert self.w2_scale is not None - # Chunk tokens so the batched-GEMM grid stays within CUDA limits. + # Chunk tokens so the batched-GEMM grid stays within CUDA limits. Each + # launch permutes into its own buffer, so only a run that fits in one + # launch can leave the top-k reduction to a deferring consumer. M = x_quant.size(0) chunk_size = self._max_supported_tokens(topk, global_num_experts) - defer = self.moe_config.should_defer_moe_finalize(M) - if defer and chunk_size < M: - # Each launch permutes into its own buffer. - raise ValueError( - f"Deferring the MoE finalize needs one kernel launch, but {M} " - f"tokens take {-(-M // chunk_size)}." - ) + defer = chunk_size >= M and self.moe_config.should_defer_moe_finalize(M) unfinalized: UnfinalizedMoEOutput | None = None for start in range(0, M, chunk_size): end = min(start + chunk_size, M) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 10477515e343..f9cca18a9c21 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -302,6 +302,16 @@ def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int: class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular): """Modular version of the implementation (just the experts).""" + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + per_token_activation: bool = False, + ): + super().__init__(moe_config, quant_config, per_token_activation) + # Each launch permutes into its own buffer, so a chunked call finalizes. + moe_config.limit_deferred_moe_finalize(self._get_chunk_size()) + @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: """The modular implementation supports all parallel configs.""" @@ -434,9 +444,12 @@ def apply( M = hidden_states.shape[0] chunk_size = self._get_chunk_size() - defer = self.moe_config.should_defer_moe_finalize(M) if chunk_size >= M: + # Each launch permutes into its own buffer, so only a run that fits + # in one launch can leave the top-k reduction to a deferring + # consumer. + defer = self.moe_config.should_defer_moe_finalize(M) return self._invoke_kernel( None if defer else output, hidden_states, @@ -448,12 +461,6 @@ def apply( global_num_experts, a1q_scale, ) - elif defer: - # Each launch permutes into its own buffer. - raise ValueError( - f"Deferring the MoE finalize needs one kernel launch, but {M} " - f"tokens take {-(-M // chunk_size)}." - ) else: for start in range(0, M, chunk_size): end = min(start + chunk_size, M) From b5e6ccfb0e94f09a048e531bd2ab88648218d515 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 25 Sep 2026 01:06:24 +0000 Subject: [PATCH 6/8] [MoE] Request deferred MoE finalize through a FusedMoEConfig method Layers call moe_config.defer_moe_finalize(max_num_tokens) instead of assigning the flag and the token cap; the flag becomes private. Kimi-K3's runner moves to the method and limit_deferred_moe_finalize, with the same behavior. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yongye Zhu --- tests/kernels/moe/utils.py | 2 +- tests/models/kimi_k3/test_latent_moe_tail.py | 12 +++- .../model_executor/layers/fused_moe/config.py | 56 ++++++++++++++----- .../kimi_k3/nvidia/latent_moe_runner.py | 14 ++--- 4 files changed, 57 insertions(+), 27 deletions(-) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 63cafe20ea93..f2bf81efdbeb 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -729,7 +729,7 @@ def check_deferred_moe_finalize( """ finalized = run() assert isinstance(finalized, torch.Tensor) - moe_config.defer_moe_finalize = True + moe_config.defer_moe_finalize() output = run() assert moe_config.should_defer_moe_finalize(finalized.shape[0]) != chunked if chunked: diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py index 5fde9ff86d4a..c638bd7a49af 100644 --- a/tests/models/kimi_k3/test_latent_moe_tail.py +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -43,12 +43,18 @@ class FakeMoEConfig: hidden_dim = LATENT_SIZE hidden_dim_unpadded = LATENT_SIZE experts_per_token = 16 - defer_moe_finalize = False + deferred = False defer_moe_finalize_max_num_tokens = -1 + def defer_moe_finalize(self) -> None: + self.deferred = True + + def limit_deferred_moe_finalize(self, max_num_tokens: int) -> None: + self.defer_moe_finalize_max_num_tokens = max_num_tokens + @property def use_deferred_moe_finalize(self) -> bool: - return self.defer_moe_finalize + return self.deferred moe_config = FakeMoEConfig() quant_method = SimpleNamespace( @@ -99,7 +105,7 @@ def fake_tail_initialize(**kwargs): latent_moe_runner.LatentMoERunner() - assert moe_config.defer_moe_finalize + assert moe_config.use_deferred_moe_finalize assert moe_config.defer_moe_finalize_max_num_tokens == 128 assert initialized_with["experts_per_token"] == 16 diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 3017513c4966..de20d13a4a3d 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import IntEnum from typing import Union @@ -1290,14 +1290,12 @@ class FusedMoEConfig: # Only honored on the non-reduced (late-AR) TP path. Default False. skip_final_all_reduce: bool = False - # When True, experts that can stop after GEMM2 are allowed to hand back an - # UnfinalizedMoEOutput instead of finalized states, leaving the top-k - # reduction to fuse into the consumer. Set by layers that have such a - # consumer; read through `use_deferred_moe_finalize`, which applies the - # guards. Kernels without the capability ignore it. Default False. - defer_moe_finalize: bool = False - # Optional consumer capacity for deferred finalize, lowered by experts that - # would split a larger call across kernel launches. Negative means unbounded. + # Requested through `defer_moe_finalize()`, read through + # `should_defer_moe_finalize()`. + _defer_moe_finalize: bool = field(default=False, init=False) + # Most tokens a deferred call covers: the consumer's capacity, lowered by + # experts that would split a larger call across kernel launches. Negative + # means unbounded. defer_moe_finalize_max_num_tokens: int = -1 # SwiGLU clamp limit. When set, backends that do not implement the clamp @@ -1423,14 +1421,13 @@ def use_ep(self): def use_deferred_moe_finalize(self) -> bool: """Whether experts may return an unfinalized output on this deployment. - Evaluated on read rather than in ``__post_init__`` because - ``defer_moe_finalize`` is set after construction, like - ``skip_final_all_reduce``. + Evaluated on read rather than in ``__post_init__`` because deferral is + requested after construction, like ``skip_final_all_reduce``. """ # The consumer fuses a TP all-reduce. Other parallel modes require a # combine or reduce-scatter after the experts and cannot defer it. return ( - self.defer_moe_finalize + self._defer_moe_finalize and self.tp_size > 1 and self.dp_size == 1 and self.ep_size == 1 @@ -1438,6 +1435,39 @@ def use_deferred_moe_finalize(self) -> bool: and not self.is_sequence_parallel ) + def defer_moe_finalize(self, max_num_tokens: int = -1) -> None: + """Ask the experts to leave the top-k reduction to the layer's consumer. + + A layer whose consumer can fuse the top-k reduction (e.g. into its TP + all-reduce) calls this once while the model is built. From then on: + + - Experts that can stop after GEMM2 (the TRTLLM-Gen ``do_finalize=False`` + path) return an ``UnfinalizedMoEOutput`` instead of finalized states + for every call ``should_defer_moe_finalize`` accepts. Experts without + the capability ignore the request and always finalize. + - ``should_defer_moe_finalize(num_tokens)`` is the answer for a call: + it requires a TP-only deployment (``use_deferred_moe_finalize``), a + non-empty call and at most ``defer_moe_finalize_max_num_tokens`` + tokens. The cap starts at ``max_num_tokens`` and only ever goes down: + experts that would split a larger call across kernel launches lower it + to their single-launch size when they are built, since each launch + permutes into its own buffer. + - The model asks ``should_defer_moe_finalize`` before each call and takes + the matching path. A deferred call runs the runner's ``_forward_impl`` + directly, since the MoE custom op returns tensors only, and the + consumer then owns the top-k reduction, the shared-expert add and the + all-reduce. ``finalize_moe_output`` runs the reduction on its own for + deferred outputs the consumer's fused kernel can't take. + + Args: + max_num_tokens: Most tokens per call the consumer can take in + deferred form. Negative means no limit of its own. + + """ + self._defer_moe_finalize = True + if max_num_tokens >= 0: + self.limit_deferred_moe_finalize(max_num_tokens) + def should_defer_moe_finalize(self, num_tokens: int) -> bool: """Return whether this invocation may defer the top-k reduction.""" max_num_tokens = self.defer_moe_finalize_max_num_tokens diff --git a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py index 8eedb5ba89ed..8678dcbfe07d 100644 --- a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py +++ b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py @@ -115,10 +115,11 @@ def __init__( experts_cls = getattr(self._quant_method, "experts_cls", None) # The kernel instance is built after weight loading, while the tail # must register its CuTeDSL warmup units during runner construction. - self.moe_config.defer_moe_finalize = ( + if ( experts_cls is TrtLlmMxfp4ExpertsMonolithic or experts_cls is TrtLlmNvFp4ExpertsMonolithic - ) and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded + ) and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded: + self.moe_config.defer_moe_finalize() from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import ( KimiK3LatentMoETailOp, ) @@ -136,20 +137,13 @@ def __init__( rms_eps=norm.variance_epsilon, ) self._k3_latent_moe_tail_op = op - self.moe_config.defer_moe_finalize_max_num_tokens = ( - op.contract.max_num_tokens - if self.moe_config.use_deferred_moe_finalize - else -1 - ) if self.moe_config.use_deferred_moe_finalize: + self.moe_config.limit_deferred_moe_finalize(op.contract.max_num_tokens) logger.info_once( "K3 latent-MoE tail fusion with deferred top-k finalization " "is enabled for up to %d tokens.", op.contract.max_num_tokens, ) - else: - self.moe_config.defer_moe_finalize = False - self.moe_config.defer_moe_finalize_max_num_tokens = -1 def _get_zero_residual( self, From e75ed4b55d2cd44fc6d036f9c6afc70dd7765552 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 25 Sep 2026 05:05:22 +0000 Subject: [PATCH 7/8] [MoE] Address review: padding, cap limits and the standalone finalize - Deferral requires hidden_dim == hidden_dim_unpadded: the consumer can't strip hidden-dim padding from GEMM2 rows, and finalize_moe_output would return the padded width. - limit_deferred_moe_finalize ignores a negative limit instead of lifting the cap, which could undo the experts' single-launch limit. - finalize_moe_output hands moe_unpermute a row count, so the kernel's num_valid read never goes through a null pointer. - defer_moe_finalize's docstring asks the layer to check for TRTLLM-Gen experts first, as Kimi-K3 does: other experts ignore the request. - Tests: compare the router weights with a copy taken before the run, run K3's fake config on the real FusedMoEConfig methods, and cover the MXFP8 FP8 experts. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yongye Zhu --- tests/kernels/moe/test_flashinfer.py | 50 +++++++++++++++---- tests/kernels/moe/utils.py | 6 ++- tests/models/kimi_k3/test_latent_moe_tail.py | 16 ++---- .../model_executor/layers/fused_moe/config.py | 42 +++++++++------- .../layers/fused_moe/moe_output.py | 3 ++ 5 files changed, 78 insertions(+), 39 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index fbc45d41a9bd..5d55876e8138 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -40,6 +40,7 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( Fp8MoeBackend, convert_to_fp8_moe_kernel_format, + make_fp8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( @@ -47,6 +48,9 @@ swap_w13_to_w31, ) from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, @@ -779,21 +783,44 @@ def _supports_quant_scheme( assert "SwiGLU" in reason +def _make_mxfp8_moe_weights( + e: int, n: int, k: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + q, s = zip( + *(mxfp8_e4m3_quantize(w[i], is_sf_swizzled_layout=False) for i in range(e)) + ) + return torch.stack(q), torch.stack(s) + + w1, w1_scale = quantize( + torch.randn(e, 2 * n, k, device="cuda", dtype=torch.bfloat16) / 10 + ) + w2, w2_scale = quantize( + torch.randn(e, k, n, device="cuda", dtype=torch.bfloat16) / 10 + ) + return w1, w2, w1_scale, w2_scale + + +@pytest.mark.parametrize("block_shape", [[128, 128], [1, 32]], ids=["block", "mxfp8"]) @pytest.mark.parametrize("m", [1, 16]) @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="Requires TRTLLM-Gen FP8 MoE (SM100)", ) -def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): - """TRTLLM-Gen block-FP8 modular experts can leave the top-k finalize to the - caller.""" +def test_trtllm_fp8_block_moe_deferred_finalize( + m: int, block_shape: list[int], workspace_init +): + """TRTLLM-Gen block-FP8 and MXFP8 modular experts can leave the top-k + finalize to the caller.""" e, topk, n, k = 32, 4, 1024, 1024 - block_shape = [128, 128] set_random_seed(7) with set_current_vllm_config(vllm_config): - (_, w1, w1_scale, _), (_, w2, w2_scale, _) = make_test_weights( - e, n, k, quant_dtype=torch.float8_e4m3fn, block_shape=block_shape - ) + if block_shape == [1, 32]: + w1, w2, w1_scale, w2_scale = _make_mxfp8_moe_weights(e, n, k) + else: + (_, w1, w1_scale, _), (_, w2, w2_scale, _) = make_test_weights( + e, n, k, quant_dtype=torch.float8_e4m3fn, block_shape=block_shape + ) w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format( Fp8MoeBackend.FLASHINFER_TRTLLM, SimpleNamespace( @@ -810,8 +837,13 @@ def test_trtllm_fp8_block_moe_deferred_finalize(m: int, workspace_init): w13_input_scale=None, w2_input_scale=None, ) - quant_config = fp8_w8a8_moe_quant_config( - w1_scale=w1_scale, w2_scale=w2_scale, block_shape=block_shape + quant_config = make_fp8_moe_quant_config( + Fp8MoeBackend.FLASHINFER_TRTLLM, + w1_scale, + w2_scale, + a1_scale=None, + a2_scale=None, + block_shape=block_shape, ) # One rank of a TP group, which deferral needs. moe_config = FusedMoEConfig( diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index f2bf81efdbeb..9fcfa974c081 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -727,6 +727,8 @@ def check_deferred_moe_finalize( the router's weights back as-is. A call the experts split across kernel launches finalizes instead. """ + # The deferred output views the router's buffer, so compare with a copy. + expected_weights = None if router_weights is None else router_weights.clone() finalized = run() assert isinstance(finalized, torch.Tensor) moe_config.defer_moe_finalize() @@ -737,9 +739,9 @@ def check_deferred_moe_finalize( return assert isinstance(output, UnfinalizedMoEOutput) - if router_weights is not None: + if expected_weights is not None: torch.testing.assert_close( - output.expert_weights, router_weights, atol=0, rtol=0 + output.expert_weights, expected_weights, atol=0, rtol=0 ) torch.testing.assert_close(finalize_moe_output(output), finalized, atol=0, rtol=0) diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py index c638bd7a49af..96876e736c3e 100644 --- a/tests/models/kimi_k3/test_latent_moe_tail.py +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -15,6 +15,7 @@ multi_process_parallel, ) from vllm.distributed import get_tp_group +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( TrtLlmMxfp4ExpertsMonolithic, ) @@ -43,18 +44,11 @@ class FakeMoEConfig: hidden_dim = LATENT_SIZE hidden_dim_unpadded = LATENT_SIZE experts_per_token = 16 - deferred = False + _defer_moe_finalize = False defer_moe_finalize_max_num_tokens = -1 - - def defer_moe_finalize(self) -> None: - self.deferred = True - - def limit_deferred_moe_finalize(self, max_num_tokens: int) -> None: - self.defer_moe_finalize_max_num_tokens = max_num_tokens - - @property - def use_deferred_moe_finalize(self) -> bool: - return self.deferred + defer_moe_finalize = FusedMoEConfig.defer_moe_finalize + limit_deferred_moe_finalize = FusedMoEConfig.limit_deferred_moe_finalize + use_deferred_moe_finalize = FusedMoEConfig.use_deferred_moe_finalize moe_config = FakeMoEConfig() quant_method = SimpleNamespace( diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index de20d13a4a3d..8b3e084653bf 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1425,7 +1425,8 @@ def use_deferred_moe_finalize(self) -> bool: requested after construction, like ``skip_final_all_reduce``. """ # The consumer fuses a TP all-reduce. Other parallel modes require a - # combine or reduce-scatter after the experts and cannot defer it. + # combine or reduce-scatter after the experts and cannot defer it, and + # the consumer has no way to strip hidden-dim padding from GEMM2 rows. return ( self._defer_moe_finalize and self.tp_size > 1 @@ -1433,25 +1434,30 @@ def use_deferred_moe_finalize(self) -> bool: and self.ep_size == 1 and self.pcp_size == 1 and not self.is_sequence_parallel + and self.hidden_dim == self.hidden_dim_unpadded ) def defer_moe_finalize(self, max_num_tokens: int = -1) -> None: """Ask the experts to leave the top-k reduction to the layer's consumer. A layer whose consumer can fuse the top-k reduction (e.g. into its TP - all-reduce) calls this once while the model is built. From then on: - - - Experts that can stop after GEMM2 (the TRTLLM-Gen ``do_finalize=False`` - path) return an ``UnfinalizedMoEOutput`` instead of finalized states - for every call ``should_defer_moe_finalize`` accepts. Experts without - the capability ignore the request and always finalize. + all-reduce) calls this once while the model is built, and only when its + experts are TRTLLM-Gen ones that can stop after GEMM2 (the + ``do_finalize=False`` path). Other experts ignore the request and + always finalize, which ``should_defer_moe_finalize`` can't see, so the + layer checks the quant method's ``experts_cls`` first, as Kimi-K3 does. + From then on: + + - The experts return an ``UnfinalizedMoEOutput`` instead of finalized + states for every call ``should_defer_moe_finalize`` accepts. - ``should_defer_moe_finalize(num_tokens)`` is the answer for a call: - it requires a TP-only deployment (``use_deferred_moe_finalize``), a - non-empty call and at most ``defer_moe_finalize_max_num_tokens`` - tokens. The cap starts at ``max_num_tokens`` and only ever goes down: - experts that would split a larger call across kernel launches lower it - to their single-launch size when they are built, since each launch - permutes into its own buffer. + it requires a TP-only deployment without hidden-dim padding + (``use_deferred_moe_finalize``), a non-empty call and at most + ``defer_moe_finalize_max_num_tokens`` tokens. The cap starts at + ``max_num_tokens`` and only ever goes down: experts that would split a + larger call across kernel launches lower it to their single-launch + size when they are built, since each launch permutes into its own + buffer. - The model asks ``should_defer_moe_finalize`` before each call and takes the matching path. A deferred call runs the runner's ``_forward_impl`` directly, since the MoE custom op returns tensors only, and the @@ -1465,8 +1471,7 @@ def defer_moe_finalize(self, max_num_tokens: int = -1) -> None: """ self._defer_moe_finalize = True - if max_num_tokens >= 0: - self.limit_deferred_moe_finalize(max_num_tokens) + self.limit_deferred_moe_finalize(max_num_tokens) def should_defer_moe_finalize(self, num_tokens: int) -> bool: """Return whether this invocation may defer the top-k reduction.""" @@ -1478,9 +1483,12 @@ def should_defer_moe_finalize(self, num_tokens: int) -> bool: ) def limit_deferred_moe_finalize(self, max_num_tokens: int) -> None: - """Finalize calls above ``max_num_tokens`` even when deferring.""" + """Finalize calls above ``max_num_tokens`` even when deferring. + + Negative means no limit, and leaves the current one in place. + """ current = self.defer_moe_finalize_max_num_tokens - if current < 0 or max_num_tokens < current: + if max_num_tokens >= 0 and (current < 0 or max_num_tokens < current): self.defer_moe_finalize_max_num_tokens = max_num_tokens @property diff --git a/vllm/model_executor/layers/fused_moe/moe_output.py b/vllm/model_executor/layers/fused_moe/moe_output.py index 814fea309fbb..0f7fab400c8d 100644 --- a/vllm/model_executor/layers/fused_moe/moe_output.py +++ b/vllm/model_executor/layers/fused_moe/moe_output.py @@ -139,11 +139,14 @@ def finalize_moe_output(routed: UnfinalizedMoEOutput) -> torch.Tensor: gemm2_permuted = routed.gemm2_permuted num_tokens = routed.expanded_idx_to_permuted_idx.shape[0] output = gemm2_permuted.new_empty(num_tokens, gemm2_permuted.shape[1]) + # The kernel reads the valid row count through this pointer either way. + num_rows = gemm2_permuted.new_full((1,), gemm2_permuted.shape[0], dtype=torch.int64) moe_unpermute( output, gemm2_permuted, routed.expert_weights.float(), routed.expanded_idx_to_permuted_idx, + expert_first_token_offset=num_rows, ) return output From 6f60836565f6b0a43c10107393ad7474625bc43c Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 25 Sep 2026 05:15:02 +0000 Subject: [PATCH 8/8] [MoE] Drop finalize_moe_output With the consumer's capacity passed as the deferral cap, calls it can't take finalize inside the experts, so no production path reduces a deferred output on its own. The deferral test calls moe_unpermute directly as its bit-exact reference. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yongye Zhu --- tests/kernels/moe/utils.py | 27 +++++++++----- .../model_executor/layers/fused_moe/config.py | 3 +- .../layers/fused_moe/moe_output.py | 35 +++---------------- 3 files changed, 23 insertions(+), 42 deletions(-) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 9fcfa974c081..9274a08af4a3 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -29,10 +29,8 @@ fused_experts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel -from vllm.model_executor.layers.fused_moe.moe_output import ( - UnfinalizedMoEOutput, - finalize_moe_output, -) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import moe_unpermute from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import ( BatchedPrepareAndFinalize, ) @@ -722,10 +720,10 @@ def check_deferred_moe_finalize( ``run`` calls the kernel on fixed inputs, first as built and then with ``moe_config`` asking to defer, which ``should_defer_moe_finalize`` must - report truthfully. A deferred output reduced by ``finalize_moe_output`` must - give the kernel's own finalized output bit for bit, and modular experts hand - the router's weights back as-is. A call the experts split across kernel - launches finalizes instead. + report truthfully. A deferred output reduced by ``moe_unpermute``, the + TRT-LLM finalize kernel, must give the kernel's own finalized output bit for + bit, and modular experts hand the router's weights back as-is. A call the + experts split across kernel launches finalizes instead. """ # The deferred output views the router's buffer, so compare with a copy. expected_weights = None if router_weights is None else router_weights.clone() @@ -743,7 +741,18 @@ def check_deferred_moe_finalize( torch.testing.assert_close( output.expert_weights, expected_weights, atol=0, rtol=0 ) - torch.testing.assert_close(finalize_moe_output(output), finalized, atol=0, rtol=0) + reference = torch.empty_like(finalized) + moe_unpermute( + reference, + output.gemm2_permuted, + output.expert_weights.float(), + output.expanded_idx_to_permuted_idx, + # The kernel reads its valid-row count through this pointer either way. + expert_first_token_offset=output.gemm2_permuted.new_full( + (1,), output.gemm2_permuted.shape[0], dtype=torch.int64 + ), + ) + torch.testing.assert_close(reference, finalized, atol=0, rtol=0) def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8): diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 8b3e084653bf..9221e9382428 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1462,8 +1462,7 @@ def defer_moe_finalize(self, max_num_tokens: int = -1) -> None: the matching path. A deferred call runs the runner's ``_forward_impl`` directly, since the MoE custom op returns tensors only, and the consumer then owns the top-k reduction, the shared-expert add and the - all-reduce. ``finalize_moe_output`` runs the reduction on its own for - deferred outputs the consumer's fused kernel can't take. + all-reduce. Args: max_num_tokens: Most tokens per call the consumer can take in diff --git a/vllm/model_executor/layers/fused_moe/moe_output.py b/vllm/model_executor/layers/fused_moe/moe_output.py index 0f7fab400c8d..8f9662b442ce 100644 --- a/vllm/model_executor/layers/fused_moe/moe_output.py +++ b/vllm/model_executor/layers/fused_moe/moe_output.py @@ -124,33 +124,6 @@ def convert_flashinfer_moe_output( ) -def finalize_moe_output(routed: UnfinalizedMoEOutput) -> torch.Tensor: - """Run the top-k reduction a deferring MoE layer left open. - - For the calls a fused consumer can't take. This is the TRT-LLM finalize as - a standalone kernel, so it reduces the same way the experts would have. - A deferring layer has every expert local (``ep_size == 1``), so the permute - map has no -1 entries. - """ - from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( - moe_unpermute, - ) - - gemm2_permuted = routed.gemm2_permuted - num_tokens = routed.expanded_idx_to_permuted_idx.shape[0] - output = gemm2_permuted.new_empty(num_tokens, gemm2_permuted.shape[1]) - # The kernel reads the valid row count through this pointer either way. - num_rows = gemm2_permuted.new_full((1,), gemm2_permuted.shape[0], dtype=torch.int64) - moe_unpermute( - output, - gemm2_permuted, - routed.expert_weights.float(), - routed.expanded_idx_to_permuted_idx, - expert_first_token_offset=num_rows, - ) - return output - - @dataclass class MoEOutput: """A MoE layer's output with its final reduction still open. @@ -163,10 +136,10 @@ class MoEOutput: unfinalized, the top-k reduction is open too and can fold into the same kernel. - Producers only leave the routed output unfinalized when the layer asked for - it -- the token ceiling and topology support are its to check. Without a - ceiling that is every call, and ``finalize_moe_output`` covers the calls a - fused consumer can't take. + Producers only leave the routed output unfinalized when a fused consumer can + actually take that form -- the token ceiling and topology support are theirs + to check -- so an ``UnfinalizedMoEOutput`` here means the fused path applies, + and a consumer need not re-derive that. """ # Un-reduced routed output, either finalized or not.