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 2, 09:45 UTC
[MoE Refactor] WNA16 MoE backend selection into oracle module (#42553)
Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
bnellnm committed
May 29, 2026 at 13:11 UTC
84b2a8a7e7ec8c7dac98f9d3a300c21ca7d2d5e7
8 files changed
+545
-376
vllm/model_executor/layers/fused_moe/config.py
+16
-51
@@ -869,19 +869,23 @@ def nvfp4_w4a16_moe_quant_config(
869
def int4_w4a16_moe_quant_config(
870
w1_scale: torch.Tensor,
871
w2_scale: torch.Tensor,
872
- w1_zp: torch.Tensor | None,
873
- w2_zp: torch.Tensor | None,
872
+ w1_zp: torch.Tensor | None = None,
873
+ w2_zp: torch.Tensor | None = None,
874
+ w1_bias: torch.Tensor | None = None,
875
+ w2_bias: torch.Tensor | None = None,
876
block_shape: list[int] | None = None,
877
+ a1_gscale: torch.Tensor | None = None,
878
+ a2_gscale: torch.Tensor | None = None,
879
) -> FusedMoEQuantConfig:
880
"""
881
Construct a quant config for 16-bit float activations and int4 weights.
882
"""
883
group_shape = GroupShape(*block_shape) if block_shape is not None else None
884
return FusedMoEQuantConfig(
881
- _a1=FusedMoEQuantDesc(shape=group_shape),
882
- _a2=FusedMoEQuantDesc(shape=group_shape),
883
- _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp),
884
- _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp),
885
+ _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale),
886
+ _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale),
887
+ _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp, w1_bias),
888
+ _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp, w2_bias),
889
)
890
891
@@ -922,19 +926,21 @@ def fp8_w8a16_moe_quant_config(
926
def int8_w8a16_moe_quant_config(
927
w1_scale: torch.Tensor,
928
w2_scale: torch.Tensor,
925
- w1_zp: torch.Tensor | None,
926
- w2_zp: torch.Tensor | None,
929
+ w1_zp: torch.Tensor | None = None,
930
+ w2_zp: torch.Tensor | None = None,
931
w1_bias: torch.Tensor | None = None,
932
w2_bias: torch.Tensor | None = None,
933
block_shape: list[int] | None = None,
934
+ a1_gscale: torch.Tensor | None = None,
935
+ a2_gscale: torch.Tensor | None = None,
936
) -> FusedMoEQuantConfig:
937
"""
938
Construct a quant config for 16-bit float activations and int8 weights.
939
"""
940
group_shape = GroupShape(*block_shape) if block_shape is not None else None
941
return FusedMoEQuantConfig(
936
- _a1=FusedMoEQuantDesc(shape=group_shape),
937
- _a2=FusedMoEQuantDesc(shape=group_shape),
942
+ _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale),
943
+ _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale),
944
_w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias),
945
_w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias),
946
)
@@ -965,47 +971,6 @@ def int4_w4afp8_moe_quant_config(
971
)
972
973
968
-def awq_marlin_moe_quant_config(
969
- w1_scale: torch.Tensor,
970
- w2_scale: torch.Tensor,
971
- w1_zp: torch.Tensor | None,
972
- w2_zp: torch.Tensor | None,
973
- weight_bits: int,
974
- group_size: int,
975
- w1_bias: torch.Tensor | None = None,
976
- w2_bias: torch.Tensor | None = None,
977
- a1_gscale: torch.Tensor | None = None,
978
- a2_gscale: torch.Tensor | None = None,
979
-) -> FusedMoEQuantConfig:
980
- """
981
- Construct a quant config for awq marlin quantization.
982
-
983
- a1_gscale / a2_gscale are optional global scales applied to activation
984
- quantization scales when Marlin runs with 8-bit activations.
985
- """
986
- from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
987
-
988
- w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size)
989
-
990
- # Activations are NOT quantized for AWQ (fp16/bf16)
991
- a_shape = w_shape # Same as weight shape for alignment
992
-
993
- # Determine weight dtype
994
- if weight_bits == 4:
995
- weight_dtype = "int4"
996
- elif weight_bits == 8:
997
- weight_dtype = torch.int8
998
- else:
999
- raise ValueError(f"Unsupported weight_bits: {weight_bits}")
1000
-
1001
- return FusedMoEQuantConfig(
1002
- _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape, alpha_or_gscale=a1_gscale),
1003
- _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape, alpha_or_gscale=a2_gscale),
1004
- _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, None, w1_zp, w1_bias),
1005
- _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, None, w2_zp, w2_bias),
1006
- )
1007
-
1008
-
974
def biased_moe_quant_config(
975
w1_bias: torch.Tensor | None,
976
w2_bias: torch.Tensor | None,
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+6
-1
@@ -44,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
44
kFp8StaticChannelSym,
45
kFp8StaticTensorSym,
46
kInt4Static,
47
+ kInt4Static32,
48
kInt8Static,
49
kMxfp4Static,
50
kMxfp8Static,
@@ -566,8 +567,9 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
567
quant_config.use_mxfp4_w4a16
568
or quant_config.use_nvfp4_w4a16
569
or quant_config.use_int4_w4a16
570
+ or quant_config.use_int8_w8a16
571
or quant_config.use_fp8_w8a16
570
- ), "Supports only {mxfp,nvfp,int}4_w4a16 or fp8_w8a16"
572
+ ), "Supports only {mxfp,nvfp,int}4_w4a16, int8_w8a16 or fp8_w8a16"
573
self.w13_g_idx = w13_g_idx
574
self.w2_g_idx = w2_g_idx
575
self.w13_g_idx_sort_indices = w13_g_idx_sort_indices
@@ -608,6 +610,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
610
kNvfp4Static,
611
kInt4Static,
612
kInt8Static,
613
+ kInt4Static32,
614
]
615
return weight_key in SUPPORTED_W
616
@@ -640,6 +643,8 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
643
if self.w1_zp is not None or self.w2_zp is not None:
644
return scalar_types.uint4.id
645
return scalar_types.uint4b8.id
646
+ elif self.quant_config.use_int8_w8a16:
647
+ return scalar_types.uint8b128.id
648
elif self.quant_config.use_mxfp4_w4a16 or self.quant_config.use_nvfp4_w4a16:
649
return scalar_types.float4_e2m1f.id
650
elif (
vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py
new
+159
@@ -0,0 +1,159 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
8
+from vllm.model_executor.layers.fused_moe.config import (
9
+ FusedMoEConfig,
10
+ FusedMoEParallelConfig,
11
+ FusedMoEQuantConfig,
12
+ RoutingMethodType,
13
+)
14
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
+ QuantKey,
16
+ kInt4Static32,
17
+)
18
+from vllm.platforms import current_platform
19
+
20
+
21
+class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
22
+ """
23
+ FlashInfer TRT-LLM MxInt4 MoE kernel. Monolithic interface
24
+ (fused router + experts).
25
+
26
+ Wraps flashinfer_trtllm_mxint4_moe().
27
+ """
28
+
29
+ def __init__(
30
+ self,
31
+ moe_config: FusedMoEConfig,
32
+ quant_config: FusedMoEQuantConfig,
33
+ ):
34
+ super().__init__(moe_config, quant_config)
35
+ self.topk = moe_config.experts_per_token
36
+ self.intermediate_size_per_partition = (
37
+ moe_config.intermediate_size_per_partition
38
+ )
39
+ self.local_num_experts = moe_config.num_local_experts
40
+ self.ep_rank = moe_config.ep_rank
41
+ self.routing_method = moe_config.routing_method
42
+
43
+ @staticmethod
44
+ def _supports_current_device() -> bool:
45
+ from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( # noqa: E501
46
+ is_flashinfer_mxint4_moe_available,
47
+ )
48
+
49
+ p = current_platform
50
+ return (
51
+ p.is_cuda()
52
+ and p.is_device_capability_family(100)
53
+ and is_flashinfer_mxint4_moe_available()
54
+ )
55
+
56
+ @staticmethod
57
+ def _supports_no_act_and_mul() -> bool:
58
+ return False
59
+
60
+ @staticmethod
61
+ def _supports_quant_scheme(
62
+ weight_key: QuantKey | None,
63
+ activation_key: QuantKey | None,
64
+ ) -> bool:
65
+ return (weight_key, activation_key) == (kInt4Static32, None)
66
+
67
+ @staticmethod
68
+ def _supports_activation(activation: MoEActivation) -> bool:
69
+ # FlashInfer MxInt4 uses a fused SwiGLU activation.
70
+ return activation == MoEActivation.SWIGLUOAI
71
+
72
+ @staticmethod
73
+ def _supports_parallel_config(
74
+ moe_parallel_config: FusedMoEParallelConfig,
75
+ ) -> bool:
76
+ return (
77
+ not moe_parallel_config.use_all2all_kernels
78
+ or moe_parallel_config.use_ag_rs_all2all_kernels
79
+ ) and not moe_parallel_config.enable_eplb
80
+
81
+ @staticmethod
82
+ def _supports_routing_method(
83
+ routing_method: RoutingMethodType,
84
+ weight_key: QuantKey | None,
85
+ activation_key: QuantKey | None,
86
+ ) -> bool:
87
+ return routing_method in [
88
+ RoutingMethodType.Renormalize,
89
+ RoutingMethodType.RenormalizeNaive,
90
+ RoutingMethodType.DeepSeekV3,
91
+ RoutingMethodType.Llama4,
92
+ RoutingMethodType.Simulated,
93
+ ]
94
+
95
+ @staticmethod
96
+ def _supports_router_logits_dtype(
97
+ router_logits_dtype: torch.dtype | None,
98
+ routing_method: RoutingMethodType,
99
+ ) -> bool:
100
+ if router_logits_dtype == torch.float32:
101
+ # DeepSeekV3 routing handles float32 logits internally.
102
+ # Simulated routing generates synthetic decisions.
103
+ return routing_method in (
104
+ RoutingMethodType.DeepSeekV3,
105
+ RoutingMethodType.Simulated,
106
+ )
107
+ return True
108
+
109
+ @staticmethod
110
+ def activation_format() -> mk.FusedMoEActivationFormat:
111
+ return mk.FusedMoEActivationFormat.Standard
112
+
113
+ def supports_expert_map(self) -> bool:
114
+ return False
115
+
116
+ @property
117
+ def expects_unquantized_inputs(self) -> bool:
118
+ # The kernel handles quantization internally.
119
+ return True
120
+
121
+ def apply(
122
+ self,
123
+ hidden_states: torch.Tensor,
124
+ w1: torch.Tensor,
125
+ w2: torch.Tensor,
126
+ router_logits: torch.Tensor,
127
+ activation: MoEActivation,
128
+ global_num_experts: int,
129
+ expert_map: torch.Tensor | None,
130
+ a1q_scale: torch.Tensor | None,
131
+ apply_router_weight_on_input: bool,
132
+ num_expert_group: int | None = None,
133
+ e_score_correction_bias: torch.Tensor | None = None,
134
+ routed_scaling_factor: float | None = None,
135
+ topk_group: int | None = None,
136
+ ) -> torch.Tensor:
137
+ from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( # noqa: E501
138
+ flashinfer_trtllm_mxint4_moe,
139
+ )
140
+
141
+ assert self.w1_scale is not None
142
+ assert self.w2_scale is not None
143
+ return flashinfer_trtllm_mxint4_moe(
144
+ x=hidden_states,
145
+ router_logits=router_logits,
146
+ w13_weight_packed=w1,
147
+ w13_weight_scale=self.w1_scale,
148
+ w2_weight_packed=w2,
149
+ w2_weight_scale=self.w2_scale,
150
+ global_num_experts=global_num_experts,
151
+ top_k=self.topk,
152
+ intermediate_size_per_partition=self.intermediate_size_per_partition,
153
+ local_num_experts=self.local_num_experts,
154
+ ep_rank=self.ep_rank,
155
+ num_expert_group=num_expert_group,
156
+ topk_group=topk_group,
157
+ e_score_correction_bias=e_score_correction_bias,
158
+ routing_method_type=self.routing_method,
159
+ )
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+220
-75
@@ -2,9 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import sys
4
from enum import Enum
5
-from typing import TYPE_CHECKING
5
+from typing import Any
6
7
import torch
8
+from compressed_tensors.quantization import (
9
+ QuantizationArgs,
10
+)
11
12
import vllm._custom_ops as ops
13
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -12,10 +15,16 @@ from vllm.logger import init_logger
15
from vllm.model_executor.layers.fused_moe.config import (
16
FusedMoEConfig,
17
FusedMoEQuantConfig,
18
+ int4_w4a16_moe_quant_config,
19
+ int8_w8a16_moe_quant_config,
20
)
21
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
22
BatchedMarlinExperts,
23
MarlinExperts,
24
+ MarlinExpertsBase,
25
+)
26
+from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
27
+ TrtLlmMxint4ExpertsMonolithic,
28
)
29
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
30
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
@@ -29,16 +38,13 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
38
)
39
from vllm.platforms import current_platform
40
32
-if TYPE_CHECKING:
33
- from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
34
- from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
35
-
41
logger = init_logger(__name__)
42
43
44
class WNA16MoEBackend(Enum):
45
MARLIN = "MARLIN"
46
BATCHED_MARLIN = "BATCHED_MARLIN"
47
+ FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
48
XPU = "XPU"
49
50
@@ -47,26 +53,17 @@ def backend_to_kernel_cls(
53
) -> list[type[mk.FusedMoEExperts]]:
54
"""Return the experts class for the given backend, or None for NONE."""
55
if backend == WNA16MoEBackend.MARLIN:
50
- from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
51
- MarlinExperts,
52
- )
53
-
56
return [MarlinExperts]
55
-
57
elif backend == WNA16MoEBackend.BATCHED_MARLIN:
57
- from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
58
- BatchedMarlinExperts,
59
- )
60
-
58
return [BatchedMarlinExperts]
62
-
59
+ elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
60
+ return [TrtLlmMxint4ExpertsMonolithic]
61
elif backend == WNA16MoEBackend.XPU:
62
from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
63
XPUExpertsWNA16,
64
)
65
66
return [XPUExpertsWNA16]
69
-
67
else:
68
raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
69
@@ -77,23 +74,25 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
74
"""
75
if current_platform.is_xpu():
76
return [WNA16MoEBackend.XPU]
80
- return [
77
+
78
+ _AVAILABLE_BACKENDS = [
79
+ WNA16MoEBackend.FLASHINFER_TRTLLM,
80
WNA16MoEBackend.MARLIN,
81
WNA16MoEBackend.BATCHED_MARLIN,
82
]
83
+ return _AVAILABLE_BACKENDS
84
85
86
def select_wna16_moe_backend(
87
config: FusedMoEConfig,
88
weight_key: QuantKey,
89
- weight_bits: int,
89
) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
90
"""Select the WNA16 MoE backend.
91
92
Args:
93
config: the shared ``FusedMoEConfig`` for this layer.
95
- weight_bits: quantization bit-width (4 or 8). 8-bit weights are not
96
- supported by the modular Marlin kernel, so ``NONE`` is returned.
94
+ weight_key: The QuantKey describing the weight quantization.
95
+ Must have int4 or int8 type.
96
97
Returns:
98
A tuple of (``WNA16MoEBackend``, experts class or ``None``).
@@ -156,15 +155,55 @@ def select_wna16_moe_backend(
155
)
156
157
158
+def make_wna16_moe_quant_config(
159
+ w1_scale: torch.Tensor,
160
+ w2_scale: torch.Tensor,
161
+ group_size: int,
162
+ num_bits: int,
163
+ w1_zp: torch.Tensor | None = None,
164
+ w2_zp: torch.Tensor | None = None,
165
+ w1_bias: torch.Tensor | None = None,
166
+ w2_bias: torch.Tensor | None = None,
167
+ a1_gscale: torch.Tensor | None = None,
168
+ a2_gscale: torch.Tensor | None = None,
169
+) -> FusedMoEQuantConfig:
170
+ """Create the FusedMoEQuantConfig for 4 or 8-bit WNA16 MoE."""
171
+ if num_bits == 4:
172
+ return int4_w4a16_moe_quant_config(
173
+ w1_scale=w1_scale,
174
+ w2_scale=w2_scale,
175
+ w1_zp=w1_zp,
176
+ w2_zp=w2_zp,
177
+ w1_bias=w1_bias,
178
+ w2_bias=w2_bias,
179
+ block_shape=[0, group_size],
180
+ a1_gscale=a1_gscale,
181
+ a2_gscale=a2_gscale,
182
+ )
183
+ else:
184
+ assert num_bits == 8
185
+ return int8_w8a16_moe_quant_config(
186
+ w1_scale=w1_scale,
187
+ w2_scale=w2_scale,
188
+ w1_zp=w1_zp,
189
+ w2_zp=w2_zp,
190
+ w1_bias=w1_bias,
191
+ w2_bias=w2_bias,
192
+ block_shape=[0, group_size],
193
+ a1_gscale=a1_gscale,
194
+ a2_gscale=a2_gscale,
195
+ )
196
+
197
+
198
def make_wna16_moe_kernel(
199
moe_quant_config: FusedMoEQuantConfig,
200
moe_config: FusedMoEConfig,
162
- experts_cls: type[mk.FusedMoEExperts] | None,
163
- is_k_full: bool,
164
- w13_g_idx: torch.Tensor | None,
165
- w2_g_idx: torch.Tensor | None,
166
- w13_g_idx_sort_indices: torch.Tensor | None,
167
- w2_g_idx_sort_indices: torch.Tensor | None,
201
+ experts_cls: type[mk.FusedMoEExperts],
202
+ is_k_full: bool = False,
203
+ w13_g_idx: torch.Tensor | None = None,
204
+ w2_g_idx: torch.Tensor | None = None,
205
+ w13_g_idx_sort_indices: torch.Tensor | None = None,
206
+ w2_g_idx_sort_indices: torch.Tensor | None = None,
207
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
208
) -> mk.FusedMoEKernel:
209
from vllm.model_executor.layers.fused_moe.all2all_utils import (
@@ -174,16 +213,37 @@ def make_wna16_moe_kernel(
213
XPUExpertsWNA16,
214
)
215
177
- assert experts_cls in (MarlinExperts, BatchedMarlinExperts, XPUExpertsWNA16)
216
+ # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts
217
+ # and BatchedMarlinExperts
218
+ assert experts_cls in (
219
+ MarlinExperts,
220
+ BatchedMarlinExperts,
221
+ TrtLlmMxint4ExpertsMonolithic,
222
+ XPUExpertsWNA16,
223
+ )
224
+
225
+ is_monolithic = experts_cls.is_monolithic()
226
227
prepare_finalize = maybe_make_prepare_finalize(
228
moe=moe_config,
229
quant_config=moe_quant_config,
230
routing_tables=routing_tables,
231
allow_new_interface=True,
232
+ use_monolithic=is_monolithic,
233
)
234
assert prepare_finalize is not None
186
- assert isinstance(prepare_finalize, mk.FusedMoEPrepareAndFinalizeModular)
235
+
236
+ logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
237
+
238
+ extra_args: dict[str, Any] = {}
239
+ if issubclass(experts_cls, MarlinExpertsBase):
240
+ extra_args = {
241
+ "w13_g_idx": w13_g_idx,
242
+ "w2_g_idx": w2_g_idx,
243
+ "w13_g_idx_sort_indices": w13_g_idx_sort_indices,
244
+ "w2_g_idx_sort_indices": w2_g_idx_sort_indices,
245
+ "is_k_full": is_k_full,
246
+ }
247
248
if experts_cls is XPUExpertsWNA16:
249
assert (
@@ -199,30 +259,20 @@ def make_wna16_moe_kernel(
259
elif (
260
prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts
261
):
202
- assert experts_cls == BatchedMarlinExperts
262
max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
263
assert max_num_tokens is not None
205
- experts = BatchedMarlinExperts(
264
+ experts = experts_cls(
265
max_num_tokens=max_num_tokens,
266
num_dispatchers=prepare_finalize.num_dispatchers(),
267
moe_config=moe_config,
268
quant_config=moe_quant_config,
210
- w13_g_idx=w13_g_idx,
211
- w2_g_idx=w2_g_idx,
212
- w13_g_idx_sort_indices=w13_g_idx_sort_indices,
213
- w2_g_idx_sort_indices=w2_g_idx_sort_indices,
214
- is_k_full=is_k_full,
269
+ **extra_args,
270
)
271
else:
217
- assert experts_cls == MarlinExperts
218
- experts = MarlinExperts(
272
+ experts = experts_cls(
273
moe_config=moe_config,
274
quant_config=moe_quant_config,
221
- w13_g_idx=w13_g_idx,
222
- w2_g_idx=w2_g_idx,
223
- w13_g_idx_sort_indices=w13_g_idx_sort_indices,
224
- w2_g_idx_sort_indices=w2_g_idx_sort_indices,
225
- is_k_full=is_k_full,
275
+ **extra_args,
276
)
277
278
return mk.FusedMoEKernel(
@@ -236,10 +286,74 @@ def make_wna16_moe_kernel(
286
# ---------------------------------------------------------------------------
287
288
289
+def _process_weights_flashinfer(
290
+ w13_qweight: torch.Tensor,
291
+ w2_qweight: torch.Tensor,
292
+ w13_scales: torch.Tensor,
293
+ w2_scales: torch.Tensor,
294
+ w13_g_idx: torch.Tensor,
295
+ w2_g_idx: torch.Tensor,
296
+ w13_bias: torch.Tensor | None = None,
297
+ w2_bias: torch.Tensor | None = None,
298
+) -> tuple[
299
+ torch.Tensor, # w13_qweight
300
+ torch.Tensor, # w2_qweight
301
+ torch.Tensor, # w13_scales
302
+ torch.Tensor, # w2_scales
303
+ torch.Tensor, # w13_g_idx
304
+ torch.Tensor, # w2_g_idx
305
+ torch.Tensor | None, # w13_g_idx_sort_indices
306
+ torch.Tensor | None, # w2_g_idx_sort_indices
307
+ torch.Tensor | None, # w13_qzeros
308
+ torch.Tensor | None, # w2_qzeros
309
+ torch.Tensor | None, # w13_input_global_scale
310
+ torch.Tensor | None, # w2_input_global_scale
311
+ torch.Tensor | None, # w13_bias
312
+ torch.Tensor | None, # w2_bias
313
+]:
314
+ """Flashinfer (TRT-LLM MXINT4) weight post-processing.
315
+
316
+ Steps
317
+ -----
318
+ 1. Transform weights/scales via ``prepare_static_weights_for_trtllm_mxint4_moe``.
319
+ 2. Return transformed tensors, passing through g_idx/bias unchanged.
320
+ """
321
+ from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
322
+ prepare_static_weights_for_trtllm_mxint4_moe,
323
+ )
324
+
325
+ dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe(
326
+ w13_qweight,
327
+ w13_scales,
328
+ w2_qweight,
329
+ w2_scales,
330
+ )
331
+
332
+ return (
333
+ dict_weights_mxint4["gemm1_weights"],
334
+ dict_weights_mxint4["gemm2_weights"],
335
+ dict_weights_mxint4["gemm1_scales"],
336
+ dict_weights_mxint4["gemm2_scales"],
337
+ w13_g_idx,
338
+ w2_g_idx,
339
+ None,
340
+ None,
341
+ None,
342
+ None,
343
+ None,
344
+ None,
345
+ w13_bias,
346
+ w2_bias,
347
+ )
348
+
349
+
350
def _process_weights_marlin(
351
layer: torch.nn.Module,
241
- quant_config: "AutoGPTQConfig",
352
input_dtype: torch.dtype | None,
353
+ num_bits: int,
354
+ pack_factor: int,
355
+ group_size: int,
356
+ actorder: str | None,
357
w13_qweight: torch.Tensor,
358
w2_qweight: torch.Tensor,
359
w13_scales: torch.Tensor,
@@ -292,6 +406,7 @@ def _process_weights_marlin(
406
407
# --- FP8 weight / scale adjustment ---
408
if input_dtype == torch.float8_e4m3fn:
409
+ # NOTE: for non-zp quantization format only
410
marlin_w13_qweight = ops.marlin_int4_fp8_preprocess(w13_qweight, inplace=False)
411
marlin_w2_qweight = ops.marlin_int4_fp8_preprocess(w2_qweight, inplace=False)
412
marlin_w13_scales = w13_scales.data * 512
@@ -303,7 +418,7 @@ def _process_weights_marlin(
418
marlin_w2_scales = w2_scales
419
420
# --- Process act_order (g_idx) ---
306
- if quant_config.desc_act:
421
+ if actorder == "group":
422
num_experts = w13_g_idx.shape[0]
423
w13_g_idx_sort_indices = torch.empty_like(w13_g_idx)
424
w2_g_idx_sort_indices = torch.empty_like(w2_g_idx)
@@ -314,6 +429,8 @@ def _process_weights_marlin(
429
w2_g_idx_sort_indices[e] = torch.argsort(w2_g_idx[e]).to(torch.int32)
430
w13_sorted_g_idx[e] = w13_g_idx[e][w13_g_idx_sort_indices[e]]
431
w2_sorted_g_idx[e] = w2_g_idx[e][w2_g_idx_sort_indices[e]]
432
+ w13_g_idx = w13_sorted_g_idx
433
+ w2_g_idx = w2_sorted_g_idx
434
else:
435
num_experts = w13_g_idx.shape[0]
436
device = w13_g_idx.device
@@ -338,17 +455,17 @@ def _process_weights_marlin(
455
marlin_w13_qweight = ops.gptq_marlin_moe_repack(
456
marlin_w13_qweight,
457
w13_g_idx_sort_indices,
341
- marlin_w13_qweight.shape[1] * quant_config.pack_factor,
458
+ marlin_w13_qweight.shape[1] * pack_factor,
459
marlin_w13_qweight.shape[2],
343
- quant_config.quant_type.size_bits,
460
+ num_bits,
461
is_a_8bit=is_a_8bit,
462
)
463
marlin_w2_qweight = ops.gptq_marlin_moe_repack(
464
marlin_w2_qweight,
465
w2_g_idx_sort_indices,
349
- marlin_w2_qweight.shape[1] * quant_config.pack_factor,
466
+ marlin_w2_qweight.shape[1] * pack_factor,
467
marlin_w2_qweight.shape[2],
351
- quant_config.quant_type.size_bits,
468
+ num_bits,
469
is_a_8bit=is_a_8bit,
470
)
471
@@ -357,19 +474,15 @@ def _process_weights_marlin(
474
s=marlin_w13_scales,
475
size_k=layer.intermediate_size_per_partition,
476
size_n=marlin_w13_scales.shape[2],
360
- group_size=quant_config.group_size,
477
+ group_size=group_size,
478
is_a_8bit=is_a_8bit,
479
)
480
+ group_size_or_pack_factor = group_size if group_size != -1 else pack_factor
481
marlin_w2_scales = marlin_moe_permute_scales(
482
s=marlin_w2_scales,
365
- size_k=marlin_w2_scales.shape[1]
366
- * (
367
- quant_config.group_size
368
- if quant_config.group_size != -1
369
- else quant_config.pack_factor
370
- ),
483
+ size_k=marlin_w2_scales.shape[1] * group_size_or_pack_factor,
484
size_n=marlin_w2_scales.shape[2],
372
- group_size=quant_config.group_size,
485
+ group_size=group_size,
486
is_a_8bit=is_a_8bit,
487
)
488
@@ -409,7 +522,9 @@ def _process_weights_marlin(
522
523
def _process_awq_weights_marlin(
524
layer: torch.nn.Module,
412
- quant_config: "AWQMarlinConfig",
525
+ weight_bits: int,
526
+ pack_factor: int,
527
+ group_size: int,
528
input_dtype: torch.dtype | None,
529
w13_qweight: torch.Tensor,
530
w2_qweight: torch.Tensor,
@@ -475,16 +590,16 @@ def _process_awq_weights_marlin(
590
w13_qweight,
591
w13_g_idx_sort_indices,
592
size_k=w13_qweight.shape[1],
478
- size_n=w13_qweight.shape[2] * quant_config.pack_factor,
479
- num_bits=quant_config.weight_bits,
593
+ size_n=w13_qweight.shape[2] * pack_factor,
594
+ num_bits=weight_bits,
595
is_a_8bit=is_a_8bit,
596
)
597
marlin_w2_qweight = ops.awq_marlin_moe_repack(
598
w2_qweight,
599
w2_g_idx_sort_indices,
600
size_k=w2_qweight.shape[1],
486
- size_n=w2_qweight.shape[2] * quant_config.pack_factor,
487
- num_bits=quant_config.weight_bits,
601
+ size_n=w2_qweight.shape[2] * pack_factor,
602
+ num_bits=weight_bits,
603
is_a_8bit=is_a_8bit,
604
)
605
@@ -492,7 +607,7 @@ def _process_awq_weights_marlin(
607
s=w13_scales,
608
size_k=layer.intermediate_size_per_partition,
609
size_n=w13_scales.shape[2],
495
- group_size=quant_config.group_size,
610
+ group_size=group_size,
611
is_a_8bit=is_a_8bit,
612
)
613
if input_dtype == torch.int8 and layer.num_groups_w13 > 1:
@@ -504,7 +619,7 @@ def _process_awq_weights_marlin(
619
s=w2_scales,
620
size_k=layer.intermediate_size_per_partition,
621
size_n=w2_scales.shape[2],
507
- group_size=quant_config.group_size,
622
+ group_size=group_size,
623
is_a_8bit=is_a_8bit,
624
)
625
if input_dtype == torch.int8 and layer.num_groups_w2 > 1:
@@ -515,15 +630,15 @@ def _process_awq_weights_marlin(
630
marlin_w13_qzeros = moe_awq_to_marlin_zero_points(
631
w13_qzeros,
632
size_k=w13_qzeros.shape[1],
518
- size_n=w13_qzeros.shape[2] * quant_config.pack_factor,
519
- num_bits=quant_config.weight_bits,
633
+ size_n=w13_qzeros.shape[2] * pack_factor,
634
+ num_bits=weight_bits,
635
is_a_8bit=is_a_8bit,
636
)
637
marlin_w2_qzeros = moe_awq_to_marlin_zero_points(
638
w2_qzeros,
639
size_k=w2_qzeros.shape[1],
525
- size_n=w2_qzeros.shape[2] * quant_config.pack_factor,
526
- num_bits=quant_config.weight_bits,
640
+ size_n=w2_qzeros.shape[2] * pack_factor,
641
+ num_bits=weight_bits,
642
is_a_8bit=is_a_8bit,
643
)
644
@@ -616,7 +731,7 @@ def _process_weights_xpu(
731
def convert_to_wna16_moe_kernel_format(
732
backend: WNA16MoEBackend,
733
layer: torch.nn.Module,
619
- quant_config: QuantizationConfig,
734
+ quant_config: QuantizationConfig | QuantizationArgs | None,
735
input_dtype: torch.dtype | None,
736
w13: torch.Tensor,
737
w2: torch.Tensor,
@@ -669,9 +784,16 @@ def convert_to_wna16_moe_kernel_format(
784
if isinstance(quant_config, AWQMarlinConfig):
785
if w13_qzeros is None or w2_qzeros is None:
786
raise ValueError("AWQ Marlin MoE requires zero-point tensors.")
787
+
788
+ weight_bits = quant_config.weight_bits
789
+ pack_factor = quant_config.pack_factor
790
+ group_size = quant_config.group_size
791
+
792
return _process_awq_weights_marlin(
793
layer,
674
- quant_config,
794
+ weight_bits,
795
+ pack_factor,
796
+ group_size,
797
input_dtype,
798
w13,
799
w2,
@@ -682,19 +804,30 @@ def convert_to_wna16_moe_kernel_format(
804
w13_bias,
805
w2_bias,
806
)
685
-
686
- if not isinstance(quant_config, AutoGPTQConfig):
807
+ elif isinstance(quant_config, AutoGPTQConfig):
808
+ num_bits = quant_config.quant_type.size_bits
809
+ pack_factor = quant_config.pack_factor
810
+ group_size = quant_config.group_size
811
+ actorder = "group" if quant_config.desc_act else None
812
+ elif isinstance(quant_config, QuantizationArgs):
813
+ num_bits = quant_config.num_bits
814
+ pack_factor = 32 // quant_config.num_bits
815
+ group_size = quant_config.group_size
816
+ actorder = quant_config.actorder
817
+ else:
818
raise TypeError(
688
- "Marlin WNA16 MoE backend requires AutoGPTQConfig or "
689
- "AWQMarlinConfig, got "
690
- f"{type(quant_config).__name__}."
819
+ "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or "
820
+ f"QuantizationArgs, got {type(quant_config).__name__}."
821
)
822
if w13_g_idx is None or w2_g_idx is None:
823
raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")
824
return _process_weights_marlin(
825
layer,
696
- quant_config,
826
input_dtype,
827
+ num_bits,
828
+ pack_factor,
829
+ group_size,
830
+ actorder,
831
w13,
832
w2,
833
w13_scale,
@@ -706,7 +839,19 @@ def convert_to_wna16_moe_kernel_format(
839
w13_bias,
840
w2_bias,
841
)
842
+ elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
843
+ return _process_weights_flashinfer(
844
+ w13,
845
+ w2,
846
+ w13_scale,
847
+ w2_scale,
848
+ w13_g_idx,
849
+ w2_g_idx,
850
+ w13_bias,
851
+ w2_bias,
852
+ )
853
elif backend == WNA16MoEBackend.XPU:
854
+ assert quant_config is not None
855
(
856
w13_xpu,
857
w2_xpu,
vllm/model_executor/layers/quantization/auto_gptq.py
+2
-1
@@ -483,7 +483,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
483
weight_key = QuantKey(quant_type, scale)
484
485
self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
486
- moe, weight_key, quant_config.weight_bits
486
+ moe,
487
+ weight_key,
488
)
489
490
def create_weights(
vllm/model_executor/layers/quantization/awq_marlin.py
+5
-7
@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.config import (
29
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
30
convert_to_wna16_moe_kernel_format,
31
make_wna16_moe_kernel,
32
+ make_wna16_moe_quant_config,
33
select_wna16_moe_backend,
34
)
35
from vllm.model_executor.layers.linear import (
@@ -521,7 +522,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
522
self.input_dtype = None
523
self.use_marlin = True
524
self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
524
- moe, kInt4Static, quant_config.weight_bits
525
+ moe,
526
+ kInt4Static,
527
)
528
529
def create_weights(
@@ -706,15 +708,11 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
708
)
709
710
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
709
- from vllm.model_executor.layers.fused_moe.config import (
710
- awq_marlin_moe_quant_config,
711
- )
712
-
713
- return awq_marlin_moe_quant_config(
711
+ return make_wna16_moe_quant_config(
712
w1_scale=layer.w13_scales,
713
w2_scale=layer.w2_scales,
716
- weight_bits=self.quant_config.weight_bits,
714
group_size=self.quant_config.group_size,
715
+ num_bits=self.quant_config.weight_bits,
716
w1_zp=getattr(layer, "w13_qzeros", None)
717
if self.quant_config.zero_point
718
else None,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+134
-241
@@ -1,16 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import enum
5
-from enum import Enum
4
+from typing import Any
5
6
import torch
7
from compressed_tensors.quantization import (
8
QuantizationArgs,
9
)
10
12
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
-from vllm import _custom_ops as ops
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.fused_moe import (
13
RoutedExperts,
@@ -19,12 +16,13 @@ from vllm.model_executor.layers.fused_moe import (
16
from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEConfig,
18
FusedMoEQuantConfig,
22
- int4_w4a16_moe_quant_config,
19
)
24
-from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
25
- BatchedMarlinExperts,
26
- MarlinExperts,
27
- fused_marlin_moe,
20
+from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
21
+ WNA16MoEBackend,
22
+ convert_to_wna16_moe_kernel_format,
23
+ make_wna16_moe_kernel,
24
+ make_wna16_moe_quant_config,
25
+ select_wna16_moe_backend,
26
)
27
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
28
CompressedTensorsMoEMethod,
@@ -32,27 +30,21 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
30
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
31
WNA16_SUPPORTED_TYPES_MAP,
32
)
35
-from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
36
- flashinfer_trtllm_mxint4_moe,
37
- is_flashinfer_mxint4_moe_available,
38
- prepare_static_weights_for_trtllm_mxint4_moe,
39
-)
33
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
34
get_marlin_input_dtype,
42
- marlin_act_int8_process_scales,
35
marlin_make_workspace_new,
44
- marlin_moe_permute_scales,
36
+)
37
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
38
+ QuantKey,
39
+ kInt4Static32GroupScale,
40
+ kInt4StaticGroupScale,
41
+ kInt8StaticGroupScale,
42
)
43
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
44
45
logger = init_logger(__name__)
46
47
51
-class GPTQMarlinState(Enum):
52
- REPACK = enum.auto()
53
- READY = enum.auto()
54
-
55
-
48
class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
49
def __init__(
50
self,
@@ -75,19 +67,27 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
67
self.actorder = weight_quant.actorder
68
69
self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
78
-
70
self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
80
- self.use_flashinfer_mxint4_moe = (
81
- is_flashinfer_mxint4_moe_available()
82
- and self.group_size == 32
83
- and weight_quant.num_bits == 4
84
- )
85
- self.kernel_backend = (
86
- "Flashinfer" if self.use_flashinfer_mxint4_moe else "Marlin"
87
- )
88
- logger.info_once(
89
- f"Using {self.kernel_backend} backend for WNA16 MoE "
90
- f"(group_size={self.group_size}, num_bits={self.num_bits})",
71
+
72
+ if self.num_bits == 4:
73
+ if self.group_size == 32:
74
+ scale = kInt4Static32GroupScale
75
+ else:
76
+ scale = kInt4StaticGroupScale
77
+ elif self.num_bits == 8:
78
+ assert self.group_size == -1
79
+ scale = kInt8StaticGroupScale
80
+ else:
81
+ raise ValueError(
82
+ "CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now."
83
+ )
84
+
85
+ weight_key = QuantKey(self.quant_type, scale)
86
+
87
+ # Select WNA16 MoE backend via oracle.
88
+ self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
89
+ config=self.moe,
90
+ weight_key=weight_key,
91
)
92
93
def get_weight_shape(
@@ -114,6 +114,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
114
"num_groups_w2 must be provided for weight scales"
115
)
116
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
117
+ is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
118
shape_map = {
119
"w13_weight": {
120
"Flashinfer": (
@@ -156,7 +157,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
157
"Marlin": (num_experts, num_groups_w2, hidden_size),
158
},
159
}
159
- return shape_map[weight_name][self.kernel_backend]
160
+ backend_key = "Flashinfer" if is_flashinfer else "Marlin"
161
+ return shape_map[weight_name][backend_key]
162
163
def create_weights(
164
self,
@@ -172,7 +174,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
174
# Will transpose the loaded weight along the
175
# intermediate and hidden dim sizes. Will
176
# shard for TP along the transposed dims
175
- is_transposed = self.kernel_backend != "Flashinfer"
177
+ is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
178
extra_weight_attrs.update(
179
{"is_transposed": is_transposed, "quant_method": self.strategy}
180
)
@@ -319,200 +321,104 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
321
322
layer.a13_scale = None
323
layer.a2_scale = None
322
- layer.marlin_state = GPTQMarlinState.REPACK
324
325
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
325
- num_experts = layer.w13_weight_g_idx.shape[0]
326
- device = layer.w13_weight_g_idx.device
327
- if self.kernel_backend == "Flashinfer":
328
- dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe(
329
- layer.w13_weight_packed,
330
- layer.w13_weight_scale,
331
- layer.w2_weight_packed,
332
- layer.w2_weight_scale,
333
- )
334
- replace_parameter(
335
- layer, "w13_weight_packed", dict_weights_mxint4["gemm1_weights"]
336
- )
337
- replace_parameter(
338
- layer, "w13_weight_scale", dict_weights_mxint4["gemm1_scales"]
339
- )
340
- replace_parameter(
341
- layer, "w2_weight_packed", dict_weights_mxint4["gemm2_weights"]
342
- )
343
- replace_parameter(
344
- layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"]
345
- )
346
- return None
347
-
348
- is_a_8bit = (
349
- self.marlin_input_dtype is not None
350
- and self.marlin_input_dtype.itemsize == 1
326
+ # Process weights using the shared oracle infrastructure
327
+ is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
328
+ (
329
+ w13_qweight,
330
+ w2_qweight,
331
+ w13_scales,
332
+ w2_scales,
333
+ w13_g_idx_processed,
334
+ w2_g_idx_processed,
335
+ w13_g_idx_sort_indices,
336
+ w2_g_idx_sort_indices,
337
+ _, # w13_qzeros
338
+ _, # w2_qzeros
339
+ w13_input_global_scale,
340
+ w2_input_global_scale,
341
+ _, # w13_bias
342
+ _, # w2_bias
343
+ ) = convert_to_wna16_moe_kernel_format(
344
+ backend=self.wna16_backend,
345
+ layer=layer,
346
+ quant_config=self.weight_quant,
347
+ input_dtype=self.marlin_input_dtype,
348
+ w13=layer.w13_weight_packed,
349
+ w2=layer.w2_weight_packed,
350
+ w13_scale=layer.w13_weight_scale,
351
+ w2_scale=layer.w2_weight_scale,
352
+ w13_g_idx=layer.w13_weight_g_idx,
353
+ w2_g_idx=layer.w2_weight_g_idx,
354
)
355
353
- if self.marlin_input_dtype == torch.float8_e4m3fn:
354
- # NOTE: for non-zp quantization format only
355
- ops.marlin_int4_fp8_preprocess(layer.w13_weight_packed, inplace=True)
356
- ops.marlin_int4_fp8_preprocess(layer.w2_weight_packed, inplace=True)
357
- layer.w13_weight_scale.data = layer.w13_weight_scale.data * 512
358
- layer.w2_weight_scale.data = layer.w2_weight_scale.data * 512
359
-
360
- # when running models with grouped act order,
361
- # resort to g_idx values provided in checkpoint
362
- if self.actorder == "group":
363
- w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx)
364
- w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx)
365
- w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx)
366
- w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx)
367
-
368
- for e in range(num_experts):
369
- w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to(
370
- torch.int32
371
- )
372
- w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to(
373
- torch.int32
374
- )
375
- w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][
376
- w13_g_idx_sort_indices[e]
377
- ]
378
- w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]]
356
+ # Replace common parameters
357
+ replace_parameter(layer, "w13_weight_packed", w13_qweight)
358
+ replace_parameter(layer, "w2_weight_packed", w2_qweight)
359
+ replace_parameter(layer, "w13_weight_scale", w13_scales)
360
+ replace_parameter(layer, "w2_weight_scale", w2_scales)
361
380
- replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx)
381
- replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx)
362
+ # Marlin-specific parameters (not needed for Flashinfer)
363
+ if not is_flashinfer:
364
+ replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
365
+ replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
366
replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
367
replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
368
385
- else:
386
- layer.w13_weight_g_idx = torch.nn.Parameter(
387
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
388
- requires_grad=False,
389
- )
390
- layer.w2_weight_g_idx = torch.nn.Parameter(
391
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
392
- requires_grad=False,
393
- )
394
- layer.w13_g_idx_sort_indices = torch.nn.Parameter(
395
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
396
- requires_grad=False,
397
- )
398
- layer.w2_g_idx_sort_indices = torch.nn.Parameter(
399
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
400
- requires_grad=False,
369
+ # Register input global scales if present
370
+ if w13_input_global_scale is not None:
371
+ layer.register_parameter(
372
+ "w13_input_global_scale",
373
+ torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
374
+ )
375
+ if w2_input_global_scale is not None:
376
+ layer.register_parameter(
377
+ "w2_input_global_scale",
378
+ torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
379
+ )
380
+
381
+ layer.workspace = marlin_make_workspace_new(
382
+ layer.w13_weight_g_idx.device, 4
383
)
384
403
- marlin_w13_qweight = ops.gptq_marlin_moe_repack(
404
- layer.w13_weight_packed,
405
- layer.w13_g_idx_sort_indices,
406
- layer.w13_weight_packed.shape[1] * self.packed_factor,
407
- layer.w13_weight_packed.shape[2],
408
- self.num_bits,
409
- is_a_8bit=is_a_8bit,
410
- )
411
- replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight)
412
-
413
- marlin_w2_qweight = ops.gptq_marlin_moe_repack(
414
- layer.w2_weight_packed,
415
- layer.w2_g_idx_sort_indices,
416
- layer.w2_weight_packed.shape[1] * self.packed_factor,
417
- layer.w2_weight_packed.shape[2],
418
- self.num_bits,
419
- is_a_8bit=is_a_8bit,
420
- )
421
- replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight)
385
+ # Alias packed weights to w13_weight/w2_weight for the modular kernel interface
386
+ layer.w13_weight = layer.w13_weight_packed
387
+ layer.w2_weight = layer.w2_weight_packed
388
423
- # Repack scales
424
- marlin_w13_scales = marlin_moe_permute_scales(
425
- s=layer.w13_weight_scale,
426
- size_k=layer.w13_weight_packed.shape[2],
427
- size_n=layer.w13_weight_scale.shape[2],
428
- group_size=self.group_size,
429
- is_a_8bit=is_a_8bit,
430
- )
431
- if self.marlin_input_dtype == torch.int8 and layer.num_groups_w13 > 1:
432
- marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
433
- marlin_w13_scales
434
- )
435
- layer.register_parameter(
436
- "w13_input_global_scale",
437
- torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
438
- )
439
- replace_parameter(layer, "w13_weight_scale", marlin_w13_scales)
389
+ assert self.experts_cls is not None
390
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
391
+ assert self.moe_quant_config is not None
392
441
- marlin_w2_scales = marlin_moe_permute_scales(
442
- s=layer.w2_weight_scale,
443
- size_k=layer.w2_weight_scale.shape[1]
444
- * (self.group_size if self.group_size != -1 else self.packed_factor),
445
- size_n=layer.w2_weight_scale.shape[2],
446
- group_size=self.group_size,
447
- is_a_8bit=is_a_8bit,
393
+ # Add Marlin-specific arguments
394
+ marlin_args: dict[str, Any] = {}
395
+ if not is_flashinfer:
396
+ marlin_args = {
397
+ "w13_g_idx": layer.w13_weight_g_idx,
398
+ "w2_g_idx": layer.w2_weight_g_idx,
399
+ "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices,
400
+ "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices,
401
+ "is_k_full": self.is_k_full,
402
+ }
403
+
404
+ self.moe_kernel = make_wna16_moe_kernel(
405
+ moe_quant_config=self.moe_quant_config,
406
+ moe_config=self.moe,
407
+ experts_cls=self.experts_cls,
408
+ routing_tables=layer._expert_routing_tables(),
409
+ **marlin_args,
410
)
449
- if self.marlin_input_dtype == torch.int8 and layer.num_groups_w2 > 1:
450
- marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
451
- marlin_w2_scales
452
- )
453
- layer.register_parameter(
454
- "w2_input_global_scale",
455
- torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
456
- )
457
- replace_parameter(layer, "w2_weight_scale", marlin_w2_scales)
458
-
459
- layer.workspace = marlin_make_workspace_new(device, 4)
411
412
def get_fused_moe_quant_config(
413
self, layer: torch.nn.Module
414
) -> FusedMoEQuantConfig | None:
464
- if self.num_bits != 4:
465
- return None
466
- return int4_w4a16_moe_quant_config(
415
+ return make_wna16_moe_quant_config(
416
w1_scale=layer.w13_weight_scale,
417
w2_scale=layer.w2_weight_scale,
469
- w1_zp=None,
470
- w2_zp=None,
471
- block_shape=[0, self.group_size],
418
+ group_size=self.group_size,
419
+ num_bits=self.num_bits,
420
)
421
474
- def select_gemm_impl(
475
- self,
476
- prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
477
- layer: torch.nn.Module,
478
- ) -> mk.FusedMoEExpertsModular:
479
- assert self.num_bits == 4, "only supporting w4"
480
- layer.w13_weight = layer.w13_weight_packed
481
- layer.w2_weight = layer.w2_weight_packed
482
- assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]])
483
- assert self.moe_quant_config is not None
484
- if (
485
- prepare_finalize.activation_format
486
- == mk.FusedMoEActivationFormat.BatchedExperts
487
- ):
488
- max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
489
- assert max_num_tokens_per_rank is not None
490
- return BatchedMarlinExperts(
491
- max_num_tokens=max_num_tokens_per_rank,
492
- num_dispatchers=prepare_finalize.num_dispatchers(),
493
- moe_config=self.moe,
494
- quant_config=self.moe_quant_config,
495
- w13_g_idx=layer.w13_weight_g_idx,
496
- w2_g_idx=layer.w2_weight_g_idx,
497
- w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
498
- w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
499
- is_k_full=self.is_k_full,
500
- )
501
- else:
502
- return MarlinExperts(
503
- moe_config=self.moe,
504
- quant_config=self.moe_quant_config,
505
- w13_g_idx=layer.w13_weight_g_idx,
506
- w2_g_idx=layer.w2_weight_g_idx,
507
- w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
508
- w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
509
- is_k_full=self.is_k_full,
510
- )
511
-
512
- @property
513
- def is_monolithic(self) -> bool:
514
- return self.kernel_backend == "Flashinfer"
515
-
422
def apply_monolithic(
423
self,
424
layer: RoutedExperts,
@@ -520,23 +426,21 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
426
router_logits: torch.Tensor,
427
input_ids: torch.Tensor | None = None,
428
) -> torch.Tensor:
523
- assert self.kernel_backend == "Flashinfer"
524
- return flashinfer_trtllm_mxint4_moe(
525
- x=x,
526
- router_logits=router_logits,
527
- w13_weight_packed=layer.w13_weight_packed,
528
- w13_weight_scale=layer.w13_weight_scale,
529
- w2_weight_packed=layer.w2_weight_packed,
530
- w2_weight_scale=layer.w2_weight_scale,
429
+ assert self.is_monolithic
430
+ assert self.moe_kernel is not None
431
+ return self.moe_kernel.apply_monolithic(
432
+ x,
433
+ layer.w13_weight,
434
+ layer.w2_weight,
435
+ router_logits,
436
+ activation=layer.activation,
437
global_num_experts=layer.global_num_experts,
532
- top_k=layer.top_k,
533
- intermediate_size_per_partition=layer.intermediate_size_per_partition,
534
- local_num_experts=layer.local_num_experts,
535
- ep_rank=layer.ep_rank,
438
+ expert_map=layer.expert_map,
439
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
440
num_expert_group=layer.num_expert_group,
441
topk_group=layer.topk_group,
442
e_score_correction_bias=layer.e_score_correction_bias,
539
- routing_method_type=layer.routing_method_type,
443
+ routed_scaling_factor=layer.routed_scaling_factor,
444
)
445
446
def apply(
@@ -548,29 +452,18 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
452
shared_experts: SharedExperts | None,
453
shared_experts_input: torch.Tensor | None,
454
) -> torch.Tensor:
551
- assert self.kernel_backend == "Marlin"
552
- return fused_marlin_moe(
455
+ assert not self.is_monolithic
456
+ assert self.moe_kernel is not None
457
+ return self.moe_kernel.apply(
458
x,
554
- layer.w13_weight_packed,
555
- layer.w2_weight_packed,
556
- None,
557
- None,
558
- layer.w13_weight_scale,
559
- layer.w2_weight_scale,
459
+ layer.w13_weight,
460
+ layer.w2_weight,
461
topk_weights,
462
topk_ids,
562
- input_global_scale1=getattr(layer, "w13_input_global_scale", None),
563
- input_global_scale2=getattr(layer, "w2_input_global_scale", None),
564
- quant_type_id=self.quant_type.id,
565
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
566
- global_num_experts=layer.global_num_experts,
463
activation=layer.activation,
464
+ global_num_experts=layer.global_num_experts,
465
expert_map=layer.expert_map,
569
- g_idx1=layer.w13_weight_g_idx,
570
- g_idx2=layer.w2_weight_g_idx,
571
- sort_indices1=layer.w13_g_idx_sort_indices,
572
- sort_indices2=layer.w2_g_idx_sort_indices,
573
- workspace=layer.workspace,
574
- input_dtype=self.marlin_input_dtype,
575
- is_k_full=self.is_k_full,
466
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
467
+ shared_experts=shared_experts,
468
+ shared_experts_input=shared_experts_input,
469
)
vllm/model_executor/layers/quantization/utils/quant_utils.py
+3
@@ -178,6 +178,9 @@ kInt4Static = QuantKey(INT4_DTYPE, scale=kInt4StaticGroupScale, symmetric=True)
178
kInt8StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1))
179
kInt8Static = QuantKey(INT8_DTYPE, scale=kInt8StaticGroupScale, symmetric=True)
180
181
+kInt4Static32GroupScale = ScaleDesc(torch.float16, True, GroupShape(1, 32))
182
+kInt4Static32 = QuantKey(INT4_DTYPE, scale=kInt4Static32GroupScale, symmetric=True)
183
+
184
kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)
185
kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)
186