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

Prefill GDN: STRATA_GDN_TC=1 runs the DeltaNet recurrence as a chunked scan on the tensor cores (sm_80+, opt-in)

Browse Source 
The chunked WY-form recurrence with every product on the tensor cores, for sm_80+ cards with 48+ SMs and 122 KB of
shared memory a block.  STRATA_GDN_TC=1 dispatches it before the STRATA_GDN_CHUNKED branch for prompt chunks of 128+
tokens (kGdnTcMin); a card below that, or a failed scratch allocation, takes the kernels as before.  Unset = the
kernels as before.  Rounding-level: the output is not the FP32 kernels' bits.  CUDA only (gdn_tc.cuh is under
!__HIPCC__; the dispatch sits in kernels.cu's CUDA branch).

src/prefill/gdn_tc.cuh (new): the mma.sync / ldmatrix wrappers, the f16 hi/lo split with power-of-two scaling, the
dynamic scale exponent, the per-chunk products a warp runs and the two kernels:
- gdn_tc_prep_kernel: gdn_chunk_prep_kernel's arithmetic, plus the k and q f16 hi / lo images x 2^8 (built in shared
  memory from the float4 registers the rows arrive in), the P images x 2^8, T as f16 hi / lo images x 2^et (et from
  the warp's largest |T|) and the token scalars (beta, exp(gamma), exp(gL - gamma), gamma, the scale 2^-et) as a
  float4 a token, into gdn_tc_scratch (60.3 MB a device).  Its Gram products k_t.k_i and q_t.k_i run as f16x3 mma
  (48 mma and 32 ldmatrix a warp in place of ~1,150 LDS a thread); after them, warps VPK + j take P and the scalars
  of value head j while warps j build A and solve with lane i's column of T in registers.
- gdn_tc_scan_kernel: a block per value head (48 blocks of 256 threads, one wave); warp w holds value columns
  16w..16w+15 of the transposed state (16 m16n8 f32 accumulator tiles) across every chunk of a super-block.
  X^T / QS^T = S^T K^T / Q^T on f16x3 (the state x 2^es), D^T = Y^T T^T on f16x3 from the prep's T images (6 of 8
  tiles, Y^T x 2^ey), o^T = e QS^T + D^T P^T on f16x3 with P's triangle skipped, S^T = eL S^T + Dd^T K on f16x3
  through ldmatrix.trans: 324 mma a warp a chunk, one __syncthreads a chunk.  Each chunk is staged with 15 cp.async
  a thread right after the chunk barrier from the prep's images.  The warp maxima run through redux.sync over short
  chains (eight chains of eight, four of four); the warp maximum and the shuffles use the full mask (a lane-divergent
  L2 prefetch before the first maximum of a chunk made __activemask() unsafe; the kernel has no early return, so
  the full mask adds no reconvergence code).  255 registers, 0 spills (prep 52).
The header is included inside the anonymous namespace of kernels.cu and of the harness, so the harness runs the
engine's text verbatim.

src/prefill/gdn_rec_parity.cu: check_variant (any launcher of gdn_rec_chunked's shape), the tc variant at
T = 1..4099, 8192, 32768 and six stress inputs (state x 0.01 / x 50 / 0, v x 50, slow gates, coherent keys) at 4099
and 32768 against gdn_rec_kh_kernel (gate 1e-4); --selftest-frag (each helper on one warp against a double
reference); --bench-tc (the tensor-core series alone, for A / B runs of two builds); --long (T = 51,820).

Measured on a CMP 170HX (GA100, 70 SMs):
- gdn_rec_parity: output <= 2.3e-6, state <= 2.3e-6 against gdn_rec_kh_kernel over every T and stress case (the
  coherent-keys case 1.8e-6 / 8.5e-7); --selftest-frag (b) 9.1e-7, (c) 4.2e-7, (d) 3.3e-7, (e) 0 words differ.
  (--long T = 51,820 at the first cooperative-staging kernel: output 1.17e-06, state 2.58e-07.)
- --bench-tc (two builds back to back, twice, medians of alternating runs), us a token at 2048 / 8192 / 32768:
  v0 0.288 / 0.287 / 0.287 (2.79x gdn_rec_kh_kernel, 0.80 us); redux maxima 0.282 / 0.281 / 0.281 (2.85x);
  full mask 0.283 / 0.282 / 0.282 (gdn_rec_kh_kernel 1.65 / 6.56 / 26.23 ms); cp.async staging from the prep's
  images 0.262 / 0.260 / 0.260 (3.07x); short max chains 0.261 / 0.260 / 0.259; the prep's Gram products on the
  tensor cores 0.238 / 0.238 / 0.238 (3.37x; prep 233 -> 153 us a 2048-token super-block under ncu); the prep's
  serial phase on six warps 0.219 / 0.218 / 0.218 (3.67x; prep 153 -> 95 us); D^T on f16x3 0.214 / 0.214 / 0.214
  (3.75x).
- Engine (Qwen3.8-Flash-Next IQ3_S, STRATA_PREFILL_TIMING), with the first engine kernel: the "gdn recurrence"
  phase of a 3,259-token prompt 105 -> 43 ms, of a 51,820-token prompt 1,705 -> 686 ms; prompts 3,849 -> 4,104
  tok/s (3K), 4,163 -> 4,439 (13K), 4,060 -> 4,332 (52K); the 11-prompt greedy set identical to the FP32 kernels.

Measured and not kept:
- stager warps with named barriers (two more warps staging a chunk ahead, 320 threads): 0.299 / 0.297 / 0.297 us a
  token against 0.282 / 0.281 / 0.281 (6 % slower): a block of more than 8 warps is capped at 168 registers a
  thread on GA100 and the consumers spill.
- the output norm fused into the scan: correct (y16 within 4.4e-4 of its largest value against
  gdn_out_norm_kernel), but 0.303 against 0.214 + 0.057 us a token at 32768 (the chunk loop ~2,010 -> ~2,900
  instructions at the 255-register cap).
- the prep overlapped with the previous super-block's scan on low / high priority streams joined by events (a
  two-slot scratch): it wins in the harness (0.214 -> 0.215 / 0.191 / 0.184 us a token at 2048 / 8192 / 32768) but
  loses in the engine: a 3,259-token prompt's GPU timeline 722 -> 760 ms, a 51,820-token one 11.79 -> 12.11 s, the
  "gdn recurrence" phase 34 -> 37 ms, from the cross-stream scheduling gaps.  Both launches stay on the caller's
  stream.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
