forked from Karylab-cklius/vllm
382bbd5144
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
526 lines
18 KiB
Python
526 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for the ROCm Hybrid W4A16 kernel (HIP skinny + Triton prefill).
|
|
|
|
Run `pytest tests/kernels/quantization/test_rdna_hybrid_w4a16.py`.
|
|
"""
|
|
|
|
import importlib
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.torch_utils import set_random_seed
|
|
|
|
if not current_platform.is_rocm():
|
|
pytest.skip("ROCm only", allow_module_level=True)
|
|
|
|
pytest.importorskip("triton")
|
|
|
|
from vllm.platforms.rocm import on_gfx1x # noqa: E402
|
|
|
|
device = "cuda"
|
|
|
|
hybrid_module = importlib.import_module(
|
|
"vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16"
|
|
)
|
|
RDNAHybridW4A16LinearKernel = hybrid_module.RDNAHybridW4A16LinearKernel
|
|
pack_int4_exllama_shuffle = hybrid_module.pack_int4_exllama_shuffle
|
|
SUPPORTED_GROUP_SIZES = hybrid_module.SUPPORTED_GROUP_SIZES
|
|
MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Reference implementation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _rdna_hybrid_w4a16_reference(
|
|
x_mk: torch.Tensor,
|
|
w_int4_nk: torch.Tensor,
|
|
scales_nkg: torch.Tensor,
|
|
zp_nkg: torch.Tensor | None,
|
|
group_size: int,
|
|
bias: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
"""Reference for the Hybrid W4A16 op.
|
|
|
|
x_mk: [M, K] fp16/bf16
|
|
w_int4_nk: [N, K] int32 with raw uint4 values in [0, 15]
|
|
scales_nkg: [N, K//G] fp16/bf16
|
|
zp_nkg: [N, K//G] fp16/bf16 raw zero points (already in act dtype),
|
|
or None for symmetric (uint4b8, dequant subtracts 8)
|
|
"""
|
|
G = group_size
|
|
N, K = w_int4_nk.shape
|
|
assert K % G == 0
|
|
s_full = scales_nkg.repeat_interleave(G, dim=1).to(torch.float32) # [N, K]
|
|
if zp_nkg is None:
|
|
z_full = torch.full((N, K), 8.0, device=x_mk.device, dtype=torch.float32)
|
|
else:
|
|
z_full = zp_nkg.repeat_interleave(G, dim=1).to(torch.float32)
|
|
w_fp = (w_int4_nk.to(torch.float32) - z_full) * s_full # [N, K]
|
|
out = x_mk.to(torch.float32) @ w_fp.t() # [M, N]
|
|
if bias is not None:
|
|
out = out + bias.to(torch.float32)
|
|
return out.to(x_mk.dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Forward correctness: decode (HIP skinny) + prefill (Triton)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)
|
|
@pytest.mark.parametrize("has_zp", [False, True])
|
|
@pytest.mark.parametrize(
|
|
"M",
|
|
[1, MAX_SKINNY_BATCH_SIZE, MAX_SKINNY_BATCH_SIZE + 1, 64],
|
|
ids=["M=1_decode", "M=5_decode", "M=6_prefill", "M=64_prefill"],
|
|
)
|
|
def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M):
|
|
"""Smoke test the registered custom op for both decode and prefill batches.
|
|
|
|
Verifies the dispatch logic in `_rdna_hybrid_w4a16_apply_impl`:
|
|
- M <= MAX_SKINNY_BATCH_SIZE: HIP wvSplitK_int4_g
|
|
- M > MAX_SKINNY_BATCH_SIZE: Triton prefill kernel
|
|
"""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA/HIP device not available")
|
|
|
|
set_random_seed(0)
|
|
|
|
K, N = 1024, 256
|
|
assert K % group_size == 0 and K % 8 == 0 and N % 8 == 0
|
|
|
|
# Activations.
|
|
x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
|
|
|
# Weights as raw uint4 in [N, K], packed to ExLlama shuffle [N, K//8].
|
|
w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
|
|
w_q_i32 = pack_int4_exllama_shuffle(w_int4_nk).contiguous() # [N, K//8] int32
|
|
w_q = w_q_i32.view(torch.int8) # same bytes viewed as int8 [N, K//2]
|
|
|
|
# Scales [N, K//G] in act dtype.
|
|
scales_nkg = (
|
|
0.05 * torch.rand((N, K // group_size), device=device, dtype=torch.float32)
|
|
).to(dtype)
|
|
|
|
# Optional raw zero points [N, K//G] in act dtype.
|
|
if has_zp:
|
|
zp_nkg = torch.randint(
|
|
0, 16, (N, K // group_size), device=device, dtype=torch.int32
|
|
).to(dtype)
|
|
else:
|
|
zp_nkg = None
|
|
|
|
from vllm.utils.platform_utils import num_compute_units
|
|
|
|
out = torch.ops.vllm.rdna_hybrid_w4a16_apply(
|
|
x_mk,
|
|
w_q,
|
|
scales_nkg,
|
|
zp_nkg,
|
|
None, # bias
|
|
num_compute_units(),
|
|
group_size,
|
|
)
|
|
|
|
ref = _rdna_hybrid_w4a16_reference(
|
|
x_mk, w_int4_nk, scales_nkg, zp_nkg, group_size, bias=None
|
|
)
|
|
|
|
torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
|
|
|
|
|
|
@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("M", [1, MAX_SKINNY_BATCH_SIZE + 1])
|
|
def test_rdna_hybrid_w4a16_apply_with_bias(dtype, M):
|
|
"""Bias is added correctly on both decode and prefill paths."""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA/HIP device not available")
|
|
|
|
set_random_seed(0)
|
|
K, N, G = 1024, 128, 128
|
|
|
|
x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
|
w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
|
|
w_q_i32 = pack_int4_exllama_shuffle(w_int4_nk).contiguous()
|
|
w_q = w_q_i32.view(torch.int8)
|
|
scales_nkg = (
|
|
0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)
|
|
).to(dtype)
|
|
bias = torch.randn(N, device=device, dtype=dtype) * 0.1
|
|
|
|
from vllm.utils.platform_utils import num_compute_units
|
|
|
|
out = torch.ops.vllm.rdna_hybrid_w4a16_apply(
|
|
x_mk,
|
|
w_q,
|
|
scales_nkg,
|
|
None,
|
|
bias,
|
|
num_compute_units(),
|
|
G,
|
|
)
|
|
ref = _rdna_hybrid_w4a16_reference(x_mk, w_int4_nk, scales_nkg, None, G, bias=bias)
|
|
|
|
torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# pack_int4_exllama_shuffle round-trips correctly
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_pack_int4_exllama_shuffle_layout():
|
|
"""Pack 8 K-values per int32 in interleave [0,2,4,6,1,3,5,7] order."""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA/HIP device not available")
|
|
set_random_seed(0)
|
|
N, K = 4, 16
|
|
w = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
|
|
packed = pack_int4_exllama_shuffle(w)
|
|
assert packed.shape == (N, K // 8) and packed.dtype == torch.int32
|
|
|
|
# Manual unshuffle using ExLlama shifts [0,16,4,20,8,24,12,28].
|
|
shifts = torch.tensor(
|
|
[0, 16, 4, 20, 8, 24, 12, 28], device=device, dtype=torch.int32
|
|
)
|
|
unshuffled = (packed.unsqueeze(-1) >> shifts) & 0xF
|
|
unshuffled = unshuffled.reshape(N, K)
|
|
torch.testing.assert_close(unshuffled, w)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# process_weights_after_loading: layout repack and zp normalization
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _pack_int4_along_k_to_ckpt(w_int4_kn: torch.Tensor) -> torch.Tensor:
|
|
"""Pack int4 values along K into CT checkpoint layout: [K,N] -> [N, K//8]."""
|
|
assert w_int4_kn.dtype == torch.int32
|
|
K, N = w_int4_kn.shape
|
|
assert K % 8 == 0
|
|
out = torch.zeros((N, K // 8), dtype=torch.int32, device=w_int4_kn.device)
|
|
for i in range(8):
|
|
out |= (w_int4_kn[i::8, :].t() & 0xF) << (i * 4)
|
|
return out.contiguous()
|
|
|
|
|
|
def _pack_int4_along_n_for_zp(zp_int4_gn: torch.Tensor) -> torch.Tensor:
|
|
"""Pack int4 zero points along N: [G, N] -> [G, N//8] int32 (CT layout)."""
|
|
assert zp_int4_gn.dtype == torch.int32
|
|
G, N = zp_int4_gn.shape
|
|
assert N % 8 == 0
|
|
shifts = torch.arange(8, device=zp_int4_gn.device, dtype=torch.int32) * 4
|
|
return torch.sum(
|
|
(zp_int4_gn.view(G, N // 8, 8) & 0xF) << shifts, dim=2, dtype=torch.int32
|
|
).contiguous()
|
|
|
|
|
|
def _build_dummy_layer(
|
|
w_ckpt_nk8: torch.Tensor,
|
|
scales_ckpt_nkg: torch.Tensor,
|
|
zeros_ckpt: torch.Tensor | None,
|
|
):
|
|
from vllm.model_executor.parameter import (
|
|
GroupQuantScaleParameter,
|
|
PackedColumnParameter,
|
|
PackedvLLMParameter,
|
|
)
|
|
|
|
weight_loader = lambda *args, **kwargs: None
|
|
|
|
class DummyLayer(torch.nn.Module):
|
|
pass
|
|
|
|
layer = DummyLayer()
|
|
layer.register_parameter(
|
|
"weight_packed",
|
|
PackedvLLMParameter(
|
|
data=w_ckpt_nk8,
|
|
weight_loader=weight_loader,
|
|
input_dim=1,
|
|
output_dim=0,
|
|
packed_factor=8,
|
|
packed_dim=1,
|
|
),
|
|
)
|
|
layer.register_parameter(
|
|
"weight_scale",
|
|
GroupQuantScaleParameter(
|
|
data=scales_ckpt_nkg,
|
|
weight_loader=weight_loader,
|
|
input_dim=1,
|
|
output_dim=0,
|
|
),
|
|
)
|
|
if zeros_ckpt is not None:
|
|
layer.register_parameter(
|
|
"weight_zero_point",
|
|
PackedColumnParameter(
|
|
data=zeros_ckpt,
|
|
weight_loader=weight_loader,
|
|
output_dim=0,
|
|
packed_factor=8,
|
|
packed_dim=0,
|
|
),
|
|
)
|
|
return layer
|
|
|
|
|
|
@pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)
|
|
def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_init):
|
|
"""uint4b8 (symmetric): w_q -> [N, K//8] int8 ExLlama shuffle, no zp param."""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA/HIP device not available")
|
|
|
|
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
|
|
MPLinearLayerConfig,
|
|
)
|
|
from vllm.scalar_type import scalar_types
|
|
|
|
set_random_seed(0)
|
|
|
|
K, N = 256, 128
|
|
G = group_size
|
|
assert K % G == 0
|
|
|
|
# Reference unpacked weights, then pack into CT checkpoint layout [N, K//8].
|
|
w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
|
w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn)
|
|
scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16)
|
|
|
|
layer = _build_dummy_layer(w_ckpt_nk8, scales_ckpt_nkg, zeros_ckpt=None)
|
|
|
|
config = MPLinearLayerConfig(
|
|
full_weight_shape=(K, N),
|
|
partition_weight_shape=(K, N),
|
|
weight_type=scalar_types.uint4b8,
|
|
act_type=torch.float16,
|
|
group_size=G,
|
|
zero_points=False,
|
|
has_g_idx=False,
|
|
)
|
|
kernel = RDNAHybridW4A16LinearKernel(
|
|
config,
|
|
w_q_param_name="weight_packed",
|
|
w_s_param_name="weight_scale",
|
|
w_zp_param_name=None,
|
|
w_gidx_param_name=None,
|
|
)
|
|
kernel.process_weights_after_loading(layer)
|
|
|
|
# Skinny weight is stored once as int8 [N, K//2]; the Triton path
|
|
# reinterprets it as int32 [N, K//8] via a view (no separate parameter).
|
|
assert layer.weight_packed.dtype == torch.int8
|
|
assert tuple(layer.weight_packed.shape) == (N, K // 2)
|
|
w_q_i32 = layer.weight_packed.view(torch.int32)
|
|
assert tuple(w_q_i32.shape) == (N, K // 8)
|
|
|
|
expected_packed = pack_int4_exllama_shuffle(w_int4_kn.t().contiguous())
|
|
torch.testing.assert_close(w_q_i32, expected_packed)
|
|
|
|
# Scales: [N, K//G] (skinny layout, no transpose since CT already had it).
|
|
assert tuple(layer.weight_scale.shape) == (N, K // G)
|
|
torch.testing.assert_close(layer.weight_scale, scales_ckpt_nkg)
|
|
|
|
|
|
@pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)
|
|
def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_init):
|
|
"""uint4 (asymmetric): zero points unpacked to raw values in act dtype."""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA/HIP device not available")
|
|
|
|
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
|
|
MPLinearLayerConfig,
|
|
)
|
|
from vllm.scalar_type import scalar_types
|
|
|
|
set_random_seed(0)
|
|
|
|
K, N = 256, 128
|
|
G = group_size
|
|
assert K % G == 0 and N % 8 == 0
|
|
|
|
w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
|
w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn)
|
|
scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16)
|
|
|
|
# CT zero-point layout is N-packed: [N//8, K//G] int32.
|
|
zeros_int4_gn = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32)
|
|
zeros_packed_gn8 = _pack_int4_along_n_for_zp(zeros_int4_gn) # [K//G, N//8]
|
|
zeros_ckpt_n8kg = zeros_packed_gn8.t().contiguous() # [N//8, K//G]
|
|
|
|
layer = _build_dummy_layer(w_ckpt_nk8, scales_ckpt_nkg, zeros_ckpt=zeros_ckpt_n8kg)
|
|
|
|
config = MPLinearLayerConfig(
|
|
full_weight_shape=(K, N),
|
|
partition_weight_shape=(K, N),
|
|
weight_type=scalar_types.uint4,
|
|
act_type=torch.float16,
|
|
group_size=G,
|
|
zero_points=True,
|
|
has_g_idx=False,
|
|
)
|
|
kernel = RDNAHybridW4A16LinearKernel(
|
|
config,
|
|
w_q_param_name="weight_packed",
|
|
w_s_param_name="weight_scale",
|
|
w_zp_param_name="weight_zero_point",
|
|
w_gidx_param_name=None,
|
|
)
|
|
kernel.process_weights_after_loading(layer)
|
|
|
|
# Zero-points: unpacked to [N, K//G], cast to act dtype, raw values [0..15].
|
|
assert layer.weight_zero_point.dtype == torch.float16
|
|
assert tuple(layer.weight_zero_point.shape) == (N, K // G)
|
|
expected_zp = zeros_int4_gn.t().to(torch.float16) # [N, K//G]
|
|
torch.testing.assert_close(layer.weight_zero_point, expected_zp)
|
|
|
|
# Quantized weights match symmetric path's layout regardless of zp.
|
|
w_q_i32 = layer.weight_packed.view(torch.int32)
|
|
assert tuple(w_q_i32.shape) == (N, K // 8)
|
|
expected_packed = pack_int4_exllama_shuffle(w_int4_kn.t().contiguous())
|
|
torch.testing.assert_close(w_q_i32, expected_packed)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# can_implement enforces the supported-group-size policy
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
|
|
@pytest.mark.parametrize(
|
|
"group_size,expected_ok", [(32, True), (64, True), (128, True), (256, False)]
|
|
)
|
|
def test_hybrid_can_implement_group_size(group_size, expected_ok):
|
|
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
|
|
MPLinearLayerConfig,
|
|
)
|
|
from vllm.scalar_type import scalar_types
|
|
|
|
K, N = 1024, 256
|
|
config = MPLinearLayerConfig(
|
|
full_weight_shape=(K, N),
|
|
partition_weight_shape=(K, N),
|
|
weight_type=scalar_types.uint4b8,
|
|
act_type=torch.float16,
|
|
group_size=group_size,
|
|
zero_points=False,
|
|
has_g_idx=False,
|
|
)
|
|
ok, _ = RDNAHybridW4A16LinearKernel.can_implement(config)
|
|
assert ok is expected_ok
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests for the HIP wvSplitK_int4_g decode kernel
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _hip_skinny_reference(
|
|
a_mk: torch.Tensor,
|
|
w_int4_nk: torch.Tensor,
|
|
scales_nkg: torch.Tensor,
|
|
*,
|
|
group_size: int,
|
|
zp_bias: int,
|
|
) -> torch.Tensor:
|
|
"""Reference for symmetric HIP skinny: C = A @ (W - zp_bias) * S."""
|
|
K = a_mk.shape[1]
|
|
N = w_int4_nk.shape[0]
|
|
num_groups = K // group_size
|
|
|
|
w_fp = (w_int4_nk.to(torch.float32) - zp_bias).view(N, num_groups, group_size)
|
|
s = scales_nkg.to(torch.float32).unsqueeze(-1)
|
|
w_dequant = (w_fp * s).view(N, K)
|
|
|
|
return (a_mk.to(torch.float32) @ w_dequant.t()).to(a_mk.dtype)
|
|
|
|
|
|
@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize(
|
|
"M,K,N,G",
|
|
[
|
|
(1, 256, 256, 32),
|
|
(1, 256, 256, 64),
|
|
(1, 512, 256, 128),
|
|
(2, 512, 256, 64),
|
|
(3, 256, 512, 64),
|
|
],
|
|
)
|
|
def test_hip_skinny_wvSplitK_int4_g(dtype, M, K, N, G):
|
|
"""Test HIP wvSplitK_int4_g kernel directly via _custom_ops."""
|
|
import vllm._custom_ops as ops
|
|
from vllm.utils.platform_utils import num_compute_units
|
|
|
|
set_random_seed(0)
|
|
|
|
a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
|
w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
|
|
|
|
b_packed_i32 = pack_int4_exllama_shuffle(w_int4_nk)
|
|
b_packed_i8 = b_packed_i32.view(torch.int8)
|
|
|
|
scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to(
|
|
dtype
|
|
)
|
|
|
|
cu_count = num_compute_units()
|
|
out = ops.wvSplitK_int4_g(b_packed_i8, a, scales, cu_count, G)
|
|
|
|
ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8)
|
|
|
|
torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests for the full hybrid dispatch (HIP decode + Triton prefill)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize(
|
|
"M,K,N,G",
|
|
[
|
|
(1, 256, 256, 64),
|
|
(1, 512, 256, 32),
|
|
(1, 512, 256, 128),
|
|
(32, 512, 256, 64),
|
|
(64, 1024, 256, 128),
|
|
],
|
|
)
|
|
def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G):
|
|
"""Test the full hybrid dispatch via the custom op."""
|
|
from vllm.utils.platform_utils import num_compute_units
|
|
|
|
set_random_seed(0)
|
|
|
|
a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
|
w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
|
|
|
|
b_packed_i32 = pack_int4_exllama_shuffle(w_int4_nk)
|
|
b_packed_i8 = b_packed_i32.view(torch.int8)
|
|
|
|
scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to(
|
|
dtype
|
|
)
|
|
|
|
cu_count = num_compute_units()
|
|
out = torch.ops.vllm.rdna_hybrid_w4a16_apply(
|
|
a, b_packed_i8, scales, None, None, cu_count, G
|
|
)
|
|
|
|
ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8)
|
|
|
|
torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2)
|