Fix Humming non-gated MoE (#49096)
Netanel Haber committed
Jul 28, 2026 at 01:56 UTC
1e34a135390b4fa3370cf55885b442667c3b8ba6
3 files changed
+111
-7
tests/kernels/moe/test_moe.py
+96
@@ -1165,6 +1165,102 @@ def test_fused_marlin_moe_non_gated(
1165
torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0)
1166
1167
1168
+@pytest.mark.parametrize(
1169
+ "activation",
1170
+ [
1171
+ MoEActivation.SILU,
1172
+ MoEActivation.RELU2_NO_MUL,
1173
+ ],
1174
+ ids=["gated", "non_gated"],
1175
+)
1176
+def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
1177
+ pytest.importorskip("humming")
1178
+ from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1179
+ HummingIndexedExperts,
1180
+ )
1181
+ from vllm.model_executor.layers.quantization.utils import humming_utils
1182
+ from vllm.utils import humming
1183
+
1184
+ top_k, num_experts = 6, 12
1185
+ hidden_size, intermediate_size = 2688, 1856
1186
+ gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size
1187
+ num_w13_stacks = 2 if activation.is_gated else 1
1188
+ moe_config = make_dummy_moe_config(
1189
+ num_experts=num_experts,
1190
+ experts_per_token=top_k,
1191
+ hidden_dim=hidden_size,
1192
+ intermediate_size=intermediate_size,
1193
+ activation=activation,
1194
+ )
1195
+
1196
+ layer = torch.nn.Module()
1197
+ layer.moe_config = moe_config
1198
+ layer.params_dtype = torch.bfloat16
1199
+
1200
+ weight_schema = humming.ModeloptNvfp4WeightSchema()
1201
+ for sublayer_name, shape_n, shape_k, stack_size in (
1202
+ ("w13", gate_up_size, hidden_size, num_w13_stacks),
1203
+ ("w2", hidden_size, intermediate_size, 1),
1204
+ ):
1205
+ tensor_attrs = weight_schema.get_tensors_attrs(
1206
+ shape_n=shape_n,
1207
+ shape_k=shape_k,
1208
+ param_dtype=layer.params_dtype,
1209
+ num_experts=num_experts,
1210
+ stack_size=stack_size,
1211
+ )
1212
+ for tensor_name, attrs in tensor_attrs.items():
1213
+ layer.register_parameter(
1214
+ f"{sublayer_name}_{tensor_name}",
1215
+ Parameter(
1216
+ torch.ones(
1217
+ attrs["shape"],
1218
+ dtype=attrs["dtype"],
1219
+ device="cuda",
1220
+ ),
1221
+ requires_grad=False,
1222
+ ),
1223
+ )
1224
+
1225
+ humming_utils.convert_to_humming_moe_kernel_format(
1226
+ layer,
1227
+ weight_schema=weight_schema,
1228
+ input_schema=humming.HummingInputSchema(a_dtype=humming.dtypes.bfloat16),
1229
+ )
1230
+
1231
+ w13_meta, w2_meta = (layer.humming_metas[name] for name in ("w13", "w2"))
1232
+ for meta in (w13_meta, w2_meta):
1233
+ assert meta.a_dtype == humming.dtypes.bfloat16
1234
+ assert meta.b_dtype == humming.dtypes.float4e2m1
1235
+
1236
+ assert w13_meta.shape_n - w13_meta.pad_shape_n == gate_up_size
1237
+ assert w2_meta.shape_k - w2_meta.pad_shape_k == intermediate_size
1238
+
1239
+ layer.local_num_experts = layer.global_num_experts = num_experts
1240
+ layer.hidden_size = hidden_size
1241
+ layer.intermediate_size_per_partition = intermediate_size
1242
+ quant_config = humming_utils.get_humming_moe_quant_config(layer)
1243
+ experts = HummingIndexedExperts(
1244
+ layer,
1245
+ moe_config,
1246
+ quant_config,
1247
+ )
1248
+
1249
+ buffer_metas, _ = experts.get_buffer_metas(
1250
+ M=1,
1251
+ topk=top_k,
1252
+ activation=moe_config.activation,
1253
+ )
1254
+ assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size
1255
+ assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size
1256
+ assert experts.moe_problem_size(
1257
+ a1=torch.empty(1, hidden_size),
1258
+ w1=torch.empty(num_experts, 1),
1259
+ w2=torch.empty(num_experts, 1),
1260
+ topk_ids=torch.empty(1, top_k, dtype=torch.long),
1261
+ ) == (num_experts, 1, intermediate_size, hidden_size, top_k)
1262
+
1263
+
1264
@pytest.mark.parametrize("ep_size", [1, 2])
1265
def test_moe_align_block_size_opcheck(ep_size):
1266
num_experts = 4
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+11
-3
@@ -289,7 +289,14 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
289
assert a1.size(0) == num_experts
290
num_tokens = a1.size(1)
291
292
- return meta1.num_experts, num_tokens, meta1.shape_n // 2, meta1.shape_k, top_k
292
+ return (
293
+ meta1.num_experts,
294
+ num_tokens,
295
+ # Logical intermediate width for both gated and non-gated activations
296
+ self.layer.intermediate_size_per_partition,
297
+ meta1.shape_k,
298
+ top_k,
299
+ )
300
301
def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
302
from vllm.utils.humming import GemmType as HummingGemmType
@@ -327,7 +334,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
334
real_shape_m = M * topk
335
output_shape = (M, K)
336
330
- down_input_size = N if activation.is_gated else (N * 2)
337
+ gate_up_size = N * (2 if activation.is_gated else 1)
338
+ down_input_size = N
339
a_dtype = self.layer.humming_metas["w13"].a_dtype
340
c_dtype = self.layer.humming_metas["w13"].c_dtype
341
num_bits = a_dtype.num_bits
@@ -347,7 +355,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
355
"dtype": torch_dtype_map[a_dtype],
356
},
357
"gate_up_output": {
350
- "shape": (real_shape_m, N * 2),
358
+ "shape": (real_shape_m, gate_up_size),
359
"dtype": torch_dtype_map[c_dtype],
360
},
361
"activation_output": {
vllm/model_executor/layers/quantization/utils/humming_utils.py
+4
-4
@@ -788,7 +788,7 @@ def _convert_sublayer_to_humming(
788
789
shape_k_stacks = [shape_k]
790
shape_n_stacks = [shape_n]
791
- if sublayer_name == "w13":
791
+ if sublayer_name == "w13" and layer.moe_config.activation.is_gated:
792
shape_n_stacks = [shape_n // 2] * 2
793
794
converted_weight_schema, converted_tensors = weight_schema.convert_humming(
@@ -991,15 +991,15 @@ def convert_to_humming_moe_kernel_format(
991
# Build sublayer configs from layer properties if not provided
992
if sublayer_configs is None:
993
is_gated = layer.moe_config.activation.is_gated
994
+ intermediate_size = layer.moe_config.intermediate_size_per_partition
995
sublayer_configs = {
996
"w13": {
996
- "shape_n": layer.moe_config.intermediate_size_per_partition * 2,
997
+ "shape_n": intermediate_size * (2 if is_gated else 1),
998
"shape_k": layer.moe_config.hidden_dim,
999
},
1000
"w2": {
1001
"shape_n": layer.moe_config.hidden_dim,
1001
- "shape_k": layer.moe_config.intermediate_size_per_partition
1002
- * (1 if is_gated else 2),
1002
+ "shape_k": intermediate_size,
1003
},
1004
}
1005