From 518c8b5ec92c9df1ad794ae71ba59e7fcc2cc54d Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:32:58 +0200 Subject: [PATCH 01/16] CUDA router logits: 64 input rows a block from 64 rows (same sums) --- zig/kernels/cuda/route_rows.cu | 54 +++++++++++++++++++++------------- zig/src/cuda/route_rows.zig | 11 ++++++- 2 files changed, 44 insertions(+), 21 deletions(-) diff --git a/zig/kernels/cuda/route_rows.cu b/zig/kernels/cuda/route_rows.cu index 9407493f75..881a1ffde7 100644 --- a/zig/kernels/cuda/route_rows.cu +++ b/zig/kernels/cuda/route_rows.cu @@ -45,19 +45,20 @@ __device__ __forceinline__ void block_best(float& v, int& i, float (*sv)[NT / 32 using namespace tf_route; -// OUT[m, e] = x[m] . w[e], a warp of 8 experts: one mma chain over K from zero, STAGES of KC staged (D % 16 == 0). -template +// OUT[m, e] = x[m] . w[e], a warp of 8 experts and RT tiles of 16 rows: one mma chain over K from zero, STAGES of KC +// staged (D % 16 == 0). A row's sums never depend on RT. +template __device__ __forceinline__ void route_logits(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, int M, int D, int E) { - constexpr int LD = KC + 8, EB = 8 * WARPS, ROWS = 16 + EB; + constexpr int LD = KC + 8, EB = 8 * WARPS, XR = 16 * RT, ROWS = XR + EB; extern __shared__ __align__(16) unsigned char smem_raw[]; - auto tile = reinterpret_cast<__nv_bfloat16 (*)[ROWS][LD]>(smem_raw); // [STAGES] 16 input rows, then EB expert rows - const int m0 = blockIdx.x * 16, e0 = blockIdx.y * EB; + auto tile = reinterpret_cast<__nv_bfloat16 (*)[ROWS][LD]>(smem_raw); // [STAGES] XR input rows, then EB expert rows + const int m0 = blockIdx.x * XR, e0 = blockIdx.y * EB; const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; const int stages = (D + KC - 1) / KC; - const int xr = min(16, M - m0); // input rows past M stay zero: written once, never staged - for (int v = threadIdx.x; v < STAGES * (16 - xr) * (KC / 8); v += blockDim.x) { - const int per = (16 - xr) * (KC / 8); + const int xr = min(XR, M - m0); // input rows past M stay zero: written once, never staged + for (int v = threadIdx.x; v < STAGES * (XR - xr) * (KC / 8); v += blockDim.x) { + const int per = (XR - xr) * (KC / 8); *reinterpret_cast(&tile[v / per][xr + (v % per) / (KC / 8)][8 * (v % (KC / 8))]) = make_uint4(0, 0, 0, 0); } auto stage = [&](int st) { @@ -65,7 +66,7 @@ __device__ __forceinline__ void route_logits(const __nv_bfloat16* __restrict__ X for (int v = threadIdx.x; v < (xr + EB) * (KC / 8); v += blockDim.x) { const int i = v / (KC / 8), k = k0 + 8 * (v % (KC / 8)); const bool x_row = i < xr; - const int row = x_row ? i : 16 + i - xr; + const int row = x_row ? i : XR + i - xr; const int r = x_row ? m0 + i : e0 + i - xr; const bool read = k < D && (x_row || r < E); const __nv_bfloat16* src = x_row ? X + static_cast(r) * x_ld + (read ? k : 0) @@ -73,7 +74,7 @@ __device__ __forceinline__ void route_logits(const __nv_bfloat16* __restrict__ X qmm_frag::cp16z(&tile[buf][row][8 * (v % (KC / 8))], src, read); } }; - float acc[4] = {}; + float acc[RT][4] = {}; for (int st = 0; st < STAGES - 1; ++st) { if (st < stages) stage(st); qmm_frag::commit(); @@ -86,21 +87,27 @@ __device__ __forceinline__ void route_logits(const __nv_bfloat16* __restrict__ X const int buf = st % STAGES; const int steps = min(KC, D - st * KC) / 16; for (int kk = 0; kk < steps; ++kk) { - uint32_t a[4], b[2]; - qmm_frag::ldmatrix4(a, &tile[buf][lane % 16][kk * 16 + (lane / 16) * 8]); - qmm_frag::ldmatrix2(b, &tile[buf][16 + warp * 8 + lane % 8][kk * 16 + ((lane / 8) % 2) * 8]); - qmm_frag::mma(acc, a, b[0], b[1]); + uint32_t b[2]; + qmm_frag::ldmatrix2(b, &tile[buf][XR + warp * 8 + lane % 8][kk * 16 + ((lane / 8) % 2) * 8]); +#pragma unroll + for (int i = 0; i < RT; ++i) { + uint32_t a[4]; + qmm_frag::ldmatrix4(a, &tile[buf][i * 16 + lane % 16][kk * 16 + (lane / 16) * 8]); + qmm_frag::mma(acc[i], a, b[0], b[1]); + } } __syncthreads(); } const int e = e0 + warp * 8 + 2 * t; #pragma unroll - for (int h = 0; h < 2; ++h) { - const int r = m0 + g + 8 * h; - if (r >= M) continue; - if (e < E) OUT[static_cast(r) * E + e] = acc[2 * h]; - if (e + 1 < E) OUT[static_cast(r) * E + e + 1] = acc[2 * h + 1]; - } + for (int i = 0; i < RT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = m0 + i * 16 + g + 8 * h; + if (r >= M) continue; + if (e < E) OUT[static_cast(r) * E + e] = acc[i][2 * h]; + if (e + 1 < E) OUT[static_cast(r) * E + e + 1] = acc[i][2 * h + 1]; + } } // route_logits over 32 experts a block (prompt chunks: the input rows are read once per 32 experts). @@ -110,6 +117,13 @@ extern "C" __global__ void __launch_bounds__(128) tf_route_logits(const __nv_bfl route_logits<4, 2, 128>(X, x_ld, W, OUT, M, D, E); } +// route_logits over 64 rows a block (from 64 rows: the router rows are read once per 64 input rows), the same sums. +extern "C" __global__ void __launch_bounds__(128) tf_route_logits_wide(const __nv_bfloat16* __restrict__ X, int x_ld, + const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, + int M, int D, int E) { + route_logits<4, 2, 128, 4>(X, x_ld, W, OUT, M, D, E); +} + // route_logits with three stages of 256 in flight (up to 16 rows, where the router rows' reads set its time). extern "C" __global__ void __launch_bounds__(128) tf_route_logits_few(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, diff --git a/zig/src/cuda/route_rows.zig b/zig/src/cuda/route_rows.zig index 665463a1df..34ebc4db09 100644 --- a/zig/src/cuda/route_rows.zig +++ b/zig/src/cuda/route_rows.zig @@ -17,19 +17,23 @@ pub const most_picks = 128; /// tf_route_logits' staged rows (16 inputs and 32 experts, padded by 8): two stages of 128, or three of 256. const many_shared: u32 = 2 * 48 * (128 + 8) * 2; const few_shared: u32 = 3 * 48 * (256 + 8) * 2; +/// tf_route_logits_wide's: 64 inputs and 32 experts, two stages of 128. +const wide_shared: u32 = 2 * 96 * (128 + 8) * 2; pub const RouteRows = struct { mod: Module, logits_fn: Function, few_fn: Function, + wide_fn: Function, topk_fn: Function, pub fn load(d: *const Driver) !RouteRows { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.route_rows); errdefer mod.unload(); - const r: RouteRows = .{ .mod = mod, .logits_fn = try mod.function("tf_route_logits"), .few_fn = try mod.function("tf_route_logits_few"), .topk_fn = try mod.function("tf_route_topk") }; + const r: RouteRows = .{ .mod = mod, .logits_fn = try mod.function("tf_route_logits"), .few_fn = try mod.function("tf_route_logits_few"), .wide_fn = try mod.function("tf_route_logits_wide"), .topk_fn = try mod.function("tf_route_topk") }; try r.few_fn.allowDynamicShared(few_shared); + try r.wide_fn.allowDynamicShared(wide_shared); return r; } @@ -49,6 +53,11 @@ pub const RouteRows = struct { for ([_]usize{ rows, d, e }) |v| a.add(@as(i32, @intCast(v))); // up to 16 rows, deeper staging (the same sums: the router rows' reads set the time) const few = rows <= 16; + if (rows >= 64) { + // 64 rows a block: each router row read once per 64 input rows (the same sums) + try launch_.launch(g.wide_fn, .{ .grid = .{ .x = @intCast((rows + 63) / 64), .y = @intCast((e + 31) / 32), .z = 1 }, .block = .{ .x = 128 }, .shared = wide_shared }, s, &a); + return; + } try launch_.launch(if (few) g.few_fn else g.logits_fn, .{ .grid = .{ .x = @intCast((rows + 15) / 16), .y = @intCast((e + 31) / 32), .z = 1 }, .block = .{ .x = 128 }, .shared = if (few) few_shared else many_shared }, s, &a); } From 2e36ac21bcbe683cb1a40f938b6586af5feb6158 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:32:58 +0200 Subject: [PATCH 02/16] Flash Next CUDA: prompt attention tiles share one buffer for keys then values, 1024 rows a launch with every chunk merged in the program (same bits) --- zig/kernels/cuda/flashnext_attn.cu | 101 ++++++++++++++++-- zig/src/families/flashnext/cuda_batch.zig | 61 ++++++----- .../families/flashnext/cuda_batch_state.zig | 10 +- zig/src/families/flashnext/cuda_glue.zig | 24 ++++- zig/src/families/flashnext/cuda_launch.zig | 6 ++ 5 files changed, 159 insertions(+), 43 deletions(-) diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu index 2f66798379..6386dc7767 100644 --- a/zig/kernels/cuda/flashnext_attn.cu +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -205,7 +205,7 @@ extern "C" __global__ void __launch_bounds__(256) fn_qsa_scores(const __nv_bfloa } // attention._chunks, program (row, KV head, chunk): the chunk's keys TK at a time, online softmax, partial o, m, l. -template +template __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, @@ -215,7 +215,8 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, typedef __nv_bfloat16 Row[HD + 8]; Row* qs = reinterpret_cast(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 + // ONE: one [TK] buffer, the tile's keys and then, once the scores are taken, its values (fewer bytes, more blocks) + __nv_bfloat16 (*ps)[TK + 8] = reinterpret_cast<__nv_bfloat16 (*)[TK + 8]>(tiles_kv + (ONE ? TK : NB * 2 * TK)); __shared__ float red_m[4][QR], red_l[4][QR]; __shared__ int key_ids[NB][TK]; const int r = blockIdx.x, hk = blockIdx.y, c = blockIdx.z; @@ -226,8 +227,10 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, sparse = SPR[r] != 0; if (sparse) n = NKR[r]; } - const int start = c * CH; - if (start >= n) return; + // FUSE: one program a row takes every chunk in order and merges them as fn_attn_merge does (no partials) + const int c_end = FUSE ? (n + CH - 1) / CH : c + 1; + if (c * CH >= n) return; + int start = c * CH; for (int e = threadIdx.x; e < QR * HD / 8; e += blockDim.x) { const int i = e / (HD / 8), v = e % (HD / 8); qmm_frag::cp16z(&qs[i][8 * v], Q + (static_cast(r) * H + hk * G + min(i, G - 1)) * HD + 8 * v, i < G); @@ -249,6 +252,18 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, qmm_frag::commit(); } }; + // ONE: the keys (pass 0) or the values (pass 1) of the tile whose ids are in key_ids[0], into the one buffer + auto load_one = [&](int pass) { + const __nv_bfloat16* src = pass == 0 ? KC : VC; + 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[0][k]; + qmm_frag::cp16z(&tiles_kv[k][8 * v], src + (static_cast(max(key, 0)) * HK + hk) * HD + 8 * v, key >= 0); + } + qmm_frag::commit(); + }; + float mm[2] = {-INFINITY, -INFINITY}, ll[2] = {0.0f, 0.0f}, oo[FUSE ? 8 : 1][4] = {}; + for (int cc = c; cc < c_end; ++cc) { + start = cc * CH; float m[2] = {-INFINITY, -INFINITY}, l[2] = {0.0f, 0.0f}; float o[8][4] = {}; const int tiles = min(n - start, CH); @@ -256,12 +271,22 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, for (int t0 = 0, it = 0; t0 < tiles; t0 += TK, ++it) { const int buf = NB == 2 ? it & 1 : 0; const bool ahead = NB == 2 && t0 + TK < tiles; // the next tile's loads go out before this tile's math - if (NB == 1) load(t0, 0); - if (ahead) load(t0 + TK, buf ^ 1); - Row* kt = tiles_kv + buf * 2 * TK; - Row* vt = kt + TK; + if constexpr (ONE) { + if (threadIdx.x < TK) { + const int ki = start + t0 + threadIdx.x; + key_ids[0][threadIdx.x] = ki < n ? (sparse ? IDS[static_cast(r) * IDW + ki] : ki) : -1; + } + __syncthreads(); + load_one(0); + } else { + if (NB == 1) load(t0, 0); + if (ahead) load(t0 + TK, buf ^ 1); + } + Row* kt = tiles_kv + (ONE ? 0 : buf * 2 * TK); + Row* vt = ONE ? kt : kt + TK; const int* keys = key_ids[buf]; - if (ahead) qmm_frag::wait<3>(); else qmm_frag::wait<1>(); + if constexpr (ONE) qmm_frag::wait<0>(); + else if (ahead) qmm_frag::wait<3>(); else qmm_frag::wait<1>(); __syncthreads(); constexpr int JT = TK / 32; // a warp's 8-key tiles float s[JT][4] = {}; @@ -278,6 +303,10 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, qmm_frag::mma(s[jt], a1, b[2], b[3]); } } + if constexpr (ONE) { // every warp has read the keys: the values take their place + __syncthreads(); + load_one(1); + } float tile_m[2] = {-INFINITY, -INFINITY}; #pragma unroll for (int jt = 0; jt < JT; ++jt) @@ -319,7 +348,7 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, rows[h] = __fadd_rn(rows[h], __shfl_xor_sync(0xffffffffu, rows[h], 2)); if (t == 0) red_l[warp][g + 8 * h] = rows[h]; } - if (ahead) qmm_frag::wait<2>(); else qmm_frag::wait<0>(); + if (ONE) qmm_frag::wait<0>(); else if (ahead) qmm_frag::wait<2>(); else qmm_frag::wait<0>(); __syncthreads(); #pragma unroll for (int h = 0; h < 2; ++h) { @@ -345,11 +374,27 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, } __syncthreads(); } + if constexpr (FUSE) { // fn_attn_merge's step for this chunk, element by element +#pragma unroll + for (int h = 0; h < 2; ++h) { + const bool active = l[h] > 0.0f; + const float next = active ? fmaxf(mm[h], m[h]) : mm[h]; + const float a = active ? (mm[h] == -INFINITY ? 0.0f : tf_exp(__fsub_rn(mm[h], next))) : 1.0f; + const float bb = active ? tf_exp(__fsub_rn(m[h], next)) : 0.0f; +#pragma unroll + for (int q = 0; q < 8; ++q) +#pragma unroll + for (int e = 0; e < 2; ++e) oo[FUSE ? q : 0][2 * h + e] = __fmaf_rn(o[q][2 * h + e], bb, __fmul_rn(oo[FUSE ? q : 0][2 * h + e], a)); + ll[h] = __fmaf_rn(l[h], bb, __fmul_rn(ll[h], a)); + mm[h] = next; + } + continue; + } #pragma unroll for (int h = 0; h < 2; ++h) { const int i = g + 8 * h; if (i >= G) continue; - const long long base = (static_cast(r) * NCH + c) * H + hk * G + i; + const long long base = (static_cast(r) * NCH + cc) * H + hk * G + i; #pragma unroll for (int q = 0; q < 8; ++q) { float2 v = make_float2(o[q][2 * h], o[q][2 * h + 1]); @@ -360,6 +405,21 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, PL[base] = l[h]; } } + } + if constexpr (FUSE) { // o / l, rounded once, as fn_attn_merge writes OUT + __nv_bfloat16* OUT = reinterpret_cast<__nv_bfloat16*>(PO); +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int i = g + 8 * h; + if (i >= G) continue; + const long long row = (static_cast(r) * H + hk * G + i) * HD; +#pragma unroll + for (int q = 0; q < 8; ++q) +#pragma unroll + for (int e = 0; e < 2; ++e) + OUT[row + warp * 64 + q * 8 + 2 * t + e] = __float2bfloat16_rn(__fdiv_rn(oo[q][2 * h + e], ll[h])); + } + } } // attn_chunks one tile at a time (prompt chunks: two blocks a multiprocessor). @@ -372,6 +432,25 @@ extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks(const __nv_bflo attn_chunks<1>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); } +// attn_chunks with one buffer and every chunk of a row merged in place (prompt rows): OUT (bf16, in PO's place), no partials. +extern "C" __global__ void __launch_bounds__(128) fn_attn_fused(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, + const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, + float* __restrict__ OUT, const int* __restrict__ IDS, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, + float SCALE, int IDW, int QSA) { + attn_chunks<1, true, true>(Q, KC, VC, POS0, OUT, nullptr, nullptr, IDS, NKR, SPR, H, HK, G, CH, 1, SCALE, IDW, QSA); +} + +// attn_chunks with one buffer, keys then values (prompt chunks: three blocks a multiprocessor), the same sums. +extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks_one(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, + 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) { + attn_chunks<1, true>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); +} + // attn_chunks with the next tile's loads in flight (a decode round's few rows): the same sums. extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks_ahead(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 611fbdc7da..6465b82418 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -331,33 +331,40 @@ pub fn Pass(comptime Tri: type) type { // 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); const nb = scoreStride(self.seq.cap); - var r0: usize = 0; - var b: usize = 0; - while (r0 < n) : ({ - r0 += att_rows; - b += 1; - }) { - const m = @min(att_rows, n - r0); - const ends = self.pos0 + r0 + m; - if (ends > self.seq.cap) return error.UnexpectedTensor; - 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))); - var a: cuda.Args = .{}; - a.add(self.seq.scores); - a.add(pos); - a.add(s.qids); - a.add(s.qnk); - a.add(s.qspr); - a.add(cint(nb)); - a.add(cint(idw)); - a.add(cint(top)); - 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.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * out_k * 2, s.qnk, s.qspr, m, kv_heads, true); + // prompt rows: selections for up to fused_rows rows, then one fused launch (chunks merged in the program) + const fused = n > att_rows; + const group: usize = if (fused) parts.fused_rows else att_rows; + var g0: usize = 0; + while (g0 < n) : (g0 += group) { + const gm = @min(group, n - g0); + var r0 = g0; + while (r0 < g0 + gm) : (r0 += att_rows) { + const m = @min(att_rows, g0 + gm - r0); + const ends = self.pos0 + r0 + m; + if (ends > self.seq.cap) return error.UnexpectedTensor; + const pos = at.pos + (r0 / att_rows) * 16; + const iq = at.iq + r0 * index_heads * index_dim * 2; + const sel = r0 - g0; // this block's rows in the selection buffers + // 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))); + var a: cuda.Args = .{}; + a.add(self.seq.scores); + a.add(pos); + a.add(s.qids + sel * idw * 4); + a.add(s.qnk + sel * 4); + a.add(s.qspr + sel * 4); + a.add(cint(nb)); + a.add(cint(idw)); + a.add(cint(top)); + a.add(cint(ratio)); + try cuda.launch.launch(self.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, self.stream, &a); + if (!fused) { + 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.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * out_k * 2, s.qnk, s.qspr, m, kv_heads, true); + } + } + if (fused) try self.tri.attnFused(at.q + g0 * q_heads * head_dim * 2, kc, vc, at.pos + (g0 / att_rows) * 16, at.ao + g0 * out_k * 2, s.qids, s.qnk, s.qspr, gm, kv_heads); } } diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index c070c68190..72c1206c9b 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -43,6 +43,8 @@ pub const pss_n = (dims / 256) * streams; pub const proj_row = @max(gdn.proj_width, proj_n); /// Attention rows a _chunks launch takes (its scratch: rows x nch x heads x dim fp32). pub const att_rows: usize = 256; +/// Prompt rows one fused attention launch takes (their selections staged att_rows at a time). +pub const fused_rows: usize = 1024; /// Rows a pass holds at most unless asked for more. pub const default_chunk: usize = 1536; const ratio = wt.index_ratio; @@ -329,7 +331,7 @@ pub const Scratch = struct { n += cap * (2 * gdn.nk * gdn.dk + 2 * gdn.nv) * 4; n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2; n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; - n += att_rows * (idw + 2) * 4; + n += fused_rows * (idw + 2) * 4; n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4; n += cap * moe_w * 6 + cap * dims * 4; n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; @@ -373,9 +375,9 @@ pub const Scratch = struct { s.po = try s.add(gpa, d, att_rows * nch * q_heads * head_dim * 4); s.pm = try s.add(gpa, d, att_rows * nch * q_heads * 4); s.pl = try s.add(gpa, d, att_rows * nch * q_heads * 4); - s.qids = try s.add(gpa, d, att_rows * idw * 4); - s.qnk = try s.add(gpa, d, att_rows * 4); - s.qspr = try s.add(gpa, d, att_rows * 4); + s.qids = try s.add(gpa, d, fused_rows * idw * 4); + s.qnk = try s.add(gpa, d, fused_rows * 4); + s.qspr = try s.add(gpa, d, fused_rows * 4); s.ao = try s.add(gpa, d, cap * out_k * 2); s.gated = try s.add(gpa, d, cap * out_k * 2); s.gxs_attn = try s.add(gpa, d, cap * (out_k / 32) * 4); diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index 643015d435..eab5fc56e1 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -11,6 +11,10 @@ const tile_heads = 16; fn chunkSmem(buffers: u32) u32 { return (tile_heads + buffers * 2 * 32) * (tile_head_dim + 8) * 2 + tile_heads * (32 + 8) * 2; } +/// The one-buffer form: a tile's keys, then its values, in the same rows of shared memory. +fn oneSmem() u32 { + return (tile_heads + 32) * (tile_head_dim + 8) * 2 + tile_heads * (32 + 8) * 2; +} /// Rows up to which attention reads the next tile ahead (fewer blocks, so latency, not occupancy, sets the time). const ahead_rows = 16; @@ -45,6 +49,8 @@ const Fns = struct { scores: cuda.Function, chunks: cuda.Function, chunks_ahead: cuda.Function, + chunks_one: cuda.Function, + fused: cuda.Function, merge: cuda.Function, }; @@ -115,10 +121,14 @@ pub const Glue = struct { .scores = try g.attn.function("fn_qsa_scores"), .chunks = try g.attn.function("fn_attn_chunks"), .chunks_ahead = try g.attn.function("fn_attn_chunks_ahead"), + .chunks_one = try g.attn.function("fn_attn_chunks_one"), + .fused = try g.attn.function("fn_attn_fused"), .merge = try g.attn.function("fn_attn_merge"), }; try g.f.chunks.allowDynamicShared(chunkSmem(1)); try g.f.chunks_ahead.allowDynamicShared(chunkSmem(2)); + try g.f.chunks_one.allowDynamicShared(oneSmem()); + try g.f.fused.allowDynamicShared(oneSmem()); return g; } @@ -194,7 +204,19 @@ pub const Glue = struct { a.add(@as(f32, 0.0625)); ints(&a, &.{ wt.index_idw, @intFromBool(qsa) }); 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); + try go(if (ahead) g.f.chunks_ahead else g.f.chunks_one, s, .{ rows, kv, chunks }, 128, if (ahead) chunkSmem(2) else oneSmem(), &a); + } + + /// Prompt rows' attention with every chunk merged in the program (fn_attn_fused): out bf16, attnChunks then + pub fn attnFused(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, out: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize) !void { + const group = wt.q_heads / kv; + if (group > tile_heads) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ q, kc, vc, pos, out, ids, nk, sparse }); + ints(&a, &.{ wt.q_heads, kv, group, 512 }); + a.add(@as(f32, 0.0625)); + ints(&a, &.{ wt.index_idw, 1 }); + try go(g.f.fused, s, .{ rows, kv, 1 }, 128, oneSmem(), &a); } pub fn attnMerge(g: *const Glue, s: cuda.Stream, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 05143c866f..21beb91525 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -249,6 +249,12 @@ pub const Tri = struct { }, &.{ ci("N", n), ci("K", k), ci("D", dims), ci("SS", streams), ci("BM", 16), ci("DB", 32), ci("GPI", 2), ci("SBN", 64) }); } + /// Prompt rows' attention, chunks merged in the program, ours only. + pub fn attnFused(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, out: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize) !void { + const g = t.glue("_chunks") orelse return error.NoGlue; + return g.attnFused(t.s, q, kc, vc, pos, out, ids, nk, sparse, rows, kv); + } + /// qmm.hc_down for one decode row. N=324 is the layer down stacked with inject. M folds when it is 1. pub fn hcDown(t: Tri, h: u64, pss: u64, scale: u64, normed: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, eps: f32, rows: usize, n: usize, k: usize, dims: usize, nc: usize, streams: usize, sk: usize) !void { if (t.glue("_qmm_hcdown")) |g| return g.hcDown(t.s, h, pss, scale, normed, weight, scales, biases, out, part, eps, rows, n, k, dims, nc, streams, sk); From 6c03e03f43a0ac94a06f06f2f0d522a22e7765f3 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 12:44:31 +0200 Subject: [PATCH 03/16] Flash Next CUDA: prompt rows' dense matmuls on a 64 x 128 tile of qmm.cu's factored arithmetic (the lane tile's bits) --- zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/flashnext_fact.cu | 260 ++++++++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_batch.zig | 11 +- .../families/flashnext/cuda_batch_state.zig | 10 + 5 files changed, 280 insertions(+), 3 deletions(-) create mode 100644 zig/kernels/cuda/flashnext_fact.cu diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 1568d5c453..87ffa3bca2 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -27,6 +27,7 @@ pub const kernels = [_]Kernel{ .{ .name = "probe", .flags = &.{"-O3"} }, .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5 .{ .name = "qmm_prefill", .flags = &.{"-O3"} }, + .{ .name = "flashnext_fact", .flags = &.{"-O3"} }, // ours: prompt rows on qmm.cu's factored arithmetic .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7 .{ .name = "experts_prefill", .flags = &.{"-O3"} }, .{ .name = "experts_pack", .flags = &.{"-O3"} }, diff --git a/zig/kernels/cuda/flashnext_fact.cu b/zig/kernels/cuda/flashnext_fact.cu new file mode 100644 index 0000000000..a29f90fb52 --- /dev/null +++ b/zig/kernels/cuda/flashnext_fact.cu @@ -0,0 +1,260 @@ +// Prompt GEMM on #952's factored arithmetic (qmm_kernel's bits): 2 groups a stage, mma per group then fp32 fmas. +#include "qmm_frag.cuh" + +namespace fact { +using namespace qmm_frag; + +template +struct Tile { + static constexpr int GS = 32; + static constexpr int THREADS = WM * WN * 32; + static constexpr int MT = BM / WM / 16; + static constexpr int NT = BN / WN / 8; + static constexpr int ROW = GS * 2; + static constexpr int CHUNKS = ROW / 16; + static constexpr int XG = BM * ROW; + static constexpr int WG = BN * GS / 2; + static constexpr int SG = BN * 2; + static constexpr int X = GPS * XG; + static constexpr int W = GPS * WG; + static constexpr int S = GPS * SG; + static constexpr int XS = GPS * BM * 4; + static constexpr int STAGE = X + W + 2 * S + XS; + static constexpr int SMEM = STAGES * STAGE; +}; + +template +__global__ void __launch_bounds__(WM * WN * 32, MINB) fact_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, + void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, int group) { + using T = Tile; + constexpr int GS = T::GS; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int KG = K / GS, cut = KG / SK, per = FOLD ? KG : cut, slice = FOLD ? 0 : blockIdx.z, g0 = slice * per; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y; + const int stages_n = per / GPS; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + auto load = [&](int s, int gfirst) { + unsigned char* p = stage(s); +#pragma unroll + for (int gi = 0; gi < GPS; ++gi) { + const int g = gfirst + gi; + unsigned char* px = p + gi * T::XG; + for (int c = tid; c < BM * T::CHUNKS; c += T::THREADS) { + const int r = c / T::CHUNKS, ch = c % T::CHUNKS; + const int row = min(m0 + r, M - 1); + cp16z(px + r * T::ROW + swz(r, ch) * 16, x + static_cast(row) * ldx + g * GS + ch * 8, + m0 + r < M); + } + unsigned char* pw = p + T::X + gi * T::WG; + constexpr int TILE_BYTES = 64 * GS / 2; + for (int c = tid; c < T::WG / 16; c += T::THREADS) { + const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + const size_t tile = static_cast(n0 / 64 + t) * KG + g; + if (n0 + t * 64 < npad) cp16(pw + c * 16, reinterpret_cast(w) + tile * TILE_BYTES + off * 16); + } + unsigned char* ps = p + T::X + T::W + gi * T::SG; + for (int c = tid; c < 2 * (T::SG / 16); c += T::THREADS) { + const int which = c / (T::SG / 16), off = c % (T::SG / 16); + const __nv_bfloat16* src = (which ? biases : scales) + static_cast(g) * npad + n0 + off * 8; + if (n0 + off * 8 < npad) cp16(ps + which * T::S + off * 16, src); + } + float* pxs = reinterpret_cast(p + T::X + T::W + 2 * T::S) + gi * BM; + for (int r = tid; r < BM; r += T::THREADS) cp4(pxs + r, xs + static_cast(min(m0 + r, M - 1)) * KG + g); + } + }; + + float acc[T::MT][T::NT][4]; + float tot[FOLD ? T::MT : 1][FOLD ? T::NT : 1][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f; + float d[2][T::MT][T::NT][4], sv[2][T::NT][2], bv[2][T::NT][2], xv[2][T::MT][2]; + + auto epilogue = [&](int b) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc[i][j][e] = __fmaf_rn(xv[b][i][e >> 1], bv[b][j][e & 1], __fmaf_rn(d[b][i][j][e], sv[b][j][e & 1], acc[i][j][e])); + }; + // a slice ends after group gl (0-based, in this CTA): fold its total in slice order and restart from zero + auto slice_end = [&](int gl) { + if constexpr (FOLD) { + if ((gl + 1) % cut == 0) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) { + tot[i][j][e] = gl + 1 == cut ? acc[i][j][e] : __fadd_rn(tot[i][j][e], acc[i][j][e]); + acc[i][j][e] = 0.0f; + } + } + } + }; + +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < stages_n) load(s, g0 + s * GPS); + commit(); + } + for (int it = 0; it < stages_n; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < stages_n) load(next % STAGES, g0 + next * GPS); + commit(); + const unsigned char* p = stage(it % STAGES); +#pragma unroll + for (int gi = 0; gi < GPS; ++gi) { + constexpr int dummy = 0; + (void)dummy; + const int b = gi & 1; + const unsigned char* px = p + gi * T::XG; + const uint32_t* pw = reinterpret_cast(p + T::X + gi * T::WG); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W + gi * T::SG); + const float* pxs = reinterpret_cast(p + T::X + T::W + 2 * T::S) + gi * BM; + uint32_t words[T::NT]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) words[j] = pw[(wn * T::NT + j) * 32 + lane]; + if constexpr (BLOCKED) { + // each (i, j) block's two k16 steps, then its epilogue: d lives a few registers, the bits are the same + uint32_t a[2][T::MT][4]; +#pragma unroll + for (int kt = 0; kt < 2; ++kt) +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = kt * 2 + (lane >> 4); + ldmatrix4(a[kt][i], px + r * T::ROW + swz(r, ch) * 16); + } + float sj[T::NT][2], bj[T::NT][2], xi[T::MT][2]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; + const __nv_bfloat162 s2 = *reinterpret_cast(ps + col); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + T::S / 2 + col); + sj[j][0] = __low2float(s2); sj[j][1] = __high2float(s2); + bj[j][0] = __low2float(b2); bj[j][1] = __high2float(b2); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int row = wm * (BM / WM) + i * 16 + (lane >> 2); + xi[i][0] = pxs[row]; + xi[i][1] = pxs[row + 8]; + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const uint32_t q0 = pair(words[j], 0), q1 = pair(words[j], 4), q2 = pair(words[j], 8), q3 = pair(words[j], 12); +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + float dd[4]; + mma0(dd, a[0][i], q0, q1); + mma(dd, a[1][i], q2, q3); +#pragma unroll + for (int e = 0; e < 4; ++e) + acc[i][j][e] = __fmaf_rn(xi[i][e >> 1], bj[j][e & 1], __fmaf_rn(dd[e], sj[j][e & 1], acc[i][j][e])); + } + } + slice_end(it * GPS + gi); + continue; + } +#pragma unroll + for (int kt = 0; kt < GS / 16; ++kt) { + uint32_t a[T::MT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = kt * 2 + (lane >> 4); + ldmatrix4(a[i], px + r * T::ROW + swz(r, ch) * 16); + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const uint32_t b0 = pair(words[j], (kt & 1) * 8); + const uint32_t b1 = pair(words[j], (kt & 1) * 8 + 4); +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[b][i][j], a[i], b0, b1); + else mma(d[b][i][j], a[i], b0, b1); + } + } + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; + const __nv_bfloat162 s2 = *reinterpret_cast(ps + col); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + T::S / 2 + col); + sv[b][j][0] = __low2float(s2); sv[b][j][1] = __high2float(s2); + bv[b][j][0] = __low2float(b2); bv[b][j][1] = __high2float(b2); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int row = wm * (BM / WM) + i * 16 + (lane >> 2); + xv[b][i][0] = pxs[row]; + xv[b][i][1] = pxs[row + 8]; + } + const int gl = it * GPS + gi; + if constexpr (DBUF) { + if (gl > 0) { + epilogue(b ^ 1); + slice_end(gl - 1); + } + } else { + epilogue(b); + slice_end(gl); + } + } + } + if constexpr (DBUF && !BLOCKED) { + if (stages_n > 0) { + epilogue((GPS - 1) & 1); + slice_end(stages_n * GPS - 1); + } + } + wait<0>(); +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = m0 + wm * (BM / WM) + i * 16 + (lane >> 2) + h * 8; + if (row >= M) continue; + const float v0 = FOLD ? tot[i][j][2 * h] : acc[i][j][2 * h]; + const float v1 = FOLD ? tot[i][j][2 * h + 1] : acc[i][j][2 * h + 1]; + if (OUT16 && !FOLD && SK == 1) { + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * N + col; + if (col + 1 < N && (N & 1) == 0) + *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1); + else { + if (col < N) dst[0] = __float2bfloat16_rn(v0); + if (col + 1 < N) dst[1] = __float2bfloat16_rn(v1); + } + continue; + } + float* dst = (SK > 1 && !FOLD) ? part + (static_cast(slice) * M + row) * N + col + : (FOLD ? part : reinterpret_cast(out)) + static_cast(row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } + } +} + +} // namespace fact + +// The prompt form the batched pass launches from 64 rows: 64 x 128 on 2 x 4 warps, bf16 out (one slice) or fp32 slices. +template __global__ void fact::fact_kernel<64, 128, 2, 4, 3, 2, false, 2, false, true, false>(const __nv_bfloat16*, + const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, float*, int, int, int, int, int, int, int); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 75722549aa..5125338a85 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -20,6 +20,7 @@ pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &. pub const qmm_group: []const u8 = if (available) &Blob("fatbin_qmm_group").bytes else &.{}; pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").bytes else &.{}; +pub const flashnext_fact: []const u8 = if (available) &Blob("fatbin_flashnext_fact").bytes else &.{}; pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 6465b82418..0dacf1a99c 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -113,8 +113,10 @@ pub fn Pass(comptime Tri: type) type { if (sk > 1) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes * (sk - 1), self.stream.handle), "cuMemsetD8Async"); return; } - const bm: usize = 16; - const bn: usize = 64; + // from 64 rows the prompt tile (flashnext_fact.cu), bit for bit the 16-row lane tile + const big = m >= 64; + const bm: usize = if (big) 64 else 16; + const bn: usize = if (big) 128 else 64; const k = w.k; const rows_t = cdiv(m, bm); const cols_t = cdiv(w.n, bn); @@ -129,7 +131,10 @@ pub fn Pass(comptime Tri: type) type { a.add(out); a.add(part); for ([_]usize{ m, w.n, k, sk, npad, k, sweep }) |v| a.add(cint(v)); - try cuda.launch.launch(self.k.lane, .{ .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(sk) }, .block = .{ .x = 128 }, .shared = qmm_smem }, self.stream, &a); + const f = if (big) self.k.fact else self.k.lane; + const block: u32 = if (big) 256 else 128; + const shared: u32 = if (big) parts.fact_smem else qmm_smem; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(sk) }, .block = .{ .x = block }, .shared = shared }, self.stream, &a); } pub fn sumSlices(self: *Self, part: u64, out: u64, bf16: bool, n: usize, sk: usize) !void { diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 72c1206c9b..474e51555b 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -52,6 +52,9 @@ const top = wt.index_top; const idw = wt.index_idw; const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; const qmm_smem: u32 = 9472; +/// Prompt rows on qmm.cu's arithmetic (flashnext_fact.cu): 64 x 128 tiles, two groups a stage, the lane tile's bits. +const cubin_fact: [:0]const u8 = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +pub const fact_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4); /// Pairs an expert item takes: the core rows kernel's tile. const rows_tile: usize = cuda.experts.RowsExperts.tile; const plan_small: usize = cuda.grouped.small; @@ -100,6 +103,8 @@ pub const Kernels = struct { router: cuda.grouped.Router, moe: cuda.experts.RowsExperts, lane: cuda.Function, + fact_mod: cuda.Module, + fact: cuda.Function, native_lane: cuda.qmmf.Lane, nvfp4: cuda.experts.Nvfp4Experts, front: cuda.Function, @@ -124,6 +129,10 @@ pub const Kernels = struct { errdefer k.batch.unload(); k.lane = try k.qmm.function(cubin_qmm); try k.lane.allowDynamicShared(qmm_smem); + k.fact_mod = try cuda.Module.load(d, cuda.kernels.flashnext_fact); + errdefer k.fact_mod.unload(); + k.fact = try k.fact_mod.function(cubin_fact); + try k.fact.allowDynamicShared(fact_smem); k.router = try cuda.grouped.Router.resolve(k.experts); k.moe = try cuda.experts.RowsExperts.resolve(k.rows); k.native_lane = try cuda.qmmf.Lane.load(d, (try d.capability(0)) / 10); @@ -150,6 +159,7 @@ pub const Kernels = struct { k.batch.unload(); k.rows.unload(); k.experts.unload(); + k.fact_mod.unload(); k.qmm.unload(); } }; From 55e3f623aeaebeff365a763f128a76f73bf7bda6 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 13:00:01 +0200 Subject: [PATCH 04/16] Flash Next CUDA: prompt rows' hyper-connection read-out on factored tiles, fn_hc_down's and fn_hc_upmix's bits --- zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/flashnext_hc.cu | 39 +++ zig/kernels/cuda/flashnext_hcf.cu | 245 ++++++++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_batch.zig | 40 ++- .../families/flashnext/cuda_batch_state.zig | 19 +- zig/src/families/flashnext/cuda_glue.zig | 23 ++ zig/src/families/flashnext/cuda_launch.zig | 11 + 8 files changed, 377 insertions(+), 2 deletions(-) create mode 100644 zig/kernels/cuda/flashnext_hcf.cu diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 87ffa3bca2..05b8192f00 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -28,6 +28,7 @@ pub const kernels = [_]Kernel{ .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5 .{ .name = "qmm_prefill", .flags = &.{"-O3"} }, .{ .name = "flashnext_fact", .flags = &.{"-O3"} }, // ours: prompt rows on qmm.cu's factored arithmetic + .{ .name = "flashnext_hcf", .flags = glue }, // ours: the hyper-connection read-out's prompt tiles, fn_hc's bits .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7 .{ .name = "experts_prefill", .flags = &.{"-O3"} }, .{ .name = "experts_pack", .flags = &.{"-O3"} }, diff --git a/zig/kernels/cuda/flashnext_hc.cu b/zig/kernels/cuda/flashnext_hc.cu index 2865adde56..78ea5662c6 100644 --- a/zig/kernels/cuda/flashnext_hc.cu +++ b/zig/kernels/cuda/flashnext_hc.cu @@ -237,6 +237,45 @@ extern "C" __global__ void __launch_bounds__(1024) fn_hc_reduce_act(const float* } } +// fn_hc_norm's 1 / rms of each row's streams [M, SS], a thread each: the column blocks' sums in order. +extern "C" __global__ void fn_hc_rinv(const float* __restrict__ PSS, float* __restrict__ RINV, int M, float eps, int D, + int NC, int SS) { + const int i = blockIdx.x * blockDim.x + threadIdx.x, r = i / SS, st = i % SS; + if (r >= M) return; + float total = 0.0f; + for (int c = 0; c < NC; ++c) total = __fadd_rn(total, PSS[(static_cast(r) * NC + c) * SS + st]); + RINV[i] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(total, static_cast(D)), eps))); +} + +// fn_hc_mix with each normed value recomputed from the streams as fn_hc_norm rounds it (no normed rows read). +extern "C" __global__ void __launch_bounds__(256) fn_hc_mix_h(const __nv_bfloat16* __restrict__ UP, const __nv_bfloat16* __restrict__ H, + const float* __restrict__ RINV, const float* __restrict__ SCALE, + __nv_bfloat16* __restrict__ MIXED, float* __restrict__ XSM, + int D, int SS) { + __shared__ float ms[256]; + const int r = blockIdx.x, d = blockIdx.y * 256 + threadIdx.x; + const long long row = static_cast(r) * SS * D; + float m = 0.0f; + if (d < D) { + float total = 0.0f; + for (int s = 0; s < SS; ++s) { + const long long at = row + static_cast(s) * D + d; + const float nv = rbf(__fmul_rn(__fmul_rn(bf(H[at]), RINV[r * SS + s]), SCALE[s * D + d])); + total = __fadd_rn(total, rbf(__fmul_rn(bsig(bf(UP[at])), nv))); + } + m = rbf(__fdiv_rn(total, static_cast(SS))); + MIXED[static_cast(r) * D + d] = __float2bfloat16_rn(m); + } + ms[threadIdx.x] = m; + __syncthreads(); + const int g = threadIdx.x; + if (g < 8 && blockIdx.y * 256 + g * 32 < D) { + float sum = 0.0f; + for (int k = 0; k < 32; ++k) sum = __fadd_rn(sum, ms[g * 32 + k]); + XSM[static_cast(r) * (D / 32) + blockIdx.y * 8 + g] = sum; + } +} + // qmm._qmm_upmix, block (16 rows, 32 dims), two warps a stream (SS 4): the up rows, then the streams mixed in order. extern "C" __global__ void __launch_bounds__(256) fn_hc_upmix(const __nv_bfloat16* __restrict__ X, const float* __restrict__ XS, const uint4* __restrict__ W, const __nv_bfloat16* __restrict__ S, diff --git a/zig/kernels/cuda/flashnext_hcf.cu b/zig/kernels/cuda/flashnext_hcf.cu new file mode 100644 index 0000000000..6ebd1b0a9d --- /dev/null +++ b/zig/kernels/cuda/flashnext_hcf.cu @@ -0,0 +1,245 @@ +// #952's hyper-connection read-out (fn_hc_down / fn_hc_upmix bits) on prompt tiles: factored groups, norm in smem. +#include "qmm_frag.cuh" +#include "glue_math.cuh" + +namespace hcf { +using namespace qmm_frag; +using tf_glue::bf; +using tf_glue::rbf; + +template +struct Tile { + static constexpr int GS = 32; + static constexpr int THREADS = WM * WN * 32; + static constexpr int MT = BM / WM / 16; + static constexpr int NT = BN / WN / 8; + static constexpr int ROW = GS * 2; + static constexpr int CHUNKS = ROW / 16; + static constexpr int XG = BM * ROW; + static constexpr int WG = BN * 16; + static constexpr int SG = BN * 2; + static constexpr int X = GPS * XG; + static constexpr int W = GPS * WG; + static constexpr int S = GPS * SG; + static constexpr int XS = GPS * BM * 4; + static constexpr int SC = GPS * GS * 4; // the norm's scale for the stage's groups + static constexpr int STAGE = X + W + 2 * S + XS + SC; + static constexpr int RINV = BM * 4 * 4; + static constexpr int SMEM = STAGES * STAGE + RINV; +}; + +// fn_hc's nibbles(w, 2t): codes 2t (low half) and 2t + 1 (high half) as exact bf16 integers. +__device__ __forceinline__ uint32_t nib2(uint32_t w, int t) { + const uint32_t v = w >> (8 * t); + const uint32_t q = (v & 0xFu) | ((v << 12) & 0xF0000u) | 0x43004300u; + uint32_t r; +#if __CUDA_ARCH__ >= 900 + asm("sub.rn.bf16x2 %0, %1, %2;\n" : "=r"(r) : "r"(q), "r"(0x43004300u)); +#else + asm("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(r) : "r"(q), "r"(0x3F803F80u), "r"(0xC300C300u)); // qmm_frag::pair's form +#endif + return r; +} + +template +__global__ void __launch_bounds__(WM * WN * 32, MINB) hcf_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint4* __restrict__ w, + const __nv_bfloat16* __restrict__ sc, const __nv_bfloat16* __restrict__ bi, void* __restrict__ out, + int M, int N, int K, int SK, int ldx, int group, const float* __restrict__ pss, const float* __restrict__ scale, + float eps, int D, int NC, int SS) { + using T = Tile; + constexpr int GS = T::GS; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g8 = lane >> 2, t4 = lane & 3; + const int wm = warp / WN, wn = warp % WN; + const int KG = K / GS, cut = KG / SK, npad = (N + 63) / 64 * 64; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y; + const int stages_n = KG / GPS; + float* rinv_s = reinterpret_cast(buf + STAGES * T::STAGE); + + if constexpr (NORM) { + for (int i = tid; i < BM * SS; i += T::THREADS) { + const int r = min(m0 + i / SS, M - 1), st = i % SS; + float total = 0.0f; + for (int c = 0; c < NC; ++c) total = __fadd_rn(total, pss[(static_cast(r) * NC + c) * SS + st]); + rinv_s[i] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(total, static_cast(D)), eps))); + } + } + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + auto load = [&](int s, int gfirst) { + unsigned char* p = stage(s); +#pragma unroll + for (int gi = 0; gi < GPS; ++gi) { + const int g = gfirst + gi; + unsigned char* px = p + gi * T::XG; + for (int c = tid; c < BM * T::CHUNKS; c += T::THREADS) { + const int r = c / T::CHUNKS, ch = c % T::CHUNKS; + const int row = min(m0 + r, M - 1); + cp16z(px + r * T::ROW + swz(r, ch) * 16, x + static_cast(row) * ldx + g * GS + ch * 8, m0 + r < M); + } + unsigned char* pw = p + T::X + gi * T::WG; + for (int c = tid; c < BN; c += T::THREADS) { + const int col = n0 + c; + if (col < npad) cp16(pw + c * 16, w + (static_cast(col / 64) * KG + g) * 64 + col % 64); + } + unsigned char* ps = p + T::X + T::W + gi * T::SG; + for (int c = tid; c < BN; c += T::THREADS) { // c < BN / 2: scale pairs, the rest bias pairs + const int which = c / (BN / 2), off = (c % (BN / 2)) * 2, col = n0 + off; + uint32_t* dst = reinterpret_cast(ps + which * T::S + off * 2); + if (col < N) cp4(dst, (which ? bi : sc) + static_cast(g) * N + col); + else *dst = 0u; + } + if constexpr (NORM) { + float* psc = reinterpret_cast(p + T::X + T::W + 2 * T::S + T::XS) + gi * GS; + for (int c = tid; c < GS / 4; c += T::THREADS) cp16(psc + c * 4, scale + g * GS + c * 4); + } else { + float* pxs = reinterpret_cast(p + T::X + T::W + 2 * T::S) + gi * BM; + for (int r = tid; r < BM; r += T::THREADS) cp4(pxs + r, xs + static_cast(min(m0 + r, M - 1)) * KG + g); + } + } + }; + + float acc[T::MT][T::NT][4]; + float tot[FOLD ? T::MT : 1][FOLD ? T::NT : 1][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f; + +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < stages_n) load(s, s * GPS); + commit(); + } + for (int it = 0; it < stages_n; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < stages_n) load(next % STAGES, next * GPS); + commit(); + unsigned char* p = stage(it % STAGES); + if constexpr (NORM) { + // x = bf16(bf16(h) * rinv[row, stream] * scale[k]) as fn_hc_down rounds it, and each row's 32-group sum in order + float* pxs = reinterpret_cast(p + T::X + T::W + 2 * T::S); + for (int c = tid; c < GPS * BM; c += T::THREADS) { + const int gi = c / BM, r = c % BM, k0 = (it * GPS + gi) * GS; + unsigned char* px = p + gi * T::XG + r * T::ROW; + const float rv = rinv_s[r * SS + k0 / D]; + float sum = 0.0f; +#pragma unroll + for (int ch = 0; ch < T::CHUNKS; ++ch) { + uint4* cell = reinterpret_cast(px + swz(r, ch) * 16); + const uint4 raw = *cell; + const __nv_bfloat16* hv = reinterpret_cast(&raw); + const float* psc = reinterpret_cast(p + T::X + T::W + 2 * T::S + T::XS) + gi * GS + ch * 8; + const float4 sa = *reinterpret_cast(psc); + const float4 sb = *reinterpret_cast(psc + 4); + const float sv[8] = {sa.x, sa.y, sa.z, sa.w, sb.x, sb.y, sb.z, sb.w}; + __nv_bfloat16 o8[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + o8[i] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(bf(hv[i]), rv), sv[i])); + sum = __fadd_rn(sum, bf(o8[i])); + } + *cell = *reinterpret_cast(o8); + } + pxs[gi * BM + r] = sum; + } + __syncthreads(); + } +#pragma unroll + for (int gi = 0; gi < GPS; ++gi) { + const unsigned char* px = p + gi * T::XG; + const uint4* pw = reinterpret_cast(p + T::X + gi * T::WG); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W + gi * T::SG); + const float* pxs = reinterpret_cast(p + T::X + T::W + 2 * T::S) + gi * BM; + uint32_t a[2][T::MT][4]; +#pragma unroll + for (int kt = 0; kt < 2; ++kt) +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = kt * 2 + (lane >> 4); + ldmatrix4(a[kt][i], px + r * T::ROW + swz(r, ch) * 16); + } + float xi[T::MT][2]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int row = wm * (BM / WM) + i * 16 + g8; + xi[i][0] = pxs[row]; + xi[i][1] = pxs[row + 8]; + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int cw = wn * (BN / WN) + j * 8; + const uint4 words = pw[cw + g8]; + const uint32_t q0 = nib2(words.x, t4), q1 = nib2(words.y, t4), q2 = nib2(words.z, t4), q3 = nib2(words.w, t4); + const __nv_bfloat162 s2 = *reinterpret_cast(ps + cw + 2 * t4); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + T::S / 2 + cw + 2 * t4); + const float s0 = __low2float(s2), s1 = __high2float(s2), b0 = __low2float(b2), b1 = __high2float(b2); +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + float dd[4]; + mma0(dd, a[0][i], q0, q1); + mma(dd, a[1][i], q2, q3); + acc[i][j][0] = __fmaf_rn(xi[i][0], b0, __fmaf_rn(dd[0], s0, acc[i][j][0])); + acc[i][j][1] = __fmaf_rn(xi[i][0], b1, __fmaf_rn(dd[1], s1, acc[i][j][1])); + acc[i][j][2] = __fmaf_rn(xi[i][1], b0, __fmaf_rn(dd[2], s0, acc[i][j][2])); + acc[i][j][3] = __fmaf_rn(xi[i][1], b1, __fmaf_rn(dd[3], s1, acc[i][j][3])); + } + } + if constexpr (FOLD) { + const int gl = it * GPS + gi; + if ((gl + 1) % cut == 0) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) { + tot[i][j][e] = gl + 1 == cut ? acc[i][j][e] : __fadd_rn(tot[i][j][e], acc[i][j][e]); + acc[i][j][e] = 0.0f; + } + } + } + } + } + wait<0>(); +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = n0 + wn * (BN / WN) + j * 8 + 2 * t4; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = m0 + wm * (BM / WM) + i * 16 + g8 + h * 8; + if (row >= M) continue; + const float v0 = FOLD ? tot[i][j][2 * h] : acc[i][j][2 * h]; + const float v1 = FOLD ? tot[i][j][2 * h + 1] : acc[i][j][2 * h + 1]; + if constexpr (OUT16) { + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * N + col; + if (col + 1 < N) *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1); + else if (col < N) dst[0] = __float2bfloat16_rn(v0); + } else { + float* dst = reinterpret_cast(out) + static_cast(row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } + } + } +} + +} // namespace hcf + +// HC down: 32 x 192 tiles, the streams normed in shared memory, the 32 K slices folded in slice order (fp32 out). +template __global__ void hcf::hcf_kernel<32, 192, 2, 4, 4, 2, true, true, false, 2>(const __nv_bfloat16*, const float*, + const uint4*, const __nv_bfloat16*, const __nv_bfloat16*, void*, int, int, int, int, int, int, const float*, const float*, + float, int, int, int); +// HC up: 64 x 128 tiles, bf16 out (the up rows fn_hc_upmix mixes). +template __global__ void hcf::hcf_kernel<64, 128, 2, 4, 3, 2, false, false, true, 2>(const __nv_bfloat16*, const float*, + const uint4*, const __nv_bfloat16*, const __nv_bfloat16*, void*, int, int, int, int, int, int, const float*, const float*, + float, int, int, int); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 5125338a85..80b9940993 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -21,6 +21,7 @@ pub const qmm_group: []const u8 = if (available) &Blob("fatbin_qmm_group").bytes pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").bytes else &.{}; pub const flashnext_fact: []const u8 = if (available) &Blob("fatbin_flashnext_fact").bytes else &.{}; +pub const flashnext_hcf: []const u8 = if (available) &Blob("fatbin_flashnext_hcf").bytes else &.{}; pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 0dacf1a99c..d24768e3c8 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -59,6 +59,11 @@ fn cint(v: usize) c_int { return @intCast(v); } +/// tile_of's sweep: row tiles a column pass keeps in 12 MiB of the input. +fn sweepOf(m: usize, bm: usize, k: usize) usize { + return @max(@as(usize, 1), @min(cdiv(m, bm), (12 << 20) / (bm * k * 2))); +} + fn cdiv(a: usize, b: usize) usize { return (a + b - 1) / b; } @@ -189,6 +194,38 @@ pub fn Pass(comptime Tri: type) type { try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", rs, n, dims, streams, 2, 10, 11, 1); } + /// Prompt rows' read-out on flashnext_hcf.cu's tiles (fn_hc_down's and fn_hc_upmix's bits): the down rows with their + /// K slices folded, the reduce on one slice, the up rows in bf16, then the mix renorming the streams on the fly. + fn readoutTiles(self: *Self, hc_w: HcW, r0: usize, n: usize) !void { + const s = self.s; + const h = s.h + r0 * wide * 2; + const pss = s.pss + r0 * pss_n * 4; + const rinv = s.rinv + r0 * streams * 4; + const act = s.act + r0 * low * 2; + const xs_act = s.xs_act + r0 * (low / 32) * 4; + const inj_ptr = if (hc_w.has_inj == 1) s.inj + r0 * streams * 2 else act; + const nc = dims / 256; + var a: cuda.Args = .{}; + for ([_]u64{ h, 0, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.hc_part }) |v| a.add(v); + for ([_]usize{ n, hc_w.n_down, wide, hc_sk, wide, sweepOf(n, 32, wide) }) |v| a.add(cint(v)); + a.add(pss); + a.add(hc_w.scale.ptr); + a.add(eps); + for ([_]usize{ dims, nc, streams }) |v| a.add(cint(v)); + try cuda.launch.launch(self.k.hcf_down, .{ .grid = .{ .x = @intCast(cdiv(n, 32) * cdiv(hc_w.n_down, 192)) }, .block = .{ .x = 256 }, .shared = parts.hcf_down_smem }, self.stream, &a); + try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, 1, n, streams, low, hc_w.n_down, hc_w.has_inj); + var b: cuda.Args = .{}; + for ([_]u64{ act, xs_act, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.proj }) |v| b.add(v); + for ([_]usize{ n, wide, low, 1, low, sweepOf(n, 64, low) }) |v| b.add(cint(v)); + b.add(@as(u64, 0)); + b.add(@as(u64, 0)); + b.add(@as(f32, 0)); + for ([_]usize{ dims, nc, streams }) |v| b.add(cint(v)); + try cuda.launch.launch(self.k.hcf_up, .{ .grid = .{ .x = @intCast(cdiv(n, 64) * cdiv(wide, 128)) }, .block = .{ .x = 256 }, .shared = parts.hcf_up_smem }, self.stream, &b); + try self.tri.hcRinv(pss, rinv, eps, n, dims, nc, streams); + try self.tri.hcMixH(s.proj, h, rinv, hc_w.scale.ptr, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); + } + /// A read-out over rows [r0, r0 + n) of the streams: mixed and its group sums, and the inject gates. pub fn readout(self: *Self, hc_w: HcW, r0: usize, n: usize) !void { const s = self.s; @@ -208,6 +245,7 @@ pub fn Pass(comptime Tri: type) type { try self.tri.hcDenseMix(s.hc_up, norm, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); return; } + if (n >= 64) return self.readoutTiles(hc_w, r0, n); try self.tri.hcDown(h, pss, hc_w.scale.ptr, norm, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn, s.hc_part, eps, n, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); const inj_ptr = if (hc_w.has_inj == 1) inj else act; try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, hc_sk, n, streams, low, hc_w.n_down, hc_w.has_inj); @@ -349,7 +387,7 @@ pub fn Pass(comptime Tri: type) type { if (ends > self.seq.cap) return error.UnexpectedTensor; const pos = at.pos + (r0 / att_rows) * 16; const iq = at.iq + r0 * index_heads * index_dim * 2; - const sel = r0 - g0; // this block's rows in the selection buffers + const sel = r0 - g0; // this block's rows in the selection buffers // 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))); var a: cuda.Args = .{}; diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 474e51555b..a2add204bf 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -54,6 +54,11 @@ const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4EL const qmm_smem: u32 = 9472; /// Prompt rows on qmm.cu's arithmetic (flashnext_fact.cu): 64 x 128 tiles, two groups a stage, the lane tile's bits. const cubin_fact: [:0]const u8 = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +/// The hyper-connection read-out's prompt tiles (flashnext_hcf.cu), fn_hc_down's and fn_hc_upmix's bits. +const cubin_hcf_down: [:0]const u8 = "_ZN3hcf10hcf_kernelILi32ELi192ELi2ELi4ELi4ELi2ELb1ELb1ELb0ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii"; +const cubin_hcf_up: [:0]const u8 = "_ZN3hcf10hcf_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELb0ELb1ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii"; +pub const hcf_down_smem: u32 = 4 * (2 * 32 * 64 + 2 * 192 * 16 + 2 * 2 * 192 * 2 + 2 * 32 * 4 + 2 * 32 * 4) + 32 * 16; +pub const hcf_up_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4 + 2 * 32 * 4) + 64 * 16; pub const fact_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4); /// Pairs an expert item takes: the core rows kernel's tile. const rows_tile: usize = cuda.experts.RowsExperts.tile; @@ -105,6 +110,9 @@ pub const Kernels = struct { lane: cuda.Function, fact_mod: cuda.Module, fact: cuda.Function, + hcf_mod: cuda.Module, + hcf_down: cuda.Function, + hcf_up: cuda.Function, native_lane: cuda.qmmf.Lane, nvfp4: cuda.experts.Nvfp4Experts, front: cuda.Function, @@ -133,6 +141,12 @@ pub const Kernels = struct { errdefer k.fact_mod.unload(); k.fact = try k.fact_mod.function(cubin_fact); try k.fact.allowDynamicShared(fact_smem); + k.hcf_mod = try cuda.Module.load(d, cuda.kernels.flashnext_hcf); + errdefer k.hcf_mod.unload(); + k.hcf_down = try k.hcf_mod.function(cubin_hcf_down); + try k.hcf_down.allowDynamicShared(hcf_down_smem); + k.hcf_up = try k.hcf_mod.function(cubin_hcf_up); + try k.hcf_up.allowDynamicShared(hcf_up_smem); k.router = try cuda.grouped.Router.resolve(k.experts); k.moe = try cuda.experts.RowsExperts.resolve(k.rows); k.native_lane = try cuda.qmmf.Lane.load(d, (try d.capability(0)) / 10); @@ -159,6 +173,7 @@ pub const Kernels = struct { k.batch.unload(); k.rows.unload(); k.experts.unload(); + k.hcf_mod.unload(); k.fact_mod.unload(); k.qmm.unload(); } @@ -221,6 +236,7 @@ pub const Scratch = struct { h: u64 = 0, pss: u64 = 0, norm: u64 = 0, + rinv: u64 = 0, dn: u64 = 0, hc_part: u64 = 0, hc_up: u64 = 0, @@ -334,7 +350,7 @@ pub const Scratch = struct { const items = cdiv(pairs, tile) + experts_n + 1; var n: usize = 0; n += cap * 4 + (cap + 1) * 16; - n += cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; + n += cap * streams * 4 + cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; n += cap * wide * 2 + 8 * cap * dims * 4; n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; @@ -362,6 +378,7 @@ pub const Scratch = struct { s.h = try s.add(gpa, d, cap * wide * 2); s.pss = try s.add(gpa, d, cap * pss_n * 4); s.norm = try s.add(gpa, d, cap * wide * 2); + s.rinv = try s.add(gpa, d, cap * streams * 4); s.dn = try s.add(gpa, d, cap * (low + 4) * 2); s.hc_part = try s.add(gpa, d, hc_sk * cap * (low + 4) * 4); s.hc_up = try s.add(gpa, d, cap * wide * 2); diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index eab5fc56e1..a5793011d1 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -38,6 +38,8 @@ const Fns = struct { down: cuda.Function, reduce_act: cuda.Function, upmix: cuda.Function, + rinv: cuda.Function, + mix_h: cuda.Function, ple_gate: cuda.Function, ple_conv: cuda.Function, add_streams: cuda.Function, @@ -110,6 +112,8 @@ pub const Glue = struct { .down = try g.hc.function("fn_hc_down"), .reduce_act = try g.hc.function("fn_hc_reduce_act"), .upmix = try g.hc.function("fn_hc_upmix"), + .rinv = try g.hc.function("fn_hc_rinv"), + .mix_h = try g.hc.function("fn_hc_mix_h"), .ple_gate = try g.hc.function("fn_ple_gate"), .ple_conv = try g.hc.function("fn_ple_conv"), .add_streams = try g.hc.function("fn_add_streams"), @@ -283,6 +287,25 @@ pub const Glue = struct { try go(g.f.reduce_act, s, .{ rows, 1, 1 }, @intCast(cdiv(@max(cols, low / 32), 32) * 32), 0, &a); } + /// fn_hc_norm's 1 / rms of each row's streams [rows, streams]. + pub fn hcRinv(g: *const Glue, s: cuda.Stream, pss: u64, rinv: u64, eps: f32, rows: usize, dims: usize, nc: usize, streams: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ pss, rinv }); + a.add(int(rows)); + a.add(eps); + ints(&a, &.{ dims, nc, streams }); + try go(g.f.rinv, s, .{ cdiv(rows * streams, 256), 1, 1 }, 256, 0, &a); + } + + /// hcMix with the normed values recomputed from the streams h, rinv and scale. + pub fn hcMixH(g: *const Glue, s: cuda.Stream, up: u64, h: u64, rinv: u64, scale: u64, mixed: u64, xsm: u64, rows: usize, dims: usize, streams: usize) !void { + if (dims % 32 != 0) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ up, h, rinv, scale, mixed, xsm }); + ints(&a, &.{ dims, streams }); + try go(g.f.mix_h, s, .{ rows, cdiv(dims, 256), 1 }, 256, 0, &a); + } + pub fn hcUpmix(g: *const Glue, s: cuda.Stream, act: u64, xs: u64, w: u64, sc: u64, b: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { if (streams != 4 or dims % 32 != 0 or k % 32 != 0) return error.RowShape; var a: cuda.Args = .{}; diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 21beb91525..12e8a548c2 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -234,6 +234,17 @@ pub const Tri = struct { } /// qmm.hc_upmix: the up projection and the stream mix. N is the wide row, K is the low rank. M folds when it is 1. + /// The read-out's 1 / rms rows, ours only. + pub fn hcRinv(t: Tri, pss: u64, rinv: u64, eps: f32, rows: usize, dims: usize, nc: usize, streams: usize) !void { + const g = t.glue("_qmm_hcdown") orelse return error.NoGlue; + return g.hcRinv(t.s, pss, rinv, eps, rows, dims, nc, streams); + } + + pub fn hcMixH(t: Tri, up: u64, h: u64, rinv: u64, scale: u64, mixed: u64, xsm: u64, rows: usize, dims: usize, streams: usize) !void { + const g = t.glue("_qmm_upmix") orelse return error.NoGlue; + return g.hcMixH(t.s, up, h, rinv, scale, mixed, xsm, rows, dims, streams); + } + pub fn hcUpmix(t: Tri, act: u64, xs: u64, weight: u64, scales: u64, biases: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { if (t.glue("_qmm_upmix")) |g| return g.hcUpmix(t.s, act, xs, weight, scales, biases, normed, mixed, xs_mixed, rows, n, k, dims, streams); try t.run("_qmm_upmix", .{ cdiv(rows, 16), u(dims / 32), 1 }, &.{ From 5a68c4ecb0cd01c63421402e8757257a80f7aab4 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 13:10:02 +0200 Subject: [PATCH 05/16] Flash Next CUDA: prompt rows' experts on factored 64-pair tiles, expert_rows' bits --- zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/flashnext_fexp.cu | 204 ++++++++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_batch.zig | 13 ++ .../families/flashnext/cuda_batch_state.zig | 15 ++ 5 files changed, 234 insertions(+) create mode 100644 zig/kernels/cuda/flashnext_fexp.cu diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 05b8192f00..df55c696b8 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -29,6 +29,7 @@ pub const kernels = [_]Kernel{ .{ .name = "qmm_prefill", .flags = &.{"-O3"} }, .{ .name = "flashnext_fact", .flags = &.{"-O3"} }, // ours: prompt rows on qmm.cu's factored arithmetic .{ .name = "flashnext_hcf", .flags = glue }, // ours: the hyper-connection read-out's prompt tiles, fn_hc's bits + .{ .name = "flashnext_fexp", .flags = &.{"-O3"} }, // ours: prompt pairs' experts on expert_rows' factored arithmetic .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7 .{ .name = "experts_prefill", .flags = &.{"-O3"} }, .{ .name = "experts_pack", .flags = &.{"-O3"} }, diff --git a/zig/kernels/cuda/flashnext_fexp.cu b/zig/kernels/cuda/flashnext_fexp.cu new file mode 100644 index 0000000000..21cbf0e628 --- /dev/null +++ b/zig/kernels/cuda/flashnext_fexp.cu @@ -0,0 +1,204 @@ +// g32 experts on prompt pair tiles with #952's factored arithmetic: expert_rows' bits at every row (64 pairs a tile). +#include +#include +#include "experts.cuh" + +namespace tf_fact_experts { + +template +struct Pre { + static constexpr int THREADS = WM * WN * 32, BM = 16 * RT * WM, SG = 64 / GS, XC = 8, WB = M * Geo::BLOCK; + static constexpr int XU = BM * XC, WU = WN * SG * WB, XSU = BM * SG / 4, SU = XU + WU + XSU; + static constexpr int XPT = (XU + THREADS - 1) / THREADS; + static constexpr int STAGE_BYTES = SU * 16; + static constexpr int STAGES = ST; + static __device__ __forceinline__ int xslot(int r, int c) { return r * XC + (c ^ ((r & 1) << 2)); } +}; + +// mma from zero (C = 0): the same sums as zeroed accumulators, no register moves. +__device__ __forceinline__ void mma_z(float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, + uint32_t b1) { + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%10, %10, %10, %10};\n" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(0.f)); +} + +// expert_rows' group_sum from the whole row: lane t's 8 inputs in order, the quad's (s0 + s1) + (s2 + s3). +__device__ __forceinline__ float row_gsum(const uint4 (&v)[4]) { + float s[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) { + float a = 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t u = comp(v[q], i); + a += lo_f(u); + a += hi_f(u); + } + s[q] = a; + } + return (s[0] + s[1]) + (s[2] + s[3]); +} + +__device__ __forceinline__ float gsum(const uint4& v) { + float s = 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t u = comp(v, i); + s += lo_f(u); + s += hi_f(u); + } + s += __shfl_xor_sync(0xffffffffu, s, 1); + s += __shfl_xor_sync(0xffffffffu, s, 2); + return s; +} + +template +__global__ void __launch_bounds__(WM * WN * 32) + fact_kernel(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, int KG, + int NB, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, void* __restrict__ out, int N, float limit) { + constexpr int GS = 32; + using G = Geo; + using P = Pre; + constexpr int STAGES = P::STAGES; + extern __shared__ uint4 sm[]; + const int nbt = (NB + WN - 1) / WN, it = blockIdx.x / nbt, cbt = blockIdx.x - it * nbt; + if (it >= __ldg(counts)) return; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, wm = warp / WN, wn = warp - wm * WN; + const int t = lane & 3, gq = lane >> 2, cb = cbt * WN + wn, cbs = min(WN, NB - cbt * WN); + const __nv_bfloat16* xsrc[P::XPT]; + int xdst[P::XPT]; +#pragma unroll + for (int i = 0; i < P::XPT; ++i) { + const int q = tid + i * P::THREADS, r = q / P::XC, c = q - r * P::XC; + xdst[i] = P::xslot(r, c); + xsrc[i] = nullptr; + if (q < P::XU && r < cnt) { + const int p = __ldg(members + first + r); + xsrc[i] = X + (size_t)(slots ? p / slots : p) * x_stride + 8 * c; + } + } + const uint4* wsrc = W + ((size_t)e * NB + cbt * WN) * (size_t)KG * P::WB; + auto stage = [&](int s, int g0) { + const int ng = min(P::SG, KG - g0); + uint4* xs = sm + s * P::SU; +#pragma unroll + for (int i = 0; i < P::XPT; ++i) + if (xsrc[i] && ((tid + i * P::THREADS) % P::XC) < ng * (GS / 8)) cp16(xs + xdst[i], xsrc[i] + (size_t)g0 * GS); + uint4* ws = xs + P::XU; + for (int q = tid; q < cbs * P::SG * P::WB; q += P::THREADS) { + const int j = q / (P::SG * P::WB), o = q - j * P::SG * P::WB; + if (o < ng * P::WB) cp16(ws + q, wsrc + ((size_t)j * KG + g0) * P::WB + o); + } + }; + const int base = 16 * RT * wm; + const int nt = max(0, min(RT, (cnt - base + 15) >> 4)); + const bool active = nt > 0 && cb < NB; + float acc[M][RT][NTW][4]; +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[m][r][j][0] = acc[m][r][j][1] = acc[m][r][j][2] = acc[m][r][j][3] = 0.f; + const int steps = (KG + P::SG - 1) / P::SG; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < steps) stage(s, s * P::SG); + cp_commit(); + } + for (int k = 0; k < steps; ++k) { + cp_wait(); + __syncthreads(); + if (k + STAGES - 1 < steps) stage((k + STAGES - 1) % STAGES, (k + STAGES - 1) * P::SG); + cp_commit(); + { // this stage's row group sums, once a CTA, into its tail + uint4* xsg = sm + (k % STAGES) * P::SU; + float* sums = reinterpret_cast(xsg + P::XU + P::WU); + for (int q = tid; q < P::BM * P::SG; q += P::THREADS) { + const int r = q / P::SG, gi = q - r * P::SG; + if (r < cnt && k * P::SG + gi < KG) { + uint4 v[4]; +#pragma unroll + for (int c = 0; c < 4; ++c) v[c] = xsg[P::xslot(r, gi * (GS / 8) + c)]; + sums[q] = row_gsum(v); + } + } + __syncthreads(); + } + if (!active) continue; + const uint4* xs = sm + (k % STAGES) * P::SU; + const float* sums = reinterpret_cast(xs + P::XU + P::WU); +#pragma unroll + for (int gi = 0; gi < P::SG; ++gi) { + if (k * P::SG + gi >= KG) break; + const uint4* ws = xs + P::XU + (wn * P::SG + gi) * P::WB; + uint4 xa[RT], xb[RT]; + float xsa[RT], xsb[RT]; +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + const int r0 = base + 16 * r + gq; + xa[r] = xs[P::xslot(r0, gi * (GS / 8) + t)]; + xb[r] = xs[P::xslot(r0 + 8, gi * (GS / 8) + t)]; + xsa[r] = sums[r0 * P::SG + gi]; + xsb[r] = sums[(r0 + 8) * P::SG + gi]; + } + float p[M][NTW][RT][4]; + uint4 sp4[M], bp4[M]; +#pragma unroll + for (int m = 0; m < M; ++m) { // every block's group products first (independent chains) + const uint4 wv = ws[m * G::BLOCK + lane]; + sp4[m] = ws[m * G::BLOCK + 32 * G::WV + t * G::SBV]; + bp4[m] = ws[m * G::BLOCK + 32 * G::WV + t * G::SBV + 1]; +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const uint32_t word = comp(wv, j); + const uint32_t q0 = nib2(word, 0), q1 = nib2(word, 4), q2 = nib2(word, 8), q3 = nib2(word, 12); +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + mma_z(p[m][j][r], comp(xa[r], 0), comp(xb[r], 0), comp(xa[r], 1), comp(xb[r], 1), q0, q1); + mma(p[m][j][r], comp(xa[r], 2), comp(xb[r], 2), comp(xa[r], 3), comp(xb[r], 3), q2, q3); + } + } + } +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const uint32_t sp = comp(sp4[m], j), bp = comp(bp4[m], j); + const float s0 = lo_f(sp), s1 = hi_f(sp), b0 = lo_f(bp), b1 = hi_f(bp); +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + float(&a)[4] = acc[m][r][j]; + a[0] = fmaf(xsa[r], b0, fmaf(p[m][j][r][0], s0, a[0])); + a[1] = fmaf(xsa[r], b1, fmaf(p[m][j][r][1], s1, a[1])); + a[2] = fmaf(xsb[r], b0, fmaf(p[m][j][r][2], s0, a[2])); + a[3] = fmaf(xsb[r], b1, fmaf(p[m][j][r][3], s1, a[3])); + } + } + } + } + if (!active) return; +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + const int m0 = base + 16 * r + gq, m1 = m0 + 8; + const bool v0 = m0 < cnt, v1 = m1 < cnt; + const int p0 = v0 ? __ldg(members + first + m0) : 0, p1 = v1 ? __ldg(members + first + m1) : 0; + epilogue(acc, r, out, N, cb * COLS + 2 * t, p0, p1, v0, v1, limit); + } +} + +} // namespace tf_fact_experts + +// gate and up with SwiGLU (bf16 out) and down to fp32, 64-pair tiles on 8 warps, 3 stages (2 CTAs an SM). +template __global__ void tf_fact_experts::fact_kernel<2, 2, 1, 4, 2, 3>(const __nv_bfloat16*, int, int, const uint4*, int, + int, const int*, const int*, const int*, void*, int, float); +template __global__ void tf_fact_experts::fact_kernel<1, 0, 2, 2, 4, 3>(const __nv_bfloat16*, int, int, const uint4*, int, + int, const int*, const int*, const int*, void*, int, float); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 80b9940993..def6b48614 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -22,6 +22,7 @@ pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").bytes else &.{}; pub const flashnext_fact: []const u8 = if (available) &Blob("fatbin_flashnext_fact").bytes else &.{}; pub const flashnext_hcf: []const u8 = if (available) &Blob("fatbin_flashnext_hcf").bytes else &.{}; +pub const flashnext_fexp: []const u8 = if (available) &Blob("fatbin_flashnext_fexp").bytes else &.{}; pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index d24768e3c8..2a4375d9b9 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -411,6 +411,11 @@ pub fn Pass(comptime Tri: type) type { } } + fn fexp(self: *Self, f: cuda.Function, wn: usize, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, plan: cuda.grouped.Plan, out: u64, n: usize, items: usize) !void { + var a = cuda.experts.Affine4Experts.argsOf(x, x_stride, slot_n, w, kg, nb, plan, out, n); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(items * cdiv(nb, wn)) }, .block = .{ .x = 256 }, .shared = parts.fexp_smem }, self.stream, &a); + } + /// Router, each row's ten experts and the shared one, the grouped expert launches: y [rows, 11, dims] fp32. pub fn moeBlock(self: *Self, lw: *const Layer) !void { const s = self.s; @@ -441,6 +446,14 @@ pub fn Pass(comptime Tri: type) type { } return; } + if (n >= 64) { + // prompt rows on 64-pair factored tiles, expert_rows' bits + try self.k.router.route(self.stream, s.picks, pairs, e, parts.fexp_tile, plan); + const items64 = cdiv(pairs, parts.fexp_tile) + e; + try self.fexp(self.k.fexp_up, 2, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items64); + try self.fexp(self.k.fexp_down, 4, s.moe_act, moe_w, 0, lw.experts_down.ptr, moe_w / 32, dims / 32, plan, s.y, dims, items64); + return; + } try self.k.router.route(self.stream, s.picks, pairs, e, rows_tile, plan); const items = cdiv(pairs, rows_tile) + e; try self.k.moe.run(self.stream, true, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index a2add204bf..83832d0a33 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -59,6 +59,11 @@ const cubin_hcf_down: [:0]const u8 = "_ZN3hcf10hcf_kernelILi32ELi192ELi2ELi4ELi4 const cubin_hcf_up: [:0]const u8 = "_ZN3hcf10hcf_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELb0ELb1ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii"; pub const hcf_down_smem: u32 = 4 * (2 * 32 * 64 + 2 * 192 * 16 + 2 * 2 * 192 * 2 + 2 * 32 * 4 + 2 * 32 * 4) + 32 * 16; pub const hcf_up_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4 + 2 * 32 * 4) + 64 * 16; +/// Prompt rows' experts on factored 64-pair tiles (flashnext_fexp.cu), expert_rows' bits: up with SwiGLU, down to fp32. +const cubin_fexp_up: [:0]const u8 = "_ZN15tf_fact_experts11fact_kernelILi2ELi2ELi1ELi4ELi2ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +const cubin_fexp_down: [:0]const u8 = "_ZN15tf_fact_experts11fact_kernelILi1ELi0ELi2ELi2ELi4ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +pub const fexp_smem: u32 = 41472; +pub const fexp_tile: usize = 64; pub const fact_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4); /// Pairs an expert item takes: the core rows kernel's tile. const rows_tile: usize = cuda.experts.RowsExperts.tile; @@ -111,6 +116,9 @@ pub const Kernels = struct { fact_mod: cuda.Module, fact: cuda.Function, hcf_mod: cuda.Module, + fexp_mod: cuda.Module, + fexp_up: cuda.Function, + fexp_down: cuda.Function, hcf_down: cuda.Function, hcf_up: cuda.Function, native_lane: cuda.qmmf.Lane, @@ -141,6 +149,12 @@ pub const Kernels = struct { errdefer k.fact_mod.unload(); k.fact = try k.fact_mod.function(cubin_fact); try k.fact.allowDynamicShared(fact_smem); + k.fexp_mod = try cuda.Module.load(d, cuda.kernels.flashnext_fexp); + errdefer k.fexp_mod.unload(); + k.fexp_up = try k.fexp_mod.function(cubin_fexp_up); + try k.fexp_up.allowDynamicShared(fexp_smem); + k.fexp_down = try k.fexp_mod.function(cubin_fexp_down); + try k.fexp_down.allowDynamicShared(fexp_smem); k.hcf_mod = try cuda.Module.load(d, cuda.kernels.flashnext_hcf); errdefer k.hcf_mod.unload(); k.hcf_down = try k.hcf_mod.function(cubin_hcf_down); @@ -174,6 +188,7 @@ pub const Kernels = struct { k.rows.unload(); k.experts.unload(); k.hcf_mod.unload(); + k.fexp_mod.unload(); k.fact_mod.unload(); k.qmm.unload(); } From 2d5c2d1a1764ff52eaceda485c987ec55ea0f460 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 13:20:33 +0200 Subject: [PATCH 06/16] Flash Next CUDA: prompt chain on two threads a state row and indexer scores 8 rows a block, both in #952's sum order (its bits) --- zig/kernels/cuda/flashnext_attn.cu | 5 + zig/kernels/cuda/flashnext_batch.cu | 9 ++ zig/kernels/cuda/flashnext_gdn_tpr.cuh | 114 ++++++++++++++++++ zig/kernels/cuda/flashnext_qsa_tile.cuh | 79 ++++++++++++ zig/src/families/flashnext/cuda_batch.zig | 7 ++ .../families/flashnext/cuda_batch_state.zig | 5 + zig/src/families/flashnext/cuda_glue.zig | 7 ++ 7 files changed, 226 insertions(+) create mode 100644 zig/kernels/cuda/flashnext_gdn_tpr.cuh create mode 100644 zig/kernels/cuda/flashnext_qsa_tile.cuh diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu index 6386dc7767..93b7e861b0 100644 --- a/zig/kernels/cuda/flashnext_attn.cu +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -487,3 +487,8 @@ extern "C" __global__ void __launch_bounds__(256) fn_attn_merge(const float* __r OUT[(static_cast(r) * H + head) * HD + d] = __float2bfloat16_rn(__fdiv_rn(o, l)); } } + +#include "flashnext_qsa_tile.cuh" + +// The indexer scores 8 rows a block, every head, the sums in fn_qsa_scores' order (its bits). +template __global__ void qsa_tile::scores<8, 4>(const __nv_bfloat16*, const __nv_bfloat16*, const int*, float*, int, int, int, int); diff --git a/zig/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu index c4ad0e2691..c31f32e44f 100644 --- a/zig/kernels/cuda/flashnext_batch.cu +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -314,3 +314,12 @@ extern "C" __global__ void __launch_bounds__(1024) fn_qsa_select(const float* __ spr[r] = 1; } } + +#include "flashnext_gdn_tpr.cuh" + +// The prompt chain on two threads a state row (gdn_tpr): fn_gdn_rows_to's bits, state_in == state_out in place. +extern "C" __global__ void __launch_bounds__(256) fn_gdn_chain_tpr(const float* qn, const float* kn, const float* v, + const float* gb, const float* state_in, float* state_out, + int rows, float* ys_out) { + gdn_tpr::chain<2, 16>(qn, kn, v, gb, state_in, state_out, rows, ys_out); +} diff --git a/zig/kernels/cuda/flashnext_gdn_tpr.cuh b/zig/kernels/cuda/flashnext_gdn_tpr.cuh new file mode 100644 index 0000000000..013359c7d7 --- /dev/null +++ b/zig/kernels/cuda/flashnext_gdn_tpr.cuh @@ -0,0 +1,114 @@ +// fn_gdn_rows' chain, a block a value head, TPR threads a state row, warp_sum's tree in the thread (same bits). +#ifndef KREG +#define KREG 0 +#endif +namespace gdn_tpr { +constexpr int NK = 16, NV = 48, DK = 128, DV = 128; +__device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } +__device__ __forceinline__ void cp16(void* dst, const void* src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(static_cast(__cvta_generic_to_shared(dst))), "l"(src)); +} +__device__ __forceinline__ void cp8(void* dst, const void* src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" ::"r"(static_cast(__cvta_generic_to_shared(dst))), "l"(src)); +} +template struct Stage { + static constexpr int ROW = (DK + DK + DV) * 4 + 16; + static constexpr int BYTES = STEPS * ROW; +}; +// warp_sum's tree over the 32 lane partials x[l'] (lane l = l' TPR + t): in-thread levels, then the threads' xor levels +template +__device__ __forceinline__ float tree(float (&x)[32 / TPR]) { +#pragma unroll + for (int h = 16 / TPR; h >= 1; h >>= 1) +#pragma unroll + for (int i = 0; i < h; ++i) x[i] = x[i] + x[i + h]; + float r = x[0]; +#pragma unroll + for (int m = TPR / 2; m >= 1; m >>= 1) r = r + __shfl_xor_sync(0xffffffffu, r, m); + return r; +} +template +__device__ __forceinline__ void chain(const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, const float* state_in, float* state_out, int rows, + float* __restrict__ ys_out) { + using S = Stage; + constexpr int L = 32 / TPR; // the lanes a thread holds + extern __shared__ __align__(16) unsigned char smem[]; + const int hv = blockIdx.x, hk = hv / (NV / NK), row = threadIdx.x / TPR, t = threadIdx.x % TPR; + const size_t sbase = (size_t)hv * DV * DK + (size_t)row * DK; + float s[L][4]; +#pragma unroll + for (int l = 0; l < L; ++l) +#pragma unroll + for (int i = 0; i < 4; ++i) s[l][i] = state_in[sbase + (l * TPR + t) * 4 + i]; + auto load = [&](int buf, int r0) { + unsigned char* base = smem + buf * S::BYTES; + const int n = min(STEPS, rows - r0); + for (int c = threadIdx.x; c < n * 97; c += blockDim.x) { + const int r = c / 97, u = c % 97, rr = r0 + r; + unsigned char* d = base + r * S::ROW; + if (u < 32) cp16(d + u * 16, kn + ((size_t)rr * NK + hk) * DK + u * 4); + else if (u < 64) cp16(d + DK * 4 + (u - 32) * 16, qn + ((size_t)rr * NK + hk) * DK + (u - 32) * 4); + else if (u < 96) cp16(d + 2 * DK * 4 + (u - 64) * 16, v + (size_t)rr * NV * DV + hv * DV + (u - 64) * 4); + else cp8(d + (2 * DK + DV) * 4, gb + ((size_t)rr * NV + hv) * 2); + } + asm volatile("cp.async.commit_group;\n" ::); + }; + const int batches = (rows + STEPS - 1) / STEPS; + if (batches > 0) load(0, 0); + for (int bi = 0; bi < batches; ++bi) { + if (bi + 1 < batches) { + load((bi + 1) & 1, (bi + 1) * STEPS); + asm volatile("cp.async.wait_group 1;\n" ::); + } else { + asm volatile("cp.async.wait_group 0;\n" ::); + } + __syncthreads(); + const unsigned char* base = smem + (bi & 1) * S::BYTES; + const int r0 = bi * STEPS, n = min(STEPS, rows - r0); + for (int rr = 0; rr < n; ++rr) { + const unsigned char* d = base + rr * S::ROW; + const float vv = reinterpret_cast(d + 2 * DK * 4)[row]; + const float2 gbv = *reinterpret_cast(d + (2 * DK + DV) * 4); + const float g = gbv.x, beta = gbv.y; + float x[L]; + float4 kr[KREG ? L : 1]; +#pragma unroll + for (int l = 0; l < L; ++l) { + const float4 k4 = *reinterpret_cast(d + (l * TPR + t) * 16); + if (KREG) kr[KREG ? l : 0] = k4; + const float kk[4] = {k4.x, k4.y, k4.z, k4.w}; + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[l][i] = s[l][i] * g; + kv = kv + s[l][i] * kk[i]; + } + x[l] = kv; + } + const float kv = tree(x); + const float delta = (vv - kv) * beta; +#pragma unroll + for (int l = 0; l < L; ++l) { + const float4 k4 = KREG ? kr[KREG ? l : 0] : *reinterpret_cast(d + (l * TPR + t) * 16); + const float4 q4 = *reinterpret_cast(d + DK * 4 + (l * TPR + t) * 16); + const float kk[4] = {k4.x, k4.y, k4.z, k4.w}, qq[4] = {q4.x, q4.y, q4.z, q4.w}; + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[l][i] = s[l][i] + kk[i] * delta; + o = o + s[l][i] * qq[i]; + } + x[l] = o; + } + const float o = bf(tree(x)); + if (t == 0) ys_out[((size_t)(r0 + rr) * NV + hv) * DV + row] = o; + } + __syncthreads(); + } +#pragma unroll + for (int l = 0; l < L; ++l) +#pragma unroll + for (int i = 0; i < 4; ++i) state_out[sbase + (l * TPR + t) * 4 + i] = s[l][i]; +} +} // namespace gdn_tpr diff --git a/zig/kernels/cuda/flashnext_qsa_tile.cuh b/zig/kernels/cuda/flashnext_qsa_tile.cuh new file mode 100644 index 0000000000..091a7bb92b --- /dev/null +++ b/zig/kernels/cuda/flashnext_qsa_tile.cuh @@ -0,0 +1,79 @@ +// fn_qsa_scores' sums without shuffles: warp_sum's tree in the thread, leaves bit-reversed, R rows a block (same bits). +namespace qsa_tile { +constexpr int DI = 128; +__device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } +__device__ __forceinline__ int rev5(int c) { return __brev(c) >> 27; } +// a binary counter of partial sums: level k holds a closed subtree of 2^k leaves until its sibling closes +struct Stack { + float v[5]; + __device__ __forceinline__ void push(int c, float x) { +#pragma unroll + for (int k = 0; k < 5; ++k) { + if (((c >> k) & 1) == 0) { v[k] = x; return; } + x = v[k] + x; + } + v[0] = x; // c == 31: the root, kept in v[0] + } +}; +template +__global__ void __launch_bounds__(256) scores(const __nv_bfloat16* __restrict__ IQ, const __nv_bfloat16* __restrict__ POOLED, + const int* __restrict__ POS0, float* __restrict__ SC, int NB, int RATIO, int TOP, + int rows) { + __shared__ float4 qs[R][HI][DI / 4]; + const int r0 = blockIdx.x * R, b = blockIdx.y * 256 + threadIdx.x; + const int pos0 = POS0[0]; + const int last = min(r0 + R, rows) - 1; + const int cmax = (pos0 + last + 1) / RATIO; + if (cmax <= TOP || blockIdx.y * 256 >= cmax) return; + for (int e = threadIdx.x; e < R * HI * DI / 4; e += blockDim.x) { + const int ri = e / (HI * DI / 4), rem = e % (HI * DI / 4), r = min(r0 + ri, rows - 1); + const __nv_bfloat162* src = reinterpret_cast(IQ + (static_cast(r) * HI * DI + rem * 4)); + const float2 a = __bfloat1622float2(src[0]), c = __bfloat1622float2(src[1]); + reinterpret_cast(qs)[e] = make_float4(a.x, a.y, c.x, c.y); + } + __syncthreads(); + if (b >= cmax) return; + Stack st[R][HI]; + const __nv_bfloat16* kb = POOLED + static_cast(b) * DI; +#pragma unroll + for (int half = 0; half < 2; ++half) { + float k[16][4]; // the 16 lanes l (bit 0 == half) a half's leaves read +#pragma unroll + for (int c = 0; c < 16; ++c) { + const int l = rev5(c + 16 * half); + const uint2 raw = *reinterpret_cast(kb + 4 * l); + const __nv_bfloat16* p = reinterpret_cast(&raw); +#pragma unroll + for (int e = 0; e < 4; ++e) k[c][e] = __bfloat162float(p[e]); + } +#pragma unroll + for (int c = 0; c < 16; ++c) { + const int cc = c + 16 * half, l = rev5(cc); +#pragma unroll + for (int ri = 0; ri < R; ++ri) +#pragma unroll + for (int h = 0; h < HI; ++h) { + const float4 q4 = qs[ri][h][l]; + float part = 0.0f; + part = __fmaf_rn(k[c][0], q4.x, part); + part = __fmaf_rn(k[c][1], q4.y, part); + part = __fmaf_rn(k[c][2], q4.z, part); + part = __fmaf_rn(k[c][3], q4.w, part); + st[ri][h].push(cc, part); + } + } + } + const float scale = __fsqrt_rn(static_cast(DI)); +#pragma unroll + for (int ri = 0; ri < R; ++ri) { + const int r = r0 + ri; + if (r >= rows) break; + const int complete = (pos0 + r + 1) / RATIO; + if (complete <= TOP || b >= complete) continue; + float total = 0.0f; +#pragma unroll + for (int h = 0; h < HI; ++h) total = __fadd_rn(total, fmaxf(st[ri][h].v[0], 0.0f)); + SC[static_cast(r) * NB + b] = __fdiv_rn(total, scale); + } +} +} // namespace qsa_tile diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 2a4375d9b9..c0e043c3ec 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -331,6 +331,13 @@ pub fn Pass(comptime Tri: type) type { a.add(gv); a.add(ggb); a.add(state); + if (rows >= 64) { + // two threads a state row (fn_gdn_chain_tpr), fn_gdn_rows' bits; in place when there is no rollback copy + a.add(if (state_out != 0) state_out else state); + a.add(cint(rows)); + a.add(ys); + return cuda.launch.launch(self.k.chain_tpr, .{ .grid = .{ .x = gdn.nv }, .block = .{ .x = 256 }, .shared = parts.chain_tpr_smem }, self.stream, &a); + } if (state_out != 0) a.add(state_out); a.add(cint(rows)); a.add(ys); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 83832d0a33..02f2796303 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -64,6 +64,8 @@ const cubin_fexp_up: [:0]const u8 = "_ZN15tf_fact_experts11fact_kernelILi2ELi2EL const cubin_fexp_down: [:0]const u8 = "_ZN15tf_fact_experts11fact_kernelILi1ELi0ELi2ELi2ELi4ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; pub const fexp_smem: u32 = 41472; pub const fexp_tile: usize = 64; +/// fn_gdn_chain_tpr's staged rows: 2 x 16 steps of q, k, v (128 fp32 each) and g, beta. +pub const chain_tpr_smem: u32 = 2 * 16 * (3 * 128 * 4 + 16); pub const fact_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4); /// Pairs an expert item takes: the core rows kernel's tile. const rows_tile: usize = cuda.experts.RowsExperts.tile; @@ -126,6 +128,7 @@ pub const Kernels = struct { front: cuda.Function, chain: cuda.Function, chain_to: cuda.Function, + chain_tpr: cuda.Function, gnorm: cuda.Function, sum: cuda.Function, shift: cuda.Function, @@ -177,6 +180,8 @@ pub const Kernels = struct { k.front = try k.batch.function("fn_gdn_front"); k.chain = try k.batch.function("fn_gdn_rows"); k.chain_to = try k.batch.function("fn_gdn_rows_to"); + k.chain_tpr = try k.batch.function("fn_gdn_chain_tpr"); + try k.chain_tpr.allowDynamicShared(chain_tpr_smem); k.gnorm = try k.batch.function("fn_gdn_norm"); return k; } diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index a5793011d1..214fb10369 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -49,6 +49,7 @@ const Fns = struct { gate: cuda.Function, pool: cuda.Function, scores: cuda.Function, + scores_tile: cuda.Function, chunks: cuda.Function, chunks_ahead: cuda.Function, chunks_one: cuda.Function, @@ -123,6 +124,7 @@ pub const Glue = struct { .gate = try g.attn.function("fn_attn_gate"), .pool = try g.attn.function("fn_qsa_pool"), .scores = try g.attn.function("fn_qsa_scores"), + .scores_tile = try g.attn.function("_ZN8qsa_tile6scoresILi8ELi4EEEvPK13__nv_bfloat16S3_PKiPfiiii"), .chunks = try g.attn.function("fn_attn_chunks"), .chunks_ahead = try g.attn.function("fn_attn_chunks_ahead"), .chunks_one = try g.attn.function("fn_attn_chunks_one"), @@ -242,6 +244,11 @@ pub const Glue = struct { if (wt.index_dim != 128 or wt.index_heads > 8) return error.RowShape; var a: cuda.Args = .{}; ptrs(&a, &.{ iq, pooled, pos, sc }); + if (wt.index_heads == 4) { + // 8 rows and 256 blocks a CTA, every head, fn_qsa_scores' sums in its order (its bits) + ints(&a, &.{ nb, wt.index_ratio, wt.index_top, rows }); + return go(g.f.scores_tile, s, .{ cdiv(rows, 8), cdiv(blocks, 256), 1 }, 256, 0, &a); + } ints(&a, &.{ nb, wt.index_heads, wt.index_ratio, wt.index_top }); try go(g.f.scores, s, .{ rows, cdiv(blocks, 64), 1 }, 256, 0, &a); } From 1d07399417cfd5af5331a346c1a1b275d15116e9 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:27:27 +0200 Subject: [PATCH 07/16] Flash Next CUDA: prompt rows' K slices folded in the matmul's CTA, their sum in slice order (the same bits) --- zig/kernels/cuda/flashnext_fact.cu | 3 ++ zig/src/families/flashnext/cuda_batch.zig | 42 ++++++++++++------- .../families/flashnext/cuda_batch_state.zig | 4 ++ zig/src/families/flashnext/cuda_mtp_ops.zig | 12 +++--- 4 files changed, 40 insertions(+), 21 deletions(-) diff --git a/zig/kernels/cuda/flashnext_fact.cu b/zig/kernels/cuda/flashnext_fact.cu index a29f90fb52..2fc25e1f19 100644 --- a/zig/kernels/cuda/flashnext_fact.cu +++ b/zig/kernels/cuda/flashnext_fact.cu @@ -258,3 +258,6 @@ __global__ void __launch_bounds__(WM * WN * 32, MINB) fact_kernel( // The prompt form the batched pass launches from 64 rows: 64 x 128 on 2 x 4 warps, bf16 out (one slice) or fp32 slices. template __global__ void fact::fact_kernel<64, 128, 2, 4, 3, 2, false, 2, false, true, false>(const __nv_bfloat16*, const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, float*, int, int, int, int, int, int, int); +// The folded form from 64 rows with K slices: every slice in the CTA, their totals added in slice order (one fp32 slice). +template __global__ void fact::fact_kernel<64, 128, 2, 4, 3, 2, false, 2, true, true, true>(const __nv_bfloat16*, + const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, float*, int, int, int, int, int, int, int); diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index c0e043c3ec..fe012d64b8 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -104,19 +104,25 @@ pub fn Pass(comptime Tri: type) type { pos0: usize = 0, /// The MoE write-back the next layer (or the head) applies first, and the layer it came from. pending: ?usize = null, + /// K slices the last out projection wrote to part (1 when the prompt tile folded them). + out_slices: usize = out_sk, prepared: bool = false, // ---- launches ---------------------------------------------------------------------------------------- /// qmm.cu's lane tile over m rows: out bf16 [m, n] (sk 1) or the unreduced fp32 slices [sk, m, n] in part. - pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !void { + /// Returns the K slices written to part: 1 when the prompt tile folded them (their sum in slice order). + pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !usize { if (w.native) |native| { - if (sk == 1) return self.k.native_lane.matmul(self.stream, x, w.k, m, native, out, self.s.native_part); + if (sk == 1) { + try self.k.native_lane.matmul(self.stream, x, w.k, m, native, out, self.s.native_part); + return 1; + } // Legacy callers consume SK slices; the native lane has already reduced its own shape-driven slices. try self.k.native_lane.matmulF32(self.stream, x, w.k, m, native, part, self.s.native_part); const bytes = m * w.n * 4; if (sk > 1) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes * (sk - 1), self.stream.handle), "cuMemsetD8Async"); - return; + return sk; } // from 64 rows the prompt tile (flashnext_fact.cu), bit for bit the 16-row lane tile const big = m >= 64; @@ -136,10 +142,13 @@ pub fn Pass(comptime Tri: type) type { a.add(out); a.add(part); for ([_]usize{ m, w.n, k, sk, npad, k, sweep }) |v| a.add(cint(v)); - const f = if (big) self.k.fact else self.k.lane; + const fold = big and sk > 1; + const f = if (fold) self.k.fact_fold else if (big) self.k.fact else self.k.lane; const block: u32 = if (big) 256 else 128; const shared: u32 = if (big) parts.fact_smem else qmm_smem; - try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(sk) }, .block = .{ .x = block }, .shared = shared }, self.stream, &a); + const z: usize = if (fold) 1 else sk; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(z) }, .block = .{ .x = block }, .shared = shared }, self.stream, &a); + return z; } pub fn sumSlices(self: *Self, part: u64, out: u64, bf16: bool, n: usize, sk: usize) !void { @@ -257,9 +266,9 @@ pub fn Pass(comptime Tri: type) type { const n = self.rows; // the rows' n-gram embeddings (host gathers from the mapped table) are already uploaded by `stage` _ = ids; - try self.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, n, 1); - try self.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, n, 4); - try self.sumSlices(s.ple_part, s.ple_vals, true, n * dims, 4); + _ = try self.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, n, 1); + const vsk = try self.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, n, 4); + try self.sumSlices(s.ple_part, s.ple_vals, true, n * dims, vsk); try self.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, n, dims, streams); if (self.lanes.len == 0) { try self.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, self.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, n, dims, streams, 4, 3); @@ -277,7 +286,7 @@ pub fn Pass(comptime Tri: type) type { fn gdnBlock(self: *Self, lw: *const Layer, li: usize, keep: LayerKeep) !void { const s = self.s; const n = self.rows; - try self.lane(s.mixed, s.xsm, lw.proj, s.proj, s.dummy, n, 1); + _ = try self.lane(s.mixed, s.xsm, lw.proj, s.proj, s.dummy, n, 1); try self.copy(keep.proj, s.proj, n * gdn.proj_width * 2); try self.copy(keep.hist, self.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2); // every row's conv, norms and gates, the recurrence, then the gated norm: the one-block chain's bits @@ -305,7 +314,7 @@ pub fn Pass(comptime Tri: type) type { } else for (self.lanes, 0..) |l, r| { try self.shiftWindow(l.seq.conv[li].ptr, s.proj + r * gdn.proj_width * 2, gdn.proj_width, 1, gdn.taps - 1, gdn.conv_dim); } - try self.lane(s.gout, s.gxs, lw.out, s.dummy, s.part, n, out_sk); + self.out_slices = try self.lane(s.gout, s.gxs, lw.out, s.dummy, s.part, n, out_sk); } /// fn_gdn_front over `rows` projection rows after the conv window `conv`. @@ -348,7 +357,7 @@ pub fn Pass(comptime Tri: type) type { const s = self.s; const n = self.rows; const pa = s.proj; - try self.lane(s.mixed, s.xsm, lw.proj, pa, s.dummy, n, 1); + _ = try self.lane(s.mixed, s.xsm, lw.proj, pa, s.dummy, n, 1); if (self.lanes.len > 0) { // a round: each row's rope, cache writes, selection and attention on its own stream's caches const own_seq = self.seq; @@ -364,7 +373,7 @@ pub fn Pass(comptime Tri: type) type { } } else try self.attnRows(lw, ai, pa, n, .{ .pos = s.pos, .q = s.q, .iq = s.iq, .ao = s.ao }); try self.tri.attnGate(s.ao, pa, s.gated, s.gxs_attn, n, proj_n, q_heads, head_dim); - try self.lane(s.gated, s.gxs_attn, lw.out, s.dummy, s.part, n, out_sk); + self.out_slices = try self.lane(s.gated, s.gxs_attn, lw.out, s.dummy, s.part, n, out_sk); } /// Where attnRows reads its rows' positions (16 bytes apart) and writes queries, indexer queries and outputs. @@ -569,7 +578,7 @@ pub fn Pass(comptime Tri: type) type { } else { try self.attnBlock(lw, wt.attnIndex(layer)); } - try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, out_sk); + try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, self.out_slices); try self.readout(lw.mlp_hc, 0, n); try self.moeBlock(lw); self.pending = layer; @@ -585,10 +594,13 @@ pub fn Pass(comptime Tri: type) type { /// Column tiles keep their own reduction order; each gathered row restores global vocabulary order. fn projectHead(self: *Self, out: u64, n: usize) !void { const s = self.s; - if (self.m.head.n == vocab) return self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); + if (self.m.head.n == vocab) { + _ = try self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); + return; + } if (self.m.head.n != vocab / 2 or n > s.cap) return error.UnexpectedTensor; const c = self.m.collective orelse return error.TensorParallelUnavailable; - try self.lane(s.mixed, s.xsm, self.m.head, s.head_local, s.dummy, n, 1); + _ = try self.lane(s.mixed, s.xsm, self.m.head, s.head_local, s.dummy, n, 1); for (0..n) |r| try c.gather(self.stream, s.head_local + r * vocab, out + r * vocab * 2, vocab / 2); } diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 02f2796303..c44eb405a0 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -66,6 +66,7 @@ pub const fexp_smem: u32 = 41472; pub const fexp_tile: usize = 64; /// fn_gdn_chain_tpr's staged rows: 2 x 16 steps of q, k, v (128 fp32 each) and g, beta. pub const chain_tpr_smem: u32 = 2 * 16 * (3 * 128 * 4 + 16); +const cubin_fact_fold: [:0]const u8 = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb1ELb1ELb1EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; pub const fact_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4); /// Pairs an expert item takes: the core rows kernel's tile. const rows_tile: usize = cuda.experts.RowsExperts.tile; @@ -117,6 +118,7 @@ pub const Kernels = struct { lane: cuda.Function, fact_mod: cuda.Module, fact: cuda.Function, + fact_fold: cuda.Function, hcf_mod: cuda.Module, fexp_mod: cuda.Module, fexp_up: cuda.Function, @@ -152,6 +154,8 @@ pub const Kernels = struct { errdefer k.fact_mod.unload(); k.fact = try k.fact_mod.function(cubin_fact); try k.fact.allowDynamicShared(fact_smem); + k.fact_fold = try k.fact_mod.function(cubin_fact_fold); + try k.fact_fold.allowDynamicShared(fact_smem); k.fexp_mod = try cuda.Module.load(d, cuda.kernels.flashnext_fexp); errdefer k.fexp_mod.unload(); k.fexp_up = try k.fexp_mod.function(cubin_fexp_up); diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index cfcd9fdd05..b5f7a24dd4 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -142,18 +142,18 @@ pub fn Ops(comptime Tri: type) type { } try @import("cuda_media.zig").injectNext(d.media, o.stream(), s.ple_keys, d.len, m); try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, m, dims, dims); - try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, m, fc_sk); - try p.sumSlices(s.part, s.branch, true, m * dims, fc_sk); + const esk = try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, m, fc_sk); + try p.sumSlices(s.part, s.branch, true, m * dims, esk); // the hidden branch: the four streams normed as one row, fc_hidden on each stream try tri.rmsnorm(rows + at * wide * 2, h.norm_h.ptr, s.norm, s.gq, eps, wide, m, wide, wide); - try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams * m, fc_sk); - try p.sumSlices(s.part, s.ple_gated, true, streams * m * dims, fc_sk); + const hsk = try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams * m, fc_sk); + try p.sumSlices(s.part, s.ple_gated, true, streams * m * dims, hsk); try tri.addStreams(s.branch, s.ple_gated, s.h, m, dims, streams); // the layer: attention over the head's own cache, its experts try p.writeback(0, m); try p.readout(h.layer.attn_hc, 0, m); try p.attnBlock(&h.layer, 0); - try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, out_sk); + try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, p.out_slices); try p.readout(h.layer.mlp_hc, 0, m); try p.moeBlock(&h.layer); p.pending = 0; @@ -168,7 +168,7 @@ pub fn Ops(comptime Tri: type) type { try p.copy(s.pss, s.pss + (m - 1) * pss_n * 4, pss_n * 4); } try p.readout(h.mixer, 0, 1); - try p.lane(s.mixed, s.xsm, h.draft, d.dlogits.ptr, s.dummy, 1, 1); + _ = try p.lane(s.mixed, s.xsm, h.draft, d.dlogits.ptr, s.dummy, 1, 1); } } } From c6ed446f8e36361c9a677e65a311ac8a82b61229 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:35:52 +0200 Subject: [PATCH 08/16] CUDA router logits: 128 rows and 64 experts a block from 64 rows (same sums) --- zig/kernels/cuda/route_rows.cu | 7 ++++--- zig/src/cuda/route_rows.zig | 8 ++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/zig/kernels/cuda/route_rows.cu b/zig/kernels/cuda/route_rows.cu index 881a1ffde7..f29d7038cc 100644 --- a/zig/kernels/cuda/route_rows.cu +++ b/zig/kernels/cuda/route_rows.cu @@ -117,11 +117,12 @@ extern "C" __global__ void __launch_bounds__(128) tf_route_logits(const __nv_bfl route_logits<4, 2, 128>(X, x_ld, W, OUT, M, D, E); } -// route_logits over 64 rows a block (from 64 rows: the router rows are read once per 64 input rows), the same sums. -extern "C" __global__ void __launch_bounds__(128) tf_route_logits_wide(const __nv_bfloat16* __restrict__ X, int x_ld, +// route_logits over 128 rows and 64 experts a block, eight warps (from 64 rows: each router row read once per 128 +// input rows), the same sums. +extern "C" __global__ void __launch_bounds__(256) tf_route_logits_wide(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, int M, int D, int E) { - route_logits<4, 2, 128, 4>(X, x_ld, W, OUT, M, D, E); + route_logits<8, 2, 64, 8>(X, x_ld, W, OUT, M, D, E); } // route_logits with three stages of 256 in flight (up to 16 rows, where the router rows' reads set its time). diff --git a/zig/src/cuda/route_rows.zig b/zig/src/cuda/route_rows.zig index 34ebc4db09..b1fc6ebe89 100644 --- a/zig/src/cuda/route_rows.zig +++ b/zig/src/cuda/route_rows.zig @@ -17,8 +17,8 @@ pub const most_picks = 128; /// tf_route_logits' staged rows (16 inputs and 32 experts, padded by 8): two stages of 128, or three of 256. const many_shared: u32 = 2 * 48 * (128 + 8) * 2; const few_shared: u32 = 3 * 48 * (256 + 8) * 2; -/// tf_route_logits_wide's: 64 inputs and 32 experts, two stages of 128. -const wide_shared: u32 = 2 * 96 * (128 + 8) * 2; +/// tf_route_logits_wide's: 128 inputs and 64 experts, two stages of 64. +const wide_shared: u32 = 2 * 192 * (64 + 8) * 2; pub const RouteRows = struct { mod: Module, @@ -54,8 +54,8 @@ pub const RouteRows = struct { // up to 16 rows, deeper staging (the same sums: the router rows' reads set the time) const few = rows <= 16; if (rows >= 64) { - // 64 rows a block: each router row read once per 64 input rows (the same sums) - try launch_.launch(g.wide_fn, .{ .grid = .{ .x = @intCast((rows + 63) / 64), .y = @intCast((e + 31) / 32), .z = 1 }, .block = .{ .x = 128 }, .shared = wide_shared }, s, &a); + // 128 rows and 64 experts a block: each router row read once per 128 input rows (the same sums) + try launch_.launch(g.wide_fn, .{ .grid = .{ .x = @intCast((rows + 127) / 128), .y = @intCast((e + 63) / 64), .z = 1 }, .block = .{ .x = 256 }, .shared = wide_shared }, s, &a); return; } try launch_.launch(if (few) g.few_fn else g.logits_fn, .{ .grid = .{ .x = @intCast((rows + 15) / 16), .y = @intCast((e + 31) / 32), .z = 1 }, .block = .{ .x = 128 }, .shared = if (few) few_shared else many_shared }, s, &a); From 1d4a43c933d6470e44c2a6b398ef452c75182855 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:52:12 +0200 Subject: [PATCH 09/16] Flash Next CUDA: TF_FLASHNEXT_CHUNK takes 1 to 16384 rows and refuses anything else at load --- zig/src/families/flashnext/cuda_batch.zig | 1 + .../families/flashnext/cuda_batch_state.zig | 2 ++ .../families/flashnext/cuda_lanes_load.zig | 3 +-- zig/src/families/flashnext/cuda_prompt.zig | 25 ++++++++++++++++--- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index fe012d64b8..48f12e3cbf 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -46,6 +46,7 @@ pub const proj_row = @max(gdn.proj_width, proj_n); pub const att_rows = parts.att_rows; /// Rows a pass holds at most unless asked for more. pub const default_chunk = parts.default_chunk; +pub const max_chunk = parts.max_chunk; const ratio = wt.index_ratio; const top = wt.index_top; const idw = wt.index_idw; diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index c44eb405a0..eccd46bbd8 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -47,6 +47,8 @@ pub const att_rows: usize = 256; pub const fused_rows: usize = 1024; /// Rows a pass holds at most unless asked for more. pub const default_chunk: usize = 1536; +/// The most rows TF_FLASHNEXT_CHUNK may ask for (the largest measured pass; its scratch is about 13 GiB more). +pub const max_chunk: usize = 16384; const ratio = wt.index_ratio; const top = wt.index_top; const idw = wt.index_idw; diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 34b84a9da0..19166bd0bc 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -29,8 +29,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (captured) |kernels| own.set = try F.comparison(gpa, io, driver, ctx, kernels, o.drafts); errdefer if (own.set) |*x| x.deinit(); if (o.tp == 2) { - const configured_chunk = F.envUsize("TF_FLASHNEXT_CHUNK", prompt.batch.default_chunk); - var id = try ranks.identity(gpa, io, dir, o.context, if (configured_chunk == 0) prompt.batch.default_chunk else configured_chunk); + var id = try ranks.identity(gpa, io, dir, o.context, try prompt.chunkRows()); id.mtp = @intFromBool(o.drafts); if (o.vision) |v| { try @import("qwen_vision").media.checkLimits(v); diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 5009a98d31..17aac6044f 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -133,7 +133,7 @@ pub const Engine = struct { batch_s: ?batch.Scratch = null, /// Block scores for the sparse rows of an attention block, grown to the longest sequence's need. scores: ?cuda.DeviceBuffer = null, - /// Rows a batched chunk takes at most (TF_FLASHNEXT_CHUNK; batch.default_chunk by default). + /// Rows a batched chunk takes at most (chunkRows). chunk: usize, round_graph: round_graph.Cache = .{}, image_factory: ?*@import("cuda_images.zig").Vision = null, @@ -204,9 +204,7 @@ pub const Engine = struct { e.batch_k = try batch.Kernels.load(driver); e.batch_s = null; e.scores = null; - e.chunk = batch.default_chunk; - if (std.c.getenv("TF_FLASHNEXT_CHUNK")) |v| e.chunk = std.fmt.parseInt(usize, std.mem.span(v), 10) catch batch.default_chunk; - if (e.chunk == 0) e.chunk = batch.default_chunk; + e.chunk = try chunkRows(); return e; } @@ -469,3 +467,22 @@ test "a failed fill frees the buffer it was filling and every one before it" { defer std.testing.allocator.free(all); try std.testing.expectEqual(@as(usize, 3), live); } + +/// TF_FLASHNEXT_CHUNK's rows a pass: unset, the default; otherwise 1 to batch.max_chunk, anything else refused. +pub fn parseChunk(v: ?[]const u8) !usize { + const s = v orelse return batch.default_chunk; + const n = std.fmt.parseInt(usize, s, 10) catch return error.BadChunk; + if (n == 0 or n > batch.max_chunk) return error.BadChunk; + return n; +} + +pub fn chunkRows() !usize { + return parseChunk(if (std.c.getenv("TF_FLASHNEXT_CHUNK")) |v| std.mem.span(v) else null); +} + +test "TF_FLASHNEXT_CHUNK takes 1 to the largest measured pass and refuses the rest" { + try std.testing.expectEqual(batch.default_chunk, try parseChunk(null)); + try std.testing.expectEqual(@as(usize, 4096), try parseChunk("4096")); + try std.testing.expectEqual(batch.max_chunk, try parseChunk("16384")); + for ([_][]const u8{ "0", "16385", "-1", "4k", "" }) |bad| try std.testing.expectError(error.BadChunk, parseChunk(bad)); +} From 371b19095bc4e2a6dec007a3ced6f22231a340cf Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:52:12 +0200 Subject: [PATCH 10/16] Flash Next CUDA: prompt tiles' launches in cuda_tiles.zig, one-line comments under 120 columns --- zig/kernels/cuda/flashnext_attn.cu | 2 +- zig/kernels/cuda/flashnext_fact.cu | 2 +- zig/kernels/cuda/flashnext_hcf.cu | 2 +- zig/kernels/cuda/route_rows.cu | 6 +- zig/src/families/flashnext/cuda_batch.zig | 91 ++------------------- zig/src/families/flashnext/cuda_launch.zig | 2 +- zig/src/families/flashnext/cuda_tiles.zig | 95 ++++++++++++++++++++++ 7 files changed, 110 insertions(+), 90 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_tiles.zig diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu index 93b7e861b0..3033fad110 100644 --- a/zig/kernels/cuda/flashnext_attn.cu +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -432,7 +432,7 @@ extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks(const __nv_bflo attn_chunks<1>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); } -// attn_chunks with one buffer and every chunk of a row merged in place (prompt rows): OUT (bf16, in PO's place), no partials. +// attn_chunks on one buffer, a row's chunks merged in place (prompt rows): OUT bf16 in PO's place, no partials. extern "C" __global__ void __launch_bounds__(128) fn_attn_fused(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, float* __restrict__ OUT, const int* __restrict__ IDS, const int* __restrict__ NKR, diff --git a/zig/kernels/cuda/flashnext_fact.cu b/zig/kernels/cuda/flashnext_fact.cu index 2fc25e1f19..bdb8b31268 100644 --- a/zig/kernels/cuda/flashnext_fact.cu +++ b/zig/kernels/cuda/flashnext_fact.cu @@ -258,6 +258,6 @@ __global__ void __launch_bounds__(WM * WN * 32, MINB) fact_kernel( // The prompt form the batched pass launches from 64 rows: 64 x 128 on 2 x 4 warps, bf16 out (one slice) or fp32 slices. template __global__ void fact::fact_kernel<64, 128, 2, 4, 3, 2, false, 2, false, true, false>(const __nv_bfloat16*, const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, float*, int, int, int, int, int, int, int); -// The folded form from 64 rows with K slices: every slice in the CTA, their totals added in slice order (one fp32 slice). +// The folded form from 64 rows: every K slice in the CTA, their totals added in slice order (one fp32 slice). template __global__ void fact::fact_kernel<64, 128, 2, 4, 3, 2, false, 2, true, true, true>(const __nv_bfloat16*, const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, float*, int, int, int, int, int, int, int); diff --git a/zig/kernels/cuda/flashnext_hcf.cu b/zig/kernels/cuda/flashnext_hcf.cu index 6ebd1b0a9d..2b5c4d7fbe 100644 --- a/zig/kernels/cuda/flashnext_hcf.cu +++ b/zig/kernels/cuda/flashnext_hcf.cu @@ -123,7 +123,7 @@ __global__ void __launch_bounds__(WM * WN * 32, MINB) hcf_kernel( commit(); unsigned char* p = stage(it % STAGES); if constexpr (NORM) { - // x = bf16(bf16(h) * rinv[row, stream] * scale[k]) as fn_hc_down rounds it, and each row's 32-group sum in order + // x = bf16(h * rinv[row, stream] * scale[k]) as fn_hc_down rounds it, each row's 32-group sum in order float* pxs = reinterpret_cast(p + T::X + T::W + 2 * T::S); for (int c = tid; c < GPS * BM; c += T::THREADS) { const int gi = c / BM, r = c % BM, k0 = (it * GPS + gi) * GS; diff --git a/zig/kernels/cuda/route_rows.cu b/zig/kernels/cuda/route_rows.cu index f29d7038cc..5489c57868 100644 --- a/zig/kernels/cuda/route_rows.cu +++ b/zig/kernels/cuda/route_rows.cu @@ -45,8 +45,7 @@ __device__ __forceinline__ void block_best(float& v, int& i, float (*sv)[NT / 32 using namespace tf_route; -// OUT[m, e] = x[m] . w[e], a warp of 8 experts and RT tiles of 16 rows: one mma chain over K from zero, STAGES of KC -// staged (D % 16 == 0). A row's sums never depend on RT. +// OUT[m, e] = x[m] . w[e], 8 experts a warp, RT tiles of 16 rows, one mma chain over K, STAGES of KC (row sums fixed). template __device__ __forceinline__ void route_logits(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, int M, int D, int E) { @@ -117,8 +116,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_route_logits(const __nv_bfl route_logits<4, 2, 128>(X, x_ld, W, OUT, M, D, E); } -// route_logits over 128 rows and 64 experts a block, eight warps (from 64 rows: each router row read once per 128 -// input rows), the same sums. +// route_logits over 128 rows and 64 experts a block, eight warps (from 64 rows), the same sums. extern "C" __global__ void __launch_bounds__(256) tf_route_logits_wide(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, int M, int D, int E) { diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 48f12e3cbf..e20dfe01e1 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -7,6 +7,7 @@ const ple = @import("cuda_ple.zig"); const wt = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); const parts = @import("cuda_batch_state.zig"); +const tiles = @import("cuda_tiles.zig"); const layers_n = wt.layers_n; const linear_n = wt.linear_n; @@ -51,7 +52,6 @@ const ratio = wt.index_ratio; const top = wt.index_top; const idw = wt.index_idw; const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; -const qmm_smem: u32 = 9472; /// Pairs an expert item takes: the core rows kernel's tile. const rows_tile: usize = cuda.experts.RowsExperts.tile; const plan_small: usize = cuda.grouped.small; @@ -60,11 +60,6 @@ fn cint(v: usize) c_int { return @intCast(v); } -/// tile_of's sweep: row tiles a column pass keeps in 12 MiB of the input. -fn sweepOf(m: usize, bm: usize, k: usize) usize { - return @max(@as(usize, 1), @min(cdiv(m, bm), (12 << 20) / (bm * k * 2))); -} - fn cdiv(a: usize, b: usize) usize { return (a + b - 1) / b; } @@ -111,45 +106,9 @@ pub fn Pass(comptime Tri: type) type { // ---- launches ---------------------------------------------------------------------------------------- - /// qmm.cu's lane tile over m rows: out bf16 [m, n] (sk 1) or the unreduced fp32 slices [sk, m, n] in part. - /// Returns the K slices written to part: 1 when the prompt tile folded them (their sum in slice order). + /// qmm.cu's arithmetic over m rows into out (bf16) or part (fp32 K slices); returns the slices (1 when folded). pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !usize { - if (w.native) |native| { - if (sk == 1) { - try self.k.native_lane.matmul(self.stream, x, w.k, m, native, out, self.s.native_part); - return 1; - } - // Legacy callers consume SK slices; the native lane has already reduced its own shape-driven slices. - try self.k.native_lane.matmulF32(self.stream, x, w.k, m, native, part, self.s.native_part); - const bytes = m * w.n * 4; - if (sk > 1) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes * (sk - 1), self.stream.handle), "cuMemsetD8Async"); - return sk; - } - // from 64 rows the prompt tile (flashnext_fact.cu), bit for bit the 16-row lane tile - const big = m >= 64; - const bm: usize = if (big) 64 else 16; - const bn: usize = if (big) 128 else 64; - const k = w.k; - const rows_t = cdiv(m, bm); - const cols_t = cdiv(w.n, bn); - const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); - const npad = (w.n + 127) / 128 * 128; - var a: cuda.Args = .{}; - a.add(x); - a.add(xs); - a.add(w.w.ptr); - a.add(w.s.ptr); - a.add(w.b.ptr); - a.add(out); - a.add(part); - for ([_]usize{ m, w.n, k, sk, npad, k, sweep }) |v| a.add(cint(v)); - const fold = big and sk > 1; - const f = if (fold) self.k.fact_fold else if (big) self.k.fact else self.k.lane; - const block: u32 = if (big) 256 else 128; - const shared: u32 = if (big) parts.fact_smem else qmm_smem; - const z: usize = if (fold) 1 else sk; - try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(z) }, .block = .{ .x = block }, .shared = shared }, self.stream, &a); - return z; + return tiles.matmul(self.k, self.stream, x, xs, w, out, part, self.s.native_part, m, sk); } pub fn sumSlices(self: *Self, part: u64, out: u64, bf16: bool, n: usize, sk: usize) !void { @@ -204,36 +163,9 @@ pub fn Pass(comptime Tri: type) type { try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", rs, n, dims, streams, 2, 10, 11, 1); } - /// Prompt rows' read-out on flashnext_hcf.cu's tiles (fn_hc_down's and fn_hc_upmix's bits): the down rows with their - /// K slices folded, the reduce on one slice, the up rows in bf16, then the mix renorming the streams on the fly. + /// Prompt rows' read-out on flashnext_hcf.cu's tiles, fn_hc_down's and fn_hc_upmix's bits (K slices folded). fn readoutTiles(self: *Self, hc_w: HcW, r0: usize, n: usize) !void { - const s = self.s; - const h = s.h + r0 * wide * 2; - const pss = s.pss + r0 * pss_n * 4; - const rinv = s.rinv + r0 * streams * 4; - const act = s.act + r0 * low * 2; - const xs_act = s.xs_act + r0 * (low / 32) * 4; - const inj_ptr = if (hc_w.has_inj == 1) s.inj + r0 * streams * 2 else act; - const nc = dims / 256; - var a: cuda.Args = .{}; - for ([_]u64{ h, 0, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.hc_part }) |v| a.add(v); - for ([_]usize{ n, hc_w.n_down, wide, hc_sk, wide, sweepOf(n, 32, wide) }) |v| a.add(cint(v)); - a.add(pss); - a.add(hc_w.scale.ptr); - a.add(eps); - for ([_]usize{ dims, nc, streams }) |v| a.add(cint(v)); - try cuda.launch.launch(self.k.hcf_down, .{ .grid = .{ .x = @intCast(cdiv(n, 32) * cdiv(hc_w.n_down, 192)) }, .block = .{ .x = 256 }, .shared = parts.hcf_down_smem }, self.stream, &a); - try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, 1, n, streams, low, hc_w.n_down, hc_w.has_inj); - var b: cuda.Args = .{}; - for ([_]u64{ act, xs_act, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.proj }) |v| b.add(v); - for ([_]usize{ n, wide, low, 1, low, sweepOf(n, 64, low) }) |v| b.add(cint(v)); - b.add(@as(u64, 0)); - b.add(@as(u64, 0)); - b.add(@as(f32, 0)); - for ([_]usize{ dims, nc, streams }) |v| b.add(cint(v)); - try cuda.launch.launch(self.k.hcf_up, .{ .grid = .{ .x = @intCast(cdiv(n, 64) * cdiv(wide, 128)) }, .block = .{ .x = 256 }, .shared = parts.hcf_up_smem }, self.stream, &b); - try self.tri.hcRinv(pss, rinv, eps, n, dims, nc, streams); - try self.tri.hcMixH(s.proj, h, rinv, hc_w.scale.ptr, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); + return tiles.readout(self.tri, self.k, self.stream, self.s, hc_w, r0, n); } /// A read-out over rows [r0, r0 + n) of the streams: mixed and its group sums, and the inject gates. @@ -342,7 +274,7 @@ pub fn Pass(comptime Tri: type) type { a.add(ggb); a.add(state); if (rows >= 64) { - // two threads a state row (fn_gdn_chain_tpr), fn_gdn_rows' bits; in place when there is no rollback copy + // two threads a state row (fn_gdn_chain_tpr), fn_gdn_rows' bits; in place without a rollback copy a.add(if (state_out != 0) state_out else state); a.add(cint(rows)); a.add(ys); @@ -405,7 +337,7 @@ pub fn Pass(comptime Tri: type) type { const pos = at.pos + (r0 / att_rows) * 16; const iq = at.iq + r0 * index_heads * index_dim * 2; const sel = r0 - g0; // this block's rows in the selection buffers - // rows past `top` complete blocks score theirs; _scores leaves the rest alone, so skip it when none is + // rows past `top` complete blocks score theirs; _scores leaves the rest, 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))); var a: cuda.Args = .{}; a.add(self.seq.scores); @@ -428,11 +360,6 @@ pub fn Pass(comptime Tri: type) type { } } - fn fexp(self: *Self, f: cuda.Function, wn: usize, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, plan: cuda.grouped.Plan, out: u64, n: usize, items: usize) !void { - var a = cuda.experts.Affine4Experts.argsOf(x, x_stride, slot_n, w, kg, nb, plan, out, n); - try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(items * cdiv(nb, wn)) }, .block = .{ .x = 256 }, .shared = parts.fexp_smem }, self.stream, &a); - } - /// Router, each row's ten experts and the shared one, the grouped expert launches: y [rows, 11, dims] fp32. pub fn moeBlock(self: *Self, lw: *const Layer) !void { const s = self.s; @@ -467,8 +394,8 @@ pub fn Pass(comptime Tri: type) type { // prompt rows on 64-pair factored tiles, expert_rows' bits try self.k.router.route(self.stream, s.picks, pairs, e, parts.fexp_tile, plan); const items64 = cdiv(pairs, parts.fexp_tile) + e; - try self.fexp(self.k.fexp_up, 2, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items64); - try self.fexp(self.k.fexp_down, 4, s.moe_act, moe_w, 0, lw.experts_down.ptr, moe_w / 32, dims / 32, plan, s.y, dims, items64); + try tiles.experts(self.stream, self.k.fexp_up, 2, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items64); + try tiles.experts(self.stream, self.k.fexp_down, 4, s.moe_act, moe_w, 0, lw.experts_down.ptr, moe_w / 32, dims / 32, plan, s.y, dims, items64); return; } try self.k.router.route(self.stream, s.picks, pairs, e, rows_tile, plan); diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 12e8a548c2..f6d266ec15 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -233,7 +233,6 @@ pub const Tri = struct { }, &.{ ci("SK", sk), ci("S", streams), ci("LOW", low), ci("LOWP", rmsBlock(low)), ci("HAS_INJ", has_inj), ci("NDN", ndn) }); } - /// qmm.hc_upmix: the up projection and the stream mix. N is the wide row, K is the low rank. M folds when it is 1. /// The read-out's 1 / rms rows, ours only. pub fn hcRinv(t: Tri, pss: u64, rinv: u64, eps: f32, rows: usize, dims: usize, nc: usize, streams: usize) !void { const g = t.glue("_qmm_hcdown") orelse return error.NoGlue; @@ -245,6 +244,7 @@ pub const Tri = struct { return g.hcMixH(t.s, up, h, rinv, scale, mixed, xsm, rows, dims, streams); } + /// qmm.hc_upmix: the up projection and the stream mix. N is the wide row, K is the low rank. M folds when it is 1. pub fn hcUpmix(t: Tri, act: u64, xs: u64, weight: u64, scales: u64, biases: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { if (t.glue("_qmm_upmix")) |g| return g.hcUpmix(t.s, act, xs, weight, scales, biases, normed, mixed, xs_mixed, rows, n, k, dims, streams); try t.run("_qmm_upmix", .{ cdiv(rows, 16), u(dims / 32), 1 }, &.{ diff --git a/zig/src/families/flashnext/cuda_tiles.zig b/zig/src/families/flashnext/cuda_tiles.zig new file mode 100644 index 0000000000..d6413f6009 --- /dev/null +++ b/zig/src/families/flashnext/cuda_tiles.zig @@ -0,0 +1,95 @@ +//! Prompt rows on tiles keeping #952's bits: the dense matmul (qmm.cu's arithmetic), the read-out and the experts. +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const parts = @import("cuda_batch_state.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const eps = wt.eps; +const hc_sk = wt.hc_sk; +const pss_n = (dims / 256) * streams; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const Kernels = parts.Kernels; +const Scratch = parts.Scratch; +/// qmm.cu's 16 x 64 lane tile's shared bytes. +const qmm_smem: u32 = 9472; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// tile_of's sweep: row tiles a column pass keeps in 12 MiB of the input. +fn sweepOf(m: usize, bm: usize, k: usize) usize { + return @max(@as(usize, 1), @min(cdiv(m, bm), (12 << 20) / (bm * k * 2))); +} + +/// x [m, k] times w: the native lane, the 16-row lane tile, or from 64 rows the prompt tile (folded with K slices). +pub fn matmul(k: *const Kernels, stream: cuda.Stream, x: u64, xs: u64, w: Q4, out: u64, part: u64, native_part: u64, m: usize, sk: usize) !usize { + if (w.native) |native| { + if (sk == 1) { + try k.native_lane.matmul(stream, x, w.k, m, native, out, native_part); + return 1; + } + // Legacy callers consume SK slices; the native lane has already reduced its own shape-driven slices. + try k.native_lane.matmulF32(stream, x, w.k, m, native, part, native_part); + const bytes = m * w.n * 4; + if (sk > 1) try stream.d.check(stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes * (sk - 1), stream.handle), "cuMemsetD8Async"); + return sk; + } + // from 64 rows the prompt tile (flashnext_fact.cu), bit for bit the 16-row lane tile + const big = m >= 64; + const bm: usize = if (big) 64 else 16; + const bn: usize = if (big) 128 else 64; + var a: cuda.Args = .{}; + for ([_]u64{ x, xs, w.w.ptr, w.s.ptr, w.b.ptr, out, part }) |v| a.add(v); + for ([_]usize{ m, w.n, w.k, sk, (w.n + 127) / 128 * 128, w.k, sweepOf(m, bm, w.k) }) |v| a.add(cint(v)); + const fold = big and sk > 1; + const f = if (fold) k.fact_fold else if (big) k.fact else k.lane; + const shared: u32 = if (big) parts.fact_smem else qmm_smem; + const z: usize = if (fold) 1 else sk; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(cdiv(m, bm) * cdiv(w.n, bn)), .z = @intCast(z) }, .block = .{ .x = if (big) 256 else 128 }, .shared = shared }, stream, &a); + return z; +} + +/// Rows [r0, r0 + n)'s read-out on flashnext_hcf.cu: down with its K slices folded, the reduce, up in bf16, the mix. +pub fn readout(tri: anytype, k: *const Kernels, stream: cuda.Stream, s: *Scratch, hc_w: HcW, r0: usize, n: usize) !void { + const h = s.h + r0 * wide * 2; + const pss = s.pss + r0 * pss_n * 4; + const rinv = s.rinv + r0 * streams * 4; + const act = s.act + r0 * low * 2; + const xs_act = s.xs_act + r0 * (low / 32) * 4; + const inj_ptr = if (hc_w.has_inj == 1) s.inj + r0 * streams * 2 else act; + const nc = dims / 256; + var a: cuda.Args = .{}; + for ([_]u64{ h, 0, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.hc_part }) |v| a.add(v); + for ([_]usize{ n, hc_w.n_down, wide, hc_sk, wide, sweepOf(n, 32, wide) }) |v| a.add(cint(v)); + a.add(pss); + a.add(hc_w.scale.ptr); + a.add(eps); + for ([_]usize{ dims, nc, streams }) |v| a.add(cint(v)); + try cuda.launch.launch(k.hcf_down, .{ .grid = .{ .x = @intCast(cdiv(n, 32) * cdiv(hc_w.n_down, 192)) }, .block = .{ .x = 256 }, .shared = parts.hcf_down_smem }, stream, &a); + try tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, 1, n, streams, low, hc_w.n_down, hc_w.has_inj); + var b: cuda.Args = .{}; + for ([_]u64{ act, xs_act, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.proj }) |v| b.add(v); + for ([_]usize{ n, wide, low, 1, low, sweepOf(n, 64, low) }) |v| b.add(cint(v)); + b.add(@as(u64, 0)); + b.add(@as(u64, 0)); + b.add(@as(f32, 0)); + for ([_]usize{ dims, nc, streams }) |v| b.add(cint(v)); + try cuda.launch.launch(k.hcf_up, .{ .grid = .{ .x = @intCast(cdiv(n, 64) * cdiv(wide, 128)) }, .block = .{ .x = 256 }, .shared = parts.hcf_up_smem }, stream, &b); + try tri.hcRinv(pss, rinv, eps, n, dims, nc, streams); + try tri.hcMixH(s.proj, h, rinv, hc_w.scale.ptr, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); +} + +/// The grouped experts on flashnext_fexp.cu's 64-pair tiles (expert_rows' bits), wn column blocks a CTA. +pub fn experts(stream: cuda.Stream, f: cuda.Function, wn: usize, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, plan: cuda.grouped.Plan, out: u64, n: usize, items: usize) !void { + var a = cuda.experts.Affine4Experts.argsOf(x, x_stride, slot_n, w, kg, nb, plan, out, n); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(items * cdiv(nb, wn)) }, .block = .{ .x = 256 }, .shared = parts.fexp_smem }, stream, &a); +} From 596f204b9f91a983251c0ad7fa5f23725e68131d Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:30:55 +0200 Subject: [PATCH 11/16] Flash Next CUDA: the prompt tiles' modules load in cuda_tile_kernels.zig (cuda_batch_state.zig under 600 lines) --- zig/src/families/flashnext/cuda_batch.zig | 9 +-- .../families/flashnext/cuda_batch_state.zig | 54 +++-------------- .../families/flashnext/cuda_tile_kernels.zig | 58 +++++++++++++++++++ zig/src/families/flashnext/cuda_tiles.zig | 11 ++-- 4 files changed, 77 insertions(+), 55 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_tile_kernels.zig diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index e20dfe01e1..5292780a12 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -8,6 +8,7 @@ const wt = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); const parts = @import("cuda_batch_state.zig"); const tiles = @import("cuda_tiles.zig"); +const tk = @import("cuda_tile_kernels.zig"); const layers_n = wt.layers_n; const linear_n = wt.linear_n; @@ -392,10 +393,10 @@ pub fn Pass(comptime Tri: type) type { } if (n >= 64) { // prompt rows on 64-pair factored tiles, expert_rows' bits - try self.k.router.route(self.stream, s.picks, pairs, e, parts.fexp_tile, plan); - const items64 = cdiv(pairs, parts.fexp_tile) + e; - try tiles.experts(self.stream, self.k.fexp_up, 2, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items64); - try tiles.experts(self.stream, self.k.fexp_down, 4, s.moe_act, moe_w, 0, lw.experts_down.ptr, moe_w / 32, dims / 32, plan, s.y, dims, items64); + try self.k.router.route(self.stream, s.picks, pairs, e, tk.fexp_tile, plan); + const items64 = cdiv(pairs, tk.fexp_tile) + e; + try tiles.experts(self.stream, self.k.tiles.fexp_up, 2, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items64); + try tiles.experts(self.stream, self.k.tiles.fexp_down, 4, s.moe_act, moe_w, 0, lw.experts_down.ptr, moe_w / 32, dims / 32, plan, s.y, dims, items64); return; } try self.k.router.route(self.stream, s.picks, pairs, e, rows_tile, plan); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index eccd46bbd8..a5fa6bc234 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const gdn = @import("cuda_gdn.zig"); const ple = @import("cuda_ple.zig"); const wt = @import("cuda_prompt_wt.zig"); +const tk = @import("cuda_tile_kernels.zig"); const load = @import("cuda_prompt_load.zig"); const layers_n = wt.layers_n; @@ -52,24 +53,10 @@ pub const max_chunk: usize = 16384; const ratio = wt.index_ratio; const top = wt.index_top; const idw = wt.index_idw; -const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; -const qmm_smem: u32 = 9472; -/// Prompt rows on qmm.cu's arithmetic (flashnext_fact.cu): 64 x 128 tiles, two groups a stage, the lane tile's bits. -const cubin_fact: [:0]const u8 = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; -/// The hyper-connection read-out's prompt tiles (flashnext_hcf.cu), fn_hc_down's and fn_hc_upmix's bits. -const cubin_hcf_down: [:0]const u8 = "_ZN3hcf10hcf_kernelILi32ELi192ELi2ELi4ELi4ELi2ELb1ELb1ELb0ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii"; -const cubin_hcf_up: [:0]const u8 = "_ZN3hcf10hcf_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELb0ELb1ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii"; -pub const hcf_down_smem: u32 = 4 * (2 * 32 * 64 + 2 * 192 * 16 + 2 * 2 * 192 * 2 + 2 * 32 * 4 + 2 * 32 * 4) + 32 * 16; -pub const hcf_up_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4 + 2 * 32 * 4) + 64 * 16; -/// Prompt rows' experts on factored 64-pair tiles (flashnext_fexp.cu), expert_rows' bits: up with SwiGLU, down to fp32. -const cubin_fexp_up: [:0]const u8 = "_ZN15tf_fact_experts11fact_kernelILi2ELi2ELi1ELi4ELi2ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; -const cubin_fexp_down: [:0]const u8 = "_ZN15tf_fact_experts11fact_kernelILi1ELi0ELi2ELi2ELi4ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; -pub const fexp_smem: u32 = 41472; -pub const fexp_tile: usize = 64; +pub const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +pub const qmm_smem: u32 = 9472; /// fn_gdn_chain_tpr's staged rows: 2 x 16 steps of q, k, v (128 fp32 each) and g, beta. pub const chain_tpr_smem: u32 = 2 * 16 * (3 * 128 * 4 + 16); -const cubin_fact_fold: [:0]const u8 = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb1ELb1ELb1EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; -pub const fact_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4); /// Pairs an expert item takes: the core rows kernel's tile. const rows_tile: usize = cuda.experts.RowsExperts.tile; const plan_small: usize = cuda.grouped.small; @@ -118,15 +105,8 @@ pub const Kernels = struct { router: cuda.grouped.Router, moe: cuda.experts.RowsExperts, lane: cuda.Function, - fact_mod: cuda.Module, - fact: cuda.Function, - fact_fold: cuda.Function, - hcf_mod: cuda.Module, - fexp_mod: cuda.Module, - fexp_up: cuda.Function, - fexp_down: cuda.Function, - hcf_down: cuda.Function, - hcf_up: cuda.Function, + /// The prompt tiles that keep #952's bits (cuda_tile_kernels.zig). + tiles: tk.TileKernels, native_lane: cuda.qmmf.Lane, nvfp4: cuda.experts.Nvfp4Experts, front: cuda.Function, @@ -152,24 +132,8 @@ pub const Kernels = struct { errdefer k.batch.unload(); k.lane = try k.qmm.function(cubin_qmm); try k.lane.allowDynamicShared(qmm_smem); - k.fact_mod = try cuda.Module.load(d, cuda.kernels.flashnext_fact); - errdefer k.fact_mod.unload(); - k.fact = try k.fact_mod.function(cubin_fact); - try k.fact.allowDynamicShared(fact_smem); - k.fact_fold = try k.fact_mod.function(cubin_fact_fold); - try k.fact_fold.allowDynamicShared(fact_smem); - k.fexp_mod = try cuda.Module.load(d, cuda.kernels.flashnext_fexp); - errdefer k.fexp_mod.unload(); - k.fexp_up = try k.fexp_mod.function(cubin_fexp_up); - try k.fexp_up.allowDynamicShared(fexp_smem); - k.fexp_down = try k.fexp_mod.function(cubin_fexp_down); - try k.fexp_down.allowDynamicShared(fexp_smem); - k.hcf_mod = try cuda.Module.load(d, cuda.kernels.flashnext_hcf); - errdefer k.hcf_mod.unload(); - k.hcf_down = try k.hcf_mod.function(cubin_hcf_down); - try k.hcf_down.allowDynamicShared(hcf_down_smem); - k.hcf_up = try k.hcf_mod.function(cubin_hcf_up); - try k.hcf_up.allowDynamicShared(hcf_up_smem); + k.tiles = try tk.TileKernels.load(d); + errdefer k.tiles.unload(); k.router = try cuda.grouped.Router.resolve(k.experts); k.moe = try cuda.experts.RowsExperts.resolve(k.rows); k.native_lane = try cuda.qmmf.Lane.load(d, (try d.capability(0)) / 10); @@ -198,9 +162,7 @@ pub const Kernels = struct { k.batch.unload(); k.rows.unload(); k.experts.unload(); - k.hcf_mod.unload(); - k.fexp_mod.unload(); - k.fact_mod.unload(); + k.tiles.unload(); k.qmm.unload(); } }; diff --git a/zig/src/families/flashnext/cuda_tile_kernels.zig b/zig/src/families/flashnext/cuda_tile_kernels.zig new file mode 100644 index 0000000000..edb68ce921 --- /dev/null +++ b/zig/src/families/flashnext/cuda_tile_kernels.zig @@ -0,0 +1,58 @@ +//! The prompt tiles' modules (flashnext_fact, _hcf and _fexp): their instances' names, shared bytes and loading. +const cuda = @import("cuda"); + +/// Prompt rows on qmm.cu's arithmetic (flashnext_fact.cu): 64 x 128 tiles, two groups a stage, the lane tile's bits. +pub const cubin_fact: [:0]const u8 = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +/// The hyper-connection read-out's prompt tiles (flashnext_hcf.cu), fn_hc_down's and fn_hc_upmix's bits. +pub const cubin_hcf_down: [:0]const u8 = "_ZN3hcf10hcf_kernelILi32ELi192ELi2ELi4ELi4ELi2ELb1ELb1ELb0ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii"; +pub const cubin_hcf_up: [:0]const u8 = "_ZN3hcf10hcf_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELb0ELb1ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii"; +pub const hcf_down_smem: u32 = 4 * (2 * 32 * 64 + 2 * 192 * 16 + 2 * 2 * 192 * 2 + 2 * 32 * 4 + 2 * 32 * 4) + 32 * 16; +pub const hcf_up_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4 + 2 * 32 * 4) + 64 * 16; +/// Prompt rows' experts on factored 64-pair tiles (flashnext_fexp.cu), expert_rows' bits: up with SwiGLU, down to fp32. +pub const cubin_fexp_up: [:0]const u8 = "_ZN15tf_fact_experts11fact_kernelILi2ELi2ELi1ELi4ELi2ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +pub const cubin_fexp_down: [:0]const u8 = "_ZN15tf_fact_experts11fact_kernelILi1ELi0ELi2ELi2ELi4ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +pub const fexp_smem: u32 = 41472; +pub const fexp_tile: usize = 64; +pub const cubin_fact_fold: [:0]const u8 = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb1ELb1ELb1EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +pub const fact_smem: u32 = 3 * (2 * 64 * 64 + 2 * 128 * 16 + 2 * 2 * 128 * 2 + 2 * 64 * 4); + +pub const TileKernels = struct { + fact_mod: cuda.Module, + fact: cuda.Function, + fact_fold: cuda.Function, + hcf_mod: cuda.Module, + hcf_down: cuda.Function, + hcf_up: cuda.Function, + fexp_mod: cuda.Module, + fexp_up: cuda.Function, + fexp_down: cuda.Function, + + pub fn load(d: *const cuda.Driver) !TileKernels { + var k: TileKernels = undefined; + k.fact_mod = try cuda.Module.load(d, cuda.kernels.flashnext_fact); + errdefer k.fact_mod.unload(); + k.fact = try k.fact_mod.function(cubin_fact); + try k.fact.allowDynamicShared(fact_smem); + k.fact_fold = try k.fact_mod.function(cubin_fact_fold); + try k.fact_fold.allowDynamicShared(fact_smem); + k.fexp_mod = try cuda.Module.load(d, cuda.kernels.flashnext_fexp); + errdefer k.fexp_mod.unload(); + k.fexp_up = try k.fexp_mod.function(cubin_fexp_up); + try k.fexp_up.allowDynamicShared(fexp_smem); + k.fexp_down = try k.fexp_mod.function(cubin_fexp_down); + try k.fexp_down.allowDynamicShared(fexp_smem); + k.hcf_mod = try cuda.Module.load(d, cuda.kernels.flashnext_hcf); + errdefer k.hcf_mod.unload(); + k.hcf_down = try k.hcf_mod.function(cubin_hcf_down); + try k.hcf_down.allowDynamicShared(hcf_down_smem); + k.hcf_up = try k.hcf_mod.function(cubin_hcf_up); + try k.hcf_up.allowDynamicShared(hcf_up_smem); + return k; + } + + pub fn unload(k: *TileKernels) void { + k.hcf_mod.unload(); + k.fexp_mod.unload(); + k.fact_mod.unload(); + } +}; diff --git a/zig/src/families/flashnext/cuda_tiles.zig b/zig/src/families/flashnext/cuda_tiles.zig index d6413f6009..a623d411af 100644 --- a/zig/src/families/flashnext/cuda_tiles.zig +++ b/zig/src/families/flashnext/cuda_tiles.zig @@ -2,6 +2,7 @@ const cuda = @import("cuda"); const wt = @import("cuda_prompt_wt.zig"); const parts = @import("cuda_batch_state.zig"); +const tk = @import("cuda_tile_kernels.zig"); const dims = wt.dims; const streams = wt.streams; @@ -51,8 +52,8 @@ pub fn matmul(k: *const Kernels, stream: cuda.Stream, x: u64, xs: u64, w: Q4, ou for ([_]u64{ x, xs, w.w.ptr, w.s.ptr, w.b.ptr, out, part }) |v| a.add(v); for ([_]usize{ m, w.n, w.k, sk, (w.n + 127) / 128 * 128, w.k, sweepOf(m, bm, w.k) }) |v| a.add(cint(v)); const fold = big and sk > 1; - const f = if (fold) k.fact_fold else if (big) k.fact else k.lane; - const shared: u32 = if (big) parts.fact_smem else qmm_smem; + const f = if (fold) k.tiles.fact_fold else if (big) k.tiles.fact else k.lane; + const shared: u32 = if (big) tk.fact_smem else qmm_smem; const z: usize = if (fold) 1 else sk; try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(cdiv(m, bm) * cdiv(w.n, bn)), .z = @intCast(z) }, .block = .{ .x = if (big) 256 else 128 }, .shared = shared }, stream, &a); return z; @@ -74,7 +75,7 @@ pub fn readout(tri: anytype, k: *const Kernels, stream: cuda.Stream, s: *Scratch a.add(hc_w.scale.ptr); a.add(eps); for ([_]usize{ dims, nc, streams }) |v| a.add(cint(v)); - try cuda.launch.launch(k.hcf_down, .{ .grid = .{ .x = @intCast(cdiv(n, 32) * cdiv(hc_w.n_down, 192)) }, .block = .{ .x = 256 }, .shared = parts.hcf_down_smem }, stream, &a); + try cuda.launch.launch(k.tiles.hcf_down, .{ .grid = .{ .x = @intCast(cdiv(n, 32) * cdiv(hc_w.n_down, 192)) }, .block = .{ .x = 256 }, .shared = tk.hcf_down_smem }, stream, &a); try tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, 1, n, streams, low, hc_w.n_down, hc_w.has_inj); var b: cuda.Args = .{}; for ([_]u64{ act, xs_act, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.proj }) |v| b.add(v); @@ -83,7 +84,7 @@ pub fn readout(tri: anytype, k: *const Kernels, stream: cuda.Stream, s: *Scratch b.add(@as(u64, 0)); b.add(@as(f32, 0)); for ([_]usize{ dims, nc, streams }) |v| b.add(cint(v)); - try cuda.launch.launch(k.hcf_up, .{ .grid = .{ .x = @intCast(cdiv(n, 64) * cdiv(wide, 128)) }, .block = .{ .x = 256 }, .shared = parts.hcf_up_smem }, stream, &b); + try cuda.launch.launch(k.tiles.hcf_up, .{ .grid = .{ .x = @intCast(cdiv(n, 64) * cdiv(wide, 128)) }, .block = .{ .x = 256 }, .shared = tk.hcf_up_smem }, stream, &b); try tri.hcRinv(pss, rinv, eps, n, dims, nc, streams); try tri.hcMixH(s.proj, h, rinv, hc_w.scale.ptr, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); } @@ -91,5 +92,5 @@ pub fn readout(tri: anytype, k: *const Kernels, stream: cuda.Stream, s: *Scratch /// The grouped experts on flashnext_fexp.cu's 64-pair tiles (expert_rows' bits), wn column blocks a CTA. pub fn experts(stream: cuda.Stream, f: cuda.Function, wn: usize, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, plan: cuda.grouped.Plan, out: u64, n: usize, items: usize) !void { var a = cuda.experts.Affine4Experts.argsOf(x, x_stride, slot_n, w, kg, nb, plan, out, n); - try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(items * cdiv(nb, wn)) }, .block = .{ .x = 256 }, .shared = parts.fexp_smem }, stream, &a); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(items * cdiv(nb, wn)) }, .block = .{ .x = 256 }, .shared = tk.fexp_smem }, stream, &a); } From f1b8c882e8b2b14c6e89e8b6b5ac686caa367d4c Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:31:01 +0200 Subject: [PATCH 12/16] Kernel contract: Flash Next's prompt tiles as replacements of 1.1.0's kernels, bit for bit, their launches registered --- zig/src/core/contract/affine.zig | 2 +- .../core/contract/flashnext_tile_records.zig | 18 + zig/src/core/contract/registry.zig | 11 +- zig/src/cuda/route_rows.zig | 6 +- zig/src/families/flashnext/cuda_launch.zig | 3 + zig/tests/cuda/contract_flashnext_fact.zig | 165 +++++++++ zig/tests/cuda/contract_flashnext_fexp.zig | 295 +++++++++++++++ zig/tests/cuda/contract_flashnext_glue.zig | 346 ++++++++++++++++++ zig/tests/cuda/contract_route_rows.zig | 120 ++++++ zig/tests/cuda/contract_tests.zig | 68 +++- 10 files changed, 1024 insertions(+), 10 deletions(-) create mode 100644 zig/src/core/contract/flashnext_tile_records.zig create mode 100644 zig/tests/cuda/contract_flashnext_fact.zig create mode 100644 zig/tests/cuda/contract_flashnext_fexp.zig create mode 100644 zig/tests/cuda/contract_flashnext_glue.zig create mode 100644 zig/tests/cuda/contract_route_rows.zig diff --git a/zig/src/core/contract/affine.zig b/zig/src/core/contract/affine.zig index 41e6ba0359..f0ad4d5eae 100644 --- a/zig/src/core/contract/affine.zig +++ b/zig/src/core/contract/affine.zig @@ -108,7 +108,7 @@ pub const Affine = struct { } /// Output (r, j) in float64 by the kernel's form, and the magnitude of the terms its sums add. - fn one(f: Affine, r: usize, j: usize) [2]f64 { + pub fn one(f: Affine, r: usize, j: usize) [2]f64 { const k = f.shape.k; var y: f64 = 0; var mag: f64 = 0; diff --git a/zig/src/core/contract/flashnext_tile_records.zig b/zig/src/core/contract/flashnext_tile_records.zig new file mode 100644 index 0000000000..a8a2f148a3 --- /dev/null +++ b/zig/src/core/contract/flashnext_tile_records.zig @@ -0,0 +1,18 @@ +//! Flash Next's prompt tiles' CUDA records (gb10), in registry.cuda: from nvcc 13.2, to re-derive with 13.0 at landing. +const CudaRecord = @import("registry.zig").CudaRecord; + +pub const cuda = [_]CudaRecord{ + .{ .launch = .{ .name = "flashnext_fact factored prompt matmul 64x128", .symbol = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii", .image = "flashnext_fact", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_fact factored prompt matmul 64x128 folded", .symbol = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb1ELb1ELb1EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii", .image = "flashnext_fact", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hcf hyper-connection down 32x192 folded", .symbol = "_ZN3hcf10hcf_kernelILi32ELi192ELi2ELi4ELi4ELi2ELb1ELb1ELb0ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii", .image = "flashnext_hcf", .threads = 256, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 126, .static = 0, .bound = 256 }, .{ .regs = 126, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hcf hyper-connection up 64x128", .symbol = "_ZN3hcf10hcf_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELb0ELb1ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii", .image = "flashnext_hcf", .threads = 256, .dynamic = 43264, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_fexp factored experts gate and up", .symbol = "_ZN15tf_fact_experts11fact_kernelILi2ELi2ELi1ELi4ELi2ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif", .image = "flashnext_fexp", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 125, .static = 0, .bound = 256 }, .{ .regs = 125, .static = 0, .bound = 256 }, .{ .regs = 125, .static = 0, .bound = 256 }, .{ .regs = 116, .static = 0, .bound = 256 }, .{ .regs = 116, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_fexp factored experts down", .symbol = "_ZN15tf_fact_experts11fact_kernelILi1ELi0ELi2ELi2ELi4ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif", .image = "flashnext_fexp", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 120, .static = 0, .bound = 256 }, .{ .regs = 120, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_chain_tpr", .symbol = "fn_gdn_chain_tpr", .image = "flashnext_batch", .threads = 256, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 152, .static = 0, .bound = 256 }, .{ .regs = 152, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn qsa_tile scores", .symbol = "_ZN8qsa_tile6scoresILi8ELi4EEEvPK13__nv_bfloat16S3_PKiPfiiii", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 224, .static = 16384, .bound = 256 }, .{ .regs = 224, .static = 16384, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_fused", .symbol = "fn_attn_fused", .image = "flashnext_attn", .threads = 128, .dynamic = 26624 }, .usage = .{ .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 121, .static = 640, .bound = 128 }, .{ .regs = 121, .static = 640, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_one", .symbol = "fn_attn_chunks_one", .image = "flashnext_attn", .threads = 128, .dynamic = 26624, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 96, .static = 640, .bound = 128 }, .{ .regs = 96, .static = 640, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_rinv", .symbol = "fn_hc_rinv", .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 = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_mix_h", .symbol = "fn_hc_mix_h", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 38, .static = 1024, .bound = 256 }, .{ .regs = 38, .static = 1024, .bound = 256 } } }, + .{ .launch = .{ .name = "route_rows tf_route_logits_wide", .symbol = "tf_route_logits_wide", .image = "route_rows", .threads = 256, .dynamic = 55296, .opt_in = true }, .usage = .{ .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 58, .static = 0, .bound = 256 }, .{ .regs = 58, .static = 0, .bound = 256 } } }, +}; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 717c770c1c..589256fe16 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -222,7 +222,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, -} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda; +} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda; /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins. pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 }; @@ -248,8 +248,8 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "qmmf FP8G fp32 out 64-row tile, fused slices", .usage = .{ .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 155, .static = 0, .bound = 128 }, .{ .regs = 155, .static = 0, .bound = 128 } } }, .{ .name = "qmmf BF16 fp32 out 64-row tile, fused slices", .usage = .{ .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 } } }, .{ .name = "flashnext_gdn flashnext_gdn_chain", .usage = .{ .{ .regs = 51, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 64, .static = 2060, .bound = 1024 }, .{ .regs = 64, .static = 2060, .bound = 1024 } } }, - .{ .name = "flashnext_attn fn_attn_chunks", .usage = .{ .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 73, .static = 640, .bound = 128 }, .{ .regs = 73, .static = 640, .bound = 128 } } }, - .{ .name = "flashnext_attn fn_attn_chunks_ahead", .usage = .{ .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 } } }, + .{ .name = "flashnext_attn fn_attn_chunks", .usage = .{ .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 } } }, + .{ .name = "flashnext_attn fn_attn_chunks_ahead", .usage = .{ .{ .regs = 70, .static = 768, .bound = 128 }, .{ .regs = 70, .static = 768, .bound = 128 }, .{ .regs = 70, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 } } }, .{ .name = "flashnext_attn fn_attn_merge", .usage = .{ .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 } } }, .{ .name = "flashnext_batch fn_gdn_front", .usage = .{ .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 30, .static = 1024, .bound = 256 }, .{ .regs = 30, .static = 1024, .bound = 256 } } }, .{ .name = "flashnext_hc fn_hc_down", .usage = .{ .{ .regs = 56, .static = 64, .bound = 128 }, .{ .regs = 54, .static = 64, .bound = 128 }, .{ .regs = 54, .static = 64, .bound = 128 }, .{ .regs = 64, .static = 64, .bound = 128 }, .{ .regs = 64, .static = 64, .bound = 128 } } }, @@ -261,8 +261,8 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "flashnext_hc fn_ple_conv", .usage = .{ .{ .regs = 23, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 36, .static = 0, .bound = 512 }, .{ .regs = 36, .static = 0, .bound = 512 } } }, .{ .name = "flashnext_hc fn_ple_gate", .usage = .{ .{ .regs = 32, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 33, .static = 128, .bound = 512 }, .{ .regs = 33, .static = 128, .bound = 512 } } }, .{ .name = "flashnext_attn fn_qsa_pool", .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 } } }, - .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 42, .static = 1160, .bound = 1024 }, .{ .regs = 42, .static = 1160, .bound = 1024 } } }, - .{ .name = "route_rows tf_route_logits", .usage = .{ .{ .regs = 31, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 } } }, + .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 42, .static = 1168, .bound = 1024 }, .{ .regs = 42, .static = 1168, .bound = 1024 } } }, + .{ .name = "route_rows tf_route_logits", .usage = .{ .{ .regs = 32, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 } } }, .{ .name = "route_rows tf_route_logits_few", .usage = .{ .{ .regs = 31, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 34, .static = 0, .bound = 128 }, .{ .regs = 34, .static = 0, .bound = 128 } } }, .{ .name = "qwen_vision tf_qv_qkv_rope", .usage = .{ .{ .regs = 32, .static = 0, .bound = 128 }, .{ .regs = 37, .static = 0, .bound = 128 }, .{ .regs = 37, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } }, .{ .name = "qwen_vision tf_qv_attention64", .usage = .{ .{ .regs = 164, .static = 32768, .bound = 64 }, .{ .regs = 164, .static = 32768, .bound = 64 }, .{ .regs = 164, .static = 32768, .bound = 64 }, .{ .regs = 162, .static = 32768, .bound = 64 }, .{ .regs = 162, .static = 32768, .bound = 64 } } }, @@ -337,6 +337,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "latent tf_super_topk_few", .usage = .{ .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 } } }, .{ .name = "ck tf_ck4_up_relu2_few", .usage = .{ .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 } } }, .{ .name = "mamba tf_nemo_scan", .usage = .{ .{ .regs = 138, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 } } }, + .{ .name = "flashnext_batch fn_gdn_norm", .usage = .{ .{ .regs = 24, .static = 528, .bound = 128 }, .{ .regs = 26, .static = 528, .bound = 128 }, .{ .regs = 26, .static = 528, .bound = 128 }, .{ .regs = 24, .static = 528, .bound = 128 }, .{ .regs = 24, .static = 528, .bound = 128 } } }, } ++ @import("paging_pending.zig").cuda_132 ++ @import("nemotron_prompt_records.zig").cuda_132 ++ @import("glm_cuda_records.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132) }, .{ .release = .{ .major = 13, .minor = 3 }, .usage = &.{ .{ .name = "router tf_nemo_topk", .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, diff --git a/zig/src/cuda/route_rows.zig b/zig/src/cuda/route_rows.zig index b1fc6ebe89..4aeb28202a 100644 --- a/zig/src/cuda/route_rows.zig +++ b/zig/src/cuda/route_rows.zig @@ -15,10 +15,10 @@ pub const most_experts = 1024; pub const most_picks = 128; /// tf_route_logits' staged rows (16 inputs and 32 experts, padded by 8): two stages of 128, or three of 256. -const many_shared: u32 = 2 * 48 * (128 + 8) * 2; -const few_shared: u32 = 3 * 48 * (256 + 8) * 2; +pub const many_shared: u32 = 2 * 48 * (128 + 8) * 2; +pub const few_shared: u32 = 3 * 48 * (256 + 8) * 2; /// tf_route_logits_wide's: 128 inputs and 64 experts, two stages of 64. -const wide_shared: u32 = 2 * 192 * (64 + 8) * 2; +pub const wide_shared: u32 = 2 * 192 * (64 + 8) * 2; pub const RouteRows = struct { mod: Module, diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index f6d266ec15..9718d4ae08 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -13,6 +13,9 @@ pub const modelopt_experts = @import("cuda_modelopt_experts.zig"); pub const hc = @import("cuda_hc.zig"); pub const draft_affine = @import("cuda_draft_affine.zig"); pub const qmm = @import("cuda_qmm.zig"); +/// The batched pass's and its prompt tiles' symbols and launch sizes, which the contract's tests launch. +pub const batch_state = @import("cuda_batch_state.zig"); +pub const tile_kernels = @import("cuda_tile_kernels.zig"); pub const gdn = @import("cuda_gdn.zig"); pub const mlp = @import("cuda_mlp.zig"); pub const moe = @import("cuda_moe.zig"); diff --git a/zig/tests/cuda/contract_flashnext_fact.zig b/zig/tests/cuda/contract_flashnext_fact.zig new file mode 100644 index 0000000000..e3bf1efe5e --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_fact.zig @@ -0,0 +1,165 @@ +//! Flash Next's prompt matmul on qmm.cu's factored arithmetic (flashnext_fact.cu) bit-equal to the 16x64 lane tile. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Affine = contract.affine.Affine; +const parts = flash.batch_state; +const tk = flash.tile_kernels; + +pub const Kind = enum { prompt, fold }; + +/// A group's 32 products from zero in two k16 steps, two fmas a group, then up to 8 slices added in order. +fn depth(s: c.Shape) usize { + return 32 + 2 * (s.k / 32) + 8; +} + +/// The engine's K slices: PLE values 4, out projections 8. +fn slicesOf(s: c.Shape) usize { + if (s.n == 2560 and s.k == 2560) return 4; + if (s.k == 6144) return 8; + for ([_]usize{ 8, 4, 2 }) |sk| if ((s.k / 32) % (2 * sk) == 0) return sk; + return 1; +} + +const tiles = [3]c.Tile{ .{ .size = 64 }, .{ .size = 128, .step = 32, .min = 32 }, .{ .size = 256, .step = 64, .min = 64 } }; + +pub fn spec(kind: Kind) c.Spec { + return switch (kind) { + .prompt => .{ .name = "cuda flashnext factored prompt matmul 64x128", .accum = .f32, .depth = depth, .ops = 2, .out = .bf16, .replaces = "qmm.cu's 16x64 lane tile", .tiles = tiles, .served = &.{ .{ .rows = 64, .n = 13952, .k = 2560 }, .{ .rows = 100, .n = 2560, .k = 2560 } } }, + .fold => .{ .name = "cuda flashnext factored prompt matmul 64x128, K slices folded", .accum = .f32, .depth = depth, .ops = 2, .out = .f32, .replaces = "the lane tile's K slices then fn_sum_slices", .tiles = tiles, .served = &.{ .{ .rows = 64, .n = 2560, .k = 2560 }, .{ .rows = 100, .n = 2560, .k = 6144 } } }, + }; +} + +pub const Fact = struct { + rig: rig_.Rig, + kind: Kind, + lane: cuda.Function, + prompt: cuda.Function, + fold: cuda.Function, + sum: cuda.Function, + f: ?Affine = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + sk: usize = 1, + npad: usize = 0, + x: u64 = 0, + xs: u64 = 0, + w: u64 = 0, + s: u64 = 0, + b: u64 = 0, + part: u64 = 0, + base: u64 = 0, + out: u64 = 0, + + pub fn deinit(self: *Fact) void { + if (self.f) |f| f.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Fact) g.Kernel { + return .{ .spec = spec(self.kind), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Fact); + + pub fn prepare(self: *Fact, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.f) |f| f.deinit(); + self.f = null; + self.rig.reset(); + // two groups a stage: K in 64s (every served K is) + if (shape.rows == 0 or shape.n == 0 or shape.n % 8 != 0 or shape.k == 0 or shape.k % 64 != 0) return false; + self.sk = if (self.kind == .prompt) 1 else slicesOf(shape); + if (self.kind == .fold and self.sk == 1) return false; + self.shape = shape; + const a = self.rig.gpu.gpa; + const k = shape.k; + const f = try Affine.init(a, shape, 4, 32); + self.f = f; + try f.fill(fill, seed); + // the producers' fp32 group sums, in order + const xs = try a.alloc(f32, shape.rows * (k / 32)); + defer a.free(xs); + for (xs, 0..) |*v, i| { + var t: f32 = 0; + for (f.x[i * 32 ..][0..32]) |h| t += c.fromBf16(h); + v.* = t; + } + const words = try a.alloc(u32, shape.n * k / 8); + defer a.free(words); + f.words(words); + var p = try flash.qmm.pack(a, std.mem.sliceAsBytes(words), std.mem.sliceAsBytes(f.scales), std.mem.sliceAsBytes(f.biases), shape.n, k / 8); + defer p.deinit(a); + self.npad = p.npad; + self.x = try self.rig.input(std.mem.sliceAsBytes(f.x)); + self.xs = try self.rig.input(std.mem.sliceAsBytes(xs)); + self.w = try self.rig.input(p.weight); + self.s = try self.rig.input(p.scales); + self.b = try self.rig.input(p.biases); + const esize = c.size(spec(self.kind).out); + self.out = try self.rig.output(shape.rows * shape.n * esize); + self.base = try self.rig.scratch(shape.rows * shape.n * esize); + if (self.kind == .fold) self.part = try self.rig.scratch(self.sk * shape.rows * shape.n * 4); + return true; + } + + pub fn run(self: *Fact) !void { + switch (self.kind) { + .prompt => try self.launch(self.prompt, 64, 128, 256, tk.fact_smem, self.out, 0, 1, 1), + .fold => try self.launch(self.fold, 64, 128, 256, tk.fact_smem, 0, self.out, self.sk, 1), + } + try self.rig.s.synchronize(); + } + + /// cuda_batch.zig's lane() launch: grid of row and column tiles, z the slices on the grid. + fn launch(self: *Fact, f: cuda.Function, bm: usize, bn: usize, block: u32, shared: u32, out: u64, part: u64, sk: usize, z: usize) !void { + const sh = self.shape; + const rows_t = (sh.rows + bm - 1) / bm; + const cols_t = (sh.n + bn - 1) / bn; + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * sh.k * 2))); + var a: cuda.Args = .{}; + for ([_]u64{ self.x, self.xs, self.w, self.s, self.b, out, part }) |v| a.add(v); + for ([_]usize{ sh.rows, sh.n, sh.k, sk, self.npad, sh.k, sweep }) |v| a.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(z) }, .block = .{ .x = block }, .shared = shared }, self.rig.s, &a); + } + + /// The lane tile on the same inputs: bf16 out, or its K slices summed in order by fn_sum_slices (fp32). + pub fn replaced(self: *Fact) ![]const u8 { + const n = self.shape.rows * self.shape.n; + switch (self.kind) { + .prompt => { + try self.launch(self.lane, 16, 64, 128, parts.qmm_smem, self.base, 0, 1, 1); + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, n * 2); + }, + .fold => { + try self.launch(self.lane, 16, 64, 128, parts.qmm_smem, 0, self.part, self.sk, self.sk); + var a: cuda.Args = .{}; + a.add(self.part); + a.add(self.base); + a.add(@as(c_int, 0)); + a.add(@as(c_longlong, @intCast(n))); + a.add(@as(c_int, @intCast(self.sk))); + try cuda.launch.launch(self.sum, .{ .grid = .{ .x = @intCast((n + 255) / 256) }, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, n * 4); + }, + } + } + + pub fn count(self: *Fact) usize { + return self.f.?.outputs(); + } + + pub fn reference(self: *Fact, want: []f64, mag: []f64) !void { + return self.f.?.reference(want, mag); + } + + pub fn simulate(self: *Fact, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.f.?.simulate(acc, order, out); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_fexp.zig b/zig/tests/cuda/contract_flashnext_fexp.zig new file mode 100644 index 0000000000..8389b4ba13 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_fexp.zig @@ -0,0 +1,295 @@ +//! Flash Next's prompt experts on 64-pair factored tiles (flashnext_fexp.cu), bit-equal to expert_rows. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Affine = contract.affine.Affine; +const parts = flash.batch_state; +const tk = flash.tile_kernels; + +pub const Kind = enum { down, up }; + +/// The packed layout's blocks (experts_pack<1>): 32 columns by 32 inputs, then their scales and biases. +const block_bytes: usize = 160 * 4; +pub const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiiiii"; +const experts_n = 4; + +/// A token's pairs in the up form (ten experts and the shared one); one at a few pairs, so planted rows stay a few. +fn slotsOf(pairs: usize) usize { + return if (pairs < 44) 1 else 11; +} + +/// A group's 32 products from zero, its input sum, two fmas a group. +fn depth(s: c.Shape) usize { + return 32 + 32 + 2 * (s.k / 32); +} + +/// The up form's sums are given (expert_rows' fp32 form): only its SwiGLU epilogue rounds. +fn epilogue(_: c.Shape) usize { + return 0; +} + +const tiles = [3]c.Tile{ .{ .size = 64 }, .{ .size = 128, .step = 32, .min = 32 }, .{ .size = 128, .step = 64, .min = 64 } }; +/// rows = expert pairs (11 a token for up), n = outputs, k = inputs: Flash Next's 640-wide experts on 2560 dims. +const down_shapes = [_]c.Shape{ .{ .rows = 704, .n = 2560, .k = 640 }, .{ .rows = 1100, .n = 2560, .k = 640 } }; +const up_shapes = [_]c.Shape{ .{ .rows = 704, .n = 640, .k = 2560 }, .{ .rows = 1100, .n = 640, .k = 2560 } }; + +pub fn spec(kind: Kind) c.Spec { + return switch (kind) { + .down => .{ .name = "cuda flashnext factored experts down 64 pairs, fp32 out", .accum = .f32, .depth = depth, .ops = 2, .out = .f32, .replaces = "expert_rows' down form", .tiles = tiles, .served = &down_shapes }, + // exp's 2 ulps, 1 + e, the quotient and the product; silu rounded to bf16 before the product + .up => .{ .name = "cuda flashnext factored experts gate and up 64 pairs, SwiGLU", .accum = .f32, .depth = epilogue, .ops = 5, .mids = 1, .out = .bf16, .replaces = "expert_rows' up form", .tiles = tiles, .served = &up_shapes }, + }; +} + +pub const Fns = struct { up: cuda.Function, down: cuda.Function, rows_up: cuda.Function, rows_down: cuda.Function, pack: cuda.Function }; + +const Plan = struct { items: u64 = 0, counts: u64 = 0, members: u64 = 0, n: usize = 0 }; + +pub const Fexp = struct { + rig: rig_.Rig, + kind: Kind, + f: *const Fns, + a: std.mem.Allocator = undefined, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + /// Each (matrix, expert)'s affine weights on the shared token rows (gate then up for the up form). + aff: std.ArrayList(Affine) = .empty, + expert: []usize = &.{}, + gs: []f32 = &.{}, + us: []f32 = &.{}, + plan64: Plan = .{}, + plan32: Plan = .{}, + dx: u64 = 0, + dw: u64 = 0, + base: u64 = 0, + out: u64 = 0, + + pub fn deinit(self: *Fexp) void { + self.drop(); + self.aff.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + fn drop(self: *Fexp) void { + const a = self.rig.gpu.gpa; + for (self.aff.items) |f| f.deinit(); + self.aff.clearRetainingCapacity(); + for ([_][]const u8{ std.mem.sliceAsBytes(self.expert), std.mem.sliceAsBytes(self.gs), std.mem.sliceAsBytes(self.us) }) |b| if (b.len > 0) a.free(b); + self.expert = &.{}; + self.gs = &.{}; + self.us = &.{}; + self.rig.reset(); + } + + pub fn kernel(self: *Fexp) g.Kernel { + return .{ .spec = spec(self.kind), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Fexp); + + fn slots(self: *const Fexp) usize { + return if (self.kind == .up) slotsOf(self.shape.rows) else 0; + } + + fn tokens(self: *const Fexp) usize { + return if (self.kind == .up) (self.shape.rows + self.slots() - 1) / self.slots() else self.shape.rows; + } + + fn tokenOf(self: *const Fexp, p: usize) usize { + return if (self.kind == .up) p / self.slots() else p; + } + + pub fn prepare(self: *Fexp, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.drop(); + // two groups a stage (K in 64s), outputs in 32-column blocks + if (shape.rows == 0 or shape.n == 0 or shape.n % 32 != 0 or shape.k == 0 or shape.k % 64 != 0) return false; + const a = self.rig.gpu.gpa; + self.a = a; + self.shape = shape; + const n = shape.n; + const k = shape.k; + const kg = k / 32; + const mats: usize = if (self.kind == .up) 2 else 1; + self.expert = try a.alloc(usize, shape.rows); + if (self.kind == .up) { + self.gs = try a.alloc(f32, shape.rows * n); + self.us = try a.alloc(f32, shape.rows * n); + } + const words = try a.alloc(u32, mats * experts_n * n * k / 8); + defer a.free(words); + const scales = try a.alloc(u16, mats * experts_n * n * kg); + defer a.free(scales); + const biases = try a.alloc(u16, scales.len); + defer a.free(biases); + // each matrix of each expert from the affine family, every one on the first one's token rows + for (0..mats * experts_n) |m| { + const f = try Affine.init(a, .{ .rows = self.tokens(), .n = n, .k = k }, 4, 32); + try self.aff.append(a, f); + try f.fill(fill, seed + m); + if (m > 0) @memcpy(f.x, self.aff.items[0].x); + f.words(words[m * n * k / 8 ..][0 .. n * k / 8]); + @memcpy(scales[m * n * kg ..][0 .. n * kg], f.scales); + @memcpy(biases[m * n * kg ..][0 .. n * kg], f.biases); + } + // pairs to experts unevenly (one expert past a tile, one with none at small counts) + for (self.expert, 0..) |*e, p| e.* = if (p % 5 < 3) 0 else if (p % 5 == 3) 1 else 3; + self.plan64 = try self.plan(64); + self.plan32 = try self.plan(cuda.experts.RowsExperts.tile); + self.dx = try self.rig.input(std.mem.sliceAsBytes(self.aff.items[0].x)); + // pack each matrix as the engine does (experts_pack<1>), gate and up blocks side by side for the up form + const per = experts_n * (n / 32) * kg * block_bytes; + const packed_w = try a.alloc(u8, mats * per); + defer a.free(packed_w); + for (0..mats) |m| try self.pack(words[m * experts_n * n * k / 8 ..][0 .. experts_n * n * k / 8], scales[m * experts_n * n * kg ..][0 .. experts_n * n * kg], biases[m * experts_n * n * kg ..][0 .. experts_n * n * kg], packed_w[m * per ..][0..per]); + if (self.kind == .up) { + try self.upSums(packed_w[0..per], self.gs); + try self.upSums(packed_w[per..], self.us); + const stacked = try a.alloc(u8, 2 * per); + defer a.free(stacked); + for (0..per / block_bytes) |b| { + @memcpy(stacked[2 * b * block_bytes ..][0..block_bytes], packed_w[b * block_bytes ..][0..block_bytes]); + @memcpy(stacked[(2 * b + 1) * block_bytes ..][0..block_bytes], packed_w[per + b * block_bytes ..][0..block_bytes]); + } + self.dw = try self.rig.input(stacked); + } else self.dw = try self.rig.input(packed_w); + const esize = c.size(spec(self.kind).out); + self.out = try self.rig.output(shape.rows * n * esize); + self.base = try self.rig.scratch(shape.rows * n * esize); + return true; + } + + /// A grouped plan of `tile` pairs an item (expert, first member, count), its members expert by expert. + fn plan(self: *Fexp, tile: usize) !Plan { + const a = self.a; + const members = try a.alloc(i32, self.shape.rows); + defer a.free(members); + var items: std.ArrayList(i32) = .empty; + defer items.deinit(a); + var at: usize = 0; + for (0..experts_n) |e| { + const first = at; + for (self.expert, 0..) |pe, p| if (pe == e) { + members[at] = @intCast(p); + at += 1; + }; + var i = first; + while (i < at) : (i += tile) try items.appendSlice(a, &.{ @intCast(e), @intCast(i), @intCast(@min(tile, at - i)) }); + } + const n_items = [_]i32{@intCast(items.items.len / 3)}; + return .{ .items = try self.rig.input(std.mem.sliceAsBytes(items.items)), .counts = try self.rig.input(std.mem.sliceAsBytes(&n_items)), .members = try self.rig.input(std.mem.sliceAsBytes(members)), .n = items.items.len / 3 }; + } + + fn upload(self: *Fexp, bytes: []const u8) !u64 { + const p = try self.rig.scratch(bytes.len); + const b: cuda.DeviceBuffer = .{ .d = self.rig.gpu.d, .ptr = p, .len = bytes.len }; + try b.upload(0, bytes); + return p; + } + + /// experts_pack<1> on the device, read back: [experts][n / 32][k / 32] blocks. + fn pack(self: *Fexp, words: []const u32, scales: []const u16, biases: []const u16, out: []u8) !void { + const n = self.shape.n; + const kg = self.shape.k / 32; + const dst = try self.rig.scratch(out.len); + var a: cuda.Args = .{}; + a.add(try self.upload(std.mem.sliceAsBytes(words))); + a.add(try self.upload(std.mem.sliceAsBytes(scales))); + a.add(try self.upload(std.mem.sliceAsBytes(biases))); + a.add(dst); + for ([_]usize{ n, self.shape.k / 8, kg, n / 32, 1, 0 }) |v| a.add(@as(c_int, @intCast(v))); // one table, the 0th + try cuda.launch.launch(self.f.pack, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(n / 32), .z = experts_n }, .block = .{ .x = 160 } }, self.rig.s, &a); + @memcpy(out, try self.rig.bytes(dst, out.len)); + } + + /// One matrix's fp32 sums [pairs, n] by expert_rows' down form on the token rows, read back. + fn upSums(self: *Fexp, w: []const u8, out: []f32) !void { + const dw = try self.upload(w); + const dst = try self.rig.scratch(out.len * 4); + try self.rows(self.f.rows_down, self.slots(), dw, dst); + @memcpy(std.mem.sliceAsBytes(out), try self.rig.bytes(dst, out.len * 4)); + } + + fn args(self: *Fexp, slot_n: usize, w: u64, p: Plan, out: u64) cuda.Args { + const sh = self.shape; + return cuda.experts.Affine4Experts.argsOf(self.dx, sh.k, slot_n, w, sh.k / 32, sh.n / 32, .{ .items = p.items, .counts = p.counts, .members = p.members, .rank = 0, .hist = 0 }, out, sh.n); + } + + /// RowsExperts.run's launch: items of 32 pairs, four column blocks a warp's share, at most 4096 blocks. + fn rows(self: *Fexp, f: cuda.Function, slot_n: usize, w: u64, out: u64) !void { + var a = self.args(slot_n, w, self.plan32, out); + const nb = self.shape.n / 32; + const grid = @min((self.plan32.n * nb + 3) / 4, 4096); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(grid) }, .block = .{ .x = 128 } }, self.rig.s, &a); + } + + /// cuda_batch.zig's fexp launch: items of 64 pairs, wn column blocks a CTA. + fn launch(self: *Fexp, out: u64) !void { + const wn: usize = if (self.kind == .up) 2 else 4; + const f = if (self.kind == .up) self.f.up else self.f.down; + var a = self.args(self.slots(), self.dw, self.plan64, out); + const nb = self.shape.n / 32; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(self.plan64.n * ((nb + wn - 1) / wn)) }, .block = .{ .x = 256 }, .shared = tk.fexp_smem }, self.rig.s, &a); + } + + pub fn run(self: *Fexp) !void { + try self.launch(self.out); + try self.rig.s.synchronize(); + } + + /// expert_rows on the same inputs and pairs (its own plan: items of 32 over the same members). + pub fn replaced(self: *Fexp) ![]const u8 { + const esize = c.size(spec(self.kind).out); + try self.rows(if (self.kind == .up) self.f.rows_up else self.f.rows_down, self.slots(), self.dw, self.base); + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, self.count() * esize); + } + + pub fn count(self: *Fexp) usize { + return self.shape.rows * self.shape.n; + } + + /// SwiGLU on the bf16 gate and up values in float64 (the kernel's rounding of silu is the spec's one mid). + fn swiglu(gv: f64, uv: f64) f64 { + return gv / (1 + @exp(-gv)) * uv; + } + + fn bf(v: f64) f64 { + return c.fromBf16(c.toBf16(@floatCast(v))); + } + + pub fn reference(self: *Fexp, want: []f64, mag: []f64) !void { + const n = self.shape.n; + if (self.kind == .up) { + for (want, mag, self.gs, self.us) |*w, *m, gv, uv| { + w.* = swiglu(bf(gv), bf(uv)); + m.* = 0; + } + return; + } + for (0..self.shape.rows) |p| for (0..n) |j| { + const o = self.aff.items[self.expert[p]].one(p, j); + want[p * n + j] = o[0]; + mag[p * n + j] = o[1]; + }; + } + + /// Each matrix's sums as an accumulator leaves them (the affine family's structures), per pair from its expert. + pub fn simulate(self: *Fexp, acc: c.Accum, order: g.Order, out: []f64) !void { + const a = self.a; + const n = self.shape.n; + const t = self.tokens(); + const all = try a.alloc(f64, self.aff.items.len * t * n); + defer a.free(all); + for (self.aff.items, 0..) |f, m| try f.simulate(acc, order, all[m * t * n ..][0 .. t * n]); + for (0..self.shape.rows) |p| for (0..n) |j| { + const at = self.tokenOf(p) * n + j; + const s0 = all[self.expert[p] * t * n + at]; + out[p * n + j] = if (self.kind == .up) swiglu(bf(s0), bf(all[(experts_n + self.expert[p]) * t * n + at])) else s0; + }; + } +}; diff --git a/zig/tests/cuda/contract_flashnext_glue.zig b/zig/tests/cuda/contract_flashnext_glue.zig new file mode 100644 index 0000000000..60469d399f --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_glue.zig @@ -0,0 +1,346 @@ +//! Flash Next's prompt glue replacing #952's: the DeltaNet chain, the indexer scores and the hyper-connection read-out. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const parts = flash.batch_state; +const tk = flash.tile_kernels; + +pub const Kind = enum { chain, scores, readout }; + +const nk = 16; +const nv = 48; +const dk = 128; +const dv = 128; +const di = 128; +const heads = 4; +const ratio = 4; +const top = 512; +const dims = 2560; +const streams = 4; +const wide = streams * dims; +const low = 320; +const n_down = low + streams; +const hc_sk = 32; +const nc = dims / 256; +const eps: f32 = 1e-6; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// rows = token rows; n = outputs a row (value heads x width, pooled blocks, or the read-out's mixed dims). +pub fn spec(kind: Kind) c.Spec { + return switch (kind) { + .chain => .{ .name = "cuda flashnext DeltaNet prompt chain fn_gdn_chain_tpr", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "ys and the state bit-equal to fn_gdn_rows_to's, non-finite inputs included" }, .replaces = "fn_gdn_rows_to", .tiles = .{ .{ .size = 16 }, .{}, .{} }, .served = &.{ .{ .rows = 64, .n = nv * dv, .k = dk }, .{ .rows = 300, .n = nv * dv, .k = dk } } }, + .scores => .{ .name = "cuda flashnext indexer scores qsa_tile::scores<8, 4>", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "every score byte equal to fn_qsa_scores', non-finite inputs included" }, .replaces = "fn_qsa_scores", .tiles = .{ .{ .size = 8 }, .{ .size = 256, .step = 1, .min = 1 }, .{} }, .served = &.{ .{ .rows = 64, .n = 600, .k = di }, .{ .rows = 300, .n = 2100, .k = di } } }, + .readout => .{ .name = "cuda flashnext hyper-connection read-out on prompt tiles (hcf down, reduce, hcf up, rinv, mix_h)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "act, xs, inj, mixed and xsm bit-equal to #952's chain, non-finite inputs included" }, .replaces = "fn_hc_down, fn_hc_reduce_act and fn_hc_upmix", .tiles = .{ .{ .size = 64 }, .{}, .{} }, .served = &.{ .{ .rows = 64, .n = dims, .k = wide }, .{ .rows = 130, .n = dims, .k = wide } } }, + }; +} + +pub const Fns = struct { tpr: cuda.Function, rows_to: cuda.Function, tile: cuda.Function, scores: cuda.Function, hcf_down: cuda.Function, hcf_up: cuda.Function }; + +pub const Glue = struct { + rig: rig_.Rig, + kind: Kind, + f: *const Fns, + glue: *const flash.Glue, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + ins: [12]u64 = @splat(0), + outs: [5]u64 = @splat(0), + lens: [5]usize = @splat(0), + nout: usize = 0, + base: u64 = 0, + pos0: usize = 0, + seen: std.ArrayList(u8) = .empty, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Glue) void { + self.seen.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *Glue) g.Kernel { + return .{ .spec = spec(self.kind), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Glue); + + fn rng(seed: u64) std.Random.DefaultPrng { + return std.Random.DefaultPrng.init(seed ^ 0x676c7565); + } + + fn f32s(self: *Glue, n: usize, r: std.Random, lo: f32, hi: f32) ![]f32 { + const v = try self.rig.gpu.gpa.alloc(f32, n); + for (v) |*x| x.* = lo + r.float(f32) * (hi - lo); + return v; + } + + fn bfs(self: *Glue, n: usize, r: std.Random, lo: f32, hi: f32) ![]u16 { + const v = try self.rig.gpu.gpa.alloc(u16, n); + for (v) |*x| x.* = c.toBf16(lo + r.float(f32) * (hi - lo)); + return v; + } + + fn in(self: *Glue, i: usize, bytes: []const u8) !void { + self.ins[i] = try self.rig.input(bytes); + } + + fn out(self: *Glue, len: usize) !void { + self.outs[self.nout] = try self.rig.output(len); + self.lens[self.nout] = len; + self.nout += 1; + } + + pub fn prepare(self: *Glue, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + self.nout = 0; + self.shape = shape; + if (shape.rows == 0) return false; + var prng = rng(seed); + const r = prng.random(); + const a = self.rig.gpu.gpa; + const bad = fill == .nonfinite; + switch (self.kind) { + .chain => { + if (shape.n != nv * dv or shape.k != dk) return false; + const m = shape.rows; + const qn = try self.f32s(m * nk * dk, r, -0.1, 0.1); + defer a.free(qn); + const kn = try self.f32s(m * nk * dk, r, -0.1, 0.1); + defer a.free(kn); + const v = try self.f32s(m * nv * dv, r, -1, 1); + defer a.free(v); + const gb = try self.f32s(m * nv * 2, r, 0, 1); + defer a.free(gb); + for (0..m * nv) |i| gb[2 * i] = -gb[2 * i]; // a log decay, then beta in (0, 1) + const st = try self.f32s(nv * dv * dk, r, -0.5, 0.5); + defer a.free(st); + if (bad) { + qn[3] = std.math.nan(f32); + v[(m / 2) * nv * dv + 5] = std.math.inf(f32); + st[dv * dk + 7] = -std.math.inf(f32); + } + try self.in(0, std.mem.sliceAsBytes(qn)); + try self.in(1, std.mem.sliceAsBytes(kn)); + try self.in(2, std.mem.sliceAsBytes(v)); + try self.in(3, std.mem.sliceAsBytes(gb)); + try self.in(4, std.mem.sliceAsBytes(st)); + try self.out(m * nv * dv * 4); + try self.out(nv * dv * dk * 4); + }, + .scores => { + if (shape.k != di or shape.n < 2 or shape.n > 4096) return false; + const m = shape.rows; + const nb = shape.n; + // the last row sees nb complete blocks of `ratio`; rows past `top` blocks score + if (nb * ratio < m) return false; + self.pos0 = nb * ratio - m; + if ((self.pos0 + m) / ratio <= top) return false; + const iq = try self.bfs(m * heads * di, r, -1, 1); + defer a.free(iq); + const pooled = try self.bfs(nb * di, r, -1, 1); + defer a.free(pooled); + if (bad) { + iq[7] = c.toBf16(std.math.nan(f32)); + pooled[di * 3 + 1] = c.toBf16(std.math.inf(f32)); + } + const pos = [4]i32{ @intCast(self.pos0), 0, 0, 0 }; + try self.in(0, std.mem.sliceAsBytes(iq)); + try self.in(1, std.mem.sliceAsBytes(pooled)); + try self.in(2, std.mem.sliceAsBytes(&pos)); + try self.out(m * nb * 4); + }, + .readout => { + if (shape.n != dims or shape.k != wide) return false; + try self.readoutInputs(shape.rows, r, bad); + }, + } + self.base = 0; + var total: usize = 0; + for (self.lens[0..self.nout]) |l| total += l; + self.base = try self.rig.scratch(total); + return true; + } + + /// The streams and their sums of squares, the norm's scale, the down rows (inject stacked) and the up rows. + fn readoutInputs(self: *Glue, m: usize, r: std.Random, bad: bool) !void { + const a = self.rig.gpu.gpa; + const h = try self.bfs(m * wide, r, -1, 1); + defer a.free(h); + if (bad) h[11] = c.toBf16(std.math.nan(f32)); + const pss = try a.alloc(f32, m * nc * streams); + defer a.free(pss); + for (0..m) |row| for (0..nc) |cc| for (0..streams) |s| { + var t: f32 = 0; + for (0..256) |i| { + const v: f32 = @floatCast(c.fromBf16(h[row * wide + s * dims + cc * 256 + i])); + t += v * v; + } + pss[(row * nc + cc) * streams + s] = t; + }; + const scale = try self.f32s(wide, r, 0.5, 1.5); + defer a.free(scale); + try self.in(0, std.mem.sliceAsBytes(h)); + try self.in(1, std.mem.sliceAsBytes(pss)); + try self.in(2, std.mem.sliceAsBytes(scale)); + for ([_][2]usize{ .{ n_down, wide }, .{ wide, low } }, 0..) |nk_, wi| { + const n = nk_[0]; + const k = nk_[1]; + const words = try a.alloc(u32, n * k / 8); + defer a.free(words); + for (words) |*w| w.* = r.int(u32); + const sc = try self.bfs(n * k / 32, r, 0.001, if (wi == 0) 0.004 else 0.05); + defer a.free(sc); + const bi = try self.bfs(n * k / 32, r, -0.01, if (wi == 0) 0.01 else 0.1); + defer a.free(bi); + const tw = try flash.hc.tileWords(a, std.mem.sliceAsBytes(words), n, k / 8); + defer a.free(tw); + const ts = try flash.hc.transposeBf16(a, std.mem.sliceAsBytes(sc), n, k / 32); + defer a.free(ts); + const tb = try flash.hc.transposeBf16(a, std.mem.sliceAsBytes(bi), n, k / 32); + defer a.free(tb); + try self.in(3 + wi * 3, tw); + try self.in(4 + wi * 3, ts); + try self.in(5 + wi * 3, tb); + } + // act, xs, inj (the reduce's), then mixed and xsm (the mix's) + try self.out(m * low * 2); + try self.out(m * (low / 32) * 4); + try self.out(m * streams * 2); + try self.out(m * dims * 2); + try self.out(m * (dims / 32) * 4); + } + + fn sweepOf(m: usize, bm: usize, k: usize) usize { + return @max(@as(usize, 1), @min((m + bm - 1) / bm, (12 << 20) / (bm * k * 2))); + } + + /// Our launches as the pass makes them, into `o` (the outputs in order). + fn ours(self: *Glue, o: []const u64) !void { + const s = self.rig.s; + const m = self.shape.rows; + switch (self.kind) { + .chain => { + var a: cuda.Args = .{}; + for ([_]u64{ self.ins[0], self.ins[1], self.ins[2], self.ins[3], self.ins[4], o[1] }) |v| a.add(v); + a.add(@as(c_int, @intCast(m))); + a.add(o[0]); + try cuda.launch.launch(self.f.tpr, .{ .grid = .{ .x = nv }, .block = .{ .x = 256 }, .shared = parts.chain_tpr_smem }, s, &a); + }, + .scores => { + var a: cuda.Args = .{}; + for ([_]u64{ self.ins[0], self.ins[1], self.ins[2], o[0] }) |v| a.add(v); + for ([_]usize{ self.shape.n, ratio, top, m }) |v| a.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(self.f.tile, .{ .grid = .{ .x = @intCast((m + 7) / 8), .y = @intCast((self.shape.n + 255) / 256) }, .block = .{ .x = 256 } }, s, &a); + }, + .readout => { + const part = try self.rig.scratch(m * n_down * 4); + const rinv = try self.rig.scratch(m * streams * 4); + const up = try self.rig.scratch(m * wide * 2); + var a: cuda.Args = .{}; + for ([_]u64{ self.ins[0], 0, self.ins[3], self.ins[4], self.ins[5], part }) |v| a.add(v); + for ([_]usize{ m, n_down, wide, hc_sk, wide, sweepOf(m, 32, wide) }) |v| a.add(@as(c_int, @intCast(v))); + a.add(self.ins[1]); + a.add(self.ins[2]); + a.add(eps); + for ([_]usize{ dims, nc, streams }) |v| a.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(self.f.hcf_down, .{ .grid = .{ .x = @intCast((m + 31) / 32 * ((n_down + 191) / 192)) }, .block = .{ .x = 256 }, .shared = tk.hcf_down_smem }, s, &a); + try self.glue.hcReduceAct(s, part, o[0], o[1], o[2], 1, m, streams, low, n_down, 1); + var b: cuda.Args = .{}; + for ([_]u64{ o[0], o[1], self.ins[6], self.ins[7], self.ins[8], up }) |v| b.add(v); + for ([_]usize{ m, wide, low, 1, low, sweepOf(m, 64, low) }) |v| b.add(@as(c_int, @intCast(v))); + b.add(@as(u64, 0)); + b.add(@as(u64, 0)); + b.add(@as(f32, 0)); + for ([_]usize{ dims, nc, streams }) |v| b.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(self.f.hcf_up, .{ .grid = .{ .x = @intCast((m + 63) / 64 * (wide / 128)) }, .block = .{ .x = 256 }, .shared = tk.hcf_up_smem }, s, &b); + try self.glue.hcRinv(s, self.ins[1], rinv, eps, m, dims, nc, streams); + try self.glue.hcMixH(s, up, self.ins[0], rinv, self.ins[2], o[3], o[4], m, dims, streams); + }, + } + } + + /// #952's launches on the same inputs, into `o`. + fn theirs(self: *Glue, o: []const u64) !void { + const s = self.rig.s; + const m = self.shape.rows; + switch (self.kind) { + .chain => { + var a: cuda.Args = .{}; + for ([_]u64{ self.ins[0], self.ins[1], self.ins[2], self.ins[3], self.ins[4], o[1] }) |v| a.add(v); + a.add(@as(c_int, @intCast(m))); + a.add(o[0]); + try cuda.launch.launch(self.f.rows_to, .{ .grid = .{ .x = nv, .y = dv / 16 }, .block = .{ .x = 128 } }, s, &a); + }, + .scores => { + var a: cuda.Args = .{}; + for ([_]u64{ self.ins[0], self.ins[1], self.ins[2], o[0] }) |v| a.add(v); + for ([_]usize{ self.shape.n, heads, ratio, top }) |v| a.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(self.f.scores, .{ .grid = .{ .x = @intCast(m), .y = @intCast((self.shape.n + 63) / 64) }, .block = .{ .x = 256 } }, s, &a); + }, + .readout => { + const normed = try self.rig.scratch(m * wide * 2); + const dn = try self.rig.scratch(m * (low + 4) * 2); + const part = try self.rig.scratch(hc_sk * m * n_down * 4); + try self.glue.hcDown(s, self.ins[0], self.ins[1], self.ins[2], normed, self.ins[3], self.ins[4], self.ins[5], dn, part, eps, m, n_down, wide, dims, nc, streams, hc_sk); + try self.glue.hcReduceAct(s, part, o[0], o[1], o[2], hc_sk, m, streams, low, n_down, 1); + try self.glue.hcUpmix(s, o[0], o[1], self.ins[6], self.ins[7], self.ins[8], normed, o[3], o[4], m, wide, low, dims, streams); + }, + } + } + + pub fn run(self: *Glue) !void { + try self.ours(self.outs[0..self.nout]); + try self.rig.s.synchronize(); + } + + /// The scratch the baseline writes: the outputs' layout back to back, poisoned like them (0xff). + fn baseline(self: *Glue) ![]const u8 { + var at = self.base; + var o: [5]u64 = undefined; + var total: usize = 0; + for (self.lens[0..self.nout], 0..) |l, i| { + o[i] = at; + at += l; + total += l; + } + const ff = try self.rig.gpu.gpa.alloc(u8, total); + defer self.rig.gpu.gpa.free(ff); + @memset(ff, 0xff); + const b: cuda.DeviceBuffer = .{ .d = self.rig.gpu.d, .ptr = self.base, .len = total }; + try b.upload(0, ff); + try self.theirs(o[0..self.nout]); + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, total); + } + + pub fn replaced(self: *Glue) ![]const u8 { + return self.baseline(); + } + + /// Gate 1's stated rule: the outputs on the planted inputs equal #952's bytes. + pub fn check(self: *Glue, gate: g.Gate) !?[]const u8 { + _ = gate; + self.seen.clearRetainingCapacity(); + try self.seen.appendSlice(self.rig.gpu.gpa, try self.rig.outputBytes()); + const want = try self.baseline(); + if (std.mem.eql(u8, self.seen.items, want)) return null; + return try std.fmt.bufPrint(&self.msg, "bytes differ from #952's from byte {d}", .{std.mem.indexOfDiff(u8, self.seen.items, want).?}); + } + + pub fn count(self: *Glue) usize { + var total: usize = 0; + for (self.lens[0..self.nout]) |l| total += l; + return total; + } + + pub fn reference(_: *Glue, _: []f64, _: []f64) !void { + return error.NoFloat64Reference; + } +}; diff --git a/zig/tests/cuda/contract_route_rows.zig b/zig/tests/cuda/contract_route_rows.zig new file mode 100644 index 0000000000..e73f8fbc17 --- /dev/null +++ b/zig/tests/cuda/contract_route_rows.zig @@ -0,0 +1,120 @@ +//! route_rows' router logits under the contract: tf_route_logits against float64, its other row forms bit-equal. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; +const rr = cuda.route_rows; + +pub const Kind = enum { logits, few, wide }; + +/// One fp32 MMA chain over the inputs from zero. +fn chain(s: c.Shape) usize { + return s.k / 16 + 16; +} + +/// Flash Next's router: 2560 inputs, 512 experts and the shared one. +const fn_e = 513; +const fn_d = 2560; + +pub fn spec(kind: Kind) c.Spec { + return switch (kind) { + .logits => .{ .name = "cuda route_rows tf_route_logits", .accum = .f32, .depth = chain, .out = .f32, .tiles = .{ .{ .size = 16 }, .{ .size = 32, .step = 1, .min = 1 }, .{ .size = 128, .step = 16, .min = 16 } }, .served = &.{ .{ .rows = 32, .n = fn_e, .k = fn_d }, .{ .rows = 17, .n = fn_e, .k = fn_d }, .{ .rows = 63, .n = fn_e, .k = fn_d } } }, + .few => .{ .name = "cuda route_rows tf_route_logits_few", .accum = .f32, .depth = chain, .out = .f32, .replaces = "tf_route_logits", .tiles = .{ .{ .size = 16 }, .{ .size = 32, .step = 1, .min = 1 }, .{ .size = 256, .step = 16, .min = 16 } }, .served = &.{ .{ .rows = 16, .n = fn_e, .k = fn_d }, .{ .rows = 1, .n = fn_e, .k = fn_d } } }, + .wide => .{ .name = "cuda route_rows tf_route_logits_wide", .accum = .f32, .depth = chain, .out = .f32, .replaces = "tf_route_logits", .tiles = .{ .{ .size = 128 }, .{ .size = 64, .step = 1, .min = 1 }, .{ .size = 64, .step = 16, .min = 16 } }, .served = &.{ .{ .rows = 64, .n = fn_e, .k = fn_d }, .{ .rows = 300, .n = fn_e, .k = fn_d } } }, + }; +} + +pub const Logits = struct { + rig: rig_.Rig, + kind: Kind, + route: *const rr.RouteRows, + d: ?Dense = null, + x: u64 = 0, + w: u64 = 0, + base: u64 = 0, + out: u64 = 0, + + pub fn deinit(self: *Logits) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Logits) g.Kernel { + return .{ .spec = spec(self.kind), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Logits); + + pub fn prepare(self: *Logits, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 16 != 0) return false; + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 1); + const w = try a.alloc(u16, shape.n * shape.k); + defer a.free(w); + var prng = std.Random.DefaultPrng.init(seed ^ 0x726f7574); + const rng = prng.random(); + const k = shape.k; + for (0..shape.n) |j| for (0..k) |i| { + // gate 4: integers mirrored into the second half, so fp32 sums are exact in any order + const v: f32 = if (fill != .stress) (rng.float(f32) - 0.5) * 0.1 else if (i < k / 2) @floatFromInt(1 + rng.uintLessThan(u32, 7)) else -c.fromBf16(w[j * k + i - k / 2]); + w[j * k + i] = c.toBf16(v); + d.w[j * k + i] = c.fromBf16(w[j * k + i]); + }; + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.w = try self.rig.input(std.mem.sliceAsBytes(w)); + self.out = try self.rig.output(shape.rows * shape.n * 4); + if (self.kind != .logits) self.base = try self.rig.scratch(shape.rows * shape.n * 4); + return true; + } + + pub fn run(self: *Logits) !void { + try self.launch(self.kind, self.out); + try self.rig.s.synchronize(); + } + + /// RouteRows.logits' launch of one form, whatever the row count. + fn launch(self: *Logits, kind: Kind, out: u64) !void { + const sh = self.d.?.shape; + var a: cuda.Args = .{}; + a.add(self.x); + a.add(@as(i32, @intCast(sh.k))); + a.add(self.w); + a.add(out); + for ([_]usize{ sh.rows, sh.k, sh.n }) |v| a.add(@as(i32, @intCast(v))); + const r = self.route; + switch (kind) { + .wide => try cuda.launch.launch(r.wide_fn, .{ .grid = .{ .x = @intCast((sh.rows + 127) / 128), .y = @intCast((sh.n + 63) / 64) }, .block = .{ .x = 256 }, .shared = rr.wide_shared }, self.rig.s, &a), + .few => try cuda.launch.launch(r.few_fn, .{ .grid = .{ .x = @intCast((sh.rows + 15) / 16), .y = @intCast((sh.n + 31) / 32) }, .block = .{ .x = 128 }, .shared = rr.few_shared }, self.rig.s, &a), + .logits => try cuda.launch.launch(r.logits_fn, .{ .grid = .{ .x = @intCast((sh.rows + 15) / 16), .y = @intCast((sh.n + 31) / 32) }, .block = .{ .x = 128 }, .shared = rr.many_shared }, self.rig.s, &a), + } + } + + /// tf_route_logits on the same inputs. + pub fn replaced(self: *Logits) ![]const u8 { + if (self.kind == .logits) return error.NotAReplacement; + try self.launch(.logits, self.base); + return self.rig.bytes(self.base, self.count() * 4); + } + + pub fn count(self: *Logits) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *Logits, want: []f64, mag: []f64) !void { + return self.d.?.reference(want, mag); + } + + pub fn simulate(self: *Logits, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.d.?.simulate(acc, order, out); + } +}; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index d5d52993cc..58bd17a101 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -23,6 +23,10 @@ const ck_k = @import("contract_ck.zig"); const scan_k = @import("contract_scan.zig"); const lane8_k = @import("contract_ck_lane.zig"); const latent_k = @import("contract_latent.zig"); +const fact_k = @import("contract_flashnext_fact.zig"); +const route_k = @import("contract_route_rows.zig"); +const fexp_k = @import("contract_flashnext_fexp.zig"); +const glue_k = @import("contract_flashnext_glue.zig"); const contract = core.contract; const g = contract.gates; @@ -41,7 +45,7 @@ const h3_image = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; - inline for (.{ "norms", "moe", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "lane_gemv", "prompt16", "qmmf", "ck_experts", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "nemotron_mamba", "sample", "qwen35", "qwen35_keep", "qwen35_kv8", "qwen35_draft", "ck_lane" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "moe", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "lane_gemv", "prompt16", "qmmf", "ck_experts", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "nemotron_mamba", "sample", "qwen35", "qwen35_keep", "qwen35_kv8", "qwen35_draft", "ck_lane", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } @@ -228,6 +232,68 @@ pub fn run(gpu: check.Gpu) !void { nlane += 1; } } + const fparts = @import("flashnext").batch_state; + const ftk = @import("flashnext").tile_kernels; + var qmm_mod = try cuda.Module.load(gpu.d, kk.qmm); + defer qmm_mod.unload(); + var fact_mod = try cuda.Module.load(gpu.d, kk.flashnext_fact); + defer fact_mod.unload(); + const flane = try qmm_mod.function(fparts.cubin_qmm); + try flane.allowDynamicShared(fparts.qmm_smem); + const fprompt = try fact_mod.function(ftk.cubin_fact); + try fprompt.allowDynamicShared(ftk.fact_smem); + const ffold = try fact_mod.function(ftk.cubin_fact_fold); + try ffold.allowDynamicShared(ftk.fact_smem); + var facts: [2]fact_k.Fact = undefined; + for (std.enums.values(fact_k.Kind), &facts) |kind, *m| m.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .lane = flane, .prompt = fprompt, .fold = ffold, .sum = try batch.function("fn_sum_slices") }; + defer for (&facts) |*m| m.deinit(); + for (&facts) |*m| try kernels.append(a, m.kernel()); + var route = try cuda.route_rows.RouteRows.load(gpu.d); + defer route.unload(); + var routes: [3]route_k.Logits = undefined; + for (std.enums.values(route_k.Kind), &routes) |kind, *r| r.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .route = &route }; + defer for (&routes) |*r| r.deinit(); + for (&routes) |*r| try kernels.append(a, r.kernel()); + var rows_mod = try cuda.Module.load(gpu.d, kk.experts_rows); + defer rows_mod.unload(); + var fexp_mod = try cuda.Module.load(gpu.d, kk.flashnext_fexp); + defer fexp_mod.unload(); + var pack_mod = try cuda.Module.load(gpu.d, kk.experts_pack); + defer pack_mod.unload(); + const fexp_fns: fexp_k.Fns = .{ + .up = try fexp_mod.function(ftk.cubin_fexp_up), + .down = try fexp_mod.function(ftk.cubin_fexp_down), + .rows_up = try rows_mod.function(cuda.experts.RowsExperts.symbols.up), + .rows_down = try rows_mod.function(cuda.experts.RowsExperts.symbols.down), + .pack = try pack_mod.function(fexp_k.pack_symbol), + }; + try fexp_fns.up.allowDynamicShared(ftk.fexp_smem); + try fexp_fns.down.allowDynamicShared(ftk.fexp_smem); + var fexps: [2]fexp_k.Fexp = undefined; + for (std.enums.values(fexp_k.Kind), &fexps) |kind, *m| m.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .f = &fexp_fns }; + defer for (&fexps) |*m| m.deinit(); + for (&fexps) |*m| try kernels.append(a, m.kernel()); + var attn = try cuda.Module.load(gpu.d, kk.flashnext_attn); + defer attn.unload(); + var hcf_mod = try cuda.Module.load(gpu.d, kk.flashnext_hcf); + defer hcf_mod.unload(); + var fglue = try @import("flashnext").Glue.load(gpu.d); + defer fglue.deinit(); + const glue_fns: glue_k.Fns = .{ + .tpr = try batch.function("fn_gdn_chain_tpr"), + .rows_to = try batch.function("fn_gdn_rows_to"), + .tile = try attn.function("_ZN8qsa_tile6scoresILi8ELi4EEEvPK13__nv_bfloat16S3_PKiPfiiii"), + .scores = try attn.function("fn_qsa_scores"), + .hcf_down = try hcf_mod.function(ftk.cubin_hcf_down), + .hcf_up = try hcf_mod.function(ftk.cubin_hcf_up), + }; + try glue_fns.tpr.allowDynamicShared(fparts.chain_tpr_smem); + try glue_fns.hcf_down.allowDynamicShared(ftk.hcf_down_smem); + try glue_fns.hcf_up.allowDynamicShared(ftk.hcf_up_smem); + var glues: [3]glue_k.Glue = undefined; + for (std.enums.values(glue_k.Kind), &glues) |kind, *m| m.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .f = &glue_fns, .glue = &fglue }; + defer for (&glues) |*m| m.deinit(); + for (&glues) |*m| try kernels.append(a, m.kernel()); var t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t); From 055baedb296863c4e912789dd3509b3aeee1385d Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 20:42:34 +0200 Subject: [PATCH 13/16] Kernel contract: Flash Next's tile records from nvcc 13.0 in the base, 13.2's and 13.3's where they differ --- .../core/contract/flashnext_tile_records.zig | 32 ++++++++++++++----- zig/src/core/contract/registry.zig | 18 +++++------ 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/zig/src/core/contract/flashnext_tile_records.zig b/zig/src/core/contract/flashnext_tile_records.zig index a8a2f148a3..080f7b60d5 100644 --- a/zig/src/core/contract/flashnext_tile_records.zig +++ b/zig/src/core/contract/flashnext_tile_records.zig @@ -1,18 +1,34 @@ -//! Flash Next's prompt tiles' CUDA records (gb10), in registry.cuda: from nvcc 13.2, to re-derive with 13.0 at landing. -const CudaRecord = @import("registry.zig").CudaRecord; +//! Flash Next's prompt tiles' CUDA records: nvcc 13.0's in registry.cuda, then 13.2's and 13.3's where they differ. +const registry = @import("registry.zig"); +const CudaRecord = registry.CudaRecord; pub const cuda = [_]CudaRecord{ - .{ .launch = .{ .name = "flashnext_fact factored prompt matmul 64x128", .symbol = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii", .image = "flashnext_fact", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_fact factored prompt matmul 64x128", .symbol = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii", .image = "flashnext_fact", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 127, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_fact factored prompt matmul 64x128 folded", .symbol = "_ZN4fact11fact_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELi2ELb1ELb1ELb1EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii", .image = "flashnext_fact", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_hcf hyper-connection down 32x192 folded", .symbol = "_ZN3hcf10hcf_kernelILi32ELi192ELi2ELi4ELi4ELi2ELb1ELb1ELb0ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii", .image = "flashnext_hcf", .threads = 256, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 126, .static = 0, .bound = 256 }, .{ .regs = 126, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_hcf hyper-connection up 64x128", .symbol = "_ZN3hcf10hcf_kernelILi64ELi128ELi2ELi4ELi3ELi2ELb0ELb0ELb1ELi2EEEvPK13__nv_bfloat16PKfPK5uint4S3_S3_PviiiiiiS5_S5_fiii", .image = "flashnext_hcf", .threads = 256, .dynamic = 43264, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_fexp factored experts gate and up", .symbol = "_ZN15tf_fact_experts11fact_kernelILi2ELi2ELi1ELi4ELi2ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif", .image = "flashnext_fexp", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 125, .static = 0, .bound = 256 }, .{ .regs = 125, .static = 0, .bound = 256 }, .{ .regs = 125, .static = 0, .bound = 256 }, .{ .regs = 116, .static = 0, .bound = 256 }, .{ .regs = 116, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "flashnext_fexp factored experts down", .symbol = "_ZN15tf_fact_experts11fact_kernelILi1ELi0ELi2ELi2ELi4ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif", .image = "flashnext_fexp", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 120, .static = 0, .bound = 256 }, .{ .regs = 120, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "flashnext_batch fn_gdn_chain_tpr", .symbol = "fn_gdn_chain_tpr", .image = "flashnext_batch", .threads = 256, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 152, .static = 0, .bound = 256 }, .{ .regs = 152, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "flashnext_attn qsa_tile scores", .symbol = "_ZN8qsa_tile6scoresILi8ELi4EEEvPK13__nv_bfloat16S3_PKiPfiiii", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 224, .static = 16384, .bound = 256 }, .{ .regs = 224, .static = 16384, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_fexp factored experts down", .symbol = "_ZN15tf_fact_experts11fact_kernelILi1ELi0ELi2ELi2ELi4ELi3EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif", .image = "flashnext_fexp", .threads = 256, .dynamic = 41472, .opt_in = true }, .usage = .{ .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 118, .static = 0, .bound = 256 }, .{ .regs = 118, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_chain_tpr", .symbol = "fn_gdn_chain_tpr", .image = "flashnext_batch", .threads = 256, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 154, .static = 0, .bound = 256 }, .{ .regs = 154, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn qsa_tile scores", .symbol = "_ZN8qsa_tile6scoresILi8ELi4EEEvPK13__nv_bfloat16S3_PKiPfiiii", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 215, .static = 16384, .bound = 256 }, .{ .regs = 216, .static = 16384, .bound = 256 }, .{ .regs = 216, .static = 16384, .bound = 256 }, .{ .regs = 224, .static = 16384, .bound = 256 }, .{ .regs = 224, .static = 16384, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_fused", .symbol = "fn_attn_fused", .image = "flashnext_attn", .threads = 128, .dynamic = 26624 }, .usage = .{ .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 121, .static = 640, .bound = 128 }, .{ .regs = 121, .static = 640, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_one", .symbol = "fn_attn_chunks_one", .image = "flashnext_attn", .threads = 128, .dynamic = 26624, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 96, .static = 640, .bound = 128 }, .{ .regs = 96, .static = 640, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_hc fn_hc_rinv", .symbol = "fn_hc_rinv", .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 = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_mix_h", .symbol = "fn_hc_mix_h", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 38, .static = 1024, .bound = 256 }, .{ .regs = 38, .static = 1024, .bound = 256 } } }, - .{ .launch = .{ .name = "route_rows tf_route_logits_wide", .symbol = "tf_route_logits_wide", .image = "route_rows", .threads = 256, .dynamic = 55296, .opt_in = true }, .usage = .{ .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 58, .static = 0, .bound = 256 }, .{ .regs = 58, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_mix_h", .symbol = "fn_hc_mix_h", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 40, .static = 1024, .bound = 256 }, .{ .regs = 40, .static = 1024, .bound = 256 } } }, + .{ .launch = .{ .name = "route_rows tf_route_logits_wide", .symbol = "tf_route_logits_wide", .image = "route_rows", .threads = 256, .dynamic = 55296, .opt_in = true }, .usage = .{ .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 61, .static = 0, .bound = 256 }, .{ .regs = 61, .static = 0, .bound = 256 } } }, +}; + +/// nvcc 13.2's records where it compiles these tiles differently from 13.0. +pub const cuda_132 = [_]registry.Override{ + .{ .name = "flashnext_fact factored prompt matmul 64x128", .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, + .{ .name = "flashnext_fexp factored experts down", .usage = .{ .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 136, .static = 0, .bound = 256 }, .{ .regs = 120, .static = 0, .bound = 256 }, .{ .regs = 120, .static = 0, .bound = 256 } } }, + .{ .name = "flashnext_batch fn_gdn_chain_tpr", .usage = .{ .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 129, .static = 0, .bound = 256 }, .{ .regs = 152, .static = 0, .bound = 256 }, .{ .regs = 152, .static = 0, .bound = 256 } } }, + .{ .name = "flashnext_attn qsa_tile scores", .usage = .{ .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 224, .static = 16384, .bound = 256 }, .{ .regs = 224, .static = 16384, .bound = 256 } } }, + .{ .name = "flashnext_hc fn_hc_mix_h", .usage = .{ .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 38, .static = 1024, .bound = 256 }, .{ .regs = 38, .static = 1024, .bound = 256 } } }, + .{ .name = "route_rows tf_route_logits_wide", .usage = .{ .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 58, .static = 0, .bound = 256 }, .{ .regs = 58, .static = 0, .bound = 256 } } }, +}; + +/// nvcc 13.3's records where it compiles these tiles differently from 13.2. +pub const cuda_133 = [_]registry.Override{ + .{ .name = "flashnext_attn fn_attn_fused", .usage = .{ .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 118, .static = 640, .bound = 128 }, .{ .regs = 122, .static = 640, .bound = 128 }, .{ .regs = 122, .static = 640, .bound = 128 } } }, }; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 589256fe16..31c5e029c9 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -21,13 +21,13 @@ 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 = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .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 = 70, .static = 768, .bound = 128 }, .{ .regs = 70, .static = 768, .bound = 128 }, .{ .regs = 70, .static = 768, .bound = 128 }, .{ .regs = 73, .static = 768, .bound = 128 }, .{ .regs = 73, .static = 768, .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_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_norm", .symbol = "fn_gdn_norm", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 24, .static = 528, .bound = 128 }, .{ .regs = 26, .static = 528, .bound = 128 }, .{ .regs = 26, .static = 528, .bound = 128 }, .{ .regs = 24, .static = 528, .bound = 128 }, .{ .regs = 24, .static = 528, .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 } } }, .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows_to", .symbol = "fn_gdn_rows_to", .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 } } }, .{ .launch = .{ .name = "flashnext_hc fn_hc_dense_mix", .symbol = "fn_hc_dense_mix", .image = "flashnext_hc", .threads = 32 }, .usage = .{ .{ .regs = 22, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 } } }, @@ -42,7 +42,7 @@ pub const cuda = [_]CudaRecord{ .{ .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_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_batch fn_qsa_select", .symbol = "fn_qsa_select", .image = "flashnext_batch", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 }, .{ .regs = 40, .static = 1168, .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 } } }, .{ .launch = .{ .name = "flashnext_hc fn_shared_scatter", .symbol = "fn_shared_scatter", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "flashnext_hc fn_shared_swiglu", .symbol = "fn_shared_swiglu", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 17, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 } } }, @@ -50,7 +50,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "flashnext_batch fn_sum_slices", .symbol = "fn_sum_slices", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 38, .static = 0, .bound = 0 }, .{ .regs = 38, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "embed_rows tf_embed_rows_bf16_copies", .symbol = "tf_embed_rows_bf16_copies", .image = "embed_rows", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "norms tf_rms_rows_sums", .symbol = "tf_rms_rows_sums", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 31, .static = 132, .bound = 1024 }, .{ .regs = 37, .static = 132, .bound = 1024 }, .{ .regs = 37, .static = 132, .bound = 1024 }, .{ .regs = 34, .static = 132, .bound = 1024 }, .{ .regs = 34, .static = 132, .bound = 1024 } } }, - .{ .launch = .{ .name = "route_rows tf_route_logits", .symbol = "tf_route_logits", .image = "route_rows", .threads = 128, .dynamic = 26112 }, .usage = .{ .{ .regs = 31, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 }, .{ .regs = 34, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "route_rows tf_route_logits", .symbol = "tf_route_logits", .image = "route_rows", .threads = 128, .dynamic = 26112 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 }, .{ .regs = 34, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "route_rows tf_route_logits_few", .symbol = "tf_route_logits_few", .image = "route_rows", .threads = 128, .dynamic = 76032, .opt_in = true }, .usage = .{ .{ .regs = 31, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 }, .{ .regs = 37, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "route_rows tf_route_topk", .symbol = "tf_route_topk", .image = "route_rows", .threads = 128 }, .usage = .{ .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 36, .static = 64, .bound = 128 }, .{ .regs = 36, .static = 64, .bound = 128 } } }, .{ .launch = .{ .name = "norms tf_rms_rows_bf16", .symbol = "tf_rms_rows_bf16", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 } } }, @@ -337,11 +337,11 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "latent tf_super_topk_few", .usage = .{ .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 } } }, .{ .name = "ck tf_ck4_up_relu2_few", .usage = .{ .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 } } }, .{ .name = "mamba tf_nemo_scan", .usage = .{ .{ .regs = 138, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 } } }, - .{ .name = "flashnext_batch fn_gdn_norm", .usage = .{ .{ .regs = 24, .static = 528, .bound = 128 }, .{ .regs = 26, .static = 528, .bound = 128 }, .{ .regs = 26, .static = 528, .bound = 128 }, .{ .regs = 24, .static = 528, .bound = 128 }, .{ .regs = 24, .static = 528, .bound = 128 } } }, - } ++ @import("paging_pending.zig").cuda_132 ++ @import("nemotron_prompt_records.zig").cuda_132 ++ @import("glm_cuda_records.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132) }, - .{ .release = .{ .major = 13, .minor = 3 }, .usage = &.{ + } ++ @import("paging_pending.zig").cuda_132 ++ @import("nemotron_prompt_records.zig").cuda_132 ++ @import("glm_cuda_records.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132) }, + .{ .release = .{ .major = 13, .minor = 3 }, .usage = &([_]Override{ .{ .name = "router tf_nemo_topk", .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, - } }, + .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 } } }, + } ++ @import("flashnext_tile_records.zig").cuda_133) }, }; pub fn newer(x: limits.Nvcc, y: limits.Nvcc) bool { From 06805fcb28fbaa977dc30d0eaa79f9cb07dd03de Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Sun, 18 Oct 2026 21:04:15 +0200 Subject: [PATCH 14/16] Flash Next CUDA tests: the fused prompt attention launch against chunks then merge, byte for byte, dense and sparse --- zig/tests/cuda/flashnext_attn_tests.zig | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/zig/tests/cuda/flashnext_attn_tests.zig b/zig/tests/cuda/flashnext_attn_tests.zig index 806a7207bf..ec86527148 100644 --- a/zig/tests/cuda/flashnext_attn_tests.zig +++ b/zig/tests/cuda/flashnext_attn_tests.zig @@ -264,6 +264,10 @@ fn attention(r: *Rig, g: *const flash.Glue, sparse: bool) !void { } }; for (counts[0 .. counts.len - 1], 0..) |m, c| try r.sameRows(name, u16, outs[c], got[0 .. m * q_heads * head_dim]); + // prompt rows' fused launch (every chunk merged in the program): the chunks then merge's bytes + const fused = try r.out(rows * q_heads * head_dim * 2); + try g.attnFused(r.s, dq, dkc, dvc, dpos, fused, dids, dnk, dspr, rows, kv_heads); + try r.sameRows(if (sparse) "attention sparse fused" else "attention dense fused", u16, fused, got); } pub fn run(gpu: check.Gpu) !void { From 4dd3cbb03632d40883ffa5ea5eeba341938fa193 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:05:06 +0200 Subject: [PATCH 15/16] Flash Next CUDA: TF_FLASHNEXT_CHUNK refused before the loader and the engine take anything, with a loader-level control --- zig/src/families/flashnext/cuda_lanes.zig | 15 +++++++++++++++ zig/src/families/flashnext/cuda_lanes_load.zig | 4 +++- zig/src/families/flashnext/cuda_prompt.zig | 4 +++- 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 67356ec88e..6be063655b 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -511,6 +511,21 @@ pub fn Namespace(comptime Tri: type) type { _ = @import("cuda_lanes_admission_test.zig").Cases(Cuda, Cuda.open, dropLane); } + test "a bad TF_FLASHNEXT_CHUNK is refused at load before the loader takes a stream, a module or a buffer" { + const env = if (@import("builtin").os.tag == .windows) @import("sys").c else struct { + extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: i32) i32; + extern "c" fn unsetenv(name: [*:0]const u8) i32; + }; + defer _ = env.unsetenv("TF_FLASHNEXT_CHUNK"); + // every driver entry point undefined: the refusal must come before the first call + var driver: cuda.Driver = undefined; + const ctx: cuda.Context = .{ .d = &driver, .device = 0, .handle = null }; + for ([_][*:0]const u8{ "0", "16385", "4k" }) |bad| { + _ = env.setenv("TF_FLASHNEXT_CHUNK", bad, 1); + try std.testing.expectError(error.BadChunk, open(std.testing.allocator, std.testing.io, &ctx, "no-such-checkpoint", null, .{ .context = 4096, .drafts = true })); + } + } + // -- the prompt cache's copies ------------------------------------------------------------------------ fn snapBytes(ptr: *anyopaque, at: u32) u64 { diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 19166bd0bc..debd9c081d 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -15,6 +15,8 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (o.context > prompt.max_tokens) return error.ContextOverWindow; if (o.segments != 1 or o.slide) return error.UnsupportedOption; if ((o.tp != 1 and o.tp != 2) or o.rank >= o.tp) return error.BadRank; + // refused before the stream, the glue or a captured set loads + const chunk = try prompt.chunkRows(); const mode = try config.mode(gpa, io, dir); if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; const driver: *cuda.Driver = @constCast(ctx.d); @@ -29,7 +31,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (captured) |kernels| own.set = try F.comparison(gpa, io, driver, ctx, kernels, o.drafts); errdefer if (own.set) |*x| x.deinit(); if (o.tp == 2) { - var id = try ranks.identity(gpa, io, dir, o.context, try prompt.chunkRows()); + var id = try ranks.identity(gpa, io, dir, o.context, chunk); id.mtp = @intFromBool(o.drafts); if (o.vision) |v| { try @import("qwen_vision").media.checkLimits(v); diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 17aac6044f..c647975ceb 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -146,6 +146,8 @@ pub const Engine = struct { pub fn initRank(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8, world: usize, rank: usize) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + // refused before anything loads: a bad TF_FLASHNEXT_CHUNK leaves no module or buffer behind + const chunk = try chunkRows(); const e = try gpa.create(Engine); errdefer gpa.destroy(e); e.gpa = gpa; @@ -204,7 +206,7 @@ pub const Engine = struct { e.batch_k = try batch.Kernels.load(driver); e.batch_s = null; e.scores = null; - e.chunk = try chunkRows(); + e.chunk = chunk; return e; } From 4637e1babebd63479f598748f582bd153ed80079 Mon Sep 17 00:00:00 2001 From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:06:37 +0200 Subject: [PATCH 16/16] Flash Next CUDA: the indexer and one-buffer attention tiles' comments say decode runs them too --- zig/kernels/cuda/flashnext_attn.cu | 2 +- zig/src/core/contract/flashnext_tile_records.zig | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu index 3033fad110..93b075f071 100644 --- a/zig/kernels/cuda/flashnext_attn.cu +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -441,7 +441,7 @@ extern "C" __global__ void __launch_bounds__(128) fn_attn_fused(const __nv_bfloa attn_chunks<1, true, true>(Q, KC, VC, POS0, OUT, nullptr, nullptr, IDS, NKR, SPR, H, HK, G, CH, 1, SCALE, IDW, QSA); } -// attn_chunks with one buffer, keys then values (prompt chunks: three blocks a multiprocessor), the same sums. +// attn_chunks with one buffer, keys then values (17 to 256 rows: prompt chunks and wide decode rounds), the same sums. extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks_one(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, diff --git a/zig/src/core/contract/flashnext_tile_records.zig b/zig/src/core/contract/flashnext_tile_records.zig index 080f7b60d5..c3909adbb2 100644 --- a/zig/src/core/contract/flashnext_tile_records.zig +++ b/zig/src/core/contract/flashnext_tile_records.zig @@ -1,4 +1,4 @@ -//! Flash Next's prompt tiles' CUDA records: nvcc 13.0's in registry.cuda, then 13.2's and 13.3's where they differ. +//! Flash Next's new tiles' CUDA records, decode's indexer and one-buffer attention among them: nvcc 13.0, 13.2, 13.3. const registry = @import("registry.zig"); const CudaRecord = registry.CudaRecord;