wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Share gfx12 K enc0 V tiles while preserving strict FP16 arithmetic

Browse Source
This commit is contained in:
xdfnx-dev committed 2026-10-02 15:03:32 +03:00
1 parent 28832eef6c
commit a2a4a7e765
8 files changed
+159 -36
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels/k/dltss_pwin_dec0_layer.hip
+6 -1
View File
@@ -1,5 +1,10 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#ifndef D4R_K_POS_MT
#define D4R_K_POS_MT 2
#elif !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
#error "experimental position token tiling requires gfx12 strict K baseline"
#endif
#include "pwin_pos.h"
POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, 2)
POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, D4R_K_POS_MT)
kernels/k/dltss_pwin_enc0_layer.hip
+15 -1
View File
@@ -1,5 +1,19 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#ifndef D4R_K_POS_MT
#define D4R_K_POS_MT 2
#elif !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
#error "experimental position token tiling requires gfx12 strict K baseline"
#endif
// Reuse V across the window's waves. Paired gfx1201 replays retain every half
// code and reduce enc0 time by about 5%; gfx11 and arithmetic experiments keep
// their previous path. D4R_K_PRIVATE_V supplies the original A/B control.
#if defined(__GFX12__) && defined(D4R_K_FP16_BASELINE) && !defined(D4R_K_PRIVATE_V) && \
!defined(D4R_K_PACKED_ACC) && !defined(D4R_K_F16_WMMA)
#ifndef POS_VSHARE
#define POS_VSHARE
#endif
#endif
#include "pwin_pos.h"
POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, 2)
POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, D4R_K_POS_MT)
kernels/k/dltss_pwin_enc1_layer.hip
+11
View File
@@ -1,5 +1,16 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#ifdef D4R_K_ENC1_POS_MT
#if !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
#error "experimental enc1 token tiling requires gfx12 strict K baseline"
#endif
#if D4R_K_ENC1_POS_MT != 1 && D4R_K_ENC1_POS_MT != 2 && D4R_K_ENC1_POS_MT != 4
#error "enc1 token tiles per wave must be 1, 2 or 4"
#endif
#include "pwin_pos.h"
POS_ENCODER(dltss_pwin_enc1_layer, 2, 64, 64, 0, D4R_K_ENC1_POS_MT)
#else
#include "pwin_layer.h"
PWIN_ENCODER_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)
#endif
kernels/k/pwin_common.h
+29 -5
View File
@@ -23,6 +23,16 @@ typedef _Float16 half_t;
typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
typedef _Float16 h16 __attribute__((ext_vector_type(16)));
typedef float f8v __attribute__((ext_vector_type(8)));
#ifdef D4R_K_PACKED_ACC
#if !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE) || defined(D4R_K_F16_WMMA)
#error "packed accumulator storage requires gfx12 strict F32 WMMA baseline"
#endif
// Keep already-rounded values packed between steps. Every matrix instruction
// still uses F32 C/D and rounds once to half afterwards, like the baseline.
typedef _Float16 acc8v __attribute__((ext_vector_type(8)));
#else
typedef f8v acc8v;
#endif
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
@@ -78,9 +88,13 @@ __device__ __forceinline__ half_t l2_sum(const half_t* row)
// one k16 step with the f16 accumulator of NVIDIA's f16 wmma (rounded after the step)
// PWIN_F32ACC: keep the accumulator in f32 through the chain (rounded to f16 where the values are used)
__device__ __forceinline__ f8v mma16(const op_t& a, const op_t& b, f8v c)
__device__ __forceinline__ acc8v mma16(const op_t& a, const op_t& b, acc8v c)
{
#ifdef D4R_K_F16_WMMA
#ifdef D4R_K_PACKED_ACC
const f8v fc = __builtin_convertvector(c, f8v);
const f8v d = wm_mma(a, b, fc);
return __builtin_convertvector(d, acc8v);
#elif defined(D4R_K_F16_WMMA)
#if !defined(__GFX12__) || D4R_WMMA_LAYOUT != 12 || !defined(D4R_K_FP16_BASELINE)
#error "experimental packed F16 WMMA requires gfx12 native layout and strict K baseline"
#endif
@@ -96,9 +110,14 @@ __device__ __forceinline__ f8v mma16(const op_t& a, const op_t& b, f8v c)
#endif
}
__device__ __forceinline__ f8v splat8(float v)
__device__ __forceinline__ acc8v splat8(float v)
{
return (f8v){v, v, v, v, v, v, v, v};
#ifdef D4R_K_PACKED_ACC
const half_t h = (half_t)v;
return (acc8v){h, h, h, h, h, h, h, h};
#else
return (acc8v){v, v, v, v, v, v, v, v};
#endif
}
__device__ __forceinline__ uint32_t other_half(uint32_t v)
@@ -127,9 +146,14 @@ __device__ __forceinline__ op_t operand_from_dt(const half_t own[8])
return wm_op_from_acc(own);
}
__device__ __forceinline__ op_t operand_from_f8(f8v d)
__device__ __forceinline__ op_t operand_from_f8(acc8v d)
{
#ifdef D4R_K_PACKED_ACC
const half_t h[8] = {d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7]};
return wm_op_from_acc(h);
#else
return wm_op_from_f8(d);
#endif
}
// element c (0..15) of a row of 16 halves (VALU code; not a WMMA operand)
kernels/k/pwin_layer.h
+12 -12
View File
@@ -163,7 +163,7 @@ __device__ __forceinline__ half_t gelu(half_t x)
}
// gelu of 8 accumulator values as 4 packed f16 pairs (the same f16 operations as gelu(), element-wise)
__device__ __forceinline__ void gelu8(const f8v& d, half_t g[8])
__device__ __forceinline__ void gelu8(const acc8v& d, half_t g[8])
{
const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
#pragma unroll
@@ -283,7 +283,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
block_sync(); // V of all 64 tokens is computed by every wave from hb
STAMP(2);
// ---- attention, one head at a time; the output projection accumulates head by head
f8v acc[KT];
acc8v acc[KT];
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
@@ -310,7 +310,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
f8v d = splat8(0.0f);
acc8v d = splat8(0.0f);
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * wv + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
op_img_store(&vsh[h][nt][0][0], wv * 16 + m, operand_from_f8(d));
@@ -333,7 +333,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
op_t qop[2];
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
f8v dq[6];
acc8v dq[6];
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = splat8(0.0f);
@@ -349,7 +349,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const f8v d = dq[2 * j + nt];
const acc8v d = dq[2 * j + nt];
if (j == 0)
qop[nt] = operand_from_f8(d);
else if (j == 1)
@@ -370,7 +370,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int kt = 0; kt < 4; ++kt)
{
const u4v bi = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
f8v d;
acc8v d;
#pragma unroll
for (int i = 0; i < 4; ++i)
{
@@ -404,7 +404,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
f8v d = splat8(0.0f);
acc8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
@@ -462,7 +462,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
#pragma unroll
for (int hn = 0; hn < 2; ++hn)
{
f8v d;
acc8v d;
{
half_t vv_[8];
@@ -555,7 +555,7 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
#pragma unroll
for (int nt = 0; nt < C / 16; ++nt)
{
f8v d;
acc8v d;
{
half_t vv_[8];
dvec8(p0.w, 32 * nt, vv_);
@@ -603,7 +603,7 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
const uint8_t* w = p.w;
for (int vt = wv; vt < L::PMT; vt += 4)
{
f8v d;
acc8v d;
{
half_t vv_[8];
@@ -753,7 +753,7 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
for (int nt = 0; nt < C / 16; ++nt)
{
const int gnt = wv * (C / 16) + nt; // expand output column tile
f8v d;
acc8v d;
{
half_t vv_[8];
@@ -790,7 +790,7 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
#pragma unroll
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
{
f8v d;
acc8v d;
{
half_t vv_[8];
kernels/k/pwin_pos.h
+10 -9
View File
@@ -8,6 +8,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
half_t (*act)[L::C], half_t (*hb)[L::C])
{
constexpr int C = L::C, KT = L::KT, H = L::H, NW = 4 / MT;
static_assert(MT == 1 || MT == 2 || MT == 4, "position token tiles per wave");
static_assert(L::POS, "position-only attention layers");
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = threadIdx.z;
@@ -37,7 +38,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
STAMP(2);
// ---- attention: O = P V per head (P from the table), output projection accumulated head by head
f8v acc[MT][KT];
acc8v acc[MT][KT];
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
@@ -64,7 +65,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
for (int mi = 0; mi < MT; ++mi)
{
const int kt = wv * MT + mi;
f8v d = splat8(0.0f);
acc8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
@@ -82,7 +83,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
f8v d = splat8(0.0f);
acc8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
@@ -101,7 +102,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
f8v d = splat8(0.0f);
acc8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
d = mma16(vtop[nt][kt], pop[kt], d);
@@ -156,7 +157,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
{
half_t bb[8];
dvec8(w, L::B1 + 64 * c + 32 * hn, bb);
f8v d[MT];
acc8v d[MT];
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
#pragma unroll
@@ -231,7 +232,7 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
{
half_t vv[8];
dvec8(p0.w, 32 * nt, vv);
f8v d;
acc8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
@@ -278,7 +279,7 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
{
half_t vv[8];
dvec8(p.w, L::PMB + 32 * vt, vv);
f8v d;
acc8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
@@ -322,7 +323,7 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
const int gnt = q * (C / 16) + nt;
half_t vv[8];
dvec8(p.w, L::EXPB + 32 * gnt, vv);
f8v d;
acc8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
@@ -353,7 +354,7 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
{
half_t vv[8];
dvec8(q.w, L::HB + 32 * nt, vv);
f8v d;
acc8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
kernels/k/pwin_wide.h
+8 -8
View File
@@ -68,7 +68,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
block_sync(); // hb complete, R free
// ---- attention: NG heads per round; the output projection accumulates head by head
f8v acc[NTL];
acc8v acc[NTL];
#pragma unroll
for (int j = 0; j < NTL; ++j)
{
@@ -93,7 +93,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
if (h < H)
{
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
f8v dq[6];
acc8v dq[6];
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = splat8(0.0f);
@@ -109,7 +109,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const f8v d = dq[2 * j + nt];
const acc8v d = dq[2 * j + nt];
if (j == 0)
qop[nt] = operand_from_f8(d);
else if (j == 1)
@@ -133,7 +133,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int kt = 0; kt < 4; ++kt)
{
const u4v bi = bias[((h * 4 + t) * 4 + kt) * 32 + l];
f8v d;
acc8v d;
#pragma unroll
for (int i = 0; i < 4; ++i)
{
@@ -166,7 +166,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
f8v d = splat8(0.0f);
acc8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
d = mma16(op_lds(&vt[g][16 * nt + m][16 * kt]), pop[kt], d);
@@ -233,7 +233,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
if (c >= L::NMLP)
break;
f8v d;
acc8v d;
{
half_t vv[8];
dvec8(w, L::B1 + 64 * c + 32 * hn, vv);
@@ -329,7 +329,7 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
const uint8_t* w = p.w;
for (int vt = wv; vt < L::PMT; vt += 4 * NG)
{
f8v d;
acc8v d;
{
half_t vv[8];
dvec8(w, L::PMB + 32 * vt, vv);
@@ -382,7 +382,7 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
for (int nt = g; nt < C / 16; nt += NG)
{
const int gnt = q * (C / 16) + nt;
f8v d;
acc8v d;
{
half_t vv[8];
dvec8(p.w, L::EXPB + 32 * gnt, vv);