Zig CUDA: one drafts-off Flash Next prompt matches

zig/build/cuda.zig

@@ -10,8 +10,10 @@ const torch_ops = &[_][]const u8{ "-O3", "--fmad=false", "--ftz=false" };
 
 const kernels = [_]Kernel{
     .{ .name = "gdn", .flags = &.{ "-O3", "--fmad=false" } }, // cuda/kernels/gdn.py, tensorfold_gdn_v2
+    .{ .name = "flashnext_gdn", .flags = &.{ "-O3", "--fmad=false" } }, // qwen4_exp/cuda/gdn.cu chain, one row
     .{ .name = "probe", .flags = &.{"-O3"} },
     .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5
+    .{ .name = "qmm", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py lane matmul, groups of 32
     .{ .name = "qmm_prefill", .flags = &.{"-O3"} },
     .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7
     .{ .name = "experts_prefill", .flags = &.{"-O3"} },
@@ -55,7 +57,7 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti
     return cuda;
 }
 
-fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module } {
+fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module, tokenizer: *std.Build.Module } {
     const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true });
     const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true });
     core.addImport("tokenizer", tokenizer);
@@ -65,7 +67,11 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin
     nemotron.addImport("core", core);
     nemotron.addImport("lanes", lanes);
     nemotron.addImport("nemotron_draft_ids", draft_ids);
-    return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer };
+    const flashnext = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_launch.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    flashnext.addImport("cuda", cuda);
+    flashnext.addImport("core", core);
+    flashnext.addImport("lanes", lanes);
+    return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .flashnext = flashnext, .tokenizer = tokenizer };
 }
 
 /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`.
@@ -100,25 +106,45 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true });
     runner.addImport("cuda", cuda);
     b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }));
-    nativeServer(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.tokenizer);
+    // not installed: one embedding row through the captured _embed cubin
+    const embed_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_embed_main.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    embed_check.addImport("cuda", cuda);
+    embed_check.addImport("core", mods.core);
+    embed_check.addImport("flashnext", mods.flashnext);
+    const embed_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-embed", .root_module = embed_check }));
+    const embed_model = b.option([]const u8, "flashnext_model", "Flash Next checkpoint directory for the embedding check");
+    const embed_kernels = b.option([]const u8, "flashnext_kernels", "captured Triton pack for the embedding check");
+    const embed_token = b.option([]const u8, "flashnext_token", "token id for the embedding check") orelse "7";
+    if (embed_model != null and embed_kernels != null) embed_run.addArgs(&.{ embed_model.?, embed_kernels.?, embed_token });
+    b.step("flashnext-embed", "Upload the Flash Next embedding table and run _embed for one token").dependOn(&embed_run.step);
+    const prompt_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_prompt_main.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    prompt_check.addImport("cuda", cuda);
+    prompt_check.addImport("core", mods.core);
+    prompt_check.addImport("flashnext", mods.flashnext);
+    const prompt_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-prompt", .root_module = prompt_check }));
+    const prompt_ids = b.option([]const u8, "flashnext_ids", "file of prompt token ids for the Flash Next prompt check");
+    const prompt_count = b.option([]const u8, "flashnext_count", "greedy tokens to emit") orelse "8";
+    if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count });
+    b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step);
+    nativeServer(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.flashnext, mods.tokenizer);
 }
 
 /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families.
-fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } {
+fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } {
     const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} });
     const mod = b.createModule(.{
         .root_source_file = b.path("zig/src/native/cuda.zig"),
         .target = target,
         .optimize = optimize,
         .link_libc = true,
-        .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron } },
+        .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron }, .{ .name = "flashnext", .module = flashnext } },
     });
     return .{ .api = api, .engines = mod };
 }
 
 /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it.
-fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module) void {
-    const m = engines(b, target, optimize, cuda, lanes, nemotron);
+fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module, tokenizer: *std.Build.Module) void {
+    const m = engines(b, target, optimize, cuda, lanes, nemotron, flashnext);
     // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's)
     const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true });
     const exe = b.addExecutable(.{ .name = "tensorfold-native", .root_module = b.createModule(.{
@@ -137,8 +163,8 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, step: *std.Build.S
     const host = b.graph.host;
     const cuda = runtime(b, host, .debug, &.{});
     const mods = family(b, host, .debug, cuda, draft_ids);
-    const native = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron).engines;
-    for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step);
+    const native = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron, mods.flashnext).engines;
+    for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.flashnext, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step);
     const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true });
     cli.addImport("cuda", cuda);
     cli.addImport("core", mods.core);

zig/kernels/cuda/experts.cu

@@ -280,3 +280,8 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const
     int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float);
 TF_EXPERT(1)
 TF_EXPERT(0)
+// Groups of 32: SwiGLU gate and up together (epilogue 2), and the down projection to fp32.
+template __global__ void tf_experts::expert_kernel<32, 2, 2, 2, 4>(const __nv_bfloat16*, \
+    int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float);
+template __global__ void tf_experts::expert_kernel<32, 1, 0, 2, 4>(const __nv_bfloat16*, \
+    int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float);

zig/kernels/cuda/experts_pack.cu

@@ -37,6 +37,8 @@ pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ sca
 
 } // namespace tf_experts_pack
 
-// Groups of 64 inputs.
+// Groups of 64 inputs, then groups of 32.
 template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t*, const uint16_t*, const uint16_t*,
     uint32_t*, int, int, int, int);
+template __global__ void tf_experts_pack::pack_kernel<1>(const uint32_t*, const uint16_t*, const uint16_t*,
+    uint32_t*, int, int, int, int);

zig/kernels/cuda/flashnext_gdn.cu

@@ -0,0 +1,143 @@
+// One value head per block: depthwise conv, the delta-rule step, then the gated RMSNorm. 16 key heads, 48 value heads.
+// Same operations as the one-row chain, compiled without FMA contraction.
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+constexpr int NK = 16, NV = 48, DK = 128, DV = 128, TAPS = 4;
+
+static __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); }
+
+static __device__ __forceinline__ float warp_sum(float x) {
+    for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o);
+    return x;
+}
+
+static __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); }
+
+static __device__ __forceinline__ float softplusf_(float x) { return x > 20.0f ? x : log1pf(expf(x)); }
+
+static __device__ __forceinline__ void update(float (&s)[4][4], const float (&kk)[4], const float* vrow, int warp, float g,
+                                       float beta) {
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        float kv = 0.0f;
+#pragma unroll
+        for (int i = 0; i < 4; ++i) {
+            s[j][i] = s[j][i] * g;
+            kv = kv + s[j][i] * kk[i];
+        }
+        kv = warp_sum(kv);
+        const float delta = (vrow[warp * 4 + j] - kv) * beta;
+#pragma unroll
+        for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta;
+    }
+}
+
+extern "C" __global__ void __launch_bounds__(1024) flashnext_gdn_chain(
+        const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs,
+        const __nv_bfloat16* __restrict__ cw, const float* __restrict__ state_in, const float* __restrict__ a_log,
+        const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ norm_w, float eps, int rows,
+        __nv_bfloat16* __restrict__ out, float* __restrict__ xs, float* __restrict__ state_out) {
+    constexpr int C = 2 * NK * DK + NV * DV;
+    constexpr int PW = C + NV * DV + 2 * NV;
+    const int hv = blockIdx.x, hk = hv / (NV / NK);
+    const int t = threadIdx.x, warp = t >> 5, lane = t & 31;
+    __shared__ float qs[DK], ks[DK], vs[DV], ys[DV];
+    __shared__ float gates[2], rinv;
+    int c = -1;
+    if (t < DK) c = hk * DK + t;
+    else if (t < 2 * DK) c = NK * DK + hk * DK + (t - DK);
+    else if (t < 2 * DK + DV) c = 2 * NK * DK + hv * DV + (t - 2 * DK);
+    float s[4][4];
+    const size_t sbase = (size_t)hv * DV * DK;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+#pragma unroll
+        for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i];
+    const __nv_bfloat16* pz = P + C + hv * DV + (t < DV ? t : 0);
+    const __nv_bfloat16* pb = P + C + NV * DV + hv;
+    const __nv_bfloat16* pa = pb + NV;
+    for (int r = 0; r < rows; ++r) {
+        if (c >= 0) {
+            float acc = 0.0f;
+#pragma unroll
+            for (int tap = 0; tap < TAPS; ++tap) {
+                const int at = r + tap;
+                const float x = at < TAPS - 1 ? __bfloat162float(cs[at * C + c])
+                                              : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c]);
+                acc = acc + __bfloat162float(cw[c * TAPS + tap]) * x;
+            }
+            const float act = bf(acc / (1.0f + expf(-acc)));
+            if (t < DK) qs[t] = act;
+            else if (t < 2 * DK) ks[t - DK] = act;
+            else vs[t - 2 * DK] = act;
+        }
+        __syncthreads();
+        if (warp < 2) {
+            float* x = warp == 0 ? qs : ks;
+            float v4[4], ss = 0.0f;
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                v4[i] = x[lane * 4 + i];
+                ss = ss + v4[i] * v4[i];
+            }
+            ss = warp_sum(ss);
+            float inv = 1.0f / sqrtf(ss + 1e-6f);
+            if (warp == 0) inv = inv * (1.0f / sqrtf((float)DK));
+            __syncwarp();
+#pragma unroll
+            for (int i = 0; i < 4; ++i) x[lane * 4 + i] = v4[i] * inv;
+        } else if (warp == 2 && lane == 0) {
+            const float b = __bfloat162float(pb[(size_t)r * PW]);
+            const float a = __bfloat162float(pa[(size_t)r * PW]);
+            gates[0] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv]));
+            gates[1] = bf(sigmoidf_(b));
+        }
+        __syncthreads();
+        const float g = gates[0], beta = gates[1];
+        float kk[4], qq[4];
+#pragma unroll
+        for (int i = 0; i < 4; ++i) {
+            kk[i] = ks[lane * 4 + i];
+            qq[i] = qs[lane * 4 + i];
+        }
+        update(s, kk, vs, warp, g, beta);
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            float o = 0.0f;
+#pragma unroll
+            for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i];
+            o = warp_sum(o);
+            if (lane == 0) ys[warp * 4 + j] = bf(o);
+        }
+        __syncthreads();
+        if (warp == 0) {
+            float ss = 0.0f;
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const float y = ys[lane * 4 + i];
+                ss = ss + y * y;
+            }
+            ss = warp_sum(ss);
+            if (lane == 0) rinv = 1.0f / sqrtf(ss / (float)DV + eps);
+        }
+        __syncthreads();
+        if (t < DV) {
+            const float yn = bf(bf(ys[t] * rinv) * __bfloat162float(norm_w[t]));
+            const float z = __bfloat162float(pz[(size_t)r * PW]);
+            const float o = bf(yn * sigmoidf_(z));
+            out[(size_t)r * NV * DV + hv * DV + t] = __float2bfloat16_rn(o);
+            const float gs = warp_sum(o);
+            if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs;
+        }
+        __syncthreads();
+    }
+    if (state_out != nullptr) {
+#pragma unroll
+        for (int j = 0; j < 4; ++j)
+#pragma unroll
+            for (int i = 0; i < 4; ++i)
+                state_out[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i];
+    }
+}