Flash Next CUDA: the prompt DeltaNet block in one kernel a value head

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"

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<float*>(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<float*>(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<float*>(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<float (*)[STEPS][DV]>(smem + 2 * STAGE);
+    float (*xb)[2][DK] = reinterpret_cast<float (*)[2][DK]>(smem + 2 * STAGE + 2 * STEPS * DV * 4);
+    float* rinv = reinterpret_cast<float*>(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<const float*>(d + 2 * DK * 4)[row];
+                const float2 gbv = *reinterpret_cast<const float2*>(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<const float4*>(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<TPR>(x);
+                const float delta = (vv - kv) * beta;
+#pragma unroll
+                for (int l = 0; l < L; ++l) {
+                    const float4 k4 = *reinterpret_cast<const float4*>(d + (l * TPR + t) * 16);
+                    const float4 q4 = *reinterpret_cast<const float4*>(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<TPR>(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);
+}

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 } } },

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| {

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;
     }