Qwen3.8-27B on CUDA: the qwen27 family's engine for NVIDIA GPUs, every stream drafted

zig/build/cuda.zig

@@ -13,6 +13,7 @@ const glue = &[_][]const u8{ "-O3", "--fmad=false", "--ftz=false", "--prec-div=t
 
 const kernels = [_]Kernel{
     .{ .name = "gdn", .flags = &.{ "-O3", "--fmad=false" } }, // cuda/kernels/gdn.py, tensorfold_gdn_v2
+    .{ .name = "gdn_prefill", .flags = &.{ "-O3", "--fmad=false" } }, // the same extension's prompt chain
     .{ .name = "probe", .flags = &.{"-O3"} },
     .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5
     .{ .name = "qmm_prefill", .flags = &.{"-O3"} },
@@ -25,6 +26,7 @@ const kernels = [_]Kernel{
     .{ .name = "affine4_pack", .flags = &.{"-O3"} }, // ours: the MLX affine-4 repack into qlinear's tiles
     .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA
     .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws
+    .{ .name = "qwen_ops", .flags = &.{"-O3"} }, // ours: the 27B's batched state copies, taps, drafter norms and rotary
     .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue
     .{ .name = "nemotron_route", .flags = glue },
     .{ .name = "nemotron_mamba", .flags = glue },
@@ -39,6 +41,7 @@ const kernels = [_]Kernel{
     .{ .name = "torch_indexing", .src = "torch_ops/indexing", .flags = torch_ops },
     .{ .name = "torch_movement", .src = "torch_ops/movement", .flags = torch_ops },
     .{ .name = "torch_nemotron_constants", .src = "torch_ops/nemotron_constants", .flags = torch_ops },
+    .{ .name = "torch_norm_rotary", .src = "torch_ops/norm_rotary", .flags = torch_ops }, // F.rms_norm of 5120 columns
 };
 
 /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define.
@@ -69,7 +72,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, heat: *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, qwen27: *std.Build.Module, tokenizer: *std.Build.Module, heat: *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);
@@ -81,7 +84,13 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin
     nemotron.addImport("lanes", lanes);
     nemotron.addImport("heat", heat);
     nemotron.addImport("nemotron_draft_ids", draft_ids);
-    return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat };
+    // the 27B's CUDA engine, beside its Metal one in the family's folder (it reads the family's config.zig)
+    const qwen27 = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen27/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    qwen27.addImport("cuda", cuda);
+    qwen27.addImport("core", core);
+    qwen27.addImport("lanes", lanes);
+    qwen27.addImport("nemotron", nemotron); // the checkpoint Source and torch ops the CUDA families share
+    return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .qwen27 = qwen27, .tokenizer = tokenizer, .heat = heat };
 }
 
 /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`.
@@ -113,7 +122,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     cli.addImport("core", mods.core);
     cli.addImport("lanes", mods.lanes);
     cli.addImport("nemotron", mods.nemotron);
-    const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron);
+    const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.qwen27);
     cli.addOptions("build_options", build_options);
     cli.addImport("checkpoint_cli", checkpointCli(b, target, optimize, strip, native.engines));
     b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli }));
@@ -122,18 +131,22 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     runner.addImport("lanes", mods.lanes);
     runner.addImport("nemotron", mods.nemotron);
     b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }));
-    nativeServer(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip;
+    const qwen_cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/qwen27_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    qwen_cli.addImport("cuda", cuda);
+    qwen_cli.addImport("qwen27", mods.qwen27);
+    b.installArtifact(b.addExecutable(.{ .name = "tf-qwen27-cuda", .root_module = qwen_cli }));
+    nativeServer(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.qwen27, mods.tokenizer, build_options, true).root_module.strip = strip;
 }
 
 /// 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, qwen27: *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 = "qwen27", .module = qwen27 } },
     });
     return .{ .api = api, .engines = mod };
 }
@@ -144,8 +157,8 @@ fn checkpointCli(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.
 }
 
 /// `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, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile {
-    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, qwen27: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile {
+    const m = engines(b, target, optimize, cuda, lanes, nemotron, qwen27);
     // 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(.{
@@ -172,7 +185,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     const host = b.graph.host;
     const cuda = runtime(b, host, .debug, &.{});
     const mods = family(b, host, .debug, cuda, draft_ids);
-    const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron);
+    const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron, mods.qwen27);
     const native = native_modules.engines;
     // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too
     const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = host, .optimize = .debug, .link_libc = true });
