@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 29c9af5211

[Kimi K3 Refactor] Add `UnfinalizedMoEOutput` proto following up for #53152 (#53310)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Aug 24, 2026 at 11:21 UTC 29c9af5211e618bfb78c4140db9e814f1a838aa7

6 files changed +21 -9

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

+11 -4

@@ -825,8 +825,13 @@ class MoERunner(MoERunnerInterface):

825 | tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]

826 ):

827 if self.do_naive_dispatch_combine:

828 + if isinstance(hidden_states, UnfinalizedMoEOutput):

829 + raise RuntimeError(

830 + "Naive expert-parallel combine cannot consume a deferred "

831 + "MoE output."

832 + )

833 hidden_states = get_ep_group().combine(

829 - cast(torch.Tensor, hidden_states),

834 + hidden_states,

835 self.moe_config.is_sequence_parallel,

836 )

837

@@ -834,9 +839,11 @@ class MoERunner(MoERunnerInterface):

839 self.moe_config.pcp_size > 1

840 and not self.moe_config.moe_parallel_config.use_all2all_kernels

841 ):

837 - hidden_states = get_pcp_group().reduce_scatter(

838 - cast(torch.Tensor, hidden_states), dim=0

839 - )

842 + if isinstance(hidden_states, UnfinalizedMoEOutput):

843 + raise RuntimeError(

844 + "PCP reduce-scatter cannot consume a deferred MoE output."

845 + )

846 + hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0)

847

848 if self.shared_experts is not None:

849 assert shared_output is not None

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+2 -1

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.config import (

17 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (

18 FusedMoEMethodBase,

19 )

20 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput

21 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (

22 UnquantizedMoeBackend,

23 convert_to_unquantized_kernel_format,

@@ -299,7 +300,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

300 x: torch.Tensor,

301 router_logits: torch.Tensor,

302 input_ids: torch.Tensor | None = None,

302 - ) -> torch.Tensor:

303 + ) -> torch.Tensor | UnfinalizedMoEOutput:

304 assert self.is_monolithic

305 assert self.moe_kernel is not None

306 return self.moe_kernel.apply_monolithic(

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py

+2 -1

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.config import (

14 FusedMoEConfig,

15 FusedMoEQuantConfig,

16 )

17 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput

18 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

19 NvFp4MoeBackend,

20 convert_to_nvfp4_moe_kernel_format,

@@ -278,7 +279,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

279 x: torch.Tensor,

280 router_logits: torch.Tensor,

281 input_ids: torch.Tensor | None = None,

281 - ) -> torch.Tensor:

282 + ) -> torch.Tensor | UnfinalizedMoEOutput:

283 assert self.is_monolithic

284 assert self.moe_kernel is not None

285 return self.moe_kernel.apply_monolithic(

vllm/model_executor/layers/quantization/modelopt.py

+2 -1

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe import (

24 RoutedExperts,

25 SharedExperts,

26 )

27 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput

28 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

29 Fp8MoeBackend,

30 convert_to_fp8_moe_kernel_format,

@@ -1624,7 +1625,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1625 x: torch.Tensor,

1626 router_logits: torch.Tensor,

1627 input_ids: torch.Tensor | None = None,

1627 - ) -> torch.Tensor:

1628 + ) -> torch.Tensor | UnfinalizedMoEOutput:

1629 assert self.is_monolithic

1630 assert self.moe_kernel is not None

1631 return self.moe_kernel.apply_monolithic(

vllm/model_executor/layers/quantization/online/moe_base.py

+2 -1

@@ -10,6 +10,7 @@ from vllm.model_executor.layers.fused_moe import (

10 RoutedExperts,

11 SharedExperts,

12 )

13 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput

14 from vllm.model_executor.model_loader.reload.layerwise import (

15 initialize_online_processing,

16 )

@@ -139,7 +140,7 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):

140 x: torch.Tensor,

141 router_logits: torch.Tensor,

142 input_ids: torch.Tensor | None = None,

142 - ) -> torch.Tensor:

143 + ) -> torch.Tensor | UnfinalizedMoEOutput:

144 assert self.is_monolithic

145 assert self.moe_kernel is not None

146 return self.moe_kernel.apply_monolithic(

vllm/model_executor/layers/quantization/quark/quark_moe.py

+2 -1

@@ -26,6 +26,7 @@ from vllm.model_executor.layers.fused_moe.config import (

26 mxfp4_w4a16_moe_quant_config,

27 ocp_mx_moe_quant_config,

28 )

29 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput

30 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

31 convert_to_fp8_moe_kernel_format,

32 make_fp8_moe_kernel,

@@ -1444,7 +1445,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1445 x: torch.Tensor,

1446 router_logits: torch.Tensor,

1447 input_ids: torch.Tensor | None = None,

1447 - ) -> torch.Tensor:

1448 + ) -> torch.Tensor | UnfinalizedMoEOutput:

1449 assert self.is_monolithic

1450 assert self.moe_kernel is not None

1451 return self.moe_kernel.apply_monolithic(