diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index b70dfe431f..df2c51f344 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -41,6 +41,10 @@ const kernels = [_]Kernel{ .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts + .{ .name = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // GLM: the Python family's glm5_next/cuda/kda.cu + .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op + .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) + .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, @@ -121,7 +125,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.core, mods.heat); cli.addOptions("build_options", build_options); cli.addImport("checkpoint_cli", checkpointCli(b, target, optimize, strip, native.engines, mods.core)); b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli })); @@ -137,19 +141,34 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen.zig"), .target = target, .optimize = optimize })); runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); + runner.addImport("core", mods.core); + const glm_cuda = glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat); + runner.addImport("glm_cuda", glm_cuda); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; + nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, build_options, true).root_module.strip = strip; +} + +/// The GLM-5.3-Flash CUDA family over the runtime, the core, the lane core, canon's shares and core/heat.zig. +fn glmCuda(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, heat: *std.Build.Module) *std.Build.Module { + const canon = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon.zig"), .target = target, .optimize = optimize }); + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/glm/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("core", core); + m.addImport("canon", canon); + m.addImport("lanes", lanes); + m.addImport("heat", heat); + return m; } /// 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, core: *std.Build.Module, heat: *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 = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) } }, }); return .{ .api = api, .engines = mod }; } @@ -160,8 +179,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, core: *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, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, heat: *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, core, heat); // 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(.{ @@ -189,7 +208,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/build/check_fatbins.zig"), .target = host }) })).step); 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.core, mods.heat); 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 }); @@ -205,6 +224,15 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st b.step("test-server-cpu", "The HTTP server's unit tests (routes, templates, tool parsing) without a GPU").dependOn(&server_test.step); // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + // the GLM-5.3-Flash CUDA family's host side + const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); + step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm })).step); + // the shares a GLM-5.3-Flash rank takes (cluster/canon.zig against the Python family's) + const shares = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon_glm_test.zig"), .target = host, .optimize = .debug }); + step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = shares })).step); + // the tool-call parser's tests, run beside the server's on Macs: the GLM calls the CUDA port serves + const parse_tests = b.createModule(.{ .root_source_file = b.path("zig/src/server/tool_parse_test.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = native_modules.api }, .{ .name = "tokenizer", .module = mods.tokenizer }, .{ .name = "template", .module = template } } }); + step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = parse_tests })).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); @@ -240,7 +268,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.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, build_options, false); } /// SASS targets every release fatbin holds: Ampere (sm_86), RTX 50 and RTX PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/kernels/cuda/glm_glue.cu b/zig/kernels/cuda/glm_glue.cu new file mode 100644 index 0000000000..06f2ba94cc --- /dev/null +++ b/zig/kernels/cuda/glm_glue.cu @@ -0,0 +1,10 @@ +// GLM-5.3-Flash's glue (the Python family's glue, latent, sparse and qmm) in CUDA, each with its Triton PTX math. +#include "glm_glue_norm.cuh" +#include "glm_glue_rows.cuh" +#include "glm_glue_hc.cuh" +#include "glm_glue_moe.cuh" +#include "glm_glue_router.cuh" +#include "glm_glue_latent.cuh" +#include "glm_glue_index.cuh" +#include "glm_glue_bmm.cuh" +#include "glm_glue_attn.cuh" diff --git a/zig/kernels/cuda/glm_glue_attn.cuh b/zig/kernels/cuda/glm_glue_attn.cuh new file mode 100644 index 0000000000..db481c0d9e --- /dev/null +++ b/zig/kernels/cuda/glm_glue_attn.cuh @@ -0,0 +1,326 @@ +// latent.py's MLA attention and sparse.py's index scores by hand: the chunk kernels' online softmax, the chunk merge. +#pragma once + +#include "glm_glue_router.cuh" + +namespace glm_attn { + +constexpr int LW = 512, KT = 32, CH = 512, KPAD = LW + 8; + +// p fp32 summed as Triton's tile does: a thread's two columns, the quad (xor 2, 1), then warps (w0 + w2) + (w1 + w3). +__device__ __forceinline__ float quad_sum(float v) { + v = v + __shfl_xor_sync(0xffffffffu, v, 2); + return v + __shfl_xor_sync(0xffffffffu, v, 1); +} +__device__ __forceinline__ float quad_max(float v) { + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, 2)); + return fmaxf(v, __shfl_xor_sync(0xffffffffu, v, 1)); +} + +__device__ __forceinline__ void cp16(void* dst, const void* src, bool ok) { + const unsigned d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(d), "l"(src), "r"(ok ? 16 : 0)); +} +__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); +} +template +__device__ __forceinline__ void ldm4(uint32_t* r, const void* p) { + const unsigned a = static_cast(__cvta_generic_to_shared(p)); + if (TRANS) + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); + else + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} + +// A 512-wide bf16 row's 16-byte unit u in shared memory, swizzled by the row so ldmatrix's eight rows hit eight banks. +__device__ __forceinline__ int wide(int row, int u) { return row * LW + ((u ^ (row & 7)) << 3); } +// p's [32][32] rows (64 bytes): unit u of row r. +__device__ __forceinline__ int narrow(int row, int u) { return row * KT + ((u ^ ((row >> 1) & 3)) << 3); } + +// latent._dense_chunks/_sparse_chunks (sparse 0/1), block (row, HB heads, chunk), 8 warps: online softmax by 32 keys. +template +constexpr int wide_smem() { + return (HB * LW + 2 * KT * LW + HB * KT) * 2 + (HB * 4 + 2 * HB) * 4; +} + +template +__device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, const int* pos, const int* tok, + const int* cnt, float* po, float* pm, float* pl, int64_t rows, int64_t heads, int64_t w, + float scale, int sparse) { + constexpr int MT = HB / 16; + extern __shared__ __align__(16) unsigned char smem[]; + __nv_bfloat16* qs = reinterpret_cast<__nv_bfloat16*>(smem); + __nv_bfloat16* kvs = qs + HB * LW; + __nv_bfloat16* ps = kvs + 2 * KT * LW; + float* red = reinterpret_cast(ps + HB * KT); + float* stat = red + HB * 4; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5, g = lane >> 2, q = lane & 3; + const int kg = warp & 3, hh = warp >> 2; + const int64_t r = blockIdx.x, h0 = (int64_t)blockIdx.y * HB, c = blockIdx.z; + const int64_t n = sparse ? cnt[r] : 0, limit = sparse ? 0 : (int64_t)pos[0] + r, start = c * CH; + const bool live = sparse ? start < n : start <= limit; + auto load_kv = [&](int tile) { + __nv_bfloat16* dst = kvs + (tile & 1) * KT * LW; + const int64_t first = start + tile * KT; + for (int idx = t; idx < KT * (LW / 8); idx += blockDim.x) { + const int key = idx / (LW / 8), u = idx % (LW / 8); + const int64_t ki = first + key; + const bool ok = sparse ? ki < n : ki <= limit; + const int64_t src = ok ? (sparse ? (int64_t)tok[r * w + ki] : ki) : 0; + cp16(dst + wide(key, u), lc + src * LW + u * 8, ok); + } + }; + float m[MT][2], l[MT][2], o[MT][8][4]; + for (int i = 0; i < MT; ++i) + for (int e = 0; e < 2; ++e) m[i][e] = -INFINITY, l[i][e] = 0.0f; + for (int i = 0; i < MT; ++i) + for (int j = 0; j < 8; ++j) + for (int e = 0; e < 4; ++e) o[i][j][e] = 0.0f; + constexpr int tiles = CH / KT; + if (live) { + for (int idx = t; idx < HB * (LW / 8); idx += blockDim.x) { + const int hd = idx / (LW / 8), u = idx % (LW / 8); + const bool ok = h0 + hd < heads; + cp16(qs + wide(hd, u), qa + (r * heads + (ok ? h0 + hd : 0)) * LW + u * 8, ok); + } + load_kv(0); + commit(); + load_kv(1); + commit(); + } + for (int tile = 0; live && tile < tiles; ++tile) { + if (tile + 1 < tiles) + wait<1>(); + else + wait<0>(); + __syncthreads(); + const __nv_bfloat16* kv = kvs + (tile & 1) * KT * LW; + const int64_t first = start + tile * KT; + float sc[4]; + if (hh < MT) { + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; +#pragma unroll 4 + for (int kk = 0; kk < LW; kk += 32) { + uint32_t a0[4], a1[4], b[4]; + ldm4(a0, qs + wide(hh * 16 + (lane & 15), (kk >> 3) + (lane >> 4))); + ldm4(a1, qs + wide(hh * 16 + (lane & 15), ((kk + 16) >> 3) + (lane >> 4))); + ldm4(b, kv + wide(kg * 8 + (lane & 7), (kk >> 3) + (lane >> 3))); + glm_router::mma16816(acc, a0, b); + glm_router::mma16816(acc, a1, b + 2); + } + for (int e = 0; e < 2; ++e) { + const int64_t ki = first + kg * 8 + q * 2 + e; + const bool ok = sparse ? ki < n : ki <= limit; + sc[e] = ok ? __fmul_rn(acc[e], scale) : -INFINITY; + sc[e + 2] = ok ? __fmul_rn(acc[e + 2], scale) : -INFINITY; + } + for (int hf = 0; hf < 2; ++hf) { + const float mx = quad_max(fmaxf(sc[hf * 2], sc[hf * 2 + 1])); + if (q == 0) red[(hh * 16 + hf * 8 + g) * 4 + kg] = mx; + } + } + __syncthreads(); + if (t < HB * 4) { + float v = red[t]; + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, 2)); + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, 1)); + if ((t & 3) == 0) stat[t / 4] = v; + } + __syncthreads(); + float alpha[MT][2]; + for (int i = 0; i < MT; ++i) + for (int hf = 0; hf < 2; ++hf) { + const float tm = stat[i * 16 + hf * 8 + g]; + const bool active = tm != -INFINITY; + const float nm = active ? fmaxf(m[i][hf], tm) : m[i][hf]; + alpha[i][hf] = active ? (m[i][hf] == -INFINITY ? 0.0f : glm_triton::exp_approx(m[i][hf] - nm)) : 1.0f; + if (i == hh) { + float p[2]; + for (int e = 0; e < 2; ++e) { + const int64_t ki = first + kg * 8 + q * 2 + e; + const bool ok = sparse ? ki < n : ki <= limit; + p[e] = active && ok ? glm_triton::exp_approx(sc[hf * 2 + e] - nm) : 0.0f; + } + const int hr = i * 16 + hf * 8 + g, col = kg * 8 + q * 2; + *reinterpret_cast<__nv_bfloat162*>(ps + narrow(hr, col >> 3) + (col & 7)) = __floats2bfloat162_rn(p[0], p[1]); + const float s = quad_sum(p[0] + p[1]); + if (q == 0) red[hr * 4 + kg] = s; + } + m[i][hf] = nm; + } + __syncthreads(); + if (t < HB * 4) { + float v = red[t]; + v = v + __shfl_xor_sync(0xffffffffu, v, 2); + v = v + __shfl_xor_sync(0xffffffffu, v, 1); + if ((t & 3) == 0) stat[HB + t / 4] = v; + } + for (int i = 0; i < MT; ++i) + for (int j = 0; j < 8; ++j) + for (int e = 0; e < 4; ++e) o[i][j][e] = __fmul_rn(o[i][j][e], alpha[i][e / 2]); +#pragma unroll + for (int k16 = 0; k16 < KT; k16 += 16) { + uint32_t a[MT][4]; + for (int i = 0; i < MT; ++i) ldm4(a[i], ps + narrow(i * 16 + (lane & 15), (k16 >> 3) + (lane >> 4))); +#pragma unroll + for (int j = 0; j < 8; j += 2) { + uint32_t b[4]; + ldm4(b, kv + wide(k16 + (lane & 15), ((warp * 64 + j * 8) >> 3) + (lane >> 4))); + for (int i = 0; i < MT; ++i) { + glm_router::mma16816(o[i][j], a[i], b); + glm_router::mma16816(o[i][j + 1], a[i], b + 2); + } + } + } + __syncthreads(); + for (int i = 0; i < MT; ++i) + for (int hf = 0; hf < 2; ++hf) l[i][hf] = __fmaf_rn(l[i][hf], alpha[i][hf], stat[HB + i * 16 + hf * 8 + g]); + if (tile + 2 < tiles) load_kv(tile + 2); + commit(); + } + for (int i = 0; i < MT; ++i) + for (int e = 0; e < 4; ++e) { + const int64_t hd = h0 + i * 16 + g + (e / 2) * 8; + if (hd >= heads) continue; + const int64_t base = (c * rows + r) * heads + hd; + for (int j = 0; j < 8; ++j) po[base * LW + warp * 64 + j * 8 + q * 2 + e % 2] = o[i][j][e]; + if (warp == 0 && q == 0 && e % 2 == 0) { + pm[base] = m[i][e / 2]; + pl[base] = l[i][e / 2]; + } + } +} + +} // namespace glm_attn + +// The 16-head form (decode windows) and 32-head form (prompt chunks), each in glm_attn::wide_smem() bytes of smem. +extern "C" __global__ void __launch_bounds__(256) + tf_glm_chunks_16(const __nv_bfloat16* qa, const __nv_bfloat16* lc, const int* pos, const int* tok, const int* cnt, + float* po, float* pm, float* pl, int64_t rows, int64_t heads, int64_t w, float scale, int sparse) { + glm_attn::chunks_wide<16>(qa, lc, pos, tok, cnt, po, pm, pl, rows, heads, w, scale, sparse); +} + +extern "C" __global__ void __launch_bounds__(256) + tf_glm_chunks_32(const __nv_bfloat16* qa, const __nv_bfloat16* lc, const int* pos, const int* tok, const int* cnt, + float* po, float* pm, float* pl, int64_t rows, int64_t heads, int64_t w, float scale, int sparse) { + glm_attn::chunks_wide<32>(qa, lc, pos, tok, cnt, po, pm, pl, rows, heads, w, scale, sparse); +} + +// latent._merge, block (row, head), 128 threads: chunks merged in order, o / l by div.full; skips sparse count-0 rows. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_merge(const float* po, const float* pm, const float* pl, __nv_bfloat16* out, const int* cnt, int64_t rows, + int64_t heads, int nch, int sparse) { + constexpr int LW = glm_attn::LW; + const int64_t r = blockIdx.x, h = blockIdx.y; + if (sparse && cnt[r] == 0) return; + const int k0 = threadIdx.x * 4; + float m = -INFINITY, l = 0.0f, o[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int c = 0; c < nch; ++c) { + const int64_t base = (c * rows + r) * heads + h; + const float cm = pm[base], cl = pl[base]; + const bool active = cl > 0.0f; + const float nm = active ? fmaxf(m, cm) : m; + const float b = active ? glm_triton::exp_approx(cm - nm) : 0.0f; + if (c == 0) { + for (int j = 0; j < 4; ++j) o[j] = __fmaf_rn(b, po[base * LW + k0 + j], 0.0f); + l = __fmaf_rn(b, cl, 0.0f); + } else { + const float a = active ? (m == -INFINITY ? 0.0f : glm_triton::exp_approx(m - nm)) : 1.0f; + for (int j = 0; j < 4; ++j) o[j] = __fmaf_rn(o[j], a, __fmul_rn(b, po[base * LW + k0 + j])); + l = __fmaf_rn(l, a, __fmul_rn(b, cl)); + } + m = nm; + } + for (int j = 0; j < 4; ++j) out[(r * heads + h) * LW + k0 + j] = __float2bfloat16_rn(glm_triton::div_full(o[j], l)); +} + +// sparse._scores, block (row, L pool blocks) of 4 warps, 16 pools a warp: sum_h w relu(dots scale), -inf past visible. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_scores(const __nv_bfloat16* qi, const __nv_bfloat16* wts, int64_t w_stride, const __nv_bfloat16* pk, + float* out, const int* pos, int64_t np, int64_t nb, float scale, float wscale, int64_t heads, + int loop) { + constexpr int D = 128, HP = 32, BP = 64; + __shared__ __align__(16) __nv_bfloat16 sm[4 * 2 * 16 * D]; // q's 32 rows, then each warp's two 16-pool buffers + auto at = [](int row, int u) { return row * D + ((u ^ (row & 7)) << 3); }; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, g = lane >> 2, q = lane & 3; + const int64_t r = blockIdx.x, npool = ((int64_t)pos[0] + r + 1) / 4; + for (int idx = threadIdx.x; idx < HP * (D / 8); idx += 128) { + const int hh = idx / (D / 8), u = idx % (D / 8); + const bool ok = hh < heads; + glm_attn::cp16(sm + at(hh, u), qi + (r * heads + (ok ? hh : 0)) * D + u * 8, ok); + } + glm_attn::commit(); + float wv[4]; + for (int i = 0; i < 4; ++i) { + const int64_t hh = g + 8 * i; + wv[i] = hh < heads ? glm_triton::fmul(__bfloat162float(wts[r * w_stride + hh]), wscale) : 0.0f; + } + // A block's pools for this warp: rows p0 .. p0 + 15 of pk, 16 chunks of 16 bytes each, 8 a lane. + auto stage = [&](int buf, int64_t p0) { +#pragma unroll + for (int t = 0; t < 8; ++t) { + const int idx = lane + t * 32, row = idx / (D / 8), u = idx % (D / 8); + const int64_t p = p0 + row; + const bool ok = p < npool; + glm_attn::cp16(sm + (warp * 2 + buf) * 16 * D + at(row, u), pk + (ok ? p : 0) * D + u * 8, ok); + } + }; + const int64_t first = (int64_t)blockIdx.y * loop; + // Blocks that score: those under nb with visible pools (pb * BP < npool); the rest only write -inf. + auto live = [&](int j) { const int64_t pb = first + j; return j < loop && pb < nb && pb * BP < npool; }; + glm_attn::wait<0>(); + __syncthreads(); + uint32_t aq[D / 16][2][4]; +#pragma unroll + for (int kq = 0; kq < D / 16; ++kq) +#pragma unroll + for (int i = 0; i < 2; ++i) glm_attn::ldm4(aq[kq][i], sm + at(i * 16 + (lane & 15), kq * 2 + (lane >> 4))); + __syncthreads(); // q's fragments are in registers: the key buffers take its place + if (live(0)) stage(0, first * BP + warp * 16); + glm_attn::commit(); + for (int j = 0; j < loop; ++j) { + const int64_t pb = first + j; + if (pb >= nb) break; + const int64_t p0 = pb * BP + warp * 16; + if (!live(j)) { + for (int e = lane; e < 16; e += 32) + if (p0 + e < np) out[r * np + p0 + e] = -INFINITY; + continue; + } + if (live(j + 1)) stage((j + 1) & 1, p0 + BP); + glm_attn::commit(); + glm_attn::wait<1>(); + __syncwarp(); + float acc[2][2][4] = {}; + const __nv_bfloat16* keys = sm + (warp * 2 + (j & 1)) * 16 * D; +#pragma unroll + for (int kq = 0; kq < D / 16; ++kq) { + uint32_t b[4]; + glm_attn::ldm4(b, keys + at((lane >> 4) * 8 + (lane & 7), kq * 2 + ((lane >> 3) & 1))); + for (int nt = 0; nt < 2; ++nt) + for (int i = 0; i < 2; ++i) glm_router::mma16816(acc[i][nt], aq[kq][i], b + 2 * nt); + } + __syncwarp(); + for (int nt = 0; nt < 2; ++nt) + for (int e = 0; e < 2; ++e) { + const float v0 = fmaxf(__fmul_rn(acc[0][nt][e], scale), 0.0f), + v1 = fmaxf(__fmul_rn(acc[0][nt][e + 2], scale), 0.0f); + const float v2 = fmaxf(__fmul_rn(acc[1][nt][e], scale), 0.0f), + v3 = fmaxf(__fmul_rn(acc[1][nt][e + 2], scale), 0.0f); + float s = glm_triton::fmul(wv[1], v1); + s = __fmaf_rn(wv[0], v0, s); + s = __fmaf_rn(wv[2], v2, s); + s = __fmaf_rn(wv[3], v3, s); + for (int o = 16; o >= 4; o >>= 1) s = s + __shfl_xor_sync(0xffffffffu, s, o); + const int64_t p = p0 + nt * 8 + q * 2 + e; + if (g == 0 && p < np) out[r * np + p] = p < npool ? s : -INFINITY; + } + } +} diff --git a/zig/kernels/cuda/glm_glue_bmm.cuh b/zig/kernels/cuda/glm_glue_bmm.cuh new file mode 100644 index 0000000000..5c5c4cc4fc --- /dev/null +++ b/zig/kernels/cuda/glm_glue_bmm.cuh @@ -0,0 +1,409 @@ +// qmm._bmm and qmm._bmm_rows by hand (the BF16 matmuls): out [m][n] = x rows . w rows over K slices. +#pragma once + +#include "glm_glue_router.cuh" + +namespace glm_bmm { + +// ROWS rows of one KS-wide K step into smem (pad 8: no bank conflicts), 16-byte loads if aligned, zeros past the end. +template +__device__ __forceinline__ void stage(__nv_bfloat16 (*dst)[KS + 8], const __nv_bfloat16* src, int64_t first, + int64_t limit, int64_t stride, int64_t kk, bool wide) { + for (int idx = threadIdx.x; idx < ROWS * (KS / 8); idx += blockDim.x) { + const int r = idx / (KS / 8), c = (idx % (KS / 8)) * 8; + const int64_t row = first + r; + uint4 v = make_uint4(0, 0, 0, 0); + if (row < limit) { + const __nv_bfloat16* p = src + row * stride + kk + c; + if (wide) { + v = __ldg(reinterpret_cast(p)); + } else { + uint32_t h[4]; + for (int e = 0; e < 4; ++e) h[e] = glm_router::pair(p + 2 * e, true); + v = make_uint4(h[0], h[1], h[2], h[3]); + } + } + *reinterpret_cast(&dst[r][c]) = v; + } +} + +// mma chain per slice from zero, K ascending, as Triton; decode writes slice z, prompt adds in order; tiles keep bits. +template +__device__ void bmm(const __nv_bfloat16* x, const __nv_bfloat16* w, void* out, float* part, int64_t m, int64_t x_stride, + int64_t n, int64_t k, int sk, int rows, int f32) { + constexpr int MT = BM / WM / 16, NT = BN / WN / 8; + __shared__ __align__(16) __nv_bfloat16 sa[BM][KS + 8]; + __shared__ __align__(16) __nv_bfloat16 sb[BN][KS + 8]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, g = lane >> 2, q = lane & 3; + const int64_t bm0 = (int64_t)blockIdx.x * BM, bn0 = (int64_t)blockIdx.y * BN; + const int wm = (warp / WN) * MT * 16, wn = (warp % WN) * NT * 8; + const bool xwide = (reinterpret_cast(x) & 15) == 0 && (x_stride & 7) == 0; + const int64_t per = k / sk; + const int first = rows ? 0 : blockIdx.z, last = rows ? sk : blockIdx.z + 1; + float total[MT][NT][4]; + for (int s = first; s < last; ++s) { + float c[MT][NT][4] = {}; + for (int64_t kk = s * per; kk < (s + 1) * per; kk += KS) { + stage(sa, x, bm0, m, x_stride, kk, xwide); + stage(sb, w, bn0, n, k, kk, true); + __syncthreads(); +#pragma unroll + for (int k16 = 0; k16 < KS; k16 += 16) { + uint32_t a[MT][4], b[NT][2]; +#pragma unroll + for (int i = 0; i < MT; ++i) { + const int r = wm + i * 16 + g; + a[i][0] = *reinterpret_cast(&sa[r][k16 + q * 2]); + a[i][1] = *reinterpret_cast(&sa[r + 8][k16 + q * 2]); + a[i][2] = *reinterpret_cast(&sa[r][k16 + q * 2 + 8]); + a[i][3] = *reinterpret_cast(&sa[r + 8][k16 + q * 2 + 8]); + } +#pragma unroll + for (int j = 0; j < NT; ++j) { + const int col = wn + j * 8 + g; + b[j][0] = *reinterpret_cast(&sb[col][k16 + q * 2]); + b[j][1] = *reinterpret_cast(&sb[col][k16 + q * 2 + 8]); + } +#pragma unroll + for (int i = 0; i < MT; ++i) +#pragma unroll + for (int j = 0; j < NT; ++j) glm_router::mma16816(c[i][j], a[i], b[j]); + } + __syncthreads(); + } +#pragma unroll + for (int i = 0; i < MT; ++i) +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) total[i][j][e] = s == first ? c[i][j][e] : total[i][j][e] + c[i][j][e]; + } +#pragma unroll + for (int i = 0; i < MT; ++i) +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int64_t r = bm0 + wm + i * 16 + g + (e / 2) * 8, col = bn0 + wn + j * 8 + q * 2 + e % 2; + if (r >= m || col >= n) continue; + const float v = total[i][j][e]; + if (!rows && sk > 1) + part[((int64_t)blockIdx.z * m + r) * n + col] = v; + else if (f32) + static_cast(out)[r * n + col] = v; + else + static_cast<__nv_bfloat16*>(out)[r * n + col] = __float2bfloat16_rn(v); + } +} + +__device__ __forceinline__ void cp16(void* dst, const void* src, bool ok) { + const unsigned d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(d), "l"(src), "r"(ok ? 16 : 0)); +} +__device__ __forceinline__ void cp16s(unsigned d, const void* src, bool ok) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(d), "l"(src), "r"(ok ? 16 : 0)); +} +__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); +} +__device__ __forceinline__ void ldm4s(uint32_t* r, unsigned a) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(a)); +} +__device__ __forceinline__ void ldm2s(uint32_t* r, unsigned a) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(a)); +} +__device__ __forceinline__ void ldm4(uint32_t* r, const void* p) { + ldm4s(r, static_cast(__cvta_generic_to_shared(p))); +} + +// Prompt rows by cp.async (STAGES buffers, swizzled rows, ldmatrix): bmm's rows-form chains. SLICE: decode slice z. +template +struct Tile { + static constexpr int LD = KS, A = BM * LD, B = BN * LD, STAGED = STAGES * (A + B) * 2, LDO = BN + 4, + LDH = BN + 8, THREADS = WM * WN * 32; + // Dynamic shared bytes: the stages, or the output tile if larger (pitches BN + 4 words, BN + 8 bf16: 16-byte rows). + static constexpr int smem(bool f32) { + const int tile = f32 ? BM * LDO * 4 : BM * LDH * 2; + return tile > STAGED ? tile : STAGED; + } + static_assert(KS == 32 || KS == 64, "swizzles for rows of 4 or 8 chunks"); + __device__ static int swz(int r) { return KS == 32 ? (r >> 1) & 3 : r & 7; } +}; + +template +__device__ void bmm_tiles(const __nv_bfloat16* x, const __nv_bfloat16* w, void* out, int64_t m, int64_t x_stride, + int64_t n, int64_t k, int sk, int f32) { + using T = Tile; + constexpr int MT = BM / WM / 16, NT = BN / WN / 8; + static_assert(MT >= 1 && (NT == 1 || NT % 2 == 0), "a warp's tile is whole m16 rows and one or pairs of n8 columns"); + extern __shared__ __align__(16) unsigned char smem[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, q = lane & 3; + const int64_t bm0 = (int64_t)blockIdx.x * BM, bn0 = (int64_t)blockIdx.y * BN; + const int wm = (warp / WN) * MT * 16, wn = (warp % WN) * NT * 8; + const int steps = static_cast((SLICE ? k / sk : k) / KS), per = static_cast(k / sk / KS); + const int64_t k0 = SLICE ? blockIdx.z * (k / sk) : 0; + // A thread copies one 16-byte column of rows lr, lr + RPI, ...; a step moves it KS along K (past the matrix: zeros). + constexpr int CPR = KS / 8, RPI = T::THREADS / CPR, AI = BM / RPI, BI = BN / RPI, STAGE = (T::A + T::B) * 2; + static_assert(T::THREADS % CPR == 0 && BM % RPI == 0 && BN % RPI == 0 && RPI % 8 == 0, "whole rows a pass"); + const int lr = tid / CPR, lc = (tid % CPR) * 8; + const unsigned sbase = static_cast(__cvta_generic_to_shared(smem)); + const unsigned soff = (lr * T::LD + ((lc / 8) ^ T::swz(lr)) * 8) * 2; // rows lr + i RPI swizzle alike (RPI % 8 == 0) + const __nv_bfloat16* ga[AI]; + const __nv_bfloat16* gb[BI]; + unsigned oka = 0, okb = 0; +#pragma unroll + for (int i = 0; i < AI; ++i) { + const int64_t row = bm0 + lr + i * RPI; + oka |= (row < m ? 1u : 0u) << i; + ga[i] = x + (row < m ? row : 0) * x_stride + k0 + lc; + } +#pragma unroll + for (int i = 0; i < BI; ++i) { + const int64_t row = bn0 + lr + i * RPI; + okb |= (row < n ? 1u : 0u) << i; + gb[i] = w + (row < n ? row : 0) * k + k0 + lc; + } + auto load = [&](int st, int kt) { + const unsigned s = sbase + st * STAGE + soff; + const int kk = kt * KS; +#pragma unroll + for (int i = 0; i < AI; ++i) cp16s(s + i * RPI * T::LD * 2, ga[i] + kk, (oka >> i) & 1); +#pragma unroll + for (int i = 0; i < BI; ++i) cp16s(s + T::A * 2 + i * RPI * T::LD * 2, gb[i] + kk, (okb >> i) & 1); + }; + // ldmatrix rows in stage 0: the warp's A rows (lane & 15, chunk + 1 for lanes 16 on) and B rows; k16s swizzle as r. + const int ra = wm + (lane & 15), rb = wn + (NT == 1 ? 0 : (lane >> 4) * 8) + (lane & 7); + const unsigned fa = sbase + ra * T::LD * 2, fb = sbase + T::A * 2 + rb * T::LD * 2; + unsigned ca[KS / 16], cb[KS / 16]; +#pragma unroll + for (int s16 = 0; s16 < KS / 16; ++s16) { + ca[s16] = ((2 * s16 + (lane >> 4)) ^ T::swz(ra)) * 16; + cb[s16] = ((2 * s16 + ((lane >> 3) & 1)) ^ T::swz(rb)) * 16; + } + float c[MT][NT][4] = {}, total[SPLIT ? MT : 1][SPLIT ? NT : 1][4] = {}; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < steps) load(s, s); + commit(); + } + int rs = 0, ws = STAGES - 1, left = per; + for (int kt = 0; kt < steps; ++kt) { + if constexpr (STAGES == 1) { + __syncthreads(); // every warp is done with the one buffer's last step + load(0, kt); + commit(); + wait<0>(); + __syncthreads(); + } else if constexpr (SLICE) { + // decode: the next stage's copies go out before this step's wait, as Triton's _bmm issues them (two barriers) + __syncthreads(); // every warp is done with step kt - 1's buffer (ws) + if (kt + STAGES - 1 < steps) load(ws, kt + STAGES - 1); + commit(); + ws = ws + 1 == STAGES ? 0 : ws + 1; + wait(); + __syncthreads(); + } else { + wait(); + __syncthreads(); + if (kt + STAGES - 1 < steps) load(ws, kt + STAGES - 1); + commit(); + ws = ws + 1 == STAGES ? 0 : ws + 1; + } + const unsigned sa = fa + rs * STAGE, sb = fb + rs * STAGE; + rs = rs + 1 == STAGES ? 0 : rs + 1; +#pragma unroll + for (int k16 = 0; k16 < KS; k16 += 16) { + uint32_t a[MT][4], b[NT][2]; +#pragma unroll + for (int i = 0; i < MT; ++i) ldm4s(a[i], sa + i * 16 * T::LD * 2 + ca[k16 / 16]); + if constexpr (NT == 1) { + ldm2s(b[0], sb + cb[k16 / 16]); + } else { +#pragma unroll + for (int j = 0; j < NT; j += 2) { + uint32_t r[4]; + ldm4s(r, sb + j * 8 * T::LD * 2 + cb[k16 / 16]); + b[j][0] = r[0], b[j][1] = r[1], b[j + 1][0] = r[2], b[j + 1][1] = r[3]; + } + } +#pragma unroll + for (int i = 0; i < MT; ++i) +#pragma unroll + for (int j = 0; j < NT; ++j) glm_router::mma16816(c[i][j], a[i], b[j]); + } + if constexpr (SPLIT) { + if (--left == 0) { // a slice's last step: its chain into the running sum, the next slice from zero + const bool first = kt + 1 == per; + left = per; +#pragma unroll + for (int i = 0; i < MT; ++i) +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) { + total[i][j][e] = first ? c[i][j][e] : total[i][j][e] + c[i][j][e]; + c[i][j][e] = 0.0f; + } + } + } + } + wait<0>(); + auto val = [&](int i, int j, int e) -> float { + if constexpr (SPLIT) return total[i][j][e]; + else return c[i][j][e]; + }; + // The tile into shared memory (Tile's pitches), then rows out 16 bytes a store. + __syncthreads(); + float* so = reinterpret_cast(smem); + __nv_bfloat16* sh = reinterpret_cast<__nv_bfloat16*>(smem); +#pragma unroll + for (int i = 0; i < MT; ++i) +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = wm + i * 16 + g + h * 8, col = wn + j * 8 + q * 2; + if (f32) + *reinterpret_cast(&so[r * T::LDO + col]) = make_float2(val(i, j, 2 * h), val(i, j, 2 * h + 1)); + else + *reinterpret_cast<__nv_bfloat162*>(&sh[r * T::LDH + col]) = + __halves2bfloat162(__float2bfloat16_rn(val(i, j, 2 * h)), __float2bfloat16_rn(val(i, j, 2 * h + 1))); + } + __syncthreads(); + const int per_row = f32 ? BN / 4 : BN / 8; + for (int idx = tid; idx < BM * per_row; idx += T::THREADS) { + const int r = idx / per_row, v = idx % per_row; + const int64_t row = bm0 + r, col = bn0 + v * (f32 ? 4 : 8); + if (row >= m || col >= n) continue; + if (f32) + *reinterpret_cast(static_cast(out) + row * n + col) = + *reinterpret_cast(&so[r * T::LDO + v * 4]); + else + *reinterpret_cast(static_cast<__nv_bfloat16*>(out) + row * n + col) = + *reinterpret_cast(&sh[r * T::LDH + v * 8]); + } +} + +} // namespace glm_bmm + +// Prompt windows on aligned rows, K sliced or at most 256: 128 x 64 a block of 8 warps, K 64 in 3 cp.async stages. +extern "C" __global__ void __launch_bounds__(256) + tf_glm_bmm_tile(const __nv_bfloat16* x, const __nv_bfloat16* w, void* out, int64_t m, int64_t x_stride, + int64_t n, int64_t k, int sk, int f32) { + glm_bmm::bmm_tiles<128, 64, 64, 3, 2, 4, true>(x, w, out, m, x_stride, n, k, sk, f32); +} + +// Prompt windows on aligned rows, one slice: 128 x 128 a block of 8 warps, K 32 in three stages, two blocks an SM. +extern "C" __global__ void __launch_bounds__(256, 2) + tf_glm_bmm_tile1(const __nv_bfloat16* x, const __nv_bfloat16* w, void* out, int64_t m, int64_t x_stride, + int64_t n, int64_t k, int sk, int f32) { + glm_bmm::bmm_tiles<128, 128, 32, 3, 2, 4, false>(x, w, out, m, x_stride, n, k, sk, f32); +} + +// glue._router_part on aligned rows: part[s][m][e] = x[m] . w[e] over K slice s, tf_glm_router_part's chain, tiled. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_router_tile(const __nv_bfloat16* x, const __nv_bfloat16* w, float* part, int64_t rows, int64_t x_stride, + int64_t d, int64_t ne, int ks) { + constexpr int BM = 64, BE = 32, KS = 64, LD = KS + 8, A = BM * LD, B = BE * LD, STAGE = (A + B) * 2, RPI = 16; + __shared__ __align__(16) __nv_bfloat16 sm[2][A + B]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, q = lane & 3; + const int64_t m0 = (int64_t)blockIdx.x * BM, e0 = (int64_t)blockIdx.y * BE, s = blockIdx.z, span = d / ks; + const int steps = static_cast(span / KS), lr = tid / 8, lc = (tid % 8) * 8; + const __nv_bfloat16* ga[BM / RPI]; + const __nv_bfloat16* gb[BE / RPI]; + unsigned oka = 0, okb = 0; +#pragma unroll + for (int i = 0; i < BM / RPI; ++i) { + const int64_t row = m0 + lr + i * RPI; + oka |= (row < rows ? 1u : 0u) << i; + ga[i] = x + (row < rows ? row : 0) * x_stride + s * span + lc; + } +#pragma unroll + for (int i = 0; i < BE / RPI; ++i) { + const int64_t e = e0 + lr + i * RPI; + okb |= (e < ne ? 1u : 0u) << i; + gb[i] = w + (e < ne ? e : 0) * d + s * span + lc; + } + const unsigned sbase = static_cast(__cvta_generic_to_shared(sm)), soff = (lr * LD + lc) * 2; + auto load = [&](int st, int kt) { + const unsigned a = sbase + st * STAGE + soff; +#pragma unroll + for (int i = 0; i < BM / RPI; ++i) glm_bmm::cp16s(a + i * RPI * LD * 2, ga[i] + kt * KS, (oka >> i) & 1); +#pragma unroll + for (int i = 0; i < BE / RPI; ++i) glm_bmm::cp16s(a + A * 2 + i * RPI * LD * 2, gb[i] + kt * KS, (okb >> i) & 1); + }; + const unsigned fa = sbase + ((warp * 16 + (lane & 15)) * LD + (lane >> 4) * 8) * 2; + const unsigned fb = sbase + A * 2 + (((lane >> 4) * 8 + (lane & 7)) * LD + ((lane >> 3) & 1) * 8) * 2; + float c[BE / 8][4] = {}; + load(0, 0); + glm_bmm::commit(); + for (int kt = 0; kt < steps; ++kt) { + if (kt + 1 < steps) load((kt + 1) & 1, kt + 1); + glm_bmm::commit(); + glm_bmm::wait<1>(); + __syncthreads(); + const unsigned st = (kt & 1) * STAGE; +#pragma unroll + for (int k16 = 0; k16 < KS; k16 += 16) { + uint32_t a[4], b[BE / 8][2]; + glm_bmm::ldm4s(a, fa + st + k16 * 2); +#pragma unroll + for (int j = 0; j < BE / 8; j += 2) { + uint32_t r[4]; + glm_bmm::ldm4s(r, fb + st + (j * 8 * LD + k16) * 2); + b[j][0] = r[0], b[j][1] = r[1], b[j + 1][0] = r[2], b[j + 1][1] = r[3]; + } +#pragma unroll + for (int j = 0; j < BE / 8; ++j) glm_router::mma16816(c[j], a, b[j]); + } + __syncthreads(); + } +#pragma unroll + for (int j = 0; j < BE / 8; ++j) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int jj = 0; jj < 2; ++jj) { + const int64_t r = m0 + warp * 16 + g + h * 8, col = e0 + j * 8 + q * 2 + jj; + if (r < rows && col < ne) part[(s * rows + r) * ne + col] = c[j][h * 2 + jj]; + } +} + +// Decode windows (under 128 rows): 16 x 64 a block of 4 warps, K 64 at a time, one K slice a block along z. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_bmm_decode(const __nv_bfloat16* x, const __nv_bfloat16* w, void* out, float* part, int64_t m, + int64_t x_stride, int64_t n, int64_t k, int sk, int f32) { + glm_bmm::bmm<16, 64, 64, 1, 4>(x, w, out, part, m, x_stride, n, k, sk, 0, f32); +} + +// Decode windows on aligned rows as Triton's _bmm stages a row (cuda_b16.decodeTile, 4 warps): 16 x BN, K 64, slice z. +namespace glm_bmm { +template +__device__ __forceinline__ void dec_tile(const __nv_bfloat16* x, const __nv_bfloat16* w, void* out, float* part, + int64_t m, int64_t x_stride, int64_t n, int64_t k, int sk, int f32) { + void* o = sk > 1 ? static_cast(part + blockIdx.z * m * n) : out; + bmm_tiles<16, BN, 64, ST, 1, 4, false, true>(x, w, o, m, x_stride, n, k, sk, sk > 1 ? 1 : f32); +} +} // namespace glm_bmm +#define TF_GLM_DEC_TILE(name, BN, ST) \ + extern "C" __global__ void __launch_bounds__(128) \ + name(const __nv_bfloat16* x, const __nv_bfloat16* w, void* out, float* part, int64_t m, int64_t x_stride, \ + int64_t n, int64_t k, int sk, int f32) { \ + glm_bmm::dec_tile(x, w, out, part, m, x_stride, n, k, sk, f32); \ + } +TF_GLM_DEC_TILE(tf_glm_bmm_d32s1, 32, 1) +TF_GLM_DEC_TILE(tf_glm_bmm_d32s2, 32, 2) +TF_GLM_DEC_TILE(tf_glm_bmm_d32s3, 32, 3) +TF_GLM_DEC_TILE(tf_glm_bmm_d64s2, 64, 2) +#undef TF_GLM_DEC_TILE + +// Prompt windows (_bmm_rows): 128 x 64 a block of 4 warps, K 32 at a time, the slices added in order in registers. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_bmm_rows(const __nv_bfloat16* x, const __nv_bfloat16* w, void* out, int64_t m, int64_t x_stride, + int64_t n, int64_t k, int sk, int f32) { + glm_bmm::bmm<128, 64, 32, 2, 2>(x, w, out, nullptr, m, x_stride, n, k, sk, 1, f32); +} diff --git a/zig/kernels/cuda/glm_glue_hc.cuh b/zig/kernels/cuda/glm_glue_hc.cuh new file mode 100644 index 0000000000..1de6bb3b71 --- /dev/null +++ b/zig/kernels/cuda/glm_glue_hc.cuh @@ -0,0 +1,200 @@ +// glue._hc_partial and glue._hc_finish by hand (hyper-connections' pre side, 4 streams of 4096). +#pragma once + +#include "glm_triton.cuh" + +using namespace glm_triton; + +// glue._hc_partial, block (4 rows, K block b) of 128 threads: a row's 24 dots and squares in Triton's fma order. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_hc_partial( + const __nv_bfloat16* x, const __nv_bfloat16* fn, float* part, int64_t wide, int nb, int64_t rows) { + constexpr int R = 4; + __shared__ float smem[R][4]; + const int64_t r0 = (int64_t)blockIdx.x * R, b = blockIdx.y, kb = wide / nb; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5, c = t & 15, m0 = (t >> 4) & 7; + float acc[R][3] = {}, ss[R] = {}; + for (int64_t step = 0; step < kb / 128; ++step) { + const int64_t base = b * kb + step * 128; + float xv[R][8]; + for (int q = 0; q < R; ++q) { // the 8 inputs in one 16-byte load (wide and kb are multiples of 128: aligned) + uint4 u = make_uint4(0, 0, 0, 0); + if (r0 + q < rows) u = *reinterpret_cast(x + (r0 + q) * wide + base + c * 8); + const __nv_bfloat16* h = reinterpret_cast(&u); + for (int j = 0; j < 8; ++j) xv[q][j] = bf(h[j]); + } + for (int i = 0; i < 3; ++i) { + const uint4 u = __ldg(reinterpret_cast(fn + (int64_t)(m0 + 8 * i) * wide + base + c * 8)); + const __nv_bfloat16* hw = reinterpret_cast(&u); + float w[8]; + for (int j = 0; j < 8; ++j) w[j] = bf(hw[j]); + for (int q = 0; q < R; ++q) { + float d = __fmul_rn(xv[q][1], w[1]); + d = __fmaf_rn(xv[q][0], w[0], d); + for (int j = 2; j < 8; ++j) d = __fmaf_rn(xv[q][j], w[j], d); + for (int o = 8; o > 0; o >>= 1) d = d + __shfl_xor_sync(0xffffffffu, d, o); + acc[q][i] = acc[q][i] + d; + } + } + for (int q = 0; q < R; ++q) { + const float xt = r0 + q < rows ? bf(x[(r0 + q) * wide + base + t]) : 0.0f; + ss[q] = __fmaf_rn(xt, xt, ss[q]); + } + } + for (int q = 0; q < R; ++q) { + for (int o = 16; o > 0; o >>= 1) ss[q] += __shfl_xor_sync(0xffffffffu, ss[q], o); + if (lane == 0) smem[q][warp] = ss[q]; + } + __syncthreads(); + for (int q = 0; q < R; ++q) { + if (r0 + q >= rows) break; + float* out = part + ((r0 + q) * nb + b) * 32; + if (c == 0) + for (int i = 0; i < 3; ++i) out[m0 + 8 * i] = acc[q][i]; + if (t == 0) out[24] = (smem[q][0] + smem[q][2]) + (smem[q][1] + smem[q][3]); + } +} + +// The same for decode windows (16 blocks a row): 8 steps of 128 at once, a 128-thread group each; same ops, same order. +extern "C" __global__ void __launch_bounds__(1024) + tf_glm_hc_partial8(const __nv_bfloat16* x, const __nv_bfloat16* fn, float* part, int64_t wide, int nb) { + __shared__ float dots[8][24]; + __shared__ float smem[4]; + const int64_t r = blockIdx.x, b = blockIdx.y, kb = wide / nb; + const int grp = threadIdx.x >> 7, t = threadIdx.x & 127, c = t & 15, m0 = (t >> 4) & 7; + const int steps = static_cast(kb / 128); + if (grp < steps) { + const int64_t base = b * kb + grp * 128; + float xv[8]; + { + const uint4 u = *reinterpret_cast(x + r * wide + base + c * 8); + const __nv_bfloat16* h = reinterpret_cast(&u); + for (int j = 0; j < 8; ++j) xv[j] = bf(h[j]); + } + for (int i = 0; i < 3; ++i) { + const uint4 u = __ldg(reinterpret_cast(fn + (int64_t)(m0 + 8 * i) * wide + base + c * 8)); + const __nv_bfloat16* w = reinterpret_cast(&u); + float d = __fmul_rn(xv[1], bf(w[1])); + d = __fmaf_rn(xv[0], bf(w[0]), d); + for (int j = 2; j < 8; ++j) d = __fmaf_rn(xv[j], bf(w[j]), d); + for (int o = 8; o > 0; o >>= 1) d = d + __shfl_xor_sync(0xffffffffu, d, o); + if (c == 0) dots[grp][m0 + 8 * i] = d; + } + } + if (grp == 0) { + float ss = 0.0f; + for (int s = 0; s < steps; ++s) { + const float xt = bf(x[r * wide + b * kb + s * 128 + t]); + ss = __fmaf_rn(xt, xt, ss); + } + for (int o = 16; o > 0; o >>= 1) ss += __shfl_xor_sync(0xffffffffu, ss, o); + if ((t & 31) == 0) smem[t >> 5] = ss; + } + __syncthreads(); + float* out = part + (r * nb + b) * 32; + if (threadIdx.x < 24) { + float acc = 0.0f; + for (int s = 0; s < steps; ++s) acc = acc + dots[s][threadIdx.x]; + out[threadIdx.x] = acc; + } else if (threadIdx.x == 32) { + out[24] = (smem[0] + smem[2]) + (smem[1] + smem[3]); + } +} + +namespace glm_hc { + +// The Sinkhorn sums of 4: lanes (i, j) = (lane / 4 % 4, lane % 4), rows over j (xor 2, 1), columns over i (xor 8, 4). +__device__ __forceinline__ float row4(float v) { + v = v + __shfl_xor_sync(0xffffffffu, v, 2); + return v + __shfl_xor_sync(0xffffffffu, v, 1); +} +__device__ __forceinline__ float col4(float v) { + v = v + __shfl_xor_sync(0xffffffffu, v, 8); + return v + __shfl_xor_sync(0xffffffffu, v, 4); +} +__device__ __forceinline__ float sigmoid(float v) { return div_full(1.0f, exp_approx(0.0f - v) + 1.0f); } + +} // namespace glm_hc + +// glue._hc_finish, 256 threads a row (D 4096, 4 streams): warp 0 sums K blocks and Sinkhorn, then collapse, RMS-norm. +extern "C" __global__ void __launch_bounds__(256) + tf_glm_hc_finish(const __nv_bfloat16* x, const float* part, const float* base, const float* scale, + const __nv_bfloat16* nw, __nv_bfloat16* out, float* xs, float* post, float* comb, float eps, + float hc_eps, int nb, int iters) { + constexpr int D = 4096, S = 4; + __shared__ float smem[8]; + __shared__ float pre[4]; + const int64_t r = blockIdx.x; + const int t = threadIdx.x, lane = t & 31; + // A thread's 8 inputs of each half are contiguous: one 16-byte load per stream and half, issued before Sinkhorn. + const __nv_bfloat16* row = x + r * S * D; + uint4 xu[2][S], wu[2]; + for (int rep = 0; rep < 2; ++rep) { + for (int st = 0; st < S; ++st) xu[rep][st] = *reinterpret_cast(row + st * D + rep * 2048 + t * 8); + wu[rep] = __ldg(reinterpret_cast(nw + rep * 2048 + t * 8)); + } + if (t < 32) { + const float* p = part + r * nb * 32; + float mix = 0.0f, ss = 0.0f; + for (int b = 0; b < nb; ++b) { + mix = mix + p[b * 32 + lane]; + ss = ss + p[b * 32 + 24]; + } + const float rinv = div_full(1.0f, sqrt_approx(eps + div_full(ss, (float)(S * D)))); + mix = __fmul_rn(mix, rinv); + const int k = lane & 3, idx = lane & 15; + const float bpre = base[k], bpost = base[4 + k], bcomb = base[8 + idx]; + const float mpre = __shfl_sync(0xffffffffu, mix, k), mpost = __shfl_sync(0xffffffffu, mix, 4 + k); + const float mcomb = __shfl_sync(0xffffffffu, mix, 8 + idx); + const float pr = hc_eps + glm_hc::sigmoid(__fmaf_rn(mpre, scale[0], bpre)); + const float sp = glm_hc::sigmoid(__fmaf_rn(mpost, scale[1], bpost)); + const float cl = __fmaf_rn(mcomb, scale[2], bcomb); + float cmax = fmaxf(cl, __shfl_xor_sync(0xffffffffu, cl, 2)); + cmax = fmaxf(cmax, __shfl_xor_sync(0xffffffffu, cmax, 1)); + const float ce = exp_approx(cl - cmax); + float cm = fadd(hc_eps, div_full(ce, glm_hc::row4(ce))); // ptxas: fma(rcp(sum), ce, hc_eps), as for Triton + cm = div_full(cm, hc_eps + glm_hc::col4(cm)); + for (int it = 1; it < iters; ++it) { + cm = div_full(cm, hc_eps + glm_hc::row4(cm)); + cm = div_full(cm, hc_eps + glm_hc::col4(cm)); + } + if (lane < 4) { + post[r * S + lane] = sp + sp; + pre[lane] = pr; + } + if (lane < 16) comb[r * 16 + lane] = cm; + } + __syncthreads(); + const float p0 = pre[0], p1 = pre[1], p2 = pre[2], p3 = pre[3]; + float c[16]; + for (int k = 0; k < 16; ++k) { + const __nv_bfloat16* h0 = reinterpret_cast(&xu[k / 8][0]); + const __nv_bfloat16* h1 = reinterpret_cast(&xu[k / 8][1]); + const __nv_bfloat16* h2 = reinterpret_cast(&xu[k / 8][2]); + const __nv_bfloat16* h3 = reinterpret_cast(&xu[k / 8][3]); + const float x0 = bf(h0[k % 8]), x1 = bf(h1[k % 8]), x2 = bf(h2[k % 8]), x3 = bf(h3[k % 8]); + float m = k % 2 == 0 ? __fmaf_rn(p1, x1, __fmul_rn(p0, x0)) : __fmaf_rn(p0, x0, __fmul_rn(p1, x1)); + m = __fmaf_rn(p2, x2, m); + c[k] = rbf(__fmaf_rn(p3, x3, m)); + } + float acc = __fmaf_rn(c[0], c[0], __fmul_rn(c[1], c[1])); + for (int k = 2; k < 16; ++k) acc = k % 2 == 0 ? __fmaf_rn(c[k], c[k], acc) : fadd(fmul(c[k], c[k]), acc); + const float total = block_sum<8>(acc, smem); + const float rinv = div_full(1.0f, sqrt_approx(eps + div_full(total, (float)D))); + for (int rep = 0; rep < 2; ++rep) { + float y[8]; + const __nv_bfloat16* wh = reinterpret_cast(&wu[rep]); + uint4 ou; + __nv_bfloat16* oh = reinterpret_cast<__nv_bfloat16*>(&ou); + for (int j = 0; j < 8; ++j) { + const __nv_bfloat16 o = __float2bfloat16_rn(__fmul_rn(bf(wh[j]), rbf(__fmul_rn(rinv, c[rep * 8 + j])))); + oh[j] = o; + y[j] = bf(o); + } + *reinterpret_cast(out + r * D + rep * 2048 + t * 8) = ou; + float g = y[0] + y[1]; + for (int j = 2; j < 8; ++j) g = y[j] + g; + for (int o = 4; o > 0; o >>= 1) g = g + __shfl_xor_sync(0xffffffffu, g, o); + if ((t & 7) == 0) xs[r * (D / 64) + rep * 32 + t / 8] = g; + } +} diff --git a/zig/kernels/cuda/glm_glue_index.cuh b/zig/kernels/cuda/glm_glue_index.cuh new file mode 100644 index 0000000000..ba1241cdec --- /dev/null +++ b/zig/kernels/cuda/glm_glue_index.cuh @@ -0,0 +1,187 @@ +// sparse.py's DSA indexer glue by hand: the index key write, the 4-token pools and the per-row top-512 pools. +#pragma once + +#include "glm_triton.cuh" + +using namespace glm_triton; + +// sparse._index_write, a warp a row (D 128): LayerNorm of k_raw (div.full, sqrt.approx) into ik, gate row into ig. +extern "C" __global__ void __launch_bounds__(32) + tf_glm_index_write(const __nv_bfloat16* kr, int64_t k_stride, const float* gr, const __nv_bfloat16* lnw, + const __nv_bfloat16* lnb, __nv_bfloat16* ik, __nv_bfloat16* ig, const int* pos, float eps) { + constexpr int D = 128; + const int64_t r = blockIdx.x, at = (int64_t)pos[0] + r; + const int d0 = threadIdx.x * 4; + float x[4]; + for (int j = 0; j < 4; ++j) x[j] = bf(kr[r * k_stride + d0 + j]); + float s = x[0] + x[1]; + s = x[2] + s; + s = x[3] + s; + for (int o = 16; o > 0; o >>= 1) s = s + __shfl_xor_sync(0xffffffffu, s, o); + const float mean = div_full(s, (float)D); + float xc[4]; + for (int j = 0; j < 4; ++j) xc[j] = __fsub_rn(x[j], mean); // sub.rn in Triton: never contracted + float v = fmul(xc[1], xc[1]); + v = __fmaf_rn(xc[0], xc[0], v); + v = __fmaf_rn(xc[2], xc[2], v); + v = __fmaf_rn(xc[3], xc[3], v); + for (int o = 16; o > 0; o >>= 1) v = v + __shfl_xor_sync(0xffffffffu, v, o); + const float sd = sqrt_approx(fadd(eps, div_full(v, (float)D))); + for (int j = 0; j < 4; ++j) { + const int d = d0 + j; + ik[at * D + d] = __float2bfloat16_rn(__fmaf_rn(div_full(xc[j], sd), bf(lnw[d]), bf(lnb[d]))); + ig[at * D + d] = __float2bfloat16_rn(gr[r * D + d]); + } +} + +// sparse._pool_keys, a warp a completed pool: softmax of its 4 gate rows (+ ape) times bf16 keys, summed in order. +extern "C" __global__ void __launch_bounds__(32) + tf_glm_pool_keys(const __nv_bfloat16* ik, const __nv_bfloat16* ig, const __nv_bfloat16* ape, __nv_bfloat16* pk, + const int* pos, int64_t rows) { + constexpr int D = 128; + const int64_t P = pos[0], p = P / 4 + blockIdx.x; + if (4 * p + 3 > P + rows - 1) return; + for (int j = 0; j < 4; ++j) { + const int d = threadIdx.x * 4 + j; + float l[4], e[4]; + for (int q = 0; q < 4; ++q) l[q] = bf(ig[(4 * p + q) * D + d]) + bf(ape[q * D + d]); + const float m = fmaxf(fmaxf(l[0], l[1]), fmaxf(l[2], l[3])); + for (int q = 0; q < 4; ++q) e[q] = exp_approx(l[q] - m); + const float s = ((e[0] + e[1]) + e[2]) + e[3]; + float t = 0.0f; + for (int q = 0; q < 4; ++q) { + const float tq = rbf(__fmul_rn(rbf(div_full(e[q], s)), bf(ik[(4 * p + q) * D + d]))); + t = q == 0 ? tq : t + tq; + } + pk[p * D + d] = __float2bfloat16_rn(t); + } +} + +namespace glm_index { + +// sparse._order_key: a score's bits as a uint32 whose unsigned order is the scores' (-0 counted as +0). +__device__ __forceinline__ uint32_t order_key(float s) { + const int32_t b = __float_as_int(s + 0.0f); + return (uint32_t)(b ^ ((b >> 31) | (int32_t)0x80000000)); +} + +} // namespace glm_index + +namespace glm_index { + +// 4 warps' exclusive prefix of v in thread order and its total, via slot of `wsum`; its barrier orders slot reuse. +__device__ __forceinline__ int block_scan(int v, int (*wsum)[4], int slot, int* total) { + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + int inc = v; + for (int o = 1; o < 32; o <<= 1) { + const int u = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) inc += u; + } + if (lane == 31) wsum[slot][warp] = inc; + __syncthreads(); + int before = inc - v; + for (int w = 0; w < warp; ++w) before += wsum[slot][w]; + *total = wsum[slot][0] + wsum[slot][1] + wsum[slot][2] + wsum[slot][3]; + return before; +} + +} // namespace glm_index + +// sparse._select_rows (VISIBLE), a row per 4-warp block: K best visible pools by 8-bit radix select, ties to lower. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_select_rows(const float* scores, int64_t* out, int64_t np, const int* pos, int k) { + __shared__ int hist[256]; + __shared__ uint32_t sh_prefix; + __shared__ int sh_need, wsum[2][4]; + const int64_t r = blockIdx.x; + const float* row = scores + r * np; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5; + const unsigned below = (1u << lane) - 1; + const int64_t vis = ((int64_t)pos[0] + r + 1) / 4; + const int64_t limit = np < (vis > k ? vis : k) ? np : (vis > k ? vis : k); + const bool vec = (reinterpret_cast(row) & 15) == 0; + // The order keys of pools c + 8 t .. + 7 (0 past limit; callers test the index). + auto keys = [&](int64_t c, uint32_t u[8]) { + const int64_t i0 = c + 8 * t; + if (vec && i0 + 8 <= limit) { + const float4 a = *reinterpret_cast(row + i0), b = *reinterpret_cast(row + i0 + 4); + const float v[8] = {a.x, a.y, a.z, a.w, b.x, b.y, b.z, b.w}; + for (int j = 0; j < 8; ++j) u[j] = glm_index::order_key(v[j]); + } else { + for (int j = 0; j < 8; ++j) u[j] = i0 + j < limit ? glm_index::order_key(row[i0 + j]) : 0u; + } + }; + uint32_t prefix = 0, fixed = 0; + int need = k; + for (int shift = 24; shift >= 0; shift -= 8) { + for (int b = t; b < 256; b += blockDim.x) hist[b] = 0; + __syncthreads(); + for (int64_t c = 0; c < limit; c += 1024) { + uint32_t u[8]; + keys(c, u); + for (int j = 0; j < 8; ++j) { + const int b = c + 8 * t + j < limit && (u[j] & fixed) == prefix ? (int)((u[j] >> shift) & 0xFF) : 256; + if (__any_sync(0xffffffffu, b < 256)) { + const unsigned same = __match_any_sync(0xffffffffu, b); + if (b < 256 && (same & below) == 0) atomicAdd(&hist[b], __popc(same)); + } + } + } + __syncthreads(); + if (warp == 0) { + // Lane l holds bins 255 - 8 l .. 248 - 8 l; the digit is in the first lane whose count reaches need, else 0. + int h[8], own = 0; + for (int j = 0; j < 8; ++j) own += h[j] = hist[255 - 8 * lane - j]; + int upto = own; + for (int o = 1; o < 32; o <<= 1) { + const int v = __shfl_up_sync(0xffffffffu, upto, o); + if (lane >= o) upto += v; + } + const unsigned hit = __ballot_sync(0xffffffffu, upto >= need); + if (lane == (hit ? __ffs(hit) - 1 : 31)) { + int above = upto - own, digit = 0; + for (int j = 0; j < 8; ++j) { + if (above + h[j] >= need) { + digit = 255 - 8 * lane - j; + break; + } + above += h[j]; + } + sh_need = need - above; + sh_prefix = prefix | ((uint32_t)digit << shift); + } + } + __syncthreads(); + need = sh_need; + prefix = sh_prefix; + fixed |= 0xFFu << shift; + __syncthreads(); + } + int run_take = 0, run_eq = 0; + for (int64_t c = 0; c < limit; c += 1024) { + uint32_t u[8]; + keys(c, u); + bool eq[8], gt[8]; + int eqs = 0; + for (int j = 0; j < 8; ++j) { + const bool in = c + 8 * t + j < limit; + eq[j] = in && u[j] == prefix; + gt[j] = in && u[j] > prefix; + eqs += eq[j]; + } + int all_eq, all_take; + int seen = run_eq + glm_index::block_scan(eqs, wsum, 0, &all_eq); + bool take[8]; + int takes = 0; + for (int j = 0; j < 8; ++j) { + take[j] = gt[j] || (eq[j] && seen < need); + seen += eq[j]; + takes += take[j]; + } + int at = run_take + glm_index::block_scan(takes, wsum, 1, &all_take); + for (int j = 0; j < 8; ++j) + if (take[j]) out[r * k + at++] = c + 8 * t + j; + run_eq += all_eq; + run_take += all_take; + } +} diff --git a/zig/kernels/cuda/glm_glue_latent.cuh b/zig/kernels/cuda/glm_glue_latent.cuh new file mode 100644 index 0000000000..12afc2daa3 --- /dev/null +++ b/zig/kernels/cuda/glm_glue_latent.cuh @@ -0,0 +1,74 @@ +// latent.py's BF16 absorb, expand and cache write by hand (the MLA layers' 512-wide latent, 256-wide heads). +#pragma once + +#include "glm_triton.cuh" + +using namespace glm_triton; + +// latent._absorb_q(_rows), block (head, 32 latent columns, row block) of 128 threads; rows per block change no bits. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_absorb_q(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, int64_t rows, int64_t heads, + int64_t rb) { + constexpr int D = 256, LW = 512; + __shared__ float smem[128]; + const int t = threadIdx.x, k0 = t >> 2, nq = t & 3, warp = t >> 5; + const int64_t h = blockIdx.x, n0 = (int64_t)blockIdx.y * 32 + nq * 8, r0 = blockIdx.z * rb; + float w[8][8]; + for (int rep = 0; rep < 8; ++rep) + for (int j = 0; j < 8; ++j) w[rep][j] = bf(wk[(h * D + k0 + 32 * rep) * LW + n0 + j]); + for (int64_t r = r0; r < r0 + rb && r < rows; ++r) { + float qv[8]; + for (int rep = 0; rep < 8; ++rep) qv[rep] = bf(q[(r * heads + h) * D + k0 + 32 * rep]); + for (int j = 0; j < 8; ++j) { + float a = fmul(w[1][j], qv[1]); + a = __fmaf_rn(w[0][j], qv[0], a); + for (int rep = 2; rep < 8; ++rep) a = __fmaf_rn(w[rep][j], qv[rep], a); + for (int o = 16; o >= 4; o >>= 1) a = a + __shfl_xor_sync(0xffffffffu, a, o); + if ((t & 28) == 0) smem[(nq * 8 + j) * 4 + warp] = a; + } + __syncthreads(); + float v = smem[t]; + v = v + __shfl_xor_sync(0xffffffffu, v, 2); + v = v + __shfl_xor_sync(0xffffffffu, v, 1); + if ((t & 3) == 0) qa[(r * heads + h) * LW + (int64_t)blockIdx.y * 32 + t / 4] = __float2bfloat16_rn(v); + __syncthreads(); + } +} + +// latent._expand_v(_rows), block (head, 16 outputs, row block) of 128 threads: o . wv by fma, lanes xor 16 .. 1. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_expand_v(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, int64_t rows, + int64_t heads, int64_t rb) { + constexpr int DV = 256, LW = 512; + __shared__ float smem[32]; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5, kb = (warp & 1) * 256 + lane * 8, nb = warp >> 1; + const int64_t h = blockIdx.x, n0 = (int64_t)blockIdx.y * 16, r0 = blockIdx.z * rb; + float w[8][8]; + for (int rep = 0; rep < 8; ++rep) + for (int j = 0; j < 8; ++j) w[rep][j] = bf(wv[(h * DV + n0 + nb + 2 * rep) * LW + kb + j]); + for (int64_t r = r0; r < r0 + rb && r < rows; ++r) { + float o[8]; + for (int j = 0; j < 8; ++j) o[j] = bf(ol[(r * heads + h) * LW + kb + j]); + for (int rep = 0; rep < 8; ++rep) { + float a = fmul(w[rep][1], o[1]); + a = __fmaf_rn(w[rep][0], o[0], a); + for (int j = 2; j < 8; ++j) a = __fmaf_rn(w[rep][j], o[j], a); + for (int s = 16; s > 0; s >>= 1) a = a + __shfl_xor_sync(0xffffffffu, a, s); + if (lane == 0) smem[(nb + 2 * rep) * 2 + (warp & 1)] = a; + } + __syncthreads(); + if (t < 32) { + float v = smem[t]; + v = v + __shfl_xor_sync(0xffffffffu, v, 1); + if ((t & 1) == 0) out[(r * heads + h) * DV + n0 + t / 2] = __float2bfloat16_rn(v); + } + __syncthreads(); + } +} + +// latent._lat_write (BF16 cache): row r of `lat` into cache slot pos + r. +extern "C" __global__ void tf_glm_lat_write(const __nv_bfloat16* lat, int64_t lat_stride, __nv_bfloat16* cache, + const int* pos, int64_t lw) { + const int64_t r = blockIdx.x, at = (int64_t)pos[0] + r; + for (int64_t k = threadIdx.x; k < lw; k += blockDim.x) cache[at * lw + k] = lat[r * lat_stride + k]; +} diff --git a/zig/kernels/cuda/glm_glue_moe.cuh b/zig/kernels/cuda/glm_glue_moe.cuh new file mode 100644 index 0000000000..7f2a27d6b0 --- /dev/null +++ b/zig/kernels/cuda/glm_glue_moe.cuh @@ -0,0 +1,113 @@ +// GLM's MLP and MoE glue by hand: glue._swiglu, glue._topk and qmm._group_sums (the router's dot is glm_glue_router). +#pragma once + +#include "glm_triton.cuh" + +using namespace glm_triton; + +namespace glm_moe { + +// A thread's S consecutive values summed from the second, then the butterfly over the 64 / S lanes of their group. +template +__device__ __forceinline__ float group64(const float* q) { + float g = q[0] + q[1]; + for (int j = 2; j < S; ++j) g = q[j] + g; + for (int o = 32 / S; o > 0; o >>= 1) g = g + __shfl_xor_sync(0xffffffffu, g, o); + return g; +} + +__device__ __forceinline__ float swiglu1(float g, float u, float limit) { + g = fminf(g, limit); + u = fminf(fmaxf(u, 0.0f - limit), limit); + const float s = rbf(div_full(g, exp_approx(0.0f - g) + 1.0f)); + return rbf(__fmul_rn(u, s)); +} + +} // namespace glm_moe + +// glue._swiglu with BLOCK 512 (128 threads, 4 values each, 64-group sums over 16 lanes), column block blockIdx.y. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_swiglu_512(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { + const int64_t r = blockIdx.x, d0 = blockIdx.y * 512 + threadIdx.x * 4; + const uint2 gv = *reinterpret_cast(gu + r * 2 * w + d0); + const uint2 uv = *reinterpret_cast(gu + r * 2 * w + w + d0); + const __nv_bfloat16* gh = reinterpret_cast(&gv); + const __nv_bfloat16* uh = reinterpret_cast(&uv); + float a[4]; + uint2 ov; + __nv_bfloat16* oh = reinterpret_cast<__nv_bfloat16*>(&ov); + for (int j = 0; j < 4; ++j) { + a[j] = glm_moe::swiglu1(bf(gh[j]), bf(uh[j]), limit); + oh[j] = __float2bfloat16_rn(a[j]); + } + *reinterpret_cast(out + r * w + d0) = ov; + const float g = glm_moe::group64<4>(a); + if ((threadIdx.x & 15) == 0) xs[r * (w / 64) + d0 / 64] = g; +} + +// glue._swiglu with BLOCK 64 (the TP=3 shared expert's 704): one value a thread, the warp's butterfly, then 2 warps. +extern "C" __global__ void __launch_bounds__(64) + tf_glm_swiglu_64(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { + __shared__ float smem[2]; + const int64_t r = blockIdx.x, d = blockIdx.y * 64 + threadIdx.x; + const float a = glm_moe::swiglu1(bf(gu[r * 2 * w + d]), bf(gu[r * 2 * w + w + d]), limit); + out[r * w + d] = __float2bfloat16_rn(a); + const float g = block_sum<2>(a, smem); + if (threadIdx.x == 0) xs[r * (w / 64) + blockIdx.y] = g; +} + +// qmm._group_sums: row blockIdx.x's 64-input sums, 8 lanes a group of 8 consecutive inputs each, 16 groups a block. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_group_sums(const __nv_bfloat16* x, float* xs, int64_t x_stride, int64_t kg) { + const int64_t m = blockIdx.x, g = blockIdx.y * 16 + threadIdx.x / 8; + float q[8]; + for (int j = 0; j < 8; ++j) q[j] = g < kg ? bf(x[m * x_stride + g * 64 + (threadIdx.x & 7) * 8 + j]) : 0.0f; + const float s = glm_moe::group64<8>(q); + if ((threadIdx.x & 7) == 0 && g < kg) xs[m * kg + g] = s; +} + +// glue._topk, a warp a row: top_k by sigmoid + bias (ties to lower id), normed if `norm`, x scale; then shared expert. +extern "C" __global__ void __launch_bounds__(32) + tf_glm_topk(const float* logits, const float* bias, int* pick, float* wts, float scale, int ne, int top_k, + int norm) { + constexpr int MAXE = 16; // 512 experts a row + const int64_t r = blockIdx.x; + const int lane = threadIdx.x; + float score[MAXE], choice[MAXE]; + for (int i = 0; i < MAXE; ++i) { + const int e = i * 32 + lane; + score[i] = e < ne ? div_full(1.0f, exp_approx(0.0f - logits[r * ne + e]) + 1.0f) : 0.0f; + choice[i] = e < ne ? score[i] + bias[e] : -INFINITY; + } + float total = 0.0f, mine = 0.0f; + int mine_id = 0; + for (int k = 0; k < top_k; ++k) { + float mx = choice[0]; + for (int i = 1; i < MAXE; ++i) mx = fmaxf(mx, choice[i]); + for (int o = 16; o > 0; o >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o)); + int id = 1 << 30; + for (int i = MAXE - 1; i >= 0; --i) + if (choice[i] == mx) id = i * 32 + lane; + for (int o = 16; o > 0; o >>= 1) id = min(id, __shfl_xor_sync(0xffffffffu, id, o)); + float sk = 0.0f; + for (int i = 0; i < MAXE; ++i) + if (i * 32 + lane == id) { + sk = score[i]; + choice[i] = -INFINITY; + } + sk = __shfl_sync(0xffffffffu, sk, id & 31); + total = total + sk; + if (lane == k) { + mine = sk; + mine_id = id; + } + } + if (lane < top_k) { + const float v = norm ? div_full(mine, total + 1e-20f) : mine; + pick[r * (top_k + 1) + lane] = mine_id; + wts[r * (top_k + 1) + lane] = __fmul_rn(v, scale); + } else if (lane == top_k) { + pick[r * (top_k + 1) + lane] = ne; + wts[r * (top_k + 1) + lane] = 1.0f; + } +} diff --git a/zig/kernels/cuda/glm_glue_norm.cuh b/zig/kernels/cuda/glm_glue_norm.cuh new file mode 100644 index 0000000000..c9d74c57f3 --- /dev/null +++ b/zig/kernels/cuda/glm_glue_norm.cuh @@ -0,0 +1,122 @@ +// glue._rmsnorm by hand (glm_glue.cu's first group): see glm_triton.cuh for the Triton ops it keeps. +#pragma once + +#include "glm_triton.cuh" + +namespace glm_glue { +using namespace glm_triton; + +// S consecutive bf16 (S 4 or 8) moved in one load or store of 2 S bytes when `vec`, else two bytes at a time. +template +struct Pack; +template <> +struct Pack<4> { + using T = uint2; +}; +template <> +struct Pack<8> { + using T = uint4; +}; +template +__device__ __forceinline__ bool packed(const void* p, int64_t stride) { + return ((reinterpret_cast(p) | static_cast(stride * 2)) & (2 * S - 1)) == 0; +} +template +__device__ __forceinline__ void load_s(const __nv_bfloat16* p, float* v, bool vec) { + if (vec) { + const typename Pack::T u = *reinterpret_cast::T*>(p); + const __nv_bfloat16* h = reinterpret_cast(&u); +#pragma unroll + for (int j = 0; j < S; ++j) v[j] = bf(h[j]); + } else { +#pragma unroll + for (int j = 0; j < S; ++j) v[j] = bf(p[j]); + } +} +template +__device__ __forceinline__ void store_s(__nv_bfloat16* p, const __nv_bfloat16* o, bool vec) { + if (vec) { + typename Pack::T u; + __nv_bfloat16* h = reinterpret_cast<__nv_bfloat16*>(&u); +#pragma unroll + for (int j = 0; j < S; ++j) h[j] = o[j]; + *reinterpret_cast::T*>(p) = u; + } else { +#pragma unroll + for (int j = 0; j < S; ++j) p[j] = o[j]; + } +} + +// glue._rmsnorm, BLOCK = D to a power of two, NW warps, S values a thread: bf16(w * bf16(x / rms)), `sums` per 64. +template +__device__ void rmsnorm(const __nv_bfloat16* x, int64_t x_stride, const __nv_bfloat16* w, __nv_bfloat16* out, + int64_t o_stride, float* xs, float eps, int sums) { + constexpr int T = NW * 32, TILE = S * T, REPS = BLOCK / TILE, N = REPS * S, G = 64 / S; + static_assert(D % S == 0, "a thread's values all inside D or all past it"); + __shared__ float smem[NW]; + const int64_t r = blockIdx.x; + const int t = threadIdx.x; + const bool xv = packed(x, x_stride), wv = packed(w, 0), ov = packed(out, o_stride); + float v[N]; +#pragma unroll + for (int rep = 0; rep < REPS; ++rep) { + const int d = rep * TILE + t * S; + if (d < D) { + load_s(x + r * x_stride + d, v + rep * S, xv); + } else { +#pragma unroll + for (int j = 0; j < S; ++j) v[rep * S + j] = 0.0f; + } + } + float acc = __fmul_rn(v[1], v[1]); + acc = __fmaf_rn(v[0], v[0], acc); +#pragma unroll + for (int k = 2; k < N; ++k) acc = __fmaf_rn(v[k], v[k], acc); + const float total = block_sum(acc, smem); + const float rinv = div_full(1.0f, sqrt_approx(eps + div_full(total, (float)D))); + float y[N]; +#pragma unroll + for (int rep = 0; rep < REPS; ++rep) { + const int d = rep * TILE + t * S; + if (d >= D) { +#pragma unroll + for (int j = 0; j < S; ++j) y[rep * S + j] = 0.0f; + continue; + } + float wf[S]; + __nv_bfloat16 o[S]; + load_s(w + d, wf, wv); +#pragma unroll + for (int j = 0; j < S; ++j) { + o[j] = __float2bfloat16_rn(__fmul_rn(wf[j], rbf(__fmul_rn(rinv, v[rep * S + j])))); + y[rep * S + j] = bf(o[j]); + } + store_s(out + r * o_stride + d, o, ov); + } + if (!sums) return; +#pragma unroll + for (int rep = 0; rep < REPS; ++rep) { + const float* q = y + rep * S; + float g = q[1]; + g = q[0] + g; +#pragma unroll + for (int j = 2; j < S; ++j) g = q[j] + g; + for (int o = G / 2; o > 0; o >>= 1) g += __shfl_xor_sync(0xffffffffu, g, o); + const int group = (rep * TILE + t * S) / 64; + if (t % G == 0 && group < D / 64) xs[r * (D / 64) + group] = g; + } +} + +} // namespace glm_glue + +#define RMSNORM(D, BLOCK, NW, S) \ + extern "C" __global__ void __launch_bounds__(NW * 32) \ + tf_glm_rmsnorm_##D(const __nv_bfloat16* x, int64_t x_stride, const __nv_bfloat16* w, __nv_bfloat16* out, \ + int64_t o_stride, float* xs, float eps, int sums) { \ + glm_glue::rmsnorm(x, x_stride, w, out, o_stride, xs, eps, sums); \ + } + +// The model's norms: hidden 4096 (8 warps), q_lora 1536 and kv_lora 512 (4 warps), as glue.rmsnorm launches them. +RMSNORM(4096, 4096, 8, 8) +RMSNORM(1536, 2048, 4, 8) +RMSNORM(512, 512, 4, 4) diff --git a/zig/kernels/cuda/glm_glue_router.cuh b/zig/kernels/cuda/glm_glue_router.cuh new file mode 100644 index 0000000000..02435cef64 --- /dev/null +++ b/zig/kernels/cuda/glm_glue_router.cuh @@ -0,0 +1,50 @@ +// glue._router_part by hand: tl.dot's chain of mma.m16n8k16 (bf16 in, fp32 accumulate) from zero, K ascending. +#pragma once + +#include "glm_triton.cuh" + +namespace glm_router { + +// Two consecutive bf16 as one mma register (the lower K in the low half), zero outside the rows or columns. +__device__ __forceinline__ uint32_t pair(const __nv_bfloat16* p, bool ok) { + if (!ok) return 0u; + return (uint32_t)__bfloat16_as_ushort(p[0]) | ((uint32_t)__bfloat16_as_ushort(p[1]) << 16); +} + + +__device__ __forceinline__ void mma16816(float* c, const uint32_t* a, const uint32_t* b) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +} // namespace glm_router + +// glue._router_part: part[s][m][e] = x[m] . w[e] over K slice s of ks (d / ks inputs), each warp one 16 x 8 tile. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_router_part(const __nv_bfloat16* x, const __nv_bfloat16* w, float* part, int64_t rows, int64_t x_stride, + int64_t d, int64_t ne, int ks) { + const int lane = threadIdx.x & 31, g = lane >> 2, q = lane & 3; + const int64_t m0 = (int64_t)blockIdx.x * 16, n = (int64_t)blockIdx.y * 32 + (threadIdx.x >> 5) * 8; + const int64_t s = blockIdx.z, span = d / ks, lo = s * span; + const int64_t ra = m0 + g, rb = ra + 8, e = n + g; + const bool oka = ra < rows, okb = rb < rows, oke = e < ne; + float c[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int64_t k = lo; k < lo + span; k += 16) { + const int64_t kq = k + q * 2; + const uint32_t a[4] = {glm_router::pair(x + ra * x_stride + kq, oka), glm_router::pair(x + rb * x_stride + kq, okb), + glm_router::pair(x + ra * x_stride + kq + 8, oka), + glm_router::pair(x + rb * x_stride + kq + 8, okb)}; + const uint32_t b[2] = {glm_router::pair(w + e * d + kq, oke), glm_router::pair(w + e * d + kq + 8, oke)}; + glm_router::mma16816(c, a, b); + } + for (int h = 0; h < 2; ++h) { + const int64_t r = h == 0 ? ra : rb; + for (int j = 0; j < 2; ++j) { + const int64_t col = n + q * 2 + j; + if (r < rows && col < ne) part[(s * rows + r) * ne + col] = c[h * 2 + j]; + } + } +} diff --git a/zig/kernels/cuda/glm_glue_rows.cuh b/zig/kernels/cuda/glm_glue_rows.cuh new file mode 100644 index 0000000000..5928630754 --- /dev/null +++ b/zig/kernels/cuda/glm_glue_rows.cuh @@ -0,0 +1,228 @@ +// GLM's element-wise glue by hand: each output value alone, in the Triton kernel's order of fp32 ops. +#pragma once + +#include "glm_triton.cuh" + +namespace glm_glue { +using namespace glm_triton; + +// Column c of a 2-D grid's row blockIdx.x, 256 threads a block along y. +__device__ __forceinline__ int64_t col() { return (int64_t)blockIdx.y * blockDim.x + threadIdx.x; } + +// The first of a thread's 8 contiguous columns: Triton's BLOCK 1024 over 4 warps, 128 threads a block along y. +__device__ __forceinline__ int64_t col8() { return ((int64_t)blockIdx.y * blockDim.x + threadIdx.x) * 8; } + +// 8 contiguous values in one 16-byte load (bf16) or two (fp32), and 8 bf16 out in one 16-byte store. +__device__ __forceinline__ void ld8(const __nv_bfloat16* p, float v[8]) { + const uint4 u = *reinterpret_cast(p); + const __nv_bfloat16* h = reinterpret_cast(&u); + for (int j = 0; j < 8; ++j) v[j] = bf(h[j]); +} +__device__ __forceinline__ void ld8(const float* p, float v[8]) { + const float4 a = *reinterpret_cast(p), b = *reinterpret_cast(p + 4); + v[0] = a.x, v[1] = a.y, v[2] = a.z, v[3] = a.w, v[4] = b.x, v[5] = b.y, v[6] = b.z, v[7] = b.w; +} +__device__ __forceinline__ void st8(__nv_bfloat16* p, const float v[8]) { + uint4 u; + __nv_bfloat16* h = reinterpret_cast<__nv_bfloat16*>(&u); + for (int j = 0; j < 8; ++j) h[j] = __float2bfloat16_rn(v[j]); + *reinterpret_cast(p) = u; +} + +// The ranks' fp32 partials at i .. i + 7, rank 0 first and `rs` apart. +__device__ __forceinline__ void ranks8(const float* g, int64_t rs, int64_t i, int world, float a[8]) { + ld8(g + i, a); + for (int k = 1; k < world; ++k) { + float b[8]; + ld8(g + k * rs + i, b); + for (int j = 0; j < 8; ++j) a[j] = a[j] + b[j]; + } +} + +} // namespace glm_glue + +using namespace glm_triton; + +// glue._residual_add: xout = bf16(x + bf16(the partials summed rank 0 first)), 8 columns a thread. +extern "C" __global__ void tf_glm_residual_add(const __nv_bfloat16* x, __nv_bfloat16* xout, const float* g, int64_t rs, + int64_t d, int world) { + const int64_t c = glm_glue::col8(); + if (c >= d) return; + const int64_t i = blockIdx.x * d + c; + float a[8], v[8]; + glm_glue::ranks8(g, rs, i, world, a); + glm_glue::ld8(x + i, v); + for (int j = 0; j < 8; ++j) v[j] = v[j] + rbf(a[j]); + glm_glue::st8(xout + i, v); +} + +// glue._stream_mean: bf16 of the s streams' sum over s (div.full, as Triton divides), 8 columns a thread. +extern "C" __global__ void tf_glm_stream_mean(const __nv_bfloat16* x, __nv_bfloat16* out, int64_t d, int s) { + const int64_t c = glm_glue::col8(); + if (c >= d) return; + const __nv_bfloat16* row = x + blockIdx.x * s * d + c; + float a[8]; + glm_glue::ld8(row, a); + for (int k = 1; k < s; ++k) { + float b[8]; + glm_glue::ld8(row + k * d, b); + for (int j = 0; j < 8; ++j) a[j] = a[j] + b[j]; + } + for (int j = 0; j < 8; ++j) a[j] = div_full(a[j], (float)s); + glm_glue::st8(out + blockIdx.x * d + c, a); +} + +// glue._rank_sum: bf16 of the partials summed rank 0 first, `me` from `own`, others from `peers`, 8 columns a thread. +extern "C" __global__ void tf_glm_rank_sum(const float* own, const float* peers, __nv_bfloat16* out, int64_t rs, + int64_t d, int world, int me) { + const int64_t c = glm_glue::col8(); + if (c >= d) return; + const int64_t i = blockIdx.x * d + c; + float a[8]; + glm_glue::ld8(me == 0 ? own + i : peers + i, a); + for (int k = 1; k < world; ++k) { + float b[8]; + glm_glue::ld8(k == me ? own + i : peers + (k > me ? k - 1 : k) * rs + i, b); + for (int j = 0; j < 8; ++j) a[j] = a[j] + b[j]; + } + glm_glue::st8(out + i, a); +} + +// glue._hc_post: stream s = bf16(post[s] branch + sum_t comb[t][s] x[t]) in LLVM's fma order; x may be xout. +extern "C" __global__ void tf_glm_hc_post(const __nv_bfloat16* x, __nv_bfloat16* xout, const float* g, + const float* post, const float* comb, int64_t rs, int64_t d, int world) { + const int64_t c = glm_glue::col8(); + if (c >= d) return; + const int64_t r = blockIdx.x; + float branch[8]; + glm_glue::ranks8(g, rs, r * d + c, world, branch); + float xv[4][8]; + for (int t = 0; t < 4; ++t) glm_glue::ld8(x + r * 4 * d + t * d + c, xv[t]); + for (int s = 0; s < 4; ++s) { + const float* cm = comb + r * 16 + s; + const float c0 = cm[0], c1 = cm[4], c2 = cm[8], c3 = cm[12], p = post[r * 4 + s]; + float o[8]; + for (int j = 0; j < 8; ++j) { + float m = __fmul_rn(xv[1][j], c1); + m = __fmaf_rn(xv[0][j], c0, m); + m = __fmaf_rn(xv[2][j], c2, m); + m = __fmaf_rn(xv[3][j], c3, m); + o[j] = __fmaf_rn(rbf(branch[j]), p, m); + } + glm_glue::st8(xout + r * 4 * d + s * d + c, o); + } +} + +// glue._combine: fp32 sum over slots in order of y[r][k] wts[r][k] (fma from 0), y fp32 or bf16, 8 columns a thread. +extern "C" __global__ void tf_glm_combine( + const void* y, int y_f32, const float* wts, float* out, int64_t d, int slots) { + const int64_t c = glm_glue::col8(); + if (c >= d) return; + const int64_t r = blockIdx.x; + float a[8] = {}; + for (int k = 0; k < slots; ++k) { + const int64_t i = (r * slots + k) * d + c; + float v[8]; + if (y_f32) + glm_glue::ld8(static_cast(y) + i, v); + else + glm_glue::ld8(static_cast(y) + i, v); + const float w = wts[r * slots + k]; + for (int j = 0; j < 8; ++j) a[j] = __fmaf_rn(v[j], w, a[j]); + } + float* o = out + r * d + c; + *reinterpret_cast(o) = make_float4(a[0], a[1], a[2], a[3]); + *reinterpret_cast(o + 4) = make_float4(a[4], a[5], a[6], a[7]); +} + +// glue._router_sum and qmm._reduce: ks slices of `total` fp32 added in order (fp32 or bf16 out), 8 values a thread. +extern "C" __global__ void tf_glm_slice_sum(const float* part, void* out, int64_t total, int ks, int f32) { + const int64_t i0 = ((int64_t)blockIdx.x * blockDim.x + threadIdx.x) * 8; + if (i0 >= total) return; + const bool aligned = ((reinterpret_cast(out) | reinterpret_cast(part)) & 15) == 0; + if (i0 + 8 <= total && (total & 3) == 0 && aligned) { + float a[8]; + glm_glue::ld8(part + i0, a); + for (int s = 1; s < ks; ++s) { + float b[8]; + glm_glue::ld8(part + s * total + i0, b); + for (int j = 0; j < 8; ++j) a[j] = a[j] + b[j]; + } + if (f32) { + float* o = static_cast(out) + i0; + *reinterpret_cast(o) = make_float4(a[0], a[1], a[2], a[3]); + *reinterpret_cast(o + 4) = make_float4(a[4], a[5], a[6], a[7]); + } else { + glm_glue::st8(static_cast<__nv_bfloat16*>(out) + i0, a); + } + return; + } + for (int64_t i = i0; i < i0 + 8 && i < total; ++i) { + float a = part[i]; + for (int s = 1; s < ks; ++s) a = a + part[s * total + i]; + if (f32) + static_cast(out)[i] = a; + else + static_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(a); + } +} + +// glue._embed_b16: token ids[r]'s bf16 row into `copies` streams of d. +extern "C" __global__ void tf_glm_embed_b16(const int* ids, const __nv_bfloat16* w, __nv_bfloat16* out, int64_t d, + int copies) { + const int64_t c = glm_glue::col(); + if (c >= d) return; + const int64_t r = blockIdx.x; + const __nv_bfloat16 v = w[(int64_t)ids[r] * d + c]; + for (int k = 0; k < copies; ++k) out[(r * copies + k) * d + c] = v; +} + +// forward._conv_shift: layer blockIdx.x's 3-row window takes rows keep .. keep + 2 of [window; projection rows]. +extern "C" __global__ void tf_glm_conv_shift(__nv_bfloat16* conv, const __nv_bfloat16* proj, int64_t keep, + int64_t conv_layer, int64_t proj_layer, int64_t proj_row, int64_t c) { + const int64_t k = glm_glue::col(); + if (k >= c) return; + const int64_t l = blockIdx.x; + __nv_bfloat16 v[3]; + for (int j = 0; j < 3; ++j) { + const int64_t src = keep + j; + v[j] = src < 3 ? conv[l * conv_layer + src * c + k] : proj[l * proj_layer + (src - 3) * proj_row + k]; + } + for (int j = 0; j < 3; ++j) conv[l * conv_layer + j * c + k] = v[j]; +} + +// decode.sample_rows' greedy pick, as the host scan: first max of n bf16 logits, NaN never wins; out fp32 bits, lo+id. +extern "C" __global__ void __launch_bounds__(1024) tf_glm_top1(const unsigned short* logits, int64_t n, unsigned lo, + unsigned* out) { + __shared__ float sv[1024]; + __shared__ unsigned si[1024]; + const int t = threadIdx.x; + float top = -__int_as_float(0x7f800000); + unsigned best = 0xffffffffu; + for (int64_t i = t; i < n; i += 1024) { + const float v = __int_as_float(static_cast(static_cast(logits[i]) << 16)); + if (v > top) top = v, best = static_cast(i); + } + sv[t] = top, si[t] = best; + __syncthreads(); + for (int h = 512; h > 0; h >>= 1) { + if (t < h) { + const float a = sv[t], b = sv[t + h]; + const unsigned ia = si[t], ib = si[t + h]; + if (b > a || (b == a && ib < ia)) sv[t] = b, si[t] = ib; + } + __syncthreads(); + } + if (t == 0) { + out[0] = __float_as_uint(sv[0]); + out[1] = lo + (si[0] == 0xffffffffu ? 0u : si[0]); + } +} + +// The leader's step for the other ranks, written to the link's buffer from launch arguments (cuda_follow's Link.send). +struct tf_glm_words { + unsigned w[96]; +}; +extern "C" __global__ void tf_glm_put_words(unsigned* dst, unsigned n, tf_glm_words v) { + for (unsigned i = threadIdx.x; i < n; i += blockDim.x) dst[i] = v.w[i]; +} diff --git a/zig/kernels/cuda/glm_kda.cu b/zig/kernels/cuda/glm_kda.cu new file mode 100644 index 0000000000..4d4eb8349a --- /dev/null +++ b/zig/kernels/cuda/glm_kda.cu @@ -0,0 +1,434 @@ +// GLM-5.3-Flash's KDA chain and replay: the Python family's kda.cu device code; -O3 --fmad=false. + +#include +#include +#include +namespace { + +constexpr int DK = 128, DV = 128, TAPS = 4; + +__device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { + for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o); + return x; +} + +__device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); } + +// One delta-rule step on this thread's 4 x 4 block of the state: per-channel decay, read (k), correct toward v. +__device__ __forceinline__ void update(float (&s)[4][4], const float (&kk)[4], const float (&gg)[4], + const float* vrow, int warp, float beta) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[j][i] = s[j][i] * gg[i]; + kv = kv + s[j][i] * kk[i]; + } + kv = warp_sum(kv); + const float delta = (vrow[warp * 4 + j] - kv) * beta; +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta; + } +} + +__device__ __forceinline__ void chain_kernel_body( + int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, + const __nv_bfloat16* __restrict__ A, int a_stride, const __nv_bfloat16* __restrict__ G, int g_stride, + const __nv_bfloat16* __restrict__ cs, const float* __restrict__ cw, + const float* __restrict__ state_in, const float* __restrict__ a_log, const float* __restrict__ dt_bias, + const __nv_bfloat16* __restrict__ norm_w, float eps, float lower, int rows, + __nv_bfloat16* __restrict__ out, float* __restrict__ state_out, + float* __restrict__ k_save, __nv_bfloat16* __restrict__ v_save, float* __restrict__ g_save, + float* __restrict__ b_save) { + const int C = 3 * H * DK; // conv channels: q | k | v + const int h = blockIdx.x; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float qs[DK], ks[DK], vs[DV], ys[DV], gs[DK]; + __shared__ float beta_s, rinv; + int c = -1; + if (t < 3 * DK) c = (t / DK) * H * DK + h * DK + (t % DK); + float s[4][4]; + const size_t sbase = (size_t)h * DV * DK; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i]; + const float decay_rate = expf(a_log[h]); + for (int r = 0; r < rows; ++r) { + if (c >= 0) { + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) { + const int at = r + tap; + const float x = at < TAPS - 1 ? __bfloat162float(cs[(size_t)at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * p_stride + c]); + acc = acc + cw[(size_t)c * TAPS + tap] * x; + } + const float act = bf(acc / (1.0f + expf(-acc))); + if (t < DK) qs[t] = act; + else if (t < 2 * DK) ks[t - DK] = act; + else vs[t - 2 * DK] = act; + } else if (t >= 512 && t < 512 + DK) { + const int i = t - 512; + const float a = __bfloat162float(A[(size_t)r * a_stride + h * DK + i]) + dt_bias[h * DK + i]; + gs[i] = expf(lower * sigmoidf_(decay_rate * a)); + } else if (t == 1023) { + beta_s = bf(sigmoidf_(__bfloat162float(P[(size_t)r * p_stride + b_off + h]))); + } + __syncthreads(); + if (warp < 2) { + float* x = warp == 0 ? qs : ks; + float v4[4], ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { v4[i] = x[lane * 4 + i]; ss = ss + v4[i] * v4[i]; } + ss = warp_sum(ss); + float inv = 1.0f / sqrtf(ss + 1e-6f); + __syncwarp(); +#pragma unroll + for (int i = 0; i < 4; ++i) { + float y = v4[i] * inv; + if (warp == 0) y = y * (1.0f / sqrtf((float)DK)); + x[lane * 4 + i] = y; + } + } + __syncthreads(); + const float beta = beta_s; + float kk[4], qq[4], gg[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { kk[i] = ks[lane * 4 + i]; qq[i] = qs[lane * 4 + i]; gg[i] = gs[lane * 4 + i]; } + update(s, kk, gg, vs, warp, beta); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i]; + o = warp_sum(o); + if (lane == 0) ys[warp * 4 + j] = bf(o); + } + if (k_save != nullptr) { + const size_t base = ((size_t)r * H + h) * DK; + if (t < DK) { k_save[base + t] = ks[t]; g_save[base + t] = gs[t]; } + if (t >= DK && t < DK + DV) v_save[base + t - DK] = __float2bfloat16_rn(vs[t - DK]); + if (t == 0) b_save[r * H + h] = beta; + } + __syncthreads(); + if (warp == 0) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { const float y = ys[lane * 4 + i]; ss = ss + y * y; } + ss = warp_sum(ss); + if (lane == 0) rinv = 1.0f / sqrtf(ss / (float)DV + eps); + } + __syncthreads(); + if (t < DV) { + const float yn = ys[t] * rinv; + const float yw = __bfloat162float(norm_w[t]) * yn; + const float gate = __bfloat162float(G[(size_t)r * g_stride + h * DV + t]); + out[(size_t)r * H * DV + h * DV + t] = __float2bfloat16_rn(yw * sigmoidf_(gate)); + } + __syncthreads(); + } + if (state_out != nullptr) { +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) state_out[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; + } +} + +__device__ __forceinline__ void replay_kernel_body( + int H, const float* __restrict__ state_in, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, + const float* __restrict__ b_save, int rows, float* __restrict__ state_out) { + const int h = blockIdx.x; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float vs[DV]; + float s[4][4]; + const size_t sbase = (size_t)h * DV * DK; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i]; + for (int r = 0; r < rows; ++r) { + const size_t base = ((size_t)r * H + h) * DK; + if (t < DV) vs[t] = __bfloat162float(v_save[base + t]); + __syncthreads(); + float kk[4], gg[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { kk[i] = k_save[base + lane * 4 + i]; gg[i] = g_save[base + lane * 4 + i]; } + update(s, kk, gg, vs, warp, b_save[r * H + h]); + __syncthreads(); + } +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) state_out[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; +} + +// All layers at once: block (layer, head); per-layer strides of the state buffers and saved rows. +__device__ __forceinline__ void replay_layers_kernel_body( + int H, const float* __restrict__ state_in, size_t state_stride, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, + size_t kv_stride, size_t b_stride, int rows, float* __restrict__ state_out) { + const int layer = blockIdx.x / H, h = blockIdx.x % H; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float vs[DV]; + float s[4][4]; + const float* sin = state_in + layer * state_stride; + float* sout = state_out + layer * state_stride; + const float* ks = k_save + layer * kv_stride; + const __nv_bfloat16* vsv = v_save + layer * kv_stride; + const float* gsv = g_save + layer * kv_stride; + const float* bsv = b_save + layer * b_stride; + const size_t sbase = (size_t)h * DV * DK; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = sin[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i]; + for (int r = 0; r < rows; ++r) { + const size_t base = ((size_t)r * H + h) * DK; + if (t < DV) vs[t] = __bfloat162float(vsv[base + t]); + __syncthreads(); + float kk[4], gg[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { kk[i] = ks[base + lane * 4 + i]; gg[i] = gsv[base + lane * 4 + i]; } + update(s, kk, gg, vs, warp, bsv[r * H + h]); + __syncthreads(); + } +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) sout[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; +} + +// Long windows: state-free work batched (prep_kernel, out_kernel), delta rule per row (step_kernel): chain_kernel bits. + +// Block (row, head), 512 threads: conv + SiLU of the 3 x 128 q|k|v channels, the decay, beta, then the q/k norms. +__device__ __forceinline__ void prep_kernel_body( + int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, + const __nv_bfloat16* __restrict__ A, int a_stride, const __nv_bfloat16* __restrict__ cs, + const float* __restrict__ cw, const float* __restrict__ a_log, const float* __restrict__ dt_bias, + float lower, float* __restrict__ q_out, float* __restrict__ k_save, __nv_bfloat16* __restrict__ v_save, + float* __restrict__ g_save, float* __restrict__ b_save) { + const int C = 3 * H * DK; + const int r = blockIdx.x, h = blockIdx.y; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float qs[DK], ks[DK]; + const size_t base = ((size_t)r * H + h) * DK; + if (t < 3 * DK) { + const int c = (t / DK) * H * DK + h * DK + (t % DK); + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) { + const int at = r + tap; + const float x = at < TAPS - 1 ? __bfloat162float(cs[(size_t)at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * p_stride + c]); + acc = acc + cw[(size_t)c * TAPS + tap] * x; + } + const float act = bf(acc / (1.0f + expf(-acc))); + if (t < DK) qs[t] = act; + else if (t < 2 * DK) ks[t - DK] = act; + else v_save[base + t - 2 * DK] = __float2bfloat16_rn(act); + } else { + const int i = t - 3 * DK; + const float decay_rate = expf(a_log[h]); + const float a = __bfloat162float(A[(size_t)r * a_stride + h * DK + i]) + dt_bias[h * DK + i]; + g_save[base + i] = expf(lower * sigmoidf_(decay_rate * a)); + if (i == 0) b_save[r * H + h] = bf(sigmoidf_(__bfloat162float(P[(size_t)r * p_stride + b_off + h]))); + } + __syncthreads(); + if (warp < 2) { + const float* x = warp == 0 ? qs : ks; + float v4[4], ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { v4[i] = x[lane * 4 + i]; ss = ss + v4[i] * v4[i]; } + ss = warp_sum(ss); + float inv = 1.0f / sqrtf(ss + 1e-6f); + float* dst = warp == 0 ? q_out : k_save; +#pragma unroll + for (int i = 0; i < 4; ++i) { + float y = v4[i] * inv; + if (warp == 0) y = y * (1.0f / sqrtf((float)DK)); + dst[base + lane * 4 + i] = y; + } + } +} + +// warp_sum of x[0..3] at once: lane l gets x[l >> 3]'s sum, warp_sum's lane pairs and bits, 6 shuffles not 20. +__device__ __forceinline__ float warp_sum4(const float (&x)[4], int lane) { + const bool hi = lane & 16, mid = lane & 8; + float a = hi ? x[2] : x[0], b = hi ? x[3] : x[1]; + const float ra = __shfl_xor_sync(0xffffffffu, hi ? x[0] : x[2], 16); + const float rb = __shfl_xor_sync(0xffffffffu, hi ? x[1] : x[3], 16); + a = a + ra; + b = b + rb; + float c = mid ? b : a; + c = c + __shfl_xor_sync(0xffffffffu, mid ? a : b, 8); + for (int o = 4; o; o >>= 1) c += __shfl_xor_sync(0xffffffffu, c, o); + return c; +} + +// ``update`` with its four warp sums through warp_sum4 (and broadcast back): the same bits. +__device__ __forceinline__ void update4(float (&s)[4][4], const float (&kk)[4], const float (&gg)[4], + const float (&vv)[4], float beta, int lane) { + float part[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[j][i] = s[j][i] * gg[i]; + kv = kv + s[j][i] * kk[i]; + } + part[j] = kv; + } + const float mine = warp_sum4(part, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float delta = (vv[j] - __shfl_sync(0xffffffffu, mine, 8 * j)) * beta; +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta; + } +} + +// Block (head, part): a warp owns 4 value rows; TR rows of saved k, q, decay, v, beta staged in smem at a time. +template +__device__ __forceinline__ void step_kernel_body( + int H, const float* __restrict__ state_in, const float* __restrict__ q_in, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, + int rows, __nv_bfloat16* __restrict__ y_out, float* __restrict__ state_out) { + constexpr int NT = WARPS * 32, VR = WARPS * 4; + const int h = blockIdx.x, v0 = blockIdx.y * VR; + const int lw = threadIdx.x >> 5, lane = threadIdx.x & 31, warp = blockIdx.y * WARPS + lw; + __shared__ __align__(16) float kt[TR][DK]; + __shared__ __align__(16) float qt[TR][DK]; + __shared__ __align__(16) float gt[TR][DK]; + __shared__ __align__(16) float vt[TR][VR]; + __shared__ float bt[TR]; + float s[4][4]; + const size_t sbase = (size_t)h * DV * DK; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i]; + for (int r0 = 0; r0 < rows; r0 += TR) { + const int n = rows - r0 < TR ? rows - r0 : TR; + __syncthreads(); + for (int x = threadIdx.x; x < n * DK; x += NT) { + const int rr = x / DK, c = x % DK; + const size_t base = ((size_t)(r0 + rr) * H + h) * DK + c; + kt[rr][c] = k_save[base]; + qt[rr][c] = q_in[base]; + gt[rr][c] = g_save[base]; + } + for (int x = threadIdx.x; x < n * VR; x += NT) { + const int rr = x / VR, c = x % VR; + vt[rr][c] = __bfloat162float(v_save[((size_t)(r0 + rr) * H + h) * DV + v0 + c]); + } + if (threadIdx.x < n) bt[threadIdx.x] = b_save[(r0 + threadIdx.x) * H + h]; + __syncthreads(); + for (int rr = 0; rr < n; ++rr) { + const float4 k4 = reinterpret_cast(kt[rr])[lane]; // one 16-byte load a lane: + const float4 q4 = reinterpret_cast(qt[rr])[lane]; // no bank conflicts + const float4 g4 = reinterpret_cast(gt[rr])[lane]; + const float4 v4 = reinterpret_cast(vt[rr])[lw]; + float kk[4] = {k4.x, k4.y, k4.z, k4.w}, qq[4] = {q4.x, q4.y, q4.z, q4.w}; + float gg[4] = {g4.x, g4.y, g4.z, g4.w}, vv[4] = {v4.x, v4.y, v4.z, v4.w}; + update4(s, kk, gg, vv, bt[rr], lane); + float part[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i]; + part[j] = o; + } + const float o = warp_sum4(part, lane); // read-out lane >> 3 + if ((lane & 7) == 0) + y_out[((size_t)(r0 + rr) * H + h) * DV + warp * 4 + (lane >> 3)] = __float2bfloat16_rn(o); + } + } + if (state_out != nullptr) { +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) state_out[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; + } +} + +// Block (row, head), 128 threads: the gated RMSNorm of the read-out. +__device__ __forceinline__ void out_kernel_body( + int H, const __nv_bfloat16* __restrict__ y_in, const __nv_bfloat16* __restrict__ G, int g_stride, + const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out) { + const int r = blockIdx.x, h = blockIdx.y; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float rinv; + const size_t base = ((size_t)r * H + h) * DV; + if (warp == 0) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { const float y = __bfloat162float(y_in[base + lane * 4 + i]); ss = ss + y * y; } + ss = warp_sum(ss); + if (lane == 0) rinv = 1.0f / sqrtf(ss / (float)DV + eps); + } + __syncthreads(); + const float yn = __bfloat162float(y_in[base + t]) * rinv; + const float yw = __bfloat162float(norm_w[t]) * yn; + const float gate = __bfloat162float(G[(size_t)r * g_stride + h * DV + t]); + out[(size_t)r * H * DV + h * DV + t] = __float2bfloat16_rn(yw * sigmoidf_(gate)); +} + +} // namespace + +extern "C" __global__ void __launch_bounds__(1024) tf_glm_kda_chain( + int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, const __nv_bfloat16* __restrict__ A, + int a_stride, const __nv_bfloat16* __restrict__ G, int g_stride, const __nv_bfloat16* __restrict__ cs, + const float* __restrict__ cw, const float* __restrict__ state_in, const float* __restrict__ a_log, + const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ norm_w, float eps, float lower, int rows, + __nv_bfloat16* __restrict__ out, float* __restrict__ state_out, float* __restrict__ k_save, + __nv_bfloat16* __restrict__ v_save, float* __restrict__ g_save, float* __restrict__ b_save) { + chain_kernel_body( + H, P, p_stride, b_off, A, a_stride, G, g_stride, cs, cw, state_in, a_log, dt_bias, norm_w, eps, lower, rows, + out, state_out, k_save, v_save, g_save, b_save); +} + +extern "C" __global__ void __launch_bounds__(1024) tf_glm_kda_replay( + int H, const float* __restrict__ state_in, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, + int rows, float* __restrict__ state_out) { + replay_kernel_body(H, state_in, k_save, v_save, g_save, b_save, rows, state_out); +} + +extern "C" __global__ void __launch_bounds__(1024) tf_glm_kda_replay_layers( + int H, const float* __restrict__ state_in, size_t state_stride, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, + size_t kv_stride, size_t b_stride, int rows, float* __restrict__ state_out) { + replay_layers_kernel_body( + H, state_in, state_stride, k_save, v_save, g_save, b_save, kv_stride, b_stride, rows, state_out); +} + +extern "C" __global__ void __launch_bounds__(512) tf_glm_kda_prep( + int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, const __nv_bfloat16* __restrict__ A, + int a_stride, const __nv_bfloat16* __restrict__ cs, const float* __restrict__ cw, + const float* __restrict__ a_log, const float* __restrict__ dt_bias, float lower, float* __restrict__ q_out, + float* __restrict__ k_save, __nv_bfloat16* __restrict__ v_save, float* __restrict__ g_save, + float* __restrict__ b_save) { + prep_kernel_body( + H, P, p_stride, b_off, A, a_stride, cs, cw, a_log, dt_bias, lower, q_out, k_save, v_save, g_save, b_save); +} + +extern "C" __global__ void __launch_bounds__(256) tf_glm_kda_step( + int H, const float* __restrict__ state_in, const float* __restrict__ q_in, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, + int rows, __nv_bfloat16* __restrict__ y_out, float* __restrict__ state_out) { + step_kernel_body<8, 16>(H, state_in, q_in, k_save, v_save, g_save, b_save, rows, y_out, state_out); +} + +extern "C" __global__ void __launch_bounds__(128) tf_glm_kda_out( + int H, const __nv_bfloat16* __restrict__ y_in, const __nv_bfloat16* __restrict__ G, int g_stride, + const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out) { + out_kernel_body(H, y_in, G, g_stride, norm_w, eps, out); +} diff --git a/zig/kernels/cuda/glm_ops.cu b/zig/kernels/cuda/glm_ops.cu new file mode 100644 index 0000000000..7805cd23a6 --- /dev/null +++ b/zig/kernels/cuda/glm_ops.cu @@ -0,0 +1,51 @@ +// Own kernels where the Python family used torch ops (integer work, copies, rounding), with the same bits. + +#include +#include + +// sparse.select_tokens' assembly, a block a row: the K pools' 4 tokens each, then the last pool's visible tail. +extern "C" __global__ void tf_glm_dsa_tokens(const long long* __restrict__ pools, const int* __restrict__ pos, int K, + int* __restrict__ tokens, int* __restrict__ counts) { + const int r = blockIdx.x; + const long long q = static_cast(*pos) + r; + const long long npool = (q + 1) / 4; + const int width = 4 * K + 3; + int* row = tokens + static_cast(r) * width; + for (int i = threadIdx.x; i < 4 * K; i += blockDim.x) + row[i] = static_cast(pools[static_cast(r) * K + i / 4] * 4 + i % 4); + if (threadIdx.x < 3) { + const long long t = npool * 4 + threadIdx.x; + row[4 * K + threadIdx.x] = t <= q ? static_cast(t) : -1; + } + if (threadIdx.x == 0) { + int tail = 0; + for (int t = 0; t < 3; ++t) tail += npool * 4 + t <= q; + counts[r] = npool > K ? 4 * K + tail : 0; + } +} + +// pos_dev + a (sparse.select_tokens' at = pos_dev + a for a row block past the first): one int32 on the device. +extern "C" __global__ void tf_glm_add_i32(const int* __restrict__ x, int a, int* __restrict__ out) { + if (threadIdx.x == 0 && blockIdx.x == 0) *out = *x + a; +} + +// forward.shared_expert's b.ey[:, slot].copy_(b.sy): fp32 rows into slot `slot`, as is or bf16 round-to-nearest-even. +extern "C" __global__ void tf_glm_put_slot(const float* __restrict__ src, void* __restrict__ dst, int rows, int d, + int slots, int slot, int bf16) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= static_cast(rows) * d) return; + const long long r = i / d, c = i % d; + const long long at = (r * slots + slot) * d + c; + if (bf16) reinterpret_cast<__nv_bfloat16*>(dst)[at] = __float2bfloat16_rn(src[i]); + else reinterpret_cast(dst)[at] = src[i]; +} + +// A strided destination's copy_ (Fp4Linear's out= column slice): `rows` rows of row_bytes between pitches, as is. +extern "C" __global__ void tf_glm_copy_rows(const unsigned char* __restrict__ src, long long src_pitch, + unsigned char* __restrict__ dst, long long dst_pitch, int rows, + int row_bytes) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= static_cast(rows) * row_bytes) return; + const long long r = i / row_bytes, c = i % row_bytes; + dst[r * dst_pitch + c] = src[r * src_pitch + c]; +} diff --git a/zig/kernels/cuda/glm_quant.cu b/zig/kernels/cuda/glm_quant.cu new file mode 100644 index 0000000000..5013270ea0 --- /dev/null +++ b/zig/kernels/cuda/glm_quant.cu @@ -0,0 +1,127 @@ +// GLM's draft quantizers the Python family ran in torch at load (experts.quantize, qmm.quantize4); x / k = x * (1/k). + +#include +#include +#include + +namespace { + +// E2M1's magnitudes in code order; experts.py's argmin keeps the first of equal distances (midpoints round down). +__device__ __constant__ float kE2M1[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + +__device__ __forceinline__ void widen8(const uint4& u, float* x) { + const __nv_bfloat16* h = reinterpret_cast(&u); +#pragma unroll + for (int j = 0; j < 8; ++j) x[j] = __bfloat162float(h[j]); +} + +// argmin over |v| - E2M1 (first minimum), plus 8 when v < 0 (-0.0 is not). +__device__ __forceinline__ unsigned e2m1(float v) { + const float a = fabsf(v); + unsigned best = 0; + float d = fabsf(a - kE2M1[0]); +#pragma unroll + for (unsigned i = 1; i < 8; ++i) { + const float di = fabsf(a - kE2M1[i]); + if (di < d) d = di, best = i; + } + return best + (v < 0.0f ? 8u : 0u); +} + +} // namespace + +// experts._quantize's per-expert scale, a block per expert: g[e] = clamp_min(amax |w[e]| * fp32(1 / 2688), 1e-30). +extern "C" __global__ void tf_glm_nvfp4_amax( + const __nv_bfloat16* __restrict__ w, long long per, float* __restrict__ g) { + const uint4* src = reinterpret_cast(w + blockIdx.x * per); + float m = 0.0f; + for (long long i = threadIdx.x; i < per / 8; i += blockDim.x) { + float x[8]; + widen8(src[i], x); +#pragma unroll + for (int j = 0; j < 8; ++j) m = fmaxf(m, fabsf(x[j])); + } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); + __shared__ float part[32]; + if (threadIdx.x % 32 == 0) part[threadIdx.x / 32] = m; + __syncthreads(); + if (threadIdx.x == 0) { + for (unsigned i = 1; i < blockDim.x / 32; ++i) m = fmaxf(m, part[i]); + g[blockIdx.x] = fmaxf(m * (1.0f / 2688.0f), (float)1e-30); + } +} + +// experts._quantize's blocks, a thread per 16 inputs: s = e4m3(min(amax * fp32(1 / 6) / g, 448)), codes and scales. +extern "C" __global__ void tf_glm_nvfp4_quantize(const __nv_bfloat16* __restrict__ w, const float* __restrict__ g, + long long per, long long total, uint8_t* __restrict__ codes, + uint8_t* __restrict__ scales) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total) return; + const float ge = g[i / per]; + const uint4* src = reinterpret_cast(w) + 2 * i; + float x[16]; + widen8(src[0], x); + widen8(src[1], x + 8); + float amax = 0.0f; +#pragma unroll + for (int j = 0; j < 16; ++j) amax = fmaxf(amax, fabsf(x[j])); + const float pre = fminf(__fdiv_rn(amax * (1.0f / 6.0f), ge), 448.0f); + __nv_fp8_e4m3 s; + s.__x = __nv_cvt_float_to_fp8(pre, __NV_SATFINITE, __NV_E4M3); + const float step = static_cast(s) * ge; + const float den = fmaxf(step, (float)1e-30); + uint32_t lo = 0, hi = 0; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const float v = step > 0.0f ? __fdiv_rn(x[j], den) : 0.0f; + const uint32_t c = e2m1(v); + if (j < 8) lo |= c << (4 * j); + else hi |= c << (4 * (j - 8)); + } + reinterpret_cast(codes)[i] = make_uint2(lo, hi); + scales[i] = s.__x; +} + +// qmm.quantize4 before make_q4's tiling, a thread per 64-group: bf16 scale (hi - lo) * fp32(1 / 15) >= 1e-8, bias lo. +extern "C" __global__ void tf_glm_quantize4(const __nv_bfloat16* __restrict__ w, long long groups, + int32_t* __restrict__ words, __nv_bfloat16* __restrict__ scales, + __nv_bfloat16* __restrict__ biases) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= groups) return; + const uint4* src = reinterpret_cast(w) + 8 * i; + float x[64]; +#pragma unroll + for (int t = 0; t < 8; ++t) widen8(src[t], x + 8 * t); + float lo = x[0], hi = x[0]; +#pragma unroll + for (int j = 1; j < 64; ++j) lo = fminf(lo, x[j]), hi = fmaxf(hi, x[j]); + // +0/-0 tie: amin keeps the zero at the largest bit-reversed position of 64; hi's sign never matters. + if (lo == 0.0f) { + int best = -1; +#pragma unroll + for (int j = 0; j < 64; ++j) { + const int r = static_cast(__brev(static_cast(j)) >> 26); + if (x[j] == 0.0f && r > best) best = r, lo = x[j]; + } + } + const __nv_bfloat16 sb = __float2bfloat16_rn(fmaxf((hi - lo) * (1.0f / 15.0f), (float)1e-8)); + const __nv_bfloat16 bb = __float2bfloat16_rn(lo); + const float scale = __bfloat162float(sb), bias = __bfloat162float(bb); + uint32_t out[8]; +#pragma unroll + for (int t = 0; t < 8; ++t) { + uint32_t word = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float q = fminf(fmaxf(rintf(__fdiv_rn(x[8 * t + j] - bias, scale)), 0.0f), 15.0f); + word |= static_cast(q) << (4 * j); + } + out[t] = word; + } + uint4* dst = reinterpret_cast(words) + 2 * i; + dst[0] = make_uint4(out[0], out[1], out[2], out[3]); + dst[1] = make_uint4(out[4], out[5], out[6], out[7]); + scales[i] = sb; + biases[i] = bb; +} diff --git a/zig/kernels/cuda/glm_triton.cuh b/zig/kernels/cuda/glm_triton.cuh new file mode 100644 index 0000000000..8c95e5049d --- /dev/null +++ b/zig/kernels/cuda/glm_triton.cuh @@ -0,0 +1,60 @@ +// What GLM's Triton kernels compile to, for our hand-written copies (glm_glue.cu): approximate PTX ops and tl.sum. +#pragma once + +#include +#include + +namespace glm_triton { + +__device__ __forceinline__ float bf(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float rbf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +__device__ __forceinline__ float sqrt_approx(float x) { + float y; + asm("sqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float div_full(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} + +// Triton's add.f32 and mul.f32 carry no rounding mode, so ptxas may contract them into an fma as it did Triton's. +__device__ __forceinline__ float fadd(float a, float b) { + float y; + asm("add.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float fmul(float a, float b) { + float y; + asm("mul.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} + +// tl.exp: ex2.approx of x times log2(e) rounded to fp32. +__device__ __forceinline__ float exp_approx(float x) { + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(__fmul_rn(x, 0x1.715476p+0f))); + return y; +} + +// tl.sum over a block: a butterfly over each warp, then over the warps' partials in the first NW lanes of warp 0. +template +__device__ __forceinline__ float block_sum(float acc, float* smem) { + for (int o = 16; o > 0; o >>= 1) acc += __shfl_xor_sync(0xffffffffu, acc, o); + if (NW == 1) return acc; + if ((threadIdx.x & 31) == 0) smem[threadIdx.x >> 5] = acc; + __syncthreads(); + float v = threadIdx.x < NW ? smem[threadIdx.x] : 0.0f; + for (int o = NW / 2; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + __syncthreads(); + if (threadIdx.x == 0) smem[0] = v; + __syncthreads(); + const float r = smem[0]; + __syncthreads(); + return r; +} + +} // namespace glm_triton diff --git a/zig/kernels/cuda/nvfp4_experts.cu b/zig/kernels/cuda/nvfp4_experts.cu index 0a92b042f4..be724dead8 100644 --- a/zig/kernels/cuda/nvfp4_experts.cu +++ b/zig/kernels/cuda/nvfp4_experts.cu @@ -291,11 +291,12 @@ __device__ __forceinline__ void prompt_body(const __nv_bfloat16* __restrict__ X, const uint4* __restrict__ W, const float* __restrict__ scale, int KG, int NB, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, void* __restrict__ out, int N, - float limit) { + float limit, int skip) { static_assert(RT == 4, "the dispatch below spells out 0 .. 4 row tiles"); const int nbt = (NB + WN - 1) / WN, it = blockIdx.x / nbt, cbt = blockIdx.x - it * nbt; if (it >= __ldg(counts)) return; const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + if (e == skip) return; // an expert the family computes itself, left alone as the decode form does const int nt = max(0, min(RT, (cnt - 16 * RT * ((threadIdx.x >> 5) / WN) + 15) >> 4)); #define TF_TILES(NT) \ prompt_tiles(X, x_stride, slots, W, scale, KG, NB, members, out, N, limit, cbt, e, first, cnt) @@ -321,9 +322,9 @@ template __global__ void nvfp4_expert_kernel<1, 1, 4>(const __nv_bfloat16* __res extern "C" __global__ void __launch_bounds__(WN * 32) NAME( \ const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, \ const float* __restrict__ scale, int KG, int NB, const int* __restrict__ items, const int* __restrict__ counts, \ - const int* __restrict__ members, void* __restrict__ out, int N, float limit) { \ + const int* __restrict__ members, void* __restrict__ out, int N, float limit, int skip) { \ tf_nvfp4_experts::prompt_body(X, x_stride, slots, W, scale, KG, NB, items, counts, members, \ - out, N, limit); \ + out, N, limit, skip); \ } TF_NVFP4_PROMPT(tf_nvfp4_prompt_relu2, 1, 1, 8) TF_NVFP4_PROMPT(tf_nvfp4_prompt_swiglu, 2, 2, 4) diff --git a/zig/kernels/cuda/prompt16.cu b/zig/kernels/cuda/prompt16.cu index c2c6c2f3ab..01749ba86d 100644 --- a/zig/kernels/cuda/prompt16.cu +++ b/zig/kernels/cuda/prompt16.cu @@ -192,5 +192,6 @@ __global__ void __launch_bounds__(WM * WN * 32) prompt_kernel( // The instance the Zig launcher loads by name (PROMPT_TILE 4 in linear.py). template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); +template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); } diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 4703273410..57a7c1e582 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -308,6 +308,13 @@ template __global__ void qmmf_kernel(c template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); @@ -323,4 +330,5 @@ template __global__ void qmmf_kernel template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); +template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); } diff --git a/zig/src/cluster/canon_glm_test.zig b/zig/src/cluster/canon_glm_test.zig new file mode 100644 index 0000000000..4218d341c7 --- /dev/null +++ b/zig/src/cluster/canon_glm_test.zig @@ -0,0 +1,40 @@ +//! canon.partition over equal ranks cuts GLM-5.3-Flash's split axes as the Python family's split_units does. + +const std = @import("std"); +const canon = @import("canon.zig"); + +/// An axis in whole units (heads one by one; MLP widths and vocabulary in 64s), the Python family's [first, end). +const Case = struct { what: []const u8, units: u32, shares: []const [2]u32 }; + +// printed by release/2.7.0's split_units for TP=2 and TP=3 +const cases = [_]Case{ + .{ .what = "KDA heads, TP=2", .units = 64, .shares = &.{ .{ 0, 32 }, .{ 32, 64 } } }, + .{ .what = "KDA heads, TP=3", .units = 64, .shares = &.{ .{ 0, 22 }, .{ 22, 43 }, .{ 43, 64 } } }, + .{ .what = "routed expert width / 64, TP=2", .units = 32, .shares = &.{ .{ 0, 16 }, .{ 16, 32 } } }, + .{ .what = "routed expert width / 64, TP=3", .units = 32, .shares = &.{ .{ 0, 11 }, .{ 11, 22 }, .{ 22, 32 } } }, + .{ .what = "dense MLP width / 64, TP=2", .units = 192, .shares = &.{ .{ 0, 96 }, .{ 96, 192 } } }, + .{ .what = "dense MLP width / 64, TP=3", .units = 192, .shares = &.{ .{ 0, 64 }, .{ 64, 128 }, .{ 128, 192 } } }, + .{ .what = "vocabulary / 64, TP=2", .units = 2420, .shares = &.{ .{ 0, 1210 }, .{ 1210, 2420 } } }, + .{ .what = "vocabulary / 64, TP=3", .units = 2420, .shares = &.{ .{ 0, 807 }, .{ 807, 1614 }, .{ 1614, 2420 } } }, +}; + +test "equal ranks get the Python family's shares of every GLM axis: units / world, the remainder to the lowest ranks" { + const ones = [_]u64{ 1, 1, 1 }; + for (cases) |c| { + var out: [3]canon.Range = undefined; + const world = c.shares.len; + canon.partition(c.units, ones[0..world], out[0..world]); + for (c.shares, out[0..world], 0..) |want, got, rank| { + if (got.begin != want[0] or got.end != want[1]) { + std.debug.print("{s}: rank {d} got [{d}, {d}), Python [{d}, {d})\n", .{ c.what, rank, got.begin, got.end, want[0], want[1] }); + return error.TestUnexpectedResult; + } + } + } +} + +test "a rank left empty: the Python family refuses the split, partition gives it nothing, so a share plan must check" { + var out: [3]canon.Range = undefined; + canon.partition(2, &.{ 1, 1, 1 }, &out); + try std.testing.expectEqual(@as(u32, 0), out[2].len()); +} diff --git a/zig/src/core/checkpoint.zig b/zig/src/core/checkpoint.zig index 1c2110a685..4425e4829d 100644 --- a/zig/src/core/checkpoint.zig +++ b/zig/src/core/checkpoint.zig @@ -122,6 +122,11 @@ pub const Checkpoint = struct { }; } + /// Every file's pages out of the page cache (safetensors.File.release). + pub fn release(self: *const Checkpoint) void { + for (self.files.items) |*f| f.release(); + } + /// Names no `get` took (the loader refuses a checkpoint with leftovers). pub fn unused(self: *const Checkpoint) usize { var n: usize = 0; @@ -224,3 +229,21 @@ test "has finds a name without taking it; skip marks a prefix used" { _ = try ck.get("backbone.x"); try std.testing.expectEqual(@as(usize, 0), ck.unused()); } + +test "after release the mapping stays readable: a later read faults the same bytes back in" { + const io = std.testing.io; + const gpa = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const image = try testImage(gpa, "x", 4096); + defer gpa.free(image); + for (image[image.len - 4096 ..], 0..) |*b, i| b.* = @truncate(i * 7 + 1); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = image }); + var path: [64]u8 = undefined; + var ck = try Checkpoint.openModel(gpa, io, try std.fmt.bufPrint(&path, ".zig-cache/tmp/{s}", .{tmp.sub_path})); + defer ck.close(); + const t = try ck.get("x"); + try std.testing.expectEqualSlices(u8, image[image.len - 4096 ..], t.bytes); + ck.release(); + try std.testing.expectEqualSlices(u8, image[image.len - 4096 ..], t.bytes); +} diff --git a/zig/src/core/fixtures/glm53_flash_nvfp4/config.json b/zig/src/core/fixtures/glm53_flash_nvfp4/config.json new file mode 100644 index 0000000000..d0e8e75de4 --- /dev/null +++ b/zig/src/core/fixtures/glm53_flash_nvfp4/config.json @@ -0,0 +1,469 @@ +{ + "architectures": [ + "Glm5NextForConditionalGeneration" + ], + "dtype": "bfloat16", + "image_end_token_id": 154831, + "image_start_token_id": 154830, + "image_token_id": 154854, + "model_type": "glm5_next", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "base_model_ep_plan": { + "layers.*.mlp.experts": "torch_a2a_experts", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.gate": "torch_a2a_router" + }, + "bos_token_id": null, + "dtype": "bfloat16", + "eos_token_id": [ + 154820, + 154827, + 154829 + ], + "first_k_dense_replace": 3, + "hc_eps": 1e-06, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "head_dim": 0, + "hidden_act": "silu", + "hidden_size": 4096, + "index_head_dim": 128, + "index_kpool": 4, + "index_kpool_always_select_tail": true, + "index_kpool_compress": true, + "index_n_heads": 32, + "index_share_for_mtp_iteration": true, + "index_topk": 2048, + "indexer_rope_interleave": true, + "indexer_types": [ + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full" + ], + "initializer_range": 0.02, + "intermediate_size": 12288, + "kv_lora_rank": 512, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention" + ], + "linear_attn_config": { + "full_attn_layers": [ + 3, + 7, + 11, + 15, + 19, + 23, + 27, + 31, + 35, + 39, + 43 + ], + "gate_lower_bound": -5.0, + "head_dim": 128, + "kda_layers": [ + 0, + 1, + 2, + 4, + 5, + 6, + 8, + 9, + 10, + 12, + 13, + 14, + 16, + 17, + 18, + 20, + 21, + 22, + 24, + 25, + 26, + 28, + 29, + 30, + 32, + 33, + 34, + 36, + 37, + 38, + 40, + 41, + 42, + 44 + ], + "num_heads": 64, + "short_conv_kernel_size": 4 + }, + "linear_conv_kernel_dim": 4, + "linear_head_dim": 128, + "linear_lower_bound": -5.0, + "linear_num_heads": 64, + "max_position_embeddings": 1048576, + "mhc": true, + "mla_use_nope": true, + "mlp_layer_types": [ + "dense", + "dense", + "dense", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse" + ], + "model_type": "glm5_next_text", + "moe_intermediate_size": 2048, + "moe_router_dtype": "float32", + "n_group": 1, + "n_routed_experts": 288, + "n_shared_experts": 1, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts_per_tok": 8, + "num_hidden_layers": 45, + "num_key_value_heads": 64, + "num_nextn_predict_layers": 1, + "output_router_logits": false, + "pad_token_id": 154820, + "q_lora_rank": 1536, + "qk_head_dim": 256, + "qk_nope_head_dim": 256, + "qk_rope_head_dim": 0, + "rms_norm_eps": 1e-05, + "routed_scaling_factor": 2.5, + "router_aux_loss_coef": 0.001, + "scoring_func": "sigmoid", + "swiglu_limit": 10.0, + "tie_word_embeddings": false, + "topk_group": 1, + "topk_method": "noaux_tc", + "use_cache": true, + "v_head_dim": 256, + "vocab_size": 154880 + }, + "tie_word_embeddings": false, + "transformers_version": "5.16.1", + "video_end_token_id": 154833, + "video_start_token_id": 154832, + "video_token_id": 154855, + "vision_config": { + "attention_bias": true, + "attention_dropout": 0.0, + "depth": 24, + "dtype": "bfloat16", + "hidden_act": "silu", + "hidden_size": 1024, + "image_size": 448, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "glm5_next_vision", + "num_heads": 16, + "out_hidden_size": 4096, + "patch_size": 14, + "projection_intermediate_size": 10240, + "rms_norm_eps": 1e-05, + "spatial_merge_size": 2, + "swiglu_limit": 10.0, + "temporal_patch_size": 2 + }, + "quantization_config": { + "config_groups": { + "group_0": { + "input_activations": { + "dynamic": false, + "num_bits": 4, + "type": "float", + "group_size": 16 + }, + "weights": { + "dynamic": false, + "num_bits": 4, + "type": "float", + "group_size": 16 + }, + "targets": [ + "Linear" + ] + } + }, + "ignore": [ + "lm_head", + "model.language_model.embed_tokens", + "model.language_model.layers.0.self_attn*", + "model.language_model.layers.1.self_attn*", + "model.language_model.layers.10.mlp.gate", + "model.language_model.layers.10.mlp.shared_experts*", + "model.language_model.layers.10.self_attn*", + "model.language_model.layers.11.mlp.gate", + "model.language_model.layers.11.mlp.shared_experts*", + "model.language_model.layers.11.self_attn*", + "model.language_model.layers.12.mlp.gate", + "model.language_model.layers.12.mlp.shared_experts*", + "model.language_model.layers.12.self_attn*", + "model.language_model.layers.13.mlp.gate", + "model.language_model.layers.13.mlp.shared_experts*", + "model.language_model.layers.13.self_attn*", + "model.language_model.layers.14.mlp.gate", + "model.language_model.layers.14.mlp.shared_experts*", + "model.language_model.layers.14.self_attn*", + "model.language_model.layers.15.mlp.gate", + "model.language_model.layers.15.mlp.shared_experts*", + "model.language_model.layers.15.self_attn*", + "model.language_model.layers.16.mlp.gate", + "model.language_model.layers.16.mlp.shared_experts*", + "model.language_model.layers.16.self_attn*", + "model.language_model.layers.17.mlp.gate", + "model.language_model.layers.17.mlp.shared_experts*", + "model.language_model.layers.17.self_attn*", + "model.language_model.layers.18.mlp.gate", + "model.language_model.layers.18.mlp.shared_experts*", + "model.language_model.layers.18.self_attn*", + "model.language_model.layers.19.mlp.gate", + "model.language_model.layers.19.mlp.shared_experts*", + "model.language_model.layers.19.self_attn*", + "model.language_model.layers.2.self_attn*", + "model.language_model.layers.20.mlp.gate", + "model.language_model.layers.20.mlp.shared_experts*", + "model.language_model.layers.20.self_attn*", + "model.language_model.layers.21.mlp.gate", + "model.language_model.layers.21.mlp.shared_experts*", + "model.language_model.layers.21.self_attn*", + "model.language_model.layers.22.mlp.gate", + "model.language_model.layers.22.mlp.shared_experts*", + "model.language_model.layers.22.self_attn*", + "model.language_model.layers.23.mlp.gate", + "model.language_model.layers.23.mlp.shared_experts*", + "model.language_model.layers.23.self_attn*", + "model.language_model.layers.24.mlp.gate", + "model.language_model.layers.24.mlp.shared_experts*", + "model.language_model.layers.24.self_attn*", + "model.language_model.layers.25.mlp.gate", + "model.language_model.layers.25.mlp.shared_experts*", + "model.language_model.layers.25.self_attn*", + "model.language_model.layers.26.mlp.gate", + "model.language_model.layers.26.mlp.shared_experts*", + "model.language_model.layers.26.self_attn*", + "model.language_model.layers.27.mlp.gate", + "model.language_model.layers.27.mlp.shared_experts*", + "model.language_model.layers.27.self_attn*", + "model.language_model.layers.28.mlp.gate", + "model.language_model.layers.28.mlp.shared_experts*", + "model.language_model.layers.28.self_attn*", + "model.language_model.layers.29.mlp.gate", + "model.language_model.layers.29.mlp.shared_experts*", + "model.language_model.layers.29.self_attn*", + "model.language_model.layers.3.mlp.gate", + "model.language_model.layers.3.mlp.shared_experts*", + "model.language_model.layers.3.self_attn*", + "model.language_model.layers.30.mlp.gate", + "model.language_model.layers.30.mlp.shared_experts*", + "model.language_model.layers.30.self_attn*", + "model.language_model.layers.31.mlp.gate", + "model.language_model.layers.31.mlp.shared_experts*", + "model.language_model.layers.31.self_attn*", + "model.language_model.layers.32.mlp.gate", + "model.language_model.layers.32.mlp.shared_experts*", + "model.language_model.layers.32.self_attn*", + "model.language_model.layers.33.mlp.gate", + "model.language_model.layers.33.mlp.shared_experts*", + "model.language_model.layers.33.self_attn*", + "model.language_model.layers.34.mlp.gate", + "model.language_model.layers.34.mlp.shared_experts*", + "model.language_model.layers.34.self_attn*", + "model.language_model.layers.35.mlp.gate", + "model.language_model.layers.35.mlp.shared_experts*", + "model.language_model.layers.35.self_attn*", + "model.language_model.layers.36.mlp.gate", + "model.language_model.layers.36.mlp.shared_experts*", + "model.language_model.layers.36.self_attn*", + "model.language_model.layers.37.mlp.gate", + "model.language_model.layers.37.mlp.shared_experts*", + "model.language_model.layers.37.self_attn*", + "model.language_model.layers.38.mlp.gate", + "model.language_model.layers.38.mlp.shared_experts*", + "model.language_model.layers.38.self_attn*", + "model.language_model.layers.39.mlp.gate", + "model.language_model.layers.39.mlp.shared_experts*", + "model.language_model.layers.39.self_attn*", + "model.language_model.layers.4.mlp.gate", + "model.language_model.layers.4.mlp.shared_experts*", + "model.language_model.layers.4.self_attn*", + "model.language_model.layers.40.mlp.gate", + "model.language_model.layers.40.mlp.shared_experts*", + "model.language_model.layers.40.self_attn*", + "model.language_model.layers.41.mlp.gate", + "model.language_model.layers.41.mlp.shared_experts*", + "model.language_model.layers.41.self_attn*", + "model.language_model.layers.42.mlp.gate", + "model.language_model.layers.42.mlp.shared_experts*", + "model.language_model.layers.42.self_attn*", + "model.language_model.layers.43.mlp.gate", + "model.language_model.layers.43.mlp.shared_experts*", + "model.language_model.layers.43.self_attn*", + "model.language_model.layers.44.mlp.gate", + "model.language_model.layers.44.mlp.shared_experts*", + "model.language_model.layers.44.self_attn*", + "model.language_model.layers.5.mlp.gate", + "model.language_model.layers.5.mlp.shared_experts*", + "model.language_model.layers.5.self_attn*", + "model.language_model.layers.6.mlp.gate", + "model.language_model.layers.6.mlp.shared_experts*", + "model.language_model.layers.6.self_attn*", + "model.language_model.layers.7.mlp.gate", + "model.language_model.layers.7.mlp.shared_experts*", + "model.language_model.layers.7.self_attn*", + "model.language_model.layers.8.mlp.gate", + "model.language_model.layers.8.mlp.shared_experts*", + "model.language_model.layers.8.self_attn*", + "model.language_model.layers.9.mlp.gate", + "model.language_model.layers.9.mlp.shared_experts*", + "model.language_model.layers.9.self_attn*", + "model.visual*" + ], + "quant_algo": "NVFP4", + "kv_cache_scheme": { + "dynamic": false, + "num_bits": 8, + "type": "float" + }, + "producer": { + "name": "modelopt", + "version": "0.47.0.dev393+ga4bc45b30.d20260828" + }, + "quant_method": "modelopt" + } +} \ No newline at end of file diff --git a/zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz b/zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz new file mode 100644 index 0000000000..db569985ec Binary files /dev/null and b/zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz differ diff --git a/zig/src/core/modelopt.zig b/zig/src/core/modelopt.zig new file mode 100644 index 0000000000..da37ad7b5e --- /dev/null +++ b/zig/src/core/modelopt.zig @@ -0,0 +1,110 @@ +//! ModelOpt NVFP4 config check (groups of 16); tensors say which linears are NVFP4, not ignore (423acf37's BF16 MTP). + +const std = @import("std"); + +/// The NVFP4 group size ModelOpt writes for this format. +pub const nvfp4_group = 16; + +/// Why the check refused the checkpoint, kept for the caller's log line (tests read it instead). +pub const Why = struct { + buf: [320]u8 = undefined, + len: usize = 0, + + /// Keep the reason, cut short if it is long. + pub fn set(self: *Why, comptime fmt: []const u8, args: anytype) void { + var w: std.Io.Writer = .fixed(&self.buf); + w.print(fmt, args) catch {}; + self.len = w.end; + } + + pub fn text(self: *const Why) []const u8 { + return self.buf[0..self.len]; + } +}; + +/// config.json's quantization_config (top level, else text_config's) checked as ModelOpt NVFP4. +pub fn admit(gpa: std.mem.Allocator, config_json: []const u8, why: *Why) !void { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const root = try std.json.parseFromSliceLeaky(std.json.Value, arena.allocator(), config_json, .{}); + const q = block(root) orelse { + why.set("config.json has no ModelOpt quantization_config", .{}); + return error.NotModelOpt; + }; + const algo = text(q.get("quant_algo")) orelse ""; + if (!std.ascii.eqlIgnoreCase(algo, "NVFP4")) { + why.set("ModelOpt {s}: this engine reads ModelOpt NVFP4 (nvidia/GLM-5.3-Flash-NVFP4's layout)", .{if (algo.len > 0) algo else "without quant_algo"}); + return error.UnsupportedQuant; + } + const groups = object(q.get("config_groups")) orelse { + why.set("ModelOpt NVFP4 without config_groups", .{}); + return error.UnsupportedQuant; + }; + var it = groups.iterator(); + while (it.next()) |g| { + const w = if (object(g.value_ptr.*)) |o| object(o.get("weights")) else null; + const bits = if (w) |o| int(o.get("num_bits")) else null; + const size = if (w) |o| int(o.get("group_size")) else null; + const kind = if (w) |o| text(o.get("type")) orelse "" else ""; + if (bits != 4 or size != nvfp4_group or !std.mem.eql(u8, kind, "float")) { + why.set("ModelOpt NVFP4 group {s}: weights {?d}-bit {s} in groups of {?d}, expected 4-bit float in groups of 16", .{ g.key_ptr.*, bits, kind, size }); + return error.UnsupportedQuant; + } + } +} + +fn block(root: std.json.Value) ?std.json.ObjectMap { + const top = object(root) orelse return null; + for ([_]?std.json.ObjectMap{ top, object(top.get("text_config")) }) |source| { + const q = object((source orelse continue).get("quantization_config")) orelse continue; + if (std.ascii.eqlIgnoreCase(text(q.get("quant_method")) orelse "", "modelopt")) return q; + } + return null; +} + +fn object(v: ?std.json.Value) ?std.json.ObjectMap { + const x = v orelse return null; + return if (x == .object) x.object else null; +} + +fn text(v: ?std.json.Value) ?[]const u8 { + const x = v orelse return null; + return if (x == .string) x.string else null; +} + +fn int(v: ?std.json.Value) ?i64 { + const x = v orelse return null; + return if (x == .integer) x.integer else null; +} + +test "nvidia/GLM-5.3-Flash-NVFP4's config is admitted" { + var why: Why = .{}; + try admit(std.testing.allocator, @embedFile("fixtures/glm53_flash_nvfp4/config.json"), &why); +} + +/// admit() must refuse `json` with `err`, its reason naming `word`. +fn expectRefused(json: []const u8, err: anyerror, word: []const u8) !void { + var why: Why = .{}; + try std.testing.expectError(err, admit(std.testing.allocator, json, &why)); + if (std.mem.indexOf(u8, why.text(), word) == null) { + std.debug.print("refusal \"{s}\" does not name \"{s}\"\n", .{ why.text(), word }); + return error.TestUnexpectedResult; + } +} + +test "other ModelOpt formats, and configs without ModelOpt, are refused by name" { + try expectRefused("{\"text_config\": {}}", error.NotModelOpt, "no ModelOpt"); + try expectRefused("{\"quantization_config\": {\"quant_method\": \"exl3\"}}", error.NotModelOpt, "no ModelOpt"); + try expectRefused("{\"quantization_config\": {\"quant_method\": \"modelopt\", \"quant_algo\": \"FP8\"}}", error.UnsupportedQuant, "ModelOpt FP8"); + // MIXED_PRECISION (the AXL re-pack) comes with 3.x + try expectRefused("{\"quantization_config\": {\"quant_method\": \"modelopt\", \"quant_algo\": \"MIXED_PRECISION\"}}", error.UnsupportedQuant, "MIXED_PRECISION"); + try expectRefused( + \\{"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "config_groups": + \\ {"g": {"weights": {"num_bits": 4, "type": "float", "group_size": 32}}}}} + , error.UnsupportedQuant, "groups of 32"); + var why: Why = .{}; + try admit(std.testing.allocator, + \\{"text_config": {"quantization_config": {"quant_method": "modelopt", "quant_algo": "nvfp4", "config_groups": + \\ {"g": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}}}}}} + , &why); +} diff --git a/zig/src/core/nvfp4.zig b/zig/src/core/nvfp4.zig new file mode 100644 index 0000000000..ed6948a003 --- /dev/null +++ b/zig/src/core/nvfp4.zig @@ -0,0 +1,123 @@ +//! NVFP4 as ModelOpt stores it: e2m1 codes two to a byte, an e4m3 scale per 16 inputs, an fp32 scale per tensor. + +const std = @import("std"); +const ckpt = @import("checkpoint.zig"); +const quant = @import("quant/quant.zig"); +const DType = ckpt.DType; + +/// Inputs sharing one e4m3 scale. +pub const group = quant.nvfp4.group; + +/// A stored tensor's dtype and shape, as a header lists it. +pub const Meta = struct { + dtype: DType, + shape: []const usize, +}; + +pub const Dims = struct { n: usize, k: usize }; + +/// The linear's [n, k] from codes U8 [n, k/2], scales F8_E4M3 [n, k/16] and fp32 scalars; else error.BadNvfp4. +pub fn dims(weight: Meta, scale: Meta, scale2: Meta, input_scale: ?Meta) !Dims { + if (weight.dtype != .u8 or weight.shape.len != 2) return error.BadNvfp4; + if (scale.dtype != .f8_e4m3 or scale.shape.len != 2) return error.BadNvfp4; + const n = weight.shape[0]; + const k = 2 * weight.shape[1]; + if (k % group != 0 or scale.shape[0] != n or scale.shape[1] * group != k) return error.BadNvfp4; + if (!scalar(scale2)) return error.BadNvfp4; + if (input_scale) |s| if (!scalar(s)) return error.BadNvfp4; + return .{ .n = n, .k = k }; +} + +fn scalar(m: Meta) bool { + return m.dtype == .f32 and (m.shape.len == 0 or (m.shape.len == 1 and m.shape[0] == 1)); +} + +/// Whether `module` (a linear's name without ".weight") is stored as NVFP4: its tensors say, not config.json. +pub fn stored(ck: *const ckpt.Checkpoint, module: []const u8) bool { + var buf: [256]u8 = undefined; + const name = std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module}) catch return false; + for (ck.files.items) |*f| if (f.names.contains(name)) return true; + return false; +} + +/// One NVFP4 linear's stored bytes, as the checkpoint maps them. +pub const Linear = struct { + n: usize, + k: usize, + codes: []const u8, // [n, k/2]: input 2j in byte j's low nibble + scales: []const u8, // [n, k/16] e4m3 + global: f32, // weight_scale_2 + + /// `module`'s four tensors, marked as used; input_scale is the W4A4 math's static scale, unread by W4A16. + pub fn read(ck: *ckpt.Checkpoint, module: []const u8) !Linear { + var buf: [256]u8 = undefined; + const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{module})); + const s = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module})); + const g = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{module})); + const a = try ck.get(try std.fmt.bufPrint(&buf, "{s}.input_scale", .{module})); + const d = dims(meta(&w), meta(&s), meta(&g), meta(&a)) catch |e| { + std.log.err("{s}: not ModelOpt NVFP4 ({t} {any}, {t} {any})", .{ module, w.dtype, w.shape[0..w.rank], s.dtype, s.shape[0..s.rank] }); + return e; + }; + return .{ .n = d.n, .k = d.k, .codes = w.bytes, .scales = s.bytes, .global = std.mem.bytesToValue(f32, g.bytes[0..4]) }; + } +}; + +fn meta(t: *const ckpt.Tensor) Meta { + return .{ .dtype = t.dtype, .shape = t.shape[0..t.rank] }; +} + +test "dims refuse what is not ModelOpt NVFP4" { + const w: Meta = .{ .dtype = .u8, .shape = &.{ 4, 16 } }; + const s: Meta = .{ .dtype = .f8_e4m3, .shape = &.{ 4, 2 } }; + const g: Meta = .{ .dtype = .f32, .shape = &.{} }; + const d = try dims(w, s, g, g); + try std.testing.expectEqual(Dims{ .n = 4, .k = 32 }, d); + try std.testing.expectError(error.BadNvfp4, dims(.{ .dtype = .bf16, .shape = &.{ 4, 16 } }, s, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, .{ .dtype = .f8_e4m3, .shape = &.{ 4, 3 } }, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, .{ .dtype = .u8, .shape = &.{ 4, 2 } }, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, s, .{ .dtype = .f32, .shape = &.{2} }, null)); +} + +test "a linear read from a checkpoint folder decodes as core/quant's reference" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + // n = 2, k = 32: codes 2 x 16 bytes, scales 2 x 2, two fp32 scalars + const header = + \\{"l.weight":{"dtype":"U8","shape":[2,16],"data_offsets":[0,32]}, + \\ "l.weight_scale":{"dtype":"F8_E4M3","shape":[2,2],"data_offsets":[32,36]}, + \\ "l.weight_scale_2":{"dtype":"F32","shape":[],"data_offsets":[36,40]}, + \\ "l.input_scale":{"dtype":"F32","shape":[],"data_offsets":[40,44]}} + ; + var data: [8 + header.len + 44]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + const body = data[8 + header.len ..]; + for (body[0..32], 0..) |*b, i| b.* = @intCast((i * 7 + 3) & 0xff); // every nibble value appears + @memcpy(body[32..36], &[_]u8{ 0x38, 0x40, 0x30, 0xb8 }); // 1.0, 2.0, 0.5, -1.0 + @memcpy(body[36..40], std.mem.asBytes(&@as(f32, 0.25))); + @memcpy(body[40..44], std.mem.asBytes(&@as(f32, 3.0))); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &data }); + const dir = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer gpa.free(dir); + var ck = try ckpt.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + try std.testing.expect(stored(&ck, "l")); + try std.testing.expect(!stored(&ck, "m")); + const l = try Linear.read(&ck, "l"); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); + try std.testing.expectEqual(@as(f32, 0.25), l.global); + var out: [64]f64 = undefined; + try quant.nvfp4.reference(.{ + .weight = .{ .dtype = .u8, .rank = 2, .shape = .{ l.n, l.k / 2, 1, 1, 1 }, .bytes = l.codes }, + .scale = .{ .dtype = .f8_e4m3, .rank = 2, .shape = .{ l.n, l.k / group, 1, 1, 1 }, .bytes = l.scales }, + .global = &[_]f32{l.global}, + }, &out); + // byte 0 is 3: input 0 is code 3 (1.5), input 1 code 0; inputs 16.. take the second scale (2.0) + try std.testing.expectEqual(@as(f64, 1.5 * 1.0 * 0.25), out[0]); + try std.testing.expectEqual(@as(f64, 0.0), out[1]); + try std.testing.expectEqual(quant.nvfp4.e2m1(body[8] & 0xf) * 2.0 * 0.25, out[16]); + try std.testing.expectEqual(quant.nvfp4.e2m1(body[16 + 15] >> 4) * -1.0 * 0.25, out[32 + 31]); // row 1, scale -1.0 +} diff --git a/zig/src/core/nvfp4_fixture_test.zig b/zig/src/core/nvfp4_fixture_test.zig new file mode 100644 index 0000000000..9f4bfa3694 --- /dev/null +++ b/zig/src/core/nvfp4_fixture_test.zig @@ -0,0 +1,69 @@ +//! nvidia/GLM-5.3-Flash-NVFP4's whole header: every linear is NVFP4 or BF16 as nvfp4.dims reads it. + +const std = @import("std"); +const nvfp4 = @import("nvfp4.zig"); +const DType = @import("safetensors.zig").DType; + +const Meta = nvfp4.Meta; + +/// name -> [dtype, shape] from the gzip'd fixture, kept in `arena`. +fn headers(arena: std.mem.Allocator) !std.json.ObjectMap { + var in: std.Io.Reader = .fixed(@embedFile("fixtures/glm53_flash_nvfp4/headers.json.gz")); + var gz: std.compress.flate.Decompress = .init(&in, .gzip, &.{}); + var out: std.Io.Writer.Allocating = .init(arena); + _ = try gz.reader.streamRemaining(&out.writer); + const v = try std.json.parseFromSliceLeaky(std.json.Value, arena, out.written(), .{}); + return v.object; +} + +fn meta(arena: std.mem.Allocator, h: std.json.ObjectMap, name: []const u8) !?Meta { + const v = h.get(name) orelse return null; + const pair = v.array.items; + const dims = pair[1].array.items; + const shape = try arena.alloc(usize, dims.len); + for (dims, shape) |d, *s| s.* = @intCast(d.integer); + return .{ .dtype = DType.parse(pair[0].string) orelse return error.UnsupportedDType, .shape = shape }; +} + +fn at(arena: std.mem.Allocator, h: std.json.ObjectMap, module: []const u8, suffix: []const u8) !?Meta { + return meta(arena, h, try std.fmt.allocPrint(arena, "{s}.{s}", .{ module, suffix })); +} + +test "every NVFP4 linear of the checkpoint reads as ModelOpt stores it, and the MTP layer stays BF16" { + var state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer state.deinit(); + const a = state.allocator(); + const h = try headers(a); + try std.testing.expectEqual(@as(usize, 147661), h.count()); + var quantized: usize = 0; + var experts: usize = 0; + var it = h.iterator(); + while (it.next()) |kv| { + const name = kv.key_ptr.*; + if (std.mem.endsWith(u8, name, ".weight_scale")) { + // no scale without its codes + try std.testing.expect((try at(a, h, name[0 .. name.len - ".weight_scale".len], "weight")) != null); + continue; + } + if (!std.mem.endsWith(u8, name, ".weight")) continue; + const module = name[0 .. name.len - ".weight".len]; + const scale = try at(a, h, module, "weight_scale") orelse continue; + const d = try nvfp4.dims((try meta(a, h, name)).?, scale, (try at(a, h, module, "weight_scale_2")).?, try at(a, h, module, "input_scale")); + quantized += 1; + try std.testing.expect(std.mem.indexOf(u8, module, ".mlp.") != null); + try std.testing.expect(std.mem.indexOf(u8, module, ".layers.45.") == null); + const down = std.mem.endsWith(u8, module, "down_proj"); + if (std.mem.indexOf(u8, module, ".experts.") != null) { + experts += 1; + try std.testing.expectEqual(if (down) nvfp4.Dims{ .n = 4096, .k = 2048 } else nvfp4.Dims{ .n = 2048, .k = 4096 }, d); + } else { + try std.testing.expectEqual(if (down) nvfp4.Dims{ .n = 4096, .k = 12288 } else nvfp4.Dims{ .n = 12288, .k = 4096 }, d); + } + } + try std.testing.expectEqual(@as(usize, 42 * 288 * 3 + 3 * 3), quantized); + try std.testing.expectEqual(@as(usize, 42 * 288 * 3), experts); + // the MTP layer's routed experts are BF16 as stored (a family may re-pack them for drafting) + const mtp = (try meta(a, h, "model.language_model.layers.45.mlp.experts.0.gate_proj.weight")).?; + try std.testing.expectEqual(DType.bf16, mtp.dtype); + try std.testing.expect((try meta(a, h, "model.language_model.layers.45.mlp.experts.0.gate_proj.weight_scale")) == null); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 736248a93d..b3d7cfbc48 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -18,6 +18,8 @@ pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); pub const act_ref = @import("act_ref.zig"); +pub const nvfp4 = @import("nvfp4.zig"); +pub const modelopt = @import("modelopt.zig"); test { _ = safetensors; @@ -37,4 +39,7 @@ test { _ = draft_depth; _ = @import("copy_index.zig"); _ = ids_json; + _ = nvfp4; + _ = modelopt; + _ = @import("nvfp4_fixture_test.zig"); } diff --git a/zig/src/core/safetensors.zig b/zig/src/core/safetensors.zig index 8c03a4f762..83bf122f47 100644 --- a/zig/src/core/safetensors.zig +++ b/zig/src/core/safetensors.zig @@ -165,6 +165,14 @@ pub const File = struct { self.* = undefined; } + /// The file's pages out of the mapping and the page cache (unified memory needs the room); Linux only, refaults. + pub fn release(self: *const File) void { + if (@import("builtin").os.tag != .linux) return; + const linux = std.os.linux; + _ = linux.madvise(self.map.memory.ptr, self.map.memory.len, linux.MADV.DONTNEED); + _ = linux.fadvise(self.file.handle, 0, 0, linux.POSIX_FADV.DONTNEED); + } + pub fn get(self: *const File, name: []const u8) ?Tensor { const e = self.names.get(name) orelse return null; return .{ .dtype = e.dtype, .rank = e.rank, .shape = e.shape, .bytes = self.map.memory[self.data + e.begin .. self.data + e.end] }; diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index c940202bea..8a6664cbd6 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -341,7 +341,6 @@ pub const Nvfp4Experts = struct { /// Prompt rows on a plan of at most `items` items of 64 pairs: up (relu^2 or SwiGLU by the layer), bf16 sums down. pub fn runPrompt(e: *const Nvfp4Experts, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { if (epi == .down_f32) return error.EpilogueNotBuilt; - if (skip >= 0) return error.OptionNotBuilt; const gate = epi == .up; const parts: usize = if (gate and !l.relu2) 2 else 1; const nb = (if (gate) l.width else l.dims) / cols; @@ -357,6 +356,7 @@ pub const Nvfp4Experts = struct { for ([_]u64{ p.items, p.counts, p.members, out }) |v| a.add(v); a.add(@as(i32, @intCast(if (gate) l.width else l.dims))); a.add(if (gate) l.limit else @as(f32, 0.0)); + a.add(skip); const fi: usize = if (!gate) 2 else if (l.relu2) 0 else 1; const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(items * ((nb + prompt_wn[fi] - 1) / prompt_wn[fi])), .y = 1, .z = 1 }, .block = .{ .x = @intCast(prompt_wn[fi] * 32), .y = 1, .z = 1 }, .shared = promptSmem(parts, prompt_wn[fi]) }; try launch_.launch(e.prompt_fns[fi], cfg, s, &a); @@ -374,11 +374,22 @@ pub const Nvfp4Experts = struct { /// One expert's words [n, k/2] (low nibble first) and e4m3 scales [n, k/16] as blocks [n/32][k/32][144] (_pack). pub fn packOne(out: []u32, codes: []const u8, scales: []const u8, n: usize, k: usize) void { + packPart(out, 1, 0, codes, scales, n, k); + } + + /// Gate and up interleaved as make() stacks them, [n/32][k/32][2][144], packed in place (a loader packs thousands). + pub fn packGateUp(out: []u32, gate: [2][]const u8, up: [2][]const u8, n: usize, k: usize) void { + packPart(out, 2, 0, gate[0], gate[1], n, k); + packPart(out, 2, 1, up[0], up[1], n, k); + } + + /// packOne into matrix `part` of `parts` stacked per (column block, k32 group). + fn packPart(out: []u32, parts: usize, part: usize, codes: []const u8, scales: []const u8, n: usize, k: usize) void { const kg = k / 32; - std.debug.assert(n % cols == 0 and k % 32 == 0 and out.len == n / cols * kg * words); + std.debug.assert(n % cols == 0 and k % 32 == 0 and out.len == n / cols * kg * parts * words and part < parts); std.debug.assert(codes.len == n * k / 2 and scales.len == n * k / 16); for (0..n / cols) |cb| for (0..kg) |g| { - const block = out[(cb * kg + g) * words ..][0..words]; + const block = out[((cb * kg + g) * parts + part) * words ..][0..words]; for (0..8) |gq| for (0..4) |t| for (0..4) |j| { const row = cb * cols + j * 8 + gq; var w: u32 = 0; @@ -397,22 +408,6 @@ pub const Nvfp4Experts = struct { }; }; } - - /// Gate and up interleaved per (column block, k32 group) as make() stacks them: [n/32][k/32][2][144]. - pub fn packGateUp(gpa: std.mem.Allocator, out: []u32, gate: [2][]const u8, up: [2][]const u8, n: usize, k: usize) !void { - const half = n / cols * (k / 32) * words; - std.debug.assert(out.len == 2 * half); - const g = try gpa.alloc(u32, half); - defer gpa.free(g); - const u = try gpa.alloc(u32, half); - defer gpa.free(u); - packOne(g, gate[0], gate[1], n, k); - packOne(u, up[0], up[1], n, k); - for (0..half / words) |b| { - @memcpy(out[(2 * b) * words ..][0..words], g[b * words ..][0..words]); - @memcpy(out[(2 * b + 1) * words ..][0..words], u[b * words ..][0..words]); - } - } }; test "a packed nvfp4 block puts each nibble and scale where nvfp4/experts.py's _pack does" { @@ -435,6 +430,25 @@ test "a packed nvfp4 block puts each nibble and scale where nvfp4/experts.py's _ try std.testing.expectEqual(scales[44 * 4 + 3], sc[((2 * 2 + 1) * 4 + 1) * 2]); } +test "packGateUp stacks packOne's gate and up blocks per (column block, k32 group)" { + const n = 64; + const k = 64; + const words = Nvfp4Experts.words; + var codes: [2][n * k / 2]u8 = undefined; + var scales: [2][n * k / 16]u8 = undefined; + for (0..2) |m| { + for (&codes[m], 0..) |*c, i| c.* = @truncate(i * 29 + 3 + m * 101); + for (&scales[m], 0..) |*c, i| c.* = @truncate(i * 7 + 1 + m * 13); + } + var one: [2][n / 32 * (k / 32) * words]u32 = undefined; + for (0..2) |m| Nvfp4Experts.packOne(&one[m], &codes[m], &scales[m], n, k); + var both: [2 * one[0].len]u32 = undefined; + Nvfp4Experts.packGateUp(&both, .{ &codes[0], &scales[0] }, .{ &codes[1], &scales[1] }, n, k); + for (0..one[0].len / words) |b| for (0..2) |m| { + try std.testing.expectEqualSlices(u32, one[m][b * words ..][0..words], both[(2 * b + m) * words ..][0..words]); + }; +} + test "nvfp4 rows of either kind are refused when the format is not loaded" { const g: Grouped = .{}; const l: Layer = .{ .format = .nvfp4, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .relu2 = true }; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..644ab7d5f5 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -33,6 +33,10 @@ pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &. pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; +pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes else &.{}; +pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; +pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; +pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_pointwise").bytes else &.{}; diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..b0d527ce75 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -25,6 +25,8 @@ pub const Api = struct { ncclAllReduce: *const fn (D, D, usize, DataType, RedOp, Comm, abi.Stream) callconv(.c) R, ncclAllGather: *const fn (D, D, usize, DataType, Comm, abi.Stream) callconv(.c) R, ncclBroadcast: *const fn (D, D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclSend: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclRecv: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, ncclGroupStart: *const fn () callconv(.c) R, ncclGroupEnd: *const fn () callconv(.c) R, }; diff --git a/zig/src/cuda/nvfp4.zig b/zig/src/cuda/nvfp4.zig index d9502895d3..e0d3c838f4 100644 --- a/zig/src/cuda/nvfp4.zig +++ b/zig/src/cuda/nvfp4.zig @@ -64,12 +64,13 @@ pub const Prompt = struct { mod: Module, fp4: Function, fp8: Function, + fp4_f32: Function, // FP4 writing the fp32 sums unrounded const bm = 128; const bn = 128; const threads = 128; - const name = "_ZN11tf_prompt1613prompt_kernelILi{d}ELi128ELi128ELi2ELi2ELi2ELb0EEEvPK13__nv_bfloat16PKhS5_fPviiiiii"; - pub const symbols = [2][:0]const u8{ std.fmt.comptimePrint(name, .{0}), std.fmt.comptimePrint(name, .{1}) }; + const name = "_ZN11tf_prompt1613prompt_kernelILi{d}ELi128ELi128ELi2ELi2ELi2ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPviiiiii"; + pub const symbols = [3][:0]const u8{ std.fmt.comptimePrint(name, .{ 0, 0 }), std.fmt.comptimePrint(name, .{ 1, 0 }), std.fmt.comptimePrint(name, .{ 0, 1 }) }; /// prompt16.cu Tile::SMEM. pub fn sharedBytes(mode: qmmf.Mode) u32 { @@ -86,7 +87,9 @@ pub const Prompt = struct { try fp4.allowDynamicShared(sharedBytes(.fp4)); const fp8 = try mod.function(symbols[1]); try fp8.allowDynamicShared(sharedBytes(.fp8)); - return .{ .mod = mod, .fp4 = fp4, .fp8 = fp8 }; + const fp4_f32 = try mod.function(symbols[2]); + try fp4_f32.allowDynamicShared(sharedBytes(.fp4)); + return .{ .mod = mod, .fp4 = fp4, .fp8 = fp8, .fp4_f32 = fp4_f32 }; } pub fn unload(p: *Prompt) void { @@ -95,9 +98,14 @@ pub const Prompt = struct { /// y [m, n] bf16 = prompt rows x [m, k] bf16 (rows `ldx` apart) times an FP4 or FP8 weight (linear._prompt). pub fn matmul(p: Prompt, s: Stream, x: u64, ldx: usize, m: usize, w: qmmf.Weight, y: u64) !void { + return p.matmulOut(s, x, ldx, m, w, y, .bf16); + } + + /// matmul() writing y as `out` says: bf16, or fp32 (FP4 only). + pub fn matmulOut(p: Prompt, s: Stream, x: u64, ldx: usize, m: usize, w: qmmf.Weight, y: u64, out: qmmf.Out) !void { const f = switch (w.mode) { - .fp4 => p.fp4, - .fp8 => p.fp8, + .fp4 => if (out == .f32) p.fp4_f32 else p.fp4, + .fp8 => if (out == .f32) return error.EpilogueNotBuilt else p.fp8, else => return error.FormatMismatch, }; const rows_t = (m + bm - 1) / bm; @@ -126,6 +134,7 @@ test "the prompt tile's shared bytes and instance names" { try std.testing.expectEqual(@as(u32, 41984), Prompt.sharedBytes(.fp4)); try std.testing.expectEqual(@as(u32, 49152), Prompt.sharedBytes(.fp8)); // 2 x (16384 + 8192): no scale tile try std.testing.expectEqualStrings("_ZN11tf_prompt1613prompt_kernelILi0ELi128ELi128ELi2ELi2ELi2ELb0EEEvPK13__nv_bfloat16PKhS5_fPviiiiii", Prompt.symbols[0]); + try std.testing.expectEqualStrings("_ZN11tf_prompt1613prompt_kernelILi0ELi128ELi128ELi2ELi2ELi2ELb1EEEvPK13__nv_bfloat16PKhS5_fPviiiiii", Prompt.symbols[2]); } test "word tiling reads each nibble where qmm.pack puts it" { diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..ff70a0d88a 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -91,16 +91,42 @@ pub const Linear = struct { } } + /// Decode rows as decode(), out (rows, n) fp32 unrounded: a tensor-parallel rank's partial rows (nvfp4 only). + pub fn decodeF32(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { + const q = switch (w) { + .nvfp4 => |q| q, + else => return error.EpilogueNotBuilt, + }; + return l.onLaneAs(s, q, .fp4, in, out, .f32, rows); + } + + /// Prompt rows as prompt(), out (rows, n) fp32 unrounded (nvfp4 only). + pub fn promptF32(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { + const q = switch (w) { + .nvfp4 => |q| q, + else => return error.EpilogueNotBuilt, + }; + return l.onPromptAs(s, q, .fp4, in, out, .f32, rows); + } + fn onPrompt(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, rows: usize) !void { + return l.onPromptAs(s, q, mode, in, out, .bf16, rows); + } + + fn onPromptAs(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, kind: qmmf.Out, rows: usize) !void { const p = l.nvfp4_prompt orelse return error.FormatNotLoaded; if (q.mode != mode) return error.FormatMismatch; - return p.matmul(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out); + return p.matmulOut(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, kind); } fn onLane(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, rows: usize) !void { + return l.onLaneAs(s, q, mode, in, out, .bf16, rows); + } + + fn onLaneAs(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, kind: qmmf.Out, rows: usize) !void { const lane = l.lane orelse return error.FormatNotLoaded; if (q.mode != mode) return error.FormatMismatch; - return lane.matmul(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, in.part); + return lane.matmulOut(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, kind, in.part); } }; @@ -331,4 +357,7 @@ test "a lane weight whose mode is not its format's is refused before any launch" const s: Stream = undefined; try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .fp8g = fp4 }, .{ .x = 0 }, 0, 1)); try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, .{ .nvfp4 = fp4 }, .{ .x = 0 }, 0, 1)); + const f8: qmmf.Weight = .{ .mode = .fp8g, .codes = 0, .scales = 0, .n = 64, .k = 64, .npad = 128 }; + try std.testing.expectError(error.EpilogueNotBuilt, l.decodeF32(s, .{ .fp8g = f8 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.EpilogueNotBuilt, l.promptF32(s, .{ .bf16 = f8 }, .{ .x = 0 }, 0, 1)); } diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 054e2c34bb..38e3e43a33 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -66,35 +66,43 @@ pub fn plan(m: usize, n: usize, k: usize, clusters: bool) Plan { } /// qmmf.cu's instance names, per mode: row tiles 16, 32, 64 (each plain, then with a cluster), then the fused kernel. -fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool) [:0]const u8 { - const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb0ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; - return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(cluster), @intFromBool(fused) }); +fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool, comptime wide: bool) [:0]const u8 { + const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb{d}ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; + return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(wide), @intFromBool(cluster), @intFromBool(fused) }); } -pub const reduce_symbol = "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif"; +pub const reduce_symbols = [2][:0]const u8{ "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif", "_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif" }; const row_tiles = [_]u32{ 16, 32, 64 }; +/// What a call writes: bf16 rows, or the fp32 sums unrounded (FP4 only, a tensor-parallel rank's partial rows). +pub const Out = enum { bf16, f32 }; + +fn loadTiles(mod: Module, comptime mode: Mode, comptime wide: bool) ![7]Function { + var set: [7]Function = undefined; + inline for (row_tiles, 0..) |bm, ti| inline for (.{ false, true }, 0..) |cl, ci| { + const f = try mod.function(symbol(mode, bm, cl, false, wide)); + try f.allowDynamicShared(sharedBytes(mode, bm)); + set[ti * 2 + ci] = f; + }; + set[6] = try mod.function(symbol(mode, 64, false, true, wide)); + try set[6].allowDynamicShared(sharedBytes(mode, 64)); + return set; +} + pub const Lane = struct { mod: Module, fns: [modes.len][7]Function, // per mode: lane16, lane16c, lane32, lane32c, lane64, lane64c, fused - reduce: Function, + fp4_f32: [7]Function, // FP4's tiles writing fp32 + reduce: [2]Function, // by Out clusters: bool, // sm_90 on: slices meet in a cluster's shared memory pub fn load(d: *const Driver, capability_major: u32) !Lane { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.qmmf); errdefer mod.unload(); - var l: Lane = .{ .mod = mod, .fns = undefined, .reduce = try mod.function(reduce_symbol), .clusters = capability_major >= 9 }; - inline for (modes, 0..) |mode, mi| { - inline for (row_tiles, 0..) |bm, ti| inline for (.{ false, true }, 0..) |cl, ci| { - const f = try mod.function(symbol(mode, bm, cl, false)); - try f.allowDynamicShared(sharedBytes(mode, bm)); - l.fns[mi][ti * 2 + ci] = f; - }; - const f = try mod.function(symbol(mode, 64, false, true)); - try f.allowDynamicShared(sharedBytes(mode, 64)); - l.fns[mi][6] = f; - } + var l: Lane = .{ .mod = mod, .fns = undefined, .fp4_f32 = try loadTiles(mod, .fp4, true), .reduce = undefined, .clusters = capability_major >= 9 }; + inline for (modes, 0..) |mode, mi| l.fns[mi] = try loadTiles(mod, mode, false); + for (&l.reduce, reduce_symbols) |*f, name| f.* = try mod.function(name); return l; } @@ -110,8 +118,14 @@ pub const Lane = struct { /// y [m, n] bf16 = x [m, k] bf16 (rows `ldx` apart) times the weight, fp32 sums; `part` holds partBytes. pub fn matmul(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64) !void { + return l.matmulOut(s, x, ldx, m, w, y, .bf16, part); + } + + /// matmul() writing y as `out` says: bf16, or fp32 (FP4 only). + pub fn matmulOut(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, out: Out, part: u64) !void { + if (out == .f32 and w.mode != .fp4) return error.EpilogueNotBuilt; const p = plan(m, w.n, w.k, l.clusters); - const set = l.fns[modeIndex(w.mode)]; + const set = if (out == .f32) l.fp4_f32 else l.fns[modeIndex(w.mode)]; const ti: usize = if (p.bm == 16) 0 else if (p.bm == 32) 1 else 2; const f = if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)]; const rows_t = (m + p.bm - 1) / p.bm; @@ -144,7 +158,7 @@ pub const Lane = struct { r.add(@as(i64, @intCast(total))); r.add(@as(i32, @intCast(p.sk))); r.add(w.scale); - try launch_.launch(l.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); + try launch_.launch(l.reduce[@backingInt(out)], .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); } } }; @@ -209,6 +223,7 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { } test "instance names are qmmf.cu's" { - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false)); - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 32, true, false, true)); } diff --git a/zig/src/families/glm/config.zig b/zig/src/families/glm/config.zig index 87ce6195e0..314521534e 100644 --- a/zig/src/families/glm/config.zig +++ b/zig/src/families/glm/config.zig @@ -13,6 +13,7 @@ pub const Config = struct { layers: u32 = 45, run: u32 = 45, // backbone layers loaded and run: all, or the first few for a check that must fit (GLM_LAYERS) vocab: u32 = 154880, + head_unit: u32 = 64, // the vocabulary's TP cut in whole units: 128 for EXL3, whose matrices split at 128 columns eps: f32 = 1e-5, dense_layers: u32 = 3, dense_inter: u32 = 12288, // this Mac's dense MLP rows: all, or its half in TP2 @@ -21,7 +22,9 @@ pub const Config = struct { inter: [2]u32 = .{ 0, 2048 }, // each routed expert's intermediate rows [lo, hi) this Mac holds (half by rows) topk: u32 = 8, moe_inter: u32 = 2048, + shared_inter: u32 = 2048, // the shared expert's width: moe_intermediate_size x n_shared_experts routed_scale: f32 = 2.5, + norm_topk: bool = true, // norm_topk_prob: the routed weights renormalized over the top-k swiglu_limit: f32 = 10.0, kda_heads: u32 = 64, // this Mac's KDA heads: all, or its share in TP2 (heads [tp_rank kda_heads, ...)) kda_dim: u32 = 128, @@ -134,7 +137,7 @@ fn float(o: std.json.ObjectMap, key: []const u8) !f64 { }; } -/// The checkpoint's text_config, refused unless it is GLM-5.3-Flash's shape (the kernels are built for it). +/// text_config with GLM-5.3-Flash's widths; smaller counts pass (synthetic), experts >= top-k, vocab a multiple of 128. pub fn parse(gpa: std.mem.Allocator, json: []const u8) !Config { const parsed = try std.json.parseFromSlice(std.json.Value, gpa, json, .{}); defer parsed.deinit(); @@ -146,6 +149,9 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8) !Config { c.layers = @intCast(try int(t, "num_hidden_layers")); c.run = c.layers; c.vocab = @intCast(try int(t, "vocab_size")); + if (root.get("quantization_config")) |q| if (q == .object) if (q.object.get("quant_method")) |m| if (m == .string and std.mem.eql(u8, m.string, "exl3")) { + c.head_unit = 128; + }; c.eps = @floatCast(try float(t, "rms_norm_eps")); c.dense_layers = @intCast(try int(t, "first_k_dense_replace")); c.dense_inter = @intCast(try int(t, "intermediate_size")); @@ -154,6 +160,8 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8) !Config { c.inter = .{ 0, c.moe_inter }; c.topk = @intCast(try int(t, "num_experts_per_tok")); c.moe_inter = @intCast(try int(t, "moe_intermediate_size")); + c.shared_inter = c.moe_inter * @as(u32, @intCast(if (t.get("n_shared_experts")) |v| v.integer else 1)); + if (t.get("norm_topk_prob")) |v| c.norm_topk = v != .bool or v.bool; c.routed_scale = @floatCast(try float(t, "routed_scaling_factor")); c.swiglu_limit = @floatCast(try float(t, "swiglu_limit")); c.mla_heads = @intCast(try int(t, "num_attention_heads")); @@ -187,15 +195,18 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8) !Config { }; const rope = t.get("qk_rope_head_dim"); const plain = (if (t.get("n_group")) |v| v.integer else 1) == 1 and (if (t.get("topk_group")) |v| v.integer else 1) == 1; - const same = c.hidden == want.hidden and c.layers == want.layers and c.vocab == want.vocab and - c.dense_layers == want.dense_layers and c.dense_inter == want.dense_inter and c.experts == want.experts and - c.topk == want.topk and c.moe_inter == want.moe_inter and c.mla_heads == want.mla_heads and + const widths = c.hidden == want.hidden and c.dense_inter == want.dense_inter and c.topk == want.topk and + c.moe_inter == want.moe_inter and c.shared_inter == want.shared_inter and c.mla_heads == want.mla_heads and c.nope == want.nope and c.v_dim == want.v_dim and c.q_lora == want.q_lora and c.kv_lora == want.kv_lora and c.i_heads == want.i_heads and c.i_dim == want.i_dim and c.i_topk == want.i_topk and c.kpool == want.kpool and c.hc == want.hc and c.sinkhorn == want.sinkhorn and c.kda_heads == want.kda_heads and - c.kda_dim == want.kda_dim and c.conv == want.conv and c.mla.eql(want.mla) and plain and + c.kda_dim == want.kda_dim and c.conv == want.conv and plain and (rope == null or (rope.? == .integer and rope.?.integer == 0)); - if (!same) return error.NotGlm53Flash; + var eos_in = true; + for (c.eos[0..c.eos_n]) |id| eos_in = eos_in and id < c.vocab; + const counts = c.layers >= 1 and c.layers <= want.layers and c.dense_layers <= c.layers and + c.experts >= c.topk and c.experts <= want.experts and c.vocab <= want.vocab and c.vocab % 128 == 0 and eos_in; + if (!widths or !counts) return error.NotGlm53Flash; return c; } @@ -245,6 +256,45 @@ test "GLM-5.3-Flash's layer kinds: 34 KDA and 11 MLA in the backbone, the MTP la try std.testing.expectEqual(@as(u32, 10), c.kindIndex(43)); } +/// GLM-5.3-Flash's text_config with the counts given (the widths as the checkpoint has them), for the tests. +pub fn textConfig(gpa: std.mem.Allocator, layers: u32, dense: u32, mla: []const u32, experts: u32, vocab: u32, hidden: u32) ![]u8 { + var list: std.ArrayList(u8) = .empty; + defer list.deinit(gpa); + for (mla, 0..) |l, i| try list.print(gpa, "{s}{d}", .{ if (i > 0) ", " else "", l }); + return std.fmt.allocPrint(gpa, + \\{{"model_type": "glm5_next", "text_config": {{"hidden_size": {d}, "num_hidden_layers": {d}, "vocab_size": {d}, + \\"rms_norm_eps": 1e-05, "first_k_dense_replace": {d}, "intermediate_size": 12288, "n_routed_experts": {d}, + \\"num_experts_per_tok": 8, "moe_intermediate_size": 2048, "n_shared_experts": 1, "norm_topk_prob": true, + \\"routed_scaling_factor": 2.5, "swiglu_limit": 10.0, "num_attention_heads": 64, "qk_nope_head_dim": 256, + \\"v_head_dim": 256, "q_lora_rank": 1536, "kv_lora_rank": 512, "index_n_heads": 32, "index_head_dim": 128, + \\"index_topk": 2048, "index_kpool": 4, "hc_mult": 4, "hc_eps": 1e-06, "hc_sinkhorn_iters": 20, + \\"num_nextn_predict_layers": 1, "eos_token_id": [{d}], "linear_attn_config": {{"num_heads": 64, "head_dim": 128, + \\"short_conv_kernel_size": 4, "gate_lower_bound": -5.0, "full_attn_layers": [{s}]}}}}}} + , .{ hidden, layers, vocab, dense, experts, vocab - 2, list.items }); +} + +test "the real counts parse to the defaults; smaller counts parse, other widths do not" { + const gpa = std.testing.allocator; + const real = try textConfig(gpa, 45, 3, &.{ 3, 7, 11, 15, 19, 23, 27, 31, 35, 39, 43 }, 288, 154880, 4096); + defer gpa.free(real); + const c = try parse(gpa, real); + try std.testing.expect(c.mla.eql((Config{}).mla)); + try std.testing.expectEqual(@as(u32, 45), c.layers); + const small = try textConfig(gpa, 2, 1, &.{1}, 16, 8192, 4096); + defer gpa.free(small); + const s = try parse(gpa, small); + try std.testing.expectEqual(Kind.kda, s.kind(0)); + try std.testing.expectEqual(Kind.mla, s.kind(1)); + try std.testing.expectEqual(Kind.mla, s.kind(2)); // the MTP layer + try std.testing.expectEqual(@as(u32, 16), s.experts); + const wide = try textConfig(gpa, 2, 1, &.{1}, 16, 8192, 2048); + defer gpa.free(wide); + try std.testing.expectError(error.NotGlm53Flash, parse(gpa, wide)); + const few = try textConfig(gpa, 2, 1, &.{1}, 4, 8192, 4096); + defer gpa.free(few); + try std.testing.expectError(error.NotGlm53Flash, parse(gpa, few)); +} + test "a layer subset keeps the MTP layer's place and refuses an empty or oversized count" { var c = Config{}; try subset(&c, 8); diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig new file mode 100644 index 0000000000..af9254de61 --- /dev/null +++ b/zig/src/families/glm/cuda.zig @@ -0,0 +1,49 @@ +//! GLM-5.3-Flash on CUDA: the family's parts, ported from the Python family's CUDA path. + +pub const experts = @import("cuda_experts.zig"); +pub const glue = @import("cuda_glue.zig"); +pub const kda = @import("cuda_kda.zig"); +pub const b16 = @import("cuda_b16.zig"); +pub const latent = @import("cuda_latent.zig"); +pub const sparse = @import("cuda_sparse.zig"); +pub const moe = @import("cuda_moe.zig"); +pub const kda_block = @import("cuda_kda_block.zig"); +pub const dsa_block = @import("cuda_dsa_block.zig"); +pub const mlp = @import("cuda_mlp.zig"); +pub const shard = @import("cuda_shard.zig"); +pub const weights = @import("cuda_weights.zig"); +pub const forward = @import("cuda_forward.zig"); +pub const config = @import("config.zig"); +pub const quant = @import("cuda_quant.zig"); +pub const q4 = @import("cuda_q4.zig"); +pub const ranks = @import("cuda_ranks.zig"); +pub const sampler = @import("cuda_sampler.zig"); +pub const engine = @import("cuda_engine.zig"); +pub const follow = @import("cuda_follow.zig"); +pub const lanes = @import("cuda_lanes.zig"); +pub const native = @import("cuda_native.zig"); + +test { + _ = experts; + _ = glue; + _ = kda; + _ = b16; + _ = latent; + _ = sparse; + _ = moe; + _ = kda_block; + _ = dsa_block; + _ = mlp; + _ = shard; + _ = weights; + _ = forward; + _ = quant; + _ = q4; + _ = ranks; + _ = sampler; + _ = engine; + _ = follow; + _ = lanes; + _ = native; + _ = @import("heat"); // upstream's heat bands, run where the GLM ranks use them +} diff --git a/zig/src/families/glm/cuda_b16.zig b/zig/src/families/glm/cuda_b16.zig new file mode 100644 index 0000000000..d935b26a64 --- /dev/null +++ b/zig/src/families/glm/cuda_b16.zig @@ -0,0 +1,106 @@ +//! GLM-5.3-Flash's BF16 matmuls on CUDA (attention, router, shared experts): the Python family's B16 path. + +const std = @import("std"); +const cuda = @import("cuda"); +const glue = @import("cuda_glue.zig"); + +/// qmm.SPLIT_TARGET, B16_SPLIT_COLS, B16_BK, geometry.MLA_B16_SLICES and MLA_B16_ROWS_FROM. +pub const split_target = 192; +pub const split_cols = 64; +pub const bk = 64; +pub const max_slices = 8; +pub const rows_from = 128; + +/// qmm.b16_split_k: K slices doubled while the 64-column tiles stay under 192 and each slice keeps 4 steps of 64. +pub fn splitK(n: usize, k: usize) usize { + const tiles = (n + split_cols - 1) / split_cols; + const steps = k / bk; + var sk: usize = 1; + while (sk < max_slices and tiles * sk < split_target and steps % (sk * 2) == 0 and steps / (sk * 2) >= 4) sk *= 2; + return sk; +} + +/// Triton's one-row _bmm (aot.json M = 1, GB10): BLOCK_N, stages by (n, k), one stage -> two; else tf_glm_bmm_decode. +const DecodeTile = struct { n: usize, k: usize, bn: u32, st: u32 }; +const decode_tiles = [_]DecodeTile{ + .{ .n = 4096, .k = 4096, .bn = 32, .st = 2 }, .{ .n = 4096, .k = 128, .bn = 32, .st = 2 }, + .{ .n = 4096, .k = 2816, .bn = 32, .st = 2 }, .{ .n = 4096, .k = 5632, .bn = 32, .st = 2 }, + .{ .n = 5376, .k = 1536, .bn = 32, .st = 2 }, .{ .n = 8192, .k = 1536, .bn = 32, .st = 2 }, + .{ .n = 12576, .k = 4096, .bn = 32, .st = 3 }, .{ .n = 4096, .k = 1024, .bn = 32, .st = 3 }, + .{ .n = 160, .k = 4096, .bn = 32, .st = 3 }, .{ .n = 1280, .k = 4096, .bn = 32, .st = 3 }, + .{ .n = 2816, .k = 128, .bn = 32, .st = 3 }, .{ .n = 2688, .k = 128, .bn = 32, .st = 3 }, + .{ .n = 4096, .k = 8192, .bn = 64, .st = 2 }, .{ .n = 4096, .k = 704, .bn = 64, .st = 2 }, + .{ .n = 4096, .k = 16384, .bn = 64, .st = 2 }, .{ .n = 4096, .k = 1536, .bn = 64, .st = 2 }, + .{ .n = 4096, .k = 2048, .bn = 64, .st = 2 }, .{ .n = 16384, .k = 1536, .bn = 64, .st = 2 }, + .{ .n = 2048, .k = 4096, .bn = 64, .st = 2 }, .{ .n = 1408, .k = 4096, .bn = 64, .st = 2 }, + .{ .n = 8192, .k = 128, .bn = 64, .st = 2 }, .{ .n = 24896, .k = 4096, .bn = 64, .st = 2 }, + .{ .n = 77440, .k = 4096, .bn = 64, .st = 2 }, .{ .n = 154880, .k = 4096, .bn = 64, .st = 2 }, +}; + +fn decodeTile(n: usize, k: usize) ?DecodeTile { + for (decode_tiles) |t| if (t.n == n and t.k == k) return t; + return null; +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// The fp32 split-K partials a window of `m` rows needs (0 when it sums in registers or takes one slice). +pub fn partBytes(m: usize, n: usize, k: usize) usize { + const sk = splitK(n, k); + return if (m >= rows_from or sk == 1) 0 else sk * m * n * 4; +} + +/// qmm.matmul for a B16 weight: out [m][n] bf16 (fp32 with `wide`) = x . w; 128+ rows sum in registers, else `part`. +pub fn matmul(set: glue.Set, stream: cuda.Stream, x: u64, x_stride: usize, w: u64, out: u64, part: u64, m: usize, n: usize, k: usize, wide: bool) !void { + if (k % bk != 0) return error.KNotMultipleOf64; + const sk = splitK(n, k); + var a: cuda.Args = .{}; + a.add(x); + a.add(w); + a.add(out); + if (m < rows_from) a.add(if (sk > 1) part else out); + inline for (.{ @as(i64, @intCast(m)), @as(i64, @intCast(x_stride)), @as(i64, @intCast(n)), @as(i64, @intCast(k)), @as(c_int, @intCast(sk)), @as(c_int, @intFromBool(wide)) }) |v| a.add(v); + if (m >= rows_from and x % 16 == 0 and x_stride % 8 == 0 and n % 8 == 0 and out % 16 == 0) { // the cp.async tiles + // One slice past 256 columns: 128 x 128 blocks, two an SM; K slices or narrow outputs: 128 x 64 blocks. + const one = sk == 1 and n > 256; + const f = if (one) set.glue.bmm_tile1 else set.glue.bmm_tile; + const bn: usize = if (one) 128 else 64; + const shared: u32 = if (one) glue.bmm_tile1_smem[@intFromBool(wide)] else glue.bmm_tile_smem; + return cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(cdiv(m, 128)), .y = @intCast(cdiv(n, bn)) }, .block = .{ .x = 256 }, .shared = shared }, stream, &a); + } + if (m >= rows_from) return cuda.launch.launch(set.glue.bmm_rows, .{ .grid = .{ .x = @intCast(cdiv(m, 128)), .y = @intCast(cdiv(n, 64)) }, .block = .{ .x = 128 } }, stream, &a); + if (decodeTile(n, k)) |t| { + if (x % 16 == 0 and x_stride % 8 == 0 and n % 8 == 0 and out % 16 == 0 and part % 16 == 0) { + const f = switch (t.bn * 10 + t.st) { + 321 => set.glue.bmm_d32s1, + 322 => set.glue.bmm_d32s2, + 323 => set.glue.bmm_d32s3, + 642 => set.glue.bmm_d64s2, + else => unreachable, + }; + const shared: u32 = t.st * (16 + t.bn) * bk * 2; // Tile<16, bn, 64, st, 1, 4>'s stages (the output tile is smaller) + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(cdiv(m, 16)), .y = @intCast(cdiv(n, t.bn)), .z = @intCast(sk) }, .block = .{ .x = 128 }, .shared = shared }, stream, &a); + if (sk > 1) try glue.sliceSum(set, stream, part, out, m * n, sk, wide); + return; + } + } + try cuda.launch.launch(set.glue.bmm_decode, .{ .grid = .{ .x = @intCast(cdiv(m, 16)), .y = @intCast(cdiv(n, 64)), .z = @intCast(sk) }, .block = .{ .x = 128 } }, stream, &a); + if (sk > 1) try glue.sliceSum(set, stream, part, out, m * n, sk, wide); +} + +/// qmm.group_sums: (m, k) bf16 rows (stride `x_stride`) -> (m, k / 64) fp32 sums of each 64 inputs. +pub fn groupSums(set: glue.Set, stream: cuda.Stream, x: u64, x_stride: usize, xs: u64, m: usize, k: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ x, xs, @as(i64, @intCast(x_stride)), @as(i64, @intCast(k / 64)) }) |v| a.add(v); + try cuda.launch.launch(set.glue.group_sums, .{ .grid = .{ .x = @intCast(m), .y = @intCast(cdiv(k / 64, 16)) }, .block = .{ .x = 128 } }, stream, &a); +} + +test "K slices as qmm.py picks them" { + try std.testing.expectEqual(@as(usize, 1), splitK(12576, 4096)); // 197 tiles: past the target at one slice + try std.testing.expectEqual(@as(usize, 4), splitK(4096, 4096)); // 64 tiles: 2, then 4 (256 past the target) + try std.testing.expectEqual(@as(usize, 8), splitK(160, 4096)); + try std.testing.expectEqual(@as(usize, 1), splitK(4096, 128)); // two steps: no slice keeps four + try std.testing.expectEqual(@as(usize, 0), partBytes(200, 4096, 4096)); +} diff --git a/zig/src/families/glm/cuda_dsa_block.zig b/zig/src/families/glm/cuda_dsa_block.zig new file mode 100644 index 0000000000..feb7402b29 --- /dev/null +++ b/zig/src/families/glm/cuda_dsa_block.zig @@ -0,0 +1,144 @@ +//! GLM-5.3-Flash's DSA block on CUDA, a rank's latent-path share (the Python family's dsa_block) to an fp32 partial. + +const std = @import("std"); +const cuda = @import("cuda"); +const glue = @import("cuda_glue.zig"); +const b16 = @import("cuda_b16.zig"); +const latent = @import("cuda_latent.zig"); +const sparse = @import("cuda_sparse.zig"); + +/// A prompt chunk's rows one dense latent attention call takes (geometry.MLA_PROMPT_ATT_ROWS). +pub const prompt_att_rows = 512; + +/// The indexer's weights, replicated: [wk | weights_proj] [d + heads][dims], gate, key LayerNorm, pool bias, wq_b. +pub const Index = struct { kw: u64, gate: u64, ln_w: u64, ln_b: u64, ape: u64, qb: u64, heads: usize, d: usize }; + +/// One DSA layer on a rank of `heads` heads: proj, norms, q_b, kv_b per head as wk/wv, o, indexer; shapes from config. +pub const Layer = struct { + proj: u64, + q_norm: u64, + kv_norm: u64, + q_b: u64, + wk: u64, + wv: u64, + o: u64, + index: Index, + heads: usize, + qk: usize, + dv: usize, + lw: usize, + q_lora: usize, + dims: usize, + eps: f32, + dense_limit: usize, // the largest context whose top-k selection keeps every visible key (Config.keyWidth) +}; + +/// The caches: latent bf16 [capacity][lw], index keys/gates bf16 [capacity][128], pooled keys bf16 [pk_rows][128]. +pub const Caches = struct { lc: u64, ik: u64, ig: u64, pk: u64, pk_rows: usize }; + +/// The block's scratch for up to `rows` rows. +pub const Scratch = struct { + dp: u64, // bf16 [rows][q_lora + lw] + qr: u64, // bf16 [rows][q_lora] + xs_qr: u64, // fp32 [rows][q_lora / 64] + lat: u64, // bf16 [rows][lw] + xs_lat: u64, // fp32 [rows][lw / 64] + q: u64, // bf16 [rows][heads qk] + ikr: u64, // bf16 [rows][160] + igr: u64, // fp32 [rows][128] + router_part: u64, // fp32 [router_ks][rows][128] + qa: u64, // bf16 [rows][heads][lw] + ol: u64, // bf16 [rows][heads][lw] + vn: u64, // bf16 [rows][heads dv] + qi: u64, // bf16 [rows][index heads d] + partials: latent.Partials, + dummy: u64, // an i32 + select: sparse.Scratch, + tokens: u64, // i32 [rows][sparse.width] + counts: u64, // i32 [rows] + mm_part: u64, // fp32 split-K partials + at: u64, // an i32: a prompt block's first position (pos + its first row) +}; + +/// The kernels: the captured Triton set and glm_ops (the token lists and position offsets). +pub const Kernels = struct { set: glue.Set, ops: cuda.Module }; + +/// forward.dsa_block (latent path) of `rows` rows at pos to an fp32 partial; `nch`: 512-key chunks, `part_rows`: block. +pub fn block(k: Kernels, stream: cuda.Stream, l: Layer, c: Caches, s: Scratch, x: u64, x_stride: usize, rows: usize, pos: u64, host_pos: usize, nch: usize, part_rows: usize, out: u64) !void { + try body(k, stream, l, c, s, x, x_stride, rows, pos, host_pos, nch, part_rows); + try fill(k.set, stream, l, s, 0, rows, out); +} + +/// forward.out_proj's fill for the DSA block: o_proj of rows lo..hi (s.vn) into out [hi - lo][dims] fp32. +pub fn fill(set: glue.Set, stream: cuda.Stream, l: Layer, s: Scratch, lo: usize, hi: usize, out: u64) !void { + const w = l.heads * l.dv; + try b16.matmul(set, stream, s.vn + lo * w * 2, w, l.o, out, s.mm_part, hi - lo, l.dims, w, true); +} + +/// The block up to its out_proj (s.vn holds the expanded rows). +pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, c: Caches, s: Scratch, x: u64, x_stride: usize, rows: usize, pos: u64, host_pos: usize, nch: usize, part_rows: usize) !void { + try prepare(k.set, stream, l, s, x, x_stride, rows, sparseRows(l, host_pos, rows)); + try attend(k, stream, l, c, s, 0, rows, pos, host_pos, nch, part_rows); + try latent.expandV(k.set, stream, s.ol, l.wv, s.vn, rows, l.heads, l.dv, l.lw); +} + +/// Whether a sequence at `host_pos` has rows past the dense limit among its next `rows` (the sparse pass). +pub fn sparseRows(l: Layer, host_pos: usize, rows: usize) bool { + return host_pos + rows - 1 >= l.dense_limit; +} + +/// The block's row-by-row part for a shared window: projections, norms, q_b, indexer, absorbed query, `qi`'s query. +pub fn prepare(set: glue.Set, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, qi: bool) !void { + const width = l.q_lora + l.lw; + try b16.matmul(set, stream, x, x_stride, l.proj, s.dp, s.mm_part, rows, width, l.dims, false); + try glue.rmsnorm(set, stream, s.dp, width, l.q_norm, s.qr, l.q_lora, s.xs_qr, l.eps, rows, l.q_lora); + try glue.rmsnorm(set, stream, s.dp + l.q_lora * 2, width, l.kv_norm, s.lat, l.lw, s.xs_lat, l.eps, rows, l.lw); + try b16.matmul(set, stream, s.qr, l.q_lora, l.q_b, s.q, s.mm_part, rows, l.heads * l.qk, l.q_lora, false); + const ix = l.index; + try b16.matmul(set, stream, x, x_stride, ix.kw, s.ikr, s.mm_part, rows, ix.d + ix.heads, l.dims, false); + try glue.router(set, stream, x, x_stride, ix.gate, s.igr, s.router_part, rows, l.dims, ix.d); + try latent.absorbQ(set, stream, s.q, l.wk, s.qa, rows, l.heads, l.qk, l.lw); + if (qi) try b16.matmul(set, stream, s.qr, l.q_lora, ix.qb, s.qi, s.mm_part, rows, ix.heads * ix.d, l.q_lora, false); +} + +/// One sequence's rows r0..r0 + rows at pos: latent and index keys written, dense then sparse attention into s.ol. +pub fn attend(k: Kernels, stream: cuda.Stream, l: Layer, c: Caches, s: Scratch, r0: usize, rows: usize, pos: u64, host_pos: usize, nch: usize, part_rows: usize) !void { + const set = k.set; + const ix = l.index; + const kw_rows = ix.d + ix.heads; + var v = s; // the scratch from row r0 + v.lat += r0 * l.lw * 2; + v.ikr += r0 * kw_rows * 2; + v.igr += r0 * ix.d * 4; + v.qa += r0 * l.heads * l.lw * 2; + v.ol += r0 * l.heads * l.lw * 2; + v.qi += r0 * ix.heads * ix.d * 2; + try latent.latentWrite(set, stream, v.lat, l.lw, c.lc, pos, rows, l.lw); + try sparse.indexUpdate(set, stream, v.ikr, kw_rows, v.igr, ix.ln_w, ix.ln_b, ix.ape, c.ik, c.ig, c.pk, pos, rows); + const scale: f32 = @floatCast(std.math.pow(f64, @floatFromInt(l.qk), -0.5)); + if (host_pos < l.dense_limit) try denseAttention(k, stream, l, c, v, pos, rows, nch, part_rows, scale); + if (sparseRows(l, host_pos, rows)) { + try sparse.selectTokens(set, k.ops, stream, v.qi, v.ikr + ix.d * 2, kw_rows, c.pk, pos, host_pos, rows, ix.heads, ix.d, c.pk_rows - 2, v.select, v.tokens, v.counts); + try latent.sparseAttention(set, stream, v.qa, c.lc, v.tokens, v.counts, sparse.width, v.partials, v.ol, rows, l.heads, l.lw, scale); + } +} + +/// forward.dense_attention: latent.attention in `part_rows`-row blocks, each at its first row's position. +fn denseAttention(k: Kernels, stream: cuda.Stream, l: Layer, c: Caches, s: Scratch, pos: u64, rows: usize, nch: usize, part_rows: usize, scale: f32) !void { + const hb = latent.headBlock(rows); + var r0: usize = 0; + while (r0 < rows) : (r0 += part_rows) { + const n = @min(part_rows, rows - r0); + var at = pos; + if (r0 > 0) { + var a: cuda.Args = .{}; + a.add(pos); + a.add(@as(c_int, @intCast(r0))); + a.add(s.at); + try cuda.launch.launch(try k.ops.function("tf_glm_add_i32"), .{ .grid = .{ .x = 1 }, .block = .{ .x = 32 } }, stream, &a); + at = s.at; + } + const off = r0 * l.heads * l.lw * 2; + try latent.attention(k.set, stream, s.qa + off, c.lc, at, s.partials, s.ol + off, s.dummy, n, l.heads, l.lw, nch, scale, hb); + } +} diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig new file mode 100644 index 0000000000..f7656a2d0b --- /dev/null +++ b/zig/src/families/glm/cuda_engine.zig @@ -0,0 +1,488 @@ +//! One TP rank's GLM-5.3-Flash CUDA engine (the Python family's decode.Engine); all ranks draw the same tokens. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const config = @import("config.zig"); +const shard = @import("cuda_shard.zig"); +const weights = @import("cuda_weights.zig"); +const fw = @import("cuda_forward.zig"); +const ranks = @import("cuda_ranks.zig"); +const sampler = @import("cuda_sampler.zig"); +const heat = @import("heat"); +const modelopt = @import("core").modelopt; + +pub const Options = struct { + capacity: usize, // tokens a sequence holds + mtp: bool, // load the MTP head (drafts) + settings: ranks.Settings, +}; + +/// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. +pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) bool }; + +/// The most drafts a chain holds (decode.DepthPolicy's top depth). +pub const max_drafts = 3; + +/// A prompt pass's chunk rows (decode.prefill's). +pub const prompt_rows = fw.prefill_rows; + +/// A shared verify's max rows and streams: under hc_pre's 64-row block change, 8 streams of 4 rows. +pub const shared_rows = 32; +pub const max_shared = 8; + +/// One stream's window of a shared verify: its slot, its ids, and its drawn tokens. +pub const Part = struct { slot: u32, ids: []const u32, out: []u32 }; + +/// A kept prompt state as the prompt cache holds it: the id every rank keeps its own copy under, and its position. +pub const Kept = struct { id: u64, at: u32 }; + +/// A stream's own sequence: state, sampling rule, MTP drafts; buffers are shared, one stream's round at a time. +const Seq = struct { + st: fw.State, + sampling: ?sampler.Sampling = null, // null: greedy + in_head: usize = 0, // drafted rows the MTP caches hold past the kept ones (st.mtp_drafted) + rows_at: u64 = 0, // its last verify's final-norm rows (bf16 [rows][hidden]), which the next draft reads +}; + +/// One rank's copy of a kept prompt state: the state arena's bytes and the host fields beside them. +const Snap = struct { buf: cuda.DeviceBuffer, pos: usize, mtp_len: usize, cur: []u1 }; + +/// Where a prompt pass resumes: a kept state (every rank holds `id`) at `at`, a multiple of the chunk rows. +pub const Resume = struct { id: u64, at: usize }; + +/// Rank 0's marks in a prompt pass: `call` may keep the state at a chunk start; other ranks copy at that chunk's gate. +pub const Marks = struct { at: []const u32, ctx: *anyopaque, call: *const fn (ctx: *anyopaque, at: u32) void }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + io: std.Io, + d: *const cuda.Driver, + c: config.Config, + share: shard.Share, + spans: []shard.Span, // every rank's vocabulary span + lib: *const cuda.nccl.Library, + comm: cuda.nccl.Comm, + settings: ranks.Settings, + m: fw.Modules, + stream: cuda.Stream, + pieces: fw.Pieces, + w: weights.Weights, + bp: fw.Buffers, // prompt chunks + bd: fw.Buffers, // verify windows + bm: fw.Buffers, // the MTP head's rows + bs: ?fw.Buffers = null, // shared verify windows (with the MTP head, as the lane core shares only then) + hold: cuda.DeviceBuffer, // the prompt's last row's hidden (decode's e.last_hidden), the first draft's source + pairs: cuda.DeviceBuffer, // every rank's (value, id) of a greedy draw + smp: sampler.Sampler, + seqs: std.ArrayList(?*Seq) = .empty, // every stream's sequence by slot (the lane backend's), slot 0 made at init + cur: *Seq = undefined, // the bound slot's: every step below runs on it + slot: u32 = 0, + capacity: usize, + mtp: bool, + heat_gate: heat.Gate, // upstream's bands (off unless set), read before each prompt chunk on the gathered maximum + shared_buf: cuda.DeviceBuffer, // one f64 every rank all-reduces to its maximum + snaps: std.AutoHashMapUnmanaged(u64, Snap) = .empty, // kept prompt states by id (the same ids on every rank) + next_snap: u64 = 1, // rank 0: the next id + announce: u64 = 0, // rank 0: a state kept at this chunk's mark, for the other ranks to copy at its gate + dropped: std.ArrayList(u32) = .empty, // rank 0: ids dropped since the other ranks last heard (cuda_follow) + + /// Opens rank `rank` of `world` over `comm`; the ranks agree on settings (cuda_ranks) before the weights load. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, ctx: *const cuda.Context, checkpoint: []const u8, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, rank: usize, world: usize, o: Options) !*Engine { + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const cfg_path = try std.fs.path.join(a, &.{ checkpoint, "config.json" }); + const text = try std.Io.Dir.cwd().readFileAlloc(io, cfg_path, a, .limited(1 << 24)); + var why: modelopt.Why = .{}; + modelopt.admit(gpa, text, &why) catch |err| { // a quant_algo other than NVFP4 is refused before anything loads + std.log.err("{s}: {s}", .{ cfg_path, why.text() }); + return err; + }; + const c = try config.parse(gpa, text); + var stream = try cuda.Stream.init(d, false); + errdefer stream.deinit(); + const fp = ranks.Fingerprint.of(world, rank, o.settings, o.capacity, try ranks.hashFile(a, io, cfg_path), try ranks.hashFile(a, io, "/proc/self/exe"), checkpoint); + try ranks.agree(d, lib, comm, stream, fp, try a.alloc(ranks.Fingerprint, world)); + const share = try shard.Share.of(&c, rank, world); + const spans = try gpa.alloc(shard.Span, world); + errdefer gpa.free(spans); + for (spans, 0..) |*sp, q| sp.* = (try shard.Share.of(&c, q, world)).vocab; + var m = try fw.Modules.load(d, ctx); + errdefer m.deinit(); + var pieces = try fw.Pieces.init(d, o.settings.pieces); + errdefer pieces.deinit(); + if (o.settings.layers > c.layers) return error.BadLayers; + const layers = try a.alloc(u32, if (o.settings.layers > 0) o.settings.layers else c.layers); + for (layers, 0..) |*l, i| l.* = @intCast(i); + e.* = .{ .gpa = gpa, .io = io, .heat_gate = .{ .bands = o.settings.heat, .root = heatRoot() }, .shared_buf = undefined, .d = d, .c = c, .share = share, .spans = spans, .lib = lib, .comm = comm, .settings = o.settings, .m = m, .stream = stream, .pieces = pieces, .w = undefined, .capacity = o.capacity, .bp = undefined, .bd = undefined, .bm = undefined, .hold = undefined, .pairs = undefined, .smp = undefined, .mtp = o.mtp }; + e.w = try weights.load(gpa, io, d, stream, e.m.tools(), checkpoint, c, share, layers, o.mtp); + errdefer e.w.deinit(); + try e.bind(0); + errdefer { + e.free(0); + e.seqs.deinit(gpa); + } + e.bp = try fw.Buffers.init(d, &e.w, fw.prefill_rows, true, o.capacity); + errdefer e.bp.deinit(); + e.bd = try fw.Buffers.init(d, &e.w, fw.window_rows, false, o.capacity); + errdefer e.bd.deinit(); + e.bm = try fw.Buffers.init(d, &e.w, fw.window_rows, false, o.capacity); + errdefer e.bm.deinit(); + if (o.mtp) e.bs = try fw.Buffers.initShared(d, &e.w, shared_rows, o.capacity); + errdefer if (e.bs) |*b| b.deinit(); + e.hold = try cuda.DeviceBuffer.alloc(d, c.hidden * 2); + errdefer e.hold.free(); + e.pairs = try cuda.DeviceBuffer.alloc(d, 8 * world); + errdefer e.pairs.free(); + e.smp = try sampler.Sampler.init(d, c.vocab, fw.window_rows); + errdefer e.smp.deinit(); + e.shared_buf = try cuda.DeviceBuffer.alloc(d, 8); + try e.stream.synchronize(); + return e; + } + + pub fn deinit(e: *Engine) void { + e.stream.synchronize() catch {}; + var it = e.snaps.valueIterator(); + while (it.next()) |s| e.freeSnap(s); + e.snaps.deinit(e.gpa); + e.dropped.deinit(e.gpa); + e.shared_buf.free(); + e.smp.deinit(); + e.pairs.free(); + e.hold.free(); + if (e.bs) |*b| b.deinit(); + e.bm.deinit(); + e.bd.deinit(); + e.bp.deinit(); + for (0..e.seqs.items.len) |i| e.free(i); + e.seqs.deinit(e.gpa); + e.w.deinit(); + e.pieces.deinit(); + e.m.deinit(); + e.stream.deinit(); + e.gpa.free(e.spans); + const gpa = e.gpa; + gpa.destroy(e); + } + + fn exchange(e: *const Engine) fw.Exchange { + return .{ .nccl = .{ .lib = e.lib, .comm = e.comm } }; + } + + fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { + out.* = .{.{ .w = &e.w, .st = &e.cur.st, .b = b }}; + return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .kda_wide = e.settings.kda_wide }; + } + + /// Binds slot `slot`'s sequence for the steps that follow, making it (a zeroed state) on first use. + pub fn bind(e: *Engine, slot: u32) !void { + if (slot >= e.seqs.items.len) try e.seqs.appendNTimes(e.gpa, null, slot + 1 - e.seqs.items.len); + if (e.seqs.items[slot] == null) { + const q = try e.gpa.create(Seq); + errdefer e.gpa.destroy(q); + q.* = .{ .st = try fw.State.init(e.gpa, e.d, &e.w, e.capacity) }; + e.seqs.items[slot] = q; + } + e.cur = e.seqs.items[slot].?; + e.slot = slot; + } + + /// Frees slot `slot`'s sequence (slot 0 stays: the first stream's, made at init) and binds slot 0. + pub fn release(e: *Engine, slot: u32) void { + if (slot != 0) e.free(slot); + e.cur = e.seqs.items[0].?; + e.slot = 0; + } + + fn free(e: *Engine, slot: usize) void { + const q = e.seqs.items[slot] orelse return; + e.stream.synchronize() catch {}; + q.st.deinit(); + e.gpa.destroy(q); + e.seqs.items[slot] = null; + } + + /// Bytes one more stream's sequence takes on this rank (its state arena). + pub fn seqBytes(e: *const Engine) u64 { + return e.seqs.items[0].?.st.arena.buf.len; + } + + /// The committed length: the pending token sits at this position, a window's row r draws position pos + 1 + r. + pub fn pos(e: *const Engine) usize { + return e.cur.st.pos; + } + + /// A fresh sequence with `sampling` (null or temperature 0: greedy). + pub fn begin(e: *Engine, sampling: ?sampler.Sampling) !void { + e.cur.sampling = sampler.check(sampling); + e.cur.in_head = 0; + try e.cur.st.reset(e.stream); + } + + /// decode.prefill: `ids` in 2,048-row chunks, then the first token from the last row; `from` restores a kept state. + pub fn prefill(e: *Engine, ids: []const u32, from: ?Resume, cancel: ?Cancel, marks: ?Marks) !u32 { + var start = try e.prefillStart(ids.len, from); + while (true) : (start += fw.prefill_rows) { + if (try e.prefillChunk(ids[start..@min(ids.len, start + fw.prefill_rows + 1)], start, ids.len, null, cancel, marks)) |first| return first; + } + } + + /// Where a prompt pass of `len` ids starts: 0, or `from`'s position with its kept state restored. + fn prefillStart(e: *Engine, len: usize, from: ?Resume) !usize { + if (len == 0 or len + fw.window_rows >= e.cur.st.capacity) return error.PromptTooLong; + const r = from orelse return 0; + if (r.at == 0 or r.at >= len or r.at % fw.prefill_rows != 0) return error.BadResume; + try e.restoreState(r.id); + return r.at; + } + + /// One chunk of a `len`-id prompt pass: `part` from `start` plus the next id; first token at the end, else null. + pub fn prefillChunk(e: *Engine, part: []const u32, start: usize, len: usize, from: ?Resume, cancel: ?Cancel, marks: ?Marks) !?u32 { + if ((from != null or start == 0) and try e.prefillStart(len, from) != start) return error.BadChunk; + if (start >= len or start % fw.prefill_rows != 0) return error.BadChunk; + const rows = @min(len - start, fw.prefill_rows); + if (part.len != @min(len - start, rows + 1)) return error.BadChunk; + var one: [1]fw.Rank = undefined; + const f = e.view(&e.bp, &one); + const as_i32: []const i32 = @ptrCast(part); + if (marks) |m| if (start > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(start)) != null) m.call(m.ctx, @intCast(start)); + try e.gate(cancel); + try f.run(as_i32[0..rows]); + const last = start + rows == len; + var first: ?u32 = null; + if (last) { + var t: [1]u32 = undefined; + try e.draw(e.bp.logits, 1, len, &t); + first = t[0]; + } + if (e.mtp) { + if (last) try e.copy(e.hold.ptr, e.bp.fnormed + (rows - 1) * e.c.hidden * 2, e.c.hidden * 2); + const nxt = as_i32[1..]; + if (nxt.len > 0) { // decode._absorb_rows + try f.mtp(nxt, &.{e.bp.fnormed}); + try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + nxt.len); + } + } + try f.commit(rows); + return first; + } + + /// The chunk gate on every rank: a kept state at the mark, the heat wait, rank 0's cancel; the id exchanged first. + fn gate(e: *Engine, cancel: ?Cancel) !void { + const kept = try e.shared(@floatFromInt(e.announce)); + e.announce = 0; + if (kept > 0 and e.share.rank != 0) try e.keepAs(@intFromFloat(kept)); + if (e.heat_gate.bands != null) { + var w: Wait = .{ .e = e, .cancel = cancel }; + const before = e.heat_gate.waited_s; + e.heat_gate.waitUntil(Wait.read, Wait.sleep, &w, .{ .ptr = &w, .check = Wait.stop }) catch |err| return w.failed orelse err; + if (e.heat_gate.waited_s > before) heat.note(e.heat_gate.waited_s - before); + } + if (try e.shared(if (cancel) |c| @floatFromInt(@intFromBool(c.check(c.ptr))) else 0) > 0) return error.Cancelled; + } + + /// Bytes one kept state takes on this rank: the whole state arena (its caches are sized for the window, not `at`). + pub fn stateBytes(e: *const Engine) u64 { + return e.cur.st.arena.buf.len; + } + + /// Rank 0, at a mark in a prompt pass: a copy of the state under a new id, named to the other ranks at the gate. + pub fn keepState(e: *Engine) !u64 { + const id = e.next_snap; + try e.keepAs(id); + e.next_snap += 1; + e.announce = id; + return id; + } + + fn keepAs(e: *Engine, id: u64) !void { + if (e.snaps.contains(id)) return error.StateKept; + var buf = try cuda.DeviceBuffer.alloc(e.d, e.cur.st.arena.buf.len); + errdefer buf.free(); + try e.copy(buf.ptr, e.cur.st.arena.buf.ptr, buf.len); + try e.stream.synchronize(); + const cur = try e.gpa.dupe(u1, e.cur.st.cur); + errdefer e.gpa.free(cur); + try e.snaps.put(e.gpa, id, .{ .buf = buf, .pos = e.cur.st.pos, .mtp_len = e.cur.st.mtp_len, .cur = cur }); + } + + /// The live state made kept state `id`'s; a rank without it refuses rather than prefill from another position. + fn restoreState(e: *Engine, id: u64) !void { + const s = e.snaps.getPtr(id) orelse return error.NoKeptState; + try e.copy(e.cur.st.arena.buf.ptr, s.buf.ptr, s.buf.len); + try e.cur.st.restoreHost(e.stream, s.pos, s.mtp_len, s.cur); + e.cur.in_head = 0; + try e.stream.synchronize(); + } + + pub fn holds(e: *const Engine, id: u64) bool { + return e.snaps.contains(id); + } + + /// Frees kept state `id`; rank 0 also notes it for the other ranks (cuda_follow sends it before its next step). + pub fn dropState(e: *Engine, id: u64) void { + var kv = e.snaps.fetchRemove(id) orelse return; + e.freeSnap(&kv.value); + if (e.share.rank == 0 and e.share.world > 1) e.dropped.append(e.gpa, @intCast(id)) catch + std.log.warn("kept state {d}: the other ranks keep their copy until they exit (no memory for the note)", .{id}); + } + + fn freeSnap(e: *Engine, s: *Snap) void { + s.buf.free(); + e.gpa.free(s.cur); + } + + /// The ranks' maximum of `v` (an all-reduce on the engine's stream). + fn shared(e: *Engine, v: f64) !f64 { + var x = v; + try e.shared_buf.uploadAsync(0, std.mem.asBytes(&x), e.stream.handle); + try e.lib.check(e.lib.api.ncclAllReduce(e.shared_buf.ptr, e.shared_buf.ptr, 1, .f64, .max, e.comm, e.stream.handle), "ncclAllReduce"); + try e.stream.synchronize(); + try e.shared_buf.download(0, std.mem.asBytes(&x)); + return x; + } + + /// heat.Gate.waitUntil's hooks over the ranks: gathered max reading, rank 0's stop: all pause together. + const Wait = struct { + e: *Engine, + cancel: ?Cancel, + failed: ?anyerror = null, + + fn read(ptr: *anyopaque) anyerror!f64 { + const w: *Wait = @ptrCast(@alignCast(ptr)); + const local: f64 = blk: { + var dir = std.Io.Dir.cwd().openDir(w.e.io, w.e.heat_gate.root, .{ .iterate = true }) catch break :blk -std.math.inf(f64); + defer dir.close(w.e.io); + break :blk heat.hottest(w.e.io, dir) catch -std.math.inf(f64); + }; + return w.e.shared(local); + } + + fn sleep(ptr: *anyopaque) void { + const w: *Wait = @ptrCast(@alignCast(ptr)); + std.Io.sleep(w.e.io, .fromMilliseconds(2000), .awake) catch {}; + } + + fn stop(ptr: *anyopaque) bool { + const w: *Wait = @ptrCast(@alignCast(ptr)); + const mine: f64 = if (w.cancel) |c| @floatFromInt(@intFromBool(c.check(c.ptr))) else 0; + const any = w.e.shared(mine) catch |err| { + w.failed = err; + return true; + }; + return any > 0; + } + }; + + /// A verify window: `ids` (pending token, then drafts) at the committed position; `out[r]` drawn at pos + 1 + r. + pub fn verify(e: *Engine, ids: []const u32, out: []u32) !void { + if (ids.len == 0 or ids.len > fw.window_rows) return error.WindowTooWide; + var one: [1]fw.Rank = undefined; + const f = e.view(&e.bd, &one); + const at = e.cur.st.pos; + try f.run(@ptrCast(ids)); + e.cur.rows_at = e.bd.fnormed; + try e.draw(e.bd.logits, ids.len, at + 1, out); + } + + /// Several streams' windows in one forward: shared row ops, per-part attention and KDA. + pub fn verifyShared(e: *Engine, parts: []const Part) !void { + const b = if (e.bs) |*x| x else return error.NoSharedWindows; + if (parts.len == 0 or parts.len > max_shared) return error.WindowTooWide; + var segs: [max_shared]fw.Seg = undefined; + var ids: [shared_rows]u32 = undefined; + var rows: usize = 0; + for (parts, 0..) |p, i| { + if (p.ids.len == 0 or p.ids.len > fw.window_rows or rows + p.ids.len > shared_rows) return error.WindowTooWide; + for (parts[0..i]) |q| if (q.slot == p.slot) return error.SlotTwice; + try e.bind(p.slot); + segs[i] = .{ .st = &e.cur.st, .row0 = rows, .rows = p.ids.len }; + @memcpy(ids[rows..][0..p.ids.len], p.ids); + rows += p.ids.len; + } + var one: [1]fw.Rank = .{.{ .w = &e.w, .st = segs[0].st, .b = b, .segs = segs[0..parts.len] }}; + const f: fw.Forward = .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = &one, .exchange = e.exchange(), .pieces = null, .split = e.settings.split, .kda_wide = e.settings.kda_wide }; + try f.run(@ptrCast(ids[0..rows])); + for (parts, segs[0..parts.len]) |p, g| { + try e.bind(p.slot); + e.cur.rows_at = b.fnormed + g.row0 * e.c.hidden * 2; + try e.draw(b.logits + g.row0 * e.share.vocab.len() * 2, g.rows, g.st.pos + 1, p.out); + } + } + + /// forward.commit of the last window's first `n` rows (the KDA replay when fewer than it ran). + pub fn keep(e: *Engine, n: usize) !void { + var one: [1]fw.Rank = undefined; + const f = e.view(&e.bd, &one); + try f.commit(n); + } + + /// decode.draft: the MTP head absorbs the kept rows (committed first) with `follow`, then chains `depth` drafts. + pub fn draft(e: *Engine, from_prompt: bool, follow: []const u32, depth: usize, out: []u32) !void { + if (!e.mtp) return error.NoDraftHead; + if (depth > max_drafts) return error.TooDeep; + var one: [1]fw.Rank = undefined; + const f = e.view(&e.bm, &one); + const D = e.c.hidden; + const hidden = if (from_prompt) e.hold.ptr else e.cur.rows_at; + try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len - e.cur.in_head); + e.cur.in_head = 0; + const next: []const i32 = @ptrCast(follow); + var s0: usize = 0; + while (s0 < next.len) : (s0 += fw.window_rows) { + const part = next[s0..@min(next.len, s0 + fw.window_rows)]; + try f.mtp(part, &.{hidden + s0 * D * 2}); + try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + part.len); + } + for (0..depth) |j| { + try e.draw(e.bm.logits, 1, e.cur.st.pos + 1 + j, out[j..][0..1]); + if (j + 1 < depth) { + try f.mtp(@ptrCast(out[j..][0..1]), &.{e.bm.fnormed}); + try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + 1); + e.cur.in_head += 1; + } + } + } + + /// `rows` logits rows of this rank's vocabulary -> the tokens at positions first, first + 1, ... + fn draw(e: *Engine, logits: u64, rows: usize, first: u64, out: []u32) !void { + if (e.cur.sampling) |s| return e.smp.draw(e.stream, e.exchange(), e.spans, &.{.{ .rank = e.share.rank, .logits = logits }}, rows, first, sampler.rule(s), out); + for (0..rows) |r| out[r] = try e.greedy(logits + r * e.share.vocab.len() * 2); + } + + /// decode.sample_rows' greedy pick: per-rank top-1 (tf_glm_top1), all-gathered, largest value then lowest id. + fn greedy(e: *Engine, logits: u64) !u32 { + const world = e.share.world; + var a: cuda.Args = .{}; + a.add(logits); + a.add(@as(i64, @intCast(e.share.vocab.len()))); + a.add(@as(u32, @intCast(e.share.vocab.lo))); + a.add(e.pairs.ptr + e.share.rank * 8); + try cuda.launch.launch(e.m.glue.top1, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, e.stream, &a); + try e.lib.check(e.lib.api.ncclAllGather(e.pairs.ptr + e.share.rank * 8, e.pairs.ptr, 2, .u32, e.comm, e.stream.handle), "ncclAllGather"); + try e.stream.synchronize(); + var all: [16][2]u32 = undefined; + try e.d.check(e.d.api.cuMemcpyDtoH_v2(&all, e.pairs.ptr, world * 8), "cuMemcpyDtoH"); + var pick = all[0]; + for (all[1..world]) |p| { + const v: f32 = @bitCast(p[0]); + const b: f32 = @bitCast(pick[0]); + if (v > b or (v == b and p[1] < pick[1])) pick = p; + } + return pick[1]; + } + + fn heatRoot() []const u8 { + const p = std.c.getenv("TF_HEAT_ROOT") orelse return heat.default_root; + const s = std.mem.span(p); + return if (s.len == 0) heat.default_root else s; + } + + fn copy(e: *const Engine, dst: u64, src: u64, bytes: usize) !void { + try e.d.check(e.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, e.stream.handle), "cuMemcpyDtoDAsync"); + } +}; diff --git a/zig/src/families/glm/cuda_experts.zig b/zig/src/families/glm/cuda_experts.zig new file mode 100644 index 0000000000..645b441675 --- /dev/null +++ b/zig/src/families/glm/cuda_experts.zig @@ -0,0 +1,57 @@ +//! GLM-5.3-Flash's routed NVFP4 experts on CUDA (the Python family's experts): grouped.Router, experts.Grouped. + +const std = @import("std"); +const cuda = @import("cuda"); + +/// Pairs an item holds on a decode plan (tensorfold.cuda.experts.TILE). +pub const tile = 16; +/// Pairs an item holds on a prompt plan (experts.py STAGED_TILE, experts.promptTile). +pub const staged_tile = 64; +pub const maxItems = cuda.grouped.maxItems; + +pub const Which = enum(u2) { gate_up, down_f32, down }; + +/// Upstream's NVFP4 experts and plan kernels. +pub const Kernels = struct { + nvfp4: cuda.experts.Nvfp4Experts, + router: cuda.grouped.Router, + + /// `plans`: the loaded experts module (its plan kernels). + pub fn load(d: *const cuda.Driver, plans: cuda.Module, sms: u32) !Kernels { + var nvfp4 = try cuda.experts.Nvfp4Experts.load(d, sms); + errdefer nvfp4.unload(); + return .{ .nvfp4 = nvfp4, .router = try cuda.grouped.Router.resolve(plans) }; + } + + pub fn unload(k: *Kernels) void { + k.nvfp4.unload(); + } +}; + +/// A grouped plan (tensorfold.cuda.experts.Plan): the router's arrays, the slots a row holds, the experts planned. +pub const Plan = struct { route: cuda.grouped.Plan, slots: usize, experts: usize, prompt: bool }; + +/// experts.route: `pairs` picks (i32) grouped by expert into the plan's items, 16 pairs a decode item, 64 a prompt one. +pub fn route(ks: *const Kernels, stream: cuda.Stream, picks: u64, pairs: usize, p: Plan) !void { + return ks.router.route(stream, picks, pairs, p.experts, if (p.prompt) staged_tile else tile, p.route); +} + +/// nvx.gate_up / down (prompt_* on a prompt plan): `x` rows -> `out` bf16 SwiGLU, fp32 or bf16; `skip` untouched. +pub fn run(ks: *const Kernels, stream: cuda.Stream, which: Which, l: cuda.experts.Layer, plan: Plan, x: u64, x_stride: usize, rows: usize, out: u64, skip: i32) !void { + const g: cuda.experts.Grouped = .{ .nvfp4 = &ks.nvfp4 }; + const epi: cuda.experts.Epi = switch (which) { + .gate_up => .up, + .down_f32 => .down_f32, + .down => .down_bf16, + }; + const in: cuda.experts.Rows = .{ .x = x, .stride = x_stride, .slots = plan.slots }; + const pairs = rows * plan.slots; + if (plan.prompt) return g.prompt(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, staged_tile), skip); + return g.decode(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, tile), skip); +} + +test "plan sizes as the Python family's" { + try std.testing.expectEqual(@as(usize, 1 + 0), maxItems(1, 257, tile)); + try std.testing.expectEqual(@as(usize, 9 + 2), maxItems(45, 9, tile)); // 9 experts used, 45 pairs: two more items of 16 + try std.testing.expectEqual(@as(usize, staged_tile), cuda.experts.promptTile(.nvfp4)); +} diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig new file mode 100644 index 0000000000..8072b7cfe1 --- /dev/null +++ b/zig/src/families/glm/cuda_follow.zig @@ -0,0 +1,279 @@ +//! TP over processes until #552: rank 0's Leader sends each step (header, ids, slot) first; other ranks `follow` it. + +const std = @import("std"); +const cuda = @import("cuda"); +const engine = @import("cuda_engine.zig"); +const sampler = @import("cuda_sampler.zig"); +const Engine = engine.Engine; + +pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared }; + +/// One step as it crosses: 64 bytes, then `n` u32 ids. +pub const Header = extern struct { + kind: u32, + n: u32 = 0, // ids that follow + depth: u32 = 0, // draft: drafts to chain; chunk: where it starts; shared: its parts + flags: u32 = 0, // begin: sampled; draft: from the prompt + a: u64 = 0, // keep: rows; prefill, chunk: the kept state it resumes (0: none) + seed: u64 = 0, // begin: the sampling seed; prefill, chunk: the kept state's position + temperature: f64 = 0, + top_p: f64 = 0, + min_p: f64 = 0, + top_k: u32 = 0, // begin: top_k; chunk: the prompt's length + slot: u32 = 0, // the stream's sequence the step runs on (Engine.bind) + + comptime { + std.debug.assert(@sizeOf(Header) == 64); + } +}; + +/// tf_glm_put_words' argument: a step's header (16 words) and up to `put_ids` ids. +pub const Words = extern struct { w: [96]u32 }; +pub const put_ids = 96 - 16; + +/// The most ids one step carries (a prompt at the engine's capacity). +pub const Link = struct { + d: *const cuda.Driver, + lib: *const cuda.nccl.Library, + comm: cuda.nccl.Comm, + stream: cuda.Stream, + buf: cuda.DeviceBuffer, // header, then ids + ids: []u32, // host copy of the last step's ids + gpa: std.mem.Allocator, + wait: bool, // TF_GLM_LINK_WAIT=1: every step uploaded and waited for, as before the put + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, stream: cuda.Stream, max_ids: usize) !Link { + const buf = try cuda.DeviceBuffer.alloc(d, 64 + max_ids * 4); + const wait = if (std.c.getenv("TF_GLM_LINK_WAIT")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; + return .{ .d = d, .lib = lib, .comm = comm, .stream = stream, .buf = buf, .ids = try gpa.alloc(u32, max_ids), .gpa = gpa, .wait = wait }; + } + + pub fn deinit(l: *Link) void { + l.gpa.free(l.ids); + l.buf.free(); + } + + fn bcast(l: *Link, bytes: usize) !void { + try l.post(bytes); + try l.stream.synchronize(); + } + + fn post(l: *Link, bytes: usize) !void { + try l.lib.check(l.lib.api.ncclBroadcast(l.buf.ptr, l.buf.ptr, bytes, .u8, 0, l.comm, l.stream.handle), "ncclBroadcast"); + } + + /// Rank 0: the step and ids to every rank; up to `put_ids` ids go by `put` unwaited, longer ones upload and wait. + pub fn send(l: *Link, put: cuda.Function, h: Header, ids: []const u32) !void { + if (ids.len > l.ids.len) return error.TooManyIds; + var hh = h; + hh.n = @intCast(ids.len); + if (!l.wait and ids.len <= put_ids) { + var w: Words = undefined; + @memcpy(w.w[0..16], std.mem.bytesAsSlice(u32, std.mem.asBytes(&hh))); + @memcpy(w.w[16..][0..ids.len], ids); + var a: cuda.Args = .{}; + a.add(l.buf.ptr); + a.add(@as(u32, @intCast(16 + ids.len))); + a.add(w); + try cuda.launch.launch(put, .{ .grid = .{ .x = 1 }, .block = .{ .x = 128 } }, l.stream, &a); + try l.post(64); + if (ids.len > 0) try l.post(64 + ids.len * 4); + return; + } + try l.buf.uploadAsync(0, std.mem.asBytes(&hh), l.stream.handle); + try l.bcast(64); + if (ids.len > 0) { + try l.buf.uploadAsync(64, std.mem.sliceAsBytes(ids), l.stream.handle); + try l.bcast(64 + ids.len * 4); + } + } + + /// Other ranks: the next step and its ids (valid until the next receive). + pub fn receive(l: *Link) !struct { h: Header, ids: []const u32 } { + try l.bcast(64); + var h: Header = undefined; + try l.buf.download(0, std.mem.asBytes(&h)); + if (h.n > l.ids.len) return error.TooManyIds; + if (h.n > 0) { + try l.bcast(64 + h.n * 4); + try l.buf.download(64, std.mem.sliceAsBytes(l.ids[0..h.n])); + } + return .{ .h = h, .ids = l.ids[0..h.n] }; + } +}; + +/// Rank 0's engine: every step goes to the other ranks first. +pub const Leader = struct { + e: *Engine, + link: *Link, + + /// The step on the bound stream's sequence (its slot) to every rank. + fn send(x: Leader, h: Header, ids: []const u32) !void { + var hh = h; + hh.slot = x.e.slot; + try x.link.send(x.e.m.glue.put_words, hh, ids); + } + + /// The stream's sequence for the steps that follow (each step names it, so the other ranks bind it too). + pub fn bind(x: Leader, slot: u32) !void { + try x.e.bind(slot); + } + + /// Frees a stream's sequence on every rank. + pub fn release(x: Leader, slot: u32) !void { + try x.flush(); + try x.e.bind(slot); + try x.send(.{ .kind = @backingInt(Kind.release) }, &.{}); + x.e.release(slot); + } + + /// States rank 0 dropped since the last step, so the other ranks free theirs (ids that fit one step at a time). + fn flush(x: Leader) !void { + const d = &x.e.dropped; + while (d.items.len > 0) { + const n = @min(d.items.len, x.link.ids.len); + try x.send(.{ .kind = @backingInt(Kind.drop) }, d.items[0..n]); + d.replaceRangeAssumeCapacity(0, n, &.{}); + } + } + + pub fn begin(x: Leader, s: ?sampler.Sampling) !void { + try x.flush(); + const c = sampler.check(s); + var h: Header = .{ .kind = @backingInt(Kind.begin) }; + if (c) |v| { + h.flags = 1; + h.seed = v.seed; + h.temperature = v.temperature; + h.top_p = v.top_p; + h.min_p = v.min_p; + h.top_k = v.top_k; + } + try x.send(h, &.{}); + try x.e.begin(c); + } + + /// `cancel` (rank 0's), asked before each chunk, decides for every rank; `from` every rank holds; `marks` rank 0's. + pub fn prefill(x: Leader, ids: []const u32, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !u32 { + try x.flush(); + try x.send(.{ .kind = @backingInt(Kind.prefill), .a = if (from) |r| r.id else 0, .seed = if (from) |r| r.at else 0 }, ids); + return x.e.prefill(ids, from, cancel, marks); + } + + /// One chunk of a prompt pass on every rank (Engine.prefillChunk); the first token once the chunk ends the prompt. + pub fn chunk(x: Leader, ids: []const u32, start: usize, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !?u32 { + if (start >= ids.len) return error.BadChunk; + try x.flush(); + const part = ids[start..@min(ids.len, start + engine.prompt_rows + 1)]; + try x.send(.{ .kind = @backingInt(Kind.chunk), .a = if (from) |r| r.id else 0, .seed = if (from) |r| r.at else 0, .depth = @intCast(start), .top_k = @intCast(ids.len) }, part); + return x.e.prefillChunk(part, start, ids.len, from, cancel, marks); + } + + pub fn verify(x: Leader, ids: []const u32, out: []u32) !void { + try x.flush(); + try x.send(.{ .kind = @backingInt(Kind.verify) }, ids); + try x.e.verify(ids, out); + } + + /// Several streams' windows in one forward on every rank (Engine.verifyShared): parts' slots and rows, then ids. + pub fn verifyShared(x: Leader, parts: []const engine.Part) !void { + if (parts.len == 0 or parts.len > engine.max_shared) return error.WindowTooWide; + try x.flush(); + var buf: [2 * engine.max_shared + engine.shared_rows]u32 = undefined; + var n: usize = 2 * parts.len; + for (parts, 0..) |p, i| { + if (n + p.ids.len > buf.len) return error.WindowTooWide; + buf[2 * i] = p.slot; + buf[2 * i + 1] = @intCast(p.ids.len); + @memcpy(buf[n..][0..p.ids.len], p.ids); + n += p.ids.len; + } + try x.send(.{ .kind = @backingInt(Kind.shared), .depth = @intCast(parts.len) }, buf[0..n]); + try x.e.verifyShared(parts); + } + + pub fn keep(x: Leader, n: usize) !void { + try x.flush(); + try x.send(.{ .kind = @backingInt(Kind.keep), .a = n }, &.{}); + try x.e.keep(n); + } + + pub fn draft(x: Leader, from_prompt: bool, follow_ids: []const u32, depth: usize, out: []u32) !void { + try x.flush(); + try x.send(.{ .kind = @backingInt(Kind.draft), .depth = @intCast(depth), .flags = @intFromBool(from_prompt) }, follow_ids); + try x.e.draft(from_prompt, follow_ids, depth, out); + } + + /// The other ranks leave their loops. + pub fn end(x: Leader) !void { + try x.send(.{ .kind = @backingInt(Kind.end) }, &.{}); + } +}; + +/// A rank other than 0 runs rank 0's steps until it ends; a cancelled prefill is error.Cancelled and the loop goes on. +pub fn follow(e: *Engine, link: *Link) !void { + var drafts: [engine.max_drafts]u32 = undefined; + var sampled: [16]u32 = undefined; + while (true) { + const step = try link.receive(); + const h = step.h; + const kind: Kind = @fromBackingInt(@intCast(h.kind)); + if (kind != .end and kind != .drop and kind != .shared) try e.bind(h.slot); + switch (kind) { + .begin => try e.begin(if (h.flags & 1 != 0) .{ .seed = h.seed, .temperature = h.temperature, .top_p = h.top_p, .min_p = h.min_p, .top_k = h.top_k } else null), + .prefill => { + const ids = try e.gpa.dupe(u32, step.ids); + defer e.gpa.free(ids); + const from: ?engine.Resume = if (h.a != 0) .{ .id = h.a, .at = h.seed } else null; + _ = e.prefill(ids, from, null, null) catch |err| switch (err) { + error.Cancelled => {}, + else => return err, + }; + }, + .chunk => { + const ids = try e.gpa.dupe(u32, step.ids); + defer e.gpa.free(ids); + const from: ?engine.Resume = if (h.a != 0) .{ .id = h.a, .at = h.seed } else null; + _ = e.prefillChunk(ids, h.depth, h.top_k, from, null, null) catch |err| switch (err) { + error.Cancelled => {}, + else => return err, + }; + }, + .verify => try e.verify(step.ids, sampled[0..step.ids.len]), + .shared => { + const k = h.depth; + if (k == 0 or k > engine.max_shared or step.ids.len < 2 * k) return error.BadStep; + var parts: [engine.max_shared]engine.Part = undefined; + var outs: [engine.shared_rows]u32 = undefined; + var at: usize = 2 * k; + var row: usize = 0; + for (parts[0..k], 0..) |*p, i| { + const n = step.ids[2 * i + 1]; + if (at + n > step.ids.len or row + n > outs.len) return error.BadStep; + p.* = .{ .slot = step.ids[2 * i], .ids = step.ids[at..][0..n], .out = outs[row..][0..n] }; + at += n; + row += n; + } + try e.verifyShared(parts[0..k]); + }, + .keep => try e.keep(@intCast(h.a)), + .draft => { + const ids = try e.gpa.dupe(u32, step.ids); + defer e.gpa.free(ids); + try e.draft(h.flags & 1 != 0, ids, h.depth, &drafts); + }, + .drop => for (step.ids) |id| e.dropState(id), + .release => e.release(h.slot), + .end => return, + } + } +} + +test "a step's header is 64 bytes and keeps its fields" { + const h: Header = .{ .kind = @backingInt(Kind.draft), .n = 3, .depth = 2, .flags = 1, .a = 7, .seed = 1 << 62, .temperature = 0.6, .top_p = 0.95, .top_k = 20 }; + const bytes = std.mem.asBytes(&h); + var back: Header = undefined; + @memcpy(std.mem.asBytes(&back), bytes); + try std.testing.expectEqual(h, back); + try std.testing.expectEqual(Kind.draft, @as(Kind, @fromBackingInt(@intCast(back.kind)))); +} diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig new file mode 100644 index 0000000000..1e57e55a08 --- /dev/null +++ b/zig/src/families/glm/cuda_forward.zig @@ -0,0 +1,569 @@ +//! GLM-5.3-Flash's TP forward on CUDA (the Python family's forward.compute and commit), partials over NCCL or copied. + +const std = @import("std"); +const cuda = @import("cuda"); +const cfg = @import("config.zig"); +const weights = @import("cuda_weights.zig"); +const glue = @import("cuda_glue.zig"); +const kda = @import("cuda_kda.zig"); +const b16 = @import("cuda_b16.zig"); +const latent = @import("cuda_latent.zig"); +const sparse = @import("cuda_sparse.zig"); +const experts = @import("cuda_experts.zig"); +const kda_block = @import("cuda_kda_block.zig"); +const dsa_block = @import("cuda_dsa_block.zig"); +const mlp = @import("cuda_mlp.zig"); +const moe = @import("cuda_moe.zig"); +const quant = @import("cuda_quant.zig"); +const q4 = @import("cuda_q4.zig"); +const state = @import("cuda_state.zig"); +const shared = @import("cuda_shared.zig"); + +pub const prefill_rows = state.prefill_rows; +pub const window_rows = state.window_rows; +pub const State = state.State; +pub const Buffers = state.Buffers; + +/// The modules a forward launches from: our fatbins. +pub const Modules = struct { + glue: glue.Kernels, + mods: [7]cuda.Module, + kda: kda.Kernels, + lane: cuda.qmmf.Lane, // the NVFP4 linears' decode rows (qlinear) + prompt: cuda.nvfp4.Prompt, // and their prompt rows + experts: experts.Kernels, + copy_rows: cuda.Function, + put_slot: cuda.Function, + pack_dense: cuda.Function, + quant: quant.Kernels, + q4: q4.Kernels, + + pub fn load(d: *const cuda.Driver, ctx: *const cuda.Context) !Modules { + const k = cuda.kernels; + var m: Modules = undefined; + var n: usize = 0; + errdefer for (m.mods[0..n]) |*x| x.unload(); + for (&m.mods, [_][]const u8{ k.glm_kda, k.glm_ops, k.experts, k.affine4_pack, k.glm_quant, k.qmm_group, k.glm_glue }) |*x, image| { + x.* = try cuda.Module.load(d, image); + n += 1; + } + m.lane = try cuda.qmmf.Lane.load(d, @intCast(try ctx.attribute(.compute_capability_major))); + errdefer m.lane.unload(); + m.prompt = try cuda.nvfp4.Prompt.load(d); + errdefer m.prompt.unload(); + m.kda = try kda.Kernels.load(m.mods[0]); + m.experts = try experts.Kernels.load(d, m.mods[2], @intCast(try ctx.attribute(.multiprocessor_count))); + errdefer m.experts.unload(); + m.copy_rows = try m.mods[1].function("tf_glm_copy_rows"); + m.put_slot = try m.mods[1].function("tf_glm_put_slot"); + m.pack_dense = try m.mods[3].function("tf_pack_dense"); + m.quant = try quant.Kernels.load(m.mods[4]); + m.q4 = try q4.Kernels.load(m.mods[5], m.mods[3], ctx); + m.glue = try glue.Kernels.load(m.mods[6]); + return m; + } + + /// What weights.load runs: the dense packing, and for the MTP head the quantizers and the 4-bit tiling. + pub fn tools(m: *const Modules) weights.Tools { + return .{ .pack_dense = m.pack_dense, .quant = &m.quant, .q4 = &m.q4 }; + } + + pub fn ops(m: *const Modules) cuda.Module { + return m.mods[1]; + } + + pub fn deinit(m: *Modules) void { + for (&m.mods) |*x| x.unload(); + m.lane.unload(); + m.prompt.unload(); + m.experts.unload(); + m.* = undefined; + } +}; + +/// overlap.ROW_STEP and PIECES_MAX: piece boundaries on the BF16 prompt matmul's 128-row blocks, at most 16 pieces. +pub const row_step = 128; +pub const max_pieces = 16; + +/// overlap.ranges: [lo, hi) row pieces of a chunk of `rows`, about equal on row_step boundaries (fewer when short). +pub fn ranges(rows: usize, pieces: usize, out: *[max_pieces][2]usize) usize { + const per = (rows + pieces - 1) / pieces; + const size = (per + row_step - 1) / row_step * row_step; + var n: usize = 0; + var lo: usize = 0; + while (lo < rows) : (lo += size) { + out[n] = .{ lo, @min(rows, lo + size) }; + n += 1; + } + return n; +} + +/// A prompt chunk's pieced exchanges (overlap.Overlap): second stream and events; `count`: TF_GLM_OVERLAP_PIECES (3). +pub const Pieces = struct { + count: usize = 4, + side: cuda.Stream, + filled: [max_pieces]cuda.Event, + gathered: [max_pieces]cuda.Event, + glues: [max_pieces]cuda.Event, // split: a piece's own rows glued + shared: [max_pieces]cuda.Event, // split: every rank's glued rows of a piece in + + pub fn init(d: *const cuda.Driver, count: usize) !Pieces { + if (count == 0 or count > max_pieces) return error.BadPieces; + var p: Pieces = .{ .count = count, .side = try cuda.Stream.init(d, false), .filled = undefined, .gathered = undefined, .glues = undefined, .shared = undefined }; + for (&p.filled, &p.gathered, &p.glues, &p.shared) |*a, *b, *g, *s| { + a.* = try cuda.Event.init(d, false); + b.* = try cuda.Event.init(d, false); + g.* = try cuda.Event.init(d, false); + s.* = try cuda.Event.init(d, false); + } + return p; + } + + pub fn deinit(p: *Pieces) void { + for (&p.filled, &p.gathered, &p.glues, &p.shared) |*a, *b, *g, *s| { + a.deinit(); + b.deinit(); + g.deinit(); + s.deinit(); + } + p.side.deinit(); + p.* = undefined; + } +}; + +/// How the ranks' partials meet: every rank in-process copying into the others' slots, or one rank over NCCL. +pub const Exchange = union(enum) { + local, + nccl: struct { lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm }, +}; + +/// One rank of the forward: its weights, its sequence's caches, the window's buffers. +pub const Rank = struct { + w: *const weights.Weights, + st: *State, + b: *Buffers, + /// A shared window (Buffers.initShared): each stream's rows in row order; empty: every row on `st`. + segs: []const Seg = &.{}, +}; + +/// One stream's rows of a shared window: rows row0..row0 + rows, on sequence `st` from its position. +pub const Seg = struct { st: *State, row0: usize, rows: usize }; + +pub const Forward = struct { + d: *const cuda.Driver, + m: *const Modules, + stream: cuda.Stream, + ranks: []const Rank, // every rank of a local check, or this process's one + exchange: Exchange, + kda_wide: bool = true, // TF_GLM_KDA_DECODE_WIDE (the Python family's default): decode windows on KDA's three-kernel chain + pieces: ?*const Pieces = null, // a prompt chunk's exchanges in row pieces on a second stream (null: whole) + split: bool = true, // TF_GLM_PREFILL_REDUCE=split (the default with send/recv): each rank glues its own rows + + pub fn set(f: *const Forward) glue.Set { + return .{ .glue = &f.m.glue }; + } + + /// reduce.sums' gather: every rank's partial of `rows` rows into every rank's gathered buffer, rank order. + fn gather(f: *const Forward, rows: usize) !void { + const c = &f.ranks[0].w.config; + const bytes = rows * c.hidden * 4; + switch (f.exchange) { + .local => for (f.ranks) |to| for (f.ranks) |from| { + if (to.w.share.rank == from.w.share.rank) continue; + const off = from.w.share.rank * bytes; + try f.d.check(f.d.api.cuMemcpyDtoDAsync_v2(to.b.gathered + off, from.b.gathered + off, bytes, f.stream.handle), "cuMemcpyDtoDAsync"); + }, + .nccl => |n| { + const r = f.ranks[0]; + const count = rows * c.hidden; + try n.lib.check(n.lib.api.ncclAllGather(r.b.gathered + r.w.share.rank * bytes, r.b.gathered, count, .f32, n.comm, f.stream.handle), "ncclAllGather"); + }, + } + } + + fn slot(r: Rank, rows: usize) u64 { + return r.b.gathered + r.w.share.rank * rows * r.w.config.hidden * 4; + } + + /// forward._pre: hc_pre of every row at a site into b.normed, xs, post and comb. + fn pre(f: *const Forward, r: Rank, hc: weights.Hc, norm: u64, rows: usize) !void { + const c = &r.w.config; + const b = r.b; + try glue.hcPre(f.set(), f.stream, b.x, hc.fn_w, hc.base, hc.scale, norm, b.normed, b.xs, b.post, b.comb, b.hcpart, c.eps, c.hc_eps, c.sinkhorn, rows, c.hidden, c.hc); + } + + fn post(f: *const Forward, r: Rank, rows: usize) !void { + const c = &r.w.config; + try glue.hcPost(f.set(), f.stream, r.b.x, r.b.x, r.b.gathered, r.b.post, r.b.comb, rows, c.hidden, c.hc, r.w.share.world); + } + + /// forward._mixer: the layer's KDA or DSA block on b.normed; `li`/`di`: its index among loaded KDA / DSA layers. + fn mixer(f: *const Forward, r: Rank, l: weights.Layer, li: usize, di: usize, rows: usize, nch: usize, out: ?u64) !void { + const c = &r.w.config; + const b = r.b; + const st = r.st; + switch (l.mixer) { + .kda => |k| { + const cur = st.cur[li]; + var s = b.kda; + s.p = if (b.prefill) b.kproj else st.proj + li * st.proj_layer * 2; + s.saves.state_out = st.recAt(1 - cur, li); + if (!b.prefill) { // a window's saved rows stay with the layer until its commit + const kv = li * window_rows * st.heads * kda.dk; + s.saves.k = st.save_k + kv * 4; + s.saves.v = st.save_v + kv * 2; + s.saves.g = st.save_g + kv * 4; + s.saves.b = st.save_b + li * window_rows * st.heads * 4; + } + const conv = st.conv + li * 3 * st.conv_c * 2; + try kda_block.body(f.set(), &f.m.kda, f.stream, k, s, b.normed, c.hidden, rows, conv, st.recAt(cur, li), !b.prefill and f.kda_wide, b.prefill); + if (out) |o| try kda_block.fill(f.set(), f.stream, k, s, 0, rows, o); + if (b.prefill) st.cur[li] = 1 - cur; // a prompt chunk keeps every row: the layer commits now + }, + .mla => |a| { + try dsa_block.body(.{ .set = f.set(), .ops = f.m.ops() }, f.stream, a, st.caches[di], b.dsa, b.normed, c.hidden, rows, st.pos_dev, st.pos, nch, b.part_rows); + if (out) |o| try dsa_block.fill(f.set(), f.stream, a, b.dsa, 0, rows, o); + }, + } + } + + /// forward._ffn: the dense MLP or MoE block on b.normed into this rank's slot. + fn ffn(f: *const Forward, r: Rank, l: weights.Layer, rows: usize, out: ?u64) !void { + const c = &r.w.config; + const b = r.b; + switch (l.ffn) { + .dense => |m| { + try mlp.body(f.mlpKernels(), f.stream, m, b.dense, b.normed, c.hidden, rows, b.prefill); + if (out) |o| try mlp.fill(f.mlpKernels(), f.stream, m, b.dense, 0, rows, b.prefill, o); + }, + .moe => |m| { + try moe.body(f.moeKernels(), f.stream, m, b.moe, b.normed, c.hidden, rows, b.prefill); + if (out) |o| try moe.fill(f.moeKernels(), f.stream, m, b.moe, b.normed, c.hidden, 0, rows, b.prefill, o); + }, + } + } + + fn mlpKernels(f: *const Forward) mlp.Kernels { + return .{ .set = f.set(), .linear = .{ .lane = &f.m.lane, .nvfp4_prompt = &f.m.prompt }, .copy_rows = f.m.copy_rows }; + } + + fn moeKernels(f: *const Forward) moe.Kernels { + return .{ .set = f.set(), .experts = &f.m.experts, .put_slot = f.m.put_slot }; + } + + const Site = enum { mixer, ffn }; + + /// forward.out_proj's fill of rows lo..hi at a site of layer `l` into `out` (this rank's rows of a piece). + fn fillRows(f: *const Forward, r: Rank, l: weights.Layer, site: Site, lo: usize, hi: usize, out: u64) !void { + const c = &r.w.config; + const b = r.b; + switch (site) { + .mixer => switch (l.mixer) { + .kda => |k| try kda_block.fill(f.set(), f.stream, k, b.kda, lo, hi, out), + .mla => |a| try dsa_block.fill(f.set(), f.stream, a, b.dsa, lo, hi, out), + }, + .ffn => switch (l.ffn) { + .dense => |m| try mlp.fill(f.mlpKernels(), f.stream, m, b.dense, lo, hi, b.prefill, out), + .moe => |m| try moe.fill(f.moeKernels(), f.stream, m, b.moe, b.normed, c.hidden, lo, hi, b.prefill, out), + }, + } + } + + /// Piece [lo, hi)'s room in a rank's gathered buffer (pieces never share bytes): [world][hi - lo][dims] fp32. + fn room(r: Rank, lo: usize) u64 { + return r.b.gathered + r.w.share.world * lo * r.w.config.hidden * 4; + } + + /// reduce.shares: [lo, hi) of the rows of a piece of `n` each rank owns: about equal, the first n % world one more. + fn share(n: usize, world: usize, rank: usize) [2]usize { + const base = n / world; + const extra = n % world; + const lo = rank * base + @min(rank, extra); + return .{ lo, lo + base + @intFromBool(rank < extra) }; + } + + /// Whether piece [lo, lo + n) splits (TF_GLM_PREFILL_REDUCE=split): at least one row for every rank. + fn splits(f: *const Forward, n: usize) bool { + return f.split and n >= f.ranks[0].w.share.world; + } + + fn copy(f: *const Forward, stream: cuda.Stream, dst: u64, src: u64, bytes: usize) !void { + try f.d.check(f.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, stream.handle), "cuMemcpyDtoDAsync"); + } + + /// overlap.Overlap.partials: each piece filled on the main stream, exchanged on the second (gathered or split). + fn exchangePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut: []const [2]usize) !void { + const D = f.ranks[0].w.config.hidden; + const world = f.ranks[0].w.share.world; + for (cut, 0..) |piece, q| { + const lo = piece[0]; + const n = piece[1] - lo; + const split = f.splits(n); + for (f.ranks) |r| try f.fillRows(r, r.w.layers[j], site, lo, piece[1], if (split) r.b.part + lo * D * 4 else room(r, lo) + r.w.share.rank * n * D * 4); + try p.filled[q].record(f.stream); + try p.side.wait(p.filled[q]); + const bytes = n * D * 4; + if (!split) switch (f.exchange) { + .local => for (f.ranks) |to| for (f.ranks) |from| { + if (to.w.share.rank == from.w.share.rank) continue; + const off = from.w.share.rank * bytes; + try f.copy(p.side, room(to, lo) + off, room(from, lo) + off, bytes); + }, + .nccl => |x| { + const r = f.ranks[0]; + try x.lib.check(x.lib.api.ncclAllGather(room(r, lo) + r.w.share.rank * bytes, room(r, lo), n * D, .f32, x.comm, p.side.handle), "ncclAllGather"); + }, + } else switch (f.exchange) { + .local => for (f.ranks) |to| { + const mine = share(n, world, to.w.share.rank); + const own = mine[1] - mine[0]; + for (f.ranks) |from| try f.copy(p.side, room(to, lo) + from.w.share.rank * own * D * 4, from.b.part + (lo + mine[0]) * D * 4, own * D * 4); + }, + .nccl => |x| { + const r = f.ranks[0]; + const me = r.w.share.rank; + const mine = share(n, world, me); + const own = mine[1] - mine[0]; + try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart"); + for (0..world) |peer| if (peer != me) { + const theirs = share(n, world, peer); + try x.lib.check(x.lib.api.ncclSend(r.b.part + (lo + theirs[0]) * D * 4, (theirs[1] - theirs[0]) * D, .f32, @intCast(peer), x.comm, p.side.handle), "ncclSend"); + try x.lib.check(x.lib.api.ncclRecv(room(r, lo) + peer * own * D * 4, own * D, .f32, @intCast(peer), x.comm, p.side.handle), "ncclRecv"); + }; + try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd"); + try f.copy(p.side, room(r, lo) + me * own * D * 4, r.b.part + (lo + mine[0]) * D * 4, own * D * 4); + }, + } + try p.gathered[q].record(p.side); + } + } + + /// hc_post of rows lo..hi from partials `g` [world][hi - lo][D] (a gathered room or a split one's own rows). + fn postRows(f: *const Forward, r: Rank, lo: usize, hi: usize, g: u64) !void { + const c = &r.w.config; + const b = r.b; + const S: usize = c.hc; + const D: usize = c.hidden; + const x = b.x + lo * S * D * 2; + try glue.hcPost(f.set(), f.stream, x, x, g, b.post + lo * S * 4, b.comb + lo * S * S * 4, hi - lo, D, S, r.w.share.world); + } + + /// forward._pre of rows lo..hi. + fn preRows(f: *const Forward, r: Rank, hc: weights.Hc, norm: u64, lo: usize, hi: usize) !void { + const c = &r.w.config; + const b = r.b; + const S: usize = c.hc; + const D: usize = c.hidden; + try glue.hcPre(f.set(), f.stream, b.x + lo * S * D * 2, hc.fn_w, hc.base, hc.scale, norm, b.normed + lo * D * 2, b.xs + lo * (D / 64) * 4, b.post + lo * S * 4, b.comb + lo * S * S * 4, b.hcpart, c.eps, c.hc_eps, c.sinkhorn, hi - lo, D, S); + } + + /// One piece's glue after a site: hc_post, then the next layer's hc_pre or (after the last layer) the stream mean. + fn glueRows(f: *const Forward, r: Rank, j: usize, site: Site, lo: usize, hi: usize, g: u64) !void { + const c = &r.w.config; + const layers = r.w.layers; + try f.postRows(r, lo, hi, g); + if (site == .mixer) return f.preRows(r, layers[j].ffn_hc, layers[j].post_norm, lo, hi); + if (j + 1 < layers.len) return f.preRows(r, layers[j + 1].attn_hc, layers[j + 1].in_norm, lo, hi); + try glue.streamMean(f.set(), f.stream, r.b.x + lo * c.hc * c.hidden * 2, r.b.hidden + lo * c.hidden * 2, hi - lo, c.hidden, c.hc); + } + + /// overlap.Overlap.glue: each piece's glue once its exchange is in, its rows then shared on the second stream. + fn gluePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut: []const [2]usize) !void { + const c = &f.ranks[0].w.config; + const D: usize = c.hidden; + const world = f.ranks[0].w.share.world; + const last = site == .ffn and j + 1 == f.ranks[0].w.layers.len; + var sent: usize = 0; + for (cut, 0..) |piece, q| { + const lo = piece[0]; + const n = piece[1] - lo; + try f.stream.wait(p.gathered[q]); + if (!f.splits(n)) { + for (f.ranks) |r| try f.glueRows(r, j, site, lo, piece[1], room(r, lo)); + continue; + } + for (f.ranks) |r| { + const mine = share(n, world, r.w.share.rank); + try f.glueRows(r, j, site, lo + mine[0], lo + mine[1], room(r, lo)); + } + try p.glues[q].record(f.stream); + try p.side.wait(p.glues[q]); + switch (f.exchange) { + .local => for (f.ranks) |from| { + const mine = share(n, world, from.w.share.rank); + const bytes = (mine[1] - mine[0]) * D * 2; + for (f.ranks) |to| if (to.w.share.rank != from.w.share.rank) { + const at = (lo + mine[0]) * D * 2; + try f.copy(p.side, (if (last) to.b.hidden else to.b.normed) + at, (if (last) from.b.hidden else from.b.normed) + at, bytes); + }; + }, + .nccl => |x| { + const r = f.ranks[0]; + const me = r.w.share.rank; + const rows_of = if (last) r.b.hidden else r.b.normed; + const mine = share(n, world, me); + try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart"); + for (0..world) |peer| if (peer != me) { + const theirs = share(n, world, peer); + try x.lib.check(x.lib.api.ncclSend(rows_of + (lo + mine[0]) * D * 2, (mine[1] - mine[0]) * D, .bf16, @intCast(peer), x.comm, p.side.handle), "ncclSend"); + try x.lib.check(x.lib.api.ncclRecv(rows_of + (lo + theirs[0]) * D * 2, (theirs[1] - theirs[0]) * D, .bf16, @intCast(peer), x.comm, p.side.handle), "ncclRecv"); + }; + try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd"); + }, + } + try p.shared[q].record(p.side); + sent |= @as(usize, 1) << @intCast(q); + } + for (0..cut.len) |q| if (sent >> @intCast(q) & 1 == 1) try f.stream.wait(p.shared[q]); + } + + /// forward.pieced_layers: a prompt chunk's layers with exchanges in row pieces; same kernels and bits as unpieced. + fn pieced(f: *const Forward, p: *const Pieces, R: usize, cut: []const [2]usize) !void { + const layers = f.ranks[0].w.layers; + for (f.ranks) |r| try f.pre(r, r.w.layers[0].attn_hc, r.w.layers[0].in_norm, R); + var li: usize = 0; + var di: usize = 0; + for (layers, 0..) |_, j| { + for (f.ranks) |r| try f.mixer(r, r.w.layers[j], li, di, R, @min(latent.chunksFor(r.st.pos + R), r.b.nch), null); + try f.exchangePieces(p, j, .mixer, cut); + try f.gluePieces(p, j, .mixer, cut); + for (f.ranks) |r| try f.ffn(r, r.w.layers[j], R, null); + try f.exchangePieces(p, j, .ffn, cut); + try f.gluePieces(p, j, .ffn, cut); + switch (layers[j].mixer) { + .kda => li += 1, + .mla => di += 1, + } + } + } + + /// forward.forward of `ids`: logits in b.logits; a prompt chunk's KDA layers committed, a window's at `commit`. + pub fn run(f: *const Forward, ids: []const i32) !void { + const R = ids.len; + for (f.ranks) |r| { + if (R == 0 or R > r.b.rows) return error.BadWindow; + if (r.segs.len == 0) { + if (r.st.pos + R > r.st.capacity) return error.PastCapacity; + r.st.pending = R; + r.st.prompt = r.b.prefill; + } else try shared.begin(r, R); + try f.d.check(f.d.api.cuMemcpyHtoDAsync_v2(r.b.ids, ids.ptr, R * 4, f.stream.handle), "cuMemcpyHtoDAsync"); + } + const c = &f.ranks[0].w.config; + for (f.ranks) |r| try glue.embedB16(f.set(), f.stream, r.b.ids, r.w.embed, r.b.x, R, c.hidden, c.hc); + var cut: [max_pieces][2]usize = undefined; + const pieces = if (f.pieces) |p| (if (f.ranks[0].b.prefill) ranges(R, p.count, &cut) else 0) else 0; + // overlap.Overlap.begin: pieced with two pieces or more, and split takes a chunk of one piece too + if (pieces > 1 or (pieces == 1 and f.splits(R))) try f.pieced(f.pieces.?, R, cut[0..pieces]) else try f.whole(R); + try f.finish(R); + } + + /// The unpieced loop over the layers (and the stream mean): each site's partials whole, then the glue. + fn whole(f: *const Forward, R: usize) !void { + const c = &f.ranks[0].w.config; + var li: usize = 0; + var di: usize = 0; + for (f.ranks[0].w.layers, 0..) |_, j| { + for (f.ranks) |r| { + const l = r.w.layers[j]; + try f.pre(r, l.attn_hc, l.in_norm, R); + if (r.segs.len > 0) try shared.mixer(f, r, l, li, di, R, slot(r, R)) else try f.mixer(r, l, li, di, R, @min(latent.chunksFor(r.st.pos + R), r.b.nch), slot(r, R)); + } + try f.gather(R); + for (f.ranks) |r| { + const l = r.w.layers[j]; + try f.post(r, R); + try f.pre(r, l.ffn_hc, l.post_norm, R); + try f.ffn(r, l, R, slot(r, R)); + } + try f.gather(R); + for (f.ranks) |r| try f.post(r, R); + switch (f.ranks[0].w.layers[j].mixer) { + .kda => li += 1, + .mla => di += 1, + } + } + for (f.ranks) |r| try glue.streamMean(f.set(), f.stream, r.b.x, r.b.hidden, R, c.hidden, c.hc); + } + + /// The final norm and this rank's rows of the head (a prompt buffer's last row only). + fn finish(f: *const Forward, R: usize) !void { + const c = &f.ranks[0].w.config; + for (f.ranks) |r| { + const b = r.b; + try glue.rmsnorm(f.set(), f.stream, b.hidden, c.hidden, r.w.norm, b.fnormed, c.hidden, b.fxs, c.eps, R, c.hidden); + const head_rows: usize = if (b.prefill) 1 else R; + const first = if (b.prefill) R - 1 else 0; + try b16.matmul(f.set(), f.stream, b.fnormed + first * c.hidden * 2, c.hidden, r.w.head, b.logits, b.kda.mm_part, head_rows, r.w.share.vocab.len(), c.hidden, false); + } + } + + /// mtp.mtp_forward: caches from mtp_len, stream in b.mx, a window's draft logits in b.logits (prompt: no head). + pub fn mtp(f: *const Forward, ids: []const i32, hidden: []const u64) !void { + const n = ids.len; + const c = &f.ranks[0].w.config; + const D: usize = c.hidden; + for (f.ranks, hidden) |r, h| { + if (r.w.mtp == null) return error.NoMtpHead; + if (n == 0 or n > r.b.rows) return error.BadWindow; + if (r.st.mtp_len + n > r.st.capacity) return error.PastCapacity; + try f.d.check(f.d.api.cuMemcpyHtoDAsync_v2(r.b.ids, ids.ptr, n * 4, f.stream.handle), "cuMemcpyHtoDAsync"); + if (h != r.b.hin) try f.d.check(f.d.api.cuMemcpyDtoDAsync_v2(r.b.hin, h, n * D * 2, f.stream.handle), "cuMemcpyDtoDAsync"); + const m = r.w.mtp.?; + const b = r.b; + try glue.embedB16(f.set(), f.stream, b.ids, r.w.embed, b.me, n, D, 1); + if (r.st.mtp_len == 0) try f.d.check(f.d.api.cuMemsetD8Async(b.me, 0, D * 2, f.stream.handle), "cuMemsetD8Async"); // position 0's embedding is zeroed + try glue.rmsnorm(f.set(), f.stream, b.me, D, m.enorm, b.mcat, 2 * D, null, c.eps, n, D); + try glue.rmsnorm(f.set(), f.stream, b.hin, D, m.hnorm, b.mcat + D * 2, 2 * D, null, c.eps, n, D); + try b16.matmul(f.set(), f.stream, b.mcat, 2 * D, m.eh, b.mx, b.kda.mm_part, n, D, 2 * D, false); + try glue.rmsnorm(f.set(), f.stream, b.mx, D, m.layer.in_norm, b.normed, D, b.xs, c.eps, n, D); + const nch = @min(latent.chunksFor(r.st.mtp_len + n), b.nch); + try dsa_block.block(.{ .set = f.set(), .ops = f.m.ops() }, f.stream, m.layer.mixer.mla, r.st.mtp_caches, b.dsa, b.normed, D, n, r.st.mtp_pos_dev, r.st.mtp_len, nch, b.part_rows, slot(r, n)); + } + try f.gather(n); + for (f.ranks) |r| { + const m = r.w.mtp.?; + const b = r.b; + try glue.residualAdd(f.set(), f.stream, b.mx, b.mx, b.gathered, n, D, r.w.share.world); + try glue.rmsnorm(f.set(), f.stream, b.mx, D, m.layer.post_norm, b.normed, D, b.xs, c.eps, n, D); + try f.ffn(r, m.layer, n, slot(r, n)); + } + try f.gather(n); + for (f.ranks) |r| { + const m = r.w.mtp.?; + const b = r.b; + try glue.residualAdd(f.set(), f.stream, b.mx, b.mx, b.gathered, n, D, r.w.share.world); + try glue.rmsnorm(f.set(), f.stream, b.mx + (n - 1) * D * 2, D, m.norm, b.fnormed, D, b.fxs, c.eps, 1, D); + if (!b.prefill) try q4.matmul(&f.m.q4, f.stream, b.fnormed, D, b.fxs, m.draft, b.logits, 1); + } + } + + /// forward.commit: keep the first `keep` rows; window KDA layers take new states (replayed if fewer), conv shifted. + pub fn commit(f: *const Forward, keep: usize) !void { + for (f.ranks) |r| { + const st = r.st; + if (keep == 0 or keep > st.pending or (st.prompt and keep != st.pending)) return error.BadCommit; + if (!st.prompt and st.kda_layers > 0) { + const cur = st.cur[0]; + for (st.cur) |x| if (x != cur) return error.KdaLayersOutOfStep; + if (keep < st.pending) { + const h = st.heads; + var a: cuda.Args = .{}; + a.add(@as(c_int, @intCast(h))); + a.add(st.recAt(cur, 0)); + a.add(@as(usize, st.rec_layer / 4)); + for ([_]u64{ st.save_k, st.save_v, st.save_g, st.save_b }) |x| a.add(x); + a.add(@as(usize, window_rows * h * kda.dk)); + a.add(@as(usize, window_rows * h)); + a.add(@as(c_int, @intCast(keep))); + a.add(st.recAt(1 - cur, 0)); + try cuda.launch.launch(f.m.kda.replay_layers, .{ .grid = .{ .x = @intCast(st.kda_layers * h) }, .block = .{ .x = 1024 } }, f.stream, &a); + } + for (st.cur) |*x| x.* = 1 - x.*; + if (st.kda_layers > 0) try glue.convShift(f.set(), f.stream, st.conv, st.proj, keep, st.proj_layer, st.proj_row, st.kda_layers, st.conv_c); + } + st.pending = 0; + try st.setPos(f.stream, st.pos + keep); + } + } +}; diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig new file mode 100644 index 0000000000..733edcdb27 --- /dev/null +++ b/zig/src/families/glm/cuda_glue.zig @@ -0,0 +1,254 @@ +//! GLM-5.3-Flash's glue on CUDA: the Python family's cuda/glue.py on glm_glue.cu, whose kernels keep Triton's bytes. + +const std = @import("std"); +const cuda = @import("cuda"); + +/// What the forward's kernels launch from: our glue's functions. +pub const Set = struct { glue: *const Kernels }; + +/// Our glm_glue.cu's functions. +pub const Kernels = struct { + rmsnorm_4096: cuda.Function, + rmsnorm_1536: cuda.Function, + rmsnorm_512: cuda.Function, + residual_add: cuda.Function, + stream_mean: cuda.Function, + rank_sum: cuda.Function, + hc_post: cuda.Function, + combine: cuda.Function, + slice_sum: cuda.Function, + top1: cuda.Function, + put_words: cuda.Function, + embed_b16: cuda.Function, + conv_shift: cuda.Function, + swiglu_512: cuda.Function, + swiglu_64: cuda.Function, + group_sums: cuda.Function, + topk: cuda.Function, + router_part: cuda.Function, + router_tile: cuda.Function, + hc_partial: cuda.Function, + hc_partial8: cuda.Function, + hc_finish: cuda.Function, + absorb_q: cuda.Function, + expand_v: cuda.Function, + lat_write: cuda.Function, + index_write: cuda.Function, + pool_keys: cuda.Function, + select_rows: cuda.Function, + bmm_decode: cuda.Function, + bmm_rows: cuda.Function, + bmm_tile: cuda.Function, + bmm_tile1: cuda.Function, + bmm_d32s1: cuda.Function, + bmm_d32s2: cuda.Function, + bmm_d32s3: cuda.Function, + bmm_d64s2: cuda.Function, + chunks_16: cuda.Function, + chunks_32: cuda.Function, + merge: cuda.Function, + scores: cuda.Function, + + pub fn load(m: cuda.Module) !Kernels { + var k: Kernels = undefined; + inline for (@typeInfo(Kernels).@"struct".field_names) |name| @field(k, name) = try m.function("tf_glm_" ++ name); + try k.bmm_tile.allowDynamicShared(bmm_tile_smem); + try k.bmm_tile1.allowDynamicShared(bmm_tile1_smem[1]); + try k.chunks_16.allowDynamicShared(chunksSmem(16)); + try k.chunks_32.allowDynamicShared(chunksSmem(32)); + return k; + } +}; + +/// An element-wise kernel's launch: a block of 256 columns, `rows` x the blocks to cover `d`. +fn rowsLaunch(f: cuda.Function, stream: cuda.Stream, a: *cuda.Args, rows: usize, d: usize) !void { + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = cdiv(d, 256) }, .block = .{ .x = 256 } }, stream, a); +} + +/// As above, 8 columns a thread (16-byte loads), 1024 a block of 128 (Triton's BLOCK); d % 8 == 0, rows 16B-aligned. +fn rows8Launch(f: cuda.Function, stream: cuda.Stream, a: *cuda.Args, rows: usize, d: usize, ptrs: anytype) !void { + if (d % 8 != 0) return error.WidthNotEights; + inline for (ptrs) |p| if (p % 16 != 0) return error.RowsUnaligned; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = cdiv(d, 1024) }, .block = .{ .x = 128 } }, stream, a); +} + +fn i64of(v: usize) i64 { + return @intCast(v); +} + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn pow2(d: usize) i64 { + return @intCast(std.math.ceilPowerOfTwoAssert(usize, d)); +} + +/// glue.rmsnorm: bf16(w * bf16(x / sqrt(mean(x^2) + eps))) into `out`, with `xs` each 64 outputs' fp32 sum. +pub fn rmsnorm(set: Set, stream: cuda.Stream, x: u64, x_stride: usize, w: u64, out: u64, o_stride: usize, xs: ?u64, eps: f32, rows: usize, d: usize) !void { + const f, const warps: u32 = switch (d) { + 4096 => .{ set.glue.rmsnorm_4096, 8 }, + 1536 => .{ set.glue.rmsnorm_1536, 4 }, + 512 => .{ set.glue.rmsnorm_512, 4 }, + else => return error.NoRmsnormWidth, + }; + var a: cuda.Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(x_stride))); + a.add(w); + a.add(out); + a.add(@as(i64, @intCast(o_stride))); + a.add(xs orelse @as(u64, 0)); + a.add(eps); + a.add(@as(c_int, @intFromBool(xs != null))); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = warps * 32 } }, stream, &a); +} + +fn cdiv(a: usize, b: usize) u32 { + return @intCast((a + b - 1) / b); +} + +/// swiglu's column block: 512 where it divides the width, else 64 (TP=3's 704 and 640, the latter 128 in glue.py). +fn swigluBlock(w: usize) !usize { + if (w % 512 == 0) return 512; + if (w % 64 == 0) return 64; + return error.NoSwigluBlock; +} + +/// glue.swiglu: [gate | up] -> bf16(bf16(silu(min(g, L))) * clip(u, -L, L)), 64-sums in `xs`. +pub fn swiglu(set: Set, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { + const block = try swigluBlock(w); + const f, const threads: u32 = if (block == 512) .{ set.glue.swiglu_512, 128 } else .{ set.glue.swiglu_64, 64 }; + var a: cuda.Args = .{}; + inline for (.{ gu, out, xs, limit, i64of(w) }) |v| a.add(v); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(w / block) }, .block = .{ .x = threads } }, stream, &a); +} + +/// glm_glue_bmm.cuh's Tile::smem(f32) of tf_glm_bmm_tile: 3 stages of 128 x and 64 w rows, 64 bf16 a row. +pub const bmm_tile_smem: u32 = 3 * (128 + 64) * 64 * 2; +/// tf_glm_bmm_tile1's [bf16, fp32 out]: 3 stages of 128 x and 128 w rows, 32 bf16 a row, or the fp32 tile if larger. +pub const bmm_tile1_smem = [2]u32{ 3 * (128 + 128) * 32 * 2, 128 * (128 + 4) * 4 }; +/// glm_glue_attn.cuh's wide_smem(): q (hb x 512), two 32-key tiles, p (hb x 32) in bf16, then max/sum scratch. +pub fn chunksSmem(comptime hb: u32) u32 { + return (hb * 512 + 2 * 32 * 512 + hb * 32) * 2 + (hb * 4 + 2 * hb) * 4; +} + +/// K slices of the router's dot (glue.ROUTER_KS): fixed by the shape, as the sliced and the single pass differ in bits. +pub const router_ks = 8; + +/// glue.router: out [m][ne] fp32 = x . w rows; d a multiple of router_ks x 64: K slices in `part`, added in order. +pub fn router(set: Set, stream: cuda.Stream, x: u64, x_stride: usize, w: u64, out: u64, part: u64, m: usize, d: usize, ne: usize) !void { + if (d % (router_ks * 64) != 0) return error.RouterWidth; // the Python family's single-pass _router: no GLM width + var a: cuda.Args = .{}; + inline for (.{ x, w, part, i64of(m), i64of(x_stride), i64of(d), i64of(ne), @as(c_int, router_ks) }) |v| a.add(v); + if (x % 16 == 0 and x_stride % 8 == 0) // rows on 16 bytes: 64-row tiles staged by cp.async (decode's rows past m load as zeros) + try cuda.launch.launch(set.glue.router_tile, .{ .grid = .{ .x = cdiv(m, 64), .y = cdiv(ne, 32), .z = router_ks }, .block = .{ .x = 128 } }, stream, &a) + else + try cuda.launch.launch(set.glue.router_part, .{ .grid = .{ .x = cdiv(m, 16), .y = cdiv(ne, 32), .z = router_ks }, .block = .{ .x = 128 } }, stream, &a); + try sliceSum(set, stream, part, out, m * ne, router_ks, true); +} + +/// glue.select: top_k by sigmoid + bias (lower id on ties), normed if `norm`, times `scale`, then the shared expert. +pub fn select(set: Set, stream: cuda.Stream, logits: u64, bias: u64, pick: u64, wts: u64, rows: usize, top_k: usize, experts: usize, scale: f32, norm: bool) !void { + if (experts > 512 or top_k >= 32) return error.TopkTooWide; + var a: cuda.Args = .{}; + inline for (.{ logits, bias, pick, wts, scale, cint(experts), cint(top_k), @as(c_int, @intFromBool(norm)) }) |v| a.add(v); + try cuda.launch.launch(set.glue.topk, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 32 } }, stream, &a); +} + +/// glue.combine: out [rows][d] fp32 = sum over slots in order of wts[r][k] y[r][k] (y fp32 in decode, bf16 in prefill). +pub fn combine(set: Set, stream: cuda.Stream, y: u64, y_f32: bool, wts: u64, out: u64, rows: usize, d: usize, slots: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ y, @as(c_int, @intFromBool(y_f32)), wts, out, i64of(d), cint(slots) }) |v| a.add(v); + try rows8Launch(set.glue.combine, stream, &a, rows, d, .{ y, out }); +} + +/// K blocks of the hyper-connection mixing dots (glue.HC_BLOCKS): fixed by the shape, like the router's slices. +pub const hc_blocks = 16; + +/// The mixing kernels' scratch: rows x hc_blocks x 32 fp32 (24 dots and the sum of squares of each K block). +pub fn hcPartBytes(rows: usize) usize { + return rows * hc_blocks * 32 * 4; +} + +/// glue.hc_pre: streams [rows][s d] -> the collapsed RMS-normed row and its 64-sums `xs`; `post`, `comb` (Sinkhorn). +pub fn hcPre(set: Set, stream: cuda.Stream, x: u64, fn_w: u64, base: u64, scale: u64, norm_w: u64, out: u64, xs: u64, post: u64, comb: u64, part: u64, eps: f32, hc_eps: f32, iters: usize, rows: usize, d: usize, s: usize) !void { + if (d != 4096 or s != 4) return error.HcShape; + var a: cuda.Args = .{}; + inline for (.{ x, fn_w, part, i64of(s * d), @as(c_int, hc_blocks) }) |v| a.add(v); + // Decode windows (16 blocks a row) take each block's 8 steps at once; prompt rows 4 rows a 128-thread block. + if (rows < 64) { + try cuda.launch.launch(set.glue.hc_partial8, .{ .grid = .{ .x = @intCast(rows), .y = hc_blocks }, .block = .{ .x = 1024 } }, stream, &a); + } else { + a.add(i64of(rows)); + try cuda.launch.launch(set.glue.hc_partial, .{ .grid = .{ .x = cdiv(rows, 4), .y = hc_blocks }, .block = .{ .x = 128 } }, stream, &a); + } + var b: cuda.Args = .{}; + inline for (.{ x, part, base, scale, norm_w, out, xs, post, comb, eps, hc_eps, @as(c_int, hc_blocks), cint(iters) }) |v| b.add(v); + try cuda.launch.launch(set.glue.hc_finish, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, stream, &b); +} + +/// glue.hc_post: branch = bf16(`g`, rank 0 first); xout[s] = bf16(post[s] branch + sum_t comb[t][s] x[t]); x may alias. +pub fn hcPost(set: Set, stream: cuda.Stream, x: u64, xout: u64, g: u64, post: u64, comb: u64, rows: usize, d: usize, s: usize, world: usize) !void { + if (s != 4) return error.StreamsNotFour; + var a: cuda.Args = .{}; + inline for (.{ x, xout, g, post, comb, i64of(rows * d), i64of(d), cint(world) }) |v| a.add(v); + try rows8Launch(set.glue.hc_post, stream, &a, rows, d, .{ x, xout, g }); +} + +/// glue.rank_sum: out bf16 = partials summed rank 0 first, `me`'s from `own`, others' from `peers` in rank order. +pub fn rankSum(set: Set, stream: cuda.Stream, own: u64, peers: u64, peer_stride: usize, out: u64, me: usize, rows: usize, d: usize, world: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ own, peers, out, i64of(peer_stride), i64of(d), cint(world), cint(me) }) |v| a.add(v); + try rows8Launch(set.glue.rank_sum, stream, &a, rows, d, .{ own, peers, out }); +} + +/// glue.residual_add (the MTP block's plain residual): xout = bf16(x + bf16(the partials `g` summed rank 0 first)). +pub fn residualAdd(set: Set, stream: cuda.Stream, x: u64, xout: u64, g: u64, rows: usize, d: usize, world: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ x, xout, g, i64of(rows * d), i64of(d), cint(world) }) |v| a.add(v); + try rows8Launch(set.glue.residual_add, stream, &a, rows, d, .{ x, xout, g }); +} + +/// glue.stream_mean: out [rows][d] = bf16 of the s streams' mean (the final collapse). +pub fn streamMean(set: Set, stream: cuda.Stream, x: u64, out: u64, rows: usize, d: usize, s: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ x, out, i64of(d), cint(s) }) |v| a.add(v); + try rows8Launch(set.glue.stream_mean, stream, &a, rows, d, .{ x, out }); +} + +/// glue.embed of a BF16 table: each row's token (ids i32) as `copies` streams of d values, out [rows][copies d]. +pub fn embedB16(set: Set, stream: cuda.Stream, ids: u64, table: u64, out: u64, rows: usize, d: usize, copies: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ ids, table, out, i64of(d), cint(copies) }) |v| a.add(v); + try rowsLaunch(set.glue.embed_b16, stream, &a, rows, d); +} + +/// forward._shift_conv: each layer's 3-row conv window becomes rows keep .. keep + 2 of [window; projection rows]. +pub fn convShift(set: Set, stream: cuda.Stream, conv: u64, proj: u64, keep: usize, proj_layer: usize, proj_row: usize, layers: usize, c: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ conv, proj, i64of(keep), i64of(3 * c), i64of(proj_layer), i64of(proj_row), i64of(c) }) |v| a.add(v); + try rowsLaunch(set.glue.conv_shift, stream, &a, layers, c); +} + +/// glue._router_sum and qmm._reduce: `total` values summed over `ks` slices in slice order, fp32 or bf16 out. +pub fn sliceSum(set: Set, stream: cuda.Stream, part: u64, out: u64, total: usize, ks: usize, wide: bool) !void { + var a: cuda.Args = .{}; + inline for (.{ part, out, i64of(total), cint(ks), @as(c_int, @intFromBool(wide)) }) |v| a.add(v); + try cuda.launch.launch(set.glue.slice_sum, .{ .grid = .{ .x = cdiv(total, 1024) }, .block = .{ .x = 128 } }, stream, &a); +} + +test "the norm's block is the next power of two of the width" { + try std.testing.expectEqual(@as(i64, 4096), pow2(4096)); + try std.testing.expectEqual(@as(i64, 8192), pow2(6144)); +} + +test "the wrappers' shape rules as glue.py's" { + try std.testing.expectEqual(@as(usize, 512), try swigluBlock(6144)); // swiglu's block at the dense TP=2 width + try std.testing.expectEqual(@as(usize, 512), try swigluBlock(1024)); // a routed expert's rows at TP=2 + try std.testing.expectEqual(@as(usize, 64), try swigluBlock(704)); // TP=3's first two ranks + try std.testing.expectEqual(@as(usize, 64), try swigluBlock(640)); // TP=3's last rank + try std.testing.expectError(error.NoSwigluBlock, swigluBlock(96)); + try std.testing.expectEqual(@as(i64, 512), pow2(288 + 1)); // _topk's BLOCK for 288 experts + try std.testing.expectEqual(@as(i64, 16), pow2(8 + 1)); // its SLOTP for top-8 and the shared slot +} diff --git a/zig/src/families/glm/cuda_kda.zig b/zig/src/families/glm/cuda_kda.zig new file mode 100644 index 0000000000..cf2388e640 --- /dev/null +++ b/zig/src/families/glm/cuda_kda.zig @@ -0,0 +1,115 @@ +//! GLM-5.3-Flash's KDA on CUDA (the Python family's kda): fused and three-kernel chains match bits; replay rebuilds. + +const std = @import("std"); +const cuda = @import("cuda"); + +pub const dk = 128; +pub const dv = 128; +/// Windows of this many rows or more (prefill chunks) take the three-kernel chain (kda.WIDE_ROWS). +pub const wide_rows = 64; + +/// glm_kda.cu's entries. +pub const Kernels = struct { + chain: cuda.Function, + replay: cuda.Function, + replay_layers: cuda.Function, + prep: cuda.Function, + step: cuda.Function, + out: cuda.Function, + + pub fn load(m: cuda.Module) !Kernels { + var k: Kernels = undefined; + inline for (.{ "chain", "replay", "replay_layers", "prep", "step", "out" }) |name| @field(k, name) = try m.function("tf_glm_kda_" ++ name); + return k; + } +}; + +/// One layer's chain inputs, `heads` heads: p rows [q|k|v|f_a|g_a|...] (beta at `b_off`), decay and gate rows, states. +pub const Chain = struct { + p: u64, + p_stride: usize, + b_off: usize, + a: u64, + a_stride: usize, + g: u64, + g_stride: usize, + conv_state: u64, + conv_w: u64, + state_in: u64, + a_log: u64, + dt_bias: u64, + norm_w: u64, + eps: f32, + lower: f32, + heads: usize, +}; + +/// Where a chain writes: out bf16, state_out like state_in, and replay's rows: k fp32, v bf16, decay and beta fp32. +pub const Saves = struct { out: u64, state_out: u64, k: u64, v: u64, g: u64, b: u64 }; + +/// The three-kernel chain's scratch: q fp32 and the read-out bf16, [rows][heads][128] each. +pub const Wide = struct { q: u64, y: u64 }; + +/// kda.chain: `rows` rows through the fused kernel, or with `wide` through prep, step and out (WARPS 8, TR 16). +pub fn chain(ks: *const Kernels, stream: cuda.Stream, c: Chain, rows: usize, s: Saves, wide: ?Wide) !void { + const h: c_int = @intCast(c.heads); + if (wide) |w| { + var a: cuda.Args = .{}; + a.add(h); + a.add(c.p); + a.add(@as(c_int, @intCast(c.p_stride))); + a.add(@as(c_int, @intCast(c.b_off))); + a.add(c.a); + a.add(@as(c_int, @intCast(c.a_stride))); + for ([_]u64{ c.conv_state, c.conv_w, c.a_log, c.dt_bias }) |v| a.add(v); + a.add(c.lower); + for ([_]u64{ w.q, s.k, s.v, s.g, s.b }) |v| a.add(v); + const grid: cuda.launch.Dim3 = .{ .x = @intCast(rows), .y = @intCast(c.heads) }; + try cuda.launch.launch(ks.prep, .{ .grid = grid, .block = .{ .x = 512 } }, stream, &a); + var b: cuda.Args = .{}; + b.add(h); + for ([_]u64{ c.state_in, w.q, s.k, s.v, s.g, s.b }) |v| b.add(v); + b.add(@as(c_int, @intCast(rows))); + b.add(w.y); + b.add(s.state_out); + try cuda.launch.launch(ks.step, .{ .grid = .{ .x = @intCast(c.heads), .y = dv / 4 / 8 }, .block = .{ .x = 8 * 32 } }, stream, &b); + var o: cuda.Args = .{}; + o.add(h); + o.add(w.y); + o.add(c.g); + o.add(@as(c_int, @intCast(c.g_stride))); + o.add(c.norm_w); + o.add(c.eps); + o.add(s.out); + return cuda.launch.launch(ks.out, .{ .grid = grid, .block = .{ .x = dv } }, stream, &o); + } + var a: cuda.Args = .{}; + a.add(h); + a.add(c.p); + a.add(@as(c_int, @intCast(c.p_stride))); + a.add(@as(c_int, @intCast(c.b_off))); + a.add(c.a); + a.add(@as(c_int, @intCast(c.a_stride))); + a.add(c.g); + a.add(@as(c_int, @intCast(c.g_stride))); + for ([_]u64{ c.conv_state, c.conv_w, c.state_in, c.a_log, c.dt_bias, c.norm_w }) |v| a.add(v); + a.add(c.eps); + a.add(c.lower); + a.add(@as(c_int, @intCast(rows))); + for ([_]u64{ s.out, s.state_out, s.k, s.v, s.g, s.b }) |v| a.add(v); + try cuda.launch.launch(ks.chain, .{ .grid = .{ .x = @intCast(c.heads) }, .block = .{ .x = 1024 } }, stream, &a); +} + +/// kda.replay: the state after the first `rows` saved rows, from state_in. +pub fn replay(ks: *const Kernels, stream: cuda.Stream, state_in: u64, k: u64, v: u64, g: u64, b: u64, rows: usize, heads: usize, state_out: u64) !void { + var a: cuda.Args = .{}; + a.add(@as(c_int, @intCast(heads))); + for ([_]u64{ state_in, k, v, g, b }) |x| a.add(x); + a.add(@as(c_int, @intCast(rows))); + a.add(state_out); + try cuda.launch.launch(ks.replay, .{ .grid = .{ .x = @intCast(heads) }, .block = .{ .x = 1024 } }, stream, &a); +} + +test "the three-kernel chain's step grid is the Python family's" { + try std.testing.expectEqual(@as(u32, 4), dv / 4 / 8); // DV / 4 / WARPS value rows a block, 8 warps of 4 rows +} diff --git a/zig/src/families/glm/cuda_kda_block.zig b/zig/src/families/glm/cuda_kda_block.zig new file mode 100644 index 0000000000..918cf48a82 --- /dev/null +++ b/zig/src/families/glm/cuda_kda_block.zig @@ -0,0 +1,67 @@ +//! GLM-5.3-Flash's KDA block on CUDA, a rank's share of the Python family's kda_block: projection, chain, o_proj. + +const std = @import("std"); +const cuda = @import("cuda"); +const glue = @import("cuda_glue.zig"); +const kda = @import("cuda_kda.zig"); +const b16 = @import("cuda_b16.zig"); + +/// One KDA layer, `heads` heads: proj bf16 [3 heads 128 + 256 + heads][dims], f_b, g_b, o, the chain's taps and norm. +pub const Layer = struct { + proj: u64, + fb: u64, + gb: u64, + o: u64, + conv_w: u64, + a_log: u64, + dt_bias: u64, + norm_w: u64, + eps: f32, + lower: f32, + heads: usize, + dims: usize, + + pub fn width(l: Layer) usize { + return 3 * l.heads * kda.dk + 2 * kda.dk + l.heads; + } +}; + +/// The block's scratch for `rows` rows: projection, decay and gate rows, chain saves and scratch, B16 split-K partials. +pub const Scratch = struct { p: u64, ka: u64, kg: u64, saves: kda.Saves, wide: kda.Wide, mm_part: u64 }; + +/// forward.kda_block: out [rows][dims] fp32 partial; `wide`: three-kernel chain; `prompt`: commits the conv window. +pub fn block(set: glue.Set, ks: *const kda.Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, conv_state: u64, state_in: u64, wide: bool, prompt: bool, out: u64) !void { + try body(set, ks, stream, l, s, x, x_stride, rows, conv_state, state_in, wide, prompt); + try fill(set, stream, l, s, 0, rows, out); +} + +/// forward.out_proj's fill for the KDA block: o_proj of rows lo..hi (s.saves.out) into out [hi - lo][dims] fp32. +pub fn fill(set: glue.Set, stream: cuda.Stream, l: Layer, s: Scratch, lo: usize, hi: usize, out: u64) !void { + const n = l.heads * kda.dk; + try b16.matmul(set, stream, s.saves.out + lo * n * 2, n, l.o, out, s.mm_part, hi - lo, l.dims, n, true); +} + +/// The block up to its out_proj: the projections, the chain, and a prompt chunk's conv window commit. +pub fn body(set: glue.Set, ks: *const kda.Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, conv_state: u64, state_in: u64, wide: bool, prompt: bool) !void { + try project(set, stream, l, s, x, x_stride, rows); + try chainRows(ks, stream, l, s, rows, conv_state, state_in, wide); + if (prompt) try glue.convShift(set, stream, conv_state, s.p, rows, rows * l.width(), l.width(), 1, 3 * l.heads * kda.dk); +} + +/// The projections of `rows` rows into s.p, s.ka (decay) and s.kg (gate), row by row so segments share one call. +pub fn project(set: glue.Set, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize) !void { + const n = l.heads * kda.dk; + const w = l.width(); + const fa = 3 * n; + try b16.matmul(set, stream, x, x_stride, l.proj, s.p, s.mm_part, rows, w, l.dims, false); + try b16.matmul(set, stream, s.p + fa * 2, w, l.fb, s.ka, s.mm_part, rows, n, kda.dk, false); + try b16.matmul(set, stream, s.p + (fa + kda.dk) * 2, w, l.gb, s.kg, s.mm_part, rows, n, kda.dk, false); +} + +/// The chain over `rows` projected rows (s.p, s.ka, s.kg) on one sequence's conv window and recurrent state. +pub fn chainRows(ks: *const kda.Kernels, stream: cuda.Stream, l: Layer, s: Scratch, rows: usize, conv_state: u64, state_in: u64, wide: bool) !void { + const n = l.heads * kda.dk; + const fa = 3 * n; + const c: kda.Chain = .{ .p = s.p, .p_stride = l.width(), .b_off = fa + 2 * kda.dk, .a = s.ka, .a_stride = n, .g = s.kg, .g_stride = n, .conv_state = conv_state, .conv_w = l.conv_w, .state_in = state_in, .a_log = l.a_log, .dt_bias = l.dt_bias, .norm_w = l.norm_w, .eps = l.eps, .lower = l.lower, .heads = l.heads }; + try kda.chain(ks, stream, c, rows, s.saves, if (wide or rows >= kda.wide_rows) s.wide else null); +} diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig new file mode 100644 index 0000000000..0106c383b0 --- /dev/null +++ b/zig/src/families/glm/cuda_lanes.zig @@ -0,0 +1,295 @@ +//! GLM-5.3-Flash's CUDA lane backend (as nemotron/cuda_lanes.zig): a slot per stream, shared verify, chunked prompts. + +const std = @import("std"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const follow = @import("cuda_follow.zig"); +const fw = @import("cuda_forward.zig"); + +const be = lanes.backend; + +/// First tokens a handle names. +const ring = 1024; + +/// The widest verify window: the pending token and a full draft chain (glm-decode-check's, the AOT row variants). +pub const max_rows = 1 + engine.max_drafts; + +/// A stream's place in the backend: its sequence's slot, the rows its last verify left uncommitted, its held drafts. +const Lane = struct { + slot: u32, + pending_rows: ?usize = null, + held: [engine.max_drafts]u32 = undefined, + held_n: usize = 0, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts +}; + +pub const Lanes = struct { + x: follow.Leader, + whole: bool = false, // prompt passes run whole (TF_GLM_PROMPT_STEPS=0): no other stream's round between chunks + alone: bool = false, // each window in its own forward (TF_GLM_SHARED=0) + shared: u64 = 0, // forwards that held several streams' windows (checks read it) + lanes: std.AutoHashMapUnmanaged(*const lanes.Stream, Lane) = .empty, + drawn: [ring]u32 = undefined, + next: u64 = 0, + + pub fn deinit(self: *Lanes) void { + self.lanes.deinit(self.x.e.gpa); + } + + pub fn backend(self: *Lanes) be.Backend { + return .{ .ptr = self, .vtable = if (self.whole) &whole_vt else &steps_vt }; + } + + const whole_vt: be.Backend.VTable = .{ + .prefill = prefillFn, + .first = firstFn, + .queue = queueFn, + .read = readFn, + .verify = verifyFn, + .keep = keepFn, + .draft = draftFn, + .release = releaseFn, + }; + const steps_vt: be.Backend.VTable = vt: { + var v = whole_vt; + v.prefill_step = prefillStepFn; + break :vt v; + }; + + /// Round loop facts: windows up to max_rows, drafts with the head loaded, shared windows unless TF_GLM_SHARED=0. + pub fn facts(self: *const Lanes) lanes.Model { + const drafting = self.x.e.mtp; + const sharing = drafting and !self.alone; + return .{ + .exact_width = if (drafting) max_rows else 1, + .gpu_tokens = false, + .mtp = drafting, + .speculate = drafting, + .speculate_early = false, + .drafts = engine.max_drafts, + .hidden_rows = sharing, + .batch_rows = if (sharing) engine.shared_rows else 0, + .max_streams = if (sharing) engine.max_shared else 1, + }; + } + + fn of(ptr: *anyopaque) *Lanes { + return @ptrCast(@alignCast(ptr)); + } + + fn take(self: *Lanes, token: u32) u64 { + const h = self.next; + self.drawn[h % ring] = token; + self.next += 1; + return h; + } + + fn value(self: *const Lanes, feed: be.Feed) u32 { + return switch (feed) { + .handle => |h| self.drawn[h % ring], + .value => |v| v, + }; + } + + /// The lowest slot no stream holds. + fn freeSlot(self: *const Lanes) u32 { + var slot: u32 = 0; + while (true) : (slot += 1) { + var it = self.lanes.valueIterator(); + const taken = while (it.next()) |l| { + if (l.slot == slot) break true; + } else false; + if (!taken) return slot; + } + } + + /// The stream's sequence bound, its last verify committed (`kept` rows, else all of them). + fn bind(self: *Lanes, s: *const lanes.Stream, kept: ?usize) !*Lane { + const l = self.lanes.getPtr(s) orelse return error.UnknownStream; + try self.x.bind(l.slot); + if (l.pending_rows) |rows| try self.x.keep(kept orelse rows); + l.pending_rows = null; + return l; + } + + fn cancelled(ptr: ?*anyopaque) bool { + return @as(*lanes.Stream, @ptrCast(@alignCast(ptr.?))).isCancelled(); + } + + // -- the vtable --------------------------------------------------------------------------------------------- + + fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + const self = of(ptr); + const from = try self.opening(s); + _ = self.take(try self.x.prefill(s.prompt(), from, .{ .ptr = s, .check = cancelled }, marksOf(s))); + } + + /// One chunk of the stream's prompt pass (fillOne runs other rounds between); true once the first token is drawn. + fn prefillStepFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!bool { + const self = of(ptr); + const going: ?usize = if (self.lanes.getPtr(s)) |l| l.fill else null; + const from = if (going == null) try self.opening(s) else null; + const l = self.lanes.getPtr(s).?; + if (going != null) try self.x.bind(l.slot); + const at = going orelse if (from) |r| r.at else 0; + errdefer l.fill = null; + if (try self.x.chunk(s.prompt(), at, from, .{ .ptr = s, .check = cancelled }, marksOf(s))) |first| { + _ = self.take(first); + l.fill = null; + return true; + } + l.fill = at + engine.prompt_rows; + return false; + } + + /// A prompt pass begins: the stream's slot bound and its sequence fresh; where the pass resumes (a kept state). + fn opening(self: *Lanes, s: *lanes.Stream) !?engine.Resume { + const ids = s.prompt(); + if (ids.len + s.max_new + max_rows >= self.x.e.capacity) return error.PromptTooLong; + const gop = try self.lanes.getOrPut(self.x.e.gpa, s); + if (!gop.found_existing) gop.value_ptr.* = .{ .slot = self.freeSlot() }; + gop.value_ptr.pending_rows = null; + gop.value_ptr.held_n = 0; + gop.value_ptr.fill = null; + try self.x.bind(gop.value_ptr.slot); + try self.x.begin(s.sampling); + s.cached = 0; + return self.resumeFrom(s, ids.len); + } + + fn marksOf(s: *lanes.Stream) ?engine.Marks { + return if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; + } + + /// The stream's kept state, where its prompt pass starts (null: none, off the chunk grid, or no longer held). + fn resumeFrom(self: *Lanes, s: *lanes.Stream, len: usize) ?engine.Resume { + const saved = s.reuse.saved orelse return null; + const k: *const engine.Kept = @ptrCast(@alignCast(saved)); + if (k.at != s.reuse.at or k.at == 0 or k.at >= len or k.at % fw.prefill_rows != 0 or !self.x.e.holds(k.id)) { + s.reuse_failed = true; + return null; + } + s.cached = k.at; + return .{ .id = k.id, .at = k.at }; + } + + /// The prompt pass stands at a mark: the lane host keeps the state there (cuda_native's snap_save). + fn keptAt(ctx: *anyopaque, at: u32) void { + const s: *lanes.Stream = @ptrCast(@alignCast(ctx)); + if (s.reuse.hook) |k| k.at(k.ptr, s, at); + } + + fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + const self = of(ptr); + if (position != s.prompt_len) return error.PositionMismatch; + return self.next - 1; + } + + fn queueFn(ptr: *anyopaque, s: *lanes.Stream, feed: be.Feed, position: u64) anyerror!u64 { + _ = ptr; + _ = s; + _ = feed; + _ = position; + return error.NotPipelined; + } + + fn readFn(ptr: *anyopaque, handle: u64) anyerror!u32 { + const self = of(ptr); + if (handle >= self.next or self.next - handle > ring) return error.NoSuchToken; + return self.drawn[handle % ring]; + } + + /// Each window in its own forward on its stream's sequence (the core sends one when no forward is shared). + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const self = of(ptr); + if (windows.len > 1) return self.verifyShared(windows, out); + for (windows, out) |w, *o| { + if (w.parents != null) return error.TreesNotBuilt; + const l = try self.bind(w.stream, null); + const rows = w.rows(); + if (rows > max_rows) return error.WindowTooWide; + if (w.held > l.held_n) return error.NoHeldDrafts; + const pos = self.x.e.pos(); + for (w.positions, 0..) |p, r| if (p != pos + 1 + r) return error.PositionMismatch; + var ids: [max_rows]u32 = undefined; + ids[0] = w.pending; + const drafts = if (w.held > 0) l.held[0..w.held] else w.tokens; + @memcpy(ids[1..][0..drafts.len], drafts); + var sampled: [max_rows]u32 = undefined; + try self.x.verify(ids[0..rows], sampled[0..rows]); + @memcpy(o.sampled[0..rows], sampled[0..rows]); + @memcpy(o.drafts[0..drafts.len], drafts); + l.pending_rows = rows; + } + } + + /// A round's windows in one forward (Engine.verifyShared), each stream's rows on its own sequence. + fn verifyShared(self: *Lanes, windows: []const be.Window, out: []be.Verified) !void { + if (windows.len > engine.max_shared) return error.WindowTooWide; + var parts: [engine.max_shared]engine.Part = undefined; + var ids: [engine.max_shared][max_rows]u32 = undefined; + var sampled: [engine.max_shared][max_rows]u32 = undefined; + var lane: [engine.max_shared]*Lane = undefined; + for (windows, 0..) |w, k| { + if (w.parents != null) return error.TreesNotBuilt; + const l = try self.bind(w.stream, null); + const rows = w.rows(); + if (rows > max_rows) return error.WindowTooWide; + if (w.held > l.held_n) return error.NoHeldDrafts; + const pos = self.x.e.pos(); + for (w.positions, 0..) |p, r| if (p != pos + 1 + r) return error.PositionMismatch; + ids[k][0] = w.pending; + const drafts = if (w.held > 0) l.held[0..w.held] else w.tokens; + @memcpy(ids[k][1..][0..drafts.len], drafts); + parts[k] = .{ .slot = l.slot, .ids = ids[k][0..rows], .out = sampled[k][0..rows] }; + lane[k] = l; + } + try self.x.verifyShared(parts[0..windows.len]); + self.shared += 1; + for (out[0..windows.len], parts[0..windows.len], lane[0..windows.len]) |*o, p, l| { + const rows = p.ids.len; + @memcpy(o.sampled[0..rows], p.out); + @memcpy(o.drafts[0 .. rows - 1], p.ids[1..rows]); + l.pending_rows = rows; + } + } + + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const self = of(ptr); + for (windows, paths) |w, path| { + if (path.len == 0) return error.EmptyPath; + for (path, 0..) |r, i| if (r != i) return error.TreesNotBuilt; + _ = try self.bind(w.stream, path.len); + } + } + + fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + const self = of(ptr); + if (!self.x.e.mtp) return error.NoDraftHead; + for (requests) |r| { + const l = try self.bind(r.stream, if (r.rows) |rows| rows.len else null); + if (r.position != self.x.e.pos() + 1) return error.PositionMismatch; + if (r.lanes != null) return error.TreesNotBuilt; + const depth = @min(r.depth, engine.max_drafts); + if (r.rows) |rows| { + for (rows, 0..) |row, i| if (row != i) return error.TreesNotBuilt; + try self.x.draft(false, r.follow[0..rows.len], depth, &l.held); + } else { + const first = self.value(r.first orelse return error.NoFirstToken); + try self.x.draft(true, &.{first}, depth, &l.held); + } + l.held_n = depth; + } + } + + /// The stream left the rounds: its sequence is freed on every rank (slot 0 stays for the next stream). + fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { + const self = of(ptr); + const kv = self.lanes.fetchRemove(s) orelse return; + self.x.release(kv.value.slot) catch |err| std.log.err("releasing slot {d}: {s}", .{ kv.value.slot, @errorName(err) }); + } +}; + +comptime { + std.debug.assert(max_rows <= fw.window_rows); +} diff --git a/zig/src/families/glm/cuda_latent.zig b/zig/src/families/glm/cuda_latent.zig new file mode 100644 index 0000000000..a41381ef3b --- /dev/null +++ b/zig/src/families/glm/cuda_latent.zig @@ -0,0 +1,98 @@ +//! GLM-5.3-Flash's MLA on the 512-wide latent cache (DSA layers): the Python family's latent on BF16, row by row. + +const std = @import("std"); +const cuda = @import("cuda"); +const glue = @import("cuda_glue.zig"); + +/// latent.CHUNK, KT, HB, HB_WIDE, PROMPT_ROWS, PROMPT_RB. +pub const chunk = 512; +pub const key_tile = 32; +pub const heads_tile = 16; +pub const heads_tile_wide = 32; +pub const prompt_rows = 64; +pub const prompt_rb = 64; + +/// latent.head_block: a rank's 32 heads a program in prompt chunks (each key read once), 16 in decode windows. +pub fn headBlock(rows: usize) usize { + return if (rows >= 64) heads_tile_wide else heads_tile; +} + +/// latent.chunks_for: 512-key chunks to cover `length` keys. +pub fn chunksFor(length: usize) usize { + return (length + chunk - 1) / chunk; +} + +fn cdiv(a: usize, b: usize) u32 { + return @intCast((a + b - 1) / b); +} + +/// latent.absorb_q (BF16 kv_b): out [rows][heads][lw] bf16 = q [rows][heads][d] . wk [heads][d][lw], fp32 sums. +pub fn absorbQ(set: glue.Set, stream: cuda.Stream, q: u64, wk: u64, out: u64, rows: usize, heads: usize, d: usize, lw: usize) !void { + if (d != 256 or lw != 512) return error.LatentShape; + try rowBlocks(set.glue.absorb_q, stream, q, wk, out, rows, heads, lw / 32); +} + +/// The absorb and expand launches: rows in one block in decode windows, prompt_rb a block in prompt chunks. +fn rowBlocks(f: cuda.Function, stream: cuda.Stream, in: u64, w: u64, out: u64, rows: usize, heads: usize, cols: usize) !void { + const rb = if (rows >= prompt_rows) prompt_rb else rows; + var a: cuda.Args = .{}; + inline for (.{ in, w, out, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(i64, @intCast(rb)) }) |v| a.add(v); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(heads), .y = @intCast(cols), .z = cdiv(rows, rb) }, .block = .{ .x = 128 } }, stream, &a); +} + +/// latent.expand_v (BF16 kv_b): out [rows][heads][dv] bf16 = o_lat [rows][heads][lw] . wv [heads][dv][lw]. +pub fn expandV(set: glue.Set, stream: cuda.Stream, ol: u64, wv: u64, out: u64, rows: usize, heads: usize, dv: usize, lw: usize) !void { + if (dv != 256 or lw != 512) return error.LatentShape; + try rowBlocks(set.glue.expand_v, stream, ol, wv, out, rows, heads, dv / 16); +} + +/// latent.latent_write: rows of `lat` (stride `lat_stride`) into slots pos .. pos + rows - 1 (pos an i32 on device). +pub fn latentWrite(set: glue.Set, stream: cuda.Stream, lat: u64, lat_stride: usize, cache: u64, pos: u64, rows: usize, lw: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ lat, @as(i64, @intCast(lat_stride)), cache, pos, @as(i64, @intCast(lw)) }) |v| a.add(v); + try cuda.launch.launch(set.glue.lat_write, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 128 } }, stream, &a); +} + +/// The chunk partials of one attention call: po fp32 [nch][rows][heads][lw], pm and pl fp32 [nch][rows][heads]. +pub const Partials = struct { po: u64, pm: u64, pl: u64 }; + +/// latent.attention: causal qa over the cache in `nch` 512-key chunks into out bf16; `dummy`: the merge's count i32. +pub fn attention(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, p: Partials, out: u64, dummy: u64, rows: usize, heads: usize, lw: usize, nch: usize, scale: f32, hb: usize) !void { + try chunks(set, stream, qa, cache, pos, 0, 0, p, rows, heads, lw, nch, 0, scale, hb, false); + try merge(set, stream, p, out, dummy, rows, heads, lw, nch, false); +} + +/// latent._dense_chunks (sparse false: keys through pos + r) and _sparse_chunks (row r's `counts` tokens of `tokens`). +fn chunks(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, tokens: u64, counts: u64, p: Partials, rows: usize, heads: usize, lw: usize, nch: usize, w: usize, scale: f32, hb: usize, sparse: bool) !void { + if (lw != 512) return error.LatentShape; + const f = switch (hb) { + heads_tile => set.glue.chunks_16, + heads_tile_wide => set.glue.chunks_32, + else => return error.HeadBlock, + }; + var a: cuda.Args = .{}; + inline for (.{ qa, cache, pos, tokens, counts, p.po, p.pm, p.pl, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(i64, @intCast(w)), scale, @as(c_int, @intFromBool(sparse)) }) |v| a.add(v); + const shared = if (hb == heads_tile_wide) glue.chunksSmem(heads_tile_wide) else glue.chunksSmem(heads_tile); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = cdiv(heads, hb), .z = @intCast(nch) }, .block = .{ .x = 256 }, .shared = shared }, stream, &a); +} + +/// latent.sparse_attention: rows with counts > 0 over their selected tokens (i32 [rows][w]); others left alone. +pub fn sparseAttention(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, tokens: u64, counts: u64, w: usize, p: Partials, out: u64, rows: usize, heads: usize, lw: usize, scale: f32) !void { + const nch = chunksFor(w); + try chunks(set, stream, qa, cache, 0, tokens, counts, p, rows, heads, lw, nch, w, scale, headBlock(rows), true); + try merge(set, stream, p, out, counts, rows, heads, lw, nch, true); +} + +fn merge(set: glue.Set, stream: cuda.Stream, p: Partials, out: u64, counts: u64, rows: usize, heads: usize, lw: usize, nch: usize, sparse: bool) !void { + if (lw != 512) return error.LatentShape; + var a: cuda.Args = .{}; + inline for (.{ p.po, p.pm, p.pl, out, counts, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(c_int, @intCast(nch)), @as(c_int, @intFromBool(sparse)) }) |v| a.add(v); + try cuda.launch.launch(set.glue.merge, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(heads) }, .block = .{ .x = 128 } }, stream, &a); +} + +test "latent.py's tiling rules" { + try std.testing.expectEqual(@as(usize, 16), headBlock(8)); + try std.testing.expectEqual(@as(usize, 32), headBlock(64)); + try std.testing.expectEqual(@as(usize, 3), chunksFor(1025)); + try std.testing.expectEqual(@as(usize, 1), chunksFor(512)); +} diff --git a/zig/src/families/glm/cuda_mlp.zig b/zig/src/families/glm/cuda_mlp.zig new file mode 100644 index 0000000000..8e8d5eba11 --- /dev/null +++ b/zig/src/families/glm/cuda_mlp.zig @@ -0,0 +1,53 @@ +//! GLM-5.3-Flash's dense MLP (layers 0-2), a rank's share of the Python family's mlp_block: gate/up, SwiGLU, down. + +const std = @import("std"); +const cuda = @import("cuda"); +const glue = @import("cuda_glue.zig"); + +/// One dense layer on a rank: gate and up [width][dims] and down [dims][width], NVFP4 as Fp4Linear holds them. +pub const Layer = struct { gate: cuda.qmmf.Weight, up: cuda.qmmf.Weight, down: cuda.qmmf.Weight, width: usize, dims: usize, limit: f32 }; + +/// The block's scratch for up to `rows` rows: one projection's rows, the [gate | up] rows, the activation and its sums. +pub const Scratch = struct { tmp: u64, gu: u64, act: u64, xs_act: u64 }; + +pub const Kernels = struct { set: glue.Set, linear: cuda.qlinear.Linear, copy_rows: cuda.Function }; + +/// Fp4Linear.__call__ (prefill on a prompt chunk): bf16 out, or fp32 with `wide`; the prompt GEMM needs `ldx` % 8 == 0. +fn mm(k: Kernels, stream: cuda.Stream, w: cuda.qmmf.Weight, x: u64, ldx: usize, rows: usize, out: u64, wide: bool, prompt: bool) !void { + const q: cuda.qlinear.Weight = .{ .nvfp4 = w }; + const in: cuda.qlinear.Rows = .{ .x = x, .ldx = ldx }; + if (prompt) { + std.debug.assert(x % 16 == 0 and (rows == 1 or ldx % 8 == 0)); + return if (wide) k.linear.promptF32(stream, q, in, out, rows) else k.linear.prompt(stream, q, in, out, rows); + } + return if (wide) k.linear.decodeF32(stream, q, in, out, rows) else k.linear.decode(stream, q, in, out, rows); +} + +/// forward.mlp_block for `rows` normed rows (bf16, stride `x_stride`): out [rows][dims] fp32, this rank's partial. +pub fn block(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool, out: u64) !void { + try body(k, stream, l, s, x, x_stride, rows, prompt); + try fill(k, stream, l, s, 0, rows, prompt, out); +} + +/// forward.out_proj's fill for the dense MLP: down of rows lo..hi (s.act) into out [hi - lo][dims] fp32. +pub fn fill(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, lo: usize, hi: usize, prompt: bool, out: u64) !void { + try mm(k, stream, l.down, s.act + lo * l.width * 2, l.width, hi - lo, out, true, prompt); +} + +/// The block up to its down projection: [gate | up] and SwiGLU into s.act. +pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool) !void { + const w = l.width; + for ([_]cuda.qmmf.Weight{ l.gate, l.up }, 0..) |proj, j| { + // Fp4Linear(out=a column slice): the rows computed contiguous, then copied into their half of each row + try mm(k, stream, proj, x, x_stride, rows, s.tmp, false, prompt); + var a: cuda.Args = .{}; + a.add(s.tmp); + a.add(@as(c_longlong, @intCast(w * 2))); + a.add(s.gu + j * w * 2); + a.add(@as(c_longlong, @intCast(2 * w * 2))); + a.add(@as(c_int, @intCast(rows))); + a.add(@as(c_int, @intCast(w * 2))); + try cuda.launch.launch(k.copy_rows, .{ .grid = .{ .x = @intCast((rows * w * 2 + 255) / 256) }, .block = .{ .x = 256 } }, stream, &a); + } + try glue.swiglu(k.set, stream, s.gu, s.act, s.xs_act, l.limit, rows, w); +} diff --git a/zig/src/families/glm/cuda_moe.zig b/zig/src/families/glm/cuda_moe.zig new file mode 100644 index 0000000000..ec860a188e --- /dev/null +++ b/zig/src/families/glm/cuda_moe.zig @@ -0,0 +1,76 @@ +//! GLM-5.3-Flash's MoE on CUDA, a rank's share of the Python family's moe_block: top-8 plus shared, grouped NVFP4. + +const std = @import("std"); +const cuda = @import("cuda"); +const glue = @import("cuda_glue.zig"); +const experts = @import("cuda_experts.zig"); +const b16 = @import("cuda_b16.zig"); + +/// One MoE layer on a rank: router (bf16 [count][dims]) and fp32 bias, routed experts, shared [gate | up] and down. +pub const Layer = struct { + router: u64, + bias: u64, + routed: cuda.experts.Layer, + count: usize, + shared_gu: u64, + shared_down: u64, + shared: usize, + top_k: usize, + scale: f32, + norm: bool, +}; + +/// The block's scratch for up to `rows` rows (forward.Buffers' MoE part). +pub const Scratch = struct { + mlog: u64, // fp32 [rows][count] + router_part: u64, // fp32 [router_ks][rows][count] + pick: u64, // i32 [rows][top_k + 1] + wts: u64, // fp32 [rows][top_k + 1] + plan: cuda.grouped.Plan, + eact: u64, // bf16 [rows (top_k + 1)][width] + ey: u64, // [rows][top_k + 1][dims]: fp32, bf16 in a prompt chunk + sgu: u64, // bf16 [rows][2 shared] + sact: u64, // bf16 [rows][shared] + sxs: u64, // fp32 [rows][shared / 64] + sy: u64, // fp32 [rows][dims] + mm_part: u64, // fp32 split-K partials of the shared expert's matmuls +}; + +/// The kernels the block runs: the captured Triton set, the grouped NVFP4 experts, the plan, and glm_ops. +pub const Kernels = struct { set: glue.Set, experts: *const experts.Kernels, put_slot: cuda.Function }; + +/// forward.moe_block: out [rows][dims] fp32 partial; `prompt`: 64-pair items, the prompt kernels, bf16 slot rows. +pub fn block(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool, out: u64) !void { + try body(k, stream, l, s, x, x_stride, rows, prompt); + try fill(k, stream, l, s, x, x_stride, 0, rows, prompt, out); +} + +/// The block up to its fill: the router, the top-k pick, the plan and the routed experts into the slot rows. +pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool) !void { + const d = l.routed.dims; + const slots = l.top_k + 1; + const shared_id: i32 = @intCast(l.count); + try glue.router(k.set, stream, x, x_stride, l.router, s.mlog, s.router_part, rows, d, l.count); + try glue.select(k.set, stream, s.mlog, l.bias, s.pick, s.wts, rows, l.top_k, l.count, l.scale, l.norm); + const plan: experts.Plan = .{ .route = s.plan, .slots = slots, .experts = l.count + 1, .prompt = prompt }; + try experts.route(k.experts, stream, s.pick, rows * slots, plan); + try experts.run(k.experts, stream, .gate_up, l.routed, plan, x, x_stride, rows, s.eact, shared_id); + try experts.run(k.experts, stream, if (prompt) .down else .down_f32, l.routed, plan, s.eact, l.routed.width, rows, s.ey, shared_id); +} + +/// moe_block's fill of rows lo..hi: shared expert into the last slot, slot sum into out fp32 (any range, same bits). +pub fn fill(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, lo: usize, hi: usize, prompt: bool, out: u64) !void { + const d = l.routed.dims; + const slots = l.top_k + 1; + const n = hi - lo; + const eb: usize = if (prompt) 2 else 4; // the slot rows: bf16 in a prompt chunk + try b16.matmul(k.set, stream, x + lo * x_stride * 2, x_stride, l.shared_gu, s.sgu + lo * 2 * l.shared * 2, s.mm_part, n, 2 * l.shared, d, false); + try glue.swiglu(k.set, stream, s.sgu + lo * 2 * l.shared * 2, s.sact + lo * l.shared * 2, s.sxs + lo * (l.shared / 64) * 4, l.routed.limit, n, l.shared); + try b16.matmul(k.set, stream, s.sact + lo * l.shared * 2, l.shared, l.shared_down, s.sy + lo * d * 4, s.mm_part, n, d, l.shared, true); + var a: cuda.Args = .{}; + a.add(s.sy + lo * d * 4); + a.add(s.ey + lo * slots * d * eb); + for ([_]usize{ n, d, slots, l.top_k, @intFromBool(prompt) }) |v| a.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(k.put_slot, .{ .grid = .{ .x = @intCast((n * d + 255) / 256) }, .block = .{ .x = 256 } }, stream, &a); + try glue.combine(k.set, stream, s.ey + lo * slots * d * eb, !prompt, s.wts + lo * slots * 4, out, n, d, slots); +} diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig new file mode 100644 index 0000000000..2e5fbef14e --- /dev/null +++ b/zig/src/families/glm/cuda_native.zig @@ -0,0 +1,205 @@ +//! GLM-5.3-Flash NVFP4 native CUDA at TP over processes (TF_GLM_RANK, NCCL_COMM_ID): rank 0 serves, others follow. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const follow = @import("cuda_follow.zig"); +const glm_lanes = @import("cuda_lanes.zig"); +const ranks = @import("cuda_ranks.zig"); +const fw = @import("cuda_forward.zig"); +const config = @import("config.zig"); +const shard = @import("cuda_shard.zig"); + +pub const model_type = "glm5_next"; +pub const formats: []const []const u8 = &.{"modelopt-nvfp4"}; +// cc-defer: a fixed window; the Python family sizes it from free memory (capacity.admit), port that when serving needs +pub const default_context: i64 = 8192; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = fw.prefill_rows; +/// Prompt kernels change arithmetic with a chunk's rows: states keep and resume only on chunk starts. +pub const snap_grid: u32 = prompt_rows; +/// `explain` words this family's open refusals too (where to set the rank and the id). +pub const explains_open = true; + +/// The server's reading of this rank's room after its load (the budget streams are admitted by); null if unreadable. +pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 }; +/// `slide` stays false: the server refuses --slide for a family without a learner (no `slide` here). +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, room: ?Room = null }; +/// The smallest room among the other ranks and whose it is. +pub const Peer = struct { bytes: u64, rank: usize }; + +pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; + +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32 = 1, + learn_bytes: usize = 0, // no learner: the streams keep all the room + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, + target: *engine.Engine, // the prompt cache keeps this engine's states (snap_*) + peer_room: ?Peer = null, // rank 0: the server admits streams and kept states by this when it is smaller +}; + +/// This rank's place: TF_GLM_RANK of TF_GLM_WORLD (both, or neither for a refusal that names them). +pub const Place = struct { rank: usize, world: usize }; + +pub fn place() !Place { + const r = std.c.getenv("TF_GLM_RANK") orelse return error.NoRank; + const w = std.c.getenv("TF_GLM_WORLD") orelse return error.NoRank; + const p: Place = .{ .rank = std.fmt.parseInt(usize, std.mem.span(r), 10) catch return error.NoRank, .world = std.fmt.parseInt(usize, std.mem.span(w), 10) catch return error.NoRank }; + if (p.world < 1 or p.rank >= p.world) return error.NoRank; + return p; +} + +/// Bytes this rank's weights take at most, for the server's admission: the checkpoint times its largest axis share. +pub fn weightShare(a: std.mem.Allocator, io: std.Io, dir: []const u8, total: u64) u64 { + const p = place() catch return total; + const path = std.fs.path.join(a, &.{ dir, "config.json" }) catch return total; + const text = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 24)) catch return total; + const c = config.parse(a, text) catch return total; + const f = (shard.Share.of(&c, p.rank, p.world) catch return total).part(&c); + return @intCast(@as(u128, total) * f[0] / f[1]); +} + +/// Slots for `context` plus a verify window, a multiple of 64 so the sparse pool keeps the AOT set. +pub fn capacity(context: usize) usize { + return std.mem.alignForward(usize, context + fw.window_rows, 64); +} + +const Owned = struct { + gpa: std.mem.Allocator, + lib: cuda.nccl.Library, + comm: cuda.nccl.Comm, + e: *engine.Engine, + link: follow.Link, + lanes: glm_lanes.Lanes, +}; + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + _ = kernels; // our own kernels: a captured set is not used + const p = try place(); + if (p.world > 1 and std.c.getenv("NCCL_COMM_ID") == null) return error.NoCommId; + var problem: []const u8 = ""; + const set = ranks.Settings.read(&ranks.environment, &problem) catch |err| { + std.log.err("{s}", .{problem}); + return err; + }; + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.lib = try cuda.nccl.Library.open(); + errdefer own.lib.close(); + var id: cuda.nccl.UniqueId = undefined; + try own.lib.check(own.lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); // NCCL_COMM_ID: the same id on every rank + own.comm = null; + try own.lib.check(own.lib.api.ncclCommInitRank(&own.comm, @intCast(p.world), id, @intCast(p.rank)), "ncclCommInitRank"); + errdefer _ = own.lib.api.ncclCommDestroy(own.comm); + const cap = capacity(o.context); + own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, &own.lib, own.comm, p.rank, p.world, .{ .capacity = cap, .mtp = o.drafts, .settings = set }); + errdefer own.e.deinit(); + own.link = try follow.Link.init(gpa, ctx.d, &own.lib, own.comm, own.e.stream, cap); + errdefer own.link.deinit(); + // Rank 0 admits by the smallest room of the ranks; its sizes are the largest (the first ranks take the remainder). + const rooms = try gpa.alloc(u64, p.world); + defer gpa.free(rooms); + const mine: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); + try ranks.gather(u64, ctx.d, &own.lib, own.comm, own.e.stream, p.rank, mine, rooms); + var peer: ?Peer = null; + for (rooms[1..], 1..) |r, q| if (peer == null or r < peer.?.bytes) { + peer = .{ .bytes = r, .rank = q }; + }; + if (p.rank != 0) { // a follower never serves: it runs rank 0's steps, then leaves the process + follow.follow(own.e, &own.link) catch |err| { + std.log.err("rank {d}: following rank 0 stopped ({s})", .{ p.rank, @errorName(err) }); + std.process.exit(1); + }; + std.log.info("rank {d}: rank 0 ended; leaving", .{p.rank}); + std.process.exit(0); + } + // Rank 0 alone reads these: TF_GLM_PROMPT_STEPS=0 passes prompts whole, TF_GLM_SHARED=0 runs windows alone. + const off = struct { + fn get(name: [*:0]const u8) bool { + return if (std.c.getenv(name)) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; + } + }.get; + own.lanes = .{ .x = .{ .e = own.e, .link = &own.link }, .whole = off("TF_GLM_PROMPT_STEPS"), .alone = off("TF_GLM_SHARED") }; + return .{ + .backend = own.lanes.backend(), + .facts = own.lanes.facts(), + .rows = glm_lanes.max_rows, + .stream_bytes = own.e.seqBytes(), // each stream past the first makes its own sequence on every rank + .free_streams = 1, // the first is the sequence the engine made + .ctx = own, + .deinit = close, + .target = own.e, + .peer_room = peer, + }; +} + +// Each rank keeps a copy under rank 0's id, in the room rank 0 admitted by (the smallest over the ranks). + +/// Bytes a state takes at any position: the whole state arena on each rank. +pub fn snap_bytes(ptr: *anyopaque, at: u32) u64 { + _ = at; + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + return e.stateBytes(); +} + +/// The prompt pass stands at `at` (a mark at a chunk start): copy the state there under a new id. +pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + _ = owner; + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + if (e.pos() != at) return error.NotAtMark; + const k = try e.gpa.create(engine.Kept); + errdefer e.gpa.destroy(k); + k.* = .{ .id = try e.keepState(), .at = at }; + return k; +} + +/// A restore outside a prompt pass (the store's `begin`) would leave the other ranks behind, so it refuses. +pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + _ = ptr; + _ = owner; + _ = saved; + return error.RestoredInPromptPass; +} + +pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + const k: *engine.Kept = @ptrCast(@alignCast(saved)); + e.dropState(k.id); + e.gpa.destroy(k); +} + +fn close(ptr: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(ptr)); + own.lanes.x.end() catch {}; + own.lanes.deinit(); + own.link.deinit(); + own.e.deinit(); + _ = own.lib.api.ncclCommDestroy(own.comm); + own.lib.close(); + own.gpa.destroy(own); +} + +/// The server's words for this family's refusals. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.NoRank => "GLM-5.3-Flash runs at TP over hosts: set TF_GLM_RANK and TF_GLM_WORLD on every host (WORLD 1: one GPU, with TF_GLM_LAYERS for a thin model)", + error.NoCommId => "GLM-5.3-Flash's ranks meet over NCCL: with TF_GLM_WORLD 2 or more, set NCCL_COMM_ID to rank 0's address:port on every host", + error.RankMismatch => "the ranks' settings differ (each rank's log names the one): start them alike", + error.NotModelOpt, error.UnsupportedQuant => "GLM-5.3-Flash on CUDA reads ModelOpt NVFP4 checkpoints (the log names what config.json has)", + else => null, + }; +} + +test "capacity keeps a verify window past the context and a multiple of 64" { + try std.testing.expectEqual(@as(usize, 8256), capacity(8192)); + try std.testing.expectEqual(@as(usize, 330560), capacity(330529)); + try std.testing.expect(capacity(330529) % 64 == 0); +} diff --git a/zig/src/families/glm/cuda_q4.zig b/zig/src/families/glm/cuda_q4.zig new file mode 100644 index 0000000000..382ca2932c --- /dev/null +++ b/zig/src/families/glm/cuda_q4.zig @@ -0,0 +1,108 @@ +//! GLM-5.3-Flash's affine 4-bit matmul for the draft head (the Python family's qmm.matmul, tile 2) and its tiling. + +const std = @import("std"); +const cuda = @import("cuda"); + +/// qmm_group.cu's tile-2 instantiation (LaneTile<64, 16, 64, 1, 4, 8>, bf16 out), as zig/kernels/cuda exports it. +const group_symbol = "_ZN12tf_qmm_group12group_kernelILi64ELi16ELi64ELi1ELi4ELi8ELb0ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfNS_5PartsEiiiii"; +const group_smem: u32 = 35328; // LaneTile<64, 16, 64, 1, 4, 8>::SMEM + +/// qmm_group.cu's Part and Parts, passed by value: four projections at most, one used here. +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 }; +const Parts = extern struct { p: [4]Part, count: c_int }; + +comptime { + std.debug.assert(@sizeOf(Part) == 56 and @sizeOf(Parts) == 232); +} + +/// A tiled 4-bit matrix [n][k] (shared.Q4): packed words, (k / 64, npad) bf16 scales and biases, npad = n up to 128s. +pub const Q4 = struct { w: u64, s: u64, b: u64, n: usize, k: usize, npad: usize }; + +pub const Kernels = struct { + group: cuda.Function, + pack_dense: cuda.Function, + transpose: cuda.Function, + gb10: bool, // qmm_group_cuda's programmatic dependent launch: on for GB10 (sm_121) + + /// From the qmm_group and affine4_pack modules. + pub fn load(group: cuda.Module, ops: cuda.Module, ctx: *const cuda.Context) !Kernels { + const k: Kernels = .{ + .group = try group.function(group_symbol), + .pack_dense = try ops.function("tf_pack_dense"), + .transpose = try ops.function("tf_transpose_pad16"), + .gb10 = try ctx.attribute(.compute_capability_major) == 12 and try ctx.attribute(.compute_capability_minor) == 1, + }; + try k.group.allowDynamicShared(group_smem); + return k; + } +}; + +pub fn npadOf(n: usize) usize { + return (n + 127) / 128 * 128; +} + +/// Device bytes of a Q4's words, scales and biases. +pub fn bytes(n: usize, k: usize) [3]usize { + const p = npadOf(n); + return .{ p * k / 2, k / 64 * p * 2, k / 64 * p * 2 }; +} + +/// 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; +} + +/// shared.pack: MLX words i32 [n][k/8] and bf16 scales and biases [n][k/64] -> `q` (caller's buffers, `bytes` long). +pub fn pack(ks: *const Kernels, stream: cuda.Stream, words: u64, scales: u64, biases: u64, q: Q4) !void { + const kg = q.k / 64; + const total: u64 = @as(u64, q.npad / 64) * kg * 512; + var a: cuda.Args = .{}; + a.add(words); + for ([_]usize{ q.n, q.k / 8, kg }) |v| a.add(@as(c_int, @intCast(v))); + a.add(q.w); + a.add(@as(c_longlong, @intCast(total))); + try cuda.launch.launch(ks.pack_dense, .{ .grid = .{ .x = @intCast((total + 255) / 256) }, .block = .{ .x = 256 } }, stream, &a); + for ([_]u64{ scales, biases }, [_]u64{ q.s, q.b }) |src, dst| { + var t: cuda.Args = .{}; + t.add(src); + for ([_]usize{ q.n, kg, q.npad }) |v| t.add(@as(c_int, @intCast(v))); + t.add(dst); + const cells = @as(u64, kg) * q.npad; + try cuda.launch.launch(ks.transpose, .{ .grid = .{ .x = @intCast((cells + 255) / 256) }, .block = .{ .x = 256 } }, stream, &t); + } +} + +/// qmm.matmul of a Q4 (tile 2, up to 16 rows): x bf16 (`ldx` apart), group sums xs fp32 [rows][k/64] -> out bf16. +pub fn matmul(ks: *const Kernels, stream: cuda.Stream, x: u64, ldx: usize, xs: u64, q: Q4, out: u64, rows: usize) !void { + if (rows == 0 or rows > 16) return error.WindowTooWide; + const sk = splitK(q.n, q.k); + 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 = @intCast(q.n), .npad = @intCast(q.npad), .sk = @intCast(sk), .tiles = @intCast(tiles), .first = 0 }; + const rows_t = (rows + 15) / 16; + const clusters = rows_t * tiles; // one part: a cluster is the sk K slices of one column tile + var a: cuda.Args = .{}; + a.add(x); + a.add(xs); + a.add(parts); + for ([_]usize{ rows, q.k, if (rows == 1) q.k else ldx, rows_t, sk }) |v| a.add(@as(c_int, @intCast(v))); // M, K, ldx, rows_t, C + const cfg: cuda.Config = .{ + .grid = .{ .x = @intCast(clusters * sk) }, + .block = .{ .x = 128 }, + .shared = group_smem, + .cluster = if (sk > 1) .{ .x = @intCast(sk) } else null, + .pdl = ks.gb10, + }; + try cuda.launch.launch(ks.group, cfg, stream, &a); +} + +test "the draft head's shape takes one K slice; npad rounds to 128s" { + try std.testing.expectEqual(@as(usize, 1), splitK(77440, 4096)); + try std.testing.expectEqual(@as(usize, 77440), npadOf(77440)); + try std.testing.expectEqual(@as(usize, 51712), npadOf(51627)); +} diff --git a/zig/src/families/glm/cuda_quant.zig b/zig/src/families/glm/cuda_quant.zig new file mode 100644 index 0000000000..08fb0c9d37 --- /dev/null +++ b/zig/src/families/glm/cuda_quant.zig @@ -0,0 +1,57 @@ +//! GLM-5.3-Flash's draft-only weight quantizers on CUDA (glm_quant.cu): the Python family's load-time ops, same bits. + +const cuda = @import("cuda"); + +/// glm_quant.cu's entries. +pub const Kernels = struct { + amax: cuda.Function, + nvfp4: cuda.Function, + q4: cuda.Function, + + pub fn load(m: cuda.Module) !Kernels { + return .{ + .amax = try m.function("tf_glm_nvfp4_amax"), + .nvfp4 = try m.function("tf_glm_nvfp4_quantize"), + .q4 = try m.function("tf_glm_quantize4"), + }; + } +}; + +fn blocks(threads: usize) u32 { + return @intCast((threads + 255) / 256); +} + +/// experts.quantize: w bf16 [e][n][k] -> codes u8 [e][n][k/2], e4m3 scales [e][n][k/16], g fp32 [e]; 16-byte aligned. +pub fn nvfp4(ks: *const Kernels, stream: cuda.Stream, w: u64, e: usize, n: usize, k: usize, codes: u64, scales: u64, g: u64) !void { + if (k % 16 != 0) return error.KNotMultipleOf16; + if (e == 0 or n == 0) return; + const per = n * k; + var a: cuda.Args = .{}; + a.add(w); + a.add(@as(c_longlong, @intCast(per))); + a.add(g); + try cuda.launch.launch(ks.amax, .{ .grid = .{ .x = @intCast(e) }, .block = .{ .x = 1024 } }, stream, &a); + const total = e * per / 16; + var b: cuda.Args = .{}; + b.add(w); + b.add(g); + b.add(@as(c_longlong, @intCast(per / 16))); + b.add(@as(c_longlong, @intCast(total))); + b.add(codes); + b.add(scales); + try cuda.launch.launch(ks.nvfp4, .{ .grid = .{ .x = blocks(total) }, .block = .{ .x = 256 } }, stream, &b); +} + +/// qmm.quantize4: w bf16 [n][k] -> MLX words i32 [n][k/8], scales and biases bf16 [n][k/64]; 16-byte aligned. +pub fn q4(ks: *const Kernels, stream: cuda.Stream, w: u64, n: usize, k: usize, words: u64, scales: u64, biases: u64) !void { + if (k % 64 != 0) return error.KNotMultipleOf64; + const groups = n * k / 64; + if (groups == 0) return; + var a: cuda.Args = .{}; + a.add(w); + a.add(@as(c_longlong, @intCast(groups))); + a.add(words); + a.add(scales); + a.add(biases); + try cuda.launch.launch(ks.q4, .{ .grid = .{ .x = blocks(groups) }, .block = .{ .x = 256 } }, stream, &a); +} diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig new file mode 100644 index 0000000000..50b70dd0c8 --- /dev/null +++ b/zig/src/families/glm/cuda_ranks.zig @@ -0,0 +1,191 @@ +//! A TP rank's settings fingerprint, all-gathered before weights load: a mismatch stops every rank. + +const std = @import("std"); +const cuda = @import("cuda"); +const heat = @import("heat"); +const Sha256 = std.crypto.hash.sha2.Sha256; + +/// The forward's knobs a rank reads from its environment (the Python family's names and defaults). +pub const Settings = struct { + split: bool = true, // TF_GLM_PREFILL_REDUCE: split (each rank glues its own rows) or gather + kda_wide: bool = true, // TF_GLM_KDA_DECODE_WIDE: 1 (decode windows on KDA's three-kernel chain) or 0 + pieces: u32 = 3, // TF_GLM_OVERLAP_PIECES: a prompt chunk's exchanges in this many row pieces (1: whole). + heat: ?heat.Bands = null, // TF_HEAT_HIGH and TF_HEAT_LOW (upstream's names): a prompt chunk waits on every rank's hottest zone + layers: u32 = 0, // TF_GLM_LAYERS: the first this many layers only, a thin model for one-GPU checks (0: all) + + /// The settings `get` (an environment lookup) gives; `problem` names a value that is not one of the choices. + pub fn read(get: *const fn ([:0]const u8) ?[]const u8, problem: *[]const u8) !Settings { + var s: Settings = .{}; + if (get("TF_GLM_PREFILL_REDUCE")) |v| { + if (std.mem.eql(u8, v, "split")) s.split = true else if (std.mem.eql(u8, v, "gather")) s.split = false else { + problem.* = "TF_GLM_PREFILL_REDUCE: split or gather"; + return error.BadSetting; + } + } + if (get("TF_GLM_KDA_DECODE_WIDE")) |v| { + if (std.mem.eql(u8, v, "1")) s.kda_wide = true else if (std.mem.eql(u8, v, "0")) s.kda_wide = false else { + problem.* = "TF_GLM_KDA_DECODE_WIDE: 0 or 1"; + return error.BadSetting; + } + } + if (get("TF_GLM_OVERLAP_PIECES")) |v| { + s.pieces = std.fmt.parseInt(u32, v, 10) catch 0; + if (s.pieces < 1 or s.pieces > 8) { + problem.* = "TF_GLM_OVERLAP_PIECES: 1 to 8"; + return error.BadSetting; + } + } + if (get("TF_GLM_LAYERS")) |v| { + s.layers = std.fmt.parseInt(u32, v, 10) catch 0; + if (s.layers == 0) { + problem.* = "TF_GLM_LAYERS: the first how many layers, 1 or more"; + return error.BadSetting; + } + } + s.heat = heat.bandsFrom(get("TF_HEAT_HIGH"), get("TF_HEAT_LOW")) catch { + problem.* = "TF_HEAT_HIGH and TF_HEAT_LOW: both or neither, in degrees, with low at or under high"; + return error.BadSetting; + }; + return s; + } +}; + +/// The process environment's value of `name` (null when unset). +pub fn environment(name: [:0]const u8) ?[]const u8 { + return std.mem.span(std.c.getenv(name) orelse return null); +} + +/// What every rank must agree on, as bytes the communicator carries. +pub const Fingerprint = extern struct { + layout: u32 = 4, // this struct's version + world: u32, + rank: u32, + split: u32, + kda_wide: u32, + pieces: u32, + layers: u32, // TF_GLM_LAYERS (0: all) + pad: u32 = 0, + capacity: u64, // tokens a sequence holds + heat_high: f64, // the heat bands in degrees (0 and 0: off): every rank must wait alike + heat_low: f64, + config: [32]u8, // sha256 of config.json + program: [32]u8, // sha256 of this executable + checkpoint: [64]u8, // the checkpoint folder's last name (the snapshot's revision), zero padded + + pub const Field = enum { layout, rank, world, split, kda_wide, pieces, layers, capacity, heat_high, heat_low, config, program, checkpoint }; + + pub fn of(world: usize, rank: usize, s: Settings, capacity: usize, config: [32]u8, program: [32]u8, checkpoint_dir: []const u8) Fingerprint { + var f: Fingerprint = .{ .world = @intCast(world), .rank = @intCast(rank), .split = @intFromBool(s.split), .kda_wide = @intFromBool(s.kda_wide), .pieces = s.pieces, .layers = s.layers, .capacity = capacity, .heat_high = if (s.heat) |b| b.high_c else 0, .heat_low = if (s.heat) |b| b.low_c else 0, .config = config, .program = program, .checkpoint = @splat(0) }; + const name = std.fs.path.basename(std.mem.trimEnd(u8, checkpoint_dir, "/")); + const n = @min(name.len, f.checkpoint.len); + @memcpy(f.checkpoint[0..n], name[0..n]); + return f; + } +}; + +pub const Mismatch = struct { field: Fingerprint.Field, rank: usize }; + +/// The first field where rank `r`'s fingerprint differs from rank 0's (or holds the wrong rank), null when all agree. +pub fn mismatch(all: []const Fingerprint) ?Mismatch { + for (all, 0..) |f, r| { + if (f.layout != all[0].layout) return .{ .field = .layout, .rank = r }; + if (f.rank != r) return .{ .field = .rank, .rank = r }; + inline for (.{ "world", "split", "kda_wide", "pieces", "layers", "capacity", "heat_high", "heat_low", "config", "program", "checkpoint" }) |name| { + if (!std.mem.eql(u8, std.mem.asBytes(&@field(f, name)), std.mem.asBytes(&@field(all[0], name)))) + return .{ .field = @field(Fingerprint.Field, name), .rank = r }; + } + } + return null; +} + +/// sha256 of a file's bytes (read whole: config.json, the executable). +pub fn hashFile(gpa: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 30)); + defer gpa.free(bytes); + return hashBytes(bytes); +} + +pub fn hashBytes(bytes: []const u8) [32]u8 { + var out: [32]u8 = undefined; + Sha256.hash(bytes, &out, .{}); + return out; +} + +/// Every rank's `mine` in rank order, over one ncclAllGather. +pub fn gather(comptime T: type, d: *const cuda.Driver, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, stream: cuda.Stream, rank: usize, mine: T, out: []T) !void { + const size = @sizeOf(T); + var dev = try cuda.DeviceBuffer.alloc(d, size * out.len); + defer dev.free(); + try dev.uploadAsync(rank * size, std.mem.asBytes(&mine), stream.handle); + try lib.check(lib.api.ncclAllGather(dev.ptr + rank * size, dev.ptr, size, .u8, comm, stream.handle), "ncclAllGather"); + try stream.synchronize(); + try dev.download(0, std.mem.sliceAsBytes(out)); +} + +/// All-gathers `mine` over `comm`; any disagreement is error.RankMismatch on every rank, printing what differs. +pub fn agree(d: *const cuda.Driver, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, stream: cuda.Stream, mine: Fingerprint, out: []Fingerprint) !void { + try gather(Fingerprint, d, lib, comm, stream, mine.rank, mine, out); + if (mismatch(out)) |m| { + std.debug.print("REFUSE rank settings: rank {d}'s {s} differs from rank 0's ({s})\n", .{ m.rank, @tagName(m.field), explain(m.field) }); + return error.RankMismatch; + } +} + +fn explain(f: Fingerprint.Field) []const u8 { + return switch (f) { + .layout => "a different build of this check", + .rank => "a rank number given twice or out of order", + .world => "the ranks were started with different world sizes", + .split => "TF_GLM_PREFILL_REDUCE", + .kda_wide => "TF_GLM_KDA_DECODE_WIDE", + .pieces => "TF_GLM_OVERLAP_PIECES", + .layers => "TF_GLM_LAYERS", + .capacity => "the sequence capacity (the context)", + .heat_high => "TF_HEAT_HIGH", + .heat_low => "TF_HEAT_LOW", + .config => "config.json", + .program => "the executable", + .checkpoint => "the checkpoint's snapshot folder", + }; +} + +fn fake(name: [:0]const u8) ?[]const u8 { + if (std.mem.eql(u8, name, "TF_GLM_PREFILL_REDUCE")) return "gather"; + if (std.mem.eql(u8, name, "TF_GLM_OVERLAP_PIECES")) return "2"; + if (std.mem.eql(u8, name, "TF_HEAT_HIGH")) return "92"; + if (std.mem.eql(u8, name, "TF_HEAT_LOW")) return "88"; + return null; +} + +fn bad(name: [:0]const u8) ?[]const u8 { + if (std.mem.eql(u8, name, "TF_GLM_PREFILL_REDUCE")) return "sum"; + return null; +} + +test "settings from the environment: the defaults, the choices, and a refusal naming a bad value" { + var problem: []const u8 = ""; + const s = try Settings.read(&fake, &problem); + try std.testing.expect(!s.split and s.kda_wide and s.pieces == 2 and s.heat.?.high_c == 92 and s.heat.?.low_c == 88); + try std.testing.expectError(error.BadSetting, Settings.read(&bad, &problem)); + try std.testing.expectEqualStrings("TF_GLM_PREFILL_REDUCE: split or gather", problem); +} + +test "fingerprints: agreeing ranks pass; the first differing field and rank are named" { + const h0: [32]u8 = @splat(1); + const h1: [32]u8 = @splat(2); + var all = [_]Fingerprint{ + Fingerprint.of(2, 0, .{}, 8192, h0, h0, "/hub/snapshots/423acf37/"), + Fingerprint.of(2, 1, .{}, 8192, h0, h0, "/hub/snapshots/423acf37"), + }; + try std.testing.expect(mismatch(&all) == null); + all[1].program = h1; + try std.testing.expectEqual(Fingerprint.Field.program, mismatch(&all).?.field); + all[1] = Fingerprint.of(2, 1, .{ .split = false }, 8192, h0, h0, "/hub/snapshots/423acf37"); + try std.testing.expectEqual(Mismatch{ .field = .split, .rank = 1 }, mismatch(&all).?); + all[1] = Fingerprint.of(2, 0, .{}, 8192, h0, h0, "/hub/snapshots/423acf37"); + try std.testing.expectEqual(Fingerprint.Field.rank, mismatch(&all).?.field); + all[1] = Fingerprint.of(2, 1, .{ .heat = .{ .high_c = 92, .low_c = 88 } }, 8192, h0, h0, "/hub/snapshots/423acf37"); + try std.testing.expectEqual(Fingerprint.Field.heat_high, mismatch(&all).?.field); + all[1] = Fingerprint.of(2, 1, .{}, 8192, h0, h0, "/hub/snapshots/other"); + try std.testing.expectEqual(Fingerprint.Field.checkpoint, mismatch(&all).?.field); +} diff --git a/zig/src/families/glm/cuda_sampler.zig b/zig/src/families/glm/cuda_sampler.zig new file mode 100644 index 0000000000..5eea093bb3 --- /dev/null +++ b/zig/src/families/glm/cuda_sampler.zig @@ -0,0 +1,118 @@ +//! Sampled draws on upstream's sample.cu (#542 option C): every rank runs tf_draw. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const Span = @import("cuda_shard.zig").Span; +const Exchange = @import("cuda_forward.zig").Exchange; + +pub const Sampling = lanes.Sampling; + +/// The request's rule as the engine runs it: null (or temperature 0) decodes greedily. +pub fn check(s: ?Sampling) ?Sampling { + const x = s orelse return null; + return if (x.temperature > 0) x else null; +} + +/// sample.cu's Rule: the seed, 1 / T, top_p, the candidates' near window, ln(min_p) and top_k (0: off). +pub const Rule = extern struct { seed: u64, inv_t: f32, top_p: f32, near: f32, min_log: f32, top_k: u32, pad: u32 = 0 }; + +/// The rule as nemotron/cuda_sampler.zig fills it. +pub fn rule(s: Sampling) Rule { + return .{ .seed = s.seed, .inv_t = @floatCast(1.0 / @max(s.temperature, 1e-6)), .top_p = @floatCast(s.top_p), .near = 20.0, .min_log = @floatCast(s.minLog()), .top_k = s.top_k }; +} + +/// A rank's logits rows: bf16 [rows][its vocabulary span]. +pub const Local = struct { rank: usize, logits: u64 }; + +pub const Sampler = struct { + d: *const cuda.Driver, + module: cuda.Module, + tf_draw: cuda.Function, + vocab: usize, + rows: usize, + full: cuda.DeviceBuffer, // bf16 [rows][vocab]: every rank's columns in place + rule: cuda.DeviceBuffer, + meta: cuda.DeviceBuffer, + tok: cuda.DeviceBuffer, + prob: cuda.DeviceBuffer, + + pub fn init(d: *const cuda.Driver, vocab: usize, rows: usize) !Sampler { + var m = try cuda.Module.load(d, cuda.kernels.sample); + errdefer m.unload(); + const f = try m.function("tf_draw"); + var full = try cuda.DeviceBuffer.alloc(d, rows * vocab * 2); + errdefer full.free(); + var r = try cuda.DeviceBuffer.alloc(d, @sizeOf(Rule)); + errdefer r.free(); + var meta = try cuda.DeviceBuffer.alloc(d, 16); + errdefer meta.free(); + var tok = try cuda.DeviceBuffer.alloc(d, rows * 4); + errdefer tok.free(); + const prob = try cuda.DeviceBuffer.alloc(d, rows * 4); + return .{ .d = d, .module = m, .tf_draw = f, .vocab = vocab, .rows = rows, .full = full, .rule = r, .meta = meta, .tok = tok, .prob = prob }; + } + + pub fn deinit(s: *Sampler) void { + s.prob.free(); + s.tok.free(); + s.meta.free(); + s.rule.free(); + s.full.free(); + s.module.unload(); + } + + /// `rows` rows of `mine` -> `out`, tokens at positions first, first + 1, ...; `spans`: every rank's, rank order. + pub fn draw(s: *Sampler, stream: cuda.Stream, exchange: Exchange, spans: []const Span, mine: []const Local, rows: usize, first: u64, r: Rule, out: []u32) !void { + std.debug.assert(rows <= s.rows and out.len >= rows and first >= 1); + const V = s.vocab; + for (mine) |l| { + const span = spans[l.rank]; + for (0..rows) |row| try s.copy(stream, s.full.ptr + (row * V + span.lo) * 2, l.logits + row * span.len() * 2, span.len() * 2); + } + switch (exchange) { + .local => {}, + .nccl => |x| { // each rank's columns of each row from that rank (spans may differ in width: TP=3) + try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart"); + for (0..rows) |row| for (spans, 0..) |span, q| { + const at = s.full.ptr + (row * V + span.lo) * 2; + try x.lib.check(x.lib.api.ncclBroadcast(at, at, span.len(), .bf16, @intCast(q), x.comm, stream.handle), "ncclBroadcast"); + }; + try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd"); + }, + } + try s.rule.uploadAsync(0, std.mem.asBytes(&r), stream.handle); + const meta = [4]i32{ @intCast(first - 1), 0, 0, 0 }; // row r keys meta[0] + r + 1 + try s.meta.uploadAsync(0, std.mem.asBytes(&meta), stream.handle); + var args: cuda.Args = .{}; + args.add(s.full.ptr); + args.add(@as(u32, @intCast(V))); + args.add(s.rule.ptr); + args.add(s.meta.ptr); + args.add(@as(i32, 0)); + args.add(s.tok.ptr); + args.add(s.prob.ptr); + try cuda.launch.launch(s.tf_draw, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 1024 } }, stream, &args); + try stream.synchronize(); + try s.tok.download(0, std.mem.sliceAsBytes(out[0..rows])); + } + + /// The whole rows a draw read (bf16 [rows][vocab]), for checks against the host rule. + pub fn rowsRead(s: *const Sampler, rows: usize, host: []u16) !void { + try s.full.download(0, std.mem.sliceAsBytes(host[0 .. rows * s.vocab])); + } + + fn copy(s: *const Sampler, stream: cuda.Stream, dst: u64, src: u64, bytes: usize) !void { + try s.d.check(s.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, stream.handle), "cuMemcpyDtoDAsync"); + } +}; + +test "rules as nemotron/cuda_sampler fills them; temperature 0 is greedy" { + const s: Sampling = .{ .seed = 1234, .temperature = 0.5, .top_k = 20, .top_p = 0.95 }; + const r = rule(check(s).?); + try std.testing.expectEqual(@as(f32, 2.0), r.inv_t); + try std.testing.expectEqual(@as(u32, 20), r.top_k); + try std.testing.expect(std.math.isInf(r.min_log)); + try std.testing.expectEqual(@as(usize, 32), @sizeOf(Rule)); + try std.testing.expectEqual(@as(?Sampling, null), check(.{ .seed = 0, .temperature = 0 })); +} diff --git a/zig/src/families/glm/cuda_shard.zig b/zig/src/families/glm/cuda_shard.zig new file mode 100644 index 0000000000..eb16059317 --- /dev/null +++ b/zig/src/families/glm/cuda_shard.zig @@ -0,0 +1,100 @@ +//! A TP rank's share of GLM-5.3-Flash on CUDA (the Python family's ShardPlan): heads one by one, widths in 64s. + +const std = @import("std"); +const canon = @import("canon"); +const Config = @import("config.zig").Config; + +/// Elements of an MLP width or the vocabulary a rank takes whole (split.UNIT). +pub const unit = 64; + +pub const Span = struct { + lo: usize, + hi: usize, + + pub fn len(s: Span) usize { + return s.hi - s.lo; + } +}; + +pub const Share = struct { + rank: usize, + world: usize, + kda: Span, // KDA heads + mla: Span, // DSA heads + dense: Span, // the dense MLP's intermediate rows + moe: Span, // each routed expert's intermediate rows + shared: Span, // the shared expert's intermediate rows + vocab: Span, // the head's rows + + /// Rank `rank` of `world`; refused where a rank would be left empty (the Python family's split_units refuses it). + pub fn of(c: *const Config, rank: usize, world: usize) !Share { + if (world == 0 or rank >= world) return error.BadRank; + return .{ + .rank = rank, + .world = world, + .kda = try span(c.kda_heads, 1, rank, world), + .mla = try span(c.mla_heads, 1, rank, world), + .dense = try span(c.dense_inter, unit, rank, world), + .moe = try span(c.moe_inter, unit, rank, world), + .shared = try span(c.shared_inter, unit, rank, world), + .vocab = try span(c.vocab, c.head_unit, rank, world), + }; + } + + /// This rank's largest fraction of an axis as (part, whole): its weights are at most that much of the checkpoint. + pub fn part(s: Share, c: *const Config) [2]u64 { + const axes = [_][2]u64{ .{ s.kda.len(), c.kda_heads }, .{ s.mla.len(), c.mla_heads }, .{ s.dense.len(), c.dense_inter }, .{ s.moe.len(), c.moe_inter }, .{ s.shared.len(), c.shared_inter }, .{ s.vocab.len(), c.vocab } }; + var best: [2]u64 = .{ 0, 1 }; + for (axes) |x| if (x[0] * best[1] > best[0] * x[1]) { + best = .{ x[0], x[1] }; + }; + return best; + } +}; + +/// Rank `rank`'s [lo, hi) of `total` elements taken in whole `per`s. +fn span(total: usize, per: usize, rank: usize, world: usize) !Span { + if (total % per != 0 or world > 256) return error.BadSplit; + const ones: [256]u64 = @splat(1); + var out: [256]canon.Range = undefined; + canon.partition(@intCast(total / per), ones[0..world], out[0..world]); + const r = out[rank]; + if (r.len() == 0) return error.EmptyShare; + return .{ .lo = r.begin * per, .hi = r.end * per }; +} + +test "a rank's part is its largest axis: TP=3 rank 0's 22 of 64 KDA heads, rank 2's 4,096 of 12,288 dense rows" { + const c = Config{}; + const p0 = (try Share.of(&c, 0, 3)).part(&c); + try std.testing.expectEqual(@as(u64, 11) * p0[1], 32 * p0[0]); + const p2 = (try Share.of(&c, 2, 3)).part(&c); + try std.testing.expectEqual(p2[1], 3 * p2[0]); + const h = (try Share.of(&c, 1, 2)).part(&c); + try std.testing.expectEqual(h[1], 2 * h[0]); +} + +test "TP=2 halves every axis; TP=3 takes the Python family's uneven shares (22/21/21 heads, 704/704/640 expert rows)" { + const c = Config{}; + const r1 = try Share.of(&c, 1, 2); + try std.testing.expectEqual(Span{ .lo = 32, .hi = 64 }, r1.kda); + try std.testing.expectEqual(Span{ .lo = 6144, .hi = 12288 }, r1.dense); + try std.testing.expectEqual(Span{ .lo = 1024, .hi = 2048 }, r1.moe); + try std.testing.expectEqual(Span{ .lo = 1024, .hi = 2048 }, r1.shared); + try std.testing.expectEqual(Span{ .lo = 77440, .hi = 154880 }, r1.vocab); + const t0 = try Share.of(&c, 0, 3); + const t2 = try Share.of(&c, 2, 3); + try std.testing.expectEqual(@as(usize, 22), t0.mla.len()); + try std.testing.expectEqual(Span{ .lo = 43, .hi = 64 }, t2.kda); + try std.testing.expectEqual(@as(usize, 704), t0.moe.len()); + try std.testing.expectEqual(@as(usize, 640), t2.moe.len()); + try std.testing.expectEqual(Span{ .lo = 103296, .hi = 154880 }, t2.vocab); + try std.testing.expectError(error.BadRank, Share.of(&c, 2, 2)); +} + +test "an EXL3 checkpoint cuts the vocabulary at 128s: TP=4 takes 303, 303, 302 and 302 blocks" { + var c = Config{}; + c.head_unit = 128; + try std.testing.expectEqual(Span{ .lo = 0, .hi = 38784 }, (try Share.of(&c, 0, 4)).vocab); + try std.testing.expectEqual(Span{ .lo = 116224, .hi = 154880 }, (try Share.of(&c, 3, 4)).vocab); + try std.testing.expectEqual(Span{ .lo = 77440, .hi = 154880 }, (try Share.of(&c, 1, 2)).vocab); +} diff --git a/zig/src/families/glm/cuda_shared.zig b/zig/src/families/glm/cuda_shared.zig new file mode 100644 index 0000000000..86b19c640d --- /dev/null +++ b/zig/src/families/glm/cuda_shared.zig @@ -0,0 +1,70 @@ +//! GLM-5.3-Flash's shared verify windows: several streams' rows in one forward, chains per segment. + +const std = @import("std"); +const cuda = @import("cuda"); +const fw = @import("cuda_forward.zig"); +const weights = @import("cuda_weights.zig"); +const kda = @import("cuda_kda.zig"); +const latent = @import("cuda_latent.zig"); +const kda_block = @import("cuda_kda_block.zig"); +const dsa_block = @import("cuda_dsa_block.zig"); + +/// Checks a shared window's segments against its `rows` rows, then each sequence holds its rows until its commit. +pub fn begin(r: fw.Rank, rows: usize) !void { + if (r.b.prefill) return error.BadWindow; + var at: usize = 0; + for (r.segs) |g| { + if (g.row0 != at or g.rows == 0 or g.rows > fw.window_rows) return error.BadWindow; + if (g.st.pos + g.rows > g.st.capacity) return error.PastCapacity; + at += g.rows; + } + if (at != rows) return error.BadWindow; + for (r.segs) |g| { + g.st.pending = g.rows; + g.st.prompt = false; + } +} + +/// Forward.mixer on a shared window: all rows' projections, each segment's chain or attention alone, all o_proj. +pub fn mixer(f: *const fw.Forward, r: fw.Rank, l: weights.Layer, li: usize, di: usize, rows: usize, out: u64) !void { + const c = &r.w.config; + const b = r.b; + switch (l.mixer) { + .kda => |k| { + const n = k.heads * kda.dk; + const w = k.width(); + var s = b.kda; + s.p = b.kproj; + try kda_block.project(f.set(), f.stream, k, s, b.normed, c.hidden, rows); + for (r.segs) |g| { + const st = g.st; + const cur = st.cur[li]; + var v = s; + v.p = st.proj + li * st.proj_layer * 2; // the window's rows stay with the layer until its commit + try f.d.check(f.d.api.cuMemcpyDtoDAsync_v2(v.p, b.kproj + g.row0 * w * 2, g.rows * w * 2, f.stream.handle), "cuMemcpyDtoDAsync"); + v.ka += g.row0 * n * 2; + v.kg += g.row0 * n * 2; + v.saves.out += g.row0 * n * 2; + v.saves.state_out = st.recAt(1 - cur, li); + const kv = li * fw.window_rows * st.heads * kda.dk; + v.saves.k = st.save_k + kv * 4; + v.saves.v = st.save_v + kv * 2; + v.saves.g = st.save_g + kv * 4; + v.saves.b = st.save_b + li * fw.window_rows * st.heads * 4; + try kda_block.chainRows(&f.m.kda, f.stream, k, v, g.rows, st.conv + li * 3 * st.conv_c * 2, st.recAt(cur, li), f.kda_wide); + } + try kda_block.fill(f.set(), f.stream, k, s, 0, rows, out); + }, + .mla => |a| { + var sparse_any = false; + for (r.segs) |g| sparse_any = sparse_any or dsa_block.sparseRows(a, g.st.pos, g.rows); + try dsa_block.prepare(f.set(), f.stream, a, b.dsa, b.normed, c.hidden, rows, sparse_any); + for (r.segs) |g| { + const nch = @min(latent.chunksFor(g.st.pos + g.rows), b.nch); + try dsa_block.attend(.{ .set = f.set(), .ops = f.m.ops() }, f.stream, a, g.st.caches[di], b.dsa, g.row0, g.rows, g.st.pos_dev, g.st.pos, nch, b.part_rows); + } + try latent.expandV(f.set(), f.stream, b.dsa.ol, a.wv, b.dsa.vn, rows, a.heads, a.dv, a.lw); + try dsa_block.fill(f.set(), f.stream, a, b.dsa, 0, rows, out); + }, + } +} diff --git a/zig/src/families/glm/cuda_sparse.zig b/zig/src/families/glm/cuda_sparse.zig new file mode 100644 index 0000000000..c72a3c125d --- /dev/null +++ b/zig/src/families/glm/cuda_sparse.zig @@ -0,0 +1,95 @@ +//! GLM-5.3-Flash's DSA indexer on CUDA (the Python family's sparse): each row keeps its best 512 4-token pools. + +const std = @import("std"); +const cuda = @import("cuda"); +const glue = @import("cuda_glue.zig"); + +/// sparse.POOL, TOPK_POOLS, SELECT_LIST, SCORE_LOOP_FROM, SCORE_LOOP; geometry.MLA_SELECT_ROWS. +pub const pool = 4; +pub const topk_pools = 512; +pub const select_list = 1024; +pub const score_loop_from = 64; +pub const score_loop = 32; +pub const select_rows = 512; +/// A row's token list: the 512 pools' tokens, then the incomplete pool's up to 3. +pub const width = topk_pools * pool + pool - 1; + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// sparse.index_update: the window's LayerNormed index keys and gates into ik and ig at pos, completed pools into pk. +pub fn indexUpdate(set: glue.Set, stream: cuda.Stream, k_raw: u64, k_stride: usize, gate: u64, ln_w: u64, ln_b: u64, ape: u64, ik: u64, ig: u64, pk: u64, pos: u64, rows: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ k_raw, @as(i64, @intCast(k_stride)), gate, ln_w, ln_b, ik, ig, pos, @as(f32, 1e-6) }) |v| a.add(v); + try cuda.launch.launch(set.glue.index_write, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 32 } }, stream, &a); + var b: cuda.Args = .{}; + inline for (.{ ik, ig, ape, pk, pos, @as(i64, @intCast(rows)) }) |v| b.add(v); + try cuda.launch.launch(set.glue.pool_keys, .{ .grid = .{ .x = @intCast(rows / 4 + 2) }, .block = .{ .x = 32 } }, stream, &b); +} + +/// sparse.pool_bucket: visible pools for rows pos .. pos + rows - 1, a power of two (at least 1024), at most `np_max`. +pub fn poolBucket(pos: usize, rows: usize, np_max: usize) usize { + const visible = (pos + rows) / pool + 1; + return @min(np_max, @max(1024, std.math.ceilPowerOfTwoAssert(usize, visible))); +} + +/// sparse.score_columns: the pool columns selection reads for those rows (the last row's complete pools, at least 512). +pub fn scoreColumns(pos: usize, rows: usize, np_max: usize) usize { + return @min(np_max, @max(topk_pools, (pos + rows) / pool)); +} + +/// select_tokens' 1/sqrt(heads) head-weight scale (32 heads: 1 / 5.656854249492381, as the Python family writes it). +pub fn headScale(heads: usize) f32 { + if (heads == 32) return @floatCast(1.0 / 5.656854249492381); + return @floatCast(1.0 / @sqrt(@as(f64, @floatFromInt(heads)))); +} + +/// The scratch of one selectTokens: scores fp32 [min(rows, 512)][np], and with np past select_list the list rows. +pub const Scratch = struct { scores: u64, ls: u64, li: u64, pools: u64, at: u64 }; + +/// sparse.select_tokens: each row's tokens (i32 [rows][width], ascending, -1 padded), count 0 under the dense limit. +pub fn selectTokens(set: glue.Set, ops: cuda.Module, stream: cuda.Stream, qi: u64, wts: u64, w_stride: usize, pk: u64, pos: u64, host_pos: usize, rows: usize, heads: usize, d: usize, np_max_cap: usize, s: Scratch, tokens: u64, counts: u64) !void { + const np_max = poolBucket(host_pos, rows, np_max_cap); + const b = @min(rows, select_rows); + const loop: usize = if (rows >= score_loop_from) score_loop else 1; + const add = try ops.function("tf_glm_add_i32"); + var a: usize = 0; + while (a < rows) : (a += b) { + const n = @min(b, rows - a); + var at = pos; + if (a > 0) { + var args: cuda.Args = .{}; + args.add(pos); + args.add(@as(c_int, @intCast(a))); + args.add(s.at); + try cuda.launch.launch(add, .{ .grid = .{}, .block = .{ .x = 32 } }, stream, &args); + at = s.at; + } + const nb = cdiv(scoreColumns(host_pos + a, n, np_max), 64); + var sa: cuda.Args = .{}; + inline for (.{ qi + a * heads * d * 2, wts + a * w_stride * 2, @as(i64, @intCast(w_stride)), pk, s.scores, at, @as(i64, @intCast(np_max)), @as(i64, @intCast(nb)), @as(f32, @floatCast(std.math.pow(f64, @floatFromInt(d), -0.5))), headScale(heads), @as(i64, @intCast(heads)), @as(c_int, @intCast(loop)) }) |v| sa.add(v); + if (d != 128 or heads > 32) return error.IndexShape; + try cuda.launch.launch(set.glue.scores, .{ .grid = .{ .x = @intCast(n), .y = @intCast(cdiv(nb, loop)) }, .block = .{ .x = 128 } }, stream, &sa); + var sel: cuda.Args = .{}; + inline for (.{ s.scores, s.pools + a * topk_pools * 8, @as(i64, @intCast(np_max)), at, @as(c_int, topk_pools) }) |v| sel.add(v); + try cuda.launch.launch(set.glue.select_rows, .{ .grid = .{ .x = @intCast(n) }, .block = .{ .x = 128 } }, stream, &sel); + } + const assemble = try ops.function("tf_glm_dsa_tokens"); + var args: cuda.Args = .{}; + args.add(s.pools); + args.add(pos); + args.add(@as(c_int, topk_pools)); + args.add(tokens); + args.add(counts); + try cuda.launch.launch(assemble, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, stream, &args); +} + +test "sparse.py's sizing rules" { + try std.testing.expectEqual(@as(usize, 1024), poolBucket(100, 1, 1 << 20)); + try std.testing.expectEqual(@as(usize, 2048), poolBucket(5000, 1, 1 << 20)); // 1251 visible pools + try std.testing.expectEqual(@as(usize, 512), scoreColumns(100, 1, 1024)); + try std.testing.expectEqual(@as(usize, 1250), scoreColumns(5000, 1, 2048)); + try std.testing.expectEqual(@as(usize, 2051), width); + try std.testing.expectEqual(@as(f32, @floatCast(1.0 / 5.656854249492381)), headScale(32)); +} diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig new file mode 100644 index 0000000000..32559505c3 --- /dev/null +++ b/zig/src/families/glm/cuda_state.zig @@ -0,0 +1,325 @@ +//! GLM-5.3-Flash's per-rank CUDA memory: forward Buffers (shared by sequences) and a sequence's State, an arena each. + +const std = @import("std"); +const cuda = @import("cuda"); +const weights = @import("cuda_weights.zig"); +const kda = @import("cuda_kda.zig"); +const b16 = @import("cuda_b16.zig"); +const latent = @import("cuda_latent.zig"); +const sparse = @import("cuda_sparse.zig"); +const experts = @import("cuda_experts.zig"); +const dsa_block = @import("cuda_dsa_block.zig"); +const moe = @import("cuda_moe.zig"); +const glue = @import("cuda_glue.zig"); +const mlp = @import("cuda_mlp.zig"); +const kda_block = @import("cuda_kda_block.zig"); + +/// A prompt chunk's rows (decode.PREFILL_ROWS). +pub const prefill_rows = 2048; +/// A decode window's rows (Engine's max_rows). +pub const window_rows = 8; +/// The context the buffers' dense attention partials reach (Buffers' min(capacity, 2560)). +pub const dense_reach = 2560; + +/// Sub-allocations of one device allocation, 256-byte aligned like the driver's own. +const Arena = struct { + buf: cuda.DeviceBuffer, + used: usize = 0, + + fn take(a: *Arena, bytes: usize) u64 { + const at = std.mem.alignForward(usize, a.used, 256); + a.used = at + bytes; + std.debug.assert(a.used <= a.buf.len); + return a.buf.ptr + at; + } +}; + +/// Allocates `sizes` (bytes) in one zeroed device buffer and hands out each field's address in order. +fn arena(d: *const cuda.Driver, sizes: []const usize, out: []u64) !Arena { + var total: usize = 0; + for (sizes) |n| total += std.mem.alignForward(usize, n, 256); + var a: Arena = .{ .buf = try cuda.DeviceBuffer.alloc(d, @max(total, 256)) }; + errdefer a.buf.free(); + try a.buf.fill8(0, null); + for (sizes, out) |n, *p| p.* = a.take(n); + return a; +} + +/// A rank's B16 matmul shapes ([n, k]) whose split-K partials the buffers hold. +fn b16Shapes(w: *const weights.Weights) [10][2]usize { + const c = &w.config; + const s = &w.share; + const n = s.kda.len() * c.kda_dim; + const kw = 3 * n + 2 * c.kda_dim + s.kda.len(); + return .{ + .{ kw, c.hidden }, .{ n, c.kda_dim }, .{ c.hidden, n }, + .{ c.q_lora + c.kv_lora, c.hidden }, .{ s.mla.len() * c.nope, c.q_lora }, .{ c.i_dim + c.i_heads, c.hidden }, + .{ c.i_heads * c.i_dim, c.q_lora }, .{ c.hidden, s.mla.len() * c.v_dim }, .{ 2 * s.shared.len(), c.hidden }, + .{ c.hidden, 2 * c.hidden }, // the MTP head's eh_proj + }; +} + +/// Scratch for windows of up to `rows` rows (the Python family's Buffers): a prompt buffer (`prefill`) or a decode's. +pub const Buffers = struct { + arena: Arena, + rows: usize, + prefill: bool, + part_rows: usize, // rows one dense attention call takes + nch: usize, // the dense pass's most 512-key chunks + ids: u64, + x: u64, + normed: u64, + xs: u64, + post: u64, + comb: u64, + hcpart: u64, + gathered: u64, // fp32 [world][rows][dims]: every rank's partial in rank order + kproj: u64, // a prompt chunk's or a shared window's KDA projection rows (a window's stay in the state, for its commit) + kda: kda_block.Scratch, + dsa: dsa_block.Scratch, + dense: mlp.Scratch, + moe: moe.Scratch, + hidden: u64, + fnormed: u64, + fxs: u64, + logits: u64, // bf16 [head rows][this rank's vocabulary] + // the MTP head (every Buffers holds them, as the Python family's) + hin: u64, // bf16 [rows][dims]: the hidden rows it reads + me: u64, // bf16 [rows][dims]: the next tokens' embedding + mcat: u64, // bf16 [rows][2 dims]: [enorm(me) | hnorm(hin)] + mx: u64, // bf16 [rows][dims]: its stream + part: u64, // fp32 [rows][dims]: this rank's partial rows a split exchange sends from + + pub fn init(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize) !Buffers { + return make(d, w, rows, prefill, capacity, rows); + } + + /// A shared window's scratch: `rows` rows of several streams, each one's attention and KDA chain on `window_rows`. + pub fn initShared(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, capacity: usize) !Buffers { + return make(d, w, rows, false, capacity, window_rows); + } + + /// `seg`: the most rows one stream's attention takes (a decode window's rows, or a shared window's segment). + fn make(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize, seg: usize) !Buffers { + const c = &w.config; + const s = &w.share; + const D: usize = c.hidden; + const S: usize = c.hc; + const LL = s.kda.len(); + const HL = s.mla.len(); + const n = LL * c.kda_dim; + const kw = 3 * n + 2 * c.kda_dim + LL; + const mw = s.dense.len(); + const ew = s.moe.len(); + const sh = s.shared.len(); + const slots: usize = c.topk + 1; + const pairs = rows * slots; + const count: usize = c.experts; + const lw: usize = c.kv_lora; + const part_rows = if (prefill) @min(rows, dsa_block.prompt_att_rows) else seg; + const nch = latent.chunksFor(@min(capacity, dense_reach) + seg); + const partial_rows = @max(nch * part_rows, latent.chunksFor(sparse.width) * seg); // dense calls, or one sparse + const np_max = capacity / sparse.pool; + var mm: usize = 4; + // windows under rows_from rows keep their K slices apart (a prompt's tail chunk on the prefill buffers too) + for (b16Shapes(w)) |sh_| mm = @max(mm, b16.partBytes(@min(rows, b16.rows_from - 1), sh_[0], sh_[1])); + const hist_len = if (pairs > cuda.grouped.small) (pairs + cuda.grouped.small - 1) / cuda.grouped.small * (count + 1) else 1; + const sizes = [_]usize{ + rows * 4, rows * S * D * 2, rows * D * 2, rows * (D / 64) * 4, + rows * S * 4, rows * S * S * 4, glue.hcPartBytes(rows), s.world * rows * D * 4, + if (prefill or seg < rows) rows * kw * 2 else 4, rows * n * 2, rows * n * 2, // kproj, ka, kg + rows * n * 2, rows * n * 4, rows * n * 2, rows * n * 4, // saves out, k, v, g + rows * LL * 4, rows * n * 4, rows * n * 2, // saves b, wide q, y + rows * (c.q_lora + lw) * 2, rows * c.q_lora * 2, rows * (c.q_lora / 64) * 4, + rows * lw * 2, rows * (lw / 64) * 4, rows * HL * c.nope * 2, + rows * (c.i_dim + c.i_heads) * 2, rows * c.i_dim * 4, glue.router_ks * rows * @max(count, c.i_dim) * 4, + rows * HL * lw * 2, rows * HL * lw * 2, rows * HL * c.v_dim * 2, + rows * c.i_heads * c.i_dim * 2, partial_rows * HL * lw * 4, partial_rows * HL * 4, + partial_rows * HL * 4, + 4, // dummy + @min(seg, sparse.select_rows) * np_max * 4, + rows * sparse.select_list * 4, + rows * sparse.select_list * 4, + rows * sparse.topk_pools * 8, + 4, + rows * sparse.width * 4, + rows * 4, + mm, + 4, // a prompt block's position + rows * mw * 2, + rows * 2 * mw * 2, + rows * mw * 2, + rows * (mw / 64) * 4, + rows * count * 4, pairs * 4, pairs * 4, pairs * 4, // mlog, pick, wts, members + experts.maxItems(pairs, count + 1, experts.tile) * 3 * 4, // items (the decode tile's count bounds the prompt's) + 8, @max(1, pairs) * 4, hist_len * 4, // counts, rank, hist + pairs * ew * 2, pairs * D * 4, rows * 2 * sh * 2, + rows * sh * 2, + rows * (sh / 64) * 4, rows * D * 4, // sxs, sy + rows * D * 2, rows * D * 2, + rows * (D / 64) * 4, (if (prefill) 1 else rows) * s.vocab.len() * 2, + rows * D * 2, rows * D * 2, rows * 2 * D * 2, rows * D * 2, // hin, me, mcat, mx + rows * D * 4, // part: this rank's fp32 partial rows a split exchange sends from + }; + var p: [sizes.len]u64 = undefined; + const a = try arena(d, &sizes, &p); + return .{ + .arena = a, + .rows = rows, + .prefill = prefill, + .part_rows = part_rows, + .nch = nch, + .ids = p[0], + .x = p[1], + .normed = p[2], + .xs = p[3], + .post = p[4], + .comb = p[5], + .hcpart = p[6], + .gathered = p[7], + .kproj = p[8], + .kda = .{ .p = 0, .ka = p[9], .kg = p[10], .saves = .{ .out = p[11], .state_out = 0, .k = p[12], .v = p[13], .g = p[14], .b = p[15] }, .wide = .{ .q = p[16], .y = p[17] }, .mm_part = p[42] }, + .dsa = .{ + .dp = p[18], + .qr = p[19], + .xs_qr = p[20], + .lat = p[21], + .xs_lat = p[22], + .q = p[23], + .ikr = p[24], + .igr = p[25], + .router_part = p[26], + .qa = p[27], + .ol = p[28], + .vn = p[29], + .qi = p[30], + .partials = .{ .po = p[31], .pm = p[32], .pl = p[33] }, + .dummy = p[34], + .select = .{ .scores = p[35], .ls = p[36], .li = p[37], .pools = p[38], .at = p[39] }, + .tokens = p[40], + .counts = p[41], + .mm_part = p[42], + .at = p[43], + }, + .dense = .{ .tmp = p[44], .gu = p[45], .act = p[46], .xs_act = p[47] }, + .moe = .{ .mlog = p[48], .router_part = p[26], .pick = p[49], .wts = p[50], .plan = .{ .members = p[51], .items = p[52], .counts = p[53], .rank = p[54], .hist = p[55] }, .eact = p[56], .ey = p[57], .sgu = p[58], .sact = p[59], .sxs = p[60], .sy = p[61], .mm_part = p[42] }, + .hidden = p[62], + .fnormed = p[63], + .fxs = p[64], + .logits = p[65], + .hin = p[66], + .me = p[67], + .mcat = p[68], + .mx = p[69], + .part = p[70], + }; + } + + pub fn deinit(b: *Buffers) void { + b.arena.buf.free(); + b.* = undefined; + } +}; + +/// A sequence's committed caches on a rank (the Python family's State): KDA conv and states, saved rows, DSA caches. +pub const State = struct { + arena: Arena, + gpa: std.mem.Allocator, + capacity: usize, + pos: usize = 0, + pending: usize = 0, // rows of the last forward, until its commit + prompt: bool = false, // whether that forward was a prompt chunk + pos_dev: u64, // i32 + conv: u64, // bf16 [kda layers][3][3 heads 128] + rec: u64, // fp32 [2][kda layers][heads][128][128] + proj: u64, // bf16 [kda layers][window rows][kda width] + save_k: u64, // fp32 [kda layers][window rows][heads][128] + save_v: u64, // bf16, as save_k + save_g: u64, // fp32, as save_k + save_b: u64, // fp32 [kda layers][window rows][heads] + cur: []u1, + kda_layers: usize, + heads: usize, // this rank's KDA heads + caches: []dsa_block.Caches, + mtp_len: usize = 0, // rows in the MTP layer's caches + mtp_pos_dev: u64, // i32 + mtp_caches: dsa_block.Caches, // the MTP layer's latent rows, index keys, gates and pooled keys (the Python family's index[-1]) + conv_c: usize, + rec_layer: usize, + proj_layer: usize, + proj_row: usize, + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, w: *const weights.Weights, capacity: usize) !State { + const c = &w.config; + const s = &w.share; + var nk: usize = 0; + var nd: usize = 0; + for (w.layers) |l| switch (l.mixer) { + .kda => nk += 1, + .mla => nd += 1, + }; + const n = s.kda.len() * c.kda_dim; + const kw = 3 * n + 2 * c.kda_dim + s.kda.len(); + const pk_rows = capacity / sparse.pool + 2; + var sizes: std.ArrayList(usize) = .empty; + defer sizes.deinit(gpa); + try sizes.appendSlice(gpa, &.{ 4, nk * 3 * 3 * n * 2, 2 * nk * s.kda.len() * kda.dk * kda.dv * 4, nk * window_rows * kw * 2 }); + try sizes.appendSlice(gpa, &.{ nk * window_rows * n * 4, nk * window_rows * n * 2, nk * window_rows * n * 4, nk * window_rows * s.kda.len() * 4 }); + for (0..nd + 1) |_| try sizes.appendSlice(gpa, &.{ capacity * c.kv_lora * 2, capacity * c.i_dim * 2, capacity * c.i_dim * 2, pk_rows * c.i_dim * 2 }); + try sizes.append(gpa, 4); // the MTP head's position + const p = try gpa.alloc(u64, sizes.items.len); + defer gpa.free(p); + var a = try arena(d, sizes.items, p); + errdefer a.buf.free(); + const cur = try gpa.alloc(u1, nk); + errdefer gpa.free(cur); + @memset(cur, 0); + const caches = try gpa.alloc(dsa_block.Caches, nd); + for (caches, 0..) |*x, j| x.* = .{ .lc = p[8 + 4 * j], .ik = p[9 + 4 * j], .ig = p[10 + 4 * j], .pk = p[11 + 4 * j], .pk_rows = pk_rows }; + return .{ .arena = a, .gpa = gpa, .capacity = capacity, .pos_dev = p[0], .conv = p[1], .rec = p[2], .proj = p[3], .save_k = p[4], .save_v = p[5], .save_g = p[6], .save_b = p[7], .cur = cur, .heads = s.kda.len(), .mtp_caches = .{ .lc = p[8 + 4 * nd], .ik = p[9 + 4 * nd], .ig = p[10 + 4 * nd], .pk = p[11 + 4 * nd], .pk_rows = pk_rows }, .mtp_pos_dev = p[12 + 4 * nd], .kda_layers = nk, .caches = caches, .conv_c = 3 * n, .rec_layer = s.kda.len() * kda.dk * kda.dv * 4, .proj_layer = window_rows * kw, .proj_row = kw }; + } + + pub fn deinit(st: *State) void { + st.arena.buf.free(); + st.gpa.free(st.cur); + st.gpa.free(st.caches); + st.* = undefined; + } + + /// State.reset: a fresh sequence (zero conv windows and states, position 0); the caches past it are never read. + pub fn reset(st: *State, stream: cuda.Stream) !void { + const d = st.arena.buf.d; + try d.check(d.api.cuMemsetD8Async(st.conv, 0, st.kda_layers * 3 * st.conv_c * 2, stream.handle), "cuMemsetD8Async"); + try d.check(d.api.cuMemsetD8Async(st.rec, 0, 2 * st.kda_layers * st.rec_layer, stream.handle), "cuMemsetD8Async"); + @memset(st.cur, 0); + st.pending = 0; + try st.setPos(stream, 0); + try st.setMtpLen(stream, 0); + } + + /// Host side of a state just copied back (a kept prompt state): position, MTP rows, KDA parity, nothing pending. + pub fn restoreHost(st: *State, stream: cuda.Stream, pos: usize, mtp_len: usize, cur: []const u1) !void { + @memcpy(st.cur, cur); + st.pending = 0; + st.prompt = false; + try st.setPos(stream, pos); + try st.setMtpLen(stream, mtp_len); + } + + /// State.set_mtp_len: rows in the MTP layer's caches (the next MTP step writes from here). + pub fn setMtpLen(st: *State, stream: cuda.Stream, n: usize) !void { + st.mtp_len = n; + const v: i32 = @intCast(n); + try st.arena.buf.d.check(st.arena.buf.d.api.cuMemsetD32Async(st.mtp_pos_dev, @bitCast(v), 1, stream.handle), "cuMemsetD32Async"); + } + + pub fn setPos(st: *State, stream: cuda.Stream, pos: usize) !void { + st.pos = pos; + const v: i32 = @intCast(pos); + try st.arena.buf.d.check(st.arena.buf.d.api.cuMemsetD32Async(st.pos_dev, @bitCast(v), 1, stream.handle), "cuMemsetD32Async"); + } + + pub fn recAt(st: *const State, copy: u1, li: usize) u64 { + return st.rec + (@as(usize, copy) * st.kda_layers + li) * st.rec_layer; + } +}; diff --git a/zig/src/families/glm/cuda_weights.zig b/zig/src/families/glm/cuda_weights.zig new file mode 100644 index 0000000000..80bc44a8e0 --- /dev/null +++ b/zig/src/families/glm/cuda_weights.zig @@ -0,0 +1,499 @@ +//! A TP rank's GLM-5.3-Flash weights on CUDA (the Python family's weights.load, NVFP4), named by its paths for digests. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const cfg = @import("config.zig"); +const shard = @import("cuda_shard.zig"); +const X = cuda.experts.Nvfp4Experts; +const kda_block = @import("cuda_kda_block.zig"); +const dsa_block = @import("cuda_dsa_block.zig"); +const mlp = @import("cuda_mlp.zig"); +const moe = @import("cuda_moe.zig"); +const quant = @import("cuda_quant.zig"); +const q4 = @import("cuda_q4.zig"); + +const Tensor = core.checkpoint.Tensor; +const prefix = "model.language_model."; + +/// A hyper-connection site: fn bf16 [s (s + 2)][s d], base and scale widened to fp32. +pub const Hc = struct { fn_w: u64, base: u64, scale: u64 }; + +pub const Mixer = union(cfg.Kind) { kda: kda_block.Layer, mla: dsa_block.Layer }; +pub const Ffn = union(enum) { dense: mlp.Layer, moe: moe.Layer }; + +pub const Layer = struct { index: u32, attn_hc: Hc, ffn_hc: Hc, in_norm: u64, post_norm: u64, mixer: Mixer, ffn: Ffn }; + +/// The MTP head (weights.MTPW): norms, eh_proj, a DSA + MoE layer (NVFP4 experts), shared_head.norm, 4-bit draft head. +pub const Mtp = struct { enorm: u64, hnorm: u64, eh: u64, norm: u64, layer: Layer, draft: q4.Q4 }; + +/// The kernels a load runs: affine4_pack's tf_pack_dense, and for the MTP head the quantizers and the 4-bit tiling. +pub const Tools = struct { pack_dense: cuda.Function, quant: ?*const quant.Kernels = null, q4: ?*const q4.Kernels = null }; + +/// A loaded tensor by the Python family's attribute path, for digests against its weights. +pub const Named = struct { name: []u8, ptr: u64, len: usize }; + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: cfg.Config, + share: shard.Share, + embed: u64 = 0, // bf16 [vocab][dims] + layers: []Layer = &.{}, + norm: u64 = 0, + head: u64 = 0, // bf16 [share.vocab][dims] + mtp: ?Mtp = null, + buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + named: std.ArrayList(Named) = .empty, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + for (w.named.items) |n| w.gpa.free(n.name); + w.buffers.deinit(w.gpa); + w.named.deinit(w.gpa); + w.gpa.free(w.layers); + w.* = undefined; + } + + /// The upload named `name` (the Python family's attribute path), or null. + pub fn find(w: *const Weights, name: []const u8) ?Named { + for (w.named.items) |n| if (std.mem.eql(u8, n.name, name)) return n; + return null; + } +}; + +/// Host bf16 from fp32: round to nearest even (torch's float -> bfloat16), NaN kept quiet. +fn toBf16(v: f32) u16 { + const b: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((b >> 16) | 0x40); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +const Loader = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + stream: cuda.Stream, + pack_dense: cuda.Function, + tools: Tools, + w: *Weights, + ck: *core.Checkpoint, + host: std.ArrayList(u8) = .empty, + name_buf: [192]u8 = undefined, + + fn get(L: *Loader, comptime fmt: []const u8, args: anytype) !Tensor { + const name = try std.fmt.bufPrint(&L.name_buf, prefix ++ fmt, args); + return L.ck.get(name); + } + + /// The host bytes gathered so far uploaded as one tensor named `fmt`. + fn put(L: *Loader, comptime fmt: []const u8, args: anytype) !u64 { + const b = try cuda.DeviceBuffer.fromHost(L.d, L.host.items); + try L.w.buffers.append(L.gpa, b); + try L.w.named.append(L.gpa, .{ .name = try std.fmt.allocPrint(L.gpa, fmt, args), .ptr = b.ptr, .len = b.len }); + L.host.clearRetainingCapacity(); + return b.ptr; + } + + /// Elements [lo, hi) of `t` flat as `to` (bf16 or fp32 from a bf16 or fp32 tensor; other dtypes only as stored). + fn add(L: *Loader, t: Tensor, lo: usize, hi: usize, to: core.checkpoint.DType) !void { + if (t.dtype == to) { + const size = t.bytes.len / t.numel(); + return L.host.appendSlice(L.gpa, t.bytes[lo * size .. hi * size]); + } + for (lo..hi) |i| { + const v: f32 = switch (t.dtype) { + .bf16 => @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)) << 16), + .f32 => std.mem.bytesToValue(f32, t.bytes[4 * i ..][0..4]), + else => return error.UnexpectedTensor, + }; + switch (to) { + .f32 => try L.host.appendSlice(L.gpa, std.mem.asBytes(&v)), + .bf16 => try L.host.appendSlice(L.gpa, std.mem.asBytes(&toBf16(v))), + else => return error.UnexpectedTensor, + } + } + } + + /// Rows [lo, hi) of `t`'s outermost axis (`total` long; a vector's elements): split.py's "row" cut. + fn rows(L: *Loader, t: Tensor, lo: usize, hi: usize, total: usize, to: core.checkpoint.DType) !void { + const per = t.numel() / total; + if (per * total != t.numel()) return error.UnexpectedTensor; + try L.add(t, lo * per, hi * per, to); + } + + fn whole(L: *Loader, t: Tensor, to: core.checkpoint.DType) !void { + try L.add(t, 0, t.numel(), to); + } + + /// Columns [lo, hi) of every row of a matrix: split.py's "col" cut. + fn cols(L: *Loader, t: Tensor, lo: usize, hi: usize, to: core.checkpoint.DType) !void { + if (t.rank != 2) return error.UnexpectedTensor; + const k = t.dim(1); + for (0..t.dim(0)) |r| try L.add(t, r * k + lo, r * k + hi, to); + } + + fn hc(L: *Loader, i: u32, comptime site: []const u8) !Hc { + try L.whole(try L.get("layers.{d}.hc_" ++ site ++ "_fn", .{i}), .bf16); + const f = try L.put("layers.{d}.{s}_hc.fn", .{ i, site }); + try L.whole(try L.get("layers.{d}.hc_" ++ site ++ "_base", .{i}), .f32); + const base = try L.put("layers.{d}.{s}_hc.base", .{ i, site }); + try L.whole(try L.get("layers.{d}.hc_" ++ site ++ "_scale", .{i}), .f32); + return .{ .fn_w = f, .base = base, .scale = try L.put("layers.{d}.{s}_hc.scale", .{ i, site }) }; + } + + /// weights.load's kda(i): [q | k | v | f_a | g_a | b], f_b/g_b by heads, o by columns, conv taps, A_log, dt_bias. + fn kda(L: *Loader, i: u32) !kda_block.Layer { + const c = &L.w.config; + const h = L.w.share.kda; + const dk = c.kda_dim; + const all = c.kda_heads * dk; + inline for (.{ "q", "k", "v" }) |x| try L.rows(try L.get("layers.{d}.self_attn." ++ x ++ "_proj.weight", .{i}), h.lo * dk, h.hi * dk, all, .bf16); + inline for (.{ "f_a", "g_a" }) |x| try L.whole(try L.get("layers.{d}.self_attn." ++ x ++ "_proj.weight", .{i}), .bf16); + try L.rows(try L.get("layers.{d}.self_attn.b_proj.weight", .{i}), h.lo, h.hi, c.kda_heads, .bf16); + const proj = try L.put("layers.{d}.kda.proj", .{i}); + try L.rows(try L.get("layers.{d}.self_attn.f_b_proj.weight", .{i}), h.lo * dk, h.hi * dk, all, .bf16); + const fb = try L.put("layers.{d}.kda.fb", .{i}); + try L.rows(try L.get("layers.{d}.self_attn.g_b_proj.weight", .{i}), h.lo * dk, h.hi * dk, all, .bf16); + const gb = try L.put("layers.{d}.kda.gb", .{i}); + inline for (.{ "q", "k", "v" }) |x| try L.rows(try L.get("layers.{d}.self_attn." ++ x ++ "_conv1d.weight", .{i}), h.lo * dk, h.hi * dk, all, .f32); + const conv = try L.put("layers.{d}.kda.conv", .{i}); + try L.rows(try L.get("layers.{d}.self_attn.A_log", .{i}), h.lo, h.hi, c.kda_heads, .f32); + const a_log = try L.put("layers.{d}.kda.a_log", .{i}); + try L.rows(try L.get("layers.{d}.self_attn.dt_bias", .{i}), h.lo * dk, h.hi * dk, all, .f32); + const dt = try L.put("layers.{d}.kda.dt_bias", .{i}); + try L.whole(try L.get("layers.{d}.self_attn.o_norm.weight", .{i}), .bf16); + const norm = try L.put("layers.{d}.kda.norm", .{i}); + try L.cols(try L.get("layers.{d}.self_attn.o_proj.weight", .{i}), h.lo * dk, h.hi * dk, .bf16); + const o = try L.put("layers.{d}.kda.o", .{i}); + return .{ .proj = proj, .fb = fb, .gb = gb, .o = o, .conv_w = conv, .a_log = a_log, .dt_bias = dt, .norm_w = norm, .eps = c.eps, .lower = c.lower_bound, .heads = h.len(), .dims = c.hidden }; + } + + /// weights.load's dsa(i) (latent): [q_a | kv_a], norms, q_b by heads, kv_b as wk/wv (AbsorbW), o by cols, indexer. + fn dsa(L: *Loader, i: u32) !dsa_block.Layer { + const c = &L.w.config; + const h = L.w.share.mla; + const qk = c.nope; + inline for (.{ "q_a_proj", "kv_a_proj_with_mqa" }) |x| try L.whole(try L.get("layers.{d}.self_attn." ++ x ++ ".weight", .{i}), .bf16); + const proj = try L.put("layers.{d}.dsa.proj", .{i}); + try L.whole(try L.get("layers.{d}.self_attn.q_a_layernorm.weight", .{i}), .bf16); + const q_norm = try L.put("layers.{d}.dsa.q_norm", .{i}); + try L.whole(try L.get("layers.{d}.self_attn.kv_a_layernorm.weight", .{i}), .bf16); + const kv_norm = try L.put("layers.{d}.dsa.kv_norm", .{i}); + try L.rows(try L.get("layers.{d}.self_attn.q_b_proj.weight", .{i}), h.lo * qk, h.hi * qk, c.mla_heads * qk, .bf16); + const q_b = try L.put("layers.{d}.dsa.q_b", .{i}); + const kv_b = try L.get("layers.{d}.self_attn.kv_b_proj.weight", .{i}); + const per = qk + c.v_dim; // a head's kv_b rows: its keys, then its values + const total = c.mla_heads * per; + for (h.lo..h.hi) |head| try L.rows(kv_b, head * per, head * per + qk, total, .bf16); + const wk = try L.put("layers.{d}.dsa.absorb.wk", .{i}); + for (h.lo..h.hi) |head| try L.rows(kv_b, head * per + qk, (head + 1) * per, total, .bf16); + const wv = try L.put("layers.{d}.dsa.absorb.wv", .{i}); + try L.cols(try L.get("layers.{d}.self_attn.o_proj.weight", .{i}), h.lo * c.v_dim, h.hi * c.v_dim, .bf16); + const o = try L.put("layers.{d}.dsa.o", .{i}); + inline for (.{ "wk", "weights_proj" }) |x| try L.whole(try L.get("layers.{d}.self_attn.indexer." ++ x ++ ".weight", .{i}), .bf16); + const kw = try L.put("layers.{d}.dsa.index.kw", .{i}); + try L.whole(try L.get("layers.{d}.self_attn.indexer.wq_b.weight", .{i}), .bf16); + const qb = try L.put("layers.{d}.dsa.index.qb", .{i}); + try L.whole(try L.get("layers.{d}.self_attn.indexer.k_norm.weight", .{i}), .bf16); + const ln_w = try L.put("layers.{d}.dsa.index.ln_w", .{i}); + try L.whole(try L.get("layers.{d}.self_attn.indexer.k_norm.bias", .{i}), .bf16); + const ln_b = try L.put("layers.{d}.dsa.index.ln_b", .{i}); + try L.whole(try L.get("layers.{d}.self_attn.indexer.index_kpool_compress_gate", .{i}), .bf16); + const gate = try L.put("layers.{d}.dsa.index.gate", .{i}); + try L.whole(try L.get("layers.{d}.self_attn.indexer.index_kpool_compress_ape", .{i}), .bf16); + const ape = try L.put("layers.{d}.dsa.index.ape", .{i}); + return .{ + .proj = proj, + .q_norm = q_norm, + .kv_norm = kv_norm, + .q_b = q_b, + .wk = wk, + .wv = wv, + .o = o, + .index = .{ .kw = kw, .gate = gate, .ln_w = ln_w, .ln_b = ln_b, .ape = ape, .qb = qb, .heads = c.i_heads, .d = c.i_dim }, + .heads = h.len(), + .qk = qk, + .dv = c.v_dim, + .lw = c.kv_lora, + .q_lora = c.q_lora, + .dims = c.hidden, + .eps = c.eps, + .dense_limit = c.keyWidth(), + }; + } + + /// An NVFP4 linear's rows [lo, hi), inputs [klo, khi) as Fp4Linear.from_checkpoint (tf_pack_dense, packScales). + fn fp4(L: *Loader, comptime fmt: []const u8, i: u32, lo: usize, hi: usize, klo: usize, khi: usize) !cuda.qmmf.Weight { + const name = try std.fmt.bufPrint(&L.name_buf, prefix ++ fmt, .{i}); + const l = try core.nvfp4.Linear.read(L.ck, name); + const n = hi - lo; + const k = khi - klo; + for (0..n) |r| try L.host.appendSlice(L.gpa, l.codes[(lo + r) * l.k / 2 + klo / 2 ..][0 .. k / 2]); + var src = try cuda.DeviceBuffer.fromHost(L.d, L.host.items); + defer src.free(); + L.host.clearRetainingCapacity(); + const npad = cuda.nvfp4.padded(n); + const kg = k / 64; + const total: u64 = @as(u64, npad / 64) * kg * 512; + const words = try cuda.DeviceBuffer.alloc(L.d, total * 4); + try L.w.buffers.append(L.gpa, words); + try L.w.named.append(L.gpa, .{ .name = try std.fmt.allocPrint(L.gpa, fmt ++ ".words", .{i}), .ptr = words.ptr, .len = words.len }); + var args: cuda.Args = .{}; + args.add(src.ptr); + for ([_]usize{ n, k / 8, kg }) |v| args.add(@as(c_int, @intCast(v))); + args.add(words.ptr); + args.add(@as(c_longlong, @intCast(total))); + try cuda.launch.launch(L.pack_dense, .{ .grid = .{ .x = @intCast((total + 255) / 256) }, .block = .{ .x = 256 } }, L.stream, &args); + try L.stream.synchronize(); + const scales = try L.gpa.alloc(u8, n * k / 16); + defer L.gpa.free(scales); + for (0..n) |r| @memcpy(scales[r * k / 16 ..][0 .. k / 16], l.scales[(lo + r) * l.k / 16 + klo / 16 ..][0 .. k / 16]); + try L.host.resize(L.gpa, npad * k / 16); + cuda.nvfp4.packScales(L.host.items, scales, n, k); + const bs = try L.put(fmt ++ ".scales", .{i}); + return cuda.nvfp4.weight(words.ptr, bs, l.global, n, k); + } + + /// weights.load's mlp for a dense layer: gate and up by the rank's rows, down by its columns (NVFP4). + fn dense(L: *Loader, i: u32) !mlp.Layer { + const c = &L.w.config; + const s = L.w.share.dense; + const d = c.hidden; + return .{ + .gate = try L.fp4("layers.{d}.mlp.gate_proj", i, s.lo, s.hi, 0, d), + .up = try L.fp4("layers.{d}.mlp.up_proj", i, s.lo, s.hi, 0, d), + .down = try L.fp4("layers.{d}.mlp.down_proj", i, 0, d, s.lo, s.hi), + .width = s.len(), + .dims = d, + .limit = c.swiglu_limit, + }; + } + + /// One device allocation the loader frees itself (staging for the GPU's quantizers). + fn scratch(L: *Loader, n: usize) !cuda.DeviceBuffer { + return cuda.DeviceBuffer.alloc(L.d, n); + } + + /// moe_nvfp4 of the MTP layer: BF16 experts cut to the rank, GPU-quantized 8 at a time, scale as weight_scale_2. + fn quantizedExperts(L: *Loader, i: u32, up: []u32, down: []u32, up_scale: []f32, down_scale: []f32) !void { + const qk = L.tools.quant orelse return error.NoQuantizer; + const c = &L.w.config; + const d: usize = c.hidden; + const cut = L.w.share.moe; + const ni = cut.len(); + const up_len = X.bytes(1, ni, d, 2) / 4; + const down_len = X.bytes(1, d, ni, 1) / 4; + const chunk = 8; + var wb = try L.scratch(chunk * ni * d * 2); + defer wb.free(); + var cb = try L.scratch(chunk * ni * d / 2); + defer cb.free(); + var sb = try L.scratch(chunk * ni * d / 16); + defer sb.free(); + var gb = try L.scratch(chunk * 4); + defer gb.free(); + const codes = try L.gpa.alloc(u8, 3 * chunk * ni * d / 2); // gate, up and down apart: packGateUp takes both + defer L.gpa.free(codes); + const scales = try L.gpa.alloc(u8, 3 * chunk * ni * d / 16); + defer L.gpa.free(scales); + var g: [3][chunk]f32 = undefined; + var e0: usize = 0; + while (e0 < c.experts) : (e0 += chunk) { + const ec = @min(chunk, c.experts - e0); + inline for (.{ "gate_proj", "up_proj", "down_proj" }, 0..) |proj, m| { + const n = if (m < 2) ni else d; // gate and up: the rank's rows; down: every row, the rank's columns + const k = if (m < 2) d else ni; + for (e0..e0 + ec) |e| { + const t = try L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".weight", .{ i, e }); + if (m < 2) try L.rows(t, cut.lo, cut.hi, c.moe_inter, .bf16) else try L.cols(t, cut.lo, cut.hi, .bf16); + } + try wb.upload(0, L.host.items); + L.host.clearRetainingCapacity(); + try quant.nvfp4(qk, L.stream, wb.ptr, ec, n, k, cb.ptr, sb.ptr, gb.ptr); + try L.stream.synchronize(); + try cb.download(0, codes[m * chunk * n * k / 2 ..][0 .. ec * n * k / 2]); + try sb.download(0, scales[m * chunk * n * k / 16 ..][0 .. ec * n * k / 16]); + try gb.download(0, std.mem.sliceAsBytes(g[m][0..ec])); + } + const mc = chunk * ni * d / 2; // one matrix's codes for the chunk; scales an eighth of that + for (0..ec) |j| { + const e = e0 + j; + const one = ni * d; + var cs: [3][]const u8 = undefined; + var ss: [3][]const u8 = undefined; + for (0..3) |m| { + cs[m] = codes[m * mc + j * one / 2 ..][0 .. one / 2]; + ss[m] = scales[m * mc / 8 + j * one / 16 ..][0 .. one / 16]; + } + X.packGateUp(up[e * up_len ..][0..up_len], .{ cs[0], ss[0] }, .{ cs[1], ss[1] }, ni, d); + X.packOne(down[e * down_len ..][0..down_len], cs[2], ss[2], d, ni); + up_scale[e * 2] = g[0][j]; + up_scale[e * 2 + 1] = g[1][j]; + down_scale[e] = g[2][j]; + } + } + } + + /// weights.load's moe(i) on NVFP4: router and bias, routed experts as Experts4, the BF16 shared expert's share. + fn moeLayer(L: *Loader, i: u32) !moe.Layer { + const c = &L.w.config; + const d: usize = c.hidden; + const cut = L.w.share.moe; + const ni = cut.len(); + try L.whole(try L.get("layers.{d}.mlp.gate.weight", .{i}), .bf16); + const router = try L.put("layers.{d}.moe.router", .{i}); + try L.whole(try L.get("layers.{d}.mlp.gate.e_score_correction_bias", .{i}), .f32); + const bias = try L.put("layers.{d}.moe.bias", .{i}); + const count: usize = c.experts; + const up_len = X.bytes(1, ni, d, 2) / 4; + const down_len = X.bytes(1, d, ni, 1) / 4; + const up = try L.gpa.alloc(u32, count * up_len); + defer L.gpa.free(up); + const down = try L.gpa.alloc(u32, count * down_len); + defer L.gpa.free(down); + const up_scale = try L.gpa.alloc(f32, count * 2); + defer L.gpa.free(up_scale); + const down_scale = try L.gpa.alloc(f32, count); + defer L.gpa.free(down_scale); + const codes = try L.gpa.alloc(u8, d * ni / 2); + defer L.gpa.free(codes); + const scales = try L.gpa.alloc(u8, d * ni / 16); + defer L.gpa.free(scales); + if (i == c.layers) try L.quantizedExperts(i, up, down, up_scale, down_scale) else for (0..count) |e| { + var ls: [3]core.nvfp4.Linear = undefined; + inline for (.{ "gate_proj", "up_proj", "down_proj" }, 0..) |proj, m| { + const name = try std.fmt.bufPrint(&L.name_buf, prefix ++ "layers.{d}.mlp.experts.{d}." ++ proj, .{ i, e }); + ls[m] = try core.nvfp4.Linear.read(L.ck, name); + } + // gate and up: the rank's rows; down: every row, the rank's columns + const gate = .{ ls[0].codes[cut.lo * d / 2 ..][0 .. ni * d / 2], ls[0].scales[cut.lo * d / 16 ..][0 .. ni * d / 16] }; + const upm = .{ ls[1].codes[cut.lo * d / 2 ..][0 .. ni * d / 2], ls[1].scales[cut.lo * d / 16 ..][0 .. ni * d / 16] }; + X.packGateUp(up[e * up_len ..][0..up_len], gate, upm, ni, d); + const l = ls[2]; + for (0..d) |r| { + @memcpy(codes[r * ni / 2 ..][0 .. ni / 2], l.codes[r * l.k / 2 + cut.lo / 2 ..][0 .. ni / 2]); + @memcpy(scales[r * ni / 16 ..][0 .. ni / 16], l.scales[r * l.k / 16 + cut.lo / 16 ..][0 .. ni / 16]); + } + X.packOne(down[e * down_len ..][0..down_len], codes, scales, d, ni); + up_scale[e * 2] = ls[0].global; + up_scale[e * 2 + 1] = ls[1].global; + down_scale[e] = l.global; + } + var ptrs: [4]u64 = undefined; + for (&ptrs, [_][]const u8{ std.mem.sliceAsBytes(up), std.mem.sliceAsBytes(down), std.mem.sliceAsBytes(up_scale), std.mem.sliceAsBytes(down_scale) }, [_][]const u8{ "up", "down", "up_scale", "down_scale" }) |*p, bytes, what| { + try L.host.appendSlice(L.gpa, bytes); + p.* = try L.put("layers.{d}.moe.experts.{s}", .{ i, what }); + } + const sh = L.w.share.shared; + inline for (.{ "gate_proj", "up_proj" }) |x| try L.rows(try L.get("layers.{d}.mlp.shared_experts." ++ x ++ ".weight", .{i}), sh.lo, sh.hi, c.shared_inter, .bf16); + const gu = try L.put("layers.{d}.moe.shared.gu", .{i}); + try L.cols(try L.get("layers.{d}.mlp.shared_experts.down_proj.weight", .{i}), sh.lo, sh.hi, .bf16); + const sd = try L.put("layers.{d}.moe.shared.down", .{i}); + return .{ + .router = router, + .bias = bias, + .routed = .{ .format = .nvfp4, .up = ptrs[0], .down = ptrs[1], .up_scale = ptrs[2], .down_scale = ptrs[3], .width = ni, .dims = d, .experts = count, .limit = c.swiglu_limit }, + .count = count, + .shared_gu = gu, + .shared_down = sd, + .shared = sh.len(), + .top_k = c.topk, + .scale = c.routed_scale, + .norm = c.norm_topk, + }; + } + + /// weights.load's MTPW: layer `layers` as plain DSA + MoE, norms, eh_proj, draft head from the rank's head rows. + fn mtp(L: *Loader) !Mtp { + const c = &L.w.config; + const i = c.layers; + try L.whole(try L.get("layers.{d}.input_layernorm.weight", .{i}), .bf16); + const in_norm = try L.put("layers.{d}.in_norm", .{i}); + try L.whole(try L.get("layers.{d}.post_attention_layernorm.weight", .{i}), .bf16); + const post_norm = try L.put("layers.{d}.post_norm", .{i}); + const none: Hc = .{ .fn_w = 0, .base = 0, .scale = 0 }; + const layer_w: Layer = .{ .index = i, .attn_hc = none, .ffn_hc = none, .in_norm = in_norm, .post_norm = post_norm, .mixer = .{ .mla = try L.dsa(i) }, .ffn = .{ .moe = try L.moeLayer(i) } }; + L.ck.release(); + var out: [4]u64 = undefined; + for (&out, [_][]const u8{ "enorm.weight", "hnorm.weight", "eh_proj.weight", "shared_head.norm.weight" }, [_][]const u8{ "enorm", "hnorm", "eh", "norm" }) |*p, name, what| { + try L.whole(try L.ck.get(try std.fmt.bufPrint(&L.name_buf, prefix ++ "layers.{d}.{s}", .{ i, name })), .bf16); + p.* = try L.put("mtp.{s}", .{what}); + } + // the draft head: quantize4 of the BF16 head's rows, then make_q4's tiling + const qk = L.tools.quant orelse return error.NoQuantizer; + const tk = L.tools.q4 orelse return error.NoQuantizer; + const n = L.w.share.vocab.len(); + const k: usize = c.hidden; + var raw: [3]cuda.DeviceBuffer = undefined; + var made: usize = 0; + defer for (raw[0..made]) |*b| b.free(); + for (&raw, [_]usize{ n * k / 2, n * k / 32, n * k / 32 }) |*b, size| { + b.* = try L.scratch(size); + made += 1; + } + try quant.q4(qk, L.stream, L.w.head, n, k, raw[0].ptr, raw[1].ptr, raw[2].ptr); + var tiled: [3]u64 = undefined; + for (&tiled, q4.bytes(n, k), [_][]const u8{ "words", "scales", "biases" }) |*p, size, what| { + const b = try cuda.DeviceBuffer.alloc(L.d, size); + try L.w.buffers.append(L.gpa, b); + try L.w.named.append(L.gpa, .{ .name = try std.fmt.allocPrint(L.gpa, "draft_head.{s}", .{what}), .ptr = b.ptr, .len = b.len }); + p.* = b.ptr; + } + const draft: q4.Q4 = .{ .w = tiled[0], .s = tiled[1], .b = tiled[2], .n = n, .k = k, .npad = q4.npadOf(n) }; + try q4.pack(tk, L.stream, raw[0].ptr, raw[1].ptr, raw[2].ptr, draft); + try L.stream.synchronize(); + return .{ .enorm = out[0], .hnorm = out[1], .eh = out[2], .norm = out[3], .layer = layer_w, .draft = draft }; + } + + fn layer(L: *Loader, i: u32) !Layer { + const c = &L.w.config; + try L.whole(try L.get("layers.{d}.input_layernorm.weight", .{i}), .bf16); + const in_norm = try L.put("layers.{d}.in_norm", .{i}); + try L.whole(try L.get("layers.{d}.post_attention_layernorm.weight", .{i}), .bf16); + const post_norm = try L.put("layers.{d}.post_norm", .{i}); + return .{ + .index = i, + .attn_hc = try L.hc(i, "attn"), + .ffn_hc = try L.hc(i, "ffn"), + .in_norm = in_norm, + .post_norm = post_norm, + .mixer = switch (c.kind(i)) { + .kda => .{ .kda = try L.kda(i) }, + .mla => .{ .mla = try L.dsa(i) }, + }, + .ffn = if (c.isMoe(i)) .{ .moe = try L.moeLayer(i) } else .{ .dense = try L.dense(i) }, + }; + } +}; + +/// weights.load without the MTP head: embedding, layers, final norm, the rank's BF16 head rows; `pack_dense`. +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, stream: cuda.Stream, tools: Tools, dir: []const u8, c: cfg.Config, share: shard.Share, layers: []const u32, with_mtp: bool) !Weights { + var w: Weights = .{ .gpa = gpa, .config = c, .share = share }; + errdefer w.deinit(); + var ck = try core.Checkpoint.openModelPrefix(gpa, io, dir, prefix); + defer ck.close(); + var head_ck = try core.Checkpoint.openModelPrefix(gpa, io, dir, "lm_head."); + defer head_ck.close(); + var L: Loader = .{ .gpa = gpa, .d = d, .stream = stream, .pack_dense = tools.pack_dense, .tools = tools, .w = &w, .ck = &ck }; + defer L.host.deinit(gpa); + w.layers = try gpa.alloc(Layer, layers.len); + for (layers, w.layers) |i, *l| { + if (i >= c.layers) return error.BadLayer; + l.* = try L.layer(i); + ck.release(); // a rank reads about half the checkpoint: its pages would crowd the GPU's allocations out + } + try L.whole(try L.get("embed_tokens.weight", .{}), .bf16); + w.embed = try L.put("embed", .{}); + try L.whole(try L.get("norm.weight", .{}), .bf16); + w.norm = try L.put("norm", .{}); + try L.rows(try head_ck.get("lm_head.weight"), share.vocab.lo, share.vocab.hi, c.vocab, .bf16); + w.head = try L.put("head", .{}); + if (with_mtp and c.mtp > 0) w.mtp = try L.mtp(); + return w; +} + +test "bf16 rounding of fp32 is to nearest even" { + try std.testing.expectEqual(@as(u16, 0x3f80), toBf16(1.0)); + try std.testing.expectEqual(@as(u16, 0x3f80), toBf16(@bitCast(@as(u32, 0x3f808000)))); // a tie to even (down) + try std.testing.expectEqual(@as(u16, 0x3f82), toBf16(@bitCast(@as(u32, 0x3f818000)))); // a tie to even (up) + try std.testing.expectEqual(@as(u16, 0x3f81), toBf16(@bitCast(@as(u32, 0x3f808001)))); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..f8fbd26d45 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -4,21 +4,15 @@ const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const glm = @import("glm_cuda"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); +const snaps = @import("cuda_snaps.zig"); const nemotron_slide = @import("nemotron_slide.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{nemotron.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 = .{ - .bytes = nemotron.native.snap_bytes, - .save = nemotron.native.snap_save, - .restore = nemotron.native.snap_restore, - .drop = nemotron.native.snap_drop, -}; +const registry = .{ nemotron.native, glm.native }; pub const backends: []const []const u8 = &.{"cuda"}; pub const families: []const api.Family = blk: { @@ -322,49 +316,75 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }, }; const before = try pool(a, io, &g.ctx, problem) orelse return null; - const weights = weightBytes(io, o.dir); + // a family split over hosts holds only its rank's share (GLM-5.3-Flash's TP) + const weights = if (@hasDecl(F, "weightShare")) F.weightShare(a, io, o.dir, weightBytes(io, o.dir)) else weightBytes(io, o.dir); const held0 = cuda.usage(false).device; + // a family over hosts reads every rank's room after its load, by this budget, and admits by the smallest + const Probe = struct { + a: Allocator, + io: std.Io, + ctx: *const cuda.Context, + held: u64, + problem: *[]const u8, + fn read(p: *anyopaque) ?u64 { + const self: *@This() = @ptrCast(@alignCast(p)); + const now = (pool(self.a, self.io, self.ctx, self.problem) catch return null) orelse return null; + return now.room(cuda.usage(false).device - self.held); + } + }; + var probe: Probe = .{ .a = a, .io = io, .ctx = &g.ctx, .held = held0, .problem = problem }; + var opts: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments, .slide = o.slide }; + if (@hasField(F.Options, "room")) opts.room = .{ .ctx = &probe, .read = Probe.read }; if (weights > before.room(held0)) { 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| { + const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, opts) catch |e| { + if (@hasDecl(F, "explains_open")) if (F.explain(null, e)) |text| { + problem.* = text; + return null; + }; 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; }; defer if (!opened) loaded.deinit(loaded.ctx); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - const room = after.room(model) -| loaded.learn_bytes; + var room = after.room(model) -| loaded.learn_bytes; + var low: []const u8 = ""; // names the rank whose room bounds the streams, when it is not this one + if (@hasField(@TypeOf(loaded), "peer_room")) if (loaded.peer_room) |pr| if (pr.bytes < room) { + room = pr.bytes; + low = try std.fmt.allocPrint(a, " on rank {d}, the smallest room of the ranks", .{pr.rank}); + }; const free = loaded.free_streams; const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { - error.NoStream => try std.fmt.allocPrint(a, "the CUDA memory budget fits no stream: one at a {d}-token window takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; lower --context, or free device memory", .{ window, toGib(loaded.stream_bytes), toGib(room), toGib(model), toGib(after.reserve) }), - error.TooMany => try std.fmt.allocPrint(a, "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, and the CUDA memory budget leaves {d:.2} GiB: serve --parallel {d}, or lower --context", .{ o.lanes, toGib(loaded.stream_bytes * (o.lanes -| free)), window, toGib(room), free + room / loaded.stream_bytes }), + error.NoStream => try std.fmt.allocPrint(a, "the CUDA memory budget fits no stream: one at a {d}-token window takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve{s}; lower --context, or free device memory", .{ window, toGib(loaded.stream_bytes), toGib(room), toGib(model), toGib(after.reserve), low }), + error.TooMany => try std.fmt.allocPrint(a, "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, and the CUDA memory budget leaves {d:.2} GiB{s}: serve --parallel {d}, or lower --context", .{ o.lanes, toGib(loaded.stream_bytes * (o.lanes -| free)), window, toGib(room), low, free + room / loaded.stream_bytes }), }; return null; }; 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)); - 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 passes of up to {d} rows{s}{s}; {s}", .{ + const cache = if (@hasDecl(F, "snap_save")) 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_save") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); - // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet - store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &nemotron_snaps }, .{ .lookahead = 1, .min_prompt = 0, .warm = .waits }, cache - try nemotron.native.snap_ready(loaded.target, cache)); + // rules and budget from the family's snap_* declarations (cuda_snaps.zig) + store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &snaps.Snaps(F).vtable }, snaps.cacheRules(F), try snaps.cacheBudget(F, loaded.target, cache)); h.store = store; } errdefer if (h.store) |store| { store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { + if (@hasDecl(F, "slide")) if (loaded.learns != null) { const Slide = nemotron_slide.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); @@ -378,13 +398,13 @@ 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); gpa.free(h.dir); }; // 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, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null and @hasDecl(F, "warm_quiet"), .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } @@ -397,6 +417,7 @@ test { _ = budget; // cuda_memory.zig's own tests, which no other test binary reaches _ = refusals; _ = lane_backend; + _ = snaps; } test "the lane core sees a family's prefill_step through the lane backend, so a served prompt can step" { diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig new file mode 100644 index 0000000000..961ab866ac --- /dev/null +++ b/zig/src/native/cuda_snaps.zig @@ -0,0 +1,62 @@ +//! A CUDA family's prompt cache in the native server: rules, budget and copies from the family's `snap_*` decls. +const std = @import("std"); +const api = @import("engine_api"); + +/// Any row, the head one token on (Nemotron's); `snap_grid`'s chunk starts and `warm_quiet`'s warm turns if declared. +pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { + var rules: api.prompt_cache.Rules = .{ .lookahead = 1, .min_prompt = 0 }; + if (@hasDecl(F, "warm_quiet")) rules.warm = .waits; + if (@hasDecl(F, "snap_grid")) { + rules.planned = true; + rules.grid = F.snap_grid; + } + return rules; +} + +/// The cache's bytes less what the family sets aside for its states (`snap_ready`, as Nemotron's spare buffers). +pub fn cacheBudget(comptime F: type, target: anytype, cache: u64) !u64 { + return if (@hasDecl(F, "snap_ready")) cache - try F.snap_ready(target, cache) else cache; +} + +/// The family's `snap_*` copies (none: no prompt cache), a container-level value so the store's pointer stays valid. +pub fn Snaps(comptime F: type) type { + return struct { + pub const vtable: api.prompt_cache.Snapshots.VTable = .{ + .bytes = F.snap_bytes, + .save = F.snap_save, + .restore = F.snap_restore, + .drop = F.snap_drop, + }; + }; +} + +test "a family with snap_grid keeps on its chunk starts, any other at any row as Nemotron does" { + const Glm = struct { + pub const snap_grid: u32 = 2048; + }; + const Other = struct {}; + const Warm = struct { + pub const warm_quiet: u64 = 1; + }; + const any = cacheRules(Other); + try std.testing.expect(!any.planned and any.min_prompt == 0 and any.lookahead == 1 and any.warm == .off); + try std.testing.expect(cacheRules(Warm).warm == .waits and cacheRules(Glm).warm == .off); + const Stub = struct { + fn bytes(_: *anyopaque, _: u32) u64 { + return 0; + } + fn save(_: *anyopaque, _: ?*anyopaque, _: u32) anyerror!api.prompt_cache.Saved { + return error.Unused; + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: api.prompt_cache.Saved) anyerror!void {} + fn drop(_: *anyopaque, _: api.prompt_cache.Saved) void {} + }; + var unused: u8 = 0; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const prompt = try arena.allocator().alloc(u32, 3547); + for (prompt, 0..) |*t, i| t.* = @intCast(i); + var store = api.prompt_cache.Store.init(std.testing.allocator, .{ .ptr = &unused, .vtable = &.{ .bytes = Stub.bytes, .save = Stub.save, .restore = Stub.restore, .drop = Stub.drop } }, cacheRules(Glm), 1 << 30); + defer store.deinit(); + try std.testing.expectEqualSlices(u32, &[_]u32{2048}, try store.marks(arena.allocator(), prompt, 0, 3540, &.{}, &.{}, &.{})); +} diff --git a/zig/tests/cuda/check.zig b/zig/tests/cuda/check.zig index 88b1f20f9c..f791cb213d 100644 --- a/zig/tests/cuda/check.zig +++ b/zig/tests/cuda/check.zig @@ -33,7 +33,11 @@ pub fn sameBytes(what: []const u8, got: []const u8, want: []const u8) Failed!voi var words: usize = 0; var i: usize = 0; while (i + 4 <= got.len) : (i += 4) { - if (!std.mem.eql(u8, got[i..][0..4], want[i..][0..4])) words += 1; + if (std.mem.eql(u8, got[i..][0..4], want[i..][0..4])) continue; + words += 1; + const g = std.mem.readInt(u32, got[i..][0..4], .little); + const w = std.mem.readInt(u32, want[i..][0..4], .little); + if (words <= 4) std.debug.print(" word {d}: {x:0>8} ({e}) for {x:0>8} ({e})\n", .{ i / 4, g, @as(f32, @bitCast(g)), w, @as(f32, @bitCast(w)) }); } std.debug.print("FAIL {s}: first difference at byte {d}, {d} of {d} words differ\n", .{ what, first, words, got.len / 4 }); return error.TestFailed; diff --git a/zig/tests/cuda/glm_lanes_tests.zig b/zig/tests/cuda/glm_lanes_tests.zig new file mode 100644 index 0000000000..02ceaa32fe --- /dev/null +++ b/zig/tests/cuda/glm_lanes_tests.zig @@ -0,0 +1,471 @@ +//! GLM-5.3-Flash's lane backend under upstream's round loop, no server: rank 0 drives lanes.Engine. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +/// A communicator of `world` ranks: rank 0 writes its unique id to `id_path` as hex, the others read it there. +pub fn openComm(gpu: Gpu, lib: *const cuda.nccl.Library, rank: c_int, world: c_int, id_path: []const u8) !cuda.nccl.Comm { + var id: cuda.nccl.UniqueId = undefined; + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + const hex = std.fmt.bytesToHex(id.internal, .lower); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = id_path, .data = &hex }); + } else { + var hex: [256]u8 = undefined; + const got = try std.Io.Dir.cwd().readFile(gpu.io, id_path, &hex); + try check.expect(got.len == 256, "unique id file {s}: {d} hex digits", .{ id_path, got.len }); + _ = try std.fmt.hexToBytes(&id.internal, &hex); + } + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, world, id, rank), "ncclCommInitRank"); + return comm; +} + +const DcKind = struct { kind: []const u8, prompt: []const i32, want: []const i32 }; + +/// `glm-lanes-cover`: CONTRIBUTING's single-stream contracts on one Engine: drafts, cancel, discard, eos, prompt cache. +pub fn coverCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_int, world: c_int, id_path: []const u8) !void { + const glm = @import("glm_cuda"); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dc = try std.json.parseFromSliceLeaky(struct { capacity: usize = 8192, kinds: []const DcKind }, a, try std.Io.Dir.cwd().readFileAlloc(gpu.io, dc_path, a, .limited(1 << 26)), .{ .ignore_unknown_fields = true }); + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + const comm = try openComm(gpu, &lib, rank, world, id_path); + defer _ = lib.api.ncclCommDestroy(comm); + var problem: []const u8 = ""; + const set = glm.ranks.Settings.read(&glm.ranks.environment, &problem) catch |err| { + std.debug.print("REFUSE {s}\n", .{problem}); + return err; + }; + const e = try glm.engine.Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, checkpoint, &lib, comm, @intCast(rank), @intCast(world), .{ .capacity = dc.capacity, .mtp = true, .settings = set }); + defer e.deinit(); + var link = try glm.follow.Link.init(gpu.gpa, gpu.d, &lib, comm, e.stream, dc.capacity); + defer link.deinit(); + if (rank != 0) { + try glm.follow.follow(e, &link); + check.pass("rank {d} followed rank 0's steps to the end, through every cancel and discard", .{rank}); + return; + } + var gl: glm.lanes.Lanes = .{ .x = .{ .e = e, .link = &link } }; + defer gl.deinit(); + var cfg = try lanes.Config.init(gpu.gpa, gl.facts(), glm.lanes.max_rows, glm.lanes.max_rows - 1); + defer cfg.deinit(gpu.gpa); + var wall: lanes.backend.WallClock = .{ .io = gpu.io }; + var le = lanes.Engine.init(gpu.gpa, &cfg, gl.backend(), wall.clock()); + defer le.deinit(); + defer gl.x.end() catch {}; + var c: Cover = .{ .gpa = gpu.gpa, .a = a, .le = &le }; + + const k = dc.kinds[0]; + const p = try a.alloc(u32, k.prompt.len); + for (p, k.prompt) |*u, i| u.* = @intCast(i); + // W, what every contract below holds to: the model's own plain greedy pass ("draft": false) + const w: []const u32 = (try c.run(.{ .prompt = p, .max_new = @intCast(k.want.len), .drafts = false })).ids; + const copy_prompt = try std.mem.concat(a, u32, &.{ p, w[0..256], p }); + + // copy drafts, greedy and sampled + for ([_]bool{ false, true }) |keyed| { + const sampling: ?lanes.Sampling = if (keyed) .{ .seed = lanes.sampling.seedFor(copy_prompt, 0) } else null; + const plain = try c.run(.{ .prompt = copy_prompt, .max_new = 256, .drafts = false, .sampling = sampling }); + const drafted = try c.run(.{ .prompt = copy_prompt, .max_new = 256, .sampling = sampling, .copies = true }); + try c.same(if (keyed) "copy drafts, sampled" else "copy drafts, greedy", drafted, plain.ids, .length); + // a thin model's sampled pass does not repeat W, so nothing is there to copy: its count is only reported + if (keyed and set.layers > 0) check.pass("copy drafts, sampled, thin model: the proposer judged {d} copied rows ({d} kept)", .{ drafted.judged, drafted.copied }) else try c.expect(drafted.judged > 0, "copy drafts{s}: the proposer judged {d} copied rows ({d} kept)", .{ if (keyed) ", sampled" else "", drafted.judged, drafted.copied }); + } + // cancels in the prompt pass (P + W[0..256] + P is three 2,048-row chunks: asks 0, 1, 2 before them, 3 after) + for ([_]u32{ 0, 1, 3 }) |at| { + var cut: Cut = .{ .at = at }; + const r = c.run(.{ .prompt = copy_prompt, .max_new = 256, .cancel = &cut }); + try c.expect(if (r) |_| false else |err| err == error.Cancelled, "a cancel at ask {d} of the prompt pass ended it with error.Cancelled", .{at}); + try c.expect(cut.asked == at + 1, "the cancel was heard at ask {d} ({d} asks)", .{ at, cut.asked }); + try c.same(try std.fmt.allocPrint(a, "the request after a cancel at prompt ask {d}", .{at}), try c.run(.{ .prompt = p, .max_new = 64 }), w[0..64], .length); + } + // discards between rounds: after 1, after 8, and right after a round that rejected a draft (its rows uncommitted) + for ([_]?u32{ 1, 8, null }) |rounds| { + const left = try c.discardAfter(p, rounds); + try c.expect(left, "a stream was discarded mid-decode ({s})", .{if (rounds) |n| try std.fmt.allocPrint(a, "after {d} rounds", .{n}) else "after a rejected draft"}); + try c.same(try std.fmt.allocPrint(a, "the request after a discard {s}", .{if (rounds) |n| try std.fmt.allocPrint(a, "after {d} rounds", .{n}) else "after a rejected draft"}), try c.run(.{ .prompt = p, .max_new = 64 }), w[0..64], .length); + } + // eos at W[j]'s first occurrence (the first at or past each of 0, 9, 33, 100; a thin model's W repeats more) + var last_j: ?usize = null; + for ([_]usize{ 0, 9, 33, 100 }) |from| { + const j = for (from..w.len) |i| { + if (std.mem.indexOfScalar(u32, w[0..i], w[i]) == null) break i; + } else continue; + if (last_j == j) continue; + last_j = j; + for ([_]bool{ true, false }) |drafts| { + const r = try c.run(.{ .prompt = p, .max_new = 512, .eos = &.{w[j]}, .drafts = drafts }); + try c.same(try std.fmt.allocPrint(a, "eos at token {d}{s}", .{ j, if (drafts) " with drafts" else "" }), r, w[0 .. j + 1], .stop); + } + } + // max_new at the window's edges + for (1..6) |n| try c.same(try std.fmt.allocPrint(a, "max_new {d} with drafts", .{n}), try c.run(.{ .prompt = p, .max_new = @intCast(n) }), w[0..n], .length); + // resumed = fresh: states kept at marks, as the lane host keeps them, then passes resumed from them + { + var keeper: Keeper = .{ .e = e }; + const marks = [_]u32{ 2048, 4096 }; + const fresh = try c.run(.{ .prompt = copy_prompt, .max_new = 64, .reuse = .{ .marks = &marks, .hook = .{ .ptr = &keeper, .at = Keeper.at } } }); + try c.expect(keeper.n == marks.len, "the pass kept {d} states at its marks ({d} asked)", .{ keeper.n, marks.len }); + for (keeper.kept[0..keeper.n], marks[0..keeper.n]) |kept, at| { + const got = try c.run(.{ .prompt = copy_prompt, .max_new = 64, .reuse = .{ .saved = kept, .at = at } }); + try c.expect(got.cached == at, "the pass resumed at {d} ({d} restored)", .{ at, got.cached }); + try c.same(try std.fmt.allocPrint(a, "resumed at {d}, against the fresh pass", .{at}), got, fresh.ids, .length); + } + const other = try std.mem.concat(a, u32, &.{ copy_prompt[0..4096], p[0..300] }); + const other_fresh = try c.run(.{ .prompt = other, .max_new = 64 }); + if (keeper.n == marks.len) { + const got = try c.run(.{ .prompt = other, .max_new = 64, .reuse = .{ .saved = keeper.kept[1], .at = 4096 } }); + try c.expect(got.cached == 4096, "another tail resumed at 4096 ({d} restored)", .{got.cached}); + try c.same("another tail after the kept 4096, against its fresh pass", got, other_fresh.ids, .length); + } + for (keeper.kept[0..keeper.n]) |kept| glm.native.snap_drop(e, kept); + } + try c.same("the last request, all of W with drafts", try c.run(.{ .prompt = p, .max_new = @intCast(w.len) }), w, .length); + if (c.failed) return error.TestFailed; +} + +/// `glm-lanes-concurrent`: each concurrent stream on one lanes.Engine equals its solo run (2+ streams a round). +pub fn concurrentCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_int, world: c_int, id_path: []const u8) !void { + const glm = @import("glm_cuda"); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dc = try std.json.parseFromSliceLeaky(struct { capacity: usize = 8192, kinds: []const DcKind }, a, try std.Io.Dir.cwd().readFileAlloc(gpu.io, dc_path, a, .limited(1 << 26)), .{ .ignore_unknown_fields = true }); + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + const comm = try openComm(gpu, &lib, rank, world, id_path); + defer _ = lib.api.ncclCommDestroy(comm); + var problem: []const u8 = ""; + const set = glm.ranks.Settings.read(&glm.ranks.environment, &problem) catch |err| { + std.debug.print("REFUSE {s}\n", .{problem}); + return err; + }; + const e = try glm.engine.Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, checkpoint, &lib, comm, @intCast(rank), @intCast(world), .{ .capacity = dc.capacity, .mtp = true, .settings = set }); + defer e.deinit(); + var link = try glm.follow.Link.init(gpu.gpa, gpu.d, &lib, comm, e.stream, dc.capacity); + defer link.deinit(); + if (rank != 0) { + try glm.follow.follow(e, &link); + check.pass("rank {d} followed rank 0's steps to the end, every stream's sequence on its slot", .{rank}); + return; + } + var gl: glm.lanes.Lanes = .{ .x = .{ .e = e, .link = &link } }; + defer gl.deinit(); + var cfg = try lanes.Config.init(gpu.gpa, gl.facts(), glm.lanes.max_rows, glm.lanes.max_rows - 1); + defer cfg.deinit(gpu.gpa); + var wall: lanes.backend.WallClock = .{ .io = gpu.io }; + var le = lanes.Engine.init(gpu.gpa, &cfg, gl.backend(), wall.clock()); + defer le.deinit(); + defer gl.x.end() catch {}; + var c: Cover = .{ .gpa = gpu.gpa, .a = a, .le = &le }; + + const n = 128; + const ps = try a.alloc([]const u32, dc.kinds.len); + for (ps, dc.kinds) |*p, k| { + const ids = try a.alloc(u32, k.prompt.len); + for (ids, k.prompt) |*u, i| u.* = @intCast(i); + p.* = ids; + } + const seeds = [_]lanes.Sampling{ .{ .seed = 11 }, .{ .seed = 22 }, .{ .seed = 33 } }; + const solo = try a.alloc([]const u32, ps.len); + const solo_s = try a.alloc([]const u32, ps.len); + for (ps, solo, solo_s, 0..) |p, *g, *gs, i| { + g.* = (try c.run(.{ .prompt = p, .max_new = n })).ids; + gs.* = (try c.run(.{ .prompt = p, .max_new = n, .sampling = seeds[i % seeds.len] })).ids; + } + const P = ps.len; + var most: usize = 0; + const cases = [_]struct { name: []const u8, joins: []const Join }{ + .{ .name = "all at once", .joins = &.{ .{ .p = 0 }, .{ .p = 1 }, .{ .p = 2 } } }, + .{ .name = "staggered by 5 rounds", .joins = &.{ .{ .p = 0 }, .{ .p = 1, .at = 5 }, .{ .p = 2, .at = 10 } } }, + .{ .name = "in reverse", .joins = &.{ .{ .p = 2 }, .{ .p = 1 }, .{ .p = 0 } } }, + .{ .name = "the same prompt twice", .joins = &.{ .{ .p = 0 }, .{ .p = 0 } } }, + .{ .name = "sampled", .joins = &.{ .{ .p = 0, .sampled = true }, .{ .p = 1, .sampled = true }, .{ .p = 2, .sampled = true } } }, + .{ .name = "the second discarded after 6 rounds", .joins = &.{ .{ .p = 0 }, .{ .p = 1, .discard_at = 6 }, .{ .p = 2 } } }, + .{ .name = "prompts a chunk at a time while the others decode", .joins = &.{ .{ .p = 0 }, .{ .p = 1, .at = 2, .filled = true }, .{ .p = 2, .at = 4, .filled = true } } }, + .{ .name = "every prompt a chunk at a time", .joins = &.{ .{ .p = 0, .filled = true }, .{ .p = 1, .filled = true }, .{ .p = 2, .filled = true } } }, + .{ .name = "sampled, prompts a chunk at a time", .joins = &.{ .{ .p = 0, .sampled = true }, .{ .p = 1, .at = 2, .sampled = true, .filled = true }, .{ .p = 2, .at = 4, .sampled = true, .filled = true } } }, + }; + for (cases) |cs| { + const got = try c.together(cs.joins, ps, seeds[0..], n); + most = @max(most, got.most); + for (cs.joins, got.ids, 0..) |j, ids, i| { + if (j.discard_at != null or j.p >= P) continue; + try c.same(try std.fmt.allocPrint(a, "{s}: stream {d} (prompt {d}{s}) against its solo run", .{ cs.name, i, j.p, if (j.sampled) ", sampled" else "" }), .{ .ids = ids, .reason = .length, .rounds = got.rounds[i], .judged = 0, .copied = 0, .cached = 0 }, if (j.sampled) solo_s[j.p] else solo[j.p], .length); + } + try c.expect(got.most >= 2, "{s}: {d} streams at most in one round, {d} rounds with two or more", .{ cs.name, got.most, got.shared }); + for (cs.joins) |j| if (j.filled) { + try c.expect(got.between > 0, "{s}: {d} rounds of the others between a prompt's chunks", .{ cs.name, got.between }); + break; + }; + } + try c.expect(gl.shared > 0, "{d} forwards held several streams' windows (Engine.verifyShared)", .{gl.shared}); + if (c.failed) return error.TestFailed; +} + +/// `glm-shared-check`: streams verified in one shared forward equal lone twins byte for byte. +pub fn sharedCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_int, world: c_int, id_path: []const u8) !void { + const glm = @import("glm_cuda"); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dc = try std.json.parseFromSliceLeaky(struct { capacity: usize = 8192, kinds: []const DcKind }, a, try std.Io.Dir.cwd().readFileAlloc(gpu.io, dc_path, a, .limited(1 << 26)), .{ .ignore_unknown_fields = true }); + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + const comm = try openComm(gpu, &lib, rank, world, id_path); + defer _ = lib.api.ncclCommDestroy(comm); + var problem: []const u8 = ""; + const set = glm.ranks.Settings.read(&glm.ranks.environment, &problem) catch |err| { + std.debug.print("REFUSE {s}\n", .{problem}); + return err; + }; + const e = try glm.engine.Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, checkpoint, &lib, comm, @intCast(rank), @intCast(world), .{ .capacity = dc.capacity, .mtp = true, .settings = set }); + defer e.deinit(); + var link = try glm.follow.Link.init(gpu.gpa, gpu.d, &lib, comm, e.stream, dc.capacity); + defer link.deinit(); + if (rank != 0) { + try glm.follow.follow(e, &link); + check.pass("rank {d} followed rank 0's shared and lone windows to the end", .{rank}); + return; + } + const x: glm.follow.Leader = .{ .e = e, .link = &link }; + defer x.end() catch {}; + const S = 5; // streams: every prompt, the first twice, the last sampled + const V = e.share.vocab.len(); + var failed = false; + var first: [S]u32 = undefined; + for (0..2 * S) |slot| { + const i = slot % S; + const k = dc.kinds[i % dc.kinds.len]; + const ids = try a.alloc(u32, k.prompt.len); + for (ids, k.prompt) |*u, v| u.* = @intCast(v); + try x.bind(@intCast(slot)); + try x.begin(if (i == S - 1) .{ .seed = 77, .temperature = 0.7, .top_k = 20, .top_p = 0.95 } else null); + const t = try x.prefill(ids, null, null, null); + if (slot >= S and t != first[i]) { + std.debug.print("FAIL shared: stream {d}'s twin drew {d} from the prompt, the stream {d}\n", .{ i, t, first[i] }); + failed = true; + } + first[i] = t; + var held: [glm.engine.max_drafts]u32 = undefined; + try x.draft(true, &.{t}, 1, &held); + } + const lone = try a.alloc(u16, 4 * V); + const mine = try a.alloc(u16, 4 * V); + var compared: usize = 0; + var round: usize = 0; + while (round < 12) : (round += 1) { + const n = 2 + round % (S - 1); // streams this round: 2 to S + var parts: [S]glm.engine.Part = undefined; + var ids: [S][4]u32 = undefined; + var outs: [S][4]u32 = undefined; + var twin: [S][4]u32 = undefined; + for (0..n) |i| { + const rows = 1 + (round + i) % 4; + for (0..rows) |r| ids[i][r] = @intCast(100 + (round * 31 + i * 7 + r * 13) % 5000); + parts[i] = .{ .slot = @intCast(i), .ids = ids[i][0..rows], .out = outs[i][0..rows] }; + } + try x.verifyShared(parts[0..n]); + var row0: usize = 0; + for (0..n) |i| { // each stream's rows of the shared forward against its twin's lone round (verify, keep, draft) + const rows = parts[i].ids.len; + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(mine.ptr, e.bs.?.logits + row0 * V * 2, rows * V * 2), "cuMemcpyDtoH"); + try x.bind(@intCast(S + i)); + try x.verify(parts[i].ids, twin[i][0..rows]); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(lone.ptr, e.bd.logits, rows * V * 2), "cuMemcpyDtoH"); + if (!std.mem.eql(u16, mine[0 .. rows * V], lone[0 .. rows * V]) or !std.mem.eql(u32, outs[i][0..rows], twin[i][0..rows])) { + std.debug.print("FAIL shared: round {d}, stream {d} ({d} rows from row {d} of the shared window): logits or tokens differ from its twin's\n", .{ round, i, rows, row0 }); + failed = true; + } + compared += rows; + row0 += rows; + const keep = 1 + (round * 7 + i) % rows; // the same rows kept and drafted from on both + var d0: [glm.engine.max_drafts]u32 = undefined; + var d1: [glm.engine.max_drafts]u32 = undefined; + for ([_]u32{ @intCast(S + i), @intCast(i) }, [_]*[glm.engine.max_drafts]u32{ &d1, &d0 }) |slot, d| { + try x.bind(slot); + try x.keep(keep); + try x.draft(false, outs[i][0..keep], 2, d); + } + if (!std.mem.eql(u32, d0[0..2], d1[0..2])) { + std.debug.print("FAIL shared: round {d}, stream {d}: drafts {any} against its twin's {any}\n", .{ round, i, d0[0..2], d1[0..2] }); + failed = true; + } + } + } + try e.stream.synchronize(); + for (0..S) |i| { // the whole states after every round's commit + const st = &e.seqs.items[i].?.st; + const tw = &e.seqs.items[S + i].?.st; + const len = st.arena.buf.len; + const h0 = try a.alloc(u8, len); + const h1 = try a.alloc(u8, len); + try st.arena.buf.download(0, h0); + try tw.arena.buf.download(0, h1); + const at = std.mem.indexOfDiff(u8, h0, h1); + if (st.pos != tw.pos or st.mtp_len != tw.mtp_len or !std.mem.eql(u1, st.cur, tw.cur) or at != null) { + std.debug.print("FAIL shared: stream {d}'s state differs from its twin's (pos {d}/{d}, first byte {?d} of {d})\n", .{ i, st.pos, tw.pos, at, len }); + failed = true; + } + } + if (failed) return error.TestFailed; + check.pass("shared windows at TP={d}: {d} rows of {d} streams over 12 rounds (2 to {d} streams a forward, 1 to 4 rows each, greedy and sampled) equal their lone twins' logits, tokens and drafts byte for byte, and every state its twin's", .{ world, compared, S, S }); +} + +/// A concurrent case's stream: prompt `p` joining at round `at`; `filled`: its prompt a chunk at a time between rounds. +const Join = struct { p: usize, at: u32 = 0, sampled: bool = false, discard_at: ?u32 = null, filled: bool = false }; + +/// A mark hook that keeps the state as the lane host's prompt cache does (cuda_native's snap_save). +const Keeper = struct { + e: *@import("glm_cuda").engine.Engine, + kept: [4]*anyopaque = undefined, + n: usize = 0, + + fn at(ptr: *anyopaque, s: *lanes.Stream, pos: u32) void { + _ = s; + const k: *Keeper = @ptrCast(@alignCast(ptr)); + k.kept[k.n] = @import("glm_cuda").native.snap_save(k.e, null, pos) catch |err| { + std.debug.print("FAIL keeping the state at {d}: {s}\n", .{ pos, @errorName(err) }); + return; + }; + k.n += 1; + } +}; + +/// A cancel hook that answers yes from its `at`-th ask on. +const Cut = struct { + at: u32, + asked: u32 = 0, + + fn check(ptr: *anyopaque) bool { + const cut: *Cut = @ptrCast(@alignCast(ptr)); + cut.asked += 1; + return cut.asked > cut.at; + } +}; + +const Cover = struct { + gpa: std.mem.Allocator, + a: std.mem.Allocator, + le: *lanes.Engine, + failed: bool = false, + + const Ask = struct { + prompt: []const u32, + max_new: u32, + eos: []const u32 = &.{}, + drafts: bool = true, + sampling: ?lanes.Sampling = null, + copies: bool = false, // the server's proposer + cancel: ?*Cut = null, + reuse: lanes.stream.Reuse = .{}, + }; + const Got = struct { ids: []const u32, reason: lanes.stream.Reason, rounds: u64, judged: i64, copied: i64, cached: u32 }; + + /// One request through the shared lanes.Engine, as the server's lane host runs it. + fn run(c: *Cover, ask: Ask) !Got { + var lookup = try lanes.SuffixLookup.init(c.gpa, .{ .min_match = 4 }); + defer lookup.deinit(); + var s = try lanes.Stream.init(c.gpa, .{ .id = "cover", .prompt = ask.prompt, .max_new = ask.max_new, .eos = ask.eos, .drafts = ask.drafts, .sampling = ask.sampling, .proposer = if (ask.copies) lookup.proposer() else null, .cancel_check = if (ask.cancel) |cut| .{ .ptr = cut, .check = Cut.check } else null, .reuse = ask.reuse }); + defer s.deinit(c.gpa); + try c.le.addStream(&s); + while (c.le.activeCount() > 0) try c.le.step(); + return .{ .ids = try c.a.dupe(u32, s.emitted()), .reason = s.reason, .rounds = s.rounds, .judged = lookup.judged_tokens, .copied = lookup.accepted_tokens, .cached = s.cached }; + } + + const Together = struct { ids: []const []const u32, rounds: []const u64, most: usize, shared: u32, between: u32 }; + + /// The streams of `joins` through the rounds together, each joining at its round; how many were live at once. + fn together(c: *Cover, joins: []const Join, ps: []const []const u32, seeds: []const lanes.Sampling, max_new: u32) !Together { + const ss = try c.a.alloc(lanes.Stream, joins.len); + const made = try c.a.alloc(bool, joins.len); + @memset(made, false); + defer for (ss, made) |*s, m| if (m) s.deinit(c.gpa); + var most: usize = 0; + var shared: u32 = 0; + var between: u32 = 0; + var round: u32 = 0; + var filling: std.ArrayList(*lanes.Stream) = .empty; // the lane host's fillOne: the oldest's next chunk a round + defer filling.deinit(c.gpa); + var began = false; + while (true) : (round += 1) { + for (joins, ss, made) |j, *s, *m| if (!m.* and j.at <= round) { + s.* = try lanes.Stream.init(c.gpa, .{ .id = "joined", .prompt = ps[j.p], .max_new = max_new, .sampling = if (j.sampled) seeds[j.p % seeds.len] else null }); + m.* = true; + if (j.filled) try filling.append(c.gpa, s) else try c.le.addStream(s); + }; + if (filling.items.len > 0) { + if (try c.le.fillStream(filling.items[0], !began)) { + _ = filling.orderedRemove(0); + began = false; + } else began = true; + } + for (joins, ss, made) |j, *s, m| if (m and j.discard_at == round and !s.finished) c.le.discard(s); + const live = c.le.activeCount(); + most = @max(most, live); + shared += @intFromBool(live >= 2); + if (live == 0) { + if (std.mem.indexOfScalar(bool, made, false) == null and filling.items.len == 0) break; + continue; + } + between += @intFromBool(began); + try c.le.step(); + } + const ids = try c.a.alloc([]const u32, joins.len); + const rounds = try c.a.alloc(u64, joins.len); + for (ss, ids, rounds) |*s, *i, *r| { + i.* = try c.a.dupe(u32, s.emitted()); + r.* = s.rounds; + } + return .{ .ids = ids, .rounds = rounds, .most = most, .shared = shared, .between = between }; + } + + /// A drafted stream over `prompt`, discarded after `rounds` (null: first rejecting round); false if it ended first. + fn discardAfter(c: *Cover, prompt: []const u32, rounds: ?u32) !bool { + var s = try lanes.Stream.init(c.gpa, .{ .id = "discarded", .prompt = prompt, .max_new = 512 }); + defer s.deinit(c.gpa); + try c.le.addStream(&s); + var n: u32 = 0; + while (c.le.activeCount() > 0) : (n += 1) { + if (rounds) |r| if (n == r) break; + const rejected = s.drafted - s.accepted; + try c.le.step(); + if (rounds == null and s.drafted - s.accepted > rejected) break; + } + if (c.le.activeCount() == 0) return false; + c.le.discard(&s); + return true; + } + + fn same(c: *Cover, what: []const u8, got: Got, want: []const u32, reason: lanes.stream.Reason) !void { + const n = @min(got.ids.len, want.len); + const first = for (got.ids[0..n], want[0..n], 0..) |x, y, i| { + if (x != y) break i; + } else if (got.ids.len != want.len) n else null; + if (first == null and got.reason == reason) { + check.pass("{s}: {d} ids equal, finished by {s} ({d} rounds)", .{ what, want.len, reason.name(), got.rounds }); + return; + } + c.failed = true; + if (first) |i| { + std.debug.print("FAIL {s}: token {d} of {d} is {d}, want {d} (emitted {d}, {s})\n", .{ what, i, want.len, if (i < got.ids.len) got.ids[i] else 0, if (i < want.len) want[i] else 0, got.ids.len, got.reason.name() }); + } else std.debug.print("FAIL {s}: finished by {s}, want {s}\n", .{ what, got.reason.name(), reason.name() }); + } + + fn expect(c: *Cover, ok: bool, comptime fmt: []const u8, args: anytype) !void { + if (ok) return check.pass(fmt, args); + c.failed = true; + std.debug.print("FAIL " ++ fmt ++ "\n", args); + } +}; + +fn seconds(io: std.Io, since: std.Io.Timestamp) f64 { + const now = std.Io.Clock.awake.now(io); + return @as(f64, @floatFromInt(now.toNanoseconds() - since.toNanoseconds())) / 1e9; +} diff --git a/zig/tests/cuda/glm_ref_attn.zig b/zig/tests/cuda/glm_ref_attn.zig new file mode 100644 index 0000000000..b2209959b0 --- /dev/null +++ b/zig/tests/cuda/glm_ref_attn.zig @@ -0,0 +1,160 @@ +//! GLM's MLA attention (chunked softmax, merge; dense, sparse) vs float64; 16/32-head blocks, 1-row windows same bits. +const std = @import("std"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const ref = @import("glm_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; +const latent = glm.latent; + +const lw = 512; +const heads = 32; + +fn partials(r: *Rig, nch: usize, rows: usize) !latent.Partials { + return .{ .po = try r.zeros(nch * rows * heads * lw * 4), .pm = try r.zeros(nch * rows * heads * 4), .pl = try r.zeros(nch * rows * heads * 4) }; +} + +/// float64 o = sum p c / sum p, p = exp(s - max); bound: bf16 out and p (2^-8 each), (512 + 8) 2^-24 scale sum |q c|. +fn reference(r: *Rig, q: []const u16, cd: []const f64, keys: []const []const u32, scale: f64) !struct { want: []f64, tol: []f64 } { + const rows = keys.len; + const want = try r.a.alloc(f64, rows * heads * lw); + const tol = try r.a.alloc(f64, rows * heads * lw); + const qd = try r.a.alloc(f64, lw); + const o = try r.a.alloc(f64, lw); + const spread = try r.a.alloc(f64, lw); + var most_n: usize = 0; + for (keys) |k| most_n = @max(most_n, k.len); + const s = try r.a.alloc(f64, most_n); + const ds = try r.a.alloc(f64, most_n); + for (keys, 0..) |set, row| for (0..heads) |h| { + const base = (row * heads + h) * lw; + for (qd, q[base..][0..lw]) |*d, v| d.* = bf(v); + var most = -std.math.inf(f64); + for (set, 0..) |key, i| { + var dot: f64 = 0; + var abs: f64 = 0; + for (qd, cd[key * lw ..][0..lw]) |a, b| { + dot += a * b; + abs += @abs(a * b); + } + s[i] = dot * scale; + ds[i] = abs * scale * 520 * std.math.pow(f64, 2, -24); + most = @max(most, s[i]); + } + @memset(o, 0); + @memset(spread, 0); + var l: f64 = 0; + for (set, 0..) |key, i| { + const p = @exp(s[i] - most); + l += p; + for (o, spread, cd[key * lw ..][0..lw]) |*od, *sd, v| { + od.* += p * v; + sd.* += p * @abs(v) * (std.math.pow(f64, 2, -8) + ds[i]); + } + } + for (0..lw) |d| { + want[base + d] = o[d] / l; + tol[base + d] = @abs(o[d] / l) * std.math.pow(f64, 2, -8) + spread[d] / l + 1e-30; + } + }; + return .{ .want = want, .tol = tol }; +} + +fn cacheOf(r: *Rig, keys: usize) !struct { dev: u64, f64s: []f64 } { + const c = try r.bfs(keys * lw, 1.0); + const cd = try r.a.alloc(f64, keys * lw); + for (cd, c) |*d, v| d.* = bf(v); + return .{ .dev = try r.dev(std.mem.sliceAsBytes(c)), .f64s = cd }; +} + +/// latent.attention of `rows` rows at `pos` (keys 0 .. pos + r each) with head block hb into a fresh output. +fn attend(r: *Rig, set: glm.glue.Set, q: u64, cache: u64, pos: usize, rows: usize, scale: f32, hb: usize) !u64 { + const nch = latent.chunksFor(pos + rows); + const out = try r.zeros(rows * heads * lw * 2); + const p0: i32 = @intCast(pos); + const at = try r.dev(std.mem.asBytes(&p0)); + try latent.attention(set, r.s, q, cache, at, try partials(r, nch, rows), out, try r.zeros(4), rows, heads, lw, nch, scale, hb); + return out; +} + +/// Row `row` of a [rows][heads][lw] bf16 output alone equals that row of the wide call's bytes. +fn sameRow(r: *Rig, name: []const u8, wide: []const u8, one: u64, row: usize) !void { + const n = heads * lw * 2; + try check.sameBytes(name, try r.back(one, n), wide[row * n ..][0..n]); +} + +/// Dense causal attention on 1,500 keys: 70 rows at 1,430 vs float64 (32-head), 16-head blocks' bytes, one-row windows. +pub fn dense(r: *Rig, set: glm.glue.Set) !void { + const pos = 1430; + const rows = 70; + const scale: f32 = 0.125; + const q = try r.bfs(rows * heads * lw, 1.0); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const c = try cacheOf(r, pos + rows); + const wide = try attend(r, set, dq, c.dev, pos, rows, scale, latent.heads_tile_wide); + const keys = try r.a.alloc([]const u32, rows); + for (keys, 0..) |*k, row| { + const all = try r.a.alloc(u32, pos + row + 1); + for (all, 0..) |*v, i| v.* = @intCast(i); + k.* = all; + } + const ref64 = try reference(r, q, c.f64s, keys, scale); + var name: [128]u8 = undefined; + try ref.closeCols(r, try std.fmt.bufPrint(&name, "latent attention {d} rows at {d}, {d} heads, 32-head blocks", .{ rows, pos, heads }), wide, false, ref64.want, ref64.tol, lw); + const narrow = try attend(r, set, dq, c.dev, pos, rows, scale, latent.heads_tile); + try ref.sameRows(r, "latent attention: 16-head blocks equal the 32-head blocks", wide, narrow, rows * heads * lw * 2); + const all = try r.back(wide, rows * heads * lw * 2); + for ([_]usize{ 0, 37, 69 }) |row| { + const one = try attend(r, set, dq + row * heads * lw * 2, c.dev, pos + row, 1, scale, latent.headBlock(1)); + try sameRow(r, try std.fmt.bufPrint(&name, "latent attention: row {d} alone equals its row in the {d}-row call", .{ row, rows }), all, one, row); + } + check.pass("latent attention over {d} keys: within float64 tolerance, both head blocks and one-row windows equal", .{pos + rows}); +} + +/// Sparse attention on indexer lists (<= 1,100 of 6,000 keys): 70 rows vs float64, 1-row windows, count 0 left alone. +pub fn sparse(r: *Rig, set: glm.glue.Set) !void { + const keys = 6000; + const rows = 70; + const w = 1100; + const scale: f32 = 0.125; + const q = try r.bfs(rows * heads * lw, 1.0); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const c = try cacheOf(r, keys); + const tok = try r.a.alloc(i32, rows * w); + @memset(tok, -1); + const cnt = try r.a.alloc(i32, rows); + const sets = try r.a.alloc([]const u32, rows); + for (sets, 0..) |*set_, row| { + const pick = try r.a.alloc(u32, 1 + r.rng.uintLessThan(usize, w)); + for (pick) |*v| v.* = r.rng.uintLessThan(u32, keys); + std.mem.sort(u32, pick, {}, std.sort.asc(u32)); + var n: usize = 0; + for (pick) |v| if (n == 0 or pick[n - 1] != v) { + pick[n] = v; + n += 1; + }; + set_.* = pick[0..n]; + cnt[row] = @intCast(n); + for (pick[0..n], tok[row * w ..][0..n]) |v, *t| t.* = @intCast(v); + } + const dtok = try r.dev(std.mem.sliceAsBytes(tok)); + const dcnt = try r.dev(std.mem.sliceAsBytes(cnt)); + const out = try r.zeros(rows * heads * lw * 2); + const nch = latent.chunksFor(w); + try latent.sparseAttention(set, r.s, dq, c.dev, dtok, dcnt, w, try partials(r, nch, rows), out, rows, heads, lw, scale); + const ref64 = try reference(r, q, c.f64s, sets, scale); + var name: [128]u8 = undefined; + try ref.closeCols(r, try std.fmt.bufPrint(&name, "latent sparse attention {d} rows, up to {d} of {d} keys", .{ rows, w, keys }), out, false, ref64.want, ref64.tol, lw); + const all = try r.back(out, rows * heads * lw * 2); + for ([_]usize{ 0, 37, 69 }) |row| { + const one = try r.zeros(heads * lw * 2); + try latent.sparseAttention(set, r.s, dq + row * heads * lw * 2, c.dev, dtok + row * w * 4, dcnt + row * 4, w, try partials(r, nch, 1), one, 1, heads, lw, scale); + try sameRow(r, try std.fmt.bufPrint(&name, "latent sparse attention: row {d} alone equals its row in the {d}-row call", .{ row, rows }), all, one, row); + } + const none: i32 = 0; + const left = try r.zeros(heads * lw * 2); + try latent.sparseAttention(set, r.s, dq, c.dev, dtok, try r.dev(std.mem.asBytes(&none)), w, try partials(r, nch, 1), left, 1, heads, lw, scale); + for (try r.back(left, heads * lw * 2)) |b| try check.expect(b == 0xff, "latent sparse attention: a row with count 0 is left alone", .{}); + check.pass("latent sparse attention over up to {d} tokens: within float64 tolerance, one-row windows equal, count 0 left alone", .{w}); +} diff --git a/zig/tests/cuda/glm_ref_hc.zig b/zig/tests/cuda/glm_ref_hc.zig new file mode 100644 index 0000000000..bbbd26707d --- /dev/null +++ b/zig/tests/cuda/glm_ref_hc.zig @@ -0,0 +1,71 @@ +//! GLM's hyper-connection dots (hcPre: 24 dots + squares a 1,024 block) vs float64; hc_partial8 = hc_partial bytes. +const std = @import("std"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const ref = @import("glm_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; +const glue = glm.glue; + +const d = 4096; +const s = 4; +const wide = s * d; +const mixes = 24; + +const Outs = struct { out: u64, xs: u64, post: u64, comb: u64, part: u64 }; + +fn pre(r: *Rig, set: glue.Set, x: u64, w: [4]u64, rows: usize) !Outs { + const o: Outs = .{ .out = try r.zeros(rows * d * 2), .xs = try r.zeros(rows * (d / 64) * 4), .post = try r.zeros(rows * s * 4), .comb = try r.zeros(rows * s * s * 4), .part = try r.zeros(glue.hcPartBytes(rows)) }; + try glue.hcPre(set, r.s, x, w[0], w[1], w[2], w[3], o.out, o.xs, o.post, o.comb, o.part, 1e-6, 1e-6, 20, rows, d, s); + return o; +} + +/// 70 rows (hc_partial) vs float64 within (1,024 + 8) 2^-24 of sum |terms|; 1, 5, 63 rows (hc_partial8) same bytes. +pub fn mixing(r: *Rig, set: glue.Set) !void { + const rows = 70; + const x = try r.bfs(rows * wide, 1.0); + const fnw = try r.bfs(mixes * wide, 0.02); + const base = try r.a.alloc(f32, mixes); + for (base) |*b| b.* = (r.rng.float(f32) * 2 - 1) * 0.5; + const scale = [3]f32{ 0.7, 0.9, 1.1 }; + const nw = try r.bfs(d, 1.0); + const w = [4]u64{ try r.dev(std.mem.sliceAsBytes(fnw)), try r.dev(std.mem.sliceAsBytes(base)), try r.dev(std.mem.asBytes(&scale)), try r.dev(std.mem.sliceAsBytes(nw)) }; + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const all = try pre(r, set, dx, w, rows); + + const nb = glue.hc_blocks; + const kb = wide / nb; + const raw = std.mem.bytesAsSlice(f32, try r.back(all.part, glue.hcPartBytes(rows))); + const col = try r.a.alloc(f64, mixes + 1); + @memset(col, 0); + const bound = @as(f64, kb + 8) * std.math.pow(f64, 2, -24); + for (0..rows) |row| for (0..nb) |b| { + const got = raw[(row * nb + b) * 32 ..][0..32]; + const xs = x[row * wide + b * kb ..][0..kb]; + for (0..mixes + 1) |m| { + var sum: f64 = 0; + var abs: f64 = 0; + for (xs, 0..) |xv, j| { + const a = bf(xv); + const t = if (m == mixes) a * a else a * bf(fnw[m * wide + b * kb + j]); + sum += t; + abs += @abs(t); + } + const q = @abs(got[m] - sum) / (abs * bound + 1e-30); + if (!(q <= col[m])) col[m] = if (std.math.isNan(q)) std.math.inf(f64) else q; + } + }; + var name: [128]u8 = undefined; + try ref.worstCol(r, try std.fmt.bufPrint(&name, "hc mixing {d} rows x {d} K blocks: 24 dots and the squares (column 24)", .{ rows, nb }), col); + + const full = [5]u64{ all.out, all.xs, all.post, all.comb, all.part }; + const per = [5]usize{ d * 2, (d / 64) * 4, s * 4, s * s * 4, glue.hcPartBytes(1) }; + const what = [5][]const u8{ "normed rows", "group sums", "post", "comb", "partials" }; + for ([_]usize{ 1, 5, 63 }) |m| { + const o = try pre(r, set, dx, w, m); + for ([5]u64{ o.out, o.xs, o.post, o.comb, o.part }, full, per, what) |got, want, n, label| + try ref.sameRows(r, try std.fmt.bufPrint(&name, "hc pre: {d} rows' {s} equal the {d}-row call's", .{ m, label, rows }), want, got, m * n); + } + check.pass("hc mixing: dots and squares within float64 tolerance, decode windows' partials and outputs equal the prompt rows'", .{}); +} diff --git a/zig/tests/cuda/glm_ref_index.zig b/zig/tests/cuda/glm_ref_index.zig new file mode 100644 index 0000000000..b60dc3c9ea --- /dev/null +++ b/zig/tests/cuda/glm_ref_index.zig @@ -0,0 +1,98 @@ +//! GLM's DSA indexer: pool scores vs float64 (worst column), radix top-512 vs the host rule on device scores (exact). +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const ref = @import("glm_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; +const sp = glm.sparse; + +const heads = 32; +const d = 128; + +/// sparse._order_key: a score's bits as a u32 whose unsigned order is the scores' (-0 counted as +0). +fn orderKey(s: f32) u32 { + const b: i32 = @bitCast(s + 0.0); + return @bitCast(b ^ ((b >> 31) | std.math.minInt(i32))); +} + +/// sparse.selectTokens for 96 rows at 6,000 (1,500 to 1,524 visible pools in a 2,048-pool bucket). +pub fn selection(r: *Rig, set: glm.glue.Set, ops: cuda.Module) !void { + const pos = 6000; + const rows = 96; + const cap = 4096; + const np = sp.poolBucket(pos, rows, cap); + const cols = @min(np, (sp.scoreColumns(pos, rows, np) + 63) / 64 * 64); + const qi = try r.bfs(rows * heads * d, 1.0); + const wts = try r.bfs(rows * heads, 1.0); + const pk = try r.bfs(np * d, 1.0); + const p0: i32 = pos; + const s: sp.Scratch = .{ .scores = try r.zeros(rows * np * 4), .ls = try r.zeros(rows * sp.select_list * 4), .li = try r.zeros(rows * sp.select_list * 4), .pools = try r.zeros(rows * sp.topk_pools * 8), .at = try r.zeros(4) }; + const tokens = try r.zeros(rows * sp.width * 4); + const counts = try r.zeros(rows * 4); + try sp.selectTokens(set, ops, r.s, try r.dev(std.mem.sliceAsBytes(qi)), try r.dev(std.mem.sliceAsBytes(wts)), heads, try r.dev(std.mem.sliceAsBytes(pk)), try r.dev(std.mem.asBytes(&p0)), pos, rows, heads, d, cap, s, tokens, counts); + const raw = try r.back(s.scores, rows * np * 4); + const got = std.mem.bytesAsSlice(f32, raw); + + // float64: sum_h w_h relu(scale q_h . k_p); bound: (128 + 8) 2^-24 scale sum |q k| a dot, 40 2^-24 sum |w_h relu|. + const scale = 1 / @sqrt(@as(f64, d)); + const wscale: f64 = sp.headScale(heads); + const kd = try r.a.alloc(f64, np * d); + for (kd, pk) |*v, b| v.* = bf(b); + const col = try r.a.alloc(f64, cols); + @memset(col, 0); + for (0..rows) |row| { + const npool = (pos + row + 1) / 4; + for (0..cols) |p| { + const g = got[row * np + p]; + if (p >= npool) { + try check.expect(g == -std.math.inf(f32), "score row {d} pool {d}: past the visible pools, -inf (got {d})", .{ row, p, g }); + continue; + } + var sum: f64 = 0; + var bound: f64 = 0; + var spread: f64 = 0; + for (0..heads) |h| { + const w = bf(wts[row * heads + h]) * wscale; + var dot: f64 = 0; + var abs: f64 = 0; + for (0..d) |i| { + const t = bf(qi[(row * heads + h) * d + i]) * kd[p * d + i]; + dot += t; + abs += @abs(t); + } + sum += w * @max(dot * scale, 0); + spread += @abs(w) * @max(dot * scale, 0); + bound += @abs(w) * abs * scale * 136 * std.math.pow(f64, 2, -24); + } + const q = @abs(g - sum) / (bound + spread * 40 * std.math.pow(f64, 2, -24) + 1e-30); + if (!(q <= col[p])) col[p] = if (std.math.isNan(q)) std.math.inf(f64) else q; + } + } + var name: [128]u8 = undefined; + try ref.worstCol(r, try std.fmt.bufPrint(&name, "index scores {d} rows at {d}, {d} heads x {d}, {d} pools", .{ rows, pos, heads, d, cols }), col); + + // Top 512 of each row's first max(visible, 512) pools (at most np) by order key, ties to the lower pool, ascending. + const picks = std.mem.bytesAsSlice(i64, try r.back(s.pools, rows * sp.topk_pools * 8)); + const order = try r.a.alloc(u32, np); + for (0..rows) |row| { + const vis = (pos + row + 1) / 4; + const limit = @min(np, @max(vis, sp.topk_pools)); + const keys = try r.a.alloc(u32, limit); + for (keys, 0..) |*k, i| k.* = orderKey(got[row * np + i]); + const idx = order[0..limit]; + for (idx, 0..) |*v, i| v.* = @intCast(i); + std.mem.sort(u32, idx, keys, struct { + fn lt(k: []const u32, a: u32, b: u32) bool { + return k[a] > k[b] or (k[a] == k[b] and a < b); + } + }.lt); + const n = @min(limit, sp.topk_pools); + std.mem.sort(u32, idx[0..n], {}, std.sort.asc(u32)); + for (idx[0..n], picks[row * sp.topk_pools ..][0..n], 0..) |want, have, i| + try check.expect(have == want, "select row {d} pick {d}: pool {d}, the host's rule gives {d}", .{ row, i, have, want }); + } + check.pass("DSA indexer {d} rows at {d}: scores within float64 tolerance, top-{d} pools equal the host's rule exactly", .{ rows, pos, sp.topk_pools }); +} diff --git a/zig/tests/cuda/glm_ref_kda.zig b/zig/tests/cuda/glm_ref_kda.zig new file mode 100644 index 0000000000..67c26575ae --- /dev/null +++ b/zig/tests/cuda/glm_ref_kda.zig @@ -0,0 +1,148 @@ +//! GLM's KDA vs a float64 recurrence rounding where the kernel does; fused = three-kernel bytes; replay rebuilds state. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const ref = @import("glm_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; +const bf16 = qr.helpers.bf16Of; + +const kda = glm.kda; +const heads = 32; +const n = heads * kda.dk; +const dk = kda.dk; +const p_stride = 3 * n + 64; +const lower = -5.0; +const eps = 1e-6; + +fn floats(r: *Rig, count: usize, lo: f32, hi: f32) ![]f32 { + const out = try r.a.alloc(f32, count); + for (out) |*v| v.* = lo + r.rng.float(f32) * (hi - lo); + return out; +} + +/// The value rounded to bf16, as a float64. +fn rb(x: f64) f64 { + return bf(bf16(@floatCast(x))); +} + +fn sigmoid(x: f64) f64 { + return 1 / (1 + @exp(-x)); +} + +const Inputs = struct { p: []u16, a: []u16, g: []u16, cs: []u16, cw: []f32, state: []f32, a_log: []f32, dt_bias: []f32, norm: []u16 }; + +/// float64 chain a head: conv+silu, L2 q/k, decay, bf16 beta, delta rule, y = bf16(S q), RMS norm, gate. Out [rows][n]. +fn reference(r: *Rig, in: Inputs, rows: usize) ![]f64 { + const out = try r.a.alloc(f64, rows * n); + const s = try r.a.alloc(f64, dk * dk); + var q: [dk]f64 = undefined; + var k: [dk]f64 = undefined; + var v: [dk]f64 = undefined; + var g: [dk]f64 = undefined; + var y: [dk]f64 = undefined; + const cc = 3 * n; + for (0..heads) |h| { + for (s, in.state[h * dk * dk ..][0 .. dk * dk]) |*d, x| d.* = x; + const rate = @exp(@as(f64, in.a_log[h])); + for (0..rows) |row| { + for (0..3) |part| for (0..dk) |i| { + const c = part * n + h * dk + i; + var acc: f64 = 0; + for (0..4) |tap| { + const at = row + tap; + const x = if (at < 3) bf(in.cs[at * cc + c]) else bf(in.p[(at - 3) * p_stride + c]); + acc += @as(f64, in.cw[c * 4 + tap]) * x; + } + const act = rb(acc * sigmoid(acc)); + switch (part) { + 0 => q[i] = act, + 1 => k[i] = act, + else => v[i] = act, + } + }; + for ([_]*[dk]f64{ &q, &k }, [_]f64{ 1 / @sqrt(@as(f64, dk)), 1 }) |x, extra| { + var ss: f64 = 0; + for (x) |e| ss += e * e; + const inv = 1 / @sqrt(ss + 1e-6) * extra; + for (x) |*e| e.* *= inv; + } + for (&g, 0..) |*e, i| e.* = @exp(lower * sigmoid(rate * (bf(in.a[row * n + h * dk + i]) + in.dt_bias[h * dk + i]))); + const beta = rb(sigmoid(bf(in.p[row * p_stride + 3 * n + h]))); + for (0..dk) |j| { + const srow = s[j * dk ..][0..dk]; + var kv: f64 = 0; + for (srow, g, k) |*e, gi, ki| { + e.* *= gi; + kv += e.* * ki; + } + const delta = (v[j] - kv) * beta; + for (srow, k) |*e, ki| e.* += ki * delta; + } + var ss: f64 = 0; + for (&y, 0..) |*e, j| { + var o: f64 = 0; + for (s[j * dk ..][0..dk], q) |e2, qi| o += e2 * qi; + e.* = rb(o); + ss += e.* * e.*; + } + const rinv = 1 / @sqrt(ss / dk + eps); + for (y, 0..) |e, t| out[row * n + h * dk + t] = bf(in.norm[t]) * e * rinv * sigmoid(bf(in.g[row * n + h * dk + t])); + } + } + return out; +} + +/// Fused and three-kernel chains, each replayed; fused out vs float64 within a bf16 ulp (boundaries) + 2^-8 of row max. +fn window(r: *Rig, ks: *const kda.Kernels, rows: usize) !void { + const in: Inputs = .{ + .p = try r.bfs(rows * p_stride, 1.0), // q | k | v, then beta's logits at 3 n + .a = try r.bfs(rows * n, 2.0), + .g = try r.bfs(rows * n, 2.0), + .cs = try r.bfs(3 * 3 * n, 1.0), + .cw = try floats(r, 3 * n * 4, -0.5, 0.5), + .state = try floats(r, heads * kda.dv * dk, -0.05, 0.05), + .a_log = try floats(r, heads, -1, 1), + .dt_bias = try floats(r, n, -0.5, 0.5), + .norm = try r.bfs(dk, 1.0), + }; + const d = [_]u64{ + try r.dev(std.mem.sliceAsBytes(in.p)), try r.dev(std.mem.sliceAsBytes(in.a)), try r.dev(std.mem.sliceAsBytes(in.g)), + try r.dev(std.mem.sliceAsBytes(in.cs)), try r.dev(std.mem.sliceAsBytes(in.cw)), try r.dev(std.mem.sliceAsBytes(in.state)), + try r.dev(std.mem.sliceAsBytes(in.a_log)), try r.dev(std.mem.sliceAsBytes(in.dt_bias)), try r.dev(std.mem.sliceAsBytes(in.norm)), + }; + const ch: kda.Chain = .{ .p = d[0], .p_stride = p_stride, .b_off = 3 * n, .a = d[1], .a_stride = n, .g = d[2], .g_stride = n, .conv_state = d[3], .conv_w = d[4], .state_in = d[5], .a_log = d[6], .dt_bias = d[7], .norm_w = d[8], .eps = eps, .lower = lower, .heads = heads }; + const sizes = [_]usize{ rows * n * 2, heads * kda.dv * dk * 4, rows * n * 4, rows * n * 2, rows * n * 4, rows * heads * 4 }; + const what = [_][]const u8{ "out", "state", "saved k", "saved v", "saved decay", "saved beta" }; + var runs: [2][6]u64 = undefined; + for (&runs, [_]bool{ false, true }) |*o, wide| { + for (o, sizes) |*b, size| b.* = try r.zeros(size); + const w: ?kda.Wide = if (wide) .{ .q = try r.zeros(rows * n * 4), .y = try r.zeros(rows * n * 2) } else null; + try kda.chain(ks, r.s, ch, rows, .{ .out = o[0], .state_out = o[1], .k = o[2], .v = o[3], .g = o[4], .b = o[5] }, w); + const replayed = try r.zeros(sizes[1]); + try kda.replay(ks, r.s, d[5], o[2], o[3], o[4], o[5], rows, heads, replayed); + try ref.sameRows(r, if (wide) "KDA three-kernel chain: the replay rebuilds its state" else "KDA fused chain: the replay rebuilds its state", o[1], replayed, sizes[1]); + } + var name: [96]u8 = undefined; + for (runs[0], runs[1], sizes, what) |fused, three, size, label| + try ref.sameRows(r, try std.fmt.bufPrint(&name, "KDA {d} rows: the three-kernel chain's {s} equal the fused chain's", .{ rows, label }), fused, three, size); + const want = try reference(r, in, rows); + const tol = try r.a.alloc(f64, want.len); + for (0..rows) |row| { + var most: f64 = 0; + for (want[row * n ..][0..n]) |e| most = @max(most, @abs(e)); + for (tol[row * n ..][0..n], want[row * n ..][0..n]) |*t, e| t.* = @abs(e) * std.math.pow(f64, 2, -8) + most * std.math.pow(f64, 2, -8) + 1e-30; + } + try ref.closeCols(r, try std.fmt.bufPrint(&name, "KDA {d} rows x {d} heads, fused chain against float64 (column = head x 128 + t)", .{ rows, heads }), runs[0][0], false, want, tol, n); +} + +/// `glm-ref`'s KDA: 70 rows (a prompt window) and 3 rows (a decode window) on 32 heads. +pub fn chains(r: *Rig) !void { + var mod = try cuda.Module.load(r.gpu.d, cuda.kernels.glm_kda); + defer mod.unload(); + const ks = try kda.Kernels.load(mod); + for ([_]usize{ 70, 3 }) |rows| try window(r, &ks, rows); + check.pass("KDA {d} heads: within float64 tolerance, fused and three-kernel chains equal at 70 and 3 rows, replays rebuild the state", .{heads}); +} diff --git a/zig/tests/cuda/glm_ref_tests.zig b/zig/tests/cuda/glm_ref_tests.zig new file mode 100644 index 0000000000..b0e7598990 --- /dev/null +++ b/zig/tests/cuda/glm_ref_tests.zig @@ -0,0 +1,207 @@ +//! GLM's projections in the kernels' reduction order vs float64 (worst column reported); narrower calls = widest bits. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; + +const b16 = glm.b16; +const glue = glm.glue; + +/// RESULT for one check from each column's worst fraction of its tolerance: the worst column named, all within 1. +pub fn worstCol(r: *Rig, name: []const u8, col: []const f64) !void { + var most: f64 = 0; + var at: usize = 0; + for (col, 0..) |q, j| if (q > most) { + most = q; + at = j; + }; + r.worst = @max(r.worst, most); + std.debug.print("RESULT {s}: worst column {d} at {e:.3} of its tolerance\n", .{ name, at, most }); + try check.expect(most <= 1, "{s} within its float64 tolerance", .{name}); +} + +/// Every element within its tolerance of `want`; RESULT names the worst column (n a row) and its worst fraction. +pub fn closeCols(r: *Rig, name: []const u8, ptr: u64, f32_out: bool, want: []const f64, tol: []const f64, n: usize) !void { + const raw = try r.back(ptr, want.len * @as(usize, if (f32_out) 4 else 2)); + const col = try r.a.alloc(f64, n); + @memset(col, 0); + for (want, tol, 0..) |w, t, i| { + const g: f64 = if (f32_out) std.mem.bytesToValue(f32, raw[4 * i ..][0..4]) else bf(std.mem.bytesToValue(u16, raw[2 * i ..][0..2])); + const q = @abs(g - w) / t; + if (!(q <= col[i % n])) col[i % n] = if (std.math.isNan(q)) std.math.inf(f64) else q; + } + try worstCol(r, name, col); +} + +/// The first `bytes` of a narrower call equal the widest call's. +pub fn sameRows(r: *Rig, name: []const u8, wide: u64, narrow: u64, bytes: usize) !void { + try check.sameBytes(name, try r.back(narrow, bytes), try r.back(wide, bytes)); +} + +/// y = x w^T in float64 over bf16 rows; bound: output rounding (2^-8 |y| for bf16) + fp32 (K + 8) 2^-24 of sum |x w|. +pub fn dots(r: *Rig, x: []const u16, w: []const u16, m: usize, n: usize, k: usize, f32_out: bool) !struct { want: []f64, tol: []f64 } { + const want = try r.a.alloc(f64, m * n); + const tol = try r.a.alloc(f64, m * n); + const wd = try r.a.alloc(f64, n * k); + for (wd, w) |*d, v| d.* = bf(v); + const xd = try r.a.alloc(f64, k); + for (0..m) |row| { + for (xd, x[row * k ..][0..k]) |*d, v| d.* = bf(v); + for (0..n) |j| { + var sum: f64 = 0; + var abs: f64 = 0; + for (xd, wd[j * k ..][0..k]) |a, b| { + sum += a * b; + abs += @abs(a * b); + } + want[row * n + j] = sum; + const chain = abs * @as(f64, @floatFromInt(k + 8)) * std.math.pow(f64, 2, -24); + tol[row * n + j] = (if (f32_out) 0 else @abs(sum) * std.math.pow(f64, 2, -8)) + chain + 1e-30; + } + } + return .{ .want = want, .tol = tol }; +} + +/// bf16 rows uploaded one element on, so they start 2 bytes past 16: every launch takes the unstaged kernels. +fn shifted(r: *Rig, x: []const u16) !u64 { + const s = try r.a.alloc(u16, x.len + 1); + s[0] = 0; + @memcpy(s[1..], x); + return (try r.dev(std.mem.sliceAsBytes(s))) + 2; +} + +const Shape = struct { n: usize, k: usize }; +const decode_rows = [_]usize{ 1, 3, 16, 17, 40, 127 }; +const wide_rows = 300; + +/// b16.matmul at 300 rows vs float64; unaligned x (bmm_rows), decode row counts (bmm_dtile, bmm_decode) vs its bytes. +fn matmuls(r: *Rig, set: glue.Set, s: Shape) !void { + const m = wide_rows; + const x = try r.bfs(m * s.k, 1.0); + const w = try r.bfs(s.n * s.k, 0.05); + const xs = [2]u64{ try r.dev(std.mem.sliceAsBytes(x)), try shifted(r, x) }; + const dw = try r.dev(std.mem.sliceAsBytes(w)); + const part = try r.zeros(b16.partBytes(decode_rows[decode_rows.len - 1], s.n, s.k)); + var name: [128]u8 = undefined; + for ([_]bool{ false, true }) |wide| { + const eb: usize = if (wide) 4 else 2; + const ref = try dots(r, x, w, m, s.n, s.k, wide); + const full = try r.zeros(m * s.n * eb); + try b16.matmul(set, r.s, xs[0], s.k, dw, full, part, m, s.n, s.k, wide); + const kind = if (wide) "fp32" else "bf16"; + try closeCols(r, try std.fmt.bufPrint(&name, "b16 matmul {d} rows x [{d}, {d}] (K in {d} slices), {s} out", .{ m, s.n, s.k, b16.splitK(s.n, s.k), kind }), full, wide, ref.want, ref.tol, s.n); + const o = try r.zeros(m * s.n * eb); + try b16.matmul(set, r.s, xs[1], s.k, dw, o, part, m, s.n, s.k, wide); + try sameRows(r, try std.fmt.bufPrint(&name, "b16 matmul [{d}, {d}] {s}: unaligned rows equal the staged tiles'", .{ s.n, s.k, kind }), full, o, m * s.n * eb); + for (decode_rows) |dm| for (xs, [_][]const u8{ "aligned", "unaligned" }) |xp, how| { + const d = try r.zeros(dm * s.n * eb); + try b16.matmul(set, r.s, xp, s.k, dw, d, part, dm, s.n, s.k, wide); + try sameRows(r, try std.fmt.bufPrint(&name, "b16 matmul [{d}, {d}] {s}: {d} {s} rows equal the {d}-row call's", .{ s.n, s.k, kind, dm, how, m }), full, d, dm * s.n * eb); + }; + } + check.pass("b16 matmul [{d}, {d}] ({d} slices): bf16 and fp32 out within float64 tolerance, every entry and row count's bits equal", .{ s.n, s.k, b16.splitK(s.n, s.k) }); +} + +/// glue.router (8 K slices) at 300 rows vs float64, fewer rows and unaligned x vs its bytes; glue.select vs float64. +fn routers(r: *Rig, set: glue.Set) !void { + const m = wide_rows; + const d = 4096; + const ne = 256; + const top_k = 8; + const scale: f32 = 2.5; + const x = try r.bfs(m * d, 1.0); + const w = try r.bfs(ne * d, 0.02); + const xs = [2]u64{ try r.dev(std.mem.sliceAsBytes(x)), try shifted(r, x) }; + const dw = try r.dev(std.mem.sliceAsBytes(w)); + const part = try r.zeros(glue.router_ks * m * ne * 4); + const full = try r.zeros(m * ne * 4); + try glue.router(set, r.s, xs[0], d, dw, full, part, m, d, ne); + const ref = try dots(r, x, w, m, ne, d, true); + var name: [128]u8 = undefined; + try closeCols(r, try std.fmt.bufPrint(&name, "router {d} rows x [{d}, {d}] (K in {d} slices)", .{ m, ne, d, glue.router_ks }), full, true, ref.want, ref.tol, ne); + for ([_]usize{ 1, 3, 63, 64, 100, m }) |dm| for (xs, [_][]const u8{ "aligned", "unaligned" }) |xp, how| { + const o = try r.zeros(dm * ne * 4); + try glue.router(set, r.s, xp, d, dw, o, part, dm, d, ne); + try sameRows(r, try std.fmt.bufPrint(&name, "router: {d} {s} rows equal the {d}-row call's", .{ dm, how, m }), full, o, dm * ne * 4); + }; + check.pass("router [{d}, {d}]: within float64 tolerance, router_tile and router_part bits equal at every row count", .{ ne, d }); + + const bias = try r.a.alloc(f32, ne); + for (bias) |*b| b.* = (r.rng.float(f32) * 2 - 1) * 0.05; + const pick = try r.zeros(m * (top_k + 1) * 4); + const wts = try r.zeros(m * (top_k + 1) * 4); + try glue.select(set, r.s, full, try r.dev(std.mem.sliceAsBytes(bias)), pick, wts, m, top_k, ne, scale, true); + const logits = try r.back(full, m * ne * 4); + const got_pick = try r.back(pick, m * (top_k + 1) * 4); + const got_wts = try r.back(wts, m * (top_k + 1) * 4); + var swaps: usize = 0; + var most: f64 = 0; + const score = try r.a.alloc(f64, ne); + const order = try r.a.alloc(usize, ne); + for (0..m) |row| { + for (score, 0..) |*sc, e| sc.* = 1 / (1 + @exp(-@as(f64, std.mem.bytesToValue(f32, logits[4 * (row * ne + e) ..][0..4])))); + for (order, 0..) |*o, e| o.* = e; + const Ctx = struct { s: []const f64, b: []const f32 }; + std.mem.sort(usize, order, Ctx{ .s = score, .b = bias }, struct { + fn lt(c: Ctx, a: usize, b: usize) bool { + const ca = c.s[a] + c.b[a]; + const cb = c.s[b] + c.b[b]; + return ca > cb or (ca == cb and a < b); + } + }.lt); + var total: f64 = 0; + for (0..top_k) |i| total += score[@intCast(std.mem.bytesToValue(i32, got_pick[4 * (row * (top_k + 1) + i) ..][0..4]))]; + for (0..top_k + 1) |i| { + const at = 4 * (row * (top_k + 1) + i); + const id: usize = @intCast(std.mem.bytesToValue(i32, got_pick[at..][0..4])); + const wt: f64 = std.mem.bytesToValue(f32, got_wts[at..][0..4]); + if (i == top_k) { + try check.expect(id == ne and wt == 1, "select row {d}: the shared expert last at weight 1", .{row}); + continue; + } + if (id != order[i]) { // a near tie in float64 may swap: the scores' approximate exp is within 2^-20 + const gap = @abs(score[id] + bias[id] - score[order[i]] - bias[order[i]]); + try check.expect(gap <= std.math.pow(f64, 2, -20), "select row {d} pick {d}: {d} where float64 ranks {d} (gap {e})", .{ row, i, id, order[i], gap }); + swaps += 1; + } + const want = score[id] / total * scale; + most = @max(most, @abs(wt - want) / (@abs(want) * std.math.pow(f64, 2, -16))); + } + } + r.worst = @max(r.worst, most); + std.debug.print("RESULT select {d} rows, top {d} of {d}: weights at {e:.3} of 2^-16 relative, {d} near-tie swaps\n", .{ m, top_k, ne, most, swaps }); + try check.expect(most <= 1, "select weights within 2^-16 of float64", .{}); + check.pass("select top-{d} of {d}: picks follow the float64 ranking (near ties aside), weights within float64 tolerance", .{ top_k, ne }); +} + +/// `glm-ref`: the checks above on synthetic inputs (seeded), no checkpoint and no golden. +pub fn run(gpu: Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(593); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer mod.unload(); + const ks = try glue.Kernels.load(mod); + const set: glue.Set = .{ .glue = &ks }; + for ([_]Shape{ .{ .n = 512, .k = 4096 }, .{ .n = 1280, .k = 1536 }, .{ .n = 1024, .k = 128 } }) |sh| try matmuls(&r, set, sh); + try routers(&r, set); + try @import("glm_ref_attn.zig").dense(&r, set); + try @import("glm_ref_attn.zig").sparse(&r, set); + try @import("glm_ref_hc.zig").mixing(&r, set); + try @import("glm_ref_kda.zig").chains(&r); + var ops = try cuda.Module.load(gpu.d, cuda.kernels.glm_ops); + defer ops.unload(); + try @import("glm_ref_index.zig").selection(&r, set, ops); + check.pass("glm-ref: every projection within {e:.3} of its float64 tolerance", .{r.worst}); +} diff --git a/zig/tests/cuda/glm_synth.zig b/zig/tests/cuda/glm_synth.zig new file mode 100644 index 0000000000..2f38f1fbd3 --- /dev/null +++ b/zig/tests/cuda/glm_synth.zig @@ -0,0 +1,211 @@ +//! `glm-synth OUT_DIR`: a 2-layer + MTP GLM-5.3-Flash checkpoint, dc.json; run TF_GLM_LAYERS=2. +const std = @import("std"); +const core = @import("core"); +const glm = @import("glm_cuda"); +const check = @import("check.zig"); + +const st = core.safetensors; +const New = core.shard_edit.New; + +const layers = 2; +const experts = 16; +const vocab = 8192; +const hidden = 4096; +/// The ModelOpt block nvidia/GLM-5.3-Flash-NVFP4 writes, which the CUDA family admits (modelopt.admit). +const modelopt_block = ",\"quantization_config\": {\"quant_method\": \"modelopt\", \"quant_algo\": \"NVFP4\", " ++ + "\"config_groups\": {\"group_0\": {\"weights\": {\"num_bits\": 4, \"type\": \"float\", \"group_size\": 16}}}}"; +const P = "model.language_model."; + +const Out = struct { + a: std.mem.Allocator, + rng: std.Random, + list: std.ArrayList(New) = .empty, + + fn put(o: *Out, name: []const u8, dtype: st.DType, shape: []const usize, bytes: []const u8) !void { + try o.list.append(o.a, .{ .name = name, .dtype = dtype, .shape = try o.a.dupe(usize, shape), .bytes = bytes }); + } + + /// bf16 values uniform in [-scale, scale). + fn bf16(o: *Out, name: []const u8, shape: []const usize, scale: f32) !void { + const v = try o.a.alloc(u16, count(shape)); + for (v) |*x| x.* = bf16Of((o.rng.float(f32) * 2 - 1) * scale); + try o.put(name, .bf16, shape, std.mem.sliceAsBytes(v)); + } + + fn bf16Const(o: *Out, name: []const u8, shape: []const usize, value: f32) !void { + const v = try o.a.alloc(u16, count(shape)); + @memset(v, bf16Of(value)); + try o.put(name, .bf16, shape, std.mem.sliceAsBytes(v)); + } + + /// f32 values uniform in [lo, hi). + fn f32s(o: *Out, name: []const u8, shape: []const usize, lo: f32, hi: f32) !void { + const v = try o.a.alloc(f32, count(shape)); + for (v) |*x| x.* = lo + o.rng.float(f32) * (hi - lo); + try o.put(name, .f32, shape, std.mem.sliceAsBytes(v)); + } + + /// A [k]-wide linear's rows scaled about 1/sqrt(k). + fn linear(o: *Out, name: []const u8, n: usize, k: usize) !void { + try o.bf16(name, &.{ n, k }, 1.0 / @sqrt(@as(f32, @floatFromInt(k)))); + } + + /// modelopt NVFP4: random e2m1 codes, e4m3 block scales 0.5 to 1.875, a global scale near 1/(6 sqrt(k)). + fn nvfp4(o: *Out, base: []const u8, n: usize, k: usize) !void { + const codes = try o.a.alloc(u8, n * k / 2); + o.rng.bytes(codes); + try o.put(try cat(o.a, base, ".weight"), .u8, &.{ n, k / 2 }, codes); + const scales = try o.a.alloc(u8, n * k / 16); + for (scales) |*x| x.* = 0x30 + o.rng.uintLessThan(u8, 16); + try o.put(try cat(o.a, base, ".weight_scale"), .f8_e4m3, &.{ n, k / 16 }, scales); + try o.scalar(try cat(o.a, base, ".weight_scale_2"), 1.0 / (6.0 * @sqrt(@as(f32, @floatFromInt(k))))); + try o.scalar(try cat(o.a, base, ".input_scale"), 1.0); + } + + fn scalar(o: *Out, name: []const u8, value: f32) !void { + const v = try o.a.alloc(f32, 1); + v[0] = value; + try o.put(name, .f32, &.{}, std.mem.sliceAsBytes(v)); + } +}; + +fn count(shape: []const usize) usize { + var n: usize = 1; + for (shape) |d| n *= d; + return n; +} + +fn bf16Of(x: f32) u16 { + const b: u32 = @bitCast(x); + return @intCast((b + 0x7FFF + ((b >> 16) & 1)) >> 16); +} + +fn cat(a: std.mem.Allocator, x: []const u8, y: []const u8) ![]u8 { + return std.mem.concat(a, u8, &.{ x, y }); +} + +fn tensor(a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) ![]u8 { + return std.fmt.allocPrint(a, P ++ fmt, args); +} + +fn hc(o: *Out, i: usize) !void { + for ([_][]const u8{ "attn", "ffn" }) |w| { + try o.bf16(try tensor(o.a, "layers.{d}.hc_{s}_fn", .{ i, w }), &.{ 24, 4 * hidden }, 1.0 / 128.0); + try o.bf16(try tensor(o.a, "layers.{d}.hc_{s}_base", .{ i, w }), &.{24}, 0.1); + try o.bf16Const(try tensor(o.a, "layers.{d}.hc_{s}_scale", .{ i, w }), &.{3}, 1.0); + } +} + +fn norms(o: *Out, i: usize) !void { + try o.bf16Const(try tensor(o.a, "layers.{d}.input_layernorm.weight", .{i}), &.{hidden}, 1.0); + try o.bf16Const(try tensor(o.a, "layers.{d}.post_attention_layernorm.weight", .{i}), &.{hidden}, 1.0); +} + +fn kda(o: *Out, i: usize) !void { + const s = try tensor(o.a, "layers.{d}.self_attn.", .{i}); + for ([_][]const u8{ "q", "k", "v" }) |m| { + try o.linear(try std.fmt.allocPrint(o.a, "{s}{s}_proj.weight", .{ s, m }), 8192, hidden); + try o.f32s(try std.fmt.allocPrint(o.a, "{s}{s}_conv1d.weight", .{ s, m }), &.{ 8192, 1, 4 }, -0.5, 0.5); + } + try o.linear(try cat(o.a, s, "f_a_proj.weight"), 128, hidden); + try o.linear(try cat(o.a, s, "g_a_proj.weight"), 128, hidden); + try o.linear(try cat(o.a, s, "b_proj.weight"), 64, hidden); + try o.linear(try cat(o.a, s, "f_b_proj.weight"), 8192, 128); + try o.linear(try cat(o.a, s, "g_b_proj.weight"), 8192, 128); + try o.f32s(try cat(o.a, s, "A_log"), &.{64}, 0.0, 2.5); + try o.f32s(try cat(o.a, s, "dt_bias"), &.{8192}, -0.5, 0.5); + try o.bf16Const(try cat(o.a, s, "o_norm.weight"), &.{128}, 1.0); + try o.linear(try cat(o.a, s, "o_proj.weight"), hidden, 8192); +} + +fn mla(o: *Out, i: usize) !void { + const s = try tensor(o.a, "layers.{d}.self_attn.", .{i}); + try o.linear(try cat(o.a, s, "q_a_proj.weight"), 1536, hidden); + try o.linear(try cat(o.a, s, "kv_a_proj_with_mqa.weight"), 512, hidden); + try o.bf16Const(try cat(o.a, s, "q_a_layernorm.weight"), &.{1536}, 1.0); + try o.bf16Const(try cat(o.a, s, "kv_a_layernorm.weight"), &.{512}, 1.0); + try o.linear(try cat(o.a, s, "q_b_proj.weight"), 16384, 1536); + try o.linear(try cat(o.a, s, "kv_b_proj.weight"), 32768, 512); + try o.linear(try cat(o.a, s, "o_proj.weight"), hidden, 16384); + try o.linear(try cat(o.a, s, "indexer.wk.weight"), 128, hidden); + try o.linear(try cat(o.a, s, "indexer.weights_proj.weight"), 32, hidden); + try o.linear(try cat(o.a, s, "indexer.wq_b.weight"), 4096, 1536); + try o.bf16Const(try cat(o.a, s, "indexer.k_norm.weight"), &.{128}, 1.0); + try o.bf16Const(try cat(o.a, s, "indexer.k_norm.bias"), &.{128}, 0.0); + try o.linear(try cat(o.a, s, "indexer.index_kpool_compress_gate"), 128, hidden); + try o.bf16(try cat(o.a, s, "indexer.index_kpool_compress_ape"), &.{ 4, 128 }, 0.02); +} + +fn dense(o: *Out, i: usize) !void { + for ([_][]const u8{ "gate", "up" }) |m| try o.nvfp4(try tensor(o.a, "layers.{d}.mlp.{s}_proj", .{ i, m }), 12288, hidden); + try o.nvfp4(try tensor(o.a, "layers.{d}.mlp.down_proj", .{i}), hidden, 12288); +} + +/// The router, the shared expert, and the routed ones in NVFP4 (backbone) or bf16 (the MTP layer). +fn moe(o: *Out, i: usize, bf: bool) !void { + try o.linear(try tensor(o.a, "layers.{d}.mlp.gate.weight", .{i}), experts, hidden); + try o.f32s(try tensor(o.a, "layers.{d}.mlp.gate.e_score_correction_bias", .{i}), &.{experts}, 0.0, 0.0); + try o.linear(try tensor(o.a, "layers.{d}.mlp.shared_experts.gate_proj.weight", .{i}), 2048, hidden); + try o.linear(try tensor(o.a, "layers.{d}.mlp.shared_experts.up_proj.weight", .{i}), 2048, hidden); + try o.linear(try tensor(o.a, "layers.{d}.mlp.shared_experts.down_proj.weight", .{i}), hidden, 2048); + for (0..experts) |e| { + const s = try tensor(o.a, "layers.{d}.mlp.experts.{d}.", .{ i, e }); + if (bf) { + try o.linear(try cat(o.a, s, "gate_proj.weight"), 2048, hidden); + try o.linear(try cat(o.a, s, "up_proj.weight"), 2048, hidden); + try o.linear(try cat(o.a, s, "down_proj.weight"), hidden, 2048); + } else { + try o.nvfp4(try cat(o.a, s, "gate_proj"), 2048, hidden); + try o.nvfp4(try cat(o.a, s, "up_proj"), 2048, hidden); + try o.nvfp4(try cat(o.a, s, "down_proj"), hidden, 2048); + } + } +} + +pub fn run(gpu: check.Gpu, out_dir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(53); + var o: Out = .{ .a = a, .rng = prng.random() }; + try o.bf16(P ++ "embed_tokens.weight", &.{ vocab, hidden }, 1.0); + try o.bf16Const(P ++ "norm.weight", &.{hidden}, 1.0); + try o.linear("lm_head.weight", vocab, hidden); + try kda(&o, 0); + try hc(&o, 0); + try norms(&o, 0); + try dense(&o, 0); + try mla(&o, 1); + try hc(&o, 1); + try norms(&o, 1); + try moe(&o, 1, false); + try mla(&o, layers); // the MTP layer: no hyper-connections, bf16 routed experts, the draft head's norms and eh_proj + try norms(&o, layers); + try moe(&o, layers, true); + for ([_][]const u8{ "enorm", "hnorm", "shared_head.norm" }) |m| try o.bf16Const(try tensor(a, "layers.{d}.{s}.weight", .{ layers, m }), &.{hidden}, 1.0); + try o.linear(try tensor(a, "layers.{d}.eh_proj.weight", .{layers}), hidden, 2 * hidden); + + try std.Io.Dir.cwd().createDirPath(gpu.io, out_dir); + const text = try glm.config.textConfig(a, layers, 1, &.{1}, experts, vocab, hidden); + const end = std.mem.lastIndexOfScalar(u8, text, '}') orelse return error.BadConfig; + const cfg = try std.mem.concat(a, u8, &.{ text[0..end], modelopt_block, text[end..] }); // its linears are NVFP4 + _ = try glm.config.parse(a, cfg); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = try std.fs.path.join(a, &.{ out_dir, "config.json" }), .data = cfg }); + try core.shard_edit.create(gpu.gpa, try std.fs.path.join(a, &.{ out_dir, "model.safetensors" }), &.{.{ "format", "pt" }}, o.list.items); + + // three prompts of the decode checks' lengths, ids under the eos at vocab - 2; want: 512 zeros (a length only) + var dc: std.Io.Writer.Allocating = .init(a); + try dc.writer.writeAll("{\"capacity\": 8192, \"kinds\": ["); + for ([_][]const u8{ "count", "prose", "code" }, [_]usize{ 2066, 2108, 2089 }, 0..) |kind, len, k| { + try dc.writer.print("{s}{{\"kind\": \"{s}\", \"prompt\": [", .{ if (k > 0) ", " else "", kind }); + for (0..len) |i| try dc.writer.print("{s}{d}", .{ if (i > 0) "," else "", o.rng.uintLessThan(u32, vocab - 128) }); + try dc.writer.writeAll("], \"want\": ["); + for (0..512) |i| try dc.writer.writeAll(if (i > 0) ",0" else "0"); + try dc.writer.writeAll("]}"); + } + try dc.writer.writeAll("]}\n"); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = try std.fs.path.join(a, &.{ out_dir, "dc.json" }), .data = dc.written() }); + var bytes: usize = 0; + for (o.list.items) |t| bytes += t.bytes.len; + check.pass("glm-synth: {d} tensors, {d} MiB, config {d} layers {d} experts vocabulary {d}, dc.json in {s}", .{ o.list.items.len, bytes >> 20, layers, experts, vocab, out_dir }); +} diff --git a/zig/tests/cuda/libs_tests.zig b/zig/tests/cuda/libs_tests.zig index 67bf07ea59..63d83779fc 100644 --- a/zig/tests/cuda/libs_tests.zig +++ b/zig/tests/cuda/libs_tests.zig @@ -67,6 +67,72 @@ pub fn cublaslt(gpu: Gpu) !void { check.pass("cuBLASLt {d}: bf16 x bf16 -> fp32 ({d}x{d}x{d}), reruns bit-equal, max abs error {e} vs fp64, workspace {d} B", .{ lt.api.cublasLtGetVersion(), m, n, k, worst, linear.workspace_need }); } +/// `nccl-pair RANK WORLD ID`: one rank of a multi-host communicator (id as hex in file ID): send/recv, sums, gathers. +pub fn ncclPair(gpu: Gpu, rank: c_int, world: c_int, id_path: []const u8) !void { + const d = gpu.d; + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + var id: cuda.nccl.UniqueId = undefined; + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + const hex = std.fmt.bytesToHex(id.internal, .lower); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = id_path, .data = &hex }); + } else { + var hex: [256]u8 = undefined; + const got = try std.Io.Dir.cwd().readFile(gpu.io, id_path, &hex); + try expect(got.len == 256, "unique id file {s}: {d} hex digits", .{ id_path, got.len }); + _ = try std.fmt.hexToBytes(&id.internal, &hex); + } + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, world, id, rank), "ncclCommInitRank"); + defer _ = lib.api.ncclCommDestroy(comm); + const count: usize = 1 << 20; + const W: usize = @intCast(world); + const R: usize = @intCast(rank); + const mine = try gpu.gpa.alloc(f32, count); + defer gpu.gpa.free(mine); + for (mine, 0..) |*v, i| v.* = @floatFromInt(R * 1000 + i % 977); + var send = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(mine)); + defer send.free(); + var recv = try cuda.DeviceBuffer.alloc(d, count * 4 * W); + defer recv.free(); + var stream = try cuda.Stream.init(d, true); + defer stream.deinit(); + try lib.check(lib.api.ncclGroupStart(), "ncclGroupStart"); + for (0..W) |p| if (p != R) { + try lib.check(lib.api.ncclSend(send.ptr, count, .f32, @intCast(p), comm, stream.handle), "ncclSend"); + try lib.check(lib.api.ncclRecv(recv.ptr + p * count * 4, count, .f32, @intCast(p), comm, stream.handle), "ncclRecv"); + }; + try lib.check(lib.api.ncclGroupEnd(), "ncclGroupEnd"); + try stream.synchronize(); + const got = try check.download(gpu, recv); + defer gpu.gpa.free(got); + const want = try gpu.gpa.alloc(f32, count); + defer gpu.gpa.free(want); + for (0..W) |p| if (p != R) { + for (want, 0..) |*v, i| v.* = @floatFromInt(p * 1000 + i % 977); + try check.sameBytes("send/recv from a peer", got[p * count * 4 ..][0 .. count * 4], std.mem.sliceAsBytes(want)); + }; + try lib.check(lib.api.ncclAllReduce(send.ptr, recv.ptr, count, .f32, .sum, comm, stream.handle), "ncclAllReduce"); + try stream.synchronize(); + try recv.download(0, got[0 .. count * 4]); + for (want, 0..) |*v, i| { + var s: f32 = 0; + for (0..W) |p| s += @floatFromInt(p * 1000 + i % 977); + v.* = s; + } + try check.sameBytes("all-reduce", got[0 .. count * 4], std.mem.sliceAsBytes(want)); + // all-gather: every rank's pattern in rank order + try lib.check(lib.api.ncclAllGather(send.ptr, recv.ptr, count, .f32, comm, stream.handle), "ncclAllGather"); + try stream.synchronize(); + try recv.download(0, got); + for (0..W) |p| { + for (want, 0..) |*v, i| v.* = @floatFromInt(p * 1000 + i % 977); + try check.sameBytes("all-gather", got[p * count * 4 ..][0 .. count * 4], std.mem.sliceAsBytes(want)); + } + check.pass("NCCL {d}: rank {d} of {d}, send/recv with every peer ({d} MiB each), all-reduce and all-gather exact", .{ try lib.version(), rank, world, count * 4 >> 20 }); +} + pub fn nccl(gpu: Gpu) !void { const d = gpu.d; var lib = try cuda.nccl.Library.open(); @@ -96,5 +162,14 @@ pub fn nccl(gpu: Gpu) !void { try stream.synchronize(); try recv.download(0, got); try check.sameBytes("one-rank all-gather", got, std.mem.sliceAsBytes(src)); - check.pass("NCCL {d}: one-rank communicator, all-reduce and all-gather exact", .{version}); + // a send to this rank and its receive, in one group (NCCL pairs them) + try recv.fill8(0, null); + try lib.check(lib.api.ncclGroupStart(), "ncclGroupStart"); + try lib.check(lib.api.ncclSend(send.ptr, count, .f32, 0, comms[0], stream.handle), "ncclSend"); + try lib.check(lib.api.ncclRecv(recv.ptr, count, .f32, 0, comms[0], stream.handle), "ncclRecv"); + try lib.check(lib.api.ncclGroupEnd(), "ncclGroupEnd"); + try stream.synchronize(); + try recv.download(0, got); + try check.sameBytes("one-rank send/recv", got, std.mem.sliceAsBytes(src)); + check.pass("NCCL {d}: one-rank communicator, all-reduce, all-gather and send/recv to itself exact", .{version}); } diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aec6b9ec1b..36382fde57 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -19,6 +19,7 @@ const fp8_tests = @import("fp8_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); +const glm_ref_tests = @import("glm_ref_tests.zig"); const carveout_tests = @import("carveout_tests.zig"); const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); @@ -29,6 +30,8 @@ const norms_tests = @import("norms_tests.zig"); const norm_rows_tests = @import("norm_rows_tests.zig"); const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); +const glm_lanes_tests = @import("glm_lanes_tests.zig"); +const glm_synth = @import("glm_synth.zig"); const usage = \\usage: tf-cuda-test @@ -42,6 +45,7 @@ const usage = \\ symbols every gdn fatbin instantiation resolves by its listed symbol \\ cublaslt bf16 GEMM through cuBLASLt against an fp64 reference \\ nccl one-rank NCCL all-reduce and all-gather + \\ nccl-pair one rank of a multi-host communicator: send/recv with every peer, all-reduce, all-gather \\ gdn-replay replay_kernel bits against the Python oracle's fixture \\ gdn-tree tree_kernel bits against the Python oracle's fixture \\ triton a Triton cubin's bits against the Python oracle's fixture @@ -69,6 +73,11 @@ const usage = \\ exlr8-bench DIR [MANIFEST] [LAYER] [SETS] a real EXLR8 layer: FP64, row windows, then GB/s by stage \\ exlr8-dense [DIR] EXLR8 dense formats: decode bit-equal, FP64, row windows; DIR adds the real matrices \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time + \\ glm-ref GLM projections (matmuls, router, attention, indexer, hc) against float64 (synthetic) + \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests + \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run + \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte + \\ glm-synth a synthetic GLM-5.3-Flash checkpoint (2 layers, 16 experts) and its dc.json for the lane checks \\ ; @@ -115,6 +124,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "symbols")) return runtime_tests.symbols(gpu); if (std.mem.eql(u8, cmd, "cublaslt")) return libs_tests.cublaslt(gpu); if (std.mem.eql(u8, cmd, "nccl")) return libs_tests.nccl(gpu); + if (std.mem.eql(u8, cmd, "nccl-pair")) return libs_tests.ncclPair(gpu, try std.fmt.parseInt(c_int, try arg(rest, 0), 10), try std.fmt.parseInt(c_int, try arg(rest, 1), 10), try arg(rest, 2)); if (std.mem.eql(u8, cmd, "gdn-replay")) return oracle_tests.gdnReplay(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "gdn-tree")) return oracle_tests.gdnTree(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "triton")) return oracle_tests.tritonKernel(gpu, try arg(rest, 0)); @@ -123,6 +133,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); + if (std.mem.eql(u8, cmd, "glm-ref")) return glm_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); @@ -146,6 +157,10 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { return carveout_tests.run(gpu, if (rest.len > 1) rest[1] else cuda.carveout.default_card, mib); } if (std.mem.eql(u8, cmd, "train-pack")) return train_pack.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 24); + if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); + if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); + if (std.mem.eql(u8, cmd, "glm-synth")) return glm_synth.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); return error.UnknownCommand; } @@ -161,3 +176,8 @@ fn info(gpu: check.Gpu) !void { cuda.kernels.available, }); } + +fn hasFlag(flags: []const [:0]const u8, name: []const u8) bool { + for (flags) |f| if (std.mem.eql(u8, f, name)) return true; + return false; +} diff --git a/zig/tests/cuda/nvfp4_experts_tests.zig b/zig/tests/cuda/nvfp4_experts_tests.zig index 7c8b7bdb2e..2031017836 100644 --- a/zig/tests/cuda/nvfp4_experts_tests.zig +++ b/zig/tests/cuda/nvfp4_experts_tests.zig @@ -39,7 +39,7 @@ pub fn experts(gpu: Gpu, dir: []const u8) !void { for (0..e) |x| { const gate = [2][]const u8{ raw[0][0][x * wb ..][0..wb], raw[0][1][x * sbytes ..][0..sbytes] }; const up = [2][]const u8{ raw[1][0][x * wb ..][0..wb], raw[1][1][x * sbytes ..][0..sbytes] }; - try fx4.packGateUp(a, pu[x * 2 * blk ..][0 .. 2 * blk], gate, up, ni, d); + fx4.packGateUp(pu[x * 2 * blk ..][0 .. 2 * blk], gate, up, ni, d); fx4.packOne(pd[x * blk ..][0..blk], raw[2][0][x * wb ..][0..wb], raw[2][1][x * sbytes ..][0..sbytes], d, ni); } inline for (.{ .{ "packed_up", pu }, .{ "packed_down", pd } }) |c| { diff --git a/zig/tests/cuda/quant_ref_experts.zig b/zig/tests/cuda/quant_ref_experts.zig index 4b9dad472d..891785aa54 100644 --- a/zig/tests/cuda/quant_ref_experts.zig +++ b/zig/tests/cuda/quant_ref_experts.zig @@ -120,6 +120,41 @@ fn batch(r: *Rig, comptime label: []const u8, g: experts.Grouped, router: groupe return .{ .act = act, .y = y }; } +/// Prompt items of expert `skip` (the shared one, as a family that computes it itself passes) are left alone. +fn skipped(r: *Rig, comptime label: []const u8, g: experts.Grouped, router: grouped.Router, l: experts.Layer, x: []const u16, rows: usize) !void { + const pairs = rows * slots; + const pk = try picks(r, rows); + const pl = try plan(r, router, try r.dev(std.mem.sliceAsBytes(pk)), pairs, experts.promptTile(l.format)); + const fill = try r.a.alloc(u8, pairs * @max(NI, D) * 2); + @memset(fill, 0xA5); + const skip: i32 = E - 1; + const up_in: experts.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(x[0 .. rows * D])), .stride = D, .slots = slots }; + const a0 = try r.zeros(pairs * NI * 2); + const a1 = try r.dev(fill[0 .. pairs * NI * 2]); + try g.prompt(r.s, .up, up_in, l, pl.p, a0, pl.items, -1); + try g.prompt(r.s, .up, up_in, l, pl.p, a1, pl.items, skip); + try leftAlone(label ++ " prompt gate-up", pk, skip, try r.back(a0, pairs * NI * 2), try r.back(a1, pairs * NI * 2), NI * 2, fill); + const down_in: experts.Rows = .{ .x = a0, .stride = NI }; + const y0 = try r.zeros(pairs * D * 2); + const y1 = try r.dev(fill[0 .. pairs * D * 2]); + try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y0, pl.items, -1); + try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y1, pl.items, skip); + try leftAlone(label ++ " prompt down", pk, skip, try r.back(y0, pairs * D * 2), try r.back(y1, pairs * D * 2), D * 2, fill); +} + +/// Each pair's row: the skipped expert's still the fill, every other one equal to the run without skip. +fn leftAlone(name: []const u8, pk: []const i32, skip: i32, plain: []const u8, got: []const u8, row: usize, fill: []const u8) !void { + var left: usize = 0; + var differ: usize = 0; + for (pk, 0..) |e, p| { + const b = got[p * row ..][0..row]; + if (e == skip) { + if (std.mem.eql(u8, b, fill[0..row])) left += 1 else differ += 1; + } else if (!std.mem.eql(u8, plain[p * row ..][0..row], b)) differ += 1; + } + try check.expect(differ == 0 and left > 0, "{s}: expert {d}'s {d} pairs left alone, every other pair as without skip ({d} differ)", .{ name, skip, left, differ }); +} + const Fp8M = struct { codes: []const u8, inv: []const f32, n: usize, k: usize }; fn fp8Val(c: Fp8M, e: usize, row: usize, col: usize) f64 { const nb = c.n / 128; @@ -181,6 +216,7 @@ pub fn run(r: *Rig) !void { const M = Matrix(Fp8M); for ([_]usize{ 1, 3 }) |rows| _ = try batch(r, "fp8", g, router, l, x, rows, false, M{ .ctx = gm, .value = fp8Val }, M{ .ctx = um, .value = fp8Val }, M{ .ctx = dm, .value = fp8Val }); _ = try batch(r, "fp8", g, router, l, x, 300, true, M{ .ctx = gm, .value = fp8Val }, M{ .ctx = um, .value = fp8Val }, M{ .ctx = dm, .value = fp8Val }); + try skipped(r, "fp8", g, router, l, x, 300); } // NVFP4 { @@ -195,7 +231,7 @@ pub fn run(r: *Rig) !void { const sb = NI * D / 16; const us = try r.a.alloc(f32, 2 * E); for (0..E) |e| { - try W.packGateUp(r.a, pu[e * 2 * blk ..][0 .. 2 * blk], .{ gm.codes[e * wb ..][0..wb], gm.scales[e * sb ..][0..sb] }, .{ um.codes[e * wb ..][0..wb], um.scales[e * sb ..][0..sb] }, NI, D); + W.packGateUp(pu[e * 2 * blk ..][0 .. 2 * blk], .{ gm.codes[e * wb ..][0..wb], gm.scales[e * sb ..][0..sb] }, .{ um.codes[e * wb ..][0..wb], um.scales[e * sb ..][0..sb] }, NI, D); W.packOne(pd[e * blk ..][0..blk], dm.codes[e * wb ..][0..wb], dm.scales[e * sb ..][0..sb], D, NI); us[2 * e] = gm.global[e]; us[2 * e + 1] = um.global[e]; @@ -207,6 +243,7 @@ pub fn run(r: *Rig) !void { const M = Matrix(Fp4M); for ([_]usize{ 1, 3, 40 }) |rows| _ = try batch(r, "nvfp4", g, router, l, x, rows, false, M{ .ctx = gm, .value = fp4Val }, M{ .ctx = um, .value = fp4Val }, M{ .ctx = dm, .value = fp4Val }); _ = try batch(r, "nvfp4", g, router, l, x, 300, true, M{ .ctx = gm, .value = fp4Val }, M{ .ctx = um, .value = fp4Val }, M{ .ctx = dm, .value = fp4Val }); + try skipped(r, "nvfp4", g, router, l, x, 300); } - check.pass("experts fp8g and nvfp4 (decode, prompt): gate-up and down within float64 tolerance", .{}); + check.pass("experts fp8g and nvfp4 (decode, prompt): gate-up and down within float64 tolerance; a skipped expert's prompt items left alone", .{}); } diff --git a/zig/tests/cuda/quant_ref_tests.zig b/zig/tests/cuda/quant_ref_tests.zig index a4552e898b..f33487c538 100644 --- a/zig/tests/cuda/quant_ref_tests.zig +++ b/zig/tests/cuda/quant_ref_tests.zig @@ -168,11 +168,32 @@ fn lanes(r: *Rig, lin: qlinear.Linear, n: usize, k: usize) !void { const path = if (prompt) "prompt" else "decode"; try r.close(try std.fmt.bufPrint(&name, "{s} {s} {d} rows x [{d}, {d}]", .{ f[0], path, wide, n, k }), outs[rows.len - 1], false, ref.want, ref.tol); for (rows[0 .. rows.len - 1], outs[0 .. rows.len - 1]) |m, o| try sameRows(r, try std.fmt.bufPrint(&name, "{s} {s}: {d} rows equal the {d}-row call's", .{ f[0], path, m, wide }), outs[rows.len - 1], o, m * n * 2); + if (comptime std.mem.eql(u8, f[0], "nvfp4")) try fp32Out(r, lin, view, dx, part, n, k, prompt, ref.want, ref.tol, outs[rows.len - 1]); } } check.pass("qlinear fp8g and nvfp4, [{d}, {d}]: decode and prompt within float64 tolerance, every row count's bits equal", .{ n, k }); } +/// nvfp4's decodeF32 and promptF32: in float64 tolerance, the same bits at any row count, the bf16 out its rounding. +fn fp32Out(r: *Rig, lin: qlinear.Linear, view: qlinear.Weight, dx: u64, part: u64, n: usize, k: usize, prompt: bool, want: []const f64, tol: []const f64, bf16_out: u64) !void { + const wide = rows[rows.len - 1]; + const path = if (prompt) "prompt" else "decode"; + var name: [96]u8 = undefined; + const outs = try r.a.alloc(u64, rows.len); + for (rows, outs) |m, *o| { + o.* = try r.zeros(m * n * 4); + const in: qlinear.Rows = .{ .x = dx, .ldx = k, .part = part }; + if (prompt) try lin.promptF32(r.s, view, in, o.*, m) else try lin.decodeF32(r.s, view, in, o.*, m); + } + try r.close(try std.fmt.bufPrint(&name, "nvfp4 {s} fp32 {d} rows x [{d}, {d}]", .{ path, wide, n, k }), outs[rows.len - 1], true, want, tol); + for (rows[0 .. rows.len - 1], outs[0 .. rows.len - 1]) |m, o| try sameRows(r, try std.fmt.bufPrint(&name, "nvfp4 {s} fp32: {d} rows equal the {d}-row call's", .{ path, m, wide }), outs[rows.len - 1], o, m * n * 4); + const f = std.mem.bytesAsSlice(f32, try r.back(outs[rows.len - 1], wide * n * 4)); + const b = std.mem.bytesAsSlice(u16, try r.back(bf16_out, wide * n * 2)); + var differ: usize = 0; + for (f, b) |x, y| differ += @intFromBool(bf16(x) != y); + try check.expect(differ == 0, "nvfp4 {s}: the bf16 out is the fp32 out rounded ({d} of {d} differ)", .{ path, differ, f.len }); +} + pub fn run(gpu: Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit();