Port K kernel tuning and M gfx1101 flags
Preset K, byte-identical to the previous output in all captured buffers and the harness image, measured on an RX 7700 XT: - LDS row padding so WMMA operand loads hit distinct banks - WGP mode for enc1/enc2/dec1/dec2, ds_bpermute half-wave exchange for enc0/dec0/dec3, more waves per window in enc4/dec4/dec5 - cheaper operand plumbing in common/wmma_layout.h (packed f16 conversions, no dynamic vector indexing), raw buffer weight loads, accumulator-init vectors prepared once - enc1 as a single-buffer recompute layer, dec1 Q/K/V chains split - dec5 as four dependent phase kernels
This commit is contained in:
ayo committed
1 parent
8ce8618181
commit
22026ede12
45 files changed
+4514
-354
No files matched your search
@@ -1,5 +1,11 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers). The WMMA half
|
||||
// exchange goes through ds_bpermute (VALU-bound layer: -1..-3%)
|
||||
// f32 -> f16 pairs converted straight into one register (WM_CVT_ASM: no v_pack copies; -1..-3 %)
|
||||
// norm / residual / embedding ReLU on packed f16 pairs (PWIN_PK_ELEM with the packed g2/b2 table; same ops, -1 %)
|
||||
// d4r-build-flags: -mcumode -DWM_HALF_BPERMUTE -DWM_CVT_ASM -DPWIN_PK_ELEM -DPWIN_PACKED_RESIDUAL
|
||||
// RDNA4: no WM_CVT_ASM (its inline asm needs the compiler's real-true16 mode on gfx12; the plain conversion
|
||||
// costs one pack per pair and gfx12 operands need no lane exchange anyway)
|
||||
// d4r-build-flags-gfx12: -mcumode -DWM_HALF_BPERMUTE -DPWIN_PK_ELEM -DPWIN_PACKED_RESIDUAL
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_pos.h"
|
||||
POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, 2)
|
||||
@@ -1,5 +1,8 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
// WGP mode: with the padded LDS rows (PWIN_PAD) this layer is 2-5% faster in WGP mode than in CU mode
|
||||
// (CU mode was 1-6% faster before the padding)
|
||||
// accumulator init vectors from the weights (the prepared f32 table is 2-4% slower here)
|
||||
// Q, K, V chains two at a time (fewer live accumulators, spills 164 -> 68 B; -2 %)
|
||||
// d4r-build-flags: -DPWIN_VEC_FROM_WEIGHTS -DPWIN_QKV_SPLIT
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_layer.h"
|
||||
PWIN_DECODER(dltss_pwin_dec1_layer, 2, 64, 64)
|
||||
@@ -1,5 +1,7 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
// WGP mode: with the padded LDS rows (PWIN_PAD) this layer is 2-5% faster in WGP mode than in CU mode
|
||||
// (CU mode was 1-6% faster before the padding)
|
||||
// weight images through plain global loads here (buffer loads were 1-3% slower in this layer)
|
||||
// d4r-build-flags: -DPWIN_NO_BUFFER_IMG -DPWIN_VEC_FROM_WEIGHTS
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_layer.h"
|
||||
PWIN_DECODER(dltss_pwin_dec2_layer, 2, 64, 96)
|
||||
@@ -1,2 +1,4 @@
|
||||
// the WMMA half exchange through ds_bpermute instead of v_permlanex16: -1.5% here
|
||||
// d4r-build-flags: -DWM_HALF_BPERMUTE
|
||||
#include "pwin_wide.h"
|
||||
WIDE_DECODER(dltss_pwin_dec3_layer, 4, 96, 128, 2)
|
||||
@@ -1,2 +1,3 @@
|
||||
// NG = 8 (32 waves per window): with the padded LDS rows faster than NG = 4 here
|
||||
#include "pwin_wide.h"
|
||||
WIDE_DECODER(dltss_pwin_dec4_layer, 4, 128, 160, 4)
|
||||
WIDE_DECODER(dltss_pwin_dec4_layer, 4, 128, 160, 8)
|
||||
@@ -1,2 +1,4 @@
|
||||
#include "pwin_wide.h"
|
||||
WIDE_PLAIN(dltss_pwin_dec5_layer, 8, 160, 4)
|
||||
// dec5 (bottleneck, H 8, C 160) as four phase kernels (pwin_phase.h). Needs the patched
|
||||
// ZLUDA hook (post kernels). MAXW 64 windows (1440p: 9).
|
||||
#include "pwin_phase.h"
|
||||
PHASE_PLAIN(dltss_pwin_dec5_layer, 8, 160, 64, 128, 96, 384, 96)
|
||||
@@ -1,5 +1,13 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers). The WMMA half
|
||||
// exchange goes through ds_bpermute (VALU-bound layer: -1..-3%)
|
||||
// f32 -> f16 pairs converted straight into one register (WM_CVT_ASM: no v_pack copies; -1..-3 %)
|
||||
// accumulator-init vectors: f32 table read with scalar loads + one select per value (PWIN_VEC_SCALAR, -2 %; replaces
|
||||
// PWIN_VEC_FROM_WEIGHTS)
|
||||
// norm / residual / embedding ReLU on packed f16 pairs (PWIN_PK_ELEM with the packed g2/b2 table; same ops, -1 %)
|
||||
// d4r-build-flags: -mcumode -DWM_HALF_BPERMUTE -DPWIN_VEC_SCALAR -DWM_CVT_ASM -DPWIN_PK_ELEM -DPWIN_PACKED_RESIDUAL
|
||||
// RDNA4: no WM_CVT_ASM (its inline asm needs the compiler's real-true16 mode on gfx12; the plain conversion
|
||||
// costs one pack per pair and gfx12 operands need no lane exchange anyway)
|
||||
// d4r-build-flags-gfx12: -mcumode -DWM_HALF_BPERMUTE -DPWIN_VEC_SCALAR -DPWIN_PK_ELEM -DPWIN_PACKED_RESIDUAL
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_pos.h"
|
||||
POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, 2)
|
||||
@@ -1,5 +1,5 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
// Exact single-buffer enc1: reload the original input for the residual, preserve WGP mode.
|
||||
// d4r-build-flags: -DPOS_VSHARE1 -DPWIN_VEC_SCALAR -DPWIN_PK_ELEM
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_layer.h"
|
||||
PWIN_ENCODER_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)
|
||||
#include "pwin_pos_recompute.h"
|
||||
POS_RECOMPUTE_ENCODER(dltss_pwin_enc1_layer, 2, 64, 64, 0, 1)
|
||||
@@ -1,5 +1,7 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
// WGP mode: with the padded LDS rows (PWIN_PAD) this layer is 2-5% faster in WGP mode than in CU mode
|
||||
// (CU mode was 1-6% faster before the padding)
|
||||
// weight images through plain global loads here (buffer loads were 1-3% slower in this layer)
|
||||
// d4r-build-flags: -DPWIN_NO_BUFFER_IMG -DPWIN_VEC_FROM_WEIGHTS
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_layer.h"
|
||||
PWIN_ENCODER(dltss_pwin_enc2_layer, 2, 64, 96)
|
||||
@@ -1,2 +1,3 @@
|
||||
// NG = 8 (32 waves per window): with the padded LDS rows faster than NG = 4 here
|
||||
#include "pwin_wide.h"
|
||||
WIDE_ENCODER(dltss_pwin_enc4_layer, 4, 128, 160, 4)
|
||||
WIDE_ENCODER(dltss_pwin_enc4_layer, 4, 128, 160, 8)
|
||||
@@ -13,7 +13,10 @@
|
||||
#if defined(D4R_ACCURACY) && defined(PWIN_F32ACC)
|
||||
#undef PWIN_F32ACC
|
||||
#endif
|
||||
// D4R_FAST_MATH: the opt-in approximate variant (FMA contraction allowed; not bit-identical to NVIDIA's rounding)
|
||||
#ifndef D4R_FAST_MATH
|
||||
#pragma clang fp contract(off)
|
||||
#endif
|
||||
|
||||
typedef _Float16 half_t;
|
||||
typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
|
||||
@@ -22,6 +25,32 @@ typedef float f8v __attribute__((ext_vector_type(8)));
|
||||
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
|
||||
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
|
||||
|
||||
__device__ __forceinline__ half_t half_reciprocal(half_t h)
|
||||
{
|
||||
#ifdef PWIN_NATIVE_HRECIP
|
||||
return __builtin_amdgcn_rcph(h);
|
||||
#else
|
||||
return (half_t)__builtin_amdgcn_rcpf((float)h);
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t half_rsqrt(half_t h)
|
||||
{
|
||||
#ifdef PWIN_NATIVE_HRECIP
|
||||
return __builtin_amdgcn_rsqh(h);
|
||||
#else
|
||||
return (half_t)__builtin_amdgcn_rsqf((float)h);
|
||||
#endif
|
||||
}
|
||||
|
||||
// LDS row padding in halves: activation rows of N halves are stored N + PWIN_PAD apart, so the 16 token rows a
|
||||
// WMMA operand load (op_lds / lds_row16) touches start on distinct banks (row stride = 16 mod 32 bytes). With
|
||||
// unpadded 128- or 256-byte rows every lane hits the same banks (measured: 6x slower WMMA feed). 0 = old layout.
|
||||
#ifndef PWIN_PAD
|
||||
#define PWIN_PAD 8
|
||||
#endif
|
||||
#define LDSR(N) ((N) + PWIN_PAD)
|
||||
|
||||
__device__ __forceinline__ uint32_t lane_id()
|
||||
{
|
||||
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
|
||||
@@ -33,11 +62,24 @@ typedef wm_op op_t; // WMMA operand: u8v (gfx11) or u4v (gfx12)
|
||||
// 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)
|
||||
{
|
||||
#if defined(PWIN_NATIVE_F16ACC) && !defined(PWIN_F32ACC) && D4R_WMMA_LAYOUT == 11
|
||||
h16 hc = {};
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
hc[2 * i] = (half_t)c[i];
|
||||
const h16 hd = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(
|
||||
__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), hc, false);
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)hd[2 * i];
|
||||
#else
|
||||
f8v d = wm_mma(a, b, c);
|
||||
#ifndef PWIN_F32ACC
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)(half_t)d[i];
|
||||
#endif
|
||||
#endif
|
||||
return d;
|
||||
}
|
||||
@@ -136,6 +178,46 @@ __device__ __forceinline__ op_t op_img(const u8v* __restrict__ img, int idx)
|
||||
{
|
||||
return wm_op_image(img, idx);
|
||||
}
|
||||
// 8 floats of image slot `slot` + this lane's half (f32 accumulator-init vectors the prep kernel stores after the
|
||||
// operand tiles: slot pair (2 v, 2 v + 1) = rows wm_acc_row_h(., h) of vector v), slot uniform
|
||||
__device__ __forceinline__ f8v img_vec(const u8v* __restrict__ img, int slot)
|
||||
{
|
||||
const uint32_t h = lane_id() >> 4;
|
||||
#if defined(PWIN_VEC_SCALAR) // both layouts: slot pair (2 v, 2 v + 1) = the two halves' rows
|
||||
// both halves' slots with wave-uniform (scalar) loads, then one select per value (instead of a per-lane
|
||||
// vector load, or 8 shifts + 8 f16 -> f32 conversions from the weights); the same f32 values
|
||||
const u8v x0 = img[slot], x1 = img[slot + 1];
|
||||
u8v x;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
x[i] = h ? x1[i] : x0[i];
|
||||
return __builtin_bit_cast(f8v, x);
|
||||
#elif D4R_WMMA_LAYOUT == 12 || defined(PWIN_NO_BUFFER_IMG)
|
||||
const u8v x = img[slot + (int)h];
|
||||
return __builtin_bit_cast(f8v, x);
|
||||
#else
|
||||
const __amdgpu_buffer_rsrc_t r = __builtin_amdgcn_make_buffer_rsrc((void*)img, (short)0, 0x7fffffff, 0x31004000);
|
||||
const uint32_t vo = h * 32u, so = 32u * (uint32_t)slot;
|
||||
const u4v a = __builtin_amdgcn_raw_buffer_load_b128(r, vo, so, 0), b = __builtin_amdgcn_raw_buffer_load_b128(r, vo + 16, so, 0);
|
||||
return __builtin_bit_cast(f8v, (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]});
|
||||
#endif
|
||||
}
|
||||
// slot (tile, lane l & 15) of an operand image in global memory: tile is wave-uniform, so the row address stays a
|
||||
// scalar base and only the lane indexes it (global loads with an SGPR base and a VGPR offset; written as one
|
||||
// index tile * 16 + m every load paid 64-bit VALU address arithmetic)
|
||||
__device__ __forceinline__ op_t op_imgt(const u8v* __restrict__ img, int tile)
|
||||
{
|
||||
#if D4R_WMMA_LAYOUT == 12 || defined(PWIN_NO_BUFFER_IMG)
|
||||
return wm_op_image(img + 16 * tile, (int)(lane_id() & 15));
|
||||
#else
|
||||
// gfx11: a raw buffer load (resource from img: base in SGPRs, 0x31004000 = gfx11 raw dword3, validated), the
|
||||
// tile offset as the scalar offset and the lane's 32-byte slot as the vector offset: no VALU address math
|
||||
const __amdgpu_buffer_rsrc_t r = __builtin_amdgcn_make_buffer_rsrc((void*)img, (short)0, 0x7fffffff, 0x31004000);
|
||||
const uint32_t vo = (lane_id() & 15) * 32u, so = 512u * (uint32_t)tile;
|
||||
const u4v a = __builtin_amdgcn_raw_buffer_load_b128(r, vo, so, 0), b = __builtin_amdgcn_raw_buffer_load_b128(r, vo + 16, so, 0);
|
||||
return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
|
||||
#endif
|
||||
}
|
||||
// operand -> slot idx of an operand image in LDS (read back with op_img)
|
||||
__device__ __forceinline__ void op_img_store(u8v* img, int idx, const op_t& v)
|
||||
{
|
||||
@@ -184,3 +266,25 @@ __device__ __forceinline__ int frag_offset(int k, int n)
|
||||
{
|
||||
return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1);
|
||||
}
|
||||
|
||||
// Diagnostic phase stamps (-DPWIN_PROF): per-wave SHADER_CYCLES deltas summed into d4r_dbg[k], wave counts into
|
||||
// d4r_dbg[32 + k] (read back by bench/native/pwbench). Region k ends at PROF(k). Off: no code.
|
||||
#ifdef PWIN_PROF
|
||||
__device__ unsigned long long d4r_dbg[64];
|
||||
#define PROF_INIT uint32_t prof_last_ = (uint32_t)__builtin_readcyclecounter()
|
||||
#define PROF(k) \
|
||||
do \
|
||||
{ \
|
||||
const uint32_t n_ = (uint32_t)__builtin_readcyclecounter(); \
|
||||
if (lane_id() == 0) \
|
||||
{ \
|
||||
__hip_atomic_fetch_add(&d4r_dbg[k], (unsigned long long)((n_ - prof_last_) & 0xfffffu), __ATOMIC_RELAXED, \
|
||||
__HIP_MEMORY_SCOPE_AGENT); \
|
||||
__hip_atomic_fetch_add(&d4r_dbg[32 + (k)], 1ull, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); \
|
||||
} \
|
||||
prof_last_ = n_; \
|
||||
} while (0)
|
||||
#else
|
||||
#define PROF_INIT do {} while (0)
|
||||
#define PROF(k) do {} while (0)
|
||||
#endif
|
||||
@@ -0,0 +1,37 @@
|
||||
// PWIN_GELU_LUT: exact f16 GELU by exhaustive lookup, populated on the prep stream.
|
||||
// Include after gelu8_math. The POS prep wrappers reserve at least 65536 work-items.
|
||||
// Every index comes from a uint16_t f16 bit pattern, so it is inside this allocation.
|
||||
#pragma once
|
||||
|
||||
__device__ __attribute__((aligned(256))) uint16_t g_gelu_lut[65536];
|
||||
|
||||
constexpr int pwin_gelu_prep_items(int n) { return n < 65536 ? 65536 : n; }
|
||||
|
||||
__device__ __forceinline__ void pwin_gelu_prep()
|
||||
{
|
||||
const unsigned i = blockIdx.x * 128u + threadIdx.x;
|
||||
if (i < 65536u)
|
||||
{
|
||||
const half_t x = __builtin_bit_cast(half_t, (uint16_t)i);
|
||||
half_t g[8];
|
||||
gelu8_math(splat8((float)x), g);
|
||||
g_gelu_lut[i] = __builtin_bit_cast(uint16_t, g[0]);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void gelu8_lut(const f8v& d, half_t g[8])
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const uint16_t bits = __builtin_bit_cast(uint16_t, (half_t)d[i]);
|
||||
#ifdef PWIN_GELU_LUT_BUFFER
|
||||
const __amdgpu_buffer_rsrc_t r = __builtin_amdgcn_make_buffer_rsrc(
|
||||
(void*)g_gelu_lut, (short)0, sizeof(g_gelu_lut), 0x31004000);
|
||||
const uint16_t value = __builtin_amdgcn_raw_buffer_load_b16(r, (unsigned)bits * 2u, 0u, 0);
|
||||
#else
|
||||
const uint16_t value = g_gelu_lut[bits];
|
||||
#endif
|
||||
g[i] = __builtin_bit_cast(half_t, value);
|
||||
}
|
||||
}
|
||||
+501
-124
@@ -38,7 +38,12 @@ template <int H_, int C_, int COUT_, bool POS_ = false> struct PwinLayout
|
||||
static constexpr int T_PM = T_W2 + NMLP * 2 * KT; // [ks C/4][nt COUT/16]
|
||||
static constexpr int T_END = T_PM + (C / 4) * PMT;
|
||||
static constexpr int NBIAS = H * 4 * 4 * 32; // u4v per lane
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + NBIAS;
|
||||
// f32 accumulator-init vectors in image slots after the tiles (see img_vec): Wo bias, MLP fc1 biases, merge bias
|
||||
static constexpr int V_BO = 0, V_B1 = V_BO + KT, V_PM = V_B1 + 2 * NMLP, NV = V_PM + PMT;
|
||||
static constexpr int VEC0 = T_END * 16, IMG_SLOTS = VEC0 + 2 * NV;
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + NBIAS + 2 * NV;
|
||||
|
||||
__device__ static const uint8_t* extra_vec(const uint8_t*, const uint8_t*, int) { return nullptr; }
|
||||
|
||||
__device__ static int src_of(int tile)
|
||||
{
|
||||
@@ -87,9 +92,101 @@ __device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
|
||||
return r;
|
||||
}
|
||||
|
||||
// one (vector, half) slot of the f32 vector table: the 8 rows wm_acc_row_h(i, h) of the 16 halves at src, as floats
|
||||
__device__ __forceinline__ void vec_slot(u8v* img, int slot, const uint8_t* src, int h)
|
||||
{
|
||||
f8v f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
f[i] = (float)wload(src, 2 * wm_acc_row_h(i, h));
|
||||
img[slot] = __builtin_bit_cast(u8v, f);
|
||||
}
|
||||
|
||||
// Optional 640-byte module table for G2/B2 in each half-wave's accumulator order.
|
||||
#if defined(PWIN_PACKED_RESIDUAL) && D4R_WMMA_LAYOUT == 11
|
||||
__device__ __attribute__((aligned(16))) u4v g_residual_vectors[40];
|
||||
#endif
|
||||
|
||||
template <class L> __device__ __forceinline__ void residual_vector_prep(const uint8_t* w, int idx)
|
||||
{
|
||||
#if defined(PWIN_PACKED_RESIDUAL) && D4R_WMMA_LAYOUT == 11
|
||||
static_assert(L::KT <= 10, "residual table capacity");
|
||||
if (idx >= 0 && idx < 4 * L::KT)
|
||||
{
|
||||
const int kind = idx / (2 * L::KT), nt = (idx >> 1) % L::KT, h = idx & 1;
|
||||
const int base = (kind ? L::B2 : L::G2) + 32 * nt;
|
||||
u4v v;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
v[j] = pack2(wload(w, base + 2 * (4 * j + h)), wload(w, base + 2 * (4 * j + 2 + h)));
|
||||
g_residual_vectors[idx] = v;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template <class L> __device__ __forceinline__ void residual_vector(const uint8_t* w, int kind, int nt, half_t out[8])
|
||||
{
|
||||
#if defined(PWIN_PACKED_RESIDUAL) && D4R_WMMA_LAYOUT == 11
|
||||
static_assert(L::KT <= 10, "residual table capacity");
|
||||
const u4v v = g_residual_vectors[2 * L::KT * kind + 2 * nt + (lane_id() >> 4)];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
out[2 * j] = lo16(v[j]);
|
||||
out[2 * j + 1] = hi16(v[j]);
|
||||
}
|
||||
#else
|
||||
dvec8(w, (kind ? L::B2 : L::G2) + 32 * nt, out);
|
||||
#endif
|
||||
}
|
||||
|
||||
// prep items of the f32 vector table (j = 0 .. 2 NV - 1): core vectors from the core weights wc, the derived
|
||||
// layouts' extra vectors from w0 (the layer's weights pointer). MERGE: the layer has a patch merge (its merge bias
|
||||
// exists; for decoders and the bottleneck that offset lies past the layer's weights, so it is never read).
|
||||
template <class L, bool MERGE> __device__ void pwin_vec_prep(const uint8_t* w0, const uint8_t* wc, u8v* img, int j)
|
||||
{
|
||||
const int v = j >> 1, h = j & 1;
|
||||
const uint8_t* src = nullptr;
|
||||
if (v < L::V_B1)
|
||||
src = wc + L::BO + 32 * (v - L::V_BO);
|
||||
else if (v < L::V_PM)
|
||||
src = wc + L::B1 + 32 * (v - L::V_B1);
|
||||
else if (v < L::V_PM + L::PMT)
|
||||
{
|
||||
if (!MERGE)
|
||||
return;
|
||||
src = wc + L::PMB + 32 * (v - L::V_PM);
|
||||
}
|
||||
else
|
||||
src = L::extra_vec(w0, wc, v);
|
||||
if (src)
|
||||
vec_slot(img, L::VEC0 + j, src, h);
|
||||
}
|
||||
|
||||
// accumulator init vector v (16 halves at w + byte as floats in accumulator row order): from the prepared f32 table
|
||||
// (one 32-byte load per lane), or with PWIN_VEC_FROM_WEIGHTS from the weights (load + 8 shifts + 8 conversions;
|
||||
// faster in some layers). The same values.
|
||||
template <class L> __device__ __forceinline__ f8v acc_vec(const u8v* img, const uint8_t* w, int byte, int v)
|
||||
{
|
||||
#ifdef PWIN_VEC_FROM_WEIGHTS
|
||||
half_t t[8];
|
||||
dvec8(w, byte, t);
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)t[i];
|
||||
(void)img; (void)v;
|
||||
return d;
|
||||
#else
|
||||
(void)w; (void)byte;
|
||||
return img_vec(img, L::VEC0 + 2 * v);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u4v* bias, int idx)
|
||||
{
|
||||
const uint8_t* w = p.w;
|
||||
residual_vector_prep<L>(w, idx);
|
||||
if (idx < L::T_END * 16)
|
||||
{
|
||||
img[idx] = img_tile(w, L::src_of(idx >> 4), idx & 15);
|
||||
@@ -134,9 +231,20 @@ template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u
|
||||
bias[b] = (u4v){pack2(v[0], v[1]), pack2(v[2], v[3]), pack2(v[4], v[5]), pack2(v[6], v[7])};
|
||||
}
|
||||
|
||||
template <class L> __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias)
|
||||
template <class L, bool MERGE = true> __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias)
|
||||
{
|
||||
pwin_prep_at<L>(p, img, bias, blockIdx.x * 128 + threadIdx.x);
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x, j = idx - (L::T_END * 16 + L::NBIAS);
|
||||
if (j >= 0)
|
||||
{
|
||||
if (j < 2 * L::NV)
|
||||
pwin_vec_prep<L, MERGE>(p.w, p.w, img, j);
|
||||
return;
|
||||
}
|
||||
pwin_prep_at<L>(p, img, bias, idx);
|
||||
}
|
||||
template <class L> __device__ void pwin_prep_nomerge(const PwinParams& p, u8v* img, u4v* bias)
|
||||
{
|
||||
pwin_prep<L, false>(p, img, bias);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t hmin(half_t a, half_t b) { return a < b ? a : b; }
|
||||
@@ -150,7 +258,11 @@ __device__ __forceinline__ half_t softmax_e(half_t s)
|
||||
const half_t u = t * t;
|
||||
const half_t v = __builtin_fmaf16(u, a, (half_t)1.0f);
|
||||
const half_t x = __builtin_fmaf16(t, v, c);
|
||||
#ifdef PWIN_NATIVE_HEXP
|
||||
return __builtin_elementwise_exp2(x);
|
||||
#else
|
||||
return (half_t)__builtin_amdgcn_exp2f((float)x);
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t gelu(half_t x)
|
||||
@@ -163,7 +275,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_math(const f8v& 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
|
||||
@@ -178,6 +290,121 @@ __device__ __forceinline__ void gelu8(const f8v& d, half_t g[8])
|
||||
}
|
||||
}
|
||||
|
||||
// Exact GELU table experiment; enabled only by the POS wrappers that initialize it in prep.
|
||||
#ifdef PWIN_GELU_LUT
|
||||
#include "pwin_gelu_lut.h"
|
||||
#else
|
||||
constexpr int pwin_gelu_prep_items(int n) { return n; }
|
||||
__device__ __forceinline__ void pwin_gelu_prep() {}
|
||||
#endif
|
||||
|
||||
__device__ __forceinline__ void gelu8(const f8v& d, half_t g[8])
|
||||
{
|
||||
#ifdef PWIN_GELU_LUT
|
||||
gelu8_lut(d, g);
|
||||
#else
|
||||
gelu8_math(d, g);
|
||||
#endif
|
||||
}
|
||||
|
||||
// gelu of the 8 accumulator values as the next GEMM's operand. With WM_CVT_ASM the pairs stay packed from the
|
||||
// conversion through the packed f16 math to the operand exchange (the same f16 operations as gelu8_math).
|
||||
__device__ __forceinline__ op_t gelu_op(const f8v& d)
|
||||
{
|
||||
#if defined(PWIN_GELU_PACKED) && defined(WM_HAVE_FROM_MINE) && !defined(PWIN_GELU_LUT)
|
||||
const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
|
||||
uint32_t mine[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
#ifdef PWIN_GELU_PACKED_C
|
||||
const hv2 x = (hv2){(half_t)d[2 * j], (half_t)d[2 * j + 1]}; // C conversions (no asm in the MLP)
|
||||
#else
|
||||
const hv2 x = __builtin_bit_cast(hv2, wm_cvt2(d[2 * j], d[2 * j + 1]));
|
||||
#endif
|
||||
const hv2 c = __builtin_elementwise_min(__builtin_elementwise_max(x, (hv2)(half_t)-2.0f), (hv2)(half_t)2.0f);
|
||||
const hv2 u = k1 - k2 * __builtin_elementwise_abs(c);
|
||||
mine[j] = __builtin_bit_cast(uint32_t, x * ((hv2)(half_t)0.5f + c * u));
|
||||
}
|
||||
return wm_op_from_mine(mine);
|
||||
#else
|
||||
half_t g[8];
|
||||
gelu8(d, g);
|
||||
return operand_from_dt(g);
|
||||
#endif
|
||||
}
|
||||
|
||||
// ---- accumulator abstraction for the POS layers (pwin_pos.h).
|
||||
// Exact builds: pacc = f8v and every helper is the operation the code used before (bit-identical).
|
||||
// PWIN_F16ACC_NATIVE (opt-in approximate variant): accumulators stay in f16 registers through
|
||||
// the native f16-output WMMA (value i in the low half of dword i); this rounds every k16 step to f16 instead of
|
||||
// keeping f32 through the chain (harness PSNR ~62 dB vs the exact output), and removes most f32 -> f16 conversions.
|
||||
#ifdef PWIN_F16ACC_NATIVE
|
||||
static_assert(D4R_WMMA_LAYOUT == 11, "f16 accumulator variant: gfx11 layout");
|
||||
typedef h16 pacc;
|
||||
__device__ __forceinline__ pacc acc_zero() { return (pacc){}; }
|
||||
template <class L> __device__ __forceinline__ pacc acc_vec_p(const u8v* img, const uint8_t* w, int byte, int v)
|
||||
{
|
||||
const f8v f = acc_vec<L>(img, w, byte, v);
|
||||
pacc r = {};
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
r[2 * i] = (half_t)f[i];
|
||||
return r;
|
||||
}
|
||||
__device__ __forceinline__ pacc mma16p(const op_t& a, const op_t& b, pacc c)
|
||||
{
|
||||
return __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c, false);
|
||||
}
|
||||
// elements 2 j, 2 j + 1 as one packed register (low halves of dwords 2 j and 2 j + 1: one byte permutation)
|
||||
__device__ __forceinline__ uint32_t acc_pair(const pacc& d, int j)
|
||||
{
|
||||
const u8v u = __builtin_bit_cast(u8v, d);
|
||||
return __builtin_amdgcn_perm(u[2 * j + 1], u[2 * j], 0x05040100u);
|
||||
}
|
||||
__device__ __forceinline__ void acc_set_pair(pacc& d, int j, hv2 y)
|
||||
{
|
||||
d[4 * j] = y[0];
|
||||
d[4 * j + 2] = y[1];
|
||||
}
|
||||
__device__ __forceinline__ op_t operand_from_acc_p(const pacc& d)
|
||||
{
|
||||
uint32_t mine[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
mine[j] = acc_pair(d, j);
|
||||
return wm_op_from_mine(mine);
|
||||
}
|
||||
__device__ __forceinline__ op_t gelu_op_p(const pacc& d)
|
||||
{
|
||||
const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
|
||||
uint32_t mine[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
const hv2 x = __builtin_bit_cast(hv2, acc_pair(d, j));
|
||||
const hv2 c = __builtin_elementwise_min(__builtin_elementwise_max(x, (hv2)(half_t)-2.0f), (hv2)(half_t)2.0f);
|
||||
const hv2 u = k1 - k2 * __builtin_elementwise_abs(c);
|
||||
mine[j] = __builtin_bit_cast(uint32_t, x * ((hv2)(half_t)0.5f + c * u));
|
||||
}
|
||||
return wm_op_from_mine(mine);
|
||||
}
|
||||
#else
|
||||
typedef f8v pacc;
|
||||
__device__ __forceinline__ pacc acc_zero() { return splat8(0.0f); }
|
||||
template <class L> __device__ __forceinline__ pacc acc_vec_p(const u8v* img, const uint8_t* w, int byte, int v) { return acc_vec<L>(img, w, byte, v); }
|
||||
__device__ __forceinline__ pacc mma16p(const op_t& a, const op_t& b, pacc c) { return mma16(a, b, c); }
|
||||
__device__ __forceinline__ uint32_t acc_pair(const pacc& d, int j) { return wm_cvt2(d[2 * j], d[2 * j + 1]); }
|
||||
__device__ __forceinline__ void acc_set_pair(pacc& d, int j, hv2 y)
|
||||
{
|
||||
d[2 * j] = (float)y[0];
|
||||
d[2 * j + 1] = (float)y[1];
|
||||
}
|
||||
__device__ __forceinline__ op_t operand_from_acc_p(const pacc& d) { return operand_from_f8(d); }
|
||||
__device__ __forceinline__ op_t gelu_op_p(const pacc& d) { return gelu_op(d); }
|
||||
#endif
|
||||
|
||||
|
||||
// store the D^T tile of token m as 16 contiguous channels (lanes of one half-wave write)
|
||||
__device__ __forceinline__ void store_row16(half_t* dst, const u8v& v)
|
||||
{
|
||||
@@ -198,6 +425,8 @@ __device__ uint64_t g_stamp[16384][16];
|
||||
g_stamp[bid][k] = __builtin_readcyclecounter(); \
|
||||
} \
|
||||
} while (0)
|
||||
#elif defined(PWIN_PROF)
|
||||
#define STAMP(k) PROF(k)
|
||||
#else
|
||||
#define STAMP(k) do {} while (0)
|
||||
#endif
|
||||
@@ -246,17 +475,59 @@ __device__ __forceinline__ float row_sum_total(float rs, const half_t (&e)[4][8]
|
||||
|
||||
// The Swin core for one window: x0 rows of this wave's 16 tokens are in act (f16, C per token).
|
||||
// On return act holds the block output y. hb: normalised input rows; K/V of one head at a time in kl/vt.
|
||||
template <class L>
|
||||
#ifdef PWIN_VT_DIRECT
|
||||
// V^T operands in the vt buffer: slot r (0..127) at byte 48 r (32 B operand + 16 B pad: 16 lanes reading 16
|
||||
// consecutive slots with b128 loads hit disjoint LDS banks; an unpadded 32 B pitch conflicts 2-way)
|
||||
#define PWIN_VT_ROWS 43 /* 43 x LDSR(64) halves = 6192 B >= 128 x 48 B */
|
||||
static_assert(PWIN_VT_ROWS * LDSR(64) * 2 >= 128 * 48, "vt operand slots");
|
||||
static_assert(sizeof(op_t) == 32, "gfx11 operand layout");
|
||||
__device__ __forceinline__ void vt_op_store(half_t* vt, int r, const op_t& v)
|
||||
{
|
||||
if (!wm_half())
|
||||
{
|
||||
u4v* q = (u4v*)((uint8_t*)vt + 48 * r);
|
||||
q[0] = (u4v){v[0], v[1], v[2], v[3]};
|
||||
q[1] = (u4v){v[4], v[5], v[6], v[7]};
|
||||
}
|
||||
}
|
||||
__device__ __forceinline__ op_t vt_op_load(const half_t* vt, int r)
|
||||
{
|
||||
const u4v* q = (const u4v*)((const uint8_t*)vt + 48 * r);
|
||||
const u4v a = q[0], b = q[1];
|
||||
return (op_t){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
|
||||
}
|
||||
#else
|
||||
#define PWIN_VT_ROWS 32
|
||||
#endif
|
||||
|
||||
struct PwinNoX0
|
||||
{
|
||||
static constexpr bool active = false;
|
||||
__device__ void operator()() const {}
|
||||
};
|
||||
template <class F, class B> struct PwinX0
|
||||
{
|
||||
static constexpr bool active = true;
|
||||
F& f;
|
||||
B b;
|
||||
__device__ void operator()() const { f(b); }
|
||||
};
|
||||
|
||||
// X0F (PWIN_DEC_RECOMPUTE): act and hb are one buffer, so the residual's x0 is recomputed after attention by x0f()
|
||||
// into x0b (rows of this layer's tokens, aliasing kl / vt, which are dead by then)
|
||||
template <class L, class X0F = PwinNoX0>
|
||||
__device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
|
||||
half_t (*act)[L::C], half_t (*hb)[L::C], half_t (*kl)[32], half_t (*vt)[64])
|
||||
half_t (*act)[LDSR(L::C)], half_t (*hb)[LDSR(L::C)], half_t (*kl)[LDSR(32)], half_t (*vt)[LDSR(64)],
|
||||
const X0F& x0f = X0F(), half_t (*x0b)[LDSR(L::C)] = nullptr)
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z);
|
||||
const uint8_t* w = p.w;
|
||||
half_t* arow = act[16 * wv + m];
|
||||
half_t* hrow = hb[16 * wv + m];
|
||||
|
||||
PROF_INIT;
|
||||
STAMP(1);
|
||||
// ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb
|
||||
float ss = 0.0f;
|
||||
@@ -264,13 +535,23 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
#pragma unroll
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_SS)
|
||||
for (int j = 0; j < 8; ++j)
|
||||
{
|
||||
const uint32_t xj = xv[j]; // (no __builtin_bit_cast of a vector subscript: clang bug, REWRITE.md)
|
||||
const hv2 v = __builtin_bit_cast(hv2, xj), q = v * v;
|
||||
ss += (float)q[0];
|
||||
ss += (float)q[1];
|
||||
}
|
||||
#else
|
||||
for (int c = 0; c < 16; ++c)
|
||||
{
|
||||
const half_t v = op_get(xv, c);
|
||||
ss += (float)(half_t)(v * v);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
|
||||
const half_t r16 = half_rsqrt((half_t)ss);
|
||||
for (int kt = (int)hf; kt < KT; kt += 2)
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
@@ -278,8 +559,15 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt); // 16 contiguous gammas
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_APPLY)
|
||||
{
|
||||
const uint32_t xj = xv[j], gj = gv[j];
|
||||
hv[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, xj) * ((hv2)r16 * __builtin_bit_cast(hv2, gj)));
|
||||
}
|
||||
#else
|
||||
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
|
||||
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
|
||||
#endif
|
||||
store_row16(hrow + 16 * kt, hv);
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
@@ -296,19 +584,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
f8v acc[KT];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(w, L::BO + 32 * nt, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
acc[nt][i] = (float)vv_[i];
|
||||
|
||||
}
|
||||
{ acc[nt] = acc_vec<L>(img, w, L::BO + 32 * nt, L::V_BO + nt); }
|
||||
for (int h = 0; h < H; ++h)
|
||||
{
|
||||
// position-only attention: V^T A operands straight from D = h1 Wv (lane = dim, 16 keys per tile)
|
||||
@@ -322,7 +598,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
{
|
||||
f8v 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);
|
||||
d = mma16(op_lds(&hb[16 * wv + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), d);
|
||||
op_img_store(&vsh[h][nt][0][0], wv * 16 + m, operand_from_f8(d));
|
||||
}
|
||||
block_sync();
|
||||
@@ -337,11 +613,59 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
{
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
pop[kt] = op_img((const u8v*)bias, ((h * 4 + wv) * 4 + kt) * 16 + m);
|
||||
pop[kt] = op_imgt((const u8v*)bias, (h * 4 + wv) * 4 + kt);
|
||||
}
|
||||
else
|
||||
{
|
||||
op_t qop[2];
|
||||
#ifdef PWIN_BIAS_PREFETCH
|
||||
// this head's score bias tiles, loaded before the QKV chains (used right after the next barrier)
|
||||
u4v bpre[4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
bpre[kt] = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
|
||||
#endif
|
||||
#ifdef PWIN_QKV_SPLIT
|
||||
// the six chains (V, K, Q x 2 n-tiles; each still over kt in order) two at a time: 2 live accumulators
|
||||
#pragma unroll
|
||||
for (int j = 2; j >= 0; --j)
|
||||
{
|
||||
f8v dj[2] = {splat8(0.0f), splat8(0.0f)};
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const op_t row = op_lds(hrow + 16 * kt);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
const op_t wt = op_imgt(img, L::T_QKV + ((h * 3 + j) * KT + kt) * 2 + nt);
|
||||
#ifdef PWIN_VT_DIRECT
|
||||
dj[nt] = j == 2 ? mma16(row, wt, dj[nt]) : mma16(wt, row, dj[nt]);
|
||||
#else
|
||||
dj[nt] = mma16(wt, row, dj[nt]);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
const f8v d = dj[nt];
|
||||
if (j == 0)
|
||||
qop[nt] = operand_from_f8(d);
|
||||
else if (j == 1)
|
||||
op_store(&kl[16 * wv + m][16 * nt], operand_from_f8(d));
|
||||
else
|
||||
{
|
||||
#ifdef PWIN_VT_DIRECT
|
||||
vt_op_store(&vt[0][0], 64 * nt + wv * 16 + m, operand_from_f8(d));
|
||||
#else
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
vt[16 * nt + wm_acc_row(i)][16 * wv + m] = (half_t)d[i];
|
||||
#endif
|
||||
}
|
||||
}
|
||||
}
|
||||
#else
|
||||
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
|
||||
f8v dq[6];
|
||||
#pragma unroll
|
||||
@@ -350,9 +674,19 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const op_t row = op_lds(hrow + 16 * kt);
|
||||
#ifdef PWIN_VT_DIRECT
|
||||
// V chains with the operands swapped (as in the POS layers): D = V^T of this wave's 16 keys, lane = dim
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 4; ++u)
|
||||
dq[u] = mma16(op_imgt(img, L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)), row, dq[u]);
|
||||
#pragma unroll
|
||||
for (int u = 4; u < 6; ++u)
|
||||
dq[u] = mma16(row, op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + kt) * 2 + (u & 1)), dq[u]);
|
||||
#else
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 6; ++u)
|
||||
dq[u] = mma16(op_img(img, (L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m), row, dq[u]);
|
||||
dq[u] = mma16(op_imgt(img, L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)), row, dq[u]);
|
||||
#endif
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 3; ++j)
|
||||
@@ -368,18 +702,28 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
}
|
||||
else
|
||||
{
|
||||
#ifdef PWIN_VT_DIRECT
|
||||
// V^T operand of key tile wv -> slot nt * 64 + wv * 16 + m of the vt buffer (128 x 48 B)
|
||||
vt_op_store(&vt[0][0], 64 * nt + wv * 16 + m, operand_from_f8(d));
|
||||
#else
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
vt[16 * nt + wm_acc_row(i)][16 * wv + m] = (half_t)d[i];
|
||||
#endif
|
||||
}
|
||||
}
|
||||
#endif
|
||||
block_sync();
|
||||
half_t e[4][8];
|
||||
float rs = 0.0f;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
#ifdef PWIN_BIAS_PREFETCH
|
||||
const u4v bi = bpre[kt];
|
||||
#else
|
||||
const u4v bi = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
|
||||
#endif
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
@@ -399,7 +743,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
}
|
||||
}
|
||||
rs = row_sum_total(rs, e);
|
||||
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
|
||||
const half_t rc = half_reciprocal((half_t)rs);
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
@@ -421,7 +765,13 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
if constexpr (L::POS)
|
||||
d = mma16(vtop[nt][kt], pop[kt], d);
|
||||
else
|
||||
{
|
||||
#ifdef PWIN_VT_DIRECT
|
||||
d = mma16(vt_op_load(&vt[0][0], 64 * nt + kt * 16 + m), pop[kt], d);
|
||||
#else
|
||||
d = mma16(op_lds(&vt[16 * nt + m][16 * kt]), pop[kt], d);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
oop[nt] = operand_from_f8(d);
|
||||
#ifdef PWIN_DBG
|
||||
@@ -434,7 +784,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
acc[nt] = mma16(op_img(img, (L::T_WO + (2 * h + ks) * KT + nt) * 16 + m), oop[ks], acc[nt]);
|
||||
acc[nt] = mma16(op_imgt(img, L::T_WO + (2 * h + ks) * KT + nt), oop[ks], acc[nt]);
|
||||
if constexpr (!L::POS)
|
||||
block_sync(); // kl / vt are reused by the next head
|
||||
}
|
||||
@@ -447,13 +797,33 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
#endif
|
||||
STAMP(3);
|
||||
// ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
|
||||
if constexpr (X0F::active)
|
||||
{
|
||||
x0f(); // every wave's attention is done (block barrier at the end of the last head)
|
||||
block_sync(); // x0 rows of all waves' tokens written
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
const op_t xv = op_lds(arow + 16 * nt);
|
||||
const op_t xv = X0F::active ? op_lds(&x0b[16 * wv + m][16 * nt]) : op_lds(arow + 16 * nt);
|
||||
half_t mv[8], g2[8], b2[8];
|
||||
dvec8(w, L::G2 + 32 * nt, g2);
|
||||
dvec8(w, L::B2 + 32 * nt, b2);
|
||||
residual_vector<L>(w, 0, nt, g2);
|
||||
residual_vector<L>(w, 1, nt, b2);
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_RESID)
|
||||
uint32_t mine[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
const hv2 x1 = __builtin_bit_cast(hv2, wm_cvt2(acc[nt][2 * j], acc[nt][2 * j + 1])) +
|
||||
__builtin_bit_cast(hv2, wm_op_acc_pair(xv, j));
|
||||
mine[j] = __builtin_bit_cast(uint32_t, x1 * (hv2){g2[2 * j], g2[2 * j + 1]});
|
||||
const hv2 y = x1 + (hv2){b2[2 * j], b2[2 * j + 1]};
|
||||
acc[nt][2 * j] = (float)y[0];
|
||||
acc[nt][2 * j + 1] = (float)y[1];
|
||||
}
|
||||
(void)mv;
|
||||
op_store(hrow + 16 * nt, wm_op_from_mine(mine));
|
||||
#else
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
@@ -462,6 +832,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
acc[nt][i] = (float)(half_t)(x1 + b2[i]);
|
||||
}
|
||||
op_store(hrow + 16 * nt, operand_from_dt(mv));
|
||||
#endif
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
|
||||
@@ -473,30 +844,16 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
for (int hn = 0; hn < 2; ++hn)
|
||||
{
|
||||
f8v d;
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(w, L::B1 + 64 * c + 32 * hn, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
d[i] = (float)vv_[i];
|
||||
|
||||
}
|
||||
{ d = acc_vec<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn); }
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
|
||||
half_t g[8];
|
||||
gelu8(d, g);
|
||||
gop[hn] = operand_from_dt(g);
|
||||
d = mma16(op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn), op_lds(hrow + 16 * kt), d);
|
||||
gop[hn] = gelu_op(d);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
acc[nt] = mma16(op_img(img, (L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m), gop[ks], acc[nt]);
|
||||
acc[nt] = mma16(op_imgt(img, L::T_W2 + (c * 2 + ks) * KT + nt), gop[ks], acc[nt]);
|
||||
}
|
||||
STAMP(4);
|
||||
// y -> act (rows of this wave)
|
||||
@@ -506,6 +863,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
op_store(arow + 16 * nt, operand_from_f8(acc[nt]));
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
PROF(5);
|
||||
}
|
||||
|
||||
// first layer: token embedding x0 = relu(in (CIN) W_e + b_e) in front of the core (core weights at CB)
|
||||
@@ -514,12 +872,22 @@ template <int H_, int C_, int COUT_, bool POS_, int CIN_> struct PwinEmbLayout :
|
||||
using B = PwinLayout<H_, C_, COUT_, POS_>;
|
||||
static constexpr int CIN = CIN_, KE = CIN / 16, CB = CIN ? 2 * C_ + 2 * CIN * C_ : 0;
|
||||
static constexpr int T_EMB = B::T_END, T_END = T_EMB + KE * (C_ / 16);
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
|
||||
// + the embedding bias (C / 16 vectors at the layer's weights pointer)
|
||||
static constexpr int V_EMB = B::NV, NV = V_EMB + (CIN ? C_ / 16 : 0);
|
||||
static constexpr int VEC0 = T_END * 16, IMG_SLOTS = VEC0 + 2 * NV;
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS + 2 * NV;
|
||||
__device__ static const uint8_t* extra_vec(const uint8_t* w0, const uint8_t*, int v) { return w0 + 32 * (v - V_EMB); }
|
||||
};
|
||||
|
||||
template <class L> __device__ void pwin_emb_prep(const PwinParams& p, u8v* img, u4v* bias)
|
||||
{
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x;
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x, j = idx - (L::T_END * 16 + L::B::NBIAS);
|
||||
if (j >= 0)
|
||||
{
|
||||
if (j < 2 * L::NV)
|
||||
pwin_vec_prep<L, true>(p.w, p.w + L::CB, img, j);
|
||||
return;
|
||||
}
|
||||
if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
|
||||
{
|
||||
const int t = (idx >> 4) - L::T_EMB, nt = t % (L::C / 16), ks = t / (L::C / 16);
|
||||
@@ -543,14 +911,15 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
|
||||
if constexpr (L::CB != 0)
|
||||
p.w = p0.w + L::CB;
|
||||
constexpr int C = L::C, COUT = L::COUT;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[64][32];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[32][64];
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[64][LDSR(32)];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[PWIN_VT_ROWS][LDSR(64)];
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
|
||||
const int bx = blockIdx.x, by = blockIdx.y;
|
||||
|
||||
PROF_INIT;
|
||||
STAMP(0);
|
||||
// input rows (mirror padded) -> act (first layer: embedding + ReLU)
|
||||
{
|
||||
@@ -566,16 +935,10 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
{
|
||||
f8v d;
|
||||
{
|
||||
half_t vv_[8];
|
||||
dvec8(p0.w, 32 * nt, vv_);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv_[i];
|
||||
}
|
||||
{ d = acc_vec<L>(img, p0.w, 32 * nt, L::V_EMB + nt); }
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EMB + ks * (C / 16) + nt) * 16 + m), xe[ks], d);
|
||||
d = mma16(op_imgt(img, L::T_EMB + ks * (C / 16) + nt), xe[ks], d);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
@@ -614,25 +977,14 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
|
||||
for (int vt = wv; vt < L::PMT; vt += 4)
|
||||
{
|
||||
f8v d;
|
||||
{
|
||||
{ d = acc_vec<L>(img, w, L::PMB + 32 * vt, L::V_PM + vt);
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(w, L::PMB + 32 * vt, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
d[i] = (float)vv_[i];
|
||||
|
||||
STAMP(7);
|
||||
}
|
||||
STAMP(7); }
|
||||
for (int ks = 0; ks < C / 4; ++ks)
|
||||
{
|
||||
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
|
||||
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
|
||||
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
|
||||
d = mma16(op_imgt(img, L::T_PM + ks * L::PMT + vt), op_lds(&act[src][16 * (ks % L::KT)]), d);
|
||||
}
|
||||
// allocated column 16 vt + c belongs to original warp g, local column 16 u + c
|
||||
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
|
||||
@@ -642,7 +994,7 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
|
||||
|
||||
#define PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, CIN) \
|
||||
using NAME##_L = PwinEmbLayout<H, C, COUT, POS, CIN>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
@@ -663,12 +1015,12 @@ template <class L>
|
||||
__device__ void pwin_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[64][32];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[32][64];
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[64][LDSR(32)];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[PWIN_VT_ROWS][LDSR(64)];
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
|
||||
const int bx = blockIdx.x, by = blockIdx.y;
|
||||
{
|
||||
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
|
||||
@@ -689,7 +1041,7 @@ __device__ void pwin_plain(const PwinParams& p, const u8v* __restrict__ img, con
|
||||
|
||||
#define PWIN_KERNEL(NAME, H, C, COUT, BODY) \
|
||||
using NAME##_L = PwinLayout<H, C, COUT>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
@@ -715,12 +1067,25 @@ struct PwinDecLayout : PwinLayout<H_, C_, 16, POS_>
|
||||
static constexpr int T_EXP = B::T_PM; // expand tiles replace the merge tiles
|
||||
static constexpr int T_HEAD = T_EXP + KL * EXPN;
|
||||
static constexpr int T_END = T_HEAD + B::KT * (NOUTA_ / 16);
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
|
||||
// + the expand biases (EXPN vectors at w0 + EXPB) and the output head bias (NOUTA / 16 at core + HB)
|
||||
static constexpr int V_EXP = B::NV, V_HB = V_EXP + EXPN, NV = V_HB + NOUTA_ / 16;
|
||||
static constexpr int VEC0 = T_END * 16, IMG_SLOTS = VEC0 + 2 * NV;
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS + 2 * NV;
|
||||
__device__ static const uint8_t* extra_vec(const uint8_t* w0, const uint8_t* wc, int v)
|
||||
{
|
||||
return v < V_HB ? w0 + EXPB + 32 * (v - V_EXP) : wc + HB + 32 * (v - V_HB);
|
||||
}
|
||||
};
|
||||
|
||||
template <class L> __device__ void pwin_dec_prep(const PwinParams& p, u8v* img, u4v* bias)
|
||||
{
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x;
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x, j = idx - (L::T_END * 16 + L::B::NBIAS);
|
||||
if (j >= 0)
|
||||
{
|
||||
if (j < 2 * L::NV)
|
||||
pwin_vec_prep<L, false>(p.w, p.w + L::CB, img, j);
|
||||
return;
|
||||
}
|
||||
if (idx >= L::T_EXP * 16 && idx < L::T_HEAD * 16)
|
||||
{
|
||||
const int t = (idx >> 4) - L::T_EXP, nt = t % L::EXPN, ks = t / L::EXPN;
|
||||
@@ -743,16 +1108,35 @@ template <class L>
|
||||
__device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
constexpr int KVR = L::POS ? 1 : 64, KVC = L::POS ? 1 : 32;
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[KVR][32];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[KVC][64];
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
|
||||
constexpr int KVR = L::POS ? 1 : 64, KVC = L::POS ? 1 : PWIN_VT_ROWS;
|
||||
#ifdef PWIN_DEC_RECOMPUTE
|
||||
// one buffer for act / hb; kl, vt and (after attention) the recomputed x0 rows share one region
|
||||
static_assert(!L::POS, "PWIN_DEC_RECOMPUTE: full-attention decoders");
|
||||
constexpr int KLB = KVR * LDSR(32) * 2, VTB = KVC * LDSR(64) * 2, X0B = 64 * LDSR(C) * 2;
|
||||
constexpr int KVXB = (KLB + VTB) > X0B ? (KLB + VTB) : X0B;
|
||||
__shared__ __attribute__((aligned(16))) uint8_t kvx[KVXB];
|
||||
half_t (*hb)[LDSR(C)] = act;
|
||||
half_t (*kl)[LDSR(32)] = (half_t (*)[LDSR(32)])kvx;
|
||||
half_t (*vt)[LDSR(64)] = (half_t (*)[LDSR(64)])(kvx + KLB);
|
||||
half_t (*x0b)[LDSR(C)] = (half_t (*)[LDSR(C)])kvx;
|
||||
#else
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[KVR][LDSR(32)];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[KVC][LDSR(64)];
|
||||
#endif
|
||||
#ifdef PWIN_LDS_PAD
|
||||
// occupancy-sensitivity probe only: PWIN_LDS_PAD bytes of unused LDS per block (one dword written)
|
||||
__shared__ uint32_t lds_pad[PWIN_LDS_PAD / 4];
|
||||
if (threadIdx.x == 0 && threadIdx.z == 0)
|
||||
((volatile uint32_t*)lds_pad)[0] = 0u;
|
||||
#endif
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), bx = blockIdx.x, by = blockIdx.y;
|
||||
PROF_INIT;
|
||||
|
||||
// ---- patch expand: wave wv produces sub-token q = wv (dy, dx) of all 16 low-resolution tokens
|
||||
{
|
||||
// ---- patch expand: wave wv produces sub-token q = wv (dy, dx) of all 16 low-resolution tokens -> dst rows
|
||||
auto expand = [&](half_t (*dst)[LDSR(C)]) {
|
||||
const int W2 = p.W / 2, H2 = p.H / 2;
|
||||
const int ly = m >> 2, lx = m & 3;
|
||||
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
|
||||
@@ -764,33 +1148,37 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
|
||||
{
|
||||
const int gnt = wv * (C / 16) + nt; // expand output column tile
|
||||
f8v d;
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(p.w, L::EXPB + 32 * gnt, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
d[i] = (float)vv_[i];
|
||||
|
||||
}
|
||||
{ d = acc_vec<L>(img, p.w, L::EXPB + 32 * gnt, L::V_EXP + gnt); }
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
|
||||
d = mma16(op_imgt(img, L::T_EXP + ks * L::EXPN + gnt), op_gload(xl + 16 * ks), d);
|
||||
const op_t sv = op_gload(sk + 16 * nt);
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_EXPAND)
|
||||
uint32_t mine[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
mine[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, wm_cvt2(d[2 * j], d[2 * j + 1])) +
|
||||
__builtin_bit_cast(hv2, wm_op_acc_pair(sv, j)));
|
||||
op_store(&dst[t][16 * nt], wm_op_from_mine(mine));
|
||||
#else
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
|
||||
op_store(&act[t][16 * nt], operand_from_dt(xo));
|
||||
op_store(&dst[t][16 * nt], operand_from_dt(xo));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
};
|
||||
expand(act);
|
||||
block_sync();
|
||||
PROF(8);
|
||||
PwinParams q = p;
|
||||
q.w = p.w + L::CB;
|
||||
#ifdef PWIN_DEC_RECOMPUTE
|
||||
pwin_core<L>(q, img, bias, act, hb, kl, vt, PwinX0<decltype(expand), half_t (*)[LDSR(C)]>{expand, x0b}, x0b);
|
||||
#else
|
||||
pwin_core<L>(q, img, bias, act, hb, kl, vt);
|
||||
#endif
|
||||
PROF(9);
|
||||
|
||||
const int tok = 16 * wv + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
|
||||
const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
|
||||
@@ -801,22 +1189,10 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
|
||||
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
|
||||
{
|
||||
f8v d;
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(q.w, L::HB + 32 * nt, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
d[i] = (float)vv_[i];
|
||||
|
||||
}
|
||||
{ d = acc_vec<L>(img, q.w, L::HB + 32 * nt, L::V_HB + nt); }
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < L::KT; ++kt)
|
||||
d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
|
||||
d = mma16(op_imgt(img, L::T_HEAD + kt * (L::NOUTA / 16) + nt), op_lds(&act[tok][16 * kt]), d);
|
||||
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
|
||||
}
|
||||
}
|
||||
@@ -826,11 +1202,12 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
|
||||
}
|
||||
PROF(10);
|
||||
}
|
||||
|
||||
#define PWIN_DECODER_EX(NAME, H, C, CL, POS, NOUT, NOUTA) \
|
||||
using NAME##_L = PwinDecLayout<H, C, CL, POS, NOUT, NOUTA>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
|
||||
@@ -0,0 +1,402 @@
|
||||
// Phase-split DltssPaddedWinLayer: the deep layers have few windows (dec5: 9), so one block
|
||||
// per window leaves most of the GPU idle. Here a layer runs as four kernels (the main kernel plus NAME_post1..3,
|
||||
// launched back to back by the patched ZLUDA hook, patches/zluda/0008-native-post-phase-kernels.patch), each a fixed-size grid
|
||||
// looping over independent units:
|
||||
// P0 (main) per (window, head): L2 norm + QKV + attention -> O (f16, global)
|
||||
// P1 (post1) per (window, out tile, tok tile): Wo chain over heads + residual -> m, a (f16, global)
|
||||
// P2 (post2) per (window, hidden tile, tok): fc1 chain + gelu -> hid (f16, global)
|
||||
// P3 (post3) per (window, out tile, tok tile): fc2 chain over hidden from a -> y (layer output)
|
||||
// Every value goes through the same f16 / f32 operations in the same order as pwin_wide.h (and pwin_layer.h), and
|
||||
// everything passed between phases is an f16 value already, so the output is bit-identical. Units are distributed
|
||||
// with compile-time grid sizes (no gridDim), the window count comes from the parameters (W, H, shift).
|
||||
// Scratch buffers are module globals sized for MAXW windows; a launch with more windows does nothing (guard
|
||||
// against out-of-bounds writes; the layer output is then not written).
|
||||
#pragma once
|
||||
#include "pwin_wide.h"
|
||||
|
||||
// layer kinds: core weights at p.w (+ CB for decoders); x0 rows from the layer input (mirror padded) or, for
|
||||
// decoders, from the expand phase's scratch rows gX[window][token][C]
|
||||
template <class L> struct PhaseCB
|
||||
{
|
||||
template <class T> static constexpr int cb(decltype(T::CB)*) { return T::CB; }
|
||||
template <class T> static constexpr int cb(...) { return 0; }
|
||||
static constexpr int value = cb<L>(nullptr);
|
||||
};
|
||||
template <class L> __device__ __forceinline__ const half_t* phase_x0_row(const PwinParams& p, const half_t* gX, int wi, int bx,
|
||||
int by, int tok)
|
||||
{
|
||||
if (gX)
|
||||
return gX + ((size_t)wi * 64 + tok) * L::C;
|
||||
const int X = mirror(8 * bx - p.sx + (tok & 7), p.W), Y = mirror(8 * by - p.sy + (tok >> 3), p.H);
|
||||
return (const half_t*)p.in + (size_t)(Y * p.W + X) * L::C;
|
||||
}
|
||||
|
||||
template <class L> struct PhaseWin
|
||||
{
|
||||
int GX, GY, n;
|
||||
__device__ PhaseWin(const PwinParams& p) : GX((p.W + p.sx + 7) / 8), GY((p.H + p.sy + 7) / 8), n(GX * GY) {}
|
||||
};
|
||||
|
||||
// P0: unit u = (window w, head h); 4 waves = the 4 token tiles of the window
|
||||
template <class L, int MAXW, int GRID>
|
||||
__device__ void phase_attn(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias, half_t* gO,
|
||||
const half_t* gX = nullptr, int ub = -1)
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H;
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[1][64][LDSR(32)];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[1][32][LDSR(64)];
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int t = __builtin_amdgcn_readfirstlane(threadIdx.z), tok = 16 * t + m;
|
||||
const PhaseWin<L> win(p);
|
||||
if (win.n > MAXW)
|
||||
return;
|
||||
const uint8_t* w = p.w + PhaseCB<L>::value;
|
||||
half_t* hrow = hb[tok];
|
||||
// ub >= 0 (single-unit launches): exactly unit ub; otherwise grid-strided over all units
|
||||
for (int u = ub >= 0 ? ub : (int)blockIdx.x; u < win.n * H; u += ub >= 0 ? 0x3fffffff : GRID)
|
||||
{
|
||||
const int wi = u / H, h = u % H, bx = wi % win.GX, by = wi / win.GX;
|
||||
const half_t* xin = phase_x0_row<L>(p, gX, wi, bx, by, tok);
|
||||
// L2 norm (wide_core): h1 = x0 * (rsqrt(sum x0^2) * gamma1)
|
||||
float ss = 0.0f;
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const u8v xv = gload_row16(xin + 16 * kt);
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 16; ++c)
|
||||
{
|
||||
const half_t v = op_get(xv, c);
|
||||
ss += (float)(half_t)(v * v);
|
||||
}
|
||||
}
|
||||
const half_t r16 = half_rsqrt((half_t)ss);
|
||||
for (int kt = (int)hf; kt < KT; kt += 2)
|
||||
{
|
||||
const u8v xv = gload_row16(xin + 16 * kt);
|
||||
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
|
||||
u8v hv;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
|
||||
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
|
||||
store_row16(hrow + 16 * kt, hv);
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
op_t qop[2], oop[2];
|
||||
wide_qkv<L, 0, 6>(img, hrow, h, 0, tok, qop, kl, vt);
|
||||
block_sync();
|
||||
{
|
||||
half_t e[4][8];
|
||||
float rs = 0.0f;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
const u4v bi = bias[((h * 4 + t) * 4 + kt) * 32 + l];
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
d[2 * i] = (float)lo16(bi[i]);
|
||||
d[2 * i + 1] = (float)hi16(bi[i]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int ds = 0; ds < 2; ++ds)
|
||||
d = mma16(op_lds(&kl[0][16 * kt + m][16 * ds]), qop[ds], d);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
e[kt][i] = softmax_e((half_t)d[i]);
|
||||
if constexpr (!kShimRowSum)
|
||||
rs += (float)e[kt][i];
|
||||
}
|
||||
}
|
||||
rs = row_sum_total(rs, e);
|
||||
const half_t rc = half_reciprocal((half_t)rs);
|
||||
op_t pop[4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
half_t pv[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
pv[i] = e[kt][i] * rc;
|
||||
pop[kt] = operand_from_dt(pv);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
f8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
d = mma16(op_lds(&vt[0][16 * nt + m][16 * kt]), pop[kt], d);
|
||||
oop[nt] = operand_from_f8(d);
|
||||
}
|
||||
}
|
||||
half_t* orow = gO + ((size_t)(wi * H + h) * 64 + tok) * 32;
|
||||
op_gstore(orow, oop[0], true);
|
||||
op_gstore(orow + 16, oop[1], true);
|
||||
block_sync(); // hb / kl / vt are rewritten by the next unit
|
||||
}
|
||||
}
|
||||
|
||||
// global wave index helpers for the per-wave-unit phases (blocks of 4 waves, GRID blocks)
|
||||
// (ub >= 0: only the 4 wave units of chunk ub, see phase_attn)
|
||||
#define PHASE_WAVE_LOOP(UNITS, GRID) \
|
||||
for (int uw = (ub >= 0 ? ub : (int)blockIdx.x) * 4 + (int)__builtin_amdgcn_readfirstlane(threadIdx.z); uw < (UNITS); \
|
||||
uw += ub >= 0 ? 0x3fffffff : (GRID) * 4)
|
||||
|
||||
// P1: unit = (window, output tile nt, token tile): Wo chain over heads, residual; writes the MLP input m and the
|
||||
// fc2 accumulator init a = f16(x1 + b2) (rows of C halves per token)
|
||||
template <class L, int MAXW, int GRID>
|
||||
__device__ void phase_wo(const PwinParams& p, const u8v* __restrict__ img, const half_t* gO, half_t* gM, half_t* gA,
|
||||
const half_t* gX = nullptr, int ub = -1)
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H;
|
||||
const uint32_t m = lane_id() & 15;
|
||||
const PhaseWin<L> win(p);
|
||||
if (win.n > MAXW)
|
||||
return;
|
||||
const uint8_t* w = p.w + PhaseCB<L>::value;
|
||||
PHASE_WAVE_LOOP(win.n * KT * 4, GRID)
|
||||
{
|
||||
const int wi = uw / (KT * 4), r = uw % (KT * 4), nt = r >> 2, tt = r & 3, tok = 16 * tt + m;
|
||||
const int bx = wi % win.GX, by = wi / win.GX;
|
||||
f8v acc = acc_vec<L>(img, w, L::BO + 32 * nt, L::V_BO + nt);
|
||||
for (int hh = 0; hh < H; ++hh)
|
||||
{
|
||||
const half_t* orow = gO + ((size_t)(wi * H + hh) * 64 + tok) * 32;
|
||||
acc = mma16(op_imgt(img, L::T_WO + (2 * hh) * KT + nt), op_gload(orow), acc);
|
||||
acc = mma16(op_imgt(img, L::T_WO + (2 * hh + 1) * KT + nt), op_gload(orow + 16), acc);
|
||||
}
|
||||
const op_t xv = op_gload(phase_x0_row<L>(p, gX, wi, bx, by, tok) + 16 * nt);
|
||||
half_t mv[8], av[8], g2[8], b2[8];
|
||||
residual_vector<L>(w, 0, nt, g2);
|
||||
residual_vector<L>(w, 1, nt, b2);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const half_t x1 = (half_t)acc[i] + wm_op_acc_elem(xv, i);
|
||||
mv[i] = x1 * g2[i];
|
||||
av[i] = (half_t)(x1 + b2[i]);
|
||||
}
|
||||
const size_t row = ((size_t)wi * 64 + tok) * C + 16 * nt;
|
||||
op_gstore(gM + row, operand_from_dt(mv), true);
|
||||
op_gstore(gA + row, operand_from_dt(av), true);
|
||||
}
|
||||
}
|
||||
|
||||
// P2: unit = (window, hidden tile j = 2 c + hn, token tile): fc1 chain over kt + gelu -> hidden rows (4C halves)
|
||||
template <class L, int MAXW, int GRID>
|
||||
__device__ void phase_fc1(const PwinParams& p, const u8v* __restrict__ img, const half_t* gM, half_t* gH, int ub = -1)
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT, NJ = 2 * L::NMLP;
|
||||
const uint32_t m = lane_id() & 15;
|
||||
const PhaseWin<L> win(p);
|
||||
if (win.n > MAXW)
|
||||
return;
|
||||
const uint8_t* w = p.w + PhaseCB<L>::value;
|
||||
PHASE_WAVE_LOOP(win.n * NJ * 4, GRID)
|
||||
{
|
||||
const int wi = uw / (NJ * 4), r = uw % (NJ * 4), j = r >> 2, tt = r & 3, tok = 16 * tt + m;
|
||||
const int c = j >> 1, hn = j & 1;
|
||||
f8v d = acc_vec<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn);
|
||||
const half_t* mrow = gM + ((size_t)wi * 64 + tok) * C;
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
d = mma16(op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn), op_gload(mrow + 16 * kt), d);
|
||||
op_gstore(gH + ((size_t)wi * 64 + tok) * (4 * C) + 16 * j, gelu_op(d), true);
|
||||
}
|
||||
}
|
||||
|
||||
// P3: unit = (window, output tile nt, token tile): fc2 chain over the hidden chunks from a; y -> layer output
|
||||
template <class L, int MAXW, int GRID>
|
||||
__device__ void phase_fc2(const PwinParams& p, const u8v* __restrict__ img, const half_t* gA, const half_t* gH, uint8_t* dst,
|
||||
half_t* gY = nullptr, int ub = -1)
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT;
|
||||
const uint32_t m = lane_id() & 15;
|
||||
const PhaseWin<L> win(p);
|
||||
if (win.n > MAXW)
|
||||
return;
|
||||
PHASE_WAVE_LOOP(win.n * KT * 4, GRID)
|
||||
{
|
||||
const int wi = uw / (KT * 4), r = uw % (KT * 4), nt = r >> 2, tt = r & 3, tok = 16 * tt + m;
|
||||
const int bx = wi % win.GX, by = wi / win.GX;
|
||||
const op_t av = op_gload(gA + ((size_t)wi * 64 + tok) * C + 16 * nt);
|
||||
f8v acc;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[i] = (float)wm_op_acc_elem(av, i);
|
||||
const half_t* hrow = gH + ((size_t)wi * 64 + tok) * (4 * C);
|
||||
for (int cc = 0; cc < L::NMLP; ++cc)
|
||||
{
|
||||
acc = mma16(op_imgt(img, L::T_W2 + (cc * 2) * KT + nt), op_gload(hrow + 32 * cc), acc);
|
||||
acc = mma16(op_imgt(img, L::T_W2 + (cc * 2 + 1) * KT + nt), op_gload(hrow + 32 * cc + 16), acc);
|
||||
}
|
||||
const int Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
|
||||
const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
|
||||
half_t* yout = (half_t*)dst + (size_t)(Y0 * p.W + X0) * C;
|
||||
const op_t y = operand_from_f8(acc);
|
||||
op_gstore(yout + 16 * nt, y, inb);
|
||||
if (gY)
|
||||
op_gstore(gY + ((size_t)wi * 64 + tok) * C + 16 * nt, y, true);
|
||||
}
|
||||
}
|
||||
|
||||
// decoder expand (wide_decoder): unit = (window, sub-token q, tile nt) per wave; lanes = the 16 low-resolution
|
||||
// tokens; x0 = patch expand of the previous decoder output + encoder skip -> gX rows (all 64 tokens of the window)
|
||||
template <class L, int MAXW, int GRID>
|
||||
__device__ void phase_expand(const PwinParams& p, const u8v* __restrict__ img, half_t* gX, int ub = -1)
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT;
|
||||
const uint32_t m = lane_id() & 15;
|
||||
const PhaseWin<L> win(p);
|
||||
if (win.n > MAXW)
|
||||
return;
|
||||
PHASE_WAVE_LOOP(win.n * 4 * KT, GRID)
|
||||
{
|
||||
const int wi = uw / (4 * KT), r = uw % (4 * KT), q = r / KT, nt = r % KT;
|
||||
const int bx = wi % win.GX, by = wi / win.GX;
|
||||
const int W2 = p.W / 2, H2 = p.H / 2;
|
||||
const int ly = m >> 2, lx = m & 3;
|
||||
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
|
||||
const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
|
||||
const int dy = q >> 1, dx = q & 1, tt = 8 * (2 * ly + dy) + 2 * lx + dx;
|
||||
const int SX = mirror(8 * bx - p.sx + (tt & 7), p.W), SY = mirror(8 * by - p.sy + (tt >> 3), p.H);
|
||||
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
|
||||
const int gnt = q * (C / 16) + nt;
|
||||
f8v d = acc_vec<L>(img, p.w, L::EXPB + 32 * gnt, L::V_EXP + gnt);
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
d = mma16(op_imgt(img, L::T_EXP + ks * L::EXPN + gnt), op_gload(xl + 16 * ks), d);
|
||||
const op_t sv = op_gload(sk + 16 * nt);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
|
||||
op_gstore(gX + ((size_t)wi * 64 + tt) * C + 16 * nt, operand_from_dt(xo), true);
|
||||
}
|
||||
}
|
||||
|
||||
// encoder patch merge (wide_encoder): unit = (window, merge tile vt) per wave; lanes = the 16 merged tokens
|
||||
template <class L, int MAXW, int GRID>
|
||||
__device__ void phase_merge(const PwinParams& p, const u8v* __restrict__ img, const half_t* gY, int ub = -1)
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT;
|
||||
const uint32_t m = lane_id() & 15;
|
||||
const PhaseWin<L> win(p);
|
||||
if (win.n > MAXW)
|
||||
return;
|
||||
const uint8_t* w = p.w;
|
||||
PHASE_WAVE_LOOP(win.n * L::PMT, GRID)
|
||||
{
|
||||
const int wi = uw / L::PMT, vt = uw % L::PMT;
|
||||
const int bx = wi % win.GX, by = wi / win.GX;
|
||||
const int my = m >> 2, mx = m & 3;
|
||||
const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
|
||||
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
|
||||
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * L::COUT;
|
||||
f8v d = acc_vec<L>(img, w, L::PMB + 32 * vt, L::V_PM + vt);
|
||||
for (int ks = 0; ks < C / 4; ++ks)
|
||||
{
|
||||
const int sub = ks / KT, dy = sub >> 1, dx = sub & 1;
|
||||
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
|
||||
d = mma16(op_imgt(img, L::T_PM + ks * L::PMT + vt), op_gload(gY + ((size_t)wi * 64 + src) * C + 16 * (ks % KT)), d);
|
||||
}
|
||||
const int gw = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
|
||||
op_gstore(mout + L::NPW * gw + 16 * u, operand_from_f8(d), minb, valid);
|
||||
}
|
||||
}
|
||||
|
||||
// bottleneck layer (dec5 shape: core only, output to +24) as four phase kernels
|
||||
#define PHASE_PLAIN(NAME, H, C, MAXW, G0, G1, G2, G3) \
|
||||
using NAME##_L = PwinLayout<H, C, 16>; \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
__device__ half_t g_O[(MAXW) * (H) * 64 * 32]; \
|
||||
__device__ half_t g_M[(MAXW) * 64 * (C)]; \
|
||||
__device__ half_t g_A[(MAXW) * 64 * (C)]; \
|
||||
__device__ half_t g_H[(MAXW) * 64 * 4 * (C)]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = (G0); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_post1_grid_x = (G1); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_post1_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_post2_grid_x = (G2); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_post2_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_post3_grid_x = (G3); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_post3_block_z = 4; \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_prep_nomerge<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) \
|
||||
{ \
|
||||
phase_attn<NAME##_L, MAXW, G0>(p, g_img, g_bias, g_O); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) \
|
||||
{ \
|
||||
phase_wo<NAME##_L, MAXW, G1>(p, g_img, g_O, g_M, g_A); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) \
|
||||
{ \
|
||||
phase_fc1<NAME##_L, MAXW, G2>(p, g_img, g_M, g_H); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) \
|
||||
{ \
|
||||
phase_fc2<NAME##_L, MAXW, G3>(p, g_img, g_A, g_H, p.out24); \
|
||||
}
|
||||
|
||||
#define PHASE_GLOBALS(NAME, H, C, MAXW) \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
__device__ half_t g_O[(MAXW) * (H) * 64 * 32]; \
|
||||
__device__ half_t g_M[(MAXW) * 64 * (C)]; \
|
||||
__device__ half_t g_A[(MAXW) * 64 * (C)]; \
|
||||
__device__ half_t g_H[(MAXW) * 64 * 4 * (C)]; \
|
||||
__device__ half_t g_X[(MAXW) * 64 * (C)]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64;
|
||||
#define PHASE_POST(K, G) \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_post##K##_grid_x = (G); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_post##K##_block_z = 4;
|
||||
|
||||
// decoder (wide shape, patch expand + skip, no output head): expand, attention, Wo, fc1, fc2
|
||||
#define PHASE_DECODER(NAME, H, C, CL, MAXW, GE, G0, G1, G2, G3) \
|
||||
using NAME##_L = PwinDecLayout<H, C, CL>; \
|
||||
PHASE_GLOBALS(NAME, H, C, MAXW) \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = (GE); \
|
||||
PHASE_POST(1, G0) PHASE_POST(2, G1) PHASE_POST(3, G2) PHASE_POST(4, G3) \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) { pwin_dec_prep<NAME##_L>(p, g_img, g_bias); } \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) { phase_expand<NAME##_L, MAXW, GE>(p, g_img, g_X); } \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) \
|
||||
{ \
|
||||
phase_attn<NAME##_L, MAXW, G0>(p, g_img, g_bias, g_O, g_X); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) \
|
||||
{ \
|
||||
phase_wo<NAME##_L, MAXW, G1>(p, g_img, g_O, g_M, g_A, g_X); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) { phase_fc1<NAME##_L, MAXW, G2>(p, g_img, g_M, g_H); } \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post4(PwinParams p) \
|
||||
{ \
|
||||
phase_fc2<NAME##_L, MAXW, G3>(p, g_img, g_A, g_H, p.out24); \
|
||||
}
|
||||
|
||||
// encoder (wide shape): attention, Wo, fc1, fc2 (+ full-resolution output), patch merge
|
||||
#define PHASE_ENCODER(NAME, H, C, COUT, MAXW, G0, G1, G2, G3, GM) \
|
||||
using NAME##_L = PwinLayout<H, C, COUT>; \
|
||||
PHASE_GLOBALS(NAME, H, C, MAXW) \
|
||||
__device__ half_t g_Y[(MAXW) * 64 * (C)]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = (G0); \
|
||||
PHASE_POST(1, G1) PHASE_POST(2, G2) PHASE_POST(3, G3) PHASE_POST(4, GM) \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) { pwin_prep<NAME##_L>(p, g_img, g_bias); } \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) { phase_attn<NAME##_L, MAXW, G0>(p, g_img, g_bias, g_O); } \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) { phase_wo<NAME##_L, MAXW, G1>(p, g_img, g_O, g_M, g_A); } \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) { phase_fc1<NAME##_L, MAXW, G2>(p, g_img, g_M, g_H); } \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) \
|
||||
{ \
|
||||
phase_fc2<NAME##_L, MAXW, G3>(p, g_img, g_A, g_H, p.out32, g_Y); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_post4(PwinParams p) { phase_merge<NAME##_L, MAXW, GM>(p, g_img, g_Y); }
|
||||
+171
-85
@@ -1,17 +1,21 @@
|
||||
// Position-only-attention layers (trait flag 2: enc0, enc1, dec0) with MT token tiles per wave, so every
|
||||
// weight operand loaded from L2 is used MT times (4 / MT waves per window). Same math as pwin_core.
|
||||
#pragma once
|
||||
#if defined(PWIN_F16ACC_NATIVE) && !defined(PWIN_PK_ELEM)
|
||||
#error "PWIN_F16ACC_NATIVE needs PWIN_PK_ELEM (the element-wise paths use acc_pair / acc_set_pair)"
|
||||
#endif
|
||||
#include "pwin_layer.h"
|
||||
|
||||
template <class L, int MT>
|
||||
__device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
|
||||
half_t (*act)[L::C], half_t (*hb)[L::C])
|
||||
half_t (*act)[LDSR(L::C)], half_t (*hb)[LDSR(L::C)])
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H, NW = 4 / MT;
|
||||
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;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z);
|
||||
const uint8_t* w = p.w;
|
||||
PROF_INIT;
|
||||
|
||||
// ---- L2 norm of every tile of this wave -> hb
|
||||
#pragma unroll
|
||||
@@ -25,13 +29,23 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
#pragma unroll
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_SS)
|
||||
for (int j = 0; j < 8; ++j)
|
||||
{
|
||||
const uint32_t xj = xv[j]; // (__builtin_bit_cast of a vector subscript reads element 0: clang bug)
|
||||
const hv2 v = __builtin_bit_cast(hv2, xj), q = v * v;
|
||||
ss += (float)q[0];
|
||||
ss += (float)q[1];
|
||||
}
|
||||
#else
|
||||
for (int c = 0; c < 16; ++c)
|
||||
{
|
||||
const half_t v = op_get(xv, c);
|
||||
ss += (float)(half_t)(v * v);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
|
||||
const half_t r16 = half_rsqrt((half_t)ss);
|
||||
for (int kt = (int)hf; kt < KT; kt += 2)
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
@@ -39,8 +53,15 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
u8v hv;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_APPLY)
|
||||
{
|
||||
const uint32_t xj = xv[j], gj = gv[j]; // (see above: no __builtin_bit_cast of a vector subscript)
|
||||
hv[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, xj) * ((hv2)r16 * __builtin_bit_cast(hv2, gj)));
|
||||
}
|
||||
#else
|
||||
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
|
||||
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
|
||||
#endif
|
||||
store_row16(hrow + 16 * kt, hv);
|
||||
}
|
||||
}
|
||||
@@ -48,38 +69,58 @@ __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];
|
||||
pacc acc[MT][KT];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(w, L::BO + 32 * nt, vv);
|
||||
const pacc bo = acc_vec_p<L>(img, w, L::BO + 32 * nt, L::V_BO + nt);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[mi][nt][i] = (float)vv[i];
|
||||
acc[mi][nt] = bo;
|
||||
}
|
||||
#ifdef POS_VSHARE
|
||||
// V^T operands of every head: each wave computes those of its own key tiles, shared through LDS
|
||||
// (one slot per head, so a head's slot is never rewritten while it is read)
|
||||
__shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
|
||||
#endif
|
||||
#ifdef POS_VSHARE1
|
||||
// V^T operands shared through ONE 4 KB slot reused per head (a barrier before it is rewritten); each wave
|
||||
// computes its own key tiles, the PV loop reads the operands from LDS when it needs them (no 8 operands live
|
||||
// in registers). Same values as computing all of V in every wave.
|
||||
__shared__ __attribute__((aligned(16))) u8v vsh1[2][4][16];
|
||||
#endif
|
||||
for (int h = 0; h < H; ++h)
|
||||
{
|
||||
op_t vtop[2][4];
|
||||
#ifdef POS_VSHARE
|
||||
#ifdef POS_VSHARE1
|
||||
if (h > 0)
|
||||
block_sync(); // every wave has finished reading the previous head's V^T
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int kt = wv * MT + mi;
|
||||
f8v d = splat8(0.0f);
|
||||
pacc d = acc_zero();
|
||||
#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);
|
||||
op_img_store(&vsh[h][nt][0][0], kt * 16 + m, operand_from_f8(d));
|
||||
d = mma16p(op_lds(&hb[16 * kt + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), d);
|
||||
op_img_store(&vsh1[nt][0][0], kt * 16 + m, operand_from_acc_p(d));
|
||||
}
|
||||
block_sync();
|
||||
(void)vtop;
|
||||
#elif defined(POS_VSHARE)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int kt = wv * MT + mi;
|
||||
pacc d = acc_zero();
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < KT; ++ks)
|
||||
d = mma16p(op_lds(&hb[16 * kt + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), d);
|
||||
op_img_store(&vsh[h][nt][0][0], kt * 16 + m, operand_from_acc_p(d));
|
||||
}
|
||||
block_sync();
|
||||
#pragma unroll
|
||||
@@ -93,11 +134,11 @@ __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);
|
||||
pacc d = acc_zero();
|
||||
#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);
|
||||
vtop[nt][kt] = operand_from_f8(d);
|
||||
d = mma16p(op_lds(&hb[16 * kt + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), d);
|
||||
vtop[nt][kt] = operand_from_acc_p(d);
|
||||
}
|
||||
#endif
|
||||
op_t oop[MT][2];
|
||||
@@ -108,15 +149,21 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
op_t pop[4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
pop[kt] = op_img((const u8v*)bias, ((h * 4 + qt) * 4 + kt) * 16 + m);
|
||||
pop[kt] = op_imgt((const u8v*)bias, (h * 4 + qt) * 4 + kt);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
f8v d = splat8(0.0f);
|
||||
pacc d = acc_zero();
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
d = mma16(vtop[nt][kt], pop[kt], d);
|
||||
oop[mi][nt] = operand_from_f8(d);
|
||||
{
|
||||
#ifdef POS_VSHARE1
|
||||
d = mma16p(op_img(&vsh1[nt][0][0], kt * 16 + m), pop[kt], d);
|
||||
#else
|
||||
d = mma16p(vtop[nt][kt], pop[kt], d);
|
||||
#endif
|
||||
}
|
||||
oop[mi][nt] = operand_from_acc_p(d);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
@@ -124,10 +171,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
{
|
||||
const op_t wo = op_img(img, (L::T_WO + (2 * h + ks) * KT + nt) * 16 + m);
|
||||
const op_t wo = op_imgt(img, L::T_WO + (2 * h + ks) * KT + nt);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
acc[mi][nt] = mma16(wo, oop[mi][ks], acc[mi][nt]);
|
||||
acc[mi][nt] = mma16p(wo, oop[mi][ks], acc[mi][nt]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -137,13 +184,27 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
half_t g2[8], b2[8];
|
||||
dvec8(w, L::G2 + 32 * nt, g2);
|
||||
dvec8(w, L::B2 + 32 * nt, b2);
|
||||
residual_vector<L>(w, 0, nt, g2);
|
||||
residual_vector<L>(w, 1, nt, b2);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int tok = 16 * (wv * MT + mi) + m;
|
||||
const op_t xv = op_lds(&act[tok][16 * nt]);
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_RESID)
|
||||
// the same f16 operations on element pairs (packed f16 math; pairs stay packed into the operand)
|
||||
uint32_t mine[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
const hv2 x1 = __builtin_bit_cast(hv2, acc_pair(acc[mi][nt], j)) +
|
||||
__builtin_bit_cast(hv2, wm_op_acc_pair(xv, j));
|
||||
mine[j] = __builtin_bit_cast(uint32_t, x1 * (hv2){g2[2 * j], g2[2 * j + 1]});
|
||||
const hv2 y = x1 + (hv2){b2[2 * j], b2[2 * j + 1]};
|
||||
acc_set_pair(acc[mi][nt], j, y);
|
||||
}
|
||||
op_store(&hb[tok][16 * nt], wm_op_from_mine(mine));
|
||||
#else
|
||||
half_t mv[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
@@ -153,6 +214,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
acc[mi][nt][i] = (float)(half_t)(x1 + b2[i]);
|
||||
}
|
||||
op_store(&hb[tok][16 * nt], operand_from_dt(mv));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
@@ -165,28 +227,23 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
#pragma unroll
|
||||
for (int hn = 0; hn < 2; ++hn)
|
||||
{
|
||||
half_t bb[8];
|
||||
dvec8(w, L::B1 + 64 * c + 32 * hn, bb);
|
||||
f8v d[MT];
|
||||
const pacc bb = acc_vec_p<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn);
|
||||
pacc d[MT];
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[mi][i] = (float)bb[i];
|
||||
d[mi] = bb;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const op_t w1 = op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m);
|
||||
const op_t w1 = op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
d[mi] = mma16(w1, op_lds(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
|
||||
d[mi] = mma16p(w1, op_lds(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
half_t g[8];
|
||||
gelu8(d[mi], g);
|
||||
gop[mi][hn] = operand_from_dt(g);
|
||||
gop[mi][hn] = gelu_op_p(d[mi]);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
@@ -194,10 +251,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
{
|
||||
const op_t w2 = op_img(img, (L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m);
|
||||
const op_t w2 = op_imgt(img, L::T_W2 + (c * 2 + ks) * KT + nt);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
acc[mi][nt] = mma16(w2, gop[mi][ks], acc[mi][nt]);
|
||||
acc[mi][nt] = mma16p(w2, gop[mi][ks], acc[mi][nt]);
|
||||
}
|
||||
}
|
||||
STAMP(5);
|
||||
@@ -207,9 +264,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_f8(acc[mi][nt]));
|
||||
op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_acc_p(acc[mi][nt]));
|
||||
}
|
||||
block_sync();
|
||||
PROF(6);
|
||||
}
|
||||
|
||||
// encoder (enc0 with embedding, enc1): core + full-resolution output + patch merge
|
||||
@@ -217,14 +275,15 @@ template <class L, int MT>
|
||||
__device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C, NW = 4 / MT;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
|
||||
PwinParams p = p0;
|
||||
if constexpr (L::CB != 0)
|
||||
p.w = p0.w + L::CB;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
|
||||
STAMP(0);
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), bx = blockIdx.x, by = blockIdx.y;
|
||||
PROF_INIT;
|
||||
STAMP(20);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
@@ -232,28 +291,34 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
|
||||
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
|
||||
if constexpr (L::CIN != 0)
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
|
||||
op_t xe[L::KE];
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
xe[ks] = op_gload(xin + 16 * ks);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(p0.w, 32 * nt, vv);
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
pacc d = acc_vec_p<L>(img, p0.w, 32 * nt, L::V_EMB + nt);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EMB + ks * (C / 16) + nt) * 16 + m), xe[ks], d);
|
||||
d = mma16p(op_imgt(img, L::T_EMB + ks * (C / 16) + nt), xe[ks], d);
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_RELU)
|
||||
uint32_t mine[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
const hv2 x = __builtin_bit_cast(hv2, acc_pair(d, j));
|
||||
mine[j] = __builtin_bit_cast(uint32_t, (hv2){hmax(x[0], (half_t)0.0f), hmax(x[1], (half_t)0.0f)});
|
||||
}
|
||||
op_store(&act[tok][16 * nt], wm_op_from_mine(mine));
|
||||
#else
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = hmax((half_t)d[i], (half_t)0.0f);
|
||||
op_store(&act[tok][16 * nt], operand_from_dt(xo));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
else
|
||||
@@ -264,9 +329,9 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
|
||||
}
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
STAMP(1);
|
||||
STAMP(21);
|
||||
pos_core<L, MT>(p, img, bias, act, hb);
|
||||
STAMP(6);
|
||||
STAMP(26);
|
||||
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
@@ -280,29 +345,24 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
|
||||
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
|
||||
}
|
||||
}
|
||||
STAMP(7);
|
||||
STAMP(27);
|
||||
const int my = m >> 2, mx = m & 3;
|
||||
const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
|
||||
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
|
||||
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * L::COUT;
|
||||
for (int vt = wv; vt < L::PMT; vt += NW)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(p.w, L::PMB + 32 * vt, vv);
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
pacc d = acc_vec_p<L>(img, p.w, L::PMB + 32 * vt, L::V_PM + vt);
|
||||
for (int ks = 0; ks < C / 4; ++ks)
|
||||
{
|
||||
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
|
||||
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
|
||||
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
|
||||
d = mma16p(op_imgt(img, L::T_PM + ks * L::PMT + vt), op_lds(&act[src][16 * (ks % L::KT)]), d);
|
||||
}
|
||||
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
|
||||
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
|
||||
op_gstore(mout + L::NPW * g + 16 * u, operand_from_acc_p(d), minb, valid);
|
||||
}
|
||||
STAMP(8);
|
||||
STAMP(28);
|
||||
}
|
||||
|
||||
// dec0: patch expand + skip, core, output head
|
||||
@@ -310,10 +370,11 @@ template <class L, int MT>
|
||||
__device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C, NW = 4 / MT;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), bx = blockIdx.x, by = blockIdx.y;
|
||||
PROF_INIT;
|
||||
{
|
||||
const int W2 = p.W / 2, H2 = p.H / 2, ly = m >> 2, lx = m & 3;
|
||||
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
|
||||
@@ -322,8 +383,24 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
xls[ks] = op_gload(xl + 16 * ks);
|
||||
for (int q = wv; q < 4; q += NW)
|
||||
#ifndef POS_NO_SKIP_PREFETCH
|
||||
// the skip rows of all this wave's sub-tokens, loaded up front with xls (their latency overlaps)
|
||||
op_t svs[MT][C / 16];
|
||||
#pragma unroll
|
||||
for (int qi = 0; qi < MT; ++qi)
|
||||
{
|
||||
const int q = wv + NW * qi, dy = q >> 1, dx = q & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
|
||||
const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
|
||||
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
svs[qi][nt] = op_gload(sk + 16 * nt);
|
||||
}
|
||||
#endif
|
||||
#pragma unroll
|
||||
for (int qi = 0; qi < MT; ++qi)
|
||||
{
|
||||
const int q = wv + NW * qi;
|
||||
const int dy = q >> 1, dx = q & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
|
||||
const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
|
||||
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
|
||||
@@ -331,28 +408,39 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
{
|
||||
const int gnt = q * (C / 16) + nt;
|
||||
half_t vv[8];
|
||||
dvec8(p.w, L::EXPB + 32 * gnt, vv);
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
pacc d = acc_vec_p<L>(img, p.w, L::EXPB + 32 * gnt, L::V_EXP + gnt);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), xls[ks], d);
|
||||
d = mma16p(op_imgt(img, L::T_EXP + ks * L::EXPN + gnt), xls[ks], d);
|
||||
#ifndef POS_NO_SKIP_PREFETCH
|
||||
const op_t sv = svs[qi][nt];
|
||||
(void)sk;
|
||||
#else
|
||||
const op_t sv = op_gload(sk + 16 * nt);
|
||||
#endif
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_EXPAND)
|
||||
uint32_t mine[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
mine[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, acc_pair(d, j)) +
|
||||
__builtin_bit_cast(hv2, wm_op_acc_pair(sv, j)));
|
||||
op_store(&act[t][16 * nt], wm_op_from_mine(mine));
|
||||
#else
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
|
||||
op_store(&act[t][16 * nt], operand_from_dt(xo));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
PROF(8);
|
||||
PwinParams q = p;
|
||||
q.w = p.w + L::CB;
|
||||
pos_core<L, MT>(q, img, bias, act, hb);
|
||||
PROF(9);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
@@ -362,29 +450,26 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(q.w, L::HB + 32 * nt, vv);
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
pacc d = acc_vec_p<L>(img, q.w, L::HB + 32 * nt, L::V_HB + nt);
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < L::KT; ++kt)
|
||||
d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
|
||||
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
|
||||
d = mma16p(op_imgt(img, L::T_HEAD + kt * (L::NOUTA / 16) + nt), op_lds(&act[tok][16 * kt]), d);
|
||||
op_gstore(hout + 16 * nt, operand_from_acc_p(d), inb, L::NOUT - 16 * nt);
|
||||
}
|
||||
}
|
||||
PROF(10);
|
||||
}
|
||||
|
||||
#define POS_ENCODER(NAME, H, C, COUT, CIN, MT) \
|
||||
using NAME##_L = PwinEmbLayout<H, C, COUT, true, CIN>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (pwin_gelu_prep_items(NAME##_L::PREP_ITEMS) + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_gelu_prep(); \
|
||||
pwin_emb_prep<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
|
||||
@@ -394,13 +479,14 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
|
||||
|
||||
#define POS_DECODER(NAME, H, C, CL, NOUT, NOUTA, MT) \
|
||||
using NAME##_L = PwinDecLayout<H, C, CL, true, NOUT, NOUTA>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (pwin_gelu_prep_items(NAME##_L::PREP_ITEMS) + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_gelu_prep(); \
|
||||
pwin_dec_prep<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
|
||||
|
||||
@@ -0,0 +1,379 @@
|
||||
// Single-buffer position-only-attention encoder experiment (trait flag 2: enc0, enc1, dec0) with MT token tiles per wave, so every
|
||||
// weight operand loaded from L2 is used MT times (4 / MT waves per window). Same math as pwin_core.
|
||||
#pragma once
|
||||
#include "pwin_layer.h"
|
||||
|
||||
template <class L, int MT>
|
||||
__device__ void pos_recompute_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
|
||||
half_t (*act)[LDSR(L::C)], half_t (*hb)[LDSR(L::C)])
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H, NW = 4 / MT;
|
||||
static_assert(L::POS, "position-only attention layers");
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z);
|
||||
const uint8_t* w = p.w;
|
||||
PROF_INIT;
|
||||
|
||||
// ---- L2 norm of every tile of this wave -> hb
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const half_t* arow = act[16 * (wv * MT + mi) + m];
|
||||
half_t* hrow = hb[16 * (wv * MT + mi) + m];
|
||||
float ss = 0.0f;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
#pragma unroll
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_SS)
|
||||
for (int j = 0; j < 8; ++j)
|
||||
{
|
||||
const uint32_t xj = xv[j];
|
||||
const hv2 v = __builtin_bit_cast(hv2, xj), q = v * v;
|
||||
ss += (float)q[0];
|
||||
ss += (float)q[1];
|
||||
}
|
||||
#else
|
||||
for (int c = 0; c < 16; ++c)
|
||||
{
|
||||
const half_t v = op_get(xv, c);
|
||||
ss += (float)(half_t)(v * v);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
const half_t r16 = half_rsqrt((half_t)ss);
|
||||
for (int kt = (int)hf; kt < KT; kt += 2)
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
|
||||
u8v hv;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_APPLY)
|
||||
{
|
||||
const uint32_t xj = xv[j], gj = gv[j];
|
||||
hv[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, xj) * ((hv2)r16 * __builtin_bit_cast(hv2, gj)));
|
||||
}
|
||||
#else
|
||||
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
|
||||
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
|
||||
#endif
|
||||
store_row16(hrow + 16 * kt, hv);
|
||||
}
|
||||
}
|
||||
block_sync(); // V of all 64 tokens is computed from hb
|
||||
STAMP(2);
|
||||
|
||||
// ---- attention: O = P V per head (P from the table), output projection accumulated head by head
|
||||
f8v acc[MT][KT];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
const f8v bo = acc_vec<L>(img, w, L::BO + 32 * nt, L::V_BO + nt);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
acc[mi][nt] = bo;
|
||||
}
|
||||
#ifdef POS_VSHARE
|
||||
// V^T operands of every head: each wave computes those of its own key tiles, shared through LDS
|
||||
// (one slot per head, so a head's slot is never rewritten while it is read)
|
||||
__shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
|
||||
#endif
|
||||
#ifdef POS_VSHARE1
|
||||
// V^T operands shared through ONE 4 KB slot reused per head (a barrier before it is rewritten); each wave
|
||||
// computes its own key tiles, the PV loop reads the operands from LDS when it needs them (no 8 operands live
|
||||
// in registers). Same values as computing all of V in every wave.
|
||||
__shared__ __attribute__((aligned(16))) u8v vsh1[2][4][16];
|
||||
#endif
|
||||
for (int h = 0; h < H; ++h)
|
||||
{
|
||||
op_t vtop[2][4];
|
||||
#ifdef POS_VSHARE1
|
||||
if (h > 0)
|
||||
block_sync(); // every wave has finished reading the previous head's V^T
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int kt = wv * MT + mi;
|
||||
f8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < KT; ++ks)
|
||||
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), d);
|
||||
op_img_store(&vsh1[nt][0][0], kt * 16 + m, operand_from_f8(d));
|
||||
}
|
||||
block_sync();
|
||||
(void)vtop;
|
||||
#elif defined(POS_VSHARE)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int kt = wv * MT + mi;
|
||||
f8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < KT; ++ks)
|
||||
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), d);
|
||||
op_img_store(&vsh[h][nt][0][0], kt * 16 + m, operand_from_f8(d));
|
||||
}
|
||||
block_sync();
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
|
||||
#else
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
f8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < KT; ++ks)
|
||||
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), d);
|
||||
vtop[nt][kt] = operand_from_f8(d);
|
||||
}
|
||||
#endif
|
||||
op_t oop[MT][2];
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int qt = wv * MT + mi;
|
||||
op_t pop[4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
pop[kt] = op_imgt((const u8v*)bias, (h * 4 + qt) * 4 + kt);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
f8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
#ifdef POS_VSHARE1
|
||||
d = mma16(op_img(&vsh1[nt][0][0], kt * 16 + m), pop[kt], d);
|
||||
#else
|
||||
d = mma16(vtop[nt][kt], pop[kt], d);
|
||||
#endif
|
||||
}
|
||||
oop[mi][nt] = operand_from_f8(d);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
{
|
||||
const op_t wo = op_imgt(img, L::T_WO + (2 * h + ks) * KT + nt);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
acc[mi][nt] = mma16(wo, oop[mi][ks], acc[mi][nt]);
|
||||
}
|
||||
}
|
||||
|
||||
STAMP(3);
|
||||
// All attention readers finish before the shared input rows are reused.
|
||||
block_sync();
|
||||
// Recover x0 exactly from the unchanged input (repeat the original embedding,
|
||||
// or reload a no-embedding layer), then write the MLP input into the same LDS rows.
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
half_t g2[8], b2[8];
|
||||
residual_vector<L>(w, 0, nt, g2);
|
||||
residual_vector<L>(w, 1, nt, b2);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int tok = 16 * (wv * MT + mi) + m;
|
||||
const int X = mirror(8 * (int)blockIdx.x - p.sx + (tok & 7), p.W);
|
||||
const int Y = mirror(8 * (int)blockIdx.y - p.sy + (tok >> 3), p.H);
|
||||
half_t x0[8];
|
||||
if constexpr (L::CIN != 0)
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
|
||||
f8v d = acc_vec<L>(img, w - L::CB, 32 * nt, L::V_EMB + nt);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
d = mma16(op_imgt(img, L::T_EMB + ks * KT + nt), op_gload(xin + 16 * ks), d);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
x0[i] = hmax((half_t)d[i], (half_t)0.0f);
|
||||
}
|
||||
else
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
|
||||
const op_t xv = op_gload(xin + 16 * nt);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
x0[i] = wm_op_acc_elem(xv, i);
|
||||
}
|
||||
half_t mv[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const half_t x1 = (half_t)acc[mi][nt][i] + x0[i];
|
||||
mv[i] = x1 * g2[i];
|
||||
acc[mi][nt][i] = (float)(half_t)(x1 + b2[i]);
|
||||
}
|
||||
op_store(&hb[tok][16 * nt], operand_from_dt(mv));
|
||||
}
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
|
||||
STAMP(4);
|
||||
// ---- MLP: C/8 chunks of 32 hidden; each weight operand serves the MT tiles
|
||||
for (int c = 0; c < L::NMLP; ++c)
|
||||
{
|
||||
op_t gop[MT][2];
|
||||
#pragma unroll
|
||||
for (int hn = 0; hn < 2; ++hn)
|
||||
{
|
||||
const f8v bb = acc_vec<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn);
|
||||
f8v d[MT];
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
d[mi] = bb;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const op_t w1 = op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
d[mi] = mma16(w1, op_lds(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
gop[mi][hn] = gelu_op(d[mi]);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
{
|
||||
const op_t w2 = op_imgt(img, L::T_W2 + (c * 2 + ks) * KT + nt);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
acc[mi][nt] = mma16(w2, gop[mi][ks], acc[mi][nt]);
|
||||
}
|
||||
}
|
||||
STAMP(5);
|
||||
// y -> act
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_f8(acc[mi][nt]));
|
||||
}
|
||||
block_sync();
|
||||
PROF(6);
|
||||
}
|
||||
|
||||
// encoder (enc0 with embedding, enc1): core + full-resolution output + patch merge
|
||||
template <class L, int MT>
|
||||
__device__ void pos_recompute_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C, NW = 4 / MT;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
|
||||
half_t (*hb)[LDSR(C)] = act; // one allocation: x0 -> norm -> MLP input -> final output
|
||||
PwinParams p = p0;
|
||||
if constexpr (L::CB != 0)
|
||||
p.w = p0.w + L::CB;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), bx = blockIdx.x, by = blockIdx.y;
|
||||
PROF_INIT;
|
||||
STAMP(20);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int tok = 16 * (wv * MT + mi) + m, ty = tok >> 3, tx = tok & 7;
|
||||
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
|
||||
if constexpr (L::CIN != 0)
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
|
||||
op_t xe[L::KE];
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
xe[ks] = op_gload(xin + 16 * ks);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
{
|
||||
f8v d = acc_vec<L>(img, p0.w, 32 * nt, L::V_EMB + nt);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
d = mma16(op_imgt(img, L::T_EMB + ks * (C / 16) + nt), xe[ks], d);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = hmax((half_t)d[i], (half_t)0.0f);
|
||||
op_store(&act[tok][16 * nt], operand_from_dt(xo));
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
|
||||
}
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
STAMP(21);
|
||||
pos_recompute_core<L, MT>(p, img, bias, act, hb);
|
||||
STAMP(26);
|
||||
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int tok = 16 * (wv * MT + mi) + m;
|
||||
const int Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
|
||||
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
|
||||
{
|
||||
half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
|
||||
}
|
||||
}
|
||||
STAMP(27);
|
||||
const int my = m >> 2, mx = m & 3;
|
||||
const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
|
||||
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
|
||||
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * L::COUT;
|
||||
for (int vt = wv; vt < L::PMT; vt += NW)
|
||||
{
|
||||
f8v d = acc_vec<L>(img, p.w, L::PMB + 32 * vt, L::V_PM + vt);
|
||||
for (int ks = 0; ks < C / 4; ++ks)
|
||||
{
|
||||
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
|
||||
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
|
||||
d = mma16(op_imgt(img, L::T_PM + ks * L::PMT + vt), op_lds(&act[src][16 * (ks % L::KT)]), d);
|
||||
}
|
||||
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
|
||||
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
|
||||
}
|
||||
STAMP(28);
|
||||
}
|
||||
|
||||
#define POS_RECOMPUTE_ENCODER(NAME, H, C, COUT, CIN, MT) \
|
||||
using NAME##_L = PwinEmbLayout<H, C, COUT, true, CIN>; \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (pwin_gelu_prep_items(NAME##_L::PREP_ITEMS) + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_gelu_prep(); \
|
||||
pwin_emb_prep<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
|
||||
{ \
|
||||
pos_recompute_encoder<NAME##_L, MT>(p, g_img, g_bias); \
|
||||
}
|
||||
|
||||
+124
-102
@@ -15,27 +15,71 @@ template <class L, int NG> struct WideDims
|
||||
// NA: heads per attention round (their K / V images share R)
|
||||
static constexpr int NA = NG < 4 ? NG : 4;
|
||||
static constexpr int C = L::C, KT = L::KT, H = L::H, NTL = (KT + NG - 1) / NG, ROUNDS = (H + NA - 1) / NA;
|
||||
static constexpr int X0_BYTES = 64 * C * 2, KV_BYTES = NA * 2 * 64 * 32 * 2;
|
||||
// K rows (64 x 32), V^T rows (32 x 64), O rows (64 x 32), hidden rows (64 x 32): all padded (LDSR)
|
||||
static constexpr int KL_BYTES = 64 * LDSR(32) * 2, VT_BYTES = 32 * LDSR(64) * 2;
|
||||
static constexpr int X0_BYTES = 64 * LDSR(C) * 2, KV_BYTES = NA * (KL_BYTES + VT_BYTES);
|
||||
// MLP hidden chunks: one per group and round, double-buffered while that fits
|
||||
static constexpr int NB = NG <= 4 ? 2 : 1, MROUNDS = (L::NMLP + NG - 1) / NG, HID_BYTES = NB * NG * 64 * 32 * 2;
|
||||
static constexpr int NB = NG <= 4 ? 2 : 1, MROUNDS = (L::NMLP + NG - 1) / NG, HID_BYTES = NB * NG * 64 * LDSR(32) * 2;
|
||||
static constexpr int R_BYTES = X0_BYTES > KV_BYTES ? (X0_BYTES > HID_BYTES ? X0_BYTES : HID_BYTES)
|
||||
: (KV_BYTES > HID_BYTES ? KV_BYTES : HID_BYTES);
|
||||
static_assert(!L::POS, "position-only attention layers use pwin_pos.h");
|
||||
};
|
||||
|
||||
// QKV chains u = U0 .. U0 + NU - 1 of head h for this wave's token tile (u = 2 j + nt; j: Q, K, V): Q stays in
|
||||
// registers (qop), K and V^T go to LDS slot s. Each chain runs over kt in order (exact).
|
||||
template <class L, int U0, int NU>
|
||||
__device__ __forceinline__ void wide_qkv(const u8v* __restrict__ img, const half_t* hrow, int h, int s, int tok,
|
||||
op_t (&qop)[2], half_t (*kl)[64][LDSR(32)], half_t (*vt)[32][LDSR(64)])
|
||||
{
|
||||
constexpr int KT = L::KT;
|
||||
const uint32_t m = lane_id() & 15;
|
||||
f8v dq[NU];
|
||||
#pragma unroll
|
||||
for (int uu = 0; uu < NU; ++uu)
|
||||
dq[uu] = splat8(0.0f);
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const op_t row = op_lds(hrow + 16 * kt);
|
||||
#pragma unroll
|
||||
for (int uu = 0; uu < NU; ++uu)
|
||||
{
|
||||
const int u = U0 + uu;
|
||||
dq[uu] = mma16(op_imgt(img, L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)), row, dq[uu]);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int uu = 0; uu < NU; ++uu)
|
||||
{
|
||||
const int u = U0 + uu, j = u >> 1, nt = u & 1;
|
||||
const f8v d = dq[uu];
|
||||
if (j == 0)
|
||||
qop[nt] = operand_from_f8(d);
|
||||
else if (j == 1)
|
||||
op_store(&kl[s][tok][16 * nt], operand_from_f8(d));
|
||||
else
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
vt[s][16 * nt + wm_acc_row(i)][tok] = (half_t)d[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Core for one window. On entry R holds x0 (64 x C); on return hb holds the block output y.
|
||||
template <class L, int NG>
|
||||
__device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
|
||||
half_t (*hb)[L::C], uint8_t* R)
|
||||
half_t (*hb)[LDSR(L::C)], uint8_t* R)
|
||||
{
|
||||
using D = WideDims<L, NG>;
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H, NTL = D::NTL;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const uint8_t* w = p.w;
|
||||
half_t(*act)[C] = (half_t(*)[C])R;
|
||||
half_t(*act)[LDSR(C)] = (half_t(*)[LDSR(C)])R;
|
||||
const half_t* arow = act[tok];
|
||||
half_t* hrow = hb[tok];
|
||||
PROF_INIT;
|
||||
PROF(1); // entry (x0 load + barrier)
|
||||
|
||||
// ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb (all groups of a token tile compute the sum)
|
||||
float ss = 0.0f;
|
||||
@@ -49,7 +93,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
ss += (float)(half_t)(v * v);
|
||||
}
|
||||
}
|
||||
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
|
||||
const half_t r16 = half_rsqrt((half_t)ss);
|
||||
for (int kt = 2 * g + (int)hf; kt < KT; kt += 2 * NG)
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
@@ -76,6 +120,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
}
|
||||
}
|
||||
block_sync(); // hb complete, R free
|
||||
PROF(2); // norm
|
||||
|
||||
// ---- attention: NG heads per round; the output projection accumulates head by head
|
||||
f8v acc[NTL];
|
||||
@@ -84,65 +129,57 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
{
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(w, L::BO + 32 * nt, vv);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[j][i] = (float)vv[i];
|
||||
}
|
||||
acc[j] = acc_vec<L>(img, w, L::BO + 32 * nt, L::V_BO + nt);
|
||||
}
|
||||
half_t(*kl)[64][32] = (half_t(*)[64][32])R; // [group][token][dim]
|
||||
half_t(*kl)[64][LDSR(32)] = (half_t(*)[64][LDSR(32)])R; // [group][token][dim]
|
||||
constexpr int NA = D::NA;
|
||||
half_t(*vt)[32][64] = (half_t(*)[32][64])(R + NA * 64 * 32 * 2); // [group][dim][token]
|
||||
half_t(*ob)[64][32] = (half_t(*)[64][32])R; // [group][token][dim] (after K / V)
|
||||
half_t(*vt)[32][LDSR(64)] = (half_t(*)[32][LDSR(64)])(R + NA * D::KL_BYTES); // [group][dim][token]
|
||||
half_t(*ob)[64][LDSR(32)] = (half_t(*)[64][LDSR(32)])R; // [group][token][dim] (after K / V)
|
||||
// NG >= 2 NA: the groups beyond the NA heads of a round would idle during QKV and attention; instead group
|
||||
// g + NA computes half of the six QKV chains of group g's head (K tile 1 and V), group g the other half
|
||||
// (Q, kept in registers for its attention, and K tile 0)
|
||||
constexpr bool SPLIT = NG >= 2 * NA;
|
||||
const int hs = SPLIT ? g % NA : g, part = SPLIT ? g / NA : 0;
|
||||
for (int r = 0; r < D::ROUNDS; ++r)
|
||||
{
|
||||
const int h = g < NA ? r * NA + g : H;
|
||||
const int h = (SPLIT ? g < 2 * NA : g < NA) ? r * NA + hs : H;
|
||||
const bool att = h < H && part == 0;
|
||||
op_t qop[2], oop[2];
|
||||
#ifdef PWIN_BIAS_PREFETCH
|
||||
u4v bpre[4];
|
||||
if (att)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
bpre[kt] = bias[((h * 4 + t) * 4 + kt) * 32 + l];
|
||||
}
|
||||
#endif
|
||||
if (h < H)
|
||||
{
|
||||
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
|
||||
f8v dq[6];
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 6; ++u)
|
||||
dq[u] = splat8(0.0f);
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
if constexpr (SPLIT)
|
||||
{
|
||||
const op_t row = op_lds(hrow + 16 * kt);
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 6; ++u)
|
||||
dq[u] = mma16(op_img(img, (L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m), row, dq[u]);
|
||||
if (part == 0)
|
||||
wide_qkv<L, 0, 3>(img, hrow, h, hs, tok, qop, kl, vt);
|
||||
else
|
||||
wide_qkv<L, 3, 3>(img, hrow, h, hs, tok, qop, kl, vt);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 3; ++j)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
const f8v d = dq[2 * j + nt];
|
||||
if (j == 0)
|
||||
qop[nt] = operand_from_f8(d);
|
||||
else if (j == 1)
|
||||
{
|
||||
op_store(&kl[g][tok][16 * nt], operand_from_f8(d));
|
||||
}
|
||||
else
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
vt[g][16 * nt + wm_acc_row(i)][tok] = (half_t)d[i];
|
||||
}
|
||||
}
|
||||
else
|
||||
wide_qkv<L, 0, 6>(img, hrow, h, hs, tok, qop, kl, vt);
|
||||
}
|
||||
block_sync();
|
||||
if (h < H)
|
||||
PROF(3); // QKV
|
||||
if (att)
|
||||
{
|
||||
half_t e[4][8];
|
||||
float rs = 0.0f;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
#ifdef PWIN_BIAS_PREFETCH
|
||||
const u4v bi = bpre[kt];
|
||||
#else
|
||||
const u4v bi = bias[((h * 4 + t) * 4 + kt) * 32 + l];
|
||||
#endif
|
||||
f8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
@@ -152,7 +189,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
}
|
||||
#pragma unroll
|
||||
for (int ds = 0; ds < 2; ++ds)
|
||||
d = mma16(op_lds(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
|
||||
d = mma16(op_lds(&kl[hs][16 * kt + m][16 * ds]), qop[ds], d);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
@@ -162,7 +199,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
}
|
||||
}
|
||||
rs = row_sum_total(rs, e);
|
||||
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
|
||||
const half_t rc = half_reciprocal((half_t)rs);
|
||||
op_t pop[4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
@@ -179,15 +216,17 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
f8v 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);
|
||||
d = mma16(op_lds(&vt[hs][16 * nt + m][16 * kt]), pop[kt], d);
|
||||
oop[nt] = operand_from_f8(d);
|
||||
}
|
||||
}
|
||||
PROF(4); // attention
|
||||
block_sync(); // K / V of the round consumed
|
||||
if (h < H)
|
||||
PROF(5); // barrier after attention
|
||||
if (att)
|
||||
{
|
||||
op_store(&ob[g][tok][0], oop[0]);
|
||||
op_store(&ob[g][tok][16], oop[1]);
|
||||
op_store(&ob[hs][tok][0], oop[0]);
|
||||
op_store(&ob[hs][tok][16], oop[1]);
|
||||
}
|
||||
block_sync();
|
||||
for (int jh = 0; jh < NA; ++jh)
|
||||
@@ -202,13 +241,14 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh) * KT + nt) * 16 + m), o0, acc[j]);
|
||||
acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m), o1, acc[j]);
|
||||
acc[j] = mma16(op_imgt(img, L::T_WO + (2 * hh) * KT + nt), o0, acc[j]);
|
||||
acc[j] = mma16(op_imgt(img, L::T_WO + (2 * hh + 1) * KT + nt), o1, acc[j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (r + 1 < D::ROUNDS)
|
||||
block_sync(); // ob is overwritten by the next round's K / V
|
||||
PROF(6); // O exchange + Wo
|
||||
}
|
||||
|
||||
// ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
|
||||
@@ -219,8 +259,8 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
if (nt < KT)
|
||||
{
|
||||
half_t mv[8], g2[8], b2[8];
|
||||
dvec8(w, L::G2 + 32 * nt, g2);
|
||||
dvec8(w, L::B2 + 32 * nt, b2);
|
||||
residual_vector<L>(w, 0, nt, g2);
|
||||
residual_vector<L>(w, 1, nt, b2);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
@@ -232,9 +272,10 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
PROF(7); // residual
|
||||
|
||||
// ---- MLP: rounds of NG chunks of 32 hidden (double-buffered in R)
|
||||
half_t(*hid)[NG][64][32] = (half_t(*)[NG][64][32])R;
|
||||
half_t(*hid)[NG][64][LDSR(32)] = (half_t(*)[NG][64][LDSR(32)])R;
|
||||
for (int rr = 0; rr < D::MROUNDS; ++rr)
|
||||
{
|
||||
const int c = rr * NG + g;
|
||||
@@ -243,21 +284,14 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
{
|
||||
if (c >= L::NMLP)
|
||||
break;
|
||||
f8v d;
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(w, L::B1 + 64 * c + 32 * hn, vv);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
}
|
||||
f8v d = acc_vec<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn);
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
|
||||
half_t gg[8];
|
||||
gelu8(d, gg);
|
||||
op_store(&hid[rr % D::NB][g][tok][16 * hn], operand_from_dt(gg));
|
||||
d = mma16(op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn), op_lds(hrow + 16 * kt), d);
|
||||
op_store(&hid[rr % D::NB][g][tok][16 * hn], gelu_op(d));
|
||||
}
|
||||
PROF(8); // MLP W1 + gelu
|
||||
block_sync();
|
||||
PROF(9); // MLP barrier
|
||||
for (int jc = 0; jc < NG; ++jc)
|
||||
{
|
||||
const int cc = rr * NG + jc;
|
||||
@@ -270,13 +304,14 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2) * KT + nt) * 16 + m), h0, acc[j]);
|
||||
acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m), h1, acc[j]);
|
||||
acc[j] = mma16(op_imgt(img, L::T_W2 + (cc * 2) * KT + nt), h0, acc[j]);
|
||||
acc[j] = mma16(op_imgt(img, L::T_W2 + (cc * 2 + 1) * KT + nt), h1, acc[j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (D::NB == 1 && rr + 1 < D::MROUNDS)
|
||||
block_sync();
|
||||
PROF(10); // MLP W2
|
||||
}
|
||||
// y -> hb (the last reads of hb were the W1 operands before the final barrier)
|
||||
#pragma unroll
|
||||
@@ -289,15 +324,16 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
PROF(11); // y store + barrier
|
||||
}
|
||||
|
||||
// x0 rows of a window (mirror padded) -> R
|
||||
template <class L, int NG> __device__ void wide_load_x0(const PwinParams& p, uint8_t* R)
|
||||
{
|
||||
constexpr int C = L::C;
|
||||
half_t(*act)[C] = (half_t(*)[C])R;
|
||||
half_t(*act)[LDSR(C)] = (half_t(*)[LDSR(C)])R;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const int X = mirror(8 * (int)blockIdx.x - p.sx + (tok & 7), p.W), Y = mirror(8 * (int)blockIdx.y - p.sy + (tok >> 3), p.H);
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
|
||||
for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
|
||||
@@ -305,10 +341,10 @@ template <class L, int NG> __device__ void wide_load_x0(const PwinParams& p, uin
|
||||
}
|
||||
|
||||
// y rows of the window's own tokens -> dst (C per token)
|
||||
template <class L, int NG> __device__ void wide_store_y(const PwinParams& p, half_t (*hb)[L::C], uint8_t* dst)
|
||||
template <class L, int NG> __device__ void wide_store_y(const PwinParams& p, half_t (*hb)[LDSR(L::C)], uint8_t* dst)
|
||||
{
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const int Y0 = 8 * (int)blockIdx.y - p.sy + (tok >> 3), X0 = 8 * (int)blockIdx.x - p.sx + (tok & 7);
|
||||
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
|
||||
{
|
||||
@@ -322,7 +358,7 @@ template <class L, int NG>
|
||||
__device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C, COUT = L::COUT;
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
|
||||
wide_load_x0<L, NG>(p, R);
|
||||
block_sync();
|
||||
@@ -331,7 +367,7 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
|
||||
|
||||
// patch merge: 16 merged tokens, K = 4 sub-tokens x C, the allocated n-tiles over the waves
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z);
|
||||
const int my = m >> 2, mx = m & 3;
|
||||
const int Ym = (8 * (int)blockIdx.y - p.sy) / 2 + my, Xm = (8 * (int)blockIdx.x - p.sx) / 2 + mx;
|
||||
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
|
||||
@@ -339,19 +375,12 @@ __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;
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(w, L::PMB + 32 * vt, vv);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
}
|
||||
f8v d = acc_vec<L>(img, w, L::PMB + 32 * vt, L::V_PM + vt);
|
||||
for (int ks = 0; ks < C / 4; ++ks)
|
||||
{
|
||||
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
|
||||
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
|
||||
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&hb[src][16 * (ks % L::KT)]), d);
|
||||
d = mma16(op_imgt(img, L::T_PM + ks * L::PMT + vt), op_lds(&hb[src][16 * (ks % L::KT)]), d);
|
||||
}
|
||||
const int gw = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
|
||||
op_gstore(mout + L::NPW * gw + 16 * u, operand_from_f8(d), minb, valid);
|
||||
@@ -362,7 +391,7 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
|
||||
template <class L, int NG>
|
||||
__device__ void wide_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][L::C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(L::C)];
|
||||
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
|
||||
wide_load_x0<L, NG>(p, R);
|
||||
block_sync();
|
||||
@@ -374,11 +403,11 @@ template <class L, int NG>
|
||||
__device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C;
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
|
||||
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
|
||||
half_t(*act)[C] = (half_t(*)[C])R;
|
||||
half_t(*act)[LDSR(C)] = (half_t(*)[LDSR(C)])R;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, q = wv & 3, g = wv >> 2, bx = blockIdx.x, by = blockIdx.y;
|
||||
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), q = wv & 3, g = wv >> 2, bx = blockIdx.x, by = blockIdx.y;
|
||||
|
||||
// ---- patch expand: waves of sub-token q = (dy, dx) produce its tiles nt = g + NG j for the 16 low-res tokens
|
||||
{
|
||||
@@ -392,16 +421,9 @@ __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;
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(p.w, L::EXPB + 32 * gnt, vv);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
}
|
||||
f8v d = acc_vec<L>(img, p.w, L::EXPB + 32 * gnt, L::V_EXP + gnt);
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
|
||||
d = mma16(op_imgt(img, L::T_EXP + ks * L::EXPN + gnt), op_gload(xl + 16 * ks), d);
|
||||
const op_t sv = op_gload(sk + 16 * nt);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
@@ -420,7 +442,7 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
|
||||
|
||||
#define WIDE_KERNEL(NAME, PREP, BODY, NG, ...) \
|
||||
using NAME##_L = __VA_ARGS__; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 * (NG); \
|
||||
@@ -435,5 +457,5 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
|
||||
}
|
||||
|
||||
#define WIDE_ENCODER(NAME, H, C, COUT, NG) WIDE_KERNEL(NAME, pwin_prep, wide_encoder, NG, PwinLayout<H, C, COUT>)
|
||||
#define WIDE_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep, wide_plain, NG, PwinLayout<H, C, 16>)
|
||||
#define WIDE_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep_nomerge, wide_plain, NG, PwinLayout<H, C, 16>)
|
||||
#define WIDE_DECODER(NAME, H, C, CL, NG) WIDE_KERNEL(NAME, pwin_dec_prep, wide_decoder, NG, PwinDecLayout<H, C, CL>)
|
||||
Reference in new issue
Block a user