@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dc917cceb8
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 18:15 UTC

[MoE Refactor] Move `select_experts` from `FusedMoEQuantMethod` -> `FusedMoE` (#31996)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Jan 22, 2026 at 18:21 UTC dc917cceb877dfd13f98c538c4c96158047d98bd
22 files changed +495 -530
tests/kernels/moe/test_moe.py
+12 -16
@@ -957,18 +957,18 @@ class MarlinMoEWeightData:
957 )
958 @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
959 def test_fused_marlin_moe(
960 - a_type,
961 - b_type,
962 - c_type,
963 - group_blocks,
964 - m,
965 - n,
966 - k,
967 - e,
968 - topk,
969 - ep_size,
970 - act_order,
971 - is_k_full,
960 + a_type: ScalarType,
961 + b_type: ScalarType,
962 + c_type: ScalarType,
963 + group_blocks: int,
964 + m: int,
965 + n: int,
966 + k: int,
967 + e: int,
968 + topk: int,
969 + ep_size: int,
970 + act_order: bool,
971 + is_k_full: bool,
972 ):
973 torch.cuda.manual_seed(1)
974 group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
@@ -1044,7 +1044,6 @@ def test_fused_marlin_moe(
1044 None,
1045 w1_data.scales,
1046 w2_data.scales,
1047 - score,
1047 topk_weights,
1048 topk_ids,
1049 global_num_experts=e,
@@ -1120,7 +1119,6 @@ def test_fused_marlin_moe_with_bias(m):
1119 w2_data.marlin_bias,
1120 w1_data.scales,
1121 w2_data.scales,
1123 - score,
1122 topk_weights,
1123 topk_ids,
1124 global_num_experts=e,
@@ -1199,7 +1197,6 @@ def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):
1197 None, # bias2
1198 w1_data.scales,
1199 w2_data.scales,
1202 - score,
1200 topk_weights,
1201 topk_ids,
1202 global_num_experts=e,
@@ -1519,7 +1516,6 @@ def test_batched_fused_marlin_moe(
1516 "bias2": None,
1517 "w1_scale": w1_data.scales,
1518 "w2_scale": w2_data.scales,
1522 - "gating_output": score,
1519 "global_num_experts": e,
1520 "expert_map": None,
1521 "global_scale1": w1_data.global_scale,
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
-8
@@ -210,7 +210,6 @@ def fused_marlin_moe(
210 bias2: torch.Tensor | None,
211 w1_scale: torch.Tensor,
212 w2_scale: torch.Tensor,
213 - gating_output: torch.Tensor | None,
213 topk_weights: torch.Tensor,
214 topk_ids: torch.Tensor,
215 quant_type_id: int,
@@ -250,8 +249,6 @@ def fused_marlin_moe(
249 - w2 (torch.Tensor): The second set of expert weights.
250 - w1_scale (torch.Tensor): Scale to be used for w1.
251 - w2_scale (torch.Tensor): Scale to be used for w2.
253 - - gating_output (torch.Tensor|None): The output of the gating
254 - operation (before softmax).
252 - g_idx1 (torch.Tensor|None): The first set of act_order indices.
253 - g_idx2 (torch.Tensor|None): The second set of act_order indices.
254 - sort_indices1 (torch.Tensor|None): The first act_order input
@@ -292,8 +289,6 @@ def fused_marlin_moe(
289 topk = topk_ids.size(1)
290
291 # Check constraints.
295 - if gating_output is not None:
296 - assert gating_output.size(0) == M, "Number of tokens mismatch"
292 assert w1.size(1) * 16 == K, "Hidden size mismatch w1"
293 assert w2.size(2) // (num_bits // 2) == K, "Hidden size mismatch w2"
294 assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
@@ -381,7 +376,6 @@ def batched_fused_marlin_moe(
376 bias2: torch.Tensor | None,
377 w1_scale: torch.Tensor,
378 w2_scale: torch.Tensor,
384 - gating_output: torch.Tensor | None,
379 quant_type_id: int,
380 apply_router_weight_on_input: bool = False,
381 global_num_experts: int = -1,
@@ -718,7 +712,6 @@ class MarlinExperts(MarlinExpertsBase):
712 bias2=self.w2_bias,
713 w1_scale=self.w1_scale,
714 w2_scale=self.w2_scale,
721 - gating_output=None,
715 topk_weights=topk_weights,
716 topk_ids=topk_ids,
717 global_scale1=self.g1_alphas,
@@ -833,7 +826,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
826 bias2=self.w2_bias,
827 w1_scale=self.w1_scale,
828 w2_scale=self.w2_scale,
836 - gating_output=None,
829 quant_type_id=self.quant_type_id,
830 apply_router_weight_on_input=apply_router_weight_on_input,
831 global_num_experts=global_num_experts,
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+15 -5
@@ -14,9 +14,6 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
14 FusedMoEPermuteExpertsUnpermute,
15 FusedMoEPrepareAndFinalize,
16 )
17 -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
18 - FusedMoERouter,
19 -)
17 from vllm.model_executor.layers.quantization.base_config import (
18 QuantizeMethodBase,
19 )
@@ -108,11 +105,24 @@ class FusedMoEMethodBase(QuantizeMethodBase):
105 def method_name(self) -> str:
106 return self.__class__.__name__
107
111 - @abstractmethod
108 + @property
109 + def is_monolithic(self) -> bool:
110 + return False
111 +
112 + # @abstractmethod
113 def apply(
114 self,
115 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
115 - router: FusedMoERouter,
116 + x: torch.Tensor,
117 + topk_weights: torch.Tensor,
118 + topk_ids: torch.Tensor,
119 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
120 + raise NotImplementedError
121 +
122 + # @abstractmethod
123 + def apply_monolithic(
124 + self,
125 + layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
126 x: torch.Tensor,
127 router_logits: torch.Tensor,
128 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+4 -13
@@ -16,9 +16,6 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
16 FusedMoEModularKernel,
17 FusedMoEPrepareAndFinalize,
18 )
19 -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
20 - FusedMoERouter,
21 -)
19
20 logger = init_logger(__name__)
21
@@ -40,6 +37,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
37 not self.fused_experts.supports_expert_map(),
38 )
39 self.old_quant_method = old_quant_method
40 + assert not self.old_quant_method.is_monolithic
41 logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)
42
43 @staticmethod
@@ -94,16 +92,11 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
92 def apply(
93 self,
94 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
97 - router: FusedMoERouter,
95 x: torch.Tensor,
99 - router_logits: torch.Tensor,
96 + topk_weights: torch.Tensor,
97 + topk_ids: torch.Tensor,
98 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
101 - topk_weights, topk_ids = router.select_experts(
102 - hidden_states=x,
103 - router_logits=router_logits,
104 - )
105 -
106 - result = self.fused_experts(
99 + return self.fused_experts(
100 hidden_states=x,
101 w1=layer.w13_weight,
102 w2=layer.w2_weight,
@@ -115,5 +108,3 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
108 apply_router_weight_on_input=layer.apply_router_weight_on_input,
109 expert_map=None if self.disable_expert_map else layer.expert_map,
110 )
118 -
119 - return result
vllm/model_executor/layers/fused_moe/layer.py
+61 -22
@@ -522,8 +522,7 @@ class FusedMoE(CustomOp):
522 self.apply_router_weight_on_input = apply_router_weight_on_input
523 self.activation = activation
524
525 - # TODO(bnell): in next PR move capture back to layer
526 - capture: Callable[[torch.Tensor], None] | None = None
525 + self.capture: Callable[[torch.Tensor], None] | None = None
526 if (
527 self.vllm_config.model_config is not None
528 and self.vllm_config.model_config.enable_return_routed_experts
@@ -531,7 +530,9 @@ class FusedMoE(CustomOp):
530 # In dummy runs, the capturer is not initialized.
531 capturer = RoutedExpertsCapturer.get_instance()
532 if capturer is not None:
534 - capture = lambda topk_ids: capturer.capture(self.layer_id, topk_ids)
533 + self.capture = lambda topk_ids: capturer.capture(
534 + self.layer_id, topk_ids
535 + )
536
537 self.router = create_fused_moe_router(
538 top_k=top_k,
@@ -550,7 +551,6 @@ class FusedMoE(CustomOp):
551 # TODO(bnell): once we can construct the MK at init time, we
552 # can make this a value.
553 indices_type_getter=lambda: self.quant_method.topk_indices_dtype,
553 - capture=capture,
554 )
555 self.routing_method_type: RoutingMethodType = self.router.routing_method_type
556
@@ -1673,12 +1673,27 @@ class FusedMoE(CustomOp):
1673 staged_router_logits.copy_(router_logits, non_blocking=True)
1674
1675 # Matrix multiply.
1676 - final_hidden_states = self.quant_method.apply(
1677 - layer=self,
1678 - router=self.router,
1679 - x=staged_hidden_states,
1680 - router_logits=staged_router_logits,
1681 - )
1676 + if self.quant_method.is_monolithic:
1677 + final_hidden_states = self.quant_method.apply_monolithic(
1678 + layer=self,
1679 + x=staged_hidden_states,
1680 + router_logits=staged_router_logits,
1681 + )
1682 + else:
1683 + topk_weights, topk_ids = self.router.select_experts(
1684 + hidden_states=staged_hidden_states,
1685 + router_logits=staged_router_logits,
1686 + )
1687 +
1688 + if self.capture is not None:
1689 + self.capture(topk_ids)
1690 +
1691 + final_hidden_states = self.quant_method.apply(
1692 + layer=self,
1693 + x=staged_hidden_states,
1694 + topk_weights=topk_weights,
1695 + topk_ids=topk_ids,
1696 + )
1697
1698 if has_separate_shared_experts:
1699 assert not isinstance(final_hidden_states, tuple)
@@ -1810,15 +1825,20 @@ class FusedMoE(CustomOp):
1825 extra_tensors=extra_tensors,
1826 )
1827 if extra_tensors is not None:
1813 - hidden_states_combined, router_logits, extra_tensors_combined = (
1814 - dispatch_res
1815 - )
1828 + (
1829 + orig_hidden_states,
1830 + router_logits,
1831 + extra_tensors_combined,
1832 + ) = dispatch_res
1833 hidden_states_combined = (
1817 - hidden_states_combined,
1834 + orig_hidden_states,
1835 extra_tensors_combined[0],
1836 )
1837 else:
1838 hidden_states_combined, router_logits = dispatch_res
1839 + orig_hidden_states = hidden_states_combined
1840 + else:
1841 + orig_hidden_states = hidden_states
1842
1843 # Run shared experts before matrix multiply.
1844 # because matrix multiply maybe modify the hidden_states.
@@ -1840,14 +1860,33 @@ class FusedMoE(CustomOp):
1860 )
1861
1862 # Matrix multiply.
1843 - final_hidden_states = self.quant_method.apply(
1844 - layer=self,
1845 - router=self.router,
1846 - x=hidden_states_combined
1847 - if do_naive_dispatch_combine
1848 - else hidden_states,
1849 - router_logits=router_logits,
1850 - )
1863 + x = hidden_states_combined if do_naive_dispatch_combine else hidden_states
1864 +
1865 + # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014).
1866 + # Figure out nicer way to do this.
1867 + x_orig = orig_hidden_states if do_naive_dispatch_combine else hidden_states
1868 +
1869 + if self.quant_method.is_monolithic:
1870 + final_hidden_states = self.quant_method.apply_monolithic(
1871 + layer=self,
1872 + x=x,
1873 + router_logits=router_logits,
1874 + )
1875 + else:
1876 + topk_weights, topk_ids = self.router.select_experts(
1877 + hidden_states=x_orig,
1878 + router_logits=router_logits,
1879 + )
1880 +
1881 + if self.capture is not None:
1882 + self.capture(topk_ids)
1883 +
1884 + final_hidden_states = self.quant_method.apply(
1885 + layer=self,
1886 + x=x, # The type signture of this is wrong due to the hack.
1887 + topk_weights=topk_weights,
1888 + topk_ids=topk_ids,
1889 + )
1890
1891 if has_separate_shared_experts:
1892 assert self.shared_experts is not None
vllm/model_executor/layers/fused_moe/router/base_router.py
-5
@@ -127,7 +127,6 @@ class BaseRouter(FusedMoERouter):
127 self.eplb_state = eplb_state
128 self.enable_eplb = enable_eplb
129 self.indices_type_getter = indices_type_getter
130 - self.capture: Callable[[torch.tensor], None] | None = None
130
131 def _validate_eplb_state(self) -> None:
132 """Validate that EPLB state is properly initialized if EPLB is enabled."""
@@ -238,8 +237,4 @@ class BaseRouter(FusedMoERouter):
237 # Step 5: Convert indices dtype
238 topk_ids = self._convert_indices_dtype(topk_ids, indices_type)
239
241 - # TODO(bnell): temporary hack until select_experts is moved into FusedMoE
242 - if self.capture is not None:
243 - self.capture(topk_ids)
244 -
240 return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/custom_routing_router.py
+3 -1
@@ -55,4 +55,6 @@ class CustomRoutingRouter(BaseRouter):
55 renormalize=self.renormalize,
56 )
57
58 - return topk_weights.to(torch.float32), topk_ids
58 + return topk_weights.to(torch.float32), topk_ids.to(
59 + torch.int32 if indices_type is None else indices_type
60 + )
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+4 -1
@@ -124,7 +124,9 @@ def fused_topk_bias(
124 topk_weights = scores.gather(1, topk_indices)
125 if renormalize:
126 topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
127 - return topk_weights.to(torch.float32), topk_indices.to(torch.int32)
127 + return topk_weights.to(torch.float32), topk_indices.to(
128 + torch.int32 if indices_type is None else indices_type
129 + )
130
131
132 class FusedTopKBiasRouter(BaseRouter):
@@ -176,6 +178,7 @@ class FusedTopKBiasRouter(BaseRouter):
178 topk=self.top_k,
179 renormalize=self.renormalize,
180 scoring_func=self.scoring_func,
181 + indices_type=indices_type,
182 )
183
184 if self.routed_scaling_factor != 1.0:
vllm/model_executor/layers/fused_moe/router/router_factory.py
+5 -20
@@ -6,7 +6,6 @@ import torch
6
7 import vllm.envs as envs
8 from vllm.distributed.eplb.eplb_state import EplbLayerState
9 -from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
9 from vllm.model_executor.layers.fused_moe.router.custom_routing_router import (
10 CustomRoutingRouter,
11 )
@@ -49,7 +48,6 @@ def create_fused_moe_router(
48 # eplb parameters
49 enable_eplb: bool = False,
50 eplb_state: EplbLayerState = EMPTY_EPLB_STATE,
52 - capture: Callable[[torch.tensor], None] | None = None,
51 ) -> FusedMoERouter:
52 """
53 Factory function to create the appropriate FusedMoERouter subclass based on
@@ -90,21 +88,16 @@ def create_fused_moe_router(
88 Returns:
89 An instance of the appropriate FusedMoERouter subclass
90 """
93 - router: BaseRouter
91
92 routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
93 if routing_strategy != "":
97 - router = RoutingSimulatorRouter(
94 + return RoutingSimulatorRouter(
95 top_k=top_k,
96 global_num_experts=global_num_experts,
97 eplb_state=eplb_state,
98 enable_eplb=enable_eplb,
99 indices_type_getter=indices_type_getter,
100 )
104 - # TODO(bnell): this is temporary until select_experts is
105 - # separated from apply.
106 - router.capture = capture
107 - return router
101
102 if use_grouped_topk:
103 assert custom_routing_function is None
@@ -113,7 +106,7 @@ def create_fused_moe_router(
106 "num_expert_group and topk_group must be provided when "
107 "use_grouped_topk is True"
108 )
116 - router = GroupedTopKRouter(
109 + return GroupedTopKRouter(
110 top_k=top_k,
111 global_num_experts=global_num_experts,
112 eplb_state=eplb_state,
@@ -127,11 +120,9 @@ def create_fused_moe_router(
120 enable_eplb=enable_eplb,
121 indices_type_getter=indices_type_getter,
122 )
130 - router.capture = capture
131 - return router
123
124 if custom_routing_function is not None:
134 - router = CustomRoutingRouter(
125 + return CustomRoutingRouter(
126 top_k=top_k,
127 global_num_experts=global_num_experts,
128 eplb_state=eplb_state,
@@ -140,11 +131,9 @@ def create_fused_moe_router(
131 enable_eplb=enable_eplb,
132 indices_type_getter=indices_type_getter,
133 )
143 - router.capture = capture
144 - return router
134
135 if e_score_correction_bias is not None:
147 - router = FusedTopKBiasRouter(
136 + return FusedTopKBiasRouter(
137 top_k=top_k,
138 global_num_experts=global_num_experts,
139 eplb_state=eplb_state,
@@ -155,10 +144,8 @@ def create_fused_moe_router(
144 enable_eplb=enable_eplb,
145 indices_type_getter=indices_type_getter,
146 )
158 - router.capture = capture
159 - return router
147
161 - router = FusedTopKRouter(
148 + return FusedTopKRouter(
149 top_k=top_k,
150 global_num_experts=global_num_experts,
151 eplb_state=eplb_state,
@@ -167,5 +154,3 @@ def create_fused_moe_router(
154 enable_eplb=enable_eplb,
155 indices_type_getter=indices_type_getter,
156 )
170 - router.capture = capture
171 - return router
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+26 -46
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from collections.abc import Callable
5
6 import torch
7 import torch.nn.functional as F
@@ -31,9 +32,6 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
32 make_unquantized_moe_kernel,
33 select_unquantized_moe_backend,
34 )
34 -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
35 - FusedMoERouter,
36 -)
35 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
36 from vllm.platforms import current_platform
37 from vllm.platforms.interface import CpuArchEnum
@@ -66,6 +64,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
64 rocm_aiter_ops.is_fused_moe_enabled() and moe.is_act_and_mul
65 )
66 self.kernel: mk.FusedMoEModularKernel | None = None
67 + self._is_monolithic = current_platform.is_cpu() or current_platform.is_xpu()
68 +
69 + @property
70 + def is_monolithic(self) -> bool:
71 + return self._is_monolithic
72
73 @property
74 def supports_eplb(self) -> bool:
@@ -212,7 +215,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
215 import intel_extension_for_pytorch as ipex
216
217 ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts
215 - layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
218 + self.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
219 layer.w13_weight,
220 layer.w2_weight,
221 use_prepack=True,
@@ -244,11 +247,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
247 )
248 assert packed_w2_weight.size() == layer.w2_weight.size()
249 layer.w2_weight.copy_(packed_w2_weight)
247 - layer.cpu_fused_moe = cpu_fused_moe.SGLFusedMOE(layer)
250 + self.cpu_fused_moe: Callable = cpu_fused_moe.SGLFusedMOE(layer)
251 else:
249 - layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
252 + self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
253 else:
251 - layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
254 + self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
255 elif current_platform.is_cuda_alike():
256 self._setup_kernel(
257 layer=layer,
@@ -259,15 +262,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
262 def apply(
263 self,
264 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
262 - router: FusedMoERouter,
265 x: torch.Tensor,
264 - router_logits: torch.Tensor,
266 + topk_weights: torch.Tensor,
267 + topk_ids: torch.Tensor,
268 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
269 return self.forward(
267 - router=router,
270 layer=layer,
271 x=x,
270 - router_logits=router_logits,
272 + topk_weights=topk_weights,
273 + topk_ids=topk_ids,
274 )
275
276 def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
@@ -282,18 +285,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
285 def forward_cuda(
286 self,
287 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
285 - router: FusedMoERouter,
288 x: torch.Tensor,
287 - router_logits: torch.Tensor,
289 + topk_weights: torch.Tensor,
290 + topk_ids: torch.Tensor,
291 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
289 - assert self.kernel
290 -
291 - topk_weights, topk_ids = router.select_experts(
292 - hidden_states=x,
293 - router_logits=router_logits,
294 - )
295 -
296 - result = self.kernel(
292 + assert self.kernel is not None
293 + return self.kernel(
294 hidden_states=x,
295 w1=layer.w13_weight,
296 w2=layer.w2_weight,
@@ -306,24 +303,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
303 expert_map=layer.expert_map,
304 )
305
309 - return result
310 -
311 - def forward_cpu(
306 + def forward_monolithic_cpu(
307 self,
308 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
314 - router: FusedMoERouter,
309 x: torch.Tensor,
310 router_logits: torch.Tensor,
311 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
318 - if (
319 - layer.enable_eplb is not False
320 - or layer.eplb_state.expert_load_view is not None
321 - or layer.eplb_state.logical_to_physical_map is not None
322 - or layer.eplb_state.logical_replica_count is not None
323 - ):
324 - raise NotImplementedError("Expert load balancing is not supported for CPU.")
325 -
326 - return layer.cpu_fused_moe(
312 + return self.cpu_fused_moe(
313 layer,
314 x,
315 layer.use_grouped_topk,
@@ -342,21 +328,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
328 layer.activation,
329 )
330
345 - def forward_xpu(
331 + def forward_monolithic_xpu(
332 self,
333 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
348 - router: FusedMoERouter,
334 x: torch.Tensor,
335 router_logits: torch.Tensor,
336 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
352 - if (
353 - layer.enable_eplb is not False
354 - or layer.eplb_state.expert_load_view is not None
355 - or layer.eplb_state.logical_to_physical_map is not None
356 - or layer.eplb_state.logical_replica_count is not None
357 - ):
358 - raise NotImplementedError("Expert load balancing is not supported for XPU.")
359 - return layer.ipex_fusion(
337 + return self.ipex_fusion(
338 x,
339 layer.use_grouped_topk,
340 layer.top_k,
@@ -368,8 +346,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
346 )
347
348 if current_platform.is_cpu():
371 - forward_native = forward_cpu
349 + forward_native: Callable = forward_monolithic_cpu
350 + apply_monolithic = forward_monolithic_cpu
351 elif current_platform.is_xpu():
373 - forward_native = forward_xpu
352 + forward_native = forward_monolithic_xpu
353 + apply_monolithic = forward_monolithic_xpu
354 else:
355 forward_native = forward_cuda
vllm/model_executor/layers/quantization/awq_marlin.py
+2 -9
@@ -10,7 +10,6 @@ from torch.nn import Parameter
10 import vllm.model_executor.layers.fused_moe # noqa
11 from vllm import _custom_ops as ops
12 from vllm.logger import init_logger
13 -from vllm.model_executor.layers.fused_moe import FusedMoERouter
13 from vllm.model_executor.layers.fused_moe.config import (
14 FusedMoEConfig,
15 FusedMoEQuantConfig,
@@ -762,15 +761,10 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
761 def apply(
762 self,
763 layer: FusedMoE,
765 - router: FusedMoERouter,
764 x: torch.Tensor,
767 - router_logits: torch.Tensor,
765 + topk_weights: torch.Tensor,
766 + topk_ids: torch.Tensor,
767 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
769 - topk_weights, topk_ids = router.select_experts(
770 - hidden_states=x,
771 - router_logits=router_logits,
772 - )
773 -
768 return fused_marlin_moe(
769 x,
770 layer.w13_qweight,
@@ -779,7 +773,6 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
773 getattr(layer, "w2_bias", None),
774 layer.w13_scales,
775 layer.w2_scales,
782 - router_logits,
776 topk_weights,
777 topk_ids,
778 input_global_scale1=getattr(layer, "w13_input_global_scale", None),
vllm/model_executor/layers/quantization/bitsandbytes.py
+2 -7
@@ -6,7 +6,6 @@ from typing import Any, Union
6 import torch
7 from packaging import version
8
9 -from vllm.model_executor.layers.fused_moe import FusedMoERouter
9 from vllm.model_executor.layers.fused_moe.config import (
10 FusedMoEConfig,
11 FusedMoEQuantConfig,
@@ -499,16 +498,12 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
498 def apply(
499 self,
500 layer: FusedMoE,
502 - router: FusedMoERouter,
501 x: torch.Tensor,
504 - router_logits: torch.Tensor,
502 + topk_weights: torch.Tensor,
503 + topk_ids: torch.Tensor,
504 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
505 from vllm.model_executor.layers.fused_moe import fused_experts
506
508 - topk_weights, topk_ids = router.select_experts(
509 - hidden_states=x,
510 - router_logits=router_logits,
511 - )
507 # TODO(bnell): Do these need to be called on the hot path?
508 if self.quant_config.load_in_8bit:
509 w13, w2 = self._apply_8bit_dequant(layer)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+106 -126
@@ -21,7 +21,6 @@ from vllm.model_executor.layers.fused_moe import (
21 FusedMoEActivationFormat,
22 FusedMoEMethodBase,
23 FusedMoEPermuteExpertsUnpermute,
24 - FusedMoERouter,
24 FusedMoeWeightScaleSupported,
25 UnquantizedFusedMoEMethod,
26 )
@@ -126,7 +125,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
125 quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
126 layer: torch.nn.Module,
127 layer_name: str,
129 - ) -> "CompressedTensorsMoEMethod":
128 + ) -> FusedMoEMethodBase:
129 # FusedMoE was made by combining multiple Linears so need to
130 # make sure quantization config for Linear can target it
131 quant_config._add_fused_moe_to_target_scheme_map()
@@ -345,19 +344,10 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
344 def apply(
345 self,
346 layer: FusedMoE,
348 - router: FusedMoERouter,
347 x: torch.Tensor,
350 - router_logits: torch.Tensor,
348 + topk_weights: torch.Tensor,
349 + topk_ids: torch.Tensor,
350 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
352 - if isinstance(x, tuple):
353 - x_routing, _ = x
354 - else:
355 - x_routing = x
356 -
357 - topk_weights, topk_ids = router.select_experts(
358 - hidden_states=x_routing,
359 - router_logits=router_logits,
360 - )
351 assert self.kernel is not None
352 return self.kernel(
353 x,
@@ -639,42 +629,48 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
629 a2_scale=layer.w2_input_scale,
630 )
631
642 - def apply(
632 + @property
633 + def is_monolithic(self) -> bool:
634 + return (
635 + self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM
636 + and not self.moe.moe_parallel_config.enable_eplb
637 + )
638 +
639 + def apply_monolithic(
640 self,
641 layer: FusedMoE,
645 - router: FusedMoERouter,
642 x: torch.Tensor,
643 router_logits: torch.Tensor,
644 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
645 + assert self.is_monolithic
646 assert layer.activation == "silu", "Only SiLU activation is supported."
650 -
651 - if (
647 + assert (
648 self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM
649 and not layer.enable_eplb
654 - ):
655 - return flashinfer_trtllm_fp4_moe(
656 - layer=layer,
657 - x=x,
658 - router_logits=router_logits,
659 - top_k=layer.top_k,
660 - activation=layer.activation,
661 - global_num_experts=layer.global_num_experts,
662 - num_expert_group=layer.num_expert_group,
663 - topk_group=layer.topk_group,
664 - custom_routing_function=layer.custom_routing_function,
665 - e_score_correction_bias=layer.e_score_correction_bias,
666 - )
667 -
668 - # Hidden_states in select_experts is only used to extract metadata
669 - if isinstance(x, tuple):
670 - x_routing, _ = x
671 - else:
672 - x_routing = x
673 - topk_weights, topk_ids = router.select_experts(
674 - hidden_states=x_routing,
650 + )
651 + return flashinfer_trtllm_fp4_moe(
652 + layer=layer,
653 + x=x,
654 router_logits=router_logits,
655 + top_k=layer.top_k,
656 + activation=layer.activation,
657 + global_num_experts=layer.global_num_experts,
658 + num_expert_group=layer.num_expert_group,
659 + topk_group=layer.topk_group,
660 + custom_routing_function=layer.custom_routing_function,
661 + e_score_correction_bias=layer.e_score_correction_bias,
662 )
663
664 + def apply(
665 + self,
666 + layer: FusedMoE,
667 + x: torch.Tensor,
668 + topk_weights: torch.Tensor,
669 + topk_ids: torch.Tensor,
670 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
671 + assert not self.is_monolithic
672 + assert layer.activation == "silu", "Only SiLU activation is supported."
673 +
674 # EPLB path
675 if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:
676 assert layer.enable_eplb
@@ -1059,70 +1055,73 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1055 block_shape=self.weight_block_size,
1056 )
1057
1062 - def apply(
1058 + @property
1059 + def is_monolithic(self) -> bool:
1060 + return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
1061 +
1062 + def apply_monolithic(
1063 self,
1064 layer: FusedMoE,
1065 - router: FusedMoERouter,
1065 x: torch.Tensor,
1066 router_logits: torch.Tensor,
1067 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1069 - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
1070 - if layer.enable_eplb:
1071 - raise NotImplementedError(
1072 - "EPLB not supported for `FlashInfer TRTLLM FP8 MoE`."
1073 - )
1074 - assert layer.activation == "silu"
1075 -
1076 - if self.block_quant:
1077 - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
1068 + assert self.is_monolithic
1069 + assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
1070 + assert layer.activation == "silu"
1071
1079 - e_score_correction_bias = (
1080 - layer.e_score_correction_bias.to(x.dtype)
1081 - if layer.e_score_correction_bias is not None
1082 - else None
1083 - )
1084 - routing_method_type = layer.routing_method_type
1085 - return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
1086 - routing_logits=router_logits.to(torch.float32)
1087 - if routing_method_type == RoutingMethodType.DeepSeekV3
1088 - else router_logits,
1089 - routing_bias=e_score_correction_bias,
1090 - x=x,
1091 - w13_weight=layer.w13_weight,
1092 - w13_weight_scale_inv=layer.w13_weight_scale,
1093 - w2_weight=layer.w2_weight,
1094 - w2_weight_scale_inv=layer.w2_weight_scale,
1095 - global_num_experts=layer.global_num_experts,
1096 - top_k=layer.top_k,
1097 - num_expert_group=layer.num_expert_group,
1098 - topk_group=layer.topk_group,
1099 - intermediate_size=layer.intermediate_size_per_partition,
1100 - expert_offset=layer.ep_rank * layer.local_num_experts,
1101 - local_num_experts=layer.local_num_experts,
1102 - block_shape=self.weight_block_size,
1103 - routing_method_type=routing_method_type,
1104 - routed_scaling=layer.routed_scaling_factor,
1105 - )
1106 - else:
1107 - return apply_fi_trtllm_fp8_per_tensor_moe(
1108 - layer=layer,
1109 - hidden_states=x,
1110 - router_logits=router_logits,
1111 - routing_bias=layer.e_score_correction_bias,
1112 - global_num_experts=layer.global_num_experts,
1113 - top_k=layer.top_k,
1114 - num_expert_group=layer.num_expert_group,
1115 - topk_group=layer.topk_group,
1116 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
1117 - )
1072 + if self.block_quant:
1073 + import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
1074
1119 - topk_weights, topk_ids = router.select_experts(
1120 - hidden_states=x,
1121 - router_logits=router_logits,
1122 - )
1075 + e_score_correction_bias = (
1076 + layer.e_score_correction_bias.to(x.dtype)
1077 + if layer.e_score_correction_bias is not None
1078 + else None
1079 + )
1080 + routing_method_type = layer.routing_method_type
1081 + return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
1082 + routing_logits=router_logits.to(torch.float32)
1083 + if routing_method_type == RoutingMethodType.DeepSeekV3
1084 + else router_logits,
1085 + routing_bias=e_score_correction_bias,
1086 + x=x,
1087 + w13_weight=layer.w13_weight,
1088 + w13_weight_scale_inv=layer.w13_weight_scale,
1089 + w2_weight=layer.w2_weight,
1090 + w2_weight_scale_inv=layer.w2_weight_scale,
1091 + global_num_experts=layer.global_num_experts,
1092 + top_k=layer.top_k,
1093 + num_expert_group=layer.num_expert_group,
1094 + topk_group=layer.topk_group,
1095 + intermediate_size=layer.intermediate_size_per_partition,
1096 + expert_offset=layer.ep_rank * layer.local_num_experts,
1097 + local_num_experts=layer.local_num_experts,
1098 + block_shape=self.weight_block_size,
1099 + routing_method_type=routing_method_type,
1100 + routed_scaling=layer.routed_scaling_factor,
1101 + )
1102 + else:
1103 + return apply_fi_trtllm_fp8_per_tensor_moe(
1104 + layer=layer,
1105 + hidden_states=x,
1106 + router_logits=router_logits,
1107 + routing_bias=layer.e_score_correction_bias,
1108 + global_num_experts=layer.global_num_experts,
1109 + top_k=layer.top_k,
1110 + num_expert_group=layer.num_expert_group,
1111 + topk_group=layer.topk_group,
1112 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
1113 + )
1114
1115 + def apply(
1116 + self,
1117 + layer: FusedMoE,
1118 + x: torch.Tensor,
1119 + topk_weights: torch.Tensor,
1120 + topk_ids: torch.Tensor,
1121 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1122 + assert not self.is_monolithic
1123 assert self.kernel is not None
1125 - result = self.kernel(
1124 + return self.kernel(
1125 x,
1126 layer.w13_weight,
1127 layer.w2_weight,
@@ -1137,8 +1136,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1136 apply_router_weight_on_input=layer.apply_router_weight_on_input,
1137 )
1138
1140 - return result
1141 -
1139 @property
1140 def supports_eplb(self) -> bool:
1141 return True
@@ -1257,17 +1254,12 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1254 def apply(
1255 self,
1256 layer: FusedMoE,
1260 - router: FusedMoERouter,
1257 x: torch.Tensor,
1262 - router_logits: torch.Tensor,
1258 + topk_weights: torch.Tensor,
1259 + topk_ids: torch.Tensor,
1260 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1261 from vllm.model_executor.layers.fused_moe import fused_experts
1262
1266 - topk_weights, topk_ids = router.select_experts(
1267 - hidden_states=x,
1268 - router_logits=router_logits,
1269 - )
1270 -
1263 return fused_experts(
1264 hidden_states=x,
1265 w1=layer.w13_weight,
@@ -1621,15 +1613,10 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1613 def apply(
1614 self,
1615 layer: FusedMoE,
1624 - router: FusedMoERouter,
1616 x: torch.Tensor,
1626 - router_logits: torch.Tensor,
1617 + topk_weights: torch.Tensor,
1618 + topk_ids: torch.Tensor,
1619 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1628 - topk_weights, topk_ids = router.select_experts(
1629 - hidden_states=x,
1630 - router_logits=router_logits,
1631 - )
1632 -
1620 return fused_marlin_moe(
1621 x,
1622 layer.w13_weight_packed,
@@ -1638,7 +1625,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1625 None,
1626 layer.w13_weight_scale,
1627 layer.w2_weight_scale,
1641 - router_logits,
1628 topk_weights,
1629 topk_ids,
1630 input_global_scale1=getattr(layer, "w13_input_global_scale", None),
@@ -1873,17 +1859,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1859 def apply(
1860 self,
1861 layer: FusedMoE,
1876 - router: FusedMoERouter,
1862 x: torch.Tensor,
1878 - router_logits: torch.Tensor,
1863 + topk_weights: torch.Tensor,
1864 + topk_ids: torch.Tensor,
1865 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1866 from vllm.model_executor.layers.fused_moe import fused_experts
1867
1882 - topk_weights, topk_ids = router.select_experts(
1883 - hidden_states=x,
1884 - router_logits=router_logits,
1885 - )
1886 -
1868 return fused_experts(
1869 x,
1870 layer.w13_weight_packed,
@@ -2172,10 +2153,13 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
2153 # fused_experts; quant config is not needed.
2154 return None
2155
2175 - def apply(
2156 + @property
2157 + def is_monolithic(self) -> bool:
2158 + return True
2159 +
2160 + def apply_monolithic(
2161 self,
2162 layer: FusedMoE,
2178 - router: FusedMoERouter,
2163 x: torch.Tensor,
2164 router_logits: torch.Tensor,
2165 ) -> torch.Tensor:
@@ -2489,19 +2473,15 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
2473 def apply(
2474 self,
2475 layer: FusedMoE,
2492 - router: FusedMoERouter,
2476 x: torch.Tensor,
2494 - router_logits: torch.Tensor,
2495 - ):
2477 + topk_weights: torch.Tensor,
2478 + topk_ids: torch.Tensor,
2479 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
2480 if layer.enable_eplb:
2481 raise NotImplementedError(
2482 "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."
2483 )
2484 assert self.moe_quant_config is not None
2501 - topk_weights, topk_ids = router.select_experts(
2502 - hidden_states=x,
2503 - router_logits=router_logits,
2504 - )
2485
2486 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
2487 cutlass_moe_w4a8_fp8,
vllm/model_executor/layers/quantization/experts_int8.py
+2 -8
@@ -10,7 +10,6 @@ from vllm.model_executor.layers.fused_moe import (
10 FusedMoE,
11 FusedMoEConfig,
12 FusedMoEMethodBase,
13 - FusedMoERouter,
13 )
14 from vllm.model_executor.layers.fused_moe.config import (
15 FusedMoEQuantConfig,
@@ -138,17 +137,12 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
137 def apply(
138 self,
139 layer: FusedMoE,
141 - router: FusedMoERouter,
140 x: torch.Tensor,
143 - router_logits: torch.Tensor,
141 + topk_weights: torch.Tensor,
142 + topk_ids: torch.Tensor,
143 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
144 from vllm.model_executor.layers.fused_moe import fused_experts
145
147 - topk_weights, topk_ids = router.select_experts(
148 - hidden_states=x,
149 - router_logits=router_logits,
150 - )
151 -
146 return fused_experts(
147 x,
148 layer.w13_weight,
vllm/model_executor/layers/quantization/fp8.py
+64 -59
@@ -22,7 +22,6 @@ from vllm.model_executor.layers.fused_moe import (
22 FusedMoEMethodBase,
23 FusedMoEPermuteExpertsUnpermute,
24 FusedMoEPrepareAndFinalize,
25 - FusedMoERouter,
25 FusedMoeWeightScaleSupported,
26 )
27 from vllm.model_executor.layers.fused_moe.config import (
@@ -968,71 +967,79 @@ class Fp8MoEMethod(FusedMoEMethodBase):
967 def allow_inplace(self) -> bool:
968 return True
969
971 - def apply(
970 + @property
971 + def is_monolithic(self) -> bool:
972 + return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
973 +
974 + def apply_monolithic(
975 self,
976 layer: FusedMoE,
974 - router: FusedMoERouter,
977 x: torch.Tensor,
978 router_logits: torch.Tensor,
979 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
978 - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
979 - # TODO(rob): convert this to MK.
980 - if layer.enable_eplb:
981 - raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.")
982 - assert layer.activation == "silu", (
983 - f"Expected 'silu' activation but got {layer.activation}"
984 - )
985 -
986 - if self.block_quant:
987 - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
980 + assert self.is_monolithic
981 + assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
982 +
983 + # TODO(rob): convert this to MK.
984 + if layer.enable_eplb:
985 + raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.")
986 + assert layer.activation == "silu", (
987 + f"Expected 'silu' activation but got {layer.activation}"
988 + )
989
989 - e_score_correction_bias = (
990 - layer.e_score_correction_bias.to(x.dtype)
991 - if layer.e_score_correction_bias is not None
992 - else None
993 - )
994 - routing_method_type = layer.routing_method_type
995 - return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
996 - routing_logits=router_logits.to(torch.float32)
997 - if routing_method_type == RoutingMethodType.DeepSeekV3
998 - else router_logits,
999 - routing_bias=e_score_correction_bias,
1000 - x=x,
1001 - w13_weight=layer.w13_weight,
1002 - w13_weight_scale_inv=layer.w13_weight_scale_inv,
1003 - w2_weight=layer.w2_weight,
1004 - w2_weight_scale_inv=layer.w2_weight_scale_inv,
1005 - global_num_experts=layer.global_num_experts,
1006 - top_k=layer.top_k,
1007 - num_expert_group=layer.num_expert_group,
1008 - topk_group=layer.topk_group,
1009 - intermediate_size=layer.intermediate_size_per_partition,
1010 - expert_offset=layer.ep_rank * layer.local_num_experts,
1011 - local_num_experts=layer.local_num_experts,
1012 - block_shape=self.weight_block_size,
1013 - routing_method_type=routing_method_type,
1014 - routed_scaling=layer.routed_scaling_factor,
1015 - )
1016 - else:
1017 - return apply_fi_trtllm_fp8_per_tensor_moe(
1018 - layer=layer,
1019 - hidden_states=x,
1020 - router_logits=router_logits,
1021 - routing_bias=layer.e_score_correction_bias,
1022 - global_num_experts=layer.global_num_experts,
1023 - top_k=layer.top_k,
1024 - num_expert_group=layer.num_expert_group,
1025 - topk_group=layer.topk_group,
1026 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
1027 - )
990 + if self.block_quant:
991 + import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
992
1029 - topk_weights, topk_ids = router.select_experts(
1030 - hidden_states=x,
1031 - router_logits=router_logits,
1032 - )
993 + e_score_correction_bias = (
994 + layer.e_score_correction_bias.to(x.dtype)
995 + if layer.e_score_correction_bias is not None
996 + else None
997 + )
998 + routing_method_type = layer.routing_method_type
999 + return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
1000 + routing_logits=router_logits.to(torch.float32)
1001 + if routing_method_type == RoutingMethodType.DeepSeekV3
1002 + else router_logits,
1003 + routing_bias=e_score_correction_bias,
1004 + x=x,
1005 + w13_weight=layer.w13_weight,
1006 + w13_weight_scale_inv=layer.w13_weight_scale_inv,
1007 + w2_weight=layer.w2_weight,
1008 + w2_weight_scale_inv=layer.w2_weight_scale_inv,
1009 + global_num_experts=layer.global_num_experts,
1010 + top_k=layer.top_k,
1011 + num_expert_group=layer.num_expert_group,
1012 + topk_group=layer.topk_group,
1013 + intermediate_size=layer.intermediate_size_per_partition,
1014 + expert_offset=layer.ep_rank * layer.local_num_experts,
1015 + local_num_experts=layer.local_num_experts,
1016 + block_shape=self.weight_block_size,
1017 + routing_method_type=routing_method_type,
1018 + routed_scaling=layer.routed_scaling_factor,
1019 + )
1020 + else:
1021 + return apply_fi_trtllm_fp8_per_tensor_moe(
1022 + layer=layer,
1023 + hidden_states=x,
1024 + router_logits=router_logits,
1025 + routing_bias=layer.e_score_correction_bias,
1026 + global_num_experts=layer.global_num_experts,
1027 + top_k=layer.top_k,
1028 + num_expert_group=layer.num_expert_group,
1029 + topk_group=layer.topk_group,
1030 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
1031 + )
1032
1033 + def apply(
1034 + self,
1035 + layer: FusedMoE,
1036 + x: torch.Tensor,
1037 + topk_weights: torch.Tensor,
1038 + topk_ids: torch.Tensor,
1039 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1040 assert self.kernel is not None
1035 - result = self.kernel(
1041 + assert not self.is_monolithic
1042 + return self.kernel(
1043 x,
1044 layer.w13_weight,
1045 layer.w2_weight,
@@ -1045,8 +1052,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1052 apply_router_weight_on_input=layer.apply_router_weight_on_input,
1053 )
1054
1048 - return result
1049 -
1055
1056 class Fp8OnlineMoEMethod(Fp8MoEMethod):
1057 """MoE method for online FP8 quantization.
vllm/model_executor/layers/quantization/gguf.py
+2 -7
@@ -12,7 +12,6 @@ from torch.nn.parameter import Parameter, UninitializedParameter
12
13 from vllm import _custom_ops as ops
14 from vllm.logger import init_logger
15 -from vllm.model_executor.layers.fused_moe import FusedMoERouter
15 from vllm.model_executor.layers.fused_moe.config import (
16 FusedMoEConfig,
17 FusedMoEQuantConfig,
@@ -633,9 +632,9 @@ class GGUFMoEMethod(FusedMoEMethodBase):
632 def apply(
633 self,
634 layer: FusedMoE,
636 - router: FusedMoERouter,
635 x: torch.Tensor,
638 - router_logits: torch.Tensor,
636 + topk_weights: torch.Tensor,
637 + topk_ids: torch.Tensor,
638 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
639 assert layer.activation == "silu", "Only SiLU activation is supported."
640 if layer.apply_router_weight_on_input:
@@ -644,10 +643,6 @@ class GGUFMoEMethod(FusedMoEMethodBase):
643 "fused GGUF MoE method."
644 )
645
647 - topk_weights, topk_ids = router.select_experts(
648 - hidden_states=x,
649 - router_logits=router_logits,
650 - )
646 return fused_moe_gguf(
647 x,
648 layer.w13_qweight,
vllm/model_executor/layers/quantization/gptq_marlin.py
+2 -9
@@ -10,7 +10,6 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
10 import vllm.model_executor.layers.fused_moe # noqa
11 from vllm import _custom_ops as ops
12 from vllm.logger import init_logger
13 -from vllm.model_executor.layers.fused_moe import FusedMoERouter
13 from vllm.model_executor.layers.fused_moe.config import (
14 FusedMoEConfig,
15 FusedMoEQuantConfig,
@@ -898,15 +897,10 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
897 def apply(
898 self,
899 layer: FusedMoE,
901 - router: FusedMoERouter,
900 x: torch.Tensor,
903 - router_logits: torch.Tensor,
901 + topk_weights: torch.Tensor,
902 + topk_ids: torch.Tensor,
903 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
905 - topk_weights, topk_ids = router.select_experts(
906 - hidden_states=x,
907 - router_logits=router_logits,
908 - )
909 -
904 return fused_marlin_moe(
905 x,
906 layer.w13_qweight,
@@ -915,7 +909,6 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
909 getattr(layer, "w2_bias", None),
910 layer.w13_scales,
911 layer.w2_scales,
918 - router_logits,
912 topk_weights,
913 topk_ids,
914 input_global_scale1=getattr(layer, "w13_input_global_scale", None),
vllm/model_executor/layers/quantization/ipex_quant.py
+5 -5
@@ -8,9 +8,6 @@ from packaging import version
8 from torch.nn import Module
9
10 from vllm._ipex_ops import ipex_ops as ops
11 -from vllm.model_executor.layers.fused_moe import (
12 - FusedMoERouter,
13 -)
11 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12 from vllm.model_executor.layers.linear import (
13 LinearBase,
@@ -384,10 +381,13 @@ class XPUFp8MoEMethod(Fp8OnlineMoEMethod):
381 ) -> FusedMoEQuantConfig | None:
382 return None
383
387 - def apply(
384 + @property
385 + def is_monolithic(self) -> bool:
386 + return True
387 +
388 + def apply_monolithic(
389 self,
390 layer: torch.nn.Module,
390 - router: FusedMoERouter,
391 x: torch.Tensor,
392 router_logits: torch.Tensor,
393 ) -> torch.Tensor:
vllm/model_executor/layers/quantization/modelopt.py
+64 -53
@@ -13,7 +13,6 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
14 from vllm.attention.layer import Attention
15 from vllm.logger import init_logger
16 -from vllm.model_executor.layers.fused_moe import FusedMoERouter
16 from vllm.model_executor.layers.fused_moe.config import (
17 FusedMoEConfig,
18 FusedMoEQuantConfig,
@@ -945,42 +944,49 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
944 a2_scale=a2_scale,
945 )
946
948 - def apply(
947 + @property
948 + def is_monolithic(self) -> bool:
949 + return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
950 +
951 + def apply_monolithic(
952 self,
953 layer: FusedMoE,
951 - router: FusedMoERouter,
954 x: torch.Tensor,
955 router_logits: torch.Tensor,
956 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
955 - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
956 - if layer.enable_eplb:
957 - raise NotImplementedError(
958 - "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend."
959 - )
960 - # TODO(rob): this validation should happen at kernel selection
961 - # time in the oracle rather than here.
962 - assert layer.activation == "silu", (
963 - f"Expected 'silu' activation but got {layer.activation}"
964 - )
965 - assert not layer.renormalize
966 - return apply_fi_trtllm_fp8_per_tensor_moe(
967 - layer=layer,
968 - hidden_states=x,
969 - router_logits=router_logits,
970 - routing_bias=layer.e_score_correction_bias,
971 - global_num_experts=layer.global_num_experts,
972 - top_k=layer.top_k,
973 - num_expert_group=layer.num_expert_group,
974 - topk_group=layer.topk_group,
975 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
957 + assert self.is_monolithic
958 + assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
959 + if layer.enable_eplb:
960 + raise NotImplementedError(
961 + "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend."
962 )
977 -
978 - # Expert selection
979 - topk_weights, topk_ids = router.select_experts(
963 + # TODO(rob): this validation should happen at kernel selection
964 + # time in the oracle rather than here.
965 + assert layer.activation == "silu", (
966 + f"Expected 'silu' activation but got {layer.activation}"
967 + )
968 + assert not layer.renormalize
969 + return apply_fi_trtllm_fp8_per_tensor_moe(
970 + layer=layer,
971 hidden_states=x,
972 router_logits=router_logits,
973 + routing_bias=layer.e_score_correction_bias,
974 + global_num_experts=layer.global_num_experts,
975 + top_k=layer.top_k,
976 + num_expert_group=layer.num_expert_group,
977 + topk_group=layer.topk_group,
978 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
979 )
980
981 + def apply(
982 + self,
983 + layer: FusedMoE,
984 + x: torch.Tensor,
985 + topk_weights: torch.Tensor,
986 + topk_ids: torch.Tensor,
987 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
988 + assert not self.is_monolithic
989 +
990 # TODO(rob): this validation should happen at kernel selection
991 # time in the oracle rather than here.
992 if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
@@ -990,7 +996,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
996 )
997
998 assert self.kernel is not None
993 - result = self.kernel(
999 + return self.kernel(
1000 x,
1001 layer.w13_weight,
1002 layer.w2_weight,
@@ -1003,8 +1009,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
1009 apply_router_weight_on_input=layer.apply_router_weight_on_input,
1010 )
1011
1006 - return result
1007 -
1012
1013 ModelOptFp8Config.LinearMethodCls = ModelOptFp8LinearMethod
1014 ModelOptFp8Config.FusedMoEMethodCls = ModelOptFp8MoEMethod
@@ -1629,40 +1633,47 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1633 def supports_eplb(self) -> bool:
1634 return True
1635
1632 - def apply(
1636 + @property
1637 + def is_monolithic(self) -> bool:
1638 + return (
1639 + self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM
1640 + and not self.moe.moe_parallel_config.enable_eplb
1641 + )
1642 +
1643 + def apply_monolithic(
1644 self,
1645 layer: FusedMoE,
1635 - router: FusedMoERouter,
1646 x: torch.Tensor,
1647 router_logits: torch.Tensor,
1648 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1639 - if (
1649 + assert self.is_monolithic
1650 + assert (
1651 self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM
1652 and not layer.enable_eplb
1642 - ):
1643 - return flashinfer_trtllm_fp4_moe(
1644 - layer=layer,
1645 - x=x,
1646 - router_logits=router_logits,
1647 - top_k=layer.top_k,
1648 - activation=layer.activation,
1649 - global_num_experts=layer.global_num_experts,
1650 - num_expert_group=layer.num_expert_group,
1651 - topk_group=layer.topk_group,
1652 - custom_routing_function=layer.custom_routing_function,
1653 - e_score_correction_bias=layer.e_score_correction_bias,
1654 - )
1653 + )
1654
1656 - # Hidden_states in select_experts is only used to extract metadata
1657 - if isinstance(x, tuple):
1658 - x_routing, _ = x
1659 - else:
1660 - x_routing = x
1661 - topk_weights, topk_ids = router.select_experts(
1662 - hidden_states=x_routing,
1655 + return flashinfer_trtllm_fp4_moe(
1656 + layer=layer,
1657 + x=x,
1658 router_logits=router_logits,
1659 + top_k=layer.top_k,
1660 + activation=layer.activation,
1661 + global_num_experts=layer.global_num_experts,
1662 + num_expert_group=layer.num_expert_group,
1663 + topk_group=layer.topk_group,
1664 + custom_routing_function=layer.custom_routing_function,
1665 + e_score_correction_bias=layer.e_score_correction_bias,
1666 )
1667
1668 + def apply(
1669 + self,
1670 + layer: FusedMoE,
1671 + x: torch.Tensor,
1672 + topk_weights: torch.Tensor,
1673 + topk_ids: torch.Tensor,
1674 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1675 + assert not self.is_monolithic
1676 +
1677 # EPLB path
1678 if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:
1679 assert layer.enable_eplb
vllm/model_executor/layers/quantization/moe_wna16.py
+2 -7
@@ -6,7 +6,6 @@ from typing import Any, Optional
6 import torch
7
8 from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
9 -from vllm.model_executor.layers.fused_moe import FusedMoERouter
9 from vllm.model_executor.layers.fused_moe.config import (
10 FusedMoEQuantConfig,
11 int4_w4a16_moe_quant_config,
@@ -365,17 +364,13 @@ class MoeWNA16Method(FusedMoEMethodBase):
364 def apply(
365 self,
366 layer: FusedMoE,
368 - router: FusedMoERouter,
367 x: torch.Tensor,
370 - router_logits: torch.Tensor,
368 + topk_weights: torch.Tensor,
369 + topk_ids: torch.Tensor,
370 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
371 from vllm.model_executor.layers.fused_moe import fused_experts
372
373 assert layer.activation == "silu", "Only SiLU activation is supported."
375 - topk_weights, topk_ids = router.select_experts(
376 - hidden_states=x,
377 - router_logits=router_logits,
378 - )
374
375 return fused_experts(
376 x,
vllm/model_executor/layers/quantization/mxfp4.py
+108 -80
@@ -14,7 +14,6 @@ from vllm.model_executor.layers.fused_moe import (
14 FusedMoE,
15 FusedMoEConfig,
16 FusedMoEMethodBase,
17 - FusedMoERouter,
17 )
18 from vllm.model_executor.layers.fused_moe import modular_kernel as mk
19 from vllm.model_executor.layers.fused_moe.config import (
@@ -890,22 +889,26 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
889 def allow_inplace(self) -> bool:
890 return True
891
892 + @property
893 + def is_monolithic(self) -> bool:
894 + return (
895 + self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
896 + or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
897 + or self.mxfp4_backend == Mxfp4Backend.TRITON
898 + )
899 +
900 def apply(
901 self,
902 layer: FusedMoE,
896 - router: FusedMoERouter,
903 x: torch.Tensor,
898 - router_logits: torch.Tensor,
904 + topk_weights: torch.Tensor,
905 + topk_ids: torch.Tensor,
906 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
907 + assert not self.is_monolithic
908 if layer.enable_eplb:
909 raise NotImplementedError("EPLB is not supported for mxfp4")
910
911 if self.mxfp4_backend == Mxfp4Backend.MARLIN:
904 - topk_weights, topk_ids = router.select_experts(
905 - hidden_states=x,
906 - router_logits=router_logits,
907 - )
908 -
912 return fused_marlin_moe(
913 x,
914 layer.w13_weight,
@@ -914,7 +917,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
917 layer.w2_bias,
918 layer.w13_weight_scale,
919 layer.w2_weight_scale,
917 - router_logits,
920 topk_weights,
921 topk_ids,
922 global_scale1=None,
@@ -942,6 +944,98 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
944 layer.eplb_state.logical_replica_count,
945 ), "MXFP4 are not supported with this configuration."
946
947 + assert (
948 + self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
949 + or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
950 + )
951 + from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
952 +
953 + # Backend-specific preparation
954 + if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS:
955 + from flashinfer import mxfp8_quantize
956 +
957 + x_quant, x_scale = mxfp8_quantize(x, True, 32)
958 +
959 + fake_input_scale = torch.ones(self.num_experts, device=x.device)
960 + quant_scales = [
961 + layer.w13_weight_scale.contiguous().view(torch.int32),
962 + fake_input_scale,
963 + layer.w2_weight_scale.contiguous().view(torch.int32),
964 + fake_input_scale,
965 + ]
966 +
967 + fi_input = x_quant
968 + extra_kwargs = dict(
969 + use_mxfp8_act_scaling=True,
970 + input_sf=x_scale,
971 + fc1_expert_weights=layer.w13_weight.contiguous().view(torch.long),
972 + fc2_expert_weights=layer.w2_weight.contiguous().view(torch.long),
973 + )
974 + elif self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16:
975 + assert x.dtype == torch.bfloat16
976 +
977 + quant_scales = [
978 + layer.w13_weight_scale,
979 + layer.w2_weight_scale,
980 + ]
981 +
982 + fi_input = x
983 + extra_kwargs = dict(
984 + use_w4_group_scaling=True,
985 + fc1_expert_weights=layer.w13_weight,
986 + fc2_expert_weights=layer.w2_weight,
987 + )
988 +
989 + output = torch.empty_like(x, dtype=torch.bfloat16)
990 +
991 + flashinfer_cutlass_fused_moe(
992 + input=fi_input,
993 + token_selected_experts=topk_ids.to(torch.int).contiguous(),
994 + token_final_scales=topk_weights,
995 + output_dtype=torch.bfloat16,
996 + output=output,
997 + quant_scales=quant_scales,
998 + fc1_expert_biases=layer.w13_bias,
999 + fc2_expert_biases=layer.w2_bias,
1000 + swiglu_alpha=layer.gemm1_alpha,
1001 + swiglu_beta=layer.gemm1_beta,
1002 + swiglu_limit=layer.gemm1_clamp_limit,
1003 + tp_size=self.moe.tp_size,
1004 + tp_rank=self.moe.tp_rank,
1005 + ep_size=self.moe.ep_size,
1006 + ep_rank=self.moe.ep_rank,
1007 + tune_max_num_tokens=max(self.max_capture_size, 1),
1008 + **extra_kwargs,
1009 + )
1010 +
1011 + return output
1012 +
1013 + def apply_monolithic(
1014 + self,
1015 + layer: FusedMoE,
1016 + x: torch.Tensor,
1017 + router_logits: torch.Tensor,
1018 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1019 + assert self.is_monolithic
1020 +
1021 + if layer.enable_eplb:
1022 + raise NotImplementedError("EPLB is not supported for mxfp4")
1023 +
1024 + assert _can_support_mxfp4(
1025 + layer.use_grouped_topk,
1026 + layer.topk_group,
1027 + layer.num_expert_group,
1028 + layer.expert_map,
1029 + layer.custom_routing_function,
1030 + layer.e_score_correction_bias,
1031 + layer.apply_router_weight_on_input,
1032 + layer.scoring_func,
1033 + layer.activation,
1034 + layer.eplb_state.expert_load_view,
1035 + layer.eplb_state.logical_to_physical_map,
1036 + layer.eplb_state.logical_replica_count,
1037 + ), "MXFP4 are not supported with this configuration."
1038 +
1039 if (
1040 self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
1041 or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
@@ -988,75 +1082,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1082 tune_max_num_tokens=max(self.max_capture_size, 1),
1083 )[0]
1084 return trtllm_gen_output
991 - elif (
992 - self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
993 - or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
994 - ):
995 - from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
996 -
997 - topk_weights, topk_ids = router.select_experts(
998 - hidden_states=x,
999 - router_logits=router_logits,
1000 - )
1001 -
1002 - # Backend-specific preparation
1003 - if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS:
1004 - from flashinfer import mxfp8_quantize
1005 -
1006 - x_quant, x_scale = mxfp8_quantize(x, True, 32)
1007 -
1008 - fake_input_scale = torch.ones(self.num_experts, device=x.device)
1009 - quant_scales = [
1010 - layer.w13_weight_scale.contiguous().view(torch.int32),
1011 - fake_input_scale,
1012 - layer.w2_weight_scale.contiguous().view(torch.int32),
1013 - fake_input_scale,
1014 - ]
1015 -
1016 - fi_input = x_quant
1017 - extra_kwargs = dict(
1018 - use_mxfp8_act_scaling=True,
1019 - input_sf=x_scale,
1020 - fc1_expert_weights=layer.w13_weight.contiguous().view(torch.long),
1021 - fc2_expert_weights=layer.w2_weight.contiguous().view(torch.long),
1022 - )
1023 - elif self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16:
1024 - assert x.dtype == torch.bfloat16
1025 -
1026 - quant_scales = [
1027 - layer.w13_weight_scale,
1028 - layer.w2_weight_scale,
1029 - ]
1030 -
1031 - fi_input = x
1032 - extra_kwargs = dict(
1033 - use_w4_group_scaling=True,
1034 - fc1_expert_weights=layer.w13_weight,
1035 - fc2_expert_weights=layer.w2_weight,
1036 - )
1037 -
1038 - output = torch.empty_like(x, dtype=torch.bfloat16)
1039 - _ = flashinfer_cutlass_fused_moe(
1040 - input=fi_input,
1041 - token_selected_experts=topk_ids.to(torch.int).contiguous(),
1042 - token_final_scales=topk_weights,
1043 - output_dtype=torch.bfloat16,
1044 - output=output,
1045 - quant_scales=quant_scales,
1046 - fc1_expert_biases=layer.w13_bias,
1047 - fc2_expert_biases=layer.w2_bias,
1048 - swiglu_alpha=layer.gemm1_alpha,
1049 - swiglu_beta=layer.gemm1_beta,
1050 - swiglu_limit=layer.gemm1_clamp_limit,
1051 - tp_size=self.moe.tp_size,
1052 - tp_rank=self.moe.tp_rank,
1053 - ep_size=self.moe.ep_size,
1054 - ep_rank=self.moe.ep_rank,
1055 - tune_max_num_tokens=max(self.max_capture_size, 1),
1056 - **extra_kwargs,
1057 - )
1058 -
1059 - return output
1085 elif self.mxfp4_backend == Mxfp4Backend.TRITON:
1086 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501
1087 triton_kernel_moe_forward,
@@ -1119,10 +1144,13 @@ class IpexMxfp4MoEMethod(Mxfp4MoEMethod):
1144 experts_start_id=ep_rank_start,
1145 )
1146
1122 - def apply(
1147 + @property
1148 + def is_monolithic(self) -> bool:
1149 + return True
1150 +
1151 + def apply_monolithic(
1152 self,
1153 layer: FusedMoE,
1125 - router: FusedMoERouter,
1154 x: torch.Tensor,
1155 router_logits: torch.Tensor,
1156 ) -> torch.Tensor:
vllm/model_executor/layers/quantization/quark/quark_moe.py
+6 -23
@@ -13,7 +13,6 @@ from vllm.model_executor.layers.fused_moe import (
13 FusedMoE,
14 FusedMoEConfig,
15 FusedMoEMethodBase,
16 - FusedMoERouter,
16 FusedMoeWeightScaleSupported,
17 )
18 from vllm.model_executor.layers.fused_moe.config import (
@@ -351,15 +350,10 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
350 def apply(
351 self,
352 layer: FusedMoE,
354 - router: FusedMoERouter,
353 x: torch.Tensor,
356 - router_logits: torch.Tensor,
354 + topk_weights: torch.Tensor,
355 + topk_ids: torch.Tensor,
356 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
358 - topk_weights, topk_ids = router.select_experts(
359 - hidden_states=x,
360 - router_logits=router_logits,
361 - )
362 -
357 if self.rocm_aiter_moe_enabled:
358 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
359 rocm_aiter_fused_experts,
@@ -388,7 +382,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
382 None,
383 layer.w13_weight_scale,
384 layer.w2_weight_scale,
391 - router_logits,
385 topk_weights,
386 topk_ids,
387 quant_type_id=scalar_types.float8_e4m3fn.id,
@@ -544,15 +537,10 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
537 def apply(
538 self,
539 layer: FusedMoE,
547 - router: FusedMoERouter,
540 x: torch.Tensor,
549 - router_logits: torch.Tensor,
541 + topk_weights: torch.Tensor,
542 + topk_ids: torch.Tensor,
543 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
551 - topk_weights, topk_ids = router.select_experts(
552 - hidden_states=x,
553 - router_logits=router_logits,
554 - )
555 -
544 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
545 rocm_aiter_fused_experts,
546 )
@@ -753,15 +741,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
741 def apply(
742 self,
743 layer: FusedMoE,
756 - router: FusedMoERouter,
744 x: torch.Tensor,
758 - router_logits: torch.Tensor,
745 + topk_weights: torch.Tensor,
746 + topk_ids: torch.Tensor,
747 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
760 - topk_weights, topk_ids = router.select_experts(
761 - hidden_states=x,
762 - router_logits=router_logits,
763 - )
764 -
748 if not self.emulate:
749 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
750 rocm_aiter_fused_experts,