[ROCm] Converge the two W4A16 linear kernels onto the upstream shell

benchmarks/kernels/sweep_int4g_kernel.py

@@ -23,7 +23,7 @@
 import torch
 
 import vllm._custom_ops as ops
-from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
+from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
     _group_stride_pad,
     _pad_group_rows,
     pack_skinny_int4,

tests/kernels/moe/test_hybrid_w4a16_moe.py

@@ -20,7 +20,7 @@
 from tests.kernels.moe.utils import make_dummy_moe_config
 from tests.kernels.utils import torch_experts
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
+from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
     pack_int4_exllama_shuffle,
 )
 from vllm.model_executor.layers.fused_moe import fused_topk

tests/kernels/quantization/test_hip_w4a16.py

@@ -596,7 +596,7 @@ def test_wvsplitk_int4_g_packed_zero_points(
     plausible-looking output, so this compares actual values.
     """
     import vllm._custom_ops as ops
-    from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
+    from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
         pack_skinny_int4,
     )
     from vllm.utils.platform_utils import num_compute_units
@@ -636,7 +636,7 @@ def test_wvsplitk_int4_g_rejects_unpacked_zero_points() -> None:
     silently accepting that shape now would read garbage nibbles.
     """
     import vllm._custom_ops as ops
-    from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
+    from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
         pack_skinny_int4,
     )
     from vllm.utils.platform_utils import num_compute_units
@@ -673,7 +673,7 @@ def test_wvsplitk_int4_g_padded_group_stride(k: int, group_size: int) -> None:
     is exercised regardless.
     """
     import vllm._custom_ops as ops
-    from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
+    from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
         pack_skinny_int4,
     )
     from vllm.utils.platform_utils import num_compute_units

tests/kernels/quantization/test_hybrid_w4a16_perf.py

@@ -456,7 +456,7 @@ def prepare_hybrid_weights(
     Scales/zp follow the activation *dtype* so the kernel exercises the
     same fp16 vs bf16 code path it takes in production.
     """
-    from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
+    from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
         _group_stride_pad,
         _pad_group_rows,
         pack_skinny_int4,
@@ -499,7 +499,7 @@ def _compute_packed_scale_zp(
     dtype: torch.dtype,
 ) -> torch.Tensor | None:
     """Pack the per-group packed scale/zp carrier into one fp32, matching the load-time
-    carrier built in ``HybridW4A16LinearKernel.process_weights_after_loading``.
+    carrier built in ``RDNAHybridW4A16LinearKernel.process_weights_after_loading``.
 
     Built ONLY for asymmetric layers (``w_zp`` given); returns None for symmetric
     (the kernel uses the constant -8 offset there, no carrier). Low 16 bits =
@@ -591,8 +591,8 @@ def measure_tflops(
     (hot), matching a real prefill where activations are freshly produced and
     weights stream from HBM.
     """
-    from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
-        _hybrid_w4a16_apply_impl,
+    from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
+        _rdna_hybrid_w4a16_apply_impl,
     )
     from vllm.triton_utils import triton
     from vllm.utils.platform_utils import num_compute_units
