diff --git a/zig/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu index 9e5bc18db1..f7064ed911 100644 --- a/zig/kernels/cuda/flashnext_batch.cu +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -341,3 +341,6 @@ extern "C" __global__ void __launch_bounds__(256) fn_gdn_chain_tpr(const float* int rows, float* ys_out) { gdn_tpr::chain<2, 16>(qn, kn, v, gb, state_in, state_out, rows, ys_out); } + +// The prompt GDN block in one kernel a value head (front, chain, gated norm; their bits), no fp32 rows through memory. +#include "flashnext_gdn_fused.cuh" diff --git a/zig/kernels/cuda/flashnext_gdn_fused.cuh b/zig/kernels/cuda/flashnext_gdn_fused.cuh new file mode 100644 index 0000000000..3b15c092e0 --- /dev/null +++ b/zig/kernels/cuda/flashnext_gdn_fused.cuh @@ -0,0 +1,204 @@ +// Flash Next prompt GDN in one kernel a value head: front, chain and gated norm, rows kept in shared memory. +namespace gdn_fused { +using namespace fn_gdn; +constexpr int STEPS = 16, TPR = 2, CH = 256, PR = 128, THREADS = CH + PR; +constexpr int ROW = (DK + DK + DV) * 4 + 16; // gdn_tpr::Stage's row: k, q, v, g, beta +constexpr int STAGE = STEPS * ROW; +constexpr int SMEM = 2 * STAGE + 2 * STEPS * DV * 4 + STEPS * 2 * DK * 4 + STEPS * 4; + +__device__ __forceinline__ void pbar() { asm volatile("bar.sync 1, 128;\n" ::: "memory"); } + +// conv_act for rows r0 .. r0 + n - 1 of channel c, each P / cs row loaded once (conv_act's taps, order and rounding) +__device__ __forceinline__ void conv_rows(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, int r0, int n, int c, float (&out)[STEPS]) { + float w[TAPS]; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) w[tap] = __bfloat162float(cw[c * TAPS + tap]); + float xw[STEPS + TAPS - 1]; +#pragma unroll + for (int k = 0; k < STEPS + TAPS - 1; ++k) { + const int at = r0 + k; // conv_act's `at` for row r0 + j, tap k - j + xw[k] = k < n + TAPS - 1 ? (at < TAPS - 1 ? __bfloat162float(cs[at * C + c]) : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c])) : 0.0f; + } +#pragma unroll + for (int j = 0; j < STEPS; ++j) { + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) acc = acc + w[tap] * xw[j + tap]; + out[j] = bf(acc / (1.0f + expf(-acc))); + } +} + +// producers: rows r0 .. r0 + n - 1's q, k (normed as fn_gdn_front), v, g, beta into the stage +__device__ __forceinline__ void front_batch(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, const float* __restrict__ a_log, + const float* __restrict__ dt_bias, int r0, int n, int hk, int hv, int p, + unsigned char* base, float (*x)[2][DK]) { + float o[STEPS]; + conv_rows(P, cs, cw, r0, n, hk * DK + p, o); // fn_gdn_front unit hk, t = p: q's channel +#pragma unroll + for (int j = 0; j < STEPS; ++j) if (j < n) x[j][0][p] = o[j]; + conv_rows(P, cs, cw, r0, n, NK * DK + hk * DK + p, o); // t = p + DK: k's channel +#pragma unroll + for (int j = 0; j < STEPS; ++j) if (j < n) x[j][1][p] = o[j]; + conv_rows(P, cs, cw, r0, n, 2 * NK * DK + hv * DV + p, o); // unit NK + hv, t = p: v +#pragma unroll + for (int j = 0; j < STEPS; ++j) if (j < n) reinterpret_cast(base + j * ROW + 2 * DK * 4)[p] = o[j]; + if (p < n) { // unit NK + hv, t == DV: row r0 + p's g and beta + const int r = r0 + p; + const float b = __bfloat162float(P[(size_t)r * PW + C + NV * DV + hv]); + const float a = __bfloat162float(P[(size_t)r * PW + C + NV * DV + NV + hv]); + float* gbp = reinterpret_cast(base + p * ROW + (2 * DK + DV) * 4); + gbp[0] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv])); + gbp[1] = bf(sigmoidf_(b)); + } + pbar(); + const int warp = p >> 5, lane = p & 31; + for (int task = warp; task < 2 * n; task += PR / 32) { // (row, q|k): fn_gdn_front's warp 0 (q) / 1 (k) + const int j = task >> 1, which = task & 1; + float v4[4], ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + v4[i] = x[j][which][lane * 4 + i]; + ss = ss + v4[i] * v4[i]; + } + ss = warp_sum(ss); + float inv = 1.0f / sqrtf(ss + 1e-6f); + if (which == 0) inv = inv * (1.0f / sqrtf((float)DK)); + float* dst = reinterpret_cast(base + j * ROW + (which == 0 ? DK * 4 : 0)); // stage: k first, then q +#pragma unroll + for (int i = 0; i < 4; ++i) dst[lane * 4 + i] = v4[i] * inv; + } +} + +// producers: fn_gdn_norm for n rows from row r0 (readouts ys[j][dv]) +__device__ __forceinline__ void norm_rows(const float (*ys)[DV], const __nv_bfloat16* __restrict__ P, + const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out, + float* __restrict__ xs, int r0, int n, int hv, int p, float* rinv) { + const int warp = p >> 5, lane = p & 31; + for (int j = warp; j < n; j += PR / 32) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float y = ys[j][lane * 4 + i]; + ss = ss + y * y; + } + ss = warp_sum(ss); + if (lane == 0) rinv[j] = 1.0f / sqrtf(ss / (float)DV + eps); + } + pbar(); + float zs[STEPS]; +#pragma unroll + for (int j = 0; j < STEPS; ++j) zs[j] = j < n ? __bfloat162float(P[(size_t)(r0 + j) * PW + C + hv * DV + p]) : 0.0f; + const float nw = __bfloat162float(norm_w[p]); +#pragma unroll + for (int j = 0; j < STEPS; ++j) { + if (j >= n) break; + const int r = r0 + j; + const float yn = bf(bf(ys[j][p] * rinv[j]) * nw); + const float z = zs[j]; + const float o = bf(yn * sigmoidf_(z)); + out[(size_t)r * NV * DV + hv * DV + p] = __float2bfloat16_rn(o); + const float gs = warp_sum(o); + if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs; + } + pbar(); +} + +__device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, const float* __restrict__ a_log, + const float* __restrict__ dt_bias, const float* state_in, float* state_out, int rows, + const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out, + float* __restrict__ xs, float* __restrict__ ys_out) { + extern __shared__ __align__(16) unsigned char smem[]; + unsigned char* stage = smem; // [2][STEPS] rows + float (*ysb)[STEPS][DV] = reinterpret_cast(smem + 2 * STAGE); + float (*xb)[2][DK] = reinterpret_cast(smem + 2 * STAGE + 2 * STEPS * DV * 4); + float* rinv = reinterpret_cast(smem + 2 * STAGE + 2 * STEPS * DV * 4 + STEPS * 2 * DK * 4); + const int hv = blockIdx.x, hk = hv / (NV / NK); + const bool producer = threadIdx.x >= CH; + const int p = threadIdx.x - CH; + const int batches = (rows + STEPS - 1) / STEPS; + auto fill = [&](int bi) { + const int r0 = bi * STEPS, n = min(STEPS, rows - r0); + front_batch(P, cs, cw, a_log, dt_bias, r0, n, hk, hv, p, stage + (bi & 1) * STAGE, xb); + }; + // the chain's registers (gdn_tpr::chain<2, 16>'s) + constexpr int L = 32 / TPR; + const int 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]; + if (!producer) { +#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]; + } else if (batches > 0) fill(0); + __syncthreads(); + for (int bi = 0; bi < batches; ++bi) { + const int r0 = bi * STEPS, n = min(STEPS, rows - r0); + if (producer) { + if (bi + 1 < batches) fill(bi + 1); + if (bi >= 1) norm_rows(ysb[(bi - 1) & 1], P, norm_w, eps, out, xs, (bi - 1) * STEPS, STEPS, hv, p, rinv); + } else { + const unsigned char* base = stage + (bi & 1) * STAGE; + for (int rr = 0; rr < n; ++rr) { + const unsigned char* d = base + rr * 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]; +#pragma unroll + for (int l = 0; l < L; ++l) { + const float4 k4 = *reinterpret_cast(d + (l * TPR + t) * 16); + 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 = gdn_tpr::tree(x); + const float delta = (vv - kv) * beta; +#pragma unroll + for (int l = 0; l < L; ++l) { + const float4 k4 = *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(gdn_tpr::tree(x)); + if (t == 0) { + ysb[bi & 1][rr][row] = o; + if (ys_out) ys_out[((size_t)(r0 + rr) * NV + hv) * DV + row] = o; + } + } + } + __syncthreads(); + } + if (producer) { + if (batches > 0) norm_rows(ysb[(batches - 1) & 1], P, norm_w, eps, out, xs, (batches - 1) * STEPS, rows - (batches - 1) * STEPS, hv, p, rinv); + } else { +#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_fused + +extern "C" __global__ void __launch_bounds__(gdn_fused::THREADS) fn_gdn_fused( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, + const float* __restrict__ a_log, const float* __restrict__ dt_bias, const float* state_in, float* state_out, int rows, + const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, + float* __restrict__ ys_out) { + gdn_fused::run(P, cs, cw, a_log, dt_bias, state_in, state_out, rows, norm_w, eps, out, xs, ys_out); +} diff --git a/zig/src/core/contract/flashnext_tile_records.zig b/zig/src/core/contract/flashnext_tile_records.zig index c3909adbb2..f7687055dc 100644 --- a/zig/src/core/contract/flashnext_tile_records.zig +++ b/zig/src/core/contract/flashnext_tile_records.zig @@ -10,6 +10,7 @@ pub const cuda = [_]CudaRecord{ .{ .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 = 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_batch fn_gdn_fused", .symbol = "fn_gdn_fused", .image = "flashnext_batch", .threads = 384, .dynamic = 82496, .opt_in = true }, .usage = .{ .{ .regs = 138, .static = 0, .bound = 384 }, .{ .regs = 142, .static = 0, .bound = 384 }, .{ .regs = 142, .static = 0, .bound = 384 }, .{ .regs = 166, .static = 0, .bound = 384 }, .{ .regs = 166, .static = 0, .bound = 384 } } }, .{ .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 } } }, diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 8858db829a..db864f2160 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const gdn = @import("cuda_gdn.zig"); +const gdn_fused = @import("cuda_gdn_fused.zig"); const ple = @import("cuda_ple.zig"); const wt = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); @@ -268,7 +269,13 @@ pub fn Pass(comptime Tri: type) type { 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 - if (self.lanes.len == 0) { + const fused = self.lanes.len == 0 and n >= 64 and gdn_fused.on(); + if (fused) { + // the three in one kernel a value head: the same per-row op sequences, no fp32 rows through memory + const st = self.seq.state[li].ptr; + const in: gdn_fused.In = .{ .proj = s.proj, .conv = self.seq.conv[li].ptr, .cw = lw.cw.ptr, .a_log = lw.a_log.ptr, .dt = lw.dt.ptr, .state = st, .state_out = if (keep.state_out != 0) keep.state_out else st, .norm = lw.norm.ptr }; + try gdn_fused.launch(self.k.gdn_fused, self.stream, in, n, s.gout, s.gxs); + } else if (self.lanes.len == 0) { try self.gdnFront(lw, s.proj, self.seq.conv[li].ptr, n, s.gq, s.gk, s.gv, s.ggb); try self.gdnChain(s.gq, s.gk, s.gv, s.ggb, self.seq.state[li].ptr, keep.state_out, n, s.ys); } else for (self.lanes, 0..) |l, r| { @@ -279,14 +286,7 @@ pub fn Pass(comptime Tri: type) type { try self.gdnFront(lw, s.proj + r * gdn.proj_width * 2, l.seq.conv[li].ptr, 1, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb); try self.gdnChain(s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, 0, 1, s.ys + v); } - var b: cuda.Args = .{}; - b.add(s.ys); - b.add(s.proj); - b.add(lw.norm.ptr); - b.add(eps); - b.add(s.gout); - b.add(s.gxs); - try cuda.launch.launch(self.k.gnorm, .{ .grid = .{ .x = @intCast(n), .y = gdn.nv }, .block = .{ .x = gdn.dv } }, self.stream, &b); + if (!fused) try gdn_fused.norm(self.k.gnorm, self.stream, s.ys, s.proj, lw.norm.ptr, s.gout, s.gxs, n); if (self.lanes.len == 0) { try self.shiftWindow(self.seq.conv[li].ptr, s.proj, gdn.proj_width, n, gdn.taps - 1, gdn.conv_dim); } else for (self.lanes, 0..) |l, r| { diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 2ae6340c25..f5bb342347 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -59,6 +59,10 @@ pub const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1EL 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); +/// fn_gdn_fused's shared memory (its SMEM): two chain stages, two readout batches, the conv rows and 16 norm scales. +pub const gdn_fused_smem: u32 = chain_tpr_smem + 2 * 16 * 128 * 4 + 16 * 2 * 128 * 4 + 16 * 4; +/// fn_gdn_fused's block: 256 chain threads and 128 producer threads. +pub const gdn_fused_threads: u32 = 384; /// 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; @@ -155,6 +159,8 @@ pub const Kernels = struct { chain: cuda.Function, chain_to: cuda.Function, chain_tpr: cuda.Function, + /// front + chain + gated norm in one kernel a value head (prompt passes of 64 rows or more). + gdn_fused: cuda.Function, gnorm: cuda.Function, sum: cuda.Function, moe_sum: cuda.Function, @@ -198,6 +204,8 @@ pub const Kernels = struct { 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.gdn_fused = try k.batch.function("fn_gdn_fused"); + try k.gdn_fused.allowDynamicShared(gdn_fused_smem); k.gnorm = try k.batch.function("fn_gdn_norm"); return k; } diff --git a/zig/src/families/flashnext/cuda_gdn_fused.zig b/zig/src/families/flashnext/cuda_gdn_fused.zig new file mode 100644 index 0000000000..fdec57027b --- /dev/null +++ b/zig/src/families/flashnext/cuda_gdn_fused.zig @@ -0,0 +1,35 @@ +//! The prompt GDN block's launches: fn_gdn_fused (front, chain, gated norm in one kernel a value head), the gated norm. +const std = @import("std"); +const cuda = @import("cuda"); +const gdn = @import("cuda_gdn.zig"); +const parts = @import("cuda_batch_state.zig"); +const eps = @import("cuda_prompt_wt.zig").eps; + +/// One kernel a value head for a lone stream's pass, unless TF_FLASHNEXT_GDN_FUSED=0 (three launches, the same bits). +pub fn on() bool { + const v = std.c.getenv("TF_FLASHNEXT_GDN_FUSED") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// fn_gdn_fused's inputs: projection rows, conv window and taps, gate parameters, the state and where it goes. +pub const In = struct { proj: u64, conv: u64, cw: u64, a_log: u64, dt: u64, state: u64, state_out: u64, norm: u64 }; + +/// fn_gdn_front, fn_gdn_chain_tpr and fn_gdn_norm over `rows` rows in one launch (their bits): out and xs. +pub fn launch(f: cuda.Function, stream: cuda.Stream, in: In, rows: usize, out: u64, xs: u64) !void { + var a: cuda.Args = .{}; + for ([_]u64{ in.proj, in.conv, in.cw, in.a_log, in.dt, in.state, in.state_out }) |v| a.add(v); + a.add(@as(c_int, @intCast(rows))); + a.add(in.norm); + a.add(eps); + for ([_]u64{ out, xs, 0 }) |v| a.add(v); // no fp32 readouts kept + try cuda.launch.launch(f, .{ .grid = .{ .x = gdn.nv }, .block = .{ .x = parts.gdn_fused_threads }, .shared = parts.gdn_fused_smem }, stream, &a); +} + +/// fn_gdn_norm over `rows` readouts ys: the gated norm into out, its sums into xs. +pub fn norm(f: cuda.Function, stream: cuda.Stream, ys: u64, proj: u64, norm_w: u64, out: u64, xs: u64, rows: usize) !void { + var b: cuda.Args = .{}; + for ([_]u64{ ys, proj, norm_w }) |v| b.add(v); + b.add(eps); + for ([_]u64{ out, xs }) |v| b.add(v); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = gdn.nv }, .block = .{ .x = gdn.dv } }, stream, &b); +} diff --git a/zig/tests/cuda/contract_flashnext_glue.zig b/zig/tests/cuda/contract_flashnext_glue.zig index 60469d399f..816db99bb0 100644 --- a/zig/tests/cuda/contract_flashnext_glue.zig +++ b/zig/tests/cuda/contract_flashnext_glue.zig @@ -11,7 +11,7 @@ const g = contract.gates; const parts = flash.batch_state; const tk = flash.tile_kernels; -pub const Kind = enum { chain, scores, readout }; +pub const Kind = enum { chain, scores, readout, fused }; const nk = 16; const nv = 48; @@ -29,6 +29,9 @@ const n_down = low + streams; const hc_sk = 32; const nc = dims / 256; const eps: f32 = 1e-6; +const taps = 4; +const conv_c = 2 * nk * dk + nv * dv; // a projection row's conv channels (q, k, v) +const proj_w = conv_c + nv * dv + 2 * nv; // a projection row: conv channels, z, b, a fn noSums(_: c.Shape) usize { return 0; @@ -39,11 +42,12 @@ 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 } } }, + .fused => .{ .name = "cuda flashnext DeltaNet prompt block fn_gdn_fused (front, chain, gated norm)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "gated-norm rows, xs, the state and the readouts bit-equal to fn_gdn_front -> fn_gdn_rows_to -> fn_gdn_norm, non-finite inputs included" }, .replaces = "fn_gdn_front, fn_gdn_chain_tpr and fn_gdn_norm", .tiles = .{ .{ .size = 16 }, .{}, .{} }, .served = &.{ .{ .rows = 64, .n = nv * dv, .k = dk }, .{ .rows = 300, .n = nv * dv, .k = dk } } }, .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 Fns = struct { tpr: cuda.Function, rows_to: cuda.Function, tile: cuda.Function, scores: cuda.Function, hcf_down: cuda.Function, hcf_up: cuda.Function, front: cuda.Function, norm: cuda.Function, fused: cuda.Function }; pub const Glue = struct { rig: rig_.Rig, @@ -160,6 +164,40 @@ pub const Glue = struct { if (shape.n != dims or shape.k != wide) return false; try self.readoutInputs(shape.rows, r, bad); }, + .fused => { + if (shape.n != nv * dv or shape.k != dk) return false; + const m = shape.rows; + const proj = try self.bfs(m * proj_w, r, -1, 1); + defer a.free(proj); + const cs = try self.bfs((taps - 1) * conv_c, r, -1, 1); + defer a.free(cs); + const cw = try self.bfs(conv_c * taps, r, -0.5, 0.5); + defer a.free(cw); + const a_log = try self.f32s(nv, r, -1, 1); + defer a.free(a_log); + const dt = try self.f32s(nv, r, -1, 1); + defer a.free(dt); + const st = try self.f32s(nv * dv * dk, r, -0.5, 0.5); + defer a.free(st); + const nw = try self.bfs(dv, r, 0.5, 1.5); + defer a.free(nw); + if (bad) { + proj[3] = c.toBf16(std.math.nan(f32)); + proj[(m / 2) * proj_w + conv_c + 7] = c.toBf16(std.math.inf(f32)); + st[dv * dk + 7] = -std.math.inf(f32); + } + try self.in(0, std.mem.sliceAsBytes(proj)); + try self.in(1, std.mem.sliceAsBytes(cs)); + try self.in(2, std.mem.sliceAsBytes(cw)); + try self.in(3, std.mem.sliceAsBytes(a_log)); + try self.in(4, std.mem.sliceAsBytes(dt)); + try self.in(5, std.mem.sliceAsBytes(st)); + try self.in(6, std.mem.sliceAsBytes(nw)); + try self.out(m * nv * dv * 2); // gated-norm rows (bf16) + try self.out(m * (nv * dv / 32) * 4); // their 32-wide sums + try self.out(nv * dv * dk * 4); // the state after the pass + try self.out(m * nv * dv * 4); // the chain's readouts + }, } self.base = 0; var total: usize = 0; @@ -239,6 +277,15 @@ pub const Glue = struct { 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); }, + .fused => { + var a: cuda.Args = .{}; + for ([_]u64{ self.ins[0], self.ins[1], self.ins[2], self.ins[3], self.ins[4], self.ins[5], o[2] }) |v| a.add(v); + a.add(@as(c_int, @intCast(m))); + a.add(self.ins[6]); + a.add(eps); + for ([_]u64{ o[0], o[1], o[3] }) |v| a.add(v); + try cuda.launch.launch(self.f.fused, .{ .grid = .{ .x = nv }, .block = .{ .x = parts.gdn_fused_threads }, .shared = parts.gdn_fused_smem }, s, &a); + }, .readout => { const part = try self.rig.scratch(m * n_down * 4); const rinv = try self.rig.scratch(m * streams * 4); @@ -284,6 +331,26 @@ pub const Glue = struct { 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); }, + .fused => { + const gq = try self.rig.scratch(m * nk * dk * 4); + const gk = try self.rig.scratch(m * nk * dk * 4); + const gv = try self.rig.scratch(m * nv * dv * 4); + const ggb = try self.rig.scratch(m * nv * 2 * 4); + var f: cuda.Args = .{}; + for ([_]u64{ self.ins[0], self.ins[1], self.ins[2], self.ins[3], self.ins[4], gq, gk, gv, ggb }) |v| f.add(v); + try cuda.launch.launch(self.f.front, .{ .grid = .{ .x = @intCast(m), .y = nk + nv }, .block = .{ .x = 256 } }, s, &f); + var a: cuda.Args = .{}; + for ([_]u64{ gq, gk, gv, ggb, self.ins[5], o[2] }) |v| a.add(v); + a.add(@as(c_int, @intCast(m))); + a.add(o[3]); + try cuda.launch.launch(self.f.rows_to, .{ .grid = .{ .x = nv, .y = dv / 16 }, .block = .{ .x = 128 } }, s, &a); + var b: cuda.Args = .{}; + for ([_]u64{ o[3], self.ins[0], self.ins[6] }) |v| b.add(v); + b.add(eps); + b.add(o[0]); + b.add(o[1]); + try cuda.launch.launch(self.f.norm, .{ .grid = .{ .x = @intCast(m), .y = nv }, .block = .{ .x = dv } }, s, &b); + }, .readout => { const normed = try self.rig.scratch(m * wide * 2); const dn = try self.rig.scratch(m * (low + 4) * 2); diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index b5dcebd49a..d201bddeba 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -398,11 +398,15 @@ pub fn run(gpu: check.Gpu) !void { .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), + .front = try batch.function("fn_gdn_front"), + .norm = try batch.function("fn_gdn_norm"), + .fused = try batch.function("fn_gdn_fused"), }; + try glue_fns.fused.allowDynamicShared(fparts.gdn_fused_smem); 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; + var glues: [4]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());