@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e74698c27a

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 10, 11:30 UTC

[Misc][Refactor] Add FusedMoERouter object (#30519)

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

bnellnm committed Jan 8, 2026 at 15:52 UTC e74698c27ad5a66bb0f1b131a9e5582a0fbf00d1

20 files changed +165 -36

tests/test_routing_simulator.py

+1 -1

@@ -127,7 +127,7 @@ def test_routing_strategy_integration(monkeypatch, device):

127 envs.environment_variables[env_name] = lambda s=strategy: s

128

129 # Test the select_experts method

130 - topk_weights, topk_ids = fused_moe.select_experts(

130 + topk_weights, topk_ids = fused_moe.router.select_experts(

131 hidden_states=hidden_states,

132 router_logits=router_logits,

133 )

vllm/model_executor/layers/fused_moe/__init__.py

@@ -11,6 +11,9 @@ from vllm.model_executor.layers.fused_moe.config import (

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

12 FusedMoEMethodBase,

13 )

14 +from vllm.model_executor.layers.fused_moe.fused_moe_router import (

15 + FusedMoERouter,

16 +)

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

18 FusedMoE,

19 FusedMoeWeightScaleSupported,

@@ -48,6 +51,7 @@ def get_config() -> dict[str, Any] | None:

52 __all__ = [

53 "FusedMoE",

54 + "FusedMoERouter",

55 "FusedMoEConfig",

56 "FusedMoEMethodBase",

57 "UnquantizedFusedMoEMethod",

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

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

10 FusedMoEConfig,

11 FusedMoEQuantConfig,

12 )

13 +from vllm.model_executor.layers.fused_moe.fused_moe_router import (

14 + FusedMoERouter,

15 +)

16 from vllm.model_executor.layers.fused_moe.modular_kernel import (

17 FusedMoEPermuteExpertsUnpermute,

18 FusedMoEPrepareAndFinalize,

@@ -109,6 +112,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):

112 def apply(

113 self,

114 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

115 + router: FusedMoERouter,

116 x: torch.Tensor,

117 router_logits: torch.Tensor,

118 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py

+3 -1

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

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

13 FusedMoEMethodBase,

14 )

15 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

16 from vllm.model_executor.layers.fused_moe.modular_kernel import (

17 FusedMoEModularKernel,

18 FusedMoEPrepareAndFinalize,

@@ -88,10 +89,11 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

89 def apply(

90 self,

91 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

92 + router: FusedMoERouter,

93 x: torch.Tensor,

94 router_logits: torch.Tensor,

95 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

94 - topk_weights, topk_ids = layer.select_experts(

96 + topk_weights, topk_ids = router.select_experts(

97 hidden_states=x,

98 router_logits=router_logits,

99 )

vllm/model_executor/layers/fused_moe/fused_moe_router.py new

+40

@@ -0,0 +1,40 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from abc import ABC, abstractmethod

4 +

5 +import torch

6 +

7 +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType

8 +

9 +

10 +class FusedMoERouter(ABC):

11 + """

12 + FusedMoERouter is an abstract class that provides a 'select_experts'

13 + method that is used for routing hidden states based on router logits.

14 + """

15 +

16 + @property

17 + @abstractmethod

18 + def routing_method_type(self) -> RoutingMethodType:

19 + raise NotImplementedError

20 +

21 + @abstractmethod

22 + def select_experts(

23 + self,

24 + hidden_states: torch.Tensor,

25 + router_logits: torch.Tensor,

26 + ) -> tuple[torch.Tensor, torch.Tensor]:

27 + """

28 + Route the input hidden states to the top-k experts based on the

29 + router logits.

30 +

31 + Returns:

32 + (topk_weights, topk_ids)

33 + (tuple[torch.Tensor, torch.Tensor]):

34 + The weights and expert ids computation result.

35 +

36 + **Compatibility**: When EPLB is not enabled, the returned ids are

37 + equivalent to global logical ids, so should be compatible with

38 + plain MoE implementations without redundant experts.

39 + """

40 + raise NotImplementedError

vllm/model_executor/layers/fused_moe/layer.py

+25 -3

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

31 FusedMoEQuantConfig,

32 RoutingMethodType,

33 )

34 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

35 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

36 init_aiter_topK_meta_data,

37 )

@@ -284,6 +285,23 @@ def maybe_roundup_hidden_size(

285 return hidden_size

286

287

288 +class FusedMoERouterImpl(FusedMoERouter):

289 + def __init__(self, layer: "FusedMoE"):

290 + super().__init__()

291 + self.layer = layer

292 +

293 + @property

294 + def routing_method_type(self) -> RoutingMethodType:

295 + return self.layer.routing_method_type

296 +

297 + def select_experts(

298 + self,

299 + hidden_states: torch.Tensor,

300 + router_logits: torch.Tensor,

301 + ) -> tuple[torch.Tensor, torch.Tensor]:

302 + return self.layer._select_experts(hidden_states, router_logits)

303 +

304 +

305 @CustomOp.register("fused_moe")

306 class FusedMoE(CustomOp):

307 """FusedMoE layer for MoE models.

@@ -339,7 +357,7 @@ class FusedMoE(CustomOp):

357 is_sequence_parallel=False,

358 expert_mapping: list[tuple[str, str, int, str]] | None = None,

359 n_shared_experts: int | None = None,

342 - routing_method_type: int | None = None,

360 + routing_method_type: RoutingMethodType | None = None,

361 router_logits_dtype: torch.dtype | None = None,

362 ):

363 super().__init__()

@@ -529,7 +547,7 @@ class FusedMoE(CustomOp):

547

548 # ToDo: Better logic to determine the routing method type

549 if routing_method_type is not None:

532 - self.routing_method_type = routing_method_type

550 + self.routing_method_type: RoutingMethodType = routing_method_type

551 else:

552 if scoring_func == "sigmoid":

553 if self.use_grouped_topk:

@@ -640,6 +658,8 @@ class FusedMoE(CustomOp):

658 self.batched_hidden_states: torch.Tensor | None = None

659 self.batched_router_logits: torch.Tensor | None = None

660

661 + self.router = FusedMoERouterImpl(self)

662 +

663 # Note: maybe_init_modular_kernel should only be called by

664 # prepare_communication_buffer_for_model.

665 # This is called after all weight loading and post-processing, so it

@@ -1503,7 +1523,7 @@ class FusedMoE(CustomOp):

1523 device=torch.cuda.current_device(),

1524 )

1525

1506 - def select_experts(

1526 + def _select_experts(

1527 self,

1528 hidden_states: torch.Tensor,

1529 router_logits: torch.Tensor,

@@ -1772,6 +1792,7 @@ class FusedMoE(CustomOp):

1792 # Matrix multiply.

1793 final_hidden_states = self.quant_method.apply(

1794 layer=self,

1795 + router=self.router,

1796 x=staged_hidden_states,

1797 router_logits=staged_router_logits,

1798 )

@@ -1944,6 +1965,7 @@ class FusedMoE(CustomOp):

1965 # Matrix multiply.

1966 final_hidden_states = self.quant_method.apply(

1967 layer=self,

1968 + router=self.router,

1969 x=hidden_states_combined

1970 if do_naive_dispatch_combine

1971 else hidden_states,

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+7 -1

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

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

23 FusedMoEMethodBase,

24 )

25 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

26 from vllm.model_executor.layers.fused_moe.modular_kernel import (

27 FusedMoEActivationFormat,

28 FusedMoEPermuteExpertsUnpermute,

@@ -285,10 +286,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

286 def apply(

287 self,

288 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

289 + router: FusedMoERouter,

290 x: torch.Tensor,

291 router_logits: torch.Tensor,

292 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

293 return self.forward(

294 + router=router,

295 layer=layer,

296 x=x,

297 router_logits=router_logits,

@@ -306,10 +309,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

309 def forward_cuda(

310 self,

311 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

312 + router: FusedMoERouter,

313 x: torch.Tensor,

314 router_logits: torch.Tensor,

315 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

312 - topk_weights, topk_ids = layer.select_experts(

316 + topk_weights, topk_ids = router.select_experts(

317 hidden_states=x,

318 router_logits=router_logits,

319 )

@@ -332,6 +336,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

336 def forward_cpu(

337 self,

338 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

339 + router: FusedMoERouter,

340 x: torch.Tensor,

341 router_logits: torch.Tensor,

342 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -365,6 +370,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

370 def forward_xpu(

371 self,

372 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

373 + router: FusedMoERouter,

374 x: torch.Tensor,

375 router_logits: torch.Tensor,

376 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

vllm/model_executor/layers/quantization/awq_marlin.py

+3 -1

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

15 FusedMoEQuantConfig,

16 )

17 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

18 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

19 from vllm.model_executor.layers.fused_moe.layer import (

20 FusedMoE,

21 FusedMoEMethodBase,

@@ -759,12 +760,13 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

760 def apply(

761 self,

762 layer: FusedMoE,

763 + router: FusedMoERouter,

764 x: torch.Tensor,

765 router_logits: torch.Tensor,

766 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

767 assert layer.activation == "silu", "Only SiLU activation is supported."

768

767 - topk_weights, topk_ids = layer.select_experts(

769 + topk_weights, topk_ids = router.select_experts(

770 hidden_states=x,

771 router_logits=router_logits,

772 )

vllm/model_executor/layers/quantization/bitsandbytes.py

+7 -2

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

10 FusedMoEConfig,

11 FusedMoEQuantConfig,

12 )

13 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEMethodBase

13 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

14 +from vllm.model_executor.layers.fused_moe.layer import (

15 + FusedMoE,

16 + FusedMoEMethodBase,

17 +)

18 from vllm.model_executor.layers.linear import (

19 LinearBase,

20 LinearMethodBase,

@@ -495,12 +499,13 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

499 def apply(

500 self,

501 layer: FusedMoE,

502 + router: FusedMoERouter,

503 x: torch.Tensor,

504 router_logits: torch.Tensor,

505 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

506 from vllm.model_executor.layers.fused_moe import fused_experts

507

503 - topk_weights, topk_ids = layer.select_experts(

508 + topk_weights, topk_ids = router.select_experts(

509 hidden_states=x,

510 router_logits=router_logits,

511 )

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

+14 -6

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

40 MarlinExperts,

41 fused_marlin_moe,

42 )

43 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

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

45 Fp8MoeBackend,

46 convert_to_fp8_moe_kernel_format,

@@ -458,6 +459,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

459 def apply(

460 self,

461 layer: FusedMoE,

462 + router: FusedMoERouter,

463 x: torch.Tensor,

464 router_logits: torch.Tensor,

465 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -484,7 +486,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

486 x_routing, _ = x

487 else:

488 x_routing = x

487 - topk_weights, topk_ids = layer.select_experts(

489 + topk_weights, topk_ids = router.select_experts(

490 hidden_states=x_routing,

491 router_logits=router_logits,

492 )

@@ -926,10 +928,11 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

928 def apply(

929 self,

930 layer: FusedMoE,

931 + router: FusedMoERouter,

932 x: torch.Tensor,

933 router_logits: torch.Tensor,

934 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

932 - topk_weights, topk_ids = layer.select_experts(

935 + topk_weights, topk_ids = router.select_experts(

936 hidden_states=x,

937 router_logits=router_logits,

938 )

@@ -1066,12 +1069,13 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1069 def apply(

1070 self,

1071 layer: FusedMoE,

1072 + router: FusedMoERouter,

1073 x: torch.Tensor,

1074 router_logits: torch.Tensor,

1075 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1076 from vllm.model_executor.layers.fused_moe import fused_experts

1077

1074 - topk_weights, topk_ids = layer.select_experts(

1078 + topk_weights, topk_ids = router.select_experts(

1079 hidden_states=x,

1080 router_logits=router_logits,

1081 )

@@ -1426,6 +1430,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1430 def apply(

1431 self,

1432 layer: FusedMoE,

1433 + router: FusedMoERouter,

1434 x: torch.Tensor,

1435 router_logits: torch.Tensor,

1436 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -1433,7 +1438,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1438 f"{layer.activation} not supported for Marlin MoE."

1439 )

1440

1436 - topk_weights, topk_ids = layer.select_experts(

1441 + topk_weights, topk_ids = router.select_experts(

1442 hidden_states=x,

1443 router_logits=router_logits,

1444 )

@@ -1677,12 +1682,13 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1682 def apply(

1683 self,

1684 layer: FusedMoE,

1685 + router: FusedMoERouter,

1686 x: torch.Tensor,

1687 router_logits: torch.Tensor,

1688 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1689 from vllm.model_executor.layers.fused_moe import fused_experts

1690

1685 - topk_weights, topk_ids = layer.select_experts(

1691 + topk_weights, topk_ids = router.select_experts(

1692 hidden_states=x,

1693 router_logits=router_logits,

1694 )

@@ -1978,6 +1984,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):

1984 def apply(

1985 self,

1986 layer: FusedMoE,

1987 + router: FusedMoERouter,

1988 x: torch.Tensor,

1989 router_logits: torch.Tensor,

1990 ) -> torch.Tensor:

@@ -2290,6 +2297,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2297 def apply(

2298 self,

2299 layer: FusedMoE,

2300 + router: FusedMoERouter,

2301 x: torch.Tensor,

2302 router_logits: torch.Tensor,

2303 ):

@@ -2298,7 +2306,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2306 "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."

2307 )

2308 assert self.moe_quant_config is not None

2301 - topk_weights, topk_ids = layer.select_experts(

2309 + topk_weights, topk_ids = router.select_experts(

2310 hidden_states=x,

2311 router_logits=router_logits,

2312 )

vllm/model_executor/layers/quantization/experts_int8.py

+3 -1

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

15 FusedMoEQuantConfig,

16 int8_w8a16_moe_quant_config,

17 )

18 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

19 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod

20 from vllm.model_executor.layers.quantization import QuantizationMethods

21 from vllm.model_executor.layers.quantization.base_config import (

@@ -137,12 +138,13 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

138 def apply(

139 self,

140 layer: FusedMoE,

141 + router: FusedMoERouter,

142 x: torch.Tensor,

143 router_logits: torch.Tensor,

144 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

145 from vllm.model_executor.layers.fused_moe import fused_experts

146

145 - topk_weights, topk_ids = layer.select_experts(

147 + topk_weights, topk_ids = router.select_experts(

148 hidden_states=x,

149 router_logits=router_logits,

150 )

vllm/model_executor/layers/quantization/fp8.py

+3 -1

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

29 FusedMoEQuantConfig,

30 RoutingMethodType,

31 )

32 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

33 from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod

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

35 Fp8MoeBackend,

@@ -997,6 +998,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

998 def apply(

999 self,

1000 layer: FusedMoE,

1001 + router: FusedMoERouter,

1002 x: torch.Tensor,

1003 router_logits: torch.Tensor,

1004 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -1051,7 +1053,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1053 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1054 )

1055

1054 - topk_weights, topk_ids = layer.select_experts(

1056 + topk_weights, topk_ids = router.select_experts(

1057 hidden_states=x,

1058 router_logits=router_logits,

1059 )

vllm/model_executor/layers/quantization/gguf.py

+7 -2

@@ -16,7 +16,11 @@ from vllm.model_executor.layers.fused_moe.config import (

16 FusedMoEConfig,

17 FusedMoEQuantConfig,

18 )

19 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEMethodBase

19 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

20 +from vllm.model_executor.layers.fused_moe.layer import (

21 + FusedMoE,

22 + FusedMoEMethodBase,

23 +)

24 from vllm.model_executor.layers.linear import (

25 LinearBase,

26 LinearMethodBase,

@@ -629,6 +633,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):

633 def apply(

634 self,

635 layer: FusedMoE,

636 + router: FusedMoERouter,

637 x: torch.Tensor,

638 router_logits: torch.Tensor,

639 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -639,7 +644,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):

644 "fused GGUF MoE method."

645 )

646

642 - topk_weights, topk_ids = layer.select_experts(

647 + topk_weights, topk_ids = router.select_experts(

648 hidden_states=x,

649 router_logits=router_logits,

650 )

vllm/model_executor/layers/quantization/gptq_marlin.py

+3 -1

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

15 FusedMoEQuantConfig,

16 )

17 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

18 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

19 from vllm.model_executor.layers.fused_moe.layer import (

20 FusedMoE,

21 FusedMoEMethodBase,

@@ -895,12 +896,13 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

896 def apply(

897 self,

898 layer: FusedMoE,

899 + router: FusedMoERouter,

900 x: torch.Tensor,

901 router_logits: torch.Tensor,

902 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

903 assert layer.activation == "silu", "Only SiLU activation is supported."

904

903 - topk_weights, topk_ids = layer.select_experts(

905 + topk_weights, topk_ids = router.select_experts(

906 hidden_states=x,

907 router_logits=router_logits,

908 )

vllm/model_executor/layers/quantization/ipex_quant.py

@@ -9,6 +9,9 @@ from torch.nn import Module

10 from vllm._ipex_ops import ipex_ops as ops

11 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

12 +from vllm.model_executor.layers.fused_moe.fused_moe_router import (

13 + FusedMoERouter,

14 +)

15 from vllm.model_executor.layers.linear import (

16 LinearBase,

17 LinearMethodBase,

@@ -384,6 +387,7 @@ class XPUFp8MoEMethod(Fp8OnlineMoEMethod):

387 def apply(

388 self,

389 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

+15 -8

@@ -14,8 +14,10 @@ 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.config import (

17 + FusedMoEConfig,

18 FusedMoEQuantConfig,

19 )

20 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

21 from vllm.model_executor.layers.fused_moe.layer import (

22 FusedMoE,

23 FusedMoEMethodBase,

@@ -200,7 +202,9 @@ class ModelOptQuantConfigBase(QuantizationConfig):

202 quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)

203 return quant_method

204 elif isinstance(layer, FusedMoE):

203 - quant_method = self.FusedMoEMethodCls(quant_config=self, layer=layer)

205 + quant_method = self.FusedMoEMethodCls(

206 + quant_config=self, moe_config=layer.moe_config

207 + )

208 if getattr(quant_method, "backend", "") == "marlin":

209 quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)

210 return quant_method

@@ -720,14 +724,14 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

724 def __init__(

725 self,

726 quant_config: ModelOptFp8Config,

723 - layer: FusedMoE,

727 + moe_config: FusedMoEConfig,

728 ) -> None:

725 - super().__init__(layer.moe_config)

729 + super().__init__(moe_config)

730 self.quant_config = quant_config

731 assert self.quant_config.is_checkpoint_fp8_serialized

732 self.fp8_backend = select_fp8_moe_backend(

733 block_quant=False,

730 - tp_size=layer.moe_parallel_config.tp_size,

734 + tp_size=moe_config.moe_parallel_config.tp_size,

735 with_lora_support=self.moe.is_lora_enabled,

736 )

737 self.kernel: mk.FusedMoEModularKernel | None = None

@@ -935,6 +939,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

939 def apply(

940 self,

941 layer: FusedMoE,

942 + router: FusedMoERouter,

943 x: torch.Tensor,

944 router_logits: torch.Tensor,

945 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -961,7 +966,8 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

966 apply_router_weight_on_input=layer.apply_router_weight_on_input,

967 )

968

964 - topk_weights, topk_ids = layer.select_experts(

969 + # Expert selection

970 + topk_weights, topk_ids = router.select_experts(

971 hidden_states=x,

972 router_logits=router_logits,

973 )

@@ -1325,9 +1331,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1331 def __init__(

1332 self,

1333 quant_config: ModelOptNvFp4Config,

1328 - layer: FusedMoE,

1334 + moe_config: FusedMoEConfig,

1335 ) -> None:

1330 - super().__init__(layer.moe_config)

1336 + super().__init__(moe_config)

1337 self.quant_config = quant_config

1338 self.nvfp4_backend = select_nvfp4_moe_backend()

1339 # TODO: move this type of check into the oracle.

@@ -1597,6 +1603,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1603 def apply(

1604 self,

1605 layer: FusedMoE,

1606 + router: FusedMoERouter,

1607 x: torch.Tensor,

1608 router_logits: torch.Tensor,

1609 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -1621,7 +1628,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1628 x_routing, _ = x

1629 else:

1630 x_routing = x

1624 - topk_weights, topk_ids = layer.select_experts(

1631 + topk_weights, topk_ids = router.select_experts(

1632 hidden_states=x_routing,

1633 router_logits=router_logits,

1634 )

vllm/model_executor/layers/quantization/moe_wna16.py

+3 -1

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

11 int4_w4a16_moe_quant_config,

12 int8_w8a16_moe_quant_config,

13 )

14 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

15 from vllm.model_executor.layers.fused_moe.layer import (

16 FusedMoE,

17 FusedMoEConfig,

@@ -364,13 +365,14 @@ class MoeWNA16Method(FusedMoEMethodBase):

365 def apply(

366 self,

367 layer: FusedMoE,

368 + router: FusedMoERouter,

369 x: torch.Tensor,

370 router_logits: torch.Tensor,

371 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

372 from vllm.model_executor.layers.fused_moe import fused_experts

373

374 assert layer.activation == "silu", "Only SiLU activation is supported."

373 - topk_weights, topk_ids = layer.select_experts(

375 + topk_weights, topk_ids = router.select_experts(

376 hidden_states=x,

377 router_logits=router_logits,

378 )

vllm/model_executor/layers/quantization/mxfp4.py

+6 -3

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

27 MarlinExperts,

28 fused_marlin_moe,

29 )

30 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

31 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (

32 OAITritonExperts,

33 UnfusedOAITritonExperts,

@@ -891,6 +892,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

892 def apply(

893 self,

894 layer: FusedMoE,

895 + router: FusedMoERouter,

896 x: torch.Tensor,

897 router_logits: torch.Tensor,

898 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -898,7 +900,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

900 raise NotImplementedError("EPLB is not supported for mxfp4")

901

902 if self.mxfp4_backend == Mxfp4Backend.MARLIN:

901 - topk_weights, topk_ids = layer.select_experts(

903 + topk_weights, topk_ids = router.select_experts(

904 hidden_states=x,

905 router_logits=router_logits,

906 )

@@ -992,7 +994,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

994 ):

995 from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe

996

995 - topk_weights, topk_ids = layer.select_experts(

997 + topk_weights, topk_ids = router.select_experts(

998 hidden_states=x,

999 router_logits=router_logits,

1000 )

@@ -1119,7 +1121,8 @@ class IpexMxfp4MoEMethod(Mxfp4MoEMethod):

1121

1122 def apply(

1123 self,

1122 - layer: torch.nn.Module,

1124 + layer: FusedMoE,

1125 + router: FusedMoERouter,

1126 x: torch.Tensor,

1127 router_logits: torch.Tensor,

1128 ) -> torch.Tensor:

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

+6 -2

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

13 FusedMoE,

14 FusedMoEConfig,

15 FusedMoEMethodBase,

16 + FusedMoERouter,

17 FusedMoeWeightScaleSupported,

18 )

19 from vllm.model_executor.layers.fused_moe.config import (

@@ -350,10 +351,11 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

351 def apply(

352 self,

353 layer: FusedMoE,

354 + router: FusedMoERouter,

355 x: torch.Tensor,

356 router_logits: torch.Tensor,

357 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

356 - topk_weights, topk_ids = layer.select_experts(

358 + topk_weights, topk_ids = router.select_experts(

359 hidden_states=x,

360 router_logits=router_logits,

361 )

@@ -542,6 +544,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

544 def apply(

545 self,

546 layer: FusedMoE,

547 + router: FusedMoERouter,

548 x: torch.Tensor,

549 router_logits: torch.Tensor,

550 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

@@ -750,10 +753,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

753 def apply(

754 self,

755 layer: FusedMoE,

756 + router: FusedMoERouter,

757 x: torch.Tensor,

758 router_logits: torch.Tensor,

759 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

756 - topk_weights, topk_ids = layer.select_experts(

760 + topk_weights, topk_ids = router.select_experts(

761 hidden_states=x,

762 router_logits=router_logits,

763 )

vllm/model_executor/layers/quantization/rtn.py

+7 -2

@@ -15,7 +15,11 @@ from vllm.model_executor.layers.fused_moe.config import (

15 FusedMoEQuantConfig,

16 )

17 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

18 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEMethodBase

18 +from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter

19 +from vllm.model_executor.layers.fused_moe.layer import (

20 + FusedMoE,

21 + FusedMoEMethodBase,

22 +)

23 from vllm.model_executor.layers.linear import (

24 LinearBase,

25 LinearMethodBase,

@@ -356,10 +360,11 @@ class RTNMoEMethod(FusedMoEMethodBase):

360 def apply(

361 self,

362 layer: FusedMoE,

363 + router: FusedMoERouter,

364 x: torch.Tensor,

365 router_logits: torch.Tensor,

366 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

362 - topk_weights, topk_ids = layer.select_experts(

367 + topk_weights, topk_ids = router.select_experts(

368 hidden_states=x,

369 router_logits=router_logits,

370 )