@@ -634,7 +634,7 @@ def measure_tflops(
     def run():
         w = bufs[idx[0] % n_buf]
         idx[0] += 1
-        return _hybrid_w4a16_apply_impl(
+        return _rdna_hybrid_w4a16_apply_impl(
             a,
             w["w_q_skinny"],
             w["w_s_skinny"],

tests/kernels/quantization/test_hybrid_w4a16_triton.py

@@ -7,7 +7,7 @@
 Tests validate:
   - Triton GEMM (triton_w4a16_skinny_fmt_gemm) for prefill path
   - HIP wvSplitK_int4_g for decode path
-  - Full hybrid dispatch (torch.ops.vllm.hybrid_w4a16_apply) routing
+  - Full hybrid dispatch (torch.ops.vllm.rdna_hybrid_w4a16_apply) routing
 
 Run `pytest tests/kernels/quantization/test_hybrid_w4a16_triton.py`.
 """
@@ -30,7 +30,7 @@
 device = "cuda"
 
 hybrid_w4a16_module = importlib.import_module(
-    "vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16"
+    "vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16"
 )
 triton_w4a16_skinny_fmt_gemm = hybrid_w4a16_module.triton_w4a16_skinny_fmt_gemm
 
@@ -151,7 +151,7 @@ def _pack_scale_zp(
     """Build the asymmetric PackedSb carrier [N, K//G] fp32 that
     triton_w4a16_skinny_fmt_gemm consumes (low 16 bits = scale; high 16 bits =
     fp16 bias_eff = -(8 + (zp-8))*scale, or bf16 integer zp). Mirrors
-    HybridW4A16LinearKernel.process_weights_after_loading.
+    RDNAHybridW4A16LinearKernel.process_weights_after_loading.
     """
     scale_u16 = scales_nkg.contiguous().view(torch.uint16).to(torch.int32) & 0xFFFF
     if dtype == torch.float16:
@@ -318,7 +318,7 @@ def test_hybrid_w4a16_dispatch(dtype, M, K, N, G, random_seed: int):
     )
 
     cu_count = num_compute_units()
-    out = torch.ops.vllm.hybrid_w4a16_apply(
+    out = torch.ops.vllm.rdna_hybrid_w4a16_apply(
         a, b_packed_i8, scales, b_packed_i32, None, None, cu_count, G
     )
 

tests/kernels/quantization/test_moe_gemm_w4a16.py

@@ -33,7 +33,7 @@
 
 pytest.importorskip("triton")
 
-from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
+from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
     pack_int4_exllama_shuffle,
 )
 from vllm.model_executor.layers.fused_moe import moe_hip_w4a16  # noqa: E402

tests/kernels/quantization/test_rdna_hybrid_w4a16.py

@@ -37,6 +37,41 @@
 # ---------------------------------------------------------------------------
 
 
+def _pack_zp_rows_for_kernel(zp_nkg: torch.Tensor) -> torch.Tensor:
+    """Pack raw uint4 zero points along N: [N, G] int32 -> [N//8, G] int32.
+
+    Mirrors the layout the HIP skinny kernel reads (and that
+    ``process_weights_after_loading`` produces): row n's nibble lives in
+    word[n//8] at bits 4*(n%8).
+    """
+    assert zp_nkg.dtype == torch.int32
+    N, G = zp_nkg.shape
+    assert N % 8 == 0
+    shifts = (torch.arange(8, device=zp_nkg.device, dtype=torch.int32) * 4)[:, None]
+    return torch.sum(
+        (zp_nkg.view(N // 8, 8, G) & 0xF) << shifts, dim=1, dtype=torch.int32
+    ).contiguous()
+
+
+def _build_packed_scale_zp(
+    scales_nkg: torch.Tensor, zp_nkg: torch.Tensor, dtype: torch.dtype
+) -> torch.Tensor:
+    """Build the Triton scale/zp carrier exactly as the layer does.
+
+    fp16 packs scale | bias_eff, bf16 packs scale | zp_int; see the kernel
+    docstring for how each is consumed.
+    """
+    scale_u16 = scales_nkg.view(torch.uint16).to(torch.int32) & 0xFFFF
+    if dtype == torch.float16:
+        w_s_f32 = scales_nkg.to(torch.float32)
+        scaled_zp_f32 = (zp_nkg.to(torch.float32) - 8.0) * w_s_f32
+        bias_eff = (-(8.0 * w_s_f32 + scaled_zp_f32)).to(dtype)
+        hi_u16 = bias_eff.contiguous().view(torch.uint16).to(torch.int32) & 0xFFFF
+    else:
+        hi_u16 = zp_nkg.to(torch.int32) & 0xFFFF
+    return ((hi_u16 << 16) | scale_u16).view(torch.float32).contiguous()
+
+
 def _rdna_hybrid_w4a16_reference(
     x_mk: torch.Tensor,
     w_int4_nk: torch.Tensor,
@@ -50,7 +85,7 @@ def _rdna_hybrid_w4a16_reference(
     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),
+    zp_nkg: [N, K//G] int32 raw zero points in [0, 15],
             or None for symmetric (uint4b8, dequant subtracts 8)
     """
     G = group_size
@@ -110,24 +145,31 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M)
         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.
+    # Optional raw zero points [N, K//G] as int32 nibbles. The kernels consume
+    # them packed 8x along N (skinny) and folded into the carrier (Triton).
     if has_zp:
         zp_nkg = torch.randint(
             0, 16, (N, K // group_size), device=device, dtype=torch.int32
-        ).to(dtype)
+        )
+        w_zp = _pack_zp_rows_for_kernel(zp_nkg)
+        packed_scale_zp = _build_packed_scale_zp(scales_nkg, zp_nkg, dtype)
     else:
         zp_nkg = None
+        w_zp = None
+        packed_scale_zp = 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,
+        w_q_i32,
+        w_zp,
         None,  # bias
         num_compute_units(),
         group_size,
+        packed_scale_zp,
     )
 
     ref = _rdna_hybrid_w4a16_reference(
@@ -163,6 +205,7 @@ def test_rdna_hybrid_w4a16_apply_with_bias(dtype, M):
         x_mk,
         w_q,
         scales_nkg,
+        w_q_i32,
         None,
         bias,
         num_compute_units(),
@@ -378,11 +421,18 @@ def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_in
     )
     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)
+    # Zero-points stay in their PACKED 4-bit form -- [N//8, K//G] int32, row n's
+    # nibble at word[n//8] bits 4*(n%8) -- which is what the HIP skinny kernel
+    # reads. Expanding them to the activation dtype would cost 4x the DRAM
+    # traffic on a memory-bound kernel (and the C++ rejects that dtype).
+    assert layer.weight_zero_point.dtype == torch.int32
+    assert tuple(layer.weight_zero_point.shape) == (N // 8, K // G)
+    torch.testing.assert_close(layer.weight_zero_point, zeros_ckpt_n8kg)
+
+    # The Triton prefill path reads the scale/zp carrier instead of two loads;
+    # asymmetric layers must have it built at load time.
+    assert hasattr(layer, "_hybrid_w_packed_scale_zp")
+    assert tuple(layer._hybrid_w_packed_scale_zp.shape) == (N, K // G)
 
     # Quantized weights match symmetric path's layout regardless of zp.
     w_q_i32 = layer.weight_packed.view(torch.int32)
@@ -517,7 +567,7 @@ def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G):
 
     cu_count = num_compute_units()
     out = torch.ops.vllm.rdna_hybrid_w4a16_apply(
-        a, b_packed_i8, scales, None, None, cu_count, G
+        a, b_packed_i8, scales, b_packed_i32, None, None, cu_count, G
     )
 
     ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8)

tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py

@@ -38,7 +38,7 @@ def test_rocm_compressed_tensors_w4a16_e2e(
     with vllm_runner(
         model_path, dtype="float16", gpu_memory_utilization=0.35
     ) as vllm_model:
-        # Note: we cannot assert HybridW4A16LinearKernel is selected here
+        # Note: we cannot assert RDNAHybridW4A16LinearKernel is selected here
         # because V1 engine runs the model in a subprocess and apply_model
         # requires serializable callables (msgpack can't serialize functions).
         # If the W4A16 kernel is broken, generate_greedy will throw.

tests/quantization/test_hip_w4a16_kernel.py

@@ -385,8 +385,8 @@ def test_hybrid_w4a16_process_weights_symmetric_with_prefill_dequant():
     from types import SimpleNamespace
 
     from vllm.config import VllmConfig, set_current_vllm_config
-    from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (  # noqa: E501
-        HybridW4A16LinearKernel,
+    from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (  # noqa: E501
+        RDNAHybridW4A16LinearKernel,
     )
 
     _ensure_single_process_model_parallel()
@@ -415,7 +415,7 @@ def test_hybrid_w4a16_process_weights_symmetric_with_prefill_dequant():
         ),
     )
 
-    kernel = HybridW4A16LinearKernel(
+    kernel = RDNAHybridW4A16LinearKernel(
         MPLinearLayerConfig(
             full_weight_shape=(K, N),
             partition_weight_shape=(K, N),

vllm/model_executor/kernels/linear/__init__.py

@@ -54,9 +54,6 @@
 from vllm.model_executor.kernels.linear.mixed_precision.humming import (
     HummingLinearKernel,
 )
-from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
-    HybridW4A16LinearKernel,
-)
 from vllm.model_executor.kernels.linear.mixed_precision.machete import (
     MacheteLinearKernel,
 )
@@ -437,7 +434,6 @@ def _filter_kernels_by_backend(
         HummingLinearKernel,
     ],
     PlatformEnum.ROCM: [
-        HybridW4A16LinearKernel,
         HipW8A16LinearKernel,
         RDNA3W4A16LinearKernel,
         RDNAHybridW4A16LinearKernel,
@@ -1192,7 +1188,6 @@ def register_linear_kernel(
     "RDNAHybridW4A16LinearKernel",
     "MacheteLinearKernel",
     "MarlinLinearKernel",
-    "HybridW4A16LinearKernel",
     "TritonW4A16LinearKernel",
     "XPUW4A8IntLinearKernel",
     "XPUwNa16LinearKernel",

vllm/model_executor/kernels/linear/mixed_precision/__init__.py

@@ -28,9 +28,6 @@
 from vllm.model_executor.kernels.linear.mixed_precision.humming import (
     HummingLinearKernel,
 )
-from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import (
-    HybridW4A16LinearKernel,
-)
 from vllm.model_executor.kernels.linear.mixed_precision.machete import (
     MacheteLinearKernel,
 )
@@ -68,7 +65,6 @@
     "Dynamic4bitLinearKernel",
     "ExllamaLinearKernel",
     "HipW4A16LinearKernel",
-    "HybridW4A16LinearKernel",
     "HipW8A16LinearKernel",
     "HummingLinearKernel",
     "MacheteLinearKernel",