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

[ROCm][Perf] W4A16: pad gfx11 weight and activation strides (#56301)

Browse Source 
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
This commit is contained in:
Matthias Gehre
2026-10-06 08:22:33 -05:00
committed by GitHub
parent 049507aa76
commit b267fe4881
4 changed files with 316 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py
+26 -11
View File
@@ -15,6 +15,9 @@ import os
import torch
from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
pack_skinny_int4,
)
from vllm.triton_utils import triton
# ---------------------------------------------------------------------------
@@ -56,20 +59,25 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"):
Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives
its int32 view from w_q_skinny, so no separate int32 buffer is returned.
Weights go through ``pack_skinny_int4``, the packer the layer uses, so the
benchmark sees the production row stride rather than an always-dense one.
"""
num_groups = K // group_size
# Random packed weights — actual values don't matter for throughput
w_q_skinny_i32 = torch.randint(
0, 2**31, (N, K // 8), dtype=torch.int32, device=device
)
w_q_skinny = w_q_skinny_i32.view(torch.int8).contiguous()
# Values don't matter for throughput, but the layout does.
unpacked = torch.randint(0, 16, (N, K), dtype=torch.int32, device=device)
w_q_skinny = pack_skinny_int4(unpacked)
del unpacked
w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01
# Raw per-group zero-points for asymmetric benchmarks
w_zp = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device).to(
torch.float16
)
# Per-group zero-points for asymmetric benchmarks, in the packed layout the
# kernels read: [N//8, num_groups] int32, row n's nibble at bits 4*(n%8).
zp_raw = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device)
shifts = (torch.arange(8, device=device, dtype=torch.int32) * 4)[:, None]
w_zp = torch.sum(
(zp_raw.view(N // 8, 8, num_groups) & 0xF) << shifts, dim=1, dtype=torch.int32
).contiguous()
# FP16 baseline for F.linear
w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01
@@ -97,6 +105,13 @@ PROVIDERS = ["torch-fp16", "hybrid-w4a16", "hybrid-w4a16-zp"]
)
)
def benchmark(batch_size, provider, N, K, group_size, weights):
"""Time one GEMM.
``do_bench`` rather than ``do_bench_cudagraph``: it zeroes a 256 MB buffer
between reps, so the weight is read from DRAM as it would be in a model.
With a cudagraph replay the weight stays resident in the 32 MB MALL, which
flatters wide-BLOCK_N tiles and hides row-stride effects entirely.
"""
M = batch_size
device = "cuda"
dtype = torch.float16
@@ -106,7 +121,7 @@ def benchmark(batch_size, provider, N, K, group_size, weights):
if provider == "torch-fp16":
w_fp16 = weights["w_fp16"]
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
ms, min_ms, max_ms = triton.testing.do_bench(
lambda: torch.nn.functional.linear(a, w_fp16),
quantiles=quantiles,
)
@@ -133,7 +148,7 @@ def benchmark(batch_size, provider, N, K, group_size, weights):
group_size,
)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
ms, min_ms, max_ms = triton.testing.do_bench(
run,
quantiles=quantiles,
)
csrc/rocm/skinny_gemms_int4.cu
+33 -28
View File
@@ -8,6 +8,7 @@
#include <stdexcept>
#include <algorithm>
#include <utility> // std::in_range
#include "../cuda_compat.h"
#include "dispatch_utils.h"
@@ -121,6 +122,7 @@ __device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed,
// GROUP_SIZE: 0 = per-channel scale [M], >0 = per-group scale [M,
// K/GROUP_SIZE].
// Requires GROUP_SIZE % A_CHUNK == 0 when GROUP_SIZE > 0.
// B_row_stride_bytes is the per-row byte stride of the packed weights.
#if defined(__HIP__GFX1X__)
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
@@ -130,9 +132,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
const scalar_t* __restrict__ A, const scalar_t* scale,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS, scalar_t* C,
const int _WvPrGrp, const int CuCount) {
const int _WvPrGrp, const int CuCount,
const int B_row_stride_bytes) {
constexpr int max_lds_len = LDS_SIZE / 2;
const int K_packed = K / 2;
union bigTypeA {
scalar_t h[A_CHUNK];
@@ -181,9 +183,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
uint32_t k_ = k + threadIdx.x * A_CHUNK;
if (k_ >= K) break;
const uint8_t* B_ = &B_packed[(m + 0) * K_packed + k_ / 2];
const uint8_t* B_ = &B_packed[(m + 0) * B_row_stride_bytes + k_ / 2];
for (int y = 0; y < YTILE; y++) {
const float* src = (const float*)(&B_[y * K_packed]);
const float* src = (const float*)(&B_[y * B_row_stride_bytes]);
#pragma unroll
for (int i = 0; i < A_CHUNK / 8; i++)
bigB[y][k2].f[i] = loadnt((float*)&src[i]);
@@ -344,14 +346,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
#else // !defined(__HIP__GFX1X__)
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
__global__ void wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx,
const int By, const uint8_t* B_packed,
const scalar_t* __restrict__ A,
const scalar_t* scale,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS,
scalar_t* C, const int _WvPrGrp,
const int CuCount) {
__global__ void wvSplitK_int4_hf_sml_(
const int K, const int M, const int Bx, const int By,
const uint8_t* B_packed, const scalar_t* __restrict__ A,
const scalar_t* scale, const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp,
const int CuCount, const int B_row_stride_bytes) {
UNREACHABLE_CODE
}
#endif // defined(__HIP__GFX1X__)
@@ -367,9 +367,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
const uint8_t* B_packed, const scalar_t* __restrict__ A,
const scalar_t* scale, const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS, scalar_t* C,
const int _WvPrGrp, const int CuCount) {
const int _WvPrGrp, const int CuCount,
const int B_row_stride_bytes) {
constexpr int max_lds_len = LDS_SIZE / 2;
const int K_packed = K / 2;
union bigTypeA {
scalar_t h[A_CHUNK];
@@ -430,9 +430,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
uint32_t k_ = k + threadIdx.x * A_CHUNK;
if (k_ >= K) break;
const uint8_t* B_ = &B_packed[(m + 0) * K_packed + k_ / 2];
const uint8_t* B_ = &B_packed[(m + 0) * B_row_stride_bytes + k_ / 2];
for (int y = 0; y < YTILE; y++) {
const float* src = (const float*)(&B_[y * K_packed]);
const float* src = (const float*)(&B_[y * B_row_stride_bytes]);
#pragma unroll
for (int i = 0; i < A_CHUNK / 8; i++)
bigB[y][k2].f[i] = loadnt((float*)&src[i]);
@@ -606,14 +606,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
#else // !defined(__HIP__GFX1X__)
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
__global__ void wvSplitK_int4_hf_(const int K, const int M, const int Bx,
const int By, const uint8_t* B_packed,
const scalar_t* __restrict__ A,
const scalar_t* scale,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS,
scalar_t* C, const int _WvPrGrp,
const int CuCount) {
__global__ void wvSplitK_int4_hf_(
const int K, const int M, const int Bx, const int By,
const uint8_t* B_packed, const scalar_t* __restrict__ A,
const scalar_t* scale, const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp,
const int CuCount, const int B_row_stride_bytes) {
UNREACHABLE_CODE
}
#endif // defined(__HIP__GFX1X__)
@@ -658,10 +656,17 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
? in_bias->size(0)
: 1;
int64_t expected_weight_bytes = M_in * K_in / 2;
int64_t actual_weight_bytes = in_a.numel() * in_a.element_size();
const int64_t expected_weight_bytes = M_in * K_in / 2;
const int64_t actual_weight_bytes = in_a.numel() * in_a.element_size();
TORCH_CHECK(actual_weight_bytes == expected_weight_bytes,
"Weight tensor must contain M*K/2 bytes for int4 packing");
const int64_t b_row_stride_bytes = in_a.stride(0) * in_a.element_size();
TORCH_CHECK(b_row_stride_bytes >= K_in / 2, "Weight row stride (",
b_row_stride_bytes, " bytes) must hold at least K/2=", K_in / 2,
" bytes per row");
TORCH_CHECK(std::in_range<int>(b_row_stride_bytes), "Weight row stride (",
b_row_stride_bytes, " bytes) exceeds int range");
const int b_row_stride_bytes_i32 = static_cast<int>(b_row_stride_bytes);
TORCH_CHECK(
in_b.dtype() == torch::kFloat16 || in_b.dtype() == torch::kBFloat16,
"Activation must be float16 or bfloat16");
@@ -722,12 +727,12 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
wvSplitK_int4_hf_sml_<fptype, _THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, \
_HAS_ZP><<<grid, block, 0, stream>>>( \
K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \
__wvPrGrp, CuCount); \
__wvPrGrp, CuCount, b_row_stride_bytes_i32); \
else \
wvSplitK_int4_hf_<fptype, _THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, \
_HAS_ZP><<<grid, block, 0, stream>>>( \
K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \
__wvPrGrp, CuCount); \
__wvPrGrp, CuCount, b_row_stride_bytes_i32); \
}
#define WVSPLITK_INT4G(_YTILE, _UNRL, _N, _GS, _HAS_ZP) \
tests/kernels/quantization/test_rdna_hybrid_w4a16.py
+179
View File
@@ -563,3 +563,182 @@ def test_wvsplitk_int4_g_rejects_unpacked_zero_points():
with pytest.raises(RuntimeError, match="Zero points must be int32 or uint32"):
ops.wvSplitK_int4_g(w, a, scales, num_compute_units(), G, zp_unpacked, None)
@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
@pytest.mark.skipif(
not hasattr(torch.ops, "_rocm_C")
or not hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"),
reason="wvSplitK_int4_g not built",
)
def test_wvsplitk_int4_g_rejects_undersized_weight_view():
"""A sufficient row stride does not prove the row holds K/2 packed bytes.
Row padding makes stride(0) > K/2 legitimate, but the logical size still
has to be checked: this view has a valid stride and half the required
bytes, and without the size check the kernel reads past it.
"""
import vllm._custom_ops as ops
from vllm.utils.platform_utils import num_compute_units
K, N, G, M = 512, 256, 128, 1
a = torch.randn((M, K), device=device, dtype=torch.float16)
# Full-size backing allocation, so a missing check misreads rather than
# running off the end of the buffer.
backing = torch.randint(0, 255, (N, K // 2), device=device, dtype=torch.uint8).view(
torch.int8
)
undersized = backing[:, : K // 4]
assert undersized.stride(0) == K // 2
scales = torch.rand((N, K // G), device=device, dtype=torch.float16)
with pytest.raises(RuntimeError, match=r"must contain M\*K/2 bytes"):
ops.wvSplitK_int4_g(undersized, a, scales, num_compute_units(), G, None, None)
# ---------------------------------------------------------------------------
# gfx11 weight row-stride padding
# ---------------------------------------------------------------------------
_weight_pad_bytes = hybrid_module._weight_pad_bytes
_act_pad_bytes = hybrid_module._act_pad_bytes
_WEIGHT_CLIFF_BYTES = hybrid_module._WEIGHT_CLIFF_BYTES
_ACT_CLIFF_BYTES = hybrid_module._ACT_CLIFF_BYTES
_STRIDE_PAD_BYTES = hybrid_module._STRIDE_PAD_BYTES
def test_weight_pad_bytes_only_moves_strides_on_the_cliff():
"""Pad 1024 B multiples, leave everything else dense."""
# On the cliff (K % 2048 == 0): 1024 B multiples.
for row_bytes in (1024, 2048, 4096, 5120, 6144, 7168, 8192):
assert _weight_pad_bytes(row_bytes) == _STRIDE_PAD_BYTES
# Off it, including strides that are 512 B multiples but not 1024 B ones:
# padding those measured as a real loss on gfx1151 (-14% at M=1 on a
# 4864 B row).
for row_bytes in (1280, 1536, 2560, 4864, 9472, 12800):
assert _weight_pad_bytes(row_bytes) == 0
# The padded stride is never back on the cliff.
for row_bytes in range(16, 16384, 16):
padded = row_bytes + _weight_pad_bytes(row_bytes)
assert padded % _WEIGHT_CLIFF_BYTES != 0 or row_bytes % _WEIGHT_CLIFF_BYTES
def test_act_pad_bytes_only_moves_strides_on_the_cliff():
"""Activations sit on a wider cliff than the packed weight: 2048 B."""
# On the cliff (K % 1024 == 0 for a 2-byte dtype): 2048 B multiples.
for row_bytes in (2048, 4096, 8192, 16384, 20480, 24576, 43008):
assert _act_pad_bytes(row_bytes) == _STRIDE_PAD_BYTES
# Off it, including 1024 B multiples that are not 2048 B ones -- the
# weight cliff must not be applied to activations.
for row_bytes in (1024, 3072, 5120, 10752, 19456, 37888):
assert _act_pad_bytes(row_bytes) == 0
# The padded stride is never back on the cliff.
for row_bytes in range(16, 65536, 16):
padded = row_bytes + _act_pad_bytes(row_bytes)
assert padded % _ACT_CLIFF_BYTES != 0 or row_bytes % _ACT_CLIFF_BYTES
def _row_padded_copy(t: torch.Tensor, pad_cols: int) -> torch.Tensor:
"""Copy of ``t`` whose stride(0) is ``t.shape[1] + pad_cols``."""
rows, cols = t.shape
buf = torch.empty((rows, cols + pad_cols), dtype=t.dtype, device=t.device)
buf[:, :cols].copy_(t)
return buf[:, :cols]
@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])
@pytest.mark.parametrize("M", [1, MAX_SKINNY_BATCH_SIZE + 1])
def test_weight_row_padding_does_not_change_results(dtype, has_zp, M):
"""Padding the weight rows is a pure layout change.
Both the HIP skinny kernel (M small) and the Triton prefill kernel take the
weight row stride from the tensor, so a padded layout must produce exactly
the same numbers as the dense one.
"""
from vllm.utils.platform_utils import num_compute_units
set_random_seed(0)
K, N, G = 512, 256, 128
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)
w_q = pack_int4_exllama_shuffle(w_int4_nk).view(torch.int8)
scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to(
dtype
)
zp = (
_pack_zp_rows_for_kernel(
torch.randint(0, 16, (N, K // G), device=device, dtype=torch.int32)
)
if has_zp
else None
)
# 32 int32 columns = 128 B of pad, the production pad size.
w_q_pad = _row_padded_copy(w_q.view(torch.int32), 32).view(torch.int8)
assert w_q_pad.stride(0) == w_q.stride(0) + _STRIDE_PAD_BYTES
cu_count = num_compute_units()
dense = torch.ops.vllm.rdna_hybrid_w4a16_apply(
a, w_q, scales, zp, None, cu_count, G
)
padded = torch.ops.vllm.rdna_hybrid_w4a16_apply(
a, w_q_pad, scales, zp, None, cu_count, G
)
torch.testing.assert_close(padded, dense, rtol=0, atol=0)
@pytest.mark.skipif(
not hybrid_module._on_gfx1151(),
reason="row-stride padding is only enabled on gfx1151",
)
@pytest.mark.parametrize(
"K,expected_weight_stride",
[
(512, 256), # 256 B row: off the cliff, stored dense
(8192, 4096 + 128), # 4096 B row: on the cliff, padded
],
)
def test_process_weights_pads_cliff_rows(K, expected_weight_stride, dist_init):
"""The stored weight row stride follows the padding rule."""
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
MPLinearLayerConfig,
)
from vllm.scalar_type import scalar_types
set_random_seed(0)
N, G = 128, 128
w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
layer = _build_dummy_layer(
_pack_int4_along_k_to_ckpt(w_int4_kn),
0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16),
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,
)
RDNAHybridW4A16LinearKernel(
config,
w_q_param_name="weight_packed",
w_s_param_name="weight_scale",
w_zp_param_name=None,
).process_weights_after_loading(layer)
# Weight rows: stride is in int8 elements, i.e. bytes.
assert layer.weight_packed.stride(0) == expected_weight_stride
# The int32 view the Triton path uses survives the padded stride.
assert tuple(layer.weight_packed.view(torch.int32).shape) == (N, K // 8)
# Metadata is untouched by this change.
assert layer.weight_scale.is_contiguous()
vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
+78 -7
View File
@@ -84,6 +84,8 @@ def _triton_w4a16_skinny_fmt_kernel(
K,
K8, # K // 8
num_groups, # K // group_size
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,
@@ -127,12 +129,12 @@ def _triton_w4a16_skinny_fmt_kernel(
offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K)
mask_k = offs_k < K
a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
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] * K8 + offs_k8[None, :]
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)
@@ -210,14 +212,16 @@ def triton_w4a16_skinny_fmt_gemm(
Output matrix [M, N], same dtype as a.
"""
assert a.is_contiguous(), "Activation matrix must be contiguous"
assert b_q.is_contiguous(), "Weight matrix must be contiguous"
assert a.stride(1) == 1, "Activation rows must be contiguous"
assert b_q.stride(1) == 1, "Weight rows must be contiguous"
assert scales.is_contiguous(), "Scales must be contiguous"
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), (
@@ -331,6 +335,8 @@ def triton_w4a16_skinny_fmt_gemm(
K,
K8,
num_groups,
stride_bn,
stride_am,
group_size=group_size,
ZP_BIAS=zp_bias,
HAS_ZP=has_zp,
@@ -368,11 +374,77 @@ def pack_int4_exllama_shuffle(w_uint4: torch.Tensor) -> torch.Tensor:
)
# gfx11 row strides alias in the vector caches when they land on a power-of-two
# multiple.
_WEIGHT_CLIFF_BYTES = 1024 # packed weight row, K/2 B
_ACT_CLIFF_BYTES = 2048 # activation row, K*2 B
_STRIDE_PAD_BYTES = 128 # one cache line
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
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
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]
# ---------------------------------------------------------------------------
# Hybrid dispatch logic
# ---------------------------------------------------------------------------
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
def _rdna_hybrid_w4a16_apply_impl(
x_2d: torch.Tensor,
w_q: torch.Tensor,
@@ -410,7 +482,7 @@ def _rdna_hybrid_w4a16_apply_impl(
)
with ctx:
output = triton_w4a16_skinny_fmt_gemm(
a=x_2d,
a=_pad_activation_rows(x_2d),
b_q=w_q.view(torch.int32),
scales=w_s,
group_size=group_size,
@@ -513,9 +585,8 @@ class RDNAHybridW4A16LinearKernel(MPLinearKernel):
if getattr(w_q_raw, "output_dim", 0) != 0:
unpacked = unpacked.t().contiguous()
shuffled = pack_int4_exllama_shuffle(unpacked)
# Store as int8; Triton reinterprets via .view(torch.int32) at apply time.
w_q_skinny = shuffled.contiguous().view(torch.int8)
w_q_skinny = pack_skinny_int4(unpacked)
permute_param_layout_(w_s_raw, input_dim=1, output_dim=0)
w_s_skinny = w_s_raw.data.contiguous()