iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] Enable medium-skinny dispatch in RDNAHybridW4A16LinearKernel (#52619)

Browse Source 
Signed-off-by: tangzzycc <3081129260@qq.com>
This commit is contained in:
00CC
2026-10-07 13:21:43 +00:00
committed by GitHub
parent efb8d85f24
commit 5281e49908
2 changed files with 88 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/quantization/test_rdna_hybrid_w4a16.py
+79 -2
View File
@@ -30,6 +30,8 @@ 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
LDS_CAPACITY_ELEMENTS = hybrid_module.LDS_CAPACITY_ELEMENTS
MEDIUM_SKINNY_LIMIT_ELEMENTS = hybrid_module.MEDIUM_SKINNY_LIMIT_ELEMENTS
# ---------------------------------------------------------------------------
@@ -100,8 +102,8 @@ 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
- supported M and K*M: HIP wvSplitK_int4_g
- otherwise: Triton prefill kernel
"""
if not torch.cuda.is_available():
pytest.skip("CUDA/HIP device not available")
@@ -742,3 +744,78 @@ def test_process_weights_pads_cliff_rows(K, expected_weight_stride, dist_init):
assert tuple(layer.weight_packed.view(torch.int32).shape) == (N, K // 8)
# Metadata is untouched by this change.
assert layer.weight_scale.is_contiguous()
@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
@pytest.mark.parametrize(
"M,K,expected_path",
[
(4, 8192, "hip"),
(4, 9728, "hip"),
(4, 9856, "triton"),
(MAX_SKINNY_BATCH_SIZE + 1, 1024, "triton"),
],
ids=["regular_limit", "medium_range", "above_medium", "batch_too_large"],
)
def test_rdna_hybrid_w4a16_dispatch_boundaries(M, K, expected_path, monkeypatch):
"""Route only supported regular and medium skinny shapes to HIP."""
import vllm._custom_ops as ops
N, G = 16, 128
x_mk = torch.empty((M, K), dtype=torch.float16)
w_q = torch.empty((N, K // 2), dtype=torch.int8)
scales = torch.empty((N, K // G), dtype=torch.float16)
called = []
def fake_hip(*args, **kwargs):
called.append("hip")
return torch.empty((M, N), dtype=x_mk.dtype)
def fake_triton(*args, **kwargs):
called.append("triton")
return torch.empty((M, N), dtype=x_mk.dtype)
monkeypatch.setattr(ops, "wvSplitK_int4_g", fake_hip)
monkeypatch.setattr(hybrid_module, "triton_w4a16_skinny_fmt_gemm", fake_triton)
output = hybrid_module._rdna_hybrid_w4a16_apply_impl(
x_mk, w_q, scales, None, None, 40, G
)
assert output.shape == (M, N)
assert called == [expected_path]
@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("has_zp", [False, True])
def test_rdna_hybrid_w4a16_medium_skinny_matches_reference(dtype, has_zp):
"""Validate the HIP medium path against an FP32 dequantization reference."""
if not torch.cuda.is_available():
pytest.skip("CUDA/HIP device not available")
from vllm.utils.platform_utils import num_compute_units
set_random_seed(0)
M, K, N, G = 4, 9728, 256, 128
assert LDS_CAPACITY_ELEMENTS < K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS
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 = pack_int4_exllama_shuffle(w_int4_nk).contiguous().view(torch.int8)
scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to(
dtype
)
zp = (
torch.randint(0, 16, (N, K // G), device=device, dtype=torch.int32)
if has_zp
else None
)
w_zp = _pack_zp_rows_for_kernel(zp) if zp is not None else None
output = torch.ops.vllm.rdna_hybrid_w4a16_apply(
x_mk, w_q, scales, w_zp, None, num_compute_units(), G
)
reference = _rdna_hybrid_w4a16_reference(x_mk, w_int4_nk, scales, zp, G, bias=None)
torch.testing.assert_close(output, reference, rtol=2e-2, atol=2e-2)
vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
+9 -7
View File
@@ -2,9 +2,10 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode.
Routes based on batch size M:
M <= MAX_SKINNY_BATCH_SIZE: HIP skinny GEMM (wvSplitK_int4_g)
M > MAX_SKINNY_BATCH_SIZE: Triton W4A16 fused dequant GEMM
Routes based on the activation shape:
M <= MAX_SKINNY_BATCH_SIZE and K*M <= MEDIUM_SKINNY_LIMIT_ELEMENTS:
HIP skinny GEMM (wvSplitK_int4_g)
otherwise: Triton W4A16 fused dequant GEMM
Stores the weights ONCE as int8 [N, K//2] (ExLlama shuffle packed). Both
paths read this single buffer: the HIP skinny kernel uses it directly, and
@@ -56,13 +57,14 @@ def _on_gfx1151() -> bool:
return on_gfx1151()
# Maximum batch size M for the HIP skinny kernel path (C++ supports N_in
# up to 5). When M is below this AND K*M fits in LDS, the skinny kernel is
# used; otherwise the Triton prefill path handles the GEMM.
# Maximum activation rows supported by HIP skinny (Python M maps to C++ N_in).
# When K*M exceeds regular LDS capacity, the C++ medium kernel caches the
# prefix in LDS and reads the remaining activation elements from global memory.
MAX_SKINNY_BATCH_SIZE = 5
# 64 KiB per-workgroup LDS limit expressed in fp16 elements.
# (AMD RDNA has 128 KiB total LDS per CU, but 64 KiB per workgroup.)
LDS_CAPACITY_ELEMENTS = 64 * 1024 // 2 # 32768 fp16 elements
MEDIUM_SKINNY_LIMIT_ELEMENTS = int(LDS_CAPACITY_ELEMENTS * 1.2)
# ---------------------------------------------------------------------------
@@ -464,7 +466,7 @@ def _rdna_hybrid_w4a16_apply_impl(
K = x_2d.shape[1]
N = w_q.shape[0]
if M <= MAX_SKINNY_BATCH_SIZE and K * M <= LDS_CAPACITY_ELEMENTS:
if M <= MAX_SKINNY_BATCH_SIZE and K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS:
# record_function is not torch.compile-safe; use nullcontext when
# compiling to keep the op traceable.
ctx = (