diff --git a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py index 2d5ec80c6a..7ce49e2e11 100644 --- a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py +++ b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py @@ -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, ) diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index a20ca41d41..4588b216fe 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -8,6 +8,7 @@ #include #include +#include // 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 @@ -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 -__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 -__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(b_row_stride_bytes), "Weight row stride (", + b_row_stride_bytes, " bytes) exceeds int range"); + const int b_row_stride_bytes_i32 = static_cast(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_<<>>( \ 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_<<>>( \ 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) \ diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index fe3db8dab3..d6045a791c 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -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() 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 86e029b5f8..9d53034699 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 @@ -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()