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 4, 18:45 UTC
[MoE refactor] refactor GPTQMarlinMoEMethod with MK (#37990)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
Kunshang Ji committed
Apr 23, 2026 at 13:21 UTC
e4ee48da2d24c502a7e16606f871e12ef1e1fa3d
4 files changed
+576
-227
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+4
@@ -40,6 +40,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
40
kFp8Static128BlockSym,
41
kFp8StaticChannelSym,
42
kFp8StaticTensorSym,
43
+ kInt4Static,
44
+ kInt8Static,
45
kMxfp4Static,
46
kMxfp8Static,
47
kNvfp4Static,
@@ -585,6 +587,8 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
587
kMxfp4Static,
588
kMxfp8Static,
589
kNvfp4Static,
590
+ kInt4Static,
591
+ kInt8Static,
592
]
593
return weight_key in SUPPORTED_W
594
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
new
+445
@@ -0,0 +1,445 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from enum import Enum
4
+from typing import TYPE_CHECKING
5
+
6
+import torch
7
+
8
+import vllm._custom_ops as ops
9
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
+from vllm.logger import init_logger
11
+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 (
16
+ BatchedMarlinExperts,
17
+ MarlinExperts,
18
+)
19
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
20
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
21
+ marlin_act_int8_process_scales,
22
+ marlin_moe_permute_scales,
23
+ marlin_permute_bias,
24
+)
25
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
26
+ QuantKey,
27
+)
28
+
29
+if TYPE_CHECKING:
30
+ from vllm.model_executor.layers.quantization.gptq_marlin import GPTQMarlinConfig
31
+
32
+logger = init_logger(__name__)
33
+
34
+
35
+class WNA16MoEBackend(Enum):
36
+ MARLIN = "MARLIN"
37
+ BATCHED_MARLIN = "BATCHED_MARLIN"
38
+
39
+
40
+def backend_to_kernel_cls(
41
+ backend: WNA16MoEBackend,
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 (
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 (
53
+ BatchedMarlinExperts,
54
+ )
55
+
56
+ return [BatchedMarlinExperts]
57
+
58
+ else:
59
+ raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
60
+
61
+
62
+def _get_priority_backends() -> list[WNA16MoEBackend]:
63
+ """
64
+ Get available backends in priority order based on platform and config.
65
+ """
66
+ _AVAILABLE_BACKENDS = [
67
+ WNA16MoEBackend.MARLIN,
68
+ WNA16MoEBackend.BATCHED_MARLIN,
69
+ ]
70
+ return _AVAILABLE_BACKENDS
71
+
72
+
73
+def select_wna16_moe_backend(
74
+ config: FusedMoEConfig,
75
+ weight_key: QuantKey,
76
+ weight_bits: int,
77
+) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
78
+ """Select the WNA16 MoE backend.
79
+
80
+ Args:
81
+ config: the shared ``FusedMoEConfig`` for this layer.
82
+ weight_bits: quantization bit-width (4 or 8). 8-bit weights are not
83
+ supported by the modular Marlin kernel, so ``NONE`` is returned.
84
+
85
+ Returns:
86
+ A tuple of (``WNA16MoEBackend``, experts class or ``None``).
87
+ """
88
+
89
+ activation_format = (
90
+ mk.FusedMoEActivationFormat.BatchedExperts
91
+ if config.moe_parallel_config.use_batched_activation_format
92
+ else mk.FusedMoEActivationFormat.Standard
93
+ )
94
+
95
+ def _make_log_backend(backend: WNA16MoEBackend):
96
+ return f"Using '{backend.value}' WNA16 MoE backend."
97
+
98
+ def _make_log_unsupported(backend: WNA16MoEBackend, reason: str | None) -> str:
99
+ if reason:
100
+ return (
101
+ f"WNA16 MoE backend '{backend.value}' does not support the "
102
+ f"deployment configuration since {reason}."
103
+ )
104
+ return (
105
+ f"WNA16 MoE backend '{backend.value}' does not support the "
106
+ "deployment configuration."
107
+ )
108
+
109
+ def _return_or_raise(
110
+ backend: WNA16MoEBackend,
111
+ config: FusedMoEConfig,
112
+ weight_key: QuantKey | None,
113
+ activation_key: QuantKey | None,
114
+ activation_format: mk.FusedMoEActivationFormat,
115
+ ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
116
+ reason: str | None = None
117
+ for k_cls in backend_to_kernel_cls(backend):
118
+ supported, reason = k_cls.is_supported_config(
119
+ k_cls, config, weight_key, activation_key, activation_format
120
+ )
121
+ if supported:
122
+ logger.info_once(_make_log_backend(backend), scope="local")
123
+ return backend, k_cls
124
+ raise ValueError(_make_log_unsupported(backend, reason))
125
+
126
+ # Select kernels in order of backend.
127
+ AVAILABLE_BACKENDS = _get_priority_backends()
128
+
129
+ for backend in AVAILABLE_BACKENDS:
130
+ activation_key = None # always BF16 activation for WNA16 MoE
131
+ for k_cls in backend_to_kernel_cls(backend):
132
+ supported, reason = k_cls.is_supported_config(
133
+ k_cls, config, weight_key, activation_key, activation_format
134
+ )
135
+ if supported:
136
+ logger.info_once(_make_log_backend(backend), scope="local")
137
+ return backend, k_cls
138
+ else:
139
+ logger.debug_once(_make_log_unsupported(backend, reason), scope="local")
140
+
141
+ raise NotImplementedError(
142
+ "No WNA16 MoE backend supports the deployment configuration."
143
+ )
144
+
145
+
146
+def make_wna16_moe_kernel(
147
+ moe_quant_config: FusedMoEQuantConfig,
148
+ moe_config: FusedMoEConfig,
149
+ experts_cls: type[mk.FusedMoEExperts] | None,
150
+ layer: torch.nn.Module,
151
+ is_k_full: bool,
152
+ w13_g_idx: torch.Tensor | None,
153
+ w2_g_idx: torch.Tensor | None,
154
+ w13_g_idx_sort_indices: torch.Tensor | None,
155
+ w2_g_idx_sort_indices: torch.Tensor | None,
156
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
157
+ shared_experts: torch.nn.Module | None = None,
158
+) -> mk.FusedMoEKernel:
159
+ # Currently, we only support MarlinExperts and BatchedMarlinExperts
160
+ assert experts_cls in (MarlinExperts, BatchedMarlinExperts)
161
+
162
+ from vllm.model_executor.layers.fused_moe.all2all_utils import (
163
+ maybe_make_prepare_finalize,
164
+ )
165
+
166
+ prepare_finalize = maybe_make_prepare_finalize(
167
+ moe=moe_config,
168
+ quant_config=moe_quant_config,
169
+ routing_tables=routing_tables,
170
+ allow_new_interface=True,
171
+ )
172
+ assert prepare_finalize is not None
173
+ assert isinstance(prepare_finalize, mk.FusedMoEPrepareAndFinalizeModular)
174
+
175
+ if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
176
+ assert experts_cls == BatchedMarlinExperts
177
+ max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
178
+ assert max_num_tokens is not None
179
+ experts: mk.FusedMoEExperts = BatchedMarlinExperts(
180
+ max_num_tokens=max_num_tokens,
181
+ num_dispatchers=prepare_finalize.num_dispatchers(),
182
+ moe_config=moe_config,
183
+ quant_config=moe_quant_config,
184
+ w13_g_idx=w13_g_idx,
185
+ w2_g_idx=w2_g_idx,
186
+ w13_g_idx_sort_indices=w13_g_idx_sort_indices,
187
+ w2_g_idx_sort_indices=w2_g_idx_sort_indices,
188
+ is_k_full=is_k_full,
189
+ )
190
+ else:
191
+ assert experts_cls == MarlinExperts
192
+ experts = MarlinExperts(
193
+ moe_config=moe_config,
194
+ quant_config=moe_quant_config,
195
+ w13_g_idx=w13_g_idx,
196
+ w2_g_idx=w2_g_idx,
197
+ w13_g_idx_sort_indices=w13_g_idx_sort_indices,
198
+ w2_g_idx_sort_indices=w2_g_idx_sort_indices,
199
+ is_k_full=is_k_full,
200
+ )
201
+
202
+ return mk.FusedMoEKernel(
203
+ prepare_finalize,
204
+ experts,
205
+ shared_experts=shared_experts,
206
+ inplace=not moe_config.disable_inplace,
207
+ )
208
+
209
+
210
+# ---------------------------------------------------------------------------
211
+# Per-backend weight post-processing
212
+# ---------------------------------------------------------------------------
213
+
214
+
215
+def _process_weights_marlin(
216
+ layer: torch.nn.Module,
217
+ quant_config: "GPTQMarlinConfig",
218
+ input_dtype: torch.dtype | None,
219
+ w13_qweight: torch.Tensor,
220
+ w2_qweight: torch.Tensor,
221
+ w13_scales: torch.Tensor,
222
+ w2_scales: torch.Tensor,
223
+ w13_g_idx: torch.Tensor,
224
+ w2_g_idx: torch.Tensor,
225
+ w13_bias: torch.Tensor | None = None,
226
+ w2_bias: torch.Tensor | None = None,
227
+) -> tuple[
228
+ torch.Tensor, # w13_qweight
229
+ torch.Tensor, # w2_qweight
230
+ torch.Tensor, # w13_scales
231
+ torch.Tensor, # w2_scales
232
+ torch.Tensor, # w13_g_idx
233
+ torch.Tensor, # w2_g_idx
234
+ torch.Tensor, # w13_g_idx_sort_indices
235
+ torch.Tensor, # w2_g_idx_sort_indices
236
+ torch.Tensor | None, # w13_input_global_scale
237
+ torch.Tensor | None, # w2_input_global_scale
238
+ torch.Tensor | None, # w13_bias
239
+ torch.Tensor | None, # w2_bias
240
+]:
241
+ """Standard Marlin weight post-processing shared by MARLIN and
242
+ BATCHED_MARLIN backends.
243
+
244
+ Steps
245
+ -----
246
+ 1. Optional FP8 preprocessing of packed weights / scales.
247
+ 2. Sort / reset g_idx tensors for act-order handling.
248
+ 3. Repack weights via ``gptq_marlin_moe_repack``.
249
+ 4. Permute scales (and optionally extract INT8 global scales).
250
+ 5. Permute bias tensors.
251
+ """
252
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
253
+
254
+ marlin_w13_qweight: torch.Tensor
255
+ marlin_w2_qweight: torch.Tensor
256
+ marlin_w13_scales: torch.Tensor
257
+ marlin_w2_scales: torch.Tensor
258
+ w13_g_idx_sort_indices: torch.Tensor | None = None
259
+ w2_g_idx_sort_indices: torch.Tensor | None = None
260
+ w13_input_global_scale: torch.Tensor | None = None
261
+ w2_input_global_scale: torch.Tensor | None = None
262
+ w13_bias_out: torch.Tensor | None = None
263
+ w2_bias_out: torch.Tensor | None = None
264
+
265
+ # --- FP8 weight / scale adjustment ---
266
+ if input_dtype == torch.float8_e4m3fn:
267
+ marlin_w13_qweight = ops.marlin_int4_fp8_preprocess(w13_qweight, inplace=False)
268
+ marlin_w2_qweight = ops.marlin_int4_fp8_preprocess(w2_qweight, inplace=False)
269
+ marlin_w13_scales = w13_scales.data * 512
270
+ marlin_w2_scales = w2_scales.data * 512
271
+ else:
272
+ marlin_w13_qweight = w13_qweight
273
+ marlin_w2_qweight = w2_qweight
274
+ marlin_w13_scales = w13_scales
275
+ marlin_w2_scales = w2_scales
276
+
277
+ # --- Process act_order (g_idx) ---
278
+ if quant_config.desc_act:
279
+ num_experts = w13_g_idx.shape[0]
280
+ w13_g_idx_sort_indices = torch.empty_like(w13_g_idx)
281
+ w2_g_idx_sort_indices = torch.empty_like(w2_g_idx)
282
+ w13_sorted_g_idx = torch.empty_like(w13_g_idx)
283
+ w2_sorted_g_idx = torch.empty_like(w2_g_idx)
284
+ for e in range(num_experts):
285
+ w13_g_idx_sort_indices[e] = torch.argsort(w13_g_idx[e]).to(torch.int32)
286
+ w2_g_idx_sort_indices[e] = torch.argsort(w2_g_idx[e]).to(torch.int32)
287
+ w13_sorted_g_idx[e] = w13_g_idx[e][w13_g_idx_sort_indices[e]]
288
+ w2_sorted_g_idx[e] = w2_g_idx[e][w2_g_idx_sort_indices[e]]
289
+ else:
290
+ num_experts = w13_g_idx.shape[0]
291
+ device = w13_g_idx.device
292
+ w13_g_idx = torch.nn.Parameter(
293
+ torch.empty((num_experts, 0), dtype=torch.int32, device=device),
294
+ requires_grad=False,
295
+ )
296
+ w2_g_idx = torch.nn.Parameter(
297
+ torch.empty((num_experts, 0), dtype=torch.int32, device=device),
298
+ requires_grad=False,
299
+ )
300
+ w13_g_idx_sort_indices = torch.nn.Parameter(
301
+ torch.empty((num_experts, 0), dtype=torch.int32, device=device),
302
+ requires_grad=False,
303
+ )
304
+ w2_g_idx_sort_indices = torch.nn.Parameter(
305
+ torch.empty((num_experts, 0), dtype=torch.int32, device=device),
306
+ requires_grad=False,
307
+ )
308
+
309
+ # --- Repack weights ---
310
+ marlin_w13_qweight = ops.gptq_marlin_moe_repack(
311
+ marlin_w13_qweight,
312
+ w13_g_idx_sort_indices,
313
+ marlin_w13_qweight.shape[1] * quant_config.pack_factor,
314
+ marlin_w13_qweight.shape[2],
315
+ quant_config.quant_type.size_bits,
316
+ is_a_8bit=is_a_8bit,
317
+ )
318
+ marlin_w2_qweight = ops.gptq_marlin_moe_repack(
319
+ marlin_w2_qweight,
320
+ w2_g_idx_sort_indices,
321
+ marlin_w2_qweight.shape[1] * quant_config.pack_factor,
322
+ marlin_w2_qweight.shape[2],
323
+ quant_config.quant_type.size_bits,
324
+ is_a_8bit=is_a_8bit,
325
+ )
326
+
327
+ # --- Permute scales ---
328
+ marlin_w13_scales = marlin_moe_permute_scales(
329
+ s=marlin_w13_scales,
330
+ size_k=layer.intermediate_size_per_partition,
331
+ size_n=marlin_w13_scales.shape[2],
332
+ group_size=quant_config.group_size,
333
+ is_a_8bit=is_a_8bit,
334
+ )
335
+ marlin_w2_scales = marlin_moe_permute_scales(
336
+ s=marlin_w2_scales,
337
+ size_k=marlin_w2_scales.shape[1]
338
+ * (
339
+ quant_config.group_size
340
+ if quant_config.group_size != -1
341
+ else quant_config.pack_factor
342
+ ),
343
+ size_n=marlin_w2_scales.shape[2],
344
+ group_size=quant_config.group_size,
345
+ is_a_8bit=is_a_8bit,
346
+ )
347
+
348
+ if input_dtype == torch.int8:
349
+ if layer.num_groups_w13 > 1:
350
+ marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
351
+ marlin_w13_scales
352
+ )
353
+ if layer.num_groups_w2 > 1:
354
+ marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
355
+ marlin_w2_scales
356
+ )
357
+
358
+ # --- Permute bias ---
359
+ if w13_bias is not None:
360
+ w13_bias_out = marlin_permute_bias(w13_bias)
361
+ if w2_bias is not None:
362
+ w2_bias_out = marlin_permute_bias(w2_bias)
363
+
364
+ return (
365
+ marlin_w13_qweight,
366
+ marlin_w2_qweight,
367
+ marlin_w13_scales,
368
+ marlin_w2_scales,
369
+ w13_g_idx,
370
+ w2_g_idx,
371
+ w13_g_idx_sort_indices,
372
+ w2_g_idx_sort_indices,
373
+ w13_input_global_scale,
374
+ w2_input_global_scale,
375
+ w13_bias_out,
376
+ w2_bias_out,
377
+ )
378
+
379
+
380
+def convert_to_wna16_moe_kernel_format(
381
+ backend: WNA16MoEBackend,
382
+ layer: torch.nn.Module,
383
+ quant_config: QuantizationConfig,
384
+ input_dtype: torch.dtype | None,
385
+ w13: torch.Tensor,
386
+ w2: torch.Tensor,
387
+ w13_scale: torch.Tensor,
388
+ w2_scale: torch.Tensor,
389
+ w13_g_idx: torch.Tensor,
390
+ w2_g_idx: torch.Tensor,
391
+ w13_bias: torch.Tensor | None = None,
392
+ w2_bias: torch.Tensor | None = None,
393
+) -> tuple[
394
+ torch.Tensor, # w13_qweight
395
+ torch.Tensor, # w2_qweight
396
+ torch.Tensor, # w13_scales
397
+ torch.Tensor, # w2_scales
398
+ torch.Tensor | None, # w13_g_idx
399
+ torch.Tensor | None, # w2_g_idx
400
+ torch.Tensor | None, # w13_g_idx_sort_indices
401
+ torch.Tensor | None, # w2_g_idx_sort_indices
402
+ torch.Tensor | None, # w13_input_global_scale
403
+ torch.Tensor | None, # w2_input_global_scale
404
+ torch.Tensor | None, # w13_bias
405
+ torch.Tensor | None, # w2_bias
406
+]:
407
+ """Dispatch weight post-processing to the appropriate per-backend handler.
408
+
409
+ To add a new backend, implement a ``_process_weights_<name>`` helper and
410
+ add a branch here.
411
+
412
+ Args:
413
+ backend: the selected ``WNA16MoEBackend``.
414
+ layer: the ``FusedMoE`` layer whose parameters are being prepared.
415
+ quant_config: the ``QuantizationConfig`` for this layer.
416
+ input_dtype: optional activation dtype, usually should be 16 bit.
417
+ """
418
+ if backend in (
419
+ WNA16MoEBackend.MARLIN,
420
+ WNA16MoEBackend.BATCHED_MARLIN,
421
+ ):
422
+ from vllm.model_executor.layers.quantization.gptq_marlin import (
423
+ GPTQMarlinConfig,
424
+ )
425
+
426
+ if not isinstance(quant_config, GPTQMarlinConfig):
427
+ raise TypeError(
428
+ "Marlin WNA16 MoE backend requires GPTQMarlinConfig, got "
429
+ f"{type(quant_config).__name__}."
430
+ )
431
+ return _process_weights_marlin(
432
+ layer,
433
+ quant_config,
434
+ input_dtype,
435
+ w13,
436
+ w2,
437
+ w13_scale,
438
+ w2_scale,
439
+ w13_g_idx,
440
+ w2_g_idx,
441
+ w13_bias,
442
+ w2_bias,
443
+ )
444
+ else:
445
+ raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}")
vllm/model_executor/layers/quantization/gptq_marlin.py
+119
-227
@@ -9,7 +9,6 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
9
from transformers import PretrainedConfig
10
11
import vllm.model_executor.layers.fused_moe # noqa
12
-from vllm import _custom_ops as ops
12
from vllm.logger import init_logger
13
from vllm.model_executor.kernels.linear import (
14
MPLinearLayerConfig,
@@ -19,13 +18,17 @@ from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEConfig,
19
FusedMoEQuantConfig,
20
)
22
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
21
from vllm.model_executor.layers.fused_moe.layer import (
22
FusedMoE,
23
FusedMoEMethodBase,
24
FusedMoeWeightScaleSupported,
25
UnquantizedFusedMoEMethod,
26
)
27
+from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
28
+ convert_to_wna16_moe_kernel_format,
29
+ make_wna16_moe_kernel,
30
+ select_wna16_moe_backend,
31
+)
32
from vllm.model_executor.layers.linear import LinearMethodBase, set_weight_attrs
33
from vllm.model_executor.layers.quantization import QuantizationMethods
34
from vllm.model_executor.layers.quantization.base_config import (
@@ -42,13 +45,15 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
45
check_marlin_supported,
46
check_moe_marlin_supports_layer,
47
get_marlin_input_dtype,
45
- marlin_act_int8_process_scales,
48
marlin_make_workspace_new,
47
- marlin_moe_permute_scales,
48
- marlin_permute_bias,
49
marlin_repeat_scales_on_all_ranks,
50
verify_marlin_supported,
51
)
52
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
53
+ QuantKey,
54
+ kInt4StaticGroupScale,
55
+ kInt8StaticGroupScale,
56
+)
57
from vllm.model_executor.parameter import (
58
ChannelQuantScaleParameter,
59
GroupQuantScaleParameter,
@@ -500,13 +505,20 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
505
super().__init__(moe)
506
self.quant_config = quant_config
507
if self.quant_config.quant_type.size_bits == 4:
503
- self.quant_type = scalar_types.uint4b8
508
+ quant_type = scalar_types.uint4b8
509
+ scale = kInt4StaticGroupScale
510
elif self.quant_config.quant_type.size_bits == 8:
505
- self.quant_type = scalar_types.uint8b128
511
+ quant_type = scalar_types.uint8b128
512
+ scale = kInt8StaticGroupScale
513
else:
514
raise ValueError("GPTQMarlinMoEMethod only supports int4 and int8 now.")
515
self.input_dtype = None
516
self.use_marlin = True
517
+ weight_key = QuantKey(quant_type, scale)
518
+
519
+ self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
520
+ moe, weight_key, quant_config.weight_bits
521
+ )
522
523
def create_weights(
524
self,
@@ -521,7 +533,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
533
is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
534
535
if is_a_8bit:
524
- assert self.quant_type == scalar_types.uint4b8, (
536
+ assert self.quant_config.quant_type.size_bits == 8, (
537
"W8A8-INT8 is not supported by marlin kernel."
538
)
539
@@ -668,134 +680,100 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
680
is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
681
682
if is_a_8bit:
671
- assert self.quant_type == scalar_types.uint4b8, (
683
+ assert self.quant_config.quant_type.size_bits == 8, (
684
"W8A8-INT8 is not supported by marlin kernel."
685
)
686
675
- if self.input_dtype == torch.float8_e4m3fn:
676
- ops.marlin_int4_fp8_preprocess(layer.w13_qweight, inplace=True)
677
- ops.marlin_int4_fp8_preprocess(layer.w2_qweight, inplace=True)
678
- layer.w13_scales.data = layer.w13_scales.data * 512
679
- layer.w2_scales.data = layer.w2_scales.data * 512
680
-
681
- # Process act_order
682
- if self.quant_config.desc_act:
683
- # Get sorting based on g_idx
684
- num_experts = layer.w13_g_idx.shape[0]
685
- w13_g_idx_sort_indices = torch.empty_like(layer.w13_g_idx)
686
- w2_g_idx_sort_indices = torch.empty_like(layer.w2_g_idx)
687
- w13_sorted_g_idx = torch.empty_like(layer.w13_g_idx)
688
- w2_sorted_g_idx = torch.empty_like(layer.w2_g_idx)
689
- for e in range(num_experts):
690
- w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_g_idx[e]).to(
691
- torch.int32
692
- )
693
- w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_g_idx[e]).to(
694
- torch.int32
695
- )
696
- w13_sorted_g_idx[e] = layer.w13_g_idx[e][w13_g_idx_sort_indices[e]]
697
- w2_sorted_g_idx[e] = layer.w2_g_idx[e][w2_g_idx_sort_indices[e]]
698
- replace_parameter(layer, "w13_g_idx", w13_sorted_g_idx)
699
- replace_parameter(layer, "w2_g_idx", w2_sorted_g_idx)
700
- replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
701
- replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
702
- else:
703
- # Reset g_idx related tensors
704
- num_experts = layer.w13_g_idx.shape[0]
705
- device = layer.w13_g_idx.device
706
- layer.w13_g_idx = torch.nn.Parameter(
707
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
708
- requires_grad=False,
709
- )
710
- layer.w2_g_idx = torch.nn.Parameter(
711
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
712
- requires_grad=False,
713
- )
714
- layer.w13_g_idx_sort_indices = torch.nn.Parameter(
715
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
716
- requires_grad=False,
717
- )
718
- layer.w2_g_idx_sort_indices = torch.nn.Parameter(
719
- torch.empty((num_experts, 0), dtype=torch.int32, device=device),
720
- requires_grad=False,
721
- )
722
- # Repack weights
723
- marlin_w13_qweight = ops.gptq_marlin_moe_repack(
724
- layer.w13_qweight,
725
- layer.w13_g_idx_sort_indices,
726
- layer.w13_qweight.shape[1] * self.quant_config.pack_factor,
727
- layer.w13_qweight.shape[2],
728
- self.quant_config.quant_type.size_bits,
729
- is_a_8bit=is_a_8bit,
730
- )
731
- replace_parameter(layer, "w13_qweight", marlin_w13_qweight)
732
- marlin_w2_qweight = ops.gptq_marlin_moe_repack(
733
- layer.w2_qweight,
734
- layer.w2_g_idx_sort_indices,
735
- layer.w2_qweight.shape[1] * self.quant_config.pack_factor,
736
- layer.w2_qweight.shape[2],
737
- self.quant_config.quant_type.size_bits,
738
- is_a_8bit=is_a_8bit,
739
- )
740
- replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
741
-
742
- # The modular kernel expects w13_weight and w2_weight,
743
- # but GPTQ uses w13_qweight and w2_qweight
744
- # Alias for modular kernel
745
- layer.w13_weight = layer.w13_qweight
746
- # Alias for modular kernel
747
- layer.w2_weight = layer.w2_qweight
748
-
749
- # Repack scales
750
- marlin_w13_scales = marlin_moe_permute_scales(
751
- s=layer.w13_scales,
752
- size_k=layer.intermediate_size_per_partition,
753
- size_n=layer.w13_scales.shape[2],
754
- group_size=self.quant_config.group_size,
755
- is_a_8bit=is_a_8bit,
687
+ (
688
+ w13,
689
+ w2,
690
+ w13_scale,
691
+ w2_scale,
692
+ w13_g_idx,
693
+ w2_g_idx,
694
+ w13_g_idx_sort_indices,
695
+ w2_g_idx_sort_indices,
696
+ w13_input_global_scale,
697
+ w2_input_global_scale,
698
+ w13_bias,
699
+ w2_bias,
700
+ ) = convert_to_wna16_moe_kernel_format(
701
+ backend=self.wna16_moe_backend,
702
+ layer=layer,
703
+ quant_config=self.quant_config,
704
+ input_dtype=self.input_dtype,
705
+ w13=layer.w13_qweight,
706
+ w2=layer.w2_qweight,
707
+ w13_scale=layer.w13_scales,
708
+ w2_scale=layer.w2_scales,
709
+ w13_g_idx=layer.w13_g_idx,
710
+ w2_g_idx=layer.w2_g_idx,
711
+ w13_bias=getattr(layer, "w13_bias", None),
712
+ w2_bias=getattr(layer, "w2_bias", None),
713
)
757
- if self.input_dtype == torch.int8 and layer.num_groups_w13 > 1:
758
- marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
759
- marlin_w13_scales
760
- )
761
- layer.register_parameter(
762
- "w13_input_global_scale",
763
- torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
764
- )
714
766
- replace_parameter(layer, "w13_scales", marlin_w13_scales)
767
- marlin_w2_scales = marlin_moe_permute_scales(
768
- s=layer.w2_scales,
769
- size_k=layer.w2_scales.shape[1]
770
- * (
771
- self.quant_config.group_size
772
- if self.quant_config.group_size != -1
773
- else self.quant_config.pack_factor
774
- ),
775
- size_n=layer.w2_scales.shape[2],
776
- group_size=self.quant_config.group_size,
777
- is_a_8bit=is_a_8bit,
778
- )
779
- if self.input_dtype == torch.int8 and layer.num_groups_w2 > 1:
780
- marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
781
- marlin_w2_scales
782
- )
783
- layer.register_parameter(
784
- "w2_input_global_scale",
785
- torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
786
- )
715
+ replace_parameter(layer, "w13_qweight", w13)
716
+ replace_parameter(layer, "w2_qweight", w2)
717
+ replace_parameter(layer, "w13_scales", w13_scale)
718
+ replace_parameter(layer, "w2_scales", w2_scale)
719
+ replace_parameter(layer, "w13_g_idx", w13_g_idx)
720
+ replace_parameter(layer, "w2_g_idx", w2_g_idx)
721
+ replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
722
+ replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
723
+ if w13_input_global_scale is not None:
724
+ if hasattr(layer, "w13_input_global_scale"):
725
+ replace_parameter(
726
+ layer, "w13_input_global_scale", w13_input_global_scale
727
+ )
728
+ else:
729
+ layer.register_parameter(
730
+ "w13_input_global_scale",
731
+ torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
732
+ )
733
+ if w2_input_global_scale is not None:
734
+ if hasattr(layer, "w2_input_global_scale"):
735
+ replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale)
736
+ else:
737
+ layer.register_parameter(
738
+ "w2_input_global_scale",
739
+ torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
740
+ )
741
+ if w13_bias is not None:
742
+ if hasattr(layer, "w13_bias"):
743
+ replace_parameter(layer, "w13_bias", w13_bias)
744
+ else:
745
+ layer.register_parameter(
746
+ "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False)
747
+ )
748
+ if w2_bias is not None:
749
+ if hasattr(layer, "w2_bias"):
750
+ replace_parameter(layer, "w2_bias", w2_bias)
751
+ else:
752
+ layer.register_parameter(
753
+ "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False)
754
+ )
755
788
- replace_parameter(layer, "w2_scales", marlin_w2_scales)
756
+ self._setup_kernel(layer)
757
790
- if hasattr(layer, "w13_bias") and layer.w13_bias is not None:
791
- layer.w13_bias.data = marlin_permute_bias(layer.w13_bias)
758
+ def _setup_kernel(self, layer: FusedMoE) -> None:
759
+ """Build the FusedMoEKernel for this layer."""
760
793
- if hasattr(layer, "w2_bias") and layer.w2_bias is not None:
794
- layer.w2_bias.data = marlin_permute_bias(layer.w2_bias)
761
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
762
+ self.moe_kernel = make_wna16_moe_kernel(
763
+ moe_quant_config=self.moe_quant_config,
764
+ moe_config=self.moe,
765
+ experts_cls=self.experts_cls,
766
+ layer=layer,
767
+ is_k_full=self.is_k_full,
768
+ w13_g_idx=layer.w13_g_idx,
769
+ w2_g_idx=layer.w2_g_idx,
770
+ w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
771
+ w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
772
+ routing_tables=layer._maybe_init_expert_routing_tables(),
773
+ shared_experts=layer.shared_experts,
774
+ )
775
796
- def get_fused_moe_quant_config(
797
- self, layer: torch.nn.Module
798
- ) -> FusedMoEQuantConfig | None:
776
+ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
777
from vllm.model_executor.layers.fused_moe.config import (
778
gptq_marlin_moe_quant_config,
779
)
@@ -820,86 +798,11 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
798
prepare_finalize,
799
layer: torch.nn.Module,
800
):
823
- """
824
- Select the GEMM implementation for GPTQ-Marlin MoE.
825
-
826
- Returns MarlinExperts configured for GPTQ quantization.
827
- This is ONLY used when LoRA is enabled.
828
- Without LoRA, GPTQ uses its own apply() method.
829
- """
830
- # Only use modular kernels when LoRA is enabled
831
- # Without LoRA, GPTQ's own apply() method works fine and is more efficient
832
- if not self.moe.is_lora_enabled:
833
- raise NotImplementedError(
834
- "GPTQ-Marlin uses its own apply() method when LoRA is not enabled. "
835
- "Modular kernels are only used for LoRA support."
836
- )
837
-
838
- # The modular marlin kernels do not support 8-bit weights.
839
- if self.quant_config.weight_bits == 8:
840
- raise NotImplementedError(
841
- "GPTQ-Marlin kernel does not support 8-bit weights."
842
- )
843
-
844
- from vllm.model_executor.layers.fused_moe import modular_kernel as mk
845
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
846
- BatchedMarlinExperts,
847
- MarlinExperts,
848
- )
849
-
850
- # Ensure quant config is initialized
851
- assert self.moe_quant_config is not None, (
852
- "moe_quant_config must be initialized before select_gemm_impl"
853
- )
854
-
855
- w13_g_idx = (
856
- getattr(layer, "w13_g_idx", None) if self.quant_config.desc_act else None
857
- )
858
- w2_g_idx = (
859
- getattr(layer, "w2_g_idx", None) if self.quant_config.desc_act else None
860
- )
861
- w13_g_idx_sort_indices = (
862
- getattr(layer, "w13_g_idx_sort_indices", None)
863
- if self.quant_config.desc_act
864
- else None
865
- )
866
- w2_g_idx_sort_indices = (
867
- getattr(layer, "w2_g_idx_sort_indices", None)
868
- if self.quant_config.desc_act
869
- else None
801
+ raise ValueError(
802
+ f"{self.__class__.__name__} uses the new modular kernel "
803
+ "initialization logic. This function should not be called."
804
)
805
872
- # Check if using batched expert format (for Expert Parallelism)
873
- if (
874
- prepare_finalize.activation_format
875
- == mk.FusedMoEActivationFormat.BatchedExperts
876
- ):
877
- # For batched format, use BatchedMarlinExperts
878
- max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
879
- assert max_num_tokens_per_rank is not None
880
- return BatchedMarlinExperts(
881
- max_num_tokens=max_num_tokens_per_rank,
882
- num_dispatchers=prepare_finalize.num_dispatchers(),
883
- moe_config=self.moe,
884
- quant_config=self.moe_quant_config,
885
- w13_g_idx=w13_g_idx,
886
- w2_g_idx=w2_g_idx,
887
- w13_g_idx_sort_indices=w13_g_idx_sort_indices,
888
- w2_g_idx_sort_indices=w2_g_idx_sort_indices,
889
- is_k_full=self.is_k_full,
890
- )
891
- else:
892
- # Standard Marlin experts for GPTQ
893
- return MarlinExperts(
894
- moe_config=self.moe,
895
- quant_config=self.moe_quant_config,
896
- w13_g_idx=w13_g_idx,
897
- w2_g_idx=w2_g_idx,
898
- w13_g_idx_sort_indices=w13_g_idx_sort_indices,
899
- w2_g_idx_sort_indices=w2_g_idx_sort_indices,
900
- is_k_full=self.is_k_full,
901
- )
902
-
806
def apply(
807
self,
808
layer: FusedMoE,
@@ -908,28 +811,17 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
811
topk_ids: torch.Tensor,
812
shared_experts_input: torch.Tensor | None,
813
) -> torch.Tensor:
911
- return fused_marlin_moe(
912
- x,
913
- layer.w13_qweight,
914
- layer.w2_qweight,
915
- getattr(layer, "w13_bias", None),
916
- getattr(layer, "w2_bias", None),
917
- layer.w13_scales,
918
- layer.w2_scales,
919
- topk_weights,
920
- topk_ids,
921
- input_global_scale1=getattr(layer, "w13_input_global_scale", None),
922
- input_global_scale2=getattr(layer, "w2_input_global_scale", None),
923
- quant_type_id=self.quant_type.id,
924
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
814
+ assert not self.is_monolithic
815
+ assert self.moe_kernel is not None
816
+ return self.moe_kernel.apply(
817
+ hidden_states=x,
818
+ w1=layer.w13_qweight,
819
+ w2=layer.w2_qweight,
820
+ topk_weights=topk_weights,
821
+ topk_ids=topk_ids,
822
+ activation=layer.activation,
823
global_num_experts=layer.global_num_experts,
824
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
825
expert_map=layer.expert_map,
927
- g_idx1=layer.w13_g_idx,
928
- g_idx2=layer.w2_g_idx,
929
- sort_indices1=layer.w13_g_idx_sort_indices,
930
- sort_indices2=layer.w2_g_idx_sort_indices,
931
- workspace=layer.workspace,
932
- is_k_full=self.is_k_full,
933
- input_dtype=self.input_dtype,
934
- inplace=not self.moe.disable_inplace,
826
+ shared_experts_input=shared_experts_input,
827
)
vllm/model_executor/layers/quantization/utils/quant_utils.py
+8
@@ -20,6 +20,8 @@ if TYPE_CHECKING:
20
FP8_DTYPE = current_platform.fp8_dtype()
21
FP4_DTYPE = torch.uint8
22
MXFP_SCALE_DTYPE = torch.uint8
23
+INT4_DTYPE = scalar_types.uint4b8
24
+INT8_DTYPE = scalar_types.uint8b128
25
26
27
def get_fp8_min_max() -> tuple[float, float]:
@@ -170,6 +172,12 @@ kMxfp8Dynamic = QuantKey(FP8_DTYPE, scale=kMxfp8DynamicGroupScale, symmetric=Tru
172
kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32))
173
kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True)
174
175
+# TODO: convert this to use SCALAR_TYPE. This is not right.
176
+kInt4StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1))
177
+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
kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)
182
kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)
183