[Bugfix][Quantization] Add Humming to the W4A8 (INT4xFP8) MoE oracle (#58427)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Sep 24, 2026 at 15:37 UTC
cd06e81c41f7d1f92a5680a616ad9679232a9e7a
5 files changed
+163
-71
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+3
-1
@@ -1326,7 +1326,9 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1326
1327
@staticmethod
1328
def _supports_current_device() -> bool:
1329
- return cutlass_group_gemm_supported()
1329
+ return (
1330
+ current_platform.is_device_capability(90) and cutlass_group_gemm_supported()
1331
+ )
1332
1333
@staticmethod
1334
def _supports_no_act_and_mul() -> bool:
vllm/model_executor/layers/fused_moe/oracle/w4a8.py
+120
-36
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Oracle for W4A8 (INT4, FP8) MoE
4
from enum import Enum
5
from typing import TYPE_CHECKING
6
@@ -22,26 +23,37 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
)
24
25
if TYPE_CHECKING:
25
- from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
26
- CutlassExpertsW4A8Fp8,
27
- )
26
+ from vllm.model_executor.layers.fused_moe import RoutedExperts
27
28
logger = init_logger(__name__)
29
30
31
class W4A8MoeBackend(Enum):
32
CUTLASS = "CUTLASS"
33
+ HUMMING = "HUMMING"
34
35
36
def backend_to_kernel_cls(
37
backend: W4A8MoeBackend,
38
-) -> list[type["CutlassExpertsW4A8Fp8"]]:
38
+) -> list[type[mk.FusedMoEExperts]]:
39
if backend == W4A8MoeBackend.CUTLASS:
40
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
41
CutlassExpertsW4A8Fp8,
42
)
43
44
return [CutlassExpertsW4A8Fp8]
45
+ elif backend == W4A8MoeBackend.HUMMING:
46
+ from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
47
+ BatchedHummingGroupedExperts,
48
+ HummingGroupedExperts,
49
+ HummingIndexedExperts,
50
+ )
51
+
52
+ return [
53
+ BatchedHummingGroupedExperts,
54
+ HummingGroupedExperts,
55
+ HummingIndexedExperts,
56
+ ]
57
else:
58
raise ValueError(f"Unknown W4A8 MoE backend: {backend.value}")
59
@@ -50,8 +62,19 @@ def select_w4a8_moe_backend(
62
config: FusedMoEConfig,
63
weight_key: QuantKey | None = kInt4Static,
64
activation_key: QuantKey | None = kFp8DynamicTokenSym,
53
-) -> tuple[W4A8MoeBackend, type["CutlassExpertsW4A8Fp8"]]:
54
- backend = W4A8MoeBackend.CUTLASS
65
+) -> tuple[W4A8MoeBackend, type[mk.FusedMoEExperts]]:
66
+ backends = [W4A8MoeBackend.CUTLASS, W4A8MoeBackend.HUMMING]
67
+ if config.moe_backend != "auto":
68
+ backend_map = {
69
+ "cutlass": W4A8MoeBackend.CUTLASS,
70
+ "humming": W4A8MoeBackend.HUMMING,
71
+ }
72
+ if config.moe_backend not in backend_map:
73
+ raise ValueError(
74
+ f"moe_backend='{config.moe_backend}' is not supported for W4A8 MoE. "
75
+ f"Expected one of {list(backend_map)}."
76
+ )
77
+ backends = [backend_map[config.moe_backend]]
78
79
activation_format = (
80
mk.FusedMoEActivationFormat.BatchedExperts
@@ -59,29 +82,33 @@ def select_w4a8_moe_backend(
82
else mk.FusedMoEActivationFormat.Standard
83
)
84
62
- last_reason: str | None = None
63
- for kernel_cls in backend_to_kernel_cls(backend):
64
- supported, reason = kernel_cls.is_supported_config(
65
- kernel_cls,
66
- config,
67
- weight_key,
68
- activation_key,
69
- activation_format,
70
- )
71
- if supported:
72
- logger.info_once("Using %s W4A8 MoE backend.", backend.value)
73
- return backend, kernel_cls
74
- last_reason = reason
85
+ reasons = []
86
+ for backend in backends:
87
+ for kernel_cls in backend_to_kernel_cls(backend):
88
+ supported, reason = kernel_cls.is_supported_config(
89
+ kernel_cls,
90
+ config,
91
+ weight_key,
92
+ activation_key,
93
+ activation_format,
94
+ )
95
+ if supported:
96
+ logger.info_once("Using %s W4A8 MoE backend.", backend.value)
97
+ return backend, kernel_cls
98
+ reasons.append(f"{kernel_cls.__name__}: {reason}")
99
100
raise NotImplementedError(
77
- f"W4A8 MoE backend {backend.value} does not support the "
78
- f"deployment configuration: {last_reason}."
101
+ "No W4A8 MoE backend supports the deployment configuration: "
102
+ + "; ".join(reasons)
103
)
104
105
106
def convert_to_w4a8_moe_kernel_format(
83
- w13_weight_packed: torch.Tensor,
84
- w2_weight_packed: torch.Tensor,
107
+ backend: W4A8MoeBackend,
108
+ layer: "RoutedExperts",
109
+ group_size: int,
110
+ w13_weight: torch.Tensor,
111
+ w2_weight: torch.Tensor,
112
w13_weight_scale: torch.Tensor,
113
w2_weight_scale: torch.Tensor,
114
) -> tuple[
@@ -89,11 +116,40 @@ def convert_to_w4a8_moe_kernel_format(
116
torch.Tensor,
117
torch.Tensor,
118
torch.Tensor,
92
- torch.Tensor,
93
- torch.Tensor,
94
- torch.Tensor,
95
- torch.Tensor,
119
+ torch.Tensor | None,
120
+ torch.Tensor | None,
121
+ torch.Tensor | None,
122
+ torch.Tensor | None,
123
]:
124
+ if backend == W4A8MoeBackend.HUMMING:
125
+ from vllm.model_executor.layers.quantization.utils.humming import (
126
+ convert_to_humming_moe_kernel_format,
127
+ )
128
+ from vllm.utils.humming import (
129
+ HummingInputSchema,
130
+ HummingWeightSchema,
131
+ dtypes,
132
+ )
133
+
134
+ convert_to_humming_moe_kernel_format(
135
+ layer,
136
+ weight_schema=HummingWeightSchema(
137
+ b_dtype=dtypes.uint4, weight_scale_group_size=group_size
138
+ ),
139
+ input_schema=HummingInputSchema(a_dtype=dtypes.float8e4m3),
140
+ allow_input_schema_fallback=False,
141
+ )
142
+ return (
143
+ layer.w13_weight,
144
+ layer.w2_weight,
145
+ layer.w13_weight_scale,
146
+ layer.w2_weight_scale,
147
+ None,
148
+ None,
149
+ None,
150
+ None,
151
+ )
152
+
153
from vllm import _custom_ops as ops
154
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
155
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -104,18 +160,18 @@ def convert_to_w4a8_moe_kernel_format(
160
161
quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
162
107
- convert_packed_uint4b8_to_signed_int4_inplace(w13_weight_packed)
163
+ convert_packed_uint4b8_to_signed_int4_inplace(w13_weight)
164
# Mirror the sync in CutlassW4A8LinearKernel; required for TP>1 correctness.
165
torch.accelerator.synchronize()
166
w13_weight_shuffled, b_strides1 = ops.cutlass_encode_and_reorder_int4b_grouped(
111
- w13_weight_packed
167
+ w13_weight
168
)
169
114
- convert_packed_uint4b8_to_signed_int4_inplace(w2_weight_packed)
170
+ convert_packed_uint4b8_to_signed_int4_inplace(w2_weight)
171
# Mirror the sync in CutlassW4A8LinearKernel; required for TP>1 correctness.
172
torch.accelerator.synchronize()
173
w2_weight_shuffled, b_strides2 = ops.cutlass_encode_and_reorder_int4b_grouped(
118
- w2_weight_packed
174
+ w2_weight
175
)
176
177
w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8(
@@ -147,11 +203,23 @@ def convert_to_w4a8_moe_kernel_format(
203
204
205
def make_w4a8_moe_quant_config(
206
+ backend: W4A8MoeBackend,
207
w1_scale: torch.Tensor,
208
w2_scale: torch.Tensor,
152
- g1_alphas: torch.Tensor,
153
- g2_alphas: torch.Tensor,
209
+ g1_alphas: torch.Tensor | None,
210
+ g2_alphas: torch.Tensor | None,
211
+ layer: torch.nn.Module | None = None,
212
) -> FusedMoEQuantConfig:
213
+ if backend == W4A8MoeBackend.HUMMING:
214
+ from vllm.model_executor.layers.fused_moe import RoutedExperts
215
+ from vllm.model_executor.layers.quantization.utils.humming import (
216
+ get_humming_moe_quant_config,
217
+ )
218
+
219
+ assert isinstance(layer, RoutedExperts)
220
+ return get_humming_moe_quant_config(layer)
221
+
222
+ assert g1_alphas is not None and g2_alphas is not None
223
return int4_w4afp8_moe_quant_config(
224
w1_scale=w1_scale,
225
w2_scale=w2_scale,
@@ -163,14 +231,30 @@ def make_w4a8_moe_quant_config(
231
232
233
def make_w4a8_moe_kernel(
234
+ backend: W4A8MoeBackend,
235
moe_quant_config: FusedMoEQuantConfig,
236
moe_config: FusedMoEConfig,
168
- experts_cls: type["CutlassExpertsW4A8Fp8"],
169
- b_strides1: torch.Tensor,
170
- b_strides2: torch.Tensor,
237
+ experts_cls: type[mk.FusedMoEExperts],
238
+ b_strides1: torch.Tensor | None,
239
+ b_strides2: torch.Tensor | None,
240
group_size: int,
241
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
242
) -> mk.FusedMoEKernel:
243
+ if backend == W4A8MoeBackend.HUMMING:
244
+ from vllm.model_executor.layers.quantization.utils.humming import (
245
+ make_humming_moe_kernel,
246
+ )
247
+
248
+ return make_humming_moe_kernel(
249
+ moe_quant_config, moe_config, experts_cls, routing_tables=routing_tables
250
+ )
251
+
252
+ from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
253
+ CutlassExpertsW4A8Fp8,
254
+ )
255
+
256
+ assert issubclass(experts_cls, CutlassExpertsW4A8Fp8)
257
+ assert b_strides1 is not None and b_strides2 is not None
258
prepare_finalize = maybe_make_prepare_finalize(
259
moe=moe_config,
260
quant_config=moe_quant_config,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+12
@@ -511,6 +511,13 @@ class CompressedTensorsConfig(QuantizationConfig):
511
def _is_dynamic_token_w4a8_int(
512
weight_quant: QuantizationArgs, input_quant: QuantizationArgs
513
) -> bool:
514
+ if not weight_quant or not input_quant:
515
+ return False
516
+ if (
517
+ weight_quant.type != QuantizationType.INT
518
+ or input_quant.type != QuantizationType.INT
519
+ ):
520
+ return False
521
is_weight_4_bits = weight_quant.num_bits == 4
522
is_activation_8_bits = input_quant.num_bits == 8
523
weight_strategy = (
@@ -575,6 +582,11 @@ class CompressedTensorsConfig(QuantizationConfig):
582
) -> bool:
583
if not weight_quant or not input_quant:
584
return False
585
+ if (
586
+ weight_quant.type != QuantizationType.INT
587
+ or input_quant.type != QuantizationType.FLOAT
588
+ ):
589
+ return False
590
is_weight_4_bits = weight_quant.num_bits == 4
591
is_activation_8_bits = input_quant.num_bits == 8
592
weight_strategy = weight_quant.strategy == QuantizationStrategy.GROUP.value
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+1
-1
@@ -164,7 +164,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
164
return CompressedTensorsW8A8Int8MoEMethod(
165
weight_quant, input_quant, layer.moe_config
166
)
167
- elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant):
167
+ elif quant_config._is_fp8_w4a8(weight_quant, input_quant):
168
from .compressed_tensors_moe_w4a8_fp8 import (
169
CompressedTensorsW4A8Fp8MoEMethod,
170
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py
+27
-33
@@ -143,57 +143,49 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
143
layer.register_parameter("w13_weight_shape", w13_weight_shape)
144
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
145
146
- w13_weight_chan_scale = torch.nn.Parameter(
147
- torch.ones(
148
- num_experts,
149
- self.moe.w13_num_shards * intermediate_size_per_partition,
150
- dtype=torch.float32,
151
- ),
152
- requires_grad=False,
153
- )
154
- layer.register_parameter("w13_weight_chan_scale", w13_weight_chan_scale)
155
-
156
- w2_weight_chan_scale = torch.nn.Parameter(
157
- torch.ones(num_experts, hidden_size, dtype=torch.float32),
158
- requires_grad=False,
159
- )
160
- layer.register_parameter("w2_weight_chan_scale", w2_weight_chan_scale)
161
-
146
# don't use input scales
147
layer.w13_input_scale = None
148
layer.w2_input_scale = None
149
150
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
151
+ for name in ("w13", "w2"):
152
+ setattr(layer, f"{name}_weight", getattr(layer, f"{name}_weight_packed"))
153
+ delattr(layer, f"{name}_weight_packed")
154
+
155
(
168
- w13_weight_packed,
169
- w2_weight_packed,
156
+ w13_weight,
157
+ w2_weight,
158
w13_weight_scale,
159
w2_weight_scale,
160
w13_weight_chan_scale,
161
w2_weight_chan_scale,
174
- b_strides1,
175
- b_strides2,
162
+ self.b_strides1,
163
+ self.b_strides2,
164
) = convert_to_w4a8_moe_kernel_format(
177
- w13_weight_packed=layer.w13_weight_packed,
178
- w2_weight_packed=layer.w2_weight_packed,
165
+ backend=self.w4a8_backend,
166
+ layer=layer,
167
+ group_size=self.group_size,
168
+ w13_weight=layer.w13_weight,
169
+ w2_weight=layer.w2_weight,
170
w13_weight_scale=layer.w13_weight_scale,
171
w2_weight_scale=layer.w2_weight_scale,
172
)
173
183
- replace_parameter(layer, "w13_weight_packed", w13_weight_packed)
184
- replace_parameter(layer, "w2_weight_packed", w2_weight_packed)
185
- replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
186
- replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
187
- replace_parameter(layer, "w13_weight_chan_scale", w13_weight_chan_scale)
188
- replace_parameter(layer, "w2_weight_chan_scale", w2_weight_chan_scale)
189
-
190
- self.b_strides1 = b_strides1
191
- self.b_strides2 = b_strides2
174
+ for name, tensor in (
175
+ ("w13_weight", w13_weight),
176
+ ("w2_weight", w2_weight),
177
+ ("w13_weight_scale", w13_weight_scale),
178
+ ("w2_weight_scale", w2_weight_scale),
179
+ ("w13_weight_chan_scale", w13_weight_chan_scale),
180
+ ("w2_weight_chan_scale", w2_weight_chan_scale),
181
+ ):
182
+ replace_parameter(layer, name, tensor)
183
184
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
185
if self.moe_quant_config is not None:
186
assert self.experts_cls is not None
187
self.moe_kernel = make_w4a8_moe_kernel(
188
+ backend=self.w4a8_backend,
189
moe_quant_config=self.moe_quant_config,
190
moe_config=self.moe,
191
experts_cls=self.experts_cls,
@@ -205,10 +197,12 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
197
198
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
199
return make_w4a8_moe_quant_config(
200
+ backend=self.w4a8_backend,
201
w1_scale=layer.w13_weight_scale,
202
w2_scale=layer.w2_weight_scale,
203
g1_alphas=layer.w13_weight_chan_scale,
204
g2_alphas=layer.w2_weight_chan_scale,
205
+ layer=layer,
206
)
207
208
def apply(
@@ -224,8 +218,8 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
218
assert self.moe_kernel is not None
219
return self.moe_kernel.apply(
220
hidden_states=x,
227
- w1=layer.w13_weight_packed,
228
- w2=layer.w2_weight_packed,
221
+ w1=layer.w13_weight,
222
+ w2=layer.w2_weight,
223
topk_weights=topk_weights,
224
topk_ids=topk_ids,
225
activation=layer.activation,