@@ -186,7 +199,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     const server_test = b.addRunArtifact(b.addTest(.{ .root_module = server_tests }));
     step.dependOn(&server_test.step);
     b.step("test-server-cpu", "The HTTP server's unit tests (routes, templates, tool parsing) without a GPU").dependOn(&server_test.step);
-    for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step);
+    for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen27, mods.heat, 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);
@@ -222,7 +235,7 @@ pub fn distServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *s
     };
     const cuda = runtime(b, target, .fast, if (prebuilt != null) &images else &.{});
     const mods = family(b, target, .fast, cuda, draft_ids);
-    return nativeServer(b, target, .fast, cuda, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false);
+    return nativeServer(b, target, .fast, cuda, mods.lanes, mods.nemotron, mods.qwen27, mods.tokenizer, build_options, false);
 }
 
 /// Validate the complete named input set, including images unused by today's server.

zig/kernels/cuda/gdn_prefill.cu

@@ -0,0 +1,150 @@
+// Device code of src/tensorfold/cuda/kernels/gdn_prefill.cu (lines 1-145, comments and ATen includes dropped).
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+namespace tf_gdn_prefill {
+
+constexpr int DK = 128;
+constexpr int DV = 128;
+
+__device__ __forceinline__ float4 widen(uint2 w) {
+    const __nv_bfloat162 a = *reinterpret_cast<const __nv_bfloat162*>(&w.x);
+    const __nv_bfloat162 b = *reinterpret_cast<const __nv_bfloat162*>(&w.y);
+    return make_float4(__low2float(a), __high2float(a), __low2float(b), __high2float(b));
+}
+
+__device__ __forceinline__ float4 widen(float4 w) { return w; }
+
+template <typename QK, int THREADS, int ROWS, int STEPS>
+struct Prefetch {
+    using Raw = typename std::conditional<std::is_same<QK, float>::value, float4, uint2>::type;
+    static constexpr int KQ = STEPS * (DK / 4) / THREADS, VV = STEPS * ROWS / 8 / THREADS;
+    static_assert(KQ * THREADS == STEPS * (DK / 4) && VV * THREADS * 8 == STEPS * ROWS, "whole slots a thread");
+    Raw kr[KQ], qr[KQ];
+    uint4 vr[VV];
+    float gr, br;
+
+    __device__ __forceinline__ void fetch(const QK* q, const QK* k, const __nv_bfloat16* v, const float* g,
+                                          const float* beta, int t0, int W, int hk, int hv, int head, int key_head,
+                                          int col0) {
+#pragma unroll
+        for (int it = 0; it < KQ; ++it) {
+            const int i = threadIdx.x + it * THREADS, st = i / (DK / 4), c = i % (DK / 4);
+            const size_t at = (static_cast<size_t>(min(t0 + st, W - 1)) * hk + key_head) * DK + 4 * c;
+            kr[it] = *reinterpret_cast<const Raw*>(k + at);
+            qr[it] = *reinterpret_cast<const Raw*>(q + at);
+        }
+#pragma unroll
+        for (int it = 0; it < VV; ++it) {
+            const int i = threadIdx.x + it * THREADS, st = i / (ROWS / 8), c = i % (ROWS / 8);
+            vr[it] = *reinterpret_cast<const uint4*>(v + (static_cast<size_t>(min(t0 + st, W - 1)) * hv + head) * DV
+                                                     + col0 + 8 * c);
+        }
+        if (threadIdx.x < STEPS) {
+            const size_t at = static_cast<size_t>(min(t0 + static_cast<int>(threadIdx.x), W - 1)) * hv + head;
+            gr = g[at];
+            br = beta[at];
+        }
+    }
+
+    __device__ __forceinline__ void store(float4 (*ks)[DK / 4], float4 (*qs)[DK / 4], float* gs, float* bs,
+                                          __nv_bfloat16 (*vs)[ROWS]) const {
+#pragma unroll
+        for (int it = 0; it < KQ; ++it) {
+            const int i = threadIdx.x + it * THREADS;
+            ks[i / (DK / 4)][i % (DK / 4)] = widen(kr[it]);
+            qs[i / (DK / 4)][i % (DK / 4)] = widen(qr[it]);
+        }
+#pragma unroll
+        for (int it = 0; it < VV; ++it) {
+            const int i = threadIdx.x + it * THREADS;
+            *reinterpret_cast<uint4*>(&vs[i / (ROWS / 8)][8 * (i % (ROWS / 8))]) = vr[it];
+        }
+        if (threadIdx.x < STEPS) {
+            gs[threadIdx.x] = gr;
+            bs[threadIdx.x] = br;
+        }
+    }
+};
+
+template <typename QK, int ROWS, int STEPS>
+__global__ void __launch_bounds__(2 * ROWS) chain_kernel(
+        const QK* __restrict__ q, const QK* __restrict__ k, const __nv_bfloat16* __restrict__ v,
+        const float* __restrict__ g, const float* __restrict__ beta, const float* __restrict__ state,
+        float* __restrict__ last, __nv_bfloat16* __restrict__ y, int W, int hk, int hv) {
+    __shared__ float4 ks[STEPS][DK / 4], qs[STEPS][DK / 4];
+    __shared__ float gs[STEPS], bs[STEPS];
+    __shared__ __align__(16) __nv_bfloat16 vs[STEPS][ROWS];
+    const int head = blockIdx.x, row = blockIdx.y * ROWS + (threadIdx.x >> 1), half = threadIdx.x & 1;
+    const int key_head = head / (hv / hk);
+    float s[64];
+    const float* s0 = state + (static_cast<size_t>(head) * DV + row) * DK;
+#pragma unroll
+    for (int j = 0; j < 16; ++j) {
+        const float4 t = *reinterpret_cast<const float4*>(s0 + 8 * j + 4 * half);
+        s[4 * j] = t.x; s[4 * j + 1] = t.y; s[4 * j + 2] = t.z; s[4 * j + 3] = t.w;
+    }
+    Prefetch<QK, 2 * ROWS, ROWS, STEPS> pf;                // each stage's loads in flight during the one before
+    pf.fetch(q, k, v, g, beta, 0, W, hk, hv, head, key_head, blockIdx.y * ROWS);
+    for (int t0 = 0; t0 < W; t0 += STEPS) {
+        const int n = min(STEPS, W - t0);
+        __syncthreads();
+        pf.store(ks, qs, gs, bs, vs);
+        __syncthreads();
+        if (t0 + STEPS < W) pf.fetch(q, k, v, g, beta, t0 + STEPS, W, hk, hv, head, key_head, blockIdx.y * ROWS);
+        for (int tt = 0; tt < n; ++tt) {
+            const float gt = gs[tt], bt = bs[tt];
+            const size_t vat = (static_cast<size_t>(t0 + tt) * hv + head) * DV + row;
+            const float vt = __bfloat162float(vs[tt][threadIdx.x >> 1]);
+            float m[4] = {0.0f, 0.0f, 0.0f, 0.0f};
+#pragma unroll
+            for (int j = 0; j < 16; ++j) {
+                const float4 kk = ks[tt][2 * j + half];
+                s[4 * j] = s[4 * j] * gt;
+                s[4 * j + 1] = s[4 * j + 1] * gt;
+                s[4 * j + 2] = s[4 * j + 2] * gt;
+                s[4 * j + 3] = s[4 * j + 3] * gt;
+                m[0] = __fmaf_rn(s[4 * j], kk.x, m[0]);
+                m[1] = __fmaf_rn(s[4 * j + 1], kk.y, m[1]);
+                m[2] = __fmaf_rn(s[4 * j + 2], kk.z, m[2]);
+                m[3] = __fmaf_rn(s[4 * j + 3], kk.w, m[3]);
+            }
+            float mem = (m[0] + m[1]) + (m[2] + m[3]);
+            mem = mem + __shfl_xor_sync(0xffffffffu, mem, 1);
+            const float delta = (vt - mem) * bt;
+            float o[4] = {0.0f, 0.0f, 0.0f, 0.0f};
+#pragma unroll
+            for (int j = 0; j < 16; ++j) {
+                const float4 kk = ks[tt][2 * j + half], qq = qs[tt][2 * j + half];
+                s[4 * j] = __fmaf_rn(kk.x, delta, s[4 * j]);
+                s[4 * j + 1] = __fmaf_rn(kk.y, delta, s[4 * j + 1]);
+                s[4 * j + 2] = __fmaf_rn(kk.z, delta, s[4 * j + 2]);
+                s[4 * j + 3] = __fmaf_rn(kk.w, delta, s[4 * j + 3]);
+                o[0] = __fmaf_rn(s[4 * j], qq.x, o[0]);
+                o[1] = __fmaf_rn(s[4 * j + 1], qq.y, o[1]);
+                o[2] = __fmaf_rn(s[4 * j + 2], qq.z, o[2]);
+                o[3] = __fmaf_rn(s[4 * j + 3], qq.w, o[3]);
+            }
+            float out = (o[0] + o[1]) + (o[2] + o[3]);
+            out = out + __shfl_xor_sync(0xffffffffu, out, 1);
+            if (half == 0) y[vat] = __float2bfloat16_rn(out);
+        }
+    }
+    float* s1 = last + (static_cast<size_t>(head) * DV + row) * DK;
+#pragma unroll
+    for (int j = 0; j < 16; ++j)
+        *reinterpret_cast<float4*>(s1 + 8 * j + 4 * half) = make_float4(s[4 * j], s[4 * j + 1], s[4 * j + 2],
+                                                                        s[4 * j + 3]);
+}
+
+} // namespace tf_gdn_prefill
+
+// The instantiations gdn_prefill_cuda launches for bf16 keys: 128 value rows a block, or 64 when heads < SMs.
+template __global__ void tf_gdn_prefill::chain_kernel<__nv_bfloat16, 128, 32>(const __nv_bfloat16*,
+    const __nv_bfloat16*, const __nv_bfloat16*, const float*, const float*, const float*, float*, __nv_bfloat16*, int,
+    int, int);
+template __global__ void tf_gdn_prefill::chain_kernel<__nv_bfloat16, 64, 32>(const __nv_bfloat16*,
+    const __nv_bfloat16*, const __nv_bfloat16*, const float*, const float*, const float*, float*, __nv_bfloat16*, int,
+    int, int);

