Skip to content

Flash Next CUDA: the streams beside the first take KV pages from the core pool #1170

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jschmied wants to merge 11 commits into from
Closed

Flash Next CUDA: the streams beside the first take KV pages from the core pool #1170

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
3f5b87d
cuda qwen3_5: attention caches readable through a page table
jschmied Oct 21, 2026
d5189b3
cuda qwen3_5: the streams beside the first take KV pages from the cor…
jschmied Oct 21, 2026
6037621
cuda qwen3_5: with a KV pool, warm the shared shapes on lanes opened …
jschmied Oct 21, 2026
5c09685
cuda qwen3_5: --kv-cache fp8 through the page table too
jschmied Oct 21, 2026
9a99d1f
cuda qwen3_5: one-line comments on row8 and the pool's reservation
jschmied Oct 21, 2026
1ac2075
cuda qwen3_5: a pool lane's recurrent state is a run of pool pages wh…
jschmied Oct 21, 2026
51f8354
cuda qwen3_5: a DFlash2 stream's ring is part of its pool lane's run;…
jschmied Oct 21, 2026
ce5b78d
cuda qwen3_5: the attention partials test the page table once a block…
jschmied Oct 21, 2026
85aeb4b
cuda flashnext: attention caches readable through a page table
jschmied Oct 21, 2026
336cb3f
cuda flashnext: the streams beside the first take KV pages from the c…
jschmied Oct 21, 2026
e0e69c5
cuda qwen3_5: the reuse tests' caches say they have no page table
jschmied Oct 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 44 additions & 27 deletions zig/kernels/cuda/flashnext_attn.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<long long>(i) * row;
return static_cast<long long>(pg.pt[i / per]) * pg.ps + static_cast<long long>(i % per) * row;
}

} // namespace fn_attn

using namespace fn_attn;
Expand All @@ -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<long long>(r) * PW;
if (head >= NQ + NKV + NI) {
if (d < IHD) IKC[static_cast<long long>(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;
Expand All @@ -94,9 +107,9 @@ __device__ __forceinline__ void fn_attn_prep_body(const __nv_bfloat16* __restric
if (head < NQ) {
Q[(static_cast<long long>(r) * NQ + head) * HD + d] = out;
} else if (!index) {
const long long slot = static_cast<long long>(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<long long>(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<long long>(r) * NI + head - NQ - NKV) * IHD + d] = out;
}
Expand All @@ -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,
Expand All @@ -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.
Expand All @@ -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<long long>(RATIO) * b * DI + d]);
for (int k = 1; k < RATIO; ++k) acc = __fadd_rn(acc, bf(IKC[(static_cast<long long>(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<long long>(k) * DI]));
const float x = rbf(__fdiv_rn(acc, static_cast<float>(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<long long>(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;
Expand All @@ -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<const uint2*>(POOLED + static_cast<long long>(b) * DI + 4 * lane) : make_uint2(0, 0);
raw[i] = b < complete ? *reinterpret_cast<const uint2*>(POOLED + pg_at(pg, b, pg.ptok / RATIO, DI) + 4 * lane) : make_uint2(0, 0);
}
#pragma unroll
for (int i = 0; i < 8; ++i) {
Expand All @@ -210,14 +224,15 @@ __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<Row*>(smem_raw); // [QR] the row's query heads
Row* tiles_kv = qs + QR; // [NB][2][TK] each buffer's keys, then values
__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;
Expand All @@ -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<long long>(r) * IDW + ki] : ki) : -1;
const int key = ki < n ? (sparse ? IDS[static_cast<long long>(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<long long>(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<long long>(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();
}
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand Down
52 changes: 30 additions & 22 deletions zig/kernels/cuda/qwen35.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
};

Expand Down Expand Up @@ -434,31 +433,40 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r
}
return;
}
const __nv_bfloat16* K = reinterpret_cast<const __nv_bfloat16*>(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<const __nv_bfloat16*>(sg.kv);
float q[DPL], o[DPL];
#pragma unroll
for (int i = 0; i < DPL; i++) {
q[i] = bf(Q[(m * QH + qh) * HD + lane * DPL + i]);
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<const uint4*>(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL);
const __nv_bfloat16* kb = reinterpret_cast<const __nv_bfloat16*>(&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<const uint4*>(V + ((size_t)h * sg.cap + pos) * HD + lane * DPL);
const __nv_bfloat16* vb = reinterpret_cast<const __nv_bfloat16*>(&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<const uint4*>(row(0, pos) + lane * DPL);
const __nv_bfloat16* kb = reinterpret_cast<const __nv_bfloat16*>(&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<const uint4*>(row(1, pos) + lane * DPL);
const __nv_bfloat16* vb = reinterpret_cast<const __nv_bfloat16*>(&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;
Expand Down
11 changes: 11 additions & 0 deletions zig/kernels/cuda/qwen35.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<const int*>(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
Expand Down
Loading
Loading