Add Humming indexed-MoE regression test (#50468)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 30, 2026 at 16:26 UTC
0eec856cc31cf9d9518b547c37a610d3f02ccddb
1 file changed
+89
-17
tests/kernels/moe/test_moe.py
+89
-17
@@ -1182,15 +1182,7 @@ def test_fused_marlin_moe_non_gated(
1182
torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0)
1183
1184
1185
-@pytest.mark.parametrize(
1186
- "activation",
1187
- [
1188
- MoEActivation.SILU,
1189
- MoEActivation.RELU2_NO_MUL,
1190
- ],
1191
- ids=["gated", "non_gated"],
1192
-)
1193
-def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
1185
+def _make_humming_indexed_experts(activation: MoEActivation):
1186
pytest.importorskip("humming")
1187
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1188
HummingIndexedExperts,
@@ -1245,14 +1237,6 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
1237
input_schema=humming.HummingInputSchema(a_dtype=humming.dtypes.bfloat16),
1238
)
1239
1248
- w13_meta, w2_meta = (layer.humming_metas[name] for name in ("w13", "w2"))
1249
- for meta in (w13_meta, w2_meta):
1250
- assert meta.a_dtype == humming.dtypes.bfloat16
1251
- assert meta.b_dtype == humming.dtypes.float4e2m1
1252
-
1253
- assert w13_meta.shape_n - w13_meta.pad_shape_n == gate_up_size
1254
- assert w2_meta.shape_k - w2_meta.pad_shape_k == intermediate_size
1255
-
1240
layer.local_num_experts = layer.global_num_experts = num_experts
1241
layer.hidden_size = hidden_size
1242
layer.intermediate_size_per_partition = intermediate_size
@@ -1262,6 +1246,36 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
1246
moe_config,
1247
quant_config,
1248
)
1249
+ return experts
1250
+
1251
+
1252
+@pytest.mark.parametrize(
1253
+ "activation",
1254
+ [
1255
+ MoEActivation.SILU,
1256
+ MoEActivation.RELU2_NO_MUL,
1257
+ ],
1258
+ ids=["gated", "non_gated"],
1259
+)
1260
+def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
1261
+ from vllm.utils import humming
1262
+
1263
+ experts = _make_humming_indexed_experts(activation)
1264
+ layer = experts.layer
1265
+ moe_config = experts.moe_config
1266
+ top_k = moe_config.experts_per_token
1267
+ num_experts = moe_config.num_experts
1268
+ hidden_size = moe_config.hidden_dim
1269
+ intermediate_size = moe_config.intermediate_size
1270
+ gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size
1271
+
1272
+ w13_meta, w2_meta = (layer.humming_metas[name] for name in ("w13", "w2"))
1273
+ for meta in (w13_meta, w2_meta):
1274
+ assert meta.a_dtype == humming.dtypes.bfloat16
1275
+ assert meta.b_dtype == humming.dtypes.float4e2m1
1276
+
1277
+ assert w13_meta.shape_n - w13_meta.pad_shape_n == gate_up_size
1278
+ assert w2_meta.shape_k - w2_meta.pad_shape_k == intermediate_size
1279
1280
buffer_metas, _ = experts.get_buffer_metas(
1281
M=1,
@@ -1278,6 +1292,64 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
1292
) == (num_experts, 1, intermediate_size, hidden_size, top_k)
1293
1294
1295
+def test_humming_indexed_writes_supplied_output_buffer():
1296
+ from vllm.forward_context import set_forward_context
1297
+
1298
+ activation = MoEActivation.SILU
1299
+ experts = _make_humming_indexed_experts(activation)
1300
+ moe_config = experts.moe_config
1301
+ num_tokens = 1
1302
+ top_k = moe_config.experts_per_token
1303
+ hidden_size = moe_config.hidden_dim
1304
+ num_experts = moe_config.num_experts
1305
+ workspace13_shape, workspace2_shape, _ = experts.workspace_shapes(
1306
+ M=num_tokens,
1307
+ N=moe_config.intermediate_size,
1308
+ K=hidden_size,
1309
+ topk=top_k,
1310
+ global_num_experts=num_experts,
1311
+ local_num_experts=num_experts,
1312
+ expert_tokens_meta=None,
1313
+ activation=activation,
1314
+ )
1315
+
1316
+ device = torch.device("cuda")
1317
+ dtype = experts.layer.params_dtype
1318
+ workspace13 = torch.empty(workspace13_shape, dtype=dtype, device=device)
1319
+ workspace2 = torch.empty(workspace2_shape, dtype=dtype, device=device)
1320
+ hidden_states = torch.ones((num_tokens, hidden_size), dtype=dtype, device=device)
1321
+ output = torch.full_like(hidden_states, torch.nan)
1322
+ topk_weights = torch.full(
1323
+ (num_tokens, top_k),
1324
+ 1 / top_k,
1325
+ dtype=dtype,
1326
+ device=device,
1327
+ )
1328
+ topk_ids = torch.arange(top_k, dtype=torch.int32, device=device).unsqueeze(0)
1329
+ unused = torch.empty((num_experts, 0), device=device)
1330
+
1331
+ with set_forward_context(None, vllm_config, num_tokens=num_tokens):
1332
+ experts.apply(
1333
+ output=output,
1334
+ hidden_states=hidden_states,
1335
+ w1=unused,
1336
+ w2=unused,
1337
+ topk_weights=topk_weights,
1338
+ topk_ids=topk_ids,
1339
+ activation=activation,
1340
+ global_num_experts=num_experts,
1341
+ expert_map=None,
1342
+ a1q_scale=None,
1343
+ a2_scale=None,
1344
+ workspace13=workspace13,
1345
+ workspace2=workspace2,
1346
+ expert_tokens_meta=None,
1347
+ apply_router_weight_on_input=False,
1348
+ )
1349
+
1350
+ assert torch.isfinite(output).all()
1351
+
1352
+
1353
@pytest.mark.parametrize("ep_size", [1, 2])
1354
def test_moe_align_block_size_opcheck(ep_size):
1355
num_experts = 4