vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16

Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode.

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 the triton kernel reinterprets it as int32 [N, K//8] via a view (and transposes tiles in-register). No dual weight storage.

Classes:

• RDNAHybridW4A16LinearKernel –

  Hybrid W4A16 kernel: HIP skinny for decode, Triton for prefill.

Functions:

• pack_int4_exllama_shuffle –

  Pack uint4 values into ExLlama shuffle format: [N, K] -> [N, K//8] int32.

• pack_skinny_int4 –

  Pack [N, K] uint4 into the skinny weight layout the kernels consume.

• triton_w4a16_skinny_fmt_gemm –

  Fused W4A16 GEMM reading from skinny weight format [N, K//8].

RDNAHybridW4A16LinearKernel

Bases: MPLinearKernel

Hybrid W4A16 kernel: HIP skinny for decode, Triton for prefill.

Stores the weights once as int8 [N, K//2] (ExLlama shuffle packed). The HIP skinny kernel reads it directly; the triton kernel reinterprets the same buffer as int32 [N, K//8] via a view, so there is no dual weight storage.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

class RDNAHybridW4A16LinearKernel(MPLinearKernel):
    """Hybrid W4A16 kernel: HIP skinny for decode, Triton for prefill.

    Stores the weights once as int8 [N, K//2] (ExLlama shuffle packed). The
    HIP skinny kernel reads it directly; the triton kernel reinterprets the
    same buffer as int32 [N, K//8] via a view, so there is no dual weight
    storage.
    """

    SUPPORTED_QUANT_TYPES = [
        scalar_types.uint4b8,  # symmetric GPTQ (bias=8)
        scalar_types.uint4,  # asymmetric (zero_points)
    ]

    @classmethod
    def get_min_capability(cls) -> int:
        # Arch filtering is handled by can_implement (_on_gfx1x check)
        return 0

    @classmethod
    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
        if not current_platform.is_rocm():
            return False, "RDNAHybridW4A16LinearKernel only targets ROCm"

        if not _on_gfx1x():
            return False, "RDNAHybridW4A16LinearKernel only targets gfx11/gfx12"

        if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
            return (
                False,
                f"Quant type {c.weight_type} not supported; "
                f"supported: {cls.SUPPORTED_QUANT_TYPES}",
            )

        if c.act_type not in (torch.float16, torch.bfloat16):
            return False, "requires float16 or bfloat16 activations"

        gs = c.group_size
        if gs not in SUPPORTED_GROUP_SIZES:
            return (
                False,
                f"Group size {gs} not supported; supported: {SUPPORTED_GROUP_SIZES}",
            )

        K = c.partition_weight_shape[0]
        if K % 16 != 0:
            return False, f"K={K} must be divisible by 16"

        if K % gs != 0:
            return (
                False,
                f"K={K} must be divisible by group_size={gs}",
            )

        return True, None

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        c = self.config

        w_q_raw = getattr(layer, self.w_q_name)

        unpacked = unpack_quantized_values_into_int32(
            w_q_raw.data, c.weight_type, packed_dim=w_q_raw.packed_dim
        )
        # AWQ weights arrive as (K, N) with output_dim=1;
        # compressed-tensors arrive as (N, K) with output_dim=0.
        if getattr(w_q_raw, "output_dim", 0) != 0:
            unpacked = unpacked.t().contiguous()

        # Store as int8; Triton reinterprets via .view(torch.int32) at apply time.
        w_q_skinny = pack_skinny_int4(unpacked)

        # Permute via `_transform_param`, so the loader params keep their dims.
        if c.zero_points:
            assert self.w_zp_name is not None
            self._transform_param(
                layer,
                self.w_zp_name,
                lambda x: permute_param_layout_(
                    x, input_dim=1, output_dim=0, packed_dim=0
                ).data.contiguous(),
            )

        self._transform_param(layer, self.w_q_name, lambda x: w_q_skinny)
        self._transform_param(
            layer,
            self.w_s_name,
            lambda x: permute_param_layout_(
                x, input_dim=1, output_dim=0
            ).data.contiguous(),
        )

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        from vllm.utils.platform_utils import num_compute_units

        c = self.config
        w_q, w_s, w_zp = self._get_weight_params(layer)

        x_2d = x.reshape(-1, x.shape[-1])
        N = w_q.shape[0]
        out_shape = x.shape[:-1] + (N,)

        cu_count = num_compute_units()
        output = torch.ops.vllm.rdna_hybrid_w4a16_apply(
            x_2d,
            w_q,
            w_s,
            w_zp,
            bias,
            cu_count,
            c.group_size,
        )
        return output.reshape(out_shape)

_act_pad_bytes(row_bytes)

Bytes to add to an activation row stride to move it off the cliff.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

def _act_pad_bytes(row_bytes: int) -> int:
    """Bytes to add to an activation row stride to move it off the cliff."""
    if row_bytes % _ACT_CLIFF_BYTES:
        return 0
    return _STRIDE_PAD_BYTES

_int4_pair_to_fp16x2(x)

Unpack two packed int4 nibbles into a uint32 holding two fp16 lanes, each equal to 1024 + nibble, with one v_and_or_b32 ((x & 0x000F000F) | 0x64006400).

OR-ing a 4-bit nibble into the low mantissa of fp16 1024.0 (0x6400) bitcasts to exactly 1024+n. Doing it on a full 32-bit lane dequants two nibbles per instruction, vs the scalar v_and_b16 + v_or_b16 pair Triton emits from the elementwise form.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

@triton.jit
def _int4_pair_to_fp16x2(x):
    """Unpack two packed int4 nibbles into a uint32 holding two fp16 lanes,
    each equal to 1024 + nibble, with one ``v_and_or_b32``
    (``(x & 0x000F000F) | 0x64006400``).

    OR-ing a 4-bit nibble into the low mantissa of fp16 1024.0 (0x6400)
    bitcasts to exactly 1024+n. Doing it on a full 32-bit lane dequants two
    nibbles per instruction, vs the scalar v_and_b16 + v_or_b16 pair Triton
    emits from the elementwise form.
    """
    mask = tl.full(x.shape, 0x000F000F, tl.int32)
    return tl.inline_asm_elementwise(
        asm="v_and_or_b32 $0, $1, $2, 0x64006400",
        constraints="=v,v,v",
        args=[x, mask],
        dtype=tl.uint32,
        is_pure=True,
        pack=1,
    )

_pad_activation_rows(x_2d)

Copy x_2d into a row-padded buffer when its row stride is on the cliff.

The packed weight is padded once at load time, but activations are produced fresh every step, so this materialises a padded copy.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

def _pad_activation_rows(x_2d: torch.Tensor) -> torch.Tensor:
    """Copy ``x_2d`` into a row-padded buffer when its row stride is on the cliff.

    The packed weight is padded once at load time, but activations are produced
    fresh every step, so this materialises a padded copy.
    """
    pad_bytes = _act_pad_bytes(x_2d.shape[1] * x_2d.element_size())
    if not (pad_bytes and _on_gfx1151()):
        return x_2d
    pad_elems = pad_bytes // x_2d.element_size()
    buf = torch.empty(
        (x_2d.shape[0], x_2d.shape[1] + pad_elems),
        dtype=x_2d.dtype,
        device=x_2d.device,
    )
    padded = buf[:, : x_2d.shape[1]]
    padded.copy_(x_2d)
    return padded

_rdna_hybrid_w4a16_apply_impl(x_2d, w_q, w_s, w_zp, bias, cu_count, group_size)

Dispatch between skinny GEMM and Triton based on batch size M.

w_zp is [N//8, K//G] int32 for asymmetric layers, None for symmetric.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

def _rdna_hybrid_w4a16_apply_impl(
    x_2d: torch.Tensor,
    w_q: torch.Tensor,
    w_s: torch.Tensor,
    w_zp: torch.Tensor | None,
    bias: torch.Tensor | None,
    cu_count: int,
    group_size: int,
) -> torch.Tensor:
    """Dispatch between skinny GEMM and Triton based on batch size M.

    ``w_zp`` is [N//8, K//G] int32 for asymmetric layers, None for symmetric.
    """
    import vllm._custom_ops as ops

    M = x_2d.shape[0]
    K = x_2d.shape[1]
    N = w_q.shape[0]

    if M <= MAX_SKINNY_BATCH_SIZE and (
        K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS or _on_gfx115x()
    ):
        # record_function is not torch.compile-safe; use nullcontext when
        # compiling to keep the op traceable.
        ctx = (
            nullcontext()
            if torch.compiler.is_compiling()
            else torch.profiler.record_function(f"wvsplitk_int4 {M}x{N}x{K}")
        )
        with ctx:
            return ops.wvSplitK_int4_g(w_q, x_2d, w_s, cu_count, group_size, w_zp, bias)

    ctx = (
        nullcontext()
        if torch.compiler.is_compiling()
        else torch.profiler.record_function(f"hybrid_triton_w4a16 {M}x{N}x{K}")
    )
    with ctx:
        output = triton_w4a16_skinny_fmt_gemm(
            a=_pad_activation_rows(x_2d),
            b_q=w_q.view(torch.int32),
            scales=w_s,
            group_size=group_size,
            zp=w_zp,
        )
        if bias is not None:
            output.add_(bias)
    return output

_select_skinny_gfx1151_config(M, N, K, group_size, dtype, has_zp=False)

Return (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, packed_dequant) for gfx1151.

The packed fp16 dequant and the scalar bf16 dequant want different tiles, so selection is dtype-aware. bf16 always uses the scalar-tuned ladder.

num_stages None means "leave Triton's default pipeline depth alone". Re-run benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py after edits.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

def _select_skinny_gfx1151_config(
    M: int,
    N: int,
    K: int,
    group_size: int,
    dtype: torch.dtype,
    has_zp: bool = False,
) -> tuple[int, int, int, int, int | None, bool]:
    """Return (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages,
    packed_dequant) for gfx1151.

    The packed fp16 dequant and the scalar bf16 dequant want different tiles,
    so selection is dtype-aware. bf16 always uses the scalar-tuned ladder.

    num_stages None means "leave Triton's default pipeline depth alone".
    Re-run benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py after edits.
    """
    num_stages: int | None = None
    # Up to this M the packed fp16 path is slower. Asymmetric weights were
    # measured to break even at a larger M.
    min_m = 256 if has_zp else 128
    if dtype == torch.float16 and min_m < M:
        # Tuned together with these tiles
        num_stages = 1
        packed_dequant = True
        if M <= 256:
            block_m, block_n, block_k, num_warps = 128, 128, 32, 8
        elif M < 2048:  # 257..2047 (mostly 512, 1024): wide distilled tile
            block_m, block_n, block_k, num_warps = 128, 256, 32, 8
        else:  # M >= 2048 (deep prefill)
            if N <= 2048 and K >= 4096:  # narrow + deep: halved BM saturates
                block_m, block_n, block_k, num_warps = 64, 256, 32, 8
            else:
                block_m, block_n, block_k, num_warps = 128, 256, 32, 8
        # Very narrow N at small/mid M: a wide BLOCK_N leaves too few N-tiles to
        # fill the CUs, so clamp it. At M>=1024 the M-tiles already saturate.
        if N <= 1024 and M <= 512:
            block_n = min(block_n, 32)
    else:
        packed_dequant = dtype == torch.float16 and not has_zp
        key = (group_size, K, N)
        override = _GFX1151_PREFILL_OVERRIDES.get(key) if M <= 128 else None
        if override is not None:
            block_m, block_n, block_k, num_warps, num_stages = override
        elif M <= 32:
            block_m, block_n, block_k, num_warps = 32, 32, 128, 4
        elif M <= 64:
            block_m, block_n, block_k, num_warps = 64, 64, 32, 4
        elif M <= 128:
            if K >= 2 * N:  # tall K (e.g. down_proj)
                block_m, block_n, block_k, num_warps = 64, 16, 64, 1
            elif N > K:  # wide N (e.g. qkv_proj, gate_up_proj)
                block_m, block_n, block_k, num_warps = 64, 64, 64, 4
            else:  # N ~= K (e.g. o_proj)
                block_m, block_n, block_k, num_warps = 64, 32, 64, 4
        elif M <= 1024:
            if K >= 2 * N:  # tall K (e.g. down_proj)
                block_m, block_n, block_k, num_warps = 64, 64, 64, 4
            elif N >= 4 * K:  # very wide N (e.g. gate_up_proj)
                block_m, block_n, block_k, num_warps = 128, 64, 64, 8
            else:
                block_m, block_n, block_k, num_warps = 64, 128, 32, 4
        else:
            if K >= 2 * N:  # tall K (e.g. down_proj)
                block_m, block_n, block_k, num_warps = 128, 512, 32, 16
            else:
                block_m, block_n, block_k, num_warps = 128, 64, 64, 8
    return block_m, block_n, block_k, num_warps, num_stages, packed_dequant

_triton_w4a16_skinny_fmt_kernel(a_ptr, b_ptr, scales_ptr, zp_ptr, c_ptr, M, N, K, K8, num_groups, stride_sn, stride_bn, stride_am, group_size, ZP_BIAS, HAS_ZP, PACKED_DEQUANT, BLOCK_M, BLOCK_N, BLOCK_K)

Fused W4A16 GEMM reading weights from skinny format [N, K//8].

B is stored as [N, K//8] int32 using ExLlama shuffle packing: each int32 packs 8 K-values with interleave [0,2,4,6,1,3,5,7]: packed = val[0] | (val[2]<<4) | (val[4]<<8) | (val[6]<<12) | (val[1]<<16) | (val[3]<<20) | (val[5]<<24) | (val[7]<<28)

Scales are [N, K//G] (skinny layout, NOT transposed). When HAS_ZP=True, zp_ptr holds [N//8, K//G] int32 with row n's raw zero-point at word[n//8] bits 4*(n%8), and dequant is (nibble - zp_raw) * scale. When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric).

With PACKED_DEQUANT the nibble arrives as b_raw = 1024 + nibble (the magic-constant unpack), so the subtrahend absorbs the 1024: the arithmetic is unchanged and every intermediate stays exact, since fp16 represents every integer below 2048.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

@triton.jit
def _triton_w4a16_skinny_fmt_kernel(
    # Pointers
    a_ptr,  # [M, K]  fp16/bf16 activations
    b_ptr,  # [N, K//8]  int32 packed (ExLlama shuffle, K is packed dim)
    scales_ptr,  # [N, K//G]  fp16/bf16 scales (skinny layout)
    zp_ptr,  # [N//8, K//G]  int32 zero-points (when HAS_ZP=True)
    c_ptr,  # [M, N]  fp16/bf16 output
    # Dimensions
    M,
    N,
    K,
    K8,  # K // 8
    num_groups,  # K // group_size
    stride_sn,  # row stride of scales_ptr and zp_ptr, in groups
    stride_bn,  # b_ptr row stride; >= K8, the rows may be padded
    stride_am,  # a_ptr row stride in elements; >= K, rows may be padded
    # Quantization parameters
    group_size,
    ZP_BIAS: tl.constexpr,
    HAS_ZP: tl.constexpr,
    PACKED_DEQUANT: tl.constexpr,
    # Block sizes
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    """Fused W4A16 GEMM reading weights from skinny format [N, K//8].

    B is stored as [N, K//8] int32 using ExLlama shuffle packing:
      each int32 packs 8 K-values with interleave [0,2,4,6,1,3,5,7]:
        packed = val[0] | (val[2]<<4) | (val[4]<<8) | (val[6]<<12)
               | (val[1]<<16) | (val[3]<<20) | (val[5]<<24) | (val[7]<<28)

    Scales are [N, K//G] (skinny layout, NOT transposed).
    When HAS_ZP=True, zp_ptr holds [N//8, K//G] int32 with row n's raw
    zero-point at word[n//8] bits 4*(n%8), and dequant is
    (nibble - zp_raw) * scale.
    When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric).

    With PACKED_DEQUANT the nibble arrives as ``b_raw`` = 1024 + nibble (the
    magic-constant unpack), so the subtrahend absorbs the 1024: the arithmetic
    is unchanged and every intermediate stays exact, since fp16 represents every
    integer below 2048.
    """
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

    # ExLlama unshuffle shifts: shift[j] = (j//2)*4 + (j%2)*16
    # For 8 values: [0, 16, 4, 20, 8, 24, 12, 28]
    exllama_shifts_row = (tl.arange(0, 8) // 2) * 4 + (tl.arange(0, 8) % 2) * 16
    shifts_1d = tl.reshape(
        tl.broadcast_to(exllama_shifts_row[None, :], (BLOCK_K // 8, 8)),
        (BLOCK_K,),
    )
    shifts_full = tl.broadcast_to(shifts_1d[None, :], (BLOCK_N, BLOCK_K))

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k_start in range(0, tl.cdiv(K, BLOCK_K)):
        offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K)
        mask_k = offs_k < K

        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :]
        mask_a = (offs_m[:, None] < M) & mask_k[None, :]
        a = tl.load(a_ptrs, mask=mask_a, other=0.0)

        offs_k8 = k_start * (BLOCK_K // 8) + tl.arange(0, BLOCK_K // 8)
        b_ptrs = b_ptr + offs_n[:, None] * stride_bn + offs_k8[None, :]
        mask_b = (offs_n[:, None] < N) & (offs_k8[None, :] < K8)
        b_packed = tl.load(b_ptrs, mask=mask_b, other=0)

        if PACKED_DEQUANT:
            # Dequants two nibbles per instruction instead of one; the dequant
            # loop is VALU-issue-bound on gfx11.
            #
            # Each ExLlama int32 holds 4 nibble pairs: pair p is val[2p] at bits
            # [4p, 4p+4) and val[2p+1] at bits [16+4p, 20+4p). Shifting right by
            # 4p moves pair p to bits [0, 4) and [16, 20). The sign fill of the
            # signed shift lands above bit 20 and is masked off below.
            shifts4 = (tl.arange(0, 4) * 4)[None, None, :]
            # [BLOCK_N, BLOCK_K // 8, 4] -> [BLOCK_N, BLOCK_K // 2]: one int32
            # per nibble pair, in K order.
            bp_shift = tl.reshape(
                b_packed[:, :, None] >> shifts4, (BLOCK_N, BLOCK_K // 2)
            )
            # One v_and_or_b32 per pair: two fp16 lanes, each 1024 + nibble.
            packed_hl = _int4_pair_to_fp16x2(bp_shift)
            # Split the lanes: lo = val[2p], hi = val[2p+1].
            lo = (packed_hl & 0xFFFF).to(tl.uint16).to(tl.float16, bitcast=True)
            hi = (packed_hl >> 16).to(tl.uint16).to(tl.float16, bitcast=True)
            # Back to K order, [BLOCK_N, BLOCK_K] fp16 of 1024 + nibble. Adjacent
            # pairs let the affine below pack into v_pk_fma_f16.
            b_raw = tl.interleave(lo, hi)
        else:
            b = tl.interleave(b_packed, b_packed)
            b = tl.interleave(b, b)
            b = tl.interleave(b, b)
            b = (b >> shifts_full) & 0xF

        group_idx = (k_start * BLOCK_K) // group_size
        scale_ptrs = scales_ptr + offs_n * stride_sn + group_idx
        scale_mask = offs_n < N
        scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0)

        if HAS_ZP:
            zp_ptrs = zp_ptr + (offs_n // 8) * stride_sn + group_idx
            zp_word = tl.load(zp_ptrs, mask=scale_mask, other=0)
            zp_raw = (zp_word >> (4 * (offs_n % 8))) & 0xF

        if PACKED_DEQUANT:
            # b_raw is 1024 + nibble, so subtract 1024 along with the zero point.
            c1024 = tl.full((), 1024.0, tl.float16)
            if HAS_ZP:
                # 1024 + zp is one fp16 per row, not per element.
                zp_off = (c1024 + zp_raw.to(tl.float16))[:, None]
                # b_raw - (1024 + zp) == nibble - zp exactly: every value is an
                # integer below 2048, which fp16 represents exactly. Only the
                # multiply by the scale rounds, as in the scalar path.
                b_fp = (b_raw - zp_off) * scales[:, None]
            else:
                # Same with the constant ZP_BIAS in place of the zero point.
                b_fp = (b_raw - (c1024 + ZP_BIAS)) * scales[:, None]
        else:
            if HAS_ZP:
                b_fp = (b - zp_raw[:, None]).to(scales.dtype) * scales[:, None]
            else:
                b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None]

        b_fp_t = tl.trans(b_fp)
        accumulator += tl.dot(a, b_fp_t, out_dtype=tl.float32)

    c = accumulator.to(c_ptr.type.element_ty)
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask_c = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask_c)

_weight_pad_bytes(row_bytes)

Bytes to add to a packed weight row stride to move it off the cliff.

Only strides that are a multiple of _WEIGHT_CLIFF_BYTES are moved, and only by one cache line; padding a stride that is already off the cliff costs memory and measures slower.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

def _weight_pad_bytes(row_bytes: int) -> int:
    """Bytes to add to a packed weight row stride to move it off the cliff.

    Only strides that are a multiple of ``_WEIGHT_CLIFF_BYTES`` are moved, and
    only by one cache line; padding a stride that is already off the cliff
    costs memory and measures slower.
    """
    if row_bytes % _WEIGHT_CLIFF_BYTES:
        return 0
    return _STRIDE_PAD_BYTES

pack_int4_exllama_shuffle(w_uint4)

Pack uint4 values into ExLlama shuffle format: [N, K] -> [N, K//8] int32.

Each int32 packs 8 K-values with interleave order [0,2,4,6,1,3,5,7].

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

def pack_int4_exllama_shuffle(w_uint4: torch.Tensor) -> torch.Tensor:
    """Pack uint4 values into ExLlama shuffle format: [N, K] -> [N, K//8] int32.

    Each int32 packs 8 K-values with interleave order [0,2,4,6,1,3,5,7].
    """
    N_dim, K_dim = w_uint4.shape
    assert K_dim % 8 == 0
    g = w_uint4.to(torch.uint8).view(N_dim, K_dim // 8, 8).to(torch.int32)
    return (
        g[:, :, 0]
        | (g[:, :, 2] << 4)
        | (g[:, :, 4] << 8)
        | (g[:, :, 6] << 12)
        | (g[:, :, 1] << 16)
        | (g[:, :, 3] << 20)
        | (g[:, :, 5] << 24)
        | (g[:, :, 7] << 28)
    )

pack_skinny_int4(unpacked)

Pack [N, K] uint4 into the skinny weight layout the kernels consume.

ExLlama shuffle to [N, K//8] int32, viewed as int8 [N, K//2]. On gfx1151 the row stride is nudged off the cliff (see _weight_pad_bytes).

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

def pack_skinny_int4(unpacked: torch.Tensor) -> torch.Tensor:
    """Pack [N, K] uint4 into the skinny weight layout the kernels consume.

    ExLlama shuffle to [N, K//8] int32, viewed as int8 [N, K//2]. On gfx1151
    the row stride is nudged off the cliff (see ``_weight_pad_bytes``).
    """
    shuffled = pack_int4_exllama_shuffle(unpacked)
    n_rows, k8 = shuffled.shape
    pad_int32 = _weight_pad_bytes(k8 * 4) // 4
    if not (pad_int32 and _on_gfx1151()):
        return shuffled.contiguous().view(torch.int8)
    padded = torch.empty(
        (n_rows, k8 + pad_int32), dtype=torch.int32, device=shuffled.device
    )
    padded[:, :k8].copy_(shuffled)
    # The int8 view keeps stride(0) = 4 * (k8 + pad_int32) bytes, and
    # .view(torch.int32) at apply time recovers the [N, K//8] int32 view.
    return padded.view(torch.int8)[:, : k8 * 4]

triton_w4a16_skinny_fmt_gemm(a, b_q, scales, group_size, zp_bias=8, zp=None)

Fused W4A16 GEMM reading from skinny weight format [N, K//8].

Parameters:

• a

  (Tensor) –

  Activation matrix [M, K], float16 or bfloat16.

• b_q

  (Tensor) –

  Packed weight matrix [N, K//8], int32 (ExLlama shuffle).

• scales

  (Tensor) –

  Per-group scales [N, K//G], same dtype as a.

• group_size

  (int) –

  Quantization group size (resolved from -1 to K by caller).

• zp_bias

  (int, default: 8 ) –

  Constant zero bias (default 8 for unsigned int4).

• zp

  (Tensor | None, default: None ) –

  Raw per-group zero-points [N//8, K//G] int32, row n at word[n//8] bits 4*(n%8) (asymmetric). When provided, dequant is (nibble - zp_raw) * scale.

Returns:

• Tensor –

  Output matrix [M, N], same dtype as a.

Source code in vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py

def triton_w4a16_skinny_fmt_gemm(
    a: torch.Tensor,  # [M, K] fp16/bf16
    b_q: torch.Tensor,  # [N, K//8] int32 (ExLlama shuffle packed)
    scales: torch.Tensor,  # [N, K//G] fp16/bf16
    group_size: int,
    zp_bias: int = 8,
    zp: torch.Tensor | None = None,  # [N//8, K//G] int32 zero-points
) -> torch.Tensor:
    """Fused W4A16 GEMM reading from skinny weight format [N, K//8].

    Args:
        a:          Activation matrix [M, K], float16 or bfloat16.
        b_q:        Packed weight matrix [N, K//8], int32 (ExLlama shuffle).
        scales:     Per-group scales [N, K//G], same dtype as a.
        group_size: Quantization group size (resolved from -1 to K by caller).
        zp_bias:    Constant zero bias (default 8 for unsigned int4).
        zp:         Raw per-group zero-points [N//8, K//G] int32, row n at
                    word[n//8] bits 4*(n%8) (asymmetric). When provided,
                    dequant is (nibble - zp_raw) * scale.

    Returns:
        Output matrix [M, N], same dtype as a.

    """
    assert a.stride(1) == 1, "Activation rows must be contiguous"
    assert b_q.stride(1) == 1, "Weight rows must be contiguous"
    assert scales.stride(1) == 1, "Scale rows must be contiguous"
    stride_sn = scales.stride(0)

    M, K = a.shape
    N = b_q.shape[0]
    K8 = K // 8
    num_groups = K // group_size
    stride_bn = b_q.stride(0)
    stride_am = a.stride(0)

    assert b_q.shape == (N, K8), f"b_q shape mismatch: {b_q.shape} vs ({N}, {K8})"
    assert scales.shape == (N, num_groups), (
        f"scales shape mismatch: {scales.shape} vs ({N}, {num_groups})"
    )
    if zp is not None:
        assert zp.stride(1) == 1, "Zero-point rows must be contiguous"
        assert zp.stride(0) == stride_sn, (
            f"zp row stride {zp.stride(0)} must match scales {stride_sn}"
        )
        assert N % 8 == 0, f"N must be divisible by 8 for packed zp, got {N}"
        assert zp.shape == (N // 8, num_groups), (
            f"zp shape mismatch: {zp.shape} vs ({N // 8}, {num_groups})"
        )
    has_zp = zp is not None

    c = torch.empty((M, N), dtype=a.dtype, device=a.device)

    # num_stages stays None unless the tile table sets it, so the generic
    # heuristics fall back to Triton's default pipeline depth.
    num_stages: int | None = None
    packed_dequant = False
    if _on_gfx12x():
        # Tuned on gfx1201 (Radeon AI PRO R9700, 32 CUs, 32-wide wavefronts)
        # using Llama-3.1-8B AWQ weight shapes with group_size=128.
        if M <= 32:
            BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 16, 16, 128, 4
        elif M <= 64:
            if K >= 2 * N:  # tall K (e.g. down_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 32, 128, 8
            elif N > K:  # wide N (e.g. qkv_proj, gate_up_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 32, 64, 8
            else:  # N ~= K (e.g. o_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 32, 64, 128, 4
        elif M <= 128:
            if K >= 2 * N:  # tall K (e.g. down_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 16, 64, 1
            elif N >= 2 * K:  # very wide N (e.g. gate_up_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 128, 64, 8
            else:  # N ~= K (e.g. o_proj, qkv_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 64, 64, 8
        elif M <= 512:
            if K >= 2 * N:  # tall K (e.g. down_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 64, 64, 8
            elif N >= 4 * K:  # very wide N (e.g. gate_up_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 128, 64, 8
            else:
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 128, 64, 8
        else:
            if K >= 2 * N:  # tall K (e.g. down_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 64, 64, 8
            elif N >= 4 * K:  # very wide N (e.g. gate_up_proj)
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 256, 64, 64, 8
            else:
                BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 128, 32, 8
    elif _on_gfx1151():
        BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, packed_dequant = (
            _select_skinny_gfx1151_config(M, N, K, group_size, a.dtype, has_zp)
        )
    else:
        num_warps = 4
        if M <= 32:
            BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32
        elif M <= 64:
            BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
        else:
            BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32

    # The kernel loads one scale per BLOCK_K tile, so BLOCK_K must not
    # exceed group_size — otherwise elements in the tile that belong to
    # a different group would get the wrong scale.
    BLOCK_K = min(BLOCK_K, group_size)

    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))

    extra_kwargs = {} if num_stages is None else {"num_stages": num_stages}
    _triton_w4a16_skinny_fmt_kernel[grid](
        a,
        b_q,
        scales,
        zp if has_zp else scales,  # dummy pointer when no zp (unused)
        c,
        M,
        N,
        K,
        K8,
        num_groups,
        stride_sn,
        stride_bn,
        stride_am,
        group_size=group_size,
        ZP_BIAS=zp_bias,
        HAS_ZP=has_zp,
        PACKED_DEQUANT=packed_dequant,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        num_warps=num_warps,
        **extra_kwargs,
    )
    return c