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, 22:45 UTC
[MoE] Move various experts classes to fused_moe/experts/ (#41979)
Signed-off-by: Jackmin801 <ongjackm@gmail.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Signed-off-by: Jackmin801 <56836461+Jackmin801@users.noreply.github.com> Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Jackmin801 <ongjackm@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Jackmin801 <56836461+Jackmin801@users.noreply.github.com>
bnellnm committed
May 10, 2026 at 19:54 UTC
1b57eb41f24166d0bee2031c53847b188089d77d
34 files changed
+586
-553
csrc/quantization/gguf/moe.cuh
+1
-1
@@ -1,7 +1,7 @@
1
#include <cstdint>
2
3
/* Adapted from ./csrc/quantization/gguf/mmq.cuh
4
- based on ./vllm/model_executor/layers/fused_moe/fused_moe.py */
4
+ based on ./vllm/model_executor/layers/fused_moe/experts/triton_moe.py */
5
template <typename scalar_t, int qk, int qr, int qi, bool need_sum,
6
typename block_q_t, int mmq_x, int mmq_y, int nwarps,
7
allocate_tiles_cuda_t allocate_tiles, load_tiles_cuda_t load_tiles,
docs/design/custom_op.md
+1
-1
@@ -138,7 +138,7 @@ For example:
138
139
--8<-- "vllm/model_executor/models/transformers/moe.py:transformers_fused_moe"
140
141
---8<-- "vllm/model_executor/layers/fused_moe/fused_moe.py:grouped_topk"
141
+--8<-- "vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py:grouped_topk"
142
```
143
144
**9. Norm:**
docs/design/moe_kernel_features.md
+3
-3
@@ -80,14 +80,14 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
80
81
| Kernel | Input act. format | Quant. types | Quant. format | Activation function | Apply Weight On Input | Modular | Source |
82
| ------ | ----------------- | ------------ | ------------- | ------------------- | --------------------- | ------- | ------ |
83
-| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
83
+| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.experts.triton_moe.TritonExperts] |
84
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
85
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
86
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
87
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
88
-| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
88
+| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe.FlashInferExperts] |
89
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
90
-| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
90
+| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] |
91
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
92
| rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
93
| cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+2
-2
@@ -237,7 +237,7 @@ if has_mori():
237
)
238
239
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
240
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
240
+ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
241
FlashInferExperts,
242
)
243
from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501
@@ -298,7 +298,7 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
298
)
299
300
if has_aiter():
301
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
301
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
302
AiterExperts,
303
)
304
tests/kernels/moe/test_flashinfer.py
+3
-3
@@ -18,12 +18,12 @@ from vllm.model_executor.layers.fused_moe.config import (
18
RoutingMethodType,
19
fp8_w8a8_moe_quant_config,
20
)
21
+from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
22
+ FlashInferExperts,
23
+)
24
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
25
TrtLlmFp8ExpertsMonolithic,
26
)
24
-from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
25
- FlashInferExperts,
26
-)
27
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
28
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
29
rotate_weights_for_fi_trtllm_fp8_per_tensor_moe,
tests/kernels/moe/test_flashinfer_moe.py
+1
-1
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
22
FusedMoEParallelConfig,
23
RoutingMethodType,
24
)
25
-from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
25
+from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
26
FlashInferExperts,
27
is_valid_flashinfer_cutlass_fused_moe,
28
)
tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py
+1
-1
@@ -5,7 +5,7 @@
5
import pytest
6
import torch
7
8
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
8
+from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
9
fused_marlin_moe,
10
)
11
from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
tests/kernels/moe/test_moe.py
+1
-1
@@ -32,7 +32,7 @@ from vllm.model_executor.layers.fused_moe.config import (
32
int4_w4a16_moe_quant_config,
33
int8_w8a16_moe_quant_config,
34
)
35
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
35
+from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
36
batched_fused_marlin_moe,
37
fused_marlin_moe,
38
)
tests/kernels/moe/test_rocm_aiter_topk.py
+1
-1
@@ -20,7 +20,7 @@ if not current_platform.is_rocm():
20
pytest.skip("This test can only run on ROCm.", allow_module_level=True)
21
22
# this import statement is needed to ensure the ops are registered
23
-import vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe # noqa: F401
23
+import vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe # noqa: F401
24
25
# need to import once to ensure the ops are registered
26
# Check if aiter package is installed
tests/kernels/moe/test_triton_moe_no_act_mul.py
+3
-3
@@ -15,7 +15,7 @@ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
15
from vllm.model_executor.layers.fused_moe.config import (
16
FUSED_MOE_UNQUANTIZED_CONFIG,
17
)
18
-from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
18
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
19
from vllm.platforms import current_platform
20
21
# Test parameters
@@ -151,7 +151,7 @@ def test_triton_experts_no_mul_activation(
151
@torch.inference_mode()
152
def test_workspace_shapes_no_mul_vs_gated():
153
"""Test that workspace shapes differ correctly between gated and non-gated."""
154
- from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
154
+ from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
155
156
M, N, K, topk = 64, 256, 128, 2
157
@@ -192,7 +192,7 @@ def test_workspace_shapes_no_mul_vs_gated():
192
@torch.inference_mode()
193
def test_adjust_n_for_activation():
194
"""Test the adjust_N_for_activation method."""
195
- from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
195
+ from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
196
197
experts = TritonExperts(
198
moe_config=make_dummy_moe_config(),
tests/kernels/moe/test_unquantized_backend_selection.py
+1
-1
@@ -158,7 +158,7 @@ def test_select_cuda_flashinfer_trtllm_backend(mock_is_supported_trtllm, monkeyp
158
return_value=(False, None),
159
)
160
@patch(
161
- "vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts.is_supported_config",
161
+ "vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe.FlashInferExperts.is_supported_config",
162
return_value=(True, None),
163
)
164
@pytest.mark.skipif(
tests/kernels/moe/utils.py
+3
-1
@@ -17,12 +17,14 @@ from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEQuantConfig,
18
RoutingMethodType,
19
)
20
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
21
+ TritonExperts,
22
+)
23
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
24
BatchedTritonExperts,
25
NaiveBatchedExperts,
26
)
27
from vllm.model_executor.layers.fused_moe.fused_moe import (
25
- TritonExperts,
28
fused_experts,
29
)
30
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
vllm/model_executor/layers/fused_moe/__init__.py
+7
-5
@@ -85,6 +85,13 @@ if HAS_TRITON:
85
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
86
DeepGemmExperts,
87
)
88
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
89
+ AiterExperts,
90
+ )
91
+ from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
92
+ TritonExperts,
93
+ TritonWNA16Experts,
94
+ )
95
from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
96
XPUExperts,
97
XPUExpertsFp8,
@@ -94,14 +101,9 @@ if HAS_TRITON:
101
BatchedTritonExperts,
102
)
103
from vllm.model_executor.layers.fused_moe.fused_moe import (
97
- TritonExperts,
98
- TritonWNA16Experts,
104
fused_experts,
105
get_config_file_name,
106
)
102
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
103
- AiterExperts,
104
- )
107
from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
108
fused_topk,
109
)
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
renamed
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
renamed
+1
-2
@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
29
from vllm.model_executor.layers.fused_moe.utils import (
30
_resize_cache,
31
disable_inplace,
32
+ swiglu_limit_func,
33
)
34
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
35
get_marlin_input_dtype,
@@ -50,8 +51,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
51
from vllm.platforms import current_platform
52
from vllm.scalar_type import ScalarType, scalar_types
53
53
-from .utils import swiglu_limit_func
54
-
54
55
def _fused_marlin_moe(
56
hidden_states: torch.Tensor,
vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py
+1
-1
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEConfig,
21
FusedMoEQuantConfig,
22
)
23
-from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
23
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
24
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
25
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
26
dequantize_to_dtype,
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
+1
-1
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEConfig,
21
FusedMoEQuantConfig,
22
)
23
-from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
23
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
24
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
25
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
26
from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
renamed
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
new
+522
@@ -0,0 +1,522 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Triton-based MoE expert implementations."""
4
+
5
+import torch
6
+
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm import _custom_ops as ops
9
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10
+from vllm.model_executor.layers.fused_moe.config import (
11
+ FusedMoEConfig,
12
+ FusedMoEParallelConfig,
13
+ FusedMoEQuantConfig,
14
+)
15
+from vllm.model_executor.layers.fused_moe.fused_moe import (
16
+ _prepare_expert_assignment,
17
+ invoke_fused_moe_triton_kernel,
18
+ invoke_fused_moe_wna16_triton_kernel,
19
+ try_get_optimal_moe_config,
20
+)
21
+from vllm.model_executor.layers.fused_moe.lora_experts_mixin import (
22
+ LoRAExpertsMixin,
23
+)
24
+from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
25
+ moe_align_block_size,
26
+)
27
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
28
+ TopKWeightAndReduceNoOP,
29
+)
30
+from vllm.model_executor.layers.fused_moe.utils import (
31
+ _resize_cache,
32
+ moe_kernel_quantize_input,
33
+)
34
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
35
+ QuantKey,
36
+ kFp8Dynamic128Sym,
37
+ kFp8DynamicTensorSym,
38
+ kFp8DynamicTokenSym,
39
+ kFp8Static128BlockSym,
40
+ kFp8StaticChannelSym,
41
+ kFp8StaticTensorSym,
42
+ kInt8DynamicTokenSym,
43
+ kInt8StaticChannelSym,
44
+)
45
+from vllm.platforms import current_platform
46
+from vllm.triton_utils import tl
47
+
48
+
49
+class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
50
+ """Triton-based fused MoE expert implementation."""
51
+
52
+ def __init__(
53
+ self,
54
+ moe_config: FusedMoEConfig,
55
+ quant_config: FusedMoEQuantConfig,
56
+ ):
57
+ # Whether quantized MOE runs natively, or through
58
+ # higher-precision + activation QDQ.
59
+ self.quantization_emulation = False
60
+ super().__init__(moe_config, quant_config)
61
+
62
+ @staticmethod
63
+ def activation_format() -> mk.FusedMoEActivationFormat:
64
+ return mk.FusedMoEActivationFormat.Standard
65
+
66
+ @staticmethod
67
+ def _supports_current_device() -> bool:
68
+ return current_platform.is_cuda_alike() or current_platform.is_xpu()
69
+
70
+ @staticmethod
71
+ def _supports_no_act_and_mul() -> bool:
72
+ return True
73
+
74
+ @staticmethod
75
+ def _supports_quant_scheme(
76
+ weight_key: QuantKey | None,
77
+ activation_key: QuantKey | None,
78
+ ) -> bool:
79
+ # INT8 requires at least 7.5 (Turing).
80
+ device_supports_int8 = (
81
+ current_platform.is_cuda()
82
+ and current_platform.has_device_capability((7, 5))
83
+ )
84
+
85
+ supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
86
+ if device_supports_int8:
87
+ supported.append((kInt8StaticChannelSym, kInt8DynamicTokenSym))
88
+ if current_platform.supports_fp8():
89
+ supported += [
90
+ (kFp8Static128BlockSym, kFp8Dynamic128Sym),
91
+ (kFp8StaticChannelSym, kFp8DynamicTokenSym),
92
+ (kFp8StaticTensorSym, kFp8DynamicTokenSym),
93
+ (kFp8StaticTensorSym, kFp8StaticTensorSym),
94
+ (kFp8StaticTensorSym, kFp8DynamicTensorSym),
95
+ ]
96
+ return (weight_key, activation_key) in supported
97
+
98
+ @staticmethod
99
+ def _supports_activation(activation: MoEActivation) -> bool:
100
+ return activation in [
101
+ MoEActivation.SILU,
102
+ MoEActivation.GELU,
103
+ MoEActivation.GELU_TANH,
104
+ MoEActivation.SWIGLUOAI,
105
+ MoEActivation.SWIGLUSTEP,
106
+ MoEActivation.SILU_NO_MUL,
107
+ MoEActivation.GELU_NO_MUL,
108
+ MoEActivation.GELU_TANH_NO_MUL,
109
+ MoEActivation.RELU2_NO_MUL,
110
+ ]
111
+
112
+ @staticmethod
113
+ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
114
+ return not (
115
+ moe_parallel_config.use_fi_nvl_two_sided_kernels
116
+ or moe_parallel_config.use_fi_nvl_one_sided_kernels
117
+ )
118
+
119
+ @staticmethod
120
+ def _supports_batch_invariance():
121
+ return True
122
+
123
+ def supports_expert_map(self) -> bool:
124
+ return True
125
+
126
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
127
+ return TopKWeightAndReduceNoOP()
128
+
129
+ def workspace_shapes(
130
+ self,
131
+ M: int,
132
+ N: int,
133
+ K: int,
134
+ topk: int,
135
+ global_num_experts: int,
136
+ local_num_experts: int,
137
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
138
+ activation: MoEActivation,
139
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
140
+ activation_out_dim = self.adjust_N_for_activation(N, activation)
141
+ workspace1 = (M, topk, max(activation_out_dim, K))
142
+ workspace2 = (M, topk, max(N, K))
143
+ output = (M, K)
144
+ return (workspace1, workspace2, output)
145
+
146
+ def apply(
147
+ self,
148
+ output: torch.Tensor,
149
+ hidden_states: torch.Tensor,
150
+ w1: torch.Tensor,
151
+ w2: torch.Tensor,
152
+ topk_weights: torch.Tensor,
153
+ topk_ids: torch.Tensor,
154
+ activation: MoEActivation,
155
+ global_num_experts: int,
156
+ expert_map: torch.Tensor | None,
157
+ a1q_scale: torch.Tensor | None,
158
+ a2_scale: torch.Tensor | None,
159
+ workspace13: torch.Tensor,
160
+ workspace2: torch.Tensor,
161
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
162
+ apply_router_weight_on_input: bool,
163
+ ):
164
+ # Check constraints.
165
+ if self.quant_config.use_int4_w4a16:
166
+ assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
167
+ else:
168
+ assert hidden_states.size(-1) == w1.size(2), (
169
+ f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
170
+ )
171
+
172
+ assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
173
+ assert hidden_states.dim() == 2
174
+ assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
175
+ assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
176
+ assert hidden_states.dtype in [
177
+ torch.float32,
178
+ torch.float16,
179
+ torch.bfloat16,
180
+ torch.float8_e4m3fn,
181
+ torch.float8_e4m3fnuz,
182
+ ]
183
+
184
+ E, num_tokens, N, K, top_k_num = self.moe_problem_size(
185
+ hidden_states, w1, w2, topk_ids
186
+ )
187
+
188
+ if global_num_experts == -1:
189
+ global_num_experts = E
190
+
191
+ config = try_get_optimal_moe_config(
192
+ w1.size(),
193
+ w2.size(),
194
+ top_k_num,
195
+ self.quant_config.config_name(hidden_states.dtype),
196
+ num_tokens,
197
+ block_shape=self.block_shape,
198
+ )
199
+
200
+ if hidden_states.dtype == torch.bfloat16:
201
+ compute_type = tl.bfloat16
202
+ elif hidden_states.dtype == torch.float16:
203
+ compute_type = tl.float16
204
+ elif hidden_states.dtype == torch.float32:
205
+ compute_type = tl.float32
206
+ elif (
207
+ hidden_states.dtype == torch.float8_e4m3fn
208
+ or hidden_states.dtype == torch.float8_e4m3fnuz
209
+ ):
210
+ compute_type = tl.bfloat16
211
+ else:
212
+ raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
213
+
214
+ # Note that the output tensor might be in workspace1
215
+ intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
216
+ cache2_dim = self.adjust_N_for_activation(N, activation)
217
+ intermediate_cache2 = _resize_cache(
218
+ workspace13, (num_tokens * top_k_num, cache2_dim)
219
+ )
220
+ intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
221
+
222
+ sorted_token_ids, expert_ids, num_tokens_post_padded = (
223
+ _prepare_expert_assignment(
224
+ topk_ids,
225
+ config,
226
+ num_tokens,
227
+ top_k_num,
228
+ global_num_experts,
229
+ expert_map,
230
+ use_int8_w8a16=self.quant_config.use_int8_w8a16,
231
+ use_int4_w4a16=self.quant_config.use_int4_w4a16,
232
+ block_shape=self.block_shape,
233
+ )
234
+ )
235
+
236
+ invoke_fused_moe_triton_kernel(
237
+ hidden_states,
238
+ w1,
239
+ intermediate_cache1,
240
+ a1q_scale,
241
+ self.w1_scale,
242
+ None, # topk_weights
243
+ sorted_token_ids,
244
+ expert_ids,
245
+ num_tokens_post_padded,
246
+ False, # mul_routed_weights
247
+ top_k_num,
248
+ config,
249
+ compute_type=compute_type,
250
+ use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
251
+ use_int8_w8a8=self.quant_config.use_int8_w8a8,
252
+ use_int8_w8a16=self.quant_config.use_int8_w8a16,
253
+ use_int4_w4a16=self.quant_config.use_int4_w4a16,
254
+ per_channel_quant=self.per_act_token_quant,
255
+ block_shape=self.block_shape,
256
+ B_bias=self.w1_bias,
257
+ )
258
+
259
+ # LoRA w13: applied to intermediate_cache1 before activation, using
260
+ # hidden_states as the lora_a input. moe_lora_align_block_size is
261
+ # called once here and results reused for the w2 LoRA below.
262
+ sorted_token_ids_lora = None
263
+ expert_ids_lora = None
264
+ num_tokens_post_padded_lora = None
265
+ token_lora_mapping = None
266
+ lora_context = self._lora_context
267
+ if lora_context is not None:
268
+ (
269
+ sorted_token_ids_lora,
270
+ expert_ids_lora,
271
+ num_tokens_post_padded_lora,
272
+ token_lora_mapping,
273
+ ) = self.apply_w13_lora(
274
+ lora_context,
275
+ y=intermediate_cache1,
276
+ x=hidden_states,
277
+ topk_ids=topk_ids,
278
+ topk_weights=topk_weights,
279
+ expert_map=expert_map,
280
+ w1=w1,
281
+ w2=w2,
282
+ num_tokens=num_tokens,
283
+ top_k_num=top_k_num,
284
+ )
285
+
286
+ self.activation(
287
+ activation, intermediate_cache2, intermediate_cache1.view(-1, N)
288
+ )
289
+
290
+ a2q_scale: torch.Tensor | None = None
291
+
292
+ qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
293
+ intermediate_cache2,
294
+ a2_scale,
295
+ self.quant_dtype,
296
+ self.per_act_token_quant,
297
+ self.block_shape,
298
+ quantization_emulation=self.quantization_emulation,
299
+ )
300
+
301
+ invoke_fused_moe_triton_kernel(
302
+ qintermediate_cache2,
303
+ w2,
304
+ intermediate_cache3,
305
+ a2q_scale,
306
+ self.w2_scale,
307
+ topk_weights,
308
+ sorted_token_ids,
309
+ expert_ids,
310
+ num_tokens_post_padded,
311
+ not apply_router_weight_on_input,
312
+ 1,
313
+ config,
314
+ compute_type=compute_type,
315
+ use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
316
+ use_int8_w8a8=self.quant_config.use_int8_w8a8,
317
+ use_int8_w8a16=self.quant_config.use_int8_w8a16,
318
+ use_int4_w4a16=self.quant_config.use_int4_w4a16,
319
+ per_channel_quant=self.per_act_token_quant,
320
+ block_shape=self.block_shape,
321
+ B_bias=self.w2_bias,
322
+ )
323
+
324
+ # LoRA w2: applied to intermediate_cache3 before moe_sum, using the
325
+ # unquantized intermediate_cache2 as the lora_a input. Reuses the
326
+ # sorted_token_ids_lora computed above.
327
+ if lora_context is not None:
328
+ self.apply_w2_lora(
329
+ lora_context,
330
+ y=intermediate_cache3,
331
+ x=intermediate_cache2,
332
+ topk_weights=topk_weights,
333
+ sorted_token_ids_lora=sorted_token_ids_lora,
334
+ expert_ids_lora=expert_ids_lora,
335
+ num_tokens_post_padded_lora=num_tokens_post_padded_lora,
336
+ token_lora_mapping=token_lora_mapping,
337
+ num_tokens=num_tokens,
338
+ w1=w1,
339
+ w2=w2,
340
+ top_k_num=top_k_num,
341
+ )
342
+
343
+ # separate function is required for MoE + LoRA
344
+ self.moe_sum(intermediate_cache3, output)
345
+
346
+ def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
347
+ ops.moe_sum(input, output)
348
+
349
+
350
+class TritonWNA16Experts(TritonExperts):
351
+ @staticmethod
352
+ def _supports_current_device() -> bool:
353
+ raise NotImplementedError(
354
+ "TritonWNA16Experts is not yet used by an Oracle. "
355
+ "This method should not be called."
356
+ )
357
+
358
+ @staticmethod
359
+ def _supports_no_act_and_mul() -> bool:
360
+ raise NotImplementedError(
361
+ "TritonWNA16Experts is not yet used by an Oracle. "
362
+ "This method should not be called."
363
+ )
364
+
365
+ @staticmethod
366
+ def _supports_quant_scheme(
367
+ weight_key: QuantKey | None,
368
+ activation_key: QuantKey | None,
369
+ ) -> bool:
370
+ raise NotImplementedError(
371
+ "TritonWNA16Experts is not yet used by an Oracle. "
372
+ "This method should not be called."
373
+ )
374
+
375
+ @staticmethod
376
+ def _supports_activation(activation: MoEActivation) -> bool:
377
+ raise NotImplementedError(
378
+ "TritonWNA16Experts is not yet used by an Oracle. "
379
+ "This method should not be called."
380
+ )
381
+
382
+ @staticmethod
383
+ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
384
+ raise NotImplementedError(
385
+ "TritonWNA16Experts is not yet used by an Oracle. "
386
+ "This method should not be called."
387
+ )
388
+
389
+ def apply(
390
+ self,
391
+ output: torch.Tensor,
392
+ hidden_states: torch.Tensor,
393
+ w1: torch.Tensor,
394
+ w2: torch.Tensor,
395
+ topk_weights: torch.Tensor,
396
+ topk_ids: torch.Tensor,
397
+ activation: MoEActivation,
398
+ global_num_experts: int,
399
+ expert_map: torch.Tensor | None,
400
+ a1q_scale: torch.Tensor | None,
401
+ a2_scale: torch.Tensor | None,
402
+ workspace13: torch.Tensor,
403
+ workspace2: torch.Tensor,
404
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
405
+ apply_router_weight_on_input: bool,
406
+ ):
407
+ # Check constraints.
408
+ if self.quant_config.use_int4_w4a16:
409
+ assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
410
+ else:
411
+ assert hidden_states.size(-1) == w1.size(2), (
412
+ f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
413
+ )
414
+
415
+ assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
416
+ assert hidden_states.dim() == 2
417
+ assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
418
+ assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
419
+ assert hidden_states.dtype in [
420
+ torch.float32,
421
+ torch.float16,
422
+ torch.bfloat16,
423
+ torch.float8_e4m3fn,
424
+ torch.float8_e4m3fnuz,
425
+ ]
426
+
427
+ E, num_tokens, N, K, top_k_num = self.moe_problem_size(
428
+ hidden_states, w1, w2, topk_ids
429
+ )
430
+
431
+ if global_num_experts == -1:
432
+ global_num_experts = E
433
+
434
+ config = try_get_optimal_moe_config(
435
+ w1.size(),
436
+ w2.size(),
437
+ top_k_num,
438
+ self.quant_config.config_name(hidden_states.dtype),
439
+ num_tokens,
440
+ block_shape=self.block_shape,
441
+ )
442
+
443
+ if hidden_states.dtype == torch.bfloat16:
444
+ compute_type = tl.bfloat16
445
+ elif hidden_states.dtype == torch.float16:
446
+ compute_type = tl.float16
447
+ elif hidden_states.dtype == torch.float32:
448
+ compute_type = tl.float32
449
+ elif (
450
+ hidden_states.dtype == torch.float8_e4m3fn
451
+ or hidden_states.dtype == torch.float8_e4m3fnuz
452
+ ):
453
+ compute_type = tl.bfloat16
454
+ else:
455
+ raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
456
+
457
+ # Note that the output tensor might be in workspace1
458
+ intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
459
+ activation_out_dim = self.adjust_N_for_activation(N, activation)
460
+ intermediate_cache2 = _resize_cache(
461
+ workspace13, (num_tokens * top_k_num, activation_out_dim)
462
+ )
463
+ intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
464
+
465
+ sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
466
+ topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map
467
+ )
468
+
469
+ invoke_fused_moe_wna16_triton_kernel(
470
+ hidden_states,
471
+ w1,
472
+ intermediate_cache1,
473
+ self.w1_scale,
474
+ self.quant_config.w1_zp,
475
+ None, # topk_weights
476
+ sorted_token_ids,
477
+ expert_ids,
478
+ num_tokens_post_padded,
479
+ False, # mul_routed_weights
480
+ top_k_num,
481
+ config,
482
+ compute_type=compute_type,
483
+ use_int8_w8a16=self.quant_config.use_int8_w8a16,
484
+ use_int4_w4a16=self.quant_config.use_int4_w4a16,
485
+ block_shape=self.block_shape,
486
+ )
487
+
488
+ self.activation(
489
+ activation, intermediate_cache2, intermediate_cache1.view(-1, N)
490
+ )
491
+
492
+ a2q_scale: torch.Tensor | None = None
493
+
494
+ qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
495
+ intermediate_cache2,
496
+ a2_scale,
497
+ self.quant_dtype,
498
+ self.per_act_token_quant,
499
+ self.block_shape,
500
+ )
501
+
502
+ invoke_fused_moe_wna16_triton_kernel(
503
+ qintermediate_cache2,
504
+ w2,
505
+ intermediate_cache3,
506
+ self.w2_scale,
507
+ self.quant_config.w2_zp,
508
+ topk_weights,
509
+ sorted_token_ids,
510
+ expert_ids,
511
+ num_tokens_post_padded,
512
+ not apply_router_weight_on_input,
513
+ 1,
514
+ config,
515
+ compute_type=compute_type,
516
+ use_int8_w8a16=self.quant_config.use_int8_w8a16,
517
+ use_int4_w4a16=self.quant_config.use_int4_w4a16,
518
+ block_shape=self.block_shape,
519
+ )
520
+
521
+ # separate function is required for MoE + LoRA
522
+ self.moe_sum(intermediate_cache3, output)
vllm/model_executor/layers/fused_moe/fused_moe.py
-494
@@ -20,34 +20,16 @@ from vllm.model_executor.layers.fused_moe.activation import (
20
)
21
from vllm.model_executor.layers.fused_moe.config import (
22
FUSED_MOE_UNQUANTIZED_CONFIG,
23
- FusedMoEConfig,
24
- FusedMoEParallelConfig,
23
FusedMoEQuantConfig,
24
_get_config_dtype_str,
25
)
28
-from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
26
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
27
moe_align_block_size,
28
)
32
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
33
- TopKWeightAndReduceNoOP,
34
-)
29
from vllm.model_executor.layers.fused_moe.utils import (
36
- _resize_cache,
30
disable_inplace,
31
moe_kernel_quantize_input,
32
)
40
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
41
- QuantKey,
42
- kFp8Dynamic128Sym,
43
- kFp8DynamicTensorSym,
44
- kFp8DynamicTokenSym,
45
- kFp8Static128BlockSym,
46
- kFp8StaticChannelSym,
47
- kFp8StaticTensorSym,
48
- kInt8DynamicTokenSym,
49
- kInt8StaticChannelSym,
50
-)
33
from vllm.platforms import current_platform
34
from vllm.triton_utils import tl, triton
35
from vllm.utils.torch_utils import direct_register_custom_op
@@ -1885,479 +1867,3 @@ def fused_experts_impl(
1867
)
1868
1869
return out_hidden_states
1888
-
1889
-
1890
-class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
1891
- """Triton-based fused MoE expert implementation."""
1892
-
1893
- def __init__(
1894
- self,
1895
- moe_config: FusedMoEConfig,
1896
- quant_config: FusedMoEQuantConfig,
1897
- ):
1898
- # Whether quantized MOE runs natively, or through
1899
- # higher-precision + activation QDQ.
1900
- self.quantization_emulation = False
1901
- super().__init__(moe_config, quant_config)
1902
-
1903
- @staticmethod
1904
- def activation_format() -> mk.FusedMoEActivationFormat:
1905
- return mk.FusedMoEActivationFormat.Standard
1906
-
1907
- @staticmethod
1908
- def _supports_current_device() -> bool:
1909
- return current_platform.is_cuda_alike() or current_platform.is_xpu()
1910
-
1911
- @staticmethod
1912
- def _supports_no_act_and_mul() -> bool:
1913
- return True
1914
-
1915
- @staticmethod
1916
- def _supports_quant_scheme(
1917
- weight_key: QuantKey | None,
1918
- activation_key: QuantKey | None,
1919
- ) -> bool:
1920
- # INT8 requires at least 7.5 (Turing).
1921
- device_supports_int8 = (
1922
- current_platform.is_cuda()
1923
- and current_platform.has_device_capability((7, 5))
1924
- )
1925
-
1926
- supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
1927
- if device_supports_int8:
1928
- supported.append((kInt8StaticChannelSym, kInt8DynamicTokenSym))
1929
- if current_platform.supports_fp8():
1930
- supported += [
1931
- (kFp8Static128BlockSym, kFp8Dynamic128Sym),
1932
- (kFp8StaticChannelSym, kFp8DynamicTokenSym),
1933
- (kFp8StaticTensorSym, kFp8DynamicTokenSym),
1934
- (kFp8StaticTensorSym, kFp8StaticTensorSym),
1935
- (kFp8StaticTensorSym, kFp8DynamicTensorSym),
1936
- ]
1937
- return (weight_key, activation_key) in supported
1938
-
1939
- @staticmethod
1940
- def _supports_activation(activation: MoEActivation) -> bool:
1941
- return activation in [
1942
- MoEActivation.SILU,
1943
- MoEActivation.GELU,
1944
- MoEActivation.GELU_TANH,
1945
- MoEActivation.SWIGLUOAI,
1946
- MoEActivation.SWIGLUSTEP,
1947
- MoEActivation.SILU_NO_MUL,
1948
- MoEActivation.GELU_NO_MUL,
1949
- MoEActivation.GELU_TANH_NO_MUL,
1950
- MoEActivation.RELU2_NO_MUL,
1951
- ]
1952
-
1953
- @staticmethod
1954
- def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
1955
- return not (
1956
- moe_parallel_config.use_fi_nvl_two_sided_kernels
1957
- or moe_parallel_config.use_fi_nvl_one_sided_kernels
1958
- )
1959
-
1960
- @staticmethod
1961
- def _supports_batch_invariance():
1962
- return True
1963
-
1964
- def supports_expert_map(self) -> bool:
1965
- return True
1966
-
1967
- def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1968
- return TopKWeightAndReduceNoOP()
1969
-
1970
- def workspace_shapes(
1971
- self,
1972
- M: int,
1973
- N: int,
1974
- K: int,
1975
- topk: int,
1976
- global_num_experts: int,
1977
- local_num_experts: int,
1978
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
1979
- activation: MoEActivation,
1980
- ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1981
- activation_out_dim = self.adjust_N_for_activation(N, activation)
1982
- workspace1 = (M, topk, max(activation_out_dim, K))
1983
- workspace2 = (M, topk, max(N, K))
1984
- output = (M, K)
1985
- return (workspace1, workspace2, output)
1986
-
1987
- def apply(
1988
- self,
1989
- output: torch.Tensor,
1990
- hidden_states: torch.Tensor,
1991
- w1: torch.Tensor,
1992
- w2: torch.Tensor,
1993
- topk_weights: torch.Tensor,
1994
- topk_ids: torch.Tensor,
1995
- activation: MoEActivation,
1996
- global_num_experts: int,
1997
- expert_map: torch.Tensor | None,
1998
- a1q_scale: torch.Tensor | None,
1999
- a2_scale: torch.Tensor | None,
2000
- workspace13: torch.Tensor,
2001
- workspace2: torch.Tensor,
2002
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
2003
- apply_router_weight_on_input: bool,
2004
- ):
2005
- # Check constraints.
2006
- if self.quant_config.use_int4_w4a16:
2007
- assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
2008
- else:
2009
- assert hidden_states.size(-1) == w1.size(2), (
2010
- f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
2011
- )
2012
-
2013
- assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
2014
- assert hidden_states.dim() == 2
2015
- assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
2016
- assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
2017
- assert hidden_states.dtype in [
2018
- torch.float32,
2019
- torch.float16,
2020
- torch.bfloat16,
2021
- torch.float8_e4m3fn,
2022
- torch.float8_e4m3fnuz,
2023
- ]
2024
-
2025
- E, num_tokens, N, K, top_k_num = self.moe_problem_size(
2026
- hidden_states, w1, w2, topk_ids
2027
- )
2028
-
2029
- if global_num_experts == -1:
2030
- global_num_experts = E
2031
-
2032
- config = try_get_optimal_moe_config(
2033
- w1.size(),
2034
- w2.size(),
2035
- top_k_num,
2036
- self.quant_config.config_name(hidden_states.dtype),
2037
- num_tokens,
2038
- block_shape=self.block_shape,
2039
- )
2040
-
2041
- if hidden_states.dtype == torch.bfloat16:
2042
- compute_type = tl.bfloat16
2043
- elif hidden_states.dtype == torch.float16:
2044
- compute_type = tl.float16
2045
- elif hidden_states.dtype == torch.float32:
2046
- compute_type = tl.float32
2047
- elif (
2048
- hidden_states.dtype == torch.float8_e4m3fn
2049
- or hidden_states.dtype == torch.float8_e4m3fnuz
2050
- ):
2051
- compute_type = tl.bfloat16
2052
- else:
2053
- raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
2054
-
2055
- # Note that the output tensor might be in workspace1
2056
- intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
2057
- cache2_dim = self.adjust_N_for_activation(N, activation)
2058
- intermediate_cache2 = _resize_cache(
2059
- workspace13, (num_tokens * top_k_num, cache2_dim)
2060
- )
2061
- intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
2062
-
2063
- sorted_token_ids, expert_ids, num_tokens_post_padded = (
2064
- _prepare_expert_assignment(
2065
- topk_ids,
2066
- config,
2067
- num_tokens,
2068
- top_k_num,
2069
- global_num_experts,
2070
- expert_map,
2071
- use_int8_w8a16=self.quant_config.use_int8_w8a16,
2072
- use_int4_w4a16=self.quant_config.use_int4_w4a16,
2073
- block_shape=self.block_shape,
2074
- )
2075
- )
2076
-
2077
- invoke_fused_moe_triton_kernel(
2078
- hidden_states,
2079
- w1,
2080
- intermediate_cache1,
2081
- a1q_scale,
2082
- self.w1_scale,
2083
- None, # topk_weights
2084
- sorted_token_ids,
2085
- expert_ids,
2086
- num_tokens_post_padded,
2087
- False, # mul_routed_weights
2088
- top_k_num,
2089
- config,
2090
- compute_type=compute_type,
2091
- use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
2092
- use_int8_w8a8=self.quant_config.use_int8_w8a8,
2093
- use_int8_w8a16=self.quant_config.use_int8_w8a16,
2094
- use_int4_w4a16=self.quant_config.use_int4_w4a16,
2095
- per_channel_quant=self.per_act_token_quant,
2096
- block_shape=self.block_shape,
2097
- B_bias=self.w1_bias,
2098
- )
2099
-
2100
- # LoRA w13: applied to intermediate_cache1 before activation, using
2101
- # hidden_states as the lora_a input. moe_lora_align_block_size is
2102
- # called once here and results reused for the w2 LoRA below.
2103
- sorted_token_ids_lora = None
2104
- expert_ids_lora = None
2105
- num_tokens_post_padded_lora = None
2106
- token_lora_mapping = None
2107
- lora_context = self._lora_context
2108
- if lora_context is not None:
2109
- (
2110
- sorted_token_ids_lora,
2111
- expert_ids_lora,
2112
- num_tokens_post_padded_lora,
2113
- token_lora_mapping,
2114
- ) = self.apply_w13_lora(
2115
- lora_context,
2116
- y=intermediate_cache1,
2117
- x=hidden_states,
2118
- topk_ids=topk_ids,
2119
- topk_weights=topk_weights,
2120
- expert_map=expert_map,
2121
- w1=w1,
2122
- w2=w2,
2123
- num_tokens=num_tokens,
2124
- top_k_num=top_k_num,
2125
- )
2126
-
2127
- self.activation(
2128
- activation, intermediate_cache2, intermediate_cache1.view(-1, N)
2129
- )
2130
-
2131
- a2q_scale: torch.Tensor | None = None
2132
-
2133
- qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
2134
- intermediate_cache2,
2135
- a2_scale,
2136
- self.quant_dtype,
2137
- self.per_act_token_quant,
2138
- self.block_shape,
2139
- quantization_emulation=self.quantization_emulation,
2140
- )
2141
-
2142
- invoke_fused_moe_triton_kernel(
2143
- qintermediate_cache2,
2144
- w2,
2145
- intermediate_cache3,
2146
- a2q_scale,
2147
- self.w2_scale,
2148
- topk_weights,
2149
- sorted_token_ids,
2150
- expert_ids,
2151
- num_tokens_post_padded,
2152
- not apply_router_weight_on_input,
2153
- 1,
2154
- config,
2155
- compute_type=compute_type,
2156
- use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
2157
- use_int8_w8a8=self.quant_config.use_int8_w8a8,
2158
- use_int8_w8a16=self.quant_config.use_int8_w8a16,
2159
- use_int4_w4a16=self.quant_config.use_int4_w4a16,
2160
- per_channel_quant=self.per_act_token_quant,
2161
- block_shape=self.block_shape,
2162
- B_bias=self.w2_bias,
2163
- )
2164
-
2165
- # LoRA w2: applied to intermediate_cache3 before moe_sum, using the
2166
- # unquantized intermediate_cache2 as the lora_a input. Reuses the
2167
- # sorted_token_ids_lora computed above.
2168
- if lora_context is not None:
2169
- self.apply_w2_lora(
2170
- lora_context,
2171
- y=intermediate_cache3,
2172
- x=intermediate_cache2,
2173
- topk_weights=topk_weights,
2174
- sorted_token_ids_lora=sorted_token_ids_lora,
2175
- expert_ids_lora=expert_ids_lora,
2176
- num_tokens_post_padded_lora=num_tokens_post_padded_lora,
2177
- token_lora_mapping=token_lora_mapping,
2178
- num_tokens=num_tokens,
2179
- w1=w1,
2180
- w2=w2,
2181
- top_k_num=top_k_num,
2182
- )
2183
-
2184
- # separate function is required for MoE + LoRA
2185
- self.moe_sum(intermediate_cache3, output)
2186
-
2187
- def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
2188
- ops.moe_sum(input, output)
2189
-
2190
-
2191
-class TritonWNA16Experts(TritonExperts):
2192
- @staticmethod
2193
- def _supports_current_device() -> bool:
2194
- raise NotImplementedError(
2195
- "TritonWNA16Experts is not yet used by an Oracle. "
2196
- "This method should not be called."
2197
- )
2198
-
2199
- @staticmethod
2200
- def _supports_no_act_and_mul() -> bool:
2201
- raise NotImplementedError(
2202
- "TritonWNA16Experts is not yet used by an Oracle. "
2203
- "This method should not be called."
2204
- )
2205
-
2206
- @staticmethod
2207
- def _supports_quant_scheme(
2208
- weight_key: QuantKey | None,
2209
- activation_key: QuantKey | None,
2210
- ) -> bool:
2211
- raise NotImplementedError(
2212
- "TritonWNA16Experts is not yet used by an Oracle. "
2213
- "This method should not be called."
2214
- )
2215
-
2216
- @staticmethod
2217
- def _supports_activation(activation: MoEActivation) -> bool:
2218
- raise NotImplementedError(
2219
- "TritonWNA16Experts is not yet used by an Oracle. "
2220
- "This method should not be called."
2221
- )
2222
-
2223
- @staticmethod
2224
- def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
2225
- raise NotImplementedError(
2226
- "TritonWNA16Experts is not yet used by an Oracle. "
2227
- "This method should not be called."
2228
- )
2229
-
2230
- def apply(
2231
- self,
2232
- output: torch.Tensor,
2233
- hidden_states: torch.Tensor,
2234
- w1: torch.Tensor,
2235
- w2: torch.Tensor,
2236
- topk_weights: torch.Tensor,
2237
- topk_ids: torch.Tensor,
2238
- activation: MoEActivation,
2239
- global_num_experts: int,
2240
- expert_map: torch.Tensor | None,
2241
- a1q_scale: torch.Tensor | None,
2242
- a2_scale: torch.Tensor | None,
2243
- workspace13: torch.Tensor,
2244
- workspace2: torch.Tensor,
2245
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
2246
- apply_router_weight_on_input: bool,
2247
- ):
2248
- # Check constraints.
2249
- if self.quant_config.use_int4_w4a16:
2250
- assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
2251
- else:
2252
- assert hidden_states.size(-1) == w1.size(2), (
2253
- f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
2254
- )
2255
-
2256
- assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
2257
- assert hidden_states.dim() == 2
2258
- assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
2259
- assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
2260
- assert hidden_states.dtype in [
2261
- torch.float32,
2262
- torch.float16,
2263
- torch.bfloat16,
2264
- torch.float8_e4m3fn,
2265
- torch.float8_e4m3fnuz,
2266
- ]
2267
-
2268
- E, num_tokens, N, K, top_k_num = self.moe_problem_size(
2269
- hidden_states, w1, w2, topk_ids
2270
- )
2271
-
2272
- if global_num_experts == -1:
2273
- global_num_experts = E
2274
-
2275
- config = try_get_optimal_moe_config(
2276
- w1.size(),
2277
- w2.size(),
2278
- top_k_num,
2279
- self.quant_config.config_name(hidden_states.dtype),
2280
- num_tokens,
2281
- block_shape=self.block_shape,
2282
- )
2283
-
2284
- if hidden_states.dtype == torch.bfloat16:
2285
- compute_type = tl.bfloat16
2286
- elif hidden_states.dtype == torch.float16:
2287
- compute_type = tl.float16
2288
- elif hidden_states.dtype == torch.float32:
2289
- compute_type = tl.float32
2290
- elif (
2291
- hidden_states.dtype == torch.float8_e4m3fn
2292
- or hidden_states.dtype == torch.float8_e4m3fnuz
2293
- ):
2294
- compute_type = tl.bfloat16
2295
- else:
2296
- raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
2297
-
2298
- # Note that the output tensor might be in workspace1
2299
- intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
2300
- activation_out_dim = self.adjust_N_for_activation(N, activation)
2301
- intermediate_cache2 = _resize_cache(
2302
- workspace13, (num_tokens * top_k_num, activation_out_dim)
2303
- )
2304
- intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
2305
-
2306
- sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
2307
- topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map
2308
- )
2309
-
2310
- invoke_fused_moe_wna16_triton_kernel(
2311
- hidden_states,
2312
- w1,
2313
- intermediate_cache1,
2314
- self.w1_scale,
2315
- self.quant_config.w1_zp,
2316
- None, # topk_weights
2317
- sorted_token_ids,
2318
- expert_ids,
2319
- num_tokens_post_padded,
2320
- False, # mul_routed_weights
2321
- top_k_num,
2322
- config,
2323
- compute_type=compute_type,
2324
- use_int8_w8a16=self.quant_config.use_int8_w8a16,
2325
- use_int4_w4a16=self.quant_config.use_int4_w4a16,
2326
- block_shape=self.block_shape,
2327
- )
2328
-
2329
- self.activation(
2330
- activation, intermediate_cache2, intermediate_cache1.view(-1, N)
2331
- )
2332
-
2333
- a2q_scale: torch.Tensor | None = None
2334
-
2335
- qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
2336
- intermediate_cache2,
2337
- a2_scale,
2338
- self.quant_dtype,
2339
- self.per_act_token_quant,
2340
- self.block_shape,
2341
- )
2342
-
2343
- invoke_fused_moe_wna16_triton_kernel(
2344
- qintermediate_cache2,
2345
- w2,
2346
- intermediate_cache3,
2347
- self.w2_scale,
2348
- self.quant_config.w2_zp,
2349
- topk_weights,
2350
- sorted_token_ids,
2351
- expert_ids,
2352
- num_tokens_post_padded,
2353
- not apply_router_weight_on_input,
2354
- 1,
2355
- config,
2356
- compute_type=compute_type,
2357
- use_int8_w8a16=self.quant_config.use_int8_w8a16,
2358
- use_int4_w4a16=self.quant_config.use_int4_w4a16,
2359
- block_shape=self.block_shape,
2360
- )
2361
-
2362
- # separate function is required for MoE + LoRA
2363
- self.moe_sum(intermediate_cache3, output)
vllm/model_executor/layers/fused_moe/layer.py
+3
-3
@@ -26,15 +26,15 @@ from vllm.model_executor.layers.fused_moe.config import (
26
FusedMoEQuantConfig,
27
RoutingMethodType,
28
)
29
+from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
30
+ init_aiter_topK_meta_data,
31
+)
32
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
33
FusedMoEMethodBase,
34
)
35
from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
36
FusedMoEModularMethod,
37
)
35
-from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
36
- init_aiter_topK_meta_data,
37
-)
38
from vllm.model_executor.layers.fused_moe.router.router_factory import (
39
create_fused_moe_router,
40
)
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+4
-4
@@ -123,7 +123,7 @@ def backend_to_kernel_cls(
123
return [TrtLlmFp8ExpertsMonolithic, TrtLlmFp8ExpertsModular]
124
125
elif backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
126
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
126
+ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501
127
FlashInferExperts,
128
)
129
@@ -144,14 +144,14 @@ def backend_to_kernel_cls(
144
return [BatchedDeepGemmExperts]
145
146
elif backend == Fp8MoeBackend.MARLIN:
147
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
147
+ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
148
MarlinExperts,
149
)
150
151
return [MarlinExperts]
152
153
elif backend == Fp8MoeBackend.TRITON:
154
- from vllm.model_executor.layers.fused_moe.fused_moe import (
154
+ from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
155
TritonExperts,
156
)
157
@@ -165,7 +165,7 @@ def backend_to_kernel_cls(
165
return [BatchedTritonExperts]
166
167
elif backend == Fp8MoeBackend.AITER:
168
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
168
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
169
AiterExperts,
170
)
171
vllm/model_executor/layers/fused_moe/oracle/int8.py
+1
-1
@@ -46,7 +46,7 @@ def backend_to_kernel_cls(
46
backend: Int8MoeBackend,
47
) -> list[type[mk.FusedMoEExperts]]:
48
if backend == Int8MoeBackend.TRITON:
49
- from vllm.model_executor.layers.fused_moe.fused_moe import (
49
+ from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
50
TritonExperts,
51
)
52
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+3
-3
@@ -12,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.config import (
12
FusedMoEConfig,
13
FusedMoEQuantConfig,
14
)
15
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
15
+from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
16
BatchedMarlinExperts,
17
MarlinExperts,
18
)
@@ -42,14 +42,14 @@ def backend_to_kernel_cls(
42
) -> list[type[mk.FusedMoEExperts]]:
43
"""Return the experts class for the given backend, or None for NONE."""
44
if backend == WNA16MoEBackend.MARLIN:
45
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
45
+ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
46
MarlinExperts,
47
)
48
49
return [MarlinExperts]
50
51
elif backend == WNA16MoEBackend.BATCHED_MARLIN:
52
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
52
+ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
53
BatchedMarlinExperts,
54
)
55
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+4
-4
@@ -124,7 +124,7 @@ def backend_to_kernel_cls(
124
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
125
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
126
):
127
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
127
+ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501
128
FlashInferExperts,
129
)
130
@@ -160,21 +160,21 @@ def backend_to_kernel_cls(
160
]
161
162
elif backend == Mxfp4MoeBackend.MARLIN:
163
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
163
+ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
164
MarlinExperts,
165
)
166
167
return [MarlinExperts]
168
169
elif backend == Mxfp4MoeBackend.BATCHED_MARLIN:
170
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
170
+ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
171
BatchedMarlinExperts,
172
)
173
174
return [BatchedMarlinExperts]
175
176
elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
177
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
177
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
178
AiterExperts,
179
)
180
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+2
-2
@@ -89,7 +89,7 @@ def backend_to_kernel_cls(
89
]
90
91
elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:
92
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
92
+ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501
93
FlashInferExperts,
94
)
95
@@ -117,7 +117,7 @@ def backend_to_kernel_cls(
117
return [CutlassExpertsFp4]
118
119
elif backend == NvFp4MoeBackend.MARLIN:
120
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
120
+ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
121
MarlinExperts,
122
)
123
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+5
-3
@@ -95,21 +95,23 @@ def backend_to_kernel_cls(
95
return TrtLlmBf16Experts
96
97
elif backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
98
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
98
+ from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501
99
FlashInferExperts,
100
)
101
102
return FlashInferExperts
103
104
elif backend == UnquantizedMoeBackend.AITER:
105
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
105
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
106
AiterExperts,
107
)
108
109
return AiterExperts
110
111
elif backend == UnquantizedMoeBackend.TRITON:
112
- from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
112
+ from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
113
+ TritonExperts,
114
+ )
115
116
return TritonExperts
117
vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
+1
-1
@@ -14,7 +14,7 @@ from vllm.model_executor.layers.fused_moe.config import (
14
RoutingMethodType,
15
get_routing_method_type,
16
)
17
-from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
17
+from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
18
rocm_aiter_grouped_topk,
19
)
20
from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py
+1
-1
@@ -11,8 +11,8 @@ from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
12
)
13
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
14
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
15
from vllm.model_executor.layers.fused_moe.fallback import FallbackExperts
15
-from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
16
from vllm.platforms import current_platform
17
18
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+1
-1
@@ -14,8 +14,8 @@ from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
14
_valid_deep_gemm,
15
_valid_deep_gemm_shape,
16
)
17
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
18
from vllm.model_executor.layers.fused_moe.fallback import FallbackExperts
18
-from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
19
from vllm.utils.deep_gemm import (
20
is_deep_gemm_e8m0_used,
21
)
vllm/model_executor/layers/quantization/awq_marlin.py
+2
-2
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEConfig,
21
FusedMoEQuantConfig,
22
)
23
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
23
+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,
@@ -764,7 +764,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
764
)
765
766
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
767
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
767
+ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
768
BatchedMarlinExperts,
769
MarlinExperts,
770
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+1
-1
@@ -17,7 +17,7 @@ from vllm.model_executor.layers.fused_moe.config import (
17
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
18
CutlassExpertsMxfp4,
19
)
20
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
20
+from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
21
MarlinExperts,
22
)
23
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+1
-1
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEQuantConfig,
21
int4_w4a16_moe_quant_config,
22
)
23
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
23
+from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
24
BatchedMarlinExperts,
25
MarlinExperts,
26
fused_marlin_moe,
vllm/model_executor/layers/quantization/quark/quark_moe.py
+4
-4
@@ -26,7 +26,7 @@ from vllm.model_executor.layers.fused_moe.config import (
26
mxfp4_w4a16_moe_quant_config,
27
ocp_mx_moe_quant_config,
28
)
29
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
29
+from vllm.model_executor.layers.fused_moe.experts.marlin_moe import fused_marlin_moe
30
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
31
TRITON_BACKENDS,
32
Mxfp4MoeBackend,
@@ -444,7 +444,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
444
shared_experts_input: torch.Tensor | None,
445
) -> torch.Tensor:
446
if self.rocm_aiter_moe_enabled:
447
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
447
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
448
rocm_aiter_fused_experts,
449
)
450
@@ -909,7 +909,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
909
topk_ids: torch.Tensor,
910
shared_experts_input: torch.Tensor | None,
911
) -> torch.Tensor:
912
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
912
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
913
rocm_aiter_fused_experts,
914
)
915
@@ -1436,7 +1436,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1436
1437
# AITER path
1438
# TODO: Refactor this to use modular MOE kernel as well.
1439
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
1439
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1440
rocm_aiter_fused_experts,
1441
)
1442