@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 382bbd5144

[ROCm][Kernel] Add HybridW4A16LinearKernel: Triton prefill + HIP skinny decode (#40977)

Signed-off-by: Matthias Gehre <matthias.gehre@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

Matthias Gehre committed Jul 14, 2026 at 05:22 UTC 382bbd51448b2f58c73b3e51d051bc352166ba91
11 files changed +2165 -5
CMakeLists.txt
+1
@@ -1364,6 +1364,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
1364 set(VLLM_ROCM_EXT_SRC
1365 "csrc/rocm/torch_bindings.cpp"
1366 "csrc/rocm/skinny_gemms.cu"
1367 + "csrc/rocm/skinny_gemms_int4.cu"
1368 "csrc/rocm/attention.cu")
1369
1370 set(VLLM_ROCM_HAS_GFX1100 OFF)
benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py new
+201
@@ -0,0 +1,201 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes.
5 +
6 +Usage:
7 + python benchmark_int4_gemm.py
8 + python benchmark_int4_gemm.py --models Qwen/Qwen3-4B
9 + python benchmark_int4_gemm.py --group-size 128
10 +"""
11 +
12 +import argparse
13 +import copy
14 +import itertools
15 +import os
16 +
17 +import torch
18 +
19 +from vllm.triton_utils import triton
20 +
21 +# ---------------------------------------------------------------------------
22 +# Weight shapes: [K, N], TP_SPLIT_DIM
23 +# ---------------------------------------------------------------------------
24 +WEIGHT_SHAPES = {
25 + "Qwen/Qwen3-4B": [
26 + ([2560, 3840], 1), # qkv_proj
27 + ([2560, 2560], 0), # o_proj
28 + ([2560, 19456], 1), # gate_up_proj
29 + ([9728, 2560], 0), # down_proj
30 + ],
31 + "Qwen/Qwen2.5-7B-Instruct": [
32 + ([3584, 4608], 1),
33 + ([3584, 3584], 0),
34 + ([3584, 37888], 1),
35 + ([18944, 3584], 0),
36 + ],
37 + "trymirai/SmolLM2-1.7B-Instruct-AWQ": [
38 + ([2048, 6144], 1), # qkv_proj
39 + ([2048, 2048], 0), # o_proj
40 + ([2048, 16384], 1), # gate_up_proj
41 + ([8192, 2048], 0), # down_proj
42 + ],
43 + "RedHatAI/Qwen3-8B-quantized.w4a16": [
44 + ([4096, 6144], 1), # qkv_proj
45 + ([4096, 4096], 0), # o_proj
46 + ([4096, 24576], 1), # gate_up_proj
47 + ([12288, 4096], 0), # down_proj
48 + ],
49 +}
50 +
51 +
52 +# ---------------------------------------------------------------------------
53 +# Weight packing
54 +# ---------------------------------------------------------------------------
55 +def prepare_hybrid_weights(K, N, group_size, device="cuda"):
56 + """Create random weights for benchmarking.
57 +
58 + Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives
59 + its int32 view from w_q_skinny, so no separate int32 buffer is returned.
60 + """
61 + num_groups = K // group_size
62 +
63 + # Random packed weights — actual values don't matter for throughput
64 + w_q_skinny_i32 = torch.randint(
65 + 0, 2**31, (N, K // 8), dtype=torch.int32, device=device
66 + )
67 + w_q_skinny = w_q_skinny_i32.view(torch.int8).contiguous()
68 + w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01
69 +
70 + # Raw per-group zero-points for asymmetric benchmarks
71 + w_zp = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device).to(
72 + torch.float16
73 + )
74 +
75 + # FP16 baseline for F.linear
76 + w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01
77 +
78 + return w_q_skinny, w_s_skinny, w_fp16, w_zp
79 +
80 +
81 +# ---------------------------------------------------------------------------
82 +# Benchmark
83 +# ---------------------------------------------------------------------------
84 +PROVIDERS = ["torch-fp16", "hybrid-w4a16", "hybrid-w4a16-zp"]
85 +
86 +
87 +@triton.testing.perf_report(
88 + triton.testing.Benchmark(
89 + x_names=["batch_size"],
90 + x_vals=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],
91 + x_log=False,
92 + line_arg="provider",
93 + line_vals=PROVIDERS,
94 + line_names=PROVIDERS,
95 + ylabel="TFLOP/s (larger is better)",
96 + plot_name="FP16 vs Hybrid W4A16",
97 + args={},
98 + )
99 +)
100 +def benchmark(batch_size, provider, N, K, group_size, weights):
101 + M = batch_size
102 + device = "cuda"
103 + dtype = torch.float16
104 + a = torch.randn((M, K), device=device, dtype=dtype)
105 +
106 + quantiles = [0.5, 0.2, 0.8]
107 +
108 + if provider == "torch-fp16":
109 + w_fp16 = weights["w_fp16"]
110 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
111 + lambda: torch.nn.functional.linear(a, w_fp16),
112 + quantiles=quantiles,
113 + )
114 + elif provider in ("hybrid-w4a16", "hybrid-w4a16-zp"):
115 + from vllm.model_executor.kernels.linear.mixed_precision import (
116 + rdna_hybrid_w4a16 as _k,
117 + )
118 +
119 + _rdna_hybrid_w4a16_apply_impl = _k._rdna_hybrid_w4a16_apply_impl
120 + from vllm.utils.platform_utils import num_compute_units
121 +
122 + w = weights
123 + cu_count = num_compute_units()
124 + use_zp = provider == "hybrid-w4a16-zp"
125 +
126 + def run():
127 + return _rdna_hybrid_w4a16_apply_impl(
128 + a,
129 + w["w_q_skinny"],
130 + w["w_s_skinny"],
131 + w["w_zp"] if use_zp else None,
132 + None, # bias
133 + cu_count,
134 + group_size,
135 + )
136 +
137 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
138 + run,
139 + quantiles=quantiles,
140 + )
141 + else:
142 + return 0.0, 0.0, 0.0
143 +
144 + to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
145 + return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
146 +
147 +
148 +def prepare_shapes(args):
149 + KN_model_names = []
150 + for model, tp_size in itertools.product(args.models, args.tp_sizes):
151 + for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
152 + KN[tp_dim] //= tp_size
153 + KN.append(model)
154 + KN_model_names.append(KN)
155 + return KN_model_names
156 +
157 +
158 +if __name__ == "__main__":
159 + parser = argparse.ArgumentParser(
160 + description="Benchmark RDNAHybridW4A16LinearKernel"
161 + )
162 + parser.add_argument(
163 + "--models",
164 + nargs="+",
165 + type=str,
166 + default=["Qwen/Qwen3-4B"],
167 + choices=list(WEIGHT_SHAPES.keys()),
168 + )
169 + parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
170 + parser.add_argument("--group-size", type=int, default=128)
171 + parser.add_argument("--save-path", type=str, default=None)
172 + args = parser.parse_args()
173 +
174 + for K, N, model in prepare_shapes(args):
175 + group_size = args.group_size
176 + print(f"\n{'=' * 70}")
177 + print(f"{model}, N={N} K={K}, group_size={group_size}")
178 + print(f"{'=' * 70}")
179 +
180 + w_q_skinny, w_s_skinny, w_fp16, w_zp = prepare_hybrid_weights(K, N, group_size)
181 +
182 + weights = {
183 + "w_q_skinny": w_q_skinny,
184 + "w_s_skinny": w_s_skinny,
185 + "w_fp16": w_fp16,
186 + "w_zp": w_zp,
187 + }
188 +
189 + save_path = args.save_path or f"bench_int4_res_n{N}_k{K}"
190 + os.makedirs(save_path, exist_ok=True)
191 + benchmark.run(
192 + print_data=True,
193 + show_plots=False,
194 + save_path=save_path,
195 + N=N,
196 + K=K,
197 + group_size=group_size,
198 + weights=weights,
199 + )
200 +
201 + print("\nBenchmark finished!")
csrc/rocm/ops.h
+6
@@ -9,6 +9,12 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
9 const std::optional<at::Tensor>& in_bias,
10 const int64_t CuCount);
11
12 +torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
13 + const at::Tensor& in_scale,
14 + const std::optional<at::Tensor>& in_zero_points,
15 + const std::optional<at::Tensor>& in_bias,
16 + const int64_t CuCount, const int64_t group_size);
17 +
18 torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
19 const std::optional<at::Tensor>& in_bias,
20 const int64_t CuCount);
csrc/rocm/skinny_gemms_int4.cu new
+795
@@ -0,0 +1,795 @@
1 +#include <torch/all.h>
2 +#include <ATen/cuda/CUDAContext.h>
3 +#include <c10/cuda/CUDAGuard.h>
4 +
5 +#include <cuda_runtime.h>
6 +#include <cuda_fp16.h>
7 +#include <cuda_bf16.h>
8 +
9 +#include <stdexcept>
10 +#include <algorithm>
11 +
12 +#include "../cuda_compat.h"
13 +#include "dispatch_utils.h"
14 +
15 +// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts
16 +#if defined(__GFX11__) || defined(__GFX12__)
17 + #define __HIP__GFX1X__
18 +#endif
19 +
20 +#define LDS_SIZE 64 * 1024
21 +
22 +static constexpr int get_lds_size_int4() { return 64 * 1024; }
23 +
24 +#if defined(NDEBUG)
25 + #undef NDEBUG
26 + #include <assert.h>
27 + #define UNREACHABLE_CODE assert(false);
28 + #define NDEBUG
29 +#else
30 + #define UNREACHABLE_CODE assert(false);
31 +#endif
32 +
33 +template <typename T>
34 +__device__ __forceinline__ float __s2float(T v);
35 +
36 +template <>
37 +__device__ __forceinline__ float __s2float(half v) {
38 + return __half2float(v);
39 +}
40 +
41 +template <>
42 +__device__ __forceinline__ float __s2float(__hip_bfloat16 v) {
43 + return __bfloat162float(v);
44 +}
45 +
46 +template <typename T>
47 +__device__ __forceinline__ T __float2s(float v);
48 +
49 +template <>
50 +__device__ __forceinline__ half __float2s(float v) {
51 + return __float2half(v);
52 +}
53 +
54 +template <>
55 +__device__ __forceinline__ __hip_bfloat16 __float2s(float v) {
56 + return __float2bfloat16(v);
57 +}
58 +
59 +template <typename T>
60 +__device__ __forceinline__ T loadnt(T* addr) {
61 + return __builtin_nontemporal_load(addr);
62 +}
63 +
64 +template <typename T>
65 +struct scalar {};
66 +
67 +template <>
68 +struct scalar<c10::Half> {
69 + using type = half;
70 +};
71 +
72 +template <>
73 +struct scalar<c10::BFloat16> {
74 + using type = __hip_bfloat16;
75 +};
76 +
77 +#define DOT2C(V0, V2, V3) \
78 + if constexpr (std::is_same_v<scalar_t, half>) { \
79 + V0 = __builtin_amdgcn_fdot2(*((half2*)(&(V2))), *((half2*)(&(V3))), V0, \
80 + false); \
81 + } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) { \
82 + typedef short __attribute__((ext_vector_type(2))) bf16x2_t; \
83 + V0 = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&(V2))), \
84 + *((bf16x2_t*)(&(V3))), V0, false); \
85 + }
86 +
87 +// DPP-based wave32 reduction for GFX1X (matches skinny_gemms.cu).
88 +// row_shr within 16-lane rows, then cross-row __shfl_xor.
89 +// Result accumulates in lane (THRDS - 1).
90 +#if defined(__HIP__GFX1X__)
91 + #define REDUCE_SUM_DPP_WAVE32(val) \
92 + do { \
93 + val += __builtin_amdgcn_mov_dpp(val, 0x118, 0xf, 0xf, 1); /* shr8 */ \
94 + val += __builtin_amdgcn_mov_dpp(val, 0x114, 0xf, 0xf, 1); /* shr4 */ \
95 + val += __builtin_amdgcn_mov_dpp(val, 0x112, 0xf, 0xf, 1); /* shr2 */ \
96 + val += __builtin_amdgcn_mov_dpp(val, 0x111, 0xf, 0xf, 1); /* shr1 */ \
97 + val += __shfl_xor(val, 16); \
98 + } while (0)
99 +#endif
100 +
101 +__device__ inline unsigned int min__(uint32_t a, uint32_t b) {
102 + return min(a, b);
103 +}
104 +
105 +// W4A16 skinny GEMM kernel: packed int4 weights, fp16/bf16 activations
106 +// Targets the "sml" case where activations fit in LDS.
107 +// A_CHUNK: number of K-elements processed per thread per step.
108 +// Weights are packed as 2 int4 values per byte: [M, K/2] bytes.
109 +// GROUP_SIZE: 0 = per-channel scale [M], >0 = per-group scale [M,
110 +// K/GROUP_SIZE].
111 +// Requires GROUP_SIZE % A_CHUNK == 0 when GROUP_SIZE > 0.
112 +#if defined(__HIP__GFX1X__)
113 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
114 + int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
115 +__global__ void __launch_bounds__(WvPrGrp* THRDS)
116 + wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By,
117 + const uint8_t* B_packed,
118 + const scalar_t* __restrict__ A, const scalar_t* scale,
119 + const scalar_t* zero_points,
120 + const scalar_t* __restrict__ BIAS, scalar_t* C,
121 + const int _WvPrGrp, const int CuCount) {
122 + constexpr int max_lds_len = LDS_SIZE / 2;
123 + const int K_packed = K / 2;
124 +
125 + union bigTypeA {
126 + scalar_t h[A_CHUNK];
127 + float f[A_CHUNK / 2];
128 + };
129 +
130 + union bigTypeW {
131 + uint8_t b[A_CHUNK / 2];
132 + uint32_t u32[A_CHUNK / 8];
133 + float f[A_CHUNK / 8];
134 + };
135 +
136 + __shared__ scalar_t s[max_lds_len];
137 +
138 + for (uint32_t k = 0; k < min__(K * N, max_lds_len);
139 + k += THRDS * WvPrGrp * A_CHUNK) {
140 + uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
141 +
142 + if (k_in >= min__(K * N, max_lds_len)) break;
143 +
144 + *((bigTypeA*)(&s[k_in])) = *((bigTypeA*)(&A[k_in]));
145 + }
146 + __syncthreads();
147 +
148 + if (threadIdx.y >= _WvPrGrp) return;
149 +
150 + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
151 +
152 + // For per-group, precompute num_groups and scale stride
153 + [[maybe_unused]] const int num_groups =
154 + (GROUP_SIZE > 0) ? (K / GROUP_SIZE) : 0;
155 +
156 + float sum[N][YTILE];
157 +
158 + while (m < M) {
159 + for (int i = 0; i < YTILE; i++)
160 + for (int n = 0; n < N; n++) sum[n][i] = 0;
161 +
162 + bigTypeA bigA[N][UNRL];
163 + bigTypeW bigB[YTILE][UNRL];
164 +
165 + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
166 + #pragma unroll
167 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
168 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
169 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
170 + if (k_ >= K) break;
171 +
172 + const uint8_t* B_ = &B_packed[(m + 0) * K_packed + k_ / 2];
173 + for (int y = 0; y < YTILE; y++) {
174 + const float* src = (const float*)(&B_[y * K_packed]);
175 + #pragma unroll
176 + for (int i = 0; i < A_CHUNK / 8; i++)
177 + bigB[y][k2].f[i] = loadnt((float*)&src[i]);
178 + }
179 + }
180 +
181 + #pragma unroll
182 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
183 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
184 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
185 + if (k_ >= K) break;
186 +
187 + for (int n = 0; n < N; n++) {
188 + bigA[n][k2] = *((const bigTypeA*)(&(s[k_ + K * n])));
189 + }
190 + }
191 +
192 + #pragma unroll
193 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
194 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
195 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
196 + if (k_ >= K) break;
197 +
198 + #pragma unroll
199 + for (uint32_t n = 0; n < N; n++) {
200 + #pragma unroll
201 + for (int y = 0; y < YTILE; y++) {
202 + bigTypeA cvtB;
203 +
204 + if constexpr (std::is_same_v<scalar_t, half>) {
205 + constexpr uint32_t FP16_MAGIC = 0x64006400u;
206 + // When HAS_ZERO_POINTS, store raw nibble values;
207 + // the zero-point subtraction below handles the full shift.
208 + // When symmetric, bake -8 into the constants.
209 + constexpr uint32_t BIAS_LO =
210 + HAS_ZERO_POINTS ? 0x64006400u : 0x64086408u;
211 + constexpr uint32_t SCALE16 = 0x2C002C00u;
212 + constexpr uint32_t BIAS_HI =
213 + HAS_ZERO_POINTS ? 0xD400D400u : 0xD480D480u;
214 + #pragma unroll
215 + for (uint32_t w = 0; w < A_CHUNK / 8; w++) {
216 + uint32_t qa = bigB[y][k2].u32[w];
217 + uint32_t lo0 = (qa & 0x000F000Fu) | FP16_MAGIC;
218 + uint32_t hi0 = (qa & 0x00F000F0u) | FP16_MAGIC;
219 + qa >>= 8;
220 + uint32_t lo1 = (qa & 0x000F000Fu) | FP16_MAGIC;
221 + uint32_t hi1 = (qa & 0x00F000F0u) | FP16_MAGIC;
222 +
223 + *(half2*)&cvtB.f[w * 4 + 0] =
224 + __hsub2(*(half2*)&lo0, *(const half2*)&BIAS_LO);
225 + *(half2*)&cvtB.f[w * 4 + 1] =
226 + __hfma2(*(half2*)&hi0, *(const half2*)&SCALE16,
227 + *(const half2*)&BIAS_HI);
228 + *(half2*)&cvtB.f[w * 4 + 2] =
229 + __hsub2(*(half2*)&lo1, *(const half2*)&BIAS_LO);
230 + *(half2*)&cvtB.f[w * 4 + 3] =
231 + __hfma2(*(half2*)&hi1, *(const half2*)&SCALE16,
232 + *(const half2*)&BIAS_HI);
233 + }
234 + } else {
235 + // bf16 path: marlin-style magic-number trick.
236 + // Skip the expensive bias subtraction (avoids fp32 round-trip
237 + // on gfx1151); use magic values directly in DOT2C and correct
238 + // the accumulator with bias * sum(activations) afterward.
239 + constexpr uint32_t BF16_MAGIC = 0x43004300u;
240 + #pragma unroll
241 + for (uint32_t w = 0; w < A_CHUNK / 8; w++) {
242 + uint32_t qa = bigB[y][k2].u32[w];
243 + *(uint32_t*)&cvtB.f[w * 4 + 0] =
244 + (qa & 0x000F000Fu) | BF16_MAGIC;
245 + qa >>= 4;
246 + *(uint32_t*)&cvtB.f[w * 4 + 1] =
247 + (qa & 0x000F000Fu) | BF16_MAGIC;
248 + qa >>= 4;
249 + *(uint32_t*)&cvtB.f[w * 4 + 2] =
250 + (qa & 0x000F000Fu) | BF16_MAGIC;
251 + qa >>= 4;
252 + *(uint32_t*)&cvtB.f[w * 4 + 3] =
253 + (qa & 0x000F000Fu) | BF16_MAGIC;
254 + }
255 + }
256 +
257 + if constexpr (!std::is_same_v<scalar_t, __hip_bfloat16>) {
258 + if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {
259 + uint32_t group_idx = k_ / GROUP_SIZE;
260 + scalar_t zp = zero_points[(m + y) * num_groups + group_idx];
261 + #pragma unroll
262 + for (uint32_t b = 0; b < A_CHUNK; b++) {
263 + cvtB.h[b] = cvtB.h[b] - zp;
264 + }
265 + }
266 + }
267 +
268 + if constexpr (GROUP_SIZE > 0) {
269 + float partial = 0;
270 + #pragma unroll
271 + for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
272 + DOT2C(partial, bigA[n][k2].f[b], cvtB.f[b])
273 + }
274 + if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
275 + constexpr uint32_t BF16_ONES = 0x3F803F80u;
276 + float act_sum = 0.0f;
277 + #pragma unroll
278 + for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
279 + DOT2C(act_sum, bigA[n][k2].f[b], *(const float*)&BF16_ONES)
280 + }
281 + if constexpr (HAS_ZERO_POINTS) {
282 + uint32_t group_idx_zp = k_ / GROUP_SIZE;
283 + float zp_f = __s2float(
284 + zero_points[(m + y) * num_groups + group_idx_zp]);
285 + partial -= (128.0f + zp_f) * act_sum;
286 + } else {
287 + partial -= 136.0f * act_sum;
288 + }
289 + }
290 + uint32_t group_idx = k_ / GROUP_SIZE;
291 + sum[n][y] +=
292 + partial * __s2float(scale[(m + y) * num_groups + group_idx]);
293 + } else {
294 + #pragma unroll
295 + for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
296 + DOT2C(sum[n][y], bigA[n][k2].f[b], cvtB.f[b])
297 + }
298 + if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
299 + constexpr float BIAS_VAL = HAS_ZERO_POINTS ? 128.0f : 136.0f;
300 + constexpr uint32_t BF16_ONES = 0x3F803F80u;
301 + float act_sum = 0.0f;
302 + #pragma unroll
303 + for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
304 + DOT2C(act_sum, bigA[n][k2].f[b], *(const float*)&BF16_ONES)
305 + }
306 + sum[n][y] -= BIAS_VAL * act_sum;
307 + }
308 + }
309 + }
310 + }
311 + }
312 + }
313 +
314 + for (int n = 0; n < N; n++)
315 + for (int y = 0; y < YTILE; y++) REDUCE_SUM_DPP_WAVE32(sum[n][y]);
316 +
317 + if (threadIdx.x == (THRDS - 1)) {
318 + for (int n = 0; n < N; n++) {
319 + for (int i = 0; i < YTILE; i++) {
320 + if constexpr (GROUP_SIZE == 0) {
321 + sum[n][i] *= __s2float(scale[m + i]);
322 + }
323 + if (BIAS) sum[n][i] += __s2float(BIAS[(m + i) % Bx + (n % By) * M]);
324 + C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
325 + }
326 + }
327 + }
328 + m += CuCount * _WvPrGrp * YTILE;
329 + }
330 +}
331 +#else // !defined(__HIP__GFX1X__)
332 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
333 + int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
334 +__global__ void wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx,
335 + const int By, const uint8_t* B_packed,
336 + const scalar_t* __restrict__ A,
337 + const scalar_t* scale,
338 + const scalar_t* zero_points,
339 + const scalar_t* __restrict__ BIAS,
340 + scalar_t* C, const int _WvPrGrp,
341 + const int CuCount) {
342 + UNREACHABLE_CODE
343 +}
344 +#endif // defined(__HIP__GFX1X__)
345 +
346 +// W4A16 skinny GEMM "medium" kernel: activation matrix marginally exceeds LDS.
347 +// Loads as much of A into LDS as fits; overflowing rows fall back to global
348 +// memory. Also handles M not divisible by YTILE via commitColumn tracking.
349 +#if defined(__HIP__GFX1X__)
350 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
351 + int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
352 +__global__ void __launch_bounds__(WvPrGrp* THRDS)
353 + wvSplitK_int4_hf_(const int K, const int M, const int Bx, const int By,
354 + const uint8_t* B_packed, const scalar_t* __restrict__ A,
355 + const scalar_t* scale, const scalar_t* zero_points,
356 + const scalar_t* __restrict__ BIAS, scalar_t* C,
357 + const int _WvPrGrp, const int CuCount) {
358 + constexpr int max_lds_len = LDS_SIZE / 2;
359 + const int K_packed = K / 2;
360 +
361 + union bigTypeA {
362 + scalar_t h[A_CHUNK];
363 + float f[A_CHUNK / 2];
364 + };
365 +
366 + union bigTypeW {
367 + uint8_t b[A_CHUNK / 2];
368 + uint32_t u32[A_CHUNK / 8];
369 + float f[A_CHUNK / 8];
370 + };
371 +
372 + __shared__ scalar_t s[max_lds_len];
373 +
374 + uint32_t commitColumn[YTILE];
375 + for (uint32_t i = 0; i < YTILE; i++) {
376 + commitColumn[i] = 1;
377 + }
378 +
379 + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
380 +
381 + if (m < M && (m + YTILE) >= M) {
382 + uint32_t startColumn = M - YTILE;
383 + for (uint32_t i = 0; i < (m - startColumn); i++) {
384 + commitColumn[i] = 0;
385 + }
386 + m = startColumn;
387 + }
388 +
389 + for (uint32_t k = 0; k < min__(K * N, max_lds_len);
390 + k += THRDS * WvPrGrp * A_CHUNK) {
391 + uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
392 +
393 + if (k_in >= min__(K * N, max_lds_len)) break;
394 +
395 + *((bigTypeA*)(&s[k_in])) = *((bigTypeA*)(&A[k_in]));
396 + }
397 + __syncthreads();
398 +
399 + if (threadIdx.y >= _WvPrGrp) return;
400 +
401 + [[maybe_unused]] const int num_groups =
402 + (GROUP_SIZE > 0) ? (K / GROUP_SIZE) : 0;
403 +
404 + float sum[N][YTILE];
405 +
406 + while (m < M) {
407 + for (int i = 0; i < YTILE; i++)
408 + for (int n = 0; n < N; n++) sum[n][i] = 0;
409 +
410 + bigTypeA bigA[N][UNRL];
411 + bigTypeW bigB[YTILE][UNRL];
412 +
413 + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
414 + #pragma unroll
415 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
416 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
417 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
418 + if (k_ >= K) break;
419 +
420 + const uint8_t* B_ = &B_packed[(m + 0) * K_packed + k_ / 2];
421 + for (int y = 0; y < YTILE; y++) {
422 + const float* src = (const float*)(&B_[y * K_packed]);
423 + #pragma unroll
424 + for (int i = 0; i < A_CHUNK / 8; i++)
425 + bigB[y][k2].f[i] = loadnt((float*)&src[i]);
426 + }
427 + }
428 +
429 + #pragma unroll
430 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
431 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
432 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
433 + if (k_ >= K) break;
434 +
435 + for (int n = 0; n < N; n++) {
436 + if (k_ + K * n < max_lds_len)
437 + bigA[n][k2] = *((const bigTypeA*)(&(s[k_ + K * n])));
438 + else
439 + bigA[n][k2] = *((const bigTypeA*)(&(A[k_ + K * n])));
440 + }
441 + }
442 +
443 + #pragma unroll
444 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
445 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
446 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
447 + if (k_ >= K) break;
448 +
449 + #pragma unroll
450 + for (uint32_t n = 0; n < N; n++) {
451 + #pragma unroll
452 + for (int y = 0; y < YTILE; y++) {
453 + bigTypeA cvtB;
454 +
455 + if constexpr (std::is_same_v<scalar_t, half>) {
456 + constexpr uint32_t FP16_MAGIC = 0x64006400u;
457 + // When HAS_ZERO_POINTS, store raw nibble values;
458 + // the zero-point subtraction below handles the full shift.
459 + // When symmetric, bake -8 into the constants.
460 + constexpr uint32_t BIAS_LO =
461 + HAS_ZERO_POINTS ? 0x64006400u : 0x64086408u;
462 + constexpr uint32_t SCALE16 = 0x2C002C00u;
463 + constexpr uint32_t BIAS_HI =
464 + HAS_ZERO_POINTS ? 0xD400D400u : 0xD480D480u;
465 + #pragma unroll
466 + for (uint32_t w = 0; w < A_CHUNK / 8; w++) {
467 + uint32_t qa = bigB[y][k2].u32[w];
468 + uint32_t lo0 = (qa & 0x000F000Fu) | FP16_MAGIC;
469 + uint32_t hi0 = (qa & 0x00F000F0u) | FP16_MAGIC;
470 + qa >>= 8;
471 + uint32_t lo1 = (qa & 0x000F000Fu) | FP16_MAGIC;
472 + uint32_t hi1 = (qa & 0x00F000F0u) | FP16_MAGIC;
473 +
474 + *(half2*)&cvtB.f[w * 4 + 0] =
475 + __hsub2(*(half2*)&lo0, *(const half2*)&BIAS_LO);
476 + *(half2*)&cvtB.f[w * 4 + 1] =
477 + __hfma2(*(half2*)&hi0, *(const half2*)&SCALE16,
478 + *(const half2*)&BIAS_HI);
479 + *(half2*)&cvtB.f[w * 4 + 2] =
480 + __hsub2(*(half2*)&lo1, *(const half2*)&BIAS_LO);
481 + *(half2*)&cvtB.f[w * 4 + 3] =
482 + __hfma2(*(half2*)&hi1, *(const half2*)&SCALE16,
483 + *(const half2*)&BIAS_HI);
484 + }
485 + } else {
486 + // bf16 path: marlin-style magic-number trick.
487 + // Skip the expensive bias subtraction (avoids fp32 round-trip
488 + // on gfx1151); use magic values directly in DOT2C and correct
489 + // the accumulator with bias * sum(activations) afterward.
490 + constexpr uint32_t BF16_MAGIC = 0x43004300u;
491 + #pragma unroll
492 + for (uint32_t w = 0; w < A_CHUNK / 8; w++) {
493 + uint32_t qa = bigB[y][k2].u32[w];
494 + *(uint32_t*)&cvtB.f[w * 4 + 0] =
495 + (qa & 0x000F000Fu) | BF16_MAGIC;
496 + qa >>= 4;
497 + *(uint32_t*)&cvtB.f[w * 4 + 1] =
498 + (qa & 0x000F000Fu) | BF16_MAGIC;
499 + qa >>= 4;
500 + *(uint32_t*)&cvtB.f[w * 4 + 2] =
501 + (qa & 0x000F000Fu) | BF16_MAGIC;
502 + qa >>= 4;
503 + *(uint32_t*)&cvtB.f[w * 4 + 3] =
504 + (qa & 0x000F000Fu) | BF16_MAGIC;
505 + }
506 + }
507 +
508 + if constexpr (!std::is_same_v<scalar_t, __hip_bfloat16>) {
509 + if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {
510 + uint32_t group_idx = k_ / GROUP_SIZE;
511 + scalar_t zp = zero_points[(m + y) * num_groups + group_idx];
512 + #pragma unroll
513 + for (uint32_t b = 0; b < A_CHUNK; b++) {
514 + cvtB.h[b] = cvtB.h[b] - zp;
515 + }
516 + }
517 + }
518 +
519 + if constexpr (GROUP_SIZE > 0) {
520 + float partial = 0;
521 + #pragma unroll
522 + for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
523 + DOT2C(partial, bigA[n][k2].f[b], cvtB.f[b])
524 + }
525 + if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
526 + constexpr uint32_t BF16_ONES = 0x3F803F80u;
527 + float act_sum = 0.0f;
528 + #pragma unroll
529 + for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
530 + DOT2C(act_sum, bigA[n][k2].f[b], *(const float*)&BF16_ONES)
531 + }
532 + if constexpr (HAS_ZERO_POINTS) {
533 + uint32_t group_idx_zp = k_ / GROUP_SIZE;
534 + float zp_f = __s2float(
535 + zero_points[(m + y) * num_groups + group_idx_zp]);
536 + partial -= (128.0f + zp_f) * act_sum;
537 + } else {
538 + partial -= 136.0f * act_sum;
539 + }
540 + }
541 + uint32_t group_idx = k_ / GROUP_SIZE;
542 + sum[n][y] +=
543 + partial * __s2float(scale[(m + y) * num_groups + group_idx]);
544 + } else {
545 + #pragma unroll
546 + for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
547 + DOT2C(sum[n][y], bigA[n][k2].f[b], cvtB.f[b])
548 + }
549 + if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
550 + constexpr float BIAS_VAL = HAS_ZERO_POINTS ? 128.0f : 136.0f;
551 + constexpr uint32_t BF16_ONES = 0x3F803F80u;
552 + float act_sum = 0.0f;
553 + #pragma unroll
554 + for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
555 + DOT2C(act_sum, bigA[n][k2].f[b], *(const float*)&BF16_ONES)
556 + }
557 + sum[n][y] -= BIAS_VAL * act_sum;
558 + }
559 + }
560 + }
561 + }
562 + }
563 + }
564 +
565 + for (int n = 0; n < N; n++)
566 + for (int y = 0; y < YTILE; y++) REDUCE_SUM_DPP_WAVE32(sum[n][y]);
567 +
568 + if (threadIdx.x == (THRDS - 1)) {
569 + for (int n = 0; n < N; n++) {
570 + for (int i = 0; i < YTILE; i++) {
571 + if (commitColumn[i]) {
572 + if constexpr (GROUP_SIZE == 0) {
573 + sum[n][i] *= __s2float(scale[m + i]);
574 + }
575 + if (BIAS) sum[n][i] += __s2float(BIAS[(m + i) % Bx + (n % By) * M]);
576 + C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
577 + }
578 + }
579 + }
580 + }
581 + m += CuCount * _WvPrGrp * YTILE;
582 +
583 + if (m < M && (m + YTILE) >= M) {
584 + uint32_t startColumn = M - YTILE;
585 + for (uint32_t i = 0; i < (m - startColumn); i++) {
586 + commitColumn[i] = 0;
587 + }
588 + m = startColumn;
589 + }
590 + }
591 +}
592 +#else // !defined(__HIP__GFX1X__)
593 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
594 + int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
595 +__global__ void wvSplitK_int4_hf_(const int K, const int M, const int Bx,
596 + const int By, const uint8_t* B_packed,
597 + const scalar_t* __restrict__ A,
598 + const scalar_t* scale,
599 + const scalar_t* zero_points,
600 + const scalar_t* __restrict__ BIAS,
601 + scalar_t* C, const int _WvPrGrp,
602 + const int CuCount) {
603 + UNREACHABLE_CODE
604 +}
605 +#endif // defined(__HIP__GFX1X__)
606 +
607 +static int mindiv_int4(int N, int div1, int div2) {
608 + int nPrRnd = div1 * div2;
609 + int limit = div2 < 13 ? div2 : 13;
610 + int rnds[16];
611 + for (int i = 0; i < limit; i++) {
612 + rnds[i] = (N + nPrRnd - 1) / nPrRnd;
613 + nPrRnd -= div1;
614 + }
615 + for (int i = limit - 1; i >= 0; i--)
616 + if (rnds[0] == rnds[i]) return (div2 - i);
617 + return 0;
618 +}
619 +
620 +// Per-group W4A16 skinny GEMM: packed int4 weights with group-wise scales.
621 +// in_a: packed int4 weights [M, K/2] (int8) or [M, K/8] (int32)
622 +// in_b: activations [N, K] (fp16/bf16)
623 +// in_scale: group scales [M, K/group_size] (fp16/bf16)
624 +// in_zero_points: optional raw zero points [M, K/group_size] (fp16/bf16)
625 +// If provided, kernel dequants as (nibble - zp_raw) * scale (asymmetric).
626 +// If absent, kernel dequants as (nibble - 8) * scale (symmetric uint4b8).
627 +// group_size: 32, 64, or 128
628 +
629 +torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
630 + const at::Tensor& in_scale,
631 + const std::optional<at::Tensor>& in_zero_points,
632 + const std::optional<at::Tensor>& in_bias,
633 + const int64_t CuCount, const int64_t group_size) {
634 + auto M_in = in_a.size(0);
635 + auto K_in = in_b.size(1);
636 + auto N_in = in_b.size(0);
637 + auto Bx_in =
638 + (in_bias.has_value() && in_bias->numel() > 0)
639 + ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
640 + : 1;
641 + auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
642 + in_bias->sizes().size() == 2)
643 + ? in_bias->size(0)
644 + : 1;
645 +
646 + int64_t expected_weight_bytes = M_in * K_in / 2;
647 + int64_t actual_weight_bytes = in_a.numel() * in_a.element_size();
648 + TORCH_CHECK(actual_weight_bytes == expected_weight_bytes,
649 + "Weight tensor must contain M*K/2 bytes for int4 packing");
650 + TORCH_CHECK(
651 + in_b.dtype() == torch::kFloat16 || in_b.dtype() == torch::kBFloat16,
652 + "Activation must be float16 or bfloat16");
653 + TORCH_CHECK(in_scale.dtype() == in_b.dtype(),
654 + "Scale dtype must match activation dtype");
655 + TORCH_CHECK(group_size == 32 || group_size == 64 || group_size == 128,
656 + "group_size must be 32, 64, or 128, got ", group_size);
657 + TORCH_CHECK(K_in % group_size == 0,
658 + "K must be divisible by group_size=", group_size);
659 + int64_t num_groups = K_in / group_size;
660 + TORCH_CHECK(in_scale.dim() == 2,
661 + "Scale must be 2D [M, K/group_size], got shape ",
662 + in_scale.sizes());
663 + TORCH_CHECK(in_scale.size(0) == M_in && in_scale.size(1) == num_groups,
664 + "Scale must be [M, K/group_size] = [", M_in, ", ", num_groups,
665 + "] but got [", in_scale.size(0), ", ", in_scale.size(1), "]");
666 + if (in_zero_points.has_value()) {
667 + TORCH_CHECK(in_zero_points->dtype() == in_b.dtype(),
668 + "Zero points dtype must match activation dtype");
669 + TORCH_CHECK(in_zero_points->dim() == 2,
670 + "Zero points must be 2D [M, K/group_size], got shape ",
671 + in_zero_points->sizes());
672 + TORCH_CHECK(in_zero_points->size(0) == M_in &&
673 + in_zero_points->size(1) == num_groups,
674 + "Zero points must be [M, K/group_size] = [", M_in, ", ",
675 + num_groups, "] but got [", in_zero_points->size(0), ", ",
676 + in_zero_points->size(1), "]");
677 + }
678 + TORCH_CHECK(K_in % 16 == 0, "K must be divisible by 16");
679 +
680 + const int max_lds_len = get_lds_size_int4() / 2;
681 + TORCH_CHECK(K_in * N_in <= (int64_t)(max_lds_len * 1.2),
682 + "K*N exceeds LDS capacity (medium limit). K=", K_in, " N=", N_in);
683 +
684 + auto out_c = torch::empty(
685 + {N_in, M_in},
686 + torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
687 +
688 + dim3 grid(CuCount);
689 +
690 + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
691 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
692 +
693 +// Dispatch macro: _HAS_ZP selects the HAS_ZERO_POINTS template parameter
694 +#define WVSPLITK_INT4G_LAUNCH(_THRDS, _YTILE, _UNRL, _N, _GS, _HAS_ZP) \
695 + { \
696 + dim3 block(_THRDS, 16); \
697 + int __wvPrGrp = mindiv_int4(M_in, CuCount * _YTILE, 16); \
698 + if (K_in * N_in <= max_lds_len && M_in % _YTILE == 0) \
699 + wvSplitK_int4_hf_sml_<fptype, _THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, \
700 + _HAS_ZP><<<grid, block, 0, stream>>>( \
701 + K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \
702 + __wvPrGrp, CuCount); \
703 + else \
704 + wvSplitK_int4_hf_<fptype, _THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, \
705 + _HAS_ZP><<<grid, block, 0, stream>>>( \
706 + K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \
707 + __wvPrGrp, CuCount); \
708 + }
709 +
710 +#define WVSPLITK_INT4G(_YTILE, _UNRL, _N, _GS, _HAS_ZP) \
711 + WVSPLITK_INT4G_LAUNCH(32, _YTILE, _UNRL, _N, _GS, _HAS_ZP)
712 +
713 +#define WVSPLIT_INT4G_GS(_YTILE, _UNRL, _N, _HAS_ZP) \
714 + if (group_size == 32) \
715 + WVSPLITK_INT4G(_YTILE, _UNRL, _N, 32, _HAS_ZP) \
716 + else if (group_size == 64) \
717 + WVSPLITK_INT4G(_YTILE, _UNRL, _N, 64, _HAS_ZP) \
718 + else \
719 + WVSPLITK_INT4G(_YTILE, _UNRL, _N, 128, _HAS_ZP)
720 +
721 +#define WVSPLIT_INT4G_TILE(_sYT, __N, _HAS_ZP) \
722 + { \
723 + if (K_in * N_in > max_lds_len) { \
724 + if (_sYT < 30) \
725 + WVSPLIT_INT4G_GS(4, 2, __N, _HAS_ZP) \
726 + else \
727 + WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \
728 + } else if (__N >= 4 && _sYT >= 480) \
729 + WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \
730 + else if (__N >= 3 && _sYT >= 40) \
731 + WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \
732 + else if (__N >= 3 && _sYT < 40 && (K_in <= 2048 || K_in >= 4096)) \
733 + WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \
734 + else if (__N >= 3 && _sYT < 40) \
735 + WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \
736 + else if (__N >= 2) \
737 + WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \
738 + else if (_sYT >= 30) \
739 + WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \
740 + else \
741 + WVSPLIT_INT4G_GS(1, 4, __N, _HAS_ZP) \
742 + }
743 +
744 +// Inner dispatch: shared by both symmetric and asymmetric paths
745 +#define WVSPLIT_INT4G_DISPATCH(_HAS_ZP) \
746 + { \
747 + int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); \
748 + switch (N_in) { \
749 + case 1: \
750 + WVSPLIT_INT4G_TILE(sYT, 1, _HAS_ZP) break; \
751 + case 2: \
752 + WVSPLIT_INT4G_TILE(sYT, 2, _HAS_ZP) break; \
753 + case 3: \
754 + WVSPLIT_INT4G_TILE(sYT, 3, _HAS_ZP) break; \
755 + case 4: \
756 + WVSPLIT_INT4G_TILE(sYT, 4, _HAS_ZP) break; \
757 + case 5: \
758 + WVSPLIT_INT4G_TILE(sYT, 5, _HAS_ZP) break; \
759 + default: \
760 + throw std::runtime_error("Unsupported N value: " + \
761 + std::to_string(N_in)); \
762 + } \
763 + }
764 +
765 + AT_DISPATCH_REDUCED_FLOATING_TYPES(
766 + in_b.scalar_type(), "wvSplitK_int4_g", [&] {
767 + using fptype = typename scalar<scalar_t>::type;
768 + const uint8_t* wptr = reinterpret_cast<const uint8_t*>(in_a.data_ptr());
769 + const fptype* aptr = reinterpret_cast<const fptype*>(in_b.data_ptr());
770 + const fptype* sptr =
771 + reinterpret_cast<const fptype*>(in_scale.data_ptr());
772 + const fptype* zpptr =
773 + in_zero_points.has_value()
774 + ? reinterpret_cast<const fptype*>(in_zero_points->data_ptr())
775 + : nullptr;
776 + const fptype* biasptr =
777 + (in_bias.has_value() && in_bias->numel() > 0)
778 + ? reinterpret_cast<const fptype*>(in_bias->data_ptr())
779 + : nullptr;
780 + fptype* cptr = reinterpret_cast<fptype*>(out_c.data_ptr());
781 +
782 + if (in_zero_points.has_value())
783 + WVSPLIT_INT4G_DISPATCH(true)
784 + else
785 + WVSPLIT_INT4G_DISPATCH(false)
786 + });
787 +
788 +#undef WVSPLITK_INT4G_LAUNCH
789 +#undef WVSPLITK_INT4G
790 +#undef WVSPLIT_INT4G_GS
791 +#undef WVSPLIT_INT4G_TILE
792 +#undef WVSPLIT_INT4G_DISPATCH
793 +
794 + return out_c;
795 +}
csrc/rocm/torch_bindings.cpp
+8
@@ -26,6 +26,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
26 "Tensor");
27 rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
28
29 + // W4A16 grouped skinny GEMM: packed int4 weights, per-group scales,
30 + // optional zero points for asymmetric quantization
31 + rocm_ops.def(
32 + "wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, "
33 + "Tensor? in_zero_points, Tensor? in_bias, int CuCount, "
34 + "int group_size) -> Tensor");
35 + rocm_ops.impl("wvSplitK_int4_g", torch::kCUDA, &wvSplitK_int4_g);
36 +
37 // Custom gemm op for skinny matrix-matrix multiplication
38 rocm_ops.def(
39 "wvSplitKrc(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> "
tests/kernels/quantization/test_rdna_hybrid_w4a16.py new
+525
@@ -0,0 +1,525 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""Tests for the ROCm Hybrid W4A16 kernel (HIP skinny + Triton prefill).
5 +
6 +Run `pytest tests/kernels/quantization/test_rdna_hybrid_w4a16.py`.
7 +"""
8 +
9 +import importlib
10 +
11 +import pytest
12 +import torch
13 +
14 +from vllm.platforms import current_platform
15 +from vllm.utils.torch_utils import set_random_seed
16 +
17 +if not current_platform.is_rocm():
18 + pytest.skip("ROCm only", allow_module_level=True)
19 +
20 +pytest.importorskip("triton")
21 +
22 +from vllm.platforms.rocm import on_gfx1x # noqa: E402
23 +
24 +device = "cuda"
25 +
26 +hybrid_module = importlib.import_module(
27 + "vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16"
28 +)
29 +RDNAHybridW4A16LinearKernel = hybrid_module.RDNAHybridW4A16LinearKernel
30 +pack_int4_exllama_shuffle = hybrid_module.pack_int4_exllama_shuffle
31 +SUPPORTED_GROUP_SIZES = hybrid_module.SUPPORTED_GROUP_SIZES
32 +MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE
33 +
34 +
35 +# ---------------------------------------------------------------------------
36 +# Reference implementation
37 +# ---------------------------------------------------------------------------
38 +
39 +
40 +def _rdna_hybrid_w4a16_reference(
41 + x_mk: torch.Tensor,
42 + w_int4_nk: torch.Tensor,
43 + scales_nkg: torch.Tensor,
44 + zp_nkg: torch.Tensor | None,
45 + group_size: int,
46 + bias: torch.Tensor | None,
47 +) -> torch.Tensor:
48 + """Reference for the Hybrid W4A16 op.
49 +
50 + x_mk: [M, K] fp16/bf16
51 + w_int4_nk: [N, K] int32 with raw uint4 values in [0, 15]
52 + scales_nkg: [N, K//G] fp16/bf16
53 + zp_nkg: [N, K//G] fp16/bf16 raw zero points (already in act dtype),
54 + or None for symmetric (uint4b8, dequant subtracts 8)
55 + """
56 + G = group_size
57 + N, K = w_int4_nk.shape
58 + assert K % G == 0
59 + s_full = scales_nkg.repeat_interleave(G, dim=1).to(torch.float32) # [N, K]
60 + if zp_nkg is None:
61 + z_full = torch.full((N, K), 8.0, device=x_mk.device, dtype=torch.float32)
62 + else:
63 + z_full = zp_nkg.repeat_interleave(G, dim=1).to(torch.float32)
64 + w_fp = (w_int4_nk.to(torch.float32) - z_full) * s_full # [N, K]
65 + out = x_mk.to(torch.float32) @ w_fp.t() # [M, N]
66 + if bias is not None:
67 + out = out + bias.to(torch.float32)
68 + return out.to(x_mk.dtype)
69 +
70 +
71 +# ---------------------------------------------------------------------------
72 +# Forward correctness: decode (HIP skinny) + prefill (Triton)
73 +# ---------------------------------------------------------------------------
74 +
75 +
76 +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
77 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
78 +@pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)
79 +@pytest.mark.parametrize("has_zp", [False, True])
80 +@pytest.mark.parametrize(
81 + "M",
82 + [1, MAX_SKINNY_BATCH_SIZE, MAX_SKINNY_BATCH_SIZE + 1, 64],
83 + ids=["M=1_decode", "M=5_decode", "M=6_prefill", "M=64_prefill"],
84 +)
85 +def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M):
86 + """Smoke test the registered custom op for both decode and prefill batches.
87 +
88 + Verifies the dispatch logic in `_rdna_hybrid_w4a16_apply_impl`:
89 + - M <= MAX_SKINNY_BATCH_SIZE: HIP wvSplitK_int4_g
90 + - M > MAX_SKINNY_BATCH_SIZE: Triton prefill kernel
91 + """
92 + if not torch.cuda.is_available():
93 + pytest.skip("CUDA/HIP device not available")
94 +
95 + set_random_seed(0)
96 +
97 + K, N = 1024, 256
98 + assert K % group_size == 0 and K % 8 == 0 and N % 8 == 0
99 +
100 + # Activations.
101 + x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
102 +
103 + # Weights as raw uint4 in [N, K], packed to ExLlama shuffle [N, K//8].
104 + w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
105 + w_q_i32 = pack_int4_exllama_shuffle(w_int4_nk).contiguous() # [N, K//8] int32
106 + w_q = w_q_i32.view(torch.int8) # same bytes viewed as int8 [N, K//2]
107 +
108 + # Scales [N, K//G] in act dtype.
109 + scales_nkg = (
110 + 0.05 * torch.rand((N, K // group_size), device=device, dtype=torch.float32)
111 + ).to(dtype)
112 +
113 + # Optional raw zero points [N, K//G] in act dtype.
114 + if has_zp:
115 + zp_nkg = torch.randint(
116 + 0, 16, (N, K // group_size), device=device, dtype=torch.int32
117 + ).to(dtype)
118 + else:
119 + zp_nkg = None
120 +
121 + from vllm.utils.platform_utils import num_compute_units
122 +
123 + out = torch.ops.vllm.rdna_hybrid_w4a16_apply(
124 + x_mk,
125 + w_q,
126 + scales_nkg,
127 + zp_nkg,
128 + None, # bias
129 + num_compute_units(),
130 + group_size,
131 + )
132 +
133 + ref = _rdna_hybrid_w4a16_reference(
134 + x_mk, w_int4_nk, scales_nkg, zp_nkg, group_size, bias=None
135 + )
136 +
137 + torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
138 +
139 +
140 +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
141 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
142 +@pytest.mark.parametrize("M", [1, MAX_SKINNY_BATCH_SIZE + 1])
143 +def test_rdna_hybrid_w4a16_apply_with_bias(dtype, M):
144 + """Bias is added correctly on both decode and prefill paths."""
145 + if not torch.cuda.is_available():
146 + pytest.skip("CUDA/HIP device not available")
147 +
148 + set_random_seed(0)
149 + K, N, G = 1024, 128, 128
150 +
151 + x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
152 + w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
153 + w_q_i32 = pack_int4_exllama_shuffle(w_int4_nk).contiguous()
154 + w_q = w_q_i32.view(torch.int8)
155 + scales_nkg = (
156 + 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)
157 + ).to(dtype)
158 + bias = torch.randn(N, device=device, dtype=dtype) * 0.1
159 +
160 + from vllm.utils.platform_utils import num_compute_units
161 +
162 + out = torch.ops.vllm.rdna_hybrid_w4a16_apply(
163 + x_mk,
164 + w_q,
165 + scales_nkg,
166 + None,
167 + bias,
168 + num_compute_units(),
169 + G,
170 + )
171 + ref = _rdna_hybrid_w4a16_reference(x_mk, w_int4_nk, scales_nkg, None, G, bias=bias)
172 +
173 + torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
174 +
175 +
176 +# ---------------------------------------------------------------------------
177 +# pack_int4_exllama_shuffle round-trips correctly
178 +# ---------------------------------------------------------------------------
179 +
180 +
181 +def test_pack_int4_exllama_shuffle_layout():
182 + """Pack 8 K-values per int32 in interleave [0,2,4,6,1,3,5,7] order."""
183 + if not torch.cuda.is_available():
184 + pytest.skip("CUDA/HIP device not available")
185 + set_random_seed(0)
186 + N, K = 4, 16
187 + w = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
188 + packed = pack_int4_exllama_shuffle(w)
189 + assert packed.shape == (N, K // 8) and packed.dtype == torch.int32
190 +
191 + # Manual unshuffle using ExLlama shifts [0,16,4,20,8,24,12,28].
192 + shifts = torch.tensor(
193 + [0, 16, 4, 20, 8, 24, 12, 28], device=device, dtype=torch.int32
194 + )
195 + unshuffled = (packed.unsqueeze(-1) >> shifts) & 0xF
196 + unshuffled = unshuffled.reshape(N, K)
197 + torch.testing.assert_close(unshuffled, w)
198 +
199 +
200 +# ---------------------------------------------------------------------------
201 +# process_weights_after_loading: layout repack and zp normalization
202 +# ---------------------------------------------------------------------------
203 +
204 +
205 +def _pack_int4_along_k_to_ckpt(w_int4_kn: torch.Tensor) -> torch.Tensor:
206 + """Pack int4 values along K into CT checkpoint layout: [K,N] -> [N, K//8]."""
207 + assert w_int4_kn.dtype == torch.int32
208 + K, N = w_int4_kn.shape
209 + assert K % 8 == 0
210 + out = torch.zeros((N, K // 8), dtype=torch.int32, device=w_int4_kn.device)
211 + for i in range(8):
212 + out |= (w_int4_kn[i::8, :].t() & 0xF) << (i * 4)
213 + return out.contiguous()
214 +
215 +
216 +def _pack_int4_along_n_for_zp(zp_int4_gn: torch.Tensor) -> torch.Tensor:
217 + """Pack int4 zero points along N: [G, N] -> [G, N//8] int32 (CT layout)."""
218 + assert zp_int4_gn.dtype == torch.int32
219 + G, N = zp_int4_gn.shape
220 + assert N % 8 == 0
221 + shifts = torch.arange(8, device=zp_int4_gn.device, dtype=torch.int32) * 4
222 + return torch.sum(
223 + (zp_int4_gn.view(G, N // 8, 8) & 0xF) << shifts, dim=2, dtype=torch.int32
224 + ).contiguous()
225 +
226 +
227 +def _build_dummy_layer(
228 + w_ckpt_nk8: torch.Tensor,
229 + scales_ckpt_nkg: torch.Tensor,
230 + zeros_ckpt: torch.Tensor | None,
231 +):
232 + from vllm.model_executor.parameter import (
233 + GroupQuantScaleParameter,
234 + PackedColumnParameter,
235 + PackedvLLMParameter,
236 + )
237 +
238 + weight_loader = lambda *args, **kwargs: None
239 +
240 + class DummyLayer(torch.nn.Module):
241 + pass
242 +
243 + layer = DummyLayer()
244 + layer.register_parameter(
245 + "weight_packed",
246 + PackedvLLMParameter(
247 + data=w_ckpt_nk8,
248 + weight_loader=weight_loader,
249 + input_dim=1,
250 + output_dim=0,
251 + packed_factor=8,
252 + packed_dim=1,
253 + ),
254 + )
255 + layer.register_parameter(
256 + "weight_scale",
257 + GroupQuantScaleParameter(
258 + data=scales_ckpt_nkg,
259 + weight_loader=weight_loader,
260 + input_dim=1,
261 + output_dim=0,
262 + ),
263 + )
264 + if zeros_ckpt is not None:
265 + layer.register_parameter(
266 + "weight_zero_point",
267 + PackedColumnParameter(
268 + data=zeros_ckpt,
269 + weight_loader=weight_loader,
270 + output_dim=0,
271 + packed_factor=8,
272 + packed_dim=0,
273 + ),
274 + )
275 + return layer
276 +
277 +
278 +@pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)
279 +def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_init):
280 + """uint4b8 (symmetric): w_q -> [N, K//8] int8 ExLlama shuffle, no zp param."""
281 + if not torch.cuda.is_available():
282 + pytest.skip("CUDA/HIP device not available")
283 +
284 + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
285 + MPLinearLayerConfig,
286 + )
287 + from vllm.scalar_type import scalar_types
288 +
289 + set_random_seed(0)
290 +
291 + K, N = 256, 128
292 + G = group_size
293 + assert K % G == 0
294 +
295 + # Reference unpacked weights, then pack into CT checkpoint layout [N, K//8].
296 + w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
297 + w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn)
298 + scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16)
299 +
300 + layer = _build_dummy_layer(w_ckpt_nk8, scales_ckpt_nkg, zeros_ckpt=None)
301 +
302 + config = MPLinearLayerConfig(
303 + full_weight_shape=(K, N),
304 + partition_weight_shape=(K, N),
305 + weight_type=scalar_types.uint4b8,
306 + act_type=torch.float16,
307 + group_size=G,
308 + zero_points=False,
309 + has_g_idx=False,
310 + )
311 + kernel = RDNAHybridW4A16LinearKernel(
312 + config,
313 + w_q_param_name="weight_packed",
314 + w_s_param_name="weight_scale",
315 + w_zp_param_name=None,
316 + w_gidx_param_name=None,
317 + )
318 + kernel.process_weights_after_loading(layer)
319 +
320 + # Skinny weight is stored once as int8 [N, K//2]; the Triton path
321 + # reinterprets it as int32 [N, K//8] via a view (no separate parameter).
322 + assert layer.weight_packed.dtype == torch.int8
323 + assert tuple(layer.weight_packed.shape) == (N, K // 2)
324 + w_q_i32 = layer.weight_packed.view(torch.int32)
325 + assert tuple(w_q_i32.shape) == (N, K // 8)
326 +
327 + expected_packed = pack_int4_exllama_shuffle(w_int4_kn.t().contiguous())
328 + torch.testing.assert_close(w_q_i32, expected_packed)
329 +
330 + # Scales: [N, K//G] (skinny layout, no transpose since CT already had it).
331 + assert tuple(layer.weight_scale.shape) == (N, K // G)
332 + torch.testing.assert_close(layer.weight_scale, scales_ckpt_nkg)
333 +
334 +
335 +@pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)
336 +def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_init):
337 + """uint4 (asymmetric): zero points unpacked to raw values in act dtype."""
338 + if not torch.cuda.is_available():
339 + pytest.skip("CUDA/HIP device not available")
340 +
341 + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
342 + MPLinearLayerConfig,
343 + )
344 + from vllm.scalar_type import scalar_types
345 +
346 + set_random_seed(0)
347 +
348 + K, N = 256, 128
349 + G = group_size
350 + assert K % G == 0 and N % 8 == 0
351 +
352 + w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
353 + w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn)
354 + scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16)
355 +
356 + # CT zero-point layout is N-packed: [N//8, K//G] int32.
357 + zeros_int4_gn = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32)
358 + zeros_packed_gn8 = _pack_int4_along_n_for_zp(zeros_int4_gn) # [K//G, N//8]
359 + zeros_ckpt_n8kg = zeros_packed_gn8.t().contiguous() # [N//8, K//G]
360 +
361 + layer = _build_dummy_layer(w_ckpt_nk8, scales_ckpt_nkg, zeros_ckpt=zeros_ckpt_n8kg)
362 +
363 + config = MPLinearLayerConfig(
364 + full_weight_shape=(K, N),
365 + partition_weight_shape=(K, N),
366 + weight_type=scalar_types.uint4,
367 + act_type=torch.float16,
368 + group_size=G,
369 + zero_points=True,
370 + has_g_idx=False,
371 + )
372 + kernel = RDNAHybridW4A16LinearKernel(
373 + config,
374 + w_q_param_name="weight_packed",
375 + w_s_param_name="weight_scale",
376 + w_zp_param_name="weight_zero_point",
377 + w_gidx_param_name=None,
378 + )
379 + kernel.process_weights_after_loading(layer)
380 +
381 + # Zero-points: unpacked to [N, K//G], cast to act dtype, raw values [0..15].
382 + assert layer.weight_zero_point.dtype == torch.float16
383 + assert tuple(layer.weight_zero_point.shape) == (N, K // G)
384 + expected_zp = zeros_int4_gn.t().to(torch.float16) # [N, K//G]
385 + torch.testing.assert_close(layer.weight_zero_point, expected_zp)
386 +
387 + # Quantized weights match symmetric path's layout regardless of zp.
388 + w_q_i32 = layer.weight_packed.view(torch.int32)
389 + assert tuple(w_q_i32.shape) == (N, K // 8)
390 + expected_packed = pack_int4_exllama_shuffle(w_int4_kn.t().contiguous())
391 + torch.testing.assert_close(w_q_i32, expected_packed)
392 +
393 +
394 +# ---------------------------------------------------------------------------
395 +# can_implement enforces the supported-group-size policy
396 +# ---------------------------------------------------------------------------
397 +
398 +
399 +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
400 +@pytest.mark.parametrize(
401 + "group_size,expected_ok", [(32, True), (64, True), (128, True), (256, False)]
402 +)
403 +def test_hybrid_can_implement_group_size(group_size, expected_ok):
404 + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
405 + MPLinearLayerConfig,
406 + )
407 + from vllm.scalar_type import scalar_types
408 +
409 + K, N = 1024, 256
410 + config = MPLinearLayerConfig(
411 + full_weight_shape=(K, N),
412 + partition_weight_shape=(K, N),
413 + weight_type=scalar_types.uint4b8,
414 + act_type=torch.float16,
415 + group_size=group_size,
416 + zero_points=False,
417 + has_g_idx=False,
418 + )
419 + ok, _ = RDNAHybridW4A16LinearKernel.can_implement(config)
420 + assert ok is expected_ok
421 +
422 +
423 +# ---------------------------------------------------------------------------
424 +# Tests for the HIP wvSplitK_int4_g decode kernel
425 +# ---------------------------------------------------------------------------
426 +
427 +
428 +def _hip_skinny_reference(
429 + a_mk: torch.Tensor,
430 + w_int4_nk: torch.Tensor,
431 + scales_nkg: torch.Tensor,
432 + *,
433 + group_size: int,
434 + zp_bias: int,
435 +) -> torch.Tensor:
436 + """Reference for symmetric HIP skinny: C = A @ (W - zp_bias) * S."""
437 + K = a_mk.shape[1]
438 + N = w_int4_nk.shape[0]
439 + num_groups = K // group_size
440 +
441 + w_fp = (w_int4_nk.to(torch.float32) - zp_bias).view(N, num_groups, group_size)
442 + s = scales_nkg.to(torch.float32).unsqueeze(-1)
443 + w_dequant = (w_fp * s).view(N, K)
444 +
445 + return (a_mk.to(torch.float32) @ w_dequant.t()).to(a_mk.dtype)
446 +
447 +
448 +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
449 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
450 +@pytest.mark.parametrize(
451 + "M,K,N,G",
452 + [
453 + (1, 256, 256, 32),
454 + (1, 256, 256, 64),
455 + (1, 512, 256, 128),
456 + (2, 512, 256, 64),
457 + (3, 256, 512, 64),
458 + ],
459 +)
460 +def test_hip_skinny_wvSplitK_int4_g(dtype, M, K, N, G):
461 + """Test HIP wvSplitK_int4_g kernel directly via _custom_ops."""
462 + import vllm._custom_ops as ops
463 + from vllm.utils.platform_utils import num_compute_units
464 +
465 + set_random_seed(0)
466 +
467 + a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
468 + w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
469 +
470 + b_packed_i32 = pack_int4_exllama_shuffle(w_int4_nk)
471 + b_packed_i8 = b_packed_i32.view(torch.int8)
472 +
473 + scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to(
474 + dtype
475 + )
476 +
477 + cu_count = num_compute_units()
478 + out = ops.wvSplitK_int4_g(b_packed_i8, a, scales, cu_count, G)
479 +
480 + ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8)
481 +
482 + torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2)
483 +
484 +
485 +# ---------------------------------------------------------------------------
486 +# Tests for the full hybrid dispatch (HIP decode + Triton prefill)
487 +# ---------------------------------------------------------------------------
488 +
489 +
490 +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
491 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
492 +@pytest.mark.parametrize(
493 + "M,K,N,G",
494 + [
495 + (1, 256, 256, 64),
496 + (1, 512, 256, 32),
497 + (1, 512, 256, 128),
498 + (32, 512, 256, 64),
499 + (64, 1024, 256, 128),
500 + ],
501 +)
502 +def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G):
503 + """Test the full hybrid dispatch via the custom op."""
504 + from vllm.utils.platform_utils import num_compute_units
505 +
506 + set_random_seed(0)
507 +
508 + a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
509 + w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
510 +
511 + b_packed_i32 = pack_int4_exllama_shuffle(w_int4_nk)
512 + b_packed_i8 = b_packed_i32.view(torch.int8)
513 +
514 + scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to(
515 + dtype
516 + )
517 +
518 + cu_count = num_compute_units()
519 + out = torch.ops.vllm.rdna_hybrid_w4a16_apply(
520 + a, b_packed_i8, scales, None, None, cu_count, G
521 + )
522 +
523 + ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8)
524 +
525 + torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2)
tests/kernels/quantization/test_w4a16_kernel_selection.py
+24 -5
@@ -16,9 +16,28 @@ from vllm.model_executor.kernels.linear import (
16 from vllm.platforms import current_platform
17 from vllm.scalar_type import scalar_types
18
19 +if current_platform.is_rocm():
20 + from vllm.platforms.rocm import on_gfx1x, on_gfx1100
21 +else:
22 + on_gfx1100 = on_gfx1x = lambda: False # noqa: E731
23 +
24 +# Group sizes the HIP skinny path of RDNAHybridW4A16LinearKernel instantiates.
25 +_HYBRID_GROUP_SIZES = (32, 64, 128)
26 +
27 +
28 +def _expected_rocm_kernel(weight_type, group_size: int) -> str:
29 + """Mirror the ROCm priority order in ``_POSSIBLE_KERNELS[ROCM]``:
30 + RDNA3 (gfx1100, symmetric uint4b8) -> Hybrid (gfx11/gfx12) -> Triton.
31 + """
32 + if on_gfx1100() and weight_type == scalar_types.uint4b8:
33 + return "RDNA3W4A16LinearKernel"
34 + if on_gfx1x() and group_size in _HYBRID_GROUP_SIZES:
35 + return "RDNAHybridW4A16LinearKernel"
36 + return "TritonW4A16LinearKernel"
37 +
38
39 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
21 -def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4b8():
40 +def test_choose_mp_linear_kernel_uint4b8():
41 # int4 weights, 16-bit activations (CT W4A16 typical config).
42 K, N = 1024, 256
43 config = MPLinearLayerConfig(
@@ -32,12 +51,12 @@ def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4b8():
51 )
52
53 kernel_type = choose_mp_linear_kernel(config)
35 - assert kernel_type.__name__ == "TritonW4A16LinearKernel"
54 + assert kernel_type.__name__ == _expected_rocm_kernel(scalar_types.uint4b8, 128)
55
56
57 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
39 -def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4_asymmetric():
40 - # Asymmetric int4 weights should also be supported (explicit zero points).
58 +def test_choose_mp_linear_kernel_uint4_asymmetric():
59 + # Asymmetric int4 weights (explicit zero points).
60 K, N = 512, 512
61 config = MPLinearLayerConfig(
62 full_weight_shape=(K, N),
@@ -50,4 +69,4 @@ def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4_asymmetric():
69 )
70
71 kernel_type = choose_mp_linear_kernel(config)
53 - assert kernel_type.__name__ == "TritonW4A16LinearKernel"
72 + assert kernel_type.__name__ == _expected_rocm_kernel(scalar_types.uint4, 64)
vllm/_custom_ops.py
+36
@@ -2077,6 +2077,42 @@ def wvSplitK(
2077 return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count)
2078
2079
2080 +def wvSplitK_int4_g(
2081 + weight: torch.Tensor,
2082 + activation: torch.Tensor,
2083 + scale: torch.Tensor,
2084 + cu_count: int,
2085 + group_size: int,
2086 + zero_points: torch.Tensor | None = None,
2087 + bias: torch.Tensor | None = None,
2088 +) -> torch.Tensor:
2089 + # NOTE: the kernel is weight-major; `weight` is the packed int4 operand
2090 + # (in_a, [out_features, K/2]) and `activation` is in_b ([num_tokens, K]).
2091 + return torch.ops._rocm_C.wvSplitK_int4_g(
2092 + weight, activation, scale, zero_points, bias, cu_count, group_size
2093 + )
2094 +
2095 +
2096 +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"):
2097 +
2098 + @register_fake("_rocm_C::wvSplitK_int4_g")
2099 + def _wvSplitK_int4_g_fake(
2100 + in_a: torch.Tensor, # packed int4 weight [out_features, K/2]
2101 + in_b: torch.Tensor, # activation [num_tokens, K]
2102 + in_scale: torch.Tensor,
2103 + in_zero_points: torch.Tensor | None,
2104 + in_bias: torch.Tensor | None,
2105 + CuCount: int,
2106 + group_size: int,
2107 + ) -> torch.Tensor:
2108 + # Kernel returns {in_b.size(0), in_a.size(0)} = [num_tokens, out_features].
2109 + num_tokens = in_b.size(0)
2110 + out_features = in_a.size(0)
2111 + return torch.empty(
2112 + (num_tokens, out_features), dtype=in_b.dtype, device=in_b.device
2113 + )
2114 +
2115 +
2116 def wvSplitKrc(
2117 a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
2118 ) -> torch.Tensor:
vllm/model_executor/kernels/linear/__init__.py
+5
@@ -57,6 +57,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
57 from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (
58 RDNA3W4A16LinearKernel,
59 )
60 +from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
61 + RDNAHybridW4A16LinearKernel,
62 +)
63 from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
64 TritonW4A16LinearKernel,
65 )
@@ -403,6 +406,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
406 ],
407 PlatformEnum.ROCM: [
408 RDNA3W4A16LinearKernel,
409 + RDNAHybridW4A16LinearKernel,
410 TritonW4A16LinearKernel,
411 ConchLinearKernel,
412 ExllamaLinearKernel,
@@ -1063,6 +1067,7 @@ __all__ = [
1067 "CutlassW4A8LinearKernel",
1068 "Dynamic4bitLinearKernel",
1069 "ExllamaLinearKernel",
1070 + "RDNAHybridW4A16LinearKernel",
1071 "MacheteLinearKernel",
1072 "MarlinLinearKernel",
1073 "TritonW4A16LinearKernel",
vllm/model_executor/kernels/linear/mixed_precision/__init__.py
+4
@@ -35,6 +35,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
35 from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (
36 RDNA3W4A16LinearKernel,
37 )
38 +from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import (
39 + RDNAHybridW4A16LinearKernel,
40 +)
41 from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
42 TritonW4A16LinearKernel,
43 )
@@ -59,6 +62,7 @@ __all__ = [
62 "MacheteLinearKernel",
63 "MarlinLinearKernel",
64 "RDNA3W4A16LinearKernel",
65 + "RDNAHybridW4A16LinearKernel",
66 "TritonW4A16LinearKernel",
67 "XPUW4A8IntLinearKernel",
68 "XPUwNa16LinearKernel",
vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py new
+560
@@ -0,0 +1,560 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode.
5 +
6 +Routes based on batch size M:
7 + M <= MAX_SKINNY_BATCH_SIZE: HIP skinny GEMM (wvSplitK_int4_g)
8 + M > MAX_SKINNY_BATCH_SIZE: Triton W4A16 fused dequant GEMM
9 +
10 +Stores the weights ONCE as int8 [N, K//2] (ExLlama shuffle packed). Both
11 +paths read this single buffer: the HIP skinny kernel uses it directly, and
12 +the triton kernel reinterprets it as int32 [N, K//8] via a view (and
13 +transposes tiles in-register). No dual weight storage.
14 +"""
15 +
16 +from contextlib import nullcontext
17 +
18 +import torch
19 +
20 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
21 + unpack_quantized_values_into_int32,
22 +)
23 +from vllm.model_executor.parameter import (
24 + permute_param_layout_,
25 +)
26 +from vllm.platforms import current_platform
27 +from vllm.scalar_type import scalar_types
28 +from vllm.triton_utils import tl, triton
29 +from vllm.utils.torch_utils import direct_register_custom_op
30 +
31 +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
32 +
33 +SUPPORTED_GROUP_SIZES = [32, 64, 128]
34 +
35 +
36 +def _on_gfx12x() -> bool:
37 + if not current_platform.is_rocm():
38 + return False
39 + from vllm.platforms.rocm import on_gfx12x
40 +
41 + return on_gfx12x()
42 +
43 +
44 +def _on_gfx1x() -> bool:
45 + if not current_platform.is_rocm():
46 + return False
47 + from vllm.platforms.rocm import on_gfx1x
48 +
49 + return on_gfx1x()
50 +
51 +
52 +def _on_gfx1151() -> bool:
53 + if not current_platform.is_rocm():
54 + return False
55 + from vllm.platforms.rocm import on_gfx1151
56 +
57 + return on_gfx1151()
58 +
59 +
60 +# Maximum batch size M for the HIP skinny kernel path (C++ supports N_in
61 +# up to 5). When M is below this AND K*M fits in LDS, the skinny kernel is
62 +# used; otherwise the Triton prefill path handles the GEMM.
63 +MAX_SKINNY_BATCH_SIZE = 5
64 +# 64 KiB per-workgroup LDS limit expressed in fp16 elements.
65 +# (AMD RDNA has 128 KiB total LDS per CU, but 64 KiB per workgroup.)
66 +LDS_CAPACITY_ELEMENTS = 64 * 1024 // 2 # 32768 fp16 elements
67 +
68 +
69 +# ---------------------------------------------------------------------------
70 +# Triton kernel for the prefill path (reads skinny-format weights [N, K//8])
71 +# ---------------------------------------------------------------------------
72 +
73 +
74 +@triton.jit
75 +def _triton_w4a16_skinny_fmt_kernel(
76 + # Pointers
77 + a_ptr, # [M, K] fp16/bf16 activations
78 + b_ptr, # [N, K//8] int32 packed (ExLlama shuffle, K is packed dim)
79 + scales_ptr, # [N, K//G] fp16/bf16 scales (skinny layout)
80 + zp_ptr, # [N, K//G] fp16/bf16 raw zero-points (when HAS_ZP=True)
81 + c_ptr, # [M, N] fp16/bf16 output
82 + # Dimensions
83 + M,
84 + N,
85 + K,
86 + K8, # K // 8
87 + num_groups, # K // group_size
88 + # Quantization parameters
89 + group_size,
90 + ZP_BIAS: tl.constexpr,
91 + HAS_ZP: tl.constexpr,
92 + # Block sizes
93 + BLOCK_M: tl.constexpr,
94 + BLOCK_N: tl.constexpr,
95 + BLOCK_K: tl.constexpr,
96 +):
97 + """
98 + Fused W4A16 GEMM reading weights from skinny format [N, K//8].
99 +
100 + B is stored as [N, K//8] int32 using ExLlama shuffle packing:
101 + each int32 packs 8 K-values with interleave [0,2,4,6,1,3,5,7]:
102 + packed = val[0] | (val[2]<<4) | (val[4]<<8) | (val[6]<<12)
103 + | (val[1]<<16) | (val[3]<<20) | (val[5]<<24) | (val[7]<<28)
104 +
105 + Scales are [N, K//G] (skinny layout, NOT transposed).
106 + When HAS_ZP=True, raw zero-points zp_raw are loaded from zp_ptr [N, K//G]
107 + and subtracted directly: (nibble - zp_raw) * scale.
108 + When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric).
109 + """
110 + pid_m = tl.program_id(0)
111 + pid_n = tl.program_id(1)
112 +
113 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
114 + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
115 +
116 + # ExLlama unshuffle shifts: shift[j] = (j//2)*4 + (j%2)*16
117 + # For 8 values: [0, 16, 4, 20, 8, 24, 12, 28]
118 + exllama_shifts_row = (tl.arange(0, 8) // 2) * 4 + (tl.arange(0, 8) % 2) * 16
119 + shifts_1d = tl.reshape(
120 + tl.broadcast_to(exllama_shifts_row[None, :], (BLOCK_K // 8, 8)),
121 + (BLOCK_K,),
122 + )
123 + shifts_full = tl.broadcast_to(shifts_1d[None, :], (BLOCK_N, BLOCK_K))
124 +
125 + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
126 +
127 + for k_start in range(0, tl.cdiv(K, BLOCK_K)):
128 + offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K)
129 + mask_k = offs_k < K
130 +
131 + a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
132 + mask_a = (offs_m[:, None] < M) & mask_k[None, :]
133 + a = tl.load(a_ptrs, mask=mask_a, other=0.0)
134 +
135 + offs_k8 = k_start * (BLOCK_K // 8) + tl.arange(0, BLOCK_K // 8)
136 + b_ptrs = b_ptr + offs_n[:, None] * K8 + offs_k8[None, :]
137 + mask_b = (offs_n[:, None] < N) & (offs_k8[None, :] < K8)
138 + b_packed = tl.load(b_ptrs, mask=mask_b, other=0)
139 +
140 + b = tl.interleave(b_packed, b_packed)
141 + b = tl.interleave(b, b)
142 + b = tl.interleave(b, b)
143 + b = (b >> shifts_full) & 0xF
144 +
145 + g_idx = (k_start * BLOCK_K) // group_size
146 + scale_ptrs = scales_ptr + offs_n * num_groups + g_idx
147 + scale_mask = offs_n < N
148 + scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0)
149 +
150 + if HAS_ZP:
151 + zp_ptrs = zp_ptr + offs_n * num_groups + g_idx
152 + zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0)
153 + b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None]
154 + else:
155 + b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None]
156 +
157 + b_fp_t = tl.trans(b_fp)
158 + accumulator += tl.dot(a, b_fp_t, out_dtype=tl.float32)
159 +
160 + c = accumulator.to(c_ptr.type.element_ty)
161 + c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
162 + mask_c = (offs_m[:, None] < M) & (offs_n[None, :] < N)
163 + tl.store(c_ptrs, c, mask=mask_c)
164 +
165 +
166 +# Per-shape (group_size, K, N) -> (BLOCK_M, BLOCK_N, BLOCK_K, num_warps,
167 +# num_stages) tile-config overrides for prefill (M <= 128) on gfx1x.
168 +# Picked by sweeping benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py + a
169 +# per-config sweep script; only added when better than the generic heuristic
170 +# by > 20% at M=128. Re-run benchmarks after edits.
171 +_GFX1X_PREFILL_OVERRIDES: dict[tuple[int, int, int], tuple[int, int, int, int, int]] = {
172 + # SmolLM2-1.7B-Instruct-AWQ (gs=32, K=2048; gs forces BLOCK_K to 32 so
173 + # widen BLOCK_M and let Triton pipeline 4 stages to amortize the small
174 + # K-tile).
175 + (32, 2048, 6144): (128, 32, 32, 4, 4), # qkv_proj
176 + (32, 2048, 2048): (128, 32, 32, 4, 4), # o_proj
177 + (32, 2048, 16384): (128, 32, 32, 4, 4), # gate_up_proj
178 + (32, 8192, 2048): (128, 64, 32, 8, 2), # down_proj
179 + # Qwen3-8B-quantized.w4a16 (gs=128, K=4096 / 12288). For these K's a
180 + # full BLOCK_K=group_size with no software pipelining beats the generic
181 + # 64x64x64 — single-stage keeps register pressure down.
182 + (128, 4096, 6144): (128, 64, 128, 8, 1), # qkv_proj
183 + (128, 4096, 4096): (128, 32, 128, 4, 1), # o_proj
184 + (128, 4096, 24576): (64, 32, 128, 2, 1), # gate_up_proj
185 + (128, 12288, 4096): (128, 64, 128, 8, 1), # down_proj
186 +}
187 +
188 +
189 +def triton_w4a16_skinny_fmt_gemm(
190 + a: torch.Tensor, # [M, K] fp16/bf16
191 + b_q: torch.Tensor, # [N, K//8] int32 (ExLlama shuffle packed)
192 + scales: torch.Tensor, # [N, K//G] fp16/bf16
193 + group_size: int,
194 + zp_bias: int = 8,
195 + zp: torch.Tensor | None = None, # [N, K//G] per-group zero-points
196 +) -> torch.Tensor:
197 + """
198 + Fused W4A16 GEMM reading from skinny weight format [N, K//8].
199 +
200 + Args:
201 + a: Activation matrix [M, K], float16 or bfloat16.
202 + b_q: Packed weight matrix [N, K//8], int32 (ExLlama shuffle).
203 + scales: Per-group scales [N, K//G], same dtype as a.
204 + group_size: Quantization group size (resolved from -1 to K by caller).
205 + zp_bias: Constant zero bias (default 8 for unsigned int4).
206 + zp: Raw per-group zero-points [N, K//G] (asymmetric),
207 + stored as zp_raw in activation dtype. When provided,
208 + dequant is (nibble - zp_raw) * scale.
209 +
210 + Returns:
211 + Output matrix [M, N], same dtype as a.
212 + """
213 + assert a.is_contiguous(), "Activation matrix must be contiguous"
214 + assert b_q.is_contiguous(), "Weight matrix must be contiguous"
215 + assert scales.is_contiguous(), "Scales must be contiguous"
216 +
217 + M, K = a.shape
218 + N = b_q.shape[0]
219 + K8 = K // 8
220 + num_groups = K // group_size
221 +
222 + assert b_q.shape == (N, K8), f"b_q shape mismatch: {b_q.shape} vs ({N}, {K8})"
223 + assert scales.shape == (N, num_groups), (
224 + f"scales shape mismatch: {scales.shape} vs ({N}, {num_groups})"
225 + )
226 + if zp is not None:
227 + assert zp.is_contiguous(), "Zero-points must be contiguous"
228 + assert zp.shape == (N, num_groups), (
229 + f"zp shape mismatch: {zp.shape} vs ({N}, {num_groups})"
230 + )
231 + has_zp = zp is not None
232 +
233 + c = torch.empty((M, N), dtype=a.dtype, device=a.device)
234 +
235 + # num_stages stays None unless a per-shape override sets it, so the
236 + # generic heuristics fall back to Triton's default pipeline depth.
237 + num_stages: int | None = None
238 + if _on_gfx12x():
239 + # Tuned on gfx1201 (Radeon AI PRO R9700, 32 CUs, 32-wide wavefronts)
240 + # using Llama-3.1-8B AWQ weight shapes with group_size=128.
241 + if M <= 32:
242 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 16, 16, 128, 4
243 + elif M <= 64:
244 + if K >= 2 * N: # tall K (e.g. down_proj)
245 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 32, 128, 8
246 + elif N > K: # wide N (e.g. qkv_proj, gate_up_proj)
247 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 32, 64, 8
248 + else: # N ~= K (e.g. o_proj)
249 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 32, 64, 128, 4
250 + elif M <= 128:
251 + if K >= 2 * N: # tall K (e.g. down_proj)
252 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 16, 64, 1
253 + elif N >= 2 * K: # very wide N (e.g. gate_up_proj)
254 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 128, 64, 8
255 + else: # N ~= K (e.g. o_proj, qkv_proj)
256 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 64, 64, 8
257 + elif M <= 512:
258 + if K >= 2 * N: # tall K (e.g. down_proj)
259 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 64, 64, 8
260 + elif N >= 4 * K: # very wide N (e.g. gate_up_proj)
261 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 128, 64, 8
262 + else:
263 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 128, 64, 8
264 + else:
265 + if K >= 2 * N: # tall K (e.g. down_proj)
266 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 64, 64, 8
267 + elif N >= 4 * K: # very wide N (e.g. gate_up_proj)
268 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 256, 64, 64, 8
269 + else:
270 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 128, 32, 8
271 + elif _on_gfx1151():
272 + # Tuned on gfx1151 (Strix Halo, 40 CUs, 32-wide wavefronts)
273 + # using Qwen3-4B weight shapes with group_size=128.
274 + # Per-shape overrides for known prefill regressions live in a small
275 + # lookup table — see _GFX1X_PREFILL_OVERRIDES below. Re-run
276 + # benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py after edits.
277 + override = (
278 + _GFX1X_PREFILL_OVERRIDES.get((group_size, K, N)) if M <= 128 else None
279 + )
280 + if override is not None:
281 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = override
282 + elif M <= 32:
283 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 32, 32, 128, 4
284 + elif M <= 64:
285 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 64, 32, 4
286 + elif M <= 128:
287 + if K >= 2 * N: # tall K (e.g. down_proj)
288 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 16, 64, 1
289 + elif N > K: # wide N (e.g. qkv_proj, gate_up_proj)
290 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 64, 64, 4
291 + else: # N ~= K (e.g. o_proj)
292 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 32, 64, 4
293 + elif M <= 1024:
294 + if K >= 2 * N: # tall K (e.g. down_proj)
295 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 64, 64, 4
296 + elif N >= 4 * K: # very wide N (e.g. gate_up_proj)
297 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 64, 64, 8
298 + else:
299 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 128, 32, 4
300 + else:
301 + if K >= 2 * N: # tall K (e.g. down_proj)
302 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 512, 32, 16
303 + else:
304 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 64, 64, 8
305 + else:
306 + num_warps = 4
307 + if M <= 32:
308 + BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32
309 + elif M <= 64:
310 + BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
311 + else:
312 + BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
313 +
314 + # The kernel loads one scale per BLOCK_K tile, so BLOCK_K must not
315 + # exceed group_size — otherwise elements in the tile that belong to
316 + # a different group would get the wrong scale.
317 + BLOCK_K = min(BLOCK_K, group_size)
318 +
319 + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
320 +
321 + extra_kwargs = {} if num_stages is None else {"num_stages": num_stages}
322 + _triton_w4a16_skinny_fmt_kernel[grid](
323 + a,
324 + b_q,
325 + scales,
326 + zp if has_zp else scales, # dummy pointer when no zp (unused)
327 + c,
328 + M,
329 + N,
330 + K,
331 + K8,
332 + num_groups,
333 + group_size=group_size,
334 + ZP_BIAS=zp_bias,
335 + HAS_ZP=has_zp,
336 + BLOCK_M=BLOCK_M,
337 + BLOCK_N=BLOCK_N,
338 + BLOCK_K=BLOCK_K,
339 + num_warps=num_warps,
340 + **extra_kwargs,
341 + )
342 + return c
343 +
344 +
345 +# ---------------------------------------------------------------------------
346 +# Weight packing
347 +# ---------------------------------------------------------------------------
348 +
349 +
350 +def pack_int4_exllama_shuffle(w_uint4: torch.Tensor) -> torch.Tensor:
351 + """Pack uint4 values into ExLlama shuffle format: [N, K] -> [N, K//8] int32.
352 +
353 + Each int32 packs 8 K-values with interleave order [0,2,4,6,1,3,5,7].
354 + """
355 + N_dim, K_dim = w_uint4.shape
356 + assert K_dim % 8 == 0
357 + g = w_uint4.to(torch.uint8).view(N_dim, K_dim // 8, 8).to(torch.int32)
358 + return (
359 + g[:, :, 0]
360 + | (g[:, :, 2] << 4)
361 + | (g[:, :, 4] << 8)
362 + | (g[:, :, 6] << 12)
363 + | (g[:, :, 1] << 16)
364 + | (g[:, :, 3] << 20)
365 + | (g[:, :, 5] << 24)
366 + | (g[:, :, 7] << 28)
367 + )
368 +
369 +
370 +# ---------------------------------------------------------------------------
371 +# Hybrid dispatch logic
372 +# ---------------------------------------------------------------------------
373 +
374 +
375 +def _rdna_hybrid_w4a16_apply_impl(
376 + x_2d: torch.Tensor,
377 + w_q: torch.Tensor,
378 + w_s: torch.Tensor,
379 + w_zp: torch.Tensor | None,
380 + bias: torch.Tensor | None,
381 + cu_count: int,
382 + group_size: int,
383 +) -> torch.Tensor:
384 + """Dispatch between skinny GEMM and Triton based on batch size M."""
385 + import vllm._custom_ops as ops
386 +
387 + M = x_2d.shape[0]
388 + K = x_2d.shape[1]
389 + N = w_q.shape[0]
390 +
391 + if M <= MAX_SKINNY_BATCH_SIZE and K * M <= LDS_CAPACITY_ELEMENTS:
392 + # record_function is not torch.compile-safe; use nullcontext when
393 + # compiling to keep the op traceable.
394 + ctx = (
395 + nullcontext()
396 + if torch.compiler.is_compiling()
397 + else torch.profiler.record_function(f"wvsplitk_int4 {M}x{N}x{K}")
398 + )
399 + with ctx:
400 + return ops.wvSplitK_int4_g(w_q, x_2d, w_s, cu_count, group_size, w_zp, bias)
401 +
402 + ctx = (
403 + nullcontext()
404 + if torch.compiler.is_compiling()
405 + else torch.profiler.record_function(f"hybrid_triton_w4a16 {M}x{N}x{K}")
406 + )
407 + with ctx:
408 + output = triton_w4a16_skinny_fmt_gemm(
409 + a=x_2d,
410 + b_q=w_q.view(torch.int32),
411 + scales=w_s,
412 + group_size=group_size,
413 + zp=w_zp,
414 + )
415 + if bias is not None:
416 + output.add_(bias)
417 + return output
418 +
419 +
420 +def _rdna_hybrid_w4a16_apply_fake(
421 + x_2d: torch.Tensor,
422 + w_q: torch.Tensor,
423 + w_s: torch.Tensor,
424 + w_zp: torch.Tensor | None,
425 + bias: torch.Tensor | None,
426 + cu_count: int,
427 + group_size: int,
428 +) -> torch.Tensor:
429 + M = x_2d.size(0)
430 + N = w_q.size(0)
431 + return torch.empty((M, N), dtype=x_2d.dtype, device=x_2d.device)
432 +
433 +
434 +direct_register_custom_op(
435 + op_name="rdna_hybrid_w4a16_apply",
436 + op_func=_rdna_hybrid_w4a16_apply_impl,
437 + mutates_args=[],
438 + fake_impl=_rdna_hybrid_w4a16_apply_fake,
439 +)
440 +
441 +
442 +class RDNAHybridW4A16LinearKernel(MPLinearKernel):
443 + """Hybrid W4A16 kernel: HIP skinny for decode, Triton for prefill.
444 +
445 + Stores the weights once as int8 [N, K//2] (ExLlama shuffle packed). The
446 + HIP skinny kernel reads it directly; the triton kernel reinterprets the
447 + same buffer as int32 [N, K//8] via a view, so there is no dual weight
448 + storage.
449 + """
450 +
451 + SUPPORTED_QUANT_TYPES = [
452 + scalar_types.uint4b8, # symmetric GPTQ (bias=8)
453 + scalar_types.uint4, # asymmetric (zero_points)
454 + ]
455 +
456 + @classmethod
457 + def get_min_capability(cls) -> int:
458 + # Arch filtering is handled by can_implement (_on_gfx1x check)
459 + return 0
460 +
461 + @classmethod
462 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
463 + if not current_platform.is_rocm():
464 + return False, "RDNAHybridW4A16LinearKernel only targets ROCm"
465 +
466 + if not _on_gfx1x():
467 + return False, "RDNAHybridW4A16LinearKernel only targets gfx11/gfx12"
468 +
469 + if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
470 + return (
471 + False,
472 + f"Quant type {c.weight_type} not supported; "
473 + f"supported: {cls.SUPPORTED_QUANT_TYPES}",
474 + )
475 +
476 + if c.act_type not in (torch.float16, torch.bfloat16):
477 + return False, "requires float16 or bfloat16 activations"
478 +
479 + if c.has_g_idx:
480 + return False, "does not support g_idx reordering"
481 +
482 + gs = c.group_size
483 + if gs not in SUPPORTED_GROUP_SIZES:
484 + return (
485 + False,
486 + f"Group size {gs} not supported; supported: {SUPPORTED_GROUP_SIZES}",
487 + )
488 +
489 + K = c.partition_weight_shape[0]
490 + if K % 16 != 0:
491 + return False, f"K={K} must be divisible by 16"
492 +
493 + if K % gs != 0:
494 + return (
495 + False,
496 + f"K={K} must be divisible by group_size={gs}",
497 + )
498 +
499 + return True, None
500 +
501 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
502 + c = self.config
503 +
504 + w_q_raw = getattr(layer, self.w_q_name)
505 + w_s_raw = getattr(layer, self.w_s_name)
506 +
507 + unpacked = unpack_quantized_values_into_int32(
508 + w_q_raw.data, c.weight_type, packed_dim=w_q_raw.packed_dim
509 + )
510 + # AWQ weights arrive as (K, N) with output_dim=1;
511 + # compressed-tensors arrive as (N, K) with output_dim=0.
512 + if getattr(w_q_raw, "output_dim", 0) != 0:
513 + unpacked = unpacked.t().contiguous()
514 +
515 + shuffled = pack_int4_exllama_shuffle(unpacked)
516 + # Store as int8; Triton reinterprets via .view(torch.int32) at apply time.
517 + w_q_skinny = shuffled.contiguous().view(torch.int8)
518 +
519 + permute_param_layout_(w_s_raw, input_dim=1, output_dim=0)
520 + w_s_skinny = w_s_raw.data.contiguous()
521 +
522 + if c.zero_points:
523 + assert self.w_zp_name is not None
524 + w_zp_raw = getattr(layer, self.w_zp_name)
525 + permute_param_layout_(w_zp_raw, input_dim=1, output_dim=0, packed_dim=0)
526 + zp_unpacked = unpack_quantized_values_into_int32(
527 + w_zp_raw.data, c.weight_type, packed_dim=0
528 + )
529 + w_zp = zp_unpacked.to(c.act_type).contiguous()
530 + self._transform_param(layer, self.w_zp_name, lambda x: w_zp)
531 +
532 + self._transform_param(layer, self.w_q_name, lambda x: w_q_skinny)
533 + self._transform_param(layer, self.w_s_name, lambda x: w_s_skinny)
534 +
535 + def apply_weights(
536 + self,
537 + layer: torch.nn.Module,
538 + x: torch.Tensor,
539 + bias: torch.Tensor | None = None,
540 + ) -> torch.Tensor:
541 + from vllm.utils.platform_utils import num_compute_units
542 +
543 + c = self.config
544 + w_q, w_s, w_zp, _ = self._get_weight_params(layer)
545 +
546 + x_2d = x.reshape(-1, x.shape[-1])
547 + N = w_q.shape[0]
548 + out_shape = x.shape[:-1] + (N,)
549 +
550 + cu_count = num_compute_units()
551 + output = torch.ops.vllm.rdna_hybrid_w4a16_apply(
552 + x_2d,
553 + w_q,
554 + w_s,
555 + w_zp,
556 + bias,
557 + cu_count,
558 + c.group_size,
559 + )
560 + return output.reshape(out_shape)