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 11, 17:30 UTC
[MoE][Refactor] Remove most arguments to FusedMoEMethodBase.apply (#29066)
Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
bnellnm committed
Dec 9, 2025 at 16:48 UTC
00e5cbb96789fd9d083b4015cdcff318e4de3808
18 files changed
+318
-872
vllm/model_executor/layers/fused_moe/__init__.py
+5
-1
@@ -4,7 +4,10 @@
4
from contextlib import contextmanager
5
from typing import Any
6
7
-from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
7
+from vllm.model_executor.layers.fused_moe.config import (
8
+ FusedMoEConfig,
9
+ RoutingMethodType,
10
+)
11
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
12
FusedMoEMethodBase,
13
)
@@ -49,6 +52,7 @@ __all__ = [
52
"FusedMoEPermuteExpertsUnpermute",
53
"FusedMoEActivationFormat",
54
"FusedMoEPrepareAndFinalize",
55
+ "RoutingMethodType",
56
"SharedFusedMoE",
57
"activation_without_mul",
58
"override_config",
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
-18
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from abc import abstractmethod
5
-from collections.abc import Callable
5
6
import torch
7
@@ -100,22 +99,5 @@ class FusedMoEMethodBase(QuantizeMethodBase):
99
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
100
x: torch.Tensor,
101
router_logits: torch.Tensor,
103
- top_k: int,
104
- renormalize: bool,
105
- use_grouped_topk: bool = False,
106
- topk_group: int | None = None,
107
- num_expert_group: int | None = None,
108
- global_num_experts: int = -1,
109
- expert_map: torch.Tensor | None = None,
110
- custom_routing_function: Callable | None = None,
111
- scoring_func: str = "softmax",
112
- routed_scaling_factor: float = 1.0,
113
- e_score_correction_bias: torch.Tensor | None = None,
114
- apply_router_weight_on_input: bool = False,
115
- activation: str = "silu",
116
- enable_eplb: bool = False,
117
- expert_load_view: torch.Tensor | None = None,
118
- logical_to_physical_map: torch.Tensor | None = None,
119
- logical_replica_count: torch.Tensor | None = None,
102
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
103
raise NotImplementedError
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+4
-22
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
5
import torch
6
@@ -97,23 +96,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
96
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
97
x: torch.Tensor,
98
router_logits: torch.Tensor,
100
- top_k: int,
101
- renormalize: bool,
102
- use_grouped_topk: bool = False,
103
- topk_group: int | None = None,
104
- num_expert_group: int | None = None,
105
- global_num_experts: int = -1,
106
- expert_map: torch.Tensor | None = None,
107
- custom_routing_function: Callable | None = None,
108
- scoring_func: str = "softmax",
109
- routed_scaling_factor: float = 1.0,
110
- e_score_correction_bias: torch.Tensor | None = None,
111
- apply_router_weight_on_input: bool = False,
112
- activation: str = "silu",
113
- enable_eplb: bool = False,
114
- expert_load_view: torch.Tensor | None = None,
115
- logical_to_physical_map: torch.Tensor | None = None,
116
- logical_replica_count: torch.Tensor | None = None,
99
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
100
topk_weights, topk_ids, zero_expert_result = layer.select_experts(
101
hidden_states=x,
@@ -127,10 +109,10 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
109
topk_weights=topk_weights,
110
topk_ids=topk_ids,
111
inplace=self.allow_inplace,
130
- activation=activation,
131
- global_num_experts=global_num_experts,
132
- apply_router_weight_on_input=apply_router_weight_on_input,
133
- expert_map=None if self.disable_expert_map else expert_map,
112
+ activation=layer.activation,
113
+ global_num_experts=layer.global_num_experts,
114
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
115
+ expert_map=None if self.disable_expert_map else layer.expert_map,
116
)
117
118
if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
vllm/model_executor/layers/fused_moe/layer.py
+18
-56
@@ -33,10 +33,6 @@ from vllm.model_executor.layers.fused_moe.config import (
33
RoutingMethodType,
34
)
35
from vllm.model_executor.layers.fused_moe.fused_moe import zero_experts_compute_triton
36
-from vllm.model_executor.layers.fused_moe.modular_kernel import (
37
- FusedMoEPermuteExpertsUnpermute,
38
- FusedMoEPrepareAndFinalize,
39
-)
36
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
37
init_aiter_topK_meta_data,
38
)
@@ -57,11 +53,8 @@ from vllm.utils.torch_utils import (
53
from vllm.v1.worker.ubatching import dbo_current_ubatch_id
54
55
if current_platform.is_cuda_alike():
60
- from .fused_moe import eplb_map_to_physical_and_record, fused_experts
56
+ from .fused_moe import eplb_map_to_physical_and_record
57
else:
62
- fused_experts = None # type: ignore
63
- FusedMoEPermuteExpertsUnpermute = object # type: ignore
64
- FusedMoEPrepareAndFinalize = object # type: ignore
58
59
def _eplb_map_to_physical_and_record(
60
topk_ids: torch.Tensor,
@@ -483,7 +476,7 @@ class FusedMoE(CustomOp):
476
enable_eplb=self.enable_eplb,
477
)
478
486
- self.expert_map: torch.Tensor | None
479
+ self._expert_map: torch.Tensor | None
480
local_num_experts, expert_map, expert_mask = determine_expert_map(
481
ep_size=self.ep_size,
482
ep_rank=self.ep_rank,
@@ -493,7 +486,7 @@ class FusedMoE(CustomOp):
486
return_expert_mask=self.rocm_aiter_fmoe_enabled,
487
)
488
self.local_num_experts = local_num_experts
496
- self.register_buffer("expert_map", expert_map)
489
+ self.register_buffer("_expert_map", expert_map)
490
self.register_buffer("expert_mask", expert_mask)
491
self._maybe_init_expert_routing_tables()
492
logger.info_once(
@@ -506,10 +499,10 @@ class FusedMoE(CustomOp):
499
self.expert_placement_strategy,
500
self.local_num_experts,
501
self.global_num_experts,
509
- get_compressed_expert_map(self.expert_map),
502
+ get_compressed_expert_map(self._expert_map),
503
)
504
else:
512
- self.local_num_experts, self.expert_map, self.expert_mask = (
505
+ self.local_num_experts, self._expert_map, self.expert_mask = (
506
self.global_num_experts,
507
None,
508
None,
@@ -781,7 +774,7 @@ class FusedMoE(CustomOp):
774
),
775
)
776
784
- if self.expert_map is None:
777
+ if self._expert_map is None:
778
return None
779
780
routing_tables = self.ensure_round_robin_expert_routing_tables(
@@ -789,7 +782,7 @@ class FusedMoE(CustomOp):
782
ep_size=self.ep_size,
783
ep_rank=self.ep_rank,
784
local_num_experts=self.local_num_experts,
792
- device=self.expert_map.device,
785
+ device=self._expert_map.device,
786
)
787
788
global_to_physical, physical_to_global, local_global = routing_tables
@@ -840,8 +833,8 @@ class FusedMoE(CustomOp):
833
834
def update_expert_map(self):
835
# ep_size and ep_rank should already be updated
843
- assert self.expert_map is not None
844
- with self.expert_map.device:
836
+ assert self._expert_map is not None
837
+ with self._expert_map.device:
838
local_num_experts, expert_map, expert_mask = determine_expert_map(
839
ep_size=self.ep_size,
840
ep_rank=self.ep_rank,
@@ -851,7 +844,7 @@ class FusedMoE(CustomOp):
844
return_expert_mask=self.rocm_aiter_fmoe_enabled,
845
)
846
self.local_num_experts = local_num_experts
854
- self.register_buffer("expert_map", expert_map)
847
+ self.register_buffer("_expert_map", expert_map)
848
self.register_buffer("expert_mask", expert_mask)
849
self._maybe_init_expert_routing_tables()
850
if self.aiter_fmoe_shared_expert_enabled:
@@ -1068,9 +1061,9 @@ class FusedMoE(CustomOp):
1061
expert_data.copy_(loaded_weight)
1062
1063
def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int:
1071
- if self.expert_map is None:
1064
+ if self._expert_map is None:
1065
return expert_id
1073
- return self.expert_map[expert_id].item()
1066
+ return self._expert_map[expert_id].item()
1067
1068
def _init_aiter_shared_experts_topK_buffer(
1069
self, vllm_config: VllmConfig, dp_size: int
@@ -1744,6 +1737,12 @@ class FusedMoE(CustomOp):
1737
reduce_output(fused_output)[..., :og_hidden_states],
1738
)
1739
1740
+ @property
1741
+ def expert_map(self) -> torch.Tensor | None:
1742
+ return (
1743
+ self._expert_map if not self.rocm_aiter_fmoe_enabled else self.expert_mask
1744
+ )
1745
+
1746
def forward_cuda(
1747
self,
1748
hidden_states: torch.Tensor,
@@ -1805,24 +1804,6 @@ class FusedMoE(CustomOp):
1804
layer=self,
1805
x=staged_hidden_states,
1806
router_logits=staged_router_logits,
1808
- top_k=self.top_k,
1809
- renormalize=self.renormalize,
1810
- use_grouped_topk=self.use_grouped_topk,
1811
- global_num_experts=self.global_num_experts,
1812
- expert_map=self.expert_map
1813
- if not self.rocm_aiter_fmoe_enabled
1814
- else self.expert_mask,
1815
- topk_group=self.topk_group,
1816
- num_expert_group=self.num_expert_group,
1817
- custom_routing_function=self.custom_routing_function,
1818
- scoring_func=self.scoring_func,
1819
- routed_scaling_factor=self.routed_scaling_factor,
1820
- e_score_correction_bias=self.e_score_correction_bias,
1821
- activation=self.activation,
1822
- enable_eplb=self.enable_eplb,
1823
- expert_load_view=self.expert_load_view,
1824
- logical_to_physical_map=self.logical_to_physical_map,
1825
- logical_replica_count=self.logical_replica_count,
1807
)
1808
1809
if has_separate_shared_experts:
@@ -1968,25 +1949,6 @@ class FusedMoE(CustomOp):
1949
if do_naive_dispatch_combine
1950
else hidden_states,
1951
router_logits=router_logits,
1971
- top_k=self.top_k,
1972
- renormalize=self.renormalize,
1973
- use_grouped_topk=self.use_grouped_topk,
1974
- global_num_experts=self.global_num_experts,
1975
- expert_map=self.expert_map
1976
- if not self.rocm_aiter_fmoe_enabled
1977
- else self.expert_mask,
1978
- topk_group=self.topk_group,
1979
- num_expert_group=self.num_expert_group,
1980
- custom_routing_function=self.custom_routing_function,
1981
- scoring_func=self.scoring_func,
1982
- routed_scaling_factor=self.routed_scaling_factor,
1983
- e_score_correction_bias=self.e_score_correction_bias,
1984
- activation=self.activation,
1985
- apply_router_weight_on_input=self.apply_router_weight_on_input,
1986
- enable_eplb=self.enable_eplb,
1987
- expert_load_view=self.expert_load_view,
1988
- logical_to_physical_map=self.logical_to_physical_map,
1989
- logical_replica_count=self.logical_replica_count,
1952
)
1953
1954
if has_separate_shared_experts:
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+60
-164
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
5
import torch
6
import torch.nn.functional as F
@@ -269,53 +268,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
268
269
def apply(
270
self,
272
- layer: torch.nn.Module,
271
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
272
x: torch.Tensor,
273
router_logits: torch.Tensor,
275
- top_k: int,
276
- renormalize: bool,
277
- use_grouped_topk: bool = False,
278
- topk_group: int | None = None,
279
- num_expert_group: int | None = None,
280
- global_num_experts: int = -1,
281
- expert_map: torch.Tensor | None = None,
282
- custom_routing_function: Callable | None = None,
283
- scoring_func: str = "softmax",
284
- routed_scaling_factor: float = 1.0,
285
- e_score_correction_bias: torch.Tensor | None = None,
286
- apply_router_weight_on_input: bool = False,
287
- activation: str = "silu",
288
- enable_eplb: bool = False,
289
- expert_load_view: torch.Tensor | None = None,
290
- logical_to_physical_map: torch.Tensor | None = None,
291
- logical_replica_count: torch.Tensor | None = None,
274
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
293
- if enable_eplb:
294
- assert expert_load_view is not None
295
- assert logical_to_physical_map is not None
296
- assert logical_replica_count is not None
297
-
275
return self.forward(
299
- x=x,
276
layer=layer,
277
+ x=x,
278
router_logits=router_logits,
302
- top_k=top_k,
303
- renormalize=renormalize,
304
- use_grouped_topk=use_grouped_topk,
305
- topk_group=topk_group,
306
- num_expert_group=num_expert_group,
307
- global_num_experts=global_num_experts,
308
- expert_map=expert_map,
309
- custom_routing_function=custom_routing_function,
310
- scoring_func=scoring_func,
311
- routed_scaling_factor=routed_scaling_factor,
312
- e_score_correction_bias=e_score_correction_bias,
313
- activation=activation,
314
- apply_router_weight_on_input=apply_router_weight_on_input,
315
- enable_eplb=enable_eplb,
316
- expert_load_view=expert_load_view,
317
- logical_to_physical_map=logical_to_physical_map,
318
- logical_replica_count=logical_replica_count,
279
)
280
281
def get_fused_moe_quant_config(
@@ -333,24 +293,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
293
self,
294
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
295
x: torch.Tensor,
336
- use_grouped_topk: bool,
337
- top_k: int,
296
router_logits: torch.Tensor,
339
- renormalize: bool,
340
- topk_group: int | None = None,
341
- num_expert_group: int | None = None,
342
- global_num_experts: int = -1,
343
- expert_map: torch.Tensor | None = None,
344
- custom_routing_function: Callable | None = None,
345
- scoring_func: str = "softmax",
346
- routed_scaling_factor: float = 1.0,
347
- e_score_correction_bias: torch.Tensor | None = None,
348
- apply_router_weight_on_input: bool = False,
349
- activation: str = "silu",
350
- enable_eplb: bool = False,
351
- expert_load_view: torch.Tensor | None = None,
352
- logical_to_physical_map: torch.Tensor | None = None,
353
- logical_replica_count: torch.Tensor | None = None,
297
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
298
topk_weights, topk_ids, zero_expert_result = layer.select_experts(
299
hidden_states=x,
@@ -364,9 +307,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
307
w2=layer.w2_weight,
308
topk_weights=topk_weights,
309
topk_ids=topk_ids,
367
- expert_map=expert_map,
368
- activation=activation,
369
- apply_router_weight_on_input=apply_router_weight_on_input,
310
+ expert_map=layer.expert_map,
311
+ activation=layer.activation,
312
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
313
)
314
elif self.flashinfer_cutlass_moe_enabled:
315
return self.flashinfer_cutlass_moe(
@@ -375,8 +318,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
318
w2=layer.w2_weight,
319
topk_weights=topk_weights,
320
topk_ids=topk_ids,
378
- activation=activation,
379
- apply_router_weight_on_input=apply_router_weight_on_input,
321
+ activation=layer.activation,
322
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
323
)
324
else:
325
result = fused_experts(
@@ -386,11 +329,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
329
topk_weights=topk_weights,
330
topk_ids=topk_ids,
331
inplace=True,
389
- activation=activation,
332
+ activation=layer.activation,
333
quant_config=self.moe_quant_config,
391
- apply_router_weight_on_input=apply_router_weight_on_input,
392
- global_num_experts=global_num_experts,
393
- expert_map=expert_map,
334
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
335
+ global_num_experts=layer.global_num_experts,
336
+ expert_map=layer.expert_map,
337
)
338
339
if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
@@ -405,148 +348,101 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
348
self,
349
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
350
x: torch.Tensor,
408
- use_grouped_topk: bool,
409
- top_k: int,
351
router_logits: torch.Tensor,
411
- renormalize: bool,
412
- topk_group: int | None = None,
413
- num_expert_group: int | None = None,
414
- global_num_experts: int = -1,
415
- expert_map: torch.Tensor | None = None,
416
- custom_routing_function: Callable | None = None,
417
- scoring_func: str = "softmax",
418
- routed_scaling_factor: float = 1.0,
419
- e_score_correction_bias: torch.Tensor | None = None,
420
- apply_router_weight_on_input: bool = False,
421
- activation: str = "silu",
422
- enable_eplb: bool = False,
423
- expert_load_view: torch.Tensor | None = None,
424
- logical_to_physical_map: torch.Tensor | None = None,
425
- logical_replica_count: torch.Tensor | None = None,
352
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
353
if (
428
- enable_eplb is not False
429
- or expert_load_view is not None
430
- or logical_to_physical_map is not None
431
- or logical_replica_count is not None
354
+ layer.enable_eplb is not False
355
+ or layer.expert_load_view is not None
356
+ or layer.logical_to_physical_map is not None
357
+ or layer.logical_replica_count is not None
358
):
359
raise NotImplementedError("Expert load balancing is not supported for CPU.")
360
+
361
return layer.cpu_fused_moe(
362
layer,
363
x,
437
- use_grouped_topk,
438
- top_k,
364
+ layer.use_grouped_topk,
365
+ layer.top_k,
366
router_logits,
440
- renormalize,
441
- topk_group,
442
- num_expert_group,
443
- global_num_experts,
444
- expert_map,
445
- custom_routing_function,
446
- scoring_func,
447
- routed_scaling_factor,
448
- e_score_correction_bias,
449
- apply_router_weight_on_input,
450
- activation,
367
+ layer.renormalize,
368
+ layer.topk_group,
369
+ layer.num_expert_group,
370
+ layer.global_num_experts,
371
+ layer.expert_map,
372
+ layer.custom_routing_function,
373
+ layer.scoring_func,
374
+ layer.routed_scaling_factor,
375
+ layer.e_score_correction_bias,
376
+ layer.apply_router_weight_on_input,
377
+ layer.activation,
378
)
379
380
def forward_xpu(
381
self,
382
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
383
x: torch.Tensor,
457
- use_grouped_topk: bool,
458
- top_k: int,
384
router_logits: torch.Tensor,
460
- renormalize: bool,
461
- topk_group: int | None = None,
462
- num_expert_group: int | None = None,
463
- global_num_experts: int = -1,
464
- expert_map: torch.Tensor | None = None,
465
- custom_routing_function: Callable | None = None,
466
- scoring_func: str = "softmax",
467
- routed_scaling_factor: float = 1.0,
468
- e_score_correction_bias: torch.Tensor | None = None,
469
- apply_router_weight_on_input: bool = False,
470
- activation: str = "silu",
471
- enable_eplb: bool = False,
472
- expert_load_view: torch.Tensor | None = None,
473
- logical_to_physical_map: torch.Tensor | None = None,
474
- logical_replica_count: torch.Tensor | None = None,
385
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
386
if (
477
- enable_eplb is not False
478
- or expert_load_view is not None
479
- or logical_to_physical_map is not None
480
- or logical_replica_count is not None
387
+ layer.enable_eplb is not False
388
+ or layer.expert_load_view is not None
389
+ or layer.logical_to_physical_map is not None
390
+ or layer.logical_replica_count is not None
391
):
392
raise NotImplementedError("Expert load balancing is not supported for XPU.")
393
return layer.ipex_fusion(
394
x,
485
- use_grouped_topk,
486
- top_k,
395
+ layer.use_grouped_topk,
396
+ layer.top_k,
397
router_logits,
488
- renormalize,
489
- topk_group,
490
- num_expert_group,
491
- custom_routing_function=custom_routing_function,
398
+ layer.renormalize,
399
+ layer.topk_group,
400
+ layer.num_expert_group,
401
+ custom_routing_function=layer.custom_routing_function,
402
)
403
404
def forward_tpu(
405
self,
406
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
407
x: torch.Tensor,
498
- use_grouped_topk: bool,
499
- top_k: int,
408
router_logits: torch.Tensor,
501
- renormalize: bool,
502
- topk_group: int | None = None,
503
- num_expert_group: int | None = None,
504
- global_num_experts: int = -1,
505
- expert_map: torch.Tensor | None = None,
506
- custom_routing_function: Callable | None = None,
507
- scoring_func: str = "softmax",
508
- routed_scaling_factor: float = 1.0,
509
- e_score_correction_bias: torch.Tensor | None = None,
510
- apply_router_weight_on_input: bool = False,
511
- activation: str = "silu",
512
- enable_eplb: bool = False,
513
- expert_load_view: torch.Tensor | None = None,
514
- logical_to_physical_map: torch.Tensor | None = None,
515
- logical_replica_count: torch.Tensor | None = None,
409
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
517
- assert not use_grouped_topk
518
- assert num_expert_group is None
519
- assert topk_group is None
520
- assert custom_routing_function is None
521
- assert apply_router_weight_on_input is False
522
- if scoring_func != "softmax":
410
+ assert not layer.use_grouped_topk
411
+ assert layer.num_expert_group is None
412
+ assert layer.topk_group is None
413
+ assert layer.custom_routing_function is None
414
+ assert layer.apply_router_weight_on_input is False
415
+ if layer.scoring_func != "softmax":
416
raise NotImplementedError(
417
"Only softmax scoring function is supported for TPU."
418
)
526
- if e_score_correction_bias is not None:
419
+ if layer.e_score_correction_bias is not None:
420
raise NotImplementedError(
421
"Expert score correction bias is not supported for TPU."
422
)
530
- assert activation == "silu", f"{activation} is not supported for TPU."
531
- assert routed_scaling_factor == 1.0, (
532
- f"routed_scaling_factor {routed_scaling_factor} is not supported for TPU."
423
+ assert layer.activation == "silu", (
424
+ f"{layer.activation} is not supported for TPU."
425
+ )
426
+ assert layer.routed_scaling_factor == 1.0, (
427
+ f"routed_scaling_factor {layer.routed_scaling_factor} is "
428
+ "not supported for TPU."
429
)
430
if (
535
- enable_eplb is not False
536
- or expert_load_view is not None
537
- or logical_to_physical_map is not None
538
- or logical_replica_count is not None
431
+ layer.enable_eplb is not False
432
+ or layer.expert_load_view is not None
433
+ or layer.logical_to_physical_map is not None
434
+ or layer.logical_replica_count is not None
435
):
436
raise NotImplementedError("Expert load balancing is not supported for TPU.")
437
return fused_moe_pallas(
438
hidden_states=x,
439
w1=layer.w13_weight,
440
w2=layer.w2_weight,
545
- topk=top_k,
441
+ topk=layer.top_k,
442
gating_output=router_logits,
547
- global_num_experts=global_num_experts,
548
- expert_map=expert_map,
549
- renormalize=renormalize,
443
+ global_num_experts=layer.global_num_experts,
444
+ expert_map=layer.expert_map,
445
+ renormalize=layer.renormalize,
446
)
447
448
if current_platform.is_tpu():
vllm/model_executor/layers/quantization/awq_marlin.py
+4
-22
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from typing import TYPE_CHECKING, Any, Optional
5
6
import torch
@@ -669,25 +668,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
668
layer: FusedMoE,
669
x: torch.Tensor,
670
router_logits: torch.Tensor,
672
- top_k: int,
673
- renormalize: bool,
674
- use_grouped_topk: bool = False,
675
- topk_group: int | None = None,
676
- num_expert_group: int | None = None,
677
- global_num_experts: int = -1,
678
- expert_map: torch.Tensor | None = None,
679
- custom_routing_function: Callable | None = None,
680
- scoring_func: str = "softmax",
681
- routed_scaling_factor: float = 1.0,
682
- e_score_correction_bias: torch.Tensor | None = None,
683
- apply_router_weight_on_input: bool = False,
684
- activation: str = "silu",
685
- enable_eplb: bool = False,
686
- expert_load_view: torch.Tensor | None = None,
687
- logical_to_physical_map: torch.Tensor | None = None,
688
- logical_replica_count: torch.Tensor | None = None,
671
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
690
- assert activation == "silu", "Only SiLU activation is supported."
672
+ assert layer.activation == "silu", "Only SiLU activation is supported."
673
674
topk_weights, topk_ids, _ = layer.select_experts(
675
hidden_states=x,
@@ -708,9 +690,9 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
690
input_global_scale1=getattr(layer, "w13_input_global_scale", None),
691
input_global_scale2=getattr(layer, "w2_input_global_scale", None),
692
quant_type_id=self.quant_type.id,
711
- apply_router_weight_on_input=apply_router_weight_on_input,
712
- global_num_experts=global_num_experts,
713
- expert_map=expert_map,
693
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
694
+ global_num_experts=layer.global_num_experts,
695
+ expert_map=layer.expert_map,
696
w1_zeros=layer.w13_qzeros,
697
w2_zeros=layer.w2_qzeros,
698
workspace=layer.workspace,
vllm/model_executor/layers/quantization/bitsandbytes.py
+4
-22
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from typing import Any, Union
5
6
import torch
@@ -498,23 +497,6 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
497
layer: FusedMoE,
498
x: torch.Tensor,
499
router_logits: torch.Tensor,
501
- top_k: int,
502
- renormalize: bool,
503
- use_grouped_topk: bool = False,
504
- topk_group: int | None = None,
505
- num_expert_group: int | None = None,
506
- global_num_experts: int = -1,
507
- expert_map: torch.Tensor | None = None,
508
- custom_routing_function: Callable | None = None,
509
- scoring_func: str = "softmax",
510
- routed_scaling_factor: float = 1.0,
511
- e_score_correction_bias: torch.Tensor | None = None,
512
- apply_router_weight_on_input: bool = False,
513
- activation: str = "silu",
514
- enable_eplb: bool = False,
515
- expert_load_view: torch.Tensor | None = None,
516
- logical_to_physical_map: torch.Tensor | None = None,
517
- logical_replica_count: torch.Tensor | None = None,
500
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
501
from vllm.model_executor.layers.fused_moe import fused_experts
502
@@ -534,10 +516,10 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
516
topk_weights=topk_weights,
517
topk_ids=topk_ids,
518
inplace=True,
537
- activation=activation,
538
- apply_router_weight_on_input=apply_router_weight_on_input,
539
- global_num_experts=global_num_experts,
540
- expert_map=expert_map,
519
+ activation=layer.activation,
520
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
521
+ global_num_experts=layer.global_num_experts,
522
+ expert_map=layer.expert_map,
523
quant_config=self.moe_quant_config,
524
)
525
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+67
-187
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import enum
5
-from collections.abc import Callable
5
from enum import Enum
6
7
import torch
@@ -558,31 +557,14 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
557
layer: FusedMoE,
558
x: torch.Tensor,
559
router_logits: torch.Tensor,
561
- top_k: int,
562
- renormalize: bool,
563
- use_grouped_topk: bool = False,
564
- topk_group: int | None = None,
565
- num_expert_group: int | None = None,
566
- global_num_experts: int = -1,
567
- expert_map: torch.Tensor | None = None,
568
- custom_routing_function: Callable | None = None,
569
- scoring_func: str = "softmax",
570
- routed_scaling_factor: float = 1.0,
571
- e_score_correction_bias: torch.Tensor | None = None,
572
- apply_router_weight_on_input: bool = False,
573
- activation: str = "silu",
574
- enable_eplb: bool = False,
575
- expert_load_view: torch.Tensor | None = None,
576
- logical_to_physical_map: torch.Tensor | None = None,
577
- logical_replica_count: torch.Tensor | None = None,
560
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
579
- assert activation == "silu", "Only SiLU activation is supported."
561
+ assert layer.activation == "silu", "Only SiLU activation is supported."
562
563
if (
564
self.allow_flashinfer
565
and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
566
):
585
- if enable_eplb:
567
+ if layer.enable_eplb:
568
raise NotImplementedError(
569
"EPLB not supported for `CompressedTensorsW4A4MoEMethod` yet."
570
)
@@ -591,12 +573,12 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
573
layer=layer,
574
x=x,
575
router_logits=router_logits,
594
- top_k=top_k,
595
- global_num_experts=global_num_experts,
596
- num_expert_group=num_expert_group,
597
- topk_group=topk_group,
598
- custom_routing_function=custom_routing_function,
599
- e_score_correction_bias=e_score_correction_bias,
576
+ top_k=layer.top_k,
577
+ global_num_experts=layer.global_num_experts,
578
+ num_expert_group=layer.num_expert_group,
579
+ topk_group=layer.topk_group,
580
+ custom_routing_function=layer.custom_routing_function,
581
+ e_score_correction_bias=layer.e_score_correction_bias,
582
)
583
584
topk_weights, topk_ids, _ = layer.select_experts(
@@ -619,9 +601,9 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
601
global_scale1=layer.w13_weight_scale_2,
602
global_scale2=layer.w2_weight_scale_2,
603
quant_type_id=scalar_types.float4_e2m1f.id,
622
- apply_router_weight_on_input=apply_router_weight_on_input,
623
- global_num_experts=global_num_experts,
624
- expert_map=expert_map,
604
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
605
+ global_num_experts=layer.global_num_experts,
606
+ expert_map=layer.expert_map,
607
input_dtype=self.marlin_input_dtype,
608
workspace=layer.workspace,
609
)
@@ -646,15 +628,15 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
628
topk_ids=topk_ids,
629
quant_config=self.moe_quant_config,
630
inplace=False, # TODO(shuw): fix later, now output is high prec
649
- activation=activation,
650
- global_num_experts=global_num_experts,
651
- expert_map=expert_map,
652
- apply_router_weight_on_input=apply_router_weight_on_input,
631
+ activation=layer.activation,
632
+ global_num_experts=layer.global_num_experts,
633
+ expert_map=layer.expert_map,
634
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
635
)
636
else:
637
from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
638
657
- assert expert_map is None, (
639
+ assert layer.expert_map is None, (
640
"Expert Parallelism / expert_map "
641
"is currently not supported for "
642
"CompressedTensorsW4A4Nvfp4MoEMethod."
@@ -670,7 +652,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
652
topk_weights=topk_weights,
653
topk_ids=topk_ids,
654
quant_config=self.moe_quant_config,
673
- apply_router_weight_on_input=apply_router_weight_on_input,
655
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
656
# TODO(bnell): derive these from arguments
657
m=x.shape[0],
658
n=layer.w2_weight.shape[2] * 2,
@@ -1188,23 +1170,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1170
layer: FusedMoE,
1171
x: torch.Tensor,
1172
router_logits: torch.Tensor,
1191
- top_k: int,
1192
- renormalize: bool,
1193
- use_grouped_topk: bool = False,
1194
- topk_group: int | None = None,
1195
- num_expert_group: int | None = None,
1196
- global_num_experts: int = -1,
1197
- expert_map: torch.Tensor | None = None,
1198
- custom_routing_function: Callable | None = None,
1199
- scoring_func: str = "softmax",
1200
- routed_scaling_factor: float = 1.0,
1201
- e_score_correction_bias: torch.Tensor | None = None,
1202
- apply_router_weight_on_input: bool = False,
1203
- activation: str = "silu",
1204
- enable_eplb: bool = False,
1205
- expert_load_view: torch.Tensor | None = None,
1206
- logical_to_physical_map: torch.Tensor | None = None,
1207
- logical_replica_count: torch.Tensor | None = None,
1173
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1174
topk_weights, topk_ids, _ = layer.select_experts(
1175
hidden_states=x,
@@ -1215,7 +1180,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1180
per_channel_quant = self.weight_quant.strategy == QuantizationStrategy.CHANNEL
1181
1182
if self.use_marlin:
1218
- assert activation == "silu", f"{activation} not supported for Marlin MoE."
1183
+ assert layer.activation == "silu", (
1184
+ f"{layer.activation} not supported for Marlin MoE."
1185
+ )
1186
return fused_marlin_moe(
1187
x,
1188
layer.w13_weight,
@@ -1228,9 +1195,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1195
topk_weights,
1196
topk_ids,
1197
quant_type_id=scalar_types.float8_e4m3fn.id,
1231
- apply_router_weight_on_input=apply_router_weight_on_input,
1232
- global_num_experts=global_num_experts,
1233
- expert_map=expert_map,
1198
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1199
+ global_num_experts=layer.global_num_experts,
1200
+ expert_map=layer.expert_map,
1201
input_dtype=self.marlin_input_dtype,
1202
workspace=layer.workspace,
1203
)
@@ -1248,9 +1215,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1215
w2=layer.w2_weight,
1216
topk_weights=topk_weights,
1217
topk_ids=topk_ids,
1251
- activation=activation,
1252
- apply_router_weight_on_input=apply_router_weight_on_input,
1253
- expert_map=expert_map,
1218
+ activation=layer.activation,
1219
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1220
+ expert_map=layer.expert_map,
1221
quant_config=self.moe_quant_config,
1222
)
1223
@@ -1270,10 +1237,12 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1237
topk_weights=topk_weights,
1238
topk_ids=topk_ids,
1239
inplace=True,
1273
- activation=activation,
1274
- apply_router_weight_on_input=apply_router_weight_on_input,
1275
- global_num_experts=global_num_experts,
1276
- expert_map=None if self.disable_expert_map else expert_map,
1240
+ activation=layer.activation,
1241
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1242
+ global_num_experts=layer.global_num_experts,
1243
+ expert_map=None
1244
+ if self.disable_expert_map
1245
+ else layer.expert_map, # ???
1246
quant_config=self.moe_quant_config,
1247
)
1248
else:
@@ -1290,9 +1259,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1259
topk_weights,
1260
topk_ids,
1261
quant_config=self.moe_quant_config,
1293
- activation=activation,
1294
- global_num_experts=global_num_experts,
1295
- expert_map=None if self.disable_expert_map else expert_map,
1262
+ activation=layer.activation,
1263
+ global_num_experts=layer.global_num_experts,
1264
+ expert_map=None if self.disable_expert_map else layer.expert_map,
1265
ab_strides1=self.ab_strides1_c_strides2,
1266
ab_strides2=self.ab_strides2,
1267
c_strides1=self.c_strides1,
@@ -1314,10 +1283,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1283
topk_weights=topk_weights,
1284
topk_ids=topk_ids,
1285
inplace=True,
1317
- activation=activation,
1318
- apply_router_weight_on_input=apply_router_weight_on_input,
1319
- global_num_experts=global_num_experts,
1320
- expert_map=expert_map,
1286
+ activation=layer.activation,
1287
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1288
+ global_num_experts=layer.global_num_experts,
1289
+ expert_map=layer.expert_map,
1290
quant_config=self.moe_quant_config,
1291
)
1292
@@ -1437,23 +1406,6 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1406
layer: FusedMoE,
1407
x: torch.Tensor,
1408
router_logits: torch.Tensor,
1440
- top_k: int,
1441
- renormalize: bool,
1442
- use_grouped_topk: bool = False,
1443
- topk_group: int | None = None,
1444
- num_expert_group: int | None = None,
1445
- global_num_experts: int = -1,
1446
- expert_map: torch.Tensor | None = None,
1447
- custom_routing_function: Callable | None = None,
1448
- scoring_func: str = "softmax",
1449
- routed_scaling_factor: float = 1.0,
1450
- e_score_correction_bias: torch.Tensor | None = None,
1451
- apply_router_weight_on_input: bool = False,
1452
- activation: str = "silu",
1453
- enable_eplb: bool = False,
1454
- expert_load_view: torch.Tensor | None = None,
1455
- logical_to_physical_map: torch.Tensor | None = None,
1456
- logical_replica_count: torch.Tensor | None = None,
1409
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1410
from vllm.model_executor.layers.fused_moe import fused_experts
1411
@@ -1469,10 +1421,10 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1421
topk_weights=topk_weights,
1422
topk_ids=topk_ids,
1423
inplace=True,
1472
- activation=activation,
1473
- apply_router_weight_on_input=apply_router_weight_on_input,
1474
- global_num_experts=global_num_experts,
1475
- expert_map=expert_map,
1424
+ activation=layer.activation,
1425
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1426
+ global_num_experts=layer.global_num_experts,
1427
+ expert_map=layer.expert_map,
1428
quant_config=self.moe_quant_config,
1429
)
1430
@@ -1814,25 +1766,10 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1766
layer: FusedMoE,
1767
x: torch.Tensor,
1768
router_logits: torch.Tensor,
1817
- top_k: int,
1818
- renormalize: bool,
1819
- use_grouped_topk: bool = False,
1820
- topk_group: int | None = None,
1821
- num_expert_group: int | None = None,
1822
- global_num_experts: int = -1,
1823
- expert_map: torch.Tensor | None = None,
1824
- custom_routing_function: Callable | None = None,
1825
- scoring_func: str = "softmax",
1826
- routed_scaling_factor: float = 1.0,
1827
- e_score_correction_bias: torch.Tensor | None = None,
1828
- apply_router_weight_on_input: bool = False,
1829
- activation: str = "silu",
1830
- enable_eplb: bool = False,
1831
- expert_load_view: torch.Tensor | None = None,
1832
- logical_to_physical_map: torch.Tensor | None = None,
1833
- logical_replica_count: torch.Tensor | None = None,
1769
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1835
- assert activation == "silu", f"{activation} not supported for Marlin MoE."
1770
+ assert layer.activation == "silu", (
1771
+ f"{layer.activation} not supported for Marlin MoE."
1772
+ )
1773
1774
topk_weights, topk_ids, _ = layer.select_experts(
1775
hidden_states=x,
@@ -1853,9 +1790,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1790
input_global_scale1=getattr(layer, "w13_input_global_scale", None),
1791
input_global_scale2=getattr(layer, "w2_input_global_scale", None),
1792
quant_type_id=self.quant_type.id,
1856
- apply_router_weight_on_input=apply_router_weight_on_input,
1857
- global_num_experts=global_num_experts,
1858
- expert_map=expert_map,
1793
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1794
+ global_num_experts=layer.global_num_experts,
1795
+ expert_map=layer.expert_map,
1796
g_idx1=layer.w13_weight_g_idx,
1797
g_idx2=layer.w2_weight_g_idx,
1798
sort_indices1=layer.w13_g_idx_sort_indices,
@@ -2057,23 +1994,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1994
layer: FusedMoE,
1995
x: torch.Tensor,
1996
router_logits: torch.Tensor,
2060
- top_k: int,
2061
- renormalize: bool,
2062
- use_grouped_topk: bool = False,
2063
- topk_group: int | None = None,
2064
- num_expert_group: int | None = None,
2065
- global_num_experts: int = -1,
2066
- expert_map: torch.Tensor | None = None,
2067
- custom_routing_function: Callable | None = None,
2068
- scoring_func: str = "softmax",
2069
- routed_scaling_factor: float = 1.0,
2070
- e_score_correction_bias: torch.Tensor | None = None,
2071
- apply_router_weight_on_input: bool = False,
2072
- activation: str = "silu",
2073
- enable_eplb: bool = False,
2074
- expert_load_view: torch.Tensor | None = None,
2075
- logical_to_physical_map: torch.Tensor | None = None,
2076
- logical_replica_count: torch.Tensor | None = None,
1997
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1998
from vllm.model_executor.layers.fused_moe import fused_experts
1999
@@ -2089,10 +2009,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
2009
topk_weights=topk_weights,
2010
topk_ids=topk_ids,
2011
inplace=True,
2092
- activation=activation,
2093
- apply_router_weight_on_input=apply_router_weight_on_input,
2094
- global_num_experts=global_num_experts,
2095
- expert_map=expert_map,
2012
+ activation=layer.activation,
2013
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
2014
+ global_num_experts=layer.global_num_experts,
2015
+ expert_map=layer.expert_map,
2016
quant_config=self.moe_quant_config,
2017
)
2018
@@ -2372,32 +2292,15 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
2292
2293
def apply(
2294
self,
2375
- layer: torch.nn.Module,
2295
+ layer: FusedMoE,
2296
x: torch.Tensor,
2297
router_logits: torch.Tensor,
2378
- top_k: int,
2379
- renormalize: bool,
2380
- use_grouped_topk: bool = False,
2381
- topk_group: int | None = None,
2382
- num_expert_group: int | None = None,
2383
- global_num_experts: int = -1,
2384
- expert_map: torch.Tensor | None = None,
2385
- custom_routing_function: Callable | None = None,
2386
- scoring_func: str = "softmax",
2387
- routed_scaling_factor: float = 1.0,
2388
- e_score_correction_bias: torch.Tensor | None = None,
2389
- apply_router_weight_on_input: bool = False,
2390
- activation: str = "silu",
2391
- enable_eplb: bool = False,
2392
- expert_load_view: torch.Tensor | None = None,
2393
- logical_to_physical_map: torch.Tensor | None = None,
2394
- logical_replica_count: torch.Tensor | None = None,
2298
) -> torch.Tensor:
2396
- assert not enable_eplb, "EPLB not supported for W4A8-int MoE yet."
2397
- assert activation in ("silu", "swigluoai", "swiglu"), (
2299
+ assert not layer.enable_eplb, "EPLB not supported for W4A8-int MoE yet."
2300
+ assert layer.activation in ("silu", "swigluoai", "swiglu"), (
2301
"Only SiLU/SwiGLUGU/SwiGLUUG are supported."
2302
)
2400
- assert expert_map is None, """expert_map/EP not implemented
2303
+ assert layer.expert_map is None, """expert_map/EP not implemented
2304
for CPU dyn-4bit MoE."""
2305
2306
def _act_kind(s: str) -> int:
@@ -2414,15 +2317,9 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
2317
topk_weights, topk_ids = select_experts(
2318
hidden_states=x,
2319
router_logits=router_logits,
2417
- use_grouped_topk=use_grouped_topk,
2418
- top_k=top_k,
2419
- renormalize=renormalize,
2420
- topk_group=topk_group,
2421
- num_expert_group=num_expert_group,
2422
- custom_routing_function=custom_routing_function,
2423
- scoring_func=scoring_func,
2424
- routed_scaling_factor=routed_scaling_factor,
2425
- e_score_correction_bias=e_score_correction_bias,
2320
+ top_k=layer.top_k,
2321
+ use_grouped_topk=layer.use_grouped_topk,
2322
+ renormalize=layer.renormalize,
2323
)
2324
2325
return torch.ops._C.dynamic_4bit_int_moe(
@@ -2435,8 +2332,8 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
2332
layer.w2_in_features,
2333
layer.w13_out_features,
2334
layer.group_size,
2438
- apply_router_weight_on_input,
2439
- int(_act_kind(activation)),
2335
+ layer.apply_router_weight_on_input,
2336
+ int(_act_kind(layer.activation)),
2337
)
2338
2339
@@ -2707,28 +2604,11 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
2604
2605
def apply(
2606
self,
2710
- layer: torch.nn.Module,
2607
+ layer: FusedMoE,
2608
x: torch.Tensor,
2609
router_logits: torch.Tensor,
2713
- top_k: int,
2714
- renormalize: bool,
2715
- use_grouped_topk: bool = False,
2716
- topk_group: int | None = None,
2717
- num_expert_group: int | None = None,
2718
- global_num_experts: int = -1,
2719
- expert_map: torch.Tensor | None = None,
2720
- custom_routing_function: Callable | None = None,
2721
- scoring_func: str = "softmax",
2722
- routed_scaling_factor: float = 1.0,
2723
- e_score_correction_bias: torch.Tensor | None = None,
2724
- apply_router_weight_on_input: bool = False,
2725
- activation: str = "silu",
2726
- enable_eplb: bool = False,
2727
- expert_load_view: torch.Tensor | None = None,
2728
- logical_to_physical_map: torch.Tensor | None = None,
2729
- logical_replica_count: torch.Tensor | None = None,
2610
):
2731
- if enable_eplb:
2611
+ if layer.enable_eplb:
2612
raise NotImplementedError(
2613
"EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."
2614
)
@@ -2749,9 +2629,9 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
2629
topk_weights,
2630
topk_ids,
2631
quant_config=self.moe_quant_config,
2752
- activation=activation,
2753
- global_num_experts=global_num_experts,
2754
- expert_map=None if self.disable_expert_map else expert_map,
2632
+ activation=layer.activation,
2633
+ global_num_experts=layer.global_num_experts,
2634
+ expert_map=None if self.disable_expert_map else layer.expert_map,
2635
a_strides1=self.a_strides1_c_strides2,
2636
a_strides2=self.a_strides2,
2637
b_strides1=self.b_strides1,
vllm/model_executor/layers/quantization/experts_int8.py
+4
-22
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from typing import Any, Optional
5
6
import torch
@@ -140,23 +139,6 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
139
layer: FusedMoE,
140
x: torch.Tensor,
141
router_logits: torch.Tensor,
143
- top_k: int,
144
- renormalize: bool,
145
- use_grouped_topk: bool = False,
146
- topk_group: int | None = None,
147
- num_expert_group: int | None = None,
148
- global_num_experts: int = -1,
149
- expert_map: torch.Tensor | None = None,
150
- custom_routing_function: Callable | None = None,
151
- scoring_func: str = "softmax",
152
- routed_scaling_factor: float = 1.0,
153
- e_score_correction_bias: torch.Tensor | None = None,
154
- apply_router_weight_on_input: bool = False,
155
- activation: str = "silu",
156
- enable_eplb: bool = False,
157
- expert_load_view: torch.Tensor | None = None,
158
- logical_to_physical_map: torch.Tensor | None = None,
159
- logical_replica_count: torch.Tensor | None = None,
142
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
143
from vllm.model_executor.layers.fused_moe import fused_experts
144
@@ -172,10 +154,10 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
154
topk_weights=topk_weights,
155
topk_ids=topk_ids,
156
inplace=True,
175
- activation=activation,
176
- apply_router_weight_on_input=apply_router_weight_on_input,
177
- global_num_experts=global_num_experts,
178
- expert_map=expert_map,
157
+ activation=layer.activation,
158
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
159
+ global_num_experts=layer.global_num_experts,
160
+ expert_map=layer.expert_map,
161
quant_config=self.moe_quant_config,
162
)
163
vllm/model_executor/layers/quantization/fp8.py
+44
-60
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from enum import Enum
5
from functools import partial
6
from typing import TYPE_CHECKING, Any, Optional
@@ -1242,41 +1241,20 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1241
layer: FusedMoE,
1242
x: torch.Tensor,
1243
router_logits: torch.Tensor,
1245
- top_k: int,
1246
- renormalize: bool,
1247
- use_grouped_topk: bool = False,
1248
- topk_group: int | None = None,
1249
- num_expert_group: int | None = None,
1250
- global_num_experts: int = -1,
1251
- expert_map: torch.Tensor | None = None,
1252
- custom_routing_function: Callable | None = None,
1253
- scoring_func: str = "softmax",
1254
- routed_scaling_factor: float = 1.0,
1255
- e_score_correction_bias: torch.Tensor | None = None,
1256
- apply_router_weight_on_input: bool = False,
1257
- activation: str = "silu",
1258
- enable_eplb: bool = False,
1259
- expert_load_view: torch.Tensor | None = None,
1260
- logical_to_physical_map: torch.Tensor | None = None,
1261
- logical_replica_count: torch.Tensor | None = None,
1244
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1263
- if enable_eplb:
1264
- assert expert_load_view is not None
1265
- assert logical_to_physical_map is not None
1266
- assert logical_replica_count is not None
1267
- assert isinstance(layer, FusedMoE)
1268
-
1245
if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
1270
- assert activation == "silu", (
1271
- f"Expected 'silu' activation but got {activation}"
1246
+ if layer.enable_eplb:
1247
+ raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.")
1248
+ assert layer.activation == "silu", (
1249
+ f"Expected 'silu' activation but got {layer.activation}"
1250
)
1251
1252
if self.block_quant:
1253
import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
1254
1255
e_score_correction_bias = (
1278
- e_score_correction_bias.to(x.dtype)
1279
- if e_score_correction_bias is not None
1256
+ layer.e_score_correction_bias.to(x.dtype)
1257
+ if layer.e_score_correction_bias is not None
1258
else None
1259
)
1260
routing_method_type = layer.routing_method_type
@@ -1290,29 +1268,31 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1268
w13_weight_scale_inv=layer.w13_weight_scale_inv,
1269
w2_weight=layer.w2_weight,
1270
w2_weight_scale_inv=layer.w2_weight_scale_inv,
1293
- global_num_experts=global_num_experts,
1294
- top_k=top_k,
1295
- num_expert_group=num_expert_group,
1296
- topk_group=topk_group,
1271
+ global_num_experts=layer.global_num_experts,
1272
+ top_k=layer.top_k,
1273
+ num_expert_group=layer.num_expert_group,
1274
+ topk_group=layer.topk_group,
1275
intermediate_size=layer.intermediate_size_per_partition,
1276
expert_offset=layer.ep_rank * layer.local_num_experts,
1277
local_num_experts=layer.local_num_experts,
1278
block_shape=self.weight_block_size,
1279
routing_method_type=routing_method_type,
1302
- routed_scaling=routed_scaling_factor,
1280
+ routed_scaling=layer.routed_scaling_factor,
1281
)
1282
else:
1305
- assert not renormalize and custom_routing_function is not None
1283
+ assert (
1284
+ not layer.renormalize and layer.custom_routing_function is not None
1285
+ )
1286
result = apply_flashinfer_per_tensor_scale_fp8(
1287
layer=layer,
1288
hidden_states=x,
1289
router_logits=router_logits,
1310
- routing_bias=e_score_correction_bias,
1311
- global_num_experts=global_num_experts,
1312
- top_k=top_k,
1313
- num_expert_group=num_expert_group,
1314
- topk_group=topk_group,
1315
- apply_router_weight_on_input=apply_router_weight_on_input,
1290
+ routing_bias=layer.e_score_correction_bias,
1291
+ global_num_experts=layer.global_num_experts,
1292
+ top_k=layer.top_k,
1293
+ num_expert_group=layer.num_expert_group,
1294
+ topk_group=layer.topk_group,
1295
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1296
)
1297
1298
select_result = layer.select_experts(
@@ -1333,13 +1313,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1313
layer.w2_weight,
1314
topk_weights=topk_weights,
1315
topk_ids=topk_ids,
1336
- activation=activation,
1337
- apply_router_weight_on_input=apply_router_weight_on_input,
1338
- expert_map=expert_map,
1316
+ activation=layer.activation,
1317
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1318
+ expert_map=layer.expert_map,
1319
quant_config=self.moe_quant_config,
1320
)
1321
elif self.use_marlin:
1342
- assert activation == "silu", f"{activation} not supported for Marlin MoE."
1322
+ assert layer.activation == "silu", (
1323
+ f"{layer.activation} not supported for Marlin MoE."
1324
+ )
1325
result = fused_marlin_moe(
1326
x,
1327
layer.w13_weight,
@@ -1352,20 +1334,22 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1334
topk_weights,
1335
topk_ids,
1336
quant_type_id=scalar_types.float8_e4m3fn.id,
1355
- apply_router_weight_on_input=apply_router_weight_on_input,
1356
- global_num_experts=global_num_experts,
1357
- expert_map=expert_map,
1337
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1338
+ global_num_experts=layer.global_num_experts,
1339
+ expert_map=layer.expert_map,
1340
input_dtype=self.marlin_input_dtype,
1341
workspace=layer.workspace,
1342
)
1343
elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
1362
- assert activation == "silu", (
1363
- f"Expected 'silu' activation but got {activation}"
1344
+ assert layer.activation == "silu", (
1345
+ f"Expected 'silu' activation but got {layer.activation}"
1346
)
1347
if not self.block_quant:
1366
- assert not renormalize and custom_routing_function is not None
1367
- assert scoring_func == "sigmoid", (
1368
- f"Expected 'sigmoid' scoring func but got {scoring_func}"
1348
+ assert (
1349
+ not layer.renormalize and layer.custom_routing_function is not None
1350
+ )
1351
+ assert layer.scoring_func == "sigmoid", (
1352
+ f"Expected 'sigmoid' scoring func but got {layer.scoring_func}"
1353
)
1354
# Delegate to CUTLASS FlashInfer path; function already bound with
1355
# use_deepseek_fp8_block_scale for block-quant when applicable
@@ -1375,10 +1359,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1359
topk_weights,
1360
topk_ids,
1361
inplace=False,
1378
- activation=activation,
1379
- global_num_experts=global_num_experts,
1380
- expert_map=expert_map,
1381
- apply_router_weight_on_input=apply_router_weight_on_input,
1362
+ activation=layer.activation,
1363
+ global_num_experts=layer.global_num_experts,
1364
+ expert_map=layer.expert_map,
1365
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1366
)
1367
else:
1368
from vllm.model_executor.layers.fused_moe import fused_experts
@@ -1390,10 +1374,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1374
topk_weights=topk_weights,
1375
topk_ids=topk_ids,
1376
inplace=True,
1393
- activation=activation,
1394
- global_num_experts=global_num_experts,
1395
- apply_router_weight_on_input=apply_router_weight_on_input,
1396
- expert_map=expert_map,
1377
+ activation=layer.activation,
1378
+ global_num_experts=layer.global_num_experts,
1379
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1380
+ expert_map=layer.expert_map,
1381
quant_config=self.moe_quant_config,
1382
allow_deep_gemm=self.allow_deep_gemm,
1383
allow_cutlass_block_scaled_grouped_gemm=(
vllm/model_executor/layers/quantization/gguf.py
+4
-21
@@ -1,7 +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, Mapping
4
+from collections.abc import Mapping
5
from types import MappingProxyType
6
from typing import Any, Optional
7
@@ -625,26 +625,9 @@ class GGUFMoEMethod(FusedMoEMethodBase):
625
layer: FusedMoE,
626
x: torch.Tensor,
627
router_logits: torch.Tensor,
628
- top_k: int,
629
- renormalize: bool,
630
- use_grouped_topk: bool = False,
631
- topk_group: int | None = None,
632
- num_expert_group: int | None = None,
633
- global_num_experts: int = -1,
634
- expert_map: torch.Tensor | None = None,
635
- custom_routing_function: Callable | None = None,
636
- scoring_func: str = "softmax",
637
- routed_scaling_factor: float = 1.0,
638
- e_score_correction_bias: torch.Tensor | None = None,
639
- apply_router_weight_on_input: bool = False,
640
- activation: str = "silu",
641
- enable_eplb: bool = False,
642
- expert_load_view: torch.Tensor | None = None,
643
- logical_to_physical_map: torch.Tensor | None = None,
644
- logical_replica_count: torch.Tensor | None = None,
628
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
646
- assert activation == "silu", "Only SiLU activation is supported."
647
- if apply_router_weight_on_input:
629
+ assert layer.activation == "silu", "Only SiLU activation is supported."
630
+ if layer.apply_router_weight_on_input:
631
raise NotImplementedError(
632
"Apply router weight on input is not supported for"
633
"fused GGUF MoE method."
@@ -662,7 +645,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):
645
topk_ids,
646
layer.w13_qweight_type.weight_type,
647
layer.w2_qweight_type.weight_type,
665
- activation,
648
+ layer.activation,
649
)
650
651
vllm/model_executor/layers/quantization/gptq_marlin.py
+4
-22
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from copy import deepcopy
5
from typing import Any, Optional
6
@@ -790,25 +789,8 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
789
layer: FusedMoE,
790
x: torch.Tensor,
791
router_logits: torch.Tensor,
793
- top_k: int,
794
- renormalize: bool,
795
- use_grouped_topk: bool = False,
796
- topk_group: int | None = None,
797
- num_expert_group: int | None = None,
798
- global_num_experts: int = -1,
799
- expert_map: torch.Tensor | None = None,
800
- custom_routing_function: Callable | None = None,
801
- scoring_func: str = "softmax",
802
- routed_scaling_factor: float = 1.0,
803
- e_score_correction_bias: torch.Tensor | None = None,
804
- apply_router_weight_on_input: bool = False,
805
- activation: str = "silu",
806
- enable_eplb: bool = False,
807
- expert_load_view: torch.Tensor | None = None,
808
- logical_to_physical_map: torch.Tensor | None = None,
809
- logical_replica_count: torch.Tensor | None = None,
792
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
811
- assert activation == "silu", "Only SiLU activation is supported."
793
+ assert layer.activation == "silu", "Only SiLU activation is supported."
794
795
topk_weights, topk_ids, _ = layer.select_experts(
796
hidden_states=x,
@@ -829,9 +811,9 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
811
input_global_scale1=getattr(layer, "w13_input_global_scale", None),
812
input_global_scale2=getattr(layer, "w2_input_global_scale", None),
813
quant_type_id=self.quant_type.id,
832
- apply_router_weight_on_input=apply_router_weight_on_input,
833
- global_num_experts=global_num_experts,
834
- expert_map=expert_map,
814
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
815
+ global_num_experts=layer.global_num_experts,
816
+ expert_map=layer.expert_map,
817
g_idx1=layer.w13_g_idx,
818
g_idx2=layer.w2_g_idx,
819
sort_indices1=layer.w13_g_idx_sort_indices,
vllm/model_executor/layers/quantization/ipex_quant.py
+6
-24
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from typing import Any, Optional
5
6
import torch
@@ -440,31 +439,14 @@ class XPUFp8MoEMethod(FusedMoEMethodBase):
439
layer: torch.nn.Module,
440
x: torch.Tensor,
441
router_logits: torch.Tensor,
443
- top_k: int,
444
- renormalize: bool,
445
- use_grouped_topk: bool = False,
446
- topk_group: int | None = None,
447
- num_expert_group: int | None = None,
448
- global_num_experts: int = -1,
449
- expert_map: torch.Tensor | None = None,
450
- custom_routing_function: Callable | None = None,
451
- scoring_func: str = "softmax",
452
- routed_scaling_factor: float = 1.0,
453
- e_score_correction_bias: torch.Tensor | None = None,
454
- apply_router_weight_on_input: bool = False,
455
- activation: str = "silu",
456
- enable_eplb: bool = False,
457
- expert_load_view: torch.Tensor | None = None,
458
- logical_to_physical_map: torch.Tensor | None = None,
459
- logical_replica_count: torch.Tensor | None = None,
442
) -> torch.Tensor:
443
return layer.ipex_fusion(
444
x,
463
- use_grouped_topk,
464
- top_k,
445
+ layer.use_grouped_topk,
446
+ layer.top_k,
447
router_logits,
466
- renormalize,
467
- topk_group,
468
- num_expert_group,
469
- custom_routing_function=custom_routing_function,
448
+ layer.renormalize,
449
+ layer.topk_group,
450
+ layer.num_expert_group,
451
+ custom_routing_function=layer.custom_routing_function,
452
)
vllm/model_executor/layers/quantization/modelopt.py
+36
-70
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from fnmatch import fnmatch
5
from typing import TYPE_CHECKING, Any, Optional
6
@@ -707,43 +706,27 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
706
layer: FusedMoE,
707
x: torch.Tensor,
708
router_logits: torch.Tensor,
710
- top_k: int,
711
- renormalize: bool,
712
- use_grouped_topk: bool = False,
713
- topk_group: int | None = None,
714
- num_expert_group: int | None = None,
715
- global_num_experts: int = -1,
716
- expert_map: torch.Tensor | None = None,
717
- custom_routing_function: Callable | None = None,
718
- scoring_func: str = "softmax",
719
- routed_scaling_factor: float = 1.0,
720
- e_score_correction_bias: torch.Tensor | None = None,
721
- apply_router_weight_on_input: bool = False,
722
- activation: str = "silu",
723
- enable_eplb: bool = False,
724
- expert_load_view: torch.Tensor | None = None,
725
- logical_to_physical_map: torch.Tensor | None = None,
726
- logical_replica_count: torch.Tensor | None = None,
709
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
710
if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
711
if layer.enable_eplb:
712
raise NotImplementedError(
713
"EPLB not supported for `ModelOptFp8MoEMethod` yet."
714
)
733
- assert activation == "silu", (
734
- f"Expected 'silu' activation but got {activation}"
715
+ assert layer.activation == "silu", (
716
+ f"Expected 'silu' activation but got {layer.activation}"
717
)
736
- assert not renormalize
718
+
719
+ assert not layer.renormalize
720
return apply_flashinfer_per_tensor_scale_fp8(
721
layer=layer,
722
hidden_states=x,
723
router_logits=router_logits,
741
- routing_bias=e_score_correction_bias,
742
- global_num_experts=global_num_experts,
743
- top_k=top_k,
744
- num_expert_group=num_expert_group,
745
- topk_group=topk_group,
746
- apply_router_weight_on_input=apply_router_weight_on_input,
724
+ routing_bias=layer.e_score_correction_bias,
725
+ global_num_experts=layer.global_num_experts,
726
+ top_k=layer.top_k,
727
+ num_expert_group=layer.num_expert_group,
728
+ topk_group=layer.topk_group,
729
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
730
)
731
732
# Expert selection
@@ -753,9 +736,9 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
736
)
737
738
if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
756
- assert activation in ("silu", "relu2_no_mul"), (
739
+ assert layer.activation in ("silu", "relu2_no_mul"), (
740
"Expected activation to be in ('silu', 'relu2_no_mul'),"
758
- f"but got {activation}"
741
+ f"but got {layer.activation}"
742
)
743
return flashinfer_cutlass_moe_fp8(
744
x,
@@ -763,10 +746,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
746
topk_weights,
747
topk_ids,
748
inplace=False,
766
- activation=activation,
767
- global_num_experts=global_num_experts,
768
- expert_map=expert_map,
769
- apply_router_weight_on_input=apply_router_weight_on_input,
749
+ activation=layer.activation,
750
+ global_num_experts=layer.global_num_experts,
751
+ expert_map=layer.expert_map,
752
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
753
)
754
else:
755
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
@@ -780,11 +763,11 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
763
topk_weights=topk_weights,
764
topk_ids=topk_ids,
765
inplace=True,
783
- activation=activation,
766
+ activation=layer.activation,
767
quant_config=self.moe_quant_config,
785
- global_num_experts=global_num_experts,
786
- expert_map=expert_map,
787
- apply_router_weight_on_input=apply_router_weight_on_input,
768
+ global_num_experts=layer.global_num_experts,
769
+ expert_map=layer.expert_map,
770
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
771
)
772
773
@@ -1504,23 +1487,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1487
layer: FusedMoE,
1488
x: torch.Tensor,
1489
router_logits: torch.Tensor,
1507
- top_k: int,
1508
- renormalize: bool,
1509
- use_grouped_topk: bool = False,
1510
- topk_group: int | None = None,
1511
- num_expert_group: int | None = None,
1512
- global_num_experts: int = -1,
1513
- expert_map: torch.Tensor | None = None,
1514
- custom_routing_function: Callable | None = None,
1515
- scoring_func: str = "softmax",
1516
- routed_scaling_factor: float = 1.0,
1517
- e_score_correction_bias: torch.Tensor | None = None,
1518
- apply_router_weight_on_input: bool = False,
1519
- activation: str = "silu",
1520
- enable_eplb: bool = False,
1521
- expert_load_view: torch.Tensor | None = None,
1522
- logical_to_physical_map: torch.Tensor | None = None,
1523
- logical_replica_count: torch.Tensor | None = None,
1490
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1491
if not self.moe.is_act_and_mul:
1492
assert (
@@ -1535,7 +1501,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1501
self.allow_flashinfer
1502
and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
1503
):
1538
- if enable_eplb:
1504
+ if layer.enable_eplb:
1505
raise NotImplementedError(
1506
"EPLB not supported for `ModelOptNvFp4FusedMoE` yet."
1507
)
@@ -1543,12 +1509,12 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1509
layer=layer,
1510
x=x,
1511
router_logits=router_logits,
1546
- top_k=top_k,
1547
- global_num_experts=global_num_experts,
1548
- num_expert_group=num_expert_group,
1549
- topk_group=topk_group,
1550
- custom_routing_function=custom_routing_function,
1551
- e_score_correction_bias=e_score_correction_bias,
1512
+ top_k=layer.top_k,
1513
+ global_num_experts=layer.global_num_experts,
1514
+ num_expert_group=layer.num_expert_group,
1515
+ topk_group=layer.topk_group,
1516
+ custom_routing_function=layer.custom_routing_function,
1517
+ e_score_correction_bias=layer.e_score_correction_bias,
1518
)
1519
1520
topk_weights, topk_ids, _ = layer.select_experts(
@@ -1571,9 +1537,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1537
global_scale1=layer.w13_weight_scale_2,
1538
global_scale2=layer.w2_weight_scale_2,
1539
quant_type_id=scalar_types.float4_e2m1f.id,
1574
- apply_router_weight_on_input=apply_router_weight_on_input,
1575
- global_num_experts=global_num_experts,
1576
- expert_map=expert_map,
1540
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1541
+ global_num_experts=layer.global_num_experts,
1542
+ expert_map=layer.expert_map,
1543
input_dtype=self.marlin_input_dtype,
1544
)
1545
@@ -1604,10 +1570,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1570
topk_ids=topk_ids,
1571
quant_config=self.moe_quant_config,
1572
inplace=False,
1607
- activation=activation,
1608
- global_num_experts=global_num_experts,
1609
- expert_map=expert_map,
1610
- apply_router_weight_on_input=apply_router_weight_on_input,
1573
+ activation=layer.activation,
1574
+ global_num_experts=layer.global_num_experts,
1575
+ expert_map=layer.expert_map,
1576
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1577
)
1578
else:
1579
# If no modular kernel is provided, use cutlass_moe_fp4 for TP case
@@ -1622,8 +1588,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1588
topk_weights=topk_weights,
1589
topk_ids=topk_ids,
1590
quant_config=self.moe_quant_config,
1625
- expert_map=expert_map,
1626
- apply_router_weight_on_input=apply_router_weight_on_input,
1591
+ expert_map=layer.expert_map,
1592
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1593
# TODO: derive from arguments
1594
m=x.shape[0],
1595
n=layer.w2_weight.shape[2] * 2,
vllm/model_executor/layers/quantization/moe_wna16.py
+4
-22
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from typing import Any, Optional
5
6
import torch
@@ -362,27 +361,10 @@ class MoeWNA16Method(FusedMoEMethodBase):
361
layer: FusedMoE,
362
x: torch.Tensor,
363
router_logits: torch.Tensor,
365
- top_k: int,
366
- renormalize: bool,
367
- use_grouped_topk: bool = False,
368
- topk_group: int | None = None,
369
- num_expert_group: int | None = None,
370
- global_num_experts: int = -1,
371
- expert_map: torch.Tensor | None = None,
372
- custom_routing_function: Callable | None = None,
373
- scoring_func: str = "softmax",
374
- routed_scaling_factor: float = 1.0,
375
- e_score_correction_bias: torch.Tensor | None = None,
376
- apply_router_weight_on_input: bool = False,
377
- activation: str = "silu",
378
- enable_eplb: bool = False,
379
- expert_load_view: torch.Tensor | None = None,
380
- logical_to_physical_map: torch.Tensor | None = None,
381
- logical_replica_count: torch.Tensor | None = None,
364
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
365
from vllm.model_executor.layers.fused_moe import fused_experts
366
385
- assert activation == "silu", "Only SiLU activation is supported."
367
+ assert layer.activation == "silu", "Only SiLU activation is supported."
368
topk_weights, topk_ids, _ = layer.select_experts(
369
hidden_states=x,
370
router_logits=router_logits,
@@ -395,9 +377,9 @@ class MoeWNA16Method(FusedMoEMethodBase):
377
topk_weights=topk_weights,
378
topk_ids=topk_ids,
379
inplace=True,
398
- apply_router_weight_on_input=apply_router_weight_on_input,
399
- global_num_experts=global_num_experts,
400
- expert_map=expert_map,
380
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
381
+ global_num_experts=layer.global_num_experts,
382
+ expert_map=layer.expert_map,
383
quant_config=self.moe_quant_config,
384
)
385
vllm/model_executor/layers/quantization/mxfp4.py
+31
-66
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Callable
3
from enum import Enum
4
from typing import Optional
5
@@ -892,25 +891,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
891
layer: FusedMoE,
892
x: torch.Tensor,
893
router_logits: torch.Tensor,
895
- top_k: int,
896
- renormalize: bool,
897
- use_grouped_topk: bool = False,
898
- topk_group: int | None = None,
899
- num_expert_group: int | None = None,
900
- global_num_experts: int = -1,
901
- expert_map: torch.Tensor | None = None,
902
- custom_routing_function: Callable | None = None,
903
- scoring_func: str = "softmax",
904
- routed_scaling_factor: float = 1.0,
905
- e_score_correction_bias: torch.Tensor | None = None,
906
- apply_router_weight_on_input: bool = False,
907
- activation: str = "silu",
908
- enable_eplb: bool = False,
909
- expert_load_view: torch.Tensor | None = None,
910
- logical_to_physical_map: torch.Tensor | None = None,
911
- logical_replica_count: torch.Tensor | None = None,
894
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
913
- if enable_eplb:
895
+ if layer.enable_eplb:
896
raise NotImplementedError("EPLB is not supported for mxfp4")
897
898
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
@@ -933,26 +915,26 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
915
global_scale1=None,
916
global_scale2=None,
917
quant_type_id=scalar_types.float4_e2m1f.id,
936
- apply_router_weight_on_input=apply_router_weight_on_input,
937
- global_num_experts=global_num_experts,
938
- activation=activation,
939
- expert_map=expert_map,
918
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
919
+ global_num_experts=layer.global_num_experts,
920
+ activation=layer.activation,
921
+ expert_map=layer.expert_map,
922
input_dtype=self.marlin_input_dtype,
923
)
924
925
assert _can_support_mxfp4(
944
- use_grouped_topk,
945
- topk_group,
946
- num_expert_group,
947
- expert_map,
948
- custom_routing_function,
949
- e_score_correction_bias,
950
- apply_router_weight_on_input,
951
- scoring_func,
952
- activation,
953
- expert_load_view,
954
- logical_to_physical_map,
955
- logical_replica_count,
926
+ layer.use_grouped_topk,
927
+ layer.topk_group,
928
+ layer.num_expert_group,
929
+ layer.expert_map,
930
+ layer.custom_routing_function,
931
+ layer.e_score_correction_bias,
932
+ layer.apply_router_weight_on_input,
933
+ layer.scoring_func,
934
+ layer.activation,
935
+ layer.expert_load_view,
936
+ layer.logical_to_physical_map,
937
+ layer.logical_replica_count,
938
), "MXFP4 are not supported with this configuration."
939
940
if (
@@ -988,8 +970,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
970
None, # output1_scale_scalar
971
None, # output1_scale_gate_scalar
972
None, # output2_scale_scalar
991
- global_num_experts,
992
- top_k,
973
+ layer.global_num_experts,
974
+ layer.top_k,
975
None, # n_group
976
None, # topk_group
977
self.intermediate_size, # padded to multiple of 256
@@ -997,7 +979,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
979
self.num_experts, # local num experts
980
None,
981
None,
1000
- 1 if renormalize else 0, # routing_method_type, renormalize
982
+ 1 if layer.renormalize else 0, # routing_method_type, renormalize
983
True, # do finalize
984
tune_max_num_tokens=max(self.max_capture_size, 1),
985
)[0]
@@ -1081,12 +1063,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1063
w1=layer.w13_weight,
1064
w2=layer.w2_weight,
1065
gating_output=router_logits,
1084
- topk=top_k,
1085
- renormalize=renormalize,
1086
- global_num_experts=global_num_experts,
1087
- expert_map=expert_map,
1066
+ topk=layer.top_k,
1067
+ renormalize=layer.renormalize,
1068
+ global_num_experts=layer.global_num_experts,
1069
+ expert_map=layer.expert_map,
1070
quant_config=self.moe_quant_config,
1089
- apply_router_weight_on_input=apply_router_weight_on_input,
1071
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1072
)
1073
else:
1074
raise ValueError(f"Unsupported backend: {self.mxfp4_backend}")
@@ -1138,37 +1120,20 @@ class IpexMxfp4MoEMethod(Mxfp4MoEMethod):
1120
layer: torch.nn.Module,
1121
x: torch.Tensor,
1122
router_logits: torch.Tensor,
1141
- top_k: int,
1142
- renormalize: bool,
1143
- use_grouped_topk: bool = False,
1144
- topk_group: int | None = None,
1145
- num_expert_group: int | None = None,
1146
- global_num_experts: int = -1,
1147
- expert_map: torch.Tensor | None = None,
1148
- custom_routing_function: Callable | None = None,
1149
- scoring_func: str = "softmax",
1150
- routed_scaling_factor: float = 1.0,
1151
- e_score_correction_bias: torch.Tensor | None = None,
1152
- apply_router_weight_on_input: bool = False,
1153
- activation: str = "silu",
1154
- enable_eplb: bool = False,
1155
- expert_load_view: torch.Tensor | None = None,
1156
- logical_to_physical_map: torch.Tensor | None = None,
1157
- logical_replica_count: torch.Tensor | None = None,
1123
) -> torch.Tensor:
1159
- assert activation == "swigluoai", (
1124
+ assert layer.activation == "swigluoai", (
1125
"Only swiglu_oai activation is supported for IPEX MXFP4 MoE"
1126
)
1127
hidden_size_pad = round_up(self.original_hidden_size, 128)
1128
x_pad = torch.nn.functional.pad(x, (0, hidden_size_pad - x.size(-1)))
1129
hidden_states = layer.ipex_fusion(
1130
x_pad,
1166
- use_grouped_topk,
1167
- top_k,
1131
+ layer.use_grouped_topk,
1132
+ layer.top_k,
1133
router_logits,
1169
- renormalize,
1170
- topk_group,
1171
- num_expert_group,
1134
+ layer.renormalize,
1135
+ layer.topk_group,
1136
+ layer.num_expert_group,
1137
activation="swiglu_oai",
1138
)
1139
hidden_states = hidden_states[..., : self.original_hidden_size].contiguous()
vllm/model_executor/layers/quantization/quark/quark_moe.py
+20
-52
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
from typing import Any
5
6
import torch
@@ -337,23 +336,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
336
layer: FusedMoE,
337
x: torch.Tensor,
338
router_logits: torch.Tensor,
340
- top_k: int,
341
- renormalize: bool,
342
- use_grouped_topk: bool = False,
343
- topk_group: int | None = None,
344
- num_expert_group: int | None = None,
345
- global_num_experts: int = -1,
346
- expert_map: torch.Tensor | None = None,
347
- custom_routing_function: Callable | None = None,
348
- scoring_func: str = "softmax",
349
- routed_scaling_factor: float = 1.0,
350
- e_score_correction_bias: torch.Tensor | None = None,
351
- apply_router_weight_on_input: bool = False,
352
- activation: str = "silu",
353
- enable_eplb: bool = False,
354
- expert_load_view: torch.Tensor | None = None,
355
- logical_to_physical_map: torch.Tensor | None = None,
356
- logical_replica_count: torch.Tensor | None = None,
339
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
340
topk_weights, topk_ids, _ = layer.select_experts(
341
hidden_states=x,
@@ -371,13 +353,15 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
353
w2=layer.w2_weight,
354
topk_weights=topk_weights,
355
topk_ids=topk_ids,
374
- activation=activation,
375
- apply_router_weight_on_input=apply_router_weight_on_input,
356
+ activation=layer.activation,
357
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
358
quant_config=self.moe_quant_config,
377
- expert_map=expert_map,
359
+ expert_map=layer.expert_map,
360
)
361
elif self.use_marlin:
380
- assert activation == "silu", f"{activation} not supported for Marlin MoE."
362
+ assert layer.activation == "silu", (
363
+ f"{layer.activation} not supported for Marlin MoE."
364
+ )
365
return fused_marlin_moe(
366
x,
367
layer.w13_weight,
@@ -390,9 +374,9 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
374
topk_weights,
375
topk_ids,
376
quant_type_id=scalar_types.float8_e4m3fn.id,
393
- apply_router_weight_on_input=apply_router_weight_on_input,
394
- global_num_experts=global_num_experts,
395
- expert_map=expert_map,
377
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
378
+ global_num_experts=layer.global_num_experts,
379
+ expert_map=layer.expert_map,
380
)
381
else:
382
from vllm.model_executor.layers.fused_moe import fused_experts
@@ -404,10 +388,10 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
388
topk_weights=topk_weights,
389
topk_ids=topk_ids,
390
inplace=True,
407
- activation=activation,
408
- apply_router_weight_on_input=apply_router_weight_on_input,
409
- global_num_experts=global_num_experts,
410
- expert_map=expert_map,
391
+ activation=layer.activation,
392
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
393
+ global_num_experts=layer.global_num_experts,
394
+ expert_map=layer.expert_map,
395
quant_config=self.moe_quant_config,
396
)
397
@@ -597,23 +581,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
581
layer: FusedMoE,
582
x: torch.Tensor,
583
router_logits: torch.Tensor,
600
- top_k: int,
601
- renormalize: bool,
602
- use_grouped_topk: bool = False,
603
- topk_group: int | None = None,
604
- num_expert_group: int | None = None,
605
- global_num_experts: int = -1,
606
- expert_map: torch.Tensor | None = None,
607
- custom_routing_function: Callable | None = None,
608
- scoring_func: str = "softmax",
609
- routed_scaling_factor: float = 1.0,
610
- e_score_correction_bias: torch.Tensor | None = None,
611
- apply_router_weight_on_input: bool = False,
612
- activation: str = "silu",
613
- enable_eplb: bool = False,
614
- expert_load_view: torch.Tensor | None = None,
615
- logical_to_physical_map: torch.Tensor | None = None,
616
- logical_replica_count: torch.Tensor | None = None,
584
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
585
topk_weights, topk_ids, _ = layer.select_experts(
586
hidden_states=x,
@@ -631,9 +598,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
598
layer.w2_weight,
599
topk_weights=topk_weights,
600
topk_ids=topk_ids,
634
- activation=activation,
601
+ activation=layer.activation,
602
quant_config=self.moe_quant_config,
636
- expert_map=expert_map,
603
+ expert_map=layer.expert_map,
604
)
605
else:
606
from vllm.model_executor.layers.fused_moe import fused_experts
@@ -645,10 +612,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
612
topk_weights=topk_weights,
613
topk_ids=topk_ids,
614
inplace=True,
648
- activation=activation,
649
- global_num_experts=global_num_experts,
650
- apply_router_weight_on_input=apply_router_weight_on_input,
651
- expert_map=expert_map,
615
+ activation=layer.activation,
616
+ global_num_experts=layer.global_num_experts,
617
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
618
+ expert_map=layer.expert_map,
619
quant_config=self.moe_quant_config,
620
)
621
+
622
return out
vllm/model_executor/layers/quantization/rtn.py
+3
-21
@@ -3,7 +3,6 @@
3
# Copyright © 2025, Oracle and/or its affiliates.
4
5
import os
6
-from collections.abc import Callable
6
from typing import Any, Optional
7
8
import numpy as np
@@ -359,23 +358,6 @@ class RTNMoEMethod(FusedMoEMethodBase):
358
layer: FusedMoE,
359
x: torch.Tensor,
360
router_logits: torch.Tensor,
362
- top_k: int,
363
- renormalize: bool,
364
- use_grouped_topk: bool = False,
365
- topk_group: int | None = None,
366
- num_expert_group: int | None = None,
367
- global_num_experts: int = -1,
368
- expert_map: torch.Tensor | None = None,
369
- custom_routing_function: Callable | None = None,
370
- scoring_func: str = "softmax",
371
- routed_scaling_factor: float = 1.0,
372
- e_score_correction_bias: torch.Tensor | None = None,
373
- apply_router_weight_on_input: bool = False,
374
- activation: str = "silu",
375
- enable_eplb: bool = False,
376
- expert_load_view: torch.Tensor | None = None,
377
- logical_to_physical_map: torch.Tensor | None = None,
378
- logical_replica_count: torch.Tensor | None = None,
361
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
362
topk_weights, topk_ids, _ = layer.select_experts(
363
hidden_states=x,
@@ -394,9 +376,9 @@ class RTNMoEMethod(FusedMoEMethodBase):
376
topk_weights,
377
topk_ids,
378
quant_type_id=self.quant_config.quant_type.id,
397
- apply_router_weight_on_input=apply_router_weight_on_input,
398
- global_num_experts=global_num_experts,
399
- expert_map=expert_map,
379
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
380
+ global_num_experts=layer.global_num_experts,
381
+ expert_map=layer.expert_map,
382
workspace=workspace,
383
)
384