From 73064103eed842cf9fb091db90c1e572b2130cdf Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Thu, 20 Aug 2026 18:22:50 +0000 Subject: [PATCH 1/3] Fuse MXFP4 top-k finalization into latent-tail Signed-off-by: yewentao256 --- .../model_executor/layers/fused_moe/config.py | 11 ++ .../fused_moe/experts/trtllm_bf16_moe.py | 2 +- .../fused_moe/experts/trtllm_mxfp4_moe.py | 40 ++-- .../fused_moe/experts/trtllm_nvfp4_moe.py | 2 +- .../kimi_k3/nvidia/latent_moe_runner.py | 41 +++- ...educe_rmsnorm_reduce_scatter_early_exit.py | 183 +++++++++++++++--- .../kimi_k3/nvidia/ops/latent_moe_tail.py | 43 ++-- 7 files changed, 269 insertions(+), 53 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 456d6c579364..61e6069c19e5 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1317,6 +1317,8 @@ class FusedMoEConfig: # consumer; read through `use_deferred_moe_finalize`, which applies the # guards. Kernels without the capability ignore it. Default False. defer_moe_finalize: bool = False + # optional consumer capacity for deferred finalize. Negative means unbounded. + defer_moe_finalize_max_num_tokens: int = -1 # SwiGLU clamp limit. When set, backends that do not implement the clamp # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle @@ -1447,6 +1449,15 @@ def use_deferred_moe_finalize(self) -> bool: # into, so the deferred form buys nothing. return self.defer_moe_finalize and self.tp_size > 1 + def should_defer_moe_finalize(self, num_tokens: int) -> bool: + """whether this invocation can leave the top-k reduction deferred""" + max_num_tokens = self.defer_moe_finalize_max_num_tokens + return ( + self.use_deferred_moe_finalize + and num_tokens > 0 + and (max_num_tokens < 0 or num_tokens <= max_num_tokens) + ) + @property def use_deepep_ht_kernels(self): return self.moe_parallel_config.use_deepep_ht_kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 4372634be8d1..7267ded1a230 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -277,7 +277,7 @@ def apply( num_tokens = hidden_states.shape[0] # The runner divides by the token count on the host, so an idle rank's # dummy 0-token forward has to keep the finalized (empty) form. - defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0 + defer = self.moe_config.should_defer_moe_finalize(num_tokens) routing_replay_out = self._maybe_make_routing_replay_buffer( num_tokens=num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 7154fce0b215..e895ab1f10f9 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 @@ -11,6 +11,9 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.moe_output import ( + UnfinalizedMoEOutput, +) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) @@ -215,7 +218,7 @@ def apply( e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: from flashinfer import trtllm_fp4_block_scale_moe if a1q_scale is not None: @@ -225,18 +228,22 @@ def apply( assert hidden_states.dtype == torch.bfloat16 x_quant = hidden_states x_scale = None - output = torch.empty( - *hidden_states.shape[:-1], - self.hidden_dim_unpadded, - dtype=torch.bfloat16, - device=hidden_states.device, - ) + num_tokens = hidden_states.shape[0] + defer = self.moe_config.should_defer_moe_finalize(num_tokens) + output = None + if not defer: + output = torch.empty( + *hidden_states.shape[:-1], + self.hidden_dim_unpadded, + dtype=torch.bfloat16, + device=hidden_states.device, + ) routing_replay_out = self._maybe_make_routing_replay_buffer( - num_tokens=hidden_states.shape[0], + num_tokens=num_tokens, device=hidden_states.device, ) - trtllm_fp4_block_scale_moe( + result = trtllm_fp4_block_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=x_quant, @@ -262,15 +269,22 @@ def apply( local_num_experts=self.local_num_experts, routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, - do_finalize=True, + do_finalize=not defer, activation_type=self._flashinfer_activation_type(activation), tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), output=output, routing_replay_out=routing_replay_out, ) - self._maybe_dispatch_routing_replay( - routing_replay_out, num_tokens=hidden_states.shape[0] - ) + self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens) + if defer: + return UnfinalizedMoEOutput( + gemm2_permuted=result[0], + expert_weights=result[1], + expanded_idx_to_permuted_idx=result[2] + .to(torch.int32) + .view(num_tokens, self.topk), + ) + assert output is not None return output 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 29ccd5770e00..d65ecbaff6ef 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 @@ -549,7 +549,7 @@ def apply( num_tokens = hidden_states.shape[0] # The runner divides by the token count on the host, so an idle rank's # dummy 0-token forward has to keep the finalized (empty) form. - defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0 + defer = self.moe_config.should_defer_moe_finalize(num_tokens) routing_replay_out = self._maybe_make_routing_replay_buffer( num_tokens=num_tokens, diff --git a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py index cc45f4d3ea2b..b45f2730027e 100644 --- a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py +++ b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py @@ -11,6 +11,7 @@ tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner, _unpack from vllm.model_executor.layers.layernorm import RMSNorm from vllm.platforms import current_platform @@ -103,6 +104,15 @@ def __init__( assert transform is not None norm = transform.norm assert norm is not None + from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( + TrtLlmMxfp4ExpertsMonolithic, + ) + + experts_cls = getattr(self._quant_method, "experts_cls", None) + self.moe_config.defer_moe_finalize = ( + experts_cls is TrtLlmMxfp4ExpertsMonolithic + and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded + ) from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import ( KimiK3LatentMoETailOp, ) @@ -110,11 +120,24 @@ def __init__( op = KimiK3LatentMoETailOp.initialize( hidden_size=transform.up_proj.weight.shape[0], latent_size=norm.weight.shape[0], + experts_per_token=( + self.moe_config.experts_per_token + if self.moe_config.use_deferred_moe_finalize + else 0 + ), dtype=norm.weight.dtype, device=norm.weight.device, rms_eps=norm.variance_epsilon, ) self._k3_latent_moe_tail_op = op + self.moe_config.defer_moe_finalize_max_num_tokens = ( + op.contract.max_num_tokens + if self.moe_config.use_deferred_moe_finalize + else -1 + ) + else: + self.moe_config.defer_moe_finalize = False + self.moe_config.defer_moe_finalize_max_num_tokens = -1 def _get_zero_residual( self, @@ -177,7 +200,7 @@ def _select_tail_tier( def _small_batch_tail( self, - fused_output: torch.Tensor, + fused_output: torch.Tensor | UnfinalizedMoEOutput, shared_output: torch.Tensor, trunc_size: int | None, ) -> torch.Tensor: @@ -304,6 +327,22 @@ def _fused_forward( ) ) + if ( + og_hidden_dim_pre_xform is None + and self.moe_config.should_defer_moe_finalize(hidden_states.shape[0]) + ): + unfinalized = self._forward_impl( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + ) + shared_output, fused_output = _unpack(unfinalized) + assert shared_output is not None + assert isinstance(fused_output, UnfinalizedMoEOutput) + result = self._small_batch_tail(fused_output, shared_output, None) + return self._maybe_add_zero_expert_output(result) + result = self._forward_entry( hidden_states, router_logits, diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py index 4f33f8920645..554d0e74f88a 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py @@ -16,6 +16,8 @@ import torch.distributed._symmetric_memory as symm_mem from cutlass import BFloat16, Float32, Int32, Int64, Uint32 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput + from .primitives import ( NUM_LAMPORT_BUFFERS, PACKED_BYTES, @@ -104,6 +106,7 @@ def __init__( fp32_internal: bool = False, include_reduce_scatter: bool = True, include_routed: bool = True, + top_k: int = 0, ): validate_shape( tp_size=tp_size, @@ -153,6 +156,7 @@ def __init__( self.fp32_internal = fp32_internal self.include_reduce_scatter = include_reduce_scatter self.include_routed = include_routed + self.top_k = top_k @cute.jit def __call__( @@ -171,6 +175,8 @@ def __call__( m: Int32, epsilon: Float32, stream: cuda.CUstream, + expert_weights: cute.Tensor, + expanded_idx_to_permuted_idx: cute.Tensor, ): self.kernel( latent_source, @@ -186,6 +192,8 @@ def __call__( shared_peer_ptrs, m, epsilon, + expert_weights, + expanded_idx_to_permuted_idx, ).launch( grid=(self.token_ctas, self.cluster_ctas, self.roles), block=(self.threads, 1, 1), @@ -211,6 +219,8 @@ def kernel( shared_peer_ptrs: cute.Tensor, m: Int32, epsilon: Float32, + expert_weights: cute.Tensor, + expanded_idx_to_permuted_idx: cute.Tensor, ): tidx, _, _ = cute.arch.thread_idx() token_cta, cta_y, role = cute.arch.block_idx() @@ -244,6 +254,8 @@ def kernel( logical_role, cluster_rank, tidx, + expert_weights, + expanded_idx_to_permuted_idx, ) token = token + self.token_ctas @@ -269,6 +281,8 @@ def _token_device( role: Int32, cluster_rank: Int32, tidx: Int32, + expert_weights: cute.Tensor, + expanded_idx_to_permuted_idx: cute.Tensor, ): if role == 0: # ---------------- routed AllReduce + RMSNorm ---------------- @@ -292,15 +306,48 @@ def _token_device( ) dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2)) - local_ptr = cute.make_ptr( - BFloat16, - (latent_source.iterator + element_offset).llvm_ptr, - cute.AddressSpace.gmem, - assumed_align=16, - ) - local_packed = sanitize_negative_zero( - load_global_u32x4(local_ptr, volatile=False) - ) + if cutlass.const_expr(self.top_k > 0): + # finalize topk reduction + local_values = cute.make_rmem_tensor( + cute.make_layout((VEC_BF16,)), BFloat16 + ) + for element in cutlass.range_constexpr(VEC_BF16): + local_values[element] = BFloat16(0.0) + for slot in cutlass.range_constexpr(self.top_k): + permuted_idx = expanded_idx_to_permuted_idx[token, slot] + if permuted_idx >= Int32(0): + permuted_element = ( + Int64(permuted_idx) * self.latent_dim + + Int64(packed_idx) * VEC_BF16 + ) + permuted_ptr = cute.make_ptr( + BFloat16, + (latent_source.iterator + permuted_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + values = packed_u32x4_to_bf16x8( + load_global_u32x4(permuted_ptr, volatile=False) + ) + weight = expert_weights[token, slot].to(Float32) + for element in cutlass.range_constexpr(VEC_BF16): + scaled = (values[element].to(Float32) * weight).to(BFloat16) + local_values[element] = ( + local_values[element].to(Float32) + scaled.to(Float32) + ).to(BFloat16) + local_packed = sanitize_negative_zero( + bf16x8_to_packed_u32x4(local_values.load()) + ) + else: + local_ptr = cute.make_ptr( + BFloat16, + (latent_source.iterator + element_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + local_packed = sanitize_negative_zero( + load_global_u32x4(local_ptr, volatile=False) + ) multicast_offset = ( Int64(current_index) * Int64(bytes_per_buffer) + ( @@ -645,6 +692,8 @@ def _compile_key( fp32_internal: bool, include_reduce_scatter: bool, include_routed: bool, + *, + top_k: int, ): return ( torch.accelerator.current_device_index(), @@ -657,6 +706,7 @@ def _compile_key( fp32_internal, include_reduce_scatter, include_routed, + top_k, ) @@ -673,6 +723,9 @@ def _runtime_args( shared_flags: torch.Tensor, shared_peer_ptrs: torch.Tensor, rms_eps: float, + *, + expert_weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, ): return ( to_cute_dynamic_m(latent_source, mode=0, assumed_align=16), @@ -686,9 +739,11 @@ def _runtime_args( to_cute(shared_workspace, 16), to_cute(shared_flags, 16), to_cute(shared_peer_ptrs, 16), - Int32(latent_source.shape[0]), + Int32(shared_source.shape[0]), Float32(rms_eps), cuda.CUstream(torch.cuda.current_stream(latent_source.device).cuda_stream), + to_cute_dynamic_m(expert_weights, mode=0, assumed_align=16), + to_cute_dynamic_m(expanded_idx_to_permuted_idx, mode=0, assumed_align=16), ) @@ -712,6 +767,7 @@ def compile_kernel( fp32_internal: bool, include_reduce_scatter: bool = True, include_routed: bool = True, + top_k: int = 0, ) -> None: """Compile the rank/M specialization without retaining caller tensors.""" @@ -725,11 +781,17 @@ def compile_kernel( fp32_internal, include_reduce_scatter, include_routed, + top_k=top_k, ) if key in _COMPILED: return device = latent_output.device - latent = torch.empty((max_m, latent_dim), dtype=torch.bfloat16, device=device) + latent_rows = max_m * max(top_k, 1) + latent = torch.empty((latent_rows, latent_dim), dtype=torch.bfloat16, device=device) + expert_weights = torch.empty( + (max_m, max(top_k, 1)), dtype=torch.bfloat16, device=device + ) + expanded_idx = torch.empty((max_m, max(top_k, 1)), dtype=torch.int32, device=device) gamma = torch.empty((latent_dim,), dtype=torch.bfloat16, device=device) shared = torch.empty((max_m, hidden_dim), dtype=torch.bfloat16, device=device) kernel = AllReduceRMSNormWithReduceScatterEarlyExit( @@ -742,6 +804,7 @@ def compile_kernel( fp32_internal=fp32_internal, include_reduce_scatter=include_reduce_scatter, include_routed=include_routed, + top_k=top_k, ) _COMPILED[key] = cute.compile( kernel, @@ -758,6 +821,8 @@ def compile_kernel( shared_flags, shared_peer_ptrs, rms_eps, + expert_weights=expert_weights, + expanded_idx_to_permuted_idx=expanded_idx, ), ) @@ -785,6 +850,9 @@ def launch( fp32_internal: bool, include_reduce_scatter: bool = True, include_routed: bool = True, + expert_weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + top_k: int = 0, ) -> None: compile_kernel( rank=rank, @@ -805,6 +873,7 @@ def launch( fp32_internal=fp32_internal, include_reduce_scatter=include_reduce_scatter, include_routed=include_routed, + top_k=top_k, ) _COMPILED[ _compile_key( @@ -817,6 +886,7 @@ def launch( fp32_internal, include_reduce_scatter, include_routed, + top_k=top_k, ) ]( *_runtime_args( @@ -832,6 +902,8 @@ def launch( shared_flags, shared_peer_ptrs, rms_eps, + expert_weights=expert_weights, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, ) ) @@ -851,6 +923,7 @@ def __init__( max_token_ctas: int, rms_eps: float, fp32_internal: bool, + top_k: int = 0, ) -> None: validate_shape( tp_size=tp_size, @@ -866,8 +939,16 @@ def __init__( self.max_token_ctas = max_token_ctas self.rms_eps = float(rms_eps) self.fp32_internal = fp32_internal + self.top_k = top_k device = torch.device("cuda", torch.accelerator.current_device_index()) + self._dummy_expert_weights = torch.empty( + (max_m, max(top_k, 1)), dtype=torch.bfloat16, device=device + ) + self._dummy_expanded_idx = torch.empty( + (max_m, max(top_k, 1)), dtype=torch.int32, device=device + ) + bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2 routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer self._routed_workspace = symm_mem.empty( @@ -947,38 +1028,91 @@ def __init__( shared_peer_ptrs=self._shared_peer_ptrs, rms_eps=self.rms_eps, fp32_internal=fp32_internal, + top_k=top_k, ) dist.barrier(group=group, device_ids=[device.index]) def __call__( self, - latent_source: torch.Tensor, + latent_source: torch.Tensor | UnfinalizedMoEOutput, shared_source: torch.Tensor, gamma: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - if latent_source.ndim != 2 or shared_source.ndim != 2: - raise ValueError("latent_source and shared_source must be rank-2") - m = latent_source.shape[0] + if shared_source.ndim != 2: + raise ValueError("shared_source must be rank-2") + expected: list[tuple[torch.Tensor, tuple[int, ...], str, torch.dtype]] + if isinstance(latent_source, UnfinalizedMoEOutput): + if self.top_k <= 0: + raise ValueError("collective was not configured for top-k finalize") + gemm2_permuted = latent_source.gemm2_permuted + expert_weights = latent_source.expert_weights + expanded_idx = latent_source.expanded_idx_to_permuted_idx + m = expanded_idx.shape[0] + expected = [ + ( + gemm2_permuted, + (gemm2_permuted.shape[0], self.latent_dim), + "gemm2_permuted", + torch.bfloat16, + ), + ( + expert_weights, + (m, self.top_k), + "expert_weights", + torch.bfloat16, + ), + ( + expanded_idx, + (m, self.top_k), + "expanded_idx_to_permuted_idx", + torch.int32, + ), + ] + else: + if self.top_k > 0: + raise ValueError("top-k collective requires an unfinalized output") + if latent_source.ndim != 2: + raise ValueError("latent_source must be rank-2") + m = latent_source.shape[0] + gemm2_permuted = latent_source + expert_weights = self._dummy_expert_weights + expanded_idx = self._dummy_expanded_idx + expected = [ + ( + latent_source, + (m, self.latent_dim), + "latent_source", + torch.bfloat16, + ), + ] device = self._routed_workspace.device - expected = ( - (latent_source, (m, self.latent_dim), "latent_source"), - (shared_source, (m, self.hidden_dim), "shared_source"), - (gamma, (self.latent_dim,), "gamma"), + expected.extend( + [ + ( + shared_source, + (m, self.hidden_dim), + "shared_source", + torch.bfloat16, + ), + (gamma, (self.latent_dim,), "gamma", torch.bfloat16), + ] ) - for tensor, shape, name in expected: + for tensor, shape, name, dtype in expected: if ( tensor.shape != shape - or tensor.dtype != torch.bfloat16 + or tensor.dtype != dtype or tensor.device != device or not tensor.is_contiguous() ): - raise ValueError(f"{name} must be contiguous CUDA BF16 {list(shape)}") + raise ValueError( + f"{name} must be contiguous CUDA {dtype} {list(shape)}" + ) if not 1 <= m <= self.max_m: raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") with torch.accelerator.device_index(device.index): launch( - latent_source, + gemm2_permuted, gamma, self._latent_output, self._routed_workspace, @@ -997,6 +1131,9 @@ def __call__( max_m=self.max_m, max_token_ctas=self.max_token_ctas, fp32_internal=self.fp32_internal, + top_k=self.top_k, + expert_weights=expert_weights, + expanded_idx_to_permuted_idx=expanded_idx, ) return ( self._latent_output[:m], diff --git a/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py b/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py index fe97c8fd3087..62c4325a4075 100644 --- a/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py +++ b/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py @@ -9,17 +9,18 @@ import torch.distributed as dist from vllm.distributed import get_tp_group +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.warmup.cutedsl_warmup import ( CuTeDSLCompileUnit, register_cutedsl_warmup_provider, ) -_MAX_NUM_TOKENS = 16 +_MAX_NUM_TOKENS = 128 _SKINNY_MAX_NUM_TOKENS = 5 _MMA_TILER_MN = (64, 32) _GEMM_CLUSTER_MN = (1, 8) _B_PRIME_STAGES = 2 -_COLLECTIVE_TOKEN_CTAS = 8 +_COLLECTIVE_TOKEN_CTAS = 32 _LAMPORT_COPY_CTAS = 32 _LAMPORT_COPY_THREADS = 224 _SUPPORTED_TP_SIZES = (8, 16) @@ -35,6 +36,7 @@ class KimiK3LatentMoETailContract: latent_size: int max_num_tokens: int rms_eps: float + experts_per_token: int class KimiK3LatentMoETailOp: @@ -53,6 +55,7 @@ def _contract_and_group( dtype: torch.dtype, device: torch.device, rms_eps: float, + experts_per_token: int = 0, ) -> tuple[KimiK3LatentMoETailContract, dist.ProcessGroup]: tp = get_tp_group() group = tp.device_group @@ -72,6 +75,7 @@ def _contract_and_group( latent_size=latent_size, max_num_tokens=_MAX_NUM_TOKENS, rms_eps=float(rms_eps), + experts_per_token=experts_per_token, ), group, ) @@ -85,10 +89,12 @@ def initialize( dtype: torch.dtype, device: torch.device, rms_eps: float, + experts_per_token: int = 0, ) -> "KimiK3LatentMoETailOp": contract, group = cls._contract_and_group( hidden_size=hidden_size, latent_size=latent_size, + experts_per_token=experts_per_token, dtype=dtype, device=device, rms_eps=rms_eps, @@ -140,6 +146,7 @@ def __init__( max_token_ctas=_COLLECTIVE_TOKEN_CTAS, rms_eps=contract.rms_eps, fp32_internal=False, + top_k=contract.experts_per_token, ) self._up_projection = AdaptiveUpProjectionKernel( group=group, @@ -191,7 +198,7 @@ def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: def __call__( self, - routed_output: torch.Tensor, + routed_output: torch.Tensor | UnfinalizedMoEOutput, shared_output: torch.Tensor, rms_weight: torch.Tensor, up_weight: torch.Tensor, @@ -202,7 +209,12 @@ def __call__( rms_weight, up_weight, ) - self._up_projection.ensure_compiled(routed_output.shape[0]) + num_tokens = ( + routed_output.expanded_idx_to_permuted_idx.shape[0] + if isinstance(routed_output, UnfinalizedMoEOutput) + else routed_output.shape[0] + ) + self._up_projection.ensure_compiled(num_tokens) latent, shared_shard = self._collective( routed_output, shared_output, @@ -221,24 +233,27 @@ def __call__( ) return self._lamport_copy( mailbox, - m=routed_output.shape[0], + m=num_tokens, ).squeeze(0) def _validate_inputs( self, - routed_output: torch.Tensor, + routed_output: torch.Tensor | UnfinalizedMoEOutput, shared_output: torch.Tensor, rms_weight: torch.Tensor, up_weight: torch.Tensor, ) -> None: contract = self.contract - if routed_output.ndim != 2: - raise ValueError("routed_output must be a 2D tensor.") - num_tokens = routed_output.shape[0] - if routed_output.shape != (num_tokens, contract.latent_size): - raise ValueError( - f"routed_output must have shape [M, {contract.latent_size}]." - ) + if isinstance(routed_output, UnfinalizedMoEOutput): + num_tokens = routed_output.expanded_idx_to_permuted_idx.shape[0] + else: + if routed_output.ndim != 2: + raise ValueError("routed_output must be a 2D tensor.") + num_tokens = routed_output.shape[0] + if routed_output.shape != (num_tokens, contract.latent_size): + raise ValueError( + f"routed_output must have shape [M, {contract.latent_size}]." + ) if shared_output.shape != (num_tokens, contract.hidden_size): raise ValueError( f"shared_output must have shape [M, {contract.hidden_size}]." @@ -256,7 +271,7 @@ def _validate_inputs( f"{contract.max_num_tokens} tokens." ) - tensors = (routed_output, shared_output, rms_weight, up_weight) + tensors = (shared_output, rms_weight, up_weight) if any(tensor.device != contract.device for tensor in tensors): raise ValueError("All inputs must be on the contract device.") if any(tensor.dtype != contract.dtype for tensor in tensors): From f5d46ef1290f39ce9cfbff637eea04a1e44bbe11 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 21 Aug 2026 06:42:20 +0000 Subject: [PATCH 2/3] Refine deferred latent MoE finalization Co-authored-by: OpenAI Codex Signed-off-by: Yongye Zhu --- .../benchmark_kimi_k3_latent_moe_tail.py | 106 ++++++++++-- tests/models/kimi_k3/test_latent_moe_tail.py | 152 ++++++++++++++++++ .../model_executor/layers/fused_moe/config.py | 17 +- .../fused_moe/experts/trtllm_bf16_moe.py | 23 ++- .../fused_moe/experts/trtllm_mxfp4_moe.py | 27 ++-- .../fused_moe/experts/trtllm_nvfp4_moe.py | 21 ++- .../layers/fused_moe/fused_moe_method_base.py | 7 +- .../layers/fused_moe/modular_kernel.py | 27 +++- .../layers/fused_moe/moe_output.py | 90 +++++++++++ .../fused_moe/prepare_finalize/no_dp_ep.py | 3 + .../layers/fused_moe/routed_experts.py | 7 +- .../layers/fused_moe/runner/moe_runner.py | 62 ++++--- .../layers/quantization/mxfp4.py | 5 +- .../kimi_k3/nvidia/latent_moe_runner.py | 15 +- 14 files changed, 463 insertions(+), 99 deletions(-) diff --git a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py index 3c1391919f88..2669dadb348d 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py +++ b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py @@ -21,6 +21,10 @@ torchrun --nproc-per-node=8 \ benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail + torchrun --nproc-per-node=8 \ + benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail \ + --backend fused --routed-input both + For multi-node runs, launch one ``torchrun`` agent per node and use a shared rendezvous endpoint. """ @@ -50,6 +54,7 @@ initialize_model_parallel, set_custom_all_reduce, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import ( @@ -165,6 +170,18 @@ def parse_args() -> argparse.Namespace: nargs="+", default=[1, 5, 8, 16], ) + whole_tail.add_argument( + "--routed-input", + choices=("finalized", "deferred", "both"), + default="finalized", + help="Select whether the fused tail performs the top-k finalize step.", + ) + whole_tail.add_argument( + "--top-k", + type=int, + default=8, + help="Expert count used by the synthetic deferred-finalize input.", + ) whole_tail.add_argument("--warmup-replays", type=int, default=20) whole_tail.add_argument("--samples", type=int, default=51) whole_tail.add_argument( @@ -586,6 +603,36 @@ def make_inputs( return routed, shared +def make_deferred_input( + routed: torch.Tensor, + top_k: int, +) -> UnfinalizedMoEOutput: + num_tokens = routed.shape[0] + gemm2_permuted = ( + routed[:, None, :] + .expand(num_tokens, top_k, LATENT_SIZE) + .contiguous() + .view(num_tokens * top_k, LATENT_SIZE) + ) + expert_weights = torch.zeros( + num_tokens, + top_k, + dtype=routed.dtype, + device=routed.device, + ) + expert_weights[:, 0] = 1 + expanded_idx = torch.arange( + num_tokens * top_k, + dtype=torch.int32, + device=routed.device, + ).view(num_tokens, top_k) + return UnfinalizedMoEOutput( + gemm2_permuted=gemm2_permuted, + expert_weights=expert_weights, + expanded_idx_to_permuted_idx=expanded_idx, + ) + + def make_reference( routed: torch.Tensor, shared: torch.Tensor, @@ -628,6 +675,8 @@ def benchmark_whole_tail(args: argparse.Namespace) -> None: raise ValueError("--num-tokens values must be in [1, 16]") if args.warmup_replays < 0 or args.samples <= 0: raise ValueError("warmup replays must be nonnegative and samples positive") + if args.top_k <= 0: + raise ValueError("--top-k must be positive") if args.skinny_max_num_tokens is not None and any( not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens ): @@ -672,6 +721,11 @@ def benchmark_whole_tail(args: argparse.Namespace) -> None: use_reference = args.backend in ("reference", "both") use_fused = args.backend in ("fused", "both") + routed_input_modes = ( + ("finalized", "deferred") + if args.routed_input == "both" + else (args.routed_input,) + ) fused_ops = [] if use_fused: production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m @@ -690,23 +744,33 @@ def config_for_m( for cutoff in cutoffs: latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff latent_moe_tail.KimiK3LatentMoETailOp._instances.clear() - fused_ops.append( - ( - cutoff, - latent_moe_tail.KimiK3LatentMoETailOp.initialize( - hidden_size=HIDDEN_SIZE, - latent_size=LATENT_SIZE, - dtype=torch.bfloat16, - device=device, - rms_eps=RMS_EPS, - ), + for routed_input_mode in routed_input_modes: + fused_ops.append( + ( + cutoff, + routed_input_mode, + latent_moe_tail.KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=RMS_EPS, + experts_per_token=( + args.top_k if routed_input_mode == "deferred" else 0 + ), + ), + ) ) - ) cutedsl_warmup() results = [] for num_tokens in args.num_tokens: routed, shared = make_inputs(num_tokens, rank, device) + routed_inputs: dict[str, torch.Tensor | UnfinalizedMoEOutput] = { + "finalized": routed + } + if "deferred" in routed_input_modes: + routed_inputs["deferred"] = make_deferred_input(routed, args.top_k) reference = make_reference( routed, shared, @@ -724,17 +788,23 @@ def config_for_m( device_group=device_group, cpu_group=cpu_group, ) - for cutoff, fused_op in fused_ops: + for cutoff, routed_input_mode, fused_op in fused_ops: + routed_input = routed_inputs[routed_input_mode] def fused( - routed: torch.Tensor = routed, + routed_input: torch.Tensor | UnfinalizedMoEOutput = routed_input, shared: torch.Tensor = shared, fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op, ) -> torch.Tensor: - return fused_op(routed, shared, rms_weight, up_weight) + return fused_op(routed_input, shared, rms_weight, up_weight) fused_graph, fused_output = capture_tail_graph(fused, cpu_group) - fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}" + fused_key_parts = ["fused"] + if routed_input_mode != "finalized" or len(routed_input_modes) > 1: + fused_key_parts.append(routed_input_mode) + if len(cutoffs) > 1: + fused_key_parts.extend(("skinny", "max", str(cutoff))) + fused_key = "_".join(fused_key_parts) result[fused_key] = benchmark_tail_graph( fused_graph, warmup_replays=args.warmup_replays, @@ -747,7 +817,7 @@ def fused( speedup = ( result["reference"]["median_us"] / result[fused_key]["median_us"] ) - if len(fused_ops) == 1: + if fused_key == "fused": result["speedup"] = speedup else: result[f"{fused_key}_speedup"] = speedup @@ -763,7 +833,9 @@ def fused( "cuda_version": torch.version.cuda, "warmup_replays": args.warmup_replays, "samples": args.samples, - "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops], + "routed_input": args.routed_input, + "top_k": args.top_k, + "skinny_max_num_tokens": sorted({cutoff for cutoff, _, _ in fused_ops}), "skinny_configs": { str(num_tokens): asdict(config) for num_tokens, config in skinny_configs.items() diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py index f18f80a36c04..f5447edbf54b 100644 --- a/tests/models/kimi_k3/test_latent_moe_tail.py +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -13,6 +13,7 @@ multi_process_parallel, ) from vllm.distributed import get_tp_group +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import KimiK3LatentMoETailOp from vllm.platforms import current_platform @@ -20,6 +21,57 @@ HIDDEN_SIZE = 7168 LATENT_SIZE = 3584 EPS = 0.1 +TOP_K = 8 + + +def _make_deferred_routed_output( + num_tokens: int, + device: torch.device, +) -> tuple[UnfinalizedMoEOutput, torch.Tensor]: + num_routes = num_tokens * TOP_K + num_permuted_rows = num_routes + 7 + expanded_output = torch.randn( + num_routes, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ).mul_(0.01) + expert_weights = torch.rand( + num_tokens, + TOP_K, + device=device, + dtype=torch.bfloat16, + ) + expert_weights.div_(expert_weights.sum(dim=-1, keepdim=True)) + expanded_idx = torch.randperm(num_permuted_rows, device=device)[:num_routes] + gemm2_permuted = torch.empty( + num_permuted_rows, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + gemm2_permuted[expanded_idx] = expanded_output + + finalized = torch.zeros( + num_tokens, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + expanded_output = expanded_output.view(num_tokens, TOP_K, LATENT_SIZE) + for slot in range(TOP_K): + finalized.add_(expanded_output[:, slot] * expert_weights[:, slot, None]) + + return ( + UnfinalizedMoEOutput( + gemm2_permuted=gemm2_permuted, + expert_weights=expert_weights, + expanded_idx_to_permuted_idx=expanded_idx.to(torch.int32).view( + num_tokens, TOP_K + ), + ), + finalized, + ) @ray.remote(num_gpus=1, max_calls=1) @@ -131,6 +183,99 @@ def _run_latent_moe_tail_test( ) +@ray.remote(num_gpus=1, max_calls=1) +def _test_deferred_finalize_parity_worker( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, + pp_size: int, + rank: int, + distributed_init_port: str, +) -> None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment( + tp_size, + pp_size, + rank, + distributed_init_port, + ) + + torch.manual_seed(1000 + rank) + rms_weight = 1 + 0.1 * torch.randn( + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + up_weight = ( + torch.randn( + HIDDEN_SIZE, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + / LATENT_SIZE**0.5 + ) + finalized_op = KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=EPS, + ) + deferred_op = KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=EPS, + experts_per_token=TOP_K, + ) + cutedsl_warmup() + + for iteration, num_tokens in enumerate((1, 5, 16)): + torch.manual_seed(2000 + 100 * iteration + rank) + deferred_output, finalized_output = _make_deferred_routed_output( + num_tokens, + device, + ) + shared_output = torch.randn( + num_tokens, + HIDDEN_SIZE, + device=device, + dtype=torch.bfloat16, + ) + + expected = finalized_op( + finalized_output, + shared_output, + rms_weight, + up_weight, + ) + actual = deferred_op( + deferred_output, + shared_output, + rms_weight, + up_weight, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +def _run_deferred_finalize_parity_test( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, +) -> None: + if not current_platform.is_device_capability_family(100): + pytest.skip("K3 latent-MoE tail fusion requires SM100") + multi_process_parallel( + monkeypatch, + tp_size, + 1, + _test_deferred_finalize_parity_worker, + ) + + @multi_gpu_test(num_gpus=8) def test_latent_moe_tail_tp8_matches_native_path( monkeypatch: pytest.MonkeyPatch, @@ -143,3 +288,10 @@ def test_latent_moe_tail_tp16_matches_native_path( monkeypatch: pytest.MonkeyPatch, ) -> None: _run_latent_moe_tail_test(monkeypatch, 16) + + +@multi_gpu_test(num_gpus=8) +def test_latent_moe_tail_deferred_finalize_matches_finalized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_deferred_finalize_parity_test(monkeypatch, 8) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 61e6069c19e5..0f90ac756030 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1317,7 +1317,7 @@ class FusedMoEConfig: # consumer; read through `use_deferred_moe_finalize`, which applies the # guards. Kernels without the capability ignore it. Default False. defer_moe_finalize: bool = False - # optional consumer capacity for deferred finalize. Negative means unbounded. + # Optional consumer capacity for deferred finalize. Negative means unbounded. defer_moe_finalize_max_num_tokens: int = -1 # SwiGLU clamp limit. When set, backends that do not implement the clamp @@ -1445,12 +1445,19 @@ def use_deferred_moe_finalize(self) -> bool: ``defer_moe_finalize`` is set after construction, like ``skip_final_all_reduce``. """ - # Without TP there is no all-reduce for the top-k reduction to fuse - # into, so the deferred form buys nothing. - return self.defer_moe_finalize and self.tp_size > 1 + # The consumer fuses a TP all-reduce. Other parallel modes require a + # combine or reduce-scatter after the experts and cannot defer it. + return ( + self.defer_moe_finalize + and self.tp_size > 1 + and self.dp_size == 1 + and self.ep_size == 1 + and self.pcp_size == 1 + and not self.is_sequence_parallel + ) def should_defer_moe_finalize(self, num_tokens: int) -> bool: - """whether this invocation can leave the top-k reduction deferred""" + """Return whether this invocation may defer the top-k reduction.""" max_num_tokens = self.defer_moe_finalize_max_num_tokens return ( self.use_deferred_moe_finalize diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 7267ded1a230..6b29a6f69d1f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -13,6 +13,7 @@ ) 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, @@ -283,7 +284,7 @@ def apply( num_tokens=num_tokens, device=hidden_states.device, ) - out = flashinfer.fused_moe.trtllm_bf16_moe( + flashinfer_output = flashinfer.fused_moe.trtllm_bf16_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, @@ -303,17 +304,11 @@ def apply( routing_replay_out=routing_replay_out, do_finalize=not defer, ) + routed_output = convert_flashinfer_moe_output( + flashinfer_output, + do_finalize=not defer, + num_tokens=num_tokens, + top_k=self.topk, + ) self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens) - if defer: - # flashinfer returns a flat permute map; the protocol wants - # [num_tokens, top_k] so consumers can read top_k from its shape. - return UnfinalizedMoEOutput( - gemm2_permuted=out[0], - expert_weights=out[1], - expanded_idx_to_permuted_idx=out[2] - .to(torch.int32) - .view(num_tokens, self.topk), - ) - # do_finalize=True yields the finalized states (a bare tensor on some - # FlashInfer versions, a single-element list on others). - return out[0] if isinstance(out, (list, tuple)) else out + return routed_output 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 e895ab1f10f9..e28e401e4c63 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 @@ -13,6 +13,7 @@ ) 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, @@ -230,9 +231,9 @@ def apply( x_scale = None num_tokens = hidden_states.shape[0] defer = self.moe_config.should_defer_moe_finalize(num_tokens) - output = None + finalized_output = None if not defer: - output = torch.empty( + finalized_output = torch.empty( *hidden_states.shape[:-1], self.hidden_dim_unpadded, dtype=torch.bfloat16, @@ -243,7 +244,7 @@ def apply( num_tokens=num_tokens, device=hidden_states.device, ) - result = trtllm_fp4_block_scale_moe( + flashinfer_output = trtllm_fp4_block_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=x_quant, @@ -272,20 +273,18 @@ def apply( do_finalize=not defer, activation_type=self._flashinfer_activation_type(activation), tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), - output=output, + output=finalized_output, routing_replay_out=routing_replay_out, ) + routed_output = convert_flashinfer_moe_output( + flashinfer_output, + do_finalize=not defer, + num_tokens=num_tokens, + top_k=self.topk, + finalized_output=finalized_output, + ) self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens) - if defer: - return UnfinalizedMoEOutput( - gemm2_permuted=result[0], - expert_weights=result[1], - expanded_idx_to_permuted_idx=result[2] - .to(torch.int32) - .view(num_tokens, self.topk), - ) - assert output is not None - return output + return routed_output class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModular): 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 d65ecbaff6ef..c4bbca266894 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 @@ -15,6 +15,7 @@ ) 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, @@ -558,7 +559,7 @@ def apply( # Invoke kernel. # NOTE: Activation padding and output # truncation are handled by the MoE runner's - result = flashinfer.fused_moe.trtllm_fp4_block_scale_moe( + flashinfer_output = flashinfer.fused_moe.trtllm_fp4_block_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, @@ -592,15 +593,11 @@ def apply( tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), routing_replay_out=routing_replay_out, ) + routed_output = convert_flashinfer_moe_output( + flashinfer_output, + do_finalize=not defer, + num_tokens=num_tokens, + top_k=self.topk, + ) self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens) - if defer: - # flashinfer returns a flat permute map; the protocol wants - # [num_tokens, top_k] so consumers can read top_k from its shape. - return UnfinalizedMoEOutput( - gemm2_permuted=result[0], - expert_weights=result[1], - expanded_idx_to_permuted_idx=result[2] - .to(torch.int32) - .view(num_tokens, self.topk), - ) - return result[0] + return routed_output 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 4d68a4ab64d2..3a7c6eedfba7 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 @@ -13,6 +13,7 @@ FusedMoEParallelConfig, FusedMoEQuantConfig, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.quantization.base_config import ( QuantizeMethodBase, ) @@ -161,7 +162,7 @@ def apply( shared_experts_input: Input for shared experts (if any) Returns: - Output tensor from routed experts + Output tensor from routed experts. """ raise NotImplementedError @@ -171,7 +172,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: """ Apply the MoE operation using monolithic kernels. @@ -181,6 +182,6 @@ def apply_monolithic( router_logits: Router logits (routing done internally) 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 04734f7fbe54..c159f059d61d 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -22,6 +22,7 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, SharedExpertsOrder, @@ -425,6 +426,9 @@ class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize): described above for the monolithic case. """ + def supports_deferred_moe_finalize(self) -> bool: + return False + @abstractmethod def prepare( self, @@ -1078,7 +1082,7 @@ def apply( e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: """ Same as apply(), except uses router_logits as opposed to the topk_ids and topk_weights. This is useful for kernels @@ -1548,7 +1552,7 @@ def apply( e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: """ Same as forward(), except uses router_logits as opposed to the topk_ids and topk_weights. This is used for kernels @@ -1579,9 +1583,14 @@ def apply( topk_group=topk_group, ) - output = self.prepare_finalize.finalize(fused_out) - - return output + if isinstance(fused_out, UnfinalizedMoEOutput): + 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.prepare_finalize.finalize(fused_out) @final @@ -1664,6 +1673,12 @@ 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() + ) + def apply_monolithic( self, hidden_states: torch.Tensor, @@ -1679,7 +1694,7 @@ def apply_monolithic( e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert isinstance(self.impl, FusedMoEKernelMonolithicImpl) 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 d3fc1a75fb55..df1a2685ecbe 100644 --- a/vllm/model_executor/layers/fused_moe/moe_output.py +++ b/vllm/model_executor/layers/fused_moe/moe_output.py @@ -33,6 +33,96 @@ class UnfinalizedMoEOutput: expanded_idx_to_permuted_idx: torch.Tensor +def convert_flashinfer_moe_output( + flashinfer_output: torch.Tensor | list[torch.Tensor], + *, + do_finalize: bool, + num_tokens: int, + top_k: int, + finalized_output: torch.Tensor | None = None, +) -> torch.Tensor | UnfinalizedMoEOutput: + """Normalize the two FlashInfer TRTLLM MoE return layouts. + + Args: + flashinfer_output: Tensor returned by the FlashInfer BF16 wrapper's + legacy finalized path, or its mode-dependent tensor list. + do_finalize: Whether FlashInfer ran its top-k finalize step. + num_tokens: Number of input tokens. + top_k: Number of routed experts per token. + finalized_output: Optional destination passed to FlashInfer's ``output`` + argument. + + Returns: + A finalized tensor or the structured deferred-finalize output. + + Raises: + ValueError: If FlashInfer returns an unexpected layout. + """ + + if num_tokens < 0: + raise ValueError("num_tokens must be non-negative.") + if top_k <= 0: + raise ValueError("top_k must be positive.") + + if do_finalize: + if isinstance(flashinfer_output, torch.Tensor): + returned_output = flashinfer_output + elif len(flashinfer_output) == 1 and isinstance( + flashinfer_output[0], torch.Tensor + ): + returned_output = flashinfer_output[0] + else: + raise ValueError( + "Finalized FlashInfer MoE output must contain exactly one tensor." + ) + if finalized_output is not None and ( + returned_output.shape != finalized_output.shape + or returned_output.dtype != finalized_output.dtype + or returned_output.device != finalized_output.device + or returned_output.data_ptr() != finalized_output.data_ptr() + ): + raise ValueError( + "Finalized FlashInfer MoE output must alias the provided destination." + ) + return returned_output + + if finalized_output is not None: + raise ValueError("Deferred FlashInfer MoE output cannot have a destination.") + if isinstance(flashinfer_output, torch.Tensor) or len(flashinfer_output) != 3: + raise ValueError( + "Deferred FlashInfer MoE output must contain GEMM2 output, routing " + "weights, and a permutation map." + ) + gemm2_permuted, expert_weights, expanded_idx = flashinfer_output + if gemm2_permuted.ndim != 2: + raise ValueError("Deferred FlashInfer GEMM2 output must be rank-2.") + if expanded_idx.dtype != torch.int32: + raise ValueError("Deferred FlashInfer permutation map must use int32.") + expected_routes = num_tokens * top_k + if expert_weights.numel() != expected_routes: + raise ValueError( + "Deferred FlashInfer routing weights must contain " + f"num_tokens * top_k ({expected_routes}) elements." + ) + if expanded_idx.numel() != expected_routes: + raise ValueError( + "Deferred FlashInfer permutation map must contain " + f"num_tokens * top_k ({expected_routes}) elements." + ) + if not (gemm2_permuted.device == expert_weights.device == expanded_idx.device): + raise ValueError("Deferred FlashInfer outputs must share a device.") + if not all( + output.is_contiguous() + for output in (gemm2_permuted, expert_weights, expanded_idx) + ): + raise ValueError("Deferred FlashInfer outputs must be contiguous.") + return UnfinalizedMoEOutput( + gemm2_permuted=gemm2_permuted, + expert_weights=expert_weights.view(num_tokens, top_k), + expanded_idx_to_permuted_idx=expanded_idx.view(num_tokens, top_k), + ) + + @dataclass class MoEOutput: """A MoE layer's output with its final reduction still open. 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 69587770c40c..d13bc20673ec 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 @@ -98,6 +98,9 @@ def finalize( class MoEPrepareAndFinalizeNoDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic): + def supports_deferred_moe_finalize(self) -> bool: + 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 b5557f4b59e9..ce37b668ced1 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -19,6 +19,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.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) @@ -1216,7 +1217,7 @@ def forward_modular( shared_experts_input: Input for shared experts (if any) Returns: - Output tensor from routed experts + Output tensor from routed experts. """ assert not self.quant_method.is_monolithic @@ -1235,7 +1236,7 @@ def forward_monolithic( x: torch.Tensor, router_logits: torch.Tensor | None = None, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: """ Execute routed experts using the quantization method's apply function. @@ -1250,7 +1251,7 @@ def forward_monolithic( input_ids: input ids for DeepSeek V4 Returns: - Output tensor from routed experts + Finalized routed states or a deferred-finalize output. """ assert self.quant_method.is_monolithic diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index aaf9ba61423a..55d50e44ab82 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable, Iterable from contextlib import nullcontext -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import torch import torch.nn.functional as F @@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.routed_experts import ( RoutedExperts, ) @@ -121,11 +122,14 @@ def _moe_forward( hidden_dim_unpadded: int, ) -> torch.Tensor: layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer._forward_impl( - hidden_states, - router_logits, - shared_experts_input, - input_ids, + return cast( + torch.Tensor, + layer._forward_impl( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + ), ) @@ -155,11 +159,14 @@ def _moe_forward_shared( hidden_dim_unpadded: int, ) -> tuple[torch.Tensor, torch.Tensor]: layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer._forward_impl( - hidden_states, - router_logits, - shared_experts_input, - input_ids, + return cast( + tuple[torch.Tensor, torch.Tensor], + layer._forward_impl( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + ), ) @@ -207,8 +214,10 @@ def _moe_forward_shared_fake( def _unpack( - result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], -) -> tuple[torch.Tensor | None, torch.Tensor]: + result: torch.Tensor + | UnfinalizedMoEOutput + | tuple[torch.Tensor, torch.Tensor | UnfinalizedMoEOutput], +) -> tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]: if isinstance(result, tuple): return result else: @@ -573,7 +582,7 @@ def _apply_quant_method( router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, input_ids: torch.Tensor | None = None, - ) -> tuple[torch.Tensor | None, torch.Tensor]: + ) -> tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]: """Run expert routing and the fused MoE kernel via the quant method. Orchestrates shared expert execution (before/after), expert selection @@ -733,6 +742,7 @@ def forward( # Extract outputs from result shared_output, fused_output = _unpack(result) + fused_output = cast(torch.Tensor, fused_output) if og_hidden_dim_pre_xform is not None: fused_output = fused_output[..., :og_hidden_dim_pre_xform] @@ -808,18 +818,25 @@ def _maybe_dispatch( def _maybe_combine( self, shared_output: torch.Tensor | None, - hidden_states: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor | None, torch.Tensor]: + hidden_states: torch.Tensor | UnfinalizedMoEOutput, + ) -> ( + torch.Tensor + | UnfinalizedMoEOutput + | tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput] + ): if self.do_naive_dispatch_combine: hidden_states = get_ep_group().combine( - hidden_states, self.moe_config.is_sequence_parallel + cast(torch.Tensor, hidden_states), + self.moe_config.is_sequence_parallel, ) if ( self.moe_config.pcp_size > 1 and not self.moe_config.moe_parallel_config.use_all2all_kernels ): - hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0) + hidden_states = get_pcp_group().reduce_scatter( + cast(torch.Tensor, hidden_states), dim=0 + ) if self.shared_experts is not None: assert shared_output is not None @@ -833,7 +850,11 @@ def _forward_impl( router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> ( + torch.Tensor + | UnfinalizedMoEOutput + | tuple[torch.Tensor, torch.Tensor | UnfinalizedMoEOutput] + ): """Entry point called by the custom op to run the MoE computation. Handles pre-dispatch setup (gate application, external shared expert @@ -844,7 +865,8 @@ def _forward_impl( - fused MoE kernel execution - shared expert computation. - Returns a single tensor of combined fused and shared output (if present). + Returns routed output, optionally paired with shared-expert output. A + fused consumer may request the routed output in deferred-finalize form. """ # TODO(bnell): this can be removed after MK migration is complete. self.routed_experts._ensure_moe_quant_config_init() diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 6506c1fce1df..9d6e85e294e7 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.config import ( mxfp4_w4a16_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, @@ -450,7 +451,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -863,7 +864,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( diff --git a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py index b45f2730027e..905b8b16948f 100644 --- a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py +++ b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import IntEnum +from typing import cast import torch @@ -107,11 +108,20 @@ def __init__( from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( TrtLlmMxfp4ExpertsMonolithic, ) + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( + TrtLlmNvFp4ExpertsMonolithic, + ) experts_cls = getattr(self._quant_method, "experts_cls", None) + moe_kernel = self._quant_method.moe_kernel self.moe_config.defer_moe_finalize = ( - experts_cls is TrtLlmMxfp4ExpertsMonolithic + ( + experts_cls is TrtLlmMxfp4ExpertsMonolithic + or experts_cls is TrtLlmNvFp4ExpertsMonolithic + ) and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded + and moe_kernel is not None + and moe_kernel.supports_deferred_moe_finalize() ) from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import ( KimiK3LatentMoETailOp, @@ -354,8 +364,7 @@ def _fused_forward( else 0, ) - shared_output, fused_output = _unpack(result) - assert shared_output is not None + shared_output, fused_output = cast(tuple[torch.Tensor, torch.Tensor], result) if og_hidden_dim_pre_xform is not None: fused_output = fused_output[..., :og_hidden_dim_pre_xform] From ddac6d2053df65c4bb13776615a7427de797ebc9 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 21 Aug 2026 07:08:28 +0000 Subject: [PATCH 3/3] Fix AMD latent MoE output typing Co-authored-by: OpenAI Codex Signed-off-by: Yongye Zhu --- vllm/models/kimi_k3/amd/latent_moe_runner.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py index db1c078b5d80..83c953e1a1d4 100644 --- a/vllm/models/kimi_k3/amd/latent_moe_runner.py +++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import cast + import torch from vllm.distributed import ( @@ -7,7 +9,7 @@ tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner, _unpack +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner logger = init_logger(__name__) @@ -139,8 +141,7 @@ def _fused_forward( else 0, ) - shared_output, fused_output = _unpack(result) - assert shared_output is not None + shared_output, fused_output = cast(tuple[torch.Tensor, torch.Tensor], result) if og_hidden_dim_pre_xform is not None: fused_output = fused_output[..., :og_hidden_dim_pre_xform]