[MoE] Defer the TRTLLM-Gen top-k finalize on the modular path

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,81 @@ 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 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
+        )
+        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, which deferral needs.
+        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,
+            ),
+            router_weights=topk_weights,
+        )

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,6 +663,154 @@ def test_trtllm_gen_mxfp4_fused_moe(
     check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8)
 
 
+@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_deferred_finalize(
+    monolithic: bool,
+    num_tokens: int,
+    chunk_size: int | None,
+    monkeypatch: pytest.MonkeyPatch,
+    workspace_init,
+):
+    """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,
+        FusedMoEParallelConfig,
+        RoutingMethodType,
+    )
+    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,
+        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, which deferral needs.
+    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.Renormalize,
+        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,
+        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(
+            moe_config,
+            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,
+        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,
+        ),
+        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])