diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index 80014d5c99..e5be841f30 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -366,3 +366,110 @@ def test_xpu_platform_supports_moe_wna16(): pytest.skip("vllm_xpu_kernels not importable outside an XPU stack") assert "moe_wna16" in XPUPlatform.supported_quantization + + +def _channelwise_int4_args() -> QuantizationArgs: + """A per-channel int4 checkpoint, which leaves ``group_size`` unset.""" + args = QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.CHANNEL, + symmetric=True, + dynamic=False, + ) + assert args.group_size is None, "premise: CHANNEL leaves group_size unset" + return args + + +@pytest.mark.skipif( + current_platform.is_rocm(), + reason="check_moe_marlin_supports_config rejects every config on ROCm", +) +@pytest.mark.parametrize("backend", [WNA16MoEBackend.MARLIN, WNA16MoEBackend.TRITON]) +def test_wna16_oracle_accepts_unset_group_size(backend): + """Both backends must *accept* a per-channel config, not just survive it. + + -1 is a supported Marlin group size and the shapes below pass the Marlin + tiling checks, while Triton never reads group_size for QuantizationArgs. + A reason string from either backend would mean the unset group_size cost + the layer its preferred kernel instead of raising TypeError. + """ + from tests.kernels.moe.utils import make_dummy_moe_config + + # hidden_dim % 128 and intermediate % 64 must hold, or the Marlin shape + # check rejects the config before group_size is ever read. + reason = _backend_incompatibility_reason( + backend=backend, + moe_config=make_dummy_moe_config( + num_experts=2, hidden_dim=256, intermediate_size=512 + ), + quant_config=_channelwise_int4_args(), + may_have_zp=False, + may_have_bias=False, + allow_tile_padding=True, + ) + + assert reason is None + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="Marlin is only a candidate WNA16 MoE backend on CUDA; elsewhere " + "__init__ takes the non-Marlin branch, which rejects channelwise", +) +def test_compressed_tensors_wna16_moe_marlin_prep_with_unset_group_size(): + """Load a per-channel checkpoint through the Marlin path end to end. + + ``__init__`` and Marlin weight prep read ``group_size`` independently, so + both have to normalise the unset value. The post-repack shapes prove prep + ran with the Marlin K/N rather than merely returning something. + """ + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 + CompressedTensorsWNA16MoEMethod, + ) + + num_experts, hidden_size, intermediate_size = 2, 256, 512 + moe_config = make_dummy_moe_config( + num_experts=num_experts, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + ) + moe_config.moe_backend = "marlin" + + method = CompressedTensorsWNA16MoEMethod(_channelwise_int4_args(), None, moe_config) + assert method.wna16_backend == WNA16MoEBackend.MARLIN + assert method.group_size == -1 + + layer = torch.nn.Module() + layer.intermediate_size_per_partition = intermediate_size + layer._expert_routing_tables = lambda: (None, None, None) + method.create_weights( + layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size, + intermediate_size_full=intermediate_size, + params_dtype=torch.bfloat16, + ) + layer.cuda() + for parameter in layer.parameters(): + parameter.data.zero_() + + method.process_weights_after_loading(layer) + + # gptq_marlin_moe_repack packs to (size_k // 16, size_n * 2) for int4; w13 + # is repacked with size_k=hidden_size, size_n=2*intermediate_size, and w2 + # the other way round. Channelwise keeps one scale group per channel. + assert layer.w13_weight_packed.shape == ( + num_experts, + hidden_size // 16, + 4 * intermediate_size, + ) + assert layer.w2_weight_packed.shape == ( + num_experts, + intermediate_size // 16, + 2 * hidden_size, + ) + assert layer.w13_weight_scale.shape == (num_experts, 1, 2 * intermediate_size) + assert layer.w2_weight_scale.shape == (num_experts, 1, hidden_size) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index c56712ac58..5d3212ae14 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -160,7 +160,11 @@ def _backend_incompatibility_reason( WNA16MoEBackend.BATCHED_MARLIN, ): if isinstance(quant_config, (AutoAWQConfig, AutoGPTQConfig, QuantizationArgs)): - group_size = quant_config.group_size + # QuantizationArgs leaves group_size unset for per-channel + # strategies; -1 is the in-tree encoding for that. + group_size = ( + -1 if quant_config.group_size is None else quant_config.group_size + ) else: return "Marlin not supported for this layer" @@ -1509,7 +1513,9 @@ def convert_to_wna16_moe_kernel_format( elif isinstance(quant_config, QuantizationArgs): num_bits = quant_config.num_bits pack_factor = 32 // quant_config.num_bits - group_size = quant_config.group_size + group_size = ( + -1 if quant_config.group_size is None else quant_config.group_size + ) actorder = quant_config.actorder else: raise TypeError( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 1ab5085145..ef8e6fe089 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -70,7 +70,11 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): self.num_bits = weight_quant.num_bits self.packed_factor = Fraction(32, weight_quant.num_bits) self.strategy = weight_quant.strategy - self.group_size = weight_quant.group_size + # Per-channel strategies leave group_size unset; -1 is the in-tree + # encoding for that, as in CompressedTensorsWNA16 and siblings. + self.group_size = ( + -1 if weight_quant.group_size is None else weight_quant.group_size + ) self.actorder = weight_quant.actorder # Extract quant_type and create weight key for oracle selection