linxuhaoandClaude Opus 5.5 committed 2026-10-22 20:52:26 +00:00
1 parent 579cbfbbeb
commit f1cabc992f
3 files changed
+1026 -23
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/prefill/gdn_rec_parity.cu
+394 -22
View File
@@ -8,10 +8,15 @@
// output, the state after it and the FP16 output must be the same bits. STRATA_GDN_CHUNKED=1's recurrence in chunks
// (copied below too) sums in another order: it must stay within 1e-4 of gdn_rec_kh_kernel (largest |difference| over
// largest |value|, the output and the state), also at T = 8192 and 32768; up to 8192 both are also measured against
// the recurrence in FP64.
// the recurrence in FP64. STRATA_GDN_TC=1's tensor-core scan (gdn_tc.cuh, included below) is held to the same 1e-4,
// also on stress inputs (the state 100x smaller, 50x larger or zero, v 50x larger, slow gates, coherent keys).
//
// Synthetic inputs of the model's shape (q/k L2-normalized per key head as gdn_l2_kernel leaves them), no model.
// gdn_rec_parity the bit checks (T from 1 to 4099, a non-zero state as after an earlier chunk)
// gdn_rec_parity --long and the tensor-core scan at T = 51820 (needs ~7 GB)
// gdn_rec_parity --selftest-frag the tensor-core scan's fragment helpers alone, one warp each, against FP64
// gdn_rec_parity --bench-tc the tensor-core series alone (gdn_rec_kh_kernel and the scan alternating, nothing else:
// for A / B runs of two builds back to back)
// gdn_rec_parity --bench and the timings at T = 2048, 8192 and 32768 (one layer, all 48 value heads); the two
// variants' launches alternate, since a card under load steps its clock down after a
// fraction of a second (see native_grouped_parity); then both at 1 to 6 blocks per
@@ -590,6 +595,41 @@ cudaError_t gdn_rec_chunked(float* state, const float* h, const float* gate, con
}
return cudaSuccess; // launch errors surface at the caller's check; an error above means nothing ran
}
// src/prefill/kernels.cu's tensor-core scan (STRATA_GDN_TC=1): the same text, included there too
#include "gdn_tc.cuh"
// src/prefill/kernels.cu's gdn_rec_tc, copied as it is there (with the raw device attributes, as gdn_rec_chunked above)
cudaError_t gdn_rec_tc(float* state, const float* h, const float* gate, const float* beta, float* y, int64_t T,
cudaStream_t s) {
if (T <= 0) return cudaSuccess;
{
int dev = 0, major = 0, sms = 0, smem = 0;
if (cudaGetDevice(&dev) != cudaSuccess ||
cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) != cudaSuccess ||
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev) != cudaSuccess ||
cudaDeviceGetAttribute(&smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev) != cudaSuccess) {
cudaGetLastError();
return cudaErrorNotSupported;
}
if (major < 8 || sms < HV || (size_t) smem < kTcScanSmem<false>) return cudaErrorNotSupported;
}
const GtcScratch sc = gdn_tc_scratch();
if (sc.ti == nullptr) return cudaErrorMemoryAllocation;
{ // once a device
static bool attrs[64] = {};
int dev = 0;
cudaGetDevice(&dev);
if (dev >= 0 && dev < 64 && !attrs[dev]) {
cudaFuncSetAttribute(gdn_tc_scan_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) kTcScanSmem<false>);
attrs[dev] = true;
}
}
for (int64_t t0 = 0; t0 < T; t0 += GSB) {
const int nch = (int) ((std::min<int64_t>(GSB, T - t0) + GCH - 1) / GCH);
gdn_tc_prep_kernel<<<dim3((unsigned) nch, HK), 256, kTcPrepSmem, s>>>(h, gate, beta, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, t0, T);
gdn_tc_scan_kernel<false><<<HV, GTC_CW * 32, kTcScanSmem<false>, s>>>(state, h, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, y, t0, nch, T);
}
return cudaSuccess;
}
// the output norm over a head's 128 columns, into the FP16 copy the out projection reads (the FP32 output before the
// norm stays in its scratch buffer)
__global__ void __launch_bounds__(S) gdn_out_norm_kernel(const float* __restrict__ z, const float* __restrict__ gamma,
@@ -650,14 +690,24 @@ __global__ void l2_rows(float* h, float eps) {
__syncthreads();
x[threadIdx.x] = v * rsqrtf(part[0] + part[1] + part[2] + part[3] + eps);
}
// gate = log of the decay (mostly near 1, some strong), beta in (0, 1)
__global__ void fill_gates(float* gate, float* beta, int64_t n, uint64_t seed) {
// gate = log of the decay (mostly near 1, some strong; slow: -0.001 u, so the state grows for long), beta in (0, 1)
__global__ void fill_gates(float* gate, float* beta, int64_t n, uint64_t seed, bool slow = false) {
for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x) {
const float u = uni(seed * 0x9E3779B97F4A7C15ULL + 7 * (uint64_t) i);
gate[i] = -0.02f - 0.8f * u * u * u;
gate[i] = slow ? -0.001f * u : -0.02f - 0.8f * u * u * u;
beta[i] = 1.0f / (1.0f + expf(-gauss(seed * 0xD1B54A32D192ED03ULL + (uint64_t) i)));
}
}
// coherent keys: every token's k row = a key head's k_0 + 0.05 x the row that was there (before l2_rows): the
// chunk's T = (I + A)^-1 then carries large entries, the reason the prep scales its T images by a power of two of their own
__global__ void coherent_k(float* h, uint64_t seed) {
float* x = h + (int64_t) blockIdx.y * C + HK * S + (int64_t) blockIdx.x * S;
x[threadIdx.x] = gauss(seed * 0x9E3779B97F4A7C15ULL + (uint64_t) (blockIdx.x * S + threadIdx.x)) + 0.05f * x[threadIdx.x];
}
__global__ void scale_v(float* h, int64_t T, float scale) {
for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < T * HV * S; i += (int64_t) gridDim.x * blockDim.x)
h[(i / (HV * S)) * C + 2 * HK * S + i % (HV * S)] *= scale;
}
__global__ void count_diff(const uint32_t* a, const uint32_t* b, int64_t n, unsigned long long* out) {
unsigned long long d = 0;
for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x)
@@ -738,10 +788,15 @@ unsigned long long diff_words(const void* a, const void* b, size_t bytes) {
return h;
}
// The model's shapes with gaussian values; the defaults are the bit checks' inputs (the state as after an earlier
// chunk, scale 0.05). The others stress the tensor-core scan's f16 range handling: the state 100x smaller, 50x larger
// or zero (a fresh conversation), v 50x larger, slow gates, coherent keys.
struct Inputs {
int64_t T = 0;
float *h = nullptr, *gate = nullptr, *beta = nullptr, *z = nullptr, *gamma = nullptr, *state0 = nullptr;
explicit Inputs(int64_t t, uint64_t seed) : T(t) {
explicit Inputs(int64_t t, uint64_t seed, float state_scale = 0.05f, float v_scale = 1.0f, bool slow_gates = false,
bool coherent_keys = false)
: T(t) {
h = dalloc<float>((size_t) T * C);
gate = dalloc<float>((size_t) T * HV);
beta = dalloc<float>((size_t) T * HV);
@@ -749,11 +804,13 @@ struct Inputs {
gamma = dalloc<float>(S);
state0 = dalloc<float>((size_t) S * HV * S);
fill_gauss<<<4096, 256>>>(h, T * C, 1.0f, seed);
if (coherent_keys) coherent_k<<<dim3(HK, (unsigned) T), S>>>(h, seed + 5);
l2_rows<<<dim3(2 * HK, (unsigned) T), S>>>(h, 1e-6f);
fill_gates<<<1024, 256>>>(gate, beta, T * HV, seed + 1);
if (v_scale != 1.0f) scale_v<<<4096, 256>>>(h, T, v_scale);
fill_gates<<<1024, 256>>>(gate, beta, T * HV, seed + 1, slow_gates);
fill_gauss<<<4096, 256>>>(z, T * HV * S, 1.0f, seed + 2);
fill_gauss<<<1, 128>>>(gamma, S, 0.1f, seed + 3, 1.0f);
fill_gauss<<<1024, 256>>>(state0, (int64_t) S * HV * S, 0.05f, seed + 4);
fill_gauss<<<1024, 256>>>(state0, (int64_t) S * HV * S, state_scale, seed + 4);
ck(cudaDeviceSynchronize(), "inputs");
}
Inputs(const Inputs&) = delete;
@@ -858,8 +915,10 @@ double max_rel_f64(const float* a, const double* ref, size_t n) { // max |a -
}
return mv > 0.0 ? md / mv : md;
}
int check_chunked(int64_t T, uint64_t seed) {
Inputs in(T, seed);
// a recurrence of gdn_rec_chunked's shape (the FP32 chunked pair, the tensor-core scan) on the inputs given
using RecLaunch = cudaError_t (*)(float*, const float*, const float*, const float*, float*, int64_t, cudaStream_t);
int check_variant(const Inputs& in, const char* label, RecLaunch launch, const char* needs) {
const int64_t T = in.T;
const size_t st = (size_t) S * HV * S, oc_n = (size_t) T * HV * S;
float* s0 = dalloc<float>(st);
float* s1 = dalloc<float>(st);
@@ -868,38 +927,307 @@ int check_chunked(int64_t T, uint64_t seed) {
ck(cudaMemcpy(s0, in.state0, st * 4, cudaMemcpyDeviceToDevice), "state copy");
ck(cudaMemcpy(s1, in.state0, st * 4, cudaMemcpyDeviceToDevice), "state copy");
gdn_rec_kh_kernel<<<HK * NCB, dim3(CB, RG)>>>(s0, in.h, in.gate, in.beta, o0, in.T);
const cudaError_t e = gdn_rec_chunked(s1, in.h, in.gate, in.beta, o1, in.T, 0);
const cudaError_t e = launch(s1, in.h, in.gate, in.beta, o1, in.T, 0);
if (e == cudaErrorNotSupported) {
std::printf(" T=%-6lld chunked: not on this card (sm_80+, 128 SMs, 93.7 KB a block)\n", (long long) T);
std::printf(" T=%-6lld %s: not on this card (%s)\n", (long long) T, label, needs);
cudaFree(s0); cudaFree(s1); cudaFree(o0); cudaFree(o1);
return 0;
}
ck(e, "chunked");
ck(cudaDeviceSynchronize(), "chunked run");
ck(e, label);
ck(cudaDeviceSynchronize(), label);
const double eo = max_rel(o0, o1, oc_n), es = max_rel(s0, s1, st);
const bool bad = !(eo < 1e-4) || !(es < 1e-4);
std::printf("%s T=%-6lld chunked (STRATA_GDN_CHUNKED): output %.2e, state %.2e (max |diff| / max |value|)\n",
bad ? "FAIL" : " ", (long long) T, eo, es);
std::printf("%s T=%-6lld %s: output %.2e, state %.2e (max |diff| / max |value|)\n", bad ? "FAIL" : " ",
(long long) T, label, eo, es);
if (T <= 8192) { // both against FP64 (its output: 8 bytes a value, so not at 32K)
double* sd = dalloc<double>(st);
double* od = dalloc<double>(oc_n);
to_f64<<<1024, 256>>>(in.state0, sd, (int64_t) st);
gdn_rec_f64_kernel<<<HV, S>>>(sd, in.h, in.gate, in.beta, od, in.T);
ck(cudaDeviceSynchronize(), "FP64 run");
std::printf(" against FP64: gdn_rec_kh_kernel output %.2e, state %.2e; chunked output %.2e, state %.2e\n",
std::printf(" against FP64: gdn_rec_kh_kernel output %.2e, state %.2e; this one output %.2e, state %.2e\n",
max_rel_f64(o0, od, oc_n), max_rel_f64(s0, sd, st), max_rel_f64(o1, od, oc_n), max_rel_f64(s1, sd, st));
cudaFree(sd); cudaFree(od);
}
cudaFree(s0); cudaFree(s1); cudaFree(o0); cudaFree(o1);
return bad ? 1 : 0;
}
int check_chunked(int64_t T, uint64_t seed) {
Inputs in(T, seed);
return check_variant(in, "chunked (STRATA_GDN_CHUNKED)", gdn_rec_chunked, "sm_80+, 128 SMs, 93.7 KB a block");
}
constexpr const char* kTcNeeds = "sm_80+, 48 SMs, 122 KB a block";
int check_tc(int64_t T, uint64_t seed) {
Inputs in(T, seed);
return check_variant(in, "tc (STRATA_GDN_TC=1)", gdn_rec_tc, kTcNeeds);
}
// the tensor-core scan on the stress inputs (Inputs): its f16 range handling has a case for each
int check_tc_stress(int64_t T, uint64_t seed) {
struct Case { const char* label; float state_scale, v_scale; bool slow_gates, coherent_keys; };
const Case cases[] = {{"tc (STRATA_GDN_TC=1), state0 x 0.01", 5e-4f, 1.0f, false, false},
{"tc (STRATA_GDN_TC=1), state0 x 50", 2.5f, 1.0f, false, false},
{"tc (STRATA_GDN_TC=1), state0 = 0", 0.0f, 1.0f, false, false},
{"tc (STRATA_GDN_TC=1), v x 50", 0.05f, 50.0f, false, false},
{"tc (STRATA_GDN_TC=1), slow gates", 0.05f, 1.0f, true, false},
{"tc (STRATA_GDN_TC=1), coherent keys", 0.05f, 1.0f, false, true}};
int fails = 0;
for (const Case& c : cases) {
Inputs in(T, seed++, c.state_scale, c.v_scale, c.slow_gates, c.coherent_keys);
fails += check_variant(in, c.label, gdn_rec_tc, kTcNeeds);
}
return fails;
}
// ------------------------------------------------------------------ the tensor-core scan's fragment self-test (--selftest-frag)
// One warp each, random inputs, the reference in double on the host: (b) a 16x128 . 128x32 product through
// gtc_prod_state (the f16 split with the state's 2^es and k's 2^8, ldmatrix), (c) a 16x32 . 32x128 update through
// gtc_update_state (ldmatrix.trans, the 2^ed scale, the accumulator's pre / post scale), (d) a 16x32 . 32x32 product
// through gtc_prod_dt on a full random T of wide range (the prep's 2^et on the B images), (e) the triangle-skipped
// tile sets of gtc_prod_dt against the full ones on triangular T / P: the same bits, since the skipped tiles add
// exact zeros. Most bugs of a tensor-core kernel are fragment-layout bugs: these catch them before the kernel does.
// ((a), the C -> A tf32 shuffle, went with the tf32 form of D^T.)
__device__ void st_stage_k(uint16_t* kh, uint16_t* kl, const float* k) { // k: 32 x 128 f32 -> the hi / lo images x 2^8
for (int p = threadIdx.x; p < GCH * S / 4; p += 32) {
const int t = p / (S / 4), r = (p % (S / 4)) * 4;
gtc_split_store4(kh + t * GTC_KP + r, kl + t * GTC_KP + r, *reinterpret_cast<const float4*>(k + t * S + r), 256.0f);
}
__syncwarp();
}
// a 16-row matrix (row stride ld) as 8-column tiles in the C layout: tile j, lane (gid, tig) holds (gid, 8j + 2tig),
// (gid, 8j + 2tig + 1), (gid + 8, 8j + 2tig), (gid + 8, 8j + 2tig + 1)
template <int NT> __device__ void st_load_c(float (&c)[NT][4], const float* m, int ld) {
const int lane = threadIdx.x & 31, gid = lane >> 2, tig = lane & 3;
#pragma unroll
for (int j = 0; j < NT; ++j) {
c[j][0] = m[gid * ld + 8 * j + 2 * tig];
c[j][1] = m[gid * ld + 8 * j + 2 * tig + 1];
c[j][2] = m[(gid + 8) * ld + 8 * j + 2 * tig];
c[j][3] = m[(gid + 8) * ld + 8 * j + 2 * tig + 1];
}
}
template <int NT> __device__ void st_store_c(const float (&c)[NT][4], float* m, int ld, float scale) {
const int lane = threadIdx.x & 31, gid = lane >> 2, tig = lane & 3;
#pragma unroll
for (int j = 0; j < NT; ++j) {
m[gid * ld + 8 * j + 2 * tig] = c[j][0] * scale;
m[gid * ld + 8 * j + 2 * tig + 1] = c[j][1] * scale;
m[(gid + 8) * ld + 8 * j + 2 * tig] = c[j][2] * scale;
m[(gid + 8) * ld + 8 * j + 2 * tig + 1] = c[j][3] * scale;
}
}
template <int NT> __device__ float st_max_abs(const float (&c)[NT][4]) {
float m = 0.0f;
#pragma unroll
for (int j = 0; j < NT; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) m = fmaxf(m, fabsf(c[j][q]));
return gtc_warp_max(m);
}
__global__ void st_state_kernel(const float* a, const float* k, float* out) { // out (16 x 32) = a (16 x 128) k^T
__shared__ __align__(16) uint16_t kh[GCH * GTC_KP], kl[GCH * GTC_KP];
st_stage_k(kh, kl, k);
float s[16][4];
st_load_c<16>(s, a, S);
const int es = gtc_scale_exp(st_max_abs<16>(s));
float xt[4][4] = {}, qs[4][4] = {};
gtc_prod_state(s, gtc_pow2(es), kh, kl, kh, kl, xt, qs);
st_store_c<4>(xt, out, GCH, gtc_pow2(-(8 + es)));
}
__global__ void st_update_kernel(const float* s0, const float* dd, const float* k, float eL, float* out) { // out (16 x 128) = eL s0 + dd (16 x 32) k
__shared__ __align__(16) uint16_t kh[GCH * GTC_KP], kl[GCH * GTC_KP];
st_stage_k(kh, kl, k);
float s[16][4], dt[4][4], ddt[4][4];
st_load_c<16>(s, s0, S);
st_load_c<4>(dt, dd, GCH);
const int ed = gtc_scale_exp(st_max_abs<4>(dt));
const float sD = gtc_pow2(ed), sPre = eL * gtc_pow2(ed + 8);
#pragma unroll
for (int j = 0; j < 4; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) ddt[j][q] = dt[j][q] * sD;
#pragma unroll
for (int j = 0; j < 16; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) s[j][q] *= sPre;
gtc_update_state(s, ddt, kh, kl);
st_store_c<16>(s, out, S, gtc_pow2(-(ed + 8)));
}
// out (16 x 32) = d p^T: p's images x 2^eb (the scan's P carry 2^8, its T the prep's 2^et), d x 2^ed from its maximum
template <bool SKIP> __global__ void st_dt_kernel(const float* d, const float* p, int eb, float* out) {
__shared__ __align__(16) uint16_t ph[GCH * GTC_PP], pl[GCH * GTC_PP];
for (int q = threadIdx.x; q < GCH * GCH / 4; q += 32) {
const int t = q / (GCH / 4), i = (q % (GCH / 4)) * 4;
gtc_split_store4(ph + t * GTC_PP + i, pl + t * GTC_PP + i, *reinterpret_cast<const float4*>(p + t * GCH + i), gtc_pow2(eb));
}
__syncwarp();
float dt[4][4], ot[4][4] = {};
st_load_c<4>(dt, d, GCH);
const int ed = gtc_scale_exp(st_max_abs<4>(dt));
gtc_prod_dt<SKIP>(dt, gtc_pow2(ed), ph, pl, ot);
st_store_c<4>(ot, out, GCH, gtc_pow2(-(ed + eb)));
}
int st_scale_exp(const std::vector<float>& m) { // the prep's gtc_scale_exp on the host: max |m| into [2^14, 2^15)
float mx = 0.0f;
for (float x : m) mx = std::max(mx, std::fabs(x));
if (mx == 0.0f) return 0;
int e = 0;
std::frexp(mx, &e); // mx = f 2^e, f in [0.5, 1): floor(log2 mx) = e - 1
e = 14 - (e - 1);
return e < -40 ? -40 : (e > 40 ? 40 : e);
}
struct StRng { // xorshift64, host side
uint64_t s;
explicit StRng(uint64_t seed) : s(seed) {}
float uni() { s ^= s << 13; s ^= s >> 7; s ^= s << 17; return ((float) (s >> 40) + 0.5f) / 16777216.0f; }
float gauss() { return std::sqrt(-2.0f * std::log(uni())) * std::cos(6.2831853f * uni()); }
float logmag(float lo, float hi) { return (uni() < 0.5f ? -1.0f : 1.0f) * std::pow(10.0f, lo + (hi - lo) * uni()); }
};
float* st_upload(const std::vector<float>& v) {
float* d = dalloc<float>(v.size());
ck(cudaMemcpy(d, v.data(), v.size() * 4, cudaMemcpyHostToDevice), "upload");
return d;
}
std::vector<float> st_download(const float* d, size_t n) {
std::vector<float> v(n);
ck(cudaDeviceSynchronize(), "self-test kernel");
ck(cudaMemcpy(v.data(), d, n * 4, cudaMemcpyDeviceToHost), "download");
return v;
}
double st_max_rel(const std::vector<float>& got, const std::vector<double>& ref) {
double md = 0.0, mv = 0.0;
for (size_t i = 0; i < ref.size(); ++i) {
md = std::max(md, std::fabs((double) got[i] - ref[i]));
mv = std::max(mv, std::fabs(ref[i]));
}
return mv > 0.0 ? md / mv : md;
}
int selftest_frag() {
std::printf("--selftest-frag: the tensor-core scan's fragment helpers, one warp each\n");
int fails = 0;
StRng rng(0x5e1f7e57ULL);
auto unit_rows = [&](std::vector<float>& k, int rows, int cols) { // gaussian rows, L2-normalized
for (int t = 0; t < rows; ++t) {
double ss = 0.0;
for (int r = 0; r < cols; ++r) { k[t * cols + r] = rng.gauss(); ss += (double) k[t * cols + r] * k[t * cols + r]; }
for (int r = 0; r < cols; ++r) k[t * cols + r] = (float) (k[t * cols + r] / std::sqrt(ss));
}
};
{ // (b) X^T = S^T K^T: f16x3 with the state's 2^es and k's 2^8, B by ldmatrix
std::vector<float> a(16 * S), k(GCH * S);
for (float& x : a) x = rng.logmag(-3.0f, 2.0f);
unit_rows(k, GCH, S);
std::vector<double> ref(16 * GCH);
for (int c = 0; c < 16; ++c)
for (int i = 0; i < GCH; ++i) {
double v = 0.0;
for (int r = 0; r < S; ++r) v += (double) a[c * S + r] * k[i * S + r];
ref[c * GCH + i] = v;
}
float* da = st_upload(a);
float* dk = st_upload(k);
float* dout = dalloc<float>(16 * GCH);
st_state_kernel<<<1, 32>>>(da, dk, dout);
const double rel = st_max_rel(st_download(dout, 16 * GCH), ref);
const bool bad = !(rel <= 2e-6);
std::printf(" (b) X^T = S^T K^T (16x128 . 128x32, f16x3, 2^es 2^8, ldmatrix): max |diff| / max |ref| %.2e %s\n", rel, bad ? "FAIL" : "ok");
fails += bad;
cudaFree(da); cudaFree(dk); cudaFree(dout);
}
{ // (c) S^T = eL S^T + Dd^T K: ldmatrix.trans, 2^ed, the accumulator scaled before and after
std::vector<float> s0(16 * S), dd(16 * GCH), k(GCH * S);
for (float& x : s0) x = rng.logmag(-3.0f, 2.0f);
for (float& x : dd) x = rng.logmag(-3.0f, 2.0f);
unit_rows(k, GCH, S);
const float eL = 0.25f + 0.75f * rng.uni();
std::vector<double> ref(16 * S);
for (int c = 0; c < 16; ++c)
for (int r = 0; r < S; ++r) {
double v = (double) eL * s0[c * S + r];
for (int i = 0; i < GCH; ++i) v += (double) dd[c * GCH + i] * k[i * S + r];
ref[c * S + r] = v;
}
float* ds = st_upload(s0);
float* dd_ = st_upload(dd);
float* dk = st_upload(k);
float* dout = dalloc<float>(16 * S);
st_update_kernel<<<1, 32>>>(ds, dd_, dk, eL, dout);
const double rel = st_max_rel(st_download(dout, 16 * S), ref);
const bool bad = !(rel <= 2e-6);
std::printf(" (c) S^T = eL S^T + Dd^T K (16x32 . 32x128, f16x3, ldmatrix.trans, 2^ed pre / post scale): %.2e %s\n", rel, bad ? "FAIL" : "ok");
fails += bad;
cudaFree(ds); cudaFree(dd_); cudaFree(dk); cudaFree(dout);
}
{ // (d) D^T = Y^T T^T on f16x3 through gtc_prod_dt, a full random T of wide range with the prep's 2^et
std::vector<float> y(16 * GCH), t(GCH * GCH);
for (float& x : y) x = rng.logmag(-3.0f, 2.0f);
for (float& x : t) x = rng.logmag(-2.0f, 3.0f);
std::vector<double> ref(16 * GCH);
for (int c = 0; c < 16; ++c)
for (int i = 0; i < GCH; ++i) {
double v = 0.0;
for (int ip = 0; ip < GCH; ++ip) v += (double) y[c * GCH + ip] * t[i * GCH + ip];
ref[c * GCH + i] = v;
}
float* dy = st_upload(y);
float* dt_ = st_upload(t);
float* dout = dalloc<float>(16 * GCH);
st_dt_kernel<false><<<1, 32>>>(dy, dt_, st_scale_exp(t), dout);
const double rel = st_max_rel(st_download(dout, 16 * GCH), ref);
const bool bad = !(rel <= 2e-6);
std::printf(" (d) D^T = Y^T T^T (16x32 . 32x32, f16x3, Y x 2^ey, T x 2^et): %.2e %s\n", rel, bad ? "FAIL" : "ok");
fails += bad;
cudaFree(dy); cudaFree(dt_); cudaFree(dout);
}
{ // (e) the triangle-skipped tile sets against the full ones: T unit lower triangular, P lower triangular
std::vector<float> y(16 * GCH), t(GCH * GCH), d(16 * GCH), p(GCH * GCH);
for (float& x : y) x = rng.logmag(-3.0f, 2.0f);
for (float& x : d) x = rng.logmag(-3.0f, 2.0f);
for (int i = 0; i < GCH; ++i)
for (int j = 0; j < GCH; ++j) {
t[i * GCH + j] = j < i ? rng.logmag(-2.0f, 3.0f) : (j == i ? 1.0f : 0.0f);
p[i * GCH + j] = j <= i ? 2.0f * rng.uni() - 1.0f : 0.0f;
}
std::vector<double> ref_d(16 * GCH), ref_o(16 * GCH);
for (int c = 0; c < 16; ++c)
for (int i = 0; i < GCH; ++i) {
double v = 0.0, w = 0.0;
for (int ip = 0; ip < GCH; ++ip) { v += (double) y[c * GCH + ip] * t[i * GCH + ip]; w += (double) d[c * GCH + ip] * p[i * GCH + ip]; }
ref_d[c * GCH + i] = v;
ref_o[c * GCH + i] = w;
}
float* dy = st_upload(y);
float* dt_ = st_upload(t);
float* dd_ = st_upload(d);
float* dp = st_upload(p);
float* full = dalloc<float>(16 * GCH);
float* skip = dalloc<float>(16 * GCH);
st_dt_kernel<false><<<1, 32>>>(dy, dt_, st_scale_exp(t), full);
st_dt_kernel<true><<<1, 32>>>(dy, dt_, st_scale_exp(t), skip);
const std::vector<float> f1 = st_download(full, 16 * GCH), s1 = st_download(skip, 16 * GCH);
const unsigned long long diff1 = std::memcmp(f1.data(), s1.data(), 16 * GCH * 4) != 0 ? diff_words(full, skip, 16 * GCH * 4) : 0;
const double rel1 = st_max_rel(s1, ref_d);
st_dt_kernel<false><<<1, 32>>>(dd_, dp, 8, full);
st_dt_kernel<true><<<1, 32>>>(dd_, dp, 8, skip);
const std::vector<float> f2 = st_download(full, 16 * GCH), s2 = st_download(skip, 16 * GCH);
const unsigned long long diff2 = std::memcmp(f2.data(), s2.data(), 16 * GCH * 4) != 0 ? diff_words(full, skip, 16 * GCH * 4) : 0;
const double rel2 = st_max_rel(s2, ref_o);
const bool bad = diff1 != 0 || diff2 != 0 || !(rel1 <= 2e-6) || !(rel2 <= 2e-6);
std::printf(" (e) triangle skipping: D^T 10 of 16 tiles %llu words differ (vs FP64 %.2e), o^T 6 of 8 tiles %llu words differ (vs FP64 %.2e) %s\n",
diff1, rel1, diff2, rel2, bad ? "FAIL" : "ok");
fails += bad;
cudaFree(dy); cudaFree(dt_); cudaFree(dd_); cudaFree(dp); cudaFree(full); cudaFree(skip);
}
std::printf("--selftest-frag: %d failures\n", fails);
return fails;
}
float median(std::vector<float> v) {
std::sort(v.begin(), v.end());
return v.empty() ? 0.0f : v[v.size() / 2];
}
void bench(int64_t T) {
// tc_only (--bench-tc): the tensor-core series alone, for A / B runs of two builds back to back
void bench(int64_t T, bool tc_only = false) {
Inputs in(T, 0xB3);
const size_t st = (size_t) S * HV * S, oc_n = (size_t) T * HV * S;
float* state = dalloc<float>(st);
@@ -920,14 +1248,14 @@ void bench(int64_t T) {
};
const int warm = 3, rounds = T >= 16384 ? 10 : 20;
std::vector<float> t_rec[kRecCount], t_norm[2];
for (int i = 0; i < warm + rounds; ++i) {
for (int i = 0; !tc_only && i < warm + rounds; ++i) {
for (int j = 0; j < kRecCount; ++j) {
const Rec v = (Rec) ((i + j) % kRecCount); // a different variant first in each round
const float ms = timed([&] { run_rec(v, state, in, oc); });
if (i >= warm) t_rec[v].push_back(ms);
}
}
for (int i = 0; i < warm + rounds; ++i) {
for (int i = 0; !tc_only && i < warm + rounds; ++i) {
for (int j = 0; j < 2; ++j) {
const int v = (i + j) % 2;
const float ms = timed([&] {
@@ -938,7 +1266,7 @@ void bench(int64_t T) {
}
}
std::vector<float> t_kc[2]; // gdn_rec_kh_kernel and the chunked recurrence, alternating
const bool can_chunk = gdn_rec_chunked(state, in.h, in.gate, in.beta, oc, in.T, 0) == cudaSuccess &&
const bool can_chunk = !tc_only && gdn_rec_chunked(state, in.h, in.gate, in.beta, oc, in.T, 0) == cudaSuccess &&
cudaDeviceSynchronize() == cudaSuccess;
for (int i = 0; can_chunk && i < warm + rounds; ++i)
for (int j = 0; j < 2; ++j) {
@@ -949,8 +1277,30 @@ void bench(int64_t T) {
});
if (i >= warm) t_kc[v].push_back(ms);
}
std::vector<float> t_tc[2]; // gdn_rec_kh_kernel and the tensor-core scan, alternating
const bool can_tc = gdn_rec_tc(state, in.h, in.gate, in.beta, oc, in.T, 0) == cudaSuccess &&
cudaDeviceSynchronize() == cudaSuccess;
for (int i = 0; can_tc && i < warm + rounds; ++i)
for (int j = 0; j < 2; ++j) {
const int v = (i + j) % 2;
const float ms = timed([&] {
if (v == 0) gdn_rec_kh_kernel<<<HK * NCB, dim3(CB, RG)>>>(state, in.h, in.gate, in.beta, oc, in.T);
else ck(gdn_rec_tc(state, in.h, in.gate, in.beta, oc, in.T, 0), "tc");
});
if (i >= warm) t_tc[v].push_back(ms);
}
const float r0 = median(t_rec[kBefore]), r1 = median(t_rec[kKeyHead]), n0 = median(t_norm[0]), n1 = median(t_norm[1]);
std::printf(" T = %lld tokens, one layer (48 value heads), medians of %d alternating runs:\n", (long long) T, rounds);
if (tc_only) {
const float kh = median(t_tc[0]), tc = median(t_tc[1]);
std::printf(" recurrence, gdn_rec_kh_kernel %8.2f ms (%.3f us per token)\n", kh, 1000.0f * kh / (float) T);
std::printf(" recurrence, tensor-core chunked (STRATA_GDN_TC=1) %8.2f ms (%.3f us per token) %.2fx gdn_rec_kh_kernel\n",
tc, 1000.0f * tc / (float) T, kh / tc);
cudaEventDestroy(e0);
cudaEventDestroy(e1);
cudaFree(state); cudaFree(oc); cudaFree(y16);
return;
}
std::printf(" recurrence, before (gdn_rec_cols_pipe_kernel) %8.2f ms (%.3f us per token)\n", r0,
1000.0f * r0 / (float) T);
std::printf(" recurrence, gdn_rec_kh_kernel %8.2f ms (%.3f us per token) %.2fx\n", r1,
@@ -960,6 +1310,11 @@ void bench(int64_t T) {
std::printf(" recurrence, chunked (STRATA_GDN_CHUNKED=1) %8.2f ms (%.3f us per token) %.2fx gdn_rec_kh_kernel\n",
kc, 1000.0f * kc / (float) T, kh / kc);
}
if (can_tc) {
const float kh = median(t_tc[0]), tc = median(t_tc[1]);
std::printf(" recurrence, tensor-core chunked (STRATA_GDN_TC=1) %8.2f ms (%.3f us per token) %.2fx gdn_rec_kh_kernel\n",
tc, 1000.0f * tc / (float) T, kh / tc);
}
std::printf(" output norm, before (FP32 and FP16 store) %8.2f ms\n", n0);
std::printf(" output norm, FP16 store only %8.2f ms %.2fx\n", n1, n0 / n1);
std::printf(" 36 DeltaNet layers per %lld-token block: recurrence + norm %.0f ms -> %.0f ms (-%.0f ms)\n",
@@ -1357,18 +1712,35 @@ void fewer_sms(int64_t T) {
} // namespace
int main(int argc, char** argv) {
bool do_bench = false;
for (int i = 1; i < argc; ++i)
bool do_bench = false, do_self = false, do_long = false, do_bench_tc = false;
for (int i = 1; i < argc; ++i) {
if (std::strcmp(argv[i], "--bench") == 0) do_bench = true;
if (std::strcmp(argv[i], "--selftest-frag") == 0) do_self = true;
if (std::strcmp(argv[i], "--long") == 0) do_long = true;
if (std::strcmp(argv[i], "--bench-tc") == 0) do_bench_tc = true;
}
cudaDeviceProp p{};
ck(cudaGetDeviceProperties(&p, 0), "device");
std::printf("gdn_rec_parity on %s (sm_%d%d)\n", p.name, p.major, p.minor);
int fails = 0;
if (do_self) { // the tensor-core helpers alone
fails = selftest_frag();
std::printf("gdn_rec_parity: %d failures\n", fails);
return fails ? 1 : 0;
}
if (do_bench_tc) { // the tensor-core timings alone
for (const int64_t T : {(int64_t) 2048, (int64_t) 8192, (int64_t) 32768}) bench(T, true);
return 0;
}
const int64_t Ts[] = {1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 33, 257, 1000, 4099};
uint64_t seed = 11;
for (const int64_t T : Ts) fails += check_bits(T, seed++);
for (const int64_t T : Ts) fails += check_chunked(T, seed++);
for (const int64_t T : {(int64_t) 8192, (int64_t) 32768}) fails += check_chunked(T, seed++);
for (const int64_t T : Ts) fails += check_tc(T, seed++);
for (const int64_t T : {(int64_t) 8192, (int64_t) 32768}) fails += check_tc(T, seed++);
for (const int64_t T : {(int64_t) 4099, (int64_t) 32768}) { fails += check_tc_stress(T, seed); seed += 8; }
if (do_long) fails += check_tc(51820, seed++); // the engine's longest prompt (about 7 GB of inputs and outputs)
if (do_bench) {
std::printf("--bench: the prompt path's DeltaNet recurrence\n");
bench(2048);
src/prefill/gdn_tc.cuh
+586
View File
@@ -0,0 +1,586 @@
// src/prefill/gdn_tc.cuh - the chunked DeltaNet recurrence's scan on the tensor cores (STRATA_GDN_TC=1), sm_80+.
//
// Included inside the anonymous namespace of src/prefill/kernels.cu and of src/prefill/gdn_rec_parity.cu, after
// gdn_chunk_scan_kernel (the harness runs the engine's kernels verbatim: one text, included twice). It uses S, HK, HV,
// C, VPK, GCH, GSB, GTP, gdn_cp16 / gdn_cp_commit / gdn_cp_wait_all from there and reads no environment.
//
// gdn_chunk_scan_kernel's chunk algebra (kernels.cu, "STRATA_GDN_CHUNKED=1"), every product transposed so that a
// warp's slice of the state is the mma A operand and never leaves its registers: a block per value head, warp w owns
// value columns 16w..16w+15 of S^T (16 columns x 128 rows = 16 m16n8 f32 accumulator tiles, 64 registers a lane) and
// all 32 tokens of a chunk, so not one product couples two warps:
// X^T (16c x 32i) = S0^T . K^T f16x3, k = state row (8 k16-steps), n = token (4 tiles)
// QS^T (16c x 32t) = S0^T . Q^T the same A fragments
// Y^T (16c x 32i) = beta_i (V^T - e_i X^T) elementwise on the accumulators
// D^T (16c x 32i) = Y^T . T^T f16x3, k = token i' (2 steps), n = token i: T is zero above its diagonal,
// so 6 of 8 tiles (as P below)
// o^T (16c x 32t) = e_t QS^T + D^T . P^T f16x3, k = i (2 steps), n = t: P zero above its diagonal (6 of 8 tiles)
// Dd^T (16c x 32i) = f_i D^T, f_i = exp(gL - gamma_i)
// S^T (16c x 128r) = eL S0^T + Dd^T . K f16x3, k = token (2 steps), n = state row (16 tiles)
// 324 mma a warp a chunk; the only shared data are the staged inputs (k, q, T and P as f16 hi / lo images, v, the
// per-token scalars), one barrier a chunk.
//
// Numerics (the harness holds the result within 1e-4 of gdn_rec_kh_kernel, as the FP32 chunked pair): an f16 operand
// is split x = hi + lo, hi = f16(x), lo = f16(x - hi), and a product takes hi.hi + hi.lo + lo.hi, dropping lo.lo: with
// both parts normal that is within 3 x 2^-22 |a||b| of the exact product, the 3xTF32 bound, at half the tensor time.
// For the low part to be normal the operand's largest entry must sit in [2^14, 2^15), which a power of two does
// exactly: k and q (unit vectors after gdn_l2_kernel) carry 2^8, P (|P| <= 1) 2^8, the state 2^es from a warp max
// per chunk, Y^T 2^ey and D^T 2^ed likewise, T = (I + A)^-1 (no a-priori bound) 2^et from the prep's warp max; the
// powers are taken out of the per-token scalars and of the accumulator before and after the mma (one FMUL an
// element, exact). (D^T = Y^T T^T on tf32x3 with T split at use, the form before this one, cost 30 mma, 32 shuffles
// and 48 tf32 conversions a warp a chunk against 18 mma now; measured in the commit that changed it.)
#if !defined(__HIPCC__)
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
#define STRATA_GDN_TC_SM80 1
#else
#define STRATA_GDN_TC_SM80 0 // device code for older cards: the wrappers trap (gdn_rec_tc never launches there)
#endif
constexpr int GTC_KP = S + 8; // f16 k / q rows: a 272-B stride puts the 8 rows of a ldmatrix 8x8 on 8 distinct 16-B lanes
constexpr int GTC_VP = S + 4; // f32 v rows: lane (gid, tig) reads word 264 tig + gid, 32 distinct banks
constexpr int GTC_PP = GCH + 8; // f16 T / P rows: an 80-B stride, 8 rows on distinct 16-B lanes
constexpr int GTC_CW = S / 16; // consumer warps a block: 16 value columns each
// dynamic shared memory a block, two buffers each: k / q hi / lo, v, (z), T hi / lo, P hi / lo, the token scalars
// (beta, exp(gamma), exp(gL - gamma), gamma) as float4, (the norm's o^2 partials), T's scale 2^-et: 124,944 B
// without the fused norm, 160,784 with it (the 166,912 opt-in of GA100)
template <bool NORM>
constexpr size_t kTcScanSmem = 2 * (4 * GCH * GTC_KP * 2 + GCH * GTC_VP * 4 + (NORM ? GCH * GTC_VP * 4 : 0) +
4 * GCH * GTC_PP * 2 + GCH * 16 + (NORM ? GCH * GTC_CW * 4 : 0)) + 16;
static_assert(kTcScanSmem<false> == 124944 && kTcScanSmem<true> == 160784, "gdn_tc.cuh: the shared-memory plan");
// the prep's scratch for a super-block (gdn_tc_scratch, one a device, 60.3 MB): the T images x 2^et as
// [chunk][vh][2: hi, lo][GCH][GCH] f16 and their scales 2^-et as [chunk][vh], the token scalars as float4
// [chunk][vh][GCH], the k and q images x 2^8 as [chunk][qh][4: k hi, k lo, q hi, q lo][GCH][S] and the P images x 2^8
// as [chunk][vh][2: hi, lo][GCH][GCH] - every row a multiple of 16 B, so the scan stages all of it by cp.async
constexpr int GTC_NB = GSB / GCH;
struct GtcScratch {
uint16_t* ti;
float* te;
float4* sc;
uint16_t* kq;
uint16_t* pi;
};
constexpr size_t kTcTiBytes = (size_t) GTC_NB * HV * 2 * GCH * GCH * 2, kTcTeBytes = (size_t) GTC_NB * HV * 4,
kTcScBytes = (size_t) GTC_NB * HV * GCH * 16, kTcKqBytes = (size_t) GTC_NB * HK * 4 * GCH * S * 2,
kTcPiBytes = (size_t) GTC_NB * HV * 2 * GCH * GCH * 2;
static_assert(kTcTiBytes + kTcTeBytes + kTcScBytes + kTcKqBytes + kTcPiBytes == 60305408, "gdn_tc.cuh: the scratch plan");
// allocated on first use and kept, as gdn_chunk_scratch (kernels.cu); ti == nullptr when the card had no room
GtcScratch gdn_tc_scratch() {
static unsigned char* bufs[64] = {};
GtcScratch r = {nullptr, nullptr, nullptr, nullptr, nullptr};
int dev = 0;
if (cudaGetDevice(&dev) != cudaSuccess || dev < 0 || dev >= 64) { cudaGetLastError(); return r; }
if (bufs[dev] == nullptr && cudaMalloc((void**) &bufs[dev], kTcTiBytes + kTcTeBytes + kTcScBytes + kTcKqBytes + kTcPiBytes) != cudaSuccess) {
cudaGetLastError();
bufs[dev] = nullptr;
return r;
}
unsigned char* p = bufs[dev];
r.ti = reinterpret_cast<uint16_t*>(p);
r.te = reinterpret_cast<float*>(p + kTcTiBytes);
r.sc = reinterpret_cast<float4*>(p + kTcTiBytes + kTcTeBytes);
r.kq = reinterpret_cast<uint16_t*>(p + kTcTiBytes + kTcTeBytes + kTcScBytes);
r.pi = reinterpret_cast<uint16_t*>(p + kTcTiBytes + kTcTeBytes + kTcScBytes + kTcKqBytes);
return r;
}
// the warp's largest of non-negative floats: their bits order as unsigned integers, so one redux.sync (sm_80). Under
// the full mask: __activemask() does not reconverge the warp (CUDA guide), and the first call of the scan follows a
// lane-divergent branch (the L2 prefetch of tokens past the end, in the chunk two before the end whenever
// 0 < T - f_t0 < 32), so with it the lanes could reduce under a partial mask and disagree on es. The harness build
// happened to close that branch (BSYNC) before the REDUX.MAX; the engine's build is a separate compile that nothing
// checks. The kernel has no early return, so the full mask costs no reconvergence code.
__device__ __forceinline__ float gtc_warp_max(float v) {
#if STRATA_GDN_TC_SM80
return __uint_as_float(__reduce_max_sync(0xffffffffu, __float_as_uint(v)));
#else
#pragma unroll
for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o));
return v;
#endif
}
__device__ __forceinline__ float gtc_pow2(int e) { return __int_as_float((e + 127) << 23); } // 2^e, |e| < 127
// the power of two that puts m = max |x| >= 0 in [2^14, 2^15): f16 hi never overflows (65504) and the low part of
// every entry within 2^-17 of the max stays normal; clamped so that the compensating powers stay far from f32's range
__device__ __forceinline__ int gtc_scale_exp(float m) {
if (m == 0.0f) return 0;
const int e = 14 - ((__float_as_int(m) >> 23) - 127);
return e < -40 ? -40 : (e > 40 ? 40 : e);
}
// ---- the tensor-core wrappers (PTX ISA "Matrix Fragments": lane = 4 gid + tig)
// m16n8k16 f16: A a0 = (row gid, cols 2tig, 2tig+1) a1 = (gid+8, ..) a2 = (gid, 2tig+8..) a3 = (gid+8, 2tig+8..);
// B b0 = (k 2tig, 2tig+1; n gid) b1 = (k 2tig+8..; n gid); C/D c0 = (gid, 2tig) c1 = (gid, 2tig+1)
// c2 = (gid+8, 2tig) c3 = (gid+8, 2tig+1) - so two adjacent C tiles ARE the next product's A fragment
// ldmatrix m8n8.x4.b16: lanes 8m..8m+7 give the row addresses of matrix m; a lane gets (row gid, cols 2tig, 2tig+1)
// of each matrix (the lower column in the low half), or with .trans (rows 2tig, 2tig+1; col gid)
__device__ __forceinline__ void gtc_ldsm_x4(uint32_t (&r)[4], const void* smem) {
#if STRATA_GDN_TC_SM80
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
: "r"((unsigned) __cvta_generic_to_shared(smem)));
#else
r[0] = r[1] = r[2] = r[3] = 0u; (void) smem; __trap();
#endif
}
__device__ __forceinline__ void gtc_ldsm_x4_trans(uint32_t (&r)[4], const void* smem) {
#if STRATA_GDN_TC_SM80
asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
: "r"((unsigned) __cvta_generic_to_shared(smem)));
#else
r[0] = r[1] = r[2] = r[3] = 0u; (void) smem; __trap();
#endif
}
__device__ __forceinline__ void gtc_mma_f16(float (&c)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) {
#if STRATA_GDN_TC_SM80
asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
: "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1));
#else
(void) c; (void) a; (void) b0; (void) b1; __trap();
#endif
}
// the three passes of a split product, small terms first (as qsa_select.cu): lo.hi, hi.lo, hi.hi
__device__ __forceinline__ void gtc_mma3_f16(float (&c)[4], const uint32_t (&ah)[4], const uint32_t (&al)[4], uint32_t bh0,
uint32_t bh1, uint32_t bl0, uint32_t bl1) {
gtc_mma_f16(c, al, bh0, bh1);
gtc_mma_f16(c, ah, bl0, bl1);
gtc_mma_f16(c, ah, bh0, bh1);
}
// x0, x1 -> packed f16 hi (x0 in the low half) and lo = f16(x - hi): x - hi is exact in f32 (hi carries x's exponent)
__device__ __forceinline__ void gtc_split_f16x2(float x0, float x1, uint32_t& hi, uint32_t& lo) {
union { __half2 h; uint32_t u; } cv;
cv.h = __floats2half2_rn(x0, x1);
hi = cv.u;
const float2 r = __half22float2(cv.h);
cv.h = __floats2half2_rn(x0 - r.x, x1 - r.y);
lo = cv.u;
}
// four consecutive elements of an f16 hi / lo row from a float4 x scale (an 8-byte store each)
__device__ __forceinline__ void gtc_split_store4(uint16_t* hi, uint16_t* lo, float4 v, float scale) {
uint32_t h0, l0, h1, l1;
gtc_split_f16x2(v.x * scale, v.y * scale, h0, l0);
gtc_split_f16x2(v.z * scale, v.w * scale, h1, l1);
*reinterpret_cast<uint2*>(hi) = make_uint2(h0, h1);
*reinterpret_cast<uint2*>(lo) = make_uint2(l0, l1);
}
// a float4 from shared memory the compiler does not keep: the per-token scalars are read in three phases of a chunk,
// a component each, and kept across them they cost 24 registers a thread
__device__ __forceinline__ float4 gtc_lds128(const float4* p) {
float4 v;
#if STRATA_GDN_TC_SM80
asm volatile("ld.shared.v4.f32 {%0,%1,%2,%3}, [%4];"
: "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w)
: "r"((unsigned) __cvta_generic_to_shared(p)));
#else
v = *p;
#endif
return v;
}
// ---- the products of a chunk, a warp each (the self-test of gdn_rec_parity --selftest-frag runs them one by one)
// X^T = S^T K^T and QS^T = S^T Q^T: A = the state fragments x sA (= 2^es) split in f16 hi / lo just in time per k16-step
// (tiles 2kk and 2kk + 1 of S^T are that step's A fragment), B = the f16 hi / lo images of k and q ([GCH][GTC_KP] in
// shared memory, rows = tokens) by ldmatrix: matrices (n-tile 2hf + (lane >> 4), k-half (lane >> 3) & 1), a lane's
// row is token 8 n-tile + (lane & 7). The accumulators come out as 2^(8 + es) times the products. 2 x 96 mma.
__device__ __forceinline__ void gtc_prod_state(const float (&s)[16][4], float sA, const uint16_t* kh, const uint16_t* kl,
const uint16_t* qh, const uint16_t* ql, float (&xt)[4][4], float (&qs)[4][4]) {
const int lane = threadIdx.x & 31;
const int off = (8 * (lane >> 4) + (lane & 7)) * GTC_KP + 8 * ((lane >> 3) & 1);
#pragma unroll
for (int kk = 0; kk < S / 16; ++kk) {
uint32_t ah[4], al[4];
gtc_split_f16x2(s[2 * kk][0] * sA, s[2 * kk][1] * sA, ah[0], al[0]);
gtc_split_f16x2(s[2 * kk][2] * sA, s[2 * kk][3] * sA, ah[1], al[1]);
gtc_split_f16x2(s[2 * kk + 1][0] * sA, s[2 * kk + 1][1] * sA, ah[2], al[2]);
gtc_split_f16x2(s[2 * kk + 1][2] * sA, s[2 * kk + 1][3] * sA, ah[3], al[3]);
#pragma unroll
for (int hf = 0; hf < 2; ++hf) {
const int o = off + 16 * hf * GTC_KP + 16 * kk;
uint32_t bh[4], bl[4];
gtc_ldsm_x4(bh, kh + o);
gtc_ldsm_x4(bl, kl + o);
gtc_mma3_f16(xt[2 * hf], ah, al, bh[0], bh[1], bl[0], bl[1]);
gtc_mma3_f16(xt[2 * hf + 1], ah, al, bh[2], bh[3], bl[2], bl[3]);
gtc_ldsm_x4(bh, qh + o);
gtc_ldsm_x4(bl, ql + o);
gtc_mma3_f16(qs[2 * hf], ah, al, bh[0], bh[1], bl[0], bl[1]);
gtc_mma3_f16(qs[2 * hf + 1], ah, al, bh[2], bh[3], bl[2], bl[3]);
}
}
}
// o^T += (sD D^T) (2^8 P^T) on f16x3: A = the D^T tiles x sD (= 2^ed) split per k16-step, B = the f16 hi / lo images
// of P x 2^8 ([GCH][GTC_PP], rows = tokens t) by ldmatrix as in gtc_prod_state. SKIP leaves out the tiles above P's
// diagonal: tokens t < 16 (n-tiles 0, 1) see no i >= 16 (k-step 1). 18 (SKIP) or 24 mma. D^T += (sY Y^T) (2^et T^T)
// is the same product on the T images (T is zero above its diagonal too), the accumulator x 2^(ey + et).
template <bool SKIP>
__device__ __forceinline__ void gtc_prod_dt(const float (&dt)[4][4], float sD, const uint16_t* ph, const uint16_t* pl,
float (&ot)[4][4]) {
const int lane = threadIdx.x & 31;
const int off = (8 * (lane >> 4) + (lane & 7)) * GTC_PP + 8 * ((lane >> 3) & 1);
#pragma unroll
for (int ks = 0; ks < GCH / 16; ++ks) {
uint32_t ah[4], al[4];
gtc_split_f16x2(dt[2 * ks][0] * sD, dt[2 * ks][1] * sD, ah[0], al[0]);
gtc_split_f16x2(dt[2 * ks][2] * sD, dt[2 * ks][3] * sD, ah[1], al[1]);
gtc_split_f16x2(dt[2 * ks + 1][0] * sD, dt[2 * ks + 1][1] * sD, ah[2], al[2]);
gtc_split_f16x2(dt[2 * ks + 1][2] * sD, dt[2 * ks + 1][3] * sD, ah[3], al[3]);
#pragma unroll
for (int hf = 0; hf < 2; ++hf) {
if (SKIP && ks == 1 && hf == 0) continue;
const int o = off + 16 * hf * GTC_PP + 16 * ks;
uint32_t bh[4], bl[4];
gtc_ldsm_x4(bh, ph + o);
gtc_ldsm_x4(bl, pl + o);
gtc_mma3_f16(ot[2 * hf], ah, al, bh[0], bh[1], bl[0], bl[1]);
gtc_mma3_f16(ot[2 * hf + 1], ah, al, bh[2], bh[3], bl[2], bl[3]);
}
}
}
// S^T += Dd^T (2^8 K) on f16x3: A = the (already scaled) Dd^T tiles split per k16-step, B = the same k images read the
// other way by ldmatrix.trans: matrices (n-tile 2hf + (lane >> 4), k-half (lane >> 3) & 1), a lane's row is token
// 16 ks + 8 k-half + (lane & 7), its 16-B piece state rows 8 n-tile.. 96 mma.
__device__ __forceinline__ void gtc_update_state(float (&s)[16][4], const float (&dd)[4][4], const uint16_t* kh,
const uint16_t* kl) {
const int lane = threadIdx.x & 31;
const int off = (8 * ((lane >> 3) & 1) + (lane & 7)) * GTC_KP + 8 * (lane >> 4);
#pragma unroll
for (int ks = 0; ks < GCH / 16; ++ks) {
uint32_t ah[4], al[4];
gtc_split_f16x2(dd[2 * ks][0], dd[2 * ks][1], ah[0], al[0]);
gtc_split_f16x2(dd[2 * ks][2], dd[2 * ks][3], ah[1], al[1]);
gtc_split_f16x2(dd[2 * ks + 1][0], dd[2 * ks + 1][1], ah[2], al[2]);
gtc_split_f16x2(dd[2 * ks + 1][2], dd[2 * ks + 1][3], ah[3], al[3]);
#pragma unroll
for (int hf = 0; hf < S / 16; ++hf) {
const int o = off + 16 * ks * GTC_KP + 16 * hf;
uint32_t bh[4], bl[4];
gtc_ldsm_x4_trans(bh, kh + o);
gtc_ldsm_x4_trans(bl, kl + o);
gtc_mma3_f16(s[2 * hf], ah, al, bh[0], bh[1], bl[0], bl[1]);
gtc_mma3_f16(s[2 * hf + 1], ah, al, bh[2], bh[3], bl[2], bl[3]);
}
}
}
// ---- the prep: gdn_chunk_prep_kernel (kernels.cu) for the tensor-core scan, a block per (chunk, key head): the same
// A, solve, T and P, but the Gram products k_t.k_i and q_t.k_i on the tensor cores (f16x3 from the k and q hi / lo
// images x 2^8, which the block builds first in shared memory straight from the float4 registers the rows arrive in,
// rows past the chunk's end zero: 48 mma a warp against ~1,150 LDS a thread for the FFMA form, which was LDS-bound at
// 233 us a 2048-token super-block), and it writes what the scan stages by cp.async with no conversion of its own: the
// same k / q images, P as f16 hi / lo images x 2^8 (the split the scan did, element by element) and the token scalars
// (beta, exp(gamma), exp(gL - gamma), gamma) as a float4 a token, T as f16 hi / lo images x 2^et, et from the warp's
// largest |T| (2^-et beside them in te). The Gram entries are within 3 x 2^-22 of exact (the scan's bound), so T and
// P differ from gdn_chunk_prep_kernel's by that.
constexpr size_t kTcPrepSmem = (size_t) 4 * GCH * GTC_KP * 2 + 2 * GCH * GTP * 4; // the images, then k_t.k_i and q_t.k_i
static_assert(VPK * GCH * GTP * 4 <= 4 * GCH * GTC_KP * 2, "gdn_tc_prep_kernel: A fits where the images were");
__global__ void __launch_bounds__(256) gdn_tc_prep_kernel(const float* __restrict__ h, const float* __restrict__ gate,
const float* __restrict__ beta, uint16_t* __restrict__ ti,
float* __restrict__ te, float4* __restrict__ sc,
uint16_t* __restrict__ kq, uint16_t* __restrict__ pi, int64_t t0,
int64_t T) {
extern __shared__ __align__(16) float gsm[];
uint16_t* im = reinterpret_cast<uint16_t*>(gsm); // [4][GCH][GTC_KP] k hi, k lo, q hi, q lo (x 2^8)
float* skk = gsm + 4 * GCH * GTC_KP / 2; // [GCH][GTP] k_t.k_i
float* sqk = skk + GCH * GTP; // [GCH][GTP] q_t.k_i
float* sa = gsm; // [VPK][GCH][GTP] A, then T below the diagonal (over the
// images: read by then)
const int qh = blockIdx.y, ch = blockIdx.x, tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
const int64_t tb = t0 + (int64_t) ch * GCH;
const int n = (int) min((int64_t) GCH, T - tb);
// after the Gram: warps j and VPK + j take value head j (lane i: token i), the first A, the solve and T, the
// second P and the token scalars; their beta and gate (past the end: gate 0, beta 0, no effect on S) are loaded
// now so that they are in by then
const int j = warp % VPK, vh = qh + j * HK, i = lane;
const bool pv = warp < 2 * VPK && i < n;
const float b = pv ? beta[(tb + i) * HV + vh] : 0.0f;
float gam = pv ? gate[(tb + i) * HV + vh] : 0.0f;
{ // k and q as float4, every load issued before the first store; the images to shared and to global
constexpr int NP = GCH * S / 4 / 256; // float4 a thread, each of k and q
float4 kv[NP], qv[NP];
#pragma unroll
for (int m = 0; m < NP; ++m) {
const int p = tid + 256 * m, t = p / (S / 4), r = (p % (S / 4)) * 4;
kv[m] = t < n ? *reinterpret_cast<const float4*>(h + (tb + t) * C + HK * S + qh * S + r) : make_float4(0.f, 0.f, 0.f, 0.f);
qv[m] = t < n ? *reinterpret_cast<const float4*>(h + (tb + t) * C + qh * S + r) : make_float4(0.f, 0.f, 0.f, 0.f);
}
uint16_t* img = kq + ((size_t) ch * HK + qh) * 4 * GCH * S;
#pragma unroll
for (int m = 0; m < NP; ++m) {
const int p = tid + 256 * m, t = p / (S / 4), r = (p % (S / 4)) * 4;
const int o = t * GTC_KP + r;
gtc_split_store4(im + o, im + GCH * GTC_KP + o, kv[m], 256.0f);
gtc_split_store4(im + 2 * GCH * GTC_KP + o, im + 3 * GCH * GTC_KP + o, qv[m], 256.0f);
gtc_split_store4(img + t * S + r, img + GCH * S + t * S + r, kv[m], 256.0f);
gtc_split_store4(img + 2 * GCH * S + t * S + r, img + 3 * GCH * S + t * S + r, qv[m], 256.0f);
}
}
__syncthreads();
{ // the Gram products: warp w takes product w >> 2 (k_t.k_i, q_t.k_i), rows 16 ((w >> 1) & 1).. (the A
// tile, by ldmatrix as a row-major [t][r] 16 x 16), columns 16 (w & 1).. (two n-tiles: B = k^T, stored
// [n][k] = the k rows, by ldmatrix as the scan's gtc_prod_state); the accumulators come out x 2^16
const int pr = warp >> 2, mt = (warp >> 1) & 1, np = warp & 1;
const uint16_t* ah_ = im + 2 * pr * GCH * GTC_KP + (16 * mt + (lane & 7) + 8 * ((lane >> 3) & 1)) * GTC_KP + 8 * (lane >> 4);
const uint16_t* bh_ = im + (16 * np + 8 * (lane >> 4) + (lane & 7)) * GTC_KP + 8 * ((lane >> 3) & 1);
float c[2][4] = {};
#pragma unroll
for (int kk = 0; kk < S / 16; ++kk) {
uint32_t ah[4], al[4], bh[4], bl[4];
gtc_ldsm_x4(ah, ah_ + 16 * kk);
gtc_ldsm_x4(al, ah_ + GCH * GTC_KP + 16 * kk);
gtc_ldsm_x4(bh, bh_ + 16 * kk);
gtc_ldsm_x4(bl, bh_ + GCH * GTC_KP + 16 * kk);
gtc_mma3_f16(c[0], ah, al, bh[0], bh[1], bl[0], bl[1]);
gtc_mma3_f16(c[1], ah, al, bh[2], bh[3], bl[2], bl[3]);
}
float* g = pr ? sqk : skk;
const int gid = lane >> 2, tig = lane & 3;
#pragma unroll
for (int j = 0; j < 2; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) {
const int t = 16 * mt + gid + 8 * (q >> 1), i = 16 * np + 8 * j + 2 * tig + (q & 1);
g[t * GTP + i] = i <= t ? c[j][q] * gtc_pow2(-16) : 0.0f;
}
}
__syncthreads();
if (warp >= 2 * VPK) return;
#pragma unroll
for (int o = 1; o < 32; o <<= 1) {
const float u = __shfl_up_sync(0xffffffffu, gam, o);
if (lane >= o) gam += u;
}
if (warp >= VPK) { // P = exp(gamma_t - gamma_i) q_t.k_i (i <= t) x 2^8 as f16 hi / lo; the scalars
uint16_t* Ph = pi + ((int64_t) ch * HV + vh) * 2 * GCH * GCH;
uint16_t* Pl = Ph + GCH * GCH;
#pragma unroll
for (int t = 0; t < GCH; ++t) {
const float gt = __shfl_sync(0xffffffffu, gam, t);
const float p = (i <= t ? expf(gt - gam) * sqk[t * GTP + i] : 0.0f) * 256.0f;
const __half hi = __float2half_rn(p);
Ph[t * GCH + i] = __half_as_ushort(hi);
Pl[t * GCH + i] = __half_as_ushort(__float2half_rn(p - __half2float(hi)));
}
const float gL = __shfl_sync(0xffffffffu, gam, GCH - 1);
sc[((int64_t) ch * HV + vh) * GCH + i] = make_float4(b, expf(gam), expf(gL - gam), gam);
return;
}
float* a = sa + j * GCH * GTP;
#pragma unroll
for (int t = 0; t < GCH; ++t) {
const float gt = __shfl_sync(0xffffffffu, gam, t), bt = __shfl_sync(0xffffffffu, b, t);
a[t * GTP + i] = i < t ? bt * expf(gt - gam) * skk[t * GTP + i] : (i == t ? 1.0f : 0.0f);
}
__syncwarp();
// T = (I + A)^-1 by forward substitution, row by row: T[t][i] = -sum_{k < t} A[t][k] T[k][i] over the finished
// rows (diagonal 1, zeros above it: the terms k < i add an exact 0, so the same bits as summing from k = i).
// Lane i keeps its column T[.][i] in registers, so a row needs only A's row t (a broadcast load a term) and A
// is never written: no store and no __syncwarp a row
float tv[GCH];
tv[0] = i == 0 ? 1.0f : 0.0f;
#pragma unroll
for (int t = 1; t < GCH; ++t) {
float v = 0.0f;
#pragma unroll
for (int k = 0; k < t; ++k) v = fmaf(-a[t * GTP + k], tv[k], v);
tv[t] = i < t ? v : (i == t ? 1.0f : 0.0f);
}
// T's images x 2^et (its largest entry in [2^14, 2^15), as the scan scales its own operands)
float mt4[4] = {};
#pragma unroll
for (int t = 0; t < GCH; ++t) mt4[t & 3] = fmaxf(mt4[t & 3], fabsf(tv[t]));
const int et = gtc_scale_exp(gtc_warp_max(fmaxf(fmaxf(mt4[0], mt4[1]), fmaxf(mt4[2], mt4[3]))));
const float sT = gtc_pow2(et);
uint16_t* Th = ti + ((int64_t) ch * HV + vh) * 2 * GCH * GCH;
uint16_t* Tl = Th + GCH * GCH;
#pragma unroll
for (int t = 0; t < GCH; ++t) {
const float x = tv[t] * sT;
const __half hi = __float2half_rn(x);
Th[t * GCH + i] = __half_as_ushort(hi);
Tl[t * GCH + i] = __half_as_ushort(__float2half_rn(x - __half2float(hi)));
}
if (i == 0) te[(int64_t) ch * HV + vh] = gtc_pow2(-et);
}
// ---- the scan: a block per value head (48 blocks, one wave on 48 of the CMP 170HX's 70 SMs), 8 consumer warps of 16
// value columns each walking every chunk of the super-block with their state slice in registers (written back at the
// end as gdn_chunk_scan_kernel does), the prep's scratch from gdn_tc_prep_kernel. Everything a chunk needs arrives by
// cp.async, issued for the next chunk right after the one __syncthreads a chunk (which publishes the buffer just
// landed and ends every read of the other one): the k / q and P f16 images, T and the token scalars as the prep
// wrote them, v from h (its rows past the chunk's end zero-filled). The 256 threads do no conversion of their own.
// Earlier forms, measured on this card: the scan splitting k / q itself after the state update (0.282 us a token;
// ~300 split instructions a warp a chunk and the staging exposed at the chunk's end); two stager warps with named
// barriers (6 % slower: registers are allocated per SM sub-partition, 16K each, and a block of more than 8 warps puts
// 3 on a sub-partition, which caps the kernel at 168 registers a thread; the consumers need about 200 and spilled.
// 8 warps, 2 a sub-partition, may have 255).
// NORM (step 5: the output norm fused, the FP16 output written here) is not implemented yet; oc_out takes the FP32
// output as gdn_chunk_scan_kernel writes it.
template <bool NORM>
__global__ void __launch_bounds__(GTC_CW * 32, 1)
gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const uint16_t* __restrict__ tig_,
const float* __restrict__ teg, const float4* __restrict__ scg, const uint16_t* __restrict__ kqg,
const uint16_t* __restrict__ pig, float* __restrict__ oc_out, int64_t t0, int nch, int64_t T) {
static_assert(!NORM, "gdn_tc_scan_kernel: the fused output norm is not implemented yet");
constexpr int NT = GTC_CW * 32;
extern __shared__ __align__(16) unsigned char gtc_sm[];
uint16_t* kh = reinterpret_cast<uint16_t*>(gtc_sm); // [2][GCH][GTC_KP] f16 of k x 2^8 (hi)
uint16_t* kl = kh + 2 * GCH * GTC_KP; // [2][GCH][GTC_KP] what hi left out (lo)
uint16_t* qh = kl + 2 * GCH * GTC_KP; // q likewise
uint16_t* ql = qh + 2 * GCH * GTC_KP;
float* sv = reinterpret_cast<float*>(ql + 2 * GCH * GTC_KP); // [2][GCH][GTC_VP] v
uint16_t* th = reinterpret_cast<uint16_t*>(sv + 2 * GCH * GTC_VP); // [2][GCH][GTC_PP] T x 2^et hi / lo
uint16_t* tl = th + 2 * GCH * GTC_PP;
uint16_t* ph = tl + 2 * GCH * GTC_PP; // [2][GCH][GTC_PP] P x 2^8 hi / lo
uint16_t* pl = ph + 2 * GCH * GTC_PP;
float4* sc = reinterpret_cast<float4*>(pl + 2 * GCH * GTC_PP); // [2][GCH] (beta, exp(gamma), exp(gL - gamma), gamma)
float* ste = reinterpret_cast<float*>(sc + 2 * GCH); // [2] 2^-et
const int vh = blockIdx.x, qhd = vh % HK;
const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31, gid = lane >> 2, tig = lane & 3;
const size_t rs = (size_t) HV * S;
// chunk ch into buffer b, this thread's 16-B pieces (one cp.async group): the four k / q images (4 x GCH rows of
// 16 pieces: 8 a thread), v (4 a thread), the two T images (1), the two P images (1), the scalars (warp 0, 1),
// T's scale (thread 0, 4 B)
auto stage = [&](int ch, int b) {
const int64_t c_t0 = t0 + (int64_t) ch * GCH;
const int n = (int) min((int64_t) GCH, T - c_t0);
const uint16_t* img = kqg + ((size_t) ch * HK + qhd) * 4 * GCH * S;
#pragma unroll
for (int m = 0; m < 4 * GCH * S / 8 / NT; ++m) {
const int im = m >> 1, t = ((m & 1) * NT + tid) / (S / 8), r = (tid % (S / 8)) * 8;
uint16_t* dst = (im == 0 ? kh : im == 1 ? kl : im == 2 ? qh : ql) + (b * GCH + t) * GTC_KP + r;
gdn_cp16(reinterpret_cast<float*>(dst), reinterpret_cast<const float*>(img + (im * GCH + t) * S + r));
}
#pragma unroll
for (int m = 0; m < GCH * S / 4 / NT; ++m) {
const int p = tid + NT * m, t = p / (S / 4), r = (p % (S / 4)) * 4;
float* dst = sv + (b * GCH + t) * GTC_VP + r;
if (t < n) gdn_cp16(dst, h + (c_t0 + t) * C + 2 * HK * S + vh * S + r);
else *reinterpret_cast<float4*>(dst) = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
}
{
const int im = tid / (GCH * GCH / 8), t = (tid % (GCH * GCH / 8)) / (GCH / 8), i = (tid % (GCH / 8)) * 8;
const int64_t src = ((((int64_t) ch * HV + vh) * 2 + im) * GCH + t) * GCH + i;
const int dst = (b * GCH + t) * GTC_PP + i;
gdn_cp16(reinterpret_cast<float*>((im ? tl : th) + dst), reinterpret_cast<const float*>(tig_ + src));
gdn_cp16(reinterpret_cast<float*>((im ? pl : ph) + dst), reinterpret_cast<const float*>(pig + src));
}
if (tid < GCH) gdn_cp16(reinterpret_cast<float*>(sc + b * GCH + tid), reinterpret_cast<const float*>(scg + ((int64_t) ch * HV + vh) * GCH + tid));
if (tid == 0) gdn_cp4(ste + b, teg + (int64_t) ch * HV + vh);
gdn_cp_commit();
};
const int c0 = 16 * warp; // the warp's value columns c0 + gid and c0 + gid + 8
// the state slice as S^T fragments: tile j = rows 8j..8j+7; s[j][0] = S[8j + 2tig][c0 + gid], [1] the row after,
// [2] column + 8, [3] both (64 scalar loads a lane, 32-B runs, once a launch)
float s[16][4];
#pragma unroll
for (int j = 0; j < 16; ++j) {
const float* p = state + (size_t) (8 * j + 2 * tig) * rs + (size_t) vh * S + c0 + gid;
s[j][0] = p[0]; s[j][1] = p[rs]; s[j][2] = p[8]; s[j][3] = p[rs + 8];
}
stage(0, 0);
const float out_scale = rsqrtf((float) S);
for (int ch = 0; ch < nch; ++ch) {
const int b = ch & 1;
const int64_t c_t0 = t0 + (int64_t) ch * GCH;
const int n = (int) min((int64_t) GCH, T - c_t0);
gdn_cp_wait_all(); // this thread's pieces of chunk ch
__syncthreads(); // every thread's: chunk ch's buffer complete; every read of the other one done
if (ch + 1 < nch) stage(ch + 1, b ^ 1);
const uint16_t* bkh = kh + b * GCH * GTC_KP;
const uint16_t* bkl = kl + b * GCH * GTC_KP;
const uint16_t* bqh = qh + b * GCH * GTC_KP;
const uint16_t* bql = ql + b * GCH * GTC_KP;
const uint16_t* bth = th + b * GCH * GTC_PP;
const uint16_t* btl = tl + b * GCH * GTC_PP;
const uint16_t* bph = ph + b * GCH * GTC_PP;
const uint16_t* bpl = pl + b * GCH * GTC_PP;
const float* bsv = sv + b * GCH * GTC_VP;
const float4* bsc = sc + b * GCH;
// (a) X^T = S0^T K^T and QS^T = S0^T Q^T, the state x 2^es (its largest entry in [2^14, 2^15)). The lane's
// maximum over eight chains (nvcc keeps the order written: one chain was 64 dependent FMNMX before the first mma)
float m8[8] = {};
#pragma unroll
for (int j = 0; j < 16; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) m8[(2 * j + (q >> 1)) & 7] = fmaxf(m8[(2 * j + (q >> 1)) & 7], fabsf(s[j][q]));
const float m = fmaxf(fmaxf(fmaxf(m8[0], m8[1]), fmaxf(m8[2], m8[3])), fmaxf(fmaxf(m8[4], m8[5]), fmaxf(m8[6], m8[7])));
const int es = gtc_scale_exp(gtc_warp_max(m));
float xt[4][4] = {}, qs[4][4] = {};
gtc_prod_state(s, gtc_pow2(es), bkh, bkl, bqh, bql, xt, qs);
// (b) Y^T = beta_i (V^T - e'_i X^T), e'_i = exp(gamma_i) 2^-(8 + es): lane (gid, tig) holds tokens
// 8nt + 2tig and + 1, columns c0 + gid and + 8
const float eS = gtc_pow2(-(8 + es));
float yt[4][4];
#pragma unroll
for (int nt = 0; nt < 4; ++nt) {
const float4 s0 = gtc_lds128(bsc + 8 * nt + 2 * tig), s1 = gtc_lds128(bsc + 8 * nt + 2 * tig + 1);
const float* v = bsv + (8 * nt + 2 * tig) * GTC_VP + c0 + gid;
const float e0 = s0.y * eS, e1 = s1.y * eS;
yt[nt][0] = s0.x * fmaf(-e0, xt[nt][0], v[0]);
yt[nt][1] = s1.x * fmaf(-e1, xt[nt][1], v[GTC_VP]);
yt[nt][2] = s0.x * fmaf(-e0, xt[nt][2], v[8]);
yt[nt][3] = s1.x * fmaf(-e1, xt[nt][3], v[GTC_VP + 8]);
}
// (c) D^T = Y^T T^T on f16x3: Y^T x 2^ey (its largest entry in [2^14, 2^15)), T's images x 2^et from the prep;
// the accumulator comes out x 2^(ey + et) (T above its diagonal skipped)
float my4[4] = {};
#pragma unroll
for (int nt = 0; nt < 4; ++nt)
#pragma unroll
for (int q = 0; q < 4; ++q) my4[q] = fmaxf(my4[q], fabsf(yt[nt][q]));
const int ey = gtc_scale_exp(gtc_warp_max(fmaxf(fmaxf(my4[0], my4[1]), fmaxf(my4[2], my4[3]))));
float dt[4][4] = {};
gtc_prod_dt<true>(yt, gtc_pow2(ey), bth, btl, dt);
const float sYT = gtc_pow2(-ey) * ste[b];
#pragma unroll
for (int nt = 0; nt < 4; ++nt)
#pragma unroll
for (int q = 0; q < 4; ++q) dt[nt][q] *= sYT;
// (d) o^T = e_t QS^T + D^T P^T: the accumulator starts as 2^(ed + 8) e_t QS^T (QS^T came out x 2^(8 + es)),
// D^T goes in x 2^ed, P x 2^8; the output is the accumulator x 2^-(ed + 8), then rsqrt(S) as gdn_chunk_scan_kernel
float md4[4] = {};
#pragma unroll
for (int nt = 0; nt < 4; ++nt)
#pragma unroll
for (int q = 0; q < 4; ++q) md4[q] = fmaxf(md4[q], fabsf(dt[nt][q]));
const int ed = gtc_scale_exp(gtc_warp_max(fmaxf(fmaxf(md4[0], md4[1]), fmaxf(md4[2], md4[3]))));
const float sD = gtc_pow2(ed), sO = gtc_pow2(ed + 8), sOi = gtc_pow2(-(ed + 8)), eO = gtc_pow2(ed - es);
float ot[4][4];
#pragma unroll
for (int nt = 0; nt < 4; ++nt) {
const float e0 = gtc_lds128(bsc + 8 * nt + 2 * tig).y * eO, e1 = gtc_lds128(bsc + 8 * nt + 2 * tig + 1).y * eO;
ot[nt][0] = qs[nt][0] * e0; ot[nt][1] = qs[nt][1] * e1; ot[nt][2] = qs[nt][2] * e0; ot[nt][3] = qs[nt][3] * e1;
}
gtc_prod_dt<true>(dt, sD, bph, bpl, ot);
#pragma unroll
for (int nt = 0; nt < 4; ++nt)
#pragma unroll
for (int q = 0; q < 4; ++q) {
const int t = 8 * nt + 2 * tig + (q & 1);
if (t < n) oc_out[(c_t0 + t) * rs + (size_t) vh * S + c0 + gid + 8 * (q >> 1)] = (ot[nt][q] * sOi) * out_scale;
}
// (e) S^T = eL S0^T + Dd^T K, Dd^T = f_i D^T: the accumulator x (eL 2^(ed + 8)) before (one rounding, as the
// s *= gL of gdn_chunk_scan_kernel), Dd^T in x 2^ed, K x 2^8, the accumulator x 2^-(ed + 8) after
const float sPre = gtc_lds128(bsc + GCH - 1).y * sO; // exp(gL)
#pragma unroll
for (int j = 0; j < 16; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) s[j][q] *= sPre;
float dd[4][4];
#pragma unroll
for (int nt = 0; nt < 4; ++nt) {
const float f0 = gtc_lds128(bsc + 8 * nt + 2 * tig).z * sD, f1 = gtc_lds128(bsc + 8 * nt + 2 * tig + 1).z * sD;
dd[nt][0] = dt[nt][0] * f0; dd[nt][1] = dt[nt][1] * f1; dd[nt][2] = dt[nt][2] * f0; dd[nt][3] = dt[nt][3] * f1;
}
gtc_update_state(s, dd, bkh, bkl);
#pragma unroll
for (int j = 0; j < 16; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) s[j][q] *= sOi;
}
#pragma unroll
for (int j = 0; j < 16; ++j) {
float* p = state + (size_t) (8 * j + 2 * tig) * rs + (size_t) vh * S + c0 + gid;
p[0] = s[j][0]; p[rs] = s[j][1]; p[8] = s[j][2]; p[rs + 8] = s[j][3];
}
}
#endif // !__HIPCC__
src/prefill/kernels.cu
+46 -1
View File
@@ -879,6 +879,7 @@ bool gdn_keyhead_ok() {
// than they save below ~100 tokens (gdn_rec_parity --bench: 32 tokens 0.58x gdn_rec_kh_kernel, 96 1.03x, 128 1.20x,
// 512 1.67x, 2048+ 1.81x)
constexpr int64_t kGdnChunkedMin = 128;
constexpr int64_t kGdnTcMin = 128; // STRATA_GDN_TC: the two launches cost ~20 us; below this the kh kernel wins
constexpr int GCH = 32; // tokens per chunk (the prep's solve is a lane per token)
constexpr int GSB = 2048; // tokens per super-block: the prep's scratch is GSB / GCH chunks of T, P and gamma
constexpr int GDV = 16; // value columns per scan block
@@ -1147,6 +1148,9 @@ __global__ void __launch_bounds__(128 * VPK) gdn_chunk_scan_kernel(float* __rest
#pragma unroll
for (int y = 0; y < 4; ++y) state[(size_t) (4 * rq + xx) * rs + (size_t) vh * S + c0 + vv + y] = s[xx][y];
}
// the same scan on the tensor cores (STRATA_GDN_TC=1, gdn_rec_tc below); src/prefill/gdn_rec_parity.cu includes the
// same text
#include "gdn_tc.cuh"
// The scratch, one a device: allocated on first use at its largest (GSB / GCH chunks of T, P and gamma: 25.6 MB) and
// kept (a device runs its prompt on one stream). A stream-ordered allocation and free per call, even from a pool that
@@ -1213,6 +1217,42 @@ cudaError_t gdn_rec_chunked(float* state, const float* h, const float* gate, con
}
return cudaSuccess; // launch errors surface at the caller's check; an error above means nothing ran
}
// STRATA_GDN_TC=1: gdn_rec_chunked with the scan on the tensor cores (gdn_tc_scan_kernel, gdn_tc.cuh) and the prep
// that also writes its f16 images (gdn_tc_prep_kernel, scratch gdn_tc_scratch), 48 blocks of 256 threads in one wave. An error return means nothing was launched: the card is below
// sm_80, has fewer than 48 SMs, or less than 122 KB of shared memory a block.
cudaError_t gdn_rec_tc(float* state, const float* h, const float* gate, const float* beta, float* y, int64_t T,
cudaStream_t s) {
if (T <= 0) return cudaSuccess;
{
int dev = 0, major = 0, sms = 0, smem = 0;
if (cudaGetDevice(&dev) != cudaSuccess ||
cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) != cudaSuccess ||
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev) != cudaSuccess ||
cudaDeviceGetAttribute(&smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev) != cudaSuccess) {
cudaGetLastError();
return cudaErrorNotSupported;
}
if (strata::cc_major_of(major) < 8 || sms < HV || (size_t) strata::smem_optin_of(smem) < kTcScanSmem<false>)
return cudaErrorNotSupported;
}
const GtcScratch sc = gdn_tc_scratch();
if (sc.ti == nullptr) return cudaErrorMemoryAllocation;
{ // once a device
static bool attrs[64] = {};
int dev = 0;
cudaGetDevice(&dev);
if (dev >= 0 && dev < 64 && !attrs[dev]) {
cudaFuncSetAttribute(gdn_tc_scan_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) kTcScanSmem<false>);
attrs[dev] = true;
}
}
for (int64_t t0 = 0; t0 < T; t0 += GSB) {
const int nch = (int) ((std::min<int64_t>(GSB, T - t0) + GCH - 1) / GCH);
gdn_tc_prep_kernel<<<dim3((unsigned) nch, HK), 256, kTcPrepSmem, s>>>(h, gate, beta, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, t0, T);
gdn_tc_scan_kernel<false><<<HV, GTC_CW * 32, kTcScanSmem<false>, s>>>(state, h, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, y, t0, nch, T);
}
return cudaSuccess;
}
#endif
#if defined(__HIPCC__) // the Aurora (gfx11) recurrence kernels: AMD builds only
// Aurora (S23): the recurrence with the column-split grid of gdn_rec_cols_pipe_kernel (four 32-column blocks per value
@@ -2709,7 +2749,12 @@ void gdn_recurrence_variant(int variant, float* state, const float* h, const flo
// STRATA_GDN_CHUNKED=1 (opt-in): the recurrence in chunks for prompt chunks of kGdnChunkedMin+ tokens (other
// bits; a shorter chunk, a card it does not fit or a failed scratch allocation takes the kernels below)
static const bool chunked = [] { const char* v = std::getenv("STRATA_GDN_CHUNKED"); return v != nullptr && std::atoi(v) != 0; }();
if (chunked && T >= kGdnChunkedMin && gdn_rec_chunked(state, h, gate, beta, y, T, (cudaStream_t) stream) == cudaSuccess) {
// STRATA_GDN_TC=1 (opt-in): the chunked recurrence with the scan on the tensor cores (gdn_tc.cuh) for prompt
// chunks of kGdnTcMin+ tokens: within 1e-4 of the kernels below (gdn_rec_parity), 3.75x gdn_rec_kh_kernel on a
// CMP 170HX; a shorter chunk, a card it does not fit or a failed scratch allocation takes the kernels below
static const bool tc = [] { const char* v = std::getenv("STRATA_GDN_TC"); return v != nullptr && std::atoi(v) != 0; }();
if (tc && T >= kGdnTcMin && gdn_rec_tc(state, h, gate, beta, y, T, (cudaStream_t) stream) == cudaSuccess) {
} else if (chunked && T >= kGdnChunkedMin && gdn_rec_chunked(state, h, gate, beta, y, T, (cudaStream_t) stream) == cudaSuccess) {
} else if (pipe && gdn_keyhead_ok()) // the value heads of a key head in one thread (same bits)
gdn_rec_kh_kernel<<<HK * NCB, dim3(CB, RG), 0, (cudaStream_t) stream>>>(state, h, gate, beta, y, T);
else