diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 6329780f9dc9..5d55876e8138 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,11 +37,20 @@ 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, + 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 ( rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, 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, @@ -767,3 +781,109 @@ def _supports_quant_scheme( assert supported == expected, reason if not expected: 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, 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 + set_random_seed(7) + with set_current_vllm_config(vllm_config): + 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( + 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 = 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( + 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, + ) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 1cd204d4b00b..e4b11c9548b2 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,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]) diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index bc0711200b2f..292e4549bb7d 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 to the caller.""" + 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, 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=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, + ), + router_weights=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..9274a08af4a3 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,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 +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, ) @@ -707,6 +710,51 @@ def check_accuracy(a, b, atol, rtol, percent): ) +def check_deferred_moe_finalize( + 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 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 ``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() + finalized = run() + assert isinstance(finalized, torch.Tensor) + moe_config.defer_moe_finalize() + output = run() + assert moe_config.should_defer_moe_finalize(finalized.shape[0]) != chunked + if chunked: + torch.testing.assert_close(output, finalized, atol=0, rtol=0) + return + + assert isinstance(output, UnfinalizedMoEOutput) + if expected_weights is not None: + torch.testing.assert_close( + output.expert_weights, expected_weights, 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): """Weight/scale layouts for mxfp4 MoE, as (layout, opts) pairs. diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py index 5fde9ff86d4a..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,12 +44,11 @@ class FakeMoEConfig: hidden_dim = LATENT_SIZE hidden_dim_unpadded = LATENT_SIZE experts_per_token = 16 - defer_moe_finalize = False + _defer_moe_finalize = False defer_moe_finalize_max_num_tokens = -1 - - @property - def use_deferred_moe_finalize(self) -> bool: - return self.defer_moe_finalize + 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( @@ -99,7 +99,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 0c82d28096dc..9221e9382428 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,13 +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. 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 @@ -1422,21 +1421,57 @@ 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. + # 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 + 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 + 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, 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 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 + consumer then owns the top-k reduction, the shared-expert add and the + all-reduce. + + 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 + 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 @@ -1446,6 +1481,15 @@ 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. + + Negative means no limit, and leaves the current one in place. + """ + current = self.defer_moe_finalize_max_num_tokens + if max_num_tokens >= 0 and (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_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..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, @@ -343,7 +355,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 +367,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 +399,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 +433,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 +451,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 +475,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..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.""" @@ -333,7 +343,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 +352,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 +370,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 +398,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 +406,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 +434,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 +446,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 +475,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, 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,