[ROCm][Bugfix] Fix Triton W4A16 bug in determining if transpose is required for GPTQ/AutoGPTQ (#48998)
Signed-off-by: Qiang Li <qiang.li2@amd.com>
qli88 committed
Aug 17, 2026 at 07:26 UTC
1d3a8b9e220f1edc77c190c3370cf0f76dfdd2fd
2 files changed
+39
-25
tests/kernels/quantization/test_triton_w4a16.py
+2
@@ -306,6 +306,7 @@ def test_triton_w4a16_process_weights_after_loading_repacks_layout():
306
307
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
308
def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout():
309
+ """AutoGPTQ qzeros are already [K//G, N//8] (output_dim=1): no transpose."""
310
if not torch.cuda.is_available():
311
pytest.skip("CUDA/HIP device not available")
312
@@ -413,6 +414,7 @@ def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout():
414
415
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
416
def test_triton_w4a16_symmetric_apply_ignores_qzeros(monkeypatch):
417
+ """For symmetric (uint4b8) layers, apply_weights must pass qzeros=None."""
418
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
419
MPLinearLayerConfig,
420
)
vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
+37
-25
@@ -400,29 +400,40 @@ class TritonW4A16LinearKernel(MPLinearKernel):
400
if self.w_zp_name is not None:
401
zp = getattr(layer, self.w_zp_name, None)
402
if zp is not None:
403
- c = self.config
404
- K, N = c.partition_weight_shape
405
- group_size = c.group_size if c.group_size != -1 else K
406
- expected_shape = (K // group_size, N // 8)
407
- transposed_shape = (N // 8, K // group_size)
408
-
409
- if tuple(zp.data.shape) == expected_shape:
410
- # GPTQ/AutoGPTQ already stores qzeros in the kernel layout.
411
- qzeros = zp.data.contiguous()
412
- elif tuple(zp.data.shape) == transposed_shape:
413
- # Compressed-tensors stores qzeros transposed from what the
414
- # kernel needs.
415
- qzeros = zp.data.t().contiguous()
403
+ # Kernel needs [K//G, N//8]:
404
+ # input(K) at dim 0, output(N) packed at dim 1.
405
+ # AutoGPTQ:
406
+ # output_dim=1 -> already [K//G, N//8], no transpose.
407
+ # compressed-tensors:
408
+ # output_dim=0 -> [N//8, K//G], needs transpose.
409
+ # None (unknown):
410
+ # infer from shape; if square (ambiguous), default to transpose.
411
+ zp_output_dim = getattr(zp, "output_dim", None)
412
+ if zp_output_dim is not None:
413
+ needs_transpose = zp_output_dim != 1
414
else:
417
- raise AssertionError(
418
- f"{self.w_zp_name} shape mismatch: {zp.data.shape}; "
419
- f"expected {expected_shape} or {transposed_shape}"
420
- )
421
-
415
+ # in case output_dim is None
416
+ c = self.config
417
+ K, N = c.partition_weight_shape
418
+ group_size = c.group_size if c.group_size != -1 else K
419
+ expected_shape = (K // group_size, N // 8)
420
+ transposed_shape = (N // 8, K // group_size)
421
+ if (
422
+ tuple(zp.data.shape) == expected_shape
423
+ and expected_shape != transposed_shape
424
+ ):
425
+ needs_transpose = False
426
+ else:
427
+ needs_transpose = True
428
+ zp_data = (
429
+ zp.data.t().contiguous()
430
+ if needs_transpose
431
+ else zp.data.contiguous()
432
+ )
433
replace_parameter(
434
layer,
435
self.w_zp_name,
425
- torch.nn.Parameter(qzeros, requires_grad=False),
436
+ torch.nn.Parameter(zp_data, requires_grad=False),
437
)
438
439
def apply_weights(
@@ -437,17 +448,18 @@ class TritonW4A16LinearKernel(MPLinearKernel):
448
K = c.partition_weight_shape[0]
449
group_size = c.group_size if c.group_size != -1 else K
450
440
- # For symmetric types (uint4b8), use the scalar bias; no zeros tensor.
441
- # Some checkpoint loaders still register qzeros parameters for GPTQ
442
- # layers, but they are not part of the symmetric kernel contract.
443
- zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0
444
- qzeros = None if c.weight_type.has_bias() else w_zp
451
+ # For symmetric types (uint4b8), use the scalar bias; no zeros tensor
452
+ if c.weight_type.has_bias():
453
+ zp_bias = c.weight_type.bias
454
+ w_zp = None # symmetric: ignore qzeros, use scalar bias instead
455
+ else:
456
+ zp_bias = 0
457
458
output = triton_w4a16_gemm(
459
a=x_2d,
460
b_q=w_q,
461
scales=w_s,
450
- qzeros=qzeros,
462
+ qzeros=w_zp,
463
group_size=group_size,
464
zp_bias=zp_bias,
465
)