wehub-resource-sync
279 lines
9.7 KiB
Python
279 lines
9.7 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Correctness tests for the ROCm RDNA3 W4A16 GPTQ kernel (gfx1100).
|
|
|
|
Exercises ``RDNA3W4A16LinearKernel`` end-to-end: it builds a layer with
|
|
GPTQ-format checkpoint parameters, runs ``process_weights_after_loading``
|
|
(weight shuffle + zero-point synthesis), then ``apply_weights``, and compares
|
|
the result against an fp32 reference dequant-and-matmul.
|
|
|
|
The kernel is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3`` and is only
|
|
built for gfx11; tests are skipped elsewhere.
|
|
|
|
Run `pytest tests/kernels/quantization/test_rdna3_w4a16.py`.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.platforms import current_platform
|
|
|
|
if not current_platform.is_rocm():
|
|
pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
|
|
|
|
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402
|
|
MPLinearLayerConfig,
|
|
)
|
|
from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402
|
|
RDNA3W4A16LinearKernel,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
|
|
pack_quantized_values_into_int32,
|
|
)
|
|
from vllm.model_executor.parameter import ( # noqa: E402
|
|
GroupQuantScaleParameter,
|
|
PackedvLLMParameter,
|
|
)
|
|
from vllm.platforms.rocm import on_gfx1100 # noqa: E402
|
|
from vllm.scalar_type import scalar_types # noqa: E402
|
|
from vllm.utils.torch_utils import set_random_seed # noqa: E402
|
|
|
|
device = "cuda"
|
|
|
|
WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8
|
|
PACK_FACTOR = 8 # 8 x 4-bit nibbles per int32
|
|
|
|
# Skip everything in this module unless we are on the only architecture the
|
|
# kernel is built/registered for.
|
|
gfx1100_only = pytest.mark.skipif(
|
|
not (
|
|
on_gfx1100()
|
|
and hasattr(torch.ops, "_rocm_C")
|
|
and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
|
|
),
|
|
reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Reference implementation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _reference(
|
|
x_mk: torch.Tensor,
|
|
q_int4_kn: torch.Tensor,
|
|
scales_gn: torch.Tensor,
|
|
zeros_gn: torch.Tensor | None,
|
|
group_size: int,
|
|
bias: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
"""fp32 reference for the RDNA3 W4A16 op.
|
|
|
|
x_mk: [M, K] fp16/bf16 activations.
|
|
q_int4_kn: [K, N] int32 raw stored nibbles in [0, 15].
|
|
scales_gn: [K//G, N] per-group scales (act dtype).
|
|
zeros_gn: [K//G, N] int32 raw stored zero points in [0, 15], or None
|
|
for the symmetric path (kernel synthesizes stored zero = 7).
|
|
group_size: G.
|
|
|
|
The kernel applies the GPTQv1 "+1" zero-point quirk, so the effective
|
|
zero is ``stored_zero + 1`` (symmetric path: 7 + 1 == bias == 8).
|
|
"""
|
|
K, N = q_int4_kn.shape
|
|
s_full = scales_gn.repeat_interleave(group_size, dim=0).to(torch.float32) # [K,N]
|
|
if zeros_gn is None:
|
|
z_full = torch.full(
|
|
(K, N), float(WEIGHT_TYPE.bias), device=x_mk.device, dtype=torch.float32
|
|
)
|
|
else:
|
|
z_full = (zeros_gn + 1).repeat_interleave(group_size, dim=0).to(torch.float32)
|
|
w_fp = (q_int4_kn.to(torch.float32) - z_full) * s_full # [K, N]
|
|
out = x_mk.to(torch.float32) @ w_fp # [M, N]
|
|
if bias is not None:
|
|
out = out + bias.to(torch.float32)
|
|
return out.to(x_mk.dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Layer construction (GPTQ checkpoint format)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _build_layer(
|
|
q_int4_kn: torch.Tensor,
|
|
scales_gn: torch.Tensor,
|
|
zeros_gn: torch.Tensor | None,
|
|
dtype: torch.dtype,
|
|
) -> torch.nn.Module:
|
|
"""Build a dummy layer carrying GPTQ-format params, as the loader would."""
|
|
no_loader = lambda *args, **kwargs: None # noqa: E731
|
|
|
|
# qweight: int4 packed along K into int32 -> [K//8, N].
|
|
qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0)
|
|
|
|
class DummyLayer(torch.nn.Module):
|
|
pass
|
|
|
|
layer = DummyLayer()
|
|
layer.register_parameter(
|
|
"qweight",
|
|
PackedvLLMParameter(
|
|
data=qweight,
|
|
weight_loader=no_loader,
|
|
input_dim=0,
|
|
output_dim=1,
|
|
packed_dim=0,
|
|
packed_factor=PACK_FACTOR,
|
|
),
|
|
)
|
|
layer.register_parameter(
|
|
"scales",
|
|
GroupQuantScaleParameter(
|
|
data=scales_gn.to(dtype),
|
|
weight_loader=no_loader,
|
|
input_dim=0,
|
|
output_dim=1,
|
|
),
|
|
)
|
|
if zeros_gn is not None:
|
|
# qzeros: int4 packed along N into int32 -> [K//G, N//8].
|
|
qzeros = pack_quantized_values_into_int32(zeros_gn, WEIGHT_TYPE, packed_dim=1)
|
|
layer.register_parameter(
|
|
"qzeros",
|
|
PackedvLLMParameter(
|
|
data=qzeros,
|
|
weight_loader=no_loader,
|
|
input_dim=0,
|
|
output_dim=1,
|
|
packed_dim=1,
|
|
packed_factor=PACK_FACTOR,
|
|
),
|
|
)
|
|
return layer
|
|
|
|
|
|
def _run_kernel(
|
|
x_mk: torch.Tensor,
|
|
q_int4_kn: torch.Tensor,
|
|
scales_gn: torch.Tensor,
|
|
zeros_gn: torch.Tensor | None,
|
|
group_size: int,
|
|
bias: torch.Tensor | None,
|
|
dtype: torch.dtype,
|
|
) -> torch.Tensor:
|
|
K, N = q_int4_kn.shape
|
|
has_zp = zeros_gn is not None
|
|
|
|
config = MPLinearLayerConfig(
|
|
full_weight_shape=(K, N),
|
|
partition_weight_shape=(K, N),
|
|
weight_type=WEIGHT_TYPE,
|
|
act_type=dtype,
|
|
group_size=group_size,
|
|
zero_points=has_zp,
|
|
has_g_idx=False,
|
|
)
|
|
ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
|
|
assert ok, f"can_implement rejected a supported config: {reason}"
|
|
|
|
layer = _build_layer(q_int4_kn, scales_gn, zeros_gn, dtype)
|
|
kernel = RDNA3W4A16LinearKernel(
|
|
config,
|
|
w_q_param_name="qweight",
|
|
w_s_param_name="scales",
|
|
w_zp_param_name="qzeros" if has_zp else None,
|
|
w_gidx_param_name=None,
|
|
)
|
|
kernel.process_weights_after_loading(layer)
|
|
return kernel.apply_weights(layer, x_mk, bias=bias)
|
|
|
|
|
|
# Relative-L2 tolerance per dtype. The bf16 path widens dequantized weights
|
|
# to fp32 and accumulates in fp32, so it matches the reference almost exactly
|
|
# (<0.4% incl. the WMMA prefill path). The fp16 path uses the exllamav2
|
|
# "+1024" bit-trick (see qdq_4_rdna3.cuh): the dequantized weight is recovered
|
|
# as the fp16 difference of two ~1024*scale magnitudes, which sheds low-order
|
|
# mantissa bits and leaves ~2-3% relative noise that accumulates over K. We
|
|
# compare on the relative Frobenius norm rather than elementwise, since the
|
|
# bit-trick noise produces large *relative* errors on individual near-zero
|
|
# outputs that carry negligible absolute weight.
|
|
_REL_L2_TOL = {torch.float16: 5e-2, torch.bfloat16: 1e-2}
|
|
|
|
|
|
def _assert_close(out: torch.Tensor, ref: torch.Tensor, dtype: torch.dtype):
|
|
rel_l2 = (out.to(torch.float32) - ref.to(torch.float32)).norm() / ref.to(
|
|
torch.float32
|
|
).norm()
|
|
tol = _REL_L2_TOL[dtype]
|
|
assert rel_l2 < tol, f"relative L2 error {rel_l2:.4f} exceeds {tol} for {dtype}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Forward correctness
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
# (M, K, N, group_size). M spans the scalar decode path (small M) and the
|
|
# WMMA prefill path (M >= 16 on the bf16 dispatch). K/N satisfy the kernel's
|
|
# divisibility constraints (K % G == 0, K % 8 == 0, N % 8 == 0).
|
|
MKNG_SHAPES = [
|
|
(1, 128, 128, 128), # single group, decode
|
|
(2, 256, 256, 128), # two groups
|
|
(8, 256, 512, 64), # M=8 scalar, smaller group
|
|
(16, 512, 256, 128), # M=16 -> WMMA path for bf16
|
|
(32, 512, 512, 64), # larger prefill
|
|
]
|
|
|
|
|
|
@gfx1100_only
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("has_zp", [False, True], ids=["no_zp", "with_zp"])
|
|
@pytest.mark.parametrize(
|
|
"M,K,N,G", MKNG_SHAPES, ids=[f"m{m}_k{k}_n{n}_g{g}" for m, k, n, g in MKNG_SHAPES]
|
|
)
|
|
def test_rdna3_w4a16_matches_reference(dtype, has_zp, M, K, N, G, dist_init):
|
|
set_random_seed(0)
|
|
assert K % G == 0 and K % PACK_FACTOR == 0 and N % PACK_FACTOR == 0
|
|
|
|
groups = K // G
|
|
x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
|
q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
|
scales_gn = (
|
|
0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
|
|
).to(dtype)
|
|
zeros_gn = (
|
|
torch.randint(0, 16, (groups, N), device=device, dtype=torch.int32)
|
|
if has_zp
|
|
else None
|
|
)
|
|
|
|
out = _run_kernel(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None, dtype)
|
|
ref = _reference(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None)
|
|
|
|
assert out.shape == (M, N) and out.dtype == dtype
|
|
_assert_close(out, ref, dtype)
|
|
|
|
|
|
@gfx1100_only
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("M", [1, 32], ids=["decode", "prefill"])
|
|
def test_rdna3_w4a16_bias(dtype, M, dist_init):
|
|
"""Bias is added on both the scalar (M=1) and WMMA (M=32) paths."""
|
|
set_random_seed(0)
|
|
K, N, G = 512, 256, 128
|
|
groups = K // G
|
|
|
|
x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
|
q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
|
scales_gn = (
|
|
0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
|
|
).to(dtype)
|
|
bias = (0.1 * torch.randn(N, device=device, dtype=torch.float32)).to(dtype)
|
|
|
|
out = _run_kernel(x_mk, q_int4_kn, scales_gn, None, G, bias, dtype)
|
|
ref = _reference(x_mk, q_int4_kn, scales_gn, None, G, bias)
|
|
|
|
_assert_close(out, ref, dtype)
|