[CT] Support Humming for WNA16 MoE (#48918)
Signed-off-by: yiliu30 <yi4.liu@intel.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Yi Liu committed
Aug 20, 2026 at 00:31 UTC
54dd98be2830bbd7b70fd60a0e917b7d0df5adab
7 files changed
+253
-24
tests/quantization/test_compressed_tensors.py
+44
@@ -43,8 +43,14 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
43
find_matched_target,
44
)
45
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
46
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
47
+ GroupShape,
48
+ QuantKey,
49
+ ScaleDesc,
50
+)
51
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
52
from vllm.platforms import current_platform
53
+from vllm.scalar_type import scalar_types
54
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
55
56
# AITER only supports per-channel-per-channel INT8 gemm
@@ -971,6 +977,44 @@ def test_wna16_moe_w2_scale_sharding(actorder, group_size, part, full, expected)
977
assert result == expected
978
979
980
+@pytest.mark.parametrize("num_bits", range(2, 9))
981
+def test_humming_supports_compressed_tensors_wna16_quant_key(num_bits):
982
+ from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
983
+ HummingExpertsBase,
984
+ )
985
+ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa: E501
986
+ WNA16_SUPPORTED_TYPES_MAP,
987
+ )
988
+
989
+ weight_key = QuantKey(
990
+ dtype=WNA16_SUPPORTED_TYPES_MAP[num_bits],
991
+ scale=ScaleDesc(
992
+ dtype=torch.float16,
993
+ static=True,
994
+ group_shape=GroupShape(row=1, col=128),
995
+ ),
996
+ symmetric=True,
997
+ )
998
+
999
+ assert HummingExpertsBase._supports_quant_scheme(weight_key, None)
1000
+
1001
+
1002
+def test_quant_key_str_supports_scalar_type_dtypes():
1003
+ quant_key = QuantKey(
1004
+ dtype=scalar_types.uint2b2,
1005
+ scale=ScaleDesc(
1006
+ dtype=torch.float16,
1007
+ static=True,
1008
+ group_shape=GroupShape(row=1, col=128),
1009
+ ),
1010
+ symmetric=True,
1011
+ )
1012
+
1013
+ assert str(quant_key) == (
1014
+ "QuantKey(uint2b2,scale(f16,static,GroupShape(row=1, col=128)),symmetric)"
1015
+ )
1016
+
1017
+
1018
@pytest.mark.skipif(
1019
not current_platform.is_cuda() or not current_platform.has_device_capability(80),
1020
reason="MXFP4 requires ampere or newer",
tests/quantization/test_moe_wna16.py
+110
@@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.moe_wna16 import (
26
MoeWNA16Config,
27
MoeWNA16Method,
28
)
29
+from vllm.platforms import current_platform
30
31
32
def test_map_wna16_backend_supports_triton():
@@ -203,3 +204,112 @@ def test_moe_wna16_uses_humming_quant_config(monkeypatch):
204
)
205
206
assert method.get_fused_moe_quant_config(layer) is quant_config
207
+
208
+
209
+@pytest.mark.skipif(
210
+ not current_platform.is_cuda(),
211
+ reason="Compressed-tensors Humming WNA16 MoE requires CUDA",
212
+)
213
+@pytest.mark.parametrize("num_bits", [3, 5, 6, 7])
214
+def test_compressed_tensors_wna16_moe_create_weights_uses_ceil_packed_shapes(
215
+ num_bits,
216
+):
217
+ pytest.importorskip("humming")
218
+
219
+ from tests.kernels.moe.utils import make_dummy_moe_config
220
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
221
+ CompressedTensorsWNA16MoEMethod,
222
+ )
223
+
224
+ quant_args = QuantizationArgs(
225
+ num_bits=num_bits,
226
+ type=QuantizationType.INT,
227
+ strategy=QuantizationStrategy.GROUP,
228
+ symmetric=True,
229
+ dynamic=False,
230
+ group_size=128,
231
+ )
232
+ moe_config = make_dummy_moe_config(
233
+ num_experts=2,
234
+ hidden_dim=256,
235
+ intermediate_size=512,
236
+ )
237
+ moe_config.moe_backend = "humming"
238
+ method = CompressedTensorsWNA16MoEMethod(quant_args, None, moe_config)
239
+ layer = torch.nn.Module()
240
+
241
+ method.create_weights(
242
+ layer,
243
+ num_experts=2,
244
+ hidden_size=256,
245
+ intermediate_size_per_partition=512,
246
+ params_dtype=torch.float16,
247
+ intermediate_size_full=512,
248
+ )
249
+
250
+ packed_hidden = (256 * num_bits + 31) // 32
251
+ packed_intermediate = (512 * num_bits + 31) // 32
252
+ assert method.wna16_backend == WNA16MoEBackend.HUMMING
253
+ assert layer.w13_weight_packed.shape == (2, 1024, packed_hidden)
254
+ assert layer.w2_weight_packed.shape == (2, 256, packed_intermediate)
255
+ assert layer.w13_weight_scale.shape == (2, 1024, 2)
256
+ assert layer.w2_weight_scale.shape == (2, 256, 4)
257
+ assert layer.w13_weight_packed.dtype is torch.int32
258
+ assert layer.w2_weight_scale.dtype is torch.float16
259
+
260
+
261
+@pytest.mark.skipif(
262
+ not current_platform.is_cuda(),
263
+ reason="Compressed-tensors Humming WNA16 MoE requires CUDA",
264
+)
265
+def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel():
266
+ pytest.importorskip("humming")
267
+
268
+ from tests.kernels.moe.utils import make_dummy_moe_config
269
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
270
+ CompressedTensorsWNA16MoEMethod,
271
+ )
272
+
273
+ quant_args = QuantizationArgs(
274
+ num_bits=3,
275
+ type=QuantizationType.INT,
276
+ strategy=QuantizationStrategy.GROUP,
277
+ symmetric=True,
278
+ dynamic=False,
279
+ group_size=128,
280
+ )
281
+ moe_config = make_dummy_moe_config(
282
+ num_experts=2,
283
+ hidden_dim=256,
284
+ intermediate_size=512,
285
+ )
286
+ moe_config.moe_backend = "humming"
287
+ method = CompressedTensorsWNA16MoEMethod(quant_args, None, moe_config)
288
+ layer = torch.nn.Module()
289
+ layer.moe_config = moe_config
290
+ layer.params_dtype = torch.bfloat16
291
+ layer.layer_name = "test.humming_moe"
292
+ layer._expert_routing_tables = lambda: (None, None, None)
293
+
294
+ method.create_weights(
295
+ layer,
296
+ num_experts=2,
297
+ hidden_size=256,
298
+ intermediate_size_per_partition=512,
299
+ intermediate_size_full=512,
300
+ params_dtype=torch.bfloat16,
301
+ )
302
+ layer.cuda()
303
+ for parameter in layer.parameters():
304
+ parameter.data.zero_()
305
+
306
+ method.process_weights_after_loading(layer)
307
+
308
+ assert method.wna16_backend == WNA16MoEBackend.HUMMING
309
+ assert method.moe_kernel is not None
310
+ assert set(layer.weight_schemas) == {"w13", "w2"}
311
+ assert set(layer.humming_configs) == {"w13", "w2"}
312
+ assert not hasattr(layer, "w13_weight_packed")
313
+ assert not hasattr(layer, "w2_weight_packed")
314
+ assert layer.w13_weight.dtype is torch.int32
315
+ assert layer.w2_weight.dtype is torch.int32
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+31
-1
@@ -40,6 +40,9 @@ from vllm.model_executor.layers.fused_moe.utils import (
40
swiglu_limit_func,
41
)
42
from vllm.model_executor.layers.quantization.utils.quant_utils import (
43
+ INT4_DTYPE,
44
+ INT8_DTYPE,
45
+ GroupShape,
46
QuantKey,
47
kFp8Dynamic128Sym,
48
kFp8DynamicTokenSym,
@@ -58,6 +61,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
61
kNvfp4Static,
62
)
63
from vllm.platforms import current_platform
64
+from vllm.scalar_type import ScalarType
65
from vllm.utils.import_utils import has_humming
66
from vllm.v1.worker.workspace import current_workspace_manager
67
@@ -76,6 +80,29 @@ if TYPE_CHECKING:
80
logger = init_logger(__name__)
81
82
83
+def _is_supported_wna16_weight_key(weight_key: QuantKey | None) -> bool:
84
+ if weight_key is None or weight_key.scale2 is not None:
85
+ return False
86
+
87
+ group_shape = weight_key.scale.group_shape
88
+ if not (
89
+ group_shape == GroupShape.PER_CHANNEL
90
+ or (group_shape.row == 1 and group_shape.col > 0)
91
+ ):
92
+ return False
93
+
94
+ dtype = weight_key.dtype
95
+ if dtype in (INT4_DTYPE, INT8_DTYPE, torch.uint8):
96
+ return True
97
+
98
+ return (
99
+ isinstance(dtype, ScalarType)
100
+ and dtype.is_integer()
101
+ and not dtype.is_signed()
102
+ and 2 <= dtype.size_bits <= 8
103
+ )
104
+
105
+
106
def get_humming_moe_gemm_type() -> str:
107
env_gemm_type: str | None = envs.VLLM_HUMMING_MOE_GEMM_TYPE
108
gemm_type = "indexed"
@@ -267,7 +294,10 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
294
# mxfp8 (compressed-tensors / modelopt / online)
295
(kMxfp8Static, kMxfp8Dynamic),
296
]
270
- return (weight_key, activation_key) in SUPPORTED_W_A
297
+ return (weight_key, activation_key) in SUPPORTED_W_A or (
298
+ activation_key in (None, kFp8DynamicTokenSym)
299
+ and _is_supported_wna16_weight_key(weight_key)
300
+ )
301
302
@property
303
def expects_unquantized_inputs(self) -> bool:
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+14
-2
@@ -1104,9 +1104,21 @@ def _humming_wna16_weight_schema(
1104
"desc_act": quant_config.desc_act,
1105
"sym": quant_config.is_sym,
1106
}
1107
+ if isinstance(quant_config, QuantizationArgs):
1108
+ quant_type = getattr(quant_config.type, "value", quant_config.type)
1109
+ quant_strategy = getattr(quant_config.strategy, "value", quant_config.strategy)
1110
+ return {
1111
+ "quant_method": "compressed-tensors",
1112
+ "format": "pack-quantized",
1113
+ "type": str(quant_type),
1114
+ "num_bits": quant_config.num_bits,
1115
+ "strategy": str(quant_strategy),
1116
+ "group_size": quant_config.group_size,
1117
+ "symmetric": quant_config.symmetric,
1118
+ }
1119
raise TypeError(
1108
- "Humming WNA16 checkpoint schema requires AutoAWQConfig or "
1109
- "AutoGPTQConfig, "
1120
+ "Humming WNA16 checkpoint schema requires AutoAWQConfig, "
1121
+ "AutoGPTQConfig or QuantizationArgs, "
1122
f"got {type(quant_config).__name__}."
1123
)
1124
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+42
-14
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
+from fractions import Fraction
6
from typing import Any
7
8
import torch
@@ -40,7 +42,9 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
42
marlin_make_workspace_new,
43
)
44
from vllm.model_executor.layers.quantization.utils.quant_utils import (
45
+ GroupShape,
46
QuantKey,
47
+ ScaleDesc,
48
kInt4Static32GroupScale,
49
kInt4StaticGroupScale,
50
kInt8StaticGroupScale,
@@ -64,7 +68,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
68
# Extract properties from weight_quant
69
self.symmetric = weight_quant.symmetric
70
self.num_bits = weight_quant.num_bits
67
- self.packed_factor = 32 // weight_quant.num_bits
71
+ self.packed_factor = Fraction(32, weight_quant.num_bits)
72
self.strategy = weight_quant.strategy
73
self.group_size = weight_quant.group_size
74
self.actorder = weight_quant.actorder
@@ -84,8 +88,14 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
88
elif self.num_bits == 8:
89
scale = kInt8StaticGroupScale
90
else:
87
- raise ValueError(
88
- "CompressedTensorsWNA16MoEMethod only supports int4 and int8 now."
91
+ scale = ScaleDesc(
92
+ dtype=torch.float16,
93
+ static=True,
94
+ group_shape=(
95
+ GroupShape.PER_CHANNEL
96
+ if self.group_size == -1
97
+ else GroupShape(row=1, col=self.group_size)
98
+ ),
99
)
100
101
weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric)
@@ -109,7 +119,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
119
WNA16MoEBackend.MARLIN,
120
WNA16MoEBackend.BATCHED_MARLIN,
121
]
112
- self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
122
+ self.is_transposed = self.wna16_backend not in (
123
+ WNA16MoEBackend.FLASHINFER_TRTLLM,
124
+ WNA16MoEBackend.HUMMING,
125
+ )
126
127
if self.is_marlin:
128
assert check_moe_marlin_supports_config(
@@ -125,6 +138,9 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
138
# Non-Marlin WNA16 always uses bf16/fp16 inputs
139
self.input_dtype = torch.bfloat16
140
141
+ def _packed_dim(self, dim: int) -> int:
142
+ return math.ceil(dim * self.num_bits / 32)
143
+
144
def get_weight_shape(
145
self,
146
weight_name: str,
@@ -149,16 +165,16 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
165
"num_groups_w2 must be provided for weight scales/zero_points"
166
)
167
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
152
- shape_map = {
168
+ shape_map: dict[str, dict[str, tuple[int, int | None, int | None]]] = {
169
"w13_weight": {
170
"Flashinfer": (
171
num_experts,
172
w13_num_shards * intermediate_size_per_partition,
157
- hidden_size // self.packed_factor,
173
+ self._packed_dim(hidden_size),
174
),
175
"Marlin": (
176
num_experts,
161
- hidden_size // self.packed_factor,
177
+ self._packed_dim(hidden_size),
178
w13_num_shards * intermediate_size_per_partition,
179
),
180
},
@@ -178,20 +194,18 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
194
"Marlin": (
195
num_experts,
196
num_groups_w13,
181
- w13_num_shards
182
- * intermediate_size_per_partition
183
- // self.packed_factor,
197
+ self._packed_dim(w13_num_shards * intermediate_size_per_partition),
198
),
199
},
200
"w2_weight": {
201
"Flashinfer": (
202
num_experts,
203
hidden_size,
190
- intermediate_size_per_partition // self.packed_factor,
204
+ self._packed_dim(intermediate_size_per_partition),
205
),
206
"Marlin": (
207
num_experts,
194
- intermediate_size_per_partition // self.packed_factor,
208
+ self._packed_dim(intermediate_size_per_partition),
209
hidden_size,
210
),
211
},
@@ -203,12 +217,14 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
217
"Marlin": (
218
num_experts,
219
num_groups_w2,
206
- hidden_size // self.packed_factor,
220
+ self._packed_dim(hidden_size),
221
),
222
},
223
}
224
backend_key = "Marlin" if self.is_transposed else "Flashinfer"
211
- return shape_map[weight_name][backend_key]
225
+ shape = shape_map[weight_name][backend_key]
226
+ assert shape[1] is not None and shape[2] is not None
227
+ return shape[0], shape[1], shape[2]
228
229
@staticmethod
230
def _w2_scale_sharding(
@@ -483,6 +499,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
499
moe_quant_config=self.moe_quant_config,
500
moe_config=self.moe,
501
experts_cls=self.experts_cls,
502
+ backend=self.wna16_backend,
503
routing_tables=layer._expert_routing_tables(),
504
**marlin_args,
505
)
@@ -586,6 +603,17 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
603
def get_fused_moe_quant_config(
604
self, layer: torch.nn.Module
605
) -> FusedMoEQuantConfig | None:
606
+ if self.wna16_backend == WNA16MoEBackend.HUMMING:
607
+ from vllm.model_executor.layers.quantization.utils.humming_utils import (
608
+ get_humming_moe_quant_config,
609
+ )
610
+
611
+ return get_humming_moe_quant_config(
612
+ layer,
613
+ gemm1_alpha=getattr(layer, "swiglu_alpha", None),
614
+ gemm1_beta=getattr(layer, "swiglu_beta", None),
615
+ gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
616
+ )
617
return make_wna16_moe_quant_config(
618
w1_scale=layer.w13_weight_scale,
619
w2_scale=layer.w2_weight_scale,
vllm/model_executor/layers/quantization/utils/humming_utils.py
+3
@@ -73,6 +73,9 @@ if has_humming():
73
humming_dtypes.uint8: INT8_DTYPE,
74
humming_dtypes.uint2: torch.uint8,
75
humming_dtypes.uint3: torch.uint8,
76
+ humming_dtypes.uint5: torch.uint8,
77
+ humming_dtypes.uint6: torch.uint8,
78
+ humming_dtypes.uint7: torch.uint8,
79
}
80
81
_HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = {
vllm/model_executor/layers/quantization/utils/quant_utils.py
+9
-7
@@ -25,6 +25,13 @@ INT4_DTYPE = scalar_types.uint4b8
25
INT8_DTYPE = scalar_types.uint8b128
26
27
28
+def _dtype_abbr(dtype: torch.dtype | ScalarType) -> str:
29
+ """Return a stable short name for torch and ScalarType dtypes."""
30
+ if isinstance(dtype, ScalarType):
31
+ return str(dtype)
32
+ return fx.graph.dtype_abbrs[dtype]
33
+
34
+
35
def weight_amax(
36
weight: torch.Tensor, *, dim: int | None = None, keepdim: bool = False
37
) -> torch.Tensor:
@@ -152,7 +159,7 @@ class ScaleDesc:
159
group_shape = d.get(self.group_shape, str(self.group_shape))
160
161
return (
155
- f"{fx.graph.dtype_abbrs[self.dtype]},"
162
+ f"{_dtype_abbr(self.dtype)},"
163
f"{'static' if self.static else 'dynamic'},{group_shape}"
164
)
165
@@ -178,13 +185,8 @@ class QuantKey:
185
186
def __str__(self):
187
scale2_str = f"scale2({self.scale2})," if self.scale2 else ""
181
- dtype_description = (
182
- fx.graph.dtype_abbrs[self.dtype]
183
- if isinstance(self.dtype, torch.dtype)
184
- else self.dtype
185
- )
188
return (
187
- f"QuantKey({dtype_description},"
189
+ f"QuantKey({_dtype_abbr(self.dtype)},"
190
f"scale({self.scale}),{scale2_str}"
191
f"{'a' if not self.symmetric else ''}symmetric)"
192
)