diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 5602af100..676bda8c8 100644 --- a/zig/build/cuda.zig +++ b/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. diff --git a/zig/kernels/cuda/gdn_prefill.cu b/zig/kernels/cuda/gdn_prefill.cu new file mode 100644 index 000000000..a892ad52d --- /dev/null +++ b/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 +#include +#include + +namespace tf_gdn_prefill { + +constexpr int DK = 128; +constexpr int DV = 128; + +__device__ __forceinline__ float4 widen(uint2 w) { + const __nv_bfloat162 a = *reinterpret_cast(&w.x); + const __nv_bfloat162 b = *reinterpret_cast(&w.y); + return make_float4(__low2float(a), __high2float(a), __low2float(b), __high2float(b)); +} + +__device__ __forceinline__ float4 widen(float4 w) { return w; } + +template +struct Prefetch { + using Raw = typename std::conditional::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(min(t0 + st, W - 1)) * hk + key_head) * DK + 4 * c; + kr[it] = *reinterpret_cast(k + at); + qr[it] = *reinterpret_cast(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(v + (static_cast(min(t0 + st, W - 1)) * hv + head) * DV + + col0 + 8 * c); + } + if (threadIdx.x < STEPS) { + const size_t at = static_cast(min(t0 + static_cast(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(&vs[i / (ROWS / 8)][8 * (i % (ROWS / 8))]) = vr[it]; + } + if (threadIdx.x < STEPS) { + gs[threadIdx.x] = gr; + bs[threadIdx.x] = br; + } + } +}; + +template +__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(head) * DV + row) * DK; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const float4 t = *reinterpret_cast(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 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(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(head) * DV + row) * DK; +#pragma unroll + for (int j = 0; j < 16; ++j) + *reinterpret_cast(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); diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 122f40449..cd9a012fc 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/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); diff --git a/zig/kernels/cuda/qmm_group.cu b/zig/kernels/cuda/qmm_group.cu index f3f65e184..9c97ddfa7 100644 --- a/zig/kernels/cuda/qmm_group.cu +++ b/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); diff --git a/zig/kernels/cuda/qmm_prefill.cu b/zig/kernels/cuda/qmm_prefill.cu index 1610b2270..b90d3eb2a 100644 --- a/zig/kernels/cuda/qmm_prefill.cu +++ b/zig/kernels/cuda/qmm_prefill.cu @@ -151,3 +151,7 @@ __global__ void __launch_bounds__(WM * WN * 32) prefill_kernel( template __global__ void tf_qmm_prefill::prefill_kernel<64, 128, 128, 2, 4, 3, false>( const __nv_bfloat16*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, int, int, int, int, int, int); +// prompt_tile 9 (the 27B's: 128x128 on 2x2 warps, 2 stages, two blocks an SM on a GB10), bf16 out. +template __global__ void tf_qmm_prefill::prefill_kernel<64, 128, 128, 2, 2, 2, false>( + const __nv_bfloat16*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, int, int, int, + int, int, int); diff --git a/zig/kernels/cuda/qwen_ops.cu b/zig/kernels/cuda/qwen_ops.cu new file mode 100644 index 000000000..778cc1f16 --- /dev/null +++ b/zig/kernels/cuda/qwen_ops.cu @@ -0,0 +1,68 @@ +// The 27B's own small kernels: batched state copies, the drafter's taps and its fp32 RMSNorm. +#include +#include + +struct TfCopy { + unsigned long long dst, src, bytes; // 16-byte aligned pointers, bytes a multiple of 16 +}; + +// One block an item, its threads striding through the item's 16-byte words. +extern "C" __global__ void __launch_bounds__(256) tf_copy_batch(const TfCopy* __restrict__ items) { + const TfCopy it = items[blockIdx.x]; + const uint4* src = reinterpret_cast(it.src); + uint4* dst = reinterpret_cast(it.dst); + const unsigned long long n = it.bytes / 16; + for (unsigned long long i = threadIdx.x + static_cast(blockIdx.y) * blockDim.x; i < n; + i += static_cast(blockDim.x) * gridDim.y) + dst[i] = src[i]; +} + +// forward.tree_forward's taps: (x.float() + pending.float()).to(bfloat16), row r into columns [col, col + D) of OUT. +extern "C" __global__ void __launch_bounds__(256) tf_tap(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ pending, + __nv_bfloat16* __restrict__ out, int D, int out_stride, int col) { + const long long row = blockIdx.x; + for (int i = threadIdx.x + blockIdx.y * blockDim.x; i < D; i += blockDim.x * gridDim.y) + out[row * out_stride + col + i] = + __float2bfloat16_rn(__bfloat162float(x[row * D + i]) + __bfloat162float(pending[row * D + i])); +} + +// The drafter's RMSNorm in fp32: y = bf16(x * rsqrt(mean(x^2) + eps) * w), a 256-thread block a row. +extern "C" __global__ void __launch_bounds__(256) tf_rms_norm(const __nv_bfloat16* __restrict__ x, int x_stride, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ y, int D, float eps) { + __shared__ float part[8]; + const long long row = blockIdx.x; + const __nv_bfloat16* xr = x + row * x_stride; + float s = 0.0f; + for (int i = threadIdx.x; i < D; i += 256) { + const float v = __bfloat162float(xr[i]); + s += v * v; + } + for (int off = 16; off > 0; off >>= 1) s += __shfl_xor_sync(0xFFFFFFFFu, s, off); + if ((threadIdx.x & 31) == 0) part[threadIdx.x >> 5] = s; + __syncthreads(); + float total = 0.0f; + for (int k = 0; k < 8; ++k) total += part[k]; + const float inv = rsqrtf(total / D + eps); + for (int i = threadIdx.x; i < D; i += 256) + y[row * D + i] = __float2bfloat16_rn(__bfloat162float(xr[i]) * inv * __bfloat162float(w[i])); +} + +// The drafter's inverse frequencies as dflash2 makes them on the GPU: 1 / theta ** (2i / dim), fp32 powf then a divide. +extern "C" __global__ void tf_inv_freq(float* __restrict__ inv, float theta, int half, int dim) { + const int i = threadIdx.x; + if (i < half) inv[i] = 1.0f / powf(theta, static_cast(i) * 2.0f / static_cast(dim)); +} + +// dflash2._rotary and add_taps_streams' tables: cos and sin of float(position) * inv_freq, fp32, a block a row. +extern "C" __global__ void tf_rotary(const int* __restrict__ pos, const float* __restrict__ inv, + float* __restrict__ cos_out, float* __restrict__ sin_out, int half) { + const int row = blockIdx.x; + const float p = static_cast(pos[row]); + for (int i = threadIdx.x; i < half; i += blockDim.x) { + const float phase = p * inv[i]; + cos_out[row * half + i] = cosf(phase); + sin_out[row * half + i] = sinf(phase); + } +} diff --git a/zig/src/cli/qwen27_cuda.zig b/zig/src/cli/qwen27_cuda.zig new file mode 100644 index 000000000..a92416ce7 --- /dev/null +++ b/zig/src/cli/qwen27_cuda.zig @@ -0,0 +1,459 @@ +//! tf-qwen27-cuda: the 27B's Zig CUDA engine against the Python engine's capture (prompts.json, results.json). + +const std = @import("std"); +const cuda = @import("cuda"); +const qwen = @import("qwen27"); + +const usage = + \\usage: tf-qwen27-cuda run MODEL --kernels DIR --prompts prompts.json [--expect results.json] + \\ [--max-tokens N] [--context N] [--device N] + \\ Greedy, drafts off: each prompt's tokens and their sha against the capture's. + \\ tf-qwen27-cuda windows MODEL --kernels DIR --prompts windows.json + \\ capture_windows.py's chains and shared rounds: each row's logits against the capture's sha256. + \\ tf-qwen27-cuda draft MODEL --kernels DIR --drafter DIR --prompts prompts.json [--expect draft.json] + \\ draft_decode with DFlash2, greedy: tokens (equal to serial) and rounds against capture_draft.py's. + \\ tf-qwen27-cuda rounds MODEL --kernels DIR --prompts windows.json [--max-tokens ROUNDS] + \\ Greedy shared rounds of one row a stream over every prompt, timed. + \\ tf-qwen27-cuda multi MODEL --kernels DIR --drafter DIR --prompts multi.json + \\ capture_multi.py's drafted shared rounds: each round's windows and every reply against the capture's. + \\ +; + +const Options = struct { + model: []const u8, + kernels: []const u8 = "", + prompts: []const u8 = "", + expect: ?[]const u8 = null, + max_tokens: usize = 128, + context: usize = 16384, + device: c_int = 0, + dump: ?[]const u8 = null, // windows: a case's per-norm rows into this folder + cases: ?[]const u8 = null, // windows: only these cases, comma separated + drafter: ?[]const u8 = null, // draft: the DFlash2 checkpoint +}; + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const arena = init.arena.allocator(); + const args = try init.minimal.args.toSlice(arena); + const windows = args.len >= 2 and std.mem.eql(u8, args[1], "windows"); + const rounds = args.len >= 2 and std.mem.eql(u8, args[1], "rounds"); + const drafting = args.len >= 2 and std.mem.eql(u8, args[1], "draft"); + const multi = args.len >= 2 and std.mem.eql(u8, args[1], "multi"); + if (args.len < 3 or !(std.mem.eql(u8, args[1], "run") or windows or rounds or drafting or multi)) { + std.debug.print("{s}", .{usage}); + return 2; + } + var o: Options = .{ .model = args[2] }; + var i: usize = 3; + while (i + 1 < args.len) : (i += 2) { + const k = args[i]; + const v = args[i + 1]; + if (std.mem.eql(u8, k, "--kernels")) o.kernels = v else if (std.mem.eql(u8, k, "--prompts")) o.prompts = v else if (std.mem.eql(u8, k, "--expect")) o.expect = v else if (std.mem.eql(u8, k, "--max-tokens")) o.max_tokens = try std.fmt.parseInt(usize, v, 10) else if (std.mem.eql(u8, k, "--context")) o.context = try std.fmt.parseInt(usize, v, 10) else if (std.mem.eql(u8, k, "--device")) o.device = try std.fmt.parseInt(c_int, v, 10) else if (std.mem.eql(u8, k, "--dump")) o.dump = v else if (std.mem.eql(u8, k, "--cases")) o.cases = v else if (std.mem.eql(u8, k, "--drafter")) o.drafter = v else { + std.debug.print("unknown option {s}\n{s}", .{ k, usage }); + return 2; + } + } + if (o.kernels.len == 0 or o.prompts.len == 0) { + std.debug.print("{s}", .{usage}); + return 2; + } + const prompts = try readJson(arena, init.io, o.prompts); + const expect = if (o.expect) |p| try readJson(arena, init.io, p) else null; + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, o.device); + defer ctx.deinit(); + const clock = std.Io.Clock.awake; + var t0 = clock.now(init.io).toNanoseconds(); + const e = try qwen.Engine.init(gpa, init.io, &ctx, o.model, o.kernels, .{ .context = o.context }); + defer e.deinit(); + std.debug.print("loaded in {d:.1}s, weights {d:.2} GiB\n", .{ @as(f64, @floatFromInt(clock.now(init.io).toNanoseconds() - t0)) / 1e9, @as(f64, @floatFromInt(e.w.bytes())) / (1 << 30) }); + if (windows) return checkWindows(gpa, arena, init.io, e, prompts, o.dump, o.cases); + if (rounds) return benchRounds(gpa, arena, init.io, e, prompts, o.max_tokens); + if (drafting) return runDrafted(gpa, arena, init.io, e, o, prompts, expect); + if (multi) return runMulti(gpa, arena, init.io, e, o, prompts); + var seq = try e.sequence(o.context); + defer seq.deinit(); + var mismatches: usize = 0; + var it = prompts.object.iterator(); + while (it.next()) |entry| { + const name = entry.key_ptr.*; + const ids = try tokenList(arena, entry.value_ptr.*); + var out: std.ArrayList(u32) = .empty; + defer out.deinit(gpa); + t0 = clock.now(init.io).toNanoseconds(); + try e.generate(&seq, ids, o.max_tokens, &out); + const seconds = @as(f64, @floatFromInt(clock.now(init.io).toNanoseconds() - t0)) / 1e9; + const sha = try sha12(arena, out.items); + var verdict: []const u8 = ""; + if (expect) |x| blk: { + const want_v = (x.object.get("results") orelse break :blk).object.get(name) orelse break :blk; + const want = try tokenList(arena, want_v.object.get("tokens").?); + const same = std.mem.eql(u32, want, out.items); + if (!same) mismatches += 1; + var first: usize = 0; + while (first < @min(want.len, out.items.len) and want[first] == out.items[first]) first += 1; + verdict = if (same) " MATCH" else try std.fmt.allocPrint(arena, " DIFFER at token {d} (python sha {s})", .{ first, try sha12(arena, want) }); + } + std.debug.print("{s}: prompt {d}, {d} tokens sha {s} in {d:.2}s{s}\n", .{ name, ids.len, out.items.len, sha, seconds, verdict }); + } + return if (mismatches == 0) 0 else 1; +} + +fn readJson(a: std.mem.Allocator, io: std.Io, path: []const u8) !std.json.Value { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 30)); + return (try std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{})); +} + +fn tokenList(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + if (v != .array) return error.BadPromptFile; + const out = try a.alloc(u32, v.array.items.len); + for (v.array.items, out) |x, *t| t.* = @intCast(x.integer); + return out; +} + +/// The capture's sha12: sha256 of json.dumps(tokens), "[1, 2, 3]", its first 12 hex digits. +fn sha12(a: std.mem.Allocator, tokens: []const u32) ![]const u8 { + var text: std.ArrayList(u8) = .empty; + try text.append(a, '['); + for (tokens, 0..) |t, i| { + if (i > 0) try text.appendSlice(a, ", "); + try text.print(a, "{d}", .{t}); + } + try text.append(a, ']'); + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(text.items, &digest, .{}); + const hex = std.fmt.bytesToHex(digest, .lower); + return a.dupe(u8, hex[0..12]); +} + +/// Each case's streams prefilled afresh, its round run, every row's logits hashed; shared cases commit and go on. +fn checkWindows(gpa: std.mem.Allocator, a: std.mem.Allocator, io: std.Io, e: *qwen.Engine, doc: std.json.Value, dump: ?[]const u8, only: ?[]const u8) !u8 { + const root = doc.object; + const prompts = root.get("prompts").?.array.items; + const cases = root.get("cases").?.array.items; + var bad: usize = 0; + const row_bytes = qwen.config.vocab * 2; + const host = try gpa.alloc(u8, row_bytes); + defer gpa.free(host); + for (cases, 0..) |case, ci| { + if (only) |list| { + var it = std.mem.splitScalar(u8, list, ','); + const listed = while (it.next()) |x| { + if (std.fmt.parseInt(usize, x, 10) catch null == ci) break true; + } else false; + if (!listed) continue; + } + const streams = case.object.get("streams").?.array.items; + std.debug.print("case {d}: {s}, {d} streams\n", .{ ci, case.object.get("kind").?.string, streams.len }); + var seqs: [qwen.state.max_streams]qwen.state.Seq = undefined; + var parts: [qwen.state.max_streams]qwen.Part = undefined; + var made: usize = 0; + defer for (seqs[0..made]) |*sq| sq.deinit(); + for (streams, 0..) |sv, k| { + const sid: usize = @intCast(sv.object.get("stream").?.integer); + const ids = try tokenList(a, prompts[sid]); + seqs[k] = try e.sequence(4096); + made += 1; + try e.prefill(&seqs[k], ids); + parts[k] = .{ .seq = &seqs[k], .tokens = try tokenList(a, sv.object.get("tokens").?) }; + if (sv.object.get("parents")) |pv| { + const ps = try a.alloc(i32, pv.array.items.len); + for (pv.array.items, ps) |x, *y| y.* = @intCast(x.integer); + parts[k].parents = ps; + } + } + const n = streams.len; + var f = e.forward(); + if (dump) |dir| f.dump = .{ .io = io, .dir = dir }; + try f.round(parts[0..n]); + f.dump = null; + bad += try compareRows(&f, host, case.object.get("rows").?.array.items, ci, "rows"); + if (case.object.get("paths")) |paths_v| { + var paths: [qwen.state.max_streams][]const u32 = undefined; + var next: [qwen.state.max_streams]u32 = undefined; + for (paths_v.array.items, 0..) |pv, k| paths[k] = try tokenList(a, pv); + try f.commit(parts[0..n], paths[0..n]); + for (case.object.get("next").?.array.items, 0..) |nv, k| { + next[k] = @intCast(nv.integer); + parts[k].tokens = next[k .. k + 1]; + parts[k].parents = null; + } + try f.round(parts[0..n]); + bad += try compareRows(&f, host, case.object.get("after").?.array.items, ci, "after"); + } + if (case.object.get("keep")) |keep_v| { + var kept: [qwen.state.max_streams]usize = undefined; + var next: [qwen.state.max_streams]u32 = undefined; + const prefix = [_]u32{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15 }; + var paths: [qwen.state.max_streams][]const u32 = undefined; + for (keep_v.array.items, 0..) |kv, k| { + kept[k] = @intCast(kv.integer); + paths[k] = prefix[0..kept[k]]; + } + try f.commit(parts[0..n], paths[0..n]); + const reply_rows = doc.object.get("replies").?.array.items; + for (streams, 0..) |sv, k| { + const sid: usize = @intCast(sv.object.get("stream").?.integer); + next[k] = @intCast(reply_rows[sid].array.items[kept[k]].integer); + parts[k].tokens = next[k .. k + 1]; + } + try f.round(parts[0..n]); + bad += try compareRows(&f, host, case.object.get("after").?.array.items, ci, "after"); + } + } + std.debug.print("{d} cases, {d} rows differ\n", .{ cases.len, bad }); + return if (bad == 0) 0 else 1; +} + +fn compareRows(f: *qwen.Forward, host: []u8, want: []const std.json.Value, case: usize, what: []const u8) !usize { + var bad: usize = 0; + for (want, 0..) |rv, r| { + try f.ops.download(host, f.s.logits + r * host.len); + try f.ops.s.synchronize(); + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(host, &digest, .{}); + const hex = std.fmt.bytesToHex(digest, .lower); + if (!std.mem.eql(u8, &hex, rv.object.get("sha").?.string)) { + std.debug.print("case {d} {s} row {d}: logits differ\n", .{ case, what, r }); + bad += 1; + } + } + return bad; +} + +/// Every prompt prefilled, then greedy shared rounds of one row a stream; ms a round and tokens a second. +fn benchRounds(gpa: std.mem.Allocator, a: std.mem.Allocator, io: std.Io, e: *qwen.Engine, doc: std.json.Value, count: usize) !u8 { + { + // the served engine's cost curve: the forward's ms by a round's rows + var lanes_ = try qwen.Lanes.init(gpa, e, qwen.state.max_streams); + defer lanes_.deinit(); + try lanes_.calibrate(io, qwen.state.max_streams); + for (lanes_.costs[0..lanes_.cost_count]) |cst| std.debug.print("{d} rows: {d:.2} ms\n", .{ cst.width, cst.ms }); + } + const prompts = doc.object.get("prompts").?.array.items; + const n = @min(prompts.len, qwen.state.max_streams); + var seqs: [qwen.state.max_streams]qwen.state.Seq = undefined; + var parts: [qwen.state.max_streams]qwen.Part = undefined; + var next: [qwen.state.max_streams]u32 = undefined; + var ones: [qwen.state.max_streams][]const u32 = @splat(&.{0}); + var made: usize = 0; + defer for (seqs[0..made]) |*sq| sq.deinit(); + var f = e.forward(); + for (prompts[0..n], 0..) |pv, k| { + seqs[k] = try e.sequence(4096); + made += 1; + try e.prefill(&seqs[k], try tokenList(a, pv)); + try f.head(&seqs[k]); + var pick: [1]u32 = undefined; + try f.picks(1, &pick); + next[k] = pick[0]; + } + const clock = std.Io.Clock.awake; + const t0 = clock.now(io).toNanoseconds(); + for (0..count) |_| { + for (0..n) |k| parts[k] = .{ .seq = &seqs[k], .tokens = next[k .. k + 1] }; + try f.round(parts[0..n]); + try f.picks(n, next[0..n]); + try f.commit(parts[0..n], ones[0..n]); + } + try f.ops.s.synchronize(); + const ms = @as(f64, @floatFromInt(clock.now(io).toNanoseconds() - t0)) / 1e6; + std.debug.print("{d} streams, {d} rounds: {d:.2} ms a round, {d:.1} tok/s\n", .{ n, count, ms / @as(f64, @floatFromInt(count)), @as(f64, @floatFromInt(n * count)) / ms * 1e3 }); + return 0; +} + +/// Each prompt drafted with DFlash2 as capture_draft.py runs it: tokens, rounds and accepted drafts. +fn runDrafted(gpa: std.mem.Allocator, a: std.mem.Allocator, io: std.Io, e: *qwen.Engine, o: Options, prompts: std.json.Value, expect: ?std.json.Value) !u8 { + const clock = std.Io.Clock.awake; + var t0 = clock.now(io).toNanoseconds(); + const f0 = e.forward(); + const d = try qwen.draft.DFlash2.init(gpa, io, f0.ops, f0.t, &e.w, o.drafter orelse return error.NoDrafter, o.model); + defer d.deinit(); + std.debug.print("drafter loaded in {d:.1}s\n", .{@as(f64, @floatFromInt(clock.now(io).toNanoseconds() - t0)) / 1e9}); + var ctx = try qwen.draft.Context.init(f0.ops.k.d); + defer ctx.deinit(); + var seq = try e.sequence(o.context); + defer seq.deinit(); + var bad: usize = 0; + var it = prompts.object.iterator(); + while (it.next()) |entry| { + const name = entry.key_ptr.*; + const ids = try tokenList(a, entry.value_ptr.*); + var out: std.ArrayList(u32) = .empty; + defer out.deinit(gpa); + var stats: qwen.lone.Stats = .{}; + t0 = clock.now(io).toNanoseconds(); + try qwen.lone.prefill(e, &seq, &ctx, d, ids); + var f = e.forward(); + try f.head(&seq); + var pick: [1]u32 = undefined; + try f.picks(1, &pick); + const t1 = clock.now(io).toNanoseconds(); + try qwen.lone.decode(e, &seq, &ctx, d, ids, pick[0], o.max_tokens, 12, &out, &stats); + const secs = @as(f64, @floatFromInt(clock.now(io).toNanoseconds() - t1)) / 1e9; + var verdict: []const u8 = ""; + if (expect) |x| blk: { + const want_v = (x.object.get("results") orelse break :blk).object.get(name) orelse break :blk; + const want = try tokenList(a, want_v.object.get("tokens").?); + const same = std.mem.eql(u32, want, out.items); + if (!same) bad += 1; + verdict = try std.fmt.allocPrint(a, " {s}; python rounds {d} accepted {d}", .{ if (same) "MATCH" else "DIFFER", want_v.object.get("rounds").?.integer, want_v.object.get("accepted").?.integer }); + } + std.debug.print("{s}: {d} tokens sha {s}, rounds {d} accepted {d}, {d:.1} tok/s{s}\n", .{ name, out.items.len, try sha12(a, out.items), stats.rounds, stats.accepted, @as(f64, @floatFromInt(out.items.len)) / secs, verdict }); + } + return if (bad == 0) 0 else 1; +} + +/// One stream of the multi check: its state, drafting, context (prompt and reply) and reply. +const MultiStream = struct { + seq: qwen.state.Seq, + d: qwen.drafts.Stream, + context: std.ArrayList(u32) = .empty, + out: std.ArrayList(u32) = .empty, + done: bool = false, + rounds: usize = 0, + accepted: usize = 0, +}; + +/// MultiDecoder's rounds as capture_multi.py runs them (greedy, no costs: whole trees, blocks of 16 rows). +fn runMulti(gpa: std.mem.Allocator, a: std.mem.Allocator, io: std.Io, e: *qwen.Engine, o: Options, doc: std.json.Value) !u8 { + const clock = std.Io.Clock.awake; + const f0 = e.forward(); + const d = try qwen.draft.DFlash2.init(gpa, io, f0.ops, f0.t, &e.w, o.drafter orelse return error.NoDrafter, o.model); + defer d.deinit(); + const prompts = doc.object.get("prompts").?.array.items; + const count: usize = @intCast(doc.object.get("max_tokens").?.integer); + const want_rounds = doc.object.get("rounds").?.array.items; + const S = prompts.len; + const ms = try a.alloc(MultiStream, S); + var made: usize = 0; + defer for (ms[0..made]) |*m| { + m.seq.deinit(); + m.d.deinit(); + m.context.deinit(gpa); + m.out.deinit(gpa); + }; + for (prompts, ms) |pv, *m| { + const ids = try tokenList(a, pv); + m.* = .{ .seq = try e.sequence(o.context), .d = qwen.drafts.Stream.init(gpa, try qwen.draft.Context.init(f0.ops.k.d)) }; + made += 1; + try qwen.lone.prefill(e, &m.seq, &m.d.ctx, d, ids); + var f = e.forward(); + try f.head(&m.seq); + var pick: [1]u32 = undefined; + try f.picks(1, &pick); + try m.context.appendSlice(gpa, ids); + try m.context.append(gpa, pick[0]); + try m.out.append(gpa, pick[0]); + } + var differ: usize = 0; + var round: usize = 0; + const t0 = clock.now(io).toNanoseconds(); + var tokens: [qwen.state.round_rows]u32 = undefined; + var parents: [qwen.state.round_rows]i32 = undefined; + var picks: [qwen.state.round_rows]u32 = undefined; + var path_rows: [qwen.state.round_rows]u32 = undefined; + var kept_rows: [qwen.state.round_rows]u32 = undefined; + while (true) : (round += 1) { + var live: [qwen.state.max_streams]usize = undefined; + var n: usize = 0; + for (ms, 0..) |m, k| if (!m.done) { + live[n] = k; + n += 1; + }; + if (n == 0) break; + var asks: [qwen.state.max_streams]qwen.drafts.Ask = undefined; + for (live[0..n], 0..) |k, j| asks[j] = .{ .x = &ms[k].d, .context = ms[k].context.items, .sampling = null }; + try qwen.drafts.propose(d, asks[0..n]); + // the windows: the pending token, then the drafts (tree parents shifted under it) + var parts: [qwen.state.max_streams]qwen.Part = undefined; + var starts: [qwen.state.max_streams + 1]usize = undefined; + var total: usize = 0; + const want = if (round < want_rounds.len) want_rounds[round].array.items else &.{}; + for (live[0..n], 0..) |k, j| { + const m = &ms[k]; + starts[j] = total; + tokens[total] = m.out.items[m.out.items.len - 1]; + parents[total] = -1; + for (m.d.tokens.items, m.d.parents.items, 1..) |t, p, r| { + tokens[total + r] = t; + parents[total + r] = if (p < 0) 0 else p + 1; + } + const rows = 1 + m.d.tokens.items.len; + parts[j] = .{ .seq = &m.seq, .tokens = tokens[total..][0..rows], .parents = parents[total..][0..rows] }; + if (j < want.len) { + const w = want[j].object; + const wt = try tokenList(a, w.get("tokens").?); + const wp = w.get("parents").?.array.items; + var same = @as(usize, @intCast(w.get("stream").?.integer)) == k and std.mem.eql(u32, wt, tokens[total..][0..rows]) and wp.len == rows; + if (same) for (wp, parents[total..][0..rows]) |x, y| { + if (x.integer != y) same = false; + }; + if (!same) { + differ += 1; + if (differ == 1) std.debug.print("round {d} stream {d}: first window that differs from Python's ({s} {d} rows, python {s} {d})\n", .{ round, k, @tagName(m.d.kind), rows, w.get("mode").?.string, wt.len }); + } + } else differ += 1; + total += rows; + } + starts[n] = total; + var f = e.forward(); + f.taps = e.scratch.taps; + try f.round(parts[0..n]); + try f.picks(total, &picks); + var paths: [qwen.state.max_streams][]const u32 = undefined; + var kept: [qwen.state.max_streams]qwen.drafts.Kept = undefined; + for (live[0..n], 0..) |k, j| { + const m = &ms[k]; + const base = starts[j]; + const rows = starts[j + 1] - base; + const room = count - m.out.items.len; + // accept: children equal to their parent's draw, within the room, none past an end token + var len: usize = 1; + path_rows[base] = 0; + var terminal = picks[base]; + while (len < room and !e.w.isEos(terminal)) { + const child = for (1..rows) |r| { + if (parents[base + r] == @as(i32, @intCast(path_rows[base + len - 1])) and tokens[base + r] == terminal) break r; + } else break; + path_rows[base + len] = @intCast(child); + len += 1; + terminal = picks[base + child]; + } + paths[j] = path_rows[base..][0..len]; + for (paths[j], kept_rows[base..][0..len]) |r, *g| g.* = @intCast(base + r); + kept[j] = .{ .x = &m.d, .rows = kept_rows[base..][0..len] }; + for (paths[j][1..]) |r| { + try m.out.append(gpa, tokens[base + r]); + try m.context.append(gpa, tokens[base + r]); + } + try m.out.append(gpa, terminal); + try m.context.append(gpa, terminal); + m.rounds += 1; + m.accepted += len - 1; + } + try f.commit(parts[0..n], paths[0..n]); + try qwen.drafts.absorb(d, f.ops, e.scratch.taps, e.scratch.copy_items, kept[0..n]); + for (live[0..n]) |k| { + const m = &ms[k]; + if (m.out.items.len >= count or e.w.isEos(m.out.items[m.out.items.len - 1])) m.done = true; + } + } + const secs = @as(f64, @floatFromInt(clock.now(io).toNanoseconds() - t0)) / 1e9; + if (round != want_rounds.len) differ += 1; + var bad: usize = 0; + var produced: usize = 0; + const results = doc.object.get("results").?.array.items; + for (ms, results, 0..) |m, rv, k| { + const want = try tokenList(a, rv.object.get("tokens").?); + const same = std.mem.eql(u32, want, m.out.items); + if (!same or m.rounds != rv.object.get("rounds").?.integer) bad += 1; + produced += m.out.items.len; + std.debug.print("stream {d}: {d} tokens sha {s} rounds {d} accepted {d} {s}; python rounds {d} accepted {d}\n", .{ k, m.out.items.len, try sha12(a, m.out.items), m.rounds, m.accepted, if (same) "MATCH" else "DIFFER", rv.object.get("rounds").?.integer, rv.object.get("accepted").?.integer }); + } + // replies and every stream's rounds must match; a tree may order its nodes otherwise where drafts round otherwise + std.debug.print("{d} rounds (python {d}), {d} windows differ from Python's, {d} replies differ; {d:.1} tok/s\n", .{ round, want_rounds.len, differ, bad, @as(f64, @floatFromInt(produced)) / secs }); + return if (bad == 0 and round == want_rounds.len) 0 else 1; +} diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index dc4e4db11..27dedfdc7 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -319,18 +319,60 @@ pub const LaneHost = struct { h.mutex.unlock(h.io); } - /// Prefills the next queued request into a free lane; false when none waits or no lane is free. - fn admitOne(h: *LaneHost) bool { + /// Every queued request the lanes have room for: a lone one to the backend's driver, the rest prefilled together. + fn admitAll(h: *LaneHost) void { + var batch: std.ArrayList(*Job) = .empty; + defer batch.deinit(h.gpa); + while (true) { + const job = switch (h.prepare()) { + .none => break, + .dropped => continue, + .job => |j| j, + }; + if (batch.items.len == 0 and h.loneFits(job)) { + _ = h.runLone(job, job.began); + continue; + } + batch.append(h.gpa, job) catch { + _ = h.drop(job, "out of memory"); + continue; + }; + } + if (batch.items.len == 0) return; + const streams = h.gpa.alloc(*lanes.Stream, batch.items.len) catch return h.dropAll(batch.items, "out of memory"); + defer h.gpa.free(streams); + const errs = h.gpa.alloc(?anyerror, batch.items.len) catch return h.dropAll(batch.items, "out of memory"); + defer h.gpa.free(errs); + for (streams, batch.items) |*s, job| s.* = &job.stream; + h.core.addStreams(streams, errs) catch |e| return h.dropAll(batch.items, h.words(e)); + for (batch.items, errs) |job, err| { + if (err) |e| { + _ = if (e == error.Cancelled) h.cancel(job) else h.drop(job, h.words(e)); + continue; + } + h.prefilled(job, job.began); + if (h.deliver(job)) h.remove(job); + } + } + + fn dropAll(h: *LaneHost, jobs: []const *Job, message: []const u8) void { + for (jobs) |job| _ = h.drop(job, message); + } + + const Prepared = union(enum) { none, dropped, job: *Job }; + + /// The oldest queued request admitted with its stream made, when the lanes have room. + fn prepare(h: *LaneHost) Prepared { h.lock(); if (h.queued.items.len == 0 or h.admitted.items.len >= h.info_.lanes) { h.unlock(); - return false; + return .none; } const job = h.queued.orderedRemove(0); h.admitted.append(h.gpa, job) catch { h.unlock(); h.finish(job, .failed, "out of memory"); - return true; + return .dropped; }; h.unlock(); const r = job.request; @@ -342,7 +384,10 @@ pub const LaneHost = struct { job.marks = l.marks; reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; } else |_| {}; - job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start"); + job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch { + _ = h.drop(job, "the drafter could not start"); + return .dropped; + }; job.stream = lanes.Stream.init(h.gpa, .{ .id = "request", .prompt = r.prompt, @@ -362,16 +407,12 @@ pub const LaneHost = struct { .reuse = reuse, }) catch { job.proposer.deinit(); - return h.drop(job, "out of memory"); + _ = h.drop(job, "out of memory"); + return .dropped; }; job.started = true; - const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); - job.began = began; - if (h.loneFits(job)) return h.runLone(job, began); - h.core.addStream(&job.stream) catch |e| return if (e == error.Cancelled) h.cancel(job) else h.drop(job, h.words(e)); - h.prefilled(job, began); - if (h.deliver(job)) h.remove(job); - return true; + job.began = std.Io.Clock.awake.now(h.io).toNanoseconds(); + return .{ .job = job }; } fn prefilled(h: *LaneHost, job: *Job, began: i96) void { @@ -480,7 +521,7 @@ pub const LaneHost = struct { fn run(h: *LaneHost) void { while (true) { h.takeCancels(); - while (h.admitOne()) {} + h.admitAll(); h.noteLive(); h.lock(); if (h.closing) { diff --git a/zig/src/core/lanes/admit.zig b/zig/src/core/lanes/admit.zig new file mode 100644 index 000000000..37ec717b8 --- /dev/null +++ b/zig/src/core/lanes/admit.zig @@ -0,0 +1,90 @@ +//! Admitting streams: prompt passes (shared, as multi.py's _fill_batch), first tokens, then first drafts. +const std = @import("std"); +const be = @import("backend.zig"); +const ev = @import("events.zig"); +const win = @import("windows.zig"); +const trail = @import("trail.zig"); +const Engine = @import("engine.zig").Engine; +const Stream = @import("stream.zig").Stream; +const LogRow = @import("logprob.zig").Row; +const Feed = be.Feed; +const f = ev.f; +const str = trail.str; +const int = trail.int; + +/// Prompt passes (shared when the backend has them), first tokens, every stream's first drafts; failures in `errs`. +pub fn streams(e: *Engine, list: []const *Stream, errs: []?anyerror) !void { + _ = e.arena.reset(.retain_capacity); + const a = e.arena.allocator(); + @memset(errs, null); + for (list) |s| try trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }); + if (list.len > 1 and e.backend.vtable.prefill_many != null) { + try e.backend.vtable.prefill_many.?(e.backend.ptr, list, errs); + } else for (list, errs) |s, *err| e.backend.prefill(s) catch |x| { + err.* = x; + }; + const feeds = try a.alloc(Feed, list.len); + const drawn = try a.alloc(u64, list.len); + const asked = try a.alloc(?u32, list.len); + var requests: std.ArrayList(be.DraftRequest) = .empty; + for (list, errs, feeds, drawn, asked) |s, *err, *feed, *handle, *ask| { + ask.* = null; + if (err.*) |x| { + if (x == error.Cancelled) e.backend.release(s); // the host finishes a cancelled stream without the core + continue; + } + if (s.isCancelled()) { // cancelled in its last chunk: no first token + e.backend.release(s); + err.* = error.Cancelled; + continue; + } + s.context.shrinkRetainingCapacity(s.prompt_len); + s.rows.clearRetainingCapacity(); + s.pending = null; + s.cache_len = s.prompt_len; + const position: u64 = s.prompt_len; + handle.* = e.backend.first(s, position) catch |x| { + err.* = x; // the host discards it, as after a failed prompt pass + continue; + }; + feed.* = .{ .handle = handle.* }; + if (try e.forcedNext(s)) |t| feed.* = .{ .value = t }; + if (e.cfg.family_mtp and s.drafts) { + // the head reads the prompt's last row and the first token, and drafts the one after it + const d: u32 = @intCast(try e.rule.depth(win.who(s))); + ask.* = d; + try requests.append(a, .{ .stream = s, .follow = &.{}, .first = feed.*, .rows = null, .start = s.prompt_len, .position = position + 1, .depth = d }); + } else if (e.cfg.pipelined and s.logprobs == null) { + try e.queueNext(s, feed.*); + } + } + if (requests.items.len > 0) try e.backend.draft(requests.items); + for (list, errs, feeds, drawn, asked) |s, err, feed, handle, ask| { + if (err != null) continue; + const position: u64 = s.prompt_len; + if (ask) |d| { + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { + s.next = held; // a tree head's first drafts as host tokens, as every later round's + }; + } + const value = try e.readFeed(feed); + if (ask) |d| try trail.event(e, &.{ f("ev", str("draft")), f("stream", str(s.id)), f("depth", int(d)), f("position", int(position + 1)), f("follow", .{ .u32s = &.{value} }), f("rows", .null) }); + if (e.log != null) { + const first = if (feed == .handle) value else try e.backend.read(handle); + try trail.event(e, &.{ f("ev", str("first")), f("stream", str(s.id)), f("position", int(position)), f("drawn", int(first)), f("token", int(value)) }); + } + var first_row: [1]LogRow = undefined; + if (s.logprobs != null) first_row[0] = (try e.backend.firstRow(s)).forToken(value); + _ = try s.commit(e.gpa, &.{value}, if (s.logprobs != null) &first_row else &.{}); + s.pending = value; + try trail.resolve(e); + if (s.finished) { + try trail.finish(e, s); + e.release(s); + continue; + } + try e.live.append(e.gpa, s); + } +} diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index a10ced867..c50ca6b10 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -65,6 +65,8 @@ pub const Backend = struct { pub const VTable = struct { /// Prefill the stream's prompt into its caches (and its draft head's). prefill: *const fn (ptr: *anyopaque, s: *Stream) anyerror!void, + /// Prefill several streams in shared passes, each with `prefill`'s bits; `errs[i]` for a failure. + prefill_many: ?*const fn (ptr: *anyopaque, streams: []const *Stream, errs: []?anyerror) anyerror!void = null, /// Draw the first token from the prompt's last row at `position`; a handle to it. first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64, /// Feed one token and queue the draw of the next at `position` (a one-token round); a handle to it. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index b752d84e8..f2ad1fd57 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -12,6 +12,7 @@ const shape = @import("shape.zig"); const plan_lanes = @import("plan_lanes.zig"); const fill = @import("fill.zig"); const trail = @import("trail.zig"); +const admit = @import("admit.zig"); const Stream = sm.Stream; const Feed = be.Feed; const LogRow = @import("logprob.zig").Row; @@ -78,55 +79,14 @@ pub const Engine = struct { /// Prefill a stream, draw its first token and ask for its first drafts; it takes part from the next round. pub fn addStream(e: *Engine, s: *Stream) !void { - _ = e.arena.reset(.retain_capacity); - try trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }); - e.backend.prefill(s) catch |err| { - if (err == error.Cancelled) e.backend.release(s); // the host finishes a cancelled stream without the core - return err; - }; - if (s.isCancelled()) { // cancelled in its last chunk: no first token - e.backend.release(s); - return error.Cancelled; - } - s.context.shrinkRetainingCapacity(s.prompt_len); - s.rows.clearRetainingCapacity(); - s.pending = null; - s.cache_len = s.prompt_len; - const position: u64 = s.prompt_len; - const drawn = try e.backend.first(s, position); - var feed: Feed = .{ .handle = drawn }; - if (try e.forcedNext(s)) |t| feed = .{ .value = t }; - var asked: ?u32 = null; - if (e.cfg.family_mtp and s.drafts) { - // the head reads the prompt's last row and the first token, and drafts the one after it - const d: u32 = @intCast(try e.rule.depth(win.who(s))); - asked = d; - try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = feed, .rows = null, .start = s.prompt_len, .position = position + 1, .depth = d }}); - s.dropHeld(e.gpa); - s.next = .{ .count = d }; - if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { - s.next = held; // a tree head's first drafts as host tokens, as every later round's - }; - } else if (e.cfg.pipelined and s.logprobs == null) { - try e.queueNext(s, feed); - } - const value = try e.readFeed(feed); - if (asked) |d| try trail.event(e, &.{ f("ev", str("draft")), f("stream", str(s.id)), f("depth", int(d)), f("position", int(position + 1)), f("follow", .{ .u32s = &.{value} }), f("rows", .null) }); - if (e.log != null) { - const first = if (feed == .handle) value else try e.backend.read(drawn); - try trail.event(e, &.{ f("ev", str("first")), f("stream", str(s.id)), f("position", int(position)), f("drawn", int(first)), f("token", int(value)) }); - } - var first_row: [1]LogRow = undefined; - if (s.logprobs != null) first_row[0] = (try e.backend.firstRow(s)).forToken(value); - _ = try s.commit(e.gpa, &.{value}, if (s.logprobs != null) &first_row else &.{}); - s.pending = value; - try trail.resolve(e); - if (s.finished) { - try trail.finish(e, s); - e.release(s); - return; - } - try e.live.append(e.gpa, s); + var errs = [_]?anyerror{null}; + try e.addStreams(&.{s}, &errs); + if (errs[0]) |err| return err; + } + + /// addStream for several streams, their prompt passes shared (admit.zig). + pub fn addStreams(e: *Engine, streams: []const *Stream, errs: []?anyerror) !void { + return admit.streams(e, streams, errs); } /// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. @@ -568,7 +528,7 @@ pub const Engine = struct { } /// Feed a token and queue the draw of the next at the new cache length. - fn queueNext(e: *Engine, s: *Stream, feed: Feed) !void { + pub fn queueNext(e: *Engine, s: *Stream, feed: Feed) !void { s.cache_len += 1; const h = try e.backend.queue(s, feed, s.cache_len); s.inflight = h; @@ -579,20 +539,20 @@ pub const Engine = struct { } /// The thinking budget's or a forced fix's token at the next position instead of the draw. - fn forcedNext(e: *Engine, s: *Stream) !?u32 { + pub fn forcedNext(e: *Engine, s: *Stream) !?u32 { if (s.popForce()) |t| return t; if (s.cutsNext()) return try s.startClose(e.gpa); return null; } - fn readFeed(e: *Engine, feed: Feed) !u32 { + pub fn readFeed(e: *Engine, feed: Feed) !u32 { return switch (feed) { .handle => |h| try e.backend.read(h), .value => |v| v, }; } - fn release(e: *Engine, s: *Stream) void { + pub fn release(e: *Engine, s: *Stream) void { s.dropRoundState(e.gpa); e.backend.release(s); } diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 8c0a38405..ac6c02628 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -13,6 +13,7 @@ fn Blob(comptime import_name: []const u8) type { pub const available = options.with_kernels; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; +pub const gdn_prefill: []const u8 = if (available) &Blob("fatbin_gdn_prefill").bytes else &.{}; 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_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").bytes else &.{}; @@ -24,6 +25,7 @@ pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; pub const lane_gemv: []const u8 = if (available) &Blob("fatbin_lane_gemv").bytes else &.{}; pub const affine4_pack: []const u8 = if (available) &Blob("fatbin_affine4_pack").bytes else &.{}; +pub const qwen_ops: []const u8 = if (available) &Blob("fatbin_qwen_ops").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const fp8_lane: []const u8 = if (available) &Blob("fatbin_fp8_lane").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; @@ -39,6 +41,7 @@ pub const nemotron_attention: []const u8 = if (available) &Blob("fatbin_nemotron pub const nemotron_keyed: []const u8 = if (available) &Blob("fatbin_nemotron_keyed").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; pub const train_mixers: []const u8 = if (available) &Blob("fatbin_train_mixers").bytes else &.{}; +pub const torch_norm_rotary: []const u8 = if (available) &Blob("fatbin_torch_norm_rotary").bytes else &.{}; /// Symbols in the gdn image as cuobjdump lists them for the built fatbin (named namespace tf_gdn). pub const gdn_symbols = struct { diff --git a/zig/src/families/nemotron/cuda.zig b/zig/src/families/nemotron/cuda.zig index 7177181a4..5bb9a91c9 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -26,6 +26,10 @@ pub const train = @import("cuda_train.zig"); pub const train_ops = @import("cuda_train_ops.zig"); pub const sites = @import("cuda_sites.zig"); pub const slide_dims = @import("slide_dims.zig"); +/// Checkpoint bytes to the GPU with direct reads in flight: shared with the other CUDA families. +pub const source = @import("cuda_source.zig"); +/// The torch-op replacements (argmax, topk, casts): shared with the other CUDA families. +pub const torch_ops = @import("cuda_torch_ops.zig"); test { _ = @import("config.zig"); diff --git a/zig/src/families/nemotron/cuda_source.zig b/zig/src/families/nemotron/cuda_source.zig index ffc4ba2ee..dcf89a40d 100644 --- a/zig/src/families/nemotron/cuda_source.zig +++ b/zig/src/families/nemotron/cuda_source.zig @@ -33,9 +33,20 @@ const Job = struct { } }; +/// Where a Source's copies go: the driver, the stream, and whether the card has its own memory. +pub const Target = struct { + d: *const cuda.Driver, + s: cuda.Stream, + discrete: bool, + + pub fn of(ops: kern.Ops) Target { + return .{ .d = ops.k.d, .s = ops.s, .discrete = ops.k.discrete }; + } +}; + pub const Source = struct { gpa: std.mem.Allocator, - ops: kern.Ops, + ops: Target, files: std.ArrayList(Mapped) = .empty, slots: [slot_count]dio.Buffer = undefined, pinned: [slot_count]?cuda.HostBuffer = @splat(null), // a discrete card's slots, page-locked: copies run at the link's speed @@ -48,21 +59,21 @@ pub const Source = struct { next: u64 = 0, /// Page-aligned slots; on a discrete card page-locked, each with an event its copies record. - pub fn init(gpa: std.mem.Allocator, ops: kern.Ops) !Source { + pub fn init(gpa: std.mem.Allocator, ops: Target) !Source { var s: Source = .{ .gpa = gpa, .ops = ops }; var made: usize = 0; errdefer s.freeSlots(made); while (made < slot_count) : (made += 1) { - if (!ops.k.discrete) { + if (!ops.discrete) { s.slots[made] = try gpa.alignedAlloc(u8, .fromByteUnits(dio.alignment), slot_bytes); continue; } - var h = try cuda.HostBuffer.alloc(ops.k.d, slot_bytes); + var h = try cuda.HostBuffer.alloc(ops.d, slot_bytes); if (@intFromPtr(h.bytes.ptr) % dio.alignment != 0) { h.free(); return error.UnalignedPinnedSlot; } - s.copied[made] = cuda.Event.init(ops.k.d, false) catch |err| { + s.copied[made] = cuda.Event.init(ops.d, false) catch |err| { h.free(); return err; }; @@ -126,7 +137,8 @@ pub const Source = struct { pub fn upload(s: *Source, dst: u64, bytes: []const u8) !void { const at = s.locate(bytes) orelse { try s.flush(); - return s.ops.upload(dst, bytes); + if (bytes.len == 0) return; + return s.ops.d.check(s.ops.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.ops.s.handle), "cuMemcpyHtoDAsync"); }; var done: usize = 0; while (done < bytes.len) { @@ -212,7 +224,7 @@ pub const Source = struct { /// Page-locked slots record their copy event; pageable slots wait for the driver to stage their bytes. fn finish(s: *Source, k: usize, p: Pending, got: []u8) !void { if (got.len == 0) return; - try s.ops.k.d.check(s.ops.k.d.api.cuMemcpyHtoDAsync_v2(p.dst, got.ptr, got.len, s.ops.s.handle), "cuMemcpyHtoDAsync"); + try s.ops.d.check(s.ops.d.api.cuMemcpyHtoDAsync_v2(p.dst, got.ptr, got.len, s.ops.s.handle), "cuMemcpyHtoDAsync"); if (s.copied[k]) |e| return e.record(s.ops.s); try s.ops.s.synchronize(); } diff --git a/zig/src/families/nemotron/cuda_weights.zig b/zig/src/families/nemotron/cuda_weights.zig index e3a6f6438..0319ba41e 100644 --- a/zig/src/families/nemotron/cuda_weights.zig +++ b/zig/src/families/nemotron/cuda_weights.zig @@ -352,7 +352,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, if (c.group_size != 64 or c.bits != 4) return error.UnsupportedQuantization; var w: Weights = .{ .gpa = gpa, .config = c }; errdefer w.deinit(); - var src = try Source.init(gpa, ops); + var src = try Source.init(gpa, .of(ops)); defer src.deinit(); var L: Loader = .{ .gpa = gpa, .ops = ops, .w = &w, .src = &src, .scratch = try cuda.DeviceBuffer.alloc(ops.k.d, 1 << 20) }; defer { diff --git a/zig/src/families/qwen27/cuda.zig b/zig/src/families/qwen27/cuda.zig new file mode 100644 index 000000000..3532c69cd --- /dev/null +++ b/zig/src/families/qwen27/cuda.zig @@ -0,0 +1,33 @@ +//! Qwen3.8-27B on CUDA in Zig: the Python 0.6.6 engine's kernels and layouts, beside the family's Metal engine. + +pub const config = @import("cuda/shape.zig"); +pub const kernels = @import("cuda/kernels.zig"); +pub const triton = @import("cuda/triton.zig"); +pub const weights = @import("cuda/weights.zig"); +pub const state = @import("cuda/state.zig"); +pub const Forward = @import("cuda/forward.zig").Forward; +pub const Part = @import("cuda/forward.zig").Part; +pub const engine = @import("cuda/engine.zig"); +pub const Engine = engine.Engine; +pub const Lanes = @import("cuda/lanes.zig").Cuda; +pub const native = @import("cuda/native.zig"); +pub const draft = @import("cuda/draft.zig"); +pub const lone = @import("cuda/lone.zig"); +pub const drafts = @import("cuda/drafts.zig"); + +test { + _ = config; + _ = kernels; + _ = triton; + _ = weights; + _ = state; + _ = engine; + _ = @import("cuda/forward.zig"); + _ = @import("cuda/lanes.zig"); + _ = @import("cuda/tree.zig"); + _ = @import("cuda/draft_load.zig"); + _ = @import("cuda/draft_policy.zig"); + _ = draft; + _ = lone; + _ = native; +} diff --git a/zig/src/families/qwen27/cuda/draft.zig b/zig/src/families/qwen27/cuda/draft.zig new file mode 100644 index 000000000..587554ef0 --- /dev/null +++ b/zig/src/families/qwen27/cuda/draft.zig @@ -0,0 +1,423 @@ +//! DFlash2 on CUDA as dflash2's fast path runs it: a context from the target's taps, a masked block, the host tree. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const aot = cuda.aot; +const c = @import("shape.zig"); +const kern = @import("kernels.zig"); +const tri = @import("triton.zig"); +const st = @import("state.zig"); +const wts = @import("weights.zig"); +const dl = @import("draft_load.zig"); +const policy = @import("draft_policy.zig"); + +const bf = 2; +const H = dl.hidden; +const D = dl.head_dim; +pub const block = 16; // masked rows a round drafts at most (dflash2.block) +/// Block rows a batch takes: every stream's block in one forward (launch_blocks). +pub const most_rows = block * st.max_streams; +/// Copy items a batch takes: every layer's keys and values of every stream, two pieces a head (or the head's rows). +const most_copies = @max(dl.layers * st.max_streams * 2 * dl.kv_heads * 2, 2 * most_rows); + +/// A stream's drafting context: each layer's keys and values ([kv heads, len, head dim], two buffers to swap). +pub const Context = struct { + arena: st.Arena, + keys: [dl.layers][2]u64, + values: [dl.layers][2]u64, + cur: u1 = 0, + len: usize = 0, // rows held (at most dl.window) + end: usize = 0, // the target position after the last absorbed row + + pub fn bytes() usize { + return dl.layers * 4 * std.mem.alignForward(usize, dl.kv_heads * dl.window * D * bf, 256); + } + + pub fn init(d: *const cuda.Driver) !Context { + var x: Context = .{ .arena = try st.Arena.init(d, bytes()), .keys = undefined, .values = undefined }; + for (0..dl.layers) |l| for (0..2) |j| { + x.keys[l][j] = try x.arena.take(dl.kv_heads * dl.window * D * bf); + x.values[l][j] = try x.arena.take(dl.kv_heads * dl.window * D * bf); + }; + return x; + } + + pub fn deinit(x: *Context) void { + x.arena.deinit(); + } + + /// dflash2.skip: `n` rows go untapped, so the context starts again after them. + pub fn skip(x: *Context, n: usize) void { + x.len = 0; + x.end += n; + } +}; + +/// One stream's share of a batch absorb: its context and how many of the gathered tap rows are its own (in turn). +pub const Absorb = struct { ctx: *Context, rows: usize }; +/// One stream's block: its context and the pending token heading it. +pub const Block = struct { ctx: *const Context, pending: u32 }; + +pub const DFlash2 = struct { + gpa: std.mem.Allocator, + w: dl.Weights, + target: *const wts.Weights, + ops: kern.Ops, + t: tri.Tri, + arena: st.Arena, + inv_freq: u64, // (D / 2) fp32 on the device, as dflash2 makes them + positions: []i32, // a call's row positions on the host + // a block's buffers (block rows) and absorbing's (up to the window's rows) + x: u64, + normed: u64, + conv: u64, + dyn: u64, + xs: u64, + qkv: u64, + q: u64, + k: u64, + v: u64, + attn: u64, + o: u64, + gate: u64, + up: u64, + act: u64, + act_xs: u64, + head_h: u64, + selected: u64, + logits: u64, + logits_tail: u64, + logits32: u64, + cand_values: u64, + cand_ids: u64, + topk_scratch: u64, + pos: u64, + cos: u64, + sin: u64, + table: u64, + copies: u64, + ids: u64, + // host copies of a round's candidates + host_ids: []i64, + host_values: []f32, + host_selected: []u16, + cands: []u32, + unary: []f64, + projected: []f64, + tree: policy.Tree = .{}, + launched: usize = 0, // the last batch's block rows a stream (0: none) + + pub fn init(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, t: tri.Tri, target: *const wts.Weights, dir: []const u8, target_dir: []const u8) !*DFlash2 { + const x = try gpa.create(DFlash2); + errdefer gpa.destroy(x); + x.* = undefined; + x.gpa = gpa; + x.ops = ops; + x.t = t; + x.target = target; + x.tree = .{}; + x.w = try dl.load(gpa, io, ops, dir, target_dir); + errdefer x.w.deinit(); + const R = dl.window + 1; // rows absorbed at once at most (the prompt's tail) + const W = most_rows; + x.arena = .{ .buf = null }; + for (0..2) |pass| { + if (pass == 1) x.arena = try st.Arena.init(ops.k.d, x.arena.used); + x.arena.used = 0; + const A = &x.arena; + x.x = try A.take(R * H * bf); + x.normed = try A.take(R * H * bf); + x.conv = try A.take(W * H * bf); + x.dyn = try A.take(W * 4 * (H / dl.conv_group) * bf); + x.xs = try A.take(R * (H * 5) / 64 * 4); + x.qkv = try A.take(R * (dl.heads + 2 * dl.kv_heads) * D * bf); + x.q = try A.take(W * dl.heads * D * bf); + x.k = try A.take(R * dl.kv_heads * D * bf); + x.v = try A.take(R * dl.kv_heads * D * bf); + x.attn = try A.take(W * dl.heads * D * bf); + x.o = try A.take(W * H * bf); + x.gate = try A.take(W * 17408 * bf); + x.up = try A.take(W * 17408 * bf); + x.act = try A.take(W * 17408 * bf); + x.act_xs = try A.take(W * 17408 / 64 * 4); + x.head_h = try A.take(W * H * bf); + x.selected = try A.take(W * dl.selector_rank * bf); + x.logits = try A.take(W * dl.head_rows * bf); + x.logits_tail = try A.take(W * (dl.spans[1][1] - dl.spans[1][0]) * bf); + x.logits32 = try A.take(W * dl.head_rows * 4); + x.cand_values = try A.take(W * policy.top * 4); + x.cand_ids = try A.take(W * policy.top * 8); + x.topk_scratch = try A.take(kern.torch_ops.topkScratchBytes(W, dl.head_rows)); + x.inv_freq = try A.take(D / 2 * 4); + x.pos = try A.take(R * 4); + x.cos = try A.take(R * D / 2 * 4); + x.sin = try A.take(R * D / 2 * 4); + x.table = try A.take(dl.layers * st.max_streams * 24); + x.copies = try A.take(most_copies * @sizeOf(kern.Copy)); + x.ids = try A.take(W * 4); + } + errdefer x.arena.deinit(); + { + var a: cuda.Args = .{}; + a.add(x.inv_freq); + a.add(@as(f32, @floatCast(dl.theta))); + a.add(@as(c_int, D / 2)); + a.add(@as(c_int, D)); + try cuda.launch.launch(ops.k.inv_freq, .{ .grid = .{ .x = 1 }, .block = .{ .x = D / 2 } }, ops.s, &a); + } + x.positions = try gpa.alloc(i32, dl.window + 1); + errdefer gpa.free(x.positions); + x.host_ids = try gpa.alloc(i64, W * policy.top); + x.host_values = try gpa.alloc(f32, W * policy.top); + x.host_selected = try gpa.alloc(u16, W * dl.selector_rank); + x.cands = try gpa.alloc(u32, W * policy.top); + x.unary = try gpa.alloc(f64, W * policy.top); + x.projected = try gpa.alloc(f64, W * dl.selector_rank); + return x; + } + + pub fn deinit(x: *DFlash2) void { + x.gpa.free(x.positions); + x.gpa.free(x.host_ids); + x.gpa.free(x.host_values); + x.gpa.free(x.host_selected); + x.gpa.free(x.cands); + x.gpa.free(x.unary); + x.gpa.free(x.projected); + x.tree.tokens.deinit(x.gpa); + x.tree.parents.deinit(x.gpa); + x.tree.scores.deinit(x.gpa); + x.arena.deinit(); + x.w.deinit(); + const gpa = x.gpa; + gpa.destroy(x); + } + + /// _lin: a packed projection of `rows` rows (each row's bits its own at any count), group sums first. + fn lin(x: *DFlash2, in: u64, q: kern.QLinear, out: u64, rows: usize) !void { + try x.t.groupSums(in, x.xs, q.k, rows, q.k); + try x.ops.dense(in, x.xs, q, out, rows); + } + + /// F.rms_norm over `rows` contiguous rows of the hidden width (the layers' and the head's norms). + fn torchNorm(x: *DFlash2, in: u64, weight: u64, out: u64, rows: usize) !void { + var a: cuda.Args = .{}; + a.add(in); + a.add(weight); + a.add(out); + a.add(@as(u64, 0)); + a.add(@as(f32, 1e-6)); + try cuda.launch.launch(x.ops.k.torch_rms, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 128 } }, x.ops.s, &a); + } + + fn rmsNorm(x: *DFlash2, in: u64, stride: usize, weight: u64, out: u64, rows: usize) !void { + var a: cuda.Args = .{}; + a.add(in); + a.add(@as(c_int, @intCast(stride))); + a.add(weight); + a.add(out); + a.add(@as(c_int, H)); + a.add(@as(f32, 1e-6)); + try cuda.launch.launch(x.ops.k.rms_norm, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, x.ops.s, &a); + } + + /// dflash2._rotary: cos and sin of each run's positions [start, start + rows) times the inverse frequencies, fp32. + fn rotary(x: *DFlash2, runs: []const [2]usize) !void { + var rows: usize = 0; + for (runs) |run| for (0..run[1]) |j| { + x.positions[rows] = @intCast(run[0] + j); + rows += 1; + }; + try x.ops.upload(x.pos, std.mem.sliceAsBytes(x.positions[0..rows])); + var a: cuda.Args = .{}; + for ([_]u64{ x.pos, x.inv_freq, x.cos, x.sin }) |v| a.add(v); + a.add(@as(c_int, D / 2)); + try cuda.launch.launch(x.ops.k.rotary, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = D / 2 } }, x.ops.s, &a); + } + + /// dflash2._prep: q (heads, rows, D), k and v (kv heads, rows, D) from [q | k | v] rows, normed and rotated. + fn prep(x: *DFlash2, qkv: u64, layer: usize, rows: usize, with_q: bool, q: u64, k: u64, v: u64) !void { + const l = x.w.layers[layer]; + const h: usize = if (with_q) dl.heads else 0; + const stride = (h + 2 * dl.kv_heads) * D; + const p = aot.ptr; + try x.t.set.run(x.t.s, "_prep_kernel", .{ @intCast(rows), @intCast(h + 2 * dl.kv_heads), 1 }, &.{ + p("QKV", "*bf16", qkv), p("QN", "*bf16", l.q_norm), p("KN", "*bf16", l.k_norm), p("COS", "*fp32", x.cos), p("SIN", "*fp32", x.sin), + p("QO", "*bf16", if (with_q) q else qkv), p("KO", "*bf16", k), p("VO", "*bf16", v), aot.int("L", @intCast(rows)), aot.int("stride", @intCast(stride)), + aot.float("eps", 1e-6), + }, &.{ aot.ci("H", @intCast(h)), aot.ci("HKV", dl.kv_heads), aot.ci("HALF", D / 2) }); + } + + /// dflash2._dconv: the two-tap dynamic conv of each `seg`-row block, plus a residual. + fn dconv(x: *DFlash2, in: u64, dyn: u64, base: u64, branch: usize, res: ?u64, out: u64, rows: usize, seg: usize) !void { + const p = aot.ptr; + try x.t.set.run(x.t.s, "_dconv_kernel", .{ @intCast(rows), (H + 1023) / 1024, 1 }, &.{ + p("X", "*bf16", in), p("DYN", "*bf16", dyn), p("BASE", "*bf16", base), p("RES", "*bf16", res orelse in), p("OUT", "*bf16", out), aot.int("SEG", @intCast(seg)), + }, &.{ aot.ci("D", H), aot.ci("G", H / dl.conv_group), aot.ci("GS", dl.conv_group), aot.ci("BRANCH", @intCast(branch)), aot.ci("HAS_RES", @intFromBool(res != null)), aot.ci("BLOCK", 1024) }); + } + + /// add_taps: the target's taps (rows of 5 x hidden bf16) projected into the context's layers. + pub fn absorb(x: *DFlash2, ctx: *Context, taps: u64, rows: usize) !void { + return x.absorbMany(&.{.{ .ctx = ctx, .rows = rows }}, taps); + } + + /// add_taps_streams: every stream's rows (in turn) through each projection once, each context its last window. + pub fn absorbMany(x: *DFlash2, items: []const Absorb, taps: u64) !void { + var total: usize = 0; + for (items) |it| total += it.rows; + if (total == 0) return; + if (total > dl.window + 1 or items.len > st.max_streams) return error.TooManyDraftTaps; + try x.lin(taps, x.w.fc, x.qkv, total); // fc's (rows, hidden) output, in a buffer of the window's rows + try x.rmsNorm(x.qkv, H, x.w.hidden_norm, x.normed, total); + var runs: [st.max_streams][2]usize = undefined; + for (items, 0..) |it, j| runs[j] = .{ it.ctx.end, it.rows }; + try x.rotary(runs[0..items.len]); + // every layer's copies: the last `keep` rows of [context | new rows], a head at a time, both halves once + var list: std.ArrayList(kern.Copy) = .empty; + defer list.deinit(x.gpa); + var starts: [dl.layers + 1]usize = undefined; + for (0..dl.layers) |l| { + starts[l] = list.items.len; + var first: usize = 0; + for (items) |it| { + const ctx = it.ctx; + const keep = @min(dl.window, ctx.len + it.rows); + const old_n = ctx.len; + const drop = old_n + it.rows - keep; // rows of [old | new] that fall out of the window + const from_old = old_n -| drop; + const new_from = drop -| old_n; + const nxt: u1 = ctx.cur ^ 1; + for ([_][2]u64{ x.ctxPair(ctx.keys[l], ctx.cur, nxt), x.ctxPair(ctx.values[l], ctx.cur, nxt) }, [_]u64{ x.k, x.v }) |pair, fresh| { + for (0..dl.kv_heads) |h| { + if (from_old > 0) try list.append(x.gpa, .{ .dst = pair[1] + h * keep * D * bf, .src = pair[0] + (h * old_n + drop) * D * bf, .bytes = from_old * D * bf }); + try list.append(x.gpa, .{ .dst = pair[1] + (h * keep + from_old) * D * bf, .src = fresh + (h * total + first + new_from) * D * bf, .bytes = (it.rows - new_from) * D * bf }); + } + } + first += it.rows; + } + } + starts[dl.layers] = list.items.len; + if (list.items.len > most_copies) return error.TooManyDraftTaps; + try x.ops.upload(x.copies, std.mem.sliceAsBytes(list.items)); + for (0..dl.layers) |l| { + try x.lin(x.normed, x.w.layers[l].kv, x.qkv, total); + try x.prep(x.qkv, l, total, false, 0, x.k, x.v); + try x.ops.copies(x.copies + starts[l] * @sizeOf(kern.Copy), starts[l + 1] - starts[l]); + } + try x.ops.s.synchronize(); + for (items) |it| { + it.ctx.cur ^= 1; + it.ctx.len = @min(dl.window, it.ctx.len + it.rows); + it.ctx.end += it.rows; + } + } + + fn ctxPair(_: *DFlash2, bufs: [2]u64, cur: u1, nxt: u1) [2]u64 { + return .{ bufs[cur], bufs[nxt] }; + } + + /// launch_blocks: each stream's pending token then masks through the layers in one forward; candidates on the host. + pub fn launchMany(x: *DFlash2, blocks: []const Block, L: usize) !void { + const S = blocks.len; + const rows = S * L; + if (S == 0 or S > st.max_streams or L < 2 or L > block) return error.NoDraftContext; + for (blocks) |b| if (b.ctx.len == 0) return error.NoDraftContext; + var ids: [most_rows]i32 = @splat(dl.mask_id); + var runs: [st.max_streams][2]usize = undefined; + // each layer's table: the streams' (keys, values) pointers, then their context lengths + const per = std.mem.alignForward(usize, S * 20, 16); // 16-byte aligned tables, as fresh torch tensors are + var table: [dl.layers * (st.max_streams * 20 + 16)]u8 align(8) = undefined; + for (blocks, 0..) |b, j| { + ids[j * L] = @intCast(b.pending); + runs[j] = .{ b.ctx.end, L }; + for (0..dl.layers) |l| { + const at = table[l * per ..]; + std.mem.writeInt(u64, at[j * 16 ..][0..8], b.ctx.keys[l][b.ctx.cur], .little); + std.mem.writeInt(u64, at[j * 16 + 8 ..][0..8], b.ctx.values[l][b.ctx.cur], .little); + std.mem.writeInt(i32, at[S * 16 + j * 4 ..][0..4], @intCast(b.ctx.len), .little); + } + } + try x.ops.upload(x.ids, std.mem.sliceAsBytes(ids[0..rows])); + try x.ops.upload(x.table, table[0 .. dl.layers * per]); + const tw = x.target; + try x.t.embed(x.ids, tw.embed.w, tw.embed.s, tw.embed.b, x.x, rows, H); + try x.rotary(runs[0..S]); + const scale: f32 = @floatCast(std.math.pow(f64, D, -0.5)); + for (x.w.layers, 0..) |l, i| { + try x.torchNorm(x.x, l.in_norm, x.normed, rows); + try x.lin(x.normed, l.attn_proj, x.dyn, rows); + try x.dconv(x.normed, x.dyn, l.attn_base, 0, null, x.conv, rows, L); + try x.lin(x.conv, l.qkv, x.qkv, rows); + try x.prep(x.qkv, i, rows, true, x.q, x.k, x.v); + const tb = x.table + i * per; + const p = aot.ptr; + try x.t.set.run(x.t.s, "_block_attention", .{ @intCast(S), dl.kv_heads, 1 }, &.{ + p("Q", "*bf16", x.q), p("KB", "*bf16", x.k), p("VB", "*bf16", x.v), p("TABLE", "*i64", tb), p("LENS", "*i32", tb + S * 16), p("O", "*bf16", x.attn), + aot.float("scale", scale), aot.int("window", dl.window), aot.int("R", @intCast(rows)), + }, &.{ aot.ci("G", dl.heads / dl.kv_heads), aot.ci("HKV", dl.kv_heads), aot.ci("L", @intCast(L)), aot.ci("LP", @intCast(@max(16, std.math.ceilPowerOfTwo(usize, L) catch unreachable))), aot.ci("D", D), aot.ci("BN", 64), aot.ci("CAUSAL", 0) }); + try x.lin(x.attn, l.o, x.o, rows); + try x.dconv(x.o, x.dyn, l.attn_base, 1, x.x, x.x, rows, L); + try x.torchNorm(x.x, l.post_norm, x.normed, rows); + try x.lin(x.normed, l.mlp_proj, x.dyn, rows); + try x.dconv(x.normed, x.dyn, l.mlp_base, 0, null, x.conv, rows, L); + try x.t.groupSums(x.conv, x.xs, H, rows, H); + try x.ops.dense(x.conv, x.xs, l.gate, x.gate, rows); + try x.ops.dense(x.conv, x.xs, l.up, x.up, rows); + try x.t.swiglu(x.gate, x.up, x.act, x.act_xs, rows, l.gate.n); + try x.ops.dense(x.act, x.act_xs, l.down, x.o, rows); + try x.dconv(x.o, x.dyn, l.mlp_base, 1, x.x, x.x, rows, L); + } + // every block but its pending row, normed into stacked depths + const depths = L - 1; + const total = S * depths; + for (0..S) |j| try x.torchNorm(x.x + (j * L + 1) * H * bf, x.w.norm, x.head_h + j * depths * H * bf, depths); + try x.lin(x.head_h, x.w.select, x.selected, total); + // matmul_rows: the head's spans at the stacked rows' K split (1), then joined column-wise + try x.t.groupSums(x.head_h, x.xs, H, total, H); + const head = tw.head; + const first: kern.QLinear = .{ .w = head.w, .s = head.s, .b = head.b, .n = dl.spans[0][1], .k = head.k, .npad = head.npad }; + try x.ops.denseSplit(x.head_h, x.xs, first, x.logits32, total, 1); // a scratch of the right size + try x.ops.denseSplit(x.head_h, x.xs, x.w.head_tail, x.logits_tail, total, 1); + var items: std.ArrayList(kern.Copy) = .empty; + defer items.deinit(x.gpa); + const na = dl.spans[0][1]; + const nb = dl.spans[1][1] - dl.spans[1][0]; + for (0..total) |r| { + try items.append(x.gpa, .{ .dst = x.logits + r * dl.head_rows * bf, .src = x.logits32 + r * na * bf, .bytes = na * bf }); + try items.append(x.gpa, .{ .dst = x.logits + (r * dl.head_rows + na) * bf, .src = x.logits_tail + r * nb * bf, .bytes = nb * bf }); + } + try x.ops.upload(x.copies, std.mem.sliceAsBytes(items.items)); + try x.ops.copies(x.copies, items.items.len); + const torch = x.ops.torch(); + try torch.toF32(x.logits, x.logits32, total * dl.head_rows); + try torch.topk(x.logits32, dl.head_rows, total, policy.top, x.cand_values, x.cand_ids, x.topk_scratch); + try x.ops.download(std.mem.sliceAsBytes(x.host_values[0 .. total * policy.top]), x.cand_values); + try x.ops.download(std.mem.sliceAsBytes(x.host_ids[0 .. total * policy.top]), x.cand_ids); + try x.ops.download(std.mem.sliceAsBytes(x.host_selected[0 .. total * dl.selector_rank]), x.selected); + try x.ops.s.synchronize(); + for (0..total * policy.top) |j| { + const local: usize = @intCast(x.host_ids[j]); + x.cands[j] = @intCast(if (local < na) local else dl.spans[1][0] + local - na); + x.unary[j] = x.host_values[j]; + } + for (0..total * dl.selector_rank) |j| x.projected[j] = @as(f32, @bitCast(@as(u32, x.host_selected[j]) << 16)); + x.launched = L; + } + + /// finish_tree for the batch's `j`th block: best_first into `out` (parents -1: under the pending row). + pub fn finish(x: *DFlash2, j: usize, pending: u32, context_length: usize, max_nodes: usize, sampling: ?lanes.Sampling, out: *policy.Tree) !void { + const depths = x.launched - 1; + const at = j * depths; + try policy.bestFirst(x.gpa, .{ .ids = x.cands[at * policy.top ..][0 .. depths * policy.top], .unary = x.unary[at * policy.top ..][0 .. depths * policy.top], .projected = x.projected[at * dl.selector_rank ..][0 .. depths * dl.selector_rank], .depths = depths }, x.w.pred, x.w.succ, pending, @min(127, max_nodes), sampling, context_length, out); + } + + /// finish_tree: best_first over the launched block's candidates; nodes and parents (-1: under the pending row). + pub fn propose(x: *DFlash2, ctx: *const Context, pending: u32, context_length: usize, max_nodes: usize, sampling: ?lanes.Sampling) !struct { tokens: []const u32, parents: []const i32 } { + if (ctx.len == 0 or max_nodes == 0) return .{ .tokens = &.{}, .parents = &.{} }; + const rows: usize = @min(block, max_nodes + 1); + try x.launchMany(&.{.{ .ctx = ctx, .pending = pending }}, rows); + try x.finish(0, pending, context_length, max_nodes, sampling, &x.tree); + return .{ .tokens = x.tree.tokens.items, .parents = x.tree.parents.items }; + } +}; diff --git a/zig/src/families/qwen27/cuda/draft_load.zig b/zig/src/families/qwen27/cuda/draft_load.zig new file mode 100644 index 000000000..8f9b0a740 --- /dev/null +++ b/zig/src/families/qwen27/cuda/draft_load.zig @@ -0,0 +1,274 @@ +//! DFlash2's weights as dflash2 holds them: projections quantize4'd and tiled, codebooks fp32 on the host. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const kern = @import("kernels.zig"); +const st = @import("state.zig"); + +pub const QLinear = kern.QLinear; +const Tensor = core.checkpoint.Tensor; + +pub const hidden = 5120; +pub const layers = 5; +pub const heads = 32; +pub const kv_heads = 8; +pub const head_dim = 128; +pub const conv_group = 16; +pub const mask_id = 248070; +pub const window = 2047; // sliding_window - 1: context rows each layer keeps +pub const taps = [_]usize{ 5, 19, 33, 47, 61 }; // dflash_config.target_layer_ids +pub const selector_rank = 256; +pub const theta: f64 = 10000000; +/// The head's rows the drafter reads: the first 98,304 tokens and the last 288 (dflash2.vocab_spans). +pub const spans = [2][2]usize{ .{ 0, 98304 }, .{ 248032, 248320 } }; +pub const head_rows = spans[0][1] - spans[0][0] + spans[1][1] - spans[1][0]; + +pub const Layer = struct { + in_norm: u64, + post_norm: u64, + attn_proj: QLinear, // attention_conv.kernel_projection: the conv's dynamic taps (2 x 2 x 320) + attn_base: u64, // attention_conv.base_kernel (2, 2, 5120) bf16 + mlp_proj: QLinear, + mlp_base: u64, + qkv: QLinear, // [q | k | v] stacked + kv: QLinear, // [k | v] stacked: the context's projection + o: QLinear, + q_norm: u64, + k_norm: u64, + gate: QLinear, + up: QLinear, + down: QLinear, +}; + +pub const Weights = struct { + gpa: std.mem.Allocator, + layers: [layers]Layer = undefined, + fc: QLinear = undefined, // five taps (25,600) to the hidden width + hidden_norm: u64 = 0, + norm: u64 = 0, + select: QLinear = undefined, // candidate_selector.hidden_projection + head_tail: QLinear = undefined, // the target head's last 288 rows, tiled on their own (rows() copies them) + pred: []f32 = &.{}, // candidate_selector.predecessor_codebook (vocab, 256) + succ: []f32 = &.{}, + buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.pred); + w.gpa.free(w.succ); + w.* = undefined; + } +}; + +fn bf16Value(word: u16) f32 { + return @bitCast(@as(u32, word) << 16); +} + +fn toBf16(x: f32) u16 { + const b: u32 = @bitCast(x); + return @intCast((b + 0x7FFF + ((b >> 16) & 1)) >> 16); +} + +/// torch.round: halves to even. +fn roundEven(x: f32) f32 { + const r = @round(x); + return if (@abs(x - @trunc(x)) == 0.5) 2.0 * @round(x / 2.0) else r; +} + +/// dflash2.quantize4 on the host, bit for bit: (n, k) bf16 to MLX affine 4-bit words, bf16 scales and biases. +pub fn quantize4(src: []const u8, n: usize, k: usize, words: []u32, scales: []u16, biases: []u16) void { + const kg = k / 64; + for (0..n) |r| for (0..kg) |gi| { + var vals: [64]f32 = undefined; + for (&vals, 0..) |*v, j| v.* = bf16Value(std.mem.readInt(u16, src[(r * k + gi * 64 + j) * 2 ..][0..2], .little)); + var lo = vals[0]; + var hi = vals[0]; + for (vals[1..]) |v| { + lo = @min(lo, v); + hi = @max(hi, v); + } + const s = toBf16(@max((hi - lo) / 15.0, 1e-8)); + const b = toBf16(lo); + scales[r * kg + gi] = s; + biases[r * kg + gi] = b; + const sf = bf16Value(s); + const bf = bf16Value(b); + for (0..8) |wi| { + var word: u32 = 0; + for (0..8) |j| { + const q = std.math.clamp(roundEven((vals[wi * 8 + j] - bf) / sf), 0.0, 15.0); + word |= @as(u32, @intFromFloat(q)) << @intCast(4 * j); + } + words[r * (k / 8) + gi * 8 + wi] = word; + } + }; +} + +const Loader = struct { + gpa: std.mem.Allocator, + ops: kern.Ops, + w: *Weights, + words: std.ArrayList(u32) = .empty, + scales: std.ArrayList(u16) = .empty, + biases: std.ArrayList(u16) = .empty, + staging: ?cuda.DeviceBuffer = null, + + fn alloc(L: *Loader, len: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(L.ops.k.d, len); + try L.w.buffers.append(L.gpa, b); + return b.ptr; + } + + fn raw(L: *Loader, t: Tensor) !u64 { + const ptr = try L.alloc(t.bytes.len); + try L.ops.upload(ptr, t.bytes); + try L.ops.s.synchronize(); + return ptr; + } + + /// qmm_fast.tile of host MLX words, scales and biases (n rows of k inputs): the lane matmul's layout. + fn tile(L: *Loader, words: []const u32, scales: []const u16, biases: []const u16, n: usize, k: usize) !QLinear { + const k8 = k / 8; + const kg = k / 64; + const npad = (n + 127) / 128 * 128; + const sizes = [3]usize{ n * k8 * 4, n * kg * 2, n * kg * 2 }; + const total_in = sizes[0] + sizes[1] + sizes[2]; + if (L.staging == null or L.staging.?.len < total_in) { + if (L.staging) |*b| b.free(); + L.staging = try cuda.DeviceBuffer.alloc(L.ops.k.d, total_in); + } + const base = L.staging.?.ptr; + try L.ops.upload(base, std.mem.sliceAsBytes(words)); + try L.ops.upload(base + sizes[0], std.mem.sliceAsBytes(scales)); + try L.ops.upload(base + sizes[0] + sizes[1], std.mem.sliceAsBytes(biases)); + const q: QLinear = .{ .w = try L.alloc(npad * k / 2), .s = try L.alloc(kg * npad * 2), .b = try L.alloc(kg * npad * 2), .n = n, .k = k, .npad = npad }; + const total: u64 = @as(u64, npad / 64) * kg * 512; + var a: cuda.Args = .{}; + a.add(base); + for ([_]usize{ n, k8, kg }) |v| a.add(@as(c_int, @intCast(v))); + a.add(q.w); + a.add(@as(c_longlong, @intCast(total))); + try cuda.launch.launch(L.ops.k.pack_dense, .{ .grid = .{ .x = @intCast((total + 255) / 256) }, .block = .{ .x = 256 } }, L.ops.s, &a); + for ([_]u64{ q.s, q.b }, [_]usize{ sizes[0], sizes[0] + sizes[1] }) |out, off| { + var t: cuda.Args = .{}; + t.add(base + off); + for ([_]usize{ n, kg, npad }) |v| t.add(@as(c_int, @intCast(v))); + t.add(out); + const cells = @as(u64, kg) * npad; + try cuda.launch.launch(L.ops.k.transpose16, .{ .grid = .{ .x = @intCast((cells + 255) / 256) }, .block = .{ .x = 256 } }, L.ops.s, &t); + } + try L.ops.s.synchronize(); // the staging buffer and host arrays are reused next + return q; + } + + /// One bf16 (n, k) checkpoint tensor (or several stacked by rows) quantized and tiled. + fn quantized(L: *Loader, parts: []const Tensor, k: usize) !QLinear { + var n: usize = 0; + for (parts) |t| { + if (t.dtype != .bf16 or t.rank != 2 or t.dim(1) != k) return error.UnexpectedDraftTensor; + n += t.dim(0); + } + try L.words.resize(L.gpa, n * k / 8); + try L.scales.resize(L.gpa, n * k / 64); + try L.biases.resize(L.gpa, n * k / 64); + var row: usize = 0; + for (parts) |t| { + const r = t.dim(0); + quantize4(t.bytes, r, k, L.words.items[row * k / 8 ..][0 .. r * k / 8], L.scales.items[row * k / 64 ..][0 .. r * k / 64], L.biases.items[row * k / 64 ..][0 .. r * k / 64]); + row += r; + } + return L.tile(L.words.items, L.scales.items, L.biases.items, n, k); + } + + fn get(ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype) !Tensor { + var buf: [160]u8 = undefined; + return ck.get(try std.fmt.bufPrint(&buf, fmt, args)); + } + + /// The codebook as fp32 on the host (dflash2 keeps tensor.float().numpy()). + fn codebook(L: *Loader, t: Tensor) ![]f32 { + if (t.dtype != .bf16 or t.rank != 2 or t.dim(1) != selector_rank) return error.UnexpectedDraftTensor; + const out = try L.gpa.alloc(f32, t.dim(0) * selector_rank); + for (out, 0..) |*o, i| o.* = bf16Value(std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)); + return out; + } +}; + +/// The drafter at `dir`, and the target head's last span tiled from the target checkpoint at `target_dir`. +pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, target_dir: []const u8) !Weights { + var w: Weights = .{ .gpa = gpa }; + errdefer w.deinit(); + var L: Loader = .{ .gpa = gpa, .ops = ops, .w = &w }; + defer { + L.words.deinit(gpa); + L.scales.deinit(gpa); + L.biases.deinit(gpa); + if (L.staging) |*b| b.free(); + } + var ck = try core.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + for (&w.layers, 0..) |*l, i| { + const q = try Loader.get(&ck, "layers.{d}.self_attn.q_proj.weight", .{i}); + const k = try Loader.get(&ck, "layers.{d}.self_attn.k_proj.weight", .{i}); + const v = try Loader.get(&ck, "layers.{d}.self_attn.v_proj.weight", .{i}); + l.qkv = try L.quantized(&.{ q, k, v }, hidden); + l.kv = try L.quantized(&.{ k, v }, hidden); + l.o = try L.quantized(&.{try Loader.get(&ck, "layers.{d}.self_attn.o_proj.weight", .{i})}, heads * head_dim); + l.q_norm = try L.raw(try Loader.get(&ck, "layers.{d}.self_attn.q_norm.weight", .{i})); + l.k_norm = try L.raw(try Loader.get(&ck, "layers.{d}.self_attn.k_norm.weight", .{i})); + l.in_norm = try L.raw(try Loader.get(&ck, "layers.{d}.input_layernorm.weight", .{i})); + l.post_norm = try L.raw(try Loader.get(&ck, "layers.{d}.post_attention_layernorm.weight", .{i})); + l.attn_proj = try L.quantized(&.{try Loader.get(&ck, "layers.{d}.attention_conv.kernel_projection.weight", .{i})}, hidden); + l.attn_base = try L.raw(try Loader.get(&ck, "layers.{d}.attention_conv.base_kernel", .{i})); + l.mlp_proj = try L.quantized(&.{try Loader.get(&ck, "layers.{d}.mlp_conv.kernel_projection.weight", .{i})}, hidden); + l.mlp_base = try L.raw(try Loader.get(&ck, "layers.{d}.mlp_conv.base_kernel", .{i})); + l.gate = try L.quantized(&.{try Loader.get(&ck, "layers.{d}.mlp.gate_proj.weight", .{i})}, hidden); + l.up = try L.quantized(&.{try Loader.get(&ck, "layers.{d}.mlp.up_proj.weight", .{i})}, hidden); + const down = try Loader.get(&ck, "layers.{d}.mlp.down_proj.weight", .{i}); + l.down = try L.quantized(&.{down}, down.dim(1)); + } + const fc = try ck.get("fc.weight"); + w.fc = try L.quantized(&.{fc}, fc.dim(1)); + w.hidden_norm = try L.raw(try ck.get("hidden_norm.weight")); + w.norm = try L.raw(try ck.get("norm.weight")); + w.select = try L.quantized(&.{try ck.get("candidate_selector.hidden_projection.weight")}, hidden); + w.pred = try L.codebook(try ck.get("candidate_selector.predecessor_codebook")); + w.succ = try L.codebook(try ck.get("candidate_selector.successor_codebook")); + if (ck.unused() != 0) return error.UnusedDraftTensors; + // dflash2's sub-head: rows() of the tiled head copies a span off the 64-row tile grid, packed on its own + var tk = try core.Checkpoint.openModelPrefix(gpa, io, target_dir, "language_model.lm_head"); + defer tk.close(); + const hw = try tk.get("language_model.lm_head.weight"); + const hs = try tk.get("language_model.lm_head.scales"); + const hb = try tk.get("language_model.lm_head.biases"); + const a = spans[1][0]; + const n = spans[1][1] - a; + const k8 = hidden / 8; + const kg = hidden / 64; + // mapped tensors may sit at odd offsets: the rows go through aligned host arrays + try L.words.resize(gpa, n * k8); + try L.scales.resize(gpa, n * kg); + try L.biases.resize(gpa, n * kg); + @memcpy(std.mem.sliceAsBytes(L.words.items), hw.bytes[a * k8 * 4 ..][0 .. n * k8 * 4]); + @memcpy(std.mem.sliceAsBytes(L.scales.items), hs.bytes[a * kg * 2 ..][0 .. n * kg * 2]); + @memcpy(std.mem.sliceAsBytes(L.biases.items), hb.bytes[a * kg * 2 ..][0 .. n * kg * 2]); + w.head_tail = try L.tile(L.words.items, L.scales.items, L.biases.items, n, hidden); + return w; +} + +test "quantize4 rounds as torch does" { + // one group: values 0..63 / 4 -> lo 0, scale bf16(63/4/15 = 1.05), each q round-half-even of v / 1.046875 + var src: [64 * 2]u8 = undefined; + for (0..64) |j| std.mem.writeInt(u16, src[2 * j ..][0..2], toBf16(@as(f32, @floatFromInt(j)) / 4.0), .little); + var words: [8]u32 = undefined; + var scales: [1]u16 = undefined; + var biases: [1]u16 = undefined; + quantize4(&src, 1, 64, &words, &scales, &biases); + try std.testing.expectEqual(toBf16(63.0 / 4.0 / 15.0), scales[0]); + try std.testing.expectEqual(@as(u16, 0), biases[0]); + try std.testing.expectEqual(@as(u32, 15), words[7] >> 28); + try std.testing.expectEqual(roundEven(2.5), 2.0); + try std.testing.expectEqual(roundEven(3.5), 4.0); +} diff --git a/zig/src/families/qwen27/cuda/draft_policy.zig b/zig/src/families/qwen27/cuda/draft_policy.zig new file mode 100644 index 000000000..83888d435 --- /dev/null +++ b/zig/src/families/qwen27/cuda/draft_policy.zig @@ -0,0 +1,122 @@ +//! draft_tree.best_first on the host: which of DFlash2's candidates a round verifies, and where each hangs. + +const std = @import("std"); +const lanes = @import("lanes"); + +pub const branch = 4; +const edge_weight = 0.6; +const noise_weight = 0.7; +const scale = 1.5; +pub const rank = 256; // the selector's codebook width +pub const top = 16; // candidates a depth + +/// A drafted block's candidates: each depth's 16 token ids, their logits (unary) and the selector's projection. +pub const Candidates = struct { ids: []const u32, unary: []const f64, projected: []const f64, depths: usize }; + +/// The tree: nodes in pop order (parents first), each node's parent (-1: hangs from the anchor), path scores. +pub const Tree = struct { tokens: std.ArrayList(u32) = .empty, parents: std.ArrayList(i32) = .empty, scores: std.ArrayList(f64) = .empty }; + +const Item = struct { score: f64, parent: i32, token: u32, depth: usize }; + +fn before(_: void, a: Item, b: Item) std.math.Order { + // heapq's tuple order: score, then parent, then token, then depth + if (a.score != b.score) return std.math.order(a.score, b.score); + if (a.parent != b.parent) return std.math.order(a.parent, b.parent); + if (a.token != b.token) return std.math.order(a.token, b.token); + return std.math.order(a.depth, b.depth); +} + +/// best_first: expand the anchor's candidates, then pop the best path score until `max_nodes` are placed. +pub fn bestFirst(gpa: std.mem.Allocator, c: Candidates, pred: []const f32, succ: []const f32, anchor: u32, max_nodes: usize, sampling: ?lanes.Sampling, first_position: u64, out: *Tree) !void { + out.tokens.clearRetainingCapacity(); + out.parents.clearRetainingCapacity(); + out.scores.clearRetainingCapacity(); + const temp: f64 = if (sampling) |s| @max(s.temperature, 1e-6) else 1.0; + var heap: std.PriorityQueue(Item, void, before) = .empty; + defer heap.deinit(gpa); + const Ctx = struct { + gpa: std.mem.Allocator, + c: Candidates, + pred: []const f32, + succ: []const f32, + temp: f64, + sampling: ?lanes.Sampling, + first_position: u64, + heap: *std.PriorityQueue(Item, void, before), + + fn expand(x: @This(), token: u32, depth: usize, parent: i32, path: f64) !void { + var values: [top]f64 = undefined; + const p = x.pred[@as(usize, token) * rank ..][0..rank]; + const proj = x.c.projected[depth * rank ..][0..rank]; + var weighted: [rank]f64 = undefined; + for (&weighted, p, proj) |*wgt, a, b| wgt.* = @as(f64, a) * b; + for (&values, 0..) |*v, i| { + const cand = x.c.ids[depth * top + i]; + const s = x.succ[@as(usize, cand) * rank ..][0..rank]; + var edge: f64 = 0; + for (s, weighted) |a, b| edge += @as(f64, a) * b; + v.* = (x.c.unary[depth * top + i] + edge_weight * edge) / x.temp; + if (x.sampling) |smp| v.* += noise_weight * -@log(-@log(lanes.sampling.uniform(smp.seed, x.first_position + depth, cand))); + v.* /= scale; + } + var most = values[0]; + for (values[1..]) |v| most = @max(most, v); + var total: f64 = 0; + for (&values) |*v| { + v.* -= most; + total += @exp(v.*); + } + const log_total = @log(total); + var order: [top]usize = undefined; + for (&order, 0..) |*o, i| o.* = i; + std.mem.sort(usize, &order, &values, struct { + fn more(vals: *const [top]f64, a: usize, b: usize) bool { + return vals[a] > vals[b] or (vals[a] == vals[b] and a < b); + } + }.more); + for (order[0..branch]) |i| { + const logp = values[i] - log_total; + if (!std.math.isFinite(logp)) break; + try x.heap.push(x.gpa, .{ .score = path - logp, .parent = parent, .token = x.c.ids[depth * top + i], .depth = depth }); + } + } + }; + const ctx: Ctx = .{ .gpa = gpa, .c = c, .pred = pred, .succ = succ, .temp = temp, .sampling = sampling, .first_position = first_position, .heap = &heap }; + try ctx.expand(anchor, 0, -1, 0.0); + while (heap.count() > 0 and out.tokens.items.len < max_nodes) { + const it = heap.pop().?; + const me: i32 = @intCast(out.tokens.items.len); + try out.tokens.append(gpa, it.token); + try out.parents.append(gpa, it.parent); + try out.scores.append(gpa, it.score); + if (it.depth + 1 < c.depths) try ctx.expand(it.token, it.depth + 1, me, it.score); + } +} + +test "best_first pops by path score, parents before children, ties by parent" { + const gpa = std.testing.allocator; + const vocab = 32; + const pred = try gpa.alloc(f32, vocab * rank); + defer gpa.free(pred); + const succ = try gpa.alloc(f32, vocab * rank); + defer gpa.free(succ); + @memset(pred, 0); + @memset(succ, 0); + var ids: [2 * top]u32 = undefined; + var unary: [2 * top]f64 = undefined; + for (0..2) |d| for (0..top) |i| { + ids[d * top + i] = @intCast(i); + unary[d * top + i] = -@as(f64, @floatFromInt(i)); + }; + const projected: [2 * rank]f64 = @splat(0); + var tree: Tree = .{}; + defer { + tree.tokens.deinit(gpa); + tree.parents.deinit(gpa); + tree.scores.deinit(gpa); + } + try bestFirst(gpa, .{ .ids = &ids, .unary = &unary, .projected = &projected, .depths = 2 }, pred, succ, 3, 5, null, 100, &tree); + // logp_i = -i / 1.5 - ln(sum): the anchor's 0 and 1 (score 0.72, 1.39) come before 0's child 0 (1.44) + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 0, 2, 1 }, tree.tokens.items); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, 0, -1, 0 }, tree.parents.items); +} diff --git a/zig/src/families/qwen27/cuda/drafts.zig b/zig/src/families/qwen27/cuda/drafts.zig new file mode 100644 index 000000000..02e365952 --- /dev/null +++ b/zig/src/families/qwen27/cuda/drafts.zig @@ -0,0 +1,103 @@ +//! DFlash2 for a shared round's streams (multi.py's drafter halves): copies first, else every block in one batch. + +const std = @import("std"); +const lanes = @import("lanes"); +const st = @import("state.zig"); +const kern = @import("kernels.zig"); +const dr = @import("draft.zig"); +const dl = @import("draft_load.zig"); +const lone = @import("lone.zig"); + +const bf = 2; +pub const taps_width = dl.taps.len * dl.hidden; +/// A shared window's drafts at most (MultiDecoder's max_rows - 1). +pub const max_nodes = st.max_rows - 1; + +pub const Kind = enum { none, copy, tree }; + +/// One stream's drafting: its drafter context, its copy index, and the drafts it holds for its next window. +pub const Stream = struct { + gpa: std.mem.Allocator, + ctx: dr.Context, + copies: lone.CopyIndex, + kind: Kind = .none, + tokens: std.ArrayList(u32) = .empty, + parents: std.ArrayList(i32) = .empty, // each draft's parent draft (-1: under the pending row) + scores: std.ArrayList(f64) = .empty, // each draft's path score (best_first's; 0 for a copy) + + pub fn init(gpa: std.mem.Allocator, ctx: dr.Context) Stream { + return .{ .gpa = gpa, .ctx = ctx, .copies = .{ .gpa = gpa } }; + } + + pub fn deinit(x: *Stream) void { + x.ctx.deinit(); + x.copies.deinit(); + x.tokens.deinit(x.gpa); + x.parents.deinit(x.gpa); + x.scores.deinit(x.gpa); + } + + fn clear(x: *Stream) void { + x.kind = .none; + x.tokens.clearRetainingCapacity(); + x.parents.clearRetainingCapacity(); + x.scores.clearRetainingCapacity(); + } +}; + +/// A stream asking for its next drafts: its context through the pending token, and its sampling. +pub const Ask = struct { x: *Stream, context: []const u32, sampling: ?lanes.Sampling }; + +/// _mode, launch_blocks and finish_tree: a copied continuation where the context repeats, else a DFlash2 tree. +pub fn propose(d: *dr.DFlash2, asks: []const Ask) !void { + var blocks: [st.max_streams]dr.Block = undefined; + var which: [st.max_streams]usize = undefined; + var n: usize = 0; + for (asks, 0..) |a, k| { + const x = a.x; + x.clear(); + const copied = try x.copies.propose(a.context, max_nodes); + if (copied.len > 0) { + x.kind = .copy; + try x.tokens.appendSlice(x.gpa, copied); + for (0..copied.len) |r| { + try x.parents.append(x.gpa, @as(i32, @intCast(r)) - 1); + try x.scores.append(x.gpa, 0); + } + } else if (x.ctx.len > 0) { + blocks[n] = .{ .ctx = &x.ctx, .pending = a.context[a.context.len - 1] }; + which[n] = k; + n += 1; + } + } + if (n == 0) return; + try d.launchMany(blocks[0..n], @min(dr.block, max_nodes + 1)); + for (blocks[0..n], which[0..n], 0..) |b, k, j| { + const a = asks[k]; + try d.finish(j, b.pending, a.context.len, max_nodes, a.sampling, &d.tree); + a.x.kind = .tree; + try a.x.tokens.appendSlice(a.x.gpa, d.tree.tokens.items); + try a.x.parents.appendSlice(a.x.gpa, d.tree.parents.items); + try a.x.scores.appendSlice(a.x.gpa, d.tree.scores.items); + } +} + +/// A stream's kept rows of the last round (rows of the round's taps, in path order). +pub const Kept = struct { x: *Stream, rows: []const u32 }; + +/// add_taps_streams: every stream's kept rows' taps gathered after the round's rows, then absorbed in one batch. +pub fn absorb(d: *dr.DFlash2, ops: kern.Ops, taps: u64, copy_items: u64, kept: []const Kept) !void { + var items: std.ArrayList(kern.Copy) = .empty; + defer items.deinit(d.gpa); + var absorbs: [st.max_streams]dr.Absorb = undefined; + const row = taps_width * bf; + const gathered = taps + st.round_rows * row; + for (kept, 0..) |k, j| { + for (k.rows) |r| try items.append(d.gpa, .{ .dst = gathered + items.items.len * row, .src = taps + r * row, .bytes = row }); + absorbs[j] = .{ .ctx = &k.x.ctx, .rows = k.rows.len }; + } + if (items.items.len == 0) return; + try ops.upload(copy_items, std.mem.sliceAsBytes(items.items)); + try ops.copies(copy_items, items.items.len); + try d.absorbMany(absorbs[0..kept.len], gathered); +} diff --git a/zig/src/families/qwen27/cuda/engine.zig b/zig/src/families/qwen27/cuda/engine.zig new file mode 100644 index 000000000..52576d7cd --- /dev/null +++ b/zig/src/families/qwen27/cuda/engine.zig @@ -0,0 +1,118 @@ +//! The 27B's CUDA engine for one stream: the weights, kernels and scratch, prefill and greedy serial decode. + +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("shape.zig"); +const kern = @import("kernels.zig"); +const tri = @import("triton.zig"); +const st = @import("state.zig"); +const wts = @import("weights.zig"); +const Forward = @import("forward.zig").Forward; +const Part = @import("forward.zig").Part; + +/// prefill.CHUNK: a prompt chunk's rows (prompt_rows on a GPU of 80 GB or more). +pub const prompt_rows = 4096; + +/// prefill.chunks: even chunks of at most `size` rows of [start, end); any bounds give the same bits. +pub fn chunkBounds(start: usize, end: usize, size: usize, j: usize) [2]usize { + const n = (end - start + size - 1) / size; + return .{ start + (end - start) * j / n, start + (end - start) * (j + 1) / n }; +} + +pub const Options = struct { context: usize }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + ctx: *const cuda.Context, + s: cuda.Stream, + k: kern.Kernels, + w: wts.Weights, + scratch: st.Scratch, + context: usize, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: []const u8, o: Options) !*Engine { + const cfg = try c.Config.read(gpa, io, dir); + const g = cfg.g; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.gpa = gpa; + e.ctx = ctx; + e.context = o.context; + e.s = try cuda.Stream.init(ctx.d, true); + errdefer e.s.deinit(); + e.k = try kern.Kernels.load(gpa, io, ctx, kernels, g.linear_v_heads); + errdefer e.k.deinit(); + e.w = try wts.load(gpa, io, e.ops(), dir, cfg); + errdefer e.w.deinit(); + e.scratch = try st.Scratch.init(gpa, ctx.d, g, prompt_rows, o.context); + errdefer e.scratch.deinit(); + try e.s.synchronize(); + return e; + } + + pub fn deinit(e: *Engine) void { + e.scratch.deinit(); + e.w.deinit(); + e.k.deinit(); + e.s.deinit(); + e.gpa.destroy(e); + } + + pub fn ops(e: *Engine) kern.Ops { + return .{ .k = &e.k, .s = e.s }; + } + + pub fn forward(e: *Engine) Forward { + return .{ .gpa = e.gpa, .w = &e.w, .ops = e.ops(), .t = .{ .set = &e.k.triton, .s = e.s }, .s = &e.scratch }; + } + + /// A stream's state for `capacity` positions (its caches, recurrences and conv windows), empty. + pub fn sequence(e: *Engine, capacity: usize) !st.Seq { + return st.Seq.init(e.gpa, e.ops(), e.w.g, capacity); + } + + /// prefill_state from an empty state: the prompt in even chunks of at most `prompt_rows`, its last row normed. + pub fn prefill(e: *Engine, seq: *st.Seq, prompt: []const u32) !void { + if (prompt.len == 0) return error.EmptyPrompt; + if (prompt.len >= seq.capacity) return error.PromptTooLong; + try seq.reset(e.ops(), e.w.g); + var f = e.forward(); + const n = (prompt.len + prompt_rows - 1) / prompt_rows; + for (0..n) |j| { + const b = chunkBounds(0, prompt.len, prompt_rows, j); + try f.chunk(prompt[b[0]..b[1]], seq, j == n - 1); + } + } + + /// decode.serial_decode, greedy: one-row windows from `pending` until `count` tokens or an end token. + pub fn decode(e: *Engine, seq: *st.Seq, pending: u32, count: usize, out: *std.ArrayList(u32)) !void { + var f = e.forward(); + try out.append(e.gpa, pending); + var pick: [1]u32 = undefined; + while (out.items.len < count and !e.w.isEos(out.items[out.items.len - 1])) { + const last = out.items[out.items.len - 1]; + const part = [_]Part{.{ .seq = seq, .tokens = &.{last} }}; + try f.round(&part); + try f.picks(1, &pick); + try f.commit(&part, &.{&.{0}}); + try out.append(e.gpa, pick[0]); + } + } + + /// The prompt, then up to `max_tokens` greedy tokens (the first included), as generate(draft=False) gives them. + pub fn generate(e: *Engine, seq: *st.Seq, prompt: []const u32, max_tokens: usize, out: *std.ArrayList(u32)) !void { + const room = @min(max_tokens, seq.capacity - prompt.len); + try e.prefill(seq, prompt); + var f = e.forward(); + try f.head(seq); + var pick: [1]u32 = undefined; + try f.picks(1, &pick); + try e.decode(seq, pick[0], @max(1, room), out); + } +}; + +test "prefill.chunks splits evenly" { + try std.testing.expectEqual([2]usize{ 0, 2539 }, chunkBounds(0, 5078, 4096, 0)); + try std.testing.expectEqual([2]usize{ 2539, 5078 }, chunkBounds(0, 5078, 4096, 1)); + try std.testing.expectEqual([2]usize{ 0, 40 }, chunkBounds(0, 40, 4096, 0)); +} diff --git a/zig/src/families/qwen27/cuda/forward.zig b/zig/src/families/qwen27/cuda/forward.zig new file mode 100644 index 000000000..4e4964be4 --- /dev/null +++ b/zig/src/families/qwen27/cuda/forward.zig @@ -0,0 +1,436 @@ +//! The 27B's forward on CUDA: verify rounds (tree_forward, multi_tree_forward), their commit, and prompt chunks. + +const std = @import("std"); +const c = @import("shape.zig"); +const kern = @import("kernels.zig"); +const tri = @import("triton.zig"); +const st = @import("state.zig"); +const wts = @import("weights.zig"); +const rnd = @import("round.zig"); + +const bf = 2; +pub const Part = rnd.Part; +/// The target layers DFlash2 reads (dflash_config.target_layer_ids). +pub const tap_layers = [_]usize{ 5, 19, 33, 47, 61 }; + +pub const Forward = struct { + gpa: std.mem.Allocator, + w: *const wts.Weights, + ops: kern.Ops, + t: tri.Tri, + s: *st.Scratch, + dump: ?Dump = null, // a debugging copy of every norm's (h, y) rows, as capture.py's hook writes them + taps: ?u64 = null, // rows of the drafter's five layer taps (5 x hidden bf16), written by round and chunk + + pub const Dump = struct { io: std.Io, dir: []const u8, n: usize = 0, rows: usize = 0 }; + + fn g(f: *const Forward) c.Geometry { + return f.w.g; + } + + fn gdnShape(f: *const Forward) tri.Gdn { + const g_ = f.g(); + return .{ .c = g_.convDim(), .kh = g_.linear_k_heads, .vh = g_.linear_v_heads, .dk = c.linear_dim, .keep = c.conv_taps - 1 }; + } + + fn heads(f: *const Forward) kern.Heads { + return .{ .hk = f.g().linear_k_heads, .hv = f.g().linear_v_heads, .dv = c.linear_dim }; + } + + fn attnShape(f: *const Forward) tri.Attn { + return .{ .heads = f.g().query_heads, .kv_heads = f.g().kv_heads, .dim = c.head_dim, .half = c.rotary_dim / 2 }; + } + + fn upload32(f: *Forward, dst: u64, values: []const i32) !void { + try f.ops.upload(dst, std.mem.sliceAsBytes(values)); + try f.ops.s.synchronize(); // the host copy must outlive the async upload + } + + /// add_rmsnorm into the scratch's other residual buffer; returns the residual stream's new home. + fn norm(f: *Forward, x: u64, r: ?u64, weight: u64, rows: usize) !u64 { + const out = if (r == null) x else if (x == f.s.x) f.s.h else f.s.x; + try f.t.addRmsnorm(x, r, weight, out, f.s.y, f.s.xs, rows, f.g().hidden, c.eps); + if (f.dump) |*d| try f.saveRows(d, out, f.s.y, rows); + return out; + } + + /// Layer `i`'s tap, when the drafter reads it: (x + pending) of the call's rows into the taps buffer. + fn tapLayer(f: *Forward, i: usize, x: u64, rows: usize) !void { + const out = f.taps orelse return; + const j = std.mem.indexOfScalar(usize, &tap_layers, i) orelse return; + const h = f.g().hidden; + try f.ops.tap(x, f.s.pending, out, rows, h, tap_layers.len * h, j * h); + } + + fn saveRows(f: *Forward, d: *Dump, h: u64, y: u64, rows: usize) !void { + const bytes = try f.gpa.alloc(u8, rows * f.g().hidden * bf); + defer f.gpa.free(bytes); + for ([_]u64{ h, y }, [_][]const u8{ "h", "y" }) |src, name| { + try f.ops.download(bytes, src); + try f.ops.s.synchronize(); + var buf: [256]u8 = undefined; + const path = try std.fmt.bufPrint(&buf, "{s}/{d:0>3}_{s}.bin", .{ d.dir, d.n, name }); + try std.Io.Dir.cwd().writeFile(d.io, .{ .sub_path = path, .data = bytes }); + } + d.n += 1; + } + + /// The window of one stream at its position (tree_forward): logits (rows, vocab) bf16 in the scratch, records kept. + pub fn window(f: *Forward, tokens: []const u32, seq: *st.Seq) !void { + return f.round(&.{.{ .seq = seq, .tokens = tokens }}); + } + + /// Several streams' chain windows in one forward (multi_tree_forward; one stream: tree_forward), rows in order. + pub fn round(f: *Forward, parts: []const Part) !void { + const s = f.s; + const g_ = f.g(); + const w = f.w; + const r = try rnd.stage(f.gpa, f.ops, f.t, s, g_, parts); + const W = r.rows; + try f.t.embed(r.ids, w.embed.w, w.embed.s, w.embed.b, s.x, W, g_.hidden); + var x = s.x; + var pending: ?u64 = null; + var lin: usize = 0; + var att: usize = 0; + for (w.layers, 0..) |layer, i| { + x = try f.norm(x, pending, layer.input_norm, W); + switch (layer.mixer) { + .delta => |d| { + const rec = s.records[i]; + try f.ops.dense(s.y, s.xs, d.qkv, rec.qkv, W); + try f.ops.dense(s.y, s.xs, d.z, s.z, W); + try f.ops.dense(s.y, s.xs, d.b, s.b, W); + try f.ops.dense(s.y, s.xs, d.a, s.a, W); + const sid: ?u64 = if (r.multi) r.sids else null; + try f.t.gdnPre(rec.qkv, r.conv(parts, s, g_, i, lin), d.conv, r.windows, sid, s.a, s.b, d.a_log, d.dt_bias, rec.q, rec.k, rec.v, rec.g, rec.beta, W, f.gdnShape()); + const state = if (r.multi) 0 else parts[0].seq.rec[i]; + const table = if (r.multi) r.table(lin) else 0; + try f.ops.gdnTree(rec.q, rec.k, rec.v, rec.g, rec.beta, state, table, if (r.multi) r.starts else 0, r.plan, s.yr, W, r.streams, r.slots, r.most, f.heads()); + try f.t.gatedNorm(s.yr, s.z, d.norm, s.gated, s.gated_xs, W, g_.linear_v_heads, c.linear_dim, c.eps); + try f.ops.dense(s.gated, s.gated_xs, d.out, s.r, W); + lin += 1; + }, + .attention => |a| { + const rec = s.attn_records[i]; + try f.ops.dense(s.y, s.xs, a.q, s.qg, W); + try f.ops.dense(s.y, s.xs, a.k, s.key, W); + try f.ops.dense(s.y, s.xs, a.v, rec.v, W); + try f.t.attnPrep(s.qg, s.key, a.q_norm, a.k_norm, r.pos, w.inv_freq, s.q_rot, rec.k, W, f.attnShape(), c.eps); + try f.t.attention(s.q_rot, rec.k, rec.v, s.origin, r.offsets(i, att), r.attn, .{ .o = s.part_o, .m = s.part_m, .l = s.part_l }, s.attn, f.attnShape()); + try f.t.gateMul(s.attn, s.qg, s.gated, s.gated_xs, W, g_.query_heads, c.head_dim); + try f.ops.dense(s.gated, s.gated_xs, a.o, s.r, W); + att += 1; + }, + } + x = try f.norm(x, s.r, layer.post_norm, W); + try f.ops.dense(s.y, s.xs, layer.gate, s.gate, W); + try f.ops.dense(s.y, s.xs, layer.up, s.up, W); + try f.t.swiglu(s.gate, s.up, s.act, s.act_xs, W, g_.intermediate); + try f.ops.dense(s.act, s.act_xs, layer.down, s.pending, W); + pending = s.pending; + try f.tapLayer(i, x, W); + } + const spare = if (x == s.x) s.h else s.x; + try f.t.addRmsnorm(x, pending, w.norm, spare, s.last_h, s.last_xs, W, g_.hidden, c.eps); + if (f.dump) |*d| try f.saveRows(d, spare, s.last_h, W); + try f.ops.dense(s.last_h, s.last_xs, w.head, s.logits, W); + } + + /// commit_streams after a round over `parts`: each stream keeps the window rows of its accepted path. + pub fn commit(f: *Forward, parts: []const Part, paths: []const []const u32) !void { + const s = f.s; + const g_ = f.g(); + const S = parts.len; + var width: usize = 0; + for (parts, paths) |p, path| { + if (path.len < 1 or path.len > p.tokens.len) return error.InvalidQwenCommit; + for (path) |row| if (row >= p.tokens.len) return error.InvalidQwenCommit; + width += p.tokens.len; + } + // gdn.replay_table: each DeltaNet layer's k, v, g, beta, then each stream's state of each layer + var table: std.ArrayList(u64) = .empty; + defer table.deinit(f.gpa); + var layers: usize = 0; + for (0..g_.layers) |i| if (c.linear(i)) { + const rc = s.records[i]; + try table.appendSlice(f.gpa, &.{ rc.k, rc.v, rc.g, rc.beta }); + layers += 1; + }; + for (parts) |p| for (0..g_.layers) |i| if (c.linear(i)) try table.append(f.gpa, p.seq.rec[i]); + try f.ops.upload(s.replay_table, std.mem.sliceAsBytes(table.items)); + // path_indices: each stream's accepted rows (round rows) padded to the round's width, then its count + var rows: std.ArrayList(i32) = .empty; + defer rows.deinit(f.gpa); + var base: usize = 0; + for (parts, paths) |p, path| { + for (0..width) |j| try rows.append(f.gpa, if (j < path.len) @intCast(base + path[j]) else 0); + try rows.append(f.gpa, @intCast(path.len)); + base += p.tokens.len; + } + try f.ops.upload(s.replay_rows, std.mem.sliceAsBytes(rows.items)); + try f.ops.s.synchronize(); + try f.ops.gdnReplay(s.replay_table, layers, S, s.replay_rows, width + 1, s.replay_rows + width * 4, width + 1, f.heads()); + const row = g_.convDim() * bf; + const keep = c.conv_taps - 1; + const kv = g_.kvInner() * bf; + // conv rows: the last three of [committed | accepted rows]; old ones saved first, then written with the rest + var save: std.ArrayList(kern.Copy) = .empty; + defer save.deinit(f.gpa); + var write: std.ArrayList(kern.Copy) = .empty; + defer write.deinit(f.gpa); + base = 0; + for (parts, paths, 0..) |p, path, k| { + const n = path.len; + var lin: usize = 0; + for (0..g_.layers) |i| { + if (c.linear(i)) { + const old = s.conv_old + (k * g_.layers + lin) * keep * row; + const stay = keep -| n; // committed rows that stay: old[n..3) go to new[0..3-n) + if (stay > 0) { + try save.append(f.gpa, .{ .dst = old, .src = p.seq.conv[i] + n * row, .bytes = stay * row }); + try write.append(f.gpa, .{ .dst = p.seq.conv[i], .src = old, .bytes = stay * row }); + } + for (path[n - (keep - stay) ..], stay..) |acc, j| try write.append(f.gpa, .{ .dst = p.seq.conv[i] + j * row, .src = s.records[i].qkv + (base + acc) * row, .bytes = row }); + lin += 1; + } else { + for (path, 0..) |acc, j| { + const at = (p.seq.pos + j) * kv; + try write.append(f.gpa, .{ .dst = p.seq.keys[i] + at, .src = s.attn_records[i].k + (base + acc) * kv, .bytes = kv }); + try write.append(f.gpa, .{ .dst = p.seq.values[i] + at, .src = s.attn_records[i].v + (base + acc) * kv, .bytes = kv }); + } + } + } + base += p.tokens.len; + p.seq.pos += n; + } + const second = s.copy_items + st.copy_items * @sizeOf(kern.Copy); + try f.ops.upload(s.copy_items, std.mem.sliceAsBytes(save.items)); + try f.ops.upload(second, std.mem.sliceAsBytes(write.items)); + try f.ops.copies(s.copy_items, save.items.len); + try f.ops.copies(second, write.items.len); + try f.ops.s.synchronize(); // the tables' host copies go when this returns + } + + /// prefill_chunk: `ids` committed at the stream's position; `last` leaves the last row normed, with group_sums. + pub fn chunk(f: *Forward, ids: []const u32, seq: *st.Seq, last: bool) !void { + const W = ids.len; + const s = f.s; + const g_ = f.g(); + const w = f.w; + if (W < 1 or W > s.rows) return error.ChunkTooLong; + if (seq.pos + W > seq.capacity) return error.PromptTooLong; + const p0 = seq.pos; + { + const host = try f.gpa.alloc(i32, W * (2 + c.conv_taps)); + defer f.gpa.free(host); + for (ids, 0..) |tok, r| { + host[r] = @intCast(tok); + host[W + r] = @intCast(p0 + r); + for (0..c.conv_taps) |j| host[2 * W + r * c.conv_taps + j] = @intCast(r + j); + } + try f.upload32(s.ids, host[0..W]); + try f.upload32(s.pos, host[W .. 2 * W]); + try f.upload32(s.windows, host[2 * W ..]); + } + try f.t.embed(s.ids, w.embed.w, w.embed.s, w.embed.b, s.x, W, g_.hidden); + var x = s.x; + var pending: ?u64 = null; + const kv = g_.kvInner() * bf; + const row = g_.convDim() * bf; + const keep = c.conv_taps - 1; + for (w.layers, 0..) |layer, i| { + x = try f.norm(x, pending, layer.input_norm, W); + switch (layer.mixer) { + .delta => |d| { + try f.ops.prefillDense(s.y, d.qkv, s.qkv, W); + try f.ops.prefillDense(s.y, d.z, s.z, W); + try f.ops.prefillDense(s.y, d.b, s.b, W); + try f.ops.prefillDense(s.y, d.a, s.a, W); + try f.t.gdnPre(s.qkv, seq.conv[i], d.conv, s.windows, null, s.a, s.b, d.a_log, d.dt_bias, s.gq, s.gk, s.gv, s.gg, s.gbeta, W, f.gdnShape()); + // each thread reads its state rows before it writes them: the chain may end in place + try f.ops.gdnPrompt(s.gq, s.gk, s.gv, s.gg, s.gbeta, seq.rec[i], seq.rec[i], s.yr, W, f.heads()); + if (W >= keep) { + try f.ops.copy(seq.conv[i], s.qkv + (W - keep) * row, keep * row); + } else { + try f.ops.copy(s.conv_tmp, seq.conv[i], keep * row); + try f.ops.copy(s.conv_tmp + keep * row, s.qkv, W * row); + try f.ops.copy(seq.conv[i], s.conv_tmp + W * row, keep * row); + } + try f.t.gatedNorm(s.yr, s.z, d.norm, s.gated, s.gated_xs, W, g_.linear_v_heads, c.linear_dim, c.eps); + try f.ops.prefillDense(s.gated, d.out, s.r, W); + }, + .attention => |a| { + try f.ops.prefillDense(s.y, a.q, s.qg, W); + try f.ops.prefillDense(s.y, a.k, s.key, W); + try f.ops.prefillDense(s.y, a.v, seq.values[i] + p0 * kv, W); + try f.t.attnPrep(s.qg, s.key, a.q_norm, a.k_norm, s.pos, w.inv_freq, s.q_rot, seq.keys[i] + p0 * kv, W, f.attnShape(), c.eps); + const scale: f32 = @floatCast(std.math.pow(f64, c.head_dim, -0.5)); + try f.ops.prefillAttention(s.q_rot, seq.keys[i], seq.values[i], s.attn, p0, W, g_.query_heads, g_.kv_heads, scale); + try f.t.gateMul(s.attn, s.qg, s.gated, s.gated_xs, W, g_.query_heads, c.head_dim); + try f.ops.prefillDense(s.gated, a.o, s.r, W); + }, + } + x = try f.norm(x, s.r, layer.post_norm, W); + try f.ops.prefillDense(s.y, layer.gate, s.gate, W); + try f.ops.prefillDense(s.y, layer.up, s.up, W); + try f.t.swiglu(s.gate, s.up, s.act, s.act_xs, W, g_.intermediate); + try f.ops.prefillDense(s.act, layer.down, s.pending, W); + pending = s.pending; + try f.tapLayer(i, x, W); + } + seq.pos = p0 + W; + if (!last) return; + const at = (W - 1) * g_.hidden * bf; + const spare = if (x == s.x) s.h else s.x; + try f.t.addRmsnorm(x + at, pending.? + at, w.norm, spare, seq.head_in, s.last_xs, 1, g_.hidden, c.eps); + try f.t.groupSums(seq.head_in, seq.head_xs, g_.hidden, 1, g_.hidden); + } + + /// One stream's share of a batched prompt pass: its ids, committed at its position; `last` normalizes its last row. + pub const Piece = struct { ids: []const u32, seq: *st.Seq, last: bool }; + + /// prefill_rows: several streams' prompt chunks in one pass (each stream's conv, DeltaNet and attention its own). + pub fn chunks(f: *Forward, pieces: []const Piece) !void { + const s = f.s; + const g_ = f.g(); + const w = f.w; + const S = pieces.len; + if (S == 0 or S > st.max_streams) return error.InvalidQwenWindows; + var starts: [st.max_streams + 1]usize = undefined; + starts[0] = 0; + for (pieces, 0..) |p, k| { + if (p.ids.len == 0) return error.ChunkTooLong; + if (p.seq.pos + p.ids.len > p.seq.capacity) return error.PromptTooLong; + starts[k + 1] = starts[k] + p.ids.len; + } + const W = starts[S]; + if (W > s.rows) return error.ChunkTooLong; + const keep = c.conv_taps - 1; + const kv = g_.kvInner() * bf; + const row = g_.convDim() * bf; + { + // ids, positions, each row's stream and its conv window: [the streams' kept rows; every prompt row] + const host = try f.gpa.alloc(i32, W * (3 + c.conv_taps)); + defer f.gpa.free(host); + for (pieces, 0..) |p, k| for (p.ids, 0..) |tok, r| { + const at = starts[k] + r; + host[at] = @intCast(tok); + host[W + at] = @intCast(p.seq.pos + r); + host[2 * W + at] = @intCast(k); + for (0..c.conv_taps) |j| host[3 * W + at * c.conv_taps + j] = @intCast(if (r + j < keep) r + j else r + j + starts[k]); + }; + try f.upload32(s.ids, host[0..W]); + try f.upload32(s.pos, host[W .. 2 * W]); + try f.upload32(s.sids, host[2 * W .. 3 * W]); + try f.upload32(s.windows, host[3 * W ..]); + } + // copies the pass needs: every stream's conv rows stacked by DeltaNet layer, then each layer's writes + var items: std.ArrayList(kern.Copy) = .empty; + defer items.deinit(f.gpa); + var lin_n: usize = 0; + for (0..g_.layers) |i| if (c.linear(i)) { + const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + for (pieces, 0..) |p, k| try items.append(f.gpa, .{ .dst = cat + k * keep * row, .src = p.seq.conv[i], .bytes = keep * row }); + lin_n += 1; + }; + const staged = items.items.len; + var at_layer: [256]usize = undefined; // each layer's first write, then the end + lin_n = 0; + for (0..g_.layers) |i| { + at_layer[i] = items.items.len; + if (c.linear(i)) { + // the last `keep` rows of [kept rows; the stream's prompt rows] + const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + for (pieces, 0..) |p, k| { + const n = p.ids.len; + const o = starts[k]; + if (n >= keep) { + try items.append(f.gpa, .{ .dst = p.seq.conv[i], .src = s.qkv + (o + n - keep) * row, .bytes = keep * row }); + } else { + try items.append(f.gpa, .{ .dst = p.seq.conv[i], .src = cat + (k * keep + n) * row, .bytes = (keep - n) * row }); + try items.append(f.gpa, .{ .dst = p.seq.conv[i] + (keep - n) * row, .src = s.qkv + o * row, .bytes = n * row }); + } + } + lin_n += 1; + } else { + for (pieces, 0..) |p, k| { + const o = starts[k]; + const n = p.ids.len; + try items.append(f.gpa, .{ .dst = p.seq.keys[i] + p.seq.pos * kv, .src = s.prompt_keys + o * kv, .bytes = n * kv }); + try items.append(f.gpa, .{ .dst = p.seq.values[i] + p.seq.pos * kv, .src = s.prompt_values + o * kv, .bytes = n * kv }); + } + } + } + at_layer[g_.layers] = items.items.len; + if (items.items.len > 2 * st.copy_items) return error.InvalidQwenWindows; + try f.ops.upload(s.copy_items, std.mem.sliceAsBytes(items.items)); + try f.ops.copies(s.copy_items, staged); + const item = @sizeOf(kern.Copy); + try f.t.embed(s.ids, w.embed.w, w.embed.s, w.embed.b, s.x, W, g_.hidden); + var x = s.x; + var pending: ?u64 = null; + lin_n = 0; + for (w.layers, 0..) |layer, i| { + x = try f.norm(x, pending, layer.input_norm, W); + switch (layer.mixer) { + .delta => |d| { + try f.ops.prefillDense(s.y, d.qkv, s.qkv, W); + try f.ops.prefillDense(s.y, d.z, s.z, W); + try f.ops.prefillDense(s.y, d.b, s.b, W); + try f.ops.prefillDense(s.y, d.a, s.a, W); + const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + try f.t.gdnPre(s.qkv, cat, d.conv, s.windows, s.sids, s.a, s.b, d.a_log, d.dt_bias, s.gq, s.gk, s.gv, s.gg, s.gbeta, W, f.gdnShape()); + for (pieces, 0..) |p, k| { + const o = starts[k]; + // each thread reads its state rows before it writes them: the chain may end in place + try f.ops.gdnPrompt(s.gq + o * g_.kInner() * bf, s.gk + o * g_.kInner() * bf, s.gv + o * g_.vInner() * bf, s.gg + o * g_.linear_v_heads * 4, s.gbeta + o * g_.linear_v_heads * 4, p.seq.rec[i], p.seq.rec[i], s.yr + o * g_.vInner() * bf, p.ids.len, f.heads()); + } + try f.ops.copies(s.copy_items + at_layer[i] * item, at_layer[i + 1] - at_layer[i]); + try f.t.gatedNorm(s.yr, s.z, d.norm, s.gated, s.gated_xs, W, g_.linear_v_heads, c.linear_dim, c.eps); + try f.ops.prefillDense(s.gated, d.out, s.r, W); + lin_n += 1; + }, + .attention => |a| { + try f.ops.prefillDense(s.y, a.q, s.qg, W); + try f.ops.prefillDense(s.y, a.k, s.key, W); + try f.ops.prefillDense(s.y, a.v, s.prompt_values, W); + try f.t.attnPrep(s.qg, s.key, a.q_norm, a.k_norm, s.pos, w.inv_freq, s.q_rot, s.prompt_keys, W, f.attnShape(), c.eps); + try f.ops.copies(s.copy_items + at_layer[i] * item, at_layer[i + 1] - at_layer[i]); + const scale: f32 = @floatCast(std.math.pow(f64, c.head_dim, -0.5)); + const qrow = g_.query_heads * c.head_dim * bf; + for (pieces, 0..) |p, k| try f.ops.prefillAttention(s.q_rot + starts[k] * qrow, p.seq.keys[i], p.seq.values[i], s.attn + starts[k] * qrow, p.seq.pos, p.ids.len, g_.query_heads, g_.kv_heads, scale); + try f.t.gateMul(s.attn, s.qg, s.gated, s.gated_xs, W, g_.query_heads, c.head_dim); + try f.ops.prefillDense(s.gated, a.o, s.r, W); + }, + } + x = try f.norm(x, s.r, layer.post_norm, W); + try f.ops.prefillDense(s.y, layer.gate, s.gate, W); + try f.ops.prefillDense(s.y, layer.up, s.up, W); + try f.t.swiglu(s.gate, s.up, s.act, s.act_xs, W, g_.intermediate); + try f.ops.prefillDense(s.act, layer.down, s.pending, W); + pending = s.pending; + try f.tapLayer(i, x, W); + } + const spare = if (x == s.x) s.h else s.x; + for (pieces, 0..) |p, k| { + p.seq.pos += p.ids.len; + if (!p.last) continue; + const at = (starts[k + 1] - 1) * g_.hidden * bf; + try f.t.addRmsnorm(x + at, pending.? + at, w.norm, spare, p.seq.head_in, s.last_xs, 1, g_.hidden, c.eps); + try f.t.groupSums(p.seq.head_in, p.seq.head_xs, g_.hidden, 1, g_.hidden); + } + try f.ops.s.synchronize(); // the copy table's host list goes when this returns + } + + /// The head over the stream's last prompt row: logits (1, vocab) in the scratch. + pub fn head(f: *Forward, seq: *const st.Seq) !void { + try f.ops.dense(seq.head_in, seq.head_xs, f.w.head, f.s.logits, 1); + } + + /// Greedy picks of the scratch's first `rows` logits rows (torch.argmax, the first maximum). + pub fn picks(f: *Forward, rows: usize, out: []u32) !void { + try f.ops.argmax(f.s.logits, c.vocab, c.vocab, f.s.picks, rows); + try f.ops.download(std.mem.sliceAsBytes(out[0..rows]), f.s.picks); + try f.ops.s.synchronize(); + } +}; diff --git a/zig/src/families/qwen27/cuda/kernels.zig b/zig/src/families/qwen27/cuda/kernels.zig new file mode 100644 index 000000000..4b486922f --- /dev/null +++ b/zig/src/families/qwen27/cuda/kernels.zig @@ -0,0 +1,365 @@ +//! The 27B's CUDA kernels: our .cu fatbins with the Python wrappers' launch logic, and the captured Triton set. + +const std = @import("std"); +const cuda = @import("cuda"); +pub const torch_ops = @import("nemotron").torch_ops; + +/// Mangled names of the instantiations the copies in zig/kernels/cuda export (cuobjdump -symbols of each fatbin). +const sym = struct { + const group = "_ZN12tf_qmm_group12group_kernelILi64ELi16ELi64ELi1ELi4ELi8ELb0ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfNS_5PartsEiiiii"; + const wide = [_][:0]const u8{ + "_ZN12tf_qmm_group12group_kernelILi64ELi32ELi64ELi1ELi4ELi4ELb0ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfNS_5PartsEiiiii", + "_ZN12tf_qmm_group12group_kernelILi64ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfNS_5PartsEiiiii", + "_ZN12tf_qmm_group12group_kernelILi64ELi64ELi128ELi2ELi4ELi3ELb0ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfNS_5PartsEiiiii", + }; + const gemv = "_ZN12tf_lane_gemv11gemv_kernelILi64ELi64ELi8EEEvPK13__nv_bfloat16PKfNS_4PartEiii"; + const prefill_mm = "_ZN14tf_qmm_prefill14prefill_kernelILi64ELi128ELi128ELi2ELi2ELi2ELb0EEEvPK13__nv_bfloat16PKjS3_S3_Pviiiiii"; + const pattn = "_ZN20tf_prefill_attention12pattn_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiif"; + const chain128 = "_ZN14tf_gdn_prefill12chain_kernelI13__nv_bfloat16Li128ELi32EEEvPKT_S4_PKS1_PKfS8_S8_PfPS1_iii"; + const chain64 = "_ZN14tf_gdn_prefill12chain_kernelI13__nv_bfloat16Li64ELi32EEEvPKT_S4_PKS1_PKfS8_S8_PfPS1_iii"; + const argmax = "tf_argmax_rows_i32_kernel"; +}; + +/// qmm_group.cu's Part and Parts, passed by value: four projections at most, one used here. +pub const Part = extern struct { w: u64, scales: u64, biases: u64, out: u64, n: c_int, npad: c_int, sk: c_int, tiles: c_int, first: c_int }; +pub const Parts = extern struct { p: [4]Part, count: c_int }; + +/// lane_gemv.cu's Part: one projection whose column tiles the CTAs share out. +pub const GemvPart = extern struct { w: u64, scales: u64, biases: u64, out: u64, n: c_int, npad: c_int, sk: c_int, tiles: c_int }; + +/// gdn.cu's Pending (the tree kernel's rows still to fold): all null here, commits replay instead. +pub const Pending = extern struct { k: u64 = 0, v: u64 = 0, g: u64 = 0, beta: u64 = 0, rows: u64 = 0, row_stride: c_int = 0, counts: u64 = 0, count_stride: c_int = 0 }; + +comptime { + std.debug.assert(@sizeOf(Part) == 56 and @sizeOf(Parts) == 232 and @sizeOf(GemvPart) == 48 and @sizeOf(Pending) == 64); +} + +pub const group_smem: u32 = 35328; // LaneTile<64, 16, 64, 1, 4, 8>::SMEM +/// qmm_group's GB10 tiles 3-5 for wider rounds: rows a tile, columns a tile, threads, LaneTile::SMEM. +pub const Wide = struct { bm: usize, bn: usize, threads: u32, smem: u32 }; +pub const wide = [_]Wide{ .{ .bm = 32, .bn = 64, .threads = 128, .smem = 26112 }, .{ .bm = 64, .bn = 64, .threads = 128, .smem = 43008 }, .{ .bm = 64, .bn = 128, .threads = 256, .smem = 39168 } }; +pub const prefill_mm_smem: u32 = 41984; // Tile<64, 128, 128, 2, 2, 2>::SMEM: prompt_tile 9, two stages +pub const pattn_smem: u32 = 98304; // two 32-key slots of 256 dims beside eight warps' staged queries + +/// One item of a batched copy: 16-byte aligned pointers, a multiple of 16 bytes. +pub const Copy = extern struct { dst: u64, src: u64, bytes: u64 }; + +/// A tiled 4-bit projection: packed words, (kg, npad) scales and biases, n outputs from k inputs. +pub const QLinear = struct { w: u64, s: u64, b: u64, n: usize, k: usize, npad: usize }; + +pub const Kernels = struct { + d: *const cuda.Driver, + mods: [15]cuda.Module, + triton: cuda.aot.Set, + group: cuda.Function, + wide: [wide.len]cuda.Function, // tiles 3-5: rows <= 32, <= 64, more + gemv: cuda.Function, + prefill_mm: cuda.Function, + pattn: cuda.Function, + trees: [cuda.kernels.gdn_variants.len]cuda.Function, // tree_kernel by gdn_variants: [0] the chain one + chip: [3]c_int, // major, minor, SMs: what dispatch_tree reads + replay: cuda.Function, + prompt_chain: cuda.Function, + prompt_rows: usize, // the prompt chain's value rows a block: 64 when heads < SMs + pack_dense: cuda.Function, + transpose16: cuda.Function, + argmax: cuda.Function, + torch: torch_ops.Functions, + copy_batch: cuda.Function, + tap: cuda.Function, + rms_norm: cuda.Function, + inv_freq: cuda.Function, + torch_rms: cuda.Function, // F.rms_norm's bits on bf16 rows of 5120 (torch_ops/norm_rotary.cu) + rotary: cuda.Function, + gemv_blocks: usize, // resident lane_gemv CTAs: per SM times SMs + gemv_rows: usize, // windows of at most this many rows take lane_gemv (none; TF_QWEN27_GEMV_ROWS) + gb10: bool, + discrete: bool, // the card has its own memory: checkpoint bytes reach it through page-locked slots + + /// Loads every module; `triton_dir` holds the captured aot.json and cubins for this GPU. + pub fn load(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, triton_dir: []const u8, heads: usize) !Kernels { + const d = ctx.d; + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + k.d = d; + const kk = cuda.kernels; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.prefill_attention, kk.gdn, kk.gdn_prefill, kk.affine4_pack, kk.torch_argmax, kk.torch_topk, kk.torch_pointwise, kk.torch_indexing, kk.torch_movement, kk.torch_nemotron_constants, kk.qwen_ops, kk.torch_norm_rotary }; + var loaded: usize = 0; + errdefer for (k.mods[0..loaded]) |*m| m.unload(); + for (images, 0..) |img, i| { + k.mods[i] = try cuda.Module.load(d, img); + loaded += 1; + } + k.group = try k.mods[0].function(sym.group); + for (&k.wide, sym.wide) |*f, name| f.* = try k.mods[0].function(name); + k.prefill_mm = try k.mods[1].function(sym.prefill_mm); + k.gemv = try k.mods[2].function(sym.gemv); + k.pattn = try k.mods[3].function(sym.pattn); + for (&k.trees, cuda.kernels.gdn_variants) |*f, v| f.* = try k.mods[4].function(v.symbol); + k.replay = try k.mods[4].function(cuda.kernels.gdn_symbols.replay_bf16); + k.pack_dense = try k.mods[6].function(cuda.qlinear.symbols.pack_words); + k.transpose16 = try k.mods[6].function(cuda.qlinear.symbols.pack_scales); + k.argmax = try k.mods[7].function(sym.argmax); + k.torch = try torch_ops.Functions.resolve(k.mods[7..13]); + k.copy_batch = try k.mods[13].function("tf_copy_batch"); + k.tap = try k.mods[13].function("tf_tap"); + k.rms_norm = try k.mods[13].function("tf_rms_norm"); + k.inv_freq = try k.mods[13].function("tf_inv_freq"); + k.rotary = try k.mods[13].function("tf_rotary"); + k.torch_rms = try k.mods[14].function("tf_dflash_rms5120_bf16_fma_kernel"); + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + // gdn_prefill_cuda: a block takes 128 value rows, or 64 when there are fewer heads than SMs + k.prompt_rows = if (heads >= sms) 128 else 64; + k.prompt_chain = try k.mods[5].function(if (k.prompt_rows == 128) sym.chain128 else sym.chain64); + k.triton = try cuda.aot.Set.load(gpa, io, d, ctx.device, triton_dir); + errdefer k.triton.deinit(); + try k.group.allowDynamicShared(group_smem); + for (k.wide, wide) |f, t| try f.allowDynamicShared(t.smem); + try k.gemv.allowDynamicShared(group_smem); + try k.prefill_mm.allowDynamicShared(prefill_mm_smem); + try k.pattn.allowDynamicShared(pattn_smem); + k.gemv_blocks = @max(1, try k.gemv.occupancy(128, group_smem)) * sms; + const major = try ctx.attribute(.compute_capability_major); + const minor = try ctx.attribute(.compute_capability_minor); + k.gb10 = major == 12 and minor == 1; + k.discrete = try ctx.attribute(.integrated) == 0; + k.gemv_rows = 0; // a GB10 runs the 27B's projections faster in Python's clusters, 1 to 16 rows alike + if (std.c.getenv("TF_QWEN27_GEMV_ROWS")) |v| k.gemv_rows = std.fmt.parseInt(usize, std.mem.span(v), 10) catch k.gemv_rows; + k.chip = .{ major, minor, @intCast(sms) }; + return k; + } + + pub fn deinit(k: *Kernels) void { + k.triton.deinit(); + for (&k.mods) |*m| m.unload(); + } +}; + +/// qmm.split_k: K slices fixed by the weight's shape, never by the row count. +pub fn splitK(n: usize, k: usize) usize { + const tiles = (n + 63) / 64; + const groups = k / 64; + var sk: usize = 1; + while (sk < 8 and tiles * sk < 192 and groups % (sk * 2) == 0 and groups / (sk * 2) >= 8) sk *= 2; + return sk; +} + +fn int(x: usize) c_int { + return @intCast(x); +} + +fn u(x: usize) u32 { + return @intCast(x); +} + +/// The GDN shapes every launch takes: key heads, value heads, value rows a head. +pub const Heads = struct { hk: usize, hv: usize, dv: usize }; + +/// Launch helpers on one stream; each mirrors the Python wrapper it replaces. +pub const Ops = struct { + k: *const Kernels, + s: cuda.Stream, + + fn go(o: Ops, f: cuda.Function, grid: [3]usize, block: u32, shared: u32, args: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = .{ .x = u(grid[0]), .y = u(grid[1]), .z = u(grid[2]) }, .block = .{ .x = block }, .shared = shared }, o.s, args); + } + + /// qmm.matmul on sm_12x: x (rows, k) bf16 with group sums xs -> out (rows, n) bf16, qmm_group's bits. + pub fn dense(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize) !void { + const sk = splitK(q.n, q.k); + if (rows > 16) return o.wideTile(x, xs, q, out, rows, sk); + // Python's clusters, a CTA a K slice; lane_gemv (byte-equal) when TF_QWEN27_GEMV_ROWS asks for it + return if (sk > 1 and rows <= o.k.gemv_rows) o.gemv(x, xs, q, out, rows, sk) else o.cluster(x, xs, q, out, rows, sk); + } + + /// qmm.matmul at a K split the caller names (matmul_rows takes the stacked weight's for each part). + pub fn denseSplit(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + if (rows > 16) return o.wideTile(x, xs, q, out, rows, sk); + return if (sk > 1 and rows <= o.k.gemv_rows) o.gemv(x, xs, q, out, rows, sk) else o.cluster(x, xs, q, out, rows, sk); + } + + /// tree_forward's taps: bf16(x + pending) of `rows` rows into columns [col, col + d) of rows `stride` wide. + pub fn tap(o: Ops, x: u64, pending: u64, out: u64, rows: usize, d: usize, stride: usize, col: usize) !void { + var a: cuda.Args = .{}; + for ([_]u64{ x, pending, out }) |v| a.add(v); + for ([_]usize{ d, stride, col }) |v| a.add(int(v)); + try o.go(o.k.tap, .{ rows, 4, 1 }, 256, 0, &a); + } + + /// lane_gemv: every K slice of a column tile in one CTA, summed in slice order, CTAs looping over the tiles. + fn gemv(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + const tiles = (q.n + 63) / 64; + var a: cuda.Args = .{}; + a.add(x); + a.add(xs); + a.add(GemvPart{ .w = q.w, .scales = q.s, .biases = q.b, .out = out, .n = int(q.n), .npad = int(q.npad), .sk = int(sk), .tiles = int(tiles) }); + for ([_]usize{ rows, q.k, q.k }) |v| a.add(int(v)); + const cfg: cuda.Config = .{ .grid = .{ .x = u(@min(tiles, o.k.gemv_blocks)) }, .block = .{ .x = 128 }, .shared = group_smem, .pdl = o.k.gb10 }; + try cuda.launch.launch(o.k.gemv, cfg, o.s, &a); + } + + /// qmm_group's tile 2 with one K slice: a CTA a column tile. + fn cluster(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + const tiles = (q.n + 63) / 64; + var parts: Parts = std.mem.zeroes(Parts); + parts.count = 1; + parts.p[0] = .{ .w = q.w, .scales = q.s, .biases = q.b, .out = out, .n = int(q.n), .npad = int(q.npad), .sk = int(sk), .tiles = int(tiles), .first = 0 }; + const rows_t = (rows + 15) / 16; + var a: cuda.Args = .{}; + a.add(x); + a.add(xs); + a.add(parts); + for ([_]usize{ rows, q.k, q.k, rows_t, sk }) |v| a.add(int(v)); + const cfg: cuda.Config = .{ + .grid = .{ .x = u(rows_t * tiles * sk) }, + .block = .{ .x = 128 }, + .shared = group_smem, + .cluster = if (sk > 1) .{ .x = u(sk) } else null, + .pdl = o.k.gb10, + }; + try cuda.launch.launch(o.k.group, cfg, o.s, &a); + } + + /// qmm_group's dispatch on a GB10 past 16 rows: tile 3, 4 or 5, a cluster of `sk` CTAs a column tile. + fn wideTile(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + const i: usize = if (rows <= 32) 0 else if (rows <= 64) 1 else 2; + const t = wide[i]; + const tiles = (q.n + t.bn - 1) / t.bn; + const rows_t = (rows + t.bm - 1) / t.bm; + var parts: Parts = std.mem.zeroes(Parts); + parts.count = 1; + parts.p[0] = .{ .w = q.w, .scales = q.s, .biases = q.b, .out = out, .n = int(q.n), .npad = int(q.npad), .sk = int(sk), .tiles = int(tiles), .first = 0 }; + var a: cuda.Args = .{}; + a.add(x); + a.add(xs); + a.add(parts); + for ([_]usize{ rows, q.k, q.k, rows_t, sk }) |v| a.add(int(v)); + const cfg: cuda.Config = .{ + .grid = .{ .x = u(rows_t * tiles * sk) }, + .block = .{ .x = t.threads }, + .shared = t.smem, + .cluster = if (sk > 1) .{ .x = u(sk) } else null, + .pdl = o.k.gb10, + }; + try cuda.launch.launch(o.k.wide[i], cfg, o.s, &a); + } + + /// qmm.prefill_matmul at prompt_tile 9: weights rounded once to bf16, one fp32 chain over K; bf16 out. + pub fn prefillDense(o: Ops, x: u64, q: QLinear, out: u64, rows: usize) !void { + const rows_t = (rows + 127) / 128; + const band = (12 << 20) / (128 * q.k * 2); + const group = @max(1, @min(rows_t, band)); + var a: cuda.Args = .{}; + for ([_]u64{ x, q.w, q.s, q.b, out }) |v| a.add(v); + for ([_]usize{ rows, q.n, q.k, q.npad, q.k, group }) |v| a.add(int(v)); + try o.go(o.k.prefill_mm, .{ rows_t * ((q.n + 127) / 128), 1, 1 }, 128, prefill_mm_smem, &a); + } + + /// gdn.tree (dispatch_tree): y from one `state`, or each stream's own by `table` and `starts`; `slots` 0: chains. + pub fn gdnTree(o: Ops, q: u64, k: u64, v: u64, g: u64, beta: u64, state: u64, table: u64, starts: u64, plan: u64, y: u64, rows: usize, streams: usize, slots: usize, most: usize, h: Heads) !void { + const variant = cuda.kernels.treeVariant(@intCast(slots), @intCast(streams), o.k.chip[0], o.k.chip[1], o.k.chip[2]); + const index = for (cuda.kernels.gdn_variants, 0..) |x, n| { + if (std.mem.eql(u8, x.symbol, variant.symbol)) break n; + } else unreachable; + var a: cuda.Args = .{}; + for ([_]u64{ q, k, v, g, beta, state, table, starts, plan }) |x| a.add(x); + a.add(int(rows)); + a.add(y); + for ([_]usize{ h.hk, h.hv, h.dv }) |x| a.add(int(x)); + a.add(Pending{}); + a.add(@as(u64, 0)); + a.add(@as(u64, 0)); + a.add(h.dv % variant.r == 0 and v % (2 * @as(u64, @min(variant.r, 8))) == 0); + const shared = 16 * @as(usize, variant.warps) * variant.slots * variant.r * 32 + (if (variant.chain) 0 else 4 * 3 * most); + try o.go(o.k.trees[index], .{ (h.dv + variant.r * variant.warps - 1) / (variant.r * variant.warps), h.hv, streams }, 32 * variant.warps, @intCast(shared), &a); + } + + /// gdn.replay in place: each stream's accepted rows stepped into the states the table names. + pub fn gdnReplay(o: Ops, table: u64, layers: usize, streams: usize, rows: u64, row_stride: usize, counts: u64, count_stride: usize, h: Heads) !void { + var a: cuda.Args = .{}; + a.add(table); + a.add(int(layers)); + a.add(rows); + a.add(int(row_stride)); + a.add(counts); + a.add(int(count_stride)); + a.add(@as(u64, 0)); + for ([_]usize{ h.hk, h.hv, h.dv }) |x| a.add(int(x)); + try o.go(o.k.replay, .{ (h.dv + 31) / 32, h.hv, layers * streams }, 128, 0, &a); + } + + /// gdn.chain: a prompt chunk's rows from `state` (read only) to `last`, chunk-invariant bits. + pub fn gdnPrompt(o: Ops, q: u64, k: u64, v: u64, g: u64, beta: u64, state: u64, last: u64, y: u64, rows: usize, h: Heads) !void { + var a: cuda.Args = .{}; + for ([_]u64{ q, k, v, g, beta, state, last, y }) |x| a.add(x); + for ([_]usize{ rows, h.hk, h.hv }) |x| a.add(int(x)); + const r = o.k.prompt_rows; + try o.go(o.k.prompt_chain, .{ h.hv, 128 / r, 1 }, u(2 * r), 0, &a); + } + + /// prefill_attention (head dim 256, two query heads a block): q (rows, heads, 256) at positions p0.. + pub fn prefillAttention(o: Ops, q: u64, kc: u64, vc: u64, out: u64, p0: usize, rows: usize, heads: usize, kv_heads: usize, scale: f32) !void { + const g = heads / kv_heads; + const hpc = 2; + var a: cuda.Args = .{}; + for ([_]u64{ q, kc, vc, out }) |v| a.add(v); + for ([_]usize{ p0, rows, heads, kv_heads, g }) |v| a.add(int(v)); + a.add(scale); + const block_rows = 16 * (8 / hpc); + try o.go(o.k.pattn, .{ (rows + block_rows - 1) / block_rows, kv_heads * (g / hpc), 1 }, 256, pattn_smem, &a); + } + + /// torch.argmax(rows of bf16 logits) as int32: the first maximum, a NaN first. + pub fn argmax(o: Ops, logits: u64, vocab: usize, ld: usize, out: u64, rows: usize) !void { + var a: cuda.Args = .{}; + a.add(logits); + a.add(out); + for ([_]usize{ rows, vocab, ld }) |v| a.add(@as(i64, @intCast(v))); + a.add(@as(i32, 0)); // dtype 0: bf16 + try o.go(o.k.argmax, .{ rows, 1, 1 }, 1024, 0, &a); + } + + /// The torch-op replacements on this stream (topk and casts for sampled draws). + pub fn torch(o: Ops) torch_ops.Torch { + return .{ .f = &o.k.torch, .s = o.s }; + } + + /// qwen_ops.cu's tf_copy_batch: `count` (dst, src, bytes) items from the device table `items`, one launch. + pub fn copies(o: Ops, items: u64, count: usize) !void { + if (count == 0) return; + var a: cuda.Args = .{}; + a.add(items); + try o.go(o.k.copy_batch, .{ count, 4, 1 }, 256, 0, &a); + } + + pub fn copy(o: Ops, dst: u64, src: u64, bytes: usize) !void { + if (bytes == 0) return; + try o.k.d.check(o.k.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, o.s.handle), "cuMemcpyDtoDAsync"); + } + + pub fn fill32(o: Ops, dst: u64, value: u32, words: usize) !void { + try o.k.d.check(o.k.d.api.cuMemsetD32Async(dst, value, words, o.s.handle), "cuMemsetD32Async"); + } + + pub fn upload(o: Ops, dst: u64, bytes: []const u8) !void { + if (bytes.len == 0) return; + try o.k.d.check(o.k.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, o.s.handle), "cuMemcpyHtoDAsync"); + } + + pub fn download(o: Ops, dst: []u8, src: u64) !void { + if (dst.len == 0) return; + try o.k.d.check(o.k.d.api.cuMemcpyDtoHAsync_v2(dst.ptr, src, dst.len, o.s.handle), "cuMemcpyDtoHAsync"); + } +}; + +test "split_k follows the 27B's shapes" { + try std.testing.expectEqual(@as(usize, 2), splitK(10240, 5120)); // qkv + try std.testing.expectEqual(@as(usize, 2), splitK(6144, 5120)); // z + try std.testing.expectEqual(@as(usize, 8), splitK(48, 5120)); // b, a + try std.testing.expectEqual(@as(usize, 1), splitK(17408, 5120)); // gate, up + try std.testing.expectEqual(@as(usize, 4), splitK(5120, 17408)); // down + try std.testing.expectEqual(@as(usize, 4), splitK(5120, 6144)); // out_proj, o_proj + try std.testing.expectEqual(@as(usize, 1), splitK(12288, 5120)); // q with its gate + try std.testing.expectEqual(@as(usize, 8), splitK(1024, 5120)); // k, v + try std.testing.expectEqual(@as(usize, 1), splitK(248320, 5120)); // head +} diff --git a/zig/src/families/qwen27/cuda/lanes.zig b/zig/src/families/qwen27/cuda/lanes.zig new file mode 100644 index 000000000..8a5f86fad --- /dev/null +++ b/zig/src/families/qwen27/cuda/lanes.zig @@ -0,0 +1,549 @@ +//! The 27B's CUDA engine behind the lane core: every stream's window in one verify, DFlash2 drafts in one batch. + +const std = @import("std"); +const lanes = @import("lanes"); +const c = @import("shape.zig"); +const st = @import("state.zig"); +const Engine = @import("engine.zig").Engine; +const Part = @import("forward.zig").Part; +const Forward = @import("forward.zig").Forward; +const prompt_rows = @import("engine.zig").prompt_rows; +const dl = @import("draft_load.zig"); +const dr = @import("draft.zig"); +const drafts = @import("drafts.zig"); +const lone = @import("lone.zig"); +const trees = @import("tree.zig"); +const be = lanes.backend; + +/// The row counts calibrate times (calibration_rows without the tile steps a GB10 skips), to 16 a stream. +const calibration = [_]usize{ 1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256 }; + +/// The last verified round, committed when the next call needs the streams' states (keep may cut it first). +const Pending = struct { + streams: [st.max_streams]*lanes.Stream = undefined, + parts: [st.max_streams]Part = undefined, + paths: [st.max_streams][]const u32 = undefined, + path_rows: [st.round_rows]u32 = undefined, // the round's kept rows, each stream's path in its own slice + tokens: [st.round_rows]u32 = undefined, + parents: [st.round_rows]i32 = undefined, + starts: [st.max_streams]usize = undefined, // each window's first row in the round + n: usize = 0, +}; + +pub const Cuda = struct { + gpa: std.mem.Allocator, + e: *Engine, + streams: u32, + seqs: std.AutoHashMapUnmanaged(*lanes.Stream, *st.Seq) = .empty, + logits: []u16 = &.{}, // a whole row on the host, for a draw with top_k off + cand_values: []f32 = &.{}, // the rows' top candidates on the host + cand_ids: []i64 = &.{}, + values: []f64 = &.{}, + pending: Pending = .{}, + drafter: ?*dr.DFlash2 = null, // DFlash2, when its checkpoint is pulled and drafts are on + drafting: std.AutoHashMapUnmanaged(*lanes.Stream, *drafts.Stream) = .empty, + costs: [calibration.len + 1]lanes.config.Cost = undefined, // the forward's ms by a round's rows (calibrate) + cost_count: usize = 0, + + pub fn init(gpa: std.mem.Allocator, e: *Engine, streams: u32) !Cuda { + return .{ + .gpa = gpa, + .e = e, + .streams = streams, + .logits = try gpa.alloc(u16, c.vocab), + .cand_values = try gpa.alloc(f32, st.round_rows * st.max_candidates), + .cand_ids = try gpa.alloc(i64, st.round_rows * st.max_candidates), + .values = try gpa.alloc(f64, st.max_candidates), + }; + } + + pub fn deinit(x: *Cuda) void { + var it = x.seqs.valueIterator(); + while (it.next()) |seq| { + seq.*.deinit(); + x.gpa.destroy(seq.*); + } + x.seqs.deinit(x.gpa); + var dit = x.drafting.valueIterator(); + while (dit.next()) |ds| { + ds.*.deinit(); + x.gpa.destroy(ds.*); + } + x.drafting.deinit(x.gpa); + x.gpa.free(x.logits); + x.gpa.free(x.cand_values); + x.gpa.free(x.cand_ids); + x.gpa.free(x.values); + } + + /// The facts the round loop reads: 16-row windows in one forward; with DFlash2, host trees with node chances. + pub fn facts(x: *const Cuda) lanes.Model { + var m: lanes.Model = .{ .exact_width = st.max_rows, .hidden_rows = true, .batch_rows = st.round_rows, .max_streams = @min(x.streams, st.max_streams), .streams_exact = true }; + if (x.drafter != null) { + m.mtp = true; + m.speculate = true; + m.speculate_early = false; + m.drafts = drafts.max_nodes; + m.draft_probabilities = true; + m.draft_streams = true; + } + const timed = x.costs[0..x.cost_count]; + var narrow: usize = 0; + while (narrow < timed.len and timed[narrow].width <= st.max_rows) narrow += 1; + m.window_costs = timed[0..narrow]; + m.shared_costs = timed; + return m; + } + + /// multi.calibrate: the forward's ms (median of three) by the rows `streams` windows bring, on a fresh state. + pub fn calibrate(x: *Cuda, io: std.Io, streams: usize) !void { + var seq = x.e.sequence(st.max_rows + 2) catch |err| return if (err == error.OutOfMemory) err else error.OutOfDeviceMemory; + defer seq.deinit(); + const zeros: [st.max_rows]u32 = @splat(0); + const most: usize = st.max_rows * @as(usize, @min(streams, st.max_streams)); + var rows: [calibration.len + 1]usize = undefined; + var n: usize = 0; + for (calibration) |r| if (r < most) { + rows[n] = r; + n += 1; + }; + rows[n] = most; + n += 1; + var f = x.e.forward(); + const clock = std.Io.Clock.awake; + for (rows[0..n], x.costs[0..n]) |r, *cost| { + const parts_n = (r + st.max_rows - 1) / st.max_rows; + var parts: [st.max_streams]Part = undefined; + for (parts[0..parts_n], 0..) |*part, i| part.* = .{ .seq = &seq, .tokens = zeros[0 .. r / parts_n + @intFromBool(i < r % parts_n)] }; + var times: [3]f64 = undefined; + for (0..4) |rep| { + try f.ops.s.synchronize(); + const t0 = clock.now(io).toNanoseconds(); + try f.round(parts[0..parts_n]); + try f.ops.s.synchronize(); + if (rep > 0) times[rep - 1] = @as(f64, @floatFromInt(clock.now(io).toNanoseconds() - t0)) / 1e6; + } + std.mem.sort(f64, ×, {}, std.sort.asc(f64)); + cost.* = .{ .width = @intCast(r), .ms = times[1] }; + } + x.cost_count = n; + } + + pub fn backend(x: *Cuda) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .probabilities = probabilities, .tree = tree, .release = release, .prefill_many = prefillMany } }; + } + + /// The stream's drafting (its DFlash2 context and copy index), made on first use. + pub fn draftingOf(x: *Cuda, s: *lanes.Stream) !*drafts.Stream { + if (x.drafting.get(s)) |ds| return ds; + const ds = try x.gpa.create(drafts.Stream); + errdefer x.gpa.destroy(ds); + const ctx = dr.Context.init(x.e.ctx.d) catch |err| return if (err == error.OutOfMemory) err else error.OutOfDeviceMemory; + ds.* = drafts.Stream.init(x.gpa, ctx); + errdefer ds.deinit(); + try x.drafting.put(x.gpa, s, ds); + return ds; + } + + fn cast(p: *anyopaque) *Cuda { + return @ptrCast(@alignCast(p)); + } + + fn seqOf(x: *Cuda, s: *lanes.Stream) !*st.Seq { + return x.seqs.get(s) orelse error.UnknownQwenStream; + } + + /// Commits the last verified round: every stream its kept rows (its whole window unless keep cut it). + fn settle(x: *Cuda) !void { + const p = &x.pending; + if (p.n == 0) return; + const n = p.n; + p.n = 0; + var f = x.e.forward(); + try f.commit(p.parts[0..n], p.paths[0..n]); + } + + /// The stream's state, the last round committed first. + fn settled(x: *Cuda, s: *lanes.Stream) !*st.Seq { + try x.settle(); + return x.seqOf(s); + } + + /// The round's logits rows drawn: greedy by one argmax, sampled from top (top_k + 8) GPU candidates or whole rows. + pub fn drawRows(x: *Cuda, rules: []const ?lanes.Sampling, positions: []const u64, out: []u32) !void { + const rows = positions.len; + var f = x.e.forward(); + const s = f.s; + try f.picks(rows, out[0..rows]); + var want: usize = 0; + for (rules) |r| if (r) |rule| if (rule.temperature > 0 and rule.top_k != 0) { + want = @max(want, @min(c.vocab, @as(usize, rule.top_k) + 8)); + }; + if (want > st.max_candidates) want = 0; // past the GPU's candidates: whole rows below + if (want > 0) { + const t = f.ops.torch(); + try t.toF32(s.logits, s.logits32, rows * c.vocab); + try t.topk(s.logits32, c.vocab, rows, want, s.cand_values, s.cand_ids, s.topk_scratch); + try f.ops.download(std.mem.sliceAsBytes(x.cand_values[0 .. rows * want]), s.cand_values); + try f.ops.download(std.mem.sliceAsBytes(x.cand_ids[0 .. rows * want]), s.cand_ids); + try f.ops.s.synchronize(); + } + for (rules, positions, out[0..rows], 0..) |rule_, pos, *o, r| { + const rule = rule_ orelse continue; + if (rule.temperature <= 0) continue; + if (rule.top_k != 0 and want > 0) { + const vals = x.values[0..want]; + for (vals, x.cand_values[r * want ..][0..want]) |*v, w| v.* = w; + const ids = x.cand_ids[r * want ..][0..want]; + o.* = @intCast(try lanes.sampling.choose(x.gpa, vals, @ptrCast(ids), pos, rule)); + continue; + } + try f.ops.download(std.mem.sliceAsBytes(x.logits), s.logits + r * c.vocab * 2); + try f.ops.s.synchronize(); + o.* = try draw(x.gpa, x.logits, rule, pos); + } + } + + /// A new stream's state, sized for its prompt and reply. + pub fn open(x: *Cuda, s: *lanes.Stream) !*st.Seq { + try x.settle(); + const ids = s.prompt(); + if (ids.len == 0 or x.seqs.contains(s)) return error.InvalidQwenPrompt; + const capacity = @min(x.e.context, ids.len + s.max_new + st.max_rows); + if (ids.len >= capacity) return error.PromptTooLong; + const seq = try x.gpa.create(st.Seq); + errdefer x.gpa.destroy(seq); + seq.* = x.e.sequence(capacity) catch |err| return if (err == error.OutOfMemory) err else error.OutOfDeviceMemory; + errdefer seq.deinit(); + try x.seqs.put(x.gpa, s, seq); + return seq; + } + + fn prefill(p: *anyopaque, s: *lanes.Stream) !void { + const x = cast(p); + const seq = try x.open(s); + errdefer x.dropSeq(s); + try x.prefillOpened(s, seq); + } + + /// prefill_state into an opened stream: the prompt in chunks, a drafting stream's last window of taps absorbed. + fn prefillOpened(x: *Cuda, s: *lanes.Stream, seq: *st.Seq) !void { + if (x.drafter) |d| if (s.drafts) { + const ds = try x.draftingOf(s); + return lone.prefill(x.e, seq, &ds.ctx, d, s.prompt()); + }; + try x.e.prefill(seq, s.prompt()); + } + + /// prefill_batch: prompts in shared passes of at most a chunk's rows (a longer one alone, in its chunks). + fn prefillMany(p: *anyopaque, streams: []const *lanes.Stream, errs: []?anyerror) !void { + const x = cast(p); + var pieces: [st.max_streams]Forward.Piece = undefined; + var which: [st.max_streams]usize = undefined; + var n: usize = 0; + var rows: usize = 0; + for (streams, errs, 0..) |s, *err, i| { + const seq = x.open(s) catch |e| { + err.* = e; + continue; + }; + const ids = s.prompt(); + if (ids.len > prompt_rows) { + x.prefillOpened(s, seq) catch |e| { + x.dropSeq(s); + err.* = e; + }; + continue; + } + if (n == st.max_streams or rows + ids.len > prompt_rows) { + x.pass(streams, errs, pieces[0..n], which[0..n]); + n = 0; + rows = 0; + } + pieces[n] = .{ .ids = ids, .seq = seq, .last = true }; + which[n] = i; + n += 1; + rows += ids.len; + } + if (n > 0) x.pass(streams, errs, pieces[0..n], which[0..n]); + } + + /// One shared prompt pass; on a failure every stream in it fails. + fn pass(x: *Cuda, streams: []const *lanes.Stream, errs: []?anyerror, pieces: []const Forward.Piece, which: []const usize) void { + x.passPieces(streams, pieces, which) catch |e| for (which) |i| { + x.dropSeq(streams[i]); + errs[i] = e; + }; + } + + fn passPieces(x: *Cuda, streams: []const *lanes.Stream, pieces: []const Forward.Piece, which: []const usize) !void { + var f = x.e.forward(); + const d = x.drafter orelse return f.chunks(pieces); + f.taps = x.e.scratch.taps; + try f.chunks(pieces); + // each drafting stream's last window of taps, consecutive whole prompts absorbed together + const row = drafts.taps_width * 2; + var group: [st.max_streams]dr.Absorb = undefined; + var g: usize = 0; + var g_rows: usize = 0; + var g_from: usize = 0; // the group's first tap row + var at: usize = 0; // this piece's first row in the pass + for (pieces, which) |pc, i| { + defer at += pc.ids.len; + const s = streams[i]; + if (!s.drafts) continue; + const ds = try x.draftingOf(s); + ds.ctx.len = 0; + ds.ctx.end = 0; + const tap_from = pc.ids.len -| dl.window; + ds.ctx.skip(tap_from); + const taken = pc.ids.len - tap_from; + if (g > 0 and (tap_from > 0 or g_from + g_rows != at or g_rows + taken > dl.window + 1)) { + try d.absorbMany(group[0..g], x.e.scratch.taps + g_from * row); + g = 0; + } + if (g == 0) { + g_from = at + tap_from; + g_rows = 0; + } + group[g] = .{ .ctx = &ds.ctx, .rows = taken }; + g += 1; + g_rows += taken; + } + if (g > 0) try d.absorbMany(group[0..g], x.e.scratch.taps + g_from * row); + } + + fn dropSeq(x: *Cuda, s: *lanes.Stream) void { + if (x.seqs.fetchRemove(s)) |entry| { + entry.value.deinit(); + x.gpa.destroy(entry.value); + } + } + + fn first(p: *anyopaque, s: *lanes.Stream, position: u64) !u64 { + const x = cast(p); + const seq = try x.settled(s); + if (position != seq.pos) return error.InvalidQwenPosition; + var f = x.e.forward(); + try f.head(seq); + var out: [1]u32 = undefined; + try x.drawRows(&.{s.sampling}, &.{position}, &out); + return out[0]; + } + + fn queue(p: *anyopaque, s: *lanes.Stream, feed: be.Feed, position: u64) !u64 { + const x = cast(p); + const seq = try x.settled(s); + if (position != seq.pos + 1) return error.InvalidQwenPosition; + const token: u32 = switch (feed) { + .value => |v| v, + .handle => |h| @intCast(h), + }; + var f = x.e.forward(); + const part = [_]Part{.{ .seq = seq, .tokens = &.{token} }}; + try f.round(&part); + var out: [1]u32 = undefined; + try x.drawRows(&.{s.sampling}, &.{position}, &out); + try f.commit(&part, &.{&.{0}}); + return out[0]; + } + + fn read(_: *anyopaque, handle: u64) !u32 { + if (handle >= c.vocab) return error.InvalidQwenToken; + return @intCast(handle); + } + + /// One forward over every window (one stream: tree_forward's, several: multi_tree_forward's), each row drawn. + fn verify(p: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { + const x = cast(p); + if (windows.len != out.len or windows.len == 0 or windows.len > st.max_streams) return error.InvalidQwenWindows; + try x.settle(); + const pend = &x.pending; + var total: usize = 0; + for (windows, out, 0..) |w, o, k| { + const rows = w.rows(); + if (w.held != 0 or w.positions.len != rows or rows > st.max_rows or total + rows > st.round_rows) return error.UnsupportedQwenWindow; + if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; + const seq = try x.seqOf(w.stream); + pend.tokens[total] = w.pending; + @memcpy(pend.tokens[total + 1 ..][0..w.tokens.len], w.tokens); + const parents: ?[]const i32 = if (w.parents) |given| blk: { + if (given.len != rows) return error.InvalidQwenWindows; + @memcpy(pend.parents[total..][0..rows], given); + break :blk pend.parents[total..][0..rows]; + } else null; + // each row's draw is keyed at its depth past the committed position + var depth: [st.max_rows]usize = undefined; + if (parents) |ps| { + try trees.depths(ps, depth[0..rows]); + } else for (depth[0..rows], 0..) |*d, r| { + d.* = r; + } + for (w.positions, depth[0..rows]) |pos, d| if (pos != seq.pos + d + 1) return error.InvalidQwenPosition; + pend.parts[k] = .{ .seq = seq, .tokens = pend.tokens[total..][0..rows], .parents = parents }; + pend.starts[k] = total; + pend.streams[k] = w.stream; + // a window kept whole: its rows in order (keep replaces them with its path) + for (pend.path_rows[total..][0..rows], 0..) |*row, r| row.* = @intCast(r); + pend.paths[k] = pend.path_rows[total..][0..rows]; + total += rows; + } + var f = x.e.forward(); + if (x.drafter != null) f.taps = x.e.scratch.taps; + try f.round(pend.parts[0..windows.len]); + pend.n = windows.len; + var rules: [st.round_rows]?lanes.Sampling = undefined; + var positions: [st.round_rows]u64 = undefined; + var drawn: [st.round_rows]u32 = undefined; + var row: usize = 0; + for (windows) |w| for (w.positions) |pos| { + rules[row] = w.stream.sampling; + positions[row] = pos; + row += 1; + }; + try x.drawRows(rules[0..total], positions[0..total], &drawn); + row = 0; + for (windows, out) |w, o| { + @memcpy(o.sampled, drawn[row..][0..o.sampled.len]); + row += o.sampled.len; + @memcpy(o.drafts, w.tokens); + } + } + + /// A window keeps its accepted path (a chain's prefix, or a tree's root-to-node rows); verify keeps all of it. + fn keep(p: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { + const x = cast(p); + const pend = &x.pending; + for (windows, paths) |w, path| { + const k = for (pend.streams[0..pend.n], 0..) |s, k| { + if (s == w.stream) break k; + } else return error.UnknownQwenStream; + const rows = pend.parts[k].tokens.len; + if (path.len == 0 or path.len > rows) return error.InvalidQwenCommit; + const slot: [*]u32 = @constCast(pend.paths[k].ptr); + for (path, 0..) |r, i| { + if (r >= rows) return error.InvalidQwenCommit; + slot[i] = r; + } + pend.paths[k] = slot[0..path.len]; + } + } + + /// The kept rows' taps into each stream's drafter context, then each stream's next drafts. + fn draft(p: *anyopaque, requests: []const be.DraftRequest) !void { + const x = cast(p); + const d = x.drafter orelse { + for (requests) |r| if (r.depth != 0) return error.QwenCheckpointHasNoDraftHead; + return; + }; + const pend = &x.pending; + var kept: [st.max_streams]drafts.Kept = undefined; + var rows: [st.round_rows]u32 = undefined; + var asks: [st.max_streams]drafts.Ask = undefined; + var firsts: [st.max_streams]std.ArrayList(u32) = @splat(.empty); + defer for (&firsts) |*l| l.deinit(x.gpa); + var nk: usize = 0; + var nr: usize = 0; + var na: usize = 0; + for (requests) |r| { + if (na == st.max_streams) return error.InvalidQwenWindows; + const ds = x.drafting.get(r.stream) orelse continue; + var context = r.stream.context.items; + if (r.rows) |path| { + const k = for (pend.streams[0..pend.n], 0..) |s, k| { + if (s == r.stream) break k; + } else return error.UnknownQwenStream; + for (path, rows[nr..][0..path.len]) |row, *g| g.* = @intCast(pend.starts[k] + row); + kept[nk] = .{ .x = ds, .rows = rows[nr..][0..path.len] }; + nk += 1; + nr += path.len; + } else if (r.first) |feed| { + // the prompt's taps went in with its prefill; the first token ends the context the copies read + const token: u32 = switch (feed) { + .value => |v| v, + .handle => |h| try read(p, h), + }; + try firsts[na].appendSlice(x.gpa, context); + try firsts[na].append(x.gpa, token); + context = firsts[na].items; + } + asks[na] = .{ .x = ds, .context = context, .sampling = r.stream.sampling }; + na += 1; + } + const f = x.e.forward(); + try drafts.absorb(d, f.ops, x.e.scratch.taps, x.e.scratch.copy_items, kept[0..nk]); + try drafts.propose(d, asks[0..na]); + } + + /// The drafts a stream holds as host tokens and parents (a copy's as a chain). + fn tree(p: *anyopaque, s: *lanes.Stream, gpa: std.mem.Allocator) !?lanes.stream.Held { + const x = cast(p); + const ds = x.drafting.get(s) orelse return null; + const tokens = try gpa.dupe(u32, ds.tokens.items); + errdefer gpa.free(tokens); + const parents: ?[]i32 = if (ds.kind == .tree) try gpa.dupe(i32, ds.parents.items) else null; + return .{ .count = @intCast(tokens.len), .tokens = tokens, .parents = parents }; + } + + /// Each held draft's chance of landing: a tree node's exp(-path score) (allocate's), a copied token's 1. + fn probabilities(p: *anyopaque, s: *lanes.Stream, out: []f64) !bool { + const x = cast(p); + const ds = x.drafting.get(s) orelse return false; + if (out.len > ds.scores.items.len) return false; + for (out, ds.scores.items[0..out.len]) |*o, score| o.* = @exp(-score); + return true; + } + + fn release(p: *anyopaque, s: *lanes.Stream) void { + const x = cast(p); + x.settle() catch |err| std.log.err("committing the last round before a release: {s}", .{@errorName(err)}); + if (x.seqs.fetchRemove(s)) |entry| { + entry.value.deinit(); + x.gpa.destroy(entry.value); + } + if (x.drafting.fetchRemove(s)) |entry| { + entry.value.deinit(); + x.gpa.destroy(entry.value); + } + } +}; + +fn value(word: u16) f64 { + return @as(f32, @bitCast(@as(u32, word) << 16)); +} + +/// The Metal backend's draw: the top-k candidates by value then id, then the shared fp64 keyed sampler. +pub fn draw(gpa: std.mem.Allocator, logits: []const u16, settings: ?lanes.Sampling, position: u64) !u32 { + var best: usize = 0; + for (logits, 0..) |word, id| { + const x = value(word); + if (!std.math.isFinite(x)) return error.NonfiniteQwenLogits; + if (x > value(logits[best])) best = id; + } + const sampling = settings orelse return @intCast(best); + if (sampling.temperature <= 0) return @intCast(best); + const k = if (sampling.top_k == 0) logits.len else @min(logits.len, sampling.top_k); + const values = try gpa.alloc(f64, k); + defer gpa.free(values); + const ids = try gpa.alloc(u64, k); + defer gpa.free(ids); + var filled: usize = 0; + for (logits, 0..) |word, id| { + const x = value(word); + if (k == logits.len) { + values[id] = x; + ids[id] = id; + continue; + } + if (filled == k and x <= values[k - 1]) continue; + var at = @min(filled, k - 1); + while (at > 0 and x > values[at - 1]) : (at -= 1) { + values[at] = values[at - 1]; + ids[at] = ids[at - 1]; + } + values[at] = x; + ids[at] = id; + filled = @min(filled + 1, k); + } + return @intCast(try lanes.sampling.choose(gpa, values, ids, position, sampling)); +} diff --git a/zig/src/families/qwen27/cuda/lone.zig b/zig/src/families/qwen27/cuda/lone.zig new file mode 100644 index 000000000..0397ea706 --- /dev/null +++ b/zig/src/families/qwen27/cuda/lone.zig @@ -0,0 +1,181 @@ +//! decode.draft_decode for one stream: copies from the context, else a DFlash2 tree; verify, walk, commit, absorb. + +const std = @import("std"); +const lanes = @import("lanes"); +const c = @import("shape.zig"); +const st = @import("state.zig"); +const tree = @import("tree.zig"); +const kern = @import("kernels.zig"); +const dr = @import("draft.zig"); +const dl = @import("draft_load.zig"); +const Engine = @import("engine.zig").Engine; +const Forward = @import("forward.zig").Forward; +const Part = @import("forward.zig").Part; +const Lanes = @import("lanes.zig").Cuda; + +const bf = 2; +const taps_width = dl.taps.len * dl.hidden; + +/// decode.CopyIndex: earlier 8-token matches of the context's suffix, the longest continuation proposed. +pub const CopyIndex = struct { + pub const n = 8; + const Key = [n]u32; + gpa: std.mem.Allocator, + positions: std.AutoHashMapUnmanaged(Key, std.ArrayList(usize)) = .empty, + indexed: usize = 0, + + pub fn deinit(x: *CopyIndex) void { + var it = x.positions.valueIterator(); + while (it.next()) |l| l.deinit(x.gpa); + x.positions.deinit(x.gpa); + } + + fn update(x: *CopyIndex, context: []const u32) !void { + if (context.len < n) return; + const last = context.len - n; // the trailing query itself stays out of the index + var start = x.indexed; + while (start < last) : (start += 1) { + const gop = try x.positions.getOrPut(x.gpa, context[start..][0..n].*); + if (!gop.found_existing) gop.value_ptr.* = .empty; + try gop.value_ptr.append(x.gpa, start); + } + x.indexed = @max(x.indexed, last); + } + + /// The longest continuation after an earlier occurrence of the last 8 tokens (latest first), at least 8 long. + pub fn propose(x: *CopyIndex, context: []const u32, max_nodes: usize) ![]const u32 { + if (context.len < 2 * n) return &.{}; + try x.update(context); + const starts = x.positions.get(context[context.len - n ..][0..n].*) orelse return &.{}; + var best: []const u32 = &.{}; + var i = starts.items.len; + while (i > 0) { + i -= 1; + const from = starts.items[i] + n; + const cont = context[from..@min(context.len, from + max_nodes)]; + if (cont.len > best.len) { + best = cont; + if (best.len == max_nodes) break; + } + } + return if (best.len >= n) best else &.{}; + } +}; + +/// decode.next_copy_rows: a copy's next window, twice as wide after a whole copy, half after a broken one. +pub fn nextCopyRows(rows: usize, whole: bool, tree_rows: usize, max_rows: usize) usize { + const first = @min(max_rows, @max(tree_rows, 16)); + return if (whole) @min(max_rows, @max(rows, first) * 2) else @max(first, @min(rows, max_rows) / 2); +} + +pub const Stats = struct { rounds: usize = 0, drafted: usize = 0, accepted: usize = 0 }; + +/// prefill_state with a drafter: the prompt in even chunks, its last `window` rows' taps into the context. +pub fn prefill(e: *Engine, seq: *st.Seq, ctx: *dr.Context, d: *dr.DFlash2, prompt: []const u32) !void { + try seq.reset(e.ops(), e.w.g); + ctx.len = 0; + ctx.end = 0; + const tap_from = prompt.len -| dl.window; + ctx.skip(tap_from); + var f = e.forward(); + f.taps = e.scratch.taps; + const rows = @import("engine.zig").prompt_rows; + const chunks = (prompt.len + rows - 1) / rows; + for (0..chunks) |j| { + const b = @import("engine.zig").chunkBounds(0, prompt.len, rows, j); + try f.chunk(prompt[b[0]..b[1]], seq, j == chunks - 1); + if (b[1] <= tap_from) continue; + const skip = tap_from -| b[0]; + try d.absorb(ctx, e.scratch.taps + skip * taps_width * bf, b[1] - b[0] - skip); + } +} + +/// draft_decode, greedy: rounds until `count` tokens or an end token; the first token is `pending`. +pub fn decode(e: *Engine, seq: *st.Seq, ctx: *dr.Context, d: *dr.DFlash2, prompt: []const u32, pending: u32, count: usize, max_rows: usize, out: *std.ArrayList(u32), stats: *Stats) !void { + const gpa = e.gpa; + var context: std.ArrayList(u32) = .empty; + defer context.deinit(gpa); + try context.appendSlice(gpa, prompt); + try out.append(gpa, pending); + try context.append(gpa, pending); + var copies: CopyIndex = .{ .gpa = gpa }; + defer copies.deinit(); + const tree_rows = max_rows; + var copy_rows = nextCopyRows(tree_rows, false, tree_rows, max_rows); + var tokens: [st.max_rows]u32 = undefined; + var parents: [st.max_rows]i32 = undefined; + var sampled: [st.max_rows]u32 = undefined; + var path: [st.max_rows]u32 = undefined; + var f = e.forward(); + f.taps = e.scratch.taps; + while (out.items.len < count and !e.w.isEos(out.items[out.items.len - 1])) { + const last = out.items[out.items.len - 1]; + tokens[0] = last; + parents[0] = -1; + var rows: usize = 1; + const copied = try copies.propose(context.items, copy_rows - 1); + if (copied.len > 0) { + for (copied, 1..) |t, r| { + tokens[r] = t; + parents[r] = @as(i32, @intCast(r)) - 1; + } + rows += copied.len; + } else if (ctx.len > 0) { + const p = try d.propose(ctx, last, context.items.len, tree_rows - 1, null); + for (p.tokens, p.parents, 1..) |t, pr, r| { + tokens[r] = t; + parents[r] = if (pr < 0) 0 else pr + 1; + } + rows += p.tokens.len; + } + const part = [_]Part{.{ .seq = seq, .tokens = tokens[0..rows], .parents = parents[0..rows] }}; + try f.round(&part); + try f.picks(rows, &sampled); + // the path: from the root, each next token the child whose token the target drew + var n: usize = 1; + path[0] = 0; + var terminal = sampled[0]; + while (out.items.len + n < count and !e.w.isEos(terminal)) { + const child = for (1..rows) |r| { + if (parents[r] == @as(i32, @intCast(path[n - 1])) and tokens[r] == terminal) break r; + } else break; + path[n] = @intCast(child); + n += 1; + terminal = sampled[child]; + } + try f.commit(&part, &.{path[0..n]}); + // the drafter reads the kept rows' taps, gathered in path order + var items: [st.max_rows]kern.Copy = undefined; + const gathered = e.scratch.taps + st.max_rows * taps_width * bf; + for (path[0..n], 0..) |row, j| items[j] = .{ .dst = gathered + j * taps_width * bf, .src = e.scratch.taps + row * taps_width * bf, .bytes = taps_width * bf }; + try f.ops.upload(e.scratch.copy_items, std.mem.sliceAsBytes(items[0..n])); + try f.ops.copies(e.scratch.copy_items, n); + try d.absorb(ctx, gathered, n); + for (path[1..n]) |row| { + try out.append(gpa, tokens[row]); + try context.append(gpa, tokens[row]); + } + try out.append(gpa, terminal); + try context.append(gpa, terminal); + if (copied.len > 0) copy_rows = nextCopyRows(copy_rows, n == rows, tree_rows, max_rows); + stats.rounds += 1; + stats.drafted += rows - 1; + stats.accepted += n - 1; + } +} + +test "copy windows grow after whole copies and shrink after broken ones" { + try std.testing.expectEqual(@as(usize, 12), nextCopyRows(12, false, 12, 12)); + try std.testing.expectEqual(@as(usize, 12), nextCopyRows(12, true, 12, 12)); + try std.testing.expectEqual(@as(usize, 16), nextCopyRows(16, false, 12, 64)); + try std.testing.expectEqual(@as(usize, 32), nextCopyRows(16, true, 12, 64)); +} + +test "copies propose the latest earlier continuation of the last 8 tokens" { + const gpa = std.testing.allocator; + var x: CopyIndex = .{ .gpa = gpa }; + defer x.deinit(); + const ctx = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 1, 2, 3, 4, 5, 6, 7, 8 }; + try std.testing.expectEqualSlices(u32, &.{ 9, 10, 11, 12, 13, 14, 15, 16 }, try x.propose(&ctx, 8)); + try std.testing.expectEqual(@as(usize, 0), (try x.propose(ctx[0..15], 8)).len); +} diff --git a/zig/src/families/qwen27/cuda/native.zig b/zig/src/families/qwen27/cuda/native.zig new file mode 100644 index 000000000..6e3f362fd --- /dev/null +++ b/zig/src/families/qwen27/cuda/native.zig @@ -0,0 +1,111 @@ +//! Qwen3.5-family dense models on CUDA behind the native family interface (native/cuda.zig's registry). + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const c = @import("shape.zig"); +const st = @import("state.zig"); +const engine = @import("engine.zig"); +const Lanes = @import("lanes.zig").Cuda; +const dr = @import("draft.zig"); + +pub const model_type = "qwen3_5"; +pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const default_context: i64 = 32768; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = engine.prompt_rows; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, drafter: ?[]const u8 = null }; + +/// A lone drafted stream's own driver (none here: one stream decodes in the lane core too, as under --parallel). +pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool) anyerror!bool; + +/// What the native server drives: the lane backend, the facts its round loop reads, and how to free it. +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + /// Device bytes each admitted stream allocates for its own state and caches. + stream_bytes: usize, + /// Streams that allocate nothing (none here: every stream's state is its own). + free_streams: u32 = 0, + /// What a --slide trainer keeps out of the streams' room (none: the 27B does not learn on CUDA yet). + learn_bytes: usize = 0, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, +}; + +const Owned = struct { gpa: std.mem.Allocator, e: *engine.Engine, lanes: Lanes, draft: ?*dr.DFlash2 = null }; + +/// The drafter the Python engine pairs with the 27B once pulled (z-lab/Qwen3.8-27B-DFlash2). +const drafter_repo = "models--z-lab--Qwen3.8-27B-DFlash2"; + +fn getenv(name: [:0]const u8) ?[]const u8 { + return std.mem.span(std.c.getenv(name) orelse return null); +} + +/// --drafter, else the pulled drafter's snapshot in the Hugging Face cache; null when neither holds one. +fn drafterDir(a: std.mem.Allocator, io: std.Io, given: ?[]const u8) !?[]const u8 { + if (given) |d| return try a.dupe(u8, d); + const hub = if (getenv("HF_HUB_CACHE")) |h| try a.dupe(u8, h) else if (getenv("HF_HOME")) |h| try std.fs.path.join(a, &.{ h, "hub" }) else blk: { + const home = getenv("HOME") orelse return null; + break :blk try std.fs.path.join(a, &.{ home, ".cache", "huggingface", "hub" }); + }; + defer a.free(hub); + const snaps = try std.fs.path.join(a, &.{ hub, drafter_repo, "snapshots" }); + defer a.free(snaps); + var dir = std.Io.Dir.cwd().openDir(io, snaps, .{ .iterate = true }) catch return null; + defer dir.close(io); + var it = dir.iterate(); + while (try it.next(io)) |entry| if (entry.kind == .directory) return try std.fs.path.join(a, &.{ snaps, entry.name }); + return null; +} + +/// The streams the lane core may admit at most; the server's memory admission picks how many fit. +const most_streams = 32; + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + // the 27B's glue runs the Python engine's Triton kernels, captured on this chip (no own glue kernels yet) + const e = try engine.Engine.init(gpa, io, ctx, dir, kernels orelse return error.QwenNeedsCapturedKernels, .{ .context = o.context }); + errdefer e.deinit(); + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.* = .{ .gpa = gpa, .e = e, .lanes = try Lanes.init(gpa, e, most_streams) }; + errdefer own.lanes.deinit(); + if (o.drafts) if (try drafterDir(gpa, io, o.drafter)) |draft_dir| { + defer gpa.free(draft_dir); + const f = e.forward(); + own.draft = try dr.DFlash2.init(gpa, io, f.ops, f.t, &e.w, draft_dir, dir); + own.lanes.drafter = own.draft; + std.log.info("DFlash2 drafts every stream: {s}", .{draft_dir}); + // the forward's ms by rows: a shared round's tree widths follow the curve + try own.lanes.calibrate(io, most_streams); + }; + return .{ + .backend = own.lanes.backend(), + .facts = own.lanes.facts(), + .rows = st.max_rows, + .stream_bytes = st.Seq.bytes(e.w.g, o.context) + if (own.draft != null) dr.Context.bytes() else 0, + .ctx = own, + .deinit = release, + }; +} + +/// A request this engine refuses, in words; null: none of its own. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + error.QwenCheckpointHasNoDraftHead => "this checkpoint has no draft head on CUDA yet: serve with --no-drafts", + else => null, + }; +} + +fn release(p: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(p)); + own.lanes.deinit(); + if (own.draft) |d| d.deinit(); + own.e.deinit(); + own.gpa.destroy(own); +} diff --git a/zig/src/families/qwen27/cuda/round.zig b/zig/src/families/qwen27/cuda/round.zig new file mode 100644 index 000000000..2b2d4e269 --- /dev/null +++ b/zig/src/families/qwen27/cuda/round.zig @@ -0,0 +1,230 @@ +//! A verify round's inputs laid out as tree_forward or multi_tree_forward lays them (Triton variants follow alignment). + +const std = @import("std"); +const c = @import("shape.zig"); +const kern = @import("kernels.zig"); +const tri = @import("triton.zig"); +const st = @import("state.zig"); +const tree = @import("tree.zig"); + +const bf = 2; + +/// One stream's window: its tokens from its committed position, and each row's parent (null: a chain). +pub const Part = struct { seq: *st.Seq, tokens: []const u32, parents: ?[]const i32 = null }; + +/// Where a round's inputs live on the device, and what the layer loop launches with. +pub const Round = struct { + rows: usize, + streams: usize, + multi: bool, + ids: u64, + pos: u64, + windows: u64, + sids: u64, // multi: each row's stream; else the windows again (gdn_pre passes WIN as SID) + plan: u64, // GDN schedule entries (rows, 3) + starts: u64, // stream row starts (streams + 1) + slots: usize, // the GDN schedule's live states (0: every window a chain) + most: usize, // the most rows a stream's window holds + tables: u64, // multi: (DeltaNet layers, streams) state pointers, a layer's row at its DeltaNet index + attn: tri.Tri.Plan, + offs: u64, // (attention layers, streams, 2) int64 cache offsets; single: indexed by layer + starts_host: [st.max_streams + 1]usize, + + /// The conv rows gdn_pre reads for layer `i` (its DeltaNet index `j`): the stream's own, or the stacked copy. + pub fn conv(r: *const Round, parts: []const Part, s: *const st.Scratch, g: c.Geometry, i: usize, j: usize) u64 { + if (!r.multi) return parts[0].seq.conv[i]; + return s.conv_cat + j * st.max_streams * (c.conv_taps - 1) * g.convDim() * bf; + } + + /// Layer `i`'s attention cache offsets: tree_forward's (one stream, by layer) or the packed (layer, streams) rows. + pub fn offsets(r: *const Round, i: usize, a: usize) u64 { + return if (r.multi) r.offs + a * r.streams * 16 else r.offs + i * 16; + } + + /// The DeltaNet state table of DeltaNet layer `j` (multi only). + pub fn table(r: *const Round, j: usize) u64 { + return r.tables + j * r.streams * 8; + } +}; + +/// A part's parents: its own tree's, or a chain's. +pub fn parentsOf(p: Part, buf: []i32) []const i32 { + if (p.parents) |x| return x; + for (buf[0..p.tokens.len], 0..) |*x, r| x.* = @as(i32, @intCast(r)) - 1; + return buf[0..p.tokens.len]; +} + +/// attention.plan_host: rows, (first, rows, keys, slots) per stream, items, then window parents (global rows). +fn attnPlan(gpa: std.mem.Allocator, out: *std.ArrayList(i32), parts: []const Part, group: usize) !usize { + for (parts, 0..) |p, s| for (p.tokens) |_| try out.append(gpa, @intCast(s)); + var start: usize = 0; + for (parts) |p| { + const w = p.tokens.len; + try out.appendSlice(gpa, &.{ @intCast(start), @intCast(w), @intCast(p.seq.pos), @intCast(tri.slots(p.seq.pos, w)) }); + start += w; + } + var n_items: usize = 0; + for (parts, 0..) |p, s| { + const w = p.tokens.len; + const folded = tri.groups(p.seq.pos, w); + const loose = p.seq.pos / tri.chunk - folded * tri.group_chunks; + for (0..folded + loose) |j| { + const code: i32 = if (j < folded) @intCast(j) else -1 - @as(i32, @intCast(j - folded)); + var first: usize = 0; + while (first < w * group) : (first += 16) { + try out.appendSlice(gpa, &.{ @intCast(s), @intCast(first), code }); + n_items += 1; + } + } + } + start = 0; + var buf: [tree.max_nodes]i32 = undefined; + for (parts) |p| { + for (parentsOf(p, &buf)) |x| try out.append(gpa, if (x < 0) -1 else @intCast(start + @as(usize, @intCast(x)))); + start += p.tokens.len; + } + return n_items; +} + +/// Each stream's key and value caches as bf16 element offsets from the scratch's origin. +fn offsetOf(s: *const st.Scratch, ptr: u64) i64 { + return @divExact(@as(i64, @bitCast(ptr -% s.origin)), bf); +} + +/// Every stream's host plans: positions (pos + depth), the GDN schedule (nodes offset by its start), conv windows. +const Host = struct { pos: std.ArrayList(i32) = .empty, entries: std.ArrayList(i32) = .empty, windows: std.ArrayList(i32) = .empty, slots: usize = 0, most: usize = 0 }; + +fn plans(gpa: std.mem.Allocator, parts: []const Part, starts: []const usize, multi: bool, h: *Host) !void { + const keep = c.conv_taps - 1; + var buf: [tree.max_nodes]i32 = undefined; + var dep: [tree.max_nodes]usize = undefined; + var win: [tree.max_nodes * c.conv_taps]i32 = undefined; + for (parts, 0..) |p, k| { + const parents = parentsOf(p, &buf); + const w = parents.len; + try tree.depths(parents, dep[0..w]); + for (dep[0..w]) |d| try h.pos.append(gpa, @intCast(p.seq.pos + d)); + const s = try tree.schedule(gpa, parents); + defer gpa.free(s.entries); + for (s.entries, 0..) |x, j| try h.entries.append(gpa, if (j % 3 == 0) x + @as(i32, @intCast(starts[k])) else x); + h.slots = @max(h.slots, s.slots); + h.most = @max(h.most, w); + // one stream's windows index [its state; its rows]; several streams' index [their states; all rows] + tree.convWindows(parents, keep, if (multi) starts[k] else 0, win[0 .. w * (keep + 1)]); + try h.windows.appendSlice(gpa, win[0 .. w * (keep + 1)]); + } +} + +/// Uploads the round's inputs and runs attention's _paths; the caller's layer loop reads the returned Round. +pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, g: c.Geometry, parts: []const Part) !Round { + if (parts.len == 0 or parts.len > st.max_streams) return error.InvalidQwenWindows; + var W: usize = 0; + var r: Round = undefined; + r.starts_host[0] = 0; + for (parts, 0..) |p, k| { + if (p.tokens.len == 0 or p.tokens.len > tree.max_nodes) return error.InvalidQwenWindows; + if (p.parents) |x| if (x.len != p.tokens.len) return error.InvalidQwenWindows; + if (p.seq.pos + p.tokens.len > p.seq.capacity) return error.PromptTooLong; + W += p.tokens.len; + r.starts_host[k + 1] = W; + } + if (W > st.round_rows) return error.WindowTooWide; + r.rows = W; + r.streams = parts.len; + r.multi = parts.len > 1; + const S = parts.len; + const group = g.query_heads / g.kv_heads; + var h: Host = .{}; + defer { + h.pos.deinit(gpa); + h.entries.deinit(gpa); + h.windows.deinit(gpa); + } + try plans(gpa, parts, r.starts_host[0 .. S + 1], r.multi, &h); + r.slots = h.slots; + r.most = h.most; + var host: std.ArrayList(i32) = .empty; + defer host.deinit(gpa); + var offs: std.ArrayList(i64) = .empty; + defer offs.deinit(gpa); + if (!r.multi) { + // tree_forward: tokens, positions, plan (entries then starts), windows and attention's plan apart + const p = parts[0]; + for (p.tokens) |tok| try host.append(gpa, @intCast(tok)); + try ops.upload(s.ids, std.mem.sliceAsBytes(host.items)); + try ops.upload(s.pos, std.mem.sliceAsBytes(h.pos.items)); + try ops.upload(s.windows, std.mem.sliceAsBytes(h.windows.items)); + try ops.s.synchronize(); + host.clearRetainingCapacity(); + try host.appendSlice(gpa, h.entries.items); + try host.appendSlice(gpa, &.{ 0, @intCast(W) }); + try ops.upload(s.plan, std.mem.sliceAsBytes(host.items)); + try ops.s.synchronize(); + host.clearRetainingCapacity(); + const n_items = try attnPlan(gpa, &host, parts, group); + try ops.upload(s.attn_plan, std.mem.sliceAsBytes(host.items)); + try ops.s.synchronize(); + for (0..g.layers) |i| try offs.appendSlice(gpa, if (c.linear(i)) &.{ 0, 0 } else &.{ offsetOf(s, p.seq.keys[i]), offsetOf(s, p.seq.values[i]) }); + try ops.upload(s.offs, std.mem.sliceAsBytes(offs.items)); + try ops.s.synchronize(); + r.ids = s.ids; + r.pos = s.pos; + r.windows = s.windows; + r.sids = s.windows; + r.plan = s.plan; + r.starts = s.plan + 3 * W * 4; + r.tables = 0; + r.offs = s.offs; + r.attn = .{ .rows = s.attn_plan, .streams = s.attn_plan + W * 4, .items = s.attn_plan + (W + 4) * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; + try t.paths(s.attn_plan + (W + 4 + 3 * n_items) * 4, s.paths, s.depths, W); + return r; + } + // multi_tree_forward: positions, stream ids, plan entries, starts, ids, windows, attention's plan in one copy + try host.appendSlice(gpa, h.pos.items); + for (parts, 0..) |p, k| for (p.tokens) |_| try host.append(gpa, @intCast(k)); + try host.appendSlice(gpa, h.entries.items); + for (r.starts_host[0 .. S + 1]) |x| try host.append(gpa, @intCast(x)); + for (parts) |p| for (p.tokens) |tok| try host.append(gpa, @intCast(tok)); + try host.appendSlice(gpa, h.windows.items); + const at_attn = host.items.len; + const n_items = try attnPlan(gpa, &host, parts, group); + try ops.upload(s.round, std.mem.sliceAsBytes(host.items)); + try ops.s.synchronize(); + const base = s.round; + r.pos = base; + r.sids = base + W * 4; + r.plan = base + 2 * W * 4; + r.starts = base + 5 * W * 4; + r.ids = base + (5 * W + S + 1) * 4; + r.windows = base + (6 * W + S + 1) * 4; + const ab = base + at_attn * 4; + r.attn = .{ .rows = ab, .streams = ab + W * 4, .items = ab + (W + 4 * S) * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; + // the state pointer table and cache offsets, one copy each; every stream's conv rows stacked by layer + var ptrs: std.ArrayList(u64) = .empty; + defer ptrs.deinit(gpa); + var stack: std.ArrayList(kern.Copy) = .empty; + defer stack.deinit(gpa); + const keep = c.conv_taps - 1; + var j: usize = 0; + for (0..g.layers) |i| { + if (!c.linear(i)) { + for (parts) |p| try offs.appendSlice(gpa, &.{ offsetOf(s, p.seq.keys[i]), offsetOf(s, p.seq.values[i]) }); + continue; + } + const cat = s.conv_cat + j * st.max_streams * keep * g.convDim() * bf; + for (parts, 0..) |p, k| { + try ptrs.append(gpa, p.seq.rec[i]); + try stack.append(gpa, .{ .dst = cat + k * keep * g.convDim() * bf, .src = p.seq.conv[i], .bytes = keep * g.convDim() * bf }); + } + j += 1; + } + try ops.upload(s.tables, std.mem.sliceAsBytes(ptrs.items)); + try ops.upload(s.offs, std.mem.sliceAsBytes(offs.items)); + try ops.upload(s.copy_items, std.mem.sliceAsBytes(stack.items)); + try ops.copies(s.copy_items, stack.items.len); + try ops.s.synchronize(); + r.tables = s.tables; + r.offs = s.offs; + try t.paths(ab + (W + 4 * S + 3 * n_items) * 4, s.paths, s.depths, W); + return r; +} diff --git a/zig/src/families/qwen27/cuda/shape.zig b/zig/src/families/qwen27/cuda/shape.zig new file mode 100644 index 000000000..b018dc22a --- /dev/null +++ b/zig/src/families/qwen27/cuda/shape.zig @@ -0,0 +1,102 @@ +//! The shapes the CUDA engine is built for: the family's parsed config (config.zig), held to the 27B's. +const std = @import("std"); +const config = @import("../config.zig"); + +pub const vocab = 248320; +pub const linear_dim = 128; // a DeltaNet key or value head's width +pub const conv_taps = 4; +pub const head_dim = 256; +pub const rotary_dim = 64; +pub const eps: f32 = 1e-6; +pub const theta: f32 = 10000000; + +/// Whether layer `index` is a DeltaNet layer: every fourth is attention. +pub fn linear(index: usize) bool { + return index % 4 != 3; +} + +/// The widths a forward sizes its buffers by. +pub const Geometry = struct { + hidden: usize, + intermediate: usize, + layers: usize, + linear_k_heads: usize, // DeltaNet key (and query) heads + linear_v_heads: usize, // DeltaNet value heads, each key head serving a run of them + query_heads: usize, + kv_heads: usize, + + /// The DeltaNet input projection's row: queries, keys, then values. + pub fn convDim(g: Geometry) usize { + return (2 * g.linear_k_heads + g.linear_v_heads) * linear_dim; + } + pub fn kInner(g: Geometry) usize { + return g.linear_k_heads * linear_dim; + } + pub fn vInner(g: Geometry) usize { + return g.linear_v_heads * linear_dim; + } + pub fn qInner(g: Geometry) usize { + return g.query_heads * head_dim; + } + pub fn kvInner(g: Geometry) usize { + return g.kv_heads * head_dim; + } + /// The widest row an activation buffer holds: the recurrence's values or the attention's heads. + pub fn inner(g: Geometry) usize { + return @max(g.vInner(), g.qInner()); + } + /// One layer's recurrent state (fp32) and its conv window (bf16). + pub fn deltaBytes(g: Geometry) usize { + return g.linear_v_heads * linear_dim * linear_dim * 4; + } + pub fn convBytes(g: Geometry) usize { + return (conv_taps - 1) * g.convDim() * 2; + } +}; + +/// Qwen3.8-27B, the one checkpoint shape the CUDA engine serves. +pub const the_27b: Geometry = .{ .hidden = 5120, .intermediate = 17408, .layers = 64, .linear_k_heads = 16, .linear_v_heads = 48, .query_heads = 24, .kv_heads = 4 }; + +/// A checkpoint's config.json (and generation_config.json's end tokens), refused unless it is the 27B's. +pub const Config = struct { + g: Geometry, + eos: [config.max_eos]u32, + eos_count: usize, + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const text = try file(gpa, io, dir, "config.json") orelse return error.MissingConfig; + defer gpa.free(text); + const generation = try file(gpa, io, dir, "generation_config.json"); + defer if (generation) |x| gpa.free(x); + return check(try config.parse(gpa, text, generation)); + } + + fn file(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) !?[]u8 { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + return std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)) catch |err| switch (err) { + error.FileNotFound => null, + else => err, + }; + } +}; + +/// The parsed config as the CUDA engine's Config, when every shape the kernels assume is the 27B's. +pub fn check(c: config.Config) !Config { + const g: Geometry = .{ .hidden = c.hidden, .intermediate = c.intermediate, .layers = c.layers, .linear_k_heads = c.k_heads, .linear_v_heads = c.v_heads, .query_heads = c.heads, .kv_heads = c.kv_heads }; + if (!std.meta.eql(g, the_27b) or c.vocab != vocab or c.head_dim != head_dim or c.dk != linear_dim or c.dv != linear_dim) return error.UnsupportedQwenGeometry; + if (c.conv_kernel != conv_taps or c.rope_dims != rotary_dim or c.eps != eps or c.rope_theta != theta) return error.UnsupportedQwenGeometry; + for (c.kinds[0..c.layers], 0..) |kind, i| if ((kind == .linear) != linear(i)) return error.UnsupportedQwenGeometry; + return .{ .g = g, .eos = c.eos, .eos_count = c.eos_count }; +} + +test "the 27B's shapes pass and another hidden size is refused" { + var c: config.Config = .{ .hidden = 5120, .intermediate = 17408, .layers = 64, .vocab = vocab, .heads = 24, .kv_heads = 4, .head_dim = head_dim, .k_heads = 16, .v_heads = 48, .dk = linear_dim, .dv = linear_dim, .conv_kernel = conv_taps, .eps = eps, .rope_dims = rotary_dim, .rope_theta = theta }; + for (c.kinds[0..64], 0..) |*k, i| k.* = if (linear(i)) .linear else .attention; + c.eos[0] = 248044; + c.eos_count = 1; + const ok = try check(c); + try std.testing.expectEqual(@as(usize, 1), ok.eos_count); + c.hidden = 2048; + try std.testing.expectError(error.UnsupportedQwenGeometry, check(c)); +} diff --git a/zig/src/families/qwen27/cuda/state.zig b/zig/src/families/qwen27/cuda/state.zig new file mode 100644 index 000000000..4cbe06c07 --- /dev/null +++ b/zig/src/families/qwen27/cuda/state.zig @@ -0,0 +1,304 @@ +//! A stream's committed state (DeltaNet recurrences, conv windows, attention caches) and the forward's scratch. + +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("shape.zig"); +const kern = @import("kernels.zig"); +const tri = @import("triton.zig"); +const torch_ops = kern.torch_ops; + +/// The most rows a stream's verify window takes (the multi-stream GDN tree kernel's limit). +pub const max_rows = 16; +/// The most streams a shared round takes. +pub const max_streams = 16; +/// The most rows a verify round takes: every stream's whole window. +pub const round_rows = max_rows * max_streams; +/// The most items a batched copy takes: every stream's layers, three rows or a key and a value each. +pub const copy_items = max_streams * 64 * 3 + round_rows * 64 * 2; +/// The most candidates a sampled row reads from the GPU: top_k up to 1,024, and exact_sampling.MARGIN past it. +pub const max_candidates = 1024 + 8; + +/// A bump allocator over one device buffer, every slice 256-byte aligned. +pub const Arena = struct { + buf: ?cuda.DeviceBuffer, // null: a measuring pass, every take counted and none backed + used: usize = 0, + + pub fn init(d: *const cuda.Driver, len: usize) !Arena { + return .{ .buf = try cuda.DeviceBuffer.alloc(d, len) }; + } + + pub fn deinit(a: *Arena) void { + if (a.buf) |*b| b.free(); + } + + pub fn take(a: *Arena, len: usize) !u64 { + const at = std.mem.alignForward(usize, a.used, 256); + a.used = at + len; + const b = a.buf orelse return 0; + if (a.used > b.len) return error.ArenaFull; + return b.ptr + at; + } +}; + +/// One stream's committed state: DeltaNet recurrences and conv windows, attention keys and values for `capacity`. +pub const Seq = struct { + arena: Arena, + rec: []u64, // by layer: (hv, dv, dk) fp32, 0 for attention layers + conv: []u64, // by layer: (3, conv dim) bf16 + keys: []u64, // by layer: (capacity, kv heads, head dim) bf16 + values: []u64, + head_in: u64, // the last prompt row normed, (hidden) bf16, and its group sums: the first draw's input + head_xs: u64, + capacity: usize, + pos: usize = 0, + gpa: std.mem.Allocator, + + /// Device bytes a stream of `capacity` positions takes. + pub fn bytes(g: c.Geometry, capacity: usize) usize { + var n: usize = std.mem.alignForward(usize, g.hidden * 2, 256) + std.mem.alignForward(usize, g.hidden / 64 * 4, 256); + const a = std.mem.alignForward; + for (0..g.layers) |i| n += if (c.linear(i)) a(usize, g.deltaBytes(), 256) + a(usize, g.convBytes(), 256) else 2 * a(usize, capacity * g.kvInner() * 2, 256); + return n; + } + + /// A fresh stream: zero recurrences and conv windows, caches for `capacity` positions. + pub fn init(gpa: std.mem.Allocator, ops: kern.Ops, g: c.Geometry, capacity: usize) !Seq { + var s: Seq = .{ .gpa = gpa, .arena = try Arena.init(ops.k.d, bytes(g, capacity)), .rec = &.{}, .conv = &.{}, .keys = &.{}, .values = &.{}, .head_in = 0, .head_xs = 0, .capacity = capacity }; + errdefer s.arena.deinit(); + s.rec = try gpa.alloc(u64, g.layers); + errdefer gpa.free(s.rec); + s.conv = try gpa.alloc(u64, g.layers); + errdefer gpa.free(s.conv); + s.keys = try gpa.alloc(u64, g.layers); + errdefer gpa.free(s.keys); + s.values = try gpa.alloc(u64, g.layers); + errdefer gpa.free(s.values); + const bf = 2; + s.head_in = try s.arena.take(g.hidden * 2); + s.head_xs = try s.arena.take(g.hidden / 64 * 4); + for (0..g.layers) |i| { + s.rec[i] = 0; + s.conv[i] = 0; + s.keys[i] = 0; + s.values[i] = 0; + const A = &s.arena; + if (c.linear(i)) { + s.rec[i] = try A.take(g.deltaBytes()); + s.conv[i] = try A.take(g.convBytes()); + } else { + s.keys[i] = try A.take(capacity * g.kvInner() * bf); + s.values[i] = try A.take(capacity * g.kvInner() * bf); + } + } + try s.reset(ops, g); + return s; + } + + pub fn deinit(s: *Seq) void { + s.arena.deinit(); + s.gpa.free(s.rec); + s.gpa.free(s.conv); + s.gpa.free(s.keys); + s.gpa.free(s.values); + s.* = undefined; + } + + /// Back to an empty prompt: zero recurrences and conv windows (State.__init__'s zeros). + pub fn reset(s: *Seq, ops: kern.Ops, g: c.Geometry) !void { + for (0..g.layers) |i| if (c.linear(i)) { + try ops.fill32(s.rec[i], 0, g.deltaBytes() / 4); + try ops.fill32(s.conv[i], 0, g.convBytes() / 4); + }; + s.pos = 0; + } +}; + +/// One DeltaNet layer's window record, kept until commit: what the replay steps and the conv window takes. +pub const GdnRecord = struct { q: u64, k: u64, v: u64, g: u64, beta: u64, qkv: u64 }; +/// One attention layer's window record: the rotated keys and the values the window's rows would add. +pub const AttnRecord = struct { k: u64, v: u64 }; + +/// The forward's device buffers, sized for `rows` rows a call (a decode window, or a prompt chunk). +pub const Scratch = struct { + arena: Arena, + rows: usize, + // the residual stream and the normed rows with their group sums + x: u64, + h: u64, + y: u64, + xs: u64, + r: u64, + pending: u64, + // DeltaNet + qkv: u64, + z: u64, + b: u64, + a: u64, + yr: u64, + gq: u64, // a prompt chunk's DeltaNet inputs (a window's live in its records) + gk: u64, + gv: u64, + gg: u64, + gbeta: u64, + gated: u64, + gated_xs: u64, + // attention + qg: u64, + key: u64, + value: u64, + q_rot: u64, + attn: u64, + // MLP + gate: u64, + up: u64, + act: u64, + act_xs: u64, + // the head's input, logits and picks + last_h: u64, + last_xs: u64, + logits: u64, + picks: u64, + logits32: u64, // the round's logits as fp32, the top-k input + cand_values: u64, // (rows, max_candidates) fp32 and int64: each row's top candidates + cand_ids: u64, + topk_scratch: u64, + // inputs a call uploads: ids, positions, conv windows, the GDN plan, the attention plan + ids: u64, + pos: u64, + sids: u64, // a batched prompt's rows' streams + prompt_keys: u64, // a batched prompt's rotated keys and values, before each stream's cache takes its rows + prompt_values: u64, + windows: u64, + plan: u64, + attn_plan: u64, + paths: u64, + depths: u64, + offs: u64, + part_o: u64, + part_m: u64, + part_l: u64, + part_chunks: usize, + origin: u64, // the bf16 origin cache offsets count from + conv_tmp: u64, + conv_old: u64, // a round's streams' conv rows that survive its commit, by stream and DeltaNet layer + copy_items: u64, // two batched-copy tables (kern.Copy), one for each phase of a commit + taps: u64, // a call's rows of the drafter's taps: (rows, 5 x hidden) bf16 + replay_table: u64, + replay_rows: u64, + round: u64, // a shared round's packed int32 inputs (multi_tree_forward's one host copy) + tables: u64, // (DeltaNet layers, streams) int64: each stream's state of each layer + conv_cat: u64, // by DeltaNet layer: every stream's conv rows, stacked + records: []GdnRecord, // by layer: the round's DeltaNet inputs, kept until commit + attn_records: []AttnRecord, // by layer: the round's rotated keys and values + gpa: std.mem.Allocator, + + /// Scratch for prompt chunks of `prompt_rows` and windows of `max_rows`, over caches of `capacity` positions. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: c.Geometry, prompt_rows: usize, capacity: usize) !Scratch { + var s: Scratch = undefined; + s.gpa = gpa; + s.records = try gpa.alloc(GdnRecord, g.layers); + errdefer gpa.free(s.records); + s.attn_records = try gpa.alloc(AttnRecord, g.layers); + errdefer gpa.free(s.attn_records); + s.arena = .{ .buf = null }; + try s.layout(g, prompt_rows, capacity); + s.arena = try Arena.init(d, s.arena.used); + errdefer s.arena.deinit(); + try s.layout(g, prompt_rows, capacity); + return s; + } + + fn layout(s: *Scratch, g: c.Geometry, prompt_rows: usize, capacity: usize) !void { + const R = @max(prompt_rows, round_rows); + const W = round_rows; + const bf = 2; + const chunks = (capacity + W + tri.chunk - 1) / tri.chunk; + // attention.plan_host: a query tile's items per committed chunk and folded group + const codes = capacity / tri.chunk + capacity / (tri.chunk * tri.group_chunks); + const items = codes * ((W * (g.query_heads / g.kv_heads) + 15) / 16); + s.rows = R; + s.part_chunks = chunks; + s.arena.used = 0; + const A = &s.arena; + s.origin = try A.take(128); + s.x = try A.take(R * g.hidden * bf); + s.h = try A.take(R * g.hidden * bf); + s.y = try A.take(R * g.hidden * bf); + s.r = try A.take(R * g.hidden * bf); + s.pending = try A.take(R * g.hidden * bf); + s.xs = try A.take(R * g.hidden / 64 * 4); + s.qkv = try A.take(R * g.convDim() * bf); + s.z = try A.take(R * g.vInner() * bf); + s.b = try A.take(R * g.linear_v_heads * bf); + s.a = try A.take(R * g.linear_v_heads * bf); + s.yr = try A.take(R * g.vInner() * bf); + s.gq = try A.take(R * g.kInner() * bf); + s.gk = try A.take(R * g.kInner() * bf); + s.gv = try A.take(R * g.vInner() * bf); + s.gg = try A.take(R * g.linear_v_heads * 4); + s.gbeta = try A.take(R * g.linear_v_heads * 4); + s.gated = try A.take(R * g.inner() * bf); + s.gated_xs = try A.take(R * g.inner() / 64 * 4); + s.qg = try A.take(R * 2 * g.qInner() * bf); + s.key = try A.take(R * g.kvInner() * bf); + s.value = try A.take(R * g.kvInner() * bf); + s.q_rot = try A.take(R * g.qInner() * bf); + s.attn = try A.take(R * g.qInner() * bf); + s.gate = try A.take(R * g.intermediate * bf); + s.up = try A.take(R * g.intermediate * bf); + s.act = try A.take(R * g.intermediate * bf); + s.act_xs = try A.take(R * g.intermediate / 64 * 4); + s.last_h = try A.take(W * g.hidden * bf); + s.last_xs = try A.take(W * g.hidden / 64 * 4); + s.logits = try A.take(W * c.vocab * bf); + s.picks = try A.take(W * 4); + s.logits32 = try A.take(W * c.vocab * 4); + s.cand_values = try A.take(W * max_candidates * 4); + s.cand_ids = try A.take(W * max_candidates * 8); + s.topk_scratch = try A.take(torch_ops.topkScratchBytes(W, c.vocab)); + s.ids = try A.take(R * 4); + s.pos = try A.take(R * 4); + s.sids = try A.take(R * 4); + s.prompt_keys = try A.take(R * g.kvInner() * bf); + s.prompt_values = try A.take(R * g.kvInner() * bf); + s.windows = try A.take(R * c.conv_taps * 4); + s.plan = try A.take(R * 3 * 4); + s.attn_plan = try A.take((W + 4 + 3 * items + W) * 4); + s.paths = try A.take(W * tri.max_nodes * 4); + s.depths = try A.take(W * 4); + s.offs = try A.take(g.layers * max_streams * 2 * 8); + s.part_o = try A.take(chunks * W * g.query_heads * c.head_dim * 4); + s.part_m = try A.take(chunks * W * g.query_heads * 4); + s.part_l = try A.take(chunks * W * g.query_heads * 4); + s.conv_tmp = try A.take((c.conv_taps - 1 + R) * g.convDim() * bf); + s.conv_old = try A.take(max_streams * g.layers * (c.conv_taps - 1) * g.convDim() * bf); + s.copy_items = try A.take(2 * copy_items * @sizeOf(kern.Copy)); + s.taps = try A.take(R * 5 * g.hidden * bf); + s.replay_table = try A.take(g.layers * (4 + max_streams) * 8); + s.replay_rows = try A.take(max_streams * (W + 1) * 4); + s.round = try A.take((6 * W + max_streams + 1 + W * c.conv_taps + W + 4 * max_streams + 3 * items * max_streams + W) * 4); + s.tables = try A.take(g.layers * max_streams * 8); + s.conv_cat = try A.take(g.layers * max_streams * (c.conv_taps - 1) * g.convDim() * bf); + for (0..g.layers) |i| { + if (c.linear(i)) { + s.records[i] = .{ + .q = try A.take(W * g.kInner() * bf), + .k = try A.take(W * g.kInner() * bf), + .v = try A.take(W * g.vInner() * bf), + .g = try A.take(W * g.linear_v_heads * 4), + .beta = try A.take(W * g.linear_v_heads * 4), + .qkv = try A.take(W * g.convDim() * bf), + }; + } else { + s.attn_records[i] = .{ .k = try A.take(W * g.kvInner() * bf), .v = try A.take(W * g.kvInner() * bf) }; + } + } + } + + pub fn deinit(s: *Scratch) void { + s.arena.deinit(); + s.gpa.free(s.records); + s.gpa.free(s.attn_records); + s.* = undefined; + } +}; diff --git a/zig/src/families/qwen27/cuda/tree.zig b/zig/src/families/qwen27/cuda/tree.zig new file mode 100644 index 000000000..f8cdf7187 --- /dev/null +++ b/zig/src/families/qwen27/cuda/tree.zig @@ -0,0 +1,181 @@ +//! A verify tree's host plans: gdn.schedule's slots and visiting order, and each node's conv window and depth. + +const std = @import("std"); + +/// The most nodes a stream's tree takes (forward._paths), and the most live states a schedule may hold. +pub const max_nodes = 128; +pub const max_slots = 32; + +/// gdn.schedule's result: (node, source, dest) entries in visiting order and the slots they use. +pub const Schedule = struct { entries: []i32, slots: usize }; + +/// forward._paths: each row's depth from the root (parents topologically sorted, -1 at row 0 only). +pub fn depths(parents: []const i32, out: []usize) !void { + if (parents.len == 0 or parents.len > max_nodes or parents[0] != -1) return error.InvalidQwenTree; + for (parents, 0..) |p, row| { + if (row == 0) { + out[0] = 0; + continue; + } + if (p < 0 or p >= row) return error.InvalidQwenTree; + out[row] = out[@intCast(p)] + 1; + } +} + +/// Whether every row's parent is the row before it: a chain. +pub fn chain(parents: []const i32) bool { + for (parents, 0..) |p, row| if (p != @as(i32, @intCast(row)) - 1) return false; + return true; +} + +/// gdn._order_slots: entries for one visiting order, and the slots it needs. +fn orderSlots(gpa: std.mem.Allocator, parents: []const i32, order: []const usize) !Schedule { + const n = parents.len; + const pos = try gpa.alloc(usize, n); + defer gpa.free(pos); + for (order, 0..) |node, i| pos[node] = i; + // each node's last child in this order, and whether all its children follow it directly + const last = try gpa.alloc(?usize, n); + defer gpa.free(last); + @memset(last, null); + const direct = try gpa.alloc(bool, n); + defer gpa.free(direct); + @memset(direct, true); + for (parents, 0..) |p, node| { + if (p < 0) continue; + const pu: usize = @intCast(p); + if (last[pu] == null or pos[node] > pos[last[pu].?]) last[pu] = node; + if (pos[node] != pos[pu] + 1) direct[pu] = false; + } + const slot_of = try gpa.alloc(?i32, n); + defer gpa.free(slot_of); + @memset(slot_of, null); + var free: std.ArrayList(i32) = .empty; + defer free.deinit(gpa); + var entries = try gpa.alloc(i32, 3 * n); + errdefer gpa.free(entries); + var used: i32 = 0; + for (order, 0..) |node, i| { + const parent = parents[node]; + var source: i32 = -1; + if (parent >= 0) { + const pu: usize = @intCast(parent); + source = if (i > 0 and order[i - 1] == pu) -2 else slot_of[pu] orelse return error.InvalidQwenTree; + // its last child reads it now: the slot may be reused + if (slot_of[pu] != null and pos[last[pu].?] == i) { + try free.append(gpa, slot_of[pu].?); + slot_of[pu] = null; + } + } + var dest: i32 = -1; + if (last[node] != null and !direct[node]) { + if (free.items.len > 0) { + std.mem.sort(i32, free.items, {}, std.sort.asc(i32)); + dest = free.orderedRemove(0); + } else { + dest = used; + used += 1; + } + slot_of[node] = dest; + } + entries[3 * i ..][0..3].* = .{ @intCast(node), source, dest }; + } + return .{ .entries = entries, .slots = @intCast(used) }; +} + +/// gdn.schedule: depth-first or level order, whichever needs fewer slots (depth-first on a tie). +pub fn schedule(gpa: std.mem.Allocator, parents: []const i32) !Schedule { + const n = parents.len; + const d = try gpa.alloc(usize, n); + defer gpa.free(d); + try depths(parents, d); + // depth-first: children in row order, as the Python stack pushes them reversed + const dfs = try gpa.alloc(usize, n); + defer gpa.free(dfs); + var stack: std.ArrayList(usize) = .empty; + defer stack.deinit(gpa); + try stack.append(gpa, 0); + var at: usize = 0; + while (stack.pop()) |node| { + dfs[at] = node; + at += 1; + var child = n; + while (child > node + 1) { + child -= 1; + if (parents[child] == @as(i32, @intCast(node))) try stack.append(gpa, child); + } + } + const level = try gpa.alloc(usize, n); + defer gpa.free(level); + for (level, 0..) |*l, i| l.* = i; + std.mem.sort(usize, level, d, struct { + fn less(ds: []const usize, a: usize, b: usize) bool { + return ds[a] < ds[b] or (ds[a] == ds[b] and a < b); + } + }.less); + const first = try orderSlots(gpa, parents, dfs); + const second = orderSlots(gpa, parents, level) catch |err| { + gpa.free(first.entries); + return err; + }; + const best, const other = if (second.slots < first.slots) .{ second, first } else .{ first, second }; + gpa.free(other.entries); + if (best.slots > max_slots) { + gpa.free(best.entries); + return error.QwenTreeTooWide; + } + return best; +} + +/// forward._conv_windows: each node's last `keep` inputs along its path, then its own row (`keep + base + row`). +pub fn convWindows(parents: []const i32, keep: usize, base: usize, out: []i32) void { + const w = keep + 1; + for (parents, 0..) |p, row| { + const dst = out[row * w ..][0..w]; + if (p < 0) { + for (0..keep) |j| dst[j] = @intCast(j); + } else { + const src = out[@as(usize, @intCast(p)) * w ..][0..w]; + @memcpy(dst[0..keep], src[1..w]); + } + dst[keep] = @intCast(keep + base + row); + } +} + +test "a chain's schedule keeps no slot and reads each predecessor" { + const gpa = std.testing.allocator; + const s = try schedule(gpa, &.{ -1, 0, 1, 2 }); + defer gpa.free(s.entries); + try std.testing.expectEqual(@as(usize, 0), s.slots); + try std.testing.expectEqualSlices(i32, &.{ 0, -1, -1, 1, -2, -1, 2, -2, -1, 3, -2, -1 }, s.entries); +} + +test "a fork saves the shared parent in a slot until its last child" { + const gpa = std.testing.allocator; + // root 0 with children 1 and 2, 1 with child 3: depth-first 0 1 3 2 needs the root kept for 2 + const s = try schedule(gpa, &.{ -1, 0, 0, 1 }); + defer gpa.free(s.entries); + try std.testing.expectEqual(@as(usize, 1), s.slots); + try std.testing.expectEqualSlices(i32, &.{ 0, -1, 0, 1, -2, -1, 3, -2, -1, 2, 0, -1 }, s.entries); +} + +test "schedules match gdn.schedule on random trees" { + const gpa = std.testing.allocator; + const cases = .{ + .{ &[_]i32{ -1, 0, 1, 2, 0, 3 }, &[_]i32{ 0, -1, 0, 1, -2, -1, 2, -2, -1, 3, -2, -1, 5, -2, -1, 4, 0, -1 }, 1 }, + .{ &[_]i32{ -1, 0, 0, 0, 0, 2, 3, 6, 3 }, &[_]i32{ 0, -1, 0, 1, -2, -1, 2, 0, -1, 5, -2, -1, 3, 0, 1, 6, -2, -1, 7, -2, -1, 8, 1, -1, 4, 0, -1 }, 2 }, + .{ &[_]i32{ -1, 0, 0, 2, 1, 0, 5, 1, 6, 4, 2, 6 }, &[_]i32{ 0, -1, 0, 1, -2, 1, 4, -2, -1, 9, -2, -1, 7, 1, -1, 2, 0, 1, 3, -2, -1, 10, 1, -1, 5, 0, -1, 6, -2, 0, 8, -2, -1, 11, 0, -1 }, 2 }, + }; + inline for (cases) |case| { + const s = try schedule(gpa, case[0]); + defer gpa.free(s.entries); + try std.testing.expectEqual(@as(usize, case[2]), s.slots); + try std.testing.expectEqualSlices(i32, case[1], s.entries); + } +} + +test "conv windows follow each node's path" { + var out: [4 * 4]i32 = undefined; + convWindows(&.{ -1, 0, 0, 1 }, 3, 0, &out); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3, 1, 2, 3, 4, 1, 2, 3, 5, 2, 3, 4, 6 }, &out); +} diff --git a/zig/src/families/qwen27/cuda/triton.zig b/zig/src/families/qwen27/cuda/triton.zig new file mode 100644 index 000000000..5206ff382 --- /dev/null +++ b/zig/src/families/qwen27/cuda/triton.zig @@ -0,0 +1,141 @@ +//! The 27B's Triton kernels from the captured cubins: each launch has the Python wrapper's grid and constexprs. + +const std = @import("std"); +const cuda = @import("cuda"); +const aot = cuda.aot; + +const p = aot.ptr; + +fn int(name: []const u8, v: usize) aot.Arg { + return aot.int(name, @intCast(v)); +} + +fn ci(name: []const u8, v: usize) aot.Const { + return aot.ci(name, @intCast(v)); +} + +fn u(x: usize) u32 { + return @intCast(x); +} + +fn cdiv(a: usize, b: usize) u32 { + return u((a + b - 1) / b); +} + +fn pow2(n: usize) usize { + return std.math.ceilPowerOfTwo(usize, n) catch unreachable; +} + +/// attention.py's layout constants: 512-key chunks, folded four to a group, up to 128 window rows. +pub const chunk = 512; +pub const group_chunks = 4; +pub const max_nodes = 128; +const merge_columns = 64; + +/// The DeltaNet shapes glue.gdn_pre takes: projection width, key and value heads, head dim, kept conv rows. +pub const Gdn = struct { c: usize, kh: usize, vh: usize, dk: usize, keep: usize }; + +/// The attention shapes: query and key heads, head dim, rotary half. +pub const Attn = struct { heads: usize, kv_heads: usize, dim: usize, half: usize }; + +pub const Tri = struct { + set: *const aot.Set, + s: cuda.Stream, + + fn run(t: Tri, name: []const u8, grid: [3]u32, args: []const aot.Arg, consts: []const aot.Const) !void { + try t.set.run(t.s, name, grid, args, consts); + } + + /// glue.embed: a 4-bit table row dequantized, one program a row and 64-group. + pub fn embed(t: Tri, ids: u64, w: u64, s: u64, b: u64, out: u64, rows: usize, d: usize) !void { + try t.run("_embed", .{ u(rows), u(d / 64), 1 }, &.{ p("IDS", "*i32", ids), p("Wt", "*i32", w), p("S", "*bf16", s), p("B", "*bf16", b), p("OUT", "*bf16", out) }, &.{ci("D", d)}); + } + + /// glue.add_rmsnorm: h = x + r (or x itself without r), y = rmsnorm(h) * w, xs = y's 64-group sums. + pub fn addRmsnorm(t: Tri, x: u64, r: ?u64, w: u64, h: u64, y: u64, xs: u64, rows: usize, d: usize, eps: f32) !void { + try t.run("_add_rmsnorm", .{ u(rows), 1, 1 }, &.{ p("X", "*bf16", x), p("R", "*bf16", r orelse x), p("W", "*bf16", w), p("H", "*bf16", if (r != null) h else x), p("Y", "*bf16", y), p("XS", "*fp32", xs), aot.float("eps", eps) }, &.{ ci("D", d), ci("BLOCK", pow2(d)), ci("HAS_R", @intFromBool(r != null)) }); + } + + /// glue.gdn_pre: conv over each row's window of [state; qkv], q/k norms, g and beta; `sid`: several streams. + pub fn gdnPre(t: Tri, qkv: u64, cs: u64, cw: u64, win: u64, sid: ?u64, a: u64, b: u64, alog: u64, dtb: u64, q: u64, k: u64, v: u64, g: u64, beta: u64, rows: usize, m: Gdn) !void { + try t.run("_gdn_pre", .{ u(rows), u(2 * m.kh + m.vh), 1 }, &.{ + p("QKV", "*bf16", qkv), p("CS", "*bf16", cs), p("CW", "*bf16", cw), p("WIN", "*i32", win), p("A", "*bf16", a), p("B", "*bf16", b), p("ALOG", "*fp32", alog), + p("DTB", "*fp32", dtb), p("Q", "*bf16", q), p("K", "*bf16", k), p("V", "*bf16", v), p("G", "*fp32", g), p("BETA", "*fp32", beta), p("SID", "*i32", sid orelse win), + }, &.{ ci("C", m.c), ci("KH", m.kh), ci("VH", m.vh), ci("DK", m.dk), ci("NKEEP", m.keep), ci("MULTI", @intFromBool(sid != null)) }); + } + + /// glue.gated_norm: silu(z) * rmsnorm(y) * w per value head, and its group sums. + pub fn gatedNorm(t: Tri, y: u64, z: u64, w: u64, out: u64, xs: u64, rows: usize, vh: usize, dv: usize, eps: f32) !void { + try t.run("_gated_norm", .{ u(rows), u(vh), 1 }, &.{ p("Yr", "*bf16", y), p("Z", "*bf16", z), p("W", "*bf16", w), p("OUT", "*bf16", out), p("XS", "*fp32", xs), aot.float("eps", eps) }, &.{ ci("VH", vh), ci("DV", dv) }); + } + + /// glue.swiglu: silu(gate) * up in 1024-column blocks, and its group sums. + pub fn swiglu(t: Tri, gate: u64, up: u64, out: u64, xs: u64, rows: usize, n: usize) !void { + try t.run("_swiglu", .{ u(rows), cdiv(n, 1024), 1 }, &.{ p("GATE", "*bf16", gate), p("UP", "*bf16", up), p("OUT", "*bf16", out), p("XS", "*fp32", xs) }, &.{ ci("N", n), ci("BLOCK", 1024) }); + } + + /// glue.attn_prep: q (from the [q | gate] rows) and k normed, then rotated at each row's position. + pub fn attnPrep(t: Tri, qg: u64, kv: u64, qn: u64, kn: u64, pos: u64, inv: u64, qout: u64, kout: u64, rows: usize, a: Attn, eps: f32) !void { + try t.run("_attn_prep", .{ u(rows), u(a.heads + a.kv_heads), 1 }, &.{ p("QG", "*bf16", qg), p("KV", "*bf16", kv), p("QN", "*bf16", qn), p("KN", "*bf16", kn), p("POS", "*i32", pos), p("INV", "*fp32", inv), p("QOUT", "*bf16", qout), p("KOUT", "*bf16", kout), aot.float("eps", eps) }, &.{ + ci("H", a.heads), ci("HKV", a.kv_heads), ci("D", a.dim), ci("HALF", a.half), ci("MROPE", 0), ci("ROWS", 0), ci("HSEC", 11), ci("WSEC", 10), + }); + } + + /// glue.gate_mul: the attention output times sigmoid(gate) from the [q | gate] rows, and its group sums. + pub fn gateMul(t: Tri, o: u64, qg: u64, out: u64, xs: u64, rows: usize, heads: usize, dim: usize) !void { + try t.run("_gate_mul", .{ u(rows), u(heads), 1 }, &.{ p("O", "*bf16", o), p("QG", "*bf16", qg), p("OUT", "*bf16", out), p("XS", "*fp32", xs) }, &.{ ci("H", heads), ci("D", dim) }); + } + + /// qmm.group_sums: each row's 64-group input sums in fp32 (rows may be strided by `ld`). + pub fn groupSums(t: Tri, x: u64, xs: u64, ld: usize, rows: usize, k: usize) !void { + try t.run("_group_sums", .{ u(rows), cdiv(k / 64, 16), 1 }, &.{ p("X", "*bf16", x), p("XS", "*fp32", xs), int("ldx", ld) }, &.{ ci("KG", k / 64), ci("GS", 64), ci("GB", 16) }); + } + + /// attention.from_packed's _paths: each row's root-to-row window rows and its depth. + pub fn paths(t: Tri, parents: u64, out: u64, depths: u64, rows: usize) !void { + try t.run("_paths", .{ u(rows), 1, 1 }, &.{ p("PARENTS", "*i32", parents), p("PATHS", "*i32", out), p("DEPTHS", "*i32", depths) }, &.{ci("MAXD", max_nodes)}); + } + + /// The layout attention.attention reads: the plan's tables and the partial buffers it writes. + pub const Plan = struct { streams: u64, items: u64, n_items: usize, rows: u64, paths: u64, depths: u64, width: usize }; + pub const Partials = struct { o: u64, m: u64, l: u64 }; + + /// attention.attention: committed keys in 512-key chunks, the row's own path in its tail, merged in order. + pub fn attention(t: Tri, q: u64, kn: u64, vn: u64, base: u64, offs: u64, plan: Plan, part: Partials, out: u64, a: Attn) !void { + const g = a.heads / a.kv_heads; + const scale: f32 = @floatCast(std.math.pow(f64, @floatFromInt(a.dim), -0.5)); + const w = plan.width; + if (plan.n_items > 0) try t.run("_shared", .{ u(plan.n_items * a.kv_heads), 1, 1 }, &.{ + p("Q", "*bf16", q), p("KC", "*bf16", base), p("VC", "*bf16", base), p("OFF", "*i64", offs), p("STREAM", "*i32", plan.streams), + p("ITEMS", "*i32", plan.items), p("PO", "*fp32", part.o), p("PM", "*fp32", part.m), p("PL", "*fp32", part.l), int("W", w), + }, &.{ ci("H", a.heads), ci("HK", a.kv_heads), ci("D", a.dim), ci("G", g), ci("CH", chunk), aot.cf("SCALE", scale), ci("GR", group_chunks) }); + const tails = 1 + (max_nodes + chunk - 1) / chunk; + try t.run("_tail", .{ u(w), u(a.kv_heads), u(tails) }, &.{ + p("Q", "*bf16", q), p("KN", "*bf16", kn), p("VN", "*bf16", vn), p("KC", "*bf16", base), p("VC", "*bf16", base), p("OFF", "*i64", offs), p("STREAM", "*i32", plan.streams), + p("ROWS", "*i32", plan.rows), p("PATHS", "*i32", plan.paths), p("DEPTHS", "*i32", plan.depths), p("PO", "*fp32", part.o), p("PM", "*fp32", part.m), p("PL", "*fp32", part.l), int("W", w), + }, &.{ ci("H", a.heads), ci("HK", a.kv_heads), ci("D", a.dim), ci("G", g), ci("CH", chunk), ci("MAXD", max_nodes), aot.cf("SCALE", scale), ci("GR", group_chunks) }); + try t.run("_merge", .{ u(w), u(a.kv_heads), u(a.dim / merge_columns) }, &.{ + p("PO", "*fp32", part.o), p("PM", "*fp32", part.m), p("PL", "*fp32", part.l), p("OUT", "*bf16", out), p("STREAM", "*i32", plan.streams), p("ROWS", "*i32", plan.rows), int("W", w), + }, &.{ ci("H", a.heads), ci("D", a.dim), ci("G", g), ci("DS", merge_columns), ci("CH", chunk), ci("GR", group_chunks) }); + } +}; + +/// attention.groups: whole committed 2048-key groups fold in their own programs only when they fill the GPU. +pub fn groups(p_: usize, w: usize) usize { + const g = p_ / (chunk * group_chunks); + return if (g * w >= 64) g else 0; +} + +/// attention.slots: one partial slot per folded group, then one per chunk through the window's last key. +pub fn slots(p_: usize, w: usize) usize { + const g = groups(p_, w); + return g + (p_ + w + chunk - 1) / chunk - g * group_chunks; +} + +test "attention slots follow the Python plan" { + try std.testing.expectEqual(@as(usize, 1), slots(0, 1)); + try std.testing.expectEqual(@as(usize, 2), slots(512, 1)); + try std.testing.expectEqual(@as(usize, 9), slots(4096, 1)); // 8 chunks + the window's own + try std.testing.expectEqual(@as(usize, 13), slots(6144, 16)); // 3 groups x 16 rows < 64: every chunk its own + try std.testing.expectEqual(@as(usize, 4), slots(6144, 32)); // 3 folded groups, then the window's chunk +} diff --git a/zig/src/families/qwen27/cuda/weights.zig b/zig/src/families/qwen27/cuda/weights.zig new file mode 100644 index 000000000..5c796bd43 --- /dev/null +++ b/zig/src/families/qwen27/cuda/weights.zig @@ -0,0 +1,252 @@ +//! The 27B's weights on the GPU in the Python engine's layouts: tiled projections, the stored embedding, fp32 gates. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const source = @import("nemotron").source; +const c = @import("shape.zig"); +const config_eos = @import("../config.zig").max_eos; +const kern = @import("kernels.zig"); + +pub const QLinear = kern.QLinear; +const Tensor = core.checkpoint.Tensor; +const prefix = "language_model."; + +/// The token table as stored (MLX words, scales, biases): a row lookup, not a matmul. +pub const Embed = struct { w: u64, s: u64, b: u64 }; +pub const Delta = struct { qkv: QLinear, z: QLinear, b: QLinear, a: QLinear, out: QLinear, conv: u64, a_log: u64, dt_bias: u64, norm: u64 }; +pub const Attention = struct { q: QLinear, k: QLinear, v: QLinear, o: QLinear, q_norm: u64, k_norm: u64 }; +pub const Mixer = union(enum) { delta: Delta, attention: Attention }; +pub const Layer = struct { input_norm: u64, post_norm: u64, mixer: Mixer, gate: QLinear, up: QLinear, down: QLinear }; + +pub const Weights = struct { + gpa: std.mem.Allocator, + g: c.Geometry, + embed: Embed = undefined, + layers: []Layer = &.{}, + norm: u64 = 0, + head: QLinear = undefined, + inv_freq: u64 = 0, // (rotary_dim / 2,) fp32 + eos: [config_eos]u32 = undefined, // config.json's end tokens, then generation_config.json's + eos_count: usize = 0, + buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.layers); + w.* = undefined; + } + + pub fn isEos(w: *const Weights, token: u32) bool { + for (w.eos[0..w.eos_count]) |e| if (e == token) return true; + return false; + } + + /// Device bytes the weights hold. + pub fn bytes(w: *const Weights) usize { + var n: usize = 0; + for (w.buffers.items) |b| n += b.len; + return n; + } +}; + +const Loader = struct { + gpa: std.mem.Allocator, + ops: kern.Ops, + w: *Weights, + src: *source.Source, + scratch: cuda.DeviceBuffer, + host: std.ArrayList(u8) = .empty, + + fn alloc(L: *Loader, len: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(L.ops.k.d, len); + try L.w.buffers.append(L.gpa, b); + return b.ptr; + } + + /// Device scratch of at least `len` bytes; growing waits for the stream so pending packs keep their inputs. + fn tmp(L: *Loader, len: usize) !u64 { + if (L.scratch.len < len) { + try L.src.flush(); + try L.ops.s.synchronize(); + L.scratch.free(); + L.scratch = try cuda.DeviceBuffer.alloc(L.ops.k.d, len); + } + return L.scratch.ptr; + } + + /// The host staging buffer, once the stream has read its last contents. + fn staging(L: *Loader, len: usize) ![]u8 { + try L.src.flush(); + try L.ops.s.synchronize(); + try L.host.resize(L.gpa, len); + return L.host.items; + } + + fn raw(L: *Loader, t: Tensor) !u64 { + const ptr = try L.alloc(t.bytes.len); + try L.src.upload(ptr, t.bytes); + return ptr; + } + + /// `.float()` of a bf16 tensor (or an f32 one as stored): exact. + fn widen(L: *Loader, t: Tensor) !u64 { + if (t.dtype == .f32) return L.raw(t); + if (t.dtype != .bf16) return error.UnexpectedTensor; + const n = t.bytes.len / 2; + const out = std.mem.bytesAsSlice(u32, try L.staging(n * 4)); + const in = try L.src.view(t.bytes); + for (out, 0..) |*o, i| o.* = @as(u32, std.mem.readInt(u16, in[2 * i ..][0..2], .little)) << 16; + const ptr = try L.alloc(n * 4); + try L.ops.upload(ptr, L.host.items); + return ptr; + } + + /// qmm_fast.tile: MLX words, scales and biases packed into the lane matmul's tiled layout. + fn tile(L: *Loader, ck: *core.Checkpoint, name: []const u8, n: usize, k: usize) !QLinear { + var parts: [3]Tensor = undefined; + var buf: [192]u8 = undefined; + for ([_][]const u8{ "weight", "scales", "biases" }, 0..) |s, i| parts[i] = try ck.get(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, s })); + const k8 = k / 8; + const kg = k / 64; + if (parts[0].dtype != .u32 or parts[1].dtype != .bf16 or parts[2].dtype != .bf16) return error.UnsupportedQuantization; + if (!parts[0].is(.u32, &.{ n, k8 }) or !parts[1].is(.bf16, &.{ n, kg }) or !parts[2].is(.bf16, &.{ n, kg })) { + std.log.err("{s}: not a 4-bit group-64 ({d}, {d}) projection", .{ name, n, k }); + return error.UnexpectedTensor; + } + const npad = (n + 127) / 128 * 128; + const sizes = [3]usize{ n * k8 * 4, n * kg * 2, n * kg * 2 }; + const base = try L.tmp(sizes[0] + sizes[1] + sizes[2]); + const at = [3]usize{ 0, sizes[0], sizes[0] + sizes[1] }; + for (0..3) |j| try L.src.upload(base + at[j], parts[j].bytes); + const q: QLinear = .{ .w = try L.alloc(npad * k / 2), .s = try L.alloc(kg * npad * 2), .b = try L.alloc(kg * npad * 2), .n = n, .k = k, .npad = npad }; + const total: u64 = @as(u64, npad / 64) * kg * 512; + try L.src.flush(); + var a: cuda.Args = .{}; + a.add(base); + for ([_]usize{ n, k8, kg }) |v| a.add(@as(c_int, @intCast(v))); + a.add(q.w); + a.add(@as(c_longlong, @intCast(total))); + try cuda.launch.launch(L.ops.k.pack_dense, .{ .grid = .{ .x = @intCast((total + 255) / 256) }, .block = .{ .x = 256 } }, L.ops.s, &a); + for ([_]u64{ q.s, q.b }, [_]usize{ at[1], at[2] }) |out, off| { + var t: cuda.Args = .{}; + t.add(base + off); + for ([_]usize{ n, kg, npad }) |v| t.add(@as(c_int, @intCast(v))); + t.add(out); + const cells = @as(u64, kg) * npad; + try cuda.launch.launch(L.ops.k.transpose16, .{ .grid = .{ .x = @intCast((cells + 255) / 256) }, .block = .{ .x = 256 } }, L.ops.s, &t); + } + return q; + } + + fn vector(L: *Loader, ck: *core.Checkpoint, name: []const u8, len: usize) !u64 { + var buf: [192]u8 = undefined; + return L.raw(try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, name }), .bf16, &.{len})); + } + + fn delta(L: *Loader, ck: *core.Checkpoint, p: []const u8, g: c.Geometry) !Delta { + var b1: [160]u8 = undefined; + var b2: [192]u8 = undefined; + const at = struct { + fn f(buf: []u8, pre: []const u8, name: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}linear_attn.{s}", .{ pre, name }); + } + }.f; + var d: Delta = undefined; + d.qkv = try L.tile(ck, try at(&b1, p, "in_proj_qkv"), g.convDim(), g.hidden); + d.z = try L.tile(ck, try at(&b1, p, "in_proj_z"), g.vInner(), g.hidden); + d.b = try L.tile(ck, try at(&b1, p, "in_proj_b"), g.linear_v_heads, g.hidden); + d.a = try L.tile(ck, try at(&b1, p, "in_proj_a"), g.linear_v_heads, g.hidden); + d.out = try L.tile(ck, try at(&b1, p, "out_proj"), g.hidden, g.vInner()); + // conv1d.weight (C, 4, 1) is (C, 4) as stored: the reshape moves no bytes + d.conv = try L.raw(try ck.expect(try std.fmt.bufPrint(&b2, "{s}{s}", .{ prefix, try at(&b1, p, "conv1d.weight") }), .bf16, &.{ g.convDim(), c.conv_taps, 1 })); + d.a_log = try L.widen(try ck.get(try std.fmt.bufPrint(&b2, "{s}{s}", .{ prefix, try at(&b1, p, "A_log") }))); + d.dt_bias = try L.widen(try ck.get(try std.fmt.bufPrint(&b2, "{s}{s}", .{ prefix, try at(&b1, p, "dt_bias") }))); + d.norm = try L.vector(ck, try at(&b1, p, "norm.weight"), c.linear_dim); + return d; + } + + fn attention(L: *Loader, ck: *core.Checkpoint, p: []const u8, g: c.Geometry) !Attention { + var b1: [160]u8 = undefined; + const at = struct { + fn f(buf: []u8, pre: []const u8, name: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}self_attn.{s}", .{ pre, name }); + } + }.f; + return .{ + .q = try L.tile(ck, try at(&b1, p, "q_proj"), 2 * g.qInner(), g.hidden), + .k = try L.tile(ck, try at(&b1, p, "k_proj"), g.kvInner(), g.hidden), + .v = try L.tile(ck, try at(&b1, p, "v_proj"), g.kvInner(), g.hidden), + .o = try L.tile(ck, try at(&b1, p, "o_proj"), g.hidden, g.qInner()), + .q_norm = try L.vector(ck, try at(&b1, p, "q_norm.weight"), c.head_dim), + .k_norm = try L.vector(ck, try at(&b1, p, "k_norm.weight"), c.head_dim), + }; + } +}; + +/// weights.inv_freq: theta ** (-i / half) in float64, rounded to fp32. +pub fn invFreq(out: []f32) void { + const half: f64 = @floatFromInt(out.len); + for (out, 0..) |*f, i| f.* = @floatCast(std.math.pow(f64, c.theta, -@as(f64, @floatFromInt(i)) / half)); +} + +/// weights.load (tiled, one GPU): the model folder's text tensors on the GPU, leftovers refused. +pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, cfg: c.Config) !Weights { + const g = cfg.g; + var w: Weights = .{ .gpa = gpa, .g = g, .eos = cfg.eos, .eos_count = cfg.eos_count }; + errdefer w.deinit(); + var src = try source.Source.init(gpa, .{ .d = ops.k.d, .s = ops.s, .discrete = ops.k.discrete }); + defer src.deinit(); + var L: Loader = .{ .gpa = gpa, .ops = ops, .w = &w, .src = &src, .scratch = try cuda.DeviceBuffer.alloc(ops.k.d, 1 << 20) }; + defer { + src.flush() catch {}; + ops.s.synchronize() catch {}; + L.scratch.free(); + L.host.deinit(gpa); + } + // the vision tower and the MTP head stay on disk, as weights._Tensors skips them + var ck = try core.Checkpoint.openModelPrefix(gpa, io, dir, prefix); + defer ck.close(); + try src.add(&ck); + w.layers = try gpa.alloc(Layer, g.layers); + for (w.layers, 0..) |*layer, i| { + var pb: [64]u8 = undefined; + var nb: [160]u8 = undefined; + const p = try std.fmt.bufPrint(&pb, "model.layers.{d}.", .{i}); + const join = struct { + fn f(buf: []u8, a: []const u8, b: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}{s}", .{ a, b }); + } + }.f; + layer.input_norm = try L.vector(&ck, try join(&nb, p, "input_layernorm.weight"), g.hidden); + layer.mixer = if (c.linear(i)) .{ .delta = try L.delta(&ck, p, g) } else .{ .attention = try L.attention(&ck, p, g) }; + layer.gate = try L.tile(&ck, try join(&nb, p, "mlp.gate_proj"), g.intermediate, g.hidden); + layer.up = try L.tile(&ck, try join(&nb, p, "mlp.up_proj"), g.intermediate, g.hidden); + layer.down = try L.tile(&ck, try join(&nb, p, "mlp.down_proj"), g.hidden, g.intermediate); + layer.post_norm = try L.vector(&ck, try join(&nb, p, "post_attention_layernorm.weight"), g.hidden); + } + w.embed = .{ + .w = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.weight", .u32, &.{ c.vocab, g.hidden / 8 })), + .s = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.scales", .bf16, &.{ c.vocab, g.hidden / 64 })), + .b = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.biases", .bf16, &.{ c.vocab, g.hidden / 64 })), + }; + w.norm = try L.vector(&ck, "model.norm.weight", g.hidden); + w.head = try L.tile(&ck, "lm_head", c.vocab, g.hidden); + if (ck.unused() != 0) return error.UnusedCheckpointTensors; + var inv: [c.rotary_dim / 2]f32 = undefined; + invFreq(&inv); + w.inv_freq = try L.alloc(@sizeOf(@TypeOf(inv))); + try ops.upload(w.inv_freq, std.mem.asBytes(&inv)); + try src.flush(); + try ops.s.synchronize(); + return w; +} + +test "inv_freq: theta ** (-i / 32) rounded to fp32" { + var inv: [32]f32 = undefined; + invFreq(&inv); + try std.testing.expectEqual(@as(f32, 1.0), inv[0]); + try std.testing.expectEqual(@as(f32, 0.6042963862419128), inv[1]); + try std.testing.expectEqual(@as(f32, 0.22067341208457947), inv[3]); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 668bd82c4..032f64c60 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -4,13 +4,14 @@ const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const qwen27 = @import("qwen27"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const nemotron_slide = @import("nemotron_slide.zig"); const Pool = budget.Pool; /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{nemotron.native}; +const registry = .{ nemotron.native, qwen27.native }; /// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ @@ -248,6 +249,12 @@ const Host = struct { return x.inner.vtable.prefill(x.inner.ptr, s); } }.f, + .prefill_many = if (v.prefill_many != null) struct { + fn f(p: *anyopaque, streams: []const *lanes.Stream, errs: []?anyerror) anyerror!void { + const x = bind(p); + return x.inner.vtable.prefill_many.?(x.inner.ptr, streams, errs); + } + }.f else null, .first = struct { fn f(p: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const x = bind(p); @@ -404,7 +411,10 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); return null; } - const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments, .slide = o.slide }) catch |e| { + var family_options: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments }; + if (@hasField(F.Options, "slide")) family_options.slide = o.slide; + if (@hasField(F.Options, "drafter")) family_options.drafter = o.drafter; + const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, family_options) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} with {s} glue kernels ({s})", .{ o.dir, kernels orelse "its own", @errorName(e) }); return null; }; @@ -423,13 +433,14 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); + // prompt-cache snapshots are Nemotron's (cuda_reuse.zig); a family without them serves without a cache + const cache = if (@hasDecl(F, "snap_bytes")) budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)) else 0; h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in {d}-row chunks{s}{s}; {s}", .{ capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (@hasDecl(F, "snap_bytes") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); // states kept and resumed on the family's grid, where a cache-off pass cuts; the head reads one token on @@ -440,7 +451,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O store.deinit(); gpa.destroy(store); }; - if (loaded.learns) |e| { + if (@hasDecl(F, "slide")) if (loaded.learns) |e| { const Slide = nemotron_slide.Adapter(F.slide); const s = try gpa.create(Slide); s.* = Slide.init(gpa, io, e); @@ -453,7 +464,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O } }.f; h.learn = s.hook(); - } + }; errdefer if (h.slide) |s| h.slide_free.?(gpa, s); // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .logprobs = loaded.backend.vtable.first_row != null }, .{ .ctx = loaded.ctx, .text = F.explain }); diff --git a/zig/tests/cuda/qwen27/capture.py b/zig/tests/cuda/qwen27/capture.py new file mode 100644 index 000000000..a21dd828e --- /dev/null +++ b/zig/tests/cuda/qwen27/capture.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +"""The 27B's Python CUDA engine as the Zig engine's oracle: Triton launches, per-layer dumps, tokens.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + +TEXTS = { + "story": "Write a short story about a lighthouse keeper who finds a message in a bottle.", + "code": "Write a Python function that merges two sorted lists, with a docstring and three tests.", + "facts": "Explain how a refrigerator moves heat out of its cabinet, step by step.", +} +LONG_BODY = 5000 # tokens of module source in the long prompt: two 4096-row prompt chunks +TEACHER = 32 # teacher-forced decode steps (one row each, from an empty state) +DUMP_STEPS = (0, 1, TEACHER - 1) + + +def chat(user: str) -> str: + return f"<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + + +def prompts(model: Path) -> dict[str, list[int]]: + from tokenizers import Tokenizer + import textwrap + + tok = Tokenizer.from_file(str(model / "tokenizer.json")) + out = {name: tok.encode(chat(text), add_special_tokens=False).ids for name, text in TEXTS.items()} + words = Path(textwrap.__file__).read_text() + body = tok.encode(words * (1 + LONG_BODY // max(1, len(words) // 3)), add_special_tokens=False).ids[:LONG_BODY] + head = tok.encode("<|im_start|>user\nReview this module and list its public functions:\n\n", + add_special_tokens=False).ids + tail = tok.encode("<|im_end|>\n<|im_start|>assistant\n\n\n\n\n", add_special_tokens=False).ids + out["long"] = head + body + tail + return out + + +def sha12(tokens: list[int]) -> str: + return hashlib.sha256(json.dumps([int(t) for t in tokens]).encode()).hexdigest()[:12] + + +def digest(t) -> str: + import torch + + raw = t.detach().contiguous() + return hashlib.sha256(raw.view(torch.uint8).cpu().numpy().tobytes()).hexdigest() + + +class Dumps: + """Raw little-endian tensors by name under the current directory; None turns dumping off.""" + + def __init__(self) -> None: + self.dir: Path | None = None + self.i = 0 + + def at(self, d: Path | None) -> None: + self.dir, self.i = d, 0 + if d is not None: + d.mkdir(parents=True, exist_ok=True) + + def save(self, name: str, t) -> None: + import torch + + if self.dir is None: + return + torch.cuda.synchronize() + t.detach().contiguous().view(torch.uint8).cpu().numpy().tofile(self.dir / f"{name}.bin") + + +def hook(dumps: Dumps) -> None: + """Every add_rmsnorm call's (h, y, xs), in call order: two a layer, then the final norm.""" + + from tensorfold.families.qwen3_5.cuda import glue + + norm = glue.add_rmsnorm + + def traced(x, r, w, eps): + h, y, xs = norm(x, r, w, eps) + if dumps.dir is not None: + i = dumps.i + dumps.i += 1 + dumps.save(f"{i:03d}_h", h) + dumps.save(f"{i:03d}_y", y) + dumps.save(f"{i:03d}_xs", xs) + return h, y, xs + + glue.add_rmsnorm = traced + + +def specialization() -> dict: + """Each launched JIT function's parameter names and which ones Triton never specializes.""" + + from triton.runtime.jit import JITFunction + import gc + + out = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + out[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + return out + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", default="", help="folder of triton_aot_manifest.py: record every Triton launch") + ap.add_argument("--max-tokens", type=int, default=128) + ap.add_argument("--repeats", type=int, default=2) + ap.add_argument("--context", type=int, default=16384) + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + rec = None + if a.tools: + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + + import torch + from contextlib import nullcontext + + def phase(name: str, detail: bool = False): + return rec.scope(name, detail) if rec is not None else nullcontext() + + model = Path(a.model) + ids = prompts(model) + (out / "prompts.json").write_text(json.dumps(ids) + "\n") + dumps = Dumps() + hook(dumps) + from tensorfold.families.qwen3_5.cuda.engine import Qwen27Engine + + t0 = time.perf_counter() + with phase("startup"): + eng = Qwen27Engine(model, None, streams=1, context=a.context, context_explicit=True) + if rec is not None: + from tensorfold.cuda.kernels import gdn, prefill_attention, qmm + + rec.wrap(qmm._ext(), ("qmm", "qmm_group", "qmm_prefill"), "qmm") + rec.wrap(gdn._ext(), ("tree", "replay", "prefill"), "gdn") + rec.wrap(prefill_attention._ext(), ("prefill_attention",), "prefill_attention") + info = {"startup_s": round(time.perf_counter() - t0, 2), "context": eng.context_window, "eos": list(eng.eos), + "prompt_rows": eng.w.prompt_rows, "torch": torch.__version__, "gpu": torch.cuda.get_device_name(0)} + print(json.dumps(info), flush=True) + + results: dict[str, dict] = {} + for name, prompt in ids.items(): + runs = [] + for _ in range(a.repeats): + toks: list[int] = [] + with phase(f"serial-{name}"): + torch.cuda.synchronize() + start = time.perf_counter() + stats = eng.generate(prompt, a.max_tokens, None, lambda new: toks.extend(map(int, new)) and False, + draft=False, stop_eos=True) + torch.cuda.synchronize() + wall = time.perf_counter() - start + runs.append({"tokens": toks, "sha": sha12(toks), "wall_s": round(wall, 4), **stats}) + last = runs[-1] + same = all(r["tokens"] == last["tokens"] for r in runs) + step = last.get("decode_s", 0.0) / max(1, len(last["tokens"]) - 1) * 1e3 + results[name] = {**last, "repeats_same": same, "ms_per_token": round(step, 4)} + print(f"serial {name}: {len(last['tokens'])} tokens sha {last['sha']} prefill {last.get('prefill_s'):.3f}s " + f"{step:.3f} ms/token same {same}", flush=True) + + from tensorfold.families.qwen3_5.cuda.decode import prefill + from tensorfold.families.qwen3_5.cuda.forward import State, commit, tree_forward + + w = eng.w + tf = ids["code"][:TEACHER] + sampled, logit_sha = [], [] + st = State(w) + with phase("teacher", detail=True): + for i, t in enumerate(tf): + dumps.at(out / "teacher" / f"step{i:03d}" if i in DUMP_STEPS else None) + logits, record = tree_forward(w, torch.tensor([t], dtype=torch.int32, device="cuda"), [-1], st) + dumps.save("logits", logits) + sampled.append(int(logits.argmax(dim=-1)[0])) + logit_sha.append(digest(logits)) + commit(st, record, [0]) + dumps.at(None) + pre_out = {} + for name in ("code", "long"): + dumps.at(out / "prefill" / name if name == "code" else None) + with phase(f"prefill-{name}", detail=name == "code"): + _, pending = prefill(w, ids[name], None, limit=a.context) + dumps.at(None) + pre_out[name] = {"pending": pending} + (out / "teacher.json").write_text(json.dumps({"tokens": tf, "sampled": sampled, "logits_sha256": logit_sha, + "dump_steps": list(DUMP_STEPS), "prefill": pre_out}) + "\n") + (out / "results.json").write_text(json.dumps({"info": info, "results": results}, indent=1) + "\n") + if rec is not None: + rec.dump(out / "launches.json") + (out / "jit.json").write_text(json.dumps(specialization(), indent=1, sort_keys=True) + "\n") + print("done", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/qwen27/capture_draft.py b/zig/tests/cuda/qwen27/capture_draft.py new file mode 100644 index 000000000..b5aa457eb --- /dev/null +++ b/zig/tests/cuda/qwen27/capture_draft.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""The 27B's DFlash2 drafting as the Zig engine's oracle: drafted replies (equal to serial), rounds, Triton launches.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + + +def sha12(tokens: list[int]) -> str: + return hashlib.sha256(json.dumps([int(t) for t in tokens]).encode()).hexdigest()[:12] + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", required=True) + ap.add_argument("--draft", required=True) + ap.add_argument("--prompts", required=True, help="capture.py's prompts.json") + ap.add_argument("--out", required=True) + ap.add_argument("--tools", default="") + ap.add_argument("--max-tokens", type=int, default=128) + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + rec = None + if a.tools: + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + import torch + + from tensorfold.families.qwen3_5.cuda.decode import draft_decode, prefill + from tensorfold.families.qwen3_5.cuda.engine import Qwen27Engine + + eng = Qwen27Engine(Path(a.model), Path(a.draft), streams=1, context=16384, context_explicit=True) + if rec is not None: + from tensorfold.cuda.kernels import gdn, prefill_attention, qmm + + rec.wrap(qmm._ext(), ("qmm", "qmm_group", "qmm_prefill"), "qmm") + rec.wrap(gdn._ext(), ("tree", "replay", "prefill"), "gdn") + rec.wrap(prefill_attention._ext(), ("prefill_attention",), "prefill_attention") + prompts = json.loads(Path(a.prompts).read_text()) + results = {} + for name, ids in prompts.items(): + for _ in range(2): # the second run is the timed one + drafter = eng.draft + drafter.restore(([None] * drafter.layers, [None] * drafter.layers, 0, 0)) + torch.cuda.synchronize() + t0 = time.perf_counter() + st, pending = prefill(eng.w, ids, None, drafter, limit=16384) + trace: list = [] + r = draft_decode(eng.w, st, ids, pending, a.max_tokens, None, drafter, max_rows=eng.max_rows, + tree_rows=eng.tree_rows, allow_copy=True, stop_eos=True, trace=trace, inplace=True) + torch.cuda.synchronize() + wall = time.perf_counter() - t0 + results[name] = {"tokens": r.tokens, "sha": sha12(r.tokens), "rounds": r.rounds, "drafted": r.drafted_rows, + "accepted": r.accepted_drafts, "decode_s": round(r.seconds, 4), "wall_s": round(wall, 4), + "widths": r.widths, "paths": [len(t.get("path", [])) for t in trace]} + print(f"{name}: {len(r.tokens)} tokens sha {sha12(r.tokens)} rounds {r.rounds} accepted {r.accepted_drafts} " + f"{len(r.tokens) / r.seconds:.1f} tok/s", flush=True) + (out / "draft.json").write_text(json.dumps({"max_rows": eng.max_rows, "results": results}) + "\n") + if rec is not None: + rec.dump(out / "launches.json") + print("done", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/qwen27/capture_multi.py b/zig/tests/cuda/qwen27/capture_multi.py new file mode 100644 index 000000000..6d1823207 --- /dev/null +++ b/zig/tests/cuda/qwen27/capture_multi.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""The 27B's drafted shared rounds as the Zig engine's oracle: MultiDecoder's windows each round, and every reply.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +from pathlib import Path + +TOPICS = ["lighthouses", "volcanoes", "bread", "rivers", "chess", "glaciers", "owls", "trains"] +CODE = '''def fibonacci(n): + """Return the nth Fibonacci number.""" + if n < 2: + return n + a, b = 0, 1 + for _ in range(n - 1): + a, b = b, a + b + return b +''' + + +def chat(user: str) -> str: + return f"<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + + +def sha12(tokens: list[int]) -> str: + return hashlib.sha256(json.dumps([int(t) for t in tokens]).encode()).hexdigest()[:12] + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", required=True) + ap.add_argument("--draft", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", default="") + ap.add_argument("--streams", type=int, default=6) + ap.add_argument("--max-tokens", type=int, default=96) + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + rec = None + if a.tools: + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + from tokenizers import Tokenizer + + from tensorfold.cuda.streams import Stream + from tensorfold.families.qwen3_5.cuda.engine import Qwen27Engine + from tensorfold.families.qwen3_5.cuda.multi import MultiDecoder + + model = Path(a.model) + tok = Tokenizer.from_file(str(model / "tokenizer.json")) + users = [f"Rewrite this function with a docstring that explains each step, keeping the code the same:\n\n{CODE}"] + users += [f"Write four sentences about {t}." for t in TOPICS] + prompts = [tok.encode(chat(u), add_special_tokens=False).ids for u in users[:a.streams]] + eng = Qwen27Engine(model, Path(a.draft), streams=1, context=8192, context_explicit=True) + if rec is not None: + from tensorfold.cuda.kernels import gdn, prefill_attention, qmm + + rec.wrap(qmm._ext(), ("qmm", "qmm_group", "qmm_prefill"), "qmm") + rec.wrap(gdn._ext(), ("tree", "replay", "prefill"), "gdn") + rec.wrap(prefill_attention._ext(), ("prefill_attention",), "prefill_attention") + # GB10's planning: no measured costs (every tree node verified), the drafter's block 16 rows a stream + md = MultiDecoder(eng.w, eng.draft, allow_copy=True, context=8192, keep=1) + md.depth = False + rounds: list[list[dict]] = [] + windows = md._windows + + def logged(plan, copied, blocks): + wins = windows(plan, copied, blocks) + rounds.append([{"stream": sid, "mode": ["copy", "tree", "one"][mode], "tokens": t, "parents": p} + for (sid, mode, _, _), (t, p) in zip(plan, wins)]) + return wins + + md._windows = logged + scores: list[list[float]] = [] # each tree's path scores, in finish order + finish = md.draft.finish_tree + + def scored(*args, **kw): + t = finish(*args, **kw) + scores.append([float(x) for x in t[2]]) + return t + + md.draft.finish_tree = scored + streams = [Stream(list(p), a.max_tokens, None, draft=True, stop_eos=True) for p in prompts] + for s in streams: + md.admit(s) + while any(not s.done for s in streams): + md.finish(md.round()) + results = [{"tokens": s.out, "sha": sha12(s.out), "rounds": s.rounds, "drafted": s.drafted, + "accepted": s.accepted} for s in streams] + for k, r in enumerate(results): + print(f"stream {k}: {len(r['tokens'])} tokens sha {r['sha']} rounds {r['rounds']} accepted {r['accepted']}", + flush=True) + (out / "multi.json").write_text(json.dumps({"prompts": prompts, "max_tokens": a.max_tokens, "results": results, + "rounds": rounds, "scores": scores}) + "\n") + if rec is not None: + rec.dump(out / "launches.json") + print(f"{len(rounds)} rounds", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/qwen27/capture_windows.py b/zig/tests/cuda/qwen27/capture_windows.py new file mode 100644 index 000000000..78b755895 --- /dev/null +++ b/zig/tests/cuda/qwen27/capture_windows.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +"""The 27B's verify windows as the Zig engine's oracle: chains of 1-16 rows and shared forwards, each row hashed.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +from pathlib import Path + +TOPICS = ["lighthouses", "volcanoes", "bread", "rivers", "chess", "glaciers", "owls", "trains", + "deserts", "bees", "clocks", "ferns", "kites", "salt", "comets", "violins"] +REPLY = 24 # serial tokens each stream decodes first: the chain windows' tokens + + +def chat(user: str) -> str: + return f"<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + + +def sha(t) -> str: + import torch + + return hashlib.sha256(t.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()).hexdigest() + + +def rows(logits) -> list[dict]: + return [{"sha": sha(logits[r]), "argmax": int(logits[r].argmax())} for r in range(logits.shape[0])] + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", default="") + ap.add_argument("--wide", action="store_true", help="also rounds past 16 rows") + ap.add_argument("--sweep", action="store_true", help="also every alignment of a shared round's packed inputs") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + rec = None + if a.tools: + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + import torch + from tokenizers import Tokenizer + + from tensorfold.families.qwen3_5.cuda.decode import clone_state, prefill, serial_decode + from tensorfold.families.qwen3_5.cuda.engine import Qwen27Engine + from tensorfold.families.qwen3_5.cuda.forward import commit_streams, multi_tree_forward, tree_forward + + model = Path(a.model) + tok = Tokenizer.from_file(str(model / "tokenizer.json")) + prompts = [tok.encode(chat(f"Write four sentences about {t}."), add_special_tokens=False).ids for t in TOPICS] + eng = Qwen27Engine(model, None, streams=1, context=4096, context_explicit=True) + if rec is not None: + from tensorfold.cuda.kernels import gdn, prefill_attention, qmm + + rec.wrap(qmm._ext(), ("qmm", "qmm_group", "qmm_prefill"), "qmm") + rec.wrap(gdn._ext(), ("tree", "replay", "prefill"), "gdn") + rec.wrap(prefill_attention._ext(), ("prefill_attention",), "prefill_attention") + w = eng.w + states, replies = [], [] + for p in prompts: + st, pending = prefill(w, p, None, limit=4096) + states.append(st) + replies.append(serial_decode(w, st, pending, REPLY, None, stop_eos=False).tokens) + dev = w.norm.device + cases = [] + + def tokens(ids): + return torch.tensor(ids, dtype=torch.int32, device=dev) + + # chains of one stream: the serial reply's first W tokens, each row's draw the next serial token + for width in range(1, 17): + st = clone_state(states[0]) + logits, _ = tree_forward(w, tokens(replies[0][:width]), list(range(-1, width - 1)), st) + cases.append({"kind": "chain", "streams": [{"stream": 0, "tokens": replies[0][:width]}], "rows": rows(logits)}) + # shared forwards: each stream its own chain from its own state, rows in stream order + layouts = [[1, 1], [1] * 4, [1] * 8, [1] * 16, [4, 4, 4, 4], [2, 5, 1, 8], [3] * 5, [16]] + if a.wide: # rounds past 16 rows: the wider matmul tiles (to 32, 64, more) + layouts += [[16, 16], [5, 16, 3], [16] * 4, [9, 16, 2, 16, 7, 16, 1, 5], [12] * 8, [16] * 16] + if a.sweep: # the packed copy's pointer alignments: rows and streams mod 4, 16-row multiples + for n in range(2, 6): + for total in (16 * n, 16 * n - 12, 16 * n - 11, 16 * n - 10, 16 * n - 9): + layouts.append([total // n + (i < total % n) for i in range(n)]) + for layout in layouts: + sts = [clone_state(states[s]) for s in range(len(layout))] + wins = [(replies[s][:n], list(range(-1, n - 1)), sts[s]) for s, n in enumerate(layout)] + logits, record, _, starts = multi_tree_forward(w, wins) + case = {"kind": "shared", "streams": [{"stream": s, "tokens": replies[s][:n]} for s, n in enumerate(layout)], + "rows": rows(logits)} + # each stream keeps a prefix of its rows (as multi._commit names them: window rows), then one more row each + keep = [list(range(max(1, n - 1))) for n in layout] + commit_streams(sts, record, [[starts[s] + r for r in k] for s, k in enumerate(keep)]) + nxt = [([replies[s][len(k)]], [-1], sts[s]) for s, k in enumerate(keep)] + after, _, _, _ = multi_tree_forward(w, nxt) + case["keep"] = [len(k) for k in keep] + case["after"] = rows(after) + cases.append(case) + # trees: each stream's serial reply as the main chain, with other tokens hung beside it + def tree(stream: int, shape: str): + r = replies[stream] + alt = lambda t: (t + 7919) % w.config.vocab # noqa: E731 + if shape == "fork": # 0 r0; 1 r1 (0); 2 alt (0); 3 r2 (1); 4 alt (1); 5 r3 (3); 6 alt (2) + return [r[0], r[1], alt(r[1]), r[2], alt(r[2]), r[3], alt(r[3])], [-1, 0, 0, 1, 1, 3, 2] + if shape == "fan": # the root's six children (r1 first), then r2 under r1 + return [r[0], r[1]] + [alt(r[1] + j) for j in range(5)] + [r[2]], [-1, 0, 0, 0, 0, 0, 0, 1] + raise ValueError(shape) + + tree_cases = [([(0, "fork")], [[0, 1, 3, 5]]), ([(1, "fan")], [[0, 1, 7]]), + ([(2, "fork"), (3, "fork")], [[0, 1, 3, 5], [0, 2, 6]])] + if a.wide: # every stream a tree in one round, 120 rows + tree_cases.append(([(s, "fork" if s % 2 else "fan") for s in range(16)], + [[0, 1, 3, 5] if s % 2 else [0, 1, 7] for s in range(16)])) + for spec, paths in tree_cases: + sts = [clone_state(states[s]) for s, _ in spec] + wins = [(*tree(s, shape), st) for (s, shape), st in zip(spec, sts)] + if len(spec) == 1: + logits, record = tree_forward(w, tokens(wins[0][0]), wins[0][1], sts[0]) + starts = [0, len(wins[0][0])] + else: + logits, record, _, starts = multi_tree_forward(w, [(t, p, st) for t, p, st in wins]) + case = {"kind": "tree", "streams": [{"stream": s, "tokens": t, "parents": p} for (s, _), (t, p, _) in + zip(spec, wins)], "rows": rows(logits), "paths": paths} + nxt = [int(logits[starts[k] + path[-1]].argmax()) for k, path in enumerate(paths)] + commit_streams(sts, record, [[starts[k] + r for r in path] for k, path in enumerate(paths)]) + after, _, _, _ = multi_tree_forward(w, [([t], [-1], st) for t, st in zip(nxt, sts)]) + case["next"] = nxt + case["after"] = rows(after) + cases.append(case) + torch.cuda.synchronize() + (out / "windows.json").write_text(json.dumps({"prompts": prompts, "replies": replies, "cases": cases}) + "\n") + if rec is not None: + rec.dump(out / "launches.json") + print(f"{len(cases)} cases", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())