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 3, 20:15 UTC
[MoE Refactor] Introduce RoutedExperts alias for FusedMoE and don't store SharedExperts in MK (#40735)
Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
bnellnm committed
May 12, 2026 at 13:37 UTC
4d591db470c0b53304f9dd2369d4feb8275a94bc
53 files changed
+369
-332
tests/kernels/moe/test_zero_expert_moe.py
+1
@@ -183,6 +183,7 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
183
x=hidden_states,
184
topk_weights=topk_weights,
185
topk_ids=topk_ids,
186
+ shared_experts=None,
187
shared_experts_input=None,
188
)
189
vllm/lora/layers/fused_moe.py
-3
@@ -40,9 +40,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
40
self.base_layer.ensure_moe_quant_config_init()
41
if getattr(self.base_layer.quant_method, "supports_internal_mk", False):
42
moe_kernel = self.base_layer.quant_method.moe_kernel
43
- # Don't let the kernel own shared experts so the runner can
44
- # overlap them with routed experts via a separate CUDA stream.
45
- moe_kernel.shared_experts = None
43
else:
44
prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular()
45
moe_kernel = FusedMoEKernel(
vllm/model_executor/layers/fused_moe/__init__.py
+14
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from contextlib import contextmanager
5
-from typing import Any
5
+from typing import Any, TypeAlias
6
7
from vllm.model_executor.layers.fused_moe.activation import (
8
MoEActivation,
@@ -11,6 +11,8 @@ from vllm.model_executor.layers.fused_moe.activation import (
11
)
12
from vllm.model_executor.layers.fused_moe.config import (
13
FusedMoEConfig,
14
+ FusedMoEParallelConfig,
15
+ FusedMoEQuantConfig,
16
RoutingMethodType,
17
)
18
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
@@ -30,6 +32,9 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
32
FusedMoERouter,
33
)
34
from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
35
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
36
+ SharedExperts,
37
+)
38
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
39
UnquantizedFusedMoEMethod,
40
)
@@ -38,6 +43,10 @@ from vllm.triton_utils import HAS_TRITON
43
_config: dict[str, Any] | None = None
44
45
46
+# Temporary alias for FusedMoE, eventually we be its own class.
47
+RoutedExperts: TypeAlias = FusedMoE
48
+
49
+
50
@contextmanager
51
def override_config(config):
52
global _config
@@ -55,6 +64,8 @@ __all__ = [
64
"FusedMoE",
65
"FusedMoERouter",
66
"FusedMoEConfig",
67
+ "FusedMoEQuantConfig",
68
+ "FusedMoEParallelConfig",
69
"FusedMoEMethodBase",
70
"MoEActivation",
71
"UnquantizedFusedMoEMethod",
@@ -64,6 +75,8 @@ __all__ = [
75
"FusedMoEPrepareAndFinalizeModular",
76
"GateLinear",
77
"RoutingMethodType",
78
+ "RoutedExperts",
79
+ "SharedExperts",
80
"activation_without_mul",
81
"apply_moe_activation",
82
"fused_moe_make_expert_params_mapping",
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
-1
@@ -1489,7 +1489,6 @@ def cutlass_moe_w4a8_fp8(
1489
quant_config=quant_config,
1490
group_size=group_size,
1491
),
1492
- shared_experts=None,
1492
inplace=False,
1493
)
1494
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+6
-2
@@ -4,7 +4,7 @@
4
5
import json
6
import math
7
-from typing import Any
7
+from typing import TYPE_CHECKING, Any
8
9
import torch
10
from humming import dtypes
@@ -38,6 +38,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
38
from vllm.platforms import current_platform
39
from vllm.v1.worker.workspace import current_workspace_manager
40
41
+if TYPE_CHECKING:
42
+ from vllm.model_executor.layers.fused_moe import RoutedExperts
43
+
44
+
45
logger = init_logger(__name__)
46
47
@@ -58,7 +62,7 @@ def get_humming_moe_gemm_type() -> str:
62
class HummingExpertsBase(mk.FusedMoEExpertsModular):
63
def __init__(
64
self,
61
- layer: torch.nn.Module,
65
+ layer: "RoutedExperts",
66
moe_config: FusedMoEConfig,
67
quant_config: FusedMoEQuantConfig,
68
max_num_tokens: int | None = None,
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+11
-4
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from abc import abstractmethod
5
+from typing import TYPE_CHECKING
6
7
import torch
8
@@ -20,6 +21,9 @@ from vllm.model_executor.layers.quantization.base_config import (
21
QuantizeMethodBase,
22
)
23
24
+if TYPE_CHECKING:
25
+ from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts
26
+
27
logger = init_logger(__name__)
28
29
@@ -37,10 +41,12 @@ class FusedMoEMethodBase(QuantizeMethodBase):
41
return self.moe_kernel is not None
42
43
@property
40
- def mk_owns_shared_expert(self) -> bool:
44
+ def mk_can_overlap_shared_experts(self) -> bool:
45
# NOTE(rob): temporary attribute to indicate support for
46
# completed migration to the new internal MK interface.
43
- return self.moe_kernel is not None and self.moe_kernel.owns_shared_experts
47
+ return (
48
+ self.moe_kernel is not None and self.moe_kernel.can_overlap_shared_experts
49
+ )
50
51
@abstractmethod
52
def create_weights(
@@ -156,17 +162,18 @@ class FusedMoEMethodBase(QuantizeMethodBase):
162
163
def apply(
164
self,
159
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
165
+ layer: "RoutedExperts", # type: ignore[name-defined] # noqa: F821
166
x: torch.Tensor,
167
topk_weights: torch.Tensor,
168
topk_ids: torch.Tensor,
169
+ shared_experts: "SharedExperts | None",
170
shared_experts_input: torch.Tensor | None,
171
) -> torch.Tensor:
172
raise NotImplementedError
173
174
def apply_monolithic(
175
self,
169
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
176
+ layer: "RoutedExperts", # type: ignore[name-defined] # noqa: F821
177
x: torch.Tensor,
178
router_logits: torch.Tensor,
179
input_ids: torch.Tensor | None = None,
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+3
-3
@@ -47,7 +47,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
47
moe_layer: torch.nn.Module,
48
old_quant_method: FusedMoEMethodBase,
49
prepare_finalize: FusedMoEPrepareAndFinalizeModular,
50
- shared_experts: SharedExperts | None,
50
inplace: bool = False,
51
) -> "FusedMoEModularMethod":
52
return FusedMoEModularMethod(
@@ -55,7 +54,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
54
FusedMoEKernel(
55
prepare_finalize,
56
old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),
58
- shared_experts=shared_experts,
57
inplace=inplace,
58
),
59
)
@@ -86,10 +84,11 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
84
85
def apply(
86
self,
89
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
87
+ layer: "RoutedExperts", # type: ignore[name-defined] # noqa: F821
88
x: torch.Tensor,
89
topk_weights: torch.Tensor,
90
topk_ids: torch.Tensor,
91
+ shared_experts: SharedExperts | None,
92
shared_experts_input: torch.Tensor | None,
93
) -> torch.Tensor:
94
assert self.moe_kernel is not None
@@ -103,5 +102,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
102
global_num_experts=layer.global_num_experts,
103
apply_router_weight_on_input=layer.apply_router_weight_on_input,
104
expert_map=None if self.disable_expert_map else layer.expert_map,
105
+ shared_experts=shared_experts,
106
shared_experts_input=shared_experts_input,
107
)
vllm/model_executor/layers/fused_moe/layer.py
-1
@@ -480,7 +480,6 @@ class FusedMoE(PluggableLayer):
480
self,
481
self.base_quant_method,
482
prepare_finalize,
483
- self.shared_experts,
483
inplace=not self.moe_config.disable_inplace,
484
)
485
)
vllm/model_executor/layers/fused_moe/modular_kernel.py
+14
-15
@@ -1018,17 +1018,10 @@ class FusedMoEKernelModularImpl:
1018
self,
1019
prepare_finalize: FusedMoEPrepareAndFinalizeModular,
1020
fused_experts: FusedMoEExpertsModular,
1021
- shared_experts: SharedExperts | None,
1021
inplace: bool = False,
1022
):
1023
self.prepare_finalize = prepare_finalize
1024
self.fused_experts = fused_experts
1026
- # Only accept shared experts if they can be run w/async.
1027
- # The MoERunner/SharedExperts class will coordinate with the MK to ensure
1028
- # that the SharedExperts are executed only once.
1029
- self.shared_experts = (
1030
- shared_experts if prepare_finalize.supports_async() else None
1031
- )
1025
self.inplace = inplace
1026
moe_parallel_config = fused_experts.moe_config.moe_parallel_config
1027
self.moe_parallel_config = moe_parallel_config
@@ -1103,11 +1096,12 @@ class FusedMoEKernelModularImpl:
1096
1097
def _maybe_apply_shared_experts(
1098
self,
1099
+ shared_experts: SharedExperts | None,
1100
shared_experts_input: torch.Tensor | None,
1101
):
1108
- if self.shared_experts is not None:
1102
+ if shared_experts is not None:
1103
assert shared_experts_input is not None
1110
- self.shared_experts.apply(
1104
+ shared_experts.apply(
1105
shared_experts_input,
1106
SharedExpertsOrder.MK_INTERNAL_OVERLAPPED,
1107
)
@@ -1289,6 +1283,7 @@ class FusedMoEKernelModularImpl:
1283
topk_weights: torch.Tensor,
1284
topk_ids: torch.Tensor,
1285
apply_router_weight_on_input: bool,
1286
+ shared_experts: SharedExperts | None,
1287
shared_experts_input: torch.Tensor | None,
1288
) -> torch.Tensor:
1289
"""
@@ -1296,6 +1291,7 @@ class FusedMoEKernelModularImpl:
1291
that handles DBO, async and shared expert overlap.
1292
1293
Args:
1294
+ shared_experts: SharedExperts | None. The shared experts if any.
1295
shared_experts_input: Optional separate input for shared experts.
1296
When latent MoE is used, hidden_states is the latent-projected
1297
tensor (smaller dimension) used by routed experts, while
@@ -1322,7 +1318,7 @@ class FusedMoEKernelModularImpl:
1318
apply_router_weight_on_input,
1319
self.fused_experts.finalize_weight_and_reduce_impl(),
1320
)
1325
- self._maybe_apply_shared_experts(shared_experts_input)
1321
+ self._maybe_apply_shared_experts(shared_experts, shared_experts_input)
1322
1323
# TODO(lucas): refactor this in the alternative schedules followup
1324
# currently unpack if we have hook + receiver pair or just
@@ -1358,6 +1354,7 @@ class FusedMoEKernelModularImpl:
1354
global_num_experts: int = -1,
1355
expert_map: torch.Tensor | None = None,
1356
apply_router_weight_on_input: bool = False,
1357
+ shared_experts: SharedExperts | None = None,
1358
shared_experts_input: torch.Tensor | None = None,
1359
) -> torch.Tensor:
1360
"""
@@ -1380,6 +1377,7 @@ class FusedMoEKernelModularImpl:
1377
- apply_router_weight_on_input (bool): When true, the topk weights are
1378
applied directly on the inputs. This is only applicable when topk is
1379
1.
1380
+ - shared_experts: SharedExperts | None. The shared experts if any.
1381
- shared_experts_input (Optional[torch.Tensor]): Optional separate
1382
input for shared experts. For latent MoE, this is the original
1383
hidden_states before latent projection.
@@ -1388,7 +1386,7 @@ class FusedMoEKernelModularImpl:
1386
- torch.Tensor: The output tensor after applying the MoE layer.
1387
"""
1388
if self.inplace:
1391
- assert self.shared_experts is None
1389
+ assert shared_experts is None
1390
assert not disable_inplace()
1391
output = hidden_states
1392
else:
@@ -1431,6 +1429,7 @@ class FusedMoEKernelModularImpl:
1429
topk_weights,
1430
topk_ids,
1431
apply_router_weight_on_input,
1432
+ shared_experts=shared_experts,
1433
shared_experts_input=shared_experts_input,
1434
)
1435
@@ -1503,7 +1502,6 @@ class FusedMoEKernel:
1502
self,
1503
prepare_finalize: FusedMoEPrepareAndFinalize,
1504
fused_experts: FusedMoEExperts,
1506
- shared_experts: SharedExperts | None = None,
1505
inplace: bool = False,
1506
):
1507
super().__init__()
@@ -1516,7 +1514,6 @@ class FusedMoEKernel:
1514
self.impl = FusedMoEKernelModularImpl(
1515
prepare_finalize,
1516
fused_experts,
1519
- shared_experts,
1517
inplace,
1518
)
1519
@@ -1539,9 +1536,9 @@ class FusedMoEKernel:
1536
self._post_init_setup()
1537
1538
@property
1542
- def owns_shared_experts(self) -> bool:
1539
+ def can_overlap_shared_experts(self) -> bool:
1540
if isinstance(self.impl, FusedMoEKernelModularImpl):
1544
- return self.impl.shared_experts is not None
1541
+ return self.impl.prepare_finalize.supports_async()
1542
else:
1543
return False
1544
@@ -1627,6 +1624,7 @@ class FusedMoEKernel:
1624
global_num_experts: int,
1625
expert_map: torch.Tensor | None,
1626
apply_router_weight_on_input: bool,
1627
+ shared_experts: SharedExperts | None = None,
1628
shared_experts_input: torch.Tensor | None = None,
1629
) -> torch.Tensor:
1630
assert isinstance(self.impl, FusedMoEKernelModularImpl)
@@ -1640,5 +1638,6 @@ class FusedMoEKernel:
1638
global_num_experts=global_num_experts,
1639
expert_map=expert_map,
1640
apply_router_weight_on_input=apply_router_weight_on_input,
1641
+ shared_experts=shared_experts,
1642
shared_experts_input=shared_experts_input,
1643
)
vllm/model_executor/layers/fused_moe/oracle/fp8.py
-8
@@ -18,9 +18,6 @@ from vllm.model_executor.layers.fused_moe.config import (
18
fp8_w8a8_moe_quant_config,
19
fp8_w8a16_moe_quant_config,
20
)
21
-from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
22
- SharedExperts,
23
-)
21
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
22
FlashinferMoeBackend,
23
get_flashinfer_moe_backend,
@@ -590,7 +587,6 @@ def make_fp8_moe_kernel(
587
experts_cls: type[mk.FusedMoEExperts],
588
fp8_backend: Fp8MoeBackend,
589
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
593
- shared_experts: SharedExperts | None = None,
590
) -> mk.FusedMoEKernel:
591
# Create Prepare/Finalize.
592
prepare_finalize = maybe_make_prepare_finalize(
@@ -620,13 +616,9 @@ def make_fp8_moe_kernel(
616
quant_config=moe_quant_config,
617
)
618
623
- # NOTE(rob): we only want the mk to control the shared_expert
624
- # if using all2all (for SBO). bnell is making this explicit in
625
- # the new MoE runner class.
619
kernel = mk.FusedMoEKernel(
620
prepare_finalize,
621
experts,
629
- shared_experts=shared_experts,
622
inplace=(
623
not moe_config.disable_inplace
624
and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS
vllm/model_executor/layers/fused_moe/oracle/int8.py
-5
@@ -17,9 +17,6 @@ from vllm.model_executor.layers.fused_moe.config import (
17
int8_w8a8_moe_quant_config,
18
int8_w8a16_moe_quant_config,
19
)
20
-from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
21
- SharedExperts,
22
-)
20
from vllm.model_executor.layers.quantization.utils.quant_utils import (
21
QuantKey,
22
kInt8DynamicTokenSym,
@@ -178,7 +175,6 @@ def make_int8_moe_kernel(
175
moe_config: FusedMoEConfig,
176
experts_cls: type[mk.FusedMoEExperts],
177
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
181
- shared_experts: SharedExperts | None = None,
178
) -> mk.FusedMoEKernel:
179
# Create Prepare/Finalize.
180
prepare_finalize = maybe_make_prepare_finalize(
@@ -211,7 +207,6 @@ def make_int8_moe_kernel(
207
kernel = mk.FusedMoEKernel(
208
prepare_finalize,
209
experts,
214
- shared_experts=shared_experts,
210
inplace=not moe_config.disable_inplace,
211
)
212
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
-2
@@ -154,7 +154,6 @@ def make_wna16_moe_kernel(
154
w13_g_idx_sort_indices: torch.Tensor | None,
155
w2_g_idx_sort_indices: torch.Tensor | None,
156
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
157
- shared_experts: torch.nn.Module | None = None,
157
) -> mk.FusedMoEKernel:
158
# Currently, we only support MarlinExperts and BatchedMarlinExperts
159
assert experts_cls in (MarlinExperts, BatchedMarlinExperts)
@@ -202,7 +201,6 @@ def make_wna16_moe_kernel(
201
return mk.FusedMoEKernel(
202
prepare_finalize,
203
experts,
205
- shared_experts=shared_experts,
204
inplace=not moe_config.disable_inplace,
205
)
206
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+6
-8
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from enum import Enum
4
-from typing import Union
4
+from typing import TYPE_CHECKING, Union
5
6
import torch
7
@@ -37,6 +37,10 @@ from vllm.platforms import current_platform
37
from vllm.utils.import_utils import has_triton_kernels
38
from vllm.utils.math_utils import round_up
39
40
+if TYPE_CHECKING:
41
+ from vllm.model_executor.layers.fused_moe import RoutedExperts
42
+
43
+
44
logger = init_logger(__name__)
45
46
if has_triton_kernels():
@@ -1479,8 +1483,7 @@ def make_mxfp4_moe_kernel(
1483
experts_cls: type[mk.FusedMoEExperts],
1484
mxfp4_backend: Mxfp4MoeBackend,
1485
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
1482
- shared_experts: torch.nn.Module | None = None,
1483
- layer: torch.nn.Module | None = None,
1486
+ layer: "RoutedExperts | None" = None,
1487
) -> mk.FusedMoEKernel:
1488
"""Create a FusedMoEKernel for the given MXFP4 backend."""
1489
is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)
@@ -1522,11 +1525,6 @@ def make_mxfp4_moe_kernel(
1525
kernel = mk.FusedMoEKernel(
1526
prepare_finalize,
1527
experts,
1525
- shared_experts=(
1526
- shared_experts
1527
- if moe_config.moe_parallel_config.use_batched_activation_format
1528
- else None
1529
- ),
1528
inplace=(
1529
not moe_config.disable_inplace and mxfp4_backend not in TRTLLM_BACKENDS
1530
),
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
-8
@@ -17,9 +17,6 @@ from vllm.model_executor.layers.fused_moe.config import (
17
nvfp4_moe_quant_config,
18
nvfp4_w4a16_moe_quant_config,
19
)
20
-from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
21
- SharedExperts,
22
-)
20
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
21
prepare_nvfp4_moe_layer_for_fi_or_cutlass,
22
prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
@@ -476,7 +473,6 @@ def make_nvfp4_moe_kernel(
473
moe_config: FusedMoEConfig,
474
experts_cls: type[mk.FusedMoEExperts],
475
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
479
- shared_experts: SharedExperts | None = None,
476
) -> mk.FusedMoEKernel:
477
# Create Prepare/Finalize.
478
prepare_finalize = maybe_make_prepare_finalize(
@@ -506,13 +502,9 @@ def make_nvfp4_moe_kernel(
502
quant_config=moe_quant_config,
503
)
504
509
- # NOTE(rob): we only want the mk to control the shared_expert
510
- # if using all2all (for SBO). bnell is making this explicit in
511
- # the new MoE runner class.
505
kernel = mk.FusedMoEKernel(
506
prepare_finalize,
507
experts,
515
- shared_experts=shared_experts,
508
inplace=False,
509
)
510
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
-5
@@ -18,9 +18,6 @@ from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEConfig,
19
FusedMoEQuantConfig,
20
)
21
-from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
22
- SharedExperts,
23
-)
21
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
22
FlashinferMoeBackend,
23
convert_moe_weights_to_flashinfer_trtllm_block_layout,
@@ -329,7 +326,6 @@ def make_unquantized_moe_kernel(
326
backend: UnquantizedMoeBackend,
327
experts_cls: type[mk.FusedMoEExperts],
328
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
332
- shared_experts: SharedExperts | None = None,
329
) -> mk.FusedMoEKernel:
330
# Create Prepare/Finalize
331
is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)
@@ -363,7 +359,6 @@ def make_unquantized_moe_kernel(
359
kernel = mk.FusedMoEKernel(
360
prepare_finalize,
361
experts,
366
- shared_experts=shared_experts,
362
inplace=(not moe_config.disable_inplace and not is_monolithic),
363
)
364
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+1
@@ -540,6 +540,7 @@ class MoERunner(MoERunnerInterface):
540
x=hidden_states,
541
topk_weights=topk_weights,
542
topk_ids=topk_ids,
543
+ shared_experts=self._shared_experts,
544
shared_experts_input=shared_experts_input,
545
)
546
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+1
-1
@@ -96,7 +96,7 @@ class SharedExperts:
96
if self._disable_shared_experts_overlap:
97
return SharedExpertsOrder.NO_OVERLAP
98
99
- if self._quant_method.mk_owns_shared_expert:
99
+ if self._quant_method.mk_can_overlap_shared_experts:
100
return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
101
102
should_run_shared_in_aux_stream = (
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+17
-5
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable
5
+from typing import TYPE_CHECKING
6
7
import torch
8
import torch.nn.functional as F
@@ -29,10 +30,16 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
30
make_unquantized_moe_kernel,
31
select_unquantized_moe_backend,
32
)
33
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
34
+ SharedExperts,
35
+)
36
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
37
from vllm.platforms import current_platform
38
from vllm.platforms.interface import CpuArchEnum
39
40
+if TYPE_CHECKING:
41
+ from vllm.model_executor.layers.fused_moe import RoutedExperts
42
+
43
logger = init_logger(__name__)
44
45
@@ -192,7 +199,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
199
backend=self.unquantized_backend,
200
experts_cls=self.experts_cls,
201
routing_tables=layer._expert_routing_tables(),
195
- shared_experts=layer.shared_experts,
202
)
203
204
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -277,10 +283,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
283
284
def apply(
285
self,
280
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
286
+ layer: "RoutedExperts",
287
x: torch.Tensor,
288
topk_weights: torch.Tensor,
289
topk_ids: torch.Tensor,
290
+ shared_experts: SharedExperts | None,
291
shared_experts_input: torch.Tensor | None,
292
) -> torch.Tensor:
293
return self.forward(
@@ -288,15 +295,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
295
x=x,
296
topk_weights=topk_weights,
297
topk_ids=topk_ids,
298
+ shared_experts=shared_experts,
299
shared_experts_input=shared_experts_input,
300
)
301
302
def forward_native(
303
self,
296
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
304
+ layer: "RoutedExperts",
305
x: torch.Tensor,
306
topk_weights: torch.Tensor,
307
topk_ids: torch.Tensor,
308
+ shared_experts: SharedExperts | None,
309
shared_experts_input: torch.Tensor | None,
310
) -> torch.Tensor:
311
assert self.moe_kernel is not None
@@ -310,15 +319,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
319
apply_router_weight_on_input=layer.apply_router_weight_on_input,
320
global_num_experts=layer.global_num_experts,
321
expert_map=layer.expert_map,
322
+ shared_experts=shared_experts,
323
shared_experts_input=shared_experts_input,
324
)
325
326
def forward_cuda(
327
self,
318
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
328
+ layer: "RoutedExperts",
329
x: torch.Tensor,
330
topk_weights: torch.Tensor,
331
topk_ids: torch.Tensor,
332
+ shared_experts: SharedExperts | None,
333
shared_experts_input: torch.Tensor | None,
334
) -> torch.Tensor:
335
return self.forward_native(
@@ -326,12 +337,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
337
x,
338
topk_weights,
339
topk_ids,
340
+ shared_experts,
341
shared_experts_input,
342
)
343
344
def apply_monolithic(
345
self,
334
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
346
+ layer: "RoutedExperts",
347
x: torch.Tensor,
348
router_logits: torch.Tensor,
349
input_ids: torch.Tensor | None = None,
vllm/model_executor/layers/quantization/awq.py
+2
-2
@@ -10,7 +10,7 @@ from transformers import PretrainedConfig
10
from vllm import _custom_ops as ops
11
from vllm import envs
12
from vllm.logger import init_logger
13
-from vllm.model_executor.layers.fused_moe.layer import FusedMoE
13
+from vllm.model_executor.layers.fused_moe import RoutedExperts
14
from vllm.model_executor.layers.linear import (
15
LinearBase,
16
LinearMethodBase,
@@ -106,7 +106,7 @@ class AWQConfig(QuantizationConfig):
106
):
107
return UnquantizedLinearMethod()
108
return AWQLinearMethod(self)
109
- elif isinstance(layer, FusedMoE):
109
+ elif isinstance(layer, RoutedExperts):
110
# Lazy import to avoid circular import.
111
from .awq_marlin import AWQMarlinConfig
112
from .moe_wna16 import MoeWNA16Config
vllm/model_executor/layers/quantization/awq_marlin.py
+14
-12
@@ -16,17 +16,18 @@ from vllm.model_executor.kernels.linear import (
16
MPLinearLayerConfig,
17
choose_mp_linear_kernel,
18
)
19
+from vllm.model_executor.layers.fused_moe import (
20
+ FusedMoEMethodBase,
21
+ FusedMoeWeightScaleSupported,
22
+ RoutedExperts,
23
+ SharedExperts,
24
+ UnquantizedFusedMoEMethod,
25
+)
26
from vllm.model_executor.layers.fused_moe.config import (
27
FusedMoEConfig,
28
FusedMoEQuantConfig,
29
)
30
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import fused_marlin_moe
24
-from vllm.model_executor.layers.fused_moe.layer import (
25
- FusedMoE,
26
- FusedMoEMethodBase,
27
- FusedMoeWeightScaleSupported,
28
- UnquantizedFusedMoEMethod,
29
-)
31
from vllm.model_executor.layers.linear import (
32
LinearBase,
33
LinearMethodBase,
@@ -286,7 +287,7 @@ class AWQMarlinConfig(QuantizationConfig):
287
quant_method = AWQMarlinLinearMethod(self)
288
quant_method.input_dtype = get_marlin_input_dtype(prefix)
289
return quant_method
289
- elif isinstance(layer, FusedMoE):
290
+ elif isinstance(layer, RoutedExperts):
291
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
292
293
if is_layer_skipped(
@@ -507,7 +508,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
508
509
def create_weights(
510
self,
510
- layer: torch.nn.Module,
511
+ layer: RoutedExperts,
512
num_experts: int,
513
hidden_size: int,
514
intermediate_size_per_partition: int,
@@ -606,7 +607,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
607
device = layer.w13_qweight.device
608
layer.workspace = marlin_make_workspace_new(device, 4)
609
609
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
610
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
611
num_experts = layer.w13_qweight.shape[0]
612
device = layer.w13_qweight.device
613
is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
@@ -723,7 +724,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
724
layer.w2_bias.data = marlin_permute_bias(layer.w2_bias)
725
726
def get_fused_moe_quant_config(
726
- self, layer: torch.nn.Module
727
+ self, layer: RoutedExperts
728
) -> FusedMoEQuantConfig | None:
729
from vllm.model_executor.layers.fused_moe.config import (
730
awq_marlin_moe_quant_config,
@@ -747,7 +748,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
748
def select_gemm_impl(
749
self,
750
prepare_finalize,
750
- layer: torch.nn.Module,
751
+ layer: RoutedExperts,
752
):
753
"""
754
Select the GEMM implementation for AWQ-Marlin MoE.
@@ -812,10 +813,11 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
813
814
def apply(
815
self,
815
- layer: FusedMoE,
816
+ layer: RoutedExperts,
817
x: torch.Tensor,
818
topk_weights: torch.Tensor,
819
topk_ids: torch.Tensor,
820
+ shared_experts: SharedExperts | None,
821
shared_experts_input: torch.Tensor | None,
822
) -> torch.Tensor:
823
return fused_marlin_moe(
vllm/model_executor/layers/quantization/bitsandbytes.py
+9
-9
@@ -6,13 +6,12 @@ from typing import Any, Union
6
import torch
7
from packaging import version
8
9
-from vllm.model_executor.layers.fused_moe.config import (
9
+from vllm.model_executor.layers.fused_moe import (
10
FusedMoEConfig,
11
- FusedMoEQuantConfig,
12
-)
13
-from vllm.model_executor.layers.fused_moe.layer import (
14
- FusedMoE,
11
FusedMoEMethodBase,
12
+ FusedMoEQuantConfig,
13
+ RoutedExperts,
14
+ SharedExperts,
15
)
16
from vllm.model_executor.layers.linear import (
17
LinearBase,
@@ -164,7 +163,7 @@ class BitsAndBytesConfig(QuantizationConfig):
163
if is_layer_skipped_bnb(prefix, self.llm_int8_skip_modules):
164
return UnquantizedLinearMethod()
165
return BitsAndBytesLinearMethod(self)
167
- elif isinstance(layer, FusedMoE):
166
+ elif isinstance(layer, RoutedExperts):
167
return BitsAndBytesMoEMethod(self, layer.moe_config)
168
return None
169
@@ -451,7 +450,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
450
451
def create_weights(
452
self,
454
- layer: torch.nn.Module,
453
+ layer: RoutedExperts,
454
num_experts: int,
455
hidden_size: int,
456
intermediate_size_per_partition: int,
@@ -472,16 +471,17 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
471
)
472
473
def get_fused_moe_quant_config(
475
- self, layer: torch.nn.Module
474
+ self, layer: RoutedExperts
475
) -> FusedMoEQuantConfig | None:
476
return None
477
478
def apply(
479
self,
481
- layer: FusedMoE,
480
+ layer: RoutedExperts,
481
x: torch.Tensor,
482
topk_weights: torch.Tensor,
483
topk_ids: torch.Tensor,
484
+ shared_experts: SharedExperts | None,
485
shared_experts_input: torch.Tensor | None,
486
) -> torch.Tensor:
487
from vllm.model_executor.layers.fused_moe import fused_experts
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+3
-3
@@ -23,7 +23,7 @@ from vllm.distributed import (
23
)
24
from vllm.logger import init_logger
25
from vllm.model_executor.layers.attention import Attention
26
-from vllm.model_executor.layers.fused_moe import FusedMoE
26
+from vllm.model_executor.layers.fused_moe import RoutedExperts
27
from vllm.model_executor.layers.linear import (
28
LinearBase,
29
LinearMethodBase,
@@ -191,7 +191,7 @@ class CompressedTensorsConfig(QuantizationConfig):
191
192
if isinstance(layer, Attention):
193
return CompressedTensorsKVCacheMethod(self)
194
- if isinstance(layer, FusedMoE):
194
+ if isinstance(layer, RoutedExperts):
195
return CompressedTensorsMoEMethod.get_moe_method(
196
self, layer, layer_name=prefix
197
)
@@ -202,7 +202,7 @@ class CompressedTensorsConfig(QuantizationConfig):
202
Helper function to update target_scheme_map
203
since linear layers get fused into FusedMoE
204
targeting 'Linear' needs to also match
205
- FusedMoE modules.
205
+ RoutedExperts modules.
206
"""
207
if (
208
"Linear" not in self.target_scheme_map
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+1
-1
@@ -32,7 +32,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
32
layer: torch.nn.Module,
33
layer_name: str,
34
) -> FusedMoEMethodBase:
35
- # FusedMoE was made by combining multiple Linears so need to
35
+ # RoutedExperts was made by combining multiple Linears so need to
36
# make sure quantization config for Linear can target it
37
quant_config._add_fused_moe_to_target_scheme_map()
38
unfused_names = [
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+6
-4
@@ -7,8 +7,9 @@ import torch
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
from vllm.logger import init_logger
9
from vllm.model_executor.layers.fused_moe import (
10
- FusedMoE,
10
FusedMoeWeightScaleSupported,
11
+ RoutedExperts,
12
+ SharedExperts,
13
)
14
from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEQuantConfig,
@@ -136,7 +137,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
137
w2_scale=layer.w2_weight_scale,
138
)
139
139
- def process_weights_after_loading(self, layer: FusedMoE) -> None:
140
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
141
layer.w13_weight = torch.nn.Parameter(
142
layer.w13_weight_packed.data, requires_grad=False
143
)
@@ -194,16 +195,16 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
195
moe_config=self.moe,
196
experts_cls=self.experts_cls,
197
mxfp4_backend=self.mxfp4_backend,
197
- shared_experts=layer.shared_experts,
198
routing_tables=layer._expert_routing_tables(),
199
)
200
201
def apply(
202
self,
203
- layer: FusedMoE,
203
+ layer: RoutedExperts,
204
x: torch.Tensor,
205
topk_weights: torch.Tensor,
206
topk_ids: torch.Tensor,
207
+ shared_experts: SharedExperts | None,
208
shared_experts_input: torch.Tensor | None,
209
) -> torch.Tensor:
210
assert self.moe_kernel is not None
@@ -217,5 +218,6 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
218
global_num_experts=layer.global_num_experts,
219
expert_map=layer.expert_map,
220
apply_router_weight_on_input=layer.apply_router_weight_on_input,
221
+ shared_experts=shared_experts,
222
shared_experts_input=shared_experts_input,
223
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+7
-5
@@ -7,8 +7,9 @@ import torch
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
from vllm.logger import init_logger
9
from vllm.model_executor.layers.fused_moe import (
10
- FusedMoE,
10
FusedMoeWeightScaleSupported,
11
+ RoutedExperts,
12
+ SharedExperts,
13
)
14
from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEConfig,
@@ -167,7 +168,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
168
)
169
set_weight_attrs(w2_input_scale, extra_weight_attrs)
170
170
- def process_weights_after_loading(self, layer: FusedMoE) -> None:
171
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
172
"""
173
Convert NVFP4 MoE weights into kernel format and setup the kernel.
174
"""
@@ -235,7 +236,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
236
moe_quant_config=self.moe_quant_config,
237
moe_config=self.moe,
238
experts_cls=self.experts_cls,
238
- shared_experts=layer.shared_experts,
239
routing_tables=layer._expert_routing_tables(),
240
)
241
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
@@ -262,7 +262,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
262
263
def apply_monolithic(
264
self,
265
- layer: FusedMoE,
265
+ layer: RoutedExperts,
266
x: torch.Tensor,
267
router_logits: torch.Tensor,
268
input_ids: torch.Tensor | None = None,
@@ -286,10 +286,11 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
286
287
def apply(
288
self,
289
- layer: FusedMoE,
289
+ layer: RoutedExperts,
290
x: torch.Tensor,
291
topk_weights: torch.Tensor,
292
topk_ids: torch.Tensor,
293
+ shared_experts: SharedExperts | None,
294
shared_experts_input: torch.Tensor | None,
295
) -> torch.Tensor:
296
assert self.moe_kernel is not None
@@ -303,5 +304,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
304
global_num_experts=layer.global_num_experts,
305
expert_map=layer.expert_map,
306
apply_router_weight_on_input=layer.apply_router_weight_on_input,
307
+ shared_experts=shared_experts,
308
shared_experts_input=shared_experts_input,
309
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py
+5
-3
@@ -11,10 +11,11 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm import _custom_ops as ops
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.fused_moe import (
14
- FusedMoE,
14
FusedMoEActivationFormat,
15
FusedMoEExpertsModular,
16
FusedMoeWeightScaleSupported,
17
+ RoutedExperts,
18
+ SharedExperts,
19
)
20
from vllm.model_executor.layers.fused_moe.config import (
21
FusedMoEConfig,
@@ -136,7 +137,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
137
requires_grad=False,
138
)
139
layer.register_parameter("w2_weight_scale", w2_weight_scale)
139
- # Add PER-GROUP quantization for FusedMoE.weight_loader.
140
+ # Add PER-GROUP quantization for RoutedExperts.weight_loader.
141
extra_weight_attrs.update(
142
{"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
143
)
@@ -303,10 +304,11 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
304
305
def apply(
306
self,
306
- layer: FusedMoE,
307
+ layer: RoutedExperts,
308
x: torch.Tensor,
309
topk_weights: torch.Tensor,
310
topk_ids: torch.Tensor,
311
+ shared_experts: SharedExperts | None,
312
shared_experts_input: torch.Tensor | None,
313
) -> torch.Tensor:
314
if layer.enable_eplb:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py
+2
-2
@@ -10,7 +10,7 @@ from compressed_tensors.quantization import (
10
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.fused_moe import (
13
- FusedMoE,
13
+ RoutedExperts,
14
)
15
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16
from vllm.model_executor.layers.fused_moe.config import (
@@ -302,7 +302,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
302
303
def apply_monolithic(
304
self,
305
- layer: FusedMoE,
305
+ layer: RoutedExperts,
306
x: torch.Tensor,
307
router_logits: torch.Tensor,
308
input_ids: torch.Tensor | None = None,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
+10
-8
@@ -12,8 +12,9 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.fused_moe import (
15
- FusedMoE,
15
FusedMoeWeightScaleSupported,
16
+ RoutedExperts,
17
+ SharedExperts,
18
)
19
from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEConfig,
@@ -191,7 +192,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
192
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
193
)
194
layer.register_parameter("w2_weight_scale", w2_weight_scale)
194
- # Add PER-TENSOR quantization for FusedMoE.weight_loader.
195
+ # Add PER-TENSOR quantization for RoutedExperts.weight_loader.
196
extra_weight_attrs.update(
197
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
198
)
@@ -214,7 +215,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
215
requires_grad=False,
216
)
217
layer.register_parameter("w2_weight_scale", w2_weight_scale)
217
- # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
218
+ # Add PER-CHANNEL quantization for RoutedExperts.weight_loader.
219
extra_weight_attrs.update(
220
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
221
)
@@ -243,7 +244,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
244
requires_grad=False,
245
)
246
layer.register_parameter("w2_weight_scale", w2_weight_scale)
246
- # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
247
+ # Add PER-CHANNEL quantization for RoutedExperts.weight_loader.
248
extra_weight_attrs.update(
249
{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
250
)
@@ -267,7 +268,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
268
layer.w13_input_scale = None
269
layer.w2_input_scale = None
270
270
- def process_weights_after_loading(self, layer: FusedMoE) -> None:
271
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
272
# Allow for accessing weights and scales in standard way.
273
w13 = layer.w13_weight
274
w2 = layer.w2_weight
@@ -337,7 +338,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
338
fp8_backend=self.fp8_backend,
339
experts_cls=self.experts_cls,
340
routing_tables=layer._expert_routing_tables(),
340
- shared_experts=layer.shared_experts,
341
)
342
343
def maybe_make_prepare_finalize(
@@ -364,7 +364,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
364
365
def apply_monolithic(
366
self,
367
- layer: FusedMoE,
367
+ layer: RoutedExperts,
368
x: torch.Tensor,
369
router_logits: torch.Tensor,
370
input_ids: torch.Tensor | None = None,
@@ -387,10 +387,11 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
387
388
def apply(
389
self,
390
- layer: FusedMoE,
390
+ layer: RoutedExperts,
391
x: torch.Tensor,
392
topk_weights: torch.Tensor,
393
topk_ids: torch.Tensor,
394
+ shared_experts: SharedExperts | None,
395
shared_experts_input: torch.Tensor | None,
396
) -> torch.Tensor:
397
assert not self.is_monolithic
@@ -407,6 +408,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
408
# https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501
409
expert_map=layer.expert_map,
410
apply_router_weight_on_input=layer.apply_router_weight_on_input,
411
+ shared_experts=shared_experts,
412
shared_experts_input=shared_experts_input,
413
)
414
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py
+7
-5
@@ -11,8 +11,9 @@ from compressed_tensors.quantization import (
11
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.fused_moe import (
14
- FusedMoE,
14
FusedMoeWeightScaleSupported,
15
+ RoutedExperts,
16
+ SharedExperts,
17
)
18
from vllm.model_executor.layers.fused_moe.config import (
19
FusedMoEConfig,
@@ -128,7 +129,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
129
requires_grad=False,
130
)
131
layer.register_parameter("w2_weight_scale", w2_weight_scale)
131
- # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
132
+ # Add PER-CHANNEL quantization for RoutedExperts.weight_loader.
133
extra_weight_attrs.update(
134
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
135
)
@@ -140,7 +141,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
141
layer.w13_input_scale = None
142
layer.w2_input_scale = None
143
143
- def process_weights_after_loading(self, layer: FusedMoE) -> None:
144
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
145
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
146
assert self.experts_cls is not None
147
self.moe_kernel = make_int8_moe_kernel(
@@ -148,7 +149,6 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
149
moe_config=self.moe,
150
experts_cls=self.experts_cls,
151
routing_tables=layer._expert_routing_tables(),
151
- shared_experts=layer.shared_experts,
152
)
153
154
def maybe_make_prepare_finalize(
@@ -171,10 +171,11 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
171
172
def apply(
173
self,
174
- layer: FusedMoE,
174
+ layer: RoutedExperts,
175
x: torch.Tensor,
176
topk_weights: torch.Tensor,
177
topk_ids: torch.Tensor,
178
+ shared_experts: SharedExperts | None,
179
shared_experts_input: torch.Tensor | None,
180
) -> torch.Tensor:
181
assert not self.is_monolithic
@@ -189,5 +190,6 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
190
global_num_experts=layer.global_num_experts,
191
expert_map=layer.expert_map,
192
apply_router_weight_on_input=layer.apply_router_weight_on_input,
193
+ shared_experts=shared_experts,
194
shared_experts_input=shared_experts_input,
195
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py
+7
-5
@@ -5,8 +5,9 @@ import torch
5
6
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
from vllm.model_executor.layers.fused_moe import (
8
- FusedMoE,
8
FusedMoeWeightScaleSupported,
9
+ RoutedExperts,
10
+ SharedExperts,
11
)
12
from vllm.model_executor.layers.fused_moe.config import (
13
FusedMoEConfig,
@@ -111,7 +112,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
112
layer.w13_input_scale = None
113
layer.w2_input_scale = None
114
114
- def process_weights_after_loading(self, layer: FusedMoE) -> None:
115
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
116
layer.weight_block_size = self.weight_block_size
117
118
w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
@@ -139,7 +140,6 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
140
fp8_backend=self.fp8_backend,
141
experts_cls=self.experts_cls,
142
routing_tables=layer._expert_routing_tables(),
142
- shared_experts=layer.shared_experts,
143
)
144
145
def get_fused_moe_quant_config(
@@ -165,7 +165,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
165
166
def apply_monolithic(
167
self,
168
- layer: FusedMoE,
168
+ layer: RoutedExperts,
169
x: torch.Tensor,
170
router_logits: torch.Tensor,
171
input_ids: torch.Tensor | None = None,
@@ -188,10 +188,11 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
188
189
def apply(
190
self,
191
- layer: FusedMoE,
191
+ layer: RoutedExperts,
192
x: torch.Tensor,
193
topk_weights: torch.Tensor,
194
topk_ids: torch.Tensor,
195
+ shared_experts: SharedExperts | None,
196
shared_experts_input: torch.Tensor | None,
197
) -> torch.Tensor:
198
assert not self.is_monolithic
@@ -206,5 +207,6 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
207
global_num_experts=layer.global_num_experts,
208
expert_map=layer.expert_map,
209
apply_router_weight_on_input=layer.apply_router_weight_on_input,
210
+ shared_experts=shared_experts,
211
shared_experts_input=shared_experts_input,
212
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+4
-2
@@ -10,7 +10,8 @@ from compressed_tensors.quantization import (
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.fused_moe import (
13
- FusedMoE,
13
+ RoutedExperts,
14
+ SharedExperts,
15
)
16
from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEConfig,
@@ -240,10 +241,11 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
241
242
def apply(
243
self,
243
- layer: FusedMoE,
244
+ layer: RoutedExperts,
245
x: torch.Tensor,
246
topk_weights: torch.Tensor,
247
topk_ids: torch.Tensor,
248
+ shared_experts: SharedExperts | None,
249
shared_experts_input: torch.Tensor | None,
250
) -> torch.Tensor:
251
from vllm.model_executor.layers.fused_moe import fused_experts
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+5
-3
@@ -13,7 +13,8 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
from vllm import _custom_ops as ops
14
from vllm.logger import init_logger
15
from vllm.model_executor.layers.fused_moe import (
16
- FusedMoE,
16
+ RoutedExperts,
17
+ SharedExperts,
18
)
19
from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEConfig,
@@ -514,7 +515,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
515
516
def apply_monolithic(
517
self,
517
- layer: FusedMoE,
518
+ layer: RoutedExperts,
519
x: torch.Tensor,
520
router_logits: torch.Tensor,
521
input_ids: torch.Tensor | None = None,
@@ -540,10 +541,11 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
541
542
def apply(
543
self,
543
- layer: FusedMoE,
544
+ layer: RoutedExperts,
545
x: torch.Tensor,
546
topk_weights: torch.Tensor,
547
topk_ids: torch.Tensor,
548
+ shared_experts: SharedExperts | None,
549
shared_experts_input: torch.Tensor | None,
550
) -> torch.Tensor:
551
assert self.kernel_backend == "Marlin"
vllm/model_executor/layers/quantization/experts_int8.py
+4
-2
@@ -5,7 +5,9 @@ from typing import Any
5
6
import torch
7
8
-from vllm.model_executor.layers.fused_moe import FusedMoE
8
+from vllm.model_executor.layers.fused_moe import (
9
+ RoutedExperts,
10
+)
11
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
12
from vllm.model_executor.layers.quantization import QuantizationMethods
13
from vllm.model_executor.layers.quantization.base_config import (
@@ -53,6 +55,6 @@ class ExpertsInt8Config(QuantizationConfig):
55
) -> "QuantizeMethodBase | None":
56
if isinstance(layer, LinearBase):
57
return UnquantizedLinearMethod()
56
- elif isinstance(layer, FusedMoE):
58
+ elif isinstance(layer, RoutedExperts):
59
return Int8OnlineMoEMethod(layer=layer)
60
return None
vllm/model_executor/layers/quantization/fp8.py
+19
-18
@@ -4,7 +4,6 @@
4
from typing import TYPE_CHECKING, Any
5
6
import torch
7
-from torch.nn import Module
7
from torch.utils._python_dispatch import TorchDispatchMode
8
9
import vllm.envs as envs
@@ -22,14 +21,15 @@ from vllm.model_executor.kernels.linear.scaled_mm import (
21
)
22
from vllm.model_executor.layers.attention import Attention
23
from vllm.model_executor.layers.fused_moe import (
25
- FusedMoE,
24
FusedMoEMethodBase,
25
FusedMoeWeightScaleSupported,
26
+ RoutedExperts,
27
+ SharedExperts,
28
+ UnquantizedFusedMoEMethod,
29
)
30
from vllm.model_executor.layers.fused_moe.config import (
31
FusedMoEQuantConfig,
32
)
32
-from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
33
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
34
Fp8MoeBackend,
35
convert_to_fp8_moe_kernel_format,
@@ -191,7 +191,7 @@ class Fp8Config(QuantizationConfig):
191
offline_method = Fp8LinearMethod(self)
192
offline_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
193
return offline_method
194
- elif isinstance(layer, FusedMoE):
194
+ elif isinstance(layer, RoutedExperts):
195
if is_layer_skipped(
196
prefix=prefix,
197
ignored_layers=self.ignored_layers,
@@ -315,7 +315,7 @@ class Fp8LinearMethod(LinearMethodBase):
315
316
def create_weights(
317
self,
318
- layer: torch.nn.Module,
318
+ layer: RoutedExperts,
319
input_size_per_partition: int,
320
output_partition_sizes: list[int],
321
input_size: int,
@@ -389,7 +389,7 @@ class Fp8LinearMethod(LinearMethodBase):
389
390
self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
391
392
- def process_weights_after_loading(self, layer: Module) -> None:
392
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
393
if self.use_marlin:
394
# Only Marlin kernels support `marlin_input_dtype`; guard to avoid
395
# AttributeError if backend selection changes.
@@ -534,7 +534,7 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
534
)
535
self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
536
537
- def process_weights_after_loading(self, layer: Module) -> None:
537
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
538
if getattr(layer, "_already_called_process_weights_after_loading", False):
539
return
540
@@ -575,7 +575,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
575
quant_config: The quantization config.
576
"""
577
578
- def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
578
+ def __init__(self, quant_config: Fp8Config, layer: RoutedExperts):
579
super().__init__(layer.moe_config)
580
self.quant_config = quant_config
581
self.weight_block_size = self.quant_config.weight_block_size
@@ -606,7 +606,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
606
607
def create_weights(
608
self,
609
- layer: Module,
609
+ layer: RoutedExperts,
610
num_experts: int,
611
hidden_size: int,
612
intermediate_size_per_partition: int,
@@ -746,7 +746,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
746
747
def _setup_kernel(
748
self,
749
- layer: FusedMoE,
749
+ layer: RoutedExperts,
750
w13: torch.Tensor,
751
w2: torch.Tensor,
752
w13_scale: torch.Tensor,
@@ -787,10 +787,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
787
fp8_backend=self.fp8_backend,
788
experts_cls=self.experts_cls,
789
routing_tables=layer._expert_routing_tables(),
790
- shared_experts=layer.shared_experts,
790
)
791
793
- def process_weights_after_loading(self, layer: Module) -> None:
792
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
793
# Allow for accessing weights and scales in standard way.
794
w13 = layer.w13_weight
795
w2 = layer.w2_weight
@@ -844,7 +843,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
843
"logic. This function should not be called."
844
)
845
847
- def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
846
+ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
847
w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
848
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
849
a1_scale = layer.w13_input_scale
@@ -877,7 +876,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
876
877
def apply_monolithic(
878
self,
880
- layer: FusedMoE,
879
+ layer: RoutedExperts,
880
x: torch.Tensor,
881
router_logits: torch.Tensor,
882
input_ids: torch.Tensor | None = None,
@@ -901,10 +900,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
900
901
def apply(
902
self,
904
- layer: FusedMoE,
903
+ layer: RoutedExperts,
904
x: torch.Tensor,
905
topk_weights: torch.Tensor,
906
topk_ids: torch.Tensor,
907
+ shared_experts: SharedExperts | None,
908
shared_experts_input: torch.Tensor | None,
909
) -> torch.Tensor:
910
assert not self.is_monolithic
@@ -919,6 +919,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
919
global_num_experts=layer.global_num_experts,
920
expert_map=layer.expert_map,
921
apply_router_weight_on_input=layer.apply_router_weight_on_input,
922
+ shared_experts=shared_experts,
923
shared_experts_input=shared_experts_input,
924
)
925
@@ -937,7 +938,7 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
938
939
uses_meta_device: bool = True
940
940
- def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
941
+ def __init__(self, quant_config: Fp8Config, layer: RoutedExperts):
942
super().__init__(quant_config, layer)
943
assert not quant_config.is_checkpoint_fp8_serialized
944
assert quant_config.activation_scheme == "dynamic"
@@ -945,7 +946,7 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
946
947
def create_weights(
948
self,
948
- layer: Module,
949
+ layer: RoutedExperts,
950
num_experts: int,
951
hidden_size: int,
952
intermediate_size_per_partition: int,
@@ -1011,7 +1012,7 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1012
1013
initialize_online_processing(layer)
1014
1014
- def process_weights_after_loading(self, layer: Module) -> None:
1015
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
1016
# TODO(@ksayers): inplace fp8 quant kernel, initialize scales with ones
1017
if getattr(layer, "_already_called_process_weights_after_loading", False):
1018
return
vllm/model_executor/layers/quantization/gguf.py
+11
-13
@@ -15,17 +15,14 @@ from torch.nn.parameter import Parameter, UninitializedParameter
15
16
from vllm import _custom_ops as ops
17
from vllm.logger import init_logger
18
-from vllm.model_executor.layers.fused_moe.activation import (
19
- MoEActivation,
20
- apply_moe_activation,
21
-)
22
-from vllm.model_executor.layers.fused_moe.config import (
18
+from vllm.model_executor.layers.fused_moe import (
19
FusedMoEConfig,
24
- FusedMoEQuantConfig,
25
-)
26
-from vllm.model_executor.layers.fused_moe.layer import (
27
- FusedMoE,
20
FusedMoEMethodBase,
21
+ FusedMoEQuantConfig,
22
+ MoEActivation,
23
+ RoutedExperts,
24
+ SharedExperts,
25
+ apply_moe_activation,
26
)
27
from vllm.model_executor.layers.linear import (
28
LinearBase,
@@ -107,7 +104,7 @@ class GGUFConfig(QuantizationConfig):
104
):
105
return UnquantizedEmbeddingMethod()
106
return GGUFEmbeddingMethod(self)
110
- elif isinstance(layer, FusedMoE):
107
+ elif isinstance(layer, RoutedExperts):
108
# TODO: Select UnquantizedFusedMoEMethod on unquantized layers.
109
return GGUFMoEMethod(self, layer.moe_config)
110
return None
@@ -578,7 +575,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):
575
576
def create_weights(
577
self,
581
- layer: torch.nn.Module,
578
+ layer: RoutedExperts,
579
num_experts: int,
580
hidden_size: int,
581
intermediate_size_per_partition: int,
@@ -639,16 +636,17 @@ class GGUFMoEMethod(FusedMoEMethodBase):
636
layer.register_parameter("w2_qweight_type", w2_qweight_type)
637
638
def get_fused_moe_quant_config(
642
- self, layer: torch.nn.Module
639
+ self, layer: RoutedExperts
640
) -> FusedMoEQuantConfig | None:
641
return None
642
643
def apply(
644
self,
648
- layer: FusedMoE,
645
+ layer: RoutedExperts,
646
x: torch.Tensor,
647
topk_weights: torch.Tensor,
648
topk_ids: torch.Tensor,
649
+ shared_experts: SharedExperts | None,
650
shared_experts_input: torch.Tensor | None,
651
) -> torch.Tensor:
652
if layer.apply_router_weight_on_input:
vllm/model_executor/layers/quantization/gptq.py
+2
-2
@@ -13,7 +13,7 @@ from transformers import PretrainedConfig
13
14
from vllm import _custom_ops as ops
15
from vllm.logger import init_logger
16
-from vllm.model_executor.layers.fused_moe.layer import FusedMoE
16
+from vllm.model_executor.layers.fused_moe import RoutedExperts
17
from vllm.model_executor.layers.linear import LinearMethodBase
18
from vllm.model_executor.layers.quantization.base_config import (
19
QuantizationConfig,
@@ -172,7 +172,7 @@ class GPTQConfig(QuantizationConfig):
172
def get_quant_method(
173
self, layer: torch.nn.Module, prefix: str
174
) -> Union["GPTQLinearMethod", "QuantizeMethodBase"] | None:
175
- if isinstance(layer, FusedMoE):
175
+ if isinstance(layer, RoutedExperts):
176
# GPTQ MoE support: fall back to MoeWNA16 for broad compatibility
177
from .moe_wna16 import MoeWNA16Config
178
vllm/model_executor/layers/quantization/gptq_marlin.py
+28
-29
@@ -14,14 +14,13 @@ from vllm.model_executor.kernels.linear import (
14
MPLinearLayerConfig,
15
choose_mp_linear_kernel,
16
)
17
-from vllm.model_executor.layers.fused_moe.config import (
17
+from vllm.model_executor.layers.fused_moe import (
18
FusedMoEConfig,
19
- FusedMoEQuantConfig,
20
-)
21
-from vllm.model_executor.layers.fused_moe.layer import (
22
- FusedMoE,
19
FusedMoEMethodBase,
20
+ FusedMoEQuantConfig,
21
FusedMoeWeightScaleSupported,
22
+ RoutedExperts,
23
+ SharedExperts,
24
UnquantizedFusedMoEMethod,
25
)
26
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
@@ -71,29 +70,28 @@ logger = init_logger(__name__)
70
71
def get_moe_quant_method(
72
config: "GPTQMarlinConfig",
74
- layer: torch.nn.Module,
73
+ layer: RoutedExperts,
74
prefix: str,
75
moe_method_cls: type,
76
):
77
cloned_config = deepcopy(config)
78
80
- if isinstance(layer, FusedMoE):
81
- # False = skip module, None = no override, else = Positive match
82
- if (
83
- get_dynamic_override( # noqa: E712
84
- cloned_config, # noqa: E712
85
- layer_name=prefix,
86
- )
87
- == False
88
- ): # noqa: E712
89
- return UnquantizedFusedMoEMethod(layer.moe_config)
79
+ assert isinstance(layer, RoutedExperts)
80
+ # False = skip module, None = no override, else = Positive match
81
+ if (
82
+ get_dynamic_override( # noqa: E712
83
+ cloned_config, # noqa: E712
84
+ layer_name=prefix,
85
+ )
86
+ == False
87
+ ): # noqa: E712
88
+ return UnquantizedFusedMoEMethod(layer.moe_config)
89
91
- if prefix:
92
- # Dynamic per module/layer rules may override base config
93
- override_config(cloned_config, prefix=prefix)
90
+ if prefix:
91
+ # Dynamic per module/layer rules may override base config
92
+ override_config(cloned_config, prefix=prefix)
93
95
- return moe_method_cls(cloned_config, layer.moe_config)
96
- return None
94
+ return moe_method_cls(cloned_config, layer.moe_config)
95
96
97
class GPTQMarlinConfig(QuantizationConfig):
@@ -247,7 +245,7 @@ class GPTQMarlinConfig(QuantizationConfig):
245
def get_quant_method(
246
self, layer: torch.nn.Module, prefix: str
247
) -> "QuantizeMethodBase | None":
250
- if isinstance(layer, FusedMoE):
248
+ if isinstance(layer, RoutedExperts):
249
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
250
251
if not check_moe_marlin_supports_layer(layer, self.group_size):
@@ -522,7 +520,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
520
521
def create_weights(
522
self,
525
- layer: torch.nn.Module,
523
+ layer: RoutedExperts,
524
num_experts: int,
525
hidden_size: int,
526
intermediate_size_per_partition: int,
@@ -676,7 +674,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
674
device = layer.w13_qweight.device
675
layer.workspace = marlin_make_workspace_new(device, 4)
676
679
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
677
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
678
is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
679
680
if is_a_8bit:
@@ -755,7 +753,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
753
754
self._setup_kernel(layer)
755
758
- def _setup_kernel(self, layer: FusedMoE) -> None:
756
+ def _setup_kernel(self, layer: RoutedExperts) -> None:
757
"""Build the FusedMoEKernel for this layer."""
758
759
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
@@ -770,10 +768,9 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
768
w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
769
w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
770
routing_tables=layer._expert_routing_tables(),
773
- shared_experts=layer.shared_experts,
771
)
772
776
- def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
773
+ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
774
from vllm.model_executor.layers.fused_moe.config import (
775
gptq_marlin_moe_quant_config,
776
)
@@ -796,7 +793,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
793
def select_gemm_impl(
794
self,
795
prepare_finalize,
799
- layer: torch.nn.Module,
796
+ layer: RoutedExperts,
797
):
798
raise ValueError(
799
f"{self.__class__.__name__} uses the new modular kernel "
@@ -805,10 +802,11 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
802
803
def apply(
804
self,
808
- layer: FusedMoE,
805
+ layer: RoutedExperts,
806
x: torch.Tensor,
807
topk_weights: torch.Tensor,
808
topk_ids: torch.Tensor,
809
+ shared_experts: SharedExperts | None,
810
shared_experts_input: torch.Tensor | None,
811
) -> torch.Tensor:
812
assert not self.is_monolithic
@@ -823,5 +821,6 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
821
global_num_experts=layer.global_num_experts,
822
apply_router_weight_on_input=layer.apply_router_weight_on_input,
823
expert_map=layer.expert_map,
824
+ shared_experts=shared_experts,
825
shared_experts_input=shared_experts_input,
826
)
vllm/model_executor/layers/quantization/humming.py
+11
-9
@@ -9,14 +9,15 @@ import regex as re
9
import torch
10
11
from vllm import envs
12
+from vllm.model_executor.layers.fused_moe import (
13
+ FusedMoEMethodBase,
14
+ RoutedExperts,
15
+ SharedExperts,
16
+)
17
from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEConfig,
19
FusedMoEQuantConfig,
20
)
16
-from vllm.model_executor.layers.fused_moe.layer import (
17
- FusedMoE,
18
- FusedMoEMethodBase,
19
-)
21
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
22
UnquantizedFusedMoEMethod,
23
)
@@ -334,20 +335,20 @@ class HummingConfig(QuantizationConfig):
335
self, layer: torch.nn.Module, prefix: str
336
) -> "QuantizeMethodBase | None":
337
layer_type = "other"
337
- if isinstance(layer, FusedMoE):
338
+ if isinstance(layer, RoutedExperts):
339
layer_type = "moe"
340
elif isinstance(layer, LinearBase):
341
layer_type = "linear"
342
343
quant_config = self.get_quant_config_for_layer(prefix, layer_type)
344
if quant_config is None:
344
- if isinstance(layer, FusedMoE):
345
+ if isinstance(layer, RoutedExperts):
346
return UnquantizedFusedMoEMethod(layer.moe_config)
347
elif isinstance(layer, LinearBase):
348
return UnquantizedLinearMethod()
349
elif isinstance(layer, LinearBase):
350
return HummingLinearMethod(quant_config)
350
- elif isinstance(layer, FusedMoE):
351
+ elif isinstance(layer, RoutedExperts):
352
return HummingMoEMethod(quant_config, layer.moe_config)
353
return None
354
@@ -764,7 +765,7 @@ class HummingMoEMethod(FusedMoEMethodBase):
765
766
return get_humming_moe_quant_config(layer)
767
767
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
768
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
769
if getattr(self, "processed", False):
770
return
771
self.processed = True
@@ -898,10 +899,11 @@ class HummingMoEMethod(FusedMoEMethodBase):
899
900
def apply(
901
self,
901
- layer: FusedMoE,
902
+ layer: RoutedExperts,
903
x: torch.Tensor,
904
topk_weights: torch.Tensor,
905
topk_ids: torch.Tensor,
906
+ shared_experts: SharedExperts | None,
907
shared_experts_input: torch.Tensor | None,
908
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
909
workspace1, workspace2, output = self.experts.make_workspaces(
vllm/model_executor/layers/quantization/inc.py
+5
-6
@@ -9,6 +9,7 @@ import torch
9
from torch.nn.parameter import Parameter
10
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.fused_moe import RoutedExperts
13
from vllm.model_executor.layers.linear import (
14
LinearBase,
15
LinearMethodBase,
@@ -234,7 +235,6 @@ class INCConfig(QuantizationConfig):
235
self.extra_config = hf_to_vllm_mapper.apply_dict(self.extra_config)
236
237
def apply_awq_quant_layer(self, layer, prefix: str, backend: str = "auto"):
237
- from vllm.model_executor.layers.fused_moe import FusedMoE
238
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
239
check_marlin_supported,
240
check_moe_marlin_supports_layer,
@@ -264,7 +264,7 @@ class INCConfig(QuantizationConfig):
264
AWQ_TYPE_MAP[weight_bits], group_size, not sym
265
)
266
267
- if isinstance(layer, FusedMoE):
267
+ if isinstance(layer, RoutedExperts):
268
use_marlin = use_marlin and check_moe_marlin_supports_layer(
269
layer, group_size
270
)
@@ -298,7 +298,7 @@ class INCConfig(QuantizationConfig):
298
zero_point=not sym,
299
)
300
301
- if isinstance(layer, FusedMoE):
301
+ if isinstance(layer, RoutedExperts):
302
if use_marlin:
303
return AWQMarlinMoEMethod(quant_args_marlin, layer.moe_config)
304
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
@@ -320,7 +320,6 @@ class INCConfig(QuantizationConfig):
320
return None
321
322
def apply_gptq_quant_layer(self, layer, prefix: str, backend: str = "auto"):
323
- from vllm.model_executor.layers.fused_moe import FusedMoE
323
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
324
check_marlin_supported,
325
check_moe_marlin_supports_layer,
@@ -349,7 +348,7 @@ class INCConfig(QuantizationConfig):
348
use_marlin = (weight_bits, sym) in GPTQ_TYPE_MAP and check_marlin_supported(
349
GPTQ_TYPE_MAP[(weight_bits, sym)], group_size, has_zp=not sym
350
)
352
- if isinstance(layer, FusedMoE):
351
+ if isinstance(layer, RoutedExperts):
352
use_marlin = use_marlin and check_moe_marlin_supports_layer(
353
layer, group_size
354
)
@@ -385,7 +384,7 @@ class INCConfig(QuantizationConfig):
384
dynamic={},
385
)
386
388
- if isinstance(layer, FusedMoE):
387
+ if isinstance(layer, RoutedExperts):
388
if use_marlin:
389
return GPTQMarlinMoEMethod(quant_args_marlin, layer.moe_config)
390
else:
vllm/model_executor/layers/quantization/modelopt.py
+37
-37
@@ -18,18 +18,15 @@ from vllm.model_executor.kernels.linear import (
18
init_nvfp4_linear_kernel,
19
)
20
from vllm.model_executor.layers.attention import Attention, MLAAttention
21
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
22
-from vllm.model_executor.layers.fused_moe.config import (
21
+from vllm.model_executor.layers.fused_moe import (
22
FusedMoEConfig,
24
- FusedMoEQuantConfig,
25
- RoutingMethodType,
26
-)
27
-from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
23
FusedMoEMethodBase,
29
-)
30
-from vllm.model_executor.layers.fused_moe.layer import (
31
- FusedMoE,
24
+ FusedMoEQuantConfig,
25
FusedMoeWeightScaleSupported,
26
+ MoEActivation,
27
+ RoutedExperts,
28
+ RoutingMethodType,
29
+ SharedExperts,
30
)
31
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
32
Fp8MoeBackend,
@@ -208,7 +205,7 @@ class ModelOptQuantConfigBase(QuantizationConfig):
205
if getattr(quant_method, "backend", "") == "marlin":
206
quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
207
return quant_method
211
- elif isinstance(layer, FusedMoE):
208
+ elif isinstance(layer, RoutedExperts):
209
quant_method = self.FusedMoEMethodCls(
210
quant_config=self, moe_config=layer.moe_config
211
)
@@ -779,7 +776,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
776
def select_gemm_impl(
777
self,
778
prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
782
- layer: torch.nn.Module,
779
+ layer: RoutedExperts,
780
) -> mk.FusedMoEExpertsModular:
781
raise ValueError(
782
f"{self.__class__.__name__} uses the new modular kernel initialization "
@@ -788,7 +785,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
785
786
def create_weights(
787
self,
791
- layer: torch.nn.Module,
788
+ layer: RoutedExperts,
789
num_experts: int,
790
hidden_size: int,
791
intermediate_size_per_partition: int,
@@ -867,7 +864,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
864
865
def _setup_kernel(
866
self,
870
- layer: FusedMoE,
867
+ layer: RoutedExperts,
868
w13: torch.Tensor,
869
w2: torch.Tensor,
870
w13_scale: torch.Tensor,
@@ -902,10 +899,9 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
899
fp8_backend=self.fp8_backend,
900
experts_cls=self.experts_cls,
901
routing_tables=layer._expert_routing_tables(),
905
- shared_experts=layer.shared_experts,
902
)
903
908
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
904
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
905
w13 = layer.w13_weight
906
w2 = layer.w2_weight
907
w13_scale = layer.w13_weight_scale
@@ -936,7 +932,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
932
layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
933
)
934
939
- def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
935
+ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
936
w1_scale = layer.w13_weight_scale
937
w2_scale = layer.w2_weight_scale
938
a1_scale = layer.w13_input_scale
@@ -952,7 +948,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
948
949
def apply_monolithic(
950
self,
955
- layer: FusedMoE,
951
+ layer: RoutedExperts,
952
x: torch.Tensor,
953
router_logits: torch.Tensor,
954
input_ids: torch.Tensor | None = None,
@@ -976,10 +972,11 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
972
973
def apply(
974
self,
979
- layer: FusedMoE,
975
+ layer: RoutedExperts,
976
x: torch.Tensor,
977
topk_weights: torch.Tensor,
978
topk_ids: torch.Tensor,
979
+ shared_experts: SharedExperts | None,
980
shared_experts_input: torch.Tensor | None,
981
) -> torch.Tensor:
982
assert not self.is_monolithic
@@ -994,6 +991,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
991
global_num_experts=layer.global_num_experts,
992
expert_map=layer.expert_map,
993
apply_router_weight_on_input=layer.apply_router_weight_on_input,
994
+ shared_experts=shared_experts,
995
shared_experts_input=shared_experts_input,
996
)
997
@@ -1421,7 +1419,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1419
1420
def create_weights(
1421
self,
1424
- layer: torch.nn.Module,
1422
+ layer: RoutedExperts,
1423
num_experts: int,
1424
hidden_size: int,
1425
intermediate_size_per_partition: int,
@@ -1535,7 +1533,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1533
)
1534
layer.register_parameter("w2_input_scale", w2_input_scale)
1535
1538
- def process_weights_after_loading(self, layer: FusedMoE) -> None:
1536
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
1537
"""
1538
Convert NVFP4 MoE weights into kernel format and setup the kernel.
1539
"""
@@ -1589,12 +1587,11 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1587
moe_quant_config=self.moe_quant_config,
1588
moe_config=self.moe,
1589
experts_cls=self.experts_cls,
1592
- shared_experts=layer.shared_experts,
1590
routing_tables=layer._expert_routing_tables(),
1591
)
1592
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
1593
1597
- def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
1594
+ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
1595
return make_nvfp4_moe_quant_config(
1596
backend=self.nvfp4_backend,
1597
w13_scale=layer.w13_weight_scale,
@@ -1611,7 +1608,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1608
1609
def apply_monolithic(
1610
self,
1614
- layer: FusedMoE,
1611
+ layer: RoutedExperts,
1612
x: torch.Tensor,
1613
router_logits: torch.Tensor,
1614
input_ids: torch.Tensor | None = None,
@@ -1635,10 +1632,11 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1632
1633
def apply(
1634
self,
1638
- layer: FusedMoE,
1635
+ layer: RoutedExperts,
1636
x: torch.Tensor,
1637
topk_weights: torch.Tensor,
1638
topk_ids: torch.Tensor,
1639
+ shared_experts: SharedExperts | None,
1640
shared_experts_input: torch.Tensor | None,
1641
) -> torch.Tensor:
1642
assert not self.is_monolithic
@@ -1653,6 +1651,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1651
global_num_experts=layer.global_num_experts,
1652
expert_map=layer.expert_map,
1653
apply_router_weight_on_input=layer.apply_router_weight_on_input,
1654
+ shared_experts=shared_experts,
1655
shared_experts_input=shared_experts_input,
1656
)
1657
@@ -1860,15 +1859,15 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1859
1860
def create_weights(
1861
self,
1863
- layer: torch.nn.Module,
1862
+ layer: RoutedExperts,
1863
num_experts: int,
1864
hidden_size: int,
1865
intermediate_size_per_partition: int,
1866
params_dtype: torch.dtype,
1867
**extra_weight_attrs,
1868
):
1870
- layer.intermediate_size_per_partition = intermediate_size_per_partition
1871
- layer.hidden_size = hidden_size
1869
+ assert layer.intermediate_size_per_partition == intermediate_size_per_partition
1870
+ assert layer.hidden_size == hidden_size
1871
layer.orig_dtype = params_dtype
1872
1873
if hidden_size % MXFP8_BLOCK_SIZE != 0:
@@ -2051,7 +2050,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2050
torch.stack(w2_scale_shuffled).contiguous(),
2051
)
2052
2054
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
2053
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
2054
if getattr(layer, "_already_called_process_weights_after_loading", False):
2055
return
2056
@@ -2071,7 +2070,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2070
def select_gemm_impl(
2071
self,
2072
prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
2074
- layer: torch.nn.Module,
2073
+ layer: RoutedExperts,
2074
) -> mk.FusedMoEExpertsModular:
2075
raise ValueError(
2076
f"{self.__class__.__name__} uses the new modular kernel initialization "
@@ -2079,7 +2078,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2078
)
2079
2080
def get_fused_moe_quant_config(
2082
- self, layer: torch.nn.Module
2081
+ self, layer: RoutedExperts
2082
) -> FusedMoEQuantConfig | None:
2083
# TRTLLM MXFP8 path is monolithic and does not use modular kernel config.
2084
return None
@@ -2090,11 +2089,11 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2089
2090
def apply_monolithic(
2091
self,
2093
- layer: FusedMoE,
2092
+ layer: RoutedExperts,
2093
x: torch.Tensor,
2094
router_logits: torch.Tensor,
2095
input_ids: torch.Tensor | None = None,
2097
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
2096
+ ) -> torch.Tensor:
2097
from flashinfer.fused_moe.core import (
2098
ActivationType,
2099
Fp8QuantizationType,
@@ -2174,12 +2173,13 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2173
2174
def apply(
2175
self,
2177
- layer: FusedMoE,
2176
+ layer: RoutedExperts,
2177
x: torch.Tensor,
2178
topk_weights: torch.Tensor,
2179
topk_ids: torch.Tensor,
2180
+ shared_experts: SharedExperts | None,
2181
shared_experts_input: torch.Tensor | None,
2182
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
2182
+ ) -> torch.Tensor:
2183
assert not self.is_monolithic
2184
raise NotImplementedError(
2185
"Non-monolithic MXFP8 MoE path is not yet implemented."
@@ -2295,7 +2295,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2295
Tries three strategies in order:
2296
1. Direct lookup in ``quantized_layers``.
2297
2. Packed/fused-layer lookup (unfuse via ``packed_modules_mapping``).
2298
- 3. Prefix-based lookup for FusedMoE (any child key starts with
2298
+ 3. Prefix-based lookup for RoutedExperts (any child key starts with
2299
``prefix + "."``).
2300
2301
Returns the upper-cased quant_algo string, or *None* if the prefix
@@ -2322,7 +2322,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2322
f"{algos}. All shards must use the same quantization."
2323
)
2324
2325
- # 3. Prefix-based lookup (for FusedMoE / parent modules)
2325
+ # 3. Prefix-based lookup (for RoutedExperts / parent modules)
2326
prefix_dot = prefix + "."
2327
for key, info in self.quantized_layers.items():
2328
if key.startswith(prefix_dot):
@@ -2356,7 +2356,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2356
# Layer not in quantized_layers — leave unquantized
2357
return UnquantizedLinearMethod()
2358
2359
- if isinstance(layer, FusedMoE):
2359
+ if isinstance(layer, RoutedExperts):
2360
if quant_algo == "FP8":
2361
return ModelOptFp8MoEMethod(
2362
quant_config=self.fp8_config,
vllm/model_executor/layers/quantization/moe_wna16.py
+13
-11
@@ -6,18 +6,19 @@ from typing import Any
6
import torch
7
8
from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
9
+from vllm.model_executor.layers.fused_moe import (
10
+ FusedMoEConfig,
11
+ FusedMoEMethodBase,
12
+ FusedMoeWeightScaleSupported,
13
+ RoutedExperts,
14
+ SharedExperts,
15
+)
16
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
17
from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEQuantConfig,
19
int4_w4a16_moe_quant_config,
20
int8_w8a16_moe_quant_config,
21
)
15
-from vllm.model_executor.layers.fused_moe.layer import (
16
- FusedMoE,
17
- FusedMoEConfig,
18
- FusedMoEMethodBase,
19
- FusedMoeWeightScaleSupported,
20
-)
22
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
23
UnquantizedFusedMoEMethod,
24
)
@@ -166,7 +167,7 @@ class MoeWNA16Config(QuantizationConfig):
167
self, layer: torch.nn.Module, prefix: str
168
) -> "QuantizeMethodBase | None":
169
if is_layer_skipped_quant(prefix, self.modules_to_not_convert):
169
- if isinstance(layer, FusedMoE):
170
+ if isinstance(layer, RoutedExperts):
171
return UnquantizedFusedMoEMethod(layer.moe_config)
172
return UnquantizedLinearMethod()
173
elif isinstance(layer, LinearBase):
@@ -202,7 +203,7 @@ class MoeWNA16Config(QuantizationConfig):
203
)
204
else:
205
raise ValueError("moe_wna16 only support gptq and awq.")
205
- elif isinstance(layer, FusedMoE):
206
+ elif isinstance(layer, RoutedExperts):
207
return MoeWNA16Method(self, layer.moe_config)
208
return None
209
@@ -224,7 +225,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
225
226
def create_weights(
227
self,
227
- layer: torch.nn.Module,
228
+ layer: RoutedExperts,
229
num_experts: int,
230
hidden_size: int,
231
intermediate_size_per_partition: int,
@@ -343,7 +344,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
344
set_weight_attrs(param, extra_weight_attrs)
345
346
def get_fused_moe_quant_config(
346
- self, layer: torch.nn.Module
347
+ self, layer: RoutedExperts
348
) -> FusedMoEQuantConfig | None:
349
weight_bits = self.quant_config.weight_bits
350
has_zp = self.quant_config.has_zp
@@ -364,10 +365,11 @@ class MoeWNA16Method(FusedMoEMethodBase):
365
366
def apply(
367
self,
367
- layer: FusedMoE,
368
+ layer: RoutedExperts,
369
x: torch.Tensor,
370
topk_weights: torch.Tensor,
371
topk_ids: torch.Tensor,
372
+ shared_experts: SharedExperts | None,
373
shared_experts_input: torch.Tensor | None,
374
) -> torch.Tensor:
375
from vllm.model_executor.layers.fused_moe import fused_experts
vllm/model_executor/layers/quantization/mxfp4.py
+22
-20
@@ -7,15 +7,14 @@ from vllm.config import get_current_vllm_config
7
from vllm.logger import init_logger
8
from vllm.model_executor.layers.attention import Attention
9
from vllm.model_executor.layers.fused_moe import (
10
- FusedMoE,
10
FusedMoEConfig,
11
FusedMoEMethodBase,
13
-)
14
-from vllm.model_executor.layers.fused_moe import modular_kernel as mk
15
-from vllm.model_executor.layers.fused_moe.config import (
12
FusedMoEParallelConfig,
13
FusedMoEQuantConfig,
14
+ RoutedExperts,
15
+ SharedExperts,
16
)
17
+from vllm.model_executor.layers.fused_moe import modular_kernel as mk
18
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
19
TRITON_BACKENDS,
20
Mxfp4MoeBackend,
@@ -87,7 +86,7 @@ class Mxfp4Config(QuantizationConfig):
86
"UnquantizedLinearMethod.",
87
)
88
return UnquantizedLinearMethod()
90
- elif isinstance(layer, FusedMoE):
89
+ elif isinstance(layer, RoutedExperts):
90
return GptOssMxfp4MoEMethod(layer.moe_config)
91
elif isinstance(layer, Attention):
92
logger.debug_once(
@@ -178,7 +177,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
177
178
def create_weights(
179
self,
181
- layer: torch.nn.Module,
180
+ layer: RoutedExperts,
181
num_experts: int,
182
hidden_size: int,
183
intermediate_size_per_partition: int,
@@ -272,7 +271,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
271
272
def _setup_kernel(
273
self,
275
- layer: FusedMoE,
274
+ layer: RoutedExperts,
275
w13: torch.Tensor,
276
w2: torch.Tensor,
277
w13_scale: torch.Tensor,
@@ -370,11 +369,10 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
369
mxfp4_backend=self.mxfp4_backend,
370
experts_cls=self.experts_cls,
371
routing_tables=layer._expert_routing_tables(),
373
- shared_experts=layer.shared_experts,
372
layer=layer,
373
)
374
377
- def process_weights_after_loading(self, layer):
375
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
376
w13 = layer.w13_weight
377
w2 = layer.w2_weight
378
w13_scale = layer.w13_weight_scale
@@ -388,7 +386,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
386
self._setup_kernel(layer, w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
387
388
def get_fused_moe_quant_config(
391
- self, layer: torch.nn.Module
389
+ self, layer: RoutedExperts
390
) -> FusedMoEQuantConfig | None:
391
w1_scale = layer.w13_weight_scale
392
w2_scale = layer.w2_weight_scale
@@ -416,7 +414,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
414
def select_gemm_impl(
415
self,
416
prepare_finalize: mk.FusedMoEPrepareAndFinalize,
419
- layer: torch.nn.Module,
417
+ layer: RoutedExperts,
418
) -> mk.FusedMoEExpertsModular:
419
raise ValueError(
420
f"{self.__class__.__name__} uses the new modular kernel "
@@ -425,10 +423,11 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
423
424
def apply(
425
self,
428
- layer: FusedMoE,
426
+ layer: RoutedExperts,
427
x: torch.Tensor,
428
topk_weights: torch.Tensor,
429
topk_ids: torch.Tensor,
430
+ shared_experts: SharedExperts | None,
431
shared_experts_input: torch.Tensor | None,
432
) -> torch.Tensor:
433
assert not self.is_monolithic
@@ -443,12 +442,13 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
442
global_num_experts=layer.global_num_experts,
443
apply_router_weight_on_input=layer.apply_router_weight_on_input,
444
expert_map=layer.expert_map,
445
+ shared_experts=shared_experts,
446
shared_experts_input=shared_experts_input,
447
)
448
449
def apply_monolithic(
450
self,
451
- layer: FusedMoE,
451
+ layer: RoutedExperts,
452
x: torch.Tensor,
453
router_logits: torch.Tensor,
454
input_ids: torch.Tensor | None = None,
@@ -511,7 +511,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
511
512
def create_weights(
513
self,
514
- layer: torch.nn.Module,
514
+ layer: RoutedExperts,
515
num_experts: int,
516
hidden_size: int,
517
intermediate_size_per_partition: int,
@@ -605,7 +605,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
605
606
def _setup_kernel(
607
self,
608
- layer: FusedMoE,
608
+ layer: RoutedExperts,
609
w13: torch.Tensor,
610
w2: torch.Tensor,
611
w13_scale: torch.Tensor,
@@ -703,7 +703,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
703
mxfp4_backend=self.mxfp4_backend,
704
experts_cls=self.experts_cls,
705
routing_tables=layer._expert_routing_tables(),
706
- shared_experts=layer.shared_experts,
706
layer=layer,
707
)
708
@@ -721,7 +720,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
720
self._setup_kernel(layer, w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
721
722
def get_fused_moe_quant_config(
724
- self, layer: torch.nn.Module
723
+ self,
724
+ layer: RoutedExperts,
725
) -> FusedMoEQuantConfig | None:
726
w1_scale = layer.w13_weight_scale
727
w2_scale = layer.w2_weight_scale
@@ -748,7 +748,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
748
def select_gemm_impl(
749
self,
750
prepare_finalize: mk.FusedMoEPrepareAndFinalize,
751
- layer: torch.nn.Module,
751
+ layer: RoutedExperts,
752
) -> mk.FusedMoEExpertsModular:
753
raise ValueError(
754
f"{self.__class__.__name__} uses the new modular kernel "
@@ -757,10 +757,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
757
758
def apply(
759
self,
760
- layer: FusedMoE,
760
+ layer: RoutedExperts,
761
x: torch.Tensor,
762
topk_weights: torch.Tensor,
763
topk_ids: torch.Tensor,
764
+ shared_experts: SharedExperts | None,
765
shared_experts_input: torch.Tensor | None,
766
) -> torch.Tensor:
767
assert not self.is_monolithic
@@ -775,12 +776,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
776
global_num_experts=layer.global_num_experts,
777
apply_router_weight_on_input=layer.apply_router_weight_on_input,
778
expert_map=layer.expert_map,
779
+ shared_experts=shared_experts,
780
shared_experts_input=shared_experts_input,
781
)
782
783
def apply_monolithic(
784
self,
783
- layer: FusedMoE,
785
+ layer: RoutedExperts,
786
x: torch.Tensor,
787
router_logits: torch.Tensor,
788
input_ids: torch.Tensor | None = None,
vllm/model_executor/layers/quantization/online/base.py
+2
-2
@@ -11,7 +11,7 @@ from vllm.config.quantization import (
11
)
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.fused_moe import (
14
- FusedMoE,
14
+ RoutedExperts,
15
)
16
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
17
UnquantizedFusedMoEMethod,
@@ -118,7 +118,7 @@ class OnlineQuantizationConfig(QuantizationConfig):
118
return Mxfp8OnlineLinearMethod()
119
else:
120
return Fp8PerTensorOnlineLinearMethod()
121
- elif isinstance(layer, FusedMoE):
121
+ elif isinstance(layer, RoutedExperts):
122
if should_ignore_layer(
123
prefix,
124
ignore=self.ignored_layers,
vllm/model_executor/layers/quantization/online/fp8.py
+2
-3
@@ -8,7 +8,6 @@ from torch.nn import Module
8
9
if TYPE_CHECKING:
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
- from vllm.model_executor.layers.fused_moe import FusedMoE
11
from vllm.model_executor.layers.fused_moe.config import (
12
FusedMoEQuantConfig,
13
)
@@ -21,6 +20,7 @@ from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
20
from vllm.model_executor.kernels.linear.scaled_mm import (
21
CutlassFP8ScaledMMLinearKernel,
22
)
23
+from vllm.model_executor.layers.fused_moe import RoutedExperts
24
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
25
select_fp8_moe_backend,
26
)
@@ -314,7 +314,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
314
315
def _setup_kernel(
316
self,
317
- layer: "FusedMoE",
317
+ layer: RoutedExperts,
318
w13: torch.Tensor,
319
w2: torch.Tensor,
320
w13_scale: torch.Tensor,
@@ -355,7 +355,6 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
355
fp8_backend=self.fp8_backend,
356
experts_cls=self.experts_cls,
357
routing_tables=layer._expert_routing_tables(),
358
- shared_experts=layer.shared_experts,
358
)
359
360
def get_fused_moe_quant_config(
vllm/model_executor/layers/quantization/online/int8.py
+2
-3
@@ -7,11 +7,11 @@ import torch
7
from torch.nn import Module
8
9
if TYPE_CHECKING:
10
- from vllm.model_executor.layers.fused_moe import FusedMoE
10
from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
12
)
13
14
+from vllm.model_executor.layers.fused_moe import RoutedExperts
15
from vllm.model_executor.layers.fused_moe.oracle.int8 import (
16
make_int8_moe_kernel,
17
make_int8_moe_quant_config,
@@ -91,7 +91,7 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
91
replace_parameter(layer, "w13_scale", w13_scale)
92
replace_parameter(layer, "w2_scale", w2_scale)
93
94
- def _setup_kernel(self, layer: "FusedMoE") -> None:
94
+ def _setup_kernel(self, layer: RoutedExperts) -> None:
95
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
96
assert self.moe_quant_config is not None
97
assert self.experts_cls is not None
@@ -100,7 +100,6 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
100
moe_config=self.moe,
101
experts_cls=self.experts_cls,
102
routing_tables=layer._expert_routing_tables(),
103
- shared_experts=layer.shared_experts,
103
)
104
105
def get_fused_moe_quant_config(
vllm/model_executor/layers/quantization/online/moe_base.py
+12
-6
@@ -6,8 +6,12 @@ from abc import abstractmethod
6
import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
-from vllm.model_executor.layers.fused_moe import FusedMoEMethodBase
10
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9
+from vllm.model_executor.layers.fused_moe import (
10
+ FusedMoEMethodBase,
11
+ FusedMoEQuantConfig,
12
+ RoutedExperts,
13
+ SharedExperts,
14
+)
15
from vllm.model_executor.model_loader.reload.layerwise import (
16
initialize_online_processing,
17
)
@@ -127,11 +131,11 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
131
132
def apply_monolithic(
133
self,
130
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
134
+ layer: RoutedExperts,
135
x: torch.Tensor,
136
router_logits: torch.Tensor,
137
input_ids: torch.Tensor | None = None,
134
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
138
+ ) -> torch.Tensor:
139
assert self.is_monolithic
140
assert self.moe_kernel is not None
141
return self.moe_kernel.apply_monolithic(
@@ -151,12 +155,13 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
155
156
def apply(
157
self,
154
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
158
+ layer: RoutedExperts,
159
x: torch.Tensor,
160
topk_weights: torch.Tensor,
161
topk_ids: torch.Tensor,
162
+ shared_experts: SharedExperts | None,
163
shared_experts_input: torch.Tensor | None,
159
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
164
+ ) -> torch.Tensor:
165
assert not self.is_monolithic
166
assert self.moe_kernel is not None
167
return self.moe_kernel.apply(
@@ -169,5 +174,6 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
174
global_num_experts=layer.global_num_experts,
175
expert_map=layer.expert_map,
176
apply_router_weight_on_input=layer.apply_router_weight_on_input,
177
+ shared_experts=shared_experts,
178
shared_experts_input=shared_experts_input,
179
)
vllm/model_executor/layers/quantization/online/mxfp8.py
+3
-4
@@ -10,9 +10,9 @@ from torch.nn import Module
10
11
if TYPE_CHECKING:
12
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
- from vllm.model_executor.layers.fused_moe import FusedMoE
14
- from vllm.model_executor.layers.fused_moe.config import (
13
+ from vllm.model_executor.layers.fused_moe import (
14
FusedMoEQuantConfig,
15
+ RoutedExperts,
16
)
17
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
18
@@ -161,7 +161,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
161
162
def _setup_kernel(
163
self,
164
- layer: "FusedMoE",
164
+ layer: "RoutedExperts",
165
w13: torch.Tensor,
166
w2: torch.Tensor,
167
w13_scale: torch.Tensor,
@@ -200,7 +200,6 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
200
fp8_backend=self.fp8_backend,
201
experts_cls=self.experts_cls,
202
routing_tables=layer._expert_routing_tables(),
203
- shared_experts=layer.shared_experts,
203
)
204
205
def get_fused_moe_quant_config(
vllm/model_executor/layers/quantization/quark/quark.py
+2
-2
@@ -8,7 +8,7 @@ import torch
8
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.attention import Attention
11
-from vllm.model_executor.layers.fused_moe import FusedMoE
11
+from vllm.model_executor.layers.fused_moe import RoutedExperts
12
from vllm.model_executor.layers.linear import (
13
LinearBase,
14
LinearMethodBase,
@@ -135,7 +135,7 @@ class QuarkConfig(QuantizationConfig):
135
if isinstance(layer, Attention):
136
return QuarkKVCacheMethod(self)
137
138
- if isinstance(layer, FusedMoE):
138
+ if isinstance(layer, RoutedExperts):
139
return QuarkMoEMethod.get_moe_method(self, module=layer, layer_name=prefix)
140
return None
141
vllm/model_executor/layers/quantization/quark/quark_moe.py
+25
-20
@@ -12,11 +12,12 @@ from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.config import get_current_vllm_config
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.fused_moe import (
15
- FusedMoE,
15
FusedMoEConfig,
16
FusedMoEMethodBase,
17
FusedMoeWeightScaleSupported,
18
MoEActivation,
19
+ RoutedExperts,
20
+ SharedExperts,
21
)
22
from vllm.model_executor.layers.fused_moe.config import (
23
FusedMoEParallelConfig,
@@ -73,7 +74,7 @@ class QuarkMoEMethod(FusedMoEMethodBase):
74
@staticmethod
75
def get_moe_method(
76
quant_config: "QuarkConfig", # type: ignore # noqa E501 # noqa F821
76
- module: torch.nn.Module,
77
+ module: RoutedExperts,
78
layer_name: str,
79
) -> "QuarkMoEMethod":
80
layer_quant_config = quant_config._find_matched_config(layer_name, module)
@@ -164,7 +165,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
165
166
def create_weights(
167
self,
167
- layer: torch.nn.Module,
168
+ layer: RoutedExperts,
169
num_experts: int,
170
hidden_size: int,
171
intermediate_size_per_partition: int,
@@ -220,7 +221,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
221
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
222
)
223
layer.register_parameter("w2_weight_scale", w2_weight_scale)
223
- # Add PER-TENSOR quantization for FusedMoE.weight_loader.
224
+ # Add PER-TENSOR quantization for RoutedExperts.weight_loader.
225
extra_weight_attrs.update(
226
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
227
)
@@ -242,7 +243,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
243
requires_grad=False,
244
)
245
layer.register_parameter("w2_weight_scale", w2_weight_scale)
245
- # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
246
+ # Add PER-CHANNEL quantization for RoutedExperts.weight_loader.
247
extra_weight_attrs.update(
248
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
249
)
@@ -287,7 +288,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
288
else:
289
layer.w13_bias, layer.w2_bias = None, None
290
290
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
291
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
292
# Fp8 moe kernels require a single activation scale.
293
# We take the max of all the scales in case they differ.
294
if self.static_input_scales:
@@ -422,7 +423,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
423
)
424
425
def get_fused_moe_quant_config(
425
- self, layer: torch.nn.Module
426
+ self, layer: RoutedExperts
427
) -> FusedMoEQuantConfig | None:
428
return fp8_w8a8_moe_quant_config(
429
w1_scale=layer.w13_weight_scale,
@@ -437,10 +438,11 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
438
439
def apply(
440
self,
440
- layer: FusedMoE,
441
+ layer: RoutedExperts,
442
x: torch.Tensor,
443
topk_weights: torch.Tensor,
444
topk_ids: torch.Tensor,
445
+ shared_experts: SharedExperts | None,
446
shared_experts_input: torch.Tensor | None,
447
) -> torch.Tensor:
448
if self.rocm_aiter_moe_enabled:
@@ -751,12 +753,13 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
753
754
def apply(
755
self,
754
- layer: FusedMoE,
756
+ layer: RoutedExperts,
757
x: torch.Tensor,
758
topk_weights: torch.Tensor,
759
topk_ids: torch.Tensor,
760
+ shared_experts: SharedExperts | None,
761
shared_experts_input: torch.Tensor | None,
759
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
762
+ ) -> torch.Tensor:
763
from vllm.model_executor.layers.fused_moe import fused_experts
764
765
return fused_experts(
@@ -791,7 +794,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
794
795
def create_weights(
796
self,
794
- layer: torch.nn.Module,
797
+ layer: RoutedExperts,
798
num_experts: int,
799
hidden_size: int,
800
intermediate_size_per_partition: int,
@@ -858,7 +861,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
861
set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
862
set_weight_attrs(w2_weight_scale_2, extra_weight_attrs)
863
861
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
864
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
865
shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
866
layer.w13_weight.data, layer.w2_weight.data
867
)
@@ -903,10 +906,11 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
906
907
def apply(
908
self,
906
- layer: FusedMoE,
909
+ layer: RoutedExperts,
910
x: torch.Tensor,
911
topk_weights: torch.Tensor,
912
topk_ids: torch.Tensor,
913
+ shared_experts: SharedExperts | None,
914
shared_experts_input: torch.Tensor | None,
915
) -> torch.Tensor:
916
from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
@@ -1093,7 +1097,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1097
1098
def create_weights(
1099
self,
1096
- layer: torch.nn.Module,
1100
+ layer: RoutedExperts,
1101
num_experts: int,
1102
hidden_size: int,
1103
intermediate_size_per_partition: int,
@@ -1198,7 +1202,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1202
layer.w13_input_scale = None
1203
layer.w2_input_scale = None
1204
1201
- def process_weights_after_loading(self, layer):
1205
+ def process_weights_after_loading(self, layer: RoutedExperts):
1206
# For MXFP4 schemes with native backend, use oracle
1207
if self.mxfp4_backend != Mxfp4MoeBackend.NONE:
1208
self._setup_kernel(layer)
@@ -1290,7 +1294,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1294
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1295
torch.accelerator.empty_cache()
1296
1293
- def _setup_kernel(self, layer: FusedMoE):
1297
+ def _setup_kernel(self, layer: RoutedExperts):
1298
"""Setup kernel using oracle functions for MXFP4 schemes (W4A16, W4A8)."""
1299
w13_bias = getattr(layer, "w13_bias", None)
1300
w2_bias = getattr(layer, "w2_bias", None)
@@ -1341,11 +1345,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1345
mxfp4_backend=self.mxfp4_backend,
1346
experts_cls=self.experts_cls,
1347
routing_tables=layer._expert_routing_tables(),
1344
- shared_experts=layer.shared_experts,
1348
)
1349
1350
def get_fused_moe_quant_config(
1348
- self, layer: torch.nn.Module
1351
+ self, layer: RoutedExperts
1352
) -> FusedMoEQuantConfig | None:
1353
# For oracle-based backends (W4A16, W4A8), use make_mxfp4_moe_quant_config
1354
if self.mxfp4_backend not in (Mxfp4MoeBackend.NONE, Mxfp4MoeBackend.EMULATION):
@@ -1413,10 +1416,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1416
1417
def apply(
1418
self,
1416
- layer: FusedMoE,
1419
+ layer: RoutedExperts,
1420
x: torch.Tensor,
1421
topk_weights: torch.Tensor,
1422
topk_ids: torch.Tensor,
1423
+ shared_experts: SharedExperts | None,
1424
shared_experts_input: torch.Tensor | None,
1425
) -> torch.Tensor:
1426
# For oracle-based kernels (W4A16, W4A8) or emulation kernel
@@ -1431,6 +1435,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1435
global_num_experts=layer.global_num_experts,
1436
apply_router_weight_on_input=layer.apply_router_weight_on_input,
1437
expert_map=layer.expert_map,
1438
+ shared_experts=shared_experts,
1439
shared_experts_input=shared_experts_input,
1440
)
1441
@@ -1454,7 +1459,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1459
1460
def apply_monolithic(
1461
self,
1457
- layer: FusedMoE,
1462
+ layer: RoutedExperts,
1463
x: torch.Tensor,
1464
router_logits: torch.Tensor,
1465
input_ids: torch.Tensor | None = None,
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+3
-3
@@ -21,7 +21,7 @@ from vllm.utils.flashinfer import (
21
)
22
23
if TYPE_CHECKING:
24
- from vllm.model_executor.layers.fused_moe.layer import FusedMoE
24
+ from vllm.model_executor.layers.fused_moe import RoutedExperts
25
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
26
NvFp4MoeBackend,
27
)
@@ -80,7 +80,7 @@ def interleave_linear_and_gate(
80
81
82
def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
83
- layer: "FusedMoE",
83
+ layer: "RoutedExperts",
84
w13: torch.Tensor,
85
w13_scale: torch.Tensor,
86
w13_scale_2: torch.Tensor,
@@ -286,7 +286,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(
286
287
def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
288
backend: "NvFp4MoeBackend",
289
- layer: "FusedMoE",
289
+ layer: "RoutedExperts",
290
w13: torch.Tensor,
291
w13_scale: torch.Tensor,
292
w13_scale_2: torch.Tensor,
vllm/model_executor/layers/quantization/utils/humming_utils.py
+3
-3
@@ -12,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.config import (
12
FusedMoEQuantConfig,
13
FusedMoEQuantDesc,
14
)
15
-from vllm.model_executor.layers.fused_moe.layer import FusedMoE
15
+from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
16
from vllm.model_executor.layers.linear import LinearBase
17
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
18
@@ -82,7 +82,7 @@ def prepare_humming_layer(layer: LinearBase, quant_config: dict):
82
HummingMethod.transform_humming_layer(layer)
83
84
85
-def prepare_humming_moe_layer(layer: FusedMoE, quant_config: dict):
85
+def prepare_humming_moe_layer(layer: RoutedExperts, quant_config: dict):
86
weight_schema = BaseWeightSchema.from_config(quant_config)
87
input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}
88
if humming_is_layer_skipped(input_quant_config, layer.layer_name):
@@ -164,7 +164,7 @@ def prepare_humming_moe_layer(layer: FusedMoE, quant_config: dict):
164
layer.register_buffer("locks", locks)
165
166
167
-def get_humming_moe_quant_config(layer: FusedMoE):
167
+def get_humming_moe_quant_config(layer: RoutedExperts):
168
input_schema = layer.input_schemas["w13"]
169
weight_schema = layer.weight_schemas["w13"]
170
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+2
-1
@@ -8,6 +8,7 @@ import torch
8
import vllm.envs as envs
9
from vllm import _custom_ops as ops
10
from vllm.logger import init_logger
11
+from vllm.model_executor.layers.fused_moe import RoutedExperts
12
from vllm.model_executor.layers.linear import LinearBase
13
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
14
from vllm.model_executor.layers.quantization.utils.int8_utils import (
@@ -226,7 +227,7 @@ def check_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool:
227
)[0]
228
229
229
-def check_moe_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool:
230
+def check_moe_marlin_supports_layer(layer: RoutedExperts, group_size: int) -> bool:
231
if current_platform.is_rocm():
232
return False
233
hidden_size = layer.hidden_size
vllm/model_executor/models/laguna.py
+5
-2
@@ -20,7 +20,10 @@ from vllm.distributed import (
20
)
21
from vllm.logger import init_logger
22
from vllm.model_executor.layers.attention import Attention
23
-from vllm.model_executor.layers.fused_moe import FusedMoE
23
+from vllm.model_executor.layers.fused_moe import (
24
+ FusedMoE,
25
+ fused_moe_make_expert_params_mapping,
26
+)
27
from vllm.model_executor.layers.layernorm import RMSNorm
28
from vllm.model_executor.layers.linear import (
29
ColumnParallelLinear,
@@ -666,7 +669,7 @@ class LagunaModel(nn.Module, EagleModelMixin):
669
Returns mapping tuples of (param_name, weight_name, expert_id, shard_id)
670
that handle both weights and quantization scales.
671
"""
669
- return FusedMoE.make_expert_params_mapping(
672
+ return fused_moe_make_expert_params_mapping(
673
self,
674
ckpt_gate_proj_name="gate_proj",
675
ckpt_down_proj_name="down_proj",