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
Revert "[MoE Refactor] Migrate MoeWNA16Method quantization to MK orac… (#44033)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
May 29, 2026 at 19:45 UTC
187457a952cbaf21e28944920e0b93a28f6cb1bd
8 files changed
+97
-332
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+20
-30
@@ -43,10 +43,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
43
kFp8Static128BlockSym,
44
kFp8StaticChannelSym,
45
kFp8StaticTensorSym,
46
- kInt4Static,
47
- kInt4Static32,
46
kInt8DynamicTokenSym,
49
- kInt8Static,
47
kInt8StaticChannelSym,
48
)
49
from vllm.platforms import current_platform
@@ -459,45 +456,40 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
456
class TritonWNA16Experts(TritonExperts):
457
@staticmethod
458
def _supports_current_device() -> bool:
462
- return current_platform.is_cuda_alike() or current_platform.is_xpu()
459
+ raise NotImplementedError(
460
+ "TritonWNA16Experts is not yet used by an Oracle. "
461
+ "This method should not be called."
462
+ )
463
464
@staticmethod
465
def _supports_no_act_and_mul() -> bool:
466
- return True
466
+ raise NotImplementedError(
467
+ "TritonWNA16Experts is not yet used by an Oracle. "
468
+ "This method should not be called."
469
+ )
470
471
@staticmethod
472
def _supports_quant_scheme(
473
weight_key: QuantKey | None,
474
activation_key: QuantKey | None,
475
) -> bool:
473
- SUPPORTED_W = [
474
- kInt4Static,
475
- kInt8Static,
476
- kInt4Static32,
477
- # other group sizes?
478
- ]
479
- return weight_key in SUPPORTED_W
476
+ raise NotImplementedError(
477
+ "TritonWNA16Experts is not yet used by an Oracle. "
478
+ "This method should not be called."
479
+ )
480
481
@staticmethod
482
def _supports_activation(activation: MoEActivation) -> bool:
483
- return activation in [
484
- MoEActivation.SILU,
485
- MoEActivation.GELU,
486
- MoEActivation.GELU_TANH,
487
- MoEActivation.SWIGLUOAI,
488
- MoEActivation.SWIGLUSTEP,
489
- MoEActivation.SILU_NO_MUL,
490
- MoEActivation.GELU_NO_MUL,
491
- MoEActivation.GELU_TANH_NO_MUL,
492
- MoEActivation.RELU2_NO_MUL,
493
- ]
483
+ raise NotImplementedError(
484
+ "TritonWNA16Experts is not yet used by an Oracle. "
485
+ "This method should not be called."
486
+ )
487
488
@staticmethod
489
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
497
- # Why?
498
- return not (
499
- moe_parallel_config.use_fi_nvl_two_sided_kernels
500
- or moe_parallel_config.use_fi_nvl_one_sided_kernels
490
+ raise NotImplementedError(
491
+ "TritonWNA16Experts is not yet used by an Oracle. "
492
+ "This method should not be called."
493
)
494
495
def apply(
@@ -520,9 +512,7 @@ class TritonWNA16Experts(TritonExperts):
512
):
513
# Check constraints.
514
if self.quant_config.use_int4_w4a16:
523
- assert hidden_states.size(-1) // 2 == w1.size(2), (
524
- f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}"
525
- )
515
+ assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
516
else:
517
assert hidden_states.size(-1) == w1.size(2), (
518
f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-6
@@ -30,7 +30,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
30
)
31
from vllm.platforms import current_platform
32
from vllm.triton_utils import tl, triton
33
-from vllm.utils.math_utils import next_power_of_2
33
from vllm.utils.torch_utils import direct_register_custom_op
34
35
logger = init_logger(__name__)
@@ -1244,11 +1243,7 @@ def get_default_config(
1243
bit = 4 if dtype == "int4_w4a16" else 8
1244
use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit)
1245
if use_moe_wna16_cuda:
1247
- config = {
1248
- "BLOCK_SIZE_M": min(16, next_power_of_2(M)),
1249
- "GROUP_SIZE_M": 1,
1250
- "SPLIT_K": 1,
1251
- }
1246
+ config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1}
1247
elif M <= 20:
1248
config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1}
1249
elif M <= 40:
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+6
-91
@@ -23,9 +23,6 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
23
MarlinExperts,
24
MarlinExpertsBase,
25
)
26
-from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
27
- TritonWNA16Experts,
28
-)
26
from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
27
TrtLlmMxint4ExpertsMonolithic,
28
)
@@ -34,7 +31,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
31
marlin_act_int8_process_scales,
32
marlin_moe_permute_scales,
33
marlin_permute_bias,
37
- marlin_zero_points,
34
moe_awq_to_marlin_zero_points,
35
)
36
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -49,7 +45,6 @@ class WNA16MoEBackend(Enum):
45
MARLIN = "MARLIN"
46
BATCHED_MARLIN = "BATCHED_MARLIN"
47
FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
52
- TRITON = "TRITON"
48
XPU = "XPU"
49
50
@@ -63,8 +58,6 @@ def backend_to_kernel_cls(
58
return [BatchedMarlinExperts]
59
elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
60
return [TrtLlmMxint4ExpertsMonolithic]
66
- elif backend == WNA16MoEBackend.TRITON:
67
- return [TritonWNA16Experts]
61
elif backend == WNA16MoEBackend.XPU:
62
from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
63
XPUExpertsWNA16,
@@ -75,38 +68,24 @@ def backend_to_kernel_cls(
68
raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
69
70
78
-def _get_priority_backends(
79
- may_have_zp: bool, may_have_bias: bool
80
-) -> list[WNA16MoEBackend]:
71
+def _get_priority_backends() -> list[WNA16MoEBackend]:
72
"""
73
Get available backends in priority order based on platform and config.
74
"""
75
if current_platform.is_xpu():
76
return [WNA16MoEBackend.XPU]
77
87
- _AVAILABLE_BACKENDS = []
88
-
89
- if not may_have_zp and not may_have_bias:
90
- _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER_TRTLLM)
91
-
92
- # Marlin supports ZP and bias
93
- _AVAILABLE_BACKENDS += [
78
+ _AVAILABLE_BACKENDS = [
79
+ WNA16MoEBackend.FLASHINFER_TRTLLM,
80
WNA16MoEBackend.MARLIN,
81
WNA16MoEBackend.BATCHED_MARLIN,
82
]
83
return _AVAILABLE_BACKENDS
84
99
- if not may_have_bias:
100
- _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON)
101
-
102
- return _AVAILABLE_BACKENDS
103
-
85
86
def select_wna16_moe_backend(
87
config: FusedMoEConfig,
88
weight_key: QuantKey,
108
- may_have_zp: bool,
109
- may_have_bias: bool,
89
) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
90
"""Select the WNA16 MoE backend.
91
@@ -157,7 +136,7 @@ def select_wna16_moe_backend(
136
raise ValueError(_make_log_unsupported(backend, reason))
137
138
# Select kernels in order of backend.
160
- AVAILABLE_BACKENDS = _get_priority_backends(may_have_zp, may_have_bias)
139
+ AVAILABLE_BACKENDS = _get_priority_backends()
140
141
for backend in AVAILABLE_BACKENDS:
142
activation_key = None # always BF16 activation for WNA16 MoE
@@ -239,7 +218,6 @@ def make_wna16_moe_kernel(
218
assert experts_cls in (
219
MarlinExperts,
220
BatchedMarlinExperts,
242
- TritonWNA16Experts,
221
TrtLlmMxint4ExpertsMonolithic,
222
XPUExpertsWNA16,
223
)
@@ -256,7 +234,6 @@ def make_wna16_moe_kernel(
234
assert prepare_finalize is not None
235
236
logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
259
- logger.info_once("Using %s", experts_cls.__name__, scope="local")
237
238
extra_args: dict[str, Any] = {}
239
if issubclass(experts_cls, MarlinExpertsBase):
@@ -426,8 +403,6 @@ def _process_weights_marlin(
403
w2_input_global_scale: torch.Tensor | None = None
404
w13_bias_out: torch.Tensor | None = None
405
w2_bias_out: torch.Tensor | None = None
429
- w13_qzeros_out: torch.Tensor | None = None
430
- w2_qzeros_out: torch.Tensor | None = None
406
407
# --- FP8 weight / scale adjustment ---
408
if input_dtype == torch.float8_e4m3fn:
@@ -527,24 +502,6 @@ def _process_weights_marlin(
502
if w2_bias is not None:
503
w2_bias_out = marlin_permute_bias(w2_bias)
504
530
- if w13_qzeros is not None:
531
- w13_qzeros_out = marlin_zero_points(
532
- w13_qzeros,
533
- size_k=layer.intermediate_size_per_partition,
534
- size_n=w13_qzeros.shape[2],
535
- num_bits=num_bits,
536
- is_a_8bit=is_a_8bit,
537
- )
538
-
539
- if w2_qzeros is not None:
540
- w2_qzeros_out = marlin_zero_points(
541
- w2_qzeros,
542
- size_k=w2_qzeros.shape[1] * group_size_or_pack_factor,
543
- size_n=w2_qzeros.shape[2],
544
- num_bits=num_bits,
545
- is_a_8bit=is_a_8bit,
546
- )
547
-
505
return (
506
marlin_w13_qweight,
507
marlin_w2_qweight,
@@ -554,8 +511,8 @@ def _process_weights_marlin(
511
w2_g_idx,
512
w13_g_idx_sort_indices,
513
w2_g_idx_sort_indices,
557
- w13_qzeros_out,
558
- w2_qzeros_out,
514
+ w13_qzeros,
515
+ w2_qzeros,
516
w13_input_global_scale,
517
w2_input_global_scale,
518
w13_bias_out,
@@ -823,9 +780,6 @@ def convert_to_wna16_moe_kernel_format(
780
from vllm.model_executor.layers.quantization.awq_marlin import (
781
AWQMarlinConfig,
782
)
826
- from vllm.model_executor.layers.quantization.moe_wna16 import (
827
- MoeWNA16Config,
828
- )
783
784
if isinstance(quant_config, AWQMarlinConfig):
785
if w13_qzeros is None or w2_qzeros is None:
@@ -860,17 +814,11 @@ def convert_to_wna16_moe_kernel_format(
814
pack_factor = 32 // quant_config.num_bits
815
group_size = quant_config.group_size
816
actorder = quant_config.actorder
863
- elif isinstance(quant_config, MoeWNA16Config):
864
- num_bits = quant_config.weight_bits
865
- pack_factor = quant_config.bit8_pack_factor
866
- group_size = quant_config.group_size
867
- actorder = None
817
else:
818
raise TypeError(
819
"Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or "
820
f"QuantizationArgs, got {type(quant_config).__name__}."
821
)
873
-
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(
@@ -902,39 +850,6 @@ def convert_to_wna16_moe_kernel_format(
850
w13_bias,
851
w2_bias,
852
)
905
- elif backend == WNA16MoEBackend.TRITON:
906
- # Convert from int32 to uint8 format for Triton kernel.
907
- # This changes the shape from (E, N, K // 8) to (E, N, K // 2) for int4,
908
- # which matches what the Triton kernel expects.
909
- w13_uint8 = w13.view(torch.uint8)
910
- w2_uint8 = w2.view(torch.uint8)
911
- return (
912
- w13_uint8,
913
- w2_uint8,
914
- w13_scale,
915
- w2_scale,
916
- None,
917
- None,
918
- None,
919
- None,
920
- w13_qzeros,
921
- w2_qzeros,
922
- None,
923
- None,
924
- w13_bias,
925
- w2_bias,
926
- )
927
- elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
928
- return _process_weights_flashinfer(
929
- w13,
930
- w2,
931
- w13_scale,
932
- w2_scale,
933
- w13_g_idx,
934
- w2_g_idx,
935
- w13_bias,
936
- w2_bias,
937
- )
853
elif backend == WNA16MoEBackend.XPU:
854
assert quant_config is not None
855
(
vllm/model_executor/layers/quantization/auto_gptq.py
+38
-26
@@ -485,8 +485,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
485
self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
486
moe,
487
weight_key,
488
- may_have_zp=True,
489
- may_have_bias=True,
488
)
489
490
def create_weights(
@@ -646,22 +644,12 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
644
layer.workspace = marlin_make_workspace_new(device, 4)
645
646
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
649
- def replace_or_register(name: str, val: torch.Tensor | None):
650
- if val is None:
651
- return
652
-
653
- if hasattr(layer, name):
654
- replace_parameter(layer, name, val)
655
- else:
656
- layer.register_parameter(
657
- name, torch.nn.Parameter(val, requires_grad=False)
658
- )
659
-
647
is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
648
662
- assert not is_a_8bit or self.quant_config.quant_type.size_bits == 8, (
663
- "W8A8-INT8 is not supported by marlin kernel."
664
- )
649
+ if is_a_8bit:
650
+ assert self.quant_config.quant_type.size_bits == 8, (
651
+ "W8A8-INT8 is not supported by marlin kernel."
652
+ )
653
654
(
655
w13,
@@ -672,8 +660,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
660
w2_g_idx,
661
w13_g_idx_sort_indices,
662
w2_g_idx_sort_indices,
675
- w13_qzeros,
676
- w2_qzeros,
663
+ _w13_qzeros,
664
+ _w2_qzeros,
665
w13_input_global_scale,
666
w2_input_global_scale,
667
w13_bias,
@@ -691,8 +679,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
679
w2_g_idx=layer.w2_g_idx,
680
w13_bias=getattr(layer, "w13_bias", None),
681
w2_bias=getattr(layer, "w2_bias", None),
694
- w13_qzeros=getattr(layer, "w13_qzeros", None),
695
- w2_qzeros=getattr(layer, "w2_qzeros", None),
682
)
683
684
replace_parameter(layer, "w13_qweight", w13)
@@ -703,12 +689,38 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
689
replace_parameter(layer, "w2_g_idx", w2_g_idx)
690
replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
691
replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
706
- replace_or_register("w13_input_global_scale", w13_input_global_scale)
707
- replace_or_register("w2_input_global_scale", w2_input_global_scale)
708
- replace_or_register("w13_bias", w13_bias)
709
- replace_or_register("w2_bias", w2_bias)
710
- replace_or_register("w13_qzeros", w13_qzeros)
711
- replace_or_register("w2_qzeros", w2_qzeros)
692
+ if w13_input_global_scale is not None:
693
+ if hasattr(layer, "w13_input_global_scale"):
694
+ replace_parameter(
695
+ layer, "w13_input_global_scale", w13_input_global_scale
696
+ )
697
+ else:
698
+ layer.register_parameter(
699
+ "w13_input_global_scale",
700
+ torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
701
+ )
702
+ if w2_input_global_scale is not None:
703
+ if hasattr(layer, "w2_input_global_scale"):
704
+ replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale)
705
+ else:
706
+ layer.register_parameter(
707
+ "w2_input_global_scale",
708
+ torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
709
+ )
710
+ if w13_bias is not None:
711
+ if hasattr(layer, "w13_bias"):
712
+ replace_parameter(layer, "w13_bias", w13_bias)
713
+ else:
714
+ layer.register_parameter(
715
+ "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False)
716
+ )
717
+ if w2_bias is not None:
718
+ if hasattr(layer, "w2_bias"):
719
+ replace_parameter(layer, "w2_bias", w2_bias)
720
+ else:
721
+ layer.register_parameter(
722
+ "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False)
723
+ )
724
725
self._setup_kernel(layer)
726
vllm/model_executor/layers/quantization/awq_marlin.py
-2
@@ -524,8 +524,6 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
524
self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
525
moe,
526
kInt4Static,
527
- may_have_zp=self.quant_config.zero_point,
528
- may_have_bias=True,
527
)
528
529
def create_weights(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+7
-10
@@ -88,13 +88,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
88
self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
89
config=self.moe,
90
weight_key=weight_key,
91
- may_have_zp=False,
92
- may_have_bias=False,
91
)
94
- self.is_marlin = self.wna16_backend in [
95
- WNA16MoEBackend.MARLIN,
96
- WNA16MoEBackend.BATCHED_MARLIN,
97
- ]
92
93
def get_weight_shape(
94
self,
@@ -120,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": (
@@ -162,7 +157,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
157
"Marlin": (num_experts, num_groups_w2, hidden_size),
158
},
159
}
165
- backend_key = "Marlin" if self.is_marlin else "Flashinfer"
160
+ backend_key = "Flashinfer" if is_flashinfer else "Marlin"
161
return shape_map[weight_name][backend_key]
162
163
def create_weights(
@@ -179,8 +174,9 @@ 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
177
+ is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
178
extra_weight_attrs.update(
183
- {"is_transposed": self.is_marlin, "quant_method": self.strategy}
179
+ {"is_transposed": is_transposed, "quant_method": self.strategy}
180
)
181
182
w13_weight = torch.nn.Parameter(
@@ -328,6 +324,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
324
325
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
326
# Process weights using the shared oracle infrastructure
327
+ is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
328
(
329
w13_qweight,
330
w2_qweight,
@@ -363,7 +360,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
360
replace_parameter(layer, "w2_weight_scale", w2_scales)
361
362
# Marlin-specific parameters (not needed for Flashinfer)
366
- if self.is_marlin:
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)
@@ -395,7 +392,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
392
393
# Add Marlin-specific arguments
394
marlin_args: dict[str, Any] = {}
398
- if self.is_marlin:
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,
vllm/model_executor/layers/quantization/moe_wna16.py
+24
-150
@@ -13,15 +13,11 @@ from vllm.model_executor.layers.fused_moe import (
13
RoutedExperts,
14
SharedExperts,
15
)
16
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
17
from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEQuantConfig,
18
-)
19
-from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
20
- WNA16MoEBackend,
21
- convert_to_wna16_moe_kernel_format,
22
- make_wna16_moe_kernel,
23
- make_wna16_moe_quant_config,
24
- select_wna16_moe_backend,
19
+ int4_w4a16_moe_quant_config,
20
+ int8_w8a16_moe_quant_config,
21
)
22
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
23
UnquantizedFusedMoEMethod,
@@ -35,15 +31,7 @@ from vllm.model_executor.layers.quantization.base_config import (
31
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
32
check_marlin_supports_layer,
33
)
38
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
39
- INT4_DTYPE,
40
- INT8_DTYPE,
41
- QuantKey,
42
- kInt4Static32GroupScale,
43
- kInt4StaticGroupScale,
44
- kInt8StaticGroupScale,
45
-)
46
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
34
+from vllm.model_executor.utils import set_weight_attrs
35
from vllm.platforms import current_platform
36
37
@@ -228,33 +216,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
216
super().__init__(moe)
217
self.quant_config = quant_config
218
231
- num_bits = self.quant_config.weight_bits
232
- group_size = self.quant_config.group_size
233
-
234
- if num_bits == 4:
235
- quant_type = INT4_DTYPE
236
- if group_size == 32:
237
- scale = kInt4Static32GroupScale
238
- else:
239
- scale = kInt4StaticGroupScale
240
- elif num_bits == 8:
241
- assert group_size == -1
242
- quant_type = INT8_DTYPE
243
- scale = kInt8StaticGroupScale
244
- else:
245
- raise ValueError("MoeWNA16Method only supports int4 and int8 now.")
246
-
247
- weight_key = QuantKey(quant_type, scale)
248
-
249
- # Select WNA16 MoE backend via oracle.
250
- # handle ZP?
251
- self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
252
- config=self.moe,
253
- weight_key=weight_key,
254
- may_have_zp=self.quant_config.has_zp,
255
- may_have_bias=False,
256
- )
257
-
219
def create_weights(
220
self,
221
layer: RoutedExperts,
@@ -375,89 +336,24 @@ class MoeWNA16Method(FusedMoEMethodBase):
336
layer.register_parameter(key, param)
337
set_weight_attrs(param, extra_weight_attrs)
338
378
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
379
- has_zp = self.quant_config.has_zp
380
- (
381
- w13_qweight,
382
- w2_qweight,
383
- w13_scales,
384
- w2_scales,
385
- w13_g_idx_processed,
386
- w2_g_idx_processed,
387
- w13_g_idx_sort_indices,
388
- w2_g_idx_sort_indices,
389
- w13_qzeros,
390
- w2_qzeros,
391
- w13_input_global_scale,
392
- w2_input_global_scale,
393
- _, # w13_bias
394
- _, # w2_bias
395
- ) = convert_to_wna16_moe_kernel_format(
396
- backend=self.wna16_backend,
397
- layer=layer,
398
- quant_config=self.quant_config,
399
- input_dtype=None,
400
- w13=layer.w13_qweight,
401
- w2=layer.w2_qweight,
402
- w13_scale=layer.w13_scales,
403
- w2_scale=layer.w2_scales,
404
- w13_g_idx=getattr(layer, "w13_g_idx", None),
405
- w2_g_idx=getattr(layer, "w2_g_idx", None),
406
- w13_qzeros=layer.w13_qzeros if has_zp else None,
407
- w2_qzeros=layer.w2_qzeros if has_zp else None,
408
- )
409
-
410
- # Replace common parameters
411
- replace_parameter(layer, "w13_qweight", w13_qweight)
412
- replace_parameter(layer, "w2_qweight", w2_qweight)
413
- replace_parameter(layer, "w13_scales", w13_scales)
414
- replace_parameter(layer, "w2_scales", w2_scales)
415
-
416
- if has_zp:
417
- assert w13_qzeros is not None and w2_qzeros is not None
418
- replace_parameter(layer, "w13_qzeros", w13_qzeros)
419
- replace_parameter(layer, "w2_qzeros", w2_qzeros)
420
-
421
- # Marlin-specific parameters (not needed for Flashinfer)
422
- if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM:
423
- replace_parameter(layer, "w13_g_idx", w13_g_idx_processed)
424
- replace_parameter(layer, "w2_g_idx", w2_g_idx_processed)
425
- replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
426
- replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
427
-
428
- # Register input global scales if present
429
- if w13_input_global_scale is not None:
430
- layer.register_parameter(
431
- "w13_input_global_scale",
432
- torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
433
- )
434
- if w2_input_global_scale is not None:
435
- layer.register_parameter(
436
- "w2_input_global_scale",
437
- torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
438
- )
439
-
440
- assert self.experts_cls is not None
441
- self.moe_quant_config = self.get_fused_moe_quant_config(layer)
442
- assert self.moe_quant_config is not None
443
- self.moe_kernel = make_wna16_moe_kernel(
444
- moe_quant_config=self.moe_quant_config,
445
- moe_config=self.moe,
446
- experts_cls=self.experts_cls,
447
- routing_tables=layer._expert_routing_tables(),
448
- )
449
-
339
def get_fused_moe_quant_config(
340
self, layer: RoutedExperts
341
) -> FusedMoEQuantConfig | None:
342
+ weight_bits = self.quant_config.weight_bits
343
has_zp = self.quant_config.has_zp
454
- return make_wna16_moe_quant_config(
344
+ assert weight_bits == 4 or weight_bits == 8
345
+ config_builder = (
346
+ int4_w4a16_moe_quant_config
347
+ if weight_bits == 4
348
+ else int8_w8a16_moe_quant_config
349
+ )
350
+
351
+ return config_builder(
352
w1_scale=layer.w13_scales,
353
w2_scale=layer.w2_scales,
354
w1_zp=layer.w13_qzeros if has_zp else None,
355
w2_zp=layer.w2_qzeros if has_zp else None,
459
- group_size=layer.group_size,
460
- num_bits=self.quant_config.weight_bits,
356
+ block_shape=[0, layer.group_size],
357
)
358
359
def apply(
@@ -469,44 +365,22 @@ class MoeWNA16Method(FusedMoEMethodBase):
365
shared_experts: SharedExperts | None,
366
shared_experts_input: torch.Tensor | None,
367
) -> torch.Tensor:
472
- assert not self.is_monolithic
473
- assert self.moe_kernel is not None
474
- return self.moe_kernel.apply(
368
+ from vllm.model_executor.layers.fused_moe import fused_experts
369
+
370
+ assert layer.activation == MoEActivation.SILU, (
371
+ f"Only SiLU activation is supported, not {layer.activation}."
372
+ )
373
+
374
+ return fused_experts(
375
x,
376
layer.w13_qweight,
377
layer.w2_qweight,
478
- topk_weights,
479
- topk_ids,
480
- activation=layer.activation,
481
- global_num_experts=layer.global_num_experts,
482
- expert_map=layer.expert_map,
378
+ topk_weights=topk_weights,
379
+ topk_ids=topk_ids,
380
apply_router_weight_on_input=layer.apply_router_weight_on_input,
484
- shared_experts=shared_experts,
485
- shared_experts_input=shared_experts_input,
486
- )
487
-
488
- def apply_monolithic(
489
- self,
490
- layer: RoutedExperts,
491
- x: torch.Tensor,
492
- router_logits: torch.Tensor,
493
- input_ids: torch.Tensor | None = None,
494
- ) -> torch.Tensor:
495
- assert self.is_monolithic
496
- assert self.moe_kernel is not None
497
- return self.moe_kernel.apply_monolithic(
498
- x,
499
- layer.w13_weight,
500
- layer.w2_weight,
501
- router_logits,
502
- activation=layer.activation,
381
global_num_experts=layer.global_num_experts,
382
expert_map=layer.expert_map,
505
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
506
- num_expert_group=layer.num_expert_group,
507
- topk_group=layer.topk_group,
508
- e_score_correction_bias=layer.e_score_correction_bias,
509
- routed_scaling_factor=layer.routed_scaling_factor,
383
+ quant_config=self.moe_quant_config,
384
)
385
386
@staticmethod
vllm/model_executor/layers/quantization/utils/gptq_utils.py
+1
-17
@@ -3,7 +3,7 @@
3
from collections.abc import Mapping
4
from copy import deepcopy
5
from types import MappingProxyType
6
-from typing import TYPE_CHECKING, Any
6
+from typing import TYPE_CHECKING
7
8
import regex as re
9
import torch
@@ -69,20 +69,6 @@ def get_dynamic_override(
69
return default_value
70
71
72
-def flatten_list(lst: list[Any]) -> list[Any]:
73
- output = []
74
-
75
- def _flatten(lst: list[Any]):
76
- for i in lst:
77
- if isinstance(i, list):
78
- _flatten(i)
79
- else:
80
- output.append(i)
81
-
82
- _flatten(lst)
83
- return output
84
-
85
-
72
def is_layer_gptq_quantized(
73
prefix: str,
74
quantized_layers: list[str],
@@ -97,8 +83,6 @@ def is_layer_gptq_quantized(
83
84
proj_name = prefix.split(".")[-1]
85
100
- quantized_layers = flatten_list(quantized_layers)
101
-
86
# Fused layers like gate_up_proj or qkv_proj will not be fused
87
# in the safetensors checkpoint. So, we convert the name
88
# from the fused version to unfused + check to make sure that