[ROCm][BugFix] Triton W4A16 handling for GPTQ/AutoGPTQ qzeros layout (#47770)
Signed-off-by: giuseppegrossi <ggrossi@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Giuseppe Grossi committed
Jul 15, 2026 at 12:55 UTC
1d99f0f42152a1459e4b54804b44fc1f989558bb
2 files changed
+182
-5
tests/kernels/quantization/test_triton_w4a16.py
+157
@@ -302,3 +302,160 @@ def test_triton_w4a16_process_weights_after_loading_repacks_layout():
302
torch.testing.assert_close(layer.weight_packed, expected_w_kn8)
303
torch.testing.assert_close(layer.weight_scale, expected_scales_gn)
304
torch.testing.assert_close(layer.weight_zero_point, expected_zeros_gn8)
305
+
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
+ if not torch.cuda.is_available():
310
+ pytest.skip("CUDA/HIP device not available")
311
+
312
+ from vllm.config import VllmConfig, set_current_vllm_config
313
+ from vllm.distributed import (
314
+ ensure_model_parallel_initialized,
315
+ init_distributed_environment,
316
+ )
317
+ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
318
+ MPLinearLayerConfig,
319
+ )
320
+ from vllm.model_executor.parameter import (
321
+ GroupQuantScaleParameter,
322
+ PackedvLLMParameter,
323
+ )
324
+ from vllm.scalar_type import scalar_types
325
+
326
+ with set_current_vllm_config(VllmConfig()):
327
+ init_distributed_environment(
328
+ world_size=1,
329
+ rank=0,
330
+ distributed_init_method="tcp://127.0.0.1:0",
331
+ local_rank=0,
332
+ )
333
+ ensure_model_parallel_initialized(1, 1)
334
+
335
+ set_random_seed(0)
336
+
337
+ K, N = 256, 256
338
+ G = 32
339
+
340
+ w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
341
+ w_gptq_k8n = _pack_int4_along_k_to_ckpt(w_int4_kn).t().contiguous()
342
+ scales_gptq_gn = 0.05 * torch.rand((K // G, N), device=device, dtype=torch.float16)
343
+ zeros_int4_gn = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32)
344
+ zeros_gptq_gn8 = _pack_int4_along_n(zeros_int4_gn)
345
+
346
+ config = MPLinearLayerConfig(
347
+ full_weight_shape=(K, N),
348
+ partition_weight_shape=(K, N),
349
+ weight_type=scalar_types.uint4,
350
+ act_type=torch.float16,
351
+ group_size=G,
352
+ zero_points=True,
353
+ has_g_idx=False,
354
+ )
355
+ kernel = TritonW4A16LinearKernel(
356
+ config,
357
+ w_q_param_name="qweight",
358
+ w_s_param_name="scales",
359
+ w_zp_param_name="qzeros",
360
+ w_gidx_param_name=None,
361
+ )
362
+
363
+ weight_loader = lambda *args, **kwargs: None
364
+
365
+ class DummyLayer(torch.nn.Module):
366
+ pass
367
+
368
+ layer = DummyLayer()
369
+ layer.register_parameter(
370
+ "qweight",
371
+ PackedvLLMParameter(
372
+ data=w_gptq_k8n,
373
+ weight_loader=weight_loader,
374
+ input_dim=0,
375
+ output_dim=1,
376
+ packed_factor=8,
377
+ packed_dim=0,
378
+ ),
379
+ )
380
+ layer.register_parameter(
381
+ "scales",
382
+ GroupQuantScaleParameter(
383
+ data=scales_gptq_gn,
384
+ weight_loader=weight_loader,
385
+ input_dim=0,
386
+ output_dim=1,
387
+ ),
388
+ )
389
+ layer.register_parameter(
390
+ "qzeros",
391
+ PackedvLLMParameter(
392
+ data=zeros_gptq_gn8,
393
+ weight_loader=weight_loader,
394
+ input_dim=0,
395
+ output_dim=1,
396
+ packed_factor=8,
397
+ packed_dim=1,
398
+ ),
399
+ )
400
+
401
+ kernel.process_weights_after_loading(layer)
402
+
403
+ expected_w_kn8 = _pack_int4_along_n(w_int4_kn)
404
+
405
+ assert tuple(layer.qweight.shape) == (K, N // 8)
406
+ assert tuple(layer.scales.shape) == (K // G, N)
407
+ assert tuple(layer.qzeros.shape) == (K // G, N // 8)
408
+
409
+ torch.testing.assert_close(layer.qweight, expected_w_kn8)
410
+ torch.testing.assert_close(layer.scales, scales_gptq_gn)
411
+ torch.testing.assert_close(layer.qzeros, zeros_gptq_gn8)
412
+
413
+
414
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
415
+def test_triton_w4a16_symmetric_apply_ignores_qzeros(monkeypatch):
416
+ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
417
+ MPLinearLayerConfig,
418
+ )
419
+ from vllm.scalar_type import scalar_types
420
+
421
+ K, N, G = 256, 256, 32
422
+ config = MPLinearLayerConfig(
423
+ full_weight_shape=(K, N),
424
+ partition_weight_shape=(K, N),
425
+ weight_type=scalar_types.uint4b8,
426
+ act_type=torch.float16,
427
+ group_size=G,
428
+ zero_points=False,
429
+ has_g_idx=False,
430
+ )
431
+ kernel = TritonW4A16LinearKernel(
432
+ config,
433
+ w_q_param_name="qweight",
434
+ w_s_param_name="scales",
435
+ w_zp_param_name="qzeros",
436
+ w_gidx_param_name=None,
437
+ )
438
+
439
+ class DummyLayer(torch.nn.Module):
440
+ pass
441
+
442
+ layer = DummyLayer()
443
+ layer.qweight = torch.empty((K, N // 8), device=device, dtype=torch.int32)
444
+ layer.scales = torch.empty((K // G, N), device=device, dtype=torch.float16)
445
+ layer.qzeros = torch.empty((1, 1), device=device, dtype=torch.int32)
446
+
447
+ captured = {}
448
+
449
+ def fake_gemm(*, a, b_q, scales, qzeros, group_size, zp_bias):
450
+ captured["qzeros"] = qzeros
451
+ captured["zp_bias"] = zp_bias
452
+ return torch.empty((a.shape[0], N), device=a.device, dtype=a.dtype)
453
+
454
+ monkeypatch.setattr(triton_w4a16_module, "triton_w4a16_gemm", fake_gemm)
455
+
456
+ x = torch.empty((2, K), device=device, dtype=torch.float16)
457
+ output = kernel.apply_weights(layer, x)
458
+
459
+ assert output.shape == (2, N)
460
+ assert captured["qzeros"] is None
461
+ assert captured["zp_bias"] == scalar_types.uint4b8.bias
vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
+25
-5
@@ -400,12 +400,29 @@ 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
- # Checkpoint: [N//8, K//G] int32 (N packed at dim 0, K//G at dim 1)
404
- # Kernel needs: [K//G, N//8] — just transpose
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()
416
+ 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
+
422
replace_parameter(
423
layer,
424
self.w_zp_name,
408
- torch.nn.Parameter(zp.data.t().contiguous(), requires_grad=False),
425
+ torch.nn.Parameter(qzeros, requires_grad=False),
426
)
427
428
def apply_weights(
@@ -420,14 +437,17 @@ class TritonW4A16LinearKernel(MPLinearKernel):
437
K = c.partition_weight_shape[0]
438
group_size = c.group_size if c.group_size != -1 else K
439
423
- # For symmetric types (uint4b8), use the scalar bias; no zeros tensor
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
445
446
output = triton_w4a16_gemm(
447
a=x_2d,
448
b_q=w_q,
449
scales=w_s,
430
- qzeros=w_zp,
450
+ qzeros=qzeros,
451
group_size=group_size,
452
zp_bias=zp_bias,
453
)