cuda: chunked WY delta-net kernel v1 (opt-in DELTA_WY_CUDA=1) + harness --cuda mode
This commit is contained in:
parent
3a7d1016b3
commit
a34feeb0f9
|
|
@ -1,5 +1,6 @@
|
|||
#include "common.cuh"
|
||||
#include "delta-net.cuh"
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
|
||||
|
|
@ -182,6 +183,266 @@ __global__ void delta_net_recurrent_f32(
|
|||
}
|
||||
}
|
||||
|
||||
|
||||
// Chunked WY parallel-scan delta net for long prefills (qwen4exp PP).
|
||||
// Validated math: tests/test-delta-chunk.cpp (candidate C) must stay green —
|
||||
// any change here needs a matching change there first. One CUDA block per
|
||||
// (seq, v-head); 64-token chunks (32-token fallback tier); per chunk:
|
||||
// scalars b/d/logP per token (shared, 3*C floats; logP sequential by thread 0)
|
||||
// L/KQ dot rows recomputed per use (no L/KQ storage; 64-way parallel dots,
|
||||
// results broadcast via shared row, kh/qt stream from global/L1)
|
||||
// triangular solve for E (E rows persist in shared; flat-partitioned updates)
|
||||
// outputs + single-GEMM assembly of S_C (valid: dispatch guarantees
|
||||
// saved_states == NULL, so no per-step states are needed)
|
||||
// guard (max decay-ratio / max |e|, nan-aware): wild dynamics fall back to
|
||||
// the exact sequential inner loop for that chunk (same block, S still
|
||||
// intact in shared). Operating regime never trips it.
|
||||
// v1 is OPT-IN (DELTA_WY_CUDA=1) and correctness-first: SIMT code, no tensor
|
||||
// cores/cp.async (v2 after measurements). Shared budget d=128/C=64 is
|
||||
// (16384+8192+192+64)*4 = 99328 bytes (opt-in attr); denial or smaller
|
||||
// budgets drop to C=32 tier, then to the sequential kernel. All silent and
|
||||
// exact in every tier.
|
||||
#define DELTA_WY_CHUNK 64
|
||||
#define DELTA_WY_CHUNK_SMALL 32
|
||||
#define WY_GUARD_D 1e4f
|
||||
#define WY_GUARD_E 1e4f
|
||||
|
||||
template <int HEAD_DIM>
|
||||
__global__ void delta_net_chunked_f32(
|
||||
const float * __restrict__ q,
|
||||
const float * __restrict__ k,
|
||||
const float * __restrict__ v,
|
||||
const float * __restrict__ g,
|
||||
const float * __restrict__ beta_in,
|
||||
const float * __restrict__ state_in,
|
||||
float * __restrict__ dst,
|
||||
float * __restrict__ state_out,
|
||||
const int64_t n_heads,
|
||||
const int64_t gqa_ratio,
|
||||
const int repeat_type,
|
||||
const int64_t n_tokens,
|
||||
const int64_t n_seqs,
|
||||
size_t vnb1, size_t vnb2, size_t vnb3,
|
||||
const int chunk) {
|
||||
const int batch_idx = (int)(blockIdx.x / n_heads);
|
||||
const int head_idx = (int)(blockIdx.x % n_heads);
|
||||
const int head_idx_kq = repeat_type == 0 ? (int)(head_idx / gqa_ratio) : (int)(head_idx % (n_heads/gqa_ratio));
|
||||
const int tid = (int)threadIdx.x;
|
||||
const int nthreads = (int)blockDim.x;
|
||||
|
||||
// identical strides to delta_net_recurrent_f32
|
||||
const int64_t qkv_stride_token = HEAD_DIM;
|
||||
const int64_t qkv_stride_head = HEAD_DIM * n_tokens;
|
||||
const int64_t qkv_stride_batch = HEAD_DIM * n_tokens * n_heads;
|
||||
const int64_t qkv_stride_batch_kq = qkv_stride_batch / gqa_ratio;
|
||||
const int64_t g_stride_batch = n_tokens * n_heads;
|
||||
const int64_t state_head_offset = head_idx * HEAD_DIM * HEAD_DIM;
|
||||
const int64_t state_batch_stride = HEAD_DIM * HEAD_DIM * n_heads;
|
||||
const int64_t out_token_stride = HEAD_DIM * n_heads;
|
||||
|
||||
const float * q_base = q + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head;
|
||||
const float * k_base = k + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head;
|
||||
const float * v_base = v + batch_idx * vnb3 + head_idx * vnb2;
|
||||
const float * g_base = g + batch_idx * g_stride_batch + head_idx;
|
||||
const float * beta_base = beta_in + batch_idx * g_stride_batch + head_idx;
|
||||
const float * state_src = state_in + batch_idx * state_batch_stride + state_head_offset;
|
||||
float * out_base = dst + batch_idx * (HEAD_DIM * n_heads * n_tokens) + head_idx * HEAD_DIM;
|
||||
float * state_dst = state_out + batch_idx * state_batch_stride + state_head_offset;
|
||||
|
||||
const float scale = rsqrtf((float)HEAD_DIM);
|
||||
|
||||
extern __shared__ float smem[];
|
||||
float * sS = smem; // HEAD_DIM*HEAD_DIM (carry)
|
||||
float * sE = sS + HEAD_DIM * HEAD_DIM; // chunk*HEAD_DIM
|
||||
float * sScl = sE + (int64_t)chunk * HEAD_DIM; // 3*chunk (b,d,logP)
|
||||
float * sRow = sScl + 3 * chunk; // chunk (dot rows; [0..7] reused by guard reduce)
|
||||
float * sBv = sScl, * sDv = sScl + chunk, * sLogP = sScl + 2 * chunk;
|
||||
|
||||
for (int64_t x = tid; x < (int64_t)HEAD_DIM * HEAD_DIM; x += nthreads) sS[x] = state_src[x];
|
||||
__syncthreads();
|
||||
|
||||
for (int64_t t0 = 0; t0 < n_tokens; t0 += chunk) {
|
||||
const int64_t rem = n_tokens - t0;
|
||||
const int C = (int)(rem < chunk ? rem : chunk);
|
||||
|
||||
// ---- scalars (logP sequential by thread 0: prefix dependency)
|
||||
for (int i = tid; i < C; i += nthreads) {
|
||||
const float g_val = g_base[(t0 + i) * n_heads];
|
||||
const float b_raw = beta_base[(t0 + i) * n_heads];
|
||||
sBv[i] = 1.0f / (1.0f + expf(-b_raw));
|
||||
sDv[i] = expf(fminf(g_val, 50.0f));
|
||||
}
|
||||
if (tid == 0) {
|
||||
float acc = 0.0f;
|
||||
for (int i = 0; i < C; ++i) {
|
||||
acc += fminf(g_base[(t0 + i) * n_heads], 50.0f);
|
||||
sLogP[i] = acc;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ---- solve for E (C sequential steps; flat-partitioned element updates)
|
||||
float thr_maxD = 0.0f, thr_maxE = 0.0f;
|
||||
for (int i = 0; i < C; ++i) {
|
||||
// L[i][j] = b_i * exp(logP[i]-logP[j]) * (kh_i.kh_j), j<i; thread j
|
||||
if (tid < i) {
|
||||
const int j = tid;
|
||||
float dot = 0.0f;
|
||||
for (int d = 0; d < HEAD_DIM; ++d) {
|
||||
dot += k_base[(t0 + i) * qkv_stride_token + d]
|
||||
* k_base[(t0 + j) * qkv_stride_token + d];
|
||||
}
|
||||
const float Dij = expf(sLogP[i] - sLogP[j]);
|
||||
if (!(Dij <= thr_maxD)) thr_maxD = Dij; // nan-aware max
|
||||
sRow[j] = sBv[i] * Dij * dot;
|
||||
}
|
||||
__syncthreads();
|
||||
// e_i[d] for owned d; ck recomputed (S rows in shared, kh streamed)
|
||||
for (int dd = tid; dd < HEAD_DIM; dd += nthreads) {
|
||||
float ck = 0.0f;
|
||||
for (int e = 0; e < HEAD_DIM; ++e) ck += sS[dd + e * HEAD_DIM] * k_base[(t0 + i) * qkv_stride_token + e];
|
||||
const float Pi = expf(sLogP[i]); // underflow to 0 is correct
|
||||
const float * vv = v_base + (t0 + i) * vnb1;
|
||||
float ev = sBv[i] * vv[dd] - sBv[i] * Pi * ck;
|
||||
for (int j = 0; j < i; ++j) ev -= sRow[j] * sE[j * HEAD_DIM + dd];
|
||||
sE[i * HEAD_DIM + dd] = ev;
|
||||
const float ae = fabsf(ev);
|
||||
if (!(ae <= thr_maxE)) thr_maxE = ae; // nan-aware max
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// ---- guard reduce (warp shuffle -> sRow[0..7] -> scalar mailbox sRow[0])
|
||||
{
|
||||
float md = thr_maxD, me = thr_maxE;
|
||||
for (int off = 16; off > 0; off >>= 1) {
|
||||
const float od = __shfl_down_sync(0xffffffff, md, off);
|
||||
const float oe = __shfl_down_sync(0xffffffff, me, off);
|
||||
if (!(od <= md)) md = od;
|
||||
if (!(oe <= me)) me = oe;
|
||||
}
|
||||
const int wid = tid / 32, lane = tid % 32;
|
||||
if (lane == 0) { sRow[wid] = md; sRow[8 + wid] = me; }
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
float gmd = 0.0f, gme = 0.0f;
|
||||
const int nwarps = nthreads / 32;
|
||||
for (int w = 0; w < nwarps; ++w) {
|
||||
if (!(sRow[w] <= gmd)) gmd = sRow[w];
|
||||
if (!(sRow[8 + w] <= gme)) gme = sRow[8 + w];
|
||||
}
|
||||
sRow[0] = (gmd <= WY_GUARD_D && gme <= WY_GUARD_E) ? 1.0f : 0.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (sRow[0] < 0.5f) {
|
||||
// ---- sequential fallback for this chunk (exact old math).
|
||||
// Fully thread-local per owned row (no shared scratch): dots are
|
||||
// per-row serial loops, only the scalar score needs broadcast.
|
||||
for (int i = 0; i < C; ++i) {
|
||||
float score = 0.0f;
|
||||
if (tid == 0) {
|
||||
float s = 0.0f;
|
||||
for (int d = 0; d < HEAD_DIM; ++d) {
|
||||
s += k_base[(t0 + i) * qkv_stride_token + d] * q_base[(t0 + i) * qkv_stride_token + d];
|
||||
}
|
||||
score = s * scale;
|
||||
sRow[0] = score;
|
||||
sRow[1] = 1.0f / (1.0f + expf(-beta_base[(t0 + i) * n_heads]));
|
||||
sRow[2] = expf(fminf(g_base[(t0 + i) * n_heads], 50.0f));
|
||||
}
|
||||
__syncthreads();
|
||||
{
|
||||
const float sc = sRow[0], bv = sRow[1], dv = sRow[2];
|
||||
float * out_t = out_base + (t0 + i) * out_token_stride;
|
||||
const float * vv = v_base + (t0 + i) * vnb1;
|
||||
const float * kk = k_base + (t0 + i) * qkv_stride_token;
|
||||
const float * qq = q_base + (t0 + i) * qkv_stride_token;
|
||||
for (int dd = tid; dd < HEAD_DIM; dd += nthreads) {
|
||||
float sk = 0.0f, sq = 0.0f;
|
||||
for (int e = 0; e < HEAD_DIM; ++e) {
|
||||
sk += sS[dd + e * HEAD_DIM] * kk[e];
|
||||
sq += sS[dd + e * HEAD_DIM] * qq[e];
|
||||
}
|
||||
const float vnew = bv * (vv[dd] - sk * dv);
|
||||
out_t[dd] = sq * dv * scale + vnew * sc;
|
||||
for (int cc = 0; cc < HEAD_DIM; ++cc) {
|
||||
float s = dv * sS[dd + cc * HEAD_DIM] + vnew * kk[cc];
|
||||
sS[dd + cc * HEAD_DIM] = fminf(fmaxf(s, -1e6f), 1e6f);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads(); // S read/write hazard across threads per token
|
||||
}
|
||||
} else {
|
||||
// ---- fast path outputs: KQ rows recomputed per step (no storage),
|
||||
// aq recomputed per owned element, ci recomputed per thread (128
|
||||
// MACs; cheaper than a broadcast round-trip).
|
||||
for (int i = 0; i < C; ++i) {
|
||||
// KQ[i][j] = kh_j . qt_i, j<i; thread j (idle beyond i)
|
||||
if (tid < i) {
|
||||
const int j = tid;
|
||||
float kq = 0.0f;
|
||||
for (int e = 0; e < HEAD_DIM; ++e) {
|
||||
kq += k_base[(t0 + j) * qkv_stride_token + e]
|
||||
* q_base[(t0 + i) * qkv_stride_token + e] * scale;
|
||||
}
|
||||
// NOTE: qt carries *scale already (matches candidate C)
|
||||
sRow[j] = kq;
|
||||
}
|
||||
__syncthreads();
|
||||
{
|
||||
// ci = kh_i . qt_i (recomputed per thread, no broadcast).
|
||||
// NOTE: pim/logPim/D[*] all come from the staged sLogP prefix
|
||||
// sums — bitwise identical to candidate C in the harness.
|
||||
float ci = 0.0f;
|
||||
for (int e = 0; e < HEAD_DIM; ++e) {
|
||||
ci += k_base[(t0 + i) * qkv_stride_token + e]
|
||||
* q_base[(t0 + i) * qkv_stride_token + e] * scale;
|
||||
}
|
||||
const float pim = i ? expf(sLogP[i - 1]) : 1.0f;
|
||||
const float dv = sDv[i];
|
||||
float * out_t = out_base + (t0 + i) * out_token_stride;
|
||||
for (int dd = tid; dd < HEAD_DIM; dd += nthreads) {
|
||||
float aq = 0.0f;
|
||||
for (int e = 0; e < HEAD_DIM; ++e) {
|
||||
aq += sS[dd + e * HEAD_DIM] * q_base[(t0 + i) * qkv_stride_token + e] * scale;
|
||||
}
|
||||
float a = pim * aq;
|
||||
for (int j = 0; j < i; ++j) {
|
||||
a += expf(sLogP[i - 1] - sLogP[j]) * sRow[j] * sE[j * HEAD_DIM + dd];
|
||||
}
|
||||
out_t[dd] = dv * a + ci * sE[i * HEAD_DIM + dd];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
// ---- GEMM assembly of S_C (single-GEMM form; valid: no per-step
|
||||
// states needed since dispatch guarantees saved_states == NULL).
|
||||
// S_C = P[C-1]*S_in + sum_i D[C-1][i] e_i kh_i^T, clamped.
|
||||
{
|
||||
const float pC = expf(sLogP[C - 1]);
|
||||
for (int64_t x = tid; x < (int64_t)HEAD_DIM * HEAD_DIM; x += nthreads) {
|
||||
const int row = (int)(x % HEAD_DIM);
|
||||
const int col = (int)(x / HEAD_DIM);
|
||||
float s = pC * sS[x];
|
||||
for (int i = 0; i < C; ++i) {
|
||||
s += expf(sLogP[C - 1] - sLogP[i]) * sE[i * HEAD_DIM + row]
|
||||
* k_base[(t0 + i) * qkv_stride_token + col];
|
||||
}
|
||||
sS[x] = fminf(fmaxf(s, -1e6f), 1e6f);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// write back final state
|
||||
for (int64_t x = tid; x < (int64_t)HEAD_DIM * HEAD_DIM; x += nthreads) state_dst[x] = sS[x];
|
||||
}
|
||||
|
||||
static void delta_net_f32_cuda(
|
||||
const float * q,
|
||||
const float * k,
|
||||
|
|
@ -209,6 +470,58 @@ static void delta_net_f32_cuda(
|
|||
GGML_ABORT("Unsupported delta net head size");
|
||||
}
|
||||
|
||||
// Chunked WY path (opt-in via DELTA_WY_CUDA=1; validated math in
|
||||
// tests/test-delta-chunk.cpp). Requirements: long enough prefill for at
|
||||
// least 2 chunks, no per-step checkpoints (GEMM assembly form), and enough
|
||||
// dynamic shared memory (else transparent fallback to sequential below).
|
||||
// Falls back silently-exact in every other case: default behavior unchanged.
|
||||
{
|
||||
const char * wy_env = getenv("DELTA_WY_CUDA");
|
||||
const bool wy_wanted = wy_env && wy_env[0] == '1';
|
||||
if (wy_wanted && n_tokens > DELTA_WY_CHUNK && saved_states == nullptr) {
|
||||
int chunk = DELTA_WY_CHUNK;
|
||||
size_t need = ((size_t)head_dim * head_dim + (size_t)chunk * head_dim + 4 * chunk) * sizeof(float);
|
||||
// shared tier selection: full chunks, then half chunks, then sequential.
|
||||
// (cudaFuncSetAttribute form mirrors fattn-new-mma.cu usage.)
|
||||
auto wy_attr_ok = [&](size_t bytes) -> bool {
|
||||
cudaError_t err = (head_dim == 64)
|
||||
? cudaFuncSetAttribute(delta_net_chunked_f32<64>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, (int)bytes)
|
||||
: cudaFuncSetAttribute(delta_net_chunked_f32<128>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, (int)bytes);
|
||||
return err == cudaSuccess;
|
||||
};
|
||||
if (!wy_attr_ok(need)) {
|
||||
chunk = DELTA_WY_CHUNK_SMALL;
|
||||
need = ((size_t)head_dim * head_dim + (size_t)chunk * head_dim + 4 * chunk) * sizeof(float);
|
||||
if (!wy_attr_ok(need)) {
|
||||
chunk = 0;
|
||||
}
|
||||
}
|
||||
if (chunk > 0) {
|
||||
static bool wy_note_done = false;
|
||||
if (!wy_note_done) {
|
||||
wy_note_done = true;
|
||||
fprintf(stderr, "delta_net: chunked WY path, C=%d, shared %zu bytes\n", chunk, need);
|
||||
}
|
||||
constexpr int wy_threads = 256;
|
||||
const int wy_blocks = (int)(n_seqs * n_heads);
|
||||
if (head_dim == 64) {
|
||||
delta_net_chunked_f32<64><<<wy_blocks, wy_threads, need, stream>>>(
|
||||
q, k, v, g, beta, state_in, dst, state_out,
|
||||
n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, chunk);
|
||||
} else {
|
||||
delta_net_chunked_f32<128><<<wy_blocks, wy_threads, need, stream>>>(
|
||||
q, k, v, g, beta, state_in, dst, state_out,
|
||||
n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, chunk);
|
||||
}
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
return;
|
||||
}
|
||||
// else: fall through to the sequential kernel below
|
||||
}
|
||||
}
|
||||
|
||||
GGML_ASSERT(head_dim % WARP_SIZE == 0);
|
||||
const int num_blocks = n_seqs * n_heads * (head_dim/WARP_SIZE);
|
||||
const size_t smem_size = 2 * head_dim * sizeof(float);
|
||||
|
|
|
|||
|
|
@ -420,12 +420,14 @@ static void solve_chunk_wy(
|
|||
const int C = c1 - c0;
|
||||
const float scale = 1.0f / sqrtf((float)hd);
|
||||
// per-token ingredients (same factoring as run_token)
|
||||
std::vector<float> kh(C * hd), qt(C * hd), bv(C), dv(C), pv(C), ck(C * hd), aq(C * hd);
|
||||
// Absolute decay products are ALWAYS exp(logP[.]) (never progressive
|
||||
// products): bitwise identical to what the CUDA kernel stages in shared.
|
||||
std::vector<float> kh(C * hd), qt(C * hd), bv(C), dv(C), ck(C * hd), aq(C * hd);
|
||||
// log-decay prefix sums: every decay RATIO is evaluated as exp(logP[i]-logP[j]).
|
||||
// Raw products P[i] underflow to 0 over 64 steps with decay<1, turning later
|
||||
// P[i]/P[j] ratios into 0/0 = nan. exp-of-difference is exact-or-zero, which
|
||||
// is the mathematically right answer (true ratio ~0). Absolute P[i] (=pv[i])
|
||||
// may still underflow to 0 in f_i/S_C/output terms — also correct there.
|
||||
// is the mathematically right answer (true ratio ~0). Absolute P[i]
|
||||
// (=exp(logP[i])) may still underflow to 0 in f_i/S_C/output terms — also correct there.
|
||||
std::vector<float> logP(C);
|
||||
for (int i = 0; i < C; ++i) {
|
||||
const int tt = c0 + i;
|
||||
|
|
@ -441,7 +443,6 @@ static void solve_chunk_wy(
|
|||
kn = 1.0f / sqrtf(kn + 1e-12f);
|
||||
bv[i] = 1.0f / (1.0f + expf(-boff));
|
||||
dv[i] = expf(fminf(goff, 50.0f));
|
||||
pv[i] = dv[i] * (i ? pv[i - 1] : 1.0f);
|
||||
logP[i] = fminf(goff, 50.0f) + (i ? logP[i - 1] : 0.0f);
|
||||
for (int d = 0; d < hd; ++d) {
|
||||
kh[i * hd + d] = k[d] * kn;
|
||||
|
|
@ -479,7 +480,7 @@ static void solve_chunk_wy(
|
|||
const int tt = c0 + i;
|
||||
const float * v = t.v.data() + (((size_t)b * c.hv + h) * c.nt + tt) * hd;
|
||||
for (int d = 0; d < hd; ++d) {
|
||||
float e = bv[i] * v[d] - bv[i] * pv[i] * ck[i * hd + d];
|
||||
float e = bv[i] * v[d] - bv[i] * expf(logP[i]) * ck[i * hd + d];
|
||||
for (int j = 0; j < i; ++j) e -= L[i * C + j] * E[j * hd + d];
|
||||
E[i * hd + d] = e;
|
||||
const float ae = fabsf(e);
|
||||
|
|
@ -511,7 +512,7 @@ static void solve_chunk_wy(
|
|||
memcpy(S.data(), s_in, stsz * sizeof(float));
|
||||
for (int i = 0; i < C; ++i) {
|
||||
const int tt = c0 + i;
|
||||
const float pim = i ? pv[i - 1] : 1.0f;
|
||||
const float pim = i ? expf(logP[i - 1]) : 1.0f;
|
||||
const float logPim = i ? logP[i - 1] : 0.0f;
|
||||
float ci = 0.0f;
|
||||
for (int d = 0; d < hd; ++d) ci += kh[i * hd + d] * qt[i * hd + d];
|
||||
|
|
@ -540,7 +541,7 @@ static void solve_chunk_wy(
|
|||
float md = 0.0f;
|
||||
for (int row = 0; row < hd; ++row) {
|
||||
for (int col = 0; col < hd; ++col) {
|
||||
float s = pv[C - 1] * s_in[row + col * hd];
|
||||
float s = expf(logP[C - 1]) * s_in[row + col * hd];
|
||||
for (int i = 0; i < C; ++i) {
|
||||
s += expf(logP[C - 1] - logP[i]) * E[i * hd + row] * kh[i * hd + col];
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue