| 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 | // --------------------------------------------------------------------------- |
| 119 | // Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS). |
| 120 | // RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those |
| 121 | // landed on gfx940 / gfx1250 respectively), so this lowers to |
| 122 | // global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to |
| 123 | // write 4 output columns per row in a single atomic operation — half the |
| 124 | // atomic instruction count and half the contention vs two 32-bit CAS calls. |
| 125 | // |
| 126 | // Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer + |
| 127 | // cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue |
| 128 | // cast pass that an fp32-accumulator design would need. |
| 129 | // |
| 130 | // 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4 |
| 131 | // (always multiple of 4), and partition_weight_shape[1] is required to be a |
| 132 | // multiple of 8 by can_implement(), so every (m, n) write target is 8-byte |
| 133 | // aligned. Required by global_atomic_cmpswap_b64. |
| 134 | // --------------------------------------------------------------------------- |
| 135 | |
| 136 | __forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01, |
| 137 | half2 v23) { |
| 138 | unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr); |
| 139 | unsigned long long old = *addr_u; |
| 140 | while (true) { |
| 141 | union { |
| 142 | unsigned long long u; |
| 143 | half2 h2[2]; |
| 144 | } cur, sum; |
| 145 | cur.u = old; |
| 146 | sum.h2[0] = __hadd2(cur.h2[0], v01); |
| 147 | sum.h2[1] = __hadd2(cur.h2[1], v23); |
| 148 | unsigned long long prev = atomicCAS(addr_u, old, sum.u); |
| 149 | if (prev == old) break; |
| 150 | old = prev; |
| 151 | } |
| 152 | } |
| 153 | |
| 154 | __forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01, |
| 155 | bf162_t v23) { |
| 156 | unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr); |
| 157 | unsigned long long old = *addr_u; |
| 158 | while (true) { |
| 159 | union { |
| 160 | unsigned long long u; |
| 161 | bf162_t b2[2]; |
| 162 | } cur, sum; |
| 163 | cur.u = old; |
| 164 | sum.b2[0] = __hadd2(cur.b2[0], v01); |
| 165 | sum.b2[1] = __hadd2(cur.b2[1], v23); |
| 166 | unsigned long long prev = atomicCAS(addr_u, old, sum.u); |
| 167 | if (prev == old) break; |
| 168 | old = prev; |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | // Load one row's worth of 4 packed zeros (column n..n+3) from a [groups, N/8] |
| 173 | // uint32 tensor. n is a multiple of 4 by construction (n = offset_n + t*4 with |
| 174 | // offset_n = blockIdx.x * 512), so the 4 nibbles always live within one or two |
| 175 | // uint32 words; in practice within one because n & 7 is 0 or 4. |
| 176 | __forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n, |
| 177 | int (&zeros)[4]) { |
| 178 | int qcol = n / 8; |
| 179 | int shift = (n & 0x07) * 4; |
| 180 | uint32_t d = qzeros_row[qcol] >> shift; |
| 181 | zeros[0] = (int)(d & 0xF); |
| 182 | zeros[1] = (int)((d >> 4) & 0xF); |
| 183 | zeros[2] = (int)((d >> 8) & 0xF); |
| 184 | zeros[3] = (int)((d >> 12) & 0xF); |
| 185 | } |
| 186 | |
| 187 | template <typename T> |
| 188 | __forceinline__ __device__ void load4_scales(const T* scales_row, int n, |
| 189 | T (&scales)[4]) { |
| 190 | scales[0] = scales_row[n + 0]; |
| 191 | scales[1] = scales_row[n + 1]; |
| 192 | scales[2] = scales_row[n + 2]; |
| 193 | scales[3] = scales_row[n + 3]; |
| 194 | } |
| 195 | |
| 196 | // --------------------------------------------------------------------------- |
| 197 | // Main kernel. |
| 198 | // --------------------------------------------------------------------------- |
| 199 | |
| 200 | template <typename T, int M_COUNT> |
| 201 | __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, |
| 202 | const uint32_t* __restrict__ b_q_weight, |
| 203 | const uint32_t* __restrict__ b_qzeros, |
| 204 | const T* __restrict__ b_scales, |
| 205 | T* __restrict__ c, const int size_m, |
| 206 | const int size_n, const int size_k, |
| 207 | const int groups, const int zero_offset) { |
| 208 | const int t = threadIdx.x; |
| 209 | const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; |
| 210 | const int offset_m = blockIdx.y * M_COUNT; |
| 211 | const int offset_k = blockIdx.z * BLOCK_KN_SIZE; |
| 212 | const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); |
| 213 | const int n = offset_n + t * 4; |
| 214 | |
| 215 | // LDS layout: [M_COUNT][BLOCK_KN_SIZE + LDS_PAD]. The PAD=8 elements per M |
| 216 | // row break the natural 256-element/512-byte alignment that would otherwise |
| 217 | // collide on the same LDS bank when a thread reads block_a[0..M_COUNT-1][k] |
| 218 | // (same k, different m). Row stride becomes 264 elements * 2B = 528B = 132 |
| 219 | // 4-byte banks, so m-stride hits banks (m*132)%32 = (m*4)%32 — distinct for |
| 220 | // all M_COUNT ≤ 8. Cost: 16B LDS per block, irrelevant. |
| 221 | constexpr int LDS_PAD = 8; |
| 222 | __shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD]; |
| 223 | |
| 224 | // Stage A: each thread loads 1 K element per M row into LDS. |
| 225 | // THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map. |
| 226 | // For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m |
| 227 | // are zero-padded so the dot product contribution is 0 (we then skip the |
| 228 | // atomic write for those rows below). |
| 229 | // |
| 230 | // M=1 fast path: skip LDS staging + __syncthreads entirely. All 256 threads |
| 231 | // read the SAME 8-element A window per inner step (a_off is uniform across |
| 232 | // the block), so the cache-line broadcast through L1 makes global reads as |
| 233 | // cheap as LDS reads. Measured: ~1% on 4B b=1, ~6% on 27B b=1 in=128. |
| 234 | static_assert(BLOCK_KN_SIZE == THREADS_X, |
| 235 | "BLOCK_KN_SIZE must equal THREADS_X (1 K element per thread)"); |
| 236 | // The M=1 fast path (skip LDS) only has a global-read code path for bf16 |
| 237 | // (the v_dot2_f32_bf16 branch). The fp16 inner loop still indexes |
| 238 | // block_a[m][a_off] unconditionally, so for fp16 we MUST stage A through |
| 239 | // LDS even at M=1 to avoid reading uninitialized shared memory. |
| 240 | constexpr bool USE_LDS_A = (M_COUNT > 1) || std::is_same<T, half>::value; |
| 241 | if constexpr (USE_LDS_A) { |
| 242 | if (offset_k + t < end_k) { |
| 243 | #pragma unroll |
| 244 | for (int m = 0; m < M_COUNT; ++m) { |
| 245 | T av; |
| 246 | if (offset_m + m < size_m) { |
| 247 | const T* a_row = a + (offset_m + m) * size_k; |
| 248 | av = a_row[offset_k + t]; |
| 249 | } else { |
| 250 | av = tzero<T>(); // zero-pad invalid M rows |
| 251 | } |
| 252 | block_a[m][t] = av; |
| 253 | } |
| 254 | } |
| 255 | |
| 256 | // Threads beyond the right edge of N have nothing to do. Note: we must NOT |
| 257 | // return before __syncthreads() if any thread in the block participates in |
| 258 | // the LDS load above — but here all THREADS_X (=256) threads always do, |
| 259 | // regardless of whether their `n` is in bounds. |
| 260 | __syncthreads(); |
| 261 | } |
| 262 | if (n >= size_n) return; |
| 263 | |
| 264 | // Group bookkeeping. We require size_k % groups == 0 (groupsize divides K). |
| 265 | const int groupsize = size_k / groups; |
| 266 | int group = offset_k / groupsize; |
| 267 | int nextgroup = (group + 1) * groupsize; |
| 268 | |
| 269 | // qweight stride: weights are [K/8, N] uint32 with K packed at dim 0. |
| 270 | int qk = offset_k / 8; |
| 271 | const uint32_t* b_ptr = b_q_weight + qk * size_n + n; |
| 272 | |
| 273 | // Per-column dequant constants. We hold one set of (z, y) pairs per column. |
| 274 | // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper- |
| 275 | // nibble-*16 trick. bf16 uses fp32 scalars (z, y) for the inline v_dot2 |
| 276 | // bias correction. |
| 277 | half2 z1z16_h[4][2], y1y16_h[4][2]; |
| 278 | float z_b_f[4], y_b_f[4]; |
| 279 | |
| 280 | auto refresh_group = [&](int g) { |
| 281 | const uint32_t* qz_row = b_qzeros + g * (size_n / 8); |
| 282 | const T* sc_row = b_scales + g * size_n; |
| 283 | int zeros[4]; |
| 284 | T scales[4]; |
| 285 | load4_zeros(qz_row, n, zeros); |
| 286 | load4_scales<T>(sc_row, n, scales); |
| 287 | if constexpr (std::is_same<T, half>::value) { |
| 288 | #pragma unroll |
| 289 | for (int i = 0; i < 4; ++i) { |
| 290 | prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i], |
| 291 | z1z16_h[i], y1y16_h[i]); |
| 292 | } |
| 293 | } else { |
| 294 | #pragma unroll |
| 295 | for (int i = 0; i < 4; ++i) { |
| 296 | prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i], |
| 297 | z_b_f[i], y_b_f[i]); |
| 298 | } |
| 299 | } |
| 300 | }; |
| 301 | |
| 302 | refresh_group(group); |
| 303 | |
| 304 | float block_c[M_COUNT][4]; |
| 305 | #pragma unroll |
| 306 | for (int m = 0; m < M_COUNT; ++m) { |
| 307 | #pragma unroll |
| 308 | for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f; |
| 309 | } |
| 310 | |
| 311 | // Note on group-transition granularity: we check `k == nextgroup` at the |
| 312 | // start of each outer iteration (which advances K by 32). This is correct |
| 313 | // when group_size >= 32 OR group_size divides 32 evenly (groupsize is one |
| 314 | // of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the |
| 315 | // inner loop would cross a group boundary between j-iterations; we require |
| 316 | // group_size >= 32 here, mirroring exllama's assumption. |
| 317 | // |
| 318 | // Software pipelining: we issue all 4 vectorized weight loads up front |
| 319 | // before any dequant/FMA depends on them. This gives the AMDGPU backend |
| 320 | // freedom to schedule the global_loads early and overlap their latency |
| 321 | // with dequant + v_pk_fma_f16 of earlier iterations. Cost: 4×int4 = 16 |
| 322 | // VGPRs in flight per thread, plenty of headroom on RDNA3. |
| 323 | int k = offset_k; |
| 324 | while (k < end_k) { |
| 325 | if (k == nextgroup) { |
| 326 | group++; |
| 327 | nextgroup += groupsize; |
| 328 | refresh_group(group); |
| 329 | } |
| 330 | |
| 331 | // Prefetch all four j-iterations' weight words. The compiler emits 4 |
| 332 | // global_load_b128 instructions back-to-back; the dependent dequant + |
| 333 | // FMA work below hides their latency. |
| 334 | int4 b_w[4]; |
| 335 | #pragma unroll |
| 336 | for (int j = 0; j < 4; ++j) { |
| 337 | b_w[j] = *(const int4*)(b_ptr + j * size_n); |
| 338 | } |
| 339 | b_ptr += 4 * size_n; |
| 340 | |
| 341 | #pragma unroll |
| 342 | for (int j = 0; j < 4; ++j) { |
| 343 | const int a_off = (k - offset_k) + 8 * j; |
| 344 | |
| 345 | if constexpr (std::is_same<T, half>::value) { |
| 346 | half2 dq[4][4]; |
| 347 | dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]); |
| 348 | dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]); |
| 349 | dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]); |
| 350 | dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]); |
| 351 | |
| 352 | #pragma unroll |
| 353 | for (int m = 0; m < M_COUNT; ++m) { |
| 354 | const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]); |
| 355 | block_c[m][0] += dot22_8_f(dq[0], a_ptr); |
| 356 | block_c[m][1] += dot22_8_f(dq[1], a_ptr); |
| 357 | block_c[m][2] += dot22_8_f(dq[2], a_ptr); |
| 358 | block_c[m][3] += dot22_8_f(dq[3], a_ptr); |
| 359 | } |
| 360 | } else if constexpr (M_COUNT == 1) { |
| 361 | // bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of |
| 362 | // Hybrid PR #40977's wvSplitK_int4 kernel exactly so clang's |
| 363 | // InstCombine cannot fold the bf16→fp32 widening (LLVM #76000): |
| 364 | // * activations and magic-value weights share a fp32-aliased |
| 365 | // union (bytes written as uint32, read as bf16x2_t for the |
| 366 | // dot — pointer-cast opacity defeats the fold) |
| 367 | // * sum_a computed via a *second* v_dot2 with bf162(1,1) as the |
| 368 | // second operand, avoiding any explicit bf16→fp32 widen of A |
| 369 | // * bias correction y_b_f * partial + z_b_f * sum_a, identical |
| 370 | // to the previous fp32-FMA-chain path |
| 371 | // |
| 372 | // Net: 20 v_dot2_f32_bf16 + 8 fp32 FMA per int32 weight vs the |
| 373 | // previous 40 fp32 FMA. v_dot2 runs at full rate on gfx1100, so |
| 374 | // the substitution is ~2× cheaper for the inner accumulator. |
| 375 | typedef short __attribute__((ext_vector_type(2))) bf16x2_t; |
| 376 | constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128) |
| 377 | constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0) |
| 378 | union pack4 { |
| 379 | float f[4]; |
| 380 | uint32_t u[4]; |
| 381 | }; |
| 382 | |
| 383 | uint32_t w[4]; |
| 384 | __builtin_memcpy(w, &b_w[j], sizeof(int4)); |
| 385 | |
| 386 | // Load 8 bf16 activations as 4 uint32s (= 4 bf16x2 pairs) into a |
| 387 | // fp32-aliased union. Storing as uint32 keeps the IR-level type |
| 388 | // opaque so the inner v_dot2 cannot be folded to fp32 widening. |
| 389 | // |
| 390 | // A is read direct from global (no LDS staging — see USE_LDS_A above). |
| 391 | pack4 a_pack; |
| 392 | { |
| 393 | const uint32_t* a_words = |
| 394 | reinterpret_cast<const uint32_t*>(a + offset_k + a_off); |
| 395 | a_pack.u[0] = a_words[0]; |
| 396 | a_pack.u[1] = a_words[1]; |
| 397 | a_pack.u[2] = a_words[2]; |
| 398 | a_pack.u[3] = a_words[3]; |
| 399 | } |
| 400 | |
| 401 | // sum_a = Σ a[i]. Computed via 4× v_dot2_f32_bf16 with bf162(1,1) as |
| 402 | // the second operand — every bf16 pair contributes 1·a_lo + 1·a_hi. |
| 403 | // No fp32 widening of activations: the bytes go straight from LDS |
| 404 | // through v_dot2 into the fp32 accumulator. |
| 405 | float sum_a = 0.0f; |
| 406 | #pragma unroll |
| 407 | for (int b = 0; b < 4; ++b) { |
| 408 | sum_a = __builtin_amdgcn_fdot2_f32_bf16( |
| 409 | *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES), |
| 410 | sum_a, /*clamp=*/false); |
| 411 | } |
| 412 | |
| 413 | // unroll 1 keeps q_pack alive only one col at a time (8 fp32 VGPRs |
| 414 | // recycled across cols), avoiding straight-line expansion that |
| 415 | // would inflate live-range to 32 VGPRs. |
| 416 | #pragma unroll 1 |
| 417 | for (int col = 0; col < 4; ++col) { |
| 418 | // Build dequant magic values bf16(128 + nibble) directly into a |
| 419 | // fp32-aliased union via uint32 stores. No fp32 in the data flow |
| 420 | // until v_dot2 consumes the bytes. |
| 421 | pack4 q_pack; |
| 422 | const uint32_t qa = w[col]; |
| 423 | q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; |
| 424 | q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; |
| 425 | q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; |
| 426 | q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; |
| 427 | |
| 428 | // partial = Σ (128 + nibble[i]) · a[i], via 4× v_dot2_f32_bf16. |
| 429 | float partial = 0.0f; |
| 430 | #pragma unroll |
| 431 | for (int b = 0; b < 4; ++b) { |
| 432 | partial = __builtin_amdgcn_fdot2_f32_bf16( |
| 433 | *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])), |
| 434 | partial, /*clamp=*/false); |
| 435 | } |
| 436 | |
| 437 | // block_c += y_b_f * partial + z_b_f * sum_a |
| 438 | // y_b_f = scale, z_b_f = -(128+zero)*scale |
| 439 | // partial holds (128 + nibble) · a; subtracting (128+zero)·sum_a |
| 440 | // and scaling yields scale · (nibble - zero) · a as required. |
| 441 | block_c[0][col] = |
| 442 | __fmaf_rn(y_b_f[col], partial, |
| 443 | __fmaf_rn(z_b_f[col], sum_a, block_c[0][col])); |
| 444 | } |
| 445 | } else { |
| 446 | // bf16 M_COUNT > 1 path with v_dot2_f32_bf16. Same opacity trick as |
| 447 | // the M=1 branch: activations + magic-value weights stored in |
| 448 | // fp32-aliased unions, dot via __builtin_amdgcn_fdot2_f32_bf16 with |
| 449 | // pointer-cast to bf16x2_t. sum_a[m] computed via second v_dot2 |
| 450 | // with BF16_ONES; bias correction (y_b_f * partial + z_b_f * sum_a) |
| 451 | // applied after the dot. Magic values built once per col and reused |
| 452 | // across all M rows — amortizes dequant cost across M_COUNT. |
| 453 | typedef short __attribute__((ext_vector_type(2))) bf16x2_t; |
| 454 | constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128) |
| 455 | constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0) |
| 456 | union pack4 { |
| 457 | float f[4]; |
| 458 | uint32_t u[4]; |
| 459 | }; |
| 460 | |
| 461 | uint32_t w[4]; |
| 462 | __builtin_memcpy(w, &b_w[j], sizeof(int4)); |
| 463 | |
| 464 | // Load M_COUNT × 8 bf16 activations as 4 uint32s each into pack4 |
| 465 | // unions. Stored as uint32 to keep IR-level types opaque (defeats |
| 466 | // InstCombine fold). At M_COUNT=8 this is 32 fp32 VGPRs — within RDNA3 |
| 467 | // budget. |
| 468 | pack4 a_pack[M_COUNT]; |
| 469 | #pragma unroll |
| 470 | for (int m = 0; m < M_COUNT; ++m) { |
| 471 | const uint32_t* a_words = |
| 472 | reinterpret_cast<const uint32_t*>(&block_a[m][a_off]); |
| 473 | a_pack[m].u[0] = a_words[0]; |
| 474 | a_pack[m].u[1] = a_words[1]; |
| 475 | a_pack[m].u[2] = a_words[2]; |
| 476 | a_pack[m].u[3] = a_words[3]; |
| 477 | } |
| 478 | |
| 479 | // sum_a[m] = Σ a[m][i] via 4× v_dot2 with bf162(1,1) — no fp32 widen. |
| 480 | float sum_a[M_COUNT]; |
| 481 | #pragma unroll |
| 482 | for (int m = 0; m < M_COUNT; ++m) { |
| 483 | float s = 0.0f; |
| 484 | #pragma unroll |
| 485 | for (int b = 0; b < 4; ++b) { |
| 486 | s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])), |
| 487 | *((const bf16x2_t*)&BF16_ONES), |
| 488 | s, /*clamp=*/false); |
| 489 | } |
| 490 | sum_a[m] = s; |
| 491 | } |
| 492 | |
| 493 | // Per col: build magic-value pack, dot against all M activations. |
| 494 | // unroll 1 keeps q_pack live one col at a time (8 fp32 VGPRs recycled) |
| 495 | // — same register-pressure trick as the previous fp32 path. |
| 496 | #pragma unroll 1 |
| 497 | for (int col = 0; col < 4; ++col) { |
| 498 | pack4 q_pack; |
| 499 | const uint32_t qa = w[col]; |
| 500 | q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; |
| 501 | q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; |
| 502 | q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; |
| 503 | q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; |
| 504 | |
| 505 | #pragma unroll |
| 506 | for (int m = 0; m < M_COUNT; ++m) { |
| 507 | float partial = 0.0f; |
| 508 | #pragma unroll |
| 509 | for (int b = 0; b < 4; ++b) { |
| 510 | partial = __builtin_amdgcn_fdot2_f32_bf16( |
| 511 | *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])), |
| 512 | partial, /*clamp=*/false); |
| 513 | } |
| 514 | // block_c += y_b_f * partial + z_b_f * sum_a (same correction as |
| 515 | // M=1) |
| 516 | block_c[m][col] = |
| 517 | __fmaf_rn(y_b_f[col], partial, |
| 518 | __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col])); |
| 519 | } |
| 520 | } |
| 521 | } |
| 522 | } |
| 523 | k += 32; // 4 weight words * 8 nibbles = 32 K elements |
| 524 | } |
| 525 | |
| 526 | // Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all |
| 527 | // four lanes in a single 64-bit CAS write directly to the T-typed output |
| 528 | // (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with |
| 529 | // a single b64 op we halve the atomic instruction count vs two b32 CAS |
| 530 | // calls, AND save the FP32 buffer + memset + cast pass entirely. |
| 531 | #pragma unroll |
| 532 | for (int m = 0; m < M_COUNT; ++m) { |
| 533 | if (offset_m + m >= size_m) continue; // skip padding rows past size_m |
| 534 | T* out = c + (offset_m + m) * size_n + n; |
| 535 | if constexpr (std::is_same<T, half>::value) { |
| 536 | half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]), |
| 537 | __float2half_rn(block_c[m][1])); |
| 538 | half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]), |
| 539 | __float2half_rn(block_c[m][3])); |
| 540 | atomic_add_pk4_f16(out, r01, r23); |
| 541 | } else { |
| 542 | bf162_t r01; |
| 543 | r01.x = __float2bfloat16(block_c[m][0]); |
| 544 | r01.y = __float2bfloat16(block_c[m][1]); |
| 545 | bf162_t r23; |
| 546 | r23.x = __float2bfloat16(block_c[m][2]); |
| 547 | r23.y = __float2bfloat16(block_c[m][3]); |
| 548 | atomic_add_pk4_bf16(out, r01, r23); |
| 549 | } |
| 550 | } |
| 551 | } |
| 552 | |
| 553 | #else // non-RDNA3 device pass: empty __global__ for symbol parity. |
| 554 | |
| 555 | template <typename T, int M_COUNT> |
| 556 | __global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*, |
| 557 | const T*, T*, const int, const int, |
| 558 | const int, const int, const int) {} |
| 559 | |
| 560 | #endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__ |
| 561 | |
| 562 | // --------------------------------------------------------------------------- |
| 563 | // Launcher. |
| 564 | // --------------------------------------------------------------------------- |
| 565 | |
| 566 | template <typename T, int M_COUNT> |
| 567 | void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight, |
| 568 | const uint32_t* b_qzeros, const T* b_scales, |
| 569 | T* c, int size_m, int size_n, int size_k, |
| 570 | int groups, int zero_offset, |
| 571 | cudaStream_t stream) { |
| 572 | dim3 block(THREADS_X); |
| 573 | dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4), |
| 574 | (size_m + M_COUNT - 1) / M_COUNT, |
| 575 | (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE); |
| 576 | |
| 577 | gemm_q4_kernel_rdna3<T, M_COUNT> |
| 578 | <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, |
| 579 | size_n, size_k, groups, zero_offset); |
| 580 | } |
| 581 | |
| 582 | // Dispatch to the largest M_COUNT template that doesn't waste more than |
| 583 | // half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is |
| 584 | // the right tool, not bigger M_COUNT in the scalar dot-product path. |
| 585 | // |
| 586 | // Tile-waste table: |
| 587 | // M=1 -> M_COUNT=1 (no waste) |
| 588 | // M=2,3 -> M_COUNT=2 (M=3 wastes 1/2 of last tile) |
| 589 | // M=4-7 -> M_COUNT=4 (worst case M=5: wastes 3/4 of last tile) |
| 590 | // M=8-15-> M_COUNT=8 (worst case M=9: wastes 7/8 of last tile) |
| 591 | // "Wasted" rows are zero-padded in LDS and skip the atomic write, so they |
| 592 | // only burn instructions on the last block, never affect correctness. |
| 593 | template <typename T> |
| 594 | void launch_gemm_q4(const T* a, const uint32_t* b_q_weight, |
| 595 | const uint32_t* b_qzeros, const T* b_scales, T* c, |
| 596 | int size_m, int size_n, int size_k, int groups, |
| 597 | bool use_v2_format, cudaStream_t stream) { |
| 598 | const int zero_offset = use_v2_format ? 0 : 1; |
| 599 | |
| 600 | if (size_m == 1) { |
| 601 | launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, c, |
| 602 | size_m, size_n, size_k, groups, zero_offset, |
| 603 | stream); |
| 604 | } else if (size_m <= 3) { |
| 605 | launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, c, |
| 606 | size_m, size_n, size_k, groups, zero_offset, |
| 607 | stream); |
| 608 | } else if (size_m <= 7) { |
| 609 | launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, c, |
| 610 | size_m, size_n, size_k, groups, zero_offset, |
| 611 | stream); |
| 612 | } else { |
| 613 | // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the |
| 614 | // WMMA path, but if it falls through we still produce correct output — |
| 615 | // just leaving 3-5× of throughput on the table for prefill workloads. |
| 616 | launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, c, |
| 617 | size_m, size_n, size_k, groups, zero_offset, |
| 618 | stream); |
| 619 | } |
| 620 | } |
| 621 | |
| 622 | } // namespace gptq_rdna3 |
| 623 | } // namespace vllm |
| 624 | |
| 625 | // --------------------------------------------------------------------------- |
| 626 | // Public entry point. |
| 627 | // --------------------------------------------------------------------------- |
| 628 | // |
| 629 | // Inputs: |
| 630 | // a [M, K] half or bfloat16 |
| 631 | // b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle) |
| 632 | // b_qzeros [groups, N/8] uint32 (packed 4-bit zeros) |
| 633 | // b_scales [groups, N] half or bfloat16 |
| 634 | // use_v2_format bool (true = GPTQv2, no +1 zero offset) |
| 635 | // |
| 636 | // Output: |
| 637 | // c [M, N] same dtype as a |
| 638 | |
| 639 | torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, |
| 640 | torch::Tensor b_qzeros, |
| 641 | torch::Tensor b_scales, bool use_v2_format); |
| 642 | |
| 643 | torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, |
| 644 | torch::Tensor b_qzeros, torch::Tensor b_scales, |
| 645 | bool use_v2_format) { |
| 646 | if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 && |
| 647 | b_q_weight.size(1) % 16 == 0 && |
| 648 | ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) || |
| 649 | (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) { |
| 650 | return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, |
| 651 | use_v2_format); |
| 652 | } |
| 653 | |
| 654 | TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); |
| 655 | TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); |
| 656 | TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor"); |
| 657 | TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor"); |
| 658 | TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]"); |
| 659 | TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]"); |
| 660 | TORCH_CHECK( |
| 661 | a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, |
| 662 | "a must be half or bfloat16"); |
| 663 | TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(), |
| 664 | "b_scales dtype must match a"); |
| 665 | |
| 666 | const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); |
| 667 | auto stream = at::cuda::getCurrentCUDAStream(); |
| 668 | |
| 669 | int size_m = (int)a.size(0); |
| 670 | int size_k = (int)a.size(1); |
| 671 | int size_n = (int)b_q_weight.size(1); |
| 672 | int groups = (int)b_qzeros.size(0); |
| 673 | |
| 674 | TORCH_CHECK(b_q_weight.size(0) * 8 == size_k, |
| 675 | "b_q_weight first dim must be K/8"); |
| 676 | TORCH_CHECK(b_scales.size(0) == groups, |
| 677 | "b_scales must have same group count as qzeros"); |
| 678 | TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N"); |
| 679 | TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)"); |
| 680 | |
| 681 | auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); |
| 682 | at::Tensor c = torch::zeros({size_m, size_n}, opts); |
| 683 | |
| 684 | if (a.scalar_type() == torch::kHalf) { |
| 685 | vllm::gptq_rdna3::launch_gemm_q4<half>( |
| 686 | (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), |
| 687 | (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(), |
| 688 | (half*)c.data_ptr(), size_m, size_n, size_k, groups, use_v2_format, |
| 689 | stream); |
| 690 | } else { |
| 691 | vllm::gptq_rdna3::launch_gemm_q4<vllm::gptq_rdna3::bf16_t>( |
| 692 | (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(), |
| 693 | (const uint32_t*)b_q_weight.data_ptr(), |
| 694 | (const uint32_t*)b_qzeros.data_ptr(), |
| 695 | (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), |
| 696 | (vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups, |
| 697 | use_v2_format, stream); |
| 698 | } |
| 699 | |
| 700 | return c; |
| 701 | } |