diff --git a/tests/kernels/quantization/test_triton_w4a16.py b/tests/kernels/quantization/test_triton_w4a16.py index 42f163dea4..3c46829116 100644 --- a/tests/kernels/quantization/test_triton_w4a16.py +++ b/tests/kernels/quantization/test_triton_w4a16.py @@ -306,6 +306,7 @@ def test_triton_w4a16_process_weights_after_loading_repacks_layout(): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout(): + """AutoGPTQ qzeros are already [K//G, N//8] (output_dim=1): no transpose.""" if not torch.cuda.is_available(): pytest.skip("CUDA/HIP device not available") @@ -413,6 +414,7 @@ def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout(): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_triton_w4a16_symmetric_apply_ignores_qzeros(monkeypatch): + """For symmetric (uint4b8) layers, apply_weights must pass qzeros=None.""" from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( MPLinearLayerConfig, ) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py index bc0a587b67..8a4a493c6c 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py @@ -400,29 +400,40 @@ class TritonW4A16LinearKernel(MPLinearKernel): if self.w_zp_name is not None: zp = getattr(layer, self.w_zp_name, None) if zp is not None: - c = self.config - K, N = c.partition_weight_shape - group_size = c.group_size if c.group_size != -1 else K - expected_shape = (K // group_size, N // 8) - transposed_shape = (N // 8, K // group_size) - - if tuple(zp.data.shape) == expected_shape: - # GPTQ/AutoGPTQ already stores qzeros in the kernel layout. - qzeros = zp.data.contiguous() - elif tuple(zp.data.shape) == transposed_shape: - # Compressed-tensors stores qzeros transposed from what the - # kernel needs. - qzeros = zp.data.t().contiguous() + # Kernel needs [K//G, N//8]: + # input(K) at dim 0, output(N) packed at dim 1. + # AutoGPTQ: + # output_dim=1 -> already [K//G, N//8], no transpose. + # compressed-tensors: + # output_dim=0 -> [N//8, K//G], needs transpose. + # None (unknown): + # infer from shape; if square (ambiguous), default to transpose. + zp_output_dim = getattr(zp, "output_dim", None) + if zp_output_dim is not None: + needs_transpose = zp_output_dim != 1 else: - raise AssertionError( - f"{self.w_zp_name} shape mismatch: {zp.data.shape}; " - f"expected {expected_shape} or {transposed_shape}" - ) - + # in case output_dim is None + c = self.config + K, N = c.partition_weight_shape + group_size = c.group_size if c.group_size != -1 else K + expected_shape = (K // group_size, N // 8) + transposed_shape = (N // 8, K // group_size) + if ( + tuple(zp.data.shape) == expected_shape + and expected_shape != transposed_shape + ): + needs_transpose = False + else: + needs_transpose = True + zp_data = ( + zp.data.t().contiguous() + if needs_transpose + else zp.data.contiguous() + ) replace_parameter( layer, self.w_zp_name, - torch.nn.Parameter(qzeros, requires_grad=False), + torch.nn.Parameter(zp_data, requires_grad=False), ) def apply_weights( @@ -437,17 +448,18 @@ class TritonW4A16LinearKernel(MPLinearKernel): K = c.partition_weight_shape[0] group_size = c.group_size if c.group_size != -1 else K - # For symmetric types (uint4b8), use the scalar bias; no zeros tensor. - # Some checkpoint loaders still register qzeros parameters for GPTQ - # layers, but they are not part of the symmetric kernel contract. - zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0 - qzeros = None if c.weight_type.has_bias() else w_zp + # For symmetric types (uint4b8), use the scalar bias; no zeros tensor + if c.weight_type.has_bias(): + zp_bias = c.weight_type.bias + w_zp = None # symmetric: ignore qzeros, use scalar bias instead + else: + zp_bias = 0 output = triton_w4a16_gemm( a=x_2d, b_q=w_q, scales=w_s, - qzeros=qzeros, + qzeros=w_zp, group_size=group_size, zp_bias=zp_bias, )