zig/kernels/cuda/prefill_attention.cu

@@ -243,3 +243,6 @@ pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric
 // Head dim 128, eight warps, eight query heads a block, eight staging slots (Nemotron's 16 heads a KV head).
 template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16*,
     const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, int, float);
+// Head dim 256, two query heads a block, two staging slots beside the staged queries (Qwen3.8's 6 heads a KV head).
+template __global__ void tf_prefill_attention::pattn_kernel<256, 8, 2, 2>(const __nv_bfloat16*,
+    const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, int, float);

zig/kernels/cuda/qmm_group.cu

@@ -339,3 +339,10 @@ __global__ void __launch_bounds__(WM * WN * 32) group_kernel(
 // The instantiations Nemotron's windows launch on sm_121 (tile 2: rows <= 16, bf16 out).
 template __global__ void tf_qmm_group::group_kernel<64, 16, 64, 1, 4, 8, false, false, false, false>(
     const __nv_bfloat16*, const float*, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int);
+// Tiles 3-5 (rows <= 32, <= 64, more), which a GB10 takes for wider verify rounds; tiles never change bits.
+template __global__ void tf_qmm_group::group_kernel<64, 32, 64, 1, 4, 4, false, false, false, false>(
+    const __nv_bfloat16*, const float*, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int);
+template __global__ void tf_qmm_group::group_kernel<64, 64, 64, 1, 4, 4, false, false, false, false>(
+    const __nv_bfloat16*, const float*, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int);
+template __global__ void tf_qmm_group::group_kernel<64, 64, 128, 2, 4, 3, false, false, false, false>(
+    const __nv_bfloat16*, const float*, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int);