From 62fbfcad3d39f4a7242af8ad05e863daf2399bd6 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 18 Jan 2026 20:35:31 -0500 Subject: [PATCH 001/207] imports Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 213 +++++++++++++++++- .../layers/fused_moe/modular_kernel.py | 23 ++ 2 files changed, 235 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 3bb5a23abb7b..aa0364ce392d 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -3,7 +3,15 @@ import torch -from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( calculate_tile_tokens_dim, @@ -190,3 +198,206 @@ def fi_trtllm_fp8_per_tensor_moe_fake( fake_impl=fi_trtllm_fp8_per_tensor_moe_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) + + +class FlashInferTrtLlmNvFp4Experts(mk.FusedMoEPermuteExpertsUnpermute): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(quant_config) + + import flashinfer + + # TODO: set this via the constructor + self.routing_method_type = flashinfer.RoutingMethodType.Renormalize + # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 + # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 + + self.routing_bias = None + self.e_score_correction_bias = None + self.topk_group = None + self.num_expert_group = None + + self.topk = moe_config.experts_per_token + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.hidden_dim = moe_config.hidden_dim + self.local_num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + # a13_scale * w13_scale_2 / a2_scale + self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale + + @property + def activation_formats( + self, + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]: + return ( + mk.FusedMoEActivationFormat.Standard, + mk.FusedMoEActivationFormat.Standard, + ) + + def supports_chunking(self) -> bool: + return False + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: str, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # The workspaces for this implementation are managed by flashinfer. + workspace1 = (0,) + workspace2 = (0,) + output = (M, K) + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + import flashinfer + + assert activation == "silu" + assert a1q_scale is not None + + # Pack topk_ids and topk_weights into format expected by the kernel. + packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to( + torch.bfloat16 + ).view(torch.int16) + + # Invoke kernel. + # TODO(avoid the copy). + out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_tensor, + routing_bias=None, + hidden_states=hidden_states, + hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c, + output1_scale_gate_scalar=self.quant_config.g1_alphas, + output2_scale_scalar=self.quant_config.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + n_group=0, + topk_group=0, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=None, + tile_tokens_dim=None, + routing_method_type=1, + do_finalize=True, + )[0] + + output.copy_(out) + + def apply_monolthic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + import flashinfer + + assert activation == "silu" + + # Quantize input to FP4 + if isinstance(hidden_states, tuple): + a1q, a1q_scale = hidden_states + else: + a1q, a1q_scale = flashinfer.fp4_quantize( + hidden_states, + self.quant_config.a1_gscale, + is_sf_swizzled_layout=False, + ) + + # Prepare routing bias + routing_bias = self.e_score_correction_bias + if routing_bias is not None: + routing_bias = routing_bias.to(torch.bfloat16) + + router_logits = ( + router_logits.to(torch.float32) + if self.routing_method_type == RoutingMethodType.DeepSeekV3 + else router_logits + ) + + # Call TRT-LLM FP4 block-scale MoE kernel + out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe( + routing_logits=router_logits, + routing_bias=routing_bias, + hidden_states=a1q, + hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c, + output1_scale_gate_scalar=self.quant_config.g1_alphas, + output2_scale_scalar=self.quant_config.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + n_group=self.num_expert_group if self.num_expert_group is not None else 0, + topk_group=self.topk_group if self.topk_group is not None else 0, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=None, + tile_tokens_dim=None, + routing_method_type=self.routing_method_type, + do_finalize=True, + )[0] + + return out diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 962d0fe78fbc..13559080e820 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -664,6 +664,29 @@ def apply( """ raise NotImplementedError + def apply_monolthic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + routing_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + """ + Same as apply(), except uses routing_logits as opposed + to the topk_ids and topk_weights. This is useful for kernels + with fused router and fused_experts (e.g. FLASHINFER_TRTLLM). + """ + raise NotImplementedError + def _slice_scales( scales: torch.Tensor | None, start: int, end: int From db2f014bf9330e66d3497e9498bb90a49e54b0a8 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 18 Jan 2026 20:41:57 -0500 Subject: [PATCH 002/207] updated Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/config.py | 1 + .../layers/fused_moe/flashinfer_trtllm_moe.py | 6 ++---- vllm/model_executor/layers/fused_moe/layer.py | 1 + .../layers/fused_moe/oracle/nvfp4.py | 14 ++++++++++++-- 4 files changed, 16 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index c8baefbd55fe..10693119406a 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1026,6 +1026,7 @@ class FusedMoEConfig: num_experts: int experts_per_token: int hidden_dim: int + intermediate_dim: int num_local_experts: int moe_parallel_config: FusedMoEParallelConfig diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index aa0364ce392d..6f7a6ba850b4 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -221,9 +221,7 @@ def __init__( self.num_expert_group = None self.topk = moe_config.experts_per_token - self.intermediate_size_per_partition = ( - moe_config.intermediate_size_per_partition - ) + self.intermediate_dim = moe_config.intermediate_dim self.hidden_dim = moe_config.hidden_dim self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank @@ -317,7 +315,7 @@ def apply( top_k=self.topk, n_group=0, topk_group=0, - intermediate_size=self.intermediate_size_per_partition, + intermediate_size=self.intermediate_dim, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, routed_scaling_factor=None, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index e24d60150d60..b12baddc1d9f 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -524,6 +524,7 @@ def __init__( num_experts=self.global_num_experts, experts_per_token=top_k, hidden_dim=hidden_size, + intermediate_dim=self.intermediate_size_per_partition, num_local_experts=self.local_num_experts, moe_parallel_config=self.moe_parallel_config, in_dtype=moe_in_dtype, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index f2d69cf09c46..1918a880b09f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -20,6 +20,9 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, ) +from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import ( + FlashInferTrtLlmNvFp4Experts, +) from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( MarlinExperts, ) @@ -246,8 +249,6 @@ def make_nvfp4_moe_kernel( assert moe_config.dp_size == 1 UNSUPPORTED_BACKENDS = [ - # TRTLLM does not use the modular kernl abstraction. - NvFp4MoeBackend.FLASHINFER_TRTLLM, # CUTEDSL is used with BATCHED (masked) format only. # TODO: add here once we support dp/ep via the oracle. NvFp4MoeBackend.FLASHINFER_CUTEDSL, @@ -256,6 +257,15 @@ def make_nvfp4_moe_kernel( if backend in UNSUPPORTED_BACKENDS: return None + elif backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: + return mk.FusedMoEModularKernel( + MoEPrepareAndFinalizeNoEP(defer_input_quant=False), + FlashInferTrtLlmNvFp4Experts( + moe_config=moe_config, + quant_config=quant_config, + ), + ) + elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS: return mk.FusedMoEModularKernel( MoEPrepareAndFinalizeNoEP(defer_input_quant=True), From ec9c645205f106590c3906ac935bc29e6285b00b Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 09:25:42 -0500 Subject: [PATCH 003/207] stash changes for remote review Signed-off-by: Robert Shaw --- .../layers/fused_moe/modular_kernel.py | 1 + .../layers/fused_moe/oracle/nvfp4.py | 6 +- .../layers/fused_moe/prepare_finalize.py | 4 +- .../layers/quantization/modelopt.py | 65 ++++--------------- 4 files changed, 19 insertions(+), 57 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 13559080e820..82107581469d 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -867,6 +867,7 @@ def _allocate_buffers( expert_tokens_meta, activation, ) + print(f"{fused_out_shape=}") # We can reuse the memory between cache1 and cache3 because by the # time we need cache3, we're done with cache1. diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 1918a880b09f..ec3b37336b96 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -217,11 +217,7 @@ def make_nvfp4_moe_quant_config( a13_scale: torch.Tensor, a2_scale: torch.Tensor, ) -> FusedMoEQuantConfig | None: - UNSUPPORTED = [NvFp4MoeBackend.FLASHINFER_TRTLLM] - if backend in UNSUPPORTED: - return None - - elif backend == NvFp4MoeBackend.MARLIN: + if backend == NvFp4MoeBackend.MARLIN: return nvfp4_w4a16_moe_quant_config( g1_alphas=w13_scale_2, g2_alphas=w2_scale_2, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 5d806fa843a3..5a001b769032 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -59,10 +59,12 @@ def prepare( a1q, a1q_scale = moe_kernel_quantize_input( a1, - quant_config.a1_scale, + # quant_config.a1_scale, + quant_config.a1_gscale, quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, + is_fp4_scale_swizzled=False, # ? ) return a1q, a1q_scale, None, None, None diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 4c9fac39ca7e..cc5072c84e77 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -52,8 +52,6 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( build_flashinfer_fp4_cutlass_moe_prepare_finalize, - flashinfer_trtllm_fp4_moe, - flashinfer_trtllm_fp4_routed_moe, select_nvfp4_gemm_impl, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( @@ -1607,59 +1605,24 @@ def apply( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if ( - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - and not layer.enable_eplb - ): - return flashinfer_trtllm_fp4_moe( - layer=layer, - x=x, - router_logits=router_logits, - top_k=layer.top_k, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - custom_routing_function=layer.custom_routing_function, - e_score_correction_bias=layer.e_score_correction_bias, - ) - - # Hidden_states in select_experts is only used to extract metadata - if isinstance(x, tuple): - x_routing, _ = x - else: - x_routing = x topk_weights, topk_ids = router.select_experts( - hidden_states=x_routing, + hidden_states=x, router_logits=router_logits, ) - # EPLB path - if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - assert layer.enable_eplb - return flashinfer_trtllm_fp4_routed_moe( - layer=layer, - x=x, - topk_ids=topk_ids, - topk_weights=topk_weights, - top_k=layer.top_k, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - ) - else: - assert self.kernel is not None - return self.kernel( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - inplace=False, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - ) + assert self.kernel is not None + return self.kernel( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + inplace=False, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ) ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod From edc84f8323349612c90e4e9ed352e16cafbc0159 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 09:26:53 -0500 Subject: [PATCH 004/207] stash Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/config.py | 2 +- vllm/model_executor/layers/fused_moe/layer.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 10693119406a..428083fa4031 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1026,7 +1026,7 @@ class FusedMoEConfig: num_experts: int experts_per_token: int hidden_dim: int - intermediate_dim: int + intermediate_size_per_partition: int num_local_experts: int moe_parallel_config: FusedMoEParallelConfig diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index b12baddc1d9f..83c8fcdc308f 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -524,7 +524,7 @@ def __init__( num_experts=self.global_num_experts, experts_per_token=top_k, hidden_dim=hidden_size, - intermediate_dim=self.intermediate_size_per_partition, + intermediate_size_per_partition=self.intermediate_size_per_partition, num_local_experts=self.local_num_experts, moe_parallel_config=self.moe_parallel_config, in_dtype=moe_in_dtype, From 32260117711ad331e1b791ac73bed6ef987bb9d3 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 09:48:58 -0500 Subject: [PATCH 005/207] we have startup, but incorrect answers Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 14 ++++++++++---- .../layers/fused_moe/modular_kernel.py | 1 - 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 6f7a6ba850b4..e8600d9bea32 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -221,7 +221,9 @@ def __init__( self.num_expert_group = None self.topk = moe_config.experts_per_token - self.intermediate_dim = moe_config.intermediate_dim + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) self.hidden_dim = moe_config.hidden_dim self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank @@ -261,7 +263,11 @@ def workspace_shapes( # The workspaces for this implementation are managed by flashinfer. workspace1 = (0,) workspace2 = (0,) - output = (M, K) + + # Hidden states are Nvfp4, packed into int8 dtype. + assert self.hidden_dim == K * 2 + output = (M, self.hidden_dim) + return (workspace1, workspace2, output) def apply( @@ -315,7 +321,7 @@ def apply( top_k=self.topk, n_group=0, topk_group=0, - intermediate_size=self.intermediate_dim, + intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, routed_scaling_factor=None, @@ -324,7 +330,7 @@ def apply( do_finalize=True, )[0] - output.copy_(out) + output.copy_(out[0]) def apply_monolthic( self, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 82107581469d..13559080e820 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -867,7 +867,6 @@ def _allocate_buffers( expert_tokens_meta, activation, ) - print(f"{fused_out_shape=}") # We can reuse the memory between cache1 and cache3 because by the # time we need cache3, we're done with cache1. From 2b296296710b7ad903e7f2bb8a4fe415c4eefc83 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 10:25:12 -0500 Subject: [PATCH 006/207] stash Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index e8600d9bea32..dc8aaa871c20 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -210,6 +210,7 @@ def __init__( import flashinfer + self.moe_config = moe_config # TODO: set this via the constructor self.routing_method_type = flashinfer.RoutingMethodType.Renormalize # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 @@ -298,6 +299,13 @@ def apply( torch.bfloat16 ).view(torch.int16) + self.x = 0 if not hasattr(self, "x") else self.x + 1 + if self.x == 0 and self.moe_config.tp_rank == 0: + print(f"{a1q_scale=}") + print(f"{self.g1_scale_c=}") + print(f"{self.quant_config.g1_alphas=}") + print(f"{self.quant_config.g2_alphas=}") + # Invoke kernel. # TODO(avoid the copy). out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( From 3507b6b54f9f7c17dc409c98ee69ec04c1f94a7c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 10:52:59 -0500 Subject: [PATCH 007/207] stash Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 13 +++++++------ .../layers/fused_moe/modular_kernel.py | 1 - 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index dc8aaa871c20..951873d89aa2 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -299,12 +299,13 @@ def apply( torch.bfloat16 ).view(torch.int16) - self.x = 0 if not hasattr(self, "x") else self.x + 1 - if self.x == 0 and self.moe_config.tp_rank == 0: - print(f"{a1q_scale=}") - print(f"{self.g1_scale_c=}") - print(f"{self.quant_config.g1_alphas=}") - print(f"{self.quant_config.g2_alphas=}") + + + print(f"{w1[-1,-1]=}") + print(f"{self.quant_config.w1_scale[-1,-1]=}") + print(f"{w2[-1,-1]=}") + print(f"{self.quant_config.w2_scale[-1,-1]=}") + # Invoke kernel. # TODO(avoid the copy). diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 13559080e820..49629a35232c 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1103,7 +1103,6 @@ def input_chunk_range(chunk_idx: int) -> tuple[int, int]: c_fused_out = self._slice_output_tensor( fused_out, chunk_idx, num_chunks, CHUNK_SIZE, M_full ) - self.fused_experts.apply( output=c_fused_out, hidden_states=a1q[s:e], From 372d131d4d3c4ddc51a47b8d99f5592e1670eb2c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 11:12:58 -0500 Subject: [PATCH 008/207] stash Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 11 +++++++---- vllm/model_executor/layers/quantization/modelopt.py | 5 ++++- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 951873d89aa2..f1a1e3115e5b 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -301,10 +301,13 @@ def apply( - print(f"{w1[-1,-1]=}") - print(f"{self.quant_config.w1_scale[-1,-1]=}") - print(f"{w2[-1,-1]=}") - print(f"{self.quant_config.w2_scale[-1,-1]=}") + # print(f"{w1[-1,-1]=}") + # print(f"{self.quant_config.w1_scale[-1,-1]=}") + # print(f"{w2[-1,-1]=}") + # print(f"{self.quant_config.w2_scale[-1,-1]=}") + # print(f"{hidden_states[-1,-1]=}") + # print(f"{packed_tensor[-1,-1]=}") + # print(f"{a1q_scale[-1,-1]=}") # Invoke kernel. diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index cc5072c84e77..acb077714cd0 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1611,7 +1611,7 @@ def apply( ) assert self.kernel is not None - return self.kernel( + out = self.kernel( x, layer.w13_weight, layer.w2_weight, @@ -1624,6 +1624,9 @@ def apply( apply_router_weight_on_input=layer.apply_router_weight_on_input, ) + print(f"{out[-1,-1]=}") + return out + ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE From 0c26ada95800e8e6711d2dbaecb6e897c05aae30 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 11:32:16 -0500 Subject: [PATCH 009/207] stash Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index f1a1e3115e5b..3d1d0e63dd16 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -312,6 +312,8 @@ def apply( # Invoke kernel. # TODO(avoid the copy). + print(f"{self.ep_rank=}, {self.local_num_experts=}") + print(f"{self.ep_rank=}, {self.local_num_experts=}") out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( topk_ids=packed_tensor, routing_bias=None, From c0dcdfbf761458cf7749ffee5554b0cc283dd221 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 12:07:24 -0500 Subject: [PATCH 010/207] stash Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 2 -- vllm/model_executor/layers/fused_moe/layer.py | 1 + vllm/model_executor/layers/quantization/modelopt.py | 6 +++++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 3d1d0e63dd16..f1a1e3115e5b 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -312,8 +312,6 @@ def apply( # Invoke kernel. # TODO(avoid the copy). - print(f"{self.ep_rank=}, {self.local_num_experts=}") - print(f"{self.ep_rank=}, {self.local_num_experts=}") out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( topk_ids=packed_tensor, routing_bias=None, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 83c8fcdc308f..ff2c10160967 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -335,6 +335,7 @@ def __init__( router_logits_dtype: torch.dtype | None = None, ): super().__init__() + self.prefix = prefix # Allow disabling of the separate shared experts stream for # debug purposes. diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index acb077714cd0..b60d8ae49e0b 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1624,7 +1624,11 @@ def apply( apply_router_weight_on_input=layer.apply_router_weight_on_input, ) - print(f"{out[-1,-1]=}") + if layer.prefix in [ + "model.layers.0.mlp.experts", + "model.layers.47.mlp.experts" + ]: + print(f"{layer.prefix=} | {x.shape=}: {x[-1,-1]=}, {out[-1,-1]=}") return out From 313638e17121c4bf9363f5015bb743995fe9a413 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 12:28:03 -0500 Subject: [PATCH 011/207] working again, had incorrect copy Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 14 ++------------ .../model_executor/layers/quantization/modelopt.py | 5 ----- 2 files changed, 2 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index f1a1e3115e5b..ab89e1326fcc 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -299,17 +299,6 @@ def apply( torch.bfloat16 ).view(torch.int16) - - - # print(f"{w1[-1,-1]=}") - # print(f"{self.quant_config.w1_scale[-1,-1]=}") - # print(f"{w2[-1,-1]=}") - # print(f"{self.quant_config.w2_scale[-1,-1]=}") - # print(f"{hidden_states[-1,-1]=}") - # print(f"{packed_tensor[-1,-1]=}") - # print(f"{a1q_scale[-1,-1]=}") - - # Invoke kernel. # TODO(avoid the copy). out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( @@ -342,7 +331,8 @@ def apply( do_finalize=True, )[0] - output.copy_(out[0]) + assert output.shape == out.shape + output.copy_(out) def apply_monolthic( self, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index b60d8ae49e0b..c0ef0307886e 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1624,11 +1624,6 @@ def apply( apply_router_weight_on_input=layer.apply_router_weight_on_input, ) - if layer.prefix in [ - "model.layers.0.mlp.experts", - "model.layers.47.mlp.experts" - ]: - print(f"{layer.prefix=} | {x.shape=}: {x[-1,-1]=}, {out[-1,-1]=}") return out From cc1a2ea8eb0b5092c1e5316c92918b781fede043 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 12:28:48 -0500 Subject: [PATCH 012/207] remove loc Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/modular_kernel.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 49629a35232c..13559080e820 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1103,6 +1103,7 @@ def input_chunk_range(chunk_idx: int) -> tuple[int, int]: c_fused_out = self._slice_output_tensor( fused_out, chunk_idx, num_chunks, CHUNK_SIZE, M_full ) + self.fused_experts.apply( output=c_fused_out, hidden_states=a1q[s:e], From 942b758708f5a0c43153a43eabf7ec1448b2dbe7 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 12:35:02 -0500 Subject: [PATCH 013/207] make trtllm use the inplace buffer Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index ab89e1326fcc..41b374eeecf3 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import flashinfer import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -208,8 +209,6 @@ def __init__( ): super().__init__(quant_config) - import flashinfer - self.moe_config = moe_config # TODO: set this via the constructor self.routing_method_type = flashinfer.RoutingMethodType.Renormalize @@ -289,8 +288,6 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - import flashinfer - assert activation == "silu" assert a1q_scale is not None @@ -300,8 +297,7 @@ def apply( ).view(torch.int16) # Invoke kernel. - # TODO(avoid the copy). - out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( + flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( topk_ids=packed_tensor, routing_bias=None, hidden_states=hidden_states, @@ -329,10 +325,8 @@ def apply( tile_tokens_dim=None, routing_method_type=1, do_finalize=True, - )[0] - - assert output.shape == out.shape - output.copy_(out) + output=output, + ) def apply_monolthic( self, @@ -350,8 +344,6 @@ def apply_monolthic( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: - import flashinfer - assert activation == "silu" # Quantize input to FP4 From 67758ba7396f867792e03a9ac9e68103b0b46372 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 12:35:57 -0500 Subject: [PATCH 014/207] remove debug loc Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/layer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index ff2c10160967..83c8fcdc308f 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -335,7 +335,6 @@ def __init__( router_logits_dtype: torch.dtype | None = None, ): super().__init__() - self.prefix = prefix # Allow disabling of the separate shared experts stream for # debug purposes. From c273804d90231c96fa9b83ad9160ebe0a970f9f6 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 12:59:44 -0500 Subject: [PATCH 015/207] fix precommits Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 15 +++++++-------- .../layers/quantization/modelopt.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 41b374eeecf3..33f0d66e7e9c 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -290,8 +290,10 @@ def apply( ): assert activation == "silu" assert a1q_scale is not None + assert self.quant_config.w1_scale is not None + assert self.quant_config.w2_scale is not None - # Pack topk_ids and topk_weights into format expected by the kernel. + # Pack topk ids and weights into format expected by the kernel. packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to( torch.bfloat16 ).view(torch.int16) @@ -346,7 +348,7 @@ def apply_monolthic( ) -> torch.Tensor: assert activation == "silu" - # Quantize input to FP4 + # Quantize input. if isinstance(hidden_states, tuple): a1q, a1q_scale = hidden_states else: @@ -356,19 +358,18 @@ def apply_monolthic( is_sf_swizzled_layout=False, ) - # Prepare routing bias + # Prepare routing bias into kernel format. routing_bias = self.e_score_correction_bias if routing_bias is not None: routing_bias = routing_bias.to(torch.bfloat16) - router_logits = ( router_logits.to(torch.float32) if self.routing_method_type == RoutingMethodType.DeepSeekV3 else router_logits ) - # Call TRT-LLM FP4 block-scale MoE kernel - out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe( + # Invoke kernel. + return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( routing_logits=router_logits, routing_bias=routing_bias, hidden_states=a1q, @@ -397,5 +398,3 @@ def apply_monolthic( routing_method_type=self.routing_method_type, do_finalize=True, )[0] - - return out diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index c0ef0307886e..2752d7a9eccd 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1349,6 +1349,11 @@ def __init__( self.nvfp4_backend ) self.kernel: mk.FusedMoEModularKernel | None = None + self.experts: mk.FusedMoEPermuteExpertsUnpermute | None = None + self.use_monolithic = ( + self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM + and not moe_config.moe_parallel_config.use_all2all_kernels + ) def maybe_make_prepare_finalize( self, @@ -1605,6 +1610,20 @@ def apply( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if self.use_monolithic: + assert self.experts is not None + out = self.experts.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + inplace=False, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ) + topk_weights, topk_ids = router.select_experts( hidden_states=x, router_logits=router_logits, From 6a4be4d1ed1f802d064264986eddcf5ff7fe39df Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:01:42 -0500 Subject: [PATCH 016/207] fix precommits Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/prepare_finalize.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 5a001b769032..99a21e6424c9 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -57,14 +57,20 @@ def prepare( if self.defer_input_quant: return a1, None, None, None, None + a1_scale = ( + quant_config.a1_gscale + if (quant_config.quant_dtype == "nvfp4") + else quant_config.a1_scale + ) a1q, a1q_scale = moe_kernel_quantize_input( a1, + a1_scale, # quant_config.a1_scale, quant_config.a1_gscale, quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, - is_fp4_scale_swizzled=False, # ? + is_fp4_scale_swizzled=False, ) return a1q, a1q_scale, None, None, None From dd32a16aa5a0c0764462e7218d6bdb7827758697 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:01:58 -0500 Subject: [PATCH 017/207] fix precommits Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/prepare_finalize.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 99a21e6424c9..e5c136f52cbe 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -65,8 +65,6 @@ def prepare( a1q, a1q_scale = moe_kernel_quantize_input( a1, a1_scale, - # quant_config.a1_scale, - quant_config.a1_gscale, quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, From 2074c44d9dfca124160673174b86b2687a92122f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:14:46 -0500 Subject: [PATCH 018/207] updates Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 5 +---- .../layers/fused_moe/oracle/nvfp4.py | 2 +- .../layers/quantization/modelopt.py | 16 +++++++--------- 3 files changed, 9 insertions(+), 14 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 33f0d66e7e9c..8cc9aaefa2f3 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -330,7 +330,7 @@ def apply( output=output, ) - def apply_monolthic( + def apply_monolithic( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -341,9 +341,6 @@ def apply_monolthic( expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, a2_scale: torch.Tensor | None, - workspace13: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: assert activation == "silu" diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index ec3b37336b96..346b957412d8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -216,7 +216,7 @@ def make_nvfp4_moe_quant_config( w2_scale_2: torch.Tensor, a13_scale: torch.Tensor, a2_scale: torch.Tensor, -) -> FusedMoEQuantConfig | None: +) -> FusedMoEQuantConfig: if backend == NvFp4MoeBackend.MARLIN: return nvfp4_w4a16_moe_quant_config( g1_alphas=w13_scale_2, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 2752d7a9eccd..f413a1672c01 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1349,7 +1349,6 @@ def __init__( self.nvfp4_backend ) self.kernel: mk.FusedMoEModularKernel | None = None - self.experts: mk.FusedMoEPermuteExpertsUnpermute | None = None self.use_monolithic = ( self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM and not moe_config.moe_parallel_config.use_all2all_kernels @@ -1560,7 +1559,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: self.moe_quant_config = self.get_fused_moe_quant_config(layer) use_dp = self.moe.dp_size > 1 - if self.moe_quant_config is not None and not use_dp: + if not use_dp: self.kernel = make_nvfp4_moe_kernel( backend=self.nvfp4_backend, quant_config=self.moe_quant_config, @@ -1586,9 +1585,7 @@ def prepare_dp_allgather_tensor( extra_tensors: list[torch.Tensor] = [hidden_states_sf] return hidden_states_fp4, extra_tensors - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: return make_nvfp4_moe_quant_config( backend=self.nvfp4_backend, w13_scale=layer.w13_weight_scale, @@ -1610,17 +1607,19 @@ def apply( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert self.kernel is not None + if self.use_monolithic: - assert self.experts is not None - out = self.experts.apply_monolithic( + out = self.kernel.fused_experts.apply_monolithic( x, layer.w13_weight, layer.w2_weight, router_logits, - inplace=False, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, + a1q_scale=None, + a2_scale=None, apply_router_weight_on_input=layer.apply_router_weight_on_input, ) @@ -1629,7 +1628,6 @@ def apply( router_logits=router_logits, ) - assert self.kernel is not None out = self.kernel( x, layer.w13_weight, From bbd9c4c5789c2af04435a36ce993a6cc7a9ec83d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:29:23 -0500 Subject: [PATCH 019/207] add interfaces Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 2 - .../layers/fused_moe/modular_kernel.py | 35 +++++++++++++--- .../layers/quantization/modelopt.py | 40 +++++++++---------- 3 files changed, 49 insertions(+), 28 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 8cc9aaefa2f3..57eabc42c645 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -339,8 +339,6 @@ def apply_monolithic( activation: str, global_num_experts: int, expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: assert activation == "silu" diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 13559080e820..4f2eb517a596 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -669,15 +669,10 @@ def apply_monolthic( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - routing_logits: torch.Tensor, + router_logits: torch.Tensor, activation: str, global_num_experts: int, expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, - workspace13: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: """ @@ -1274,3 +1269,31 @@ def forward( topk_ids, apply_router_weight_on_input, ) + + def forward_monolithic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + """ + Same as forward(), except uses routing_logits as opposed + to the topk_ids and topk_weights. This is used for kernels + that have fused router + experts (e.g. FLASHINFER_TRTLLM). + """ + + return self.fused_experts.apply_monolthic( + hidden_states=hidden_states, + w1=w1, + w2=w2, + router_logits=router_logits, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + ) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index f413a1672c01..ec3083ca83ff 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1610,7 +1610,8 @@ def apply( assert self.kernel is not None if self.use_monolithic: - out = self.kernel.fused_experts.apply_monolithic( + # In monolithic case, router is fused with expert. + out = self.kernel.forward_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -1618,28 +1619,27 @@ def apply( activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, - a1q_scale=None, - a2_scale=None, apply_router_weight_on_input=layer.apply_router_weight_on_input, ) + else: + # Otherwise, expert selection is separate. + topk_weights, topk_ids = router.select_experts( + hidden_states=x, + router_logits=router_logits, + ) - topk_weights, topk_ids = router.select_experts( - hidden_states=x, - router_logits=router_logits, - ) - - out = self.kernel( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - inplace=False, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - ) + out = self.kernel( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + inplace=False, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ) return out From 22733bef221959fbd5b74cdce57760daf147c340 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:29:44 -0500 Subject: [PATCH 020/207] nits Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/modular_kernel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 4f2eb517a596..d220752706c2 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -676,7 +676,7 @@ def apply_monolthic( apply_router_weight_on_input: bool, ) -> torch.Tensor: """ - Same as apply(), except uses routing_logits as opposed + Same as apply(), except uses router_logits as opposed to the topk_ids and topk_weights. This is useful for kernels with fused router and fused_experts (e.g. FLASHINFER_TRTLLM). """ @@ -1282,7 +1282,7 @@ def forward_monolithic( apply_router_weight_on_input: bool, ) -> torch.Tensor: """ - Same as forward(), except uses routing_logits as opposed + Same as forward(), except uses router_logits as opposed to the topk_ids and topk_weights. This is used for kernels that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ From 75627c8a662e2802c2c3acab9f0ea6d71c3d332a Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:41:32 -0500 Subject: [PATCH 021/207] update comments Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 7 +++++-- vllm/model_executor/layers/fused_moe/modular_kernel.py | 4 ++-- vllm/model_executor/layers/quantization/modelopt.py | 10 ++++++---- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 57eabc42c645..4c480cafd344 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -228,7 +228,9 @@ def __init__( self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank - # a13_scale * w13_scale_2 / a2_scale + # g1_alpha_s = a13_scale * w13_scale_2 + # a2_gscale = (1 / a2_scale) + # g1_scale_c = a13_scale * w13_scale_2 / a2_scale self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale @property @@ -264,7 +266,8 @@ def workspace_shapes( workspace1 = (0,) workspace2 = (0,) - # Hidden states are Nvfp4, packed into int8 dtype. + # Hidden states are Nvfp4, packed into int8 dtype, so we + # need to multiply K by 2 to get the output shape right. assert self.hidden_dim == K * 2 output = (M, self.hidden_dim) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index d220752706c2..8acbad8f2fc5 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -664,7 +664,7 @@ def apply( """ raise NotImplementedError - def apply_monolthic( + def apply_monolithic( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -1287,7 +1287,7 @@ def forward_monolithic( that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ - return self.fused_experts.apply_monolthic( + return self.fused_experts.apply_monolithic( hidden_states=hidden_states, w1=w1, w2=w2, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index ec3083ca83ff..0092a5f7d94c 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1349,10 +1349,6 @@ def __init__( self.nvfp4_backend ) self.kernel: mk.FusedMoEModularKernel | None = None - self.use_monolithic = ( - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - and not moe_config.moe_parallel_config.use_all2all_kernels - ) def maybe_make_prepare_finalize( self, @@ -1566,6 +1562,11 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: moe_config=self.moe, ) + self.use_monolithic = ( + layer.enable_eplb + and self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM + ) + def prepare_dp_allgather_tensor( self, layer: FusedMoE, @@ -1609,6 +1610,7 @@ def apply( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert self.kernel is not None + logger.info_once(f"{self.use_monolithic=}", scope="local") if self.use_monolithic: # In monolithic case, router is fused with expert. out = self.kernel.forward_monolithic( From 6c44f2aec41e9e611efb7c950c5d62f69f129b14 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:45:39 -0500 Subject: [PATCH 022/207] remove debug logging Signed-off-by: Robert Shaw --- vllm/model_executor/layers/quantization/modelopt.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 0092a5f7d94c..0db935eea970 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1563,7 +1563,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) self.use_monolithic = ( - layer.enable_eplb + not layer.enable_eplb and self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM ) @@ -1610,7 +1610,6 @@ def apply( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert self.kernel is not None - logger.info_once(f"{self.use_monolithic=}", scope="local") if self.use_monolithic: # In monolithic case, router is fused with expert. out = self.kernel.forward_monolithic( @@ -1624,7 +1623,6 @@ def apply( apply_router_weight_on_input=layer.apply_router_weight_on_input, ) else: - # Otherwise, expert selection is separate. topk_weights, topk_ids = router.select_experts( hidden_states=x, router_logits=router_logits, From 22933cebfbdbf3c35ee9b9ec968dcd32fe269e2b Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:55:13 -0500 Subject: [PATCH 023/207] apply to compressed-tensors Signed-off-by: Robert Shaw --- .../compressed_tensors_moe.py | 72 ++++++++----------- .../layers/quantization/modelopt.py | 6 +- 2 files changed, 34 insertions(+), 44 deletions(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 13a123ba6026..01ff1ecf685d 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -63,8 +63,6 @@ ) from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( build_flashinfer_fp4_cutlass_moe_prepare_finalize, - flashinfer_trtllm_fp4_moe, - flashinfer_trtllm_fp4_routed_moe, select_nvfp4_gemm_impl, ) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( @@ -574,13 +572,18 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Initialize the kernel that will be called in apply(). self.moe_quant_config = self.get_fused_moe_quant_config(layer) use_dp = self.moe.dp_size > 1 - if self.moe_quant_config is not None and not use_dp: + if not use_dp: self.kernel = make_nvfp4_moe_kernel( backend=self.nvfp4_backend, quant_config=self.moe_quant_config, moe_config=self.moe, ) + self.use_monolithic = ( + not layer.enable_eplb + and self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM + ) + def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, @@ -616,9 +619,7 @@ def select_gemm_impl( logger.debug_once("Using %s", experts.__class__.__name__) return experts - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: return make_nvfp4_moe_quant_config( backend=self.nvfp4_backend, w13_scale=layer.w13_weight_scale, @@ -637,49 +638,32 @@ def apply( router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert layer.activation == "silu", "Only SiLU activation is supported." + assert self.kernel is not None - if ( - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - and not layer.enable_eplb - ): - return flashinfer_trtllm_fp4_moe( - layer=layer, - x=x, - router_logits=router_logits, - top_k=layer.top_k, + if self.use_monolithic: + # In monolithic case, router is fused with expert. + out = self.kernel.forward_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - custom_routing_function=layer.custom_routing_function, - e_score_correction_bias=layer.e_score_correction_bias, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, ) - - # Hidden_states in select_experts is only used to extract metadata - if isinstance(x, tuple): - x_routing, _ = x else: - x_routing = x - topk_weights, topk_ids = router.select_experts( - hidden_states=x_routing, - router_logits=router_logits, - ) - - # EPLB path - if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - assert layer.enable_eplb - return flashinfer_trtllm_fp4_routed_moe( - layer=layer, - x=x, - topk_ids=topk_ids, - topk_weights=topk_weights, - top_k=layer.top_k, - activation=layer.activation, - global_num_experts=layer.global_num_experts, + # Hidden_states in select_experts is only used to extract metadata + if isinstance(x, tuple): + x_routing, _ = x + else: + x_routing = x + topk_weights, topk_ids = router.select_experts( + hidden_states=x_routing, + router_logits=router_logits, ) - else: - assert self.kernel is not None - return self.kernel( + + out = self.kernel( x, layer.w13_weight, layer.w2_weight, @@ -692,6 +676,8 @@ def apply( apply_router_weight_on_input=layer.apply_router_weight_on_input, ) + return out + class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): def __init__( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 0db935eea970..4ffb434e58dc 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1623,8 +1623,12 @@ def apply( apply_router_weight_on_input=layer.apply_router_weight_on_input, ) else: + if isinstance(x, tuple): + x_routing, _ = x + else: + x_routing = x topk_weights, topk_ids = router.select_experts( - hidden_states=x, + hidden_states=x_routing, router_logits=router_logits, ) From 80edbadf1b33dab3c746da7ed5808ac32f5cb0a0 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:57:29 -0500 Subject: [PATCH 024/207] updated Signed-off-by: Robert Shaw --- .../quantization/compressed_tensors/compressed_tensors_moe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 01ff1ecf685d..70efa6473962 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -653,7 +653,6 @@ def apply( apply_router_weight_on_input=layer.apply_router_weight_on_input, ) else: - # Hidden_states in select_experts is only used to extract metadata if isinstance(x, tuple): x_routing, _ = x else: From 30c99e4ff20d6e81cbdc7ea481d90420703c27b1 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 13:59:44 -0500 Subject: [PATCH 025/207] remove the trtllm entrypoints Signed-off-by: Robert Shaw --- .../quantization/utils/flashinfer_fp4_moe.py | 288 +----------------- 1 file changed, 1 insertion(+), 287 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 272b13861fee..2b201cd85e7b 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -12,7 +12,6 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEQuantConfig, - RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import ( FlashInferCuteDSLExperts, @@ -23,9 +22,6 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501 create_flashinfer_prepare_finalize, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - swizzle_blockscale, -) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( has_flashinfer_cutedsl_grouped_gemm_nt_masked, @@ -33,9 +29,7 @@ ) if TYPE_CHECKING: - from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( - NvFp4MoeBackend, - ) + pass logger = init_logger(__name__) @@ -248,283 +242,3 @@ def prepare_static_weights_for_trtllm_fp4_moe( gemm2_weights_fp4_shuffled, gemm2_scales_fp4_shuffled, ) - - -def flashinfer_trtllm_fp4_moe( - layer: torch.nn.Module, - x: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - router_logits: torch.Tensor, - top_k: int, - activation: str, - global_num_experts: int, - num_expert_group: int | None, - topk_group: int | None, - custom_routing_function: object | None, - e_score_correction_bias: torch.Tensor | None, -) -> torch.Tensor: - """ - Apply FlashInfer TensorRT-LLM FP4 MoE kernel. - - Args: - layer: The MoE layer with weights and scales - x: Input tensor - router_logits: Router logits for expert selection - top_k: Number of experts to select per token - activation: Activation function to use - global_num_experts: Total number of experts across all ranks - num_expert_group: Number of expert groups (for grouped routing) - topk_group: Top-k within each group - custom_routing_function: Custom routing function (e.g., Llama4) - e_score_correction_bias: Optional routing bias correction - - Returns: - Output tensor from the MoE layer - """ - import flashinfer - - from vllm.model_executor.models.llama4 import Llama4MoE - - # https://github.com/flashinfer-ai/flashinfer/blob/f0277fd1bff90e309e5c19cab36c5dae056d685d/flashinfer/fused_moe/core.py#L2404 - assert activation == "silu", ( - "Only SiLU activation is supported for FlashInfer TRTLLM FP4 MoE. " - f"{activation} found instead." - ) - - # Quantize input to FP4 - if isinstance(x, tuple): - hidden_states_fp4, hidden_states_scale_linear_fp4 = x - else: - # hidden_states is the already quantized - (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize( - x, - layer.a1_gscale, - is_sf_swizzled_layout=False, - ) - - # Determine routing method type - use_llama4_routing = custom_routing_function is Llama4MoE.custom_routing_function - routing_method_type = layer.routing_method_type - if use_llama4_routing: - routing_method_type = flashinfer.RoutingMethodType.Llama4 - - # Prepare routing bias - routing_bias = e_score_correction_bias - if routing_bias is not None: - routing_bias = routing_bias.to(torch.bfloat16) - - router_logits = ( - router_logits.to(torch.float32) - if routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits - ) - - # Call TRT-LLM FP4 block-scale MoE kernel - out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe( - routing_logits=router_logits, - routing_bias=routing_bias, - hidden_states=hidden_states_fp4, - hidden_states_scale=hidden_states_scale_linear_fp4.view( - torch.float8_e4m3fn - ).flatten(), - gemm1_weights=layer.w13_weight.data, - gemm1_weights_scale=layer.w13_weight_scale.data.view(torch.float8_e4m3fn), - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=layer.w2_weight.data, - gemm2_weights_scale=layer.w2_weight_scale.data.view(torch.float8_e4m3fn), - gemm2_bias=None, - output1_scale_scalar=layer.g1_scale_c.data, - output1_scale_gate_scalar=layer.g1_alphas.data, - output2_scale_scalar=layer.g2_alphas.data, - num_experts=global_num_experts, - top_k=top_k, - n_group=num_expert_group if num_expert_group is not None else 0, - topk_group=topk_group if topk_group is not None else 0, - intermediate_size=layer.intermediate_size_per_partition, - local_expert_offset=layer.ep_rank * layer.local_num_experts, - local_num_experts=layer.local_num_experts, - routed_scaling_factor=None, - tile_tokens_dim=None, - routing_method_type=routing_method_type, - do_finalize=True, - )[0] - - return out - - -def flashinfer_trtllm_fp4_routed_moe( - layer: torch.nn.Module, - x: torch.Tensor, - topk_ids: torch.Tensor, - topk_weights: torch.Tensor, - top_k: int, - activation: str, - global_num_experts: int, -) -> torch.Tensor: - """ - Apply FlashInfer TensorRT-LLM FP4 MoE kernel. Uses packed - input top k expert indices and scores rather than computing - top k expert indices from scores. - - Args: - layer: The MoE layer with weights and scales - x: Input tensor - topk_ids: Ids of selected experts - top_k: Number of experts to select per token - activation: Activation function to use - global_num_experts: Total number of experts across all ranks - - Returns: - Output tensor from the MoE layer - """ - import flashinfer - - # https://github.com/flashinfer-ai/flashinfer/blob/f0277fd1bff90e309e5c19cab36c5dae056d685d/flashinfer/fused_moe/core.py#L2535 - assert activation == "silu", ( - "Only SiLU activation is supported for FlashInfer TRTLLM FP4 Routed MoE. " - f"{activation} found instead." - ) - - # Pack top k ids and expert weights into a single int32 tensor, as - # required by TRT-LLM - packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to( - torch.bfloat16 - ).view(torch.int16) - - if isinstance(x, tuple): - # Hidden_states is the already quantized - hidden_states_fp4, hidden_states_scale_linear_fp4 = x - else: - # Quantize input to FP4 - (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize( - x, - layer.a1_gscale, - is_sf_swizzled_layout=False, - ) - - # Call TRT-LLM FP4 block-scale MoE kernel - out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( - topk_ids=packed_tensor, - routing_bias=None, - hidden_states=hidden_states_fp4, - hidden_states_scale=hidden_states_scale_linear_fp4.view( - torch.float8_e4m3fn - ).flatten(), - gemm1_weights=layer.w13_weight.data, - gemm1_weights_scale=layer.w13_weight_scale.data.view(torch.float8_e4m3fn), - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=layer.w2_weight.data, - gemm2_weights_scale=layer.w2_weight_scale.data.view(torch.float8_e4m3fn), - gemm2_bias=None, - output1_scale_scalar=layer.g1_scale_c.data, - output1_scale_gate_scalar=layer.g1_alphas.data, - output2_scale_scalar=layer.g2_alphas.data, - num_experts=global_num_experts, - top_k=top_k, - n_group=0, - topk_group=0, - intermediate_size=layer.intermediate_size_per_partition, - local_expert_offset=layer.ep_rank * layer.local_num_experts, - local_num_experts=layer.local_num_experts, - routed_scaling_factor=None, - tile_tokens_dim=None, - routing_method_type=1, - do_finalize=True, - )[0] - - return out - - -def prepare_nvfp4_moe_layer_for_fi_or_cutlass( - backend: "NvFp4MoeBackend", - layer: torch.nn.Module, - w13: torch.Tensor, - w13_scale: torch.Tensor, - w13_scale_2: torch.Tensor, - a13_scale: torch.Tensor, - w2: torch.Tensor, - w2_scale: torch.Tensor, - w2_scale_2: torch.Tensor, - a2_scale: torch.Tensor, - is_act_and_mul: bool, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - # Delayed import for circular dependency avoidance. - from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( - NvFp4MoeBackend, - is_global_sf_supported_for_nvfp4_backend, - ) - - assert backend in [ - NvFp4MoeBackend.VLLM_CUTLASS, - NvFp4MoeBackend.FLASHINFER_CUTLASS, - NvFp4MoeBackend.FLASHINFER_TRTLLM, - NvFp4MoeBackend.FLASHINFER_CUTEDSL, - ] - - # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels. - if is_act_and_mul and backend in [ - NvFp4MoeBackend.FLASHINFER_CUTLASS, - NvFp4MoeBackend.FLASHINFER_TRTLLM, - ]: - w13, w13_scale = reorder_w1w3_to_w3w1(w13, w13_scale) - - # For some FI kernels, the input scales are shared by all experts. - if is_global_sf_supported_for_nvfp4_backend(backend): - num_experts = w13.shape[0] - a13_scale = a13_scale.max().to(torch.float32).expand(num_experts) - a2_scale = a2_scale.max().to(torch.float32).expand(num_experts) - else: - a13_scale = a13_scale.max(dim=1).values.to(torch.float32) - - # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels. - if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - w13, w13_scale, w2, w2_scale = prepare_static_weights_for_trtllm_fp4_moe( - w13, - w2, - w13_scale, - w2_scale, - w2.size(-2), # hidden_size - w13.size(-2) // 2, # intermediate_size - w13.size(0), # num_experts - ) - - # We do not need to make this a parameter, because - # it is not used during the weight (re)-loading process. - layer.g1_scale_c = a13_scale * w13_scale_2 / a2_scale - layer.a1_gscale = 1.0 / a13_scale - layer.g1_alphas = a13_scale * w13_scale_2 - layer.g2_alphas = a2_scale * w2_scale_2 - else: - # Swizzle the block scales for other FI NVFP4 MoE kernels. - w13_scale = swizzle_blockscale(w13_scale) - - # Apply padding if needed. - pad_size = w13_scale.size(1) - w13.size(1) - if pad_size > 0: - if is_act_and_mul: - raise NotImplementedError( - "Intermediate size padding for w1 and w3, for %s " - "NvFp4 backend, but this is not currently supported", - backend.value, - ) - w13 = torch.nn.functional.pad(w13, (0, 0, 0, pad_size)) - w2 = torch.nn.functional.pad(w2, (0, pad_size // 2, 0, 0)) - w2_scale = torch.nn.functional.pad(w2_scale, (0, pad_size // 16)) - - w2_scale = swizzle_blockscale(w2_scale) - - return w13, w13_scale, w13_scale_2, a13_scale, w2, w2_scale, w2_scale_2, a2_scale From 61b501299ffe8cd5856594e930d702ef1f07f35c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 14:00:59 -0500 Subject: [PATCH 026/207] remove typing Signed-off-by: Robert Shaw --- .../layers/quantization/utils/flashinfer_fp4_moe.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 2b201cd85e7b..4c06dc1b7b27 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility helpers for NVFP4 + FlashInfer fused-MoE path""" -from typing import TYPE_CHECKING - import torch import vllm.envs as envs @@ -28,9 +26,6 @@ has_flashinfer_cutlass_fused_moe, ) -if TYPE_CHECKING: - pass - logger = init_logger(__name__) From 8bec44c347b6c07c28c63d69142754a43ab0b1b0 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 14:14:05 -0500 Subject: [PATCH 027/207] rename and move to a separate file Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 204 ----------------- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 212 ++++++++++++++++++ .../layers/fused_moe/oracle/nvfp4.py | 2 +- 3 files changed, 213 insertions(+), 205 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index 4c480cafd344..c66a30623a3b 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -1,18 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import flashinfer import torch -import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEQuantConfig, RoutingMethodType, ) -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceNoOP, -) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( calculate_tile_tokens_dim, @@ -199,200 +192,3 @@ def fi_trtllm_fp8_per_tensor_moe_fake( fake_impl=fi_trtllm_fp8_per_tensor_moe_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) - - -class FlashInferTrtLlmNvFp4Experts(mk.FusedMoEPermuteExpertsUnpermute): - def __init__( - self, - moe_config: FusedMoEConfig, - quant_config: FusedMoEQuantConfig, - ): - super().__init__(quant_config) - - self.moe_config = moe_config - # TODO: set this via the constructor - self.routing_method_type = flashinfer.RoutingMethodType.Renormalize - # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 - # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 - - self.routing_bias = None - self.e_score_correction_bias = None - self.topk_group = None - self.num_expert_group = None - - self.topk = moe_config.experts_per_token - self.intermediate_size_per_partition = ( - moe_config.intermediate_size_per_partition - ) - self.hidden_dim = moe_config.hidden_dim - self.local_num_experts = moe_config.num_local_experts - self.ep_rank = moe_config.moe_parallel_config.ep_rank - - # g1_alpha_s = a13_scale * w13_scale_2 - # a2_gscale = (1 / a2_scale) - # g1_scale_c = a13_scale * w13_scale_2 / a2_scale - self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale - - @property - def activation_formats( - self, - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]: - return ( - mk.FusedMoEActivationFormat.Standard, - mk.FusedMoEActivationFormat.Standard, - ) - - def supports_chunking(self) -> bool: - return False - - def supports_expert_map(self) -> bool: - return False - - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - return TopKWeightAndReduceNoOP() - - def workspace_shapes( - self, - M: int, - N: int, - K: int, - topk: int, - global_num_experts: int, - local_num_experts: int, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - activation: str, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - # The workspaces for this implementation are managed by flashinfer. - workspace1 = (0,) - workspace2 = (0,) - - # Hidden states are Nvfp4, packed into int8 dtype, so we - # need to multiply K by 2 to get the output shape right. - assert self.hidden_dim == K * 2 - output = (M, self.hidden_dim) - - return (workspace1, workspace2, output) - - def apply( - self, - output: torch.Tensor, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: str, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, - workspace13: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - apply_router_weight_on_input: bool, - ): - assert activation == "silu" - assert a1q_scale is not None - assert self.quant_config.w1_scale is not None - assert self.quant_config.w2_scale is not None - - # Pack topk ids and weights into format expected by the kernel. - packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to( - torch.bfloat16 - ).view(torch.int16) - - # Invoke kernel. - flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( - topk_ids=packed_tensor, - routing_bias=None, - hidden_states=hidden_states, - hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), - gemm1_weights=w1, - gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=w2, - gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), - gemm2_bias=None, - output1_scale_scalar=self.g1_scale_c, - output1_scale_gate_scalar=self.quant_config.g1_alphas, - output2_scale_scalar=self.quant_config.g2_alphas, - num_experts=global_num_experts, - top_k=self.topk, - n_group=0, - topk_group=0, - intermediate_size=self.intermediate_size_per_partition, - local_expert_offset=self.ep_rank * self.local_num_experts, - local_num_experts=self.local_num_experts, - routed_scaling_factor=None, - tile_tokens_dim=None, - routing_method_type=1, - do_finalize=True, - output=output, - ) - - def apply_monolithic( - self, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - router_logits: torch.Tensor, - activation: str, - global_num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, - ) -> torch.Tensor: - assert activation == "silu" - - # Quantize input. - if isinstance(hidden_states, tuple): - a1q, a1q_scale = hidden_states - else: - a1q, a1q_scale = flashinfer.fp4_quantize( - hidden_states, - self.quant_config.a1_gscale, - is_sf_swizzled_layout=False, - ) - - # Prepare routing bias into kernel format. - routing_bias = self.e_score_correction_bias - if routing_bias is not None: - routing_bias = routing_bias.to(torch.bfloat16) - router_logits = ( - router_logits.to(torch.float32) - if self.routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits - ) - - # Invoke kernel. - return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( - routing_logits=router_logits, - routing_bias=routing_bias, - hidden_states=a1q, - hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), - gemm1_weights=w1, - gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=w2, - gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), - gemm2_bias=None, - output1_scale_scalar=self.g1_scale_c, - output1_scale_gate_scalar=self.quant_config.g1_alphas, - output2_scale_scalar=self.quant_config.g2_alphas, - num_experts=global_num_experts, - top_k=self.topk, - n_group=self.num_expert_group if self.num_expert_group is not None else 0, - topk_group=self.topk_group if self.topk_group is not None else 0, - intermediate_size=self.intermediate_size_per_partition, - local_expert_offset=self.ep_rank * self.local_num_experts, - local_num_experts=self.local_num_experts, - routed_scaling_factor=None, - tile_tokens_dim=None, - routing_method_type=self.routing_method_type, - do_finalize=True, - )[0] diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py new file mode 100644 index 000000000000..3843c3610626 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import flashinfer +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) + + +class FlashInferTrtLlmNvFp4Experts(mk.FusedMoEPermuteExpertsUnpermute): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(quant_config) + + self.moe_config = moe_config + # TODO: set this via the constructor + self.routing_method_type = flashinfer.RoutingMethodType.Renormalize + # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 + # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 + + self.routing_bias = None + self.e_score_correction_bias = None + self.topk_group = None + self.num_expert_group = None + + self.topk = moe_config.experts_per_token + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.hidden_dim = moe_config.hidden_dim + self.local_num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + # g1_alpha_s = a13_scale * w13_scale_2 + # a2_gscale = (1 / a2_scale) + # g1_scale_c = a13_scale * w13_scale_2 / a2_scale + self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale + + @property + def activation_formats( + self, + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]: + return ( + mk.FusedMoEActivationFormat.Standard, + mk.FusedMoEActivationFormat.Standard, + ) + + def supports_chunking(self) -> bool: + return False + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: str, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # The workspaces for this implementation are managed by flashinfer. + workspace1 = (0,) + workspace2 = (0,) + + # Hidden states are Nvfp4, packed into int8 dtype, so we + # need to multiply K by 2 to get the output shape right. + assert self.hidden_dim == K * 2 + output = (M, self.hidden_dim) + + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert activation == "silu" + assert a1q_scale is not None + assert self.quant_config.w1_scale is not None + assert self.quant_config.w2_scale is not None + + # Pack topk ids and weights into format expected by the kernel. + packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to( + torch.bfloat16 + ).view(torch.int16) + + # Invoke kernel. + flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_tensor, + routing_bias=None, + hidden_states=hidden_states, + hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c, + output1_scale_gate_scalar=self.quant_config.g1_alphas, + output2_scale_scalar=self.quant_config.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + n_group=0, + topk_group=0, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=None, + tile_tokens_dim=None, + routing_method_type=1, + do_finalize=True, + output=output, + ) + + def apply_monolithic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + assert activation == "silu" + + # Quantize input. + if isinstance(hidden_states, tuple): + a1q, a1q_scale = hidden_states + else: + a1q, a1q_scale = flashinfer.fp4_quantize( + hidden_states, + self.quant_config.a1_gscale, + is_sf_swizzled_layout=False, + ) + + # Prepare routing bias into kernel format. + routing_bias = self.e_score_correction_bias + if routing_bias is not None: + routing_bias = routing_bias.to(torch.bfloat16) + router_logits = ( + router_logits.to(torch.float32) + if self.routing_method_type == RoutingMethodType.DeepSeekV3 + else router_logits + ) + + # Invoke kernel. + return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( + routing_logits=router_logits, + routing_bias=routing_bias, + hidden_states=a1q, + hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c, + output1_scale_gate_scalar=self.quant_config.g1_alphas, + output2_scale_scalar=self.quant_config.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + n_group=self.num_expert_group if self.num_expert_group is not None else 0, + topk_group=self.topk_group if self.topk_group is not None else 0, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=None, + tile_tokens_dim=None, + routing_method_type=self.routing_method_type, + do_finalize=True, + )[0] diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 346b957412d8..2317c0ae47b2 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, ) -from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import ( +from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_nvfp4_moe import ( FlashInferTrtLlmNvFp4Experts, ) from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( From 7e914f8a784f3cdee1399ea51ecc11bb72c46649 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 14:15:08 -0500 Subject: [PATCH 028/207] rename to fp8 moe Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 194 ------------------ 1 file changed, 194 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py deleted file mode 100644 index c66a30623a3b..000000000000 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ /dev/null @@ -1,194 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch - -from vllm.model_executor.layers.fused_moe.config import ( - RoutingMethodType, -) -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - calculate_tile_tokens_dim, -) -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - per_token_group_quant_fp8, -) -from vllm.utils.torch_utils import direct_register_custom_op - - -def flashinfer_fused_moe_blockscale_fp8( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - x: torch.Tensor, - w13_weight: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - global_num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - expert_offset: int, - local_num_experts: int, - block_shape: list[int], - routing_method_type: int = RoutingMethodType.DeepSeekV3, - routed_scaling: float | None = 1.0, -) -> torch.Tensor: - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe - - topk_group = topk_group if topk_group is not None else 0 - assert top_k <= global_num_experts - assert top_k <= 10 - assert global_num_experts % 4 == 0 - assert block_shape == [128, 128] - # Routing kernel expects #experts <= #threads 512 - assert global_num_experts <= 512 - - a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1]) - # NOTE: scales of hidden states have to be transposed! - a_sf_t = a_sf.t().contiguous() - return flashinfer_trtllm_fp8_block_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=a_q, - hidden_states_scale=a_sf_t, - gemm1_weights=w13_weight, - gemm1_weights_scale=w13_weight_scale_inv, - gemm2_weights=w2_weight, - gemm2_weights_scale=w2_weight_scale_inv, - num_experts=global_num_experts, - top_k=top_k, - n_group=num_expert_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=expert_offset, - local_num_experts=local_num_experts, - routed_scaling_factor=routed_scaling, - tile_tokens_dim=None, - routing_method_type=routing_method_type, - use_shuffled_weight=False, - ) - - -def flashinfer_fused_moe_blockscale_fp8_fake( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - x: torch.Tensor, - w13_weight: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - global_num_experts: int, - top_k: int, - num_expert_group: int, - topk_group: int, - intermediate_size: int, - expert_offset: int, - local_num_experts: int, - block_shape: list[int], - routing_method_type: int, - routed_scaling: float = 1.0, -) -> torch.Tensor: - return torch.empty_like(x) - - -# TODO(bnell): Does this really need to be a torch.op? -direct_register_custom_op( - op_name="flashinfer_fused_moe_blockscale_fp8", - op_func=flashinfer_fused_moe_blockscale_fp8, - fake_impl=flashinfer_fused_moe_blockscale_fp8_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -def fi_trtllm_fp8_per_tensor_moe( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - hidden_states: torch.Tensor, - input_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm2_weights: torch.Tensor, - output1_scales_scalar: torch.Tensor, - output1_scales_gate_scalar: torch.Tensor, - output2_scales_scalar: torch.Tensor, - num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - use_routing_scales_on_input: bool, - routing_method_type: int, - routed_scaling_factor: float = 1.0, -) -> torch.Tensor: - num_expert_group = num_expert_group if num_expert_group is not None else 0 - topk_group = topk_group if topk_group is not None else 0 - - quant_hidden_states, _ = moe_kernel_quantize_input( - hidden_states, - input_scale, - quant_dtype=torch.float8_e4m3fn, - per_act_token_quant=False, - ) - - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe - - return flashinfer_trtllm_fp8_per_tensor_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=quant_hidden_states, - gemm1_weights=gemm1_weights, - output1_scales_scalar=output1_scales_scalar, - output1_scales_gate_scalar=output1_scales_gate_scalar, - gemm2_weights=gemm2_weights, - output2_scales_scalar=output2_scales_scalar, - num_experts=num_experts, - top_k=top_k, - n_group=num_expert_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=local_expert_offset, - local_num_experts=local_num_experts, - routed_scaling_factor=routed_scaling_factor, - use_routing_scales_on_input=use_routing_scales_on_input, - tile_tokens_dim=calculate_tile_tokens_dim( - hidden_states.shape[0], top_k, num_experts - ), - routing_method_type=routing_method_type, - ) - - -def fi_trtllm_fp8_per_tensor_moe_fake( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - hidden_states: torch.Tensor, - input_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm2_weights: torch.Tensor, - output1_scales_scalar: torch.Tensor, - output1_scales_gate_scalar: torch.Tensor, - output2_scales_scalar: torch.Tensor, - num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - use_routing_scales_on_input: bool, - routing_method_type: int, - routed_scaling_factor: float = 1.0, -) -> torch.Tensor: - return torch.empty_like(hidden_states) - - -# TODO(bnell): Does this really need to be a torch.op? -direct_register_custom_op( - op_name="fi_trtllm_fp8_per_tensor_moe", - op_func=fi_trtllm_fp8_per_tensor_moe, - mutates_args=["hidden_states"], - fake_impl=fi_trtllm_fp8_per_tensor_moe_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) From 3d7abb60cc9e6451ccbbf6a63440f776eaab31e0 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 14:39:49 -0500 Subject: [PATCH 029/207] pre-commit Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 3843c3610626..58cc28a478ce 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -45,6 +45,8 @@ def __init__( # g1_alpha_s = a13_scale * w13_scale_2 # a2_gscale = (1 / a2_scale) # g1_scale_c = a13_scale * w13_scale_2 / a2_scale + assert self.quant_config.g1_alphas is not None + assert self.quant_config.a2_gscale is not None self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale @property From 457bd9d89826ad0f51b5b3986a3b5e5fcd872ed5 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 14:41:14 -0500 Subject: [PATCH 030/207] add back missing file Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 194 ++++++++++++++++++ 1 file changed, 194 insertions(+) create mode 100644 vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py new file mode 100644 index 000000000000..c66a30623a3b --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + calculate_tile_tokens_dim, +) +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, +) +from vllm.utils.torch_utils import direct_register_custom_op + + +def flashinfer_fused_moe_blockscale_fp8( + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + x: torch.Tensor, + w13_weight: torch.Tensor, + w13_weight_scale_inv: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale_inv: torch.Tensor, + global_num_experts: int, + top_k: int, + num_expert_group: int | None, + topk_group: int | None, + intermediate_size: int, + expert_offset: int, + local_num_experts: int, + block_shape: list[int], + routing_method_type: int = RoutingMethodType.DeepSeekV3, + routed_scaling: float | None = 1.0, +) -> torch.Tensor: + from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe + + topk_group = topk_group if topk_group is not None else 0 + assert top_k <= global_num_experts + assert top_k <= 10 + assert global_num_experts % 4 == 0 + assert block_shape == [128, 128] + # Routing kernel expects #experts <= #threads 512 + assert global_num_experts <= 512 + + a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1]) + # NOTE: scales of hidden states have to be transposed! + a_sf_t = a_sf.t().contiguous() + return flashinfer_trtllm_fp8_block_scale_moe( + routing_logits=routing_logits, + routing_bias=routing_bias, + hidden_states=a_q, + hidden_states_scale=a_sf_t, + gemm1_weights=w13_weight, + gemm1_weights_scale=w13_weight_scale_inv, + gemm2_weights=w2_weight, + gemm2_weights_scale=w2_weight_scale_inv, + num_experts=global_num_experts, + top_k=top_k, + n_group=num_expert_group, + topk_group=topk_group, + intermediate_size=intermediate_size, + local_expert_offset=expert_offset, + local_num_experts=local_num_experts, + routed_scaling_factor=routed_scaling, + tile_tokens_dim=None, + routing_method_type=routing_method_type, + use_shuffled_weight=False, + ) + + +def flashinfer_fused_moe_blockscale_fp8_fake( + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + x: torch.Tensor, + w13_weight: torch.Tensor, + w13_weight_scale_inv: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale_inv: torch.Tensor, + global_num_experts: int, + top_k: int, + num_expert_group: int, + topk_group: int, + intermediate_size: int, + expert_offset: int, + local_num_experts: int, + block_shape: list[int], + routing_method_type: int, + routed_scaling: float = 1.0, +) -> torch.Tensor: + return torch.empty_like(x) + + +# TODO(bnell): Does this really need to be a torch.op? +direct_register_custom_op( + op_name="flashinfer_fused_moe_blockscale_fp8", + op_func=flashinfer_fused_moe_blockscale_fp8, + fake_impl=flashinfer_fused_moe_blockscale_fp8_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +def fi_trtllm_fp8_per_tensor_moe( + routing_logits: torch.Tensor, + routing_bias: torch.Tensor | None, + hidden_states: torch.Tensor, + input_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm2_weights: torch.Tensor, + output1_scales_scalar: torch.Tensor, + output1_scales_gate_scalar: torch.Tensor, + output2_scales_scalar: torch.Tensor, + num_experts: int, + top_k: int, + num_expert_group: int | None, + topk_group: int | None, + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + use_routing_scales_on_input: bool, + routing_method_type: int, + routed_scaling_factor: float = 1.0, +) -> torch.Tensor: + num_expert_group = num_expert_group if num_expert_group is not None else 0 + topk_group = topk_group if topk_group is not None else 0 + + quant_hidden_states, _ = moe_kernel_quantize_input( + hidden_states, + input_scale, + quant_dtype=torch.float8_e4m3fn, + per_act_token_quant=False, + ) + + from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe + + return flashinfer_trtllm_fp8_per_tensor_scale_moe( + routing_logits=routing_logits, + routing_bias=routing_bias, + hidden_states=quant_hidden_states, + gemm1_weights=gemm1_weights, + output1_scales_scalar=output1_scales_scalar, + output1_scales_gate_scalar=output1_scales_gate_scalar, + gemm2_weights=gemm2_weights, + output2_scales_scalar=output2_scales_scalar, + num_experts=num_experts, + top_k=top_k, + n_group=num_expert_group, + topk_group=topk_group, + intermediate_size=intermediate_size, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + routed_scaling_factor=routed_scaling_factor, + use_routing_scales_on_input=use_routing_scales_on_input, + tile_tokens_dim=calculate_tile_tokens_dim( + hidden_states.shape[0], top_k, num_experts + ), + routing_method_type=routing_method_type, + ) + + +def fi_trtllm_fp8_per_tensor_moe_fake( + routing_logits: torch.Tensor, + routing_bias: torch.Tensor | None, + hidden_states: torch.Tensor, + input_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm2_weights: torch.Tensor, + output1_scales_scalar: torch.Tensor, + output1_scales_gate_scalar: torch.Tensor, + output2_scales_scalar: torch.Tensor, + num_experts: int, + top_k: int, + num_expert_group: int | None, + topk_group: int | None, + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + use_routing_scales_on_input: bool, + routing_method_type: int, + routed_scaling_factor: float = 1.0, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +# TODO(bnell): Does this really need to be a torch.op? +direct_register_custom_op( + op_name="fi_trtllm_fp8_per_tensor_moe", + op_func=fi_trtllm_fp8_per_tensor_moe, + mutates_args=["hidden_states"], + fake_impl=fi_trtllm_fp8_per_tensor_moe_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) From aa92d528cb93ecc18d5f429189753f269cdc07fc Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 15:22:40 -0500 Subject: [PATCH 031/207] add back missing file Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 149 ++++++++++++++++++ .../layers/fused_moe/oracle/fp8.py | 20 ++- .../model_executor/layers/quantization/fp8.py | 90 +++++------ .../quantization/utils/flashinfer_fp4_moe.py | 100 ++++++++++++ .../quantization/utils/flashinfer_utils.py | 2 +- 5 files changed, 310 insertions(+), 51 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index c66a30623a3b..72cda6a2660a 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import flashinfer import torch +import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input @@ -192,3 +196,148 @@ def fi_trtllm_fp8_per_tensor_moe_fake( fake_impl=fi_trtllm_fp8_per_tensor_moe_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) + + +class FlashInferTrtLlmFp8Experts(mk.FusedMoEPermuteExpertsUnpermute): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(quant_config) + + self.moe_config = moe_config + # TODO: set this via the constructor + self.routing_method_type = flashinfer.RoutingMethodType.Renormalize + # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 + # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 + + self.routing_bias = None + # TODO: to: in_dtype.shape + self.e_score_correction_bias = None + self.topk_group = None + self.num_expert_group = None + self.routing_scaling_factor = None + + self.topk = moe_config.experts_per_token + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.hidden_dim = moe_config.hidden_dim + self.local_num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + @property + def activation_formats( + self, + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]: + return ( + mk.FusedMoEActivationFormat.Standard, + mk.FusedMoEActivationFormat.Standard, + ) + + def supports_chunking(self) -> bool: + return False + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + raise NotImplementedError( + f"{self.__class__.__name__} only supports the apply_monolithic interface." + ) + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: str, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + raise NotImplementedError( + f"{self.__class__.__name__} only supports the apply_monolithic interface." + ) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + raise NotImplementedError( + f"{self.__class__.__name__} only supports the apply_monolithic interface." + ) + + def apply_monolithic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + assert activation == "silu" + assert ( + self.e_score_correction_bias is None + or self.e_score_correction_bias.dtype == hidden_states.dtype + ) + + if self.routing_method_type == RoutingMethodType.DeepSeekV3: + router_logits = router_logits.to(torch.float32) + + topk_group = self.topk_group if self.topk_group is not None else 0 + + assert self.topk <= global_num_experts + assert self.topk <= 10 + assert global_num_experts % 4 == 0 + assert self.quant_config.block_shape == [128, 128] + # Routing kernel expects #experts <= #threads 512 + assert global_num_experts <= 512 + + a1_q, a1q_scale = per_token_group_quant_fp8( + hidden_states, self.quant_config.block_shape[1] + ) + # Kernel requires transposed hidden state scales + # TODO: can we avoid this transpose in the kernel? + a1q_scale_t = a1q_scale.t().contiguous() + + return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( + routing_logits=router_logits, + routing_bias=self.e_score_correction_bias, + hidden_states=a1_q, + hidden_states_scale=a1q_scale_t, + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale, + num_experts=global_num_experts, + top_k=self.topk, + n_group=self.num_expert_group, + topk_group=topk_group, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=self.routing_scaling_factor, + tile_tokens_dim=None, + routing_method_type=self.routing_method_type, + use_shuffled_weight=False, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 6872b542f492..b989a96e0a42 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -210,7 +210,7 @@ def make_fp8_moe_quant_config( a1_scale: torch.Tensor | None, a2_scale: torch.Tensor | None, block_shape: list[int] | None = None, -) -> FusedMoEQuantConfig | None: +) -> FusedMoEQuantConfig: """ Create FusedMoEQuantConfig for the specifed FP8 Backend. The FusedMoEQuantConfig holds the scales that are used @@ -223,9 +223,6 @@ def make_fp8_moe_quant_config( In a future PR, we will have this function should be a method of the modular kernel itself. """ - # TRTLLM does not use Modular Kernel abstraction yet. - if fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM: - return None # MARLIN is mixed precision W8A16 config. if fp8_backend == Fp8MoeBackend.MARLIN: @@ -284,7 +281,20 @@ def make_fp8_moe_kernel( # via the same code path (i.e. via maybe_init_modular_kernel). # NOTE(rob): in progress migrating all into this format. use_inplace = True - if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS: + if fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM: + from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( + FlashInferTrtLlmFp8Experts, + ) + + kernel = mk.FusedMoEModularKernel( + MoEPrepareAndFinalizeNoEP(), + FlashInferTrtLlmFp8Experts( + moe_config=moe_config, + quant_config=moe_quant_config, + ), + ) + + elif fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS: from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, ) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 14ed28630680..ca1537472b69 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -28,7 +28,6 @@ ) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, - RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( @@ -819,13 +818,12 @@ def _setup_kernel( # Setup modular kernel for TP case. self.moe_quant_config = self.get_fused_moe_quant_config(layer) - if self.moe_quant_config: - self.kernel, self.use_inplace = make_fp8_moe_kernel( - layer=layer, - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - fp8_backend=self.fp8_backend, - ) + self.kernel, self.use_inplace = make_fp8_moe_kernel( + layer=layer, + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + ) def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): @@ -967,13 +965,7 @@ def select_gemm_impl( ) return TritonExperts(self.moe_quant_config) - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - # TRTLLM does not use Modular Kernel. - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM: - return None - + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") a1_scale = layer.w13_input_scale @@ -1007,40 +999,48 @@ def apply( # TODO(rob): convert this to MK. if layer.enable_eplb: raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.") - assert layer.activation == "silu", ( - f"Expected 'silu' activation but got {layer.activation}" - ) if self.block_quant: - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401 - - e_score_correction_bias = ( - layer.e_score_correction_bias.to(x.dtype) - if layer.e_score_correction_bias is not None - else None - ) - routing_method_type = layer.routing_method_type - return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8( - routing_logits=router_logits.to(torch.float32) - if routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits, - routing_bias=e_score_correction_bias, - x=x, - w13_weight=layer.w13_weight, - w13_weight_scale_inv=layer.w13_weight_scale_inv, - w2_weight=layer.w2_weight, - w2_weight_scale_inv=layer.w2_weight_scale_inv, + assert self.kernel is not None + return self.kernel.forward_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, global_num_experts=layer.global_num_experts, - top_k=layer.top_k, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - intermediate_size=layer.intermediate_size_per_partition, - expert_offset=layer.ep_rank * layer.local_num_experts, - local_num_experts=layer.local_num_experts, - block_shape=self.weight_block_size, - routing_method_type=routing_method_type, - routed_scaling=layer.routed_scaling_factor, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, ) + # import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe # noqa: E501, F401 + + # e_score_correction_bias = ( + # layer.e_score_correction_bias.to(x.dtype) + # if layer.e_score_correction_bias is not None + # else None + # ) + # routing_method_type = layer.routing_method_type + # return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8( + # routing_logits=router_logits.to(torch.float32) + # if routing_method_type == RoutingMethodType.DeepSeekV3 + # else router_logits, + # routing_bias=e_score_correction_bias, + # x=x, + # w13_weight=layer.w13_weight, + # w13_weight_scale_inv=layer.w13_weight_scale_inv, + # w2_weight=layer.w2_weight, + # w2_weight_scale_inv=layer.w2_weight_scale_inv, + # global_num_experts=layer.global_num_experts, + # top_k=layer.top_k, + # num_expert_group=layer.num_expert_group, + # topk_group=layer.topk_group, + # intermediate_size=layer.intermediate_size_per_partition, + # expert_offset=layer.ep_rank * layer.local_num_experts, + # local_num_experts=layer.local_num_experts, + # block_shape=self.weight_block_size, + # routing_method_type=routing_method_type, + # routed_scaling=layer.routed_scaling_factor, + # ) else: result = apply_fi_trtllm_fp8_per_tensor_moe( layer=layer, diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 4c06dc1b7b27..4b9dd2b10963 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility helpers for NVFP4 + FlashInfer fused-MoE path""" +from typing import TYPE_CHECKING + import torch import vllm.envs as envs @@ -20,12 +22,20 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501 create_flashinfer_prepare_finalize, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + swizzle_blockscale, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( has_flashinfer_cutedsl_grouped_gemm_nt_masked, has_flashinfer_cutlass_fused_moe, ) +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + ) + logger = init_logger(__name__) @@ -237,3 +247,93 @@ def prepare_static_weights_for_trtllm_fp4_moe( gemm2_weights_fp4_shuffled, gemm2_scales_fp4_shuffled, ) + + +def prepare_nvfp4_moe_layer_for_fi_or_cutlass( + backend: "NvFp4MoeBackend", + layer: torch.nn.Module, + w13: torch.Tensor, + w13_scale: torch.Tensor, + w13_scale_2: torch.Tensor, + a13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + w2_scale_2: torch.Tensor, + a2_scale: torch.Tensor, + is_act_and_mul: bool, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + # Delayed import for circular dependency avoidance. + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + is_global_sf_supported_for_nvfp4_backend, + ) + + assert backend in [ + NvFp4MoeBackend.VLLM_CUTLASS, + NvFp4MoeBackend.FLASHINFER_CUTLASS, + NvFp4MoeBackend.FLASHINFER_TRTLLM, + NvFp4MoeBackend.FLASHINFER_CUTEDSL, + ] + + # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels. + if is_act_and_mul and backend in [ + NvFp4MoeBackend.FLASHINFER_CUTLASS, + NvFp4MoeBackend.FLASHINFER_TRTLLM, + ]: + w13, w13_scale = reorder_w1w3_to_w3w1(w13, w13_scale) + + # For some FI kernels, the input scales are shared by all experts. + if is_global_sf_supported_for_nvfp4_backend(backend): + num_experts = w13.shape[0] + a13_scale = a13_scale.max().to(torch.float32).expand(num_experts) + a2_scale = a2_scale.max().to(torch.float32).expand(num_experts) + else: + a13_scale = a13_scale.max(dim=1).values.to(torch.float32) + + # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels. + if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: + w13, w13_scale, w2, w2_scale = prepare_static_weights_for_trtllm_fp4_moe( + w13, + w2, + w13_scale, + w2_scale, + w2.size(-2), # hidden_size + w13.size(-2) // 2, # intermediate_size + w13.size(0), # num_experts + ) + + # We do not need to make this a parameter, because + # it is not used during the weight (re)-loading process. + layer.g1_scale_c = a13_scale * w13_scale_2 / a2_scale + layer.a1_gscale = 1.0 / a13_scale + layer.g1_alphas = a13_scale * w13_scale_2 + layer.g2_alphas = a2_scale * w2_scale_2 + else: + # Swizzle the block scales for other FI NVFP4 MoE kernels. + w13_scale = swizzle_blockscale(w13_scale) + + # Apply padding if needed. + pad_size = w13_scale.size(1) - w13.size(1) + if pad_size > 0: + if is_act_and_mul: + raise NotImplementedError( + "Intermediate size padding for w1 and w3, for %s " + "NvFp4 backend, but this is not currently supported", + backend.value, + ) + w13 = torch.nn.functional.pad(w13, (0, 0, 0, pad_size)) + w2 = torch.nn.functional.pad(w2, (0, pad_size // 2, 0, 0)) + w2_scale = torch.nn.functional.pad(w2_scale, (0, pad_size // 16)) + + w2_scale = swizzle_blockscale(w2_scale) + + return w13, w13_scale, w13_scale_2, a13_scale, w2, w2_scale, w2_scale_2, a2_scale diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 799854479823..bfab3d70b781 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -138,7 +138,7 @@ def apply_fi_trtllm_fp8_per_tensor_moe( ) -> torch.Tensor: from flashinfer.fused_moe import RoutingMethodType - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401 + import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe # noqa: E501, F401 from vllm.model_executor.models.llama4 import Llama4MoE # Added to the layer by: register_scales_for_trtllm_fp8_per_tensor_moe From 8abae8c8db00ed51a8807ec077b35d55aa08a789 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 15:26:48 -0500 Subject: [PATCH 032/207] able to launch with fp8 Signed-off-by: Robert Shaw --- .../model_executor/layers/quantization/fp8.py | 29 ------------------- 1 file changed, 29 deletions(-) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index ca1537472b69..495f57981405 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -1012,35 +1012,6 @@ def apply( expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, ) - # import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe # noqa: E501, F401 - - # e_score_correction_bias = ( - # layer.e_score_correction_bias.to(x.dtype) - # if layer.e_score_correction_bias is not None - # else None - # ) - # routing_method_type = layer.routing_method_type - # return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8( - # routing_logits=router_logits.to(torch.float32) - # if routing_method_type == RoutingMethodType.DeepSeekV3 - # else router_logits, - # routing_bias=e_score_correction_bias, - # x=x, - # w13_weight=layer.w13_weight, - # w13_weight_scale_inv=layer.w13_weight_scale_inv, - # w2_weight=layer.w2_weight, - # w2_weight_scale_inv=layer.w2_weight_scale_inv, - # global_num_experts=layer.global_num_experts, - # top_k=layer.top_k, - # num_expert_group=layer.num_expert_group, - # topk_group=layer.topk_group, - # intermediate_size=layer.intermediate_size_per_partition, - # expert_offset=layer.ep_rank * layer.local_num_experts, - # local_num_experts=layer.local_num_experts, - # block_shape=self.weight_block_size, - # routing_method_type=routing_method_type, - # routed_scaling=layer.routed_scaling_factor, - # ) else: result = apply_fi_trtllm_fp8_per_tensor_moe( layer=layer, From f195aeaeaa0fc602e47daddd385d0de44373a7f4 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 15:33:16 -0500 Subject: [PATCH 033/207] remove the custom op for fp8 block Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 86 ------------------- 1 file changed, 86 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 72cda6a2660a..92549e507be8 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -20,92 +20,6 @@ from vllm.utils.torch_utils import direct_register_custom_op -def flashinfer_fused_moe_blockscale_fp8( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - x: torch.Tensor, - w13_weight: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - global_num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - expert_offset: int, - local_num_experts: int, - block_shape: list[int], - routing_method_type: int = RoutingMethodType.DeepSeekV3, - routed_scaling: float | None = 1.0, -) -> torch.Tensor: - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe - - topk_group = topk_group if topk_group is not None else 0 - assert top_k <= global_num_experts - assert top_k <= 10 - assert global_num_experts % 4 == 0 - assert block_shape == [128, 128] - # Routing kernel expects #experts <= #threads 512 - assert global_num_experts <= 512 - - a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1]) - # NOTE: scales of hidden states have to be transposed! - a_sf_t = a_sf.t().contiguous() - return flashinfer_trtllm_fp8_block_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=a_q, - hidden_states_scale=a_sf_t, - gemm1_weights=w13_weight, - gemm1_weights_scale=w13_weight_scale_inv, - gemm2_weights=w2_weight, - gemm2_weights_scale=w2_weight_scale_inv, - num_experts=global_num_experts, - top_k=top_k, - n_group=num_expert_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=expert_offset, - local_num_experts=local_num_experts, - routed_scaling_factor=routed_scaling, - tile_tokens_dim=None, - routing_method_type=routing_method_type, - use_shuffled_weight=False, - ) - - -def flashinfer_fused_moe_blockscale_fp8_fake( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - x: torch.Tensor, - w13_weight: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - global_num_experts: int, - top_k: int, - num_expert_group: int, - topk_group: int, - intermediate_size: int, - expert_offset: int, - local_num_experts: int, - block_shape: list[int], - routing_method_type: int, - routed_scaling: float = 1.0, -) -> torch.Tensor: - return torch.empty_like(x) - - -# TODO(bnell): Does this really need to be a torch.op? -direct_register_custom_op( - op_name="flashinfer_fused_moe_blockscale_fp8", - op_func=flashinfer_fused_moe_blockscale_fp8, - fake_impl=flashinfer_fused_moe_blockscale_fp8_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - def fi_trtllm_fp8_per_tensor_moe( routing_logits: torch.Tensor, routing_bias: torch.Tensor | None, From 2196a0143d7e182dcc6f3c3841260d9220a6dbdb Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 17:15:00 -0500 Subject: [PATCH 034/207] scaffolding of per-tensor kernel Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 104 +++++++++++++++++- vllm/model_executor/models/llama4.py | 1 + 2 files changed, 101 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 92549e507be8..64f1900f64ab 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -122,8 +122,8 @@ def __init__( self.moe_config = moe_config # TODO: set this via the constructor - self.routing_method_type = flashinfer.RoutingMethodType.Renormalize - # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 + # self.routing_method_type = flashinfer.RoutingMethodType.Renormalize + self.routing_method_type = flashinfer.RoutingMethodType.Llama4 # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 self.routing_bias = None @@ -141,6 +141,18 @@ def __init__( self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + make_fp8_moe_alpha_scales_for_fi, + ) + + self.g1_alphas, self.g2_alphas = make_fp8_moe_alpha_scales_for_fi( + w13_scale=self.quant_config.w1_scale, + w13_input_scale=self.quant_config.a1_scale, + w2_scale=self.quant_config.w2_scale, + w2_input_scale=self.quant_config.a2_scale, + ) + self.g1_scale_c = self.g1_alphas / self.quant_config.a2_scale + @property def activation_formats( self, @@ -198,7 +210,7 @@ def apply( f"{self.__class__.__name__} only supports the apply_monolithic interface." ) - def apply_monolithic( + def _apply_per_block_monolithic( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -231,7 +243,7 @@ def apply_monolithic( hidden_states, self.quant_config.block_shape[1] ) # Kernel requires transposed hidden state scales - # TODO: can we avoid this transpose in the kernel? + # TODO: fuse into the quant kernel. a1q_scale_t = a1q_scale.t().contiguous() return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( @@ -255,3 +267,87 @@ def apply_monolithic( routing_method_type=self.routing_method_type, use_shuffled_weight=False, ) + + def _apply_per_tensor_monolithic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + assert self.routing_method_type == RoutingMethodType.Llama4 + assert apply_router_weight_on_input + + a1q, _ = moe_kernel_quantize_input( + hidden_states, + self.quant_config.a1_scale, + quant_dtype=self.quant_config.quant_dtype, + per_act_token_quant=self.quant_config.per_act_token_quant, + ) + + return flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( + routing_logits=router_logits, + routing_bias=self.e_score_correction_bias, + hidden_states=a1q, + gemm1_weights=w1, + output1_scales_scalar=self.g1_alphas, + output1_scales_gate_scalar=self.g1_scale_c, + gemm2_weights=w2, + output2_scales_scalar=self.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + num_expert_group=self.num_expert_group, + topk_group=self.topk_group, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=self.routing_scaling_factor, + use_routing_scales_on_input=apply_router_weight_on_input, + tile_tokens_dim=calculate_tile_tokens_dim( + hidden_states.shape[0], self.topk, self.local_num_experts + ), + routing_method_type=self.routing_method_type, + ) + + def apply_monolithic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + if self.quant_config.block_shape is not None: + return self._apply_per_block_monolithic( + hidden_states, + w1, + w2, + router_logits, + activation, + global_num_experts, + expert_map, + apply_router_weight_on_input, + ) + elif self.quant_config.is_per_tensor: + return self._apply_per_tensor_monolithic( + hidden_states, + w1, + w2, + router_logits, + activation, + global_num_experts, + expert_map, + apply_router_weight_on_input, + ) + else: + raise NotImplementedError( + "Only per-block and per-tensor quantization are supported in " + f"{self.__class__.__name__}." + ) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index dde6db7c204b..7cc170638428 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -434,6 +434,7 @@ def load_moe_expert_weights( # Whether the MoE expert weights are loaded successfully. expert_param_loaded = False + loaded_weight = loaded_weight.to("cuda") # If fused is True, the loaded weight is in the layout of: # [num_experts, hidden_in, hidden_out], so we must transpose the last From 42a328efd3704d1f2b8f1ab8520cb718dd4451f9 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 17:28:19 -0500 Subject: [PATCH 035/207] basic poc with llama scout modelopt Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_fp8_moe.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 64f1900f64ab..3bccc21206f3 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -145,13 +145,13 @@ def __init__( make_fp8_moe_alpha_scales_for_fi, ) - self.g1_alphas, self.g2_alphas = make_fp8_moe_alpha_scales_for_fi( + self._g1_alphas, self._g2_alphas = make_fp8_moe_alpha_scales_for_fi( w13_scale=self.quant_config.w1_scale, w13_input_scale=self.quant_config.a1_scale, w2_scale=self.quant_config.w2_scale, w2_input_scale=self.quant_config.a2_scale, ) - self.g1_scale_c = self.g1_alphas / self.quant_config.a2_scale + self.g1_scale_c = self._g1_alphas / self.quant_config.a2_scale @property def activation_formats( @@ -294,10 +294,10 @@ def _apply_per_tensor_monolithic( routing_bias=self.e_score_correction_bias, hidden_states=a1q, gemm1_weights=w1, - output1_scales_scalar=self.g1_alphas, + output1_scales_scalar=self._g1_alphas, output1_scales_gate_scalar=self.g1_scale_c, gemm2_weights=w2, - output2_scales_scalar=self.g2_alphas, + output2_scales_scalar=self._g2_alphas, num_experts=global_num_experts, top_k=self.topk, num_expert_group=self.num_expert_group, From e814e52c972510bc1b5daa1776319f306c14c5a6 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 18:27:24 -0500 Subject: [PATCH 036/207] fix import Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 24 ++++++++---- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 16 ++++---- .../model_executor/layers/quantization/fp8.py | 37 ++++++------------- .../layers/quantization/modelopt.py | 36 +++++++----------- 4 files changed, 49 insertions(+), 64 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index bb2f6b873941..828397794754 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -19,7 +19,6 @@ MoEPrepareAndFinalizeNoEP, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - apply_fi_trtllm_fp8_per_tensor_moe, register_scales_for_trtllm_fp8_per_tensor_moe, rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, swap_w13_to_w31, @@ -207,15 +206,26 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( quant_config=quant_config, ) - flashinfer_output = apply_fi_trtllm_fp8_per_tensor_moe( - layer=td.layer, + from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( + FlashInferTrtLlmFp8Experts, + ) + + kernel = mk.FusedMoEModularKernel( + MoEPrepareAndFinalizeNoEP(), + FlashInferTrtLlmFp8Experts( + moe_config=td.layer.moe, + quant_config=quant_config, + ), + ) + + flashinfer_output = kernel.apply_monolithic( hidden_states=td.hidden_states, + w1=td.layer.w13_weight, + w2=td.layer.w2_weight, router_logits=score, - routing_bias=None, + activation="silu", global_num_experts=e, - top_k=topk, - num_expert_group=None, - topk_group=None, + expert_map=None, apply_router_weight_on_input=True, ) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 3bccc21206f3..1309fc3b88ae 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( calculate_tile_tokens_dim, + make_fp8_moe_alpha_scales_for_fi, ) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, @@ -141,10 +142,6 @@ def __init__( self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank - from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - make_fp8_moe_alpha_scales_for_fi, - ) - self._g1_alphas, self._g2_alphas = make_fp8_moe_alpha_scales_for_fi( w13_scale=self.quant_config.w1_scale, w13_input_scale=self.quant_config.a1_scale, @@ -265,7 +262,6 @@ def _apply_per_block_monolithic( routed_scaling_factor=self.routing_scaling_factor, tile_tokens_dim=None, routing_method_type=self.routing_method_type, - use_shuffled_weight=False, ) def _apply_per_tensor_monolithic( @@ -282,6 +278,8 @@ def _apply_per_tensor_monolithic( assert self.routing_method_type == RoutingMethodType.Llama4 assert apply_router_weight_on_input + topk_group = self.topk_group if self.topk_group is not None else 0 + a1q, _ = moe_kernel_quantize_input( hidden_states, self.quant_config.a1_scale, @@ -294,14 +292,14 @@ def _apply_per_tensor_monolithic( routing_bias=self.e_score_correction_bias, hidden_states=a1q, gemm1_weights=w1, - output1_scales_scalar=self._g1_alphas, - output1_scales_gate_scalar=self.g1_scale_c, + output1_scales_scalar=self.g1_scale_c, + output1_scales_gate_scalar=self._g1_alphas, gemm2_weights=w2, output2_scales_scalar=self._g2_alphas, num_experts=global_num_experts, top_k=self.topk, - num_expert_group=self.num_expert_group, - topk_group=self.topk_group, + n_group=self.num_expert_group, + topk_group=topk_group, intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 495f57981405..9f3308fdc8b6 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -49,7 +49,6 @@ ) from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - apply_fi_trtllm_fp8_per_tensor_moe, build_flashinfer_fp8_cutlass_moe_prepare_finalize, select_cutlass_fp8_gemm_impl, ) @@ -996,34 +995,20 @@ def apply( router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM: - # TODO(rob): convert this to MK. if layer.enable_eplb: raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.") - if self.block_quant: - assert self.kernel is not None - return self.kernel.forward_monolithic( - x, - layer.w13_weight, - layer.w2_weight, - router_logits, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - ) - else: - result = apply_fi_trtllm_fp8_per_tensor_moe( - layer=layer, - hidden_states=x, - router_logits=router_logits, - routing_bias=layer.e_score_correction_bias, - global_num_experts=layer.global_num_experts, - top_k=layer.top_k, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - ) + assert self.kernel is not None + return self.kernel.forward_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ) topk_weights, topk_ids = router.select_experts( hidden_states=x, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 4ffb434e58dc..13c741cbb553 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -55,7 +55,6 @@ select_nvfp4_gemm_impl, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - apply_fi_trtllm_fp8_per_tensor_moe, build_flashinfer_fp8_cutlass_moe_prepare_finalize, select_cutlass_fp8_gemm_impl, ) @@ -877,13 +876,12 @@ def _setup_kernel( # Setup modular kernel for TP case. self.moe_quant_config = self.get_fused_moe_quant_config(layer) - if self.moe_quant_config: - self.kernel, self.use_inplace = make_fp8_moe_kernel( - layer=layer, - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - fp8_backend=self.fp8_backend, - ) + self.kernel, self.use_inplace = make_fp8_moe_kernel( + layer=layer, + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w13 = layer.w13_weight @@ -944,21 +942,15 @@ def apply( raise NotImplementedError( "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend." ) - # TODO(rob): this validation should happen at kernel selection - # time in the oracle rather than here. - assert layer.activation == "silu", ( - f"Expected 'silu' activation but got {layer.activation}" - ) - assert not layer.renormalize - return apply_fi_trtllm_fp8_per_tensor_moe( - layer=layer, - hidden_states=x, - router_logits=router_logits, - routing_bias=layer.e_score_correction_bias, + assert self.kernel is not None + return self.kernel.forward_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, global_num_experts=layer.global_num_experts, - top_k=layer.top_k, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, + expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, ) From c0ce7548741590941bb3bfd5c9a30c313265ced2 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 18:32:29 -0500 Subject: [PATCH 037/207] messing around with typrs Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 828397794754..7bdc87450201 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -8,6 +8,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, FusedMoEQuantConfig, fp8_w8a8_moe_quant_config, ) @@ -89,7 +90,13 @@ class TestData: @staticmethod def make_moe_tensors_8bit( - m: int, k: int, n: int, e: int, is_trtllm: bool, activation: str = "silu" + m: int, + k: int, + n: int, + e: int, + is_trtllm: bool, + activation: str = "silu", + topk: int = 1, ) -> "TestData": is_gated = activation != "relu2_no_mul" @@ -146,6 +153,17 @@ def make_moe_tensors_8bit( layer.ep_rank = 0 layer.local_num_experts = e + layer.moe = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size_per_partition=n, + num_local_experts=e, + moe_parallel_config=layer.moe_parallel_config, + in_dtype=hidden_states.dtype, + is_act_and_mul=is_gated, + ) + return TestData( hidden_states=hidden_states, w13_quantized=w13_quantized, @@ -218,7 +236,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( ), ) - flashinfer_output = kernel.apply_monolithic( + flashinfer_output = kernel.forward_monolithic( hidden_states=td.hidden_states, w1=td.layer.w13_weight, w2=td.layer.w2_weight, From ff99bbebddcd7fe46b37bea0ea91e6723f491913 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 19 Jan 2026 18:33:36 -0500 Subject: [PATCH 038/207] remove stray namings Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 8 -- .../quantization/utils/flashinfer_utils.py | 84 +------------------ 2 files changed, 1 insertion(+), 91 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 7bdc87450201..b15a5952f127 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -20,7 +20,6 @@ MoEPrepareAndFinalizeNoEP, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - register_scales_for_trtllm_fp8_per_tensor_moe, rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, swap_w13_to_w31, ) @@ -141,13 +140,6 @@ def make_moe_tensors_8bit( rotate_weights_for_fi_trtllm_fp8_per_tensor_moe( layer.w13_weight, layer.w2_weight ) - register_scales_for_trtllm_fp8_per_tensor_moe( - layer, - layer.w13_weight_scale, - layer.w13_input_scale, - layer.w2_weight_scale, - layer.w2_input_scale, - ) layer.custom_routing_function = Llama4MoE.custom_routing_function layer.intermediate_size_per_partition = n layer.ep_rank = 0 diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index bfab3d70b781..5b7296df3642 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -105,80 +105,6 @@ def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe( ) -def register_scales_for_trtllm_fp8_per_tensor_moe( - layer: torch.nn.Module, - w13_scale: torch.Tensor, - w13_input_scale: torch.Tensor, - w2_scale: torch.Tensor, - w2_input_scale: torch.Tensor, -) -> None: - """Register necessary scales for FlashInfer TRTLLM FP8 MoE kernel""" - g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi( - w13_scale=w13_scale, - w13_input_scale=w13_input_scale, - w2_scale=w2_scale, - w2_input_scale=w2_input_scale, - ) - layer.w2_input_scale_inv = 1.0 / w2_input_scale - layer.output1_scales_gate_scalar = g1_alphas - layer.output1_scales_scalar = g1_alphas * layer.w2_input_scale_inv - layer.output2_scales_scalar = g2_alphas - - -def apply_fi_trtllm_fp8_per_tensor_moe( - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - global_num_experts: int, - apply_router_weight_on_input: bool, -) -> torch.Tensor: - from flashinfer.fused_moe import RoutingMethodType - - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe # noqa: E501, F401 - from vllm.model_executor.models.llama4 import Llama4MoE - - # Added to the layer by: register_scales_for_trtllm_fp8_per_tensor_moe - assert ( - hasattr(layer, "output1_scales_scalar") - and hasattr(layer, "output1_scales_gate_scalar") - and hasattr(layer, "output2_scales_scalar") - ) - - # Added to the layer by: register_scales_for_trtllm_fp8_per_tensor_moe - assert ( - hasattr(layer, "output1_scales_scalar") - and hasattr(layer, "output1_scales_gate_scalar") - and hasattr(layer, "output2_scales_scalar") - ) - - is_llama4 = layer.custom_routing_function == Llama4MoE.custom_routing_function - assert is_llama4, "FusedMoE flashinfer kernels are only supported for Llama4" - return torch.ops.vllm.fi_trtllm_fp8_per_tensor_moe( - routing_logits=router_logits, - routing_bias=routing_bias, - hidden_states=hidden_states, - input_scale=layer.w13_input_scale, - gemm1_weights=layer.w13_weight, - gemm2_weights=layer.w2_weight, - output1_scales_scalar=layer.output1_scales_scalar, - output1_scales_gate_scalar=layer.output1_scales_gate_scalar, - output2_scales_scalar=layer.output2_scales_scalar, - num_experts=global_num_experts, - top_k=top_k, - num_expert_group=num_expert_group, - topk_group=topk_group, - intermediate_size=layer.intermediate_size_per_partition, - local_expert_offset=layer.ep_rank * layer.local_num_experts, - local_num_experts=layer.local_num_experts, - use_routing_scales_on_input=apply_router_weight_on_input, - routing_method_type=RoutingMethodType.Llama4, - ) - - def make_fp8_moe_alpha_scales_for_fi( w13_scale: torch.Tensor, w13_input_scale: torch.Tensor, @@ -352,19 +278,11 @@ def prepare_fp8_moe_layer_for_fi( w13_scale = swap_w13_to_w31(w13_scale) # FI TRT-LLM FP8 per-tensor MoE kernel requires weight shuffle - # and registration of alpha scales. Note that we do not register - # as nn.Parameters since they are not needed for weight-reloading. + # and registration of alpha scales. if is_trtllm and not block_quant: assert w13_input_scale is not None assert w2_input_scale is not None rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(w13, w2) - register_scales_for_trtllm_fp8_per_tensor_moe( - layer, - w13_scale=w13_scale, - w13_input_scale=w13_input_scale, - w2_scale=w2_scale, - w2_input_scale=w2_input_scale, - ) return w13, w2, w13_scale From 0532d55192c3773b7841356a5a05831027269b7b Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 25 Jan 2026 19:26:09 -0500 Subject: [PATCH 039/207] stash Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 7 --- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 62 ++++++++++++++++--- 2 files changed, 54 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 1309fc3b88ae..db6b7eac417c 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -12,7 +12,6 @@ ) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - calculate_tile_tokens_dim, make_fp8_moe_alpha_scales_for_fi, ) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( @@ -72,9 +71,6 @@ def fi_trtllm_fp8_per_tensor_moe( local_num_experts=local_num_experts, routed_scaling_factor=routed_scaling_factor, use_routing_scales_on_input=use_routing_scales_on_input, - tile_tokens_dim=calculate_tile_tokens_dim( - hidden_states.shape[0], top_k, num_experts - ), routing_method_type=routing_method_type, ) @@ -305,9 +301,6 @@ def _apply_per_tensor_monolithic( local_num_experts=self.local_num_experts, routed_scaling_factor=self.routing_scaling_factor, use_routing_scales_on_input=apply_router_weight_on_input, - tile_tokens_dim=calculate_tile_tokens_dim( - hidden_states.shape[0], self.topk, self.local_num_experts - ), routing_method_type=self.routing_method_type, ) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 58cc28a478ce..7126512e7e9f 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -7,12 +7,19 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, + FusedMoEParallelConfig, FusedMoEQuantConfig, RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.platforms import current_platform class FlashInferTrtLlmNvFp4Experts(mk.FusedMoEPermuteExpertsUnpermute): @@ -49,14 +56,53 @@ def __init__( assert self.quant_config.a2_gscale is not None self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale - @property - def activation_formats( - self, - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]: - return ( - mk.FusedMoEActivationFormat.Standard, - mk.FusedMoEActivationFormat.Standard, - ) + @staticmethod + def _supports_current_device() -> bool: + """Supports only Blackwell-family GPUs.""" + p = current_platform + return p.is_cuda() and p.is_device_capability_family(100) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + """Does not support non-gated MoE (i.e. Nemotron-Nano).""" + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Supports Nvfp4 quantization.""" + SUPPORTED_W_A = [ + (kNvfp4Static, kNvfp4Dynamic), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: str) -> bool: + """Supports only SiLU activation.""" + return activation in ["silu"] + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """Supports EP and TP.""" + return True + + @staticmethod + def _supports_routing_method( + routing_method_type: RoutingMethodType, + ) -> bool: + # NOTE(rob): this is a conservative list. + return routing_method_type in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.Llama4, + ] + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard def supports_chunking(self) -> bool: return False From a6258a7bd0c4cb5de60cf4e9e827733dc0a3938a Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 25 Jan 2026 21:17:36 -0500 Subject: [PATCH 040/207] nitA Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/config.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 89ed39982c8e..9a28c3193587 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1061,7 +1061,6 @@ class FusedMoEConfig: experts_per_token: int hidden_dim: int intermediate_size_per_partition: int - num_local_experts: int activation: str device: torch.device | str From 2e5c741fc5eb0ba2c7739ba89fca5ad85601061d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 25 Jan 2026 21:18:24 -0500 Subject: [PATCH 041/207] pre-commit Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_fp8_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index db6b7eac417c..abd2fc0cc957 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -115,7 +115,7 @@ def __init__( moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, ): - super().__init__(quant_config) + super().__init__(moe_config, quant_config) self.moe_config = moe_config # TODO: set this via the constructor From a1abf66d6ffc88ffcf55d72829022b3c151dee8c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 25 Jan 2026 21:47:26 -0500 Subject: [PATCH 042/207] fix typing Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 74 ++++- .../layers/fused_moe/flashinfer_trtllm_moe.py | 289 ------------------ .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 4 +- .../layers/fused_moe/modular_kernel.py | 18 ++ 4 files changed, 87 insertions(+), 298 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index abd2fc0cc957..b4edd7d4b21a 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -7,6 +7,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, + FusedMoEParallelConfig, FusedMoEQuantConfig, RoutingMethodType, ) @@ -17,7 +18,14 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + kFp8StaticTensorSym, +) from vllm.utils.torch_utils import direct_register_custom_op +from vllm.v1.engine.utils import current_platform def fi_trtllm_fp8_per_tensor_moe( @@ -146,14 +154,64 @@ def __init__( ) self.g1_scale_c = self._g1_alphas / self.quant_config.a2_scale - @property - def activation_formats( - self, - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]: - return ( - mk.FusedMoEActivationFormat.Standard, - mk.FusedMoEActivationFormat.Standard, - ) + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + """Supports only Blackwell-family GPUs.""" + p = current_platform + # Add check flashinfer trtllm is available + return p.is_cuda() and p.is_device_capability_family(100) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + """Does not support non-gated MoE (i.e. Nanotron-Mini).""" + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Supports Fp8 per-tensor and Fp8 block.""" + SUPPORTED_W_A = [ + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: str) -> bool: + """Supports silu activation only.""" + return activation in ["silu"] + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Monolithic kernels need to express router support.""" + if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): + # NOTE(rob): potentially allow others here. This is a conservative list. + return routing_method in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym): + # NOTE(rob): kernel requires Llama4. + return routing_method == RoutingMethodType.Llama4 + + else: + raise ValueError("Unsupported quantization scheme.") + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """Supports TRTLLM Kernel does not support EPLB.""" + return not moe_parallel_config.enable_eplb def supports_chunking(self) -> bool: return False diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py deleted file mode 100644 index 647108cc44fd..000000000000 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ /dev/null @@ -1,289 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEParallelConfig, - RoutingMethodType, -) -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - per_token_group_quant_fp8, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - QuantKey, - kFp8Dynamic128Sym, - kFp8Static128BlockSym, - kFp8StaticTensorSym, -) -from vllm.platforms import current_platform -from vllm.utils.torch_utils import direct_register_custom_op - -# -# Methods used by the oracle for kernel selection. -# - - -def _supports_current_device() -> bool: - """Supports only Blackwell-family GPUs.""" - p = current_platform - # Add check flashinfer trtllm is available - return p.is_cuda() and p.is_device_capability_family(100) - - -def _supports_no_act_and_mul() -> bool: - """Does not support non-gated MoE (i.e. Nanotron-Mini).""" - return False - - -def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, -) -> bool: - """Supports Fp8 per-tensor and Fp8 block.""" - SUPPORTED_W_A = [ - (kFp8Static128BlockSym, kFp8Dynamic128Sym), - (kFp8StaticTensorSym, kFp8StaticTensorSym), - ] - return (weight_key, activation_key) in SUPPORTED_W_A - - -def _supports_activation(activation: str) -> bool: - """Supports silu activation only.""" - return activation in ["silu"] - - -def _supports_routing_method( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - routing_method: RoutingMethodType, -) -> bool: - """Monolithic kernels need to express router support.""" - if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): - # NOTE(rob): potentially allow others here. This is a conservative list. - return routing_method in [ - RoutingMethodType.DeepSeekV3, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ] - elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym): - # NOTE(rob): kernel requires Llama4. - return routing_method == RoutingMethodType.Llama4 - - else: - raise ValueError("Unsupported quantization scheme.") - - -def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """Supports TRTLLM Kernel does not support EPLB.""" - return not moe_parallel_config.enable_eplb - - -def is_supported_config_trtllm( - moe_config: FusedMoEConfig, - weight_key: QuantKey | None, - activation_key: QuantKey | None, - activation_format: mk.FusedMoEActivationFormat, -) -> tuple[bool, str | None]: - """ - This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config - """ - - def _make_reason(reason: str) -> str: - return f"kernel does not support {reason}" - - if not _supports_current_device(): - return False, _make_reason("current device") - elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()): - return False, _make_reason("no act_and_mul MLP layer") - elif not _supports_activation(moe_config.activation): - return False, _make_reason(f"{moe_config.activation} activation") - elif not _supports_quant_scheme(weight_key, activation_key): - return False, _make_reason("quantization scheme") - elif not _supports_parallel_config(moe_config.moe_parallel_config): - return False, _make_reason("parallel config") - elif not _supports_routing_method( - weight_key, activation_key, moe_config.routing_method - ): - return False, _make_reason("routing method") - elif activation_format != mk.FusedMoEActivationFormat.Standard: - return False, _make_reason("activation format") - - return True, None - - -def flashinfer_fused_moe_blockscale_fp8( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - x: torch.Tensor, - w13_weight: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - global_num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - expert_offset: int, - local_num_experts: int, - block_shape: list[int], - routing_method_type: int = int(RoutingMethodType.DeepSeekV3), - routed_scaling: float | None = 1.0, -) -> torch.Tensor: - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe - - topk_group = topk_group if topk_group is not None else 0 - assert top_k <= global_num_experts - assert top_k <= 10 - assert global_num_experts % 4 == 0 - assert block_shape == [128, 128] - # Routing kernel expects #experts <= #threads 512 - assert global_num_experts <= 512 - - a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1]) - # NOTE: scales of hidden states have to be transposed! - a_sf_t = a_sf.t().contiguous() - return flashinfer_trtllm_fp8_block_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=a_q, - hidden_states_scale=a_sf_t, - gemm1_weights=w13_weight, - gemm1_weights_scale=w13_weight_scale_inv, - gemm2_weights=w2_weight, - gemm2_weights_scale=w2_weight_scale_inv, - num_experts=global_num_experts, - top_k=top_k, - n_group=num_expert_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=expert_offset, - local_num_experts=local_num_experts, - routed_scaling_factor=routed_scaling, - routing_method_type=routing_method_type, - use_shuffled_weight=False, - ) - - -def flashinfer_fused_moe_blockscale_fp8_fake( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - x: torch.Tensor, - w13_weight: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - global_num_experts: int, - top_k: int, - num_expert_group: int, - topk_group: int, - intermediate_size: int, - expert_offset: int, - local_num_experts: int, - block_shape: list[int], - routing_method_type: int, - routed_scaling: float = 1.0, -) -> torch.Tensor: - return torch.empty_like(x) - - -# TODO(bnell): Does this really need to be a torch.op? -direct_register_custom_op( - op_name="flashinfer_fused_moe_blockscale_fp8", - op_func=flashinfer_fused_moe_blockscale_fp8, - fake_impl=flashinfer_fused_moe_blockscale_fp8_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -def fi_trtllm_fp8_per_tensor_moe( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - hidden_states: torch.Tensor, - input_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm2_weights: torch.Tensor, - output1_scales_scalar: torch.Tensor, - output1_scales_gate_scalar: torch.Tensor, - output2_scales_scalar: torch.Tensor, - num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - use_routing_scales_on_input: bool, - routing_method_type: int, - routed_scaling_factor: float = 1.0, -) -> torch.Tensor: - num_expert_group = num_expert_group if num_expert_group is not None else 0 - topk_group = topk_group if topk_group is not None else 0 - - quant_hidden_states, _ = moe_kernel_quantize_input( - hidden_states, - input_scale, - quant_dtype=torch.float8_e4m3fn, - per_act_token_quant=False, - ) - - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe - - return flashinfer_trtllm_fp8_per_tensor_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=quant_hidden_states, - gemm1_weights=gemm1_weights, - output1_scales_scalar=output1_scales_scalar, - output1_scales_gate_scalar=output1_scales_gate_scalar, - gemm2_weights=gemm2_weights, - output2_scales_scalar=output2_scales_scalar, - num_experts=num_experts, - top_k=top_k, - n_group=num_expert_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=local_expert_offset, - local_num_experts=local_num_experts, - routed_scaling_factor=routed_scaling_factor, - use_routing_scales_on_input=use_routing_scales_on_input, - routing_method_type=routing_method_type, - ) - - -def fi_trtllm_fp8_per_tensor_moe_fake( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - hidden_states: torch.Tensor, - input_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm2_weights: torch.Tensor, - output1_scales_scalar: torch.Tensor, - output1_scales_gate_scalar: torch.Tensor, - output2_scales_scalar: torch.Tensor, - num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - use_routing_scales_on_input: bool, - routing_method_type: int, - routed_scaling_factor: float = 1.0, -) -> torch.Tensor: - return torch.empty_like(hidden_states) - - -# TODO(bnell): Does this really need to be a torch.op? -direct_register_custom_op( - op_name="fi_trtllm_fp8_per_tensor_moe", - op_func=fi_trtllm_fp8_per_tensor_moe, - mutates_args=["hidden_states"], - fake_impl=fi_trtllm_fp8_per_tensor_moe_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 7126512e7e9f..97c57d0fc74b 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -28,7 +28,7 @@ def __init__( moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, ): - super().__init__(quant_config) + super().__init__(moe_config=moe_config, quant_config=quant_config) self.moe_config = moe_config # TODO: set this via the constructor @@ -91,6 +91,8 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo @staticmethod def _supports_routing_method( routing_method_type: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, ) -> bool: # NOTE(rob): this is a conservative list. return routing_method_type in [ diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index f67d9dbb4504..e7ab6567f698 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -16,6 +16,7 @@ FusedMoEConfig, FusedMoEParallelConfig, FusedMoEQuantConfig, + RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, @@ -498,6 +499,10 @@ def _make_reason(reason: str) -> str: return False, _make_reason("quantization scheme") elif not cls._supports_parallel_config(moe_config.moe_parallel_config): return False, _make_reason("parallel config") + elif not cls._supports_routing_method( + moe_config.routing_method, weight_key, activation_key + ): + return False, _make_reason("routing method") elif activation_format != cls.activation_format(): return False, _make_reason(f"{activation_format.value} activation format") return True, None @@ -544,6 +549,19 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo """ raise NotImplementedError + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """ + Whether the kernel supports a routing method. Can be overriden + by monolithic kernels that excute the router in addition to the + fused experts. + """ + return True + # # Various helpers for accessing quantization parameters from the # quant_config. From 311462dea89d7b8c3d455a3734fb3ce2a675b9dd Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 08:41:51 -0500 Subject: [PATCH 043/207] remove FI TRTLLM specific logic from fp8 oracle Signed-off-by: Robert Shaw --- .../layers/fused_moe/oracle/fp8.py | 97 ++++++------------- 1 file changed, 32 insertions(+), 65 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 1788b5e10d32..62e1ad55b276 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -14,9 +14,6 @@ fp8_w8a8_moe_quant_config, fp8_w8a16_moe_quant_config, ) -from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import ( - is_supported_config_trtllm, -) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -209,59 +206,38 @@ def _return_or_raise( elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"): # If user is explicit about backend, validate it. fi_backend = get_flashinfer_moe_backend() - - if fi_backend == FlashinferMoeBackend.TENSORRT_LLM: - backend = Fp8MoeBackend.FLASHINFER_TRTLLM - supported, reason = is_supported_config_trtllm( - config, weight_key, activation_key, activation_format - ) - if supported: - logger.info_once(_make_log_backend(backend)) - return backend, None - else: - raise ValueError(_make_log_unsupported(backend, reason)) - - elif fi_backend == FlashinferMoeBackend.CUTLASS: + if fi_backend == FlashinferMoeBackend.CUTLASS: backend = Fp8MoeBackend.FLASHINFER_CUTLASS - return _return_or_raise( - backend, config, weight_key, activation_key, activation_format - ) - + elif fi_backend == FlashinferMoeBackend.TENSORRT_LLM: + backend = Fp8MoeBackend.FLASHINFER_TRTLLM else: - assert fi_backend == FlashinferMoeBackend.CUTEDSL - raise ValueError("FlashInfer MaskedGEMM not supported for FP8") - + raise ValueError( + f"FlashInfer MOE backend {fi_backend} does not support FP8 MoE." + ) + k_cls = backend_to_kernel_cls(backend) + return _return_or_raise( + backend, config, weight_key, activation_key, activation_format + ) else: # If the user is not explicit about the backend, try both. for backend in [ Fp8MoeBackend.FLASHINFER_TRTLLM, Fp8MoeBackend.FLASHINFER_CUTLASS, ]: - if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: - k_cls = None - supported, reason = is_supported_config_trtllm( - config, - weight_key, - activation_key, - activation_format, - ) - else: - k_cls = backend_to_kernel_cls(backend) - supported, reason = k_cls.is_supported_config( - k_cls, - config, - weight_key, - activation_key, - activation_format, - ) - - if supported: - logger.info_once(_make_log_backend(backend), scope="local") - return backend, k_cls - else: - logger.debug_once( - _make_log_unsupported(backend, reason), scope="local" - ) + k_cls = backend_to_kernel_cls(backend) + supported, reason = k_cls.is_supported_config( + k_cls, + config, + weight_key, + activation_key, + activation_format, + ) + + if supported: + logger.info_once(_make_log_backend(backend), scope="local") + return backend, k_cls + else: + logger.debug_once(_make_log_unsupported(backend, reason), scope="local") raise NotImplementedError( "Found VLLM_USE_FLASHINFER_MOE_FP8=1, but no " @@ -306,23 +282,14 @@ def _return_or_raise( # Select kernels in order of backend. for backend in AVAILABLE_BACKENDS: - if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: - k_cls = None - supported, reason = is_supported_config_trtllm( - config, - weight_key, - activation_key, - activation_format, - ) - else: - k_cls = backend_to_kernel_cls(backend) - supported, reason = k_cls.is_supported_config( - k_cls, - config, - weight_key, - activation_key, - activation_format, - ) + k_cls = backend_to_kernel_cls(backend) + supported, reason = k_cls.is_supported_config( + k_cls, + config, + weight_key, + activation_key, + activation_format, + ) if supported: logger.info_once(_make_log_backend(backend), scope="local") From 5adcf51dd10e447a3749bb59e7951b9ed24d7bb2 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 08:44:13 -0500 Subject: [PATCH 044/207] remove FI TRTLLM specific logic from nvfp4 oracle Signed-off-by: Robert Shaw --- .../layers/fused_moe/oracle/nvfp4.py | 72 ++++++------------- 1 file changed, 20 insertions(+), 52 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 0c5d5fb02e5d..d13ecab59965 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -18,7 +18,6 @@ MoEPrepareAndFinalizeNoEP, ) from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( - is_supported_config_trtllm, prepare_nvfp4_moe_layer_for_fi_or_cutlass, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( @@ -178,43 +177,21 @@ def _return_or_raise( elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"): # If user is explicit about backend, validate it. - fi_backend = get_flashinfer_moe_backend() - - if fi_backend == FlashinferMoeBackend.TENSORRT_LLM: - backend = NvFp4MoeBackend.FLASHINFER_TRTLLM - supported, reason = is_supported_config_trtllm( - config, weight_key, activation_key, activation_format - ) - if supported: - logger.info_once(_make_log_backend(backend)) - return backend, None - else: - raise ValueError(_make_log_unsupported(backend, reason)) - else: - backend = fi_2_vllm_backend_map[fi_backend] - return _return_or_raise( - backend, config, weight_key, activation_key, activation_format - ) + backend = fi_2_vllm_backend_map[get_flashinfer_moe_backend()] + return _return_or_raise( + backend, config, weight_key, activation_key, activation_format + ) else: # If the user is not explicit about the backend, try each. for backend in FLASHINFER_NVFP4_MOE_BACKENDS: - if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - k_cls = None - supported, reason = is_supported_config_trtllm( - config, - weight_key, - activation_key, - activation_format, - ) - else: - k_cls = backend_to_kernel_cls(backend) - supported, reason = k_cls.is_supported_config( - k_cls, - config, - weight_key, - activation_key, - activation_format, - ) + k_cls = backend_to_kernel_cls(backend) + supported, reason = k_cls.is_supported_config( + k_cls, + config, + weight_key, + activation_key, + activation_format, + ) if supported: logger.info_once(_make_log_backend(backend), scope="local") return backend, None @@ -236,23 +213,14 @@ def _return_or_raise( # Select kernels in order of backend. for backend in AVAILABLE_BACKENDS: - if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - k_cls = None # type: ignore[assignment] - supported, reason = is_supported_config_trtllm( - config, - weight_key, - activation_key, - activation_format, - ) - else: - k_cls = backend_to_kernel_cls(backend) - supported, reason = k_cls.is_supported_config( - k_cls, - config, - weight_key, - activation_key, - activation_format, - ) + k_cls = backend_to_kernel_cls(backend) + supported, reason = k_cls.is_supported_config( + k_cls, + config, + weight_key, + activation_key, + activation_format, + ) if supported: logger.info_once(_make_log_backend(backend), scope="local") From 41b6de867c2188512226a5013f970b34164575d3 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 08:46:02 -0500 Subject: [PATCH 045/207] add trtllm to backend_to_kernel_cls Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/oracle/fp8.py | 6 +++++- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 8 +++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 62e1ad55b276..faac55af41ad 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -54,7 +54,11 @@ def backend_to_kernel_cls( backend: Fp8MoeBackend, ) -> type[mk.FusedMoEPermuteExpertsUnpermute]: if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: - raise NotImplementedError + from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( + FlashInferTrtLlmFp8Experts, + ) + + return FlashInferTrtLlmFp8Experts elif backend == Fp8MoeBackend.FLASHINFER_CUTLASS: from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index d13ecab59965..6f4476b5c2b6 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -67,10 +67,12 @@ def backend_to_kernel_cls( backend: NvFp4MoeBackend, ) -> type[mk.FusedMoEPermuteExpertsUnpermute]: if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - raise NotImplementedError( - "FLASHINFER_TRTLLM doesn't support Modular Kernel Interface" + from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_nvfp4_moe import ( + FlashInferTrtLlmFp4Experts, ) + return FlashInferTrtLlmFp4Experts + elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS: from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, @@ -165,7 +167,7 @@ def _return_or_raise( k_cls, config, weight_key, activation_key, activation_format ) if supported: - logger.info_once(_make_log_backend(backend)) + logger.info_once(_make_log_backend(backend), scope="local") return backend, k_cls raise ValueError(_make_log_unsupported(backend, reason)) From 6a53c7ebc0d8dc8e72de5cbb7e6da672542d49aa Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 09:05:41 -0500 Subject: [PATCH 046/207] things are working with blockfp8 Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 25 ++++++++----------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index b4edd7d4b21a..b6db55f267ea 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -126,10 +126,8 @@ def __init__( super().__init__(moe_config, quant_config) self.moe_config = moe_config - # TODO: set this via the constructor - # self.routing_method_type = flashinfer.RoutingMethodType.Renormalize - self.routing_method_type = flashinfer.RoutingMethodType.Llama4 - # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 + + self.routing_method_type = moe_config.routing_method self.routing_bias = None # TODO: to: in_dtype.shape @@ -146,13 +144,14 @@ def __init__( self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank - self._g1_alphas, self._g2_alphas = make_fp8_moe_alpha_scales_for_fi( - w13_scale=self.quant_config.w1_scale, - w13_input_scale=self.quant_config.a1_scale, - w2_scale=self.quant_config.w2_scale, - w2_input_scale=self.quant_config.a2_scale, - ) - self.g1_scale_c = self._g1_alphas / self.quant_config.a2_scale + if self.quant_config.is_per_tensor: + self._g1_alphas, self._g2_alphas = make_fp8_moe_alpha_scales_for_fi( + w13_scale=self.quant_config.w1_scale, + w13_input_scale=self.quant_config.a1_scale, + w2_scale=self.quant_config.w2_scale, + w2_input_scale=self.quant_config.a2_scale, + ) + self.g1_scale_c = self._g1_alphas / self.quant_config.a2_scale @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -194,15 +193,14 @@ def _supports_routing_method( activation_key: QuantKey | None, ) -> bool: """Monolithic kernels need to express router support.""" + # NOTE(rob): potentially allow others here. This is a conservative list. if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): - # NOTE(rob): potentially allow others here. This is a conservative list. return routing_method in [ RoutingMethodType.DeepSeekV3, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, ] elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym): - # NOTE(rob): kernel requires Llama4. return routing_method == RoutingMethodType.Llama4 else: @@ -314,7 +312,6 @@ def _apply_per_block_monolithic( local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, routed_scaling_factor=self.routing_scaling_factor, - tile_tokens_dim=None, routing_method_type=self.routing_method_type, ) From 8465084d66db50f534bd97611cfb54d2e7bc2593 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 09:13:16 -0500 Subject: [PATCH 047/207] remove torch op Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_fp8_moe.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index b6db55f267ea..f589d84201ab 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -125,10 +125,7 @@ def __init__( ): super().__init__(moe_config, quant_config) - self.moe_config = moe_config - self.routing_method_type = moe_config.routing_method - self.routing_bias = None # TODO: to: in_dtype.shape self.e_score_correction_bias = None @@ -144,6 +141,7 @@ def __init__( self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank + # Make additional scales for per-tensor interface. if self.quant_config.is_per_tensor: self._g1_alphas, self._g2_alphas = make_fp8_moe_alpha_scales_for_fi( w13_scale=self.quant_config.w1_scale, @@ -208,7 +206,7 @@ def _supports_routing_method( @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """Supports TRTLLM Kernel does not support EPLB.""" + """TRTLLMGenKernel does not support EPLB.""" return not moe_parallel_config.enable_eplb def supports_chunking(self) -> bool: From 7c8778c26239f132eb18bd809257de4e10f8b166 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 09:13:22 -0500 Subject: [PATCH 048/207] remove torch op Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 90 ------------------- 1 file changed, 90 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index f589d84201ab..d61ff34442d6 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -24,99 +24,9 @@ kFp8Static128BlockSym, kFp8StaticTensorSym, ) -from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.engine.utils import current_platform -def fi_trtllm_fp8_per_tensor_moe( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - hidden_states: torch.Tensor, - input_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm2_weights: torch.Tensor, - output1_scales_scalar: torch.Tensor, - output1_scales_gate_scalar: torch.Tensor, - output2_scales_scalar: torch.Tensor, - num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - use_routing_scales_on_input: bool, - routing_method_type: int, - routed_scaling_factor: float = 1.0, -) -> torch.Tensor: - num_expert_group = num_expert_group if num_expert_group is not None else 0 - topk_group = topk_group if topk_group is not None else 0 - - quant_hidden_states, _ = moe_kernel_quantize_input( - hidden_states, - input_scale, - quant_dtype=torch.float8_e4m3fn, - per_act_token_quant=False, - ) - - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe - - return flashinfer_trtllm_fp8_per_tensor_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=quant_hidden_states, - gemm1_weights=gemm1_weights, - output1_scales_scalar=output1_scales_scalar, - output1_scales_gate_scalar=output1_scales_gate_scalar, - gemm2_weights=gemm2_weights, - output2_scales_scalar=output2_scales_scalar, - num_experts=num_experts, - top_k=top_k, - n_group=num_expert_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=local_expert_offset, - local_num_experts=local_num_experts, - routed_scaling_factor=routed_scaling_factor, - use_routing_scales_on_input=use_routing_scales_on_input, - routing_method_type=routing_method_type, - ) - - -def fi_trtllm_fp8_per_tensor_moe_fake( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - hidden_states: torch.Tensor, - input_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm2_weights: torch.Tensor, - output1_scales_scalar: torch.Tensor, - output1_scales_gate_scalar: torch.Tensor, - output2_scales_scalar: torch.Tensor, - num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - use_routing_scales_on_input: bool, - routing_method_type: int, - routed_scaling_factor: float = 1.0, -) -> torch.Tensor: - return torch.empty_like(hidden_states) - - -# TODO(bnell): Does this really need to be a torch.op? -direct_register_custom_op( - op_name="fi_trtllm_fp8_per_tensor_moe", - op_func=fi_trtllm_fp8_per_tensor_moe, - mutates_args=["hidden_states"], - fake_impl=fi_trtllm_fp8_per_tensor_moe_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - class FlashInferTrtLlmFp8Experts(mk.FusedMoEPermuteExpertsUnpermute): def __init__( self, From e2dadb879907cea9db251217eb8d6a695de8aa2c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 09:15:46 -0500 Subject: [PATCH 049/207] remove Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 29 +++++++++---------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index d61ff34442d6..8895cdf2a1ea 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -12,12 +12,6 @@ RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - make_fp8_moe_alpha_scales_for_fi, -) -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - per_token_group_quant_fp8, -) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, @@ -53,12 +47,12 @@ def __init__( # Make additional scales for per-tensor interface. if self.quant_config.is_per_tensor: - self._g1_alphas, self._g2_alphas = make_fp8_moe_alpha_scales_for_fi( - w13_scale=self.quant_config.w1_scale, - w13_input_scale=self.quant_config.a1_scale, - w2_scale=self.quant_config.w2_scale, - w2_input_scale=self.quant_config.a2_scale, - ) + self._g1_alphas = ( + self.quant_config.w1_scale * self.quant_config.a1_scale + ).squeeze() + self._g2_alphas = ( + self.quant_config.w2_scale * self.quant_config.a2_scale + ).squeeze() self.g1_scale_c = self._g1_alphas / self.quant_config.a2_scale @staticmethod @@ -178,6 +172,7 @@ def _apply_per_block_monolithic( expert_map: torch.Tensor | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: + assert not apply_router_weight_on_input assert activation == "silu" assert ( self.e_score_correction_bias is None @@ -196,9 +191,13 @@ def _apply_per_block_monolithic( # Routing kernel expects #experts <= #threads 512 assert global_num_experts <= 512 - a1_q, a1q_scale = per_token_group_quant_fp8( - hidden_states, self.quant_config.block_shape[1] + a1q, a1q_scale = moe_kernel_quantize_input( + hidden_states, + self.quant_config.a1_scale, + quant_dtype=self.quant_config.quant_dtype, + per_act_token_quant=self.quant_config.per_act_token_quant, ) + # Kernel requires transposed hidden state scales # TODO: fuse into the quant kernel. a1q_scale_t = a1q_scale.t().contiguous() @@ -206,7 +205,7 @@ def _apply_per_block_monolithic( return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( routing_logits=router_logits, routing_bias=self.e_score_correction_bias, - hidden_states=a1_q, + hidden_states=a1q, hidden_states_scale=a1q_scale_t, gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale, From 9226abc1da4eff9b88b44bea29b999382dd8825d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 10:16:37 -0500 Subject: [PATCH 050/207] updated Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 26 +++---- .../layers/fused_moe/modular_kernel.py | 68 +++++++++++++++++-- .../layers/fused_moe/prepare_finalize.py | 30 ++++++++ 3 files changed, 101 insertions(+), 23 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 8895cdf2a1ea..485019247899 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -11,7 +11,6 @@ FusedMoEQuantConfig, RoutingMethodType, ) -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, @@ -170,6 +169,7 @@ def _apply_per_block_monolithic( activation: str, global_num_experts: int, expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: assert not apply_router_weight_on_input @@ -191,21 +191,15 @@ def _apply_per_block_monolithic( # Routing kernel expects #experts <= #threads 512 assert global_num_experts <= 512 - a1q, a1q_scale = moe_kernel_quantize_input( - hidden_states, - self.quant_config.a1_scale, - quant_dtype=self.quant_config.quant_dtype, - per_act_token_quant=self.quant_config.per_act_token_quant, - ) - # Kernel requires transposed hidden state scales # TODO: fuse into the quant kernel. + assert a1q_scale is not None a1q_scale_t = a1q_scale.t().contiguous() return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( routing_logits=router_logits, routing_bias=self.e_score_correction_bias, - hidden_states=a1q, + hidden_states=hidden_states, hidden_states_scale=a1q_scale_t, gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale, @@ -231,24 +225,19 @@ def _apply_per_tensor_monolithic( activation: str, global_num_experts: int, expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: assert self.routing_method_type == RoutingMethodType.Llama4 assert apply_router_weight_on_input + assert a1q_scale is None topk_group = self.topk_group if self.topk_group is not None else 0 - a1q, _ = moe_kernel_quantize_input( - hidden_states, - self.quant_config.a1_scale, - quant_dtype=self.quant_config.quant_dtype, - per_act_token_quant=self.quant_config.per_act_token_quant, - ) - return flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( routing_logits=router_logits, routing_bias=self.e_score_correction_bias, - hidden_states=a1q, + hidden_states=hidden_states, gemm1_weights=w1, output1_scales_scalar=self.g1_scale_c, output1_scales_gate_scalar=self._g1_alphas, @@ -275,6 +264,7 @@ def apply_monolithic( activation: str, global_num_experts: int, expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: if self.quant_config.block_shape is not None: @@ -286,6 +276,7 @@ def apply_monolithic( activation, global_num_experts, expert_map, + a1q_scale, apply_router_weight_on_input, ) elif self.quant_config.is_per_tensor: @@ -297,6 +288,7 @@ def apply_monolithic( activation, global_num_experts, expert_map, + a1q_scale, apply_router_weight_on_input, ) else: diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index e7ab6567f698..63a81719eada 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -152,6 +152,18 @@ def apply( torch.Tensor | None, ] +# +# PrepareResultType is a tuple of: +# - quantized + dispatched a. +# - quantized + dispatched a1_scales. +# +# See `prepare_monolithic` method below. +# +PrepareMonolithicResultType = tuple[ + torch.Tensor, + torch.Tensor | None, +] + ReceiverType = Callable[[], PrepareResultType] @@ -171,6 +183,13 @@ def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"): """ return + def supports_async(self) -> bool: + """ + Indicates whether or not this class implements prepare_async and + finalize_async. + """ + return False + @abstractmethod def prepare( self, @@ -205,12 +224,33 @@ def prepare( """ raise NotImplementedError - def supports_async(self) -> bool: + def prepare_monolithic( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + ) -> PrepareMonolithicResultType: """ - Indicates whether or not this class implements prepare_async and - finalize_async. + Perform any quantization (and/or) dispatching needed for this kernel. + - a1: The (unquantized) input to the MoE layer. + - router_logits: the logits from the router. + - num_experts: The total number of experts in the global expert space. + - expert_map: A tensor mapping expert indices from the global expert + space to the local expert space of the expert parallel shard. + - apply_router_weight_on_input: When True, apply the weights to the + activations, before quantization + dispatching. + - quant_config: Quantization info provided by the fused experts. + + Returns a tuple of: + - quantized + dispatched a. + - Optional quantized + dispatched a1_scales. """ - return False + raise NotImplementedError( + f"prepare_monolithic not supported for {self.__class__.__name__}" + ) def prepare_async( self, @@ -800,6 +840,7 @@ def apply_monolithic( activation: str, global_num_experts: int, expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: """ @@ -1414,8 +1455,17 @@ def forward_monolithic( that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ - return self.fused_experts.apply_monolithic( - hidden_states=hidden_states, + a1q, a1q_scale = self.prepare_finalize.prepare_monolithic( + hidden_states, + router_logits, + global_num_experts, + expert_map, + apply_router_weight_on_input, + self.fused_experts.quant_config, + ) + + fused_out = self.fused_experts.apply_monolithic( + hidden_states=a1q, w1=w1, w2=w2, router_logits=router_logits, @@ -1423,4 +1473,10 @@ def forward_monolithic( global_num_experts=global_num_experts, expert_map=expert_map, apply_router_weight_on_input=apply_router_weight_on_input, + a1q_scale=a1q_scale, ) + + # TODO(rob): once naive P/F Dp/Ep lands, we will need to + # add support here for finalizing over just the hidden states. + + return fused_out diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index e5c136f52cbe..d1d4f5075cbc 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -73,6 +73,36 @@ def prepare( return a1q, a1q_scale, None, None, None + def prepare_monolithic( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # Defer input quant to moe kernel for backends (e.g. AITER, FI) + # which use a single kernel call for quant + experts. + if self.defer_input_quant: + return a1, None + + a1_scale = ( + quant_config.a1_gscale + if (quant_config.quant_dtype == "nvfp4") + else quant_config.a1_scale + ) + a1q, a1q_scale = moe_kernel_quantize_input( + a1, + a1_scale, + quant_config.quant_dtype, + quant_config.per_act_token_quant, + quant_config.block_shape, + is_fp4_scale_swizzled=False, + ) + + return a1q, a1q_scale + def finalize( self, output: torch.Tensor, From 9f0a8bed85d41b6e3c714f10f58e1892cd0a35d5 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 10:32:13 -0500 Subject: [PATCH 051/207] updated interface Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_fp8_moe.py | 1 - .../layers/fused_moe/modular_kernel.py | 14 -------------- .../layers/fused_moe/prepare_finalize.py | 6 ------ 3 files changed, 21 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 485019247899..58d49afc6a99 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -230,7 +230,6 @@ def _apply_per_tensor_monolithic( ) -> torch.Tensor: assert self.routing_method_type == RoutingMethodType.Llama4 assert apply_router_weight_on_input - assert a1q_scale is None topk_group = self.topk_group if self.topk_group is not None else 0 diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 63a81719eada..27c349fa9ceb 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -227,21 +227,11 @@ def prepare( def prepare_monolithic( self, a1: torch.Tensor, - router_logits: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, quant_config: FusedMoEQuantConfig, ) -> PrepareMonolithicResultType: """ Perform any quantization (and/or) dispatching needed for this kernel. - a1: The (unquantized) input to the MoE layer. - - router_logits: the logits from the router. - - num_experts: The total number of experts in the global expert space. - - expert_map: A tensor mapping expert indices from the global expert - space to the local expert space of the expert parallel shard. - - apply_router_weight_on_input: When True, apply the weights to the - activations, before quantization + dispatching. - quant_config: Quantization info provided by the fused experts. Returns a tuple of: @@ -1457,10 +1447,6 @@ def forward_monolithic( a1q, a1q_scale = self.prepare_finalize.prepare_monolithic( hidden_states, - router_logits, - global_num_experts, - expert_map, - apply_router_weight_on_input, self.fused_experts.quant_config, ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index d1d4f5075cbc..1ddea2c81318 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -68,7 +68,6 @@ def prepare( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, - is_fp4_scale_swizzled=False, ) return a1q, a1q_scale, None, None, None @@ -76,10 +75,6 @@ def prepare( def prepare_monolithic( self, a1: torch.Tensor, - router_logits: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, quant_config: FusedMoEQuantConfig, ) -> tuple[torch.Tensor, torch.Tensor | None]: # Defer input quant to moe kernel for backends (e.g. AITER, FI) @@ -98,7 +93,6 @@ def prepare_monolithic( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, - is_fp4_scale_swizzled=False, ) return a1q, a1q_scale From 2a61a6b3e6dd53c8258e4a3b62a3b2fd4e2efe56 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 10:35:16 -0500 Subject: [PATCH 052/207] updated interface Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 14 +++----------- .../layers/fused_moe/prepare_finalize.py | 1 + 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 97c57d0fc74b..414de453f578 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -206,19 +206,11 @@ def apply_monolithic( activation: str, global_num_experts: int, expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, ) -> torch.Tensor: assert activation == "silu" - - # Quantize input. - if isinstance(hidden_states, tuple): - a1q, a1q_scale = hidden_states - else: - a1q, a1q_scale = flashinfer.fp4_quantize( - hidden_states, - self.quant_config.a1_gscale, - is_sf_swizzled_layout=False, - ) + assert a1q_scale is not None # Prepare routing bias into kernel format. routing_bias = self.e_score_correction_bias @@ -234,7 +226,7 @@ def apply_monolithic( return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( routing_logits=router_logits, routing_bias=routing_bias, - hidden_states=a1q, + hidden_states=hidden_states, hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 1ddea2c81318..5cead286e4b2 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -93,6 +93,7 @@ def prepare_monolithic( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, + is_fp4_scale_swizzled=False, # TODO: figure out how to propogate this flag ) return a1q, a1q_scale From c6386279bfe7c0322a19964edbcd315968b2f913 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 10:39:00 -0500 Subject: [PATCH 053/207] updated interface Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 17 +++++++++++------ .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 4 +--- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 58d49afc6a99..a95aeb03aa08 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -46,12 +46,17 @@ def __init__( # Make additional scales for per-tensor interface. if self.quant_config.is_per_tensor: - self._g1_alphas = ( - self.quant_config.w1_scale * self.quant_config.a1_scale - ).squeeze() - self._g2_alphas = ( - self.quant_config.w2_scale * self.quant_config.a2_scale - ).squeeze() + w1_scale = self.quant_config.w1_scale + assert w1_scale is not None + a1_scale = self.quant_config.a1_scale + assert a1_scale is not None + w2_scale = self.quant_config.w2_scale + assert w2_scale is not None + a2_scale = self.quant_config.a2_scale + assert a2_scale is not None + + self._g1_alphas = (w1_scale * a1_scale).squeeze() + self._g2_alphas = (w2_scale * a2_scale).squeeze() self.g1_scale_c = self._g1_alphas / self.quant_config.a2_scale @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 414de453f578..c992f8345bc2 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -30,9 +30,7 @@ def __init__( ): super().__init__(moe_config=moe_config, quant_config=quant_config) - self.moe_config = moe_config - # TODO: set this via the constructor - self.routing_method_type = flashinfer.RoutingMethodType.Renormalize + self.routing_method_type = self.moe_config.routing_method_type # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 From e744eebb23153038620639f7ba74090556f93cf2 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 10:40:22 -0500 Subject: [PATCH 054/207] update comment Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/modular_kernel.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 27c349fa9ceb..5c5c67fed2bf 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -586,9 +586,10 @@ def _supports_routing_method( activation_key: QuantKey | None, ) -> bool: """ - Whether the kernel supports a routing method. Can be overriden - by monolithic kernels that excute the router in addition to the - fused experts. + Whether the kernel supports a routing method (e.g. GroupedTopK). + + Can be overriden by monolithic kernels that execute the router + in addition to the experts if certain routers are not supported. """ return True From 1b4a43f3e756f4955616484ad33f6f2a9c1ea22a Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 10:57:50 -0500 Subject: [PATCH 055/207] add skip swizzle Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/config.py | 6 ++++++ vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 3 +++ vllm/model_executor/layers/fused_moe/prepare_finalize.py | 3 ++- 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 9a28c3193587..1cf36dcc03c3 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -204,6 +204,7 @@ class FusedMoEQuantConfig: _a2: FusedMoEQuantDesc _w1: FusedMoEQuantDesc _w2: FusedMoEQuantDesc + skip_nvfp4_swizzle: bool = False def __post_init__(self): assert not self.per_act_token_quant or self.block_shape is None, ( @@ -443,6 +444,7 @@ def make( w1_zp: torch.Tensor | None = None, w2_zp: torch.Tensor | None = None, weight_dtype: torch.dtype | str | None = None, + skip_nvfp4_swizzle: bool = False, ) -> "FusedMoEQuantConfig": """ General builder function for a FusedMoEQuantConfig. @@ -472,6 +474,7 @@ def make( - w2_bias: Optional biases for w1 (GPT OSS Triton). - w1_zp: Optional w1 zero points for int4/int8 quantization. - w2_zp: Optional w2 zero points for int4/int8 quantization. + - skip_nvfp4_swizzle: Whether to skip nvfp4 scale swizzling. """ assert not isinstance(quant_dtype, str) or quant_dtype in { "nvfp4", @@ -502,6 +505,7 @@ def make( _w2=FusedMoEQuantDesc( weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias ), + skip_nvfp4_swizzle=skip_nvfp4_swizzle, ) assert quant_config.per_act_token_quant == per_act_token_quant assert quant_config.per_out_ch_quant == per_out_ch_quant @@ -673,6 +677,7 @@ def nvfp4_moe_quant_config( a2_gscale: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor, + skip_nvfp4_swizzle: bool = False, ) -> FusedMoEQuantConfig: """ Construct a quant config for mxfp4 activations and nvp4 weights. @@ -688,6 +693,7 @@ def nvfp4_moe_quant_config( per_act_token_quant=False, per_out_ch_quant=False, block_shape=None, + skip_nvfp4_swizzle=skip_nvfp4_swizzle, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 6f4476b5c2b6..69312f70c1a8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -354,6 +354,9 @@ def make_nvfp4_moe_quant_config( a2_gscale=(1.0 / a2_scale), w1_scale=w13_scale, w2_scale=w2_scale, + # NOTE(rob): this is a hack until the expets_cls() + # constructs the quant configs. + skip_nvfp4_swizzle=(backend == NvFp4MoeBackend.FLASHINFER_TRTLLM), ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 5cead286e4b2..2095c7502229 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -68,6 +68,7 @@ def prepare( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, + is_fp4_scale_swizzled=(not quant_config.skip_nvfp4_swizzle), ) return a1q, a1q_scale, None, None, None @@ -93,7 +94,7 @@ def prepare_monolithic( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, - is_fp4_scale_swizzled=False, # TODO: figure out how to propogate this flag + is_fp4_scale_swizzled=(not quant_config.skip_nvfp4_swizzle), ) return a1q, a1q_scale From 5519fa2b83eb94c38daba2433134380379e2d678 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 11:13:32 -0500 Subject: [PATCH 056/207] updated Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/config.py | 6 ------ .../layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py | 4 +--- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 7 ++----- vllm/model_executor/layers/fused_moe/prepare_finalize.py | 3 +-- 4 files changed, 4 insertions(+), 16 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 1cf36dcc03c3..9a28c3193587 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -204,7 +204,6 @@ class FusedMoEQuantConfig: _a2: FusedMoEQuantDesc _w1: FusedMoEQuantDesc _w2: FusedMoEQuantDesc - skip_nvfp4_swizzle: bool = False def __post_init__(self): assert not self.per_act_token_quant or self.block_shape is None, ( @@ -444,7 +443,6 @@ def make( w1_zp: torch.Tensor | None = None, w2_zp: torch.Tensor | None = None, weight_dtype: torch.dtype | str | None = None, - skip_nvfp4_swizzle: bool = False, ) -> "FusedMoEQuantConfig": """ General builder function for a FusedMoEQuantConfig. @@ -474,7 +472,6 @@ def make( - w2_bias: Optional biases for w1 (GPT OSS Triton). - w1_zp: Optional w1 zero points for int4/int8 quantization. - w2_zp: Optional w2 zero points for int4/int8 quantization. - - skip_nvfp4_swizzle: Whether to skip nvfp4 scale swizzling. """ assert not isinstance(quant_dtype, str) or quant_dtype in { "nvfp4", @@ -505,7 +502,6 @@ def make( _w2=FusedMoEQuantDesc( weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias ), - skip_nvfp4_swizzle=skip_nvfp4_swizzle, ) assert quant_config.per_act_token_quant == per_act_token_quant assert quant_config.per_out_ch_quant == per_out_ch_quant @@ -677,7 +673,6 @@ def nvfp4_moe_quant_config( a2_gscale: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor, - skip_nvfp4_swizzle: bool = False, ) -> FusedMoEQuantConfig: """ Construct a quant config for mxfp4 activations and nvp4 weights. @@ -693,7 +688,6 @@ def nvfp4_moe_quant_config( per_act_token_quant=False, per_out_ch_quant=False, block_shape=None, - skip_nvfp4_swizzle=skip_nvfp4_swizzle, ) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index c992f8345bc2..8b11eae43653 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -30,7 +30,7 @@ def __init__( ): super().__init__(moe_config=moe_config, quant_config=quant_config) - self.routing_method_type = self.moe_config.routing_method_type + self.routing_method_type = self.moe_config.routing_method # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 @@ -189,7 +189,6 @@ def apply( local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, routed_scaling_factor=None, - tile_tokens_dim=None, routing_method_type=1, do_finalize=True, output=output, @@ -246,7 +245,6 @@ def apply_monolithic( local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, routed_scaling_factor=None, - tile_tokens_dim=None, routing_method_type=self.routing_method_type, do_finalize=True, )[0] diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 69312f70c1a8..44139886ad57 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -68,10 +68,10 @@ def backend_to_kernel_cls( ) -> type[mk.FusedMoEPermuteExpertsUnpermute]: if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_nvfp4_moe import ( - FlashInferTrtLlmFp4Experts, + FlashInferTrtLlmNvFp4Experts, ) - return FlashInferTrtLlmFp4Experts + return FlashInferTrtLlmNvFp4Experts elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS: from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( @@ -354,9 +354,6 @@ def make_nvfp4_moe_quant_config( a2_gscale=(1.0 / a2_scale), w1_scale=w13_scale, w2_scale=w2_scale, - # NOTE(rob): this is a hack until the expets_cls() - # constructs the quant configs. - skip_nvfp4_swizzle=(backend == NvFp4MoeBackend.FLASHINFER_TRTLLM), ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 2095c7502229..3f816812c2a0 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -68,7 +68,6 @@ def prepare( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, - is_fp4_scale_swizzled=(not quant_config.skip_nvfp4_swizzle), ) return a1q, a1q_scale, None, None, None @@ -94,7 +93,7 @@ def prepare_monolithic( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, - is_fp4_scale_swizzled=(not quant_config.skip_nvfp4_swizzle), + is_fp4_scale_swizzled=False, ) return a1q, a1q_scale From af048427b13c21068485773ea4783c4babeaa6be Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 12:14:50 -0500 Subject: [PATCH 057/207] updated Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/config.py | 6 ++++++ vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 4 ++++ vllm/model_executor/layers/fused_moe/prepare_finalize.py | 3 ++- 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 9a28c3193587..a55030edf8e6 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -204,6 +204,7 @@ class FusedMoEQuantConfig: _a2: FusedMoEQuantDesc _w1: FusedMoEQuantDesc _w2: FusedMoEQuantDesc + is_nvfp4_scale_swizzled: bool = True def __post_init__(self): assert not self.per_act_token_quant or self.block_shape is None, ( @@ -443,6 +444,7 @@ def make( w1_zp: torch.Tensor | None = None, w2_zp: torch.Tensor | None = None, weight_dtype: torch.dtype | str | None = None, + is_nvfp4_scale_swizzled: bool = True, ) -> "FusedMoEQuantConfig": """ General builder function for a FusedMoEQuantConfig. @@ -472,6 +474,7 @@ def make( - w2_bias: Optional biases for w1 (GPT OSS Triton). - w1_zp: Optional w1 zero points for int4/int8 quantization. - w2_zp: Optional w2 zero points for int4/int8 quantization. + - is_nvfp4_scale_swizzled: Whether to swizzle the nvfp4 scale swizzling. """ assert not isinstance(quant_dtype, str) or quant_dtype in { "nvfp4", @@ -502,6 +505,7 @@ def make( _w2=FusedMoEQuantDesc( weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias ), + is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled, ) assert quant_config.per_act_token_quant == per_act_token_quant assert quant_config.per_out_ch_quant == per_out_ch_quant @@ -673,6 +677,7 @@ def nvfp4_moe_quant_config( a2_gscale: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor, + is_nvfp4_scale_swizzled: bool = True, ) -> FusedMoEQuantConfig: """ Construct a quant config for mxfp4 activations and nvp4 weights. @@ -688,6 +693,7 @@ def nvfp4_moe_quant_config( per_act_token_quant=False, per_out_ch_quant=False, block_shape=None, + is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 44139886ad57..67c996a12d04 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -354,6 +354,10 @@ def make_nvfp4_moe_quant_config( a2_gscale=(1.0 / a2_scale), w1_scale=w13_scale, w2_scale=w2_scale, + # NOTE(rob): this is a hack until the MoE kernels + # create their own quant configs. TRTLLM kernel + # does not accept swizzled input quant scales. + is_nvfp4_scale_swizzled=(backend != NvFp4MoeBackend.FLASHINFER_TRTLLM), ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 3f816812c2a0..b7184b12414d 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -68,6 +68,7 @@ def prepare( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, ) return a1q, a1q_scale, None, None, None @@ -93,7 +94,7 @@ def prepare_monolithic( quant_config.quant_dtype, quant_config.per_act_token_quant, quant_config.block_shape, - is_fp4_scale_swizzled=False, + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, ) return a1q, a1q_scale From 89fa74f98cc861ce78b67d3f8016a507041e7b76 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 12:19:22 -0500 Subject: [PATCH 058/207] updated Signed-off-by: Robert Shaw --- .../layers/fused_moe/modular_kernel.py | 20 ++++++++++++ .../layers/fused_moe/prepare_finalize.py | 31 ++----------------- 2 files changed, 22 insertions(+), 29 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 5c5c67fed2bf..e16106ea9d2a 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -23,6 +23,7 @@ apply_moe_activation, count_expert_num_tokens, disable_inplace, + moe_kernel_quantize_input, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -190,6 +191,25 @@ def supports_async(self) -> bool: """ return False + def _quantize_input( + self, + a1: torch.Tensor, + quant_config: FusedMoEQuantConfig, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + a1_scale = ( + quant_config.a1_gscale + if quant_config.use_nvfp4_w4a4 + else quant_config.a1_scale + ) + return moe_kernel_quantize_input( + a1, + a1_scale, + quant_config.quant_dtype, + quant_config.per_act_token_quant, + quant_config.block_shape, + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, + ) + @abstractmethod def prepare( self, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index b7184b12414d..b748cb63ce02 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -9,7 +9,6 @@ TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate, ) -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize): @@ -57,19 +56,7 @@ def prepare( if self.defer_input_quant: return a1, None, None, None, None - a1_scale = ( - quant_config.a1_gscale - if (quant_config.quant_dtype == "nvfp4") - else quant_config.a1_scale - ) - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - a1_scale, - quant_config.quant_dtype, - quant_config.per_act_token_quant, - quant_config.block_shape, - is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, - ) + a1q, a1q_scale = self._quantize_input(a1, quant_config) return a1q, a1q_scale, None, None, None @@ -83,21 +70,7 @@ def prepare_monolithic( if self.defer_input_quant: return a1, None - a1_scale = ( - quant_config.a1_gscale - if (quant_config.quant_dtype == "nvfp4") - else quant_config.a1_scale - ) - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - a1_scale, - quant_config.quant_dtype, - quant_config.per_act_token_quant, - quant_config.block_shape, - is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, - ) - - return a1q, a1q_scale + return self._quantize_input(a1, quant_config) def finalize( self, From ad19472faabbd71bd66ba7ecd3ba6191c47a69ed Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 13:03:21 -0500 Subject: [PATCH 059/207] convert deepep ht to use the scheme Signed-off-by: Robert Shaw --- .../fused_moe/deepep_ht_prepare_finalize.py | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py index 929cff79980c..56c41d352a0d 100644 --- a/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py @@ -11,7 +11,6 @@ TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate, ) -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, @@ -229,13 +228,7 @@ def _receiver( # Quantize after dispatch. expert_x_scale = None if expert_x.numel() != 0: - expert_x, expert_x_scale = moe_kernel_quantize_input( - expert_x, - a1_scale, - quant_dtype=quant_config.quant_dtype, - per_act_token_quant=False, - block_shape=quant_config.block_shape, - ) + expert_x, expert_x_scale = self._quantize_input(expert_x, quant_config) return ( expert_x, @@ -268,13 +261,7 @@ def prepare_async( if quant_config.is_block_quantized: # Quant and Dispatch - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - quant_config.a1_scale, - quant_dtype=quant_config.quant_dtype, - per_act_token_quant=quant_config.per_act_token_quant, - block_shape=quant_config.block_shape, - ) + a1q, a1q_scale = self._quantize_input(a1, quant_config) if a1q_scale is not None and a1q_scale.numel() == 1: a1q_scale = a1q_scale.view(1, 1) a1_post_scale = None From 76b0f8642e9608a19e4bd730111d5ca9d2062e1f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 13:15:40 -0500 Subject: [PATCH 060/207] convert some kernels to use _quantize_input Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_cutlass_prepare_finalize.py | 1 + .../layers/fused_moe/pplx_prepare_finalize.py | 11 ++--------- .../compressed_tensors/compressed_tensors_moe.py | 3 ++- vllm/model_executor/layers/quantization/modelopt.py | 3 ++- 4 files changed, 7 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py index dfff860750d6..21d80373971d 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py @@ -106,6 +106,7 @@ def prepare( if not self.use_dp: # Non-DP case: quantize activations unless using block-scale path + # TODO(rob): fix this one the P/F merged. if not self.use_deepseek_fp8_block_scale: a1q, a1q_scale = moe_kernel_quantize_input( a1, diff --git a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py index 8c4da97113ce..39bceafe7c1e 100644 --- a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py @@ -13,7 +13,6 @@ ) from vllm.model_executor.layers.fused_moe.utils import ( _validate_scale_shape, - moe_kernel_quantize_input, ) from vllm.utils.math_utils import cdiv, round_up @@ -135,14 +134,8 @@ def prepare_async( repeat_cols = 4 repeat_rows = 1 if quant_config.per_act_token_quant else a1.size(0) - # TODO(bnell): always pass quant_config.a1_scale? - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - (None if quant_config.per_act_token_quant else quant_config.a1_scale), - quant_dtype=quant_config.quant_dtype, - per_act_token_quant=quant_config.per_act_token_quant, - block_shape=quant_config.block_shape, - ) + + a1q, a1q_scale = self._quantize_input(a1, quant_config) _validate_scale_shape( a1q, a1q_scale, quant_config.per_act_token_quant, quant_config.block_shape diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index ffde738c25da..cb4f9a032634 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -625,7 +625,8 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon def is_monolithic(self) -> bool: return ( self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - and not self.moe.moe_parallel_config.enable_eplb + # NOTE(rob): this will not work until the Naive P/F is merged. + and not self.moe.moe_parallel_config.use_all2all_kernels ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 0c01b6b55b54..270042858450 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1626,7 +1626,8 @@ def supports_eplb(self) -> bool: def is_monolithic(self) -> bool: return ( self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - and not self.moe.moe_parallel_config.enable_eplb + # NOTE(rob): this will not work until the Naive P/F is merged. + and not self.moe.moe_parallel_config.use_all2all_kernels ) def apply_monolithic( From ffc6621e44dc79d6afd84b5f98a66d0f938c8df5 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 13:19:36 -0500 Subject: [PATCH 061/207] edit comment Signed-off-by: Robert Shaw --- .../quantization/compressed_tensors/compressed_tensors_moe.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index cb4f9a032634..45d82b58c7e1 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -1074,8 +1074,6 @@ def apply( inplace=self.use_inplace, activation=layer.activation, global_num_experts=layer.global_num_experts, - # TODO(rob): investigate the disable_expert_map introduced by: - # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501 expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, ) From 23b347e4bea0fa0e253d48fb2c6829f54220fc5b Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 13:38:57 -0500 Subject: [PATCH 062/207] update how we pass around router data for monolithic case Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 55 ++++++++++++------- .../layers/fused_moe/modular_kernel.py | 15 +++++ .../compressed_tensors_moe.py | 16 +++--- .../layers/quantization/modelopt.py | 25 +++------ 4 files changed, 66 insertions(+), 45 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index a95aeb03aa08..3a6b11746b4a 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -29,13 +29,6 @@ def __init__( super().__init__(moe_config, quant_config) self.routing_method_type = moe_config.routing_method - self.routing_bias = None - # TODO: to: in_dtype.shape - self.e_score_correction_bias = None - self.topk_group = None - self.num_expert_group = None - self.routing_scaling_factor = None - self.topk = moe_config.experts_per_token self.intermediate_size_per_partition = ( moe_config.intermediate_size_per_partition @@ -176,19 +169,22 @@ def _apply_per_block_monolithic( expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, ) -> torch.Tensor: assert not apply_router_weight_on_input assert activation == "silu" assert ( - self.e_score_correction_bias is None - or self.e_score_correction_bias.dtype == hidden_states.dtype + e_score_correction_bias is None + or e_score_correction_bias.dtype == hidden_states.dtype ) if self.routing_method_type == RoutingMethodType.DeepSeekV3: router_logits = router_logits.to(torch.float32) - topk_group = self.topk_group if self.topk_group is not None else 0 - assert self.topk <= global_num_experts assert self.topk <= 10 assert global_num_experts % 4 == 0 @@ -203,7 +199,7 @@ def _apply_per_block_monolithic( return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( routing_logits=router_logits, - routing_bias=self.e_score_correction_bias, + routing_bias=e_score_correction_bias, hidden_states=hidden_states, hidden_states_scale=a1q_scale_t, gemm1_weights=w1, @@ -212,12 +208,12 @@ def _apply_per_block_monolithic( gemm2_weights_scale=self.quant_config.w2_scale, num_experts=global_num_experts, top_k=self.topk, - n_group=self.num_expert_group, - topk_group=topk_group, + n_group=num_expert_group, + topk_group=(topk_group or 0), intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, - routed_scaling_factor=self.routing_scaling_factor, + routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, ) @@ -232,15 +228,23 @@ def _apply_per_tensor_monolithic( expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, ) -> torch.Tensor: assert self.routing_method_type == RoutingMethodType.Llama4 assert apply_router_weight_on_input - topk_group = self.topk_group if self.topk_group is not None else 0 + # Should only have Llama4 routing here. + assert routed_scaling_factor is None + assert e_score_correction_bias is None + assert num_expert_group is None return flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( routing_logits=router_logits, - routing_bias=self.e_score_correction_bias, + routing_bias=e_score_correction_bias, hidden_states=hidden_states, gemm1_weights=w1, output1_scales_scalar=self.g1_scale_c, @@ -249,12 +253,12 @@ def _apply_per_tensor_monolithic( output2_scales_scalar=self._g2_alphas, num_experts=global_num_experts, top_k=self.topk, - n_group=self.num_expert_group, - topk_group=topk_group, + n_group=num_expert_group, + topk_group=topk_group or 0, intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, - routed_scaling_factor=self.routing_scaling_factor, + routed_scaling_factor=routed_scaling_factor, use_routing_scales_on_input=apply_router_weight_on_input, routing_method_type=self.routing_method_type, ) @@ -270,6 +274,11 @@ def apply_monolithic( expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, ) -> torch.Tensor: if self.quant_config.block_shape is not None: return self._apply_per_block_monolithic( @@ -282,6 +291,9 @@ def apply_monolithic( expert_map, a1q_scale, apply_router_weight_on_input, + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, ) elif self.quant_config.is_per_tensor: return self._apply_per_tensor_monolithic( @@ -294,6 +306,9 @@ def apply_monolithic( expert_map, a1q_scale, apply_router_weight_on_input, + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, ) else: 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 e16106ea9d2a..2ed62816bbee 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -853,6 +853,11 @@ def apply_monolithic( expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, ) -> torch.Tensor: """ Same as apply(), except uses router_logits as opposed @@ -1459,6 +1464,11 @@ def forward_monolithic( global_num_experts: int, expert_map: torch.Tensor | None, apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, ) -> torch.Tensor: """ Same as forward(), except uses router_logits as opposed @@ -1481,6 +1491,11 @@ def forward_monolithic( expert_map=expert_map, apply_router_weight_on_input=apply_router_weight_on_input, a1q_scale=a1q_scale, + # grouped topk + fused topk bias parameters + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + topk_group=topk_group, ) # TODO(rob): once naive P/F Dp/Ep lands, we will need to diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 45d82b58c7e1..8a741baed808 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -635,8 +635,6 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.is_monolithic - assert layer.activation == "silu", "Only SiLU activation is supported." assert self.kernel is not None return self.kernel.forward_monolithic( x, @@ -647,6 +645,10 @@ def apply_monolithic( global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) def apply( @@ -656,10 +658,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert not self.is_monolithic - assert layer.activation == "silu", "Only SiLU activation is supported." assert self.kernel is not None - return self.kernel( x, layer.w13_weight, @@ -1039,7 +1038,6 @@ def apply_monolithic( router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert self.is_monolithic - assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM assert layer.activation == "silu" assert self.kernel is not None @@ -1050,10 +1048,12 @@ def apply_monolithic( router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, - # TODO(rob): investigate the disable_expert_map introduced by: - # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501 expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) def apply( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 270042858450..fbb7bb447fdb 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -949,8 +949,6 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.is_monolithic - assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM if layer.enable_eplb: raise NotImplementedError( "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend." @@ -965,6 +963,10 @@ def apply_monolithic( global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) def apply( @@ -974,16 +976,6 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert not self.is_monolithic - - # TODO(rob): this validation should happen at kernel selection - # time in the oracle rather than here. - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS: - assert layer.activation in ("silu", "relu2_no_mul"), ( - "Expected activation to be in ('silu', 'relu2_no_mul')," - f"but got {layer.activation}" - ) - assert self.kernel is not None return self.kernel( x, @@ -1636,10 +1628,6 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - assert not layer.enable_eplb - - # In monolithic case, router is fused with expert. assert self.kernel is not None return self.kernel.forward_monolithic( x, @@ -1650,6 +1638,10 @@ def apply_monolithic( global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) def apply( @@ -1659,7 +1651,6 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert not self.is_monolithic assert self.kernel is not None return self.kernel( x, From 261b20ae97bd816df30ad1d505c7a7c3f3aeff36 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 14:01:58 -0500 Subject: [PATCH 063/207] remove special case quant config Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_cutlass_moe.py | 15 +++++++++---- .../layers/fused_moe/oracle/fp8.py | 21 ------------------- .../compressed_tensors_moe.py | 20 +++++++----------- 3 files changed, 18 insertions(+), 38 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index 8d598587505e..2555db82b8e0 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -78,6 +78,13 @@ def __init__( # - skip input activation quantization (kernel applies scaling) self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized + # Preprocess scales for FlashInfer CUTLASS kernel. + if self.quant_config.is_per_tensor: + self._a1_gscale = 1.0 / quant_config.a1_scale + self._a2_gscale = 1.0 / quant_config.a2_scale + self._g1_alphas = (quant_config.w1_scale * quant_config.a1_scale).squeeze() + self._g2_alphas = (quant_config.w2_scale * quant_config.a2_scale).squeeze() + @staticmethod def expects_unquantized_inputs( moe_config: mk.FusedMoEConfig, quant_config: FusedMoEQuantConfig @@ -235,10 +242,10 @@ def apply( ): # FP8 per-tensor path: use global alphas/scales; do not pass input_sf quant_scales = [ - self.g1_alphas, # w13_weight_scale * w13_input_scale - self.a2_gscale, # 1.0 / w2_input_scale - self.g2_alphas, # w2_weight_scale * w2_input_scale - self.a1_scale, + self._g1_alphas, # w13_weight_scale * w13_input_scale + self._a2_gscale, # 1.0 / w2_input_scale + self._g2_alphas, # w2_weight_scale * w2_input_scale + self._a1_gscale, ] a1q_scale = None # not passing input_sf in fp8 diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index faac55af41ad..e7d2dbe65015 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -20,7 +20,6 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( FlashinferMoeBackend, get_flashinfer_moe_backend, - make_fp8_moe_alpha_scales_for_fi, prepare_fp8_moe_layer_for_fi, ) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( @@ -393,26 +392,6 @@ def make_fp8_moe_quant_config( block_shape=block_shape, ) - # Flashinfer CUTLASS per-tensor uses single dq scale - # (alpha = w_scale * a_scale) and inverse a2 scale. - if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS and block_shape is None: - assert a1_scale is not None and a2_scale is not None - g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi( - w1_scale, - a1_scale, - w2_scale, - a2_scale, - ) - return fp8_w8a8_moe_quant_config( - w1_scale=w1_scale, - w2_scale=w2_scale, - a1_scale=a1_scale, - a2_scale=a2_scale, - a1_gscale=(1.0 / a1_scale), - a2_gscale=(1.0 / a2_scale), - g1_alphas=g1_alphas, - g2_alphas=g2_alphas, - ) # All other backends use normal config. return fp8_w8a8_moe_quant_config( w1_scale=w1_scale, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 8a741baed808..a2161e5858ec 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -1009,21 +1009,15 @@ def select_gemm_impl( ) def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: - w1_scale = layer.w13_weight_scale - w2_scale = layer.w2_weight_scale - a1_scale = layer.w13_input_scale - a2_scale = layer.w2_input_scale - + is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN return make_fp8_moe_quant_config( fp8_backend=self.fp8_backend, - w1_scale=w1_scale, - w2_scale=w2_scale, - a1_scale=a1_scale, - a2_scale=a2_scale, - per_act_token_quant=( - self.input_quant.strategy == QuantizationStrategy.TOKEN - ), - per_out_ch_quant=(self.input_quant.strategy == QuantizationStrategy.TOKEN), + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + per_act_token_quant=is_per_token, + per_out_ch_quant=is_per_token, block_shape=self.weight_block_size, ) From 7ed64ef74b6e0de32896c50dfc1e598e9e288c26 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 14:04:59 -0500 Subject: [PATCH 064/207] remove flashinfer special case for fp8 config Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 4 ---- vllm/model_executor/layers/fused_moe/config.py | 8 -------- 2 files changed, 12 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 475ff5c7ecd2..bf7147ce66b1 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -266,13 +266,9 @@ def test_flashinfer_cutlass_moe_fp8_no_graph( quant_config = fp8_w8a8_moe_quant_config( w1_scale=td.w13_weight_scale, - g1_alphas=(td.w13_weight_scale * td.a1_scale).squeeze(), w2_scale=td.w2_weight_scale, - g2_alphas=(td.w2_weight_scale * td.a2_scale).squeeze(), a1_scale=td.a1_scale, - a1_gscale=td.a1_scale, a2_scale=td.a2_scale, - a2_gscale=1.0 / td.a2_scale, per_act_token_quant=False, ) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index a55030edf8e6..4a4eb819232c 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -521,10 +521,6 @@ def fp8_w8a8_moe_quant_config( per_act_token_quant: bool = False, per_out_ch_quant: bool = False, block_shape: list[int] | None = None, - a1_gscale: torch.Tensor | None = None, - a2_gscale: torch.Tensor | None = None, - g1_alphas: torch.Tensor | None = None, - g2_alphas: torch.Tensor | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for fp8 activations and fp8 weights. @@ -532,13 +528,9 @@ def fp8_w8a8_moe_quant_config( return FusedMoEQuantConfig.make( torch.float8_e4m3fn, w1_scale=w1_scale, - g1_alphas=g1_alphas, w2_scale=w2_scale, - g2_alphas=g2_alphas, a1_scale=a1_scale, - a1_gscale=a1_gscale, a2_scale=a2_scale, - a2_gscale=a2_gscale, per_act_token_quant=per_act_token_quant, per_out_ch_quant=per_out_ch_quant, block_shape=block_shape, From 855663780f3ca45832741c576e22b3826c56e4ab Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 14:05:26 -0500 Subject: [PATCH 065/207] updated docstring Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/config.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 4a4eb819232c..dddaf607a8af 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -463,10 +463,8 @@ def make( - a2_scale: Optional scale to be used for a2. - g1_alphas: Optional global quantization scales for w1 (for nvfp4). Optional per-channel scales for w1 (for W4A8 FP8). - Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8). - g2_alphas: Optional global quantization scales for w2 (for nvfp4). Optional per-channel scales for w2 (for W4A8 FP8). - Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8). - a1_gscale: Optional global quantization scales for a1 (1.0 /a2_scale). - a2_gscale: Optional global quantization scales for a2 (1.0 /a2_scale). From 7fccb8a419d249aaf3c04799273bdaa9ff2b3106 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 14:13:45 -0500 Subject: [PATCH 066/207] updated docstring Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_fp8_moe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 3a6b11746b4a..4380a43c80ee 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -50,7 +50,7 @@ def __init__( self._g1_alphas = (w1_scale * a1_scale).squeeze() self._g2_alphas = (w2_scale * a2_scale).squeeze() - self.g1_scale_c = self._g1_alphas / self.quant_config.a2_scale + self._g1_scale_c = self._g1_alphas / self.quant_config.a2_scale @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -247,7 +247,7 @@ def _apply_per_tensor_monolithic( routing_bias=e_score_correction_bias, hidden_states=hidden_states, gemm1_weights=w1, - output1_scales_scalar=self.g1_scale_c, + output1_scales_scalar=self._g1_scale_c, output1_scales_gate_scalar=self._g1_alphas, gemm2_weights=w2, output2_scales_scalar=self._g2_alphas, From 6c970153a3f52ffaea38e4bc2786c9ffeb0e2774 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 14:36:02 -0500 Subject: [PATCH 067/207] remove stray function Signed-off-by: Robert Shaw --- .../layers/quantization/utils/flashinfer_utils.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 35d58b63a0ff..cf9156c6916a 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -77,18 +77,6 @@ def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe( ) -def make_fp8_moe_alpha_scales_for_fi( - w13_scale: torch.Tensor, - w13_input_scale: torch.Tensor, - w2_scale: torch.Tensor, - w2_input_scale: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - g1_alphas = (w13_scale * w13_input_scale).squeeze() - g2_alphas = (w2_scale * w2_input_scale).squeeze() - - return g1_alphas, g2_alphas - - def build_flashinfer_fp8_cutlass_moe_prepare_finalize( moe: FusedMoEConfig | None, use_deepseek_fp8_block_scale: bool = False ) -> mk.FusedMoEPrepareAndFinalize: From 91067941d611a9bd6d5ea9a42a98d5cd56a16080 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 22:24:40 -0500 Subject: [PATCH 068/207] do the merge Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/prepare_finalize.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 6c43d26fa70a..60fb02a25e7b 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -90,7 +90,11 @@ def prepare( a1q, topk_weights, topk_ids, scales = res assert scales is not None and len(scales) == 1 a1q_scale = scales[0] - if quant_config.quant_dtype == "nvfp4": + # Apply swizzling after a2a if the MoE kernel needs it. + if ( + quant_config.quant_dtype == "nvfp4" + and quant_config.is_nvfp4_scale_swizzled + ): assert a1q_scale is not None if a1q_scale.element_size() == 1: a1q_scale = a1q_scale.view(torch.uint8) From 8e1f210f02ea4cfc1b872538b6b83c43467e521d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 22:27:11 -0500 Subject: [PATCH 069/207] remove old info Signed-off-by: Robert Shaw --- .../quantization/utils/flashinfer_fp4_moe.py | 89 ------------------- 1 file changed, 89 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 4b47c1791e8d..d062d662689d 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -7,17 +7,8 @@ import torch import vllm.envs as envs -import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEParallelConfig, - RoutingMethodType, -) from vllm.model_executor.layers.quantization.utils.quant_utils import ( - QuantKey, - kNvfp4Dynamic, - kNvfp4Static, swizzle_blockscale, ) from vllm.platforms import current_platform @@ -40,86 +31,6 @@ "reorder_w1w3_to_w3w1", ] -# -# Methods used by the oracle for kernel selection. -# - - -def _supports_current_device() -> bool: - """Supports only Blackwell-family GPUs.""" - p = current_platform - return p.is_cuda() and p.is_device_capability_family(100) - - -def _supports_no_act_and_mul() -> bool: - """Does not support non-gated MoE (i.e. Nemotron-Nano).""" - return False - - -def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, -) -> bool: - """Supports Nvfp4 quantization.""" - SUPPORTED_W_A = [ - (kNvfp4Static, kNvfp4Dynamic), - ] - return (weight_key, activation_key) in SUPPORTED_W_A - - -def _supports_activation(activation: str) -> bool: - """Supports silu activation only.""" - return activation in ["silu"] - - -def _supports_routing_method( - routing_method: RoutingMethodType, -) -> bool: - """Monolithic kernels need to express router support.""" - # NOTE(rob): potentially allow others here. This is a conservative list. - return routing_method in [ - RoutingMethodType.DeepSeekV3, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - RoutingMethodType.Llama4, - ] - - -def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """Supports EP.""" - return True - - -def is_supported_config_trtllm( - moe_config: FusedMoEConfig, - weight_key: QuantKey | None, - activation_key: QuantKey | None, - activation_format: mk.FusedMoEActivationFormat, -) -> tuple[bool, str | None]: - """ - This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config - """ - - def _make_reason(reason: str) -> str: - return f"kernel does not support {reason}" - - if not _supports_current_device(): - return False, _make_reason("current device") - elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()): - return False, _make_reason("no act_and_mul MLP layer") - elif not _supports_activation(moe_config.activation): - return False, _make_reason(f"{moe_config.activation} activation") - elif not _supports_quant_scheme(weight_key, activation_key): - return False, _make_reason("quantization scheme") - elif not _supports_parallel_config(moe_config.moe_parallel_config): - return False, _make_reason("parallel config") - elif not _supports_routing_method(moe_config.routing_method): - return False, _make_reason("routing method") - elif activation_format != mk.FusedMoEActivationFormat.Standard: - return False, _make_reason("activation format") - - return True, None - def is_flashinfer_fp4_cutlass_moe_available() -> bool: """Return `True` when FlashInfer CUTLASS NV-FP4 kernels can be used.""" From 7c1960eeef1687acfb18c286fd4ab76181979d7c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 22:39:15 -0500 Subject: [PATCH 070/207] remove the nvfp4 dp/ep hack for trtllm Signed-off-by: Robert Shaw --- .../layers/fused_moe/fused_moe_method_base.py | 12 ------ vllm/model_executor/layers/fused_moe/layer.py | 43 ++----------------- .../compressed_tensors_moe.py | 1 - .../layers/quantization/modelopt.py | 24 ----------- 4 files changed, 4 insertions(+), 76 deletions(-) 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 3ad56cc4c2de..9f99a8a24976 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 @@ -85,18 +85,6 @@ def select_gemm_impl( "implementation based on the prepare_finalize" ) - def prepare_dp_allgather_tensor( - self, - layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, list[torch.Tensor]]: - """Hook to prepare tensors and extra tensors for DP allgather + EP dispatch.""" - raise NotImplementedError( - "Method 'prepare_dp_allgather_tensor' is not implemented in " - f"{self.__class__.__name__}." - ) - @abstractmethod def get_fused_moe_quant_config( self, layer: torch.nn.Module diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index e3827500484b..521d4bc9f0d8 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -1797,44 +1797,13 @@ def forward_impl( ) with sp_ctx: - extra_tensors = None if do_naive_dispatch_combine: - post_quant_allgather = ( - self.quant_method is not None - and self.dp_size > 1 - and self.use_ep - and getattr(self.quant_method, "do_post_quant_allgather", False) - ) - if post_quant_allgather: - hidden_states_to_dispatch, extra_tensors = ( - self.quant_method.prepare_dp_allgather_tensor( - self, hidden_states, router_logits - ) - ) - else: - hidden_states_to_dispatch = hidden_states - dispatch_res = get_ep_group().dispatch_router_logits( - hidden_states_to_dispatch, + hidden_states, router_logits, self.is_sequence_parallel, - extra_tensors=extra_tensors, ) - if extra_tensors is not None: - ( - orig_hidden_states, - router_logits, - extra_tensors_combined, - ) = dispatch_res - hidden_states_combined = ( - orig_hidden_states, - extra_tensors_combined[0], - ) - else: - hidden_states_combined, router_logits = dispatch_res - orig_hidden_states = hidden_states_combined - else: - orig_hidden_states = hidden_states + hidden_states_combined, router_logits = dispatch_res # Run shared experts before matrix multiply. # because matrix multiply maybe modify the hidden_states. @@ -1858,10 +1827,6 @@ def forward_impl( # Matrix multiply. x = hidden_states_combined if do_naive_dispatch_combine else hidden_states - # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014). - # Figure out nicer way to do this. - x_orig = orig_hidden_states if do_naive_dispatch_combine else hidden_states - if self.quant_method.is_monolithic: final_hidden_states = self.quant_method.apply_monolithic( layer=self, @@ -1870,7 +1835,7 @@ def forward_impl( ) else: topk_weights, topk_ids = self.router.select_experts( - hidden_states=x_orig, + hidden_states=x, router_logits=router_logits, ) @@ -1879,7 +1844,7 @@ def forward_impl( final_hidden_states = self.quant_method.apply( layer=self, - x=x, # The type signture of this is wrong due to the hack. + x=x, topk_weights=topk_weights, topk_ids=topk_ids, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index c03ddb9b0e5e..acac3e0fca6b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -593,7 +593,6 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon def is_monolithic(self) -> bool: return ( self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - # NOTE(rob): this will not work until the Naive P/F is merged. and not self.moe.moe_parallel_config.use_all2all_kernels ) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 0cd6023f0c97..3f50f0b122c9 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1536,29 +1536,6 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: def do_post_quant_allgather(self): return self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - def prepare_dp_allgather_tensor( - self, - layer: FusedMoE, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, list[torch.Tensor]]: - """Optionally prepare extra tensors to carry through DP allgather/EP.""" - if self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM: - raise RuntimeError( - "prepare_dp_allgather_tensor is only supported for " - "FlashInfer TRTLLM NVFP4 MoE backend." - ) - - import flashinfer - - hidden_states_fp4, hidden_states_sf = flashinfer.fp4_quantize( - hidden_states, - layer.a1_gscale, - is_sf_swizzled_layout=False, - ) - extra_tensors: list[torch.Tensor] = [hidden_states_sf] - return hidden_states_fp4, extra_tensors - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: return make_nvfp4_moe_quant_config( backend=self.nvfp4_backend, @@ -1578,7 +1555,6 @@ def supports_eplb(self) -> bool: def is_monolithic(self) -> bool: return ( self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - # NOTE(rob): this will not work until the Naive P/F is merged. and not self.moe.moe_parallel_config.use_all2all_kernels ) From 9da1cc5de2ea58d48b8fa5462c1bc42b86118b30 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 22:51:50 -0500 Subject: [PATCH 071/207] apply new signature to nvfp4 case Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 8b11eae43653..b914b3c1d477 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -31,14 +31,6 @@ def __init__( super().__init__(moe_config=moe_config, quant_config=quant_config) self.routing_method_type = self.moe_config.routing_method - # self.routing_method_type = flashinfer.RoutingMethodType.Llama4 - # self.routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3 - - self.routing_bias = None - self.e_score_correction_bias = None - self.topk_group = None - self.num_expert_group = None - self.topk = moe_config.experts_per_token self.intermediate_size_per_partition = ( moe_config.intermediate_size_per_partition @@ -205,6 +197,11 @@ def apply_monolithic( expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, ) -> torch.Tensor: assert activation == "silu" assert a1q_scale is not None @@ -239,8 +236,8 @@ def apply_monolithic( output2_scale_scalar=self.quant_config.g2_alphas, num_experts=global_num_experts, top_k=self.topk, - n_group=self.num_expert_group if self.num_expert_group is not None else 0, - topk_group=self.topk_group if self.topk_group is not None else 0, + n_group=(num_expert_group or 0), + topk_group=(topk_group or 0), intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, From 301cf4930d7480747ece5769cedf80101c473ffb Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 22:54:18 -0500 Subject: [PATCH 072/207] make MNNVL work with TRTLLM Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_a2a_prepare_finalize.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py index 39b373861d03..473ccd3b6f0e 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py @@ -185,8 +185,8 @@ def flashinfer_alltoall_dispatch( ep_size, ) - # Swizzle after the A2A if nvfp4. - if quant_config.quant_dtype == "nvfp4": + # Swizzle after the A2A if nvfp4 and the MoE kernel expects swizzled scales. + if quant_config.quant_dtype == "nvfp4" and quant_config.is_nvfp4_scale_swizzled: if x_sf.element_size() == 1: x_sf = x_sf.view(torch.uint8) x_sf = nvfp4_block_scale_interleave(x_sf) From da40161c5dec5895ddafe4c92a8cc23e3883d128 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 23:02:33 -0500 Subject: [PATCH 073/207] revert FI cutlass quant config changes Signed-off-by: Robert Shaw --- .../model_executor/layers/fused_moe/config.py | 10 ++++++++++ .../fused_moe/flashinfer_cutlass_moe.py | 20 ++++--------------- .../layers/fused_moe/oracle/fp8.py | 15 ++++++++++++++ 3 files changed, 29 insertions(+), 16 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 447311549131..e8cbbd68467d 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -462,8 +462,10 @@ def make( - a2_scale: Optional scale to be used for a2. - g1_alphas: Optional global quantization scales for w1 (for nvfp4). Optional per-channel scales for w1 (for W4A8 FP8). + Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8). - g2_alphas: Optional global quantization scales for w2 (for nvfp4). Optional per-channel scales for w2 (for W4A8 FP8). + Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8). - a1_gscale: Optional global quantization scales for a1 (1.0 /a2_scale). - a2_gscale: Optional global quantization scales for a2 (1.0 /a2_scale). @@ -518,6 +520,10 @@ def fp8_w8a8_moe_quant_config( per_act_token_quant: bool = False, per_out_ch_quant: bool = False, block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, + g1_alphas: torch.Tensor | None = None, + g2_alphas: torch.Tensor | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for fp8 activations and fp8 weights. @@ -525,9 +531,13 @@ def fp8_w8a8_moe_quant_config( return FusedMoEQuantConfig.make( torch.float8_e4m3fn, w1_scale=w1_scale, + g1_alphas=g1_alphas, w2_scale=w2_scale, + g2_alphas=g2_alphas, a1_scale=a1_scale, + a1_gscale=a1_gscale, a2_scale=a2_scale, + a2_gscale=a2_gscale, per_act_token_quant=per_act_token_quant, per_out_ch_quant=per_out_ch_quant, block_shape=block_shape, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index 5f269925ceb4..ba19f090ebe3 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -78,18 +78,6 @@ def __init__( # - skip input activation quantization (kernel applies scaling) self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized - # Preprocess scales for FlashInfer CUTLASS kernel. - if self.quant_config.is_per_tensor: - assert quant_config.w1_scale is not None - assert quant_config.w2_scale is not None - assert quant_config.a1_scale is not None - assert quant_config.a2_scale is not None - - self._a1_gscale = 1.0 / quant_config.a1_scale - self._a2_gscale = 1.0 / quant_config.a2_scale - self._g1_alphas = (quant_config.w1_scale * quant_config.a1_scale).squeeze() - self._g2_alphas = (quant_config.w2_scale * quant_config.a2_scale).squeeze() - @property def expects_unquantized_inputs(self) -> bool: return self.quant_config.use_fp8_w8a8 and self.quant_config.is_block_quantized @@ -239,10 +227,10 @@ def apply( ): # FP8 per-tensor path: use global alphas/scales; do not pass input_sf quant_scales = [ - self._g1_alphas, # w13_weight_scale * w13_input_scale - self._a2_gscale, # 1.0 / w2_input_scale - self._g2_alphas, # w2_weight_scale * w2_input_scale - self._a1_gscale, + self.g1_alphas, # w13_weight_scale * w13_input_scale + self.a2_gscale, # 1.0 / w2_input_scale + self.g2_alphas, # w2_weight_scale * w2_input_scale + self.a1_gscale, ] a1q_scale = None # not passing input_sf in fp8 diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index f91ee0f12502..eb9ebff9c116 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -400,6 +400,21 @@ def make_fp8_moe_quant_config( block_shape=block_shape, ) + # Flashinfer CUTLASS per-tensor uses single dq scale + # (alpha = w_scale * a_scale) and inverse a2 scale. + if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS and block_shape is None: + assert a1_scale is not None and a2_scale is not None + return fp8_w8a8_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + a1_gscale=(1.0 / a1_scale), + a2_gscale=(1.0 / a2_scale), + g1_alphas=(w1_scale * a1_scale).squeeze(), + g2_alphas=(w2_scale * a2_scale).squeeze(), + ) + # All other backends use normal config. return fp8_w8a8_moe_quant_config( w1_scale=w1_scale, From 1bb5cc6d635517efb3ac4f58615088d2f042ba4c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 23:04:05 -0500 Subject: [PATCH 074/207] tests nit Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 8cd438c3f2db..113665058069 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -17,6 +17,9 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, ) +from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( + FlashInferTrtLlmFp8Experts, +) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, @@ -207,10 +210,6 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( quant_config=quant_config, ) - from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( - FlashInferTrtLlmFp8Experts, - ) - kernel = mk.FusedMoEModularKernel( MoEPrepareAndFinalizeNoEP(), FlashInferTrtLlmFp8Experts( From 7ba5feb79b25967f2580ca4398a2b0588e7dff4a Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 23:04:55 -0500 Subject: [PATCH 075/207] tests nit Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 113665058069..0bbf4ae2687f 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -265,9 +265,13 @@ def test_flashinfer_cutlass_moe_fp8_no_graph( quant_config = fp8_w8a8_moe_quant_config( w1_scale=td.w13_weight_scale, + g1_alphas=(td.w13_weight_scale * td.a1_scale).squeeze(), w2_scale=td.w2_weight_scale, + g2_alphas=(td.w2_weight_scale * td.a2_scale).squeeze(), a1_scale=td.a1_scale, + a1_gscale=td.a1_scale, a2_scale=td.a2_scale, + a2_gscale=1.0 / td.a2_scale, per_act_token_quant=False, ) From 999de41035f79d1c68e52d4260754df75326a407 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 23:05:54 -0500 Subject: [PATCH 076/207] tests nit Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_a2a_prepare_finalize.py | 2 +- vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py index 473ccd3b6f0e..2857e51a3a72 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py @@ -185,7 +185,7 @@ def flashinfer_alltoall_dispatch( ep_size, ) - # Swizzle after the A2A if nvfp4 and the MoE kernel expects swizzled scales. + # Swizzle after the A2A if MoE kernel expects swizzled scales. if quant_config.quant_dtype == "nvfp4" and quant_config.is_nvfp4_scale_swizzled: if x_sf.element_size() == 1: x_sf = x_sf.view(torch.uint8) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index ba19f090ebe3..faa654ea3a2f 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -230,7 +230,7 @@ def apply( self.g1_alphas, # w13_weight_scale * w13_input_scale self.a2_gscale, # 1.0 / w2_input_scale self.g2_alphas, # w2_weight_scale * w2_input_scale - self.a1_gscale, + self.a1_scale, ] a1q_scale = None # not passing input_sf in fp8 From 648e03082052ade385b91be03c7b7e62ef9e9fa8 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 26 Jan 2026 23:53:30 -0500 Subject: [PATCH 077/207] make the monolithic kernels work with naive P/F Signed-off-by: Robert Shaw --- .../layers/fused_moe/modular_kernel.py | 38 ++++- .../layers/fused_moe/prepare_finalize.py | 132 ++++++++++++++---- 2 files changed, 134 insertions(+), 36 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 9842c274a0ee..8cc39d86e0f9 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -163,6 +163,7 @@ def apply( PrepareMonolithicResultType = tuple[ torch.Tensor, torch.Tensor | None, + torch.Tensor, ] ReceiverType = Callable[[], PrepareResultType] @@ -255,10 +256,14 @@ def prepare( def prepare_monolithic( self, a1: torch.Tensor, + router_logits: torch.Tensor, quant_config: FusedMoEQuantConfig, defer_input_quant: bool = False, ) -> PrepareMonolithicResultType: """ + Optional method for subclasses compatible with monolithic + FusedMoEPermuteExpertsUnpermute kernels. + Perform any quantization (and/or) dispatching needed for this kernel. - a1: The (unquantized) input to the MoE layer. - quant_config: Quantization info provided by the fused experts. @@ -350,6 +355,24 @@ def finalize( """ raise NotImplementedError + def finalize_monolithic( + self, + fused_expert_output: torch.Tensor, + weight_and_reduce_impl: TopKWeightAndReduce, + ) -> torch.Tensor: + """ + Optional method for subclasses compatible with monolithic + FusedMoEPermuteExpertsUnpermute kernels. + + Perform any combine plus apply weights and perform a reduction on the + fused experts output. + - fused_expert_output: The unweighted, unreduced output of the fused + experts, it will have (M, topk, K) shape. + - weight_and_reduce_impl: An optional TopKWeightAndReduce + implementation. + """ + raise NotImplementedError + def finalize_async( self, output: torch.Tensor, @@ -1490,10 +1513,11 @@ def forward_monolithic( to the topk_ids and topk_weights. This is used for kernels that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ - - a1q, a1q_scale = self.prepare_finalize.prepare_monolithic( + # TODO(rob): add inplace support. + a1q, a1q_scale, router_logits = self.prepare_finalize.prepare_monolithic( hidden_states, - self.fused_experts.quant_config, + router_logits=router_logits, + quant_config=self.fused_experts.quant_config, defer_input_quant=self.fused_experts.expects_unquantized_inputs, ) @@ -1514,7 +1538,9 @@ def forward_monolithic( topk_group=topk_group, ) - # TODO(rob): once naive P/F Dp/Ep lands, we will need to - # add support here for finalizing over just the hidden states. + output = self.prepare_finalize.finalize_monolithic( + fused_expert_output=fused_out, + weight_and_reduce_impl=self.fused_experts.finalize_weight_and_reduce_impl(), + ) - return fused_out + return output diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 60fb02a25e7b..878419ef8df3 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -9,6 +9,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate, + TopKWeightAndReduceNoOP, ) from vllm.utils.flashinfer import nvfp4_block_scale_interleave @@ -39,25 +40,12 @@ def num_dispatchers(self) -> int: def output_is_reduced(self) -> bool: return False - def prepare( + def _quantize_and_setup_dispatch( self, a1: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, quant_config: FusedMoEQuantConfig, defer_input_quant: bool = False, - ) -> mk.PrepareResultType: - if apply_router_weight_on_input: - topk = topk_ids.size(1) - assert topk == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - # Note: do not use inplace for shared experts overlap - a1 = a1 * topk_weights.to(a1.dtype) - + ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: # Defer input quantization to the MoE kernel. if defer_input_quant: a1q = a1 @@ -77,6 +65,77 @@ def prepare( skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0 scales = None if skip_gather_scales else [a1q_scale] + return a1q, scales + + def _unwrap_scale_and_prepare_for_moe( + self, + scales: list[torch.Tensor] | None, + quant_config: FusedMoEQuantConfig, + ) -> torch.Tensor: + assert scales is not None and len(scales) == 1 + a1q_scale = scales[0] + # Apply swizzling after a2a if the MoE kernel needs it. + if quant_config.quant_dtype == "nvfp4" and quant_config.is_nvfp4_scale_swizzled: + assert a1q_scale is not None + if a1q_scale.element_size() == 1: + a1q_scale = a1q_scale.view(torch.uint8) + a1q_scale = nvfp4_block_scale_interleave(a1q_scale) + + return a1q_scale + + def prepare_monolithic( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareMonolithicResultType: + """Quantize and Dispatch Topk Weights and Topk Ids.""" + + a1q, scales = self._quantize_and_setup_dispatch( + a1, quant_config, defer_input_quant + ) + + res = get_ep_group().dispatch_router_logits( + a1q, + router_logits, + is_sequence_parallel=self.is_sequence_parallel, + extra_tensors=scales, + ) + + if scales is None: + a1q, router_logits = res + else: + a1q, router_logits, scales = res + a1q_scale = self._unwrap_scale_and_prepare_for_moe(scales, quant_config) + + return a1q, a1q_scale, router_logits + + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareResultType: + """Quantize and Dispatch Topk Weights and Topk Ids.""" + + if apply_router_weight_on_input: + topk = topk_ids.size(1) + assert topk == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + # Note: do not use inplace for shared experts overlap + a1 = a1 * topk_weights.to(a1.dtype) + + a1q, scales = self._quantize_and_setup_dispatch( + a1, quant_config, defer_input_quant + ) + res = get_ep_group().dispatch( a1q, topk_weights, @@ -84,21 +143,12 @@ def prepare( is_sequence_parallel=self.is_sequence_parallel, extra_tensors=scales, ) - if skip_gather_scales: + + if scales is None: a1q, topk_weights, topk_ids = res else: a1q, topk_weights, topk_ids, scales = res - assert scales is not None and len(scales) == 1 - a1q_scale = scales[0] - # Apply swizzling after a2a if the MoE kernel needs it. - if ( - quant_config.quant_dtype == "nvfp4" - and quant_config.is_nvfp4_scale_swizzled - ): - assert a1q_scale is not None - if a1q_scale.element_size() == 1: - a1q_scale = a1q_scale.view(torch.uint8) - a1q_scale = nvfp4_block_scale_interleave(a1q_scale) + a1q_scale = self._unwrap_scale_and_prepare_for_moe(scales, quant_config) return a1q, a1q_scale, None, topk_ids, topk_weights @@ -126,6 +176,18 @@ def finalize( get_ep_group().combine(out, is_sequence_parallel=self.is_sequence_parallel) ) + def finalize_monolithic( + self, + fused_expert_output: torch.Tensor, + weight_and_reduce_impl: mk.TopKWeightAndReduce, + ) -> torch.Tensor: + assert weight_and_reduce_impl == TopKWeightAndReduceNoOP + out = get_ep_group().combine( + fused_expert_output, is_sequence_parallel=self.is_sequence_parallel + ) + assert isinstance(out, torch.Tensor) + return out + class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize): @property @@ -176,15 +238,17 @@ def prepare( def prepare_monolithic( self, a1: torch.Tensor, + router_logits: torch.Tensor, quant_config: FusedMoEQuantConfig, defer_input_quant: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None]: + ) -> mk.PrepareMonolithicResultType: # Defer input quant to moe kernel for backends (e.g. AITER, FI) # which use a single kernel call for quant + experts. if defer_input_quant: - return a1, None + return a1, None, router_logits - return self._quantize_input(a1, quant_config) + a1q, a1q_scale = self._quantize_input(a1, quant_config) + return a1q, a1q_scale, router_logits def finalize( self, @@ -204,3 +268,11 @@ def finalize( topk_ids=topk_ids, apply_router_weight_on_input=apply_router_weight_on_input, ) + + def finalize_monolithic( + self, + fused_expert_output: torch.Tensor, + weight_and_reduce_impl: mk.TopKWeightAndReduce, + ) -> torch.Tensor: + assert weight_and_reduce_impl == TopKWeightAndReduceNoOP + return fused_expert_output From b0c0fb48645aebca1f00567c6e38ed897690e61d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 00:05:48 -0500 Subject: [PATCH 078/207] update guards for monolithic mk Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_fp8_moe.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 4380a43c80ee..62f65a5e1e27 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -107,8 +107,11 @@ def _supports_routing_method( @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """TRTLLMGenKernel does not support EPLB.""" - return not moe_parallel_config.enable_eplb + """TRTLLMGenKernel is monolithic, so it only supports TP or naive DP/EP.""" + return not moe_parallel_config.use_all2all_kernels or ( + moe_parallel_config.use_naive_all2all_kernels + and not moe_parallel_config.enable_eplb + ) def supports_chunking(self) -> bool: return False From bbd9190180cfed3830b867cc26507f6c6cb500d3 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 00:14:44 -0500 Subject: [PATCH 079/207] revert quantize input Signed-off-by: Robert Shaw --- .../fused_moe/deepep_ht_prepare_finalize.py | 18 ++++++++++-- .../layers/fused_moe/modular_kernel.py | 24 ---------------- .../layers/fused_moe/pplx_prepare_finalize.py | 9 +++++- .../layers/fused_moe/prepare_finalize.py | 28 ++++++++++++++++--- 4 files changed, 48 insertions(+), 31 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py index 386d7726edf6..514aa205a3cb 100644 --- a/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py @@ -11,6 +11,7 @@ TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate, ) +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, @@ -231,7 +232,14 @@ def _receiver( # Quantize after dispatch. expert_x_scale = None if expert_x.numel() != 0: - expert_x, expert_x_scale = self._quantize_input(expert_x, quant_config) + # TODO: support per_act_token_quant, + expert_x, expert_x_scale = moe_kernel_quantize_input( + expert_x, + a1_scale, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=False, + block_shape=quant_config.block_shape, + ) return ( expert_x, @@ -269,7 +277,13 @@ def prepare_async( # * For expert kernels that require unquantized inputs, # defer quantization to FusedMoEExpertsPermuteUnpermute. if quant_config.is_block_quantized and not defer_input_quant: - a1q, a1q_scale = self._quantize_input(a1, quant_config) + a1q, a1q_scale = moe_kernel_quantize_input( + a1, + quant_config.a1_scale, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + ) if a1q_scale is not None and a1q_scale.numel() == 1: a1q_scale = a1q_scale.view(1, 1) a1_post_scale = None diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 8cc39d86e0f9..4e5d165cb8bf 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -23,7 +23,6 @@ apply_moe_activation, count_expert_num_tokens, disable_inplace, - moe_kernel_quantize_input, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -192,29 +191,6 @@ def supports_async(self) -> bool: """ return False - def _quantize_input( - self, - a1: torch.Tensor, - quant_config: FusedMoEQuantConfig, - skip_fp4_swizzle: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - a1_scale = ( - quant_config.a1_gscale - if quant_config.use_nvfp4_w4a4 - else quant_config.a1_scale - ) - - return moe_kernel_quantize_input( - a1, - a1_scale, - quant_config.quant_dtype, - quant_config.per_act_token_quant, - quant_config.block_shape, - is_fp4_scale_swizzled=( - quant_config.is_nvfp4_scale_swizzled and not skip_fp4_swizzle - ), - ) - @abstractmethod def prepare( self, diff --git a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py index 26a95f818aeb..72545ed084a8 100644 --- a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py @@ -13,6 +13,7 @@ ) from vllm.model_executor.layers.fused_moe.utils import ( _validate_scale_shape, + moe_kernel_quantize_input, ) from vllm.utils.math_utils import cdiv, round_up @@ -142,7 +143,13 @@ def prepare_async( repeat_cols = 4 repeat_rows = 1 if quant_config.per_act_token_quant else a1.size(0) - a1q, a1q_scale = self._quantize_input(a1, quant_config) + a1q, a1q_scale = moe_kernel_quantize_input( + a1, + quant_config.a1_scale, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + ) _validate_scale_shape( a1q, a1q_scale, quant_config.per_act_token_quant, quant_config.block_shape diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 878419ef8df3..242c7bc64d43 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -11,6 +11,7 @@ TopKWeightAndReduceDelegate, TopKWeightAndReduceNoOP, ) +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.flashinfer import nvfp4_block_scale_interleave @@ -55,8 +56,13 @@ def _quantize_and_setup_dispatch( # which makes the scales tensor different shape than # the hidden states, breaking the A2A kernel. So, we # delay the swizzling until after the A2A. - a1q, a1q_scale = self._quantize_input( - a1, quant_config, skip_fp4_swizzle=True + a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input( + a1, + quant_config.a1_scale, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + is_fp4_scale_swizzled=False, ) # Skip gathering scales if we have static quantization @@ -231,7 +237,14 @@ def prepare( if defer_input_quant: return a1, None, None, None, None - a1q, a1q_scale = self._quantize_input(a1, quant_config) + a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input( + a1, + quant_config.a1_scale, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, + ) return a1q, a1q_scale, None, None, None @@ -247,7 +260,14 @@ def prepare_monolithic( if defer_input_quant: return a1, None, router_logits - a1q, a1q_scale = self._quantize_input(a1, quant_config) + a1q, a1q_scale = moe_kernel_quantize_input( + a1, + quant_config.a1_scale, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, + ) return a1q, a1q_scale, router_logits def finalize( From 15d2338421a1a9d70551457aa22bd5a143436f92 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 00:19:27 -0500 Subject: [PATCH 080/207] reduce LOC changed Signed-off-by: Robert Shaw --- .../layers/fused_moe/modular_kernel.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 4e5d165cb8bf..2d1adb4d1f0c 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -184,13 +184,6 @@ def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"): """ return - def supports_async(self) -> bool: - """ - Indicates whether or not this class implements prepare_async and - finalize_async. - """ - return False - @abstractmethod def prepare( self, @@ -254,6 +247,13 @@ def prepare_monolithic( f"prepare_monolithic not supported for {self.__class__.__name__}" ) + def supports_async(self) -> bool: + """ + Indicates whether or not this class implements prepare_async and + finalize_async. + """ + return False + def prepare_async( self, a1: torch.Tensor, From 71c1bdec7efb07e3852d4c7404e219d75f7f7386 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 00:20:16 -0500 Subject: [PATCH 081/207] reduce LOC changed Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py index 72545ed084a8..78b941498062 100644 --- a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py @@ -142,10 +142,10 @@ def prepare_async( repeat_cols = 4 repeat_rows = 1 if quant_config.per_act_token_quant else a1.size(0) - + # TODO(bnell): always pass quant_config.a1_scale? a1q, a1q_scale = moe_kernel_quantize_input( a1, - quant_config.a1_scale, + (None if quant_config.per_act_token_quant else quant_config.a1_scale), quant_dtype=quant_config.quant_dtype, per_act_token_quant=quant_config.per_act_token_quant, block_shape=quant_config.block_shape, From b6963ced12727991d07b9aad20099abc31af33ba Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 00:22:13 -0500 Subject: [PATCH 082/207] reduce LOC changed Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 84027e6e152f..27e7e788c338 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -167,7 +167,7 @@ def _return_or_raise( k_cls, config, weight_key, activation_key, activation_format ) if supported: - logger.info_once(_make_log_backend(backend), scope="local") + logger.info_once(_make_log_backend(backend)) return backend, k_cls raise ValueError(_make_log_unsupported(backend, reason)) From 19b2a48e6d77259732253ffafca792546cc755f0 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 16:46:04 -0500 Subject: [PATCH 083/207] fix pre-ciommits Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py | 4 +++- .../compressed_tensors/compressed_tensors_moe.py | 8 ++++---- vllm/model_executor/layers/quantization/fp8.py | 4 ++-- vllm/model_executor/layers/quantization/ipex_quant.py | 6 ------ vllm/model_executor/layers/quantization/modelopt.py | 8 ++++---- 5 files changed, 13 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index b914b3c1d477..9724489e9884 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -205,9 +205,11 @@ def apply_monolithic( ) -> torch.Tensor: assert activation == "silu" assert a1q_scale is not None + assert self.quant_config.w1_scale is not None + assert self.quant_config.w2_scale is not None # Prepare routing bias into kernel format. - routing_bias = self.e_score_correction_bias + routing_bias = e_score_correction_bias if routing_bias is not None: routing_bias = routing_bias.to(torch.bfloat16) router_logits = ( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index acac3e0fca6b..5e31b3feda01 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -602,8 +602,8 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.kernel is not None - return self.kernel.forward_monolithic( + assert self.moe_mk is not None + return self.moe_mk.forward_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -977,8 +977,8 @@ def apply_monolithic( assert self.is_monolithic assert layer.activation == "silu" - assert self.kernel is not None - return self.kernel.forward_monolithic( + assert self.moe_mk is not None + return self.moe_mk.forward_monolithic( x, layer.w13_weight, layer.w2_weight, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index c4433853a77c..06c44d9fd662 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -940,8 +940,8 @@ def apply_monolithic( if layer.enable_eplb: raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.") - assert self.kernel is not None - return self.kernel.forward_monolithic( + assert self.moe_mk is not None + return self.moe_mk.forward_monolithic( x, layer.w13_weight, layer.w2_weight, diff --git a/vllm/model_executor/layers/quantization/ipex_quant.py b/vllm/model_executor/layers/quantization/ipex_quant.py index 119fb2ef82d8..ac56a3dbf917 100644 --- a/vllm/model_executor/layers/quantization/ipex_quant.py +++ b/vllm/model_executor/layers/quantization/ipex_quant.py @@ -8,7 +8,6 @@ from torch.nn import Module from vllm._ipex_ops import ipex_ops as ops -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.linear import ( LinearBase, LinearMethodBase, @@ -376,11 +375,6 @@ def process_weights_after_loading(self, layer: Module) -> None: experts_start_id=ep_rank_start, ) - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - return None - @property def is_monolithic(self) -> bool: return True diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 3f50f0b122c9..a4dbf69ffac0 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -925,8 +925,8 @@ def apply_monolithic( raise NotImplementedError( "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend." ) - assert self.kernel is not None - return self.kernel.forward_monolithic( + assert self.moe_mk is not None + return self.moe_mk.forward_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -1564,8 +1564,8 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.kernel is not None - return self.kernel.forward_monolithic( + assert self.moe_mk is not None + return self.moe_mk.forward_monolithic( x, layer.w13_weight, layer.w2_weight, From e9e85e261d3a364bc84926938660568f2b4a67d8 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 17:10:48 -0500 Subject: [PATCH 084/207] qwen nvfp4 working across all cases Signed-off-by: Robert Shaw --- .../layers/fused_moe/oracle/fp8.py | 1 - .../layers/fused_moe/prepare_finalize.py | 26 +++++++++++++++---- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index eb9ebff9c116..6b44c4df8096 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -414,7 +414,6 @@ def make_fp8_moe_quant_config( g1_alphas=(w1_scale * a1_scale).squeeze(), g2_alphas=(w2_scale * a2_scale).squeeze(), ) - # All other backends use normal config. return fp8_w8a8_moe_quant_config( w1_scale=w1_scale, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 242c7bc64d43..5fd5882d3d0f 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -52,13 +52,19 @@ def _quantize_and_setup_dispatch( a1q = a1 a1q_scale = None else: + input_sf = ( + quant_config.a1_gscale + if quant_config.use_nvfp4_w4a4 + else quant_config.a1_scale + ) + # NOTE: swizzling pads the scales to multiple of 128 # which makes the scales tensor different shape than # the hidden states, breaking the A2A kernel. So, we # delay the swizzling until after the A2A. a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input( a1, - quant_config.a1_scale, + input_sf, quant_dtype=quant_config.quant_dtype, per_act_token_quant=quant_config.per_act_token_quant, block_shape=quant_config.block_shape, @@ -187,7 +193,7 @@ def finalize_monolithic( fused_expert_output: torch.Tensor, weight_and_reduce_impl: mk.TopKWeightAndReduce, ) -> torch.Tensor: - assert weight_and_reduce_impl == TopKWeightAndReduceNoOP + assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP) out = get_ep_group().combine( fused_expert_output, is_sequence_parallel=self.is_sequence_parallel ) @@ -237,9 +243,14 @@ def prepare( if defer_input_quant: return a1, None, None, None, None + input_sf = ( + quant_config.a1_gscale + if quant_config.use_nvfp4_w4a4 + else quant_config.a1_scale + ) a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input( a1, - quant_config.a1_scale, + input_sf, quant_dtype=quant_config.quant_dtype, per_act_token_quant=quant_config.per_act_token_quant, block_shape=quant_config.block_shape, @@ -260,9 +271,14 @@ def prepare_monolithic( if defer_input_quant: return a1, None, router_logits + input_sf = ( + quant_config.a1_gscale + if quant_config.use_nvfp4_w4a4 + else quant_config.a1_scale + ) a1q, a1q_scale = moe_kernel_quantize_input( a1, - quant_config.a1_scale, + input_sf, quant_dtype=quant_config.quant_dtype, per_act_token_quant=quant_config.per_act_token_quant, block_shape=quant_config.block_shape, @@ -294,5 +310,5 @@ def finalize_monolithic( fused_expert_output: torch.Tensor, weight_and_reduce_impl: mk.TopKWeightAndReduce, ) -> torch.Tensor: - assert weight_and_reduce_impl == TopKWeightAndReduceNoOP + assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP) return fused_expert_output From d737e4b10af3e46a5ff803099a26cdc935f0fe91 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 17:17:52 -0500 Subject: [PATCH 085/207] fix deepep high throughput with nvfp4 Signed-off-by: Robert Shaw --- .../layers/fused_moe/deepep_ht_prepare_finalize.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py index 514aa205a3cb..c4a227223703 100644 --- a/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py @@ -239,6 +239,7 @@ def _receiver( quant_dtype=quant_config.quant_dtype, per_act_token_quant=False, block_shape=quant_config.block_shape, + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, ) return ( From 564ad9bb23d091ed91ce8845b3463c43f059d55e Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 18:15:48 -0500 Subject: [PATCH 086/207] updated Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 20 +++++-- .../layers/fused_moe/prepare_finalize.py | 57 ++++++++++--------- 2 files changed, 43 insertions(+), 34 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 62f65a5e1e27..ad9ab20b24a5 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -11,6 +11,9 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, @@ -28,6 +31,12 @@ def __init__( ): super().__init__(moe_config, quant_config) + if moe_config.moe_parallel_config.use_ep and quant_config.is_per_tensor: + raise NotImplementedError( + "EP parallelism is not supported with TRTLLM" + "per-tensor FP8 quantization." + ) + self.routing_method_type = moe_config.routing_method self.topk = moe_config.experts_per_token self.intermediate_size_per_partition = ( @@ -120,9 +129,7 @@ def supports_expert_map(self) -> bool: return False def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - raise NotImplementedError( - f"{self.__class__.__name__} only supports the apply_monolithic interface." - ) + return TopKWeightAndReduceNoOP() def workspace_shapes( self, @@ -241,11 +248,11 @@ def _apply_per_tensor_monolithic( assert apply_router_weight_on_input # Should only have Llama4 routing here. - assert routed_scaling_factor is None + assert routed_scaling_factor is not None assert e_score_correction_bias is None assert num_expert_group is None - return flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( + out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, @@ -256,7 +263,7 @@ def _apply_per_tensor_monolithic( output2_scales_scalar=self._g2_alphas, num_experts=global_num_experts, top_k=self.topk, - n_group=num_expert_group, + n_group=num_expert_group or 0, topk_group=topk_group or 0, intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, @@ -265,6 +272,7 @@ def _apply_per_tensor_monolithic( use_routing_scales_on_input=apply_router_weight_on_input, routing_method_type=self.routing_method_type, ) + return out def apply_monolithic( self, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 5fd5882d3d0f..ee7035f84fe8 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -95,34 +95,6 @@ def _unwrap_scale_and_prepare_for_moe( return a1q_scale - def prepare_monolithic( - self, - a1: torch.Tensor, - router_logits: torch.Tensor, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> mk.PrepareMonolithicResultType: - """Quantize and Dispatch Topk Weights and Topk Ids.""" - - a1q, scales = self._quantize_and_setup_dispatch( - a1, quant_config, defer_input_quant - ) - - res = get_ep_group().dispatch_router_logits( - a1q, - router_logits, - is_sequence_parallel=self.is_sequence_parallel, - extra_tensors=scales, - ) - - if scales is None: - a1q, router_logits = res - else: - a1q, router_logits, scales = res - a1q_scale = self._unwrap_scale_and_prepare_for_moe(scales, quant_config) - - return a1q, a1q_scale, router_logits - def prepare( self, a1: torch.Tensor, @@ -188,6 +160,35 @@ def finalize( get_ep_group().combine(out, is_sequence_parallel=self.is_sequence_parallel) ) + def prepare_monolithic( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareMonolithicResultType: + """Quantize and Dispatch Router Logits.""" + + a1q, scales = self._quantize_and_setup_dispatch( + a1, quant_config, defer_input_quant + ) + + res = get_ep_group().dispatch_router_logits( + a1q, + router_logits, + is_sequence_parallel=self.is_sequence_parallel, + extra_tensors=scales, + ) + + if scales is None: + a1q, router_logits = res + a1q_scale = None + else: + a1q, router_logits, scales = res + a1q_scale = self._unwrap_scale_and_prepare_for_moe(scales, quant_config) + + return a1q, a1q_scale, router_logits + def finalize_monolithic( self, fused_expert_output: torch.Tensor, From 55d395a38f450fd53b6ba1b402f2bf08565811b9 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 18:23:59 -0500 Subject: [PATCH 087/207] reorder some stuff Signed-off-by: Robert Shaw --- .../layers/fused_moe/modular_kernel.py | 86 +++++++++---------- 1 file changed, 43 insertions(+), 43 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 2d1adb4d1f0c..2a07c20ace4b 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -222,31 +222,6 @@ def prepare( """ raise NotImplementedError - def prepare_monolithic( - self, - a1: torch.Tensor, - router_logits: torch.Tensor, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> PrepareMonolithicResultType: - """ - Optional method for subclasses compatible with monolithic - FusedMoEPermuteExpertsUnpermute kernels. - - Perform any quantization (and/or) dispatching needed for this kernel. - - a1: The (unquantized) input to the MoE layer. - - quant_config: Quantization info provided by the fused experts. - - defer_input_quant: Runtime parameter indicating whether or not to - defer input quantization to the FusedMoEPermuteExpertsUnpermute - - Returns a tuple of: - - quantized + dispatched a. - - Optional quantized + dispatched a1_scales. - """ - raise NotImplementedError( - f"prepare_monolithic not supported for {self.__class__.__name__}" - ) - def supports_async(self) -> bool: """ Indicates whether or not this class implements prepare_async and @@ -331,24 +306,6 @@ def finalize( """ raise NotImplementedError - def finalize_monolithic( - self, - fused_expert_output: torch.Tensor, - weight_and_reduce_impl: TopKWeightAndReduce, - ) -> torch.Tensor: - """ - Optional method for subclasses compatible with monolithic - FusedMoEPermuteExpertsUnpermute kernels. - - Perform any combine plus apply weights and perform a reduction on the - fused experts output. - - fused_expert_output: The unweighted, unreduced output of the fused - experts, it will have (M, topk, K) shape. - - weight_and_reduce_impl: An optional TopKWeightAndReduce - implementation. - """ - raise NotImplementedError - def finalize_async( self, output: torch.Tensor, @@ -391,6 +348,49 @@ def finalize_async( """ raise NotImplementedError + def prepare_monolithic( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> PrepareMonolithicResultType: + """ + Optional method for subclasses compatible with monolithic + FusedMoEPermuteExpertsUnpermute kernels. + + Perform any quantization (and/or) dispatching needed for this kernel. + - a1: The (unquantized) input to the MoE layer. + - quant_config: Quantization info provided by the fused experts. + - defer_input_quant: Runtime parameter indicating whether or not to + defer input quantization to the FusedMoEPermuteExpertsUnpermute + + Returns a tuple of: + - quantized + dispatched a. + - Optional quantized + dispatched a1_scales. + """ + raise NotImplementedError( + f"prepare_monolithic not supported for {self.__class__.__name__}" + ) + + def finalize_monolithic( + self, + fused_expert_output: torch.Tensor, + weight_and_reduce_impl: TopKWeightAndReduce, + ) -> torch.Tensor: + """ + Optional method for subclasses compatible with monolithic + FusedMoEPermuteExpertsUnpermute kernels. + + Perform any combine plus apply weights and perform a reduction on the + fused experts output. + - fused_expert_output: The unweighted, unreduced output of the fused + experts, it will have (M, topk, K) shape. + - weight_and_reduce_impl: An optional TopKWeightAndReduce + implementation. + """ + raise NotImplementedError + @property @abstractmethod def activation_format(self) -> FusedMoEActivationFormat: From def9c8612833eafc1e342edbf17dfb78890cd422 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 27 Jan 2026 18:44:11 -0500 Subject: [PATCH 088/207] Remove debug cruft Signed-off-by: Robert Shaw --- vllm/model_executor/models/llama4.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index edb424f382e2..0cdb4989ec73 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -436,7 +436,6 @@ def load_moe_expert_weights( # Whether the MoE expert weights are loaded successfully. expert_param_loaded = False - loaded_weight = loaded_weight.to("cuda") # If fused is True, the loaded weight is in the layout of: # [num_experts, hidden_in, hidden_out], so we must transpose the last From 89465f3afea7213fde1ac8f7a269ec9f8a8b07d1 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Thu, 29 Jan 2026 13:33:33 -0500 Subject: [PATCH 089/207] updated typing Signed-off-by: Robert Shaw --- .../layers/fused_moe/modular_kernel.py | 355 +++++++++++------- 1 file changed, 227 insertions(+), 128 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 2a07c20ace4b..0dce874059d2 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -168,14 +168,13 @@ def apply( ReceiverType = Callable[[], PrepareResultType] -# TODO: pass FusedMoEParallelConfig in as ctor parameter? -class FusedMoEPrepareAndFinalize(ABC): +class FusedMoEPrepareAndFinalizeBase(ABC): """ - An abstract base class for the [Quantize-Prepare] and [Finalize] steps - described above. + An abstract base class for FusedMoEPrepareAndFinalize (Modular) + and FusedMoEPrepareAndFinalizeMonolithic (Monolithic) implementations. """ - def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"): + def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermuteBase"): """ Initialize FusedMoEPrepareAndFinalize settings that depend on FusedMoEPermuteExpertsUnpermute experts object. @@ -184,6 +183,56 @@ def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"): """ return + @property + @abstractmethod + def activation_format(self) -> FusedMoEActivationFormat: + """ + A property indicating the output format of the activations for the + 'prepare' method. + """ + raise NotImplementedError + + @abstractmethod + def topk_indices_dtype(self) -> torch.dtype | None: + """ + The PrepareFinalize All2All implementations generally constrain the + dtype of the topk_ids they support. This function returns the + required topk indices dtype so it can be respected. + Return None if there are no such restrictions. + """ + raise NotImplementedError + + @abstractmethod + def max_num_tokens_per_rank(self) -> int | None: + """ + Some PrepareFinalize All2All implementations are batched. Meaning, + they can process only as set of tokens at a time. This + function returns the batch size i.e the maximum number of tokens + the implementation can process at a time. + Return None if there are no such restrictions. + """ + raise NotImplementedError + + @abstractmethod + def num_dispatchers(self) -> int: + raise NotImplementedError + + @abstractmethod + def output_is_reduced(self) -> bool: + """ + Indicates whether or not the output of finalize is reduced across all + ranks. + """ + raise NotImplementedError + + +# TODO: pass FusedMoEParallelConfig in as ctor parameter? +class FusedMoEPrepareAndFinalize(FusedMoEPrepareAndFinalizeBase): + """ + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above. + """ + @abstractmethod def prepare( self, @@ -348,6 +397,15 @@ def finalize_async( """ raise NotImplementedError + +class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalizeBase): + """ + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above but for the monolithic interface (accepts router logits + rather than topk ids and weights). + """ + + @abstractmethod def prepare_monolithic( self, a1: torch.Tensor, @@ -373,11 +431,7 @@ def prepare_monolithic( f"prepare_monolithic not supported for {self.__class__.__name__}" ) - def finalize_monolithic( - self, - fused_expert_output: torch.Tensor, - weight_and_reduce_impl: TopKWeightAndReduce, - ) -> torch.Tensor: + def finalize_monolithic(self, fused_expert_output: torch.Tensor) -> torch.Tensor: """ Optional method for subclasses compatible with monolithic FusedMoEPermuteExpertsUnpermute kernels. @@ -386,61 +440,12 @@ def finalize_monolithic( fused experts output. - fused_expert_output: The unweighted, unreduced output of the fused experts, it will have (M, topk, K) shape. - - weight_and_reduce_impl: An optional TopKWeightAndReduce - implementation. - """ - raise NotImplementedError - - @property - @abstractmethod - def activation_format(self) -> FusedMoEActivationFormat: - """ - A property indicating the output format of the activations for the - 'prepare' method. - """ - raise NotImplementedError - - @abstractmethod - def topk_indices_dtype(self) -> torch.dtype | None: - """ - The PrepareFinalize All2All implementations generally constrain the - dtype of the topk_ids they support. This function returns the - required topk indices dtype so it can be respected. - Return None if there are no such restrictions. - """ - raise NotImplementedError - - @abstractmethod - def max_num_tokens_per_rank(self) -> int | None: - """ - Some PrepareFinalize All2All implementations are batched. Meaning, - they can process only as set of tokens at a time. This - function returns the batch size i.e the maximum number of tokens - the implementation can process at a time. - Return None if there are no such restrictions. - """ - raise NotImplementedError - - @abstractmethod - def num_dispatchers(self) -> int: - raise NotImplementedError - - @abstractmethod - def output_is_reduced(self) -> bool: - """ - Indicates whether or not the output of finalize is reduced across all - ranks. """ raise NotImplementedError # TODO: add supported activations method (return string) -class FusedMoEPermuteExpertsUnpermute(ABC): - """ - An abstract base class for the [Permute-Experts-Unpermute] step described - above. - """ - +class FusedMoEPermuteExpertsUnpermuteBase(ABC): def __init__( self, moe_config: FusedMoEConfig, @@ -492,50 +497,8 @@ def activation_format() -> FusedMoEActivationFormat: """ raise NotImplementedError - def moe_problem_size( - self, - a1: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_ids: torch.Tensor, - ) -> tuple[int, int, int, int, int]: - """ - Extract the MoE problem size from the given tensor arguments: - - a: The hidden states, input to the MoE layer. - - w1: The first set of expert weights. - - w2: The second set of expert weights. - - topk_ids: The topk ids. - - Note: extracting the problem shape from the weight and activation - tensors is not obvious. It needs to be done this way specifically - due to subtle issues with particular kernels, e.g. the int4 kernels - divide the trailing dimension by two, so it's not "correct" to - extract N or K from the trailing dimension of w1 or w2. Similarly, - some kernels transpose the weights, so this needs to be kept in mind. - - Note: This implementation covers most cases. However, if experts - require a specialized implementation, like MarlinExperts, they are free - to override this function. - """ - assert w1.dim() == 3 and w2.dim() == 3 - E, N, _ = w1.size() - K = a1.size(-1) - - if a1.dim() == 2: - # Make sure we are using the correct a1 (pre-permute). - assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" - M = a1.size(0) - else: - assert a1.dim() == 3 - assert a1.size(0) == E, f"{a1.size(0)} == {E}" - M = a1.size(1) # This is max_num_tokens - - assert topk_ids.dim() == 2 - topk = topk_ids.size(1) - - return E, M, N, K, topk + # - # # Various helpers for registering support for various features. # Used by the oracle to select a particular kernel for a deployment. # @@ -717,6 +680,61 @@ def supports_packed_ue8m0_act_scales(self) -> bool: """ return False + def enable_chunking(self): + return ( + envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and self.supports_chunking() + ) + + +class FusedMoEPermuteExpertsUnpermute(FusedMoEPermuteExpertsUnpermuteBase): + """ + An abstract base class for the [Permute-Experts-Unpermute] step described + above. + """ + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + """ + Extract the MoE problem size from the given tensor arguments: + - a: The hidden states, input to the MoE layer. + - w1: The first set of expert weights. + - w2: The second set of expert weights. + - topk_ids: The topk ids. + + Note: extracting the problem shape from the weight and activation + tensors is not obvious. It needs to be done this way specifically + due to subtle issues with particular kernels, e.g. the int4 kernels + divide the trailing dimension by two, so it's not "correct" to + extract N or K from the trailing dimension of w1 or w2. Similarly, + some kernels transpose the weights, so this needs to be kept in mind. + + Note: This implementation covers most cases. However, if experts + require a specialized implementation, like MarlinExperts, they are free + to override this function. + """ + assert w1.dim() == 3 and w2.dim() == 3 + E, N, _ = w1.size() + K = a1.size(-1) + + if a1.dim() == 2: + # Make sure we are using the correct a1 (pre-permute). + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + M = a1.size(0) + else: + assert a1.dim() == 3 + assert a1.size(0) == E, f"{a1.size(0)} == {E}" + M = a1.size(1) # This is max_num_tokens + + assert topk_ids.dim() == 2 + topk = topk_ids.size(1) + + return E, M, N, K, topk + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: """ Workspace type: The dtype to use for the workspace tensors. @@ -790,11 +808,7 @@ def activation( ) -> None: apply_moe_activation(activation, output, input) - def enable_chunking(self): - return ( - envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and self.supports_chunking() - ) - + @abstractmethod def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce: raise NotImplementedError @@ -854,6 +868,14 @@ def apply( """ raise NotImplementedError + +class FusedMoEPermuteExpertsUnpermuteMonolithic(FusedMoEPermuteExpertsUnpermuteBase): + """ + An abstract base class for the [Permute-Experts-Unpermute] step described + above, but with the monolithic interface (accepts router logits + rather than topk ids and weights). + """ + def apply_monolithic( self, hidden_states: torch.Tensor, @@ -890,24 +912,11 @@ def _slice_scales( return None -@final -class FusedMoEModularKernel(torch.nn.Module): - """ - This class combines a FusedMoEPrepareAndFinalize instance and - a FusedMoEPermuteExpertsUnpermute to provide an interface that - is compatible with the `fused_experts` function in fused_moe.py. - - It takes care of managing any required scratch space. - - Note: Instances of this class should only be used for a single model - layer due to any layer specific state that may be used by the component - objects. - """ - +class FusedMoEModularKernelBase(torch.nn.Module): def __init__( self, - prepare_finalize: FusedMoEPrepareAndFinalize, - fused_experts: FusedMoEPermuteExpertsUnpermute, + prepare_finalize: FusedMoEPrepareAndFinalizeBase, + fused_experts: FusedMoEPermuteExpertsUnpermuteBase, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, ): @@ -927,6 +936,22 @@ def __init__( and moe_parallel_config.use_ep ) + # Confirm P/F and Experts kernels are consistent with eachother. + if not ( + ( + isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) + and isinstance(fused_experts, FusedMoEPermuteExpertsUnpermuteMonolithic) + ) + or ( + isinstance(prepare_finalize, FusedMoEPrepareAndFinalize) + and isinstance(fused_experts, FusedMoEPermuteExpertsUnpermute) + ) + ): + raise ValueError( + "prepare_finalize and fused_experts must both be either " + "monolithic or non-monolithic" + ) + self._post_init_setup() assert ( prepare_finalize.activation_format == fused_experts.activation_format() @@ -957,6 +982,47 @@ def output_is_reduced(self) -> bool: """ return self.prepare_finalize.output_is_reduced() + +@final +class FusedMoEModularKernel(FusedMoEModularKernelBase): + """ + This class combines a FusedMoEPrepareAndFinalize instance and + a FusedMoEPermuteExpertsUnpermute to provide an interface that + is compatible with the `fused_experts` function in fused_moe.py. + + It takes care of managing any required scratch space. + + Note: Instances of this class should only be used for a single model + layer due to any layer specific state that may be used by the component + objects. + """ + + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalize, + fused_experts: FusedMoEPermuteExpertsUnpermute, + shared_experts: torch.nn.Module | None = None, + moe_parallel_config: FusedMoEParallelConfig | None = None, + ): + if not isinstance(prepare_finalize, FusedMoEPrepareAndFinalize): + raise TypeError( + "prepare_finalize must be an instance of FusedMoEPrepareAndFinalize" + ) + if not isinstance(fused_experts, FusedMoEPermuteExpertsUnpermute): + raise TypeError( + "fused_experts must be an instance of FusedMoEPermuteExpertsUnpermute" + ) + + super().__init__( + prepare_finalize, + fused_experts, + shared_experts, + moe_parallel_config, + ) + + self.prepare_finalize: FusedMoEPrepareAndFinalize = prepare_finalize + self.fused_experts: FusedMoEPermuteExpertsUnpermute = fused_experts + def _chunk_info(self, M: int) -> tuple[int, int]: """ Compute number of chunks and chunk size for given M. @@ -999,6 +1065,7 @@ def _allocate_buffers( See `workspace_shapes` for a description of the remainder of arguments. Returns a tuple of (workspace13, workspace2, output) tensors. """ + assert isinstance(self.fused_experts, FusedMoEPermuteExpertsUnpermute) assert M_full > 0 and M_chunk > 0 num_chunks, _ = self._chunk_info(M_full) @@ -1154,6 +1221,8 @@ def _prepare( The _prepare method is a wrapper around self.prepare_finalize.prepare that handles DBO and async. """ + assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalize) + if not self.prepare_finalize.supports_async(): # We shouldn't be running an a2a kernel that doesn't # support async prepare/finalize @@ -1239,6 +1308,8 @@ def _fused_experts( apply_router_weight_on_input: bool, expert_tokens_meta: ExpertTokensMetadata | None, ) -> torch.Tensor: + assert isinstance(self.fused_experts, FusedMoEPermuteExpertsUnpermuteBase) + _, M_full, N, K, top_k = self.fused_experts.moe_problem_size( a1q, w1, w2, topk_ids ) @@ -1468,6 +1539,34 @@ def forward( apply_router_weight_on_input, ) + +@final +class FusedMoEModularKernelMonolithic(FusedMoEModularKernelBase): + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic, + fused_experts: FusedMoEPermuteExpertsUnpermuteMonolithic, + shared_experts: torch.nn.Module | None = None, + moe_parallel_config: FusedMoEParallelConfig | None = None, + ): + if not isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic): + raise TypeError( + "prepare_finalize must be an instance of " + "FusedMoEPrepareAndFinalizeMonolithic" + ) + if not isinstance(fused_experts, FusedMoEPermuteExpertsUnpermuteMonolithic): + raise TypeError( + "fused_experts must be an instance of " + "FusedMoEPermuteExpertsUnpermuteMonolithic" + ) + + super().__init__( + prepare_finalize, + fused_experts, + shared_experts, + moe_parallel_config, + ) + def forward_monolithic( self, hidden_states: torch.Tensor, @@ -1489,6 +1588,9 @@ def forward_monolithic( to the topk_ids and topk_weights. This is used for kernels that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ + assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) + assert isinstance(self.fused_experts, FusedMoEPermuteExpertsUnpermuteMonolithic) + # TODO(rob): add inplace support. a1q, a1q_scale, router_logits = self.prepare_finalize.prepare_monolithic( hidden_states, @@ -1514,9 +1616,6 @@ def forward_monolithic( topk_group=topk_group, ) - output = self.prepare_finalize.finalize_monolithic( - fused_expert_output=fused_out, - weight_and_reduce_impl=self.fused_experts.finalize_weight_and_reduce_impl(), - ) + output = self.prepare_finalize.finalize_monolithic(fused_out) return output From 39aa9c413a2834be8f05e2bd89be405cabf94f92 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Thu, 29 Jan 2026 19:13:23 -0500 Subject: [PATCH 090/207] stash Signed-off-by: Robert Shaw --- .../moe/modular_kernel_tools/common.py | 2 +- .../layers/fused_moe/all2all_utils.py | 1 + .../layers/fused_moe/cutlass_moe.py | 2 +- .../layers/fused_moe/fused_moe.py | 2 +- .../layers/fused_moe/fused_moe_method_base.py | 2 +- .../layers/fused_moe/modular_kernel.py | 34 +++++++++++++++++++ .../layers/fused_moe/oracle/fp8.py | 4 +-- .../layers/fused_moe/oracle/nvfp4.py | 6 ++-- .../layers/fused_moe/oracle/unquantized.py | 6 ++-- .../fused_moe/unquantized_fused_moe_method.py | 2 +- 10 files changed, 48 insertions(+), 13 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 4ee18e3428ef..470d200fd7a6 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -620,7 +620,7 @@ def next_power_of_2(x): config.N, ) - modular_kernel = mk.FusedMoEModularKernel( + modular_kernel = mk.FusedMoEModularKernel.make_mk( prepare_finalize=prepare_finalize, fused_experts=fused_experts, ) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index bf8ec2dc6f20..97ee5a59fc44 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -81,6 +81,7 @@ def maybe_make_prepare_finalize( quant_config: FusedMoEQuantConfig | None, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, allow_new_interface: bool = False, + use_monolithic: bool = False, ) -> FusedMoEPrepareAndFinalize | None: # NOTE(rob): we are migrating each quant_method to hold the MK # in all cases. The allow_new_interface=False flag allow us to fall diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 86edbe303ef2..f190c70267cc 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -1144,7 +1144,7 @@ def cutlass_moe_w4a8_fp8( num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0) - fn = mk.FusedMoEModularKernel( + fn = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsW4A8Fp8( out_dtype=a.dtype, diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 0206e19def4f..f4121eb3d5fa 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -2306,7 +2306,7 @@ def modular_triton_fused_moe( quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None, ) -> mk.FusedMoEModularKernel: - return mk.FusedMoEModularKernel( + return mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), TritonExperts(moe_config, quant_config), shared_experts, 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 9f99a8a24976..5eee7723f20e 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 @@ -27,7 +27,7 @@ def __init__(self, moe: FusedMoEConfig): super().__init__() self.moe: FusedMoEConfig = moe self.moe_quant_config: FusedMoEQuantConfig | None = None - self.moe_mk: mk.FusedMoEModularKernel | None = None + self.moe_mk: mk.FusedMoEModularKernelBase | None = None @property def supports_internal_mk(self) -> bool: diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 0dce874059d2..ea6b56a6ddbc 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -962,6 +962,40 @@ def __init__( f"{fused_experts.activation_format()}" ) + @staticmethod + def make_mk( + prepare_finalize: FusedMoEPrepareAndFinalizeBase, + fused_experts: FusedMoEPermuteExpertsUnpermuteBase, + shared_experts: torch.nn.Module | None = None, + moe_parallel_config: FusedMoEParallelConfig | None = None, + ) -> "FusedMoEModularKernelBase": + """ + Factory method to create a FusedMoEModularKernelBase instance. + """ + if isinstance( + prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic + ) and isinstance(fused_experts, FusedMoEPermuteExpertsUnpermuteMonolithic): + return FusedMoEModularKernelMonolithic( + prepare_finalize, + fused_experts, + shared_experts, + moe_parallel_config, + ) + elif isinstance(prepare_finalize, FusedMoEPrepareAndFinalize) and isinstance( + fused_experts, FusedMoEPermuteExpertsUnpermute + ): + return FusedMoEModularKernel( + prepare_finalize, + fused_experts, + shared_experts, + moe_parallel_config, + ) + else: + raise ValueError( + "prepare_finalize and fused_experts must both be either " + "monolithic or non-monolithic" + ) + def _post_init_setup(self): """ Resolve any leftover setup dependencies between self.prepare_finalize diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 6b44c4df8096..ea2f6546f40b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -433,7 +433,7 @@ def make_fp8_moe_kernel( fp8_backend: Fp8MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, shared_experts: torch.nn.Module | None = None, -) -> tuple[mk.FusedMoEModularKernel, bool]: +) -> tuple[mk.FusedMoEModularKernelBase, bool]: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( moe=moe_config, @@ -464,7 +464,7 @@ def make_fp8_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernelBase.make_mk( prepare_finalize, experts, shared_experts=( diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 27e7e788c338..6fc98b5aaa1d 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -364,10 +364,10 @@ def make_nvfp4_moe_quant_config( def make_nvfp4_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, - experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute], + experts_cls: type[mk.FusedMoEPermuteExpertsUnpermuteBase], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, shared_experts: torch.nn.Module | None = None, -) -> mk.FusedMoEModularKernel: +) -> mk.FusedMoEModularKernelBase: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( moe=moe_config, @@ -398,7 +398,7 @@ def make_nvfp4_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( prepare_finalize, experts, shared_experts=( diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 14c3f84e64ed..dc50b7c96db1 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -137,7 +137,7 @@ def make_unquantized_moe_kernel( FlashInferExperts, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernelBase.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, @@ -150,7 +150,7 @@ def make_unquantized_moe_kernel( AiterExperts, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernelBase.make_mk( MoEPrepareAndFinalizeNoEP(), AiterExperts( moe_config=moe_config, @@ -160,7 +160,7 @@ def make_unquantized_moe_kernel( elif backend == UnquantizedMoeBackend.TRITON: from vllm.model_executor.layers.fused_moe import TritonExperts - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernelBase.make_mk( MoEPrepareAndFinalizeNoEP(), TritonExperts( moe_config=moe_config, diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 4b85cc5c2d06..2baf420daab2 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -63,7 +63,7 @@ def __init__(self, moe: FusedMoEConfig): self.rocm_aiter_moe_enabled = ( rocm_aiter_ops.is_fused_moe_enabled() and moe.is_act_and_mul ) - self.kernel: mk.FusedMoEModularKernel | None = None + self.kernel: mk.FusedMoEModularKernelBase | None = None self._is_monolithic = current_platform.is_cpu() or current_platform.is_xpu() @property From 8970309548cb6333a54547703cff63027ac8c683 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Thu, 29 Jan 2026 20:18:51 -0500 Subject: [PATCH 091/207] update naming Signed-off-by: Robert Shaw --- docs/design/fused_moe_modular_kernel.md | 56 ++++++------- docs/design/moe_kernel_features.md | 4 +- .../moe/modular_kernel_tools/cli_args.py | 2 +- .../moe/modular_kernel_tools/common.py | 2 +- .../moe/modular_kernel_tools/mk_objects.py | 8 +- .../moe/test_modular_kernel_combinations.py | 4 +- .../layers/fused_moe/__init__.py | 4 +- .../layers/fused_moe/batched_deep_gemm_moe.py | 2 +- .../layers/fused_moe/cutlass_moe.py | 6 +- .../layers/fused_moe/deep_gemm_moe.py | 2 +- .../fused_moe/deepep_ll_prepare_finalize.py | 2 +- .../layers/fused_moe/fallback.py | 12 +-- .../fused_moe/flashinfer_cutedsl_moe.py | 2 +- .../fused_moe/flashinfer_cutlass_moe.py | 2 +- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 2 +- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 2 +- .../layers/fused_moe/fused_batched_moe.py | 4 +- .../layers/fused_moe/fused_marlin_moe.py | 2 +- .../layers/fused_moe/fused_moe.py | 4 +- .../layers/fused_moe/fused_moe_method_base.py | 6 +- .../fused_moe/gpt_oss_triton_kernels_moe.py | 2 +- .../layers/fused_moe/modular_kernel.py | 79 +++++++++---------- .../layers/fused_moe/oracle/fp8.py | 14 ++-- .../layers/fused_moe/oracle/nvfp4.py | 10 +-- .../layers/fused_moe/oracle/unquantized.py | 6 +- .../layers/fused_moe/rocm_aiter_fused_moe.py | 2 +- .../fused_moe/topk_weight_and_reduce.py | 2 +- .../layers/fused_moe/triton_cutlass_moe.py | 6 +- .../layers/fused_moe/triton_deep_gemm_moe.py | 6 +- .../layers/fused_moe/trtllm_moe.py | 2 +- .../fused_moe/unquantized_fused_moe_method.py | 6 +- .../compressed_tensors_moe.py | 14 ++-- .../model_executor/layers/quantization/fp8.py | 4 +- .../layers/quantization/modelopt.py | 4 +- .../layers/quantization/mxfp4.py | 2 +- 35 files changed, 143 insertions(+), 144 deletions(-) diff --git a/docs/design/fused_moe_modular_kernel.md b/docs/design/fused_moe_modular_kernel.md index 975df8ba29dc..32012a90ca2d 100644 --- a/docs/design/fused_moe_modular_kernel.md +++ b/docs/design/fused_moe_modular_kernel.md @@ -38,19 +38,19 @@ FusedMoEModularKernel splits the FusedMoE operation into 3 parts, 1. TopKWeightAndReduce 2. FusedMoEPrepareAndFinalize -3. FusedMoEPermuteExpertsUnpermute +3. FusedMoEModularExperts ### TopKWeightAndReduce -The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEPermuteExpertsUnpermute` is responsible for the Unpermute and `FusedMoEPrepareAndFinalize` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEPermuteExpertsUnpermute`. But some implementations choose to do it `FusedMoEPrepareAndFinalize`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class. +The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEModularExperts` is responsible for the Unpermute and `FusedMoEPrepareAndFinalize` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEModularExperts`. But some implementations choose to do it `FusedMoEPrepareAndFinalize`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class. Please find the implementations of TopKWeightAndReduce [here](../../vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py). `FusedMoEPrepareAndFinalize::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method. -The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEPermuteExpertsUnpermute` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. +The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEModularExperts` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. -* `FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEPermuteExpertsUnpermute` implementation does the weight application and reduction itself. -* `FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEPermuteExpertsUnpermute` implementation needs the `FusedMoEPrepareAndFinalize::finalize()` to do the weight application and reduction. +* `FusedMoEModularExperts::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEModularExperts` implementation does the weight application and reduction itself. +* `FusedMoEModularExperts::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEModularExperts` implementation needs the `FusedMoEPrepareAndFinalize::finalize()` to do the weight application and reduction. ### FusedMoEPrepareAndFinalize @@ -59,9 +59,9 @@ The `prepare` function is responsible for input activation Quantization and All2 ![FusedMoEPrepareAndFinalize Blocks](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png) -### FusedMoEPermuteExpertsUnpermute +### FusedMoEModularExperts -The `FusedMoEPermuteExpertsUnpermute` class is where the crux of the MoE operations happen. The `FusedMoEPermuteExpertsUnpermute` abstract class exposes a few important functions, +The `FusedMoEModularExperts` class is where the crux of the MoE operations happen. The `FusedMoEModularExperts` abstract class exposes a few important functions, * apply() * workspace_shapes() @@ -81,25 +81,25 @@ The `apply` method is where the implementations perform #### workspace_shapes() -The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEPermuteExpertsUnpermute::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. +The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEModularExperts::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. #### finalize_weight_and_reduce_impl() -It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEPermuteExpertsUnpermute::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section. -`FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalize::finalize()` to use. +It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEModularExperts::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section. +`FusedMoEModularExperts::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalize::finalize()` to use. -![FusedMoEPermuteExpertsUnpermute Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png) +![FusedMoEModularExperts Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png) ### FusedMoEModularKernel -`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` objects. +`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalize` and `FusedMoEModularExperts` objects. `FusedMoEModularKernel` pseudocode/sketch, ```py class FusedMoEModularKernel: def __init__(self, prepare_finalize: FusedMoEPrepareAndFinalize, - fused_experts: FusedMoEPermuteExpertsUnpermute): + fused_experts: FusedMoEModularExperts): self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts @@ -162,20 +162,20 @@ This section describes the significance of the various functions exposed by the We suggest picking an already existing `FusedMoEPrepareAndFinalize` implementation that matches your All2All implementation closely and using it as a reference. -### How To Add a FusedMoEPermuteExpertsUnpermute Type +### How To Add a FusedMoEModularExperts Type -FusedMoEPermuteExpertsUnpermute performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows, +FusedMoEModularExperts performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows, -`FusedMoEPermuteExpertsUnpermute::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format. +`FusedMoEModularExperts::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format. -`FusedMoEPermuteExpertsUnpermute::supports_chunking()`: Return True if the implementation supports chunking. Typically +`FusedMoEModularExperts::supports_chunking()`: Return True if the implementation supports chunking. Typically implementations that input `FusedMoEActivationFormat.Standard` support chunking and `FusedMoEActivationFormat.BatchedExperts` do not. -`FusedMoEPermuteExpertsUnpermute::supports_expert_map()`: Return True if the implementation supports expert map. +`FusedMoEModularExperts::supports_expert_map()`: Return True if the implementation supports expert map. -`FusedMoEPermuteExpertsUnpermute::workspace_shapes()` / -`FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` / -`FusedMoEPermuteExpertsUnpermute::apply`: Refer to `FusedMoEPermuteExpertsUnpermute` section above. +`FusedMoEModularExperts::workspace_shapes()` / +`FusedMoEModularExperts::finalize_weight_and_reduce_impl` / +`FusedMoEModularExperts::apply`: Refer to `FusedMoEModularExperts` section above. ### FusedMoEModularKernel Initialization @@ -194,7 +194,7 @@ Please refer to the implementations in, #### select_gemm_impl -The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEPermuteExpertsUnpermute` object. +The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEModularExperts` object. Please refer to the implementations in, * `UnquantizedFusedMoEMethod` @@ -206,7 +206,7 @@ derived classes. #### init_prepare_finalize -Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalize` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEPermuteExpertsUnpermute` object and builds the `FusedMoEModularKernel` object +Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalize` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEModularExperts` object and builds the `FusedMoEModularKernel` object Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188). **Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used. @@ -217,7 +217,7 @@ We have `FusedMoEModularKernel` unit tests at [test_modular_kernel_combinations. The unit test iterates through all combinations of `FusedMoEPrepareAndFinalize` and `FusedMoEPremuteExpertsUnpermute` types and if they are compatible, runs some correctness tests. -If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEPermuteExpertsUnpermute` implementations, +If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEModularExperts` implementations, 1. Add the implementation type to `MK_ALL_PREPARE_FINALIZE_TYPES` and `MK_FUSED_EXPERT_TYPES` in [mk_objects.py](../../tests/kernels/moe/modular_kernel_tools/mk_objects.py) respectively. 2. Update `Config::is_batched_prepare_finalize()`, `Config::is_batched_fused_experts()`, `Config::is_standard_fused_experts()`, @@ -226,24 +226,24 @@ If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEPermuteExpertsUnp Doing this will add the new implementation to the test suite. -### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` Compatibility +### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEModularExperts` Compatibility The unit test file [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py) can also be executed as a standalone script. Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` -As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` compatibility. When invoked +As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEModularExperts` compatibility. When invoked with incompatible types, the script will error. ### How To Profile Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py) The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible -`FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` types. +`FusedMoEPrepareAndFinalize` and `FusedMoEModularExperts` types. Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` ## FusedMoEPrepareAndFinalize Implementations See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-modular-all2all-backends) for a list of all the available modular prepare and finalize subclasses. -## FusedMoEPermuteExpertsUnpermute +## FusedMoEModularExperts See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-experts-kernels) for a list of all the available modular experts. diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 75ebee6ecd25..731f635c2555 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -76,7 +76,7 @@ Each experts kernel supports one or more activation functions, e.g. silu or gelu As with the backends, some experts support applying topk weights on the input activations. The entries in the column in this table only apply to the non-modular experts. -Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEPermuteExpertsUnpermute`. +Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEModularExperts`. To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats. @@ -107,7 +107,7 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels The following table shows "families" of modular kernels that are intended to work together. There are some combinations which may work but have not yet been tested, e.g. flashinfer with other fp8 experts. Note that the "naive" backend will work with any non-modular experts. -| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses | +| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEModularExperts` subclasses | |---------|-----------------------------------------|----------------------------------------------| | deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,
`TritonExperts`,
`TritonOrDeepGemmExperts`,
`CutlassExpertsFp8`,
`MarlinExperts` | | deepep_low_latency,
pplx | `DeepEPLLPrepareAndFinalize`,
`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,
`BatchedTritonExperts`,
`CutlassBatchedExpertsFp8`,
`BatchedMarlinExperts` | diff --git a/tests/kernels/moe/modular_kernel_tools/cli_args.py b/tests/kernels/moe/modular_kernel_tools/cli_args.py index 34c6ca1f999c..ddc5a95cd08e 100644 --- a/tests/kernels/moe/modular_kernel_tools/cli_args.py +++ b/tests/kernels/moe/modular_kernel_tools/cli_args.py @@ -23,7 +23,7 @@ def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalize: return pf raise ValueError(f"Cannot find a PrepareFinalize type that matches {s}") - def to_experts_class_type(s: str) -> mk.FusedMoEPermuteExpertsUnpermute: + def to_experts_class_type(s: str) -> mk.FusedMoEModularExperts: for fe in MK_FUSED_EXPERT_TYPES: if fe.__name__ == s: return fe diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 470d200fd7a6..432e62ccfb9b 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -66,7 +66,7 @@ class Config: quant_config: TestMoEQuantConfig | None prepare_finalize_type: mk.FusedMoEPrepareAndFinalize - fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute + fused_experts_type: mk.FusedMoEModularExperts fused_moe_chunk_size: int | None world_size: int diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index d215d2ab69b1..d2f8e5b5cf27 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -73,11 +73,11 @@ class ExpertInfo: PREPARE_FINALIZE_INFO: dict[mk.FusedMoEPrepareAndFinalize, PrepareFinalizeInfo] = {} -EXPERT_INFO: dict[mk.FusedMoEPermuteExpertsUnpermute, ExpertInfo] = {} +EXPERT_INFO: dict[mk.FusedMoEModularExperts, ExpertInfo] = {} MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] -MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEPermuteExpertsUnpermute] = [] +MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEModularExperts] = [] standard_format = mk.FusedMoEActivationFormat.Standard batched_format = mk.FusedMoEActivationFormat.BatchedExperts @@ -444,12 +444,12 @@ def make_cutlass_strides( def make_fused_experts( - fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute, + fused_experts_type: mk.FusedMoEModularExperts, moe: FusedMoEConfig, quant_config: FusedMoEQuantConfig, num_dispatchers: int, N: int, -) -> mk.FusedMoEPermuteExpertsUnpermute: +) -> mk.FusedMoEModularExperts: if ( fused_experts_type.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index ec31e66140a1..dad4948ce659 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -259,7 +259,7 @@ def test_modular_kernel_combinations_multigpu( dtype: torch.dtype, quant_config: TestMoEQuantConfig | None, prepare_finalize_type: mk.FusedMoEPrepareAndFinalize, - fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute, + fused_experts_type: mk.FusedMoEModularExperts, chunk_size: int | None, world_size: int, pytestconfig, @@ -301,7 +301,7 @@ def test_modular_kernel_combinations_singlegpu( dtype: torch.dtype, quant_config: TestMoEQuantConfig | None, prepare_finalize_type: mk.FusedMoEPrepareAndFinalize, - fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute, + fused_experts_type: mk.FusedMoEModularExperts, chunk_size: int | None, world_size: int, pytestconfig, diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 03be6f8b6a43..32b18a767e6f 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -17,7 +17,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEActivationFormat, - FusedMoEPermuteExpertsUnpermute, + FusedMoEModularExperts, FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( @@ -56,7 +56,7 @@ def get_config() -> dict[str, Any] | None: "FusedMoEMethodBase", "UnquantizedFusedMoEMethod", "FusedMoeWeightScaleSupported", - "FusedMoEPermuteExpertsUnpermute", + "FusedMoEModularExperts", "FusedMoEActivationFormat", "FusedMoEPrepareAndFinalize", "RoutingMethodType", diff --git a/vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py index ac37cff9329a..1bd5d986c62c 100644 --- a/vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py @@ -260,7 +260,7 @@ def persistent_masked_m_silu_mul_quant( return y_q, y_s -class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute): +class BatchedDeepGemmExperts(mk.FusedMoEModularExperts): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index f190c70267cc..9d216425d20a 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -259,7 +259,7 @@ def run_cutlass_moe_fp8( ) -class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute): +class CutlassExpertsFp8Base(mk.FusedMoEModularExperts): def __init__( self, moe_config: FusedMoEConfig, @@ -650,7 +650,7 @@ def run_cutlass_moe_fp4( return -class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute): +class CutlassExpertsFp4(mk.FusedMoEModularExperts): @property def expects_unquantized_inputs(self) -> bool: return True @@ -902,7 +902,7 @@ def run_cutlass_moe_w4a8_fp8( ) -class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute): +class CutlassExpertsW4A8Fp8(mk.FusedMoEModularExperts): def __init__( self, out_dtype: torch.dtype | None, diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py index fafcf6de6140..d98c1216a048 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py @@ -112,7 +112,7 @@ def _valid_deep_gemm( return True -class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute): +class DeepGemmExperts(mk.FusedMoEModularExperts): def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig): super().__init__(moe_config=moe_config, quant_config=quant_config) assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout() diff --git a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py index f5a3da438781..dd6276afeda3 100644 --- a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py @@ -119,7 +119,7 @@ def _maybe_cast(tensor: torch.Tensor | None) -> torch.Tensor | None: # time. This setting is handled by post_init_setup. self.use_ue8m0_dispatch = False - def post_init_setup(self, fused_experts: mk.FusedMoEPermuteExpertsUnpermute): + def post_init_setup(self, fused_experts: mk.FusedMoEExperts): if not fused_experts.supports_packed_ue8m0_act_scales(): # Early exit. return diff --git a/vllm/model_executor/layers/fused_moe/fallback.py b/vllm/model_executor/layers/fused_moe/fallback.py index 07e5b80059f0..18391e2e83b9 100644 --- a/vllm/model_executor/layers/fused_moe/fallback.py +++ b/vllm/model_executor/layers/fused_moe/fallback.py @@ -10,13 +10,13 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey -class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC): +class FallbackExperts(mk.FusedMoEModularExperts, ABC): """Base class for runtime dispatching of expert implementations.""" def __init__( self, - experts: mk.FusedMoEPermuteExpertsUnpermute, - fallback_experts: mk.FusedMoEPermuteExpertsUnpermute, + experts: mk.FusedMoEModularExperts, + fallback_experts: mk.FusedMoEModularExperts, ): super().__init__( moe_config=experts.moe_config, quant_config=experts.quant_config @@ -26,8 +26,8 @@ def __init__( @staticmethod def get_clses() -> tuple[ - type[mk.FusedMoEPermuteExpertsUnpermute], - type[mk.FusedMoEPermuteExpertsUnpermute], + type[mk.FusedMoEModularExperts], + type[mk.FusedMoEModularExperts], ]: """ Get the cls for the experts and fallback experts. @@ -148,7 +148,7 @@ def _select_experts_impl( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: raise NotImplementedError def apply( diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py index 036ee2a2ec41..68582f2b6638 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py @@ -29,7 +29,7 @@ logger = init_logger(__name__) -class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute): +class FlashInferCuteDSLExperts(mk.FusedMoEModularExperts): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index faa654ea3a2f..af3115a5ec23 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -56,7 +56,7 @@ def is_valid_flashinfer_cutlass_fused_moe( return True -class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute): +class FlashInferExperts(mk.FusedMoEModularExperts): def __init__( self, moe_config: mk.FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index ad9ab20b24a5..173d73d64405 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -23,7 +23,7 @@ from vllm.v1.engine.utils import current_platform -class FlashInferTrtLlmFp8Experts(mk.FusedMoEPermuteExpertsUnpermute): +class FlashInferTrtLlmFp8Experts(mk.FusedMoEModularExperts): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 9724489e9884..a11b711bcfdc 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -22,7 +22,7 @@ from vllm.platforms import current_platform -class FlashInferTrtLlmNvFp4Experts(mk.FusedMoEPermuteExpertsUnpermute): +class FlashInferTrtLlmNvFp4Experts(mk.FusedMoEModularExperts): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py index c681e083acb8..f8ca531f507b 100644 --- a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py @@ -644,7 +644,7 @@ def finalize( ) -class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute): +class NaiveBatchedExperts(mk.FusedMoEModularExperts): """ A reference MoE expert class that operates on expert batched format, i.e. E x max_num_tokens x K. This is the format that the pplx @@ -876,7 +876,7 @@ def batched_moe_kernel_quantize_input( return A_q, A_q_scale -class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute): +class BatchedTritonExperts(mk.FusedMoEModularExperts): """ A Triton based MoE expert class that operates on expert batched format, i.e. E x max_num_tokens x K. This is the format that the pplx diff --git a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py index 2e5167bdfd53..dd617be36f88 100644 --- a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py @@ -524,7 +524,7 @@ def batched_fused_marlin_moe( return output -class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute): +class MarlinExpertsBase(mk.FusedMoEModularExperts): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index f4121eb3d5fa..24cff704ae54 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -1703,7 +1703,7 @@ def fused_experts_impl( intermediate_cache3 = cache13[: M * top_k_num * K].view(M, top_k_num, K) # This needs separate memory since it's used concurrently with cache1 - activation_out_dim = mk.FusedMoEPermuteExpertsUnpermute.adjust_N_for_activation( + activation_out_dim = mk.FusedMoEModularExperts.adjust_N_for_activation( N, activation ) intermediate_cache2 = torch.empty( @@ -1896,7 +1896,7 @@ def fused_experts_impl( return out_hidden_states -class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute): +class TritonExperts(mk.FusedMoEModularExperts): def __init__( self, moe_config: FusedMoEConfig, 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 5eee7723f20e..ec52da8d3b51 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 @@ -12,7 +12,7 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEPermuteExpertsUnpermute, + FusedMoEModularExperts, FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.quantization.base_config import ( @@ -27,7 +27,7 @@ def __init__(self, moe: FusedMoEConfig): super().__init__() self.moe: FusedMoEConfig = moe self.moe_quant_config: FusedMoEQuantConfig | None = None - self.moe_mk: mk.FusedMoEModularKernelBase | None = None + self.moe_mk: mk.FusedMoEKernel | None = None @property def supports_internal_mk(self) -> bool: @@ -77,7 +77,7 @@ def select_gemm_impl( self, prepare_finalize: FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> FusedMoEPermuteExpertsUnpermute: + ) -> FusedMoEModularExperts: # based on the all2all implementation, select the appropriate # gemm implementation raise NotImplementedError( diff --git a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py index b209820cdfa9..3101290b7f61 100644 --- a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py @@ -244,7 +244,7 @@ def make_routing_data( return routing_data, gather_indx, scatter_indx -class BaseOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute): +class BaseOAITritonExperts(mk.FusedMoEModularExperts): @staticmethod def _supports_current_device() -> bool: 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 ea6b56a6ddbc..26ad081c201e 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -56,18 +56,18 @@ # * FusedMoEPrepareAndFinalize - an abstract base class for preparation of MoE # inputs (e.g. quantization, distribution) and finalization of Moe outputs. # The prepare method must take care of any needed quantization and the -# finalize method, informed by the FusedMoEPermuteExpertsUnpermute method, +# finalize method, informed by the FusedMoEModularExperts method, # may apply weights and/or do the final reduction of the output. -# * FusedMoEPermuteExpertsUnpermute - an abstract base class for the main fused +# * FusedMoEModularExperts - an abstract base class for the main fused # MoE operation, i.e matmul + act_mul + optionally quant + matmul. -# Some FusedMoEPermuteExpertsUnpermute implementations may choose to do +# Some FusedMoEModularExperts implementations may choose to do # the weight application and/or reduction. The class communicates this # to [Finalize] via a TopKWeightAndReduce object. # * FusedMoEModularKernel - an interface class that combines a -# FusedMoEPrepareAndFinalize and a FusedMoEPermuteExpertsUnpermute to +# FusedMoEPrepareAndFinalize and a FusedMoEModularExperts to # provide the standard fused MoE kernel interface. # * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen -# by the FusedMoEPermuteExpertsUnpermute implementation that is passed +# by the FusedMoEModularExperts implementation that is passed # on to [Finalize]. # # [Quantize-Prepare] and [Finalize] functionality are bundled into a single @@ -174,10 +174,10 @@ class FusedMoEPrepareAndFinalizeBase(ABC): and FusedMoEPrepareAndFinalizeMonolithic (Monolithic) implementations. """ - def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermuteBase"): + def post_init_setup(self, fused_experts: "FusedMoEExperts"): """ Initialize FusedMoEPrepareAndFinalize settings that depend on - FusedMoEPermuteExpertsUnpermute experts object. + FusedMoEModularExperts experts object. The FusedMoEPrepareAndFinalize implementations that have such dependencies may choose to override this function. """ @@ -257,7 +257,7 @@ def prepare( activations, before quantization + dispatching. - quant_config: Quantization info provided by the fused experts. - defer_input_quant: Runtime parameter indicating whether or not to - defer input quantization to the FusedMoEPermuteExpertsUnpermute + defer input quantization to the FusedMoEModularExperts in cases where the compute kernel expects unquantized inputs Returns a tuple of: @@ -304,7 +304,7 @@ def prepare_async( - apply_router_weight_on_input: When True, apply the weights to the activations, before quantization + dispatching. - defer_input_quant: Runtime parameter indicating whether or not to - defer input quantization to the FusedMoEPermuteExpertsUnpermute + defer input quantization to the FusedMoEModularExperts in cases where the compute kernel expects unquantized inputs Returns a callback or a hook callback pair that when invoked waits for @@ -415,13 +415,13 @@ def prepare_monolithic( ) -> PrepareMonolithicResultType: """ Optional method for subclasses compatible with monolithic - FusedMoEPermuteExpertsUnpermute kernels. + FusedMoEModularExperts kernels. Perform any quantization (and/or) dispatching needed for this kernel. - a1: The (unquantized) input to the MoE layer. - quant_config: Quantization info provided by the fused experts. - defer_input_quant: Runtime parameter indicating whether or not to - defer input quantization to the FusedMoEPermuteExpertsUnpermute + defer input quantization to the FusedMoEModularExperts Returns a tuple of: - quantized + dispatched a. @@ -434,7 +434,7 @@ def prepare_monolithic( def finalize_monolithic(self, fused_expert_output: torch.Tensor) -> torch.Tensor: """ Optional method for subclasses compatible with monolithic - FusedMoEPermuteExpertsUnpermute kernels. + FusedMoEModularExperts kernels. Perform any combine plus apply weights and perform a reduction on the fused experts output. @@ -445,7 +445,7 @@ def finalize_monolithic(self, fused_expert_output: torch.Tensor) -> torch.Tensor # TODO: add supported activations method (return string) -class FusedMoEPermuteExpertsUnpermuteBase(ABC): +class FusedMoEExperts(ABC): def __init__( self, moe_config: FusedMoEConfig, @@ -505,7 +505,7 @@ def activation_format() -> FusedMoEActivationFormat: @staticmethod def is_supported_config( - cls: type["FusedMoEPermuteExpertsUnpermute"], + cls: type["FusedMoEModularExperts"], moe_config: FusedMoEConfig, weight_key: QuantKey | None, activation_key: QuantKey | None, @@ -686,7 +686,7 @@ def enable_chunking(self): ) -class FusedMoEPermuteExpertsUnpermute(FusedMoEPermuteExpertsUnpermuteBase): +class FusedMoEModularExperts(FusedMoEExperts): """ An abstract base class for the [Permute-Experts-Unpermute] step described above. @@ -869,7 +869,7 @@ def apply( raise NotImplementedError -class FusedMoEPermuteExpertsUnpermuteMonolithic(FusedMoEPermuteExpertsUnpermuteBase): +class FusedMoEMonolithicExperts(FusedMoEExperts): """ An abstract base class for the [Permute-Experts-Unpermute] step described above, but with the monolithic interface (accepts router logits @@ -912,11 +912,11 @@ def _slice_scales( return None -class FusedMoEModularKernelBase(torch.nn.Module): +class FusedMoEKernel(torch.nn.Module): def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalizeBase, - fused_experts: FusedMoEPermuteExpertsUnpermuteBase, + fused_experts: FusedMoEExperts, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, ): @@ -940,11 +940,11 @@ def __init__( if not ( ( isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) - and isinstance(fused_experts, FusedMoEPermuteExpertsUnpermuteMonolithic) + and isinstance(fused_experts, FusedMoEMonolithicExperts) ) or ( isinstance(prepare_finalize, FusedMoEPrepareAndFinalize) - and isinstance(fused_experts, FusedMoEPermuteExpertsUnpermute) + and isinstance(fused_experts, FusedMoEModularExperts) ) ): raise ValueError( @@ -965,24 +965,24 @@ def __init__( @staticmethod def make_mk( prepare_finalize: FusedMoEPrepareAndFinalizeBase, - fused_experts: FusedMoEPermuteExpertsUnpermuteBase, + fused_experts: FusedMoEExperts, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, - ) -> "FusedMoEModularKernelBase": + ) -> "FusedMoEKernel": """ - Factory method to create a FusedMoEModularKernelBase instance. + Factory method to create a FusedMoEKernel instance. """ if isinstance( prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic - ) and isinstance(fused_experts, FusedMoEPermuteExpertsUnpermuteMonolithic): - return FusedMoEModularKernelMonolithic( + ) and isinstance(fused_experts, FusedMoEMonolithicExperts): + return FusedMoEMonolithicKernel( prepare_finalize, fused_experts, shared_experts, moe_parallel_config, ) elif isinstance(prepare_finalize, FusedMoEPrepareAndFinalize) and isinstance( - fused_experts, FusedMoEPermuteExpertsUnpermute + fused_experts, FusedMoEModularExperts ): return FusedMoEModularKernel( prepare_finalize, @@ -1018,10 +1018,10 @@ def output_is_reduced(self) -> bool: @final -class FusedMoEModularKernel(FusedMoEModularKernelBase): +class FusedMoEModularKernel(FusedMoEKernel): """ This class combines a FusedMoEPrepareAndFinalize instance and - a FusedMoEPermuteExpertsUnpermute to provide an interface that + a FusedMoEModularExperts to provide an interface that is compatible with the `fused_experts` function in fused_moe.py. It takes care of managing any required scratch space. @@ -1034,7 +1034,7 @@ class FusedMoEModularKernel(FusedMoEModularKernelBase): def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalize, - fused_experts: FusedMoEPermuteExpertsUnpermute, + fused_experts: FusedMoEModularExperts, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, ): @@ -1042,9 +1042,9 @@ def __init__( raise TypeError( "prepare_finalize must be an instance of FusedMoEPrepareAndFinalize" ) - if not isinstance(fused_experts, FusedMoEPermuteExpertsUnpermute): + if not isinstance(fused_experts, FusedMoEModularExperts): raise TypeError( - "fused_experts must be an instance of FusedMoEPermuteExpertsUnpermute" + "fused_experts must be an instance of FusedMoEModularExperts" ) super().__init__( @@ -1055,7 +1055,7 @@ def __init__( ) self.prepare_finalize: FusedMoEPrepareAndFinalize = prepare_finalize - self.fused_experts: FusedMoEPermuteExpertsUnpermute = fused_experts + self.fused_experts: FusedMoEModularExperts = fused_experts def _chunk_info(self, M: int) -> tuple[int, int]: """ @@ -1099,7 +1099,7 @@ def _allocate_buffers( See `workspace_shapes` for a description of the remainder of arguments. Returns a tuple of (workspace13, workspace2, output) tensors. """ - assert isinstance(self.fused_experts, FusedMoEPermuteExpertsUnpermute) + assert isinstance(self.fused_experts, FusedMoEModularExperts) assert M_full > 0 and M_chunk > 0 num_chunks, _ = self._chunk_info(M_full) @@ -1342,7 +1342,7 @@ def _fused_experts( apply_router_weight_on_input: bool, expert_tokens_meta: ExpertTokensMetadata | None, ) -> torch.Tensor: - assert isinstance(self.fused_experts, FusedMoEPermuteExpertsUnpermuteBase) + assert isinstance(self.fused_experts, FusedMoEExperts) _, M_full, N, K, top_k = self.fused_experts.moe_problem_size( a1q, w1, w2, topk_ids @@ -1575,11 +1575,11 @@ def forward( @final -class FusedMoEModularKernelMonolithic(FusedMoEModularKernelBase): +class FusedMoEMonolithicKernel(FusedMoEKernel): def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic, - fused_experts: FusedMoEPermuteExpertsUnpermuteMonolithic, + fused_experts: FusedMoEMonolithicExperts, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, ): @@ -1588,10 +1588,9 @@ def __init__( "prepare_finalize must be an instance of " "FusedMoEPrepareAndFinalizeMonolithic" ) - if not isinstance(fused_experts, FusedMoEPermuteExpertsUnpermuteMonolithic): + if not isinstance(fused_experts, FusedMoEMonolithicExperts): raise TypeError( - "fused_experts must be an instance of " - "FusedMoEPermuteExpertsUnpermuteMonolithic" + "fused_experts must be an instance of FusedMoEMonolithicExperts" ) super().__init__( @@ -1623,7 +1622,7 @@ def forward_monolithic( that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) - assert isinstance(self.fused_experts, FusedMoEPermuteExpertsUnpermuteMonolithic) + assert isinstance(self.fused_experts, FusedMoEMonolithicExperts) # TODO(rob): add inplace support. a1q, a1q_scale, router_logits = self.prepare_finalize.prepare_monolithic( diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index ea2f6546f40b..86eddc1105da 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -52,7 +52,7 @@ class Fp8MoeBackend(Enum): def backend_to_kernel_cls( backend: Fp8MoeBackend, -) -> type[mk.FusedMoEPermuteExpertsUnpermute]: +) -> type[mk.FusedMoEModularExperts]: if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( FlashInferTrtLlmFp8Experts, @@ -132,12 +132,12 @@ def select_fp8_moe_backend( weight_key: QuantKey | None, activation_key: QuantKey | None, allow_vllm_cutlass: bool = False, -) -> tuple[Fp8MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute] | None]: +) -> tuple[Fp8MoeBackend, type[mk.FusedMoEModularExperts] | None]: """ Select the primary FP8 MoE backend Note: Shape-specific fallbacks may still occur at runtime. """ - k_cls: type[mk.FusedMoEPermuteExpertsUnpermute] | None = None + k_cls: type[mk.FusedMoEModularExperts] | None = None if config.is_lora_enabled: return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON) @@ -190,7 +190,7 @@ def _return_or_raise( weight_key: QuantKey | None, activation_key: QuantKey | None, activation_format: mk.FusedMoEActivationFormat, - ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute]]: + ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEModularExperts]]: k_cls = backend_to_kernel_cls(backend) supported, reason = k_cls.is_supported_config( k_cls, config, weight_key, activation_key, activation_format @@ -429,11 +429,11 @@ def make_fp8_moe_quant_config( def make_fp8_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, - experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute], + experts_cls: type[mk.FusedMoEModularExperts], fp8_backend: Fp8MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, shared_experts: torch.nn.Module | None = None, -) -> tuple[mk.FusedMoEModularKernelBase, bool]: +) -> tuple[mk.FusedMoEKernel, bool]: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( moe=moe_config, @@ -464,7 +464,7 @@ def make_fp8_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - kernel = mk.FusedMoEModularKernelBase.make_mk( + kernel = mk.FusedMoEKernel.make_mk( prepare_finalize, experts, shared_experts=( diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 6fc98b5aaa1d..58de11a0740a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -65,7 +65,7 @@ def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool: def backend_to_kernel_cls( backend: NvFp4MoeBackend, -) -> type[mk.FusedMoEPermuteExpertsUnpermute]: +) -> type[mk.FusedMoEModularExperts]: if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_nvfp4_moe import ( FlashInferTrtLlmNvFp4Experts, @@ -108,7 +108,7 @@ def select_nvfp4_moe_backend( config: FusedMoEConfig, weight_key: QuantKey | None, activation_key: QuantKey | None, -) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute] | None]: +) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEModularExperts] | None]: """ Select the primary NvFP4 MoE backend Note: Shape-specific fallbacks may still occur at runtime. @@ -161,7 +161,7 @@ def _return_or_raise( weight_key: QuantKey | None, activation_key: QuantKey | None, activation_format: mk.FusedMoEActivationFormat, - ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute]]: + ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEModularExperts]]: k_cls = backend_to_kernel_cls(backend) supported, reason = k_cls.is_supported_config( k_cls, config, weight_key, activation_key, activation_format @@ -364,10 +364,10 @@ def make_nvfp4_moe_quant_config( def make_nvfp4_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, - experts_cls: type[mk.FusedMoEPermuteExpertsUnpermuteBase], + experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, shared_experts: torch.nn.Module | None = None, -) -> mk.FusedMoEModularKernelBase: +) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( moe=moe_config, diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index dc50b7c96db1..2c9139474dd5 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -137,7 +137,7 @@ def make_unquantized_moe_kernel( FlashInferExperts, ) - kernel = mk.FusedMoEModularKernelBase.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, @@ -150,7 +150,7 @@ def make_unquantized_moe_kernel( AiterExperts, ) - kernel = mk.FusedMoEModularKernelBase.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), AiterExperts( moe_config=moe_config, @@ -160,7 +160,7 @@ def make_unquantized_moe_kernel( elif backend == UnquantizedMoeBackend.TRITON: from vllm.model_executor.layers.fused_moe import TritonExperts - kernel = mk.FusedMoEModularKernelBase.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), TritonExperts( moe_config=moe_config, diff --git a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py index 33150da6f910..9daba081c431 100644 --- a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py @@ -286,7 +286,7 @@ def rocm_aiter_fused_experts( ) -class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute): +class AiterExperts(mk.FusedMoEModularExperts): @property def expects_unquantized_inputs(self) -> bool: return True diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index 99d4038ec381..a8beb8aceacd 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -10,7 +10,7 @@ class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce): """ - Useful in the case when some FusedMoEPermuteExpertsUnpermute + Useful in the case when some FusedMoEModularExperts implementation does not perform weight application and reduction but cannot address the needs of all the compatible PrepareAndFinalize implementations. diff --git a/vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py index f537f2f99ade..5ed257a7bed9 100644 --- a/vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py @@ -31,8 +31,8 @@ def __init__( @staticmethod def get_clses() -> tuple[ - type[mk.FusedMoEPermuteExpertsUnpermute], - type[mk.FusedMoEPermuteExpertsUnpermute], + type[mk.FusedMoEModularExperts], + type[mk.FusedMoEModularExperts], ]: return (CutlassExpertsFp8, TritonExperts) @@ -76,7 +76,7 @@ def _select_experts_impl( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: # Small batch fallback for sm100. if self.is_sm100 and hidden_states.shape[0] <= 8: return self.fallback_experts diff --git a/vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py index 7e41269dc538..7bf81b59056e 100644 --- a/vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py @@ -31,8 +31,8 @@ def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig @staticmethod def get_clses() -> tuple[ - type[mk.FusedMoEPermuteExpertsUnpermute], - type[mk.FusedMoEPermuteExpertsUnpermute], + type[mk.FusedMoEModularExperts], + type[mk.FusedMoEModularExperts], ]: return (DeepGemmExperts, TritonExperts) @@ -78,7 +78,7 @@ def _select_experts_impl( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: if is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2): return self.experts else: diff --git a/vllm/model_executor/layers/fused_moe/trtllm_moe.py b/vllm/model_executor/layers/fused_moe/trtllm_moe.py index aa7185040adf..90044bb9c28e 100644 --- a/vllm/model_executor/layers/fused_moe/trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/trtllm_moe.py @@ -17,7 +17,7 @@ ) -class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute): +class TrtLlmGenExperts(mk.FusedMoEModularExperts): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 2baf420daab2..708e6d1643e9 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -23,7 +23,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEActivationFormat, - FusedMoEPermuteExpertsUnpermute, + FusedMoEModularExperts, FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( @@ -63,7 +63,7 @@ def __init__(self, moe: FusedMoEConfig): self.rocm_aiter_moe_enabled = ( rocm_aiter_ops.is_fused_moe_enabled() and moe.is_act_and_mul ) - self.kernel: mk.FusedMoEModularKernelBase | None = None + self.kernel: mk.FusedMoEKernel | None = None self._is_monolithic = current_platform.is_cpu() or current_platform.is_xpu() @property @@ -91,7 +91,7 @@ def select_gemm_impl( self, prepare_finalize: FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> FusedMoEPermuteExpertsUnpermute: + ) -> FusedMoEModularExperts: assert self.moe_quant_config is not None if ( prepare_finalize.activation_format diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 5e31b3feda01..00f9a189bb6b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -20,7 +20,7 @@ FusedMoE, FusedMoEActivationFormat, FusedMoEMethodBase, - FusedMoEPermuteExpertsUnpermute, + FusedMoEModularExperts, FusedMoeWeightScaleSupported, UnquantizedFusedMoEMethod, ) @@ -572,7 +572,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -945,7 +945,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -1455,7 +1455,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: assert self.num_bits == 4, "only supporting w4" layer.w13_weight = layer.w13_weight_packed layer.w2_weight = layer.w2_weight_packed @@ -1714,7 +1714,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: if self.moe.is_lora_enabled: assert self.moe_quant_config is not None from vllm.triton_utils import HAS_TRITON @@ -2316,7 +2316,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: assert self.moe_quant_config is not None assert ( prepare_finalize.activation_format == FusedMoEActivationFormat.Standard @@ -2324,7 +2324,7 @@ def select_gemm_impl( from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8 - experts: FusedMoEPermuteExpertsUnpermute + experts: FusedMoEModularExperts logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__) experts = CutlassExpertsW4A8Fp8( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index e75531503bdf..1f0a5038438c 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoE, FusedMoEMethodBase, - FusedMoEPermuteExpertsUnpermute, + FusedMoEModularExperts, FusedMoEPrepareAndFinalize, FusedMoeWeightScaleSupported, ) @@ -897,7 +897,7 @@ def select_gemm_impl( self, prepare_finalize: FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> FusedMoEPermuteExpertsUnpermute: + ) -> FusedMoEModularExperts: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index b5998828bf86..95f49503bb7a 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -742,7 +742,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -1347,7 +1347,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 4ade9bc05e33..e42050e54da6 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -839,7 +839,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEPermuteExpertsUnpermute: + ) -> mk.FusedMoEModularExperts: if ( prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts From 849125461ba7afeb2fd5211c3aeed509132e9b2f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Thu, 29 Jan 2026 20:20:36 -0500 Subject: [PATCH 092/207] update naming Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/modular_kernel.py | 2 +- vllm/model_executor/layers/fused_moe/oracle/unquantized.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 26ad081c201e..3960779372f3 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -968,7 +968,7 @@ def make_mk( fused_experts: FusedMoEExperts, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, - ) -> "FusedMoEKernel": + ) -> "FusedMoEMonolithicKernel" | "FusedMoEModularKernel": """ Factory method to create a FusedMoEKernel instance. """ diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 2c9139474dd5..63d8727b90fa 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -137,7 +137,7 @@ def make_unquantized_moe_kernel( FlashInferExperts, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEModularKernel( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, @@ -150,7 +150,7 @@ def make_unquantized_moe_kernel( AiterExperts, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEModularKernel( MoEPrepareAndFinalizeNoEP(), AiterExperts( moe_config=moe_config, From 68a21b247adaaf33eb225115e34c1a7628085813 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Thu, 29 Jan 2026 20:21:17 -0500 Subject: [PATCH 093/207] update str Signed-off-by: Robert Shaw --- benchmarks/kernels/benchmark_cutlass_moe_fp8.py | 2 +- benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py | 4 ++-- benchmarks/kernels/benchmark_grouped_gemm_cutlass.py | 4 ++-- benchmarks/kernels/benchmark_moe.py | 2 +- tests/kernels/moe/test_block_fp8.py | 2 +- tests/kernels/moe/test_cutlass_moe.py | 4 ++-- tests/kernels/moe/test_deepgemm.py | 2 +- tests/kernels/moe/test_flashinfer.py | 4 ++-- tests/kernels/moe/test_nvfp4_moe.py | 2 +- vllm/model_executor/layers/fused_moe/oracle/unquantized.py | 4 ++-- 10 files changed, 15 insertions(+), 15 deletions(-) diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index f1234d821347..78b4339d0375 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -136,7 +136,7 @@ def bench_run( per_out_ch_quant=per_out_ch, ) - fn = mk.FusedMoEModularKernel( + fn = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index c1f4f0aa9fce..5372f9ab5a7b 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -196,7 +196,7 @@ def run_cutlass_moe_fp4( g2_alphas=w2_gs, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( make_dummy_moe_config(), @@ -241,7 +241,7 @@ def run_cutlass_from_graph( g2_alphas=w2_gs, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( make_dummy_moe_config(), diff --git a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py index 7b5daa62eb34..4e404ee9b33d 100644 --- a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py +++ b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py @@ -132,7 +132,7 @@ def run_cutlass_moe( per_act_token_quant=per_act_token, ) - fn = mk.FusedMoEModularKernel( + fn = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( @@ -164,7 +164,7 @@ def run_cutlass_from_graph( per_act_token_quant=per_act_token, ) - fn = mk.FusedMoEModularKernel( + fn = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index db1468f82073..eefeebcd94cc 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -203,7 +203,7 @@ def run(): deep_gemm_experts = None if use_deep_gemm: - deep_gemm_experts = mk.FusedMoEModularKernel( + deep_gemm_experts = mk.FusedMoEModularKernel.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=FusedMoEConfig( diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 508df9e328a9..230968a8b5fe 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -255,7 +255,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) block_shape=block_size, ) - deep_gemm_experts = mk.FusedMoEModularKernel( + deep_gemm_experts = mk.FusedMoEModularKernel.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 3a5a66a383dc..354d27bcee47 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -196,7 +196,7 @@ def slice_experts(): for kwargs, new_quant_config in slice_experts(): w2 = kwargs["w2"] a = kwargs["hidden_states"] - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( @@ -255,7 +255,7 @@ def run_8_bit( num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined] with_ep = num_local_experts is not None or num_local_experts == num_experts if not with_ep: - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 729b54753b0a..8363b6465adb 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -109,7 +109,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size): block_shape=block_size, ) - deep_gemm_experts = mk.FusedMoEModularKernel( + deep_gemm_experts = mk.FusedMoEModularKernel.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 0bbf4ae2687f..ebf5dadbb5d3 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -210,7 +210,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( quant_config=quant_config, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferTrtLlmFp8Experts( moe_config=td.layer.moe, @@ -311,7 +311,7 @@ def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig: routing_method=RoutingMethodType.TopK, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index a22b2088bb08..b2e2916703f3 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -89,7 +89,7 @@ def test_cutlass_fp4_moe_no_graph( w2_scale=w2_blockscale, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( moe_config=make_dummy_moe_config(), diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 63d8727b90fa..05ed39ccdcf6 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -137,7 +137,7 @@ def make_unquantized_moe_kernel( FlashInferExperts, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, @@ -150,7 +150,7 @@ def make_unquantized_moe_kernel( AiterExperts, ) - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEModularKernel.make_mk( MoEPrepareAndFinalizeNoEP(), AiterExperts( moe_config=moe_config, From e8bc729ed7c5467c8148a85975edbe7c1fd785e9 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Thu, 29 Jan 2026 20:27:16 -0500 Subject: [PATCH 094/207] stash changes Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 2 +- .../layers/fused_moe/fused_moe_method_base.py | 12 +++++---- .../fused_moe/fused_moe_modular_method.py | 8 +++--- vllm/model_executor/layers/fused_moe/layer.py | 2 +- .../compressed_tensors_moe.py | 26 +++++++++---------- .../model_executor/layers/quantization/fp8.py | 10 +++---- .../layers/quantization/modelopt.py | 20 +++++++------- .../model_executor/warmup/deep_gemm_warmup.py | 2 +- 8 files changed, 42 insertions(+), 40 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index ebf5dadbb5d3..eae09359d957 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -218,7 +218,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( ), ) - flashinfer_output = kernel.forward_monolithic( + flashinfer_output = kernel( hidden_states=td.hidden_states, w1=td.layer.w13_weight, w2=td.layer.w2_weight, 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 ec52da8d3b51..dce28407f83f 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 @@ -27,19 +27,21 @@ def __init__(self, moe: FusedMoEConfig): super().__init__() self.moe: FusedMoEConfig = moe self.moe_quant_config: FusedMoEQuantConfig | None = None - self.moe_mk: mk.FusedMoEKernel | None = None + self.moe_kernel: mk.FusedMoEKernel | None = None @property def supports_internal_mk(self) -> bool: # NOTE(rob): temporary attribute to indicate support for # completed migration to the new internal MK interface. - return self.moe_mk is not None + return self.moe_kernel is not None @property def mk_owns_shared_expert(self) -> bool: # NOTE(rob): temporary attribute to indicate support for # completed migration to the new internal MK interface. - return self.moe_mk is not None and self.moe_mk.shared_experts is not None + return ( + self.moe_kernel is not None and self.moe_kernel.shared_experts is not None + ) @abstractmethod def create_weights( @@ -93,8 +95,8 @@ def get_fused_moe_quant_config( @property def topk_indices_dtype(self) -> torch.dtype | None: - if self.moe_mk is not None: - return self.moe_mk.prepare_finalize.topk_indices_dtype() + if self.moe_kernel is not None: + return self.moe_kernel.prepare_finalize.topk_indices_dtype() return None @property diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 7a2244a9bc1d..69bb47be0117 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -30,11 +30,11 @@ def __init__( ): super().__init__(old_quant_method.moe) self.moe_quant_config = old_quant_method.moe_quant_config - self.moe_mk = experts + self.moe_kernel = experts self.disable_expert_map = getattr( old_quant_method, "disable_expert_map", - not self.moe_mk.supports_expert_map(), + not self.moe_kernel.supports_expert_map(), ) self.old_quant_method = old_quant_method assert not self.old_quant_method.is_monolithic @@ -92,8 +92,8 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( hidden_states=x, w1=layer.w13_weight, w2=layer.w2_weight, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 521d4bc9f0d8..c050653aefe2 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -1526,7 +1526,7 @@ def must_reduce_shared_expert_outputs(self) -> bool: assert self.quant_method is not None return ( isinstance(self.quant_method, FusedMoEModularMethod) - and self.quant_method.moe_mk.output_is_reduced() # type: ignore[union-attr] + and self.quant_method.moe_kernel.output_is_reduced() # type: ignore[union-attr] ) def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 00f9a189bb6b..463a1118f385 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -323,7 +323,7 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config is not None: - self.moe_mk = make_nvfp4_moe_kernel( + self.moe_kernel = make_nvfp4_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, @@ -338,8 +338,8 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -551,7 +551,7 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: # Setup modular kernel. self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None - self.moe_mk = make_nvfp4_moe_kernel( + self.moe_kernel = make_nvfp4_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, @@ -602,8 +602,8 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk.forward_monolithic( + assert self.moe_kernel is not None + return self.moe_kernel.forward_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -625,8 +625,8 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -923,7 +923,7 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config: assert self.experts_cls is not None - self.moe_mk, self.use_inplace = make_fp8_moe_kernel( + self.moe_kernel, self.use_inplace = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, fp8_backend=self.fp8_backend, @@ -977,8 +977,8 @@ def apply_monolithic( assert self.is_monolithic assert layer.activation == "silu" - assert self.moe_mk is not None - return self.moe_mk.forward_monolithic( + assert self.moe_kernel is not None + return self.moe_kernel.forward_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -1001,8 +1001,8 @@ def apply( topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert not self.is_monolithic - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 1f0a5038438c..dfc651d6b2ee 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -827,7 +827,7 @@ def _setup_kernel( self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None - self.moe_mk, self.use_inplace = make_fp8_moe_kernel( + self.moe_kernel, self.use_inplace = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, fp8_backend=self.fp8_backend, @@ -940,8 +940,8 @@ def apply_monolithic( if layer.enable_eplb: raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.") - assert self.moe_mk is not None - return self.moe_mk.forward_monolithic( + assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -959,9 +959,9 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None + assert self.moe_kernel is not None assert not self.is_monolithic - return self.moe_mk( + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 95f49503bb7a..e5413980d463 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -858,7 +858,7 @@ def _setup_kernel( # Setup modular kernel. self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None - self.moe_mk, self.use_inplace = make_fp8_moe_kernel( + self.moe_kernel, self.use_inplace = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, fp8_backend=self.fp8_backend, @@ -926,8 +926,8 @@ def apply_monolithic( raise NotImplementedError( "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend." ) - assert self.moe_mk is not None - return self.moe_mk.forward_monolithic( + assert self.moe_kernel is not None + return self.moe_kernel.forward_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -949,8 +949,8 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -1525,7 +1525,7 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: # Setup modular kernel. self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None - self.moe_mk = make_nvfp4_moe_kernel( + self.moe_kernel = make_nvfp4_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, @@ -1565,8 +1565,8 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk.forward_monolithic( + assert self.moe_kernel is not None + return self.moe_kernel.forward_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -1588,8 +1588,8 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index cd4efe1ca776..aa2f63953bb0 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -170,7 +170,7 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool: # Further check if the ModularKernel implementation uses the DeepGemmExperts return isinstance( - module.quant_method.moe_mk, (DeepGemmExperts, TritonOrDeepGemmExperts) + module.quant_method.moe_kernel, (DeepGemmExperts, TritonOrDeepGemmExperts) ) From d28138ab1df5a54df474f514315d536bed23d58c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sat, 31 Jan 2026 17:03:09 -0500 Subject: [PATCH 095/207] updated class heirarchy Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 59 ++------ .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 77 +++++++--- .../layers/fused_moe/modular_kernel.py | 30 ++-- .../layers/fused_moe/prepare_finalize.py | 136 +++++++++++------- 4 files changed, 170 insertions(+), 132 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 173d73d64405..ba3818875ddb 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -11,9 +11,6 @@ FusedMoEQuantConfig, RoutingMethodType, ) -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceNoOP, -) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, @@ -23,7 +20,11 @@ from vllm.v1.engine.utils import current_platform -class FlashInferTrtLlmFp8Experts(mk.FusedMoEModularExperts): +class FlashInferTrtLlmFp8Experts(mk.FusedMoEMonolithicExperts): + """ + Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface. + """ + def __init__( self, moe_config: FusedMoEConfig, @@ -128,47 +129,7 @@ def supports_chunking(self) -> bool: def supports_expert_map(self) -> bool: return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - return TopKWeightAndReduceNoOP() - - def workspace_shapes( - self, - M: int, - N: int, - K: int, - topk: int, - global_num_experts: int, - local_num_experts: int, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - activation: str, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - raise NotImplementedError( - f"{self.__class__.__name__} only supports the apply_monolithic interface." - ) - - def apply( - self, - output: torch.Tensor, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: str, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, - workspace13: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - apply_router_weight_on_input: bool, - ): - raise NotImplementedError( - f"{self.__class__.__name__} only supports the apply_monolithic interface." - ) - - def _apply_per_block_monolithic( + def _apply_per_block( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -227,7 +188,7 @@ def _apply_per_block_monolithic( routing_method_type=self.routing_method_type, ) - def _apply_per_tensor_monolithic( + def _apply_per_tensor( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -274,7 +235,7 @@ def _apply_per_tensor_monolithic( ) return out - def apply_monolithic( + def apply( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -292,7 +253,7 @@ def apply_monolithic( topk_group: int | None = None, ) -> torch.Tensor: if self.quant_config.block_shape is not None: - return self._apply_per_block_monolithic( + return self._apply_per_block( hidden_states, w1, w2, @@ -307,7 +268,7 @@ def apply_monolithic( routed_scaling_factor=routed_scaling_factor, ) elif self.quant_config.is_per_tensor: - return self._apply_per_tensor_monolithic( + return self._apply_per_tensor( hidden_states, w1, w2, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index a11b711bcfdc..36333cbbdd7c 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -22,7 +22,11 @@ from vllm.platforms import current_platform -class FlashInferTrtLlmNvFp4Experts(mk.FusedMoEModularExperts): +class FlashInferTrtLlmNvFp4ExpertsBase(mk.FusedMoEExperts): + """ + NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface. + """ + def __init__( self, moe_config: FusedMoEConfig, @@ -73,6 +77,24 @@ def _supports_activation(activation: str) -> bool: """Supports only SiLU activation.""" return activation in ["silu"] + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def supports_chunking(self) -> bool: + return False + + def supports_expert_map(self) -> bool: + return False + + +class FlashInferTrtLlmNvFp4ExpertsModular( + FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEModularExperts +): + """ + Modular version of the implementation (just the experts). + """ + @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: """Supports EP and TP.""" @@ -84,26 +106,7 @@ def _supports_routing_method( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - # NOTE(rob): this is a conservative list. - return routing_method_type in [ - RoutingMethodType.DeepSeekV3, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - RoutingMethodType.Llama4, - ] - - @staticmethod - def activation_format() -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.Standard - - def supports_chunking(self) -> bool: - return False - - def supports_expert_map(self) -> bool: - return False - - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - return TopKWeightAndReduceNoOP() + return True def workspace_shapes( self, @@ -127,6 +130,9 @@ def workspace_shapes( return (workspace1, workspace2, output) + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + def apply( self, output: torch.Tensor, @@ -186,7 +192,34 @@ def apply( output=output, ) - def apply_monolithic( + +class FlashInferTrtLlmNvFp4Experts( + FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEMonolithicExperts +): + """ + Monolithic version of the kernel (router + experts). + """ + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """The modular implementation should be used for the Dp/Ep case""" + return not moe_parallel_config.use_all2all_kernels + + @staticmethod + def _supports_routing_method( + routing_method_type: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # NOTE(rob): this is a conservative list. + return routing_method_type in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.Llama4, + ] + + def apply( self, hidden_states: torch.Tensor, w1: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 3960779372f3..e0bb4b3353cd 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -406,7 +406,7 @@ class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalizeBase): """ @abstractmethod - def prepare_monolithic( + def prepare( self, a1: torch.Tensor, router_logits: torch.Tensor, @@ -427,11 +427,10 @@ def prepare_monolithic( - quantized + dispatched a. - Optional quantized + dispatched a1_scales. """ - raise NotImplementedError( - f"prepare_monolithic not supported for {self.__class__.__name__}" - ) + raise NotImplementedError - def finalize_monolithic(self, fused_expert_output: torch.Tensor) -> torch.Tensor: + @abstractmethod + def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor: """ Optional method for subclasses compatible with monolithic FusedMoEModularExperts kernels. @@ -497,8 +496,7 @@ def activation_format() -> FusedMoEActivationFormat: """ raise NotImplementedError - # - + # # Various helpers for registering support for various features. # Used by the oracle to select a particular kernel for a deployment. # @@ -692,6 +690,10 @@ class FusedMoEModularExperts(FusedMoEExperts): above. """ + @staticmethod + def is_monolithic() -> bool: + return False + def moe_problem_size( self, a1: torch.Tensor, @@ -876,7 +878,11 @@ class FusedMoEMonolithicExperts(FusedMoEExperts): rather than topk ids and weights). """ - def apply_monolithic( + @staticmethod + def is_monolithic() -> bool: + return False + + def apply( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -1600,7 +1606,7 @@ def __init__( moe_parallel_config, ) - def forward_monolithic( + def forward( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -1625,14 +1631,14 @@ def forward_monolithic( assert isinstance(self.fused_experts, FusedMoEMonolithicExperts) # TODO(rob): add inplace support. - a1q, a1q_scale, router_logits = self.prepare_finalize.prepare_monolithic( + a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( hidden_states, router_logits=router_logits, quant_config=self.fused_experts.quant_config, defer_input_quant=self.fused_experts.expects_unquantized_inputs, ) - fused_out = self.fused_experts.apply_monolithic( + fused_out = self.fused_experts.apply( hidden_states=a1q, w1=w1, w2=w2, @@ -1649,6 +1655,6 @@ def forward_monolithic( topk_group=topk_group, ) - output = self.prepare_finalize.finalize_monolithic(fused_out) + output = self.prepare_finalize.finalize(fused_out) return output diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index ee7035f84fe8..df336fcbc7c1 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -9,13 +9,26 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate, - TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.flashinfer import nvfp4_block_scale_interleave -class MoEPrepareAndFinalizeNaiveEP(mk.FusedMoEPrepareAndFinalize): +class MoEPrepareAndFinalizeNaiveEPBase(mk.FusedMoEPrepareAndFinalizeBase): + """ + Base class for Naive Prepare/Finalize for Dp/Ep with two subclasses: + * Modular Case + * Monolithic Case + + In modular case, a separate router runs *before* and we dispatch + the topk weights and ids. + + In monolithic case, the router runs *inside* the MoE kernel so we + dispatch the router logits. + + In both cases, the quantization of X happens prior to dispatching. + """ + def __init__( self, is_sequence_parallel: bool = False, @@ -95,6 +108,17 @@ def _unwrap_scale_and_prepare_for_moe( return a1q_scale + +class MoEPrepareAndFinalizeNaiveEP( + MoEPrepareAndFinalizeNaiveEPBase, mk.FusedMoEPrepareAndFinalize +): + """ + Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. + + Uses Torch AR/RS or AR for dispatch/combine operations, applied + to the topk weights and ids. + """ + def prepare( self, a1: torch.Tensor, @@ -160,6 +184,17 @@ def finalize( get_ep_group().combine(out, is_sequence_parallel=self.is_sequence_parallel) ) + +class MoEPrepareAndFinalizeNaiveEPMonolithic( + MoEPrepareAndFinalizeNaiveEPBase, mk.FusedMoEPrepareAndFinalizeMonolithic +): + """ + Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. + + Uses Torch AR/RS or AR for dispatch/combine operations, applied + to the router logits (the MoE kernel runs the router internally). + """ + def prepare_monolithic( self, a1: torch.Tensor, @@ -189,12 +224,10 @@ def prepare_monolithic( return a1q, a1q_scale, router_logits - def finalize_monolithic( + def finalize( self, fused_expert_output: torch.Tensor, - weight_and_reduce_impl: mk.TopKWeightAndReduce, ) -> torch.Tensor: - assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP) out = get_ep_group().combine( fused_expert_output, is_sequence_parallel=self.is_sequence_parallel ) @@ -202,7 +235,13 @@ def finalize_monolithic( return out -class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize): +class MoEPrepareAndFinalizeNoEPBase(mk.FusedMoEPrepareAndFinalizeBase): + """ + Base class for TP case Prepare/Finalize. + * prepare: applies input quantization + * finalize: applies the reduction (if needed) + """ + @property def activation_format(self) -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard @@ -219,37 +258,23 @@ def num_dispatchers(self) -> int: def output_is_reduced(self) -> bool: return False - def prepare( + def _quantize_input( self, a1: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, quant_config: FusedMoEQuantConfig, defer_input_quant: bool = False, - ) -> mk.PrepareResultType: - if apply_router_weight_on_input: - topk = topk_ids.size(1) - # TODO: this only works for topK=1, will need to update for topK>1 - assert topk == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - # Note: do not use inplace for shared experts overlap - a1 = a1 * topk_weights.to(a1.dtype) - + ) -> tuple[torch.Tensor, torch.Tensor | None]: # Defer input quant to moe kernel for backends (e.g. AITER, FI) # which use a single kernel call for quant + experts. if defer_input_quant: - return a1, None, None, None, None + return a1, None input_sf = ( quant_config.a1_gscale if quant_config.use_nvfp4_w4a4 else quant_config.a1_scale ) - a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input( + a1q, a1q_scale = moe_kernel_quantize_input( a1, input_sf, quant_dtype=quant_config.quant_dtype, @@ -258,34 +283,35 @@ def prepare( is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, ) - return a1q, a1q_scale, None, None, None + return a1q, a1q_scale - def prepare_monolithic( + +class MoEPrepareAndFinalizeNoEPMonolithic( + mk.FusedMoEPrepareAndFinalize, MoEPrepareAndFinalizeNoEPBase +): + def prepare( self, a1: torch.Tensor, - router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, quant_config: FusedMoEQuantConfig, defer_input_quant: bool = False, - ) -> mk.PrepareMonolithicResultType: - # Defer input quant to moe kernel for backends (e.g. AITER, FI) - # which use a single kernel call for quant + experts. - if defer_input_quant: - return a1, None, router_logits + ) -> mk.PrepareResultType: + if apply_router_weight_on_input: + topk = topk_ids.size(1) + # TODO: this only works for topK=1, will need to update for topK>1 + assert topk == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + # Note: do not use inplace for shared experts overlap + a1 = a1 * topk_weights.to(a1.dtype) - input_sf = ( - quant_config.a1_gscale - if quant_config.use_nvfp4_w4a4 - else quant_config.a1_scale - ) - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - input_sf, - quant_dtype=quant_config.quant_dtype, - per_act_token_quant=quant_config.per_act_token_quant, - block_shape=quant_config.block_shape, - is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, - ) - return a1q, a1q_scale, router_logits + a1q, a1q_scale = self._quantize_input(a1, quant_config, defer_input_quant) + + return a1q, a1q_scale, None, None, None def finalize( self, @@ -306,10 +332,22 @@ def finalize( apply_router_weight_on_input=apply_router_weight_on_input, ) - def finalize_monolithic( + +class MoEPrepareAndFinalizeNoEP( + mk.FusedMoEPrepareAndFinalizeMonolithic, MoEPrepareAndFinalizeNoEPBase +): + def prepare( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareMonolithicResultType: + a1q, a1q_scale = self._quantize_input(a1, quant_config, defer_input_quant) + return a1q, a1q_scale, router_logits + + def finalize( self, fused_expert_output: torch.Tensor, - weight_and_reduce_impl: mk.TopKWeightAndReduce, ) -> torch.Tensor: - assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP) return fused_expert_output From edd45a17c7e3cc39416f763b8b80614e96fc8576 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sat, 31 Jan 2026 17:27:52 -0500 Subject: [PATCH 096/207] update class heirarchy Signed-off-by: Robert Shaw --- .../layers/fused_moe/all2all_utils.py | 19 ++-- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 2 +- .../layers/fused_moe/oracle/nvfp4.py | 88 ++++++++++--------- .../layers/fused_moe/prepare_finalize.py | 29 +++++- 4 files changed, 85 insertions(+), 53 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 97ee5a59fc44..172acfbc394a 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -17,11 +17,11 @@ FlashInferA2APrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEPrepareAndFinalize, + FusedMoEPrepareAndFinalizeBase, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNaiveEP, - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNaiveEPBase, + MoEPrepareAndFinalizeNoEPBase, ) from vllm.platforms import current_platform from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx @@ -82,7 +82,7 @@ def maybe_make_prepare_finalize( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, allow_new_interface: bool = False, use_monolithic: bool = False, -) -> FusedMoEPrepareAndFinalize | None: +) -> FusedMoEPrepareAndFinalizeBase | None: # NOTE(rob): we are migrating each quant_method to hold the MK # in all cases. The allow_new_interface=False flag allow us to fall # back to the old method for methods that have not yet been migrated. @@ -107,19 +107,19 @@ def maybe_make_prepare_finalize( "Detected DP deployment with no --enable-expert-parallel. " "Falling back to AllGather+ReduceScatter dispatch/combine." ) - return MoEPrepareAndFinalizeNaiveEP( + return MoEPrepareAndFinalizeNaiveEPBase.make( is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, num_dispatchers=( get_ep_group().device_communicator.all2all_manager.world_size ), ) else: - return MoEPrepareAndFinalizeNoEP() + return MoEPrepareAndFinalizeNoEPBase.make(use_monolithic) all2all_manager = get_ep_group().device_communicator.all2all_manager assert all2all_manager is not None - prepare_finalize: FusedMoEPrepareAndFinalize | None = None + prepare_finalize: FusedMoEPrepareAndFinalizeBase | None = None if moe.use_pplx_kernels: assert quant_config is not None @@ -247,8 +247,9 @@ def maybe_make_prepare_finalize( ) elif moe.use_naive_all2all_kernels and allow_new_interface: - prepare_finalize = MoEPrepareAndFinalizeNaiveEP( - is_sequence_parallel=(moe.moe_parallel_config.is_sequence_parallel), + prepare_finalize = MoEPrepareAndFinalizeNaiveEPBase.make( + use_monolithic=use_monolithic, + is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, num_dispatchers=all2all_manager.world_size, ) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 36333cbbdd7c..a0d348137c68 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -193,7 +193,7 @@ def apply( ) -class FlashInferTrtLlmNvFp4Experts( +class FlashInferTrtLlmNvFp4ExpertsMonolithic( FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEMonolithicExperts ): """ diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 58de11a0740a..497ca5a618d9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -65,41 +65,46 @@ def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool: def backend_to_kernel_cls( backend: NvFp4MoeBackend, -) -> type[mk.FusedMoEModularExperts]: +) -> list[type[mk.FusedMoEModularExperts]]: if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_nvfp4_moe import ( - FlashInferTrtLlmNvFp4Experts, + FlashInferTrtLlmNvFp4ExpertsModular, + FlashInferTrtLlmNvFp4ExpertsMonolithic, ) - return FlashInferTrtLlmNvFp4Experts + # NOTE: prefer Monolthic > Modular, so return Monolithic first. + return [ + FlashInferTrtLlmNvFp4ExpertsMonolithic, + FlashInferTrtLlmNvFp4ExpertsModular, + ] elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS: from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, ) - return FlashInferExperts + return [FlashInferExperts] elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL: from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import ( FlashInferCuteDSLExperts, ) - return FlashInferCuteDSLExperts + return [FlashInferCuteDSLExperts] elif backend == NvFp4MoeBackend.VLLM_CUTLASS: from vllm.model_executor.layers.fused_moe.cutlass_moe import ( CutlassExpertsFp4, ) - return CutlassExpertsFp4 + return [CutlassExpertsFp4] elif backend == NvFp4MoeBackend.MARLIN: from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( MarlinExperts, ) - return MarlinExperts + return [MarlinExperts] else: raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}") @@ -162,13 +167,14 @@ def _return_or_raise( activation_key: QuantKey | None, activation_format: mk.FusedMoEActivationFormat, ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEModularExperts]]: - k_cls = backend_to_kernel_cls(backend) - supported, reason = k_cls.is_supported_config( - k_cls, config, weight_key, activation_key, activation_format - ) - if supported: - logger.info_once(_make_log_backend(backend)) - return backend, k_cls + for k_cls in backend_to_kernel_cls(backend): + supported, reason = k_cls.is_supported_config( + k_cls, config, weight_key, activation_key, activation_format + ) + if supported: + logger.info_once(_make_log_backend(backend)) + return backend, k_cls + raise ValueError(_make_log_unsupported(backend, reason)) if envs.is_set("VLLM_USE_FLASHINFER_MOE_FP4"): @@ -186,21 +192,21 @@ def _return_or_raise( else: # If the user is not explicit about the backend, try each. for backend in FLASHINFER_NVFP4_MOE_BACKENDS: - k_cls = backend_to_kernel_cls(backend) - supported, reason = k_cls.is_supported_config( - k_cls, - config, - weight_key, - activation_key, - activation_format, - ) - if supported: - logger.info_once(_make_log_backend(backend), scope="local") - return backend, None - else: - logger.debug_once( - _make_log_unsupported(backend, reason), scope="local" + for k_cls in backend_to_kernel_cls(backend): + supported, reason = k_cls.is_supported_config( + k_cls, + config, + weight_key, + activation_key, + activation_format, ) + if supported: + logger.info_once(_make_log_backend(backend), scope="local") + return backend, None + else: + logger.debug_once( + _make_log_unsupported(backend, reason), scope="local" + ) raise NotImplementedError( "Found VLLM_USE_FLASHINFER_MOE_FP4=1, but no " @@ -215,20 +221,20 @@ def _return_or_raise( # Select kernels in order of backend. for backend in AVAILABLE_BACKENDS: - k_cls = backend_to_kernel_cls(backend) - supported, reason = k_cls.is_supported_config( - k_cls, - config, - weight_key, - activation_key, - activation_format, - ) + for k_cls in backend_to_kernel_cls(backend): + supported, reason = k_cls.is_supported_config( + k_cls, + config, + weight_key, + activation_key, + activation_format, + ) - if supported: - logger.info_once(_make_log_backend(backend), scope="local") - return backend, k_cls - else: - logger.debug_once(_make_log_unsupported(backend, reason), scope="local") + if supported: + logger.info_once(_make_log_backend(backend), scope="local") + return backend, k_cls + else: + logger.debug_once(_make_log_unsupported(backend, reason), scope="local") raise NotImplementedError( "No NvFp4 MoE backend supports the deployment configuration." diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index df336fcbc7c1..75f280d5a45d 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -29,6 +29,21 @@ class MoEPrepareAndFinalizeNaiveEPBase(mk.FusedMoEPrepareAndFinalizeBase): In both cases, the quantization of X happens prior to dispatching. """ + @staticmethod + def make( + is_sequence_parallel: bool = False, + num_dispatchers: int = 1, + use_monolithic: bool = False, + ) -> "MoEPrepareAndFinalizeNaiveEP" | "MoEPrepareAndFinalizeNaiveEPMonolithic": + cls = ( + MoEPrepareAndFinalizeNaiveEPMonolithic + if use_monolithic + else MoEPrepareAndFinalizeNaiveEP + ) + return cls( + is_sequence_parallel=is_sequence_parallel, num_dispatchers=num_dispatchers + ) + def __init__( self, is_sequence_parallel: bool = False, @@ -242,6 +257,16 @@ class MoEPrepareAndFinalizeNoEPBase(mk.FusedMoEPrepareAndFinalizeBase): * finalize: applies the reduction (if needed) """ + @staticmethod + def make( + use_monolithic: bool, + ) -> "MoEPrepareAndFinalizeNoEP" | "MoEPrepareAndFinalizeNoEPMonolithic": + return ( + MoEPrepareAndFinalizeNoEPMonolithic() + if use_monolithic + else MoEPrepareAndFinalizeNoEP() + ) + @property def activation_format(self) -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard @@ -286,7 +311,7 @@ def _quantize_input( return a1q, a1q_scale -class MoEPrepareAndFinalizeNoEPMonolithic( +class MoEPrepareAndFinalizeNoEP( mk.FusedMoEPrepareAndFinalize, MoEPrepareAndFinalizeNoEPBase ): def prepare( @@ -333,7 +358,7 @@ def finalize( ) -class MoEPrepareAndFinalizeNoEP( +class MoEPrepareAndFinalizeNoEPMonolithic( mk.FusedMoEPrepareAndFinalizeMonolithic, MoEPrepareAndFinalizeNoEPBase ): def prepare( From 0bb9770d706e4a53d820ac8e6d5e5e0c36e617f2 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sat, 31 Jan 2026 18:33:21 -0500 Subject: [PATCH 097/207] seems to be working properly Signed-off-by: Robert Shaw --- docs/design/fused_moe_modular_kernel.md | 56 +++++++-------- docs/design/moe_kernel_features.md | 4 +- .../moe/modular_kernel_tools/cli_args.py | 2 +- .../moe/modular_kernel_tools/common.py | 2 +- .../moe/modular_kernel_tools/mk_objects.py | 8 +-- .../moe/test_modular_kernel_combinations.py | 4 +- .../layers/fused_moe/__init__.py | 4 +- .../layers/fused_moe/batched_deep_gemm_moe.py | 2 +- .../layers/fused_moe/cutlass_moe.py | 6 +- .../layers/fused_moe/deep_gemm_moe.py | 2 +- .../layers/fused_moe/fallback.py | 12 ++-- .../fused_moe/flashinfer_cutedsl_moe.py | 2 +- .../fused_moe/flashinfer_cutlass_moe.py | 2 +- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 2 +- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 4 +- .../layers/fused_moe/fused_batched_moe.py | 4 +- .../layers/fused_moe/fused_marlin_moe.py | 2 +- .../layers/fused_moe/fused_moe.py | 4 +- .../layers/fused_moe/fused_moe_method_base.py | 4 +- .../fused_moe/gpt_oss_triton_kernels_moe.py | 2 +- .../layers/fused_moe/modular_kernel.py | 69 ++++++++++--------- .../layers/fused_moe/oracle/fp8.py | 11 ++- .../layers/fused_moe/oracle/nvfp4.py | 7 +- .../layers/fused_moe/prepare_finalize.py | 9 +-- .../layers/fused_moe/rocm_aiter_fused_moe.py | 2 +- .../fused_moe/topk_weight_and_reduce.py | 2 +- .../layers/fused_moe/triton_cutlass_moe.py | 6 +- .../layers/fused_moe/triton_deep_gemm_moe.py | 6 +- .../layers/fused_moe/trtllm_moe.py | 2 +- .../fused_moe/unquantized_fused_moe_method.py | 4 +- .../compressed_tensors_moe.py | 25 +++---- .../model_executor/layers/quantization/fp8.py | 4 +- .../layers/quantization/modelopt.py | 10 +-- .../layers/quantization/mxfp4.py | 2 +- 34 files changed, 142 insertions(+), 145 deletions(-) diff --git a/docs/design/fused_moe_modular_kernel.md b/docs/design/fused_moe_modular_kernel.md index 32012a90ca2d..011422d03f37 100644 --- a/docs/design/fused_moe_modular_kernel.md +++ b/docs/design/fused_moe_modular_kernel.md @@ -38,19 +38,19 @@ FusedMoEModularKernel splits the FusedMoE operation into 3 parts, 1. TopKWeightAndReduce 2. FusedMoEPrepareAndFinalize -3. FusedMoEModularExperts +3. FusedMoEExpertsModular ### TopKWeightAndReduce -The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEModularExperts` is responsible for the Unpermute and `FusedMoEPrepareAndFinalize` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEModularExperts`. But some implementations choose to do it `FusedMoEPrepareAndFinalize`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class. +The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEExpertsModular` is responsible for the Unpermute and `FusedMoEPrepareAndFinalize` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEExpertsModular`. But some implementations choose to do it `FusedMoEPrepareAndFinalize`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class. Please find the implementations of TopKWeightAndReduce [here](../../vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py). `FusedMoEPrepareAndFinalize::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method. -The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEModularExperts` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. +The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEExpertsModular` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. -* `FusedMoEModularExperts::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEModularExperts` implementation does the weight application and reduction itself. -* `FusedMoEModularExperts::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEModularExperts` implementation needs the `FusedMoEPrepareAndFinalize::finalize()` to do the weight application and reduction. +* `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEExpertsModular` implementation does the weight application and reduction itself. +* `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEExpertsModular` implementation needs the `FusedMoEPrepareAndFinalize::finalize()` to do the weight application and reduction. ### FusedMoEPrepareAndFinalize @@ -59,9 +59,9 @@ The `prepare` function is responsible for input activation Quantization and All2 ![FusedMoEPrepareAndFinalize Blocks](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png) -### FusedMoEModularExperts +### FusedMoEExpertsModular -The `FusedMoEModularExperts` class is where the crux of the MoE operations happen. The `FusedMoEModularExperts` abstract class exposes a few important functions, +The `FusedMoEExpertsModular` class is where the crux of the MoE operations happen. The `FusedMoEExpertsModular` abstract class exposes a few important functions, * apply() * workspace_shapes() @@ -81,25 +81,25 @@ The `apply` method is where the implementations perform #### workspace_shapes() -The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEModularExperts::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. +The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. #### finalize_weight_and_reduce_impl() -It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEModularExperts::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section. -`FusedMoEModularExperts::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalize::finalize()` to use. +It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEExpertsModular::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section. +`FusedMoEExpertsModular::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalize::finalize()` to use. -![FusedMoEModularExperts Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png) +![FusedMoEExpertsModular Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png) ### FusedMoEModularKernel -`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalize` and `FusedMoEModularExperts` objects. +`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalize` and `FusedMoEExpertsModular` objects. `FusedMoEModularKernel` pseudocode/sketch, ```py class FusedMoEModularKernel: def __init__(self, prepare_finalize: FusedMoEPrepareAndFinalize, - fused_experts: FusedMoEModularExperts): + fused_experts: FusedMoEExpertsModular): self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts @@ -162,20 +162,20 @@ This section describes the significance of the various functions exposed by the We suggest picking an already existing `FusedMoEPrepareAndFinalize` implementation that matches your All2All implementation closely and using it as a reference. -### How To Add a FusedMoEModularExperts Type +### How To Add a FusedMoEExpertsModular Type -FusedMoEModularExperts performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows, +FusedMoEExpertsModular performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows, -`FusedMoEModularExperts::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format. +`FusedMoEExpertsModular::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format. -`FusedMoEModularExperts::supports_chunking()`: Return True if the implementation supports chunking. Typically +`FusedMoEExpertsModular::supports_chunking()`: Return True if the implementation supports chunking. Typically implementations that input `FusedMoEActivationFormat.Standard` support chunking and `FusedMoEActivationFormat.BatchedExperts` do not. -`FusedMoEModularExperts::supports_expert_map()`: Return True if the implementation supports expert map. +`FusedMoEExpertsModular::supports_expert_map()`: Return True if the implementation supports expert map. -`FusedMoEModularExperts::workspace_shapes()` / -`FusedMoEModularExperts::finalize_weight_and_reduce_impl` / -`FusedMoEModularExperts::apply`: Refer to `FusedMoEModularExperts` section above. +`FusedMoEExpertsModular::workspace_shapes()` / +`FusedMoEExpertsModular::finalize_weight_and_reduce_impl` / +`FusedMoEExpertsModular::apply`: Refer to `FusedMoEExpertsModular` section above. ### FusedMoEModularKernel Initialization @@ -194,7 +194,7 @@ Please refer to the implementations in, #### select_gemm_impl -The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEModularExperts` object. +The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEExpertsModular` object. Please refer to the implementations in, * `UnquantizedFusedMoEMethod` @@ -206,7 +206,7 @@ derived classes. #### init_prepare_finalize -Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalize` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEModularExperts` object and builds the `FusedMoEModularKernel` object +Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalize` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEModularKernel` object Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188). **Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used. @@ -217,7 +217,7 @@ We have `FusedMoEModularKernel` unit tests at [test_modular_kernel_combinations. The unit test iterates through all combinations of `FusedMoEPrepareAndFinalize` and `FusedMoEPremuteExpertsUnpermute` types and if they are compatible, runs some correctness tests. -If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEModularExperts` implementations, +If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEExpertsModular` implementations, 1. Add the implementation type to `MK_ALL_PREPARE_FINALIZE_TYPES` and `MK_FUSED_EXPERT_TYPES` in [mk_objects.py](../../tests/kernels/moe/modular_kernel_tools/mk_objects.py) respectively. 2. Update `Config::is_batched_prepare_finalize()`, `Config::is_batched_fused_experts()`, `Config::is_standard_fused_experts()`, @@ -226,24 +226,24 @@ If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEModularExperts` i Doing this will add the new implementation to the test suite. -### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEModularExperts` Compatibility +### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEExpertsModular` Compatibility The unit test file [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py) can also be executed as a standalone script. Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` -As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEModularExperts` compatibility. When invoked +As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEExpertsModular` compatibility. When invoked with incompatible types, the script will error. ### How To Profile Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py) The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible -`FusedMoEPrepareAndFinalize` and `FusedMoEModularExperts` types. +`FusedMoEPrepareAndFinalize` and `FusedMoEExpertsModular` types. Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` ## FusedMoEPrepareAndFinalize Implementations See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-modular-all2all-backends) for a list of all the available modular prepare and finalize subclasses. -## FusedMoEModularExperts +## FusedMoEExpertsModular See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-experts-kernels) for a list of all the available modular experts. diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 731f635c2555..df54255e45b6 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -76,7 +76,7 @@ Each experts kernel supports one or more activation functions, e.g. silu or gelu As with the backends, some experts support applying topk weights on the input activations. The entries in the column in this table only apply to the non-modular experts. -Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEModularExperts`. +Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEExpertsModular`. To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats. @@ -107,7 +107,7 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels The following table shows "families" of modular kernels that are intended to work together. There are some combinations which may work but have not yet been tested, e.g. flashinfer with other fp8 experts. Note that the "naive" backend will work with any non-modular experts. -| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEModularExperts` subclasses | +| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEExpertsModular` subclasses | |---------|-----------------------------------------|----------------------------------------------| | deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,
`TritonExperts`,
`TritonOrDeepGemmExperts`,
`CutlassExpertsFp8`,
`MarlinExperts` | | deepep_low_latency,
pplx | `DeepEPLLPrepareAndFinalize`,
`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,
`BatchedTritonExperts`,
`CutlassBatchedExpertsFp8`,
`BatchedMarlinExperts` | diff --git a/tests/kernels/moe/modular_kernel_tools/cli_args.py b/tests/kernels/moe/modular_kernel_tools/cli_args.py index ddc5a95cd08e..28be65127780 100644 --- a/tests/kernels/moe/modular_kernel_tools/cli_args.py +++ b/tests/kernels/moe/modular_kernel_tools/cli_args.py @@ -23,7 +23,7 @@ def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalize: return pf raise ValueError(f"Cannot find a PrepareFinalize type that matches {s}") - def to_experts_class_type(s: str) -> mk.FusedMoEModularExperts: + def to_experts_class_type(s: str) -> mk.FusedMoEExpertsModular: for fe in MK_FUSED_EXPERT_TYPES: if fe.__name__ == s: return fe diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 432e62ccfb9b..50174695b8c3 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -66,7 +66,7 @@ class Config: quant_config: TestMoEQuantConfig | None prepare_finalize_type: mk.FusedMoEPrepareAndFinalize - fused_experts_type: mk.FusedMoEModularExperts + fused_experts_type: mk.FusedMoEExpertsModular fused_moe_chunk_size: int | None world_size: int diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index d2f8e5b5cf27..8abcf656de80 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -73,11 +73,11 @@ class ExpertInfo: PREPARE_FINALIZE_INFO: dict[mk.FusedMoEPrepareAndFinalize, PrepareFinalizeInfo] = {} -EXPERT_INFO: dict[mk.FusedMoEModularExperts, ExpertInfo] = {} +EXPERT_INFO: dict[mk.FusedMoEExpertsModular, ExpertInfo] = {} MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] -MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEModularExperts] = [] +MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEExpertsModular] = [] standard_format = mk.FusedMoEActivationFormat.Standard batched_format = mk.FusedMoEActivationFormat.BatchedExperts @@ -444,12 +444,12 @@ def make_cutlass_strides( def make_fused_experts( - fused_experts_type: mk.FusedMoEModularExperts, + fused_experts_type: mk.FusedMoEExpertsModular, moe: FusedMoEConfig, quant_config: FusedMoEQuantConfig, num_dispatchers: int, N: int, -) -> mk.FusedMoEModularExperts: +) -> mk.FusedMoEExpertsModular: if ( fused_experts_type.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index dad4948ce659..8c72e7c8df75 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -259,7 +259,7 @@ def test_modular_kernel_combinations_multigpu( dtype: torch.dtype, quant_config: TestMoEQuantConfig | None, prepare_finalize_type: mk.FusedMoEPrepareAndFinalize, - fused_experts_type: mk.FusedMoEModularExperts, + fused_experts_type: mk.FusedMoEExpertsModular, chunk_size: int | None, world_size: int, pytestconfig, @@ -301,7 +301,7 @@ def test_modular_kernel_combinations_singlegpu( dtype: torch.dtype, quant_config: TestMoEQuantConfig | None, prepare_finalize_type: mk.FusedMoEPrepareAndFinalize, - fused_experts_type: mk.FusedMoEModularExperts, + fused_experts_type: mk.FusedMoEExpertsModular, chunk_size: int | None, world_size: int, pytestconfig, diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 32b18a767e6f..96da5e4f9069 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -17,7 +17,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEActivationFormat, - FusedMoEModularExperts, + FusedMoEExpertsModular, FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( @@ -56,7 +56,7 @@ def get_config() -> dict[str, Any] | None: "FusedMoEMethodBase", "UnquantizedFusedMoEMethod", "FusedMoeWeightScaleSupported", - "FusedMoEModularExperts", + "FusedMoEExpertsModular", "FusedMoEActivationFormat", "FusedMoEPrepareAndFinalize", "RoutingMethodType", diff --git a/vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py index 1bd5d986c62c..13d4807871c2 100644 --- a/vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py @@ -260,7 +260,7 @@ def persistent_masked_m_silu_mul_quant( return y_q, y_s -class BatchedDeepGemmExperts(mk.FusedMoEModularExperts): +class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 9d216425d20a..3f15ab7ecf06 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -259,7 +259,7 @@ def run_cutlass_moe_fp8( ) -class CutlassExpertsFp8Base(mk.FusedMoEModularExperts): +class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular): def __init__( self, moe_config: FusedMoEConfig, @@ -650,7 +650,7 @@ def run_cutlass_moe_fp4( return -class CutlassExpertsFp4(mk.FusedMoEModularExperts): +class CutlassExpertsFp4(mk.FusedMoEExpertsModular): @property def expects_unquantized_inputs(self) -> bool: return True @@ -902,7 +902,7 @@ def run_cutlass_moe_w4a8_fp8( ) -class CutlassExpertsW4A8Fp8(mk.FusedMoEModularExperts): +class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular): def __init__( self, out_dtype: torch.dtype | None, diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py index d98c1216a048..edcefa1760b9 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py @@ -112,7 +112,7 @@ def _valid_deep_gemm( return True -class DeepGemmExperts(mk.FusedMoEModularExperts): +class DeepGemmExperts(mk.FusedMoEExpertsModular): def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig): super().__init__(moe_config=moe_config, quant_config=quant_config) assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout() diff --git a/vllm/model_executor/layers/fused_moe/fallback.py b/vllm/model_executor/layers/fused_moe/fallback.py index 18391e2e83b9..4eebe41df4c5 100644 --- a/vllm/model_executor/layers/fused_moe/fallback.py +++ b/vllm/model_executor/layers/fused_moe/fallback.py @@ -10,13 +10,13 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey -class FallbackExperts(mk.FusedMoEModularExperts, ABC): +class FallbackExperts(mk.FusedMoEExpertsModular, ABC): """Base class for runtime dispatching of expert implementations.""" def __init__( self, - experts: mk.FusedMoEModularExperts, - fallback_experts: mk.FusedMoEModularExperts, + experts: mk.FusedMoEExpertsModular, + fallback_experts: mk.FusedMoEExpertsModular, ): super().__init__( moe_config=experts.moe_config, quant_config=experts.quant_config @@ -26,8 +26,8 @@ def __init__( @staticmethod def get_clses() -> tuple[ - type[mk.FusedMoEModularExperts], - type[mk.FusedMoEModularExperts], + type[mk.FusedMoEExpertsModular], + type[mk.FusedMoEExpertsModular], ]: """ Get the cls for the experts and fallback experts. @@ -148,7 +148,7 @@ def _select_experts_impl( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: raise NotImplementedError def apply( diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py index 68582f2b6638..288fe8a625e2 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py @@ -29,7 +29,7 @@ logger = init_logger(__name__) -class FlashInferCuteDSLExperts(mk.FusedMoEModularExperts): +class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index af3115a5ec23..7d9a016bf1b2 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -56,7 +56,7 @@ def is_valid_flashinfer_cutlass_fused_moe( return True -class FlashInferExperts(mk.FusedMoEModularExperts): +class FlashInferExperts(mk.FusedMoEExpertsModular): def __init__( self, moe_config: mk.FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index ba3818875ddb..beb97478201f 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -20,7 +20,7 @@ from vllm.v1.engine.utils import current_platform -class FlashInferTrtLlmFp8Experts(mk.FusedMoEMonolithicExperts): +class FlashInferTrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic): """ Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface. """ diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index a0d348137c68..03e335478b95 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -89,7 +89,7 @@ def supports_expert_map(self) -> bool: class FlashInferTrtLlmNvFp4ExpertsModular( - FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEModularExperts + FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular ): """ Modular version of the implementation (just the experts). @@ -194,7 +194,7 @@ def apply( class FlashInferTrtLlmNvFp4ExpertsMonolithic( - FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEMonolithicExperts + FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic ): """ Monolithic version of the kernel (router + experts). diff --git a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py index f8ca531f507b..3e7cb2f55a90 100644 --- a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py @@ -644,7 +644,7 @@ def finalize( ) -class NaiveBatchedExperts(mk.FusedMoEModularExperts): +class NaiveBatchedExperts(mk.FusedMoEExpertsModular): """ A reference MoE expert class that operates on expert batched format, i.e. E x max_num_tokens x K. This is the format that the pplx @@ -876,7 +876,7 @@ def batched_moe_kernel_quantize_input( return A_q, A_q_scale -class BatchedTritonExperts(mk.FusedMoEModularExperts): +class BatchedTritonExperts(mk.FusedMoEExpertsModular): """ A Triton based MoE expert class that operates on expert batched format, i.e. E x max_num_tokens x K. This is the format that the pplx diff --git a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py index dd617be36f88..abf55027a10e 100644 --- a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py @@ -524,7 +524,7 @@ def batched_fused_marlin_moe( return output -class MarlinExpertsBase(mk.FusedMoEModularExperts): +class MarlinExpertsBase(mk.FusedMoEExpertsModular): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 24cff704ae54..ab39d95d8246 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -1703,7 +1703,7 @@ def fused_experts_impl( intermediate_cache3 = cache13[: M * top_k_num * K].view(M, top_k_num, K) # This needs separate memory since it's used concurrently with cache1 - activation_out_dim = mk.FusedMoEModularExperts.adjust_N_for_activation( + activation_out_dim = mk.FusedMoEExpertsModular.adjust_N_for_activation( N, activation ) intermediate_cache2 = torch.empty( @@ -1896,7 +1896,7 @@ def fused_experts_impl( return out_hidden_states -class TritonExperts(mk.FusedMoEModularExperts): +class TritonExperts(mk.FusedMoEExpertsModular): def __init__( self, moe_config: FusedMoEConfig, 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 dce28407f83f..b3e3b0269514 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 @@ -12,7 +12,7 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEModularExperts, + FusedMoEExpertsModular, FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.quantization.base_config import ( @@ -79,7 +79,7 @@ def select_gemm_impl( self, prepare_finalize: FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> FusedMoEModularExperts: + ) -> FusedMoEExpertsModular: # based on the all2all implementation, select the appropriate # gemm implementation raise NotImplementedError( diff --git a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py index 3101290b7f61..1c000e2f0291 100644 --- a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py @@ -244,7 +244,7 @@ def make_routing_data( return routing_data, gather_indx, scatter_indx -class BaseOAITritonExperts(mk.FusedMoEModularExperts): +class BaseOAITritonExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_current_device() -> bool: 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 e0bb4b3353cd..6096f3394949 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -56,18 +56,18 @@ # * FusedMoEPrepareAndFinalize - an abstract base class for preparation of MoE # inputs (e.g. quantization, distribution) and finalization of Moe outputs. # The prepare method must take care of any needed quantization and the -# finalize method, informed by the FusedMoEModularExperts method, +# finalize method, informed by the FusedMoEExpertsModular method, # may apply weights and/or do the final reduction of the output. -# * FusedMoEModularExperts - an abstract base class for the main fused +# * FusedMoEExpertsModular - an abstract base class for the main fused # MoE operation, i.e matmul + act_mul + optionally quant + matmul. -# Some FusedMoEModularExperts implementations may choose to do +# Some FusedMoEExpertsModular implementations may choose to do # the weight application and/or reduction. The class communicates this # to [Finalize] via a TopKWeightAndReduce object. # * FusedMoEModularKernel - an interface class that combines a -# FusedMoEPrepareAndFinalize and a FusedMoEModularExperts to +# FusedMoEPrepareAndFinalize and a FusedMoEExpertsModular to # provide the standard fused MoE kernel interface. # * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen -# by the FusedMoEModularExperts implementation that is passed +# by the FusedMoEExpertsModular implementation that is passed # on to [Finalize]. # # [Quantize-Prepare] and [Finalize] functionality are bundled into a single @@ -156,6 +156,7 @@ def apply( # PrepareResultType is a tuple of: # - quantized + dispatched a. # - quantized + dispatched a1_scales. +# - dispatched router logits. # # See `prepare_monolithic` method below. # @@ -177,7 +178,7 @@ class FusedMoEPrepareAndFinalizeBase(ABC): def post_init_setup(self, fused_experts: "FusedMoEExperts"): """ Initialize FusedMoEPrepareAndFinalize settings that depend on - FusedMoEModularExperts experts object. + FusedMoEExpertsModular experts object. The FusedMoEPrepareAndFinalize implementations that have such dependencies may choose to override this function. """ @@ -257,7 +258,7 @@ def prepare( activations, before quantization + dispatching. - quant_config: Quantization info provided by the fused experts. - defer_input_quant: Runtime parameter indicating whether or not to - defer input quantization to the FusedMoEModularExperts + defer input quantization to the FusedMoEExpertsModular in cases where the compute kernel expects unquantized inputs Returns a tuple of: @@ -304,7 +305,7 @@ def prepare_async( - apply_router_weight_on_input: When True, apply the weights to the activations, before quantization + dispatching. - defer_input_quant: Runtime parameter indicating whether or not to - defer input quantization to the FusedMoEModularExperts + defer input quantization to the FusedMoEExpertsModular in cases where the compute kernel expects unquantized inputs Returns a callback or a hook callback pair that when invoked waits for @@ -415,13 +416,13 @@ def prepare( ) -> PrepareMonolithicResultType: """ Optional method for subclasses compatible with monolithic - FusedMoEModularExperts kernels. + FusedMoEExpertsModular kernels. Perform any quantization (and/or) dispatching needed for this kernel. - a1: The (unquantized) input to the MoE layer. - quant_config: Quantization info provided by the fused experts. - defer_input_quant: Runtime parameter indicating whether or not to - defer input quantization to the FusedMoEModularExperts + defer input quantization to the FusedMoEExpertsModular Returns a tuple of: - quantized + dispatched a. @@ -433,7 +434,7 @@ def prepare( def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor: """ Optional method for subclasses compatible with monolithic - FusedMoEModularExperts kernels. + FusedMoEExpertsModular kernels. Perform any combine plus apply weights and perform a reduction on the fused experts output. @@ -503,7 +504,7 @@ def activation_format() -> FusedMoEActivationFormat: @staticmethod def is_supported_config( - cls: type["FusedMoEModularExperts"], + cls: type["FusedMoEExperts"], moe_config: FusedMoEConfig, weight_key: QuantKey | None, activation_key: QuantKey | None, @@ -684,7 +685,7 @@ def enable_chunking(self): ) -class FusedMoEModularExperts(FusedMoEExperts): +class FusedMoEExpertsModular(FusedMoEExperts): """ An abstract base class for the [Permute-Experts-Unpermute] step described above. @@ -871,7 +872,7 @@ def apply( raise NotImplementedError -class FusedMoEMonolithicExperts(FusedMoEExperts): +class FusedMoEExpertsMonolithic(FusedMoEExperts): """ An abstract base class for the [Permute-Experts-Unpermute] step described above, but with the monolithic interface (accepts router logits @@ -946,16 +947,17 @@ def __init__( if not ( ( isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) - and isinstance(fused_experts, FusedMoEMonolithicExperts) + and isinstance(fused_experts, FusedMoEExpertsMonolithic) ) or ( isinstance(prepare_finalize, FusedMoEPrepareAndFinalize) - and isinstance(fused_experts, FusedMoEModularExperts) + and isinstance(fused_experts, FusedMoEExpertsModular) ) ): raise ValueError( - "prepare_finalize and fused_experts must both be either " - "monolithic or non-monolithic" + "prepare_finalize and fused_experts must both be either monolithic " + f"or non-monolithic but got {prepare_finalize.__class__.__name__} " + "and {fused_experts.__class__.__name__}" ) self._post_init_setup() @@ -974,13 +976,13 @@ def make_mk( fused_experts: FusedMoEExperts, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, - ) -> "FusedMoEMonolithicKernel" | "FusedMoEModularKernel": + ) -> "FusedMoEKernel": """ Factory method to create a FusedMoEKernel instance. """ if isinstance( prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic - ) and isinstance(fused_experts, FusedMoEMonolithicExperts): + ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): return FusedMoEMonolithicKernel( prepare_finalize, fused_experts, @@ -988,7 +990,7 @@ def make_mk( moe_parallel_config, ) elif isinstance(prepare_finalize, FusedMoEPrepareAndFinalize) and isinstance( - fused_experts, FusedMoEModularExperts + fused_experts, FusedMoEExpertsModular ): return FusedMoEModularKernel( prepare_finalize, @@ -998,8 +1000,9 @@ def make_mk( ) else: raise ValueError( - "prepare_finalize and fused_experts must both be either " - "monolithic or non-monolithic" + "prepare_finalize and fused_experts must both be either monolithic " + f"or non-monolithic but got {prepare_finalize.__class__.__name__} " + "and {fused_experts.__class__.__name__}" ) def _post_init_setup(self): @@ -1027,7 +1030,7 @@ def output_is_reduced(self) -> bool: class FusedMoEModularKernel(FusedMoEKernel): """ This class combines a FusedMoEPrepareAndFinalize instance and - a FusedMoEModularExperts to provide an interface that + a FusedMoEExpertsModular to provide an interface that is compatible with the `fused_experts` function in fused_moe.py. It takes care of managing any required scratch space. @@ -1040,7 +1043,7 @@ class FusedMoEModularKernel(FusedMoEKernel): def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalize, - fused_experts: FusedMoEModularExperts, + fused_experts: FusedMoEExpertsModular, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, ): @@ -1048,9 +1051,9 @@ def __init__( raise TypeError( "prepare_finalize must be an instance of FusedMoEPrepareAndFinalize" ) - if not isinstance(fused_experts, FusedMoEModularExperts): + if not isinstance(fused_experts, FusedMoEExpertsModular): raise TypeError( - "fused_experts must be an instance of FusedMoEModularExperts" + "fused_experts must be an instance of FusedMoEExpertsModular" ) super().__init__( @@ -1061,7 +1064,7 @@ def __init__( ) self.prepare_finalize: FusedMoEPrepareAndFinalize = prepare_finalize - self.fused_experts: FusedMoEModularExperts = fused_experts + self.fused_experts: FusedMoEExpertsModular = fused_experts def _chunk_info(self, M: int) -> tuple[int, int]: """ @@ -1105,7 +1108,7 @@ def _allocate_buffers( See `workspace_shapes` for a description of the remainder of arguments. Returns a tuple of (workspace13, workspace2, output) tensors. """ - assert isinstance(self.fused_experts, FusedMoEModularExperts) + assert isinstance(self.fused_experts, FusedMoEExpertsModular) assert M_full > 0 and M_chunk > 0 num_chunks, _ = self._chunk_info(M_full) @@ -1585,7 +1588,7 @@ class FusedMoEMonolithicKernel(FusedMoEKernel): def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic, - fused_experts: FusedMoEMonolithicExperts, + fused_experts: FusedMoEExpertsMonolithic, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, ): @@ -1594,9 +1597,9 @@ def __init__( "prepare_finalize must be an instance of " "FusedMoEPrepareAndFinalizeMonolithic" ) - if not isinstance(fused_experts, FusedMoEMonolithicExperts): + if not isinstance(fused_experts, FusedMoEExpertsMonolithic): raise TypeError( - "fused_experts must be an instance of FusedMoEMonolithicExperts" + "fused_experts must be an instance of FusedMoEExpertsMonolithic" ) super().__init__( @@ -1628,7 +1631,7 @@ def forward( that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) - assert isinstance(self.fused_experts, FusedMoEMonolithicExperts) + assert isinstance(self.fused_experts, FusedMoEExpertsMonolithic) # TODO(rob): add inplace support. a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 86eddc1105da..8ac5661cf756 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -52,7 +52,7 @@ class Fp8MoeBackend(Enum): def backend_to_kernel_cls( backend: Fp8MoeBackend, -) -> type[mk.FusedMoEModularExperts]: +) -> type[mk.FusedMoEExperts]: if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( FlashInferTrtLlmFp8Experts, @@ -132,13 +132,11 @@ def select_fp8_moe_backend( weight_key: QuantKey | None, activation_key: QuantKey | None, allow_vllm_cutlass: bool = False, -) -> tuple[Fp8MoeBackend, type[mk.FusedMoEModularExperts] | None]: +) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts] | None]: """ Select the primary FP8 MoE backend Note: Shape-specific fallbacks may still occur at runtime. """ - k_cls: type[mk.FusedMoEModularExperts] | None = None - if config.is_lora_enabled: return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON) @@ -190,7 +188,7 @@ def _return_or_raise( weight_key: QuantKey | None, activation_key: QuantKey | None, activation_format: mk.FusedMoEActivationFormat, - ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEModularExperts]]: + ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: k_cls = backend_to_kernel_cls(backend) supported, reason = k_cls.is_supported_config( k_cls, config, weight_key, activation_key, activation_format @@ -429,7 +427,7 @@ def make_fp8_moe_quant_config( def make_fp8_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, - experts_cls: type[mk.FusedMoEModularExperts], + experts_cls: type[mk.FusedMoEExperts], fp8_backend: Fp8MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, shared_experts: torch.nn.Module | None = None, @@ -440,6 +438,7 @@ def make_fp8_moe_kernel( quant_config=moe_quant_config, routing_tables=routing_tables, allow_new_interface=True, + use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic), ) assert prepare_finalize is not None diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 497ca5a618d9..7fa75a2efebc 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -65,7 +65,7 @@ def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool: def backend_to_kernel_cls( backend: NvFp4MoeBackend, -) -> list[type[mk.FusedMoEModularExperts]]: +) -> list[type[mk.FusedMoEExperts]]: if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_nvfp4_moe import ( FlashInferTrtLlmNvFp4ExpertsModular, @@ -113,7 +113,7 @@ def select_nvfp4_moe_backend( config: FusedMoEConfig, weight_key: QuantKey | None, activation_key: QuantKey | None, -) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEModularExperts] | None]: +) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts] | None]: """ Select the primary NvFP4 MoE backend Note: Shape-specific fallbacks may still occur at runtime. @@ -166,7 +166,7 @@ def _return_or_raise( weight_key: QuantKey | None, activation_key: QuantKey | None, activation_format: mk.FusedMoEActivationFormat, - ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEModularExperts]]: + ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts]]: for k_cls in backend_to_kernel_cls(backend): supported, reason = k_cls.is_supported_config( k_cls, config, weight_key, activation_key, activation_format @@ -380,6 +380,7 @@ def make_nvfp4_moe_kernel( quant_config=moe_quant_config, routing_tables=routing_tables, allow_new_interface=True, + use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic), ) assert prepare_finalize is not None diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index 75f280d5a45d..a55557342475 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -34,7 +34,7 @@ def make( is_sequence_parallel: bool = False, num_dispatchers: int = 1, use_monolithic: bool = False, - ) -> "MoEPrepareAndFinalizeNaiveEP" | "MoEPrepareAndFinalizeNaiveEPMonolithic": + ) -> "MoEPrepareAndFinalizeNaiveEPBase": cls = ( MoEPrepareAndFinalizeNaiveEPMonolithic if use_monolithic @@ -210,7 +210,7 @@ class MoEPrepareAndFinalizeNaiveEPMonolithic( to the router logits (the MoE kernel runs the router internally). """ - def prepare_monolithic( + def prepare( self, a1: torch.Tensor, router_logits: torch.Tensor, @@ -246,7 +246,6 @@ def finalize( out = get_ep_group().combine( fused_expert_output, is_sequence_parallel=self.is_sequence_parallel ) - assert isinstance(out, torch.Tensor) return out @@ -258,9 +257,7 @@ class MoEPrepareAndFinalizeNoEPBase(mk.FusedMoEPrepareAndFinalizeBase): """ @staticmethod - def make( - use_monolithic: bool, - ) -> "MoEPrepareAndFinalizeNoEP" | "MoEPrepareAndFinalizeNoEPMonolithic": + def make(use_monolithic: bool) -> "MoEPrepareAndFinalizeNoEPBase": return ( MoEPrepareAndFinalizeNoEPMonolithic() if use_monolithic diff --git a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py index 9daba081c431..c6243422bf7f 100644 --- a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py @@ -286,7 +286,7 @@ def rocm_aiter_fused_experts( ) -class AiterExperts(mk.FusedMoEModularExperts): +class AiterExperts(mk.FusedMoEExpertsModular): @property def expects_unquantized_inputs(self) -> bool: return True diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index a8beb8aceacd..69e3da4aa21e 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -10,7 +10,7 @@ class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce): """ - Useful in the case when some FusedMoEModularExperts + Useful in the case when some FusedMoEExpertsModular implementation does not perform weight application and reduction but cannot address the needs of all the compatible PrepareAndFinalize implementations. diff --git a/vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py index 5ed257a7bed9..579aaf033153 100644 --- a/vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py @@ -31,8 +31,8 @@ def __init__( @staticmethod def get_clses() -> tuple[ - type[mk.FusedMoEModularExperts], - type[mk.FusedMoEModularExperts], + type[mk.FusedMoEExpertsModular], + type[mk.FusedMoEExpertsModular], ]: return (CutlassExpertsFp8, TritonExperts) @@ -76,7 +76,7 @@ def _select_experts_impl( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: # Small batch fallback for sm100. if self.is_sm100 and hidden_states.shape[0] <= 8: return self.fallback_experts diff --git a/vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py index 7bf81b59056e..4c0b365daaa9 100644 --- a/vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py @@ -31,8 +31,8 @@ def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig @staticmethod def get_clses() -> tuple[ - type[mk.FusedMoEModularExperts], - type[mk.FusedMoEModularExperts], + type[mk.FusedMoEExpertsModular], + type[mk.FusedMoEExpertsModular], ]: return (DeepGemmExperts, TritonExperts) @@ -78,7 +78,7 @@ def _select_experts_impl( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: if is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2): return self.experts else: diff --git a/vllm/model_executor/layers/fused_moe/trtllm_moe.py b/vllm/model_executor/layers/fused_moe/trtllm_moe.py index 90044bb9c28e..2397134b74da 100644 --- a/vllm/model_executor/layers/fused_moe/trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/trtllm_moe.py @@ -17,7 +17,7 @@ ) -class TrtLlmGenExperts(mk.FusedMoEModularExperts): +class TrtLlmGenExperts(mk.FusedMoEExpertsModular): def __init__( self, moe_config: FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 708e6d1643e9..11d9834232dc 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -23,7 +23,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEActivationFormat, - FusedMoEModularExperts, + FusedMoEExpertsModular, FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( @@ -91,7 +91,7 @@ def select_gemm_impl( self, prepare_finalize: FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> FusedMoEModularExperts: + ) -> FusedMoEExpertsModular: assert self.moe_quant_config is not None if ( prepare_finalize.activation_format diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 463a1118f385..fbcc605bc642 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -19,8 +19,8 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoE, FusedMoEActivationFormat, + FusedMoEExpertsModular, FusedMoEMethodBase, - FusedMoEModularExperts, FusedMoeWeightScaleSupported, UnquantizedFusedMoEMethod, ) @@ -572,7 +572,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -602,8 +602,8 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_kernel is not None - return self.moe_kernel.forward_monolithic( + assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -945,7 +945,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -974,11 +974,8 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.is_monolithic - assert layer.activation == "silu" - - assert self.moe_kernel is not None - return self.moe_kernel.forward_monolithic( + assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -1455,7 +1452,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: assert self.num_bits == 4, "only supporting w4" layer.w13_weight = layer.w13_weight_packed layer.w2_weight = layer.w2_weight_packed @@ -1714,7 +1711,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: if self.moe.is_lora_enabled: assert self.moe_quant_config is not None from vllm.triton_utils import HAS_TRITON @@ -2316,7 +2313,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: assert self.moe_quant_config is not None assert ( prepare_finalize.activation_format == FusedMoEActivationFormat.Standard @@ -2324,7 +2321,7 @@ def select_gemm_impl( from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8 - experts: FusedMoEModularExperts + experts: FusedMoEExpertsModular logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__) experts = CutlassExpertsW4A8Fp8( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index dfc651d6b2ee..11d504b764e6 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -19,8 +19,8 @@ ) from vllm.model_executor.layers.fused_moe import ( FusedMoE, + FusedMoEExpertsModular, FusedMoEMethodBase, - FusedMoEModularExperts, FusedMoEPrepareAndFinalize, FusedMoeWeightScaleSupported, ) @@ -897,7 +897,7 @@ def select_gemm_impl( self, prepare_finalize: FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> FusedMoEModularExperts: + ) -> FusedMoEExpertsModular: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e5413980d463..21456ff4b2a4 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -742,7 +742,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -926,8 +926,8 @@ def apply_monolithic( raise NotImplementedError( "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend." ) - assert self.moe_kernel is not None - return self.moe_kernel.forward_monolithic( + assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -1347,7 +1347,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -1565,7 +1565,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_kernel is not None + assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) return self.moe_kernel.forward_monolithic( x, layer.w13_weight, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index e42050e54da6..f2cc28270c09 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -839,7 +839,7 @@ def select_gemm_impl( self, prepare_finalize: mk.FusedMoEPrepareAndFinalize, layer: torch.nn.Module, - ) -> mk.FusedMoEModularExperts: + ) -> mk.FusedMoEExpertsModular: if ( prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts From 610686b7ca7a303be89e996e35a7481fba37c863 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sat, 31 Jan 2026 18:38:42 -0500 Subject: [PATCH 098/207] change the type names Signed-off-by: Robert Shaw --- docs/design/dbo.md | 2 +- docs/design/fused_moe_modular_kernel.md | 66 +++++++++---------- docs/design/moe_kernel_features.md | 12 ++-- .../moe/modular_kernel_tools/cli_args.py | 2 +- .../moe/modular_kernel_tools/common.py | 2 +- .../moe/modular_kernel_tools/mk_objects.py | 10 +-- .../moe/test_modular_kernel_combinations.py | 4 +- .../layers/fused_moe/__init__.py | 4 +- .../layers/fused_moe/all2all_utils.py | 6 +- .../fused_moe/deepep_ht_prepare_finalize.py | 2 +- .../fused_moe/deepep_ll_prepare_finalize.py | 2 +- .../flashinfer_a2a_prepare_finalize.py | 2 +- .../layers/fused_moe/fused_batched_moe.py | 2 +- .../fused_moe/fused_moe_modular_method.py | 4 +- .../layers/fused_moe/modular_kernel.py | 48 +++++++------- .../layers/fused_moe/mori_prepare_finalize.py | 2 +- .../layers/fused_moe/pplx_prepare_finalize.py | 2 +- .../layers/fused_moe/prepare_finalize.py | 8 +-- .../layers/fused_moe/router/base_router.py | 2 +- .../fused_moe/unquantized_fused_moe_method.py | 6 +- .../compressed_tensors_moe.py | 16 ++--- .../model_executor/layers/quantization/fp8.py | 6 +- .../layers/quantization/modelopt.py | 8 +-- .../layers/quantization/mxfp4.py | 2 +- 24 files changed, 113 insertions(+), 107 deletions(-) diff --git a/docs/design/dbo.md b/docs/design/dbo.md index f2d98ccd063f..43b3ce0bb5a7 100644 --- a/docs/design/dbo.md +++ b/docs/design/dbo.md @@ -81,7 +81,7 @@ The current implementation has all `dbo_yield` and `dbo_maybe_run_recv_hook` cal The `make_ubatch_context` function initializes two `UBatchContexts`, one for each UBatch thread. It takes two CUDA streams, the preexisting `ForwardContexts` and a CPU thread barrier. This function should be used exclusively to instantiate `UBatchContexts`. It will handle all of the event initialization. -The `dbo_register_recv_hook` method registers a callback that can be returned by the `FusedMoEPrepareAndFinalize` class in the other UBatch thread’s `UBatchContext`. The callback will be run when the other thread calls `dbo_maybe_run_recv_hook`. This is typically used to wait on an all-to-all kernel. +The `dbo_register_recv_hook` method registers a callback that can be returned by the `FusedMoEPrepareAndFinalizeModular` class in the other UBatch thread’s `UBatchContext`. The callback will be run when the other thread calls `dbo_maybe_run_recv_hook`. This is typically used to wait on an all-to-all kernel. The `dbo_maybe_run_recv_hook` method runs a callback that’s set by the `dbo_register_recv_hook` function if that callback exists. diff --git a/docs/design/fused_moe_modular_kernel.md b/docs/design/fused_moe_modular_kernel.md index 011422d03f37..75321a22f02a 100644 --- a/docs/design/fused_moe_modular_kernel.md +++ b/docs/design/fused_moe_modular_kernel.md @@ -37,27 +37,27 @@ The rest of the document will focus on the Contiguous / Non-Batched case. Extrap FusedMoEModularKernel splits the FusedMoE operation into 3 parts, 1. TopKWeightAndReduce -2. FusedMoEPrepareAndFinalize +2. FusedMoEPrepareAndFinalizeModular 3. FusedMoEExpertsModular ### TopKWeightAndReduce -The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEExpertsModular` is responsible for the Unpermute and `FusedMoEPrepareAndFinalize` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEExpertsModular`. But some implementations choose to do it `FusedMoEPrepareAndFinalize`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class. +The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEExpertsModular` is responsible for the Unpermute and `FusedMoEPrepareAndFinalizeModular` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEExpertsModular`. But some implementations choose to do it `FusedMoEPrepareAndFinalizeModular`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class. Please find the implementations of TopKWeightAndReduce [here](../../vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py). -`FusedMoEPrepareAndFinalize::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method. +`FusedMoEPrepareAndFinalizeModular::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method. The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEExpertsModular` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. * `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEExpertsModular` implementation does the weight application and reduction itself. -* `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEExpertsModular` implementation needs the `FusedMoEPrepareAndFinalize::finalize()` to do the weight application and reduction. +* `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEExpertsModular` implementation needs the `FusedMoEPrepareAndFinalizeModular::finalize()` to do the weight application and reduction. -### FusedMoEPrepareAndFinalize +### FusedMoEPrepareAndFinalizeModular -The `FusedMoEPrepareAndFinalize` abstract class exposes `prepare`, `prepare_no_receive` and `finalize` functions. -The `prepare` function is responsible for input activation Quantization and All2All Dispatch. If implemented, The `prepare_no_receive` is like `prepare` except it does not wait to receive results from other workers. Instead it returns a "receiver" callback that must be invoked to wait for the final results of worker. It is not required that this method is supported by all `FusedMoEPrepareAndFinalize` classes, but if it is available, it can be used to interleave work with the initial all to all communication, e.g. interleaving shared experts with fused experts. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section) +The `FusedMoEPrepareAndFinalizeModular` abstract class exposes `prepare`, `prepare_no_receive` and `finalize` functions. +The `prepare` function is responsible for input activation Quantization and All2All Dispatch. If implemented, The `prepare_no_receive` is like `prepare` except it does not wait to receive results from other workers. Instead it returns a "receiver" callback that must be invoked to wait for the final results of worker. It is not required that this method is supported by all `FusedMoEPrepareAndFinalizeModular` classes, but if it is available, it can be used to interleave work with the initial all to all communication, e.g. interleaving shared experts with fused experts. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section) -![FusedMoEPrepareAndFinalize Blocks](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png) +![FusedMoEPrepareAndFinalizeModular Blocks](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png) ### FusedMoEExpertsModular @@ -86,19 +86,19 @@ The core FusedMoE implementation performs a series of operations. It would be in #### finalize_weight_and_reduce_impl() It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEExpertsModular::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section. -`FusedMoEExpertsModular::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalize::finalize()` to use. +`FusedMoEExpertsModular::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalizeModular::finalize()` to use. ![FusedMoEExpertsModular Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png) ### FusedMoEModularKernel -`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalize` and `FusedMoEExpertsModular` objects. +`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` objects. `FusedMoEModularKernel` pseudocode/sketch, ```py class FusedMoEModularKernel: def __init__(self, - prepare_finalize: FusedMoEPrepareAndFinalize, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular): self.prepare_finalize = prepare_finalize @@ -128,9 +128,9 @@ class FusedMoEModularKernel: ## How-To -### How To Add a FusedMoEPrepareAndFinalize Type +### How To Add a FusedMoEPrepareAndFinalizeModular Type -Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & Combine implementation / kernel. For example, +Typically a FusedMoEPrepareAndFinalizeModular type is backed by an All2All Dispatch & Combine implementation / kernel. For example, * PplxPrepareAndFinalize type is backed by Pplx All2All kernels, * DeepEPHTPrepareAndFinalize type is backed by DeepEP High-Throughput All2All kernels, and @@ -138,29 +138,29 @@ Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & C #### Step 1: Add an All2All manager -The purpose of the All2All Manager is to set up the All2All kernel implementations. The `FusedMoEPrepareAndFinalize` implementations typically fetch a kernel-implementation "handle" from the All2All Manager to invoke the Dispatch and Combine functions. Please look at the All2All Manager implementations [here](../../vllm/distributed/device_communicators/all2all.py). +The purpose of the All2All Manager is to set up the All2All kernel implementations. The `FusedMoEPrepareAndFinalizeModular` implementations typically fetch a kernel-implementation "handle" from the All2All Manager to invoke the Dispatch and Combine functions. Please look at the All2All Manager implementations [here](../../vllm/distributed/device_communicators/all2all.py). -#### Step 2: Add a FusedMoEPrepareAndFinalize Type +#### Step 2: Add a FusedMoEPrepareAndFinalizeModular Type -This section describes the significance of the various functions exposed by the `FusedMoEPrepareAndFinalize` abstract class. +This section describes the significance of the various functions exposed by the `FusedMoEPrepareAndFinalizeModular` abstract class. -`FusedMoEPrepareAndFinalize::prepare()`: The prepare method implements the Quantization and All2All Dispatch. Typically the Dispatch function from the relevant All2All Manager is invoked. +`FusedMoEPrepareAndFinalizeModular::prepare()`: The prepare method implements the Quantization and All2All Dispatch. Typically the Dispatch function from the relevant All2All Manager is invoked. -`FusedMoEPrepareAndFinalize::has_prepare_no_receive()`: Indicates whether or not this subclass implements `prepare_no_receive`. Defaults to False. +`FusedMoEPrepareAndFinalizeModular::has_prepare_no_receive()`: Indicates whether or not this subclass implements `prepare_no_receive`. Defaults to False. -`FusedMoEPrepareAndFinalize::prepare_no_receive()`: The prepare_no_receive method implements the Quantization and All2All Dispatch. It does not wait for the result of the dispatch operation but instead returns a thunk that can be invoked to wait for the final results. Typically the Dispatch function from the relevant All2All Manager is invoked. +`FusedMoEPrepareAndFinalizeModular::prepare_no_receive()`: The prepare_no_receive method implements the Quantization and All2All Dispatch. It does not wait for the result of the dispatch operation but instead returns a thunk that can be invoked to wait for the final results. Typically the Dispatch function from the relevant All2All Manager is invoked. -`FusedMoEPrepareAndFinalize::finalize()`: Maybe perform TopK Weight Application and Reduction and All2All Combine. Typically the Combine function from the relevant All2AllManager is invoked. +`FusedMoEPrepareAndFinalizeModular::finalize()`: Maybe perform TopK Weight Application and Reduction and All2All Combine. Typically the Combine function from the relevant All2AllManager is invoked. -`FusedMoEPrepareAndFinalize::activation_format()`: Return `FusedMoEActivationFormat.BatchedExperts` if the output of the prepare method (i.e. the All2All dispatch) is Batched. Return `FusedMoEActivationFormat.Standard` otherwise. +`FusedMoEPrepareAndFinalizeModular::activation_format()`: Return `FusedMoEActivationFormat.BatchedExperts` if the output of the prepare method (i.e. the All2All dispatch) is Batched. Return `FusedMoEActivationFormat.Standard` otherwise. -`FusedMoEPrepareAndFinalize::topk_indices_dtype()`: Data type of the TopK ids. Some All2All kernels have strict requirements pertaining to the data type of the TopK ids. This requirement is passed on to the `FusedMoe::select_experts` function so it could be respected. If there are no strict requirements return None. +`FusedMoEPrepareAndFinalizeModular::topk_indices_dtype()`: Data type of the TopK ids. Some All2All kernels have strict requirements pertaining to the data type of the TopK ids. This requirement is passed on to the `FusedMoe::select_experts` function so it could be respected. If there are no strict requirements return None. -`FusedMoEPrepareAndFinalize::max_num_tokens_per_rank()`: This is the maximum number of tokens that would be submitted to the All2All Dispatch at once. +`FusedMoEPrepareAndFinalizeModular::max_num_tokens_per_rank()`: This is the maximum number of tokens that would be submitted to the All2All Dispatch at once. -`FusedMoEPrepareAndFinalize::num_dispatchers()`: Total number of dispatching units. This value determines the size of the Dispatch output. The Dispatch output is of shape (num_local_experts, max_num_tokens, K). Here max_num_tokens = num_dispatchers() * max_num_tokens_per_rank(). +`FusedMoEPrepareAndFinalizeModular::num_dispatchers()`: Total number of dispatching units. This value determines the size of the Dispatch output. The Dispatch output is of shape (num_local_experts, max_num_tokens, K). Here max_num_tokens = num_dispatchers() * max_num_tokens_per_rank(). -We suggest picking an already existing `FusedMoEPrepareAndFinalize` implementation that matches your All2All implementation closely and using it as a reference. +We suggest picking an already existing `FusedMoEPrepareAndFinalizeModular` implementation that matches your All2All implementation closely and using it as a reference. ### How To Add a FusedMoEExpertsModular Type @@ -187,7 +187,7 @@ implementations that input `FusedMoEActivationFormat.Standard` support chunking #### maybe_make_prepare_finalize -The `maybe_make_prepare_finalize` method is responsible for constructing an instance of `FusedMoEPrepareAndFinalize` when appropriate based on the current all2all backend, e.g. when EP + DP is enabled. The base class method currently constructs all the `FusedMoEPrepareAndFinalize` objects for the EP+DP case. Derived classes can override this method to construct prepare/finalize objects for different scenarios, e.g. `ModelOptNvFp4FusedMoE` can construct a `FlashInferCutlassMoEPrepareAndFinalize` for the EP+TP case. +The `maybe_make_prepare_finalize` method is responsible for constructing an instance of `FusedMoEPrepareAndFinalizeModular` when appropriate based on the current all2all backend, e.g. when EP + DP is enabled. The base class method currently constructs all the `FusedMoEPrepareAndFinalizeModular` objects for the EP+DP case. Derived classes can override this method to construct prepare/finalize objects for different scenarios, e.g. `ModelOptNvFp4FusedMoE` can construct a `FlashInferCutlassMoEPrepareAndFinalize` for the EP+TP case. Please refer to the implementations in, * `ModelOptNvFp4FusedMoE` @@ -206,7 +206,7 @@ derived classes. #### init_prepare_finalize -Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalize` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEModularKernel` object +Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalizeModular` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEModularKernel` object Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188). **Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used. @@ -215,9 +215,9 @@ Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vl We have `FusedMoEModularKernel` unit tests at [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py). -The unit test iterates through all combinations of `FusedMoEPrepareAndFinalize` and `FusedMoEPremuteExpertsUnpermute` types and if they are +The unit test iterates through all combinations of `FusedMoEPrepareAndFinalizeModular` and `FusedMoEPremuteExpertsUnpermute` types and if they are compatible, runs some correctness tests. -If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEExpertsModular` implementations, +If you are adding some `FusedMoEPrepareAndFinalizeModular` / `FusedMoEExpertsModular` implementations, 1. Add the implementation type to `MK_ALL_PREPARE_FINALIZE_TYPES` and `MK_FUSED_EXPERT_TYPES` in [mk_objects.py](../../tests/kernels/moe/modular_kernel_tools/mk_objects.py) respectively. 2. Update `Config::is_batched_prepare_finalize()`, `Config::is_batched_fused_experts()`, `Config::is_standard_fused_experts()`, @@ -226,21 +226,21 @@ If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEExpertsModular` i Doing this will add the new implementation to the test suite. -### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEExpertsModular` Compatibility +### How To Check `FusedMoEPrepareAndFinalizeModular` & `FusedMoEExpertsModular` Compatibility The unit test file [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py) can also be executed as a standalone script. Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` -As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEExpertsModular` compatibility. When invoked +As a side effect, this script can be used to test `FusedMoEPrepareAndFinalizeModular` & `FusedMoEExpertsModular` compatibility. When invoked with incompatible types, the script will error. ### How To Profile Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py) The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible -`FusedMoEPrepareAndFinalize` and `FusedMoEExpertsModular` types. +`FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` types. Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` -## FusedMoEPrepareAndFinalize Implementations +## FusedMoEPrepareAndFinalizeModular Implementations See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-modular-all2all-backends) for a list of all the available modular prepare and finalize subclasses. diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index df54255e45b6..e56b2f915284 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -4,17 +4,17 @@ The purpose of this document is to provide an overview of the various MoE kernel ## Fused MoE Modular All2All backends -There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalize` subclasses provide an interface for each all2all backend. +There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalizeModular` subclasses provide an interface for each all2all backend. The following table describes the relevant features of each backend, i.e. activation format, supported quantization schemes and async support. -The output activation format (standard or batched) corresponds to the output of the prepare step of the `FusedMoEPrepareAndFinalize` subclass, and the finalize step requires the same format. All the backend `prepare` methods expect activations in the standard format and all the `finalize` methods return activations in standard format. More details on the formats can be found in the [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) document. +The output activation format (standard or batched) corresponds to the output of the prepare step of the `FusedMoEPrepareAndFinalizeModular` subclass, and the finalize step requires the same format. All the backend `prepare` methods expect activations in the standard format and all the `finalize` methods return activations in standard format. More details on the formats can be found in the [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) document. -The quantization types and formats enumerate which quantization schemes are supported by each `FusedMoEPrepareAndFinalize` class. The quantization can happen before or after the dispatch based on the format the all2all backend supports, e.g. deepep_high_throughput supports only block-quantized fp8 format. Any other format will result in dispatching in higher precision and quantizing afterwards. The output of the prepare step for each backend is the quantized type. The finalize step generally requires the same input type as the original activations, e.g. if the original input is bfloat16 and the quantization scheme is fp8 with per-tensor scales, `prepare` will return fp8/per-tensor scale activations and `finalize` will take bfloat16 activations. See the diagrams in [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) for more details on the types and formats of activations at each step of the MoE process. If no quantization type is specified, the kernel operates on float16 and/or bfloat16. +The quantization types and formats enumerate which quantization schemes are supported by each `FusedMoEPrepareAndFinalizeModular` class. The quantization can happen before or after the dispatch based on the format the all2all backend supports, e.g. deepep_high_throughput supports only block-quantized fp8 format. Any other format will result in dispatching in higher precision and quantizing afterwards. The output of the prepare step for each backend is the quantized type. The finalize step generally requires the same input type as the original activations, e.g. if the original input is bfloat16 and the quantization scheme is fp8 with per-tensor scales, `prepare` will return fp8/per-tensor scale activations and `finalize` will take bfloat16 activations. See the diagrams in [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) for more details on the types and formats of activations at each step of the MoE process. If no quantization type is specified, the kernel operates on float16 and/or bfloat16. Async backends support the use of DBO (Dual Batch Overlap) and shared expert overlap (where shared experts are computed during the combine step). -Certain models require the topk weights to be applied to the input activations rather than the output activations when topk==1, e.g. Llama. For modular kernels, this feature is supported by the `FusedMoEPrepareAndFinalize` subclass. For non-modular kernels, it is up to the experts function to deal with this flag. +Certain models require the topk weights to be applied to the input activations rather than the output activations when topk==1, e.g. Llama. For modular kernels, this feature is supported by the `FusedMoEPrepareAndFinalizeModular` subclass. For non-modular kernels, it is up to the experts function to deal with this flag. Unless otherwise specified, backends are controlled via the `--all2all-backend` command-line argument (or the `all2all_backend` parameter in `ParallelConfig`). All backends except `flashinfer` only work with EP+DP or EP+TP. `Flashinfer` can work with EP or DP without EP. @@ -78,7 +78,7 @@ As with the backends, some experts support applying topk weights on the input ac Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEExpertsModular`. -To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats. +To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats. | Kernel | Input act. format | Quant. types | Quant. format | Activation function | Apply Weight On Input | Modular | Source | |--------|-------------------|--------------|---------------|---------------------|-----------------------|---------|--------| @@ -107,7 +107,7 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels The following table shows "families" of modular kernels that are intended to work together. There are some combinations which may work but have not yet been tested, e.g. flashinfer with other fp8 experts. Note that the "naive" backend will work with any non-modular experts. -| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEExpertsModular` subclasses | +| backend | `FusedMoEPrepareAndFinalizeModular` subclasses | `FusedMoEExpertsModular` subclasses | |---------|-----------------------------------------|----------------------------------------------| | deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,
`TritonExperts`,
`TritonOrDeepGemmExperts`,
`CutlassExpertsFp8`,
`MarlinExperts` | | deepep_low_latency,
pplx | `DeepEPLLPrepareAndFinalize`,
`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,
`BatchedTritonExperts`,
`CutlassBatchedExpertsFp8`,
`BatchedMarlinExperts` | diff --git a/tests/kernels/moe/modular_kernel_tools/cli_args.py b/tests/kernels/moe/modular_kernel_tools/cli_args.py index 28be65127780..375dfa748956 100644 --- a/tests/kernels/moe/modular_kernel_tools/cli_args.py +++ b/tests/kernels/moe/modular_kernel_tools/cli_args.py @@ -17,7 +17,7 @@ def make_config_arg_parser(description: str): - def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalize: + def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalizeModular: for pf in MK_ALL_PREPARE_FINALIZE_TYPES: if pf.__name__ == s: return pf diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 50174695b8c3..ee7ed731c42f 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -65,7 +65,7 @@ class Config: dtype: torch.dtype quant_config: TestMoEQuantConfig | None - prepare_finalize_type: mk.FusedMoEPrepareAndFinalize + prepare_finalize_type: mk.FusedMoEPrepareAndFinalizeModular fused_experts_type: mk.FusedMoEExpertsModular fused_moe_chunk_size: int | None diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 8abcf656de80..c31d2ce38c61 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -72,11 +72,13 @@ class ExpertInfo: needs_aiter: bool = False -PREPARE_FINALIZE_INFO: dict[mk.FusedMoEPrepareAndFinalize, PrepareFinalizeInfo] = {} +PREPARE_FINALIZE_INFO: dict[ + mk.FusedMoEPrepareAndFinalizeModular, PrepareFinalizeInfo +] = {} EXPERT_INFO: dict[mk.FusedMoEExpertsModular, ExpertInfo] = {} -MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] -MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] -MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = [] +MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = [] +MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = [] +MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = [] MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEExpertsModular] = [] standard_format = mk.FusedMoEActivationFormat.Standard diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 8c72e7c8df75..d0d238b4e912 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -258,7 +258,7 @@ def test_modular_kernel_combinations_multigpu( e: int, dtype: torch.dtype, quant_config: TestMoEQuantConfig | None, - prepare_finalize_type: mk.FusedMoEPrepareAndFinalize, + prepare_finalize_type: mk.FusedMoEPrepareAndFinalizeModular, fused_experts_type: mk.FusedMoEExpertsModular, chunk_size: int | None, world_size: int, @@ -300,7 +300,7 @@ def test_modular_kernel_combinations_singlegpu( e: int, dtype: torch.dtype, quant_config: TestMoEQuantConfig | None, - prepare_finalize_type: mk.FusedMoEPrepareAndFinalize, + prepare_finalize_type: mk.FusedMoEPrepareAndFinalizeModular, fused_experts_type: mk.FusedMoEExpertsModular, chunk_size: int | None, world_size: int, diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 96da5e4f9069..556f0ee46255 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -18,7 +18,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEActivationFormat, FusedMoEExpertsModular, - FusedMoEPrepareAndFinalize, + FusedMoEPrepareAndFinalizeModular, ) from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, @@ -58,7 +58,7 @@ def get_config() -> dict[str, Any] | None: "FusedMoeWeightScaleSupported", "FusedMoEExpertsModular", "FusedMoEActivationFormat", - "FusedMoEPrepareAndFinalize", + "FusedMoEPrepareAndFinalizeModular", "RoutingMethodType", "SharedFusedMoE", "ZeroExpertFusedMoE", diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 172acfbc394a..205edfee3e30 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -17,7 +17,7 @@ FlashInferA2APrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEPrepareAndFinalizeBase, + FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNaiveEPBase, @@ -82,7 +82,7 @@ def maybe_make_prepare_finalize( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, allow_new_interface: bool = False, use_monolithic: bool = False, -) -> FusedMoEPrepareAndFinalizeBase | None: +) -> FusedMoEPrepareAndFinalize | None: # NOTE(rob): we are migrating each quant_method to hold the MK # in all cases. The allow_new_interface=False flag allow us to fall # back to the old method for methods that have not yet been migrated. @@ -119,7 +119,7 @@ def maybe_make_prepare_finalize( all2all_manager = get_ep_group().device_communicator.all2all_manager assert all2all_manager is not None - prepare_finalize: FusedMoEPrepareAndFinalizeBase | None = None + prepare_finalize: FusedMoEPrepareAndFinalize | None = None if moe.use_pplx_kernels: assert quant_config is not None diff --git a/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py index c4a227223703..63312557d85d 100644 --- a/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py @@ -25,7 +25,7 @@ ) -class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): +class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): """ Prepare/Finalize using DeepEP High-Throughput kernels. """ diff --git a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py index dd6276afeda3..41b5b309b921 100644 --- a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py @@ -49,7 +49,7 @@ def dequant_fp8( return (expert_x_fp32 * expert_x_scales).view(expert_x_fp8.size()) -class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): +class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): """ Prepare/Finalize using DeepEP low-latency kernels. """ diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py index 2857e51a3a72..465d0ae8f2c4 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py @@ -18,7 +18,7 @@ def get_local_sizes(): return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank() -class FlashInferA2APrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): +class FlashInferA2APrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): """Base class for FlashInfer MoE prepare and finalize operations.""" def __init__( diff --git a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py index 3e7cb2f55a90..d36cdc424a62 100644 --- a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py @@ -488,7 +488,7 @@ def invoke_moe_batched_triton_kernel( ) -class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): +class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): """ A reference prepare/finalize class that reorganizes the tokens into expert batched format, i.e. E x max_num_tokens x K. This is the format diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 69bb47be0117..093b9d2e2d4e 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -14,7 +14,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEModularKernel, - FusedMoEPrepareAndFinalize, + FusedMoEPrepareAndFinalizeModular, ) logger = init_logger(__name__) @@ -44,7 +44,7 @@ def __init__( def make( moe_layer: torch.nn.Module, old_quant_method: FusedMoEMethodBase, - prepare_finalize: FusedMoEPrepareAndFinalize, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, shared_experts: torch.nn.Module | None, ) -> "FusedMoEModularMethod": return FusedMoEModularMethod( diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 6096f3394949..b3bcb344d4fa 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -53,7 +53,7 @@ # MoE kernel implementations. # # The following main classes are defined: -# * FusedMoEPrepareAndFinalize - an abstract base class for preparation of MoE +# * FusedMoEPrepareAndFinalizeModular - an abstract base class for preparation of MoE # inputs (e.g. quantization, distribution) and finalization of Moe outputs. # The prepare method must take care of any needed quantization and the # finalize method, informed by the FusedMoEExpertsModular method, @@ -64,14 +64,14 @@ # the weight application and/or reduction. The class communicates this # to [Finalize] via a TopKWeightAndReduce object. # * FusedMoEModularKernel - an interface class that combines a -# FusedMoEPrepareAndFinalize and a FusedMoEExpertsModular to +# FusedMoEPrepareAndFinalizeModular and a FusedMoEExpertsModular to # provide the standard fused MoE kernel interface. # * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen # by the FusedMoEExpertsModular implementation that is passed # on to [Finalize]. # # [Quantize-Prepare] and [Finalize] functionality are bundled into a single -# class `FusedMoEPrepareAndFinalize` since they could use collective +# class `FusedMoEPrepareAndFinalizeModular` since they could use collective # communication mechanisms that need to be consistent. # @@ -169,17 +169,17 @@ def apply( ReceiverType = Callable[[], PrepareResultType] -class FusedMoEPrepareAndFinalizeBase(ABC): +class FusedMoEPrepareAndFinalize(ABC): """ - An abstract base class for FusedMoEPrepareAndFinalize (Modular) + An abstract base class for FusedMoEPrepareAndFinalizeModular (Modular) and FusedMoEPrepareAndFinalizeMonolithic (Monolithic) implementations. """ def post_init_setup(self, fused_experts: "FusedMoEExperts"): """ - Initialize FusedMoEPrepareAndFinalize settings that depend on + Initialize FusedMoEPrepareAndFinalizeModular settings that depend on FusedMoEExpertsModular experts object. - The FusedMoEPrepareAndFinalize implementations that have such + The FusedMoEPrepareAndFinalizeModular implementations that have such dependencies may choose to override this function. """ return @@ -228,7 +228,7 @@ def output_is_reduced(self) -> bool: # TODO: pass FusedMoEParallelConfig in as ctor parameter? -class FusedMoEPrepareAndFinalize(FusedMoEPrepareAndFinalizeBase): +class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): """ An abstract base class for the [Quantize-Prepare] and [Finalize] steps described above. @@ -399,7 +399,7 @@ def finalize_async( raise NotImplementedError -class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalizeBase): +class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize): """ An abstract base class for the [Quantize-Prepare] and [Finalize] steps described above but for the monolithic interface (accepts router logits @@ -922,7 +922,7 @@ def _slice_scales( class FusedMoEKernel(torch.nn.Module): def __init__( self, - prepare_finalize: FusedMoEPrepareAndFinalizeBase, + prepare_finalize: FusedMoEPrepareAndFinalize, fused_experts: FusedMoEExperts, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, @@ -950,7 +950,7 @@ def __init__( and isinstance(fused_experts, FusedMoEExpertsMonolithic) ) or ( - isinstance(prepare_finalize, FusedMoEPrepareAndFinalize) + isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular) and isinstance(fused_experts, FusedMoEExpertsModular) ) ): @@ -972,7 +972,7 @@ def __init__( @staticmethod def make_mk( - prepare_finalize: FusedMoEPrepareAndFinalizeBase, + prepare_finalize: FusedMoEPrepareAndFinalize, fused_experts: FusedMoEExperts, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, @@ -989,9 +989,9 @@ def make_mk( shared_experts, moe_parallel_config, ) - elif isinstance(prepare_finalize, FusedMoEPrepareAndFinalize) and isinstance( - fused_experts, FusedMoEExpertsModular - ): + elif isinstance( + prepare_finalize, FusedMoEPrepareAndFinalizeModular + ) and isinstance(fused_experts, FusedMoEExpertsModular): return FusedMoEModularKernel( prepare_finalize, fused_experts, @@ -1029,7 +1029,7 @@ def output_is_reduced(self) -> bool: @final class FusedMoEModularKernel(FusedMoEKernel): """ - This class combines a FusedMoEPrepareAndFinalize instance and + This class combines a FusedMoEPrepareAndFinalizeModular instance and a FusedMoEExpertsModular to provide an interface that is compatible with the `fused_experts` function in fused_moe.py. @@ -1042,18 +1042,22 @@ class FusedMoEModularKernel(FusedMoEKernel): def __init__( self, - prepare_finalize: FusedMoEPrepareAndFinalize, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, ): - if not isinstance(prepare_finalize, FusedMoEPrepareAndFinalize): + if not isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular): raise TypeError( - "prepare_finalize must be an instance of FusedMoEPrepareAndFinalize" + "prepare_finalize must be an instance of " + "FusedMoEPrepareAndFinalizeModular, but got " + f"{prepare_finalize.__class__.__name__}." ) if not isinstance(fused_experts, FusedMoEExpertsModular): raise TypeError( - "fused_experts must be an instance of FusedMoEExpertsModular" + "fused_experts must be an instance of " + "FusedMoEExpertsModular, but got " + f"{fused_experts.__class__.__name__}." ) super().__init__( @@ -1063,7 +1067,7 @@ def __init__( moe_parallel_config, ) - self.prepare_finalize: FusedMoEPrepareAndFinalize = prepare_finalize + self.prepare_finalize: FusedMoEPrepareAndFinalizeModular = prepare_finalize self.fused_experts: FusedMoEExpertsModular = fused_experts def _chunk_info(self, M: int) -> tuple[int, int]: @@ -1264,7 +1268,7 @@ def _prepare( The _prepare method is a wrapper around self.prepare_finalize.prepare that handles DBO and async. """ - assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalize) + assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeModular) if not self.prepare_finalize.supports_async(): # We shouldn't be running an a2a kernel that doesn't diff --git a/vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py index dc0f32dc1992..164605dde3c0 100644 --- a/vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py @@ -12,7 +12,7 @@ logger = init_logger(__name__) -class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): +class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): """ Prepare/Finalize using MoRI kernels. """ diff --git a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py index 78b941498062..5fcdf06436b6 100644 --- a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py @@ -62,7 +62,7 @@ def pplx_hidden_dim_scale_bytes( ) -class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): +class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): def __init__( self, a2a: pplx.AllToAll, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index a55557342475..f25da1f298b8 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -14,7 +14,7 @@ from vllm.utils.flashinfer import nvfp4_block_scale_interleave -class MoEPrepareAndFinalizeNaiveEPBase(mk.FusedMoEPrepareAndFinalizeBase): +class MoEPrepareAndFinalizeNaiveEPBase(mk.FusedMoEPrepareAndFinalize): """ Base class for Naive Prepare/Finalize for Dp/Ep with two subclasses: * Modular Case @@ -125,7 +125,7 @@ def _unwrap_scale_and_prepare_for_moe( class MoEPrepareAndFinalizeNaiveEP( - MoEPrepareAndFinalizeNaiveEPBase, mk.FusedMoEPrepareAndFinalize + MoEPrepareAndFinalizeNaiveEPBase, mk.FusedMoEPrepareAndFinalizeModular ): """ Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. @@ -249,7 +249,7 @@ def finalize( return out -class MoEPrepareAndFinalizeNoEPBase(mk.FusedMoEPrepareAndFinalizeBase): +class MoEPrepareAndFinalizeNoEPBase(mk.FusedMoEPrepareAndFinalize): """ Base class for TP case Prepare/Finalize. * prepare: applies input quantization @@ -309,7 +309,7 @@ def _quantize_input( class MoEPrepareAndFinalizeNoEP( - mk.FusedMoEPrepareAndFinalize, MoEPrepareAndFinalizeNoEPBase + mk.FusedMoEPrepareAndFinalizeModular, MoEPrepareAndFinalizeNoEPBase ): def prepare( self, diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index 9969818abfd6..a1ef9df4e4a6 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -64,7 +64,7 @@ def eplb_map_to_physical_and_record( # TODO(bowen): When using `FusedMoEModularKernel`, this # can be done in a more unified way, since - # `FusedMoEPrepareAndFinalize` will return the expert + # `FusedMoEPrepareAndFinalizeModular` will return the expert # token count, in some cases directly from the kernel. # However, now there are many code paths not using # the modular kernel, e.g. calling `fused_experts`, diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 11d9834232dc..c8168c091e44 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -24,7 +24,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEActivationFormat, FusedMoEExpertsModular, - FusedMoEPrepareAndFinalize, + FusedMoEPrepareAndFinalizeModular, ) from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, @@ -81,7 +81,7 @@ def allow_inplace(self) -> bool: def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> FusedMoEPrepareAndFinalize | None: + ) -> FusedMoEPrepareAndFinalizeModular | None: if self.unquantized_backend == UnquantizedMoeBackend.AITER: return None else: @@ -89,7 +89,7 @@ def maybe_make_prepare_finalize( def select_gemm_impl( self, - prepare_finalize: FusedMoEPrepareAndFinalize, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> FusedMoEExpertsModular: assert self.moe_quant_config is not None diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index fbcc605bc642..1c4db3b8cc1a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -562,7 +562,7 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalize | None: + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -570,7 +570,7 @@ def maybe_make_prepare_finalize( def select_gemm_impl( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> mk.FusedMoEExpertsModular: raise ValueError( @@ -935,7 +935,7 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalize | None: + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -943,7 +943,7 @@ def maybe_make_prepare_finalize( def select_gemm_impl( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> mk.FusedMoEExpertsModular: raise ValueError( @@ -1450,7 +1450,7 @@ def get_fused_moe_quant_config( def select_gemm_impl( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> mk.FusedMoEExpertsModular: assert self.num_bits == 4, "only supporting w4" @@ -1709,7 +1709,7 @@ def get_fused_moe_quant_config( def select_gemm_impl( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> mk.FusedMoEExpertsModular: if self.moe.is_lora_enabled: @@ -2290,7 +2290,7 @@ def process_weights_after_loading(self, layer): def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalize | None: + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: return super().maybe_make_prepare_finalize(routing_tables) def get_fused_moe_quant_config( @@ -2311,7 +2311,7 @@ def get_fused_moe_quant_config( def select_gemm_impl( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> mk.FusedMoEExpertsModular: assert self.moe_quant_config is not None diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 11d504b764e6..4f5cba747385 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -21,7 +21,7 @@ FusedMoE, FusedMoEExpertsModular, FusedMoEMethodBase, - FusedMoEPrepareAndFinalize, + FusedMoEPrepareAndFinalizeModular, FusedMoeWeightScaleSupported, ) from vllm.model_executor.layers.fused_moe.config import ( @@ -887,7 +887,7 @@ def process_weights_after_loading(self, layer: Module) -> None: def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalize | None: + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -895,7 +895,7 @@ def maybe_make_prepare_finalize( def select_gemm_impl( self, - prepare_finalize: FusedMoEPrepareAndFinalize, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> FusedMoEExpertsModular: raise ValueError( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 21456ff4b2a4..d09a1c03ae09 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -732,7 +732,7 @@ def __init__( def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalize | None: + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -740,7 +740,7 @@ def maybe_make_prepare_finalize( def select_gemm_impl( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> mk.FusedMoEExpertsModular: raise ValueError( @@ -1337,7 +1337,7 @@ def __init__( def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalize | None: + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " "logic. This function should not be called." @@ -1345,7 +1345,7 @@ def maybe_make_prepare_finalize( def select_gemm_impl( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> mk.FusedMoEExpertsModular: raise ValueError( diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index f2cc28270c09..1ffe94a5aaeb 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -837,7 +837,7 @@ def get_fused_moe_quant_config( def select_gemm_impl( self, - prepare_finalize: mk.FusedMoEPrepareAndFinalize, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> mk.FusedMoEExpertsModular: if ( From e2b4f86b30c25dd367e8c812880910962d2e6a42 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sat, 31 Jan 2026 18:40:16 -0500 Subject: [PATCH 099/207] nits Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index eae09359d957..c8d76a623b87 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoEPMonolithic, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, @@ -211,7 +212,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( ) kernel = mk.FusedMoEModularKernel.make_mk( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoEPMonolithic(), FlashInferTrtLlmFp8Experts( moe_config=td.layer.moe, quant_config=quant_config, From f709290f1c51cca415d095f0d270cfa8bc663924 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sat, 31 Jan 2026 18:50:56 -0500 Subject: [PATCH 100/207] improve tping Signed-off-by: Robert Shaw --- tests/kernels/moe/modular_kernel_tools/common.py | 4 ++-- tests/kernels/moe/test_modular_kernel_combinations.py | 8 ++++---- .../layers/fused_moe/fused_moe_method_base.py | 10 ++++++---- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index ee7ed731c42f..a83e85942d02 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -65,8 +65,8 @@ class Config: dtype: torch.dtype quant_config: TestMoEQuantConfig | None - prepare_finalize_type: mk.FusedMoEPrepareAndFinalizeModular - fused_experts_type: mk.FusedMoEExpertsModular + prepare_finalize_type: mk.FusedMoEPrepareAndFinalize + fused_experts_type: mk.FusedMoEExperts fused_moe_chunk_size: int | None world_size: int diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index d0d238b4e912..7bdbaf0c8a52 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -258,8 +258,8 @@ def test_modular_kernel_combinations_multigpu( e: int, dtype: torch.dtype, quant_config: TestMoEQuantConfig | None, - prepare_finalize_type: mk.FusedMoEPrepareAndFinalizeModular, - fused_experts_type: mk.FusedMoEExpertsModular, + prepare_finalize_type: mk.FusedMoEPrepareAndFinalize, + fused_experts_type: mk.FusedMoEExperts, chunk_size: int | None, world_size: int, pytestconfig, @@ -300,8 +300,8 @@ def test_modular_kernel_combinations_singlegpu( e: int, dtype: torch.dtype, quant_config: TestMoEQuantConfig | None, - prepare_finalize_type: mk.FusedMoEPrepareAndFinalizeModular, - fused_experts_type: mk.FusedMoEExpertsModular, + prepare_finalize_type: mk.FusedMoEPrepareAndFinalize, + fused_experts_type: mk.FusedMoEExperts, chunk_size: int | None, world_size: int, pytestconfig, 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 b3e3b0269514..6384ab2f5fc1 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,7 +13,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEExpertsModular, - FusedMoEPrepareAndFinalize, + FusedMoEPrepareAndFinalizeModular, ) from vllm.model_executor.layers.quantization.base_config import ( QuantizeMethodBase, @@ -68,16 +68,18 @@ def uses_weight_scale_2_pattern(self) -> bool: def maybe_make_prepare_finalize( self, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> FusedMoEPrepareAndFinalize | None: + ) -> FusedMoEPrepareAndFinalizeModular | None: from .all2all_utils import maybe_make_prepare_finalize - return maybe_make_prepare_finalize( + pf = maybe_make_prepare_finalize( self.moe, self.moe_quant_config, routing_tables ) + assert pf is None or isinstance(pf, FusedMoEPrepareAndFinalizeModular) + return pf def select_gemm_impl( self, - prepare_finalize: FusedMoEPrepareAndFinalize, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, layer: torch.nn.Module, ) -> FusedMoEExpertsModular: # based on the all2all implementation, select the appropriate From b305b1bf57c1dff0f191cd4b596c26f79bba037d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sat, 31 Jan 2026 18:53:48 -0500 Subject: [PATCH 101/207] updated comments Signed-off-by: Robert Shaw --- .../layers/fused_moe/modular_kernel.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index b3bcb344d4fa..4d1e19f93581 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -171,8 +171,12 @@ def apply( class FusedMoEPrepareAndFinalize(ABC): """ - An abstract base class for FusedMoEPrepareAndFinalizeModular (Modular) - and FusedMoEPrepareAndFinalizeMonolithic (Monolithic) implementations. + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above. + + There are two variants of this class: + * FusedMoEPrepareAndFinalizeModular - this operates on topk ids and weights + * FusedMoEPrepareAndFinalizeMonolithic - the operates on router_logits """ def post_init_setup(self, fused_experts: "FusedMoEExperts"): @@ -231,7 +235,7 @@ def output_is_reduced(self) -> bool: class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): """ An abstract base class for the [Quantize-Prepare] and [Finalize] steps - described above. + described above for the Modular case. """ @abstractmethod @@ -402,8 +406,7 @@ def finalize_async( class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize): """ An abstract base class for the [Quantize-Prepare] and [Finalize] steps - described above but for the monolithic interface (accepts router logits - rather than topk ids and weights). + described above for the monolithic case. """ @abstractmethod From 07eff72684adb93da88c3e7a4feb448ddef0a830 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sat, 31 Jan 2026 18:59:47 -0500 Subject: [PATCH 102/207] update names Signed-off-by: Robert Shaw --- .../kernels/benchmark_cutlass_moe_fp8.py | 2 +- .../kernels/benchmark_cutlass_moe_nvfp4.py | 4 +- .../kernels/benchmark_grouped_gemm_cutlass.py | 4 +- benchmarks/kernels/benchmark_moe.py | 2 +- docs/design/dbo.md | 4 +- docs/design/fused_moe_modular_kernel.md | 28 ++++---- .../moe/modular_kernel_tools/common.py | 4 +- tests/kernels/moe/test_batched_deepgemm.py | 6 +- tests/kernels/moe/test_block_fp8.py | 2 +- tests/kernels/moe/test_cutlass_moe.py | 4 +- tests/kernels/moe/test_deepep_deepgemm_moe.py | 16 ++--- tests/kernels/moe/test_deepep_moe.py | 8 +-- tests/kernels/moe/test_deepgemm.py | 2 +- tests/kernels/moe/test_flashinfer.py | 4 +- tests/kernels/moe/test_flashinfer_moe.py | 4 +- .../moe/test_modular_oai_triton_moe.py | 4 +- tests/kernels/moe/test_nvfp4_moe.py | 2 +- tests/kernels/moe/test_pplx_cutlass_moe.py | 4 +- tests/kernels/moe/test_pplx_moe.py | 4 +- tests/kernels/moe/utils.py | 6 +- vllm/lora/layers/fused_moe.py | 4 +- .../layers/fused_moe/cutlass_moe.py | 2 +- .../layers/fused_moe/fused_moe.py | 4 +- .../fused_moe/fused_moe_modular_method.py | 6 +- .../layers/fused_moe/modular_kernel.py | 70 +++---------------- .../layers/fused_moe/oracle/nvfp4.py | 2 +- .../layers/fused_moe/oracle/unquantized.py | 6 +- .../layers/fused_moe/router/base_router.py | 2 +- .../compressed_tensors_moe.py | 4 +- .../model_executor/layers/quantization/fp8.py | 2 +- .../layers/quantization/modelopt.py | 4 +- .../layers/quantization/mxfp4.py | 2 +- 32 files changed, 85 insertions(+), 137 deletions(-) diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index 78b4339d0375..117a30457ab9 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -136,7 +136,7 @@ def bench_run( per_out_ch_quant=per_out_ch, ) - fn = mk.FusedMoEModularKernel.make_mk( + fn = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index 5372f9ab5a7b..7f72e99b8ee5 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -196,7 +196,7 @@ def run_cutlass_moe_fp4( g2_alphas=w2_gs, ) - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( make_dummy_moe_config(), @@ -241,7 +241,7 @@ def run_cutlass_from_graph( g2_alphas=w2_gs, ) - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( make_dummy_moe_config(), diff --git a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py index 4e404ee9b33d..e35d28a08d8a 100644 --- a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py +++ b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py @@ -132,7 +132,7 @@ def run_cutlass_moe( per_act_token_quant=per_act_token, ) - fn = mk.FusedMoEModularKernel.make_mk( + fn = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( @@ -164,7 +164,7 @@ def run_cutlass_from_graph( per_act_token_quant=per_act_token, ) - fn = mk.FusedMoEModularKernel.make_mk( + fn = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index eefeebcd94cc..bbd8812c6fe4 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -203,7 +203,7 @@ def run(): deep_gemm_experts = None if use_deep_gemm: - deep_gemm_experts = mk.FusedMoEModularKernel.make_mk( + deep_gemm_experts = mk.FusedMoEKernelModular.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=FusedMoEConfig( diff --git a/docs/design/dbo.md b/docs/design/dbo.md index 43b3ce0bb5a7..b9c40d00f487 100644 --- a/docs/design/dbo.md +++ b/docs/design/dbo.md @@ -6,7 +6,7 @@ The core motivation of the DBO system in vLLM is to overlap the sparse all-to-al ## Introduction -The Dual Batch Overlap system works by splitting the batch in the model runner, creating two worker threads, and then running the model on each of these worker threads. When DBO is enabled, yield points within the `FusedMoEModularKernel` allow the two CPU worker threads (also called UBatch threads) to ping-pong between each other so that when one is running compute, the other is waiting on communication. Throughout the code, ubatch may be used as a short form of microbatch; this is an ASCII-friendly version of the short form µ-batch. +The Dual Batch Overlap system works by splitting the batch in the model runner, creating two worker threads, and then running the model on each of these worker threads. When DBO is enabled, yield points within the `FusedMoEKernelModular` allow the two CPU worker threads (also called UBatch threads) to ping-pong between each other so that when one is running compute, the other is waiting on communication. Throughout the code, ubatch may be used as a short form of microbatch; this is an ASCII-friendly version of the short form µ-batch. The DBO system includes modifications to `GpuModelRunner` and `ModularKernel`, and defines two utility classes: `UBatchWrapper` and `UBatchContext`. `UBatchWrapper` manages thread lifecycle and CUDA graph execution of the model. `UBatchContext` wraps `ForwardContext` to coordinate synchronization between the two UBatch threads. @@ -75,7 +75,7 @@ The `UBatchContext` class is a `ForwardContext` wrapper class that is used by th When one of the UBatch threads reaches a `dbo_yield` call, it pauses, and starts the other thread which will run until it reaches the same `dbo_yield` call. This "ping-pong" dynamic continues, with threads swapping at each `dbo_yield call`, until the model's execution is complete. -The current implementation has all `dbo_yield` and `dbo_maybe_run_recv_hook` calls in the `FusedMoEModularKernel.forward` method. +The current implementation has all `dbo_yield` and `dbo_maybe_run_recv_hook` calls in the `FusedMoEKernelModular.forward` method. #### Interfaces diff --git a/docs/design/fused_moe_modular_kernel.md b/docs/design/fused_moe_modular_kernel.md index 75321a22f02a..c18f012805a7 100644 --- a/docs/design/fused_moe_modular_kernel.md +++ b/docs/design/fused_moe_modular_kernel.md @@ -2,7 +2,7 @@ ## Introduction -FusedMoEModularKernel is implemented [here](../../vllm/model_executor/layers/fused_moe/modular_kernel.py) +FusedMoEKernelModular is implemented [here](../../vllm/model_executor/layers/fused_moe/modular_kernel.py) Based on the format of the input activations, FusedMoE implementations are broadly classified into 2 types. @@ -34,7 +34,7 @@ The rest of the document will focus on the Contiguous / Non-Batched case. Extrap ## ModularKernel Components -FusedMoEModularKernel splits the FusedMoE operation into 3 parts, +FusedMoEKernelModular splits the FusedMoE operation into 3 parts, 1. TopKWeightAndReduce 2. FusedMoEPrepareAndFinalizeModular @@ -47,7 +47,7 @@ The TopK Weight Application and Reduction components happen right after the Unpe Please find the implementations of TopKWeightAndReduce [here](../../vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py). `FusedMoEPrepareAndFinalizeModular::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method. -The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEExpertsModular` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. +The `FusedMoEKernelModular` acts as a bridge between the `FusedMoEExpertsModular` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. * `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEExpertsModular` implementation does the weight application and reduction itself. * `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEExpertsModular` implementation needs the `FusedMoEPrepareAndFinalizeModular::finalize()` to do the weight application and reduction. @@ -81,7 +81,7 @@ The `apply` method is where the implementations perform #### workspace_shapes() -The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. +The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEKernelModular::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. #### finalize_weight_and_reduce_impl() @@ -90,13 +90,13 @@ It is sometimes efficient to perform TopK weight application and Reduction insid ![FusedMoEExpertsModular Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png) -### FusedMoEModularKernel +### FusedMoEKernelModular -`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` objects. -`FusedMoEModularKernel` pseudocode/sketch, +`FusedMoEKernelModular` is composed of the `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` objects. +`FusedMoEKernelModular` pseudocode/sketch, ```py -class FusedMoEModularKernel: +class FusedMoEKernelModular: def __init__(self, prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular): @@ -177,9 +177,9 @@ implementations that input `FusedMoEActivationFormat.Standard` support chunking `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` / `FusedMoEExpertsModular::apply`: Refer to `FusedMoEExpertsModular` section above. -### FusedMoEModularKernel Initialization +### FusedMoEKernelModular Initialization -`FusedMoEMethodBase` class has 3 methods that are collectively responsible in creating the `FusedMoEModularKernel` object. They are, +`FusedMoEMethodBase` class has 3 methods that are collectively responsible in creating the `FusedMoEKernelModular` object. They are, * maybe_make_prepare_finalize, * select_gemm_impl, and @@ -206,14 +206,14 @@ derived classes. #### init_prepare_finalize -Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalizeModular` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEModularKernel` object +Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalizeModular` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEKernelModular` object Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188). -**Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used. +**Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEKernelModular` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used. ### How To Unit Test -We have `FusedMoEModularKernel` unit tests at [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py). +We have `FusedMoEKernelModular` unit tests at [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py). The unit test iterates through all combinations of `FusedMoEPrepareAndFinalizeModular` and `FusedMoEPremuteExpertsUnpermute` types and if they are compatible, runs some correctness tests. @@ -236,7 +236,7 @@ with incompatible types, the script will error. ### How To Profile Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py) -The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible +The script can be used to generate Torch traces for a single `FusedMoEKernelModular::forward()` call for any compatible `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` types. Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index a83e85942d02..2fb120b2693b 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -575,7 +575,7 @@ def make_modular_kernel( config: Config, vllm_config: VllmConfig, quant_config: FusedMoEQuantConfig, -) -> mk.FusedMoEModularKernel: +) -> mk.FusedMoEKernelModular: def next_power_of_2(x): import math @@ -620,7 +620,7 @@ def next_power_of_2(x): config.N, ) - modular_kernel = mk.FusedMoEModularKernel.make_mk( + modular_kernel = mk.FusedMoEKernelModular.make_mk( prepare_finalize=prepare_finalize, fused_experts=fused_experts, ) diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py index 081a5fd0b93c..d81c6bd7ed99 100644 --- a/tests/kernels/moe/test_batched_deepgemm.py +++ b/tests/kernels/moe/test_batched_deepgemm.py @@ -12,7 +12,7 @@ BatchedPrepareAndFinalize, BatchedTritonExperts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported from .test_deepgemm import make_block_quant_fp8_weights @@ -74,7 +74,7 @@ def test_batched_deepgemm_vs_triton( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk_triton = FusedMoEModularKernel(prep_finalize, triton_experts) + mk_triton = FusedMoEKernelModular(prep_finalize, triton_experts) out_triton = mk_triton( hidden_states=a, @@ -93,7 +93,7 @@ def test_batched_deepgemm_vs_triton( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk_deepgemm = FusedMoEModularKernel(prep_finalize, deepgemm_experts) + mk_deepgemm = FusedMoEKernelModular(prep_finalize, deepgemm_experts) out_deepgemm = mk_deepgemm( hidden_states=a, diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 230968a8b5fe..21658f9064d7 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -255,7 +255,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) block_shape=block_size, ) - deep_gemm_experts = mk.FusedMoEModularKernel.make_mk( + deep_gemm_experts = mk.FusedMoEKernelModular.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 354d27bcee47..c6abe10b9f63 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -196,7 +196,7 @@ def slice_experts(): for kwargs, new_quant_config in slice_experts(): w2 = kwargs["w2"] a = kwargs["hidden_states"] - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( @@ -255,7 +255,7 @@ def run_8_bit( num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined] with_ep = num_local_experts is not None or num_local_experts == num_experts if not with_ep: - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index 1bf5ced2e84c..7a7cc0181427 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -21,7 +21,7 @@ fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular from vllm.utils.deep_gemm import ( get_mk_alignment_for_contiguous_layout, is_deep_gemm_e8m0_used, @@ -169,7 +169,7 @@ def make_ll_modular_kernel( q_dtype: torch.dtype | None, test_config: TestConfig, quant_config: FusedMoEQuantConfig, -) -> FusedMoEModularKernel: +) -> FusedMoEKernelModular: assert test_config.low_latency assert test_config.use_fp8_dispatch is not None @@ -194,7 +194,7 @@ def make_ll_modular_kernel( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) + mk = FusedMoEKernelModular(prepare_finalize=a2a, fused_experts=fused_experts) return mk @@ -206,7 +206,7 @@ def make_ht_modular_kernel( q_dtype: torch.dtype | None, test_config: TestConfig, quant_config: FusedMoEQuantConfig, -) -> FusedMoEModularKernel: +) -> FusedMoEKernelModular: assert not test_config.low_latency assert test_config.use_fp8_dispatch is None @@ -224,7 +224,7 @@ def make_ht_modular_kernel( moe_config=make_dummy_moe_config(), quant_config=quant_config, ) - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) + mk = FusedMoEKernelModular(prepare_finalize=a2a, fused_experts=fused_experts) return mk @@ -235,11 +235,11 @@ def make_modular_kernel( num_local_experts: int, test_tensors: TestTensors, quant_config: FusedMoEQuantConfig, -) -> FusedMoEModularKernel: +) -> FusedMoEKernelModular: q_dtype = torch.float8_e4m3fn test_config = test_tensors.config - mk: FusedMoEModularKernel + mk: FusedMoEKernelModular # Make modular kernel if test_config.low_latency: max_tokens_per_rank = max(64, next_power_of_2(test_tensors.rank_tokens.size(0))) @@ -300,7 +300,7 @@ def build_expert_map(): ) # Make modular kernel - mk: FusedMoEModularKernel = make_modular_kernel( + mk: FusedMoEKernelModular = make_modular_kernel( pg=pg, pgi=pgi, dp_size=dp_size, diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index f740f5bf9585..5c6a6af2181f 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -19,7 +19,7 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -134,7 +134,7 @@ def make_modular_kernel( q_dtype: torch.dtype | None, use_fp8_dispatch: bool, quant_config: FusedMoEQuantConfig, -) -> FusedMoEModularKernel: +) -> FusedMoEKernelModular: ht_args: DeepEPHTArgs | None = None ll_args: DeepEPLLArgs | None = None @@ -179,7 +179,7 @@ def make_modular_kernel( quant_config=quant_config, ) - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) + mk = FusedMoEKernelModular(prepare_finalize=a2a, fused_experts=fused_experts) return mk @@ -237,7 +237,7 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False): ) # Make modular kernel - mk: FusedMoEModularKernel = make_modular_kernel( + mk: FusedMoEKernelModular = make_modular_kernel( pg, pgi, low_latency_mode, diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 8363b6465adb..95dfe72ffaf0 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -109,7 +109,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size): block_shape=block_size, ) - deep_gemm_experts = mk.FusedMoEModularKernel.make_mk( + deep_gemm_experts = mk.FusedMoEKernelModular.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index c8d76a623b87..1ff05b49b1f0 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -211,7 +211,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( quant_config=quant_config, ) - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEPMonolithic(), FlashInferTrtLlmFp8Experts( moe_config=td.layer.moe, @@ -312,7 +312,7 @@ def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig: routing_method=RoutingMethodType.TopK, ) - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index 9bb61ddfa0fe..3bbde1541070 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -22,7 +22,7 @@ FlashInferExperts, is_valid_flashinfer_cutlass_fused_moe, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -105,7 +105,7 @@ def test_flashinfer_fp4_moe_no_graph( routing_method=RoutingMethodType.TopK, ) - flashinfer_experts = FusedMoEModularKernel( + flashinfer_experts = FusedMoEKernelModular( MoEPrepareAndFinalizeNoEP(), FlashInferExperts(moe_config=moe_config, quant_config=quant_config), ) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 38022e0e61b7..d58670e61180 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -28,7 +28,7 @@ OAITritonExperts, UnfusedOAITritonExperts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -180,7 +180,7 @@ def oai_triton_moe_impl( else: fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config) - mk = FusedMoEModularKernel(MoEPrepareAndFinalizeNoEP(), fused_experts) + mk = FusedMoEKernelModular(MoEPrepareAndFinalizeNoEP(), fused_experts) return mk.forward( hidden_states=x, diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index b2e2916703f3..b531bd72dcd9 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -89,7 +89,7 @@ def test_cutlass_fp4_moe_no_graph( w2_scale=w2_blockscale, ) - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_pplx_cutlass_moe.py b/tests/kernels/moe/test_pplx_cutlass_moe.py index ef37c1c74434..9b30e1fc858f 100644 --- a/tests/kernels/moe/test_pplx_cutlass_moe.py +++ b/tests/kernels/moe/test_pplx_cutlass_moe.py @@ -16,7 +16,7 @@ fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassBatchedExpertsFp8 -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import set_random_seed @@ -169,7 +169,7 @@ def make_moe_config() -> FusedMoEConfig: num_dispatchers=num_dispatchers, ) - fused_cutlass_experts = FusedMoEModularKernel( + fused_cutlass_experts = FusedMoEKernelModular( prepare_finalize, experts, ) diff --git a/tests/kernels/moe/test_pplx_moe.py b/tests/kernels/moe/test_pplx_moe.py index 08519087e1ce..fa885fde0384 100644 --- a/tests/kernels/moe/test_pplx_moe.py +++ b/tests/kernels/moe/test_pplx_moe.py @@ -41,7 +41,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceDelegate, ) @@ -588,7 +588,7 @@ def pplx_moe( moe_config=make_dummy_moe_config(), ) - fused_experts = FusedMoEModularKernel( + fused_experts = FusedMoEKernelModular( prepare_finalize, experts, shared_experts, diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 4883085cb836..4f4ec1977ec8 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -19,7 +19,7 @@ BatchedTritonExperts, NaiveBatchedExperts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.deep_gemm import per_block_cast_to_fp8 from vllm.utils.math_utils import round_up @@ -106,7 +106,7 @@ def batched_moe( a2_scale=a2_scale, ) - fused_experts = FusedMoEModularKernel( + fused_experts = FusedMoEKernelModular( BatchedPrepareAndFinalize( max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 ), @@ -147,7 +147,7 @@ def naive_batched_moe( a2_scale=a2_scale, ) - fused_experts = FusedMoEModularKernel( + fused_experts = FusedMoEKernelModular( BatchedPrepareAndFinalize( max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 ), diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 8f5cb14c3d09..c699e86db3e5 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -32,7 +32,7 @@ UnfusedOAITritonExperts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEModularKernel, + FusedMoEKernelModular, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, @@ -131,7 +131,7 @@ def _inject_lora_into_fused_moe(self): quant_config = self.base_layer.quant_method.moe_quant_config prepare_finalize = MoEPrepareAndFinalizeNoEP() - m_fused_moe_fn = FusedMoEModularKernel( + m_fused_moe_fn = FusedMoEKernelModular( prepare_finalize, self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 3f15ab7ecf06..28fd339d5c61 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -1144,7 +1144,7 @@ def cutlass_moe_w4a8_fp8( num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0) - fn = mk.FusedMoEModularKernel.make_mk( + fn = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsW4A8Fp8( out_dtype=a.dtype, diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index ab39d95d8246..ede2269eb677 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -2305,8 +2305,8 @@ def modular_triton_fused_moe( moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None, -) -> mk.FusedMoEModularKernel: - return mk.FusedMoEModularKernel.make_mk( +) -> mk.FusedMoEKernelModular: + return mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), TritonExperts(moe_config, quant_config), shared_experts, diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 093b9d2e2d4e..3060102791c3 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -13,7 +13,7 @@ FusedMoEMethodBase, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEModularKernel, + FusedMoEKernelModular, FusedMoEPrepareAndFinalizeModular, ) @@ -26,7 +26,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): # --8<-- [end:modular_fused_moe] def __init__( - self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEModularKernel + self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEKernelModular ): super().__init__(old_quant_method.moe) self.moe_quant_config = old_quant_method.moe_quant_config @@ -49,7 +49,7 @@ def make( ) -> "FusedMoEModularMethod": return FusedMoEModularMethod( old_quant_method, - FusedMoEModularKernel( + FusedMoEKernelModular( prepare_finalize, old_quant_method.select_gemm_impl(prepare_finalize, moe_layer), shared_experts, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 4d1e19f93581..9212b03cc2e7 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -63,7 +63,7 @@ # Some FusedMoEExpertsModular implementations may choose to do # the weight application and/or reduction. The class communicates this # to [Finalize] via a TopKWeightAndReduce object. -# * FusedMoEModularKernel - an interface class that combines a +# * FusedMoEKernelModular - an interface class that combines a # FusedMoEPrepareAndFinalizeModular and a FusedMoEExpertsModular to # provide the standard fused MoE kernel interface. # * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen @@ -986,7 +986,7 @@ def make_mk( if isinstance( prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): - return FusedMoEMonolithicKernel( + return FusedMoEKernelMonolithic( prepare_finalize, fused_experts, shared_experts, @@ -995,7 +995,7 @@ def make_mk( elif isinstance( prepare_finalize, FusedMoEPrepareAndFinalizeModular ) and isinstance(fused_experts, FusedMoEExpertsModular): - return FusedMoEModularKernel( + return FusedMoEKernelModular( prepare_finalize, fused_experts, shared_experts, @@ -1030,7 +1030,7 @@ def output_is_reduced(self) -> bool: @final -class FusedMoEModularKernel(FusedMoEKernel): +class FusedMoEKernelModular(FusedMoEKernel): """ This class combines a FusedMoEPrepareAndFinalizeModular instance and a FusedMoEExpertsModular to provide an interface that @@ -1043,36 +1043,6 @@ class FusedMoEModularKernel(FusedMoEKernel): objects. """ - def __init__( - self, - prepare_finalize: FusedMoEPrepareAndFinalizeModular, - fused_experts: FusedMoEExpertsModular, - shared_experts: torch.nn.Module | None = None, - moe_parallel_config: FusedMoEParallelConfig | None = None, - ): - if not isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular): - raise TypeError( - "prepare_finalize must be an instance of " - "FusedMoEPrepareAndFinalizeModular, but got " - f"{prepare_finalize.__class__.__name__}." - ) - if not isinstance(fused_experts, FusedMoEExpertsModular): - raise TypeError( - "fused_experts must be an instance of " - "FusedMoEExpertsModular, but got " - f"{fused_experts.__class__.__name__}." - ) - - super().__init__( - prepare_finalize, - fused_experts, - shared_experts, - moe_parallel_config, - ) - - self.prepare_finalize: FusedMoEPrepareAndFinalizeModular = prepare_finalize - self.fused_experts: FusedMoEExpertsModular = fused_experts - def _chunk_info(self, M: int) -> tuple[int, int]: """ Compute number of chunks and chunk size for given M. @@ -1122,7 +1092,7 @@ def _allocate_buffers( workspace_dtype = self.fused_experts.workspace_dtype(out_dtype) # Force worst-case allocation in profiling run for - # "mk.FusedMoEModularKernel.Standard" formats where this is only bounded + # "mk.FusedMoEKernelModular.Standard" formats where this is only bounded # by `VLLM_FUSED_MOE_CHUNK_SIZE` and may not be seen during profiling with # DP+EP due to the random token routing. is_profile_run = ( @@ -1358,7 +1328,7 @@ def _fused_experts( apply_router_weight_on_input: bool, expert_tokens_meta: ExpertTokensMetadata | None, ) -> torch.Tensor: - assert isinstance(self.fused_experts, FusedMoEExperts) + assert isinstance(self.fused_experts, FusedMoEExpertsModular) _, M_full, N, K, top_k = self.fused_experts.moe_problem_size( a1q, w1, w2, topk_ids @@ -1451,6 +1421,8 @@ def _finalize( The _finalize method is a wrapper around self.prepare_finalize.finalize that handles DBO, async and shared expert overlap. """ + assert isinstance(self.fused_experts, FusedMoEExpertsModular) + assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeModular) shared_output: torch.Tensor | None = None if not self.prepare_finalize.supports_async(): @@ -1591,31 +1563,7 @@ def forward( @final -class FusedMoEMonolithicKernel(FusedMoEKernel): - def __init__( - self, - prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic, - fused_experts: FusedMoEExpertsMonolithic, - shared_experts: torch.nn.Module | None = None, - moe_parallel_config: FusedMoEParallelConfig | None = None, - ): - if not isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic): - raise TypeError( - "prepare_finalize must be an instance of " - "FusedMoEPrepareAndFinalizeMonolithic" - ) - if not isinstance(fused_experts, FusedMoEExpertsMonolithic): - raise TypeError( - "fused_experts must be an instance of FusedMoEExpertsMonolithic" - ) - - super().__init__( - prepare_finalize, - fused_experts, - shared_experts, - moe_parallel_config, - ) - +class FusedMoEKernelMonolithic(FusedMoEKernel): def forward( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 7fa75a2efebc..4c1cb2045082 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -405,7 +405,7 @@ def make_nvfp4_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( prepare_finalize, experts, shared_experts=( diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 05ed39ccdcf6..a4f77337aba8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -126,7 +126,7 @@ def make_unquantized_moe_kernel( backend: UnquantizedMoeBackend, quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, -) -> tuple[mk.FusedMoEModularKernel | None, bool]: +) -> tuple[mk.FusedMoEKernelModular | None, bool]: use_inplace = True if backend in UNSUPPORTED_BACKEND: @@ -137,7 +137,7 @@ def make_unquantized_moe_kernel( FlashInferExperts, ) - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, @@ -150,7 +150,7 @@ def make_unquantized_moe_kernel( AiterExperts, ) - kernel = mk.FusedMoEModularKernel.make_mk( + kernel = mk.FusedMoEKernelModular.make_mk( MoEPrepareAndFinalizeNoEP(), AiterExperts( moe_config=moe_config, diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index a1ef9df4e4a6..f8d551f53806 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -62,7 +62,7 @@ def eplb_map_to_physical_and_record( # 2. Record expert load metrics. - # TODO(bowen): When using `FusedMoEModularKernel`, this + # TODO(bowen): When using `FusedMoEKernelModular`, this # can be done in a more unified way, since # `FusedMoEPrepareAndFinalizeModular` will return the expert # token count, in some cases directly from the kernel. diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 1c4db3b8cc1a..e7bb5bf52c87 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -602,7 +602,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) return self.moe_kernel( x, layer.w13_weight, @@ -974,7 +974,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) return self.moe_kernel( x, layer.w13_weight, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 4f5cba747385..61456ce027ab 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -940,7 +940,7 @@ def apply_monolithic( if layer.enable_eplb: raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.") - assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) return self.moe_kernel( x, layer.w13_weight, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index d09a1c03ae09..c165a82458b1 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -926,7 +926,7 @@ def apply_monolithic( raise NotImplementedError( "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend." ) - assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) return self.moe_kernel( x, layer.w13_weight, @@ -1565,7 +1565,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEMonolithicKernel) + assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) return self.moe_kernel.forward_monolithic( x, layer.w13_weight, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 1ffe94a5aaeb..6fe16d4c08c8 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -751,7 +751,7 @@ def _interleave_mxfp4_cutlass_sm90(w): layer.w13_bias = Parameter(w13_bias, requires_grad=False) layer.w2_bias = Parameter(w2_bias, requires_grad=False) - # Ideally we'd use FusedMoEModularKernel.prepare_finalize object + # Ideally we'd use FusedMoEKernelModular.prepare_finalize object # (stored in self.fused_experts) to determine if the MoE has a # batched activation format. As self.fused_experts is not # initialized at this point, we resort to checking the MoE config From e00bec1c29a699ad8625d4c0a1db922c09ba0156 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Feb 2026 19:56:43 -0500 Subject: [PATCH 103/207] add bf16 config back Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_moe.py | 253 ------------------ 1 file changed, 253 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index a066535c51eb..c7420e5dbd9d 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -9,16 +9,6 @@ FusedMoEParallelConfig, RoutingMethodType, ) -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - per_token_group_quant_fp8, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - QuantKey, - kFp8Dynamic128Sym, - kFp8Static128BlockSym, - kFp8StaticTensorSym, -) from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op @@ -38,47 +28,11 @@ def _supports_no_act_and_mul() -> bool: return False -def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, -) -> bool: - """Supports Fp8 per-tensor and Fp8 block.""" - SUPPORTED_W_A = [ - (kFp8Static128BlockSym, kFp8Dynamic128Sym), - (kFp8StaticTensorSym, kFp8StaticTensorSym), - ] - return (weight_key, activation_key) in SUPPORTED_W_A - - def _supports_activation(activation: str) -> bool: """Supports silu activation only.""" return activation in ["silu"] -def _supports_routing_method( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - routing_method: RoutingMethodType, -) -> bool: - """Monolithic kernels need to express router support.""" - if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): - # NOTE(rob): potentially allow others here. This is a conservative list. - return routing_method in [ - RoutingMethodType.DeepSeekV3, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ] - elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym): - # NOTE(dbari): as above, potentially allow others here. - return routing_method in [ - RoutingMethodType.Llama4, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ] - else: - raise ValueError("Unsupported quantization scheme.") - - def _supports_routing_method_bf16( routing_method: RoutingMethodType, ) -> bool: @@ -96,39 +50,6 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo return not moe_parallel_config.enable_eplb -def is_supported_config_trtllm( - moe_config: FusedMoEConfig, - weight_key: QuantKey | None, - activation_key: QuantKey | None, - activation_format: mk.FusedMoEActivationFormat, -) -> tuple[bool, str | None]: - """ - This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config - """ - - def _make_reason(reason: str) -> str: - return f"kernel does not support {reason}" - - if not _supports_current_device(): - return False, _make_reason("current device") - elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()): - return False, _make_reason("no act_and_mul MLP layer") - elif not _supports_activation(moe_config.activation): - return False, _make_reason(f"{moe_config.activation} activation") - elif not _supports_quant_scheme(weight_key, activation_key): - return False, _make_reason("quantization scheme") - elif not _supports_parallel_config(moe_config.moe_parallel_config): - return False, _make_reason("parallel config") - elif not _supports_routing_method( - weight_key, activation_key, moe_config.routing_method - ): - return False, _make_reason("routing method") - elif activation_format != mk.FusedMoEActivationFormat.Standard: - return False, _make_reason("activation format") - - return True, None - - def is_supported_config_trtllm_bf16( moe_config: FusedMoEConfig, activation_format: mk.FusedMoEActivationFormat, @@ -157,180 +78,6 @@ def _make_reason(reason: str) -> str: return True, None -def flashinfer_fused_moe_blockscale_fp8( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - x: torch.Tensor, - w13_weight: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - global_num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - expert_offset: int, - local_num_experts: int, - block_shape: list[int], - routing_method_type: int = int(RoutingMethodType.DeepSeekV3), - routed_scaling: float | None = 1.0, -) -> torch.Tensor: - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe - - topk_group = topk_group if topk_group is not None else 0 - assert top_k <= global_num_experts - assert top_k <= 10 - assert global_num_experts % 4 == 0 - assert block_shape == [128, 128] - # Routing kernel expects #experts <= #threads 512 - assert global_num_experts <= 512 - - a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1]) - # NOTE: scales of hidden states have to be transposed! - a_sf_t = a_sf.t().contiguous() - return flashinfer_trtllm_fp8_block_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=a_q, - hidden_states_scale=a_sf_t, - gemm1_weights=w13_weight, - gemm1_weights_scale=w13_weight_scale_inv, - gemm2_weights=w2_weight, - gemm2_weights_scale=w2_weight_scale_inv, - num_experts=global_num_experts, - top_k=top_k, - n_group=num_expert_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=expert_offset, - local_num_experts=local_num_experts, - routed_scaling_factor=routed_scaling, - routing_method_type=routing_method_type, - use_shuffled_weight=False, - ) - - -def flashinfer_fused_moe_blockscale_fp8_fake( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - x: torch.Tensor, - w13_weight: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - global_num_experts: int, - top_k: int, - num_expert_group: int, - topk_group: int, - intermediate_size: int, - expert_offset: int, - local_num_experts: int, - block_shape: list[int], - routing_method_type: int, - routed_scaling: float = 1.0, -) -> torch.Tensor: - return torch.empty_like(x) - - -# TODO(bnell): Does this really need to be a torch.op? -direct_register_custom_op( - op_name="flashinfer_fused_moe_blockscale_fp8", - op_func=flashinfer_fused_moe_blockscale_fp8, - fake_impl=flashinfer_fused_moe_blockscale_fp8_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -def fi_trtllm_fp8_per_tensor_moe( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - hidden_states: torch.Tensor, - input_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm2_weights: torch.Tensor, - output1_scales_scalar: torch.Tensor, - output1_scales_gate_scalar: torch.Tensor, - output2_scales_scalar: torch.Tensor, - num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - use_routing_scales_on_input: bool, - routing_method_type: int, - routed_scaling_factor: float = 1.0, -) -> torch.Tensor: - num_expert_group = num_expert_group if num_expert_group is not None else 0 - topk_group = topk_group if topk_group is not None else 0 - - quant_hidden_states, _ = moe_kernel_quantize_input( - hidden_states, - input_scale, - quant_dtype=torch.float8_e4m3fn, - per_act_token_quant=False, - ) - - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe - - return flashinfer_trtllm_fp8_per_tensor_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=quant_hidden_states, - gemm1_weights=gemm1_weights, - output1_scales_scalar=output1_scales_scalar, - output1_scales_gate_scalar=output1_scales_gate_scalar, - gemm2_weights=gemm2_weights, - output2_scales_scalar=output2_scales_scalar, - num_experts=num_experts, - top_k=top_k, - n_group=num_expert_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=local_expert_offset, - local_num_experts=local_num_experts, - routed_scaling_factor=routed_scaling_factor, - use_routing_scales_on_input=use_routing_scales_on_input, - routing_method_type=routing_method_type, - ) - - -def fi_trtllm_fp8_per_tensor_moe_fake( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor | None, - hidden_states: torch.Tensor, - input_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm2_weights: torch.Tensor, - output1_scales_scalar: torch.Tensor, - output1_scales_gate_scalar: torch.Tensor, - output2_scales_scalar: torch.Tensor, - num_experts: int, - top_k: int, - num_expert_group: int | None, - topk_group: int | None, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - use_routing_scales_on_input: bool, - routing_method_type: int, - routed_scaling_factor: float = 1.0, -) -> torch.Tensor: - return torch.empty_like(hidden_states) - - -# TODO(bnell): Does this really need to be a torch.op? -direct_register_custom_op( - op_name="fi_trtllm_fp8_per_tensor_moe", - op_func=fi_trtllm_fp8_per_tensor_moe, - mutates_args=["hidden_states"], - fake_impl=fi_trtllm_fp8_per_tensor_moe_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - def flashinfer_fused_moe_bf16( routing_logits: torch.Tensor, routing_bias: torch.Tensor | None, From 81bed133ae9144af64aefab02d2e122ec3bea172 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Feb 2026 20:01:58 -0500 Subject: [PATCH 104/207] fix test_flashinfer Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 3401d100f665..f107492ee320 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -151,6 +151,9 @@ def make_moe_tensors_8bit( moe_parallel_config=layer.moe_parallel_config, in_dtype=hidden_states.dtype, is_act_and_mul=is_gated, + routing_method=layer.routing_method_type, + activation=activation, + device=layer.device, ) return TestData( From 408258eb5afd59eb648151c6e5665662b5397d44 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Feb 2026 20:02:52 -0500 Subject: [PATCH 105/207] fix tests Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index f107492ee320..c6e57916f5b1 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -153,7 +153,7 @@ def make_moe_tensors_8bit( is_act_and_mul=is_gated, routing_method=layer.routing_method_type, activation=activation, - device=layer.device, + device=w13_quantized.device, ) return TestData( From 332b5bc2e1795990f58ea5c9dd353e99d5321ff8 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Feb 2026 20:13:31 -0500 Subject: [PATCH 106/207] fix naive ep Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 1 + vllm/model_executor/layers/fused_moe/prepare_finalize.py | 1 + vllm/model_executor/layers/quantization/modelopt.py | 2 +- 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index c6e57916f5b1..4eb270bea03b 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -233,6 +233,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( global_num_experts=e, expert_map=None, apply_router_weight_on_input=True, + routed_scaling_factor=1.0, ) torch.testing.assert_close(output, flashinfer_output, atol=5.5e-2, rtol=1e-2) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py index f25da1f298b8..b2bdd109a448 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py @@ -169,6 +169,7 @@ def prepare( if scales is None: a1q, topk_weights, topk_ids = res + a1q_scale = None else: a1q, topk_weights, topk_ids, scales = res a1q_scale = self._unwrap_scale_and_prepare_for_moe(scales, quant_config) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 254b9c6e2b4c..e1fe956a2ec2 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1438,7 +1438,7 @@ def apply_monolithic( router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) - return self.moe_kernel.forward_monolithic( + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, From 517e61d3ef3b318f4eade34ff6a924beb1a40a11 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Feb 2026 20:15:09 -0500 Subject: [PATCH 107/207] fix import Signed-off-by: Robert Shaw --- .../layers/quantization/utils/flashinfer_fp4_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 2507b50c1479..230258240052 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -8,7 +8,7 @@ import vllm.envs as envs from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.quant_utils import ( +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( swizzle_blockscale, ) from vllm.platforms import current_platform From ab7557274e3705261770f031df718937cb7b97c5 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Feb 2026 20:18:57 -0500 Subject: [PATCH 108/207] fix up typing Signed-off-by: Robert Shaw --- .../kernels/benchmark_cutlass_moe_fp8.py | 2 +- .../kernels/benchmark_cutlass_moe_nvfp4.py | 4 +-- .../kernels/benchmark_grouped_gemm_cutlass.py | 4 +-- benchmarks/kernels/benchmark_moe.py | 2 +- docs/design/dbo.md | 4 +-- docs/design/fused_moe_modular_kernel.md | 28 +++++++++---------- .../moe/modular_kernel_tools/common.py | 4 +-- tests/kernels/moe/test_batched_deepgemm.py | 6 ++-- tests/kernels/moe/test_block_fp8.py | 2 +- tests/kernels/moe/test_cutlass_moe.py | 4 +-- tests/kernels/moe/test_deepep_deepgemm_moe.py | 16 +++++------ tests/kernels/moe/test_deepep_moe.py | 8 +++--- tests/kernels/moe/test_deepgemm.py | 2 +- tests/kernels/moe/test_flashinfer.py | 4 +-- tests/kernels/moe/test_flashinfer_moe.py | 4 +-- .../moe/test_modular_oai_triton_moe.py | 4 +-- tests/kernels/moe/test_nvfp4_moe.py | 2 +- tests/kernels/moe/test_pplx_cutlass_moe.py | 4 +-- tests/kernels/moe/test_pplx_moe.py | 4 +-- tests/kernels/moe/utils.py | 6 ++-- vllm/lora/layers/fused_moe.py | 4 +-- .../layers/fused_moe/cutlass_moe.py | 2 +- .../layers/fused_moe/fused_moe.py | 4 +-- .../fused_moe/fused_moe_modular_method.py | 6 ++-- .../layers/fused_moe/modular_kernel.py | 12 ++++---- .../layers/fused_moe/oracle/nvfp4.py | 2 +- .../layers/fused_moe/oracle/unquantized.py | 6 ++-- .../layers/fused_moe/router/base_router.py | 2 +- .../compressed_tensors_moe.py | 4 +-- .../model_executor/layers/quantization/fp8.py | 2 +- .../layers/quantization/modelopt.py | 4 +-- .../layers/quantization/mxfp4.py | 2 +- 32 files changed, 82 insertions(+), 82 deletions(-) diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index 117a30457ab9..4b7c0d23ced0 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -136,7 +136,7 @@ def bench_run( per_out_ch_quant=per_out_ch, ) - fn = mk.FusedMoEKernelModular.make_mk( + fn = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index 7f72e99b8ee5..ebd79f8705ce 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -196,7 +196,7 @@ def run_cutlass_moe_fp4( g2_alphas=w2_gs, ) - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( make_dummy_moe_config(), @@ -241,7 +241,7 @@ def run_cutlass_from_graph( g2_alphas=w2_gs, ) - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( make_dummy_moe_config(), diff --git a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py index e35d28a08d8a..7a8f8be4fa18 100644 --- a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py +++ b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py @@ -132,7 +132,7 @@ def run_cutlass_moe( per_act_token_quant=per_act_token, ) - fn = mk.FusedMoEKernelModular.make_mk( + fn = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( @@ -164,7 +164,7 @@ def run_cutlass_from_graph( per_act_token_quant=per_act_token, ) - fn = mk.FusedMoEKernelModular.make_mk( + fn = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index df5d85bff40e..04271b168f46 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -202,7 +202,7 @@ def run(): deep_gemm_experts = None if use_deep_gemm: - deep_gemm_experts = mk.FusedMoEKernelModular.make_mk( + deep_gemm_experts = mk.FusedMoEKernel.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=FusedMoEConfig( diff --git a/docs/design/dbo.md b/docs/design/dbo.md index b9c40d00f487..43b3ce0bb5a7 100644 --- a/docs/design/dbo.md +++ b/docs/design/dbo.md @@ -6,7 +6,7 @@ The core motivation of the DBO system in vLLM is to overlap the sparse all-to-al ## Introduction -The Dual Batch Overlap system works by splitting the batch in the model runner, creating two worker threads, and then running the model on each of these worker threads. When DBO is enabled, yield points within the `FusedMoEKernelModular` allow the two CPU worker threads (also called UBatch threads) to ping-pong between each other so that when one is running compute, the other is waiting on communication. Throughout the code, ubatch may be used as a short form of microbatch; this is an ASCII-friendly version of the short form µ-batch. +The Dual Batch Overlap system works by splitting the batch in the model runner, creating two worker threads, and then running the model on each of these worker threads. When DBO is enabled, yield points within the `FusedMoEModularKernel` allow the two CPU worker threads (also called UBatch threads) to ping-pong between each other so that when one is running compute, the other is waiting on communication. Throughout the code, ubatch may be used as a short form of microbatch; this is an ASCII-friendly version of the short form µ-batch. The DBO system includes modifications to `GpuModelRunner` and `ModularKernel`, and defines two utility classes: `UBatchWrapper` and `UBatchContext`. `UBatchWrapper` manages thread lifecycle and CUDA graph execution of the model. `UBatchContext` wraps `ForwardContext` to coordinate synchronization between the two UBatch threads. @@ -75,7 +75,7 @@ The `UBatchContext` class is a `ForwardContext` wrapper class that is used by th When one of the UBatch threads reaches a `dbo_yield` call, it pauses, and starts the other thread which will run until it reaches the same `dbo_yield` call. This "ping-pong" dynamic continues, with threads swapping at each `dbo_yield call`, until the model's execution is complete. -The current implementation has all `dbo_yield` and `dbo_maybe_run_recv_hook` calls in the `FusedMoEKernelModular.forward` method. +The current implementation has all `dbo_yield` and `dbo_maybe_run_recv_hook` calls in the `FusedMoEModularKernel.forward` method. #### Interfaces diff --git a/docs/design/fused_moe_modular_kernel.md b/docs/design/fused_moe_modular_kernel.md index c18f012805a7..75321a22f02a 100644 --- a/docs/design/fused_moe_modular_kernel.md +++ b/docs/design/fused_moe_modular_kernel.md @@ -2,7 +2,7 @@ ## Introduction -FusedMoEKernelModular is implemented [here](../../vllm/model_executor/layers/fused_moe/modular_kernel.py) +FusedMoEModularKernel is implemented [here](../../vllm/model_executor/layers/fused_moe/modular_kernel.py) Based on the format of the input activations, FusedMoE implementations are broadly classified into 2 types. @@ -34,7 +34,7 @@ The rest of the document will focus on the Contiguous / Non-Batched case. Extrap ## ModularKernel Components -FusedMoEKernelModular splits the FusedMoE operation into 3 parts, +FusedMoEModularKernel splits the FusedMoE operation into 3 parts, 1. TopKWeightAndReduce 2. FusedMoEPrepareAndFinalizeModular @@ -47,7 +47,7 @@ The TopK Weight Application and Reduction components happen right after the Unpe Please find the implementations of TopKWeightAndReduce [here](../../vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py). `FusedMoEPrepareAndFinalizeModular::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method. -The `FusedMoEKernelModular` acts as a bridge between the `FusedMoEExpertsModular` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. +The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEExpertsModular` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens. * `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEExpertsModular` implementation does the weight application and reduction itself. * `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEExpertsModular` implementation needs the `FusedMoEPrepareAndFinalizeModular::finalize()` to do the weight application and reduction. @@ -81,7 +81,7 @@ The `apply` method is where the implementations perform #### workspace_shapes() -The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEKernelModular::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. +The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation. #### finalize_weight_and_reduce_impl() @@ -90,13 +90,13 @@ It is sometimes efficient to perform TopK weight application and Reduction insid ![FusedMoEExpertsModular Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png) -### FusedMoEKernelModular +### FusedMoEModularKernel -`FusedMoEKernelModular` is composed of the `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` objects. -`FusedMoEKernelModular` pseudocode/sketch, +`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` objects. +`FusedMoEModularKernel` pseudocode/sketch, ```py -class FusedMoEKernelModular: +class FusedMoEModularKernel: def __init__(self, prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular): @@ -177,9 +177,9 @@ implementations that input `FusedMoEActivationFormat.Standard` support chunking `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` / `FusedMoEExpertsModular::apply`: Refer to `FusedMoEExpertsModular` section above. -### FusedMoEKernelModular Initialization +### FusedMoEModularKernel Initialization -`FusedMoEMethodBase` class has 3 methods that are collectively responsible in creating the `FusedMoEKernelModular` object. They are, +`FusedMoEMethodBase` class has 3 methods that are collectively responsible in creating the `FusedMoEModularKernel` object. They are, * maybe_make_prepare_finalize, * select_gemm_impl, and @@ -206,14 +206,14 @@ derived classes. #### init_prepare_finalize -Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalizeModular` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEKernelModular` object +Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalizeModular` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEModularKernel` object Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188). -**Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEKernelModular` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used. +**Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used. ### How To Unit Test -We have `FusedMoEKernelModular` unit tests at [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py). +We have `FusedMoEModularKernel` unit tests at [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py). The unit test iterates through all combinations of `FusedMoEPrepareAndFinalizeModular` and `FusedMoEPremuteExpertsUnpermute` types and if they are compatible, runs some correctness tests. @@ -236,7 +236,7 @@ with incompatible types, the script will error. ### How To Profile Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py) -The script can be used to generate Torch traces for a single `FusedMoEKernelModular::forward()` call for any compatible +The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` types. Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 2fb120b2693b..9a8159b55ee6 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -575,7 +575,7 @@ def make_modular_kernel( config: Config, vllm_config: VllmConfig, quant_config: FusedMoEQuantConfig, -) -> mk.FusedMoEKernelModular: +) -> mk.FusedMoEModularKernel: def next_power_of_2(x): import math @@ -620,7 +620,7 @@ def next_power_of_2(x): config.N, ) - modular_kernel = mk.FusedMoEKernelModular.make_mk( + modular_kernel = mk.FusedMoEKernel.make_mk( prepare_finalize=prepare_finalize, fused_experts=fused_experts, ) diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py index d81c6bd7ed99..081a5fd0b93c 100644 --- a/tests/kernels/moe/test_batched_deepgemm.py +++ b/tests/kernels/moe/test_batched_deepgemm.py @@ -12,7 +12,7 @@ BatchedPrepareAndFinalize, BatchedTritonExperts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported from .test_deepgemm import make_block_quant_fp8_weights @@ -74,7 +74,7 @@ def test_batched_deepgemm_vs_triton( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk_triton = FusedMoEKernelModular(prep_finalize, triton_experts) + mk_triton = FusedMoEModularKernel(prep_finalize, triton_experts) out_triton = mk_triton( hidden_states=a, @@ -93,7 +93,7 @@ def test_batched_deepgemm_vs_triton( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk_deepgemm = FusedMoEKernelModular(prep_finalize, deepgemm_experts) + mk_deepgemm = FusedMoEModularKernel(prep_finalize, deepgemm_experts) out_deepgemm = mk_deepgemm( hidden_states=a, diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 21658f9064d7..b8e9846f864f 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -255,7 +255,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) block_shape=block_size, ) - deep_gemm_experts = mk.FusedMoEKernelModular.make_mk( + deep_gemm_experts = mk.FusedMoEKernel.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index c6abe10b9f63..06cf1b35fbb5 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -196,7 +196,7 @@ def slice_experts(): for kwargs, new_quant_config in slice_experts(): w2 = kwargs["w2"] a = kwargs["hidden_states"] - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( @@ -255,7 +255,7 @@ def run_8_bit( num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined] with_ep = num_local_experts is not None or num_local_experts == num_experts if not with_ep: - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index 7a7cc0181427..1bf5ced2e84c 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -21,7 +21,7 @@ fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel from vllm.utils.deep_gemm import ( get_mk_alignment_for_contiguous_layout, is_deep_gemm_e8m0_used, @@ -169,7 +169,7 @@ def make_ll_modular_kernel( q_dtype: torch.dtype | None, test_config: TestConfig, quant_config: FusedMoEQuantConfig, -) -> FusedMoEKernelModular: +) -> FusedMoEModularKernel: assert test_config.low_latency assert test_config.use_fp8_dispatch is not None @@ -194,7 +194,7 @@ def make_ll_modular_kernel( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk = FusedMoEKernelModular(prepare_finalize=a2a, fused_experts=fused_experts) + mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) return mk @@ -206,7 +206,7 @@ def make_ht_modular_kernel( q_dtype: torch.dtype | None, test_config: TestConfig, quant_config: FusedMoEQuantConfig, -) -> FusedMoEKernelModular: +) -> FusedMoEModularKernel: assert not test_config.low_latency assert test_config.use_fp8_dispatch is None @@ -224,7 +224,7 @@ def make_ht_modular_kernel( moe_config=make_dummy_moe_config(), quant_config=quant_config, ) - mk = FusedMoEKernelModular(prepare_finalize=a2a, fused_experts=fused_experts) + mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) return mk @@ -235,11 +235,11 @@ def make_modular_kernel( num_local_experts: int, test_tensors: TestTensors, quant_config: FusedMoEQuantConfig, -) -> FusedMoEKernelModular: +) -> FusedMoEModularKernel: q_dtype = torch.float8_e4m3fn test_config = test_tensors.config - mk: FusedMoEKernelModular + mk: FusedMoEModularKernel # Make modular kernel if test_config.low_latency: max_tokens_per_rank = max(64, next_power_of_2(test_tensors.rank_tokens.size(0))) @@ -300,7 +300,7 @@ def build_expert_map(): ) # Make modular kernel - mk: FusedMoEKernelModular = make_modular_kernel( + mk: FusedMoEModularKernel = make_modular_kernel( pg=pg, pgi=pgi, dp_size=dp_size, diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index 5c6a6af2181f..f740f5bf9585 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -19,7 +19,7 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -134,7 +134,7 @@ def make_modular_kernel( q_dtype: torch.dtype | None, use_fp8_dispatch: bool, quant_config: FusedMoEQuantConfig, -) -> FusedMoEKernelModular: +) -> FusedMoEModularKernel: ht_args: DeepEPHTArgs | None = None ll_args: DeepEPLLArgs | None = None @@ -179,7 +179,7 @@ def make_modular_kernel( quant_config=quant_config, ) - mk = FusedMoEKernelModular(prepare_finalize=a2a, fused_experts=fused_experts) + mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts) return mk @@ -237,7 +237,7 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False): ) # Make modular kernel - mk: FusedMoEKernelModular = make_modular_kernel( + mk: FusedMoEModularKernel = make_modular_kernel( pg, pgi, low_latency_mode, diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 95dfe72ffaf0..fe8f5cd007c4 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -109,7 +109,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size): block_shape=block_size, ) - deep_gemm_experts = mk.FusedMoEKernelModular.make_mk( + deep_gemm_experts = mk.FusedMoEKernel.make_mk( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 4eb270bea03b..6b435e0869a9 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -216,7 +216,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( quant_config=quant_config, ) - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEPMonolithic(), FlashInferTrtLlmFp8Experts( moe_config=td.layer.moe, @@ -318,7 +318,7 @@ def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig: routing_method=RoutingMethodType.TopK, ) - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index 3bbde1541070..9bb61ddfa0fe 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -22,7 +22,7 @@ FlashInferExperts, is_valid_flashinfer_cutlass_fused_moe, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -105,7 +105,7 @@ def test_flashinfer_fp4_moe_no_graph( routing_method=RoutingMethodType.TopK, ) - flashinfer_experts = FusedMoEKernelModular( + flashinfer_experts = FusedMoEModularKernel( MoEPrepareAndFinalizeNoEP(), FlashInferExperts(moe_config=moe_config, quant_config=quant_config), ) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index d58670e61180..38022e0e61b7 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -28,7 +28,7 @@ OAITritonExperts, UnfusedOAITritonExperts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -180,7 +180,7 @@ def oai_triton_moe_impl( else: fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config) - mk = FusedMoEKernelModular(MoEPrepareAndFinalizeNoEP(), fused_experts) + mk = FusedMoEModularKernel(MoEPrepareAndFinalizeNoEP(), fused_experts) return mk.forward( hidden_states=x, diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index b531bd72dcd9..30fabe2ddb18 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -89,7 +89,7 @@ def test_cutlass_fp4_moe_no_graph( w2_scale=w2_blockscale, ) - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_pplx_cutlass_moe.py b/tests/kernels/moe/test_pplx_cutlass_moe.py index 9b30e1fc858f..ef37c1c74434 100644 --- a/tests/kernels/moe/test_pplx_cutlass_moe.py +++ b/tests/kernels/moe/test_pplx_cutlass_moe.py @@ -16,7 +16,7 @@ fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassBatchedExpertsFp8 -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import set_random_seed @@ -169,7 +169,7 @@ def make_moe_config() -> FusedMoEConfig: num_dispatchers=num_dispatchers, ) - fused_cutlass_experts = FusedMoEKernelModular( + fused_cutlass_experts = FusedMoEModularKernel( prepare_finalize, experts, ) diff --git a/tests/kernels/moe/test_pplx_moe.py b/tests/kernels/moe/test_pplx_moe.py index fa885fde0384..08519087e1ce 100644 --- a/tests/kernels/moe/test_pplx_moe.py +++ b/tests/kernels/moe/test_pplx_moe.py @@ -41,7 +41,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceDelegate, ) @@ -588,7 +588,7 @@ def pplx_moe( moe_config=make_dummy_moe_config(), ) - fused_experts = FusedMoEKernelModular( + fused_experts = FusedMoEModularKernel( prepare_finalize, experts, shared_experts, diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 4f4ec1977ec8..4883085cb836 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -19,7 +19,7 @@ BatchedTritonExperts, NaiveBatchedExperts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernelModular +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.deep_gemm import per_block_cast_to_fp8 from vllm.utils.math_utils import round_up @@ -106,7 +106,7 @@ def batched_moe( a2_scale=a2_scale, ) - fused_experts = FusedMoEKernelModular( + fused_experts = FusedMoEModularKernel( BatchedPrepareAndFinalize( max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 ), @@ -147,7 +147,7 @@ def naive_batched_moe( a2_scale=a2_scale, ) - fused_experts = FusedMoEKernelModular( + fused_experts = FusedMoEModularKernel( BatchedPrepareAndFinalize( max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 ), diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index e6722925842d..be1fd7cdb663 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -32,7 +32,7 @@ UnfusedOAITritonExperts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEKernelModular, + FusedMoEModularKernel, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, @@ -131,7 +131,7 @@ def _inject_lora_into_fused_moe(self): quant_config = self.base_layer.quant_method.moe_quant_config prepare_finalize = MoEPrepareAndFinalizeNoEP() - m_fused_moe_fn = FusedMoEKernelModular( + m_fused_moe_fn = FusedMoEModularKernel( prepare_finalize, self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index fcc7a4042fa8..238c7b73b635 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -1149,7 +1149,7 @@ def cutlass_moe_w4a8_fp8( num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0) - fn = mk.FusedMoEKernelModular.make_mk( + fn = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), CutlassExpertsW4A8Fp8( out_dtype=a.dtype, diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 672101ee1b13..17093794903a 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -2297,8 +2297,8 @@ def modular_triton_fused_moe( moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None, -) -> mk.FusedMoEKernelModular: - return mk.FusedMoEKernelModular.make_mk( +) -> mk.FusedMoEModularKernel: + return mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), TritonExperts(moe_config, quant_config), shared_experts, diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 3060102791c3..093b9d2e2d4e 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -13,7 +13,7 @@ FusedMoEMethodBase, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEKernelModular, + FusedMoEModularKernel, FusedMoEPrepareAndFinalizeModular, ) @@ -26,7 +26,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): # --8<-- [end:modular_fused_moe] def __init__( - self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEKernelModular + self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEModularKernel ): super().__init__(old_quant_method.moe) self.moe_quant_config = old_quant_method.moe_quant_config @@ -49,7 +49,7 @@ def make( ) -> "FusedMoEModularMethod": return FusedMoEModularMethod( old_quant_method, - FusedMoEKernelModular( + FusedMoEModularKernel( prepare_finalize, old_quant_method.select_gemm_impl(prepare_finalize, moe_layer), shared_experts, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 9212b03cc2e7..718c77968e34 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -63,7 +63,7 @@ # Some FusedMoEExpertsModular implementations may choose to do # the weight application and/or reduction. The class communicates this # to [Finalize] via a TopKWeightAndReduce object. -# * FusedMoEKernelModular - an interface class that combines a +# * FusedMoEModularKernel - an interface class that combines a # FusedMoEPrepareAndFinalizeModular and a FusedMoEExpertsModular to # provide the standard fused MoE kernel interface. # * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen @@ -986,7 +986,7 @@ def make_mk( if isinstance( prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): - return FusedMoEKernelMonolithic( + return FusedMoEKMonolithicKernel( prepare_finalize, fused_experts, shared_experts, @@ -995,7 +995,7 @@ def make_mk( elif isinstance( prepare_finalize, FusedMoEPrepareAndFinalizeModular ) and isinstance(fused_experts, FusedMoEExpertsModular): - return FusedMoEKernelModular( + return FusedMoEModularKernel( prepare_finalize, fused_experts, shared_experts, @@ -1030,7 +1030,7 @@ def output_is_reduced(self) -> bool: @final -class FusedMoEKernelModular(FusedMoEKernel): +class FusedMoEModularKernel(FusedMoEKernel): """ This class combines a FusedMoEPrepareAndFinalizeModular instance and a FusedMoEExpertsModular to provide an interface that @@ -1092,7 +1092,7 @@ def _allocate_buffers( workspace_dtype = self.fused_experts.workspace_dtype(out_dtype) # Force worst-case allocation in profiling run for - # "mk.FusedMoEKernelModular.Standard" formats where this is only bounded + # "mk.FusedMoEModularKernel.Standard" formats where this is only bounded # by `VLLM_FUSED_MOE_CHUNK_SIZE` and may not be seen during profiling with # DP+EP due to the random token routing. is_profile_run = ( @@ -1563,7 +1563,7 @@ def forward( @final -class FusedMoEKernelMonolithic(FusedMoEKernel): +class FusedMoEKMonolithicKernel(FusedMoEKernel): def forward( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 4c1cb2045082..44b24391ef31 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -405,7 +405,7 @@ def make_nvfp4_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( prepare_finalize, experts, shared_experts=( diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index d3aacee0e262..af6fd3994955 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -155,7 +155,7 @@ def make_unquantized_moe_kernel( backend: UnquantizedMoeBackend, quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, -) -> tuple[mk.FusedMoEKernelModular | None, bool]: +) -> tuple[mk.FusedMoEModularKernel | None, bool]: use_inplace = True if backend in UNSUPPORTED_BACKEND: @@ -166,7 +166,7 @@ def make_unquantized_moe_kernel( FlashInferExperts, ) - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, @@ -179,7 +179,7 @@ def make_unquantized_moe_kernel( AiterExperts, ) - kernel = mk.FusedMoEKernelModular.make_mk( + kernel = mk.FusedMoEKernel.make_mk( MoEPrepareAndFinalizeNoEP(), AiterExperts( moe_config=moe_config, diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index fb6a82724c0f..6332827d1d09 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -62,7 +62,7 @@ def eplb_map_to_physical_and_record( # 2. Record expert load metrics. - # TODO(bowen): When using `FusedMoEKernelModular`, this + # TODO(bowen): When using `FusedMoEModularKernel`, this # can be done in a more unified way, since # `FusedMoEPrepareAndFinalizeModular` will return the expert # token count, in some cases directly from the kernel. diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index b8b1a06aa8bc..7133a1c98db1 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -607,7 +607,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) + assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) return self.moe_kernel( x, layer.w13_weight, @@ -979,7 +979,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) + assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) return self.moe_kernel( x, layer.w13_weight, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 90b1d4918597..5965ec9161cd 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -938,7 +938,7 @@ def apply_monolithic( if layer.enable_eplb: raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.") - assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) + assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) return self.moe_kernel( x, layer.w13_weight, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e1fe956a2ec2..37198fec3717 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -913,7 +913,7 @@ def apply_monolithic( raise NotImplementedError( "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend." ) - assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) + assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) return self.moe_kernel( x, layer.w13_weight, @@ -1437,7 +1437,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEKernelMonolithic) + assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) return self.moe_kernel( x, layer.w13_weight, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index c5165027ad52..3306de956be9 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -747,7 +747,7 @@ def _interleave_mxfp4_cutlass_sm90(w): layer.w13_bias = Parameter(w13_bias, requires_grad=False) layer.w2_bias = Parameter(w2_bias, requires_grad=False) - # Ideally we'd use FusedMoEKernelModular.prepare_finalize object + # Ideally we'd use FusedMoEModularKernel.prepare_finalize object # (stored in self.fused_experts) to determine if the MoE has a # batched activation format. As self.fused_experts is not # initialized at this point, we resort to checking the MoE config From caa511df0b13c8d9d2f3d017433dbdf11ab0c105 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Feb 2026 20:24:49 -0500 Subject: [PATCH 109/207] fix up naive EP issue Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/all2all_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 205edfee3e30..cb1cac390594 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -112,6 +112,7 @@ def maybe_make_prepare_finalize( num_dispatchers=( get_ep_group().device_communicator.all2all_manager.world_size ), + use_monolithic=use_monolithic, ) else: return MoEPrepareAndFinalizeNoEPBase.make(use_monolithic) From 5fc36098f029ae7cee1cb9a1e56223af123a99bf Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Feb 2026 22:30:32 -0500 Subject: [PATCH 110/207] remove non-tuping Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 44b24391ef31..9115d5b2cb78 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -113,7 +113,7 @@ def select_nvfp4_moe_backend( config: FusedMoEConfig, weight_key: QuantKey | None, activation_key: QuantKey | None, -) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts] | None]: +) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts]]: """ Select the primary NvFP4 MoE backend Note: Shape-specific fallbacks may still occur at runtime. @@ -202,7 +202,7 @@ def _return_or_raise( ) if supported: logger.info_once(_make_log_backend(backend), scope="local") - return backend, None + return backend, k_cls else: logger.debug_once( _make_log_unsupported(backend, reason), scope="local" From 0e718bbd936146f9f81858b470a0e27ff3ee7b33 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 13:07:18 -0500 Subject: [PATCH 111/207] updated moe_mk -> moe_kernel Signed-off-by: Robert Shaw --- .../layers/fused_moe/runner/default_moe_runner.py | 4 ++-- .../compressed_tensors/compressed_tensors_moe.py | 8 ++++---- vllm/model_executor/layers/quantization/fp8.py | 2 +- vllm/model_executor/layers/quantization/modelopt.py | 8 ++++---- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index e68d35b31f04..7da36592e3db 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -308,8 +308,8 @@ def must_reduce_shared_expert_outputs(self) -> bool: """ assert self.quant_method is not None return ( - self.quant_method.moe_mk is not None - and self.quant_method.moe_mk.output_is_reduced() + self.quant_method.moe_kernel is not None + and self.quant_method.moe_kernel.output_is_reduced() ) def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index efabdc6eac48..f08be2313daa 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -610,7 +610,7 @@ def apply_monolithic( router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) - return self.moe_mk( + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -633,8 +633,8 @@ def apply( topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -985,7 +985,7 @@ def apply_monolithic( router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) - return self.moe_mk( + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index c134561e54dd..28539504f551 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -838,7 +838,7 @@ def _setup_kernel( self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config: assert self.experts_cls is not None - self.moe_mk = make_fp8_moe_kernel( + self.moe_kernel = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, fp8_backend=self.fp8_backend, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index a64e9f74412d..6a9c42150026 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -946,8 +946,8 @@ def apply( topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, @@ -1471,8 +1471,8 @@ def apply( topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_mk is not None - return self.moe_mk( + assert self.moe_kernel is not None + return self.moe_kernel( x, layer.w13_weight, layer.w2_weight, From af6d9954972078859a97a4f8219cdd2f8cae567d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 13:22:06 -0500 Subject: [PATCH 112/207] fix moe activation pre-commit Signed-off-by: Robert Shaw --- vllm/lora/layers/fused_moe.py | 2 +- .../layers/fused_moe/flashinfer_trtllm_fp8_moe.py | 7 ++++--- .../layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py | 7 ++++--- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index ed33452bf55d..371cacfa2b12 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -132,7 +132,7 @@ def _inject_lora_into_fused_moe(self): if getattr(self.base_layer.quant_method, "supports_internal_mk", False): # Use the existing modular kernel from the quant method - m_fused_moe_fn = self.base_layer.quant_method.moe_mk + m_fused_moe_fn = self.base_layer.quant_method.moe_kernel else: # Create a new modular kernel via select_gemm_impl prepare_finalize = MoEPrepareAndFinalizeNoEP() diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index beb97478201f..8579ec079be6 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -5,6 +5,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, @@ -91,9 +92,9 @@ def _supports_quant_scheme( return (weight_key, activation_key) in SUPPORTED_W_A @staticmethod - def _supports_activation(activation: str) -> bool: - """Supports silu activation only.""" - return activation in ["silu"] + def _supports_activation(activation: MoEActivation) -> bool: + """Supports only SiLU and RELU^2 non-gated activation.""" + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] @staticmethod def _supports_routing_method( diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 03e335478b95..600eee7fbd89 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -5,6 +5,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, @@ -73,9 +74,9 @@ def _supports_quant_scheme( return (weight_key, activation_key) in SUPPORTED_W_A @staticmethod - def _supports_activation(activation: str) -> bool: - """Supports only SiLU activation.""" - return activation in ["silu"] + def _supports_activation(activation: MoEActivation) -> bool: + """Supports only SiLU and RELU^2 non-gated activation.""" + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: From 62471c062cb03ca33da9376eca056c407cebe219 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 13:30:10 -0500 Subject: [PATCH 113/207] clean up Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 13 +++++++++--- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 20 ++++-------------- .../layers/fused_moe/modular_kernel.py | 21 +++++++++++++++++-- 3 files changed, 33 insertions(+), 21 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 8579ec079be6..c954d4af411a 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -103,16 +103,23 @@ def _supports_routing_method( activation_key: QuantKey | None, ) -> bool: """Monolithic kernels need to express router support.""" - # NOTE(rob): potentially allow others here. This is a conservative list. + # NOTE(dbari): TopK routing could also be enabled, but need to validate models + # NOTE(dbari): Default is not implemented and should not be enabled until it is if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): + # NOTE(rob): potentially allow others here. This is a conservative list. return routing_method in [ RoutingMethodType.DeepSeekV3, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, ] elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym): - return routing_method == RoutingMethodType.Llama4 - + # NOTE(dbari): as above, potentially allow others here. + return routing_method in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Llama4, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] else: raise ValueError("Unsupported quantization scheme.") diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index 600eee7fbd89..d977b0fefb2e 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -23,7 +23,7 @@ from vllm.platforms import current_platform -class FlashInferTrtLlmNvFp4ExpertsBase(mk.FusedMoEExperts): +class FlashInferTrtLlmNvFp4ExpertsBase: """ NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface. """ @@ -33,7 +33,8 @@ def __init__( moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, ): - super().__init__(moe_config=moe_config, quant_config=quant_config) + self.moe_config = moe_config + self.quant_config = quant_config self.routing_method_type = self.moe_config.routing_method self.topk = moe_config.experts_per_token @@ -59,7 +60,7 @@ def _supports_current_device() -> bool: @staticmethod def _supports_no_act_and_mul() -> bool: - """Does not support non-gated MoE (i.e. Nemotron-Nano).""" + """Supports non-gated MoE (i.e. Nemotron-Nano).""" return True @staticmethod @@ -96,19 +97,6 @@ class FlashInferTrtLlmNvFp4ExpertsModular( Modular version of the implementation (just the experts). """ - @staticmethod - def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """Supports EP and TP.""" - return True - - @staticmethod - def _supports_routing_method( - routing_method_type: RoutingMethodType, - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - return True - def workspace_shapes( self, M: int, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index e0d350a881c8..06ebb92e39ba 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -575,7 +575,10 @@ def _supports_activation(activation: MoEActivation) -> bool: @abstractmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: """ - Whether the kernel supports deployment in expert parallel. + Whether the kernel supports deployment in particular parallel config. + + Can be overriden if a kernel does not support EP, SP or some other + configuration. """ raise NotImplementedError @@ -591,7 +594,7 @@ def _supports_routing_method( Can be overriden by monolithic kernels that execute the router in addition to the experts if certain routers are not supported. """ - return True + raise NotImplementedError # # Various helpers for accessing quantization parameters from the @@ -697,6 +700,20 @@ class FusedMoEExpertsModular(FusedMoEExperts): above. """ + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """ + Whether the kernel supports a routing method (e.g. GroupedTopK). + + Modular kernels support all routing methods, since the Expert + kernel does not apply the activation. + """ + return True + @staticmethod def is_monolithic() -> bool: return False From 15f4576e074e2888b10793576d1ea851370771f9 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 13:38:49 -0500 Subject: [PATCH 114/207] updated Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 3 + .../layers/fused_moe/flashinfer_trtllm_moe.py | 16 ----- .../layers/fused_moe/modular_kernel.py | 59 ++++++++++++++----- 3 files changed, 46 insertions(+), 32 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index c954d4af411a..20ce58cf7e63 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -131,6 +131,9 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo and not moe_parallel_config.enable_eplb ) + # @staticmethod + # def _supports_router_logits_dtype() + def supports_chunking(self) -> bool: return False diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py index ffab43ca0824..c064526efbb9 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py @@ -50,22 +50,6 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo return not moe_parallel_config.enable_eplb -def _supports_router_logits_dtype( - router_logits_dtype: torch.dtype | None, - routing_method: RoutingMethodType, -) -> bool: - """ - The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default. - Only DeepSeekV3 routing supports float32 router_logits (which is converted - internally in the kernel). - """ - if router_logits_dtype == torch.float32: - # Only DeepSeekV3 routing handles float32 logits - # https://github.com/flashinfer-ai/flashinfer/issues/2469 - return routing_method == RoutingMethodType.DeepSeekV3 - return True - - def is_supported_config_trtllm_bf16( moe_config: FusedMoEConfig, activation_format: mk.FusedMoEActivationFormat, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 06ebb92e39ba..2f9fd848eb30 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -533,6 +533,11 @@ def _make_reason(reason: str) -> str: moe_config.routing_method, weight_key, activation_key ): return False, _make_reason("routing method") + elif not cls._supports_router_logits_dtype( + moe_config.router_logits_dtype, + moe_config.routing_method, + ): + return False, _make_reason("router logits dtype") elif activation_format != cls.activation_format(): return False, _make_reason(f"{activation_format.value} activation format") return True, None @@ -594,7 +599,19 @@ def _supports_routing_method( Can be overriden by monolithic kernels that execute the router in addition to the experts if certain routers are not supported. """ - raise NotImplementedError + return True + + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + Whether a kernel supports a particular dtype for router logits input. + + Can be overriden by monolithic kernels that execute the router + in addition to the experts if certain dtypes are not supported. + """ + return True # # Various helpers for accessing quantization parameters from the @@ -700,20 +717,6 @@ class FusedMoEExpertsModular(FusedMoEExperts): above. """ - @staticmethod - def _supports_routing_method( - routing_method: RoutingMethodType, - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - """ - Whether the kernel supports a routing method (e.g. GroupedTopK). - - Modular kernels support all routing methods, since the Expert - kernel does not apply the activation. - """ - return True - @staticmethod def is_monolithic() -> bool: return False @@ -901,9 +904,33 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts): rather than topk ids and weights). """ + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """ + Whether the kernel supports a routing method (e.g. GroupedTopK). + + Monolithic kernels should explicitly opt-in to support. + """ + raise NotImplementedError + + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + Whether the kernel supports a dtype for router logits. + + Modular kernels should opt-in to support. + """ + raise NotImplementedError + @staticmethod def is_monolithic() -> bool: - return False + return True def apply( self, From 7e78317eb88e3f0963d104c4d8e6ee966f1081d7 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 13:50:38 -0500 Subject: [PATCH 115/207] updated Signed-off-by: Robert Shaw --- .../fused_moe/flashinfer_trtllm_fp8_moe.py | 27 ++++++++++++++----- .../fused_moe/flashinfer_trtllm_nvfp4_moe.py | 24 +++++++++++++++-- 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index 20ce58cf7e63..fcd5120f78db 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -125,14 +125,27 @@ def _supports_routing_method( @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """TRTLLMGenKernel is monolithic, so it only supports TP or naive DP/EP.""" - return not moe_parallel_config.use_all2all_kernels or ( - moe_parallel_config.use_naive_all2all_kernels - and not moe_parallel_config.enable_eplb - ) + """Monolithic kernel so only use with naive DP/EP and TP.""" + return ( + not moe_parallel_config.use_all2all_kernels + or moe_parallel_config.use_naive_all2all_kernels + ) and not moe_parallel_config.enable_eplb - # @staticmethod - # def _supports_router_logits_dtype() + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default. + Only DeepSeekV3 routing supports float32 router_logits (which is converted + internally in the kernel). + """ + if router_logits_dtype == torch.float32: + # Only DeepSeekV3 routing handles float32 logits + # https://github.com/flashinfer-ai/flashinfer/issues/2469 + return routing_method == RoutingMethodType.DeepSeekV3 + return True def supports_chunking(self) -> bool: return False diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index d977b0fefb2e..ccd52b094777 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -191,8 +191,11 @@ class FlashInferTrtLlmNvFp4ExpertsMonolithic( @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """The modular implementation should be used for the Dp/Ep case""" - return not moe_parallel_config.use_all2all_kernels + """The modular implementation should be used for the Dp/Ep or EPLB case.""" + return ( + not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + ) @staticmethod def _supports_routing_method( @@ -208,6 +211,23 @@ def _supports_routing_method( RoutingMethodType.Llama4, ] + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + The FlashInfer TRTLLM NVFp4 kernel expects bfloat16 router_logits by default. + Only DeepSeekV3 routing supports float32 router_logits (which is converted + internally in the kernel). + """ + # TODO: check this + if router_logits_dtype == torch.float32: + # Only DeepSeekV3 routing handles float32 logits + # https://github.com/flashinfer-ai/flashinfer/issues/2469 + return routing_method == RoutingMethodType.DeepSeekV3 + return True + def apply( self, hidden_states: torch.Tensor, From 691f84c804de27d9ae47f452cb7e66505cd3525a Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 16:23:55 -0500 Subject: [PATCH 116/207] convert to composition over inheritance Signed-off-by: Robert Shaw --- .../kernels/benchmark_cutlass_moe_fp8.py | 2 +- .../kernels/benchmark_cutlass_moe_nvfp4.py | 4 +- .../kernels/benchmark_grouped_gemm_cutlass.py | 4 +- benchmarks/kernels/benchmark_moe.py | 2 +- .../moe/modular_kernel_tools/common.py | 4 +- tests/kernels/moe/test_block_fp8.py | 2 +- tests/kernels/moe/test_cutlass_moe.py | 4 +- tests/kernels/moe/test_deepgemm.py | 2 +- tests/kernels/moe/test_flashinfer.py | 4 +- tests/kernels/moe/test_nvfp4_moe.py | 2 +- .../layers/fused_moe/cutlass_moe.py | 2 +- .../layers/fused_moe/modular_kernel.py | 266 ++++++++++-------- .../layers/fused_moe/oracle/fp8.py | 2 +- .../layers/fused_moe/oracle/nvfp4.py | 2 +- .../layers/fused_moe/oracle/unquantized.py | 10 +- .../fused_moe/unquantized_fused_moe_method.py | 2 +- 16 files changed, 173 insertions(+), 141 deletions(-) diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index 3f3ca817bcb7..92032efb94a9 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -137,7 +137,7 @@ def bench_run( per_out_ch_quant=per_out_ch, ) - fn = mk.FusedMoEKernel.make_mk( + fn = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index ebd79f8705ce..e0b6cb6d8210 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -196,7 +196,7 @@ def run_cutlass_moe_fp4( g2_alphas=w2_gs, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( make_dummy_moe_config(), @@ -241,7 +241,7 @@ def run_cutlass_from_graph( g2_alphas=w2_gs, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( make_dummy_moe_config(), diff --git a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py index 7a8f8be4fa18..fd6c5216420e 100644 --- a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py +++ b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py @@ -132,7 +132,7 @@ def run_cutlass_moe( per_act_token_quant=per_act_token, ) - fn = mk.FusedMoEKernel.make_mk( + fn = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( @@ -164,7 +164,7 @@ def run_cutlass_from_graph( per_act_token_quant=per_act_token, ) - fn = mk.FusedMoEKernel.make_mk( + fn = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 995be39649e7..064d72f2884c 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -242,7 +242,7 @@ def run(): deep_gemm_experts = None if use_deep_gemm: - deep_gemm_experts = mk.FusedMoEKernel.make_mk( + deep_gemm_experts = mk.FusedMoEKernel( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=FusedMoEConfig( diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 3fcbc116c796..0c0699edb17a 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -573,7 +573,7 @@ def make_modular_kernel( config: Config, vllm_config: VllmConfig, quant_config: FusedMoEQuantConfig, -) -> mk.FusedMoEModularKernel: +) -> mk.FusedMoEKernel: def next_power_of_2(x): import math @@ -620,7 +620,7 @@ def next_power_of_2(x): config.N, ) - modular_kernel = mk.FusedMoEKernel.make_mk( + modular_kernel = mk.FusedMoEKernel( prepare_finalize=prepare_finalize, fused_experts=fused_experts, inplace=False, diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 2f0ea248ff43..3b00cdd7326f 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -253,7 +253,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) block_shape=block_size, ) - deep_gemm_experts = mk.FusedMoEKernel.make_mk( + deep_gemm_experts = mk.FusedMoEKernel( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 0adec5c044a7..65887313d6d3 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -197,7 +197,7 @@ def slice_experts(): for kwargs, new_quant_config in slice_experts(): w2 = kwargs["w2"] a = kwargs["hidden_states"] - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( @@ -257,7 +257,7 @@ def run_8_bit( num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined] with_ep = num_local_experts is not None or num_local_experts == num_experts if not with_ep: - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index a7d311ef1d06..eedffae7ead5 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -109,7 +109,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size): block_shape=block_size, ) - deep_gemm_experts = mk.FusedMoEKernel.make_mk( + deep_gemm_experts = mk.FusedMoEKernel( prepare_finalize=MoEPrepareAndFinalizeNoEP(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index a5d6185790ce..a486e0c884b4 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -240,7 +240,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( quant_config=quant_config, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEPMonolithic(), FlashInferTrtLlmFp8Experts( moe_config=td.layer.moe, @@ -347,7 +347,7 @@ def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig: routing_method=RoutingMethodType.TopK, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index 0a7c9965ddf3..abf0a737480a 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -89,7 +89,7 @@ def test_cutlass_fp4_moe_no_graph( w2_scale=w2_blockscale, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsFp4( moe_config=make_dummy_moe_config(), diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index acc083c67fcc..ef9937befba7 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -1166,7 +1166,7 @@ def cutlass_moe_w4a8_fp8( num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0) - fn = mk.FusedMoEKernel.make_mk( + fn = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), CutlassExpertsW4A8Fp8( out_dtype=a.dtype, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 2f9fd848eb30..12cb8137fedf 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -968,129 +968,32 @@ def _slice_scales( return None -class FusedMoEKernel(torch.nn.Module): +################################################################################ +# TODO: make the below a separate file. +################################################################################ + + +@final +class FusedMoEKernelModularImpl: def __init__( self, - prepare_finalize: FusedMoEPrepareAndFinalize, - fused_experts: FusedMoEExperts, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, + fused_experts: FusedMoEExpertsModular, shared_experts: torch.nn.Module | None = None, moe_parallel_config: FusedMoEParallelConfig | None = None, inplace: bool = False, ): - super().__init__() self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts self.shared_experts = shared_experts + self.moe_parallel_config = moe_parallel_config self.inplace = inplace - - # prefer an explicit FusedMoEParallelConfig when available (from - # FusedMoE layers / tests). - # if not provided, assume this kernel is - # running in a non-DP+EP context - self.moe_parallel_config: FusedMoEParallelConfig | None = moe_parallel_config self.is_dp_ep = ( moe_parallel_config is not None and moe_parallel_config.dp_size > 1 and moe_parallel_config.use_ep ) - # Confirm P/F and Experts kernels are consistent with eachother. - if not ( - ( - isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) - and isinstance(fused_experts, FusedMoEExpertsMonolithic) - ) - or ( - isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular) - and isinstance(fused_experts, FusedMoEExpertsModular) - ) - ): - raise ValueError( - "prepare_finalize and fused_experts must both be either monolithic " - f"or non-monolithic but got {prepare_finalize.__class__.__name__} " - "and {fused_experts.__class__.__name__}" - ) - - self._post_init_setup() - assert ( - prepare_finalize.activation_format == fused_experts.activation_format() - ), ( - f"{prepare_finalize.__class__.__name__}." - f"{prepare_finalize.activation_format} == " - f"{fused_experts.__class__.__name__}." - f"{fused_experts.activation_format()}" - ) - - @staticmethod - def make_mk( - prepare_finalize: FusedMoEPrepareAndFinalize, - fused_experts: FusedMoEExperts, - shared_experts: torch.nn.Module | None = None, - moe_parallel_config: FusedMoEParallelConfig | None = None, - ) -> "FusedMoEKernel": - """ - Factory method to create a FusedMoEKernel instance. - """ - if isinstance( - prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic - ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): - return FusedMoEKMonolithicKernel( - prepare_finalize, - fused_experts, - shared_experts, - moe_parallel_config, - ) - elif isinstance( - prepare_finalize, FusedMoEPrepareAndFinalizeModular - ) and isinstance(fused_experts, FusedMoEExpertsModular): - return FusedMoEModularKernel( - prepare_finalize, - fused_experts, - shared_experts, - moe_parallel_config, - ) - else: - raise ValueError( - "prepare_finalize and fused_experts must both be either monolithic " - f"or non-monolithic but got {prepare_finalize.__class__.__name__} " - "and {fused_experts.__class__.__name__}" - ) - - def _post_init_setup(self): - """ - Resolve any leftover setup dependencies between self.prepare_finalize - and self.fused_experts here. - """ - self.prepare_finalize.post_init_setup(self.fused_experts) - - def supports_expert_map(self) -> bool: - """ - A flag indicating whether or not this class supports expert maps. - """ - return self.fused_experts.supports_expert_map() - - def output_is_reduced(self) -> bool: - """ - Indicates whether or not the output of fused MoE kernel - is reduced across all ranks. - """ - return self.prepare_finalize.output_is_reduced() - - -@final -class FusedMoEModularKernel(FusedMoEKernel): - """ - This class combines a FusedMoEPrepareAndFinalizeModular instance and - a FusedMoEExpertsModular to provide an interface that - is compatible with the `fused_experts` function in fused_moe.py. - - It takes care of managing any required scratch space. - - Note: Instances of this class should only be used for a single model - layer due to any layer specific state that may be used by the component - objects. - """ - def _chunk_info(self, M: int) -> tuple[int, int]: """ Compute number of chunks and chunk size for given M. @@ -1140,7 +1043,7 @@ def _allocate_buffers( workspace_dtype = self.fused_experts.workspace_dtype(out_dtype) # Force worst-case allocation in profiling run for - # "mk.FusedMoEModularKernel.Standard" formats where this is only bounded + # "mk.FusedMoEKernel.Standard" formats where this is only bounded # by `VLLM_FUSED_MOE_CHUNK_SIZE` and may not be seen during profiling with # DP+EP due to the random token routing. is_profile_run = ( @@ -1540,13 +1443,13 @@ def _finalize( assert shared_output is not None return shared_output, output - def forward( + def apply( self, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - topk_weights: torch.Tensor, topk_ids: torch.Tensor, + topk_weights: torch.Tensor, activation: MoEActivation = MoEActivation.SILU, global_num_experts: int = -1, expert_map: torch.Tensor | None = None, @@ -1561,8 +1464,7 @@ def forward( - hidden_states: (torch.Tensor): The input tensor to the MoE layer. - w1 (torch.Tensor): The first set of expert weights. - w2 (torch.Tensor): The second set of expert weights. - - topk_weights (torch.Tensor): The topk weights applied at the end of - the layer. + - topk_weights (torch.Tensor): The topk weights applied at the end of the layer. - topk_ids (torch.Tensor): A map of row to expert id. - activation (MoEActivation): The activation function to apply after the first MoE layer. @@ -1581,7 +1483,6 @@ def forward( Returns: - torch.Tensor: The output tensor after applying the MoE layer. """ - if self.inplace: assert self.shared_experts is None assert not disable_inplace() @@ -1630,8 +1531,16 @@ def forward( @final -class FusedMoEKMonolithicKernel(FusedMoEKernel): - def forward( +class FusedMoEKernelMonolithicImpl: + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic, + fused_experts: FusedMoEExpertsMonolithic, + ): + self.prepare_finalize = prepare_finalize + self.fused_experts = fused_experts + + def apply( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -1652,8 +1561,6 @@ def forward( to the topk_ids and topk_weights. This is used for kernels that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ - assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) - assert isinstance(self.fused_experts, FusedMoEExpertsMonolithic) # TODO(rob): add inplace support. a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( @@ -1683,3 +1590,128 @@ def forward( output = self.prepare_finalize.finalize(fused_out) return output + + +@final +class FusedMoEKernel(torch.nn.Module): + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalize, + fused_experts: FusedMoEExperts, + shared_experts: torch.nn.Module | None = None, + moe_parallel_config: FusedMoEParallelConfig | None = None, + inplace: bool = False, + ): + super().__init__() + self.shared_experts = shared_experts # NOTE: check if we can remove + + # Initialize the implementation (monolithic or modular). + self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl + if isinstance( + prepare_finalize, FusedMoEPrepareAndFinalizeModular + ) and isinstance(fused_experts, FusedMoEExpertsModular): + self.impl = FusedMoEKernelModularImpl( + prepare_finalize, + fused_experts, + shared_experts, + moe_parallel_config, + inplace, + ) + + elif isinstance( + prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic + ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): + assert shared_experts is None + assert not inplace + self.impl = FusedMoEKernelMonolithicImpl( + prepare_finalize, + fused_experts, + ) + + else: + raise ValueError( + "prepare_finalize and fused_experts must both be either monolithic " + f"or non-monolithic but got {prepare_finalize.__class__.__name__} " + "and {fused_experts.__class__.__name__}" + ) + + self._post_init_setup() + + def _post_init_setup(self): + """ + Resolve any leftover setup dependencies between self.prepare_finalize + and self.fused_experts here. + """ + self.impl.prepare_finalize.post_init_setup(self.impl.fused_experts) + assert ( + self.impl.prepare_finalize.activation_format + == self.impl.fused_experts.activation_format() + ) + + def supports_expert_map(self) -> bool: + """ + A flag indicating whether or not this class supports expert maps. + """ + return self.impl.fused_experts.supports_expert_map() + + def output_is_reduced(self) -> bool: + """ + Indicates whether or not the output of fused MoE kernel + is reduced across all ranks. + """ + return self.impl.prepare_finalize.output_is_reduced() + + def forward( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + routing_input: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + if isinstance(self.impl, FusedMoEKernelModularImpl): + assert isinstance(routing_input, tuple) + topk_ids, topk_weights = routing_input + assert num_expert_group is None + assert e_score_correction_bias is None + assert routed_scaling_factor is None + assert topk_group is None + + return self.impl.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + elif isinstance(self.impl, FusedMoEKernelMonolithicImpl): + assert isinstance(routing_input, torch.Tensor) + router_logits = routing_input + + return self.impl.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + router_logits=router_logits, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + topk_group=topk_group, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 127cc12a299d..ed1a57e0c2c0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -473,7 +473,7 @@ def make_fp8_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - return mk.FusedMoEKernel.make_mk( + return mk.FusedMoEKernel( prepare_finalize, experts, shared_experts=( diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 47b7f1bb28dd..97232e9bdba7 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -405,7 +405,7 @@ def make_nvfp4_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - return mk.FusedMoEKernel.make_mk( + return mk.FusedMoEKernel( prepare_finalize, experts, shared_experts=( diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index e93df2af56dc..6a7d783d4208 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -165,7 +165,7 @@ def make_unquantized_moe_kernel( backend: UnquantizedMoeBackend, quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, -) -> mk.FusedMoEModularKernel | None: +) -> mk.FusedMoEKernel | None: if backend in UNSUPPORTED_BACKEND: return None @@ -174,7 +174,7 @@ def make_unquantized_moe_kernel( FlashInferExperts, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), FlashInferExperts( moe_config=moe_config, @@ -187,7 +187,7 @@ def make_unquantized_moe_kernel( AiterExperts, ) - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), AiterExperts( moe_config=moe_config, @@ -197,7 +197,7 @@ def make_unquantized_moe_kernel( elif backend == UnquantizedMoeBackend.TRITON: from vllm.model_executor.layers.fused_moe import TritonExperts - kernel = mk.FusedMoEKernel.make_mk( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), TritonExperts( moe_config=moe_config, @@ -207,7 +207,7 @@ def make_unquantized_moe_kernel( elif backend == UnquantizedMoeBackend.XPU: from vllm.model_executor.layers.fused_moe import XPUExperts - kernel = mk.FusedMoEModularKernel( + kernel = mk.FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), XPUExperts( moe_config=moe_config, diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index d3c41deb8fe3..532a5375a00b 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -70,7 +70,7 @@ def __init__(self, moe: FusedMoEConfig): self.rocm_aiter_moe_enabled = ( rocm_aiter_ops.is_fused_moe_enabled() and moe.is_act_and_mul ) - self.kernel: mk.FusedMoEModularKernel | None = None + self.kernel: mk.FusedMoEKernel | None = None self._is_monolithic = ( current_platform.is_cpu() or self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM From b62f64174fe3fa6921e297d8f6d99de08e3be1c8 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 16:28:17 -0500 Subject: [PATCH 117/207] revert inplace changes Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/oracle/unquantized.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 6a7d783d4208..1ae3d1a01d94 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -180,6 +180,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + inplace=False, ) elif backend == UnquantizedMoeBackend.AITER: @@ -193,6 +194,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + inplace=not moe_config.disable_inplace, ) elif backend == UnquantizedMoeBackend.TRITON: from vllm.model_executor.layers.fused_moe import TritonExperts @@ -203,6 +205,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + inplace=not moe_config.disable_inplace, ) elif backend == UnquantizedMoeBackend.XPU: from vllm.model_executor.layers.fused_moe import XPUExperts @@ -213,5 +216,6 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + inplace=not moe_config.disable_inplace, ) return kernel From 66ea6bd6172c4587f1d7f1495f2d671db7d1a50d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 16:29:43 -0500 Subject: [PATCH 118/207] revert inplace changes Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/oracle/fp8.py | 8 +++++++- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 6 +++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index ed1a57e0c2c0..d07e35df8100 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -473,7 +473,7 @@ def make_fp8_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - return mk.FusedMoEKernel( + kernel = mk.FusedMoEKernel( prepare_finalize, experts, shared_experts=( @@ -482,4 +482,10 @@ def make_fp8_moe_kernel( else None ), moe_parallel_config=moe_config.moe_parallel_config, + inplace=( + not moe_config.disable_inplace + and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS + ), ) + + return kernel diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 97232e9bdba7..3db4f3cf7a95 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -405,7 +405,7 @@ def make_nvfp4_moe_kernel( # NOTE(rob): we only want the mk to control the shared_expert # if using all2all (for SBO). bnell is making this explict in # the new MoE runner class. - return mk.FusedMoEKernel( + kernel = mk.FusedMoEKernel( prepare_finalize, experts, shared_experts=( @@ -414,4 +414,8 @@ def make_nvfp4_moe_kernel( else None ), moe_parallel_config=moe_config.moe_parallel_config, + inplace=False, ) + + # TODO(rob): update inplace logic to be part of the kernel. + return kernel From 1f2bb4d2e54d09bf8fde6c97f774d0d9ef41d1ec Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 16:43:24 -0500 Subject: [PATCH 119/207] remove forward Signed-off-by: Robert Shaw --- .../fused_moe/fused_moe_modular_method.py | 10 +-- .../layers/fused_moe/modular_kernel.py | 84 ++++++++++--------- .../compressed_tensors_moe.py | 17 ++-- .../model_executor/layers/quantization/fp8.py | 13 ++- .../layers/quantization/modelopt.py | 22 ++--- 5 files changed, 73 insertions(+), 73 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 3600873a2b8d..63a720123997 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -13,7 +13,7 @@ FusedMoEMethodBase, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEModularKernel, + FusedMoEKernel, FusedMoEPrepareAndFinalizeModular, ) @@ -25,9 +25,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): # --8<-- [end:modular_fused_moe] - def __init__( - self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEModularKernel - ): + def __init__(self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEKernel): super().__init__(old_quant_method.moe) self.moe_quant_config = old_quant_method.moe_quant_config self.moe_kernel = experts @@ -49,7 +47,7 @@ def make( ) -> "FusedMoEModularMethod": return FusedMoEModularMethod( old_quant_method, - FusedMoEModularKernel( + FusedMoEKernel( prepare_finalize, old_quant_method.select_gemm_impl(prepare_finalize, moe_layer), shared_experts, @@ -91,7 +89,7 @@ def apply( shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert self.moe_kernel is not None - return self.moe_kernel( + return self.moe_kernel.apply( hidden_states=x, w1=layer.w13_weight, w2=layer.w2_weight, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 12cb8137fedf..56958176ec8c 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1593,7 +1593,7 @@ def apply( @final -class FusedMoEKernel(torch.nn.Module): +class FusedMoEKernel: def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalize, @@ -1661,12 +1661,12 @@ def output_is_reduced(self) -> bool: """ return self.impl.prepare_finalize.output_is_reduced() - def forward( + def apply_monolithic( self, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - routing_input: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + router_logits: torch.Tensor | tuple[torch.Tensor, torch.Tensor], activation: str, global_num_experts: int, expert_map: torch.Tensor | None, @@ -1677,41 +1677,45 @@ def forward( routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: - if isinstance(self.impl, FusedMoEKernelModularImpl): - assert isinstance(routing_input, tuple) - topk_ids, topk_weights = routing_input - assert num_expert_group is None - assert e_score_correction_bias is None - assert routed_scaling_factor is None - assert topk_group is None - - return self.impl.apply( - hidden_states=hidden_states, - w1=w1, - w2=w2, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=activation, - global_num_experts=global_num_experts, - expert_map=expert_map, - apply_router_weight_on_input=apply_router_weight_on_input, - ) - - elif isinstance(self.impl, FusedMoEKernelMonolithicImpl): - assert isinstance(routing_input, torch.Tensor) - router_logits = routing_input + assert isinstance(self.impl, FusedMoEKernelMonolithicImpl) + return self.impl.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + router_logits=router_logits, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + topk_group=topk_group, + ) - return self.impl.apply( - hidden_states=hidden_states, - w1=w1, - w2=w2, - router_logits=router_logits, - activation=activation, - global_num_experts=global_num_experts, - expert_map=expert_map, - apply_router_weight_on_input=apply_router_weight_on_input, - num_expert_group=num_expert_group, - e_score_correction_bias=e_score_correction_bias, - routed_scaling_factor=routed_scaling_factor, - topk_group=topk_group, - ) + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + assert isinstance(self.impl, FusedMoEKernelModularImpl) + return self.impl.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + shared_experts_input=shared_experts_input, + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index f08be2313daa..95bf8bd6540a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -346,7 +346,7 @@ def apply( shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert self.moe_kernel is not None - return self.moe_kernel( + return self.moe_kernel.apply( x, layer.w13_weight, layer.w2_weight, @@ -609,8 +609,9 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) - return self.moe_kernel( + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -634,17 +635,17 @@ def apply( shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert self.moe_kernel is not None - return self.moe_kernel( + return self.moe_kernel.apply( x, layer.w13_weight, layer.w2_weight, topk_weights, topk_ids, - inplace=False, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, ) @@ -984,8 +985,8 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) - return self.moe_kernel( + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -1010,7 +1011,7 @@ def apply( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert not self.is_monolithic assert self.moe_kernel is not None - return self.moe_kernel( + return self.moe_kernel.apply( x, layer.w13_weight, layer.w2_weight, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 28539504f551..a131f7b1d882 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -946,12 +946,9 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM - if layer.enable_eplb: - raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.") - - assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) - return self.moe_kernel( + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -970,9 +967,9 @@ def apply( topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.moe_kernel is not None assert not self.is_monolithic - return self.moe_kernel( + assert self.moe_kernel is not None + return self.moe_kernel.apply( x, layer.w13_weight, layer.w2_weight, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 6a9c42150026..bb19c6cee053 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -918,12 +918,9 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if layer.enable_eplb: - raise NotImplementedError( - "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend." - ) - assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) - return self.moe_kernel( + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -946,8 +943,9 @@ def apply( topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert not self.is_monolithic assert self.moe_kernel is not None - return self.moe_kernel( + return self.moe_kernel.apply( x, layer.w13_weight, layer.w2_weight, @@ -1447,8 +1445,9 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert isinstance(self.moe_kernel, mk.FusedMoEKMonolithicKernel) - return self.moe_kernel( + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( x, layer.w13_weight, layer.w2_weight, @@ -1471,18 +1470,19 @@ def apply( topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert not self.is_monolithic assert self.moe_kernel is not None - return self.moe_kernel( + return self.moe_kernel.apply( x, layer.w13_weight, layer.w2_weight, topk_weights, topk_ids, - inplace=False, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, ) From d8db16f54175eb4d828f50bec9bba9a0d6787240 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 16:44:18 -0500 Subject: [PATCH 120/207] remove forward Signed-off-by: Robert Shaw --- .../layers/fused_moe/fused_moe_modular_method.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 63a720123997..0065c11f3163 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -25,10 +25,12 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): # --8<-- [end:modular_fused_moe] - def __init__(self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEKernel): + def __init__( + self, old_quant_method: FusedMoEMethodBase, moe_kernel: FusedMoEKernel + ): super().__init__(old_quant_method.moe) self.moe_quant_config = old_quant_method.moe_quant_config - self.moe_kernel = experts + self.moe_kernel = moe_kernel self.disable_expert_map = getattr( old_quant_method, "disable_expert_map", From 658285e880c015f8f55de69840175b77d925d58c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 16:47:35 -0500 Subject: [PATCH 121/207] updated Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/modular_kernel.py | 4 ++++ .../compressed_tensors/compressed_tensors_moe.py | 10 ++++------ vllm/model_executor/layers/quantization/fp8.py | 4 ++-- vllm/model_executor/layers/quantization/modelopt.py | 10 ++++------ 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 56958176ec8c..a153f6e898de 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1637,6 +1637,10 @@ def __init__( self._post_init_setup() + @property + def is_monolithic(self) -> bool: + return isinstance(self.impl, FusedMoEKernelMonolithicImpl) + def _post_init_setup(self): """ Resolve any leftover setup dependencies between self.prepare_finalize diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 95bf8bd6540a..d607e5aa182a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -40,7 +40,6 @@ fused_marlin_moe, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( - Fp8MoeBackend, convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, @@ -598,10 +597,8 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon @property def is_monolithic(self) -> bool: - return ( - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - and not self.moe.moe_parallel_config.use_all2all_kernels - ) + assert self.moe_kernel is not None + return self.moe_kernel.is_monolithic def apply_monolithic( self, @@ -977,7 +974,8 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon @property def is_monolithic(self) -> bool: - return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM + assert self.moe_kernel is not None + return self.moe_kernel.is_monolithic def apply_monolithic( self, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index a131f7b1d882..d3874982b909 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -29,7 +29,6 @@ ) from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( - Fp8MoeBackend, convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, @@ -938,7 +937,8 @@ def supports_eplb(self) -> bool: @property def is_monolithic(self) -> bool: - return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM + assert self.moe_kernel is not None + return self.moe_kernel.is_monolithic def apply_monolithic( self, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index bb19c6cee053..e77482a22a2c 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -20,7 +20,6 @@ FusedMoeWeightScaleSupported, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( - Fp8MoeBackend, convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, @@ -910,7 +909,8 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon @property def is_monolithic(self) -> bool: - return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM + assert self.moe_kernel is not None + return self.moe_kernel.is_monolithic def apply_monolithic( self, @@ -1434,10 +1434,8 @@ def supports_eplb(self) -> bool: @property def is_monolithic(self) -> bool: - return ( - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - and not self.moe.moe_parallel_config.use_all2all_kernels - ) + assert self.moe_kernel is not None + return self.moe_kernel.is_monolithic def apply_monolithic( self, From 8ab8d93960cddfe086bd9753408ca20ef1da6f92 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 16:48:46 -0500 Subject: [PATCH 122/207] revert comment Signed-off-by: Robert Shaw --- .../quantization/compressed_tensors/compressed_tensors_moe.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index d607e5aa182a..308e5c97f559 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -1017,6 +1017,8 @@ def apply( topk_ids, activation=layer.activation, global_num_experts=layer.global_num_experts, + # TODO(rob): investigate the disable_expert_map introduced by: + # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501 expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, shared_experts_input=shared_experts_input, From 6e93977fdfc889f224658ed45cf7fd8cd25ed887 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 17:02:03 -0500 Subject: [PATCH 123/207] remove the flashinfer prepare_dp_all_gather Signed-off-by: Robert Shaw --- .../fused_moe/runner/default_moe_runner.py | 65 ++++--------------- .../layers/quantization/modelopt.py | 5 -- 2 files changed, 14 insertions(+), 56 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 7da36592e3db..ecdccbe988d7 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -623,45 +623,6 @@ def forward_impl( ) with sp_ctx: - extra_tensors = None - if do_naive_dispatch_combine: - post_quant_allgather = ( - self.quant_method is not None - and self.moe_config.dp_size > 1 - and self.moe_config.use_ep - and getattr(self.quant_method, "do_post_quant_allgather", False) - ) - if post_quant_allgather: - hidden_states_to_dispatch, extra_tensors = ( - self.quant_method.prepare_dp_allgather_tensor( - layer, hidden_states, router_logits - ) - ) - else: - hidden_states_to_dispatch = hidden_states - - dispatch_res = get_ep_group().dispatch_router_logits( - hidden_states_to_dispatch, - router_logits, - self.moe_config.is_sequence_parallel, - extra_tensors=extra_tensors, - ) - if extra_tensors is not None: - ( - orig_hidden_states, - router_logits, - extra_tensors_combined, - ) = dispatch_res - hidden_states_combined = ( - orig_hidden_states, - extra_tensors_combined[0], - ) - else: - hidden_states_combined, router_logits = dispatch_res - orig_hidden_states = hidden_states_combined - else: - orig_hidden_states = hidden_states - # Run shared experts before matrix multiply. # because matrix multiply maybe modify the hidden_states. if has_separate_shared_experts and not use_shared_experts_stream: @@ -671,6 +632,17 @@ def forward_impl( ) shared_output = self.shared_experts(shared_input) + # For naive dispatch/combine Dp/Ep, dispatch the hidden states and + # router logits to all experts. + # NOTE: this will be removed once all kernels are migrated into the + # MoEKernel framework. + if do_naive_dispatch_combine: + hidden_states, router_logits = get_ep_group().dispatch_router_logits( + hidden_states, + router_logits, + self.moe_config.is_sequence_parallel, + ) + # NOTE: Similar with DP, PCP also needs dispatch and combine. For # simplicity, AgRsAll2All was added separately for PCP here. Maybe # we should modify All2AllManager abstract to better support PCP. @@ -684,31 +656,22 @@ def forward_impl( dim=0, ) - # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014). - # Figure out nicer way to do this. - if do_naive_dispatch_combine: - x = hidden_states_combined - x_orig = orig_hidden_states - else: - x = hidden_states - x_orig = hidden_states - # Matrix multiply. if self.quant_method.is_monolithic: final_hidden_states = self.quant_method.apply_monolithic( layer=layer, - x=x, + x=hidden_states, router_logits=router_logits, ) else: topk_weights, topk_ids = self.router.select_experts( - hidden_states=x_orig, + hidden_states=hidden_states, router_logits=router_logits, ) final_hidden_states = self.quant_method.apply( layer=layer, - x=x, # The type signture of this is wrong due to the hack. + x=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, shared_experts_input=shared_input, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e77482a22a2c..21093c7ff42d 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -26,7 +26,6 @@ select_fp8_moe_backend, ) from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( - NvFp4MoeBackend, convert_to_nvfp4_moe_kernel_format, is_global_sf_supported_for_nvfp4_backend, make_nvfp4_moe_kernel, @@ -1413,10 +1412,6 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: routing_tables=layer._maybe_init_expert_routing_tables(), ) - @property - def do_post_quant_allgather(self): - return self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: return make_nvfp4_moe_quant_config( backend=self.nvfp4_backend, From 36f354509dffadbb93c14a34539417c4717275a8 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 17:10:04 -0500 Subject: [PATCH 124/207] update to FusedMoEKernel Signed-off-by: Robert Shaw --- tests/kernels/moe/test_batched_deepgemm.py | 6 +++--- tests/kernels/moe/test_deepep_deepgemm_moe.py | 16 ++++++++-------- tests/kernels/moe/test_deepep_moe.py | 8 ++++---- tests/kernels/moe/test_flashinfer_moe.py | 4 ++-- tests/kernels/moe/test_modular_oai_triton_moe.py | 4 ++-- tests/kernels/moe/test_pplx_cutlass_moe.py | 4 ++-- tests/kernels/moe/test_pplx_moe.py | 4 ++-- tests/kernels/moe/utils.py | 10 +++++----- vllm/lora/layers/fused_moe.py | 4 ++-- .../layers/fused_moe/xpu_fused_moe.py | 2 +- 10 files changed, 31 insertions(+), 31 deletions(-) diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py index 2c6c45a5f234..8aaebcee1e6e 100644 --- a/tests/kernels/moe/test_batched_deepgemm.py +++ b/tests/kernels/moe/test_batched_deepgemm.py @@ -12,7 +12,7 @@ BatchedPrepareAndFinalize, BatchedTritonExperts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported from .test_deepgemm import make_block_quant_fp8_weights @@ -74,7 +74,7 @@ def test_batched_deepgemm_vs_triton( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk_triton = FusedMoEModularKernel( + mk_triton = FusedMoEKernel( prep_finalize, triton_experts, inplace=False, @@ -96,7 +96,7 @@ def test_batched_deepgemm_vs_triton( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - mk_deepgemm = FusedMoEModularKernel( + mk_deepgemm = FusedMoEKernel( prep_finalize, deepgemm_experts, inplace=False, diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index 2b8240482829..e208fd0eed7f 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -22,7 +22,7 @@ fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.utils.deep_gemm import ( get_mk_alignment_for_contiguous_layout, is_deep_gemm_e8m0_used, @@ -170,7 +170,7 @@ def make_ll_modular_kernel( q_dtype: torch.dtype | None, test_config: TestConfig, quant_config: FusedMoEQuantConfig, -) -> FusedMoEModularKernel: +) -> FusedMoEKernel: assert test_config.low_latency assert test_config.use_fp8_dispatch is not None @@ -195,7 +195,7 @@ def make_ll_modular_kernel( quant_config=quant_config, moe_config=make_dummy_moe_config(), ) - return FusedMoEModularKernel( + return FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, inplace=False, @@ -210,7 +210,7 @@ def make_ht_modular_kernel( q_dtype: torch.dtype | None, test_config: TestConfig, quant_config: FusedMoEQuantConfig, -) -> FusedMoEModularKernel: +) -> FusedMoEKernel: assert not test_config.low_latency assert test_config.use_fp8_dispatch is None @@ -228,7 +228,7 @@ def make_ht_modular_kernel( moe_config=make_dummy_moe_config(), quant_config=quant_config, ) - return FusedMoEModularKernel( + return FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, inplace=False, @@ -242,11 +242,11 @@ def make_modular_kernel( num_local_experts: int, test_tensors: TestTensors, quant_config: FusedMoEQuantConfig, -) -> FusedMoEModularKernel: +) -> FusedMoEKernel: q_dtype = torch.float8_e4m3fn test_config = test_tensors.config - mk: FusedMoEModularKernel + mk: FusedMoEKernel # Make modular kernel if test_config.low_latency: max_tokens_per_rank = max(64, next_power_of_2(test_tensors.rank_tokens.size(0))) @@ -307,7 +307,7 @@ def build_expert_map(): ) # Make modular kernel - mk: FusedMoEModularKernel = make_modular_kernel( + mk: FusedMoEKernel = make_modular_kernel( pg=pg, pgi=pgi, dp_size=dp_size, diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index 01f340730af3..e648291e351d 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -20,7 +20,7 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -135,7 +135,7 @@ def make_modular_kernel( q_dtype: torch.dtype | None, use_fp8_dispatch: bool, quant_config: FusedMoEQuantConfig, -) -> FusedMoEModularKernel: +) -> FusedMoEKernel: ht_args: DeepEPHTArgs | None = None ll_args: DeepEPLLArgs | None = None @@ -180,7 +180,7 @@ def make_modular_kernel( quant_config=quant_config, ) - mk = FusedMoEModularKernel( + mk = FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, inplace=False, @@ -242,7 +242,7 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False): ) # Make modular kernel - mk: FusedMoEModularKernel = make_modular_kernel( + mk: FusedMoEKernel = make_modular_kernel( pg, pgi, low_latency_mode, diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index 1f1349cff841..4bd4a1bb5b7a 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -23,7 +23,7 @@ FlashInferExperts, is_valid_flashinfer_cutlass_fused_moe, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -107,7 +107,7 @@ def test_flashinfer_fp4_moe_no_graph( routing_method=RoutingMethodType.TopK, ) - flashinfer_experts = FusedMoEModularKernel( + flashinfer_experts = FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), FlashInferExperts(moe_config=moe_config, quant_config=quant_config), inplace=False, diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index cf9ff18634d0..0d233b9f5794 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -29,7 +29,7 @@ OAITritonExperts, UnfusedOAITritonExperts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -181,7 +181,7 @@ def oai_triton_moe_impl( else: fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config) - mk = FusedMoEModularKernel( + mk = FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), fused_experts, inplace=False, diff --git a/tests/kernels/moe/test_pplx_cutlass_moe.py b/tests/kernels/moe/test_pplx_cutlass_moe.py index d8a6600743e2..bac1be6db3b5 100644 --- a/tests/kernels/moe/test_pplx_cutlass_moe.py +++ b/tests/kernels/moe/test_pplx_cutlass_moe.py @@ -17,7 +17,7 @@ fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassBatchedExpertsFp8 -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import set_random_seed @@ -171,7 +171,7 @@ def make_moe_config() -> FusedMoEConfig: num_dispatchers=num_dispatchers, ) - fused_cutlass_experts = FusedMoEModularKernel( + fused_cutlass_experts = FusedMoEKernel( prepare_finalize, experts, inplace=False, diff --git a/tests/kernels/moe/test_pplx_moe.py b/tests/kernels/moe/test_pplx_moe.py index deb3b9eb4d76..b473e08c8737 100644 --- a/tests/kernels/moe/test_pplx_moe.py +++ b/tests/kernels/moe/test_pplx_moe.py @@ -41,7 +41,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceDelegate, ) @@ -588,7 +588,7 @@ def pplx_moe( moe_config=make_dummy_moe_config(), ) - fused_experts = FusedMoEModularKernel( + fused_experts = FusedMoEKernel( prepare_finalize, experts, shared_experts, diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index ef72b96bead3..2ef74d4bd9d9 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -23,7 +23,7 @@ TritonExperts, fused_experts, ) -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) @@ -115,7 +115,7 @@ def batched_moe( a2_scale=a2_scale, ) - fused_experts = FusedMoEModularKernel( + fused_experts = FusedMoEKernel( BatchedPrepareAndFinalize( max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 ), @@ -157,7 +157,7 @@ def naive_batched_moe( a2_scale=a2_scale, ) - fused_experts = FusedMoEModularKernel( + fused_experts = FusedMoEKernel( BatchedPrepareAndFinalize( max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 ), @@ -571,8 +571,8 @@ def modular_triton_fused_moe( moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None, -) -> FusedMoEModularKernel: - return FusedMoEModularKernel( +) -> FusedMoEKernel: + return FusedMoEKernel( MoEPrepareAndFinalizeNoEP(), TritonExperts(moe_config, quant_config), shared_experts, diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 371cacfa2b12..b89cfd469d4f 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -32,7 +32,7 @@ UnfusedOAITritonExperts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( - FusedMoEModularKernel, + FusedMoEKernel, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, @@ -136,7 +136,7 @@ def _inject_lora_into_fused_moe(self): else: # Create a new modular kernel via select_gemm_impl prepare_finalize = MoEPrepareAndFinalizeNoEP() - m_fused_moe_fn = FusedMoEModularKernel( + m_fused_moe_fn = FusedMoEKernel( prepare_finalize, self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer diff --git a/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py index e6f8b8efa804..95a69662e38a 100644 --- a/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py @@ -23,7 +23,7 @@ from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe -class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute): +class XPUExperts(mk.FusedMoEKernel): def __init__( self, moe_config: FusedMoEConfig, From 2e15cf26d35e04ca41d945daf9555ca24598ef41 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 17:15:26 -0500 Subject: [PATCH 125/207] remove select_gemm_impl Signed-off-by: Robert Shaw --- .../layers/fused_moe/fused_moe_method_base.py | 8 ++++---- .../compressed_tensors_moe.py | 20 ------------------- .../model_executor/layers/quantization/fp8.py | 12 ----------- .../layers/quantization/modelopt.py | 10 ---------- 4 files changed, 4 insertions(+), 46 deletions(-) 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 f45a8cf6e48e..1b16f2184971 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 @@ -84,9 +84,9 @@ def select_gemm_impl( ) -> FusedMoEExpertsModular: # based on the all2all implementation, select the appropriate # gemm implementation - raise NotImplementedError( - f"{self.__class__.__name__} must select appropriate gemm " - "implementation based on the prepare_finalize" + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel initialization " + "logic. This function should not be called." ) @abstractmethod @@ -98,7 +98,7 @@ def get_fused_moe_quant_config( @property def topk_indices_dtype(self) -> torch.dtype | None: if self.moe_kernel is not None: - return self.moe_kernel.prepare_finalize.topk_indices_dtype() + return self.moe_kernel.impl.prepare_finalize.topk_indices_dtype() return None @property diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 308e5c97f559..1344c10fd94e 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -574,16 +574,6 @@ def maybe_make_prepare_finalize( "logic. This function should not be called." ) - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic. This function should not be called." - ) - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: return make_nvfp4_moe_quant_config( backend=self.nvfp4_backend, @@ -949,16 +939,6 @@ def maybe_make_prepare_finalize( "logic. This function should not be called." ) - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic. This function should not be called." - ) - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN return make_fp8_moe_quant_config( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index d3874982b909..18dc061a0457 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -19,9 +19,7 @@ ) from vllm.model_executor.layers.fused_moe import ( FusedMoE, - FusedMoEExpertsModular, FusedMoEMethodBase, - FusedMoEPrepareAndFinalizeModular, FusedMoeWeightScaleSupported, ) from vllm.model_executor.layers.fused_moe.config import ( @@ -906,16 +904,6 @@ def maybe_make_prepare_finalize( "logic. This function should not be called." ) - def select_gemm_impl( - self, - prepare_finalize: FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> FusedMoEExpertsModular: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic. This function should not be called." - ) - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 21093c7ff42d..999b4839f690 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1222,16 +1222,6 @@ def maybe_make_prepare_finalize( "logic. This function should not be called." ) - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic. This function should not be called." - ) - def uses_weight_scale_2_pattern(self) -> bool: """ FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales. From d897152eb126837b9b52449d39ddfdd8d17b9840 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 17:20:50 -0500 Subject: [PATCH 126/207] stash Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/modular_kernel.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index a153f6e898de..4c0745df5fcb 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1502,6 +1502,8 @@ def apply( expert_map, apply_router_weight_on_input, ) + # print(f"{a1q.dtype=}") + # print(f"{a1q_scale.dtype=}") fused_out = self._fused_experts( in_dtype=hidden_states.dtype, From 0c29d090e6f5cfb48a684aaf99ffc142fb4936cb Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 17:31:56 -0500 Subject: [PATCH 127/207] fix bad tping Signed-off-by: Robert Shaw --- .../layers/fused_moe/flashinfer_trtllm_fp8_moe.py | 4 ++-- .../layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py | 8 ++++---- vllm/model_executor/layers/fused_moe/modular_kernel.py | 6 ++++-- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py index fcd5120f78db..5181987d1135 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py @@ -159,7 +159,7 @@ def _apply_per_block( w1: torch.Tensor, w2: torch.Tensor, router_logits: torch.Tensor, - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, @@ -171,7 +171,7 @@ def _apply_per_block( topk_group: int | None = None, ) -> torch.Tensor: assert not apply_router_weight_on_input - assert activation == "silu" + assert activation == MoEActivation.SILU assert ( e_score_correction_bias is None or e_score_correction_bias.dtype == hidden_states.dtype diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py index ccd52b094777..8427b59da4bd 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py @@ -130,7 +130,7 @@ def apply( w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, @@ -140,7 +140,7 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - assert activation == "silu" + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] assert a1q_scale is not None assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None @@ -234,7 +234,7 @@ def apply( w1: torch.Tensor, w2: torch.Tensor, router_logits: torch.Tensor, - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, @@ -245,7 +245,7 @@ def apply( routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: - assert activation == "silu" + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] assert a1q_scale is not None assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 4c0745df5fcb..a063565ee676 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -601,6 +601,7 @@ def _supports_routing_method( """ return True + @staticmethod def _supports_router_logits_dtype( router_logits_dtype: torch.dtype | None, routing_method: RoutingMethodType, @@ -917,6 +918,7 @@ def _supports_routing_method( """ raise NotImplementedError + @staticmethod def _supports_router_logits_dtype( router_logits_dtype: torch.dtype | None, routing_method: RoutingMethodType, @@ -938,7 +940,7 @@ def apply( w1: torch.Tensor, w2: torch.Tensor, router_logits: torch.Tensor, - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, @@ -1704,8 +1706,8 @@ def apply( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - topk_ids: torch.Tensor, topk_weights: torch.Tensor, + topk_ids: torch.Tensor, activation: str, global_num_experts: int, expert_map: torch.Tensor | None, From 1deda4ae00035b8294187b586c846c523922edc0 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 17:59:17 -0500 Subject: [PATCH 128/207] update import pathways Signed-off-by: Robert Shaw --- .../kernels/benchmark_cutlass_moe_fp8.py | 6 +- .../kernels/benchmark_cutlass_moe_nvfp4.py | 6 +- .../kernels/benchmark_grouped_gemm_cutlass.py | 6 +- benchmarks/kernels/benchmark_moe.py | 2 +- docs/design/moe_kernel_features.md | 2 +- .../moe/modular_kernel_tools/mk_objects.py | 4 +- tests/kernels/moe/test_block_fp8.py | 4 +- tests/kernels/moe/test_cutlass_moe.py | 6 +- tests/kernels/moe/test_deepgemm.py | 4 +- tests/kernels/moe/test_flashinfer.py | 8 +- tests/kernels/moe/test_flashinfer_moe.py | 4 +- .../moe/test_modular_oai_triton_moe.py | 4 +- tests/kernels/moe/test_nvfp4_moe.py | 4 +- tests/kernels/moe/utils.py | 4 +- vllm/lora/layers/fused_moe.py | 4 +- .../layers/fused_moe/all2all_utils.py | 10 +- .../layers/fused_moe/cutlass_moe.py | 6 +- .../layers/fused_moe/modular_kernel.py | 12 + .../layers/fused_moe/oracle/unquantized.py | 10 +- .../layers/fused_moe/prepare_finalize.py | 378 ------------------ .../fused_moe/prepare_finalize/__init__.py | 22 + .../fused_moe/prepare_finalize/naive_dp_ep.py | 253 ++++++++++++ .../fused_moe/prepare_finalize/no_dp_ep.py | 141 +++++++ .../fused_moe/topk_weight_and_reduce.py | 2 +- 24 files changed, 476 insertions(+), 426 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/prepare_finalize.py create mode 100644 vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py create mode 100644 vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py create mode 100644 vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index 92032efb94a9..ad5ed13a38dd 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -15,8 +15,8 @@ from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.platforms import current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -138,7 +138,7 @@ def bench_run( ) fn = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( num_experts=num_experts, diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index e0b6cb6d8210..beba7a08a247 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -24,7 +24,7 @@ ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.scalar_type import scalar_types from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -197,7 +197,7 @@ def run_cutlass_moe_fp4( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsFp4( make_dummy_moe_config(), quant_config=quant_config, @@ -242,7 +242,7 @@ def run_cutlass_from_graph( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsFp4( make_dummy_moe_config(), quant_config=quant_config, diff --git a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py index fd6c5216420e..113babac1427 100644 --- a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py +++ b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py @@ -16,7 +16,7 @@ fused_topk, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.worker.workspace import init_workspace_manager @@ -133,7 +133,7 @@ def run_cutlass_moe( ) fn = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( num_experts=w2.shape[0], @@ -165,7 +165,7 @@ def run_cutlass_from_graph( ) fn = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( num_experts=w2.shape[0], diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 064d72f2884c..729cfd4459d9 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -243,7 +243,7 @@ def run(): deep_gemm_experts = None if use_deep_gemm: deep_gemm_experts = mk.FusedMoEKernel( - prepare_finalize=MoEPrepareAndFinalizeNoEP(), + prepare_finalize=MoEPrepareAndFinalizeNoDPEPModular(), fused_experts=TritonOrDeepGemmExperts( moe_config=FusedMoEConfig( num_experts=num_experts, diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index ce7c34743018..7998273ae0c6 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -37,7 +37,7 @@ th { | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] | | flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] | -| MoEPrepareAndFinalizeNoEP5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoEP`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoEP] | +| MoEPrepareAndFinalizeNoDPEPModular5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoDPEPModular`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoDPEPModular] | | BatchedPrepareAndFinalize5 | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] | !!! info "Table key" diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index d232d10bf804..b403eade02c1 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -20,7 +20,7 @@ NaiveBatchedExperts, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -165,7 +165,7 @@ def expert_info(kind) -> ExpertInfo: register_prepare_and_finalize( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, standard_format, common_float_types, blocked_quantization_support=True, diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 3b00cdd7326f..cda032a5ddb7 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -28,7 +28,7 @@ _valid_deep_gemm_shape, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -254,7 +254,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) ) deep_gemm_experts = mk.FusedMoEKernel( - prepare_finalize=MoEPrepareAndFinalizeNoEP(), + prepare_finalize=MoEPrepareAndFinalizeNoDPEPModular(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), quant_config=quant_config, diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 65887313d6d3..eaed0782725a 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -23,7 +23,7 @@ run_cutlass_moe_fp8, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.platforms import current_platform @@ -198,7 +198,7 @@ def slice_experts(): w2 = kwargs["w2"] a = kwargs["hidden_states"] kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( num_experts=w2.shape[0], @@ -258,7 +258,7 @@ def run_8_bit( with_ep = num_local_experts is not None or num_local_experts == num_experts if not with_ep: kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsFp8( moe_config=make_dummy_moe_config( num_experts=moe_tensors.w2_q.shape[0], # type: ignore[union-attr] diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index eedffae7ead5..65a52a055c10 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -19,7 +19,7 @@ ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -110,7 +110,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size): ) deep_gemm_experts = mk.FusedMoEKernel( - prepare_finalize=MoEPrepareAndFinalizeNoEP(), + prepare_finalize=MoEPrepareAndFinalizeNoDPEPModular(), fused_experts=TritonOrDeepGemmExperts( moe_config=make_dummy_moe_config(), quant_config=quant_config, diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index a486e0c884b4..5fe741b63616 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -23,8 +23,8 @@ ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, - MoEPrepareAndFinalizeNoEPMonolithic, + MoEPrepareAndFinalizeNoDPEPModular, + MoEPrepareAndFinalizeNoDPEPModularMonolithic, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, @@ -241,7 +241,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEPMonolithic(), + MoEPrepareAndFinalizeNoDPEPModularMonolithic(), FlashInferTrtLlmFp8Experts( moe_config=td.layer.moe, quant_config=quant_config, @@ -348,7 +348,7 @@ def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig: ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), FlashInferExperts( moe_config=moe_config, quant_config=quant_config, diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index 4bd4a1bb5b7a..df1273047563 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -25,7 +25,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe @@ -108,7 +108,7 @@ def test_flashinfer_fp4_moe_no_graph( ) flashinfer_experts = FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), FlashInferExperts(moe_config=moe_config, quant_config=quant_config), inplace=False, ) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 0d233b9f5794..a9513fc4ee14 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -31,7 +31,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.utils import shuffle_weight from vllm.platforms import current_platform @@ -182,7 +182,7 @@ def oai_triton_moe_impl( fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config) mk = FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), fused_experts, inplace=False, ) diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index abf0a737480a..6804827bfc32 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -19,7 +19,7 @@ CutlassExpertsFp4, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -90,7 +90,7 @@ def test_cutlass_fp4_moe_no_graph( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsFp4( moe_config=make_dummy_moe_config(), quant_config=quant_config, diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 2ef74d4bd9d9..a74575d16a1c 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -25,7 +25,7 @@ ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input @@ -573,7 +573,7 @@ def modular_triton_fused_moe( shared_experts: torch.nn.Module | None = None, ) -> FusedMoEKernel: return FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), TritonExperts(moe_config, quant_config), shared_experts, inplace=False, diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index b89cfd469d4f..0916988c0b74 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -35,7 +35,7 @@ FusedMoEKernel, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from .utils import _get_lora_device, try_get_optimal_moe_lora_config @@ -135,7 +135,7 @@ def _inject_lora_into_fused_moe(self): m_fused_moe_fn = self.base_layer.quant_method.moe_kernel else: # Create a new modular kernel via select_gemm_impl - prepare_finalize = MoEPrepareAndFinalizeNoEP() + prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular() m_fused_moe_fn = FusedMoEKernel( prepare_finalize, self.base_layer.quant_method.select_gemm_impl( diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index cb1cac390594..7d318d1a5fc9 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -20,8 +20,8 @@ FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNaiveEPBase, - MoEPrepareAndFinalizeNoEPBase, + make_moe_prepare_and_finalize_naive_dp_ep, + make_moe_prepare_and_finalize_no_dp_ep, ) from vllm.platforms import current_platform from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx @@ -107,7 +107,7 @@ def maybe_make_prepare_finalize( "Detected DP deployment with no --enable-expert-parallel. " "Falling back to AllGather+ReduceScatter dispatch/combine." ) - return MoEPrepareAndFinalizeNaiveEPBase.make( + return make_moe_prepare_and_finalize_naive_dp_ep( is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, num_dispatchers=( get_ep_group().device_communicator.all2all_manager.world_size @@ -115,7 +115,7 @@ def maybe_make_prepare_finalize( use_monolithic=use_monolithic, ) else: - return MoEPrepareAndFinalizeNoEPBase.make(use_monolithic) + return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic) all2all_manager = get_ep_group().device_communicator.all2all_manager assert all2all_manager is not None @@ -248,7 +248,7 @@ def maybe_make_prepare_finalize( ) elif moe.use_naive_all2all_kernels and allow_new_interface: - prepare_finalize = MoEPrepareAndFinalizeNaiveEPBase.make( + prepare_finalize = make_moe_prepare_and_finalize_naive_dp_ep( use_monolithic=use_monolithic, is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, num_dispatchers=all2all_manager.world_size, diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index ef9937befba7..54f686939a94 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -21,7 +21,7 @@ moe_unpermute, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceDelegate, @@ -1167,7 +1167,7 @@ def cutlass_moe_w4a8_fp8( num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0) fn = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), CutlassExpertsW4A8Fp8( out_dtype=a.dtype, a_strides1=a_strides1, @@ -1184,7 +1184,7 @@ def cutlass_moe_w4a8_fp8( ), ) - return fn( + return fn.apply( a, w1_q, w2_q, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index a063565ee676..4309377a05f9 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -171,6 +171,11 @@ def apply( ReceiverType = Callable[[], PrepareResultType] +################################################################################ +# Prepare/Finalize +# TODO: make the below a separate file. +################################################################################ + class FusedMoEPrepareAndFinalize(ABC): """ @@ -450,6 +455,12 @@ def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor: raise NotImplementedError +################################################################################ +# Experts +# TODO: make the below a separate file. +################################################################################ + + # TODO: add supported activations method (return string) class FusedMoEExperts(ABC): def __init__( @@ -971,6 +982,7 @@ def _slice_scales( ################################################################################ +# MoEKernel # TODO: make the below a separate file. ################################################################################ diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 1ae3d1a01d94..015866bf44b9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -18,7 +18,7 @@ is_supported_config_trtllm_bf16, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoEP, + MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( swap_w13_to_w31, @@ -175,7 +175,7 @@ def make_unquantized_moe_kernel( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), FlashInferExperts( moe_config=moe_config, quant_config=quant_config, @@ -189,7 +189,7 @@ def make_unquantized_moe_kernel( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), AiterExperts( moe_config=moe_config, quant_config=quant_config, @@ -200,7 +200,7 @@ def make_unquantized_moe_kernel( from vllm.model_executor.layers.fused_moe import TritonExperts kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), TritonExperts( moe_config=moe_config, quant_config=quant_config, @@ -211,7 +211,7 @@ def make_unquantized_moe_kernel( from vllm.model_executor.layers.fused_moe import XPUExperts kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoEP(), + MoEPrepareAndFinalizeNoDPEPModular(), XPUExperts( moe_config=moe_config, quant_config=quant_config, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py deleted file mode 100644 index 8d263aa85e21..000000000000 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py +++ /dev/null @@ -1,378 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.distributed import get_ep_group -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceContiguous, - TopKWeightAndReduceDelegate, -) -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input -from vllm.utils.flashinfer import nvfp4_block_scale_interleave - - -class MoEPrepareAndFinalizeNaiveEPBase(mk.FusedMoEPrepareAndFinalize): - """ - Base class for Naive Prepare/Finalize for Dp/Ep with two subclasses: - * Modular Case - * Monolithic Case - - In modular case, a separate router runs *before* and we dispatch - the topk weights and ids. - - In monolithic case, the router runs *inside* the MoE kernel so we - dispatch the router logits. - - In both cases, the quantization of X happens prior to dispatching. - """ - - @staticmethod - def make( - is_sequence_parallel: bool = False, - num_dispatchers: int = 1, - use_monolithic: bool = False, - ) -> "MoEPrepareAndFinalizeNaiveEPBase": - cls = ( - MoEPrepareAndFinalizeNaiveEPMonolithic - if use_monolithic - else MoEPrepareAndFinalizeNaiveEP - ) - return cls( - is_sequence_parallel=is_sequence_parallel, num_dispatchers=num_dispatchers - ) - - def __init__( - self, - is_sequence_parallel: bool = False, - num_dispatchers: int = 1, - ) -> None: - super().__init__() - self.is_sequence_parallel = is_sequence_parallel - self._num_dispatchers = num_dispatchers - - @property - def activation_format(self) -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.Standard - - def max_num_tokens_per_rank(self) -> int | None: - return None - - def topk_indices_dtype(self) -> torch.dtype | None: - return None - - def num_dispatchers(self) -> int: - return self._num_dispatchers - - def output_is_reduced(self) -> bool: - return False - - def _quantize_and_setup_dispatch( - self, - a1: torch.Tensor, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: - # Defer input quantization to the MoE kernel. - if defer_input_quant: - a1q = a1 - a1q_scale = None - else: - input_sf = ( - quant_config.a1_gscale - if quant_config.use_nvfp4_w4a4 - else quant_config.a1_scale - ) - - # NOTE: swizzling pads the scales to multiple of 128 - # which makes the scales tensor different shape than - # the hidden states, breaking the A2A kernel. So, we - # delay the swizzling until after the A2A. - a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input( - a1, - input_sf, - quant_dtype=quant_config.quant_dtype, - per_act_token_quant=quant_config.per_act_token_quant, - block_shape=quant_config.block_shape, - is_fp4_scale_swizzled=False, - ) - - # Skip gathering scales if we have static quantization - # (the scale is a scalar, replicated on all ranks) or - # if quantization is deferred. - skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0 - scales = None if skip_gather_scales else [a1q_scale] - - return a1q, scales - - def _unwrap_scale_and_prepare_for_moe( - self, - scales: list[torch.Tensor] | None, - quant_config: FusedMoEQuantConfig, - ) -> torch.Tensor: - assert scales is not None and len(scales) == 1 - a1q_scale = scales[0] - # Apply swizzling after a2a if the MoE kernel needs it. - if quant_config.quant_dtype == "nvfp4" and quant_config.is_nvfp4_scale_swizzled: - assert a1q_scale is not None - if a1q_scale.element_size() == 1: - a1q_scale = a1q_scale.view(torch.uint8) - a1q_scale = nvfp4_block_scale_interleave(a1q_scale) - - return a1q_scale - - -class MoEPrepareAndFinalizeNaiveEP( - MoEPrepareAndFinalizeNaiveEPBase, mk.FusedMoEPrepareAndFinalizeModular -): - """ - Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. - - Uses Torch AR/RS or AR for dispatch/combine operations, applied - to the topk weights and ids. - """ - - def prepare( - self, - a1: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> mk.PrepareResultType: - """Quantize and Dispatch Topk Weights and Topk Ids.""" - - if apply_router_weight_on_input: - topk = topk_ids.size(1) - assert topk == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - # Note: do not use inplace for shared experts overlap - a1 = a1 * topk_weights.to(a1.dtype) - - a1q, scales = self._quantize_and_setup_dispatch( - a1, quant_config, defer_input_quant - ) - - res = get_ep_group().dispatch( - a1q, - topk_weights, - topk_ids, - is_sequence_parallel=self.is_sequence_parallel, - extra_tensors=scales, - ) - - if scales is None: - a1q, topk_weights, topk_ids = res - a1q_scale = None - else: - a1q, topk_weights, topk_ids, scales = res - a1q_scale = self._unwrap_scale_and_prepare_for_moe(scales, quant_config) - - return a1q, a1q_scale, None, topk_ids, topk_weights - - def finalize( - self, - output: torch.Tensor, - fused_expert_output: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - apply_router_weight_on_input: bool, - weight_and_reduce_impl: mk.TopKWeightAndReduce, - ) -> None: - if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): - weight_and_reduce_impl = TopKWeightAndReduceContiguous() - - out = weight_and_reduce_impl.apply( - output=None, - fused_expert_output=fused_expert_output, - topk_weights=topk_weights, - topk_ids=topk_ids, - apply_router_weight_on_input=apply_router_weight_on_input, - ) - - output.copy_( - get_ep_group().combine(out, is_sequence_parallel=self.is_sequence_parallel) - ) - - -class MoEPrepareAndFinalizeNaiveEPMonolithic( - MoEPrepareAndFinalizeNaiveEPBase, mk.FusedMoEPrepareAndFinalizeMonolithic -): - """ - Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. - - Uses Torch AR/RS or AR for dispatch/combine operations, applied - to the router logits (the MoE kernel runs the router internally). - """ - - def prepare( - self, - a1: torch.Tensor, - router_logits: torch.Tensor, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> mk.PrepareMonolithicResultType: - """Quantize and Dispatch Router Logits.""" - - a1q, scales = self._quantize_and_setup_dispatch( - a1, quant_config, defer_input_quant - ) - - res = get_ep_group().dispatch_router_logits( - a1q, - router_logits, - is_sequence_parallel=self.is_sequence_parallel, - extra_tensors=scales, - ) - - if scales is None: - a1q, router_logits = res - a1q_scale = None - else: - a1q, router_logits, scales = res - a1q_scale = self._unwrap_scale_and_prepare_for_moe(scales, quant_config) - - return a1q, a1q_scale, router_logits - - def finalize( - self, - fused_expert_output: torch.Tensor, - ) -> torch.Tensor: - out = get_ep_group().combine( - fused_expert_output, is_sequence_parallel=self.is_sequence_parallel - ) - return out - - -class MoEPrepareAndFinalizeNoEPBase(mk.FusedMoEPrepareAndFinalize): - """MoE prepare and finalize without expert parallelism.""" - - """ - Base class for TP case Prepare/Finalize. - * prepare: applies input quantization - * finalize: applies the reduction (if needed) - """ - - @staticmethod - def make(use_monolithic: bool) -> "MoEPrepareAndFinalizeNoEPBase": - return ( - MoEPrepareAndFinalizeNoEPMonolithic() - if use_monolithic - else MoEPrepareAndFinalizeNoEP() - ) - - @property - def activation_format(self) -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.Standard - - def max_num_tokens_per_rank(self) -> int | None: - return None - - def topk_indices_dtype(self) -> torch.dtype | None: - return None - - def num_dispatchers(self) -> int: - return 1 - - def output_is_reduced(self) -> bool: - return False - - def _quantize_input( - self, - a1: torch.Tensor, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - # Defer input quant to moe kernel for backends (e.g. AITER, FI) - # which use a single kernel call for quant + experts. - if defer_input_quant: - return a1, None - - input_sf = ( - quant_config.a1_gscale - if quant_config.use_nvfp4_w4a4 - else quant_config.a1_scale - ) - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - input_sf, - quant_dtype=quant_config.quant_dtype, - per_act_token_quant=quant_config.per_act_token_quant, - block_shape=quant_config.block_shape, - is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, - ) - - return a1q, a1q_scale - - -class MoEPrepareAndFinalizeNoEP( - mk.FusedMoEPrepareAndFinalizeModular, MoEPrepareAndFinalizeNoEPBase -): - def prepare( - self, - a1: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> mk.PrepareResultType: - if apply_router_weight_on_input: - topk = topk_ids.size(1) - # TODO: this only works for topK=1, will need to update for topK>1 - assert topk == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - # Note: do not use inplace for shared experts overlap - a1 = a1 * topk_weights.to(a1.dtype) - - a1q, a1q_scale = self._quantize_input(a1, quant_config, defer_input_quant) - - return a1q, a1q_scale, None, None, None - - def finalize( - self, - output: torch.Tensor, - fused_expert_output: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - apply_router_weight_on_input: bool, - weight_and_reduce_impl: mk.TopKWeightAndReduce, - ) -> None: - if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): - weight_and_reduce_impl = TopKWeightAndReduceContiguous() - weight_and_reduce_impl.apply( - output=output, - fused_expert_output=fused_expert_output, - topk_weights=topk_weights, - topk_ids=topk_ids, - apply_router_weight_on_input=apply_router_weight_on_input, - ) - - -class MoEPrepareAndFinalizeNoEPMonolithic( - mk.FusedMoEPrepareAndFinalizeMonolithic, MoEPrepareAndFinalizeNoEPBase -): - def prepare( - self, - a1: torch.Tensor, - router_logits: torch.Tensor, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> mk.PrepareMonolithicResultType: - a1q, a1q_scale = self._quantize_input(a1, quant_config, defer_input_quant) - return a1q, a1q_scale, router_logits - - def finalize( - self, - fused_expert_output: torch.Tensor, - ) -> torch.Tensor: - return fused_expert_output diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py new file mode 100644 index 000000000000..51f28547bc5e --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from naive_dp_ep import ( + MoEPrepareAndFinalizeNaiveDPEPModular, + MoEPrepareAndFinalizeNaiveDPEPMonolithic, + make_moe_prepare_and_finalize_naive_dp_ep, +) +from no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, + MoEPrepareAndFinalizeNoDPEPMonolithic, + make_moe_prepare_and_finalize_no_dp_ep, +) + +__all__ = [ + "MoEPrepareAndFinalizeNaiveDPEPMonolithic", + "MoEPrepareAndFinalizeNaiveDPEPModular", + make_moe_prepare_and_finalize_naive_dp_ep, + "MoEPrepareAndFinalizeNoDPEPMonolithic", + "MoEPrepareAndFinalizeNoDPEPModular", + "make_moe_prepare_and_finalize_no_dp_ep", +] diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py new file mode 100644 index 000000000000..6dc9f6958048 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.distributed import get_ep_group +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceContiguous, + TopKWeightAndReduceDelegate, +) +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.utils.flashinfer import nvfp4_block_scale_interleave + + +def _quantize_and_setup_dispatch( + a1: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, +) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + # Defer input quantization to the MoE kernel. + if defer_input_quant: + a1q = a1 + a1q_scale = None + else: + input_sf = ( + quant_config.a1_gscale + if quant_config.use_nvfp4_w4a4 + else quant_config.a1_scale + ) + + # NOTE: swizzling pads the scales to multiple of 128 + # which makes the scales tensor different shape than + # the hidden states, breaking the A2A kernel. So, we + # delay the swizzling until after the A2A. + a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input( + a1, + input_sf, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + is_fp4_scale_swizzled=False, + ) + + # Skip gathering scales if we have static quantization + # (the scale is a scalar, replicated on all ranks) or + # if quantization is deferred. + skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0 + scales = None if skip_gather_scales else [a1q_scale] + + return a1q, scales + + +def _unwrap_scale_and_prepare_for_moe( + scales: list[torch.Tensor] | None, + quant_config: FusedMoEQuantConfig, +) -> torch.Tensor: + assert scales is not None and len(scales) == 1 + a1q_scale = scales[0] + # Apply swizzling after a2a if the MoE kernel needs it. + if quant_config.quant_dtype == "nvfp4" and quant_config.is_nvfp4_scale_swizzled: + assert a1q_scale is not None + if a1q_scale.element_size() == 1: + a1q_scale = a1q_scale.view(torch.uint8) + a1q_scale = nvfp4_block_scale_interleave(a1q_scale) + + return a1q_scale + + +class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular): + """ + Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. + + Uses Torch AR/RS or AR for dispatch/combine operations, applied + to the topk weights and ids. + """ + + def __init__( + self, + is_sequence_parallel: bool = False, + num_dispatchers: int = 1, + ) -> None: + super().__init__() + self.is_sequence_parallel = is_sequence_parallel + self._num_dispatchers = num_dispatchers + + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def max_num_tokens_per_rank(self) -> int | None: + return None + + def topk_indices_dtype(self) -> torch.dtype | None: + return None + + def num_dispatchers(self) -> int: + return self._num_dispatchers + + def output_is_reduced(self) -> bool: + return False + + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareResultType: + """Quantize and Dispatch Topk Weights and Topk Ids.""" + + if apply_router_weight_on_input: + topk = topk_ids.size(1) + assert topk == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + # Note: do not use inplace for shared experts overlap + a1 = a1 * topk_weights.to(a1.dtype) + + a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant) + + res = get_ep_group().dispatch( + a1q, + topk_weights, + topk_ids, + is_sequence_parallel=self.is_sequence_parallel, + extra_tensors=scales, + ) + + if scales is None: + a1q, topk_weights, topk_ids = res + a1q_scale = None + else: + a1q, topk_weights, topk_ids, scales = res + a1q_scale = _unwrap_scale_and_prepare_for_moe(scales, quant_config) + + return a1q, a1q_scale, None, topk_ids, topk_weights + + def finalize( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: mk.TopKWeightAndReduce, + ) -> None: + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): + weight_and_reduce_impl = TopKWeightAndReduceContiguous() + + out = weight_and_reduce_impl.apply( + output=None, + fused_expert_output=fused_expert_output, + topk_weights=topk_weights, + topk_ids=topk_ids, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + output.copy_( + get_ep_group().combine(out, is_sequence_parallel=self.is_sequence_parallel) + ) + + +class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic): + """ + Naive Prepare/Finalize for Dp/Ep case for Modular Kernels. + + Uses Torch AR/RS or AR for dispatch/combine operations, applied + to the router logits (the MoE kernel runs the router internally). + """ + + def __init__( + self, + is_sequence_parallel: bool = False, + num_dispatchers: int = 1, + ) -> None: + super().__init__() + self.is_sequence_parallel = is_sequence_parallel + self._num_dispatchers = num_dispatchers + + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def max_num_tokens_per_rank(self) -> int | None: + return None + + def topk_indices_dtype(self) -> torch.dtype | None: + return None + + def num_dispatchers(self) -> int: + return self._num_dispatchers + + def output_is_reduced(self) -> bool: + return False + + def prepare( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareMonolithicResultType: + """Quantize and Dispatch Router Logits.""" + + a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant) + + res = get_ep_group().dispatch_router_logits( + a1q, + router_logits, + is_sequence_parallel=self.is_sequence_parallel, + extra_tensors=scales, + ) + + if scales is None: + a1q, router_logits = res + a1q_scale = None + else: + a1q, router_logits, scales = res + a1q_scale = _unwrap_scale_and_prepare_for_moe(scales, quant_config) + + return a1q, a1q_scale, router_logits + + def finalize( + self, + fused_expert_output: torch.Tensor, + ) -> torch.Tensor: + out = get_ep_group().combine( + fused_expert_output, is_sequence_parallel=self.is_sequence_parallel + ) + return out + + +def make_moe_prepare_and_finalize_naive_dp_ep( + use_monolithic: bool, + is_sequence_parallel: bool = False, + num_dispatchers: int = 1, +) -> MoEPrepareAndFinalizeNaiveDPEPModular | MoEPrepareAndFinalizeNaiveDPEPMonolithic: + return ( + MoEPrepareAndFinalizeNaiveDPEPMonolithic( + is_sequence_parallel=is_sequence_parallel, + num_dispatchers=num_dispatchers, + ) + if use_monolithic + else MoEPrepareAndFinalizeNaiveDPEPModular( + is_sequence_parallel=is_sequence_parallel, + num_dispatchers=num_dispatchers, + ) + ) 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 new file mode 100644 index 000000000000..b9d57da08326 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceContiguous, + TopKWeightAndReduceDelegate, +) +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input + + +def _quantize_input( + a1: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + # Defer input quant to moe kernel for backends (e.g. AITER, FI) + # which use a single kernel call for quant + experts. + if defer_input_quant: + return a1, None + + input_sf = ( + quant_config.a1_gscale if quant_config.use_nvfp4_w4a4 else quant_config.a1_scale + ) + a1q, a1q_scale = moe_kernel_quantize_input( + a1, + input_sf, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled, + ) + + return a1q, a1q_scale + + +class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular): + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def max_num_tokens_per_rank(self) -> int | None: + return None + + def topk_indices_dtype(self) -> torch.dtype | None: + return None + + def num_dispatchers(self) -> int: + return 1 + + def output_is_reduced(self) -> bool: + return False + + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareResultType: + if apply_router_weight_on_input: + topk = topk_ids.size(1) + # TODO: this only works for topK=1, will need to update for topK>1 + assert topk == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + # Note: do not use inplace for shared experts overlap + a1 = a1 * topk_weights.to(a1.dtype) + + a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant) + + return a1q, a1q_scale, None, None, None + + def finalize( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: mk.TopKWeightAndReduce, + ) -> None: + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): + weight_and_reduce_impl = TopKWeightAndReduceContiguous() + weight_and_reduce_impl.apply( + output=output, + fused_expert_output=fused_expert_output, + topk_weights=topk_weights, + topk_ids=topk_ids, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + +class MoEPrepareAndFinalizeNoDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic): + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def max_num_tokens_per_rank(self) -> int | None: + return None + + def topk_indices_dtype(self) -> torch.dtype | None: + return None + + def num_dispatchers(self) -> int: + return 1 + + def output_is_reduced(self) -> bool: + return False + + def prepare( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareMonolithicResultType: + a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant) + return a1q, a1q_scale, router_logits + + def finalize( + self, + fused_expert_output: torch.Tensor, + ) -> torch.Tensor: + return fused_expert_output + + +def make_moe_prepare_and_finalize_no_dp_ep( + use_monolithic: bool, +) -> MoEPrepareAndFinalizeNoDPEPModular | MoEPrepareAndFinalizeNoDPEPMonolithic: + return ( + MoEPrepareAndFinalizeNoDPEPMonolithic() + if use_monolithic + else MoEPrepareAndFinalizeNoDPEPModular() + ) diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index 69e3da4aa21e..4034512e133f 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -61,7 +61,7 @@ def apply( if output is None: return fused_expert_output - # MoEPrepareAndFinalizeNoEP needs the output to be in the `output` + # MoEPrepareAndFinalizeNoDPEPModular needs the output to be in the `output` # tensor. assert output.size() == fused_expert_output.size(), ( "output shape is expected to match the fused_expert_output shape. " From 38aebe908a4d64ae693cae00d5ae521978ca6cf4 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:01:17 -0500 Subject: [PATCH 129/207] revert relative import Signed-off-by: Robert Shaw --- .../layers/fused_moe/prepare_finalize/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py index 51f28547bc5e..03fea7c6d78b 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from naive_dp_ep import ( +from vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep import ( MoEPrepareAndFinalizeNaiveDPEPModular, MoEPrepareAndFinalizeNaiveDPEPMonolithic, make_moe_prepare_and_finalize_naive_dp_ep, ) -from no_dp_ep import ( +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( MoEPrepareAndFinalizeNoDPEPModular, MoEPrepareAndFinalizeNoDPEPMonolithic, make_moe_prepare_and_finalize_no_dp_ep, @@ -15,7 +15,7 @@ __all__ = [ "MoEPrepareAndFinalizeNaiveDPEPMonolithic", "MoEPrepareAndFinalizeNaiveDPEPModular", - make_moe_prepare_and_finalize_naive_dp_ep, + "make_moe_prepare_and_finalize_naive_dp_ep", "MoEPrepareAndFinalizeNoDPEPMonolithic", "MoEPrepareAndFinalizeNoDPEPModular", "make_moe_prepare_and_finalize_no_dp_ep", From fb1cae447b43a19485b6a2e44e4e644e32048979 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:09:01 -0500 Subject: [PATCH 130/207] move flashinfer monolithic kernels to experts folder Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 6 +++--- vllm/model_executor/layers/fused_moe/experts/__init__.py | 0 .../fused_moe/{ => experts}/flashinfer_trtllm_fp8_moe.py | 0 .../fused_moe/{ => experts}/flashinfer_trtllm_nvfp4_moe.py | 0 vllm/model_executor/layers/fused_moe/oracle/fp8.py | 2 +- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 2 +- 6 files changed, 5 insertions(+), 5 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/__init__.py rename vllm/model_executor/layers/fused_moe/{ => experts}/flashinfer_trtllm_fp8_moe.py (100%) rename vllm/model_executor/layers/fused_moe/{ => experts}/flashinfer_trtllm_nvfp4_moe.py (100%) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 5fe741b63616..d9c0db28b1c0 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -15,12 +15,12 @@ RoutingMethodType, fp8_w8a8_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_trtllm_fp8_moe import ( + FlashInferTrtLlmFp8Experts, +) from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, ) -from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( - FlashInferTrtLlmFp8Experts, -) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoDPEPModular, diff --git a/vllm/model_executor/layers/fused_moe/experts/__init__.py b/vllm/model_executor/layers/fused_moe/experts/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_fp8_moe.py similarity index 100% rename from vllm/model_executor/layers/fused_moe/flashinfer_trtllm_fp8_moe.py rename to vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_fp8_moe.py diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_nvfp4_moe.py similarity index 100% rename from vllm/model_executor/layers/fused_moe/flashinfer_trtllm_nvfp4_moe.py rename to vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_nvfp4_moe.py diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index d07e35df8100..e66930bded72 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -55,7 +55,7 @@ def backend_to_kernel_cls( backend: Fp8MoeBackend, ) -> type[mk.FusedMoEExperts]: if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: - from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_fp8_moe import ( + from vllm.model_executor.layers.fused_moe.experts.flashinfer_trtllm_fp8_moe import ( # noqa: E501 FlashInferTrtLlmFp8Experts, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 3db4f3cf7a95..866b6db4178c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -67,7 +67,7 @@ def backend_to_kernel_cls( backend: NvFp4MoeBackend, ) -> list[type[mk.FusedMoEExperts]]: if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_nvfp4_moe import ( + from vllm.model_executor.layers.fused_moe.experts.flashinfer_trtllm_nvfp4_moe import ( # noqa: E501 FlashInferTrtLlmNvFp4ExpertsModular, FlashInferTrtLlmNvFp4ExpertsMonolithic, ) From e1339833a8e0e670aa3d560094a0271d5cb8b45f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:15:10 -0500 Subject: [PATCH 131/207] reduce length of names Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 10 +- .../experts/flashinfer_trtllm_fp8_moe.py | 313 ------------------ .../experts/flashinfer_trtllm_nvfp4_moe.py | 291 ---------------- .../layers/fused_moe/oracle/fp8.py | 6 +- .../layers/fused_moe/oracle/nvfp4.py | 8 +- 5 files changed, 12 insertions(+), 616 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_fp8_moe.py delete mode 100644 vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_nvfp4_moe.py diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index d9c0db28b1c0..7adecc31d2db 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -15,8 +15,8 @@ RoutingMethodType, fp8_w8a8_moe_quant_config, ) -from vllm.model_executor.layers.fused_moe.experts.flashinfer_trtllm_fp8_moe import ( - FlashInferTrtLlmFp8Experts, +from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( + TrtLlmFp8Experts, ) from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, @@ -24,7 +24,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoDPEPModular, - MoEPrepareAndFinalizeNoDPEPModularMonolithic, + MoEPrepareAndFinalizeNoDPEPMonolithic, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, @@ -241,8 +241,8 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModularMonolithic(), - FlashInferTrtLlmFp8Experts( + MoEPrepareAndFinalizeNoDPEPMonolithic(), + TrtLlmFp8Experts( moe_config=td.layer.moe, quant_config=quant_config, ), diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_fp8_moe.py deleted file mode 100644 index 5181987d1135..000000000000 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_fp8_moe.py +++ /dev/null @@ -1,313 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import flashinfer -import torch - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEParallelConfig, - FusedMoEQuantConfig, - RoutingMethodType, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - QuantKey, - kFp8Dynamic128Sym, - kFp8Static128BlockSym, - kFp8StaticTensorSym, -) -from vllm.v1.engine.utils import current_platform - - -class FlashInferTrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic): - """ - Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface. - """ - - def __init__( - self, - moe_config: FusedMoEConfig, - quant_config: FusedMoEQuantConfig, - ): - super().__init__(moe_config, quant_config) - - if moe_config.moe_parallel_config.use_ep and quant_config.is_per_tensor: - raise NotImplementedError( - "EP parallelism is not supported with TRTLLM" - "per-tensor FP8 quantization." - ) - - self.routing_method_type = moe_config.routing_method - self.topk = moe_config.experts_per_token - self.intermediate_size_per_partition = ( - moe_config.intermediate_size_per_partition - ) - self.hidden_dim = moe_config.hidden_dim - self.local_num_experts = moe_config.num_local_experts - self.ep_rank = moe_config.moe_parallel_config.ep_rank - - # Make additional scales for per-tensor interface. - if self.quant_config.is_per_tensor: - w1_scale = self.quant_config.w1_scale - assert w1_scale is not None - a1_scale = self.quant_config.a1_scale - assert a1_scale is not None - w2_scale = self.quant_config.w2_scale - assert w2_scale is not None - a2_scale = self.quant_config.a2_scale - assert a2_scale is not None - - self._g1_alphas = (w1_scale * a1_scale).squeeze() - self._g2_alphas = (w2_scale * a2_scale).squeeze() - self._g1_scale_c = self._g1_alphas / self.quant_config.a2_scale - - @staticmethod - def activation_format() -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.Standard - - @staticmethod - def _supports_current_device() -> bool: - """Supports only Blackwell-family GPUs.""" - p = current_platform - # Add check flashinfer trtllm is available - return p.is_cuda() and p.is_device_capability_family(100) - - @staticmethod - def _supports_no_act_and_mul() -> bool: - """Does not support non-gated MoE (i.e. Nanotron-Mini).""" - return False - - @staticmethod - def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - """Supports Fp8 per-tensor and Fp8 block.""" - SUPPORTED_W_A = [ - (kFp8Static128BlockSym, kFp8Dynamic128Sym), - (kFp8StaticTensorSym, kFp8StaticTensorSym), - ] - return (weight_key, activation_key) in SUPPORTED_W_A - - @staticmethod - def _supports_activation(activation: MoEActivation) -> bool: - """Supports only SiLU and RELU^2 non-gated activation.""" - return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] - - @staticmethod - def _supports_routing_method( - routing_method: RoutingMethodType, - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - """Monolithic kernels need to express router support.""" - # NOTE(dbari): TopK routing could also be enabled, but need to validate models - # NOTE(dbari): Default is not implemented and should not be enabled until it is - if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): - # NOTE(rob): potentially allow others here. This is a conservative list. - return routing_method in [ - RoutingMethodType.DeepSeekV3, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ] - elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym): - # NOTE(dbari): as above, potentially allow others here. - return routing_method in [ - RoutingMethodType.DeepSeekV3, - RoutingMethodType.Llama4, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ] - else: - raise ValueError("Unsupported quantization scheme.") - - @staticmethod - def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """Monolithic kernel so only use with naive DP/EP and TP.""" - return ( - not moe_parallel_config.use_all2all_kernels - or moe_parallel_config.use_naive_all2all_kernels - ) and not moe_parallel_config.enable_eplb - - @staticmethod - def _supports_router_logits_dtype( - router_logits_dtype: torch.dtype | None, - routing_method: RoutingMethodType, - ) -> bool: - """ - The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default. - Only DeepSeekV3 routing supports float32 router_logits (which is converted - internally in the kernel). - """ - if router_logits_dtype == torch.float32: - # Only DeepSeekV3 routing handles float32 logits - # https://github.com/flashinfer-ai/flashinfer/issues/2469 - return routing_method == RoutingMethodType.DeepSeekV3 - return True - - def supports_chunking(self) -> bool: - return False - - def supports_expert_map(self) -> bool: - return False - - def _apply_per_block( - self, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - router_logits: torch.Tensor, - activation: MoEActivation, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - assert not apply_router_weight_on_input - assert activation == MoEActivation.SILU - assert ( - e_score_correction_bias is None - or e_score_correction_bias.dtype == hidden_states.dtype - ) - - if self.routing_method_type == RoutingMethodType.DeepSeekV3: - router_logits = router_logits.to(torch.float32) - - assert self.topk <= global_num_experts - assert self.topk <= 10 - assert global_num_experts % 4 == 0 - assert self.quant_config.block_shape == [128, 128] - # Routing kernel expects #experts <= #threads 512 - assert global_num_experts <= 512 - - # Kernel requires transposed hidden state scales - # TODO: fuse into the quant kernel. - assert a1q_scale is not None - a1q_scale_t = a1q_scale.t().contiguous() - - return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( - routing_logits=router_logits, - routing_bias=e_score_correction_bias, - hidden_states=hidden_states, - hidden_states_scale=a1q_scale_t, - gemm1_weights=w1, - gemm1_weights_scale=self.quant_config.w1_scale, - gemm2_weights=w2, - gemm2_weights_scale=self.quant_config.w2_scale, - num_experts=global_num_experts, - top_k=self.topk, - n_group=num_expert_group, - topk_group=(topk_group or 0), - intermediate_size=self.intermediate_size_per_partition, - local_expert_offset=self.ep_rank * self.local_num_experts, - local_num_experts=self.local_num_experts, - routed_scaling_factor=routed_scaling_factor, - routing_method_type=self.routing_method_type, - ) - - def _apply_per_tensor( - self, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - router_logits: torch.Tensor, - activation: str, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - assert self.routing_method_type == RoutingMethodType.Llama4 - assert apply_router_weight_on_input - - # Should only have Llama4 routing here. - assert routed_scaling_factor is not None - assert e_score_correction_bias is None - assert num_expert_group is None - - out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( - routing_logits=router_logits, - routing_bias=e_score_correction_bias, - hidden_states=hidden_states, - gemm1_weights=w1, - output1_scales_scalar=self._g1_scale_c, - output1_scales_gate_scalar=self._g1_alphas, - gemm2_weights=w2, - output2_scales_scalar=self._g2_alphas, - num_experts=global_num_experts, - top_k=self.topk, - n_group=num_expert_group or 0, - topk_group=topk_group or 0, - intermediate_size=self.intermediate_size_per_partition, - local_expert_offset=self.ep_rank * self.local_num_experts, - local_num_experts=self.local_num_experts, - routed_scaling_factor=routed_scaling_factor, - use_routing_scales_on_input=apply_router_weight_on_input, - routing_method_type=self.routing_method_type, - ) - return out - - def apply( - self, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - router_logits: torch.Tensor, - activation: str, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - if self.quant_config.block_shape is not None: - return self._apply_per_block( - hidden_states, - w1, - w2, - router_logits, - activation, - global_num_experts, - expert_map, - a1q_scale, - apply_router_weight_on_input, - num_expert_group=num_expert_group, - e_score_correction_bias=e_score_correction_bias, - routed_scaling_factor=routed_scaling_factor, - ) - elif self.quant_config.is_per_tensor: - return self._apply_per_tensor( - hidden_states, - w1, - w2, - router_logits, - activation, - global_num_experts, - expert_map, - a1q_scale, - apply_router_weight_on_input, - num_expert_group=num_expert_group, - e_score_correction_bias=e_score_correction_bias, - routed_scaling_factor=routed_scaling_factor, - ) - else: - raise NotImplementedError( - "Only per-block and per-tensor quantization are supported in " - f"{self.__class__.__name__}." - ) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_nvfp4_moe.py deleted file mode 100644 index 8427b59da4bd..000000000000 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_trtllm_nvfp4_moe.py +++ /dev/null @@ -1,291 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import flashinfer -import torch - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEParallelConfig, - FusedMoEQuantConfig, - RoutingMethodType, -) -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceNoOP, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - QuantKey, - kNvfp4Dynamic, - kNvfp4Static, -) -from vllm.platforms import current_platform - - -class FlashInferTrtLlmNvFp4ExpertsBase: - """ - NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface. - """ - - def __init__( - self, - moe_config: FusedMoEConfig, - quant_config: FusedMoEQuantConfig, - ): - self.moe_config = moe_config - self.quant_config = quant_config - - self.routing_method_type = self.moe_config.routing_method - self.topk = moe_config.experts_per_token - self.intermediate_size_per_partition = ( - moe_config.intermediate_size_per_partition - ) - self.hidden_dim = moe_config.hidden_dim - self.local_num_experts = moe_config.num_local_experts - self.ep_rank = moe_config.moe_parallel_config.ep_rank - - # g1_alpha_s = a13_scale * w13_scale_2 - # a2_gscale = (1 / a2_scale) - # g1_scale_c = a13_scale * w13_scale_2 / a2_scale - assert self.quant_config.g1_alphas is not None - assert self.quant_config.a2_gscale is not None - self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale - - @staticmethod - def _supports_current_device() -> bool: - """Supports only Blackwell-family GPUs.""" - p = current_platform - return p.is_cuda() and p.is_device_capability_family(100) - - @staticmethod - def _supports_no_act_and_mul() -> bool: - """Supports non-gated MoE (i.e. Nemotron-Nano).""" - return True - - @staticmethod - def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - """Supports Nvfp4 quantization.""" - SUPPORTED_W_A = [ - (kNvfp4Static, kNvfp4Dynamic), - ] - return (weight_key, activation_key) in SUPPORTED_W_A - - @staticmethod - def _supports_activation(activation: MoEActivation) -> bool: - """Supports only SiLU and RELU^2 non-gated activation.""" - return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] - - @staticmethod - def activation_format() -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.Standard - - def supports_chunking(self) -> bool: - return False - - def supports_expert_map(self) -> bool: - return False - - -class FlashInferTrtLlmNvFp4ExpertsModular( - FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular -): - """ - Modular version of the implementation (just the experts). - """ - - def workspace_shapes( - self, - M: int, - N: int, - K: int, - topk: int, - global_num_experts: int, - local_num_experts: int, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - activation: str, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - # The workspaces for this implementation are managed by flashinfer. - workspace1 = (0,) - workspace2 = (0,) - - # Hidden states are Nvfp4, packed into int8 dtype, so we - # need to multiply K by 2 to get the output shape right. - assert self.hidden_dim == K * 2 - output = (M, self.hidden_dim) - - return (workspace1, workspace2, output) - - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - return TopKWeightAndReduceNoOP() - - def apply( - self, - output: torch.Tensor, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: MoEActivation, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, - workspace13: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - apply_router_weight_on_input: bool, - ): - assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] - assert a1q_scale is not None - assert self.quant_config.w1_scale is not None - assert self.quant_config.w2_scale is not None - - # Pack topk ids and weights into format expected by the kernel. - packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to( - torch.bfloat16 - ).view(torch.int16) - - # Invoke kernel. - flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( - topk_ids=packed_tensor, - routing_bias=None, - hidden_states=hidden_states, - hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), - gemm1_weights=w1, - gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=w2, - gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), - gemm2_bias=None, - output1_scale_scalar=self.g1_scale_c, - output1_scale_gate_scalar=self.quant_config.g1_alphas, - output2_scale_scalar=self.quant_config.g2_alphas, - num_experts=global_num_experts, - top_k=self.topk, - n_group=0, - topk_group=0, - intermediate_size=self.intermediate_size_per_partition, - local_expert_offset=self.ep_rank * self.local_num_experts, - local_num_experts=self.local_num_experts, - routed_scaling_factor=None, - routing_method_type=1, - do_finalize=True, - output=output, - ) - - -class FlashInferTrtLlmNvFp4ExpertsMonolithic( - FlashInferTrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic -): - """ - Monolithic version of the kernel (router + experts). - """ - - @staticmethod - def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - """The modular implementation should be used for the Dp/Ep or EPLB case.""" - return ( - not moe_parallel_config.use_all2all_kernels - and not moe_parallel_config.enable_eplb - ) - - @staticmethod - def _supports_routing_method( - routing_method_type: RoutingMethodType, - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - # NOTE(rob): this is a conservative list. - return routing_method_type in [ - RoutingMethodType.DeepSeekV3, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - RoutingMethodType.Llama4, - ] - - @staticmethod - def _supports_router_logits_dtype( - router_logits_dtype: torch.dtype | None, - routing_method: RoutingMethodType, - ) -> bool: - """ - The FlashInfer TRTLLM NVFp4 kernel expects bfloat16 router_logits by default. - Only DeepSeekV3 routing supports float32 router_logits (which is converted - internally in the kernel). - """ - # TODO: check this - if router_logits_dtype == torch.float32: - # Only DeepSeekV3 routing handles float32 logits - # https://github.com/flashinfer-ai/flashinfer/issues/2469 - return routing_method == RoutingMethodType.DeepSeekV3 - return True - - def apply( - self, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - router_logits: torch.Tensor, - activation: MoEActivation, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] - assert a1q_scale is not None - assert self.quant_config.w1_scale is not None - assert self.quant_config.w2_scale is not None - - # Prepare routing bias into kernel format. - routing_bias = e_score_correction_bias - if routing_bias is not None: - routing_bias = routing_bias.to(torch.bfloat16) - router_logits = ( - router_logits.to(torch.float32) - if self.routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits - ) - - # Invoke kernel. - return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( - routing_logits=router_logits, - routing_bias=routing_bias, - hidden_states=hidden_states, - hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), - gemm1_weights=w1, - gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=w2, - gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), - gemm2_bias=None, - output1_scale_scalar=self.g1_scale_c, - output1_scale_gate_scalar=self.quant_config.g1_alphas, - output2_scale_scalar=self.quant_config.g2_alphas, - num_experts=global_num_experts, - top_k=self.topk, - n_group=(num_expert_group or 0), - topk_group=(topk_group or 0), - intermediate_size=self.intermediate_size_per_partition, - local_expert_offset=self.ep_rank * self.local_num_experts, - local_num_experts=self.local_num_experts, - routed_scaling_factor=None, - routing_method_type=self.routing_method_type, - do_finalize=True, - )[0] diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index e66930bded72..c6ed4bcdfebd 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -55,11 +55,11 @@ def backend_to_kernel_cls( backend: Fp8MoeBackend, ) -> type[mk.FusedMoEExperts]: if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: - from vllm.model_executor.layers.fused_moe.experts.flashinfer_trtllm_fp8_moe import ( # noqa: E501 - FlashInferTrtLlmFp8Experts, + from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( # noqa: E501 + TrtLlmFp8Experts, ) - return FlashInferTrtLlmFp8Experts + return TrtLlmFp8Experts elif backend == Fp8MoeBackend.FLASHINFER_CUTLASS: from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 866b6db4178c..9213fab566f4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -68,14 +68,14 @@ def backend_to_kernel_cls( ) -> list[type[mk.FusedMoEExperts]]: if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.experts.flashinfer_trtllm_nvfp4_moe import ( # noqa: E501 - FlashInferTrtLlmNvFp4ExpertsModular, - FlashInferTrtLlmNvFp4ExpertsMonolithic, + TrtLlmNvFp4ExpertsModular, + TrtLlmNvFp4ExpertsMonolithic, ) # NOTE: prefer Monolthic > Modular, so return Monolithic first. return [ - FlashInferTrtLlmNvFp4ExpertsMonolithic, - FlashInferTrtLlmNvFp4ExpertsModular, + TrtLlmNvFp4ExpertsMonolithic, + TrtLlmNvFp4ExpertsModular, ] elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS: From c196803094355a8117a829366e0bfc7d6d62c230 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:18:52 -0500 Subject: [PATCH 132/207] remove commebnts Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/modular_kernel.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 4309377a05f9..4fa3b9fdafbb 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -173,7 +173,6 @@ def apply( ################################################################################ # Prepare/Finalize -# TODO: make the below a separate file. ################################################################################ @@ -457,7 +456,6 @@ def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor: ################################################################################ # Experts -# TODO: make the below a separate file. ################################################################################ @@ -982,8 +980,7 @@ def _slice_scales( ################################################################################ -# MoEKernel -# TODO: make the below a separate file. +# Kernel ################################################################################ @@ -1516,8 +1513,6 @@ def apply( expert_map, apply_router_weight_on_input, ) - # print(f"{a1q.dtype=}") - # print(f"{a1q_scale.dtype=}") fused_out = self._fused_experts( in_dtype=hidden_states.dtype, From 361ec4b96213cb59bda8861e586f40e2a93f330c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:19:15 -0500 Subject: [PATCH 133/207] missed hitting save Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 9213fab566f4..d65b52a374b0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -67,7 +67,7 @@ def backend_to_kernel_cls( backend: NvFp4MoeBackend, ) -> list[type[mk.FusedMoEExperts]]: if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - from vllm.model_executor.layers.fused_moe.experts.flashinfer_trtllm_nvfp4_moe import ( # noqa: E501 + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( TrtLlmNvFp4ExpertsModular, TrtLlmNvFp4ExpertsMonolithic, ) From be3db44a0d20c1951dd658e928a246f6bdfdc3ef Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:19:57 -0500 Subject: [PATCH 134/207] add back missing files Signed-off-by: Robert Shaw --- .../fused_moe/experts/trtllm_fp8_moe.py | 313 ++++++++++++++++++ .../fused_moe/experts/trtllm_nvfp4_moe.py | 294 ++++++++++++++++ 2 files changed, 607 insertions(+) create mode 100644 vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py create mode 100644 vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py 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 new file mode 100644 index 000000000000..3b614120561d --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import flashinfer +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + kFp8StaticTensorSym, +) +from vllm.v1.engine.utils import current_platform + + +class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic): + """ + Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + + if moe_config.moe_parallel_config.use_ep and quant_config.is_per_tensor: + raise NotImplementedError( + "EP parallelism is not supported with TRTLLM" + "per-tensor FP8 quantization." + ) + + self.routing_method_type = moe_config.routing_method + self.topk = moe_config.experts_per_token + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.hidden_dim = moe_config.hidden_dim + self.local_num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + # Make additional scales for per-tensor interface. + if self.quant_config.is_per_tensor: + w1_scale = self.quant_config.w1_scale + assert w1_scale is not None + a1_scale = self.quant_config.a1_scale + assert a1_scale is not None + w2_scale = self.quant_config.w2_scale + assert w2_scale is not None + a2_scale = self.quant_config.a2_scale + assert a2_scale is not None + + self._g1_alphas = (w1_scale * a1_scale).squeeze() + self._g2_alphas = (w2_scale * a2_scale).squeeze() + self._g1_scale_c = self._g1_alphas / self.quant_config.a2_scale + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + """Supports only Blackwell-family GPUs.""" + p = current_platform + # Add check flashinfer trtllm is available + return p.is_cuda() and p.is_device_capability_family(100) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + """Does not support non-gated MoE (i.e. Nanotron-Mini).""" + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Supports Fp8 per-tensor and Fp8 block.""" + SUPPORTED_W_A = [ + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + """Supports only SiLU and RELU^2 non-gated activation.""" + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Monolithic kernels need to express router support.""" + # NOTE(dbari): TopK routing could also be enabled, but need to validate models + # NOTE(dbari): Default is not implemented and should not be enabled until it is + if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): + # NOTE(rob): potentially allow others here. This is a conservative list. + return routing_method in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym): + # NOTE(dbari): as above, potentially allow others here. + return routing_method in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Llama4, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + else: + raise ValueError("Unsupported quantization scheme.") + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """Monolithic kernel so only use with naive DP/EP and TP.""" + return ( + not moe_parallel_config.use_all2all_kernels + or moe_parallel_config.use_naive_all2all_kernels + ) and not moe_parallel_config.enable_eplb + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default. + Only DeepSeekV3 routing supports float32 router_logits (which is converted + internally in the kernel). + """ + if router_logits_dtype == torch.float32: + # Only DeepSeekV3 routing handles float32 logits + # https://github.com/flashinfer-ai/flashinfer/issues/2469 + return routing_method == RoutingMethodType.DeepSeekV3 + return True + + def supports_chunking(self) -> bool: + return False + + def supports_expert_map(self) -> bool: + return False + + def _apply_per_block( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + assert not apply_router_weight_on_input + assert activation == MoEActivation.SILU + assert ( + e_score_correction_bias is None + or e_score_correction_bias.dtype == hidden_states.dtype + ) + + if self.routing_method_type == RoutingMethodType.DeepSeekV3: + router_logits = router_logits.to(torch.float32) + + assert self.topk <= global_num_experts + assert self.topk <= 10 + assert global_num_experts % 4 == 0 + assert self.quant_config.block_shape == [128, 128] + # Routing kernel expects #experts <= #threads 512 + assert global_num_experts <= 512 + + # Kernel requires transposed hidden state scales + # TODO: fuse into the quant kernel. + assert a1q_scale is not None + a1q_scale_t = a1q_scale.t().contiguous() + + return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( + routing_logits=router_logits, + routing_bias=e_score_correction_bias, + hidden_states=hidden_states, + hidden_states_scale=a1q_scale_t, + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale, + num_experts=global_num_experts, + top_k=self.topk, + n_group=num_expert_group, + topk_group=(topk_group or 0), + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=routed_scaling_factor, + routing_method_type=self.routing_method_type, + ) + + def _apply_per_tensor( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + assert self.routing_method_type == RoutingMethodType.Llama4 + assert apply_router_weight_on_input + + # Should only have Llama4 routing here. + assert routed_scaling_factor is not None + assert e_score_correction_bias is None + assert num_expert_group is None + + out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( + routing_logits=router_logits, + routing_bias=e_score_correction_bias, + hidden_states=hidden_states, + gemm1_weights=w1, + output1_scales_scalar=self._g1_scale_c, + output1_scales_gate_scalar=self._g1_alphas, + gemm2_weights=w2, + output2_scales_scalar=self._g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + n_group=num_expert_group or 0, + topk_group=topk_group or 0, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=routed_scaling_factor, + use_routing_scales_on_input=apply_router_weight_on_input, + routing_method_type=self.routing_method_type, + ) + return out + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + if self.quant_config.block_shape is not None: + return self._apply_per_block( + hidden_states, + w1, + w2, + router_logits, + activation, + global_num_experts, + expert_map, + a1q_scale, + apply_router_weight_on_input, + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + elif self.quant_config.is_per_tensor: + return self._apply_per_tensor( + hidden_states, + w1, + w2, + router_logits, + activation, + global_num_experts, + expert_map, + a1q_scale, + apply_router_weight_on_input, + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + else: + raise NotImplementedError( + "Only per-block and per-tensor quantization are supported in " + f"{self.__class__.__name__}." + ) 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 new file mode 100644 index 000000000000..be358a125ac5 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import flashinfer +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.platforms import current_platform + + +class TrtLlmNvFp4ExpertsBase: + """ + NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + self.moe_config = moe_config + self.quant_config = quant_config + + self.routing_method_type = self.moe_config.routing_method + self.topk = moe_config.experts_per_token + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.hidden_dim = moe_config.hidden_dim + self.local_num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + # g1_alpha_s = a13_scale * w13_scale_2 + # a2_gscale = (1 / a2_scale) + # g1_scale_c = a13_scale * w13_scale_2 / a2_scale + assert self.quant_config.g1_alphas is not None + assert self.quant_config.a2_gscale is not None + self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale + + @staticmethod + def _supports_current_device() -> bool: + """Supports only Blackwell-family GPUs.""" + p = current_platform + return p.is_cuda() and p.is_device_capability_family(100) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + """Supports non-gated MoE (i.e. Nemotron-Nano).""" + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Supports Nvfp4 quantization.""" + SUPPORTED_W_A = [ + (kNvfp4Static, kNvfp4Dynamic), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + """Supports only SiLU and RELU^2 non-gated activation.""" + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def supports_chunking(self) -> bool: + return False + + def supports_expert_map(self) -> bool: + return False + + +class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular): + """ + Modular version of the implementation (just the experts). + """ + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """The modular implementation supports all parallel configs.""" + return True + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: str, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # The workspaces for this implementation are managed by flashinfer. + workspace1 = (0,) + workspace2 = (0,) + + # Hidden states are Nvfp4, packed into int8 dtype, so we + # need to multiply K by 2 to get the output shape right. + assert self.hidden_dim == K * 2 + output = (M, self.hidden_dim) + + return (workspace1, workspace2, output) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + assert a1q_scale is not None + assert self.quant_config.w1_scale is not None + assert self.quant_config.w2_scale is not None + + # Pack topk ids and weights into format expected by the kernel. + packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to( + torch.bfloat16 + ).view(torch.int16) + + # Invoke kernel. + flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_tensor, + routing_bias=None, + hidden_states=hidden_states, + hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c, + output1_scale_gate_scalar=self.quant_config.g1_alphas, + output2_scale_scalar=self.quant_config.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + n_group=0, + topk_group=0, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=None, + routing_method_type=1, + do_finalize=True, + output=output, + ) + + +class TrtLlmNvFp4ExpertsMonolithic( + TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic +): + """ + Monolithic version of the kernel (router + experts). + """ + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """The modular implementation should be used for the Dp/Ep or EPLB case.""" + return ( + not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + ) + + @staticmethod + def _supports_routing_method( + routing_method_type: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # NOTE(rob): this is a conservative list. + return routing_method_type in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.Llama4, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + The FlashInfer TRTLLM NVFp4 kernel expects bfloat16 router_logits by default. + Only DeepSeekV3 routing supports float32 router_logits (which is converted + internally in the kernel). + """ + # TODO: check this + if router_logits_dtype == torch.float32: + # Only DeepSeekV3 routing handles float32 logits + # https://github.com/flashinfer-ai/flashinfer/issues/2469 + return routing_method == RoutingMethodType.DeepSeekV3 + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + assert a1q_scale is not None + assert self.quant_config.w1_scale is not None + assert self.quant_config.w2_scale is not None + + # Prepare routing bias into kernel format. + routing_bias = e_score_correction_bias + if routing_bias is not None: + routing_bias = routing_bias.to(torch.bfloat16) + router_logits = ( + router_logits.to(torch.float32) + if self.routing_method_type == RoutingMethodType.DeepSeekV3 + else router_logits + ) + + # Invoke kernel. + return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( + routing_logits=router_logits, + routing_bias=routing_bias, + hidden_states=hidden_states, + hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).flatten(), + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c, + output1_scale_gate_scalar=self.quant_config.g1_alphas, + output2_scale_scalar=self.quant_config.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + n_group=(num_expert_group or 0), + topk_group=(topk_group or 0), + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.ep_rank * self.local_num_experts, + local_num_experts=self.local_num_experts, + routed_scaling_factor=None, + routing_method_type=self.routing_method_type, + do_finalize=True, + )[0] From 0c0e943a70b9d37271cd1e4f7b038f300810ed13 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:23:35 -0500 Subject: [PATCH 135/207] update to remove making a NoDPEPModular Signed-off-by: Robert Shaw --- benchmarks/kernels/benchmark_cutlass_moe_fp8.py | 6 +++--- benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py | 6 +++--- benchmarks/kernels/benchmark_grouped_gemm_cutlass.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index ad5ed13a38dd..2f7822ed9706 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -15,8 +15,8 @@ from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk -from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( - MoEPrepareAndFinalizeNoDPEPModular, +from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + make_moe_prepare_and_finalize_no_dp_ep, ) from vllm.platforms import current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -138,7 +138,7 @@ def bench_run( ) fn = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + make_moe_prepare_and_finalize_no_dp_ep(False), CutlassExpertsFp8( moe_config=make_dummy_moe_config( num_experts=num_experts, diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index beba7a08a247..bd8eaf8a3652 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -24,7 +24,7 @@ ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, + make_moe_prepare_and_finalize_no_dp_ep, ) from vllm.scalar_type import scalar_types from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -197,7 +197,7 @@ def run_cutlass_moe_fp4( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + make_moe_prepare_and_finalize_no_dp_ep(False), CutlassExpertsFp4( make_dummy_moe_config(), quant_config=quant_config, @@ -242,7 +242,7 @@ def run_cutlass_from_graph( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + make_moe_prepare_and_finalize_no_dp_ep(False), CutlassExpertsFp4( make_dummy_moe_config(), quant_config=quant_config, diff --git a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py index 113babac1427..f8fc46781928 100644 --- a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py +++ b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py @@ -16,7 +16,7 @@ fused_topk, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, + make_moe_prepare_and_finalize_no_dp_ep, ) from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.worker.workspace import init_workspace_manager @@ -133,7 +133,7 @@ def run_cutlass_moe( ) fn = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + make_moe_prepare_and_finalize_no_dp_ep(use_monolithic=False), CutlassExpertsFp8( moe_config=make_dummy_moe_config( num_experts=w2.shape[0], @@ -165,7 +165,7 @@ def run_cutlass_from_graph( ) fn = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + make_moe_prepare_and_finalize_no_dp_ep(False), CutlassExpertsFp8( moe_config=make_dummy_moe_config( num_experts=w2.shape[0], From 2d6581c472051f3efd9d88e547c19a54900e6bd8 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:30:15 -0500 Subject: [PATCH 136/207] convert to using maybe_make_prepare_finalize in the tests Signed-off-by: Robert Shaw --- .../kernels/benchmark_cutlass_moe_fp8.py | 26 +++++++---- .../kernels/benchmark_cutlass_moe_nvfp4.py | 31 ++++++++++--- .../kernels/benchmark_grouped_gemm_cutlass.py | 46 ++++++++++++------- 3 files changed, 69 insertions(+), 34 deletions(-) diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index 2f7822ed9706..bd116e36a716 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -12,12 +12,12 @@ from tests.kernels.moe.utils import make_dummy_moe_config from vllm import _custom_ops as ops from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - make_moe_prepare_and_finalize_no_dp_ep, -) from vllm.platforms import current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.worker.workspace import init_workspace_manager @@ -137,15 +137,21 @@ def bench_run( per_out_ch_quant=per_out_ch, ) + moe_config = make_dummy_moe_config( + num_experts=num_experts, + hidden_dim=k, + intermediate_size_per_partition=n, + in_dtype=a.dtype, + ) fn = mk.FusedMoEKernel( - make_moe_prepare_and_finalize_no_dp_ep(False), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), CutlassExpertsFp8( - moe_config=make_dummy_moe_config( - num_experts=num_experts, - hidden_dim=k, - intermediate_size_per_partition=n, - in_dtype=a.dtype, - ), + moe_config=moe_config, quant_config=quant_config, ), ) diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index bd8eaf8a3652..cfb1489dadf2 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -15,6 +15,9 @@ from tests.kernels.moe.utils import make_dummy_moe_config from vllm import _custom_ops as ops from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import ( fp8_w8a8_moe_quant_config, nvfp4_moe_quant_config, @@ -23,9 +26,6 @@ CutlassExpertsFp4, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - make_moe_prepare_and_finalize_no_dp_ep, -) from vllm.scalar_type import scalar_types from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.worker.workspace import init_workspace_manager @@ -196,10 +196,21 @@ def run_cutlass_moe_fp4( g2_alphas=w2_gs, ) + moe_config = make_dummy_moe_config( + num_experts=num_experts, + hidden_dim=k, + intermediate_size_per_partition=n, + in_dtype=a.dtype, + ) kernel = mk.FusedMoEKernel( - make_moe_prepare_and_finalize_no_dp_ep(False), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), CutlassExpertsFp4( - make_dummy_moe_config(), + moe_config=moe_config, quant_config=quant_config, ), ) @@ -240,11 +251,17 @@ def run_cutlass_from_graph( g1_alphas=w1_gs, g2_alphas=w2_gs, ) + moe_config = make_dummy_moe_config() kernel = mk.FusedMoEKernel( - make_moe_prepare_and_finalize_no_dp_ep(False), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), CutlassExpertsFp4( - make_dummy_moe_config(), + moe_config=moe_config, quant_config=quant_config, ), ) diff --git a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py index f8fc46781928..60ec94b878ce 100644 --- a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py +++ b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py @@ -9,15 +9,15 @@ from tests.kernels.moe.utils import make_dummy_moe_config from vllm import _custom_ops as ops from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8 from vllm.model_executor.layers.fused_moe.fused_moe import ( fused_experts, fused_topk, ) -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - make_moe_prepare_and_finalize_no_dp_ep, -) from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.worker.workspace import init_workspace_manager @@ -131,16 +131,22 @@ def run_cutlass_moe( w2_scale=w2_scale, per_act_token_quant=per_act_token, ) + moe_config = make_dummy_moe_config( + num_experts=w2.shape[0], + hidden_dim=w2.shape[1], + intermediate_size_per_partition=w2.shape[2], + in_dtype=a.dtype, + ) fn = mk.FusedMoEKernel( - make_moe_prepare_and_finalize_no_dp_ep(use_monolithic=False), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), CutlassExpertsFp8( - moe_config=make_dummy_moe_config( - num_experts=w2.shape[0], - hidden_dim=w2.shape[1], - intermediate_size_per_partition=w2.shape[2], - in_dtype=a.dtype, - ), + moe_config=moe_config, quant_config=quant_config, ), ) @@ -163,16 +169,22 @@ def run_cutlass_from_graph( w2_scale=w2_scale, per_act_token_quant=per_act_token, ) + moe_config = make_dummy_moe_config( + num_experts=w2.shape[0], + hidden_dim=w2.shape[1], + intermediate_size_per_partition=w2.shape[2], + in_dtype=a.dtype, + ) fn = mk.FusedMoEKernel( - make_moe_prepare_and_finalize_no_dp_ep(False), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), CutlassExpertsFp8( - moe_config=make_dummy_moe_config( - num_experts=w2.shape[0], - hidden_dim=w2.shape[1], - intermediate_size_per_partition=w2.shape[2], - in_dtype=a.dtype, - ), + moe_config=moe_config, quant_config=quant_config, ), ) From 3cf2b66a095d9ff13bdc13b86c2aeaaf73bf82b0 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:33:54 -0500 Subject: [PATCH 137/207] convert to using maybe_make_prepare_finalize in the tests Signed-off-by: Robert Shaw --- benchmarks/kernels/benchmark_moe.py | 39 ++++++++++++++++++----------- tests/kernels/moe/test_block_fp8.py | 16 ++++++++---- 2 files changed, 36 insertions(+), 19 deletions(-) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 729cfd4459d9..c1a3e8d9cef6 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -17,6 +17,9 @@ from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, @@ -242,22 +245,30 @@ def run(): deep_gemm_experts = None if use_deep_gemm: + moe_config = ( + FusedMoEConfig( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=hidden_size, + intermediate_size_per_partition=shard_intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + activation=MoEActivation.SILU, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=init_dtype, + routing_method=RoutingMethodType.TopK, + device="cuda", + ), + ) deep_gemm_experts = mk.FusedMoEKernel( - prepare_finalize=MoEPrepareAndFinalizeNoDPEPModular(), + prepare_finalize=maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), fused_experts=TritonOrDeepGemmExperts( - moe_config=FusedMoEConfig( - num_experts=num_experts, - experts_per_token=topk, - hidden_dim=hidden_size, - intermediate_size_per_partition=shard_intermediate_size, - num_local_experts=num_experts, - num_logical_experts=num_experts, - activation=MoEActivation.SILU, - moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), - in_dtype=init_dtype, - routing_method=RoutingMethodType.TopK, - device="cuda", - ), + moe_config=moe_config, quant_config=quant_config, ), ) diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index cda032a5ddb7..1c70027d48b8 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -21,15 +21,15 @@ fused_experts, fused_topk, ) +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import ( fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.deep_gemm_moe import ( _valid_deep_gemm_shape, ) -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, -) from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, ) @@ -252,11 +252,17 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) w2_scale=w2_s, block_shape=block_size, ) + moe_config = make_dummy_moe_config() deep_gemm_experts = mk.FusedMoEKernel( - prepare_finalize=MoEPrepareAndFinalizeNoDPEPModular(), + prepare_finalize=maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), fused_experts=TritonOrDeepGemmExperts( - moe_config=make_dummy_moe_config(), + moe_config=moe_config, quant_config=quant_config, ), inplace=False, From ad457cc712ce74940d193f265d3b0473fee39526 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:38:18 -0500 Subject: [PATCH 138/207] convert to using maybe_make_prepare_finalize in the tests Signed-off-by: Robert Shaw --- tests/kernels/moe/test_cutlass_moe.py | 46 ++++++++++++------- tests/kernels/moe/test_deepgemm.py | 16 +++++-- tests/kernels/moe/test_flashinfer.py | 21 ++++++--- tests/kernels/moe/test_flashinfer_moe.py | 13 ++++-- .../moe/test_modular_oai_triton_moe.py | 18 +++++--- tests/kernels/moe/test_nvfp4_moe.py | 16 +++++-- tests/kernels/moe/utils.py | 13 ++++-- 7 files changed, 96 insertions(+), 47 deletions(-) diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index eaed0782725a..81c6a74fc7c0 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -13,6 +13,9 @@ from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import ( FUSED_MOE_UNQUANTIZED_CONFIG, FusedMoEQuantConfig, @@ -22,9 +25,6 @@ CutlassExpertsFp8, run_cutlass_moe_fp8, ) -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, -) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -197,15 +197,21 @@ def slice_experts(): for kwargs, new_quant_config in slice_experts(): w2 = kwargs["w2"] a = kwargs["hidden_states"] + moe_config = make_dummy_moe_config( + num_experts=w2.shape[0], + hidden_dim=w2.shape[1], + intermediate_size_per_partition=w2.shape[2], + in_dtype=a.dtype, + ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=new_quant_config, + allow_new_interface=True, + use_monolithic=False, + ), CutlassExpertsFp8( - moe_config=make_dummy_moe_config( - num_experts=w2.shape[0], - hidden_dim=w2.shape[1], - intermediate_size_per_partition=w2.shape[2], - in_dtype=a.dtype, - ), + moe_config=moe_config, quant_config=new_quant_config, ), inplace=False, @@ -257,15 +263,21 @@ def run_8_bit( num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined] with_ep = num_local_experts is not None or num_local_experts == num_experts if not with_ep: + moe_config = make_dummy_moe_config( + num_experts=moe_tensors.w2_q.shape[0], # type: ignore[union-attr] + hidden_dim=moe_tensors.w2_q.shape[1], # type: ignore[union-attr] + intermediate_size_per_partition=moe_tensors.w2_q.shape[2], # type: ignore[union-attr] + in_dtype=moe_tensors.a.dtype, + ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), CutlassExpertsFp8( - moe_config=make_dummy_moe_config( - num_experts=moe_tensors.w2_q.shape[0], # type: ignore[union-attr] - hidden_dim=moe_tensors.w2_q.shape[1], # type: ignore[union-attr] - intermediate_size_per_partition=moe_tensors.w2_q.shape[2], # type: ignore[union-attr] - in_dtype=moe_tensors.a.dtype, - ), + moe_config=moe_config, quant_config=quant_config, ), inplace=False, diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 65a52a055c10..7851d9ea0abe 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -14,13 +14,13 @@ # vLLM fused-expert reference (Triton fallback + DeepGEMM option) import vllm.model_executor.layers.fused_moe.modular_kernel as mk from tests.kernels.moe.utils import make_dummy_moe_config +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import ( fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, -) from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, ) @@ -108,11 +108,17 @@ def run_single_case(m, n, k, topk, num_experts, block_size): a1_scale=a1_scale, block_shape=block_size, ) + moe_config = make_dummy_moe_config() deep_gemm_experts = mk.FusedMoEKernel( - prepare_finalize=MoEPrepareAndFinalizeNoDPEPModular(), + prepare_finalize=maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), fused_experts=TritonOrDeepGemmExperts( - moe_config=make_dummy_moe_config(), + moe_config=moe_config, quant_config=quant_config, ), inplace=False, diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 7adecc31d2db..6be439208e83 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -8,6 +8,9 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk 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 ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, @@ -22,10 +25,6 @@ FlashInferExperts, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, - MoEPrepareAndFinalizeNoDPEPMonolithic, -) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, swap_w13_to_w31, @@ -241,7 +240,12 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPMonolithic(), + maybe_make_prepare_finalize( + moe=td.layer.moe, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=True, + ), TrtLlmFp8Experts( moe_config=td.layer.moe, quant_config=quant_config, @@ -348,7 +352,12 @@ def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig: ) kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), FlashInferExperts( moe_config=moe_config, quant_config=quant_config, diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index df1273047563..8b08eb4b35e8 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -14,6 +14,9 @@ from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, @@ -24,9 +27,6 @@ is_valid_flashinfer_cutlass_fused_moe, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, -) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe from vllm.utils.torch_utils import set_random_seed @@ -108,7 +108,12 @@ def test_flashinfer_fp4_moe_no_graph( ) flashinfer_experts = FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), FlashInferExperts(moe_config=moe_config, quant_config=quant_config), inplace=False, ) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index a9513fc4ee14..2991feea2ec2 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -24,15 +24,15 @@ from triton_kernels.testing import assert_close from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( OAITritonExperts, UnfusedOAITritonExperts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, -) from vllm.model_executor.layers.utils import shuffle_weight from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -175,14 +175,20 @@ def oai_triton_moe_impl( w1_scale=w1_scale, w2_scale=w2_scale, ) + moe_config = make_dummy_moe_config() if unfused: - fused_experts = UnfusedOAITritonExperts(make_dummy_moe_config(), quant_config) + fused_experts = UnfusedOAITritonExperts(moe_config, quant_config) else: - fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config) + fused_experts = OAITritonExperts(moe_config, quant_config) mk = FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), fused_experts, inplace=False, ) diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index 6804827bfc32..6ae50412f44f 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -14,13 +14,13 @@ from vllm import _custom_ops as ops from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config from vllm.model_executor.layers.fused_moe.cutlass_moe import ( CutlassExpertsFp4, ) -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, -) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -88,11 +88,17 @@ def test_cutlass_fp4_moe_no_graph( w1_scale=w1_blockscale, w2_scale=w2_blockscale, ) + moe_config = make_dummy_moe_config() kernel = mk.FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), CutlassExpertsFp4( - moe_config=make_dummy_moe_config(), + moe_config=moe_config, quant_config=quant_config, ), inplace=False, diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index a74575d16a1c..7d089964d584 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -8,6 +8,9 @@ from tests.kernels.quantization.nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, @@ -24,9 +27,6 @@ fused_experts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel -from vllm.model_executor.layers.fused_moe.prepare_finalize import ( - MoEPrepareAndFinalizeNoDPEPModular, -) from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.deep_gemm import per_block_cast_to_fp8 @@ -573,7 +573,12 @@ def modular_triton_fused_moe( shared_experts: torch.nn.Module | None = None, ) -> FusedMoEKernel: return FusedMoEKernel( - MoEPrepareAndFinalizeNoDPEPModular(), + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), TritonExperts(moe_config, quant_config), shared_experts, inplace=False, From 37974406ac37fe0ef1701b550922de4359c445e3 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:53:33 -0500 Subject: [PATCH 139/207] making tests pass Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 1 + vllm/model_executor/layers/fused_moe/modular_kernel.py | 7 ------- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 6be439208e83..cdc475cfcede 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -167,6 +167,7 @@ def make_moe_tensors_8bit( hidden_dim=k, intermediate_size_per_partition=n, num_local_experts=e, + num_logical_experts=e, moe_parallel_config=layer.moe_parallel_config, in_dtype=hidden_states.dtype, is_act_and_mul=is_gated, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 4fa3b9fdafbb..071cba8495be 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1047,7 +1047,6 @@ def _allocate_buffers( See `workspace_shapes` for a description of the remainder of arguments. Returns a tuple of (workspace13, workspace2, output) tensors. """ - assert isinstance(self.fused_experts, FusedMoEExpertsModular) assert M_full > 0 and M_chunk > 0 num_chunks, _ = self._chunk_info(M_full) @@ -1203,8 +1202,6 @@ def _prepare( The _prepare method is a wrapper around self.prepare_finalize.prepare that handles DBO and async. """ - assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeModular) - if not self.prepare_finalize.supports_async(): # We shouldn't be running an a2a kernel that doesn't # support async prepare/finalize @@ -1290,8 +1287,6 @@ def _fused_experts( apply_router_weight_on_input: bool, expert_tokens_meta: ExpertTokensMetadata | None, ) -> torch.Tensor: - assert isinstance(self.fused_experts, FusedMoEExpertsModular) - _, M_full, N, K, top_k = self.fused_experts.moe_problem_size( a1q, w1, w2, topk_ids ) @@ -1391,8 +1386,6 @@ def _finalize( shared_experts_input is the original hidden_states (full dimension) needed by the shared expert MLP. """ - assert isinstance(self.fused_experts, FusedMoEExpertsModular) - assert isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeModular) shared_output: torch.Tensor | None = None # For latent MoE: shared experts need the original hidden_states From 8b92a0cbe2969d585d309b6cf93aba2df25c1018 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 18:55:02 -0500 Subject: [PATCH 140/207] making tests pass Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index cdc475cfcede..710ea4a89400 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -253,12 +253,12 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( ), ) - flashinfer_output = kernel( + flashinfer_output = kernel.apply_monolithic( hidden_states=td.hidden_states, w1=td.layer.w13_weight, w2=td.layer.w2_weight, router_logits=score, - activation="silu", + activation=activation, global_num_experts=e, expert_map=None, apply_router_weight_on_input=True, From 2ed5dd3600c7a9c4ff5e63cdc25cf5980a1cd47f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 19:01:29 -0500 Subject: [PATCH 141/207] plumb the activation type through Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_fp8_moe.py | 11 +++++++++-- .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 7 +++++++ 2 files changed, 16 insertions(+), 2 deletions(-) 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 3b614120561d..a4aff1c47455 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,9 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + activation_to_flashinfer_int, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, @@ -218,7 +221,7 @@ def _apply_per_tensor( w1: torch.Tensor, w2: torch.Tensor, router_logits: torch.Tensor, - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, @@ -229,6 +232,9 @@ def _apply_per_tensor( routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + activation_type = activation_to_flashinfer_int(activation) + assert self.routing_method_type == RoutingMethodType.Llama4 assert apply_router_weight_on_input @@ -256,6 +262,7 @@ def _apply_per_tensor( routed_scaling_factor=routed_scaling_factor, use_routing_scales_on_input=apply_router_weight_on_input, routing_method_type=self.routing_method_type, + activation_type=activation_type, ) return out @@ -265,7 +272,7 @@ def apply( w1: torch.Tensor, w2: torch.Tensor, router_logits: torch.Tensor, - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, 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 be358a125ac5..a19ae8c903da 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,9 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + activation_to_flashinfer_int, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kNvfp4Dynamic, @@ -153,6 +156,9 @@ def apply( torch.bfloat16 ).view(torch.int16) + # Determine activation type + activation_type = activation_to_flashinfer_int(activation) + # Invoke kernel. flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( topk_ids=packed_tensor, @@ -181,6 +187,7 @@ def apply( routed_scaling_factor=None, routing_method_type=1, do_finalize=True, + activation_type=activation_type, output=output, ) From d4e77463a89190e918e7ba8b8d38b138f6d81f90 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 19:09:33 -0500 Subject: [PATCH 142/207] stash Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_fp8_moe.py | 4 ++-- .../layers/quantization/utils/flashinfer_utils.py | 10 +++++++++- 2 files changed, 11 insertions(+), 3 deletions(-) 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 a4aff1c47455..ff774bfcf4c4 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 @@ -13,7 +13,7 @@ RoutingMethodType, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - activation_to_flashinfer_int, + activation_to_flashinfer_type, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -233,7 +233,7 @@ def _apply_per_tensor( topk_group: int | None = None, ) -> torch.Tensor: assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] - activation_type = activation_to_flashinfer_int(activation) + activation_type = activation_to_flashinfer_type(activation) assert self.routing_method_type == RoutingMethodType.Llama4 assert apply_router_weight_on_input diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 0cf47302cdef..29f815d400a4 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum +from typing import TYPE_CHECKING import torch @@ -10,6 +11,9 @@ from vllm.platforms import current_platform from vllm.utils.math_utils import round_up +if TYPE_CHECKING: + from flashinfer.fused_moe.core import ActivationType + logger = init_logger(__name__) @@ -20,6 +24,10 @@ class FlashinferMoeBackend(Enum): def activation_to_flashinfer_int(activation: MoEActivation) -> int: + return activation_to_flashinfer_type(activation).value + + +def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType": from flashinfer.fused_moe.core import ActivationType # silu and gelu are mapped to their gated versions SwiGLU and GeGLU respectively @@ -30,7 +38,7 @@ def activation_to_flashinfer_int(activation: MoEActivation) -> int: MoEActivation.GELU: ActivationType.Geglu, MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, } - return ACTIVATION_TO_FI_ACTIVATION[activation].value + return ACTIVATION_TO_FI_ACTIVATION[activation] def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor: From 44f90d2394c9122bd18a46abe418a52e8bfc6c48 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 19:25:00 -0500 Subject: [PATCH 143/207] fix .apply call Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 2 +- .../layers/fused_moe/experts/trtllm_fp8_moe.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 710ea4a89400..f8538223c3b2 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -366,7 +366,7 @@ def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig: inplace=False, ) - flashinfer_cutlass_output = kernel( + flashinfer_cutlass_output = kernel.apply( td.hidden_states, td.layer.w13_weight, td.layer.w2_weight, 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 ff774bfcf4c4..cbf8c3cb6fba 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 @@ -64,7 +64,11 @@ def __init__( self._g1_alphas = (w1_scale * a1_scale).squeeze() self._g2_alphas = (w2_scale * a2_scale).squeeze() - self._g1_scale_c = self._g1_alphas / self.quant_config.a2_scale + self._g1_scale_c = ( + self._g1_alphas / self.quant_config.a2_scale + if moe_config.is_act_and_mul + else torch.ones_like(self._g1_alphas) * self.quant_config.a2_scale + ) @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: From a72ba70f9d9dd490399f79926a02f2664cf9f37b Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 19:29:44 -0500 Subject: [PATCH 144/207] clean up Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index f8538223c3b2..599eb2598e2b 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -480,25 +480,26 @@ def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch): routing_logits = torch.randn((m, e), device="cuda", dtype=torch.float32) routing_bias = torch.randn(e, device="cuda", dtype=torch.float32) - # This should NOT crash with num_expert_group=None - output = torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8( + # This should NOT crash with n_group=None + import flashinfer + + output = flashinfer.fused_moe.trtllm_fp8_block_scale_moe( routing_logits=routing_logits, routing_bias=routing_bias, - x=x, - w13_weight=w13_fp8, - w13_weight_scale_inv=w13_scale, - w2_weight=w2_fp8, - w2_weight_scale_inv=w2_scale, - global_num_experts=e, + hidden_states=x, + gemm1_weights=w13_fp8, + gemm1_weights_scale=w13_scale, + gemm2_weights=w2_fp8, + gemm2_weights_scale=w2_scale, + num_experts=e, top_k=topk, - num_expert_group=None, + n_group=None, topk_group=None, intermediate_size=n, expert_offset=0, local_num_experts=e, - block_shape=block_shape, + routed_scaling_factor=1.0, routing_method_type=RoutingMethodType.DeepSeekV3, - routed_scaling=1.0, ) assert output is not None From 08259775a943e31cfe878a2c524c76eed2490cee Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 19:34:23 -0500 Subject: [PATCH 145/207] clean up Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 599eb2598e2b..7071f05bf7bf 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -439,7 +439,6 @@ def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch): if not current_platform.has_device_capability(100): pytest.skip("Test requires SM >= 100 (Blackwell)") - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401 from tests.kernels.quant_utils import native_per_token_group_quant_fp8 set_random_seed(7) @@ -483,10 +482,20 @@ def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch): # This should NOT crash with n_group=None import flashinfer + from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input + + a1q, a1q_scale = moe_kernel_quantize_input( + A=x, + A_scale=None, + quant_dtype=torch.float8_e4m3fn, + per_act_token_quant=False, + block_shape=block_shape, + ) output = flashinfer.fused_moe.trtllm_fp8_block_scale_moe( routing_logits=routing_logits, routing_bias=routing_bias, - hidden_states=x, + hidden_states=a1q, + hidden_states_scale=a1q_scale, gemm1_weights=w13_fp8, gemm1_weights_scale=w13_scale, gemm2_weights=w2_fp8, @@ -496,7 +505,7 @@ def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch): n_group=None, topk_group=None, intermediate_size=n, - expert_offset=0, + local_expert_offset=0, local_num_experts=e, routed_scaling_factor=1.0, routing_method_type=RoutingMethodType.DeepSeekV3, From cb21b90d8a832c6cfff740dcf466f43b7ced72ad Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 19:51:30 -0500 Subject: [PATCH 146/207] stash Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 73 ++++++++++++------- .../fused_moe/experts/trtllm_fp8_moe.py | 9 +-- 2 files changed, 51 insertions(+), 31 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 7071f05bf7bf..ab3a5d399491 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -479,36 +479,57 @@ def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch): routing_logits = torch.randn((m, e), device="cuda", dtype=torch.float32) routing_bias = torch.randn(e, device="cuda", dtype=torch.float32) - # This should NOT crash with n_group=None - import flashinfer - - from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input + activation = MoEActivation.SILU + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size_per_partition=n, + num_local_experts=e, + num_logical_experts=e, + activation=activation, + device="cuda", + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=torch.bfloat16, + is_act_and_mul=activation.is_gated, + routing_method=RoutingMethodType.DeepSeekV3, + ) - a1q, a1q_scale = moe_kernel_quantize_input( - A=x, - A_scale=None, - quant_dtype=torch.float8_e4m3fn, - per_act_token_quant=False, + quant_config = fp8_w8a8_moe_quant_config( + w1_scale=w13_scale, + w2_scale=w2_scale, block_shape=block_shape, ) - output = flashinfer.fused_moe.trtllm_fp8_block_scale_moe( - routing_logits=routing_logits, - routing_bias=routing_bias, - hidden_states=a1q, - hidden_states_scale=a1q_scale, - gemm1_weights=w13_fp8, - gemm1_weights_scale=w13_scale, - gemm2_weights=w2_fp8, - gemm2_weights_scale=w2_scale, - num_experts=e, - top_k=topk, - n_group=None, - topk_group=None, - intermediate_size=n, - local_expert_offset=0, - local_num_experts=e, + + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=True, + ), + TrtLlmFp8Experts( + moe_config=moe_config, + quant_config=quant_config, + ), + inplace=False, + ) + + # This should NOT crash with n_group=None + output = kernel.apply_monolithic( + hidden_states=x, + w1=w13_fp8, + w2=w2_fp8, + router_logits=routing_logits, + activation=activation, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + # grouped topk + fused topk bias parameters + num_expert_group=None, + e_score_correction_bias=routing_bias, routed_scaling_factor=1.0, - routing_method_type=RoutingMethodType.DeepSeekV3, + topk_group=None, ) assert output is not None 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 cbf8c3cb6fba..a553f5c18f43 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 @@ -179,10 +179,9 @@ def _apply_per_block( ) -> torch.Tensor: assert not apply_router_weight_on_input assert activation == MoEActivation.SILU - assert ( - e_score_correction_bias is None - or e_score_correction_bias.dtype == hidden_states.dtype - ) + + if e_score_correction_bias is not None: + e_score_correction_bias.to(hidden_states.dtype) if self.routing_method_type == RoutingMethodType.DeepSeekV3: router_logits = router_logits.to(torch.float32) @@ -210,7 +209,7 @@ def _apply_per_block( gemm2_weights_scale=self.quant_config.w2_scale, num_experts=global_num_experts, top_k=self.topk, - n_group=num_expert_group, + n_group=(num_expert_group or 0), topk_group=(topk_group or 0), intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, From 72c173fd21518e21aa59db1c5ec578716ef1c432 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 19:53:50 -0500 Subject: [PATCH 147/207] updated Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index ab3a5d399491..565a0682be6d 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -428,6 +428,9 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout( assert w2_converted.shape[0] == num_experts +@pytest.mark.skip( + reason="This test is failing on main. See: https://github.com/vllm-project/vllm/pull/34494#issuecomment-3911242744" +) # noqa: E501 def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch): """Test that flashinfer_fused_moe_blockscale_fp8 handles num_expert_group=None. From b691c8510bdbd348ba5dc97a24c7e426758e7046 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 19:56:04 -0500 Subject: [PATCH 148/207] clean up test flashinfer moe Signed-off-by: Robert Shaw --- tests/kernels/moe/test_flashinfer_moe.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index 8b08eb4b35e8..a3fb474f1517 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -118,13 +118,16 @@ def test_flashinfer_fp4_moe_no_graph( inplace=False, ) - flashinfer_output = flashinfer_experts( + flashinfer_output = flashinfer_experts.apply( hidden_states=a, w1=w1_q, w2=w2_q, topk_weights=topk_weights, topk_ids=topk_ids, activation=activation, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, ) # Reference check: From e8a0e71deadcd931bb8222c7ede8b5e23868bcd0 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:01:44 -0500 Subject: [PATCH 149/207] fix test batched deepgemm Signed-off-by: Robert Shaw --- tests/kernels/moe/test_batched_deepgemm.py | 11 +++++++++-- tests/kernels/moe/test_deepep_moe.py | 2 +- .../model_executor/layers/fused_moe/modular_kernel.py | 6 +++--- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py index 8aaebcee1e6e..20763b91dfd9 100644 --- a/tests/kernels/moe/test_batched_deepgemm.py +++ b/tests/kernels/moe/test_batched_deepgemm.py @@ -4,6 +4,7 @@ import pytest import torch +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import ( BatchedDeepGemmExperts, ) @@ -80,13 +81,16 @@ def test_batched_deepgemm_vs_triton( inplace=False, ) - out_triton = mk_triton( + out_triton = mk_triton.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, ) # deepgemm @@ -102,13 +106,16 @@ def test_batched_deepgemm_vs_triton( inplace=False, ) - out_deepgemm = mk_deepgemm( + out_deepgemm = mk_deepgemm.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, ) diff = calc_diff(out_deepgemm, out_triton) diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index e648291e351d..362b71a40f2d 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -255,7 +255,7 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False): quant_config, ) - out = mk.forward( + out = mk.apply( hidden_states=rank_tokens_chunk, w1=w1, w2=w2, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 071cba8495be..1a0b7042c286 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1550,7 +1550,7 @@ def apply( w1: torch.Tensor, w2: torch.Tensor, router_logits: torch.Tensor, - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, apply_router_weight_on_input: bool, @@ -1675,7 +1675,7 @@ def apply_monolithic( w1: torch.Tensor, w2: torch.Tensor, router_logits: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, apply_router_weight_on_input: bool, @@ -1708,7 +1708,7 @@ def apply( w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - activation: str, + activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, apply_router_weight_on_input: bool, From b6b0d55108436709781771d2a6d98f31810fceb3 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:07:40 -0500 Subject: [PATCH 150/207] fix test batched_moe Signed-off-by: Robert Shaw --- tests/kernels/moe/test_block_fp8.py | 8 +++++++- tests/kernels/moe/utils.py | 15 +++++++++++++-- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 1c70027d48b8..bfa1c1cb330a 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -21,6 +21,9 @@ fused_experts, fused_topk, ) +from vllm.model_executor.layers.fused_moe.activation import ( + MoEActivation, +) from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, ) @@ -269,12 +272,15 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) ) def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids): - return deep_gemm_experts( + return deep_gemm_experts.apply( hidden_states=a, w1=w1, w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + expert_map=False, ) # Set the context to avoid lots of warning spam. diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 7d089964d584..69f08f363b62 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -156,6 +156,7 @@ def naive_batched_moe( a1_scale=a1_scale, a2_scale=a2_scale, ) + moe_config = make_dummy_moe_config() fused_experts = FusedMoEKernel( BatchedPrepareAndFinalize( @@ -165,12 +166,22 @@ def naive_batched_moe( max_num_tokens=max_num_tokens, num_dispatchers=1, quant_config=quant_config, - moe_config=make_dummy_moe_config(), + moe_config=moe_config, ), inplace=False, ) - return fused_experts(a, w1, w2, topk_weight, topk_ids) + return fused_experts.apply( + a, + w1, + w2, + topk_weight, + topk_ids, + global_num_experts=moe_config.num_experts, + activation=moe_config.activation, + apply_router_weight_on_input=False, + expert_map=None, + ) def chunk_scales( From af5112593be501164d1e2bdce30c92173a85f802 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:08:50 -0500 Subject: [PATCH 151/207] fix test batched_moe Signed-off-by: Robert Shaw --- tests/kernels/moe/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 69f08f363b62..5b36fc0f8d19 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -177,7 +177,7 @@ def naive_batched_moe( w2, topk_weight, topk_ids, - global_num_experts=moe_config.num_experts, + global_num_experts=w1.shape[0], activation=moe_config.activation, apply_router_weight_on_input=False, expert_map=None, From 38c8370064fb8b80784f1e878465577e65dba61b Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:10:05 -0500 Subject: [PATCH 152/207] fix test batched_moe Signed-off-by: Robert Shaw --- tests/kernels/moe/utils.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 5b36fc0f8d19..abcc3430931c 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -115,6 +115,8 @@ def batched_moe( a2_scale=a2_scale, ) + moe_config = make_dummy_moe_config() + fused_experts = FusedMoEKernel( BatchedPrepareAndFinalize( max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 @@ -123,12 +125,22 @@ def batched_moe( max_num_tokens=max_num_tokens, num_dispatchers=1, quant_config=quant_config, - moe_config=make_dummy_moe_config(), + moe_config=moe_config, ), inplace=False, ) - return fused_experts(a, w1, w2, topk_weight, topk_ids) + return fused_experts.apply( + a, + w1, + w2, + topk_weight, + topk_ids, + global_num_experts=w1.shape[0], + activation=moe_config.activation, + apply_router_weight_on_input=False, + expert_map=None, + ) def naive_batched_moe( From 083c3aeb151dc44cc0baf0bf085a615e8e2b6498 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:28:14 -0500 Subject: [PATCH 153/207] fix test cutlass moe Signed-off-by: Robert Shaw --- tests/kernels/moe/test_block_fp8.py | 17 +++++++++++++---- tests/kernels/moe/test_cutlass_moe.py | 6 +++++- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index bfa1c1cb330a..a74e739c55e4 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -21,9 +21,7 @@ fused_experts, fused_topk, ) -from vllm.model_executor.layers.fused_moe.activation import ( - MoEActivation, -) +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, ) @@ -196,7 +194,17 @@ def test_w8a8_block_fp8_fused_moe( a, w1, w2, topk_weights, topk_ids, quant_config=quant_config ) - m_out = m_fused_moe(a, w1, w2, topk_weights, topk_ids) + m_out = m_fused_moe.apply( + a, + w1, + w2, + topk_weights, + topk_ids, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + expert_map=None, + global_num_experts=w1.shape[0], + ) # 0.039 only needed for M >= 8192 tol = 0.035 if M < 8192 else 0.039 @@ -278,6 +286,7 @@ def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids): w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, + global_num_experts=E, activation=MoEActivation.SILU, apply_router_weight_on_input=False, expert_map=False, diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 81c6a74fc7c0..958cfa8ba59d 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -258,6 +258,10 @@ def run_8_bit( "w2": moe_tensors.w2_q, # type: ignore[union-attr] "topk_weights": topk_weights, "topk_ids": topk_ids, + "global_num_experts": moe_tensors.w1_q.shape[0], # type: ignore[union-attr] + "activation": MoEActivation.SILU, + "expert_map": None, + "apply_router_weight_on_input": False, } num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined] @@ -282,7 +286,7 @@ def run_8_bit( ), inplace=False, ) - return kernel(**kwargs) + return kernel.apply(**kwargs) assert num_local_experts is not None return run_with_expert_maps( From a8c9e3b9b5109ea3570a88140ddca8b4d2477ea5 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:45:39 -0500 Subject: [PATCH 154/207] making progress on fixing kernel tests Signed-off-by: Robert Shaw --- benchmarks/kernels/benchmark_moe.py | 13 +++++++++++-- tests/kernels/moe/test_deepgemm.py | 9 ++++++++- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index c1a3e8d9cef6..4abeaefd774a 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -271,6 +271,7 @@ def run(): moe_config=moe_config, quant_config=quant_config, ), + inplace=not disable_inplace(), ) with override_config(config): @@ -280,8 +281,16 @@ def run(): inplace = not disable_inplace() if use_deep_gemm: - return deep_gemm_experts( - x, w1, w2, topk_weights, topk_ids, inplace=inplace + return deep_gemm_experts.apply( + x, + w1, + w2, + topk_weights, + topk_ids, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + apply_router_weight_on_input=False, + expert_map=False, ) return fused_experts( x, diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 7851d9ea0abe..c718414ed208 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -14,6 +14,9 @@ # vLLM fused-expert reference (Triton fallback + DeepGEMM option) import vllm.model_executor.layers.fused_moe.modular_kernel as mk from tests.kernels.moe.utils import make_dummy_moe_config +from vllm.model_executor.layers.fused_moe.activation import ( + MoEActivation, +) from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, ) @@ -136,12 +139,16 @@ def run_single_case(m, n, k, topk, num_experts, block_size): ) # DeepGemm - out_deepgemm = deep_gemm_experts( + out_deepgemm = deep_gemm_experts.apply( hidden_states=tokens_bf16, w1=w1, w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, + global_num_experts=num_experts, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + expert_map=False, ) diff = calc_diff(out_deepgemm, out_triton) assert diff < 0.001, f"Diff exceeded 1%: {diff}" From dd1dfc01a89d6fbc688d0006276473a07d1e658e Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:50:42 -0500 Subject: [PATCH 155/207] make deepgemm pass Signed-off-by: Robert Shaw --- tests/kernels/moe/test_deepgemm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index c718414ed208..c2949391c798 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -148,7 +148,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size): global_num_experts=num_experts, activation=MoEActivation.SILU, apply_router_weight_on_input=False, - expert_map=False, + expert_map=None, ) diff = calc_diff(out_deepgemm, out_triton) assert diff < 0.001, f"Diff exceeded 1%: {diff}" From d4919e33e2c304cbce4fddf7ac8d70f09e0003a4 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:53:43 -0500 Subject: [PATCH 156/207] fix test marlin vs trtllm Signed-off-by: Robert Shaw --- .../moe/test_marlin_vs_trtllm_mxint4.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py index d6735b126e2f..8687356eb3fc 100644 --- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py +++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py @@ -221,16 +221,16 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group ) marlin_output = fused_marlin_moe( - a, - w1_marlin, - w2_marlin, - None, - None, - w1_scales_marlin, - w2_scales_marlin, - None, # gating_output not needed when topk_weights/ids provided - topk_weights, - topk_ids, + hidden_states=a, + w1=w1_marlin, + w2=w2_marlin, + bias1=None, + bias2=None, + w1_scales=w1_scales_marlin, + w2_scales=w2_scales_marlin, + topk_weights=topk_weights, + topk_ids=topk_ids, + quant_type_id=scalar_types.uint4b8.id, global_num_experts=e, expert_map=None, global_scale1=None, @@ -244,7 +244,6 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group w1_zeros=None, w2_zeros=None, input_dtype=dtype, - quant_type_id=scalar_types.uint4b8.id, is_k_full=True, ) From 026fdd220082fa16470948f53710104bcc4e6c8d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 21:57:27 -0500 Subject: [PATCH 157/207] fix up more tests Signed-off-by: Robert Shaw --- tests/kernels/moe/modular_kernel_tools/common.py | 2 +- .../moe/modular_kernel_tools/profile_modular_kernel.py | 2 +- tests/kernels/moe/test_deepep_deepgemm_moe.py | 2 +- tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py | 4 ++-- tests/kernels/moe/test_modular_oai_triton_moe.py | 2 +- tests/kernels/moe/test_moe.py | 8 ++++++-- tests/kernels/moe/test_nvfp4_moe.py | 6 +++++- 7 files changed, 17 insertions(+), 9 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 0c0699edb17a..1ec13aec2ff9 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -691,6 +691,6 @@ def run_modular_kernel( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - out = mk.forward(**mk_kwargs) + out = mk.apply(**mk_kwargs) return out diff --git a/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py b/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py index 3cdc7b82130b..e64cdb528cd3 100644 --- a/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py +++ b/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py @@ -72,7 +72,7 @@ def profile_modular_kernel( "apply_router_weight_on_input": config.topk == 1, } - do_profile(mk.forward, mk_kwargs, pgi, config) + do_profile(mk.apply, mk_kwargs, pgi, config) def rank_worker( diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index e208fd0eed7f..a01fb1a452ea 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -319,7 +319,7 @@ def build_expert_map(): with with_dp_metadata( M=test_tensors.rank_tokens.size(0), world_size=pgi.world_size ): - out = mk.forward( + out = mk.apply( hidden_states=test_tensors.rank_tokens, w1=w1, w2=w2, diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py index 8687356eb3fc..aaf255ca8b6a 100644 --- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py +++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py @@ -226,8 +226,8 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group w2=w2_marlin, bias1=None, bias2=None, - w1_scales=w1_scales_marlin, - w2_scales=w2_scales_marlin, + w1_scale=w1_scales_marlin, + w2_scale=w2_scales_marlin, topk_weights=topk_weights, topk_ids=topk_ids, quant_type_id=scalar_types.uint4b8.id, diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 2991feea2ec2..33a875a8c9ac 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -193,7 +193,7 @@ def oai_triton_moe_impl( inplace=False, ) - return mk.forward( + return mk.apply( hidden_states=x, w1=w1, w2=w2, diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index eb3d9f8a8f6b..cda0b5c11040 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -346,14 +346,16 @@ def m_fused_moe( expert_map: torch.Tensor | None = None, ) -> torch.Tensor: topk_weights, topk_ids, _ = fused_topk(a, score, topk, False) - return m_fused_moe_fn( + return m_fused_moe_fn.apply( a, w1, w2, topk_weights, topk_ids, + activation=MoEActivation.SILU, global_num_experts=global_num_experts, expert_map=expert_map, + apply_router_weight_on_input=False, ) fused_moe_fn = functools.partial(fused_moe, renormalize=False) @@ -500,14 +502,16 @@ def m_fused_moe( expert_map: torch.Tensor | None = None, ) -> torch.Tensor: topk_weights, topk_ids, _ = fused_topk(a, score, topk, False) - return m_fused_moe_fn( + return m_fused_moe_fn.apply( a, w1, w2, topk_weights, topk_ids, + activation=MoEActivation.SILU, global_num_experts=global_num_experts, expert_map=expert_map, + apply_router_weight_on_input=False, ) fused_moe_fn = functools.partial(fused_moe, renormalize=False) diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index 6ae50412f44f..28227fca2e35 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -104,12 +104,16 @@ def test_cutlass_fp4_moe_no_graph( inplace=False, ) - cutlass_output = kernel( + cutlass_output = kernel.apply( hidden_states=a, w1=w1_q, w2=w2_q, topk_weights=topk_weights, topk_ids=topk_ids, + global_num_experts=e, + activation=mk.MoEActivation.SILU, + apply_router_weight_on_input=False, + expert_map=None, ) # Reference check: From e642af846f83c4e18e4f1e64018841e6e560809d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 22:03:45 -0500 Subject: [PATCH 158/207] fix mixtral Signed-off-by: Robert Shaw --- .../layers/fused_moe/unquantized_fused_moe_method.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 532a5375a00b..95b6f7b77fa0 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -325,7 +325,7 @@ def forward_cuda( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: assert self.kernel is not None - return self.kernel( + return self.kernel.apply( hidden_states=x, w1=layer.w13_weight, w2=layer.w2_weight, From 58ac35f5324021672d19fa842e482b8cc5329360 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 22:10:39 -0500 Subject: [PATCH 159/207] trying to fix pplx texts Signed-off-by: Robert Shaw --- tests/kernels/moe/test_pplx_moe.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_pplx_moe.py b/tests/kernels/moe/test_pplx_moe.py index b473e08c8737..3cbe56d0cfd2 100644 --- a/tests/kernels/moe/test_pplx_moe.py +++ b/tests/kernels/moe/test_pplx_moe.py @@ -38,6 +38,7 @@ from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk, override_config +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config @@ -608,13 +609,15 @@ def pplx_moe( else: _fused_experts = fused_experts - out = _fused_experts( + out = _fused_experts.apply( a_chunk, w1_chunk, w2_chunk, chunk_topk_weight, chunk_topk_ids, global_num_experts=num_experts, + activation=MoEActivation.SILU, + expert_map=None, ) if use_cudagraphs: From 4211f57bbabffaae1ee808d809c7ad9cc6695185 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 22:12:12 -0500 Subject: [PATCH 160/207] fix as many tests as possible Signed-off-by: Robert Shaw --- tests/kernels/moe/test_pplx_cutlass_moe.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_pplx_cutlass_moe.py b/tests/kernels/moe/test_pplx_cutlass_moe.py index bac1be6db3b5..d55f43188dbb 100644 --- a/tests/kernels/moe/test_pplx_cutlass_moe.py +++ b/tests/kernels/moe/test_pplx_cutlass_moe.py @@ -183,14 +183,16 @@ def make_moe_config() -> FusedMoEConfig: chunk_by_rank(topk_ids, rank, world_size).to(torch.uint32).to(device) ) - out = fused_cutlass_experts( + out = fused_cutlass_experts.apply( a_chunk, chunk_by_rank(w1, rank, world_size), chunk_by_rank(w2, rank, world_size), chunk_topk_weight, chunk_topk_ids, + activation=MoEActivation.SILU, global_num_experts=num_experts, expert_map=None, # TODO + apply_router_weight_on_input=False, ) torch.cuda.synchronize() From 812cdfa510698e627db46d1959ddca27f96c3a94 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 22:27:08 -0500 Subject: [PATCH 161/207] fix pre-commit Signed-off-by: Robert Shaw --- vllm/lora/layers/fused_moe.py | 9 +++++---- vllm/model_executor/layers/fused_moe/modular_kernel.py | 2 +- .../layers/fused_moe/pplx_prepare_finalize.py | 2 +- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 0916988c0b74..6b6ed4d6e3ba 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -146,10 +146,11 @@ def _inject_lora_into_fused_moe(self): if quant_config.use_mxfp4_w4a16: assert isinstance( - m_fused_moe_fn.fused_experts, (MarlinExperts, UnfusedOAITritonExperts) + m_fused_moe_fn.impl.fused_experts, + (MarlinExperts, UnfusedOAITritonExperts), ) else: - assert isinstance(m_fused_moe_fn.fused_experts, TritonExperts) + assert isinstance(m_fused_moe_fn.impl.fused_experts, TritonExperts) def fwd_decorator(layer, func): def wrapper(*args, **kwargs): @@ -329,9 +330,9 @@ def wrapper(*args, **kwargs): return wrapper - fused_experts = m_fused_moe_fn.fused_experts + fused_experts = m_fused_moe_fn.impl.fused_experts - m_fused_moe_fn.forward = fwd_decorator(self.base_layer, m_fused_moe_fn.forward) + m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply) fused_experts.activation = act_decorator( self.base_layer, fused_experts.activation ) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 1a0b7042c286..a74724ec40a1 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1636,7 +1636,7 @@ def __init__( raise ValueError( "prepare_finalize and fused_experts must both be either monolithic " f"or non-monolithic but got {prepare_finalize.__class__.__name__} " - "and {fused_experts.__class__.__name__}" + f"and {fused_experts.__class__.__name__}" ) self._post_init_setup() diff --git a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py index 289ac0d1413d..c87d948baf92 100644 --- a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py @@ -62,7 +62,7 @@ def pplx_hidden_dim_scale_bytes( ) -class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): +class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): """PPLX-based prepare and finalize for expert parallelism.""" def __init__( From 049e2c2e5cd747fb53a4b16d7e888004f08e039a Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 22:31:58 -0500 Subject: [PATCH 162/207] fix xpuexperts typing Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/xpu_fused_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py index 95a69662e38a..0693a25468fd 100644 --- a/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py @@ -23,7 +23,7 @@ from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe -class XPUExperts(mk.FusedMoEKernel): +class XPUExperts(mk.FusedMoEExpertsModular): def __init__( self, moe_config: FusedMoEConfig, From ac8287520b09cd8b883ade1283813bf1dc10b6af Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 16 Feb 2026 22:33:03 -0500 Subject: [PATCH 163/207] fix workspacde shapes typing Signed-off-by: Robert Shaw --- .../model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 a19ae8c903da..3fe909995256 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 @@ -112,7 +112,7 @@ def workspace_shapes( global_num_experts: int, local_num_experts: int, expert_tokens_meta: mk.ExpertTokensMetadata | None, - activation: str, + activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: # The workspaces for this implementation are managed by flashinfer. workspace1 = (0,) From 827e837095381183bb640c99ed83fd89d0408abf Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 10:25:05 -0500 Subject: [PATCH 164/207] fix blockfp8 Signed-off-by: Robert Shaw --- .../model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py | 1 + vllm/model_executor/layers/quantization/fp8.py | 4 ++++ 2 files changed, 5 insertions(+) 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 e50f23070879..2ddf03bf85a2 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 @@ -303,6 +303,7 @@ def apply( num_expert_group=num_expert_group, e_score_correction_bias=e_score_correction_bias, routed_scaling_factor=routed_scaling_factor, + topk_group=topk_group, ) elif self.quant_config.is_per_tensor: return self._apply_per_tensor( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 18dc061a0457..158bd786f9c5 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -945,6 +945,10 @@ def apply_monolithic( global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) def apply( From bf9f9bb93947320e586b8f97238171d43af42e45 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 10:41:06 -0500 Subject: [PATCH 165/207] updated to hopefully fix modular kernel combinations test Signed-off-by: Robert Shaw --- tests/kernels/moe/modular_kernel_tools/common.py | 3 ++- tests/kernels/moe/modular_kernel_tools/mk_objects.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 1ec13aec2ff9..6fbb0503af08 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -666,7 +666,7 @@ def run_modular_kernel( # impls might update the tensor in place hidden_states = rank_tensors.hidden_states.clone() - topk_ids = rank_tensors.topk_ids.to(mk.prepare_finalize.topk_indices_dtype()) + topk_ids = rank_tensors.topk_ids.to(mk.impl.prepare_finalize.topk_indices_dtype()) mk_kwargs = { "hidden_states": hidden_states, @@ -674,6 +674,7 @@ def run_modular_kernel( "w2": rank_weights.w2, "topk_weights": rank_tensors.topk_weights, "topk_ids": topk_ids, + "activation": MoEActivation.SILU, "expert_map": rank_tensors.expert_map, "global_num_experts": config.E, "apply_router_weight_on_input": config.topk == 1 diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index b403eade02c1..86bf254a9cd2 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -255,14 +255,14 @@ def expert_info(kind) -> ExpertInfo: if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100): from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501 - FlashInferCutlassMoEPrepareAndFinalize, + FlashInferA2APrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( FlashInferExperts, ) register_prepare_and_finalize( - FlashInferCutlassMoEPrepareAndFinalize, + FlashInferA2APrepareAndFinalize, standard_format, nvfp4_types + fp8_types, blocked_quantization_support=True, From c71ea163557bac07238b86f523c0bf722d577dad Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 10:58:56 -0500 Subject: [PATCH 166/207] hopefully fix mk Signed-off-by: Robert Shaw --- tests/kernels/moe/test_modular_kernel_combinations.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 7bdbaf0c8a52..45e3e03759c2 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -168,7 +168,6 @@ def run(config: Config, verbose: bool): def is_nyi_config(config: Config) -> bool: # We know these configs to be legitimate. but still fail. info = expert_info(config.fused_experts_type) - if info.needs_matching_quant: # The triton kernels expect both per-act-token-quant and # per-out-ch-quant or neither. From ae069f87633dc9e5d77e85eedc4294211a6fb80d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 11:31:30 -0500 Subject: [PATCH 167/207] add nemotron to blackwell moe test Signed-off-by: Robert Shaw --- tests/quantization/test_blackwell_moe.py | 60 ++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index 07da2b454e6f..12e997c2c7d1 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -124,6 +124,12 @@ def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch): can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) +def test_deepseek_fp8_block_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "0") + monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "0") + can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) + + @pytest.mark.skip( reason=( "Known issue: lack of kernel support. " @@ -142,6 +148,11 @@ def test_deepseek_fp8_block_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatc can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) +def test_deepseek_nvfp4_moe_flashinfer_vllm(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "0") + can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2", hf_overrides=HF_OVERRIDE_TEXT) + + def test_deepseek_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1") monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") @@ -186,3 +197,52 @@ def test_gptoss_eager(monkeypatch: pytest.MonkeyPatch): def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1") can_initialize("Qwen/Qwen3-Next-80B-A3B-Instruct", hf_overrides=HF_OVERRIDE_TEXT) + + +## NemoTron ## + + +def test_nemotron_fp8_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") + monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", hf_overrides=HF_OVERRIDE_TEXT + ) + + +def test_nemotron_fp8_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") + monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency") + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", hf_overrides=HF_OVERRIDE_TEXT + ) + + +def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "0") + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", hf_overrides=HF_OVERRIDE_TEXT + ) + + +def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1") + monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", hf_overrides=HF_OVERRIDE_TEXT + ) + + +def test_nemotron_fp4_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1") + monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency") + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", hf_overrides=HF_OVERRIDE_TEXT + ) + + +def test_nemotron_fp4_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "0") + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", hf_overrides=HF_OVERRIDE_TEXT + ) From d4f08dea39a76ae20606eb93ba402f61555bb994 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 11:38:10 -0500 Subject: [PATCH 168/207] add hopper moe tests Signed-off-by: Robert Shaw --- .buildkite/test_areas/quantization.yaml | 17 ++++ tests/quantization/test_hopper_moe.py | 130 ++++++++++++++++++++++++ 2 files changed, 147 insertions(+) create mode 100644 tests/quantization/test_hopper_moe.py diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 5ee2e5186966..d355b9bee23d 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -37,6 +37,23 @@ steps: commands: - pytest -s -v tests/quantization/test_blackwell_moe.py +- label: Quantized MoE Test (H100) + timeout_in_minutes: 60 + working_dir: "/vllm-workspace/" + device: h100 + source_file_dependencies: + - tests/quantization/test_hopper_moe.py + - vllm/model_executor/models/deepseek_v2.py + - vllm/model_executor/models/gpt_oss.py + - vllm/model_executor/models/llama4.py + - vllm/model_executor/layers/fused_moe + - vllm/model_executor/layers/quantization/compressed_tensors + - vllm/model_executor/layers/quantization/modelopt.py + - vllm/model_executor/layers/quantization/mxfp4.py + - vllm/v1/attention/backends/flashinfer.py + commands: + - pytest -s -v tests/quantization/test_hopper_moe.py + - label: Quantized Models Test timeout_in_minutes: 60 source_file_dependencies: diff --git a/tests/quantization/test_hopper_moe.py b/tests/quantization/test_hopper_moe.py new file mode 100644 index 000000000000..cb758dda9266 --- /dev/null +++ b/tests/quantization/test_hopper_moe.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import os +from typing import Any + +import pytest + +from tests.utils import RemoteOpenAIServer +from vllm.platforms import current_platform + +if not current_platform.is_device_capability_family(90): + pytest.skip("This test only runs on Hopper GPUs (SM90x).", allow_module_level=True) + + +@pytest.fixture(scope="module", autouse=True) +def set_test_environment(): + """Sets environment variables required for this test module.""" + # Make sure TRTLLM attention is available + os.environ["VLLM_HAS_FLASHINFER_CUBIN"] = "1" + # Set compilation threads to 16 to speed up startup + os.environ["FLASHINFER_NVCC_THREADS"] = "16" + + +# Overide the backbone layers to 4 for faster startup +HF_OVERRIDE_TEXT = { + "num_layers": 4, + "num_hidden_layers": 4, +} +HF_OVERRIDE_MM = { + "text_config": {"num_layers": 4, "num_hidden_layers": 4}, +} + + +def can_initialize( + model: str, + hf_overrides: dict[str, Any] | None = None, + extra_args: list[str] | None = None, +): + # Server arguments + extra_args = extra_args if extra_args is not None else [] + server_args = [ + "--max-model-len", + "2048", + "--max-num-batched-tokens", + "256", + "--load-format", + "dummy", + "--trust-remote-code", + "--limit-mm-per-prompt", + json.dumps({"image": 0}), + *extra_args, + ] + + # Launch server and make a simple request + with RemoteOpenAIServer( + model, + server_args, + max_wait_seconds=1500, # Due to FlashInfer compile + override_hf_configs=hf_overrides, + ) as server: + client = server.get_client() + # Make a simple request to verify the server works + completion = client.completions.create( + model=model, + prompt=["Hello, World!"], + temperature=0, + max_tokens=2, + ) + print(completion) + assert completion.choices[0].text is not None + + +## Llama4 ## + + +def test_llama4_fp8_tensor_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "0") + can_initialize( + "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM + ) + + +def test_llama4_fp8_tensor_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") + monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") + can_initialize( + "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM + ) + + +## DeepSeekV3 ## + + +def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1") + can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) + + +def test_deepseek_fp8_block_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "0") + monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "0") + can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) + + +## Qwen3 Next ## + + +def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1") + can_initialize("Qwen/Qwen3-Next-80B-A3B-Instruct", hf_overrides=HF_OVERRIDE_TEXT) + + +## NemoTron ## + + +def test_nemotron_fp8_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") + monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", hf_overrides=HF_OVERRIDE_TEXT + ) + + +def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "0") + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", hf_overrides=HF_OVERRIDE_TEXT + ) From ac123204f9edf6b5822faac5d9a8b617dcbac819 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 12:09:42 -0500 Subject: [PATCH 169/207] skip fp8 nemotron Signed-off-by: Robert Shaw --- tests/quantization/test_blackwell_moe.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index 12e997c2c7d1..02a87cbc332f 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -210,6 +210,13 @@ def test_nemotron_fp8_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch) ) +@pytest.mark.skip( + reason=( + "FP8 MoE backend FLASHINFER_TRTLLM does not support the " + "deployment configuration since kernel does not support " + "no act_and_mul MLP layer." + ) +) def test_nemotron_fp8_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency") From 3d974c94050597ba3fb22f329bc9edf5ec95ad7f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 12:59:20 -0500 Subject: [PATCH 170/207] updated Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_fp8_moe.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 2ddf03bf85a2..8db804c663f3 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 @@ -21,7 +21,7 @@ kFp8Static128BlockSym, kFp8StaticTensorSym, ) -from vllm.v1.engine.utils import current_platform +from vllm.platforms import current_platform class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic): @@ -83,7 +83,7 @@ def _supports_current_device() -> bool: @staticmethod def _supports_no_act_and_mul() -> bool: - """Does not support non-gated MoE (i.e. Nanotron-Mini).""" + """Does not support non-gated MoE (i.e. Nanotron-3-Nano).""" return False @staticmethod @@ -181,7 +181,7 @@ def _apply_per_block( assert activation == MoEActivation.SILU if e_score_correction_bias is not None: - e_score_correction_bias.to(hidden_states.dtype) + e_score_correction_bias = e_score_correction_bias.to(hidden_states.dtype) if self.routing_method_type == RoutingMethodType.DeepSeekV3: router_logits = router_logits.to(torch.float32) From 89b0b6bfa32f771cca86b4cd92d09342db9fe2ee Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 13:22:46 -0500 Subject: [PATCH 171/207] assert in trtllm nvfp4 Signed-off-by: Robert Shaw --- .../model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py | 2 ++ 1 file changed, 2 insertions(+) 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 3fe909995256..c3af7820dcb5 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 @@ -259,6 +259,8 @@ def apply( assert a1q_scale is not None assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None + assert routed_scaling_factor is None + assert not apply_router_weight_on_input # Prepare routing bias into kernel format. routing_bias = e_score_correction_bias From 7acbb50bf06cf08675744b1cc190eb7a580e2e1e Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 17 Feb 2026 13:26:03 -0500 Subject: [PATCH 172/207] remove .impl Signed-off-by: Robert Shaw --- .../kernels/moe/modular_kernel_tools/common.py | 2 +- .../layers/fused_moe/fused_moe_method_base.py | 2 +- .../layers/fused_moe/modular_kernel.py | 18 +++++++++++++----- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 6fbb0503af08..c05c31c8f501 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -666,7 +666,7 @@ def run_modular_kernel( # impls might update the tensor in place hidden_states = rank_tensors.hidden_states.clone() - topk_ids = rank_tensors.topk_ids.to(mk.impl.prepare_finalize.topk_indices_dtype()) + topk_ids = rank_tensors.topk_ids.to(mk.prepare_finalize.topk_indices_dtype()) mk_kwargs = { "hidden_states": hidden_states, 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 1b16f2184971..c08bee7da7cf 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 @@ -98,7 +98,7 @@ def get_fused_moe_quant_config( @property def topk_indices_dtype(self) -> torch.dtype | None: if self.moe_kernel is not None: - return self.moe_kernel.impl.prepare_finalize.topk_indices_dtype() + return self.moe_kernel.prepare_finalize.topk_indices_dtype() return None @property diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index a74724ec40a1..746b0d6a563d 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1645,29 +1645,37 @@ def __init__( def is_monolithic(self) -> bool: return isinstance(self.impl, FusedMoEKernelMonolithicImpl) + @property + def prepare_finalize(self) -> FusedMoEPrepareAndFinalize: + return self.impl.prepare_finalize + + @property + def fused_experts(self) -> FusedMoEExperts: + return self.impl.fused_experts + def _post_init_setup(self): """ Resolve any leftover setup dependencies between self.prepare_finalize and self.fused_experts here. """ - self.impl.prepare_finalize.post_init_setup(self.impl.fused_experts) + self.prepare_finalize.post_init_setup(self.impl.fused_experts) assert ( - self.impl.prepare_finalize.activation_format - == self.impl.fused_experts.activation_format() + self.prepare_finalize.activation_format + == self.fused_experts.activation_format() ) def supports_expert_map(self) -> bool: """ A flag indicating whether or not this class supports expert maps. """ - return self.impl.fused_experts.supports_expert_map() + return self.fused_experts.supports_expert_map() def output_is_reduced(self) -> bool: """ Indicates whether or not the output of fused MoE kernel is reduced across all ranks. """ - return self.impl.prepare_finalize.output_is_reduced() + return self.prepare_finalize.output_is_reduced() def apply_monolithic( self, From 626607e2c055654bb17c4971725a44d685f041c3 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 22 Feb 2026 17:09:32 -0500 Subject: [PATCH 173/207] fix routed scaling factor Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 a09fafc75956..bfe8db2f14e9 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 @@ -264,7 +264,6 @@ def apply( assert a1q_scale is not None assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None - assert routed_scaling_factor is None assert not apply_router_weight_on_input # Prepare routing bias into kernel format. @@ -302,7 +301,7 @@ def apply( intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, - routed_scaling_factor=None, + routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, do_finalize=True, )[0] From 9bcb3c4384b5e444ce11bf3cfa80d57bd56119d8 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 14:43:51 -0500 Subject: [PATCH 174/207] fix assertion about routing method type Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) 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 bfe8db2f14e9..29238c7ba09d 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 @@ -232,7 +232,7 @@ def _supports_router_logits_dtype( routing_method: RoutingMethodType, ) -> bool: """ - The FlashInfer TRTLLM NVFp4 kernel expects bfloat16 router_logits by default. + The FlashInfer TRTLLM NvFp4 kernel expects bfloat16 router_logits by default. Only DeepSeekV3 routing supports float32 router_logits (which is converted internally in the kernel). """ @@ -264,7 +264,13 @@ def apply( assert a1q_scale is not None assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None - assert not apply_router_weight_on_input + assert ( + apply_router_weight_on_input + and self.routing_method_type == RoutingMethodType.Llama4 + ) or ( + not apply_router_weight_on_input + and self.routing_method_type != RoutingMethodType.Llama4 + ) # Prepare routing bias into kernel format. routing_bias = e_score_correction_bias From 31ced8610911423a0b61af896c154182f998e18d Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 14:47:06 -0500 Subject: [PATCH 175/207] fix use monolithic Signed-off-by: Robert Shaw --- tests/kernels/moe/test_nvfp4_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index f66094bd2e30..46ead4b6068e 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -221,7 +221,7 @@ def test_cutlass_fp4_moe_swiglustep( ) kernel = mk.FusedMoEKernel( - make_moe_prepare_and_finalize_no_dp_ep(is_monolithic=True), + make_moe_prepare_and_finalize_no_dp_ep(use_monolithic=True), CutlassExpertsFp4( moe_config=make_dummy_moe_config(), quant_config=quant_config, From 08b798719752055588f4c450549b0898409045de Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 15:51:56 -0500 Subject: [PATCH 176/207] potentially fix flashinfer nvfp4 Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) 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 29238c7ba09d..6535b311b986 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 @@ -53,7 +53,13 @@ def __init__( # g1_scale_c = a13_scale * w13_scale_2 / a2_scale assert self.quant_config.g1_alphas is not None assert self.quant_config.a2_gscale is not None - self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale + if moe_config.is_act_and_mul: + self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale + else: + self.g1_scale_c = ( + torch.ones_like(self.quant_config.a1_scale) + * self.quant_config.a2_gscale + ) @staticmethod def _supports_current_device() -> bool: @@ -236,7 +242,6 @@ def _supports_router_logits_dtype( Only DeepSeekV3 routing supports float32 router_logits (which is converted internally in the kernel). """ - # TODO: check this if router_logits_dtype == torch.float32: # Only DeepSeekV3 routing handles float32 logits # https://github.com/flashinfer-ai/flashinfer/issues/2469 From 5d277fee05556b801ab38b9270a7bb27397012d5 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 15:53:56 -0500 Subject: [PATCH 177/207] fix nemotron nvfp4 Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 6535b311b986..50da6dd5030b 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 @@ -48,16 +48,16 @@ def __init__( self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank - # g1_alpha_s = a13_scale * w13_scale_2 - # a2_gscale = (1 / a2_scale) - # g1_scale_c = a13_scale * w13_scale_2 / a2_scale assert self.quant_config.g1_alphas is not None assert self.quant_config.a2_gscale is not None if moe_config.is_act_and_mul: + # g1_alpha_s = a13_scale * w13_scale_2 + # a2_gscale = (1 / a2_scale) + # g1_scale_c = a13_scale * w13_scale_2 / a2_scale self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale else: self.g1_scale_c = ( - torch.ones_like(self.quant_config.a1_scale) + torch.ones_like(self.quant_config.a1_gscale) * self.quant_config.a2_gscale ) From 571b36b941d361c796224a25bbbf1000aa831e30 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 17:26:22 -0500 Subject: [PATCH 178/207] stash Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 3 +++ vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 1 + .../layers/fused_moe/prepare_finalize/no_dp_ep.py | 1 + 3 files changed, 5 insertions(+) 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 50da6dd5030b..5b2d88ff967f 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 @@ -287,6 +287,9 @@ def apply( else router_logits ) + print(f"{hidden_states.shape=}") + print(f"{a1q_scale.dtype=}") + # Invoke kernel. return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( routing_logits=router_logits, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index d65b52a374b0..072357bd4f45 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -353,6 +353,7 @@ def make_nvfp4_moe_quant_config( g1_alphas = a13_scale * w13_scale_2 g2_alphas = a2_scale * w2_scale_2 + print(f"{backend=}") return nvfp4_moe_quant_config( g1_alphas=g1_alphas, g2_alphas=g2_alphas, 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 b9d57da08326..a480e9802ea9 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 @@ -24,6 +24,7 @@ def _quantize_input( input_sf = ( quant_config.a1_gscale if quant_config.use_nvfp4_w4a4 else quant_config.a1_scale ) + print(f"{quant_config.is_nvfp4_scale_swizzled=}") a1q, a1q_scale = moe_kernel_quantize_input( a1, input_sf, From c7da363227f3a6eef511413aa8f90e23db26b2c9 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 17:26:43 -0500 Subject: [PATCH 179/207] revert Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 3 --- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 1 - .../layers/fused_moe/prepare_finalize/no_dp_ep.py | 1 - 3 files changed, 5 deletions(-) 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 5b2d88ff967f..50da6dd5030b 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 @@ -287,9 +287,6 @@ def apply( else router_logits ) - print(f"{hidden_states.shape=}") - print(f"{a1q_scale.dtype=}") - # Invoke kernel. return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( routing_logits=router_logits, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 072357bd4f45..d65b52a374b0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -353,7 +353,6 @@ def make_nvfp4_moe_quant_config( g1_alphas = a13_scale * w13_scale_2 g2_alphas = a2_scale * w2_scale_2 - print(f"{backend=}") return nvfp4_moe_quant_config( g1_alphas=g1_alphas, g2_alphas=g2_alphas, 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 a480e9802ea9..b9d57da08326 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 @@ -24,7 +24,6 @@ def _quantize_input( input_sf = ( quant_config.a1_gscale if quant_config.use_nvfp4_w4a4 else quant_config.a1_scale ) - print(f"{quant_config.is_nvfp4_scale_swizzled=}") a1q, a1q_scale = moe_kernel_quantize_input( a1, input_sf, From 830a97410ff85202d462c774e9de85a7c9b03f21 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 17:36:18 -0500 Subject: [PATCH 180/207] fix typo for flashinfer latency _supports_shape Signed-off-by: Robert Shaw --- tests/quantization/test_blackwell_moe.py | 7 +++++++ .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index 02a87cbc332f..2dd9381f5ac4 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -240,6 +240,13 @@ def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch) ) +@pytest.mark.skip( + reason=( + "FP4 MoE backend FLASHINFER_TRTLLM does not support the " + "deployment configuration since kernel does not support " + "hidden_dim % 512 != 0." + ) +) def test_nemotron_fp4_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1") monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency") 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 50da6dd5030b..ad67a4503ff5 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 @@ -89,7 +89,7 @@ def _supports_activation(activation: MoEActivation) -> bool: return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] @staticmethod - def supports_shape(hidden_dim: int) -> bool: + def _supports_shape(hidden_dim: int) -> bool: """Requires hidden dim to be multiple of 512.""" return hidden_dim % 512 == 0 From cb9614dc1327317dfaf638e7706dfb82d4787df7 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 17:47:45 -0500 Subject: [PATCH 181/207] fix fp8 oracle Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/oracle/fp8.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 12181e4eec84..a9980e198bff 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -277,11 +277,13 @@ def _return_or_raise( activation_format, ) - if supported: - logger.info_once(_make_log_backend(backend), scope="local") - return backend, k_cls - else: - logger.debug_once(_make_log_unsupported(backend, reason), scope="local") + if supported: + logger.info_once(_make_log_backend(backend), scope="local") + return backend, k_cls + else: + logger.debug_once( + _make_log_unsupported(backend, reason), scope="local" + ) raise NotImplementedError( "Found VLLM_USE_FLASHINFER_MOE_FP8=1, but no " From 2d0dcc5f11b60cf1883cb6fa20bded17bfb37fb2 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 17:59:38 -0500 Subject: [PATCH 182/207] fix fp8 trtllm for nemotron Signed-off-by: Robert Shaw --- vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 8db804c663f3..8d3b497dd562 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 @@ -84,7 +84,7 @@ def _supports_current_device() -> bool: @staticmethod def _supports_no_act_and_mul() -> bool: """Does not support non-gated MoE (i.e. Nanotron-3-Nano).""" - return False + return True @staticmethod def _supports_quant_scheme( From 5c4f163b554db45732c8c954bf9d2b348c41f601 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 18:50:06 -0500 Subject: [PATCH 183/207] update is_monolithic resolution time Signed-off-by: Robert Shaw --- .../layers/fused_moe/fused_moe_method_base.py | 7 ++++++- vllm/model_executor/layers/fused_moe/modular_kernel.py | 4 ++++ .../compressed_tensors/compressed_tensors_moe.py | 10 ---------- vllm/model_executor/layers/quantization/fp8.py | 5 ----- vllm/model_executor/layers/quantization/modelopt.py | 10 ---------- 5 files changed, 10 insertions(+), 26 deletions(-) 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 c08bee7da7cf..88cd173fe6a8 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 @@ -111,7 +111,12 @@ def method_name(self) -> str: @property def is_monolithic(self) -> bool: - return False + if self.moe_kernel is None: + if hasattr(self, "experts_cls"): + return self.experts_cls.is_monolithic() + else: + return False + return self.moe_kernel.is_monolithic def apply( self, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index f1695fb71d7b..4837566c450d 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -493,6 +493,10 @@ def __init__( self.max_num_tokens = max_num_tokens self.num_dispatchers = num_dispatchers + @staticmethod + def is_monolithic() -> bool: + raise NotImplementedError("Implemented by subclasses.") + @property def expects_unquantized_inputs(self) -> bool: """ diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 1344c10fd94e..8b7fc57d0409 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -585,11 +585,6 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon a2_scale=layer.w2_input_scale, ) - @property - def is_monolithic(self) -> bool: - assert self.moe_kernel is not None - return self.moe_kernel.is_monolithic - def apply_monolithic( self, layer: FusedMoE, @@ -952,11 +947,6 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon block_shape=self.weight_block_size, ) - @property - def is_monolithic(self) -> bool: - assert self.moe_kernel is not None - return self.moe_kernel.is_monolithic - def apply_monolithic( self, layer: FusedMoE, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 85881f60e1b4..de445b4ea0c6 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -925,11 +925,6 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon def supports_eplb(self) -> bool: return True - @property - def is_monolithic(self) -> bool: - assert self.moe_kernel is not None - return self.moe_kernel.is_monolithic - def apply_monolithic( self, layer: FusedMoE, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 999b4839f690..0567f8da0d86 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -906,11 +906,6 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon a2_scale=a2_scale, ) - @property - def is_monolithic(self) -> bool: - assert self.moe_kernel is not None - return self.moe_kernel.is_monolithic - def apply_monolithic( self, layer: FusedMoE, @@ -1417,11 +1412,6 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon def supports_eplb(self) -> bool: return True - @property - def is_monolithic(self) -> bool: - assert self.moe_kernel is not None - return self.moe_kernel.is_monolithic - def apply_monolithic( self, layer: FusedMoE, From 761ffd3cbcf505b5106ea7196cbceb4210c507cc Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 19:08:14 -0500 Subject: [PATCH 184/207] make DS example pass CI Signed-off-by: Robert Shaw --- tests/quantization/test_hopper_moe.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/quantization/test_hopper_moe.py b/tests/quantization/test_hopper_moe.py index cb758dda9266..dc7a3549d857 100644 --- a/tests/quantization/test_hopper_moe.py +++ b/tests/quantization/test_hopper_moe.py @@ -47,6 +47,10 @@ def can_initialize( "256", "--load-format", "dummy", + # FIXME: OOM at 0.8 with 4 layer model - needs investigation. + # This happens while capturing CUDAGraphs. + "--gpu-memory-utilization", + "0.85", "--trust-remote-code", "--limit-mm-per-prompt", json.dumps({"image": 0}), From 2cc0e48799677f2207dc325731cd6ac521690a7f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 19:09:33 -0500 Subject: [PATCH 185/207] updated Signed-off-by: Robert Shaw --- tests/quantization/test_hopper_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/quantization/test_hopper_moe.py b/tests/quantization/test_hopper_moe.py index dc7a3549d857..dcc0e790d5bb 100644 --- a/tests/quantization/test_hopper_moe.py +++ b/tests/quantization/test_hopper_moe.py @@ -50,7 +50,7 @@ def can_initialize( # FIXME: OOM at 0.8 with 4 layer model - needs investigation. # This happens while capturing CUDAGraphs. "--gpu-memory-utilization", - "0.85", + "0.80", "--trust-remote-code", "--limit-mm-per-prompt", json.dumps({"image": 0}), From 903fcbb2de730d5351fdee1fa9b0681d8d142d53 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 21:04:19 -0500 Subject: [PATCH 186/207] update how flashinfer moe size padding is saved Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_fp8_moe.py | 9 ++++++--- .../layers/quantization/utils/flashinfer_fp4_moe.py | 1 + .../layers/quantization/utils/flashinfer_utils.py | 1 + 3 files changed, 8 insertions(+), 3 deletions(-) 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 8d3b497dd562..9d6573661767 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 @@ -13,7 +13,7 @@ RoutingMethodType, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - activation_to_flashinfer_type, + activation_to_flashinfer_int, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -67,7 +67,7 @@ def __init__( self._g1_scale_c = ( self._g1_alphas / self.quant_config.a2_scale if moe_config.is_act_and_mul - else torch.ones_like(self._g1_alphas) * self.quant_config.a2_scale + else torch.ones_like(self._g1_alphas) / self.quant_config.a2_scale ) @staticmethod @@ -216,6 +216,7 @@ def _apply_per_block( local_num_experts=self.local_num_experts, routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, + use_shuffled_weight=False, ) def _apply_per_tensor( @@ -237,7 +238,9 @@ def _apply_per_tensor( ) -> torch.Tensor: # Confirm supported activation function. assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] - activation_type = activation_to_flashinfer_type(activation) + from flashinfer.fused_moe.core import ActivationType + + activation_type = ActivationType(activation_to_flashinfer_int(activation)) # Confirm Llama-4 routing is proper. if self.routing_method_type == RoutingMethodType.Llama4: diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index c44e3ebfc27e..9a3043038f13 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -266,6 +266,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( ) ) layer.intermediate_size_per_partition = padded_intermediate + layer.moe_config.intermediate_size_per_partition = padded_intermediate w13, w13_scale, w2, w2_scale = prepare_static_weights_for_trtllm_fp4_moe( w13, diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index d124c7d625ce..a8be1d61ac24 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -342,6 +342,7 @@ def prepare_fp8_moe_layer_for_fi( min_alignment, ) layer.intermediate_size_per_partition = new_intermediate + layer.moe_config.intermediate_size_per_partition = new_intermediate # FI kernels require W31 layout rather than W13. if layer.moe_config.is_act_and_mul: From 6e11e547c8de01109e6e8e7be79af4d590ec6591 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 23 Feb 2026 22:52:05 -0500 Subject: [PATCH 187/207] hopefully fix docs build Signed-off-by: Robert Shaw --- docs/design/moe_kernel_features.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 7998273ae0c6..39647b01457c 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -37,7 +37,7 @@ th { | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] | | flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] | -| MoEPrepareAndFinalizeNoDPEPModular5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoDPEPModular`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoDPEPModular] | +| MoEPrepareAndFinalizeNoDPEPModular5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoDPEPModular`][vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep.MoEPrepareAndFinalizeNoDPEPModular] | | BatchedPrepareAndFinalize5 | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] | !!! info "Table key" From 329ac6862a7ffcbe3477b6c666c8afeb9c37d1ae Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 24 Feb 2026 11:29:15 -0500 Subject: [PATCH 188/207] add wait for memory to clear Signed-off-by: Robert Shaw --- tests/kernels/moe/test_modular_oai_triton_moe.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 33a875a8c9ac..9e359a6ce96c 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -6,6 +6,7 @@ import pytest import torch +from vllm.tests.utils import wait_for_gpu_memory_to_clear from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.utils.import_utils import has_triton_kernels @@ -169,6 +170,8 @@ def oai_triton_moe_impl( topk_ids: torch.Tensor, unfused: bool = False, ) -> torch.Tensor: + wait_for_gpu_memory_to_clear() + quant_config = mxfp4_w4a16_moe_quant_config( w1_bias=w1_bias, w2_bias=w2_bias, From f07d0a32683d399421a14caa5ab3ab726c81760b Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 24 Feb 2026 11:31:05 -0500 Subject: [PATCH 189/207] try wait for memory to clear Signed-off-by: Robert Shaw --- tests/kernels/moe/test_modular_oai_triton_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 9e359a6ce96c..ec4d9c3866b0 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -6,8 +6,8 @@ import pytest import torch -from vllm.tests.utils import wait_for_gpu_memory_to_clear +from tests.utils import wait_for_gpu_memory_to_clear from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.utils.import_utils import has_triton_kernels From 957a292f84c225f0aa8782dd55485514e1867434 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 24 Feb 2026 11:31:46 -0500 Subject: [PATCH 190/207] try wait for memory to clear Signed-off-by: Robert Shaw --- tests/kernels/moe/test_modular_oai_triton_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index ec4d9c3866b0..bd4171fa5b92 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -170,7 +170,7 @@ def oai_triton_moe_impl( topk_ids: torch.Tensor, unfused: bool = False, ) -> torch.Tensor: - wait_for_gpu_memory_to_clear() + wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) quant_config = mxfp4_w4a16_moe_quant_config( w1_bias=w1_bias, From fcdf3e0032c441887901eb321e15281992fa1b24 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 24 Feb 2026 11:34:46 -0500 Subject: [PATCH 191/207] try wait for memory to clear Signed-off-by: Robert Shaw --- tests/kernels/moe/test_modular_oai_triton_moe.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index bd4171fa5b92..b730faf56c87 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -170,8 +170,6 @@ def oai_triton_moe_impl( topk_ids: torch.Tensor, unfused: bool = False, ) -> torch.Tensor: - wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) - quant_config = mxfp4_w4a16_moe_quant_config( w1_bias=w1_bias, w2_bias=w2_bias, @@ -227,6 +225,7 @@ def test_oai_triton_moe( unfused: bool, workspace_init, ): + wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) set_random_seed(0) ( w1, From cdc34e4007652ef506fbc981056b6c5c282fc103 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 24 Feb 2026 12:07:34 -0500 Subject: [PATCH 192/207] fix readthedocs Signed-off-by: Robert Shaw --- docs/design/moe_kernel_features.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 39647b01457c..c027914d7425 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -37,8 +37,7 @@ th { | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] | | flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] | -| MoEPrepareAndFinalizeNoDPEPModular5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoDPEPModular`][vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep.MoEPrepareAndFinalizeNoDPEPModular] | -| BatchedPrepareAndFinalize5 | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] | + !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 From 73957d281f6fa9e3111ad5b3f096050672c853d9 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 24 Feb 2026 15:45:55 -0500 Subject: [PATCH 193/207] fix monolithic Signed-off-by: Robert Shaw --- docs/design/moe_kernel_features.md | 1 - tests/kernels/moe/test_nvfp4_moe.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index c027914d7425..e7bb0f507a1f 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -38,7 +38,6 @@ th { | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] | | flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] | - !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 2. A,T quantization occurs after dispatch. diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index 46ead4b6068e..2b7ccdd5af6c 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -221,7 +221,7 @@ def test_cutlass_fp4_moe_swiglustep( ) kernel = mk.FusedMoEKernel( - make_moe_prepare_and_finalize_no_dp_ep(use_monolithic=True), + make_moe_prepare_and_finalize_no_dp_ep(use_monolithic=False), CutlassExpertsFp4( moe_config=make_dummy_moe_config(), quant_config=quant_config, From 49746c45819bf46e3423a533e67f467c9cf759dd Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Tue, 24 Feb 2026 16:03:48 -0500 Subject: [PATCH 194/207] split modular oai triton into separate launch Signed-off-by: Robert Shaw --- .buildkite/test_areas/kernels.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index afc8fc49a2aa..11dec086b242 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -44,7 +44,8 @@ steps: - vllm/envs.py - vllm/config commands: - - pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 2 - label: Kernels Mamba Test From 482d62251956612a96b432d954b4948e947233a5 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 14:16:12 -0500 Subject: [PATCH 195/207] use global state to get around flashinfer autotuning Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_nvfp4_moe.py | 10 +++++++--- vllm/model_executor/warmup/kernel_warmup.py | 15 +++++++++++++-- vllm/utils/flashinfer.py | 1 + 3 files changed, 21 insertions(+), 5 deletions(-) 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 5253fe409edd..502671766400 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 @@ -167,8 +167,12 @@ def apply( torch.bfloat16 ).view(torch.int16) - # Determine activation type - activation_type = activation_to_flashinfer_int(activation) + # trtllm_fp4_block_scale_routed_moe does not support autotuning + # so skip this kernel during dummy run for autotuning. + import vllm.utils.flashinfer as fi_utils + + if fi_utils._is_fi_autotuning: + return hidden_states # Invoke kernel. flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( @@ -200,7 +204,7 @@ def apply( routed_scaling_factor=None, routing_method_type=1, do_finalize=True, - activation_type=activation_type, + activation_type=activation_to_flashinfer_int(activation), output=output, ) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 1ba5981906ca..b14e1cc93641 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -36,6 +36,7 @@ def kernel_warmup(worker: "Worker"): max_tokens = worker.scheduler_config.max_num_batched_tokens deep_gemm_warmup(model, max_tokens) + model = worker.get_model() enable_flashinfer_autotune = ( worker.vllm_config.kernel_config.enable_flashinfer_autotune ) @@ -78,6 +79,9 @@ def _is_flashinfer_backend(backend): ) +_is_fi_autotuning: bool = False + + def flashinfer_autotune(runner: "GPUModelRunner") -> None: """ Autotune FlashInfer operations. @@ -88,9 +92,14 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: Without autotuning, FlashInfer will rely on heuristics, which may be significantly slower. """ - from vllm.utils.flashinfer import autotune + import vllm.utils.flashinfer as fi_utils + + with torch.inference_mode(), fi_utils.autotune(): + # Certain FlashInfer kernels (e.g. nvfp4 routed moe) are + # incompatible with autotuning. This state is used to skip + # those kernels during the autotuning process. + fi_utils._is_fi_autotuning = True - with torch.inference_mode(), autotune(): # We skip EPLB here since we don't want to record dummy metrics # When autotuning with number of tokens m, flashinfer will autotune # operations for all number of tokens up to m. @@ -100,3 +109,5 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: skip_eplb=True, is_profile=True, ) + + fi_utils._is_fi_autotuning = False diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 333e66f68a87..3ff3bb8d89b0 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -140,6 +140,7 @@ def wrapper(*args, **kwargs): "autotune", fallback_fn=lambda *args, **kwargs: contextlib.nullcontext(), ) +_is_fi_autotuning: bool = False @functools.cache From 2eb1970536cdf6cc377f412e6d56af11e89266e1 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 14:27:08 -0500 Subject: [PATCH 196/207] updated Signed-off-by: Robert Shaw --- tests/kernels/moe/test_nvfp4_moe.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index 2b7ccdd5af6c..e12659729c9c 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -229,13 +229,16 @@ def test_cutlass_fp4_moe_swiglustep( inplace=False, ) - cutlass_output = kernel( + cutlass_output = kernel.apply( hidden_states=a, w1=w1_q, w2=w2_q, topk_weights=topk_weights, topk_ids=topk_ids, activation=MoEActivation.SWIGLUSTEP, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, ) # Reference: dequantize everything and run torch_moe with swiglustep From 120b46612c7afa405cb077aac041199b49c46490 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 20:05:46 +0000 Subject: [PATCH 197/207] fix nits Signed-off-by: Robert Shaw --- vllm/model_executor/warmup/kernel_warmup.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index b14e1cc93641..70abd8a6c503 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -36,7 +36,6 @@ def kernel_warmup(worker: "Worker"): max_tokens = worker.scheduler_config.max_num_batched_tokens deep_gemm_warmup(model, max_tokens) - model = worker.get_model() enable_flashinfer_autotune = ( worker.vllm_config.kernel_config.enable_flashinfer_autotune ) @@ -79,9 +78,6 @@ def _is_flashinfer_backend(backend): ) -_is_fi_autotuning: bool = False - - def flashinfer_autotune(runner: "GPUModelRunner") -> None: """ Autotune FlashInfer operations. From de8c00fe127c5068ccae481c9f631796fddb5e0f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 20:58:20 +0000 Subject: [PATCH 198/207] updated Signed-off-by: Robert Shaw --- tests/quantization/test_blackwell_moe.py | 8 ---- tests/quantization/test_hopper_moe.py | 50 ++++++++++++------------ 2 files changed, 26 insertions(+), 32 deletions(-) diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index bab3227c6fd3..ebc821bee183 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -273,11 +273,3 @@ def test_nemotron_fp4_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch): hf_overrides=HF_OVERRIDE_TEXT, extra_args=["--moe-backend=flashinfer_trtllm"], ) - - -def test_nemotron_fp4_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): - can_initialize( - "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", - hf_overrides=HF_OVERRIDE_TEXT, - extra_args=["--moe-backend=triton"], - ) diff --git a/tests/quantization/test_hopper_moe.py b/tests/quantization/test_hopper_moe.py index dcc0e790d5bb..fc58e3cc4336 100644 --- a/tests/quantization/test_hopper_moe.py +++ b/tests/quantization/test_hopper_moe.py @@ -80,17 +80,18 @@ def can_initialize( def test_llama4_fp8_tensor_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "0") can_initialize( - "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM + "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=cutlass"], ) def test_llama4_fp8_tensor_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") can_initialize( - "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM + "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=flashinfer_cutlass"], ) @@ -98,37 +99,38 @@ def test_llama4_fp8_tensor_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatc def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1") - can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) + can_initialize( + "deepseek-ai/DeepSeek-V3.1", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=deep_gemm"], + ) def test_deepseek_fp8_block_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "0") - monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "0") - can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) + can_initialize( + "deepseek-ai/DeepSeek-V3.1", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=tritont"], + ) ## Qwen3 Next ## -def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1") - can_initialize("Qwen/Qwen3-Next-80B-A3B-Instruct", hf_overrides=HF_OVERRIDE_TEXT) +def test_qwen3_next_bf16_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): + can_initialize( + "Qwen/Qwen3-Next-80B-A3B-Instruct", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=flashinfer_cutlass"], + ) ## NemoTron ## -def test_nemotron_fp8_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") - can_initialize( - "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", hf_overrides=HF_OVERRIDE_TEXT - ) - - -def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "0") +def test_nemotron_fp8_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): can_initialize( - "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", hf_overrides=HF_OVERRIDE_TEXT + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=flashinfer_cutlass"], ) From c084de2a913cf7158064170ba834ae67623de826 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 20:59:54 +0000 Subject: [PATCH 199/207] fix hopper moe Signed-off-by: Robert Shaw --- tests/quantization/test_hopper_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/quantization/test_hopper_moe.py b/tests/quantization/test_hopper_moe.py index fc58e3cc4336..711737ad08e7 100644 --- a/tests/quantization/test_hopper_moe.py +++ b/tests/quantization/test_hopper_moe.py @@ -83,7 +83,7 @@ def test_llama4_fp8_tensor_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): can_initialize( "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_TEXT, - extra_args=["--moe-backend=cutlass"], + extra_args=["--moe-backend=triton"], ) From 5925b134afcb28ee8c3fe49888664c8bc48d3197 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 21:02:23 +0000 Subject: [PATCH 200/207] fix hopper test issue Signed-off-by: Robert Shaw --- tests/quantization/test_hopper_moe.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/quantization/test_hopper_moe.py b/tests/quantization/test_hopper_moe.py index 711737ad08e7..9aef07945b5d 100644 --- a/tests/quantization/test_hopper_moe.py +++ b/tests/quantization/test_hopper_moe.py @@ -47,10 +47,6 @@ def can_initialize( "256", "--load-format", "dummy", - # FIXME: OOM at 0.8 with 4 layer model - needs investigation. - # This happens while capturing CUDAGraphs. - "--gpu-memory-utilization", - "0.80", "--trust-remote-code", "--limit-mm-per-prompt", json.dumps({"image": 0}), @@ -82,7 +78,7 @@ def can_initialize( def test_llama4_fp8_tensor_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): can_initialize( "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", - hf_overrides=HF_OVERRIDE_TEXT, + hf_overrides=HF_OVERRIDE_MM, extra_args=["--moe-backend=triton"], ) @@ -90,7 +86,7 @@ def test_llama4_fp8_tensor_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): def test_llama4_fp8_tensor_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): can_initialize( "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", - hf_overrides=HF_OVERRIDE_TEXT, + hf_overrides=HF_OVERRIDE_MM, extra_args=["--moe-backend=flashinfer_cutlass"], ) From 86e5ca9d35802b2c410c61379da8bc9cd06625cc Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 21:05:53 +0000 Subject: [PATCH 201/207] tweak Signed-off-by: Robert Shaw --- tests/quantization/test_hopper_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/quantization/test_hopper_moe.py b/tests/quantization/test_hopper_moe.py index 9aef07945b5d..8dedda200e55 100644 --- a/tests/quantization/test_hopper_moe.py +++ b/tests/quantization/test_hopper_moe.py @@ -106,7 +106,7 @@ def test_deepseek_fp8_block_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): can_initialize( "deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT, - extra_args=["--moe-backend=tritont"], + extra_args=["--moe-backend=triton"], ) From c2030824eaa2b5a8dd941a7b73a769f3909d7e6c Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 22:08:51 +0000 Subject: [PATCH 202/207] remove hopper tests Signed-off-by: Robert Shaw --- .buildkite/test_areas/quantization.yaml | 17 --- tests/quantization/test_hopper_moe.py | 132 ------------------------ 2 files changed, 149 deletions(-) delete mode 100644 tests/quantization/test_hopper_moe.py diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index d355b9bee23d..5ee2e5186966 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -37,23 +37,6 @@ steps: commands: - pytest -s -v tests/quantization/test_blackwell_moe.py -- label: Quantized MoE Test (H100) - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/" - device: h100 - source_file_dependencies: - - tests/quantization/test_hopper_moe.py - - vllm/model_executor/models/deepseek_v2.py - - vllm/model_executor/models/gpt_oss.py - - vllm/model_executor/models/llama4.py - - vllm/model_executor/layers/fused_moe - - vllm/model_executor/layers/quantization/compressed_tensors - - vllm/model_executor/layers/quantization/modelopt.py - - vllm/model_executor/layers/quantization/mxfp4.py - - vllm/v1/attention/backends/flashinfer.py - commands: - - pytest -s -v tests/quantization/test_hopper_moe.py - - label: Quantized Models Test timeout_in_minutes: 60 source_file_dependencies: diff --git a/tests/quantization/test_hopper_moe.py b/tests/quantization/test_hopper_moe.py deleted file mode 100644 index 8dedda200e55..000000000000 --- a/tests/quantization/test_hopper_moe.py +++ /dev/null @@ -1,132 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import json -import os -from typing import Any - -import pytest - -from tests.utils import RemoteOpenAIServer -from vllm.platforms import current_platform - -if not current_platform.is_device_capability_family(90): - pytest.skip("This test only runs on Hopper GPUs (SM90x).", allow_module_level=True) - - -@pytest.fixture(scope="module", autouse=True) -def set_test_environment(): - """Sets environment variables required for this test module.""" - # Make sure TRTLLM attention is available - os.environ["VLLM_HAS_FLASHINFER_CUBIN"] = "1" - # Set compilation threads to 16 to speed up startup - os.environ["FLASHINFER_NVCC_THREADS"] = "16" - - -# Overide the backbone layers to 4 for faster startup -HF_OVERRIDE_TEXT = { - "num_layers": 4, - "num_hidden_layers": 4, -} -HF_OVERRIDE_MM = { - "text_config": {"num_layers": 4, "num_hidden_layers": 4}, -} - - -def can_initialize( - model: str, - hf_overrides: dict[str, Any] | None = None, - extra_args: list[str] | None = None, -): - # Server arguments - extra_args = extra_args if extra_args is not None else [] - server_args = [ - "--max-model-len", - "2048", - "--max-num-batched-tokens", - "256", - "--load-format", - "dummy", - "--trust-remote-code", - "--limit-mm-per-prompt", - json.dumps({"image": 0}), - *extra_args, - ] - - # Launch server and make a simple request - with RemoteOpenAIServer( - model, - server_args, - max_wait_seconds=1500, # Due to FlashInfer compile - override_hf_configs=hf_overrides, - ) as server: - client = server.get_client() - # Make a simple request to verify the server works - completion = client.completions.create( - model=model, - prompt=["Hello, World!"], - temperature=0, - max_tokens=2, - ) - print(completion) - assert completion.choices[0].text is not None - - -## Llama4 ## - - -def test_llama4_fp8_tensor_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): - can_initialize( - "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", - hf_overrides=HF_OVERRIDE_MM, - extra_args=["--moe-backend=triton"], - ) - - -def test_llama4_fp8_tensor_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): - can_initialize( - "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", - hf_overrides=HF_OVERRIDE_MM, - extra_args=["--moe-backend=flashinfer_cutlass"], - ) - - -## DeepSeekV3 ## - - -def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch): - can_initialize( - "deepseek-ai/DeepSeek-V3.1", - hf_overrides=HF_OVERRIDE_TEXT, - extra_args=["--moe-backend=deep_gemm"], - ) - - -def test_deepseek_fp8_block_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): - can_initialize( - "deepseek-ai/DeepSeek-V3.1", - hf_overrides=HF_OVERRIDE_TEXT, - extra_args=["--moe-backend=triton"], - ) - - -## Qwen3 Next ## - - -def test_qwen3_next_bf16_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): - can_initialize( - "Qwen/Qwen3-Next-80B-A3B-Instruct", - hf_overrides=HF_OVERRIDE_TEXT, - extra_args=["--moe-backend=flashinfer_cutlass"], - ) - - -## NemoTron ## - - -def test_nemotron_fp8_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): - can_initialize( - "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", - hf_overrides=HF_OVERRIDE_TEXT, - extra_args=["--moe-backend=flashinfer_cutlass"], - ) From 95f79b3feb35b13340a63b6a2abdc555f552b77f Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 17:09:04 -0500 Subject: [PATCH 203/207] fix blackwell tests Signed-off-by: Robert Shaw --- tests/quantization/test_blackwell_moe.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index ebc821bee183..fe44017a04ee 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -244,6 +244,13 @@ def test_nemotron_fp8_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch): ) +@pytest.mark.skip( + reason=( + "FP8 MoE backend TRITON does not support the " + "deployment configuration since kernel does not support " + "no act_and_mul MLP layer." + ) +) def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): can_initialize( "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", From 27e142e125b885cf38dfbe63dab3b6f2ba982c30 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Sun, 1 Mar 2026 17:14:35 -0500 Subject: [PATCH 204/207] remove spurious Signed-off-by: Robert Shaw --- .../layers/quantization/utils/flashinfer_fp4_moe.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 9a3043038f13..42677a5927b3 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -16,7 +16,6 @@ ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( - has_flashinfer_cutedsl_grouped_gemm_nt_masked, has_flashinfer_cutlass_fused_moe, ) @@ -44,16 +43,6 @@ def is_flashinfer_fp4_cutlass_moe_available() -> bool: ) -def is_flashinfer_fp4_cutedsl_moe_available() -> bool: - """Return ``True`` when FlashInfer CUTEDSL NV-FP4 kernels can be used.""" - return ( - envs.VLLM_USE_FLASHINFER_MOE_FP4 - and has_flashinfer_cutedsl_grouped_gemm_nt_masked() - and current_platform.is_cuda() - and current_platform.is_device_capability_family(100) - ) - - def reorder_w1w3_to_w3w1( weight: torch.Tensor, scale: torch.Tensor, dim: int = -2 ) -> tuple[torch.Tensor, torch.Tensor]: From 7540885185a897e5229c39239fdcbd94ba714172 Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 2 Mar 2026 10:51:29 -0500 Subject: [PATCH 205/207] fix flashinfer import Signed-off-by: Robert Shaw --- .../layers/fused_moe/experts/trtllm_fp8_moe.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) 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 9d6573661767..febb3b2ef0d7 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 @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import flashinfer import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -177,6 +176,9 @@ def _apply_per_block( routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: + # Delay import for non-CUDA. + import flashinfer + assert not apply_router_weight_on_input assert activation == MoEActivation.SILU @@ -236,9 +238,12 @@ def _apply_per_tensor( routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: + # Delay import for non-CUDA. + import flashinfer + from flashinfer.fused_moe.core import ActivationType + # Confirm supported activation function. assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] - from flashinfer.fused_moe.core import ActivationType activation_type = ActivationType(activation_to_flashinfer_int(activation)) From 0dac45ae126e546b0f2b96fc4acd64e378b4da6b Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 2 Mar 2026 15:59:29 -0500 Subject: [PATCH 206/207] fix the shuffle_weights Signed-off-by: Robert Shaw --- tests/kernels/moe/test_modular_oai_triton_moe.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index b730faf56c87..b071e72dafbb 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -34,11 +34,10 @@ UnfusedOAITritonExperts, ) from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel -from vllm.model_executor.layers.utils import shuffle_weight from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed -from .utils import make_dummy_moe_config +from .utils import make_dummy_moe_config, shuffle_weight MNK = [ (1, 512, 384), From 4df577e12ea4e5de800b79e158b0bbfe0284f02a Mon Sep 17 00:00:00 2001 From: Robert Shaw Date: Mon, 2 Mar 2026 16:01:48 -0500 Subject: [PATCH 207/207] remove llama 4 scout Signed-off-by: Robert Shaw --- tests/evals/gsm8k/configs/moe-refactor/config-h100.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt b/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt index f8ef6211051c..6cce0ad36324 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt +++ b/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt @@ -8,7 +8,5 @@ Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml Qwen3-30B-A3B-Fp8-CT-Channel-vllm-cutlass.yaml -Llama-4-Scout-Fp8-ModelOpt-marlin.yaml -Llama-4-Scout-Fp8-ModelOpt-triton.yaml Qwen3-30B-A3B-BF16-fi-cutlass.yaml Qwen3-30B-A3B-BF16-triton.yaml \ No newline at end of file