diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu index 2f66798379..2ee410ebdc 100644 --- a/zig/kernels/cuda/flashnext_attn.cu +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -53,6 +53,19 @@ __device__ __forceinline__ int partner(int d, int half) { return d < half ? d + half : (d < 2 * half ? d - half : d); } +// A stream's cache pages (cuda_pages.zig's layout): pt[i] holds positions [i * ptok, (i + 1) * ptok); null: contiguous. +struct Pg { + const int* pt; + long long ps; + int ptok, pad; +}; + +// Element offset of item i (a position, or a block with per = ptok / RATIO) of `row` elements, contiguous or paged. +__device__ __forceinline__ long long pg_at(const Pg& pg, int i, int per, long long row) { + if (pg.pt == nullptr) return static_cast(i) * row; + return static_cast(pg.pt[i / per]) * pg.ps + static_cast(i % per) * row; +} + } // namespace fn_attn using namespace fn_attn; @@ -70,13 +83,13 @@ __device__ __forceinline__ void fn_attn_prep_body(const __nv_bfloat16* __restric __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, - int NI, int IHD, int HALF, const int* rope3, int n, int delta) { + int NI, int IHD, int HALF, const int* rope3, int n, int delta, Pg pg) { __shared__ float scratch[32]; const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; const int pos = POS0[0] + r; const long long row = static_cast(r) * PW; if (head >= NQ + NKV + NI) { - if (d < IHD) IKC[static_cast(pos) * IHD + d] = P[row + NQ * 2 * HD + 2 * NKV * HD + NI * IHD + d]; + if (d < IHD) IKC[pg_at(pg, pos, pg.ptok, IHD) + d] = P[row + NQ * 2 * HD + 2 * NKV * HD + NI * IHD + d]; return; } const bool index = head >= NQ + NKV; @@ -94,9 +107,9 @@ __device__ __forceinline__ void fn_attn_prep_body(const __nv_bfloat16* __restric if (head < NQ) { Q[(static_cast(r) * NQ + head) * HD + d] = out; } else if (!index) { - const long long slot = static_cast(pos) * NKV + head - NQ; - KC[slot * HD + d] = out; - VC[slot * HD + d] = P[row + NQ * 2 * HD + NKV * HD + (head - NQ) * HD + d]; + const long long slot = pg_at(pg, pos, pg.ptok, static_cast(NKV) * HD) + (head - NQ) * HD; + KC[slot + d] = out; + VC[slot + d] = P[row + NQ * 2 * HD + NKV * HD + (head - NQ) * HD + d]; } else { IQ[(static_cast(r) * NI + head - NQ - NKV) * IHD + d] = out; } @@ -108,8 +121,8 @@ extern "C" __global__ void __launch_bounds__(256) fn_attn_prep(const __nv_bfloat __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, - int NI, int IHD, int HALF) { - fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, nullptr, 0, 0); + int NI, int IHD, int HALF, Pg pg) { + fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, nullptr, 0, 0, pg); } extern "C" __global__ void __launch_bounds__(256) fn_attn_prep_media(const __nv_bfloat16* __restrict__ P, const int* __restrict__ POS0, @@ -118,8 +131,8 @@ extern "C" __global__ void __launch_bounds__(256) fn_attn_prep_media(const __nv_ __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, - int NI, int IHD, int HALF, const int* rope3, int n, int delta) { - fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, rope3, n, delta); + int NI, int IHD, int HALF, const int* rope3, int n, int delta, Pg pg) { + fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, rope3, n, delta, pg); } // glue._attn_gate, program (row, head) of HD threads: bf16(o / (1 + e^-gate)) and its 32-group sums. @@ -138,41 +151,42 @@ extern "C" __global__ void __launch_bounds__(256) fn_attn_gate(const __nv_bfloat __device__ __forceinline__ void fn_qsa_pool_body(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, float eps, int R, int DI, int HALF, - int RATIO, const int* rope3, int n, int delta) { + int RATIO, const int* rope3, int n, int delta, Pg pg) { __shared__ float scratch[32]; __shared__ float xs[256]; const int p0 = POS0[0]; const int b = p0 / RATIO + blockIdx.x, d = threadIdx.x; if (RATIO * b + RATIO > p0 + R) return; - float acc = bf(IKC[static_cast(RATIO) * b * DI + d]); - for (int k = 1; k < RATIO; ++k) acc = __fadd_rn(acc, bf(IKC[(static_cast(RATIO) * b + k) * DI + d])); + const long long at = pg_at(pg, RATIO * b, pg.ptok, DI) + d; // a block's RATIO rows share a page + float acc = bf(IKC[at]); + for (int k = 1; k < RATIO; ++k) acc = __fadd_rn(acc, bf(IKC[at + static_cast(k) * DI])); const float x = rbf(__fdiv_rn(acc, static_cast(RATIO))); xs[d] = x; const float rinv = inv_rms(block_sum(__fmul_rn(x, x), scratch), DI, eps); const int p = partner(d, HALF); const float xn = rbf(__fmul_rn(__fmul_rn(x, rinv), W[d])); const float xpn = rbf(__fmul_rn(__fmul_rn(xs[p], rinv), W[p])); - POOLED[static_cast(b) * DI + d] = __float2bfloat16_rn(rope(xn, xpn, d, HALF, media_pos(RATIO * b, d, HALF, rope3, n, delta), INV)); + POOLED[pg_at(pg, b, pg.ptok / RATIO, DI) + d] = __float2bfloat16_rn(rope(xn, xpn, d, HALF, media_pos(RATIO * b, d, HALF, rope3, n, delta), INV)); } extern "C" __global__ void __launch_bounds__(128) fn_qsa_pool(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, float eps, int R, int DI, int HALF, - int RATIO) { - fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, nullptr, 0, 0); + int RATIO, Pg pg) { + fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, nullptr, 0, 0, pg); } extern "C" __global__ void __launch_bounds__(128) fn_qsa_pool_media(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, float eps, int R, int DI, int HALF, - int RATIO, const int* rope3, int n, int delta) { - fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, rope3, n, delta); + int RATIO, const int* rope3, int n, int delta, Pg pg) { + fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, rope3, n, delta, pg); } // attention._scores, program (row, 64 blocks): the heads' sum of relu(q . k) / sqrt(DI), rows past TOP (DI 128). extern "C" __global__ void __launch_bounds__(256) fn_qsa_scores(const __nv_bfloat16* __restrict__ IQ, const __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, float* __restrict__ SC, int NB, int HI, - int RATIO, int TOP) { + int RATIO, int TOP, Pg pg) { constexpr int DI = 128; __shared__ float qs[8][DI]; const int r = blockIdx.x, j = blockIdx.y; @@ -186,7 +200,7 @@ extern "C" __global__ void __launch_bounds__(256) fn_qsa_scores(const __nv_bfloa #pragma unroll for (int i = 0; i < 8; ++i) { const int b = j * 64 + warp * 8 + i; - raw[i] = b < complete ? *reinterpret_cast(POOLED + static_cast(b) * DI + 4 * lane) : make_uint2(0, 0); + raw[i] = b < complete ? *reinterpret_cast(POOLED + pg_at(pg, b, pg.ptok / RATIO, DI) + 4 * lane) : make_uint2(0, 0); } #pragma unroll for (int i = 0; i < 8; ++i) { @@ -210,7 +224,7 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, - int H, int HK, int G, int CH, int NCH, float SCALE, int IDW, int QSA) { + int H, int HK, int G, int CH, int NCH, float SCALE, int IDW, int QSA, Pg pg) { extern __shared__ __align__(16) unsigned char smem_raw[]; typedef __nv_bfloat16 Row[HD + 8]; Row* qs = reinterpret_cast(smem_raw); // [QR] the row's query heads @@ -218,6 +232,7 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, __nv_bfloat16 (*ps)[TK + 8] = reinterpret_cast<__nv_bfloat16 (*)[TK + 8]>(tiles_kv + NB * 2 * TK); // [QR] bf16 p __shared__ float red_m[4][QR], red_l[4][QR]; __shared__ int key_ids[NB][TK]; + __shared__ long long key_at[NB][TK]; // each key's row offset, its page looked up once a tile const int r = blockIdx.x, hk = blockIdx.y, c = blockIdx.z; const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; int n = POS0[0] + r + 1; @@ -236,15 +251,17 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, auto load = [&](int t0, int buf) { if (threadIdx.x < TK) { const int ki = start + t0 + threadIdx.x; - key_ids[buf][threadIdx.x] = ki < n ? (sparse ? IDS[static_cast(r) * IDW + ki] : ki) : -1; + const int key = ki < n ? (sparse ? IDS[static_cast(r) * IDW + ki] : ki) : -1; + key_ids[buf][threadIdx.x] = key; + key_at[buf][threadIdx.x] = pg_at(pg, max(key, 0), pg.ptok, static_cast(HK) * HD) + hk * HD; } __syncthreads(); for (int pass = 0; pass < 2; ++pass) { const __nv_bfloat16* src = pass == 0 ? KC : VC; Row* dst = tiles_kv + (buf * 2 + pass) * TK; for (int e = threadIdx.x; e < TK * HD / 8; e += blockDim.x) { - const int k = e / (HD / 8), v = e % (HD / 8), key = key_ids[buf][k]; - qmm_frag::cp16z(&dst[k][8 * v], src + (static_cast(max(key, 0)) * HK + hk) * HD + 8 * v, key >= 0); + const int k = e / (HD / 8), v = e % (HD / 8); + qmm_frag::cp16z(&dst[k][8 * v], src + key_at[buf][k] + 8 * v, key_ids[buf][k] >= 0); } qmm_frag::commit(); } @@ -368,8 +385,8 @@ extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks(const __nv_bflo float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, - float SCALE, int IDW, int QSA) { - attn_chunks<1>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); + float SCALE, int IDW, int QSA, Pg pg) { + attn_chunks<1>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA, pg); } // attn_chunks with the next tile's loads in flight (a decode round's few rows): the same sums. @@ -378,8 +395,8 @@ extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks_ahead(const __n float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, - float SCALE, int IDW, int QSA) { - attn_chunks<2>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); + float SCALE, int IDW, int QSA, Pg pg) { + attn_chunks<2>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA, pg); } // attention._merge, program (row, head) of HD threads: the head's chunk partials merged in chunk order, o / l. diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 103b766320..2b657e642c 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -394,10 +394,9 @@ extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfl struct StoreBf16 { int layer, kvh; __device__ __forceinline__ void operator()(const Seg& sg, int h, int slot, int d, __nv_bfloat16 k, __nv_bfloat16 v) const { - __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * kvh * sg.cap * HD; - const size_t dst = ((size_t)h * sg.cap + slot) * HD + d; - keys[dst] = k; - keys[(size_t)kvh * sg.cap * HD + dst] = v; + __nv_bfloat16* kv = reinterpret_cast<__nv_bfloat16*>(sg.kv); + kv[kv_row(sg, kvh, layer, 0, h, slot) + d] = k; + kv[kv_row(sg, kvh, layer, 1, h, slot) + d] = v; } }; @@ -434,8 +433,7 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r } return; } - const __nv_bfloat16* K = reinterpret_cast(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; - const __nv_bfloat16* V = K + (size_t)KVH * sg.cap * HD; + const __nv_bfloat16* KV = reinterpret_cast(sg.kv); float q[DPL], o[DPL]; #pragma unroll for (int i = 0; i < DPL; i++) { @@ -443,22 +441,32 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (int t = k0; t < k1; t++) { - const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); - const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); - const __nv_bfloat16* kb = reinterpret_cast(&kw); - float dot = 0.0f; -#pragma unroll - for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(kb[i]), dot); - const float s = warp_sum(dot) * scale; - const float mn = fmaxf(mx, s); - const float a = __expf(mx - mn), b = __expf(s - mn); - l = l * a + b; - const uint4 vw = *reinterpret_cast(V + ((size_t)h * sg.cap + pos) * HD + lane * DPL); - const __nv_bfloat16* vb = reinterpret_cast(&vw); -#pragma unroll - for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); - mx = mn; + // the keys in order; a key's and value's rows by `row` (one branch here, not one a key) + const auto scan = [&](auto row) { + for (int t = k0; t < k1; t++) { + const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); + const uint4 kw = *reinterpret_cast(row(0, pos) + lane * DPL); + const __nv_bfloat16* kb = reinterpret_cast(&kw); + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(kb[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const uint4 vw = *reinterpret_cast(row(1, pos) + lane * DPL); + const __nv_bfloat16* vb = reinterpret_cast(&vw); +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); + mx = mn; + } + }; + if (sg.pt == 0) { + const __nv_bfloat16* K = KV + kv_row(sg, KVH, layer, 0, h, 0); + const __nv_bfloat16* V = KV + kv_row(sg, KVH, layer, 1, h, 0); + scan([&](int s, int pos) { return (s == 0 ? K : V) + (size_t)pos * HD; }); + } else { + scan([&](int s, int pos) { return KV + kv_row(sg, KVH, layer, s, h, pos); }); } if (lane == 0) { PM[slot] = mx; diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh index 728614bd49..0615636caa 100644 --- a/zig/kernels/cuda/qwen35.cuh +++ b/zig/kernels/cuda/qwen35.cuh @@ -34,8 +34,19 @@ struct Seg { unsigned long long before; // exact pre-round state for a one-row in-place update int cap, len, row0, rows, tap_rows, taps; int delta; // an image prompt's rotary shift: rows past it rotate at position + delta + unsigned long long pt; // the stream's KV page table (int a page; 0: kv holds cap contiguous rows) + unsigned long long ps; // elements a KV page holds (bf16, or bytes for fp8): every attention layer's ptok rows + int ptok; // positions a KV page holds }; +// Seg.kv's element of layer a's key (s 0) or value (s 1) at pos, head h: cap rows a head, or pos's page's row. +__device__ __forceinline__ size_t kv_row(const Seg& sg, int kvh, int a, int s, int h, int pos) { + const size_t blk = ((size_t)a * 2 + s) * kvh + h; + if (sg.pt == 0) return (blk * sg.cap + pos) * HD; + const int* pt = reinterpret_cast(sg.pt); + return (size_t)pt[pos / sg.ptok] * sg.ps + (blk * sg.ptok + pos % sg.ptok) * HD; +} + // A window row's record of one linear layer (Seg.log). struct LogRow { float* g; // [nv] decay diff --git a/zig/kernels/cuda/qwen35_keep.cu b/zig/kernels/cuda/qwen35_keep.cu index 20d00fc212..6bca4e89e0 100644 --- a/zig/kernels/cuda/qwen35_keep.cu +++ b/zig/kernels/cuda/qwen35_keep.cu @@ -44,14 +44,13 @@ extern "C" __global__ void tf_q35k_conv(Geo geo, Seg sg, int last) { extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact(Geo geo, Seg sg, Path path) { const int KVH = geo.kvh; const int layer = blockIdx.x, h = blockIdx.y; - __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + ((size_t)layer * 2 * KVH + h) * sg.cap * HD; - __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + __nv_bfloat16* kv = reinterpret_cast<__nv_bfloat16*>(sg.kv); const int d = threadIdx.x; for (int j = 0; j < path.n; j++) { const int from = sg.len + path.rows[j], to = sg.len + j; if (from != to) { - keys[(size_t)to * HD + d] = keys[(size_t)from * HD + d]; - values[(size_t)to * HD + d] = values[(size_t)from * HD + d]; + kv[kv_row(sg, KVH, layer, 0, h, to) + d] = kv[kv_row(sg, KVH, layer, 0, h, from) + d]; + kv[kv_row(sg, KVH, layer, 1, h, to) + d] = kv[kv_row(sg, KVH, layer, 1, h, from) + d]; } __syncthreads(); } diff --git a/zig/kernels/cuda/qwen35_kv8.cu b/zig/kernels/cuda/qwen35_kv8.cu index ed7ccdfab1..3069f149b1 100644 --- a/zig/kernels/cuda/qwen35_kv8.cu +++ b/zig/kernels/cuda/qwen35_kv8.cu @@ -6,9 +6,26 @@ using namespace tf_q35; -// A layer's block in Seg.kv: kvh * cap rows of HD bytes for keys then values, then kvh * cap fp32 scales for each. -__device__ __forceinline__ unsigned char* layer_at(const Seg& sg, int layer, int kvh) { - return reinterpret_cast(sg.kv) + (size_t)layer * 2 * kvh * sg.cap * (HD + 4); +// A key's and value's codes and scales. +struct Row8 { + unsigned char *k, *v; + float *ks, *vs; +}; + +// Codes and scales at (layer, h, pos): a layer's block of n = cap rows a head, or of ptok rows in pos's page. +__device__ __forceinline__ Row8 row8(const Seg& sg, int layer, int kvh, int h, int pos) { + unsigned char* base = reinterpret_cast(sg.kv); + int n = sg.cap, p = pos; + if (sg.pt != 0) { + base += (size_t)reinterpret_cast(sg.pt)[pos / sg.ptok] * sg.ps; + n = sg.ptok; + p = pos % sg.ptok; + } + unsigned char* keys = base + (size_t)layer * 2 * kvh * n * (HD + 4); + unsigned char* values = keys + (size_t)kvh * n * HD; + float* kscale = reinterpret_cast(values + (size_t)kvh * n * HD); + const size_t row = (size_t)h * n + p; + return {keys + row * HD, values + row * HD, kscale + row, kscale + (size_t)kvh * n + row}; } // The block's 256 threads' row as e4m3 into q[d], its scale (1 for a row of zeros) to *s; tf_kv_quant's bits. @@ -32,13 +49,9 @@ struct StoreFp8 { int layer, kvh; float* red; __device__ __forceinline__ void operator()(const Seg& sg, int h, int slot, int d, __nv_bfloat16 k, __nv_bfloat16 v) const { - unsigned char* keys = layer_at(sg, layer, kvh); - unsigned char* values = keys + (size_t)kvh * sg.cap * HD; - float* kscale = reinterpret_cast(values + (size_t)kvh * sg.cap * HD); - float* vscale = kscale + (size_t)kvh * sg.cap; - const size_t row = (size_t)h * sg.cap + slot; - quant_row(bf(k), d, keys + row * HD, kscale + row, red); - quant_row(bf(v), d, values + row * HD, vscale + row, red); + const Row8 r = row8(sg, layer, kvh, h, slot); + quant_row(bf(k), d, r.k, r.ks, red); + quant_row(bf(v), d, r.v, r.vs, red); } }; @@ -86,10 +99,6 @@ extern "C" __global__ void tf_q35_attn_partial8(Geo geo, const __nv_bfloat16* __ } return; } - const unsigned char* K = layer_at(sg, layer, KVH); - const unsigned char* V = K + (size_t)KVH * sg.cap * HD; - const float* KS = reinterpret_cast(V + (size_t)KVH * sg.cap * HD); - const float* VS = KS + (size_t)KVH * sg.cap; float q[DPL], o[DPL]; #pragma unroll for (int i = 0; i < DPL; i++) { @@ -97,23 +106,32 @@ extern "C" __global__ void tf_q35_attn_partial8(Geo geo, const __nv_bfloat16* __ o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (int t = k0; t < k1; t++) { - const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); - const size_t row = (size_t)h * sg.cap + pos; - float kf[DPL], vf[DPL]; - widen8(*reinterpret_cast(K + row * HD + lane * DPL), kf); - float dot = 0.0f; + // the keys in order; a row's codes and scales by `at` (one branch here, not one a key) + const auto scan = [&](auto at) { + for (int t = k0; t < k1; t++) { + const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); + const Row8 r = at(pos); + float kf[DPL], vf[DPL]; + widen8(*reinterpret_cast(r.k + lane * DPL), kf); + float dot = 0.0f; #pragma unroll - for (int i = 0; i < DPL; i++) dot = fmaf(q[i], kf[i], dot); - const float s = warp_sum(dot) * KS[row] * scale; - const float mn = fmaxf(mx, s); - const float a = __expf(mx - mn), b = __expf(s - mn); - l = l * a + b; - widen8(*reinterpret_cast(V + row * HD + lane * DPL), vf); - const float bv = b * VS[row]; + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], kf[i], dot); + const float s = warp_sum(dot) * *r.ks * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + widen8(*reinterpret_cast(r.v + lane * DPL), vf); + const float bv = b * *r.vs; #pragma unroll - for (int i = 0; i < DPL; i++) o[i] = fmaf(bv, vf[i], o[i] * a); - mx = mn; + for (int i = 0; i < DPL; i++) o[i] = fmaf(bv, vf[i], o[i] * a); + mx = mn; + } + }; + if (sg.pt == 0) { + const Row8 r0 = row8(sg, layer, KVH, h, 0); // row p is row 0's plus p rows + scan([&](int pos) { return Row8{r0.k + (size_t)pos * HD, r0.v + (size_t)pos * HD, r0.ks + pos, r0.vs + pos}; }); + } else { + scan([&](int pos) { return row8(sg, layer, KVH, h, pos); }); } if (lane == 0) { PM[slot] = mx; @@ -127,20 +145,16 @@ extern "C" __global__ void tf_q35_attn_partial8(Geo geo, const __nv_bfloat16* __ extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact8(Geo geo, Seg sg, Path path) { const int KVH = geo.kvh; const int layer = blockIdx.x, h = blockIdx.y; - unsigned char* keys = layer_at(sg, layer, KVH); - unsigned char* values = keys + (size_t)KVH * sg.cap * HD; - float* kscale = reinterpret_cast(values + (size_t)KVH * sg.cap * HD); - float* vscale = kscale + (size_t)KVH * sg.cap; - const size_t base = (size_t)h * sg.cap; const int d = threadIdx.x; for (int j = 0; j < path.n; j++) { const int from = sg.len + path.rows[j], to = sg.len + j; if (from != to) { - keys[(base + to) * HD + d] = keys[(base + from) * HD + d]; - values[(base + to) * HD + d] = values[(base + from) * HD + d]; + const Row8 a = row8(sg, layer, KVH, h, from), b = row8(sg, layer, KVH, h, to); + b.k[d] = a.k[d]; + b.v[d] = a.v[d]; if (d == 0) { - kscale[base + to] = kscale[base + from]; - vscale[base + to] = vscale[base + from]; + *b.ks = *a.ks; + *b.vs = *a.vs; } } __syncthreads(); diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..f1c855e47f 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -20,11 +20,11 @@ const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "flashnext_gdn flashnext_gdn_chain", .symbol = "flashnext_gdn_chain", .image = "flashnext_gdn", .threads = 1024 }, .usage = .{ .{ .regs = 51, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 62, .static = 2060, .bound = 1024 }, .{ .regs = 62, .static = 2060, .bound = 1024 } } }, .{ .launch = .{ .name = "flashnext_hc fn_add_streams", .symbol = "fn_add_streams", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = .{ .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 71, .static = 768, .bound = 128 }, .{ .regs = 71, .static = 768, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = .{ .{ .regs = 71, .static = 896, .bound = 128 }, .{ .regs = 71, .static = 896, .bound = 128 }, .{ .regs = 71, .static = 896, .bound = 128 }, .{ .regs = 73, .static = 896, .bound = 128 }, .{ .regs = 73, .static = 896, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 1280, .bound = 128 }, .{ .regs = 72, .static = 1280, .bound = 128 }, .{ .regs = 72, .static = 1280, .bound = 128 }, .{ .regs = 71, .static = 1280, .bound = 128 }, .{ .regs = 71, .static = 1280, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_gate", .symbol = "fn_attn_gate", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_merge", .symbol = "fn_attn_merge", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_batch fn_gdn_front", .symbol = "fn_gdn_front", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 32, .static = 1024, .bound = 256 }, .{ .regs = 32, .static = 1024, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_batch fn_gdn_norm", .symbol = "fn_gdn_norm", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 24, .static = 516, .bound = 128 }, .{ .regs = 26, .static = 516, .bound = 128 }, .{ .regs = 26, .static = 516, .bound = 128 }, .{ .regs = 24, .static = 516, .bound = 128 }, .{ .regs = 24, .static = 516, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows", .symbol = "fn_gdn_rows", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, @@ -39,7 +39,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_slices", .symbol = "fn_hc_writeback_slices", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 256, .bound = 256 }, .{ .regs = 38, .static = 256, .bound = 256 }, .{ .regs = 38, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_hc fn_ple_conv", .symbol = "fn_ple_conv", .image = "flashnext_hc", .threads = 512 }, .usage = .{ .{ .regs = 23, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 } } }, .{ .launch = .{ .name = "flashnext_hc fn_ple_gate", .symbol = "fn_ple_gate", .image = "flashnext_hc", .threads = 512 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 35, .static = 128, .bound = 512 }, .{ .regs = 35, .static = 128, .bound = 512 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool", .symbol = "fn_qsa_pool", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool", .symbol = "fn_qsa_pool", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_attn fn_qsa_scores", .symbol = "fn_qsa_scores", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 46, .static = 4096, .bound = 256 }, .{ .regs = 46, .static = 4096, .bound = 256 }, .{ .regs = 48, .static = 4096, .bound = 256 }, .{ .regs = 48, .static = 4096, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_batch fn_qsa_select", .symbol = "fn_qsa_select", .image = "flashnext_batch", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 44, .static = 1160, .bound = 1024 }, .{ .regs = 40, .static = 1160, .bound = 1024 } } }, .{ .launch = .{ .name = "flashnext_mtp fn_row_best", .symbol = "fn_row_best", .image = "flashnext_mtp", .threads = 1024 }, .usage = .{ .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 } } }, @@ -129,7 +129,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "qwen_vision tf_qv_bias_cast", .symbol = "tf_qv_bias_cast", .image = "qwen_vision", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "qwen_vision tf_qv_splice", .symbol = "tf_qv_splice", .image = "qwen_vision", .threads = 256 }, .usage = .{ .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_prep_media", .symbol = "fn_attn_prep_media", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool_media", .symbol = "fn_qsa_pool_media", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool_media", .symbol = "fn_qsa_pool_media", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 } } }, // EXL3: the linear's instantiations at their most rows' shared bytes, the routed experts, Dense's tile candidate .{ .launch = .{ .name = "exl3_linear tf_exl3_rot_in", .symbol = "tf_exl3_rot_in", .image = "exl3_linear", .threads = 128 }, .usage = .{ .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 22, .static = 0, .bound = 128 }, .{ .regs = 22, .static = 0, .bound = 128 }, .{ .regs = 22, .static = 0, .bound = 128 }, .{ .regs = 22, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "exl3_linear tf_exl3_linear_2_2_4", .symbol = "tf_exl3_linear_2_2_4", .image = "exl3_linear", .threads = 128, .dynamic = 16384 }, .usage = .{ .{ .regs = 168, .static = 16, .bound = 128 }, .{ .regs = 168, .static = 16, .bound = 128 }, .{ .regs = 168, .static = 16, .bound = 128 }, .{ .regs = 152, .static = 16, .bound = 128 }, .{ .regs = 152, .static = 16, .bound = 128 } } }, diff --git a/zig/src/core/kv_pool.zig b/zig/src/core/kv_pool.zig index 8e9d55c923..8f64bcac05 100644 --- a/zig/src/core/kv_pool.zig +++ b/zig/src/core/kv_pool.zig @@ -12,6 +12,8 @@ pub const Layout = struct { page_tokens: u32, /// Bytes of one kept snapshot's slot, outside the pages. snap_bytes: u64, + /// Bytes a running stream takes from the pages beside its keys and values, as one run of whole pages (0: none). + run_bytes: u64 = 0, pub fn pageBytes(l: Layout) u64 { return l.token_bytes * l.page_tokens; @@ -22,9 +24,14 @@ pub const Layout = struct { return @intCast((tokens + l.page_tokens - 1) / l.page_tokens); } - /// Pages a request reserves at admission: its prompt, its reply and one last verify window. + /// Pages a request reserves at admission: its prompt, its reply and one last verify window, and its run. pub fn need(l: Layout, prompt: u64, max_tokens: u64, window: u64) usize { - return l.pages(prompt + max_tokens + window); + return l.pages(prompt + max_tokens + window) + l.runPages(); + } + + /// Consecutive pages a running stream's `run_bytes` take. + pub fn runPages(l: Layout) usize { + return @intCast((l.run_bytes + l.pageBytes() - 1) / l.pageBytes()); } }; @@ -63,13 +70,13 @@ pub fn plan(l: Layout, b: Budget) error{ NoStream, TooMany }!Plan { if (b.engine > b.room or b.lanes == 0) return error.NoStream; var left = b.room - b.engine; const extra: u64 = b.lanes - 1; - const each = l.fixed_bytes + l.pageBytes(); + const each = l.fixed_bytes + (1 + l.runPages()) * l.pageBytes(); const fit = @min(extra, left / each); if (fit < extra and b.fixed) return error.TooMany; left -= fit * l.fixed_bytes; - const c = split(@min(b.cache, left - fit * l.pageBytes()), l.snap_bytes, b.snaps); + const c = split(@min(b.cache, left - fit * (1 + l.runPages()) * l.pageBytes()), l.snap_bytes, b.snaps); left -= c.snaps * l.snap_bytes + c.copied; - const most: u64 = if (b.window > 0) fit * l.pages(b.window) else std.math.maxInt(u64); + const most: u64 = if (b.window > 0) fit * (l.pages(b.window) + l.runPages()) else std.math.maxInt(u64); return .{ .lanes = @intCast(1 + fit), .snaps = c.snaps, .copied = c.copied, .pages = @intCast(@min(left / l.pageBytes(), most)) }; } @@ -77,7 +84,7 @@ pub fn plan(l: Layout, b: Budget) error{ NoStream, TooMany }!Plan { pub fn holds(l: Layout, b: Budget) bool { if (b.engine > b.room or b.lanes == 0) return false; const left = b.room - b.engine; - const each = l.fixed_bytes + l.pageBytes(); + const each = l.fixed_bytes + (1 + l.runPages()) * l.pageBytes(); return left - @min(@as(u64, b.lanes - 1), left / each) * each >= b.cache; } @@ -263,6 +270,49 @@ pub const Pool = struct { } } + /// The first id of `k` consecutive free pages, highest first (the free stack hands out low ids), or null. + pub fn findRun(p: *const Pool, k: usize) ?u32 { + if (k == 0 or k > p.refs.len) return null; + var end = p.refs.len; + while (end >= k) { + const start = end - k; + const busy = for (p.refs[start..end], start..) |r, i| { + if (r != 0) break i; + } else null; + if (busy) |i| end = i else return @intCast(start); + } + return null; + } + + /// `k` consecutive pages from `l`'s promise, held as one run until releaseRun; null: no such run is free now. + pub fn claimRun(p: *Pool, l: *Lease, k: usize) ?u32 { + if (l.promised < k) return null; + const start = p.findRun(k) orelse return null; + var n: usize = 0; + for (p.free[0..p.nfree]) |id| if (id < start or id >= start + k) { + p.free[n] = id; + n += 1; + }; + std.debug.assert(n + k == p.nfree); + p.nfree = n; + @memset(p.refs[start..][0..k], 1); + @memset(p.leased[start..][0..k], 1); // held as a lease holds its pages: never the store's to reclaim + l.promised -= k; + p.promised -= k; + return start; + } + + /// A run claimed at `start` goes back to the free pages. + pub fn releaseRun(p: *Pool, start: u32, k: usize) void { + for (start..start + k) |id| { + if (p.refs[id] != 1 or p.leased[id] != 1) std.debug.panic("kv pool: run page {d} released while not a run's", .{id}); + p.refs[id] = 0; + p.leased[id] = 0; + p.free[p.nfree] = @intCast(id); + p.nfree += 1; + } + } + /// The stream ended (finished, cancelled or failed): its promise and every page it holds go back. pub fn finish(p: *Pool, l: *Lease) void { p.unpromise(l); diff --git a/zig/src/core/kv_pool_test.zig b/zig/src/core/kv_pool_test.zig index 47196e1856..f6ef4e4fe5 100644 --- a/zig/src/core/kv_pool_test.zig +++ b/zig/src/core/kv_pool_test.zig @@ -216,3 +216,50 @@ test "a promise that does not fit asks the store to reclaim first, and fails onl p.finish(&b); try expect(p.idle()); } + +test "a run: consecutive pages from the top, from the stream's promise, back whole; KV pages fill from the bottom" { + const a = std.testing.allocator; + var p = try Pool.init(a, 10, 1); + defer p.deinit(); + const l: kv.Layout = .{ .fixed_bytes = 0, .token_bytes = 1, .page_tokens = 4, .snap_bytes = 0, .run_bytes = 9 }; + try expectEqual(@as(usize, 3), l.runPages()); + try expectEqual(@as(usize, 1 + 3), l.need(1, 0, 0)); + var x: Lease = .{}; + try p.promise(&x, l.need(5, 0, 0)); // two KV pages and a run of three + try expectEqual(@as(?u32, 7), p.claimRun(&x, 3)); + try expectEqual(@as(usize, 2), x.promised); + try p.cover(&x, 2, null); + try expectEqual(@as(u32, 0), x.ids.items[0]); // KV from the bottom + try expectEqual(@as(u32, 1), x.ids.items[1]); + try expectEqual(@as(usize, 5), p.available()); + try expectEqual(@as(usize, 0), p.reclaimable()); // a run is never the store's + // a second stream: the next run below the first, the promise checked first + var y: Lease = .{}; + try expectEqual(@as(?u32, null), p.claimRun(&y, 3)); + try p.promise(&y, 4); + try expectEqual(@as(?u32, 4), p.claimRun(&y, 3)); + // fragmentation: a run that is not free waits, even with pages to spare + var z: Lease = .{}; + try p.promise(&z, 1); + try p.cover(&z, 1, null); // page 2 + try expectEqual(@as(?u32, null), p.findRun(2)); // only page 3 is free + p.releaseRun(7, 3); + try expectEqual(@as(?u32, 7), p.findRun(3)); + p.releaseRun(4, 3); + p.finish(&x); + p.finish(&y); + p.finish(&z); + try expect(p.idle()); + try expectEqual(@as(usize, 10), p.available()); +} + +test "the plan with a run: each stream pays its run's pages beside its first page, the window bound adds them" { + const g: u64 = 1 << 30; + var l = nemotron; + l.run_bytes = 10 * l.pageBytes() - 1; + const p = try kv.plan(l, .{ .room = 40 * g, .engine = g, .lanes = 4, .window = 16384 + 16 }); + try expectEqual(@as(usize, 3 * (33 + 10)), p.pages); + const small = 2 * g + 2 * (64 * mib + 11 * l.pageBytes()) + mib; + try expectEqual(@as(u32, 3), (try kv.plan(l, .{ .room = small, .engine = 2 * g, .lanes = 8 })).lanes); + try expectEqual(nemotron.need(2048, 4096, 16) + 10, l.need(2048, 4096, 16)); +} diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index e93e1a2e30..965e5a7bcf 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -258,7 +258,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, defer ik_b.free(); var dummy = try cuda.DeviceBuffer.alloc(driver, 256); defer dummy.free(); - try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0); + try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0, .{}); var po_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); defer po_b.free(); @@ -268,7 +268,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, defer pl_b.free(); var o_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); defer o_b.free(); - try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1, false); + try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1, false, .{}); try tri.attnMerge(po_b.ptr, pm_b.ptr, pl_b.ptr, pos_b.ptr, o_b.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, false); var gated_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 5e79eee832..ae28ffa519 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -382,9 +382,10 @@ pub fn Pass(comptime Tri: type) type { const vc = self.seq.vc[ai].ptr; const ikc = self.seq.ikc[ai].ptr; const pooled = self.seq.pooled[ai].ptr; - if (self.seq.media) |m| try m.attn(self.stream, pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, eps, n) else try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0); + const pg = self.seq.pg; + if (self.seq.media) |m| try m.attn(self.stream, pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, eps, n, pg) else try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0, pg); // attention.qsa_pool for the blocks these rows complete, then each block of rows' scores and attention - if (self.seq.media) |m| try m.pooled(self.stream, ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n) else try self.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n); + if (self.seq.media) |m| try m.pooled(self.stream, ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n, pg) else try self.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n, pg); const nb = scoreStride(self.seq.cap); var r0: usize = 0; var b: usize = 0; @@ -398,7 +399,7 @@ pub fn Pass(comptime Tri: type) type { const pos = at.pos + b * 16; const iq = at.iq + r0 * index_heads * index_dim * 2; // rows past `top` complete blocks score theirs; _scores leaves the rest alone, so skip it when none is - if (ends / ratio > top) try self.tri.qsaScores(iq, pooled, pos, self.seq.scores, nb, m, @min(nb, cdiv(ends, ratio))); + if (ends / ratio > top) try self.tri.qsaScores(iq, pooled, pos, self.seq.scores, nb, m, @min(nb, cdiv(ends, ratio)), pg); var a: cuda.Args = .{}; a.add(self.seq.scores); a.add(pos); @@ -411,7 +412,7 @@ pub fn Pass(comptime Tri: type) type { a.add(cint(ratio)); try cuda.launch.launch(self.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, self.stream, &a); const chunks = cdiv(@min(ends, wt.dense_keys), 512); - try self.tri.attnChunks(at.q + r0 * q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true); + try self.tri.attnChunks(at.q + r0 * q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true, pg); try self.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * out_k * 2, s.qnk, s.qspr, m, kv_heads, true); } } diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index dc1749c657..d92c8d276c 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -497,6 +497,8 @@ pub const SeqState = struct { cap: usize, /// Block scores for an attention block's rows [min(rows, att_rows), scoreStride(cap)] fp32 (the engine's). scores: u64, + /// The attention caches' pages (pt 0: kc, vc, ikc and pooled are contiguous rows). + pg: @import("cuda_glue.zig").Pg = .{}, }; /// The block scores' row stride for a cache of `cap` tokens: its blocks rounded up to 16 (the captured _scores'). diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index e000c8f5df..cce104ba8d 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -22,6 +22,9 @@ fn cdiv(a: usize, b: usize) usize { return (a + b - 1) / b; } +/// flashnext_attn.cu's Pg: an int table a page, elements and positions a page (pt 0: contiguous rows). +pub const Pg = extern struct { pt: u64 = 0, ps: i64 = 0, ptok: c_int = 0, pad: c_int = 0 }; + /// The native kernels, by their extern "C" names. const Fns = struct { wb0: cuda.Function, @@ -177,17 +180,18 @@ pub const Glue = struct { } /// Text positions, a bf16 cache: the model's heads (cuda_prompt_wt.zig) normed, roped and stored. - pub fn attnPrep(g: *const Glue, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize) !void { + pub fn attnPrep(g: *const Glue, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, pg: Pg) !void { if (wt.head_dim != tile_head_dim) return error.RowShape; var a: cuda.Args = .{}; ptrs(&a, &.{ proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc }); a.add(eps); ints(&a, &.{ wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half }); + a.add(pg); try go(g.f.prep, s, .{ rows, wt.q_heads + wt.kv_heads + wt.index_heads + 1, 1 }, tile_head_dim, 0, &a); } /// Attention over 512-key chunks; `qsa` the sparse form (NCH 5, rows' key lists), else dense (NCH 3). - pub fn attnChunks(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { + pub fn attnChunks(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool, pg: Pg) !void { const group = wt.q_heads / kv; if (group > tile_heads) return error.RowShape; var a: cuda.Args = .{}; @@ -195,6 +199,7 @@ pub const Glue = struct { ints(&a, &.{ wt.q_heads, kv, group, 512, if (qsa) wt.nch else wt.dense_nch }); a.add(@as(f32, 0.0625)); ints(&a, &.{ wt.index_idw, @intFromBool(qsa) }); + a.add(pg); const ahead = rows <= ahead_rows; try go(if (ahead) g.f.chunks_ahead else g.f.chunks, s, .{ rows, kv, chunks }, 128, chunkSmem(if (ahead) 2 else 1), &a); } @@ -206,19 +211,21 @@ pub const Glue = struct { try go(g.f.merge, s, .{ rows, wt.q_heads, 1 }, tile_head_dim, 0, &a); } - pub fn qsaPool(g: *const Glue, s: cuda.Stream, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { + pub fn qsaPool(g: *const Glue, s: cuda.Stream, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { var a: cuda.Args = .{}; ptrs(&a, &.{ ikc, pooled, pos, k_scale, inv }); a.add(eps); ints(&a, &.{ rows, wt.index_dim, wt.half, wt.index_ratio }); + a.add(pg); try go(g.f.pool, s, .{ rows / wt.index_ratio + 2, 1, 1 }, wt.index_dim, 0, &a); } - pub fn qsaScores(g: *const Glue, s: cuda.Stream, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { + pub fn qsaScores(g: *const Glue, s: cuda.Stream, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize, pg: Pg) !void { if (wt.index_dim != 128 or wt.index_heads > 8) return error.RowShape; var a: cuda.Args = .{}; ptrs(&a, &.{ iq, pooled, pos, sc }); ints(&a, &.{ nb, wt.index_heads, wt.index_ratio, wt.index_top }); + a.add(pg); try go(g.f.scores, s, .{ rows, cdiv(blocks, 64), 1 }, 256, 0, &a); } diff --git a/zig/src/families/flashnext/cuda_kv_lanes.zig b/zig/src/families/flashnext/cuda_kv_lanes.zig new file mode 100644 index 0000000000..d4d988a1e0 --- /dev/null +++ b/zig/src/families/flashnext/cuda_kv_lanes.zig @@ -0,0 +1,268 @@ +//! Flash Next's streams beside the first take attention-cache pages from the core pool (TF_FLASHNEXT_KV_PAGES=n). +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const prompt = @import("cuda_prompt.zig"); +const pages = @import("cuda_pages.zig"); +const mtp = @import("cuda_mtp.zig"); +const snap = @import("cuda_snap.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +/// Rows past a request's prompt and reply its pages hold: one drafted window's. +pub const window_rows: u32 = mtp.max_rows; + +/// TF_FLASHNEXT_KV_PAGES=n: streams beside the first take n-position pages from one pool (unset or 0: none). +pub fn paging() !usize { + const v = std.c.getenv("TF_FLASHNEXT_KV_PAGES") orelse return 0; + const n = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadFlashNextKvPages; + if (n != 0 and !pages.valid(n)) return error.BadFlashNextKvPages; + return n; +} + +/// TF_FLASHNEXT_KV_PAGES_ALL=1: every stream on pages, none on its own contiguous caches (checks). +pub fn pagedOnly() bool { + const v = std.c.getenv("TF_FLASHNEXT_KV_PAGES_ALL") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// Pages a stream maps at most: a full window and one more drafted window. +fn mostPages(capacity: usize, tokens: usize) usize { + return (capacity + window_rows + tokens - 1) / tokens; +} + +fn al(n: usize) usize { + return std.mem.alignForward(usize, n, 256); +} + +const conv_bytes = mtp.conv_bytes; +const state_bytes = wt.state_n * 4; +const tail_bytes = wt.ple_tail_n * wt.wide * 2; + +/// A running stream's own states in its run: each DeltaNet layer's convolution tail and state, the n-gram tail. +fn seqRunBytes() usize { + return wt.linear_n * (al(conv_bytes) + al(state_bytes)) + al(tail_bytes); +} + +/// The MTP head a drafting lane runs: its draft vocabulary and the widest prompt chunk whose streams it keeps. +pub const Head = struct { draft_n: usize, rows: usize }; + +/// Fixed: a lane's table and `extra` (image rows); paged: the attention caches (the head's too); a run: the states. +pub fn layout(capacity: usize, tokens: usize, head: ?Head, extra: usize) kv.Layout { + const layers = wt.attn_n + @intFromBool(head != null); + const run = seqRunBytes() + if (head) |h| mtp.Draft.runBytes(h.draft_n, h.rows) else 0; + return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + extra, .token_bytes = pages.pageBytesOf(tokens, layers) / tokens, .page_tokens = @intCast(tokens), .snap_bytes = snap.bytes(0, head != null), .run_bytes = run }; +} + +/// Where a stream's caches came from: no pool, the one stream on its own caches beside a pool, or a pool lane. +pub const From = enum { none, own, pool }; + +/// The backend's paging: the pool `attach` made, its layout from load, and whether the own stream is taken. +pub const State = struct { + pool: ?*Pool = null, + layout: ?kv.Layout = null, + head: ?Head = null, + paged_only: bool = false, + own_busy: bool = false, + + pub fn choose(s: *const State) From { + if (s.pool == null) return .none; + return if (!s.paged_only and !s.own_busy) .own else .pool; + } + + pub fn opened(s: *State, from: From) void { + if (from == .own) s.own_busy = true; + } + + /// A stream's caches given back: a lane's to the pool (true), else the caller frees them (false). + pub fn close(s: *State, from: From, seq: *prompt.Seq) bool { + switch (from) { + .pool => return s.pool.?.close(seq), + .own => s.own_busy = false, + .none => {}, + } + return false; + } + + /// The lane host's admission: the own stream is free, or a lane and the request's pages are. + pub fn fits(s: *const State, prompt_rows: u32, max_tokens: u32) bool { + const p = s.pool orelse return true; + if (!s.paged_only and !s.own_busy) return true; + return p.fits(@as(usize, prompt_rows) + max_tokens + window_rows); + } +}; + +/// One pool lane: its cache's views (its run's states, its stream's pages), that stream's lease and its run. +const Lane = struct { seq: prompt.Seq, lease: kv.Lease = .{}, used: bool = false, run: u32 = 0 }; + +pub const Pool = struct { + gpa: std.mem.Allocator, + e: *prompt.Engine, + layout: kv.Layout, + capacity: usize, + head: ?Head, + pool: kv.Pool, + arena: cuda.DeviceBuffer, + lanes: []Lane, + + /// `streams` lanes and `count` pages of layout `l`, made and touched now; a lane holds only its table idle. + pub fn init(gpa: std.mem.Allocator, e: *prompt.Engine, l: kv.Layout, capacity: usize, head: ?Head, streams: usize, count: usize) !*Pool { + const p = try gpa.create(Pool); + errdefer gpa.destroy(p); + var arena = try cuda.DeviceBuffer.alloc(e.driver, count * l.pageBytes()); + errdefer arena.free(); + try arena.fill8(0, e.stream.handle); // committed now (the driver commits on first touch) + p.* = .{ .gpa = gpa, .e = e, .layout = l, .capacity = capacity, .head = head, .pool = try kv.Pool.init(gpa, count, l.pageBytes()), .arena = arena, .lanes = &.{} }; + errdefer p.pool.deinit(); + p.lanes = try gpa.alloc(Lane, streams); + errdefer gpa.free(p.lanes); + var made: usize = 0; + errdefer for (p.lanes[0..made]) |*x| x.seq.deinit(gpa); + for (p.lanes) |*x| { + x.* = .{ .seq = try p.laneSeq() }; + made += 1; + } + return p; + } + + /// A lane's cache: views only (states mapped from its run at open, attention caches on the arena) and a table. + fn laneSeq(p: *Pool) !prompt.Seq { + const gpa = p.gpa; + const d = p.e.driver; + var table = try cuda.DeviceBuffer.alloc(d, mostPages(p.capacity, p.layout.page_tokens) * 4); + errdefer table.free(); + try table.fill8(0, p.e.stream.handle); + const paged: pages.Paged = .{ .base = p.arena.ptr, .table = table, .ids = &.{}, .tokens = p.layout.page_tokens, .own_ids = false, .layers = wt.attn_n + @intFromBool(p.head != null) }; + const conv = try gpa.alloc(cuda.DeviceBuffer, wt.linear_n); + errdefer gpa.free(conv); + const state = try gpa.alloc(cuda.DeviceBuffer, wt.linear_n); + errdefer gpa.free(state); + const kc = try paged.views(gpa, d, .key); + errdefer gpa.free(kc); + const vc = try paged.views(gpa, d, .value); + errdefer gpa.free(vc); + const ikc = try paged.views(gpa, d, .index); + errdefer gpa.free(ikc); + const pooled = try paged.views(gpa, d, .pooled); + errdefer gpa.free(pooled); + const none: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0, .borrowed = true }; + @memset(conv, none); + @memset(state, none); + return .{ .conv = conv, .state = state, .kc = kc, .vc = vc, .ikc = ikc, .pooled = pooled, .paged = paged, .tail = none, .history = .{ p.e.eos, p.e.eos }, .fed = 0, .logits = try gpa.alloc(u8, prompt.vocabulary * 2), .cap = 0 }; + } + + pub fn deinit(p: *Pool) void { + p.e.stream.synchronize() catch std.log.warn("flashnext: the pool's stream did not drain at close", .{}); + for (p.lanes) |*x| { + if (x.used) { + p.pool.releaseRun(x.run, p.layout.runPages()); + p.pool.finish(&x.lease); + } + x.seq.deinit(p.gpa); + } + p.gpa.free(p.lanes); + p.arena.free(); + p.pool.deinit(); + p.gpa.destroy(p); + } + + /// Whether a lane is free, the pool can promise a cache of `cap` rows and a run now, and a run is free. + pub fn fits(p: *const Pool, cap: usize) bool { + const k = p.layout.runPages(); + for (p.lanes) |x| if (!x.used) return p.pool.fits(p.layout.pages(cap) + k) and p.pool.findRun(k) != null; + return false; + } + + /// A free lane's cleared cache (a copy) with `cap` rows' pages and its run; LaneMemory: not now. + pub fn open(p: *Pool, cap: usize) !prompt.Seq { + const x = for (p.lanes) |*x| { + if (!x.used) break x; + } else return error.LaneMemory; + const k = p.layout.runPages(); + const n = p.layout.pages(cap); + if (n > mostPages(p.capacity, p.layout.page_tokens)) return error.PromptTooLong; + try p.pool.promise(&x.lease, n + k); + errdefer p.pool.finish(&x.lease); + // the run first: pages fill from the bottom, runs from the top; a fragmented pool waits (LaneMemory) + const run = p.pool.claimRun(&x.lease, k) orelse return error.LaneMemory; + errdefer p.pool.releaseRun(run, k); + var up: Upload = .{ .p = p, .table = x.seq.paged.?.table }; + try p.pool.cover(&x.lease, n, .{ .ctx = &up, .write = Upload.write }); + x.seq.paged.?.ids = x.lease.ids.items; + x.run = run; + p.mapStates(&x.seq); + try p.e.clearSeq(&x.seq, cap); + x.used = true; + return x.seq; + } + + fn runBase(p: *const Pool, x: *const Lane) u64 { + return p.arena.ptr + @as(u64, x.run) * p.layout.pageBytes(); + } + + /// The lane's states as views of its run, laid out afresh (a kept window swaps state buffers with the head's). + fn mapStates(p: *Pool, seq: *prompt.Seq) void { + const x = p.laneOf(seq).?; + var at = p.runBase(x); + for (seq.conv) |*b| { + b.* = .{ .d = p.e.driver, .ptr = at, .len = conv_bytes, .borrowed = true }; + at += al(conv_bytes); + } + for (seq.state) |*b| { + b.* = .{ .d = p.e.driver, .ptr = at, .len = state_bytes, .borrowed = true }; + at += al(state_bytes); + } + seq.tail = .{ .d = p.e.driver, .ptr = at, .len = tail_bytes, .borrowed = true }; + } + + /// The MTP head of the stream on lane cache `seq`: its cache on the stream's pages, its buffers in the run. + pub fn draft(p: *Pool, seq: *const prompt.Seq, cap: usize) !mtp.Draft { + const h = p.head orelse return error.FlashNextRunTooShort; + const x = p.laneOf(seq) orelse return error.UnknownStream; + const run_len = p.layout.runPages() * p.layout.pageBytes(); + return mtp.Draft.initIn(p.gpa, p.e.driver, cap, h.draft_n, h.rows, &x.seq.paged.?, p.runBase(x) + seqRunBytes(), run_len - seqRunBytes()); + } + + fn laneOf(p: *Pool, seq: *const prompt.Seq) ?*Lane { + const t = (seq.paged orelse return null).table.ptr; + for (p.lanes) |*x| if (x.seq.paged.?.table.ptr == t) return x; + return null; + } + + /// A lane's stream ended: its pages and run go back once the forwards reading them are done. False: not a lane's. + pub fn close(p: *Pool, seq: *prompt.Seq) bool { + const x = p.laneOf(seq) orelse return false; + if (!x.used) return true; + // a page another stream takes next may still be read by queued work on the engine's stream + p.e.stream.synchronize() catch std.log.warn("flashnext: a closed stream's forwards did not drain; its pages go back", .{}); + p.pool.releaseRun(x.run, p.layout.runPages()); + p.pool.finish(&x.lease); + if (seq.media) |*m| m.deinit(); // the stream's image state, held by its copy + seq.media = null; + x.seq.paged.?.ids = &.{}; + x.seq.cap = 0; + x.used = false; + return true; + } + + /// A lane's new table entries uploaded to its device table. + const Upload = struct { + p: *Pool, + table: cuda.DeviceBuffer, + + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const u: *Upload = @ptrCast(@alignCast(ctx)); + try u.table.uploadAsync(from * 4, std.mem.sliceAsBytes(ids[from..]), u.p.e.stream.handle); + try u.p.e.stream.synchronize(); // the ids slice may move as the lease grows + } + }; +}; + +test "a stream's pages: its rows and a drafted window, no more than a full window's" { + const l = layout(40960, 256, null, 0); + try std.testing.expectEqual(pages.pageBytes(256), l.pageBytes()); + const h = layout(40960, 256, .{ .draft_n = 32768, .rows = 1536 }, 0); + try std.testing.expectEqual(pages.pageBytesOf(256, wt.attn_n + 1), h.pageBytes()); + try std.testing.expect(h.run_bytes >= seqRunBytes() + mtp.Draft.runBytes(32768, 1536)); + try std.testing.expectEqual(@as(usize, 161), mostPages(40960, 256)); + try std.testing.expect(l.pages(40960 + window_rows) <= mostPages(40960, 256)); +} diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 226c988b55..f940142611 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -13,6 +13,7 @@ const Glue = @import("cuda_glue.zig").Glue; const living = @import("living"); const slide_mod = @import("cuda_slide.zig"); const train = @import("cuda_train.zig"); +const kvl = @import("cuda_kv_lanes.zig"); const be = lanes.backend; const ring = 1024; @@ -40,6 +41,17 @@ pub fn Namespace(comptime Tri: type) type { pub const snap_save = snapSave; pub const snap_restore = snapRestore; pub const snap_drop = snapDrop; + /// TF_FLASHNEXT_KV_PAGES: the pool's lanes, the pages of a request's rows and a drafted window (kvl). + pub const window_rows = kvl.window_rows; + pub fn attach(p: *anyopaque, streams: u32, count: usize) !void { + const own = of(p); + own.kv.pool = try kvl.Pool.init(own.gpa, own.engine, own.kv.layout.?, own.context, own.kv.head, streams, count); + own.kv.paged_only = kvl.pagedOnly() and streams > 0; + } + pub fn fits(p: *anyopaque, rows: u32, max_tokens: u32) bool { + return of(p).kv.fits(rows, max_tokens); + } + /// The pass ignores its cuts, so states keep at any token; the MTP head reads one token on (lookahead 1). pub const cache_rules = struct { pub const lookahead: u32 = 1; @@ -85,6 +97,8 @@ pub fn Namespace(comptime Tri: type) type { /// Whether our own glue runs (no captured set given, or one without this family's entry points). own_glue: bool = true, follow: ?*const fn (*anyopaque) anyerror!void = null, + /// TF_FLASHNEXT_KV_PAGES: the streams beside the first take pages of one pool (null: their own caches). + layout: ?@import("core").kv_pool.Layout = null, }; /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. @@ -128,6 +142,8 @@ pub fn Namespace(comptime Tri: type) type { return switch (err) { error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + error.BadFlashNextKvPages, error.BadFlashNextPageTest => "TF_FLASHNEXT_KV_PAGES and _PAGE_TEST take positions a page, a multiple of 4", + error.FlashNextPagesUnsupported => "TF_FLASHNEXT_KV_PAGES needs one rank, our own glue and no --slide: unset it", error.UnsupportedCheckpoint => "the checkpoint differs from the supported Flash Next CUDA layout; the log names its config field", error.TensorRankDraftSnapshotUnavailable => "tensor-rank MTP snapshots are unavailable: disable the prompt cache for this configuration", error.RankDisconnected, error.RankTimeout, error.RemoteRank, error.BadRankFrame => "a tensor-parallel rank failed; restart both ranks from the same checkpoint before retrying", @@ -145,6 +161,8 @@ pub fn Namespace(comptime Tri: type) type { chain: mtp.Chain = .{}, /// The prompt's next row while its pass runs a chunk a call; null once the prompt is in. filled: ?usize = 0, + /// Where its caches came from: the pool's lanes, or its own (the one a pool leaves beside it). + from: kvl.From = .none, }; const Ops = mtp.Ops(Tri); @@ -169,6 +187,7 @@ pub fn Namespace(comptime Tri: type) type { next: u64 = 0, /// --slide: the pocket's layout, from config.json at load. lay: living.Layout = undefined, + kv: kvl.State = .{}, fn tri(self: *Cuda) Tri { return .{ .set = if (self.set) |*x| x else null, .own = &self.glue, .s = self.stream }; @@ -273,17 +292,19 @@ pub fn Namespace(comptime Tri: type) type { } // a window's rows past the reply's last token (its drafts) still need cache slots const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); - gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else self.engine.newSeq(cap)) catch |err| { + const from = self.kv.choose(); + gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else if (from == .pool) self.kv.pool.?.open(cap) else self.engine.newSeq(cap)) catch |err| { self.lanes.removeByPtr(gop.key_ptr); return err; - } }; + }, .from = from }; + self.kv.opened(from); const l = gop.value_ptr; l.seq.untaught = s.untaught; errdefer { dropLane(self, l); self.lanes.removeByPtr(gop.key_ptr); } - if (drafting) l.draft = try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); + if (drafting) l.draft = if (from == .pool) try self.kv.pool.?.draft(&l.seq, cap) else try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); if (s.images.len != 0) { const v = self.vision orelse return error.VisionNotEnabled; l.seq.media = try v.state(ids, s.images); @@ -370,6 +391,7 @@ pub fn Namespace(comptime Tri: type) type { if (own.head) |h| h.deinit(); if (own.vision) |v| v.deinit(); if (own.ranks) |r| r.deinit(); + if (own.kv.pool) |kp| kp.deinit(); own.engine.deinit(); if (own.set) |*x| x.deinit(); own.glue.deinit(); @@ -379,7 +401,7 @@ pub fn Namespace(comptime Tri: type) type { fn dropLane(self: *Cuda, l: *Lane) void { if (l.draft) |*d| d.deinit(); - if (self.ranks) |r| r.drop(&l.seq) else l.seq.deinit(self.gpa); + if (self.ranks) |r| r.drop(&l.seq) else if (!self.kv.close(l.from, &l.seq)) l.seq.deinit(self.gpa); } fn of(ptr: *anyopaque) *Cuda { diff --git a/zig/src/families/flashnext/cuda_lanes_admission_test.zig b/zig/src/families/flashnext/cuda_lanes_admission_test.zig index 5589587c80..898c0bd9c6 100644 --- a/zig/src/families/flashnext/cuda_lanes_admission_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_admission_test.zig @@ -31,6 +31,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, coun self.ranks = null; self.head = &head; self.lanes = .empty; + self.kv = .{}; defer { var it = self.lanes.valueIterator(); while (it.next()) |lane| drop(&self, lane); diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 69c300ef69..5de4cba5da 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -19,6 +19,8 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; if (o.slide and (mode != .affine or o.tp != 1)) return error.UnsupportedLearningFormat; if (o.slide and captured != null) return error.TrainingOnCapturedGlue; + const page_tokens = try @import("cuda_kv_lanes.zig").paging(); + if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null)) return error.FlashNextPagesUnsupported; const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); @@ -79,7 +81,14 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); if (o.vision) |v| stream_bytes += @as(usize, v.image_tokens) * @import("cuda_prompt_wt.zig").dims * 2 + 16 * o.context; + if (page_tokens > 0) { + // the lane's head runs at the engine's chunk: its cache on the stream's pages, its buffers in the run + own.kv.head = if (own.head) |h| .{ .draft_n = h.ids.len, .rows = own.engine.chunk } else null; + const head_bytes = if (own.head) |h| mtp.Draft.bytes(window, h.ids.len, own.engine.chunk) else 0; + own.kv.layout = @import("cuda_kv_lanes.zig").layout(o.context, page_tokens, own.kv.head, stream_bytes - prompt.seqBytes(window) - head_bytes); + } return .{ + .layout = own.kv.layout, .vision = own.vision, .backend = own.backend(), .facts = own.facts(), diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 99e25c0c12..5f69440023 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -33,6 +33,7 @@ pub const native = @import("cuda_lanes.zig").Namespace(Tri); /// Our own glue kernels, which every launch takes when no captured set is loaded pub const Glue = glue_.Glue; const glue_ = @import("cuda_glue.zig"); +pub const Pg = glue_.Pg; const p = aot.ptr; @@ -131,8 +132,9 @@ pub const Tri = struct { } /// glue.attn_prep for text positions, bf16 KV (BITS 0, the scale pointers unread); `length` unused at MODE 0. - pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { - if (t.glue("_attn_prep")) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows); + pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize, pg: Pg) !void { + if (t.glue("_attn_prep")) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows, pg); + if (pg.pt != 0) return error.PagedKvNeedsGlue; try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ p("P", "*bf16", proj), p("POS0", "*i32", pos), @@ -155,8 +157,9 @@ pub const Tri = struct { } /// attention._chunks over 512-key chunks, bf16 KV; `qsa` the sparse form (NCH 5, QSA 1), else dense (NCH 3, QSA 0). - pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { - if (t.glue("_chunks")) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa); + pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool, pg: Pg) !void { + if (t.glue("_chunks")) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa, pg); + if (pg.pt != 0) return error.PagedKvNeedsGlue; try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ p("Q", "*bf16", q), p("KC", "*bf16", kc), @@ -188,8 +191,9 @@ pub const Tri = struct { } /// attention.qsa_pool for text rows (MODE 0): every complete block of 4 indexer keys pooled, normed and roped. - pub fn qsaPool(t: Tri, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { - if (t.glue("_pool")) |g| return g.qsaPool(t.s, ikc, pooled, pos, k_scale, inv, eps, rows); + pub fn qsaPool(t: Tri, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { + if (t.glue("_pool")) |g| return g.qsaPool(t.s, ikc, pooled, pos, k_scale, inv, eps, rows, pg); + if (pg.pt != 0) return error.PagedKvNeedsGlue; try t.run("_pool", .{ u(rows / 4 + 2), 1, 1 }, &.{ p("IKC", "*bf16", ikc), p("POOLED", "*bf16", pooled), @@ -203,8 +207,9 @@ pub const Tri = struct { } /// attention._scores: rows' complete blocks scored; `nb` is the row stride (a multiple of 16), `blocks` the grid. - pub fn qsaScores(t: Tri, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { - if (t.glue("_scores")) |g| return g.qsaScores(t.s, iq, pooled, pos, sc, nb, rows, blocks); + pub fn qsaScores(t: Tri, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize, pg: Pg) !void { + if (t.glue("_scores")) |g| return g.qsaScores(t.s, iq, pooled, pos, sc, nb, rows, blocks, pg); + if (pg.pt != 0) return error.PagedKvNeedsGlue; try t.run("_scores", .{ u(rows), cdiv(blocks, 64), 1 }, &.{ p("IQ", "*bf16", iq), p("POOLED", "*bf16", pooled), @@ -384,6 +389,8 @@ test { _ = mtp; _ = native; _ = @import("cuda_snap.zig"); + _ = @import("cuda_pages.zig"); + _ = @import("cuda_kv_lanes.zig"); _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); diff --git a/zig/src/families/flashnext/cuda_media.zig b/zig/src/families/flashnext/cuda_media.zig index 6e66da20a0..014df50784 100644 --- a/zig/src/families/flashnext/cuda_media.zig +++ b/zig/src/families/flashnext/cuda_media.zig @@ -2,6 +2,7 @@ const cuda = @import("cuda"); const shared = @import("qwen_vision"); const wt = @import("cuda_prompt_wt.zig"); +const Pg = @import("cuda_glue.zig").Pg; pub const State = struct { image: shared.media.Media, prep: cuda.Function, @@ -17,7 +18,7 @@ pub const View = struct { image: shared.media.View, prep: cuda.Function, pool: cuda.Function, - pub fn attn(v: View, stream: cuda.Stream, proj: u64, pos: u64, qw: u64, kw: u64, iw: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize) !void { + pub fn attn(v: View, stream: cuda.Stream, proj: u64, pos: u64, qw: u64, kw: u64, iw: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, pg: Pg) !void { var a: cuda.Args = .{}; for ([_]u64{ proj, pos, qw, kw, iw, inv, q, kc, vc, iq, ikc }) |x| a.add(x); a.add(eps); @@ -25,9 +26,10 @@ pub const View = struct { a.add(v.image.rope); a.add(@as(i32, @intCast(v.image.n))); a.add(v.image.delta); + a.add(pg); try cuda.launch.launch(v.prep, .{ .grid = .{ .x = @intCast(rows), .y = wt.q_heads + wt.kv_heads + wt.index_heads + 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, stream, &a); } - pub fn pooled(v: View, stream: cuda.Stream, ikc: u64, out: u64, pos: u64, w: u64, inv: u64, eps: f32, rows: usize) !void { + pub fn pooled(v: View, stream: cuda.Stream, ikc: u64, out: u64, pos: u64, w: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { var a: cuda.Args = .{}; for ([_]u64{ ikc, out, pos, w, inv }) |x| a.add(x); a.add(eps); @@ -35,6 +37,7 @@ pub const View = struct { a.add(v.image.rope); a.add(@as(i32, @intCast(v.image.n))); a.add(v.image.delta); + a.add(pg); try cuda.launch.launch(v.pool, .{ .grid = .{ .x = @intCast(rows / wt.index_ratio + 2), .y = 1, .z = 1 }, .block = .{ .x = wt.index_dim, .y = 1, .z = 1 } }, stream, &a); } }; diff --git a/zig/src/families/flashnext/cuda_media_test.zig b/zig/src/families/flashnext/cuda_media_test.zig index 0f58af79ba..302a670334 100644 --- a/zig/src/families/flashnext/cuda_media_test.zig +++ b/zig/src/families/flashnext/cuda_media_test.zig @@ -3,10 +3,12 @@ const std = @import("std"); const cuda = @import("cuda"); const media = @import("cuda_media.zig"); const wt = @import("cuda_prompt_wt.zig"); +const Pg = @import("cuda_glue.zig").Pg; const Capture = struct { var calls: usize = 0; var values: [8][21]u64 = undefined; var configs: [8][7]u32 = undefined; + var pages: [8]Pg = undefined; fn launch(f: cuda.abi.Function, gx: c_uint, gy: c_uint, gz: c_uint, bx: c_uint, by: c_uint, bz: c_uint, sm: c_uint, s: cuda.abi.Stream, args: ?[*]?*anyopaque, extra: ?[*]?*anyopaque) callconv(.c) c_int { std.debug.assert(extra == null and s == @as(cuda.abi.Stream, @ptrFromInt(99))); const id = @intFromPtr(f); @@ -15,6 +17,7 @@ const Capture = struct { const wide = if (id == 1) i < 3 else if (id == 2) i < 11 or i == 18 else i < 5 or i == 10; values[calls][i] = if (wide) @as(*const u64, @ptrCast(@alignCast(args.?[i].?))).* else @as(*const u32, @ptrCast(@alignCast(args.?[i].?))).*; } + if (id != 1) pages[calls] = @as(*const Pg, @ptrCast(@alignCast(args.?[count].?))).*; configs[calls] = .{ gx, gy, gz, bx, by, bz, sm }; calls += 1; return 0; @@ -32,14 +35,16 @@ test "the splice, attention and pooled indexer launches get every pointer, scala try m.image.inject(s, 4000, 2, 3, 4); try std.testing.expectEqualSlices(u64, &.{ 1000, 2000, 4000, 2, 20, 2560, 4 }, Capture.values[0][0..7]); try std.testing.expectEqual([7]u32{ 3, 40, 1, 256, 1, 1, 0 }, Capture.configs[0]); - try m.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 3); + try m.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 3, .{ .pt = 12, .ps = 13, .ptok = 16 }); const eps: u32 = @bitCast(@as(f32, 1e-6)); const delta: u32 = @bitCast(@as(i32, -2)); try std.testing.expectEqualSlices(u64, &.{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, eps, wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half, 3000, 20, delta }, &Capture.values[1]); try std.testing.expectEqual([7]u32{ 3, wt.q_heads + wt.kv_heads + wt.index_heads + 1, 1, 256, 1, 1, 0 }, Capture.configs[1]); - try m.pooled(s, 1, 2, 3, 4, 5, 1e-6, 8); + try std.testing.expectEqual(Pg{ .pt = 12, .ps = 13, .ptok = 16 }, Capture.pages[1]); + try m.pooled(s, 1, 2, 3, 4, 5, 1e-6, 8, .{}); try std.testing.expectEqualSlices(u64, &.{ 1, 2, 3, 4, 5, eps, 8, wt.index_dim, wt.half, wt.index_ratio, 3000, 20, delta }, Capture.values[2][0..13]); try std.testing.expectEqual([7]u32{ 8 / wt.index_ratio + 2, 1, 1, 128, 1, 1, 0 }, Capture.configs[2]); + try std.testing.expectEqual(Pg{}, Capture.pages[2]); const Row = struct { seq: struct { media: ?media.View }, pos: usize }; const rows = [_]Row{ .{ .seq = .{ .media = m }, .pos = 2 }, .{ .seq = .{ .media = null }, .pos = 2 }, .{ .seq = .{ .media = view(&d, 1100) }, .pos = 5 } }; try media.inject(null, &rows, s, 9000, 0, 3); diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index 1ae18ae2ab..ddabdbd187 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -7,6 +7,7 @@ const batch = @import("cuda_batch.zig"); const prompt = @import("cuda_prompt.zig"); const gdn = @import("cuda_gdn.zig"); const mtp = @import("cuda_mtp.zig"); +const pages = @import("cuda_pages.zig"); const dims = wt.dims; const streams = wt.streams; @@ -51,6 +52,8 @@ pub const Draft = struct { ikc: cuda.DeviceBuffer, /// Its pooled indexer keys [cap / 4, 128] bf16, for sparse attention past 2,048 keys (Python State.mtp_pooled). pooled: cuda.DeviceBuffer, + /// The cache's pages (a pool lane's: its stream's, the layer after the model's), or contiguous. + pg: @import("cuda_glue.zig").Pg = .{}, /// Entries the head's cache holds (Python mtp_len), the last `drafted` of them a chain's. len: usize = 0, drafted: usize = 0, @@ -124,6 +127,12 @@ pub const Draft = struct { errdefer gpa.free(d.host); // zeroed on the compute stream, ahead of the passes that read them for ([_]cuda.DeviceBuffer{ d.kc, d.vc, d.ikc, d.pooled }) |b| try b.fill8(0, stream.handle); + d.point(); + return d; + } + + /// The keep's history and projection rows in `saved`, and the end states into the spares. + fn point(d: *Draft) void { var li: usize = 0; for (&d.keep.layers, 0..) |*lk, layer| { if (!wt.linearLayer(layer)) continue; @@ -132,6 +141,39 @@ pub const Draft = struct { li += 1; } d.aim(); + } + + /// Device bytes `initIn` takes from a run: every buffer but the head's cache, each on a 256-byte boundary. + pub fn runBytes(draft_n: usize, rows: usize) usize { + return bytes(0, draft_n, rows) + (6 + linear_n) * 256; + } + + /// A pool lane's head: its cache on `paged`'s pages after the model's layers, its buffers from `base`'s run. + pub fn initIn(gpa: std.mem.Allocator, driver: *cuda.Driver, cap: usize, draft_n: usize, rows: usize, paged: *const pages.Paged, base: u64, len: usize) !Draft { + if (len < runBytes(draft_n, rows)) return error.FlashNextRunTooShort; + var at = base; + const Carve = struct { + fn take(d: *const cuda.Driver, p: *u64, n: usize) cuda.DeviceBuffer { + const b: cuda.DeviceBuffer = .{ .d = d, .ptr = p.*, .len = n, .borrowed = true }; + p.* += std.mem.alignForward(usize, n, 256); + return b; + } + fn page(d: *const cuda.Driver, ptr: u64) cuda.DeviceBuffer { + return .{ .d = d, .ptr = ptr, .len = 0, .borrowed = true }; + } + }; + const head = wt.attn_n; + var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = Carve.page(driver, paged.at(head, .key)), .vc = Carve.page(driver, paged.at(head, .value)), .ikc = Carve.page(driver, paged.at(head, .index)), .pooled = Carve.page(driver, paged.at(head, .pooled)), .pg = paged.pg(), .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; + d.streams = Carve.take(driver, &at, d.rows_cap * wide * 2); + d.out = Carve.take(driver, &at, wide * 2); + d.logits = Carve.take(driver, &at, max_rows * vocab * 2); + d.dlogits = Carve.take(driver, &at, draft_n * 2); + d.best = Carve.take(driver, &at, max_rows * 16); + d.saved = Carve.take(driver, &at, saved_bytes); + for (&d.spare) |*b| b.* = Carve.take(driver, &at, state_n * 4); + std.debug.assert(at - base <= len); + d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); + d.point(); return d; } diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index cfcd9fdd05..eab7e37b97 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -124,7 +124,7 @@ pub fn Ops(comptime Tri: type) type { .k = &e.batch_k, .s = s, .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram }, - .seq = .{ .media = d.media, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, + .seq = .{ .media = d.media, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .pg = d.pg, .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, .hooks = .{ .logits = .none }, }; var at: usize = 0; diff --git a/zig/src/families/flashnext/cuda_pages.zig b/zig/src/families/flashnext/cuda_pages.zig new file mode 100644 index 0000000000..a59c02555c --- /dev/null +++ b/zig/src/families/flashnext/cuda_pages.zig @@ -0,0 +1,139 @@ +//! Flash Next's attention caches as pages: per attention layer, keys, values, indexer keys, then pooled rows (bf16). + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const Pg = @import("cuda_glue.zig").Pg; + +const kv_row = wt.kv_heads * wt.head_dim * 2; +const index_row = wt.index_dim * 2; +const ratio = wt.index_ratio; + +pub const Part = enum { key, value, index, pooled }; + +/// Bytes one row of `part` takes: a position's (a block's for pooled). +pub fn rowBytes(part: Part) usize { + return switch (part) { + .key, .value => kv_row, + .index, .pooled => index_row, + }; +} + +/// Rows of `part` a page of `tokens` positions holds. +pub fn rows(part: Part, tokens: usize) usize { + return if (part == .pooled) tokens / ratio else tokens; +} + +/// Bytes one attention layer's share of a page takes. +pub fn layerBytes(tokens: usize) usize { + return tokens * (2 * kv_row + index_row) + tokens / ratio * index_row; +} + +/// Bytes a page takes: every attention layer's share. +pub fn pageBytes(tokens: usize) usize { + return pageBytesOf(tokens, wt.attn_n); +} + +/// Bytes a page of `layers` attention layers takes (the model's, and the MTP head's cache when it drafts). +pub fn pageBytesOf(tokens: usize, layers: usize) usize { + return layers * layerBytes(tokens); +} + +/// Where `part` of attention layer `layer` starts inside a page. +pub fn offset(tokens: usize, layer: usize, part: Part) usize { + const at = layer * layerBytes(tokens); + return at + switch (part) { + .key => 0, + .value => tokens * kv_row, + .index => 2 * tokens * kv_row, + .pooled => 2 * tokens * kv_row + tokens * index_row, + }; +} + +/// A page size the kernels take: whole blocks of the indexer's ratio. +pub fn valid(tokens: usize) bool { + return tokens > 0 and tokens % ratio == 0; +} + +/// TF_FLASHNEXT_KV_PAGE_TEST=n: each stream's caches read through reversed n-position pages (checks); 0 when unset. +pub fn pageTest() !usize { + const v = std.c.getenv("TF_FLASHNEXT_KV_PAGE_TEST") orelse return 0; + const n = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadFlashNextPageTest; + if (!valid(n)) return error.BadFlashNextPageTest; + return n; +} + +/// One stream's pages: page i of its positions at `base` + ids[i] pages. +pub const Paged = struct { + /// Where page 0 of the arena starts. + base: u64, + /// The stream's table on the device, an int a page (`ids` uploaded). + table: cuda.DeviceBuffer, + ids: []u32, + /// Positions a page holds. + tokens: usize, + /// The paged state owns `ids` (the page test's); a pool's lane maps its lease's. + own_ids: bool = true, + /// The page test's own arena. + arena: ?cuda.DeviceBuffer = null, + /// Attention layers a page holds: the model's, and the MTP head's after them in a pool's pages. + layers: usize = wt.attn_n, + + pub fn pg(p: *const Paged) Pg { + return .{ .pt = p.table.ptr, .ps = @intCast(pageBytesOf(p.tokens, p.layers) / 2), .ptok = @intCast(p.tokens) }; + } + + /// `part` of attention layer `layer` in page 0: the cache pointer a kernel takes beside pg(). + pub fn at(p: *const Paged, layer: usize, part: Part) u64 { + return p.base + offset(p.tokens, layer, part); + } + + /// Byte address of row `row` of `part` of `layer` (a position, or a block for pooled). + pub fn rowAt(p: *const Paged, layer: usize, part: Part, row: usize) u64 { + const per = rows(part, p.tokens); + return p.base + @as(u64, p.ids[row / per]) * pageBytesOf(p.tokens, p.layers) + offset(p.tokens, layer, part) + (row % per) * rowBytes(part); + } + + /// Borrowed views of each attention layer's `part` in page 0 (what kc, vc, ikc and pooled hold when paged). + pub fn views(p: *const Paged, gpa: std.mem.Allocator, d: *const cuda.Driver, part: Part) ![]cuda.DeviceBuffer { + const out = try gpa.alloc(cuda.DeviceBuffer, wt.attn_n); + for (out, 0..) |*b, layer| b.* = .{ .d = d, .ptr = p.at(layer, part), .len = 0, .borrowed = true }; + return out; + } + + pub fn deinit(p: *Paged, gpa: std.mem.Allocator) void { + p.table.free(); + if (p.own_ids) gpa.free(p.ids); + if (p.arena) |*a| a.free(); + } +}; + +/// The page test's zeroed pages for `cap` positions in one buffer, in reverse order (no contiguous layout matches). +pub fn testPages(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize, tokens: usize) !Paged { + const n = (cap + tokens - 1) / tokens; + var arena = try cuda.DeviceBuffer.alloc(d, n * pageBytes(tokens)); + errdefer arena.free(); + try arena.fill8(0, stream.handle); + const ids = try gpa.alloc(u32, n); + errdefer gpa.free(ids); + for (ids, 0..) |*id, i| id.* = @intCast(n - 1 - i); + var table = try cuda.DeviceBuffer.alloc(d, n * 4); + errdefer table.free(); + try table.upload(0, std.mem.sliceAsBytes(ids)); + return .{ .base = arena.ptr, .table = table, .ids = ids, .tokens = tokens, .arena = arena }; +} + +test "a page's parts tile it in order and pooled rows follow their positions' page" { + const t = 16; + try std.testing.expectEqual(@as(usize, 0), offset(t, 0, .key)); + try std.testing.expectEqual(t * kv_row, offset(t, 0, .value)); + try std.testing.expectEqual(offset(t, 0, .pooled) + t / ratio * index_row, offset(t, 1, .key)); + try std.testing.expectEqual(pageBytes(t), offset(t, wt.attn_n - 1, .pooled) + t / ratio * index_row); + try std.testing.expect(!valid(6) and valid(8) and !valid(0)); + var ids = [_]u32{ 2, 0, 1 }; + const p: Paged = .{ .base = 1 << 20, .table = undefined, .ids = &ids, .tokens = t }; + // position 17 is page 1 (id 0) row 1; block 5 (positions 20..23) is page 1 row 1 of the pooled rows + try std.testing.expectEqual(p.base + offset(t, 3, .value) + kv_row, p.rowAt(3, .value, 17)); + try std.testing.expectEqual(p.base + offset(t, 3, .pooled) + index_row, p.rowAt(3, .pooled, 5)); + try std.testing.expectEqual(p.base + 2 * pageBytes(t) + offset(t, 0, .index), p.rowAt(0, .index, 0)); +} diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 0ff498c22d..883acaf326 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -113,6 +113,7 @@ pub fn pooledBytes(cap: usize) usize { } const round_graph = @import("cuda_round_graph.zig"); +const pages = @import("cuda_pages.zig"); /// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. pub const Engine = struct { @@ -287,6 +288,14 @@ pub const Engine = struct { /// A fresh cache. `cap` is how many tokens it can hold, prompt plus generated. pub fn newSeq(e: *Engine, cap: usize) !Seq { + const test_pages = try pages.pageTest(); + return e.newSeqOn(cap, if (test_pages > 0) try pages.testPages(e.gpa, e.driver, e.stream.*, cap, test_pages) else null); + } + + /// A fresh cache whose attention caches are `given` pages (owned by the cache from here), else contiguous. + pub fn newSeqOn(e: *Engine, cap: usize, given: ?pages.Paged) !Seq { + var paged = given; + errdefer if (paged) |*p| p.deinit(e.gpa); e.round_graph.clear(); if (cap == 0) return error.UnexpectedTensor; const logits = try e.gpa.alloc(u8, vocab * 2); @@ -295,13 +304,13 @@ pub const Engine = struct { errdefer freeOwned(e.gpa, conv); const state = try zeros(e.gpa, e.driver, e.stream.*, linear_n, state_n * 4); errdefer freeOwned(e.gpa, state); - const kc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + const kc = if (paged) |*p| try p.views(e.gpa, e.driver, .key) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, kc); - const vc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + const vc = if (paged) |*p| try p.views(e.gpa, e.driver, .value) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, vc); - const ikc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); + const ikc = if (paged) |*p| try p.views(e.gpa, e.driver, .index) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); errdefer freeOwned(e.gpa, ikc); - const pooled = try zeros(e.gpa, e.driver, e.stream.*, attn_n, pooledBytes(cap)); + const pooled = if (paged) |*p| try p.views(e.gpa, e.driver, .pooled) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, pooledBytes(cap)); errdefer freeOwned(e.gpa, pooled); var tail = try cuda.DeviceBuffer.alloc(e.driver, ple_tail_n * wide * 2); errdefer tail.free(); @@ -313,6 +322,7 @@ pub const Engine = struct { .vc = vc, .ikc = ikc, .pooled = pooled, + .paged = paged, .tail = tail, .history = .{ e.eos, e.eos }, .fed = 0, @@ -321,6 +331,18 @@ pub const Engine = struct { }; } + /// A pool lane's cache made fresh for its next stream: zero states and tail, nothing fed, `cap` rows. + pub fn clearSeq(e: *Engine, s: *Seq, cap: usize) !void { + e.round_graph.clear(); + for (s.conv) |b| try b.fill8(0, e.stream.handle); + for (s.state) |b| try b.fill8(0, e.stream.handle); + try s.tail.fill8(0, e.stream.handle); + s.history = .{ e.eos, e.eos }; + s.fed = 0; + s.cap = cap; + s.untaught = false; + } + /// Run `ids` on from `seq.fed` and leave the next token's scores in `seq.logits`. pub fn predict(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32) !void { return e.forward(Tri, tri, seq, ids, .{}); @@ -449,6 +471,8 @@ pub const Seq = struct { vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, pooled: []cuda.DeviceBuffer, + /// kc, vc, ikc and pooled as borrowed views of page 0 of these pages (null: each its own contiguous buffer). + paged: ?pages.Paged = null, tail: cuda.DeviceBuffer, history: [2]i64, fed: usize, @@ -459,7 +483,12 @@ pub const Seq = struct { /// The caches a batched pass reads and leaves; `scores` the engine's block scores (Engine.reserveScores). pub fn view(s: *Seq, scores: u64) batch.SeqState { - return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores }; + return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores, .pg = s.pg() }; + } + + /// The attention kernels' page argument (zero: contiguous caches). + pub fn pg(s: *const Seq) @import("cuda_glue.zig").Pg { + return if (s.paged) |*p| p.pg() else .{}; } pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { @@ -476,6 +505,7 @@ pub const Seq = struct { gpa.free(s.vc); gpa.free(s.ikc); gpa.free(s.pooled); + if (s.paged) |*p| p.deinit(gpa); if (s.tail.ptr != 0) s.tail.free(); gpa.free(s.logits); } diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig index 36d364152f..58f76a512f 100644 --- a/zig/src/families/flashnext/cuda_prompt_step.zig +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -235,16 +235,16 @@ pub fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Strea try qmm.partials(driver, stream.*, s.gout.ptr, s.gxs.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, gdn.value_dim, out_sk); } -pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { +pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32, pg: @import("cuda_glue.zig").Pg) !void { try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); // stream-ordered (queued kernels still read s.pos), then waited for: `pos` lives on this frame try s.pos.uploadAsync(0, std.mem.asBytes(&pos), stream.handle); try stream.synchronize(); - try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); + try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0, pg); // attention.attention, dense: chunks of 512 keys up to the row's position (the row-by-row path stays dense) const keys: usize = @as(usize, @intCast(pos)) + 1; if (keys > wt.dense_nch * 512) return error.UnexpectedTensor; - try tri.attnChunks(s.q.ptr, kc, vc, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512, false); + try tri.attnChunks(s.q.ptr, kc, vc, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512, false, pg); try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, false); try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); diff --git a/zig/src/families/flashnext/cuda_rows.zig b/zig/src/families/flashnext/cuda_rows.zig index 89095b53ee..49df3faa45 100644 --- a/zig/src/families/flashnext/cuda_rows.zig +++ b/zig/src/families/flashnext/cuda_rows.zig @@ -75,7 +75,7 @@ pub const Rows = struct { // dense attention only: the reference stops where the indexer would start choosing if (seq.fed + ids.len > weights.dense_nch * 512) return error.UnexpectedTensor; try r.grow(e.stream.*, ids.len); - var pass = Pass{ .gpa = e.gpa, .io = e.io, .driver = e.driver, .stream = e.stream, .model_dir = e.model_dir, .s = &r.scratch, .table = e.table, .mode = r.mode, .history = &seq.history, .conv = seq.conv, .state = seq.state, .kc = seq.kc, .vc = seq.vc, .ikc = seq.ikc, .tail = seq.tail, .mixer = e.mixer, .head_q = e.head_q, .logits = seq.logits, .plan_fn = r.plan_fn, .up_fn = r.up_fn, .down_fn = r.down_fn }; + var pass = Pass{ .gpa = e.gpa, .io = e.io, .driver = e.driver, .stream = e.stream, .model_dir = e.model_dir, .s = &r.scratch, .table = e.table, .mode = r.mode, .history = &seq.history, .conv = seq.conv, .state = seq.state, .kc = seq.kc, .vc = seq.vc, .ikc = seq.ikc, .pg = seq.pg(), .tail = seq.tail, .mixer = e.mixer, .head_q = e.head_q, .logits = seq.logits, .plan_fn = r.plan_fn, .up_fn = r.up_fn, .down_fn = r.down_fn }; @memset(pass.mode, 0); for (ids, 0..) |token, t| try embedRow(Tri, &pass, tri, token, pass.s.h.ptr + @as(u64, t) * wide * 2); for (e.layers[0..e.n_layers], 0..) |*lw, layer| try runLayer(Tri, &pass, tri, lw, layer, ids, seq.fed); @@ -99,6 +99,7 @@ const Pass = struct { kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, + pg: @import("cuda_glue.zig").Pg, tail: cuda.DeviceBuffer, mixer: HcW, head_q: Q4, @@ -137,7 +138,7 @@ fn runLayer(comptime Tri: type, p: *Pass, tri: Tri, lw: *Layer, layer: usize, id try step.gdnStep(p.gpa, p.driver, p.stream, p.s, lw.proj, lw.out, p.conv[li], lw.cw.ptr, lw.a_log.ptr, lw.dt.ptr, lw.norm.ptr, &p.state[li]); } else { const ai = weights.attnIndex(layer); - try step.attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); + try step.attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t), p.pg); } try step.writeback(tri, p.s, h, 4, inj, y, wt); try step.readout(tri, p.s, lw.mlp_hc, h, inj); diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index 0e551809b3..96871a3cc1 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -6,6 +6,7 @@ const prompt = @import("cuda_prompt.zig"); const wt = @import("cuda_prompt_wt.zig"); const gdn = @import("cuda_gdn.zig"); const Draft = @import("cuda_mtp_draft.zig").Draft; +const pages = @import("cuda_pages.zig"); const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; const state_bytes = wt.state_n * 4; @@ -59,8 +60,28 @@ fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, pieces[7..11].* = attnPieces((&h.kc)[0..1], (&h.vc)[0..1], (&h.ikc)[0..1], (&h.pooled)[0..1], at); count = 11; } - for (pieces[0..count]) |p| for (p.bufs) |b| { + for (pieces[0..count], 0..) |p, pi| for (p.bufs, 0..) |b, layer| { if (p.n == 0) continue; + // the model's caches on the stream's pages; the head's too in a pool lane (the layer after the model's) + const head = pi >= 7 and d != null and d.?.pg.pt != 0; + if ((pi >= 2 and pi < 6) or head) if (seq.paged) |*pg| { + // the same bytes in the same order as a contiguous cache's, a run a page + const part: pages.Part = @fromBackingInt(@intCast(if (head) pi - 7 else pi - 2)); + const at_layer = if (head) wt.attn_n else layer; + const row = pages.rowBytes(part); + const per = pages.rows(part, pg.tokens); + const n = p.n / row; + var r: usize = 0; + while (r < n) { + const run = @min(per - r % per, n - r); + const at_row = pg.rowAt(at_layer, part, r); + const dst, const src = if (way == .save) .{ base + off, at_row } else .{ at_row, base + off }; + try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, run * row, stream.handle), "cuMemcpyDtoDAsync"); + off += run * row; + r += run; + } + continue; + }; const dst, const src = if (way == .save) .{ base + off, b.ptr } else .{ b.ptr, base + off }; try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, p.n, stream.handle), "cuMemcpyDtoDAsync"); off += p.n; diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 0dd79a8acd..2fae132c6d 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -26,6 +26,7 @@ test { _ = @import("cuda_vision_prompt.zig"); _ = reuse; _ = native; + _ = @import("cuda_pages.zig"); } pub const slide = @import("cuda_slide.zig"); diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig index 925d8183bb..0655054135 100644 --- a/zig/src/families/qwen3_5/cuda_cache.zig +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -11,11 +11,41 @@ const DeviceBuffer = cuda.DeviceBuffer; pub const tap_rows = 2048 + engine.max_rows; /// qwen35.cuh's Seg and Path. -pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, before: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int, delta: c_int }; +pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, before: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int, delta: c_int, pt: u64 = 0, ps: u64 = 0, ptok: c_int = 0 }; pub const Path = extern struct { n: c_int, rows: [engine.most_window]c_int }; +/// A cache's keys and values as pages (qwen35.cuh kv_row): page i of the stream's positions at `base` + ids[i] pages. +pub const Paged = struct { + /// Where page 0 of the arena starts: the core pool's, or the cache's own buffer under TF_QWEN_KV_PAGE_TEST. + base: u64, + /// The stream's page table on the device, an int a page, `ids` uploaded. + table: DeviceBuffer, + ids: []u32, + /// Positions a page holds. + tokens: usize, + /// The cache owns `ids` (TF_QWEN_KV_PAGE_TEST); a pool's lane maps its lease's. + own_ids: bool = true, + + /// Elements a page holds (bf16, or bytes for fp8 rows and their scales): every attention layer's of `tokens` rows. + pub fn elems(p: Paged, g: geometry.Geometry, kv8: bool) usize { + return if (kv8) g.kvBytes(p.tokens, true) else g.kvBytes(p.tokens, false) / 2; + } +}; + +/// TF_QWEN_KV_PAGE_TEST=n: each cache's own buffer read through a table of n-position pages in reverse order (checks). +pub fn pageTest() !usize { + const v = std.c.getenv("TF_QWEN_KV_PAGE_TEST") orelse return 0; + const n = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadQwenPageTest; + if (n == 0) return error.BadQwenPageTest; + return n; +} + pub const Cache = struct { untaught: bool = false, + /// Keys and values through a page table (null: cap contiguous rows a head). + paged: ?Paged = null, + /// A pool lane's drafter ring in its run (the drafter borrows it for the lane's stream). + drafter_ring: ?u64 = null, delta: i32 = 0, state: DeviceBuffer, /// The exact state before a recorded one-row round advances its stable slot in place. @@ -34,14 +64,80 @@ pub const Cache = struct { last: ?struct { start: usize, rows: usize } = null, /// An image prompt's rotary shift (get_rope_index's rope_delta): rows after its prompt rotate at position + delta. pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { + const test_pages = try pageTest(); + const rows = if (test_pages > 0) (cap + test_pages - 1) / test_pages * test_pages else cap; + var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(rows, e.kv8)); + errdefer kv.free(); + try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) + var paged: ?Paged = null; + errdefer if (paged) |*p| freePaged(gpa, p); + if (test_pages > 0) { + const n = rows / test_pages; + const ids = try gpa.alloc(u32, n); + errdefer gpa.free(ids); + for (ids, 0..) |*id, i| id.* = @intCast(n - 1 - i); // a table no contiguous layout matches + var table = try DeviceBuffer.alloc(e.d, n * 4); + errdefer table.free(); + try table.upload(0, std.mem.sliceAsBytes(ids)); + paged = .{ .base = kv.ptr, .table = table, .ids = ids, .tokens = test_pages }; + } + return fixed(gpa, e, cap, kv, paged); + } + + /// A pool lane's cache: a table of `pages` pages at `base`; its state, log and tap ring a run mapped at open. + pub fn initPaged(gpa: std.mem.Allocator, e: *const Engine, pages: usize, tokens: usize, base: u64) !Cache { + var table = try DeviceBuffer.alloc(e.d, pages * 4); + errdefer table.free(); + try table.fill8(0, null); + const none = try DeviceBuffer.alloc(e.d, 0); + return .{ .state = none, .before = none, .kv = none, .paged = .{ .base = base, .table = table, .ids = &.{}, .tokens = tokens, .own_ids = false }, .log = none, .cap = 0, .logits = try gpa.alloc(u16, e.g.vocab) }; + } + + /// Bytes of a pool lane's run: state, before, log, tap ring and a drafter's ring (`ring`), each 256-aligned. + pub fn runBytes(e: *const Engine, ring: usize) usize { + var n: usize = 0; + for (runParts(e, ring)) |len| n += std.mem.alignForward(usize, len, 256); + return n; + } + + fn runParts(e: *const Engine, ring: usize) [5]usize { + return .{ e.g.stateBytes(), e.g.stateBytes(), e.g.logBytes(engine.most_window), e.tapBytes(), ring }; + } + + /// A lane's state, before, log, taps and drafter ring as views of the run at `base`; a short run is refused. + pub fn mapRun(self: *Cache, e: *const Engine, base: u64, len: usize, ring: usize) !void { + if (runBytes(e, ring) > len) return error.QwenRunTooShort; + var at = base; + var views: [5]DeviceBuffer = undefined; + for (runParts(e, ring), &views) |n, *v| { + v.* = .{ .d = e.d, .ptr = at, .len = n, .borrowed = true }; + at += std.mem.alignForward(usize, n, 256); + } + self.state = views[0]; + self.before = views[1]; + self.log = views[2]; + self.tap = if (e.n_taps > 0) views[3] else null; + self.drafter_ring = if (ring > 0) views[4].ptr else null; + for ([_]DeviceBuffer{ views[2], views[3], views[4] }) |v| if (v.len > 0) try v.fill8(0, e.stream.handle); + } + + /// A closed lane's views dropped: its run went back to the pool. + pub fn unmapRun(self: *Cache, e: *const Engine) void { + const none: DeviceBuffer = .{ .d = e.d, .ptr = 0, .len = 0, .borrowed = true }; + self.state = none; + self.before = none; + self.log = none; + self.tap = null; + self.drafter_ring = null; + } + + /// The parts every cache holds itself beside `kv` and `paged` (the caller's to free on failure). + fn fixed(gpa: std.mem.Allocator, e: *const Engine, cap: usize, kv: DeviceBuffer, paged: ?Paged) !Cache { var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); errdefer state.free(); try state.fill8(0, null); var before = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); errdefer before.free(); - var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap, e.kv8)); - errdefer kv.free(); - try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) var log = try DeviceBuffer.alloc(e.d, e.g.logBytes(engine.most_window)); errdefer log.free(); try log.fill8(0, null); @@ -51,7 +147,7 @@ pub const Cache = struct { tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); try tap.?.fill8(0, null); } - return .{ .state = state, .before = before, .kv = kv, .tap = tap, .log = log, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + return .{ .state = state, .before = before, .kv = kv, .paged = paged, .tap = tap, .log = log, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; } /// Publish a completed forward's metadata, including one-row slots that already advanced in place. @@ -75,6 +171,7 @@ pub const Cache = struct { self.state.free(); if (self.before) |*b| b.free(); self.kv.free(); + if (self.paged) |*p| freePaged(gpa, p); self.log.free(); if (self.tap) |*t| t.free(); gpa.free(self.logits); @@ -83,7 +180,14 @@ pub const Cache = struct { /// The stream's window of `rows` rows from forward row `row0`, after `len` cached positions. pub fn seg(self: *const Cache, e: *const Engine, len: usize, row0: usize, rows: usize) Seg { const g = e.g; - return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .before = if (self.before) |b| b.ptr else 0, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps), .delta = self.delta }; + var sg: Seg = .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .before = if (self.before) |b| b.ptr else 0, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps), .delta = self.delta }; + if (self.paged) |p| { + sg.kv = p.base; + sg.pt = p.table.ptr; + sg.ps = p.elems(g, e.kv8); + sg.ptok = @intCast(p.tokens); + } + return sg; } /// Keeps `path` (window rows, ascending from 0): the state as that chain would leave it, a tree's keys moved down. @@ -125,3 +229,8 @@ pub const Cache = struct { try self.keep(e, all[0..last.rows]); } }; + +fn freePaged(gpa: std.mem.Allocator, p: *Paged) void { + p.table.free(); + if (p.own_ids) gpa.free(p.ids); +} diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index 6883c19fc7..f8124174c9 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -100,6 +100,10 @@ pub const Draft = struct { frame_streams: usize = 1, scratch_one: usize = 0, target_taps: usize = 0, + /// TF_QWEN_KV_PAGES: a stream's ring is a view of its target lane's run (null: the contexts own theirs). + ring_of: ?RingOf = null, + + pub const RingOf = struct { ctx: *anyopaque, ring: *const fn (*anyopaque, *Stream) ?u64 }; const words_at = frame.words_at; const hp_at = frame.hp_at; @@ -235,8 +239,9 @@ pub const Draft = struct { const started = profile.now(d.io); const previous_times = d.startup; try memory.resize(d, n); - try d.pool.ensureTotalCapacity(d.gpa, n); - while (d.pool.items.len < n) d.pool.appendAssumeCapacity(try d.create()); + const own = if (d.ring_of != null) @min(n, 1) else n; // with a pool the other streams' rings are lane runs + try d.pool.ensureTotalCapacity(d.gpa, own); + while (d.pool.items.len < own) d.pool.appendAssumeCapacity(try d.create()); d.startup.reserve += profile.elapsed(d.io, started) - (d.startup.warmups - previous_times.warmups) - (d.startup.reserve - previous_times.reserve); } @@ -376,6 +381,12 @@ pub const Draft = struct { fn openFn(p: *anyopaque, s: *Stream) anyerror!void { const d = cast(p); if (d.streams.contains(s)) return error.StreamOpen; + if (d.ring_of) |r| if (r.ring(r.ctx, s)) |ptr| { // the lane's run holds it: a context borrowing it + const ctx = try d.gpa.create(Ctx); + errdefer d.gpa.destroy(ctx); + ctx.* = .{ .ring = .{ .d = d.e.d, .ptr = ptr, .len = d.c.ringBytes(), .borrowed = true } }; + return d.streams.put(d.gpa, s, ctx); + }; const ctx = d.pool.pop() orelse try d.create(); errdefer d.pool.append(d.gpa, ctx) catch d.drop(ctx); ctx.* = .{ .ring = ctx.ring }; @@ -444,7 +455,10 @@ pub const Draft = struct { fn releaseFn(p: *anyopaque, s: *Stream) void { const d = cast(p); - if (d.streams.fetchRemove(s)) |kv| d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + if (d.streams.fetchRemove(s)) |kv| { + if (kv.value.ring.borrowed) return d.gpa.destroy(kv.value); // its lane's run goes back with the lane + d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + } } }; diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index d783449cb0..fc0c4d668e 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -13,6 +13,8 @@ const Draft = @import("cuda_draft.zig").Draft; const vision_prompt = @import("cuda_vision_prompt.zig"); const reuse = @import("cuda_reuse.zig"); const Mtp = @import("cuda_mtp.zig").Mtp; +const pages_mod = @import("cuda_pages.zig"); +const kv_pool = @import("core").kv_pool; const be = lanes.backend; pub const vision = @import("cuda_vision.zig"); pub const slide = @import("cuda_slide.zig"); @@ -25,6 +27,8 @@ pub const max_segments: u32 = 1; pub const prompt_rows: u32 = E.max_rows; /// --drafter loads the DFlash2 draft model (openDrafter); the checkpoints have no head of their own. pub const takes_drafter = true; +/// A request's pages hold one more forward past its prompt and reply (TF_QWEN_KV_PAGES). +pub const window_rows: u32 = pages_mod.window_rows; /// Kept prompt states for the shared cache (cuda_reuse.zig), in pinned host memory, on the zero-anchored chunk grid. pub const snap_bytes = reuse.bytes; @@ -58,11 +62,44 @@ pub const Loaded = struct { vision: ?*vision.Vision = null, /// What the prompt cache's saves copy: the engine and the streams' caches. target: *reuse.Target, + /// TF_QWEN_KV_PAGES: the streams beside the first take pages of one pool (null: a window each). + layout: ?kv_pool.Layout = null, }; +/// The pool's layout with a drafter (`extra` bytes a stream): DFlash2's ring and taps in the run, the rest fixed. +pub fn pagedLayout(p: *anyopaque, extra: u64) ?kv_pool.Layout { + const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.page_tokens == 0) return null; + const ring = l.ringBytes(); + const in_run = ring + if (l.draft != null) l.e.tapBytes() else 0; + l.extra = @intCast(extra -| in_run); + return pages_mod.layout(l.e, l.capacity, l.page_tokens, ring, l.extra); +} + +/// The pool beside the first stream's window: `streams` lanes and `pages` pages, made now. +pub fn attach(p: *anyopaque, streams: u32, pages: usize) !void { + const l: *Lanes = @ptrCast(@alignCast(p)); + const ring = l.ringBytes(); + l.pages = try pages_mod.Pages.init(l.gpa, l.e, pages_mod.layout(l.e, l.capacity, l.page_tokens, ring, l.extra), l.capacity, ring, streams, pages); + if (l.draft) |d| if (streams > 0) { + d.ring_of = .{ .ctx = l, .ring = Lanes.ringOf }; + }; + l.paged_only = pages_mod.pagedOnly() and streams > 0; // checks: every stream on pages, the reserved window idle +} + +/// The lane host's admission: the first stream's window is free, or a lane and the request's pages are. +pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { + const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.pool.items.len > 0 and !l.paged_only) return true; + const pg = l.pages orelse return true; + return pg.fits(prompt, max_tokens); +} + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.QwenHasNoCapturedKernels; if (o.slide and o.kv8) return error.UnsupportedLearningFormat; + const page_tokens = try pages_mod.paging(); + if (page_tokens > 0 and o.slide) return error.QwenPagesNoLearner; const e = try E.Engine.init(gpa, io, ctx, dir, o.kv8); errdefer e.deinit(); if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary @@ -71,7 +108,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ const v: ?*vision.Vision = if (o.vision) |lim| try vision.Vision.open(gpa, io, ctx.d, e.stream, dir, e.g.hidden, lim) else null; errdefer if (v) |x| x.deinit(); const l = try gpa.create(Lanes); - l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .cost_context = ctx, .startup = e.startup, .vision = v, .target = undefined }; + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .page_tokens = page_tokens, .cost_context = ctx, .startup = e.startup, .vision = v, .target = undefined }; l.target = .{ .e = e, .caches = &l.caches, .seen = &l.seen }; return .{ .learns = if (o.slide) e else null, @@ -85,6 +122,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .reserve = reserve, .vision = v, .target = &l.target, + .layout = if (page_tokens > 0) pages_mod.layout(e, o.context, page_tokens, 0, 0) else null, }; } @@ -100,7 +138,7 @@ pub fn laneFacts(window: u32, verify: u32) lanes.Model { fn reserve(p: *anyopaque, streams: u32) anyerror!void { const l: *Lanes = @ptrCast(@alignCast(p)); - try reservation.target(E.Cache, E.Part, l, streams); + try reservation.target(E.Cache, E.Part, l, streams); // with a pool: the first window, the shapes on its lanes } /// The DFlash2 drafter in `dir` over engine `p`, 4-bit only, and the target's windows timed for the round loop. pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void, prepare: struct { ctx: *anyopaque, run: *const fn (*anyopaque) anyerror!lanes.drafter.Facts, bytes: usize } } { @@ -113,6 +151,7 @@ pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []con d.target_taps = l.e.tapBytes(); cost_adapter.identify(l, dir, .dflash); l.calibration = d; + l.draft = d; return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter, .prepare = .{ .ctx = l, .run = calibrateDrafter, .bytes = if (l.capacity < l.e.window_rows) l.e.streamBytes(l.e.window_rows) else 0 } }; } @@ -168,6 +207,8 @@ fn closeDrafter(p: *anyopaque) void { pub fn explain(ctx: ?*anyopaque, err: anyerror) ?[]const u8 { if (ctx == null) return switch (err) { error.NoQwenVision => "--vision: this checkpoint has no image tower (no vision_config): serve it without --vision", + error.QwenPagesNoLearner => "TF_QWEN_KV_PAGES: the learner reads whole windows, not pages; drop --slide or unset it", + error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0 or unset: a window each)", error.UnsupportedQwenVision => "--vision: this checkpoint's image tower or text M-RoPE is not one the CUDA engine reads", else => null, }; @@ -224,6 +265,14 @@ const Lanes = struct { io: std.Io, e: *E.Engine, capacity: usize, + /// TF_QWEN_KV_PAGES's positions a page (0: no pool) and the pool `attach` made. + page_tokens: usize = 0, + pages: ?*pages_mod.Pages = null, + paged_only: bool = false, + /// The DFlash2 drafter, when one is attached: its rings go into the pool lanes' runs. + draft: ?*Draft = null, + /// A drafter's bytes a stream outside the lanes' runs (pagedLayout). + extra: usize = 0, cost_context: ?*const cuda.Context = null, cost_identity: ?[32]u8 = null, cost_kind: cost_adapter.Kind = .dflash, @@ -241,13 +290,14 @@ const Lanes = struct { fn deinit(self: *Lanes) void { var it = self.caches.valueIterator(); - while (it.next()) |cache| self.drop(cache.*); + while (it.next()) |cache| self.give(cache.*); // a pool lane's pages back, a window to the reserve self.caches.deinit(self.gpa); var seen = self.seen.valueIterator(); while (seen.next()) |v| self.gpa.free(v.*); self.seen.deinit(self.gpa); for (self.pool.items) |cache| self.drop(cache); self.pool.deinit(self.gpa); + if (self.pages) |p| p.deinit(); if (self.vision) |v| v.deinit(); self.e.deinit(); self.gpa.destroy(self); @@ -258,12 +308,13 @@ const Lanes = struct { self.gpa.destroy(cache); } - /// A cache for a new stream: a pooled one reset, else one more (a window beyond the reserve). - fn take(self: *Lanes) !*E.Cache { - if (self.pool.pop()) |cache| { + /// A cache for a new stream: a pooled one reset, a pool lane with its pages, else one more window. + fn take(self: *Lanes, prompt: usize, max_new: usize) !*E.Cache { + if (!self.paged_only) if (self.pool.pop()) |cache| { try cache.reset(self.e); return cache; - } + }; + if (self.pages) |p| return p.open(prompt, max_new); // LaneMemory: the lane host waits for a stream to end const cache = try self.gpa.create(E.Cache); errdefer self.gpa.destroy(cache); cache.* = try E.Cache.init(self.gpa, self.e, self.capacity); @@ -271,6 +322,7 @@ const Lanes = struct { } fn give(self: *Lanes, cache: *E.Cache) void { + if (self.pages) |p| if (p.close(cache)) return; self.pool.append(self.gpa, cache) catch self.drop(cache); } @@ -283,6 +335,18 @@ const Lanes = struct { return .{ .ptr = self, .first_before_draft = true, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .prepare_features = prepareFeaturesFn, .features = featuresFn, .release = releaseFn } }; } + /// A drafter ring's bytes in each lane's run (0: no DFlash2 drafter). + fn ringBytes(self: *const Lanes) usize { + return if (self.draft) |d| d.c.ringBytes() else 0; + } + + /// The ring in the run of the lane `stream`'s cache holds; null: its cache is a reserved window's. + fn ringOf(ptr: *anyopaque, stream: *lanes.Stream) ?u64 { + const self = cast(ptr); + const cache = self.caches.get(stream) orelse return null; + return cache.drafter_ring; + } + fn cast(ptr: *anyopaque) *Lanes { return @ptrCast(@alignCast(ptr)); } @@ -309,7 +373,7 @@ const Lanes = struct { return err; }; errdefer if (self.seen.fetchRemove(stream)) |kv| self.gpa.free(kv.value); - const cache = try self.take(); + const cache = try self.take(ids.len, stream.max_new); errdefer self.give(cache); // the stream's cache is found by a save at a mark (cuda_reuse.zig) while its pass runs try self.caches.put(self.gpa, stream, cache); diff --git a/zig/src/families/qwen3_5/cuda_native_moe.zig b/zig/src/families/qwen3_5/cuda_native_moe.zig index 5d0f4b0c82..3bc6155bd5 100644 --- a/zig/src/families/qwen3_5/cuda_native_moe.zig +++ b/zig/src/families/qwen3_5/cuda_native_moe.zig @@ -10,6 +10,10 @@ pub const prompt_rows = native.prompt_rows; pub const takes_drafter = true; pub const Options = native.Options; pub const open = native.open; +pub const attach = native.attach; +pub const pagedLayout = native.pagedLayout; // TF_QWEN_KV_PAGES: the same pool as the dense models' +pub const fits = native.fits; +pub const window_rows = native.window_rows; pub const openDrafter = native.openMtp; pub const explain = native.explain; /// The dense family's kept prompt states, its open refusals' words and its image tower: the MLP is the only change. diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig new file mode 100644 index 0000000000..e4436ae022 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -0,0 +1,169 @@ +//! Qwen3.5's attention caches from the core pool: the streams beside the first take pages of one device arena. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const engine = @import("cuda_engine.zig"); +const caches = @import("cuda_cache.zig"); + +const Engine = engine.Engine; +const Cache = caches.Cache; + +/// Rows past a request's prompt and reply its pages hold: one more forward's (a verify window or a prompt chunk). +pub const window_rows: u32 = engine.max_rows; + +/// TF_QWEN_KV_PAGES=n: streams beside the first take n-position pages from one pool (unset or 0: a window each). +pub fn paging() !usize { + const v = std.c.getenv("TF_QWEN_KV_PAGES") orelse return 0; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch error.BadQwenKvPages; +} + +/// TF_QWEN_KV_PAGES_ALL=1: every stream on pages, the first stream's reserved window left idle (checks). +pub fn pagedOnly() bool { + const v = std.c.getenv("TF_QWEN_KV_PAGES_ALL") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// Pages a stream maps at most: a full window and one more forward. +fn mostPages(capacity: usize, tokens: usize) usize { + return (capacity + window_rows + tokens - 1) / tokens; +} + +/// A lane's table and `fixed` bytes; its keys and values a token; a run of its state, log, taps and a `ring`-byte ring. +pub fn layout(e: *const Engine, capacity: usize, tokens: usize, ring: usize, fixed: usize) kv.Layout { + return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + fixed, .token_bytes = e.g.kvBytes(1, e.kv8), .page_tokens = @intCast(tokens), .snap_bytes = e.g.stateBytes(), .run_bytes = Cache.runBytes(e, ring) }; +} + +/// One pool lane: a cache whose table maps its stream's pages, that stream's lease and its run's first page. +const Lane = struct { cache: *Cache, lease: kv.Lease = .{}, used: bool = false, run: u32 = 0 }; + +pub const Pages = struct { + gpa: std.mem.Allocator, + e: *Engine, + layout: kv.Layout, + capacity: usize, + ring: usize, + pool: kv.Pool, + arena: cuda.DeviceBuffer, + lanes: []Lane, + + /// `streams` lanes and `pages` pages of layout `l`, made and touched now; each run holds a `ring`-byte ring. + pub fn init(gpa: std.mem.Allocator, e: *Engine, l: kv.Layout, capacity: usize, ring: usize, streams: usize, pages: usize) !*Pages { + const tokens: usize = l.page_tokens; + const p = try gpa.create(Pages); + errdefer gpa.destroy(p); + var arena = try cuda.DeviceBuffer.alloc(e.d, pages * l.pageBytes()); + errdefer arena.free(); + try arena.fill8(0, null); // committed now (the driver commits on first touch) + p.* = .{ .gpa = gpa, .e = e, .layout = l, .capacity = capacity, .ring = ring, .pool = try kv.Pool.init(gpa, pages, l.pageBytes()), .arena = arena, .lanes = &.{} }; + errdefer p.pool.deinit(); + p.lanes = try gpa.alloc(Lane, streams); + errdefer gpa.free(p.lanes); + var made: usize = 0; + errdefer for (p.lanes[0..made]) |x| { + x.cache.deinit(gpa); + gpa.destroy(x.cache); + }; + for (p.lanes) |*x| { + const c = try gpa.create(Cache); + errdefer gpa.destroy(c); + c.* = try Cache.initPaged(gpa, e, mostPages(capacity, tokens), tokens, arena.ptr); + x.* = .{ .cache = c }; + made += 1; + } + return p; + } + + pub fn deinit(p: *Pages) void { + p.e.stream.synchronize() catch std.log.warn("qwen3_5: the pool's stream did not drain at close", .{}); + for (p.lanes) |*x| { + if (x.used) { + p.pool.releaseRun(x.run, p.layout.runPages()); + p.pool.finish(&x.lease); + } + x.cache.deinit(p.gpa); + p.gpa.destroy(x.cache); + } + p.gpa.free(p.lanes); + p.arena.free(); + p.pool.deinit(); + p.gpa.destroy(p); + } + + fn need(p: *const Pages, prompt: usize, max_tokens: usize) usize { + return p.layout.need(prompt, max_tokens, window_rows); + } + + /// Whether a lane is free, the pool can promise a request's pages now and a run of them is free. + pub fn fits(p: *const Pages, prompt: usize, max_tokens: usize) bool { + for (p.lanes) |x| if (!x.used) return p.pool.fits(p.need(prompt, max_tokens)) and p.pool.findRun(p.layout.runPages()) != null; + return false; + } + + /// A free lane's cache, reset, holding the pages of `prompt` + `max_tokens` rows and a window; LaneMemory: not now. + pub fn open(p: *Pages, prompt: usize, max_tokens: usize) !*Cache { + const x = for (p.lanes) |*x| { + if (!x.used) break x; + } else return error.LaneMemory; + const k = p.layout.runPages(); + const n = p.need(prompt, max_tokens) - k; // keys and values + if (n > mostPages(p.capacity, p.layout.page_tokens)) return error.InvalidQwenPrompt; + try p.pool.promise(&x.lease, n + k); + errdefer p.pool.finish(&x.lease); + // the run first: KV pages fill from the bottom, runs from the top; a fragmented pool waits (LaneMemory) + const run = p.pool.claimRun(&x.lease, k) orelse return error.LaneMemory; + errdefer p.pool.releaseRun(run, k); + var up: Upload = .{ .p = p, .cache = x.cache }; + try p.pool.cover(&x.lease, n, .{ .ctx = &up, .write = Upload.write }); + try x.cache.mapRun(p.e, p.arena.ptr + run * p.layout.pageBytes(), k * p.layout.pageBytes(), p.ring); + try x.cache.reset(p.e); + x.cache.paged.?.ids = x.lease.ids.items; + x.cache.cap = n * p.layout.page_tokens; + x.run = run; + x.used = true; + return x.cache; + } + + /// The stream on `cache` ended: its pages go back once the forwards that read them are done. False: not a lane's. + pub fn close(p: *Pages, cache: *Cache) bool { + const x = for (p.lanes) |*x| { + if (x.cache == cache) break x; + } else return false; + if (!x.used) return true; + // a page another stream takes next may still be read by queued work on the engine's stream + p.e.stream.synchronize() catch std.log.warn("qwen3_5: a closed stream's forwards did not drain; its pages go back", .{}); + p.pool.releaseRun(x.run, p.layout.runPages()); + p.pool.finish(&x.lease); + cache.unmapRun(p.e); + cache.paged.?.ids = &.{}; + cache.cap = 0; + x.used = false; + return true; + } + + /// Whether `cache` is one of the pool's lanes. + pub fn owns(p: *const Pages, cache: *const Cache) bool { + for (p.lanes) |x| if (x.cache == cache) return true; + return false; + } + + /// A lane's new table entries uploaded to its device table. + const Upload = struct { + p: *Pages, + cache: *Cache, + + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const u: *Upload = @ptrCast(@alignCast(ctx)); + const t = u.cache.paged.?.table; + try t.uploadAsync(from * 4, std.mem.sliceAsBytes(ids[from..]), u.p.e.stream.handle); + try u.p.e.stream.synchronize(); // the ids slice may move as the lease grows + } + }; +}; + +test "a stream's pages: its prompt, reply and one more forward, no more than a full window's" { + const l: kv.Layout = .{ .fixed_bytes = 1, .token_bytes = 64 << 10, .page_tokens = 256, .snap_bytes = 1 }; + try std.testing.expectEqual(@as(usize, 1), l.need(1, 0, window_rows)); + try std.testing.expectEqual(@as(usize, 2), l.need(100, 29, window_rows)); + try std.testing.expectEqual(@as(usize, 129), mostPages(32768, 256)); + try std.testing.expect(l.need(32768 - 1, 1, window_rows) <= mostPages(32768, 256)); +} diff --git a/zig/src/families/qwen3_5/cuda_reserve.zig b/zig/src/families/qwen3_5/cuda_reserve.zig index c6fa09e916..71221f8f45 100644 --- a/zig/src/families/qwen3_5/cuda_reserve.zig +++ b/zig/src/families/qwen3_5/cuda_reserve.zig @@ -2,11 +2,15 @@ const warm = @import("cuda_warm.zig"); const E = @import("cuda_engine.zig"); +/// `streams` windows reserved and warmed; with a pool the first only, the shapes warmed on one-row lanes. pub fn target(comptime Cache: type, comptime Part: type, l: anytype, streams: u32) !void { const started = warm.profile.now(l.io); var warm_ns: i96 = 0; - try l.pool.ensureTotalCapacity(l.gpa, streams); - while (l.pool.items.len < streams) { + const paged = comptime @hasField(@TypeOf(l.*), "pages"); + const pooled = if (paged) l.pages != null else false; + const windows: u32 = if (pooled) @min(streams, 1) else streams; + try l.pool.ensureTotalCapacity(l.gpa, windows); + while (l.pool.items.len < windows) { const cache = try l.gpa.create(Cache); errdefer l.gpa.destroy(cache); cache.* = try Cache.init(l.gpa, l.e, l.capacity); @@ -23,18 +27,17 @@ pub fn target(comptime Cache: type, comptime Part: type, l: anytype, streams: u3 var parts: [E.max_segs]Part = undefined; var ids: [@max(E.max_segs, E.most_window)]u32 = @splat(0); var picks: [@max(E.max_segs, E.most_window)]E.Draw = undefined; - for (0..@min(streams, E.max_segs)) |n| { - parts[n] = .{ .cache = l.pool.items[n], .rows = 1 }; - picks[n] = .{ .rule = null, .position = 1 }; - const shape = E.Engine.Shape.init(n + 1, n + 1, true, l.e.n_taps > 0, 1); - if (l.warmed.contains(shape)) continue; - const warm_start = warm.profile.now(l.io); - try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); - try l.e.stream.synchronize(); - warm_ns += warm.profile.elapsed(l.io, warm_start); - for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); - l.warmed.mark(shape); - } + var warming: [E.max_segs]*Cache = undefined; + var opened: usize = 0; + const most = @min(streams, E.max_segs); + for (l.pool.items[0..@min(windows, most)], 0..) |cache, i| warming[i] = cache; + if (comptime paged) if (l.pages) |p| { + while (windows + opened < most) : (opened += 1) warming[windows + opened] = p.open(1, 0) catch break; + }; + defer if (comptime paged) if (l.pages) |p| for (warming[windows..][0..opened]) |cache| { + _ = p.close(cache); + }; + try warmShapes(Part, l, warming[0 .. @min(windows, most) + opened], &parts, &ids, &picks, &warm_ns); // Pending calibration warms on its capacity-safe cache during synchronous open. if (l.e.n_taps > 0 and l.calibration == null) for (2..l.e.window_rows + 1) |rows| { const shape = E.Engine.Shape.init(rows, 1, true, true, 1); @@ -51,4 +54,20 @@ pub fn target(comptime Cache: type, comptime Part: type, l: anytype, streams: u3 l.startup.reserve += warm.profile.elapsed(l.io, started) - warm_ns; } +/// Each shape of 1..caches.len one-row segments run once on `caches`, then reset. +fn warmShapes(comptime Part: type, l: anytype, caches: anytype, parts: []Part, ids: []u32, picks: []E.Draw, warm_ns: *i96) !void { + for (caches, 0..) |cache, n| { + parts[n] = .{ .cache = cache, .rows = 1 }; + picks[n] = .{ .rule = null, .position = 1 }; + const shape = E.Engine.Shape.init(n + 1, n + 1, true, l.e.n_taps > 0, 1); + if (l.warmed.contains(shape)) continue; + const warm_start = warm.profile.now(l.io); + try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); + try l.e.stream.synchronize(); + warm_ns.* += warm.profile.elapsed(l.io, warm_start); + for (caches[0 .. n + 1]) |c| try c.reset(l.e); + l.warmed.mark(shape); + } +} + pub const Test = if (@import("builtin").is_test) @import("cuda_reserve_test.zig") else struct {}; diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig index 8b7eb120ee..6d715b1d5e 100644 --- a/zig/src/families/qwen3_5/cuda_reuse.zig +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -4,6 +4,7 @@ const cuda = @import("cuda"); const lanes = @import("lanes"); const geometry = @import("cuda_geometry.zig"); const engine = @import("cuda_engine.zig"); +const caches = @import("cuda_cache.zig"); const Engine = engine.Engine; const Cache = engine.Cache; @@ -90,6 +91,42 @@ const Layout = struct { } }; +/// A piece's device spans: one, or with a page table a run a page (the host copy keeps the contiguous order). +const Runs = struct { + base: u64, + pages: ?caches.Paged, + page_bytes: usize, + in_page: usize, + span: usize = 0, // a full page's run of the piece + len: usize, + page: usize = 0, + done: usize = 0, + + const Run = struct { dev: u64, at: usize, len: usize }; + + fn of(cache: *const Cache, l: Layout, i: usize, p: Piece) Runs { + if (!p.kv or cache.paged == null) return .{ .base = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off, .pages = null, .page_bytes = 0, .in_page = 0, .len = p.len }; + const pg = cache.paged.?; + // the piece's offset inside a page: its layout at cap = tokens (Layout.piece), a run of tokens rows a page + const in_page: Layout = .{ .g = l.g, .format = l.format, .cap = pg.tokens, .at = pg.tokens }; + const q = in_page.piece(i); + return .{ .base = pg.base, .pages = pg, .page_bytes = in_page.kvBytes(), .in_page = q.off, .span = q.len, .len = p.len }; + } + + fn next(r: *Runs) ?Run { + if (r.done >= r.len) return null; + const pg = r.pages orelse { + r.done = r.len; + return .{ .dev = r.base, .at = 0, .len = r.len }; + }; + const n = @min(r.span, r.len - r.done); + const run: Run = .{ .dev = r.base + pg.ids[r.page] * r.page_bytes + r.in_page, .at = r.done, .len = n }; + r.page += 1; + r.done += n; + return run; + } +}; + /// Whether a request's images match a kept state's: the same identities at the same rows, for every image before `at`. pub fn sameImages(kept: []const Seen, request: []const Seen, at: usize) bool { var n: usize = 0; @@ -135,8 +172,8 @@ pub fn saveCache(e: *Engine, cache: *const Cache, at: usize, seen: []const Seen) var off: usize = 0; for (0..l.count()) |i| { const p = l.piece(i); - const dev = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off; - if (p.len > 0) try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off, dev, p.len, e.stream.handle), "cuMemcpyDtoHAsync"); + var it = Runs.of(cache, l, i, p); + while (it.next()) |r| try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off + r.at, r.dev, r.len, e.stream.handle), "cuMemcpyDtoHAsync"); off += p.len; } try e.stream.synchronize(); @@ -174,8 +211,8 @@ pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque, seen: []const var off: usize = 0; for (0..l.count()) |i| { const p = l.piece(i); - const dev = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off; - if (p.len > 0) try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(dev, copy.host.bytes.ptr + off, p.len, e.stream.handle), "cuMemcpyHtoDAsync"); + var it = Runs.of(cache, l, i, p); + while (it.next()) |r| try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(r.dev, copy.host.bytes.ptr + off + r.at, r.len, e.stream.handle), "cuMemcpyHtoDAsync"); off += p.len; } try e.stream.synchronize(); // the host copy may be dropped once this returns @@ -294,6 +331,7 @@ test "a restore refuses another cache format, another length, a position past th if (comptime @hasField(Engine, "kv8")) e.kv8 = false; var cache: Cache = undefined; cache.cap = 64; + cache.paged = null; // contiguous: restoreCache reads it, and undefined bytes can read as a page table const d: u64 = 9; var images = [_]Seen{.{ .at = 2, .id = d }}; const l: Layout = .{ .g = e.g, .format = .bf16, .cap = 64, .at = 10 }; @@ -352,6 +390,7 @@ test "a failed restore's host copy stays through a drop until a wait, the restor if (comptime @hasField(Engine, "kv8")) e.kv8 = false; var cache: Cache = undefined; cache.cap = 64; + cache.paged = null; // contiguous: restoreCache reads it, and undefined bytes can read as a page table cache.kv = .{ .d = &d, .ptr = 0x10000, .len = 0 }; cache.state = .{ .d = &d, .ptr = 0x20000, .len = 0 }; const l: Layout = .{ .g = e.g, .format = .bf16, .cap = 64, .at = 10 }; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index a8463051a5..860fcf65c4 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -346,10 +346,11 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; const free = loaded.free_streams; const grow = budget.growRoom(room, !@hasField(F.Options, "lanes")); - const layout = kv_pages.layout(loaded); + const wrapped = kv_pages.wrapped(F, loaded, free, served.stream_bytes, room); + const layout = wrapped.layout; var paged: ?kv_pages.Plan = null; const streams = if (layout) |l| blk: { - paged = try kv_pages.plan(F, a, l, room, o, @intCast(window), problem, host_cache) orelse return null; + paged = try kv_pages.plan(F, a, l, wrapped.room, o, @intCast(window), problem, host_cache) orelse return null; break :blk paged.?.lanes; } else budget.admit(grow, served.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { diff --git a/zig/src/native/cuda_kv_pool.zig b/zig/src/native/cuda_kv_pool.zig index 754602de1d..17f52e2bc2 100644 --- a/zig/src/native/cuda_kv_pool.zig +++ b/zig/src/native/cuda_kv_pool.zig @@ -12,6 +12,32 @@ pub fn layout(loaded: anytype) ?kv_pool.Layout { return if (comptime @hasField(@TypeOf(loaded), "layout")) loaded.layout else null; } +/// The plan's view of a wrapped family: drafter buffers join the fixed part; no free stream pays its window first. +pub fn served(l: ?kv_pool.Layout, free_streams: u32, family_bytes: u64, stream_bytes: u64, room: u64) struct { layout: ?kv_pool.Layout, room: u64 } { + var out = l orelse return .{ .layout = null, .room = room }; + out.fixed_bytes += stream_bytes -| family_bytes; + return .{ .layout = out, .room = if (free_streams == 0) room -| stream_bytes else room }; +} + +test "a wrapped family's plan: drafter buffers a stream, a first window without a free stream" { + const l: kv_pool.Layout = .{ .fixed_bytes = 10, .token_bytes = 1, .page_tokens = 4, .snap_bytes = 1 }; + const nemotron = served(l, 1, 500, 500, 1000); // its own sequence is in the model's bytes: nothing moves + try std.testing.expectEqual(@as(u64, 10), nemotron.layout.?.fixed_bytes); + try std.testing.expectEqual(@as(u64, 1000), nemotron.room); + const qwen = served(l, 0, 500, 530, 1000); // a drafter's 30 bytes a stream, the first stream's 530 first + try std.testing.expectEqual(@as(u64, 40), qwen.layout.?.fixed_bytes); + try std.testing.expectEqual(@as(u64, 470), qwen.room); + try std.testing.expectEqual(@as(?kv_pool.Layout, null), served(null, 0, 1, 2, 3).layout); +} + +/// served() on the family's layout; a family that places a drafter's buffers itself (pagedLayout) plans it again. +pub fn wrapped(comptime F: type, loaded: anytype, free: u32, stream_bytes: u64, room: u64) @TypeOf(served(null, 0, 0, 0, 0)) { + const own = comptime @hasDecl(F, "pagedLayout"); + const l = layout(loaded); + const planned = if (own) (if (l != null) F.pagedLayout(loaded.ctx, stream_bytes - loaded.stream_bytes) else null) else l; + return served(planned, free, if (own) stream_bytes else loaded.stream_bytes, stream_bytes, room); +} + /// The real native startup room before page planning and cache allocation. pub fn startupRoom(comptime F: type, a: Allocator, loaded: anytype, o: api.Open, available: u64, host_budget: u64, problem: *[]const u8) !?u64 { const memory = @import("cuda_family_memory.zig"); diff --git a/zig/tests/cuda/contract_fn_pages.zig b/zig/tests/cuda/contract_fn_pages.zig new file mode 100644 index 0000000000..713670f7b8 --- /dev/null +++ b/zig/tests/cuda/contract_fn_pages.zig @@ -0,0 +1,344 @@ +//! flashnext_attn.cu's caches through a page table: prep, pooling, scores and attention, paged == contiguous, float64. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const q_heads = 24; +const kv_heads = 2; +const hd = 256; +const index_heads = 4; +const index_dim = 128; +const half = 32; +const proj_n = 13952; +const ratio = 4; +const idw = 2052; +const nch = 5; +const top = 512; +const eps: f32 = 1e-6; +const layers = 2; // a page's attention layers: the kernels run layer 1, so the layer offset counts + +const kv_row = kv_heads * hd * 2; +const index_row = index_dim * 2; + +const Part = enum { key, value, index, pooled }; + +fn rowBytes(p: Part) usize { + return if (p == .key or p == .value) kv_row else index_row; +} + +fn perPage(p: Part, t: usize) usize { + return if (p == .pooled) t / ratio else t; +} + +/// A page's layout, written out apart from cuda_pages.zig: per layer keys, values, indexer keys, then pooled rows. +fn layerBytes(t: usize) usize { + return t * (2 * kv_row + index_row) + t / ratio * index_row; +} + +fn partOffset(t: usize, p: Part) usize { + return layerBytes(t) * 1 + switch (p) { + .key => 0, + .value => t * kv_row, + .index => 2 * t * kv_row, + .pooled => 2 * t * kv_row + t * index_row, + }; +} + +/// Two streams' cached positions, the rows each adds, positions a page. +const Case = struct { lens: [2]usize, rows: usize, ptok: usize }; +const cases = [_]Case{ + .{ .lens = .{ 0, 3 }, .rows = 1, .ptok = 16 }, + .{ .lens = .{ 127, 128 }, .rows = 4, .ptok = 16 }, + .{ .lens = .{ 700, 33 }, .rows = 3, .ptok = 48 }, // a page no 512-key chunk or 32-key tile divides + .{ .lens = .{ 2100, 2300 }, .rows = 16, .ptok = 128 }, // block scores and dense attention together + .{ .lens = .{ 4095, 5000 }, .rows = 17, .ptok = 256 }, + .{ .lens = .{ 9000, 300 }, .rows = 2, .ptok = 512 }, +}; + +const Fixture = struct { + c: Case, + cap: usize, + caches: [2][4][]u8, // each stream's logical key, value, indexer key and pooled rows + p: [2][]u16, // projections [rows][proj_n] + qw: []f32, + kw: []f32, + iw: []f32, + ikw: []f32, + inv: []f32, + ids: [2][]i32, // sparse rows' key lists [rows][idw] + nk: [2][]i32, + ids_pages: [2][]const u32, // each stream's pages in the shared arena, shuffled across both + pages: usize, // the arena's pages (two spare) + + fn pageBytes(f: Fixture) usize { + return layers * layerBytes(f.c.ptok); + } + + /// Byte offset of `row` of `part` of stream `st` in the arena under `ids`. + fn paged(f: Fixture, ids: [2][]const u32, st: usize, part: Part, row: usize) usize { + const per = perPage(part, f.c.ptok); + return ids[st][row / per] * f.pageBytes() + partOffset(f.c.ptok, part) + (row % per) * rowBytes(part); + } + + fn rowsOf(f: Fixture, part: Part) usize { + return perPage(part, f.cap); + } + + fn make(a: std.mem.Allocator, c: Case, seed: u64) !Fixture { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const most = @max(c.lens[0], c.lens[1]) + c.rows; + const cap = (most + c.ptok - 1) / c.ptok * c.ptok; + var f: Fixture = undefined; + f.c = c; + f.cap = cap; + for (0..2) |st| { + for (std.enums.values(Part), 0..) |part, pi| { + f.caches[st][pi] = try a.alloc(u8, f.rowsOf(part) * rowBytes(part)); + fillBf(rng, std.mem.bytesAsSlice(u16, f.caches[st][pi]), 1); + } + f.p[st] = try a.alloc(u16, c.rows * proj_n); + fillBf(rng, f.p[st], 2); + f.ids[st] = try a.alloc(i32, c.rows * idw); + f.nk[st] = try a.alloc(i32, c.rows); + for (0..c.rows) |r| { + const keys = c.lens[st] + r + 1; + f.nk[st][r] = @intCast(@max(1, @min(idw, keys) -| r % 3)); + for (0..idw) |i| f.ids[st][r * idw + i] = @intCast((i * 1237 + r * 31 + st * 7) % keys); + } + } + f.qw = try floats(a, rng, hd, 0.5, 1.5); + f.kw = try floats(a, rng, hd, 0.5, 1.5); + f.iw = try floats(a, rng, index_dim, 0.5, 1.5); + f.ikw = try floats(a, rng, index_dim, 0.5, 1.5); + f.inv = try a.alloc(f32, half); + for (f.inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / half)); + f.pages = 2 * (cap / c.ptok) + 2; + const order = try a.alloc(u32, f.pages); + for (order, 0..) |*o, i| o.* = @intCast(i); + rng.shuffle(u32, order); + const n = cap / c.ptok; + f.ids_pages = .{ order[0..n], order[n .. 2 * n] }; + return f; + } + + /// Both streams' logical caches at their pages under `ids`; every other byte (layer 0, spare pages) random. + fn arena(f: Fixture, a: std.mem.Allocator, rng: std.Random, ids: [2][]const u32) ![]u8 { + const out = try a.alloc(u8, f.pages * f.pageBytes()); + rng.bytes(out); + for (0..2) |st| for (std.enums.values(Part), 0..) |part, pi| { + const n = rowBytes(part); + for (0..f.rowsOf(part)) |row| @memcpy(out[f.paged(ids, st, part, row)..][0..n], f.caches[st][pi][row * n ..][0..n]); + }; + return out; + } + + /// Stream `st`'s logical caches read back out of the arena's pages. + fn unpage(f: Fixture, a: std.mem.Allocator, all: []const u8, ids: [2][]const u32, st: usize) ![4][]u8 { + var out: [4][]u8 = undefined; + for (std.enums.values(Part), 0..) |part, pi| { + const n = rowBytes(part); + out[pi] = try a.alloc(u8, f.rowsOf(part) * n); + for (0..f.rowsOf(part)) |row| @memcpy(out[pi][row * n ..][0..n], all[f.paged(ids, st, part, row)..][0..n]); + } + return out; + } +}; + +fn fillBf(rng: std.Random, out: []align(1) u16, scale: f32) void { + for (out) |*x| x.* = @truncate(@as(u32, @bitCast((rng.float(f32) * 2 - 1) * scale)) >> 16); +} + +fn floats(a: std.mem.Allocator, rng: std.Random, n: usize, lo: f32, hi: f32) ![]f32 { + const v = try a.alloc(f32, n); + for (v) |*x| x.* = lo + rng.float(f32) * (hi - lo); + return v; +} + +fn bf(x: u16) f64 { + return @as(f32, @bitCast(@as(u32, x) << 16)); +} + +fn scoreStride(cap: usize) usize { + return ((cap + ratio - 1) / ratio + 15) / 16 * 16; +} + +/// One layout's results a stream: queries, indexer queries, block scores, attention rows and the caches after. +const Result = struct { q: [2][]u8, iq: [2][]u8, sc: [2][]u8, dense: [2]?[]u8, sparse: [2][]u8, caches: [2][4][]u8 }; + +/// Both streams' rows through prep, pool, scores and attention on layer 1: contiguous, or paged by `tables` (`given`). +fn launch(gpu: check.Gpu, s: cuda.Stream, g: *const flash.Glue, f: Fixture, tables: ?[2][]const u32, given: ?[2][]const u32, a: std.mem.Allocator, seed: u64) !Result { + var rig = rig_.Rig{ .gpu = gpu, .s = s }; + defer rig.deinit(); + var prng = std.Random.DefaultPrng.init(seed); + const rows = f.c.rows; + const nb = scoreStride(f.cap); + const qw = try rig.input(std.mem.sliceAsBytes(f.qw)); + const kw = try rig.input(std.mem.sliceAsBytes(f.kw)); + const iw = try rig.input(std.mem.sliceAsBytes(f.iw)); + const ikw = try rig.input(std.mem.sliceAsBytes(f.ikw)); + const inv = try rig.input(std.mem.sliceAsBytes(f.inv)); + const po = try rig.scratch(rows * nch * q_heads * hd * 4); + const pm = try rig.scratch(rows * nch * q_heads * 4); + const pl = try rig.scratch(rows * nch * q_heads * 4); + var arena_ptr: u64 = 0; + if (tables) |t| { + const host = try f.arena(a, prng.random(), t); + arena_ptr = try rig.scratch(host.len); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = arena_ptr, .len = host.len }).upload(0, host); + } + var res: Result = undefined; + var flat: [2][4]u64 = undefined; + var outs: [2][5]u64 = undefined; + for (0..2) |st| { + const pos0: i32 = @intCast(f.c.lens[st]); + const ends = f.c.lens[st] + rows; + const pos = try rig.input(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); + const p = try rig.input(std.mem.sliceAsBytes(f.p[st])); + const ids = try rig.input(std.mem.sliceAsBytes(f.ids[st])); + const nk = try rig.input(std.mem.sliceAsBytes(f.nk[st])); + const ones = try a.alloc(i32, rows); + @memset(ones, 1); + const spr1 = try rig.input(std.mem.sliceAsBytes(ones)); + @memset(ones, 0); + const spr0 = try rig.input(std.mem.sliceAsBytes(ones)); + var pg: flash.Pg = .{}; + var at: [4]u64 = undefined; + if (tables) |t| { + pg = .{ .pt = try rig.input(std.mem.sliceAsBytes((given orelse t)[st])), .ps = @intCast(f.pageBytes() / 2), .ptok = @intCast(f.c.ptok) }; + for (std.enums.values(Part), 0..) |part, pi| at[pi] = arena_ptr + partOffset(f.c.ptok, part); + } else for (0..4) |pi| { + at[pi] = try rig.scratch(f.caches[st][pi].len); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = at[pi], .len = f.caches[st][pi].len }).upload(0, f.caches[st][pi]); + } + flat[st] = at; + outs[st] = .{ try rig.output(rows * q_heads * hd * 2), try rig.output(rows * index_heads * index_dim * 2), try rig.output(rows * nb * 4), try rig.output(rows * q_heads * hd * 2), try rig.output(rows * q_heads * hd * 2) }; + const o = outs[st]; + try g.attnPrep(s, p, pos, qw, kw, iw, inv, o[0], at[0], at[1], o[1], at[2], eps, rows, pg); + try g.qsaPool(s, at[2], at[3], pos, ikw, inv, eps, rows, pg); + if (ends / ratio > top) try g.qsaScores(s, o[1], at[3], pos, o[2], nb, rows, @min(nb, (ends + ratio - 1) / ratio), pg); + if (ends <= nch * 512) { + try g.attnChunks(s, o[0], at[0], at[1], pos, po, pm, pl, ids, nk, spr0, rows, kv_heads, (ends + 511) / 512, true, pg); + try g.attnMerge(s, po, pm, pl, pos, o[3], nk, spr0, rows, kv_heads, true); + } + try g.attnChunks(s, o[0], at[0], at[1], pos, po, pm, pl, ids, nk, spr1, rows, kv_heads, (idw + 511) / 512, true, pg); + try g.attnMerge(s, po, pm, pl, pos, o[4], nk, spr1, rows, kv_heads, true); + } + for (0..2) |st| { + const o = outs[st]; + res.q[st] = try a.dupe(u8, try rig.bytes(o[0], rows * q_heads * hd * 2)); + res.iq[st] = try a.dupe(u8, try rig.bytes(o[1], rows * index_heads * index_dim * 2)); + res.sc[st] = try a.dupe(u8, try rig.bytes(o[2], rows * nb * 4)); + res.dense[st] = if (f.c.lens[st] + rows <= nch * 512) try a.dupe(u8, try rig.bytes(o[3], rows * q_heads * hd * 2)) else null; + res.sparse[st] = try a.dupe(u8, try rig.bytes(o[4], rows * q_heads * hd * 2)); + } + if (tables) |t| { + const all = try a.dupe(u8, try rig.bytes(arena_ptr, f.pages * f.pageBytes())); + for (0..2) |st| res.caches[st] = try f.unpage(a, all, t, st); + } else for (0..2) |st| for (0..4) |pi| { + res.caches[st][pi] = try a.dupe(u8, try rig.bytes(flat[st][pi], f.caches[st][pi].len)); + }; + try check.expect(try rig.intact() and !rig.overlap(), "flashnext pages: a launch wrote past a buffer", .{}); + return res; +} + +/// Each attention row against float64 softmax(q k / 16) v over the caches the pass left (p . v takes p as bf16). +fn agrees(f: Fixture, r: Result, a: std.mem.Allocator) !bool { + const most = @max(f.c.lens[0], f.c.lens[1]) + f.c.rows + idw; + const sc = try a.alloc(f64, most); + for (0..2) |st| for ([_]bool{ false, true }) |sparse| { + const out_bytes = if (sparse) r.sparse[st] else r.dense[st] orelse continue; + const out: []align(1) const u16 = std.mem.bytesAsSlice(u16, out_bytes); + const q: []align(1) const u16 = std.mem.bytesAsSlice(u16, r.q[st]); + const kc: []align(1) const u16 = std.mem.bytesAsSlice(u16, r.caches[st][0]); + const vc: []align(1) const u16 = std.mem.bytesAsSlice(u16, r.caches[st][1]); + for (0..f.c.rows) |row| for (0..q_heads) |h| { + const kvh = h / (q_heads / kv_heads); + const n: usize = if (sparse) @intCast(f.nk[st][row]) else f.c.lens[st] + row + 1; + var hi = -std.math.inf(f64); + for (0..n) |j| { + const key: usize = if (sparse) @intCast(f.ids[st][row * idw + j]) else j; + var dot: f64 = 0; + for (0..hd) |d| dot += bf(q[(row * q_heads + h) * hd + d]) * bf(kc[(key * kv_heads + kvh) * hd + d]); + sc[j] = dot * 0.0625; + hi = @max(hi, sc[j]); + } + var l: f64 = 0; + for (sc[0..n]) |*v| { + v.* = @exp(v.* - hi); + l += v.*; + } + for (0..hd) |d| { + var o: f64 = 0; + var abs: f64 = 0; + for (0..n) |j| { + const key: usize = if (sparse) @intCast(f.ids[st][row * idw + j]) else j; + const v = bf(vc[(key * kv_heads + kvh) * hd + d]); + o += sc[j] * v; + abs += sc[j] * @abs(v); + } + const want = o / l; + const got = bf(out[(row * q_heads + h) * hd + d]); + if (!(@abs(got - want) <= abs / l * 0x1p-8 + 0.5 * @abs(want) * 0x1p-7 + 1e-6)) return false; + } + }; + }; + return true; +} + +fn same(x: []const u8, y: []const u8) bool { + return std.mem.eql(u8, x, y); +} + +fn sameResult(x: Result, y: Result) bool { + for (0..2) |st| { + if (!same(x.q[st], y.q[st]) or !same(x.iq[st], y.iq[st]) or !same(x.sc[st], y.sc[st]) or !same(x.sparse[st], y.sparse[st])) return false; + if ((x.dense[st] == null) != (y.dense[st] == null)) return false; + if (x.dense[st]) |d| if (!same(d, y.dense[st].?)) return false; + for (0..4) |pi| if (!same(x.caches[st][pi], y.caches[st][pi])) return false; + } + return true; +} + +pub fn run(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + var red_table = false; + var red_bound = false; + for (cases, 0..) |c, i| { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const f = try Fixture.make(a, c, 0xf9a6 + i); + const flat = try launch(gpu, s, &g, f, null, null, a, 1); + try check.expect(try agrees(f, flat, a), "flashnext pages: contiguous attention off float64 at lens {any}, rows {d}", .{ c.lens, c.rows }); + for (0..3) |rep| { + const pg = try launch(gpu, s, &g, f, f.ids_pages, null, a, 2 + rep); + try check.expect(sameResult(pg, flat), "flashnext pages: paged bytes differ from contiguous at lens {any}, rows {d}, page {d} (repeat {d})", .{ c.lens, c.rows, c.ptok, rep }); + } + // red controls: stream 1's first and last pages swapped in the table only (the pages stay); a planted row + const n1 = (c.lens[1] + c.rows + c.ptok - 1) / c.ptok; + if (n1 >= 2) { + const swapped = try a.dupe(u32, f.ids_pages[1]); + std.mem.swap(u32, &swapped[0], &swapped[n1 - 1]); + const bad = try launch(gpu, s, &g, f, f.ids_pages, .{ f.ids_pages[0], swapped }, a, 9); + const moved = !same(bad.sparse[1], flat.sparse[1]); + red_table = red_table or moved; + try check.expect(moved, "flashnext pages: a swapped page table still gave contiguous bytes (control missed) at lens {any}", .{c.lens}); + } + var bad_r = flat; + bad_r.sparse[0] = try a.dupe(u8, flat.sparse[0]); + const planted: usize = if (c.rows >= 2) q_heads * hd + 5 else 5; + std.mem.writeInt(u16, bad_r.sparse[0][2 * planted ..][0..2], 0x4040, .little); // 3.0 where |o| <= 1 + const caught = !(try agrees(f, bad_r, a)); + red_bound = red_bound or caught; + try check.expect(caught, "flashnext pages: a planted attention row passed float64 (control missed)", .{}); + check.pass("flashnext pages: lens {any}, {d} rows a stream, {d}-position pages: paged == contiguous (queries, indexer queries, block scores, dense and sparse rows, stores, pooled), three repeats, float64 bound, guards", .{ c.lens, c.rows, c.ptok }); + } + try check.expect(red_table and red_bound, "flashnext pages: a control never fired", .{}); + check.pass("flashnext pages: controls RED as planted (swapped table, planted row)", .{}); +} diff --git a/zig/tests/cuda/contract_qwen_pages.zig b/zig/tests/cuda/contract_qwen_pages.zig new file mode 100644 index 0000000000..ddae593d81 --- /dev/null +++ b/zig/tests/cuda/contract_qwen_pages.zig @@ -0,0 +1,354 @@ +//! qwen35.cu's and qwen35_kv8.cu's attention caches through a page table: rope, attention, keep, paged == contiguous. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const hd = 256; +const max_chunks = 64; +const attentions = 2; // the fixture's attention layers: the kernels run layer 1, so a page's layer offset counts +const q_heads = 24; // the 27B's heads +const kv_heads = 4; +const e4m3 = core.quant.fp8.e4m3; + +// qwen35.cuh's Geo, Seg, Ancestor and Path (src/families/qwen3_5/cuda_cache.zig mirrors them for the engine) +const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int, nl: c_int }; +const Seg = extern struct { rec: u64 = 0, conv: u64 = 0, kv: u64, tap: u64 = 0, log: u64 = 0, before: u64 = 0, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int = 0, taps: c_int = 0, delta: c_int = 0, pt: u64 = 0, ps: u64 = 0, ptok: c_int = 0 }; +const Ancestor = extern struct { bits: u64, depth: u32, chain: u32 }; +const Path = extern struct { n: c_int, rows: [64]c_int }; + +/// bf16 rows (qwen35.cu), or e4m3 rows with an fp32 scale each (qwen35_kv8.cu, --kv-cache fp8). +const Format = enum { bf16, fp8 }; + +/// Two streams' cached positions, the window rows each adds, positions a page. +const Case = struct { lens: [2]usize, rows: usize, ptok: usize }; +const cases = [_]Case{ + .{ .lens = .{ 0, 3 }, .rows = 1, .ptok = 16 }, + .{ .lens = .{ 127, 128 }, .rows = 4, .ptok = 16 }, + .{ .lens = .{ 129, 1000 }, .rows = 8, .ptok = 128 }, + .{ .lens = .{ 700, 33 }, .rows = 2, .ptok = 48 }, // a page no chunk width divides + .{ .lens = .{ 4095, 5000 }, .rows = 16, .ptok = 256 }, + .{ .lens = .{ 9000, 300 }, .rows = 3, .ptok = 512 }, +}; + +const Span = struct { off: usize, len: usize }; + +/// Bytes of a layer's (or a whole block's, `layers`) rows, `n` positions a head. +fn blockBytes(f: Format, n: usize, layers: usize) usize { + return layers * 2 * kv_heads * n * @as(usize, if (f == .fp8) hd + 4 else hd * 2); +} + +/// Layer a's key (s 0) or value (s 1) at pos, head h, in a block of n positions a head: its row, and fp8's scale. +fn spans(f: Format, n: usize, a: usize, s: usize, h: usize, pos: usize) [2]Span { + const blk = s * kv_heads + h; + const layer = blockBytes(f, n, a); + return switch (f) { + .bf16 => .{ .{ .off = layer + (blk * n + pos) * hd * 2, .len = hd * 2 }, .{ .off = 0, .len = 0 } }, + .fp8 => .{ .{ .off = layer + (blk * n + pos) * hd, .len = hd }, .{ .off = layer + 2 * kv_heads * n * hd + (blk * n + pos) * 4, .len = 4 } }, + }; +} + +/// One case's host data: the logical caches (cap rows a head), the window's projections, the page tables. +const Fixture = struct { + c: Case, + format: Format, + cap: usize, + kv: [2][]u8, + q: []u16, // [rows][q_heads][hd] + k: []u16, // [rows][kv_heads][hd] + v: []u16, + ids: [2][]u32, // each stream's pages in the shared arena, shuffled across both + pages: usize, // the arena's pages (two spare) + + fn streamBytes(f: Fixture) usize { + return blockBytes(f.format, f.cap, attentions); + } + + fn pageBytes(f: Fixture) usize { + return blockBytes(f.format, f.c.ptok, attentions); + } + + fn logical(f: Fixture, a: usize, s: usize, h: usize, pos: usize) [2]Span { + return spans(f.format, f.cap, a, s, h, pos); + } + + /// Where logical (a, s, h, pos) of stream `st` sits in the arena under `ids`. + fn paged(f: Fixture, ids: [2][]const u32, st: usize, a: usize, s: usize, h: usize, pos: usize) [2]Span { + var r = spans(f.format, f.c.ptok, a, s, h, pos % f.c.ptok); + for (&r) |*x| x.off += ids[st][pos / f.c.ptok] * f.pageBytes(); + return r; + } + + fn make(a: std.mem.Allocator, c: Case, format: Format, seed: u64) !Fixture { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const most = @max(c.lens[0], c.lens[1]) + c.rows; + const cap = (most + c.ptok - 1) / c.ptok * c.ptok; + var f: Fixture = .{ .c = c, .format = format, .cap = cap, .kv = undefined, .q = undefined, .k = undefined, .v = undefined, .ids = undefined, .pages = 2 * (cap / c.ptok) + 2 }; + for (&f.kv) |*x| { + x.* = try a.alloc(u8, f.streamBytes()); + for (0..attentions) |at| for (0..2) |s| for (0..kv_heads) |h| for (0..cap) |pos| fill(f, rng, x.*, f.logical(at, s, h, pos)); + } + f.q = try randoms(a, rng, c.rows * 2 * q_heads * hd); + f.k = try randoms(a, rng, c.rows * 2 * kv_heads * hd); + f.v = try randoms(a, rng, c.rows * 2 * kv_heads * hd); + const order = try a.alloc(u32, f.pages); + for (order, 0..) |*o, i| o.* = @intCast(i); + rng.shuffle(u32, order); + const n = cap / c.ptok; + f.ids = .{ order[0..n], order[n .. 2 * n] }; + return f; + } + + /// A row's bytes: bf16 in [-1, 1); e4m3 codes but NaN's, a scale in [1/1792, 1/448) (values within 1). + fn fill(f: Fixture, rng: std.Random, out: []u8, sp: [2]Span) void { + switch (f.format) { + .bf16 => for (std.mem.bytesAsSlice(u16, out[sp[0].off..][0..sp[0].len])) |*x| { + x.* = @truncate(@as(u32, @bitCast(rng.float(f32) * 2 - 1)) >> 16); + }, + .fp8 => { + for (out[sp[0].off..][0..sp[0].len]) |*b| { + b.* = rng.int(u8); + if (b.* & 0x7f == 0x7f) b.* ^= 1; + } + const sc: f32 = (1 + 3 * rng.float(f32)) / 1792; + @memcpy(out[sp[1].off..][0..4], std.mem.asBytes(&sc)); + }, + } + } + + /// The arena holding both streams' logical caches at their pages under `ids`; spare pages random. + fn arena(f: Fixture, a: std.mem.Allocator, rng: std.Random, ids: [2][]const u32) ![]u8 { + const out = try a.alloc(u8, f.pages * f.pageBytes()); + rng.bytes(out); + for (0..2) |st| for (0..attentions) |at| for (0..2) |s| for (0..kv_heads) |h| for (0..f.cap) |pos| { + const from = f.logical(at, s, h, pos); + const to = f.paged(ids, st, at, s, h, pos); + for (from, to) |x, y| @memcpy(out[y.off..][0..y.len], f.kv[st][x.off..][0..x.len]); + }; + return out; + } + + /// The arena's pages read back as stream `st`'s logical cache. + fn unpage(f: Fixture, a: std.mem.Allocator, arena_: []const u8, ids: [2][]const u32, st: usize) ![]u8 { + const out = try a.alloc(u8, f.streamBytes()); + for (0..attentions) |at| for (0..2) |s| for (0..kv_heads) |h| for (0..f.cap) |pos| { + const to = f.logical(at, s, h, pos); + const from = f.paged(ids, st, at, s, h, pos); + for (from, to) |x, y| @memcpy(out[y.off..][0..y.len], arena_[x.off..][0..x.len]); + }; + return out; + } + + /// Logical value d of (a, s, h, pos) in stream cache `cache`. + fn value(f: Fixture, cache: []const u8, a: usize, s: usize, h: usize, pos: usize, d: usize) f64 { + const sp = f.logical(a, s, h, pos); + return switch (f.format) { + .bf16 => bf(std.mem.readInt(u16, cache[sp[0].off + 2 * d ..][0..2], .little)), + .fp8 => e4m3(cache[sp[0].off + d]) * @as(f64, @as(f32, @bitCast(std.mem.readInt(u32, cache[sp[1].off..][0..4], .little)))), + }; + } +}; + +/// bf16 in [-1, 1): finite, so a sum's bytes are its order's alone. +fn randoms(a: std.mem.Allocator, rng: std.Random, n: usize) ![]u16 { + const out = try a.alloc(u16, n); + for (out) |*x| x.* = @truncate(@as(u32, @bitCast(rng.float(f32) * 2 - 1)) >> 16); + return out; +} + +fn bf(x: u16) f64 { + return @as(f32, @bitCast(@as(u32, x) << 16)); +} + +/// One layout's results: the rotated queries, the attention rows, each stream's cache after rope and after keep. +const Result = struct { q: []u8, out: []u8, after_rope: [2][]u8, after_keep: [2][]u8 }; + +const Fns = struct { rope: cuda.Function, partial: cuda.Function, merge: cuda.Function, compact: cuda.Function }; + +/// rope, attention, keep on layer 1 of both streams: contiguous, or pages placed by `tables`, read through `given`. +fn launch(gpu: check.Gpu, s: cuda.Stream, fns: Fns, f: Fixture, tables: ?[2][]const u32, given: ?[2][]const u32, a: std.mem.Allocator, seed: u64) !Result { + var rig = rig_.Rig{ .gpu = gpu, .s = s }; + defer rig.deinit(); + var prng = std.Random.DefaultPrng.init(seed); + const m = 2 * f.c.rows; + const geo: Geo = .{ .hid = 0, .nk = 0, .nv = 0, .qh = q_heads, .kvh = kv_heads, .nl = 0 }; + var segs: [2]Seg = undefined; + var kv: [2]u64 = undefined; // a stream's buffer, or the arena twice + if (tables) |t| { + const host = try f.arena(a, prng.random(), t); + const ptr = try rig.scratch(host.len); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = ptr, .len = host.len }).upload(0, host); + kv = .{ ptr, ptr }; + } else for (0..2) |st| { + kv[st] = try rig.scratch(f.streamBytes()); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = kv[st], .len = f.streamBytes() }).upload(0, f.kv[st]); + } + for (0..2) |st| { + segs[st] = .{ .kv = kv[st], .cap = @intCast(f.cap), .len = @intCast(f.c.lens[st]), .row0 = @intCast(st * f.c.rows), .rows = @intCast(f.c.rows) }; + if (given orelse tables) |t| { + segs[st].pt = try rig.input(std.mem.sliceAsBytes(t[st])); + segs[st].ps = if (f.format == .fp8) f.pageBytes() else f.pageBytes() / 2; + segs[st].ptok = @intCast(f.c.ptok); + } + } + const rowseg = try a.alloc(i32, m); + const rowanc = try a.alloc(Ancestor, m); + for (0..m) |r| { + rowseg[r] = @intCast(r / f.c.rows); + rowanc[r] = .{ .bits = 0, .depth = @intCast(r % f.c.rows), .chain = 1 }; + } + const segs_d = try rig.input(std.mem.sliceAsBytes(&segs)); + const rowseg_d = try rig.input(std.mem.sliceAsBytes(rowseg)); + const rowanc_d = try rig.input(std.mem.sliceAsBytes(rowanc)); + const qn = try rig.scratch(f.q.len * 2); // rotated in place + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = qn, .len = f.q.len * 2 }).upload(0, std.mem.sliceAsBytes(f.q)); + const kn = try rig.input(std.mem.sliceAsBytes(f.k)); + const vr = try rig.input(std.mem.sliceAsBytes(f.v)); + const slots = m * q_heads * max_chunks; + const pm = try rig.scratch(slots * 4); + const pl = try rig.scratch(slots * 4); + const po = try rig.scratch(slots * hd * 4); + const out = try rig.output(m * q_heads * hd * 2); + const before = try a.dupe(u8, try rig.inputBytes()); + const layer: c_int = 1; + const none: u64 = 0; + var args: cuda.Args = .{}; + inline for (.{ geo, qn, kn, vr, segs_d, rowseg_d, rowanc_d, layer, none, none }) |x| args.add(x); + try cuda.launch.launch(fns.rope, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = hd } }, s, &args); + args = .{}; + inline for (.{ geo, qn, segs_d, rowseg_d, rowanc_d, layer, none, @as(f32, 1.0 / 16.0), pm, pl, po }) |x| args.add(x); + try cuda.launch.launch(fns.partial, .{ .grid = .{ .x = max_chunks, .y = kv_heads, .z = @intCast(m) }, .block = .{ .x = 32 * (q_heads / kv_heads) } }, s, &args); + args = .{}; + inline for (.{ geo, pm, pl, po, segs_d, rowseg_d, rowanc_d, none, out }) |x| args.add(x); + try cuda.launch.launch(fns.merge, .{ .grid = .{ .x = q_heads, .y = @intCast(m) }, .block = .{ .x = 32 } }, s, &args); + var r: Result = undefined; + r.q = try a.dupe(u8, try rig.bytes(qn, f.q.len * 2)); + r.out = try a.dupe(u8, try rig.bytes(out, m * q_heads * hd * 2)); + try readCaches(&rig, a, f, tables, kv, &r.after_rope); + // keep: stream 0's window rows 0, 2 and the last (a tree's path), moved down after its cache + if (f.c.rows >= 3) { + var path: Path = .{ .n = 3, .rows = @splat(0) }; + path.rows[1] = 2; + path.rows[2] = @intCast(f.c.rows - 1); + args = .{}; + inline for (.{ geo, segs[0], path }) |x| args.add(x); + try cuda.launch.launch(fns.compact, .{ .grid = .{ .x = attentions, .y = kv_heads }, .block = .{ .x = hd } }, s, &args); + } + try readCaches(&rig, a, f, tables, kv, &r.after_keep); + try check.expect(try rig.intact() and !rig.overlap() and std.mem.eql(u8, before, try rig.inputBytes()), "qwen35 pages: a launch changed an input or guard", .{}); + return r; +} + +fn readCaches(rig: *rig_.Rig, a: std.mem.Allocator, f: Fixture, tables: ?[2][]const u32, kv: [2]u64, out: *[2][]u8) !void { + if (tables) |t| { + const all = try a.dupe(u8, try rig.bytes(kv[0], f.pages * f.pageBytes())); + for (0..2) |st| out[st] = try f.unpage(a, all, t, st); + } else for (0..2) |st| out[st] = try a.dupe(u8, try rig.bytes(kv[st], f.streamBytes())); +} + +/// Each attention row against float64 softmax(q k / 16) v: a bound (fp32 and fast exp), bf16 rounding of its size. +fn agrees(a: std.mem.Allocator, f: Fixture, r: Result) !bool { + const q: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, r.q)); + const out: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, r.out)); + const keys_most = @max(f.c.lens[0], f.c.lens[1]) + f.c.rows; + const p = try a.alloc(f64, keys_most); + defer a.free(p); + const vals = try a.alloc(f64, keys_most * hd); + defer a.free(vals); + for (0..2 * f.c.rows) |row| { + const st = row / f.c.rows; + const keys = f.c.lens[st] + row % f.c.rows + 1; // the chain's rows are positions after the cache + const cache = r.after_rope[st]; + for (0..q_heads) |qh| { + const h = qh / (q_heads / kv_heads); + const qv = q[(row * q_heads + qh) * hd ..][0..hd]; + var top: f64 = -std.math.inf(f64); + for (0..keys) |t| { + var dot: f64 = 0; + for (0..hd) |d| dot += bf(qv[d]) * f.value(cache, 1, 0, h, t, d); + p[t] = dot / 16.0; + top = @max(top, p[t]); + for (0..hd) |d| vals[t * hd + d] = f.value(cache, 1, 1, h, t, d); + } + var sum: f64 = 0; + for (p[0..keys]) |*x| { + x.* = @exp(x.* - top); + sum += x.*; + } + for (0..hd) |d| { + var o: f64 = 0; + var mag: f64 = 0; + for (0..keys) |t| { + o += p[t] * vals[t * hd + d]; + mag += p[t] * @abs(vals[t * hd + d]); + } + o /= sum; + mag /= sum; + const got = bf(out[(row * q_heads + qh) * hd + d]); + if (!(@abs(got - o) <= std.math.ldexp(@as(f64, 1), -8) * (@abs(o) + mag))) return false; + } + } + } + return true; +} + +fn eqCaches(x: [2][]u8, y: [2][]u8) bool { + return std.mem.eql(u8, x[0], y[0]) and std.mem.eql(u8, x[1], y[1]); +} + +pub fn run(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.qwen35); + defer mod.unload(); + var keep = try cuda.Module.load(gpu.d, cuda.kernels.qwen35_keep); + defer keep.unload(); + var kv8 = try cuda.Module.load(gpu.d, cuda.kernels.qwen35_kv8); + defer kv8.unload(); + const merge = try mod.function("tf_q35_attn_merge"); + for (std.enums.values(Format)) |format| { + const fns: Fns = switch (format) { + .bf16 => .{ .rope = try mod.function("tf_q35_rope"), .partial = try mod.function("tf_q35_attn_partial"), .merge = merge, .compact = try keep.function("tf_q35k_compact") }, + .fp8 => .{ .rope = try kv8.function("tf_q35_rope8"), .partial = try kv8.function("tf_q35_attn_partial8"), .merge = merge, .compact = try kv8.function("tf_q35k_compact8") }, + }; + var red_table = false; + var red_bound = false; + for (cases, 0..) |c, i| { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const f = try Fixture.make(a, c, format, 0x9a6e + i); + const flat = try launch(gpu, s, fns, f, null, null, a, 1); + try check.expect(try agrees(a, f, flat), "qwen35 pages {t}: contiguous attention off float64 at lens {any}, rows {d}", .{ format, c.lens, c.rows }); + for (0..3) |rep| { + const pg = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, null, a, 2 + rep); + try check.sameBytes("qwen35 pages: rotated queries", pg.q, flat.q); + try check.sameBytes("qwen35 pages: attention rows, paged against contiguous", pg.out, flat.out); + try check.expect(eqCaches(pg.after_rope, flat.after_rope), "qwen35 pages {t}: rope's stores differ at lens {any}, page {d}", .{ format, c.lens, c.ptok }); + try check.expect(eqCaches(pg.after_keep, flat.after_keep), "qwen35 pages {t}: keep's moves differ at lens {any}, page {d}", .{ format, c.lens, c.ptok }); + } + // red controls: two of stream 1's pages swapped in the table only (the pages stay); a planted row + const n1 = (c.lens[1] + c.rows + c.ptok - 1) / c.ptok; + if (n1 >= 2) { + const swapped = try a.dupe(u32, f.ids[1]); + std.mem.swap(u32, &swapped[0], &swapped[n1 - 1]); + const bad = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, .{ f.ids[0], swapped }, a, 9); + const moved = !std.mem.eql(u8, bad.out, flat.out); + red_table = red_table or moved; + try check.expect(moved, "qwen35 pages {t}: a swapped page table still gave contiguous bytes (control missed) at lens {any}", .{ format, c.lens }); + } + const o = try a.dupe(u8, flat.out); + std.mem.writeInt(u16, o[2 * (q_heads * hd + 5) ..][0..2], 0x4040, .little); // 3.0 where |o| <= 1 + var bad_r = flat; + bad_r.out = o; + const caught = !(try agrees(a, f, bad_r)); + red_bound = red_bound or caught; + try check.expect(caught, "qwen35 pages {t}: a planted attention row passed float64 (control missed)", .{format}); + std.debug.print("PASS qwen35 pages {t}: lens {any}, {d} rows a stream, {d}-position pages: paged == contiguous (queries, rows, stores, keep), three repeats, float64 bound, guards\n", .{ format, c.lens, c.rows, c.ptok }); + } + try check.expect(red_table and red_bound, "qwen35 pages {t}: a control never fired", .{format}); + std.debug.print("PASS qwen35 pages {t}: controls RED as planted (swapped table, planted row)\n", .{format}); + } +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 86aa6eb966..4b772b43ed 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -281,7 +281,9 @@ pub fn run(gpu: check.Gpu) !void { std.debug.print("PASS gradient component suites (FP64 and finite-difference checks); resource receipts still required\n", .{}); const t1 = check.now(gpu.io); try @import("contract_qwen_served.zig").run(gpu); + try @import("contract_qwen_pages.zig").run(gpu); try @import("contract_qwen_learning.zig").numerical(gpu); + try @import("contract_fn_pages.zig").run(gpu); const lim = try live(gpu); const t2 = check.now(gpu.io); std.debug.print("SUMMARY {d} kernels: {d} pass gates 1-4, {d} fail, {d} known findings; {d} controls: {d} red as planted, {d} missed; {d} limit problems; gates {d} ms, limits {d} ms\n", .{ t.kernels, t.passed, t.failed, t.known, t.controls, t.caught, t.missed, lim, @divTrunc(t1 - t0, std.time.ns_per_ms), @divTrunc(t2 - t1, std.time.ns_per_ms) }); diff --git a/zig/tests/cuda/flashnext_attn_tests.zig b/zig/tests/cuda/flashnext_attn_tests.zig index 806a7207bf..744792d380 100644 --- a/zig/tests/cuda/flashnext_attn_tests.zig +++ b/zig/tests/cuda/flashnext_attn_tests.zig @@ -77,7 +77,7 @@ fn prep(r: *Rig, g: *const flash.Glue) !void { kcs[c] = try r.out(cap * kv_heads * head_dim * 2); vcs[c] = try r.out(cap * kv_heads * head_dim * 2); ikcs[c] = try r.out(cap * index_dim * 2); - try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, qs[c], kcs[c], vcs[c], iqs[c], ikcs[c], eps, m); + try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, qs[c], kcs[c], vcs[c], iqs[c], ikcs[c], eps, m, .{}); } const last = counts.len - 1; const q = try r.back(u16, qs[last], most * q_heads * head_dim); @@ -142,7 +142,7 @@ fn indexer(r: *Rig, g: *const flash.Glue) !void { const dikc = try r.dev(std.mem.sliceAsBytes(ikc)); const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); const pooled = try r.out((cap / 4) * index_dim * 2); - try g.qsaPool(r.s, dikc, pooled, dpos, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(inv)), eps, rows); + try g.qsaPool(r.s, dikc, pooled, dpos, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(inv)), eps, rows, .{}); const got = try r.back(u16, pooled, (cap / 4) * index_dim); var b: usize = pos0 / 4; while (4 * b + 4 <= pos0 + rows) : (b += 1) { @@ -171,7 +171,7 @@ fn indexer(r: *Rig, g: *const flash.Glue) !void { var scs: [2]u64 = undefined; for ([_]usize{ 1, rows }, 0..) |m, c| { scs[c] = try r.out(m * nb * 4); - try g.qsaScores(r.s, diq, dkeys, dpos, scs[c], nb, m, (pos0 + rows) / 4); + try g.qsaScores(r.s, diq, dkeys, dpos, scs[c], nb, m, (pos0 + rows) / 4, .{}); } const sc = try r.back(f32, scs[1], rows * nb); for (0..rows) |row| { @@ -227,7 +227,7 @@ fn attention(r: *Rig, g: *const flash.Glue, sparse: bool) !void { for (counts, 0..) |m, c| { outs[c] = try r.out(m * q_heads * head_dim * 2); const chunks = if (sparse) (idw + 511) / 512 else (pos0 + m + 511) / 512; - try g.attnChunks(r.s, dq, dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, m, kv_heads, chunks, true); + try g.attnChunks(r.s, dq, dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, m, kv_heads, chunks, true, .{}); try g.attnMerge(r.s, po, pm, pl, dpos, outs[c], dnk, dspr, m, kv_heads, true); } const got = try r.back(u16, outs[counts.len - 1], rows * q_heads * head_dim); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 837f1bb070..d9f648f416 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -106,6 +106,7 @@ const usage = \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) + \\ flashnext-pages the attention glue's caches through a page table: paged == contiguous, float64, controls \\ info driver, device and library versions \\ smoke copies, fills, launches, argument packing, module globals \\ shared-context Context.shared and Stream.borrowed as a library inside a host's process sees them @@ -188,6 +189,7 @@ const usage = \\ exl3-tiles [DIR] EXL3's tile candidate: codebooks and tiles bit-equal, FP64, rows through qlinear and Dense \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) + \\ qwen-pages qwen35.cu's attention cache through a page table: paged == contiguous, float64, red controls \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) \\ glm-vision MODEL F32 BF16 IMG GLM's image tower and preparation against vision_ref.py's torch runs @@ -252,6 +254,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-pages")) return @import("contract_fn_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "info")) return info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); if (std.mem.eql(u8, cmd, "shared-context")) return runtime_tests.sharedContext(gpu); @@ -357,6 +360,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "q27-gemm")) return q27_gemm_tests.run(gpu); if (std.mem.eql(u8, cmd, "gdn-back-27b")) return gdn_back_tests.runShape(gpu, 5120, 64); if (std.mem.eql(u8, cmd, "kernel-contract")) return contract_tests.run(gpu); + if (std.mem.eql(u8, cmd, "qwen-pages")) return @import("contract_qwen_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "tree")) return tree_tests.run(gpu); if (std.mem.eql(u8, cmd, "block")) return block_tests.run(gpu);