A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 11:00 UTC
[Kernel][ROCm] Native W4A16 kernel for AMD RDNA3 (gfx1100) — fp16 + bf16 (#41394)
Signed-off-by: JartX <sagformas@epdcenter.es>
JartX committed
May 29, 2026 at 13:04 UTC
0cff0741ff8853beb2cae862c0daa8133cb1aa43
13 files changed
+3846
CMakeLists.txt
+12
@@ -1284,6 +1284,14 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
1284
"csrc/rocm/skinny_gemms.cu"
1285
"csrc/rocm/attention.cu")
1286
1287
+ set(VLLM_ROCM_HAS_GFX1100 OFF)
1288
+ if(VLLM_GPU_ARCHES MATCHES "gfx1100")
1289
+ set(VLLM_ROCM_HAS_GFX1100 ON)
1290
+ list(APPEND VLLM_ROCM_EXT_SRC
1291
+ "csrc/rocm/q_gemm_rdna3.cu"
1292
+ "csrc/rocm/q_gemm_rdna3_wmma.cu")
1293
+ endif()
1294
+
1295
define_extension_target(
1296
_rocm_C
1297
DESTINATION vllm
@@ -1293,6 +1301,10 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
1301
ARCHITECTURES ${VLLM_GPU_ARCHES}
1302
USE_SABI 3
1303
WITH_SOABI)
1304
+
1305
+ if(VLLM_ROCM_HAS_GFX1100)
1306
+ target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
1307
+ endif()
1308
endif()
1309
1310
# Must run after the last HIP `define_extension_target` so every extension
csrc/rocm/ops.h
+9
@@ -18,6 +18,15 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
18
const at::Tensor& scale_a, const at::Tensor& scale_b,
19
const int64_t CuCount);
20
21
+torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
22
+ torch::Tensor b_qzeros, torch::Tensor b_scales,
23
+ torch::Tensor b_g_idx, bool use_v2_format);
24
+
25
+torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
26
+ torch::Tensor b_qzeros,
27
+ torch::Tensor b_scales,
28
+ torch::Tensor b_g_idx, bool use_v2_format);
29
+
30
void paged_attention(
31
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
32
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
csrc/rocm/q_gemm_rdna3.cu
new
+780
@@ -0,0 +1,780 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+//
4
+// W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the
5
+// activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit
6
+// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes:
7
+//
8
+// 1. Direct write to the T-typed output via packed CAS-loop on a 64-bit
9
+// word (atomic_add_pk4_{f16,bf16}). gfx11 has no native
10
+// v_global_atomic_pk_add_{f16,bf16}, so the kernel emulates one with
11
+// global_atomic_cmpswap_b64. This avoids the M*N*4-byte FP32 scratch
12
+// buffer + memset + cast-pass that an fp32-accumulator design would
13
+// need; the caller passes a zero-initialised T-typed output tensor
14
+// and every block atomically adds its partial sum into it.
15
+//
16
+// 2. The bf16 path uses a dedicated bit-trick that avoids the fp16-only
17
+// "upper nibble * 16" trick, which would overflow the 7-bit bf16
18
+// mantissa. See qdq_4_rdna3.cuh for details.
19
+//
20
+// 3. Wave32 geometry sized for high CU saturation: THREADS_X=256
21
+// (8 waves per block) and BLOCK_KN_SIZE=256, with each thread
22
+// computing 4 N output columns. gridDim.z = K / BLOCK_KN_SIZE
23
+// splits K and the output is atomically accumulated. fp16 uses
24
+// v_dot2_f32_f16 (__builtin_amdgcn_fdot2) for the inner dot;
25
+// bf16 widens to fp32 (no v_pk_fma_bf16 on gfx11) and accumulates
26
+// with v_fma_f32. M_COUNT ∈ {1,2,4,8} is selected at launch
27
+// based on size_m.
28
+//
29
+// 4. The bf16 dispatch with M >= 16 forwards to the WMMA kernel in
30
+// q_gemm_rdna3_wmma.cu (separate translation unit) where
31
+// v_wmma_f32_16x16x16_bf16_w32 wins. The fp16 path always stays
32
+// scalar (the bit-trick dequant beats WMMA below M=64).
33
+
34
+#include <cstdint>
35
+#include <cstdio>
36
+
37
+#include <torch/all.h>
38
+#include <c10/cuda/CUDAGuard.h>
39
+#include <ATen/cuda/CUDAContext.h>
40
+
41
+#include <hip/hip_runtime.h>
42
+#include <hip/hip_bf16.h>
43
+#include <hip/hip_fp16.h>
44
+
45
+#include "qdq_4_rdna3.cuh"
46
+
47
+#if defined(__HIPCC__) && defined(__gfx1100__)
48
+ #define __HIP__RDNA3__
49
+#endif
50
+
51
+namespace vllm {
52
+namespace gptq_rdna3 {
53
+
54
+// BLOCK_KN_SIZE = 256 (was 128 in exllama). Each block covers 256 K
55
+// elements and THREADS_X*4 = 1024 N columns. For Qwen-class K=4096 this
56
+// halves gridDim.z (32 → 16) and therefore halves the atomic count per
57
+// output position vs the exllama default. THREADS_X=256 = 8 waves on RDNA3
58
+// wave32; with ~32 wave slots per CU we still fit 4 blocks per CU at peak.
59
+//
60
+// We tried BLOCK_KN_SIZE=512 (microbench on Qwen3.6-27B): bf16 improved
61
+// 5-10% at large M (atomic CAS halved), but fp16 decode regressed up to
62
+// +40% on qkv-square (32 → 45 μs at M=1). Cause: 16 waves/block × 16
63
+// total blocks for [M=1, K=N=4096] only saturates ~8 of the 96 CUs,
64
+// breaking memory-latency hiding for the fp16 path which is already
65
+// memory-bound. Reverted to 256; bf16 keeps most of its gains from the
66
+// fp32 dequant rewrite alone.
67
+#define BLOCK_KN_SIZE 256
68
+#define THREADS_X 256
69
+
70
+// Device code below is RDNA3-only; non-RDNA3 device passes fall through to
71
+// the empty __global__ stub at the #else below for symbol parity.
72
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
73
+
74
+// ---------------------------------------------------------------------------
75
+// Per-dtype helpers. We avoid heavy template metaprogramming and just provide
76
+// overloaded inline functions; the kernel below selects via `if constexpr`.
77
+// ---------------------------------------------------------------------------
78
+
79
+// Type-generic zero — both half and bf16_t in HIP/ROCm have a converting
80
+// constructor from float, but going through __float2half_rn / __float2bfloat16
81
+// is the unambiguously correct path on every ROCm version.
82
+template <typename T>
83
+__forceinline__ __device__ T tzero();
84
+
85
+template <>
86
+__forceinline__ __device__ half tzero<half>() {
87
+ return __float2half_rn(0.0f);
88
+}
89
+
90
+template <>
91
+__forceinline__ __device__ bf16_t tzero<bf16_t>() {
92
+ return __float2bfloat16(0.0f);
93
+}
94
+
95
+__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
96
+ // RDNA3 has v_dot2_f32_f16 (`__builtin_amdgcn_fdot2`) which computes
97
+ // fp32 += a.x*b.x + a.y*b.y in a single instruction with the accumulator
98
+ // staying in fp32 throughout. hipcc 7.2 does NOT peephole the obvious
99
+ // `__hfma2 + cast + add` pattern into v_dot2 (verified by ISA
100
+ // disassembly: 0 v_dot2_f32_f16 vs 256 v_cvt_f32_f16 + 218 v_add_f32 in
101
+ // the M_COUNT=8 kernel before this change), so we issue the builtin
102
+ // explicitly. Saves the trailing 2× v_cvt_f32_f16 + v_add_f32 (3 ops)
103
+ // per dot22_8_f call vs the half2-accumulator form. With 128 calls per
104
+ // K=32 step that's ~384 ops/K-step less issue pressure on the VALU.
105
+ //
106
+ // Numerical bonus: accumulator stays fp32 throughout the dot. The old
107
+ // form accumulated 8 muladds in fp16 (10-bit mantissa) before casting,
108
+ // which could lose ~3 bits of precision on borderline magnitudes.
109
+ float result = 0.0f;
110
+ const half2* a2_ptr = (const half2*)a_ptr;
111
+ #pragma unroll
112
+ for (int i = 0; i < 4; i++) {
113
+ result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
114
+ }
115
+ return result;
116
+}
117
+
118
+__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4],
119
+ const bf16_t* a_ptr) {
120
+ // RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in
121
+ // the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc
122
+ // therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised
123
+ // fallback (single-element FMAs or fp32 round-trips), which empirically
124
+ // runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16
125
+ // decode path was paying that tax in full, scaling linearly with M (the
126
+ // fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate
127
+ // and the kernel becomes memory-bound).
128
+ //
129
+ // Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs)
130
+ // and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA
131
+ // count, but each FMA is fast. Bonus: the accumulator is now fp32
132
+ // throughout instead of bf16, which is also numerically more accurate
133
+ // (no compounding bf16-rounding inside the dot loop).
134
+ float result = 0.0f;
135
+ #pragma unroll
136
+ for (int i = 0; i < 4; i++) {
137
+ uint32_t aw, dw;
138
+ __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
139
+ __builtin_memcpy(&dw, &dq[i], sizeof(uint32_t));
140
+ // bf16 in low 16 bits → fp32 by left-shifting into the upper half.
141
+ // bf16 in high 16 bits → already aligned with fp32's upper half.
142
+ float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
143
+ float a_y = __uint_as_float(aw & 0xFFFF0000u);
144
+ float d_x = __uint_as_float((dw & 0xFFFFu) << 16);
145
+ float d_y = __uint_as_float(dw & 0xFFFF0000u);
146
+ result = __fmaf_rn(d_x, a_x, result);
147
+ result = __fmaf_rn(d_y, a_y, result);
148
+ }
149
+ return result;
150
+}
151
+
152
+// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already
153
+// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t
154
+// overload above does for dq (still need to widen A from bf16). Wins more
155
+// at high N: the bf162_t version's per-call widening cost scales with the
156
+// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only
157
+// for A widening (M_COUNT × 4 × 4 widens, half as many).
158
+__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
159
+ const bf16_t* a_ptr) {
160
+ float result = 0.0f;
161
+ #pragma unroll
162
+ for (int i = 0; i < 4; i++) {
163
+ uint32_t aw;
164
+ __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
165
+ float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
166
+ float a_y = __uint_as_float(aw & 0xFFFF0000u);
167
+ result = __fmaf_rn(dq[2 * i + 0], a_x, result);
168
+ result = __fmaf_rn(dq[2 * i + 1], a_y, result);
169
+ }
170
+ return result;
171
+}
172
+
173
+// ---------------------------------------------------------------------------
174
+// Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS).
175
+// RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those
176
+// landed on gfx940 / gfx1250 respectively), so this lowers to
177
+// global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to
178
+// write 4 output columns per row in a single atomic operation — half the
179
+// atomic instruction count and half the contention vs two 32-bit CAS calls.
180
+//
181
+// Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer +
182
+// cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue
183
+// cast pass that an fp32-accumulator design would need.
184
+//
185
+// 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4
186
+// (always multiple of 4), and partition_weight_shape[1] is required to be a
187
+// multiple of 8 by can_implement(), so every (m, n) write target is 8-byte
188
+// aligned. Required by global_atomic_cmpswap_b64.
189
+// ---------------------------------------------------------------------------
190
+
191
+__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
192
+ half2 v23) {
193
+ unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
194
+ unsigned long long old = *addr_u;
195
+ while (true) {
196
+ union {
197
+ unsigned long long u;
198
+ half2 h2[2];
199
+ } cur, sum;
200
+ cur.u = old;
201
+ sum.h2[0] = __hadd2(cur.h2[0], v01);
202
+ sum.h2[1] = __hadd2(cur.h2[1], v23);
203
+ unsigned long long prev = atomicCAS(addr_u, old, sum.u);
204
+ if (prev == old) break;
205
+ old = prev;
206
+ }
207
+}
208
+
209
+__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
210
+ bf162_t v23) {
211
+ unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
212
+ unsigned long long old = *addr_u;
213
+ while (true) {
214
+ union {
215
+ unsigned long long u;
216
+ bf162_t b2[2];
217
+ } cur, sum;
218
+ cur.u = old;
219
+ sum.b2[0] = __hadd2(cur.b2[0], v01);
220
+ sum.b2[1] = __hadd2(cur.b2[1], v23);
221
+ unsigned long long prev = atomicCAS(addr_u, old, sum.u);
222
+ if (prev == old) break;
223
+ old = prev;
224
+ }
225
+}
226
+
227
+// Load one row's worth of 4 packed zeros (column n..n+3) from a [groups, N/8]
228
+// uint32 tensor. n is a multiple of 4 by construction (n = offset_n + t*4 with
229
+// offset_n = blockIdx.x * 512), so the 4 nibbles always live within one or two
230
+// uint32 words; in practice within one because n & 7 is 0 or 4.
231
+__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
232
+ int (&zeros)[4]) {
233
+ int qcol = n / 8;
234
+ int shift = (n & 0x07) * 4;
235
+ uint32_t d = qzeros_row[qcol] >> shift;
236
+ zeros[0] = (int)(d & 0xF);
237
+ zeros[1] = (int)((d >> 4) & 0xF);
238
+ zeros[2] = (int)((d >> 8) & 0xF);
239
+ zeros[3] = (int)((d >> 12) & 0xF);
240
+}
241
+
242
+template <typename T>
243
+__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
244
+ T (&scales)[4]) {
245
+ scales[0] = scales_row[n + 0];
246
+ scales[1] = scales_row[n + 1];
247
+ scales[2] = scales_row[n + 2];
248
+ scales[3] = scales_row[n + 3];
249
+}
250
+
251
+// ---------------------------------------------------------------------------
252
+// Main kernel.
253
+// ---------------------------------------------------------------------------
254
+
255
+template <typename T, int M_COUNT>
256
+__global__ void gemm_q4_kernel_rdna3(
257
+ const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
258
+ const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
259
+ T* __restrict__ c, const int size_m, const int size_n, const int size_k,
260
+ const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
261
+ const int t = threadIdx.x;
262
+ const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
263
+ const int offset_m = blockIdx.y * M_COUNT;
264
+ const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
265
+ const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
266
+ const int n = offset_n + t * 4;
267
+
268
+ // LDS layout: [M_COUNT][BLOCK_KN_SIZE + LDS_PAD]. The PAD=8 elements per M
269
+ // row break the natural 256-element/512-byte alignment that would otherwise
270
+ // collide on the same LDS bank when a thread reads block_a[0..M_COUNT-1][k]
271
+ // (same k, different m). Row stride becomes 264 elements * 2B = 528B = 132
272
+ // 4-byte banks, so m-stride hits banks (m*132)%32 = (m*4)%32 — distinct for
273
+ // all M_COUNT ≤ 8. Cost: 16B LDS per block, irrelevant.
274
+ constexpr int LDS_PAD = 8;
275
+ __shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD];
276
+
277
+ // Stage A: each thread loads 1 K element per M row into LDS (with optional
278
+ // act-order permutation). THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map.
279
+ // For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m
280
+ // are zero-padded so the dot product contribution is 0 (we then skip the
281
+ // atomic write for those rows below).
282
+ //
283
+ // M=1 fast path: skip LDS staging + __syncthreads entirely. All 256 threads
284
+ // read the SAME 8-element A window per inner step (a_off is uniform across
285
+ // the block), so the cache-line broadcast through L1 makes global reads as
286
+ // cheap as LDS reads. Measured: ~1% on 4B b=1, ~6% on 27B b=1 in=128.
287
+ static_assert(BLOCK_KN_SIZE == THREADS_X,
288
+ "BLOCK_KN_SIZE must equal THREADS_X (1 K element per thread)");
289
+ // The M=1 fast path (skip LDS) only has a global-read code path for bf16
290
+ // (the v_dot2_f32_bf16 branch). The fp16 inner loop still indexes
291
+ // block_a[m][a_off] unconditionally, so for fp16 we MUST stage A through
292
+ // LDS even at M=1 to avoid reading uninitialized shared memory.
293
+ constexpr bool USE_LDS_A = (M_COUNT > 1) || std::is_same<T, half>::value;
294
+ if constexpr (USE_LDS_A) {
295
+ if (offset_k + t < end_k) {
296
+ #pragma unroll
297
+ for (int m = 0; m < M_COUNT; ++m) {
298
+ T av;
299
+ if (offset_m + m < size_m) {
300
+ const T* a_row = a + (offset_m + m) * size_k;
301
+ if (b_q_perm)
302
+ av = a_row[b_q_perm[offset_k + t]];
303
+ else
304
+ av = a_row[offset_k + t];
305
+ } else {
306
+ av = tzero<T>(); // zero-pad invalid M rows
307
+ }
308
+ block_a[m][t] = av;
309
+ }
310
+ }
311
+
312
+ // Threads beyond the right edge of N have nothing to do. Note: we must NOT
313
+ // return before __syncthreads() if any thread in the block participates in
314
+ // the LDS load above — but here all THREADS_X (=256) threads always do,
315
+ // regardless of whether their `n` is in bounds.
316
+ __syncthreads();
317
+ } else if (b_q_perm) {
318
+ // bf16 M=1 fast path skips LDS, but its global read below is sequential
319
+ // and cannot apply act-order. When a permutation is present, stage the
320
+ // single A row through LDS (as fp16 / M>1 do) so the read picks it up.
321
+ // b_q_perm is block-uniform, so the __syncthreads is non-divergent.
322
+ if (offset_k + t < end_k)
323
+ block_a[0][t] = a[offset_m * size_k + b_q_perm[offset_k + t]];
324
+ __syncthreads();
325
+ }
326
+ if (n >= size_n) return;
327
+
328
+ // Group bookkeeping. We require size_k % groups == 0 (groupsize divides K).
329
+ const int groupsize = size_k / groups;
330
+ int group = offset_k / groupsize;
331
+ int nextgroup = (group + 1) * groupsize;
332
+
333
+ // qweight stride: weights are [K/8, N] uint32 with K packed at dim 0.
334
+ int qk = offset_k / 8;
335
+ const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
336
+
337
+ // Per-column dequant constants. We hold one set of (z, y) pairs per column.
338
+ // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper-
339
+ // nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant
340
+ // produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA
341
+ // bypass for the missing v_pk_fma_bf16 on gfx11.
342
+ half2 z1z16_h[4][2], y1y16_h[4][2];
343
+ float z_b_f[4], y_b_f[4];
344
+
345
+ auto refresh_group = [&](int g) {
346
+ const uint32_t* qz_row = b_qzeros + g * (size_n / 8);
347
+ const T* sc_row = b_scales + g * size_n;
348
+ int zeros[4];
349
+ T scales[4];
350
+ load4_zeros(qz_row, n, zeros);
351
+ load4_scales<T>(sc_row, n, scales);
352
+ if constexpr (std::is_same<T, half>::value) {
353
+ #pragma unroll
354
+ for (int i = 0; i < 4; ++i) {
355
+ prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i],
356
+ z1z16_h[i], y1y16_h[i]);
357
+ }
358
+ } else {
359
+ #pragma unroll
360
+ for (int i = 0; i < 4; ++i) {
361
+ prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i],
362
+ z_b_f[i], y_b_f[i]);
363
+ }
364
+ }
365
+ };
366
+
367
+ refresh_group(group);
368
+
369
+ float block_c[M_COUNT][4];
370
+ #pragma unroll
371
+ for (int m = 0; m < M_COUNT; ++m) {
372
+ #pragma unroll
373
+ for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
374
+ }
375
+
376
+ // Note on group-transition granularity: we check `k == nextgroup` at the
377
+ // start of each outer iteration (which advances K by 32). This is correct
378
+ // when group_size >= 32 OR group_size divides 32 evenly (groupsize is one
379
+ // of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the
380
+ // inner loop would cross a group boundary between j-iterations; we require
381
+ // group_size >= 32 here, mirroring exllama's assumption.
382
+ //
383
+ // Software pipelining: we issue all 4 vectorized weight loads up front
384
+ // before any dequant/FMA depends on them. This gives the AMDGPU backend
385
+ // freedom to schedule the global_loads early and overlap their latency
386
+ // with dequant + v_pk_fma_f16 of earlier iterations. Cost: 4×int4 = 16
387
+ // VGPRs in flight per thread, plenty of headroom on RDNA3.
388
+ int k = offset_k;
389
+ while (k < end_k) {
390
+ if (k == nextgroup) {
391
+ group++;
392
+ nextgroup += groupsize;
393
+ refresh_group(group);
394
+ }
395
+
396
+ // Prefetch all four j-iterations' weight words. The compiler emits 4
397
+ // global_load_b128 instructions back-to-back; the dependent dequant +
398
+ // FMA work below hides their latency.
399
+ int4 b_w[4];
400
+ #pragma unroll
401
+ for (int j = 0; j < 4; ++j) {
402
+ b_w[j] = *(const int4*)(b_ptr + j * size_n);
403
+ }
404
+ b_ptr += 4 * size_n;
405
+
406
+ #pragma unroll
407
+ for (int j = 0; j < 4; ++j) {
408
+ const int a_off = (k - offset_k) + 8 * j;
409
+
410
+ if constexpr (std::is_same<T, half>::value) {
411
+ half2 dq[4][4];
412
+ dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]);
413
+ dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]);
414
+ dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]);
415
+ dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]);
416
+
417
+ #pragma unroll
418
+ for (int m = 0; m < M_COUNT; ++m) {
419
+ const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
420
+ block_c[m][0] += dot22_8_f(dq[0], a_ptr);
421
+ block_c[m][1] += dot22_8_f(dq[1], a_ptr);
422
+ block_c[m][2] += dot22_8_f(dq[2], a_ptr);
423
+ block_c[m][3] += dot22_8_f(dq[3], a_ptr);
424
+ }
425
+ } else if constexpr (M_COUNT == 1) {
426
+ // bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of
427
+ // Hybrid PR #40977's wvSplitK_int4 kernel exactly so clang's
428
+ // InstCombine cannot fold the bf16→fp32 widening (LLVM #76000):
429
+ // * activations and magic-value weights share a fp32-aliased
430
+ // union (bytes written as uint32, read as bf16x2_t for the
431
+ // dot — pointer-cast opacity defeats the fold)
432
+ // * sum_a computed via a *second* v_dot2 with bf162(1,1) as the
433
+ // second operand, avoiding any explicit bf16→fp32 widen of A
434
+ // * bias correction y_b_f * partial + z_b_f * sum_a, identical
435
+ // to the previous fp32-FMA-chain path
436
+ //
437
+ // Net: 20 v_dot2_f32_bf16 + 8 fp32 FMA per int32 weight vs the
438
+ // previous 40 fp32 FMA. v_dot2 runs at full rate on gfx1100, so
439
+ // the substitution is ~2× cheaper for the inner accumulator.
440
+ typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
441
+ constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128)
442
+ constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0)
443
+ union pack4 {
444
+ float f[4];
445
+ uint32_t u[4];
446
+ };
447
+
448
+ uint32_t w[4];
449
+ __builtin_memcpy(w, &b_w[j], sizeof(int4));
450
+
451
+ // Load 8 bf16 activations as 4 uint32s (= 4 bf16x2 pairs) into a
452
+ // fp32-aliased union. Storing as uint32 keeps the IR-level type
453
+ // opaque so the inner v_dot2 cannot be folded to fp32 widening.
454
+ //
455
+ // A is read direct from global (no LDS staging — see USE_LDS_A above),
456
+ // except under act-order, where it comes from the permuted LDS copy.
457
+ pack4 a_pack;
458
+ {
459
+ const uint32_t* a_words =
460
+ b_q_perm
461
+ ? reinterpret_cast<const uint32_t*>(&block_a[0][a_off])
462
+ : reinterpret_cast<const uint32_t*>(a + offset_k + a_off);
463
+ a_pack.u[0] = a_words[0];
464
+ a_pack.u[1] = a_words[1];
465
+ a_pack.u[2] = a_words[2];
466
+ a_pack.u[3] = a_words[3];
467
+ }
468
+
469
+ // sum_a = Σ a[i]. Computed via 4× v_dot2_f32_bf16 with bf162(1,1) as
470
+ // the second operand — every bf16 pair contributes 1·a_lo + 1·a_hi.
471
+ // No fp32 widening of activations: the bytes go straight from LDS
472
+ // through v_dot2 into the fp32 accumulator.
473
+ float sum_a = 0.0f;
474
+ #pragma unroll
475
+ for (int b = 0; b < 4; ++b) {
476
+ sum_a = __builtin_amdgcn_fdot2_f32_bf16(
477
+ *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
478
+ sum_a, /*clamp=*/false);
479
+ }
480
+
481
+ // unroll 1 keeps q_pack alive only one col at a time (8 fp32 VGPRs
482
+ // recycled across cols), avoiding straight-line expansion that
483
+ // would inflate live-range to 32 VGPRs.
484
+ #pragma unroll 1
485
+ for (int col = 0; col < 4; ++col) {
486
+ // Build dequant magic values bf16(128 + nibble) directly into a
487
+ // fp32-aliased union via uint32 stores. No fp32 in the data flow
488
+ // until v_dot2 consumes the bytes.
489
+ pack4 q_pack;
490
+ const uint32_t qa = w[col];
491
+ q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
492
+ q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
493
+ q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
494
+ q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
495
+
496
+ // partial = Σ (128 + nibble[i]) · a[i], via 4× v_dot2_f32_bf16.
497
+ float partial = 0.0f;
498
+ #pragma unroll
499
+ for (int b = 0; b < 4; ++b) {
500
+ partial = __builtin_amdgcn_fdot2_f32_bf16(
501
+ *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
502
+ partial, /*clamp=*/false);
503
+ }
504
+
505
+ // block_c += y_b_f * partial + z_b_f * sum_a
506
+ // y_b_f = scale, z_b_f = -(128+zero)*scale
507
+ // partial holds (128 + nibble) · a; subtracting (128+zero)·sum_a
508
+ // and scaling yields scale · (nibble - zero) · a as required.
509
+ block_c[0][col] =
510
+ __fmaf_rn(y_b_f[col], partial,
511
+ __fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
512
+ }
513
+ } else {
514
+ // bf16 M_COUNT > 1 path with v_dot2_f32_bf16. Same opacity trick as
515
+ // the M=1 branch: activations + magic-value weights stored in
516
+ // fp32-aliased unions, dot via __builtin_amdgcn_fdot2_f32_bf16 with
517
+ // pointer-cast to bf16x2_t. sum_a[m] computed via second v_dot2
518
+ // with BF16_ONES; bias correction (y_b_f * partial + z_b_f * sum_a)
519
+ // applied after the dot. Magic values built once per col and reused
520
+ // across all M rows — amortizes dequant cost across M_COUNT.
521
+ typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
522
+ constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128)
523
+ constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0)
524
+ union pack4 {
525
+ float f[4];
526
+ uint32_t u[4];
527
+ };
528
+
529
+ uint32_t w[4];
530
+ __builtin_memcpy(w, &b_w[j], sizeof(int4));
531
+
532
+ // Load M_COUNT × 8 bf16 activations as 4 uint32s each into pack4
533
+ // unions. Stored as uint32 to keep IR-level types opaque (defeats
534
+ // InstCombine fold). At M_COUNT=8 this is 32 fp32 VGPRs — within RDNA3
535
+ // budget.
536
+ pack4 a_pack[M_COUNT];
537
+ #pragma unroll
538
+ for (int m = 0; m < M_COUNT; ++m) {
539
+ const uint32_t* a_words =
540
+ reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
541
+ a_pack[m].u[0] = a_words[0];
542
+ a_pack[m].u[1] = a_words[1];
543
+ a_pack[m].u[2] = a_words[2];
544
+ a_pack[m].u[3] = a_words[3];
545
+ }
546
+
547
+ // sum_a[m] = Σ a[m][i] via 4× v_dot2 with bf162(1,1) — no fp32 widen.
548
+ float sum_a[M_COUNT];
549
+ #pragma unroll
550
+ for (int m = 0; m < M_COUNT; ++m) {
551
+ float s = 0.0f;
552
+ #pragma unroll
553
+ for (int b = 0; b < 4; ++b) {
554
+ s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
555
+ *((const bf16x2_t*)&BF16_ONES),
556
+ s, /*clamp=*/false);
557
+ }
558
+ sum_a[m] = s;
559
+ }
560
+
561
+ // Per col: build magic-value pack, dot against all M activations.
562
+ // unroll 1 keeps q_pack live one col at a time (8 fp32 VGPRs recycled)
563
+ // — same register-pressure trick as the previous fp32 path.
564
+ #pragma unroll 1
565
+ for (int col = 0; col < 4; ++col) {
566
+ pack4 q_pack;
567
+ const uint32_t qa = w[col];
568
+ q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
569
+ q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
570
+ q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
571
+ q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
572
+
573
+ #pragma unroll
574
+ for (int m = 0; m < M_COUNT; ++m) {
575
+ float partial = 0.0f;
576
+ #pragma unroll
577
+ for (int b = 0; b < 4; ++b) {
578
+ partial = __builtin_amdgcn_fdot2_f32_bf16(
579
+ *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
580
+ partial, /*clamp=*/false);
581
+ }
582
+ // block_c += y_b_f * partial + z_b_f * sum_a (same correction as
583
+ // M=1)
584
+ block_c[m][col] =
585
+ __fmaf_rn(y_b_f[col], partial,
586
+ __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
587
+ }
588
+ }
589
+ }
590
+ }
591
+ k += 32; // 4 weight words * 8 nibbles = 32 K elements
592
+ }
593
+
594
+ // Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all
595
+ // four lanes in a single 64-bit CAS write directly to the T-typed output
596
+ // (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with
597
+ // a single b64 op we halve the atomic instruction count vs two b32 CAS
598
+ // calls, AND save the FP32 buffer + memset + cast pass entirely.
599
+ #pragma unroll
600
+ for (int m = 0; m < M_COUNT; ++m) {
601
+ if (offset_m + m >= size_m) continue; // skip padding rows past size_m
602
+ T* out = c + (offset_m + m) * size_n + n;
603
+ if constexpr (std::is_same<T, half>::value) {
604
+ half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
605
+ __float2half_rn(block_c[m][1]));
606
+ half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
607
+ __float2half_rn(block_c[m][3]));
608
+ atomic_add_pk4_f16(out, r01, r23);
609
+ } else {
610
+ bf162_t r01;
611
+ r01.x = __float2bfloat16(block_c[m][0]);
612
+ r01.y = __float2bfloat16(block_c[m][1]);
613
+ bf162_t r23;
614
+ r23.x = __float2bfloat16(block_c[m][2]);
615
+ r23.y = __float2bfloat16(block_c[m][3]);
616
+ atomic_add_pk4_bf16(out, r01, r23);
617
+ }
618
+ }
619
+}
620
+
621
+#else // non-RDNA3 device pass: empty __global__ for symbol parity.
622
+
623
+template <typename T, int M_COUNT>
624
+__global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*,
625
+ const T*, T*, const int, const int,
626
+ const int, const int, const int,
627
+ const int*) {}
628
+
629
+#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
630
+
631
+// ---------------------------------------------------------------------------
632
+// Launcher.
633
+// ---------------------------------------------------------------------------
634
+
635
+template <typename T, int M_COUNT>
636
+void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight,
637
+ const uint32_t* b_qzeros, const T* b_scales,
638
+ const int* b_q_perm, T* c, int size_m,
639
+ int size_n, int size_k, int groups,
640
+ int zero_offset, cudaStream_t stream) {
641
+ dim3 block(THREADS_X);
642
+ dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
643
+ (size_m + M_COUNT - 1) / M_COUNT,
644
+ (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
645
+
646
+ gemm_q4_kernel_rdna3<T, M_COUNT><<<grid, block, 0, stream>>>(
647
+ a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
648
+ zero_offset, b_q_perm);
649
+}
650
+
651
+// Dispatch to the largest M_COUNT template that doesn't waste more than
652
+// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is
653
+// the right tool, not bigger M_COUNT in the scalar dot-product path.
654
+//
655
+// Tile-waste table:
656
+// M=1 -> M_COUNT=1 (no waste)
657
+// M=2,3 -> M_COUNT=2 (M=3 wastes 1/2 of last tile)
658
+// M=4-7 -> M_COUNT=4 (worst case M=5: wastes 3/4 of last tile)
659
+// M=8-15-> M_COUNT=8 (worst case M=9: wastes 7/8 of last tile)
660
+// "Wasted" rows are zero-padded in LDS and skip the atomic write, so they
661
+// only burn instructions on the last block, never affect correctness.
662
+template <typename T>
663
+void launch_gemm_q4(const T* a, const uint32_t* b_q_weight,
664
+ const uint32_t* b_qzeros, const T* b_scales,
665
+ const int* b_q_perm, T* c, int size_m, int size_n,
666
+ int size_k, int groups, bool use_v2_format,
667
+ cudaStream_t stream) {
668
+ const int zero_offset = use_v2_format ? 0 : 1;
669
+
670
+ if (size_m == 1) {
671
+ launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
672
+ c, size_m, size_n, size_k, groups,
673
+ zero_offset, stream);
674
+ } else if (size_m <= 3) {
675
+ launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
676
+ c, size_m, size_n, size_k, groups,
677
+ zero_offset, stream);
678
+ } else if (size_m <= 7) {
679
+ launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
680
+ c, size_m, size_n, size_k, groups,
681
+ zero_offset, stream);
682
+ } else {
683
+ // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the
684
+ // WMMA path, but if it falls through we still produce correct output —
685
+ // just leaving 3-5× of throughput on the table for prefill workloads.
686
+ launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
687
+ c, size_m, size_n, size_k, groups,
688
+ zero_offset, stream);
689
+ }
690
+}
691
+
692
+} // namespace gptq_rdna3
693
+} // namespace vllm
694
+
695
+// ---------------------------------------------------------------------------
696
+// Public entry point.
697
+// ---------------------------------------------------------------------------
698
+//
699
+// Inputs:
700
+// a [M, K] half or bfloat16
701
+// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle)
702
+// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros)
703
+// b_scales [groups, N] half or bfloat16
704
+// b_g_idx [K] or empty int32 (act-order permutation; empty=identity)
705
+// use_v2_format bool (true = GPTQv2, no +1 zero offset)
706
+//
707
+// Output:
708
+// c [M, N] same dtype as a
709
+
710
+torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
711
+ torch::Tensor b_qzeros,
712
+ torch::Tensor b_scales,
713
+ torch::Tensor b_g_idx, bool use_v2_format);
714
+
715
+torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
716
+ torch::Tensor b_qzeros, torch::Tensor b_scales,
717
+ torch::Tensor b_g_idx, bool use_v2_format) {
718
+ if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 &&
719
+ b_q_weight.size(1) % 16 == 0 &&
720
+ ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) ||
721
+ (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) {
722
+ return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, b_g_idx,
723
+ use_v2_format);
724
+ }
725
+
726
+ TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
727
+ TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
728
+ TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor");
729
+ TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor");
730
+ TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]");
731
+ TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]");
732
+ TORCH_CHECK(
733
+ a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
734
+ "a must be half or bfloat16");
735
+ TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
736
+ "b_scales dtype must match a");
737
+
738
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
739
+ auto stream = at::cuda::getCurrentCUDAStream();
740
+
741
+ int size_m = (int)a.size(0);
742
+ int size_k = (int)a.size(1);
743
+ int size_n = (int)b_q_weight.size(1);
744
+ int groups = (int)b_qzeros.size(0);
745
+
746
+ TORCH_CHECK(b_q_weight.size(0) * 8 == size_k,
747
+ "b_q_weight first dim must be K/8");
748
+ TORCH_CHECK(b_scales.size(0) == groups,
749
+ "b_scales must have same group count as qzeros");
750
+ TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N");
751
+ TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)");
752
+
753
+ auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device());
754
+ at::Tensor c = torch::zeros({size_m, size_n}, opts);
755
+
756
+ const int* g_idx_ptr = nullptr;
757
+ if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) {
758
+ TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32,
759
+ "b_g_idx must be int32");
760
+ g_idx_ptr = (const int*)b_g_idx.data_ptr();
761
+ }
762
+
763
+ if (a.scalar_type() == torch::kHalf) {
764
+ vllm::gptq_rdna3::launch_gemm_q4<half>(
765
+ (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(),
766
+ (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(),
767
+ g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups,
768
+ use_v2_format, stream);
769
+ } else {
770
+ vllm::gptq_rdna3::launch_gemm_q4<vllm::gptq_rdna3::bf16_t>(
771
+ (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(),
772
+ (const uint32_t*)b_q_weight.data_ptr(),
773
+ (const uint32_t*)b_qzeros.data_ptr(),
774
+ (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), g_idx_ptr,
775
+ (vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups,
776
+ use_v2_format, stream);
777
+ }
778
+
779
+ return c;
780
+}
csrc/rocm/q_gemm_rdna3_wmma.cu
new
+2165
@@ -0,0 +1,2165 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+//
4
+// W4A16 GPTQ WMMA prefill kernel for AMD RDNA3 (gfx1100). This is the
5
+// matrix-instruction path for M >= 16; small-M decode lives in the sibling
6
+// file q_gemm_rdna3.cu and is exposed via a different op
7
+// (`gptq_gemm_rdna3`). Keeping the two paths in separate translation units
8
+// is intentional: an earlier attempt at putting WMMA in the same TU as the
9
+// scalar dot-product kernel introduced a compile-time interaction that
10
+// silently miscompiled the M=1 path even though the WMMA template was never
11
+// instantiated for M=1. Hipcc's optimizer appears to scope some decisions
12
+// at the TU level (likely register file / SGPR pressure heuristics across
13
+// all kernels in the TU), so we isolate.
14
+//
15
+// Hardware notes (RDNA3 / gfx1100 / RX 7900 XTX):
16
+// * v_wmma_f32_16x16x16_{f16,bf16}_w32 — 16×16×16 GEMM in one instruction
17
+// (~16 cycles per WMMA on wave32). Accumulator dtype is FP32; inputs
18
+// are fp16 or bf16. There is NO 16x16x16 with fp16/bf16 accumulator on
19
+// gfx11 that we'd want here (we always need fp32 accum to avoid loss
20
+// across many K iterations).
21
+// * Wave32 input fragment storage is "doubled" — lanes 16..31 hold a
22
+// copy of lanes 0..15 for the A and B fragments. The output C
23
+// fragment uses a different mapping: lane t holds COLUMN n=lane_lo
24
+// of the 16x16 output, with 8 elements alternating M rows by lane_hi
25
+// (lanes 0..15 = even rows, lanes 16..31 = odd rows). See the layout
26
+// diagram on `gemm_q4_wmma_kernel_16x16_1w` below for the full mapping.
27
+// * No native v_global_atomic_pk_add_{f16,bf16} on gfx11; the K-split
28
+// epilogue (gridDim.z > 1) emulates packed atomic add via a CAS-32
29
+// retry loop on a uint32 word covering 2 fp16/bf16 lanes. Within a
30
+// block we shuffle adjacent lanes via shfl_xor first so each pair of
31
+// output cols goes through a single atomic — no intra-block
32
+// contention. K_SPLIT == 1 keeps the original direct-write path with
33
+// each block owning its 16M × 16N output tile.
34
+
35
+#include <cstdint>
36
+
37
+#include <torch/all.h>
38
+#include <c10/cuda/CUDAGuard.h>
39
+#include <ATen/cuda/CUDAContext.h>
40
+
41
+#include <hip/hip_runtime.h>
42
+#include <hip/hip_bf16.h>
43
+#include <hip/hip_fp16.h>
44
+
45
+#include "qdq_4_rdna3.cuh"
46
+
47
+#if defined(__HIPCC__) && defined(__gfx1100__)
48
+ #define __HIP__RDNA3__
49
+#endif
50
+
51
+namespace vllm {
52
+namespace gptq_rdna3_wmma {
53
+
54
+// Pull dequant types from the sibling namespace.
55
+using vllm::gptq_rdna3::bf162_t;
56
+using vllm::gptq_rdna3::bf16_t;
57
+
58
+// Device code below uses RDNA3-only __builtin_amdgcn_wmma_* intrinsics;
59
+// non-RDNA3 device passes fall through to empty __global__ stubs at the
60
+// #else block at the end of this TU.
61
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
62
+
63
+// PRECISE dequant variants live HERE, not in the shared qdq_4_rdna3.cuh
64
+// header. Reason: hipcc takes different register/scheduling decisions when
65
+// the shared header grows, which caused a measurable decode-tk/s regression
66
+// in the scalar kernel even when these new functions were never called from
67
+// it. Keeping them in the WMMA TU only restores scalar's binary identity
68
+// to its tuned baseline.
69
+//
70
+// Numerics: the classic fp16 bit-trick (FMA form: q*scale +
71
+// (-(1024+zero)*scale)) loses up to ~0.025 per cell at scale=0.1 because
72
+// fp16(scale) ≠ scale and the FMA amplifies that by 1024× without cancelling
73
+// against the precomputed (1024+zero)*scale (which is rounded to fp16 BEFORE
74
+// the FMA).
75
+//
76
+// Fix: subtract (1024+zero) as an integer FIRST — exact in fp16 because
77
+// integers in [1024, 2047] are exactly representable — then multiply by
78
+// scale, incurring at most one half-ULP rounding. Costs one extra
79
+// instruction per dequant pair (sub+mul vs single FMA), worth it for the
80
+// WMMA path because K can be small (16) and errors don't average.
81
+__forceinline__ __device__ void prep_zero_scale_fp16_precise(uint32_t zero,
82
+ half scale,
83
+ half2& z_prep,
84
+ half2& y_prep) {
85
+ union {
86
+ uint16_t u;
87
+ half h;
88
+ } zu;
89
+ zu.u = (uint16_t)(0x6400 | zero);
90
+ z_prep = __half2half2(zu.h);
91
+ y_prep = __half2half2(scale);
92
+}
93
+
94
+__forceinline__ __device__ void dequant_4bit_8_fp16_precise(uint32_t qa,
95
+ half2 (&dq)[4],
96
+ half2 z_prep,
97
+ half2 y_prep) {
98
+ const uint32_t c0 = 0x64006400;
99
+ union {
100
+ uint32_t u;
101
+ half2 h2;
102
+ } q0, q1, q2, q3;
103
+ q0.u = ((qa >> 0) & 0x000F000F) | c0;
104
+ q1.u = ((qa >> 4) & 0x000F000F) | c0;
105
+ q2.u = ((qa >> 8) & 0x000F000F) | c0;
106
+ q3.u = ((qa >> 12) & 0x000F000F) | c0;
107
+ dq[0] = __hmul2(__hsub2(q0.h2, z_prep), y_prep);
108
+ dq[1] = __hmul2(__hsub2(q1.h2, z_prep), y_prep);
109
+ dq[2] = __hmul2(__hsub2(q2.h2, z_prep), y_prep);
110
+ dq[3] = __hmul2(__hsub2(q3.h2, z_prep), y_prep);
111
+}
112
+
113
+// fp32 prep for the bf16→bf16 fp32-internal dequant. Returns:
114
+// z_prep = -(128 + zero) * scale (folded bias for FMA)
115
+// y_prep = scale
116
+// Per-element FMA `q*y_prep + z_prep` then yields scale * (nibble - zero).
117
+__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
118
+ bf16_t scale,
119
+ float& z_prep,
120
+ float& y_prep) {
121
+ float scale_f = __bfloat162float(scale);
122
+ z_prep = -(128.0f + (float)zero) * scale_f;
123
+ y_prep = scale_f;
124
+}
125
+
126
+// fp32 → bf16 narrow that skips the defensive NaN-canonicalisation hipcc
127
+// emits for __float2bfloat16. Round-half-to-even via add (0x7FFF + lsb)
128
+// + truncate; no NaN check, since dequant outputs are bounded products
129
+// of (nibble - zero) ∈ [-15, 15] and bf16 scales — never NaN/Inf.
130
+__forceinline__ __device__ bf16_t f32_to_bf16_no_canon(float f) {
131
+ uint32_t fu = __float_as_uint(f);
132
+ uint32_t lsb = (fu >> 16) & 1u;
133
+ uint16_t out_u = (uint16_t)((fu + 0x7FFFu + lsb) >> 16);
134
+ bf16_t out;
135
+ __builtin_memcpy(&out, &out_u, sizeof(out));
136
+ return out;
137
+}
138
+
139
+// 4-bit GPTQ dequant for the bf16 WMMA path: 8 nibbles per int32 in qa,
140
+// outputs bf162_t dq[4]. Implementation rationale:
141
+//
142
+// gfx11 has no v_pk_fma_bf16, so __hmul2/__hsub2 on bf16 lower to a
143
+// widen-fp32-op-narrow chain that hipcc decorates with NaN canonicalisation
144
+// (28 extra VALU ops per call observed in the v5 ISA dump: v_cmp_u_f32 +
145
+// v_cndmask_b32 around every bf16 sub/mul). Doing the math in fp32 directly
146
+// — bit-cast widening (`__uint_as_float((bf16_bits) << 16)` is just a shift,
147
+// no NaN canon), one fused FMA per element, single `__float2bfloat16` narrow
148
+// at the end — eliminates the canon and is also strictly more precise (one
149
+// rounding step instead of two).
150
+//
151
+// Inputs match prep_zero_scale_bf16_f32:
152
+// z_prep = -(128 + zero) * scale (folded bias for FMA)
153
+// y_prep = scale
154
+// Per-element FMA `q*y_prep + z_prep` then yields scale * (nibble - zero).
155
+__forceinline__ __device__ void dequant_4bit_8_bf16_to_bf16(uint32_t qa,
156
+ bf162_t (&dq)[4],
157
+ float z_prep,
158
+ float y_prep) {
159
+ const uint32_t c0 = 0x43004300;
160
+ const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
161
+ const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
162
+ const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
163
+ const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
164
+ // bf16(128+nibble) bits → fp32 via left-shift by 16 (zero-extends mantissa).
165
+ const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16);
166
+ const float q0y = __uint_as_float(q0 & 0xFFFF0000u);
167
+ const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16);
168
+ const float q1y = __uint_as_float(q1 & 0xFFFF0000u);
169
+ const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16);
170
+ const float q2y = __uint_as_float(q2 & 0xFFFF0000u);
171
+ const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16);
172
+ const float q3y = __uint_as_float(q3 & 0xFFFF0000u);
173
+ // r = q*scale + (-(128+zero)*scale) = (nibble - zero)*scale, then narrow.
174
+ dq[0].x = f32_to_bf16_no_canon(__fmaf_rn(q0x, y_prep, z_prep));
175
+ dq[0].y = f32_to_bf16_no_canon(__fmaf_rn(q0y, y_prep, z_prep));
176
+ dq[1].x = f32_to_bf16_no_canon(__fmaf_rn(q1x, y_prep, z_prep));
177
+ dq[1].y = f32_to_bf16_no_canon(__fmaf_rn(q1y, y_prep, z_prep));
178
+ dq[2].x = f32_to_bf16_no_canon(__fmaf_rn(q2x, y_prep, z_prep));
179
+ dq[2].y = f32_to_bf16_no_canon(__fmaf_rn(q2y, y_prep, z_prep));
180
+ dq[3].x = f32_to_bf16_no_canon(__fmaf_rn(q3x, y_prep, z_prep));
181
+ dq[3].y = f32_to_bf16_no_canon(__fmaf_rn(q3y, y_prep, z_prep));
182
+}
183
+
184
+// ---------------------------------------------------------------------------
185
+// Packed atomic-add helpers used by the K-split epilogue.
186
+//
187
+// When the kernel is launched with gridDim.z > 1, multiple K-segments
188
+// accumulate into the same 16x16 output tile and need atomic write-back.
189
+// gfx11 has no native v_global_atomic_pk_add_{f16,bf16}, so we issue a
190
+// CAS-loop on a 32-bit word covering 2 packed fp16/bf16 lanes. Within a
191
+// block the kernel pairs adjacent lanes via shfl_xor first, so each pair
192
+// of cols (n=lane_lo even, lane_lo+1) goes through a SINGLE atomic — no
193
+// intra-block contention on the same uint32 target. Inter-block
194
+// contention from gridDim.z (4-way at K_SPLIT=4) is the residual cost.
195
+// ---------------------------------------------------------------------------
196
+
197
+__forceinline__ __device__ void atomic_add_pk_f16(half2* addr, half2 val) {
198
+ uint32_t* addr_u = reinterpret_cast<uint32_t*>(addr);
199
+ uint32_t old = *addr_u;
200
+ while (true) {
201
+ half2 cur;
202
+ __builtin_memcpy(&cur, &old, sizeof(cur));
203
+ half2 sum = __hadd2(cur, val);
204
+ uint32_t sum_u;
205
+ __builtin_memcpy(&sum_u, &sum, sizeof(sum_u));
206
+ uint32_t prev = atomicCAS(addr_u, old, sum_u);
207
+ if (prev == old) break;
208
+ old = prev;
209
+ }
210
+}
211
+
212
+__forceinline__ __device__ void atomic_add_pk_bf16(bf162_t* addr, bf162_t val) {
213
+ uint32_t* addr_u = reinterpret_cast<uint32_t*>(addr);
214
+ uint32_t old = *addr_u;
215
+ while (true) {
216
+ bf162_t cur;
217
+ __builtin_memcpy(&cur, &old, sizeof(cur));
218
+ bf162_t sum = __hadd2(cur, val);
219
+ uint32_t sum_u;
220
+ __builtin_memcpy(&sum_u, &sum, sizeof(sum_u));
221
+ uint32_t prev = atomicCAS(addr_u, old, sum_u);
222
+ if (prev == old) break;
223
+ old = prev;
224
+ }
225
+}
226
+
227
+#endif // helpers guard; K-split heuristics below are pure host/device
228
+ // arithmetic, called from launch_* on non-RDNA3 device passes too.
229
+
230
+// K-split factor heuristic. Returns the gridDim.z to use for a given K.
231
+// Aim: each block does at least ~16 K-tiles (= K=256) so the per-block
232
+// constant overhead (LDS init, kernel prologue) is amortised. Upper
233
+// bound K_SPLIT=4 to cap inter-block atomic contention to 4-way.
234
+//
235
+// For typical Qwen-class shapes K ∈ {4096, 5120, 11008}, all return 4.
236
+// Smaller K (e.g., embedding lookups) fall back to 1 (no split, no
237
+// atomic). K must be divisible by (K_SPLIT × 16) for the split to be
238
+// valid; the heuristic checks divisibility before raising the factor.
239
+__host__ __device__ static inline int compute_wmma_k_split(int size_k) {
240
+ if (size_k >= 1024 && size_k % 64 == 0) return 4;
241
+ if (size_k >= 512 && size_k % 32 == 0) return 2;
242
+ return 1;
243
+}
244
+
245
+// M-and-N-aware K-split heuristic for the v3/v4/v5 launchers.
246
+//
247
+// The original `compute_wmma_k_split` was K-only and always returns 4 for
248
+// Qwen-class K, which over-subscribes wave slots and pays the atomic CAS
249
+// epilogue once-per-K-segment per output cell. With v3/v4/v5's larger
250
+// tiles (64M × 16/32/64N) and 4 resident waves per block, the no-split
251
+// grid is often already well-saturated on gfx1100's 96 CUs / 3072 wave
252
+// slots — adding gridDim.z just adds atomic overhead.
253
+//
254
+// Heuristic: compute the no-split block count gridDim.x × gridDim.y, then
255
+// pick the smallest K_SPLIT that brings total waves to at least
256
+// ~2× over-subscription (~6000 waves for our 3072 slots, i.e. 1500 blocks
257
+// at 4 waves/block). Above that threshold, K_SPLIT=1 — direct write, no
258
+// atomic.
259
+//
260
+// Args:
261
+// size_m, size_n, size_k — GEMM dims
262
+// m_tile, n_tile — block-level M and N tile (64×16 for v3,
263
+// 64×32 for v4, 64×64 for v5)
264
+//
265
+// Returns: gridDim.z divisor (1, 2, or 4), respecting K-divisibility.
266
+__host__ __device__ static inline int compute_wmma_k_split_mn(
267
+ int size_m, int size_n, int size_k, int m_tile, int n_tile) {
268
+ const int blocks_xy =
269
+ ((size_n + n_tile - 1) / n_tile) * ((size_m + m_tile - 1) / m_tile);
270
+ // Target: enough blocks to keep ~2× oversubscription on 96 CUs / 3072
271
+ // wave slots at 4 waves/block ⇒ ~1500 blocks no-split.
272
+ constexpr int kTargetBlocksXY = 1500;
273
+ if (blocks_xy >= kTargetBlocksXY) return 1;
274
+ if (blocks_xy * 2 >= kTargetBlocksXY && size_k >= 512 && size_k % 32 == 0)
275
+ return 2;
276
+ if (blocks_xy * 4 >= kTargetBlocksXY && size_k >= 1024 && size_k % 64 == 0)
277
+ return 4;
278
+ // Fall back to the K-only heuristic when blocks are very few (small
279
+ // models with small N): K-split is the only way to add parallelism.
280
+ return compute_wmma_k_split(size_k);
281
+}
282
+
283
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
284
+
285
+// Native AMDGPU vector types expected by the WMMA built-ins.
286
+using v16fp16 = _Float16 __attribute__((ext_vector_type(16)));
287
+using v16bf16 = __bf16 __attribute__((ext_vector_type(16)));
288
+using v8fp32 = float __attribute__((ext_vector_type(8)));
289
+
290
+__device__ __forceinline__ v8fp32 wmma_mma(v16fp16 a, v16fp16 b, v8fp32 c) {
291
+ return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
292
+}
293
+__device__ __forceinline__ v8fp32 wmma_mma(v16bf16 a, v16bf16 b, v8fp32 c) {
294
+ return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a, b, c);
295
+}
296
+
297
+// Map HIP wrapper types (half, __hip_bfloat16) to native compiler types
298
+// (_Float16, __bf16) used by the WMMA built-ins. Bitcast is a register
299
+// reinterpret in practice.
300
+template <typename T>
301
+struct WmmaNative;
302
+template <>
303
+struct WmmaNative<half> {
304
+ using elem = _Float16;
305
+ using v16 = v16fp16;
306
+};
307
+template <>
308
+struct WmmaNative<bf16_t> {
309
+ using elem = __bf16;
310
+ using v16 = v16bf16;
311
+};
312
+
313
+template <typename FROM, typename TO>
314
+__device__ __forceinline__ TO bitcast_elem(FROM x) {
315
+ static_assert(sizeof(FROM) == sizeof(TO),
316
+ "bitcast_elem requires equal-sized types");
317
+ TO r;
318
+ __builtin_memcpy(&r, &x, sizeof(TO));
319
+ return r;
320
+}
321
+
322
+// Per-T tzero (matches the helper in the scalar TU).
323
+template <typename T>
324
+__device__ __forceinline__ T tzero();
325
+template <>
326
+__device__ __forceinline__ half tzero<half>() {
327
+ return __float2half_rn(0.0f);
328
+}
329
+template <>
330
+__device__ __forceinline__ bf16_t tzero<bf16_t>() {
331
+ return __float2bfloat16(0.0f);
332
+}
333
+
334
+#endif // helpers guard (each __global__ below has its own guard so launch_*
335
+ // host code remains visible to the parser on non-RDNA3 device passes)
336
+
337
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
338
+
339
+// ===========================================================================
340
+// WMMA kernel: 16M × 16N tile per block, 1 wave, full K traversal.
341
+//
342
+// Wave32 fragment layout (verified empirically with all-modes diagnostic
343
+// against a random A and B and the eight candidate output mappings):
344
+//
345
+// * A frag (row-major in M, K in slot):
346
+// lane t, slot i → A[lane_lo][k = i]
347
+// Lane axis encodes M (A's row), slot encodes K.
348
+// * B frag (col-major in N, K in slot):
349
+// lane t, slot i → B[k = i][lane_lo]
350
+// Lane axis encodes N (B's column), slot encodes K. K-axis aligns
351
+// with A's K-axis (same slot index).
352
+// * C frag (output, lane = N, slot = M with hi-bit interleave):
353
+// lane t, slot i → C[m = 2*i + lane_hi][n = lane_lo]
354
+// Lane axis encodes N (C's column). Each lane holds 8 elements of
355
+// its output column, alternating rows: lanes 0..15 (hi=0) hold even
356
+// rows m=0,2,4,...,14; lanes 16..31 (hi=1) hold odd rows
357
+// m=1,3,5,...,15.
358
+// * Both halves of the wave (lanes 0..15 and 16..31) hold IDENTICAL input
359
+// fragments (AMD's "doubled" wave32 input layout). Output is split
360
+// between halves via lane_hi.
361
+//
362
+// History note: an earlier version of this kernel loaded B row-major in K
363
+// and assumed the output was C[lane_lo][2*i+lane_hi]. That layout passes
364
+// all-A=identity tests because A=I makes the K-axis sum collapse, but
365
+// implements C = A @ B^T for non-trivial A — the bug only shows up against
366
+// random A. A layout probe iterating all four
367
+// {row,col} × {row,col} loadings identified mode 1 (A row, B col) with
368
+// output [m=2*i+hi][n=lane_lo] as the unique mapping that yields A @ B.
369
+// ===========================================================================
370
+
371
+template <typename T>
372
+__global__ void gemm_q4_wmma_kernel_16x16_1w(
373
+ const T* __restrict__ a, const uint32_t* __restrict__ b_q,
374
+ const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
375
+ T* __restrict__ c, const int size_m, const int size_n, const int size_k,
376
+ const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
377
+ using E = typename WmmaNative<T>::elem;
378
+ using V16 = typename WmmaNative<T>::v16;
379
+
380
+ const int m_tile = blockIdx.y * 16;
381
+ const int n_tile = blockIdx.x * 16;
382
+ if (m_tile >= size_m || n_tile >= size_n) return;
383
+
384
+ const int lane = threadIdx.x; // 0..31
385
+ const int lane_lo = lane & 15; // row index within fragment
386
+ const int lane_hi = lane >> 4; // 0 or 1
387
+
388
+ v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0};
389
+
390
+ const int groupsize = size_k / groups;
391
+
392
+ // K-split: each block in the gridDim.z dimension processes a contiguous
393
+ // K-segment [k_start, k_end). With gridDim.z > 1, multiple blocks
394
+ // accumulate into the same output tile and need atomic write-back at
395
+ // the end. With gridDim.z == 1 the kernel falls back to the original
396
+ // behaviour: full K range, single writer per cell, direct write.
397
+ //
398
+ // The split multiplies the wave count by gridDim.z and proportionally
399
+ // raises CU saturation — this is the dominant lever for closing the
400
+ // throughput gap to the fp16 scalar kernel at M >= 16, which already
401
+ // uses K-split natively (gridDim.z = K/256). At gridDim.z = 4 with
402
+ // K=4096, WMMA jumps from 17% to ~67% wave-slot saturation.
403
+ const int k_per_split = size_k / gridDim.z;
404
+ const int k_start = blockIdx.z * k_per_split;
405
+ const int k_end = k_start + k_per_split;
406
+
407
+ // LDS tile of dequantized B. 16 K rows × 16 N cols.
408
+ __shared__ T b_lds[16][16];
409
+
410
+ for (int k_tile = k_start; k_tile < k_end; k_tile += 16) {
411
+ // ---- Dequant 16x16 B tile into LDS ----
412
+ // 32 lanes split 16 N cols × 2 K-octets per col = 32 dequant tasks.
413
+ const int my_n = lane_lo;
414
+ const int my_k_octet = lane_hi; // 0 → K[0..7], 1 → K[8..15]
415
+ const int actual_n = n_tile + my_n;
416
+
417
+ if (actual_n < size_n) {
418
+ const int qk_row = (k_tile / 8) + my_k_octet;
419
+ const uint32_t qa = b_q[qk_row * size_n + actual_n];
420
+
421
+ const int g = k_tile / groupsize;
422
+ const int qz_idx = g * (size_n / 8) + actual_n / 8;
423
+ const int qz_shift = (actual_n & 7) * 4;
424
+ const uint32_t zero_v =
425
+ ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset;
426
+ const T scale_t = b_scales[g * size_n + actual_n];
427
+
428
+ const int k_base = my_k_octet * 8;
429
+
430
+ if constexpr (std::is_same<T, half>::value) {
431
+ half2 z_prep, y_prep;
432
+ prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep);
433
+ half2 dq[4];
434
+ dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep);
435
+ b_lds[k_base + 0][my_n] = __low2half(dq[0]);
436
+ b_lds[k_base + 1][my_n] = __high2half(dq[0]);
437
+ b_lds[k_base + 2][my_n] = __low2half(dq[1]);
438
+ b_lds[k_base + 3][my_n] = __high2half(dq[1]);
439
+ b_lds[k_base + 4][my_n] = __low2half(dq[2]);
440
+ b_lds[k_base + 5][my_n] = __high2half(dq[2]);
441
+ b_lds[k_base + 6][my_n] = __low2half(dq[3]);
442
+ b_lds[k_base + 7][my_n] = __high2half(dq[3]);
443
+ } else {
444
+ float z_f, y_f;
445
+ prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f);
446
+ bf162_t dq[4];
447
+ dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f);
448
+ b_lds[k_base + 0][my_n] = dq[0].x;
449
+ b_lds[k_base + 1][my_n] = dq[0].y;
450
+ b_lds[k_base + 2][my_n] = dq[1].x;
451
+ b_lds[k_base + 3][my_n] = dq[1].y;
452
+ b_lds[k_base + 4][my_n] = dq[2].x;
453
+ b_lds[k_base + 5][my_n] = dq[2].y;
454
+ b_lds[k_base + 6][my_n] = dq[3].x;
455
+ b_lds[k_base + 7][my_n] = dq[3].y;
456
+ }
457
+ }
458
+
459
+ // No __syncthreads() needed: the launch is `dim3 block(32)` = exactly one
460
+ // wave32, so there is no inter-wave concurrency. Within a wave the
461
+ // compiler emits `s_waitcnt lgkmcnt(0)` between dependent ds_write/ds_read
462
+ // pairs automatically, so cross-lane LDS reads (lane 0 reading what
463
+ // lane 16 wrote into b_lds[8..15][0]) still observe the writes. Keeping
464
+ // the explicit `__syncthreads()` would emit a wave-level `s_barrier` that
465
+ // costs ~10-20 cycles every iteration but provides no semantic guarantee
466
+ // we don't already have for free in single-wave mode.
467
+
468
+ // ---- Build A and B fragments, run WMMA ----
469
+ V16 a_frag, b_frag;
470
+ const int m_row = m_tile + lane_lo;
471
+
472
+ if (m_row < size_m) {
473
+ const T* a_row = a + m_row * size_k;
474
+ if (b_q_perm) {
475
+ // Permuted (act-order): scattered global reads, no vectorization.
476
+ #pragma unroll
477
+ for (int i = 0; i < 16; i++) {
478
+ T v = a_row[b_q_perm[k_tile + i]];
479
+ a_frag[i] = bitcast_elem<T, E>(v);
480
+ }
481
+ } else {
482
+ // Sequential A reads: replace 16 single-element global_load_b16 with
483
+ // a bulk 32-byte copy. The AMDGPU backend lowers a memcpy of this
484
+ // size + alignment to two `global_load_b128` instructions. size_k is
485
+ // a multiple of 16 (TORCH_CHECK above) and k_tile increments by 16,
486
+ // so k_tile + 16 is always within bounds — no tail handling needed.
487
+ // Note: we memcpy into the whole vector (`&a_frag`) rather than
488
+ // `&a_frag[0]`; ext_vector_type element addresses aren't reliably
489
+ // valid C pointers across compiler versions.
490
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes (16 × 2)");
491
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
492
+ }
493
+ } else {
494
+ #pragma unroll
495
+ for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
496
+ }
497
+
498
+ // B fragment: lane t holds COLUMN n=lane_lo of the B tile (K-axis in
499
+ // slot, N-axis in lane). This is the AMD WMMA convention for the right
500
+ // operand of a matrix multiply — K-axis aligns with A's K-axis (also
501
+ // in slot), enabling per-lane inner products.
502
+ #pragma unroll
503
+ for (int i = 0; i < 16; i++) {
504
+ b_frag[i] = bitcast_elem<T, E>(b_lds[i][lane_lo]);
505
+ }
506
+
507
+ #ifdef VLLM_WMMA_LAYOUT_DEBUG
508
+ // Diagnostic: skip WMMA, force c_acc to encode (lane, slot) so the
509
+ // store pattern reveals the C-output lane→matrix mapping. Compile with
510
+ // -DVLLM_WMMA_LAYOUT_DEBUG to enable. Output: c[m][n] = lane + slot/16.
511
+ (void)a_frag;
512
+ (void)b_frag;
513
+ #pragma unroll
514
+ for (int i = 0; i < 8; i++) {
515
+ c_acc[i] = (float)lane + (float)i / 16.0f;
516
+ }
517
+ // Run only one K iteration in debug mode so c_acc isn't overwritten.
518
+ if (k_tile == 0) {
519
+ k_tile = size_k; // exit loop on next check
520
+ }
521
+ #else
522
+ c_acc = wmma_mma(a_frag, b_frag, c_acc);
523
+ #endif
524
+
525
+ // No __syncthreads() needed before the next iter overwrites b_lds:
526
+ // single-wave block, and the next iter's ds_write to b_lds is preceded
527
+ // by a `s_waitcnt lgkmcnt(0)` from the compiler that ensures the WMMA's
528
+ // ds_read of b_frag has completed before the new ds_write issues.
529
+ }
530
+
531
+ // ---- Store C ----
532
+ // Lane t holds column n=lane_lo of the output tile, 8 rows determined by
533
+ // slot i and lane_hi:
534
+ // lane_hi == 0 → rows 0, 2, 4, ..., 14 (even rows)
535
+ // lane_hi == 1 → rows 1, 3, 5, ..., 15 (odd rows)
536
+ // c_acc[i] corresponds to actual row m = 2*i + lane_hi at column lane_lo.
537
+ if (gridDim.z > 1) {
538
+ // K-split path: 4 (or whatever the split factor is) K-segments per
539
+ // output cell contend → atomic accumulation. Caller has zero-init'd c.
540
+ //
541
+ // Pair-shuffle to avoid intra-block CAS contention: lanes lane_lo and
542
+ // lane_lo+1 share the same uint32 atomic target (4 bytes = 2 fp16),
543
+ // so without pairing they'd hammer the same word. Instead, swap the
544
+ // c_acc[i] value with the lane_lo+1 neighbour via shfl_xor and have
545
+ // ONLY the even lane issue a single packed CAS. Inter-block contention
546
+ // (gridDim.z-way per cell) remains and is the residual atomic cost.
547
+ const bool is_even_lane = (lane_lo & 1) == 0;
548
+ const int out_n_pair = n_tile + lane_lo; // valid only on even lane
549
+ #pragma unroll
550
+ for (int i = 0; i < 8; i++) {
551
+ // Wave-wide shuffle: every lane participates so the side-effect is
552
+ // visible. Only even lanes use the result. shfl_xor with mask 1
553
+ // swaps with the lane_lo XOR 1 neighbour (same lane_hi → same row).
554
+ float other_f = __shfl_xor(c_acc[i], 1);
555
+ if (!is_even_lane) continue;
556
+
557
+ const int out_m = m_tile + 2 * i + lane_hi;
558
+ if (out_m >= size_m || out_n_pair >= size_n) continue;
559
+
560
+ T* dst = c + out_m * size_n + out_n_pair;
561
+ if constexpr (std::is_same<T, half>::value) {
562
+ // Pack: .x = mine (col=lane_lo even), .y = neighbour (col=lane_lo+1)
563
+ half2 packed =
564
+ __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f));
565
+ atomic_add_pk_f16(reinterpret_cast<half2*>(dst), packed);
566
+ } else {
567
+ bf162_t packed;
568
+ packed.x = __float2bfloat16(c_acc[i]);
569
+ packed.y = __float2bfloat16(other_f);
570
+ atomic_add_pk_bf16(reinterpret_cast<bf162_t*>(dst), packed);
571
+ }
572
+ }
573
+ } else {
574
+ // gridDim.z == 1: single writer per cell, direct non-atomic write.
575
+ // Caller can leave c uninitialised (torch::empty) since every cell is
576
+ // assigned exactly once.
577
+ const int out_n = n_tile + lane_lo;
578
+ if (out_n < size_n) {
579
+ #pragma unroll
580
+ for (int i = 0; i < 8; i++) {
581
+ const int out_m = m_tile + 2 * i + lane_hi;
582
+ if (out_m < size_m) {
583
+ T* dst = c + out_m * size_n + out_n;
584
+ if constexpr (std::is_same<T, half>::value) {
585
+ *dst = __float2half_rn(c_acc[i]);
586
+ } else {
587
+ *dst = __float2bfloat16(c_acc[i]);
588
+ }
589
+ }
590
+ }
591
+ }
592
+ }
593
+}
594
+
595
+#else // non-RDNA3 device pass: empty kernel for symbol parity.
596
+template <typename T>
597
+__global__ void gemm_q4_wmma_kernel_16x16_1w(const T*, const uint32_t*,
598
+ const uint32_t*, const T*, T*,
599
+ const int, const int, const int,
600
+ const int, const int, const int*) {
601
+}
602
+#endif
603
+
604
+template <typename T>
605
+void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight,
606
+ const uint32_t* b_qzeros, const T* b_scales,
607
+ const int* b_q_perm, T* c, int size_m,
608
+ int size_n, int size_k, int groups,
609
+ int zero_offset, cudaStream_t stream) {
610
+ // 1 wave per block (32 lanes), 16x16 C tile per block. gridDim.z splits
611
+ // K so that more blocks (and therefore more waves) are in flight; with
612
+ // K_SPLIT > 1 the kernel switches to atomic write-back at the epilogue.
613
+ const int k_split = compute_wmma_k_split(size_k);
614
+ dim3 block(32);
615
+ dim3 grid((size_n + 15) / 16, (size_m + 15) / 16, k_split);
616
+ gemm_q4_wmma_kernel_16x16_1w<T><<<grid, block, 0, stream>>>(
617
+ a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
618
+ zero_offset, b_q_perm);
619
+}
620
+
621
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
622
+
623
+// ===========================================================================
624
+// 32x16_2w kernel: 2 waves per block, 32M × 16N tile, double-buffered LDS.
625
+//
626
+// Targets the bf16-WMMA prefill regime (M >= 128) where the v1 single-wave
627
+// kernel saturates only ~24% of WMMA peak because each wave does roughly
628
+// one v_wmma every ~30-40 cycles (16-cycle wmma latency + dequant + LDS +
629
+// global A load all serial inside the single resident wave).
630
+//
631
+// Two structural changes vs v1:
632
+//
633
+// * 2 waves per block (64 threads). Both waves cooperate on a 32M×16N
634
+// output tile: wave 0 produces rows [0..15], wave 1 rows [16..31].
635
+// The B-tile in LDS is shared (only wave 0 dequants); each wave loads
636
+// its own A slice from global. With two resident waves the SIMD
637
+// scheduler can keep the WMMA pipeline full by interleaving wmmas
638
+// from the two waves while the other does dequant / LDS / load work.
639
+//
640
+// * Double-buffered LDS B-tile (b_lds[2][16][16]). Wave 0 dequants
641
+// the K-tile for iter k+1 while both waves consume the K-tile for
642
+// iter k. Pulls dequant out of the WMMA-critical path. Costs ~512 B
643
+// extra LDS per block (irrelevant — gfx1100 has 64 KB LDS/CU).
644
+//
645
+// One __syncthreads() per K-iter remains: it ensures wave 0's dequant of
646
+// the next K-tile has committed AND both waves have finished reading the
647
+// current K-tile before wave 0 wraps around and overwrites it.
648
+//
649
+// The v2 launcher (`launch_gemm_q4_wmma_32x16_2w`) is the production entry on
650
+// the WMMA path and falls back to v1 internally for size_m < 32 — see the
651
+// comment at the top of `launch_gemm_q4_wmma_32x16_2w` for the M=16 regression
652
+// rationale that justifies the fallback.
653
+// ===========================================================================
654
+
655
+template <typename T>
656
+__global__ void gemm_q4_wmma_kernel_32x16_2w(
657
+ const T* __restrict__ a, const uint32_t* __restrict__ b_q,
658
+ const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
659
+ T* __restrict__ c, const int size_m, const int size_n, const int size_k,
660
+ const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
661
+ using E = typename WmmaNative<T>::elem;
662
+ using V16 = typename WmmaNative<T>::v16;
663
+
664
+ const int m_tile = blockIdx.y * 32; // 32-row stride per block
665
+ const int n_tile = blockIdx.x * 16;
666
+ if (m_tile >= size_m || n_tile >= size_n) return;
667
+
668
+ const int tid = threadIdx.x; // 0..63
669
+ const int wave_id = tid >> 5; // 0 or 1
670
+ const int lane = tid & 31;
671
+ const int lane_lo = lane & 15;
672
+ const int lane_hi = lane >> 4;
673
+
674
+ v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0};
675
+
676
+ const int groupsize = size_k / groups;
677
+
678
+ // K-split: each block handles a contiguous K-segment when gridDim.z > 1.
679
+ const int k_per_split = size_k / gridDim.z;
680
+ const int k_start = blockIdx.z * k_per_split;
681
+ const int k_end = k_start + k_per_split;
682
+
683
+ // Double-buffered LDS B-tile. 2 × 16K × 16N × sizeof(T) = 1024 B for
684
+ // fp16/bf16.
685
+ __shared__ T b_lds[2][16][16];
686
+
687
+ // Dequant a 16K × 16N B-tile into b_lds[buf]. Only wave 0 participates
688
+ // (32 lanes do 32 dequant tasks: 16 N-cols × 2 K-octets per col).
689
+ auto dequant_into = [&](int buf, int k_tile) {
690
+ if (wave_id != 0) return;
691
+
692
+ const int my_n = lane_lo;
693
+ const int my_k_octet = lane_hi;
694
+ const int actual_n = n_tile + my_n;
695
+
696
+ if (actual_n >= size_n) return;
697
+
698
+ const int qk_row = (k_tile / 8) + my_k_octet;
699
+ const uint32_t qa = b_q[qk_row * size_n + actual_n];
700
+
701
+ const int g = k_tile / groupsize;
702
+ const int qz_idx = g * (size_n / 8) + actual_n / 8;
703
+ const int qz_shift = (actual_n & 7) * 4;
704
+ const uint32_t zero_v =
705
+ ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset;
706
+ const T scale_t = b_scales[g * size_n + actual_n];
707
+
708
+ const int k_base = my_k_octet * 8;
709
+
710
+ if constexpr (std::is_same<T, half>::value) {
711
+ half2 z_prep, y_prep;
712
+ prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep);
713
+ half2 dq[4];
714
+ dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep);
715
+ b_lds[buf][k_base + 0][my_n] = __low2half(dq[0]);
716
+ b_lds[buf][k_base + 1][my_n] = __high2half(dq[0]);
717
+ b_lds[buf][k_base + 2][my_n] = __low2half(dq[1]);
718
+ b_lds[buf][k_base + 3][my_n] = __high2half(dq[1]);
719
+ b_lds[buf][k_base + 4][my_n] = __low2half(dq[2]);
720
+ b_lds[buf][k_base + 5][my_n] = __high2half(dq[2]);
721
+ b_lds[buf][k_base + 6][my_n] = __low2half(dq[3]);
722
+ b_lds[buf][k_base + 7][my_n] = __high2half(dq[3]);
723
+ } else {
724
+ float z_f, y_f;
725
+ prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f);
726
+ bf162_t dq[4];
727
+ dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f);
728
+ b_lds[buf][k_base + 0][my_n] = dq[0].x;
729
+ b_lds[buf][k_base + 1][my_n] = dq[0].y;
730
+ b_lds[buf][k_base + 2][my_n] = dq[1].x;
731
+ b_lds[buf][k_base + 3][my_n] = dq[1].y;
732
+ b_lds[buf][k_base + 4][my_n] = dq[2].x;
733
+ b_lds[buf][k_base + 5][my_n] = dq[2].y;
734
+ b_lds[buf][k_base + 6][my_n] = dq[3].x;
735
+ b_lds[buf][k_base + 7][my_n] = dq[3].y;
736
+ }
737
+ };
738
+
739
+ // Pre-fill buffer 0 with the first K-tile so iter 0 has data to consume.
740
+ dequant_into(0, k_start);
741
+ __syncthreads();
742
+
743
+ int cur_buf = 0;
744
+ for (int k_tile = k_start; k_tile < k_end; k_tile += 16) {
745
+ const int next_buf = 1 - cur_buf;
746
+ const int k_next = k_tile + 16;
747
+
748
+ // Issue dequant of next K-tile (wave 0 only) — overlaps with current
749
+ // iter's WMMA work below. The LDS write is non-blocking; the sync at
750
+ // end of iter ensures wave 1 sees it before iter k+1.
751
+ if (k_next < k_end) {
752
+ dequant_into(next_buf, k_next);
753
+ }
754
+
755
+ // Load A: each wave loads its own M slice in parallel.
756
+ const int m_row = m_tile + wave_id * 16 + lane_lo;
757
+ V16 a_frag, b_frag;
758
+ if (m_row < size_m) {
759
+ const T* a_row = a + m_row * size_k;
760
+ if (b_q_perm) {
761
+ #pragma unroll
762
+ for (int i = 0; i < 16; i++) {
763
+ T v = a_row[b_q_perm[k_tile + i]];
764
+ a_frag[i] = bitcast_elem<T, E>(v);
765
+ }
766
+ } else {
767
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
768
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
769
+ }
770
+ } else {
771
+ #pragma unroll
772
+ for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
773
+ }
774
+
775
+ // Load B from current buffer (both waves read identical data).
776
+ #pragma unroll
777
+ for (int i = 0; i < 16; i++) {
778
+ b_frag[i] = bitcast_elem<T, E>(b_lds[cur_buf][i][lane_lo]);
779
+ }
780
+
781
+ // Each wave issues its own WMMA against its own a_frag + shared b_frag.
782
+ c_acc = wmma_mma(a_frag, b_frag, c_acc);
783
+
784
+ // Sync ensures: (a) wave 0's dequant into next_buf has committed,
785
+ // (b) both waves are done reading cur_buf — so next iter's overwrite
786
+ // (cur_buf becomes the previous next_buf and gets reused two iters
787
+ // later) is race-free.
788
+ __syncthreads();
789
+ cur_buf = next_buf;
790
+ }
791
+
792
+ // ---- Store C ----
793
+ // Each wave owns rows [m_tile + wave_id*16 .. + 16) of the output tile.
794
+ const int m_tile_wave = m_tile + wave_id * 16;
795
+
796
+ if (gridDim.z > 1) {
797
+ // K-split atomic path. Pair-shuffle within wave to halve atomic count.
798
+ // shfl_xor here is wave-local (wave32 semantics) so each wave does its
799
+ // own pairing — the two waves don't interact during the store.
800
+ const bool is_even_lane = (lane_lo & 1) == 0;
801
+ const int out_n_pair = n_tile + lane_lo;
802
+ #pragma unroll
803
+ for (int i = 0; i < 8; i++) {
804
+ float other_f = __shfl_xor(c_acc[i], 1);
805
+ if (!is_even_lane) continue;
806
+
807
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
808
+ if (out_m >= size_m || out_n_pair >= size_n) continue;
809
+
810
+ T* dst = c + out_m * size_n + out_n_pair;
811
+ if constexpr (std::is_same<T, half>::value) {
812
+ half2 packed =
813
+ __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f));
814
+ atomic_add_pk_f16(reinterpret_cast<half2*>(dst), packed);
815
+ } else {
816
+ bf162_t packed;
817
+ packed.x = __float2bfloat16(c_acc[i]);
818
+ packed.y = __float2bfloat16(other_f);
819
+ atomic_add_pk_bf16(reinterpret_cast<bf162_t*>(dst), packed);
820
+ }
821
+ }
822
+ } else {
823
+ // Single writer per cell, direct non-atomic write.
824
+ const int out_n = n_tile + lane_lo;
825
+ if (out_n < size_n) {
826
+ #pragma unroll
827
+ for (int i = 0; i < 8; i++) {
828
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
829
+ if (out_m < size_m) {
830
+ T* dst = c + out_m * size_n + out_n;
831
+ if constexpr (std::is_same<T, half>::value) {
832
+ *dst = __float2half_rn(c_acc[i]);
833
+ } else {
834
+ *dst = __float2bfloat16(c_acc[i]);
835
+ }
836
+ }
837
+ }
838
+ }
839
+ }
840
+}
841
+
842
+#else // non-RDNA3 device pass: empty kernel for symbol parity.
843
+template <typename T>
844
+__global__ void gemm_q4_wmma_kernel_32x16_2w(const T*, const uint32_t*,
845
+ const uint32_t*, const T*, T*,
846
+ const int, const int, const int,
847
+ const int, const int, const int*) {
848
+}
849
+#endif
850
+
851
+template <typename T>
852
+void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight,
853
+ const uint32_t* b_qzeros, const T* b_scales,
854
+ const int* b_q_perm, T* c, int size_m,
855
+ int size_n, int size_k, int groups,
856
+ int zero_offset, cudaStream_t stream) {
857
+ // Fallback to v1 for size_m < 32. With M-tile=32 the v2 block has 2 waves
858
+ // working on rows [0..15] and [16..31]; at M < 32 the second wave processes
859
+ // out-of-range M rows (zero-padded a_frag → wmma produces nothing useful)
860
+ // and just wastes SIMD cycles. Bench measured a +47 % regression at M=16
861
+ // vs v1 for this reason. The M < 32 case is rare in serving (decode at
862
+ // max-num-seqs=32 lands at M≈32 steady-state; the M=16 sliver is edge),
863
+ // but the fallback costs nothing and is the right shape.
864
+ if (size_m < 32) {
865
+ launch_gemm_q4_wmma_16x16_1w<T>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
866
+ c, size_m, size_n, size_k, groups,
867
+ zero_offset, stream);
868
+ return;
869
+ }
870
+
871
+ // 2 waves per block (64 threads), 32M × 16N C tile per block.
872
+ // K-split heuristic shared with v1. With M-tile=32, the natural
873
+ // grid blocks are halved on Y vs v1 — but each block does 2× the work,
874
+ // so total wave count is unchanged at the same M when K_SPLIT is equal.
875
+ const int k_split = compute_wmma_k_split(size_k);
876
+ dim3 block(64);
877
+ dim3 grid((size_n + 15) / 16, (size_m + 31) / 32, k_split);
878
+ gemm_q4_wmma_kernel_32x16_2w<T><<<grid, block, 0, stream>>>(
879
+ a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
880
+ zero_offset, b_q_perm);
881
+}
882
+
883
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
884
+
885
+// ===========================================================================
886
+// 64x16_4w kernel: 4 waves per block, 64M × 16N tile, double-buffered LDS.
887
+//
888
+// Targets the prefill plateau observed at M >= 128 in v2 (~144 K tk/s bf16,
889
+// ~28 % of WMMA peak). The bottleneck is wmma issue rate per resident wave:
890
+// each wave issues at most 1 wmma per ~30-40 cycles. With only 2 waves per
891
+// block, two wmmas overlap; the wmma pipeline (16-cycle latency) is mostly
892
+// idle.
893
+//
894
+// Doubling the resident wave count (2 → 4) targets ~2× wmma throughput by
895
+// keeping the pipeline closer to full. 64M tile keeps N-tile at 16 (so the
896
+// b_lds layout, dequant pattern, and store mapping carry over from v2) — the
897
+// only structural changes are:
898
+//
899
+// * 4 waves cooperate on the 64M × 16N output tile. Wave w produces rows
900
+// [16w .. 16w+15] of the M tile.
901
+// * Dequant remains on wave 0 only (32 lanes do 32 dequant slots, identical
902
+// to v2). Waves 1-3 idle through the dequant phase but their wmmas can
903
+// issue concurrently with wave 0's dequant of the *next* K-tile thanks
904
+// to the double buffer — net wave occupancy is dominated by the wmma
905
+// phase, not the dequant phase.
906
+// * One __syncthreads() per K-iter still required (same race as v2).
907
+//
908
+// Costs: same LDS as v2 (1024 B for the b_lds double buffer). Block has 128
909
+// threads vs 64 in v2; gfx1100 supports up to 1024 threads/block so this is
910
+// well within budget. Doubles VGPR pressure slightly because the four waves
911
+// each hold their own a_frag + c_acc — but each wave's working set is
912
+// independent so per-thread VGPR is unchanged.
913
+// ===========================================================================
914
+
915
+template <typename T>
916
+__global__ void gemm_q4_wmma_kernel_64x16_4w(
917
+ const T* __restrict__ a, const uint32_t* __restrict__ b_q,
918
+ const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
919
+ T* __restrict__ c, const int size_m, const int size_n, const int size_k,
920
+ const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
921
+ using E = typename WmmaNative<T>::elem;
922
+ using V16 = typename WmmaNative<T>::v16;
923
+
924
+ const int m_tile =
925
+ blockIdx.y * 64; // 64-row stride per block (4 waves × 16M)
926
+ const int n_tile = blockIdx.x * 16;
927
+ if (m_tile >= size_m || n_tile >= size_n) return;
928
+
929
+ const int tid = threadIdx.x; // 0..127
930
+ const int wave_id = tid >> 5; // 0..3
931
+ const int lane = tid & 31;
932
+ const int lane_lo = lane & 15;
933
+ const int lane_hi = lane >> 4;
934
+
935
+ v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0};
936
+
937
+ const int groupsize = size_k / groups;
938
+
939
+ // K-split: each block handles a contiguous K-segment when gridDim.z > 1.
940
+ const int k_per_split = size_k / gridDim.z;
941
+ const int k_start = blockIdx.z * k_per_split;
942
+ const int k_end = k_start + k_per_split;
943
+
944
+ // Double-buffered LDS B-tile. Same layout as v2.
945
+ __shared__ T b_lds[2][16][16];
946
+
947
+ // Dequant a 16K × 16N B-tile into b_lds[buf]. Only wave 0's 32 lanes
948
+ // participate (16 N-cols × 2 K-octets = 32 dequant slots). Waves 1-3
949
+ // skip — their wmma can run concurrently with wave 0's next-iter dequant
950
+ // through the double buffer.
951
+ auto dequant_into = [&](int buf, int k_tile) {
952
+ if (wave_id != 0) return;
953
+
954
+ const int my_n = lane_lo;
955
+ const int my_k_octet = lane_hi;
956
+ const int actual_n = n_tile + my_n;
957
+
958
+ if (actual_n >= size_n) return;
959
+
960
+ const int qk_row = (k_tile / 8) + my_k_octet;
961
+ const uint32_t qa = b_q[qk_row * size_n + actual_n];
962
+
963
+ const int g = k_tile / groupsize;
964
+ const int qz_idx = g * (size_n / 8) + actual_n / 8;
965
+ const int qz_shift = (actual_n & 7) * 4;
966
+ const uint32_t zero_v =
967
+ ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset;
968
+ const T scale_t = b_scales[g * size_n + actual_n];
969
+
970
+ const int k_base = my_k_octet * 8;
971
+
972
+ if constexpr (std::is_same<T, half>::value) {
973
+ half2 z_prep, y_prep;
974
+ prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep);
975
+ half2 dq[4];
976
+ dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep);
977
+ b_lds[buf][k_base + 0][my_n] = __low2half(dq[0]);
978
+ b_lds[buf][k_base + 1][my_n] = __high2half(dq[0]);
979
+ b_lds[buf][k_base + 2][my_n] = __low2half(dq[1]);
980
+ b_lds[buf][k_base + 3][my_n] = __high2half(dq[1]);
981
+ b_lds[buf][k_base + 4][my_n] = __low2half(dq[2]);
982
+ b_lds[buf][k_base + 5][my_n] = __high2half(dq[2]);
983
+ b_lds[buf][k_base + 6][my_n] = __low2half(dq[3]);
984
+ b_lds[buf][k_base + 7][my_n] = __high2half(dq[3]);
985
+ } else {
986
+ float z_f, y_f;
987
+ prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f);
988
+ bf162_t dq[4];
989
+ dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f);
990
+ b_lds[buf][k_base + 0][my_n] = dq[0].x;
991
+ b_lds[buf][k_base + 1][my_n] = dq[0].y;
992
+ b_lds[buf][k_base + 2][my_n] = dq[1].x;
993
+ b_lds[buf][k_base + 3][my_n] = dq[1].y;
994
+ b_lds[buf][k_base + 4][my_n] = dq[2].x;
995
+ b_lds[buf][k_base + 5][my_n] = dq[2].y;
996
+ b_lds[buf][k_base + 6][my_n] = dq[3].x;
997
+ b_lds[buf][k_base + 7][my_n] = dq[3].y;
998
+ }
999
+ };
1000
+
1001
+ // Pre-fill buffer 0 with the first K-tile so iter 0 has data to consume.
1002
+ dequant_into(0, k_start);
1003
+ __syncthreads();
1004
+
1005
+ int cur_buf = 0;
1006
+ for (int k_tile = k_start; k_tile < k_end; k_tile += 16) {
1007
+ const int next_buf = 1 - cur_buf;
1008
+ const int k_next = k_tile + 16;
1009
+
1010
+ // Issue dequant of next K-tile (wave 0 only) — overlaps with the wmma
1011
+ // work below across all 4 waves.
1012
+ if (k_next < k_end) {
1013
+ dequant_into(next_buf, k_next);
1014
+ }
1015
+
1016
+ // Each wave loads its own 16M slice of A (wave w handles M-rows
1017
+ // [m_tile + 16w .. m_tile + 16w + 15]).
1018
+ const int m_row = m_tile + wave_id * 16 + lane_lo;
1019
+ V16 a_frag, b_frag;
1020
+ if (m_row < size_m) {
1021
+ const T* a_row = a + m_row * size_k;
1022
+ if (b_q_perm) {
1023
+ #pragma unroll
1024
+ for (int i = 0; i < 16; i++) {
1025
+ T v = a_row[b_q_perm[k_tile + i]];
1026
+ a_frag[i] = bitcast_elem<T, E>(v);
1027
+ }
1028
+ } else {
1029
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1030
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1031
+ }
1032
+ } else {
1033
+ #pragma unroll
1034
+ for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
1035
+ }
1036
+
1037
+ // Load B from current buffer (all 4 waves read identical data).
1038
+ #pragma unroll
1039
+ for (int i = 0; i < 16; i++) {
1040
+ b_frag[i] = bitcast_elem<T, E>(b_lds[cur_buf][i][lane_lo]);
1041
+ }
1042
+
1043
+ // Each wave issues its own WMMA against its own a_frag + shared b_frag.
1044
+ // 4 wmmas in flight per block per K-iter.
1045
+ c_acc = wmma_mma(a_frag, b_frag, c_acc);
1046
+
1047
+ __syncthreads();
1048
+ cur_buf = next_buf;
1049
+ }
1050
+
1051
+ // ---- Store C ---- Each wave owns 16 M-rows of the output tile.
1052
+ const int m_tile_wave = m_tile + wave_id * 16;
1053
+
1054
+ if (gridDim.z > 1) {
1055
+ // K-split atomic path. Pair-shuffle within wave to halve atomic count.
1056
+ const bool is_even_lane = (lane_lo & 1) == 0;
1057
+ const int out_n_pair = n_tile + lane_lo;
1058
+ #pragma unroll
1059
+ for (int i = 0; i < 8; i++) {
1060
+ float other_f = __shfl_xor(c_acc[i], 1);
1061
+ if (!is_even_lane) continue;
1062
+
1063
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1064
+ if (out_m >= size_m || out_n_pair >= size_n) continue;
1065
+
1066
+ T* dst = c + out_m * size_n + out_n_pair;
1067
+ if constexpr (std::is_same<T, half>::value) {
1068
+ half2 packed =
1069
+ __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f));
1070
+ atomic_add_pk_f16(reinterpret_cast<half2*>(dst), packed);
1071
+ } else {
1072
+ bf162_t packed;
1073
+ packed.x = __float2bfloat16(c_acc[i]);
1074
+ packed.y = __float2bfloat16(other_f);
1075
+ atomic_add_pk_bf16(reinterpret_cast<bf162_t*>(dst), packed);
1076
+ }
1077
+ }
1078
+ } else {
1079
+ // Single writer per cell, direct non-atomic write.
1080
+ const int out_n = n_tile + lane_lo;
1081
+ if (out_n < size_n) {
1082
+ #pragma unroll
1083
+ for (int i = 0; i < 8; i++) {
1084
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1085
+ if (out_m < size_m) {
1086
+ T* dst = c + out_m * size_n + out_n;
1087
+ if constexpr (std::is_same<T, half>::value) {
1088
+ *dst = __float2half_rn(c_acc[i]);
1089
+ } else {
1090
+ *dst = __float2bfloat16(c_acc[i]);
1091
+ }
1092
+ }
1093
+ }
1094
+ }
1095
+ }
1096
+}
1097
+
1098
+#else // non-RDNA3 device pass: empty kernel for symbol parity.
1099
+template <typename T>
1100
+__global__ void gemm_q4_wmma_kernel_64x16_4w(const T*, const uint32_t*,
1101
+ const uint32_t*, const T*, T*,
1102
+ const int, const int, const int,
1103
+ const int, const int, const int*) {
1104
+}
1105
+#endif
1106
+
1107
+template <typename T>
1108
+void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight,
1109
+ const uint32_t* b_qzeros, const T* b_scales,
1110
+ const int* b_q_perm, T* c, int size_m,
1111
+ int size_n, int size_k, int groups,
1112
+ int zero_offset, cudaStream_t stream) {
1113
+ // Fall back to v2 for M < 64 (would waste 1+ waves on out-of-range rows).
1114
+ if (size_m < 64) {
1115
+ launch_gemm_q4_wmma_32x16_2w<T>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
1116
+ c, size_m, size_n, size_k, groups,
1117
+ zero_offset, stream);
1118
+ return;
1119
+ }
1120
+
1121
+ // 4 waves per block (128 threads), 64M × 16N tile per block.
1122
+ const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 16);
1123
+ dim3 block(128);
1124
+ dim3 grid((size_n + 15) / 16, (size_m + 63) / 64, k_split);
1125
+ gemm_q4_wmma_kernel_64x16_4w<T><<<grid, block, 0, stream>>>(
1126
+ a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
1127
+ zero_offset, b_q_perm);
1128
+}
1129
+
1130
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
1131
+
1132
+// ===========================================================================
1133
+// 64x32_4w kernel: 4 waves per block, 64M × 32N tile, double-buffered LDS.
1134
+//
1135
+// Builds on v3 by doubling the N-tile from 16 → 32. Each wave now issues
1136
+// 2 wmmas per K-iter (one for cols 0-15, one for cols 16-31, sharing the
1137
+// same a_frag). With 4 waves × 2 wmmas = 8 wmmas in flight per K-iter,
1138
+// the wmma pipeline gets ~2× more in-flight work than v3 — targeting the
1139
+// remaining wmma issue gap on Qwen-class shapes (gate/up, down) where v3
1140
+// plateaus at ~22 TFLOPS effective.
1141
+//
1142
+// Costs:
1143
+// * LDS B-tile doubles (2 × 16K × 32N × sizeof(T) = 2048 B for fp16/bf16).
1144
+// * Dequant doubles (32 N-cols × 2 K-octets = 64 slots/K-tile). Distributed
1145
+// across waves 0-1: wave 0 dequants n=[0..15], wave 1 dequants n=[16..31].
1146
+// Waves 2-3 idle on dequant but do wmma work.
1147
+// * Per-wave registers: 2 × v8fp32 accumulator (16 fp32 = 32 VGPRs) plus
1148
+// b_frag0 + b_frag1 (32 VGPRs total). Within budget.
1149
+//
1150
+// Mapping invariant (wave-id → output tile slice):
1151
+// * Wave w produces M rows [m_tile + 16w .. m_tile + 16w + 15]
1152
+// * c_acc0 holds N cols [n_tile + 0 .. n_tile + 15]
1153
+// * c_acc1 holds N cols [n_tile + 16 .. n_tile + 31]
1154
+// ===========================================================================
1155
+
1156
+template <typename T>
1157
+__global__ void gemm_q4_wmma_kernel_64x32_4w(
1158
+ const T* __restrict__ a, const uint32_t* __restrict__ b_q,
1159
+ const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
1160
+ T* __restrict__ c, const int size_m, const int size_n, const int size_k,
1161
+ const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
1162
+ using E = typename WmmaNative<T>::elem;
1163
+ using V16 = typename WmmaNative<T>::v16;
1164
+
1165
+ const int m_tile = blockIdx.y * 64;
1166
+ const int n_tile = blockIdx.x * 32; // 32-col stride per block (was 16 in v3)
1167
+ if (m_tile >= size_m || n_tile >= size_n) return;
1168
+
1169
+ const int tid = threadIdx.x; // 0..127
1170
+ const int wave_id = tid >> 5; // 0..3
1171
+ const int lane = tid & 31;
1172
+ const int lane_lo = lane & 15;
1173
+ const int lane_hi = lane >> 4;
1174
+
1175
+ // Two accumulators per wave: c_acc0 covers cols [n_tile..n_tile+15],
1176
+ // c_acc1 covers cols [n_tile+16..n_tile+31].
1177
+ v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0};
1178
+ v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0};
1179
+
1180
+ const int groupsize = size_k / groups;
1181
+
1182
+ const int k_per_split = size_k / gridDim.z;
1183
+ const int k_start = blockIdx.z * k_per_split;
1184
+ const int k_end = k_start + k_per_split;
1185
+
1186
+ // Doubled LDS B-tile: [buf][k][n_in_tile=0..31].
1187
+ __shared__ T b_lds[2][16][32];
1188
+
1189
+ // Dequant: 64 slots per K-tile (32 N-cols × 2 K-octets). Distributed
1190
+ // across waves 0,1: wave w handles n_in_tile in [16w..16w+15], k_oct in
1191
+ // {0,1}. Each dequanting wave has 32 lanes for 32 dequant slots — perfect
1192
+ // mapping, identical layout to v2/v3 dequant per wave. Waves 2,3 stay idle
1193
+ // during dequant; they catch up via the double buffer overlapping the
1194
+ // wmma of iter k with dequant of iter k+1.
1195
+ //
1196
+ // Tried distributing dequant across all 4 waves (16 slots/wave): regressed
1197
+ // 3% on gate/up M=2048 (357K → 347K tk/s). The dequant is not on the
1198
+ // critical path; spreading it just adds LDS bank pressure with no gain.
1199
+ auto dequant_into = [&](int buf, int k_tile) {
1200
+ if (wave_id >= 2) return;
1201
+
1202
+ const int my_n_local = lane_lo; // 0..15 (within wave's N-half)
1203
+ const int my_n_in_tile = wave_id * 16 + my_n_local; // 0..31 (in 32N tile)
1204
+ const int my_k_octet = lane_hi;
1205
+ const int actual_n = n_tile + my_n_in_tile;
1206
+
1207
+ if (actual_n >= size_n) return;
1208
+
1209
+ const int qk_row = (k_tile / 8) + my_k_octet;
1210
+ const uint32_t qa = b_q[qk_row * size_n + actual_n];
1211
+
1212
+ const int g = k_tile / groupsize;
1213
+ const int qz_idx = g * (size_n / 8) + actual_n / 8;
1214
+ const int qz_shift = (actual_n & 7) * 4;
1215
+ const uint32_t zero_v =
1216
+ ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset;
1217
+ const T scale_t = b_scales[g * size_n + actual_n];
1218
+
1219
+ const int k_base = my_k_octet * 8;
1220
+
1221
+ if constexpr (std::is_same<T, half>::value) {
1222
+ half2 z_prep, y_prep;
1223
+ prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep);
1224
+ half2 dq[4];
1225
+ dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep);
1226
+ b_lds[buf][k_base + 0][my_n_in_tile] = __low2half(dq[0]);
1227
+ b_lds[buf][k_base + 1][my_n_in_tile] = __high2half(dq[0]);
1228
+ b_lds[buf][k_base + 2][my_n_in_tile] = __low2half(dq[1]);
1229
+ b_lds[buf][k_base + 3][my_n_in_tile] = __high2half(dq[1]);
1230
+ b_lds[buf][k_base + 4][my_n_in_tile] = __low2half(dq[2]);
1231
+ b_lds[buf][k_base + 5][my_n_in_tile] = __high2half(dq[2]);
1232
+ b_lds[buf][k_base + 6][my_n_in_tile] = __low2half(dq[3]);
1233
+ b_lds[buf][k_base + 7][my_n_in_tile] = __high2half(dq[3]);
1234
+ } else {
1235
+ float z_f, y_f;
1236
+ prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f);
1237
+ bf162_t dq[4];
1238
+ dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f);
1239
+ b_lds[buf][k_base + 0][my_n_in_tile] = dq[0].x;
1240
+ b_lds[buf][k_base + 1][my_n_in_tile] = dq[0].y;
1241
+ b_lds[buf][k_base + 2][my_n_in_tile] = dq[1].x;
1242
+ b_lds[buf][k_base + 3][my_n_in_tile] = dq[1].y;
1243
+ b_lds[buf][k_base + 4][my_n_in_tile] = dq[2].x;
1244
+ b_lds[buf][k_base + 5][my_n_in_tile] = dq[2].y;
1245
+ b_lds[buf][k_base + 6][my_n_in_tile] = dq[3].x;
1246
+ b_lds[buf][k_base + 7][my_n_in_tile] = dq[3].y;
1247
+ }
1248
+ };
1249
+
1250
+ dequant_into(0, k_start);
1251
+ __syncthreads();
1252
+
1253
+ int cur_buf = 0;
1254
+ for (int k_tile = k_start; k_tile < k_end; k_tile += 16) {
1255
+ const int next_buf = 1 - cur_buf;
1256
+ const int k_next = k_tile + 16;
1257
+
1258
+ if (k_next < k_end) {
1259
+ dequant_into(next_buf, k_next);
1260
+ }
1261
+
1262
+ // Load A: each wave loads its 16M slice, shared across both wmmas.
1263
+ const int m_row = m_tile + wave_id * 16 + lane_lo;
1264
+ V16 a_frag, b_frag0, b_frag1;
1265
+ if (m_row < size_m) {
1266
+ const T* a_row = a + m_row * size_k;
1267
+ if (b_q_perm) {
1268
+ #pragma unroll
1269
+ for (int i = 0; i < 16; i++) {
1270
+ T v = a_row[b_q_perm[k_tile + i]];
1271
+ a_frag[i] = bitcast_elem<T, E>(v);
1272
+ }
1273
+ } else {
1274
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1275
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1276
+ }
1277
+ } else {
1278
+ #pragma unroll
1279
+ for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
1280
+ }
1281
+
1282
+ // Load B for cols [0..15] and cols [16..31]. Both halves of the 32N tile.
1283
+ #pragma unroll
1284
+ for (int i = 0; i < 16; i++) {
1285
+ b_frag0[i] = bitcast_elem<T, E>(b_lds[cur_buf][i][lane_lo]);
1286
+ b_frag1[i] = bitcast_elem<T, E>(b_lds[cur_buf][i][lane_lo + 16]);
1287
+ }
1288
+
1289
+ // Two wmmas per wave per K-iter, sharing a_frag.
1290
+ c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0);
1291
+ c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1);
1292
+
1293
+ __syncthreads();
1294
+ cur_buf = next_buf;
1295
+ }
1296
+
1297
+ // ---- Store C ----
1298
+ // Each wave owns 16M rows × 32N cols. c_acc0 → cols [n_tile..n_tile+15],
1299
+ // c_acc1 → cols [n_tile+16..n_tile+31].
1300
+ const int m_tile_wave = m_tile + wave_id * 16;
1301
+
1302
+ // Helper: store one v8fp32 accumulator's 8 outputs (covers 16 N-cols at
1303
+ // n_base via lane_lo + interleaved M rows m_tile_wave + 2i + lane_hi).
1304
+ auto store_acc = [&](const v8fp32& acc, int n_base) {
1305
+ if (gridDim.z > 1) {
1306
+ const bool is_even_lane = (lane_lo & 1) == 0;
1307
+ const int out_n_pair = n_base + lane_lo;
1308
+ #pragma unroll
1309
+ for (int i = 0; i < 8; i++) {
1310
+ float other_f = __shfl_xor(acc[i], 1);
1311
+ if (!is_even_lane) continue;
1312
+
1313
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1314
+ if (out_m >= size_m || out_n_pair >= size_n) continue;
1315
+
1316
+ T* dst = c + out_m * size_n + out_n_pair;
1317
+ if constexpr (std::is_same<T, half>::value) {
1318
+ half2 packed =
1319
+ __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f));
1320
+ atomic_add_pk_f16(reinterpret_cast<half2*>(dst), packed);
1321
+ } else {
1322
+ bf162_t packed;
1323
+ packed.x = __float2bfloat16(acc[i]);
1324
+ packed.y = __float2bfloat16(other_f);
1325
+ atomic_add_pk_bf16(reinterpret_cast<bf162_t*>(dst), packed);
1326
+ }
1327
+ }
1328
+ } else {
1329
+ const int out_n = n_base + lane_lo;
1330
+ if (out_n >= size_n) return;
1331
+ #pragma unroll
1332
+ for (int i = 0; i < 8; i++) {
1333
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1334
+ if (out_m < size_m) {
1335
+ T* dst = c + out_m * size_n + out_n;
1336
+ if constexpr (std::is_same<T, half>::value) {
1337
+ *dst = __float2half_rn(acc[i]);
1338
+ } else {
1339
+ *dst = __float2bfloat16(acc[i]);
1340
+ }
1341
+ }
1342
+ }
1343
+ }
1344
+ };
1345
+
1346
+ store_acc(c_acc0, n_tile);
1347
+ store_acc(c_acc1, n_tile + 16);
1348
+}
1349
+
1350
+#else // non-RDNA3 device pass: empty kernel for symbol parity.
1351
+template <typename T>
1352
+__global__ void gemm_q4_wmma_kernel_64x32_4w(const T*, const uint32_t*,
1353
+ const uint32_t*, const T*, T*,
1354
+ const int, const int, const int,
1355
+ const int, const int, const int*) {
1356
+}
1357
+#endif
1358
+
1359
+template <typename T>
1360
+void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight,
1361
+ const uint32_t* b_qzeros, const T* b_scales,
1362
+ const int* b_q_perm, T* c, int size_m,
1363
+ int size_n, int size_k, int groups,
1364
+ int zero_offset, cudaStream_t stream) {
1365
+ // Fall back to v3 when M < 64 (small-M decode/prefill stays on the
1366
+ // narrower 64M × 16N path) or when N < 32 (tile would waste a wave on
1367
+ // out-of-range cols).
1368
+ if (size_m < 64 || size_n < 32) {
1369
+ launch_gemm_q4_wmma_64x16_4w<T>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
1370
+ c, size_m, size_n, size_k, groups,
1371
+ zero_offset, stream);
1372
+ return;
1373
+ }
1374
+
1375
+ // 4 waves per block (128 threads), 64M × 32N tile per block.
1376
+ const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 32);
1377
+ dim3 block(128);
1378
+ dim3 grid((size_n + 31) / 32, (size_m + 63) / 64, k_split);
1379
+ gemm_q4_wmma_kernel_64x32_4w<T><<<grid, block, 0, stream>>>(
1380
+ a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
1381
+ zero_offset, b_q_perm);
1382
+}
1383
+
1384
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
1385
+
1386
+// ===========================================================================
1387
+// 64x64_4w kernel: 4 waves per block, 64M × 64N tile, 4 wmmas per wave per
1388
+// K-iter.
1389
+//
1390
+// Doubles the N-tile from 32 → 64. Each wave issues 4 wmmas per K-iter
1391
+// (cols 0-15, 16-31, 32-47, 48-63), all sharing the same a_frag. With
1392
+// 4 waves × 4 wmmas = 16 wmmas in flight per K-iter, the wmma pipeline
1393
+// is fully saturated (16-cycle latency × 1 issue/cycle = 16 wmmas in
1394
+// flight at peak).
1395
+//
1396
+// Costs:
1397
+// * LDS B-tile: 2 × 16K × 64N × sizeof(T) = 4 KB. Within budget.
1398
+// * Dequant: 64 N × 2 K-oct = 128 slots/K-tile. Distributed across all
1399
+// 4 waves: 32 slots/wave (16 N-cols × 2 K-octets per wave) — perfect
1400
+// 32-lane fit, full lane utilization on dequant.
1401
+// * Per-wave registers: 4 × v8fp32 acc (64 VGPRs) + 4 b_frag (64 VGPRs)
1402
+// + a_frag (16 VGPRs) ≈ 144 VGPRs/thread + locals ≈ ~170 total.
1403
+// Under the 192-VGPR gfx1100 cap.
1404
+//
1405
+// Mapping invariant:
1406
+// * Wave w produces M rows [m_tile + 16w .. m_tile + 16w + 15]
1407
+// * c_acc[i] holds N cols [n_tile + 16i .. n_tile + 16i + 15] for i=0..3
1408
+// ===========================================================================
1409
+
1410
+template <typename T>
1411
+__global__ void gemm_q4_wmma_kernel_64x64_4w(
1412
+ const T* __restrict__ a, const uint32_t* __restrict__ b_q,
1413
+ const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
1414
+ T* __restrict__ c, const int size_m, const int size_n, const int size_k,
1415
+ const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
1416
+ using E = typename WmmaNative<T>::elem;
1417
+ using V16 = typename WmmaNative<T>::v16;
1418
+
1419
+ const int m_tile = blockIdx.y * 64;
1420
+ const int n_tile = blockIdx.x * 64; // 64-col stride per block
1421
+ if (m_tile >= size_m || n_tile >= size_n) return;
1422
+
1423
+ const int tid = threadIdx.x; // 0..127
1424
+ const int wave_id = tid >> 5; // 0..3
1425
+ const int lane = tid & 31;
1426
+ const int lane_lo = lane & 15;
1427
+ const int lane_hi = lane >> 4;
1428
+
1429
+ // Four accumulators per wave, each covering 16 N-cols.
1430
+ v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0};
1431
+ v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0};
1432
+ v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0};
1433
+ v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0};
1434
+
1435
+ const int groupsize = size_k / groups;
1436
+
1437
+ const int k_per_split = size_k / gridDim.z;
1438
+ const int k_start = blockIdx.z * k_per_split;
1439
+ const int k_end = k_start + k_per_split;
1440
+
1441
+ // Larger LDS B-tile: [buf][k][n_in_tile=0..63].
1442
+ __shared__ T b_lds[2][16][64];
1443
+
1444
+ // Dequant: 128 slots per K-tile (64 N-cols × 2 K-octets). All 4 waves
1445
+ // participate; wave w covers n_in_tile in [16w..16w+15], k_oct in {0,1}
1446
+ // — 32 slots per wave (16 N-cols × 2 K-octets), perfect 32-lane fit.
1447
+ auto dequant_into = [&](int buf, int k_tile) {
1448
+ const int my_n_local = lane_lo; // 0..15 within wave
1449
+ const int my_n_in_tile = wave_id * 16 + my_n_local; // 0..63
1450
+ const int my_k_octet = lane_hi;
1451
+ const int actual_n = n_tile + my_n_in_tile;
1452
+
1453
+ if (actual_n >= size_n) return;
1454
+
1455
+ const int qk_row = (k_tile / 8) + my_k_octet;
1456
+ const uint32_t qa = b_q[qk_row * size_n + actual_n];
1457
+
1458
+ const int g = k_tile / groupsize;
1459
+ const int qz_idx = g * (size_n / 8) + actual_n / 8;
1460
+ const int qz_shift = (actual_n & 7) * 4;
1461
+ const uint32_t zero_v =
1462
+ ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset;
1463
+ const T scale_t = b_scales[g * size_n + actual_n];
1464
+
1465
+ const int k_base = my_k_octet * 8;
1466
+
1467
+ if constexpr (std::is_same<T, half>::value) {
1468
+ half2 z_prep, y_prep;
1469
+ prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep);
1470
+ half2 dq[4];
1471
+ dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep);
1472
+ b_lds[buf][k_base + 0][my_n_in_tile] = __low2half(dq[0]);
1473
+ b_lds[buf][k_base + 1][my_n_in_tile] = __high2half(dq[0]);
1474
+ b_lds[buf][k_base + 2][my_n_in_tile] = __low2half(dq[1]);
1475
+ b_lds[buf][k_base + 3][my_n_in_tile] = __high2half(dq[1]);
1476
+ b_lds[buf][k_base + 4][my_n_in_tile] = __low2half(dq[2]);
1477
+ b_lds[buf][k_base + 5][my_n_in_tile] = __high2half(dq[2]);
1478
+ b_lds[buf][k_base + 6][my_n_in_tile] = __low2half(dq[3]);
1479
+ b_lds[buf][k_base + 7][my_n_in_tile] = __high2half(dq[3]);
1480
+ } else {
1481
+ float z_f, y_f;
1482
+ prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f);
1483
+ bf162_t dq[4];
1484
+ dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f);
1485
+ b_lds[buf][k_base + 0][my_n_in_tile] = dq[0].x;
1486
+ b_lds[buf][k_base + 1][my_n_in_tile] = dq[0].y;
1487
+ b_lds[buf][k_base + 2][my_n_in_tile] = dq[1].x;
1488
+ b_lds[buf][k_base + 3][my_n_in_tile] = dq[1].y;
1489
+ b_lds[buf][k_base + 4][my_n_in_tile] = dq[2].x;
1490
+ b_lds[buf][k_base + 5][my_n_in_tile] = dq[2].y;
1491
+ b_lds[buf][k_base + 6][my_n_in_tile] = dq[3].x;
1492
+ b_lds[buf][k_base + 7][my_n_in_tile] = dq[3].y;
1493
+ }
1494
+ };
1495
+
1496
+ dequant_into(0, k_start);
1497
+ __syncthreads();
1498
+
1499
+ int cur_buf = 0;
1500
+ for (int k_tile = k_start; k_tile < k_end; k_tile += 16) {
1501
+ const int next_buf = 1 - cur_buf;
1502
+ const int k_next = k_tile + 16;
1503
+
1504
+ if (k_next < k_end) {
1505
+ dequant_into(next_buf, k_next);
1506
+ }
1507
+
1508
+ // Load A: each wave loads its 16M slice, shared across all 4 wmmas.
1509
+ const int m_row = m_tile + wave_id * 16 + lane_lo;
1510
+ V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3;
1511
+ if (m_row < size_m) {
1512
+ const T* a_row = a + m_row * size_k;
1513
+ if (b_q_perm) {
1514
+ #pragma unroll
1515
+ for (int i = 0; i < 16; i++) {
1516
+ T v = a_row[b_q_perm[k_tile + i]];
1517
+ a_frag[i] = bitcast_elem<T, E>(v);
1518
+ }
1519
+ } else {
1520
+ static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes");
1521
+ __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag));
1522
+ }
1523
+ } else {
1524
+ #pragma unroll
1525
+ for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
1526
+ }
1527
+
1528
+ // Load B for all four 16-col halves.
1529
+ #pragma unroll
1530
+ for (int i = 0; i < 16; i++) {
1531
+ b_frag0[i] = bitcast_elem<T, E>(b_lds[cur_buf][i][lane_lo + 0]);
1532
+ b_frag1[i] = bitcast_elem<T, E>(b_lds[cur_buf][i][lane_lo + 16]);
1533
+ b_frag2[i] = bitcast_elem<T, E>(b_lds[cur_buf][i][lane_lo + 32]);
1534
+ b_frag3[i] = bitcast_elem<T, E>(b_lds[cur_buf][i][lane_lo + 48]);
1535
+ }
1536
+
1537
+ // Four wmmas per wave per K-iter, sharing a_frag.
1538
+ c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0);
1539
+ c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1);
1540
+ c_acc2 = wmma_mma(a_frag, b_frag2, c_acc2);
1541
+ c_acc3 = wmma_mma(a_frag, b_frag3, c_acc3);
1542
+
1543
+ __syncthreads();
1544
+ cur_buf = next_buf;
1545
+ }
1546
+
1547
+ // ---- Store C ---- Each wave owns 16M × 64N. Helper writes one acc slice.
1548
+ const int m_tile_wave = m_tile + wave_id * 16;
1549
+ auto store_acc = [&](const v8fp32& acc, int n_base) {
1550
+ if (gridDim.z > 1) {
1551
+ const bool is_even_lane = (lane_lo & 1) == 0;
1552
+ const int out_n_pair = n_base + lane_lo;
1553
+ #pragma unroll
1554
+ for (int i = 0; i < 8; i++) {
1555
+ float other_f = __shfl_xor(acc[i], 1);
1556
+ if (!is_even_lane) continue;
1557
+
1558
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1559
+ if (out_m >= size_m || out_n_pair >= size_n) continue;
1560
+
1561
+ T* dst = c + out_m * size_n + out_n_pair;
1562
+ if constexpr (std::is_same<T, half>::value) {
1563
+ half2 packed =
1564
+ __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f));
1565
+ atomic_add_pk_f16(reinterpret_cast<half2*>(dst), packed);
1566
+ } else {
1567
+ bf162_t packed;
1568
+ packed.x = __float2bfloat16(acc[i]);
1569
+ packed.y = __float2bfloat16(other_f);
1570
+ atomic_add_pk_bf16(reinterpret_cast<bf162_t*>(dst), packed);
1571
+ }
1572
+ }
1573
+ } else {
1574
+ const int out_n = n_base + lane_lo;
1575
+ if (out_n >= size_n) return;
1576
+ #pragma unroll
1577
+ for (int i = 0; i < 8; i++) {
1578
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1579
+ if (out_m < size_m) {
1580
+ T* dst = c + out_m * size_n + out_n;
1581
+ if constexpr (std::is_same<T, half>::value) {
1582
+ *dst = __float2half_rn(acc[i]);
1583
+ } else {
1584
+ *dst = __float2bfloat16(acc[i]);
1585
+ }
1586
+ }
1587
+ }
1588
+ }
1589
+ };
1590
+
1591
+ store_acc(c_acc0, n_tile + 0);
1592
+ store_acc(c_acc1, n_tile + 16);
1593
+ store_acc(c_acc2, n_tile + 32);
1594
+ store_acc(c_acc3, n_tile + 48);
1595
+}
1596
+
1597
+// ===========================================================================
1598
+// 128x64_k16 kernel: 8 waves per block, 128M × 64N tile, K=16 per iteration.
1599
+//
1600
+// Doubles M-tile from 64 → 128. Each B-tile in LDS is reused by 8 waves
1601
+// (8 independent A-row slices) instead of 4, halving the effective B-load
1602
+// cost per output element. This matches Hybrid Triton's BLOCK_M=128.
1603
+//
1604
+// Dequant: same 128 slots as V5 (64N × 2 K-octets). Only waves 0-3
1605
+// participate in dequant (same mapping). Waves 4-7 are pure compute.
1606
+// LDS: [2][16][64] × sizeof(T) — unchanged from V5.
1607
+// ===========================================================================
1608
+
1609
+template <typename T>
1610
+__global__ void gemm_q4_wmma_kernel_128x64_k16(
1611
+ const T* __restrict__ a, const uint32_t* __restrict__ b_q,
1612
+ const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
1613
+ T* __restrict__ c, const int size_m, const int size_n, const int size_k,
1614
+ const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
1615
+ using E = typename WmmaNative<T>::elem;
1616
+ using V16 = typename WmmaNative<T>::v16;
1617
+
1618
+ const int m_tile = blockIdx.y * 128; // 128-row M tile
1619
+ const int n_tile = blockIdx.x * 64;
1620
+ if (m_tile >= size_m || n_tile >= size_n) return;
1621
+
1622
+ const int tid = threadIdx.x; // 0..255
1623
+ const int wave_id = tid >> 5; // 0..7
1624
+ const int lane = tid & 31;
1625
+ const int lane_lo = lane & 15;
1626
+ const int lane_hi = lane >> 4;
1627
+
1628
+ v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0};
1629
+ v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0};
1630
+ v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0};
1631
+ v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0};
1632
+
1633
+ const int groupsize = size_k / groups;
1634
+
1635
+ const int k_per_split = size_k / gridDim.z;
1636
+ const int k_start = blockIdx.z * k_per_split;
1637
+ const int k_end = k_start + k_per_split;
1638
+
1639
+ // LDS with K-contiguous layout [buf][N][K] — enables vectorized ds_load_b128.
1640
+ // Triton uses this layout to read 8 bf16 per instruction instead of 1.
1641
+ __shared__ T b_lds[2][64][16];
1642
+
1643
+ // Dequant with scale/zero caching: 1 global_load per iter (7/8 of the time).
1644
+ const bool dq_ok =
1645
+ (wave_id < 4) && (n_tile + wave_id * 16 + lane_lo < size_n);
1646
+ const int dq_n = wave_id * 16 + lane_lo;
1647
+ const int dq_an = n_tile + dq_n;
1648
+ const int dq_oct = lane_hi;
1649
+ const int dq_kb = dq_oct * 8;
1650
+ half2 ch_z = {}, ch_y = {};
1651
+ float cf_z = 0, cf_y = 0;
1652
+ int cached_g = -1;
1653
+
1654
+ auto dequant_into = [&](int buf, int k_tile) __attribute__((always_inline)) {
1655
+ if (!dq_ok) return;
1656
+
1657
+ // Reload scale/zero only on group boundary (every groupsize/16 iters).
1658
+ const int g = k_tile / groupsize;
1659
+ if (g != cached_g) {
1660
+ cached_g = g;
1661
+ const int qz_idx = g * (size_n / 8) + dq_an / 8;
1662
+ const uint32_t zero_v = ((b_qzeros[qz_idx] >> ((dq_an & 7) * 4)) & 0xF) +
1663
+ (uint32_t)zero_offset;
1664
+ const T sc = b_scales[g * size_n + dq_an];
1665
+ if constexpr (std::is_same<T, half>::value)
1666
+ prep_zero_scale_fp16_precise(zero_v, sc, ch_z, ch_y);
1667
+ else
1668
+ prep_zero_scale_bf16_f32(zero_v, sc, cf_z, cf_y);
1669
+ }
1670
+
1671
+ const int qk_row = (k_tile / 8) + dq_oct;
1672
+ const uint32_t qa = b_q[qk_row * size_n + dq_an];
1673
+ const int k_base = dq_kb;
1674
+
1675
+ if constexpr (std::is_same<T, half>::value) {
1676
+ half2 dq[4];
1677
+ dequant_4bit_8_fp16_precise(qa, dq, ch_z, ch_y);
1678
+ b_lds[buf][dq_n][k_base + 0] = __low2half(dq[0]);
1679
+ b_lds[buf][dq_n][k_base + 1] = __high2half(dq[0]);
1680
+ b_lds[buf][dq_n][k_base + 2] = __low2half(dq[1]);
1681
+ b_lds[buf][dq_n][k_base + 3] = __high2half(dq[1]);
1682
+ b_lds[buf][dq_n][k_base + 4] = __low2half(dq[2]);
1683
+ b_lds[buf][dq_n][k_base + 5] = __high2half(dq[2]);
1684
+ b_lds[buf][dq_n][k_base + 6] = __low2half(dq[3]);
1685
+ b_lds[buf][dq_n][k_base + 7] = __high2half(dq[3]);
1686
+ } else {
1687
+ bf162_t dq[4];
1688
+ dequant_4bit_8_bf16_to_bf16(qa, dq, cf_z, cf_y);
1689
+ b_lds[buf][dq_n][k_base + 0] = dq[0].x;
1690
+ b_lds[buf][dq_n][k_base + 1] = dq[0].y;
1691
+ b_lds[buf][dq_n][k_base + 2] = dq[1].x;
1692
+ b_lds[buf][dq_n][k_base + 3] = dq[1].y;
1693
+ b_lds[buf][dq_n][k_base + 4] = dq[2].x;
1694
+ b_lds[buf][dq_n][k_base + 5] = dq[2].y;
1695
+ b_lds[buf][dq_n][k_base + 6] = dq[3].x;
1696
+ b_lds[buf][dq_n][k_base + 7] = dq[3].y;
1697
+ }
1698
+ };
1699
+
1700
+ dequant_into(0, k_start);
1701
+ __syncthreads();
1702
+
1703
+ int cur_buf = 0;
1704
+ const int m_row = m_tile + wave_id * 16 + lane_lo;
1705
+ const T* a_row_ptr = (m_row < size_m) ? (a + m_row * size_k) : nullptr;
1706
+
1707
+ for (int k_tile = k_start; k_tile < k_end; k_tile += 16) {
1708
+ const int next_buf = 1 - cur_buf;
1709
+ const int k_next = k_tile + 16;
1710
+
1711
+ if (k_next < k_end) {
1712
+ dequant_into(next_buf, k_next);
1713
+ }
1714
+
1715
+ // A-load: vectorized 256-bit. b_q_perm branch removed from V7 hot path
1716
+ // to eliminate ~450 ISA instructions of dead code (6× icache bloat).
1717
+ V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3;
1718
+ if (a_row_ptr) {
1719
+ __builtin_memcpy(&a_frag, a_row_ptr + k_tile, sizeof(a_frag));
1720
+ } else {
1721
+ #pragma unroll
1722
+ for (int i = 0; i < 16; i++) a_frag[i] = (E)0;
1723
+ }
1724
+
1725
+ // Vectorized LDS reads: 32 bytes (16 bf16) per b_frag → ds_load_b128 × 2.
1726
+ static_assert(sizeof(V16) == 32, "V16 must be 32 bytes for memcpy");
1727
+ __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][0], 32);
1728
+ __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][0], 32);
1729
+ __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][0], 32);
1730
+ __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][0], 32);
1731
+
1732
+ c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0);
1733
+ c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1);
1734
+ c_acc2 = wmma_mma(a_frag, b_frag2, c_acc2);
1735
+ c_acc3 = wmma_mma(a_frag, b_frag3, c_acc3);
1736
+
1737
+ __syncthreads();
1738
+ cur_buf = next_buf;
1739
+ }
1740
+
1741
+ // ---- Store C ---- Each wave owns 16M × 64N.
1742
+ const int m_tile_wave = m_tile + wave_id * 16;
1743
+ auto store_acc = [&](const v8fp32& acc, int n_base) {
1744
+ if (gridDim.z > 1) {
1745
+ const bool is_even_lane = (lane_lo & 1) == 0;
1746
+ const int out_n_pair = n_base + lane_lo;
1747
+ #pragma unroll
1748
+ for (int i = 0; i < 8; i++) {
1749
+ float other_f = __shfl_xor(acc[i], 1);
1750
+ if (!is_even_lane) continue;
1751
+
1752
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1753
+ if (out_m >= size_m || out_n_pair >= size_n) continue;
1754
+
1755
+ T* dst = c + out_m * size_n + out_n_pair;
1756
+ if constexpr (std::is_same<T, half>::value) {
1757
+ half2 packed =
1758
+ __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f));
1759
+ atomic_add_pk_f16(reinterpret_cast<half2*>(dst), packed);
1760
+ } else {
1761
+ bf162_t packed;
1762
+ packed.x = __float2bfloat16(acc[i]);
1763
+ packed.y = __float2bfloat16(other_f);
1764
+ atomic_add_pk_bf16(reinterpret_cast<bf162_t*>(dst), packed);
1765
+ }
1766
+ }
1767
+ } else {
1768
+ const int out_n = n_base + lane_lo;
1769
+ if (out_n >= size_n) return;
1770
+ #pragma unroll
1771
+ for (int i = 0; i < 8; i++) {
1772
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1773
+ if (out_m < size_m) {
1774
+ T* dst = c + out_m * size_n + out_n;
1775
+ if constexpr (std::is_same<T, half>::value) {
1776
+ *dst = __float2half_rn(acc[i]);
1777
+ } else {
1778
+ *dst = __float2bfloat16(acc[i]);
1779
+ }
1780
+ }
1781
+ }
1782
+ }
1783
+ };
1784
+
1785
+ store_acc(c_acc0, n_tile + 0);
1786
+ store_acc(c_acc1, n_tile + 16);
1787
+ store_acc(c_acc2, n_tile + 32);
1788
+ store_acc(c_acc3, n_tile + 48);
1789
+}
1790
+
1791
+// ===========================================================================
1792
+// 128x64_k32 kernel: K=32 per iteration, all 8 waves dequant.
1793
+//
1794
+// Same 128M × 64N tile as V7, but processes 32 K-elements per iteration
1795
+// instead of 16. Halves iteration count and __syncthreads() calls.
1796
+//
1797
+// Dequant mapping:
1798
+// Waves 0-3 dequant N[wave*16 +: 16] × K[0:15] (same as V7)
1799
+// Waves 4-7 dequant N[(wave-4)*16 +: 16] × K[16:31] (new)
1800
+//
1801
+// LDS layout: b_lds[2][64][34] — 2 extra padding elements per row to
1802
+// avoid 8-way bank conflicts (row stride 68 bytes gives 16 unique banks
1803
+// across the 16 lanes).
1804
+//
1805
+// Per iteration: 8 WMMAs (2 A-frags × 4 N-groups) vs V7's 4.
1806
+// Requires K divisible by 32 and groupsize ≥ 32.
1807
+// ===========================================================================
1808
+template <typename T>
1809
+__global__ void gemm_q4_wmma_kernel_128x64_k32(
1810
+ const T* __restrict__ a, const uint32_t* __restrict__ b_q,
1811
+ const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
1812
+ T* __restrict__ c, const int size_m, const int size_n, const int size_k,
1813
+ const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
1814
+ using E = typename WmmaNative<T>::elem;
1815
+ using V16 = typename WmmaNative<T>::v16;
1816
+
1817
+ const int m_tile = blockIdx.y * 128;
1818
+ const int n_tile = blockIdx.x * 64;
1819
+ if (m_tile >= size_m || n_tile >= size_n) return;
1820
+
1821
+ const int tid = threadIdx.x; // 0..255
1822
+ const int wave_id = tid >> 5; // 0..7
1823
+ const int lane = tid & 31;
1824
+ const int lane_lo = lane & 15;
1825
+ const int lane_hi = lane >> 4;
1826
+
1827
+ v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0};
1828
+ v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0};
1829
+ v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0};
1830
+ v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0};
1831
+
1832
+ const int groupsize = size_k / groups;
1833
+
1834
+ const int k_per_split = size_k / gridDim.z;
1835
+ const int k_start = blockIdx.z * k_per_split;
1836
+ const int k_end = k_start + k_per_split;
1837
+
1838
+ // Padded LDS: +2 elements per row to break bank conflicts.
1839
+ // Row stride = 34 elements × 2 bytes = 68 bytes → 16 unique banks.
1840
+ __shared__ T b_lds[2][64][34];
1841
+
1842
+ // All 8 waves dequant. Waves 0-3 fill K[0:15], waves 4-7 fill K[16:31].
1843
+ const int dq_wave4 = wave_id & 3; // 0-3 for both halves
1844
+ const int dq_k_half = (wave_id >= 4) ? 1 : 0;
1845
+ const int dq_n = dq_wave4 * 16 + lane_lo; // N position 0..63
1846
+ const int dq_an = n_tile + dq_n;
1847
+ const bool dq_ok = (dq_an < size_n);
1848
+ const int dq_oct = lane_hi + dq_k_half * 2; // K octet 0-3
1849
+ const int dq_kb = dq_oct * 8; // K base 0,8,16,24
1850
+ half2 ch_z = {}, ch_y = {};
1851
+ float cf_z = 0, cf_y = 0;
1852
+ int cached_g = -1;
1853
+
1854
+ auto dequant_into = [&](int buf, int k_tile) __attribute__((always_inline)) {
1855
+ if (!dq_ok) return;
1856
+
1857
+ const int g = k_tile / groupsize;
1858
+ if (g != cached_g) {
1859
+ cached_g = g;
1860
+ const int qz_idx = g * (size_n / 8) + dq_an / 8;
1861
+ const uint32_t zero_v = ((b_qzeros[qz_idx] >> ((dq_an & 7) * 4)) & 0xF) +
1862
+ (uint32_t)zero_offset;
1863
+ const T sc = b_scales[g * size_n + dq_an];
1864
+ if constexpr (std::is_same<T, half>::value)
1865
+ prep_zero_scale_fp16_precise(zero_v, sc, ch_z, ch_y);
1866
+ else
1867
+ prep_zero_scale_bf16_f32(zero_v, sc, cf_z, cf_y);
1868
+ }
1869
+
1870
+ const int qk_row = (k_tile / 8) + dq_oct;
1871
+ const uint32_t qa = b_q[qk_row * size_n + dq_an];
1872
+
1873
+ if constexpr (std::is_same<T, half>::value) {
1874
+ half2 dq[4];
1875
+ dequant_4bit_8_fp16_precise(qa, dq, ch_z, ch_y);
1876
+ b_lds[buf][dq_n][dq_kb + 0] = __low2half(dq[0]);
1877
+ b_lds[buf][dq_n][dq_kb + 1] = __high2half(dq[0]);
1878
+ b_lds[buf][dq_n][dq_kb + 2] = __low2half(dq[1]);
1879
+ b_lds[buf][dq_n][dq_kb + 3] = __high2half(dq[1]);
1880
+ b_lds[buf][dq_n][dq_kb + 4] = __low2half(dq[2]);
1881
+ b_lds[buf][dq_n][dq_kb + 5] = __high2half(dq[2]);
1882
+ b_lds[buf][dq_n][dq_kb + 6] = __low2half(dq[3]);
1883
+ b_lds[buf][dq_n][dq_kb + 7] = __high2half(dq[3]);
1884
+ } else {
1885
+ bf162_t dq[4];
1886
+ dequant_4bit_8_bf16_to_bf16(qa, dq, cf_z, cf_y);
1887
+ b_lds[buf][dq_n][dq_kb + 0] = dq[0].x;
1888
+ b_lds[buf][dq_n][dq_kb + 1] = dq[0].y;
1889
+ b_lds[buf][dq_n][dq_kb + 2] = dq[1].x;
1890
+ b_lds[buf][dq_n][dq_kb + 3] = dq[1].y;
1891
+ b_lds[buf][dq_n][dq_kb + 4] = dq[2].x;
1892
+ b_lds[buf][dq_n][dq_kb + 5] = dq[2].y;
1893
+ b_lds[buf][dq_n][dq_kb + 6] = dq[3].x;
1894
+ b_lds[buf][dq_n][dq_kb + 7] = dq[3].y;
1895
+ }
1896
+ };
1897
+
1898
+ dequant_into(0, k_start);
1899
+ __syncthreads();
1900
+
1901
+ int cur_buf = 0;
1902
+ const int m_row = m_tile + wave_id * 16 + lane_lo;
1903
+ const T* a_row_ptr = (m_row < size_m) ? (a + m_row * size_k) : nullptr;
1904
+
1905
+ for (int k_tile = k_start; k_tile < k_end; k_tile += 32) {
1906
+ const int next_buf = 1 - cur_buf;
1907
+ const int k_next = k_tile + 32;
1908
+
1909
+ if (k_next < k_end) {
1910
+ dequant_into(next_buf, k_next);
1911
+ }
1912
+
1913
+ // A-load: two 16-element fragments for K[0:15] and K[16:31].
1914
+ V16 a_frag_lo, a_frag_hi;
1915
+ V16 b_frag0, b_frag1, b_frag2, b_frag3;
1916
+ if (a_row_ptr) {
1917
+ __builtin_memcpy(&a_frag_lo, a_row_ptr + k_tile, sizeof(V16));
1918
+ __builtin_memcpy(&a_frag_hi, a_row_ptr + k_tile + 16, sizeof(V16));
1919
+ } else {
1920
+ #pragma unroll
1921
+ for (int i = 0; i < 16; i++) a_frag_lo[i] = (E)0;
1922
+ #pragma unroll
1923
+ for (int i = 0; i < 16; i++) a_frag_hi[i] = (E)0;
1924
+ }
1925
+
1926
+ // --- Lower K half [0:15]: 4 WMMAs ---
1927
+ static_assert(sizeof(V16) == 32, "V16 must be 32 bytes for memcpy");
1928
+ __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][0], 32);
1929
+ __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][0], 32);
1930
+ __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][0], 32);
1931
+ __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][0], 32);
1932
+
1933
+ c_acc0 = wmma_mma(a_frag_lo, b_frag0, c_acc0);
1934
+ c_acc1 = wmma_mma(a_frag_lo, b_frag1, c_acc1);
1935
+ c_acc2 = wmma_mma(a_frag_lo, b_frag2, c_acc2);
1936
+ c_acc3 = wmma_mma(a_frag_lo, b_frag3, c_acc3);
1937
+
1938
+ // --- Upper K half [16:31]: 4 WMMAs ---
1939
+ __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][16], 32);
1940
+ __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][16], 32);
1941
+ __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][16], 32);
1942
+ __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][16], 32);
1943
+
1944
+ c_acc0 = wmma_mma(a_frag_hi, b_frag0, c_acc0);
1945
+ c_acc1 = wmma_mma(a_frag_hi, b_frag1, c_acc1);
1946
+ c_acc2 = wmma_mma(a_frag_hi, b_frag2, c_acc2);
1947
+ c_acc3 = wmma_mma(a_frag_hi, b_frag3, c_acc3);
1948
+
1949
+ __syncthreads();
1950
+ cur_buf = next_buf;
1951
+ }
1952
+
1953
+ // ---- Store C ---- Same as V7.
1954
+ const int m_tile_wave = m_tile + wave_id * 16;
1955
+ auto store_acc = [&](const v8fp32& acc, int n_base) {
1956
+ if (gridDim.z > 1) {
1957
+ const bool is_even_lane = (lane_lo & 1) == 0;
1958
+ const int out_n_pair = n_base + lane_lo;
1959
+ #pragma unroll
1960
+ for (int i = 0; i < 8; i++) {
1961
+ float other_f = __shfl_xor(acc[i], 1);
1962
+ if (!is_even_lane) continue;
1963
+
1964
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1965
+ if (out_m >= size_m || out_n_pair >= size_n) continue;
1966
+
1967
+ T* dst = c + out_m * size_n + out_n_pair;
1968
+ if constexpr (std::is_same<T, half>::value) {
1969
+ half2 packed =
1970
+ __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f));
1971
+ atomic_add_pk_f16(reinterpret_cast<half2*>(dst), packed);
1972
+ } else {
1973
+ bf162_t packed;
1974
+ packed.x = __float2bfloat16(acc[i]);
1975
+ packed.y = __float2bfloat16(other_f);
1976
+ atomic_add_pk_bf16(reinterpret_cast<bf162_t*>(dst), packed);
1977
+ }
1978
+ }
1979
+ } else {
1980
+ const int out_n = n_base + lane_lo;
1981
+ if (out_n >= size_n) return;
1982
+ #pragma unroll
1983
+ for (int i = 0; i < 8; i++) {
1984
+ const int out_m = m_tile_wave + 2 * i + lane_hi;
1985
+ if (out_m < size_m) {
1986
+ T* dst = c + out_m * size_n + out_n;
1987
+ if constexpr (std::is_same<T, half>::value) {
1988
+ *dst = __float2half_rn(acc[i]);
1989
+ } else {
1990
+ *dst = __float2bfloat16(acc[i]);
1991
+ }
1992
+ }
1993
+ }
1994
+ }
1995
+ };
1996
+
1997
+ store_acc(c_acc0, n_tile + 0);
1998
+ store_acc(c_acc1, n_tile + 16);
1999
+ store_acc(c_acc2, n_tile + 32);
2000
+ store_acc(c_acc3, n_tile + 48);
2001
+}
2002
+
2003
+#else // non-RDNA3 device pass: empty kernels for symbol parity (covers the
2004
+ // three kernels that share this launcher).
2005
+template <typename T>
2006
+__global__ void gemm_q4_wmma_kernel_64x64_4w(const T*, const uint32_t*,
2007
+ const uint32_t*, const T*, T*,
2008
+ const int, const int, const int,
2009
+ const int, const int, const int*) {
2010
+}
2011
+template <typename T>
2012
+__global__ void gemm_q4_wmma_kernel_128x64_k16(const T*, const uint32_t*,
2013
+ const uint32_t*, const T*, T*,
2014
+ const int, const int, const int,
2015
+ const int, const int,
2016
+ const int*) {}
2017
+template <typename T>
2018
+__global__ void gemm_q4_wmma_kernel_128x64_k32(const T*, const uint32_t*,
2019
+ const uint32_t*, const T*, T*,
2020
+ const int, const int, const int,
2021
+ const int, const int,
2022
+ const int*) {}
2023
+#endif
2024
+
2025
+template <typename T>
2026
+void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight,
2027
+ const uint32_t* b_qzeros, const T* b_scales,
2028
+ const int* b_q_perm, T* c, int size_m,
2029
+ int size_n, int size_k, int groups,
2030
+ int zero_offset, cudaStream_t stream) {
2031
+ // Fall back to v4 when N < 64 (would waste 1+ waves on out-of-range cols).
2032
+ if (size_m < 64 || size_n < 64) {
2033
+ launch_gemm_q4_wmma_64x32_4w<T>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
2034
+ c, size_m, size_n, size_k, groups,
2035
+ zero_offset, stream);
2036
+ return;
2037
+ }
2038
+
2039
+ // V8 (128M × 64N, K=32/iter, 8-wave dequant) when K%32==0 and gs≥32.
2040
+ // Falls back to V7 otherwise. V7/V8 read A sequentially, so act-order
2041
+ // (b_q_perm != null) must skip them and use v5, which honors the perm.
2042
+ if (size_m >= 128 && b_q_perm == nullptr) {
2043
+ const int k_split =
2044
+ compute_wmma_k_split_mn(size_m, size_n, size_k, 128, 64);
2045
+ const int groupsize = size_k / groups;
2046
+ dim3 block(256);
2047
+ dim3 grid((size_n + 63) / 64, (size_m + 127) / 128, k_split);
2048
+ if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) {
2049
+ gemm_q4_wmma_kernel_128x64_k32<T><<<grid, block, 0, stream>>>(
2050
+ a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
2051
+ zero_offset, b_q_perm);
2052
+ } else {
2053
+ gemm_q4_wmma_kernel_128x64_k16<T><<<grid, block, 0, stream>>>(
2054
+ a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
2055
+ zero_offset, b_q_perm);
2056
+ }
2057
+ return;
2058
+ }
2059
+
2060
+ // 4 waves per block (128 threads), 64M × 64N tile per block.
2061
+ const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 64);
2062
+ dim3 block(128);
2063
+ dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, k_split);
2064
+ gemm_q4_wmma_kernel_64x64_4w<T><<<grid, block, 0, stream>>>(
2065
+ a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
2066
+ zero_offset, b_q_perm);
2067
+}
2068
+
2069
+} // namespace gptq_rdna3_wmma
2070
+} // namespace vllm
2071
+
2072
+// ---------------------------------------------------------------------------
2073
+// Public entry point.
2074
+// ---------------------------------------------------------------------------
2075
+//
2076
+// Inputs:
2077
+// a [M, K] half or bfloat16
2078
+// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle)
2079
+// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros)
2080
+// b_scales [groups, N] half or bfloat16
2081
+// b_g_idx [K] or empty int32 (act-order permutation; empty=identity)
2082
+// use_v2_format bool (true = GPTQv2, no +1 zero offset)
2083
+//
2084
+// Output:
2085
+// c [M, N] same dtype as a
2086
+//
2087
+// Requirements:
2088
+// * size_m >= 16 (otherwise prefer the scalar gptq_gemm_rdna3 op)
2089
+// * size_n % 16 == 0 (WMMA tile size)
2090
+// * size_k % 16 == 0 (WMMA tile size)
2091
+
2092
+torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
2093
+ torch::Tensor b_qzeros,
2094
+ torch::Tensor b_scales,
2095
+ torch::Tensor b_g_idx, bool use_v2_format) {
2096
+ TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
2097
+ TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
2098
+ TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor");
2099
+ TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor");
2100
+ TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]");
2101
+ TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]");
2102
+ TORCH_CHECK(
2103
+ a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
2104
+ "a must be half or bfloat16");
2105
+ TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
2106
+ "b_scales dtype must match a");
2107
+
2108
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
2109
+ auto stream = at::cuda::getCurrentCUDAStream();
2110
+
2111
+ int size_m = (int)a.size(0);
2112
+ int size_k = (int)a.size(1);
2113
+ int size_n = (int)b_q_weight.size(1);
2114
+ int groups = (int)b_qzeros.size(0);
2115
+
2116
+ TORCH_CHECK(b_q_weight.size(0) * 8 == size_k,
2117
+ "b_q_weight first dim must be K/8");
2118
+ TORCH_CHECK(b_scales.size(0) == groups,
2119
+ "b_scales must have same group count as qzeros");
2120
+ TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N");
2121
+ TORCH_CHECK(size_n % 16 == 0, "WMMA path requires N % 16 == 0");
2122
+ TORCH_CHECK(size_k % 16 == 0, "WMMA path requires K % 16 == 0");
2123
+
2124
+ auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device());
2125
+ // Always zero-init the output: some V3-V8 boundary threads may exit
2126
+ // without writing their output cell (e.g. out_m >= size_m), leaving
2127
+ // uninitialized garbage when torch::empty is used. The cost is
2128
+ // negligible (< 1.5% of prefill time on gfx1100).
2129
+ at::Tensor c = torch::zeros({size_m, size_n}, opts);
2130
+
2131
+ const int* g_idx_ptr = nullptr;
2132
+ if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) {
2133
+ TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32,
2134
+ "b_g_idx must be int32");
2135
+ g_idx_ptr = (const int*)b_g_idx.data_ptr();
2136
+ }
2137
+
2138
+ const int zero_offset = use_v2_format ? 0 : 1;
2139
+
2140
+ // launch_gemm_q4_wmma_64x64_4w dispatches:
2141
+ // M >= 128 → 128x64_k32 / 128x64_k16 (8 waves, K=32/16 per iter)
2142
+ // 64 <= M < 128 & N >= 64 → 64x64_4w (4 waves, 4 wmma/wave/K-iter)
2143
+ // M >= 64 && 32 <= N < 64 → 64x32_4w (4 waves, 2 wmma/wave/K-iter)
2144
+ // M >= 64 && N < 32 → 64x16_4w (4 waves)
2145
+ // 32 <= M < 64 → 32x16_2w (2 waves)
2146
+ // M < 32 → 16x16_1w (1 wave)
2147
+ if (a.scalar_type() == torch::kHalf) {
2148
+ vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w<half>(
2149
+ (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(),
2150
+ (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(),
2151
+ g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups,
2152
+ zero_offset, stream);
2153
+ } else {
2154
+ vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w<
2155
+ vllm::gptq_rdna3_wmma::bf16_t>(
2156
+ (const vllm::gptq_rdna3_wmma::bf16_t*)a.data_ptr(),
2157
+ (const uint32_t*)b_q_weight.data_ptr(),
2158
+ (const uint32_t*)b_qzeros.data_ptr(),
2159
+ (const vllm::gptq_rdna3_wmma::bf16_t*)b_scales.data_ptr(), g_idx_ptr,
2160
+ (vllm::gptq_rdna3_wmma::bf16_t*)c.data_ptr(), size_m, size_n, size_k,
2161
+ groups, zero_offset, stream);
2162
+ }
2163
+
2164
+ return c;
2165
+}
csrc/rocm/qdq_4_rdna3.cuh
new
+239
@@ -0,0 +1,239 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+//
4
+// W4A16 dequant primitives for RDNA3 (gfx1100/gfx1101/gfx1102), templated on
5
+// the activation/scale dtype (half or __hip_bfloat16). The fp16 path reuses
6
+// the classic exllamav2 bit-trick:
7
+//
8
+// (qa & 0x000F000F) | 0x64006400 -> half2(1024+q_lo, 1024+q_hi)
9
+// (qa & 0x00F000F0) | 0x64006400 -> half2(1024+q_lo*16, 1024+q_hi*16)
10
+//
11
+// The "*16 then divide by 16 in the FMA" trick for the upper-nibble pairs
12
+// works in fp16 because the mantissa (10 bits) is wide enough to hold a value
13
+// shifted by 4 bits. In bf16 the mantissa is only 7 bits, so shifting an upper
14
+// nibble into bits [7:4] would spill into the exponent. To avoid that, the
15
+// bf16 path shifts each pair of nibbles down to bits [3:0]/[19:16] with a
16
+// single right-shift before the OR with 0x43004300 (= bf162(128, 128)).
17
+
18
+#ifndef _qdq_4_rdna3_cuh
19
+#define _qdq_4_rdna3_cuh
20
+
21
+#include <cstdint>
22
+
23
+#include <hip/hip_bf16.h>
24
+#include <hip/hip_fp16.h>
25
+
26
+namespace vllm {
27
+namespace gptq_rdna3 {
28
+
29
+using bf16_t = __hip_bfloat16;
30
+using bf162_t = __hip_bfloat162;
31
+
32
+// Bit-shuffle for an int32 holding 8 sequential 4-bit weights q[0..7]:
33
+// in: q[7] q[6] q[5] q[4] q[3] q[2] q[1] q[0] (LSB first)
34
+// out: q[7] q[5] q[3] q[1] q[6] q[4] q[2] q[0] (even/odd interleaved)
35
+//
36
+// After shuffle, q[2k] sits at bits [4k : 4k+3] (lower 16)
37
+// q[2k+1] sits at bits [16+4k: 16+4k+3] (upper 16)
38
+// so a single mask 0x000F000F selects the matching even/odd pair, ready to
39
+// bitcast to half2 / bfloat162 after OR-ing with the magic constant.
40
+__forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) {
41
+ uint32_t qa = q[0];
42
+ uint32_t qb = 0;
43
+#pragma unroll
44
+ for (int i = 0; i < 4; i++) {
45
+ uint32_t qa0 = qa & 0x0F;
46
+ uint32_t qa1 = (qa & 0xF0) >> 4;
47
+ qa >>= 8;
48
+ qb |= (qa1 << (i * 4 + 16));
49
+ qb |= (qa0 << (i * 4));
50
+ }
51
+ q[0] = qb;
52
+}
53
+
54
+// ---------------------------------------------------------------------------
55
+// fp16 path
56
+// ---------------------------------------------------------------------------
57
+
58
+// Precompute scale-baked constants for a single zero/scale pair.
59
+// z1z16[0] = scale * (-1024 - zero) (used for "low" pairs)
60
+// z1z16[1] = scale * (-64 - zero) (used for "high" pairs)
61
+// y1y16[0] = scale * 1 (low pairs are q + 1024)
62
+// y1y16[1] = scale * (1/16) (high pairs are q*16 + 1024)
63
+__forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale,
64
+ half2 (&z1z16)[2],
65
+ half2 (&y1y16)[2]) {
66
+ // half(-1024 - zero) via the exllamav2 bit-trick:
67
+ // half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it.
68
+ union {
69
+ uint16_t u;
70
+ half h;
71
+ } z1u;
72
+ z1u.u = (uint16_t)(0xE400 | zero);
73
+ half z1 = z1u.h;
74
+ half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero));
75
+
76
+ half2 scale2 = __half2half2(scale);
77
+ z1z16[0] = __hmul2(scale2, __half2half2(z1));
78
+ z1z16[1] = __hmul2(scale2, __half2half2(z16));
79
+
80
+ half y1 = __float2half_rn(1.0f);
81
+ half y16 = __float2half_rn(1.0f / 16.0f);
82
+ y1y16[0] = __hmul2(scale2, __half2half2(y1));
83
+ y1y16[1] = __hmul2(scale2, __half2half2(y16));
84
+}
85
+
86
+// Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs:
87
+// dq[0] = (q[0], q[1]) * scale - zero*scale
88
+// dq[1] = (q[2], q[3]) * scale - zero*scale
89
+// dq[2] = (q[4], q[5]) * scale - zero*scale
90
+// dq[3] = (q[6], q[7]) * scale - zero*scale
91
+__forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4],
92
+ half2 (&z1z16)[2],
93
+ half2 (&y1y16)[2]) {
94
+ const uint32_t c0 = 0x64006400;
95
+
96
+ union {
97
+ uint32_t u;
98
+ half2 h2;
99
+ } q0, q1, q2, q3;
100
+ q0.u = (qa & 0x000F000F) | c0; // half2(q[0]+1024, q[1]+1024)
101
+ q1.u = (qa & 0x00F000F0) | c0; // half2(q[2]*16+1024, q[3]*16+1024)
102
+ uint32_t qa_hi = qa >> 8;
103
+ q2.u = (qa_hi & 0x000F000F) | c0; // half2(q[4]+1024, q[5]+1024)
104
+ q3.u = (qa_hi & 0x00F000F0) | c0; // half2(q[6]*16+1024, q[7]*16+1024)
105
+
106
+ dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]);
107
+ dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]);
108
+ dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]);
109
+ dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]);
110
+}
111
+
112
+// ---------------------------------------------------------------------------
113
+// bf16 path
114
+// ---------------------------------------------------------------------------
115
+
116
+// Bit-trick magic for bf16:
117
+// bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0).
118
+// For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because
119
+// bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So
120
+// ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi).
121
+//
122
+// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble
123
+// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16]
124
+// via a single 4/8/12-bit right-shift before the OR. That costs one extra
125
+// shift per pair vs fp16, but keeps the FMA structure identical.
126
+__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero,
127
+ bf16_t scale,
128
+ bf162_t& z_prep,
129
+ bf162_t& y_prep) {
130
+ // z = scale * -(128 + zero); y = scale.
131
+ float scale_f = __bfloat162float(scale);
132
+ float zf = -(128.0f + (float)zero) * scale_f;
133
+ bf16_t zb = __float2bfloat16(zf);
134
+ z_prep = __bfloat162bfloat162(zb);
135
+ y_prep = __bfloat162bfloat162(scale);
136
+}
137
+
138
+__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa,
139
+ bf162_t (&dq)[4],
140
+ bf162_t z_prep,
141
+ bf162_t y_prep) {
142
+ const uint32_t c0 = 0x43004300;
143
+
144
+ union {
145
+ uint32_t u;
146
+ bf162_t b2;
147
+ } q0, q1, q2, q3;
148
+ q0.u = ((qa >> 0) & 0x000F000F) | c0; // bf162(128+q[0], 128+q[1])
149
+ q1.u = ((qa >> 4) & 0x000F000F) | c0; // bf162(128+q[2], 128+q[3])
150
+ q2.u = ((qa >> 8) & 0x000F000F) | c0; // bf162(128+q[4], 128+q[5])
151
+ q3.u = ((qa >> 12) & 0x000F000F) | c0; // bf162(128+q[6], 128+q[7])
152
+
153
+ // dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale
154
+ dq[0] = __hfma2(q0.b2, y_prep, z_prep);
155
+ dq[1] = __hfma2(q1.b2, y_prep, z_prep);
156
+ dq[2] = __hfma2(q2.b2, y_prep, z_prep);
157
+ dq[3] = __hfma2(q3.b2, y_prep, z_prep);
158
+}
159
+
160
+// ---------------------------------------------------------------------------
161
+// bf16-input → fp32-output dequant (RDNA3 scalar path).
162
+//
163
+// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow
164
+// fallback. Rather than computing dq in bf16 and widening at FMA time in
165
+// the dot product, we widen to fp32 here once (a free left-shift by 16) and
166
+// emit the (q - zero) * scale FMA directly in fp32. This:
167
+// * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32.
168
+// * Eliminates 4× bf16→fp32 widens that the dot product would do.
169
+// * Keeps the dot product accumulator in fp32 without a roundtrip.
170
+//
171
+// Output: fp32 dq[8], one element per K position (consumed by the
172
+// fp32-overload of dot22_8_f in q_gemm_rdna3.cu).
173
+__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
174
+ bf16_t scale,
175
+ float& z_prep,
176
+ float& y_prep) {
177
+ float scale_f = __bfloat162float(scale);
178
+ z_prep = -(128.0f + (float)zero) * scale_f;
179
+ y_prep = scale_f;
180
+}
181
+
182
+// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32
183
+// values 128+nibble, without folding scale/zero. The caller folds scale/zb
184
+// into the accumulator outside the inner loop using a precomputed sum_a,
185
+// which saves ~27% of the FMA count vs the per-col-dequant approach above
186
+// (only beneficial at M_COUNT=1; break-even at M_COUNT=2).
187
+//
188
+// Cost: 0 FMAs (pure bit-trick + as_float reinterprets).
189
+__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa,
190
+ float (&q_f32)[8]) {
191
+ const uint32_t c0 = 0x43004300;
192
+ const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
193
+ const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
194
+ const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
195
+ const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
196
+ q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16);
197
+ q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u);
198
+ q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16);
199
+ q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u);
200
+ q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16);
201
+ q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u);
202
+ q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16);
203
+ q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u);
204
+}
205
+
206
+__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa,
207
+ float (&dq)[8],
208
+ float z_prep,
209
+ float y_prep) {
210
+ const uint32_t c0 = 0x43004300;
211
+ const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
212
+ const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
213
+ const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
214
+ const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
215
+ // bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16
216
+ // (just zero-extends the mantissa from 7 to 23 bits; exponent preserved).
217
+ const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16);
218
+ const float q0y = __uint_as_float(q0 & 0xFFFF0000u);
219
+ const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16);
220
+ const float q1y = __uint_as_float(q1 & 0xFFFF0000u);
221
+ const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16);
222
+ const float q2y = __uint_as_float(q2 & 0xFFFF0000u);
223
+ const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16);
224
+ const float q3y = __uint_as_float(q3 & 0xFFFF0000u);
225
+ // dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale
226
+ dq[0] = __fmaf_rn(q0x, y_prep, z_prep);
227
+ dq[1] = __fmaf_rn(q0y, y_prep, z_prep);
228
+ dq[2] = __fmaf_rn(q1x, y_prep, z_prep);
229
+ dq[3] = __fmaf_rn(q1y, y_prep, z_prep);
230
+ dq[4] = __fmaf_rn(q2x, y_prep, z_prep);
231
+ dq[5] = __fmaf_rn(q2y, y_prep, z_prep);
232
+ dq[6] = __fmaf_rn(q3x, y_prep, z_prep);
233
+ dq[7] = __fmaf_rn(q3y, y_prep, z_prep);
234
+}
235
+
236
+} // namespace gptq_rdna3
237
+} // namespace vllm
238
+
239
+#endif // _qdq_4_rdna3_cuh
csrc/rocm/torch_bindings.cpp
+13
@@ -39,6 +39,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
39
" Tensor scale_b, int CuCount) -> ()");
40
rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
41
42
+#ifdef VLLM_ROCM_GFX1100
43
+ // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
44
+ rocm_ops.def(
45
+ "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
46
+ "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
47
+ rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3);
48
+
49
+ rocm_ops.def(
50
+ "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
51
+ "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
52
+ rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
53
+#endif
54
+
55
// Custom attention op
56
// Compute the attention between an input query and the cached
57
// keys/values using PagedAttention.
tests/kernels/quantization/test_rdna3_w4a16.py
new
+278
@@ -0,0 +1,278 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""Correctness tests for the ROCm RDNA3 W4A16 GPTQ kernel (gfx1100).
5
+
6
+Exercises ``RDNA3W4A16LinearKernel`` end-to-end: it builds a layer with
7
+GPTQ-format checkpoint parameters, runs ``process_weights_after_loading``
8
+(weight shuffle + zero-point synthesis), then ``apply_weights``, and compares
9
+the result against an fp32 reference dequant-and-matmul.
10
+
11
+The kernel is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3`` and is only
12
+built for gfx11; tests are skipped elsewhere.
13
+
14
+Run `pytest tests/kernels/quantization/test_rdna3_w4a16.py`.
15
+"""
16
+
17
+import pytest
18
+import torch
19
+
20
+from vllm.platforms import current_platform
21
+
22
+if not current_platform.is_rocm():
23
+ pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
24
+
25
+from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402
26
+ MPLinearLayerConfig,
27
+)
28
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402
29
+ RDNA3W4A16LinearKernel,
30
+)
31
+from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
32
+ pack_quantized_values_into_int32,
33
+)
34
+from vllm.model_executor.parameter import ( # noqa: E402
35
+ GroupQuantScaleParameter,
36
+ PackedvLLMParameter,
37
+)
38
+from vllm.platforms.rocm import on_gfx1100 # noqa: E402
39
+from vllm.scalar_type import scalar_types # noqa: E402
40
+from vllm.utils.torch_utils import set_random_seed # noqa: E402
41
+
42
+device = "cuda"
43
+
44
+WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8
45
+PACK_FACTOR = 8 # 8 x 4-bit nibbles per int32
46
+
47
+# Skip everything in this module unless we are on the only architecture the
48
+# kernel is built/registered for.
49
+gfx1100_only = pytest.mark.skipif(
50
+ not (
51
+ on_gfx1100()
52
+ and hasattr(torch.ops, "_rocm_C")
53
+ and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
54
+ ),
55
+ reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
56
+)
57
+
58
+
59
+# ---------------------------------------------------------------------------
60
+# Reference implementation
61
+# ---------------------------------------------------------------------------
62
+
63
+
64
+def _reference(
65
+ x_mk: torch.Tensor,
66
+ q_int4_kn: torch.Tensor,
67
+ scales_gn: torch.Tensor,
68
+ zeros_gn: torch.Tensor | None,
69
+ group_size: int,
70
+ bias: torch.Tensor | None,
71
+) -> torch.Tensor:
72
+ """fp32 reference for the RDNA3 W4A16 op.
73
+
74
+ x_mk: [M, K] fp16/bf16 activations.
75
+ q_int4_kn: [K, N] int32 raw stored nibbles in [0, 15].
76
+ scales_gn: [K//G, N] per-group scales (act dtype).
77
+ zeros_gn: [K//G, N] int32 raw stored zero points in [0, 15], or None
78
+ for the symmetric path (kernel synthesizes stored zero = 7).
79
+ group_size: G.
80
+
81
+ The kernel applies the GPTQv1 "+1" zero-point quirk, so the effective
82
+ zero is ``stored_zero + 1`` (symmetric path: 7 + 1 == bias == 8).
83
+ """
84
+ K, N = q_int4_kn.shape
85
+ s_full = scales_gn.repeat_interleave(group_size, dim=0).to(torch.float32) # [K,N]
86
+ if zeros_gn is None:
87
+ z_full = torch.full(
88
+ (K, N), float(WEIGHT_TYPE.bias), device=x_mk.device, dtype=torch.float32
89
+ )
90
+ else:
91
+ z_full = (zeros_gn + 1).repeat_interleave(group_size, dim=0).to(torch.float32)
92
+ w_fp = (q_int4_kn.to(torch.float32) - z_full) * s_full # [K, N]
93
+ out = x_mk.to(torch.float32) @ w_fp # [M, N]
94
+ if bias is not None:
95
+ out = out + bias.to(torch.float32)
96
+ return out.to(x_mk.dtype)
97
+
98
+
99
+# ---------------------------------------------------------------------------
100
+# Layer construction (GPTQ checkpoint format)
101
+# ---------------------------------------------------------------------------
102
+
103
+
104
+def _build_layer(
105
+ q_int4_kn: torch.Tensor,
106
+ scales_gn: torch.Tensor,
107
+ zeros_gn: torch.Tensor | None,
108
+ dtype: torch.dtype,
109
+) -> torch.nn.Module:
110
+ """Build a dummy layer carrying GPTQ-format params, as the loader would."""
111
+ no_loader = lambda *args, **kwargs: None # noqa: E731
112
+
113
+ # qweight: int4 packed along K into int32 -> [K//8, N].
114
+ qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0)
115
+
116
+ class DummyLayer(torch.nn.Module):
117
+ pass
118
+
119
+ layer = DummyLayer()
120
+ layer.register_parameter(
121
+ "qweight",
122
+ PackedvLLMParameter(
123
+ data=qweight,
124
+ weight_loader=no_loader,
125
+ input_dim=0,
126
+ output_dim=1,
127
+ packed_dim=0,
128
+ packed_factor=PACK_FACTOR,
129
+ ),
130
+ )
131
+ layer.register_parameter(
132
+ "scales",
133
+ GroupQuantScaleParameter(
134
+ data=scales_gn.to(dtype),
135
+ weight_loader=no_loader,
136
+ input_dim=0,
137
+ output_dim=1,
138
+ ),
139
+ )
140
+ if zeros_gn is not None:
141
+ # qzeros: int4 packed along N into int32 -> [K//G, N//8].
142
+ qzeros = pack_quantized_values_into_int32(zeros_gn, WEIGHT_TYPE, packed_dim=1)
143
+ layer.register_parameter(
144
+ "qzeros",
145
+ PackedvLLMParameter(
146
+ data=qzeros,
147
+ weight_loader=no_loader,
148
+ input_dim=0,
149
+ output_dim=1,
150
+ packed_dim=1,
151
+ packed_factor=PACK_FACTOR,
152
+ ),
153
+ )
154
+ return layer
155
+
156
+
157
+def _run_kernel(
158
+ x_mk: torch.Tensor,
159
+ q_int4_kn: torch.Tensor,
160
+ scales_gn: torch.Tensor,
161
+ zeros_gn: torch.Tensor | None,
162
+ group_size: int,
163
+ bias: torch.Tensor | None,
164
+ dtype: torch.dtype,
165
+) -> torch.Tensor:
166
+ K, N = q_int4_kn.shape
167
+ has_zp = zeros_gn is not None
168
+
169
+ config = MPLinearLayerConfig(
170
+ full_weight_shape=(K, N),
171
+ partition_weight_shape=(K, N),
172
+ weight_type=WEIGHT_TYPE,
173
+ act_type=dtype,
174
+ group_size=group_size,
175
+ zero_points=has_zp,
176
+ has_g_idx=False,
177
+ )
178
+ ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
179
+ assert ok, f"can_implement rejected a supported config: {reason}"
180
+
181
+ layer = _build_layer(q_int4_kn, scales_gn, zeros_gn, dtype)
182
+ kernel = RDNA3W4A16LinearKernel(
183
+ config,
184
+ w_q_param_name="qweight",
185
+ w_s_param_name="scales",
186
+ w_zp_param_name="qzeros" if has_zp else None,
187
+ w_gidx_param_name=None,
188
+ )
189
+ kernel.process_weights_after_loading(layer)
190
+ return kernel.apply_weights(layer, x_mk, bias=bias)
191
+
192
+
193
+# Relative-L2 tolerance per dtype. The bf16 path widens dequantized weights
194
+# to fp32 and accumulates in fp32, so it matches the reference almost exactly
195
+# (<0.4% incl. the WMMA prefill path). The fp16 path uses the exllamav2
196
+# "+1024" bit-trick (see qdq_4_rdna3.cuh): the dequantized weight is recovered
197
+# as the fp16 difference of two ~1024*scale magnitudes, which sheds low-order
198
+# mantissa bits and leaves ~2-3% relative noise that accumulates over K. We
199
+# compare on the relative Frobenius norm rather than elementwise, since the
200
+# bit-trick noise produces large *relative* errors on individual near-zero
201
+# outputs that carry negligible absolute weight.
202
+_REL_L2_TOL = {torch.float16: 5e-2, torch.bfloat16: 1e-2}
203
+
204
+
205
+def _assert_close(out: torch.Tensor, ref: torch.Tensor, dtype: torch.dtype):
206
+ rel_l2 = (out.to(torch.float32) - ref.to(torch.float32)).norm() / ref.to(
207
+ torch.float32
208
+ ).norm()
209
+ tol = _REL_L2_TOL[dtype]
210
+ assert rel_l2 < tol, f"relative L2 error {rel_l2:.4f} exceeds {tol} for {dtype}"
211
+
212
+
213
+# ---------------------------------------------------------------------------
214
+# Forward correctness
215
+# ---------------------------------------------------------------------------
216
+
217
+
218
+# (M, K, N, group_size). M spans the scalar decode path (small M) and the
219
+# WMMA prefill path (M >= 16 on the bf16 dispatch). K/N satisfy the kernel's
220
+# divisibility constraints (K % G == 0, K % 8 == 0, N % 8 == 0).
221
+MKNG_SHAPES = [
222
+ (1, 128, 128, 128), # single group, decode
223
+ (2, 256, 256, 128), # two groups
224
+ (8, 256, 512, 64), # M=8 scalar, smaller group
225
+ (16, 512, 256, 128), # M=16 -> WMMA path for bf16
226
+ (32, 512, 512, 64), # larger prefill
227
+]
228
+
229
+
230
+@gfx1100_only
231
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
232
+@pytest.mark.parametrize("has_zp", [False, True], ids=["no_zp", "with_zp"])
233
+@pytest.mark.parametrize(
234
+ "M,K,N,G", MKNG_SHAPES, ids=[f"m{m}_k{k}_n{n}_g{g}" for m, k, n, g in MKNG_SHAPES]
235
+)
236
+def test_rdna3_w4a16_matches_reference(dtype, has_zp, M, K, N, G, dist_init):
237
+ set_random_seed(0)
238
+ assert K % G == 0 and K % PACK_FACTOR == 0 and N % PACK_FACTOR == 0
239
+
240
+ groups = K // G
241
+ x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
242
+ q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
243
+ scales_gn = (
244
+ 0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
245
+ ).to(dtype)
246
+ zeros_gn = (
247
+ torch.randint(0, 16, (groups, N), device=device, dtype=torch.int32)
248
+ if has_zp
249
+ else None
250
+ )
251
+
252
+ out = _run_kernel(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None, dtype)
253
+ ref = _reference(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None)
254
+
255
+ assert out.shape == (M, N) and out.dtype == dtype
256
+ _assert_close(out, ref, dtype)
257
+
258
+
259
+@gfx1100_only
260
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
261
+@pytest.mark.parametrize("M", [1, 32], ids=["decode", "prefill"])
262
+def test_rdna3_w4a16_bias(dtype, M, dist_init):
263
+ """Bias is added on both the scalar (M=1) and WMMA (M=32) paths."""
264
+ set_random_seed(0)
265
+ K, N, G = 512, 256, 128
266
+ groups = K // G
267
+
268
+ x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
269
+ q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
270
+ scales_gn = (
271
+ 0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
272
+ ).to(dtype)
273
+ bias = (0.1 * torch.randn(N, device=device, dtype=torch.float32)).to(dtype)
274
+
275
+ out = _run_kernel(x_mk, q_int4_kn, scales_gn, None, G, bias, dtype)
276
+ ref = _reference(x_mk, q_int4_kn, scales_gn, None, G, bias)
277
+
278
+ _assert_close(out, ref, dtype)
tests/kernels/quantization/test_rdna3_w4a16_selection.py
new
+89
@@ -0,0 +1,89 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""Kernel-selection / gating tests for the ROCm RDNA3 W4A16 GPTQ kernel.
5
+
6
+Verifies that ``choose_mp_linear_kernel`` resolves a supported W4A16 GPTQ
7
+config to ``RDNA3W4A16LinearKernel`` on gfx1100 (it is registered ahead of
8
+``TritonW4A16LinearKernel`` in the ROCm priority list), and that
9
+``RDNA3W4A16LinearKernel.can_implement`` rejects the configs it does not
10
+support so selection falls through to the next kernel.
11
+
12
+Run `pytest tests/kernels/quantization/test_rdna3_w4a16_selection.py`.
13
+"""
14
+
15
+import pytest
16
+import torch
17
+
18
+from vllm.platforms import current_platform
19
+
20
+if not current_platform.is_rocm():
21
+ pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
22
+
23
+from vllm.model_executor.kernels.linear import ( # noqa: E402
24
+ choose_mp_linear_kernel,
25
+)
26
+from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402
27
+ MPLinearLayerConfig,
28
+)
29
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402
30
+ RDNA3W4A16LinearKernel,
31
+)
32
+from vllm.platforms.rocm import on_gfx1100 # noqa: E402
33
+from vllm.scalar_type import scalar_types # noqa: E402
34
+
35
+WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8
36
+
37
+# The kernel is only selectable when running on gfx1100 with the custom op
38
+# compiled in; otherwise can_implement rejects and selection falls through.
39
+gfx1100_only = pytest.mark.skipif(
40
+ not (
41
+ on_gfx1100()
42
+ and hasattr(torch.ops, "_rocm_C")
43
+ and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
44
+ ),
45
+ reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
46
+)
47
+
48
+
49
+@gfx1100_only
50
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
51
+def test_selection_prefers_rdna3(dtype):
52
+ """A supported W4A16 GPTQ config resolves to the RDNA3 kernel on gfx1100."""
53
+ config = MPLinearLayerConfig(
54
+ full_weight_shape=(1024, 256),
55
+ partition_weight_shape=(1024, 256),
56
+ weight_type=WEIGHT_TYPE,
57
+ act_type=dtype,
58
+ group_size=128,
59
+ zero_points=False,
60
+ has_g_idx=False,
61
+ )
62
+ assert choose_mp_linear_kernel(config).__name__ == "RDNA3W4A16LinearKernel"
63
+
64
+
65
+@gfx1100_only
66
+@pytest.mark.parametrize(
67
+ "weight_type,group_size,N,full_k,expected_ok",
68
+ [
69
+ (scalar_types.uint4b8, 128, 256, 1024, True), # nominal: supported
70
+ (scalar_types.uint4b8, -1, 256, 1024, False), # channelwise unsupported
71
+ (scalar_types.uint4b8, 128, 252, 1024, False), # N not a multiple of 8
72
+ (scalar_types.uint4b8, 96, 256, 1024, False), # group does not divide K
73
+ (scalar_types.uint8b128, 128, 256, 1024, False), # wrong quant type
74
+ ],
75
+ ids=["ok", "channelwise", "bad_n", "group_ndiv_k", "wrong_qtype"],
76
+)
77
+def test_can_implement(weight_type, group_size, N, full_k, expected_ok):
78
+ """can_implement gates on quant type, group size, and N divisibility."""
79
+ config = MPLinearLayerConfig(
80
+ full_weight_shape=(full_k, N),
81
+ partition_weight_shape=(full_k, N),
82
+ weight_type=weight_type,
83
+ act_type=torch.float16,
84
+ group_size=group_size,
85
+ zero_points=False,
86
+ has_g_idx=False,
87
+ )
88
+ ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
89
+ assert ok is expected_ok, reason
vllm/_custom_ops.py
+45
@@ -650,6 +650,51 @@ def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None
650
torch.ops._C.gptq_shuffle(q_weight, q_perm, bit)
651
652
653
+def gptq_gemm_rdna3(
654
+ a: torch.Tensor,
655
+ b_q_weight: torch.Tensor,
656
+ b_qzeros: torch.Tensor,
657
+ b_scales: torch.Tensor,
658
+ b_g_idx: torch.Tensor,
659
+ use_v2_format: bool,
660
+) -> torch.Tensor:
661
+ return torch.ops._rocm_C.gptq_gemm_rdna3(
662
+ a, b_q_weight, b_qzeros, b_scales, b_g_idx, use_v2_format
663
+ )
664
+
665
+
666
+if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3"):
667
+
668
+ @register_fake("_rocm_C::gptq_gemm_rdna3")
669
+ def _gptq_gemm_rdna3_fake(
670
+ a: torch.Tensor,
671
+ b_q_weight: torch.Tensor,
672
+ b_qzeros: torch.Tensor,
673
+ b_scales: torch.Tensor,
674
+ b_g_idx: torch.Tensor,
675
+ use_v2_format: bool,
676
+ ) -> torch.Tensor:
677
+ return torch.empty(
678
+ (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
679
+ )
680
+
681
+
682
+if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3_wmma"):
683
+
684
+ @register_fake("_rocm_C::gptq_gemm_rdna3_wmma")
685
+ def _gptq_gemm_rdna3_wmma_fake(
686
+ a: torch.Tensor,
687
+ b_q_weight: torch.Tensor,
688
+ b_qzeros: torch.Tensor,
689
+ b_scales: torch.Tensor,
690
+ b_g_idx: torch.Tensor,
691
+ use_v2_format: bool,
692
+ ) -> torch.Tensor:
693
+ return torch.empty(
694
+ (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
695
+ )
696
+
697
+
698
if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
699
700
@register_fake("_C::allspark_w8a16_gemm")
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -51,6 +51,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.machete import (
51
from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
52
MarlinLinearKernel,
53
)
54
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (
55
+ RDNA3W4A16LinearKernel,
56
+)
57
from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
58
TritonW4A16LinearKernel,
59
)
@@ -339,6 +342,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
342
TritonW4A16LinearKernel,
343
],
344
PlatformEnum.ROCM: [
345
+ RDNA3W4A16LinearKernel,
346
TritonW4A16LinearKernel,
347
ConchLinearKernel,
348
ExllamaLinearKernel,
vllm/model_executor/kernels/linear/mixed_precision/__init__.py
+4
@@ -29,6 +29,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
29
MPLinearKernel,
30
MPLinearLayerConfig,
31
)
32
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (
33
+ RDNA3W4A16LinearKernel,
34
+)
35
from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
36
TritonW4A16LinearKernel,
37
)
@@ -48,6 +51,7 @@ __all__ = [
51
"ExllamaLinearKernel",
52
"MacheteLinearKernel",
53
"MarlinLinearKernel",
54
+ "RDNA3W4A16LinearKernel",
55
"TritonW4A16LinearKernel",
56
"XPUW4A8IntLinearKernel",
57
"XPUwNa16LinearKernel",
vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py
new
+193
@@ -0,0 +1,193 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""W4A16 GPTQ kernel for AMD RDNA3 (gfx1100) — fp16 + bf16.
4
+
5
+Drop-in replacement for ExllamaLinearKernel on RDNA3 that adds native bf16
6
+support. The HIP kernel lives in ``csrc/rocm/q_gemm_rdna3.cu``
7
+and is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3``.
8
+
9
+Registered ahead of TritonW4A16LinearKernel for the ROCm-RDNA3 path; falls
10
+through to the Triton kernel on non-RDNA3 ROCm devices (e.g. CDNA/MI300).
11
+"""
12
+
13
+import torch
14
+
15
+from vllm import _custom_ops as ops
16
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
+ pack_quantized_values_into_int32,
18
+)
19
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
20
+from vllm.platforms import current_platform
21
+from vllm.scalar_type import scalar_types
22
+
23
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
24
+
25
+
26
+class RDNA3W4A16LinearKernel(MPLinearKernel):
27
+ SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8]
28
+
29
+ @classmethod
30
+ def get_min_capability(cls) -> int:
31
+ # ROCm gates via on_gfx1100() in can_implement.
32
+ return 60
33
+
34
+ @classmethod
35
+ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
36
+ if not current_platform.is_rocm():
37
+ return False, "RDNA3 W4A16 kernel is ROCm-only"
38
+
39
+ from vllm.platforms.rocm import on_gfx1100
40
+
41
+ if not on_gfx1100():
42
+ return False, "RDNA3 W4A16 kernel requires gfx1100"
43
+
44
+ # The HIP op is registered by the C++ extension; if a user is running
45
+ # against a vLLM build that doesn't include it (e.g. partial rebuild),
46
+ # fall through gracefully to the next kernel in the registry.
47
+ if not (
48
+ hasattr(torch.ops, "_rocm_C")
49
+ and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
50
+ ):
51
+ return (
52
+ False,
53
+ "torch.ops._rocm_C.gptq_gemm_rdna3 missing — rebuild C++ extension",
54
+ )
55
+
56
+ if c.act_type not in (torch.float16, torch.bfloat16):
57
+ return False, "RDNA3 W4A16 kernel only supports fp16 and bf16"
58
+
59
+ if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
60
+ return (
61
+ False,
62
+ f"Quant type ({c.weight_type}) not supported by "
63
+ f"RDNA3 W4A16 kernel; supported: {cls.SUPPORTED_QUANT_TYPES}",
64
+ )
65
+
66
+ if c.group_size <= 0:
67
+ return (
68
+ False,
69
+ "RDNA3 W4A16 kernel does not support channelwise quantization",
70
+ )
71
+
72
+ if c.full_weight_shape[0] % c.group_size != 0:
73
+ return (
74
+ False,
75
+ f"Group size ({c.group_size}) does not evenly divide K "
76
+ f"({c.full_weight_shape[0]})",
77
+ )
78
+
79
+ # Output features must be a multiple of the pack factor (8 nibbles per
80
+ # int32) and of 8 so that qzeros (packed 4-bit per col) align cleanly
81
+ # against the BLOCK_KN_SIZE*4 = 512 N-stride and per-thread 4 columns.
82
+ if c.partition_weight_shape[1] % 8 != 0:
83
+ return (
84
+ False,
85
+ "Output features must be a multiple of 8 for the RDNA3 "
86
+ "W4A16 kernel (qzeros packing)",
87
+ )
88
+
89
+ if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
90
+ return (
91
+ False,
92
+ "Act-order with TP-partitioned input features is not "
93
+ "supported by the RDNA3 W4A16 kernel",
94
+ )
95
+
96
+ return True, None
97
+
98
+ # ----- Weight prep (identical layout/shuffle as ExllamaLinearKernel) -----
99
+
100
+ def process_weights_after_loading(self, layer: torch.nn.Module):
101
+ c = self.config
102
+ device = getattr(layer, self.w_q_name).device
103
+
104
+ # Synthesize zero points if the checkpoint doesn't carry them.
105
+ if not c.zero_points:
106
+ self.w_zp_name = "qzeros"
107
+ groups = c.partition_weight_shape[0] // c.group_size
108
+ out_features = c.partition_weight_shape[1]
109
+
110
+ if c.weight_type.has_bias():
111
+ # GPTQv1 quirk: the kernel adds 1 to the stored zero, so we
112
+ # encode (bias - 1) here. See exllama.py for the link to the
113
+ # documentation of this checkpoint-format wart.
114
+ zeros = torch.full(
115
+ (groups, out_features),
116
+ c.weight_type.bias - 1,
117
+ dtype=torch.int32,
118
+ device=device,
119
+ )
120
+ else:
121
+ raise NotImplementedError(
122
+ "RDNA3 W4A16 kernel: zero-bias 4-bit quant requires "
123
+ "explicit zero points (GPTQv1 +1 quirk)."
124
+ )
125
+ zeros = pack_quantized_values_into_int32(zeros, c.weight_type, packed_dim=1)
126
+ setattr(
127
+ layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False)
128
+ )
129
+
130
+ # Act-order: convert g_idx to the inverse permutation array exllama
131
+ # expects (kernel reads a[perm[k]] instead of using groups indirected
132
+ # by g_idx[k]).
133
+ if c.has_g_idx:
134
+
135
+ def transform_w_g_idx(x):
136
+ return torch.argsort(x).to(torch.int)
137
+
138
+ self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore
139
+ else:
140
+ self.w_gidx_name = "g_idx"
141
+ empty_g_idx = torch.nn.Parameter(
142
+ torch.empty((0,), dtype=torch.int, device=device),
143
+ requires_grad=False,
144
+ )
145
+ setattr(layer, self.w_gidx_name, empty_g_idx)
146
+
147
+ def transform_w_q(x):
148
+ assert isinstance(x, BasevLLMParameter)
149
+ assert self.w_gidx_name is not None
150
+ g_idx = getattr(layer, self.w_gidx_name)
151
+
152
+ permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
153
+ x_cont = x.data.contiguous()
154
+ # Same 4-bit shuffle as exllama. The RDNA3 kernel reads weights in
155
+ # the same shuffled int32 layout and uses the (qa & 0x000F000F)
156
+ # bit-trick on top.
157
+ ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits)
158
+ return x_cont
159
+
160
+ def transform_w_s(x):
161
+ assert isinstance(x, BasevLLMParameter)
162
+ permute_param_layout_(x, input_dim=0, output_dim=1)
163
+ x.data = x.data.contiguous()
164
+ # Keep scales in the activation dtype (fp16 OR bf16) — the kernel
165
+ # branches on dtype internally.
166
+ return x.to(dtype=c.act_type)
167
+
168
+ self._transform_param(layer, self.w_q_name, transform_w_q)
169
+ self._transform_param(layer, self.w_s_name, transform_w_s)
170
+
171
+ # ----- Forward --------------------------------------------------------
172
+
173
+ def apply_weights(
174
+ self,
175
+ layer: torch.nn.Module,
176
+ x: torch.Tensor,
177
+ bias: torch.Tensor | None = None,
178
+ ) -> torch.Tensor:
179
+ c = self.config
180
+
181
+ x_2d = x.reshape(-1, x.shape[-1])
182
+ out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
183
+
184
+ w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer)
185
+
186
+ assert w_zp is not None, "Zero points are required by RDNA3 W4A16"
187
+ assert w_g_idx is not None, "g_idx tensor (possibly empty) required"
188
+
189
+ output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, w_g_idx, False)
190
+
191
+ if bias is not None:
192
+ output.add_(bias)
193
+ return output.reshape(out_shape)
vllm/platforms/rocm.py
+15
@@ -190,6 +190,9 @@ def _get_gcn_arch() -> str:
190
_GCN_ARCH = _get_gcn_arch()
191
192
_ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"])
193
+_ON_GFX11 = "gfx11" in _GCN_ARCH
194
+_ON_GFX1100 = "gfx1100" in _GCN_ARCH
195
+_ON_GFX1151 = "gfx1151" in _GCN_ARCH
196
_ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"])
197
_ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"])
198
_ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
@@ -273,6 +276,18 @@ def on_gfx1x() -> bool:
276
return _ON_GFX1X
277
278
279
+def on_gfx11() -> bool:
280
+ return _ON_GFX11
281
+
282
+
283
+def on_gfx1100() -> bool:
284
+ return _ON_GFX1100
285
+
286
+
287
+def on_gfx1151() -> bool:
288
+ return _ON_GFX1151
289
+
290
+
291
def on_gfx12x() -> bool:
292
return _ON_GFX12X
293