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 12, 14:45 UTC
[MoE][Refactor] Make select_experts a non-static method (#29067)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Nov 24, 2025 at 13:38 UTC
8f066146c395dfadb86914c88d9a0f3173f8fa39
18 files changed
+163
-472
tests/kernels/moe/test_flashinfer.py
+6
-13
@@ -11,7 +11,6 @@ from vllm.model_executor.layers.fused_moe.config import (
11
fp8_w8a8_moe_quant_config,
12
)
13
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
14
-from vllm.model_executor.layers.fused_moe.layer import FusedMoE
14
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
15
apply_flashinfer_per_tensor_scale_fp8,
16
flashinfer_cutlass_moe_fp8,
@@ -151,14 +150,11 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
150
td = TestData.make_moe_tensors_8bit(m, k, n, e, reorder=True)
151
152
score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
154
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
153
+ topk_weights, topk_ids = Llama4MoE.custom_routing_function(
154
hidden_states=td.hidden_states,
156
- router_logits=score,
157
- use_grouped_topk=False,
158
- top_k=topk,
155
+ gating_output=score,
156
+ topk=topk,
157
renormalize=False,
160
- custom_routing_function=Llama4MoE.custom_routing_function,
161
- scoring_func="softmax",
158
)
159
160
quant_config = fp8_w8a8_moe_quant_config(
@@ -219,14 +215,11 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
215
)
216
217
score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
222
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
218
+ topk_weights, topk_ids = Llama4MoE.custom_routing_function(
219
hidden_states=td.hidden_states,
224
- router_logits=score,
225
- use_grouped_topk=False,
226
- top_k=topk,
220
+ gating_output=score,
221
+ topk=topk,
222
renormalize=False,
228
- custom_routing_function=Llama4MoE.custom_routing_function,
229
- scoring_func="softmax",
223
)
224
225
quant_config = fp8_w8a8_moe_quant_config(
tests/test_routing_simulator.py
+30
-5
@@ -9,9 +9,16 @@ different routing strategies and analyze their performance, including
9
integration tests with FusedMoE layer.
10
"""
11
12
+import tempfile
13
+
14
import pytest
15
import torch
16
17
+from vllm.config import VllmConfig, set_current_vllm_config
18
+from vllm.distributed import (
19
+ init_distributed_environment,
20
+ initialize_model_parallel,
21
+)
22
from vllm.model_executor.layers.fused_moe.routing_simulator import (
23
DistributionBasedRouting,
24
RoutingSimulator,
@@ -89,6 +96,28 @@ def test_routing_strategy_integration(monkeypatch, device):
96
# Test different routing strategies
97
strategies = RoutingSimulator.get_available_strategies()
98
99
+ vllm_config = VllmConfig()
100
+ with set_current_vllm_config(vllm_config):
101
+ temp_file = tempfile.mkstemp()[1]
102
+ init_distributed_environment(
103
+ world_size=1,
104
+ rank=0,
105
+ local_rank=0,
106
+ distributed_init_method=f"file://{temp_file}",
107
+ )
108
+ initialize_model_parallel(
109
+ tensor_model_parallel_size=1,
110
+ pipeline_model_parallel_size=1,
111
+ )
112
+ fused_moe = FusedMoE(
113
+ num_experts=num_experts,
114
+ top_k=top_k,
115
+ hidden_size=hidden_size,
116
+ intermediate_size=0,
117
+ use_grouped_topk=False,
118
+ renormalize=True,
119
+ )
120
+
121
for strategy in strategies:
122
# Set environment variable
123
env_name = "VLLM_MOE_ROUTING_SIMULATION_STRATEGY"
@@ -98,13 +127,9 @@ def test_routing_strategy_integration(monkeypatch, device):
127
envs.environment_variables[env_name] = lambda s=strategy: s
128
129
# Test the select_experts method
101
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
130
+ topk_weights, topk_ids, _ = fused_moe.select_experts(
131
hidden_states=hidden_states,
132
router_logits=router_logits,
104
- top_k=top_k,
105
- use_grouped_topk=False,
106
- renormalize=True,
107
- indices_type=torch.long,
133
)
134
135
# Verify output shapes
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+5
-1
@@ -90,10 +90,14 @@ class FusedMoEMethodBase(QuantizeMethodBase):
90
def allow_inplace(self) -> bool:
91
return False
92
93
+ @property
94
+ def method_name(self) -> str:
95
+ return self.__class__.__name__
96
+
97
@abstractmethod
98
def apply(
99
self,
96
- layer: torch.nn.Module,
100
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
101
x: torch.Tensor,
102
router_logits: torch.Tensor,
103
top_k: int,
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+6
-35
@@ -66,6 +66,10 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
66
def allow_inplace(self) -> bool:
67
return self.old_quant_method.allow_inplace
68
69
+ @property
70
+ def method_name(self) -> str:
71
+ return self.old_quant_method.method_name
72
+
73
def create_weights(
74
self,
75
layer: torch.nn.Module,
@@ -84,7 +88,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
88
89
def apply(
90
self,
87
- layer: torch.nn.Module,
91
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
92
x: torch.Tensor,
93
router_logits: torch.Tensor,
94
top_k: int,
@@ -105,42 +109,9 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
109
logical_to_physical_map: torch.Tensor | None = None,
110
logical_replica_count: torch.Tensor | None = None,
111
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
108
- # Is getattr needed?
109
- zero_expert_num = getattr(layer, "zero_expert_num", 0)
110
- zero_expert_type = getattr(layer, "zero_expert_type", None)
111
-
112
- if enable_eplb:
113
- if self.supports_eplb:
114
- assert expert_load_view is not None
115
- assert logical_to_physical_map is not None
116
- assert logical_replica_count is not None
117
- else:
118
- raise NotImplementedError(
119
- "EPLB is not supported for "
120
- f"{self.old_quant_method.__class__.__name__}."
121
- )
122
-
112
topk_weights, topk_ids, zero_expert_result = layer.select_experts(
113
hidden_states=x,
114
router_logits=router_logits,
126
- use_grouped_topk=use_grouped_topk,
127
- top_k=top_k,
128
- renormalize=renormalize,
129
- topk_group=topk_group,
130
- num_expert_group=num_expert_group,
131
- custom_routing_function=custom_routing_function,
132
- scoring_func=scoring_func,
133
- routed_scaling_factor=routed_scaling_factor,
134
- e_score_correction_bias=e_score_correction_bias,
135
- indices_type=self.topk_indices_dtype,
136
- enable_eplb=enable_eplb,
137
- expert_map=expert_map,
138
- expert_load_view=expert_load_view,
139
- logical_to_physical_map=logical_to_physical_map,
140
- logical_replica_count=logical_replica_count,
141
- global_num_experts=global_num_experts,
142
- zero_expert_num=zero_expert_num,
143
- zero_expert_type=zero_expert_type,
115
)
116
117
result = self.fused_experts(
@@ -156,7 +127,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
127
expert_map=None if self.disable_expert_map else expert_map,
128
)
129
159
- if zero_expert_num != 0 and zero_expert_type is not None:
130
+ if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
131
assert not isinstance(result, tuple), (
132
"Shared + zero experts are mutually exclusive not yet supported"
133
)
vllm/model_executor/layers/fused_moe/layer.py
+58
-60
@@ -1510,30 +1510,11 @@ class FusedMoE(CustomOp):
1510
logits_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
1511
)
1512
1513
- @staticmethod
1513
def select_experts(
1514
+ self,
1515
hidden_states: torch.Tensor,
1516
router_logits: torch.Tensor,
1517
- top_k: int,
1518
- use_grouped_topk: bool,
1519
- renormalize: bool,
1520
- topk_group: int | None = None,
1521
- num_expert_group: int | None = None,
1522
- custom_routing_function: Callable | None = None,
1523
- scoring_func: str = "softmax",
1524
- routed_scaling_factor: float = 1.0,
1525
- e_score_correction_bias: torch.Tensor | None = None,
1526
- indices_type: torch.dtype | None = None,
1527
- enable_eplb: bool = False,
1528
- expert_map: torch.Tensor | None = None,
1529
- expert_load_view: torch.Tensor | None = None,
1530
- logical_to_physical_map: torch.Tensor | None = None,
1531
- logical_replica_count: torch.Tensor | None = None,
1532
- global_num_experts: int | None = None,
1533
- zero_expert_num: int | None = None,
1534
- zero_expert_type: str | None = None,
1535
- num_fused_shared_experts: int = 0,
1536
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1517
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
1518
"""
1519
Route the input hidden states to the top-k experts based on the
1520
router logits.
@@ -1552,6 +1533,27 @@ class FusedMoE(CustomOp):
1533
fused_topk_bias,
1534
)
1535
1536
+ if self.enable_eplb:
1537
+ if self.quant_method.supports_eplb:
1538
+ if self.expert_load_view is None:
1539
+ raise ValueError(
1540
+ "enable_eplb=True requiere expert_load_view != None"
1541
+ )
1542
+ if self.logical_to_physical_map is None:
1543
+ raise ValueError(
1544
+ "enable_eplb=True requiere logical_to_physical_map != None"
1545
+ )
1546
+ if self.logical_replica_count is None:
1547
+ raise ValueError(
1548
+ "enable_eplb=True requiere logical_replica_count != None"
1549
+ )
1550
+ else:
1551
+ raise NotImplementedError(
1552
+ f"EPLB is not supported for {self.quant_method.method_name}."
1553
+ )
1554
+
1555
+ indices_type = self.quant_method.topk_indices_dtype
1556
+
1557
# Check if we should use a routing simulation strategy
1558
routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
1559
if routing_strategy != "":
@@ -1559,20 +1561,20 @@ class FusedMoE(CustomOp):
1561
hidden_states=hidden_states,
1562
router_logits=router_logits,
1563
strategy_name=routing_strategy,
1562
- top_k=top_k,
1564
+ top_k=self.top_k,
1565
indices_type=indices_type,
1566
)
1567
1568
# DeepSeekv2 uses grouped_top_k
1567
- elif use_grouped_topk:
1568
- assert topk_group is not None
1569
- assert num_expert_group is not None
1569
+ elif self.use_grouped_topk:
1570
+ assert self.topk_group is not None
1571
+ assert self.num_expert_group is not None
1572
if rocm_aiter_ops.is_fused_moe_enabled():
1573
if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
1572
- assert num_fused_shared_experts == 0
1574
+ assert self.num_fused_shared_experts == 0
1575
grouped_topk_impl = partial(
1576
rocm_aiter_grouped_topk,
1575
- num_fused_shared_experts=num_fused_shared_experts,
1577
+ num_fused_shared_experts=self.num_fused_shared_experts,
1578
)
1579
else:
1580
grouped_topk_impl = grouped_topk
@@ -1580,50 +1582,46 @@ class FusedMoE(CustomOp):
1582
topk_weights, topk_ids = grouped_topk_impl(
1583
hidden_states=hidden_states,
1584
gating_output=router_logits,
1583
- topk=top_k,
1584
- renormalize=renormalize,
1585
- num_expert_group=num_expert_group,
1586
- topk_group=topk_group,
1587
- scoring_func=scoring_func,
1588
- routed_scaling_factor=routed_scaling_factor,
1589
- e_score_correction_bias=e_score_correction_bias,
1585
+ topk=self.top_k,
1586
+ renormalize=self.renormalize,
1587
+ num_expert_group=self.num_expert_group,
1588
+ topk_group=self.topk_group,
1589
+ scoring_func=self.scoring_func,
1590
+ routed_scaling_factor=self.routed_scaling_factor,
1591
+ e_score_correction_bias=self.e_score_correction_bias,
1592
)
1591
- elif e_score_correction_bias is not None:
1593
+ elif self.e_score_correction_bias is not None:
1594
topk_weights, topk_ids = fused_topk_bias(
1595
hidden_states=hidden_states,
1596
gating_output=router_logits,
1595
- e_score_correction_bias=e_score_correction_bias.data,
1596
- topk=top_k,
1597
- renormalize=renormalize,
1597
+ e_score_correction_bias=self.e_score_correction_bias.data,
1598
+ topk=self.top_k,
1599
+ renormalize=self.renormalize,
1600
)
1599
- if routed_scaling_factor != 1.0:
1600
- topk_weights *= routed_scaling_factor
1601
- elif custom_routing_function is None:
1601
+ if self.routed_scaling_factor != 1.0:
1602
+ topk_weights *= self.routed_scaling_factor
1603
+ elif self.custom_routing_function is None:
1604
topk_weights, topk_ids, token_expert_indices = fused_topk(
1605
hidden_states=hidden_states,
1606
gating_output=router_logits,
1605
- topk=top_k,
1606
- renormalize=renormalize,
1607
+ topk=self.top_k,
1608
+ renormalize=self.renormalize,
1609
indices_type=indices_type,
1610
)
1611
else:
1610
- topk_weights, topk_ids = custom_routing_function(
1612
+ topk_weights, topk_ids = self.custom_routing_function(
1613
hidden_states=hidden_states,
1614
gating_output=router_logits,
1613
- topk=top_k,
1614
- renormalize=renormalize,
1615
+ topk=self.top_k,
1616
+ renormalize=self.renormalize,
1617
)
1618
1617
- if enable_eplb:
1618
- assert expert_load_view is not None
1619
- assert logical_to_physical_map is not None
1620
- assert logical_replica_count is not None
1621
-
1619
+ if self.enable_eplb:
1620
topk_ids = eplb_map_to_physical_and_record(
1621
topk_ids=topk_ids,
1624
- expert_load_view=expert_load_view,
1625
- logical_to_physical_map=logical_to_physical_map,
1626
- logical_replica_count=logical_replica_count,
1622
+ expert_load_view=self.expert_load_view,
1623
+ logical_to_physical_map=self.logical_to_physical_map,
1624
+ logical_replica_count=self.logical_replica_count,
1625
)
1626
1627
if (indices_type is not None) and topk_ids.dtype != indices_type:
@@ -1633,16 +1631,16 @@ class FusedMoE(CustomOp):
1631
1632
# Compute zero expert result if needed
1633
if (
1636
- zero_expert_num is not None
1637
- and zero_expert_num > 0
1638
- and zero_expert_type is not None
1639
- and global_num_experts is not None
1634
+ self.zero_expert_num is not None
1635
+ and self.zero_expert_num > 0
1636
+ and self.zero_expert_type is not None
1637
+ and self.global_num_experts is not None
1638
):
1639
zero_expert_result = zero_experts_compute_triton(
1640
expert_indices=topk_ids,
1641
expert_scales=topk_weights,
1644
- num_experts=global_num_experts,
1645
- zero_expert_type=zero_expert_type,
1642
+ num_experts=self.global_num_experts,
1643
+ zero_expert_type=self.zero_expert_type,
1644
hidden_states=hidden_states,
1645
)
1646
else:
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+5
-27
@@ -331,7 +331,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
331
332
def forward_cuda(
333
self,
334
- layer: torch.nn.Module,
334
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
335
x: torch.Tensor,
336
use_grouped_topk: bool,
337
top_k: int,
@@ -352,31 +352,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
352
logical_to_physical_map: torch.Tensor | None = None,
353
logical_replica_count: torch.Tensor | None = None,
354
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
355
- zero_expert_num = getattr(layer, "zero_expert_num", 0)
356
- zero_expert_type = getattr(layer, "zero_expert_type", None)
357
-
355
topk_weights, topk_ids, zero_expert_result = layer.select_experts(
356
hidden_states=x,
357
router_logits=router_logits,
361
- use_grouped_topk=use_grouped_topk,
362
- top_k=top_k,
363
- renormalize=renormalize,
364
- topk_group=topk_group,
365
- num_expert_group=num_expert_group,
366
- custom_routing_function=custom_routing_function,
367
- scoring_func=scoring_func,
368
- routed_scaling_factor=routed_scaling_factor,
369
- e_score_correction_bias=e_score_correction_bias,
370
- indices_type=self.topk_indices_dtype,
371
- enable_eplb=enable_eplb,
372
- expert_map=expert_map,
373
- expert_load_view=expert_load_view,
374
- logical_to_physical_map=logical_to_physical_map,
375
- logical_replica_count=logical_replica_count,
376
- global_num_experts=global_num_experts,
377
- zero_expert_num=zero_expert_num,
378
- zero_expert_type=zero_expert_type,
379
- num_fused_shared_experts=layer.num_fused_shared_experts,
358
)
359
360
if self.rocm_aiter_moe_enabled:
@@ -415,7 +393,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
393
expert_map=expert_map,
394
)
395
418
- if zero_expert_num != 0 and zero_expert_type is not None:
396
+ if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
397
assert not isinstance(result, tuple), (
398
"Shared + zero experts are mutually exclusive not yet supported"
399
)
@@ -425,7 +403,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
403
404
def forward_cpu(
405
self,
428
- layer: torch.nn.Module,
406
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
407
x: torch.Tensor,
408
use_grouped_topk: bool,
409
top_k: int,
@@ -474,7 +452,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
452
453
def forward_xpu(
454
self,
477
- layer: torch.nn.Module,
455
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
456
x: torch.Tensor,
457
use_grouped_topk: bool,
458
top_k: int,
@@ -515,7 +493,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
493
494
def forward_tpu(
495
self,
518
- layer: torch.nn.Module,
496
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
497
x: torch.Tensor,
498
use_grouped_topk: bool,
499
top_k: int,
vllm/model_executor/layers/quantization/awq_marlin.py
+2
-15
@@ -597,7 +597,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
597
598
def apply(
599
self,
600
- layer: torch.nn.Module,
600
+ layer: FusedMoE,
601
x: torch.Tensor,
602
router_logits: torch.Tensor,
603
top_k: int,
@@ -618,24 +618,11 @@ class AWQMoEMethod(FusedMoEMethodBase):
618
logical_to_physical_map: torch.Tensor | None = None,
619
logical_replica_count: torch.Tensor | None = None,
620
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
621
- if enable_eplb:
622
- raise NotImplementedError("EPLB not supported for `AWQMoEMethod` yet.")
623
-
621
assert activation == "silu", "Only SiLU activation is supported."
622
626
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
623
+ topk_weights, topk_ids, _ = layer.select_experts(
624
hidden_states=x,
625
router_logits=router_logits,
629
- use_grouped_topk=use_grouped_topk,
630
- top_k=top_k,
631
- renormalize=renormalize,
632
- topk_group=topk_group,
633
- num_expert_group=num_expert_group,
634
- custom_routing_function=custom_routing_function,
635
- scoring_func=scoring_func,
636
- routed_scaling_factor=routed_scaling_factor,
637
- e_score_correction_bias=e_score_correction_bias,
638
- indices_type=self.topk_indices_dtype,
626
)
627
628
return fused_marlin_moe(
vllm/model_executor/layers/quantization/bitsandbytes.py
+3
-17
@@ -495,7 +495,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
495
496
def apply(
497
self,
498
- layer: torch.nn.Module,
498
+ layer: FusedMoE,
499
x: torch.Tensor,
500
router_logits: torch.Tensor,
501
top_k: int,
@@ -518,25 +518,11 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
518
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
519
from vllm.model_executor.layers.fused_moe import fused_experts
520
521
- if enable_eplb:
522
- raise NotImplementedError(
523
- "EPLB not supported for `BitsAndBytesMoEMethod` yet."
524
- )
525
-
526
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
521
+ topk_weights, topk_ids, _ = layer.select_experts(
522
hidden_states=x,
523
router_logits=router_logits,
529
- use_grouped_topk=use_grouped_topk,
530
- top_k=top_k,
531
- renormalize=renormalize,
532
- topk_group=topk_group,
533
- num_expert_group=num_expert_group,
534
- custom_routing_function=custom_routing_function,
535
- scoring_func=scoring_func,
536
- routed_scaling_factor=routed_scaling_factor,
537
- e_score_correction_bias=e_score_correction_bias,
538
- indices_type=self.topk_indices_dtype,
524
)
525
+ # TODO(bnell): Do these need to be called on the hot path?
526
if self.quant_config.load_in_8bit:
527
w13, w2 = self._apply_8bit_dequant(layer)
528
else:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+15
-108
@@ -511,7 +511,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
511
512
def apply(
513
self,
514
- layer: torch.nn.Module,
514
+ layer: FusedMoE,
515
x: torch.Tensor,
516
router_logits: torch.Tensor,
517
top_k: int,
@@ -532,16 +532,17 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
532
logical_to_physical_map: torch.Tensor | None = None,
533
logical_replica_count: torch.Tensor | None = None,
534
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
535
- if enable_eplb:
536
- raise NotImplementedError(
537
- "EPLB not supported for `CompressedTensorsW4A4MoeMethod` yet."
538
- )
535
assert activation == "silu", "Only SiLU activation is supported."
536
537
if (
538
self.allow_flashinfer
539
and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
540
):
541
+ if enable_eplb:
542
+ raise NotImplementedError(
543
+ "EPLB not supported for `CompressedTensorsW4A4MoeMethod` yet."
544
+ )
545
+
546
return flashinfer_trtllm_fp4_moe(
547
layer=layer,
548
x=x,
@@ -554,19 +555,9 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
555
e_score_correction_bias=e_score_correction_bias,
556
)
557
557
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
558
+ topk_weights, topk_ids, _ = layer.select_experts(
559
hidden_states=x,
560
router_logits=router_logits,
560
- use_grouped_topk=use_grouped_topk,
561
- top_k=top_k,
562
- renormalize=renormalize,
563
- topk_group=topk_group,
564
- num_expert_group=num_expert_group,
565
- custom_routing_function=custom_routing_function,
566
- scoring_func=scoring_func,
567
- routed_scaling_factor=routed_scaling_factor,
568
- e_score_correction_bias=e_score_correction_bias,
569
- indices_type=self.topk_indices_dtype,
561
)
562
563
if self.use_marlin:
@@ -1109,7 +1100,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1100
1101
def apply(
1102
self,
1112
- layer: torch.nn.Module,
1103
+ layer: FusedMoE,
1104
x: torch.Tensor,
1105
router_logits: torch.Tensor,
1106
top_k: int,
@@ -1130,31 +1121,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1121
logical_to_physical_map: torch.Tensor | None = None,
1122
logical_replica_count: torch.Tensor | None = None,
1123
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1133
- if enable_eplb:
1134
- assert expert_load_view is not None
1135
- assert logical_to_physical_map is not None
1136
- assert logical_replica_count is not None
1137
- assert isinstance(layer, FusedMoE)
1138
-
1139
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
1124
+ topk_weights, topk_ids, _ = layer.select_experts(
1125
hidden_states=x,
1126
router_logits=router_logits,
1142
- use_grouped_topk=use_grouped_topk,
1143
- top_k=top_k,
1144
- renormalize=renormalize,
1145
- topk_group=topk_group,
1146
- num_expert_group=num_expert_group,
1147
- custom_routing_function=custom_routing_function,
1148
- scoring_func=scoring_func,
1149
- routed_scaling_factor=routed_scaling_factor,
1150
- e_score_correction_bias=e_score_correction_bias,
1151
- indices_type=self.topk_indices_dtype,
1152
- num_fused_shared_experts=layer.num_fused_shared_experts,
1153
- enable_eplb=enable_eplb,
1154
- expert_map=expert_map,
1155
- expert_load_view=expert_load_view,
1156
- logical_to_physical_map=logical_to_physical_map,
1157
- logical_replica_count=logical_replica_count,
1127
)
1128
1129
per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
@@ -1377,7 +1346,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1346
1347
def apply(
1348
self,
1380
- layer: torch.nn.Module,
1349
+ layer: FusedMoE,
1350
x: torch.Tensor,
1351
router_logits: torch.Tensor,
1352
top_k: int,
@@ -1398,26 +1367,11 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1367
logical_to_physical_map: torch.Tensor | None = None,
1368
logical_replica_count: torch.Tensor | None = None,
1369
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1401
- if enable_eplb:
1402
- raise NotImplementedError(
1403
- "EPLB not supported for `CompressedTensorsW8A8Int8MoEMethod` yet."
1404
- )
1405
-
1370
from vllm.model_executor.layers.fused_moe import fused_experts
1371
1408
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
1372
+ topk_weights, topk_ids, _ = layer.select_experts(
1373
hidden_states=x,
1374
router_logits=router_logits,
1411
- use_grouped_topk=use_grouped_topk,
1412
- top_k=top_k,
1413
- renormalize=renormalize,
1414
- topk_group=topk_group,
1415
- num_expert_group=num_expert_group,
1416
- custom_routing_function=custom_routing_function,
1417
- scoring_func=scoring_func,
1418
- routed_scaling_factor=routed_scaling_factor,
1419
- e_score_correction_bias=e_score_correction_bias,
1420
- indices_type=self.topk_indices_dtype,
1375
)
1376
1377
return fused_experts(
@@ -1738,7 +1692,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1692
1693
def apply(
1694
self,
1741
- layer: torch.nn.Module,
1695
+ layer: FusedMoE,
1696
x: torch.Tensor,
1697
router_logits: torch.Tensor,
1698
top_k: int,
@@ -1759,26 +1713,11 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1713
logical_to_physical_map: torch.Tensor | None = None,
1714
logical_replica_count: torch.Tensor | None = None,
1715
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1762
- if enable_eplb:
1763
- raise NotImplementedError(
1764
- "EPLB not supported for `CompressedTensorsWNA16MarlinMoEMethod` yet."
1765
- )
1766
-
1716
assert activation == "silu", f"{activation} not supported for Marlin MoE."
1717
1769
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
1718
+ topk_weights, topk_ids, _ = layer.select_experts(
1719
hidden_states=x,
1720
router_logits=router_logits,
1772
- use_grouped_topk=use_grouped_topk,
1773
- top_k=top_k,
1774
- renormalize=renormalize,
1775
- topk_group=topk_group,
1776
- num_expert_group=num_expert_group,
1777
- custom_routing_function=custom_routing_function,
1778
- scoring_func=scoring_func,
1779
- routed_scaling_factor=routed_scaling_factor,
1780
- e_score_correction_bias=e_score_correction_bias,
1781
- indices_type=self.topk_indices_dtype,
1721
)
1722
1723
return fused_marlin_moe(
@@ -2001,7 +1940,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1940
1941
def apply(
1942
self,
2004
- layer: torch.nn.Module,
1943
+ layer: FusedMoE,
1944
x: torch.Tensor,
1945
router_logits: torch.Tensor,
1946
top_k: int,
@@ -2022,43 +1961,11 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1961
logical_to_physical_map: torch.Tensor | None = None,
1962
logical_replica_count: torch.Tensor | None = None,
1963
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
2025
- if enable_eplb:
2026
- if expert_load_view is None:
2027
- raise ValueError("enable_eplb=True requiere expert_load_view != None")
2028
- if logical_to_physical_map is None:
2029
- raise ValueError(
2030
- "enable_eplb=True requiere logical_to_physical_map != None"
2031
- )
2032
- if logical_replica_count is None:
2033
- raise ValueError(
2034
- "enable_eplb=True requiere logical_replica_count != None"
2035
- )
2036
- if not isinstance(layer, FusedMoE):
2037
- raise TypeError(
2038
- "EPLB is only supported when `layer` is a instance of FusedMoE."
2039
- )
2040
-
1964
from vllm.model_executor.layers.fused_moe import fused_experts
1965
2043
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
1966
+ topk_weights, topk_ids, _ = layer.select_experts(
1967
hidden_states=x,
1968
router_logits=router_logits,
2046
- use_grouped_topk=use_grouped_topk,
2047
- top_k=top_k,
2048
- renormalize=renormalize,
2049
- topk_group=topk_group,
2050
- num_expert_group=num_expert_group,
2051
- custom_routing_function=custom_routing_function,
2052
- scoring_func=scoring_func,
2053
- routed_scaling_factor=routed_scaling_factor,
2054
- e_score_correction_bias=e_score_correction_bias,
2055
- indices_type=self.topk_indices_dtype,
2056
- num_fused_shared_experts=getattr(layer, "num_fused_shared_experts", 0),
2057
- enable_eplb=enable_eplb,
2058
- expert_map=expert_map,
2059
- expert_load_view=expert_load_view,
2060
- logical_to_physical_map=logical_to_physical_map,
2061
- logical_replica_count=logical_replica_count,
1969
)
1970
1971
return fused_experts(
vllm/model_executor/layers/quantization/experts_int8.py
+2
-17
@@ -137,7 +137,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
137
138
def apply(
139
self,
140
- layer: torch.nn.Module,
140
+ layer: FusedMoE,
141
x: torch.Tensor,
142
router_logits: torch.Tensor,
143
top_k: int,
@@ -158,26 +158,11 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
158
logical_to_physical_map: torch.Tensor | None = None,
159
logical_replica_count: torch.Tensor | None = None,
160
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
161
- if enable_eplb:
162
- raise NotImplementedError(
163
- "EPLB not supported for `ExpertsInt8MoEMethod` yet."
164
- )
165
-
161
from vllm.model_executor.layers.fused_moe import fused_experts
162
168
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
163
+ topk_weights, topk_ids, _ = layer.select_experts(
164
hidden_states=x,
165
router_logits=router_logits,
171
- use_grouped_topk=use_grouped_topk,
172
- top_k=top_k,
173
- renormalize=renormalize,
174
- topk_group=topk_group,
175
- num_expert_group=num_expert_group,
176
- custom_routing_function=custom_routing_function,
177
- scoring_func=scoring_func,
178
- routed_scaling_factor=routed_scaling_factor,
179
- e_score_correction_bias=e_score_correction_bias,
180
- indices_type=self.topk_indices_dtype,
166
)
167
168
return fused_experts(
vllm/model_executor/layers/quantization/fp8.py
+4
-25
@@ -1140,7 +1140,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1140
1141
def apply(
1142
self,
1143
- layer: torch.nn.Module,
1143
+ layer: FusedMoE,
1144
x: torch.Tensor,
1145
router_logits: torch.Tensor,
1146
top_k: int,
@@ -1216,31 +1216,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1216
apply_router_weight_on_input=apply_router_weight_on_input,
1217
)
1218
1219
- zero_expert_num = getattr(layer, "zero_expert_num", 0)
1220
- zero_expert_type = getattr(layer, "zero_expert_type", None)
1221
-
1222
- select_result = FusedMoE.select_experts(
1219
+ select_result = layer.select_experts(
1220
hidden_states=x,
1221
router_logits=router_logits,
1225
- use_grouped_topk=use_grouped_topk,
1226
- top_k=top_k,
1227
- renormalize=renormalize,
1228
- topk_group=topk_group,
1229
- num_expert_group=num_expert_group,
1230
- custom_routing_function=custom_routing_function,
1231
- scoring_func=scoring_func,
1232
- routed_scaling_factor=routed_scaling_factor,
1233
- e_score_correction_bias=e_score_correction_bias,
1234
- indices_type=self.topk_indices_dtype,
1235
- enable_eplb=enable_eplb,
1236
- expert_map=expert_map,
1237
- expert_load_view=expert_load_view,
1238
- logical_to_physical_map=logical_to_physical_map,
1239
- logical_replica_count=logical_replica_count,
1240
- global_num_experts=global_num_experts,
1241
- zero_expert_num=zero_expert_num,
1242
- zero_expert_type=zero_expert_type,
1243
- num_fused_shared_experts=layer.num_fused_shared_experts,
1222
)
1223
1224
topk_weights, topk_ids, zero_expert_result = select_result
@@ -1322,7 +1300,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1300
self.allow_cutlass_block_scaled_grouped_gemm
1301
),
1302
)
1325
- if zero_expert_num != 0 and zero_expert_type is not None:
1303
+
1304
+ if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
1305
assert not isinstance(result, tuple), (
1306
"Shared + zero experts are mutually exclusive not yet supported"
1307
)
vllm/model_executor/layers/quantization/gguf.py
+2
-15
@@ -621,7 +621,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):
621
622
def apply(
623
self,
624
- layer: torch.nn.Module,
624
+ layer: FusedMoE,
625
x: torch.Tensor,
626
router_logits: torch.Tensor,
627
top_k: int,
@@ -642,9 +642,6 @@ class GGUFMoEMethod(FusedMoEMethodBase):
642
logical_to_physical_map: torch.Tensor | None = None,
643
logical_replica_count: torch.Tensor | None = None,
644
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
645
- if enable_eplb:
646
- raise NotImplementedError("EPLB not supported for `GGUFMoEMethod` yet.")
647
-
645
assert activation == "silu", "Only SiLU activation is supported."
646
if apply_router_weight_on_input:
647
raise NotImplementedError(
@@ -652,19 +649,9 @@ class GGUFMoEMethod(FusedMoEMethodBase):
649
"fused GGUF MoE method."
650
)
651
655
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
652
+ topk_weights, topk_ids, _ = layer.select_experts(
653
hidden_states=x,
654
router_logits=router_logits,
658
- use_grouped_topk=use_grouped_topk,
659
- top_k=top_k,
660
- renormalize=renormalize,
661
- topk_group=topk_group,
662
- num_expert_group=num_expert_group,
663
- custom_routing_function=custom_routing_function,
664
- scoring_func=scoring_func,
665
- routed_scaling_factor=routed_scaling_factor,
666
- e_score_correction_bias=e_score_correction_bias,
667
- indices_type=self.topk_indices_dtype,
655
)
656
return fused_moe_gguf(
657
x,
vllm/model_executor/layers/quantization/gptq_marlin.py
+2
-17
@@ -722,7 +722,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
722
723
def apply(
724
self,
725
- layer: torch.nn.Module,
725
+ layer: FusedMoE,
726
x: torch.Tensor,
727
router_logits: torch.Tensor,
728
top_k: int,
@@ -743,26 +743,11 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
743
logical_to_physical_map: torch.Tensor | None = None,
744
logical_replica_count: torch.Tensor | None = None,
745
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
746
- if enable_eplb:
747
- raise NotImplementedError(
748
- "EPLB not supported for `GPTQMarlinMoEMethod` yet."
749
- )
750
-
746
assert activation == "silu", "Only SiLU activation is supported."
747
753
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
748
+ topk_weights, topk_ids, _ = layer.select_experts(
749
hidden_states=x,
750
router_logits=router_logits,
756
- use_grouped_topk=use_grouped_topk,
757
- top_k=top_k,
758
- renormalize=renormalize,
759
- topk_group=topk_group,
760
- num_expert_group=num_expert_group,
761
- custom_routing_function=custom_routing_function,
762
- scoring_func=scoring_func,
763
- routed_scaling_factor=routed_scaling_factor,
764
- e_score_correction_bias=e_score_correction_bias,
765
- indices_type=self.topk_indices_dtype,
751
)
752
753
return fused_marlin_moe(
vllm/model_executor/layers/quantization/modelopt.py
+12
-33
@@ -696,7 +696,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
696
697
def apply(
698
self,
699
- layer: torch.nn.Module,
699
+ layer: FusedMoE,
700
x: torch.Tensor,
701
router_logits: torch.Tensor,
702
top_k: int,
@@ -717,12 +717,11 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
717
logical_to_physical_map: torch.Tensor | None = None,
718
logical_replica_count: torch.Tensor | None = None,
719
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
720
- if enable_eplb:
721
- raise NotImplementedError(
722
- "EPLB not supported for `ModelOptFp8MoEMethod` yet."
723
- )
724
-
720
if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
721
+ if layer.enable_eplb:
722
+ raise NotImplementedError(
723
+ "EPLB not supported for `ModelOptFp8MoEMethod` yet."
724
+ )
725
assert activation == "silu", (
726
f"Expected 'silu' activation but got {activation}"
727
)
@@ -740,19 +739,9 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
739
)
740
741
# Expert selection
743
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
742
+ topk_weights, topk_ids, _ = layer.select_experts(
743
hidden_states=x,
744
router_logits=router_logits,
746
- use_grouped_topk=use_grouped_topk,
747
- top_k=top_k,
748
- renormalize=renormalize,
749
- topk_group=topk_group,
750
- num_expert_group=num_expert_group,
751
- custom_routing_function=custom_routing_function,
752
- scoring_func=scoring_func,
753
- routed_scaling_factor=routed_scaling_factor,
754
- e_score_correction_bias=e_score_correction_bias,
755
- indices_type=self.topk_indices_dtype,
745
)
746
747
if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
@@ -1459,7 +1448,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1448
1449
def apply(
1450
self,
1462
- layer: torch.nn.Module,
1451
+ layer: FusedMoE,
1452
x: torch.Tensor,
1453
router_logits: torch.Tensor,
1454
top_k: int,
@@ -1480,16 +1469,16 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1469
logical_to_physical_map: torch.Tensor | None = None,
1470
logical_replica_count: torch.Tensor | None = None,
1471
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1483
- if enable_eplb:
1484
- raise NotImplementedError(
1485
- "EPLB not supported for `ModelOptNvFp4FusedMoE` yet."
1486
- )
1472
assert activation == "silu", "Only SiLU activation is supported."
1473
1474
if (
1475
self.allow_flashinfer
1476
and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
1477
):
1478
+ if enable_eplb:
1479
+ raise NotImplementedError(
1480
+ "EPLB not supported for `ModelOptNvFp4FusedMoE` yet."
1481
+ )
1482
return flashinfer_trtllm_fp4_moe(
1483
layer=layer,
1484
x=x,
@@ -1502,19 +1491,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1491
e_score_correction_bias=e_score_correction_bias,
1492
)
1493
1505
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
1494
+ topk_weights, topk_ids, _ = layer.select_experts(
1495
hidden_states=x,
1496
router_logits=router_logits,
1508
- use_grouped_topk=use_grouped_topk,
1509
- top_k=top_k,
1510
- renormalize=renormalize,
1511
- topk_group=topk_group,
1512
- num_expert_group=num_expert_group,
1513
- custom_routing_function=custom_routing_function,
1514
- scoring_func=scoring_func,
1515
- routed_scaling_factor=routed_scaling_factor,
1516
- e_score_correction_bias=e_score_correction_bias,
1517
- indices_type=self.topk_indices_dtype,
1497
)
1498
1499
if self.use_marlin:
vllm/model_executor/layers/quantization/moe_wna16.py
+2
-15
@@ -359,7 +359,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
359
360
def apply(
361
self,
362
- layer: torch.nn.Module,
362
+ layer: FusedMoE,
363
x: torch.Tensor,
364
router_logits: torch.Tensor,
365
top_k: int,
@@ -380,25 +380,12 @@ class MoeWNA16Method(FusedMoEMethodBase):
380
logical_to_physical_map: torch.Tensor | None = None,
381
logical_replica_count: torch.Tensor | None = None,
382
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
383
- if enable_eplb:
384
- raise NotImplementedError("EPLB not supported for `MoeWNA16Method` yet.")
385
-
383
from vllm.model_executor.layers.fused_moe import fused_experts
384
385
assert activation == "silu", "Only SiLU activation is supported."
389
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
386
+ topk_weights, topk_ids, _ = layer.select_experts(
387
hidden_states=x,
388
router_logits=router_logits,
392
- use_grouped_topk=use_grouped_topk,
393
- top_k=top_k,
394
- renormalize=renormalize,
395
- topk_group=topk_group,
396
- num_expert_group=num_expert_group,
397
- custom_routing_function=custom_routing_function,
398
- scoring_func=scoring_func,
399
- routed_scaling_factor=routed_scaling_factor,
400
- e_score_correction_bias=e_score_correction_bias,
401
- indices_type=self.topk_indices_dtype,
389
)
390
391
return fused_experts(
vllm/model_executor/layers/quantization/mxfp4.py
+3
-20
@@ -862,7 +862,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
862
863
def apply(
864
self,
865
- layer: torch.nn.Module,
865
+ layer: FusedMoE,
866
x: torch.Tensor,
867
router_logits: torch.Tensor,
868
top_k: int,
@@ -887,18 +887,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
887
raise NotImplementedError("EPLB is not supported for mxfp4")
888
889
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
890
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
890
+ topk_weights, topk_ids, _ = layer.select_experts(
891
hidden_states=x,
892
router_logits=router_logits,
893
- use_grouped_topk=use_grouped_topk,
894
- top_k=top_k,
895
- renormalize=renormalize,
896
- topk_group=topk_group,
897
- num_expert_group=num_expert_group,
898
- custom_routing_function=custom_routing_function,
899
- scoring_func=scoring_func,
900
- routed_scaling_factor=routed_scaling_factor,
901
- e_score_correction_bias=e_score_correction_bias,
893
)
894
895
return fused_marlin_moe(
@@ -989,17 +980,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
980
):
981
from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
982
992
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
983
+ topk_weights, topk_ids, _ = layer.select_experts(
984
hidden_states=x,
985
router_logits=router_logits,
995
- use_grouped_topk=use_grouped_topk,
996
- top_k=top_k,
997
- renormalize=renormalize,
998
- topk_group=topk_group,
999
- num_expert_group=num_expert_group,
1000
- custom_routing_function=custom_routing_function,
1001
- scoring_func=scoring_func,
1002
- e_score_correction_bias=e_score_correction_bias,
986
)
987
988
# Backend-specific preparation
vllm/model_executor/layers/quantization/quark/quark_moe.py
+4
-34
@@ -334,7 +334,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
334
335
def apply(
336
self,
337
- layer: torch.nn.Module,
337
+ layer: FusedMoE,
338
x: torch.Tensor,
339
router_logits: torch.Tensor,
340
top_k: int,
@@ -355,24 +355,9 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
355
logical_to_physical_map: torch.Tensor | None = None,
356
logical_replica_count: torch.Tensor | None = None,
357
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
358
- if enable_eplb:
359
- raise NotImplementedError(
360
- "EPLB not supported for `QuarkW8A8Fp8MoEMethod` yet."
361
- )
362
-
363
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
358
+ topk_weights, topk_ids, _ = layer.select_experts(
359
hidden_states=x,
360
router_logits=router_logits,
366
- use_grouped_topk=use_grouped_topk,
367
- top_k=top_k,
368
- renormalize=renormalize,
369
- topk_group=topk_group,
370
- num_expert_group=num_expert_group,
371
- custom_routing_function=custom_routing_function,
372
- scoring_func=scoring_func,
373
- routed_scaling_factor=routed_scaling_factor,
374
- e_score_correction_bias=e_score_correction_bias,
375
- indices_type=self.topk_indices_dtype,
361
)
362
363
if self.rocm_aiter_moe_enabled:
@@ -609,7 +594,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
594
595
def apply(
596
self,
612
- layer: torch.nn.Module,
597
+ layer: FusedMoE,
598
x: torch.Tensor,
599
router_logits: torch.Tensor,
600
top_k: int,
@@ -630,24 +615,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
615
logical_to_physical_map: torch.Tensor | None = None,
616
logical_replica_count: torch.Tensor | None = None,
617
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
633
- if enable_eplb:
634
- raise NotImplementedError(
635
- "EPLB not supported for `QuarkOCP_MX_MoEMethod` yet."
636
- )
637
-
638
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
618
+ topk_weights, topk_ids, _ = layer.select_experts(
619
hidden_states=x,
620
router_logits=router_logits,
641
- use_grouped_topk=use_grouped_topk,
642
- top_k=top_k,
643
- renormalize=renormalize,
644
- topk_group=topk_group,
645
- num_expert_group=num_expert_group,
646
- custom_routing_function=custom_routing_function,
647
- scoring_func=scoring_func,
648
- routed_scaling_factor=routed_scaling_factor,
649
- e_score_correction_bias=e_score_correction_bias,
650
- indices_type=self.topk_indices_dtype,
621
)
622
623
if not self.emulate:
vllm/model_executor/layers/quantization/rtn.py
+2
-15
@@ -356,7 +356,7 @@ class RTNMoEMethod(FusedMoEMethodBase):
356
357
def apply(
358
self,
359
- layer: torch.nn.Module,
359
+ layer: FusedMoE,
360
x: torch.Tensor,
361
router_logits: torch.Tensor,
362
top_k: int,
@@ -377,22 +377,9 @@ class RTNMoEMethod(FusedMoEMethodBase):
377
logical_to_physical_map: torch.Tensor | None = None,
378
logical_replica_count: torch.Tensor | None = None,
379
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
380
- if enable_eplb:
381
- raise NotImplementedError("EPLB not supported for `RTNMoEMethod` yet.")
382
-
383
- topk_weights, topk_ids, _ = FusedMoE.select_experts(
380
+ topk_weights, topk_ids, _ = layer.select_experts(
381
hidden_states=x,
382
router_logits=router_logits,
386
- use_grouped_topk=use_grouped_topk,
387
- top_k=top_k,
388
- renormalize=renormalize,
389
- topk_group=topk_group,
390
- num_expert_group=num_expert_group,
391
- custom_routing_function=custom_routing_function,
392
- scoring_func=scoring_func,
393
- routed_scaling_factor=routed_scaling_factor,
394
- e_score_correction_bias=e_score_correction_bias,
395
- indices_type=self.topk_indices_dtype,
383
)
384
385
return fused_marlin_moe(