[Bugfix][Quantization][MoE] Normalise an unset group_size on the compressed-tensors WNA16 MoE path (#53163)
Signed-off-by: Artur Fierka <artur.fierka@intel.com> Co-authored-by: Claude <noreply@anthropic.com>
Artur Fierka committed
Sep 8, 2026 at 03:42 UTC
9dbdf8e9af660e3b6876b27c2d1590bcab2e2de3
3 files changed
+120
-3
tests/quantization/test_moe_wna16.py
+107
@@ -366,3 +366,110 @@ def test_xpu_platform_supports_moe_wna16():
366
pytest.skip("vllm_xpu_kernels not importable outside an XPU stack")
367
368
assert "moe_wna16" in XPUPlatform.supported_quantization
369
+
370
+
371
+def _channelwise_int4_args() -> QuantizationArgs:
372
+ """A per-channel int4 checkpoint, which leaves ``group_size`` unset."""
373
+ args = QuantizationArgs(
374
+ num_bits=4,
375
+ type=QuantizationType.INT,
376
+ strategy=QuantizationStrategy.CHANNEL,
377
+ symmetric=True,
378
+ dynamic=False,
379
+ )
380
+ assert args.group_size is None, "premise: CHANNEL leaves group_size unset"
381
+ return args
382
+
383
+
384
+@pytest.mark.skipif(
385
+ current_platform.is_rocm(),
386
+ reason="check_moe_marlin_supports_config rejects every config on ROCm",
387
+)
388
+@pytest.mark.parametrize("backend", [WNA16MoEBackend.MARLIN, WNA16MoEBackend.TRITON])
389
+def test_wna16_oracle_accepts_unset_group_size(backend):
390
+ """Both backends must *accept* a per-channel config, not just survive it.
391
+
392
+ -1 is a supported Marlin group size and the shapes below pass the Marlin
393
+ tiling checks, while Triton never reads group_size for QuantizationArgs.
394
+ A reason string from either backend would mean the unset group_size cost
395
+ the layer its preferred kernel instead of raising TypeError.
396
+ """
397
+ from tests.kernels.moe.utils import make_dummy_moe_config
398
+
399
+ # hidden_dim % 128 and intermediate % 64 must hold, or the Marlin shape
400
+ # check rejects the config before group_size is ever read.
401
+ reason = _backend_incompatibility_reason(
402
+ backend=backend,
403
+ moe_config=make_dummy_moe_config(
404
+ num_experts=2, hidden_dim=256, intermediate_size=512
405
+ ),
406
+ quant_config=_channelwise_int4_args(),
407
+ may_have_zp=False,
408
+ may_have_bias=False,
409
+ allow_tile_padding=True,
410
+ )
411
+
412
+ assert reason is None
413
+
414
+
415
+@pytest.mark.skipif(
416
+ not current_platform.is_cuda(),
417
+ reason="Marlin is only a candidate WNA16 MoE backend on CUDA; elsewhere "
418
+ "__init__ takes the non-Marlin branch, which rejects channelwise",
419
+)
420
+def test_compressed_tensors_wna16_moe_marlin_prep_with_unset_group_size():
421
+ """Load a per-channel checkpoint through the Marlin path end to end.
422
+
423
+ ``__init__`` and Marlin weight prep read ``group_size`` independently, so
424
+ both have to normalise the unset value. The post-repack shapes prove prep
425
+ ran with the Marlin K/N rather than merely returning something.
426
+ """
427
+ from tests.kernels.moe.utils import make_dummy_moe_config
428
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
429
+ CompressedTensorsWNA16MoEMethod,
430
+ )
431
+
432
+ num_experts, hidden_size, intermediate_size = 2, 256, 512
433
+ moe_config = make_dummy_moe_config(
434
+ num_experts=num_experts,
435
+ hidden_dim=hidden_size,
436
+ intermediate_size=intermediate_size,
437
+ )
438
+ moe_config.moe_backend = "marlin"
439
+
440
+ method = CompressedTensorsWNA16MoEMethod(_channelwise_int4_args(), None, moe_config)
441
+ assert method.wna16_backend == WNA16MoEBackend.MARLIN
442
+ assert method.group_size == -1
443
+
444
+ layer = torch.nn.Module()
445
+ layer.intermediate_size_per_partition = intermediate_size
446
+ layer._expert_routing_tables = lambda: (None, None, None)
447
+ method.create_weights(
448
+ layer,
449
+ num_experts=num_experts,
450
+ hidden_size=hidden_size,
451
+ intermediate_size_per_partition=intermediate_size,
452
+ intermediate_size_full=intermediate_size,
453
+ params_dtype=torch.bfloat16,
454
+ )
455
+ layer.cuda()
456
+ for parameter in layer.parameters():
457
+ parameter.data.zero_()
458
+
459
+ method.process_weights_after_loading(layer)
460
+
461
+ # gptq_marlin_moe_repack packs to (size_k // 16, size_n * 2) for int4; w13
462
+ # is repacked with size_k=hidden_size, size_n=2*intermediate_size, and w2
463
+ # the other way round. Channelwise keeps one scale group per channel.
464
+ assert layer.w13_weight_packed.shape == (
465
+ num_experts,
466
+ hidden_size // 16,
467
+ 4 * intermediate_size,
468
+ )
469
+ assert layer.w2_weight_packed.shape == (
470
+ num_experts,
471
+ intermediate_size // 16,
472
+ 2 * hidden_size,
473
+ )
474
+ assert layer.w13_weight_scale.shape == (num_experts, 1, 2 * intermediate_size)
475
+ assert layer.w2_weight_scale.shape == (num_experts, 1, hidden_size)
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+8
-2
@@ -160,7 +160,11 @@ def _backend_incompatibility_reason(
160
WNA16MoEBackend.BATCHED_MARLIN,
161
):
162
if isinstance(quant_config, (AutoAWQConfig, AutoGPTQConfig, QuantizationArgs)):
163
- group_size = quant_config.group_size
163
+ # QuantizationArgs leaves group_size unset for per-channel
164
+ # strategies; -1 is the in-tree encoding for that.
165
+ group_size = (
166
+ -1 if quant_config.group_size is None else quant_config.group_size
167
+ )
168
else:
169
return "Marlin not supported for this layer"
170
@@ -1509,7 +1513,9 @@ def convert_to_wna16_moe_kernel_format(
1513
elif isinstance(quant_config, QuantizationArgs):
1514
num_bits = quant_config.num_bits
1515
pack_factor = 32 // quant_config.num_bits
1512
- group_size = quant_config.group_size
1516
+ group_size = (
1517
+ -1 if quant_config.group_size is None else quant_config.group_size
1518
+ )
1519
actorder = quant_config.actorder
1520
else:
1521
raise TypeError(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+5
-1
@@ -70,7 +70,11 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
70
self.num_bits = 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
73
+ # Per-channel strategies leave group_size unset; -1 is the in-tree
74
+ # encoding for that, as in CompressedTensorsWNA16 and siblings.
75
+ self.group_size = (
76
+ -1 if weight_quant.group_size is None else weight_quant.group_size
77
+ )
78
self.actorder = weight_quant.actorder
79
80
# Extract quant_type and create weight key for oracle selection