From 0cff0741ff8853beb2cae862c0daa8133cb1aa43 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 29 May 2026 13:04:40 +0200 Subject: [PATCH] =?UTF-8?q?[Kernel][ROCm]=20Native=20W4A16=20kernel=20for?= =?UTF-8?q?=20AMD=20RDNA3=20(gfx1100)=20=E2=80=94=20fp16=20+=20bf16=20(#41?= =?UTF-8?q?394)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: JartX --- CMakeLists.txt | 12 + csrc/rocm/ops.h | 9 + csrc/rocm/q_gemm_rdna3.cu | 780 ++++++ csrc/rocm/q_gemm_rdna3_wmma.cu | 2165 +++++++++++++++++ csrc/rocm/qdq_4_rdna3.cuh | 239 ++ csrc/rocm/torch_bindings.cpp | 13 + .../kernels/quantization/test_rdna3_w4a16.py | 278 +++ .../test_rdna3_w4a16_selection.py | 89 + vllm/_custom_ops.py | 45 + .../model_executor/kernels/linear/__init__.py | 4 + .../linear/mixed_precision/__init__.py | 4 + .../linear/mixed_precision/rdna3_w4a16.py | 193 ++ vllm/platforms/rocm.py | 15 + 13 files changed, 3846 insertions(+) create mode 100644 csrc/rocm/q_gemm_rdna3.cu create mode 100644 csrc/rocm/q_gemm_rdna3_wmma.cu create mode 100644 csrc/rocm/qdq_4_rdna3.cuh create mode 100644 tests/kernels/quantization/test_rdna3_w4a16.py create mode 100644 tests/kernels/quantization/test_rdna3_w4a16_selection.py create mode 100644 vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 5a3f95faee0..3f571170401 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1284,6 +1284,14 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/skinny_gemms.cu" "csrc/rocm/attention.cu") + set(VLLM_ROCM_HAS_GFX1100 OFF) + if(VLLM_GPU_ARCHES MATCHES "gfx1100") + set(VLLM_ROCM_HAS_GFX1100 ON) + list(APPEND VLLM_ROCM_EXT_SRC + "csrc/rocm/q_gemm_rdna3.cu" + "csrc/rocm/q_gemm_rdna3_wmma.cu") + endif() + define_extension_target( _rocm_C DESTINATION vllm @@ -1293,6 +1301,10 @@ if(VLLM_GPU_LANG STREQUAL "HIP") ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) + + if(VLLM_ROCM_HAS_GFX1100) + target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100) + endif() endif() # Must run after the last HIP `define_extension_target` so every extension diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index dbc466f036e..73197d8a5e2 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -18,6 +18,15 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b, const at::Tensor& scale_a, const at::Tensor& scale_b, const int64_t CuCount); +torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format); + +torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, + torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format); + void paged_attention( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu new file mode 100644 index 00000000000..fb178f6d3b1 --- /dev/null +++ b/csrc/rocm/q_gemm_rdna3.cu @@ -0,0 +1,780 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the +// activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit +// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes: +// +// 1. Direct write to the T-typed output via packed CAS-loop on a 64-bit +// word (atomic_add_pk4_{f16,bf16}). gfx11 has no native +// v_global_atomic_pk_add_{f16,bf16}, so the kernel emulates one with +// global_atomic_cmpswap_b64. This avoids the M*N*4-byte FP32 scratch +// buffer + memset + cast-pass that an fp32-accumulator design would +// need; the caller passes a zero-initialised T-typed output tensor +// and every block atomically adds its partial sum into it. +// +// 2. The bf16 path uses a dedicated bit-trick that avoids the fp16-only +// "upper nibble * 16" trick, which would overflow the 7-bit bf16 +// mantissa. See qdq_4_rdna3.cuh for details. +// +// 3. Wave32 geometry sized for high CU saturation: THREADS_X=256 +// (8 waves per block) and BLOCK_KN_SIZE=256, with each thread +// computing 4 N output columns. gridDim.z = K / BLOCK_KN_SIZE +// splits K and the output is atomically accumulated. fp16 uses +// v_dot2_f32_f16 (__builtin_amdgcn_fdot2) for the inner dot; +// bf16 widens to fp32 (no v_pk_fma_bf16 on gfx11) and accumulates +// with v_fma_f32. M_COUNT ∈ {1,2,4,8} is selected at launch +// based on size_m. +// +// 4. The bf16 dispatch with M >= 16 forwards to the WMMA kernel in +// q_gemm_rdna3_wmma.cu (separate translation unit) where +// v_wmma_f32_16x16x16_bf16_w32 wins. The fp16 path always stays +// scalar (the bit-trick dequant beats WMMA below M=64). + +#include +#include + +#include +#include +#include + +#include +#include +#include + +#include "qdq_4_rdna3.cuh" + +#if defined(__HIPCC__) && defined(__gfx1100__) + #define __HIP__RDNA3__ +#endif + +namespace vllm { +namespace gptq_rdna3 { + +// BLOCK_KN_SIZE = 256 (was 128 in exllama). Each block covers 256 K +// elements and THREADS_X*4 = 1024 N columns. For Qwen-class K=4096 this +// halves gridDim.z (32 → 16) and therefore halves the atomic count per +// output position vs the exllama default. THREADS_X=256 = 8 waves on RDNA3 +// wave32; with ~32 wave slots per CU we still fit 4 blocks per CU at peak. +// +// We tried BLOCK_KN_SIZE=512 (microbench on Qwen3.6-27B): bf16 improved +// 5-10% at large M (atomic CAS halved), but fp16 decode regressed up to +// +40% on qkv-square (32 → 45 μs at M=1). Cause: 16 waves/block × 16 +// total blocks for [M=1, K=N=4096] only saturates ~8 of the 96 CUs, +// breaking memory-latency hiding for the fp16 path which is already +// memory-bound. Reverted to 256; bf16 keeps most of its gains from the +// fp32 dequant rewrite alone. +#define BLOCK_KN_SIZE 256 +#define THREADS_X 256 + +// Device code below is RDNA3-only; non-RDNA3 device passes fall through to +// the empty __global__ stub at the #else below for symbol parity. +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// --------------------------------------------------------------------------- +// Per-dtype helpers. We avoid heavy template metaprogramming and just provide +// overloaded inline functions; the kernel below selects via `if constexpr`. +// --------------------------------------------------------------------------- + +// Type-generic zero — both half and bf16_t in HIP/ROCm have a converting +// constructor from float, but going through __float2half_rn / __float2bfloat16 +// is the unambiguously correct path on every ROCm version. +template +__forceinline__ __device__ T tzero(); + +template <> +__forceinline__ __device__ half tzero() { + return __float2half_rn(0.0f); +} + +template <> +__forceinline__ __device__ bf16_t tzero() { + return __float2bfloat16(0.0f); +} + +__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { + // RDNA3 has v_dot2_f32_f16 (`__builtin_amdgcn_fdot2`) which computes + // fp32 += a.x*b.x + a.y*b.y in a single instruction with the accumulator + // staying in fp32 throughout. hipcc 7.2 does NOT peephole the obvious + // `__hfma2 + cast + add` pattern into v_dot2 (verified by ISA + // disassembly: 0 v_dot2_f32_f16 vs 256 v_cvt_f32_f16 + 218 v_add_f32 in + // the M_COUNT=8 kernel before this change), so we issue the builtin + // explicitly. Saves the trailing 2× v_cvt_f32_f16 + v_add_f32 (3 ops) + // per dot22_8_f call vs the half2-accumulator form. With 128 calls per + // K=32 step that's ~384 ops/K-step less issue pressure on the VALU. + // + // Numerical bonus: accumulator stays fp32 throughout the dot. The old + // form accumulated 8 muladds in fp16 (10-bit mantissa) before casting, + // which could lose ~3 bits of precision on borderline magnitudes. + float result = 0.0f; + const half2* a2_ptr = (const half2*)a_ptr; + #pragma unroll + for (int i = 0; i < 4; i++) { + result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false); + } + return result; +} + +__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4], + const bf16_t* a_ptr) { + // RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in + // the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc + // therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised + // fallback (single-element FMAs or fp32 round-trips), which empirically + // runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16 + // decode path was paying that tax in full, scaling linearly with M (the + // fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate + // and the kernel becomes memory-bound). + // + // Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs) + // and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA + // count, but each FMA is fast. Bonus: the accumulator is now fp32 + // throughout instead of bf16, which is also numerically more accurate + // (no compounding bf16-rounding inside the dot loop). + float result = 0.0f; + #pragma unroll + for (int i = 0; i < 4; i++) { + uint32_t aw, dw; + __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t)); + __builtin_memcpy(&dw, &dq[i], sizeof(uint32_t)); + // bf16 in low 16 bits → fp32 by left-shifting into the upper half. + // bf16 in high 16 bits → already aligned with fp32's upper half. + float a_x = __uint_as_float((aw & 0xFFFFu) << 16); + float a_y = __uint_as_float(aw & 0xFFFF0000u); + float d_x = __uint_as_float((dw & 0xFFFFu) << 16); + float d_y = __uint_as_float(dw & 0xFFFF0000u); + result = __fmaf_rn(d_x, a_x, result); + result = __fmaf_rn(d_y, a_y, result); + } + return result; +} + +// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already +// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t +// overload above does for dq (still need to widen A from bf16). Wins more +// at high N: the bf162_t version's per-call widening cost scales with the +// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only +// for A widening (M_COUNT × 4 × 4 widens, half as many). +__forceinline__ __device__ float dot22_8_f(float (&dq)[8], + const bf16_t* a_ptr) { + float result = 0.0f; + #pragma unroll + for (int i = 0; i < 4; i++) { + uint32_t aw; + __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t)); + float a_x = __uint_as_float((aw & 0xFFFFu) << 16); + float a_y = __uint_as_float(aw & 0xFFFF0000u); + result = __fmaf_rn(dq[2 * i + 0], a_x, result); + result = __fmaf_rn(dq[2 * i + 1], a_y, result); + } + return result; +} + +// --------------------------------------------------------------------------- +// Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS). +// RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those +// landed on gfx940 / gfx1250 respectively), so this lowers to +// global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to +// write 4 output columns per row in a single atomic operation — half the +// atomic instruction count and half the contention vs two 32-bit CAS calls. +// +// Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer + +// cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue +// cast pass that an fp32-accumulator design would need. +// +// 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4 +// (always multiple of 4), and partition_weight_shape[1] is required to be a +// multiple of 8 by can_implement(), so every (m, n) write target is 8-byte +// aligned. Required by global_atomic_cmpswap_b64. +// --------------------------------------------------------------------------- + +__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01, + half2 v23) { + unsigned long long* addr_u = reinterpret_cast(addr); + unsigned long long old = *addr_u; + while (true) { + union { + unsigned long long u; + half2 h2[2]; + } cur, sum; + cur.u = old; + sum.h2[0] = __hadd2(cur.h2[0], v01); + sum.h2[1] = __hadd2(cur.h2[1], v23); + unsigned long long prev = atomicCAS(addr_u, old, sum.u); + if (prev == old) break; + old = prev; + } +} + +__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01, + bf162_t v23) { + unsigned long long* addr_u = reinterpret_cast(addr); + unsigned long long old = *addr_u; + while (true) { + union { + unsigned long long u; + bf162_t b2[2]; + } cur, sum; + cur.u = old; + sum.b2[0] = __hadd2(cur.b2[0], v01); + sum.b2[1] = __hadd2(cur.b2[1], v23); + unsigned long long prev = atomicCAS(addr_u, old, sum.u); + if (prev == old) break; + old = prev; + } +} + +// Load one row's worth of 4 packed zeros (column n..n+3) from a [groups, N/8] +// uint32 tensor. n is a multiple of 4 by construction (n = offset_n + t*4 with +// offset_n = blockIdx.x * 512), so the 4 nibbles always live within one or two +// uint32 words; in practice within one because n & 7 is 0 or 4. +__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n, + int (&zeros)[4]) { + int qcol = n / 8; + int shift = (n & 0x07) * 4; + uint32_t d = qzeros_row[qcol] >> shift; + zeros[0] = (int)(d & 0xF); + zeros[1] = (int)((d >> 4) & 0xF); + zeros[2] = (int)((d >> 8) & 0xF); + zeros[3] = (int)((d >> 12) & 0xF); +} + +template +__forceinline__ __device__ void load4_scales(const T* scales_row, int n, + T (&scales)[4]) { + scales[0] = scales_row[n + 0]; + scales[1] = scales_row[n + 1]; + scales[2] = scales_row[n + 2]; + scales[3] = scales_row[n + 3]; +} + +// --------------------------------------------------------------------------- +// Main kernel. +// --------------------------------------------------------------------------- + +template +__global__ void gemm_q4_kernel_rdna3( + const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int t = threadIdx.x; + const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; + const int offset_m = blockIdx.y * M_COUNT; + const int offset_k = blockIdx.z * BLOCK_KN_SIZE; + const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); + const int n = offset_n + t * 4; + + // LDS layout: [M_COUNT][BLOCK_KN_SIZE + LDS_PAD]. The PAD=8 elements per M + // row break the natural 256-element/512-byte alignment that would otherwise + // collide on the same LDS bank when a thread reads block_a[0..M_COUNT-1][k] + // (same k, different m). Row stride becomes 264 elements * 2B = 528B = 132 + // 4-byte banks, so m-stride hits banks (m*132)%32 = (m*4)%32 — distinct for + // all M_COUNT ≤ 8. Cost: 16B LDS per block, irrelevant. + constexpr int LDS_PAD = 8; + __shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD]; + + // Stage A: each thread loads 1 K element per M row into LDS (with optional + // act-order permutation). THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map. + // For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m + // are zero-padded so the dot product contribution is 0 (we then skip the + // atomic write for those rows below). + // + // M=1 fast path: skip LDS staging + __syncthreads entirely. All 256 threads + // read the SAME 8-element A window per inner step (a_off is uniform across + // the block), so the cache-line broadcast through L1 makes global reads as + // cheap as LDS reads. Measured: ~1% on 4B b=1, ~6% on 27B b=1 in=128. + static_assert(BLOCK_KN_SIZE == THREADS_X, + "BLOCK_KN_SIZE must equal THREADS_X (1 K element per thread)"); + // The M=1 fast path (skip LDS) only has a global-read code path for bf16 + // (the v_dot2_f32_bf16 branch). The fp16 inner loop still indexes + // block_a[m][a_off] unconditionally, so for fp16 we MUST stage A through + // LDS even at M=1 to avoid reading uninitialized shared memory. + constexpr bool USE_LDS_A = (M_COUNT > 1) || std::is_same::value; + if constexpr (USE_LDS_A) { + if (offset_k + t < end_k) { + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + T av; + if (offset_m + m < size_m) { + const T* a_row = a + (offset_m + m) * size_k; + if (b_q_perm) + av = a_row[b_q_perm[offset_k + t]]; + else + av = a_row[offset_k + t]; + } else { + av = tzero(); // zero-pad invalid M rows + } + block_a[m][t] = av; + } + } + + // Threads beyond the right edge of N have nothing to do. Note: we must NOT + // return before __syncthreads() if any thread in the block participates in + // the LDS load above — but here all THREADS_X (=256) threads always do, + // regardless of whether their `n` is in bounds. + __syncthreads(); + } else if (b_q_perm) { + // bf16 M=1 fast path skips LDS, but its global read below is sequential + // and cannot apply act-order. When a permutation is present, stage the + // single A row through LDS (as fp16 / M>1 do) so the read picks it up. + // b_q_perm is block-uniform, so the __syncthreads is non-divergent. + if (offset_k + t < end_k) + block_a[0][t] = a[offset_m * size_k + b_q_perm[offset_k + t]]; + __syncthreads(); + } + if (n >= size_n) return; + + // Group bookkeeping. We require size_k % groups == 0 (groupsize divides K). + const int groupsize = size_k / groups; + int group = offset_k / groupsize; + int nextgroup = (group + 1) * groupsize; + + // qweight stride: weights are [K/8, N] uint32 with K packed at dim 0. + int qk = offset_k / 8; + const uint32_t* b_ptr = b_q_weight + qk * size_n + n; + + // Per-column dequant constants. We hold one set of (z, y) pairs per column. + // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper- + // nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant + // produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA + // bypass for the missing v_pk_fma_bf16 on gfx11. + half2 z1z16_h[4][2], y1y16_h[4][2]; + float z_b_f[4], y_b_f[4]; + + auto refresh_group = [&](int g) { + const uint32_t* qz_row = b_qzeros + g * (size_n / 8); + const T* sc_row = b_scales + g * size_n; + int zeros[4]; + T scales[4]; + load4_zeros(qz_row, n, zeros); + load4_scales(sc_row, n, scales); + if constexpr (std::is_same::value) { + #pragma unroll + for (int i = 0; i < 4; ++i) { + prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i], + z1z16_h[i], y1y16_h[i]); + } + } else { + #pragma unroll + for (int i = 0; i < 4; ++i) { + prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i], + z_b_f[i], y_b_f[i]); + } + } + }; + + refresh_group(group); + + float block_c[M_COUNT][4]; + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + #pragma unroll + for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f; + } + + // Note on group-transition granularity: we check `k == nextgroup` at the + // start of each outer iteration (which advances K by 32). This is correct + // when group_size >= 32 OR group_size divides 32 evenly (groupsize is one + // of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the + // inner loop would cross a group boundary between j-iterations; we require + // group_size >= 32 here, mirroring exllama's assumption. + // + // Software pipelining: we issue all 4 vectorized weight loads up front + // before any dequant/FMA depends on them. This gives the AMDGPU backend + // freedom to schedule the global_loads early and overlap their latency + // with dequant + v_pk_fma_f16 of earlier iterations. Cost: 4×int4 = 16 + // VGPRs in flight per thread, plenty of headroom on RDNA3. + int k = offset_k; + while (k < end_k) { + if (k == nextgroup) { + group++; + nextgroup += groupsize; + refresh_group(group); + } + + // Prefetch all four j-iterations' weight words. The compiler emits 4 + // global_load_b128 instructions back-to-back; the dependent dequant + + // FMA work below hides their latency. + int4 b_w[4]; + #pragma unroll + for (int j = 0; j < 4; ++j) { + b_w[j] = *(const int4*)(b_ptr + j * size_n); + } + b_ptr += 4 * size_n; + + #pragma unroll + for (int j = 0; j < 4; ++j) { + const int a_off = (k - offset_k) + 8 * j; + + if constexpr (std::is_same::value) { + half2 dq[4][4]; + dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]); + dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]); + dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]); + dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]); + + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + const half* a_ptr = reinterpret_cast(&block_a[m][a_off]); + block_c[m][0] += dot22_8_f(dq[0], a_ptr); + block_c[m][1] += dot22_8_f(dq[1], a_ptr); + block_c[m][2] += dot22_8_f(dq[2], a_ptr); + block_c[m][3] += dot22_8_f(dq[3], a_ptr); + } + } else if constexpr (M_COUNT == 1) { + // bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of + // Hybrid PR #40977's wvSplitK_int4 kernel exactly so clang's + // InstCombine cannot fold the bf16→fp32 widening (LLVM #76000): + // * activations and magic-value weights share a fp32-aliased + // union (bytes written as uint32, read as bf16x2_t for the + // dot — pointer-cast opacity defeats the fold) + // * sum_a computed via a *second* v_dot2 with bf162(1,1) as the + // second operand, avoiding any explicit bf16→fp32 widen of A + // * bias correction y_b_f * partial + z_b_f * sum_a, identical + // to the previous fp32-FMA-chain path + // + // Net: 20 v_dot2_f32_bf16 + 8 fp32 FMA per int32 weight vs the + // previous 40 fp32 FMA. v_dot2 runs at full rate on gfx1100, so + // the substitution is ~2× cheaper for the inner accumulator. + typedef short __attribute__((ext_vector_type(2))) bf16x2_t; + constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128) + constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0) + union pack4 { + float f[4]; + uint32_t u[4]; + }; + + uint32_t w[4]; + __builtin_memcpy(w, &b_w[j], sizeof(int4)); + + // Load 8 bf16 activations as 4 uint32s (= 4 bf16x2 pairs) into a + // fp32-aliased union. Storing as uint32 keeps the IR-level type + // opaque so the inner v_dot2 cannot be folded to fp32 widening. + // + // A is read direct from global (no LDS staging — see USE_LDS_A above), + // except under act-order, where it comes from the permuted LDS copy. + pack4 a_pack; + { + const uint32_t* a_words = + b_q_perm + ? reinterpret_cast(&block_a[0][a_off]) + : reinterpret_cast(a + offset_k + a_off); + a_pack.u[0] = a_words[0]; + a_pack.u[1] = a_words[1]; + a_pack.u[2] = a_words[2]; + a_pack.u[3] = a_words[3]; + } + + // sum_a = Σ a[i]. Computed via 4× v_dot2_f32_bf16 with bf162(1,1) as + // the second operand — every bf16 pair contributes 1·a_lo + 1·a_hi. + // No fp32 widening of activations: the bytes go straight from LDS + // through v_dot2 into the fp32 accumulator. + float sum_a = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + sum_a = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES), + sum_a, /*clamp=*/false); + } + + // unroll 1 keeps q_pack alive only one col at a time (8 fp32 VGPRs + // recycled across cols), avoiding straight-line expansion that + // would inflate live-range to 32 VGPRs. + #pragma unroll 1 + for (int col = 0; col < 4; ++col) { + // Build dequant magic values bf16(128 + nibble) directly into a + // fp32-aliased union via uint32 stores. No fp32 in the data flow + // until v_dot2 consumes the bytes. + pack4 q_pack; + const uint32_t qa = w[col]; + q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; + + // partial = Σ (128 + nibble[i]) · a[i], via 4× v_dot2_f32_bf16. + float partial = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + partial = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])), + partial, /*clamp=*/false); + } + + // block_c += y_b_f * partial + z_b_f * sum_a + // y_b_f = scale, z_b_f = -(128+zero)*scale + // partial holds (128 + nibble) · a; subtracting (128+zero)·sum_a + // and scaling yields scale · (nibble - zero) · a as required. + block_c[0][col] = + __fmaf_rn(y_b_f[col], partial, + __fmaf_rn(z_b_f[col], sum_a, block_c[0][col])); + } + } else { + // bf16 M_COUNT > 1 path with v_dot2_f32_bf16. Same opacity trick as + // the M=1 branch: activations + magic-value weights stored in + // fp32-aliased unions, dot via __builtin_amdgcn_fdot2_f32_bf16 with + // pointer-cast to bf16x2_t. sum_a[m] computed via second v_dot2 + // with BF16_ONES; bias correction (y_b_f * partial + z_b_f * sum_a) + // applied after the dot. Magic values built once per col and reused + // across all M rows — amortizes dequant cost across M_COUNT. + typedef short __attribute__((ext_vector_type(2))) bf16x2_t; + constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128) + constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0) + union pack4 { + float f[4]; + uint32_t u[4]; + }; + + uint32_t w[4]; + __builtin_memcpy(w, &b_w[j], sizeof(int4)); + + // Load M_COUNT × 8 bf16 activations as 4 uint32s each into pack4 + // unions. Stored as uint32 to keep IR-level types opaque (defeats + // InstCombine fold). At M_COUNT=8 this is 32 fp32 VGPRs — within RDNA3 + // budget. + pack4 a_pack[M_COUNT]; + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + const uint32_t* a_words = + reinterpret_cast(&block_a[m][a_off]); + a_pack[m].u[0] = a_words[0]; + a_pack[m].u[1] = a_words[1]; + a_pack[m].u[2] = a_words[2]; + a_pack[m].u[3] = a_words[3]; + } + + // sum_a[m] = Σ a[m][i] via 4× v_dot2 with bf162(1,1) — no fp32 widen. + float sum_a[M_COUNT]; + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + float s = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])), + *((const bf16x2_t*)&BF16_ONES), + s, /*clamp=*/false); + } + sum_a[m] = s; + } + + // Per col: build magic-value pack, dot against all M activations. + // unroll 1 keeps q_pack live one col at a time (8 fp32 VGPRs recycled) + // — same register-pressure trick as the previous fp32 path. + #pragma unroll 1 + for (int col = 0; col < 4; ++col) { + pack4 q_pack; + const uint32_t qa = w[col]; + q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; + + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + float partial = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + partial = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])), + partial, /*clamp=*/false); + } + // block_c += y_b_f * partial + z_b_f * sum_a (same correction as + // M=1) + block_c[m][col] = + __fmaf_rn(y_b_f[col], partial, + __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col])); + } + } + } + } + k += 32; // 4 weight words * 8 nibbles = 32 K elements + } + + // Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all + // four lanes in a single 64-bit CAS write directly to the T-typed output + // (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with + // a single b64 op we halve the atomic instruction count vs two b32 CAS + // calls, AND save the FP32 buffer + memset + cast pass entirely. + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + if (offset_m + m >= size_m) continue; // skip padding rows past size_m + T* out = c + (offset_m + m) * size_n + n; + if constexpr (std::is_same::value) { + half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]), + __float2half_rn(block_c[m][1])); + half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]), + __float2half_rn(block_c[m][3])); + atomic_add_pk4_f16(out, r01, r23); + } else { + bf162_t r01; + r01.x = __float2bfloat16(block_c[m][0]); + r01.y = __float2bfloat16(block_c[m][1]); + bf162_t r23; + r23.x = __float2bfloat16(block_c[m][2]); + r23.y = __float2bfloat16(block_c[m][3]); + atomic_add_pk4_bf16(out, r01, r23); + } + } +} + +#else // non-RDNA3 device pass: empty __global__ for symbol parity. + +template +__global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*, + const T*, T*, const int, const int, + const int, const int, const int, + const int*) {} + +#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__ + +// --------------------------------------------------------------------------- +// Launcher. +// --------------------------------------------------------------------------- + +template +void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + dim3 block(THREADS_X); + dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4), + (size_m + M_COUNT - 1) / M_COUNT, + (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE); + + gemm_q4_kernel_rdna3<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +// Dispatch to the largest M_COUNT template that doesn't waste more than +// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is +// the right tool, not bigger M_COUNT in the scalar dot-product path. +// +// Tile-waste table: +// M=1 -> M_COUNT=1 (no waste) +// M=2,3 -> M_COUNT=2 (M=3 wastes 1/2 of last tile) +// M=4-7 -> M_COUNT=4 (worst case M=5: wastes 3/4 of last tile) +// M=8-15-> M_COUNT=8 (worst case M=9: wastes 7/8 of last tile) +// "Wasted" rows are zero-padded in LDS and skip the atomic write, so they +// only burn instructions on the last block, never affect correctness. +template +void launch_gemm_q4(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, int size_n, + int size_k, int groups, bool use_v2_format, + cudaStream_t stream) { + const int zero_offset = use_v2_format ? 0 : 1; + + if (size_m == 1) { + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + } else if (size_m <= 3) { + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + } else if (size_m <= 7) { + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + } else { + // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the + // WMMA path, but if it falls through we still produce correct output — + // just leaving 3-5× of throughput on the table for prefill workloads. + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + } +} + +} // namespace gptq_rdna3 +} // namespace vllm + +// --------------------------------------------------------------------------- +// Public entry point. +// --------------------------------------------------------------------------- +// +// Inputs: +// a [M, K] half or bfloat16 +// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle) +// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros) +// b_scales [groups, N] half or bfloat16 +// b_g_idx [K] or empty int32 (act-order permutation; empty=identity) +// use_v2_format bool (true = GPTQv2, no +1 zero offset) +// +// Output: +// c [M, N] same dtype as a + +torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, + torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format); + +torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format) { + if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 && + b_q_weight.size(1) % 16 == 0 && + ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) || + (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) { + return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, b_g_idx, + use_v2_format); + } + + TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); + TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); + TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor"); + TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor"); + TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]"); + TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]"); + TORCH_CHECK( + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, + "a must be half or bfloat16"); + TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(), + "b_scales dtype must match a"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + auto stream = at::cuda::getCurrentCUDAStream(); + + int size_m = (int)a.size(0); + int size_k = (int)a.size(1); + int size_n = (int)b_q_weight.size(1); + int groups = (int)b_qzeros.size(0); + + TORCH_CHECK(b_q_weight.size(0) * 8 == size_k, + "b_q_weight first dim must be K/8"); + TORCH_CHECK(b_scales.size(0) == groups, + "b_scales must have same group count as qzeros"); + TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N"); + TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)"); + + auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + at::Tensor c = torch::zeros({size_m, size_n}, opts); + + const int* g_idx_ptr = nullptr; + if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) { + TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32, + "b_g_idx must be int32"); + g_idx_ptr = (const int*)b_g_idx.data_ptr(); + } + + if (a.scalar_type() == torch::kHalf) { + vllm::gptq_rdna3::launch_gemm_q4( + (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), + (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(), + g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups, + use_v2_format, stream); + } else { + vllm::gptq_rdna3::launch_gemm_q4( + (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(), + (const uint32_t*)b_q_weight.data_ptr(), + (const uint32_t*)b_qzeros.data_ptr(), + (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), g_idx_ptr, + (vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups, + use_v2_format, stream); + } + + return c; +} diff --git a/csrc/rocm/q_gemm_rdna3_wmma.cu b/csrc/rocm/q_gemm_rdna3_wmma.cu new file mode 100644 index 00000000000..966d5df403f --- /dev/null +++ b/csrc/rocm/q_gemm_rdna3_wmma.cu @@ -0,0 +1,2165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// W4A16 GPTQ WMMA prefill kernel for AMD RDNA3 (gfx1100). This is the +// matrix-instruction path for M >= 16; small-M decode lives in the sibling +// file q_gemm_rdna3.cu and is exposed via a different op +// (`gptq_gemm_rdna3`). Keeping the two paths in separate translation units +// is intentional: an earlier attempt at putting WMMA in the same TU as the +// scalar dot-product kernel introduced a compile-time interaction that +// silently miscompiled the M=1 path even though the WMMA template was never +// instantiated for M=1. Hipcc's optimizer appears to scope some decisions +// at the TU level (likely register file / SGPR pressure heuristics across +// all kernels in the TU), so we isolate. +// +// Hardware notes (RDNA3 / gfx1100 / RX 7900 XTX): +// * v_wmma_f32_16x16x16_{f16,bf16}_w32 — 16×16×16 GEMM in one instruction +// (~16 cycles per WMMA on wave32). Accumulator dtype is FP32; inputs +// are fp16 or bf16. There is NO 16x16x16 with fp16/bf16 accumulator on +// gfx11 that we'd want here (we always need fp32 accum to avoid loss +// across many K iterations). +// * Wave32 input fragment storage is "doubled" — lanes 16..31 hold a +// copy of lanes 0..15 for the A and B fragments. The output C +// fragment uses a different mapping: lane t holds COLUMN n=lane_lo +// of the 16x16 output, with 8 elements alternating M rows by lane_hi +// (lanes 0..15 = even rows, lanes 16..31 = odd rows). See the layout +// diagram on `gemm_q4_wmma_kernel_16x16_1w` below for the full mapping. +// * No native v_global_atomic_pk_add_{f16,bf16} on gfx11; the K-split +// epilogue (gridDim.z > 1) emulates packed atomic add via a CAS-32 +// retry loop on a uint32 word covering 2 fp16/bf16 lanes. Within a +// block we shuffle adjacent lanes via shfl_xor first so each pair of +// output cols goes through a single atomic — no intra-block +// contention. K_SPLIT == 1 keeps the original direct-write path with +// each block owning its 16M × 16N output tile. + +#include + +#include +#include +#include + +#include +#include +#include + +#include "qdq_4_rdna3.cuh" + +#if defined(__HIPCC__) && defined(__gfx1100__) + #define __HIP__RDNA3__ +#endif + +namespace vllm { +namespace gptq_rdna3_wmma { + +// Pull dequant types from the sibling namespace. +using vllm::gptq_rdna3::bf162_t; +using vllm::gptq_rdna3::bf16_t; + +// Device code below uses RDNA3-only __builtin_amdgcn_wmma_* intrinsics; +// non-RDNA3 device passes fall through to empty __global__ stubs at the +// #else block at the end of this TU. +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// PRECISE dequant variants live HERE, not in the shared qdq_4_rdna3.cuh +// header. Reason: hipcc takes different register/scheduling decisions when +// the shared header grows, which caused a measurable decode-tk/s regression +// in the scalar kernel even when these new functions were never called from +// it. Keeping them in the WMMA TU only restores scalar's binary identity +// to its tuned baseline. +// +// Numerics: the classic fp16 bit-trick (FMA form: q*scale + +// (-(1024+zero)*scale)) loses up to ~0.025 per cell at scale=0.1 because +// fp16(scale) ≠ scale and the FMA amplifies that by 1024× without cancelling +// against the precomputed (1024+zero)*scale (which is rounded to fp16 BEFORE +// the FMA). +// +// Fix: subtract (1024+zero) as an integer FIRST — exact in fp16 because +// integers in [1024, 2047] are exactly representable — then multiply by +// scale, incurring at most one half-ULP rounding. Costs one extra +// instruction per dequant pair (sub+mul vs single FMA), worth it for the +// WMMA path because K can be small (16) and errors don't average. +__forceinline__ __device__ void prep_zero_scale_fp16_precise(uint32_t zero, + half scale, + half2& z_prep, + half2& y_prep) { + union { + uint16_t u; + half h; + } zu; + zu.u = (uint16_t)(0x6400 | zero); + z_prep = __half2half2(zu.h); + y_prep = __half2half2(scale); +} + +__forceinline__ __device__ void dequant_4bit_8_fp16_precise(uint32_t qa, + half2 (&dq)[4], + half2 z_prep, + half2 y_prep) { + const uint32_t c0 = 0x64006400; + union { + uint32_t u; + half2 h2; + } q0, q1, q2, q3; + q0.u = ((qa >> 0) & 0x000F000F) | c0; + q1.u = ((qa >> 4) & 0x000F000F) | c0; + q2.u = ((qa >> 8) & 0x000F000F) | c0; + q3.u = ((qa >> 12) & 0x000F000F) | c0; + dq[0] = __hmul2(__hsub2(q0.h2, z_prep), y_prep); + dq[1] = __hmul2(__hsub2(q1.h2, z_prep), y_prep); + dq[2] = __hmul2(__hsub2(q2.h2, z_prep), y_prep); + dq[3] = __hmul2(__hsub2(q3.h2, z_prep), y_prep); +} + +// fp32 prep for the bf16→bf16 fp32-internal dequant. Returns: +// z_prep = -(128 + zero) * scale (folded bias for FMA) +// y_prep = scale +// Per-element FMA `q*y_prep + z_prep` then yields scale * (nibble - zero). +__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero, + bf16_t scale, + float& z_prep, + float& y_prep) { + float scale_f = __bfloat162float(scale); + z_prep = -(128.0f + (float)zero) * scale_f; + y_prep = scale_f; +} + +// fp32 → bf16 narrow that skips the defensive NaN-canonicalisation hipcc +// emits for __float2bfloat16. Round-half-to-even via add (0x7FFF + lsb) +// + truncate; no NaN check, since dequant outputs are bounded products +// of (nibble - zero) ∈ [-15, 15] and bf16 scales — never NaN/Inf. +__forceinline__ __device__ bf16_t f32_to_bf16_no_canon(float f) { + uint32_t fu = __float_as_uint(f); + uint32_t lsb = (fu >> 16) & 1u; + uint16_t out_u = (uint16_t)((fu + 0x7FFFu + lsb) >> 16); + bf16_t out; + __builtin_memcpy(&out, &out_u, sizeof(out)); + return out; +} + +// 4-bit GPTQ dequant for the bf16 WMMA path: 8 nibbles per int32 in qa, +// outputs bf162_t dq[4]. Implementation rationale: +// +// gfx11 has no v_pk_fma_bf16, so __hmul2/__hsub2 on bf16 lower to a +// widen-fp32-op-narrow chain that hipcc decorates with NaN canonicalisation +// (28 extra VALU ops per call observed in the v5 ISA dump: v_cmp_u_f32 + +// v_cndmask_b32 around every bf16 sub/mul). Doing the math in fp32 directly +// — bit-cast widening (`__uint_as_float((bf16_bits) << 16)` is just a shift, +// no NaN canon), one fused FMA per element, single `__float2bfloat16` narrow +// at the end — eliminates the canon and is also strictly more precise (one +// rounding step instead of two). +// +// Inputs match prep_zero_scale_bf16_f32: +// z_prep = -(128 + zero) * scale (folded bias for FMA) +// y_prep = scale +// Per-element FMA `q*y_prep + z_prep` then yields scale * (nibble - zero). +__forceinline__ __device__ void dequant_4bit_8_bf16_to_bf16(uint32_t qa, + bf162_t (&dq)[4], + float z_prep, + float y_prep) { + const uint32_t c0 = 0x43004300; + const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0; + const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0; + const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0; + const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0; + // bf16(128+nibble) bits → fp32 via left-shift by 16 (zero-extends mantissa). + const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16); + const float q0y = __uint_as_float(q0 & 0xFFFF0000u); + const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16); + const float q1y = __uint_as_float(q1 & 0xFFFF0000u); + const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16); + const float q2y = __uint_as_float(q2 & 0xFFFF0000u); + const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16); + const float q3y = __uint_as_float(q3 & 0xFFFF0000u); + // r = q*scale + (-(128+zero)*scale) = (nibble - zero)*scale, then narrow. + dq[0].x = f32_to_bf16_no_canon(__fmaf_rn(q0x, y_prep, z_prep)); + dq[0].y = f32_to_bf16_no_canon(__fmaf_rn(q0y, y_prep, z_prep)); + dq[1].x = f32_to_bf16_no_canon(__fmaf_rn(q1x, y_prep, z_prep)); + dq[1].y = f32_to_bf16_no_canon(__fmaf_rn(q1y, y_prep, z_prep)); + dq[2].x = f32_to_bf16_no_canon(__fmaf_rn(q2x, y_prep, z_prep)); + dq[2].y = f32_to_bf16_no_canon(__fmaf_rn(q2y, y_prep, z_prep)); + dq[3].x = f32_to_bf16_no_canon(__fmaf_rn(q3x, y_prep, z_prep)); + dq[3].y = f32_to_bf16_no_canon(__fmaf_rn(q3y, y_prep, z_prep)); +} + +// --------------------------------------------------------------------------- +// Packed atomic-add helpers used by the K-split epilogue. +// +// When the kernel is launched with gridDim.z > 1, multiple K-segments +// accumulate into the same 16x16 output tile and need atomic write-back. +// gfx11 has no native v_global_atomic_pk_add_{f16,bf16}, so we issue a +// CAS-loop on a 32-bit word covering 2 packed fp16/bf16 lanes. Within a +// block the kernel pairs adjacent lanes via shfl_xor first, so each pair +// of cols (n=lane_lo even, lane_lo+1) goes through a SINGLE atomic — no +// intra-block contention on the same uint32 target. Inter-block +// contention from gridDim.z (4-way at K_SPLIT=4) is the residual cost. +// --------------------------------------------------------------------------- + +__forceinline__ __device__ void atomic_add_pk_f16(half2* addr, half2 val) { + uint32_t* addr_u = reinterpret_cast(addr); + uint32_t old = *addr_u; + while (true) { + half2 cur; + __builtin_memcpy(&cur, &old, sizeof(cur)); + half2 sum = __hadd2(cur, val); + uint32_t sum_u; + __builtin_memcpy(&sum_u, &sum, sizeof(sum_u)); + uint32_t prev = atomicCAS(addr_u, old, sum_u); + if (prev == old) break; + old = prev; + } +} + +__forceinline__ __device__ void atomic_add_pk_bf16(bf162_t* addr, bf162_t val) { + uint32_t* addr_u = reinterpret_cast(addr); + uint32_t old = *addr_u; + while (true) { + bf162_t cur; + __builtin_memcpy(&cur, &old, sizeof(cur)); + bf162_t sum = __hadd2(cur, val); + uint32_t sum_u; + __builtin_memcpy(&sum_u, &sum, sizeof(sum_u)); + uint32_t prev = atomicCAS(addr_u, old, sum_u); + if (prev == old) break; + old = prev; + } +} + +#endif // helpers guard; K-split heuristics below are pure host/device + // arithmetic, called from launch_* on non-RDNA3 device passes too. + +// K-split factor heuristic. Returns the gridDim.z to use for a given K. +// Aim: each block does at least ~16 K-tiles (= K=256) so the per-block +// constant overhead (LDS init, kernel prologue) is amortised. Upper +// bound K_SPLIT=4 to cap inter-block atomic contention to 4-way. +// +// For typical Qwen-class shapes K ∈ {4096, 5120, 11008}, all return 4. +// Smaller K (e.g., embedding lookups) fall back to 1 (no split, no +// atomic). K must be divisible by (K_SPLIT × 16) for the split to be +// valid; the heuristic checks divisibility before raising the factor. +__host__ __device__ static inline int compute_wmma_k_split(int size_k) { + if (size_k >= 1024 && size_k % 64 == 0) return 4; + if (size_k >= 512 && size_k % 32 == 0) return 2; + return 1; +} + +// M-and-N-aware K-split heuristic for the v3/v4/v5 launchers. +// +// The original `compute_wmma_k_split` was K-only and always returns 4 for +// Qwen-class K, which over-subscribes wave slots and pays the atomic CAS +// epilogue once-per-K-segment per output cell. With v3/v4/v5's larger +// tiles (64M × 16/32/64N) and 4 resident waves per block, the no-split +// grid is often already well-saturated on gfx1100's 96 CUs / 3072 wave +// slots — adding gridDim.z just adds atomic overhead. +// +// Heuristic: compute the no-split block count gridDim.x × gridDim.y, then +// pick the smallest K_SPLIT that brings total waves to at least +// ~2× over-subscription (~6000 waves for our 3072 slots, i.e. 1500 blocks +// at 4 waves/block). Above that threshold, K_SPLIT=1 — direct write, no +// atomic. +// +// Args: +// size_m, size_n, size_k — GEMM dims +// m_tile, n_tile — block-level M and N tile (64×16 for v3, +// 64×32 for v4, 64×64 for v5) +// +// Returns: gridDim.z divisor (1, 2, or 4), respecting K-divisibility. +__host__ __device__ static inline int compute_wmma_k_split_mn( + int size_m, int size_n, int size_k, int m_tile, int n_tile) { + const int blocks_xy = + ((size_n + n_tile - 1) / n_tile) * ((size_m + m_tile - 1) / m_tile); + // Target: enough blocks to keep ~2× oversubscription on 96 CUs / 3072 + // wave slots at 4 waves/block ⇒ ~1500 blocks no-split. + constexpr int kTargetBlocksXY = 1500; + if (blocks_xy >= kTargetBlocksXY) return 1; + if (blocks_xy * 2 >= kTargetBlocksXY && size_k >= 512 && size_k % 32 == 0) + return 2; + if (blocks_xy * 4 >= kTargetBlocksXY && size_k >= 1024 && size_k % 64 == 0) + return 4; + // Fall back to the K-only heuristic when blocks are very few (small + // models with small N): K-split is the only way to add parallelism. + return compute_wmma_k_split(size_k); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// Native AMDGPU vector types expected by the WMMA built-ins. +using v16fp16 = _Float16 __attribute__((ext_vector_type(16))); +using v16bf16 = __bf16 __attribute__((ext_vector_type(16))); +using v8fp32 = float __attribute__((ext_vector_type(8))); + +__device__ __forceinline__ v8fp32 wmma_mma(v16fp16 a, v16fp16 b, v8fp32 c) { + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c); +} +__device__ __forceinline__ v8fp32 wmma_mma(v16bf16 a, v16bf16 b, v8fp32 c) { + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a, b, c); +} + +// Map HIP wrapper types (half, __hip_bfloat16) to native compiler types +// (_Float16, __bf16) used by the WMMA built-ins. Bitcast is a register +// reinterpret in practice. +template +struct WmmaNative; +template <> +struct WmmaNative { + using elem = _Float16; + using v16 = v16fp16; +}; +template <> +struct WmmaNative { + using elem = __bf16; + using v16 = v16bf16; +}; + +template +__device__ __forceinline__ TO bitcast_elem(FROM x) { + static_assert(sizeof(FROM) == sizeof(TO), + "bitcast_elem requires equal-sized types"); + TO r; + __builtin_memcpy(&r, &x, sizeof(TO)); + return r; +} + +// Per-T tzero (matches the helper in the scalar TU). +template +__device__ __forceinline__ T tzero(); +template <> +__device__ __forceinline__ half tzero() { + return __float2half_rn(0.0f); +} +template <> +__device__ __forceinline__ bf16_t tzero() { + return __float2bfloat16(0.0f); +} + +#endif // helpers guard (each __global__ below has its own guard so launch_* + // host code remains visible to the parser on non-RDNA3 device passes) + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// WMMA kernel: 16M × 16N tile per block, 1 wave, full K traversal. +// +// Wave32 fragment layout (verified empirically with all-modes diagnostic +// against a random A and B and the eight candidate output mappings): +// +// * A frag (row-major in M, K in slot): +// lane t, slot i → A[lane_lo][k = i] +// Lane axis encodes M (A's row), slot encodes K. +// * B frag (col-major in N, K in slot): +// lane t, slot i → B[k = i][lane_lo] +// Lane axis encodes N (B's column), slot encodes K. K-axis aligns +// with A's K-axis (same slot index). +// * C frag (output, lane = N, slot = M with hi-bit interleave): +// lane t, slot i → C[m = 2*i + lane_hi][n = lane_lo] +// Lane axis encodes N (C's column). Each lane holds 8 elements of +// its output column, alternating rows: lanes 0..15 (hi=0) hold even +// rows m=0,2,4,...,14; lanes 16..31 (hi=1) hold odd rows +// m=1,3,5,...,15. +// * Both halves of the wave (lanes 0..15 and 16..31) hold IDENTICAL input +// fragments (AMD's "doubled" wave32 input layout). Output is split +// between halves via lane_hi. +// +// History note: an earlier version of this kernel loaded B row-major in K +// and assumed the output was C[lane_lo][2*i+lane_hi]. That layout passes +// all-A=identity tests because A=I makes the K-axis sum collapse, but +// implements C = A @ B^T for non-trivial A — the bug only shows up against +// random A. A layout probe iterating all four +// {row,col} × {row,col} loadings identified mode 1 (A row, B col) with +// output [m=2*i+hi][n=lane_lo] as the unique mapping that yields A @ B. +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_16x16_1w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 16; + const int n_tile = blockIdx.x * 16; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int lane = threadIdx.x; // 0..31 + const int lane_lo = lane & 15; // row index within fragment + const int lane_hi = lane >> 4; // 0 or 1 + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + // K-split: each block in the gridDim.z dimension processes a contiguous + // K-segment [k_start, k_end). With gridDim.z > 1, multiple blocks + // accumulate into the same output tile and need atomic write-back at + // the end. With gridDim.z == 1 the kernel falls back to the original + // behaviour: full K range, single writer per cell, direct write. + // + // The split multiplies the wave count by gridDim.z and proportionally + // raises CU saturation — this is the dominant lever for closing the + // throughput gap to the fp16 scalar kernel at M >= 16, which already + // uses K-split natively (gridDim.z = K/256). At gridDim.z = 4 with + // K=4096, WMMA jumps from 17% to ~67% wave-slot saturation. + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // LDS tile of dequantized B. 16 K rows × 16 N cols. + __shared__ T b_lds[16][16]; + + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + // ---- Dequant 16x16 B tile into LDS ---- + // 32 lanes split 16 N cols × 2 K-octets per col = 32 dequant tasks. + const int my_n = lane_lo; + const int my_k_octet = lane_hi; // 0 → K[0..7], 1 → K[8..15] + const int actual_n = n_tile + my_n; + + if (actual_n < size_n) { + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[k_base + 0][my_n] = __low2half(dq[0]); + b_lds[k_base + 1][my_n] = __high2half(dq[0]); + b_lds[k_base + 2][my_n] = __low2half(dq[1]); + b_lds[k_base + 3][my_n] = __high2half(dq[1]); + b_lds[k_base + 4][my_n] = __low2half(dq[2]); + b_lds[k_base + 5][my_n] = __high2half(dq[2]); + b_lds[k_base + 6][my_n] = __low2half(dq[3]); + b_lds[k_base + 7][my_n] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[k_base + 0][my_n] = dq[0].x; + b_lds[k_base + 1][my_n] = dq[0].y; + b_lds[k_base + 2][my_n] = dq[1].x; + b_lds[k_base + 3][my_n] = dq[1].y; + b_lds[k_base + 4][my_n] = dq[2].x; + b_lds[k_base + 5][my_n] = dq[2].y; + b_lds[k_base + 6][my_n] = dq[3].x; + b_lds[k_base + 7][my_n] = dq[3].y; + } + } + + // No __syncthreads() needed: the launch is `dim3 block(32)` = exactly one + // wave32, so there is no inter-wave concurrency. Within a wave the + // compiler emits `s_waitcnt lgkmcnt(0)` between dependent ds_write/ds_read + // pairs automatically, so cross-lane LDS reads (lane 0 reading what + // lane 16 wrote into b_lds[8..15][0]) still observe the writes. Keeping + // the explicit `__syncthreads()` would emit a wave-level `s_barrier` that + // costs ~10-20 cycles every iteration but provides no semantic guarantee + // we don't already have for free in single-wave mode. + + // ---- Build A and B fragments, run WMMA ---- + V16 a_frag, b_frag; + const int m_row = m_tile + lane_lo; + + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + // Permuted (act-order): scattered global reads, no vectorization. + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + // Sequential A reads: replace 16 single-element global_load_b16 with + // a bulk 32-byte copy. The AMDGPU backend lowers a memcpy of this + // size + alignment to two `global_load_b128` instructions. size_k is + // a multiple of 16 (TORCH_CHECK above) and k_tile increments by 16, + // so k_tile + 16 is always within bounds — no tail handling needed. + // Note: we memcpy into the whole vector (`&a_frag`) rather than + // `&a_frag[0]`; ext_vector_type element addresses aren't reliably + // valid C pointers across compiler versions. + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes (16 × 2)"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // B fragment: lane t holds COLUMN n=lane_lo of the B tile (K-axis in + // slot, N-axis in lane). This is the AMD WMMA convention for the right + // operand of a matrix multiply — K-axis aligns with A's K-axis (also + // in slot), enabling per-lane inner products. + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag[i] = bitcast_elem(b_lds[i][lane_lo]); + } + + #ifdef VLLM_WMMA_LAYOUT_DEBUG + // Diagnostic: skip WMMA, force c_acc to encode (lane, slot) so the + // store pattern reveals the C-output lane→matrix mapping. Compile with + // -DVLLM_WMMA_LAYOUT_DEBUG to enable. Output: c[m][n] = lane + slot/16. + (void)a_frag; + (void)b_frag; + #pragma unroll + for (int i = 0; i < 8; i++) { + c_acc[i] = (float)lane + (float)i / 16.0f; + } + // Run only one K iteration in debug mode so c_acc isn't overwritten. + if (k_tile == 0) { + k_tile = size_k; // exit loop on next check + } + #else + c_acc = wmma_mma(a_frag, b_frag, c_acc); + #endif + + // No __syncthreads() needed before the next iter overwrites b_lds: + // single-wave block, and the next iter's ds_write to b_lds is preceded + // by a `s_waitcnt lgkmcnt(0)` from the compiler that ensures the WMMA's + // ds_read of b_frag has completed before the new ds_write issues. + } + + // ---- Store C ---- + // Lane t holds column n=lane_lo of the output tile, 8 rows determined by + // slot i and lane_hi: + // lane_hi == 0 → rows 0, 2, 4, ..., 14 (even rows) + // lane_hi == 1 → rows 1, 3, 5, ..., 15 (odd rows) + // c_acc[i] corresponds to actual row m = 2*i + lane_hi at column lane_lo. + if (gridDim.z > 1) { + // K-split path: 4 (or whatever the split factor is) K-segments per + // output cell contend → atomic accumulation. Caller has zero-init'd c. + // + // Pair-shuffle to avoid intra-block CAS contention: lanes lane_lo and + // lane_lo+1 share the same uint32 atomic target (4 bytes = 2 fp16), + // so without pairing they'd hammer the same word. Instead, swap the + // c_acc[i] value with the lane_lo+1 neighbour via shfl_xor and have + // ONLY the even lane issue a single packed CAS. Inter-block contention + // (gridDim.z-way per cell) remains and is the residual atomic cost. + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_tile + lane_lo; // valid only on even lane + #pragma unroll + for (int i = 0; i < 8; i++) { + // Wave-wide shuffle: every lane participates so the side-effect is + // visible. Only even lanes use the result. shfl_xor with mask 1 + // swaps with the lane_lo XOR 1 neighbour (same lane_hi → same row). + float other_f = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + // Pack: .x = mine (col=lane_lo even), .y = neighbour (col=lane_lo+1) + half2 packed = + __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(c_acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + // gridDim.z == 1: single writer per cell, direct non-atomic write. + // Caller can leave c uninitialised (torch::empty) since every cell is + // assigned exactly once. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(c_acc[i]); + } else { + *dst = __float2bfloat16(c_acc[i]); + } + } + } + } + } +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_q4_wmma_kernel_16x16_1w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +#endif + +template +void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // 1 wave per block (32 lanes), 16x16 C tile per block. gridDim.z splits + // K so that more blocks (and therefore more waves) are in flight; with + // K_SPLIT > 1 the kernel switches to atomic write-back at the epilogue. + const int k_split = compute_wmma_k_split(size_k); + dim3 block(32); + dim3 grid((size_n + 15) / 16, (size_m + 15) / 16, k_split); + gemm_q4_wmma_kernel_16x16_1w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// 32x16_2w kernel: 2 waves per block, 32M × 16N tile, double-buffered LDS. +// +// Targets the bf16-WMMA prefill regime (M >= 128) where the v1 single-wave +// kernel saturates only ~24% of WMMA peak because each wave does roughly +// one v_wmma every ~30-40 cycles (16-cycle wmma latency + dequant + LDS + +// global A load all serial inside the single resident wave). +// +// Two structural changes vs v1: +// +// * 2 waves per block (64 threads). Both waves cooperate on a 32M×16N +// output tile: wave 0 produces rows [0..15], wave 1 rows [16..31]. +// The B-tile in LDS is shared (only wave 0 dequants); each wave loads +// its own A slice from global. With two resident waves the SIMD +// scheduler can keep the WMMA pipeline full by interleaving wmmas +// from the two waves while the other does dequant / LDS / load work. +// +// * Double-buffered LDS B-tile (b_lds[2][16][16]). Wave 0 dequants +// the K-tile for iter k+1 while both waves consume the K-tile for +// iter k. Pulls dequant out of the WMMA-critical path. Costs ~512 B +// extra LDS per block (irrelevant — gfx1100 has 64 KB LDS/CU). +// +// One __syncthreads() per K-iter remains: it ensures wave 0's dequant of +// the next K-tile has committed AND both waves have finished reading the +// current K-tile before wave 0 wraps around and overwrites it. +// +// The v2 launcher (`launch_gemm_q4_wmma_32x16_2w`) is the production entry on +// the WMMA path and falls back to v1 internally for size_m < 32 — see the +// comment at the top of `launch_gemm_q4_wmma_32x16_2w` for the M=16 regression +// rationale that justifies the fallback. +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_32x16_2w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 32; // 32-row stride per block + const int n_tile = blockIdx.x * 16; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..63 + const int wave_id = tid >> 5; // 0 or 1 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + // K-split: each block handles a contiguous K-segment when gridDim.z > 1. + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Double-buffered LDS B-tile. 2 × 16K × 16N × sizeof(T) = 1024 B for + // fp16/bf16. + __shared__ T b_lds[2][16][16]; + + // Dequant a 16K × 16N B-tile into b_lds[buf]. Only wave 0 participates + // (32 lanes do 32 dequant tasks: 16 N-cols × 2 K-octets per col). + auto dequant_into = [&](int buf, int k_tile) { + if (wave_id != 0) return; + + const int my_n = lane_lo; + const int my_k_octet = lane_hi; + const int actual_n = n_tile + my_n; + + if (actual_n >= size_n) return; + + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[buf][k_base + 0][my_n] = __low2half(dq[0]); + b_lds[buf][k_base + 1][my_n] = __high2half(dq[0]); + b_lds[buf][k_base + 2][my_n] = __low2half(dq[1]); + b_lds[buf][k_base + 3][my_n] = __high2half(dq[1]); + b_lds[buf][k_base + 4][my_n] = __low2half(dq[2]); + b_lds[buf][k_base + 5][my_n] = __high2half(dq[2]); + b_lds[buf][k_base + 6][my_n] = __low2half(dq[3]); + b_lds[buf][k_base + 7][my_n] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[buf][k_base + 0][my_n] = dq[0].x; + b_lds[buf][k_base + 1][my_n] = dq[0].y; + b_lds[buf][k_base + 2][my_n] = dq[1].x; + b_lds[buf][k_base + 3][my_n] = dq[1].y; + b_lds[buf][k_base + 4][my_n] = dq[2].x; + b_lds[buf][k_base + 5][my_n] = dq[2].y; + b_lds[buf][k_base + 6][my_n] = dq[3].x; + b_lds[buf][k_base + 7][my_n] = dq[3].y; + } + }; + + // Pre-fill buffer 0 with the first K-tile so iter 0 has data to consume. + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + // Issue dequant of next K-tile (wave 0 only) — overlaps with current + // iter's WMMA work below. The LDS write is non-blocking; the sync at + // end of iter ensures wave 1 sees it before iter k+1. + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // Load A: each wave loads its own M slice in parallel. + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag; + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Load B from current buffer (both waves read identical data). + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo]); + } + + // Each wave issues its own WMMA against its own a_frag + shared b_frag. + c_acc = wmma_mma(a_frag, b_frag, c_acc); + + // Sync ensures: (a) wave 0's dequant into next_buf has committed, + // (b) both waves are done reading cur_buf — so next iter's overwrite + // (cur_buf becomes the previous next_buf and gets reused two iters + // later) is race-free. + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- + // Each wave owns rows [m_tile + wave_id*16 .. + 16) of the output tile. + const int m_tile_wave = m_tile + wave_id * 16; + + if (gridDim.z > 1) { + // K-split atomic path. Pair-shuffle within wave to halve atomic count. + // shfl_xor here is wave-local (wave32 semantics) so each wave does its + // own pairing — the two waves don't interact during the store. + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_tile + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(c_acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + // Single writer per cell, direct non-atomic write. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(c_acc[i]); + } else { + *dst = __float2bfloat16(c_acc[i]); + } + } + } + } + } +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_q4_wmma_kernel_32x16_2w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +#endif + +template +void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // Fallback to v1 for size_m < 32. With M-tile=32 the v2 block has 2 waves + // working on rows [0..15] and [16..31]; at M < 32 the second wave processes + // out-of-range M rows (zero-padded a_frag → wmma produces nothing useful) + // and just wastes SIMD cycles. Bench measured a +47 % regression at M=16 + // vs v1 for this reason. The M < 32 case is rare in serving (decode at + // max-num-seqs=32 lands at M≈32 steady-state; the M=16 sliver is edge), + // but the fallback costs nothing and is the right shape. + if (size_m < 32) { + launch_gemm_q4_wmma_16x16_1w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + return; + } + + // 2 waves per block (64 threads), 32M × 16N C tile per block. + // K-split heuristic shared with v1. With M-tile=32, the natural + // grid blocks are halved on Y vs v1 — but each block does 2× the work, + // so total wave count is unchanged at the same M when K_SPLIT is equal. + const int k_split = compute_wmma_k_split(size_k); + dim3 block(64); + dim3 grid((size_n + 15) / 16, (size_m + 31) / 32, k_split); + gemm_q4_wmma_kernel_32x16_2w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// 64x16_4w kernel: 4 waves per block, 64M × 16N tile, double-buffered LDS. +// +// Targets the prefill plateau observed at M >= 128 in v2 (~144 K tk/s bf16, +// ~28 % of WMMA peak). The bottleneck is wmma issue rate per resident wave: +// each wave issues at most 1 wmma per ~30-40 cycles. With only 2 waves per +// block, two wmmas overlap; the wmma pipeline (16-cycle latency) is mostly +// idle. +// +// Doubling the resident wave count (2 → 4) targets ~2× wmma throughput by +// keeping the pipeline closer to full. 64M tile keeps N-tile at 16 (so the +// b_lds layout, dequant pattern, and store mapping carry over from v2) — the +// only structural changes are: +// +// * 4 waves cooperate on the 64M × 16N output tile. Wave w produces rows +// [16w .. 16w+15] of the M tile. +// * Dequant remains on wave 0 only (32 lanes do 32 dequant slots, identical +// to v2). Waves 1-3 idle through the dequant phase but their wmmas can +// issue concurrently with wave 0's dequant of the *next* K-tile thanks +// to the double buffer — net wave occupancy is dominated by the wmma +// phase, not the dequant phase. +// * One __syncthreads() per K-iter still required (same race as v2). +// +// Costs: same LDS as v2 (1024 B for the b_lds double buffer). Block has 128 +// threads vs 64 in v2; gfx1100 supports up to 1024 threads/block so this is +// well within budget. Doubles VGPR pressure slightly because the four waves +// each hold their own a_frag + c_acc — but each wave's working set is +// independent so per-thread VGPR is unchanged. +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_64x16_4w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = + blockIdx.y * 64; // 64-row stride per block (4 waves × 16M) + const int n_tile = blockIdx.x * 16; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..127 + const int wave_id = tid >> 5; // 0..3 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + // K-split: each block handles a contiguous K-segment when gridDim.z > 1. + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Double-buffered LDS B-tile. Same layout as v2. + __shared__ T b_lds[2][16][16]; + + // Dequant a 16K × 16N B-tile into b_lds[buf]. Only wave 0's 32 lanes + // participate (16 N-cols × 2 K-octets = 32 dequant slots). Waves 1-3 + // skip — their wmma can run concurrently with wave 0's next-iter dequant + // through the double buffer. + auto dequant_into = [&](int buf, int k_tile) { + if (wave_id != 0) return; + + const int my_n = lane_lo; + const int my_k_octet = lane_hi; + const int actual_n = n_tile + my_n; + + if (actual_n >= size_n) return; + + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[buf][k_base + 0][my_n] = __low2half(dq[0]); + b_lds[buf][k_base + 1][my_n] = __high2half(dq[0]); + b_lds[buf][k_base + 2][my_n] = __low2half(dq[1]); + b_lds[buf][k_base + 3][my_n] = __high2half(dq[1]); + b_lds[buf][k_base + 4][my_n] = __low2half(dq[2]); + b_lds[buf][k_base + 5][my_n] = __high2half(dq[2]); + b_lds[buf][k_base + 6][my_n] = __low2half(dq[3]); + b_lds[buf][k_base + 7][my_n] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[buf][k_base + 0][my_n] = dq[0].x; + b_lds[buf][k_base + 1][my_n] = dq[0].y; + b_lds[buf][k_base + 2][my_n] = dq[1].x; + b_lds[buf][k_base + 3][my_n] = dq[1].y; + b_lds[buf][k_base + 4][my_n] = dq[2].x; + b_lds[buf][k_base + 5][my_n] = dq[2].y; + b_lds[buf][k_base + 6][my_n] = dq[3].x; + b_lds[buf][k_base + 7][my_n] = dq[3].y; + } + }; + + // Pre-fill buffer 0 with the first K-tile so iter 0 has data to consume. + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + // Issue dequant of next K-tile (wave 0 only) — overlaps with the wmma + // work below across all 4 waves. + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // Each wave loads its own 16M slice of A (wave w handles M-rows + // [m_tile + 16w .. m_tile + 16w + 15]). + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag; + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Load B from current buffer (all 4 waves read identical data). + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo]); + } + + // Each wave issues its own WMMA against its own a_frag + shared b_frag. + // 4 wmmas in flight per block per K-iter. + c_acc = wmma_mma(a_frag, b_frag, c_acc); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- Each wave owns 16 M-rows of the output tile. + const int m_tile_wave = m_tile + wave_id * 16; + + if (gridDim.z > 1) { + // K-split atomic path. Pair-shuffle within wave to halve atomic count. + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_tile + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(c_acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + // Single writer per cell, direct non-atomic write. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(c_acc[i]); + } else { + *dst = __float2bfloat16(c_acc[i]); + } + } + } + } + } +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_q4_wmma_kernel_64x16_4w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +#endif + +template +void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // Fall back to v2 for M < 64 (would waste 1+ waves on out-of-range rows). + if (size_m < 64) { + launch_gemm_q4_wmma_32x16_2w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + return; + } + + // 4 waves per block (128 threads), 64M × 16N tile per block. + const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 16); + dim3 block(128); + dim3 grid((size_n + 15) / 16, (size_m + 63) / 64, k_split); + gemm_q4_wmma_kernel_64x16_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// 64x32_4w kernel: 4 waves per block, 64M × 32N tile, double-buffered LDS. +// +// Builds on v3 by doubling the N-tile from 16 → 32. Each wave now issues +// 2 wmmas per K-iter (one for cols 0-15, one for cols 16-31, sharing the +// same a_frag). With 4 waves × 2 wmmas = 8 wmmas in flight per K-iter, +// the wmma pipeline gets ~2× more in-flight work than v3 — targeting the +// remaining wmma issue gap on Qwen-class shapes (gate/up, down) where v3 +// plateaus at ~22 TFLOPS effective. +// +// Costs: +// * LDS B-tile doubles (2 × 16K × 32N × sizeof(T) = 2048 B for fp16/bf16). +// * Dequant doubles (32 N-cols × 2 K-octets = 64 slots/K-tile). Distributed +// across waves 0-1: wave 0 dequants n=[0..15], wave 1 dequants n=[16..31]. +// Waves 2-3 idle on dequant but do wmma work. +// * Per-wave registers: 2 × v8fp32 accumulator (16 fp32 = 32 VGPRs) plus +// b_frag0 + b_frag1 (32 VGPRs total). Within budget. +// +// Mapping invariant (wave-id → output tile slice): +// * Wave w produces M rows [m_tile + 16w .. m_tile + 16w + 15] +// * c_acc0 holds N cols [n_tile + 0 .. n_tile + 15] +// * c_acc1 holds N cols [n_tile + 16 .. n_tile + 31] +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_64x32_4w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 64; + const int n_tile = blockIdx.x * 32; // 32-col stride per block (was 16 in v3) + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..127 + const int wave_id = tid >> 5; // 0..3 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + // Two accumulators per wave: c_acc0 covers cols [n_tile..n_tile+15], + // c_acc1 covers cols [n_tile+16..n_tile+31]. + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Doubled LDS B-tile: [buf][k][n_in_tile=0..31]. + __shared__ T b_lds[2][16][32]; + + // Dequant: 64 slots per K-tile (32 N-cols × 2 K-octets). Distributed + // across waves 0,1: wave w handles n_in_tile in [16w..16w+15], k_oct in + // {0,1}. Each dequanting wave has 32 lanes for 32 dequant slots — perfect + // mapping, identical layout to v2/v3 dequant per wave. Waves 2,3 stay idle + // during dequant; they catch up via the double buffer overlapping the + // wmma of iter k with dequant of iter k+1. + // + // Tried distributing dequant across all 4 waves (16 slots/wave): regressed + // 3% on gate/up M=2048 (357K → 347K tk/s). The dequant is not on the + // critical path; spreading it just adds LDS bank pressure with no gain. + auto dequant_into = [&](int buf, int k_tile) { + if (wave_id >= 2) return; + + const int my_n_local = lane_lo; // 0..15 (within wave's N-half) + const int my_n_in_tile = wave_id * 16 + my_n_local; // 0..31 (in 32N tile) + const int my_k_octet = lane_hi; + const int actual_n = n_tile + my_n_in_tile; + + if (actual_n >= size_n) return; + + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[buf][k_base + 0][my_n_in_tile] = __low2half(dq[0]); + b_lds[buf][k_base + 1][my_n_in_tile] = __high2half(dq[0]); + b_lds[buf][k_base + 2][my_n_in_tile] = __low2half(dq[1]); + b_lds[buf][k_base + 3][my_n_in_tile] = __high2half(dq[1]); + b_lds[buf][k_base + 4][my_n_in_tile] = __low2half(dq[2]); + b_lds[buf][k_base + 5][my_n_in_tile] = __high2half(dq[2]); + b_lds[buf][k_base + 6][my_n_in_tile] = __low2half(dq[3]); + b_lds[buf][k_base + 7][my_n_in_tile] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[buf][k_base + 0][my_n_in_tile] = dq[0].x; + b_lds[buf][k_base + 1][my_n_in_tile] = dq[0].y; + b_lds[buf][k_base + 2][my_n_in_tile] = dq[1].x; + b_lds[buf][k_base + 3][my_n_in_tile] = dq[1].y; + b_lds[buf][k_base + 4][my_n_in_tile] = dq[2].x; + b_lds[buf][k_base + 5][my_n_in_tile] = dq[2].y; + b_lds[buf][k_base + 6][my_n_in_tile] = dq[3].x; + b_lds[buf][k_base + 7][my_n_in_tile] = dq[3].y; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // Load A: each wave loads its 16M slice, shared across both wmmas. + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag0, b_frag1; + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Load B for cols [0..15] and cols [16..31]. Both halves of the 32N tile. + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag0[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo]); + b_frag1[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 16]); + } + + // Two wmmas per wave per K-iter, sharing a_frag. + c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- + // Each wave owns 16M rows × 32N cols. c_acc0 → cols [n_tile..n_tile+15], + // c_acc1 → cols [n_tile+16..n_tile+31]. + const int m_tile_wave = m_tile + wave_id * 16; + + // Helper: store one v8fp32 accumulator's 8 outputs (covers 16 N-cols at + // n_base via lane_lo + interleaved M rows m_tile_wave + 2i + lane_hi). + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(acc[i]); + } else { + *dst = __float2bfloat16(acc[i]); + } + } + } + } + }; + + store_acc(c_acc0, n_tile); + store_acc(c_acc1, n_tile + 16); +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_q4_wmma_kernel_64x32_4w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +#endif + +template +void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // Fall back to v3 when M < 64 (small-M decode/prefill stays on the + // narrower 64M × 16N path) or when N < 32 (tile would waste a wave on + // out-of-range cols). + if (size_m < 64 || size_n < 32) { + launch_gemm_q4_wmma_64x16_4w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + return; + } + + // 4 waves per block (128 threads), 64M × 32N tile per block. + const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 32); + dim3 block(128); + dim3 grid((size_n + 31) / 32, (size_m + 63) / 64, k_split); + gemm_q4_wmma_kernel_64x32_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// 64x64_4w kernel: 4 waves per block, 64M × 64N tile, 4 wmmas per wave per +// K-iter. +// +// Doubles the N-tile from 32 → 64. Each wave issues 4 wmmas per K-iter +// (cols 0-15, 16-31, 32-47, 48-63), all sharing the same a_frag. With +// 4 waves × 4 wmmas = 16 wmmas in flight per K-iter, the wmma pipeline +// is fully saturated (16-cycle latency × 1 issue/cycle = 16 wmmas in +// flight at peak). +// +// Costs: +// * LDS B-tile: 2 × 16K × 64N × sizeof(T) = 4 KB. Within budget. +// * Dequant: 64 N × 2 K-oct = 128 slots/K-tile. Distributed across all +// 4 waves: 32 slots/wave (16 N-cols × 2 K-octets per wave) — perfect +// 32-lane fit, full lane utilization on dequant. +// * Per-wave registers: 4 × v8fp32 acc (64 VGPRs) + 4 b_frag (64 VGPRs) +// + a_frag (16 VGPRs) ≈ 144 VGPRs/thread + locals ≈ ~170 total. +// Under the 192-VGPR gfx1100 cap. +// +// Mapping invariant: +// * Wave w produces M rows [m_tile + 16w .. m_tile + 16w + 15] +// * c_acc[i] holds N cols [n_tile + 16i .. n_tile + 16i + 15] for i=0..3 +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_64x64_4w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 64; + const int n_tile = blockIdx.x * 64; // 64-col stride per block + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..127 + const int wave_id = tid >> 5; // 0..3 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + // Four accumulators per wave, each covering 16 N-cols. + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Larger LDS B-tile: [buf][k][n_in_tile=0..63]. + __shared__ T b_lds[2][16][64]; + + // Dequant: 128 slots per K-tile (64 N-cols × 2 K-octets). All 4 waves + // participate; wave w covers n_in_tile in [16w..16w+15], k_oct in {0,1} + // — 32 slots per wave (16 N-cols × 2 K-octets), perfect 32-lane fit. + auto dequant_into = [&](int buf, int k_tile) { + const int my_n_local = lane_lo; // 0..15 within wave + const int my_n_in_tile = wave_id * 16 + my_n_local; // 0..63 + const int my_k_octet = lane_hi; + const int actual_n = n_tile + my_n_in_tile; + + if (actual_n >= size_n) return; + + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[buf][k_base + 0][my_n_in_tile] = __low2half(dq[0]); + b_lds[buf][k_base + 1][my_n_in_tile] = __high2half(dq[0]); + b_lds[buf][k_base + 2][my_n_in_tile] = __low2half(dq[1]); + b_lds[buf][k_base + 3][my_n_in_tile] = __high2half(dq[1]); + b_lds[buf][k_base + 4][my_n_in_tile] = __low2half(dq[2]); + b_lds[buf][k_base + 5][my_n_in_tile] = __high2half(dq[2]); + b_lds[buf][k_base + 6][my_n_in_tile] = __low2half(dq[3]); + b_lds[buf][k_base + 7][my_n_in_tile] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[buf][k_base + 0][my_n_in_tile] = dq[0].x; + b_lds[buf][k_base + 1][my_n_in_tile] = dq[0].y; + b_lds[buf][k_base + 2][my_n_in_tile] = dq[1].x; + b_lds[buf][k_base + 3][my_n_in_tile] = dq[1].y; + b_lds[buf][k_base + 4][my_n_in_tile] = dq[2].x; + b_lds[buf][k_base + 5][my_n_in_tile] = dq[2].y; + b_lds[buf][k_base + 6][my_n_in_tile] = dq[3].x; + b_lds[buf][k_base + 7][my_n_in_tile] = dq[3].y; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // Load A: each wave loads its 16M slice, shared across all 4 wmmas. + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3; + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Load B for all four 16-col halves. + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag0[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 0]); + b_frag1[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 16]); + b_frag2[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 32]); + b_frag3[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 48]); + } + + // Four wmmas per wave per K-iter, sharing a_frag. + c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag, b_frag3, c_acc3); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- Each wave owns 16M × 64N. Helper writes one acc slice. + const int m_tile_wave = m_tile + wave_id * 16; + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(acc[i]); + } else { + *dst = __float2bfloat16(acc[i]); + } + } + } + } + }; + + store_acc(c_acc0, n_tile + 0); + store_acc(c_acc1, n_tile + 16); + store_acc(c_acc2, n_tile + 32); + store_acc(c_acc3, n_tile + 48); +} + +// =========================================================================== +// 128x64_k16 kernel: 8 waves per block, 128M × 64N tile, K=16 per iteration. +// +// Doubles M-tile from 64 → 128. Each B-tile in LDS is reused by 8 waves +// (8 independent A-row slices) instead of 4, halving the effective B-load +// cost per output element. This matches Hybrid Triton's BLOCK_M=128. +// +// Dequant: same 128 slots as V5 (64N × 2 K-octets). Only waves 0-3 +// participate in dequant (same mapping). Waves 4-7 are pure compute. +// LDS: [2][16][64] × sizeof(T) — unchanged from V5. +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_128x64_k16( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 128; // 128-row M tile + const int n_tile = blockIdx.x * 64; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..255 + const int wave_id = tid >> 5; // 0..7 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // LDS with K-contiguous layout [buf][N][K] — enables vectorized ds_load_b128. + // Triton uses this layout to read 8 bf16 per instruction instead of 1. + __shared__ T b_lds[2][64][16]; + + // Dequant with scale/zero caching: 1 global_load per iter (7/8 of the time). + const bool dq_ok = + (wave_id < 4) && (n_tile + wave_id * 16 + lane_lo < size_n); + const int dq_n = wave_id * 16 + lane_lo; + const int dq_an = n_tile + dq_n; + const int dq_oct = lane_hi; + const int dq_kb = dq_oct * 8; + half2 ch_z = {}, ch_y = {}; + float cf_z = 0, cf_y = 0; + int cached_g = -1; + + auto dequant_into = [&](int buf, int k_tile) __attribute__((always_inline)) { + if (!dq_ok) return; + + // Reload scale/zero only on group boundary (every groupsize/16 iters). + const int g = k_tile / groupsize; + if (g != cached_g) { + cached_g = g; + const int qz_idx = g * (size_n / 8) + dq_an / 8; + const uint32_t zero_v = ((b_qzeros[qz_idx] >> ((dq_an & 7) * 4)) & 0xF) + + (uint32_t)zero_offset; + const T sc = b_scales[g * size_n + dq_an]; + if constexpr (std::is_same::value) + prep_zero_scale_fp16_precise(zero_v, sc, ch_z, ch_y); + else + prep_zero_scale_bf16_f32(zero_v, sc, cf_z, cf_y); + } + + const int qk_row = (k_tile / 8) + dq_oct; + const uint32_t qa = b_q[qk_row * size_n + dq_an]; + const int k_base = dq_kb; + + if constexpr (std::is_same::value) { + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, ch_z, ch_y); + b_lds[buf][dq_n][k_base + 0] = __low2half(dq[0]); + b_lds[buf][dq_n][k_base + 1] = __high2half(dq[0]); + b_lds[buf][dq_n][k_base + 2] = __low2half(dq[1]); + b_lds[buf][dq_n][k_base + 3] = __high2half(dq[1]); + b_lds[buf][dq_n][k_base + 4] = __low2half(dq[2]); + b_lds[buf][dq_n][k_base + 5] = __high2half(dq[2]); + b_lds[buf][dq_n][k_base + 6] = __low2half(dq[3]); + b_lds[buf][dq_n][k_base + 7] = __high2half(dq[3]); + } else { + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, cf_z, cf_y); + b_lds[buf][dq_n][k_base + 0] = dq[0].x; + b_lds[buf][dq_n][k_base + 1] = dq[0].y; + b_lds[buf][dq_n][k_base + 2] = dq[1].x; + b_lds[buf][dq_n][k_base + 3] = dq[1].y; + b_lds[buf][dq_n][k_base + 4] = dq[2].x; + b_lds[buf][dq_n][k_base + 5] = dq[2].y; + b_lds[buf][dq_n][k_base + 6] = dq[3].x; + b_lds[buf][dq_n][k_base + 7] = dq[3].y; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + const int m_row = m_tile + wave_id * 16 + lane_lo; + const T* a_row_ptr = (m_row < size_m) ? (a + m_row * size_k) : nullptr; + + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // A-load: vectorized 256-bit. b_q_perm branch removed from V7 hot path + // to eliminate ~450 ISA instructions of dead code (6× icache bloat). + V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3; + if (a_row_ptr) { + __builtin_memcpy(&a_frag, a_row_ptr + k_tile, sizeof(a_frag)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Vectorized LDS reads: 32 bytes (16 bf16) per b_frag → ds_load_b128 × 2. + static_assert(sizeof(V16) == 32, "V16 must be 32 bytes for memcpy"); + __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][0], 32); + __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][0], 32); + __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][0], 32); + __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][0], 32); + + c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag, b_frag3, c_acc3); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- Each wave owns 16M × 64N. + const int m_tile_wave = m_tile + wave_id * 16; + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(acc[i]); + } else { + *dst = __float2bfloat16(acc[i]); + } + } + } + } + }; + + store_acc(c_acc0, n_tile + 0); + store_acc(c_acc1, n_tile + 16); + store_acc(c_acc2, n_tile + 32); + store_acc(c_acc3, n_tile + 48); +} + +// =========================================================================== +// 128x64_k32 kernel: K=32 per iteration, all 8 waves dequant. +// +// Same 128M × 64N tile as V7, but processes 32 K-elements per iteration +// instead of 16. Halves iteration count and __syncthreads() calls. +// +// Dequant mapping: +// Waves 0-3 dequant N[wave*16 +: 16] × K[0:15] (same as V7) +// Waves 4-7 dequant N[(wave-4)*16 +: 16] × K[16:31] (new) +// +// LDS layout: b_lds[2][64][34] — 2 extra padding elements per row to +// avoid 8-way bank conflicts (row stride 68 bytes gives 16 unique banks +// across the 16 lanes). +// +// Per iteration: 8 WMMAs (2 A-frags × 4 N-groups) vs V7's 4. +// Requires K divisible by 32 and groupsize ≥ 32. +// =========================================================================== +template +__global__ void gemm_q4_wmma_kernel_128x64_k32( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 128; + const int n_tile = blockIdx.x * 64; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..255 + const int wave_id = tid >> 5; // 0..7 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Padded LDS: +2 elements per row to break bank conflicts. + // Row stride = 34 elements × 2 bytes = 68 bytes → 16 unique banks. + __shared__ T b_lds[2][64][34]; + + // All 8 waves dequant. Waves 0-3 fill K[0:15], waves 4-7 fill K[16:31]. + const int dq_wave4 = wave_id & 3; // 0-3 for both halves + const int dq_k_half = (wave_id >= 4) ? 1 : 0; + const int dq_n = dq_wave4 * 16 + lane_lo; // N position 0..63 + const int dq_an = n_tile + dq_n; + const bool dq_ok = (dq_an < size_n); + const int dq_oct = lane_hi + dq_k_half * 2; // K octet 0-3 + const int dq_kb = dq_oct * 8; // K base 0,8,16,24 + half2 ch_z = {}, ch_y = {}; + float cf_z = 0, cf_y = 0; + int cached_g = -1; + + auto dequant_into = [&](int buf, int k_tile) __attribute__((always_inline)) { + if (!dq_ok) return; + + const int g = k_tile / groupsize; + if (g != cached_g) { + cached_g = g; + const int qz_idx = g * (size_n / 8) + dq_an / 8; + const uint32_t zero_v = ((b_qzeros[qz_idx] >> ((dq_an & 7) * 4)) & 0xF) + + (uint32_t)zero_offset; + const T sc = b_scales[g * size_n + dq_an]; + if constexpr (std::is_same::value) + prep_zero_scale_fp16_precise(zero_v, sc, ch_z, ch_y); + else + prep_zero_scale_bf16_f32(zero_v, sc, cf_z, cf_y); + } + + const int qk_row = (k_tile / 8) + dq_oct; + const uint32_t qa = b_q[qk_row * size_n + dq_an]; + + if constexpr (std::is_same::value) { + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, ch_z, ch_y); + b_lds[buf][dq_n][dq_kb + 0] = __low2half(dq[0]); + b_lds[buf][dq_n][dq_kb + 1] = __high2half(dq[0]); + b_lds[buf][dq_n][dq_kb + 2] = __low2half(dq[1]); + b_lds[buf][dq_n][dq_kb + 3] = __high2half(dq[1]); + b_lds[buf][dq_n][dq_kb + 4] = __low2half(dq[2]); + b_lds[buf][dq_n][dq_kb + 5] = __high2half(dq[2]); + b_lds[buf][dq_n][dq_kb + 6] = __low2half(dq[3]); + b_lds[buf][dq_n][dq_kb + 7] = __high2half(dq[3]); + } else { + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, cf_z, cf_y); + b_lds[buf][dq_n][dq_kb + 0] = dq[0].x; + b_lds[buf][dq_n][dq_kb + 1] = dq[0].y; + b_lds[buf][dq_n][dq_kb + 2] = dq[1].x; + b_lds[buf][dq_n][dq_kb + 3] = dq[1].y; + b_lds[buf][dq_n][dq_kb + 4] = dq[2].x; + b_lds[buf][dq_n][dq_kb + 5] = dq[2].y; + b_lds[buf][dq_n][dq_kb + 6] = dq[3].x; + b_lds[buf][dq_n][dq_kb + 7] = dq[3].y; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + const int m_row = m_tile + wave_id * 16 + lane_lo; + const T* a_row_ptr = (m_row < size_m) ? (a + m_row * size_k) : nullptr; + + for (int k_tile = k_start; k_tile < k_end; k_tile += 32) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 32; + + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // A-load: two 16-element fragments for K[0:15] and K[16:31]. + V16 a_frag_lo, a_frag_hi; + V16 b_frag0, b_frag1, b_frag2, b_frag3; + if (a_row_ptr) { + __builtin_memcpy(&a_frag_lo, a_row_ptr + k_tile, sizeof(V16)); + __builtin_memcpy(&a_frag_hi, a_row_ptr + k_tile + 16, sizeof(V16)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag_lo[i] = (E)0; + #pragma unroll + for (int i = 0; i < 16; i++) a_frag_hi[i] = (E)0; + } + + // --- Lower K half [0:15]: 4 WMMAs --- + static_assert(sizeof(V16) == 32, "V16 must be 32 bytes for memcpy"); + __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][0], 32); + __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][0], 32); + __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][0], 32); + __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][0], 32); + + c_acc0 = wmma_mma(a_frag_lo, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag_lo, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag_lo, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag_lo, b_frag3, c_acc3); + + // --- Upper K half [16:31]: 4 WMMAs --- + __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][16], 32); + __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][16], 32); + __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][16], 32); + __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][16], 32); + + c_acc0 = wmma_mma(a_frag_hi, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag_hi, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag_hi, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag_hi, b_frag3, c_acc3); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- Same as V7. + const int m_tile_wave = m_tile + wave_id * 16; + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(acc[i]); + } else { + *dst = __float2bfloat16(acc[i]); + } + } + } + } + }; + + store_acc(c_acc0, n_tile + 0); + store_acc(c_acc1, n_tile + 16); + store_acc(c_acc2, n_tile + 32); + store_acc(c_acc3, n_tile + 48); +} + +#else // non-RDNA3 device pass: empty kernels for symbol parity (covers the + // three kernels that share this launcher). +template +__global__ void gemm_q4_wmma_kernel_64x64_4w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +template +__global__ void gemm_q4_wmma_kernel_128x64_k16(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, + const int*) {} +template +__global__ void gemm_q4_wmma_kernel_128x64_k32(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, + const int*) {} +#endif + +template +void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // Fall back to v4 when N < 64 (would waste 1+ waves on out-of-range cols). + if (size_m < 64 || size_n < 64) { + launch_gemm_q4_wmma_64x32_4w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + return; + } + + // V8 (128M × 64N, K=32/iter, 8-wave dequant) when K%32==0 and gs≥32. + // Falls back to V7 otherwise. V7/V8 read A sequentially, so act-order + // (b_q_perm != null) must skip them and use v5, which honors the perm. + if (size_m >= 128 && b_q_perm == nullptr) { + const int k_split = + compute_wmma_k_split_mn(size_m, size_n, size_k, 128, 64); + const int groupsize = size_k / groups; + dim3 block(256); + dim3 grid((size_n + 63) / 64, (size_m + 127) / 128, k_split); + if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) { + gemm_q4_wmma_kernel_128x64_k32<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); + } else { + gemm_q4_wmma_kernel_128x64_k16<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); + } + return; + } + + // 4 waves per block (128 threads), 64M × 64N tile per block. + const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 64); + dim3 block(128); + dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, k_split); + gemm_q4_wmma_kernel_64x64_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +} // namespace gptq_rdna3_wmma +} // namespace vllm + +// --------------------------------------------------------------------------- +// Public entry point. +// --------------------------------------------------------------------------- +// +// Inputs: +// a [M, K] half or bfloat16 +// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle) +// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros) +// b_scales [groups, N] half or bfloat16 +// b_g_idx [K] or empty int32 (act-order permutation; empty=identity) +// use_v2_format bool (true = GPTQv2, no +1 zero offset) +// +// Output: +// c [M, N] same dtype as a +// +// Requirements: +// * size_m >= 16 (otherwise prefer the scalar gptq_gemm_rdna3 op) +// * size_n % 16 == 0 (WMMA tile size) +// * size_k % 16 == 0 (WMMA tile size) + +torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, + torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format) { + TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); + TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); + TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor"); + TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor"); + TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]"); + TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]"); + TORCH_CHECK( + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, + "a must be half or bfloat16"); + TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(), + "b_scales dtype must match a"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + auto stream = at::cuda::getCurrentCUDAStream(); + + int size_m = (int)a.size(0); + int size_k = (int)a.size(1); + int size_n = (int)b_q_weight.size(1); + int groups = (int)b_qzeros.size(0); + + TORCH_CHECK(b_q_weight.size(0) * 8 == size_k, + "b_q_weight first dim must be K/8"); + TORCH_CHECK(b_scales.size(0) == groups, + "b_scales must have same group count as qzeros"); + TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N"); + TORCH_CHECK(size_n % 16 == 0, "WMMA path requires N % 16 == 0"); + TORCH_CHECK(size_k % 16 == 0, "WMMA path requires K % 16 == 0"); + + auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + // Always zero-init the output: some V3-V8 boundary threads may exit + // without writing their output cell (e.g. out_m >= size_m), leaving + // uninitialized garbage when torch::empty is used. The cost is + // negligible (< 1.5% of prefill time on gfx1100). + at::Tensor c = torch::zeros({size_m, size_n}, opts); + + const int* g_idx_ptr = nullptr; + if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) { + TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32, + "b_g_idx must be int32"); + g_idx_ptr = (const int*)b_g_idx.data_ptr(); + } + + const int zero_offset = use_v2_format ? 0 : 1; + + // launch_gemm_q4_wmma_64x64_4w dispatches: + // M >= 128 → 128x64_k32 / 128x64_k16 (8 waves, K=32/16 per iter) + // 64 <= M < 128 & N >= 64 → 64x64_4w (4 waves, 4 wmma/wave/K-iter) + // M >= 64 && 32 <= N < 64 → 64x32_4w (4 waves, 2 wmma/wave/K-iter) + // M >= 64 && N < 32 → 64x16_4w (4 waves) + // 32 <= M < 64 → 32x16_2w (2 waves) + // M < 32 → 16x16_1w (1 wave) + if (a.scalar_type() == torch::kHalf) { + vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w( + (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), + (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(), + g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups, + zero_offset, stream); + } else { + vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w< + vllm::gptq_rdna3_wmma::bf16_t>( + (const vllm::gptq_rdna3_wmma::bf16_t*)a.data_ptr(), + (const uint32_t*)b_q_weight.data_ptr(), + (const uint32_t*)b_qzeros.data_ptr(), + (const vllm::gptq_rdna3_wmma::bf16_t*)b_scales.data_ptr(), g_idx_ptr, + (vllm::gptq_rdna3_wmma::bf16_t*)c.data_ptr(), size_m, size_n, size_k, + groups, zero_offset, stream); + } + + return c; +} diff --git a/csrc/rocm/qdq_4_rdna3.cuh b/csrc/rocm/qdq_4_rdna3.cuh new file mode 100644 index 00000000000..f4be668c1f0 --- /dev/null +++ b/csrc/rocm/qdq_4_rdna3.cuh @@ -0,0 +1,239 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// W4A16 dequant primitives for RDNA3 (gfx1100/gfx1101/gfx1102), templated on +// the activation/scale dtype (half or __hip_bfloat16). The fp16 path reuses +// the classic exllamav2 bit-trick: +// +// (qa & 0x000F000F) | 0x64006400 -> half2(1024+q_lo, 1024+q_hi) +// (qa & 0x00F000F0) | 0x64006400 -> half2(1024+q_lo*16, 1024+q_hi*16) +// +// The "*16 then divide by 16 in the FMA" trick for the upper-nibble pairs +// works in fp16 because the mantissa (10 bits) is wide enough to hold a value +// shifted by 4 bits. In bf16 the mantissa is only 7 bits, so shifting an upper +// nibble into bits [7:4] would spill into the exponent. To avoid that, the +// bf16 path shifts each pair of nibbles down to bits [3:0]/[19:16] with a +// single right-shift before the OR with 0x43004300 (= bf162(128, 128)). + +#ifndef _qdq_4_rdna3_cuh +#define _qdq_4_rdna3_cuh + +#include + +#include +#include + +namespace vllm { +namespace gptq_rdna3 { + +using bf16_t = __hip_bfloat16; +using bf162_t = __hip_bfloat162; + +// Bit-shuffle for an int32 holding 8 sequential 4-bit weights q[0..7]: +// in: q[7] q[6] q[5] q[4] q[3] q[2] q[1] q[0] (LSB first) +// out: q[7] q[5] q[3] q[1] q[6] q[4] q[2] q[0] (even/odd interleaved) +// +// After shuffle, q[2k] sits at bits [4k : 4k+3] (lower 16) +// q[2k+1] sits at bits [16+4k: 16+4k+3] (upper 16) +// so a single mask 0x000F000F selects the matching even/odd pair, ready to +// bitcast to half2 / bfloat162 after OR-ing with the magic constant. +__forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) { + uint32_t qa = q[0]; + uint32_t qb = 0; +#pragma unroll + for (int i = 0; i < 4; i++) { + uint32_t qa0 = qa & 0x0F; + uint32_t qa1 = (qa & 0xF0) >> 4; + qa >>= 8; + qb |= (qa1 << (i * 4 + 16)); + qb |= (qa0 << (i * 4)); + } + q[0] = qb; +} + +// --------------------------------------------------------------------------- +// fp16 path +// --------------------------------------------------------------------------- + +// Precompute scale-baked constants for a single zero/scale pair. +// z1z16[0] = scale * (-1024 - zero) (used for "low" pairs) +// z1z16[1] = scale * (-64 - zero) (used for "high" pairs) +// y1y16[0] = scale * 1 (low pairs are q + 1024) +// y1y16[1] = scale * (1/16) (high pairs are q*16 + 1024) +__forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale, + half2 (&z1z16)[2], + half2 (&y1y16)[2]) { + // half(-1024 - zero) via the exllamav2 bit-trick: + // half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it. + union { + uint16_t u; + half h; + } z1u; + z1u.u = (uint16_t)(0xE400 | zero); + half z1 = z1u.h; + half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero)); + + half2 scale2 = __half2half2(scale); + z1z16[0] = __hmul2(scale2, __half2half2(z1)); + z1z16[1] = __hmul2(scale2, __half2half2(z16)); + + half y1 = __float2half_rn(1.0f); + half y16 = __float2half_rn(1.0f / 16.0f); + y1y16[0] = __hmul2(scale2, __half2half2(y1)); + y1y16[1] = __hmul2(scale2, __half2half2(y16)); +} + +// Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs: +// dq[0] = (q[0], q[1]) * scale - zero*scale +// dq[1] = (q[2], q[3]) * scale - zero*scale +// dq[2] = (q[4], q[5]) * scale - zero*scale +// dq[3] = (q[6], q[7]) * scale - zero*scale +__forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4], + half2 (&z1z16)[2], + half2 (&y1y16)[2]) { + const uint32_t c0 = 0x64006400; + + union { + uint32_t u; + half2 h2; + } q0, q1, q2, q3; + q0.u = (qa & 0x000F000F) | c0; // half2(q[0]+1024, q[1]+1024) + q1.u = (qa & 0x00F000F0) | c0; // half2(q[2]*16+1024, q[3]*16+1024) + uint32_t qa_hi = qa >> 8; + q2.u = (qa_hi & 0x000F000F) | c0; // half2(q[4]+1024, q[5]+1024) + q3.u = (qa_hi & 0x00F000F0) | c0; // half2(q[6]*16+1024, q[7]*16+1024) + + dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]); + dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]); + dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]); + dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]); +} + +// --------------------------------------------------------------------------- +// bf16 path +// --------------------------------------------------------------------------- + +// Bit-trick magic for bf16: +// bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0). +// For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because +// bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So +// ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi). +// +// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble +// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16] +// via a single 4/8/12-bit right-shift before the OR. That costs one extra +// shift per pair vs fp16, but keeps the FMA structure identical. +__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero, + bf16_t scale, + bf162_t& z_prep, + bf162_t& y_prep) { + // z = scale * -(128 + zero); y = scale. + float scale_f = __bfloat162float(scale); + float zf = -(128.0f + (float)zero) * scale_f; + bf16_t zb = __float2bfloat16(zf); + z_prep = __bfloat162bfloat162(zb); + y_prep = __bfloat162bfloat162(scale); +} + +__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa, + bf162_t (&dq)[4], + bf162_t z_prep, + bf162_t y_prep) { + const uint32_t c0 = 0x43004300; + + union { + uint32_t u; + bf162_t b2; + } q0, q1, q2, q3; + q0.u = ((qa >> 0) & 0x000F000F) | c0; // bf162(128+q[0], 128+q[1]) + q1.u = ((qa >> 4) & 0x000F000F) | c0; // bf162(128+q[2], 128+q[3]) + q2.u = ((qa >> 8) & 0x000F000F) | c0; // bf162(128+q[4], 128+q[5]) + q3.u = ((qa >> 12) & 0x000F000F) | c0; // bf162(128+q[6], 128+q[7]) + + // dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale + dq[0] = __hfma2(q0.b2, y_prep, z_prep); + dq[1] = __hfma2(q1.b2, y_prep, z_prep); + dq[2] = __hfma2(q2.b2, y_prep, z_prep); + dq[3] = __hfma2(q3.b2, y_prep, z_prep); +} + +// --------------------------------------------------------------------------- +// bf16-input → fp32-output dequant (RDNA3 scalar path). +// +// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow +// fallback. Rather than computing dq in bf16 and widening at FMA time in +// the dot product, we widen to fp32 here once (a free left-shift by 16) and +// emit the (q - zero) * scale FMA directly in fp32. This: +// * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32. +// * Eliminates 4× bf16→fp32 widens that the dot product would do. +// * Keeps the dot product accumulator in fp32 without a roundtrip. +// +// Output: fp32 dq[8], one element per K position (consumed by the +// fp32-overload of dot22_8_f in q_gemm_rdna3.cu). +__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero, + bf16_t scale, + float& z_prep, + float& y_prep) { + float scale_f = __bfloat162float(scale); + z_prep = -(128.0f + (float)zero) * scale_f; + y_prep = scale_f; +} + +// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32 +// values 128+nibble, without folding scale/zero. The caller folds scale/zb +// into the accumulator outside the inner loop using a precomputed sum_a, +// which saves ~27% of the FMA count vs the per-col-dequant approach above +// (only beneficial at M_COUNT=1; break-even at M_COUNT=2). +// +// Cost: 0 FMAs (pure bit-trick + as_float reinterprets). +__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa, + float (&q_f32)[8]) { + const uint32_t c0 = 0x43004300; + const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0; + const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0; + const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0; + const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0; + q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16); + q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u); + q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16); + q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u); + q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16); + q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u); + q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16); + q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u); +} + +__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa, + float (&dq)[8], + float z_prep, + float y_prep) { + const uint32_t c0 = 0x43004300; + const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0; + const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0; + const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0; + const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0; + // bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16 + // (just zero-extends the mantissa from 7 to 23 bits; exponent preserved). + const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16); + const float q0y = __uint_as_float(q0 & 0xFFFF0000u); + const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16); + const float q1y = __uint_as_float(q1 & 0xFFFF0000u); + const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16); + const float q2y = __uint_as_float(q2 & 0xFFFF0000u); + const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16); + const float q3y = __uint_as_float(q3 & 0xFFFF0000u); + // dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale + dq[0] = __fmaf_rn(q0x, y_prep, z_prep); + dq[1] = __fmaf_rn(q0y, y_prep, z_prep); + dq[2] = __fmaf_rn(q1x, y_prep, z_prep); + dq[3] = __fmaf_rn(q1y, y_prep, z_prep); + dq[4] = __fmaf_rn(q2x, y_prep, z_prep); + dq[5] = __fmaf_rn(q2y, y_prep, z_prep); + dq[6] = __fmaf_rn(q3x, y_prep, z_prep); + dq[7] = __fmaf_rn(q3y, y_prep, z_prep); +} + +} // namespace gptq_rdna3 +} // namespace vllm + +#endif // _qdq_4_rdna3_cuh diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index b0b44964c24..1e589598c74 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -39,6 +39,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { " Tensor scale_b, int CuCount) -> ()"); rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ); +#ifdef VLLM_ROCM_GFX1100 + // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). + rocm_ops.def( + "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " + "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor"); + rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3); + + rocm_ops.def( + "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " + "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor"); + rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma); +#endif + // Custom attention op // Compute the attention between an input query and the cached // keys/values using PagedAttention. diff --git a/tests/kernels/quantization/test_rdna3_w4a16.py b/tests/kernels/quantization/test_rdna3_w4a16.py new file mode 100644 index 00000000000..b70a2b9a86e --- /dev/null +++ b/tests/kernels/quantization/test_rdna3_w4a16.py @@ -0,0 +1,278 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the ROCm RDNA3 W4A16 GPTQ kernel (gfx1100). + +Exercises ``RDNA3W4A16LinearKernel`` end-to-end: it builds a layer with +GPTQ-format checkpoint parameters, runs ``process_weights_after_loading`` +(weight shuffle + zero-point synthesis), then ``apply_weights``, and compares +the result against an fp32 reference dequant-and-matmul. + +The kernel is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3`` and is only +built for gfx11; tests are skipped elsewhere. + +Run `pytest tests/kernels/quantization/test_rdna3_w4a16.py`. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True) + +from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402 + MPLinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402 + RDNA3W4A16LinearKernel, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 + pack_quantized_values_into_int32, +) +from vllm.model_executor.parameter import ( # noqa: E402 + GroupQuantScaleParameter, + PackedvLLMParameter, +) +from vllm.platforms.rocm import on_gfx1100 # noqa: E402 +from vllm.scalar_type import scalar_types # noqa: E402 +from vllm.utils.torch_utils import set_random_seed # noqa: E402 + +device = "cuda" + +WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8 +PACK_FACTOR = 8 # 8 x 4-bit nibbles per int32 + +# Skip everything in this module unless we are on the only architecture the +# kernel is built/registered for. +gfx1100_only = pytest.mark.skipif( + not ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3") + ), + reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in", +) + + +# --------------------------------------------------------------------------- +# Reference implementation +# --------------------------------------------------------------------------- + + +def _reference( + x_mk: torch.Tensor, + q_int4_kn: torch.Tensor, + scales_gn: torch.Tensor, + zeros_gn: torch.Tensor | None, + group_size: int, + bias: torch.Tensor | None, +) -> torch.Tensor: + """fp32 reference for the RDNA3 W4A16 op. + + x_mk: [M, K] fp16/bf16 activations. + q_int4_kn: [K, N] int32 raw stored nibbles in [0, 15]. + scales_gn: [K//G, N] per-group scales (act dtype). + zeros_gn: [K//G, N] int32 raw stored zero points in [0, 15], or None + for the symmetric path (kernel synthesizes stored zero = 7). + group_size: G. + + The kernel applies the GPTQv1 "+1" zero-point quirk, so the effective + zero is ``stored_zero + 1`` (symmetric path: 7 + 1 == bias == 8). + """ + K, N = q_int4_kn.shape + s_full = scales_gn.repeat_interleave(group_size, dim=0).to(torch.float32) # [K,N] + if zeros_gn is None: + z_full = torch.full( + (K, N), float(WEIGHT_TYPE.bias), device=x_mk.device, dtype=torch.float32 + ) + else: + z_full = (zeros_gn + 1).repeat_interleave(group_size, dim=0).to(torch.float32) + w_fp = (q_int4_kn.to(torch.float32) - z_full) * s_full # [K, N] + out = x_mk.to(torch.float32) @ w_fp # [M, N] + if bias is not None: + out = out + bias.to(torch.float32) + return out.to(x_mk.dtype) + + +# --------------------------------------------------------------------------- +# Layer construction (GPTQ checkpoint format) +# --------------------------------------------------------------------------- + + +def _build_layer( + q_int4_kn: torch.Tensor, + scales_gn: torch.Tensor, + zeros_gn: torch.Tensor | None, + dtype: torch.dtype, +) -> torch.nn.Module: + """Build a dummy layer carrying GPTQ-format params, as the loader would.""" + no_loader = lambda *args, **kwargs: None # noqa: E731 + + # qweight: int4 packed along K into int32 -> [K//8, N]. + qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0) + + class DummyLayer(torch.nn.Module): + pass + + layer = DummyLayer() + layer.register_parameter( + "qweight", + PackedvLLMParameter( + data=qweight, + weight_loader=no_loader, + input_dim=0, + output_dim=1, + packed_dim=0, + packed_factor=PACK_FACTOR, + ), + ) + layer.register_parameter( + "scales", + GroupQuantScaleParameter( + data=scales_gn.to(dtype), + weight_loader=no_loader, + input_dim=0, + output_dim=1, + ), + ) + if zeros_gn is not None: + # qzeros: int4 packed along N into int32 -> [K//G, N//8]. + qzeros = pack_quantized_values_into_int32(zeros_gn, WEIGHT_TYPE, packed_dim=1) + layer.register_parameter( + "qzeros", + PackedvLLMParameter( + data=qzeros, + weight_loader=no_loader, + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=PACK_FACTOR, + ), + ) + return layer + + +def _run_kernel( + x_mk: torch.Tensor, + q_int4_kn: torch.Tensor, + scales_gn: torch.Tensor, + zeros_gn: torch.Tensor | None, + group_size: int, + bias: torch.Tensor | None, + dtype: torch.dtype, +) -> torch.Tensor: + K, N = q_int4_kn.shape + has_zp = zeros_gn is not None + + config = MPLinearLayerConfig( + full_weight_shape=(K, N), + partition_weight_shape=(K, N), + weight_type=WEIGHT_TYPE, + act_type=dtype, + group_size=group_size, + zero_points=has_zp, + has_g_idx=False, + ) + ok, reason = RDNA3W4A16LinearKernel.can_implement(config) + assert ok, f"can_implement rejected a supported config: {reason}" + + layer = _build_layer(q_int4_kn, scales_gn, zeros_gn, dtype) + kernel = RDNA3W4A16LinearKernel( + config, + w_q_param_name="qweight", + w_s_param_name="scales", + w_zp_param_name="qzeros" if has_zp else None, + w_gidx_param_name=None, + ) + kernel.process_weights_after_loading(layer) + return kernel.apply_weights(layer, x_mk, bias=bias) + + +# Relative-L2 tolerance per dtype. The bf16 path widens dequantized weights +# to fp32 and accumulates in fp32, so it matches the reference almost exactly +# (<0.4% incl. the WMMA prefill path). The fp16 path uses the exllamav2 +# "+1024" bit-trick (see qdq_4_rdna3.cuh): the dequantized weight is recovered +# as the fp16 difference of two ~1024*scale magnitudes, which sheds low-order +# mantissa bits and leaves ~2-3% relative noise that accumulates over K. We +# compare on the relative Frobenius norm rather than elementwise, since the +# bit-trick noise produces large *relative* errors on individual near-zero +# outputs that carry negligible absolute weight. +_REL_L2_TOL = {torch.float16: 5e-2, torch.bfloat16: 1e-2} + + +def _assert_close(out: torch.Tensor, ref: torch.Tensor, dtype: torch.dtype): + rel_l2 = (out.to(torch.float32) - ref.to(torch.float32)).norm() / ref.to( + torch.float32 + ).norm() + tol = _REL_L2_TOL[dtype] + assert rel_l2 < tol, f"relative L2 error {rel_l2:.4f} exceeds {tol} for {dtype}" + + +# --------------------------------------------------------------------------- +# Forward correctness +# --------------------------------------------------------------------------- + + +# (M, K, N, group_size). M spans the scalar decode path (small M) and the +# WMMA prefill path (M >= 16 on the bf16 dispatch). K/N satisfy the kernel's +# divisibility constraints (K % G == 0, K % 8 == 0, N % 8 == 0). +MKNG_SHAPES = [ + (1, 128, 128, 128), # single group, decode + (2, 256, 256, 128), # two groups + (8, 256, 512, 64), # M=8 scalar, smaller group + (16, 512, 256, 128), # M=16 -> WMMA path for bf16 + (32, 512, 512, 64), # larger prefill +] + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_zp", [False, True], ids=["no_zp", "with_zp"]) +@pytest.mark.parametrize( + "M,K,N,G", MKNG_SHAPES, ids=[f"m{m}_k{k}_n{n}_g{g}" for m, k, n, g in MKNG_SHAPES] +) +def test_rdna3_w4a16_matches_reference(dtype, has_zp, M, K, N, G, dist_init): + set_random_seed(0) + assert K % G == 0 and K % PACK_FACTOR == 0 and N % PACK_FACTOR == 0 + + groups = K // G + x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + scales_gn = ( + 0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01 + ).to(dtype) + zeros_gn = ( + torch.randint(0, 16, (groups, N), device=device, dtype=torch.int32) + if has_zp + else None + ) + + out = _run_kernel(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None, dtype) + ref = _reference(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None) + + assert out.shape == (M, N) and out.dtype == dtype + _assert_close(out, ref, dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("M", [1, 32], ids=["decode", "prefill"]) +def test_rdna3_w4a16_bias(dtype, M, dist_init): + """Bias is added on both the scalar (M=1) and WMMA (M=32) paths.""" + set_random_seed(0) + K, N, G = 512, 256, 128 + groups = K // G + + x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + scales_gn = ( + 0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01 + ).to(dtype) + bias = (0.1 * torch.randn(N, device=device, dtype=torch.float32)).to(dtype) + + out = _run_kernel(x_mk, q_int4_kn, scales_gn, None, G, bias, dtype) + ref = _reference(x_mk, q_int4_kn, scales_gn, None, G, bias) + + _assert_close(out, ref, dtype) diff --git a/tests/kernels/quantization/test_rdna3_w4a16_selection.py b/tests/kernels/quantization/test_rdna3_w4a16_selection.py new file mode 100644 index 00000000000..b53e1663781 --- /dev/null +++ b/tests/kernels/quantization/test_rdna3_w4a16_selection.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kernel-selection / gating tests for the ROCm RDNA3 W4A16 GPTQ kernel. + +Verifies that ``choose_mp_linear_kernel`` resolves a supported W4A16 GPTQ +config to ``RDNA3W4A16LinearKernel`` on gfx1100 (it is registered ahead of +``TritonW4A16LinearKernel`` in the ROCm priority list), and that +``RDNA3W4A16LinearKernel.can_implement`` rejects the configs it does not +support so selection falls through to the next kernel. + +Run `pytest tests/kernels/quantization/test_rdna3_w4a16_selection.py`. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True) + +from vllm.model_executor.kernels.linear import ( # noqa: E402 + choose_mp_linear_kernel, +) +from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402 + MPLinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402 + RDNA3W4A16LinearKernel, +) +from vllm.platforms.rocm import on_gfx1100 # noqa: E402 +from vllm.scalar_type import scalar_types # noqa: E402 + +WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8 + +# The kernel is only selectable when running on gfx1100 with the custom op +# compiled in; otherwise can_implement rejects and selection falls through. +gfx1100_only = pytest.mark.skipif( + not ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3") + ), + reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in", +) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_selection_prefers_rdna3(dtype): + """A supported W4A16 GPTQ config resolves to the RDNA3 kernel on gfx1100.""" + config = MPLinearLayerConfig( + full_weight_shape=(1024, 256), + partition_weight_shape=(1024, 256), + weight_type=WEIGHT_TYPE, + act_type=dtype, + group_size=128, + zero_points=False, + has_g_idx=False, + ) + assert choose_mp_linear_kernel(config).__name__ == "RDNA3W4A16LinearKernel" + + +@gfx1100_only +@pytest.mark.parametrize( + "weight_type,group_size,N,full_k,expected_ok", + [ + (scalar_types.uint4b8, 128, 256, 1024, True), # nominal: supported + (scalar_types.uint4b8, -1, 256, 1024, False), # channelwise unsupported + (scalar_types.uint4b8, 128, 252, 1024, False), # N not a multiple of 8 + (scalar_types.uint4b8, 96, 256, 1024, False), # group does not divide K + (scalar_types.uint8b128, 128, 256, 1024, False), # wrong quant type + ], + ids=["ok", "channelwise", "bad_n", "group_ndiv_k", "wrong_qtype"], +) +def test_can_implement(weight_type, group_size, N, full_k, expected_ok): + """can_implement gates on quant type, group size, and N divisibility.""" + config = MPLinearLayerConfig( + full_weight_shape=(full_k, N), + partition_weight_shape=(full_k, N), + weight_type=weight_type, + act_type=torch.float16, + group_size=group_size, + zero_points=False, + has_g_idx=False, + ) + ok, reason = RDNA3W4A16LinearKernel.can_implement(config) + assert ok is expected_ok, reason diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 3222913d15a..974828175dc 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -650,6 +650,51 @@ def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None torch.ops._C.gptq_shuffle(q_weight, q_perm, bit) +def gptq_gemm_rdna3( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, +) -> torch.Tensor: + return torch.ops._rocm_C.gptq_gemm_rdna3( + a, b_q_weight, b_qzeros, b_scales, b_g_idx, use_v2_format + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3"): + + @register_fake("_rocm_C::gptq_gemm_rdna3") + def _gptq_gemm_rdna3_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, + ) -> torch.Tensor: + return torch.empty( + (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3_wmma"): + + @register_fake("_rocm_C::gptq_gemm_rdna3_wmma") + def _gptq_gemm_rdna3_wmma_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, + ) -> torch.Tensor: + return torch.empty( + (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device + ) + + if hasattr(torch.ops._C, "allspark_w8a16_gemm"): @register_fake("_C::allspark_w8a16_gemm") diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index cfeb40fecf6..3a764e15657 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -51,6 +51,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.machete import ( from vllm.model_executor.kernels.linear.mixed_precision.marlin import ( MarlinLinearKernel, ) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( + RDNA3W4A16LinearKernel, +) from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import ( TritonW4A16LinearKernel, ) @@ -339,6 +342,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { TritonW4A16LinearKernel, ], PlatformEnum.ROCM: [ + RDNA3W4A16LinearKernel, TritonW4A16LinearKernel, ConchLinearKernel, ExllamaLinearKernel, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py index 4d659b36042..b95197db426 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py @@ -29,6 +29,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( MPLinearKernel, MPLinearLayerConfig, ) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( + RDNA3W4A16LinearKernel, +) from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import ( TritonW4A16LinearKernel, ) @@ -48,6 +51,7 @@ __all__ = [ "ExllamaLinearKernel", "MacheteLinearKernel", "MarlinLinearKernel", + "RDNA3W4A16LinearKernel", "TritonW4A16LinearKernel", "XPUW4A8IntLinearKernel", "XPUwNa16LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py new file mode 100644 index 00000000000..268728f4bf6 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""W4A16 GPTQ kernel for AMD RDNA3 (gfx1100) — fp16 + bf16. + +Drop-in replacement for ExllamaLinearKernel on RDNA3 that adds native bf16 +support. The HIP kernel lives in ``csrc/rocm/q_gemm_rdna3.cu`` +and is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3``. + +Registered ahead of TritonW4A16LinearKernel for the ROCm-RDNA3 path; falls +through to the Triton kernel on non-RDNA3 ROCm devices (e.g. CDNA/MI300). +""" + +import torch + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + pack_quantized_values_into_int32, +) +from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ +from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types + +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig + + +class RDNA3W4A16LinearKernel(MPLinearKernel): + SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8] + + @classmethod + def get_min_capability(cls) -> int: + # ROCm gates via on_gfx1100() in can_implement. + return 60 + + @classmethod + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: + if not current_platform.is_rocm(): + return False, "RDNA3 W4A16 kernel is ROCm-only" + + from vllm.platforms.rocm import on_gfx1100 + + if not on_gfx1100(): + return False, "RDNA3 W4A16 kernel requires gfx1100" + + # The HIP op is registered by the C++ extension; if a user is running + # against a vLLM build that doesn't include it (e.g. partial rebuild), + # fall through gracefully to the next kernel in the registry. + if not ( + hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3") + ): + return ( + False, + "torch.ops._rocm_C.gptq_gemm_rdna3 missing — rebuild C++ extension", + ) + + if c.act_type not in (torch.float16, torch.bfloat16): + return False, "RDNA3 W4A16 kernel only supports fp16 and bf16" + + if c.weight_type not in cls.SUPPORTED_QUANT_TYPES: + return ( + False, + f"Quant type ({c.weight_type}) not supported by " + f"RDNA3 W4A16 kernel; supported: {cls.SUPPORTED_QUANT_TYPES}", + ) + + if c.group_size <= 0: + return ( + False, + "RDNA3 W4A16 kernel does not support channelwise quantization", + ) + + if c.full_weight_shape[0] % c.group_size != 0: + return ( + False, + f"Group size ({c.group_size}) does not evenly divide K " + f"({c.full_weight_shape[0]})", + ) + + # Output features must be a multiple of the pack factor (8 nibbles per + # int32) and of 8 so that qzeros (packed 4-bit per col) align cleanly + # against the BLOCK_KN_SIZE*4 = 512 N-stride and per-thread 4 columns. + if c.partition_weight_shape[1] % 8 != 0: + return ( + False, + "Output features must be a multiple of 8 for the RDNA3 " + "W4A16 kernel (qzeros packing)", + ) + + if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]: + return ( + False, + "Act-order with TP-partitioned input features is not " + "supported by the RDNA3 W4A16 kernel", + ) + + return True, None + + # ----- Weight prep (identical layout/shuffle as ExllamaLinearKernel) ----- + + def process_weights_after_loading(self, layer: torch.nn.Module): + c = self.config + device = getattr(layer, self.w_q_name).device + + # Synthesize zero points if the checkpoint doesn't carry them. + if not c.zero_points: + self.w_zp_name = "qzeros" + groups = c.partition_weight_shape[0] // c.group_size + out_features = c.partition_weight_shape[1] + + if c.weight_type.has_bias(): + # GPTQv1 quirk: the kernel adds 1 to the stored zero, so we + # encode (bias - 1) here. See exllama.py for the link to the + # documentation of this checkpoint-format wart. + zeros = torch.full( + (groups, out_features), + c.weight_type.bias - 1, + dtype=torch.int32, + device=device, + ) + else: + raise NotImplementedError( + "RDNA3 W4A16 kernel: zero-bias 4-bit quant requires " + "explicit zero points (GPTQv1 +1 quirk)." + ) + zeros = pack_quantized_values_into_int32(zeros, c.weight_type, packed_dim=1) + setattr( + layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False) + ) + + # Act-order: convert g_idx to the inverse permutation array exllama + # expects (kernel reads a[perm[k]] instead of using groups indirected + # by g_idx[k]). + if c.has_g_idx: + + def transform_w_g_idx(x): + return torch.argsort(x).to(torch.int) + + self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore + else: + self.w_gidx_name = "g_idx" + empty_g_idx = torch.nn.Parameter( + torch.empty((0,), dtype=torch.int, device=device), + requires_grad=False, + ) + setattr(layer, self.w_gidx_name, empty_g_idx) + + def transform_w_q(x): + assert isinstance(x, BasevLLMParameter) + assert self.w_gidx_name is not None + g_idx = getattr(layer, self.w_gidx_name) + + permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) + x_cont = x.data.contiguous() + # Same 4-bit shuffle as exllama. The RDNA3 kernel reads weights in + # the same shuffled int32 layout and uses the (qa & 0x000F000F) + # bit-trick on top. + ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits) + return x_cont + + def transform_w_s(x): + assert isinstance(x, BasevLLMParameter) + permute_param_layout_(x, input_dim=0, output_dim=1) + x.data = x.data.contiguous() + # Keep scales in the activation dtype (fp16 OR bf16) — the kernel + # branches on dtype internally. + return x.to(dtype=c.act_type) + + self._transform_param(layer, self.w_q_name, transform_w_q) + self._transform_param(layer, self.w_s_name, transform_w_s) + + # ----- Forward -------------------------------------------------------- + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + c = self.config + + x_2d = x.reshape(-1, x.shape[-1]) + out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) + + w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer) + + assert w_zp is not None, "Zero points are required by RDNA3 W4A16" + assert w_g_idx is not None, "g_idx tensor (possibly empty) required" + + output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, w_g_idx, False) + + if bias is not None: + output.add_(bias) + return output.reshape(out_shape) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index c75d68954f8..89471e844d8 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -190,6 +190,9 @@ def _get_gcn_arch() -> str: _GCN_ARCH = _get_gcn_arch() _ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"]) +_ON_GFX11 = "gfx11" in _GCN_ARCH +_ON_GFX1100 = "gfx1100" in _GCN_ARCH +_ON_GFX1151 = "gfx1151" in _GCN_ARCH _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) _ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) @@ -273,6 +276,18 @@ def on_gfx1x() -> bool: return _ON_GFX1X +def on_gfx11() -> bool: + return _ON_GFX11 + + +def on_gfx1100() -> bool: + return _ON_GFX1100 + + +def on_gfx1151() -> bool: + return _ON_GFX1151 + + def on_gfx12x() -> bool: return _ON_GFX12X