@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ec59c1579f

[MoE Refactor] Migrate MoeWNA16Method quantization method over to using the new MK oracle scheme. (#44120)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mgoin <mgoin64@gmail.com>

bnellnm committed Jul 21, 2026 at 22:20 UTC ec59c1579f582ac1f8f9ac5871490138f9a84b14
11 files changed +805 -195
tests/quantization/test_auto_gptq.py
+79
@@ -5,13 +5,17 @@
5 Run `pytest tests/quantization/test_auto_gptq.py -v -s`.
6 """
7
8 +from types import SimpleNamespace
9 +
10 import pytest
11 import torch
12
13 from tests.quantization.utils import is_quant_method_supported
14 +from vllm.model_executor.layers.fused_moe import RoutedExperts
15 from vllm.model_executor.layers.quantization.auto_gptq import (
16 AutoGPTQConfig,
17 AutoGPTQLinearMethod,
18 + AutoGPTQMoEMethod,
19 )
20
21 PROMPT = "On the surface of Mars, we found"
@@ -54,3 +58,78 @@ def test_auto_gptq_quantization_method(vllm_runner, model_id: str, monkeypatch):
58 def test_auto_gptq_config_get_name():
59 """Test that AutoGPTQConfig.get_name() returns 'auto_gptq'."""
60 assert AutoGPTQConfig.get_name() == "auto_gptq"
61 +
62 +
63 +def test_auto_gptq_moe_creates_zero_initialized_expert_biases():
64 + method = object.__new__(AutoGPTQMoEMethod)
65 + method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
66 + method.input_dtype = None
67 + method.experts_cls = None
68 + layer = torch.nn.Module()
69 +
70 + method.create_weights(
71 + layer=layer,
72 + num_experts=2,
73 + hidden_size=8,
74 + intermediate_size_per_partition=4,
75 + params_dtype=torch.float16,
76 + intermediate_size_full=4,
77 + weight_loader=lambda *args, **kwargs: None,
78 + )
79 +
80 + assert layer.w13_bias.shape == (2, 8)
81 + assert layer.w2_bias.shape == (2, 8)
82 + assert torch.count_nonzero(layer.w13_bias) == 0
83 + assert torch.count_nonzero(layer.w2_bias) == 0
84 +
85 +
86 +def test_routed_experts_loads_per_expert_biases():
87 + class Loader:
88 + quant_config = None
89 + quant_method = object()
90 + moe_config = SimpleNamespace(
91 + is_act_and_mul=True,
92 + tp_rank=0,
93 + moe_parallel_config=SimpleNamespace(tp_size=1),
94 + )
95 + _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim)
96 + _narrow_expert_data_for_padding = staticmethod(
97 + RoutedExperts._narrow_expert_data_for_padding
98 + )
99 + _load_w13 = RoutedExperts._load_w13
100 + _loaded_expert_biases = set()
101 +
102 + @staticmethod
103 + def _map_global_expert_id_to_local_expert_id(expert_id):
104 + return expert_id
105 +
106 + loader = Loader()
107 + w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False)
108 + w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False)
109 +
110 + for shard_id, loaded in (
111 + ("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])),
112 + ("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])),
113 + ):
114 + assert RoutedExperts.weight_loader(
115 + loader,
116 + w13_bias,
117 + loaded,
118 + weight_name="model.layers.0.mlp.experts.w13_bias",
119 + shard_id=shard_id,
120 + expert_id=0,
121 + return_success=True,
122 + )
123 +
124 + assert RoutedExperts.weight_loader(
125 + loader,
126 + w2_bias,
127 + torch.tensor([9.0, 10.0, 11.0, 12.0]),
128 + weight_name="model.layers.0.mlp.experts.w2_bias",
129 + shard_id="w2",
130 + expert_id=0,
131 + return_success=True,
132 + )
133 + assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8))
134 + assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4))
135 + assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"}
tests/quantization/test_moe_wna16.py
+179 -30
@@ -5,45 +5,194 @@ from types import SimpleNamespace
5
6 import pytest
7 import torch
8 +from compressed_tensors.quantization import (
9 + ActivationOrdering,
10 + QuantizationArgs,
11 + QuantizationStrategy,
12 + QuantizationType,
13 +)
14
9 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10 -from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method
11 -from vllm.platforms import current_platform
15 +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
16 + WNA16MoEBackend,
17 + _backend_incompatibility_reason,
18 + _convert_moe_wna16_humming_tensors,
19 + convert_to_wna16_moe_kernel_format,
20 + map_wna16_backend,
21 +)
22 +from vllm.model_executor.layers.quantization import moe_wna16
23 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
24 +from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
25 +from vllm.model_executor.layers.quantization.moe_wna16 import (
26 + MoeWNA16Config,
27 + MoeWNA16Method,
28 +)
29
30
14 -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
15 -def test_moe_wna16_apply_passes_layer_activation(monkeypatch):
16 - captured_kwargs = {}
31 +def test_map_wna16_backend_supports_triton():
32 + assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON
33
18 - def fake_fused_experts(*args, **kwargs):
19 - captured_kwargs.update(kwargs)
20 - return torch.empty(1, 2)
34
22 - monkeypatch.setattr(
23 - "vllm.model_executor.layers.fused_moe.fused_experts",
24 - fake_fused_experts,
35 +@pytest.mark.parametrize(
36 + ("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"),
37 + [
38 + (
39 + WNA16MoEBackend.TRITON,
40 + AutoAWQConfig(4, 128, True, False),
41 + True,
42 + False,
43 + "AutoAWQ weight layout",
44 + ),
45 + (
46 + WNA16MoEBackend.TRITON,
47 + AutoGPTQConfig(4, 128, True, True, False, {}, {}),
48 + False,
49 + False,
50 + "activation ordering",
51 + ),
52 + (
53 + WNA16MoEBackend.TRITON,
54 + QuantizationArgs(
55 + num_bits=4,
56 + type=QuantizationType.INT,
57 + strategy=QuantizationStrategy.GROUP,
58 + symmetric=True,
59 + dynamic=False,
60 + group_size=128,
61 + actorder=ActivationOrdering.GROUP,
62 + ),
63 + False,
64 + False,
65 + "activation ordering",
66 + ),
67 + (
68 + WNA16MoEBackend.TRITON,
69 + AutoGPTQConfig(4, 128, False, True, False, {}, {}),
70 + False,
71 + True,
72 + "bias",
73 + ),
74 + (
75 + WNA16MoEBackend.MARLIN,
76 + MoeWNA16Config(
77 + linear_quant_method="gptq",
78 + weight_bits=4,
79 + group_size=128,
80 + has_zp=False,
81 + lm_head_quantized=False,
82 + modules_to_not_convert=None,
83 + full_config={},
84 + ),
85 + False,
86 + False,
87 + "MoeWNA16 checkpoint layout",
88 + ),
89 + ],
90 +)
91 +def test_wna16_oracle_rejects_incompatible_quant_structures(
92 + backend, quant_config, may_have_zp, may_have_bias, expected
93 +):
94 + reason = _backend_incompatibility_reason(
95 + backend=backend,
96 + quant_config=quant_config,
97 + may_have_zp=may_have_zp,
98 + may_have_bias=may_have_bias,
99 + )
100 +
101 + assert reason is not None
102 + assert expected in reason
103 +
104 +
105 +def test_compressed_tensors_weights_are_transposed_for_triton():
106 + quant_config = QuantizationArgs(
107 + num_bits=4,
108 + type=QuantizationType.INT,
109 + strategy=QuantizationStrategy.GROUP,
110 + symmetric=True,
111 + dynamic=False,
112 + group_size=32,
113 )
114 + w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8)
115 + w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6)
116 + w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8)
117 + w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6)
118
119 + converted = convert_to_wna16_moe_kernel_format(
120 + backend=WNA16MoEBackend.TRITON,
121 + layer=torch.nn.Module(),
122 + quant_config=quant_config,
123 + input_dtype=None,
124 + w13=w13,
125 + w2=w2,
126 + w13_scale=w13_scale,
127 + w2_scale=w2_scale,
128 + )
129 +
130 + assert converted is not None
131 + assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8))
132 + assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8))
133 + assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous())
134 + assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous())
135 +
136 +
137 +def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):
138 method = object.__new__(MoeWNA16Method)
28 - method.moe = SimpleNamespace(disable_inplace=False)
29 - method.moe_quant_config = object()
30 - layer = SimpleNamespace(
31 - w13_qweight=torch.empty(1, 2),
32 - w2_qweight=torch.empty(1, 2),
33 - activation=MoEActivation.GELU_TANH,
34 - apply_router_weight_on_input=False,
35 - global_num_experts=1,
36 - expert_map=None,
139 + method.experts_cls = object
140 + method.wna16_backend = WNA16MoEBackend.HUMMING
141 + method.moe = object()
142 + quant_config = object()
143 + method.get_fused_moe_quant_config = lambda layer: quant_config
144 + layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None))
145 + captured = {}
146 + kernel = object()
147 +
148 + def fake_make_wna16_moe_kernel(**kwargs):
149 + captured.update(kwargs)
150 + return kernel
151 +
152 + monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel)
153 +
154 + method._setup_kernel(layer)
155 +
156 + assert method.moe_kernel is kernel
157 + assert captured["backend"] == WNA16MoEBackend.HUMMING
158 + assert captured["layer"] is layer
159 +
160 +
161 +def test_moe_wna16_humming_adapter_repacks_uint8_tensors():
162 + qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8)
163 + scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4)
164 + qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2)
165 +
166 + converted = _convert_moe_wna16_humming_tensors(
167 + {"qweight": qweight, "scales": scales, "qzeros": qzeros},
168 + has_zero_point=True,
169 )
170
39 - output = method.apply(
40 - layer,
41 - x=torch.empty(1, 2),
42 - topk_weights=torch.empty(1, 1),
43 - topk_ids=torch.empty(1, 1, dtype=torch.int32),
44 - shared_experts=None,
45 - shared_experts_input=None,
171 + assert torch.equal(converted["weight"], qweight.view(torch.int32))
172 + assert converted["weight"].shape == (1, 4, 2)
173 + assert torch.equal(converted["weight_scale"], scales)
174 + expected_qzeros = (
175 + qzeros.transpose(-1, -2)
176 + .contiguous()
177 + .view(torch.int32)
178 + .transpose(-1, -2)
179 + .contiguous()
180 + )
181 + assert torch.equal(converted["zero_point"], expected_qzeros)
182 + assert converted["zero_point"].shape == (1, 2, 2)
183 +
184 +
185 +def test_moe_wna16_uses_humming_quant_config(monkeypatch):
186 + from vllm.model_executor.layers.quantization.utils import humming_utils
187 +
188 + method = object.__new__(MoeWNA16Method)
189 + method.wna16_backend = WNA16MoEBackend.HUMMING
190 + layer = object()
191 + quant_config = object()
192 + monkeypatch.setattr(
193 + humming_utils,
194 + "get_humming_moe_quant_config",
195 + lambda actual_layer: quant_config if actual_layer is layer else None,
196 )
197
48 - assert output.shape == (1, 2)
49 - assert captured_kwargs["activation"] is MoEActivation.GELU_TANH
198 + assert method.get_fused_moe_quant_config(layer) is quant_config
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+30 -20
@@ -43,7 +43,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
43 kFp8Static128BlockSym,
44 kFp8StaticChannelSym,
45 kFp8StaticTensorSym,
46 + kInt4Static,
47 + kInt4Static32,
48 kInt8DynamicTokenSym,
49 + kInt8Static,
50 kInt8StaticChannelSym,
51 )
52 from vllm.platforms import current_platform
@@ -542,40 +545,45 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
545 class TritonWNA16Experts(TritonExperts):
546 @staticmethod
547 def _supports_current_device() -> bool:
545 - raise NotImplementedError(
546 - "TritonWNA16Experts is not yet used by an Oracle. "
547 - "This method should not be called."
548 - )
548 + return current_platform.is_cuda_alike() or current_platform.is_xpu()
549
550 @staticmethod
551 def _supports_no_act_and_mul() -> bool:
552 - raise NotImplementedError(
553 - "TritonWNA16Experts is not yet used by an Oracle. "
554 - "This method should not be called."
555 - )
552 + return True
553
554 @staticmethod
555 def _supports_quant_scheme(
556 weight_key: QuantKey | None,
557 activation_key: QuantKey | None,
558 ) -> bool:
562 - raise NotImplementedError(
563 - "TritonWNA16Experts is not yet used by an Oracle. "
564 - "This method should not be called."
565 - )
559 + SUPPORTED_W = [
560 + kInt4Static,
561 + kInt8Static,
562 + kInt4Static32,
563 + # other group sizes?
564 + ]
565 + return weight_key in SUPPORTED_W
566
567 @staticmethod
568 def _supports_activation(activation: MoEActivation) -> bool:
569 - raise NotImplementedError(
570 - "TritonWNA16Experts is not yet used by an Oracle. "
571 - "This method should not be called."
572 - )
569 + return activation in [
570 + MoEActivation.SILU,
571 + MoEActivation.GELU,
572 + MoEActivation.GELU_TANH,
573 + MoEActivation.SWIGLUOAI,
574 + MoEActivation.SWIGLUSTEP,
575 + MoEActivation.SILU_NO_MUL,
576 + MoEActivation.GELU_NO_MUL,
577 + MoEActivation.GELU_TANH_NO_MUL,
578 + MoEActivation.RELU2_NO_MUL,
579 + ]
580
581 @staticmethod
582 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
576 - raise NotImplementedError(
577 - "TritonWNA16Experts is not yet used by an Oracle. "
578 - "This method should not be called."
583 + # Why?
584 + return not (
585 + moe_parallel_config.use_fi_nvl_two_sided_kernels
586 + or moe_parallel_config.use_fi_nvl_one_sided_kernels
587 )
588
589 def apply(
@@ -598,7 +606,9 @@ class TritonWNA16Experts(TritonExperts):
606 ):
607 # Check constraints.
608 if self.quant_config.use_int4_w4a16:
601 - assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
609 + assert hidden_states.size(-1) // 2 == w1.size(2), (
610 + f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}"
611 + )
612 else:
613 assert hidden_states.size(-1) == w1.size(2), (
614 f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
vllm/model_executor/layers/fused_moe/fused_moe.py
+6 -1
@@ -31,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
31 )
32 from vllm.platforms import current_platform
33 from vllm.triton_utils import tl, triton
34 +from vllm.utils.math_utils import next_power_of_2
35 from vllm.utils.platform_utils import get_device_name_as_file_name
36 from vllm.utils.torch_utils import direct_register_custom_op
37
@@ -1298,7 +1299,11 @@ def get_default_config(
1299 bit = 4 if dtype == "int4_w4a16" else 8
1300 use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit)
1301 if use_moe_wna16_cuda:
1301 - config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1}
1302 + config = {
1303 + "BLOCK_SIZE_M": min(16, next_power_of_2(M)),
1304 + "GROUP_SIZE_M": 1,
1305 + "SPLIT_K": 1,
1306 + }
1307 elif M <= 20:
1308 config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1}
1309 elif M <= 40:
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+231 -71
@@ -24,6 +24,9 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
24 MarlinExperts,
25 MarlinExpertsBase,
26 )
27 +from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
28 + TritonWNA16Experts,
29 +)
30 from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
31 TrtLlmMxint4ExpertsMonolithic,
32 )
@@ -50,6 +53,7 @@ class WNA16MoEBackend(Enum):
53 HUMMING = "HUMMING"
54 CPU = "CPU"
55 FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
56 + TRITON = "TRITON"
57 XPU = "XPU"
58 EMULATION = "EMULATION"
59
@@ -76,6 +80,8 @@ def backend_to_kernel_cls(
80 return [BatchedMarlinExperts]
81 elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
82 return [TrtLlmMxint4ExpertsMonolithic]
83 + elif backend == WNA16MoEBackend.TRITON:
84 + return [TritonWNA16Experts]
85 elif backend == WNA16MoEBackend.XPU:
86 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
87 XPUExpertsWNA16,
@@ -107,19 +113,56 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
113 if current_platform.is_xpu():
114 return [WNA16MoEBackend.XPU]
115
110 - _AVAILABLE_BACKENDS = [
116 + return [
117 WNA16MoEBackend.FLASHINFER_TRTLLM,
118 WNA16MoEBackend.MARLIN,
119 WNA16MoEBackend.BATCHED_MARLIN,
120 + WNA16MoEBackend.TRITON,
121 WNA16MoEBackend.HUMMING,
122 WNA16MoEBackend.EMULATION,
123 ]
117 - return _AVAILABLE_BACKENDS
124 +
125 +
126 +def _backend_incompatibility_reason(
127 + backend: WNA16MoEBackend,
128 + quant_config: QuantizationConfig | QuantizationArgs,
129 + may_have_zp: bool,
130 + may_have_bias: bool,
131 +) -> str | None:
132 + if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias):
133 + return "zero points and bias are not supported"
134 +
135 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
136 + from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
137 + from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
138 +
139 + if backend == WNA16MoEBackend.TRITON:
140 + if may_have_bias:
141 + return "expert bias is not supported"
142 + if isinstance(quant_config, AutoAWQConfig):
143 + return "the AutoAWQ weight layout is not supported"
144 + if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act:
145 + return "GPTQ activation ordering is not supported"
146 + if (
147 + isinstance(quant_config, QuantizationArgs)
148 + and quant_config.actorder == "group"
149 + ):
150 + return "group activation ordering is not supported"
151 +
152 + if isinstance(quant_config, MoeWNA16Config) and backend in (
153 + WNA16MoEBackend.MARLIN,
154 + WNA16MoEBackend.BATCHED_MARLIN,
155 + WNA16MoEBackend.EMULATION,
156 + ):
157 + return "the MoeWNA16 checkpoint layout is not supported"
158 +
159 + return None
160
161
162 def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend:
163 """Map user's MoEBackend to WNA16MoEBackend."""
164 mapping = {
165 + "triton": WNA16MoEBackend.TRITON,
166 "marlin": WNA16MoEBackend.MARLIN,
167 "humming": WNA16MoEBackend.HUMMING,
168 "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM,
@@ -136,6 +179,9 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend:
179 def select_wna16_moe_backend(
180 config: FusedMoEConfig,
181 weight_key: QuantKey,
182 + quant_config: QuantizationConfig | QuantizationArgs,
183 + may_have_zp: bool,
184 + may_have_bias: bool,
185 ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
186 """Select the WNA16 MoE backend.
187
@@ -143,6 +189,9 @@ def select_wna16_moe_backend(
189 config: the shared ``FusedMoEConfig`` for this layer.
190 weight_key: The QuantKey describing the weight quantization.
191 Must have int4 or int8 type.
192 + quant_config: Quantization structure and checkpoint format description.
193 + may_have_zp: Whether the integration can provide weight zero points.
194 + may_have_bias: Whether the integration can provide expert bias.
195
196 Returns:
197 A tuple of (``WNA16MoEBackend``, experts class or ``None``).
@@ -189,6 +238,11 @@ def select_wna16_moe_backend(
238 runner_backend = config.moe_backend
239 if runner_backend != "auto":
240 requested_backend = map_wna16_backend(runner_backend)
241 + reason = _backend_incompatibility_reason(
242 + requested_backend, quant_config, may_have_zp, may_have_bias
243 + )
244 + if reason is not None:
245 + raise ValueError(_make_log_unsupported(requested_backend, reason))
246 return _return_or_raise(
247 requested_backend, config, weight_key, None, activation_format
248 )
@@ -197,6 +251,12 @@ def select_wna16_moe_backend(
251 AVAILABLE_BACKENDS = _get_priority_backends()
252
253 for backend in AVAILABLE_BACKENDS:
254 + reason = _backend_incompatibility_reason(
255 + backend, quant_config, may_have_zp, may_have_bias
256 + )
257 + if reason is not None:
258 + logger.debug_once(_make_log_unsupported(backend, reason), scope="local")
259 + continue
260 activation_key = None # always BF16 activation for WNA16 MoE
261 for k_cls in backend_to_kernel_cls(backend):
262 supported, reason = k_cls.is_supported_config(
@@ -294,6 +354,7 @@ def make_wna16_moe_kernel(
354 allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = (
355 MarlinExperts,
356 BatchedMarlinExperts,
357 + TritonWNA16Experts,
358 TrtLlmMxint4ExpertsMonolithic,
359 XPUExpertsWNA16,
360 CPUExpertsInt4,
@@ -315,6 +376,7 @@ def make_wna16_moe_kernel(
376 assert prepare_finalize is not None
377
378 logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
379 + logger.info_once("Using %s", experts_cls.__name__, scope="local")
380
381 extra_args: dict[str, Any] = {}
382 if backend == WNA16MoEBackend.HUMMING:
@@ -329,35 +391,17 @@ def make_wna16_moe_kernel(
391 "is_k_full": is_k_full,
392 }
393
332 - if experts_cls is XPUExpertsWNA16:
333 - assert (
334 - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.Standard
335 - ), (
336 - "XPUExpertsWNA16 only supports the Standard activation format; "
337 - "xpu_fused_moe(is_int4=True) does not implement BatchedExperts."
338 - )
339 - experts: mk.FusedMoEExperts = XPUExpertsWNA16(
340 - moe_config=moe_config,
341 - quant_config=moe_quant_config,
342 - )
343 - elif (
344 - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts
345 - ):
394 + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
395 max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
396 assert max_num_tokens is not None
348 - experts = experts_cls(
349 - max_num_tokens=max_num_tokens,
350 - num_dispatchers=prepare_finalize.num_dispatchers(),
351 - moe_config=moe_config,
352 - quant_config=moe_quant_config,
353 - **extra_args,
354 - )
355 - else:
356 - experts = experts_cls(
357 - moe_config=moe_config,
358 - quant_config=moe_quant_config,
359 - **extra_args,
360 - )
397 + extra_args["max_num_tokens"] = max_num_tokens
398 + extra_args["num_dispatchers"] = prepare_finalize.num_dispatchers()
399 +
400 + experts = experts_cls(
401 + moe_config=moe_config,
402 + quant_config=moe_quant_config,
403 + **extra_args,
404 + )
405
406 return mk.FusedMoEKernel(
407 prepare_finalize,
@@ -1034,11 +1078,68 @@ def _humming_wna16_weight_schema(
1078 "sym": quant_config.is_sym,
1079 }
1080 raise TypeError(
1037 - "Humming WNA16 MoE requires AutoAWQConfig or AutoGPTQConfig, "
1081 + "Humming WNA16 checkpoint schema requires AutoAWQConfig or "
1082 + "AutoGPTQConfig, "
1083 f"got {type(quant_config).__name__}."
1084 )
1085
1086
1087 +def _convert_moe_wna16_humming_tensors(
1088 + tensors: dict[str, torch.Tensor], has_zero_point: bool
1089 +) -> dict[str, torch.Tensor]:
1090 + """Convert MoeWNA16's N-first uint8 packing to Humming's int32 packing."""
1091 + if sys.byteorder != "little":
1092 + raise NotImplementedError(
1093 + "MoeWNA16 to Humming conversion requires a little-endian host."
1094 + )
1095 +
1096 + output = {
1097 + "weight": tensors["qweight"].contiguous().view(torch.int32),
1098 + "weight_scale": tensors["scales"],
1099 + }
1100 + if has_zero_point:
1101 + qzeros = tensors["qzeros"]
1102 + output["zero_point"] = (
1103 + qzeros.transpose(-1, -2)
1104 + .contiguous()
1105 + .view(torch.int32)
1106 + .transpose(-1, -2)
1107 + .contiguous()
1108 + )
1109 + return output
1110 +
1111 +
1112 +class _MoeWNA16HummingWeightSchema:
1113 + """Adapter from MoeWNA16's generic packed layout to Humming's layout."""
1114 +
1115 + def __init__(self, bits: int, group_size: int, has_zero_point: bool) -> None:
1116 + self.bits = bits
1117 + self.group_size = group_size
1118 + self.has_zero_point = has_zero_point
1119 +
1120 + def convert_humming(
1121 + self,
1122 + tensors: dict[str, torch.Tensor],
1123 + shape_n_stacks: list[int],
1124 + shape_k_stacks: list[int],
1125 + param_dtype: torch.dtype,
1126 + num_experts: int | None = None,
1127 + ) -> tuple[Any, dict[str, torch.Tensor]]:
1128 + del shape_n_stacks, shape_k_stacks, num_experts
1129 + from vllm.utils.humming import HummingWeightSchema, dtypes
1130 +
1131 + output = _convert_moe_wna16_humming_tensors(
1132 + tensors, has_zero_point=self.has_zero_point
1133 + )
1134 + output["weight_scale"] = output["weight_scale"].to(param_dtype)
1135 + schema = HummingWeightSchema(
1136 + b_dtype=dtypes.DataType.from_str(f"uint{self.bits}"),
1137 + weight_scale_group_size=self.group_size,
1138 + has_zero_point=self.has_zero_point,
1139 + )
1140 + return schema, output
1141 +
1142 +
1143 def _unpack_and_dequant_int4_gptq(
1144 w_int32: torch.Tensor,
1145 scale: torch.Tensor,
@@ -1323,13 +1424,27 @@ def convert_to_wna16_moe_kernel_format(
1424 input_dtype: optional activation dtype, usually should be 16 bit.
1425 """
1426 if backend == WNA16MoEBackend.HUMMING:
1427 + from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
1428 from vllm.model_executor.layers.quantization.utils.humming_utils import (
1429 convert_to_humming_moe_kernel_format,
1430 )
1431
1330 - convert_to_humming_moe_kernel_format(
1331 - layer, quant_config=_humming_wna16_weight_schema(quant_config)
1332 - )
1432 + if isinstance(quant_config, MoeWNA16Config):
1433 + from vllm.utils.humming import HummingInputSchema
1434 +
1435 + convert_to_humming_moe_kernel_format(
1436 + layer,
1437 + weight_schema=_MoeWNA16HummingWeightSchema(
1438 + bits=quant_config.weight_bits,
1439 + group_size=layer.group_size,
1440 + has_zero_point=quant_config.has_zp,
1441 + ),
1442 + input_schema=HummingInputSchema(),
1443 + )
1444 + else:
1445 + convert_to_humming_moe_kernel_format(
1446 + layer, quant_config=_humming_wna16_weight_schema(quant_config)
1447 + )
1448 return None
1449
1450 if backend in (
@@ -1344,16 +1459,32 @@ def convert_to_wna16_moe_kernel_format(
1459 )
1460
1461 if isinstance(quant_config, AutoAWQConfig):
1347 - if w13_qzeros is None or w2_qzeros is None:
1348 - raise ValueError("AWQ Marlin MoE requires zero-point tensors.")
1349 -
1350 - weight_bits = quant_config.weight_bits
1462 + num_bits = quant_config.weight_bits
1463 pack_factor = quant_config.pack_factor
1464 group_size = quant_config.group_size
1465 + elif isinstance(quant_config, AutoGPTQConfig):
1466 + num_bits = quant_config.quant_type.size_bits
1467 + pack_factor = quant_config.pack_factor
1468 + group_size = quant_config.group_size
1469 + actorder = "group" if quant_config.desc_act else None
1470 + elif isinstance(quant_config, QuantizationArgs):
1471 + num_bits = quant_config.num_bits
1472 + pack_factor = 32 // quant_config.num_bits
1473 + group_size = quant_config.group_size
1474 + actorder = quant_config.actorder
1475 + else:
1476 + raise TypeError(
1477 + "Marlin WNA16 MoE backend requires AutoGPTQConfig, AutoAWQConfig or "
1478 + f"QuantizationArgs, got {type(quant_config).__name__}."
1479 + )
1480 +
1481 + if isinstance(quant_config, AutoAWQConfig):
1482 + if w13_qzeros is None or w2_qzeros is None:
1483 + raise ValueError("AWQ Marlin MoE requires zero-point tensors.")
1484
1485 return _process_awq_weights_marlin(
1486 layer,
1356 - weight_bits,
1487 + num_bits,
1488 pack_factor,
1489 group_size,
1490 input_dtype,
@@ -1366,41 +1497,28 @@ def convert_to_wna16_moe_kernel_format(
1497 w13_bias,
1498 w2_bias,
1499 )
1369 - elif isinstance(quant_config, AutoGPTQConfig):
1370 - num_bits = quant_config.quant_type.size_bits
1371 - pack_factor = quant_config.pack_factor
1372 - group_size = quant_config.group_size
1373 - actorder = "group" if quant_config.desc_act else None
1374 - elif isinstance(quant_config, QuantizationArgs):
1375 - num_bits = quant_config.num_bits
1376 - pack_factor = 32 // quant_config.num_bits
1377 - group_size = quant_config.group_size
1378 - actorder = quant_config.actorder
1500 else:
1380 - raise TypeError(
1381 - "Marlin WNA16 MoE backend requires AutoAWQConfig, AutoGPTQConfig or "
1382 - f"QuantizationArgs, got {type(quant_config).__name__}."
1501 + if w13_g_idx is None or w2_g_idx is None:
1502 + raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")
1503 +
1504 + return _process_weights_marlin(
1505 + layer,
1506 + input_dtype,
1507 + num_bits,
1508 + pack_factor,
1509 + group_size,
1510 + actorder,
1511 + w13,
1512 + w2,
1513 + w13_scale,
1514 + w2_scale,
1515 + w13_g_idx,
1516 + w2_g_idx,
1517 + w13_qzeros,
1518 + w2_qzeros,
1519 + w13_bias,
1520 + w2_bias,
1521 )
1384 - if w13_g_idx is None or w2_g_idx is None:
1385 - raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")
1386 - return _process_weights_marlin(
1387 - layer,
1388 - input_dtype,
1389 - num_bits,
1390 - pack_factor,
1391 - group_size,
1392 - actorder,
1393 - w13,
1394 - w2,
1395 - w13_scale,
1396 - w2_scale,
1397 - w13_g_idx,
1398 - w2_g_idx,
1399 - w13_qzeros,
1400 - w2_qzeros,
1401 - w13_bias,
1402 - w2_bias,
1403 - )
1522 elif backend == WNA16MoEBackend.CPU:
1523 return _process_weights_cpu(
1524 quant_config,
@@ -1482,5 +1600,47 @@ def convert_to_wna16_moe_kernel_format(
1600 w13_qzeros,
1601 w2_qzeros,
1602 )
1603 + elif backend == WNA16MoEBackend.TRITON:
1604 + # Two possible input layouts depending on the quantization source:
1605 + #
1606 + # MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first
1607 + # → just view as uint8 (no-op)
1608 + #
1609 + # AutoGPTQ/compressed-tensors (int32, K-first):
1610 + # (E, K // pack32, N_out)
1611 + # → transpose to N-first, then view as uint8 to get
1612 + # (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s]
1613 + # Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs)
1614 + from vllm.model_executor.layers.quantization.auto_gptq import (
1615 + AutoGPTQConfig,
1616 + )
1617 +
1618 + if isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)):
1619 + # These integrations build in K-first format even when the Triton
1620 + # backend is selected. Transpose to N-first first.
1621 + w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8)
1622 + w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8)
1623 + w13_scale = w13_scale.transpose(1, 2).contiguous()
1624 + w2_scale = w2_scale.transpose(1, 2).contiguous()
1625 + else:
1626 + # MoeWNA16 uses N-first uint8 weights and scales.
1627 + w13_uint8 = w13.view(torch.uint8)
1628 + w2_uint8 = w2.view(torch.uint8)
1629 + return (
1630 + w13_uint8,
1631 + w2_uint8,
1632 + w13_scale,
1633 + w2_scale,
1634 + None,
1635 + None,
1636 + None,
1637 + None,
1638 + w13_qzeros,
1639 + w2_qzeros,
1640 + None,
1641 + None,
1642 + w13_bias,
1643 + w2_bias,
1644 + )
1645 else:
1646 raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}")
vllm/model_executor/layers/fused_moe/routed_experts.py
+20
@@ -114,6 +114,7 @@ class RoutedExperts(PluggableLayer):
114 self.e_score_correction_bias = e_score_correction_bias
115 self.apply_router_weight_on_input = apply_router_weight_on_input
116 # End random parameters
117 + self._loaded_expert_biases: set[str] = set()
118
119 self.quant_method = self._get_quant_method(
120 self.layer_name,
@@ -691,6 +692,25 @@ class RoutedExperts(PluggableLayer):
692
693 expert_data = param.data if full_load else param.data[expert_id]
694
695 + if "bias" in weight_name:
696 + self._loaded_expert_biases.add(weight_name.rsplit(".", 1)[-1])
697 + if shard_id == "w2":
698 + expert_data = self._narrow_expert_data_for_padding(
699 + expert_data,
700 + loaded_weight,
701 + hidden_dim=0,
702 + )
703 + expert_data.copy_(loaded_weight)
704 + else:
705 + self._load_w13(
706 + shard_id=shard_id,
707 + shard_dim=0,
708 + loaded_weight=loaded_weight,
709 + expert_data=expert_data,
710 + tp_rank=self.moe_config.tp_rank,
711 + )
712 + return True if return_success else None
713 +
714 # Case input scale: input_scale loading is only supported for fp8
715 if "input_scale" in weight_name:
716 # this is needed for compressed-tensors only
vllm/model_executor/layers/quantization/auto_awq.py
+3
@@ -560,6 +560,9 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
560 self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
561 moe,
562 kInt4Static,
563 + quant_config=self.quant_config,
564 + may_have_zp=self.quant_config.zero_point,
565 + may_have_bias=True,
566 )
567
568 def create_weights(
vllm/model_executor/layers/quantization/auto_gptq.py
+60 -46
@@ -489,6 +489,9 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
489 self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
490 moe,
491 weight_key,
492 + quant_config=self.quant_config,
493 + may_have_zp=True,
494 + may_have_bias=True,
495 )
496
497 def create_weights(
@@ -577,25 +580,25 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
580 set_weight_attrs(w2_scales, extra_weight_attrs)
581 # don't shard the w2 scales when running act order
582 set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act})
580 - # up_proj scales
583 + # up_proj zero points
584 w13_qzeros = torch.nn.Parameter(
585 torch.empty(
586 num_experts,
587 scales_size13,
588 2 * intermediate_size_per_partition // self.quant_config.pack_factor,
586 - dtype=params_dtype,
589 + dtype=torch.int32,
590 ),
591 requires_grad=False,
592 )
593 layer.register_parameter("w13_qzeros", w13_qzeros)
594 set_weight_attrs(w13_qzeros, extra_weight_attrs)
592 - # down_proj scales
595 + # down_proj zero points
596 w2_qzeros = torch.nn.Parameter(
597 torch.empty(
598 num_experts,
599 scales_size2,
600 hidden_size // self.quant_config.pack_factor,
598 - dtype=params_dtype,
601 + dtype=torch.int32,
602 ),
603 requires_grad=False,
604 )
@@ -644,6 +647,26 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
647 layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
648 set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
649
650 + # Some GPTQ checkpoints contain expert biases even when the model
651 + # architecture does not declare them. Zero initialization keeps
652 + # checkpoints without biases equivalent to the bias-free path.
653 + w13_bias = torch.nn.Parameter(
654 + torch.zeros(
655 + num_experts,
656 + 2 * intermediate_size_per_partition,
657 + dtype=params_dtype,
658 + ),
659 + requires_grad=False,
660 + )
661 + layer.register_parameter("w13_bias", w13_bias)
662 + set_weight_attrs(w13_bias, extra_weight_attrs)
663 + w2_bias = torch.nn.Parameter(
664 + torch.zeros(num_experts, hidden_size, dtype=params_dtype),
665 + requires_grad=False,
666 + )
667 + layer.register_parameter("w2_bias", w2_bias)
668 + set_weight_attrs(w2_bias, extra_weight_attrs)
669 +
670 if self.experts_cls is not None and issubclass(
671 self.experts_cls, FusedMoEExpertsModular
672 ):
@@ -651,12 +674,31 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
674 layer.workspace = marlin_make_workspace_new(device, 4)
675
676 def process_weights_after_loading(self, layer: RoutedExperts) -> None:
677 + def replace_or_register(name: str, val: torch.Tensor | None):
678 + if val is None:
679 + return
680 +
681 + if hasattr(layer, name):
682 + replace_parameter(layer, name, val)
683 + else:
684 + layer.register_parameter(
685 + name, torch.nn.Parameter(val, requires_grad=False)
686 + )
687 +
688 is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
689
656 - if is_a_8bit:
657 - assert self.quant_config.quant_type.size_bits == 8, (
658 - "W8A8-INT8 is not supported by marlin kernel."
659 - )
690 + assert not is_a_8bit or self.quant_config.quant_type.size_bits == 8, (
691 + "W8A8-INT8 is not supported by marlin kernel."
692 + )
693 +
694 + w13_bias = getattr(layer, "w13_bias", None)
695 + if "w13_bias" not in layer._loaded_expert_biases:
696 + layer.register_parameter("w13_bias", None)
697 + w13_bias = None
698 + w2_bias = getattr(layer, "w2_bias", None)
699 + if "w2_bias" not in layer._loaded_expert_biases:
700 + layer.register_parameter("w2_bias", None)
701 + w2_bias = None
702
703 converted = convert_to_wna16_moe_kernel_format(
704 backend=self.wna16_moe_backend,
@@ -669,8 +711,10 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
711 w2_scale=layer.w2_scales,
712 w13_g_idx=layer.w13_g_idx,
713 w2_g_idx=layer.w2_g_idx,
672 - w13_bias=getattr(layer, "w13_bias", None),
673 - w2_bias=getattr(layer, "w2_bias", None),
714 + w13_bias=w13_bias,
715 + w2_bias=w2_bias,
716 + w13_qzeros=getattr(layer, "w13_qzeros", None),
717 + w2_qzeros=getattr(layer, "w2_qzeros", None),
718 )
719
720 if converted is None:
@@ -703,42 +747,12 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
747 replace_parameter(layer, "w2_g_idx", w2_g_idx)
748 replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
749 replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
706 - if w13_qzeros is not None:
707 - replace_parameter(layer, "w13_qzeros", w13_qzeros)
708 - if w2_qzeros is not None:
709 - replace_parameter(layer, "w2_qzeros", w2_qzeros)
710 - if w13_input_global_scale is not None:
711 - if hasattr(layer, "w13_input_global_scale"):
712 - replace_parameter(
713 - layer, "w13_input_global_scale", w13_input_global_scale
714 - )
715 - else:
716 - layer.register_parameter(
717 - "w13_input_global_scale",
718 - torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
719 - )
720 - if w2_input_global_scale is not None:
721 - if hasattr(layer, "w2_input_global_scale"):
722 - replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale)
723 - else:
724 - layer.register_parameter(
725 - "w2_input_global_scale",
726 - torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
727 - )
728 - if w13_bias is not None:
729 - if hasattr(layer, "w13_bias"):
730 - replace_parameter(layer, "w13_bias", w13_bias)
731 - else:
732 - layer.register_parameter(
733 - "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False)
734 - )
735 - if w2_bias is not None:
736 - if hasattr(layer, "w2_bias"):
737 - replace_parameter(layer, "w2_bias", w2_bias)
738 - else:
739 - layer.register_parameter(
740 - "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False)
741 - )
750 + replace_or_register("w13_input_global_scale", w13_input_global_scale)
751 + replace_or_register("w2_input_global_scale", w2_input_global_scale)
752 + replace_or_register("w13_bias", w13_bias)
753 + replace_or_register("w2_bias", w2_bias)
754 + replace_or_register("w13_qzeros", w13_qzeros)
755 + replace_or_register("w2_qzeros", w2_qzeros)
756
757 # The modular kernel reads w13_weight/w2_weight; marlin keeps *_qweight.
758 layer.w13_weight = layer.w13_qweight
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+14 -7
@@ -92,7 +92,15 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
92 self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
93 config=self.moe,
94 weight_key=weight_key,
95 + quant_config=self.weight_quant,
96 + may_have_zp=not self.symmetric,
97 + may_have_bias=False,
98 )
99 + self.is_marlin = self.wna16_backend in [
100 + WNA16MoEBackend.MARLIN,
101 + WNA16MoEBackend.BATCHED_MARLIN,
102 + ]
103 + self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
104
105 def get_weight_shape(
106 self,
@@ -118,7 +126,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
126 "num_groups_w2 must be provided for weight scales/zero_points"
127 )
128 w13_num_shards = 2 if self.moe.is_act_and_mul else 1
121 - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
129 shape_map = {
130 "w13_weight": {
131 "Flashinfer": (
@@ -177,7 +184,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
184 ),
185 },
186 }
180 - backend_key = "Flashinfer" if is_flashinfer else "Marlin"
187 + backend_key = "Marlin" if self.is_transposed else "Flashinfer"
188 return shape_map[weight_name][backend_key]
189
190 @staticmethod
@@ -217,9 +224,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
224 # Will transpose the loaded weight along the
225 # intermediate and hidden dim sizes. Will
226 # shard for TP along the transposed dims
220 - is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
227 extra_weight_attrs.update(
222 - {"is_transposed": is_transposed, "quant_method": self.strategy}
228 + {"is_transposed": self.is_transposed, "quant_method": self.strategy}
229 )
230
231 w13_weight = torch.nn.Parameter(
@@ -415,7 +421,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
421
422 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
423 # Process weights using the shared oracle infrastructure
418 - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
424 converted = convert_to_wna16_moe_kernel_format(
425 backend=self.wna16_backend,
426 layer=layer,
@@ -430,6 +435,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
435 w13_qzeros=getattr(layer, "w13_weight_zero_point", None),
436 w2_qzeros=getattr(layer, "w2_weight_zero_point", None),
437 )
438 +
439 if converted is None:
440 # In-place backends (e.g. Humming) are not wired through this
441 # marlin-only method; fail clearly rather than unpacking None.
@@ -437,6 +443,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
443 f"{type(self).__name__} does not support the "
444 f"{self.wna16_backend.value} MoE backend."
445 )
446 +
447 (
448 w13_qweight,
449 w2_qweight,
@@ -466,7 +473,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
473 replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
474
475 # Marlin-specific parameters (not needed for Flashinfer)
469 - if not is_flashinfer:
476 + if self.is_marlin:
477 if w13_g_idx_processed is not None:
478 replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
479 if w2_g_idx_processed is not None:
@@ -510,7 +517,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
517
518 # Add Marlin-specific arguments
519 marlin_args: dict[str, Any] = {}
513 - if not is_flashinfer:
520 + if self.is_marlin:
521 marlin_args = {
522 "w13_g_idx": layer.w13_weight_g_idx,
523 "w2_g_idx": layer.w2_weight_g_idx,
vllm/model_executor/layers/quantization/moe_wna16.py
+166 -19
@@ -15,8 +15,13 @@ from vllm.model_executor.layers.fused_moe import (
15 )
16 from vllm.model_executor.layers.fused_moe.config import (
17 FusedMoEQuantConfig,
18 - int4_w4a16_moe_quant_config,
19 - int8_w8a16_moe_quant_config,
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,
25 )
26 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
27 UnquantizedFusedMoEMethod,
@@ -27,7 +32,15 @@ from vllm.model_executor.layers.quantization.base_config import (
32 QuantizationConfig,
33 QuantizeMethodBase,
34 )
30 -from vllm.model_executor.utils import set_weight_attrs
35 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
36 + INT4_DTYPE,
37 + INT8_DTYPE,
38 + QuantKey,
39 + kInt4Static32GroupScale,
40 + kInt4StaticGroupScale,
41 + kInt8StaticGroupScale,
42 +)
43 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
44 from vllm.platforms import current_platform
45
46
@@ -199,6 +212,34 @@ class MoeWNA16Method(FusedMoEMethodBase):
212 super().__init__(moe)
213 self.quant_config = quant_config
214
215 + num_bits = self.quant_config.weight_bits
216 + group_size = self.quant_config.group_size
217 +
218 + if num_bits == 4:
219 + quant_type = INT4_DTYPE
220 + if group_size == 32:
221 + scale = kInt4Static32GroupScale
222 + else:
223 + scale = kInt4StaticGroupScale
224 + elif num_bits == 8:
225 + assert group_size == -1
226 + quant_type = INT8_DTYPE
227 + scale = kInt8StaticGroupScale
228 + else:
229 + raise ValueError("MoeWNA16Method only supports int4 and int8 now.")
230 +
231 + weight_key = QuantKey(quant_type, scale)
232 +
233 + # Select WNA16 MoE backend via oracle.
234 + # handle ZP?
235 + self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
236 + config=self.moe,
237 + weight_key=weight_key,
238 + quant_config=self.quant_config,
239 + may_have_zp=self.quant_config.has_zp,
240 + may_have_bias=False,
241 + )
242 +
243 def create_weights(
244 self,
245 layer: RoutedExperts,
@@ -322,23 +363,104 @@ class MoeWNA16Method(FusedMoEMethodBase):
363 def get_fused_moe_quant_config(
364 self, layer: RoutedExperts
365 ) -> FusedMoEQuantConfig | None:
325 - weight_bits = self.quant_config.weight_bits
326 - has_zp = self.quant_config.has_zp
327 - assert weight_bits == 4 or weight_bits == 8
328 - config_builder = (
329 - int4_w4a16_moe_quant_config
330 - if weight_bits == 4
331 - else int8_w8a16_moe_quant_config
332 - )
366 + if self.wna16_backend == WNA16MoEBackend.HUMMING:
367 + from vllm.model_executor.layers.quantization.utils.humming_utils import (
368 + get_humming_moe_quant_config,
369 + )
370 +
371 + return get_humming_moe_quant_config(layer)
372
334 - return config_builder(
373 + has_zp = self.quant_config.has_zp
374 + return make_wna16_moe_quant_config(
375 w1_scale=layer.w13_scales,
376 w2_scale=layer.w2_scales,
377 w1_zp=layer.w13_qzeros if has_zp else None,
378 w2_zp=layer.w2_qzeros if has_zp else None,
339 - block_shape=[0, layer.group_size],
379 + group_size=layer.group_size,
380 + num_bits=self.quant_config.weight_bits,
381 )
382
383 + def _setup_kernel(self, layer: RoutedExperts):
384 + assert self.experts_cls is not None
385 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
386 + assert self.moe_quant_config is not None
387 + self.moe_kernel = make_wna16_moe_kernel(
388 + moe_quant_config=self.moe_quant_config,
389 + moe_config=self.moe,
390 + experts_cls=self.experts_cls,
391 + backend=self.wna16_backend,
392 + layer=layer,
393 + routing_tables=layer._expert_routing_tables(),
394 + )
395 +
396 + def process_weights_after_loading(self, layer: RoutedExperts) -> None:
397 + has_zp = self.quant_config.has_zp
398 + converted = convert_to_wna16_moe_kernel_format(
399 + backend=self.wna16_backend,
400 + layer=layer,
401 + quant_config=self.quant_config,
402 + input_dtype=None,
403 + w13=layer.w13_qweight,
404 + w2=layer.w2_qweight,
405 + w13_scale=layer.w13_scales,
406 + w2_scale=layer.w2_scales,
407 + w13_g_idx=None,
408 + w2_g_idx=None,
409 + w13_qzeros=layer.w13_qzeros if has_zp else None,
410 + w2_qzeros=layer.w2_qzeros if has_zp else None,
411 + )
412 +
413 + if converted is None:
414 + # Backend rewrote the layer's params in place (e.g. Humming).
415 + self._setup_kernel(layer)
416 + return
417 +
418 + (
419 + w13_qweight,
420 + w2_qweight,
421 + w13_scales,
422 + w2_scales,
423 + _,
424 + _,
425 + _,
426 + _,
427 + w13_qzeros,
428 + w2_qzeros,
429 + w13_input_global_scale,
430 + w2_input_global_scale,
431 + _, # w13_bias
432 + _, # w2_bias
433 + ) = converted
434 +
435 + # Replace common parameters
436 + replace_parameter(layer, "w13_qweight", w13_qweight)
437 + replace_parameter(layer, "w2_qweight", w2_qweight)
438 + replace_parameter(layer, "w13_scales", w13_scales)
439 + replace_parameter(layer, "w2_scales", w2_scales)
440 + layer.w13_weight = layer.w13_qweight
441 + layer.w2_weight = layer.w2_qweight
442 +
443 + if has_zp:
444 + assert w13_qzeros is not None and w2_qzeros is not None
445 + replace_parameter(layer, "w13_qzeros", w13_qzeros)
446 + replace_parameter(layer, "w2_qzeros", w2_qzeros)
447 +
448 + # Marlin-specific parameters (not needed for Flashinfer)
449 + if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM:
450 + # Register input global scales if present
451 + if w13_input_global_scale is not None:
452 + layer.register_parameter(
453 + "w13_input_global_scale",
454 + torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
455 + )
456 + if w2_input_global_scale is not None:
457 + layer.register_parameter(
458 + "w2_input_global_scale",
459 + torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
460 + )
461 +
462 + self._setup_kernel(layer)
463 +
464 def apply(
465 self,
466 layer: RoutedExperts,
@@ -348,19 +470,44 @@ class MoeWNA16Method(FusedMoEMethodBase):
470 shared_experts: SharedExperts | None,
471 shared_experts_input: torch.Tensor | None,
472 ) -> torch.Tensor:
351 - from vllm.model_executor.layers.fused_moe import fused_experts
352 -
353 - return fused_experts(
473 + assert not self.is_monolithic
474 + assert self.moe_kernel is not None
475 + return self.moe_kernel.apply(
476 x,
355 - layer.w13_qweight,
356 - layer.w2_qweight,
477 + layer.w13_weight,
478 + layer.w2_weight,
479 topk_weights=topk_weights,
480 topk_ids=topk_ids,
481 activation=layer.activation,
482 + global_num_experts=layer.global_num_experts,
483 + expert_map=layer.expert_map,
484 apply_router_weight_on_input=layer.apply_router_weight_on_input,
485 + shared_experts=shared_experts,
486 + shared_experts_input=shared_experts_input,
487 + )
488 +
489 + def apply_monolithic(
490 + self,
491 + layer: RoutedExperts,
492 + x: torch.Tensor,
493 + router_logits: torch.Tensor,
494 + input_ids: torch.Tensor | None = None,
495 + ) -> torch.Tensor:
496 + assert self.is_monolithic
497 + assert self.moe_kernel is not None
498 + return self.moe_kernel.apply_monolithic(
499 + x,
500 + layer.w13_weight,
501 + layer.w2_weight,
502 + router_logits,
503 + activation=layer.activation,
504 global_num_experts=layer.global_num_experts,
505 expert_map=layer.expert_map,
363 - quant_config=self.moe_quant_config,
506 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
507 + num_expert_group=layer.num_expert_group,
508 + topk_group=layer.topk_group,
509 + e_score_correction_bias=layer.e_score_correction_bias,
510 + routed_scaling_factor=layer.routed_scaling_factor,
511 )
512
513 @staticmethod
vllm/model_executor/layers/quantization/utils/gptq_utils.py
+17 -1
@@ -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
6 +from typing import TYPE_CHECKING, Any
7
8 import regex as re
9 import torch
@@ -69,6 +69,20 @@ 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 +
86 def is_layer_gptq_quantized(
87 prefix: str,
88 quantized_layers: list[str],
@@ -83,6 +97,8 @@ def is_layer_gptq_quantized(
97
98 proj_name = prefix.split(".")[-1]
99
100 + quantized_layers = flatten_list(quantized_layers)
101 +
102 # Fused layers like gate_up_proj or qkv_proj will not be fused
103 # in the safetensors checkpoint. So, we convert the name
104 # from the fused version to unfused + check to make sure that