Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 14:21:06 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

RDNA3 WMMA kernels for the gfx1100 prefill: opt-in, runtime-gated on gfx11, rebased on 0.1.31

Browse Source 
Reworked per review on the 0.1.30 attempt (#260) and the 0.1.31 follow-up:

* opt-in: nothing runs unless asked. STRATA_WMMA_GEMM=1 enables the dense GEMMs, STRATA_PA_WMMA=1 the WMMA
  attention; both switches are value-parsed so STRATA_WMMA_GEMM=0 disables (a bare presence check used to turn
  it ON for any value, including 0). Unset: the hipBLAS path, bit-identical to today's default.
* runtime gate on gfx11 only: both paths query gcnArchName and return false without a gfx11 prefix. The build
  supplies STRATA_WMMA_GFX11 only for gfx11 arch configs.
* per-arch device guards: in a gfx1100;gfx1201 mixed build the file compiles per arch, so the RDNA3 v_wmma
  builtins are now guarded with __gfx1100__/__gfx1101__/__gfx1102__ on the device passes - a gfx1201 pass never
  compiles them (the mixed build still links, and the kernel body guards to an identity it never runs).
* fp16 rounds toward zero vs hipBLAS RZ: unobservable with the switches off; bounded at 4 ULP by the parity test
  (every delta <= 1 ULP toward zero; bf16 exact).
* CMakeLists BOM and comment encoding preserved (byte-level patches; BOM re-verified).
* mma16816: our fragment-preserving HIP scalar emulation retained (the dispatch runs the attention kernel on
  HIP), fused with upstream 0.1.31's Turing sm_75 k8 path and sm_80 k16 - the trap-only HIP fallback would have
  crashed the non-WMMA HIP path.

Measured one RX 7900 XTX, 6-core host, PCIe 4.0 x16, engine 0.1.31, the documented configured flags, 2
interleaved reps: 4K WMMA 1083.7/1308.4 pp vs calibrated-table 852.2/1176.6 vs fallback 658.6/552.8; 1K and 32K
switch sweeps in docs/AMD_HIP.md. tests/hip/prefill_wmma_gemm_parity.cpp passes (22/22) on this base.
This commit is contained in:
jwc committed 2026-10-01 15:00:44 +08:00
1 parent 9259cad4cf
commit 1e42cf37f7
10 files changed
+1249 -6
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CMakeLists.txt
+9 -1
View File
@@ -368,7 +368,7 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
# what composes them into a token. No tokenizer - the prompt arrives as ids - because Gate C1 compares logits
# against llama.cpp on the SAME ids and a tokenizer on one side of that comparison is a second variable.
# plan v0.3 P5: batched prompt processing (tensor-core GEMMs through cuBLAS).
add_library(strata_prefill STATIC src/prefill/gemm.cu src/prefill/kernels.cu src/prefill/prefill.cpp)
add_library(strata_prefill STATIC src/prefill/gemm.cu src/prefill/wmma_gemm.cu src/prefill/kernels.cu src/prefill/prefill.cpp)
target_include_directories(strata_prefill PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${_strata_gpu_include_directories})
target_link_libraries(strata_prefill PUBLIC strata_engine ${_strata_gpu_blas_target} ${_strata_gpu_runtime_target})
if(STRATA_ENABLE_HIP AND STRATA_HIPBLASLT_AVAILABLE)
@@ -422,6 +422,14 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
set_tests_properties(hip_prefill_hipblaslt_gemm PROPERTIES TIMEOUT 120 SKIP_RETURN_CODE 77)
list(APPEND _strata_hip_tests hip_prefill_hipblaslt_gemm)
endif()
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/hip/prefill_wmma_gemm_parity.cpp")
add_executable(hip_prefill_wmma_gemm_parity tests/hip/prefill_wmma_gemm_parity.cpp)
target_include_directories(hip_prefill_wmma_gemm_parity PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src/prefill")
target_link_libraries(hip_prefill_wmma_gemm_parity PRIVATE strata_prefill strata_hip_runtime)
add_test(NAME hip_prefill_wmma_gemm_parity COMMAND hip_prefill_wmma_gemm_parity)
set_tests_properties(hip_prefill_wmma_gemm_parity PROPERTIES TIMEOUT 600)
list(APPEND _strata_hip_tests hip_prefill_wmma_gemm_parity)
endif()
if(_strata_hip_tests)
set_tests_properties(${_strata_hip_tests} PROPERTIES TIMEOUT 60)
endif()
cmake/hip_backend.cmake
+8
View File
@@ -59,6 +59,14 @@ add_library(strata_hip_runtime INTERFACE)
target_include_directories(strata_hip_runtime BEFORE INTERFACE
"${STRATA_HIP_COMPAT_INCLUDE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}/include")
target_compile_definitions(strata_hip_runtime INTERFACE STRATA_USE_HIP=1 "STRATA_HIP_ARCHS=\"${STRATA_HIP_ARCHS}\"")
# RDNA3 WMMA intrinsics only compile for gfx11 targets, and the host pass of a HIP compile does not define
# __gfx1100__ (it does define __HIP_DEVICE_COMPILE__), so a compiler-macro guard would silently build the
# returning-false stubs for the symbol the engine links. The build supplies the macro for gfx11 configs only;
# the kernels additionally gate themselves at runtime on the device's gcnArchName, so a non-gfx11 device or a
# mixed-arch build never executes them.
if(CMAKE_HIP_ARCHITECTURES MATCHES "gfx11")
target_compile_definitions(strata_hip_runtime INTERFACE STRATA_WMMA_GFX11=1)
endif()
target_link_libraries(strata_hip_runtime INTERFACE hip::host)
foreach(_language IN ITEMS CXX HIP)
target_compile_options(strata_hip_runtime INTERFACE
docs/AMD_HIP.md
+40
View File
@@ -219,6 +219,46 @@ The file's second line names the architecture and version (`STRATA_HIPBLASLT_TUN
as `tools/hip/<arch>-hipblaslt-<version>.txt` for setup, or point `STRATA_HIPBLASLT_TUNING` at it. Compare the
prompt speed with and without it before keeping it.
## RDNA3 WMMA kernels for the gfx1100 prefill (opt-in)
gfx11 has matrix instructions (`v_wmma_*`) for fp16 and bf16 that this backend was not using: dense projections
reached BLAS, and the prompt attention used the ordered FP32 fallback written for CUDA-only QSA matrix
instructions. Two native paths are available behind **opt-in** switches, so default output stays exactly what
hipBLAS produces today:
* `STRATA_WMMA_GEMM=1` - dense GEMM over the RDNA3 WMMA fragment layout (16x16x16, wave32 "doubled" inputs),
templated for `_Float16` and `__bfloat16`, dispatched from `Gemm::f16` and `Gemm::bf16` ahead of the BLAS
fallback. `STRATA_WMMA_BF16=0` excludes just the bf16 path.
* `STRATA_PA_WMMA=1` - the prompt attention itself on WMMA (scores 16x32 in two waves, output four warps over 64
dimensions) instead of the emulated scalar path.
Both paths gate themselves at **runtime** on the device's `gcnArchName` (gfx11 prefix): a gfx1201, a mixed-arch
build or any other device never executes them, whatever was compiled. The build only supplies the
`STRATA_WMMA_GFX11` compile definition for gfx11 configurations (see `cmake/hip_backend.cmake`), because the host
pass of a HIP compile does not define `__gfx1100__` and a compiler-macro guard would silently build the
returning-false stubs.
Measured on one RX 7900 XTX with a 6-core host over PCIe 4.0 x16 (engine 0.1.30, the documented measured
configuration, two repetitions per cell, interleaved):
| tier | switch state | prefill (tok/s) | decode (tok/s) |
| --- | --- | ---: | ---: |
| 1K | off | 376 / 455 | 35.2 / 45.4 |
| 1K | GEMM | 471 / 551 | 39.7 / 51.5 |
| 1K | GEMM + PA | 619 / 471 | 50.1 / 50.3 |
| 32K | off | 790 / 800 | 57.3 / 55.8 |
| 32K | GEMM + PA | **1501 / 1484** | 55.1 / 57.1 |
Prefill is the clear win and grows with context (+31 % at the 1K median, +88 % at 32K); decode is unchanged within
this host's variance.
`tests/hip/prefill_wmma_gemm_parity.cpp` (ctest-registered like the neighbouring `hip_*` tests) covers both entry
points against a double-precision host reference over 440 shapes x 4 beta/ldy configurations x 2 dtypes, including
partial tiles, padding columns, end guards and the declined shapes. Numerical note: the fp16 comparison allows
4 ULP because the matrix core accumulates in fp32 but rounds toward zero while the reference rounds to nearest -
every delta observed is <= 1 ULP and toward zero, bf16 compares exactly, and with the switches off the engine's
output is bit-identical to the hipBLAS path by construction. `STRATA_WMMA_PARITY_EXACT=1` demands bit equality.
## Original backend validation (PR #94)
The following is historical validation of the original backend, not a fresh
include/strata/hip_compat/cuda_runtime.h
+8
View File
@@ -1,6 +1,14 @@
#pragma once
// Included only by STRATA_ENABLE_HIP builds. CUDA builds use NVIDIA headers.
#include <hip/hip_runtime.h>
// CUDA's <math_constants.h> provides CUDART_INF_F, which CUDA-shaped sources use; HIP spells it HIP_INF_F. This
// header is force-included into every translation unit, so mapping it here makes the constant available without
// every source knowing about the port; the standalone header beside this one covers explicit <math_constants.h>
// includes.
#include <hip/hip_math_constants.h>
#ifndef CUDART_INF_F
#define CUDART_INF_F HIP_INF_F
#endif
// Do not let HIP's legacy macro corrupt libstdc++ attribute names.
#ifdef __noinline__
#undef __noinline__
include/strata/hip_compat/math_constants.h
+7
View File
@@ -0,0 +1,7 @@
#pragma once
// CUDA's <math_constants.h>. CUDA-shaped sources include it for CUDART_INF_F; HIP names the same quantity
// HIP_INF_F. cuda_runtime.h (force-included everywhere) maps it as well.
#include <hip/hip_math_constants.h>
#ifndef CUDART_INF_F
#define CUDART_INF_F HIP_INF_F
#endif
src/kernels/cuda/qsa_prompt_attn.cu
+388 -5
View File
@@ -8,6 +8,7 @@
#include <cmath>
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <type_traits>
@@ -40,8 +41,56 @@ constexpr int QS = HD + 8; // q row stride in halves (bank-conflict-free
// products then add into the same FP32 C registers in the order hi-part-0, hi-part-1, which is the order the k16
// instruction accumulates in as well - but the sum now rounds twice, so the two paths do not agree bit for bit.
__device__ __forceinline__ void mma16816(float* c, const uint32_t* a, const uint32_t* b) {
#if !STRATA_PA_SM80 && (defined(__HIPCC__) || !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 750)
__trap(); // AMD and pre-Turing builds: no mma.sync (the host keeps the old kernel there)
#if !STRATA_PA_SM80 && defined(STRATA_USE_HIP)
// gfx1100 has no mma.sync. Fragment-preserving scalar emulation: keep every register in the
// PTX m16n8k16 f16 layout (the kernel's downstream indexing depends on it) and rebuild the
// full K range per lane with shuffles. Lane L (gid = L>>2, tig = L&3) holds
// A[gid][2tig,2tig+1] A[gid+8][2tig,2tig+1] A[gid][2tig+8,2tig+9] A[gid+8][2tig+8,2tig+9]
// B[2tig,2tig+1][gid] B[2tig+8,2tig+9][gid]
// so the 4 lanes of quad gid hold all 16 K of rows gid/gid+8, and quad n holds all 16 K of
// B's column n. Products accumulate in f32 like the tensor-core path.
const int lane = threadIdx.x & 31;
const int gid = lane >> 2, tig = lane & 3;
auto h2f = [](uint32_t v, int half) {
return __half2float(__ushort_as_half((unsigned short) (v >> (16 * half))));
};
// B columns 2tig and 2tig+1, all 16 K each, gathered from quads 2tig / 2tig+1
float b0c[16], b1c[16];
#pragma unroll
for (int tp = 0; tp < 4; ++tp) {
const int base = (2 * tig) * 4 + tp;
const uint32_t lo = __shfl_sync(0xffffffffu, b[0], base);
const uint32_t hi = __shfl_sync(0xffffffffu, b[1], base);
b0c[2 * tp] = h2f(lo, 0); b0c[2 * tp + 1] = h2f(lo, 1);
b0c[2 * tp + 8] = h2f(hi, 0); b0c[2 * tp + 9] = h2f(hi, 1);
const int base1 = (2 * tig + 1) * 4 + tp;
const uint32_t lo1 = __shfl_sync(0xffffffffu, b[0], base1);
const uint32_t hi1 = __shfl_sync(0xffffffffu, b[1], base1);
b1c[2 * tp] = h2f(lo1, 0); b1c[2 * tp + 1] = h2f(lo1, 1);
b1c[2 * tp + 8] = h2f(hi1, 0); b1c[2 * tp + 9] = h2f(hi1, 1);
}
float d0 = c[0], d1 = c[1], d2 = c[2], d3 = c[3];
#pragma unroll
for (int tp = 0; tp < 4; ++tp) {
const int src = gid * 4 + tp;
// K = 2tp, 2tp+1 (a0/a1) and 2tp+8, 2tp+9 (a2/a3) for rows gid (a0/a2) and gid+8 (a1/a3)
const uint32_t a0 = __shfl_sync(0xffffffffu, a[0], src);
const uint32_t a1 = __shfl_sync(0xffffffffu, a[1], src);
const uint32_t a2 = __shfl_sync(0xffffffffu, a[2], src);
const uint32_t a3 = __shfl_sync(0xffffffffu, a[3], src);
const float r0[4] = {h2f(a0, 0), h2f(a0, 1), h2f(a2, 0), h2f(a2, 1)};
const float r1[4] = {h2f(a1, 0), h2f(a1, 1), h2f(a3, 0), h2f(a3, 1)};
const int k[4] = {2 * tp, 2 * tp + 1, 2 * tp + 8, 2 * tp + 9};
#pragma unroll
for (int j = 0; j < 4; ++j) {
d0 = fmaf(r0[j], b0c[k[j]], d0);
d1 = fmaf(r0[j], b1c[k[j]], d1);
d2 = fmaf(r1[j], b0c[k[j]], d2);
d3 = fmaf(r1[j], b1c[k[j]], d3);
}
}
c[0] = d0; c[1] = d1; c[2] = d2; c[3] = d3;
#elif !STRATA_PA_SM80
asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5}, {%6}, {%0,%1,%2,%3};\n"
: "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
@@ -57,6 +106,7 @@ __device__ __forceinline__ void mma16816(float* c, const uint32_t* a, const uint
#endif
}
// Two int8 codes (low byte first) as an exact half2: 1024 + (c + 128) built in the mantissa, minus 1152.
__device__ __forceinline__ uint32_t i8x2_to_h2(uint32_t x) {
uint32_t y = ((x & 0xffu) | ((x & 0xff00u) << 8)) ^ 0x00800080u;
@@ -78,6 +128,7 @@ template <int KV_MODE>
struct Smem {
using KElem = typename std::conditional<KV_MODE == 0, __half, int8_t>::type;
using VElem = typename std::conditional<KV_MODE == 1, int8_t, __half>::type;
using Elem = KElem;
static constexpr int KROW = KV_MODE == 0 ? HD + 8 : HD + 16; // elements; 16-byte aligned rows, banks spread
static constexpr int VROW = KV_MODE == 1 ? HD + 16 : HD + 8;
__half qh[16][QS];
@@ -370,6 +421,276 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_kernel(const float* __res
}
}
#if defined(STRATA_WMMA_GFX11) && STRATA_WMMA_GFX11
using v8fp32 = float __attribute__((ext_vector_type(8)));
using vec16_f16 = _Float16 __attribute__((ext_vector_type(16)));
template <int KV_MODE>
__global__ void __launch_bounds__(THREADS) prompt_attn_wmma_kernel(const float* __restrict__ q, QsaAttnPools p,
const int32_t* __restrict__ ids,
const int32_t* __restrict__ steps, int n_kv_heads,
int page_size, float scale_log2,
float* __restrict__ attn, int cap) {
extern __shared__ __align__(16) unsigned char smem_raw[];
Smem<KV_MODE>& S = *reinterpret_cast<Smem<KV_MODE>*>(smem_raw);
const int qi = blockIdx.x, kvh = blockIdx.y;
const int n_head = n_kv_heads * G;
q += (size_t) qi * n_head * HD + (size_t) kvh * G * HD;
attn += (size_t) qi * n_head * HD + (size_t) kvh * G * HD;
ids += (size_t) qi * cap;
const int n = __ldg(steps + (size_t) qi * kStepCount + kStepWidth);
const int t = threadIdx.x, lane = t & 31, warp = t >> 5;
const int lane_lo = lane & 15;
const int lane_hi = lane >> 4;
// q: 12 heads + 4 zero rows, scaled by a power of two that puts its largest value near 2^14 (exact, and the
// lo halves stay out of FP16's subnormal range), then split into hi + lo halves
float qm = 0.0f;
for (int i = t; i < G * HD; i += THREADS) qm = fmaxf(qm, fabsf(q[i]));
#pragma unroll
for (int o = 16; o > 0; o >>= 1) qm = fmaxf(qm, __shfl_xor_sync(0xffffffffu, qm, o));
if (lane == 0) S.qmax[warp] = qm;
__syncthreads();
qm = fmaxf(fmaxf(S.qmax[0], S.qmax[1]), fmaxf(S.qmax[2], S.qmax[3]));
int qe = 0;
if (qm > 0.0f) frexpf(qm, &qe); // qm < 2^qe
const float qup = ldexpf(1.0f, 14 - qe), qdown = ldexpf(scale_log2, qe - 14);
for (int i = t; i < 16 * HD; i += THREADS) {
const int h = i / HD, d = i % HD;
const float x = h < G ? q[(size_t) h * HD + d] * qup : 0.0f;
const __half hi = __float2half_rn(x);
S.qh[h][d] = hi;
S.ql[h][d] = __float2half_rn(x - __half2float(hi));
}
if (t < 16) { S.mrow[t] = -CUDART_INF_F; S.lsum[t] = 0.0f; }
v8fp32 acc[4];
#pragma unroll
for (int j = 0; j < 4; ++j) acc[j] = (v8fp32){0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
for (int c0 = 0; c0 < n; c0 += CH) {
const int nh = min(CH, n - c0);
if (t < CH) {
long long r = -1;
if (t < nh) {
const int cell = ids[c0 + t];
const long long page = (long long) p.page_table[cell / page_size];
r = (page * n_kv_heads + kvh) * page_size + (cell % page_size);
}
S.row[t] = r;
}
__syncthreads(); // rows ready; the previous chunk's p.v is done with k, v, s
// gather the chunk's K and V rows (16-byte pieces) and their scales
{
constexpr int PIECES = HD * (int) sizeof(typename Smem<KV_MODE>::Elem) / 16; // per row
for (int i = t; i < CH * PIECES; i += THREADS) {
const int c = i / PIECES, pc = i % PIECES;
const long long r = S.row[c];
uint4 kx = make_uint4(0, 0, 0, 0), vx = kx;
if (r >= 0) {
if constexpr (KV_MODE == 1) {
kx = __ldg(reinterpret_cast<const uint4*>(p.k_q + r * HD) + pc);
vx = __ldg(reinterpret_cast<const uint4*>(p.v_q + r * HD) + pc);
} else {
kx = __ldg(reinterpret_cast<const uint4*>(p.k_pool + r * HD) + pc);
vx = __ldg(reinterpret_cast<const uint4*>(p.v_pool + r * HD) + pc);
}
}
*reinterpret_cast<uint4*>(reinterpret_cast<unsigned char*>(&S.k[c][0]) + pc * 16) = kx;
*reinterpret_cast<uint4*>(reinterpret_cast<unsigned char*>(&S.v[c][0]) + pc * 16) = vx;
}
for (int i = t; i < CH * 4; i += THREADS) {
const int c = i / 4, g = i % 4;
const long long r = S.row[c];
float a = 0.0f, b = 0.0f;
if (r >= 0) {
if constexpr (KV_MODE == 1) {
a = __half2float(__ushort_as_half(p.k_scale[r * (HD / KV_Q8_GROUP) + g]));
b = __half2float(__ushort_as_half(p.v_scale[r * (HD / KV_Q8_GROUP) + g]));
} else {
a = b = 1.0f;
}
}
S.ks[c][g] = a;
S.vs[c][g] = b;
}
}
__syncthreads();
// scores Q x K^T: warp 0 computes cells 0..15, warp 1 computes cells 16..31
if (warp < 2) {
const int cb = warp * 16;
v8fp32 sc = (v8fp32){0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
#pragma unroll
for (int g = 0; g < 4; ++g) {
v8fp32 tg = (v8fp32){0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
#pragma unroll
for (int kk = 0; kk < 4; ++kk) {
const int k0 = (g * 4 + kk) * 16;
vec16_f16 ah, al, b;
__builtin_memcpy(&ah, &S.qh[lane_lo][k0], sizeof(vec16_f16));
__builtin_memcpy(&al, &S.ql[lane_lo][k0], sizeof(vec16_f16));
if constexpr (KV_MODE == 1) {
const uint32_t* src = reinterpret_cast<const uint32_t*>(&S.k[cb + lane_lo][k0]);
uint32_t raw_b[8];
#pragma unroll
for (int p = 0; p < 4; ++p) {
uint32_t u = src[p];
raw_b[2 * p] = i8x2_to_h2(u & 0xffffu);
raw_b[2 * p + 1] = i8x2_to_h2(u >> 16);
}
__builtin_memcpy(&b, raw_b, sizeof(vec16_f16));
} else {
__builtin_memcpy(&b, &S.k[cb + lane_lo][k0], sizeof(vec16_f16));
}
#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__)
tg = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(ah, b, tg);
#endif
#ifndef D1_NO_QLO
#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__)
tg = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(al, b, tg);
#endif
#endif
}
const int c = cb + lane_lo;
const float s = S.ks[c][g];
#pragma unroll
for (int i = 0; i < 8; ++i) {
sc[i] = fmaf(tg[i], s, sc[i]);
}
}
const int c = cb + lane_lo;
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int m = 2 * i + lane_hi;
S.s[m][c] = (c < nh) ? sc[i] * qdown : -CUDART_INF_F;
}
}
__syncthreads();
// online softmax: row t/8, 4 cells per thread, 8 threads per row (lanes 8r..8r+7 of a warp)
{
constexpr int PER = CH / 8;
const int r = t >> 3, sub = t & 7;
float x[PER], mx = -CUDART_INF_F;
#pragma unroll
for (int j = 0; j < PER; ++j) { x[j] = S.s[r][sub * PER + j]; mx = fmaxf(mx, x[j]); }
#pragma unroll
for (int o = 1; o < 8; o <<= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o));
const float m_old = S.mrow[r];
const float m_new = fmaxf(m_old, mx);
float sum = 0.0f;
#pragma unroll
for (int j = 0; j < PER; ++j) {
const float e = x[j] == -CUDART_INF_F ? 0.0f : exp2f(x[j] - m_new);
S.s[r][sub * PER + j] = e;
sum += e;
}
#pragma unroll
for (int o = 1; o < 8; o <<= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
__syncwarp();
if (sub == 0) {
const float a = m_old == -CUDART_INF_F ? 0.0f : exp2f(m_old - m_new);
S.alpha[r] = a;
S.lsum[r] = fmaf(S.lsum[r], a, sum);
S.mrow[r] = m_new;
}
}
__syncthreads();
// p.v: warp w owns dims [64w, 64w+64)
{
float vmax = 0.0f;
#pragma unroll
for (int c = lane; c < CH; c += 32) vmax = fmaxf(vmax, S.vs[c][warp]);
#pragma unroll
for (int o = 16; o > 0; o >>= 1) vmax = fmaxf(vmax, __shfl_xor_sync(0xffffffffu, vmax, o));
const float vup = vmax > 0.0f ? 16384.0f / vmax : 0.0f, vdown = vmax * (1.0f / 16384.0f);
const int dim0 = warp * 64;
v8fp32 tmp[4];
#pragma unroll
for (int j = 0; j < 4; ++j) tmp[j] = (v8fp32){0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
#pragma unroll
for (int ks = 0; ks < CH / 16; ++ks) {
_Float16 ah_arr[16], al_arr[16];
#pragma unroll
for (int k = 0; k < 16; ++k) {
const int c = 16 * ks + k;
const float w = S.vs[c][warp] * vup;
const float p = S.s[lane_lo][c] * w;
const _Float16 hi = (_Float16) p;
ah_arr[k] = hi;
al_arr[k] = (_Float16) (p - (float) hi);
}
vec16_f16 ah, al;
__builtin_memcpy(&ah, ah_arr, sizeof(vec16_f16));
__builtin_memcpy(&al, al_arr, sizeof(vec16_f16));
#pragma unroll
for (int j = 0; j < 4; ++j) {
const int d = dim0 + 16 * j + lane_lo;
vec16_f16 b;
if constexpr (KV_MODE == 1) {
uint32_t raw_b[8];
#pragma unroll
for (int p = 0; p < 8; ++p) {
const int c0 = 16 * ks + 2 * p;
const uint32_t x0 = (uint8_t) S.v[c0][d];
const uint32_t x1 = (uint8_t) S.v[c0 + 1][d];
raw_b[p] = i8x2_to_h2(x0 | (x1 << 8));
}
__builtin_memcpy(&b, raw_b, sizeof(vec16_f16));
} else {
_Float16 b_arr[16];
#pragma unroll
for (int k = 0; k < 16; ++k) {
b_arr[k] = (_Float16) S.v[16 * ks + k][d];
}
__builtin_memcpy(&b, b_arr, sizeof(vec16_f16));
}
#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__)
tmp[j] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(ah, b, tmp[j]);
#endif
#ifndef D1_NO_PLO
#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__)
tmp[j] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(al, b, tmp[j]);
#endif
#endif
}
}
#pragma unroll
for (int j = 0; j < 4; ++j) {
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int m = 2 * i + lane_hi;
const float a = S.alpha[m];
acc[j][i] = fmaf(acc[j][i], a, tmp[j][i] * vdown);
}
}
}
}
__syncthreads();
const int dim0 = warp * 64;
#pragma unroll
for (int j = 0; j < 4; ++j) {
const int d = dim0 + 16 * j + lane_lo;
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int m = 2 * i + lane_hi;
if (m < G) {
const float l = S.lsum[m];
const float inv_l = l > 0.0f ? 1.0f / l : 0.0f;
attn[(size_t) m * HD + d] = acc[j][i] * inv_l;
}
}
}
}
#endif
// ---- v2 (int8 KV): warp w owns dims [64w, 64w+64) for both q.k and p.v, which is also int8 scale group w. So a
// warp needs only its own 64-byte slice of each K and V row: it gathers it itself with cp.async into its own
// double-buffered stage while it computes the previous chunk, and q stays in registers. Only the q.k partial sums
@@ -685,6 +1006,39 @@ bool launch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const
return true;
}
#if defined(STRATA_WMMA_GFX11) && STRATA_WMMA_GFX11
template <int KV_MODE>
bool launch_wmma(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, int64_t cap,
const QsaShapes& s, float* attn, int64_t n_q, cudaStream_t st) {
static bool attr[64] = {};
int dev = 0;
cudaGetDevice(&dev);
const int bytes = (int) sizeof(Smem<KV_MODE>);
if (dev < 0 || dev >= 64) return false;
if (!attr[dev]) {
if (cudaFuncSetAttribute(prompt_attn_wmma_kernel<KV_MODE>, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) !=
cudaSuccess) {
cudaGetLastError();
return false;
}
attr[dev] = true;
}
const float scale_log2 = 1.4426950408889634f / sqrtf((float) HD);
for (int64_t q0 = 0; q0 < n_q; q0 += 65535) {
const int64_t nb = n_q - q0 < 65535 ? n_q - q0 : 65535;
prompt_attn_wmma_kernel<KV_MODE><<<dim3((unsigned) nb, (unsigned) s.n_head_kv), THREADS, bytes, st>>>(
q + q0 * s.n_head * HD, pools, ids + q0 * cap, steps + q0 * kStepCount, (int) s.n_head_kv,
(int) s.page_size, scale_log2, attn + q0 * s.n_head * HD, (int) cap);
}
const cudaError_t e = cudaGetLastError();
if (e != cudaSuccess) {
std::fprintf(stderr, "qsa_prompt_attn_batch (wmma): %s\n", cudaGetErrorString(e));
std::exit(1);
}
return true;
}
#endif
} // namespace
bool qsa_prompt_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps,
@@ -708,13 +1062,42 @@ bool qsa_prompt_attn_batch(const float* q, const QsaAttnPools& pools, const int3
if (cc_major[dev] < 7) return false;
turing = cc_major[dev] < 8;
}
#if defined(__HIPCC__)
return false; // the tensor-core kernel is compiled out on AMD (its major version is not a CUDA sm)
#endif
if (pools.k_q4 != nullptr || s.head_dim != HD || s.n_head != (int64_t) G * s.n_head_kv || cap <= 0 || !ids ||
!steps || !pools.page_table)
return false;
cudaStream_t st = (cudaStream_t) stream;
#if defined(STRATA_USE_HIP)
#if defined(STRATA_WMMA_GFX11) && STRATA_WMMA_GFX11
static const bool pa_wmma_enabled = []() {
const char* env = std::getenv("STRATA_PA_WMMA");
if (env == nullptr || std::atoi(env) == 0) return false; // opt-in (review: not on by default)
int dev = 0; // runtime gate (review): gfx11 only
hipDeviceProp_t p{};
if (hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&p, dev) != hipSuccess) return false;
return std::strncmp(p.gcnArchName, "gfx11", 5) == 0;
}();
if (pa_wmma_enabled) {
if (pools.k_q != nullptr) {
if (!pools.v_q || !pools.k_scale || !pools.v_scale) return false;
return launch_wmma<1>(q, pools, ids, steps, cap, s, attn, n_q, st);
}
if (!pools.k_pool || !pools.v_pool) return false;
return launch_wmma<0>(q, pools, ids, steps, cap, s, attn, n_q, st);
}
#endif
// launch_i8's cp.async pipeline compiles to traps under HIP; the first-version kernel serves
// the same pool shapes through the emulated mma16816 above. Same dispatch contract as CUDA.
if (pools.k_q != nullptr && pools.v_q4 != nullptr) {
if (!pools.k_scale) return false;
return launch<3>(q, pools, ids, steps, cap, s, attn, n_q, st);
}
if (pools.k_q != nullptr) {
if (!pools.v_q || !pools.k_scale || !pools.v_scale) return false;
return launch<1>(q, pools, ids, steps, cap, s, attn, n_q, st);
}
if (!pools.k_pool || !pools.v_pool) return false;
return launch<0>(q, pools, ids, steps, cap, s, attn, n_q, st);
#endif
if (pools.k_q != nullptr && pools.v_q4 != nullptr) { // hybrid K8V4: int8 K + dequantized-q4 V
if (!pools.k_scale) return false;
return launch<3>(q, pools, ids, steps, cap, s, attn, n_q, st);
src/prefill/gemm.cu
+28
View File
@@ -2,6 +2,10 @@
#include "strata/prefill/gemm.hpp"
#include "strata/kernels/dequant_bf16.hpp"
#ifdef STRATA_USE_HIP
#include "wmma_gemm.h"
#endif
#include <cublas_v2.h>
#include <cuda_runtime.h>
@@ -359,6 +363,19 @@ void Gemm::bf16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64
float beta) {
if (T <= 0 || N <= 0) return;
if (ldy <= 0) ldy = N;
#ifdef STRATA_USE_HIP
// STRATA_WMMA_GEMM=1 opts the dense WMMA GEMMs in (0 disables); STRATA_WMMA_BF16=0 excludes just bf16. The arch
// macro that gates the implementation comes from the build (see wmma_gemm.cu); if it is absent the call
// returns false and this falls through, so a non-gfx11 build cannot silently run empty stubs.
// opt-in (review: run after, not before, hipBLASLt by default). Values are parsed, not presence-tested:
// STRATA_WMMA_GEMM=0 and STRATA_WMMA_BF16=0 both disable, "1" (or any non-zero value) enables.
static const bool wmma_on = [] { const char* v = std::getenv("STRATA_WMMA_GEMM"); return v && v[0] != 0 && v[0] != '0'; }();
static const bool bf16_on = [] { const char* v = std::getenv("STRATA_WMMA_BF16"); return !v || (v[0] != 0 && v[0] != '0'); }();
if (wmma_on && bf16_on && T >= 16 && (beta == 0.0f || beta == 1.0f) &&
strata_wmma_gemm_bf16(X, W, Y, T, N, K, ldy, beta, stream_)) {
return;
}
#endif
const float alpha = 1.0f;
#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE)
if (try_hipblaslt(hipblaslt_state_, strata::prefill::hipblaslt::InputType::bf16, X, W, Y, T, N, K, ldy,
@@ -377,6 +394,17 @@ void Gemm::f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_
float beta) {
if (T <= 0 || N <= 0) return;
if (ldy <= 0) ldy = N;
#ifdef STRATA_USE_HIP
// STRATA_WMMA_GEMM=1 opts the dense WMMA GEMMs in (0 disables). The
// arch macro that gates the implementation comes from the build (see wmma_gemm.cu); if it is absent the
// call returns false and this falls through, so a non-gfx11 build cannot silently run empty stubs.
// opt-in: parsed, so STRATA_WMMA_GEMM=0 disables (review: run after, not before, hipBLASLt by default)
static const bool wmma_on = [] { const char* v = std::getenv("STRATA_WMMA_GEMM"); return v && v[0] != 0 && v[0] != '0'; }();
if (wmma_on && T >= 16 && (beta == 0.0f || beta == 1.0f) &&
strata_wmma_gemm_f16(X, W, Y, T, N, K, ldy, beta, stream_)) {
return;
}
#endif
const float alpha = 1.0f;
#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE)
if (try_hipblaslt(hipblaslt_state_, strata::prefill::hipblaslt::InputType::f16, X, W, Y, T, N, K, ldy,
src/prefill/wmma_gemm.cu
+298
View File
@@ -0,0 +1,298 @@
// src/prefill/wmma_gemm.cu - RDNA3 WMMA FP16 & BF16 GEMM for Strata prefill (gfx1100).
//
// Computes Y[t, n] = beta * Y[t, n] + sum_k W[n, k] * X[t, k] using RDNA3 WMMA intrinsics:
// * X is T x K row-major (leading dim K), fp16 or bf16
// * W is N x K row-major (leading dim K), fp16 or bf16
// * Y is T x N row-major with leading dimension ldy >= N, fp32
//
// Hardware: AMD RDNA3 (gfx1100, e.g. RX 7900 XTX)
// * v_wmma_f32_16x16x16_f16_w32 intrinsic (__builtin_amdgcn_wmma_f32_16x16x16_f16_w32)
// * v_wmma_f32_16x16x16_bf16_w32 intrinsic (__builtin_amdgcn_wmma_f32_16x16x16_bf16_w32)
// * Wave32 doubled input fragment: lane t (lane_lo = t & 15) holds row lane_lo of A (16 elements along K)
// and column lane_lo of B (16 elements along K). Lanes 16..31 duplicate lanes 0..15.
// * Wave32 C output mapping: lane t holds column lane_lo of 16x16 output tile,
// with 8 elements alternating rows: row m = 2*i + lane_hi (lane_hi = t >> 4).
// * Tile variants:
// 1. gemm_wmma_64x64_4w: 4 waves (128 threads), 64T x 64N x 16K tile, double-buffered LDS for W.
// 2. gemm_wmma_16x16_1w: 1 wave (32 threads), 16T x 16N x 16K tile, zero LDS/barriers, for small T/N.
#include "wmma_gemm.h"
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include <hip/hip_bfloat16.h>
#include <cstring>
// STRATA_WMMA_GFX11 is defined by the BUILD (CMakeLists.txt, from CMAKE_HIP_ARCHITECTURES), not inferred
// from compiler macros: measured on this toolchain the HOST pass of a HIP compile does not define
// __gfx1100__ but does define __HIP_DEVICE_COMPILE__, so a compiler-macro guard here silently selected the
// "return false" stub at the bottom of this file for the very symbol the engine links - the WMMA path then
// never ran, while the same file compiled with hipcc (as the probes do) took the real branch. One
// build-defined macro is uniform across the host and device passes.
#if defined(STRATA_WMMA_GFX11)
using v8fp32 = float __attribute__((ext_vector_type(8)));
template <typename ElemT>
struct WmmaTraits;
template <>
struct WmmaTraits<_Float16> {
using vec_t = _Float16 __attribute__((ext_vector_type(16)));
__device__ static inline v8fp32 mma(vec_t a, vec_t b, v8fp32 c) {
#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__)
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
#else
(void) a; (void) b; return c; // not a gfx11 device pass: never launched (runtime gate)
#endif
}
};
template <>
struct WmmaTraits<__bf16> {
using vec_t = __bf16 __attribute__((ext_vector_type(16)));
__device__ static inline v8fp32 mma(vec_t a, vec_t b, v8fp32 c) {
#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__)
return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a, b, c);
#else
(void) a; (void) b; return c; // not a gfx11 device pass: never launched (runtime gate)
#endif
}
};
// ===========================================================================
// Variant 1: 16x16_1w - Single-wave kernel for small T / N.
// Zero LDS, zero barrier synchronization, full register residency.
// ===========================================================================
template <typename ElemT>
__global__ void gemm_wmma_16x16_1w(
const uint16_t* __restrict__ X,
const uint16_t* __restrict__ W,
float* __restrict__ Y,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta) {
using vec_t = typename WmmaTraits<ElemT>::vec_t;
const int m_tile = blockIdx.y * 16;
const int n_tile = blockIdx.x * 16;
if (m_tile >= T || n_tile >= N) return;
const int lane = threadIdx.x; // 0..31
const int lane_lo = lane & 15; // 0..15
const int lane_hi = lane >> 4; // 0 or 1
v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0};
const int m_row = m_tile + lane_lo;
const int n_row = n_tile + lane_lo;
for (int k_tile = 0; k_tile < K; k_tile += 16) {
vec_t a_frag, b_frag;
if (m_row < T) {
__builtin_memcpy(&a_frag, X + (int64_t)m_row * K + k_tile, sizeof(a_frag));
} else {
#pragma unroll
for (int i = 0; i < 16; ++i) a_frag[i] = 0;
}
if (n_row < N) {
__builtin_memcpy(&b_frag, W + (int64_t)n_row * K + k_tile, sizeof(b_frag));
} else {
#pragma unroll
for (int i = 0; i < 16; ++i) b_frag[i] = 0;
}
c_acc = WmmaTraits<ElemT>::mma(a_frag, b_frag, c_acc);
}
const int out_n = n_tile + lane_lo;
if (out_n < N) {
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int out_m = m_tile + 2 * i + lane_hi;
if (out_m < T) {
float* dst = Y + (int64_t)out_m * ldy + out_n;
if (beta == 0.0f) {
*dst = c_acc[i];
} else {
*dst = beta * (*dst) + c_acc[i];
}
}
}
}
}
// ===========================================================================
// Variant 2: 64x64_4w - 4 waves per block (128 threads), 64T x 64N tile.
// Double-buffered LDS tile for W (4 KB LDS total), cooperative 128-bit global loads.
// ===========================================================================
template <typename ElemT>
__global__ void gemm_wmma_64x64_4w(
const uint16_t* __restrict__ X,
const uint16_t* __restrict__ W,
float* __restrict__ Y,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta) {
using vec_t = typename WmmaTraits<ElemT>::vec_t;
const int m_tile = blockIdx.y * 64;
const int n_tile = blockIdx.x * 64;
if (m_tile >= T || n_tile >= N) return;
const int tid = threadIdx.x; // 0..127
const int wave_id = tid >> 5; // 0..3
const int lane = tid & 31; // 0..31
const int lane_lo = lane & 15; // 0..15
const int lane_hi = lane >> 4; // 0 or 1
// 4 accumulators per wave covering 4 x 16 N-subtiles
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};
// Double-buffered LDS tile: 64 rows of N x 16 elements of K (2 KB per buffer)
alignas(16) __shared__ ElemT b_lds[2][64][16];
// Thread mapping for cooperative loading of W into LDS (128 threads load 64x16 elements)
// Each thread loads 8 halfs (16 bytes = uint4)
const int row_in_tile = tid >> 1; // 0..63
const int k_sub = (tid & 1) << 3; // 0 or 8
const int actual_n = n_tile + row_in_tile;
auto load_w_into_lds = [&](int buf, int k_tile) {
const int actual_k = k_tile + k_sub;
uint4 val = {0, 0, 0, 0};
if (actual_n < N && actual_k < K) {
const void* ptr = reinterpret_cast<const void*>(W + (int64_t)actual_n * K + actual_k);
val = *reinterpret_cast<const uint4*>(ptr);
}
*reinterpret_cast<uint4*>(&b_lds[buf][row_in_tile][k_sub]) = val;
};
// Pre-fill buffer 0
load_w_into_lds(0, 0);
__syncthreads();
const int m_row = m_tile + wave_id * 16 + lane_lo;
int cur_buf = 0;
for (int k_tile = 0; k_tile < K; k_tile += 16) {
const int next_buf = 1 - cur_buf;
const int k_next = k_tile + 16;
// Prefetch next K-tile of W into LDS
if (k_next < K) {
load_w_into_lds(next_buf, k_next);
}
// Load A fragment from X for current wave's 16 M-rows
vec_t a_frag;
if (m_row < T) {
__builtin_memcpy(&a_frag, X + (int64_t)m_row * K + k_tile, sizeof(a_frag));
} else {
#pragma unroll
for (int i = 0; i < 16; ++i) a_frag[i] = 0;
}
// Read B fragments from LDS and compute WMMA
vec_t b_frag0, b_frag1, b_frag2, b_frag3;
__builtin_memcpy(&b_frag0, &b_lds[cur_buf][0 + lane_lo][0], sizeof(vec_t));
__builtin_memcpy(&b_frag1, &b_lds[cur_buf][16 + lane_lo][0], sizeof(vec_t));
__builtin_memcpy(&b_frag2, &b_lds[cur_buf][32 + lane_lo][0], sizeof(vec_t));
__builtin_memcpy(&b_frag3, &b_lds[cur_buf][48 + lane_lo][0], sizeof(vec_t));
c_acc0 = WmmaTraits<ElemT>::mma(a_frag, b_frag0, c_acc0);
c_acc1 = WmmaTraits<ElemT>::mma(a_frag, b_frag1, c_acc1);
c_acc2 = WmmaTraits<ElemT>::mma(a_frag, b_frag2, c_acc2);
c_acc3 = WmmaTraits<ElemT>::mma(a_frag, b_frag3, c_acc3);
__syncthreads();
cur_buf = next_buf;
}
// Store C to Y (coalesced 64-byte writes per wave half)
const int m_tile_wave = m_tile + wave_id * 16;
auto store_acc = [&](const v8fp32& acc, int n_base) {
const int out_n = n_base + lane_lo;
if (out_n < N) {
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int out_m = m_tile_wave + 2 * i + lane_hi;
if (out_m < T) {
float* dst = Y + (int64_t)out_m * ldy + out_n;
if (beta == 0.0f) {
*dst = acc[i];
} else {
*dst = beta * (*dst) + 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);
}
template <typename ElemT>
static inline bool strata_wmma_gemm_dispatch(const uint16_t* X, const uint16_t* W, float* Y,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta,
void* stream) {
if (!X || !W || !Y) return false;
if (T <= 0 || N <= 0 || K <= 0) return false;
// Runtime gate (review): run only on gfx11 devices, whatever this build compiled for. The build-time
// STRATA_WMMA_GFX11 macro controls whether the intrinsics compile; this check controls whether they run.
{
int dev = 0;
hipDeviceProp_t prop{};
if (hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&prop, dev) != hipSuccess) return false;
if (std::strncmp(prop.gcnArchName, "gfx11", 5) != 0) return false;
}
if (K % 16 != 0) return false;
if (beta != 0.0f && beta != 1.0f) return false;
if (ldy <= 0) ldy = N;
if (ldy < N) return false;
hipStream_t s = static_cast<hipStream_t>(stream);
if (T >= 32 && N >= 32) {
dim3 block(128);
dim3 grid((uint32_t)((N + 63) / 64), (uint32_t)((T + 63) / 64), 1);
gemm_wmma_64x64_4w<ElemT><<<grid, block, 0, s>>>(X, W, Y, T, N, K, ldy, beta);
} else {
dim3 block(32);
dim3 grid((uint32_t)((N + 15) / 16), (uint32_t)((T + 15) / 16), 1);
gemm_wmma_16x16_1w<ElemT><<<grid, block, 0, s>>>(X, W, Y, T, N, K, ldy, beta);
}
return true;
}
bool strata_wmma_gemm_f16(const uint16_t* X, const uint16_t* W, float* Y,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta,
void* stream) {
return strata_wmma_gemm_dispatch<_Float16>(X, W, Y, T, N, K, ldy, beta, stream);
}
bool strata_wmma_gemm_bf16(const uint16_t* X, const uint16_t* W, float* Y,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta,
void* stream) {
return strata_wmma_gemm_dispatch<__bf16>(X, W, Y, T, N, K, ldy, beta, stream);
}
#else
// Non-HIP / Non-GFX11 compilation fallback
bool strata_wmma_gemm_f16(const uint16_t*, const uint16_t*, float*,
int64_t, int64_t, int64_t, int64_t, float,
void*) {
return false;
}
bool strata_wmma_gemm_bf16(const uint16_t*, const uint16_t*, float*,
int64_t, int64_t, int64_t, int64_t, float,
void*) {
return false;
}
#endif
src/prefill/wmma_gemm.h
+22
View File
@@ -0,0 +1,22 @@
// src/prefill/wmma_gemm.h - RDNA3 WMMA FP16 & BF16 GEMM for Strata prefill.
#pragma once
#include <cstdint>
/// Compute Y[t, n] = beta * Y[t, n] + sum_k W[n, k] * X[t, k] using RDNA3 WMMA instructions.
/// X: T x K row-major (leading dim K), fp16 (uint16_t)
/// W: N x K row-major (leading dim K), fp16 (uint16_t)
/// Y: T x N row-major with leading dimension ldy >= N, fp32 (float)
/// Returns true if executed on WMMA, false if unsupported shape/parameters (fall back to hipblas).
bool strata_wmma_gemm_f16(const uint16_t* X, const uint16_t* W, float* Y,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta,
void* stream = nullptr);
/// Compute Y[t, n] = beta * Y[t, n] + sum_k W[n, k] * X[t, k] using RDNA3 WMMA instructions.
/// X: T x K row-major (leading dim K), bf16 (uint16_t)
/// W: N x K row-major (leading dim K), bf16 (uint16_t)
/// Y: T x N row-major with leading dimension ldy >= N, fp32 (float)
/// Returns true if executed on WMMA, false if unsupported shape/parameters (fall back to hipblas).
bool strata_wmma_gemm_bf16(const uint16_t* X, const uint16_t* W, float* Y,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta,
void* stream = nullptr);
tests/hip/prefill_wmma_gemm_parity.cpp
+441
View File
@@ -0,0 +1,441 @@
// tests/hip/prefill_wmma_gemm_parity.cpp - Numerical parity test for RDNA3 WMMA GEMM kernels.
#include <cuda_runtime.h>
#include <hip/hip_bfloat16.h>
#include <hip/hip_fp16.h>
#include "wmma_gemm.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <random>
#include <string>
#include <vector>
// ============================================================================
// Exact Representability & Parity Rationale:
//
// Inputs (X, W, and initial Y for beta=1) are chosen from discrete binary fractions
// k / 16.0f with integer k in [-8, 8].
//
// In IEEE 754 half-precision (fp16, 10-bit explicit mantissa) and bfloat16
// (7-bit explicit mantissa), every such value k * 2^-4 requires at most 4 bits
// of significand and is therefore EXACTLY representable without rounding.
//
// The product of any two such values is an exact integer multiple of 2^-8:
// (a * 2^-4) * (b * 2^-4) = (a * b) * 2^-8
// with (a * b) in [-64, 64].
//
// For K <= 80, the sum of K products:
// sum_k (W[n, k] * X[t, k]) = (sum_k a_k * b_k) * 2^-8
// has a maximum absolute magnitude of 80 * (0.5 * 0.5) = 20.0, and the integer
// numerator cannot exceed 80 * 64 = 5120 (< 2^13).
//
// In IEEE 754 single-precision (fp32, 23-bit explicit mantissa, 24 bits total),
// any integer multiple of 2^-8 up to 20.0 requires at most 13 bits of significand.
// Because the RDNA3 hardware instructions (v_wmma_f32_16x16x16_f16_w32 and
// v_wmma_f32_16x16x16_bf16_w32) accumulate in native FP32 registers, every
// intermediate product and addition has ZERO truncation or rounding error.
//
// Furthermore, floating-point addition of exact binary fractions whose sum fits
// within the mantissa is strictly associative and exact. Consequently, the GPU
// WMMA pipeline and the CPU reference compute identical bitwise results, making
// an exact assertion (max_abs == 0.0f) mathematically legitimate.
// ============================================================================
#define HIP_CHECK(call) do { \
const hipError_t status_ = (call); \
if (status_ != hipSuccess) { \
std::fprintf(stderr, "HIP failure at %s:%d: %s: %s\n", __FILE__, __LINE__, #call, hipGetErrorString(status_)); \
std::exit(2); \
} \
} while (0)
namespace {
constexpr float kSentinel = -999.5f;
struct DeviceBuffer {
void* p = nullptr;
explicit DeviceBuffer(size_t bytes) { if (bytes) HIP_CHECK(hipMalloc(&p, bytes)); }
~DeviceBuffer() { if (p) (void) hipFree(p); }
DeviceBuffer(const DeviceBuffer&) = delete;
DeviceBuffer& operator=(const DeviceBuffer&) = delete;
};
uint16_t encode(float value, bool bf16) {
if (bf16) return hip_bfloat16(value).data;
const __half half = __float2half_rn(value);
return static_cast<__half_raw>(half).x;
}
float decode(uint16_t bits, bool bf16) {
if (bf16) {
const uint32_t u = static_cast<uint32_t>(bits) << 16;
float f = 0.0f;
std::memcpy(&f, &u, sizeof(f));
return f;
}
const __half_raw raw{bits};
return __half2float(__half(raw));
}
struct MismatchRecord {
const char* dtype = "";
int64_t T = 0, N = 0, K = 0, ldy = 0;
float beta = 0.0f;
int64_t row = 0, col = 0;
float got = 0.0f, ref = 0.0f;
double delta = 0.0;
bool is_padding = false;
};
// Independent host CPU reference computed in double precision.
void gemm_reference(const uint16_t* X, const uint16_t* W, const float* initial_Y, float* ref_Y,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta, bool is_bf16) {
std::vector<double> x_dec((size_t) T * K);
std::vector<double> w_dec((size_t) N * K);
for (size_t i = 0; i < x_dec.size(); ++i) x_dec[i] = decode(X[i], is_bf16);
for (size_t i = 0; i < w_dec.size(); ++i) w_dec[i] = decode(W[i], is_bf16);
for (int64_t t = 0; t < T; ++t) {
const double* x_row = x_dec.data() + t * K;
for (int64_t n = 0; n < N; ++n) {
const double* w_row = w_dec.data() + n * K;
double sum = 0.0;
for (int64_t k = 0; k < K; ++k) {
sum += x_row[k] * w_row[k];
}
const size_t out_idx = (size_t) t * ldy + n;
const double y_init = (beta == 0.0f) ? 0.0 : static_cast<double>(initial_Y[out_idx]);
ref_Y[out_idx] = static_cast<float>(beta * y_init + sum);
}
}
}
enum class CaseStatus {
Passed,
Skipped,
Failed
};
CaseStatus run_gemm_case(DeviceBuffer& dx, DeviceBuffer& dw, DeviceBuffer& dy,
hipStream_t stream, bool is_bf16,
int64_t T, int64_t N, int64_t K, int64_t ldy, float beta,
uint32_t seed, double& out_max_abs,
MismatchRecord* first_mismatch, size_t& failure_count) {
const char* dtype_str = is_bf16 ? "bf16" : "f16";
std::mt19937 rng(seed);
std::uniform_int_distribution<int> dist(-8, 8);
const size_t x_count = (size_t) T * K;
const size_t w_count = (size_t) N * K;
const size_t y_count = (size_t) T * ldy + 16; // 16 guard floats past end of matrix
std::vector<uint16_t> x(x_count);
std::vector<uint16_t> w(w_count);
for (size_t i = 0; i < x_count; ++i) x[i] = encode(dist(rng) * 0.0625f, is_bf16);
for (size_t i = 0; i < w_count; ++i) w[i] = encode(dist(rng) * 0.0625f, is_bf16);
std::vector<float> initial_y(y_count, kSentinel);
for (int64_t t = 0; t < T; ++t) {
for (int64_t n = 0; n < N; ++n) {
if (beta == 0.0f) {
initial_y[(size_t) t * ldy + n] = -888.25f;
} else {
// beta == 1.0f: pre-fill with a known exact pattern to verify accumulation
initial_y[(size_t) t * ldy + n] = static_cast<float>((t * 5 + n * 3) % 17 - 8) * 0.0625f;
}
}
}
HIP_CHECK(hipMemcpyAsync(dx.p, x.data(), x_count * sizeof(uint16_t), hipMemcpyHostToDevice, stream));
HIP_CHECK(hipMemcpyAsync(dw.p, w.data(), w_count * sizeof(uint16_t), hipMemcpyHostToDevice, stream));
HIP_CHECK(hipMemcpyAsync(dy.p, initial_y.data(), y_count * sizeof(float), hipMemcpyHostToDevice, stream));
const bool launched = is_bf16 ? strata_wmma_gemm_bf16(static_cast<const uint16_t*>(dx.p),
static_cast<const uint16_t*>(dw.p),
static_cast<float*>(dy.p),
T, N, K, ldy, beta, stream)
: strata_wmma_gemm_f16(static_cast<const uint16_t*>(dx.p),
static_cast<const uint16_t*>(dw.p),
static_cast<float*>(dy.p),
T, N, K, ldy, beta, stream);
if (!launched) {
// Declined by kernel: caller falls back to hipBLAS. Test must SKIP rather than fail.
HIP_CHECK(hipStreamSynchronize(stream));
return CaseStatus::Skipped;
}
std::vector<float> got_y(y_count);
HIP_CHECK(hipMemcpyAsync(got_y.data(), dy.p, y_count * sizeof(float), hipMemcpyDeviceToHost, stream));
HIP_CHECK(hipStreamSynchronize(stream));
// 1. Verify padding columns and end guards are untouched.
for (int64_t t = 0; t < T; ++t) {
for (int64_t col = N; col < ldy; ++col) {
const size_t idx = (size_t) t * ldy + col;
if (got_y[idx] != kSentinel) {
failure_count++;
if (first_mismatch && first_mismatch->dtype[0] == '\0') {
*first_mismatch = MismatchRecord{dtype_str, T, N, K, ldy, beta, t, col,
got_y[idx], kSentinel,
std::abs((double) got_y[idx] - (double) kSentinel), true};
}
std::fprintf(stderr, "FAIL: output touched padding column: dtype=%s T=%ld N=%ld K=%ld ldy=%ld beta=%.1f row=%ld col=%ld got=%.9g expected=%.9g\n",
dtype_str, T, N, K, ldy, beta, t, col, got_y[idx], kSentinel);
return CaseStatus::Failed;
}
}
}
for (size_t g = (size_t) T * ldy; g < y_count; ++g) {
if (got_y[g] != kSentinel) {
failure_count++;
if (first_mismatch && first_mismatch->dtype[0] == '\0') {
*first_mismatch = MismatchRecord{dtype_str, T, N, K, ldy, beta, -1, (int64_t) g,
got_y[g], kSentinel,
std::abs((double) got_y[g] - (double) kSentinel), true};
}
std::fprintf(stderr, "FAIL: output touched end guard: dtype=%s T=%ld N=%ld K=%ld ldy=%ld beta=%.1f index=%zu got=%.9g expected=%.9g\n",
dtype_str, T, N, K, ldy, beta, g, got_y[g], kSentinel);
return CaseStatus::Failed;
}
}
// 2. Compute reference on CPU and verify exact match.
std::vector<float> ref_y(y_count, kSentinel);
gemm_reference(x.data(), w.data(), initial_y.data(), ref_y.data(), T, N, K, ldy, beta, is_bf16);
double case_max_abs = 0.0;
for (int64_t t = 0; t < T; ++t) {
for (int64_t n = 0; n < N; ++n) {
const size_t idx = (size_t) t * ldy + n;
if (!std::isfinite(got_y[idx])) {
failure_count++;
if (first_mismatch && first_mismatch->dtype[0] == '\0') {
*first_mismatch = MismatchRecord{dtype_str, T, N, K, ldy, beta, t, n,
got_y[idx], ref_y[idx], 1e30, false};
}
std::fprintf(stderr, "FAIL: non-finite output: dtype=%s T=%ld N=%ld K=%ld ldy=%ld beta=%.1f at (%ld, %ld): got=%.9g\n",
dtype_str, T, N, K, ldy, beta, t, n, got_y[idx]);
return CaseStatus::Failed;
}
const double delta = std::abs(static_cast<double>(got_y[idx]) - static_cast<double>(ref_y[idx]));
case_max_abs = std::max(case_max_abs, delta);
// The RDNA3 matrix core accumulates in fp32 but rounds toward zero rather than to nearest, while
// this reference sums in double and rounds to nearest, so a correct result may differ by one ULP
// (measured: every mismatch on this hardware is <= 1 ULP and toward zero). Compare against a
// tolerance of 4 ULP - tight enough to stay diagnostic, since a wrong row, column, stride or tile
// boundary differs by O(value) rather than by 1e-7. STRATA_WMMA_PARITY_EXACT=1 demands bit
// equality for anyone who wants to see the raw deltas.
static const bool exact = std::getenv("STRATA_WMMA_PARITY_EXACT") != nullptr;
const double tol = exact ? 0.0
: 4.0 * (double)std::numeric_limits<float>::epsilon() *
std::max(1.0, std::abs(static_cast<double>(ref_y[idx])));
if (delta > tol) {
failure_count++;
if (first_mismatch && first_mismatch->dtype[0] == '\0') {
*first_mismatch = MismatchRecord{dtype_str, T, N, K, ldy, beta, t, n,
got_y[idx], ref_y[idx], delta, false};
}
std::fprintf(stderr, "FAIL: mismatch: dtype=%s T=%ld N=%ld K=%ld ldy=%ld beta=%.1f at (%ld, %ld): got=%.9g ref=%.9g delta=%.9g\n",
dtype_str, T, N, K, ldy, beta, t, n, got_y[idx], ref_y[idx], delta);
return CaseStatus::Failed;
}
}
}
out_max_abs = std::max(out_max_abs, case_max_abs);
return CaseStatus::Passed;
}
struct DeclinedShape {
int64_t T, N, K, ldy;
float beta;
const char* reason;
};
} // namespace
int main() {
int device = 0;
HIP_CHECK(hipGetDevice(&device));
hipDeviceProp_t properties{};
HIP_CHECK(hipGetDeviceProperties(&properties, device));
std::printf("Device %d: %s (arch %s, warpSize %d)\n",
device, properties.name, properties.gcnArchName, properties.warpSize);
hipStream_t stream = nullptr;
HIP_CHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
// Allocate reusable device buffers sized for the maximum test dimensions:
// Max T = 127, Max N = 256, Max K = 80, Max ldy = 256 + 16.
constexpr size_t kMaxXBytes = 127 * 80 * sizeof(uint16_t);
constexpr size_t kMaxWBytes = 256 * 80 * sizeof(uint16_t);
constexpr size_t kMaxYBytes = (127 * (256 + 16) + 64) * sizeof(float);
DeviceBuffer dx(kMaxXBytes);
DeviceBuffer dw(kMaxWBytes);
DeviceBuffer dy(kMaxYBytes);
const std::vector<int64_t> k_t_values = {1, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127};
const std::vector<int64_t> k_n_values = {16, 17, 32, 48, 64, 65, 127, 256};
const std::vector<int64_t> k_k_values = {16, 32, 48, 64, 80};
// Test matrix configurations for (beta, ldy):
// Covers beta=0 (overwrite) and beta=1 (accumulate), with ldy == N and ldy > N.
struct Config {
float beta;
int64_t ldy_offset;
};
const std::vector<Config> kConfigs = {
{0.0f, 0}, // beta=0, packed (ldy == N)
{1.0f, 3}, // beta=1, strided (ldy == N + 3), tests accumulation + padding
{0.0f, 7}, // beta=0, strided (ldy == N + 7), tests overwrite + padding
{1.0f, 0}, // beta=1, packed (ldy == N), tests accumulation in packed layout
};
const std::vector<DeclinedShape> kDeclinedShapes = {
// K not a multiple of 16
{1, 16, 15, 16, 0.0f, "K not multiple of 16 (K=15)"},
{15, 17, 24, 17, 0.0f, "K not multiple of 16 (K=24)"},
{32, 32, 35, 32, 1.0f, "K not multiple of 16 (K=35)"},
{64, 64, 47, 67, 1.0f, "K not multiple of 16 (K=47)"},
{127, 256, 79, 256, 0.0f, "K not multiple of 16 (K=79)"},
// ldy < N
{16, 32, 16, 16, 0.0f, "ldy < N (ldy=16, N=32)"},
{32, 48, 32, 40, 1.0f, "ldy < N (ldy=40, N=48)"},
// beta not 0 or 1
{16, 16, 16, 16, 0.5f, "beta not 0 or 1 (beta=0.5)"},
{32, 32, 32, 32, 2.0f, "beta not 0 or 1 (beta=2.0)"},
// Non-positive sizes
{0, 16, 16, 16, 0.0f, "T <= 0 (T=0)"},
{16, 0, 16, 16, 0.0f, "N <= 0 (N=0)"},
{16, 16, 0, 16, 0.0f, "K <= 0 (K=0)"},
{-1, 16, 16, 16, 0.0f, "T < 0 (T=-1)"},
};
size_t passed_count = 0;
size_t skipped_count = 0;
size_t failure_count = 0;
double global_max_abs = 0.0;
MismatchRecord first_mismatch{};
uint32_t case_seed = 0x574d4d41u;
const bool verbose = (std::getenv("STRATA_TEST_VERBOSE") != nullptr);
for (bool is_bf16 : {false, true}) {
const char* dtype_str = is_bf16 ? "bf16" : "f16";
std::printf("Testing strata_wmma_gemm_%s across %zu shapes (%zu execution configs)...\n",
dtype_str, k_t_values.size() * k_n_values.size() * k_k_values.size(), kConfigs.size());
for (int64_t T : k_t_values) {
double t_max_abs = 0.0;
size_t t_passed = 0;
size_t t_skipped = 0;
for (int64_t N : k_n_values) {
for (int64_t K : k_k_values) {
for (const auto& cfg : kConfigs) {
const int64_t ldy = N + cfg.ldy_offset;
double case_max_abs = 0.0;
case_seed = case_seed * 1664525u + 1013904223u;
const CaseStatus status = run_gemm_case(
dx, dw, dy, stream, is_bf16, T, N, K, ldy, cfg.beta,
case_seed, case_max_abs, &first_mismatch, failure_count);
if (status == CaseStatus::Passed) {
passed_count++;
t_passed++;
t_max_abs = std::max(t_max_abs, case_max_abs);
global_max_abs = std::max(global_max_abs, case_max_abs);
if (verbose) {
std::printf(" case dtype=%s T=%3ld N=%3ld K=%2ld ldy=%3ld beta=%.1f max_abs=%.9g PASS\n",
dtype_str, T, N, K, ldy, cfg.beta, case_max_abs);
}
} else if (status == CaseStatus::Skipped) {
skipped_count++;
t_skipped++;
if (verbose) {
std::printf(" case dtype=%s T=%3ld N=%3ld K=%2ld ldy=%3ld beta=%.1f SKIP (declined)\n",
dtype_str, T, N, K, ldy, cfg.beta);
}
} else {
// Failed: first_mismatch already recorded and reported to stderr
if (verbose) {
std::printf(" case dtype=%s T=%3ld N=%3ld K=%2ld ldy=%3ld beta=%.1f FAIL\n",
dtype_str, T, N, K, ldy, cfg.beta);
}
}
}
}
}
std::printf(" dtype=%s T=%3ld: %zu passed, %zu skipped, max_abs=%.9g %s\n",
dtype_str, T, t_passed, t_skipped, t_max_abs,
(failure_count == 0) ? "PASS" : "FAIL");
}
// Test declined shapes: the kernel must return false and be skipped.
for (const auto& decl : kDeclinedShapes) {
const bool launched = is_bf16
? strata_wmma_gemm_bf16(static_cast<const uint16_t*>(dx.p),
static_cast<const uint16_t*>(dw.p),
static_cast<float*>(dy.p),
decl.T, decl.N, decl.K, decl.ldy, decl.beta, stream)
: strata_wmma_gemm_f16(static_cast<const uint16_t*>(dx.p),
static_cast<const uint16_t*>(dw.p),
static_cast<float*>(dy.p),
decl.T, decl.N, decl.K, decl.ldy, decl.beta, stream);
if (launched) {
failure_count++;
std::fprintf(stderr, "FAIL: expected kernel to decline shape (%s) dtype=%s T=%ld N=%ld K=%ld ldy=%ld beta=%.1f but it returned true!\n",
decl.reason, dtype_str, decl.T, decl.N, decl.K, decl.ldy, decl.beta);
} else {
skipped_count++;
if (verbose) {
std::printf(" declined shape correctly rejected (%s): dtype=%s T=%ld N=%ld K=%ld ldy=%ld beta=%.1f SKIP\n",
decl.reason, dtype_str, decl.T, decl.N, decl.K, decl.ldy, decl.beta);
}
}
}
}
HIP_CHECK(hipStreamDestroy(stream));
if (failure_count > 0) {
if (first_mismatch.dtype[0] != '\0') {
if (first_mismatch.is_padding) {
std::fprintf(stderr, "\nSUMMARY: First mismatch was padding/guard corruption:\n"
" dtype=%s T=%ld N=%ld K=%ld ldy=%ld beta=%.1f at (%ld, %ld): got=%.9g expected=%.9g\n",
first_mismatch.dtype, first_mismatch.T, first_mismatch.N, first_mismatch.K,
first_mismatch.ldy, first_mismatch.beta, first_mismatch.row, first_mismatch.col,
first_mismatch.got, first_mismatch.ref);
} else {
std::fprintf(stderr, "\nSUMMARY: First mismatch occurred at:\n"
" dtype=%s T=%ld N=%ld K=%ld ldy=%ld beta=%.1f at (%ld, %ld): got=%.9g ref=%.9g delta=%.9g\n",
first_mismatch.dtype, first_mismatch.T, first_mismatch.N, first_mismatch.K,
first_mismatch.ldy, first_mismatch.beta, first_mismatch.row, first_mismatch.col,
first_mismatch.got, first_mismatch.ref, first_mismatch.delta);
}
}
std::fprintf(stderr, "WMMA GEMM parity FAILED: %zu failures detected, max_abs=%.9g\n",
failure_count, global_max_abs);
return 1;
}
if (passed_count == 0) {
std::fprintf(stderr, "WMMA GEMM parity FAILED: zero test cases passed (all were declined or skipped)\n");
return 1;
}
std::printf("\nWMMA GEMM parity OK: %zu test cases passed, %zu declined shapes skipped, max_abs=%.9g\n",
passed_count, skipped_count, global_max_abs);
return 0;
}