From 5281e49908960f490c8faa789eee221d4ab5001d Mon Sep 17 00:00:00 2001
From: 00CC <83797617+tangzzycc@users.noreply.github.com>
Date: Wed, 7 Oct 2026 21:21:43 +0800
Subject: [PATCH] [ROCm][Perf] Enable medium-skinny dispatch in
 RDNAHybridW4A16LinearKernel (#52619)

Signed-off-by: tangzzycc <3081129260@qq.com>
---
 .../quantization/test_rdna_hybrid_w4a16.py    | 81 ++++++++++++++++++-
 .../mixed_precision/rdna_hybrid_w4a16.py      | 16 ++--
 2 files changed, 88 insertions(+), 9 deletions(-)

diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
index d6045a791c..da7ef2d437 100644
--- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
+++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
@@ -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)
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
index 9d53034699..bed62e89af 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
@@ -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 = (