@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 378504a544

[MoE] Defer the TRTLLM-Gen top-k finalize on the modular path (#58635)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Yongye Zhu committed Sep 25, 2026 at 00:38 UTC 378504a5442b8a9240b359dae3e6f75f35c38f19
26 files changed +652 -117
tests/kernels/moe/test_flashinfer.py
+122 -2
@@ -1,14 +1,19 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import multiprocessing
4 -from dataclasses import dataclass
4 +from dataclasses import dataclass, replace
5 +from types import SimpleNamespace
6 from unittest.mock import MagicMock
7
8 import pytest
9 import torch
10
11 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 -from tests.kernels.moe.utils import make_dummy_moe_config
12 +from tests.kernels.moe.utils import (
13 + check_deferred_moe_finalize,
14 + make_dummy_moe_config,
15 + make_test_weights,
16 +)
17 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
18 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19 from vllm.model_executor.layers.fused_moe.all2all_utils import (
@@ -32,11 +37,20 @@ from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
37 TrtLlmFp8ExpertsMonolithic,
38 )
39 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
40 +from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
41 + Fp8MoeBackend,
42 + convert_to_fp8_moe_kernel_format,
43 + make_fp8_moe_quant_config,
44 +)
45 +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
46 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
47 rotate_weights_for_fi_trtllm_fp8_per_tensor_moe,
48 swap_w13_to_w31,
49 )
50 from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
51 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
52 + mxfp8_e4m3_quantize,
53 +)
54 from vllm.model_executor.layers.quantization.utils.quant_utils import (
55 QuantKey,
56 kFp8Dynamic128Sym,
@@ -767,3 +781,109 @@ def test_trtllm_fp8_swiglu_clamp_support(
781 assert supported == expected, reason
782 if not expected:
783 assert "SwiGLU" in reason
784 +
785 +
786 +def _make_mxfp8_moe_weights(
787 + e: int, n: int, k: int
788 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
789 + def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
790 + q, s = zip(
791 + *(mxfp8_e4m3_quantize(w[i], is_sf_swizzled_layout=False) for i in range(e))
792 + )
793 + return torch.stack(q), torch.stack(s)
794 +
795 + w1, w1_scale = quantize(
796 + torch.randn(e, 2 * n, k, device="cuda", dtype=torch.bfloat16) / 10
797 + )
798 + w2, w2_scale = quantize(
799 + torch.randn(e, k, n, device="cuda", dtype=torch.bfloat16) / 10
800 + )
801 + return w1, w2, w1_scale, w2_scale
802 +
803 +
804 +@pytest.mark.parametrize("block_shape", [[128, 128], [1, 32]], ids=["block", "mxfp8"])
805 +@pytest.mark.parametrize("m", [1, 16])
806 +@pytest.mark.skipif(
807 + not current_platform.is_device_capability_family(100),
808 + reason="Requires TRTLLM-Gen FP8 MoE (SM100)",
809 +)
810 +def test_trtllm_fp8_block_moe_deferred_finalize(
811 + m: int, block_shape: list[int], workspace_init
812 +):
813 + """TRTLLM-Gen block-FP8 and MXFP8 modular experts can leave the top-k
814 + finalize to the caller."""
815 + e, topk, n, k = 32, 4, 1024, 1024
816 + set_random_seed(7)
817 + with set_current_vllm_config(vllm_config):
818 + if block_shape == [1, 32]:
819 + w1, w2, w1_scale, w2_scale = _make_mxfp8_moe_weights(e, n, k)
820 + else:
821 + (_, w1, w1_scale, _), (_, w2, w2_scale, _) = make_test_weights(
822 + e, n, k, quant_dtype=torch.float8_e4m3fn, block_shape=block_shape
823 + )
824 + w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format(
825 + Fp8MoeBackend.FLASHINFER_TRTLLM,
826 + SimpleNamespace(
827 + weight_block_size=block_shape,
828 + moe_config=SimpleNamespace(
829 + is_act_and_mul=True, intermediate_size_per_partition=n
830 + ),
831 + activation=MoEActivation.SILU,
832 + ),
833 + w1,
834 + w2,
835 + w1_scale,
836 + w2_scale,
837 + w13_input_scale=None,
838 + w2_input_scale=None,
839 + )
840 + quant_config = make_fp8_moe_quant_config(
841 + Fp8MoeBackend.FLASHINFER_TRTLLM,
842 + w1_scale,
843 + w2_scale,
844 + a1_scale=None,
845 + a2_scale=None,
846 + block_shape=block_shape,
847 + )
848 + # One rank of a TP group, which deferral needs.
849 + moe_config = FusedMoEConfig(
850 + num_experts=e,
851 + experts_per_token=topk,
852 + hidden_dim=k,
853 + intermediate_size=2 * n,
854 + num_local_experts=e,
855 + num_logical_experts=e,
856 + activation=MoEActivation.SILU,
857 + device="cuda",
858 + moe_parallel_config=replace(
859 + FusedMoEParallelConfig.make_no_parallel(), tp_size=2
860 + ),
861 + in_dtype=torch.bfloat16,
862 + routing_method=RoutingMethodType.TopK,
863 + max_num_tokens=next_power_of_2(m),
864 + )
865 + kernel = mk.FusedMoEKernel(
866 + maybe_make_prepare_finalize(
867 + moe=moe_config, quant_config=quant_config, allow_new_interface=True
868 + ),
869 + TrtLlmFp8ExpertsModular(moe_config=moe_config, quant_config=quant_config),
870 + )
871 +
872 + a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
873 + score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
874 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
875 + check_deferred_moe_finalize(
876 + moe_config,
877 + lambda: kernel.apply(
878 + hidden_states=a,
879 + w1=w1,
880 + w2=w2,
881 + topk_weights=topk_weights,
882 + topk_ids=topk_ids,
883 + activation=MoEActivation.SILU,
884 + global_num_experts=e,
885 + expert_map=None,
886 + apply_router_weight_on_input=False,
887 + ),
888 + router_weights=topk_weights,
889 + )
tests/kernels/moe/test_ocp_mx_moe.py
+149 -1
@@ -7,7 +7,7 @@ from dataclasses import dataclass, replace
7 import pytest
8 import torch
9
10 -from tests.kernels.moe.utils import check_accuracy
10 +from tests.kernels.moe.utils import check_accuracy, check_deferred_moe_finalize
11 from vllm._aiter_ops import (
12 is_aiter_found,
13 is_aiter_found_and_supported,
@@ -663,6 +663,154 @@ def test_trtllm_gen_mxfp4_fused_moe(
663 check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8)
664
665
666 +@pytest.mark.parametrize(
667 + "monolithic,num_tokens,chunk_size",
668 + [(False, 1, None), (False, 16, None), (False, 16, 5), (True, 16, None)],
669 +)
670 +@pytest.mark.skipif(
671 + not TRTLLM_GEN_MXFP4_AVAILABLE,
672 + reason="nvidia gpu and compute capability sm100 is required for this test",
673 +)
674 +def test_trtllm_mxfp4_deferred_finalize(
675 + monolithic: bool,
676 + num_tokens: int,
677 + chunk_size: int | None,
678 + monkeypatch: pytest.MonkeyPatch,
679 + workspace_init,
680 +):
681 + """TRTLLM-Gen MXFP4 experts can leave the top-k finalize to the caller."""
682 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
683 + from vllm.model_executor.layers.fused_moe.config import (
684 + FusedMoEConfig,
685 + FusedMoEParallelConfig,
686 + RoutingMethodType,
687 + )
688 + from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import (
689 + TrtLlmMxfp4ExpertsModular,
690 + TrtLlmMxfp4ExpertsMonolithic,
691 + )
692 + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
693 + Mxfp4MoeBackend,
694 + convert_weight_to_mxfp4_moe_kernel_format,
695 + make_mxfp4_moe_kernel,
696 + make_mxfp4_moe_quant_config,
697 + )
698 +
699 + num_experts, topk, hidden_size, intermediate_size = 32, 4, 1024, 512
700 + device = torch.device("cuda:0")
701 + torch.manual_seed(0)
702 + if chunk_size is not None:
703 + monkeypatch.setattr(
704 + TrtLlmMxfp4ExpertsModular,
705 + "_max_supported_tokens",
706 + lambda self, top_k, global_num_experts: chunk_size,
707 + )
708 +
709 + def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
710 + w_q, w_scale = fp4_quantize(
711 + w,
712 + torch.tensor(1.0, device=device),
713 + 32,
714 + sf_use_ue8m0=True,
715 + is_sf_swizzled_layout=False,
716 + )
717 + return (
718 + w_q.reshape(*w.shape[:-1], -1),
719 + w_scale.view(torch.uint8).reshape(*w.shape[:-1], -1),
720 + )
721 +
722 + w13, w13_scale = quantize(
723 + torch.randn(
724 + num_experts,
725 + 2 * intermediate_size,
726 + hidden_size,
727 + device=device,
728 + dtype=torch.bfloat16,
729 + )
730 + )
731 + w2, w2_scale = quantize(
732 + torch.randn(
733 + num_experts,
734 + hidden_size,
735 + intermediate_size,
736 + device=device,
737 + dtype=torch.bfloat16,
738 + )
739 + )
740 + backend = Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8
741 + w13, w2, w13_scale, w2_scale, _, _ = convert_weight_to_mxfp4_moe_kernel_format(
742 + backend,
743 + torch.nn.Module(),
744 + w13,
745 + w2,
746 + w13_scale,
747 + w2_scale,
748 + _cache_permute_indices={},
749 + )
750 + # One rank of a TP group, which deferral needs.
751 + moe_config = FusedMoEConfig(
752 + num_experts=num_experts,
753 + experts_per_token=topk,
754 + hidden_dim=hidden_size,
755 + intermediate_size=2 * intermediate_size,
756 + num_local_experts=num_experts,
757 + num_logical_experts=num_experts,
758 + activation=MoEActivation.SILU,
759 + device=device,
760 + routing_method=RoutingMethodType.Renormalize,
761 + moe_parallel_config=replace(
762 + FusedMoEParallelConfig.make_no_parallel(), tp_size=2
763 + ),
764 + in_dtype=torch.bfloat16,
765 + )
766 + kernel = make_mxfp4_moe_kernel(
767 + make_mxfp4_moe_quant_config(backend, w1_scale=w13_scale, w2_scale=w2_scale),
768 + moe_config,
769 + TrtLlmMxfp4ExpertsMonolithic if monolithic else TrtLlmMxfp4ExpertsModular,
770 + backend,
771 + )
772 +
773 + hidden_states = torch.randn(
774 + num_tokens, hidden_size, device=device, dtype=torch.bfloat16
775 + )
776 + if monolithic:
777 + router_logits = torch.randn(num_tokens, num_experts, device=device)
778 + check_deferred_moe_finalize(
779 + moe_config,
780 + lambda: kernel.apply_monolithic(
781 + hidden_states,
782 + w13,
783 + w2,
784 + router_logits,
785 + activation=MoEActivation.SILU,
786 + global_num_experts=num_experts,
787 + expert_map=None,
788 + apply_router_weight_on_input=False,
789 + ),
790 + )
791 + return
792 +
793 + topk_ids = torch.rand(num_tokens, num_experts, device=device).argsort(dim=-1)
794 + topk_ids = topk_ids[:, :topk].to(torch.int32)
795 + topk_weights = torch.rand(num_tokens, topk, device=device).softmax(dim=-1)
796 + check_deferred_moe_finalize(
797 + moe_config,
798 + lambda: kernel.apply(
799 + hidden_states=hidden_states,
800 + w1=w13,
801 + w2=w2,
802 + topk_weights=topk_weights,
803 + topk_ids=topk_ids,
804 + activation=MoEActivation.SILU,
805 + global_num_experts=num_experts,
806 + expert_map=None,
807 + apply_router_weight_on_input=False,
808 + ),
809 + router_weights=topk_weights,
810 + chunked=chunk_size is not None,
811 + )
812 +
813 +
814 @pytest.mark.parametrize("topk", [1, 4])
815 @pytest.mark.parametrize("num_experts", [32])
816 @pytest.mark.parametrize("num_tokens", [1, 128])
tests/kernels/moe/test_trtllm_nvfp4_moe.py
+102 -28
@@ -8,11 +8,13 @@ and GELU — including a Gemma4-shaped case (128 experts, top-k 8,
8 intermediate_size 704) that exercises the non-256-aligned padding path.
9 """
10
11 +from dataclasses import replace
12 +
13 import pytest
14 import torch
15
16 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
15 -from tests.kernels.moe.utils import make_test_quant_config
17 +from tests.kernels.moe.utils import check_deferred_moe_finalize, make_test_quant_config
18 from tests.kernels.quantization.nvfp4_utils import (
19 FLOAT4_E2M1_MAX,
20 FLOAT8_E4M3_MAX,
@@ -31,6 +33,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
33 from vllm.model_executor.layers.fused_moe.config import (
34 FusedMoEConfig,
35 FusedMoEParallelConfig,
36 + FusedMoEQuantConfig,
37 RoutingMethodType,
38 )
39 from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
@@ -92,6 +95,36 @@ SILU_WITH_CLAMP = op_registry[_CLAMP_OP_NAME]
95 SITU = op_registry[_SITU_OP_NAME]
96
97
98 +def _make_trtllm_fp4_moe_kernel(
99 + moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig
100 +) -> mk.FusedMoEKernel:
101 + experts = TrtLlmNvFp4ExpertsModular(
102 + moe_config=moe_config, quant_config=quant_config
103 + )
104 + # Mimic the production weight-loader path so per-expert tensors that
105 + # are normally precomputed in process_weights_after_loading (g1_scale_c
106 + # and the rescaled gemm1_clamp_limit) get materialized. The test's
107 + # synthetic quant_config has g1_alphas/g2_alphas already at their
108 + # post-fusion values, so we set w13_weight_scale_2 to alias g1_alphas
109 + # (same tensor) and use input_scale=1 to make the in-place
110 + # weight_scale_2 *= input_scale step a no-op.
111 + fake_layer = torch.nn.Module()
112 + fake_layer.w13_weight_scale_2 = quant_config.g1_alphas
113 + fake_layer.w2_weight_scale_2 = quant_config.g2_alphas
114 + fake_layer.w13_input_scale = torch.ones_like(quant_config.g1_alphas)
115 + fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas)
116 + experts.process_weights_after_loading(fake_layer)
117 + return mk.FusedMoEKernel(
118 + maybe_make_prepare_finalize(
119 + moe=moe_config,
120 + quant_config=quant_config,
121 + allow_new_interface=True,
122 + use_monolithic=False,
123 + ),
124 + experts,
125 + )
126 +
127 +
128 ACTIVATION_CASES = [
129 pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"),
130 pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"),
@@ -191,34 +224,11 @@ def test_trtllm_fp4_moe_no_graph(
224 ),
225 )
226
194 - trtllm_inner = TrtLlmNvFp4ExpertsModular(
195 - moe_config=moe_config, quant_config=quant_config
196 - )
197 - # Mimic the production weight-loader path so per-expert tensors that
198 - # are normally precomputed in process_weights_after_loading (g1_scale_c
199 - # and the rescaled gemm1_clamp_limit) get materialized. The test's
200 - # synthetic quant_config has g1_alphas/g2_alphas already at their
201 - # post-fusion values, so we set w13_weight_scale_2 to alias g1_alphas
202 - # (same tensor) and use input_scale=1 to make the in-place
203 - # weight_scale_2 *= input_scale step a no-op.
204 - fake_layer = torch.nn.Module()
205 - fake_layer.w13_weight_scale_2 = quant_config.g1_alphas
206 - fake_layer.w2_weight_scale_2 = quant_config.g2_alphas
207 - fake_layer.w13_input_scale = torch.ones_like(quant_config.g1_alphas)
208 - fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas)
209 - trtllm_inner.process_weights_after_loading(fake_layer)
227 + trtllm_experts = _make_trtllm_fp4_moe_kernel(moe_config, quant_config)
228 if activation == MoEActivation.SITU:
211 - torch.testing.assert_close(trtllm_inner.g1_scale_c, quant_config.a2_gscale)
212 -
213 - trtllm_experts = mk.FusedMoEKernel(
214 - maybe_make_prepare_finalize(
215 - moe=moe_config,
216 - quant_config=quant_config,
217 - allow_new_interface=True,
218 - use_monolithic=False,
219 - ),
220 - trtllm_inner,
221 - )
229 + torch.testing.assert_close(
230 + trtllm_experts.fused_experts.g1_scale_c, quant_config.a2_gscale
231 + )
232
233 trtllm_output = trtllm_experts.apply(
234 hidden_states=a,
@@ -331,6 +341,70 @@ def test_trtllm_fp4_moe_reprocess_does_not_refold_swiglu_params():
341 torch.testing.assert_close(experts.gemm1_beta, beta * expected)
342
343
344 +@pytest.mark.parametrize("m,chunk_size", [(1, None), (16, None), (16, 5)])
345 +@torch.inference_mode()
346 +def test_trtllm_fp4_moe_deferred_finalize(
347 + m: int,
348 + chunk_size: int | None,
349 + monkeypatch: pytest.MonkeyPatch,
350 + workspace_init,
351 +):
352 + """TRTLLM-Gen NvFP4 modular experts can leave the top-k finalize to the caller."""
353 + e, topk, n, k = 128, 8, 1024, 1024
354 + dtype = torch.bfloat16
355 + if chunk_size is not None:
356 + monkeypatch.setattr(
357 + TrtLlmNvFp4ExpertsModular, "_get_chunk_size", lambda self: chunk_size
358 + )
359 +
360 + set_random_seed(7)
361 + with set_current_vllm_config(
362 + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
363 + ):
364 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
365 + w1_q, w2_q, quant_config = make_test_quant_config(
366 + e, n, k, in_dtype=dtype, quant_dtype="nvfp4", is_scale_swizzled=False
367 + )
368 + score = torch.randn((m, e), device="cuda", dtype=dtype)
369 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
370 +
371 + # One rank of a TP group, which deferral needs.
372 + moe_config = FusedMoEConfig(
373 + num_experts=e,
374 + experts_per_token=topk,
375 + hidden_dim=k,
376 + intermediate_size=2 * n,
377 + num_local_experts=e,
378 + num_logical_experts=e,
379 + activation=MoEActivation.SILU,
380 + device="cuda",
381 + moe_parallel_config=replace(
382 + FusedMoEParallelConfig.make_no_parallel(), tp_size=2
383 + ),
384 + in_dtype=dtype,
385 + routing_method=RoutingMethodType.TopK,
386 + max_num_tokens=next_power_of_2(m),
387 + )
388 + kernel = _make_trtllm_fp4_moe_kernel(moe_config, quant_config)
389 +
390 + check_deferred_moe_finalize(
391 + moe_config,
392 + lambda: kernel.apply(
393 + hidden_states=a,
394 + w1=w1_q,
395 + w2=w2_q,
396 + topk_weights=topk_weights,
397 + topk_ids=topk_ids,
398 + activation=MoEActivation.SILU,
399 + global_num_experts=e,
400 + expert_map=None,
401 + apply_router_weight_on_input=False,
402 + ),
403 + router_weights=topk_weights,
404 + chunked=chunk_size is not None,
405 + )
406 +
407 +
408 if __name__ == "__main__":
409 test_trtllm_fp4_moe_no_graph(
410 64,
tests/kernels/moe/utils.py
+48
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4
5 import torch
6
@@ -28,6 +29,8 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
29 fused_experts,
30 )
31 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
32 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
33 +from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import moe_unpermute
34 from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (
35 BatchedPrepareAndFinalize,
36 )
@@ -707,6 +710,51 @@ def check_accuracy(a, b, atol, rtol, percent):
710 )
711
712
713 +def check_deferred_moe_finalize(
714 + moe_config: FusedMoEConfig,
715 + run: Callable[[], torch.Tensor | UnfinalizedMoEOutput],
716 + router_weights: torch.Tensor | None = None,
717 + chunked: bool = False,
718 +) -> None:
719 + """Check a kernel that defers its finalize against the finalize it skips.
720 +
721 + ``run`` calls the kernel on fixed inputs, first as built and then with
722 + ``moe_config`` asking to defer, which ``should_defer_moe_finalize`` must
723 + report truthfully. A deferred output reduced by ``moe_unpermute``, the
724 + TRT-LLM finalize kernel, must give the kernel's own finalized output bit for
725 + bit, and modular experts hand the router's weights back as-is. A call the
726 + experts split across kernel launches finalizes instead.
727 + """
728 + # The deferred output views the router's buffer, so compare with a copy.
729 + expected_weights = None if router_weights is None else router_weights.clone()
730 + finalized = run()
731 + assert isinstance(finalized, torch.Tensor)
732 + moe_config.defer_moe_finalize()
733 + output = run()
734 + assert moe_config.should_defer_moe_finalize(finalized.shape[0]) != chunked
735 + if chunked:
736 + torch.testing.assert_close(output, finalized, atol=0, rtol=0)
737 + return
738 +
739 + assert isinstance(output, UnfinalizedMoEOutput)
740 + if expected_weights is not None:
741 + torch.testing.assert_close(
742 + output.expert_weights, expected_weights, atol=0, rtol=0
743 + )
744 + reference = torch.empty_like(finalized)
745 + moe_unpermute(
746 + reference,
747 + output.gemm2_permuted,
748 + output.expert_weights.float(),
749 + output.expanded_idx_to_permuted_idx,
750 + # The kernel reads its valid-row count through this pointer either way.
751 + expert_first_token_offset=output.gemm2_permuted.new_full(
752 + (1,), output.gemm2_permuted.shape[0], dtype=torch.int64
753 + ),
754 + )
755 + torch.testing.assert_close(reference, finalized, atol=0, rtol=0)
756 +
757 +
758 def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8):
759 """Weight/scale layouts for mxfp4 MoE, as (layout, opts) pairs.
760
tests/models/kimi_k3/test_latent_moe_tail.py
+6 -6
@@ -15,6 +15,7 @@ from tests.utils import (
15 multi_process_parallel,
16 )
17 from vllm.distributed import get_tp_group
18 +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
19 from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import (
20 TrtLlmMxfp4ExpertsMonolithic,
21 )
@@ -43,12 +44,11 @@ def test_deferred_finalize_enabled_before_moe_kernel_setup(
44 hidden_dim = LATENT_SIZE
45 hidden_dim_unpadded = LATENT_SIZE
46 experts_per_token = 16
46 - defer_moe_finalize = False
47 + _defer_moe_finalize = False
48 defer_moe_finalize_max_num_tokens = -1
48 -
49 - @property
50 - def use_deferred_moe_finalize(self) -> bool:
51 - return self.defer_moe_finalize
49 + defer_moe_finalize = FusedMoEConfig.defer_moe_finalize
50 + limit_deferred_moe_finalize = FusedMoEConfig.limit_deferred_moe_finalize
51 + use_deferred_moe_finalize = FusedMoEConfig.use_deferred_moe_finalize
52
53 moe_config = FakeMoEConfig()
54 quant_method = SimpleNamespace(
@@ -99,7 +99,7 @@ def test_deferred_finalize_enabled_before_moe_kernel_setup(
99
100 latent_moe_runner.LatentMoERunner()
101
102 - assert moe_config.defer_moe_finalize
102 + assert moe_config.use_deferred_moe_finalize
103 assert moe_config.defer_moe_finalize_max_num_tokens == 128
104 assert initialized_with["experts_per_token"] == 16
105
vllm/model_executor/layers/fused_moe/config.py
+57 -13
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from dataclasses import dataclass
3 +from dataclasses import dataclass, field
4 from enum import IntEnum
5 from typing import Union
6
@@ -1290,13 +1290,12 @@ class FusedMoEConfig:
1290 # Only honored on the non-reduced (late-AR) TP path. Default False.
1291 skip_final_all_reduce: bool = False
1292
1293 - # When True, experts that can stop after GEMM2 are allowed to hand back an
1294 - # UnfinalizedMoEOutput instead of finalized states, leaving the top-k
1295 - # reduction to fuse into the consumer. Set by layers that have such a
1296 - # consumer; read through `use_deferred_moe_finalize`, which applies the
1297 - # guards. Kernels without the capability ignore it. Default False.
1298 - defer_moe_finalize: bool = False
1299 - # Optional consumer capacity for deferred finalize. Negative means unbounded.
1293 + # Requested through `defer_moe_finalize()`, read through
1294 + # `should_defer_moe_finalize()`.
1295 + _defer_moe_finalize: bool = field(default=False, init=False)
1296 + # Most tokens a deferred call covers: the consumer's capacity, lowered by
1297 + # experts that would split a larger call across kernel launches. Negative
1298 + # means unbounded.
1299 defer_moe_finalize_max_num_tokens: int = -1
1300
1301 # SwiGLU clamp limit. When set, backends that do not implement the clamp
@@ -1422,21 +1421,57 @@ class FusedMoEConfig:
1421 def use_deferred_moe_finalize(self) -> bool:
1422 """Whether experts may return an unfinalized output on this deployment.
1423
1425 - Evaluated on read rather than in ``__post_init__`` because
1426 - ``defer_moe_finalize`` is set after construction, like
1427 - ``skip_final_all_reduce``.
1424 + Evaluated on read rather than in ``__post_init__`` because deferral is
1425 + requested after construction, like ``skip_final_all_reduce``.
1426 """
1427 # The consumer fuses a TP all-reduce. Other parallel modes require a
1430 - # combine or reduce-scatter after the experts and cannot defer it.
1428 + # combine or reduce-scatter after the experts and cannot defer it, and
1429 + # the consumer has no way to strip hidden-dim padding from GEMM2 rows.
1430 return (
1432 - self.defer_moe_finalize
1431 + self._defer_moe_finalize
1432 and self.tp_size > 1
1433 and self.dp_size == 1
1434 and self.ep_size == 1
1435 and self.pcp_size == 1
1436 and not self.is_sequence_parallel
1437 + and self.hidden_dim == self.hidden_dim_unpadded
1438 )
1439
1440 + def defer_moe_finalize(self, max_num_tokens: int = -1) -> None:
1441 + """Ask the experts to leave the top-k reduction to the layer's consumer.
1442 +
1443 + A layer whose consumer can fuse the top-k reduction (e.g. into its TP
1444 + all-reduce) calls this once while the model is built, and only when its
1445 + experts are TRTLLM-Gen ones that can stop after GEMM2 (the
1446 + ``do_finalize=False`` path). Other experts ignore the request and
1447 + always finalize, which ``should_defer_moe_finalize`` can't see, so the
1448 + layer checks the quant method's ``experts_cls`` first, as Kimi-K3 does.
1449 + From then on:
1450 +
1451 + - The experts return an ``UnfinalizedMoEOutput`` instead of finalized
1452 + states for every call ``should_defer_moe_finalize`` accepts.
1453 + - ``should_defer_moe_finalize(num_tokens)`` is the answer for a call:
1454 + it requires a TP-only deployment without hidden-dim padding
1455 + (``use_deferred_moe_finalize``), a non-empty call and at most
1456 + ``defer_moe_finalize_max_num_tokens`` tokens. The cap starts at
1457 + ``max_num_tokens`` and only ever goes down: experts that would split a
1458 + larger call across kernel launches lower it to their single-launch
1459 + size when they are built, since each launch permutes into its own
1460 + buffer.
1461 + - The model asks ``should_defer_moe_finalize`` before each call and takes
1462 + the matching path. A deferred call runs the runner's ``_forward_impl``
1463 + directly, since the MoE custom op returns tensors only, and the
1464 + consumer then owns the top-k reduction, the shared-expert add and the
1465 + all-reduce.
1466 +
1467 + Args:
1468 + max_num_tokens: Most tokens per call the consumer can take in
1469 + deferred form. Negative means no limit of its own.
1470 +
1471 + """
1472 + self._defer_moe_finalize = True
1473 + self.limit_deferred_moe_finalize(max_num_tokens)
1474 +
1475 def should_defer_moe_finalize(self, num_tokens: int) -> bool:
1476 """Return whether this invocation may defer the top-k reduction."""
1477 max_num_tokens = self.defer_moe_finalize_max_num_tokens
@@ -1446,6 +1481,15 @@ class FusedMoEConfig:
1481 and (max_num_tokens < 0 or num_tokens <= max_num_tokens)
1482 )
1483
1484 + def limit_deferred_moe_finalize(self, max_num_tokens: int) -> None:
1485 + """Finalize calls above ``max_num_tokens`` even when deferring.
1486 +
1487 + Negative means no limit, and leaves the current one in place.
1488 + """
1489 + current = self.defer_moe_finalize_max_num_tokens
1490 + if max_num_tokens >= 0 and (current < 0 or max_num_tokens < current):
1491 + self.defer_moe_finalize_max_num_tokens = max_num_tokens
1492 +
1493 @property
1494 def use_deepep_ht_kernels(self):
1495 return self.moe_parallel_config.use_deepep_ht_kernels
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+20 -3
@@ -12,6 +12,10 @@ from vllm.model_executor.layers.fused_moe.config import (
12 FusedMoEQuantConfig,
13 RoutingMethodType,
14 )
15 +from vllm.model_executor.layers.fused_moe.moe_output import (
16 + UnfinalizedMoEOutput,
17 + convert_flashinfer_moe_output,
18 +)
19 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
20 TopKWeightAndReduceNoOP,
21 )
@@ -273,7 +277,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
277 workspace2: torch.Tensor,
278 expert_tokens_meta: mk.ExpertTokensMetadata | None,
279 apply_router_weight_on_input: bool,
276 - ):
280 + ) -> UnfinalizedMoEOutput | None:
281 import flashinfer
282 from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout
283
@@ -297,7 +301,9 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
301 weight_layout = WeightLayout.BlockMajorK
302 hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale)
303
300 - flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
304 + num_tokens = hidden_states.shape[0]
305 + defer = self.moe_config.should_defer_moe_finalize(num_tokens)
306 + flashinfer_output = flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
307 topk_ids=(topk_ids, topk_weights),
308 routing_bias=None,
309 hidden_states=hidden_states,
@@ -321,9 +327,20 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
327 use_shuffled_weight=use_shuffled_weight,
328 weight_layout=weight_layout,
329 fp8_quantization_type=fp8_quant_type,
324 - output=output,
330 + do_finalize=not defer,
331 + output=None if defer else output,
332 tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
333 )
334 + if not defer:
335 + return None
336 + routed_output = convert_flashinfer_moe_output(
337 + flashinfer_output,
338 + do_finalize=False,
339 + num_tokens=num_tokens,
340 + top_k=topk_ids.size(1),
341 + )
342 + assert isinstance(routed_output, UnfinalizedMoEOutput)
343 + return routed_output
344
345
346 class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolithic):
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+36 -9
@@ -289,6 +289,18 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
289 Moved from trtllm_moe.py.
290 """
291
292 + def __init__(
293 + self,
294 + moe_config: FusedMoEConfig,
295 + quant_config: FusedMoEQuantConfig,
296 + **kwargs,
297 + ):
298 + super().__init__(moe_config, quant_config, **kwargs)
299 + # Each launch permutes into its own buffer, so a chunked call finalizes.
300 + moe_config.limit_deferred_moe_finalize(
301 + self._max_supported_tokens(self.topk, moe_config.num_experts)
302 + )
303 +
304 @staticmethod
305 def _supports_parallel_config(
306 moe_parallel_config: FusedMoEParallelConfig,
@@ -343,7 +355,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
355
356 def _invoke_kernel(
357 self,
346 - output: torch.Tensor,
358 + output: torch.Tensor | None,
359 x_quant: torch.Tensor,
360 x_scale: torch.Tensor | None,
361 topk_ids: torch.Tensor,
@@ -355,10 +367,11 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
367 local_num_experts: int,
368 local_expert_offset: int,
369 topk: int,
358 - ) -> None:
370 + ) -> UnfinalizedMoEOutput | None:
371 + """Finalize into ``output``, or stop after GEMM2 when it is None."""
372 from flashinfer import trtllm_fp4_block_scale_routed_moe
373
361 - trtllm_fp4_block_scale_routed_moe(
374 + flashinfer_output = trtllm_fp4_block_scale_routed_moe(
375 topk_ids=(topk_ids, topk_weights),
376 routing_bias=None,
377 hidden_states=x_quant,
@@ -386,12 +399,22 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
399 # Modular kernel receives pre-routed tokens, so routing is already
400 # done. Use Renormalize as a safe default the TRTLLM kernel supports.
401 routing_method_type=RoutingMethodType.Renormalize,
389 - do_finalize=True,
402 + do_finalize=output is not None,
403 enable_pdl=True,
404 activation_type=self._flashinfer_activation_type(activation),
405 output=output,
406 tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
407 )
408 + if output is not None:
409 + return None
410 + routed_output = convert_flashinfer_moe_output(
411 + flashinfer_output,
412 + do_finalize=False,
413 + num_tokens=x_quant.shape[0],
414 + top_k=topk,
415 + )
416 + assert isinstance(routed_output, UnfinalizedMoEOutput)
417 + return routed_output
418
419 def apply(
420 self,
@@ -410,7 +433,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
433 workspace2: torch.Tensor,
434 expert_tokens_meta: mk.ExpertTokensMetadata | None,
435 apply_router_weight_on_input: bool,
413 - ):
436 + ) -> UnfinalizedMoEOutput | None:
437 topk_ids = topk_ids.to(dtype=torch.int32)
438
439 topk = topk_ids.size(-1)
@@ -428,13 +451,17 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
451 assert self.w1_scale is not None
452 assert self.w2_scale is not None
453
431 - # Chunk tokens so the batched-GEMM grid stays within CUDA limits.
454 + # Chunk tokens so the batched-GEMM grid stays within CUDA limits. Each
455 + # launch permutes into its own buffer, so only a run that fits in one
456 + # launch can leave the top-k reduction to a deferring consumer.
457 M = x_quant.size(0)
458 chunk_size = self._max_supported_tokens(topk, global_num_experts)
459 + defer = chunk_size >= M and self.moe_config.should_defer_moe_finalize(M)
460 + unfinalized: UnfinalizedMoEOutput | None = None
461 for start in range(0, M, chunk_size):
462 end = min(start + chunk_size, M)
436 - self._invoke_kernel(
437 - output[start:end],
463 + unfinalized = self._invoke_kernel(
464 + None if defer else output[start:end],
465 x_quant[start:end],
466 None if x_scale is None else x_scale[start:end],
467 topk_ids[start:end],
@@ -448,4 +475,4 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
475 topk,
476 )
477
451 - return output
478 + return unfinalized
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+33 -7
@@ -302,6 +302,16 @@ class TrtLlmNvFp4ExpertsBase:
302 class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular):
303 """Modular version of the implementation (just the experts)."""
304
305 + def __init__(
306 + self,
307 + moe_config: FusedMoEConfig,
308 + quant_config: FusedMoEQuantConfig,
309 + per_token_activation: bool = False,
310 + ):
311 + super().__init__(moe_config, quant_config, per_token_activation)
312 + # Each launch permutes into its own buffer, so a chunked call finalizes.
313 + moe_config.limit_deferred_moe_finalize(self._get_chunk_size())
314 +
315 @staticmethod
316 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
317 """The modular implementation supports all parallel configs."""
@@ -333,7 +343,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
343
344 def _invoke_kernel(
345 self,
336 - output: torch.Tensor,
346 + output: torch.Tensor | None,
347 hidden_states: torch.Tensor,
348 w1: torch.Tensor,
349 w2: torch.Tensor,
@@ -342,7 +352,8 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
352 activation: MoEActivation,
353 global_num_experts: int,
354 a1q_scale: torch.Tensor | None,
345 - ):
355 + ) -> UnfinalizedMoEOutput | None:
356 + """Finalize into ``output``, or stop after GEMM2 when it is None."""
357 import flashinfer
358
359 assert self.quant_config.w1_scale is not None
@@ -359,7 +370,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
370 output1_scale_gate_scalar = self.quant_config.g1_alphas
371
372 # Invoke kernel.
362 - flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(
373 + flashinfer_output = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(
374 topk_ids=(topk_ids, topk_weights),
375 routing_bias=None,
376 hidden_states=hidden_states,
@@ -387,7 +398,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
398 local_num_experts=self.local_num_experts,
399 routed_scaling_factor=None,
400 routing_method_type=1, # not used
390 - do_finalize=True,
401 + do_finalize=output is not None,
402 activation_type=activation_to_flashinfer_int(activation),
403 per_token_scale=per_token_scale,
404 output=output,
@@ -395,6 +406,16 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
406 fi_moe_largest_bucket(self.moe_config), self._get_chunk_size()
407 ),
408 )
409 + if output is not None:
410 + return None
411 + routed_output = convert_flashinfer_moe_output(
412 + flashinfer_output,
413 + do_finalize=False,
414 + num_tokens=hidden_states.shape[0],
415 + top_k=self.topk,
416 + )
417 + assert isinstance(routed_output, UnfinalizedMoEOutput)
418 + return routed_output
419
420 def apply(
421 self,
@@ -413,7 +434,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
434 workspace2: torch.Tensor,
435 expert_tokens_meta: mk.ExpertTokensMetadata | None,
436 apply_router_weight_on_input: bool,
416 - ):
437 + ) -> UnfinalizedMoEOutput | None:
438 assert self._supports_activation(activation)
439 # Per-token defers input quant to _invoke_kernel, so a1q_scale is None.
440 assert a1q_scale is not None or self.per_token_activation
@@ -425,8 +446,12 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
446 chunk_size = self._get_chunk_size()
447
448 if chunk_size >= M:
428 - self._invoke_kernel(
429 - output,
449 + # Each launch permutes into its own buffer, so only a run that fits
450 + # in one launch can leave the top-k reduction to a deferring
451 + # consumer.
452 + defer = self.moe_config.should_defer_moe_finalize(M)
453 + return self._invoke_kernel(
454 + None if defer else output,
455 hidden_states,
456 w1,
457 w2,
@@ -450,6 +475,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
475 global_num_experts,
476 None if a1q_scale is None else a1q_scale[start:end],
477 )
478 + return None
479
480
481 class TrtLlmNvFp4ExpertsMonolithic(
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+2 -2
@@ -148,7 +148,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
148 topk_ids: torch.Tensor,
149 shared_experts: "SharedExperts | None",
150 shared_experts_input: torch.Tensor | None,
151 - ) -> torch.Tensor:
151 + ) -> torch.Tensor | UnfinalizedMoEOutput:
152 """Apply the MoE operation using modular kernels.
153
154 Args:
@@ -160,7 +160,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
160 shared_experts: The shared experts module, if any
161
162 Returns:
163 - Output tensor from routed experts.
163 + Finalized routed states or a deferred-finalize output.
164
165 """
166 raise NotImplementedError
vllm/model_executor/layers/fused_moe/modular_kernel.py
+38 -13
@@ -238,6 +238,15 @@ class FusedMoEPrepareAndFinalize(ABC):
238 """
239 return False
240
241 + def supports_deferred_moe_finalize(self) -> bool:
242 + """Whether ``finalize`` can be skipped for a deferring consumer.
243 +
244 + An implementation opts in only if everything it does in ``finalize``
245 + -- the top-k reduction, and any combine or reduce-scatter -- is work
246 + the consumer takes over.
247 + """
248 + return False
249 +
250
251 # TODO: pass FusedMoEParallelConfig in as ctor parameter?
252 class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):
@@ -404,9 +413,6 @@ class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):
413 described above for the monolithic case.
414 """
415
407 - def supports_deferred_moe_finalize(self) -> bool:
408 - return False
409 -
416 @abstractmethod
417 def prepare(
418 self,
@@ -895,10 +901,14 @@ class FusedMoEExpertsModular(FusedMoEExperts):
901 workspace2: torch.Tensor,
902 expert_tokens_meta: ExpertTokensMetadata | None,
903 apply_router_weight_on_input: bool,
898 - ) -> None:
904 + ) -> UnfinalizedMoEOutput | None:
905 """This function computes the intermediate result of a Mixture of Experts
906 (MoE) layer using two sets of weights, w1 and w2.
907
908 + Writes into `output` and returns None, unless the implementation stopped
909 + after GEMM2 and left the top-k reduction to a fused consumer, in which
910 + case `output` is untouched and the unfinalized result is returned.
911 +
912 Args:
913 output: (torch.Tensor): The unweighted, unreduced output tensor.
914 hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE
@@ -1270,7 +1280,7 @@ class FusedMoEKernelModularImpl:
1280 apply_router_weight_on_input: bool,
1281 expert_tokens_meta: ExpertTokensMetadata | None,
1282 output_alias: torch.Tensor | None = None,
1273 - ) -> torch.Tensor:
1283 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1284 _, M_full, N, K, top_k = self.fused_experts.moe_problem_size(
1285 a1q, w1, w2, topk_ids
1286 )
@@ -1318,7 +1328,7 @@ class FusedMoEKernelModularImpl:
1328 elif use_output_alias:
1329 fused_out = output_alias
1330
1321 - self.fused_experts.apply(
1331 + unfinalized = self.fused_experts.apply(
1332 output=fused_out,
1333 hidden_states=a1q,
1334 w1=w1,
@@ -1336,6 +1346,11 @@ class FusedMoEKernelModularImpl:
1346 apply_router_weight_on_input=apply_router_weight_on_input,
1347 )
1348
1349 + # Experts that stopped after GEMM2 wrote nothing into `fused_out`; the
1350 + # top-k reduction they left open belongs to whoever takes this.
1351 + if isinstance(unfinalized, UnfinalizedMoEOutput):
1352 + return unfinalized
1353 +
1354 return fused_out
1355
1356 def _finalize(
@@ -1426,7 +1441,7 @@ class FusedMoEKernelModularImpl:
1441 apply_router_weight_on_input: bool = False,
1442 shared_experts: SharedExperts | None = None,
1443 shared_experts_input: torch.Tensor | None = None,
1429 - ) -> torch.Tensor:
1444 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1445 """This function computes a Mixture of Experts (MoE) layer using two sets
1446 of weights, w1 and w2, and top-k gating mechanism.
1447
@@ -1453,7 +1468,9 @@ class FusedMoEKernelModularImpl:
1468 hidden_states before latent projection.
1469
1470 Returns:
1456 - torch.Tensor: The output tensor after applying the MoE layer.
1471 + torch.Tensor: The output tensor after applying the MoE layer, or
1472 + the unfinalized output when the experts left the top-k reduction to
1473 + a fused consumer.
1474
1475 """
1476 output = torch.empty_like(hidden_states)
@@ -1498,6 +1515,17 @@ class FusedMoEKernelModularImpl:
1515 if lora_ctx is not None:
1516 lora_ctx.original_hidden_states = None
1517
1518 + if isinstance(fused_out, UnfinalizedMoEOutput):
1519 + # Nothing below can run on an unfinalized output: finalize is the
1520 + # local top-k reduction (and, for DP/EP, the combine) that the
1521 + # consumer takes over.
1522 + if not self.prepare_finalize.supports_deferred_moe_finalize():
1523 + raise RuntimeError(
1524 + f"{type(self.prepare_finalize).__name__} cannot pass through "
1525 + "a deferred MoE output."
1526 + )
1527 + return fused_out
1528 +
1529 return self._finalize(
1530 output,
1531 fused_out,
@@ -1653,10 +1681,7 @@ class FusedMoEKernel:
1681 return self.prepare_finalize.output_is_reduced()
1682
1683 def supports_deferred_moe_finalize(self) -> bool:
1656 - return (
1657 - isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic)
1658 - and self.prepare_finalize.supports_deferred_moe_finalize()
1659 - )
1684 + return self.prepare_finalize.supports_deferred_moe_finalize()
1685
1686 def apply_monolithic(
1687 self,
@@ -1703,7 +1728,7 @@ class FusedMoEKernel:
1728 apply_router_weight_on_input: bool,
1729 shared_experts: SharedExperts | None = None,
1730 shared_experts_input: torch.Tensor | None = None,
1706 - ) -> torch.Tensor:
1731 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1732 assert isinstance(self.impl, FusedMoEKernelModularImpl)
1733 return self.impl.apply(
1734 hidden_states=hidden_states,
vllm/model_executor/layers/fused_moe/moe_output.py
+6 -5
@@ -9,7 +9,7 @@ import torch
9
10 @dataclass
11 class UnfinalizedMoEOutput:
12 - """Unfinalized output of a monolithic MoE kernel.
12 + """Unfinalized output of a MoE kernel.
13
14 Kernels that can stop after GEMM2 (the TRTLLM-Gen ``do_finalize=False`` path)
15 hand back their permuted, unweighted output plus the routing weights and the
@@ -24,10 +24,11 @@ class UnfinalizedMoEOutput:
24
25 # [num_permuted_rows, hidden_dim] permuted, unweighted GEMM2 output.
26 gemm2_permuted: torch.Tensor
27 - # [num_tokens, top_k] routing weights, in the activation dtype -- consumers
28 - # index them as such, and a buffer typed wider than its contents reads as
29 - # garbage. These already carry routed_scaling_factor for routing methods that
30 - # fold it in.
27 + # [num_tokens, top_k] routing weights, in the activation dtype when the
28 + # kernel routed itself and in the router's dtype for modular experts --
29 + # consumers index them as such, and a buffer typed wider than its contents
30 + # reads as garbage. These already carry routed_scaling_factor for routing
31 + # methods that fold it in.
32 expert_weights: torch.Tensor
33 # [num_tokens, top_k] int32 permute map; -1 = expert not local to this rank.
34 expanded_idx_to_permuted_idx: torch.Tensor
vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py
+5
@@ -38,6 +38,11 @@ def _quantize_input(
38
39
40 class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular):
41 + def supports_deferred_moe_finalize(self) -> bool:
42 + # finalize() here is the top-k weight-and-reduce and nothing else,
43 + # which is exactly the work a deferring consumer takes over.
44 + return True
45 +
46 @property
47 def activation_format(self) -> mk.FusedMoEActivationFormat:
48 return mk.FusedMoEActivationFormat.Standard
vllm/model_executor/layers/fused_moe/routed_experts.py
+2 -2
@@ -1224,7 +1224,7 @@ class RoutedExperts(PluggableLayer):
1224 topk_ids: torch.Tensor,
1225 shared_experts: "SharedExperts | None" = None,
1226 shared_experts_input: torch.Tensor | None = None,
1227 - ) -> torch.Tensor:
1227 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1228 """Execute routed experts using the quantization method's apply function.
1229
1230 This is called by the runner after router selection (for modular kernels)
@@ -1239,7 +1239,7 @@ class RoutedExperts(PluggableLayer):
1239 shared_experts_input: Input for shared experts (if any)
1240
1241 Returns:
1242 - Output tensor from routed experts.
1242 + Finalized routed states or a deferred-finalize output.
1243
1244 """
1245 assert not self.quant_method.is_monolithic
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+2 -1
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
24 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
25 XPUExpertsMxFp4,
26 )
27 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
28 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
29 B12X_BACKENDS,
30 Mxfp4MoeBackend,
@@ -227,7 +228,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
228 topk_ids: torch.Tensor,
229 shared_experts: SharedExperts | None,
230 shared_experts_input: torch.Tensor | None,
230 - ) -> torch.Tensor:
231 + ) -> torch.Tensor | UnfinalizedMoEOutput:
232 assert self.moe_kernel is not None
233 return self.moe_kernel.apply(
234 x,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+1 -1
@@ -303,7 +303,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
303 topk_ids: torch.Tensor,
304 shared_experts: SharedExperts | None,
305 shared_experts_input: torch.Tensor | None,
306 - ) -> torch.Tensor:
306 + ) -> torch.Tensor | UnfinalizedMoEOutput:
307 assert self.moe_kernel is not None
308 return self.moe_kernel.apply(
309 x,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
+2 -1
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.config import (
18 FusedMoEConfig,
19 FusedMoEQuantConfig,
20 )
21 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
22 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
23 convert_to_fp8_moe_kernel_format,
24 make_fp8_moe_kernel,
@@ -375,7 +376,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
376 topk_ids: torch.Tensor,
377 shared_experts: SharedExperts | None,
378 shared_experts_input: torch.Tensor | None,
378 - ) -> torch.Tensor:
379 + ) -> torch.Tensor | UnfinalizedMoEOutput:
380 assert not self.is_monolithic
381 assert self.moe_kernel is not None
382 return self.moe_kernel.apply(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py
+2 -1
@@ -12,6 +12,7 @@ from vllm.model_executor.layers.fused_moe.config import (
12 FusedMoEConfig,
13 FusedMoEQuantConfig,
14 )
15 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
16 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
17 convert_to_fp8_moe_kernel_format,
18 make_fp8_moe_kernel,
@@ -188,7 +189,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
189 topk_ids: torch.Tensor,
190 shared_experts: SharedExperts | None,
191 shared_experts_input: torch.Tensor | None,
191 - ) -> torch.Tensor:
192 + ) -> torch.Tensor | UnfinalizedMoEOutput:
193 assert not self.is_monolithic
194 assert self.moe_kernel is not None
195 return self.moe_kernel.apply(
vllm/model_executor/layers/quantization/fp8.py
+2 -1
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe import (
24 UnquantizedFusedMoEMethod,
25 )
26 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
27 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
28 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
29 convert_to_fp8_moe_kernel_format,
30 make_fp8_moe_kernel,
@@ -827,7 +828,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
828 topk_ids: torch.Tensor,
829 shared_experts: SharedExperts | None,
830 shared_experts_input: torch.Tensor | None,
830 - ) -> torch.Tensor:
831 + ) -> torch.Tensor | UnfinalizedMoEOutput:
832 assert not self.is_monolithic
833 assert self.moe_kernel is not None
834 return self.moe_kernel.apply(
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
+2 -1
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
23 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
24 FusedMoEMethodBase,
25 )
26 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
27 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
28 B12X_BACKENDS,
29 Mxfp4MoeBackend,
@@ -228,7 +229,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase):
229 topk_ids: torch.Tensor,
230 shared_experts: SharedExperts | None,
231 shared_experts_input: torch.Tensor | None,
231 - ) -> torch.Tensor:
232 + ) -> torch.Tensor | UnfinalizedMoEOutput:
233 assert self.moe_kernel is not None
234 return self.moe_kernel.apply(
235 x,
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_moe.py
+2 -1
@@ -12,6 +12,7 @@ from vllm.model_executor.layers.fused_moe import (
12 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
13 FusedMoEMethodBase,
14 )
15 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
16 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
17 convert_to_fp8_moe_kernel_format,
18 make_fp8_moe_kernel,
@@ -179,7 +180,7 @@ class INCMxfp8MoEMethod(FusedMoEMethodBase):
180 topk_ids: torch.Tensor,
181 shared_experts: SharedExperts | None,
182 shared_experts_input: torch.Tensor | None,
182 - ) -> torch.Tensor:
183 + ) -> torch.Tensor | UnfinalizedMoEOutput:
184 assert self.moe_kernel is not None
185 return self.moe_kernel.apply(
186 x,
vllm/model_executor/layers/quantization/modelopt.py
+3 -3
@@ -690,7 +690,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
690 topk_ids: torch.Tensor,
691 shared_experts: SharedExperts | None,
692 shared_experts_input: torch.Tensor | None,
693 - ) -> torch.Tensor:
693 + ) -> torch.Tensor | UnfinalizedMoEOutput:
694 assert not self.is_monolithic
695 assert self.moe_kernel is not None
696 return self.moe_kernel.apply(
@@ -1102,7 +1102,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1102 topk_ids: torch.Tensor,
1103 shared_experts: SharedExperts | None,
1104 shared_experts_input: torch.Tensor | None,
1105 - ) -> torch.Tensor:
1105 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1106 assert not self.is_monolithic
1107 assert self.moe_kernel is not None
1108 return self.moe_kernel.apply(
@@ -1471,7 +1471,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1471 topk_ids: torch.Tensor,
1472 shared_experts: SharedExperts | None,
1473 shared_experts_input: torch.Tensor | None,
1474 - ) -> torch.Tensor:
1474 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1475 assert not self.is_monolithic
1476 assert self.moe_kernel is not None
1477 return self.moe_kernel.apply(
vllm/model_executor/layers/quantization/mxfp4.py
+2 -2
@@ -433,7 +433,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
433 topk_ids: torch.Tensor,
434 shared_experts: SharedExperts | None,
435 shared_experts_input: torch.Tensor | None,
436 - ) -> torch.Tensor:
436 + ) -> torch.Tensor | UnfinalizedMoEOutput:
437 assert not self.is_monolithic
438 assert self.moe_kernel is not None
439 return self.moe_kernel.apply(
@@ -860,7 +860,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
860 topk_ids: torch.Tensor,
861 shared_experts: SharedExperts | None,
862 shared_experts_input: torch.Tensor | None,
863 - ) -> torch.Tensor:
863 + ) -> torch.Tensor | UnfinalizedMoEOutput:
864 assert not self.is_monolithic
865 assert self.moe_kernel is not None
866 return self.moe_kernel.apply(
vllm/model_executor/layers/quantization/online/moe_base.py
+1 -1
@@ -166,7 +166,7 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
166 topk_ids: torch.Tensor,
167 shared_experts: SharedExperts | None,
168 shared_experts_input: torch.Tensor | None,
169 - ) -> torch.Tensor:
169 + ) -> torch.Tensor | UnfinalizedMoEOutput:
170 assert not self.is_monolithic
171 assert self.moe_kernel is not None
172 return self.moe_kernel.apply(
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3 -3
@@ -976,7 +976,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
976 topk_ids: torch.Tensor,
977 shared_experts: SharedExperts | None,
978 shared_experts_input: torch.Tensor | None,
979 - ) -> torch.Tensor:
979 + ) -> torch.Tensor | UnfinalizedMoEOutput:
980 assert self.moe_kernel is not None
981 return self.moe_kernel.apply(
982 hidden_states=x,
@@ -1935,7 +1935,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1935 topk_ids: torch.Tensor,
1936 shared_experts: SharedExperts | None,
1937 shared_experts_input: torch.Tensor | None,
1938 - ) -> torch.Tensor:
1938 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1939 assert self.moe_kernel is not None
1940 return self.moe_kernel.apply(
1941 hidden_states=x,
@@ -2182,7 +2182,7 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):
2182 topk_ids: torch.Tensor,
2183 shared_experts: SharedExperts | None,
2184 shared_experts_input: torch.Tensor | None,
2185 - ) -> torch.Tensor:
2185 + ) -> torch.Tensor | UnfinalizedMoEOutput:
2186 assert self.moe_kernel is not None
2187 return self.moe_kernel.apply(
2188 x,
vllm/models/kimi_k3/nvidia/latent_moe_runner.py
+4 -10
@@ -115,10 +115,11 @@ class LatentMoERunner(MoERunner):
115 experts_cls = getattr(self._quant_method, "experts_cls", None)
116 # The kernel instance is built after weight loading, while the tail
117 # must register its CuTeDSL warmup units during runner construction.
118 - self.moe_config.defer_moe_finalize = (
118 + if (
119 experts_cls is TrtLlmMxfp4ExpertsMonolithic
120 or experts_cls is TrtLlmNvFp4ExpertsMonolithic
121 - ) and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded
121 + ) and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded:
122 + self.moe_config.defer_moe_finalize()
123 from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import (
124 KimiK3LatentMoETailOp,
125 )
@@ -136,20 +137,13 @@ class LatentMoERunner(MoERunner):
137 rms_eps=norm.variance_epsilon,
138 )
139 self._k3_latent_moe_tail_op = op
139 - self.moe_config.defer_moe_finalize_max_num_tokens = (
140 - op.contract.max_num_tokens
141 - if self.moe_config.use_deferred_moe_finalize
142 - else -1
143 - )
140 if self.moe_config.use_deferred_moe_finalize:
141 + self.moe_config.limit_deferred_moe_finalize(op.contract.max_num_tokens)
142 logger.info_once(
143 "K3 latent-MoE tail fusion with deferred top-k finalization "
144 "is enabled for up to %d tokens.",
145 op.contract.max_num_tokens,
146 )
150 - else:
151 - self.moe_config.defer_moe_finalize = False
152 - self.moe_config.defer_moe_finalize_max_num_tokens = -1
147
148 def _get_zero_residual(
149 self,