diff --git a/docs/recipes/exl3.md b/docs/recipes/exl3.md index 2b36354f58..fe7441334d 100644 --- a/docs/recipes/exl3.md +++ b/docs/recipes/exl3.md @@ -172,4 +172,32 @@ reproduces the numpy oracle on every codebook and width and on real checkpoint b (the oracle vectors ship beside the test). Its real checkpoint bytes are one 128x128 block from `UnstableLlama/Qwen3.5-2B-exl3-8.00bpw` (mcg codebook, 8.00 bpw), in `zig/tests/fixtures/exl3_oracle.bin`. -The scalar decoder and its oracle tests do not register an EXL3 model or provide a native GPU projection. +The scalar decoder and its oracle tests do not register an EXL3 model on their own; the native CUDA kernels below +are checked against them. + +## Native CUDA (Zig) + +The Zig engine carries the same arithmetic as hand-written CUDA, ported from the Python engine's kernels so a +row's bits match it: + +- `zig/kernels/cuda/exl3_decode.cuh`, `exl3_linear.cu`, `zig/src/cuda/exl3.zig`: the input rotation and the linear + for 1 to 128 rows, every codebook and width (the `mul1` half steps included), with the K split planned from the + layer's shape alone. `exl3.Dense` runs prompt rows in 128-row blocks. A block of more than 16 rows on a `mul1` + weight takes the prompt kernel: 64-row passes that decode each weight tile once, keeping the decode-row kernel's + sums in its order. A prompt row's bits therefore equal its decode-row bits wherever its chunk starts. + `qlinear.Weight` has an `.exl3` arm for both. +- `zig/kernels/cuda/exl3_experts.cu`, `zig/src/cuda/exl3_experts.zig`: grouped routed experts (member lists, + each expert's rotation, gate and up, GLM's SwiGLU roundings, down and the weighted combine), widths mixed within + a layer; an expert's member count never changes a row's bits. +- `zig/src/core/exl3_load.zig`: every trellis group of a checkpoint from its safetensors headers; + `zig/src/cuda/exl3_upload.zig` puts a group, or a layer's experts, on the device. + +`tf-cuda-test exl3 DIR` checks them: every width and codebook against float64 references built on +`exl3_format.zig`, row and chunk invariance, and, with `DIR` from +`zig/tests/cuda/oracle/exl3_oracle.py` (run with the Python engine on a CUDA box), real checkpoint layers bit for +bit against the Python engine, two whole layers of RED-SNOW-5.3-Flash (a KDA and a DSA layer, all 288 experts under +real routing) included. `tf-cuda-test exl3-tp2 RANK HOST:PORT DIR` runs the GLM split rules on two ranks over NCCL +(output-column cuts, input-row cuts with fp32 partials summed rank 0 first, and the experts' intermediate cut) and +checks every rank's bytes and the sums the same way. + +No family serves EXL3 natively yet: a family's engine takes these pieces through `qlinear` and `exl3_experts`. diff --git a/docs/recipes/glm-5.3-flash.md b/docs/recipes/glm-5.3-flash.md index 4c182e2fa6..b750f487b9 100644 --- a/docs/recipes/glm-5.3-flash.md +++ b/docs/recipes/glm-5.3-flash.md @@ -3,7 +3,8 @@ This recipe describes the Python 0.6.6 engine, maintained on [`python-0.6`](https://github.com/ashhart/TensorFold/tree/python-0.6). Its CUDA EXL3 reader, two-rank flags and DFlash2 policies do not apply to the native Zig binary. -The native binary serves GLM on two M5 Ultra Macs with MLX affine 4-bit/group-64 weights; it has no GLM CUDA backend yet. +The native binary serves GLM on two M5 Ultra Macs with MLX affine 4-bit/group-64 weights; it does not serve GLM on +CUDA yet ([its CUDA engine](#native-cuda-engine-zig) runs under `tf-cuda-test` only). That is not an engine-only replacement for a CUDA EXL3 deployment. For native serving, use the [qualified model table](../../README.md#qualified-models) and [paired Metal runbook](../../RUNBOOK.md#flash-next-and-paired-metal-serving). @@ -61,6 +62,38 @@ The expert decoder and BF16 target matmul keep row arithmetic fixed. A quantized propose drafts, but target verification retains the BF16 head. EXL3 speed, capacity and long-context qualification are TBD [release-0.3.5]. +### Native CUDA engine (Zig) + +The Zig side carries a GLM-5.3-Flash engine for two CUDA ranks that reads +`vcruz305/RED-SNOW-5.3-FLASH-EXL3-SAGE-4.91bpw` (EXL3 `mul1` at mixed widths, unsplit). Each rank takes its +shard at load. Its kernels are hand-written CUDA ported from the Python engine on this repository's `redsnow-v1` +branch. They cover the bf16 matmul and glue, KDA's decode chain and chunked prompt kernels, the DSA indexer and +top-pool selection, latent attention, hyper-connection mixing, the MoE router and EXL3 experts, the dense MLP of +layers 0-2, and the head with its greedy pick. Each one keeps the Python engine's order of operations, so a +greedy reply is token for token the Python engine's `"draft": false` reply. Prompts run in 2,048-row chunks. +Decoding is serial (no drafts): each step replays a captured CUDA graph, one per KDA parity in the dense window and +one per pool bucket past it. + +On two GB10s, over the eight reference prompts (221 to 16,416 prompt tokens; 24,099 in all), it read prompts at +435.0 tok/s and decoded at 24.7 tok/s. The Python engine decodes serially at 24.6 tok/s. + +Two commands check it, both built by `zig build -Dnvcc=/usr/local/cuda/bin/nvcc`: + +```bash +# fixtures from the Python engine (run on a CUDA box with redsnow-v1 installed), then the kernels and layers +python zig/tests/cuda/oracle/glm_oracle.py --model MODEL_DIR --out FIXTURES +tf-cuda-test glm FIXTURES +# the whole model at TP=2, a process a rank (rank 1 first), against recorded replies +tf-cuda-test glm-model 1 HOST:PORT MODEL_DIR reference.json +tf-cuda-test glm-model 0 HOST:PORT MODEL_DIR reference.json +``` + +`tf-cuda-test glm` holds every kernel against float64 references and the Python engine's bytes, then whole KDA +and DSA layers of the real checkpoint at TP=2. `reference.json` is +`{"cases": [{"name", "prompt_ids", "ids"}], "eos_ids"}`: the Python engine's greedy replies with `"draft": false`. +A case passes when the Zig reply's ids equal `ids`. Rank 0 prints a PASS or FAIL line for each case and +`RESULT glm.model.prefill_tps` / `decode_tps`. + ### Draft policies For the affine checkpoint with its MTP head loaded, the default `auto` policy uses MTP for sampled requests. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index b70dfe431f..ab205be9ad 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -41,6 +41,16 @@ 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 = "exl3_linear", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/linear.cu's device code, tensorfold_exl3_linear_v5 + .{ .name = "exl3_experts", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/experts.cu's device code, tensorfold_exl3_experts + .{ .name = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // glm5_next/cuda/kda.cu's device code (chain, replay, wide trio, conv shift) + .{ .name = "glm_kda_chunk", .flags = &.{ "-O3", "--fmad=false" } }, // glm5_next/cuda/kda_chunk.cu's device code (32-row WY sub-chunks) + .{ .name = "glm_bmm", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // glm5_next/cuda's bf16 matmul and router dots (qmm._bmm, glue._router) + .{ .name = "glm_glue", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // glm5_next/cuda/glue.py's Triton kernels as device code + .{ .name = "glm_glue2", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // glue.py's 1536/512 rmsnorms and the fp32-y combine (glm_glue.cu is at the file cap) + .{ .name = "glm_sparse", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // glm5_next/cuda/sparse.py's Triton kernels as device code + .{ .name = "glm_latent", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // latent.py's cache write and kv_b projections + .{ .name = "glm_latent_att", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // latent.py's chunk attention and merge (glm_latent.cu is at the file cap) .{ .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 }, @@ -77,7 +87,7 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return cuda; } -fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { +fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, glm: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); core.addImport("tokenizer", tokenizer); @@ -89,7 +99,10 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin nemotron.addImport("lanes", lanes); nemotron.addImport("heat", heat); nemotron.addImport("nemotron_draft_ids", draft_ids); - return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat }; + const glm = b.createModule(.{ .root_source_file = b.path("zig/src/families/glm/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + glm.addImport("cuda", cuda); + glm.addImport("core", core); + return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .glm = glm, .tokenizer = tokenizer, .heat = heat }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -137,6 +150,7 @@ 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("glm", mods.glm); 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; } @@ -204,7 +218,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&server_test.step); b.step("test-server-cpu", "The HTTP server's unit tests (routes, templates, tool parsing) without a GPU").dependOn(&server_test.step); // 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); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.glm, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); diff --git a/zig/kernels/cuda/exl3_decode.cuh b/zig/kernels/cuda/exl3_decode.cuh new file mode 100644 index 0000000000..8b84470827 --- /dev/null +++ b/zig/kernels/cuda/exl3_decode.cuh @@ -0,0 +1,141 @@ +// EXL3 tiles: lane L decodes values 8L..8L+7 straight into its two mma B fragments; port of tensorfold's decode.cuh. +#pragma once + +#include +#include + +namespace tf_exl3 { + +enum Codebook : int { CB_3INST = 0, CB_MCG = 1, CB_MUL1 = 2 }; + +// 32-bit words of a tile: 256 values of K2 / 2 bits. +template +__host__ __device__ constexpr int tile_words() { + return 4 * K2; +} + +// E(p): one past the last stream bit of value p's 16-bit window (exl3_format.zig's streamEnds). +template +__host__ __device__ constexpr int stream_end(int p) { + return (K2 & 1) ? ((p + 1) * K2 - ((p + 1) & 1)) / 2 : (p + 1) * (K2 / 2); +} + +// Words a lane reads for its 8 values: the most any lane's windows span (2, or 3 for 3.5 bits and 5 to 8 bits). +template +__host__ __device__ constexpr int lane_words() { + int most = 0; + for (int l = 0; l < 32; ++l) { + const int first = (stream_end(8 * l) - 16 + 128 * K2) % 32; + const int need = first + stream_end(8 * l + 7) - stream_end(8 * l) + 16; + most = need > most ? need : most; + } + return (most + 31) / 32; +} + +// The index of the first word lane `lane` reads, and the bit offset of its first window in that word. +template +__device__ __forceinline__ void lane_start(int lane, int& word, int& offset) { + const int first = 4 * lane * K2 + K2 / 2 - 16 + 128 * K2; // E(8 lane) - 16, made non-negative + word = (first >> 5) % tile_words(); + offset = first & 31; +} + +// The same from global memory through the read-only cache. +template +__device__ __forceinline__ void ldg_lane_words(const uint32_t* tile, int lane, uint32_t (&w)[lane_words()]) { + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int i = 0; i < lane_words(); ++i) w[i] = __ldg(tile + (word + i) % tile_words()); +} + +// Bits [d, d + 16) of the 96-bit big-endian stream hi:mid:lo (d a compile-time constant after unrolling). +__device__ __forceinline__ uint32_t window16(uint32_t hi, uint32_t mid, uint32_t lo, int d) { + if (d <= 16) return (hi >> (16 - d)) & 0xffffu; + if (d < 32) return __funnelshift_l(mid, hi, d) >> 16; + if (d <= 48) return (mid >> (48 - d)) & 0xffffu; + return __funnelshift_l(lo, mid, d - 32) >> 16; +} + +// The 16-bit states of this lane's 8 values (8 lane + j, j = 0..7) from its words. +template +__device__ __forceinline__ void lane_states(const uint32_t (&w)[lane_words()], int lane, uint32_t (&s)[8]) { + int word, offset; + lane_start(lane, word, offset); + const uint32_t c = lane_words() > 2 ? w[lane_words() > 2 ? 2 : 0] : 0u; + const uint32_t hi = __funnelshift_l(w[1], w[0], offset); // the stream from the lane's first window on + const uint32_t mid = __funnelshift_l(c, w[1], offset); + const uint32_t lo = c << offset; +#pragma unroll + for (int j = 0; j < 8; ++j) s[j] = window16(hi, mid, lo, stream_end(j) - stream_end(0)); +} + +// Two codebook values from two 16-bit states, as a half2 in a uint32 (the first state in the low half). +template +__device__ __forceinline__ uint32_t decode2(uint32_t s0, uint32_t s1) { + if constexpr (CB == CB_MUL1) { + const uint32_t x0 = s0 * 0x83DCD12Du, x1 = s1 * 0x83DCD12Du; + const uint32_t h0 = __dp4a(x0, 0x01010101u, 0x6400u); // fp16 bits of 1024 + the byte sum + const uint32_t h1 = __dp4a(x1, 0x01010101u, 0x6400u); + const uint32_t hh = (h0 & 0xffffu) | (h1 << 16); + const half2 r = __hfma2(*reinterpret_cast(&hh), __half2half2(__ushort_as_half(0x1eee)), + __half2half2(__ushort_as_half(0xc931))); + return *reinterpret_cast(&r); + } else { + uint32_t x0, x1; + if constexpr (CB == CB_MCG) { + x0 = s0 * 0xCBAC1FEDu; + x1 = s1 * 0xCBAC1FEDu; + } else { + x0 = s0 * 89226354u + 64248484u; + x1 = s1 * 89226354u + 64248484u; + } + x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u; + x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u; + const uint32_t lo = __byte_perm(x0, x1, 0x5410); + const uint32_t hi = __byte_perm(x0, x1, 0x7632); + const half2 r = __hadd2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + return *reinterpret_cast(&r); + } +} + +// This lane's 8 values of a tile as the B fragments of the tile's two n8 halves (columns 0-7 in b0, 8-15 in b1). +template +__device__ __forceinline__ void decode_lane(const uint32_t (&w)[lane_words()], int lane, uint32_t (&b0)[2], + uint32_t (&b1)[2]) { + uint32_t s[8]; + lane_states(w, lane, s); + b0[0] = decode2(s[0], s[1]); + b0[1] = decode2(s[2], s[3]); + b1[0] = decode2(s[4], s[5]); + b1[1] = decode2(s[6], s[7]); +} + +// mma.m16n8k16, fp16 inputs, fp32 accumulators: d += a @ b. +__device__ __forceinline__ void mma16816(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Unscaled Walsh-Hadamard transform of 128 fp32 values, 4 a lane, in a fixed butterfly order (multiply by 1/sqrt(128)). +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + const float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3]; + v[0] = a + c; + v[1] = b + d; + v[2] = a - c; + v[3] = b - d; +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +constexpr float HAD_SCALE = 0.08838834764831845f; // 1 / sqrt(128) + +} // namespace tf_exl3 diff --git a/zig/kernels/cuda/exl3_experts.cu b/zig/kernels/cuda/exl3_experts.cu new file mode 100644 index 0000000000..c2ac128e02 --- /dev/null +++ b/zig/kernels/cuda/exl3_experts.cu @@ -0,0 +1,477 @@ +// EXL3 routed experts: grouping, grouped GEMMs, epilogues; port of tensorfold's experts.cu (ExLlamaV3, Turboderp). + +#include +#include +#include + +#include "exl3_decode.cuh" + +namespace tf_exl3 { + +// One projection's grouped arguments; the fields must stay in launch.zig add order. +struct GroupedArgs { + const half* x0; // fp16 [P, K], the rotated gate (or down) inputs + const half* x1; // fp16 [P, K], the rotated up inputs + const long long* tp0; // int64 [E], one trellis pointer per expert + const long long* tp1; + const int* k2_0; // int32 [E], half-bits a value + const int* k2_1; + const int* uids; // int32 [nexp_max], distinct routed experts in id order + const int* ucount; // int32 [1] + const int* members; // int32 [nexp_max, maxm], row * 32 + slot, -1 after the last member + float* z; // fp32 partials, [(mat * sk + split) * P + p] * N + col + int k, n, p, sk, maxm, slots, nexp_max, mats, ns; +}; + +// 16 bytes of trellis words, read-only and not kept in L1 (experts.py TF_GLM_EXL3_LOADS's ld.global.nc). +__device__ __forceinline__ uint4 ldg_nc_v4(const uint32_t* p) { + uint4 v; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p)); + return v; +} + +// A warp's k steps of NT tiles into acc; even widths: 16-byte loads (PD in flight) staged, each lane's words read back. +template +__device__ __forceinline__ void warp_tiles(const uint32_t* T, int ntiles, int kt0, int nkt, int nt0, const half* x0, + const half* x1, bool ok0, bool ok1, int lane, float (&acc)[NT][2][4], + uint32_t* stage) { + constexpr int TW = tile_words(), LW = lane_words(); + const size_t kstride = (size_t)ntiles * TW; + auto step = [&](int i, const uint32_t (&w)[NT][LW]) { + const int k = (kt0 + i) * 16; + const uint32_t a[4] = {ok0 ? *reinterpret_cast(x0 + k) : 0u, + ok1 ? *reinterpret_cast(x1 + k) : 0u, + ok0 ? *reinterpret_cast(x0 + k + 8) : 0u, + ok1 ? *reinterpret_cast(x1 + k + 8) : 0u}; +#pragma unroll + for (int j = 0; j < NT; ++j) { + uint32_t b0[2], b1[2]; + decode_lane(w[j], lane, b0, b1); + mma16816(acc[j][0], a, b0); + mma16816(acc[j][1], a, b1); + } + }; + if constexpr (PD > 0 && TW % 8 == 0) { + constexpr int NV = (NT * TW + 127) / 128; // 16-byte loads a lane a step + int word, offset; + lane_start(lane, word, offset); + const uint32_t* tb = T + (size_t)kt0 * kstride + (size_t)nt0 * TW + 4 * lane; + uint4 rv[PD][NV]; +#pragma unroll + for (int d = 0; d < PD; ++d) + if (d < nkt) +#pragma unroll + for (int v = 0; v < NV; ++v) + if (128 * v + 4 * lane < NT * TW) rv[d][v] = ldg_nc_v4(tb + d * kstride + 128 * v); + for (int ib = 0; ib < nkt; ib += PD) { +#pragma unroll + for (int d = 0; d < PD; ++d) { + const int it = ib + d; + if (it >= nkt) break; + __syncwarp(); // every lane has read the previous step back +#pragma unroll + for (int v = 0; v < NV; ++v) + if (128 * v + 4 * lane < NT * TW) *reinterpret_cast(stage + 128 * v + 4 * lane) = rv[d][v]; + if (it + PD < nkt) +#pragma unroll + for (int v = 0; v < NV; ++v) + if (128 * v + 4 * lane < NT * TW) + rv[d][v] = ldg_nc_v4(tb + (size_t)(it + PD) * kstride + 128 * v); + __syncwarp(); // the step is staged before any lane reads it + uint32_t w[NT][LW]; +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int i = 0; i < LW; ++i) w[j][i] = stage[j * TW + (word + i) % TW]; + step(it, w); + } + } + __syncwarp(); // `stage` is the caller's red[warp]: every read done before it is written + } else { + for (int i = 0; i < nkt; ++i) { + const uint32_t* t0 = T + (size_t)(kt0 + i) * kstride + (size_t)nt0 * TW; + uint32_t w[NT][LW]; +#pragma unroll + for (int j = 0; j < NT; ++j) ldg_lane_words(t0 + j * TW, lane, w[j]); + step(i, w); + } + } +} + +// Program (expert u, n block, split and matrix): its member tiles of 16 in order (each as alone), to the first empty. +template +__device__ __forceinline__ void grouped_body(const GroupedArgs a, int u, int nblock, int zz) { + if (u >= a.ucount[0]) return; + const int MT = (a.maxm + 15) / 16; + const int split = zz % a.sk; + const int mat = zz / a.sk; + const half* X = mat ? a.x1 : a.x0; + const int e = a.uids[u]; + const uint32_t* T = reinterpret_cast(mat ? a.tp1[e] : a.tp0[e]); + const int k2 = mat ? a.k2_1[e] : a.k2_0[e]; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int NTILES = a.ns > 0 ? a.ns : (a.n >> 4); // row stride of the trellis, in tiles + const int per_warp = (a.k >> 4) / a.sk / W; + const int kt0 = split * (per_warp * W) + warp * per_warp; + const int nt0 = nblock * NT; + __shared__ int rows_sh[16]; + __shared__ __align__(16) float red[W][16][NT * 16]; + + for (int mtile = 0; mtile < MT; ++mtile) { + if (threadIdx.x < 16) { + const int m = mtile * 16 + threadIdx.x; + const int code = m < a.maxm ? a.members[u * a.maxm + m] : -1; + rows_sh[threadIdx.x] = code >= 0 ? (code >> 5) * a.slots + (code & 31) : -1; + } + __syncthreads(); + if (rows_sh[0] < 0) return; // this tile and every later one are empty + const int r0 = rows_sh[g], r1 = rows_sh[g + 8]; + const half* x0 = X + (size_t)(r0 < 0 ? 0 : r0) * a.k + 2 * t; + const half* x1 = X + (size_t)(r1 < 0 ? 0 : r1) * a.k + 2 * t; + + float acc[NT][2][4]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][h][c] = 0.f; + + switch (k2) { +#define TF_EXL3X_CASE(K2_) \ + case K2_: \ + warp_tiles(T, NTILES, kt0, per_warp, nt0, x0, x1, r0 >= 0, r1 >= 0, lane, acc, \ + reinterpret_cast(red[warp])); \ + break; + TF_EXL3X_CASE(2) + TF_EXL3X_CASE(3) + TF_EXL3X_CASE(4) + TF_EXL3X_CASE(5) + TF_EXL3X_CASE(6) + TF_EXL3X_CASE(7) + TF_EXL3X_CASE(8) + TF_EXL3X_CASE(9) + TF_EXL3X_CASE(10) + TF_EXL3X_CASE(11) + TF_EXL3X_CASE(12) + TF_EXL3X_CASE(13) + TF_EXL3X_CASE(14) + TF_EXL3X_CASE(15) + TF_EXL3X_CASE(16) +#undef TF_EXL3X_CASE + default: + __trap(); + } + +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + red[warp][g][col] = acc[i][h][0]; + red[warp][g][col + 1] = acc[i][h][1]; + red[warp][g + 8][col] = acc[i][h][2]; + red[warp][g + 8][col + 1] = acc[i][h][3]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 16 * NT * 16; idx += W * 32) { + const int row = idx / (NT * 16), col = idx % (NT * 16); + const int r = rows_sh[row]; + if (r < 0) continue; + float s = red[0][row][col]; +#pragma unroll + for (int w = 1; w < W; ++w) s += red[w][row][col]; + a.z[(((size_t)mat * a.sk + split) * a.p + r) * a.n + nt0 * 16 + col] = s; + } + __syncthreads(); // rows_sh and red are rewritten by the next tile + } +} + +} // namespace tf_exl3 + +// --------------------------------------------------------------------------------------------------------------- + +// One per codebook (3inst, mcg, mul1: RED-SNOW's experts), any width; _nc: staged 16-byte loads (decode rows). +#define TF_EXL3X_GROUPED(CB) \ + extern "C" __global__ void tf_exl3_grouped_##CB(tf_exl3::GroupedArgs a) { \ + tf_exl3::grouped_body(a, blockIdx.x, blockIdx.y, blockIdx.z); \ + } \ + extern "C" __global__ void tf_exl3_grouped_nc_##CB(tf_exl3::GroupedArgs a) { \ + tf_exl3::grouped_body(a, blockIdx.x, blockIdx.y, blockIdx.z); \ + } +TF_EXL3X_GROUPED(0) +TF_EXL3X_GROUPED(1) +TF_EXL3X_GROUPED(2) +#undef TF_EXL3X_GROUPED + +namespace tf_exl3 { + +// Grouping in one block: distinct experts (< E) in id order, members row * 32 + slot in row order, -1 after the last. +constexpr int GROUP_THREADS = 1024; +constexpr int GROUP_PER_THREAD = 4; +constexpr int GROUP_WARPS = GROUP_THREADS / 32; + +struct GroupArgs { + const int* pick; // int32 [rows * slots] + int* uids; // int32 [nexp_max] + int* ucount; // int32 [1] + int* members; // int32 [nexp_max, maxm] + int rows, slots, experts, maxm; +}; + +// A stable counting sort: warps count their slices of picks, a prefix over warps places them, they land in order. +extern "C" __global__ void __launch_bounds__(GROUP_THREADS) tf_exl3_group(const GroupArgs a) { + extern __shared__ int sh[]; // cnt [warps, E], place [E], total by place [E] + __shared__ int warp_tot[GROUP_WARPS]; + __shared__ int nu; // distinct experts + const int E = a.experts, n = a.rows * a.slots; + int* cnt = sh; + int* place = sh + GROUP_WARPS * E; + int* total = place + E; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + for (int i = threadIdx.x; i < GROUP_WARPS * E; i += GROUP_THREADS) cnt[i] = 0; + __syncthreads(); + const int seg = (n + GROUP_WARPS - 1) / GROUP_WARPS, lo = warp * seg, hi = min(n, lo + seg); + for (int i = lo + lane; i < hi; i += 32) { + const int e = a.pick[i]; + if (e >= 0 && e < E) atomicAdd(&cnt[warp * E + e], 1); + } + __syncthreads(); + int tot[GROUP_PER_THREAD]; + int used = 0; +#pragma unroll + for (int q = 0; q < GROUP_PER_THREAD; ++q) { + const int e = threadIdx.x * GROUP_PER_THREAD + q; + int sum = 0; + if (e < E) + for (int w = 0; w < GROUP_WARPS; ++w) { // each slice's first member index: the earlier slices' counts + const int c = cnt[w * E + e]; + cnt[w * E + e] = sum; + sum += c; + } + tot[q] = sum; + used += sum > 0; + } + // exclusive scan of `used` over threads + int inc = used; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + int v = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) inc += v; + } + if (lane == 31) warp_tot[warp] = inc; + __syncthreads(); + if (warp == 0) { + int v = warp_tot[lane]; + int s = v; + for (int o = 1; o < 32; o <<= 1) { + int x = __shfl_up_sync(0xffffffffu, s, o); + if (lane >= o) s += x; + } + warp_tot[lane] = s - v; // exclusive per warp + if (lane == 31) { + a.ucount[0] = s; + nu = s; + } + } + __syncthreads(); + int at = warp_tot[warp] + inc - used; +#pragma unroll + for (int q = 0; q < GROUP_PER_THREAD; ++q) { + const int e = threadIdx.x * GROUP_PER_THREAD + q; + if (e >= E) continue; + place[e] = tot[q] > 0 ? at : -1; + if (tot[q] > 0) { + a.uids[at] = e; + total[at] = tot[q]; + ++at; + } + } + __syncthreads(); + for (int b = lo; b < hi; b += 32) { // a slice's picks in order, 32 at a time + const int i = b + lane; + int e = i < hi ? a.pick[i] : -1; + if (e >= E) e = -1; + const unsigned peers = __match_any_sync(0xffffffffu, e); + if (e >= 0) { + const int j = cnt[warp * E + e] + __popc(peers & ((1u << lane) - 1u)); + if (j < a.maxm) a.members[place[e] * a.maxm + j] = (i / a.slots) * 32 + (i % a.slots); + } + __syncwarp(); + if (e >= 0 && lane == 31 - __clz(peers)) cnt[warp * E + e] += __popc(peers); + __syncwarp(); + } + for (int idx = threadIdx.x; idx < nu * a.maxm; idx += GROUP_THREADS) { + const int u = idx / a.maxm, j = idx % a.maxm; + if (j >= total[u]) a.members[idx] = -1; + } +} + +// Program (slot, 128-block of K, matrix): Xh = fp16((x * suh) @ H) for gate and up of every routed slot (pick < E). +struct RotInArgs { + const void* x; // bf16 or fp16 [rows, x_stride] + int x_stride; + int x_dtype; // 0: fp16, 1: bf16 (exl3.zig DType) + const int* pick; // int32 [rows * slots] + const half* suh0; // fp16 [E, K], gate + const half* suh1; // fp16 [E, K], up + half* out0; // fp16 [P, K] + half* out1; + int k, slots, experts; +}; + +extern "C" __global__ void tf_exl3_experts_rot_in(const RotInArgs a) { + const int p = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int e = a.pick[p]; + if (e < 0 || e >= a.experts) return; + const int lane = threadIdx.x; + const int row = p / a.slots; + const half* suh = (mat ? a.suh1 : a.suh0) + (size_t)e * a.k + blk * 128 + 4 * lane; + const uint16_t* xr = reinterpret_cast(a.x) + (size_t)row * a.x_stride + blk * 128 + 4 * lane; + float v[4]; + if (a.x_dtype == 1) { + const __nv_bfloat16* xv = reinterpret_cast(xr); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xv[j]) * __half2float(suh[j]); + } else { + const half* xv = reinterpret_cast(xr); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __half2float(xv[j]) * __half2float(suh[j]); + } + fwht128(v, lane); + half* o = (mat ? a.out1 : a.out0) + (size_t)p * a.k + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +// Program (slot, 128-block of the width): splits summed in order, rotated, * svh, SwiGLU (0: bf16, 1: fp32), then Xd. +struct GateUpArgs { + const float* z; // [(mat * sk + split) * P + p] * N + col + const int* pick; // int32 [P] + const half* svh_g; // fp16 [E, N] + const half* svh_u; // fp16 [E, N] + const half* suh_d; // fp16 [E, N] + half* xd; // fp16 [P, N] + int p, n, sk, experts; + float limit; + int act_mode; +}; + +extern "C" __global__ void tf_exl3_gateup_epilogue(const GateUpArgs a) { + const int p = blockIdx.x, blk = blockIdx.y; + const int e = a.pick[p]; + if (e < 0 || e >= a.experts) return; + const int lane = threadIdx.x; + const int n = blk * 128 + 4 * lane; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float sg = 0.f, su = 0.f; + for (int s = 0; s < a.sk; ++s) { + sg += a.z[((size_t)(0 * a.sk + s) * a.p + p) * a.n + n + j]; + su += a.z[((size_t)(1 * a.sk + s) * a.p + p) * a.n + n + j]; + } + gv[j] = sg; + uv[j] = su; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float act; + if (a.act_mode == 0) { + float gg = fminf(bf16r(gv[j] * HAD_SCALE * __half2float(a.svh_g[(size_t)e * a.n + n + j])), a.limit); + float uu = fminf(fmaxf(bf16r(uv[j] * HAD_SCALE * __half2float(a.svh_u[(size_t)e * a.n + n + j])), + -a.limit), a.limit); + act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + } else { + float gg = fminf(gv[j] * HAD_SCALE * __half2float(a.svh_g[(size_t)e * a.n + n + j]), a.limit); + float uu = fminf(fmaxf(uv[j] * HAD_SCALE * __half2float(a.svh_u[(size_t)e * a.n + n + j]), -a.limit), + a.limit); + act = gg / (1.f + expf(-gg)) * uu; + } + v[j] = act * __half2float(a.suh_d[(size_t)e * a.n + n + j]); + } + fwht128(v, lane); + half* o = a.xd + (size_t)p * a.n + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (member row, 128-block of the width): down epilogue and combine in one launch, the same bits as the two. +struct DownArgs { + const float* z; // [split * P + p] * D + col + const int* pick; // int32 [P] + const half* svh_d; // fp16 [E, D] + float* y; // fp32 [P, D] + const float* wts; // fp32 [rows, slots] + float* out; // fp32 [rows, D] + int p, d, sk, experts, slots; +}; + +extern "C" __global__ void tf_exl3_down_combine(const DownArgs a) { + __shared__ float4 part[32][32]; // [slot][lane]: the slot's 4 outputs of the lane + const int r = blockIdx.x, blk = blockIdx.y; + const int k = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int n = blk * 128 + 4 * lane; + const int p = r * a.slots + k; + const int e = a.pick[p]; + float o[4]; + if (e >= 0 && e < a.experts) { + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int q = 0; q < a.sk; ++q) s += a.z[((size_t)q * a.p + p) * a.d + n + j]; + v[j] = s; + } + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { + o[j] = v[j] * HAD_SCALE * __half2float(a.svh_d[(size_t)e * a.d + n + j]); + a.y[(size_t)p * a.d + n + j] = o[j]; + } + } else { +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = a.y[(size_t)p * a.d + n + j]; + } + part[k][lane] = make_float4(o[0], o[1], o[2], o[3]); + __syncthreads(); + if (k != 0) return; + float acc[4] = {0.f, 0.f, 0.f, 0.f}; + for (int q = 0; q < a.slots; ++q) { + const float w = a.wts[r * a.slots + q]; + const float4 u = part[q][lane]; + acc[0] = fmaf(w, u.x, acc[0]); + acc[1] = fmaf(w, u.y, acc[1]); + acc[2] = fmaf(w, u.z, acc[2]); + acc[3] = fmaf(w, u.w, acc[3]); + } +#pragma unroll + for (int j = 0; j < 4; ++j) a.out[(size_t)r * a.d + n + j] = acc[j]; +} + +// The down projection's epilogue alone (experts.cu down_epilogue_kernel): y [P, D] fp32, invalid picks untouched. +extern "C" __global__ void tf_exl3_down_epilogue(const DownArgs a) { + const int p = blockIdx.x, blk = blockIdx.y, lane = threadIdx.x; + const int n = blk * 128 + 4 * lane; + const int e = a.pick[p]; + if (e < 0 || e >= a.experts) return; + float v[4]; + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int q = 0; q < a.sk; ++q) s += a.z[((size_t)q * a.p + p) * a.d + n + j]; + v[j] = s; + } + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) a.y[(size_t)p * a.d + n + j] = v[j] * HAD_SCALE * __half2float(a.svh_d[(size_t)e * a.d + n + j]); +} + +} // namespace tf_exl3 diff --git a/zig/kernels/cuda/exl3_linear.cu b/zig/kernels/cuda/exl3_linear.cu new file mode 100644 index 0000000000..0a52a1e716 --- /dev/null +++ b/zig/kernels/cuda/exl3_linear.cu @@ -0,0 +1,471 @@ +// EXL3 decode-row linear: any codebook, width, 1-128 rows, a row's bits only its own; port of tensorfold's linear.cu. +#include +#include +#include + +#include "exl3_decode.cuh" + +namespace tf_exl3 { + +enum DType : int { F16 = 0, BF16 = 1, F32 = 2 }; + +__device__ __forceinline__ int imin(int a, int b) { return a < b ? a : b; } + +struct LinearArgs { + const void* xh; // fp16 [m, k], the hadamard-rotated inputs + const uint32_t* t; // trellis words + long long stride_k; // words between k tiles + long long stride_nb; // words between 128-column blocks + const half* svh; // fp16 [n], output hadamard scales + const half* bias; // fp16 [n] or null + void* y; // output + float* z; // fp32 split partials [sk, m, n] or null + int* counters; // per (pass, block) arrival counts, zeroed before the launch + int y_dtype; + int n, sk, m, k; +}; + +struct RotInArgs { + const void* x; // f16 / bf16 / f32 [m, k] + const half* suh; // fp16 [k], input hadamard scales + half* xh; // fp16 [m, k] out + int x_dtype; + int m, k; +}; + +__device__ __forceinline__ void load4(const void* p, int dtype, size_t i, float (&v)[4]) { + if (dtype == F32) { + const float4 u = *reinterpret_cast(static_cast(p) + i); + v[0] = u.x; v[1] = u.y; v[2] = u.z; v[3] = u.w; + } else if (dtype == BF16) { + const uint2 u = *reinterpret_cast(static_cast(p) + i); + const float2 a = __bfloat1622float2(*reinterpret_cast(&u.x)); + const float2 b = __bfloat1622float2(*reinterpret_cast(&u.y)); + v[0] = a.x; v[1] = a.y; v[2] = b.x; v[3] = b.y; + } else { + const uint2 u = *reinterpret_cast(static_cast(p) + i); + const float2 a = __half22float2(*reinterpret_cast(&u.x)); + const float2 b = __half22float2(*reinterpret_cast(&u.y)); + v[0] = a.x; v[1] = a.y; v[2] = b.x; v[3] = b.y; + } +} + +__device__ __forceinline__ void store4(void* p, int dtype, size_t i, const float (&v)[4]) { + if (dtype == F32) { + *reinterpret_cast(static_cast(p) + i) = make_float4(v[0], v[1], v[2], v[3]); + } else if (dtype == BF16) { + __nv_bfloat162 a = __floats2bfloat162_rn(v[0], v[1]), b = __floats2bfloat162_rn(v[2], v[3]); + uint2 u; + u.x = *reinterpret_cast(&a); + u.y = *reinterpret_cast(&b); + *reinterpret_cast(static_cast<__nv_bfloat16*>(p) + i) = u; + } else { + half2 a = __floats2half2_rn(v[0], v[1]), b = __floats2half2_rn(v[2], v[3]); + uint2 u; + u.x = *reinterpret_cast(&a); + u.y = *reinterpret_cast(&b); + *reinterpret_cast(static_cast(p) + i) = u; + } +} + +// The finished outputs of one row's 128 columns from their fp32 sums (4 a lane): H / sqrt(128), * svh, + bias. +__device__ __forceinline__ void finish(float (&v)[4], int lane, const half* svh, const half* bias, int col) { + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { + v[j] = v[j] * HAD_SCALE * __half2float(__ldg(svh + col + j)); + // Speculative loads must have a valid address even when the split-K result has no bias. + const half bv = __ldg((bias ? bias : svh) + col + j); + if (bias) v[j] += __half2float(bv); + } +} + +// A k step's words for one lane: the step is 8 tiles back to back, each lane takes the same slice of every tile. +template +__device__ __forceinline__ void load_step(const uint32_t* step, int lane, uint32_t (&raw)[8 * lane_words()]) { + constexpr int TW = tile_words(), LW = lane_words(); + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int j = 0; j < 8; ++j) +#pragma unroll + for (int q = 0; q < LW; ++q) raw[j * LW + q] = __ldg(step + j * TW + (word + q) % TW); +} + +// 16 bytes of trellis words, read-only and not kept in L1. +__device__ __forceinline__ uint4 ldg_nc_v4(const uint32_t* p) { + uint4 v; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p)); + return v; +} + +// Tile j's lane words from a loaded step. +template +__device__ __forceinline__ void step_lane_words(const uint32_t (&raw)[8 * lane_words()], int j, + uint32_t (&w)[lane_words()]) { +#pragma unroll + for (int q = 0; q < lane_words(); ++q) w[q] = raw[j * lane_words() + q]; +} + +// xh = fp16(((x * suh) @ H) / sqrt(128)): one warp per 128 columns, one block of 4 warps per 512. +extern "C" __global__ void __launch_bounds__(128) tf_exl3_rot_in(RotInArgs a) { + const int blk = blockIdx.x * 4 + (threadIdx.x >> 5), row = blockIdx.y, lane = threadIdx.x & 31; + if (blk * 128 >= a.k) return; + const int k = blk * 128 + 4 * lane; + float v[4], s[4]; + load4(a.x, a.x_dtype, (size_t)row * a.k + k, v); + load4(a.suh, F16, k, s); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] *= s[j]; + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] *= HAD_SCALE; + store4(a.xh, F16, (size_t)row * a.k + k, v); +} + + +// One program: (N / 128) * SK blocks, split-major like a (N / 128, SK) grid, WK warps each. +template +__device__ __forceinline__ void linear_body(LinearArgs a) { + constexpr int TW = tile_words(), LW = lane_words(); + constexpr bool ST = TW % 8 == 0; // even widths: staged 16-byte loads, PD steps in flight + constexpr int PD = 2; + constexpr bool PF = !ST && K2 <= 12; // odd widths: the next step's lane words in registers + constexpr int SR = 8 * LW; + extern __shared__ __align__(16) float red[]; // WK * RH * 128 floats; even widths stage steps here + __shared__ int last; + const int M = a.m, K = a.k, N = a.n, SK = a.sk, NB = N >> 7; + const int RH = imin(M, 8); // rows of red a warp + const int nb = (int)blockIdx.x % NB, split = (int)blockIdx.x / NB; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int per_warp = (K >> 4) / SK / WK; + const int kt0 = split * (per_warp * WK) + warp * per_warp; + const uint32_t* tiles = a.t + nb * a.stride_nb; + const int col0 = nb * 128; + + for (int m0 = 0, pass = 0; m0 < M; m0 += 16, ++pass) { + const int R = imin(16, M - m0); + + float acc[8][2][4]; +#pragma unroll + for (int i = 0; i < 8; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][h][c] = 0.f; + + // the walk's rows for the mma, clamped so rows past the pass read inside the buffer (their outputs are dropped) + const int r0 = m0 + (g < R ? g : R - 1), r1 = m0 + (g + 8 < R ? g + 8 : R - 1); + const half* x0 = static_cast(a.xh) + (size_t)r0 * K; + const half* x1 = static_cast(a.xh) + (size_t)r1 * K; + const uint32_t* tile = tiles + (size_t)kt0 * a.stride_k; + uint32_t cur[PF ? SR : 1], nxt[PF ? SR : 1]; + if constexpr (PF) load_step(tile, lane, cur); // the weights do not wait for the previous kernel + constexpr int NV = (8 * TW + 127) / 128; // a staged step's 16-byte loads a lane + uint4 rv[ST ? PD : 1][ST ? NV : 1]; + if constexpr (ST) { +#pragma unroll + for (int d = 0; d < PD; ++d) + if (d < per_warp) +#pragma unroll + for (int v = 0; v < NV; ++v) + if (128 * v + 4 * lane < 8 * TW) + rv[d][v] = ldg_nc_v4(tile + d * a.stride_k + 128 * v + 4 * lane); + } + + // x's fragments of k step kt: rows g (and g + 8): [0] [2] k 2t.., 2t + 8..; [1] [3] the same of row g + 8 + auto load_x = [&](int kt, uint32_t (&u)[4]) { + u[0] = __ldg(reinterpret_cast(x0 + kt * 16 + 2 * t)); + u[2] = __ldg(reinterpret_cast(x0 + kt * 16 + 2 * t + 8)); + u[1] = __ldg(reinterpret_cast(x1 + kt * 16 + 2 * t)); + u[3] = __ldg(reinterpret_cast(x1 + kt * 16 + 2 * t + 8)); + }; + if constexpr (ST) { + uint32_t* stage = reinterpret_cast(red) + warp * 8 * 64; // a warp's step: 8 tiles, 64 words most + int word, offset; + lane_start(lane, word, offset); +#pragma unroll 1 + for (int ib = 0; ib < per_warp; ib += PD) { +#pragma unroll + for (int d = 0; d < PD; ++d) { + const int i = ib + d; + if (i >= per_warp) break; + __syncwarp(); // every lane has read the previous step back +#pragma unroll + for (int v = 0; v < NV; ++v) + if (128 * v + 4 * lane < 8 * TW) + *reinterpret_cast(stage + 128 * v + 4 * lane) = rv[d][v]; + if (i + PD < per_warp) +#pragma unroll + for (int v = 0; v < NV; ++v) + if (128 * v + 4 * lane < 8 * TW) + rv[d][v] = ldg_nc_v4(tile + (size_t)(i + PD) * a.stride_k + 128 * v + 4 * lane); + __syncwarp(); // the step is staged before any lane reads it + uint32_t u[4]; + load_x(kt0 + i, u); +#pragma unroll + for (int j = 0; j < 8; ++j) { + uint32_t w[LW]; +#pragma unroll + for (int q = 0; q < LW; ++q) w[q] = stage[j * TW + (word + q) % TW]; + uint32_t b0[2], b1[2]; + decode_lane(w, lane, b0, b1); + mma16816(acc[j][0], u, b0); + mma16816(acc[j][1], u, b1); + } + } + } + } +#pragma unroll 1 + for (int i = 0; i < (ST ? 0 : per_warp); ++i) { + uint32_t u[4]; + load_x(kt0 + i, u); + if constexpr (PF) { + if (i + 1 < per_warp) load_step(tile + (size_t)(i + 1) * a.stride_k, lane, nxt); + } +#pragma unroll + for (int j = 0; j < 8; ++j) { + uint32_t w[LW]; + if constexpr (PF) step_lane_words(cur, j, w); + else ldg_lane_words(tile + (size_t)i * a.stride_k + j * TW, lane, w); + uint32_t b0[2], b1[2]; + decode_lane(w, lane, b0, b1); + mma16816(acc[j][0], u, b0); + mma16816(acc[j][1], u, b1); + } + if constexpr (PF) { +#pragma unroll + for (int q = 0; q < SR; ++q) cur[q] = nxt[q]; + } + } + + // the warps' sums, added in warp order, rows 0-7 of the pass and then rows 8-15 + for (int rlo = 0; rlo < R; rlo += 8) { + const int rn = imin(R - rlo, 8); + __syncthreads(); // red is reused by every half and pass + if (g < RH) { +#pragma unroll + for (int i = 0; i < 8; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + *reinterpret_cast(red + (warp * RH + g) * 128 + col) = + rlo ? make_float2(acc[i][h][2], acc[i][h][3]) + : make_float2(acc[i][h][0], acc[i][h][1]); + } + } + __syncthreads(); + + if (SK == 1) { + for (int r = warp; r < rn; r += WK) { + float v[4]; + const float4 u = *reinterpret_cast(red + r * 128 + 4 * lane); + v[0] = u.x; v[1] = u.y; v[2] = u.z; v[3] = u.w; +#pragma unroll + for (int w = 1; w < WK; ++w) { + const float4 q = *reinterpret_cast(red + (w * RH + r) * 128 + 4 * lane); + v[0] += q.x; v[1] += q.y; v[2] += q.z; v[3] += q.w; + } + finish(v, lane, a.svh, a.bias, col0 + 4 * lane); + store4(a.y, a.y_dtype, (size_t)(m0 + rlo + r) * N + col0 + 4 * lane, v); + } + } else { + for (int idx = threadIdx.x; idx < rn * 32; idx += WK * 32) { + const int r = idx >> 5, c = 4 * (idx & 31); + float4 s = *reinterpret_cast(red + r * 128 + c); +#pragma unroll + for (int w = 1; w < WK; ++w) { + const float4 q = *reinterpret_cast(red + (w * RH + r) * 128 + c); + s.x += q.x; s.y += q.y; s.z += q.z; s.w += q.w; + } + *reinterpret_cast(a.z + ((size_t)split * M + m0 + rlo + r) * N + col0 + c) = s; + } + } + } + if (SK > 1) { + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) last = atomicAdd(a.counters + pass * NB + nb, 1) == SK - 1; + __syncthreads(); + if (last) { + __threadfence(); + for (int r = warp; r < R; r += WK) { + const size_t at = ((size_t)m0 + r) * N + col0 + 4 * lane; + float4 s = __ldcg(reinterpret_cast(a.z + at)); + for (int q = 1; q < SK; ++q) { + const float4 u = __ldcg(reinterpret_cast(a.z + (size_t)q * M * N + at)); + s.x += u.x; s.y += u.y; s.z += u.z; s.w += u.w; + } + float v[4] = {s.x, s.y, s.z, s.w}; + finish(v, lane, a.svh, a.bias, col0 + 4 * lane); + store4(a.y, a.y_dtype, (size_t)(m0 + r) * N + col0 + 4 * lane, v); + } + if (threadIdx.x == 0) a.counters[pass * NB + nb] = 0; // every program of the block has arrived + } + } + __syncthreads(); // red is reused in the next pass + } +} + +// A prompt block of 64 rows (mul1): warp j owns column tile j for every K range, so x's tiles are staged once in shared +constexpr int PR = 64; // rows a block +constexpr int PS = 132; // psum's row stride (floats): rows apart in banks, 16-byte aligned +constexpr int PKC = 2; // k tiles a stage +constexpr int PXS = 24; // a staged x row's stride (halves): ldmatrix rows on distinct banks +constexpr int PROMPT_SHARED = PR * PS * 4 + 2 * PKC * PR * PXS * 2; + +__device__ __forceinline__ void cp_async16(void* dst, const void* src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"((unsigned)__cvta_generic_to_shared(dst)), "l"(src)); +} + +// The same partial sums per K range as linear_body, added into psum in range order: its bits, each tile decoded once. +template +__device__ __forceinline__ void prompt_body(LinearArgs a) { + constexpr int TW = tile_words(), LW = lane_words(); + extern __shared__ __align__(16) float psum[]; // [PR][PS], then x stages [2][PKC][PR][PXS] fp16 + half* xs = reinterpret_cast(psum + PR * PS); + __shared__ int last; + const int M = a.m, K = a.k, N = a.n, SK = a.sk, NB = N >> 7; + const int row0 = blockIdx.y * PR, MB = imin(PR, M - row0), MT = (MB + 15) >> 4; + const int nb = (int)blockIdx.x % NB, split = (int)blockIdx.x / NB; + const int j = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + const int per_warp = (K >> 4) / SK / WK, nkt = per_warp * WK, kt0 = split * nkt; + const uint32_t* tiles = a.t + nb * a.stride_nb + j * TW; + const half* xh = static_cast(a.xh); + const int col0 = nb * 128; + auto stage = [&](int c) { // k tiles c * PKC.. into buffer c & 1: one 16-byte piece a thread + const int kk = threadIdx.x >> 7, row = (threadIdx.x >> 1) & (PR - 1), part = threadIdx.x & 1; + const half* src = xh + (size_t)(row0 + imin(row, MB - 1)) * K + (kt0 + c * PKC + kk) * 16 + part * 8; + cp_async16(xs + (((c & 1) * PKC + kk) * PR + row) * PXS + part * 8, src); + asm volatile("cp.async.commit_group;"); + }; + float acc[4][2][4]; + uint32_t w[LW], wn[LW]; + ldg_lane_words(tiles + (size_t)kt0 * a.stride_k, lane, w); + stage(0); + const int chunks = nkt / PKC; +#pragma unroll 1 + for (int c = 0; c < chunks; ++c) { + if (c + 1 < chunks) { + stage(c + 1); + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + __syncthreads(); +#pragma unroll + for (int kk = 0; kk < PKC; ++kk) { + const int i = c * PKC + kk; + if (i % per_warp == 0) { +#pragma unroll + for (int q = 0; q < 4; ++q) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[q][h][e] = 0.f; + } + if (i + 1 < nkt) ldg_lane_words(tiles + (size_t)(kt0 + i + 1) * a.stride_k, lane, wn); + uint32_t b0[2], b1[2]; + decode_lane(w, lane, b0, b1); + const half* xb = xs + (((c & 1) * PKC + kk) * PR + (lane & 15)) * PXS + (lane >> 4) * 8; +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (q < MT) { + uint32_t u[4]; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(u[0]), "=r"(u[1]), "=r"(u[2]), "=r"(u[3]) + : "r"((unsigned)__cvta_generic_to_shared(xb + q * 16 * PXS))); + mma16816(acc[q][0], u, b0); + mma16816(acc[q][1], u, b1); + } + } +#pragma unroll + for (int q = 0; q < LW; ++q) w[q] = wn[q]; + if ((i + 1) % per_warp == 0) { // a K range done: into psum, the first range stored, the rest added + const bool first = i + 1 == per_warp; +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (q >= MT) continue; +#pragma unroll + for (int h = 0; h < 2; ++h) { + float* p0 = psum + (16 * q + g) * PS + j * 16 + h * 8 + 2 * t; + float* p1 = p0 + 8 * PS; + p0[0] = first ? acc[q][h][0] : p0[0] + acc[q][h][0]; + p0[1] = first ? acc[q][h][1] : p0[1] + acc[q][h][1]; + p1[0] = first ? acc[q][h][2] : p1[0] + acc[q][h][2]; + p1[1] = first ? acc[q][h][3] : p1[1] + acc[q][h][3]; + } + } + } + } + __syncthreads(); // the stage is refilled two chunks on + } + if (SK == 1) { + for (int r = j; r < MB; r += 8) { + const float4 u = *reinterpret_cast(psum + r * PS + 4 * lane); + float v[4] = {u.x, u.y, u.z, u.w}; + finish(v, lane, a.svh, a.bias, col0 + 4 * lane); + store4(a.y, a.y_dtype, (size_t)(row0 + r) * N + col0 + 4 * lane, v); + } + return; + } + for (int r = j; r < MB; r += 8) + *reinterpret_cast(a.z + ((size_t)split * M + row0 + r) * N + col0 + 4 * lane) = + *reinterpret_cast(psum + r * PS + 4 * lane); + __threadfence(); + __syncthreads(); + int* counter = a.counters + blockIdx.y * NB + nb; + if (threadIdx.x == 0) last = atomicAdd(counter, 1) == SK - 1; + __syncthreads(); + if (!last) return; + __threadfence(); + for (int r = j; r < MB; r += 8) { + const size_t at = (size_t)(row0 + r) * N + col0 + 4 * lane; + float4 s = __ldcg(reinterpret_cast(a.z + at)); + for (int q = 1; q < SK; ++q) { + const float4 u = __ldcg(reinterpret_cast(a.z + (size_t)q * M * N + at)); + s.x += u.x; s.y += u.y; s.z += u.z; s.w += u.w; + } + float v[4] = {s.x, s.y, s.z, s.w}; + finish(v, lane, a.svh, a.bias, col0 + 4 * lane); + store4(a.y, a.y_dtype, at, v); + } + if (threadIdx.x == 0) *counter = 0; // every program of the block has arrived +} + +// One extern "C" kernel per instantiation: the fatbin carries __global__ functions, the host resolves them by name. +#define TF_EXL3_LINEAR(K2, CB, WK) \ + extern "C" __global__ void __launch_bounds__((WK) * 32) \ + tf_exl3_linear_##K2##_##CB##_##WK(tf_exl3::LinearArgs a) { \ + tf_exl3::linear_body(a); \ + } + +// Every width the format serves (linear.cu TF_EXL3_ALL plus 4.5-6.5): whole bits, all codebooks; 7.5 is not a width. +#define TF_EXL3_WIDTHS(K2, CB) TF_EXL3_LINEAR(K2, CB, 4) TF_EXL3_LINEAR(K2, CB, 8) +#define TF_EXL3_CODEBOOK(CB) \ + TF_EXL3_WIDTHS(2, CB) TF_EXL3_WIDTHS(4, CB) TF_EXL3_WIDTHS(6, CB) TF_EXL3_WIDTHS(8, CB) \ + TF_EXL3_WIDTHS(10, CB) TF_EXL3_WIDTHS(12, CB) TF_EXL3_WIDTHS(14, CB) TF_EXL3_WIDTHS(16, CB) +TF_EXL3_CODEBOOK(2) // mul1 +TF_EXL3_WIDTHS(3, 2) TF_EXL3_WIDTHS(5, 2) TF_EXL3_WIDTHS(7, 2) TF_EXL3_WIDTHS(9, 2) +TF_EXL3_WIDTHS(11, 2) TF_EXL3_WIDTHS(13, 2) +TF_EXL3_CODEBOOK(1) // mcg +TF_EXL3_CODEBOOK(0) // inst3 +#undef TF_EXL3_CODEBOOK +#undef TF_EXL3_WIDTHS + +// The prompt kernels: mul1, every width, 4 and 8 warps' K ranges, 8 warps a block. +#define TF_EXL3_PROMPT(K2) \ + extern "C" __global__ void __launch_bounds__(256, 2) tf_exl3_prompt_##K2##_4(tf_exl3::LinearArgs a) { \ + tf_exl3::prompt_body(a); \ + } \ + extern "C" __global__ void __launch_bounds__(256, 2) tf_exl3_prompt_##K2##_8(tf_exl3::LinearArgs a) { \ + tf_exl3::prompt_body(a); \ + } +TF_EXL3_PROMPT(2) TF_EXL3_PROMPT(3) TF_EXL3_PROMPT(4) TF_EXL3_PROMPT(5) TF_EXL3_PROMPT(6) TF_EXL3_PROMPT(7) +TF_EXL3_PROMPT(8) TF_EXL3_PROMPT(9) TF_EXL3_PROMPT(10) TF_EXL3_PROMPT(11) TF_EXL3_PROMPT(12) TF_EXL3_PROMPT(13) +TF_EXL3_PROMPT(14) TF_EXL3_PROMPT(16) +#undef TF_EXL3_PROMPT +} // namespace tf_exl3 diff --git a/zig/kernels/cuda/glm_bmm.cu b/zig/kernels/cuda/glm_bmm.cu new file mode 100644 index 0000000000..0d055e8273 --- /dev/null +++ b/zig/kernels/cuda/glm_bmm.cu @@ -0,0 +1,316 @@ +// GLM's BF16 products (qmm.py _bmm, glue.py router): mma.sync m16n8k16, four K=16 chunks per 64-wide step. +#include +#include + +namespace tf_glm_bmm { + +using bf16 = __nv_bfloat16; + +struct BmmArgs { + const bf16* x; // bf16 [m, k], rows `xs` apart + long long xs; // x's row stride in elements + const bf16* w; // bf16 [n, k] row-major + void* out; // bf16 or fp32 [m, n] when one slice + float* part; // fp32 [sk, m, n] partials when split + int m, n, k, sk; +}; + +struct RouterArgs { + const bf16* x; // bf16 [m, d], rows `xs` apart + long long xs; + const bf16* w; // bf16 [ne, d] row-major + float* out; // fp32 [m, ne] when the whole K ran in one pass + float* part; // fp32 [ks, m, ne] partials, one per slice + int m, ne, d, ks; +}; + +__device__ __forceinline__ unsigned pack2(bf16 lo, bf16 hi) { + return static_cast(__bfloat16_as_ushort(lo)) | (static_cast(__bfloat16_as_ushort(hi)) << 16); +} + +__device__ __forceinline__ void mma_16n8k16(float (&d)[4], const unsigned (&a)[4], const unsigned (&b)[2]) { + 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};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// One 16x16 A tile of x (row-major, `xs` elements a row): four fragment words, rows at or past `m` zeroed. +__device__ __forceinline__ void load_a16(const bf16* x, long long xs, int row, int col, int m, unsigned (&a)[4]) { + const int lane = static_cast(threadIdx.x) & 31; + const int r0 = lane >> 2, c0 = (lane & 3) * 2; + const bf16* base = x + static_cast(row) * xs + col; + const bool ok0 = row + r0 < m, ok1 = row + r0 + 8 < m; + a[0] = ok0 ? *reinterpret_cast(base + static_cast(r0) * xs + c0) : 0u; + a[1] = ok1 ? *reinterpret_cast(base + static_cast(r0 + 8) * xs + c0) : 0u; + a[2] = ok0 ? *reinterpret_cast(base + static_cast(r0) * xs + c0 + 8) : 0u; + a[3] = ok1 ? *reinterpret_cast(base + static_cast(r0 + 8) * xs + c0 + 8) : 0u; +} + +// One 16x8 B tile of w [n, k] (weight rows, K contiguous): two fragment words, columns at or past `n` zeroed. +__device__ __forceinline__ void load_b16(const bf16* w, long long k_stride, int n0, int k0, int n, unsigned (&b)[2]) { + const int lane = static_cast(threadIdx.x) & 31; + const int kk = (lane & 3) * 2, n_row = lane >> 2; + const bool ok = n0 + n_row < n; + const bf16* base = w + static_cast(n0 + n_row) * k_stride + k0 + kk; + b[0] = ok ? *reinterpret_cast(base) : 0u; + b[1] = ok ? *reinterpret_cast(base + 8) : 0u; +} + +// A warp's 16x8 result tile: `row` is where it lands in the output, `mrow` the unmasked input row. +__device__ __forceinline__ void store_tile(void* out, bool f32_out, long long out_stride, int row, int mrow, int col, + int m, int n, const float (&d)[4]) { + const int lane = static_cast(threadIdx.x) & 31; + const int r0 = lane >> 2, c0 = (lane & 3) * 2; + const bool ok0 = mrow + r0 < m, ok1 = mrow + r0 + 8 < m; + const bool k0 = col + c0 < n, k1 = col + c0 + 1 < n; + if (f32_out) { + float* base = static_cast(out) + static_cast(row) * out_stride + col; + if (ok0 && k0) base[static_cast(r0) * out_stride + c0] = d[0]; + if (ok0 && k1) base[static_cast(r0) * out_stride + c0 + 1] = d[1]; + if (ok1 && k0) base[static_cast(r0 + 8) * out_stride + c0] = d[2]; + if (ok1 && k1) base[static_cast(r0 + 8) * out_stride + c0 + 1] = d[3]; + } else { + bf16* base = static_cast(out) + static_cast(row) * out_stride + col; + if (ok0 && k0 && k1) { + *reinterpret_cast(base + static_cast(r0) * out_stride + c0) = + pack2(__float2bfloat16_rn(d[0]), __float2bfloat16_rn(d[1])); + } + if (ok1 && k0 && k1) { + *reinterpret_cast(base + static_cast(r0 + 8) * out_stride + c0) = + pack2(__float2bfloat16_rn(d[2]), __float2bfloat16_rn(d[3])); + } + } +} + +// A packed w's 16x8 B tiles of one 64-wide K step: lane's (kk, half) words, two 16-byte loads (fragment-order layout). +__device__ __forceinline__ void load_b_packed(const bf16* w, int k, int n0, int k0, int n, unsigned (&b)[8]) { + const int lane = static_cast(threadIdx.x) & 31; + if (n0 >= n) { +#pragma unroll + for (int i = 0; i < 8; i++) b[i] = 0u; + return; + } + const uint4* p = reinterpret_cast(w) + (static_cast(n0 >> 3) * (k >> 6) + (k0 >> 6)) * 64 + lane * 2; + const uint4 u0 = __ldg(p), u1 = __ldg(p + 1); + b[0] = u0.x; b[1] = u0.y; b[2] = u0.z; b[3] = u0.w; + b[4] = u1.x; b[5] = u1.y; b[6] = u1.z; b[7] = u1.w; +} + +// The bf16 matmul: y = x [m, k] through w [n, k], 64-wide K steps in order, `sk` slices; PK: w in fragment order. +template +__device__ __forceinline__ void bmm_body(const BmmArgs a) { + constexpr int nt = 64 / (WN * 8); // 16x8 B tiles per warp + constexpr int mt = BM / 16; + const int warp = static_cast(threadIdx.x) >> 5; + const int per = a.k / a.sk; // one slice's K + const int col0 = blockIdx.y * 64 + warp * (64 / WN); + const int row0 = blockIdx.x * BM; + const int k0_base = blockIdx.z * per; + + float acc[mt][nt][4]; +#pragma unroll + for (int r = 0; r < mt; r++) +#pragma unroll + for (int c = 0; c < nt; c++) +#pragma unroll + for (int i = 0; i < 4; i++) acc[r][c][i] = 0.0f; + + for (int k0 = k0_base; k0 < k0_base + per; k0 += 64) { + unsigned bp[PK ? nt : 1][8]; + if constexpr (PK) { +#pragma unroll + for (int c = 0; c < nt; c++) load_b_packed(a.w, a.k, col0 + c * 8, k0, a.n, bp[c]); + } +#pragma unroll + for (int kk = 0; kk < 4; kk++) { + const int kcol = k0 + kk * 16; + unsigned av[mt][4]; +#pragma unroll + for (int r = 0; r < mt; r++) load_a16(a.x, a.xs, row0 + r * 16, kcol, a.m, av[r]); +#pragma unroll + for (int c = 0; c < nt; c++) { + unsigned b[2]; + if constexpr (PK) { + b[0] = bp[c][2 * kk]; + b[1] = bp[c][2 * kk + 1]; + } else { + load_b16(a.w, a.k, col0 + c * 8, kcol, a.n, b); + } +#pragma unroll + for (int r = 0; r < mt; r++) mma_16n8k16(acc[r][c], av[r], b); + } + } + } + + if (a.sk > 1) { // fp32 partials, slice s at rows [s*m, s*m + m) +#pragma unroll + for (int r = 0; r < mt; r++) +#pragma unroll + for (int c = 0; c < nt; c++) + store_tile(a.part, true, a.n, blockIdx.z * a.m + row0 + r * 16, row0 + r * 16, col0 + c * 8, a.m, a.n, + acc[r][c]); + } else { +#pragma unroll + for (int r = 0; r < mt; r++) +#pragma unroll + for (int c = 0; c < nt; c++) + store_tile(a.out, F32_OUT, a.n, row0 + r * 16, row0 + r * 16, col0 + c * 8, a.m, a.n, acc[r][c]); + } +} + +// _bmm's split partials added in slice order, then one bf16 rounding (or fp32 through). +__device__ __forceinline__ void reduce_body(const float* part, void* out, int total, int sk, int f32) { + const int i = blockIdx.x * 1024 + static_cast(threadIdx.x); + if (i >= total) return; + float acc = part[i]; + for (int s = 1; s < sk; s++) acc += part[s * total + i]; + if (f32) static_cast(out)[i] = acc; + else static_cast(out)[i] = __float2bfloat16_rn(acc); +} + +// glue._router: OUT[m, e] fp32 = x[m] . w[e] over the whole K, 64-wide steps in order (the unsliced choice). +template +__device__ __forceinline__ void router_body(const RouterArgs a) { + constexpr int mt = BM / 16; + const int warp = static_cast(threadIdx.x) >> 5; + const int col0 = blockIdx.y * 64 + warp * 16; + const int row0 = blockIdx.x * BM; + + float acc[mt][2][4]; +#pragma unroll + for (int r = 0; r < mt; r++) +#pragma unroll + for (int i = 0; i < 8; i++) acc[r][i >> 2][i & 3] = 0.0f; + for (int k0 = 0; k0 < a.d; k0 += 64) { + unsigned bp[PK ? 2 : 1][8]; + if constexpr (PK) { + load_b_packed(a.w, a.d, col0, k0, a.ne, bp[0]); + load_b_packed(a.w, a.d, col0 + 8, k0, a.ne, bp[PK ? 1 : 0]); + } +#pragma unroll + for (int kk = 0; kk < 4; kk++) { + unsigned av[mt][4]; +#pragma unroll + for (int r = 0; r < mt; r++) load_a16(a.x, a.xs, row0 + r * 16, k0 + kk * 16, a.m, av[r]); + unsigned b[2]; + if constexpr (PK) { + b[0] = bp[0][2 * kk]; + b[1] = bp[0][2 * kk + 1]; + } else { + load_b16(a.w, a.d, col0, k0 + kk * 16, a.ne, b); + } +#pragma unroll + for (int r = 0; r < mt; r++) mma_16n8k16(acc[r][0], av[r], b); + if constexpr (PK) { + b[0] = bp[PK ? 1 : 0][2 * kk]; + b[1] = bp[PK ? 1 : 0][2 * kk + 1]; + } else { + load_b16(a.w, a.d, col0 + 8, k0 + kk * 16, a.ne, b); + } +#pragma unroll + for (int r = 0; r < mt; r++) mma_16n8k16(acc[r][1], av[r], b); + } + } +#pragma unroll + for (int r = 0; r < mt; r++) +#pragma unroll + for (int c = 0; c < 2; c++) + store_tile(a.out, true, a.ne, row0 + r * 16, row0 + r * 16, col0 + c * 8, a.m, a.ne, acc[r][c]); +} + +// glue._router_part: the same over one of `ks` equal K slices; PART[s, m, e] fp32. +template +__device__ __forceinline__ void router_part_body(const RouterArgs a) { + constexpr int mt = BM / 16; + const int warp = static_cast(threadIdx.x) >> 5; + const int col0 = blockIdx.y * 32 + warp * 8; + const int row0 = blockIdx.x * BM; + const int per = a.d / a.ks; + const int k0_base = blockIdx.z * per; + + float acc[mt][4]; +#pragma unroll + for (int r = 0; r < mt; r++) +#pragma unroll + for (int i = 0; i < 4; i++) acc[r][i] = 0.0f; + for (int k0 = k0_base; k0 < k0_base + per; k0 += 64) { + unsigned bp[8]; + if constexpr (PK) load_b_packed(a.w, a.d, col0, k0, a.ne, bp); +#pragma unroll + for (int kk = 0; kk < 4; kk++) { + unsigned av[mt][4]; +#pragma unroll + for (int r = 0; r < mt; r++) load_a16(a.x, a.xs, row0 + r * 16, k0 + kk * 16, a.m, av[r]); + unsigned b[2]; + if constexpr (PK) { + b[0] = bp[2 * kk]; + b[1] = bp[2 * kk + 1]; + } else { + load_b16(a.w, a.d, col0, k0 + kk * 16, a.ne, b); + } +#pragma unroll + for (int r = 0; r < mt; r++) mma_16n8k16(acc[r], av[r], b); + } + } +#pragma unroll + for (int r = 0; r < mt; r++) + store_tile(a.part, true, a.ne, blockIdx.z * a.m + row0 + r * 16, row0 + r * 16, col0, a.m, a.ne, acc[r]); +} + +// glue._router_sum: the `ks` slice partials added in slice order. +__device__ __forceinline__ void router_sum_body(const float* part, float* out, int total, int ks) { + const int i = blockIdx.x * 1024 + static_cast(threadIdx.x); + if (i >= total) return; + float acc = part[i]; + for (int s = 1; s < ks; s++) acc += part[s * total + i]; + out[i] = acc; +} + +// One extern "C" kernel per instantiation: the fatbin carries __global__ functions, the host resolves them by name. +#define TF_GLM_BMM(BM, WN, F32) \ + extern "C" __global__ void __launch_bounds__((WN) *32) \ + tf_glm_bmm_b16_##BM##_##F32(tf_glm_bmm::BmmArgs a) { \ + tf_glm_bmm::bmm_body(a); \ + } \ + extern "C" __global__ void __launch_bounds__((WN) *32) \ + tf_glm_bmm_p16_##BM##_##F32(tf_glm_bmm::BmmArgs a) { \ + tf_glm_bmm::bmm_body(a); \ + } +TF_GLM_BMM(16, 4, 0) +TF_GLM_BMM(16, 4, 1) +TF_GLM_BMM(32, 4, 0) +TF_GLM_BMM(32, 4, 1) +TF_GLM_BMM(64, 4, 0) +TF_GLM_BMM(64, 4, 1) +TF_GLM_BMM(128, 8, 0) +TF_GLM_BMM(128, 8, 1) + +#define TF_GLM_ROUTER(BM) \ + extern "C" __global__ void __launch_bounds__(128) tf_glm_router_##BM(tf_glm_bmm::RouterArgs a) { \ + tf_glm_bmm::router_body(a); \ + } \ + extern "C" __global__ void __launch_bounds__(128) tf_glm_router_part_##BM(tf_glm_bmm::RouterArgs a) { \ + tf_glm_bmm::router_part_body(a); \ + } \ + extern "C" __global__ void __launch_bounds__(128) tf_glm_prouter_##BM(tf_glm_bmm::RouterArgs a) { \ + tf_glm_bmm::router_body(a); \ + } \ + extern "C" __global__ void __launch_bounds__(128) tf_glm_prouter_part_##BM(tf_glm_bmm::RouterArgs a) { \ + tf_glm_bmm::router_part_body(a); \ + } +TF_GLM_ROUTER(16) +TF_GLM_ROUTER(32) +TF_GLM_ROUTER(64) +TF_GLM_ROUTER(128) + +extern "C" __global__ void tf_glm_bmm_reduce(const float* part, void* out, int total, int sk, int f32) { + tf_glm_bmm::reduce_body(part, out, total, sk, f32); +} + +extern "C" __global__ void tf_glm_router_sum(const float* part, float* out, int total, int ks) { + tf_glm_bmm::router_sum_body(part, out, total, ks); +} + +} // namespace tf_glm_bmm diff --git a/zig/kernels/cuda/glm_glue.cu b/zig/kernels/cuda/glm_glue.cu new file mode 100644 index 0000000000..fb634faefa --- /dev/null +++ b/zig/kernels/cuda/glm_glue.cu @@ -0,0 +1,560 @@ +// GLM glue kernels: bit-exact ports of the Triton kernels in families/glm5_next/cuda/glue.py (PTX-ordered). + +#include +#include + +typedef unsigned short u16; +typedef unsigned int u32; +typedef unsigned long long u64; +typedef long long i64; + +__device__ __forceinline__ float f32of(u16 v) { return __uint_as_float((u32)v << 16); } + +__device__ __forceinline__ u16 bf16_rn(float v) { u16 r; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(r) : "f"(v)); return r; } + +__device__ __forceinline__ float divf(float a, float b) { + float r; + asm("div.full.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float rcpf(float x) { + float r; + asm("rcp.approx.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +__device__ __forceinline__ float sqrtf_ftz(float x) { + float r; + asm("sqrt.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +__device__ __forceinline__ float ex2f(float x) { + float r; + asm("ex2.approx.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +#define NEG_LOG2E __int_as_float((int)0xBFB8AA3Bu) +#define POS_LOG2E __int_as_float(0x3FB8AA3Bu) +#define TINY1E20 __int_as_float(0x1E3CE508u) + +__device__ __forceinline__ float sigmoid(float x) { return divf(1.0f, 1.0f + ex2f(x * NEG_LOG2E)); } + +__device__ __forceinline__ void load2bf(const u16* p, float* v) { + u32 w = *(const u32*)p; + v[0] = f32of((u16)w); + v[1] = f32of((u16)(w >> 16)); +} + +__device__ __forceinline__ void load4bf(const u16* p, float* v) { + u32 w0 = *(const u32*)p; + u32 w1 = *(const u32*)(p + 2); + v[0] = f32of((u16)w0); + v[1] = f32of((u16)(w0 >> 16)); + v[2] = f32of((u16)w1); + v[3] = f32of((u16)(w1 >> 16)); +} + +__device__ __forceinline__ void store4bf(u16* p, const u16* v) { + *(u32*)p = (u32)v[0] | ((u32)v[1] << 16); + *(u32*)(p + 2) = (u32)v[2] | ((u32)v[3] << 16); +} + +// -- embedding (glue.py _embed_b16) ------------------------------------------------------------ +__device__ __forceinline__ void embed_b16(const u16* ids, const u16* w, u16* out, int copies) { + const int d = gridDim.y * 64; + int r = blockIdx.x, g = blockIdx.y, t = threadIdx.x; + u32 tok = ((const u32*)ids)[r]; + u32 v = *(const u32*)(w + (u64)tok * d + g * 64 + 2 * t); + u16* dst = out + (u64)r * copies * d + g * 64; + for (int c = 0; c < copies; ++c) ((u32*)(dst + (u64)c * d))[t] = v; +} + +// -- RMSNorm (glue.py _rmsnorm, D power of two, 256 threads) ------------------------------------ +template +__device__ __forceinline__ void rmsnorm(const u16* x, i64 xs, const u16* w, u16* out, i64 os, float* sums, float eps) { + __shared__ float warp_sums[8]; + __shared__ float gs[64]; + int r = blockIdx.x, t = threadIdx.x, lane = t & 31, warp = t >> 5; + const u16* xr = x + (i64)r * xs; + float a[16]; + load4bf(xr + 8 * t, a); + load4bf(xr + 8 * t + 4, a + 4); + load4bf(xr + 2048 + 8 * t, a + 8); + load4bf(xr + 2048 + 8 * t + 4, a + 12); + float p02 = fmaf(a[0], a[0], a[2] * a[2]), p13 = fmaf(a[1], a[1], a[3] * a[3]); + float p46 = fmaf(a[4], a[4], a[6] * a[6]), p57 = fmaf(a[5], a[5], a[7] * a[7]); + float p8a = fmaf(a[8], a[8], a[10] * a[10]), p9b = fmaf(a[9], a[9], a[11] * a[11]); + float pca = fmaf(a[12], a[12], a[14] * a[14]), pdb = fmaf(a[13], a[13], a[15] * a[15]); + float ssum = (((p02 + p46) + (p8a + pca)) + ((p13 + p57) + (p9b + pdb))); + for (int i = 16; i; i >>= 1) ssum += __shfl_xor_sync(0xffffffffu, ssum, i); + if (lane == 0) warp_sums[warp] = ssum; + __syncthreads(); + float v0 = warp_sums[0], v1 = warp_sums[1], v2 = warp_sums[2], v3 = warp_sums[3]; + float v4 = warp_sums[4], v5 = warp_sums[5], v6 = warp_sums[6], v7 = warp_sums[7]; + float total = ((v0 + v4) + (v2 + v6)) + ((v1 + v5) + (v3 + v7)); + float rinv = divf(1.0f, sqrtf_ftz(divf(total, (float)D) + eps)); + float ya[8], yb[8]; + u16 yva[8], yvb[8]; + const u16* wr = w + 8 * t; + for (int i = 0; i < 8; ++i) { + float wa = f32of(wr[i]), wb = f32of(wr[i + 2048]); + yva[i] = bf16_rn(wa * f32of(bf16_rn(rinv * a[i]))); + yvb[i] = bf16_rn(wb * f32of(bf16_rn(rinv * a[i + 8]))); + ya[i] = f32of(yva[i]); + yb[i] = f32of(yvb[i]); + } + store4bf(out + (i64)r * os + 8 * t, yva); + store4bf(out + (i64)r * os + 8 * t + 4, yva + 4); + store4bf(out + (i64)r * os + 2048 + 8 * t, yvb); + store4bf(out + (i64)r * os + 2048 + 8 * t + 4, yvb + 4); + if (sums) { + float pa = ((ya[0] + ya[2]) + (ya[4] + ya[6])) + ((ya[1] + ya[3]) + (ya[5] + ya[7])); + float pb = ((yb[0] + yb[2]) + (yb[4] + yb[6])) + ((yb[1] + yb[3]) + (yb[5] + yb[7])); + for (int i = 4; i; i >>= 1) pa += __shfl_xor_sync(0xffffffffu, pa, i); + for (int i = 4; i; i >>= 1) pb += __shfl_xor_sync(0xffffffffu, pb, i); + if ((t & 7) == 0) { + gs[t >> 3] = pa; + gs[32 + (t >> 3)] = pb; + } + __syncthreads(); + if (t < 64) sums[r * (D / 64) + t] = gs[t]; + } +} + +// -- hyper-connection partial dots (glue.py _hc_partial; WIDE x NB blocks of 1024) --------------- +template // block (row, b, j): warp w's dot 4 j + w in t order (split by j, same bits); j 0: squares +__device__ __forceinline__ void hc_partial(const u16* x, const float* fn, float* part) { + __shared__ float ssw[4]; + int r = blockIdx.x, b = blockIdx.y, j = blockIdx.z, w = threadIdx.x >> 5, l = threadIdx.x & 31; + const int m = 4 * j + w; + const u16* xr = x + (i64)r * WIDE; + float acc = 0.0f, ss = 0.0f; + for (int t = 0; t < 8; ++t) { + int base = b * 1024 + t * 128; + float xq[4]; + load2bf(xr + base + 4 * l, xq); + load2bf(xr + base + 4 * l + 2, xq + 2); + if (j == 0) { + float xo = f32of(xr[base + threadIdx.x]); + ss += xo * xo; + } + float wv[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + if (m < 24) { + const float* wp = fn + (i64)m * WIDE + base + 4 * l; + wv[0] = wp[0]; + wv[1] = wp[1]; + wv[2] = wp[2]; + wv[3] = wp[3]; + } + float p0 = xq[0] * wv[0], p1 = xq[1] * wv[1]; + float p2 = xq[2] * wv[2], p3 = xq[3] * wv[3]; + float q = (p0 + p2) + (p1 + p3); + for (int i = 16; i; i >>= 1) q += __shfl_xor_sync(0xffffffffu, q, i); + acc += q; + } + float* out = part + ((i64)r * NB + b) * 32; + if (l == 0 && m < 24) out[m] = acc; + if (j == 0) { + for (int i = 16; i; i >>= 1) ss += __shfl_xor_sync(0xffffffffu, ss, i); + if (l == 0) ssw[w] = ss; + __syncthreads(); + if (threadIdx.x == 0) { + out[24] = (ssw[0] + ssw[2]) + (ssw[1] + ssw[3]); + for (int k = 25; k < 32; ++k) out[k] = 0.0f; + } + } +} + +// -- hyper-connection finish (glue.py _hc_finish; S = 4 streams) --------------------------------- +template +__device__ __forceinline__ void hc_finish(const u16* x, const float* part, const float* base, const float* scale, const u16* nw, + u16* out, float* xs, float* post, float* comb, float eps_norm, float hc_eps, int iters) { + __shared__ float warp_sums[8]; + __shared__ float gs[64]; + int r = blockIdx.x, t = threadIdx.x, m = t & 31; + float mix = 0.0f, ss = 0.0f; + for (int b = 0; b < 16; ++b) mix += part[((i64)r * 16 + b) * 32 + m]; + for (int b = 0; b < 16; ++b) ss += part[((i64)r * 16 + b) * 32 + 24]; + float rinv = divf(1.0f, sqrtf_ftz(divf(ss, (float)(4 * D)) + eps_norm)); + mix *= rinv; + float base_m = (m < 24) ? base[m] : 0.0f; + float lpre = fmaf(mix, scale[0], base_m); + float lpost = fmaf(mix, scale[1], base_m); + float lcomb = fmaf(mix, scale[2], base_m); + float pre = 0.0f, combv = 0.0f; + if (m < 4) pre = sigmoid(lpre) + hc_eps; + else if (m < 8) pre = sigmoid(lpost) + sigmoid(lpost); + if (m >= 8 && m < 24) { + const unsigned SM = 0x00ffff00u; + int rr = (m - 8) >> 2, cc = (m - 8) & 3; + float cmax = lcomb; + cmax = fmaxf(cmax, __shfl_xor_sync(SM, cmax, 2)); + cmax = fmaxf(cmax, __shfl_xor_sync(SM, cmax, 1)); + float ce = ex2f((lcomb - cmax) * POS_LOG2E); + // Row/col sums are (a+c)+(b+d) butterflies; divisions lower to fmaf(rcp(rs),ce,eps) then rcp(denom)*ce. + float rs = (ce + __shfl_xor_sync(SM, ce, 2)) + (__shfl_xor_sync(SM, ce, 1) + __shfl_xor_sync(SM, ce, 3)); + ce = fmaf(rcpf(rs), ce, hc_eps); + float b1 = __shfl_sync(SM, ce, 8 + ((rr ^ 1) << 2) + cc); + float b2 = __shfl_sync(SM, ce, 8 + ((rr ^ 2) << 2) + cc); + float b3 = __shfl_sync(SM, ce, 8 + ((rr ^ 3) << 2) + cc); + ce = rcpf(((ce + b2) + (b1 + b3)) + hc_eps) * ce; + for (int it = 1; it < iters; ++it) { + rs = (ce + __shfl_xor_sync(SM, ce, 2)) + (__shfl_xor_sync(SM, ce, 1) + __shfl_xor_sync(SM, ce, 3)); + ce = rcpf(rs + hc_eps) * ce; + b1 = __shfl_sync(SM, ce, 8 + ((rr ^ 1) << 2) + cc); + b2 = __shfl_sync(SM, ce, 8 + ((rr ^ 2) << 2) + cc); + b3 = __shfl_sync(SM, ce, 8 + ((rr ^ 3) << 2) + cc); + ce = rcpf(((ce + b2) + (b1 + b3)) + hc_eps) * ce; + } + combv = ce; + } + if (t < 128) { + if (m >= 4 && m < 8) post[(i64)r * 4 + (m - 4)] = pre; + if (m >= 8 && m < 24) comb[(i64)r * 16 + (m - 8)] = combv; + } + float p[4]; + for (int sv = 0; sv < 4; ++sv) { + float v = (m == sv) ? pre : 0.0f; + for (int i = 16; i; i >>= 1) v += __shfl_xor_sync(0xffffffffu, v, i); + p[sv] = v; + } + const u16* xr = x + (i64)r * (4 * D); + float c[16]; + for (int half = 0; half < 2; ++half) { + int off = half * 2048 + 8 * t; + float x0[8], x1[8], x2[8], x3[8]; + load4bf(xr + off, x0); + load4bf(xr + off + 4, x0 + 4); + load4bf(xr + D + off, x1); + load4bf(xr + D + off + 4, x1 + 4); + load4bf(xr + 2 * D + off, x2); + load4bf(xr + 2 * D + off + 4, x2 + 4); + load4bf(xr + 3 * D + off, x3); + load4bf(xr + 3 * D + off + 4, x3 + 4); + for (int i = 0; i < 8; ++i) { + // Triton's f32x2 pairs: even columns fma(p1, x1, p0 * x0), odd ones fma(p0, x0, p1 * x1) (probed) + float v = (i & 1) ? fmaf(p[0], x0[i], p[1] * x1[i]) : fmaf(p[1], x1[i], p[0] * x0[i]); + v = fmaf(p[2], x2[i], v); + v = fmaf(p[3], x3[i], v); + c[half * 8 + i] = f32of(bf16_rn(v)); + } + } + float q1o = fmaf(c[1], c[1], c[3] * c[3]), q1e = fmaf(c[0], c[0], c[2] * c[2]); + float q2o = fmaf(c[5], c[5], c[7] * c[7]), q2e = fmaf(c[4], c[4], c[6] * c[6]); + float q3o = fmaf(c[9], c[9], c[11] * c[11]), q3e = fmaf(c[8], c[8], c[10] * c[10]); + float q4o = fmaf(c[13], c[13], c[15] * c[15]), q4e = fmaf(c[12], c[12], c[14] * c[14]); + float ssum = ((q1o + q2o) + (q3o + q4o)) + ((q1e + q2e) + (q3e + q4e)); + for (int i = 16; i; i >>= 1) ssum += __shfl_xor_sync(0xffffffffu, ssum, i); + if ((t & 31) == 0) warp_sums[t >> 5] = ssum; + __syncthreads(); + float v0 = warp_sums[0], v1 = warp_sums[1], v2 = warp_sums[2], v3 = warp_sums[3]; + float v4 = warp_sums[4], v5 = warp_sums[5], v6 = warp_sums[6], v7 = warp_sums[7]; + float total = ((v0 + v4) + (v2 + v6)) + ((v1 + v5) + (v3 + v7)); + float rinv2 = divf(1.0f, sqrtf_ftz(divf(total, (float)D) + eps_norm)); + u16 yva[8], yvb[8]; + const u16* wr = nw + 8 * t; + for (int i = 0; i < 8; ++i) { + yva[i] = bf16_rn(f32of(wr[i]) * f32of(bf16_rn(rinv2 * c[i]))); + yvb[i] = bf16_rn(f32of(wr[i + 2048]) * f32of(bf16_rn(rinv2 * c[i + 8]))); + } + store4bf(out + (i64)r * D + 8 * t, yva); + store4bf(out + (i64)r * D + 8 * t + 4, yva + 4); + store4bf(out + (i64)r * D + 2048 + 8 * t, yvb); + store4bf(out + (i64)r * D + 2048 + 8 * t + 4, yvb + 4); + float ya[8], yb[8]; + for (int i = 0; i < 8; ++i) { + ya[i] = f32of(yva[i]); + yb[i] = f32of(yvb[i]); + } + float pa = ((ya[0] + ya[2]) + (ya[4] + ya[6])) + ((ya[1] + ya[3]) + (ya[5] + ya[7])); + float pb = ((yb[0] + yb[2]) + (yb[4] + yb[6])) + ((yb[1] + yb[3]) + (yb[5] + yb[7])); + for (int i = 4; i; i >>= 1) pa += __shfl_xor_sync(0xffffffffu, pa, i); + for (int i = 4; i; i >>= 1) pb += __shfl_xor_sync(0xffffffffu, pb, i); + if ((t & 7) == 0) { + gs[t >> 3] = pa; + gs[32 + (t >> 3)] = pb; + } + __syncthreads(); + if (t < 64) xs[(i64)r * (D / 64) + t] = gs[t]; +} + +// -- hyper-connection post (glue.py _hc_post; world fp32 partials at element stride rs) ----------- +__device__ __forceinline__ void hc_post(const u16* x, u16* xout, const float* g, const float* post, const float* comb, i64 rs, + int world, int s) { + const int d = gridDim.y * 1024; + int r = blockIdx.x, cb = blockIdx.y, t = threadIdx.x; + int d0 = cb * 1024 + 8 * t; + const float* gr = g + (i64)r * d + d0; + float branch[8]; + for (int i = 0; i < 8; ++i) branch[i] = gr[i]; + for (int k = 1; k < world; ++k) { + const float* gk = g + (i64)k * rs + (i64)r * d + d0; + for (int i = 0; i < 8; ++i) branch[i] += gk[i]; + } + const u16* xr = x + (i64)r * (s * d) + d0; + float x0[8], x1[8], x2[8], x3[8]; + load4bf(xr, x0); + load4bf(xr + 4, x0 + 4); + load4bf(xr + d, x1); + load4bf(xr + d + 4, x1 + 4); + load4bf(xr + 2 * d, x2); + load4bf(xr + 2 * d + 4, x2 + 4); + load4bf(xr + 3 * d, x3); + load4bf(xr + 3 * d + 4, x3 + 4); + for (int sv = 0; sv < s; ++sv) { + float c0v = comb[(i64)r * 16 + 0 * 4 + sv]; + float c1v = comb[(i64)r * 16 + 1 * 4 + sv]; + float c2v = comb[(i64)r * 16 + 2 * 4 + sv]; + float c3v = comb[(i64)r * 16 + 3 * 4 + sv]; + float psv = post[(i64)r * s + sv]; + u16 res[8]; + for (int i = 0; i < 8; ++i) { + // Triton fuses the first product into the rounded second: fma(c0, x0, c1 * x1) (glm.layer.kda, 14 words) + float mixed = fmaf(c0v, x0[i], c1v * x1[i]); + mixed = fmaf(c2v, x2[i], mixed); + mixed = fmaf(c3v, x3[i], mixed); + res[i] = bf16_rn(fmaf(psv, f32of(bf16_rn(branch[i])), mixed)); + } + store4bf(xout + (i64)r * (s * d) + (i64)sv * d + d0, res); + store4bf(xout + (i64)r * (s * d) + (i64)sv * d + d0 + 4, res + 4); + } +} + +// -- residual add (glue.py _residual_add) -------------------------------------------------------- +__device__ __forceinline__ void residual_add(const u16* x, u16* xout, const float* g, i64 rs, int world) { + const int d = gridDim.y * 1024; + int r = blockIdx.x, cb = blockIdx.y, t = threadIdx.x; + int d0 = cb * 1024 + 8 * t; + const float* gr = g + (i64)r * d + d0; + float branch[8]; + for (int i = 0; i < 8; ++i) branch[i] = gr[i]; + for (int k = 1; k < world; ++k) { + const float* gk = g + (i64)k * rs + (i64)r * d + d0; + for (int i = 0; i < 8; ++i) branch[i] += gk[i]; + } + const u16* xr = x + (i64)r * d + d0; + u16* outr = xout + (i64)r * d + d0; + for (int i = 0; i < 8; i += 4) { + float xv[4]; + load4bf(xr + i, xv); + u16 res[4]; + for (int j = 0; j < 4; ++j) res[j] = bf16_rn(xv[j] + f32of(bf16_rn(branch[i + j]))); + store4bf(outr + i, res); + } +} + +// -- stream mean (glue.py _stream_mean) ---------------------------------------------------------- +__device__ __forceinline__ void stream_mean(const u16* x, u16* out, int s) { + const int d = gridDim.y * 1024; + int r = blockIdx.x, cb = blockIdx.y, t = threadIdx.x; + int d0 = cb * 1024 + 8 * t; + const u16* xr = x + (i64)r * (s * d) + d0; + float acc[8]; + load4bf(xr, acc); + load4bf(xr + 4, acc + 4); + for (int k = 1; k < s; ++k) { + float v[8]; + load4bf(xr + (i64)k * d, v); + load4bf(xr + (i64)k * d + 4, v + 4); + for (int i = 0; i < 8; ++i) acc[i] += v[i]; + } + u16 res[8]; + for (int i = 0; i < 8; ++i) res[i] = bf16_rn(divf(acc[i], (float)s)); + store4bf(out + (i64)r * d + d0, res); + store4bf(out + (i64)r * d + d0 + 4, res + 4); +} + +// -- dense MLP swiglu (glue.py _swiglu; BLOCK 512, 64-group sums) --------------------------------- +__device__ __forceinline__ void swiglu(const u16* gu, u16* out, float* xs, float lim) { + __shared__ float sm[8]; + int r = blockIdx.x, cb = blockIdx.y, t = threadIdx.x; + int w = gridDim.y * 512; + const u16* gur = gu + (i64)r * (2 * w) + cb * 512; + float g[4], u[4]; + load4bf(gur + 4 * t, g); + load4bf(gur + w + 4 * t, u); + for (int i = 0; i < 4; ++i) { + g[i] = fminf(g[i], lim); + u[i] = fminf(fmaxf(u[i], -lim), lim); + float d = ex2f(g[i] * NEG_LOG2E); + float si = f32of(bf16_rn(divf(g[i], d + 1.0f))); + g[i] = si * u[i]; + } + u16 a[4]; + for (int i = 0; i < 4; ++i) a[i] = bf16_rn(g[i]); + store4bf(out + (i64)r * w + cb * 512 + 4 * t, a); + float p = (f32of(a[1]) + f32of(a[3])) + (f32of(a[0]) + f32of(a[2])); + for (int i = 8; i; i >>= 1) p += __shfl_xor_sync(0xffffffffu, p, i); + if ((t & 15) == 0) sm[t >> 4] = p; + __syncthreads(); + if (t < 8) xs[(i64)r * (w / 64) + cb * 8 + t] = sm[t]; +} + +// -- MoE router top-k (glue.py _topk; BLOCK = next_pow2(ne+1) <= 512) ----------------------------- +__device__ __forceinline__ void topk(const float* lg, const float* bias, int* pick, float* wts, float scale, int ne, int topk) { + __shared__ float mxs[4]; + __shared__ float ivs[4]; // its own slots: a warp's index must not land in mxs while another still reads the max + __shared__ float sks[4]; + int r = blockIdx.x, t = threadIdx.x, w = t >> 5, lane = t & 31; + const int per = 4; // BLOCK <= 512, 128 threads + float choice[per], score[per]; + int eid[per]; + for (int i = 0; i < per; ++i) { + int e = t * per + i; + eid[i] = e; + bool ok = e < ne; + float s = ok ? sigmoid(lg[(i64)r * ne + e]) : 0.0f; + score[i] = s; + choice[i] = ok ? s + bias[e] : -__int_as_float(0x7f800000u); + } + float total = 0.0f, wk[16]; // every thread holds the picks' scores: no global read-back while others normalize + for (int k = 0; k < topk; ++k) { + float mx = choice[0]; + for (int i = 1; i < per; ++i) mx = fmaxf(mx, choice[i]); + for (int i = 16; i; i >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, i)); + if (lane == 0) mxs[w] = mx; + __syncthreads(); + mx = fmaxf(fmaxf(mxs[0], mxs[2]), fmaxf(mxs[1], mxs[3])); + float iv = (float)(per * 128); + for (int i = 0; i < per; ++i) + if (choice[i] == mx) iv = fminf(iv, (float)eid[i]); + for (int i = 16; i; i >>= 1) iv = fminf(iv, __shfl_xor_sync(0xffffffffu, iv, i)); + if (lane == 0) ivs[w] = iv; + __syncthreads(); + int idx = (int)fminf(fminf(ivs[0], ivs[2]), fminf(ivs[1], ivs[3])); + float sk = 0.0f; + for (int i = 0; i < per; ++i) + if (eid[i] == idx) sk += score[i]; + for (int i = 16; i; i >>= 1) sk += __shfl_xor_sync(0xffffffffu, sk, i); + if (lane == 0) sks[w] = sk; + __syncthreads(); + sk = (sks[0] + sks[2]) + (sks[1] + sks[3]); + pick[(i64)r * (topk + 1) + k] = idx; + wk[k] = sk; + total += sk; + for (int i = 0; i < per; ++i) + if (eid[i] == idx) choice[i] = -__int_as_float(0x7f800000u); + __syncthreads(); + } + for (int k = 0; k < topk; ++k) + if (t == k) wts[(i64)r * (topk + 1) + k] = scale * divf(wk[k], total + TINY1E20); + if (t < topk + 1) { + pick[(i64)r * (topk + 1) + topk] = ne; + wts[(i64)r * (topk + 1) + topk] = 1.0f; + } +} + +// -- MoE combine (glue.py _combine; per-thread 8 consecutive d: pos 0,7 fma order, pos 1..6 fma slot 0, S-1, 1..S-2) +__device__ __forceinline__ void combine(const u16* y, const float* wts, float* out, int slots) { + const int d = gridDim.y * 1024; + int r = blockIdx.x, cb = blockIdx.y, t = threadIdx.x; + int d0 = cb * 1024 + 8 * t; + const u16* yrow = y + (i64)r * slots * d + d0; + const float* wrow = wts + (i64)r * slots; + float* orow = out + (i64)r * d + d0; + for (int i = 0; i < 8; ++i) { + bool edge = (i == 0 || i == 7); + float acc = fmaf(f32of(yrow[i]), wrow[0], edge ? -0.0f : 0.0f); + if (edge) { + for (int k = 1; k < slots; ++k) acc = fmaf(f32of(yrow[(i64)k * d + i]), wrow[k], acc); + } else { + acc = fmaf(f32of(yrow[(i64)(slots - 1) * d + i]), wrow[slots - 1], acc); + for (int k = 1; k < slots - 1; ++k) acc = fmaf(f32of(yrow[(i64)k * d + i]), wrow[k], acc); + } + orow[i] = acc; + } +} + +// -- greedy sampling over a rank's logits (decode.py sample_rows greedy path) ---------------------- +__device__ __forceinline__ void greedy(const u16* lg, int* tok, float* packed, int n, int offset) { + __shared__ float sv[8]; + __shared__ int si[8]; + int r = blockIdx.x, t = threadIdx.x; + const u16* row = lg + (i64)r * n; + float best = -__int_as_float(0x7f800000u); + int bi = 0; + for (int i = t; i < n; i += 256) { + float v = f32of(row[i]); + if (v > best || (v == best && i < bi)) { + best = v; + bi = i; + } + } + for (int off = 16; off; off >>= 1) { + float ov = __shfl_xor_sync(0xffffffffu, best, off); + int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (ov > best || (ov == best && oi < bi)) { + best = ov; + bi = oi; + } + } + if ((t & 31) == 0) { + sv[t >> 5] = best; + si[t >> 5] = bi; + } + __syncthreads(); + if (t == 0) { + float bv = sv[0]; + int bidx = si[0]; + for (int i = 1; i < 8; ++i) + if (sv[i] > bv || (sv[i] == bv && si[i] < bidx)) { + bv = sv[i]; + bidx = si[i]; + } + if (tok) tok[r] = bidx; + if (packed) { // sample_rows' all-gather payload: the value, then the vocab id's bits + packed[2 * r] = bv; + packed[2 * r + 1] = __int_as_float(bidx + offset); + } + } +} + +// -- entry kernels (zig launches these; grids in zig/src/cuda/glm_glue.zig) ---------------------- +extern "C" __global__ void tf_glm_embed_b16(const u16* ids, const u16* w, u16* out, int copies) { + embed_b16(ids, w, out, copies); +} + +extern "C" __global__ void tf_glm_rmsnorm(const u16* x, i64 xs, const u16* w, u16* out, i64 os, float* sums, + float eps) { + rmsnorm<4096>(x, xs, w, out, os, sums, eps); +} + +extern "C" __global__ void tf_glm_hc_partial(const u16* x, const float* fn, float* part) { + hc_partial<16384, 16>(x, fn, part); +} + +extern "C" __global__ void tf_glm_hc_finish(const u16* x, const float* part, const float* base, const float* scale, + const u16* nw, u16* out, float* xs, float* post, float* comb, + float eps_norm, float hc_eps, int iters) { + hc_finish<4096>(x, part, base, scale, nw, out, xs, post, comb, eps_norm, hc_eps, iters); +} + +extern "C" __global__ void tf_glm_hc_post(const u16* x, u16* xout, const float* g, const float* post, + const float* comb, i64 rs, int world, int s) { + hc_post(x, xout, g, post, comb, rs, world, s); +} + +extern "C" __global__ void tf_glm_residual_add(const u16* x, u16* xout, const float* g, i64 rs, int world) { + residual_add(x, xout, g, rs, world); +} + +extern "C" __global__ void tf_glm_stream_mean(const u16* x, u16* out, int s) { stream_mean(x, out, s); } + +extern "C" __global__ void tf_glm_swiglu(const u16* gu, u16* out, float* xs, float lim) { swiglu(gu, out, xs, lim); } + +extern "C" __global__ void tf_glm_topk(const float* lg, const float* bias, int* pick, float* wts, float scale, int ne, + int tk) { + topk(lg, bias, pick, wts, scale, ne, tk); +} + +extern "C" __global__ void tf_glm_combine(const u16* y, const float* wts, float* out, int slots) { + combine(y, wts, out, slots); +} + +extern "C" __global__ void tf_glm_greedy(const u16* lg, int* tok, int n) { greedy(lg, tok, nullptr, n, 0); } +extern "C" __global__ void tf_glm_greedy_pack(const u16* lg, float* packed, int n, int offset) { + greedy(lg, nullptr, packed, n, offset); +} diff --git a/zig/kernels/cuda/glm_glue2.cu b/zig/kernels/cuda/glm_glue2.cu new file mode 100644 index 0000000000..760c3c7602 --- /dev/null +++ b/zig/kernels/cuda/glm_glue2.cu @@ -0,0 +1,188 @@ +// glm5_next/cuda/glue.py's Triton kernels glm_glue.cu has no room for: the 1536/512 rmsnorms, the fp32-y combine +#include + +typedef unsigned short u16; +typedef unsigned int u32; +typedef long long i64; + +__device__ __forceinline__ float f32of(u16 v) { return __uint_as_float((u32)v << 16); } + +__device__ __forceinline__ u16 bf16_rn(float v) { u16 r; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(r) : "f"(v)); return r; } + +__device__ __forceinline__ float divf(float a, float b) { + float r; + asm("div.full.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float sqrtf_ftz(float x) { + float r; + asm("sqrt.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +__device__ __forceinline__ void load4bf(const u16* p, float* v) { + u32 w0 = *(const u32*)p, w1 = *(const u32*)(p + 2); + v[0] = f32of((u16)w0); + v[1] = f32of((u16)(w0 >> 16)); + v[2] = f32of((u16)w1); + v[3] = f32of((u16)(w1 >> 16)); +} + +__device__ __forceinline__ void store4bf(u16* p, const u16* v) { + *(u32*)p = (u32)v[0] | ((u32)v[1] << 16); + *(u32*)(p + 2) = (u32)v[2] | ((u32)v[3] << 16); +} + +// -- RMSNorm, D = 512 (glue.py _rmsnorm, BLOCK 512, 4 warps): pair sums, 16-lane xor 8/4/2/1, (w0+w2)+(w1+w3) +__device__ __forceinline__ void rms512(const u16* x, i64 xs, const u16* w, u16* out, i64 os, float* sums, float eps) { + __shared__ float warp_sums[4]; + __shared__ float gs[8]; + int r = blockIdx.x, t = threadIdx.x; + const u16* xr = x + (i64)r * xs; + float a[4]; + load4bf(xr + 4 * t, a); + float p02 = fmaf(a[0], a[0], a[2] * a[2]), p13 = fmaf(a[1], a[1], a[3] * a[3]); + float ssum = p02 + p13; + for (int i = 16; i; i >>= 1) ssum += __shfl_xor_sync(0xffffffffu, ssum, i); + if ((t & 31) == 0) warp_sums[t >> 5] = ssum; + __syncthreads(); + float total = (warp_sums[0] + warp_sums[2]) + (warp_sums[1] + warp_sums[3]); + float rinv = divf(1.0f, sqrtf_ftz(divf(total, 512.0f) + eps)); + u16 yv[4]; + float ya[4]; + const u16* wr = w + 4 * t; + for (int i = 0; i < 4; ++i) { + yv[i] = bf16_rn(f32of(wr[i]) * f32of(bf16_rn(rinv * a[i]))); + ya[i] = f32of(yv[i]); + } + store4bf(out + (i64)r * os + 4 * t, yv); + if (sums) { + float pa = (ya[0] + ya[2]) + (ya[1] + ya[3]); + for (int i = 8; i; i >>= 1) pa += __shfl_xor_sync(0xffffffffu, pa, i); + gs[t >> 4] = pa; + __syncthreads(); + if (t < 8) sums[(i64)r * 8 + t] = gs[t]; + } +} + +// -- RMSNorm, D = 1536 (BLOCK 2048 masked, 4 warps): 8 elements a half, second half valid while t < 64 +__device__ __forceinline__ void rms1536(const u16* x, i64 xs, const u16* w, u16* out, i64 os, float* sums, float eps) { + __shared__ float warp_sums[4]; + __shared__ float gs[32]; + int r = blockIdx.x, t = threadIdx.x; + const u16* xr = x + (i64)r * xs; + float a[8], b[8]; + bool second = t < 64; + load4bf(xr + 8 * t, a); + load4bf(xr + 8 * t + 4, a + 4); + for (int i = 0; i < 8; ++i) b[i] = 0.0f; + if (second) { + load4bf(xr + 1024 + 8 * t, b); + load4bf(xr + 1024 + 8 * t + 4, b + 4); + } + float p02 = fmaf(a[0], a[0], a[2] * a[2]), p13 = fmaf(a[1], a[1], a[3] * a[3]); + float p46 = fmaf(a[4], a[4], a[6] * a[6]), p57 = fmaf(a[5], a[5], a[7] * a[7]); + float p8a = fmaf(b[0], b[0], b[2] * b[2]), p9b = fmaf(b[1], b[1], b[3] * b[3]); + float pca = fmaf(b[4], b[4], b[6] * b[6]), pdb = fmaf(b[5], b[5], b[7] * b[7]); + float ssum = (((p02 + p46) + (p8a + pca)) + ((p13 + p57) + (p9b + pdb))); + for (int i = 16; i; i >>= 1) ssum += __shfl_xor_sync(0xffffffffu, ssum, i); + if ((t & 31) == 0) warp_sums[t >> 5] = ssum; + __syncthreads(); + float total = (warp_sums[0] + warp_sums[2]) + (warp_sums[1] + warp_sums[3]); + // Triton folds `/ D + eps` into one FFMA with fp32 1/1536 (0x3a2aaaab): a power-of-two D divides exactly, 1536 not + float rinv = divf(1.0f, sqrtf_ftz(fmaf(total, __int_as_float(0x3a2aaaab), eps))); + u16 yva[8], yvb[8]; + float ya[8], yb[8]; + const u16* wr = w + 8 * t; + for (int i = 0; i < 8; ++i) { + yva[i] = bf16_rn(f32of(wr[i]) * f32of(bf16_rn(rinv * a[i]))); + ya[i] = f32of(yva[i]); + } + for (int i = 0; i < 8; ++i) { + yvb[i] = second ? bf16_rn(f32of(wr[1024 + i]) * f32of(bf16_rn(rinv * b[i]))) : bf16_rn(0.0f * rinv); + yb[i] = second ? f32of(yvb[i]) : 0.0f; + } + store4bf(out + (i64)r * os + 8 * t, yva); + store4bf(out + (i64)r * os + 8 * t + 4, yva + 4); + if (second) { + store4bf(out + (i64)r * os + 1024 + 8 * t, yvb); + store4bf(out + (i64)r * os + 1024 + 8 * t + 4, yvb + 4); + } + if (sums) { + float pa = ((ya[0] + ya[2]) + (ya[4] + ya[6])) + ((ya[1] + ya[3]) + (ya[5] + ya[7])); + float pb = ((yb[0] + yb[2]) + (yb[4] + yb[6])) + ((yb[1] + yb[3]) + (yb[5] + yb[7])); + for (int i = 4; i; i >>= 1) pa += __shfl_xor_sync(0xffffffffu, pa, i); + for (int i = 4; i; i >>= 1) pb += __shfl_xor_sync(0xffffffffu, pb, i); + gs[2 * (t >> 3)] = pa; + gs[2 * (t >> 3) + 1] = pb; + __syncthreads(); + if (t < 16) sums[(i64)r * 24 + t] = gs[2 * t]; + if (t >= 16 && t < 24) sums[(i64)r * 24 + t] = gs[2 * (t - 16) + 1]; + } +} + +// -- MoE combine, fp32 y: warps 0-1 sum their columns 0/4 in slot order 0..8, the rest 0,8,1..7; warps 2-3 shift to 3/7 +__device__ __forceinline__ void combine_f32(const float* y, const float* wts, float* out, int slots) { + const int d = gridDim.y * 1024; + int r = blockIdx.x, cb = blockIdx.y, t = threadIdx.x; + int d0 = cb * 1024 + 8 * t; + const float* yrow = y + (i64)r * slots * d + d0; + const float* wrow = wts + (i64)r * slots; + float* orow = out + (i64)r * d + d0; + if (slots == 9) { // GLM's 8 routed slots and the shared one: every load issued first, as 16-byte loads + float v[9][8], w[9]; +#pragma unroll + for (int k = 0; k < 9; ++k) { + const float4 lo = *reinterpret_cast(yrow + (i64)k * d); + const float4 hi = *reinterpret_cast(yrow + (i64)k * d + 4); + v[k][0] = lo.x; v[k][1] = lo.y; v[k][2] = lo.z; v[k][3] = lo.w; + v[k][4] = hi.x; v[k][5] = hi.y; v[k][6] = hi.z; v[k][7] = hi.w; + w[k] = wrow[k]; + } + float o[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + bool edge = (t >> 5 < 2) ? (i == 0 || i == 4) : (i == 3 || i == 7); + float acc = fmaf(v[0][i], w[0], 0.0f); + if (edge) { +#pragma unroll + for (int k = 1; k < 9; ++k) acc = fmaf(v[k][i], w[k], acc); + } else { + acc = fmaf(v[8][i], w[8], acc); +#pragma unroll + for (int k = 1; k < 8; ++k) acc = fmaf(v[k][i], w[k], acc); + } + o[i] = acc; + } + *reinterpret_cast(orow) = make_float4(o[0], o[1], o[2], o[3]); + *reinterpret_cast(orow + 4) = make_float4(o[4], o[5], o[6], o[7]); + return; + } + for (int i = 0; i < 8; ++i) { + bool edge = (t >> 5 < 2) ? (i == 0 || i == 4) : (i == 3 || i == 7); + float acc = fmaf(yrow[i], wrow[0], 0.0f); + if (edge) { + for (int k = 1; k < slots; ++k) acc = fmaf(yrow[(i64)k * d + i], wrow[k], acc); + } else { + acc = fmaf(yrow[(i64)(slots - 1) * d + i], wrow[slots - 1], acc); + for (int k = 1; k < slots - 1; ++k) acc = fmaf(yrow[(i64)k * d + i], wrow[k], acc); + } + orow[i] = acc; + } +} + +// -- entry kernels (zig launches these; grids in zig/src/cuda/glm_glue2.zig) ---------------------- +extern "C" __global__ void tf_glm_rmsnorm_512(const u16* x, i64 xs, const u16* w, u16* out, i64 os, float* sums, + float eps) { + rms512(x, xs, w, out, os, sums, eps); +} + +extern "C" __global__ void tf_glm_rmsnorm_1536(const u16* x, i64 xs, const u16* w, u16* out, i64 os, float* sums, + float eps) { + rms1536(x, xs, w, out, os, sums, eps); +} + +extern "C" __global__ void tf_glm_combine_f32(const float* y, const float* wts, float* out, int slots) { + combine_f32(y, wts, out, slots); +} diff --git a/zig/kernels/cuda/glm_kda.cu b/zig/kernels/cuda/glm_kda.cu new file mode 100644 index 0000000000..447f92983a --- /dev/null +++ b/zig/kernels/cuda/glm_kda.cu @@ -0,0 +1,394 @@ +// GLM-5.3-Flash's KDA (Kimi delta attention) on CUDA: one block of 1024 threads per head, a chain of R rows. + +// + +// Per row, following the Hugging Face definition: the depthwise conv over [conv state; q|k|v rows] (4 taps, + +// fp32) with SiLU and one bf16 rounding; fp32 L2 norms of q and k (eps inside the sum, q times DK^-0.5); the + +// per-channel decay g_i = exp(lower * sigmoid(exp(A_log) * (a_i + dt_bias_i))) in fp32; beta = bf16(sigmoid(b)); + +// the delta rule in fp32 (decay along the key channel, read with k, correct toward v, read out with q); + +// read-out rounded to bf16; then the gated RMSNorm bf16(w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(gate)). + +// Warp w owns value rows 4w .. 4w + 3, lane l the key columns 4l .. 4l + 3. The state update is one routine + +// (``update``) shared with ``replay``, compiled without FMA contraction, so a replayed prefix of a window gives + +// the bits of the serial steps. + +#include +#include +#include + + +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; + } +} + +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 __nv_bfloat16* __restrict__ cw, + const float* __restrict__ state_in, const float* __restrict__ a_log, const float* __restrict__ dt_bias, + const __nv_bfloat16* __restrict__ norm_w, float eps, 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 + __bfloat162float(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]; + } +} + +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) { + 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. +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) { + 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-independent work for all rows (prep_kernel, out_kernel); the delta rule row by row. + + +// Block (row, head), 512 threads: conv + SiLU of the 3 x 128 q|k|v channels, the decay, beta, then the q/k norms. +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 __nv_bfloat16* __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 + __bfloat162float(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; + } + } +} + +// Block (head, part): a warp owns 4 value rows; TR rows of saved k, q, decay, v, beta staged in shared memory. +template +__device__ __forceinline__ void step_impl( + 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}; + update(s, kk, gg, vv, 0, bt[rr]); +#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) y_out[((size_t)(r0 + rr) * H + h) * DV + warp * 4 + j] = __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]; + } +} +// The fatbin's step entry: the one instantiation the engine runs (4 warps, 16 staged rows). +extern "C" __global__ void __launch_bounds__(128) 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_impl<4, 16>(H, state_in, q_in, k_save, v_save, g_save, b_save, rows, y_out, state_out); +} + + + +// Block (row, head), 128 threads: the gated RMSNorm of the read-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) { + 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)); +} + +// The conv window shift: conv [L, 3, C] takes rows keep .. keep + 2 of [old window; proj rows] (proj [R, W >= C]). +extern "C" __global__ void __launch_bounds__(1024) tf_glm_kda_conv_shift( + __nv_bfloat16* conv, const __nv_bfloat16* __restrict__ proj, int keep, long conv_stride, + long proj_stride, long proj_row, int C) { + const int l = blockIdx.x; + const long cbase = (long)l * conv_stride; + const int c0 = blockIdx.y * 1024 + (int)threadIdx.x; + const bool ok = c0 < C; + __nv_bfloat16 w[3] = {}; + if (ok) { + for (int d = 0; d < 3; ++d) { + const int src = keep + d; + w[d] = src < 3 ? conv[cbase + (long)src * C + c0] + : proj[(long)l * proj_stride + (long)(src - 3) * proj_row + c0]; + } + } + __syncthreads(); + if (ok) { + for (int d = 0; d < 3; ++d) conv[cbase + (long)d * C + c0] = w[d]; + } +} + + diff --git a/zig/kernels/cuda/glm_kda_chunk.cu b/zig/kernels/cuda/glm_kda_chunk.cu new file mode 100644 index 0000000000..fa89f14bbe --- /dev/null +++ b/zig/kernels/cuda/glm_kda_chunk.cu @@ -0,0 +1,502 @@ +// GLM-5.3-Flash's KDA chunked prompt chunk kernel: the sub-chunk pipeline over the shared helpers' layout. + +#include "glm_kda_chunk_shared.cuh" + +__device__ long tf_glm_kda_chunk_smem_v = (long)sizeof(Smem); + +extern "C" __global__ void __launch_bounds__(NT, 1) tf_glm_kda_chunk( + 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 __nv_bfloat16* __restrict__ cw, + const float* __restrict__ state_in, const float* __restrict__ a_log, const float* __restrict__ dt_bias, + const __nv_bfloat16* __restrict__ norm_w, float eps, float lower, int rows, int off, + __nv_bfloat16* __restrict__ out, float* __restrict__ state_out) { + extern __shared__ __align__(16) unsigned char smem_raw[]; + Smem& sm = *reinterpret_cast(smem_raw); + const int h = blockIdx.x, tid = threadIdx.x; + const int vg = tid / KP, kp = tid % KP; + const int lane = tid & 31, warp = tid >> 5; + const int par = kp < 4 ? 1 : 0; // parity of the rows this thread finishes + const int ov = VPT * vg + (kp & 3); // their value column + const size_t sbase = (size_t)h * DV * DK; + const bool vec_p = (p_stride % 8 == 0) && ((uintptr_t)P % 16 == 0) && ((uintptr_t)cs % 16 == 0); + const bool vec_a = (a_stride % 8 == 0) && ((uintptr_t)A % 16 == 0); + + float s[VPT][NQ][4]; +#pragma unroll + for (int vv = 0; vv < VPT; ++vv) +#pragma unroll + for (int i = 0; i < NQ; ++i) { + const float4 x = + reinterpret_cast(state_in + sbase + (size_t)(VPT * vg + vv) * DK)[kp + KP * i]; + s[vv][i][0] = x.x; s[vv][i][1] = x.y; s[vv][i][2] = x.z; s[vv][i][3] = x.w; + } + const float rate = expf(a_log[h]); + const float dtb = dt_bias[h * DK + (tid & (DK - 1))]; // the decay channel this thread computes (tid % 128) + float nw[4]; // the norm weights of value columns 4 lane .. 4 lane + 3 +#pragma unroll + for (int x = 0; x < 4; ++x) nw[x] = ldb(norm_w + 4 * lane + x); + const bool vec_g = (g_stride % 4 == 0) && ((uintptr_t)G % 8 == 0); + const bool vec_o = ((H * DV) % 4 == 0) && ((uintptr_t)out % 8 == 0); + const int nsub = (rows + off + BT - 1) / BT; + for (int x = tid; x < PCH * TAPS; x += NT) { + const int ch = x / TAPS; + sm.cw[ch][x % TAPS] = cw[(size_t)((ch / DK) * H * DK + h * DK + ch % DK) * TAPS + x % TAPS]; + } + + for (int c = 0; c < nsub; ++c) { + const int r0 = c * BT - off; // the prompt chunk's row at the sub-chunk's row 0 + + // 0. stage this sub-chunk's rows (warmed in L2 during the previous one), warm the next one's and this + + // one's output gates + stage(sm, tid, h, H, r0, rows, P, p_stride, b_off, A, a_stride, cs, vec_p, vec_a); + asm volatile("cp.async.wait_all;" ::: "memory"); + __syncthreads(); + warm(tid, h, H, r0 + BT, c + 1 < nsub, rows, r0, P, p_stride, b_off, A, a_stride, G, g_stride); + + // 1. conv + SiLU (glm_kda.cu's arithmetic) in 4-row blocks: warp w takes rows 4 w .. 4 w + 3 of q, then of + + // k, then of v, lane l the channels 4 l .. 4 l + 3 (glm_kda.cu's norm lanes), and the L2 norms of its q and + + // k rows (glm_kda.cu's sums); then the decays and beta. +#pragma unroll 1 + for (int part = 0; part < 3; ++part) { + const int rb = 4 * warp, ch = part * DK + 4 * lane; + float wt[4][TAPS]; +#pragma unroll + for (int x = 0; x < 4; ++x) +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) wt[x][tap] = __bfloat162float(sm.cw[ch + x][tap]); + float xs[4 + HALO][4]; +#pragma unroll + for (int j = 0; j < 4 + HALO; ++j) { + const uint2 raw = *reinterpret_cast(&sm.u.raw_p[rb + j][ch]); + const __nv_bfloat162 lo = *reinterpret_cast(&raw.x); + const __nv_bfloat162 hi = *reinterpret_cast(&raw.y); + xs[j][0] = __low2float(lo); xs[j][1] = __high2float(lo); + xs[j][2] = __low2float(hi); xs[j][3] = __high2float(hi); + } + float acc[4][4], act[4][4]; + bool slow = false; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int x = 0; x < 4; ++x) { + acc[i][x] = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) acc[i][x] = acc[i][x] + wt[x][tap] * xs[i + tap][x]; + const float den = 1.0f + expf(-acc[i][x]); + act[i][x] = div_rn(acc[i][x], den); + slow |= !div_ok(acc[i][x], den); + } + if (slow) { // glm_kda.cu's IEEE division for the rare operands +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int x = 0; x < 4; ++x) act[i][x] = acc[i][x] / (1.0f + expf(-acc[i][x])); + } +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int r = r0 + rb + i; +#pragma unroll + for (int x = 0; x < 4; ++x) act[i][x] = (r >= 0 && r < rows) ? bf(act[i][x]) : 0.0f; + } + if (part < 2) { // warp-uniform + float ss[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + ss[i] = 0.0f; +#pragma unroll + for (int x = 0; x < 4; ++x) ss[i] = ss[i] + act[i][x] * act[i][x]; + } +#pragma unroll + for (int o = 16; o; o >>= 1) // warp_sum's tree, the rows interleaved +#pragma unroll + for (int i = 0; i < 4; ++i) ss[i] += __shfl_xor_sync(0xffffffffu, ss[i], o); + float (*dst)[DK] = part ? sm.k : sm.q; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float inv = 1.0f / sqrtf(ss[i] + 1e-6f); + float y[4]; +#pragma unroll + for (int x = 0; x < 4; ++x) { + y[x] = act[i][x] * inv; + if (part == 0) y[x] = y[x] * (1.0f / sqrtf((float)DK)); + } + *reinterpret_cast(&dst[rb + i][4 * lane]) = make_float4(y[0], y[1], y[2], y[3]); + } + } else { +#pragma unroll + for (int i = 0; i < 4; ++i) { + __nv_bfloat162 lo = __floats2bfloat162_rn(act[i][0], act[i][1]); + __nv_bfloat162 hi = __floats2bfloat162_rn(act[i][2], act[i][3]); + uint2 o; + o.x = *reinterpret_cast(&lo); o.y = *reinterpret_cast(&hi); + *reinterpret_cast(&sm.v[rb + i][4 * lane]) = o; + } + } + } + { + const int cc = tid & (DK - 1); +#pragma unroll + for (int j = tid >> 7; j < BT; j += NT / DK) { + const int r = r0 + j; + float gv = 1.0f; + if (r >= 0 && r < rows) { + const float a = __bfloat162float(sm.raw_a[j][cc]) + dtb; + const float den = 1.0f + expf(-(rate * a)); + const float sg = div_ok(1.0f, den) ? div_rn(1.0f, den) : 1.0f / den; // glm_kda.cu's sigmoidf_ + gv = expf(lower * sg); + } + sm.g[j][cc] = gv; + } + if (tid < BT) { + const int r = r0 + tid; + sm.beta[tid] = (r >= 0 && r < rows) ? bf(sigmoidf_(sm.raw_b[tid])) : 0.0f; + } + } + __syncthreads(); + + // 3. L and M. Group pr = tid / 16 takes column t = pr (rows s > pr), then t = 31 - pr (31 entries in all, + + // every group alike), in batches of 8; lane cg holds the key quads cg and cg + 16, k_t E(t -> s) multiplied + + // row after row; a batch's 16 sums reduce over the 16 lanes by a transposing tree. + { + const int pr = tid >> 4, cg = tid & 15; + const int n1 = BT - 1 - pr; // entries of column pr + float kt[8], e[8]; + auto load_t = [&](int tt) { +#pragma unroll + for (int u = 0; u < 2; ++u) { + const float4 x = reinterpret_cast(sm.k[tt])[cg + 16 * u]; + kt[4 * u] = x.x; kt[4 * u + 1] = x.y; kt[4 * u + 2] = x.z; kt[4 * u + 3] = x.w; + } +#pragma unroll + for (int i = 0; i < 8; ++i) e[i] = kt[i]; // k_t E(t -> s), multiplied row after row + }; + { // the diagonal of M for both columns + float dg[2]; +#pragma unroll + for (int hh = 0; hh < 2; ++hh) { + const int tt = hh ? BT - 1 - pr : pr; + load_t(tt); + float a = 0.0f; +#pragma unroll + for (int u = 0; u < 2; ++u) { + const float4 x = reinterpret_cast(sm.q[tt])[cg + 16 * u]; + a = fma_(x.x, kt[4 * u], a); a = fma_(x.y, kt[4 * u + 1], a); + a = fma_(x.z, kt[4 * u + 2], a); a = fma_(x.w, kt[4 * u + 3], a); + } + dg[hh] = a; + } +#pragma unroll + for (int o = 8; o; o >>= 1) { + dg[0] += __shfl_xor_sync(0xffffffffu, dg[0], o); + dg[1] += __shfl_xor_sync(0xffffffffu, dg[1], o); + } + if (cg == 0) { + sm.u.x.mt[pr][pr] = dg[0]; + sm.u.x.mt[BT - 1 - pr][BT - 1 - pr] = dg[1]; + } + } + load_t(pr); +#pragma unroll + for (int b = 0; b < 4; ++b) { + float val[16]; +#pragma unroll + for (int ii = 0; ii < 8; ++ii) { + const int i = 8 * b + ii; + if (i == n1 && n1 < BT - 1) load_t(BT - 1 - pr); + const int s2 = i < n1 ? pr + 1 + i : BT - pr + (i - n1); + float ak = 0.0f, aq = 0.0f; + if (i < BT - 1) { +#pragma unroll + for (int u = 0; u < 2; ++u) { + const float4 gx = reinterpret_cast(sm.g[s2])[cg + 16 * u]; + const float4 kx = reinterpret_cast(sm.k[s2])[cg + 16 * u]; + const float4 qx = reinterpret_cast(sm.q[s2])[cg + 16 * u]; + const float gg[4] = {gx.x, gx.y, gx.z, gx.w}, kk[4] = {kx.x, kx.y, kx.z, kx.w}; + const float qq[4] = {qx.x, qx.y, qx.z, qx.w}; +#pragma unroll + for (int x = 0; x < 4; ++x) { + e[4 * u + x] = e[4 * u + x] * gg[x]; + ak = fma_(kk[x], e[4 * u + x], ak); + aq = fma_(qq[x], e[4 * u + x], aq); + } + } + } + val[2 * ii] = ak; + val[2 * ii + 1] = aq; + } + const float sum = tree16(val, lane, 8, 4, 2, 1); // entry cg / 2 of the batch, L (even) or M + const int i = 8 * b + (cg >> 1); + if (i < BT - 1) { + const int tt = i < n1 ? pr : BT - 1 - pr; + const int s2 = i < n1 ? pr + 1 + i : BT - pr + (i - n1); + if (cg & 1) { + sm.u.x.mt[tt][s2] = sum; + sm.u.x.mt[s2][tt] = 0.0f; // M[tt][s2], above the diagonal + } else { + sm.u.x.L[s2][tt] = sm.beta[s2] * sum; + } + } + } + } + __syncthreads(); + + // 4. T = (I + L)^-1 by warp 0: lane l inverts column l % 16 of diagonal block l / 16 by forward + + // substitution, then block (1, 0) = -T11 (L10 T00). Warps 1-7 per key channel: k E(s ->), then k E(s) in + + // place, and E(31) (threads 32-159); q E(s) in place (threads 160-255, then 160-191 for the last 32 channels) + if (warp == 0) { + const int blk = lane >> 4, j = lane & 15, base = 16 * blk; + float col[16]; +#pragma unroll + for (int s2 = 0; s2 < 16; ++s2) { + float acc = 0.0f; +#pragma unroll + for (int m = 0; m < s2; ++m) acc = fma_(sm.u.x.L[base + s2][base + m], col[m], acc); + col[s2] = s2 == j ? 1.0f : (s2 > j ? -acc : 0.0f); + } +#pragma unroll + for (int s2 = 0; s2 < 16; ++s2) sm.u.x.tt[base + j][base + s2] = col[s2]; + if (blk == 1) { // block (0, 1), above the diagonal: zeros +#pragma unroll + for (int s2 = 0; s2 < 16; ++s2) sm.u.x.tt[16 + j][s2] = 0.0f; + } + __syncwarp(); + if (blk == 0) { + float xr[16]; // (L10 T00)[:, j] +#pragma unroll + for (int s2 = 0; s2 < 16; ++s2) { + float acc = 0.0f; +#pragma unroll + for (int m = 0; m < 16; ++m) acc = fma_(sm.u.x.L[16 + s2][m], col[m], acc); + xr[s2] = acc; + } +#pragma unroll + for (int s2 = 0; s2 < 16; ++s2) { + float acc = 0.0f; +#pragma unroll + for (int m = 0; m <= s2; ++m) // T11[s2][m] + acc = fma_(sm.u.x.tt[16 + m][16 + s2], xr[m], acc); + sm.u.x.tt[j][16 + s2] = -acc; + } + } + } else { + const int t2 = tid - 32; + if (t2 < DK) { + const int cc = t2; + float rr = 1.0f; +#pragma unroll 8 + for (int j = BT - 1; j >= 0; --j) { sm.u.x.ke[j][cc] = sm.k[j][cc] * rr; rr = rr * sm.g[j][cc]; } + float pp = 1.0f; +#pragma unroll 8 + for (int j = 0; j < BT; ++j) { pp = pp * sm.g[j][cc]; sm.k[j][cc] = sm.k[j][cc] * pp; } + sm.egl[cc] = pp; + } else { + for (int cc = t2 - DK; cc < DK; cc += NT - 32 - DK) { + float pp = 1.0f; +#pragma unroll 8 + for (int j = 0; j < BT; ++j) { pp = pp * sm.g[j][cc]; sm.q[j][cc] = sm.q[j][cc] * pp; } + } + } + } + __syncthreads(); + + // 5. Every thread: its sums of (k E) S0^T and (q E) S0^T; Y = beta (v - (k E) S0^T) into g, the (q E) S0^T + + // of the rows it finishes kept in o1. + float o1[OWN]; + { + float4 ck[NQ], cq[NQ], nk[NQ], nq[NQ]; + auto load_row = [&](float4 (&kx)[NQ], float4 (&qx)[NQ], int r) { +#pragma unroll + for (int i = 0; i < NQ; ++i) { + kx[i] = reinterpret_cast(sm.k[r])[kp + KP * i]; + qx[i] = reinterpret_cast(sm.q[r])[kp + KP * i]; + } + }; + // row r's 8 sums (k sums at ok, q sums at 4 - ok), reduced: entry kp, value column 4 vg + kp % 4 + auto row_sums = [&](const float4 (&kx)[NQ], const float4 (&qx)[NQ], int ok) { + float a[8]; +#pragma unroll + for (int x = 0; x < 8; ++x) a[x] = 0.0f; + const int oq = 4 - ok; +#pragma unroll + for (int i = 0; i < NQ; ++i) +#pragma unroll + for (int vv = 0; vv < VPT; ++vv) { + a[ok + vv] = fma_(kx[i].x, s[vv][i][0], a[ok + vv]); + a[ok + vv] = fma_(kx[i].y, s[vv][i][1], a[ok + vv]); + a[ok + vv] = fma_(kx[i].z, s[vv][i][2], a[ok + vv]); + a[ok + vv] = fma_(kx[i].w, s[vv][i][3], a[ok + vv]); + a[oq + vv] = fma_(qx[i].x, s[vv][i][0], a[oq + vv]); + a[oq + vv] = fma_(qx[i].y, s[vv][i][1], a[oq + vv]); + a[oq + vv] = fma_(qx[i].z, s[vv][i][2], a[oq + vv]); + a[oq + vv] = fma_(qx[i].w, s[vv][i][3], a[oq + vv]); + } + return tree8(a, lane); + }; + load_row(ck, cq, 0); +#pragma unroll 1 + for (int rg = 0; rg < BT / 8; ++rg) { // 8 rows: pairs 2 rp (even: k sums first), 2 rp + 1 +#pragma unroll + for (int u = 0; u < 4; ++u) { + const int r = 8 * rg + 2 * u; + load_row(nk, nq, r + 1); + const float se = row_sums(ck, cq, 0); + load_row(ck, cq, r + 2 < BT ? r + 2 : r); // (the last pair reloads a row: no predicate) + const float so = row_sums(nk, nq, 4); + // this thread keeps the q sum of its parity's row and writes Y of the other + const float qs = par ? so : se, ks = par ? se : so; + const int ry = par ? r : r + 1; + sm.g[ry][ov] = (__bfloat162float(sm.v[ry][ov]) - ks) * sm.beta[ry]; + o1[OWN - 4 + u] = qs; + } + if (rg + 1 < BT / 8) { // after the loop o1[j] is row 2 j + par's +#pragma unroll + for (int j = 0; j + 4 < OWN; ++j) o1[j] = o1[j + 4]; + } + } + } + __syncthreads(); + + // 6. load this sub-chunk's output gates (warm in L2); (q E) S0^T into k (free now); V' = T Y into q. Tiles: + + // warp w the rows 4 w' .. 4 w' + 3 with w' = w (w < 4) or 11 - w (the two warps of a scheduler, w and w + 4, + + // then run 36 steps of t together), lane l the value columns 4 l .. 4 l + 3 (T is zero above its diagonal: + + // t runs to the tile's last row) + const int rq = 4 * (warp < 4 ? warp : 11 - warp), vq = 4 * lane; + float gate[4][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int r = r0 + rq + i; + const bool ok = r >= 0 && r < rows; + const __nv_bfloat16* gp = G + (size_t)(ok ? r : 0) * g_stride + h * DV + vq; + if (vec_g) { + const uint2 x = ok ? *reinterpret_cast(gp) : make_uint2(0, 0); + const __nv_bfloat162 g0 = *reinterpret_cast(&x.x); + const __nv_bfloat162 g1 = *reinterpret_cast(&x.y); + gate[i][0] = __low2float(g0); gate[i][1] = __high2float(g0); + gate[i][2] = __low2float(g1); gate[i][3] = __high2float(g1); + } else { +#pragma unroll + for (int x = 0; x < 4; ++x) gate[i][x] = ok ? ldb(gp + x) : 0.0f; + } + } +#pragma unroll + for (int j = 0; j < OWN; ++j) sm.k[2 * j + par][ov] = o1[j]; + { + float acc[4][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int x = 0; x < 4; ++x) acc[i][x] = 0.0f; +#pragma unroll 4 + for (int t = 0; t < rq + 4; ++t) { + const float4 tq = *reinterpret_cast(&sm.u.x.tt[t][rq]); + const float4 y = *reinterpret_cast(&sm.g[t][vq]); + const float tv[4] = {tq.x, tq.y, tq.z, tq.w}, yv[4] = {y.x, y.y, y.z, y.w}; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int x = 0; x < 4; ++x) acc[i][x] = fma_(tv[i], yv[x], acc[i][x]); + } +#pragma unroll + for (int i = 0; i < 4; ++i) + *reinterpret_cast(&sm.q[rq + i][vq]) = make_float4(acc[i][0], acc[i][1], acc[i][2], acc[i][3]); + } + __syncthreads(); + + // 7. read-outs O = (q E) S0^T + M V' of the same tile, bf16, the gated RMSNorm (glm_kda.cu's formula; a + + // warp holds its rows' 128 values); the state update + { + float o[4][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float4 x = *reinterpret_cast(&sm.k[rq + i][vq]); + o[i][0] = x.x; o[i][1] = x.y; o[i][2] = x.z; o[i][3] = x.w; + } +#pragma unroll 4 + for (int t = 0; t < rq + 4; ++t) { + const float4 mq = *reinterpret_cast(&sm.u.x.mt[t][rq]); + const float4 vn = *reinterpret_cast(&sm.q[t][vq]); + const float mv[4] = {mq.x, mq.y, mq.z, mq.w}, vv4[4] = {vn.x, vn.y, vn.z, vn.w}; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int x = 0; x < 4; ++x) o[i][x] = fma_(mv[i], vv4[x], o[i][x]); + } + float ss[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + ss[i] = 0.0f; +#pragma unroll + for (int x = 0; x < 4; ++x) { o[i][x] = bf(o[i][x]); ss[i] = ss[i] + o[i][x] * o[i][x]; } + } +#pragma unroll + for (int m = 16; m; m >>= 1) +#pragma unroll + for (int i = 0; i < 4; ++i) ss[i] += __shfl_xor_sync(0xffffffffu, ss[i], m); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int r = r0 + rq + i; + if (r < 0 || r >= rows) continue; + const float rinv = 1.0f / sqrtf(ss[i] / (float)DV + eps); + __nv_bfloat16 ob[4]; +#pragma unroll + for (int x = 0; x < 4; ++x) { + const float yn = o[i][x] * rinv; + const float yw = nw[x] * yn; + const float den = 1.0f + expf(-gate[i][x]); + const float sg = div_ok(1.0f, den) ? div_rn(1.0f, den) : 1.0f / den; // glm_kda.cu's sigmoidf_ + ob[x] = __float2bfloat16_rn(yw * sg); + } + __nv_bfloat16* op = out + (size_t)r * H * DV + h * DV + vq; + if (vec_o) *reinterpret_cast(op) = *reinterpret_cast(ob); + else { +#pragma unroll + for (int x = 0; x < 4; ++x) op[x] = ob[x]; + } + } + } +#pragma unroll + for (int i = 0; i < NQ; ++i) { + const float4 ge = reinterpret_cast(sm.egl)[kp + KP * i]; +#pragma unroll + for (int vv = 0; vv < VPT; ++vv) { + s[vv][i][0] = s[vv][i][0] * ge.x; s[vv][i][1] = s[vv][i][1] * ge.y; + s[vv][i][2] = s[vv][i][2] * ge.z; s[vv][i][3] = s[vv][i][3] * ge.w; + } + } +#pragma unroll 4 + for (int t = 0; t < BT; ++t) { + const float4 vn = *reinterpret_cast(&sm.q[t][VPT * vg]); + const float vs[4] = {vn.x, vn.y, vn.z, vn.w}; +#pragma unroll + for (int i = 0; i < NQ; ++i) { + const float4 kx = reinterpret_cast(sm.u.x.ke[t])[kp + KP * i]; +#pragma unroll + for (int vv = 0; vv < VPT; ++vv) { + s[vv][i][0] = fma_(vs[vv], kx.x, s[vv][i][0]); s[vv][i][1] = fma_(vs[vv], kx.y, s[vv][i][1]); + s[vv][i][2] = fma_(vs[vv], kx.z, s[vv][i][2]); s[vv][i][3] = fma_(vs[vv], kx.w, s[vv][i][3]); + } + } + } + __syncthreads(); + } +#pragma unroll + for (int vv = 0; vv < VPT; ++vv) +#pragma unroll + for (int i = 0; i < NQ; ++i) + reinterpret_cast(state_out + sbase + (size_t)(VPT * vg + vv) * DK)[kp + KP * i] = + make_float4(s[vv][i][0], s[vv][i][1], s[vv][i][2], s[vv][i][3]); +} + diff --git a/zig/kernels/cuda/glm_kda_chunk_shared.cuh b/zig/kernels/cuda/glm_kda_chunk_shared.cuh new file mode 100644 index 0000000000..3373f8b34f --- /dev/null +++ b/zig/kernels/cuda/glm_kda_chunk_shared.cuh @@ -0,0 +1,234 @@ +// GLM-5.3-Flash's KDA prompt chunks in chunked (WY) form: one block of 256 threads per head runs the prompt + +// chunk's 32-row sub-chunks in order (sub-chunks at absolute positions), the 128 x 128 state in registers. + +// + +// Per row the inputs are glm_kda.cu's (conv + SiLU rounded to bf16, fp32 L2 norms, q times 128^-0.5, beta = + +// bf16(sigmoid(b)), the per-channel decay g = e^l, l = lower * sigmoid(exp(A_log) * (a + dt_bias))); the delta rule + +// S <- S diag(g); S <- S + beta (v - S k) k^T; o = S q, o rounded to bf16, then the gated RMSNorm. S is [value, key]. + +// + +// A sub-chunk of rows 0..31 entering state S0, E(t -> s) = g_{t+1} ... g_s per key channel (empty: 1): + +// + +// L[s, t] = beta_s sum_c k_s k_t E(t -> s) t < s M[s, t] = sum_c q_s k_t E(t -> s) t <= s + +// T = (I + L)^-1 Y = beta (v - (k E(0 -> row)) S0^T), V' = T Y + +// O = (q E(0 -> row)) S0^T + M V' S_end = S0 diag(E(0 -> 31)) + V'^T (k E(row -> 31)) + +// + +// Every decay is a running product of per-row decays <= 1, multiplied in row order as the serial chain decays its + +// state (no exponent of a difference of cumulative sums). All products run in fp32 FMAs in a fixed order, so the + +// kernel is deterministic, and a row's bits depend only on its sub-chunk's rows and the entering state: the same + +// whichever prompt chunk the row came in, as long as prompt chunks start on multiples of the sub-chunk. Rows + +// outside the prompt chunk (a short first or last sub-chunk) are zero rows: decay 1, no key, no value, beta 0. + +// + +// Layout. The state: thread (vg, kp) = (tid / 8, tid % 8) holds value rows 4 vg .. 4 vg + 3 and key quads kp + 8 i + +// (i < 4): an 8-lane group reads 128 contiguous bytes of a key row per load, and products with S reduce over the 8 + +// lanes by a transposing shuffle tree (7 shuffles for 8 sums). Read-outs and corrections: thread (vg, kp) finishes + +// value column 4 vg + kp % 4 of the rows of one parity (odd for kp < 4, even otherwise), 16 rows. + +// This header: the constants, the shared-memory layout and the sub-chunk helpers; glm_kda_chunk.cu the kernel. + +#include +#include + + +constexpr int DK = 128, DV = 128, TAPS = 4; +constexpr int BT = 32; // rows of a sub-chunk +constexpr int HALO = TAPS - 1; // conv rows before a sub-chunk +constexpr int NT = 256, NW = NT / 32; // threads, warps +constexpr int VPT = 4; // value rows of the state a thread holds +constexpr int KP = NT * VPT / DV; // threads sharing those value rows (8), each 16 keys +constexpr int NQ = DK / 4 / KP; // key quads a thread holds (4) +constexpr int OWN = BT / 2; // rows a thread finishes (16, one parity) +constexpr int LS = BT + 4; // row stride of L +constexpr int PCH = 3 * DK; // projection channels of a head (q | k | v) +constexpr int PF_A = BT * DK / 8 / NT; // decay prefetch loads a thread issues (2) +static_assert(KP == 8 && NQ == 4 && PF_A * NT * 8 == BT * DK, "layout"); + +struct Smem { + float q[BT][DK]; // q -> q E(s) -> V' + float k[BT][DK]; // k -> k E(s) -> (q E) S0^T + float g[BT][DK]; // decays -> Y + union { // staged projection rows (conv input) -> k E(s ->), L, T, M + __nv_bfloat16 raw_p[BT + HALO][PCH]; + struct { + float ke[BT][DK]; // k E(s ->) + float L[BT][LS]; // L[s][t] + float tt[BT][BT]; // T transposed: tt[t][s] = T[s][t] + float mt[BT][BT]; // M transposed: mt[t][s] = M[s][t] + } x; + } u; + __nv_bfloat16 v[BT][DV]; + __nv_bfloat16 raw_a[BT][DK]; // staged decay rows + __nv_bfloat16 cw[PCH][TAPS]; // this head's conv weights + float raw_b[BT]; + float beta[BT]; + float egl[DK]; // E(31) +}; + +__device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } +__device__ __forceinline__ float ldb(const __nv_bfloat16* p) { return __bfloat162float(*p); } +__device__ __forceinline__ float fma_(float a, float b, float c) { return __fmaf_rn(a, b, c); } +__device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); } + +// a / b with IEEE rounding, branch-free: the fast path of CUDA's division (a refined reciprocal, then one + +// correction), which is the correctly rounded quotient whenever ``div_ok`` holds; callers take a / b for the rare + +// operands outside (a huge denominator, a tiny numerator). +__device__ __forceinline__ float div_rn(float a, float b) { + float r; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(b)); + const float e = __fmaf_rn(-b, r, 1.0f); + r = __fmaf_rn(e, r, r); + const float q = a * r; + const float rem = __fmaf_rn(-b, q, a); + return __fmaf_rn(rem, r, q); +} + +__device__ __forceinline__ bool div_ok(float a, float b) { + return (a == 0.0f || fabsf(a) > 1e-36f) && b >= 1.0f && b < 8e37f; +} + +__device__ __forceinline__ float warp_sum(float x) { + for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o); + return x; +} + +// 8 bf16 from p: one 16-byte load when aligned, else 8 loads. +__device__ __forceinline__ uint4 ld8(const __nv_bfloat16* p, bool vec) { + if (vec) return *reinterpret_cast(p); + const unsigned short* s = reinterpret_cast(p); + uint4 r; + r.x = s[0] | ((unsigned)s[1] << 16); r.y = s[2] | ((unsigned)s[3] << 16); + r.z = s[4] | ((unsigned)s[5] << 16); r.w = s[6] | ((unsigned)s[7] << 16); + return r; +} + +// One level of a transposing shuffle tree: this lane keeps the half of a[0..N) its ``hi`` bit selects, adds the + +// partner's (lane ^ mask) copy of that half, and sends the other half. +template +__device__ __forceinline__ void tree_level(float (&a)[N], int mask, bool hi) { +#pragma unroll + for (int j = 0; j < N / 2; ++j) { + const float send = hi ? a[j] : a[j + N / 2]; + const float keep = hi ? a[j + N / 2] : a[j]; + a[j] = keep + __shfl_xor_sync(0xffffffffu, send, mask); + } +} + +// Sums of a[0..8) over the lanes xor 4, 2, 1: lane bits (2, 1, 0) = i gets the sum of a[i]. +__device__ __forceinline__ float tree8(float (&a)[8], int lane) { + tree_level<8>(a, 4, lane & 4); + float b[4] = {a[0], a[1], a[2], a[3]}; + tree_level<4>(b, 2, lane & 2); + float c[2] = {b[0], b[1]}; + tree_level<2>(c, 1, lane & 1); + return c[0]; +} + +// Sums of a[0..16) over the lanes xor m0 > m1 > m2 > m3: the lane whose bits (m0, m1, m2, m3) spell i gets a[i]'s. +__device__ __forceinline__ float tree16(float (&a)[16], int lane, int m0, int m1, int m2, int m3) { + tree_level<16>(a, m0, lane & m0); + float b[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) b[j] = a[j]; + tree_level<8>(b, m1, lane & m1); + float c[4] = {b[0], b[1], b[2], b[3]}; + tree_level<4>(c, m2, lane & m2); + float d[2] = {c[0], c[1]}; + tree_level<2>(d, m3, lane & m3); + return d[0]; +} + +__device__ __forceinline__ void prefetch_l2(const void* p) { + asm volatile("prefetch.global.L2 [%0];" :: "l"(p)); +} + +__device__ __forceinline__ void cp16(void* dst, const void* src) { + const unsigned d = (unsigned)__cvta_generic_to_shared(dst); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" :: "r"(d), "l"(src)); +} + +// Warm L2 with the rows of the sub-chunk whose row 0 is prompt-chunk row r0 (projection rows r0 - 3 .. r0 + 31 of + +// this head's q | k | v channels, its decay rows and beta logits; not when ``next`` is false) and the output gate + +// rows from gate_r0. +__device__ __forceinline__ void warm(int tid, int h, int H, int r0, bool next, int rows, int gate_r0, + 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) { + constexpr int LP = (BT + HALO) * 3 * 2, LA = BT * 2, LB = BT, LG = BT * 2; // 128-byte lines + for (int x = next ? tid : tid + LP + LA + LB; x < LP + LA + LB + LG; x += NT) { + if (x < LP) { + const int row = x / 6, part = (x % 6) / 2, half = x % 2, r = r0 - HALO + row; + if (r >= 0 && r < rows) prefetch_l2(P + (size_t)r * p_stride + part * H * DK + h * DK + 64 * half); + } else if (x < LP + LA) { + const int y = x - LP, r = r0 + y / 2; + if (r >= 0 && r < rows) prefetch_l2(A + (size_t)r * a_stride + h * DK + 64 * (y % 2)); + } else if (x < LP + LA + LB) { + const int r = r0 + (x - LP - LA); + if (r >= 0 && r < rows) prefetch_l2(P + (size_t)r * p_stride + b_off + h); + } else { + const int y = x - LP - LA - LB, r = gate_r0 + y / 2; + if (r >= 0 && r < rows) prefetch_l2(G + (size_t)r * g_stride + h * DV + 64 * (y % 2)); + } + } +} + +// Stage the rows of the sub-chunk whose row 0 is prompt-chunk row r0 into shared memory: projection rows r0 - 3 .. + +// r0 + 31 (the conv state before row 0, zero rows outside), decay rows and the beta logits. Global to shared + +// directly (cp.async) when the rows are 16-byte aligned; the caller waits (cp.async.wait_all) and syncs. +__device__ __forceinline__ void stage(Smem& sm, int tid, int h, int H, int r0, int rows, + const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, + const __nv_bfloat16* __restrict__ A, int a_stride, + const __nv_bfloat16* __restrict__ cs, bool vec_p, bool vec_a) { + const int C = 3 * H * DK; + for (int item = tid; item < (BT + HALO) * PCH / 8; item += NT) { + const int row = item / (PCH / 8), q8 = item % (PCH / 8); + const int ch = (q8 / (DK / 8)) * H * DK + h * DK + 8 * (q8 % (DK / 8)); + const int r = r0 - HALO + row; + uint4* dst = reinterpret_cast(&sm.u.raw_p[row][0]) + q8; + const __nv_bfloat16* src = nullptr; + if (r >= 0 && r < rows) src = P + (size_t)r * p_stride + ch; + else if (r < 0 && r >= -HALO) src = cs + (size_t)(r + HALO) * C + ch; + if (src == nullptr) *dst = make_uint4(0, 0, 0, 0); + else if (vec_p) cp16(dst, src); + else *dst = ld8(src, false); + } + for (int item = tid; item < BT * DK / 8; item += NT) { + const int row = item / (DK / 8), c8 = item % (DK / 8); + const int r = r0 + row; + uint4* dst = reinterpret_cast(&sm.raw_a[row][0]) + c8; + const __nv_bfloat16* src = A + (size_t)r * a_stride + h * DK + 8 * c8; + if (r < 0 || r >= rows) *dst = make_uint4(0, 0, 0, 0); + else if (vec_a) cp16(dst, src); + else *dst = ld8(src, false); + } + if (tid < BT) { + const int r = r0 + tid; + sm.raw_b[tid] = (r >= 0 && r < rows) ? ldb(P + (size_t)r * p_stride + b_off + h) : 0.0f; + } +} diff --git a/zig/kernels/cuda/glm_latent.cu b/zig/kernels/cuda/glm_latent.cu new file mode 100644 index 0000000000..64d21c9d85 --- /dev/null +++ b/zig/kernels/cuda/glm_latent.cu @@ -0,0 +1,216 @@ +// GLM-5.3-Flash latent.py part 1: the fp8/bf16 cache write and the kv_b absorb/expand, in Triton's SASS orders. + +#include +#include + +namespace tf_glm_latent { + +using bf16 = __nv_bfloat16; + +constexpr int H = 32; // heads +constexpr int D = 256; // qk dims +constexpr int LW = 512; // latent width (qa / ol) +constexpr int V = 256; // v dims (the expand output) + +__device__ __forceinline__ float fmax32(float a, float b) { + float r; + asm("max.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float div_full(float a, float b) { + float r; + asm("div.full.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ unsigned pack_bf16x2(float lo, float hi) { + unsigned r; + asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(hi), "f"(lo)); + return r; +} + +__device__ __forceinline__ unsigned cvt_e4m3x2(float hi, float lo) { + unsigned short r; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(r) : "f"(hi), "f"(lo)); + return r; +} + +__device__ __forceinline__ float bf16_of(unsigned u, int half) { + __nv_bfloat16_raw hr; + hr.x = half ? (unsigned short)(u >> 16) : (unsigned short)u; + return __bfloat162float(__nv_bfloat16(hr)); +} + +struct LatWriteArgs { + const bf16* lat; // bf16 [R, LW] (lat_stride elements a row) + long long lat_stride; + void* lc; // bf16 or fp8 [cap, LW] + const int* pos; // int32 [1] + float* sc; // fp32 [cap] (fp8 only) + int fp8; +}; + +// Row r: amax / 448 scale (fp8) or a plain copy (bf16) into the cache slot the device-side pos names. +extern "C" __global__ void tf_glm_lat_write(LatWriteArgs a) { + const int r = blockIdx.x; + const int t = threadIdx.x; + const long long p = a.pos[0]; + const uint2 raw = *reinterpret_cast(a.lat + (long long)r * a.lat_stride + 4 * t); + if (a.fp8 == 0) { + *reinterpret_cast(reinterpret_cast(a.lc) + (p + r) * LW + 4 * t) = raw; + return; + } + const float e0 = bf16_of(raw.x, 0), e1 = bf16_of(raw.x, 1), e2 = bf16_of(raw.y, 0), e3 = bf16_of(raw.y, 1); + float amax = fmax32(fabsf(e0), fabsf(e1)); + amax = fmax32(amax, fabsf(e2)); + amax = fmax32(amax, fabsf(e3)); + const unsigned all = 0xffffffffu; + amax = fmax32(amax, __shfl_xor_sync(all, amax, 16)); + amax = fmax32(amax, __shfl_xor_sync(all, amax, 8)); + amax = fmax32(amax, __shfl_xor_sync(all, amax, 4)); + amax = fmax32(amax, __shfl_xor_sync(all, amax, 2)); + amax = fmax32(amax, __shfl_xor_sync(all, amax, 1)); + __shared__ float wmax[4]; + if ((t & 31) == 0) wmax[t >> 5] = amax; + __syncthreads(); + amax = fmax32(fmax32(wmax[0], wmax[2]), fmax32(wmax[1], wmax[3])); + const float s = amax > 0.0f ? div_full(amax, 448.0f) : 1.0f; + if (t == 0) a.sc[p + r] = s; + const unsigned codes = cvt_e4m3x2(div_full(e1, s), div_full(e0, s)) | + (cvt_e4m3x2(div_full(e3, s), div_full(e2, s)) << 16); + *reinterpret_cast(reinterpret_cast(a.lc) + (p + r) * LW + 4 * t) = codes; +} + +struct AbsorbArgs { + const bf16* q; // bf16 [R, H, D] + const bf16* wk; // bf16 [H, D, LW] + bf16* qa; // bf16 [R, H, LW] + int R; +}; + +// qa = q . wk: per col tree ((j0+j1)+(j2+j3))+((j4+j5)+(j6+j7)), shfl 16/8/4, cross-warp (P0+P2)+(P1+P3). +extern "C" __global__ void __launch_bounds__(128) tf_glm_absorb_q(AbsorbArgs a) { + const int h = blockIdx.x; + const int n0 = blockIdx.y * 32; + const int t = threadIdx.x; + const int l = t & 31; + uint4 w8[8]; + for (int j = 0; j < 8; ++j) + w8[j] = *reinterpret_cast(a.wk + ((long long)(h * D + (t >> 2)) + 32 * j) * LW + n0 + (l & 3) * 8); + __shared__ float xw[4][32][8]; + __shared__ unsigned col_sums[16]; + for (int r = 0; r < a.R; ++r) { + const bf16* qrow = a.q + ((long long)r * H + h) * D; + float qv[8]; + for (int j = 0; j < 8; ++j) qv[j] = __bfloat162float(qrow[(t >> 2) + 32 * j]); + float sa[4], sb[4]; + for (int pc = 0; pc < 4; ++pc) { + float pa[8], pb[8]; + for (int j = 0; j < 8; ++j) { + const unsigned u = (pc == 0 ? w8[j].x : pc == 1 ? w8[j].y : pc == 2 ? w8[j].z : w8[j].w); + pa[j] = __fmul_rn(qv[j], bf16_of(u, 0)); + pb[j] = __fmul_rn(qv[j], bf16_of(u, 1)); + } + sa[pc] = __fadd_rn(__fadd_rn(__fadd_rn(pa[0], pa[1]), __fadd_rn(pa[2], pa[3])), + __fadd_rn(__fadd_rn(pa[4], pa[5]), __fadd_rn(pa[6], pa[7]))); + sb[pc] = __fadd_rn(__fadd_rn(__fadd_rn(pb[0], pb[1]), __fadd_rn(pb[2], pb[3])), + __fadd_rn(__fadd_rn(pb[4], pb[5]), __fadd_rn(pb[6], pb[7]))); + } + for (int pc = 0; pc < 4; ++pc) { + sa[pc] = __fadd_rn(sa[pc], __shfl_xor_sync(0xffffffffu, sa[pc], 16)); + sa[pc] = __fadd_rn(sa[pc], __shfl_xor_sync(0xffffffffu, sa[pc], 8)); + sa[pc] = __fadd_rn(sa[pc], __shfl_xor_sync(0xffffffffu, sa[pc], 4)); + sb[pc] = __fadd_rn(sb[pc], __shfl_xor_sync(0xffffffffu, sb[pc], 16)); + sb[pc] = __fadd_rn(sb[pc], __shfl_xor_sync(0xffffffffu, sb[pc], 8)); + sb[pc] = __fadd_rn(sb[pc], __shfl_xor_sync(0xffffffffu, sb[pc], 4)); + xw[t >> 5][l][2 * pc] = sa[pc]; + xw[t >> 5][l][2 * pc + 1] = sb[pc]; + } + __syncthreads(); + float fin[8]; + for (int i = 0; i < 8; ++i) + fin[i] = __fadd_rn(__fadd_rn(xw[0][l][i], xw[2][l][i]), __fadd_rn(xw[1][l][i], xw[3][l][i])); + for (int pc = 0; pc < 4; ++pc) col_sums[(l & 3) * 4 + pc] = pack_bf16x2(fin[2 * pc], fin[2 * pc + 1]); + __syncthreads(); + if (t < 32) { + const unsigned cv = col_sums[(t & 31) >> 1]; + __nv_bfloat16_raw hr; + hr.x = (t & 1) ? (unsigned short)(cv >> 16) : (unsigned short)cv; + a.qa[((long long)r * H + h) * LW + n0 + (t & 31)] = __nv_bfloat16(hr); + } + } +} + +struct ExpandArgs { + const bf16* ol; // bf16 [R, H, LW] + const bf16* wv; // bf16 [H, V, LW] + bf16* out; // bf16 [R, H, V] + int R; +}; + +// out = ol . wv^T: even/odd k pairs (8th row fused), 5 chained shfl adds, warp pairs sum their k halves. +extern "C" __global__ void __launch_bounds__(128) tf_glm_expand_v(ExpandArgs a) { + const int h = blockIdx.x; + const int n0 = blockIdx.y * 16; + const int t = threadIdx.x; + const int l = t & 31; + const int w0 = (t >> 6) & 1, w1 = (t >> 5) & 1; + uint4 wv8[8]; + for (int i = 0; i < 8; ++i) + wv8[i] = *reinterpret_cast(a.wv + ((long long)(h * V + n0 + w0) + 2 * i) * LW + (l + 32 * w1) * 8); + __shared__ float xw[4][8]; + __shared__ unsigned short sh_bf[16]; + for (int r = 0; r < a.R; ++r) { + const uint4 oraw = *reinterpret_cast(a.ol + ((long long)r * H + h) * LW + (l + 32 * w1) * 8); + float ov[8]; + for (int cc = 0; cc < 4; ++cc) { + ov[2 * cc] = bf16_of((&oraw.x)[cc], 0); + ov[2 * cc + 1] = bf16_of((&oraw.x)[cc], 1); + } + float parts[8]; + for (int i = 0; i < 8; ++i) { + const unsigned* wr = &wv8[i].x; + float wf[8]; + for (int cc = 0; cc < 4; ++cc) { + wf[2 * cc] = bf16_of(wr[cc], 0); + wf[2 * cc + 1] = bf16_of(wr[cc], 1); + } + if (i < 7) { + const float pe = __fadd_rn(__fadd_rn(__fmul_rn(ov[0], wf[0]), __fmul_rn(ov[2], wf[2])), + __fadd_rn(__fmul_rn(ov[4], wf[4]), __fmul_rn(ov[6], wf[6]))); + const float po = __fadd_rn(__fadd_rn(__fmul_rn(ov[1], wf[1]), __fmul_rn(ov[3], wf[3])), + __fadd_rn(__fmul_rn(ov[5], wf[5]), __fmul_rn(ov[7], wf[7]))); + parts[i] = __fadd_rn(pe, po); + } else { + const float pe = __fadd_rn(__fmaf_rn(ov[0], wf[0], __fmul_rn(ov[2], wf[2])), + __fmaf_rn(ov[4], wf[4], __fmul_rn(ov[6], wf[6]))); + const float po = __fadd_rn(__fmaf_rn(ov[1], wf[1], __fmul_rn(ov[3], wf[3])), + __fmaf_rn(ov[5], wf[5], __fmul_rn(ov[7], wf[7]))); + parts[7] = __fadd_rn(pe, po); + } + } + for (int i = 0; i < 8; ++i) { + float v = parts[i]; + v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, 16)); + v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, 8)); + v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, 4)); + v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, 2)); + v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, 1)); + if (l == 0) xw[t >> 5][i] = v; + } + __syncthreads(); + if (l == 0) + for (int i = 0; i < 8; ++i) + sh_bf[w0 + 2 * i] = + __bfloat16_as_ushort(__float2bfloat16_rn(__fadd_rn(xw[2 * w0][i], xw[2 * w0 + 1][i]))); + __syncthreads(); + if (t < 16) { + __nv_bfloat16_raw hr; + hr.x = sh_bf[t]; + a.out[((long long)r * H + h) * V + n0 + t] = __nv_bfloat16(hr); + } + } +} + +} // namespace tf_glm_latent diff --git a/zig/kernels/cuda/glm_latent_att.cu b/zig/kernels/cuda/glm_latent_att.cu new file mode 100644 index 0000000000..e204fcb6e7 --- /dev/null +++ b/zig/kernels/cuda/glm_latent_att.cu @@ -0,0 +1,332 @@ +// GLM-5.3-Flash latent.py part 2: 512-key chunk attention over the fp8/bf16 cache and the merge, Triton's orders. + +#include +#include + +namespace tf_glm_latent { + +using bf16 = __nv_bfloat16; + +constexpr int H = 32; // heads +constexpr int LW = 512; // latent width +constexpr unsigned LOG2E_BITS = 0x3FB8AA3B; + +__device__ __forceinline__ float neg_inf() { return __uint_as_float(0xFF800000u); } + +__device__ __forceinline__ float fmax32(float a, float b) { + float r; + asm("max.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float ex2(float x) { + float r; + asm("ex2.approx.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +__device__ __forceinline__ float div_full(float a, float b) { + float r; + asm("div.full.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float log2e_mul(float x) { + return __fmul_rn(x, __uint_as_float(LOG2E_BITS)); +} + +__device__ __forceinline__ unsigned pack_bf16x2(float lo, float hi) { + unsigned r; + asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(hi), "f"(lo)); + return r; +} + +__device__ __forceinline__ unsigned cvt_e4m3x2_u16(unsigned short pair) { + unsigned r; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(r) : "h"(pair)); + return r; +} + +__device__ __forceinline__ float f16_of(unsigned u, int half) { + __half_raw hr; + hr.x = half ? (unsigned short)(u >> 16) : (unsigned short)u; + return __half2float(__half(hr)); +} + +__device__ __forceinline__ void mma_16n8k16(float (&d)[4], const unsigned (&a)[4], const unsigned (&b)[2]) { + 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};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +struct ChunksArgs { + const bf16* qa; // bf16 [R, H, LW] (u32 view: 256 a row) + const void* cache; // bf16 or fp8 [cap, LW] + const float* sc; // fp32 [cap] (fp8 only) + const int* pos; // int32 [1] (dense) + const int* tok; // int32 [R, W] (sparse) + const int* cnt; // int32 [R] (sparse) + float* po; // fp32 [nch, R, H, LW] + float* pm; // fp32 [nch, R, H] + float* pl; // fp32 [nch, R, H] + int R; + int tokw; // W, the sparse token stride + float scale; // score scale (dense 0.0625, sparse 1/sqrt(128)) +}; + +// One block per (row, head group, 512-key chunk): staged keys, 32 score mma steps, value mma, o seeded o*alpha. +template +__device__ __forceinline__ void attn_chunks(const ChunksArgs a) { + const int r = blockIdx.x; + const int t0 = blockIdx.y; + const int c = blockIdx.z; + const int t = threadIdx.x; + const int w = t >> 5, l = t & 31, m = t & 3; + const int cbase = c * 512; + const bool run = SPARSE ? (cbase < a.cnt[r]) : (cbase <= a.pos[0] + r); + __shared__ unsigned kv[32][256]; // 32 keys x 512 latents as bf16 pairs + __shared__ unsigned psh[32][16]; + __shared__ float tmx[32][2][4]; // per row, a/b tile max over the warp's 8-key slice (w&3) + __shared__ float lbuf[32][2][4]; // per row, a/b rowsum over the warp's 8-key slice + extern __shared__ __align__(16) unsigned qsh[]; // the block's HBT heads of qa (u32 view), staged once + const int rowA = (HBT == 32 ? 16 * (w >> 2) : 16 * t0) + (l >> 2); + const int rowB = rowA + 8; + const int kbase = (w & 3) * 8; + const int pw = 4 * (w & 3); + const int on_base = (HBT == 32 ? 128 * (w & 3) : 64 * w); + const int NT = (HBT == 32) ? 16 : 8; + float m_a = neg_inf(), m_b = neg_inf(), l_a = 0.0f, l_b = 0.0f; + float oacc[16][4]; + for (int j = 0; j < NT; ++j) { + oacc[j][0] = 0.0f; + oacc[j][1] = 0.0f; + oacc[j][2] = 0.0f; + oacc[j][3] = 0.0f; + } + const int hb0 = HBT == 32 ? 0 : 16 * t0; // the block's first head + if (run) { + const uint4* qsrc = reinterpret_cast(a.qa) + ((long long)r * H + hb0) * 64; + for (int i = t; i < HBT * 64; i += 256) reinterpret_cast(qsh)[i] = qsrc[i]; // read after the first sync + } + for (int tile = 0; run && tile < 16; ++tile) { // 16 tiles of 32 keys = the 512-key chunk + // a tile past the last visible key leaves m, l and o as they are (alpha 1, p 0): stop at the first one + if (SPARSE ? cbase + tile * 32 >= a.cnt[r] : cbase + tile * 32 > a.pos[0] + r) break; + // stage 32 keys x 512 latents: thread t holds key (t>>3), u32 columns (t&7)*32.. + const int kk = t >> 3; + int rowbase; + float s = 1.0f; + if (SPARSE) { + const int idx = cbase + tile * 32 + kk; + const bool ok = idx < a.cnt[r]; + rowbase = ok ? a.tok[(long long)r * a.tokw + idx] : 0; + s = ok ? a.sc[rowbase] : 1.0f; + } else { + rowbase = cbase + tile * 32 + kk; + if (FP8) s = a.sc[rowbase]; + } + unsigned* dst = &kv[kk][(t & 7) * 32]; + if (FP8) { + const uint4* src = reinterpret_cast(reinterpret_cast(a.cache) + + (long long)rowbase * LW) + (t & 7) * 4; + unsigned cw[16]; // the thread's 64 bytes as four 16-byte loads +#pragma unroll + for (int v = 0; v < 4; ++v) { + const uint4 u = src[v]; + cw[4 * v] = u.x; + cw[4 * v + 1] = u.y; + cw[4 * v + 2] = u.z; + cw[4 * v + 3] = u.w; + } +#pragma unroll + for (int i = 0; i < 16; ++i) { + const unsigned c4 = cw[i]; + const unsigned p01 = cvt_e4m3x2_u16((unsigned short)c4); + const unsigned p23 = cvt_e4m3x2_u16((unsigned short)(c4 >> 16)); + const float g0 = __fmul_rn(f16_of(p01, 0), s); + const float g1 = __fmul_rn(f16_of(p01, 1), s); + const float g2 = __fmul_rn(f16_of(p23, 0), s); + const float g3 = __fmul_rn(f16_of(p23, 1), s); + dst[2 * i] = pack_bf16x2(g0, g1); + dst[2 * i + 1] = pack_bf16x2(g2, g3); + } + } else { + const uint4* src = reinterpret_cast(reinterpret_cast(a.cache) + + (long long)rowbase * 2 * LW); + for (int i = 0; i < 8; ++i) { + const uint4 v = src[(t & 7) * 8 + i]; + unsigned* d = dst + 4 * i; + d[0] = v.x; + d[1] = v.y; + d[2] = v.z; + d[3] = v.w; + } + } + __syncthreads(); + // scores: 32 ksteps of m16n8k16, A from the qa u32 view, B from shared + const unsigned* qrowA = qsh + (rowA - hb0) * 256; + const unsigned* qrowB = qrowA + 2048; + float sacc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int kstep = 0; kstep < 32; ++kstep) { + // Triton's dot operand kWidth: bf16 2 (sequential 16-latent window), fp8 4 (even/odd pairs of a 32 window) + const unsigned abase = FP8 ? 16 * (kstep >> 1) + 2 * m + (kstep & 1) : 8 * kstep + m; + const unsigned boff = FP8 ? 8 : 4; + const unsigned av[4] = {qrowA[abase], qrowB[abase], + qrowA[abase + boff], qrowB[abase + boff]}; + const unsigned* krow = kv[kbase + (l >> 2)]; + const unsigned bv[2] = {krow[abase], krow[abase + boff]}; + mma_16n8k16(sacc, av, bv); + } + // scale then the validity mask (where(valid, scores, -inf)) + const int key0 = cbase + tile * 32 + kbase + 2 * m; + bool valid[4]; + if (SPARSE) { + valid[0] = key0 < a.cnt[r]; + valid[1] = key0 + 1 < a.cnt[r]; + } else { + const int lim = a.pos[0] + r; + valid[0] = key0 <= lim; + valid[1] = key0 + 1 <= lim; + } + valid[2] = valid[0]; + valid[3] = valid[1]; + for (int i = 0; i < 4; ++i) { + sacc[i] = __fmul_rn(sacc[i], a.scale); + if (!valid[i]) sacc[i] = neg_inf(); + } + // tile max: in-warp xor 2/1, then the cross-warp exchange (order-free) + float ta = fmax32(sacc[0], sacc[1]); + float tb = fmax32(sacc[2], sacc[3]); + ta = fmax32(ta, __shfl_xor_sync(0xffffffffu, ta, 2)); + ta = fmax32(ta, __shfl_xor_sync(0xffffffffu, ta, 1)); + tb = fmax32(tb, __shfl_xor_sync(0xffffffffu, tb, 2)); + tb = fmax32(tb, __shfl_xor_sync(0xffffffffu, tb, 1)); + if ((l & 3) == 0) { // the row's first lane writes this warp's slice max into the row's slot + tmx[rowA][0][w & 3] = ta; + tmx[rowA][1][w & 3] = tb; + } + __syncthreads(); + const float tile_a = fmax32(fmax32(tmx[rowA][0][0], tmx[rowA][0][1]), fmax32(tmx[rowA][0][2], tmx[rowA][0][3])); + const float tile_b = fmax32(fmax32(tmx[rowA][1][0], tmx[rowA][1][1]), fmax32(tmx[rowA][1][2], tmx[rowA][1][3])); + const bool act_a = tile_a != neg_inf(); + const bool act_b = tile_b != neg_inf(); + const float nma = act_a ? fmax32(m_a, tile_a) : m_a; + const float nmb = act_b ? fmax32(m_b, tile_b) : m_b; + const float alpha_a = act_a ? (m_a == neg_inf() ? 0.0f : ex2(log2e_mul(__fsub_rn(m_a, nma)))) : 1.0f; + const float alpha_b = act_b ? (m_b == neg_inf() ? 0.0f : ex2(log2e_mul(__fsub_rn(m_b, nmb)))) : 1.0f; + const float p0 = (valid[0] && act_a) ? ex2(log2e_mul(__fsub_rn(sacc[0], nma))) : 0.0f; + const float p1 = (valid[1] && act_a) ? ex2(log2e_mul(__fsub_rn(sacc[1], nma))) : 0.0f; + const float p2 = (valid[2] && act_b) ? ex2(log2e_mul(__fsub_rn(sacc[2], nmb))) : 0.0f; + const float p3 = (valid[3] && act_b) ? ex2(log2e_mul(__fsub_rn(sacc[3], nmb))) : 0.0f; + for (int j = 0; j < NT; ++j) { + oacc[j][0] = __fmul_rn(oacc[j][0], alpha_a); + oacc[j][1] = __fmul_rn(oacc[j][1], alpha_a); + oacc[j][2] = __fmul_rn(oacc[j][2], alpha_b); + oacc[j][3] = __fmul_rn(oacc[j][3], alpha_b); + } + psh[rowA][pw + m] = pack_bf16x2(p0, p1); + psh[rowB][pw + m] = pack_bf16x2(p2, p3); + __syncthreads(); + // value mma: A from psh (key pairs pu0/pu1 as in the scores map), B the u16 latents of those keys + for (int j = 0; j < NT; ++j) { + for (int kstep = 0; kstep < 2; ++kstep) { + const int pu0 = FP8 ? 2 * m + kstep : 8 * kstep + m; + const int pu1 = pu0 + (FP8 ? 8 : 4); + const int ncol = on_base + 8 * j + (l >> 2); + const unsigned pav[4] = {psh[rowA][pu0], psh[rowB][pu0], psh[rowA][pu1], psh[rowB][pu1]}; + const unsigned cu = (unsigned)(ncol >> 1); + const unsigned sh = (unsigned)(ncol & 1) << 4; + const unsigned k0 = kv[2 * pu0][cu], k1 = kv[2 * pu0 + 1][cu]; + const unsigned k8 = kv[2 * pu1][cu], k9 = kv[2 * pu1 + 1][cu]; + const unsigned b0 = ((k0 >> sh) & 0xffffu) | (((k1 >> sh) & 0xffffu) << 16); + const unsigned b1 = ((k8 >> sh) & 0xffffu) | (((k9 >> sh) & 0xffffu) << 16); + const unsigned bv[2] = {b0, b1}; + mma_16n8k16(oacc[j], pav, bv); + } + } + // l rowsum: per-lane pair sums, in-warp xor 2/1, cross-warp (P0+P2)+(P1+P3), l = l*alpha + rowsum + const float t1a = __fadd_rn(p0, p1); + const float t1b = __fadd_rn(p2, p3); + float xa = __fadd_rn(t1a, __shfl_xor_sync(0xffffffffu, t1a, 2)); + xa = __fadd_rn(xa, __shfl_xor_sync(0xffffffffu, xa, 1)); + float xb = __fadd_rn(t1b, __shfl_xor_sync(0xffffffffu, t1b, 2)); + xb = __fadd_rn(xb, __shfl_xor_sync(0xffffffffu, xb, 1)); + if ((l & 3) == 0) { // per-row rowsum slots, same exchange discipline as tmx + lbuf[rowA][0][w & 3] = xa; + lbuf[rowA][1][w & 3] = xb; + } + __syncthreads(); + // SASS-confirmed tree: ((P0+P2)+(P1+P3)) over the warp key-slices (P_w = keys 8w..8w+7) + const float rsa = __fadd_rn(__fadd_rn(lbuf[rowA][0][0], lbuf[rowA][0][2]), + __fadd_rn(lbuf[rowA][0][1], lbuf[rowA][0][3])); + const float rsb = __fadd_rn(__fadd_rn(lbuf[rowA][1][0], lbuf[rowA][1][2]), + __fadd_rn(lbuf[rowA][1][1], lbuf[rowA][1][3])); + l_a = __fmaf_rn(l_a, alpha_a, rsa); + l_b = __fmaf_rn(l_b, alpha_b, rsb); + m_a = nma; + m_b = nmb; + } + const long long pbase = ((long long)c * a.R + r) * H; + float* porow = a.po + (pbase + rowA) * LW; + float* porowB = a.po + (pbase + rowB) * LW; + for (int j = 0; j < NT; ++j) { + porow[on_base + 8 * j + 2 * m] = oacc[j][0]; + porow[on_base + 8 * j + 2 * m + 1] = oacc[j][1]; + porowB[on_base + 8 * j + 2 * m] = oacc[j][2]; + porowB[on_base + 8 * j + 2 * m + 1] = oacc[j][3]; + } + if (HBT == 32 || w < 4) { + a.pm[pbase + rowA] = m_a; + a.pl[pbase + rowA] = l_a; + a.pm[pbase + rowB] = m_b; + a.pl[pbase + rowB] = l_b; + } +} + +extern "C" __global__ void __launch_bounds__(256) tf_glm_attn_d16_bf16(ChunksArgs a) { attn_chunks<16, false, false>(a); } +extern "C" __global__ void __launch_bounds__(256) tf_glm_attn_d16_fp8(ChunksArgs a) { attn_chunks<16, true, false>(a); } +extern "C" __global__ void __launch_bounds__(256) tf_glm_attn_d32_bf16(ChunksArgs a) { attn_chunks<32, false, false>(a); } +extern "C" __global__ void __launch_bounds__(256) tf_glm_attn_d32_fp8(ChunksArgs a) { attn_chunks<32, true, false>(a); } +extern "C" __global__ void __launch_bounds__(256) tf_glm_attn_s16_fp8(ChunksArgs a) { attn_chunks<16, true, true>(a); } +extern "C" __global__ void __launch_bounds__(256) tf_glm_attn_s32_fp8(ChunksArgs a) { attn_chunks<32, true, true>(a); } + +struct MergeArgs { + const float* po; // fp32 [nch, R, H, LW] + const float* pm; // fp32 [nch, R, H] + const float* pl; // fp32 [nch, R, H] + bf16* out; // bf16 [R, H, LW] + const int* cnt; // int32 [R] (sparse) + int R; + int nch; + int sparse; +}; + +// Merge one (row, head)'s chunk partials: running max and scale, fma the scaled partials, div.full to bf16. +extern "C" __global__ void tf_glm_merge(MergeArgs a) { + const int r = blockIdx.x; + const int h = blockIdx.y; + const int t = threadIdx.x; + if (a.sparse != 0 && a.cnt[r] == 0) return; + float m = neg_inf(), lsum = 0.0f, o[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int c = 0; c < a.nch; ++c) { + const long long base = ((long long)c * a.R + r) * H + h; + const float cm = a.pm[base]; + const float cl = a.pl[base]; + const bool active = cl > 0.0f; + const float nma = active ? fmax32(m, cm) : m; + const float alpha = active ? (m == neg_inf() ? 0.0f : ex2(log2e_mul(__fsub_rn(m, nma)))) : 1.0f; + const float beta = active ? ex2(log2e_mul(__fsub_rn(cm, nma))) : 0.0f; + const float4 raw = *reinterpret_cast(a.po + base * LW + 4 * t); + const float co[4] = {raw.x, raw.y, raw.z, raw.w}; + for (int i = 0; i < 4; ++i) o[i] = __fmaf_rn(o[i], alpha, __fmul_rn(co[i], beta)); + lsum = __fmaf_rn(lsum, alpha, __fmul_rn(cl, beta)); + m = nma; + } + unsigned ov[2]; + for (int i = 0; i < 2; ++i) ov[i] = pack_bf16x2(div_full(o[2 * i], lsum), div_full(o[2 * i + 1], lsum)); + *reinterpret_cast(reinterpret_cast(a.out) + ((long long)r * H + h) * LW + 4 * t) = + make_uint2(ov[0], ov[1]); +} + +} // namespace tf_glm_latent diff --git a/zig/kernels/cuda/glm_sparse.cu b/zig/kernels/cuda/glm_sparse.cu new file mode 100644 index 0000000000..7e0a0f1419 --- /dev/null +++ b/zig/kernels/cuda/glm_sparse.cu @@ -0,0 +1,450 @@ +// GLM-5.3-Flash's DSA sparse indexer (sparse.py): bit-exact via inline PTX ops, m16n8k16 dots, PTX sum orders. + +#include +#include + +namespace tf_glm_sparse { + +using bf16 = __nv_bfloat16; + +struct IndexArgs { + const bf16* kr; // bf16 [R, D] row-major (k_stride elements a row) + long long k_stride; + const float* gr; // fp32 [R, D] + const unsigned short* lnw; // fp16 [D] (the checkpoint's k_norm rows) + const unsigned short* lnb; // fp16 [D] + bf16* ik; // bf16 [cap, D] + bf16* ig; // bf16 [cap, D] + int pos; + float eps; + const int* posp; // the position on the device (a captured decode step), else null and pos +}; + +struct PoolArgs { + const bf16* ik; // bf16 [cap, D] + const bf16* ig; // bf16 [cap, D] + const bf16* ape; // bf16 [4, D] + bf16* pk; // bf16 [cap / 4 + 2, D] + int pos; + int r; + const int* posp; // as IndexArgs +}; + +struct ScoresArgs { + const bf16* qi; // bf16 [R, H * D] row-major (the block's rows) + long long qi_stride; // qi's row stride in elements (160: D + H) + const bf16* w; // bf16 [R, H] row-major (w_stride elements a row) + long long w_stride; + const bf16* pk; // bf16 [np, D] + float* out; // fp32 [R, np] + int pos; + int np; + float scale; + float wscale; + const int* posp; // as IndexArgs +}; + +struct SelectArgs { + const float* s; // fp32 [R, np] + long long s_stride; // np + long long out_stride; // K + long long* out; // int64 pools [R, K] + int np; + int pos; + int k; + int visible; // 1: limit = min(np, max(k, (pos + r + 1) / 4)); 0: limit = np + const int* posp; // as IndexArgs +}; + +struct TokensArgs { + const long long* pools; // int64 [R, K] + int* tokens; // int32 [R, 512 * 4 + 3] + int* counts; // int32 [R] + int pos; + const int* posp; // as IndexArgs +}; + +constexpr int D = 128; // index key dims +constexpr int H = 32; // index heads (the serve model's index_heads) +constexpr int BP = 64; // pools a scores block +constexpr int L = 1; // pool blocks a scores program takes (speed only: a block's sums are its own) +constexpr int LOG2E_BITS = 0x3FB8AA3B; // float log2(e)'s bits, as Triton's exp lowering multiplies by + +__device__ __forceinline__ unsigned pack2(bf16 lo, bf16 hi) { + return static_cast(__bfloat16_as_ushort(lo)) | (static_cast(__bfloat16_as_ushort(hi)) << 16); +} + +__device__ __forceinline__ float rcp_approx(float x) { + float r; + asm("rcp.approx.f32 %0, %1;\n" : "=f"(r) : "f"(x)); + return r; +} + +// Triton's exp: ex2.approx(x * log2e), but past -126 the input is halved and the result squared, keeping subnormals. +__device__ __forceinline__ float exp_approx(float x) { + float t = __fmul_rn(x, __int_as_float(LOG2E_BITS)); + const bool sub = t < -126.0f; + if (sub) t = __fmul_rn(t, 0.5f); + float r; + asm("ex2.approx.f32 %0, %1;\n" : "=f"(r) : "f"(t)); + if (sub) r = __fmul_rn(r, r); + return r; +} + +__device__ __forceinline__ float sqrt_approx_ftz(float x) { + float r; + asm("sqrt.approx.ftz.f32 %0, %1;\n" : "=f"(r) : "f"(x)); + return r; +} + +__device__ __forceinline__ void mma_16n8k16(float (&d)[4], const unsigned (&a)[4], const unsigned (&b)[2]) { + 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};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// A float32 as a uint32 whose unsigned order is the score's order (-0 counted as +0), as sparse.py's _order_key. +__device__ __forceinline__ unsigned order_key(float s) { + const int bits = __float_as_int(__fadd_rn(s, 0.0f)); + return static_cast(bits ^ ((bits >> 31) | static_cast(0x80000000))); +} + +// The xor butterfly 16, 8, 4, 2, 1 over a warp's lanes (Triton's shfl.bfly + add stages). +__device__ __forceinline__ float warp_bfly(float v) { + const unsigned all = 0xffffffffu; + v = __fadd_rn(v, __shfl_xor_sync(all, v, 16)); + v = __fadd_rn(v, __shfl_xor_sync(all, v, 8)); + v = __fadd_rn(v, __shfl_xor_sync(all, v, 4)); + v = __fadd_rn(v, __shfl_xor_sync(all, v, 2)); + v = __fadd_rn(v, __shfl_xor_sync(all, v, 1)); + return v; +} + +// One warp: the lane's four elements ((x0 + x1) + (x2 + x3)), then the full butterfly. +__device__ __forceinline__ float warp_sum128(float x0, float x1, float x2, float x3) { + return warp_bfly(__fadd_rn(__fadd_rn(x0, x1), __fadd_rn(x2, x3))); // Triton's f32x2 pairing +} + +// Row r: LayerNorm(kr) -> bf16 into ik[pos + r]; the gate row -> bf16 into ig[pos + r]. One warp per row. +extern "C" __global__ void tf_glm_index_write(IndexArgs a) { + const int r = blockIdx.x; + const int lane = threadIdx.x; + const bf16* x_in = a.kr + static_cast(r) * a.k_stride + lane * 4; + float x[4]; +#pragma unroll + for (int i = 0; i < 4; i++) x[i] = __bfloat162float(x_in[i]); + const float mean = __fmul_rn(warp_sum128(x[0], x[1], x[2], x[3]), 0.0078125f); // Triton's FMUL by 1/128, not div + float xc[4]; +#pragma unroll + for (int i = 0; i < 4; i++) xc[i] = __fsub_rn(x[i], mean); + // Triton's variance sums (x0^2 + x2^2) + (x1^2 + x3^2), each first square fused into the FFMA + const float s02 = __fmaf_rn(xc[0], xc[0], __fmul_rn(xc[2], xc[2])); + const float s13 = __fmaf_rn(xc[1], xc[1], __fmul_rn(xc[3], xc[3])); + const float var_eps = __fmaf_rn(warp_bfly(__fadd_rn(s02, s13)), 0.0078125f, a.eps); + const float d = sqrt_approx_ftz(var_eps); + const unsigned short* w = a.lnw + lane * 4; + const unsigned short* b = a.lnb + lane * 4; + const float wf[4] = {__half2float(*reinterpret_cast(w)), __half2float(*reinterpret_cast(w + 1)), + __half2float(*reinterpret_cast(w + 2)), __half2float(*reinterpret_cast(w + 3))}; + const float bf4[4] = {__half2float(*reinterpret_cast(b)), __half2float(*reinterpret_cast(b + 1)), + __half2float(*reinterpret_cast(b + 2)), __half2float(*reinterpret_cast(b + 3))}; + unsigned y[2], gv[2]; + float rcp_d; + asm("rcp.approx.f32 %0, %1;" : "=f"(rcp_d) : "f"(d)); // Triton's div.full lowers to RCP + FMUL (no refinement) +#pragma unroll + for (int i = 0; i < 4; i++) { + const float t = __fmul_rn(xc[i], rcp_d); + reinterpret_cast(y)[i] = __float2bfloat16_rn(__fmaf_rn(t, wf[i], bf4[i])); + reinterpret_cast(gv)[i] = __float2bfloat16_rn(a.gr[static_cast(r) * D + lane * 4 + i]); + } + const long long p = static_cast(a.posp ? *a.posp : a.pos) + r; + *reinterpret_cast(a.ik + p * D + lane * 4) = y[0]; + *reinterpret_cast(a.ik + p * D + lane * 4 + 2) = y[1]; + *reinterpret_cast(a.ig + p * D + lane * 4) = gv[0]; + *reinterpret_cast(a.ig + p * D + lane * 4 + 2) = gv[1]; +} + +// Program i: pool p = pos / 4 + i if it is complete within the window. One warp, four dims a lane. +extern "C" __global__ void tf_glm_pool_keys(PoolArgs a) { + const int i = blockIdx.x; + const long long pos = a.posp ? *a.posp : a.pos; + const long long p = pos / 4 + i; + if (4 * p + 3 > pos + a.r - 1) return; + const int lane = threadIdx.x; + const int d = lane * 4; + float l[4][4], e[4][4]; // pool member j, lane dim t +#pragma unroll + for (int j = 0; j < 4; j++) { + const bf16* g = a.ig + (4 * p + j) * D + d; + const bf16* ap = a.ape + j * D + d; +#pragma unroll + for (int t = 0; t < 4; t++) + l[j][t] = __fadd_rn(__bfloat162float(g[t]), __bfloat162float(ap[t])); + } +#pragma unroll + for (int t = 0; t < 4; t++) { + const float m = fmaxf(fmaxf(l[0][t], l[1][t]), fmaxf(l[2][t], l[3][t])); // exact: any order +#pragma unroll + for (int j = 0; j < 4; j++) e[j][t] = exp_approx(__fsub_rn(l[j][t], m)); + } + bf16 prod[4][4]; // pool member j, lane dim t +#pragma unroll + for (int t = 0; t < 4; t++) { + // Triton's s = ((e0 + e1) + e2) + e3 and pk = ((t0 + t1) + t2) + t3: strictly sequential, not a balanced tree + float s = e[0][t]; + s = __fadd_rn(s, e[1][t]); + s = __fadd_rn(s, e[2][t]); + s = __fadd_rn(s, e[3][t]); +#pragma unroll + for (int j = 0; j < 4; j++) { + // Triton's div.full lowers to RCP + FMUL (the divisor is a sum of exps in [1, 4]: no denormal fixups) + const float q = __bfloat162float(__float2bfloat16_rn(__fmul_rn(e[j][t], rcp_approx(s)))); + const bf16* k = a.ik + (4 * p + j) * D + d; + prod[j][t] = __float2bfloat16_rn(__fmul_rn(q, __bfloat162float(k[t]))); + } + } +#pragma unroll + for (int t = 0; t < 4; t++) { + float sum = __bfloat162float(prod[0][t]); + sum = __fadd_rn(sum, __bfloat162float(prod[1][t])); + sum = __fadd_rn(sum, __bfloat162float(prod[2][t])); + sum = __fadd_rn(sum, __bfloat162float(prod[3][t])); + reinterpret_cast(a.pk + p * D + d)[t] = __float2bfloat16_rn(sum); + } +} + +// Scores: s_p = sum_h w_h relu(scale*qi_h.pool_p), -inf past visible; warp w owns cols 16w..16w+15, K ascending chunks. +extern "C" __global__ void __launch_bounds__(128) tf_glm_scores(ScoresArgs a) { + __shared__ bf16 q_sh[H][D]; + __shared__ float sc_sh[BP]; + __shared__ float w_sh[H]; + const int r = blockIdx.x; + const int g = blockIdx.y; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + for (int i = tid; i < H * D; i += 128) { // q row r, all 32 heads (H == 32 == HP, no padding rows) + const int hh = i / D, dd = i % D; + q_sh[hh][dd] = a.qi[static_cast(r) * a.qi_stride + hh * D + dd]; + } + if (tid < H) w_sh[tid] = __bfloat162float(a.w[static_cast(r) * a.w_stride + tid]) * a.wscale; + __syncthreads(); + + const int g4 = lane >> 2; // the fragment's first row + const int t2 = (lane & 3) * 2; // the fragment's first col + const int npool = ((a.posp ? *a.posp : a.pos) + r + 1) / 4; + const float neg_inf = -__int_as_float(0x7F800000); + for (int j = 0; j < L; j++) { + const int pb = g * L + j; + const int p0 = pb * BP; + if (p0 >= a.np) break; + if (p0 >= npool) { // whole block past the row's last visible pool: -inf without dots + for (int e = tid; e < BP && p0 + e < a.np; e += 128) + a.out[static_cast(r) * a.np + p0 + e] = neg_inf; + continue; + } + // dots: 8 ascending K-chunks, 2 m-tiles, 2 n-tiles per warp + float c[2][2][4]; +#pragma unroll + for (int mt = 0; mt < 2; mt++) +#pragma unroll + for (int nt = 0; nt < 2; nt++) +#pragma unroll + for (int i = 0; i < 4; i++) c[mt][nt][i] = 0.0f; +#pragma unroll + for (int ks = 0; ks < 8; ks++) { +#pragma unroll + for (int mt = 0; mt < 2; mt++) { + unsigned av[4]; + const int row0 = g4 + 16 * mt; + const int col0 = 16 * ks + t2; + av[0] = *reinterpret_cast(&q_sh[row0][col0]); + av[1] = *reinterpret_cast(&q_sh[row0 + 8][col0]); + av[2] = *reinterpret_cast(&q_sh[row0][col0 + 8]); + av[3] = *reinterpret_cast(&q_sh[row0 + 8][col0 + 8]); +#pragma unroll + for (int nt = 0; nt < 2; nt++) { + const int pool = p0 + 16 * warp + (lane >> 2) + 8 * nt; + unsigned bv[2]; + if (pool < npool) { + const bf16* base = a.pk + static_cast(pool) * D + 16 * ks + (lane & 3) * 2; + bv[0] = *reinterpret_cast(base); + bv[1] = *reinterpret_cast(base + 8); + } else { + bv[0] = 0u; + bv[1] = 0u; + } + mma_16n8k16(c[mt][nt], av, bv); + } + } + } + // per col: fma(w_g, relu_g, w_g+8*relu_g+8) + fma(w_g+16, relu_g+16, w_g+24*relu_g+24); ptxas contracts it + float v[4]; // col slots 2t2, 2t2+1, 2t2+8, 2t2+9 within the warp's 16 cols +#pragma unroll + for (int nt = 0; nt < 2; nt++) +#pragma unroll + for (int half = 0; half < 2; half++) { + const float r0 = fmaxf(__fmul_rn(c[0][nt][half], a.scale), 0.0f); + const float r1 = fmaxf(__fmul_rn(c[0][nt][2 + half], a.scale), 0.0f); + const float r2 = fmaxf(__fmul_rn(c[1][nt][half], a.scale), 0.0f); + const float r3 = fmaxf(__fmul_rn(c[1][nt][2 + half], a.scale), 0.0f); + const float p1 = __fmaf_rn(w_sh[g4], r0, __fmul_rn(w_sh[g4 + 8], r1)); + const float p2 = __fmaf_rn(w_sh[g4 + 16], r2, __fmul_rn(w_sh[g4 + 24], r3)); + v[2 * nt + half] = __fadd_rn(p1, p2); + } + // butterfly: lane l's col value gains lanes l^16, l^8, l^4 (same col slot) in that order + const unsigned all = 0xffffffffu; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = __fadd_rn(v[i], __shfl_xor_sync(all, v[i], 16)); + v[i] = __fadd_rn(v[i], __shfl_xor_sync(all, v[i], 8)); + v[i] = __fadd_rn(v[i], __shfl_xor_sync(all, v[i], 4)); + } + // lanes 0..3 hold complete sums for cols 16w + 2s + {0,1,8,9}: stage them, then thread e stores col e + if (lane < 4) { + sc_sh[16 * warp + 2 * lane] = v[0]; + sc_sh[16 * warp + 2 * lane + 1] = v[1]; + sc_sh[16 * warp + 2 * lane + 8] = v[2]; + sc_sh[16 * warp + 2 * lane + 9] = v[3]; + } + __syncthreads(); + const int e0 = p0 + tid; + if (tid < BP && e0 < a.np) + a.out[static_cast(r) * a.np + e0] = e0 < npool ? sc_sh[tid] : neg_inf; + __syncthreads(); + } +} + +// Block-wide exclusive prefix sum of the threads' counts, in thread order (each thread owns 8 consecutive elements). +__device__ __forceinline__ int block_prefix(int v, int* warp_totals, int* total) { + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + int p = v; +#pragma unroll + for (int off = 1; off < 32; off <<= 1) { + const int u = __shfl_up_sync(0xffffffffu, p, off); + if (lane >= off) p += u; + } + if (lane == 31) warp_totals[warp] = p; + __syncthreads(); + if (threadIdx.x == 0) { + int run = 0; + for (int w = 0; w < 4; w++) { + const int t = warp_totals[w]; + warp_totals[w] = run; + run += t; + } + *total = run; + } + __syncthreads(); + return warp_totals[warp] + (lane == 0 ? 0 : p - v); +} + +// Select: a radix select (8 bits/pass) finds the K-th best key, then one pass writes pools above it plus lowest ties. +extern "C" __global__ void __launch_bounds__(128) tf_glm_select_rows(SelectArgs a) { + __shared__ int hist[256]; + __shared__ int warp_totals[4]; + __shared__ int total; + const int r = blockIdx.x; + const long long row = static_cast(r) * a.s_stride; + const int limit = a.visible ? min(a.np, max(a.k, ((a.posp ? *a.posp : a.pos) + r + 1) / 4)) : a.np; + unsigned prefix = 0, fixed = 0; + int need = a.k; +#pragma unroll 1 + for (int pass = 0; pass < 4; pass++) { + const int shift = 24 - 8 * pass; + for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0; + __syncthreads(); + for (int c = 0; c < limit; c += 1024) { + const int i = c + threadIdx.x * 8; + for (int t = 0; t < 8; t++) { + const int idx = i + t; + if (idx < limit) { + const unsigned u = order_key(a.s[row + idx]); + if ((u & fixed) == prefix) atomicAdd(&hist[(u >> shift) & 0xFF], 1); + } + } + } + __syncthreads(); + int total_hist = 0; + for (int d = 0; d < 256; d++) total_hist += hist[d]; + int excl = 0, digit = 0, above = 0; + for (int d = 0; d < 256; d++) { + const int h = hist[d]; + if (total_hist - excl >= need) { + digit = d; + above = total_hist - excl - h; // bins strictly above digit + } + excl += h; + } + need -= above; + prefix |= static_cast(digit) << shift; + fixed |= 0xFFu << shift; + __syncthreads(); + } + // write pass: pools above the K-th key, then the lowest ties, ascending; a block-exclusive scan places them + int written = 0, equal_seen = 0; + for (int c = 0; c < limit; c += 1024) { + const int i = c + threadIdx.x * 8; + int eq8 = 0; + for (int t = 0; t < 8; t++) { + const int idx = i + t; + if (idx < limit && order_key(a.s[row + idx]) == prefix) eq8++; + } + const int eq_prefix = block_prefix(eq8, warp_totals, &total); + const int eq_tot = total; + __syncthreads(); + // every element above the K-th key is taken plus lowest ties by global tie rank; positions count all takes + int take8 = 0, eq_before = 0; + for (int t = 0; t < 8; t++) { + const int idx = i + t; + if (idx < limit) { + const unsigned u = order_key(a.s[row + idx]); + if (u > prefix || (u == prefix && equal_seen + eq_prefix + eq_before < need)) take8++; + if (u == prefix) eq_before++; + } + } + const int take_prefix = block_prefix(take8, warp_totals, &total); + const int take_tot = total; + __syncthreads(); + int run = written + take_prefix; + eq_before = 0; + for (int t = 0; t < 8; t++) { + const int idx = i + t; + if (idx < limit) { + const unsigned u = order_key(a.s[row + idx]); + if (u > prefix || (u == prefix && equal_seen + eq_prefix + eq_before < need)) { + a.out[static_cast(r) * a.out_stride + run] = idx; + run++; + } + if (u == prefix) eq_before++; + } + } + __syncthreads(); + written += take_tot; + equal_seen += eq_tot; + } +} + +// Tokens/counts: each pool's four tokens ascending, then the incomplete last pool's; counts past the dense limit. +extern "C" __global__ void tf_glm_tokens_counts(TokensArgs a) { + const int r = blockIdx.x; + const int t = threadIdx.x; // 512 threads, four tokens a thread + const long long q = static_cast(a.posp ? *a.posp : a.pos) + r; + const long long npool = (q + 1) / 4; + const long long p = a.pools[static_cast(r) * 512 + t]; // every row's 512 selected pools' tokens + for (int j = 0; j < 4; j++) + a.tokens[static_cast(r) * 2051 + t * 4 + j] = static_cast(p * 4 + j); + if (t < 3) { + const long long tail = npool * 4 + t; + a.tokens[static_cast(r) * 2051 + 2048 + t] = tail <= q ? static_cast(tail) : -1; + } + if (t == 0) { + int tail_ok = 0; + for (int j = 0; j < 3; j++) + if (npool * 4 + j <= q) tail_ok++; + a.counts[r] = npool > 512 ? 2048 + tail_ok : 0; + } +} +} // namespace tf_glm_sparse diff --git a/zig/src/core/exl3_load.zig b/zig/src/core/exl3_load.zig new file mode 100644 index 0000000000..1bf23fa2d8 --- /dev/null +++ b/zig/src/core/exl3_load.zig @@ -0,0 +1,227 @@ +//! EXL3 checkpoint groups: every `{prefix}.trellis` and its parts, admitted from the safetensors headers (format.py). +const std = @import("std"); +const checkpoint = @import("checkpoint.zig"); +const st = @import("safetensors.zig"); + +pub const Codebook = enum { inst3, mcg, mul1 }; + +/// format.py's MARKERS: a scalar int32 tensor per codebook, whose value must be the magic. +pub const Marker = struct { name: []const u8, codebook: Codebook, value: u32 }; +pub const markers = [_]Marker{ + .{ .name = "mcg", .codebook = .mcg, .value = 0xCBAC1FED }, + .{ .name = "mul1", .codebook = .mul1, .value = 0x83DCD12D }, +}; + +/// One admitted group: a trellis int16 [K/16, N/16, 16 * bits], its scales, and the optional bias. +pub const Group = struct { + prefix: []const u8, + halves: u16, // 2 * bits; odd halves are the mul1 half steps + codebook: Codebook, + k: usize, + n: usize, + suh: bool, // fp16 scales; false: the packed sign words su and sv + bias: bool, +}; + +/// Counts of the admitted groups, by halves and for the routed experts' down projections. +pub const Summary = struct { + groups: usize = 0, + inst3: usize = 0, + mcg: usize = 0, + mul1: usize = 0, + suh_svh: usize = 0, + su_sv: usize = 0, + bias: usize = 0, + bits: [17]usize = @splat(0), // by halves, 2 to 16 + experts_down: [17]usize = @splat(0), // *.experts.*.down_proj by halves + + fn bump(s: *Summary, g: Group) void { + s.groups += 1; + switch (g.codebook) { + .inst3 => s.inst3 += 1, + .mcg => s.mcg += 1, + .mul1 => s.mul1 += 1, + } + if (g.suh) s.suh_svh += 1 else s.su_sv += 1; + if (g.bias) s.bias += 1; + s.bits[g.halves] += 1; + if (std.mem.endsWith(u8, g.prefix, ".down_proj") and std.mem.indexOf(u8, g.prefix, ".experts.") != null) s.experts_down[g.halves] += 1; + } +}; + +/// One group's tensors: the trellis words, its scales, and the optional bias. +pub const Parts = struct { + trellis: st.Tensor, + suh: st.Tensor, + svh: st.Tensor, + bias: ?st.Tensor = null, + suh_packed: bool = false, // su/sv sign words instead of fp16 scales +}; + +pub const Scan = struct { + gpa: std.mem.Allocator, + ck: checkpoint.Checkpoint, + groups: []Group, + summary: Summary = .{}, + bad: usize = 0, // trellis tensors refused: another format's tensors, a broken or half-written group + + pub fn deinit(s: *Scan) void { + for (s.groups) |g| s.gpa.free(g.prefix); + s.gpa.free(s.groups); + s.ck.close(); + } + + /// The group a tensor prefix names, if any. + pub fn find(s: *const Scan, prefix: []const u8) ?usize { + for (s.groups, 0..) |g, i| if (std.mem.eql(u8, g.prefix, prefix)) return i; + return null; + } + + /// One group's tensors, their layouts checked on the way out. + pub fn parts(s: *Scan, a: std.mem.Allocator, gi: usize) !Parts { + const g = s.groups[gi]; + const trellis = try s.ck.get(try std.mem.concat(a, u8, &.{ g.prefix, ".trellis" })); + if (!trellis.is(.i16, &.{ g.k / 16, g.n / 16, @as(usize, g.halves) * 8 })) return error.BadTrellis; + const su_name = if (g.suh) "suh" else "su"; + const sv_name = if (g.suh) "svh" else "sv"; + const suh = try s.ck.get(try std.mem.concat(a, u8, &.{ g.prefix, ".", su_name })); + const svh = try s.ck.get(try std.mem.concat(a, u8, &.{ g.prefix, ".", sv_name })); + const in_shape: []const usize = if (g.suh) &.{g.k} else &.{g.k / 16}; + const out_shape: []const usize = if (g.suh) &.{g.n} else &.{g.n / 16}; + if (!suh.is(.f16, in_shape) or !svh.is(.f16, out_shape)) return error.BadScales; + var bias: ?st.Tensor = null; + if (g.bias) { + bias = try s.ck.get(try std.mem.concat(a, u8, &.{ g.prefix, ".bias" })); + if (!bias.?.is(.f16, &.{g.n})) return error.BadBias; + } + return .{ .trellis = trellis, .suh = suh, .svh = svh, .bias = bias, .suh_packed = !g.suh }; + } +}; + +/// Every trellis tensor of `dir`'s shards, admitted or counted as refused. +pub fn scan(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Scan { + var ck = try checkpoint.Checkpoint.openModel(gpa, io, dir); + errdefer ck.close(); + var groups: std.ArrayList(Group) = .empty; + errdefer groups.deinit(gpa); + var summary: Summary = .{}; + var bad: usize = 0; + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + for (ck.files.items) |*f| for (f.names.keys(), f.names.values()) |name, e| { + if (!std.mem.endsWith(u8, name, ".trellis")) continue; + const prefix = name[0 .. name.len - ".trellis".len]; + if (admit(a, &ck, prefix, e)) |g| { + try groups.append(gpa, .{ .prefix = try gpa.dupe(u8, g.prefix), .halves = g.halves, .codebook = g.codebook, .k = g.k, .n = g.n, .suh = g.suh, .bias = g.bias }); + summary.bump(g); + } else |_| { + bad += 1; + } + }; + return .{ .gpa = gpa, .ck = ck, .groups = try groups.toOwnedSlice(gpa), .summary = summary, .bad = bad }; +} + +/// parse_group's admission: the trellis layout, the codebook marker's value, and the group's parts. +fn admit(a: std.mem.Allocator, ck: *checkpoint.Checkpoint, prefix: []const u8, e: st.Entry) !Group { + if (e.dtype != .i16 or e.rank != 3) return error.Refuse; + const halves = e.shape[2] / 8; // the trellis's last dim is 16 * bits + if (e.shape[2] % 8 != 0 or halves < 2 or halves > 16) return error.Refuse; + const k = e.shape[0] * 16; + const n = e.shape[1] * 16; + if (k % 128 != 0 or n % 128 != 0) return error.Refuse; + var codebook: Codebook = .inst3; + var marker: ?Marker = null; + for (markers) |m| if (ck.has(try std.mem.concat(a, u8, &.{ prefix, ".", m.name }))) { + if (marker != null) return error.Refuse; // both codebooks' markers: not a loadable group + marker = m; + codebook = m.codebook; + }; + if (halves % 2 == 1 and codebook != .mul1) return error.Refuse; // odd halves are the mul1 half steps + if (marker) |m| { + const t = try ck.get(try std.mem.concat(a, u8, &.{ prefix, ".", m.name })); + if (t.bytes.len != 4 or std.mem.readInt(u32, t.bytes[0..4], .little) != m.value) return error.Refuse; + } + const su = try scale(a, ck, prefix, "suh", k, "su", k / 16); + const sv = try scale(a, ck, prefix, "svh", n, "sv", n / 16); + if (su == null or sv == null or su.? != sv.?) return error.Refuse; // one fp16, one packed: a half-written group + const bias = try scale(a, ck, prefix, "bias", n, "", 0); + return .{ .prefix = prefix, .halves = @intCast(halves), .codebook = codebook, .k = k, .n = n, .suh = su.? == .f16, .bias = bias != null }; +} + +const Scales = enum { f16, words }; // fp16 per row, or the packed sign words + +/// One part's admission: the fp16 name, or the packed-words fallback; a missing input scale (or fallback) is refused. +fn scale(a: std.mem.Allocator, ck: *checkpoint.Checkpoint, prefix: []const u8, name: []const u8, len: usize, packed_name: []const u8, packed_len: usize) !?Scales { + const full = try std.mem.concat(a, u8, &.{ prefix, ".", name }); + if (!ck.has(full)) { + if (packed_name.len == 0) return null; // an optional part + const packed_full = try std.mem.concat(a, u8, &.{ prefix, ".", packed_name }); + if (!ck.has(packed_full)) return error.Refuse; + const p = try ck.get(packed_full); + if (p.dtype != .i16 or p.numel() != packed_len) return error.Refuse; + return .words; + } + const t = try ck.get(full); + if (t.dtype != .f16 or t.numel() != len) return error.Refuse; + return .f16; +} + +test "a hand-built safetensors file admits its one group and refuses the scaleless one" { + const io = std.testing.io; + const gpa = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + // one full group (trellis [8, 8, 32] i16, fp16 scales, the mul1 marker) and one without scales + const trellis: [4096]u8 = @splat(0); + const scales: [256]u8 = @splat(0); + const marker = [_]u8{ 0x2D, 0xD1, 0xDC, 0x83 }; // 0x83DCD12D little-endian + const Entry = st.Entry; + const entries = [_]Entry{ + .{ .dtype = .i16, .rank = 3, .shape = .{ 8, 8, 32, 1, 1 }, .begin = 0, .end = 4096 }, + .{ .dtype = .f16, .rank = 1, .shape = .{ 128, 1, 1, 1, 1 }, .begin = 4096, .end = 4352 }, + .{ .dtype = .f16, .rank = 1, .shape = .{ 128, 1, 1, 1, 1 }, .begin = 4352, .end = 4608 }, + .{ .dtype = .i32, .rank = 1, .shape = .{ 1, 1, 1, 1, 1 }, .begin = 4608, .end = 4612 }, + .{ .dtype = .i16, .rank = 3, .shape = .{ 8, 8, 32, 1, 1 }, .begin = 4612, .end = 4612 + 4096 }, + }; + const names = [_][]const u8{ "a.trellis", "a.suh", "a.svh", "a.mul1", "b.trellis" }; + const dts = [_][]const u8{ "I16", "F16", "F16", "I32", "I16" }; + const data = try std.mem.concat(gpa, u8, &.{ &trellis, &scales, &scales, &marker, &trellis }); + defer gpa.free(data); + var header: std.ArrayListUnmanaged(u8) = .empty; + defer header.deinit(gpa); + try header.appendSlice(gpa, "{"); + for (names, entries, dts, 0..) |name, e, dt, i| { + if (i != 0) try header.appendSlice(gpa, ","); + try header.print(gpa, "\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ name, dt }); + for (e.shape[0..e.rank], 0..) |d, j| { + if (j != 0) try header.appendSlice(gpa, ","); + try header.print(gpa, "{d}", .{d}); + } + try header.print(gpa, "],\"data_offsets\":[{d},{d}]}}", .{ e.begin, e.end }); + } + try header.appendSlice(gpa, "}"); + var file: std.ArrayListUnmanaged(u8) = .empty; + defer file.deinit(gpa); + try file.appendNTimes(gpa, 0, 8); + std.mem.writeInt(u64, file.items[0..8], header.items.len, .little); + try file.appendSlice(gpa, header.items); + try file.appendSlice(gpa, data); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = file.items }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + + var s = try scan(gpa, io, dir); + defer s.deinit(); + try std.testing.expectEqual(@as(usize, 1), s.bad); + try std.testing.expectEqual(@as(usize, 1), s.summary.mul1); + try std.testing.expectEqual(@as(usize, 1), s.summary.suh_svh); + try std.testing.expectEqual(@as(usize, 1), s.summary.bits[4]); // a 32-word trellis is 4 halves + const gi = s.find("a") orelse return error.TestUnexpectedResult; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const p = try s.parts(arena.allocator(), gi); + try std.testing.expectEqual(@as(usize, 128), p.trellis.dim(0) * 16); + try std.testing.expect(!p.suh_packed); + try std.testing.expectEqual(@as(usize, 128), p.svh.numel()); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 736248a93d..a901666ffd 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -4,6 +4,7 @@ pub const safetensors = @import("safetensors.zig"); pub const exl3_format = @import("exl3_format.zig"); pub const exlr8 = @import("exlr8.zig"); pub const checkpoint = @import("checkpoint.zig"); +pub const exl3_load = @import("exl3_load.zig"); pub const direct_io = @import("direct_io.zig"); pub const Checkpoint = checkpoint.Checkpoint; pub const draft_depth = @import("draft_depth.zig"); @@ -31,6 +32,7 @@ test { _ = act_ref; _ = exl3_format; _ = exlr8; + _ = exl3_load; _ = @import("exl3_rect_test.zig"); _ = checkpoint; _ = direct_io; diff --git a/zig/src/cuda/exl3.zig b/zig/src/cuda/exl3.zig new file mode 100644 index 0000000000..5da96fac86 --- /dev/null +++ b/zig/src/cuda/exl3.zig @@ -0,0 +1,349 @@ +//! The EXL3 decode-row linear (exl3/linear.cu's device code): the plan, a layer's tensors, and the two kernels. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const memory = @import("memory.zig"); +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// x's and y's element types (linear.cu DType). +pub const DType = enum(i32) { f16 = 0, bf16 = 1, f32 = 2 }; + +/// The three codebooks (exl3_format.zig's, numbered as the kernel's). +pub const Codebook = enum(u3) { inst3 = 0, mcg = 1, mul1 = 2 }; + +/// One compiled kernel: 2 * bits halves, a codebook, and the warps one block runs. +pub const Variant = struct { halves: u32, codebook: Codebook, warps: u32 }; + +/// Every instantiation of exl3_linear.cu: every width (whole bits all codebooks, the mul1 half steps), 4 and 8 warps. +pub const variants = blk: { + var vs: [60]Variant = undefined; + var i: usize = 0; + for (0..2) |w| { + const warps: u32 = if (w == 0) 4 else 8; + for ([_]u32{ 2, 1, 0 }) |cbi| { + const cb: Codebook = @fromBackingInt(@intCast(cbi)); + var halves: u32 = 2; + while (halves <= 16) : (halves += 2) { + vs[i] = .{ .halves = halves, .codebook = cb, .warps = warps }; + i += 1; + } + if (cb != .mul1) continue; // the half steps exist only with mul1 + halves = 3; + while (halves <= 13) : (halves += 2) { // the mul1 half steps, 1.5 to 6.5 bits (format.py's BITS) + vs[i] = .{ .halves = halves, .codebook = cb, .warps = warps }; + i += 1; + } + } + } + break :blk vs; +}; + +/// The kernel's extern "C" name for an instantiation (linear.cu's TF_EXL3_LINEAR macro). +pub fn symbol(comptime halves: u32, comptime codebook: Codebook, comptime warps: u32) [:0]const u8 { + return std.fmt.comptimePrint("tf_exl3_linear_{d}_{d}_{d}", .{ halves, @backingInt(codebook), warps }); +} + +pub const rot_in_symbol = "tf_exl3_rot_in"; + +/// The prompt kernels' widths (mul1, 4 and 8 warps each): every half step but 15. +pub const prompt_halves = [_]u32{ 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 16 }; + +/// A prompt kernel's shared bytes: psum [64][132] fp32 and two stages of x [2][64][24] fp16. +pub const prompt_shared = 64 * 132 * 4 + 2 * 2 * 64 * 24 * 2; + +fn promptIndex(halves: u32, warps: u32) ?usize { + for (prompt_halves, 0..) |h, i| if (h == halves) return 2 * i + @intFromBool(warps == 8); + return null; +} + +pub fn variantIndex(halves: u32, codebook: Codebook, warps: u32) ?usize { + for (variants, 0..) |v, i| { + if (v.halves == halves and v.codebook == codebook and v.warps == warps) return i; + } + return null; +} + +/// What one shape launches, from the shape alone so a row's bits never depend on the row count (linear.py plan). +pub fn plan(k: usize, n: usize, blocks: usize, min_tiles: usize) struct { sk: u32, wk: u32 } { + const kt = k / 16; + const nb = n / 128; + var sk: usize = 1; + var wk: usize = if (nb >= 64) 8 else 4; + if (nb < 64) { + while (nb * sk < blocks and sk < 64) sk *= 2; + } + while ((wk > 4 or sk > 1) and (kt % (sk * wk) != 0 or kt / (sk * wk) < min_tiles)) { + if (wk > 4) wk = 4 else if (sk > 1) sk /= 2 else break; + } + return .{ .sk = @intCast(sk), .wk = @intCast(wk) }; +} + +/// Words between k tiles and 128-column blocks of one layout: strips [N/128, K/16, 8, W] or stored [K/16, N/16, W]. +pub fn strides(k: usize, n: usize, halves: usize, strips: bool) struct { k: usize, nb: usize } { + const words = halves * 4; + return if (strips) + .{ .k = 8 * words, .nb = (k / 16) * 8 * words } + else + .{ .k = (n / 16) * words, .nb = 8 * words }; +} + +/// One layer's tensors on the device: the trellis words, the Hadamard scales, the optional bias, shape and layout. +pub const Weight = struct { + trellis: u64, // u32 words, `stride_k` and `stride_nb` apart + suh: u64, // fp16 [k] + svh: u64, // fp16 [n] + bias: u64 = 0, // fp16 [n], 0 when the layer has none + k: u32, + n: u32, + halves: u32, // 2 * bits, odd for the mul1 half steps + codebook: Codebook, + stride_k: u64, + stride_nb: u64, +}; + +/// The module and its kernels. +pub const Linear = struct { + mod: Module, + fns: [variants.len]Function, + prompt_fns: [2 * prompt_halves.len]Function, // [halves][4 or 8 warps] + rot_in: Function, + + pub fn load(d: *const Driver) !Linear { + if (!kernels.available) return error.BuiltWithoutKernels; + var l: Linear = .{ .mod = try Module.load(d, kernels.exl3_linear), .fns = undefined, .prompt_fns = undefined, .rot_in = undefined }; + errdefer l.mod.unload(); + l.rot_in = try l.mod.function(rot_in_symbol); + inline for (0..variants.len) |i| { + l.fns[i] = try l.mod.function(symbol(variants[i].halves, variants[i].codebook, variants[i].warps)); + } + inline for (prompt_halves, 0..) |h, i| inline for (.{ 4, 8 }, 0..) |wk, j| { + const f = try l.mod.function(std.fmt.comptimePrint("tf_exl3_prompt_{d}_{d}", .{ h, wk })); + try f.allowDynamicShared(prompt_shared); + l.prompt_fns[2 * i + j] = f; + }; + return l; + } + + pub fn unload(l: *Linear) void { + l.mod.unload(); + } + + /// Shared bytes one call needs: WK * min(m, 8) * 128 floats of warp sums, at most 32 KB. + pub fn sharedBytes(w: Weight, m: usize) u32 { + const p = plan(w.k, w.n, 192, 8); + const red = p.wk * @min(m, 8) * 128 * 4; + const stage = if (w.halves % 2 == 0) p.wk * 8 * 64 * 4 else 0; // even widths stage a step a warp in red + return @intCast(@max(red, stage)); + } + + /// Split partials one call needs (fp32 [sk, m, n]) when the plan splits K, else 0. + pub fn zBytes(w: Weight, m: usize) usize { + const p = plan(w.k, w.n, 192, 8); + return if (p.sk > 1) @as(usize, p.sk) * m * w.n * 4 else 0; + } + + /// Counter bytes one call needs: per pass and block, zeroed (fill8(0), never zeros()). + pub fn counterBytes(w: Weight) usize { + return 8 * (w.n / 128) * 4; + } + + /// y [m, n] = x [m, k] through the weight; `xh` is fp16 [m, k] scratch, `z` the split partials, `counters` zeroed. + pub fn matmul(l: *const Linear, s: Stream, w: Weight, x: u64, x_dtype: DType, m: usize, y: u64, y_dtype: DType, xh: u64, z: u64, counters: u64) !void { + // Shapes the loader admits (linear.py's k % 128 and n % 128, 1..128 rows) are returned errors here too. + if (m < 1 or m > 128) return error.Exl3Shape; + if (w.k % 128 != 0 or w.n % 128 != 0) return error.Exl3Shape; + const p = plan(w.k, w.n, 192, 8); + if ((w.k / 16) % (p.sk * p.wk) != 0) return error.Exl3Shape; + const idx = variantIndex(w.halves, w.codebook, p.wk) orelse return error.UnsupportedExl3Width; + if (p.sk > 1 and z == 0) return error.Exl3Scratch; + + var r: launch_.Args = .{}; + r.add(x); + r.add(w.suh); + r.add(xh); + r.add(@as(i32, @backingInt(x_dtype))); + r.add(@as(i32, @intCast(m))); + r.add(@as(i32, @intCast(w.k))); + try launch_.launch(l.rot_in, .{ + .grid = .{ .x = @intCast((w.k / 128 + 3) / 4), .y = @intCast(m), .z = 1 }, + .block = .{ .x = 128, .y = 1, .z = 1 }, + .shared = 0, + }, s, &r); + + var args: launch_.Args = .{}; + args.add(xh); + args.add(w.trellis); + args.add(@as(i64, @intCast(w.stride_k))); + args.add(@as(i64, @intCast(w.stride_nb))); + args.add(w.svh); + args.add(w.bias); + args.add(y); + args.add(if (p.sk > 1) z else @as(u64, 0)); + args.add(counters); + args.add(@as(i32, @backingInt(y_dtype))); + args.add(@as(i32, @intCast(w.n))); + args.add(@as(i32, @intCast(p.sk))); + args.add(@as(i32, @intCast(m))); + args.add(@as(i32, @intCast(w.k))); + try launch_.launch(l.fns[idx], .{ + .grid = .{ .x = @intCast(w.n / 128 * p.sk), .y = 1, .z = 1 }, + .block = .{ .x = @intCast(p.wk * 32), .y = 1, .z = 1 }, + .shared = sharedBytes(w, m), + }, s, &args); + } + + /// matmul for a prompt block of up to 128 rows on a mul1 weight: each tile decoded once for every row. + pub fn prompt(l: *const Linear, s: Stream, w: Weight, x: u64, x_dtype: DType, m: usize, y: u64, y_dtype: DType, xh: u64, z: u64, counters: u64) !void { + if (m < 1 or m > 128 or w.codebook != .mul1) return error.Exl3Shape; + if (w.k % 128 != 0 or w.n % 128 != 0) return error.Exl3Shape; + const p = plan(w.k, w.n, 192, 8); + if ((w.k / 16) % (p.sk * p.wk) != 0) return error.Exl3Shape; + const idx = promptIndex(w.halves, p.wk) orelse return error.UnsupportedExl3Width; + if (p.sk > 1 and z == 0) return error.Exl3Scratch; + var r: launch_.Args = .{}; + r.add(x); + r.add(w.suh); + r.add(xh); + r.add(@as(i32, @backingInt(x_dtype))); + r.add(@as(i32, @intCast(m))); + r.add(@as(i32, @intCast(w.k))); + try launch_.launch(l.rot_in, .{ + .grid = .{ .x = @intCast((w.k / 128 + 3) / 4), .y = @intCast(m), .z = 1 }, + .block = .{ .x = 128, .y = 1, .z = 1 }, + .shared = 0, + }, s, &r); + var args: launch_.Args = .{}; + args.add(xh); + args.add(w.trellis); + args.add(@as(i64, @intCast(w.stride_k))); + args.add(@as(i64, @intCast(w.stride_nb))); + args.add(w.svh); + args.add(w.bias); + args.add(y); + args.add(if (p.sk > 1) z else @as(u64, 0)); + args.add(counters); + args.add(@as(i32, @backingInt(y_dtype))); + args.add(@as(i32, @intCast(w.n))); + args.add(@as(i32, @intCast(p.sk))); + args.add(@as(i32, @intCast(m))); + args.add(@as(i32, @intCast(w.k))); + try launch_.launch(l.prompt_fns[idx], .{ + .grid = .{ .x = @intCast(w.n / 128 * p.sk), .y = @intCast((m + 63) / 64), .z = 1 }, + .block = .{ .x = 256, .y = 1, .z = 1 }, + .shared = prompt_shared, + }, s, &args); + } +}; + +/// The loaded kernels plus the scratch one stream's dense calls share; prompt rows run in 128-row blocks (forward.py). +pub const Dense = struct { + d: *const Driver, + lin: Linear, + xh: ?memory.DeviceBuffer = null, // the rotated inputs, 128 rows of the widest layer at most + z: ?memory.DeviceBuffer = null, // the split partials, 128 rows at most + counters: ?memory.DeviceBuffer = null, // per pass and block: zeroed once, every launch leaves them zero + + pub fn load(d: *const Driver) !Dense { + return .{ .d = d, .lin = try Linear.load(d) }; + } + + pub fn unload(de: *Dense) void { + de.lin.unload(); + const slots = [_]*?memory.DeviceBuffer{ &de.xh, &de.z, &de.counters }; + for (slots) |slot| { + if (slot.* != null) slot.*.?.free(); + slot.* = null; + } + } + + /// A scratch buffer's address, grown when a layer needs more than it holds. + fn scratch(de: *Dense, slot: *?memory.DeviceBuffer, bytes: usize) !u64 { + if (slot.* != null and slot.*.?.len >= bytes) return slot.*.?.ptr; + if (slot.* != null) slot.*.?.free(); + slot.* = try memory.DeviceBuffer.alloc(de.d, bytes); + if (slot == &de.counters) try slot.*.?.fill8(0, null); // zero once: each launch's last block resets its own + return slot.*.?.ptr; + } + + /// Sizes the scratch for `w` at its widest (128-row blocks) now: a captured graph's addresses never move later. + pub fn reserve(de: *Dense, w: Weight) !void { + const p = plan(w.k, w.n, 192, 8); + if (p.sk > 1) _ = try de.scratch(&de.z, Linear.zBytes(w, 128)); + _ = try de.scratch(&de.counters, Linear.counterBytes(w)); + _ = try de.scratch(&de.xh, 128 * w.k * 2); + } + + /// Decode rows: one call (the kernel takes 1..128), any count above in 128-row blocks. + pub fn decode(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void { + return de.blocks(s, w, x, x_dtype, rows, y, y_dtype); + } + + /// Prompt rows: 128-row blocks of `matmul`, so a row's bits never depend on where its chunk starts. + pub fn prompt(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void { + return de.blocks(s, w, x, x_dtype, rows, y, y_dtype); + } + fn blocks(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void { + const p = plan(w.k, w.n, 192, 8); + const z: u64 = if (p.sk > 1) try de.scratch(&de.z, Linear.zBytes(w, 128)) else 0; + const counters = try de.scratch(&de.counters, Linear.counterBytes(w)); + const xe: u64 = if (x_dtype == .f32) 4 else 2; + const ye: u64 = if (y_dtype == .f32) 4 else 2; + var r: usize = 0; + while (r < rows) : (r += 128) { + const m = @min(rows - r, 128); + const xh = try de.scratch(&de.xh, m * w.k * 2); + const xr = x + @as(u64, r) * w.k * xe; + const yr = y + @as(u64, r) * w.n * ye; + if (m > 16 and w.codebook == .mul1) { // prompt rows: each tile decoded once for the whole block + try de.lin.prompt(s, w, xr, x_dtype, m, yr, y_dtype, xh, z, counters); + } else { + try de.lin.matmul(s, w, xr, x_dtype, m, yr, y_dtype, xh, z, counters); + } + } + } +}; + +test "the plan matches linear.py's for known shapes" { + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 8, .wk = 4 }), plan(4096, 1024, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 2, .wk = 4 }), plan(1024, 4096, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 4, .wk = 4 }), plan(2048, 512, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 1, .wk = 4 }), plan(512, 1024, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 1, .wk = 4 }), plan(512, 512, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 4, .wk = 4 }), plan(2048, 4096, 192, 8)); +} + +test "the stored and strips strides of one layout" { + try std.testing.expectEqual(@as(@TypeOf(strides(0, 0, 0, false)), .{ .k = 4096, .nb = 256 }), strides(4096, 2048, 8, false)); + try std.testing.expectEqual(@as(@TypeOf(strides(0, 0, 0, true)), .{ .k = 256, .nb = 65536 }), strides(4096, 2048, 8, true)); +} + +test "every variant's symbol is its instantiation's name" { + try std.testing.expectEqualStrings("tf_exl3_linear_8_2_4", symbol(8, .mul1, 4)); + try std.testing.expectEqualStrings("tf_exl3_linear_7_2_8", symbol(7, .mul1, 8)); + try std.testing.expectEqualStrings("tf_exl3_linear_8_0_4", symbol(8, .inst3, 4)); + try std.testing.expectEqual(@as(?usize, 5), variantIndex(12, .mul1, 4)); + try std.testing.expectEqual(@as(?usize, 2), variantIndex(6, .mul1, 4)); + try std.testing.expectEqual(@as(?usize, null), variantIndex(3, .mcg, 4)); // the half steps are mul1's alone +} + +test "matmul refuses shapes the loader would refuse, before any launch" { + const l: Linear = .{ .mod = undefined, .fns = undefined, .prompt_fns = undefined, .rot_in = undefined }; + const s: Stream = undefined; // never reached: the refusal comes first + const st = strides(4096, 1024, 8, false); + const w: Weight = .{ .trellis = 0, .suh = 0, .svh = 0, .bias = 0, .k = 4096, .n = 1024, .halves = 8, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; + try std.testing.expectError(error.Exl3Shape, l.matmul(s, w, 0, .bf16, 0, 0, .bf16, 0, 0, 0)); + try std.testing.expectError(error.Exl3Shape, l.matmul(s, w, 0, .bf16, 129, 0, .bf16, 0, 0, 0)); + var kbad = w; + kbad.k = 144; // k % 16 == 0 but not a multiple of 128 + try std.testing.expectError(error.Exl3Shape, l.matmul(s, kbad, 0, .bf16, 8, 0, .bf16, 0, 0, 0)); + var nbad = w; + nbad.n = 100; + try std.testing.expectError(error.Exl3Shape, l.matmul(s, nbad, 0, .bf16, 8, 0, .bf16, 0, 0, 0)); + try std.testing.expectError(error.Exl3Scratch, l.matmul(s, w, 0, .bf16, 8, 0, .bf16, 0, 0, 0)); // plan(4096, 1024) splits K +} diff --git a/zig/src/cuda/exl3_experts.zig b/zig/src/cuda/exl3_experts.zig new file mode 100644 index 0000000000..8301d85586 --- /dev/null +++ b/zig/src/cuda/exl3_experts.zig @@ -0,0 +1,342 @@ +//! The EXL3 routed experts (exl3/experts.cu's device code): the group, the grouped GEMMs, the epilogues. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const exl3 = @import("exl3.zig"); +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +pub const DType = exl3.DType; +pub const Codebook = exl3.Codebook; + +/// One layer's routed experts (experts.py's Exl3RoutedExperts): per-expert trellis pointers, stacked scales, the clamp. +pub const Layer = struct { + gate: u64, // int64 [E], one trellis pointer per expert + up: u64, + down: u64, + gate_k2: u64, // int32 [E], half-bits a value (2 * bits) + up_k2: u64, + down_k2: u64, + suh_g: u64, // fp16 [E, D] + suh_u: u64, + svh_g: u64, // fp16 [E, I] + svh_u: u64, + suh_d: u64, // fp16 [E, I] + svh_d: u64, // fp16 [E, D] + experts: usize, // E + dims: usize, // D (model width) + width: usize, // I (expert width) + codebook: Codebook, + gu_stride: usize = 0, // gate/up trellis row stride in tiles (0: contiguous per projection) + limit: f32 = 0, // the SwiGLU clamp (0: none) + act_mode: i32 = 0, // 0: GLM's bf16 roundings (the route RED-SNOW runs), 1: fp32 +}; + +/// What one projection launches (experts.py default_config): the 8-tile 4-warp K splits, shape-fixed; n % 128 required. +pub fn plan(k: usize, n: usize, gateup: bool) !u32 { + const cands = [_]usize{ if (gateup) 4 else 1, 2, 1 }; + for (cands) |sk| { + if (k % (16 * sk * 4) == 0 and n % 128 == 0) return @intCast(sk); + } + return error.TileNotBuilt; +} + +/// One call's buffers (exl3/experts.py Scratch), laid out in scratch()'s order, each 256-byte aligned. +pub const Scratch = struct { + xg: u64, // fp16 [P, D], the rotated gate inputs + xu: u64, // fp16 [P, D], the rotated up inputs + xd: u64, // fp16 [P, I], the rotated gate and up outputs (the down inputs) + z: u64, // fp32, gate and up write 2 * sk_gu * P * I partials, down sk_d * P * D + y: u64, // fp32 [P, D], the per-slot outputs + ids: u64, // int32 [min(P, E)], distinct routed experts in id order + count: u64, // int32 [1] + members: u64, // int32 [min(P, E), rows], row * 32 + slot, -1 after the last member + rows: usize, // the row capacity the buffers hold + slots: usize, // slots a row, at most 32 +}; + +/// The byte size of the scratch buffers for a layer and row capacity, or error.TileNotBuilt when plan() fails. +pub fn scratchBytes(l: Layer, rows: usize, slots: usize) !usize { + const g_sk = try plan(l.dims, l.width, true); + const d_sk = try plan(l.width, l.dims, false); + const p = rows * slots; + const maxu = @min(p, l.experts); + return align256(p * l.dims * 2) + align256(p * l.dims * 2) + align256(p * l.width * 2) + + align256(@max(2 * g_sk * l.width, d_sk * l.dims) * p * 4) + align256(p * l.dims * 4) + + align256(maxu * 4) + 256 + align256(maxu * rows * 4); +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +/// The pointers into one allocation of scratchBytes(), or error.TileNotBuilt when plan() fails. +pub fn scratch(base: u64, l: Layer, rows: usize, slots: usize) !Scratch { + const p = rows * slots; + const g_sk = try plan(l.dims, l.width, true); + const d_sk = try plan(l.width, l.dims, false); + const maxu = @min(p, l.experts); + var at = base; + const take = struct { + fn f(a: *u64, bytes: usize) u64 { + const r = a.*; + a.* = r + align256(bytes); + return r; + } + }.f; + const xg = take(&at, p * l.dims * 2); + const xu = take(&at, p * l.dims * 2); + const xd = take(&at, p * l.width * 2); + const z = take(&at, @max(2 * g_sk * l.width, d_sk * l.dims) * p * 4); + const y = take(&at, p * l.dims * 4); + const ids = take(&at, maxu * 4); + const count = take(&at, 4); + const members = take(&at, maxu * rows * 4); + return .{ .xg = xg, .xu = xu, .xd = xd, .z = z, .y = y, .ids = ids, .count = count, .members = members, .rows = rows, .slots = slots }; +} + +/// The group kernel's dynamic shared memory at most (34 ints an expert: up to 722 experts). +const group_shared_max = 96 * 1024; + +/// The loaded kernels. +pub const Experts = struct { + mod: Module, + group: Function, + rot_in: Function, + grouped: [3]Function, // one per codebook + grouped_nc: [3]Function, // the same with staged 16-byte trellis loads (decode windows) + gateup: Function, + down_combine: Function, + down_epilogue: Function, + pub fn load(d: *const Driver) !Experts { + if (!kernels.available) return error.BuiltWithoutKernels; + var e: Experts = .{ + .mod = try Module.load(d, kernels.exl3_experts), + .group = undefined, + .rot_in = undefined, + .grouped = undefined, + .grouped_nc = undefined, + .gateup = undefined, + .down_combine = undefined, + .down_epilogue = undefined, + }; + errdefer e.mod.unload(); + e.group = try e.mod.function("tf_exl3_group"); + try e.group.allowDynamicShared(group_shared_max); + e.rot_in = try e.mod.function("tf_exl3_experts_rot_in"); + e.grouped = .{ + try e.mod.function("tf_exl3_grouped_0"), + try e.mod.function("tf_exl3_grouped_1"), + try e.mod.function("tf_exl3_grouped_2"), + }; + e.grouped_nc = .{ + try e.mod.function("tf_exl3_grouped_nc_0"), + try e.mod.function("tf_exl3_grouped_nc_1"), + try e.mod.function("tf_exl3_grouped_nc_2"), + }; + e.gateup = try e.mod.function("tf_exl3_gateup_epilogue"); + e.down_combine = try e.mod.function("tf_exl3_down_combine"); + e.down_epilogue = try e.mod.function("tf_exl3_down_epilogue"); + return e; + } + + pub fn unload(e: *Experts) void { + e.mod.unload(); + } + + /// The pipeline through the grouped down partials; run and runPerSlot each add their epilogue launch. + fn pipeline(e: *const Experts, s: Stream, l: Layer, x: u64, x_stride: usize, x_dtype: DType, pick: u64, rows: usize, sc: Scratch) !void { + const d = l.dims; + const w = l.width; + const slots = sc.slots; + std.debug.assert(rows >= 1 and rows <= sc.rows); + std.debug.assert(slots >= 1 and slots <= 32); + std.debug.assert(l.experts >= 1 and l.experts <= 4096); + const g_sk = try plan(d, w, true); + const d_sk = try plan(w, d, false); + const p = rows * slots; + const maxu = @min(p, l.experts); + const gfn = (if (rows <= 16) e.grouped_nc else e.grouped)[@backingInt(l.codebook)]; // same bits either way + const i32d = struct { + fn f(v: usize) i32 { + return @intCast(v); + } + }.f; + + var a: launch_.Args = .{}; + a.add(pick); + a.add(sc.ids); + a.add(sc.count); + a.add(sc.members); + a.add(i32d(rows)); + a.add(i32d(slots)); + a.add(i32d(l.experts)); + a.add(i32d(rows)); + const group_shared: u32 = @intCast(34 * l.experts * 4); // per-warp counts, places, totals + if (group_shared > group_shared_max) return error.BadExperts; + try launch_.launch(e.group, .{ .grid = .{ .x = 1, .y = 1, .z = 1 }, .block = .{ .x = 1024, .y = 1, .z = 1 }, .shared = group_shared }, s, &a); + + var r: launch_.Args = .{}; + r.add(x); + r.add(i32d(x_stride)); + r.add(@as(i32, @backingInt(x_dtype))); + r.add(pick); + r.add(l.suh_g); + r.add(l.suh_u); + r.add(sc.xg); + r.add(sc.xu); + r.add(i32d(d)); + r.add(i32d(slots)); + r.add(i32d(l.experts)); + try launch_.launch(e.rot_in, .{ .grid = .{ .x = @intCast(p), .y = @intCast(d / 128), .z = 2 }, .block = .{ .x = 32, .y = 1, .z = 1 } }, s, &r); + + var g: launch_.Args = .{}; + g.add(sc.xg); + g.add(sc.xu); + g.add(l.gate); + g.add(l.up); + g.add(l.gate_k2); + g.add(l.up_k2); + g.add(sc.ids); + g.add(sc.count); + g.add(sc.members); + g.add(sc.z); + g.add(i32d(d)); + g.add(i32d(w)); + g.add(i32d(p)); + g.add(i32d(g_sk)); + g.add(i32d(rows)); + g.add(i32d(slots)); + g.add(i32d(maxu)); + g.add(@as(i32, 2)); + g.add(i32d(l.gu_stride)); + try launch_.launch(gfn, .{ .grid = .{ .x = @intCast(maxu), .y = @intCast(w / 128), .z = @intCast(2 * g_sk) }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &g); + + var ep: launch_.Args = .{}; + ep.add(sc.z); + ep.add(pick); + ep.add(l.svh_g); + ep.add(l.svh_u); + ep.add(l.suh_d); + ep.add(sc.xd); + ep.add(i32d(p)); + ep.add(i32d(w)); + ep.add(i32d(g_sk)); + ep.add(i32d(l.experts)); + ep.add(l.limit); + ep.add(l.act_mode); + try launch_.launch(e.gateup, .{ .grid = .{ .x = @intCast(p), .y = @intCast(w / 128), .z = 1 }, .block = .{ .x = 32, .y = 1, .z = 1 } }, s, &ep); + + var dn: launch_.Args = .{}; + dn.add(sc.xd); + dn.add(sc.xd); + dn.add(l.down); + dn.add(l.down); + dn.add(l.down_k2); + dn.add(l.down_k2); + dn.add(sc.ids); + dn.add(sc.count); + dn.add(sc.members); + dn.add(sc.z); + dn.add(i32d(w)); + dn.add(i32d(d)); + dn.add(i32d(p)); + dn.add(i32d(d_sk)); + dn.add(i32d(rows)); + dn.add(i32d(slots)); + dn.add(i32d(maxu)); + dn.add(@as(i32, 1)); + dn.add(@as(i32, 0)); + try launch_.launch(gfn, .{ .grid = .{ .x = @intCast(maxu), .y = @intCast(d / 128), .z = @intCast(d_sk) }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &dn); + } + + /// The routed experts of `rows` rows: the pipeline, then the wts-weighted sum (down_combine). + pub fn run(e: *const Experts, s: Stream, l: Layer, x: u64, x_stride: usize, x_dtype: DType, pick: u64, wts: u64, rows: usize, out: u64, sc: Scratch) !void { + try e.pipeline(s, l, x, x_stride, x_dtype, pick, rows, sc); + const d = l.dims; + const slots = sc.slots; + const p = rows * slots; + const d_sk = try plan(l.width, l.dims, false); + var cb: launch_.Args = .{}; + cb.add(sc.z); + cb.add(pick); + cb.add(l.svh_d); + cb.add(sc.y); + cb.add(wts); + cb.add(out); + cb.add(@as(i32, @intCast(p))); + cb.add(@as(i32, @intCast(d))); + cb.add(@as(i32, @intCast(d_sk))); + cb.add(@as(i32, @intCast(l.experts))); + cb.add(@as(i32, @intCast(slots))); + try launch_.launch(e.down_combine, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(d / 128), .z = 1 }, .block = .{ .x = @intCast(32 * slots), .y = 1, .z = 1 } }, s, &cb); + } + + /// The pipeline, but the last stage writes each slot's own fp32 y ([rows * slots, d]); invalid picks untouched. + pub fn runPerSlot(e: *const Experts, s: Stream, l: Layer, x: u64, x_stride: usize, x_dtype: DType, pick: u64, rows: usize, y: u64, sc: Scratch) !void { + try e.pipeline(s, l, x, x_stride, x_dtype, pick, rows, sc); + const d = l.dims; + const p = rows * sc.slots; + const d_sk = try plan(l.width, l.dims, false); + var ep: launch_.Args = .{}; + ep.add(sc.z); + ep.add(pick); + ep.add(l.svh_d); + ep.add(y); + ep.add(@as(u64, 0)); // wts, unused + ep.add(@as(u64, 0)); // out, unused + ep.add(@as(i32, @intCast(p))); + ep.add(@as(i32, @intCast(d))); + ep.add(@as(i32, @intCast(d_sk))); + ep.add(@as(i32, @intCast(l.experts))); + ep.add(@as(i32, @intCast(sc.slots))); + try launch_.launch(e.down_epilogue, .{ .grid = .{ .x = @intCast(p), .y = @intCast(d / 128), .z = 1 }, .block = .{ .x = 32, .y = 1, .z = 1 } }, s, &ep); + } +}; + +test "the plan matches experts.py's for known shapes" { + try std.testing.expectEqual(@as(u32, 4), try plan(4096, 2048, true)); // GLM_GATEUP + try std.testing.expectEqual(@as(u32, 1), try plan(2048, 4096, false)); // GLM_DOWN + try std.testing.expectEqual(@as(u32, 4), try plan(512, 256, true)); + try std.testing.expectEqual(@as(u32, 1), try plan(256, 512, false)); + try std.testing.expectError(error.TileNotBuilt, plan(300, 256, true)); +} + +test "the scratch layout is 256-aligned and covered by scratchBytes" { + const l: Layer = .{ + .gate = 0, + .up = 0, + .down = 0, + .gate_k2 = 0, + .up_k2 = 0, + .down_k2 = 0, + .suh_g = 0, + .suh_u = 0, + .svh_g = 0, + .svh_u = 0, + .suh_d = 0, + .svh_d = 0, + .experts = 9, + .dims = 512, + .width = 256, + .codebook = .mul1, + }; + const bytes = try scratchBytes(l, 64, 8); + var buf: [4096]u8 align(256) = undefined; + const sc = try scratch(@intFromPtr(&buf), l, 64, 8); + try std.testing.expectEqual(@as(u64, @intFromPtr(&buf)), sc.xg); + try std.testing.expectEqual(@as(u64, 0), sc.xg % 256); + try std.testing.expectEqual(@as(u64, 0), sc.xu % 256); + try std.testing.expectEqual(@as(u64, 0), sc.xd % 256); + try std.testing.expectEqual(@as(u64, 0), sc.z % 256); + try std.testing.expectEqual(@as(u64, 0), sc.y % 256); + try std.testing.expectEqual(@as(u64, 0), sc.ids % 256); + try std.testing.expectEqual(@as(u64, 0), sc.count % 256); + try std.testing.expectEqual(@as(u64, 0), sc.members % 256); + try std.testing.expectEqual(@as(usize, @intFromPtr(&buf) + bytes), sc.members + @as(usize, 9) * 64 * 4); +} diff --git a/zig/src/cuda/exl3_upload.zig b/zig/src/cuda/exl3_upload.zig new file mode 100644 index 0000000000..ef0e434a09 --- /dev/null +++ b/zig/src/cuda/exl3_upload.zig @@ -0,0 +1,105 @@ +//! Uploading core exl3_load groups to the device as the exl3 kernels' weights (the engine's load path). +const std = @import("std"); +const driver = @import("driver.zig"); +const memory = @import("memory.zig"); +const exl3 = @import("exl3.zig"); +const exl3e = @import("exl3_experts.zig"); +const Driver = driver.Driver; + +/// The tensors one upload holds: appended to `bufs` and freed together. +pub const Uploader = struct { + d: *const Driver, + gpa: std.mem.Allocator, + bufs: std.ArrayList(memory.DeviceBuffer) = .empty, + bytes: usize = 0, // the tensors' total device bytes + + pub fn deinit(u: *Uploader) void { + for (u.bufs.items) |*b| b.free(); + u.bufs.deinit(u.gpa); + } + + pub fn dev(u: *Uploader, bytes: []const u8) !u64 { + const b = try memory.DeviceBuffer.fromHost(u.d, bytes); + try u.bufs.append(u.gpa, b); + u.bytes += bytes.len; + return b.ptr; + } + + /// One dense group: `parts` (exl3_load.Parts) and `g` (exl3_load.Group) as the Weight; the scales must be fp16. + pub fn dense(u: *Uploader, parts: anytype, g: anytype) !exl3.Weight { + if (parts.suh_packed) return error.PackedScales; + const st = exl3.strides(g.k, g.n, g.halves, false); + return .{ + .trellis = try u.dev(parts.trellis.bytes), + .suh = try u.dev(parts.suh.bytes), + .svh = try u.dev(parts.svh.bytes), + .bias = if (parts.bias) |b| try u.dev(b.bytes) else 0, + .k = @intCast(g.k), + .n = @intCast(g.n), + .halves = g.halves, + .codebook = switch (g.codebook) { + .inst3 => .inst3, + .mcg => .mcg, + .mul1 => .mul1, + }, + .stride_k = st.k, + .stride_nb = st.nb, + }; + } +}; + +/// One expert projection's tensors (fp16 scales) and half-bit width: the exl3_load Parts' bytes, the group's halves. +pub const Proj = struct { trellis: []const u8, suh: []const u8, svh: []const u8, k2: u32 }; + +/// One layer's routed experts: trellises in place (pointers stacked), scales stacked [E, n] fp16, widths int32 [E]. +pub fn layer(u: *Uploader, a: std.mem.Allocator, dims: usize, width: usize, cb: exl3.Codebook, limit: f32, act_mode: i32, gate: []const Proj, up: []const Proj, down: []const Proj) !exl3e.Layer { + const e = gate.len; + if (up.len != e or down.len != e or e == 0) return error.BadExperts; + const gt = try pointers(u, a, gate); + const ut = try pointers(u, a, up); + const dt = try pointers(u, a, down); + return .{ + .gate = try u.dev(std.mem.sliceAsBytes(gt)), + .up = try u.dev(std.mem.sliceAsBytes(ut)), + .down = try u.dev(std.mem.sliceAsBytes(dt)), + .gate_k2 = try k2s(u, a, gate), + .up_k2 = try k2s(u, a, up), + .down_k2 = try k2s(u, a, down), + .suh_g = try stacked(u, a, gate, dims, "suh"), + .suh_u = try stacked(u, a, up, dims, "suh"), + .svh_g = try stacked(u, a, gate, width, "svh"), + .svh_u = try stacked(u, a, up, width, "svh"), + .suh_d = try stacked(u, a, down, width, "suh"), + .svh_d = try stacked(u, a, down, dims, "svh"), + .experts = e, + .dims = dims, + .width = width, + .codebook = cb, + .gu_stride = 0, + .limit = limit, + .act_mode = act_mode, + }; +} + +fn pointers(u: *Uploader, a: std.mem.Allocator, projs: []const Proj) ![]u64 { + const out = try a.alloc(u64, projs.len); + for (projs, 0..) |p, i| out[i] = try u.dev(p.trellis); + return out; +} + +fn k2s(u: *Uploader, a: std.mem.Allocator, projs: []const Proj) !u64 { + const host = try a.alloc(i32, projs.len); + for (projs, 0..) |p, i| host[i] = @intCast(p.k2); + return u.dev(std.mem.sliceAsBytes(host)); +} + +/// Every expert's fp16 scale of `name` ("suh" or "svh"), `n` values each, stacked [E, n]. +fn stacked(u: *Uploader, a: std.mem.Allocator, projs: []const Proj, n: usize, comptime name: []const u8) !u64 { + const host = try a.alloc(u8, projs.len * n * 2); + for (projs, 0..) |p, i| { + const s = @field(p, name); + if (s.len != n * 2) return error.BadScales; + @memcpy(host[i * n * 2 ..][0..s.len], s); + } + return u.dev(host); +} diff --git a/zig/src/cuda/glm_bmm.zig b/zig/src/cuda/glm_bmm.zig new file mode 100644 index 0000000000..e704ceccdb --- /dev/null +++ b/zig/src/cuda/glm_bmm.zig @@ -0,0 +1,227 @@ +//! GLM's BF16 matmul (qmm.py's `_bmm`) and router dots (glue.py): the Python engine's K slices and mma order. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// ROUTER_KS: the router's shape-fixed slice count; sliced and single-pass reductions differ in their last bits. +pub const router_ks: u32 = 8; + +/// A BF16 matrix [n, k] as the checkpoint stores it (qmm.B16); a call multiplies x [m, k] through it. +pub const B16 = struct { + w: u64, // bf16 [n, k] row-major, device; `frag`: in fragment order (pack) + n: u32, + k: u32, + frag: bool = false, +}; + +/// w [n, k] bf16 -> fragment order: per 8-row tile and 64-wide K step, each lane's eight mma B words side by side. +pub fn pack(out: []u8, w: []const u8, n: usize, k: usize) !void { + if (n % 8 != 0 or k % 64 != 0 or w.len != n * k * 2 or out.len != w.len) return error.GlmBmmShape; + const src: []align(1) const u32 = std.mem.bytesAsSlice(u32, w); + const dst: []align(1) u32 = std.mem.bytesAsSlice(u32, out); + const steps = k / 64; + for (0..n / 8) |j| for (0..steps) |s| for (0..32) |l| for (0..4) |kk| for (0..2) |h| { + const row = 8 * j + l / 4; + const col = 64 * s + 16 * kk + 2 * (l % 4) + 8 * h; + dst[((j * steps + s) * 32 + l) * 8 + kk * 2 + h] = src[(row * k + col) / 2]; + }; +} + +/// qmm.b16_split_k: K slices fixed by the shape alone, in units of 64, so every row's bits hold at any row count. +pub fn b16SplitK(n: usize, k: usize) u32 { + const tiles = (n + 63) / 64; + const steps = k / 64; + var sk: u32 = 1; + while (sk < 8 and tiles * sk < 192 and steps % (@as(usize, sk) * 2) == 0 and steps / (@as(usize, sk) * 2) >= 4) { + sk *= 2; + } + return sk; +} + +/// qmm.bucket: the row-block size one call runs at; no bucket changes a row's bits. +pub fn bucket(m: usize) u32 { + for ([_]usize{ 16, 32, 64, 128 }) |b| { + if (m <= b) return @intCast(b); + } + return 128; +} + +fn bmIndex(bm: u32) usize { + return switch (bm) { + 16 => 0, + 32 => 1, + 64 => 2, + else => 3, + }; +} + +const bm_names = [_][]const u8{ "16", "32", "64", "128" }; + +/// The loaded fatbin: eight matmul instantiations, the router's two passes and the partial reductions. +pub const Bmm = struct { + mod: Module, + bmm: [2][4]Function, // [f32 out][bm index] + pbmm: [2][4]Function, // the same over packed weights + routers: [4]Function, + router_part: [4]Function, + prouters: [4]Function, // the same over packed weights + prouter_part: [4]Function, + reduce: Function, + router_sum: Function, + + pub fn load(d: *const driver.Driver) !Bmm { + if (!kernels.available) return error.BuiltWithoutKernels; + var b: Bmm = .{ + .mod = try Module.load(d, kernels.glm_bmm), + .bmm = undefined, + .pbmm = undefined, + .routers = undefined, + .router_part = undefined, + .prouters = undefined, + .prouter_part = undefined, + .reduce = undefined, + .router_sum = undefined, + }; + errdefer b.mod.unload(); + inline for (0..2) |f| { + inline for (0..4) |i| { + b.bmm[f][i] = try b.mod.function(std.fmt.comptimePrint("tf_glm_bmm_b16_{s}_{d}", .{ bm_names[i], f })); + b.pbmm[f][i] = try b.mod.function(std.fmt.comptimePrint("tf_glm_bmm_p16_{s}_{d}", .{ bm_names[i], f })); + } + } + inline for (0..4) |i| { + b.routers[i] = try b.mod.function(std.fmt.comptimePrint("tf_glm_router_{s}", .{bm_names[i]})); + b.router_part[i] = try b.mod.function(std.fmt.comptimePrint("tf_glm_router_part_{s}", .{bm_names[i]})); + b.prouters[i] = try b.mod.function(std.fmt.comptimePrint("tf_glm_prouter_{s}", .{bm_names[i]})); + b.prouter_part[i] = try b.mod.function(std.fmt.comptimePrint("tf_glm_prouter_part_{s}", .{bm_names[i]})); + } + b.reduce = try b.mod.function("tf_glm_bmm_reduce"); + b.router_sum = try b.mod.function("tf_glm_router_sum"); + return b; + } + + pub fn unload(b: *Bmm) void { + b.mod.unload(); + } + + /// fp32 partial bytes one split call needs (qmm._matmul_b16's `part`), 0 when the shape runs unsplit. + pub fn partBytes(w: B16, m: usize) usize { + const sk = b16SplitK(w.n, w.k); + return if (sk > 1) @as(usize, sk) * m * w.n * 4 else 0; + } + + /// out [m, n] = x [m, k] through w; `part` (sk*m*n*4 bytes) is required whenever b16SplitK splits K. + pub fn matmul(b: *const Bmm, s: Stream, w: B16, x: u64, xs: i64, m: usize, out: u64, is_f32: bool, part: u64) !void { + if (m < 1 or w.k == 0 or w.k % 64 != 0 or w.n == 0) return error.GlmBmmShape; + const sk = b16SplitK(w.n, w.k); + if (sk > 1 and part == 0) return error.GlmBmmScratch; + const bm = bucket(m); + var r: launch_.Args = .{}; + r.add(x); + r.add(xs); + r.add(w.w); + r.add(out); + r.add(part); + r.add(@as(i32, @intCast(m))); + r.add(@as(i32, @intCast(w.n))); + r.add(@as(i32, @intCast(w.k))); + r.add(@as(i32, @intCast(sk))); + if (w.frag and w.n % 8 != 0) return error.GlmBmmShape; + const fns = if (w.frag) &b.pbmm else &b.bmm; + try launch_.launch(fns[@intFromBool(is_f32)][bmIndex(bm)], .{ + .grid = .{ .x = @intCast((m + bm - 1) / bm), .y = @intCast((w.n + 63) / 64), .z = @intCast(sk) }, + .block = .{ .x = @as(u32, if (bm == 128) 256 else 128), .y = 1, .z = 1 }, + .shared = 0, + }, s, &r); + if (sk > 1) { + var rr: launch_.Args = .{}; + rr.add(part); + rr.add(out); + rr.add(@as(i32, @intCast(m * w.n))); + rr.add(@as(i32, @intCast(sk))); + rr.add(@as(i32, @intFromBool(is_f32))); + try launch_.launch(b.reduce, .{ + .grid = .{ .x = @intCast(((m * w.n) + 1023) / 1024), .y = 1, .z = 1 }, + .block = .{ .x = 1024, .y = 1, .z = 1 }, + .shared = 0, + }, s, &rr); + } + } + + /// glue.router: fp32 [m, ne] logits = x [m, d] through the gate [ne, d]; the shape's sliced path, else one pass. + pub fn router(b: *const Bmm, s: Stream, x: u64, xs: i64, gate: B16, m: usize, out: u64, part: u64) !void { + const w = gate.w; + const ne: usize = gate.n; + const d: usize = gate.k; + if (m < 1 or ne == 0 or d == 0 or (gate.frag and (ne % 8 != 0 or d % 64 != 0))) return error.GlmBmmShape; + const bm: u32 = if (m <= 16) 16 else if (m <= 32) 32 else if (m <= 64) 64 else 128; + const i = bmIndex(bm); + if (router_ks <= 1 or d % (router_ks * 64) != 0) { + var r: launch_.Args = .{}; + r.add(x); + r.add(xs); + r.add(w); + r.add(out); + r.add(@as(u64, 0)); + r.add(@as(i32, @intCast(m))); + r.add(@as(i32, @intCast(ne))); + r.add(@as(i32, @intCast(d))); + r.add(@as(i32, 1)); + try launch_.launch((if (gate.frag) b.prouters else b.routers)[i], .{ + .grid = .{ .x = @intCast((m + bm - 1) / bm), .y = @intCast((ne + 63) / 64), .z = 1 }, + .block = .{ .x = 128, .y = 1, .z = 1 }, + .shared = 0, + }, s, &r); + return; + } + if (part == 0) return error.GlmBmmScratch; + var r: launch_.Args = .{}; + r.add(x); + r.add(xs); + r.add(w); + r.add(out); + r.add(part); + r.add(@as(i32, @intCast(m))); + r.add(@as(i32, @intCast(ne))); + r.add(@as(i32, @intCast(d))); + r.add(@as(i32, @intCast(router_ks))); + try launch_.launch((if (gate.frag) b.prouter_part else b.router_part)[i], .{ + .grid = .{ .x = @intCast((m + bm - 1) / bm), .y = @intCast((ne + 31) / 32), .z = @intCast(router_ks) }, + .block = .{ .x = 128, .y = 1, .z = 1 }, + .shared = 0, + }, s, &r); + var rr: launch_.Args = .{}; + rr.add(part); + rr.add(out); + rr.add(@as(i32, @intCast(m * ne))); + rr.add(@as(i32, @intCast(router_ks))); + try launch_.launch(b.router_sum, .{ + .grid = .{ .x = @intCast(((m * ne) + 1023) / 1024), .y = 1, .z = 1 }, + .block = .{ .x = 1024, .y = 1, .z = 1 }, + .shared = 0, + }, s, &rr); + } +}; + +test "pack puts a lane's fragment words of one tile and K step side by side" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const n = 16; + const k = 128; + const w = try a.alloc(u16, n * k); + for (w, 0..) |*v, i| v.* = @intCast(i); + const out = try a.alloc(u8, n * k * 2); + try pack(out, std.mem.sliceAsBytes(w), n, k); + const o: []align(1) const u32 = std.mem.bytesAsSlice(u32, out); + // tile 1, step 1, lane 5 (row 9, columns 2 + ...), chunk 3, half 1: column 64 + 48 + 2 + 8 = 122 + const at = ((1 * 2 + 1) * 32 + 5) * 8 + 3 * 2 + 1; + try std.testing.expectEqual(@as(u32, w[9 * k + 122]) | (@as(u32, w[9 * k + 123]) << 16), o[at]); + try std.testing.expectError(error.GlmBmmShape, pack(out[0..2], std.mem.sliceAsBytes(w), n, k)); +} diff --git a/zig/src/cuda/glm_dsa.zig b/zig/src/cuda/glm_dsa.zig new file mode 100644 index 0000000000..40c388fa7e --- /dev/null +++ b/zig/src/cuda/glm_dsa.zig @@ -0,0 +1,233 @@ +//! GLM-5.3-Flash's DSA block on one rank (forward.py dsa_block, latent path, long context, fp8 latent cache). +const std = @import("std"); +const driver = @import("driver.zig"); +const memory = @import("memory.zig"); +const stream_ = @import("stream.zig"); +const exl3 = @import("exl3.zig"); +const glue2_ = @import("glm_glue2.zig"); +const bmm_ = @import("glm_bmm.zig"); +const latent_ = @import("glm_latent.zig"); +const sparse_ = @import("glm_sparse.zig"); +const Stream = stream_.Stream; +const B16 = bmm_.B16; + +pub const hidden = 4096; +pub const q_lora = 1536; +pub const kv_lora = latent_.lw; +pub const heads = latent_.heads; // the rank's attention heads +pub const qk_dim = latent_.d; +pub const index_dim = sparse_.d; +pub const ikr_w = index_dim + sparse_.heads; // ikr's row: the index key, then the 32 head weights +pub const chunk = 512; // keys a dense chunk (attention.CHUNK) +pub const dense_limit = sparse_.width; // the largest context dense attention covers (2,051) +pub const prompt_att_rows = 512; // rows one dense call takes in a prompt chunk (forward.PROMPT_ATT_ROWS) +pub const select_rows = 512; // rows scored at once (sparse.SELECT_ROWS) +const sparse_nch = (sparse_.width + chunk - 1) / chunk; + +/// The kernel families the block launches; the caller owns them (a layer's Kernels holds the values). +pub const Kernels = struct { + glue2: *const glue2_.Glue2, + bmm: *const bmm_.Bmm, + dense: *exl3.Dense, + latent: *latent_.Latent, + sparse: *sparse_.Sparse, +}; + +/// The model's constants the block reads (weights.Config; the indexer's LayerNorm eps is sparse.py's 1e-6). +pub const Config = struct { eps: f32, index_eps: f32, scale: f32, index_scale: f32, wscale: f32 }; + +/// The replicated indexer: wk | weights_proj [160, D], wq_b [4096, 1536], the gate [128, D], LayerNorm, pool ape. +pub const IndexW = struct { kw: B16, qb: B16, gate: B16, ln_w: u64, ln_b: u64, ape: u64 }; + +/// The rank's DSA weights: q_a/kv_a/q_b and o_proj's rows (EXL3), the bf16 norms, kv_b as AbsorbW [H, 256, 512]. +pub const DsaW = struct { + q_a: exl3.Weight, + kv_a: exl3.Weight, + q_b: exl3.Weight, + o: exl3.Weight, + q_norm: u64, + kv_norm: u64, + wk: u64, + wv: u64, + index: IndexW, +}; + +/// This layer's committed caches: fp8 latents [cap, 512] with f32 scales [cap], ik/ig [cap, 128], pk [cap/4 + 2, 128]. +pub const DsaState = struct { lc: u64, scales: u64, ik: u64, ig: u64, pk: u64, pos_dev: u64, cap: usize }; + +const Lay = struct { + at: u64, + fn take(l: *Lay, bytes: usize) u64 { + const r = l.at; + l.at += (bytes + 255) / 256 * 256; + return r; + } +}; + +/// Buffers for windows of up to `rows` rows (forward.Buffers' DSA slice); torch.empty's: nothing initialized. +pub const Scratch = struct { + buf: memory.DeviceBuffer, + rows: usize, + part_rows: usize, // rows a dense call takes: the window's in decode, 512 in a prompt chunk + nch: usize, // the dense pass's chunk cap (LatentScratch.nch) + np_max: usize, // pools a selection may score (the caches' cap / 4) + dqa: u64, + dkv: u64, + sums: u64, // the rmsnorms' dead group sums + qr: u64, + lat: u64, + q: u64, + ikr: u64, + igr: u64, + mpart: u64, // split-K partials of the bf16 matmuls and the router (one at a time) + qa: u64, + ol: u64, + vn: u64, + qi: u64, + po: u64, + pm: u64, + pl: u64, + scores: u64, + pools: u64, + tokens: u64, + counts: u64, + posv: u64, // a prompt chunk's dense blocks' positions (pos + r0), device i32 each + + fn place(l: *Lay, w: *const DsaW, rows: usize, prefill: bool, cap: usize) !Scratch { + if (rows == 0 or cap < 4 * select_rows) return error.GlmDsaScratch; + var sc: Scratch = undefined; + sc.rows = rows; + sc.part_rows = if (prefill) @min(rows, prompt_att_rows) else rows; + sc.nch = (@min(cap, 2560) + rows + chunk - 1) / chunk; + sc.np_max = cap / 4; + const qn: usize = heads * qk_dim; + const parts = @max(sc.nch * sc.part_rows, sparse_nch * rows) * heads; + var mp: usize = @as(usize, bmm_.router_ks) * rows * index_dim * 4; + for ([_]B16{ w.index.kw, w.index.qb }) |b| mp = @max(mp, bmm_.Bmm.partBytes(b, rows)); + sc.dqa = l.take(rows * q_lora * 2); + sc.dkv = l.take(rows * kv_lora * 2); + sc.sums = l.take(rows * q_lora / 64 * 4); + sc.qr = l.take(rows * q_lora * 2); + sc.lat = l.take(rows * kv_lora * 2); + sc.q = l.take(rows * qn * 2); + sc.ikr = l.take(rows * ikr_w * 2); + sc.igr = l.take(rows * index_dim * 4); + sc.mpart = l.take(mp); + sc.qa = l.take(rows * heads * kv_lora * 2); + sc.ol = l.take(rows * heads * kv_lora * 2); + sc.vn = l.take(rows * qn * 2); + sc.qi = l.take(rows * sparse_.heads * index_dim * 2); + sc.po = l.take(parts * kv_lora * 4); + sc.pm = l.take(parts * 4); + sc.pl = l.take(parts * 4); + sc.scores = l.take(@min(rows, select_rows) * sc.np_max * 4); + sc.pools = l.take(rows * sparse_.topk_pools * 8); + sc.tokens = l.take(rows * sparse_.width * 4); + sc.counts = l.take(rows * 4); + sc.posv = l.take((rows + prompt_att_rows - 1) / prompt_att_rows * 4); + return sc; + } + + /// `rows`: the buffer's (Buffers.rows: they cap the dense chunks, as LatentScratch.nch does); `cap`: the caches'. + pub fn init(d: *const driver.Driver, w: *const DsaW, rows: usize, prefill: bool, cap: usize) !Scratch { + var probe: Lay = .{ .at = 0 }; + _ = try place(&probe, w, rows, prefill, cap); + var buf = try memory.DeviceBuffer.alloc(d, probe.at); + errdefer buf.free(); + var l: Lay = .{ .at = buf.ptr }; + var sc = try place(&l, w, rows, prefill, cap); + sc.buf = buf; + return sc; + } + + pub fn deinit(sc: *Scratch) void { + sc.buf.free(); + } +}; + +/// sparse.pool_bucket: the visible pools rounded up to a power of two, at least 1,024, at most the caches'. +pub fn poolBucket(pos: usize, rows: usize, np_max: usize) usize { + const visible = (pos + rows) / sparse_.pool + 1; + return @min(np_max, @max(1024, std.math.ceilPowerOfTwoAssert(usize, visible))); +} + +/// dsa_front: q_a and kv_a on normed, their rmsnorms (qr, lat), then q_b into q [rows, 32, 256]. +fn front(k: Kernels, s: *Stream, c: Config, w: *const DsaW, sc: *const Scratch, normed: u64, rows: usize) !void { + try k.dense.prompt(s.*, w.q_a, normed, .bf16, rows, sc.dqa, .bf16); + try k.dense.prompt(s.*, w.kv_a, normed, .bf16, rows, sc.dkv, .bf16); + try k.glue2.rmsnorm2(s.*, sc.dqa, q_lora, w.q_norm, sc.qr, q_lora, sc.sums, c.eps, rows, q_lora); + try k.glue2.rmsnorm2(s.*, sc.dkv, kv_lora, w.kv_norm, sc.lat, kv_lora, sc.sums, c.eps, rows, kv_lora); + try k.dense.prompt(s.*, w.q_b, sc.qr, .bf16, rows, sc.q, .bf16); +} + +/// forward.dense_attention: part_rows-row calls, each at its first row's position, over the capped chunk count. +fn denseAttn(k: Kernels, s: *Stream, c: Config, st: DsaState, sc: *const Scratch, rows: usize, pos: usize, all: bool) !void { + const nch: u32 = @intCast(if (all) sc.nch else @min((pos + rows + chunk - 1) / chunk, sc.nch)); + const row_qa: u64 = heads * kv_lora * 2; + var r0: usize = 0; + while (r0 < rows) : (r0 += sc.part_rows) { + const n: u32 = @intCast(@min(sc.part_rows, rows - r0)); + var at = st.pos_dev; + if (r0 > 0) { // pos_dev + r0: a device word written in stream order + at = sc.posv + r0 / sc.part_rows * 4; + const slot: memory.DeviceBuffer = .{ .d = sc.buf.d, .ptr = at, .len = 4, .borrowed = true }; + try slot.fill32(@intCast(pos + r0), s.handle); + } + try k.latent.attnChunks(s, sc.qa + r0 * row_qa, st.lc, st.scales, at, 0, 0, sc.po, sc.pm, sc.pl, n, 0, nch, true, false, c.scale); + try k.latent.merge(s, sc.po, sc.pm, sc.pl, sc.ol + r0 * row_qa, 0, n, nch, false); + } +} + +/// select_tokens and sparse_attention: rows past the dense limit attend their top-512 pools' tokens. +fn sparseAttn(k: Kernels, s: *Stream, c: Config, w: *const DsaW, st: DsaState, sc: *const Scratch, rows: usize, pos: usize, capture: bool) !void { + const posp: u64 = if (capture) st.pos_dev else 0; // a captured step is one row: its position is pos_dev's + try k.bmm.matmul(s.*, w.index.qb, sc.qr, q_lora, rows, sc.qi, false, sc.mpart); + const np = poolBucket(pos, rows, sc.np_max); + if (np < sparse_.topk_pools) return error.GlmDsaCapacity; // sparse.top_pools' torch fallback is not ported + const qi_row: u64 = sparse_.heads * index_dim * 2; + var a: usize = 0; + while (a < rows) : (a += select_rows) { + const n: u32 = @intCast(@min(select_rows, rows - a)); + const at: i32 = @intCast(pos + a); + try k.sparse.scores(s, sc.qi + a * qi_row, sparse_.heads * index_dim, sc.ikr + (a * ikr_w + index_dim) * 2, ikr_w, st.pk, sc.scores, at, posp, n, @intCast(np), c.index_scale, c.wscale); + try k.sparse.selectRows(s, sc.scores, @intCast(np), sc.pools + a * sparse_.topk_pools * 8, at, posp, sparse_.topk_pools, n, true); + } + const r32: u32 = @intCast(rows); + try k.sparse.tokensCounts(s, sc.pools, sc.tokens, sc.counts, @intCast(pos), posp, r32); + try k.latent.attnChunks(s, sc.qa, st.lc, st.scales, 0, sc.tokens, sc.counts, sc.po, sc.pm, sc.pl, r32, sparse_.width, sparse_nch, true, true, c.scale); + try k.latent.merge(s, sc.po, sc.pm, sc.pl, sc.ol, sc.counts, r32, sparse_nch, true); +} + +/// dsa_block at pos (= st.pos_dev) into st and part; `capture`: a graph's one-row step (pool bucket fixed by pos). +pub fn dsaBlock(k: Kernels, s: *Stream, c: Config, w: *const DsaW, st: DsaState, sc: *const Scratch, normed: u64, part: u64, rows: usize, pos: usize, capture: bool) !void { + if (rows == 0 or rows > sc.rows) return error.GlmDsaRows; + if (pos + rows > st.cap or st.cap / 4 != sc.np_max) return error.GlmDsaCapacity; + const r32: u32 = @intCast(rows); + const p32: i32 = @intCast(pos); + try front(k, s, c, w, sc, normed, rows); + try k.latent.latWrite(s, sc.lat, kv_lora, st.lc, st.pos_dev, st.scales, r32, true); + const all_sparse = pos >= dense_limit; + const sparse_rows = pos + rows - 1 >= dense_limit; + if (capture and (rows != 1 or sparse_rows != all_sparse)) return error.GlmDsaCapture; + // the indexer runs on every call: it fills the caches while the window is still dense + const ix = w.index; + try k.bmm.matmul(s.*, ix.kw, normed, hidden, rows, sc.ikr, false, sc.mpart); + if (ix.gate.n != index_dim or ix.gate.k != hidden) return error.GlmDsaShape; + try k.bmm.router(s.*, normed, hidden, ix.gate, rows, sc.igr, sc.mpart); + try k.sparse.indexWrite(s, sc.ikr, ikr_w, sc.igr, ix.ln_w, ix.ln_b, st.ik, st.ig, p32, st.pos_dev, c.index_eps, r32); + try k.sparse.poolKeys(s, st.ik, st.ig, ix.ape, st.pk, p32, st.pos_dev, r32); + try k.latent.absorbQ(s, sc.q, w.wk, sc.qa, r32); + // rows within the limit keep the dense result; the sparse merge overwrites only rows with tokens + if (!all_sparse) try denseAttn(k, s, c, st, sc, rows, pos, capture); + if (sparse_rows) try sparseAttn(k, s, c, w, st, sc, rows, pos, capture); + try k.latent.expandV(s, sc.ol, w.wv, sc.vn, r32); + try k.dense.prompt(s.*, w.o, sc.vn, .bf16, rows, part, .f32); +} + +test "pool buckets follow sparse.pool_bucket" { + try std.testing.expectEqual(@as(usize, 1024), poolBucket(100, 1, 2560)); + try std.testing.expectEqual(@as(usize, 2048), poolBucket(8000, 8, 2560)); + try std.testing.expectEqual(@as(usize, 2560), poolBucket(9000, 8, 2560)); + try std.testing.expectEqual(@as(usize, 1024), poolBucket(4092, 3, 4096)); + try std.testing.expectEqual(@as(usize, 2048), poolBucket(4092, 4, 4096)); +} diff --git a/zig/src/cuda/glm_glue.zig b/zig/src/cuda/glm_glue.zig new file mode 100644 index 0000000000..4aaf5dfd04 --- /dev/null +++ b/zig/src/cuda/glm_glue.zig @@ -0,0 +1,231 @@ +//! GLM glue kernels: the Python engine's Triton helpers (glue.py) and greedy sampling (decode.py). +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +fn grid1(x: usize) launch_.Dim3 { + return .{ .x = @intCast(x), .y = 1, .z = 1 }; +} + +fn grid2(x: usize, y: usize) launch_.Dim3 { + return .{ .x = @intCast(x), .y = @intCast(y), .z = 1 }; +} + +fn block(x: u32) launch_.Dim3 { + return .{ .x = x, .y = 1, .z = 1 }; +} + +/// The loaded fatbin: embedding, norms, hyper-connections, MLP glue, router top-k, combine, greedy. +pub const Glue = struct { + mod: Module, + embed_b16_fn: Function, + rmsnorm_fn: Function, + hc_partial_fn: Function, + hc_finish_fn: Function, + hc_post_fn: Function, + residual_add_fn: Function, + stream_mean_fn: Function, + swiglu_fn: Function, + topk_fn: Function, + combine_fn: Function, + greedy_fn: Function, + greedy_pack_fn: Function, + + pub fn load(d: *const driver.Driver) !Glue { + if (!kernels.available) return error.BuiltWithoutKernels; + var g: Glue = .{ + .mod = try Module.load(d, kernels.glm_glue), + .embed_b16_fn = undefined, + .rmsnorm_fn = undefined, + .hc_partial_fn = undefined, + .hc_finish_fn = undefined, + .hc_post_fn = undefined, + .residual_add_fn = undefined, + .stream_mean_fn = undefined, + .swiglu_fn = undefined, + .topk_fn = undefined, + .combine_fn = undefined, + .greedy_fn = undefined, + .greedy_pack_fn = undefined, + }; + errdefer g.mod.unload(); + g.embed_b16_fn = try g.mod.function("tf_glm_embed_b16"); + g.rmsnorm_fn = try g.mod.function("tf_glm_rmsnorm"); + g.hc_partial_fn = try g.mod.function("tf_glm_hc_partial"); + g.hc_finish_fn = try g.mod.function("tf_glm_hc_finish"); + g.hc_post_fn = try g.mod.function("tf_glm_hc_post"); + g.residual_add_fn = try g.mod.function("tf_glm_residual_add"); + g.stream_mean_fn = try g.mod.function("tf_glm_stream_mean"); + g.swiglu_fn = try g.mod.function("tf_glm_swiglu"); + g.topk_fn = try g.mod.function("tf_glm_topk"); + g.combine_fn = try g.mod.function("tf_glm_combine"); + g.greedy_fn = try g.mod.function("tf_glm_greedy"); + g.greedy_pack_fn = try g.mod.function("tf_glm_greedy_pack"); + return g; + } + + pub fn unload(g: *Glue) void { + g.mod.unload(); + } + + /// glue.embed_b16: out [rows, copies, d] bf16 = embedding rows `ids` copied `copies` times. + pub fn embedB16(g: *const Glue, s: Stream, ids: u64, w: u64, out: u64, rows: usize, d: usize, copies: usize) !void { + if (rows == 0 or d == 0 or d % 64 != 0 or copies == 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(ids); + r.add(w); + r.add(out); + r.add(@as(u32, @intCast(copies))); // the kernel reads rows and d from its grid + try launch_.launch(g.embed_b16_fn, .{ .grid = grid2(rows, d / 64), .block = block(32), .shared = 0 }, s, &r); + } + + /// glue.rmsnorm: out/xs [rows, d] and per-64 group sums; d must be 4096 (the served shape). + pub fn rmsnorm(g: *const Glue, s: Stream, x: u64, xs: i64, w: u64, out: u64, os: i64, sums: u64, eps: f32, rows: usize, d: usize) !void { + if (rows == 0 or d != 4096) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(x); + r.add(xs); + r.add(w); + r.add(out); + r.add(os); + r.add(sums); + r.add(eps); + r.add(@as(u32, @intCast(rows))); + try launch_.launch(g.rmsnorm_fn, .{ .grid = grid1(rows), .block = block(256), .shared = 0 }, s, &r); + } + + /// glue.hc_partial: part [rows, 16, 32] fp32 (24 fn dots + ss at slot 24); x [rows, 16384] bf16. + pub fn hcPartial(g: *const Glue, s: Stream, x: u64, fn_w: u64, part: u64, rows: usize) !void { + if (rows == 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(x); + r.add(fn_w); + r.add(part); + r.add(@as(u32, @intCast(rows))); + try launch_.launch(g.hc_partial_fn, .{ .grid = .{ .x = @intCast(rows), .y = 16, .z = 6 }, .block = block(128), .shared = 0 }, s, &r); + } + + /// glue.hc_finish: mixed/normed output, group sums, post gates [rows, 4] and comb [rows, 16]. + pub fn hcFinish(g: *const Glue, s: Stream, x: u64, part: u64, base: u64, scale: u64, nw: u64, out: u64, xs: u64, post: u64, comb: u64, eps_norm: f32, hc_eps: f32, iters: u32, rows: usize, d: usize) !void { + if (rows == 0 or d != 4096) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(x); + r.add(part); + r.add(base); + r.add(scale); + r.add(nw); + r.add(out); + r.add(xs); + r.add(post); + r.add(comb); + r.add(eps_norm); + r.add(hc_eps); + r.add(@as(i32, @intCast(iters))); + r.add(@as(u32, @intCast(rows))); + try launch_.launch(g.hc_finish_fn, .{ .grid = grid1(rows), .block = block(256), .shared = 0 }, s, &r); + } + + /// glue.hc_post: xout [rows, s, d] = gates mixing s streams; rank-1 partials sit at element stride rs. + pub fn hcPost(g: *const Glue, s: Stream, x: u64, xout: u64, gr: u64, post: u64, comb: u64, rs: i64, world: usize, rows: usize, d: usize, n: usize) !void { + if (rows == 0 or d == 0 or d % 1024 != 0 or n == 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(x); + r.add(xout); + r.add(gr); + r.add(post); + r.add(comb); + r.add(rs); + r.add(@as(u32, @intCast(world))); + r.add(@as(u32, @intCast(n))); + try launch_.launch(g.hc_post_fn, .{ .grid = grid2(rows, d / 1024), .block = block(128), .shared = 0 }, s, &r); + } + + /// glue.residual_add: xout = x + bf16(sum of the world's fp32 partials at stride rs). + pub fn residualAdd(g: *const Glue, s: Stream, x: u64, xout: u64, gr: u64, rs: i64, world: usize, rows: usize, d: usize) !void { + if (rows == 0 or d == 0 or d % 1024 != 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(x); + r.add(xout); + r.add(gr); + r.add(rs); + r.add(@as(u32, @intCast(world))); + r.add(@as(u32, @intCast(rows))); + r.add(@as(u32, @intCast(d))); + try launch_.launch(g.residual_add_fn, .{ .grid = grid2(rows, d / 1024), .block = block(128), .shared = 0 }, s, &r); + } + + /// glue.stream_mean: out [rows, d] = bf16(mean of the s input streams). + pub fn streamMean(g: *const Glue, s: Stream, x: u64, out: u64, rows: usize, d: usize, n: usize) !void { + if (rows == 0 or d == 0 or d % 1024 != 0 or n == 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(x); + r.add(out); + r.add(@as(u32, @intCast(n))); + try launch_.launch(g.stream_mean_fn, .{ .grid = grid2(rows, d / 1024), .block = block(128), .shared = 0 }, s, &r); + } + + /// glue.swiglu: out [rows, w] bf16 and group sums xs from the fused gate/up bf16 pair [rows, 2w]. + pub fn swiglu(g: *const Glue, s: Stream, gu: u64, out: u64, xs: u64, lim: f32, rows: usize, w: usize) !void { + if (rows == 0 or w == 0 or w % 512 != 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(gu); + r.add(out); + r.add(xs); + r.add(lim); + r.add(@as(u32, @intCast(rows))); + r.add(@as(u32, @intCast(w))); + try launch_.launch(g.swiglu_fn, .{ .grid = grid2(rows, w / 512), .block = block(128), .shared = 0 }, s, &r); + } + + /// glue.topk: picks/wts [rows, topk+1]; the last slot is the shared-expert one-hot (ne, 1.0). + pub fn topk(g: *const Glue, s: Stream, lg: u64, bias: u64, pick: u64, wts: u64, scale: f32, rows: usize, ne: usize, tk: usize) !void { + if (rows == 0 or ne == 0 or ne > 511 or tk == 0 or tk > 16) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(lg); + r.add(bias); + r.add(pick); + r.add(wts); + r.add(scale); + r.add(@as(i32, @intCast(ne))); + r.add(@as(i32, @intCast(tk))); + try launch_.launch(g.topk_fn, .{ .grid = grid1(rows), .block = block(128), .shared = 0 }, s, &r); + } + + /// glue.combine: out [rows, d] fp32 = sum over slots of wts[k] * y[r, k, :]. + pub fn combine(g: *const Glue, s: Stream, y: u64, wts: u64, out: u64, rows: usize, d: usize, slots: usize) !void { + if (rows == 0 or d == 0 or d % 1024 != 0 or slots == 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(y); + r.add(wts); + r.add(out); + r.add(@as(u32, @intCast(slots))); + try launch_.launch(g.combine_fn, .{ .grid = grid2(rows, d / 1024), .block = block(128), .shared = 0 }, s, &r); + } + + /// decode greedy: tok [rows] int32 argmax of bf16 logits, ties to the lowest token id. + pub fn greedy(g: *const Glue, s: Stream, lg: u64, tok: u64, rows: usize, n: usize) !void { + if (rows == 0 or n == 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(lg); + r.add(tok); + r.add(@as(u32, @intCast(n))); + try launch_.launch(g.greedy_fn, .{ .grid = grid1(rows), .block = block(256), .shared = 0 }, s, &r); + } + + /// sample_rows' greedy candidate per row: packed [rows, 2] f32 of the top value and its vocab id's bits (+offset). + pub fn greedyPack(g: *const Glue, s: Stream, lg: u64, packed_: u64, rows: usize, n: usize, offset: u32) !void { + if (rows == 0 or n == 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(lg); + r.add(packed_); + r.add(@as(u32, @intCast(n))); + r.add(offset); + try launch_.launch(g.greedy_pack_fn, .{ .grid = grid1(rows), .block = block(256), .shared = 0 }, s, &r); + } +}; diff --git a/zig/src/cuda/glm_glue2.zig b/zig/src/cuda/glm_glue2.zig new file mode 100644 index 0000000000..41571f4fa7 --- /dev/null +++ b/zig/src/cuda/glm_glue2.zig @@ -0,0 +1,83 @@ +//! glue.py's 1536/512 rmsnorms and the fp32-y combine: glm_glue.cu is at the file cap (glm_glue2.cu). +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +fn grid1(x: usize) launch_.Dim3 { + return .{ .x = @intCast(x), .y = 1, .z = 1 }; +} + +fn grid2(x: usize, y: usize) launch_.Dim3 { + return .{ .x = @intCast(x), .y = @intCast(y), .z = 1 }; +} + +fn block(x: u32) launch_.Dim3 { + return .{ .x = x, .y = 1, .z = 1 }; +} + +/// The loaded fatbin: the two narrow rmsnorms and the fp32 combine. +pub const Glue2 = struct { + mod: Module, + rmsnorm_512_fn: Function, + rmsnorm_1536_fn: Function, + combine_f32_fn: Function, + + pub fn load(d: *const driver.Driver) !Glue2 { + if (!kernels.available) return error.BuiltWithoutKernels; + var g: Glue2 = .{ + .mod = try Module.load(d, kernels.glm_glue2), + .rmsnorm_512_fn = undefined, + .rmsnorm_1536_fn = undefined, + .combine_f32_fn = undefined, + }; + errdefer g.mod.unload(); + g.rmsnorm_512_fn = try g.mod.function("tf_glm_rmsnorm_512"); + g.rmsnorm_1536_fn = try g.mod.function("tf_glm_rmsnorm_1536"); + g.combine_f32_fn = try g.mod.function("tf_glm_combine_f32"); + return g; + } + + pub fn unload(g: *Glue2) void { + g.mod.unload(); + } + + /// glue.rmsnorm at d 512 or 1536: row-strided x/out ([rows, 2048] slices), sums [rows, d / 64]. + pub fn rmsnorm2(g: *const Glue2, s: Stream, x: u64, xs: i64, w: u64, out: u64, os: i64, sums: u64, eps: f32, rows: usize, d: usize) !void { + if (rows == 0) return error.GlmGlueShape; + const fn_ = switch (d) { + 512 => g.rmsnorm_512_fn, + 1536 => g.rmsnorm_1536_fn, + else => return error.GlmGlueShape, + }; + var r: launch_.Args = .{}; + r.add(x); + r.add(xs); + r.add(w); + r.add(out); + r.add(os); + r.add(sums); + r.add(eps); + r.add(@as(u32, @intCast(rows))); + try launch_.launch(fn_, .{ .grid = grid1(rows), .block = block(128), .shared = 0 }, s, &r); + } + + /// glue.combine over fp32 per-slot outputs: out [rows, d] fp32 = sum over slots of wts[k] * y[r, k, :]. + pub fn combineF32(g: *const Glue2, s: Stream, y: u64, wts: u64, out: u64, rows: usize, d: usize, slots: usize) !void { + if (rows == 0 or d == 0 or d % 1024 != 0 or slots == 0) return error.GlmGlueShape; + var r: launch_.Args = .{}; + r.add(y); + r.add(wts); + r.add(out); + r.add(@as(u32, @intCast(slots))); + try launch_.launch(g.combine_f32_fn, .{ .grid = grid2(rows, d / 1024), .block = block(128), .shared = 0 }, s, &r); + } +}; + +test { + _ = Glue2; +} diff --git a/zig/src/cuda/glm_head.zig b/zig/src/cuda/glm_head.zig new file mode 100644 index 0000000000..c7959b32f0 --- /dev/null +++ b/zig/src/cuda/glm_head.zig @@ -0,0 +1,90 @@ +//! GLM-5.3-Flash's head on one rank (forward.compute's end, decode.sample_rows greedy): logits and the token pick. +const std = @import("std"); +const driver = @import("driver.zig"); +const memory = @import("memory.zig"); +const stream_ = @import("stream.zig"); +const glue_ = @import("glm_glue.zig"); +const bmm_ = @import("glm_bmm.zig"); +const Stream = stream_.Stream; + +/// The rank's head: the final norm [D] bf16, its lm_head rows [vocab / world, D] bf16 from vocab id `offset`. +pub const HeadW = struct { norm: u64, head: bmm_.B16, offset: u32 }; + +/// One forward's buffers for up to `rows` rows, one allocation. +pub const Scratch = struct { + buf: memory.DeviceBuffer, + rows: usize, + hidden: u64, // [rows, D] bf16: the streams' mean + fnormed: u64, // [rows, D] bf16 + fxs: u64, // [rows, D / 64] f32: the norm's group sums (the head's matmul reads none) + logits: u64, // [rows, n] bf16 + part: u64, // the matmul's split-K partials + cand: u64, // [rows, 2] f32: each row's top value and the bits of its vocab id + + pub fn init(d: *const driver.Driver, s: Stream, w: HeadW, dims: usize, rows: usize) !Scratch { + if (rows == 0 or dims == 0 or dims % 64 != 0) return error.GlmHeadShape; + const n: usize = w.head.n; + const sizes = [_]usize{ rows * dims * 2, rows * dims * 2, rows * dims / 64 * 4, rows * n * 2, bmm_.Bmm.partBytes(w.head, rows), rows * 2 * 4 }; + var total: usize = 0; + for (sizes) |z| total += std.mem.alignForward(usize, z, 256); + var buf = try memory.DeviceBuffer.alloc(d, total); + errdefer buf.free(); + try buf.fill8(0, s.handle); + var at: [sizes.len]u64 = undefined; + var p = buf.ptr; + for (sizes, &at) |z, *a| { + a.* = p; + p += std.mem.alignForward(usize, z, 256); + } + return .{ .buf = buf, .rows = rows, .hidden = at[0], .fnormed = at[1], .fxs = at[2], .logits = at[3], .part = if (sizes[4] > 0) at[4] else 0, .cand = at[5] }; + } + + pub fn deinit(sc: *Scratch) void { + sc.buf.free(); + } +}; + +/// Into sc.cand: streams' mean, final norm, logits, each row's top-1; `last` (a prompt chunk) reads its last row. +pub fn forward(g: *const glue_.Glue, b: *const bmm_.Bmm, s: Stream, eps: f32, dims: usize, streams: usize, w: HeadW, sc: *const Scratch, x: u64, rows: usize, last: bool) !void { + if (rows == 0 or rows > sc.rows) return error.GlmHeadRows; + if (w.head.k != dims) return error.GlmHeadShape; + const d: u64 = dims; + try g.streamMean(s, x, sc.hidden, rows, dims, streams); + try g.rmsnorm(s, sc.hidden, @intCast(dims), w.norm, sc.fnormed, @intCast(dims), sc.fxs, eps, rows, dims); + const first: usize = if (last) rows - 1 else 0; + const m = rows - first; + try b.matmul(s, w.head, sc.fnormed + first * d * 2, @intCast(dims), m, sc.logits, false, sc.part); + try g.greedyPack(s, sc.logits, sc.cand, m, w.head.n, w.offset); +} + +/// sample_rows' greedy pick from every rank's gathered candidates [world, rows, 2]: highest value, ties to lowest id. +pub fn choose(gathered: []const f32, world: usize, rows: usize, out: []u32) !void { + if (gathered.len != world * rows * 2 or out.len < rows) return error.GlmHeadShape; + for (0..rows) |i| { + var best: f32 = -std.math.inf(f32); + var id: u32 = std.math.maxInt(u32); + for (0..world) |r| { + const v = gathered[(r * rows + i) * 2]; + const t: u32 = @bitCast(gathered[(r * rows + i) * 2 + 1]); + if (v > best or (v == best and t < id)) { + best = v; + id = t; + } + } + out[i] = id; + } +} + +test "the choice takes the highest value across ranks, ties to the lowest id" { + const f = struct { + fn id(t: u32) f32 { + return @bitCast(t); + } + }.id; + // rows: rank 1 wins; a tie the lower id (rank 1's) wins; rank 0 wins + const g = [_]f32{ 1.0, f(5), 2.0, f(900), 3.0, f(7), 2.0, f(100), 2.0, f(800), 1.0, f(77) }; + var out: [3]u32 = undefined; + try choose(&g, 2, 3, &out); + try std.testing.expectEqualSlices(u32, &.{ 100, 800, 7 }, &out); + try std.testing.expectError(error.GlmHeadShape, choose(g[0..10], 2, 3, &out)); +} diff --git a/zig/src/cuda/glm_kda.zig b/zig/src/cuda/glm_kda.zig new file mode 100644 index 0000000000..b704df834f --- /dev/null +++ b/zig/src/cuda/glm_kda.zig @@ -0,0 +1,283 @@ +// CUDA wrappers for GLM-5.3-Flash's KDA decode chain (glm5_next/cuda/kda.cu ported to kernels/cuda/glm_kda.cu). + +// One 1024-thread block per head runs a window's rows in order; a replay runs saved rows into a state; a long + +// window splits the state-independent work (all rows) from the row-by-row delta rule; conv windows shift in place. + +// Every function takes device pointers, shapes and a stream. + +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +pub const d_k = 128; // keys and values per head +pub const taps = 4; // the conv's window length (3 rows of state + this row's projection) + +/// The saved row inputs a replay needs (fp32 k and decays, bf16 v, fp32 beta; the wide path also q and y). +pub const Scratch = struct { + k: u64 = 0, + v: u64 = 0, + g: u64 = 0, + b: u64 = 0, + q: u64 = 0, + y: u64 = 0, + + /// Byte counts for H heads and R rows (6 views, laid out back to back by the caller). + pub fn bytes(h: u32, rows: u32) u64 { + const per: u64 = @as(u64, h) * d_k; + const r: u64 = rows; + return r * per * (4 + 2 + 4 + 4 + 2) + r * @as(u64, h) * 4; + } +}; + +pub const Kda = struct { + d: *const Driver, + m: Module, + chain_fn: Function, + prep_fn: Function, + step_fn: Function, + out_fn: Function, + replay_fn: Function, + replay_layers_fn: Function, + loaded: bool = false, + shift_fn: Function, + + pub fn load(dr: *const Driver) !Kda { + var k: Kda = .{ + .d = dr, + .m = undefined, + .chain_fn = undefined, + .prep_fn = undefined, + .step_fn = undefined, + .out_fn = undefined, + .replay_fn = undefined, + .replay_layers_fn = undefined, + .shift_fn = undefined, + }; + errdefer k.unload(); + k.m = try Module.load(dr, kernels.glm_kda); + k.chain_fn = try k.m.function("tf_glm_kda_chain"); + k.prep_fn = try k.m.function("tf_glm_kda_prep"); + k.step_fn = try k.m.function("tf_glm_kda_step"); + k.out_fn = try k.m.function("tf_glm_kda_out"); + k.replay_fn = try k.m.function("tf_glm_kda_replay"); + k.replay_layers_fn = try k.m.function("tf_glm_kda_replay_layers"); + k.shift_fn = try k.m.function("tf_glm_kda_conv_shift"); + k.loaded = true; + return k; + } + + pub fn unload(self: *Kda) void { + if (!self.loaded) return; + self.m.unload(); + self.loaded = false; + } + + // One head per block, R rows in order: P's rows hold q|k|v (3 * H * 128 channels) then the beta logit at b_off; + + // A/G's rows the decays' pre-activations and gates. Zeroed Scratch views skip the saves; zero state_out skips + + // the state write (state_in stays the input). + pub fn chain( + self: *Kda, + s: *Stream, + h: u32, + p: u64, + p_stride: i64, + b_off: i32, + a: u64, + a_stride: i64, + g: u64, + g_stride: i64, + conv_state: u64, + conv_w: u64, + state_in: u64, + a_log: u64, + dt_bias: u64, + norm_w: u64, + eps: f32, + lower: f32, + rows: u32, + out: u64, + state_out: u64, + scratch: Scratch, + ) !void { + var args = launch_.Args{}; + args.add(@as(i32, @intCast(h))); + args.add(p); + args.add(@as(i32, @intCast(p_stride))); + args.add(b_off); + args.add(a); + args.add(@as(i32, @intCast(a_stride))); + args.add(g); + args.add(@as(i32, @intCast(g_stride))); + args.add(conv_state); + args.add(conv_w); + args.add(state_in); + args.add(a_log); + args.add(dt_bias); + args.add(norm_w); + args.add(eps); + args.add(lower); + args.add(@as(i32, @intCast(rows))); + args.add(out); + args.add(state_out); + args.add(scratch.k); + args.add(scratch.v); + args.add(scratch.g); + args.add(scratch.b); + try launch_.launch(self.chain_fn, .{ .grid = .{ .x = h, .y = 1, .z = 1 }, .block = .{ .x = 1024, .y = 1, .z = 1 } }, s.*, &args); + } + + // The wide path: prep and out kernels over all rows, the delta rule row by row; q/y live in scratch.q/.y. + pub fn chainWide( + self: *Kda, + s: *Stream, + h: u32, + p: u64, + p_stride: i64, + b_off: i32, + a: u64, + a_stride: i64, + g: u64, + g_stride: i64, + conv_state: u64, + conv_w: u64, + state_in: u64, + a_log: u64, + dt_bias: u64, + norm_w: u64, + eps: f32, + lower: f32, + rows: u32, + out: u64, + state_out: u64, + scratch: Scratch, + ) !void { + // The engine's host sequence: prep writes all rows' q, k, v, decay, beta; step runs the delta rule. + var args = launch_.Args{}; + args.add(@as(i32, @intCast(h))); + args.add(p); + args.add(@as(i32, @intCast(p_stride))); + args.add(b_off); + args.add(a); + args.add(@as(i32, @intCast(a_stride))); + args.add(conv_state); + args.add(conv_w); + args.add(a_log); + args.add(dt_bias); + args.add(lower); + args.add(scratch.q); + args.add(scratch.k); + args.add(scratch.v); + args.add(scratch.g); + args.add(scratch.b); + try launch_.launch(self.prep_fn, .{ .grid = .{ .x = rows, .y = h, .z = 1 }, .block = .{ .x = 512, .y = 1, .z = 1 } }, s.*, &args); + var sargs = launch_.Args{}; + sargs.add(@as(i32, @intCast(h))); + sargs.add(state_in); + sargs.add(scratch.q); + sargs.add(scratch.k); + sargs.add(scratch.v); + sargs.add(scratch.g); + sargs.add(scratch.b); + sargs.add(@as(i32, @intCast(rows))); + sargs.add(scratch.y); + sargs.add(state_out); + try launch_.launch(self.step_fn, .{ .grid = .{ .x = h, .y = 8, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s.*, &sargs); + var oargs = launch_.Args{}; + oargs.add(@as(i32, @intCast(h))); + oargs.add(scratch.y); + oargs.add(g); + oargs.add(@as(i32, @intCast(g_stride))); + oargs.add(norm_w); + oargs.add(eps); + oargs.add(out); + try launch_.launch(self.out_fn, .{ .grid = .{ .x = rows, .y = h, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s.*, &oargs); + } + + /// Rows saved by a chain (or chain_wide) into a fresh state, skipping their conv and norms. + pub fn replay( + self: *Kda, + s: *Stream, + h: u32, + state_in: u64, + scratch: Scratch, + rows: u32, + state_out: u64, + ) !void { + var args = launch_.Args{}; + args.add(@as(i32, @intCast(h))); + args.add(state_in); + args.add(scratch.k); + args.add(scratch.v); + args.add(scratch.g); + args.add(scratch.b); + args.add(@as(i32, @intCast(rows))); + args.add(state_out); + try launch_.launch(self.replay_fn, .{ .grid = .{ .x = h, .y = 1, .z = 1 }, .block = .{ .x = 1024, .y = 1, .z = 1 } }, s.*, &args); + } + + /// All layers at once: layers * H blocks, per-layer state strides and per-layer saved-row strides. + pub fn replayLayers( + self: *Kda, + s: *Stream, + h: u32, + layers: u32, + state_in: u64, + state_stride: u64, + scratch_k: u64, + scratch_v: u64, + scratch_g: u64, + scratch_b: u64, + kv_stride: u64, + b_stride: u64, + rows: u32, + state_out: u64, + ) !void { + var args = launch_.Args{}; + args.add(@as(i32, @intCast(h))); + args.add(state_in); + args.add(state_stride); + args.add(scratch_k); + args.add(scratch_v); + args.add(scratch_g); + args.add(scratch_b); + args.add(kv_stride); + args.add(b_stride); + args.add(@as(i32, @intCast(rows))); + args.add(state_out); + try launch_.launch(self.replay_layers_fn, .{ .grid = .{ .x = layers * h, .y = 1, .z = 1 }, .block = .{ .x = 1024, .y = 1, .z = 1 } }, s.*, &args); + } + + /// conv [L, 3, C] in place takes rows keep .. keep + 2 of [old window; proj rows] (proj [R, W >= C]). + pub fn convShift( + self: *Kda, + s: *Stream, + conv: u64, + proj: u64, + layers: u32, + keep: u32, + conv_stride: i64, + proj_stride: i64, + proj_row: i64, + c: u32, + ) !void { + var args = launch_.Args{}; + args.add(conv); + args.add(proj); + args.add(@as(i32, @intCast(keep))); + args.add(conv_stride); + args.add(proj_stride); + args.add(proj_row); + args.add(@as(i32, @intCast(c))); + try launch_.launch(self.shift_fn, .{ .grid = .{ .x = layers, .y = (c + 1023) / 1024, .z = 1 }, .block = .{ .x = 1024, .y = 1, .z = 1 } }, s.*, &args); + } +}; diff --git a/zig/src/cuda/glm_kda_chunk.zig b/zig/src/cuda/glm_kda_chunk.zig new file mode 100644 index 0000000000..46fbf3a939 --- /dev/null +++ b/zig/src/cuda/glm_kda_chunk.zig @@ -0,0 +1,107 @@ +// CUDA wrappers for GLM-5.3-Flash's KDA prompt chunks in chunked (WY) form (glm5_next/cuda/kda_chunk.cu ported to + +// kernels/cuda/glm_kda_chunk.cu): one 256-thread block per head runs a prompt chunk's 32-row sub-chunks in order. + +// The 128 x 128 state lives in registers; every function takes device pointers, shapes and a stream. + +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const kernels = @import("kernels.zig"); +const stream_ = @import("stream.zig"); +const memory = @import("memory.zig"); +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +pub const sub_chunk = 32; // rows of a sub-chunk + +pub const Chunk = struct { + d: *const Driver, + m: Module, + fn_: Function, + smem: u32, + loaded: bool = false, + + fn smemBytes(dr: *const Driver, m: Module) !u32 { + const g = try m.global("tf_glm_kda_chunk_smem_v"); + var buf = try memory.DeviceBuffer.alloc(dr, 8); + defer buf.free(); + try buf.copyFrom(0, g.ptr, 8, null); + var out: [8]u8 = undefined; + try buf.download(0, &out); + return @intCast(std.mem.readInt(i64, &out, .little)); + } + + pub fn load(dr: *const Driver) !Chunk { + var c: Chunk = .{ .d = dr, .m = undefined, .fn_ = undefined, .smem = 0 }; + errdefer c.unload(); + c.m = try Module.load(dr, kernels.glm_kda_chunk); + c.fn_ = try c.m.function("tf_glm_kda_chunk"); + c.smem = try smemBytes(dr, c.m); + try c.fn_.allowDynamicShared(c.smem); + c.loaded = true; + return c; + } + + pub fn unload(self: *Chunk) void { + if (!self.loaded) return; + self.m.unload(); + self.loaded = false; + } + + // One head per block over the prompt chunk's rows: off = pos % 32 rows ahead may be zero-padded implicitly + + // (rows = the chunk's real rows, pos = the chunk's first absolute position). Zeroed out or state_out skips + + // that write. + pub fn chain( + self: *Chunk, + s: *Stream, + h: u32, + p: u64, + p_stride: i64, + b_off: i32, + a: u64, + a_stride: i64, + g: u64, + g_stride: i64, + conv_state: u64, + conv_w: u64, + state_in: u64, + a_log: u64, + dt_bias: u64, + norm_w: u64, + eps: f32, + lower: f32, + rows: u32, + off: u32, + out: u64, + state_out: u64, + ) !void { + var args = launch_.Args{}; + args.add(@as(i32, @intCast(h))); + args.add(p); + args.add(@as(i32, @intCast(p_stride))); + args.add(b_off); + args.add(a); + args.add(@as(i32, @intCast(a_stride))); + args.add(g); + args.add(@as(i32, @intCast(g_stride))); + args.add(conv_state); + args.add(conv_w); + args.add(state_in); + args.add(a_log); + args.add(dt_bias); + args.add(norm_w); + args.add(eps); + args.add(lower); + args.add(@as(i32, @intCast(rows))); + args.add(@as(i32, @intCast(off))); + args.add(out); + args.add(state_out); + try launch_.launch(self.fn_, .{ .grid = .{ .x = h, .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 }, .shared = self.smem }, s.*, &args); + } +}; diff --git a/zig/src/cuda/glm_latent.zig b/zig/src/cuda/glm_latent.zig new file mode 100644 index 0000000000..f06706f3ae --- /dev/null +++ b/zig/src/cuda/glm_latent.zig @@ -0,0 +1,176 @@ +//! GLM-5.3-Flash latent attention on CUDA: cache write, kv_b absorb/expand, 512-key chunks and their merge. + +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const module = @import("module.zig"); +const driver = @import("driver.zig"); +const kernels = @import("kernels.zig"); +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +pub const heads = 32; // index heads +pub const d = 256; // qk dims +pub const lw = 512; // latent width +pub const v = 256; // v dims + +/// The chunk kernels' dynamic shared memory: the block's `hbt` heads of qa (512 bf16 each), staged once. +fn qshBytes(hbt: u32) u32 { + return hbt * 512 * 2; +} + +pub const Latent = struct { + d: *const Driver, + m: Module, + ma: Module, + lat_fn: Function, + absorb_fn: Function, + expand_fn: Function, + merge_fn: Function, + chunks_fns: [6]Function, // d16_bf16, d16_fp8, d32_bf16, d32_fp8, s16_fp8, s32_fp8 + loaded: bool = false, + + pub fn load(dr: *const Driver) !Latent { + var k: Latent = .{ + .d = dr, + .m = undefined, + .ma = undefined, + .lat_fn = undefined, + .absorb_fn = undefined, + .expand_fn = undefined, + .merge_fn = undefined, + .chunks_fns = undefined, + }; + errdefer k.unload(); + k.m = try Module.load(dr, kernels.glm_latent); + k.ma = try Module.load(dr, kernels.glm_latent_att); + k.lat_fn = try k.m.function("tf_glm_lat_write"); + k.absorb_fn = try k.m.function("tf_glm_absorb_q"); + k.expand_fn = try k.m.function("tf_glm_expand_v"); + k.chunks_fns[0] = try k.ma.function("tf_glm_attn_d16_bf16"); + k.chunks_fns[1] = try k.ma.function("tf_glm_attn_d16_fp8"); + k.chunks_fns[2] = try k.ma.function("tf_glm_attn_d32_bf16"); + k.chunks_fns[3] = try k.ma.function("tf_glm_attn_d32_fp8"); + k.chunks_fns[4] = try k.ma.function("tf_glm_attn_s16_fp8"); + k.chunks_fns[5] = try k.ma.function("tf_glm_attn_s32_fp8"); + for (k.chunks_fns, 0..) |f, i| try f.allowDynamicShared(qshBytes(if (i == 2 or i == 3 or i == 5) 32 else 16)); + k.merge_fn = try k.ma.function("tf_glm_merge"); + k.loaded = true; + return k; + } + + pub fn unload(self: *Latent) void { + if (!self.loaded) return; + self.m.unload(); + self.ma.unload(); + self.loaded = false; + } + + /// Write rows into the cache slot the device-side pos names; fp8 rows get an amax/448 scale. + pub fn latWrite( + self: *Latent, + s: *Stream, + lat: u64, + lat_stride: i64, + lc: u64, + pos: u64, + sc: u64, + rows: u32, + fp8: bool, + ) !void { + var args = launch_.Args{}; + args.add(lat); + args.add(lat_stride); + args.add(lc); + args.add(pos); + args.add(sc); + args.add(@as(i32, if (fp8) 1 else 0)); + try launch_.launch(self.lat_fn, .{ .grid = .{ .x = rows, .y = 1, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s.*, &args); + } + + /// qa = q @ wk (the absorbed kv_b qk side): a head a block-x, a 32-latent col group a block-y; rows looped. + pub fn absorbQ(self: *Latent, s: *Stream, q: u64, wk: u64, qa: u64, rows: u32) !void { + var args = launch_.Args{}; + args.add(q); + args.add(wk); + args.add(qa); + args.add(@as(i32, @intCast(rows))); + try launch_.launch(self.absorb_fn, .{ .grid = .{ .x = heads, .y = 512 / 32, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s.*, &args); + } + + /// out = ol @ wv (the expanded kv_b v side): a head a block-x, a 16-v col group a block-y; rows looped. + pub fn expandV(self: *Latent, s: *Stream, ol: u64, wv: u64, out: u64, rows: u32) !void { + var args = launch_.Args{}; + args.add(ol); + args.add(wv); + args.add(out); + args.add(@as(i32, @intCast(rows))); + try launch_.launch(self.expand_fn, .{ .grid = .{ .x = heads, .y = 256 / 16, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s.*, &args); + } + + /// One 512-key chunk a block (z), a 16/32-row qk group a block (y; 32 qk rows when rows >= 64). + pub fn attnChunks( + self: *Latent, + s: *Stream, + qa: u64, + cache: u64, + sc: u64, + pos: u64, + tok: u64, + cnt: u64, + po: u64, + pm: u64, + pl: u64, + rows: u32, + tokw: i32, + nch: u32, + fp8: bool, + sparse: bool, + scale: f32, + ) !void { + const hbt: u32 = if (rows >= 64) 32 else 16; + if (sparse and !fp8) return error.SparseNeedsFp8; + const wide: u32 = if (hbt == 32) 1 else 0; + const idx: u32 = if (sparse) 4 + wide else 2 * wide + (if (fp8) @as(u32, 1) else 0); + var args = launch_.Args{}; + args.add(qa); + args.add(cache); + args.add(sc); + args.add(pos); + args.add(tok); + args.add(cnt); + args.add(po); + args.add(pm); + args.add(pl); + args.add(@as(i32, @intCast(rows))); + args.add(tokw); + args.add(scale); + try launch_.launch(self.chunks_fns[idx], .{ .grid = .{ .x = rows, .y = 32 / hbt, .z = nch }, .block = .{ .x = 256, .y = 1, .z = 1 }, .shared = qshBytes(hbt) }, s.*, &args); + } + + /// Combine the per-chunk partials into the bf16 output rows; sparse rows with no tokens keep their input. + pub fn merge( + self: *Latent, + s: *Stream, + po: u64, + pm: u64, + pl: u64, + out: u64, + cnt: u64, + rows: u32, + nch: u32, + sparse: bool, + ) !void { + var args = launch_.Args{}; + args.add(po); + args.add(pm); + args.add(pl); + args.add(out); + args.add(cnt); + args.add(@as(i32, @intCast(rows))); + args.add(@as(i32, @intCast(nch))); + args.add(@as(i32, if (sparse) 1 else 0)); + try launch_.launch(self.merge_fn, .{ .grid = .{ .x = rows, .y = heads, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s.*, &args); + } +}; diff --git a/zig/src/cuda/glm_layer.zig b/zig/src/cuda/glm_layer.zig new file mode 100644 index 0000000000..cadcb37034 --- /dev/null +++ b/zig/src/cuda/glm_layer.zig @@ -0,0 +1,409 @@ +//! One GLM-5.3-Flash layer on one rank (forward.py layer_forward): HC glue, the KDA or DSA mixer, MoE or dense MLP. +const std = @import("std"); +const driver = @import("driver.zig"); +const memory = @import("memory.zig"); +const stream_ = @import("stream.zig"); +const exl3 = @import("exl3.zig"); +const exl3e = @import("exl3_experts.zig"); +const glue_ = @import("glm_glue.zig"); +const glue2_ = @import("glm_glue2.zig"); +const bmm_ = @import("glm_bmm.zig"); +const kda_ = @import("glm_kda.zig"); +const chunk_ = @import("glm_kda_chunk.zig"); +const latent_ = @import("glm_latent.zig"); +const sparse_ = @import("glm_sparse.zig"); +const dsa_ = @import("glm_dsa.zig"); +const Stream = stream_.Stream; +const B16 = bmm_.B16; + +/// Every kernel family a layer launches, loaded once. +pub const Kernels = struct { + d: *const driver.Driver, + glue: glue_.Glue, + glue2: glue2_.Glue2, + bmm: bmm_.Bmm, + dense: exl3.Dense, + kda: kda_.Kda, + chunk: chunk_.Chunk, + experts: exl3e.Experts, + latent: latent_.Latent, + sparse: sparse_.Sparse, + + pub fn load(d: *const driver.Driver) !Kernels { + var glue = try glue_.Glue.load(d); + errdefer glue.unload(); + var glue2 = try glue2_.Glue2.load(d); + errdefer glue2.unload(); + var bmm = try bmm_.Bmm.load(d); + errdefer bmm.unload(); + var dense = try exl3.Dense.load(d); + errdefer dense.unload(); + var kda = try kda_.Kda.load(d); + errdefer kda.unload(); + var chunk = try chunk_.Chunk.load(d); + errdefer chunk.unload(); + var experts = try exl3e.Experts.load(d); + errdefer experts.unload(); + var latent = try latent_.Latent.load(d); + errdefer latent.unload(); + const sparse = try sparse_.Sparse.load(d); + return .{ .d = d, .glue = glue, .glue2 = glue2, .bmm = bmm, .dense = dense, .kda = kda, .chunk = chunk, .experts = experts, .latent = latent, .sparse = sparse }; + } + + pub fn unload(k: *Kernels) void { + k.sparse.unload(); + k.latent.unload(); + k.experts.unload(); + k.chunk.unload(); + k.kda.unload(); + k.dense.unload(); + k.bmm.unload(); + k.glue2.unload(); + k.glue.unload(); + } + + /// The DSA block's view of these kernels. + pub fn dsa(k: *Kernels) dsa_.Kernels { + return .{ .glue2 = &k.glue2, .bmm = &k.bmm, .dense = &k.dense, .latent = &k.latent, .sparse = &k.sparse }; + } +}; + +/// The model's constants a layer reads (weights.Config). +pub const Config = struct { + eps: f32, + hc_eps: f32, + hc_iters: u32, + lower: f32, + limit: f32, + routed_scale: f32, + top_k: u32, + experts: u32, + d: u32 = 4096, + streams: u32 = 4, + dsa: ?dsa_.Config = null, // a DSA layer's constants +}; + +/// One hyper-connection site: fn [24, S * D] fp32, base [24] fp32, scale [3] fp32. +pub const Hc = struct { fn_w: u64, base: u64, scale: u64 }; + +/// The rank's KDA weights: fused q|k|v EXL3, the bf16 pair and b, conv taps, decays, gated norm, o_proj's rows. +pub const KdaW = struct { + qkv: exl3.Weight, + fab: B16, // f_a | g_a | b stacked by rows: one matmul writes p's columns fa_off.. (the three are adjacent) + fb: B16, + gb: B16, + o: exl3.Weight, + conv: u64, // bf16 [3 H 128, 4] + a_log: u64, // fp32 [H] + dt_bias: u64, // fp32 [H 128] + norm: u64, // bf16 [128] + heads: u32, + fa_off: u32, // p's columns: q|k|v, then f_a, g_a, b + ga_off: u32, + b_off: u32, + p_stride: u32, +}; + +/// The rank's MoE: router [E, D] bf16 and bias [E] fp32, routed EXL3 experts, the shared expert's three linears. +pub const MoeW = struct { router: B16, bias: u64, experts: exl3e.Layer, gate: exl3.Weight, up: exl3.Weight, down: exl3.Weight }; + +/// A dense layer's MLP (layers 0-2): gate and up (the rank's columns), down (the rank's rows, an fp32 partial). +pub const MlpW = struct { gate: exl3.Weight, up: exl3.Weight, down: exl3.Weight }; + +/// The layer's mixer: the KDA block's weights or the DSA block's. +pub const Mixer = union(enum) { kda: KdaW, dsa: dsa_.DsaW }; + +/// The layer's feed-forward: routed experts plus the shared one, or a dense SwiGLU MLP. +pub const Ffn = union(enum) { moe: MoeW, mlp: MlpW }; + +pub const LayerW = struct { attn_hc: Hc, ffn_hc: Hc, in_norm: u64, post_norm: u64, mixer: Mixer, ffn: Ffn }; + +/// This layer's KDA state for one forward: the conv window [3, 3 H 128] bf16, the fp32 state read and written. +pub const KdaState = struct { conv: u64, rec_in: u64, rec_out: u64 }; + +/// Decode windows run the serial chain (commit shifts the conv later); prompt chunks the chunked form at `pos`. +pub const Mode = union(enum) { decode, prompt: u64 }; + +/// A DSA call: the layer's caches, its buffers (dsa.Scratch, sized by the engine's rows) and the host position. +pub const DsaRun = struct { st: dsa_.DsaState, sc: *const dsa_.Scratch, pos: usize, capture: bool = false }; + +/// The mixer's state for one forward, matching the layer's Mixer. +pub const MixState = union(enum) { kda: struct { st: KdaState, mode: Mode }, dsa: DsaRun }; + +pub const Site = enum { attn, ffn }; + +/// Fills gath [world, rows, D] fp32 with every rank's partial in rank order (forward.gather: NCCL, or a test's). +pub const Gather = struct { + ctx: *anyopaque, + world: u32, + run: *const fn (ctx: *anyopaque, s: Stream, site: Site, part: u64, gath: u64, rows: usize) anyerror!void, +}; + +const Lay = struct { + at: u64, + fn take(l: *Lay, bytes: usize) u64 { + const r = l.at; + l.at += (bytes + 255) / 256 * 256; + return r; + } +}; + +/// What a Scratch holds: the KDA mixer's buffers (null: none), the routed experts' (null: none), the widest SwiGLU. +pub const Shape = struct { + kda: ?*const KdaW = null, + experts: ?exl3e.Layer = null, + width: u32, + + /// One layer's needs. + pub fn of(w: *const LayerW) Shape { + return .{ + .kda = if (w.mixer == .kda) &w.mixer.kda else null, + .experts = if (w.ffn == .moe) w.ffn.moe.experts else null, + .width = switch (w.ffn) { + .moe => |m| m.gate.n, + .mlp => |m| m.gate.n, + }, + }; + } + + /// Both shapes' needs: one Scratch then serves layers of either (every KDA layer has the same heads). + pub fn widen(a: Shape, b: Shape) Shape { + return .{ .kda = a.kda orelse b.kda, .experts = a.experts orelse b.experts, .width = @max(a.width, b.width) }; + } +}; + +/// A forward's buffers for up to `rows` rows (forward.Buffers' slice for the layers of a Shape), one allocation. +pub const Scratch = struct { + buf: memory.DeviceBuffer, + rows: usize, + world: u32, + normed: u64, + xs: u64, + post: u64, + comb: u64, + hcpart: u64, + p: u64, // [rows, p_stride] bf16: the decode window's st.proj rows, the prompt chunk's b.kproj + qkv: u64, + fab: u64, // [rows, fab.n] bf16: f_a | g_a | b + ka: u64, + kg: u64, + mpart: u64, // split-K partials of the bf16 matmuls and the router (one at a time) + kout: u64, + kda: kda_.Scratch, + part: u64, + gath: u64, + mlog: u64, + pick: u64, + wts: u64, + ey: u64, // [rows, slots, D] fp32: routed slots, then the shared expert in the last + esc: ?exl3e.Scratch, // the routed experts' (null: the Shape has none) + esc_len: usize, + width: usize, // the widest SwiGLU the next buffers hold + smg: u64, + smu: u64, + sgu: u64, + sact: u64, + sxs: u64, + sy: u64, + + fn place(l: *Lay, c: Config, sh: Shape, rows: usize, world: u32) !Scratch { + const d: usize = c.d; + // a DSA layer takes no KDA buffers (their sizes are 0; dsa.Scratch holds its own); a dense MLP none of MoE's + const kd = sh.kda; + const per: usize = if (kd) |v| @as(usize, v.heads) * kda_.d_k else 0; + const heads: usize = if (kd) |v| v.heads else 0; + const p_stride: usize = if (kd) |v| v.p_stride else 0; + const slots: usize = c.top_k + 1; + const ne: usize = if (sh.experts != null) c.experts else 0; + const er: usize = if (sh.experts != null) rows else 0; + const sw: usize = sh.width; + var mp: usize = @as(usize, bmm_.router_ks) * rows * ne * 4; + if (kd) |v| for ([_]B16{ v.fab, v.fb, v.gb }) |q| { + mp = @max(mp, bmm_.Bmm.partBytes(q, rows)); + }; + var sc: Scratch = undefined; + sc.rows = rows; + sc.world = world; + sc.normed = l.take(rows * d * 2); + sc.xs = l.take(rows * d / 64 * 4); + sc.post = l.take(rows * c.streams * 4); + sc.comb = l.take(rows * c.streams * c.streams * 4); + sc.hcpart = l.take(rows * 16 * 32 * 4); + sc.p = l.take(rows * p_stride * 2); + sc.qkv = l.take(rows * @as(usize, if (kd) |v| v.qkv.n else 0) * 2); + sc.fab = l.take(rows * @as(usize, if (kd) |v| v.fab.n else 0) * 2); + sc.ka = l.take(rows * per * 2); + sc.kg = l.take(rows * per * 2); + sc.mpart = l.take(mp); + sc.kout = l.take(rows * per * 2); + sc.kda = .{ .k = l.take(rows * per * 4), .v = l.take(rows * per * 2), .g = l.take(rows * per * 4), .b = l.take(rows * heads * 4), .q = l.take(rows * per * 4), .y = l.take(rows * per * 2) }; + sc.part = l.take(rows * d * 4); + sc.gath = l.take(@as(usize, world) * rows * d * 4); + sc.mlog = l.take(rows * ne * 4); + sc.pick = l.take(er * slots * 4); + sc.wts = l.take(er * slots * 4); + sc.ey = l.take(er * slots * d * 4); + sc.esc = null; + sc.esc_len = 0; + if (sh.experts) |ex| { + sc.esc_len = try exl3e.scratchBytes(ex, rows, slots); + sc.esc = try exl3e.scratch(l.take(sc.esc_len), ex, rows, slots); + } + sc.width = sw; + sc.smg = l.take(rows * sw * 2); + sc.smu = l.take(rows * sw * 2); + sc.sgu = l.take(rows * 2 * sw * 2); + sc.sact = l.take(rows * sw * 2); + sc.sxs = l.take(rows * sw / 64 * 4); + sc.sy = l.take(er * d * 4); + if (rows == 1) { // one row: the copies into p, sgu and ey's shared slot become the outputs themselves + if (kd) |v| { + sc.qkv = sc.p; + sc.fab = sc.p + @as(u64, v.fa_off) * 2; + } + sc.smg = sc.sgu; + sc.smu = sc.sgu + sw * 2; + if (er > 0) sc.sy = sc.ey + (slots - 1) * d * 4; + } + return sc; + } + + /// Zeroed buffers (torch.zeros' experts scratch, its members list at -1); `world` ranks' partials gathered. + pub fn initFor(d: *const driver.Driver, s: Stream, c: Config, sh: Shape, rows: usize, world: u32) !Scratch { + var probe: Lay = .{ .at = 0 }; + _ = try place(&probe, c, sh, rows, world); + var buf = try memory.DeviceBuffer.alloc(d, probe.at); + errdefer buf.free(); + try buf.fill8(0, s.handle); + var l: Lay = .{ .at = buf.ptr }; + var sc = try place(&l, c, sh, rows, world); + sc.buf = buf; + if (sc.esc) |esc| { + const mb: memory.DeviceBuffer = .{ .d = d, .ptr = esc.members, .len = esc.xg + sc.esc_len - esc.members, .borrowed = true }; + try mb.fill8(0xff, s.handle); + } + return sc; + } + + /// One layer's buffers (Shape.of). + pub fn init(d: *const driver.Driver, s: Stream, c: Config, w: *const LayerW, rows: usize, world: u32) !Scratch { + return initFor(d, s, c, Shape.of(w), rows, world); + } + + pub fn deinit(sc: *Scratch) void { + sc.buf.free(); + } +}; + +/// rows rows of `bytes` each from src (row stride ss bytes) to dst (row stride ds bytes): torch's strided copy_. +fn copyRows(d: *const driver.Driver, s: Stream, dst: u64, ds: usize, src: u64, ss: usize, bytes: usize, rows: usize) !void { + if (dst == src and (rows == 1 or ds == ss)) return; // a one-row Scratch's aliases + try memory.DeviceBuffer.copyRows(d, dst, ds, src, ss, bytes, rows, s.handle); +} + +/// glue.hc_pre: normed, xs, post and comb of the S streams x [rows, S * D] at one site. +pub fn hcPre(k: *Kernels, s: Stream, c: Config, h: Hc, norm: u64, x: u64, sc: *const Scratch, rows: usize) !void { + try k.glue.hcPartial(s, x, h.fn_w, sc.hcpart, rows); + try k.glue.hcFinish(s, x, sc.hcpart, h.base, h.scale, norm, sc.normed, sc.xs, sc.post, sc.comb, c.eps, c.hc_eps, c.hc_iters, rows, c.d); +} + +/// kda_block: p's projections, the chain into kout and rec_out (a prompt chunk commits its conv), o_proj's partial. +pub fn kdaBlock(k: *Kernels, s: *Stream, c: Config, w: *const KdaW, st: KdaState, sc: *const Scratch, rows: usize, mode: Mode) !void { + if (rows == 0 or rows > sc.rows) return error.GlmLayerRows; + const ps: usize = @as(usize, w.p_stride) * 2; + try k.dense.prompt(s.*, w.qkv, sc.normed, .bf16, rows, sc.qkv, .bf16); + try copyRows(k.d, s.*, sc.p, ps, sc.qkv, @as(usize, w.qkv.n) * 2, @as(usize, w.qkv.n) * 2, rows); + if (w.fa_off + w.fab.n != w.p_stride) return error.GlmLayerShape; + try k.bmm.matmul(s.*, w.fab, sc.normed, c.d, rows, sc.fab, false, sc.mpart); + try copyRows(k.d, s.*, sc.p + @as(u64, w.fa_off) * 2, ps, sc.fab, @as(usize, w.fab.n) * 2, @as(usize, w.fab.n) * 2, rows); + try k.bmm.matmul(s.*, w.fb, sc.p + @as(u64, w.fa_off) * 2, w.p_stride, rows, sc.ka, false, sc.mpart); + try k.bmm.matmul(s.*, w.gb, sc.p + @as(u64, w.ga_off) * 2, w.p_stride, rows, sc.kg, false, sc.mpart); + const per: i64 = @as(i64, w.heads) * kda_.d_k; + const r32: u32 = @intCast(rows); + const b_off: i32 = @intCast(w.b_off); + switch (mode) { + .decode => if (rows >= 64) + try k.kda.chainWide(s, w.heads, sc.p, w.p_stride, b_off, sc.ka, per, sc.kg, per, st.conv, w.conv, st.rec_in, w.a_log, w.dt_bias, w.norm, c.eps, c.lower, r32, sc.kout, st.rec_out, sc.kda) + else + try k.kda.chain(s, w.heads, sc.p, w.p_stride, b_off, sc.ka, per, sc.kg, per, st.conv, w.conv, st.rec_in, w.a_log, w.dt_bias, w.norm, c.eps, c.lower, r32, sc.kout, st.rec_out, sc.kda), + .prompt => |pos| { + const off: u32 = @intCast(pos % chunk_.sub_chunk); + try k.chunk.chain(s, w.heads, sc.p, w.p_stride, b_off, sc.ka, per, sc.kg, per, st.conv, w.conv, st.rec_in, w.a_log, w.dt_bias, w.norm, c.eps, c.lower, r32, off, sc.kout, st.rec_out); + try shiftConv(k, s, w, st, sc, rows); + }, + } + try k.dense.prompt(s.*, w.o, sc.kout, .bf16, rows, sc.part, .f32); +} + +/// _shift_conv for this layer: the window becomes rows keep .. keep + 2 of [old window; p rows]. +fn shiftConv(k: *Kernels, s: *Stream, w: *const KdaW, st: KdaState, sc: *const Scratch, keep: usize) !void { + const ch: u32 = 3 * w.heads * kda_.d_k; + try k.kda.convShift(s, st.conv, sc.p, 1, @intCast(keep), 3 * @as(i64, ch), @as(i64, @intCast(sc.rows)) * w.p_stride, w.p_stride, ch); +} + +/// commit(rows, keep) after a decode window: dropped rows replayed out of rec_out, then the conv window shifted. +pub fn commit(k: *Kernels, s: *Stream, w: *const KdaW, st: KdaState, sc: *const Scratch, rows: usize, keep: usize) !void { + if (keep < 1 or keep > rows or rows > sc.rows) return error.GlmLayerKeep; + if (keep < rows) try k.kda.replay(s, w.heads, st.rec_in, sc.kda, @intCast(keep), st.rec_out); + try shiftConv(k, s, w, st, sc, keep); +} + +/// A SwiGLU MLP on normed (shared expert or dense layer): EXL3 gate and up, the fused pair, down into out fp32. +fn swigluMlp(k: *Kernels, s: Stream, c: Config, gate: exl3.Weight, up: exl3.Weight, down: exl3.Weight, sc: *const Scratch, rows: usize, out: u64) !void { + const sw: usize = gate.n; + if (sw > sc.width or up.n != sw or down.k != sw) return error.GlmLayerWidth; + try k.dense.prompt(s, gate, sc.normed, .bf16, rows, sc.smg, .bf16); + try k.dense.prompt(s, up, sc.normed, .bf16, rows, sc.smu, .bf16); + try copyRows(k.d, s, sc.sgu, 4 * sw, sc.smg, 2 * sw, 2 * sw, rows); + try copyRows(k.d, s, sc.sgu + 2 * sw, 4 * sw, sc.smu, 2 * sw, 2 * sw, rows); + try k.glue.swiglu(s, sc.sgu, sc.sact, sc.sxs, c.limit, rows, sw); + try k.dense.prompt(s, down, sc.sact, .bf16, rows, out, .f32); +} + +/// moe_block on the EXL3 route: router, select, routed per-slot y, the shared expert into the last slot, combine. +pub fn moeBlock(k: *Kernels, s: Stream, c: Config, w: *const MoeW, sc: *const Scratch, rows: usize) !void { + if (rows == 0 or rows > sc.rows) return error.GlmLayerRows; + const esc = sc.esc orelse return error.GlmLayerScratch; + const d: usize = c.d; + const slots: usize = c.top_k + 1; + if (w.router.n != c.experts or w.router.k != d) return error.GlmLayerShape; + try k.bmm.router(s, sc.normed, c.d, w.router, rows, sc.mlog, sc.mpart); + try k.glue.topk(s, sc.mlog, w.bias, sc.pick, sc.wts, c.routed_scale, rows, c.experts, c.top_k); + try k.experts.runPerSlot(s, w.experts, sc.normed, d, .bf16, sc.pick, rows, sc.ey, esc); + try swigluMlp(k, s, c, w.gate, w.up, w.down, sc, rows, sc.sy); + try copyRows(k.d, s, sc.ey + (slots - 1) * d * 4, slots * d * 4, sc.sy, d * 4, d * 4, rows); + try k.glue2.combineF32(s, sc.ey, sc.wts, sc.part, rows, d, slots); +} + +/// mlp_block: a dense layer's MLP, the rank's fp32 partial into part. +pub fn mlpBlock(k: *Kernels, s: Stream, c: Config, w: *const MlpW, sc: *const Scratch, rows: usize) !void { + if (rows == 0 or rows > sc.rows) return error.GlmLayerRows; + try swigluMlp(k, s, c, w.gate, w.up, w.down, sc, rows, sc.part); +} + +/// hc_post of the gathered partials (rank 0's first) into x in place. +fn post(k: *Kernels, s: Stream, c: Config, x: u64, sc: *const Scratch, g: Gather, site: Site, rows: usize) !void { + if (g.world != sc.world) return error.GlmLayerWorld; + try g.run(g.ctx, s, site, sc.part, sc.gath, rows); + try k.glue.hcPost(s, x, x, sc.gath, sc.post, sc.comb, @intCast(rows * c.d), g.world, rows, c.d, c.streams); +} + +/// layer_forward: x [rows, S * D] bf16 in place; a KDA state per its mode (a decode window commits later), or DSA. +pub fn forward(k: *Kernels, s: *Stream, c: Config, w: *const LayerW, mix: MixState, sc: *const Scratch, x: u64, rows: usize, g: Gather) !void { + if (!std.mem.eql(u8, @tagName(w.mixer), @tagName(mix))) return error.GlmLayerMixer; + try hcPre(k, s.*, c, w.attn_hc, w.in_norm, x, sc, rows); + switch (mix) { + .kda => |m| try kdaBlock(k, s, c, &w.mixer.kda, m.st, sc, rows, m.mode), + .dsa => |m| { + const dc = c.dsa orelse return error.GlmLayerConfig; + try dsa_.dsaBlock(k.dsa(), s, dc, &w.mixer.dsa, m.st, m.sc, sc.normed, sc.part, rows, m.pos, m.capture); + }, + } + try post(k, s.*, c, x, sc, g, .attn, rows); + try hcPre(k, s.*, c, w.ffn_hc, w.post_norm, x, sc, rows); + switch (w.ffn) { + .moe => |*m| try moeBlock(k, s.*, c, m, sc, rows), + .mlp => |*m| try mlpBlock(k, s.*, c, m, sc, rows), + } + try post(k, s.*, c, x, sc, g, .ffn, rows); +} diff --git a/zig/src/cuda/glm_sparse.zig b/zig/src/cuda/glm_sparse.zig new file mode 100644 index 0000000000..cbecf88651 --- /dev/null +++ b/zig/src/cuda/glm_sparse.zig @@ -0,0 +1,183 @@ +// CUDA wrappers for GLM-5.3-Flash's DSA sparse indexer (glm5_next/cuda/sparse.py ported to kernels/cuda/glm_sparse.cu). + +// The index update writes LayerNorm'd keys and gate rows into the caches and re-pools completed 4-token pools; the + +// scores kernel ranks every visible pool per row; the select runs a radix select a row a block; tokens and counts + +// assemble each row's attended tokens. Every function takes device pointers, shapes and a stream. + +const std = @import("std"); + +const launch_ = @import("launch.zig"); + +const stream_ = @import("stream.zig"); + +const module = @import("module.zig"); + +const driver = @import("driver.zig"); + +const kernels = @import("kernels.zig"); + +const Driver = driver.Driver; + +const Module = module.Module; + +const Function = module.Function; + +const Stream = stream_.Stream; + +pub const d = 128; // index key dims + +pub const heads = 32; // index heads (index_heads) + +pub const pool = 4; // tokens a pool + +pub const topk_pools = 512; // pools a row keeps + +pub const width = topk_pools * pool + pool - 1; // tokens a row keeps (2051) + +pub const score_loop = 1; // pool blocks a scores program takes (the kernel's L; speed only) + +pub const block_pools = 64; // pools a scores block + +pub const select_block = 1024; // scores a select pass reads a chunk + +pub const Sparse = struct { + d: *const Driver, + m: Module, + index_fn: Function, + pool_fn: Function, + scores_fn: Function, + select_fn: Function, + tokens_fn: Function, + loaded: bool = false, + + pub fn load(dr: *const Driver) !Sparse { + var k: Sparse = .{ + .d = dr, + .m = undefined, + .index_fn = undefined, + .pool_fn = undefined, + .scores_fn = undefined, + .select_fn = undefined, + .tokens_fn = undefined, + }; + errdefer k.unload(); + k.m = try Module.load(dr, kernels.glm_sparse); + k.index_fn = try k.m.function("tf_glm_index_write"); + k.pool_fn = try k.m.function("tf_glm_pool_keys"); + k.scores_fn = try k.m.function("tf_glm_scores"); + k.select_fn = try k.m.function("tf_glm_select_rows"); + k.tokens_fn = try k.m.function("tf_glm_tokens_counts"); + k.loaded = true; + return k; + } + + pub fn unload(self: *Sparse) void { + if (!self.loaded) return; + self.m.unload(); + self.loaded = false; + } + + /// Row writes: kr [R, d] bf16 rows k_stride apart, gr [R, d] fp32 -> ik/ig rows pos.. (bf16). + pub fn indexWrite( + self: *Sparse, + s: *Stream, + kr: u64, + k_stride: i64, + gr: u64, + lnw: u64, + lnb: u64, + ik: u64, + ig: u64, + pos: i32, + posp: u64, // the position on the device instead (0: pos) + eps: f32, + rows: u32, + ) !void { + var args = launch_.Args{}; + args.add(kr); + args.add(k_stride); + args.add(gr); + args.add(lnw); + args.add(lnb); + args.add(ik); + args.add(ig); + args.add(pos); + args.add(eps); + args.add(posp); + try launch_.launch(self.index_fn, .{ .grid = .{ .x = rows, .y = 1, .z = 1 }, .block = .{ .x = 32, .y = 1, .z = 1 } }, s.*, &args); + } + + /// Pool keys for the pools the window completes: pk rows (pos / 4).., capped by the window's last full pool. + pub fn poolKeys(self: *Sparse, s: *Stream, ik: u64, ig: u64, ape: u64, pk: u64, pos: i32, posp: u64, rows: u32) !void { + var args = launch_.Args{}; + args.add(ik); + args.add(ig); + args.add(ape); + args.add(pk); + args.add(pos); + args.add(@as(i32, @intCast(rows))); + args.add(posp); + try launch_.launch(self.pool_fn, .{ .grid = .{ .x = rows / 4 + 2, .y = 1, .z = 1 }, .block = .{ .x = 32, .y = 1, .z = 1 } }, s.*, &args); + } + + /// Pool scores: qi rows r.. (qi_stride a row), weights w (w_stride a row), pos first row's position; out f32. + pub fn scores( + self: *Sparse, + s: *Stream, + qi: u64, + qi_stride: i64, + w: u64, + w_stride: i64, + pk: u64, + out: u64, + pos: i32, + posp: u64, // the position on the device instead (0: pos) + rows: u32, + np: u32, + scale: f32, + wscale: f32, + ) !void { + var args = launch_.Args{}; + args.add(qi); + args.add(qi_stride); + args.add(w); + args.add(w_stride); + args.add(pk); + args.add(out); + args.add(pos); + args.add(@as(i32, @intCast(np))); + args.add(scale); + args.add(wscale); + args.add(posp); + const blocks_y = (np + block_pools * score_loop - 1) / (block_pools * score_loop); + try launch_.launch(self.scores_fn, .{ .grid = .{ .x = rows, .y = @intCast(blocks_y), .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s.*, &args); + } + + /// Radix top-pool select per row: visible bounds the per-row pool limit; out [rows, k] int64 pools ascending. + pub fn selectRows(self: *Sparse, s: *Stream, sc: u64, np: u32, out: u64, pos: i32, posp: u64, k: u32, rows: u32, visible: bool) !void { + var args = launch_.Args{}; + args.add(sc); + args.add(@as(i64, @intCast(np))); + args.add(@as(i64, @intCast(k))); + args.add(out); + args.add(@as(i32, @intCast(np))); + args.add(pos); + args.add(@as(i32, @intCast(k))); + args.add(@as(i32, if (visible) 1 else 0)); + args.add(posp); + try launch_.launch(self.select_fn, .{ .grid = .{ .x = rows, .y = 1, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s.*, &args); + } + + /// Tokens/counts from pools: tokens [rows, 2051] i32 (pools ascending then the partial tail), counts [rows]. + pub fn tokensCounts(self: *Sparse, s: *Stream, pools: u64, tokens: u64, counts: u64, pos: i32, posp: u64, rows: u32) !void { + var args = launch_.Args{}; + args.add(pools); + args.add(tokens); + args.add(counts); + args.add(pos); + args.add(posp); + try launch_.launch(self.tokens_fn, .{ .grid = .{ .x = rows, .y = 1, .z = 1 }, .block = .{ .x = 512, .y = 1, .z = 1 } }, s.*, &args); + } +}; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..2a430aea93 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -33,6 +33,16 @@ 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 exl3_linear: []const u8 = if (available) &Blob("fatbin_exl3_linear").bytes else &.{}; +pub const exl3_experts: []const u8 = if (available) &Blob("fatbin_exl3_experts").bytes else &.{}; +pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; +pub const glm_kda_chunk: []const u8 = if (available) &Blob("fatbin_glm_kda_chunk").bytes else &.{}; +pub const glm_bmm: []const u8 = if (available) &Blob("fatbin_glm_bmm").bytes else &.{}; +pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes else &.{}; +pub const glm_glue2: []const u8 = if (available) &Blob("fatbin_glm_glue2").bytes else &.{}; +pub const glm_sparse: []const u8 = if (available) &Blob("fatbin_glm_sparse").bytes else &.{}; +pub const glm_latent: []const u8 = if (available) &Blob("fatbin_glm_latent").bytes else &.{}; +pub const glm_latent_att: []const u8 = if (available) &Blob("fatbin_glm_latent_att").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/memory.zig b/zig/src/cuda/memory.zig index 2c469bae2c..eb50d8e7cd 100644 --- a/zig/src/cuda/memory.zig +++ b/zig/src/cuda/memory.zig @@ -127,6 +127,13 @@ pub const DeviceBuffer = struct { } } + /// `rows` rows of `width` bytes from src (row pitch `sp`) to dst (row pitch `dp`) in one strided copy. + pub fn copyRows(d: *const Driver, dst: abi.DevicePtr, dp: usize, src: abi.DevicePtr, sp: usize, width: usize, rows: usize, stream: abi.Stream) Error!void { + if (rows == 0 or width == 0) return; + const m: abi.Memcpy2D = .{ .src = src, .src_pitch = sp, .dst = dst, .dst_pitch = dp, .width = width, .height = rows }; + try d.check(d.api.cuMemcpy2DAsync_v2(&m, stream), "cuMemcpy2DAsync"); + } + pub fn fill8(self: DeviceBuffer, value: u8, stream: ?abi.Stream) Error!void { if (self.len == 0) return; if (stream) |s| { diff --git a/zig/src/cuda/nccl_id.zig b/zig/src/cuda/nccl_id.zig new file mode 100644 index 0000000000..9226cea193 --- /dev/null +++ b/zig/src/cuda/nccl_id.zig @@ -0,0 +1,115 @@ +//! Two ranks' NCCL communicator: rank 0 serves its 128-byte unique id on host:port, rank 1 connects (started first). +const std = @import("std"); +const posix = std.posix; +const nccl = @import("nccl.zig"); + +/// The pair's communicator for `rank` (0 or 1); the id crosses the cable over one TCP connection. +pub fn pair(lib: *const nccl.Library, io: std.Io, rank: u8, addr: []const u8) !nccl.Comm { + if (rank > 1) return error.BadRank; + var id: nccl.UniqueId = undefined; + const bytes = std.mem.asBytes(&id); + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try serve(addr, bytes); + } else { + try receive(io, addr, bytes); + } + var comm: nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, 2, id, @intCast(rank)), "ncclCommInitRank"); + return comm; +} + +/// Rank 0 binds host:port and sends its unique id to the one peer (70 s for rank 1's retries). +fn serve(addr_text: []const u8, bytes: []const u8) !void { + const a = try parseAddr(addr_text); + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = try posixLen(a.host, a.port, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const fd: posix.socket_t = @intCast(rc); + defer _ = posix.system.close(fd); + const one: c_int = 1; + try posix.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, std.mem.asBytes(&one)); + switch (posix.errno(posix.system.bind(fd, &storage.any, len))) { + .SUCCESS => {}, + .ADDRINUSE => return error.AddressInUse, + else => return error.BindFailed, + } + if (posix.errno(posix.system.listen(fd, 1)) != .SUCCESS) return error.ListenFailed; + var waited: u64 = 0; + while (true) { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, 500) catch return error.AcceptFailed; + if (ready > 0) break; + waited += 500; + if (waited > 70_000) return error.PeerTimeout; + } + var peer: std.Io.Threaded.PosixAddress = undefined; + var peer_len: posix.socklen_t = @sizeOf(std.Io.Threaded.PosixAddress); + const rc2 = posix.system.accept(fd, &peer.any, &peer_len); + if (posix.errno(rc2) != .SUCCESS) return error.AcceptFailed; + const cfd: posix.socket_t = @intCast(rc2); + defer _ = posix.system.close(cfd); + var off: usize = 0; + while (off < bytes.len) { + const n = posix.system.write(cfd, bytes.ptr + off, bytes.len - off); + if (posix.errno(n) != .SUCCESS) return error.WriteFailed; + off += @intCast(n); + } +} + +/// Rank 1 connects to host:port, retrying for 60 s (rank 0 may not have bound yet), and reads the id. +fn receive(io: std.Io, addr_text: []const u8, bytes: []u8) !void { + const a = try parseAddr(addr_text); + const ip = try std.Io.net.IpAddress.parseIp4(a.host, a.port); + var fd: posix.socket_t = -1; + defer { + if (fd >= 0) _ = posix.system.close(fd); + } + var waited: u64 = 0; + while (fd < 0) { + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&ip, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const c: posix.socket_t = @intCast(rc); + if (posix.errno(posix.system.connect(c, &storage.any, len)) == .SUCCESS) { + fd = c; + } else { + _ = posix.system.close(c); + waited += 200; + if (waited > 60_000) return error.PeerTimeout; + std.Io.sleep(io, .fromMilliseconds(200), .awake) catch {}; + } + } + var off: usize = 0; + var read_waited: u64 = 0; + while (off < bytes.len) { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, 500) catch return error.ReadFailed; + if (ready == 0) { + read_waited += 500; + if (read_waited > 60_000) return error.PeerTimeout; + continue; + } + const n = posix.system.read(fd, bytes.ptr + off, bytes.len - off); + switch (posix.errno(n)) { + .SUCCESS => { + if (n == 0) return error.PeerClosed; // the peer died before its id arrived + off += @intCast(n); + }, + .INTR => {}, + else => return error.ReadFailed, + } + } +} + +fn posixLen(host: []const u8, port: u16, storage: *std.Io.Threaded.PosixAddress) !posix.socklen_t { + const ip = try std.Io.net.IpAddress.parseIp4(host, port); + return std.Io.Threaded.addressToPosix(&ip, storage); +} + +fn parseAddr(text: []const u8) !struct { host: []const u8, port: u16 } { + const at = std.mem.lastIndexOfScalar(u8, text, ':') orelse return error.BadAddress; + return .{ .host = text[0..at], .port = try std.fmt.parseInt(u16, text[at + 1 ..], 10) }; +} diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..0ce5060684 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -7,6 +7,7 @@ const memory = @import("memory.zig"); const stream_ = @import("stream.zig"); const qmmf = @import("qmmf.zig"); const nvfp4 = @import("nvfp4.zig"); +const exl3 = @import("exl3.zig"); const Driver = driver.Driver; const Module = module.Module; @@ -36,11 +37,13 @@ pub const Weight = union(enum) { nvfp4: qmmf.Weight, fp8: qmmf.Weight, // e4m3 with one fp32 scale (ModelOpt FP8) bf16: qmmf.Weight, // unquantized: the checkpoint's bf16 values + exl3: exl3.Weight, // the trellis: decode and prompt rows through the same row-invariant kernel pub fn outputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.n, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.n, + .exl3 => |q| q.n, }; } @@ -48,6 +51,7 @@ pub const Weight = union(enum) { return switch (w) { .affine4 => |q| q.k, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.k, + .exl3 => |q| q.k, }; } }; @@ -60,6 +64,7 @@ pub const Linear = struct { affine4: ?*const Affine4Kernels = null, lane: ?*const qmmf.Lane = null, // decode rows of every lane format; fp8g and bf16 prompt rows too nvfp4_prompt: ?*const nvfp4.Prompt = null, // the prompt GEMM: nvfp4 and fp8 prompt rows (linear.py's prefill) + exl3: ?*exl3.Dense = null, // the EXL3 kernels and the scratch the dense calls share /// Decode rows: out (rows, n) bf16, each row's bits the same at every row count the format takes. pub fn decode(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { @@ -73,6 +78,11 @@ pub const Linear = struct { .nvfp4 => |q| return l.onLane(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onLane(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + .exl3 => |q| { + const e = l.exl3 orelse return error.FormatNotLoaded; + if (in.ldx != 0 and in.ldx != q.k) return error.RowStride; + return e.decode(s, q, in.x, .bf16, rows, out, .bf16); + }, } } @@ -88,6 +98,11 @@ pub const Linear = struct { .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + .exl3 => |q| { + const e = l.exl3 orelse return error.FormatNotLoaded; + if (in.ldx != 0 and in.ldx != q.k) return error.RowStride; + return e.prompt(s, q, in.x, .bf16, rows, out, .bf16); + }, } } @@ -309,6 +324,16 @@ test "a weight view reports its shape in every format" { try std.testing.expectEqual(@as(usize, 2560), f.inputs()); } +test "the exl3 weight view reports its shape and is refused when not loaded" { + const w: Weight = .{ .exl3 = .{ .trellis = 0, .suh = 0, .svh = 0, .k = 4096, .n = 1024, .halves = 8, .codebook = .mul1, .stride_k = 4096, .stride_nb = 256 } }; + try std.testing.expectEqual(@as(usize, 1024), w.outputs()); + try std.testing.expectEqual(@as(usize, 4096), w.inputs()); + const l: Linear = .{}; + const s: Stream = undefined; // never reached: the refusal comes first + try std.testing.expectError(error.FormatNotLoaded, l.decode(s, w, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, w, .{ .x = 0 }, 0, 1)); +} + test "a format that is not loaded is refused before any launch" { const l: Linear = .{}; const w: Weight = .{ .affine4 = .{ .w = 0, .s = 0, .b = 0, .n = 64, .k = 64, .npad = 64 } }; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..1b558498ef 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -21,6 +21,7 @@ pub const Dim3 = launch.Dim3; pub const graph = @import("graph.zig"); pub const cublaslt = @import("cublaslt.zig"); pub const nccl = @import("nccl.zig"); +pub const nccl_id = @import("nccl_id.zig"); pub const triton = @import("triton.zig"); pub const aot = @import("aot.zig"); pub const kernels = @import("kernels.zig"); @@ -36,6 +37,19 @@ pub const embed_rows = @import("embed_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); +pub const exl3 = @import("exl3.zig"); +pub const exl3_experts = @import("exl3_experts.zig"); +pub const exl3_upload = @import("exl3_upload.zig"); +pub const glm_kda = @import("glm_kda.zig"); +pub const glm_kda_chunk = @import("glm_kda_chunk.zig"); +pub const glm_bmm = @import("glm_bmm.zig"); +pub const glm_glue = @import("glm_glue.zig"); +pub const glm_glue2 = @import("glm_glue2.zig"); +pub const glm_sparse = @import("glm_sparse.zig"); +pub const glm_latent = @import("glm_latent.zig"); +pub const glm_layer = @import("glm_layer.zig"); +pub const glm_head = @import("glm_head.zig"); +pub const glm_dsa = @import("glm_dsa.zig"); test { _ = @import("memory.zig"); @@ -53,5 +67,18 @@ test { _ = @import("fp8.zig"); _ = carveout; _ = @import("nvfp4.zig"); + _ = exl3; + _ = exl3_experts; + _ = exl3_upload; + _ = glm_kda; + _ = glm_kda_chunk; + _ = glm_bmm; + _ = glm_glue; + _ = glm_glue2; + _ = glm_sparse; + _ = glm_latent; + _ = glm_layer; + _ = glm_head; + _ = glm_dsa; _ = @import("modelopt.zig"); } diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig new file mode 100644 index 0000000000..2307c34167 --- /dev/null +++ b/zig/src/families/glm/cuda.zig @@ -0,0 +1,11 @@ +//! GLM-5.3-Flash on CUDA in Zig at TP=2: the rank loader over the unsplit EXL3 checkpoint, the engine, greedy replies. + +pub const Config = @import("config.zig").Config; +pub const load = @import("cuda_load.zig"); +pub const Weights = @import("cuda_model.zig").Weights; +pub const Engine = @import("cuda_model.zig").Engine; + +test { + _ = @import("cuda_load.zig"); + _ = @import("cuda_model.zig"); +} diff --git a/zig/src/families/glm/cuda_load.zig b/zig/src/families/glm/cuda_load.zig new file mode 100644 index 0000000000..e30b914a6f --- /dev/null +++ b/zig/src/families/glm/cuda_load.zig @@ -0,0 +1,496 @@ +//! GLM-5.3-Flash's weights for one of two ranks from the unsplit EXL3 checkpoint (weights.py load, split.py cuts). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const config = @import("config.zig"); +const lay = cuda.glm_layer; +const dsa = cuda.glm_dsa; +const head_ = cuda.glm_head; +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const B16 = cuda.glm_bmm.B16; +const Tensor = core.safetensors.Tensor; + +pub const world = 2; +const chunk_bytes = 1 << 30; + +/// Device memory: small tensors packed in 1 GiB chunks (256-byte aligned), large ones allocated alone. +pub const Pool = struct { + d: *const cuda.Driver, + gpa: std.mem.Allocator, + chunks: std.ArrayList(cuda.DeviceBuffer) = .empty, + at: u64 = 0, + end: u64 = 0, + bytes: usize = 0, + + pub fn deinit(p: *Pool) void { + for (p.chunks.items) |*b| b.free(); + p.chunks.deinit(p.gpa); + } + + /// `n` bytes of device memory, unwritten; a span past 1/16 of a chunk gets its own allocation (no stranded tail). + pub fn reserve(p: *Pool, n: usize) !u64 { + const len = std.mem.alignForward(usize, @max(n, 1), 256); + if (len > chunk_bytes / 16) { + try p.chunks.ensureUnusedCapacity(p.gpa, 1); + const b = try cuda.DeviceBuffer.alloc(p.d, len); + p.chunks.appendAssumeCapacity(b); + p.bytes += n; + return b.ptr; + } + if (p.at + len > p.end) { + try p.chunks.ensureUnusedCapacity(p.gpa, 1); + const b = try cuda.DeviceBuffer.alloc(p.d, chunk_bytes); + p.chunks.appendAssumeCapacity(b); + p.at = b.ptr; + p.end = b.ptr + b.len; + } + const ptr = p.at; + p.at += len; + p.bytes += n; + return ptr; + } + + pub fn write(p: *Pool, ptr: u64, bytes: []const u8) !void { + const view: cuda.DeviceBuffer = .{ .d = p.d, .ptr = ptr, .len = bytes.len, .borrowed = true }; + try view.upload(0, bytes); + } + + pub fn put(p: *Pool, bytes: []const u8) !u64 { + const ptr = try p.reserve(bytes.len); + try p.write(ptr, bytes); + return ptr; + } +}; + +/// The bytes of rows [rank n / 2, (rank + 1) n / 2) of a tensor's first axis (split.py's "row" cut). +pub fn rowHalf(t: Tensor, rank: u32) ![]const u8 { + if (t.rank == 0 or t.shape[0] % world != 0) return error.GlmCut; + const half = t.bytes.len / world; + return t.bytes[rank * half ..][0..half]; +} + +/// Axis 1 of [outer, n, inner bytes] in `parts` blocks, each block's rank half in order: dim1 (1) or fused runs (3). +pub fn cutAxis1(a: std.mem.Allocator, src: []const u8, outer: usize, n: usize, inner: usize, parts: usize, rank: u32) ![]u8 { + const out = try a.alloc(u8, src.len / world); + try cutInto(out, src, outer, n, inner, parts, rank); + return out; +} + +/// cutAxis1 into `out` (src.len / 2 bytes). +pub fn cutInto(out: []u8, src: []const u8, outer: usize, n: usize, inner: usize, parts: usize, rank: u32) !void { + if (n % (parts * world) != 0 or src.len != outer * n * inner or out.len * world != src.len) return error.GlmCut; + const per = n / parts; + const half = per / world; + for (0..outer) |o| for (0..parts) |j| { + @memcpy(out[(o * parts + j) * half * inner ..][0 .. half * inner], src[(o * n + j * per + rank * half) * inner ..][0 .. half * inner]); + }; +} + +/// Trellis words [K/16, N/16, W] -> [N/128, K/16, 8, W] (linear.py strips): each 128-column block's tiles in k order. +pub fn strips(a: std.mem.Allocator, src: []const u8, kt: usize, nt: usize, wb: usize) ![]u8 { + if (nt % 8 != 0 or src.len != kt * nt * wb) return error.GlmCut; + const out = try a.alloc(u8, src.len); + for (0..nt / 8) |jb| for (0..kt) |i| { + @memcpy(out[(jb * kt + i) * 8 * wb ..][0 .. 8 * wb], src[(i * nt + jb * 8) * wb ..][0 .. 8 * wb]); + }; + return out; +} + +fn bf16Bits(v: f32) u16 { + const u: u32 = @bitCast(v); + return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +/// A tensor's values as bf16 bytes (torch's .to(bfloat16): round to nearest even); bf16 bytes pass through. +pub fn toBf16(a: std.mem.Allocator, dtype: core.safetensors.DType, bytes: []const u8) ![]const u8 { + switch (dtype) { + .bf16 => return bytes, + .f16 => { + const out = try a.alloc(u16, bytes.len / 2); + for (out, 0..) |*o, i| o.* = bf16Bits(@as(f16, @bitCast(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)))); + return std.mem.sliceAsBytes(out); + }, + .f32 => { + const out = try a.alloc(u16, bytes.len / 4); + for (out, 0..) |*o, i| o.* = bf16Bits(@bitCast(std.mem.readInt(u32, bytes[4 * i ..][0..4], .little))); + return std.mem.sliceAsBytes(out); + }, + else => return error.GlmDtype, + } +} + +/// How an EXL3 group's tensors split (split.py EXL3_RULES): whole, output columns, input rows, or fused q|k|v runs. +pub const Cut = enum { rep, out, in, qkv }; + +/// One rank's reads: headers mapped, bytes read past the page cache (GB10's allocations do not reclaim it). +pub const Reader = struct { + ck: core.Checkpoint, + rank: u32, + tmp: std.heap.ArenaAllocator, + pool: *Pool, + files: []?core.direct_io.File, // ck.files' direct handles, opened on first read + + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, rank: u32, pool: *Pool) !Reader { + var ck = try core.Checkpoint.openModel(gpa, io, dir); + errdefer ck.close(); + const files = try gpa.alloc(?core.direct_io.File, ck.files.items.len); + @memset(files, null); + return .{ .ck = ck, .rank = rank, .tmp = .init(gpa), .pool = pool, .files = files }; + } + + pub fn close(r: *Reader) void { + for (r.files) |*f| if (f.*) |*x| x.close(); + r.ck.gpa.free(r.files); + r.tmp.deinit(); + r.ck.close(); + } + + fn a(r: *Reader) std.mem.Allocator { + return r.tmp.allocator(); + } + + /// `p` ++ `name`'s header (p: the layer's prefix), its dtype checked; the bytes are not read. + fn head(r: *Reader, p: []const u8, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const full = try std.fmt.allocPrint(r.a(), "{s}{s}", .{ p, name }); + const x = try r.ck.get(full); + if (x.dtype != dtype) { + std.log.err("{s}: {t}, expected {t}", .{ full, x.dtype, dtype }); + return error.GlmDtype; + } + return x; + } + + /// `head` with the bytes read into the arena. + fn t(r: *Reader, p: []const u8, name: []const u8, dtype: core.safetensors.DType) !Tensor { + var x = try r.head(p, name, dtype); + var buf: core.direct_io.Buffer = &.{}; + x.bytes = try r.readInto(&buf, x); + return x; + } + + /// A tensor's bytes through `buf` (grown in the arena when short); valid until `buf`'s next read. + fn readInto(r: *Reader, buf: *core.direct_io.Buffer, x: Tensor) ![]const u8 { + const p = @intFromPtr(x.bytes.ptr); + for (r.ck.files.items, r.files) |*f, *h| { + const base = @intFromPtr(f.map.memory.ptr); + if (p < base or p + x.bytes.len > base + f.map.memory.len) continue; + if (h.* == null) h.* = try core.direct_io.File.open(f.path); + const off = p - base; + const need = core.direct_io.span(off, x.bytes.len); + if (buf.len < need) buf.* = try r.a().alignedAlloc(u8, .fromByteUnits(core.direct_io.alignment), need); + return h.*.?.read(buf.*, off, x.bytes.len); + } + return error.GlmTensorFile; + } + + fn put(r: *Reader, bytes: []const u8) !u64 { + return r.pool.put(bytes); + } + + /// A plain tensor whole, converted to bf16. + fn bf(r: *Reader, p: []const u8, name: []const u8, dtype: core.safetensors.DType) !u64 { + const x = try r.t(p, name, dtype); + return r.put(try toBf16(r.a(), x.dtype, x.bytes)); + } + + /// A plain matrix [n, k] as the bf16 matmul's weight; `rows`: the rank's half of its rows. + fn b16(r: *Reader, p: []const u8, name: []const u8, rows: bool) !B16 { + const x = try r.t(p, name, .f16); + const bytes = if (rows) try rowHalf(x, r.rank) else x.bytes; + const n = if (rows) x.shape[0] / world else x.shape[0]; + return r.packed_(try toBf16(r.a(), x.dtype, bytes), n, x.shape[1]); + } + + /// bf16 [n, k] row-major uploaded in the matmul's fragment order. + fn packed_(r: *Reader, bytes: []const u8, n: usize, k: usize) !B16 { + const out = try r.a().alloc(u8, bytes.len); + try cuda.glm_bmm.pack(out, bytes, n, k); + return .{ .w = try r.put(out), .n = @intCast(n), .k = @intCast(k), .frag = true }; + } + + /// One EXL3 group (Exl3Linear.from_tensors, strips layout) cut for this rank. + fn linear(r: *Reader, p: []const u8, group: []const u8, cut: Cut) !exl3.Weight { + const g = try std.fmt.allocPrint(r.a(), "{s}{s}", .{ p, group }); + _ = try r.head(g, ".mul1", .i32); + const tr = try r.t(g, ".trellis", .i16); + var suh = (try r.t(g, ".suh", .f16)).bytes; + const svh_t = try r.t(g, ".svh", .f16); + var svh = svh_t.bytes; + if (tr.rank != 3 or tr.shape[2] % 8 != 0) return error.GlmTrellis; + var kt = tr.shape[0]; + var nt = tr.shape[1]; + const wb = tr.shape[2] * 2; + var words = tr.bytes; + switch (cut) { + .rep => {}, + .out, .qkv => { + const parts: usize = if (cut == .qkv) 3 else 1; + words = try cutAxis1(r.a(), words, kt, nt, wb, parts, r.rank); + svh = if (cut == .qkv) try cutAxis1(r.a(), svh, 1, svh.len / 2, 2, 3, r.rank) else try rowHalf(svh_t, r.rank); + nt /= world; + }, + .in => { + words = try rowHalf(tr, r.rank); + suh = suh[r.rank * suh.len / world ..][0 .. suh.len / world]; + kt /= world; + }, + } + if (kt % 8 != 0 or nt % 8 != 0 or suh.len != kt * 32 or svh.len != nt * 32) return error.GlmTrellis; + const k = 16 * kt; + const n = 16 * nt; + const halves: u16 = @intCast(tr.shape[2] / 8); + const st = exl3.strides(k, n, halves, true); + return .{ + .trellis = try r.put(try strips(r.a(), words, kt, nt, wb)), + .suh = try r.put(suh), + .svh = try r.put(svh), + .k = @intCast(k), + .n = @intCast(n), + .halves = halves, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; + } + + fn hc(r: *Reader, p: []const u8, site: []const u8) !lay.Hc { + const n = try std.fmt.allocPrint(r.a(), "hc_{s}_", .{site}); + const f = struct { + fn get(rd: *Reader, pre: []const u8, base: []const u8, field: []const u8) !u64 { + const full = try std.fmt.allocPrint(rd.a(), "{s}{s}", .{ base, field }); + return rd.put((try rd.t(pre, full, .f32)).bytes); + } + }.get; + return .{ .fn_w = try f(r, p, n, "fn"), .base = try f(r, p, n, "base"), .scale = try f(r, p, n, "scale") }; + } + + fn kda(r: *Reader, p: []const u8, heads: u32) !lay.KdaW { + const qkv = try r.linear(p, "self_attn.qkv_proj", .qkv); + // f_a, g_a and the rank's b rows stacked: one matmul (each column's sum is its own) + const fa = try r.t(p, "self_attn.f_a_proj.weight", .f16); + const ga = try r.t(p, "self_attn.g_a_proj.weight", .f16); + const bt = try r.t(p, "self_attn.b_proj.weight", .f16); + if (fa.shape[1] != ga.shape[1] or fa.shape[1] != bt.shape[1]) return error.GlmShape; + const b_rows = try rowHalf(bt, r.rank); + const stacked = try std.mem.concat(r.a(), u8, &.{ fa.bytes, ga.bytes, b_rows }); + const fa_n: u32 = @intCast(fa.shape[0]); + const ga_n: u32 = @intCast(ga.shape[0]); + const fab = try r.packed_(try toBf16(r.a(), .f16, stacked), fa.shape[0] + ga.shape[0] + bt.shape[0] / world, fa.shape[1]); + const conv = try r.t(p, "self_attn.conv1d.weight", .bf16); + const taps = conv.shape[2]; + const fa_off = 3 * heads * 128; + return .{ + .qkv = qkv, + .fab = fab, + .fb = try r.b16(p, "self_attn.f_b_proj.weight", true), + .gb = try r.b16(p, "self_attn.g_b_proj.weight", true), + .o = try r.linear(p, "self_attn.o_proj", .in), + .conv = try r.put(try cutAxis1(r.a(), conv.bytes, 1, conv.shape[0], taps * 2, 3, r.rank)), + .a_log = try r.put(try rowHalf(try r.t(p, "self_attn.A_log", .f32), r.rank)), + .dt_bias = try r.put(try rowHalf(try r.t(p, "self_attn.dt_bias", .f32), r.rank)), + .norm = try r.put((try r.t(p, "self_attn.o_norm.weight", .bf16)).bytes), + .heads = heads, + .fa_off = fa_off, + .ga_off = fa_off + fa_n, + .b_off = fa_off + fa_n + ga_n, + .p_stride = qkv.n + fab.n, + }; + } + + fn dsaW(r: *Reader, p: []const u8, c: *const config.Config) !dsa.DsaW { + // kv_b's rank rows, head by head: qk_dim key rows then v_dim value rows (latent.AbsorbW [H, dim, 512] each) + const kvb = try r.t(p, "self_attn.kv_b_proj.weight", .f16); + const mine = try toBf16(r.a(), .f16, try rowHalf(kvb, r.rank)); + const lw: usize = c.kv_lora; + const heads: usize = c.mla_heads / world; + const qk: usize = c.nope; + const vd: usize = c.v_dim; + if (kvb.shape[1] != lw or mine.len != heads * (qk + vd) * lw * 2) return error.GlmShape; + const wk = try r.a().alloc(u8, heads * qk * lw * 2); + const wv = try r.a().alloc(u8, heads * vd * lw * 2); + for (0..heads) |h| { + const at = mine[h * (qk + vd) * lw * 2 ..]; + @memcpy(wk[h * qk * lw * 2 ..][0 .. qk * lw * 2], at[0 .. qk * lw * 2]); + @memcpy(wv[h * vd * lw * 2 ..][0 .. vd * lw * 2], at[qk * lw * 2 ..][0 .. vd * lw * 2]); + } + const wkx = try r.t(p, "self_attn.indexer.wk.weight", .f16); + const wpx = try r.t(p, "self_attn.indexer.weights_proj.weight", .f16); + if (wkx.shape[1] != wpx.shape[1]) return error.GlmShape; + const kw = try std.mem.concat(r.a(), u8, &.{ try toBf16(r.a(), .f16, wkx.bytes), try toBf16(r.a(), .f16, wpx.bytes) }); + return .{ + .q_a = try r.linear(p, "self_attn.q_a_proj", .rep), + .kv_a = try r.linear(p, "self_attn.kv_a_proj_with_mqa", .rep), + .q_b = try r.linear(p, "self_attn.q_b_proj", .out), + .o = try r.linear(p, "self_attn.o_proj", .in), + .q_norm = try r.put((try r.t(p, "self_attn.q_a_layernorm.weight", .bf16)).bytes), + .kv_norm = try r.put((try r.t(p, "self_attn.kv_a_layernorm.weight", .bf16)).bytes), + .wk = try r.put(wk), + .wv = try r.put(wv), + .index = .{ + .kw = try r.packed_(kw, wkx.shape[0] + wpx.shape[0], wkx.shape[1]), + .qb = try r.b16(p, "self_attn.indexer.wq_b.weight", false), + .gate = try r.b16(p, "self_attn.indexer.index_kpool_compress_gate", false), + .ln_w = try r.put((try r.t(p, "self_attn.indexer.k_norm.weight", .f16)).bytes), + .ln_b = try r.put((try r.t(p, "self_attn.indexer.k_norm.bias", .f16)).bytes), + .ape = try r.bf(p, "self_attn.indexer.index_kpool_compress_ape", .f32), + }, + }; + } + + /// The routed experts (experts.prepare): each projection's rank-cut trellises packed in one span, pointers stacked. + fn experts(r: *Reader, p: []const u8, c: *const config.Config) !exl3e.Layer { + const e_n: usize = c.experts; + const dims: usize = c.hidden; + const width: usize = c.moe_inter / world; + var tables: [3]u64 = undefined; + var k2s: [3]u64 = undefined; + var scales: [3][2][]u8 = undefined; // suh, svh stacked [E, n] fp16 + inline for (.{ "gate", "up", "down" }, 0..) |proj, j| { + const down = j == 2; + const n_suh = if (down) width else dims; + const n_svh = if (down) dims else width; + scales[j] = .{ try r.a().alloc(u8, e_n * n_suh * 2), try r.a().alloc(u8, e_n * n_svh * 2) }; + const trs = try r.a().alloc(Tensor, e_n); + const k2 = try r.a().alloc(i32, e_n); + var total: usize = 0; + for (0..e_n) |e| { + const g = try std.fmt.allocPrint(r.a(), "{s}mlp.experts.{d}.{s}_proj", .{ p, e, proj }); + _ = try r.head(g, ".mul1", .i32); + const tr = try r.head(g, ".trellis", .i16); + const suh = try r.t(g, ".suh", .f16); + const svh = try r.t(g, ".svh", .f16); + if (tr.rank != 3 or tr.shape[2] % 8 != 0) return error.GlmTrellis; + k2[e] = @intCast(tr.shape[2] / 8); + trs[e] = tr; + const su = if (down) try rowHalf(suh, r.rank) else suh.bytes; + const sv = if (down) svh.bytes else try rowHalf(svh, r.rank); + if (su.len != n_suh * 2 or sv.len != n_svh * 2) return error.GlmShape; + @memcpy(scales[j][0][e * n_suh * 2 ..][0..su.len], su); + @memcpy(scales[j][1][e * n_svh * 2 ..][0..sv.len], sv); + total += tr.bytes.len / world; + } + const base = try r.pool.reserve(total); + const ptrs = try r.a().alloc(u64, e_n); + var raw: core.direct_io.Buffer = &.{}; + var buf: []u8 = &.{}; + var off: usize = 0; + for (trs, ptrs) |hd, *q| { // each trellis read, then cut, through two reused host buffers + var tr = hd; + tr.bytes = try r.readInto(&raw, hd); + const half = tr.bytes.len / world; + if (buf.len < half) buf = try r.a().alloc(u8, half); + const bytes = if (down) try rowHalf(tr, r.rank) else blk: { + try cutInto(buf[0..half], tr.bytes, tr.shape[0], tr.shape[1], tr.shape[2] * 2, 1, r.rank); + break :blk buf[0..half]; + }; + q.* = base + off; + try r.pool.write(q.*, bytes); + off += half; + } + tables[j] = try r.put(std.mem.sliceAsBytes(ptrs)); + k2s[j] = try r.put(std.mem.sliceAsBytes(k2)); + } + return .{ + .gate = tables[0], + .up = tables[1], + .down = tables[2], + .gate_k2 = k2s[0], + .up_k2 = k2s[1], + .down_k2 = k2s[2], + .suh_g = try r.put(scales[0][0]), + .suh_u = try r.put(scales[1][0]), + .svh_g = try r.put(scales[0][1]), + .svh_u = try r.put(scales[1][1]), + .suh_d = try r.put(scales[2][0]), + .svh_d = try r.put(scales[2][1]), + .experts = e_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .limit = c.swiglu_limit, + .act_mode = 0, + }; + } + + /// Layer `i` as the engine holds it (weights.load's layer()). + pub fn layer(r: *Reader, i: u32, c: *const config.Config) !lay.LayerW { + defer _ = r.tmp.reset(.retain_capacity); + const p = try std.fmt.allocPrint(r.a(), "model.language_model.layers.{d}.", .{i}); + return .{ + .attn_hc = try r.hc(p, "attn"), + .ffn_hc = try r.hc(p, "ffn"), + .in_norm = try r.put((try r.t(p, "input_layernorm.weight", .bf16)).bytes), + .post_norm = try r.put((try r.t(p, "post_attention_layernorm.weight", .bf16)).bytes), + .mixer = switch (c.kind(i)) { + .kda => .{ .kda = try r.kda(p, c.kda_heads / world) }, + .mla => .{ .dsa = try r.dsaW(p, c) }, + }, + .ffn = if (c.isMoe(i)) .{ .moe = .{ + .router = try r.b16(p, "mlp.gate.weight", false), + .bias = try r.put((try r.t(p, "mlp.gate.e_score_correction_bias", .f32)).bytes), + .experts = try r.experts(p, c), + .gate = try r.linear(p, "mlp.shared_experts.gate_proj", .out), + .up = try r.linear(p, "mlp.shared_experts.up_proj", .out), + .down = try r.linear(p, "mlp.shared_experts.down_proj", .in), + } } else .{ .mlp = .{ + .gate = try r.linear(p, "mlp.gate_proj", .out), + .up = try r.linear(p, "mlp.up_proj", .out), + .down = try r.linear(p, "mlp.down_proj", .in), + } }, + }; + } + + /// The replicated embedding table, the final norm and this rank's half of lm_head's rows. + pub fn globals(r: *Reader, c: *const config.Config) !struct { embed: u64, head: head_.HeadW } { + defer _ = r.tmp.reset(.retain_capacity); + const emb = try r.t("", "model.language_model.embed_tokens.weight", .bf16); + const lm = try r.t("", "lm_head.weight", .bf16); + if (lm.shape[0] != c.vocab or lm.shape[1] != c.hidden) return error.GlmShape; + const n: u32 = c.vocab / world; + return .{ .embed = try r.put(emb.bytes), .head = .{ + .norm = try r.put((try r.t("", "model.language_model.norm.weight", .bf16)).bytes), + .head = try r.packed_(try rowHalf(lm, r.rank), n, c.hidden), + .offset = r.rank * n, + } }; + } +}; + +test "a fused q|k|v group's rank runs and a dim1 cut" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + // axis 1 holds 12 one-byte columns: q 0-3, k 4-7, v 8-11; rank 1 takes the second half of each part + const src = [_]u8{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 }; + const q = try cutAxis1(a, &src, 1, 12, 1, 3, 1); + try std.testing.expectEqualSlices(u8, &.{ 2, 3, 6, 7, 10, 11 }, q); + // two rows of four: rank 0 keeps each row's first half + const d = try cutAxis1(a, &.{ 0, 1, 2, 3, 4, 5, 6, 7 }, 2, 4, 1, 1, 0); + try std.testing.expectEqualSlices(u8, &.{ 0, 1, 4, 5 }, d); + try std.testing.expectError(error.GlmCut, cutAxis1(a, &src, 1, 12, 1, 4, 0)); +} + +test "strips put each 128-column block's tiles in k order" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + // kt 2, nt 16 one-byte tiles: tile (i, j) holds 16 i + j + var src: [32]u8 = undefined; + for (&src, 0..) |*v, x| v.* = @intCast(x); + const out = try strips(a, &src, 2, 16, 1); + try std.testing.expectEqualSlices(u8, &.{ 0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23 }, out[0..16]); + try std.testing.expectEqualSlices(u8, &.{ 8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31 }, out[16..32]); +} + +test "f16 and f32 round to bf16 to nearest even" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + // f32 1 + 2^-8 (a tie) rounds down to 1.0; 1 + 3 * 2^-9 rounds up + const f = [_]u32{ 0x3F808000, 0x3F80C000 }; + const out = try toBf16(a, .f32, std.mem.sliceAsBytes(&f)); + try std.testing.expectEqualSlices(u16, &.{ 0x3F80, 0x3F81 }, @alignCast(std.mem.bytesAsSlice(u16, out))); + const h = [_]u16{0x3C01}; // f16 1 + 2^-10: below bf16's half step, rounds to 1.0 + const outh = try toBf16(a, .f16, std.mem.sliceAsBytes(&h)); + try std.testing.expectEqualSlices(u16, &.{0x3F80}, @alignCast(std.mem.bytesAsSlice(u16, outh))); +} diff --git a/zig/src/families/glm/cuda_model.zig b/zig/src/families/glm/cuda_model.zig new file mode 100644 index 0000000000..21e850b5e5 --- /dev/null +++ b/zig/src/families/glm/cuda_model.zig @@ -0,0 +1,336 @@ +//! GLM-5.3-Flash on CUDA at TP=2: one rank's weights, one sequence's state, prompt chunks and serial greedy steps. +const std = @import("std"); +const cuda = @import("cuda"); +const config = @import("config.zig"); +const load_ = @import("cuda_load.zig"); +const lay = cuda.glm_layer; +const dsa = cuda.glm_dsa; +const head_ = cuda.glm_head; +const nccl = cuda.nccl; + +pub const world = load_.world; +pub const prefill_rows = 2048; // decode.PREFILL_ROWS: rows of a prompt chunk +const index_eps = 1e-6; // the DSA indexer's key LayerNorm (sparse.py) + +/// One rank's weights on the device (weights.load without the MTP layer). +pub const Weights = struct { + cfg: config.Config, + pool: load_.Pool, + layers: []lay.LayerW, + embed: u64, + head: head_.HeadW, + gpa: std.mem.Allocator, + + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8, rank: u32) !Weights { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(text); + const cfg = try config.parse(gpa, text); + var pool: load_.Pool = .{ .d = d, .gpa = gpa }; + errdefer pool.deinit(); + var r = try load_.Reader.open(gpa, io, dir, rank, &pool); + defer r.close(); + const layers = try gpa.alloc(lay.LayerW, cfg.layers); + errdefer gpa.free(layers); + for (layers, 0..) |*l, i| l.* = try r.layer(@intCast(i), &cfg); + const g = try r.globals(&cfg); + return .{ .cfg = cfg, .pool = pool, .layers = layers, .embed = g.embed, .head = g.head, .gpa = gpa }; + } + + pub fn deinit(w: *Weights) void { + w.gpa.free(w.layers); + w.pool.deinit(); + } + + /// The constants every layer reads (weights.Config; the DSA scales as forward.py computes them). + pub fn layerConfig(w: *const Weights) lay.Config { + const c = &w.cfg; + return .{ + .eps = c.eps, + .hc_eps = c.hc_eps, + .hc_iters = c.sinkhorn, + .lower = c.lower_bound, + .limit = c.swiglu_limit, + .routed_scale = c.routed_scale, + .top_k = c.topk, + .experts = c.experts, + .d = c.hidden, + .streams = c.hc, + .dsa = .{ + .eps = c.eps, + .index_eps = index_eps, + .scale = @floatCast(std.math.pow(f64, @floatFromInt(c.nope), -0.5)), + .index_scale = @floatCast(std.math.pow(f64, @floatFromInt(c.i_dim), -0.5)), + .wscale = @floatCast(1.0 / @sqrt(@as(f64, @floatFromInt(c.i_heads)))), + }, + }; + } + + fn firstDsa(w: *const Weights) !*const dsa.DsaW { + for (w.layers) |*l| if (l.mixer == .dsa) return &l.mixer.dsa; + return error.GlmNoDsa; + } +}; + +/// One forward's buffers for up to `rows` rows (forward.Buffers): the layers', the DSA block's, the head's. +const Bufs = struct { + rows: usize, + sc: lay.Scratch, + dsc: dsa.Scratch, + hs: head_.Scratch, + x: cuda.DeviceBuffer, // [rows, S * D] bf16: the hyper-connection streams + ids: cuda.DeviceBuffer, // [rows] u32 + + fn init(d: *const cuda.Driver, s: cuda.Stream, w: *const Weights, rows: usize, prefill: bool, cap: usize) !Bufs { + const lc = w.layerConfig(); + var sh = lay.Shape.of(&w.layers[0]); + for (w.layers[1..]) |*l| sh = sh.widen(lay.Shape.of(l)); + var sc = try lay.Scratch.initFor(d, s, lc, sh, rows, world); + errdefer sc.deinit(); + var dsc = try dsa.Scratch.init(d, try w.firstDsa(), rows, prefill, cap); + errdefer dsc.deinit(); + var hs = try head_.Scratch.init(d, s, w.head, lc.d, rows); + errdefer hs.deinit(); + var x = try cuda.DeviceBuffer.alloc(d, rows * @as(usize, lc.streams) * lc.d * 2); + errdefer x.free(); + const ids = try cuda.DeviceBuffer.alloc(d, rows * 4); + return .{ .rows = rows, .sc = sc, .dsc = dsc, .hs = hs, .x = x, .ids = ids }; + } + + fn deinit(b: *Bufs) void { + b.ids.free(); + b.x.free(); + b.hs.deinit(); + b.dsc.deinit(); + b.sc.deinit(); + } +}; + +/// A KDA layer's state: the conv window and the recurrent state's two buffers (forward.State conv, rec). +const KdaSlot = struct { conv: u64, rec: [2]u64 }; + +/// One sequence on one rank: the weights' forward over committed caches, the two ranks in lockstep over NCCL. +pub const Engine = struct { + d: *const cuda.Driver, + w: *const Weights, + k: lay.Kernels, + s: cuda.Stream, + lc: lay.Config, + lib: *const nccl.Library, + comm: nccl.Comm, + cap: usize, + dec: Bufs, + pre: Bufs, + state: cuda.DeviceBuffer, // every layer's caches: the KDA states first (reset zeroes them), then DSA's + kda_bytes: usize, + kda: []KdaSlot, + dsa: []dsa.DsaState, + pos_dev: cuda.DeviceBuffer, // forward.State.pos_dev: i32 + cand: cuda.DeviceBuffer, // [world, 1, 2] f32: every rank's top candidate + gpa: std.mem.Allocator, + graphs: [graph_slots][2]?cuda.graph.Exec = @splat(.{ null, null }), // decode steps by window slot and parity + cur: u1 = 0, + pos: usize = 0, + + /// `cap`: cache slots; the DSA pool buckets equal the Python engine's while cap / 4 covers every bucket used. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, w: *const Weights, lib: *const nccl.Library, comm: nccl.Comm, cap: usize) !Engine { + if (cap % 4 != 0 or cap < prefill_rows) return error.GlmCapacity; + const c = &w.cfg; + var e: Engine = undefined; + e.d = d; + e.w = w; + e.lib = lib; + e.comm = comm; + e.cap = cap; + e.gpa = gpa; + e.cur = 0; + e.pos = 0; + e.graphs = @splat(.{ null, null }); + e.lc = w.layerConfig(); + e.k = try lay.Kernels.load(d); + errdefer e.k.unload(); + e.s = try cuda.Stream.init(d, false); + errdefer e.s.deinit(); + e.dec = try Bufs.init(d, e.s, w, 1, false, cap); + errdefer e.dec.deinit(); + e.pre = try Bufs.init(d, e.s, w, prefill_rows, true, cap); + errdefer e.pre.deinit(); + const n_kda = c.countKind(.kda); + const n_dsa = c.countKind(.mla); // the backbone's (the MTP layer is not loaded) + const heads: usize = c.kda_heads / world; + const conv_b = (c.conv - 1) * 3 * heads * c.kda_dim * 2; + const rec_b = heads * c.kda_dim * c.kda_dim * 4; + e.kda_bytes = n_kda * (conv_b + 2 * rec_b); + const ix: usize = c.i_dim * 2; + const dsa_b = cap * c.kv_lora + cap * 4 + 2 * cap * ix + (cap / 4 + 2) * ix; + e.state = try cuda.DeviceBuffer.alloc(d, e.kda_bytes + n_dsa * dsa_b); + errdefer e.state.free(); + try e.state.fill8(0, e.s.handle); + e.pos_dev = try cuda.DeviceBuffer.alloc(d, 4); + errdefer e.pos_dev.free(); + e.cand = try cuda.DeviceBuffer.alloc(d, world * 2 * 4); + errdefer e.cand.free(); + e.kda = try gpa.alloc(KdaSlot, n_kda); + errdefer gpa.free(e.kda); + e.dsa = try gpa.alloc(dsa.DsaState, n_dsa); + errdefer gpa.free(e.dsa); + for (w.layers) |*l| try reserve(&e.k.dense, l); + var at = e.state.ptr; + for (e.kda) |*sl| { + sl.* = .{ .conv = at, .rec = .{ at + conv_b, at + conv_b + rec_b } }; + at += conv_b + 2 * rec_b; + } + for (e.dsa) |*st| { + st.* = .{ .lc = at, .scales = at + cap * c.kv_lora, .ik = at + cap * (c.kv_lora + 4), .ig = at + cap * (c.kv_lora + 4) + cap * ix, .pk = at + cap * (c.kv_lora + 4) + 2 * cap * ix, .pos_dev = e.pos_dev.ptr, .cap = cap }; + at += dsa_b; + } + return e; + } + + pub fn deinit(e: *Engine) void { + for (&e.graphs) |*pair| for (pair) |*g| if (g.*) |*x| x.deinit(); + e.gpa.free(e.dsa); + e.gpa.free(e.kda); + e.cand.free(); + e.pos_dev.free(); + e.state.free(); + e.pre.deinit(); + e.dec.deinit(); + e.s.deinit(); + e.k.unload(); + } + + /// State.reset: KDA states zeroed, position 0 (the DSA caches keep their bytes; positions past pos are never read). + pub fn reset(e: *Engine) !void { + const kv: cuda.DeviceBuffer = .{ .d = e.d, .ptr = e.state.ptr, .len = e.kda_bytes, .borrowed = true }; + try kv.fill8(0, e.s.handle); + e.cur = 0; + e.pos = 0; + } + + /// decode.prefill without drafts: the prompt in prompt chunks, then its first token. + pub fn prefill(e: *Engine, prompt: []const u32) !u32 { + if (prompt.len == 0) return error.GlmEmptyPrompt; + try e.reset(); + var tok: u32 = undefined; + var at: usize = 0; + while (at < prompt.len) : (at += prefill_rows) tok = try e.forward(prompt[at..@min(prompt.len, at + prefill_rows)], true); + return tok; + } + + /// serial_decode's step, the next token: a graph captured on first use for its window slot and KDA parity. + pub fn step(e: *Engine, tok: u32) !u32 { + if (e.pos + 1 > e.cap) return error.GlmCapacity; + const slot = graphSlot(e.pos, e.cap) orelse return e.forward(&.{tok}, false); + try e.dec.ids.upload(0, std.mem.asBytes(&tok)); + try e.pos_dev.fill32(@intCast(e.pos), e.s.handle); + const g = &e.graphs[slot][e.cur]; + if (g.* == null) g.* = try e.capture(); + try g.*.?.launchOn(e.s); + return e.finish(1); + } + + /// The decode step's GPU work recorded once (nothing runs): position and parity come from the replay's state. + fn capture(e: *Engine) !cuda.graph.Exec { + try cuda.graph.beginCapture(e.s, .global); + const recorded = e.body(&e.dec, 1, false, true); + var g = try cuda.graph.endCapture(e.s); + defer g.deinit(); + try recorded; + const x = try g.instantiate(); + try x.upload(e.s); + return x; + } + + /// One layer of the forward on its state; a decode row's KDA layer commits right after (its rows are reused next). + fn layer(e: *Engine, l: *const lay.LayerW, b: *Bufs, rows: usize, prompt: bool, capture_: bool, at: *[2]usize, g: lay.Gather) !void { + switch (l.mixer) { + .kda => |*kw| { + const sl = e.kda[at[0]]; + at[0] += 1; + const st: lay.KdaState = .{ .conv = sl.conv, .rec_in = sl.rec[e.cur], .rec_out = sl.rec[1 - e.cur] }; + const mode: lay.Mode = if (prompt) .{ .prompt = e.pos } else .decode; + try lay.forward(&e.k, &e.s, e.lc, l, .{ .kda = .{ .st = st, .mode = mode } }, &b.sc, b.x.ptr, rows, g); + if (!prompt) try lay.commit(&e.k, &e.s, kw, st, &b.sc, rows, rows); + }, + .dsa => { + const run: lay.DsaRun = .{ .st = e.dsa[at[1]], .sc = &b.dsc, .pos = e.pos, .capture = capture_ }; + at[1] += 1; + try lay.forward(&e.k, &e.s, e.lc, l, .{ .dsa = run }, &b.sc, b.x.ptr, rows, g); + }, + } + } + + fn gather(ctx: *anyopaque, s: cuda.Stream, site: lay.Site, part: u64, gath: u64, rows: usize) anyerror!void { + _ = site; + const e: *Engine = @ptrCast(@alignCast(ctx)); + try e.lib.check(e.lib.api.ncclAllGather(part, gath, rows * e.lc.d, .f32, e.comm, s.handle), "ncclAllGather"); + } + + /// compute + sample + commit for `ids` at the current position; the greedy token of the last row. + fn forward(e: *Engine, ids: []const u32, prompt: bool) !u32 { + const b = if (prompt) &e.pre else &e.dec; + const rows = ids.len; + if (rows == 0 or rows > b.rows or (!prompt and rows != 1)) return error.GlmRows; + if (e.pos + rows > e.cap) return error.GlmCapacity; + try b.ids.upload(0, std.mem.sliceAsBytes(ids)); + try e.pos_dev.fill32(@intCast(e.pos), e.s.handle); + try e.body(b, rows, prompt, false); + return e.finish(rows); + } + + /// The forward's GPU work from b.ids and pos_dev to every rank's candidate in e.cand; no host waits (capturable). + fn body(e: *Engine, b: *Bufs, rows: usize, prompt: bool, capture_: bool) !void { + const lc = e.lc; + try e.k.glue.embedB16(e.s, b.ids.ptr, e.w.embed, b.x.ptr, rows, lc.d, lc.streams); + const g: lay.Gather = .{ .ctx = e, .world = world, .run = gather }; + var at: [2]usize = .{ 0, 0 }; // the next KDA and DSA state + for (e.w.layers) |*l| try e.layer(l, b, rows, prompt, capture_, &at, g); + try head_.forward(&e.k.glue, &e.k.bmm, e.s, lc.eps, lc.d, lc.streams, e.w.head, &b.hs, b.x.ptr, rows, prompt); + try e.lib.check(e.lib.api.ncclAllGather(b.hs.cand, e.cand.ptr, 2, .f32, e.comm, e.s.handle), "ncclAllGather"); + } + + /// sample_rows' greedy pick from the gathered candidates, then the commit's position and KDA parity. + fn finish(e: *Engine, rows: usize) !u32 { + var got: [world * 2]f32 = undefined; + try e.s.synchronize(); + try e.cand.download(0, std.mem.sliceAsBytes(&got)); + var tok: [1]u32 = undefined; + try head_.choose(&got, world, 1, &tok); + e.cur = 1 - e.cur; + e.pos += rows; + return tok[0]; + } +}; + +const graph_slots = 8; + +/// A decode row's graph slot (forward.py's keys): 0 in the dense window, else 1 + log2(pool bucket / 1024); null: none. +fn graphSlot(pos: usize, cap: usize) ?usize { + if (pos + 1 <= dsa.dense_limit) return 0; + const np = dsa.poolBucket(pos, 1, cap / 4); + if (np < 1024) return null; + const slot = 1 + std.math.log2_int(usize, np / 1024); + return if (slot < graph_slots and np == @as(usize, 1024) << @intCast(slot - 1)) slot else null; +} + +test "decode graph slots follow the dense window and the pool buckets" { + try std.testing.expectEqual(@as(?usize, 0), graphSlot(2050, 32768)); + try std.testing.expectEqual(@as(?usize, 1), graphSlot(2051, 32768)); + try std.testing.expectEqual(@as(?usize, 2), graphSlot(4425, 32768)); + try std.testing.expectEqual(@as(?usize, 4), graphSlot(16416, 32768)); + try std.testing.expectEqual(@as(?usize, 4), graphSlot(32000, 32768)); +} + +/// The EXL3 dense scratch at every layer's widest, before any capture. +fn reserve(de: *cuda.exl3.Dense, l: *const lay.LayerW) !void { + switch (l.mixer) { + .kda => |k| for ([_]cuda.exl3.Weight{ k.qkv, k.o }) |w| try de.reserve(w), + .dsa => |a| for ([_]cuda.exl3.Weight{ a.q_a, a.kv_a, a.q_b, a.o }) |w| try de.reserve(w), + } + switch (l.ffn) { + .moe => |m| for ([_]cuda.exl3.Weight{ m.gate, m.up, m.down }) |w| try de.reserve(w), + .mlp => |m| for ([_]cuda.exl3.Weight{ m.gate, m.up, m.down }) |w| try de.reserve(w), + } +} diff --git a/zig/tests/cuda/exl3_dense_tests.zig b/zig/tests/cuda/exl3_dense_tests.zig new file mode 100644 index 0000000000..877d9ed791 --- /dev/null +++ b/zig/tests/cuda/exl3_dense_tests.zig @@ -0,0 +1,329 @@ +//! EXL3 dense prompt path: qlinear's dispatch, float64 references, chunk invariance, the engine's prompt bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3fmt = @import("core").exl3_format; +const exl3_load = @import("core").exl3_load; +const exl3 = cuda.exl3; +const qlinear = cuda.qlinear; +const exl3_tests = @import("exl3_tests.zig"); +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const bf = quant_ref_tests.helpers.bfOf; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + var dense = try exl3.Dense.load(r.gpu.d); + defer dense.unload(); + const ql: qlinear.Linear = .{ .exl3 = &dense }; + { + var ok = true; + qlinearCheck(r, &ql, &dense.lin) catch |e| { + state.fail(e, "exl3.qlinear"); + ok = false; + }; + if (ok) check.pass("exl3.qlinear", .{}); + } + { + var ok = true; + promptF64(r, &dense.lin) catch |e| { + state.fail(e, "exl3.prompt.linear.f64"); + ok = false; + }; + if (ok) check.pass("exl3.prompt.linear.f64", .{}); + } + { + var ok = true; + chunks(r, &dense) catch |e| { + state.fail(e, "exl3.prompt.linear.chunks"); + ok = false; + }; + if (ok) check.pass("exl3.prompt.linear.chunks", .{}); + } + if (pythonPrompt(r, &dense, dir) catch |e| blk: { + state.fail(e, "exl3.prompt.python"); + break :blk false; + }) check.pass("exl3.prompt.python", .{}); + if (loadRedsnow(r, &dense, dir) catch |e| blk: { + state.fail(e, "exl3.load.redsnow"); + break :blk false; + }) check.pass("exl3.load.redsnow", .{}); +} + +/// Through qlinear.Linear's decode and prompt dispatch, 1, 17 and 300 rows byte-equal the direct block loop. +fn qlinearCheck(r: *Rig, ql: *const qlinear.Linear, lin: *const exl3.Linear) !void { + const k = 1024; + const n = 512; + const trellis, const suh, const svh = try tensors(r, k, n, .{ .halves = 8 }); + const w = try weight(r, trellis, suh, svh, k, n, 8, .mul1); + const x = try r.bfs(300 * k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const counters = try exl3_tests.zeroed(r, exl3.Linear.counterBytes(w)); + const xh = try exl3_tests.zeroed(r, 128 * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try exl3_tests.zeroed(r, exl3.Linear.zBytes(w, 128)) else 0; + for ([_]usize{ 1, 17, 300 }) |m| { + const want = try r.zeros(m * n * 2); + var off: usize = 0; + while (off < m) : (off += 128) { + try lin.matmul(r.s, w, dx + off * k * 2, .bf16, @min(m - off, 128), want + off * n * 2, .bf16, xh, z, counters); + } + const wantraw = try r.back(want, m * n * 2); + const Kind = enum { decode, prompt }; + for ([_]Kind{ .decode, .prompt }) |kind| { + const got = try r.zeros(m * n * 2); + switch (kind) { + .decode => try ql.decode(r.s, .{ .exl3 = w }, .{ .x = dx }, got, m), + .prompt => try ql.prompt(r.s, .{ .exl3 = w }, .{ .x = dx }, got, m), + } + const raw = try r.back(got, m * n * 2); + const id = try std.fmt.allocPrint(r.a, "exl3.qlinear.{s} m={d}", .{ @tagName(kind), m }); + try check.sameBytes(id, raw, wantraw); + } + } +} + +/// 300 and 1000 prompt rows against the float64 reference (exl3_tests' reference with m a parameter). +fn promptF64(r: *Rig, lin: *const exl3.Linear) !void { + const cfgs = [_]struct { halves: u32, codebook: exl3fmt.Codebook, k: usize, n: usize }{ + .{ .halves = 8, .codebook = .mul1, .k = 512, .n = 512 }, + .{ .halves = 10, .codebook = .mul1, .k = 1024, .n = 512 }, + }; + for (cfgs) |c| { + const width: exl3fmt.Width = .{ .halves = @intCast(c.halves) }; + const trellis, const suh, const svh = try tensors(r, c.k, c.n, width); + const w = try weight(r, trellis, suh, svh, c.k, c.n, c.halves, cbOf(@tagName(c.codebook))); + const counters = try exl3_tests.zeroed(r, exl3.Linear.counterBytes(w)); + const xh = try exl3_tests.zeroed(r, 128 * c.k * 2); + const z: u64 = if (exl3.plan(c.k, c.n, 192, 8).sk > 1) try exl3_tests.zeroed(r, exl3.Linear.zBytes(w, 128)) else 0; + for ([_]usize{ 300, 1000 }) |m| { + const x = try r.bfs(m * c.k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const ref = try promptReference(r, width, c.codebook, trellis, c.k, c.n, x, suh, svh, m); + const id = try std.fmt.allocPrint(r.a, "exl3.prompt.linear.f64 k={d} n={d} m={d}", .{ c.k, c.n, m }); + const dy = try r.zeros(m * c.n * 2); + var off: usize = 0; + while (off < m) : (off += 128) { + try lin.matmul(r.s, w, dx + off * c.k * 2, .bf16, @min(m - off, 128), dy + off * c.n * 2, .bf16, xh, z, counters); + } + try exl3_tests.closeOut(r, id, dy, m * c.n, ref.want, ref.tol); + } + } +} + +/// A row's bytes are identical whether the prompt runs as one call, 128-row calls, a split at 37, or 2048 chunks. +fn chunks(r: *Rig, de: *exl3.Dense) !void { + const k = 512; + const n = 512; + const rows = 2185; // 17 whole 128-blocks and a 9-row tail + const trellis, const suh, const svh = try tensors(r, k, n, .{ .halves = 8 }); + const w = try weight(r, trellis, suh, svh, k, n, 8, .mul1); + const x = try r.bfs(rows * k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const y0 = try r.zeros(rows * n * 2); + try de.prompt(r.s, w, dx, .bf16, rows, y0, .bf16); + const want = try r.back(y0, rows * n * 2); + var b128: [17]usize = undefined; + for (&b128, 1..) |*v, i| v.* = i * 128; + for ([_][]const usize{ &.{}, &b128, &.{37}, &.{2048} }) |bounds| { + const y = try r.zeros(rows * n * 2); + var prev: usize = 0; + for (bounds) |b| { + try de.prompt(r.s, w, dx + prev * k * 2, .bf16, b - prev, y + prev * n * 2, .bf16); + prev = b; + } + try de.prompt(r.s, w, dx + prev * k * 2, .bf16, rows - prev, y + prev * n * 2, .bf16); + const got = try r.back(y, rows * n * 2); + const id = try std.fmt.allocPrint(r.a, "exl3.prompt.linear.chunks {d} splits", .{bounds.len}); + try check.sameBytes(id, got, want); + } +} + +/// The Python engine's recorded prompt bytes for one real dense projection, bit for bit. +fn pythonPrompt(r: *Rig, de: *exl3.Dense, dir: []const u8) !bool { + if (std.mem.eql(u8, dir, "none")) { + std.debug.print("SKIP exl3.prompt.python\n", .{}); + return false; + } + const d = (try denseEntry(r, dir)) orelse { + std.debug.print("SKIP exl3.prompt.python: no dense entry\n", .{}); + return false; + }; + const k: usize = @intCast(d.get("k").?.integer); + const n: usize = @intCast(d.get("n").?.integer); + const bits: u32 = bitsOf(d.get("bits").?); + const m: usize = @intCast(d.get("rows").?.integer); + const bias = d.get("bias").?; + const st = exl3.strides(k, n, 2 * bits, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(try exl3_tests.readFile(r, dir, d.get("trellis").?.string)), + .suh = try r.dev(try exl3_tests.readFile(r, dir, d.get("suh").?.string)), + .svh = try r.dev(try exl3_tests.readFile(r, dir, d.get("svh").?.string)), + .bias = if (std.meta.activeTag(bias) == .null) 0 else try r.dev(try exl3_tests.readFile(r, dir, bias.string)), + .k = @intCast(k), + .n = @intCast(n), + .halves = 2 * bits, + .codebook = cbOf(d.get("codebook").?.string), + .stride_k = st.k, + .stride_nb = st.nb, + }; + const x = try exl3_tests.readFile(r, dir, d.get("x").?.string); + const want = try exl3_tests.readFile(r, dir, d.get("y").?.string); + const dy = try r.zeros(m * n * 2); + try de.prompt(r.s, w, try r.dev(x), .bf16, m, dy, .bf16); + const got = try r.back(dy, want.len); + try check.sameBytes("exl3.prompt.python", got, want); + return true; +} + +/// The checkpoint loader over real RED-SNOW: every EXL3 group admitted, one layer's prompt vs the fixture bytes. +fn loadRedsnow(r: *Rig, de: *exl3.Dense, dir: []const u8) !bool { + if (std.mem.eql(u8, dir, "none")) { + std.debug.print("SKIP exl3.load.redsnow\n", .{}); + return false; + } + const o = try manifest(r, dir); + const model = o.get("model") orelse return error.TestFailed; // the fixture names its model folder + const d = (try denseEntry(r, dir)) orelse { + std.debug.print("SKIP exl3.load.redsnow: no dense entry\n", .{}); + return false; + }; + var s = try exl3_load.scan(r.gpu.gpa, r.gpu.io, model.string); + defer s.deinit(); + const sum = s.summary; + // the inspect tool's counts: every group suh/svh mul1, the whole-bits histogram, the experts' downs + try check.expect(sum.groups == 37408, "exl3.load.redsnow: {d} groups, expected 37408", .{sum.groups}); + try check.expect(sum.mul1 == 37408, "exl3.load.redsnow: {d} mul1, expected 37408", .{sum.mul1}); + try check.expect(sum.suh_svh == 37408, "exl3.load.redsnow: {d} suh/svh, expected 37408", .{sum.suh_svh}); + const hist = [_]usize{ 15425, 11801, 10054, 43, 85 }; // whole bits 4 to 8 + const downs = [_]usize{ 10283, 1448, 653, 0, 0 }; + for (hist, downs, 4..) |want, down, bits| { + try check.expect(sum.bits[2 * bits] == want, "exl3.load.redsnow: {d} groups at {d} bits, expected {d}", .{ sum.bits[2 * bits], bits, want }); + try check.expect(sum.experts_down[2 * bits] == down, "exl3.load.redsnow: {d} expert down projections at {d} bits, expected {d}", .{ sum.experts_down[2 * bits], bits, down }); + } + try check.expect(s.bad == 0, "exl3.load.redsnow: {d} tensors refused", .{s.bad}); + const gi = s.find(d.get("name").?.string) orelse return error.TestFailed; // the fixture's layer + const g = s.groups[gi]; + const p = try s.parts(r.a, gi); + const st = exl3.strides(g.k, g.n, g.halves, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(p.trellis.bytes), + .suh = try r.dev(p.suh.bytes), + .svh = try r.dev(p.svh.bytes), + .bias = if (p.bias) |b| try r.dev(b.bytes) else 0, + .k = @intCast(g.k), + .n = @intCast(g.n), + .halves = g.halves, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; + try check.expect(g.halves == 2 * bitsOf(d.get("bits").?) and g.k == @as(usize, @intCast(d.get("k").?.integer)) and + g.n == @as(usize, @intCast(d.get("n").?.integer)), "exl3.load.redsnow: the loaded layer is not the fixture's", .{}); + const m: usize = @intCast(d.get("rows").?.integer); + const x = try exl3_tests.readFile(r, dir, d.get("x").?.string); + const want = try exl3_tests.readFile(r, dir, d.get("y").?.string); + const dy = try r.zeros(m * g.n * 2); + try de.prompt(r.s, w, try r.dev(x), .bf16, m, dy, .bf16); + const got = try r.back(dy, want.len); + try check.sameBytes("exl3.load.redsnow", got, want); + return true; +} + +/// The manifest's "dense" entry, null when an older fixture dir has none. +fn denseEntry(r: *Rig, dir: []const u8) !?std.json.ObjectMap { + const o = try manifest(r, dir); + return if (o.get("dense")) |v| v.object else null; +} + +fn manifest(r: *Rig, dir: []const u8) !std.json.ObjectMap { + const text = try std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, "manifest.json" }), r.a, .limited(1 << 24)); + const parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + return parsed.value.object; +} + +fn bitsOf(v: std.json.Value) u32 { + return switch (v) { + .integer => |i| @intCast(i), + .float => |f| @intFromFloat(f), + else => unreachable, + }; +} + +fn cbOf(name: []const u8) exl3.Codebook { + if (std.mem.eql(u8, name, "mcg")) return .mcg; + if (std.mem.eql(u8, name, "inst3")) return .inst3; + return .mul1; +} + +/// Random trellis, scales at one shape (the same layout as exl3_tests' f64 cfgs). +fn tensors(r: *Rig, k: usize, n: usize, width: exl3fmt.Width) !struct { []i16, []u16, []u16 } { + const trellis = try r.a.alloc(i16, (k / 16) * (n / 16) * width.tileWords()); + for (trellis) |*v| v.* = r.rng.int(i16); + const suh = try r.a.alloc(u16, k); + const svh = try r.a.alloc(u16, n); + for (suh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + for (svh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + return .{ trellis, suh, svh }; +} + +fn weight(r: *Rig, trellis: []const i16, suh: []const u16, svh: []const u16, k: usize, n: usize, halves: u32, cb: exl3.Codebook) !exl3.Weight { + const st = exl3.strides(k, n, halves, false); + return .{ + .trellis = try r.dev(std.mem.sliceAsBytes(trellis)), + .suh = try r.dev(std.mem.sliceAsBytes(suh)), + .svh = try r.dev(std.mem.sliceAsBytes(svh)), + .bias = 0, + .k = @intCast(k), + .n = @intCast(n), + .halves = halves, + .codebook = cb, + .stride_k = st.k, + .stride_nb = st.nb, + }; +} + +/// The Python engine's f64 reference, with the row count a parameter. +fn promptReference(r: *Rig, width: exl3fmt.Width, cb: exl3fmt.Codebook, trellis: []const i16, k: usize, n: usize, x: []const u16, suh: []const u16, svh: []const u16, m: usize) !struct { want: []f64, tol: []f64 } { + const wq = try r.a.alloc(f16, k * n); + exl3fmt.unpack(width, cb, trellis, k / 16, n / 16, wq); + const xs = try r.a.alloc(f64, m * k); + for (0..m) |row| { + for (0..k) |i| xs[row * k + i] = bf(x[row * k + i]) * @as(f64, @floatCast(@as(f16, @bitCast(suh[i])))); + } + exl3fmt.rotate(xs, m, k, 1); + const xh = try r.a.alloc(f16, m * k); + for (xh, xs) |*o, v| o.* = @floatCast(v); + const ypre = try r.a.alloc(f64, m * n); // [n][m], rotated along n + @memset(ypre, 0); + const s1 = try r.a.alloc(f64, m * n); // each output's sum of |xh * wq| over k + @memset(s1, 0); + var xv: [1000]f64 = undefined; + var xa: [1000]f64 = undefined; + for (0..k) |ki| { + for (0..m) |row| { + xv[row] = @floatCast(xh[row * k + ki]); + xa[row] = @abs(xv[row]); + } + for (0..n) |col| { + const wv: f64 = @floatCast(wq[ki * n + col]); + const wa = @abs(wv); + const base = col * m; + for (0..m) |row| { + ypre[base + row] += xv[row] * wv; + s1[base + row] += xa[row] * wa; + } + } + } + exl3fmt.rotate(ypre, n, m, 0); + const want = try r.a.alloc(f64, m * n); + const tol = try r.a.alloc(f64, m * n); + for (0..n) |col| { + const scale: f64 = @floatCast(@as(f16, @bitCast(svh[col]))); + for (0..m) |row| { + want[row * n + col] = ypre[col * m + row] * scale; + // bf16 output ulp (2^-7 of the value) plus fp32 accumulation and fp16 xh cases, 2^-11 of the |t| sum + tol[row * n + col] = @abs(want[row * n + col]) / 128 + s1[row * n + col] / 2048 + 1e-30; + } + } + return .{ .want = want, .tol = tol }; +} diff --git a/zig/tests/cuda/exl3_experts_f64_tests.zig b/zig/tests/cuda/exl3_experts_f64_tests.zig new file mode 100644 index 0000000000..6e0c162cbb --- /dev/null +++ b/zig/tests/cuda/exl3_experts_f64_tests.zig @@ -0,0 +1,463 @@ +//! The EXL3 routed experts' float64 reference: the mirrored chain and each output's worst-case error bound. +const std = @import("std"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const cuda = @import("cuda"); +const exl3fmt = @import("core").exl3_format; +const exl3e = cuda.exl3_experts; +const exl3t = @import("exl3_tests.zig"); +const Rig = quant_ref_tests.RigT; +const bf = quant_ref_tests.helpers.bfOf; +const Run = exl3t.Run; + +pub const E = 9; // experts of a synthetic layer +pub const D = 512; // model width +const I = 256; // expert width +pub const SLOTS = 8; +const LIMIT: f32 = 10.0; // the SwiGLU clamp +const CAP = 300; // rows of x, pick and wts the builders make (the f64 checks use the first 128) +const rows = [_]usize{ 1, 3, 8, 16, 17, 64, 128 }; // the f64 checks' row counts (the reference builds 128 rows) + +const f64_ids = [_][]const u8{ "exl3.experts.f64.mul1.4", "exl3.experts.f64.mul1.5", "exl3.experts.f64.mul1.6" }; + +/// One synthetic layer: the device tensors, the host trellises the reference unpacks, and the seeded routing. +pub const Built = struct { + layer: exl3e.Layer, + bits: []const u32, + gate: [E][]i16, + up: [E][]i16, + down: [E][]i16, + suh_g: []u16, // fp16 [E, D] + suh_u: []u16, + svh_g: []u16, // fp16 [E, I] + svh_u: []u16, + suh_d: []u16, // fp16 [E, I] + svh_d: []u16, // fp16 [E, D] + dx: u64, + dpick: u64, + dwts: u64, + x: []u16, // bf16 [CAP, D] + pick: []i32, // [CAP, SLOTS] + wts: []f32, // [CAP, SLOTS] +}; + +/// The f64 id checks: every width alone, then the mixed layer folded into the .6 id (the contract's ids). +pub fn checks(r: *Rig, exp: *const exl3e.Experts, state: *Run) void { + for (f64_ids, 0..) |id, i| { + var ok = true; + f64Check(r, exp, &allBits(@intCast(i + 4))) catch |e| { + state.fail(e, id); + ok = false; + }; + if (ok) check.pass("{s}", .{id}); + } + { + var ok = true; + f64Check(r, exp, &.{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }) catch |e| { + state.fail(e, f64_ids[2]); + ok = false; + }; + if (ok) check.pass("{s}", .{f64_ids[2]}); + } +} + +fn allBits(bits: u32) [E]u32 { + return .{ bits, bits, bits, bits, bits, bits, bits, bits, bits }; +} + +fn widthOf(bits: u32) exl3fmt.Width { + return .{ .halves = @intCast(2 * bits) }; +} + +fn f16v(v: u16) f64 { + return @floatCast(@as(f16, @bitCast(v))); +} + +/// bf16 round-to-nearest-even of an f64 (the kernel rounds its fp32 value; the double round is within the tolerance). +fn bf16r(v: f64) f64 { + const u: u32 = @bitCast(@as(f32, @floatCast(v))); + return bf(@intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16)); +} + +pub fn buildLayer(r: *Rig, bits: []const u32) !Built { + var b: Built = .{ + .layer = undefined, + .bits = bits, + .gate = undefined, + .up = undefined, + .down = undefined, + .dx = 0, + .dpick = 0, + .dwts = 0, + .suh_g = undefined, + .suh_u = undefined, + .svh_g = undefined, + .svh_u = undefined, + .suh_d = undefined, + .svh_d = undefined, + .x = undefined, + .pick = undefined, + .wts = undefined, + }; + var gate_p: [E]u64 = undefined; + var up_p: [E]u64 = undefined; + var down_p: [E]u64 = undefined; + var gk2: [E]i32 = undefined; + var uk2: [E]i32 = undefined; + var dk2: [E]i32 = undefined; + const suh_g = try r.a.alloc(u16, E * D); + const suh_u = try r.a.alloc(u16, E * D); + const svh_g = try r.a.alloc(u16, E * I); + const svh_u = try r.a.alloc(u16, E * I); + const suh_d = try r.a.alloc(u16, E * I); + const svh_d = try r.a.alloc(u16, E * D); + for (0..E) |e| { + b.gate[e] = try randI16(r, (D / 16) * (I / 16) * widthOf(bits[e]).tileWords()); + b.up[e] = try randI16(r, (D / 16) * (I / 16) * widthOf(bits[e]).tileWords()); + b.down[e] = try randI16(r, (I / 16) * (D / 16) * widthOf(bits[e]).tileWords()); + gate_p[e] = try r.dev(std.mem.sliceAsBytes(b.gate[e])); + up_p[e] = try r.dev(std.mem.sliceAsBytes(b.up[e])); + down_p[e] = try r.dev(std.mem.sliceAsBytes(b.down[e])); + gk2[e] = @intCast(2 * bits[e]); + uk2[e] = gk2[e]; + dk2[e] = gk2[e]; + randScale(r, suh_g[e * D ..][0..D]); + randScale(r, suh_u[e * D ..][0..D]); + randScale(r, svh_g[e * I ..][0..I]); + randScale(r, svh_u[e * I ..][0..I]); + randScale(r, suh_d[e * I ..][0..I]); + randScale(r, svh_d[e * D ..][0..D]); + } + b.suh_g = suh_g; + b.suh_u = suh_u; + b.svh_g = svh_g; + b.svh_u = svh_u; + b.suh_d = suh_d; + b.svh_d = svh_d; + b.layer = .{ + .gate = try r.dev(std.mem.sliceAsBytes(gate_p[0..])), + .up = try r.dev(std.mem.sliceAsBytes(up_p[0..])), + .down = try r.dev(std.mem.sliceAsBytes(down_p[0..])), + .gate_k2 = try r.dev(std.mem.sliceAsBytes(gk2[0..])), + .up_k2 = try r.dev(std.mem.sliceAsBytes(uk2[0..])), + .down_k2 = try r.dev(std.mem.sliceAsBytes(dk2[0..])), + .suh_g = try r.dev(std.mem.sliceAsBytes(suh_g)), + .suh_u = try r.dev(std.mem.sliceAsBytes(suh_u)), + .svh_g = try r.dev(std.mem.sliceAsBytes(svh_g)), + .svh_u = try r.dev(std.mem.sliceAsBytes(svh_u)), + .suh_d = try r.dev(std.mem.sliceAsBytes(suh_d)), + .svh_d = try r.dev(std.mem.sliceAsBytes(svh_d)), + .experts = E, + .dims = D, + .width = I, + .codebook = .mul1, + .limit = LIMIT, + .act_mode = 0, + }; + b.x = try r.bfs(CAP * D, 0.5); // half-range x, with the small scales, keeps the SwiGLU inputs moderate + b.pick = try picks(r, CAP); + b.wts = try wtsOf(r, CAP); + b.dx = try r.dev(std.mem.sliceAsBytes(b.x)); + b.dpick = try r.dev(std.mem.sliceAsBytes(b.pick)); + b.dwts = try r.dev(std.mem.sliceAsBytes(b.wts)); + return b; +} + +fn randI16(r: *Rig, n: usize) ![]i16 { + const out = try r.a.alloc(i16, n); + for (out) |*v| v.* = r.rng.int(i16); + return out; +} + +/// Small scales (with the half-range x): the |t||w| sums and the SwiGLU inputs stay moderate for the bound. +fn randScale(r: *Rig, out: []u16) void { + for (out) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 0.25 - 0.125))); +} + +/// Distinct experts a row: a partial Fisher-Yates over the layer's experts, seeded. +fn picks(r: *Rig, cap: usize) ![]i32 { + const out = try r.a.alloc(i32, cap * SLOTS); + for (0..cap) |row| { + var idx: [E]u32 = undefined; + for (0..E, 0..) |_, i| idx[i] = @intCast(i); + var n: usize = E; + for (0..SLOTS) |k| { + const j = r.rng.uintLessThan(usize, n); + out[row * SLOTS + k] = @intCast(idx[j]); + idx[j] = idx[n - 1]; + n -= 1; + } + } + return out; +} + +fn wtsOf(r: *Rig, cap: usize) ![]f32 { + const out = try r.a.alloc(f32, cap * SLOTS); + for (out) |*v| v.* = 0.1 + 0.9 * r.rng.float(f32); + return out; +} + +/// One routed call over `rows` rows starting at row 0 of the built routing; the scratch and out are fresh canaries. +pub fn runRows(r: *Rig, exp: *const exl3e.Experts, b: *const Built, rows_n: usize) !u64 { + return (try runSlice(r, exp, b.layer, b.dx, b.dpick, b.dwts, rows_n)).out; +} + +pub fn runSlice(r: *Rig, exp: *const exl3e.Experts, layer: exl3e.Layer, x: u64, pick: u64, wts: u64, rows_n: usize) !struct { out: u64, y: u64 } { + const sb = try r.zeros(try exl3e.scratchBytes(layer, rows_n, SLOTS)); + const sc = try exl3e.scratch(sb, layer, rows_n, SLOTS); + const out = try r.zeros(rows_n * D * 4); + try exp.run(r.s, layer, x, D, .bf16, pick, wts, rows_n, out, sc); + return .{ .out = out, .y = sc.y }; +} + +fn f64Check(r: *Rig, exp: *const exl3e.Experts, bits: []const u32) !void { + const b = try buildLayer(r, bits); + const want = try reference(r, &b, 128); + for (rows) |m| try closeOut(r, b.bits, m, (try runRows(r, exp, &b, m)), want.want, want.tol); +} + +/// fp32 outputs within the f64 reference's tolerance, with the exact FAIL line. +fn closeOut(r: *Rig, bits: []const u32, m: usize, ptr: u64, want: []const f64, tol: []const f64) !void { + const raw = try r.back(ptr, m * D * 4); + var most: f64 = 0; + var at: usize = 0; + for (want[0 .. m * D], tol[0 .. m * D], 0..) |wv, t, i| { + const g: f64 = @floatCast(std.mem.bytesToValue(f32, raw[4 * i ..][0..4])); + const q = @abs(g - wv) / t; + if (q > most) { + most = q; + at = i; + } + } + std.debug.print("RESULT exl3.experts.f64.mul1.{d} m={d}: largest difference {d:.3} of its tolerance (element {d})\n", .{ bits[0], m, most, at }); + if (most > 1) { + std.debug.print("FAIL exl3.experts.f64.mul1.{d}: largest difference {d:.3} of its tolerance\n", .{ bits[0], most }); + return error.TestFailed; + } +} + +/// A 300-row chunk against the f64 reference (the mixed-width layer). +pub fn prompt(r: *Rig, exp: *const exl3e.Experts) !void { + var bits: [E]u32 = .{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }; + const b = try buildLayer(r, &bits); + const want = try reference(r, &b, 300); + const res = try runRows(r, exp, &b, 300); + const raw = try r.back(res, 300 * D * 4); + var most: f64 = 0; + var at: usize = 0; + for (want.want, want.tol, 0..) |wv, t, i| { + const g: f64 = @floatCast(std.mem.bytesToValue(f32, raw[4 * i ..][0..4])); + const q = @abs(g - wv) / t; + if (q > most) { + most = q; + at = i; + } + } + std.debug.print("RESULT exl3.prompt.experts.f64: largest difference {d:.3} of its tolerance (element {d})\n", .{ most, at }); + if (most > 1) return error.TestFailed; +} + +/// The f64 reference: the kernel's chain mirrored in f64, each output's error its per-element worst case. +fn reference(r: *Rig, b: *const Built, rows_n: usize) !struct { want: []f64, tol: []f64 } { + const p = rows_n * SLOTS; + // xg: fp16((x * suh_g) @ H) per slot; |xg| and its error <= V, the block's sum of |x * suh| over sqrt(128) + const xs = try r.a.alloc(f64, p * D); + for (0..p) |s0| { + const ex: usize = @intCast(b.pick[s0]); + for (0..D) |i| xs[s0 * D + i] = bf(b.x[(s0 / SLOTS) * D + i]) * f16v(b.suh_g[ex * D + i]); + } + exl3fmt.rotate(xs, p, D, 1); + const xg = try r.a.alloc(f16, p * D); + for (xg, xs) |*o, v| o.* = @floatCast(v); + const xgerr = try r.a.alloc(f64, p * D); + for (xgerr, xs) |*o, v| o.* = @abs(v); + absRotateCols(xgerr, p, D); + for (xgerr) |*v| v.* *= 1.0 / 2048.0 + 8.0 / 16777216.0; // the fp16 rounding (2^-11) plus the fp32 rotation's (8 * 2^-24), both of V + + const wq = try r.a.alloc(f16, @max(D, I) * @max(D, I)); + const zg = try r.a.alloc(f64, I * p); // [n][p] + const zu = try r.a.alloc(f64, I * p); + const s1g = try r.a.alloc(f64, I * p); + const s1u = try r.a.alloc(f64, I * p); + const eg = try r.a.alloc(f64, I * p); // the GEMM's propagated input error, before the output rotation + const eu = try r.a.alloc(f64, I * p); + @memset(zg, 0); + @memset(zu, 0); + @memset(s1g, 0); + @memset(s1u, 0); + @memset(eg, 0); + @memset(eu, 0); + for (0..E) |ex| { + exl3fmt.unpack(widthOf(b.bits[ex]), .mul1, b.gate[ex], D / 16, I / 16, wq[0 .. D * I]); + for (0..p) |s0| { + if (b.pick[s0] != @as(i32, @intCast(ex))) continue; + gemmK(xg[s0 * D ..][0..D], wq, zg[s0..], s1g[s0..], eg[s0..], xgerr[s0 * D ..][0..D], D, I, p); + } + exl3fmt.unpack(widthOf(b.bits[ex]), .mul1, b.up[ex], D / 16, I / 16, wq[0 .. D * I]); + for (0..p) |s0| { + if (b.pick[s0] != @as(i32, @intCast(ex))) continue; + gemmK(xg[s0 * D ..][0..D], wq, zu[s0..], s1u[s0..], eu[s0..], xgerr[s0 * D ..][0..D], D, I, p); + } + } + // the fp16 rounding (2^-11) and the fp32 accumulation (2^-17 of 128 roundings), both on the |t||w| sum + for (eg, s1g) |*e, s| e.* += s * (1.0 / 2048.0 + 1.0 / 131072.0); + for (eu, s1u) |*e, s| e.* += s * (1.0 / 2048.0 + 1.0 / 131072.0); + exl3fmt.rotate(zg, I, p, 0); + exl3fmt.rotate(zu, I, p, 0); + // the output Hadamard: worst case is the block's sum of |errors| over sqrt(128), plus the fp32 rotation's own + absRotate(eg, I, p); + absRotate(eu, I, p); + const zgerr = try r.a.alloc(f64, I * p); + for (zgerr, zg) |*o, v| o.* = @abs(v); + absRotate(zgerr, I, p); + for (zgerr) |*v| v.* *= 8.0 / 16777216.0; + const zuerr = try r.a.alloc(f64, I * p); + for (zuerr, zu) |*o, v| o.* = @abs(v); + absRotate(zuerr, I, p); + for (zuerr) |*v| v.* *= 8.0 / 16777216.0; + for (eg, zgerr) |*e, v| e.* += v; + for (eu, zuerr) |*e, v| e.* += v; + + // the SwiGLU (bf16 roundings, f64 exp): g's error is the rotated GEMM error through svh plus a bf16 double rounding + const v = try r.a.alloc(f64, I * p); + const verr = try r.a.alloc(f64, I * p); + for (0..I) |n| { + for (0..p) |s0| { + const ex: usize = @intCast(b.pick[s0]); + const vg = zg[n * p + s0] * f16v(b.svh_g[ex * I + n]); + const gg = @min(bf16r(vg), @as(f64, LIMIT)); + const vu = zu[n * p + s0] * f16v(b.svh_u[ex * I + n]); + const uu = @min(@max(bf16r(vu), -@as(f64, LIMIT)), @as(f64, LIMIT)); + const inner = bf16r(gg / (1.0 + std.math.exp(-gg))); + const act = bf16r(inner * uu); + v[n * p + s0] = act * f16v(b.suh_d[ex * I + n]); + const egk = eg[n * p + s0] * @abs(f16v(b.svh_g[ex * I + n])) + (@abs(vg) + @abs(gg)) / 256; + const euk = eu[n * p + s0] * @abs(f16v(b.svh_u[ex * I + n])) + (@abs(vu) + @abs(uu)) / 256; + const einner = egk + (@abs(gg) + 1) / 524288 + @abs(inner) / 256; + verr[n * p + s0] = (@abs(uu) * einner + @abs(inner) * euk + @abs(act) / 256) * @abs(f16v(b.suh_d[ex * I + n])) + + @abs(v[n * p + s0]) / 16777216; // act's: |silu'| (<= 1) g's, |silu(g)| (<= |g|) u's, expf and roundings + } + } + // the down input: the fp32 rotation's roundings (of the block's sum of |v|) and the fp16 store (2^-11) + exl3fmt.rotate(v, I, p, 0); + absRotate(verr, I, p); + const vrot = try r.a.alloc(f64, I * p); + for (vrot, v) |*o, val| o.* = @abs(val); + absRotate(vrot, I, p); + for (vrot) |*o| o.* *= 8.0 / 16777216.0; + const xd = try r.a.alloc(f16, p * I); + const xderr = try r.a.alloc(f64, p * I); + for (0..p) |s0| { + for (0..I) |ki| { + xd[s0 * I + ki] = @floatCast(v[ki * p + s0]); + xderr[s0 * I + ki] = verr[ki * p + s0] + vrot[ki * p + s0] + @abs(xd[s0 * I + ki]) / 2048; + } + } + + // down: the GEMM, its |t| and |w| sums, and the propagated input error + const zd = try r.a.alloc(f64, D * p); // [d][p] + const s1d = try r.a.alloc(f64, D * p); + const de = try r.a.alloc(f64, D * p); + @memset(zd, 0); + @memset(s1d, 0); + @memset(de, 0); + for (0..E) |ex| { + exl3fmt.unpack(widthOf(b.bits[ex]), .mul1, b.down[ex], I / 16, D / 16, wq[0 .. I * D]); + for (0..p) |s0| { + if (b.pick[s0] != @as(i32, @intCast(ex))) continue; + for (0..I) |ki| { + const t: f64 = @floatCast(xd[s0 * I + ki]); + const ta = @abs(t); + const row = wq[ki * D ..][0..D]; + const xek = xderr[s0 * I + ki]; + for (0..D) |d| { + const wv: f64 = @floatCast(row[d]); + zd[d * p + s0] += t * wv; + s1d[d * p + s0] += ta * @abs(wv); + de[d * p + s0] += @abs(wv) * xek; + } + } + } + } + // the fp16 xd rounding (2^-11 of the |t||w| sum) and the fp32 accumulation (2^-17), then the output Hadamard's mix + for (de, s1d) |*e, s| e.* += s * (1.0 / 2048.0 + 1.0 / 131072.0); + exl3fmt.rotate(zd, D, p, 0); + absRotate(de, D, p); + const zderr = try r.a.alloc(f64, D * p); + for (zderr, zd) |*o, val| o.* = @abs(val); + absRotate(zderr, D, p); + for (zderr) |*o| o.* *= 8.0 / 16777216.0; + for (de, zderr) |*e, val| e.* += val; + + // the per-slot outputs and the weighted combine: the fp32 fma chain (8 roundings of the running sum) and the store + const want = try r.a.alloc(f64, rows_n * D); + const tol = try r.a.alloc(f64, rows_n * D); + for (0..rows_n) |row| { + for (0..D) |d| { + var acc: f32 = 0; + var err: f64 = 0; + var swy: f64 = 0; + for (0..SLOTS) |k| { + const s0 = row * SLOTS + k; + const ex: usize = @intCast(b.pick[s0]); + const y: f32 = @floatCast(zd[d * p + s0] * f16v(b.svh_d[ex * D + d])); + acc = @mulAdd(f32, b.wts[row * SLOTS + k], y, acc); + const yerr = de[d * p + s0] * @abs(f16v(b.svh_d[ex * D + d])) + @abs(y) / 16777216; + err += @abs(b.wts[row * SLOTS + k]) * yerr; + swy += @abs(b.wts[row * SLOTS + k] * y); + } + want[row * D + d] = acc; + tol[row * D + d] = err + swy / 2097152 + @abs(want[row * D + d]) / 16777216 + 1e-30; + } + } + return .{ .want = want, .tol = tol }; +} + +/// One output column's GEMM over k: the sums into [n][p], their |t| sums, the propagated input error. +fn gemmK(xrow: []const f16, wq: []const f16, z: []f64, s1: []f64, e: []f64, xerr: []const f64, k: usize, n: usize, p: usize) void { + for (0..k) |ki| { + const t: f64 = @floatCast(xrow[ki]); + const ta = @abs(t); + const wrow = wq[ki * n ..][0..n]; + const xek = xerr[ki]; + for (0..n) |col| { + const wv: f64 = @floatCast(wrow[col]); + z[col * p] += t * wv; + s1[col * p] += ta * @abs(wv); + e[col * p] += @abs(wv) * xek; + } + } +} + +/// Each 128-block along the row axis replaced by the sum of its |entries| over sqrt(128) (the worst case). +fn absRotate(x: []f64, nrows: usize, cols: usize) void { + const inv = 1.0 / @sqrt(@as(f64, 128.0)); + var r: usize = 0; + while (r < nrows) : (r += 128) { + for (0..cols) |c| { + var s: f64 = 0; + for (0..128) |n| s += @abs(x[(r + n) * cols + c]); + s *= inv; + for (0..128) |n| x[(r + n) * cols + c] = s; + } + } +} + +test "the ids are the contract's" { + try std.testing.expectEqualStrings("exl3.experts.f64.mul1.4", f64_ids[0]); + try std.testing.expectEqualStrings("exl3.experts.f64.mul1.6", f64_ids[2]); +} + +/// The column-axis twin: each 128-run along the columns replaced by the sum of its |entries| over sqrt(128). +fn absRotateCols(x: []f64, nrows: usize, cols: usize) void { + const inv = 1.0 / @sqrt(@as(f64, 128.0)); + for (0..nrows) |row| { + var c: usize = 0; + while (c < cols) : (c += 128) { + var s: f64 = 0; + for (0..128) |n| s += @abs(x[row * cols + c + n]); + s *= inv; + for (0..128) |n| x[row * cols + c + n] = s; + } + } +} diff --git a/zig/tests/cuda/exl3_experts_tests.zig b/zig/tests/cuda/exl3_experts_tests.zig new file mode 100644 index 0000000000..20710d303f --- /dev/null +++ b/zig/tests/cuda/exl3_experts_tests.zig @@ -0,0 +1,197 @@ +//! The EXL3 routed experts: the f64 checks, row and chunk invariance, and the Python engine's recorded bytes. +const std = @import("std"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const cuda = @import("cuda"); +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const exl3t = @import("exl3_tests.zig"); +const ref = @import("exl3_experts_f64_tests.zig"); +const Rig = quant_ref_tests.RigT; +const Run = exl3t.Run; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + var exp = try exl3e.Experts.load(r.gpu.d); + defer exp.unload(); + + // every expert one width: 4, 5 and 6 bits (the mixed run folds into .6) + ref.checks(r, &exp, state); + { + var ok = true; + rowsInvariant(r, &exp) catch |e| { + state.fail(e, "exl3.experts.rows"); + ok = false; + }; + if (ok) check.pass("exl3.experts.rows", .{}); + } + { + var ok = true; + ref.prompt(r, &exp) catch |e| { + state.fail(e, "exl3.prompt.experts.f64"); + ok = false; + }; + if (ok) check.pass("exl3.prompt.experts.f64", .{}); + } + { + var ok = true; + chunks(r, &exp) catch |e| { + state.fail(e, "exl3.prompt.experts.chunks"); + ok = false; + }; + if (ok) check.pass("exl3.prompt.experts.chunks", .{}); + } + { + var ok = true; + python(r, &exp, dir, state) catch |e| { + state.fail(e, "exl3.experts.python"); + ok = false; + }; + if (ok) check.pass("exl3.experts.python", .{}); + } +} + +/// A row's output bytes are its own: solo, with 3 rows, in 64-row batches, shared or not. +fn rowsInvariant(r: *Rig, exp: *const exl3e.Experts) !void { + var bits: [ref.E]u32 = .{ 6, 4, 5, 6, 4, 5, 6, 4, 5 }; + const b = try ref.buildLayer(r, &bits); + const solo = try r.back(try ref.runRows(r, exp, &b, 1), ref.D * 4); + // 3 rows: two more rows on the same picks as row 0 + const p3 = try r.a.alloc(i32, 3 * ref.SLOTS); + @memcpy(p3[0 .. 2 * ref.SLOTS], b.pick[0 .. 2 * ref.SLOTS]); + @memcpy(p3[2 * ref.SLOTS ..], b.pick[0..ref.SLOTS]); + const w3 = try r.a.alloc(f32, 3 * ref.SLOTS); + @memcpy(w3, b.wts[0 .. 3 * ref.SLOTS]); + const got3 = try r.back((try ref.runSlice(r, exp, b.layer, b.dx, try r.dev(std.mem.sliceAsBytes(p3)), try r.dev(std.mem.sliceAsBytes(w3)), 3)).out, ref.D * 4); + try check.sameBytes("exl3.experts.rows 3 rows", got3, solo); + // 64 rows: the others share row 0's experts (the same picks), or run their own + const p64 = try r.a.alloc(i32, 64 * ref.SLOTS); + const w64 = try r.a.alloc(f32, 64 * ref.SLOTS); + @memcpy(p64[0..ref.SLOTS], b.pick[0..ref.SLOTS]); + @memcpy(w64[0..ref.SLOTS], b.wts[0..ref.SLOTS]); + for (1..64) |row| { + @memcpy(p64[row * ref.SLOTS ..][0..ref.SLOTS], b.pick[0..ref.SLOTS]); + @memcpy(w64[row * ref.SLOTS ..][0..ref.SLOTS], b.wts[0..ref.SLOTS]); + } + const shared = try r.back((try ref.runSlice(r, exp, b.layer, b.dx, try r.dev(std.mem.sliceAsBytes(p64)), try r.dev(std.mem.sliceAsBytes(w64)), 64)).out, ref.D * 4); + try check.sameBytes("exl3.experts.rows 64 rows shared", shared[0 .. ref.D * 4], solo); + const own = try r.back((try ref.runSlice(r, exp, b.layer, b.dx, b.dpick, b.dwts, 64)).out, ref.D * 4); + try check.sameBytes("exl3.experts.rows 64 rows own", own[0 .. ref.D * 4], solo); + // and a 64-row batch whose other rows interleave with row 0's experts in both orders + const p64b = try r.a.alloc(i32, 64 * ref.SLOTS); + const w64b = try r.a.alloc(f32, 64 * ref.SLOTS); + @memcpy(p64b[0..ref.SLOTS], b.pick[0..ref.SLOTS]); + @memcpy(w64b[0..ref.SLOTS], b.wts[0..ref.SLOTS]); + for (1..64) |row| { + const src = if (row % 2 == 0) @as(usize, 0) else row % ref.SLOTS; + @memcpy(p64b[row * ref.SLOTS ..][0..ref.SLOTS], b.pick[src * ref.SLOTS ..][0..ref.SLOTS]); + @memcpy(w64b[row * ref.SLOTS ..][0..ref.SLOTS], b.wts[src * ref.SLOTS ..][0..ref.SLOTS]); + } + const mixed = try r.back((try ref.runSlice(r, exp, b.layer, b.dx, try r.dev(std.mem.sliceAsBytes(p64b)), try r.dev(std.mem.sliceAsBytes(w64b)), 64)).out, ref.D * 4); + try check.sameBytes("exl3.experts.rows 64 rows mixed", mixed[0 .. ref.D * 4], solo); +} + +/// The same 300 rows as one call, split at 128 and at 37: a row's bytes never move. +fn chunks(r: *Rig, exp: *const exl3e.Experts) !void { + var bits: [ref.E]u32 = .{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }; + const b = try ref.buildLayer(r, &bits); + const whole = try r.back(try ref.runRows(r, exp, &b, 300), 300 * ref.D * 4); + var at: usize = 0; + while (at < 300) : (at += 128) { + const n: usize = @min(300 - at, 128); + const res = try ref.runSlice(r, exp, b.layer, b.dx + @as(u64, at) * ref.D * 2, b.dpick + @as(u64, at) * ref.SLOTS * 4, b.dwts + @as(u64, at) * ref.SLOTS * 4, n); + const got = try r.back(res.out, n * ref.D * 4); + try check.sameBytes("exl3.prompt.experts.chunks 128", got, whole[at * ref.D * 4 ..][0 .. n * ref.D * 4]); + } + at = 0; + while (at < 300) : (at += 37) { + const n: usize = @min(300 - at, 37); + const res = try ref.runSlice(r, exp, b.layer, b.dx + @as(u64, at) * ref.D * 2, b.dpick + @as(u64, at) * ref.SLOTS * 4, b.dwts + @as(u64, at) * ref.SLOTS * 4, n); + const got = try r.back(res.out, n * ref.D * 4); + try check.sameBytes("exl3.prompt.experts.chunks 37", got, whole[at * ref.D * 4 ..][0 .. n * ref.D * 4]); + } +} + +/// The Python engine's recorded y and out, bit for bit (exl3_oracle_experts.py's manifest and raw files). +fn python(r: *Rig, exp: *const exl3e.Experts, dir: []const u8, state: *Run) !void { + if (std.mem.eql(u8, dir, "none")) { + std.debug.print("SKIP exl3.experts.python\n", .{}); + return; + } + const manifest_path = try std.fs.path.join(r.a, &.{ dir, "manifest.json" }); + const text = std.Io.Dir.cwd().readFileAlloc(r.gpu.io, manifest_path, r.a, .limited(1 << 24)) catch |e| { + std.debug.print("FAIL exl3.experts.python: {t}\n", .{e}); + state.failed = true; + return; + }; + var parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + const o = parsed.value.object.get("experts").?.object; + const e_count: usize = @intCast(o.get("count").?.integer); + const dims: usize = @intCast(o.get("dims").?.integer); + const width: usize = @intCast(o.get("width").?.integer); + const slots: usize = @intCast(o.get("slots").?.integer); + const cb: exl3.Codebook = if (std.mem.eql(u8, o.get("codebook").?.string, "mcg")) + .mcg + else if (std.mem.eql(u8, o.get("codebook").?.string, "inst3")) + .inst3 + else + .mul1; + const trellis = o.get("trellis").?.object; + var gate_p = try r.a.alloc(u64, e_count); + var up_p = try r.a.alloc(u64, e_count); + var down_p = try r.a.alloc(u64, e_count); + var gk2 = try r.a.alloc(i32, e_count); + var uk2 = try r.a.alloc(i32, e_count); + var dk2 = try r.a.alloc(i32, e_count); + for (o.get("gate_k2").?.array.items, 0..) |v, i| gk2[i] = @intCast(v.integer); + for (o.get("up_k2").?.array.items, 0..) |v, i| uk2[i] = @intCast(v.integer); + for (o.get("down_k2").?.array.items, 0..) |v, i| dk2[i] = @intCast(v.integer); + for (0..e_count) |i| { + gate_p[i] = try r.dev(try readFile(r, dir, trellis.get("gate").?.array.items[i].string)); + up_p[i] = try r.dev(try readFile(r, dir, trellis.get("up").?.array.items[i].string)); + down_p[i] = try r.dev(try readFile(r, dir, trellis.get("down").?.array.items[i].string)); + } + const layer: exl3e.Layer = .{ + .gate = try r.dev(std.mem.sliceAsBytes(gate_p)), + .up = try r.dev(std.mem.sliceAsBytes(up_p)), + .down = try r.dev(std.mem.sliceAsBytes(down_p)), + .gate_k2 = try r.dev(std.mem.sliceAsBytes(gk2)), + .up_k2 = try r.dev(std.mem.sliceAsBytes(uk2)), + .down_k2 = try r.dev(std.mem.sliceAsBytes(dk2)), + .suh_g = try r.dev(try readFile(r, dir, o.get("suh_g").?.string)), + .suh_u = try r.dev(try readFile(r, dir, o.get("suh_u").?.string)), + .svh_g = try r.dev(try readFile(r, dir, o.get("svh_g").?.string)), + .svh_u = try r.dev(try readFile(r, dir, o.get("svh_u").?.string)), + .suh_d = try r.dev(try readFile(r, dir, o.get("suh_d").?.string)), + .svh_d = try r.dev(try readFile(r, dir, o.get("svh_d").?.string)), + .experts = e_count, + .dims = dims, + .width = width, + .codebook = cb, + .limit = @floatCast(o.get("limit").?.float), + .act_mode = @intCast(o.get("act_mode").?.integer), + }; + const xmap = o.get("x").?.object; + const pickmap = o.get("pick").?.object; + const wtsmap = o.get("wts").?.object; + const ymap = o.get("y").?.object; + const outmap = o.get("out").?.object; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const dx = try r.dev(try readFile(r, dir, xmap.get(key).?.string)); + const dpick = try r.dev(try readFile(r, dir, pickmap.get(key).?.string)); + const dwts = try r.dev(try readFile(r, dir, wtsmap.get(key).?.string)); + const sb = try r.zeros(try exl3e.scratchBytes(layer, m, slots)); + const sc = try exl3e.scratch(sb, layer, m, slots); + const out = try r.zeros(m * dims * 4); + try exp.run(r.s, layer, dx, dims, .bf16, dpick, dwts, m, out, sc); + const got_y = try r.back(sc.y, m * slots * dims * 4); + try check.sameBytes("exl3.experts.python y", got_y, try readFile(r, dir, ymap.get(key).?.string)); + const got_out = try r.back(out, m * dims * 4); + try check.sameBytes("exl3.experts.python out", got_out, try readFile(r, dir, outmap.get(key).?.string)); + } +} + +fn readFile(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + return std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, name }), r.a, .limited(1 << 32)); +} diff --git a/zig/tests/cuda/exl3_layer_tests.zig b/zig/tests/cuda/exl3_layer_tests.zig new file mode 100644 index 0000000000..4282f83ecd --- /dev/null +++ b/zig/tests/cuda/exl3_layer_tests.zig @@ -0,0 +1,156 @@ +//! EXL3 whole layers: a KDA and a DSA MoE layer of real RED-SNOW through core/exl3_load vs the engine's fixture bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_load = @import("core").exl3_load; +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const exl3_upload = cuda.exl3_upload; +const exl3_tests = @import("exl3_tests.zig"); +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + std.debug.print("SKIP exl3.layer.python\n", .{}); + return; + } + const o = manifest(r, dir) catch |e| { + std.debug.print("FAIL exl3.layer.python: {t}\n", .{e}); + state.failed = true; + return; + }; + const layers = o.get("layers") orelse { + std.debug.print("SKIP exl3.layer.python: no layers entry\n", .{}); + return; + }; + const model = o.get("model").?.string; + var scan = exl3_load.scan(r.gpu.gpa, r.gpu.io, model) catch |e| { + std.debug.print("FAIL exl3.layer.python: {t}\n", .{e}); + state.failed = true; + return; + }; + defer scan.deinit(); + var idx = std.StringHashMapUnmanaged(usize).empty; + for (scan.groups, 0..) |g, i| try idx.put(r.a, g.prefix, i); + var dense = try exl3.Dense.load(r.gpu.d); + defer dense.unload(); + var exp = try exl3e.Experts.load(r.gpu.d); + defer exp.unload(); + var ok = true; + layerChecks(r, &scan, &idx, &dense, &exp, dir, layers.array.items) catch |e| { + state.fail(e, "exl3.layer.python"); + ok = false; + }; + if (ok) check.pass("exl3.layer.python", .{}); +} + +fn layerChecks(r: *Rig, scan: *exl3_load.Scan, idx: *std.StringHashMapUnmanaged(usize), dense: *exl3.Dense, exp: *const exl3e.Experts, dir: []const u8, layers: []std.json.Value) !void { + const t0 = check.now(r.gpu.io); + var total: usize = 0; + for (layers) |item| { + const o = item.object; + const id = o.get("id").?.string; + var up: exl3_upload.Uploader = .{ .d = r.gpu.d, .gpa = r.gpu.gpa }; + defer up.deinit(); + for (o.get("groups").?.array.items) |gitem| try denseGroup(r, scan, idx, dense, dir, &up, gitem.object); + try expertLayer(r, scan, idx, exp, dir, &up, o.get("experts").?.object); + std.debug.print("RESULT exl3.layer.python.{s}: {d:.1} MiB of weights on the device\n", .{ id, @as(f64, @floatFromInt(up.bytes)) / (1 << 20) }); + total += up.bytes; + } + const dt = @as(f64, @floatFromInt(check.now(r.gpu.io) - t0)) / 1e9; + std.debug.print("RESULT exl3.layer.python: {d} layers, {d:.1} MiB uploaded, {d:.2} s\n", .{ layers.len, @as(f64, @floatFromInt(total)) / (1 << 20), dt }); +} + +/// One non-routed group: checkpoint tensors through the uploader, fixture inputs and bytes through the prompt path. +fn denseGroup(r: *Rig, scan: *exl3_load.Scan, idx: *std.StringHashMapUnmanaged(usize), dense: *exl3.Dense, dir: []const u8, up: *exl3_upload.Uploader, o: std.json.ObjectMap) !void { + const name = o.get("name").?.string; + const gi = idx.get(name) orelse return error.TestFailed; + const g = scan.groups[gi]; + try check.expect(g.halves == 2 * bitsOf(o.get("bits").?) and g.k == @as(usize, @intCast(o.get("k").?.integer)) and + g.n == @as(usize, @intCast(o.get("n").?.integer)), "exl3.layer.python {s}: the checkpoint group is not the fixture's", .{name}); + const parts = try scan.parts(r.a, gi); + const w = try up.dense(parts, g); + const xmap = o.get("x").?.object; + const ymap = o.get("y").?.object; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const dx = try r.dev(try readFile(r, dir, xmap.get(key).?.string)); + const want = try readFile(r, dir, ymap.get(key).?.string); + const dy = try r.zeros(m * g.n * 2); + try dense.prompt(r.s, w, dx, .bf16, m, dy, .bf16); + const got = try r.back(dy, want.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "exl3.layer.python {s} {d} rows", .{ name, m }), got, want); + } +} + +/// One layer's routed experts: all E experts' projections uploaded, the fixture's routing and inputs per row count. +fn expertLayer(r: *Rig, scan: *exl3_load.Scan, idx: *std.StringHashMapUnmanaged(usize), exp: *const exl3e.Experts, dir: []const u8, up: *exl3_upload.Uploader, o: std.json.ObjectMap) !void { + const base = o.get("name").?.string; + const e_count: usize = @intCast(o.get("count").?.integer); + const dims: usize = @intCast(o.get("dims").?.integer); + const width: usize = @intCast(o.get("width").?.integer); + const slots: usize = @intCast(o.get("slots").?.integer); + var gate = try r.a.alloc(exl3_upload.Proj, e_count); + var upj = try r.a.alloc(exl3_upload.Proj, e_count); + var down = try r.a.alloc(exl3_upload.Proj, e_count); + const shorts = [_][]const u8{ "gate", "up", "down" }; + const parts_names = [_][]const u8{ "gate_proj", "up_proj", "down_proj" }; + for (shorts, parts_names) |short, part| { + const k2map = o.get(try std.fmt.allocPrint(r.a, "{s}_k2", .{short})).?.array; + const list: *[]exl3_upload.Proj = switch (short[0]) { + 'g' => &gate, + 'u' => &upj, + else => &down, + }; + for (0..e_count) |e| { + const prefix = try std.fmt.allocPrint(r.a, "{s}.{d}.{s}", .{ base, e, part }); + const gi = idx.get(prefix) orelse return error.TestFailed; + const g = scan.groups[gi]; + try check.expect(g.halves == bitsOf(k2map.items[e]), "exl3.layer.python {s}: expert {d} is {d} halves, the fixture says {d}", .{ prefix, e, g.halves, bitsOf(k2map.items[e]) }); + const parts = try scan.parts(r.a, gi); + list.*[e] = .{ .trellis = parts.trellis.bytes, .suh = parts.suh.bytes, .svh = parts.svh.bytes, .k2 = g.halves }; + } + } + const l = try exl3_upload.layer(up, r.a, dims, width, .mul1, @floatCast(o.get("limit").?.float), @intCast(o.get("act_mode").?.integer), gate, upj, down); + const xmap = o.get("x").?.object; + const pickmap = o.get("pick").?.object; + const wtsmap = o.get("wts").?.object; + const ymap = o.get("y").?.object; + const outmap = o.get("out").?.object; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const dx = try r.dev(try readFile(r, dir, xmap.get(key).?.string)); + const dpick = try r.dev(try readFile(r, dir, pickmap.get(key).?.string)); + const dwts = try r.dev(try readFile(r, dir, wtsmap.get(key).?.string)); + const sb = try r.zeros(try exl3e.scratchBytes(l, m, slots)); + const sc = try exl3e.scratch(sb, l, m, slots); + const out = try r.zeros(m * dims * 4); + try exp.run(r.s, l, dx, dims, .bf16, dpick, dwts, m, out, sc); + const got_out = try r.back(out, m * dims * 4); + try check.sameBytes(try std.fmt.allocPrint(r.a, "exl3.layer.python {s} experts {d} rows out", .{ base, m }), got_out, try readFile(r, dir, outmap.get(key).?.string)); + const got_y = try r.back(sc.y, m * slots * dims * 4); + try check.sameBytes(try std.fmt.allocPrint(r.a, "exl3.layer.python {s} experts {d} rows y", .{ base, m }), got_y, try readFile(r, dir, ymap.get(key).?.string)); + } +} + +fn manifest(r: *Rig, dir: []const u8) !std.json.ObjectMap { + const text = try std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, "manifest.json" }), r.a, .limited(1 << 24)); + const parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + return parsed.value.object; +} + +fn bitsOf(v: std.json.Value) u32 { + return switch (v) { + .integer => |i| @intCast(i), + .float => |f| @intFromFloat(f), + else => unreachable, + }; +} + +fn readFile(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + return std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, name }), r.a, .limited(1 << 30)); +} diff --git a/zig/tests/cuda/exl3_tests.zig b/zig/tests/cuda/exl3_tests.zig new file mode 100644 index 0000000000..2f21c5e16e --- /dev/null +++ b/zig/tests/cuda/exl3_tests.zig @@ -0,0 +1,400 @@ +//! EXL3: the decode-row linear against float64 references, row invariance, the engine's fixture bytes, bandwidth. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3fmt = @import("core").exl3_format; +const exl3 = cuda.exl3; +const Gpu = check.Gpu; +const Rig = quant_ref_tests.RigT; +const bf = quant_ref_tests.helpers.bfOf; + +const rows = [_]usize{ 1, 3, 8, 16, 17, 64, 128 }; // every row count in the f64 checks (128 covers the rest) + +const Cfg = struct { id: []const u8, halves: u32, codebook: exl3fmt.Codebook, k: usize, n: usize }; + +const f64_cfgs = [_]Cfg{ + .{ .id = "exl3.linear.f64.mul1.4", .halves = 8, .codebook = .mul1, .k = 4096, .n = 1024 }, + .{ .id = "exl3.linear.f64.mul1.4", .halves = 2, .codebook = .mul1, .k = 512, .n = 512 }, // 1 bit + .{ .id = "exl3.linear.f64.mul1.4", .halves = 4, .codebook = .mul1, .k = 512, .n = 512 }, // 2 bits + .{ .id = "exl3.linear.f64.mul1.5", .halves = 10, .codebook = .mul1, .k = 1024, .n = 4096 }, + .{ .id = "exl3.linear.f64.mul1.5", .halves = 6, .codebook = .mul1, .k = 512, .n = 512 }, // 3 bits + .{ .id = "exl3.linear.f64.mul1.6", .halves = 12, .codebook = .mul1, .k = 2048, .n = 8192 }, // 8 warps a block + .{ .id = "exl3.linear.f64.mul1.7", .halves = 14, .codebook = .mul1, .k = 512, .n = 1024 }, + .{ .id = "exl3.linear.f64.mul1.8", .halves = 16, .codebook = .mul1, .k = 2048, .n = 4096 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 8, .codebook = .mcg, .k = 1024, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 2, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 4, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 6, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 10, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 12, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 14, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 16, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 8, .codebook = .inst3, .k = 1024, .n = 1024 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 2, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 4, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 6, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 10, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 12, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 14, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 16, .codebook = .inst3, .k = 512, .n = 512 }, + // the mul1 half steps, 1.5 to 6.5 bits, under one check + .{ .id = "exl3.linear.f64.halfstep", .halves = 3, .codebook = .mul1, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 5, .codebook = .mul1, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 7, .codebook = .mul1, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 9, .codebook = .mul1, .k = 1024, .n = 1024 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 11, .codebook = .mul1, .k = 2048, .n = 1024 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 13, .codebook = .mul1, .k = 512, .n = 512 }, +}; + +/// One check's outcome: a failure prints its FAIL line (TestFailed ones already printed theirs) and fails the run. +pub const Run = struct { + failed: bool = false, + + pub fn fail(self: *Run, e: anyerror, id: []const u8) void { + if (e != error.TestFailed) std.debug.print("FAIL {s}: {t}\n", .{ id, e }); + self.failed = true; + } +}; + +pub fn run(gpu: Gpu, dir: []const u8) !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(11); + 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 lin = try exl3.Linear.load(gpu.d); + defer lin.unload(); + + var state = Run{}; + var last: ?[]const u8 = null; + var ok = true; + for (f64_cfgs) |c| { + if (last) |prev| { + if (!std.mem.eql(u8, prev, c.id)) { + if (ok) check.pass("{s}", .{prev}); + ok = true; + } + } + f64Linear(&r, &lin, c) catch |e| { + state.fail(e, c.id); + ok = false; + }; + last = c.id; + } + if (last) |prev| { + if (ok) check.pass("{s}", .{prev}); + } + { + var rows_ok = true; + rowsInvariant(&r, &lin) catch |e| { + state.fail(e, "exl3.linear.rows"); + rows_ok = false; + }; + if (rows_ok) check.pass("exl3.linear.rows", .{}); + } + pythonChecks(&r, &lin, dir, &state) catch |e| { + state.fail(e, "exl3.linear.python"); + }; + try @import("exl3_dense_tests.zig").run(&r, &state, dir); + try @import("exl3_experts_tests.zig").run(&r, &state, dir); + try @import("exl3_layer_tests.zig").run(&r, &state, dir); + if (state.failed) return error.TestFailed; + try bench(&r, &lin); +} + +/// A zeroed device buffer (counters must read zero, unlike zeros()'s 0xff canaries). +pub fn zeroed(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(bytes, 16)); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0, r.s.handle); + return b.ptr; +} + +/// One codebook and width at one shape: random trellis, scales and 128 rows of x, then every row count of `rows`. +fn f64Linear(r: *Rig, lin: *const exl3.Linear, c: Cfg) !void { + const width = exl3fmt.Width{ .halves = @intCast(c.halves) }; + const cb: exl3.Codebook = switch (c.codebook) { + .inst3 => .inst3, + .mcg => .mcg, + .mul1 => .mul1, + }; + const k = c.k; + const n = c.n; + const trellis = try r.a.alloc(i16, (k / 16) * (n / 16) * width.tileWords()); + for (trellis) |*v| v.* = r.rng.int(i16); + const suh = try r.a.alloc(u16, k); + const svh = try r.a.alloc(u16, n); + for (suh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + for (svh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + const x = try r.bfs(128 * k, 1.0); + const want = try reference(r, width, c.codebook, trellis, k, n, x, suh, svh); + const st = exl3.strides(k, n, c.halves, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(std.mem.sliceAsBytes(trellis)), + .suh = try r.dev(std.mem.sliceAsBytes(suh)), + .svh = try r.dev(std.mem.sliceAsBytes(svh)), + .bias = 0, + .k = @intCast(k), + .n = @intCast(n), + .halves = c.halves, + .codebook = cb, + .stride_k = st.k, + .stride_nb = st.nb, + }; + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const counters = try zeroed(r, exl3.Linear.counterBytes(w)); + for (rows) |m| { + const dy = try r.zeros(m * n * 2); + const xh = try r.zeros(m * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, m)) else 0; + try lin.matmul(r.s, w, dx, .bf16, m, dy, .bf16, xh, z, counters); + try closeOut(r, c.id, dy, m * n, want.want, want.tol); + } +} + +/// The Python engine's f64 reference: Wq fp16 from the trellis, fp16 xh, the y sums, and each element's |t| sum. +fn reference(r: *Rig, width: exl3fmt.Width, cb: exl3fmt.Codebook, trellis: []const i16, k: usize, n: usize, x: []const u16, suh: []const u16, svh: []const u16) !struct { want: []f64, tol: []f64 } { + const m = 128; + const wq = try r.a.alloc(f16, k * n); + exl3fmt.unpack(width, cb, trellis, k / 16, n / 16, wq); + const xs = try r.a.alloc(f64, m * k); + for (0..m) |row| { + for (0..k) |i| xs[row * k + i] = bf(x[row * k + i]) * @as(f64, @floatCast(@as(f16, @bitCast(suh[i])))); + } + exl3fmt.rotate(xs, m, k, 1); + const xh = try r.a.alloc(f16, m * k); + for (xh, xs) |*o, v| o.* = @floatCast(v); + const ypre = try r.a.alloc(f64, m * n); // [n][m], rotated along n + @memset(ypre, 0); + const s1 = try r.a.alloc(f64, m * n); // each output's sum of |xh * wq| over k + @memset(s1, 0); + var xv: [128]f64 = undefined; + var xa: [128]f64 = undefined; + for (0..k) |ki| { + for (0..m) |row| { + xv[row] = @floatCast(xh[row * k + ki]); + xa[row] = @abs(xv[row]); + } + for (0..n) |col| { + const wv: f64 = @floatCast(wq[ki * n + col]); + const wa = @abs(wv); + const base = col * m; + for (0..m) |row| { + ypre[base + row] += xv[row] * wv; + s1[base + row] += xa[row] * wa; + } + } + } + exl3fmt.rotate(ypre, n, m, 0); + const want = try r.a.alloc(f64, m * n); + const tol = try r.a.alloc(f64, m * n); + for (0..n) |col| { + const scale: f64 = @floatCast(@as(f16, @bitCast(svh[col]))); + for (0..m) |row| { + want[row * n + col] = ypre[col * m + row] * scale; + // bf16 output ulp (2^-7 of the value) plus fp32 accumulation and fp16 xh cases, 2^-11 of the |t| sum + tol[row * n + col] = @abs(want[row * n + col]) / 128 + s1[row * n + col] / 2048 + 1e-30; + } + } + return .{ .want = want, .tol = tol }; +} + +/// bf16 outputs within the f64 reference's tolerance, with the exact FAIL line. +pub fn closeOut(r: *Rig, id: []const u8, ptr: u64, len: usize, want: []const f64, tol: []const f64) !void { + const raw = try r.back(ptr, len * 2); + var most: f64 = 0; + var at: usize = 0; + for (want[0..len], tol[0..len], 0..) |wv, t, i| { + const g = bf(std.mem.bytesToValue(u16, raw[2 * i ..][0..2])); + const q = @abs(g - wv) / t; + if (q > most) { + most = q; + at = i; + } + } + std.debug.print("RESULT {s}: largest difference {d:.3} of its tolerance (element {d})\n", .{ id, most, at }); + if (most > 1) { + std.debug.print("FAIL {s}: largest difference {d:.3} of its tolerance\n", .{ id, most }); + return error.TestFailed; + } +} + +/// The first m rows of a 128-row call byte-equal every narrower call, at a split shape and a single-split one. +fn rowsInvariant(r: *Rig, lin: *const exl3.Linear) !void { + const shapes = [_][2]usize{ .{ 4096, 1024 }, .{ 512, 4096 } }; + for (shapes) |sh| { + const k = sh[0]; + const n = sh[1]; + const width: exl3fmt.Width = .{ .halves = 8 }; + const trellis = try r.a.alloc(i16, (k / 16) * (n / 16) * width.tileWords()); + for (trellis) |*v| v.* = r.rng.int(i16); + const suh = try r.a.alloc(u16, k); + const svh = try r.a.alloc(u16, n); + for (suh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + for (svh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + const st = exl3.strides(k, n, 8, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(std.mem.sliceAsBytes(trellis)), + .suh = try r.dev(std.mem.sliceAsBytes(suh)), + .svh = try r.dev(std.mem.sliceAsBytes(svh)), + .bias = 0, + .k = @intCast(k), + .n = @intCast(n), + .halves = 8, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; + const x = try r.bfs(128 * k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const counters = try zeroed(r, exl3.Linear.counterBytes(w)); + const y128 = try r.zeros(128 * n * 2); + { + const xh = try r.zeros(128 * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, 128)) else 0; + try lin.matmul(r.s, w, dx, .bf16, 128, y128, .bf16, xh, z, counters); + } + const got128 = try r.back(y128, 128 * n * 2); + for ([_]usize{ 1, 2, 3, 16, 17, 64 }) |m| { + const ym = try r.zeros(m * n * 2); + const xh = try r.zeros(m * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, m)) else 0; + try lin.matmul(r.s, w, dx, .bf16, m, ym, .bf16, xh, z, counters); + const got = try r.back(ym, m * n * 2); + const id = try std.fmt.allocPrint(r.a, "exl3.linear.rows k={d} n={d} m={d}", .{ k, n, m }); + try check.sameBytes(id, got, got128[0 .. m * n * 2]); + } + } +} + +/// The Python engine's recorded outputs, bit for bit (exl3_oracle.py's manifest and raw files). +fn pythonChecks(r: *Rig, lin: *const exl3.Linear, dir: []const u8, state: *Run) !void { + if (std.mem.eql(u8, dir, "none")) { + for (4..9) |bits| std.debug.print("SKIP exl3.linear.python.{d}\n", .{bits}); + return; + } + const manifest_path = try std.fs.path.join(r.a, &.{ dir, "manifest.json" }); + const text = std.Io.Dir.cwd().readFileAlloc(r.gpu.io, manifest_path, r.a, .limited(1 << 24)) catch |e| { + for (4..9) |bits| { + std.debug.print("FAIL exl3.linear.python.{d}: {t}\n", .{ bits, e }); + state.failed = true; + } + return; + }; + var parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + var seen = [_]bool{ false, false, false, false, false, false, false, false, false }; + for (parsed.value.object.get("linears").?.array.items) |item| { + const o = item.object; + const bits: u32 = bitsOf(o.get("bits").?); + seen[bits] = true; + const id = try std.fmt.allocPrint(r.a, "exl3.linear.python.{d}", .{bits}); + var ok = true; + pythonLinear(r, lin, dir, o, id) catch |e| { + state.fail(e, id); + ok = false; + }; + if (ok) check.pass("{s}", .{id}); + } + for (4..9) |bits| { + if (!seen[bits]) std.debug.print("SKIP exl3.linear.python.{d}\n", .{bits}); + } +} + +fn bitsOf(v: std.json.Value) u32 { + return switch (v) { + .integer => |i| @intCast(i), + .float => |f| @intFromFloat(f), + else => unreachable, + }; +} + +fn pythonLinear(r: *Rig, lin: *const exl3.Linear, dir: []const u8, o: std.json.ObjectMap, id: []const u8) !void { + const k: usize = @intCast(o.get("k").?.integer); + const n: usize = @intCast(o.get("n").?.integer); + const bits: u32 = bitsOf(o.get("bits").?); + const cb: exl3.Codebook = if (std.mem.eql(u8, o.get("codebook").?.string, "mcg")) + .mcg + else if (std.mem.eql(u8, o.get("codebook").?.string, "inst3")) + .inst3 + else + .mul1; + const bias = o.get("bias").?; + const st = exl3.strides(k, n, 2 * bits, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(try readFile(r, dir, o.get("trellis").?.string)), + .suh = try r.dev(try readFile(r, dir, o.get("suh").?.string)), + .svh = try r.dev(try readFile(r, dir, o.get("svh").?.string)), + .bias = if (std.meta.activeTag(bias) == .null) 0 else try r.dev(try readFile(r, dir, bias.string)), + .k = @intCast(k), + .n = @intCast(n), + .halves = 2 * bits, + .codebook = cb, + .stride_k = st.k, + .stride_nb = st.nb, + }; + const xmap = o.get("x").?.object; + const ymap = o.get("y").?.object; + const counters = try zeroed(r, exl3.Linear.counterBytes(w)); + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const x = try readFile(r, dir, xmap.get(key).?.string); + const want = try readFile(r, dir, ymap.get(key).?.string); + const dy = try r.zeros(m * n * 2); + const xh = try r.zeros(m * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, m)) else 0; + try lin.matmul(r.s, w, try r.dev(x), .bf16, m, dy, .bf16, xh, z, counters); + const got = try r.back(dy, want.len); + try check.sameBytes(id, got, want); + } +} + +pub fn readFile(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + return std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, name }), r.a, .limited(1 << 30)); +} + +/// Trellis bytes a second: one row of a 4-bit 4096x32768 layer (64 MiB, past L2), decode-dominated. +fn bench(r: *Rig, lin: *const exl3.Linear) !void { + const k = 4096; + const n = 32768; + const bytes = k * n / 2; + const st = exl3.strides(k, n, 8, false); + const suh = try r.a.alloc(u16, k); + const svh = try r.a.alloc(u16, n); + @memset(suh, 0x3C00); // 1.0 + @memset(svh, 0x3C00); + const x = try r.bfs(k, 1.0); + const w: exl3.Weight = .{ + .trellis = try r.zeros(bytes), + .suh = try r.dev(std.mem.sliceAsBytes(suh)), + .svh = try r.dev(std.mem.sliceAsBytes(svh)), + .bias = 0, + .k = k, + .n = n, + .halves = 8, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dy = try r.zeros(n * 2); + const xh = try r.zeros(k * 2); + const z = try r.zeros(exl3.Linear.zBytes(w, 1)); + const counters = try zeroed(r, exl3.Linear.counterBytes(w)); + for (0..3) |_| try lin.matmul(r.s, w, dx, .bf16, 1, dy, .bf16, xh, z, counters); + try r.s.synchronize(); + const reps = 200; + const t0 = check.now(r.gpu.io); + for (0..reps) |_| try lin.matmul(r.s, w, dx, .bf16, 1, dy, .bf16, xh, z, counters); + try r.s.synchronize(); + const t1 = check.now(r.gpu.io); + const sec = @as(f64, @floatFromInt(t1 - t0)) / 1e9; + const gb = @as(f64, @floatFromInt(bytes * reps)) / sec / 1e9; + std.debug.print("RESULT exl3.bench.decode {d:.1}\n", .{gb}); +} diff --git a/zig/tests/cuda/exl3_tp2_tests.zig b/zig/tests/cuda/exl3_tp2_tests.zig new file mode 100644 index 0000000000..ef57f1eb1c --- /dev/null +++ b/zig/tests/cuda/exl3_tp2_tests.zig @@ -0,0 +1,325 @@ +//! EXL3 tensor parallel: each rank's shard of real RED-SNOW projections and experts vs the engine's, all-gathered. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const exl3t = @import("exl3_tests.zig"); +const Gpu = check.Gpu; +const Rig = quant_ref_tests.RigT; + +/// Both ranks over NCCL (nccl_id.pair: rank 0 serves the unique id on host:port, the benchmark starts rank 1 first). +pub fn run(gpu: Gpu, rank_arg: []const u8, addr: []const u8, dir: []const u8) !void { + const rank = try std.fmt.parseInt(u8, rank_arg, 10); + if (rank > 1) { + std.debug.print("usage: tf-cuda-test exl3-tp2 \n", .{}); + return error.BadArguments; + } + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + const comm = try cuda.nccl_id.pair(&lib, gpu.io, rank, addr); + defer _ = lib.api.ncclCommDestroy(comm); + + 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(11); + 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 lin = try exl3.Linear.load(gpu.d); + defer lin.unload(); + var exp = try exl3e.Experts.load(gpu.d); + defer exp.unload(); + + var failed = false; + { + var ok = true; + commSanity(&r, &lib, comm) catch |e| { + fail0(rank, "exl3.tp2.comm", e); + ok = false; + }; + if (!ok) failed = true; + } + if (std.mem.eql(u8, dir, "none")) { + if (rank == 0) { + std.debug.print("SKIP exl3.tp2.linear\n", .{}); + std.debug.print("SKIP exl3.tp2.experts\n", .{}); + } + } else { + linearChecks(&r, &lib, comm, &lin, dir, rank, &failed); + expertsCheck(&r, &lib, comm, &exp, dir, rank, &failed); + } + if (failed) return error.TestFailed; + std.debug.print("RESULT exl3.tp2 rank {d}: every collective and comparison agreed\n", .{rank}); +} + +/// One f32 all-gather: the communicator moves bits before anything else is trusted. +fn commSanity(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm) !void { + const one = [_]u8{ 0, 0, 0x80, 0x3f }; + const send = try r.dev(&one); + const recv = try r.zeros(8); + try lib.check(lib.api.ncclAllGather(send, recv, 1, .f32, comm, r.s.handle), "ncclAllGather"); + try r.s.synchronize(); + const got = try r.back(recv, 8); + try same(0, "exl3.tp2.comm all-gather", got, &.{ 0, 0, 0x80, 0x3f, 0, 0, 0x80, 0x3f }); +} + +fn fail0(rank: u8, id: []const u8, e: anyerror) void { + if (rank == 0) std.debug.print("FAIL {s}: {t}\n", .{ id, e }); +} + +/// Equal bytes, or the first difference; the FAIL line is rank 0's (the contract's one printer). +fn same(rank: u8, what: []const u8, got: []const u8, want: []const u8) !void { + if (std.mem.eql(u8, got, want)) return; + if (rank == 0) { + if (got.len != want.len) { + std.debug.print("FAIL {s}: {d} bytes, expected {d}\n", .{ what, got.len, want.len }); + } else { + std.debug.print("FAIL {s}: first difference at byte {d}\n", .{ what, std.mem.indexOfDiff(u8, got, want).? }); + } + } + return error.TestFailed; +} + +fn bitsOf(v: std.json.Value) u32 { + return switch (v) { + .integer => |i| @intCast(i), + .float => |f| @intFromFloat(f), + else => unreachable, + }; +} + +fn codebook(name: []const u8) exl3.Codebook { + if (std.mem.eql(u8, name, "mcg")) return .mcg; + if (std.mem.eql(u8, name, "inst3")) return .inst3; + return .mul1; +} + +fn tp2Section(r: *Rig, dir: []const u8) !std.json.ObjectMap { + const text = try std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, "manifest.json" }), r.a, .limited(1 << 24)); + var parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + return parsed.value.object.get("tp2").?.object; +} + +/// The two fixture linears, the out-split ('dim1') and the in-split ('row') one: each reports its own line. +fn linearChecks(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, lin: *const exl3.Linear, dir: []const u8, rank: u8, failed: *bool) void { + const tp2 = tp2Section(r, dir) catch |e| { + fail0(rank, "exl3.tp2.linear", e); + failed.* = true; + return; + }; + var all = true; + for (tp2.get("linears").?.array.items) |item| { + const o = item.object; + const id = std.fmt.allocPrint(r.a, "exl3.tp2.linear.{s}", .{o.get("id").?.string}) catch unreachable; + var ok = true; + linearCheck(r, lib, comm, lin, dir, rank, o, id) catch |e| { + if (e != error.TestFailed) fail0(rank, id, e); + ok = false; + }; + if (ok and rank == 0) check.pass("{s}", .{id}); + if (!ok) all = false; + } + if (all and rank == 0) check.pass("exl3.tp2.linear", .{}); + if (!all) failed.* = true; +} + +/// One projection: the rank's trellis/scale shard of the EXL3_RULES cut, its output, and the all-gather. +fn linearCheck(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, lin: *const exl3.Linear, dir: []const u8, rank: u8, o: std.json.ObjectMap, id: []const u8) !void { + const dim1 = std.mem.eql(u8, o.get("split").?.string, "dim1"); + const k: usize = @intCast(o.get("k").?.integer); + const n: usize = @intCast(o.get("n").?.integer); + const bits: u32 = bitsOf(o.get("bits").?); + const halves: usize = 2 * bits; + const kt = k / 16; + const nt = n / 16; + const tw = halves * 16; // bytes a trellis tile (8 * halves i16 values) + const full = try exl3t.readFile(r, dir, o.get("trellis").?.string); + const suh_full = try exl3t.readFile(r, dir, o.get("suh").?.string); + const svh_full = try exl3t.readFile(r, dir, o.get("svh").?.string); + var k_use = k; + var n_use = n; + // dim1: trellis tile columns and svh rows in half, suh whole; row: tile rows, suh rows and x's columns + const trellis = if (dim1) try cutDim1(r, full, kt, nt, tw, rank) else full[rank * (kt / 2) * nt * tw ..][0 .. (kt / 2) * nt * tw]; + const suh = if (dim1) suh_full else suh_full[rank * (k / 2) * 2 ..][0 .. (k / 2) * 2]; + const svh = if (dim1) svh_full[rank * (n / 2) * 2 ..][0 .. (n / 2) * 2] else svh_full; + if (dim1) n_use = n / 2 else k_use = k / 2; + const biasv = o.get("bias").?; + var bias: u64 = 0; + if (std.meta.activeTag(biasv) == .string) { + const bfull = try exl3t.readFile(r, dir, biasv.string); + bias = try r.dev(if (dim1) bfull[rank * (n / 2) * 2 ..][0 .. (n / 2) * 2] else bfull); + } + const st = exl3.strides(k_use, n_use, halves, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(trellis), + .suh = try r.dev(suh), + .svh = try r.dev(svh), + .bias = bias, + .k = @intCast(k_use), + .n = @intCast(n_use), + .halves = @intCast(halves), + .codebook = codebook(o.get("codebook").?.string), + .stride_k = st.k, + .stride_nb = st.nb, + }; + const counters = try exl3t.zeroed(r, exl3.Linear.counterBytes(w)); + const ydt: exl3.DType = if (dim1) .bf16 else .f32; + const esz: usize = if (dim1) 2 else 4; + const dt: cuda.nccl.DataType = if (dim1) .bf16 else .f32; + const xmap = o.get("x").?.object; + + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const what = try std.fmt.allocPrint(r.a, "{s} m={d}", .{ id, m }); + const xf = try exl3t.readFile(r, dir, xmap.get(key).?.string); + const x_use = if (dim1) xf else try cutDim1(r, xf, m, k, 2, rank); + const count = m * n_use; // elements this rank all-gathers + const send = try r.zeros(count * esz); + const xh = try r.zeros(m * k_use * 2); + const z: u64 = if (exl3.plan(k_use, n_use, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, m)) else 0; + try lin.matmul(r.s, w, try r.dev(x_use), .bf16, m, send, ydt, xh, z, counters); + const recv = try r.zeros(2 * count * esz); + try lib.check(lib.api.ncclAllGather(send, recv, count, dt, comm, r.s.handle), "ncclAllGather"); + try r.s.synchronize(); + const got = try r.back(recv, 2 * count * esz); + if (dim1) { + // the two ranks' bf16 halves concatenated are the full output + const it = o.get("shards").?.object.get(key).?.array.items; + const want = try std.mem.concat(r.a, u8, &.{ + try exl3t.readFile(r, dir, it[0].string), + try exl3t.readFile(r, dir, it[1].string), + }); + try same(rank, what, got, want); + } else { + // fp32 partials bit for bit, then bf16(sum rank 0 first) as forward.py's hc_post rounds + const it = o.get("partials").?.object.get(key).?.array.items; + try same(rank, what, got[0 .. m * n * 4], try exl3t.readFile(r, dir, it[0].string)); + try same(rank, what, got[m * n * 4 ..], try exl3t.readFile(r, dir, it[1].string)); + const want = try r.a.alloc(u8, m * n * 2); + for (0..m * n) |i| { + const g0 = std.mem.bytesToValue(f32, got[4 * i ..][0..4]); + const g1 = std.mem.bytesToValue(f32, got[4 * (m * n + i) ..][0..4]); + std.mem.writeInt(u16, want[2 * i ..][0..2], bf16r(g0 + g1), .little); + } + try same(rank, what, want, try exl3t.readFile(r, dir, o.get("y").?.object.get(key).?.string)); + } + } +} + +/// The routed experts of one MoE layer: each rank's gate/up/down shards, its fp32 partials, the all-gathered sum. +fn expertsCheck(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, exp: *const exl3e.Experts, dir: []const u8, rank: u8, failed: *bool) void { + const id = "exl3.tp2.experts"; + const o = blk: { + const tp2 = tp2Section(r, dir) catch |e| { + fail0(rank, id, e); + failed.* = true; + return; + }; + break :blk tp2.get("experts").?.object; + }; + var ok = true; + expertsRun(r, lib, comm, exp, dir, rank, o, id) catch |e| { + if (e != error.TestFailed) fail0(rank, id, e); + ok = false; + }; + if (ok and rank == 0) check.pass("{s}", .{id}); + if (!ok) failed.* = true; +} + +fn expertsRun(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, exp: *const exl3e.Experts, dir: []const u8, rank: u8, o: std.json.ObjectMap, id: []const u8) !void { + const e_count: usize = @intCast(o.get("count").?.integer); + const dims: usize = @intCast(o.get("dims").?.integer); + const width: usize = @intCast(o.get("width").?.integer); // the full layer's; each rank holds width / 2 + const slots: usize = @intCast(o.get("slots").?.integer); + const trellis = o.get("trellis").?.object; + var gate_p = try r.a.alloc(u64, e_count); + var up_p = try r.a.alloc(u64, e_count); + var down_p = try r.a.alloc(u64, e_count); + var gk2 = try r.a.alloc(i32, e_count); + var uk2 = try r.a.alloc(i32, e_count); + var dk2 = try r.a.alloc(i32, e_count); + for (o.get("gate_k2").?.array.items, 0..) |v, i| gk2[i] = @intCast(v.integer); + for (o.get("up_k2").?.array.items, 0..) |v, i| uk2[i] = @intCast(v.integer); + for (o.get("down_k2").?.array.items, 0..) |v, i| dk2[i] = @intCast(v.integer); + // gate/up: trellis tile columns and svh rows in half; down: tile rows and suh rows in half + for (0..e_count) |i| { + const gw = @as(usize, @intCast(gk2[i])) * 16; // bytes a gate/up tile + const dw = @as(usize, @intCast(dk2[i])) * 16; + gate_p[i] = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, trellis.get("gate").?.array.items[i].string), dims / 16, width / 16, gw, rank)); + up_p[i] = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, trellis.get("up").?.array.items[i].string), dims / 16, width / 16, gw, rank)); + const dfull = try exl3t.readFile(r, dir, trellis.get("down").?.array.items[i].string); + down_p[i] = try r.dev(dfull[rank * (width / 32) * (dims / 16) * dw ..][0 .. (width / 32) * (dims / 16) * dw]); + } + const layer: exl3e.Layer = .{ + .gate = try r.dev(std.mem.sliceAsBytes(gate_p)), + .up = try r.dev(std.mem.sliceAsBytes(up_p)), + .down = try r.dev(std.mem.sliceAsBytes(down_p)), + .gate_k2 = try r.dev(std.mem.sliceAsBytes(gk2)), + .up_k2 = try r.dev(std.mem.sliceAsBytes(uk2)), + .down_k2 = try r.dev(std.mem.sliceAsBytes(dk2)), + .suh_g = try r.dev(try exl3t.readFile(r, dir, o.get("suh_g").?.string)), + .suh_u = try r.dev(try exl3t.readFile(r, dir, o.get("suh_u").?.string)), + .svh_g = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, o.get("svh_g").?.string), e_count, width, 2, rank)), + .svh_u = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, o.get("svh_u").?.string), e_count, width, 2, rank)), + .suh_d = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, o.get("suh_d").?.string), e_count, width, 2, rank)), + .svh_d = try r.dev(try exl3t.readFile(r, dir, o.get("svh_d").?.string)), + .experts = e_count, + .dims = dims, + .width = width / 2, + .codebook = codebook(o.get("codebook").?.string), + .limit = @floatCast(o.get("limit").?.float), + .act_mode = @intCast(o.get("act_mode").?.integer), + }; + const xmap = o.get("x").?.object; + const pickmap = o.get("pick").?.object; + const wtsmap = o.get("wts").?.object; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const what = try std.fmt.allocPrint(r.a, "{s} m={d}", .{ id, m }); + const dx = try r.dev(try exl3t.readFile(r, dir, xmap.get(key).?.string)); + const dpick = try r.dev(try exl3t.readFile(r, dir, pickmap.get(key).?.string)); + const dwts = try r.dev(try exl3t.readFile(r, dir, wtsmap.get(key).?.string)); + const sb = try r.zeros(try exl3e.scratchBytes(layer, m, slots)); + const sc = try exl3e.scratch(sb, layer, m, slots); + const out = try r.zeros(m * dims * 4); + try exp.run(r.s, layer, dx, dims, .bf16, dpick, dwts, m, out, sc); + const recv = try r.zeros(2 * m * dims * 4); + try lib.check(lib.api.ncclAllGather(out, recv, m * dims, .f32, comm, r.s.handle), "ncclAllGather"); + try r.s.synchronize(); + const got = try r.back(recv, 2 * m * dims * 4); + const it = o.get("partials").?.object.get(key).?.array.items; + try same(rank, what, got[0 .. m * dims * 4], try exl3t.readFile(r, dir, it[0].string)); + try same(rank, what, got[m * dims * 4 ..], try exl3t.readFile(r, dir, it[1].string)); + const want = try r.a.alloc(u8, m * dims * 2); + for (0..m * dims) |i| { + const g0 = std.mem.bytesToValue(f32, got[4 * i ..][0..4]); + const g1 = std.mem.bytesToValue(f32, got[4 * (m * dims + i) ..][0..4]); + std.mem.writeInt(u16, want[2 * i ..][0..2], bf16r(g0 + g1), .little); + } + try same(rank, what, want, try exl3t.readFile(r, dir, o.get("y").?.object.get(key).?.string)); + } +} + +/// The rank's half of an even 2-D row-major tensor's second axis, repacked contiguous (split.py's dim1 cut). +fn cutDim1(r: *Rig, src: []const u8, rows: usize, cols: usize, item: usize, rank: u8) ![]u8 { + const per = cols / 2 * item; + const out = try r.a.alloc(u8, rows * per); + for (0..rows) |i| { + @memcpy(out[i * per ..][0..per], src[i * cols * item + @as(usize, rank) * per ..][0..per]); + } + return out; +} + +/// bf16 round-to-nearest-even of an fp32 sum (hc_post's `acc.to(torch.bfloat16)`). +fn bf16r(v: f32) u16 { + const u: u32 = @bitCast(v); + return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} diff --git a/zig/tests/cuda/glm_dsa_block_tests.zig b/zig/tests/cuda/glm_dsa_block_tests.zig new file mode 100644 index 0000000000..1cbee028ce --- /dev/null +++ b/zig/tests/cuda/glm_dsa_block_tests.zig @@ -0,0 +1,187 @@ +//! glm.dsa.block: a real RED-SNOW DSA layer at TP=2 rank 0 vs the Python engine's dsa_block, bit for bit. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const lref = @import("glm_layer_ref.zig"); +const dref = @import("glm_dsa_ref.zig"); +const dsa = cuda.glm_dsa; +const exl3 = cuda.exl3; + +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const Obj = std.json.ObjectMap; +const id = "glm.dsa.block"; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + block(r, dir) catch |e| state.fail(e, id); +} + +fn skip() void { + std.debug.print("SKIP {s}\n", .{id}); +} + +fn file(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + return ref.readFile(r.a, r.gpu.io, dir, name); +} + +/// Device bytes back to the host 8-aligned (the f64 chain reads them as u16/f32/i32 slices). +fn back(r: *Rig, ptr: u64, len: usize) ![]u8 { + const out = try r.a.alignedAlloc(u8, .@"8", len); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = len }; + try b.download(0, out); + return out; +} + +fn linW(r: *Rig, l: lref.Lin) !exl3.Weight { + const st = exl3.strides(l.k, l.n, l.halves, true); + return .{ .trellis = try r.dev(l.words), .suh = try r.dev(l.suh), .svh = try r.dev(l.svh), .k = @intCast(l.k), .n = @intCast(l.n), .halves = l.halves, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; +} + +fn b16W(r: *Rig, m: lref.Mat) !cuda.glm_bmm.B16 { + return .{ .w = try r.dev(m.w), .n = @intCast(m.n), .k = @intCast(m.k) }; +} + +/// The caches every window starts from: FILL tokens the real block wrote (kc fp8, scales, ik, ig, pk). +const names = [_][]const u8{ "kc", "scales", "ik", "ig", "pk" }; +const row_bytes = [_]usize{ dsa.kv_lora, 4, dsa.index_dim * 2, dsa.index_dim * 2, dsa.index_dim * 2 }; + +fn block(r: *Rig, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) return skip(); + var man = try ref.Manifest.open(r.gpu.gpa, r.gpu.io, dir); + defer man.deinit(); + const o = (man.root().get("dsa") orelse return skip()).object; + const p = o.get("params").?.object; + const c: dsa.Config = .{ + .eps = @floatCast(ref.float(p, "eps")), + .index_eps = @floatCast(ref.float(p, "index_eps")), + .scale = @floatCast(ref.float(p, "scale")), + .index_scale = @floatCast(ref.float(p, "index_scale")), + .wscale = @floatCast(ref.float(p, "wscale")), + }; + const p64: dref.Params = .{ .eps = c.eps, .index_eps = c.index_eps, .scale = c.scale, .index_scale = c.index_scale, .wscale = c.wscale }; + const cap: usize = @intCast(ref.int(p, "cap")); + const host = try dref.Host.read(r.a, r.gpu.io, dir, o.get("weights").?.object); + const w: dsa.DsaW = .{ + .q_a = try linW(r, host.q_a), + .kv_a = try linW(r, host.kv_a), + .q_b = try linW(r, host.q_b), + .o = try linW(r, host.o), + .q_norm = try r.dev(host.q_norm), + .kv_norm = try r.dev(host.kv_norm), + .wk = try r.dev(host.wk), + .wv = try r.dev(host.wv), + .index = .{ .kw = try b16W(r, host.kw), .qb = try b16W(r, host.qb), .gate = .{ .w = try r.dev(host.gate), .n = dsa.index_dim, .k = dsa.hidden }, .ln_w = try r.dev(host.ln_w), .ln_b = try r.dev(host.ln_b), .ape = try r.dev(host.ape) }, + }; + var pre: [names.len][]u8 = undefined; + for (names, 0..) |n, i| pre[i] = try file(r, dir, ref.str(o.get("pre").?.object, n)); + var glue2 = try cuda.glm_glue2.Glue2.load(r.gpu.d); + defer glue2.unload(); + var bmm = try cuda.glm_bmm.Bmm.load(r.gpu.d); + defer bmm.unload(); + var dense = try exl3.Dense.load(r.gpu.d); + defer dense.unload(); + var lat = try cuda.glm_latent.Latent.load(r.gpu.d); + defer lat.unload(); + var sp = try cuda.glm_sparse.Sparse.load(r.gpu.d); + defer sp.unload(); + const k: dsa.Kernels = .{ .glue2 = &glue2, .bmm = &bmm, .dense = &dense, .latent = &lat, .sparse = &sp }; + var worst: f64 = 0; + const cases = o.get("cases").?.object; + var it = cases.iterator(); + while (it.next()) |kv| { + const t0 = check.now(r.gpu.io); + // a window's host bytes live in its own arena + var arena = std.heap.ArenaAllocator.init(r.gpu.gpa); + defer arena.deinit(); + const saved = r.a; + r.a = arena.allocator(); + defer r.a = saved; + const q = try window(r, k, c, p64, &w, &host, &pre, cap, dir, kv.key_ptr.*, kv.value_ptr.object); + worst = @max(worst, q); + const dt = @as(f64, @floatFromInt(check.now(r.gpu.io) - t0)) / 1e9; + std.debug.print("RESULT {s} {s}: bit-exact, f64 worst {d:.4} of tolerance ({d:.1} s)\n", .{ id, kv.key_ptr.*, q, dt }); + } + std.debug.print("RESULT {s}: {d} windows bit-exact, f64 worst {d:.4} of tolerance\n", .{ id, cases.count(), worst }); + check.pass("{s}", .{id}); +} + +fn same(r: *Rig, tag: []const u8, what: []const u8, got: []const u8, want: []const u8) !void { + try check.sameBytes(try std.fmt.allocPrint(r.a, "{s} {s} {s}", .{ id, tag, what }), got, want); +} + +/// One window from the shared pre-state: every stage's bits, the fp32 partial, the caches after; then the f64 chain. +fn window(r: *Rig, k: dsa.Kernels, c: dsa.Config, p64: dref.Params, w: *const dsa.DsaW, host: *const dref.Host, pre: *const [names.len][]u8, cap: usize, dir: []const u8, tag: []const u8, e: Obj) !f64 { + const rows: usize = @intCast(ref.int(e, "rows")); + const pos: usize = @intCast(ref.int(e, "pos")); + var sc = try dsa.Scratch.init(r.gpu.d, w, @intCast(ref.int(e, "buf_rows")), e.get("prefill").?.bool, cap); + defer sc.deinit(); + try sc.buf.fill8(0xff, r.s.handle); // torch.empty's scratch: a read of an unwritten word shows + var caches: [names.len]u64 = undefined; + for (pre, 0..) |bytes, i| caches[i] = try r.dev(bytes); + var pos_b: [4]u8 = undefined; + std.mem.writeInt(i32, &pos_b, @intCast(pos), .little); + const st: dsa.DsaState = .{ .lc = caches[0], .scales = caches[1], .ik = caches[2], .ig = caches[3], .pk = caches[4], .pos_dev = try r.dev(&pos_b), .cap = cap }; + const normed = try file(r, dir, ref.str(e, "normed")); + const part = try r.zeros(rows * dsa.hidden * 4); + try dsa.dsaBlock(k, &r.s, c, w, st, &sc, try r.dev(normed), part, rows, pos, false); + // every stage in the block's order: the first that differs names the kernel + const stages = [_]struct { []const u8, u64 }{ + .{ "dqa", sc.dqa }, .{ "dkv", sc.dkv }, .{ "qr", sc.qr }, .{ "lat", sc.lat }, + .{ "q", sc.q }, .{ "ikr", sc.ikr }, .{ "igr", sc.igr }, .{ "qa", sc.qa }, + .{ "qi", sc.qi }, .{ "tokens", sc.tokens }, .{ "counts", sc.counts }, .{ "ol", sc.ol }, + .{ "vn", sc.vn }, .{ "part", part }, + }; + var got: [stages.len]?[]const u8 = undefined; + for (stages, 0..) |stg, i| { + got[i] = null; + const name = (e.get(stg[0]) orelse continue).string; + const want = try file(r, dir, name); + const g = try back(r, stg[1], want.len); + try same(r, tag, stg[0], g, want); + got[i] = g; + } + // the caches: the window's rows as Python wrote them, every other row as the pre-state left it + var post: [names.len][]const u8 = undefined; + for (names, 0..) |n, i| { + const co = e.get(n).?.object; + const lo: usize = @intCast(ref.int(co, "lo")); + const hi: usize = @intCast(ref.int(co, "hi")); + const want = try r.a.dupe(u8, pre[i]); + const win = try file(r, dir, ref.str(co, "file")); + try check.expect(win.len == (hi - lo) * row_bytes[i], "{s} {s} {s}: window bytes", .{ id, tag, n }); + @memcpy(want[lo * row_bytes[i] ..][0..win.len], win); + post[i] = try back(r, caches[i], want.len); + try same(r, tag, n, post[i], want); + } + const np = dsa.poolBucket(pos, rows, cap / 4); + const last = (rows - 1) / dsa.select_rows * dsa.select_rows; + const g: dref.Got = .{ + .normed = normed, + .dqa = got[0].?, + .dkv = got[1].?, + .qr = got[2].?, + .lat = got[3].?, + .q = got[4].?, + .ikr = got[5].?, + .igr = got[6].?, + .qa = got[7].?, + .qi = got[8], + .tokens = got[9], + .counts = got[10], + .ol = got[11].?, + .vn = got[12].?, + .part = got[13].?, + .scores = if (got[8] != null) try back(r, sc.scores, (rows - last) * np * 4) else null, + .np = np, + .kc = post[0], + .scales = post[1], + .ik = post[2], + .ig = post[3], + .pk = post[4], + }; + return dref.chain(r.a, host, p64, g, tag, rows, pos); +} diff --git a/zig/tests/cuda/glm_dsa_ref.zig b/zig/tests/cuda/glm_dsa_ref.zig new file mode 100644 index 0000000000..4c5ecb4c45 --- /dev/null +++ b/zig/tests/cuda/glm_dsa_ref.zig @@ -0,0 +1,293 @@ +//! glm.dsa.block's float64 part: every kernel's f64 reference on the block's own stage inputs, chained in order. +const std = @import("std"); +const ref = @import("glm_kda_ref.zig"); +const sref = @import("glm_sparse_ref.zig"); +const latref = @import("glm_latent_ref.zig"); +const lref = @import("glm_layer_ref.zig"); + +const Obj = std.json.ObjectMap; +const id = "glm.dsa.block"; +const H = 32; // the rank's heads +const QK = 256; +const LW = 512; +const IKR = 160; +const ID = 128; + +/// Rank 0's DSA tensors from the fixture (manifest "dsa"."weights"), read once. +pub const Host = struct { + q_a: lref.Lin, + kv_a: lref.Lin, + q_b: lref.Lin, + o: lref.Lin, + kw: lref.Mat, + qb: lref.Mat, + q_norm: []const u8, + kv_norm: []const u8, + wk: []const u8, + wv: []const u8, + gate: []const u8, + ln_w: []const u8, + ln_b: []const u8, + ape: []const u8, + + pub fn read(a: std.mem.Allocator, io: std.Io, dir: []const u8, o: Obj) !Host { + const R = struct { + a: std.mem.Allocator, + io: std.Io, + dir: []const u8, + o: Obj, + fn f(s: @This(), key: []const u8) ![]const u8 { + return ref.readFile(s.a, s.io, s.dir, ref.str(s.o, key)); + } + fn lin(s: @This(), key: []const u8) !lref.Lin { + const l = s.o.get(key).?.object; + if (!std.mem.eql(u8, ref.str(l, "codebook"), "mul1")) return error.GlmDsaCodebook; + const rd = ref.readFile; + return .{ .words = try rd(s.a, s.io, s.dir, ref.str(l, "words")), .suh = try rd(s.a, s.io, s.dir, ref.str(l, "suh")), .svh = try rd(s.a, s.io, s.dir, ref.str(l, "svh")), .k = @intCast(ref.int(l, "k")), .n = @intCast(ref.int(l, "n")), .halves = @intCast(ref.int(l, "k2")) }; + } + fn mat(s: @This(), key: []const u8) !lref.Mat { + const m = s.o.get(key).?.object; + return .{ .w = try ref.readFile(s.a, s.io, s.dir, ref.str(m, "w")), .n = @intCast(ref.int(m, "n")), .k = @intCast(ref.int(m, "k")) }; + } + }; + const s: R = .{ .a = a, .io = io, .dir = dir, .o = o }; + return .{ + .q_a = try s.lin("q_a"), + .kv_a = try s.lin("kv_a"), + .q_b = try s.lin("q_b"), + .o = try s.lin("o"), + .kw = try s.mat("kw"), + .qb = try s.mat("qb"), + .q_norm = try s.f("q_norm"), + .kv_norm = try s.f("kv_norm"), + .wk = try s.f("wk"), + .wv = try s.f("wv"), + .gate = try s.f("gate"), + .ln_w = try s.f("ln_w"), + .ln_b = try s.f("ln_b"), + .ape = try s.f("ape"), + }; + } +}; + +pub const Params = struct { eps: f64, index_eps: f64, scale: f64, index_scale: f64, wscale: f64 }; + +/// One window's stage bytes (already bit-equal to Python's) and the caches after it; scores: the last select block. +pub const Got = struct { + normed: []const u8, + dqa: []const u8, + dkv: []const u8, + qr: []const u8, + lat: []const u8, + q: []const u8, + ikr: []const u8, + igr: []const u8, + qa: []const u8, + ol: []const u8, + vn: []const u8, + part: []const u8, + qi: ?[]const u8, + tokens: ?[]const u8, + counts: ?[]const u8, + scores: ?[]const u8, + np: usize, + kc: []const u8, + scales: []const u8, + ik: []const u8, + ig: []const u8, + pk: []const u8, +}; + +fn u16s(b: []const u8) []const u16 { + return @alignCast(std.mem.bytesAsSlice(u16, b)); +} +fn bfAt(b: []const u8, i: usize) f64 { + return ref.bf(std.mem.readInt(u16, b[2 * i ..][0..2], .little)); +} +fn f32At(b: []const u8, i: usize) f64 { + return @as(f32, @bitCast(std.mem.readInt(u32, b[4 * i ..][0..4], .little))); +} +fn i32At(b: []const u8, i: usize) i32 { + return std.mem.readInt(i32, b[4 * i ..][0..4], .little); +} + +/// The worst |got - want| / tol of one stage, NaN counting as infinite. +const Bound = struct { + worst: f64 = 0, + at: usize = 0, + + fn add(b: *Bound, got: f64, want: f64, tol: f64, i: usize) void { + const q = if (got == want) 0 else @abs(got - want) / tol; + if (!(q <= b.worst)) { + b.worst = if (std.math.isNan(q)) std.math.inf(f64) else q; + b.at = i; + } + } + + fn rel(b: *Bound, got: f64, want: f64, scale: f64, floor: f64, i: usize) void { + b.add(got, want, @abs(want) * scale + floor, i); + } +}; + +/// The stages' verdicts: the first stage past its tolerance fails the window. +const Report = struct { + tag: []const u8, + worst: f64 = 0, + + fn stage(rp: *Report, name: []const u8, b: Bound) !void { + if (b.worst > 1.0) { + std.debug.print("FAIL {s} {s} {s} f64: worst {d:.4} of tolerance at element {d}\n", .{ id, rp.tag, name, b.worst, b.at }); + return error.TestFailed; + } + rp.worst = @max(rp.worst, b.worst); + } +}; + +/// An EXL3 linear on sample rows (glm_layer_ref's reference): |y|/128 (bf16 out) or |y|/1024 (fp32) + s1/2048. +fn exl3Stage(rp: *Report, a: std.mem.Allocator, name: []const u8, l: lref.Lin, x: []const u8, rows: []const usize, got: []const u8, f32_out: bool) !void { + const xs = try a.alloc(f64, rows.len * l.k); + for (rows, 0..) |row, i| for (0..l.k) |j| { + xs[i * l.k + j] = bfAt(x, row * l.k + j); + }; + const res = try lref.exl3Ref(a, l, xs, rows.len); + var b = Bound{}; + for (rows, 0..) |row, i| for (0..l.n) |col| { + const want = res.want[i * l.n + col]; + const g = if (f32_out) f32At(got, row * l.n + col) else bfAt(got, row * l.n + col); + b.add(g, want, @abs(want) / @as(f64, if (f32_out) 1024 else 128) + res.s1[i * l.n + col] / 2048 + 1e-30, row * l.n + col); + }; + try rp.stage(name, b); +} + +/// A plain bf16 matmul on sample rows: x [rows, xk] bf16 (row stride xk), w [n, k]; bf16 or fp32 out [rows, n]. +fn b16Stage(rp: *Report, name: []const u8, m: lref.Mat, x: []const u8, rows: []const usize, got: []const u8, f32_out: bool) !void { + var b = Bound{}; + for (rows) |row| for (0..m.n) |col| { + var acc: f64 = 0; + var s: f64 = 0; + for (0..m.k) |j| { + const t = bfAt(x, row * m.k + j) * bfAt(m.w, col * m.k + j); + acc += t; + s += @abs(t); + } + const g = if (f32_out) f32At(got, row * m.n + col) else bfAt(got, row * m.n + col); + b.add(g, acc, @abs(acc) / @as(f64, if (f32_out) 1024 else 128) + s / 2048 + 1e-30, row * m.n + col); + }; + try rp.stage(name, b); +} + +/// glue.rmsnorm in f64: w x / sqrt(mean(x^2) + eps); two bf16 roundings: |y|/64. +fn rmsStage(rp: *Report, name: []const u8, x: []const u8, w: []const u8, got: []const u8, d: usize, rows: usize, eps: f64) !void { + var b = Bound{}; + for (0..rows) |row| { + var ss: f64 = 0; + for (0..d) |i| ss += bfAt(x, row * d + i) * bfAt(x, row * d + i); + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(d)) + eps); + for (0..d) |i| b.rel(bfAt(got, row * d + i), bfAt(w, i) * bfAt(x, row * d + i) * rinv, 1.0 / 64.0, 1e-30, row * d + i); + } + try rp.stage(name, b); +} + +/// The chained float64 part of one window; returns its worst stage as a fraction of tolerance. +pub fn chain(a: std.mem.Allocator, h: *const Host, p: Params, g: Got, tag: []const u8, rows: usize, pos: usize) !f64 { + var rp: Report = .{ .tag = tag }; + var sbuf: [8]usize = undefined; + const sp = lref.samples(rows); + @memcpy(sbuf[0..sp.len], sp); + const smp = sbuf[0..sp.len]; + // dsa_front + try exl3Stage(&rp, a, "q_a", h.q_a, g.normed, smp, g.dqa, false); + try exl3Stage(&rp, a, "kv_a", h.kv_a, g.normed, smp, g.dkv, false); + try rmsStage(&rp, "q_norm", g.dqa, h.q_norm, g.qr, h.q_a.n, rows, p.eps); + try rmsStage(&rp, "kv_norm", g.dkv, h.kv_norm, g.lat, h.kv_a.n, rows, p.eps); + try exl3Stage(&rp, a, "q_b", h.q_b, g.qr, smp, g.q, false); + // the latent write: scale amax / 448, codes ~ lat / scale (e4m3's 3 bits) + var bw = Bound{}; + var bs = Bound{}; + for (0..rows) |r| { + var amax: f64 = 0; + for (0..LW) |i| amax = @max(amax, @abs(bfAt(g.lat, r * LW + i))); + const s = if (amax > 0) amax / 448.0 else 1.0; + bs.rel(f32At(g.scales, pos + r), s, 1e-6, 1e-33, pos + r); + for (0..LW) |i| bw.rel(sref.fp8e4m3(g.kc[(pos + r) * LW + i]), bfAt(g.lat, r * LW + i) / s, 1.0 / 8.0, 0.01, r * LW + i); + } + try rp.stage("latent scales", bs); + try rp.stage("latent codes", bw); + // the indexer: keys | head weights, the gates (fp32), LayerNorm'd keys, gates as bf16, completed pools + try b16Stage(&rp, "index kw", h.kw, g.normed, smp, g.ikr, false); + try b16Stage(&rp, "index gate", .{ .w = h.gate, .n = ID, .k = h.kw.k }, g.normed, smp, g.igr, true); + const ik64 = try sref.indexRef(a, u16s(g.ikr), IKR, u16s(h.ln_w), u16s(h.ln_b), p.index_eps, rows); + var bi = Bound{}; + var bg = Bound{}; + for (0..rows * ID) |i| { + bi.rel(bfAt(g.ik, pos * ID + i), ik64[i], 1.0 / 128.0, 2e-3, i); + bg.rel(bfAt(g.ig, pos * ID + i), f32At(g.igr, i), 1.0 / 128.0, 1e-30, i); + } + try rp.stage("index keys", bi); + try rp.stage("index gates", bg); + const p0 = pos / 4; + const p1 = (pos + rows) / 4; + const pk64 = try sref.poolRef(a, u16s(g.ig), u16s(h.ape), u16s(g.ik), p0, p1); + var bp = Bound{}; + for (0..(p1 - p0) * ID) |i| bp.rel(bfAt(g.pk, p0 * ID + i), pk64[i], 1.0 / 64.0, 2e-3, i); + try rp.stage("pool keys", bp); + try attention(&rp, a, h, p, g, smp, pos); + // expand_v and o_proj's fp32 partial + var be = Bound{}; + for (smp) |r| { + const want = try latref.expandRef(a, u16s(g.ol)[r * H * LW ..][0 .. H * LW], u16s(h.wv), 1, H, LW, QK); + for (want, 0..) |w, i| be.rel(bfAt(g.vn, r * H * QK + i), w, 1.0 / 128.0, 2e-3, r * H * QK + i); + } + try rp.stage("expand", be); + try exl3Stage(&rp, a, "o_proj", h.o, g.vn, smp, g.part, true); + return rp.worst; +} + +/// absorb_q, then each sample row's attention: dense through its position, or sparse over its selected tokens. +fn attention(rp: *Report, a: std.mem.Allocator, h: *const Host, p: Params, g: Got, smp: []const usize, pos: usize) !void { + var ba = Bound{}; + for (smp) |r| { + const want = try latref.absorbRef(a, u16s(g.q)[r * H * QK ..][0 .. H * QK], u16s(h.wk), 1, H, QK, LW); + for (want, 0..) |w, i| ba.rel(bfAt(g.qa, r * H * LW + i), w, 1.0 / 128.0, 2e-3, r * H * LW + i); + } + try rp.stage("absorb", ba); + const rows = g.normed.len / (2 * 4096); + if (g.qi) |qi| try b16Stage(rp, "index qb", h.qb, g.qr, smp, qi, false); + if (g.scores) |sc| { // the selection: f64 pool scores, then the order-key sort of the kernel's own fp32 scores + const n = sc.len / (4 * g.np); + const a0 = rows - n; + var bsc = Bound{}; + for (smp) |r| { + if (r < a0) continue; + const want = try sref.scoresRef(a, u16s(g.qi.?)[r * 4096 ..], 4096, u16s(g.ikr)[r * IKR + ID ..], IKR, u16s(g.pk), 1, g.np, p.index_scale, p.wscale, @intCast(pos + r)); + for (want, 0..) |w, i| { + const got = f32At(sc, (r - a0) * g.np + i); + if (std.math.isInf(w)) bsc.add(got, w, 1, i) else bsc.rel(got, w, 1.0 / 64.0, 2e-3, (r - a0) * g.np + i); + } + } + try rp.stage("pool scores", bsc); + const tok = try a.alloc(i32, n * 2051); + const cnt = try a.alloc(i32, n); + try sref.selectRef(a, @alignCast(std.mem.bytesAsSlice(f32, sc)), g.np, 512, @intCast(pos + a0), n, tok, cnt); + for (tok, 0..) |t, i| if (t != i32At(g.tokens.?, a0 * 2051 + i)) { + std.debug.print("FAIL {s} {s} tokens f64: token {d} is {d}, the order-key sort's {d}\n", .{ id, rp.tag, i, i32At(g.tokens.?, a0 * 2051 + i), t }); + return error.TestFailed; + }; + for (cnt, 0..) |c, i| if (c != i32At(g.counts.?, a0 + i)) return error.TestFailed; + } + var bd = Bound{}; + var bsp = Bound{}; + const scl: []const f32 = @alignCast(std.mem.bytesAsSlice(f32, g.scales)); + for (smp) |r| { + const qa = u16s(g.qa)[r * H * LW ..][0 .. H * LW]; + const sparse_row = g.counts != null and i32At(g.counts.?, r) > 0; + const want = if (sparse_row) + try sref.sparseRef(a, qa, g.kc, scl, @alignCast(std.mem.bytesAsSlice(i32, g.tokens.?[r * 2051 * 4 ..][0 .. 2051 * 4])), &[_]i32{i32At(g.counts.?, r)}, H, LW, p.scale) + else + try latref.attnRef(a, qa, g.kc, scl, @intCast(pos + r), 1, H, LW, p.scale, true, 1); + const b = if (sparse_row) &bsp else &bd; + for (want, 0..) |w, i| b.rel(bfAt(g.ol, r * H * LW + i), w, 1.0 / 64.0, 2e-3, r * H * LW + i); + } + try rp.stage("dense attention", bd); + try rp.stage("sparse attention", bsp); +} diff --git a/zig/tests/cuda/glm_gaps_tests.zig b/zig/tests/cuda/glm_gaps_tests.zig new file mode 100644 index 0000000000..2c2e6df9b7 --- /dev/null +++ b/zig/tests/cuda/glm_gaps_tests.zig @@ -0,0 +1,194 @@ +//! The glm_glue2 kernels and MoE per-slot epilogue against the Python engine's bytes (manifest "gaps"). +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const glue = cuda.glm_glue; +const glue2 = cuda.glm_glue2; +const exl3e = cuda.exl3_experts; + +const Rig = quant_ref_tests.RigT; +const Tol = @import("glm_glue_tests.zig").Tol; +const bf = quant_ref_tests.helpers.bfOf; + +/// The manifest's "gaps" object, or null (the caller skips) when the fixture dir has none. +pub fn section(r: *Rig, dir: []const u8) !?std.json.ObjectMap { + if (std.mem.eql(u8, dir, "none")) return null; + var man = try ref.Manifest.open(r.a, r.gpu.io, dir); + return if (ref.Manifest.root(&man).get("gaps")) |g| g.object else null; +} + +fn up(r: *Rig, dir: []const u8, name: []const u8) !u64 { + return r.dev(try ref.readFile(r.a, r.gpu.io, dir, name)); +} + +fn f32s(bytes: []const u8) []const f32 { + return ref.f32s(bytes); +} + +/// bf16 bytes as f64 values (exact). +fn bf4(a: std.mem.Allocator, bytes: []const u8) ![]f64 { + const u = std.mem.bytesAsSlice(u16, bytes); + const out = try a.alloc(f64, u.len); + for (u, 0..) |v, i| out[i] = bf(v); + return out; +} + +// -- glm.glue.rmsnorm, the 1536 and 512 wide DSA layernorms ---------------------------------------- +pub fn rmsnorm(r: *Rig, dir: []const u8) !void { + const gp = (try section(r, dir)) orelse return; + const o = gp.get("rms").?.object; + var k = try glue2.Glue2.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const eps: f32 = @floatCast(ref.float(o, "eps")); + const dp = try up(r, dir, ref.str(o, "dp")); + for ([_]struct { d: usize, xoff: usize, nw: []const u8, out: []const u8, sums: []const u8 }{ + .{ .d = 1536, .xoff = 0, .nw = ref.str(o, "nw1"), .out = ref.str(o, "out1"), .sums = ref.str(o, "sums1") }, + .{ .d = 512, .xoff = 1536, .nw = ref.str(o, "nw2"), .out = ref.str(o, "out2"), .sums = ref.str(o, "sums2") }, + }) |case| { + const w = try up(r, dir, case.nw); + const want = try ref.readFile(r.a, r.gpu.io, dir, case.out); + const want_sums = try ref.readFile(r.a, r.gpu.io, dir, case.sums); + const out = try r.zeros(want.len); + const sums = try r.zeros(want_sums.len); + try k.rmsnorm2(r.s, dp + case.xoff * 2, 2048, w, out, @intCast(case.d), sums, eps, rows, case.d); + try check.sameBytes("glm.glue.rmsnorm", try r.back(out, want.len), want); + try check.sameBytes("glm.glue.rmsnorm sums", try r.back(sums, want_sums.len), want_sums); + // f64: x starts at element case.xoff of dp's 2048-wide rows; w is the case's own. + const dpv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "dp"))); + const wv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, case.nw)); + const gotv = try bf4(r.a, try r.back(out, want.len)); + const off: usize = if (case.d == 1536) 0 else 1536; // 512's x starts at element 1536 of the 2048-wide row + var t = Tol{}; + for (0..rows) |i| { + var ss: f64 = 0; + for (0..case.d) |j| ss += dpv[i * 2048 + off + j] * dpv[i * 2048 + off + j]; + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(case.d)) + eps); + for (0..case.d) |j| t.f64v(gotv[i * case.d + j], wv[j] * dpv[i * 2048 + off + j] * rinv, 1.0 / 64.0, 3e-3); + } + try t.done("glm.glue.rmsnorm gaps"); + check.pass("glm.glue.rmsnorm", .{}); + } +} + +// -- glm.glue.combine, fp32 y ----------------------------------------------------------------------- +pub fn combine(r: *Rig, dir: []const u8) !void { + const gp = (try section(r, dir)) orelse return; + const o = gp.get("combine").?.object; + var k = try glue2.Glue2.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const y = try up(r, dir, ref.str(o, "y")); + const wts = try up(r, dir, ref.str(o, "wts")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "out")); + const out = try r.zeros(want.len); + try k.combineF32(r.s, y, wts, out, rows, d, 9); + try check.sameBytes("glm.glue.combine", try r.back(out, want.len), want); + const yv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "y"))); + const wv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "wts"))); + const gotv = f32s(try r.back(out, want.len)); + var t = Tol{}; + for (0..rows) |i| for (0..d) |j| { + var acc: f64 = 0; + for (0..9) |kk| acc += wv[i * 9 + kk] * yv[(i * 9 + kk) * d + j]; + t.f32v(gotv, i * d + j, acc, 1.0 / 1024.0, 1e-5); + }; + try t.done("glm.glue.combine gaps"); + check.pass("glm.glue.combine", .{}); +} + +// -- glm.glue.combine, the real MoE layer: per-slot down epilogue, then the fp32 combine +pub fn combineMoe(r: *Rig, dir: []const u8) !void { + const gp = (try section(r, dir)) orelse return; + const o = gp.get("moe").?.object; + const count: usize = @intCast(ref.int(o, "experts")); + const dims: usize = @intCast(ref.int(o, "dims")); + const width: usize = @intCast(ref.int(o, "width")); + const slots: usize = @intCast(ref.int(o, "slots")); + var exp = try exl3e.Experts.load(r.gpu.d); + defer exp.unload(); + var gate_p = try r.a.alloc(u64, count); + var up_p = try r.a.alloc(u64, count); + var down_p = try r.a.alloc(u64, count); + var gk2 = try r.a.alloc(i32, count); + var uk2 = try r.a.alloc(i32, count); + var dk2 = try r.a.alloc(i32, count); + for (o.get("gate_k2").?.array.items, 0..) |v, i| gk2[i] = @intCast(v.integer); + for (o.get("up_k2").?.array.items, 0..) |v, i| uk2[i] = @intCast(v.integer); + for (o.get("down_k2").?.array.items, 0..) |v, i| dk2[i] = @intCast(v.integer); + for (0..count) |i| { + gate_p[i] = try up(r, dir, o.get("gate").?.array.items[i].string); + up_p[i] = try up(r, dir, o.get("up").?.array.items[i].string); + down_p[i] = try up(r, dir, o.get("down").?.array.items[i].string); + } + const layer: exl3e.Layer = .{ + .gate = try r.dev(std.mem.sliceAsBytes(gate_p)), + .up = try r.dev(std.mem.sliceAsBytes(up_p)), + .down = try r.dev(std.mem.sliceAsBytes(down_p)), + .gate_k2 = try r.dev(std.mem.sliceAsBytes(gk2)), + .up_k2 = try r.dev(std.mem.sliceAsBytes(uk2)), + .down_k2 = try r.dev(std.mem.sliceAsBytes(dk2)), + .suh_g = try up(r, dir, ref.str(o, "suh_g")), + .suh_u = try up(r, dir, ref.str(o, "suh_u")), + .svh_g = try up(r, dir, ref.str(o, "svh_g")), + .svh_u = try up(r, dir, ref.str(o, "svh_u")), + .suh_d = try up(r, dir, ref.str(o, "suh_d")), + .svh_d = try up(r, dir, ref.str(o, "svh_d")), + .experts = count, + .dims = dims, + .width = width, + .codebook = .mul1, + .limit = @floatCast(ref.float(o, "limit")), + .act_mode = @intCast(ref.int(o, "act_mode")), + }; + var k = try glue2.Glue2.load(r.gpu.d); + defer k.unload(); + const xmap = o.get("x").?.object; + const pickmap = o.get("pick").?.object; + const wtsmap = o.get("wts").?.object; + const ymap = o.get("y").?.object; + const symap = o.get("sy").?.object; + const partmap = o.get("part").?.object; + const key = struct { + fn f(a: std.mem.Allocator, m: i64) ![]const u8 { + return std.fmt.allocPrint(a, "{d}", .{m}); + } + }.f; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const kk = try key(r.a, rv.integer); + const dx = try up(r, dir, xmap.get(kk).?.string); + const dpick = try up(r, dir, pickmap.get(kk).?.string); + const dwts = try up(r, dir, wtsmap.get(kk).?.string); + // ey: 0xff everywhere except the shared expert's slot, which starts at the fixture's sy + const eyb = try r.a.alloc(u8, m * slots * dims * 4); + @memset(eyb, 0xff); + const sy = try ref.readFile(r.a, r.gpu.io, dir, symap.get(kk).?.string); + for (0..m) |row| @memcpy(eyb[(row * slots + slots - 1) * dims * 4 ..][0 .. dims * 4], sy[row * dims * 4 ..][0 .. dims * 4]); + const dey = try r.dev(eyb); + const sb = try r.zeros(try exl3e.scratchBytes(layer, m, slots)); + const sc = try exl3e.scratch(sb, layer, m, slots); + try exp.runPerSlot(r.s, layer, dx, dims, .bf16, dpick, m, dey, sc); + const want_y = try ref.readFile(r.a, r.gpu.io, dir, ymap.get(kk).?.string); + try check.sameBytes("glm.glue.combine moe ey", try r.back(dey, want_y.len), want_y); + const want = try ref.readFile(r.a, r.gpu.io, dir, partmap.get(kk).?.string); + const out = try r.zeros(want.len); + try k.combineF32(r.s, dey, dwts, out, m, dims, slots); + try check.sameBytes("glm.glue.combine moe", try r.back(out, want.len), want); + // f64 over the fixture's y (fp32, exact): the same ordered sum in float64 + const yv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ymap.get(kk).?.string)); + const wv = f32s(try ref.readFile(r.a, r.gpu.io, dir, wtsmap.get(kk).?.string)); + const gotv = f32s(try r.back(out, want.len)); + var t = Tol{}; + for (0..m) |i| for (0..dims) |j| { + var acc: f64 = 0; + for (0..slots) |q| acc += wv[i * slots + q] * yv[(i * slots + q) * dims + j]; + t.f32v(gotv, i * dims + j, acc, 1.0 / 1024.0, 1e-5); + }; + try t.done("glm.glue.combine moe f64"); + check.pass("glm.glue.combine", .{}); + } +} diff --git a/zig/tests/cuda/glm_glue_tests.zig b/zig/tests/cuda/glm_glue_tests.zig new file mode 100644 index 0000000000..392ae0f634 --- /dev/null +++ b/zig/tests/cuda/glm_glue_tests.zig @@ -0,0 +1,539 @@ +//! GLM glue kernels: each against a float64 reference and the Python engine's fixture bytes (glue.py, qmm.py). +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const gaps = @import("glm_gaps_tests.zig"); +const glue = cuda.glm_glue; + +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const Failed = check.Failed; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + embedCheck(r, state, dir) catch |e| state.fail(e, "glm.glue.embed"); + rmsnorm(r, state, dir) catch |e| state.fail(e, "glm.glue.rmsnorm"); + try rmsnormGaps(r, state, dir); + hcPre(r, state, dir) catch |e| state.fail(e, "glm.glue.hc_pre"); + hcPost(r, state, dir) catch |e| state.fail(e, "glm.glue.hc_post"); + residual(r, state, dir) catch |e| state.fail(e, "glm.glue.residual"); + swiglu(r, state, dir) catch |e| state.fail(e, "glm.glue.swiglu"); + router(r, state, dir) catch |e| state.fail(e, "glm.glue.router"); + select(r, state, dir) catch |e| state.fail(e, "glm.glue.select"); + combine(r, state, dir) catch |e| state.fail(e, "glm.glue.combine"); + try combineGaps(r, state, dir); + bmm(r, state, dir) catch |e| state.fail(e, "glm.bmm"); + greedy(r, state, dir) catch |e| state.fail(e, "glm.sample.greedy"); +} + +fn skip(id: []const u8) void { + std.debug.print("SKIP {s}\n", .{id}); +} + +/// The fixture family's manifest object, or null (SKIP printed) when there is no fixture dir or family. +fn section(r: *Rig, dir: []const u8, id: []const u8) !?std.json.ObjectMap { + if (std.mem.eql(u8, dir, "none")) { + skip(id); + return null; + } + var man = try ref.Manifest.open(r.a, r.gpu.io, dir); + const g = ref.Manifest.root(&man).get("glue") orelse { + skip(id); + return null; + }; + return g.object; +} + +fn up(r: *Rig, dir: []const u8, name: []const u8) !u64 { + return r.dev(try ref.readFile(r.a, r.gpu.io, dir, name)); +} + +fn f32s(bytes: []const u8) []const f32 { + return ref.f32s(bytes); +} + +/// bf16 bytes as f64 values (exact). +fn bf4(a: std.mem.Allocator, bytes: []const u8) ![]f64 { + const u = std.mem.bytesAsSlice(u16, bytes); + const out = try a.alloc(f64, u.len); + for (u, 0..) |v, i| out[i] = ref.bf(v); + return out; +} + +fn sigmoid(x: f64) f64 { + return 1.0 / (1.0 + @exp(-x)); +} + +pub const Tol = struct { + worst: f64 = 0, + + pub fn f64v(t: *Tol, got: f64, want: f64, scale: f64, floor: f64) void { + const q = @abs(got - want) / (@abs(want) * scale + floor); + if (q > t.worst) t.worst = q; + } + + pub fn f32v(t: *Tol, arr: []const f32, i: usize, want: f64, scale: f64, floor: f64) void { + t.f64v(arr[i], want, scale, floor); + } + + pub fn done(t: *Tol, name: []const u8) Failed!void { + if (t.worst > 1.0) { + std.debug.print("FAIL {s}: f64 worst {d:.4} of tolerance\n", .{ name, t.worst }); + return error.TestFailed; + } + std.debug.print("RESULT {s}: f64 worst {d:.4} of tolerance\n", .{ name, t.worst }); + } +}; + +// -- glm.glue.embed ----------------------------------------------------------------------------- +fn embedCheck(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.embed")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const ids = try up(r, dir, ref.str(o, "ids")); + const tab = try up(r, dir, ref.str(o, "e_tab")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "e_out")); + const out = try r.zeros(want.len); + try k.embedB16(r.s, ids, tab, out, rows, d, 4); + try check.sameBytes("glm.glue.embed", try r.back(out, want.len), want); + // f64: the rows are exact copies of the table's bf16 values. + const idv = std.mem.bytesAsSlice(u32, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "ids"))); + const tabv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "e_tab"))); + const gotv = try bf4(r.a, try r.back(out, want.len)); + var t = Tol{}; + for (0..rows) |i| for (0..4) |c| for (0..d) |j| { + t.f64v(gotv[(i * 4 + c) * d + j], tabv[@as(usize, @intCast(idv[i])) * d + j], 0.0, 1e-6); + }; + try t.done("glm.glue.embed"); + check.pass("glm.glue.embed", .{}); +} + +// -- glm.glue.rmsnorm --------------------------------------------------------------------------- +fn rmsnorm(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.rmsnorm")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const eps: f32 = @floatCast(ref.float(o, "eps")); + const x = try up(r, dir, ref.str(o, "r_x")); + const w = try up(r, dir, ref.str(o, "nw")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "r_out")); + const want_sums = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "r_sums")); + const out = try r.zeros(want.len); + const sums = try r.zeros(want_sums.len); + try k.rmsnorm(r.s, x, @intCast(d), w, out, @intCast(d), sums, eps, rows, d); + try check.sameBytes("glm.glue.rmsnorm", try r.back(out, want.len), want); + try check.sameBytes("glm.glue.rmsnorm sums", try r.back(sums, want_sums.len), want_sums); + const xv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "r_x"))); + const wv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "nw"))); + const gotv = try bf4(r.a, try r.back(out, want.len)); + var t = Tol{}; + for (0..rows) |i| { + var ss: f64 = 0; + for (0..d) |j| ss += xv[i * d + j] * xv[i * d + j]; + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(d)) + eps); + for (0..d) |j| t.f64v(gotv[i * d + j], wv[j] * xv[i * d + j] * rinv, 1.0 / 64.0, 3e-3); + } + try t.done("glm.glue.rmsnorm"); + check.pass("glm.glue.rmsnorm", .{}); +} + +// -- glm.glue.hc_pre ---------------------------------------------------------------------------- +fn hcPre(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.hc_pre")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const eps: f32 = @floatCast(ref.float(o, "eps")); + const hc_eps: f32 = @floatCast(ref.float(o, "hc_eps")); + const iters: u32 = @intCast(ref.int(o, "iters")); + const x = try up(r, dir, ref.str(o, "h_x")); + const fn_w = try up(r, dir, ref.str(o, "fn")); + const base = try up(r, dir, ref.str(o, "base")); + const scale = try up(r, dir, ref.str(o, "hc_scale")); + const nw = try up(r, dir, ref.str(o, "nw")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "h_out")); + const want_xs = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "h_xs")); + const want_post = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "h_post")); + const want_comb = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "h_comb")); + const part = try r.zeros(rows * 16 * 32 * 4); + const out = try r.zeros(want.len); + const xs = try r.zeros(want_xs.len); + const post = try r.zeros(want_post.len); + const comb = try r.zeros(want_comb.len); + try k.hcPartial(r.s, x, fn_w, part, rows); + try k.hcFinish(r.s, x, part, base, scale, nw, out, xs, post, comb, eps, hc_eps, iters, rows, d); + try check.sameBytes("glm.glue.hc_pre", try r.back(out, want.len), want); + try check.sameBytes("glm.glue.hc_pre xs", try r.back(xs, want_xs.len), want_xs); + try check.sameBytes("glm.glue.hc_pre post", try r.back(post, want_post.len), want_post); + try check.sameBytes("glm.glue.hc_pre part", try r.back(part, rows * 16 * 32 * 4), try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "h_part"))); + try check.sameBytes("glm.glue.hc_pre comb", try r.back(comb, want_comb.len), want_comb); + // f64: the dots, gates, sinkhorn and collapsed row in double precision. + const xv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "h_x"))); + const fnv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "fn"))); + const basev = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "base"))); + const scv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "hc_scale"))); + const nwv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "nw"))); + const wide = 4 * d; + const gotv = try bf4(r.a, try r.back(out, want.len)); + const post_got = f32s(try r.back(post, want_post.len)); + const comb_got = f32s(try r.back(comb, want_comb.len)); + var t = Tol{}; + for (0..rows) |i| { + var mix: [24]f64 = @splat(0); + var ss: f64 = 0; + for (0..wide) |j| { + const xj: f64 = xv[i * wide + j]; + ss += xj * xj; + for (0..24) |m| mix[m] += fnv[m * wide + j] * xj; + } + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(wide)) + eps); + var p: [4]f64 = undefined; + for (0..4) |sv| { + p[sv] = sigmoid(mix[sv] * rinv * scv[0] + basev[sv]) + hc_eps; + t.f32v(post_got, i * 4 + sv, 2.0 * sigmoid(mix[4 + sv] * rinv * scv[1] + basev[4 + sv]), 1.0 / 1024.0, 1e-6); + } + var ce: [4][4]f64 = undefined; + for (0..4) |a| { + var row: [4]f64 = undefined; + var cmax: f64 = -std.math.inf(f64); + for (0..4) |b| { + row[b] = mix[8 + a * 4 + b] * rinv * scv[2] + basev[8 + a * 4 + b]; + cmax = @max(cmax, row[b]); + } + for (0..4) |b| ce[a][b] = @exp(row[b] - cmax); + } + for (0..4) |a| { + var rs: f64 = 0; + for (0..4) |b| rs += ce[a][b]; + for (0..4) |b| ce[a][b] = ce[a][b] / rs + hc_eps; + } + for (0..4) |b| { + var cs: f64 = 0; + for (0..4) |a| cs += ce[a][b]; + for (0..4) |a| ce[a][b] = ce[a][b] / (cs + hc_eps); + } + for (1..iters) |_| { + for (0..4) |a| { + var rs: f64 = 0; + for (0..4) |b| rs += ce[a][b]; + for (0..4) |b| ce[a][b] = ce[a][b] / (rs + hc_eps); + } + for (0..4) |b| { + var cs: f64 = 0; + for (0..4) |a| cs += ce[a][b]; + for (0..4) |a| ce[a][b] = ce[a][b] / (cs + hc_eps); + } + } + for (0..4) |a| for (0..4) |b| t.f32v(comb_got, i * 16 + a * 4 + b, ce[a][b], 1.0 / 512.0, 1e-6); + // the collapsed row and its norm + const c = try r.a.alloc(f64, d); + for (0..d) |j| c[j] = p[0] * xv[i * wide + j] + p[1] * xv[i * wide + d + j] + p[2] * xv[i * wide + 2 * d + j] + p[3] * xv[i * wide + 3 * d + j]; + var cs2: f64 = 0; + for (0..d) |j| cs2 += c[j] * c[j]; + const rinv2 = 1.0 / @sqrt(cs2 / @as(f64, @floatFromInt(d)) + eps); + for (0..d) |j| t.f64v(gotv[i * d + j], nwv[j] * c[j] * rinv2, 1.0 / 64.0, 3e-3); + } + try t.done("glm.glue.hc_pre"); + check.pass("glm.glue.hc_pre", .{}); +} + +// -- glm.glue.hc_post --------------------------------------------------------------------------- +fn hcPost(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.hc_post")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const x = try up(r, dir, ref.str(o, "p_x")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "p_out")); + const want_pair = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "p_pair_out")); + const post = try up(r, dir, ref.str(o, "h_post")); + const comb = try up(r, dir, ref.str(o, "h_comb")); + const g_all = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "p_g")); + const xout = try r.zeros(want.len); + try k.hcPost(r.s, x, xout, try r.dev(g_all), post, comb, @intCast(rows * d), 2, rows, d, 4); + try check.sameBytes("glm.glue.hc_post", try r.back(xout, want.len), want); + // the pair form: the two ranks' partials held apart, rank 1 read at its element distance + const row_bytes = g_all.len / 2; + const g0 = try r.dev(g_all[0..row_bytes]); + const g1 = try r.dev(g_all[row_bytes..]); + const pout = try r.zeros(want_pair.len); + try k.hcPost(r.s, x, pout, g0, post, comb, @intCast((g1 - g0) / 4), 2, rows, d, 4); + try check.sameBytes("glm.glue.hc_post pair", try r.back(pout, want_pair.len), want_pair); + // f64: the branch sum, the gate mix + const xv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "p_x"))); + const gv = f32s(g_all); + const postv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "h_post"))); + const combv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "h_comb"))); + const gotv = try bf4(r.a, try r.back(xout, want.len)); + var t = Tol{}; + for (0..rows) |i| for (0..4) |sv| for (0..d) |j| { + const branch: f64 = bf16f64(gv[i * d + j] + gv[(rows + i) * d + j]); + var mixed: f64 = combv[i * 16 + 0 * 4 + sv] * xv[(i * 4 + 0) * d + j]; + mixed += combv[i * 16 + 1 * 4 + sv] * xv[(i * 4 + 1) * d + j]; + mixed += combv[i * 16 + 2 * 4 + sv] * xv[(i * 4 + 2) * d + j]; + mixed += combv[i * 16 + 3 * 4 + sv] * xv[(i * 4 + 3) * d + j]; + t.f64v(gotv[(i * 4 + sv) * d + j], postv[i * 4 + sv] * branch + mixed, 1.0 / 64.0, 3e-3); + }; + try t.done("glm.glue.hc_post"); + check.pass("glm.glue.hc_post", .{}); +} + +// -- glm.glue.residual (residual_add and stream_mean) -------------------------------------------- +fn residual(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.residual")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const x = try up(r, dir, ref.str(o, "res_x")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "res_out")); + const g_all = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "p_g")); + const xout = try r.zeros(want.len); + try k.residualAdd(r.s, x, xout, try r.dev(g_all), @intCast(rows * d), 2, rows, d); + try check.sameBytes("glm.glue.residual", try r.back(xout, want.len), want); + const smx = try up(r, dir, ref.str(o, "sm_x")); + const want_sm = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "sm_out")); + const smout = try r.zeros(want_sm.len); + try k.streamMean(r.s, smx, smout, rows, d, 4); + try check.sameBytes("glm.glue.residual mean", try r.back(smout, want_sm.len), want_sm); + // f64 + const xv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "res_x"))); + const gv = f32s(g_all); + const gotv = try bf4(r.a, try r.back(xout, want.len)); + var t = Tol{}; + for (0..rows * d) |i| t.f64v(gotv[i], bf16f64(xv[i] + bf16f64(gv[i] + gv[rows * d + i])), 1.0 / 64.0, 3e-3); + const smv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "sm_x"))); + const got_sm = try bf4(r.a, try r.back(smout, want_sm.len)); + for (0..rows) |i| for (0..d) |j| { + var acc: f64 = 0; + for (0..4) |sv| acc += smv[(i * 4 + sv) * d + j]; + t.f64v(got_sm[i * d + j], acc / 4.0, 1.0 / 64.0, 3e-3); + }; + try t.done("glm.glue.residual"); + check.pass("glm.glue.residual", .{}); +} + +// -- glm.glue.swiglu ----------------------------------------------------------------------------- +fn swiglu(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.swiglu")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const w: usize = 6144; + const lim: f32 = @floatCast(ref.float(o, "lim")); + const gu = try up(r, dir, ref.str(o, "s_gu")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "s_out")); + const want_xs = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "s_xs")); + const out = try r.zeros(want.len); + const xs = try r.zeros(want_xs.len); + try k.swiglu(r.s, gu, out, xs, lim, rows, w); + try check.sameBytes("glm.glue.swiglu", try r.back(out, want.len), want); + try check.sameBytes("glm.glue.swiglu xs", try r.back(xs, want_xs.len), want_xs); + const guv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "s_gu"))); + const gotv = try bf4(r.a, try r.back(out, want.len)); + var t = Tol{}; + for (0..rows) |i| for (0..w) |j| { + const g = @min(guv[i * 2 * w + j], lim); + const u = @min(@max(guv[i * 2 * w + w + j], -lim), lim); + t.f64v(gotv[i * w + j], bf16f64(bf16f64(g / (1.0 + @exp(-g))) * u), 1.0 / 64.0, 3e-3); + }; + try t.done("glm.glue.swiglu"); + check.pass("glm.glue.swiglu", .{}); +} + +// -- glm.glue.router ----------------------------------------------------------------------------- +fn router(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.router")) orelse return; + _ = state; + var k = try cuda.glm_bmm.Bmm.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const ne: usize = @intCast(ref.int(o, "ne")); + const x = try up(r, dir, ref.str(o, "router_x")); + const w = try up(r, dir, ref.str(o, "router_w")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "router_out")); + const part = try r.zeros(8 * rows * ne * 4); + const out = try r.zeros(want.len); + try k.router(r.s, x, @intCast(d), .{ .w = w, .n = @intCast(ne), .k = @intCast(d) }, rows, out, part); + try check.sameBytes("glm.glue.router", try r.back(out, want.len), want); + const xv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "router_x"))); + const wv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "router_w"))); + const gotv = f32s(try r.back(out, want.len)); + var t = Tol{}; + for (0..rows) |i| for (0..ne) |e| { + var acc: f64 = 0; + for (0..d) |j| acc += xv[i * d + j] * wv[e * d + j]; + t.f32v(gotv, i * ne + e, acc, 1.0 / 1024.0, 1e-5); + }; + try t.done("glm.glue.router"); + check.pass("glm.glue.router", .{}); +} + +// -- glm.glue.select ----------------------------------------------------------------------------- +fn select(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.select")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const ne: usize = @intCast(ref.int(o, "ne")); + const topk: usize = @intCast(ref.int(o, "topk")); + const scale: f32 = @floatCast(ref.float(o, "scale")); + const lg = try up(r, dir, ref.str(o, "sel_logits")); + const bias = try up(r, dir, ref.str(o, "sel_bias")); + const want_pick = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "sel_pick")); + const want_wts = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "sel_wts")); + const pick = try r.zeros(want_pick.len); + const wts = try r.zeros(want_wts.len); + try k.topk(r.s, lg, bias, pick, wts, scale, rows, ne, topk); + try check.sameBytes("glm.glue.select", try r.back(pick, want_pick.len), want_pick); + try check.sameBytes("glm.glue.select wts", try r.back(wts, want_wts.len), want_wts); + // f64: the same picks by biased score, lowest id on ties, then the normalized weights + const lgv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "sel_logits"))); + const bv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "sel_bias"))); + const got_wts = f32s(try r.back(wts, want_wts.len)); + var t = Tol{}; + for (0..rows) |i| { + const score = try r.a.alloc(f64, ne); + const used = try r.a.alloc(bool, ne); + var total: f64 = 0; + for (0..ne) |e| { + score[e] = sigmoid(lgv[i * ne + e]); + used[e] = false; + } + const picks = try r.a.alloc(usize, topk); + for (0..topk) |kk| { + var best: f64 = -std.math.inf(f64); + var idx: usize = ne; + for (0..ne) |e| if (!used[e] and score[e] + bv[e] > best) { + best = score[e] + bv[e]; + idx = e; + }; + used[idx] = true; + picks[kk] = idx; + total += score[idx]; + } + for (0..topk) |kk| t.f32v(got_wts, i * (topk + 1) + kk, score[picks[kk]] / (total + 1e-20) * scale, 1.0 / 1024.0, 1e-6); + } + try t.done("glm.glue.select"); + check.pass("glm.glue.select", .{}); +} + +fn bf16f64(v: f64) f64 { + const u: u32 = @bitCast(@as(f32, @floatCast(v))); + const ru: u32 = u +% (0x7FFF + ((u >> 16) & 1)); + return @as(f64, @as(f32, @bitCast(ru & 0xFFFF0000))); +} + +// -- glm.glue.combine ---------------------------------------------------------------------------- +fn combine(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.glue.combine")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const y = try up(r, dir, ref.str(o, "c_y")); + const wts = try up(r, dir, ref.str(o, "c_wts")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "c_out")); + const out = try r.zeros(want.len); + try k.combine(r.s, y, wts, out, rows, d, 9); + try check.sameBytes("glm.glue.combine", try r.back(out, want.len), want); + std.debug.print("DBG c_y={s} c_wts={s} rows={d} d={d}\n", .{ ref.str(o, "c_y"), ref.str(o, "c_wts"), rows, d }); + const yv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "c_y"))); + std.debug.print("DBG yv len={d}\n", .{yv.len}); + const wv = f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "c_wts"))); + std.debug.print("DBG wv len={d} ptr={*}\n", .{ wv.len, wv.ptr }); + const gotv = f32s(try r.back(out, want.len)); + var t = Tol{}; + for (0..rows) |i| for (0..d) |j| { + var acc: f64 = 0; + for (0..9) |kk| acc += wv[i * 9 + kk] * yv[(i * 9 + kk) * d + j]; + t.f32v(gotv, i * d + j, acc, 1.0 / 1024.0, 1e-5); + }; + try t.done("glm.glue.combine"); + check.pass("glm.glue.combine", .{}); +} + +// -- glm.bmm ------------------------------------------------------------------------------------- +fn bmm(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.bmm")) orelse return; + _ = state; + var k = try cuda.glm_bmm.Bmm.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = @intCast(ref.int(o, "d")); + const x = try up(r, dir, ref.str(o, "bmm_x")); + const w = try up(r, dir, ref.str(o, "bmm_w")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "bmm_out")); + const w16 = cuda.glm_bmm.B16{ .w = w, .n = 4096, .k = @intCast(d) }; + const part = try r.zeros(cuda.glm_bmm.Bmm.partBytes(w16, rows)); + const out = try r.zeros(want.len); + try k.matmul(r.s, w16, x, @intCast(d), rows, out, false, part); + try check.sameBytes("glm.bmm", try r.back(out, want.len), want); + const xv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "bmm_x"))); + const wv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "bmm_w"))); + const gotv = try bf4(r.a, try r.back(out, want.len)); + var t = Tol{}; + for (0..rows) |i| for (0..4096) |e| { + var acc: f64 = 0; + for (0..d) |j| acc += xv[i * d + j] * wv[e * d + j]; + t.f64v(gotv[i * 4096 + e], acc, 1.0 / 64.0, 3e-3); + }; + try t.done("glm.bmm"); + check.pass("glm.bmm", .{}); +} + +// -- glm.sample.greedy --------------------------------------------------------------------------- +fn greedy(r: *Rig, state: *Run, dir: []const u8) !void { + const o = (try section(r, dir, "glm.sample.greedy")) orelse return; + _ = state; + var k = try glue.Glue.load(r.gpu.d); + defer k.unload(); + const rows: usize = @intCast(ref.int(o, "rows")); + const d: usize = 4096; + const lg = try up(r, dir, ref.str(o, "g_logits")); + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "g_tok")); + const tok = try r.zeros(want.len); + try k.greedy(r.s, lg, tok, rows, d); + try check.sameBytes("glm.sample.greedy", try r.back(tok, want.len), want); + // f64: the highest logit, lowest id on ties + const lgv = try bf4(r.a, try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "g_logits"))); + const got = std.mem.bytesAsSlice(i32, try r.back(tok, want.len)); + for (0..rows) |i| { + var best: f64 = lgv[i * d]; + var bi: usize = 0; + for (0..d) |j| if (lgv[i * d + j] > best) { + best = lgv[i * d + j]; + bi = j; + }; + if (got[i] != @as(i32, @intCast(bi))) return error.TestFailed; + } + check.pass("glm.sample.greedy", .{}); +} + +// -- the gap fixtures (glm_oracle_gaps.py's "gaps"): folded into the rmsnorm and combine ids above --- + +fn rmsnormGaps(r: *Rig, state: *Run, dir: []const u8) !void { + gaps.rmsnorm(r, dir) catch |e| state.fail(e, "glm.glue.rmsnorm"); +} + +fn combineGaps(r: *Rig, state: *Run, dir: []const u8) !void { + gaps.combine(r, dir) catch |e| state.fail(e, "glm.glue.combine"); + gaps.combineMoe(r, dir) catch |e| state.fail(e, "glm.glue.combine"); +} diff --git a/zig/tests/cuda/glm_head_tests.zig b/zig/tests/cuda/glm_head_tests.zig new file mode 100644 index 0000000000..9f9f72c114 --- /dev/null +++ b/zig/tests/cuda/glm_head_tests.zig @@ -0,0 +1,116 @@ +//! glm.head: the model's tail on both ranks' head halves vs the Python engine (compute's end, sample_rows greedy). +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const glm_tests = @import("glm_tests.zig"); +const glue_tests = @import("glm_glue_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const head = cuda.glm_head; +const Tol = glue_tests.Tol; + +const id = "glm.head"; + +pub fn run(r: *glm_tests.Rig, state: *glm_tests.Run, dir: []const u8) !void { + check_(r, dir) catch |e| state.fail(e, id); +} + +fn bfv(a: std.mem.Allocator, bytes: []const u8) ![]f64 { + const u = std.mem.bytesAsSlice(u16, bytes); + const out = try a.alloc(f64, u.len); + for (u, out) |v, *o| o.* = ref.bf(v); + return out; +} + +fn check_(r: *glm_tests.Rig, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) return std.debug.print("SKIP {s}\n", .{id}); + var man = try ref.Manifest.open(r.gpu.gpa, r.gpu.io, dir); + defer man.deinit(); + const o = (man.root().get("head") orelse return std.debug.print("SKIP {s}\n", .{id})).object; + const p = o.get("params").?.object; + const eps: f32 = @floatCast(ref.float(p, "eps")); + const d: usize = @intCast(ref.int(p, "hidden")); + const streams: usize = @intCast(ref.int(p, "streams")); + var g = try cuda.glm_glue.Glue.load(r.gpu.d); + defer g.unload(); + var b = try cuda.glm_bmm.Bmm.load(r.gpu.d); + defer b.unload(); + const norm_bytes = try ref.readFile(r.a, r.gpu.io, dir, ref.str(o, "norm")); + const norm = try r.dev(norm_bytes); + const normv = try bfv(r.a, norm_bytes); + const hs = o.get("heads").?.array.items; + if (hs.len != 2) return error.GlmHeadShape; + var ws: [2]head.HeadW = undefined; + var hostw: [2][]const u8 = undefined; + for (hs, &ws, &hostw) |hv, *w, *hw| { + const h = hv.object; + hw.* = try ref.readFile(r.a, r.gpu.io, dir, ref.str(h, "w")); + w.* = .{ .norm = norm, .head = .{ .w = try r.dev(hw.*), .n = @intCast(ref.int(h, "n")), .k = @intCast(ref.int(h, "k")) }, .offset = @intCast(ref.int(h, "offset")) }; + } + var t = Tol{}; + var it = o.get("cases").?.object.iterator(); + var count: usize = 0; + while (it.next()) |kv| : (count += 1) { + const tag = kv.key_ptr.*; + const e = kv.value_ptr.object; + const rows: usize = @intCast(ref.int(e, "rows")); + const last = e.get("prefill").?.bool; + const m: usize = if (last) 1 else rows; + const xb = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "x")); + const x = try r.dev(xb); + const lgs = e.get("logits").?.array.items; + const gathered = try r.a.alloc(f32, 2 * m * 2); + for (0..2) |rank| { + var sc = try head.Scratch.init(r.gpu.d, r.s, ws[rank], d, rows); + defer sc.deinit(); + try head.forward(&g, &b, r.s, eps, d, streams, ws[rank], &sc, x, rows, last); + const nm = try std.fmt.allocPrint(r.a, "{s} {s}", .{ id, tag }); + if (rank == 0) for ([_][]const u8{ "hidden", "fnormed", "fxs" }, [_]u64{ sc.hidden, sc.fnormed, sc.fxs }) |key, ptr| { + const want = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, key)); + try check.sameBytes(try std.fmt.allocPrint(r.a, "{s} {s}", .{ nm, key }), try r.back(ptr, want.len), want); + }; + const want_lg = try ref.readFile(r.a, r.gpu.io, dir, lgs[rank].string); + const got_lg = try r.back(sc.logits, want_lg.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "{s} logits rank {d}", .{ nm, rank }), got_lg, want_lg); + const cand = std.mem.bytesAsSlice(f32, try r.back(sc.cand, m * 8)); + @memcpy(gathered[rank * m * 2 ..][0 .. m * 2], cand); + if (rank == 0) try f64Stages(r.a, &t, xb, try r.back(sc.hidden, rows * d * 2), try r.back(sc.fnormed, rows * d * 2), normv, hostw[0], got_lg, eps, rows, d, streams, m, ws[0].head.n); + } + const toks = try r.a.alloc(u32, m); + try head.choose(gathered, 2, m, toks); + const want_tok = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "tok")); + try check.sameBytes(try std.fmt.allocPrint(r.a, "{s} {s} tokens", .{ id, tag }), std.mem.sliceAsBytes(toks), want_tok); + std.debug.print("RESULT {s} {s}: {d} rows, {d} tokens bit-exact across both ranks\n", .{ id, tag, rows, m }); + } + try t.done(id); + std.debug.print("RESULT {s}: {d} cases bit-exact\n", .{ id, count }); + check.pass("{s}", .{id}); +} + +/// Rank 0's stages in f64 from each stage's own inputs: the streams' mean, the final norm, sampled logit columns. +fn f64Stages(a: std.mem.Allocator, t: *Tol, xb: []const u8, hb: []const u8, fb: []const u8, normv: []const f64, wb: []const u8, lb: []const u8, eps: f32, rows: usize, d: usize, streams: usize, m: usize, n: usize) !void { + const xv = try bfv(a, xb); + const hv = try bfv(a, hb); + const fv = try bfv(a, fb); + const lv = try bfv(a, lb); + const wv = std.mem.bytesAsSlice(u16, wb); + for (0..rows) |i| { + var ss: f64 = 0; + for (0..d) |j| { + var acc: f64 = 0; + for (0..streams) |s| acc += xv[(i * streams + s) * d + j]; + t.f64v(hv[i * d + j], acc / @as(f64, @floatFromInt(streams)), 1.0 / 64.0, 3e-3); + ss += hv[i * d + j] * hv[i * d + j]; + } + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(d)) + eps); + for (0..d) |j| t.f64v(fv[i * d + j], normv[j] * hv[i * d + j] * rinv, 1.0 / 64.0, 3e-3); + } + const first = rows - m; + for (0..m) |i| { + var col: usize = i * 7; + while (col < n) : (col += 997) { + var acc: f64 = 0; + for (0..d) |j| acc += fv[(first + i) * d + j] * ref.bf(wv[col * d + j]); + t.f64v(lv[i * n + col], acc, 1.0 / 64.0, 3e-3); + } + } +} diff --git a/zig/tests/cuda/glm_kda_block.zig b/zig/tests/cuda/glm_kda_block.zig new file mode 100644 index 0000000000..ff8da3360a --- /dev/null +++ b/zig/tests/cuda/glm_kda_block.zig @@ -0,0 +1,228 @@ +//! The KDA attention block on real tensors: qkv and the bf16 pair, the chunked chain and its cut, o_proj's partials. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const base = @import("glm_kda_tests.zig"); +const exl3 = cuda.exl3; +const exl3_load = @import("core").exl3_load; +const exl3_upload = cuda.exl3_upload; +const bmm_mod = cuda.glm_bmm; +const chunked_mod = cuda.glm_kda_chunk; + +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; + +/// Three [start, end) runs: the fused q|k|v group's per-rank trellis columns and svh rows. +fn runs(v: std.json.Value) ![3][2]usize { + if (v.array.items.len != 3) return error.BadRuns; + var out: [3][2]usize = undefined; + for (v.array.items, 0..) |item, i| { + out[i][0] = @intCast(item.array.items[0].integer); + out[i][1] = @intCast(item.array.items[1].integer); + } + return out; +} + +fn groupBySuffix(scan: *exl3_load.Scan, suffix: []const u8) !usize { + for (scan.groups, 0..) |grp, i| if (std.mem.endsWith(u8, grp.prefix, suffix)) return i; + return error.GroupMissing; +} + +/// The qkv group's trellis: the three runs' columns stacked back into one [K/16, n/16, words] tensor. +fn slicedTrellis(a: std.mem.Allocator, t: exl3_load.Parts, cols: [3][2]usize) ![]u8 { + const wcol = t.trellis.dim(2) * 2; + const full_w = t.trellis.dim(1); + var width: usize = 0; + for (cols) |rn| width += rn[1] - rn[0]; + const out = try a.alloc(u8, t.trellis.dim(0) * width * wcol); + for (0..t.trellis.dim(0)) |row| { + var dst: usize = 0; + for (cols) |rn| { + const len = (rn[1] - rn[0]) * wcol; + const src = (row * full_w + rn[0]) * wcol; + @memcpy(out[row * width * wcol + dst ..][0..len], t.trellis.bytes[src..][0..len]); + dst += len; + } + } + return out; +} + +/// The qkv group's svh: the three runs' rows stacked back into one [n] fp16 scale row. +fn slicedRows(a: std.mem.Allocator, t: exl3_load.Parts, rows: [3][2]usize) ![]u8 { + var total: usize = 0; + for (rows) |rn| total += rn[1] - rn[0]; + const out = try a.alloc(u8, total * 2); + var dst: usize = 0; + for (rows) |rn| { + const len = (rn[1] - rn[0]) * 2; + @memcpy(out[dst..][0..len], t.svh.bytes[rn[0] * 2 ..][0..len]); + dst += len; + } + return out; +} + +pub fn block(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + base.skip("glm.kda.block"); + return; + } + var man = try ref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + _ = state; + const o = man.kda() orelse { + base.skip("glm.kda.block"); + return; + }; + const q = try base.Params.read(o); + const model_dir = ref.str(man.root(), "model"); + const e = o.get("block").?.object; + const rows: u32 = @intCast(ref.int(e, "rows")); + const cut: u32 = @intCast(ref.int(e, "cut")); + const qkv_o = e.get("qkv").?.object; + const qkv_n: usize = @intCast(ref.int(qkv_o, "n")); + + var scan = try exl3_load.scan(r.gpu.gpa, r.gpu.io, model_dir); + defer scan.deinit(); + var up = exl3_upload.Uploader{ .d = r.gpu.d, .gpa = r.gpu.gpa }; + defer up.deinit(); + var dense = try exl3.Dense.load(r.gpu.d); + defer dense.unload(); + + // the fused q|k|v group: the rank's three runs stacked back into one [6912, 4096] linear + const gi = try groupBySuffix(&scan, ref.str(qkv_o, "prefix")); + const g = scan.groups[gi]; + const parts = try scan.parts(r.a, gi); + const t_cols = try runs(qkv_o.get("t_runs").?); + const s_rows = try runs(qkv_o.get("s_runs").?); + var t2 = parts.trellis; + t2.bytes = try slicedTrellis(r.a, parts, t_cols); + var s2 = parts.svh; + s2.bytes = try slicedRows(r.a, parts, s_rows); + const w = try up.dense(.{ .trellis = t2, .suh = parts.suh, .svh = s2, .bias = parts.bias, .suh_packed = parts.suh_packed }, .{ .k = @as(usize, @intCast(ref.int(qkv_o, "k"))), .n = qkv_n, .halves = @as(u32, @intCast(g.halves)), .codebook = g.codebook }); + try check.expect(w.k == @as(u32, @intCast(ref.int(qkv_o, "k"))) and w.n == @as(u32, @intCast(qkv_n)), "glm.kda.block: qkv weight shape", .{}); + + const normed = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "normed"))); + const qkv_tmp = try r.zeros(rows * qkv_n * 2); + try dense.prompt(r.s, w, normed, .bf16, rows, qkv_tmp, .bf16); + + // the plain pair and b: bf16 [n, k] exactly as the engine holds them + var bmm = try bmm_mod.Bmm.load(r.gpu.d); + defer bmm.unload(); + const mms = e.get("mms").?.object; + const fa_w = try mm16(r, dir, mms, "fa"); + const ga_w = try mm16(r, dir, mms, "ga"); + const fb_w = try mm16(r, dir, mms, "fb"); + const gb_w = try mm16(r, dir, mms, "gb"); + const b_w = try mm16(r, dir, mms, "b"); + + const fa_tmp = try r.zeros(rows * @as(usize, fa_w.n) * 2); + const ga_tmp = try r.zeros(rows * @as(usize, ga_w.n) * 2); + const b_tmp = try r.zeros(rows * @as(usize, b_w.n) * 2); + const fa_part = try r.zeros(bmm_mod.Bmm.partBytes(fa_w, rows)); + const ga_part = try r.zeros(bmm_mod.Bmm.partBytes(ga_w, rows)); + const b_part = try r.zeros(bmm_mod.Bmm.partBytes(b_w, rows)); + try bmm.matmul(r.s, fa_w, normed, @intCast(fa_w.k), rows, fa_tmp, false, fa_part); + try bmm.matmul(r.s, ga_w, normed, @intCast(ga_w.k), rows, ga_tmp, false, ga_part); + try bmm.matmul(r.s, b_w, normed, @intCast(b_w.k), rows, b_tmp, false, b_part); + const a_rows: usize = @as(usize, fb_w.n); + const a_buf = try r.zeros(rows * a_rows * 2); + const fa_part2 = try r.zeros(bmm_mod.Bmm.partBytes(fb_w, rows)); + try bmm.matmul(r.s, fb_w, fa_tmp, @intCast(fa_w.n), rows, a_buf, false, fa_part2); + const g_buf = try r.zeros(rows * a_rows * 2); + const ga_part2 = try r.zeros(bmm_mod.Bmm.partBytes(gb_w, rows)); + try bmm.matmul(r.s, gb_w, ga_tmp, @intCast(ga_w.n), rows, g_buf, false, ga_part2); + + // p's rows: q|k|v | f_a | g_a | b; the fixture's p holds the assembled bits + const stride: usize = @intCast(q.p_stride); + const p_bytes = rows * stride * 2; + const p_buf = try r.zeros(p_bytes); + const pbuf: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = p_buf, .len = p_bytes }; + for (0..rows) |ri| { + try pbuf.copyFrom(ri * stride * 2, qkv_tmp + ri * qkv_n * 2, qkv_n * 2, r.s.handle); + try pbuf.copyFrom(ri * stride * 2 + qkv_n * 2, fa_tmp + ri * @as(usize, fa_w.n) * 2, @as(usize, fa_w.n) * 2, r.s.handle); + try pbuf.copyFrom(ri * stride * 2 + (qkv_n + 128) * 2, ga_tmp + ri * @as(usize, ga_w.n) * 2, @as(usize, ga_w.n) * 2, r.s.handle); + try pbuf.copyFrom(ri * stride * 2 + @as(usize, @intCast(q.b_off)) * 2, b_tmp + ri * @as(usize, b_w.n) * 2, @as(usize, b_w.n) * 2, r.s.handle); + } + const want_p = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "p")); + // the chunked chain over the assembled p, and its cut + try check.sameBytes("glm.kda.block p", try r.back(p_buf, p_bytes), want_p); + var chunk = try chunked_mod.Chunk.load(r.gpu.d); + defer chunk.unload(); + const a_log = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_a_log.bin")); + const dt_bias = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_dt_bias.bin")); + const norm_w = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_norm.bin")); + const conv_w = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_conv.bin")); + const conv0 = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "conv0"))); + const state0 = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state0"))); + const want_out = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "out")); + const want_state = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state_out")); + const out = try r.zeros(want_out.len); + const st = try r.zeros(want_state.len); + try chunk.chain(&r.s, q.h, p_buf, q.p_stride, @as(i32, @intCast(q.b_off)), a_buf, @as(i64, @intCast(q.per)), g_buf, @as(i64, @intCast(q.per)), conv0, conv_w, state0, a_log, dt_bias, norm_w, q.eps, q.lower, rows, 0, out, st); + try check.sameBytes("glm.kda.block out", try r.back(out, want_out.len), want_out); + try check.sameBytes("glm.kda.block state", try r.back(st, want_state.len), want_state); + + const cut_out1 = try r.zeros(@as(usize, cut) * q.per * 2); + const cut_st1 = try r.zeros(want_state.len); + try chunk.chain(&r.s, q.h, p_buf, q.p_stride, @as(i32, @intCast(q.b_off)), a_buf, @as(i64, @intCast(q.per)), g_buf, @as(i64, @intCast(q.per)), conv0, conv_w, state0, a_log, dt_bias, norm_w, q.eps, q.lower, cut, 0, cut_out1, cut_st1); + const ch: usize = q.per * 3; + const conv_cut = try r.a.alloc(u8, 3 * ch * 2); + for (0..3) |ri| @memcpy(conv_cut[ri * ch * 2 ..][0 .. ch * 2], want_p[(cut - 3 + ri) * stride * 2 ..][0 .. ch * 2]); + const cut_out2 = try r.zeros((@as(usize, rows) - cut) * q.per * 2); + const cut_st2 = try r.zeros(want_state.len); + try chunk.chain(&r.s, q.h, p_buf + @as(u64, cut) * @as(u64, stride) * 2, q.p_stride, @as(i32, @intCast(q.b_off)), a_buf + @as(u64, cut) * @as(u64, q.per) * 2, @as(i64, @intCast(q.per)), g_buf + @as(u64, cut) * @as(u64, q.per) * 2, @as(i64, @intCast(q.per)), try r.dev(conv_cut), conv_w, cut_st1, a_log, dt_bias, norm_w, q.eps, q.lower, rows - cut, 0, cut_out2, cut_st2); + const got_c1 = try r.back(cut_out1, @as(usize, cut) * q.per * 2); + const got_c2 = try r.back(cut_out2, (@as(usize, rows) - cut) * q.per * 2); + const cat = try r.a.alloc(u8, got_c1.len + got_c2.len); + @memcpy(cat[0..got_c1.len], got_c1); + @memcpy(cat[got_c1.len..], got_c2); + try check.sameBytes("glm.kda.block cut rows", cat, want_out); + try check.sameBytes("glm.kda.block cut state", try r.back(cut_st2, want_state.len), want_state); + + // o_proj: the engine's fp32 partials (no rounding), through the group or the plain weight + const o_o = e.get("o_proj").?.object; + const want_part = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "part")); + const part = try r.zeros(want_part.len); + if (o_o.get("w") == null) { + const gi_o = try groupBySuffix(&scan, ref.str(o_o, "prefix")); + const go = scan.groups[gi_o]; + const parts_o = try scan.parts(r.a, gi_o); + const half = parts_o.trellis.dim(0) / 2; + const suh_half = parts_o.suh.dim(0) / 2; + var t3 = parts_o.trellis; + t3.bytes = parts_o.trellis.bytes[0 .. half * t3.dim(1) * t3.dim(2) * 2]; + var suh3 = parts_o.suh; + suh3.bytes = parts_o.suh.bytes[0 .. suh_half * 2]; + const w_o = try up.dense(.{ .trellis = t3, .suh = suh3, .svh = parts_o.svh, .bias = parts_o.bias, .suh_packed = parts_o.suh_packed }, .{ .k = half * 16, .n = @as(usize, @intCast(ref.int(o_o, "n"))), .halves = @as(u32, @intCast(go.halves)), .codebook = go.codebook }); + try check.expect(w_o.k == @as(u32, @intCast(ref.int(o_o, "k"))), "glm.kda.block: o_proj k is the rank's rows", .{}); + try dense.prompt(r.s, w_o, out, .bf16, rows, part, .f32); + } else { + const w_o = bmm_mod.B16{ .w = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(o_o, "w"))), .n = @intCast(ref.int(o_o, "n")), .k = @intCast(ref.int(o_o, "k")) }; + const part_scr = try r.zeros(bmm_mod.Bmm.partBytes(w_o, rows)); + try bmm.matmul(r.s, w_o, out, @intCast(w_o.k), rows, part, true, part_scr); + } + try check.sameBytes("glm.kda.block part", try r.back(part, want_part.len), want_part); + + // the block's own float64 reference over its prep outputs + const out64 = try r.a.alloc(f64, @as(usize, rows) * q.per); + const st64 = try r.a.alloc(f64, q.per * ref.DK); + try ref.chunkedRef(r.a, ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "prep_q"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "prep_k"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "prep_v"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "prep_lg"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "prep_beta"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state0"))), ref.bf16s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "g"))), ref.bf16s(try ref.readFile(r.a, r.gpu.io, dir, "kda_norm.bin")), q.eps, q.h, rows, 0, out64, st64); + var t = base.T{}; + try t.close("glm.kda.block out f64", try r.back(out, want_out.len), out64, false, 1.0 / 128.0, 2e-3); + try t.close("glm.kda.block state f64", try r.back(st, want_state.len), st64, true, 2e-2, 2e-3); + std.debug.print("RESULT glm.kda.block: f64 worst {d:.4} of tolerance\n", .{t.worst}); + check.pass("glm.kda.block", .{}); +} + +fn mm16(r: *Rig, dir: []const u8, mms: std.json.ObjectMap, key: []const u8) !bmm_mod.B16 { + const mo = mms.get(key).?.object; + return .{ + .w = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(mo, "file"))), + .n = @intCast(ref.int(mo, "n")), + .k = @intCast(ref.int(mo, "k")), + }; +} diff --git a/zig/tests/cuda/glm_kda_ref.zig b/zig/tests/cuda/glm_kda_ref.zig new file mode 100644 index 0000000000..7617a14502 --- /dev/null +++ b/zig/tests/cuda/glm_kda_ref.zig @@ -0,0 +1,306 @@ +//! The KDA checks' helpers: manifest, fixture bytes, f64 references (kda_chunked's recurrent and chunked ones). + +const std = @import("std"); + +pub const DK = 128; +const BT = 32; // rows of a sub-chunk + +pub fn bf(u: u16) f64 { + const f: f32 = @floatCast(@as(f32, @bitCast(@as(u32, u) << 16))); + return f; +} + +pub fn f32Of(u: u16) f32 { + return @floatCast(@as(f32, @bitCast(@as(u32, u) << 16))); +} + +pub const Manifest = struct { + parsed: std.json.Parsed(std.json.Value), + + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Manifest { + const path = try std.fs.path.join(gpa, &.{ dir, "manifest.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + return .{ .parsed = parsed }; + } + + pub fn root(m: *const Manifest) std.json.ObjectMap { + return m.parsed.value.object; + } + + pub fn deinit(m: *Manifest) void { + m.parsed.deinit(); + } + + pub fn kda(m: *const Manifest) ?std.json.ObjectMap { + const v = m.parsed.value.object.get("kda") orelse return null; + return v.object; + } +}; + +/// One array file's raw bytes (little-endian); the caller frees them. +pub fn readFile(a: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) ![]u8 { + const path = try std.fs.path.join(a, &.{ dir, name }); + defer a.free(path); + const raw = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 32)); + defer a.free(raw); + const out = try a.alignedAlloc(u8, .@"8", raw.len); + @memcpy(out, raw); + return out; +} + +pub fn f32s(bytes: []const u8) []const f32 { + return @alignCast(std.mem.bytesAsSlice(f32, bytes)); +} + +/// A bf16 array's values, widened (exact in f32). +pub fn bf16s(bytes: []const u8) []const f32 { + const n = bytes.len / 2; + const src = std.mem.bytesAsSlice(u16, bytes); + const o = std.heap.page_allocator.alloc(f32, n) catch unreachable; + for (0..n) |i| o[i] = @floatCast(@as(f32, @bitCast(@as(u32, src[i]) << 16))); + return o; +} + +/// The gated RMSNorm the out kernel applies: norm_w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(gate), f64. +pub fn gatedNorm(out: []f64, norm_w: []const f32, gate: []const f32, eps: f64, h: usize, rows: usize) void { + var r: usize = 0; + while (r < rows) : (r += 1) { + var head: usize = 0; + while (head < h) : (head += 1) { + const base = (r * h + head) * DK; + var ss: f64 = 0; + for (0..DK) |vv| ss += out[base + vv] * out[base + vv]; + const rinv = 1.0 / @sqrt(ss / @as(f64, DK) + eps); + for (0..DK) |vv| { + const gt = 1.0 / (1.0 + @exp(-@as(f64, gate[base + vv]))); + out[base + vv] = @as(f64, norm_w[vv]) * out[base + vv] * rinv * gt; + } + } + } +} + +pub fn str(o: std.json.ObjectMap, key: []const u8) []const u8 { + return o.get(key).?.string; +} + +pub fn int(o: std.json.ObjectMap, key: []const u8) i64 { + return o.get(key).?.integer; +} + +pub fn float(o: std.json.ObjectMap, key: []const u8) f64 { + return switch (o.get(key).?) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => unreachable, + }; +} + +/// recurrent_reference in f64: q k lg v fp32 [rows, H, 128], beta [rows, H], state [H, 128, 128] -> out, state f64. +pub fn serialRef( + a: std.mem.Allocator, + q: []const f32, + k: []const f32, + v: []const f32, + lg: []const f32, + beta: []const f32, + state_in: []const f32, + gate: []const f32, + norm_w: []const f32, + eps: f32, + h: usize, + rows: usize, + out: []f64, + state: []f64, +) !void { + const y = try a.alloc(f64, rows * h * DK); + const S = try a.alloc(f64, h * DK * DK); + defer a.free(S); + defer a.free(y); + for (S, state_in) |*x, s| x.* = s; + const g = try a.alloc(f64, DK); + defer a.free(g); + const delta = try a.alloc(f64, DK); + defer a.free(delta); + var r: usize = 0; + while (r < rows) : (r += 1) { + for (0..h) |head| { + const base = (r * h + head) * DK; + for (0..DK) |c| g[c] = @exp(@as(f64, lg[base + c])); + for (0..DK) |c| { + for (0..DK) |vv| S[head * DK * DK + vv * DK + c] *= g[c]; + } + for (0..DK) |vv| { + var kv: f64 = 0; + for (0..DK) |c| kv += S[head * DK * DK + vv * DK + c] * @as(f64, k[base + c]); + delta[vv] = (@as(f64, v[base + vv]) - kv) * @as(f64, beta[r * h + head]); + } + for (0..DK) |c| { + for (0..DK) |vv| S[head * DK * DK + vv * DK + c] += @as(f64, k[base + c]) * delta[vv]; + } + for (0..DK) |vv| { + var o: f64 = 0; + for (0..DK) |c| o += S[head * DK * DK + vv * DK + c] * @as(f64, q[base + c]); + y[base + vv] = o; + } + } + } + for (state, S) |*x, s| x.* = s; + gatedNorm(y, norm_w, gate, @as(f64, eps), h, rows); + @memcpy(out, y); +} + +/// chunked_reference in f64: serialRef's inputs in sub-chunks at absolute pos; padded rows decay 1, no k, v or beta. +pub fn chunkedRef( + a: std.mem.Allocator, + q: []const f32, + k: []const f32, + v: []const f32, + lg: []const f32, + beta: []const f32, + state_in: []const f32, + gate: []const f32, + norm_w: []const f32, + eps: f32, + h: usize, + rows: usize, + pos: usize, + out: []f64, + state: []f64, +) !void { + const off = pos % BT; + const nc = (rows + off + BT - 1) / BT; + + const S = try a.alloc(f64, h * DK * DK); + defer a.free(S); + for (S, state_in) |*x, s| x.* = s; + const y = try a.alloc(f64, rows * h * DK); + defer a.free(y); + + const gx = try a.alloc(f64, BT * DK); // [t][c] = g of row t (1 outside) + defer a.free(gx); + const fwd = try a.alloc(f64, BT * DK); // [t][c] = E(0 -> t) + defer a.free(fwd); + const back = try a.alloc(f64, BT * DK); // [t][c] = E(t -> 31) + defer a.free(back); + const Vn = try a.alloc(f64, BT * DK); + defer a.free(Vn); + const L = try a.alloc(f64, BT * BT); + defer a.free(L); + const M = try a.alloc(f64, BT * BT); + defer a.free(M); + const T = try a.alloc(f64, BT * BT); + defer a.free(T); + @memset(L, 0); + const e_run = try a.alloc(f64, DK); // E(t -> s) of the pair being accumulated + defer a.free(e_run); + @memset(M, 0); + @memset(T, 0); + const Y = try a.alloc(f64, BT * DK); + defer a.free(Y); + + var c: usize = 0; + while (c < nc) : (c += 1) { + for (0..h) |head| { + for (0..BT) |t| { + const row_i = c * BT + t; + const ok = row_i >= off and row_i < off + rows; + const row = row_i -| off; + for (0..DK) |cc| gx[t * DK + cc] = if (ok) @exp(@as(f64, lg[(row * h + head) * DK + cc])) else 1.0; + } + for (0..DK) |cc| { + fwd[cc] = gx[cc]; + for (1..BT) |t| fwd[t * DK + cc] = fwd[(t - 1) * DK + cc] * gx[t * DK + cc]; + back[(BT - 1) * DK + cc] = 1.0; + var t: usize = BT - 1; + while (t > 0) : (t -= 1) back[(t - 1) * DK + cc] = back[t * DK + cc] * gx[t * DK + cc]; + } + // L[s][t] (t < s) and M[s][t] (t <= s); E(t -> s) multiplied up from 1 per t, exactly the engine's order + for (0..BT) |t| { + for (0..DK) |cc| e_run[cc] = 1.0; + var s: usize = t; + while (s < BT) : (s += 1) { + if (s > t) { + for (0..DK) |cc| e_run[cc] *= gx[s * DK + cc]; + } + var akk: f64 = 0; + var aqk: f64 = 0; + for (0..DK) |cc| { + const kt = rowAt(k, h, off, rows, c * BT + t, head, cc); + akk += rowAt(k, h, off, rows, c * BT + s, head, cc) * e_run[cc] * kt; + aqk += rowAt(q, h, off, rows, c * BT + s, head, cc) * e_run[cc] * kt; + } + const bs = rowBeta(beta, h, off, rows, c * BT + s, head); + if (s > t) L[s * BT + t] = bs * akk; + M[s * BT + t] = aqk; + } + } + // T = (I + L)^-1 by forward substitution: T[s] = -(L[s] T) with T[s][s] = 1 + for (0..BT) |s| { + for (0..BT) |t| T[s * BT + t] = if (s == t) 1.0 else 0.0; + } + for (1..BT) |s| { + for (0..BT) |m| { + var acc: f64 = 0; + for (0..BT) |mm| acc += L[s * BT + mm] * T[mm * BT + m]; + T[s * BT + m] = -acc; + } + T[s * BT + s] = 1.0; + } + // Y[t][vv] = beta_t (v_t[vv] - (k_t fwd_t) S^T) + for (0..BT) |t| { + const bt = rowBeta(beta, h, off, rows, c * BT + t, head); + for (0..DK) |vv| { + var kv: f64 = 0; + for (0..DK) |cc| kv += rowAt(k, h, off, rows, c * BT + t, head, cc) * fwd[t * DK + cc] * + S[head * DK * DK + vv * DK + cc]; + Y[t * DK + vv] = bt * (rowAt(v, h, off, rows, c * BT + t, head, vv) - kv); + } + } + // Vn[t][vv] = sum_s T[t][s] Y[s][vv], then the rows: (q fwd) S^T + M V' + for (0..BT) |t| for (0..DK) |vv| { + var vn: f64 = 0; + for (0..BT) |s| vn += T[t * BT + s] * Y[s * DK + vv]; + Vn[t * DK + vv] = vn; + }; + for (0..BT) |t| { + if (c * BT + t < off) continue; + const gr = c * BT + t - off; + if (gr >= rows) continue; + const row = gr; + for (0..DK) |vv| { + var o: f64 = 0; + for (0..DK) |cc| o += rowAt(q, h, off, rows, c * BT + t, head, cc) * fwd[t * DK + cc] * + S[head * DK * DK + vv * DK + cc]; + for (0..BT) |s| o += M[t * BT + s] * Vn[s * DK + vv]; + y[(row * h + head) * DK + vv] = o; + } + } + // the state: S diag(E(0 -> 31)) + V'^T (k back) + for (0..DK) |cc| { + for (0..DK) |vv| S[head * DK * DK + vv * DK + cc] *= fwd[(BT - 1) * DK + cc]; + } + for (0..BT) |t| { + for (0..DK) |cc| { + const kb = rowAt(k, h, off, rows, c * BT + t, head, cc) * back[t * DK + cc]; + for (0..DK) |vv| S[head * DK * DK + vv * DK + cc] += kb * Vn[t * DK + vv]; + } + } + } + } + for (state, S) |*x, s| x.* = s; + gatedNorm(y, norm_w, gate, @as(f64, eps), h, rows); + @memcpy(out, y); +} + +fn rowAt(arr: []const f32, h: usize, off: usize, rows: usize, gr: usize, head: usize, c: usize) f64 { + if (gr < off or gr >= off + rows) return 0; + return arr[((gr - off) * h + head) * DK + c]; +} + +fn rowBeta(beta: []const f32, h: usize, off: usize, rows: usize, gr: usize, head: usize) f64 { + if (gr < off or gr >= off + rows) return 0; + return beta[(gr - off) * h + head]; +} diff --git a/zig/tests/cuda/glm_kda_tests.zig b/zig/tests/cuda/glm_kda_tests.zig new file mode 100644 index 0000000000..5ad324139f --- /dev/null +++ b/zig/tests/cuda/glm_kda_tests.zig @@ -0,0 +1,252 @@ +//! GLM-5.3-Flash's KDA on CUDA: decode chain and prompt chunks against float64 and the Python engine's bytes. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const block_tests = @import("glm_kda_block.zig"); +const kda = cuda.glm_kda; +const chunked_mod = cuda.glm_kda_chunk; + +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const Failed = check.Failed; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + chain(r, state, dir) catch |e| state.fail(e, "glm.kda.chain"); + chunked(r, state, dir) catch |e| state.fail(e, "glm.kda.chunked"); + block_tests.block(r, state, dir) catch |e| state.fail(e, "glm.kda.block"); +} + +pub fn skip(id: []const u8) void { + std.debug.print("SKIP {s}\n", .{id}); +} + +/// The largest difference as a fraction of its tolerance, over one check's float64 comparisons. +pub const T = struct { + worst: f64 = 0, + at: usize = 0, + + pub fn close(self: *T, name: []const u8, got: []const u8, want: []const f64, is_f32: bool, scale: f64, floor: f64) Failed!void { + for (want, 0..) |w, i| { + const g: f64 = if (is_f32) + @as(f32, @bitCast(std.mem.readInt(u32, got[4 * i ..][0..4], .little))) + else + ref.bf(std.mem.readInt(u16, got[2 * i ..][0..2], .little)); + const q = @abs(g - w) / (@abs(w) * scale + floor); + if (q > self.worst) { + self.worst = q; + self.at = i; + } + } + if (self.worst > 1.0) { + std.debug.print("FAIL {s}: f64 worst {d:.4} of tolerance at element {d}\n", .{ name, self.worst, self.at }); + return error.TestFailed; + } + } +}; + +/// The fixture's params (shared by every check). +pub const Params = struct { + h: u32, + b_off: i32, + p_stride: i64, + eps: f32, + lower: f32, + per: usize, // h * dk + + pub fn read(o: std.json.ObjectMap) !Params { + const p = o.get("params").?.object; + const h: u32 = @intCast(ref.int(p, "H")); + return .{ + .h = h, + .b_off = @intCast(ref.int(p, "b_off")), + .p_stride = ref.int(p, "p_stride"), + .eps = @floatCast(ref.float(p, "eps")), + .lower = @floatCast(ref.float(p, "lower")), + .per = @as(usize, h) * ref.DK, + }; + } +}; + +fn chain(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + _ = state; + skip("glm.kda.chain"); + return; + } + var man = try ref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + const o = man.kda() orelse { + skip("glm.kda.chain"); + return; + }; + const q = try Params.read(o); + var k = try kda.Kda.load(r.gpu.d); + defer k.unload(); + + const a_log = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_a_log.bin")); + const dt_bias = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_dt_bias.bin")); + const norm_w = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_norm.bin")); + const conv_w = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_conv.bin")); + + var t = T{}; + const chains = o.get("chain").?.object; + for ([_][]const u8{ "1", "3", "8", "64", "128" }) |key| { + const rows: u32 = @intCast(try std.fmt.parseInt(u32, key, 10)); + const e = chains.get(key).?.object; + const p = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "p"))); + const a = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "a"))); + const g = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "g"))); + const conv0 = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "conv0"))); + const state0 = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state0"))); + const want_out = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "out")); + const want_state = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state_out")); + const out = try r.zeros(want_out.len); + const st_out = try r.zeros(want_state.len); + + const sc = kda.Scratch{ + .k = try r.zeros(@as(usize, rows) * q.per * 4), + .v = try r.zeros(@as(usize, rows) * q.per * 2), + .g = try r.zeros(@as(usize, rows) * q.per * 4), + .b = try r.zeros(@as(usize, rows) * @as(usize, q.h) * 4), + .q = try r.zeros(@as(usize, rows) * q.per * 4), + .y = try r.zeros(@as(usize, rows) * q.per * 2), + }; + if (rows >= 64) { + try k.chainWide(&r.s, q.h, p, q.p_stride, @as(i32, @intCast(q.b_off)), a, @as(i64, @intCast(q.per)), g, @as(i64, @intCast(q.per)), conv0, conv_w, state0, a_log, dt_bias, norm_w, q.eps, q.lower, rows, out, st_out, sc); + } else { + try k.chain(&r.s, q.h, p, q.p_stride, @as(i32, @intCast(q.b_off)), a, @as(i64, @intCast(q.per)), g, @as(i64, @intCast(q.per)), conv0, conv_w, state0, a_log, dt_bias, norm_w, q.eps, q.lower, rows, out, st_out, sc); + } + const got_out = try r.back(out, want_out.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.kda.chain r{d} out", .{rows}), got_out, want_out); + const got_state = try r.back(st_out, want_state.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.kda.chain r{d} state", .{rows}), got_state, want_state); + + if (rows == 8) { + const st2 = try r.zeros(want_state.len); + try k.replay(&r.s, q.h, state0, .{ .k = sc.k, .v = sc.v, .g = sc.g, .b = sc.b }, 8, st2); + const got_replay = try r.back(st2, want_state.len); + try check.sameBytes("glm.kda.chain replay state", got_replay, want_state); + const sk8 = e.get("scratch").?.object; + inline for (.{ "k", "v", "g", "b" }) |sk_name| { + const want_sc = try ref.readFile(r.a, r.gpu.io, dir, ref.str(sk8, sk_name)); + const got_sc = try r.back(@field(sc, sk_name), want_sc.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.kda.chain scratch {s}", .{sk_name}), got_sc, want_sc); + } + try rowInvariant(r, dir, &k, q, p, a, g, conv0, conv_w, state0, a_log, dt_bias, norm_w, want_out, want_state); + } + + const out64 = try r.a.alloc(f64, @as(usize, rows) * q.per); + const st64 = try r.a.alloc(f64, q.per * ref.DK); + try ref.serialRef(r.a, ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "q"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "k"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "v"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "lg"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "beta"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state0"))), ref.bf16s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "g"))), ref.bf16s(try ref.readFile(r.a, r.gpu.io, dir, "kda_norm.bin")), q.eps, q.h, rows, out64, st64); + try t.close(try std.fmt.allocPrint(r.a, "glm.kda.chain r{d} out f64", .{rows}), got_out, out64, false, 1.0 / 128.0, 2e-3); + try t.close(try std.fmt.allocPrint(r.a, "glm.kda.chain r{d} state f64", .{rows}), got_state, st64, true, 2e-2, 2e-3); + } + std.debug.print("RESULT glm.kda.chain: f64 worst {d:.4} of tolerance\n", .{t.worst}); + check.pass("glm.kda.chain", .{}); +} + +/// The same 8 rows one at a time, conv windows shifted in between: every row's bits and the final state. +fn rowInvariant(r: *Rig, dir: []const u8, k: *kda.Kda, q: Params, p: u64, a: u64, g: u64, conv0: u64, conv_w: u64, state0: u64, a_log: u64, dt_bias: u64, norm_w: u64, want_out: []const u8, want_state: []const u8) !void { + const row_bytes = want_out.len / 8; + var s_in = try r.zeros(want_state.len); + var s_out = try r.zeros(want_state.len); + const buf: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = s_in, .len = want_state.len }; + try buf.copyFrom(0, state0, want_state.len, r.s.handle); + const conv = try r.zeros(9 * q.per * 2); + const cbuf: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = conv, .len = 9 * q.per * 2 }; + try cbuf.copyFrom(0, conv0, 9 * q.per * 2, r.s.handle); + for (0..8) |i| { + const out_i = try r.zeros(row_bytes); + try k.chain(&r.s, q.h, p + @as(u64, @intCast(@as(i64, @intCast(i)) * q.p_stride * 2)), q.p_stride, @as(i32, @intCast(q.b_off)), a + @as(u64, @intCast(i * q.per * 2)), @as(i64, @intCast(q.per)), g + @as(u64, @intCast(i * q.per * 2)), @as(i64, @intCast(q.per)), conv, conv_w, s_in, a_log, dt_bias, norm_w, q.eps, q.lower, 1, out_i, s_out, .{}); + const got = try r.back(out_i, row_bytes); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.kda.chain row {d} (one row at a time)", .{i}), got, want_out[i * row_bytes ..][0..row_bytes]); + const t = s_in; + s_in = s_out; + s_out = t; + if (i < 7) try k.convShift(&r.s, conv, p + @as(u64, @intCast(@as(i64, @intCast(i)) * q.p_stride * 2)), 1, 1, 9 * @as(i64, @intCast(q.per)), q.p_stride, q.p_stride, @intCast(q.per * 3)); + } + const rw_state = try ref.readFile(r.a, r.gpu.io, dir, "chain_r8_rowwise_state.bin"); + const got_state = try r.back(s_in, rw_state.len); + try check.sameBytes("glm.kda.chain rowwise state", got_state, rw_state); +} + +fn chunked(r: *Rig, state: *Run, dir: []const u8) !void { + _ = state; + if (std.mem.eql(u8, dir, "none")) { + skip("glm.kda.chunked"); + return; + } + var man = try ref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + const o = man.kda() orelse { + skip("glm.kda.chunked"); + return; + }; + const q = try Params.read(o); + var c = try chunked_mod.Chunk.load(r.gpu.d); + defer c.unload(); + + const a_log = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_a_log.bin")); + const dt_bias = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_dt_bias.bin")); + const norm_w = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_norm.bin")); + const conv_w = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, "kda_conv.bin")); + + const e = o.get("chunked").?.object; + const rows: u32 = @intCast(ref.int(e, "rows")); + const p_bytes = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "p")); + const p = try r.dev(p_bytes); + const a = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "a"))); + const g = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "g"))); + const conv0 = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "conv0"))); + const state0 = try r.dev(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state0"))); + const want_out = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "out")); + const want_state = try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state_out")); + const out = try r.zeros(want_out.len); + const st = try r.zeros(want_state.len); + try c.chain(&r.s, q.h, p, q.p_stride, @as(i32, @intCast(q.b_off)), a, @as(i64, @intCast(q.per)), g, @as(i64, @intCast(q.per)), conv0, conv_w, state0, a_log, dt_bias, norm_w, q.eps, q.lower, rows, 0, out, st); + const got_out = try r.back(out, want_out.len); + try check.sameBytes("glm.kda.chunked out", got_out, want_out); + const got_state = try r.back(st, want_state.len); + try check.sameBytes("glm.kda.chunked state", got_state, want_state); + + const out64 = try r.a.alloc(f64, @as(usize, rows) * q.per); + const st64 = try r.a.alloc(f64, q.per * ref.DK); + try ref.chunkedRef(r.a, ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "q"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "k"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "v"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "lg"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "beta"))), ref.f32s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "state0"))), ref.bf16s(try ref.readFile(r.a, r.gpu.io, dir, ref.str(e, "g"))), ref.bf16s(try ref.readFile(r.a, r.gpu.io, dir, "kda_norm.bin")), q.eps, q.h, rows, 0, out64, st64); + var t = T{}; + try t.close("glm.kda.chunked out f64", got_out, out64, false, 1.0 / 128.0, 2e-3); + try t.close("glm.kda.chunked state f64", got_state, st64, true, 2e-2, 2e-3); + + // the cuts: segment 1's state, then rows n.. from it through the shifted conv window reproduce the whole run + const splits = e.get("splits").?.object; + for ([_][]const u8{ "32", "96", "1024", "2016" }) |key| { + const n: u32 = @intCast(try std.fmt.parseInt(u32, key, 10)); + const want_st1 = try ref.readFile(r.a, r.gpu.io, dir, ref.str(splits.get(key).?.object, "st1")); + const out1 = try r.zeros(@as(usize, n) * q.per * 2); + const st1 = try r.zeros(want_state.len); + try c.chain(&r.s, q.h, p, q.p_stride, @as(i32, @intCast(q.b_off)), a, @as(i64, @intCast(q.per)), g, @as(i64, @intCast(q.per)), conv0, conv_w, state0, a_log, dt_bias, norm_w, q.eps, q.lower, n, 0, out1, st1); + const got_st1 = try r.back(st1, want_st1.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.kda.chunked cut {d} state", .{n}), got_st1, want_st1); + + const stride: usize = @intCast(q.p_stride); + const ch: usize = q.per * 3; + const conv_mid = try r.a.alloc(u8, 3 * ch * 2); + for (0..3) |ri| @memcpy(conv_mid[ri * ch * 2 ..][0 .. ch * 2], p_bytes[(n - 3 + ri) * stride * 2 ..][0 .. ch * 2]); + const out2 = try r.zeros((@as(usize, rows) - n) * q.per * 2); + const st2 = try r.zeros(want_state.len); + try c.chain(&r.s, q.h, p + @as(u64, n) * @as(u64, stride) * 2, q.p_stride, q.b_off, a + @as(u64, n) * @as(u64, q.per) * 2, @as(i64, @intCast(q.per)), g + @as(u64, n) * @as(u64, q.per) * 2, @as(i64, @intCast(q.per)), try r.dev(conv_mid), conv_w, st1, a_log, dt_bias, norm_w, q.eps, q.lower, rows - n, 0, out2, st2); + const got_out1 = try r.back(out1, @as(usize, n) * q.per * 2); + const got_out2 = try r.back(out2, (@as(usize, rows) - n) * q.per * 2); + const cat = try r.a.alloc(u8, got_out1.len + got_out2.len); + @memcpy(cat[0..got_out1.len], got_out1); + @memcpy(cat[got_out1.len..], got_out2); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.kda.chunked cut {d} rows", .{n}), cat, want_out); + const got_st2 = try r.back(st2, want_state.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.kda.chunked cut {d} final state", .{n}), got_st2, want_state); + } + std.debug.print("RESULT glm.kda.chunked: f64 worst {d:.4} of tolerance\n", .{t.worst}); + check.pass("glm.kda.chunked", .{}); +} diff --git a/zig/tests/cuda/glm_latent_ref.zig b/zig/tests/cuda/glm_latent_ref.zig new file mode 100644 index 0000000000..c32aeca6ee --- /dev/null +++ b/zig/tests/cuda/glm_latent_ref.zig @@ -0,0 +1,101 @@ +//! Float64 references for the GLM-5.3-Flash latent attention checks (absorb/expand dots and dense attention). + +const std = @import("std"); +const kref = @import("glm_kda_ref.zig"); +const sref = @import("glm_sparse_ref.zig"); + +pub const fp8c = sref.fp8e4m3; +pub const bf = kref.bf; + +/// A dequantized fp8 cache value as the engine sees it: code times the f32 scale, rounded through bf16. +fn cacheVal(cache: []const u8, scales: ?[]const f32, fp8: bool, row: usize, i: usize) f64 { + if (!fp8) return bf(u16At(cache, row * 512 + i)); + return sref.bf16rn(fp8c(cache[row * 512 + i]) * @as(f64, scales.?[row])); +} + +fn u16At(bytes: []const u8, i: usize) u16 { + return std.mem.readInt(u16, bytes[2 * i ..][0..2], .little); +} + +/// qa[r,h,n] = sum_k q[r,h,k] * wk[h,k,n] over bf16 operands. +pub fn absorbRef(a: std.mem.Allocator, q: []const u16, wk: []const u16, rows: usize, heads: usize, d: usize, lw: usize) ![]f64 { + const out = try a.alloc(f64, rows * heads * lw); + var r: usize = 0; + while (r < rows) : (r += 1) { + var h: usize = 0; + while (h < heads) : (h += 1) { + var n: usize = 0; + while (n < lw) : (n += 1) { + var sum: f64 = 0; + var k: usize = 0; + while (k < d) : (k += 1) sum += bf(q[(r * heads + h) * d + k]) * bf(wk[(h * d + k) * lw + n]); + out[(r * heads + h) * lw + n] = sum; + } + } + } + return out; +} + +/// out[r,h,n] = sum_k ol[r,h,k] * wv[h,n,k] over bf16 operands. +pub fn expandRef(a: std.mem.Allocator, ol: []const u16, wv: []const u16, rows: usize, heads: usize, lw: usize, v: usize) ![]f64 { + const out = try a.alloc(f64, rows * heads * v); + var r: usize = 0; + while (r < rows) : (r += 1) { + var h: usize = 0; + while (h < heads) : (h += 1) { + var n: usize = 0; + while (n < v) : (n += 1) { + var sum: f64 = 0; + var k: usize = 0; + while (k < lw) : (k += 1) sum += bf(ol[(r * heads + h) * lw + k]) * bf(wv[(h * v + n) * lw + k]); + out[(r * heads + h) * v + n] = sum; + } + } + } + return out; +} + +/// Dense latent attention: row r attends keys 0..pos+r (the causal window) over the fp8/bf16 cache. +pub fn attnRef( + a: std.mem.Allocator, + qa: []const u16, + cache: []const u8, + scales: ?[]const f32, + pos: i64, + rows: usize, + heads: usize, + lw: usize, + scale: f64, + fp8: bool, + cap_rows: usize, +) ![]f64 { + const nrows = @min(rows, cap_rows); + const out = try a.alloc(f64, nrows * heads * lw); + var r: usize = 0; + while (r < nrows) : (r += 1) { + const keys: usize = @intCast(pos + @as(i64, @intCast(r)) + 1); + var h: usize = 0; + while (h < heads) : (h += 1) { + var m: f64 = -std.math.inf(f64); + var l: f64 = 0; + const o = try a.alloc(f64, lw); + defer a.free(o); + @memset(o, 0); + const qrow = qa[(r * heads + h) * lw ..][0..lw]; + var j: usize = 0; + while (j < keys) : (j += 1) { + var dot: f64 = 0; + for (0..lw) |i| dot += bf(qrow[i]) * cacheVal(cache, scales, fp8, j, i); + const s = dot * scale; + const nm = @max(m, s); + const al = if (m == -std.math.inf(f64)) 0 else @exp(m - nm); + const p = @exp(s - nm); + for (0..lw) |i| o[i] = o[i] * al + p * cacheVal(cache, scales, fp8, j, i); + l = l * al + p; + m = nm; + } + for (0..lw) |i| out[(r * heads + h) * lw + i] = o[i] / l; + } + } + return out; +} diff --git a/zig/tests/cuda/glm_latent_tests.zig b/zig/tests/cuda/glm_latent_tests.zig new file mode 100644 index 0000000000..b5153fdeef --- /dev/null +++ b/zig/tests/cuda/glm_latent_tests.zig @@ -0,0 +1,264 @@ +//! GLM-5.3-Flash latent attention on CUDA: cache write, absorb/expand, dense and sparse chunks vs float64 and Python. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const sref = @import("glm_sparse_ref.zig"); +const lref = @import("glm_latent_ref.zig"); +const ktests = @import("glm_kda_tests.zig"); +const latent = cuda.glm_latent; + +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const T = ktests.T; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + write(r, state, dir) catch |e| state.fail(e, "glm.dsa.latent"); + absorb(r, state, dir) catch |e| state.fail(e, "glm.dsa.absorb"); + dense(r, state, dir) catch |e| state.fail(e, "glm.dsa.dense"); + sparseAttn(r, state, dir) catch |e| state.fail(e, "glm.dsa.sparse"); +} + +fn skip(id: []const u8) void { + std.debug.print("SKIP {s}\n", .{id}); +} + +fn u16s(bytes: []const u8) []const u16 { + return @alignCast(std.mem.bytesAsSlice(u16, bytes)); +} +fn f32s(bytes: []const u8) []const f32 { + return @alignCast(std.mem.bytesAsSlice(f32, bytes)); +} +fn i32s(bytes: []const u8) []const i32 { + return @alignCast(std.mem.bytesAsSlice(i32, bytes)); +} + +/// A device-side copy of the scalar window position the kernels read. +fn posDev(r: *Rig, pos: i64) !u64 { + const bytes = try r.a.alloc(u8, 4); + std.mem.writeInt(i32, bytes[0..4], @intCast(pos), .little); + return r.dev(bytes); +} + +/// The cache write: fp8 codes with a per-token amax/448 scale, and a plain bf16 copy. +fn write(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + _ = state; + skip("glm.dsa.latent"); + return; + } + var man = try sref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + const fam = man.root().get("latent") orelse { + skip("glm.dsa.latent"); + return; + }; + var k = try latent.Latent.load(r.gpu.d); + defer k.unload(); + var t = T{}; + var it = fam.object.get("write").?.object.iterator(); + while (it.next()) |kv| { + const e = kv.value_ptr.object; + const tag = kv.key_ptr.*; + const rows: u32 = @intCast(ref.int(e, "R")); + const pos: i64 = ref.int(e, "pos"); + const fp8 = e.get("fp8").?.bool; + const lat = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "lat")); + const want_cache = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "cache")); + const cache = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "cache_pre"))); + const pos_d = try posDev(r, pos); + var want_s: []f64 = &.{}; + if (fp8) { + const scales = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "scales_pre"))); + try k.latWrite(&r.s, try r.dev(lat), 512, cache, pos_d, scales, rows, true); + const want_scales = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "scales")); + const got_scales = try r.back(scales, want_scales.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.latent {s} scales", .{tag}), got_scales, want_scales); + // float64: the scale is the row's amax / 448 (the engine's div.full rounding aside) + want_s = try r.a.alloc(f64, rows); + for (0..rows) |row| { + var amax: f64 = 0; + for (0..512) |i| amax = @max(amax, @abs(lref.bf(u16s(lat)[row * 512 + i]))); + want_s[row] = if (amax > 0) amax / 448.0 else 1.0; + } + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.latent {s} scales f64", .{tag}), got_scales[@as(usize, @intCast(pos)) * 4 ..], want_s, true, 1.0e-6, 1.0e-33); + } else { + try k.latWrite(&r.s, try r.dev(lat), 512, cache, pos_d, 0, rows, false); + } + const got = try r.back(cache, want_cache.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.latent {s} cache", .{tag}), got, want_cache); + // float64: the fp8 codes decode to lat / s; the bf16 copy widens exactly + for (0..rows) |row| { + const want_row = try r.a.alloc(f64, 512); + for (0..512) |i| { + const x = lref.bf(u16s(lat)[row * 512 + i]); + want_row[i] = if (fp8) x / want_s[row] else x; + } + const tol: f64 = if (fp8) 1.0 / 8.0 else 1.0 / 128.0; + const off = (@as(usize, @intCast(pos)) + row) * (if (fp8) @as(usize, 512) else 1024); + var gbuf: []u8 = got[off..]; + if (fp8) { + // close reads bf16: widen each stored code's exact fp8 value to bf16 bytes + gbuf = try r.a.alloc(u8, 1024); + for (0..512) |i| std.mem.writeInt(u16, gbuf[2 * i ..][0..2], quant_ref_tests.helpers.bf16Of(@floatCast(lref.fp8c(got[off + i]))), .little); + } else gbuf = gbuf[0..1024]; + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.latent {s} r{d} f64", .{ tag, row }), gbuf, want_row, false, tol, 0.01); + } + } + std.debug.print("RESULT glm.dsa.latent: f64 worst {d:.4} of tolerance\n", .{t.worst}); + check.pass("glm.dsa.latent", .{}); +} + +/// The kv_b absorb and expand projections over the layer's bf16 weights. +fn absorb(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + _ = state; + skip("glm.dsa.absorb"); + return; + } + var man = try sref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + const fam = man.root().get("latent") orelse { + skip("glm.dsa.absorb"); + return; + }; + var k = try latent.Latent.load(r.gpu.d); + defer k.unload(); + const wk = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, "wk.bin")); + const wv = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, "wv.bin")); + var t = T{}; + var it = fam.object.get("absorb").?.object.iterator(); + while (it.next()) |kv| { + const e = kv.value_ptr.object; + const tag = kv.key_ptr.*; + const rows: u32 = @intCast(ref.int(e, "R")); + const q = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "q")); + const want_qa = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "qa")); + const ol = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "ol")); + const want_out = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "out")); + const qa = try r.zeros(want_qa.len); + const out = try r.zeros(want_out.len); + try k.absorbQ(&r.s, try r.dev(q), wk, qa, rows); + try k.expandV(&r.s, try r.dev(ol), wv, out, rows); + const got_qa = try r.back(qa, want_qa.len); + const got_out = try r.back(out, want_out.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.absorb {s} qa", .{tag}), got_qa, want_qa); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.absorb {s} out", .{tag}), got_out, want_out); + if (rows <= 64) { + const want64 = try lref.absorbRef(r.a, u16s(q), u16s(try sref.readFile(r.a, r.gpu.io, dir, "wk.bin")), rows, 32, 256, 512); + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.absorb {s} qa f64", .{tag}), got_qa, want64, false, 1.0 / 128.0, 2.0e-3); + const want_o = try lref.expandRef(r.a, u16s(ol), u16s(try sref.readFile(r.a, r.gpu.io, dir, "wv.bin")), rows, 32, 512, 256); + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.absorb {s} out f64", .{tag}), got_out, want_o, false, 1.0 / 128.0, 2.0e-3); + } + } + std.debug.print("RESULT glm.dsa.absorb: f64 worst {d:.4} of tolerance\n", .{t.worst}); + check.pass("glm.dsa.absorb", .{}); +} + +/// Dense attention: 512-key chunks with the flash update, then the cross-chunk merge. +fn dense(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + _ = state; + skip("glm.dsa.dense"); + return; + } + var man = try sref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + const fam = man.root().get("latent") orelse { + skip("glm.dsa.dense"); + return; + }; + const scale = ref.float(fam.object.get("params").?.object, "scale"); + var k = try latent.Latent.load(r.gpu.d); + defer k.unload(); + var t = T{}; + var it = fam.object.get("dense").?.object.iterator(); + while (it.next()) |kv| { + const e = kv.value_ptr.object; + const tag = kv.key_ptr.*; + const rows: u32 = @intCast(ref.int(e, "R")); + const pos: i64 = ref.int(e, "pos"); + const nch: u32 = @intCast(ref.int(e, "nch")); + const fp8 = e.get("fp8").?.bool; + const qa = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "qa")); + const cache_pre = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "cache_pre")); + const want_out = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "out")); + const qa_dev = try r.dev(qa); + const cache_dev = try r.dev(cache_pre); + const pos_d = try posDev(r, pos); + var scales_dev: u64 = 0; + if (fp8) scales_dev = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "scales_pre"))); + const po = try r.zeros(@as(usize, nch) * rows * 32 * 512 * 4); + const pm = try r.zeros(@as(usize, nch) * rows * 32 * 4); + const pl = try r.zeros(@as(usize, nch) * rows * 32 * 4); + const out = try r.zeros(want_out.len); + try k.attnChunks(&r.s, qa_dev, cache_dev, scales_dev, pos_d, 0, 0, po, pm, pl, rows, 0, nch, fp8, false, @floatCast(scale)); + try k.merge(&r.s, po, pm, pl, out, 0, rows, nch, false); + const got = try r.back(out, want_out.len); + if (rows <= 8) { + const want64 = try lref.attnRef(r.a, u16s(qa), cache_pre, if (fp8) f32s(try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "scales_pre"))) else null, pos, rows, 32, 512, scale, fp8, 8); + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.dense {s} out f64", .{tag}), got, want64, false, 1.0 / 64.0, 2.0e-3); + } + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.dense {s} out", .{tag}), got, want_out); + } + std.debug.print("RESULT glm.dsa.dense: f64 worst {d:.4} of tolerance\n", .{t.worst}); + check.pass("glm.dsa.dense", .{}); +} + +/// The sparse path: each row attends its selected tokens (count 0 rows keep their input). +fn sparseAttn(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + _ = state; + skip("glm.dsa.sparse"); + return; + } + var man = try sref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + const sp = man.root().get("sparse") orelse { + skip("glm.dsa.sparse"); + return; + }; + const scale = ref.float(sp.object.get("params").?.object, "scale"); + var k = try latent.Latent.load(r.gpu.d); + defer k.unload(); + var t = T{}; + var it = sp.object.get("cases").?.object.iterator(); + while (it.next()) |kv| { + const e = kv.value_ptr.object; + const tag = kv.key_ptr.*; + const rows: u32 = @intCast(ref.int(e, "R")); + const tokens = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "tokens")); + const counts = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "counts")); + const w = i32s(tokens).len / rows; + const nch: u32 = @intCast((w + 511) / 512); + const qa = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "qa")); + const cache = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "cache")); + const scales = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "scales")); + const want_out = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "out")); + const qa_dev = try r.dev(qa); + const out = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "out_in"))); + const tok_dev = try r.dev(tokens); + const cnt_dev = try r.dev(counts); + const scales_dev = try r.dev(scales); + const po = try r.zeros(@as(usize, nch) * rows * 32 * 512 * 4); + const pm = try r.zeros(@as(usize, nch) * rows * 32 * 4); + const pl = try r.zeros(@as(usize, nch) * rows * 32 * 4); + try k.attnChunks(&r.s, qa_dev, try r.dev(cache), scales_dev, 0, tok_dev, cnt_dev, po, pm, pl, rows, @intCast(w), nch, true, true, @floatCast(scale)); + try k.merge(&r.s, po, pm, pl, out, cnt_dev, rows, nch, true); + const got = try r.back(out, want_out.len); + if (rows <= 8) { + const cnts = i32s(counts); + const want64 = try sref.sparseRef(r.a, u16s(qa), cache, f32s(scales), i32s(tokens), cnts, 32, 512, scale); + for (0..rows) |row| { + if (cnts[row] == 0) continue; + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.sparse {s} r{d} f64", .{ tag, row }), got[row * 32 * 512 * 2 ..][0 .. 32 * 512 * 2], want64[row * 32 * 512 ..][0 .. 32 * 512], false, 1.0 / 64.0, 2.0e-3); + } + } + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.sparse {s} out", .{tag}), got, want_out); + } + std.debug.print("RESULT glm.dsa.sparse: f64 worst {d:.4} of tolerance\n", .{t.worst}); + check.pass("glm.dsa.sparse", .{}); +} diff --git a/zig/tests/cuda/glm_layer_dsa_ref.zig b/zig/tests/cuda/glm_layer_dsa_ref.zig new file mode 100644 index 0000000000..bb0311d50d --- /dev/null +++ b/zig/tests/cuda/glm_layer_dsa_ref.zig @@ -0,0 +1,38 @@ +//! glm.layer.dsa's float64 part: the layer's HC glue and MoE (glm_layer_ref), the DSA block's chain (glm_dsa_ref). +const std = @import("std"); +const cuda = @import("cuda"); +const ref = @import("glm_kda_ref.zig"); +const lref = @import("glm_layer_ref.zig"); +const dref = @import("glm_dsa_ref.zig"); +const lay = cuda.glm_layer; + +const Obj = std.json.ObjectMap; +const id = "glm.layer.dsa"; + +/// The whole layer in f64, stage by stage from each stage's inputs; returns the worst fraction of tolerance. +pub fn layerF64(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, c: lay.Config, h: *const lref.Host, dh: *const dref.Host, p: dref.Params, g: dref.Got, tag: []const u8, e: Obj) !f64 { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const w: lref.Win = .{ .a = a, .io = io, .dir = dir, .e = e }; + var rp: lref.Report = .{ .a = a, .tag = tag, .id = id }; + const rows: usize = @intCast(ref.int(e, "rows")); + const pos: usize = @intCast(ref.int(e, "pos")); + const d: usize = c.d; + // the attention site: hc_pre, the DSA block, hc_post + const x0 = try w.bf(null, "x0"); + const n0 = try w.bf("s0", "normed"); + try lref.hcPreStage(&rp, a, "attn hc_pre", c, h.attn_hc, h.in_norm, x0, rows, n0, try w.f("s0", "post"), try w.f("s0", "comb")); + const qd = try dref.chain(a, dh, p, g, tag, rows, pos); + rp.worst = @max(rp.worst, qd); + rp.line = try std.fmt.allocPrint(a, "{s} dsa block {d:.3};", .{ rp.line, qd }); + const x1 = try w.bf("s1", "x"); + try lref.hcPostStage(&rp, "attn hc_post", d, x0, try w.f("s0", "part"), try w.f("s0", "peer"), try w.f("s0", "post"), try w.f("s0", "comb"), rows, x1); + // the MoE site: hc_pre, router, select, routed and shared experts, combine, hc_post + const n1 = try w.bf("s1", "normed"); + try lref.hcPreStage(&rp, a, "ffn hc_pre", c, h.ffn_hc, h.post_norm, x1, rows, n1, try w.f("s1", "post"), try w.f("s1", "comb")); + try lref.moeStage(&rp, a, w, c, h, n1, rows); + try lref.hcPostStage(&rp, "ffn hc_post", d, x1, try w.f("s1", "part"), try w.f("s1", "peer"), try w.f("s1", "post"), try w.f("s1", "comb"), rows, try w.bf(null, "x")); + std.debug.print("RESULT {s} {s} f64 per stage (fraction of tolerance):{s}\n", .{ id, tag, rp.line }); + return rp.worst; +} diff --git a/zig/tests/cuda/glm_layer_dsa_tests.zig b/zig/tests/cuda/glm_layer_dsa_tests.zig new file mode 100644 index 0000000000..9c0f459e01 --- /dev/null +++ b/zig/tests/cuda/glm_layer_dsa_tests.zig @@ -0,0 +1,204 @@ +//! glm.layer.dsa: a real RED-SNOW DSA MoE layer at TP=2 rank 0 vs the Python engine's layer_forward, bit for bit. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const lref = @import("glm_layer_ref.zig"); +const dref = @import("glm_dsa_ref.zig"); +const ldref = @import("glm_layer_dsa_ref.zig"); +const lt = @import("glm_layer_tests.zig"); +const lay = cuda.glm_layer; +const dsa = cuda.glm_dsa; + +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const Obj = std.json.ObjectMap; +const Stage = lt.Stage; +const id = "glm.layer.dsa"; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + dsaLayer(r, dir) catch |e| state.fail(e, id); +} + +fn skip() void { + std.debug.print("SKIP {s}\n", .{id}); +} + +/// The caches every window starts from (kc fp8, scales, ik, ig, pk), as the real prefill of this layer left them. +const names = [_][]const u8{ "kc", "scales", "ik", "ig", "pk" }; +const row_bytes = [_]usize{ dsa.kv_lora, 4, dsa.index_dim * 2, dsa.index_dim * 2, dsa.index_dim * 2 }; + +/// The fixed inputs of every window. +const Layer = struct { + c: lay.Config, + p64: dref.Params, + cap: usize, + host: lref.Host, + dh: dref.Host, + w: lay.LayerW, + pre: [names.len][]u8, +}; + +fn dsaLayer(r: *Rig, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) return skip(); + var man = try ref.Manifest.open(r.gpu.gpa, r.gpu.io, dir); + defer man.deinit(); + const o = (man.root().get("layer_dsa") orelse return skip()).object; + const p = o.get("params").?.object; + var c = lt.config(p); + c.dsa = .{ + .eps = c.eps, + .index_eps = @floatCast(ref.float(p, "index_eps")), + .scale = @floatCast(ref.float(p, "scale")), + .index_scale = @floatCast(ref.float(p, "index_scale")), + .wscale = @floatCast(ref.float(p, "wscale")), + }; + const dc = c.dsa.?; + const wo = o.get("weights").?.object; + var l: Layer = .{ + .c = c, + .p64 = .{ .eps = dc.eps, .index_eps = dc.index_eps, .scale = dc.scale, .index_scale = dc.index_scale, .wscale = dc.wscale }, + .cap = @intCast(ref.int(p, "cap")), + .host = try lref.Host.read(r.a, r.gpu.io, dir, wo), + .dh = try dref.Host.read(r.a, r.gpu.io, dir, wo.get("dsa").?.object), + .w = undefined, + .pre = undefined, + }; + const h = &l.host; + const dh = &l.dh; + l.w = .{ + .attn_hc = .{ .fn_w = try r.dev(h.attn_hc[0]), .base = try r.dev(h.attn_hc[1]), .scale = try r.dev(h.attn_hc[2]) }, + .ffn_hc = .{ .fn_w = try r.dev(h.ffn_hc[0]), .base = try r.dev(h.ffn_hc[1]), .scale = try r.dev(h.ffn_hc[2]) }, + .in_norm = try r.dev(h.in_norm), + .post_norm = try r.dev(h.post_norm), + .mixer = .{ .dsa = .{ + .q_a = try lt.linW(r, dh.q_a), + .kv_a = try lt.linW(r, dh.kv_a), + .q_b = try lt.linW(r, dh.q_b), + .o = try lt.linW(r, dh.o), + .q_norm = try r.dev(dh.q_norm), + .kv_norm = try r.dev(dh.kv_norm), + .wk = try r.dev(dh.wk), + .wv = try r.dev(dh.wv), + .index = .{ .kw = try lt.b16W(r, dh.kw), .qb = try lt.b16W(r, dh.qb), .gate = .{ .w = try r.dev(dh.gate), .n = dsa.index_dim, .k = dsa.hidden }, .ln_w = try r.dev(dh.ln_w), .ln_b = try r.dev(dh.ln_b), .ape = try r.dev(dh.ape) }, + } }, + .ffn = .{ .moe = try lt.moeW(r, h, c) }, + }; + for (names, 0..) |n, i| l.pre[i] = try lt.file(r, dir, ref.str(o.get("pre").?.object, n)); + var k = try lay.Kernels.load(r.gpu.d); + defer k.unload(); + var worst: f64 = 0; + const cases = o.get("cases").?.object; + var it = cases.iterator(); + while (it.next()) |kv| { + const t0 = check.now(r.gpu.io); + // a window's host bytes live in its own arena + var arena = std.heap.ArenaAllocator.init(r.gpu.gpa); + defer arena.deinit(); + const saved = r.a; + r.a = arena.allocator(); + defer r.a = saved; + const q = try window(r, &k, &l, dir, kv.key_ptr.*, kv.value_ptr.object); + worst = @max(worst, q); + const dt = @as(f64, @floatFromInt(check.now(r.gpu.io) - t0)) / 1e9; + std.debug.print("RESULT {s} {s}: bit-exact, f64 worst {d:.4} of tolerance ({d:.1} s)\n", .{ id, kv.key_ptr.*, q, dt }); + } + std.debug.print("RESULT {s}: {d} windows bit-exact, f64 worst {d:.4} of tolerance\n", .{ id, cases.count(), worst }); + check.pass("{s}", .{id}); +} + +/// Device bytes back to the host 8-aligned (the f64 chain reads them as u16/f32/i32 slices). +fn back(r: *Rig, ptr: u64, len: usize) ![]u8 { + const out = try r.a.alignedAlloc(u8, .@"8", len); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = len }; + try b.download(0, out); + return out; +} + +/// The fixture's bytes 8-aligned. +fn file8(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + const raw = try lt.file(r, dir, name); + const out = try r.a.alignedAlloc(u8, .@"8", raw.len); + @memcpy(out, raw); + return out; +} + +/// One window from the pre-state: forward with every stage's bits checked at the gathers, x and the caches after. +fn window(r: *Rig, k: *lay.Kernels, l: *const Layer, dir: []const u8, tag: []const u8, e: Obj) !f64 { + const rows: usize = @intCast(ref.int(e, "rows")); + const pos: usize = @intCast(ref.int(e, "pos")); + const prefill = e.get("prefill").?.bool; + const w = &l.w; + var sc = try lay.Scratch.init(r.gpu.d, r.s, l.c, w, rows, 2); + defer sc.deinit(); + var dsc = try dsa.Scratch.init(r.gpu.d, &w.mixer.dsa, @intCast(ref.int(e, "buf_rows")), prefill, l.cap); + defer dsc.deinit(); + try dsc.buf.fill8(0xff, r.s.handle); // torch.empty's scratch: a read of an unwritten word shows + var caches: [names.len]u64 = undefined; + for (l.pre, 0..) |bytes, i| caches[i] = try r.dev(bytes); + var pos_b: [4]u8 = undefined; + std.mem.writeInt(i32, &pos_b, @intCast(pos), .little); + const st: dsa.DsaState = .{ .lc = caches[0], .scales = caches[1], .ik = caches[2], .ig = caches[3], .pk = caches[4], .pos_dev = try r.dev(&pos_b), .cap = l.cap }; + const x = try lt.up(r, dir, ref.str(e, "x0")); + // the DSA block's stages in its order: the first that differs names the kernel + const attn = [_]Stage{ + .{ "dqa", dsc.dqa }, .{ "dkv", dsc.dkv }, .{ "qr", dsc.qr }, .{ "lat", dsc.lat }, + .{ "q", dsc.q }, .{ "ikr", dsc.ikr }, .{ "igr", dsc.igr }, .{ "qa", dsc.qa }, + .{ "qi", dsc.qi }, .{ "tokens", dsc.tokens }, .{ "counts", dsc.counts }, .{ "ol", dsc.ol }, + .{ "vn", dsc.vn }, + }; + var ctx: lt.Ctx = .{ .r = r, .id = id, .dir = dir, .tag = tag, .e = e, .sc = &sc, .attn = &attn, .x = x, .d = l.c.d }; + const g: lay.Gather = .{ .ctx = &ctx, .world = 2, .run = lt.Ctx.gather }; + try lay.forward(k, &r.s, l.c, w, .{ .dsa = .{ .st = st, .sc = &dsc, .pos = pos } }, &sc, x, rows, g); + try lt.same(r, dir, try std.fmt.allocPrint(r.a, "{s} {s} x", .{ id, tag }), x, ref.str(e, "x")); + // the caches: the window's rows as Python wrote them, every other row as the pre-state left it + var post: [names.len][]const u8 = undefined; + for (names, 0..) |n, i| { + const co = e.get(n).?.object; + const lo: usize = @intCast(ref.int(co, "lo")); + const hi: usize = @intCast(ref.int(co, "hi")); + const want = try r.a.dupe(u8, l.pre[i]); + const win = try lt.file(r, dir, ref.str(co, "file")); + try check.expect(win.len == (hi - lo) * row_bytes[i], "{s} {s} {s}: window bytes", .{ id, tag, n }); + @memcpy(want[lo * row_bytes[i] ..][0..win.len], win); + post[i] = try back(r, caches[i], want.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "{s} {s} {s}", .{ id, tag, n }), post[i], want); + } + // the oracle ran the chunk row-split too (HC_SPLIT=1, PREFILL_OVERLAP=2): own rows and every cache equal unsplit + if (prefill) try check.expect(if (e.get("split_equal")) |v| v.bool else false, "{s} {s}: no split check", .{ id, tag }); + // the f64 chain reads the block's stages (bit-equal by now) from the fixture + const s0 = e.get("s0").?.object; + var got: [attn.len + 2]?[]const u8 = undefined; + const keys = [_][]const u8{ "dqa", "dkv", "qr", "lat", "q", "ikr", "igr", "qa", "qi", "tokens", "counts", "ol", "vn", "normed", "part" }; + for (keys, 0..) |key, i| got[i] = if (s0.get(key)) |v| try file8(r, dir, v.string) else null; + const np = dsa.poolBucket(pos, rows, l.cap / 4); + const last = (rows - 1) / dsa.select_rows * dsa.select_rows; + const dg: dref.Got = .{ + .normed = got[13].?, + .dqa = got[0].?, + .dkv = got[1].?, + .qr = got[2].?, + .lat = got[3].?, + .q = got[4].?, + .ikr = got[5].?, + .igr = got[6].?, + .qa = got[7].?, + .qi = got[8], + .tokens = got[9], + .counts = got[10], + .ol = got[11].?, + .vn = got[12].?, + .part = got[14].?, + .scores = if (got[8] != null) try back(r, dsc.scores, (rows - last) * np * 4) else null, + .np = np, + .kc = post[0], + .scales = post[1], + .ik = post[2], + .ig = post[3], + .pk = post[4], + }; + return ldref.layerF64(r.gpu.gpa, r.gpu.io, dir, l.c, &l.host, &l.dh, l.p64, dg, tag, e); +} diff --git a/zig/tests/cuda/glm_layer_ref.zig b/zig/tests/cuda/glm_layer_ref.zig new file mode 100644 index 0000000000..5a1b873945 --- /dev/null +++ b/zig/tests/cuda/glm_layer_ref.zig @@ -0,0 +1,586 @@ +//! glm.layer.kda and glm.layer.mlp's float64 part: every stage of the layer recomputed in f64 from its fixture inputs. +const std = @import("std"); +const cuda = @import("cuda"); +const ref = @import("glm_kda_ref.zig"); +const exl3fmt = @import("core").exl3_format; +const lay = cuda.glm_layer; + +const Obj = std.json.ObjectMap; +const id = "glm.layer.kda"; + +/// One EXL3 linear's host bytes: words (the engine's strips layout, or stored), fp16 suh/svh, shape and width. +pub const Lin = struct { words: []const u8, strips: bool = true, suh: []const u8, svh: []const u8, k: usize, n: usize, halves: u16 }; + +/// One plain bf16 matrix [n, k]. +pub const Mat = struct { w: []const u8, n: usize, k: usize }; + +pub const KdaHost = struct { + qkv: Lin, + o: Lin, + fa: Mat, + ga: Mat, + b: Mat, + fb: Mat, + gb: Mat, + conv: []const u8, + a_log: []const u8, + dt_bias: []const u8, + norm: []const u8, + heads: u32, + fa_off: u32, + ga_off: u32, + b_off: u32, + p_stride: u32, +}; + +/// The routed experts: per projection one file of every expert's stored trellis, byte offsets, widths, scales. +pub const ExpertsHost = struct { + proj: [3][]const u8, + offs: [3][]u64, + k2: [3][]i32, + suh_g: []const u8, + suh_u: []const u8, + svh_g: []const u8, + svh_u: []const u8, + suh_d: []const u8, + svh_d: []const u8, + count: usize, + dims: usize, + width: usize, + + /// Expert e's projection j (gate, up, down) as a stored-layout linear. + pub fn lin(x: *const ExpertsHost, j: usize, e: usize) Lin { + const k = if (j == 2) x.width else x.dims; + const n = if (j == 2) x.dims else x.width; + const halves: u16 = @intCast(x.k2[j][e]); + const len = (k / 16) * (n / 16) * @as(usize, halves) * 16; + const scales = [3][2][]const u8{ .{ x.suh_g, x.svh_g }, .{ x.suh_u, x.svh_u }, .{ x.suh_d, x.svh_d } }; + return .{ .words = x.proj[j][x.offs[j][e]..][0..len], .strips = false, .suh = scales[j][0][e * k * 2 ..][0 .. k * 2], .svh = scales[j][1][e * n * 2 ..][0 .. n * 2], .k = k, .n = n, .halves = halves }; + } +}; + +/// Rank 0's layer tensors from the fixture (manifest section's "weights"), read once; ffn MoE or dense MLP. +pub const Host = struct { + attn_hc: [3][]const u8, + ffn_hc: [3][]const u8, + in_norm: []const u8, + post_norm: []const u8, + router: []const u8, + bias: []const u8, + kda: ?KdaHost, + ex: ExpertsHost, // undefined for a dense layer + sh: [3]Lin, // the shared expert's gate, up, down; a dense layer's MLP + dense: bool, + + pub fn read(a: std.mem.Allocator, io: std.Io, dir: []const u8, o: Obj) !Host { + const R = struct { + a: std.mem.Allocator, + io: std.Io, + dir: []const u8, + fn f(s: @This(), obj: Obj, key: []const u8) ![]const u8 { + return ref.readFile(s.a, s.io, s.dir, ref.str(obj, key)); + } + fn lin(s: @This(), obj: Obj) !Lin { + if (!std.mem.eql(u8, ref.str(obj, "codebook"), "mul1")) return error.GlmLayerCodebook; + return .{ .words = try s.f(obj, "words"), .suh = try s.f(obj, "suh"), .svh = try s.f(obj, "svh"), .k = @intCast(ref.int(obj, "k")), .n = @intCast(ref.int(obj, "n")), .halves = @intCast(ref.int(obj, "k2")) }; + } + fn mat(s: @This(), obj: Obj) !Mat { + return .{ .w = try s.f(obj, "w"), .n = @intCast(ref.int(obj, "n")), .k = @intCast(ref.int(obj, "k")) }; + } + fn hc(s: @This(), obj: Obj) ![3][]const u8 { + return .{ try s.f(obj, "fn"), try s.f(obj, "base"), try s.f(obj, "scale") }; + } + }; + const s: R = .{ .a = a, .io = io, .dir = dir }; + const mo = o.get("moe"); + const m = (mo orelse o.get("mlp").?).object; + const sh = if (mo != null) m.get("shared").?.object else m; + var h: Host = .{ + .attn_hc = try s.hc(o.get("attn_hc").?.object), + .ffn_hc = try s.hc(o.get("ffn_hc").?.object), + .in_norm = try s.f(o, "in_norm"), + .post_norm = try s.f(o, "post_norm"), + .router = if (mo != null) try s.f(m, "router") else "", + .bias = if (mo != null) try s.f(m, "bias") else "", + .kda = if (o.get("kda")) |kv| kda: { // a DSA layer's fixture has none + const k = kv.object; + break :kda .{ + .qkv = try s.lin(k.get("qkv").?.object), + .o = try s.lin(k.get("o").?.object), + .fa = try s.mat(k.get("fa").?.object), + .ga = try s.mat(k.get("ga").?.object), + .b = try s.mat(k.get("b").?.object), + .fb = try s.mat(k.get("fb").?.object), + .gb = try s.mat(k.get("gb").?.object), + .conv = try s.f(k, "conv"), + .a_log = try s.f(k, "a_log"), + .dt_bias = try s.f(k, "dt_bias"), + .norm = try s.f(k, "norm"), + .heads = @intCast(ref.int(k, "heads")), + .fa_off = @intCast(ref.int(k, "fa_off")), + .ga_off = @intCast(ref.int(k, "ga_off")), + .b_off = @intCast(ref.int(k, "b_off")), + .p_stride = @intCast(ref.int(k, "p_stride")), + }; + } else null, + .ex = undefined, + .sh = .{ try s.lin(sh.get("gate").?.object), try s.lin(sh.get("up").?.object), try s.lin(sh.get("down").?.object) }, + .dense = mo == null, + }; + if (h.dense) return h; + const ex = m.get("experts").?.object; + if (ref.int(ex, "cb") != 2) return error.GlmLayerCodebook; + const count: usize = @intCast(ref.int(ex, "count")); + h.ex = .{ .proj = undefined, .offs = undefined, .k2 = undefined, .suh_g = try s.f(ex, "suh_g"), .suh_u = try s.f(ex, "suh_u"), .svh_g = try s.f(ex, "svh_g"), .svh_u = try s.f(ex, "svh_u"), .suh_d = try s.f(ex, "suh_d"), .svh_d = try s.f(ex, "svh_d"), .count = count, .dims = @intCast(ref.int(ex, "dims")), .width = @intCast(ref.int(ex, "width")) }; + const names = [3][3][]const u8{ .{ "gate", "gate_offs", "gate_k2" }, .{ "up", "up_offs", "up_k2" }, .{ "down", "down_offs", "down_k2" } }; + for (names, 0..) |nm, j| { + h.ex.proj[j] = try s.f(ex, nm[0]); + h.ex.offs[j] = try a.alloc(u64, count); + h.ex.k2[j] = try a.alloc(i32, count); + for (ex.get(nm[1]).?.array.items, ex.get(nm[2]).?.array.items, 0..) |off, k2, e| { + h.ex.offs[j][e] = @intCast(off.integer); + h.ex.k2[j][e] = @intCast(k2.integer); + } + } + return h; + } +}; + +pub fn bfs(a: std.mem.Allocator, bytes: []const u8) ![]f64 { + const u = std.mem.bytesAsSlice(u16, bytes); + const out = try a.alloc(f64, u.len); + for (u, 0..) |v, i| out[i] = ref.bf(v); + return out; +} + +pub fn f32s(bytes: []const u8) []const f32 { + return ref.f32s(bytes); +} + +fn f16At(bytes: []const u8, i: usize) f64 { + return @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)))); +} + +fn bf16r(v: f64) f64 { + const u: u32 = @bitCast(@as(f32, @floatCast(v))); + return @as(f32, @bitCast((u +% (0x7FFF + ((u >> 16) & 1))) & 0xFFFF0000)); +} + +fn sigmoid(x: f64) f64 { + return 1.0 / (1.0 + @exp(-x)); +} + +/// The worst |got - want| / tol of one stage, NaN counting as infinite. +const Bound = struct { + worst: f64 = 0, + at: usize = 0, + + fn add(b: *Bound, got: f64, want: f64, tol: f64, i: usize) void { + const q = @abs(got - want) / tol; + if (!(q <= b.worst)) { + b.worst = if (std.math.isNan(q)) std.math.inf(f64) else q; + b.at = i; + } + } + + fn rel(b: *Bound, got: f64, want: f64, scale: f64, floor: f64, i: usize) void { + b.add(got, want, @abs(want) * scale + floor, i); + } +}; + +/// The stages' verdicts: the first stage past its tolerance fails the window; `line` lists every stage's worst. +pub const Report = struct { + a: std.mem.Allocator, + tag: []const u8, + id: []const u8 = id, + worst: f64 = 0, + line: []const u8 = "", + + fn stage(rp: *Report, name: []const u8, b: Bound) !void { + if (b.worst > 1.0) { + std.debug.print("FAIL {s} {s} {s} f64: worst {d:.4} of tolerance at element {d}\n", .{ rp.id, rp.tag, name, b.worst, b.at }); + return error.TestFailed; + } + rp.worst = @max(rp.worst, b.worst); + rp.line = try std.fmt.allocPrint(rp.a, "{s} {s} {d:.3};", .{ rp.line, name, b.worst }); + } +}; + +/// y [m <= 16, n] of x [m, k] through an EXL3 linear (exl3_dense_tests' reference) and each output's |svh| sum |xh wq|. +pub fn exl3Ref(a: std.mem.Allocator, l: Lin, x: []const f64, m: usize) !struct { want: []f64, s1: []f64 } { + if (m == 0 or m > 16) return error.GlmLayerRows; + const kt = l.k / 16; + const nt = l.n / 16; + const tw: usize = @as(usize, l.halves) * 8; + const src: []const i16 = @alignCast(std.mem.bytesAsSlice(i16, l.words)); + var tr = src; + if (l.strips) { // [N/128, K/16, 8, W] back to the stored [K/16, N/16, W] + const st = try a.alloc(i16, src.len); + for (0..nt / 8) |nb| for (0..kt) |ki| for (0..8) |j| { + @memcpy(st[(ki * nt + nb * 8 + j) * tw ..][0..tw], src[((nb * kt + ki) * 8 + j) * tw ..][0..tw]); + }; + tr = st; + } + const wq = try a.alloc(f16, l.k * l.n); + exl3fmt.unpack(.{ .halves = l.halves }, .mul1, tr, kt, nt, wq); + const xs = try a.alloc(f64, m * l.k); + for (0..m) |row| for (0..l.k) |i| { + xs[row * l.k + i] = x[row * l.k + i] * f16At(l.suh, i); + }; + exl3fmt.rotate(xs, m, l.k, 1); + const ypre = try a.alloc(f64, m * l.n); + const s1 = try a.alloc(f64, m * l.n); + @memset(ypre, 0); + @memset(s1, 0); + var xv: [16]f64 = undefined; + for (0..l.k) |ki| { + for (0..m) |row| xv[row] = @floatCast(@as(f16, @floatCast(xs[row * l.k + ki]))); + for (0..l.n) |col| { + const wv: f64 = @floatCast(wq[ki * l.n + col]); + for (0..m) |row| { + ypre[col * m + row] += xv[row] * wv; + s1[col * m + row] += @abs(xv[row] * wv); + } + } + } + exl3fmt.rotate(ypre, l.n, m, 0); + const want = try a.alloc(f64, m * l.n); + const s = try a.alloc(f64, m * l.n); + for (0..l.n) |col| { + const sv = f16At(l.svh, col); + for (0..m) |row| { + want[row * l.n + col] = ypre[col * m + row] * sv; + s[row * l.n + col] = s1[col * m + row] * @abs(sv); + } + } + return .{ .want = want, .s1 = s }; +} + +/// The rows the EXL3 stages recompute: every decode row, four rows of a prompt chunk. +pub fn samples(rows: usize) []const usize { + const all = [_]usize{ 0, 1, 2, 3, 4, 5, 6, 7 }; + if (rows <= all.len) return all[0..rows]; + const S = struct { + var four: [4]usize = undefined; + }; + S.four = .{ 0, 1, rows / 2, rows - 1 }; + return &S.four; +} + +/// x's chosen rows [m, k]. +fn pick(a: std.mem.Allocator, x: []const f64, k: usize, stride: usize, off: usize, rows: []const usize) ![]f64 { + const out = try a.alloc(f64, rows.len * k); + for (rows, 0..) |row, i| @memcpy(out[i * k ..][0..k], x[row * stride + off ..][0..k]); + return out; +} + +/// An EXL3 stage on sample rows: got (bf16 or fp32 [rows, stride] at column off) vs the f64 reference. +fn exl3Stage(rp: *Report, a: std.mem.Allocator, name: []const u8, l: Lin, x: []const f64, rows: []const usize, got: []const f64, stride: usize, off: usize, f32_out: bool) !void { + const ref_ = try exl3Ref(a, l, try pick(a, x, l.k, l.k, 0, rows), rows.len); + var b = Bound{}; + for (rows, 0..) |row, i| for (0..l.n) |col| { + const want = ref_.want[i * l.n + col]; + const tol = @abs(want) / @as(f64, if (f32_out) 1024 else 128) + ref_.s1[i * l.n + col] / 2048 + 1e-30; + b.add(got[row * stride + off + col], want, tol, row * stride + off + col); + }; + try rp.stage(name, b); +} + +/// A plain bf16 matmul over every row: x [rows, k] at (xs, xo), w [n, k], got bf16 [rows, n]. +fn b16Stage(rp: *Report, a: std.mem.Allocator, name: []const u8, mt: Mat, x: []const f64, xs: usize, xo: usize, got: []const f64, gs: usize, go: usize, rows: usize) !void { + const wv = try bfs(a, mt.w); + var b = Bound{}; + for (0..rows) |row| for (0..mt.n) |col| { + var acc: f64 = 0; + var s: f64 = 0; + for (0..mt.k) |j| { + const t = x[row * xs + xo + j] * wv[col * mt.k + j]; + acc += t; + s += @abs(t); + } + b.add(got[row * gs + go + col], acc, @abs(acc) / 128 + s / 1024 + 1e-30, row * gs + go + col); + }; + try rp.stage(name, b); +} + +/// glue.hc_pre in f64: the 24 mixes, gates, Sinkhorn, the collapsed row's norm (glm_glue_tests' reference). +pub fn hcPreStage(rp: *Report, a: std.mem.Allocator, name: []const u8, c: lay.Config, hc: [3][]const u8, nw_b: []const u8, x: []const f64, rows: usize, normed: []const f64, post: []const f32, comb: []const f32) !void { + const d: usize = c.d; + const wide = 4 * d; + const fnv = f32s(hc[0]); + const basev = f32s(hc[1]); + const scv = f32s(hc[2]); + const nwv = try bfs(a, nw_b); + const hc_eps: f64 = c.hc_eps; + var b = Bound{}; + var bp = Bound{}; + var bc = Bound{}; + const cv = try a.alloc(f64, d); + for (0..rows) |i| { + var mix: [24]f64 = @splat(0); + var ss: f64 = 0; + for (0..wide) |j| { + const xj = x[i * wide + j]; + ss += xj * xj; + for (0..24) |m| mix[m] += fnv[m * wide + j] * xj; + } + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(wide)) + c.eps); + var p: [4]f64 = undefined; + for (0..4) |sv| { + p[sv] = sigmoid(mix[sv] * rinv * scv[0] + basev[sv]) + hc_eps; + bp.rel(post[i * 4 + sv], 2.0 * sigmoid(mix[4 + sv] * rinv * scv[1] + basev[4 + sv]), 1.0 / 1024.0, 1e-6, i * 4 + sv); + } + var ce: [4][4]f64 = undefined; + for (0..4) |r| { + var cmax: f64 = -std.math.inf(f64); + for (0..4) |q| { + ce[r][q] = mix[8 + r * 4 + q] * rinv * scv[2] + basev[8 + r * 4 + q]; + cmax = @max(cmax, ce[r][q]); + } + var rs: f64 = 0; + for (0..4) |q| { + ce[r][q] = @exp(ce[r][q] - cmax); + rs += ce[r][q]; + } + for (0..4) |q| ce[r][q] = ce[r][q] / rs + hc_eps; + } + for (0..c.hc_iters) |it| { + if (it > 0) for (0..4) |r| { + var rs: f64 = 0; + for (0..4) |q| rs += ce[r][q]; + for (0..4) |q| ce[r][q] /= rs + hc_eps; + }; + for (0..4) |q| { + var cs: f64 = 0; + for (0..4) |r| cs += ce[r][q]; + for (0..4) |r| ce[r][q] /= cs + hc_eps; + } + } + for (0..4) |r| for (0..4) |q| bc.rel(comb[i * 16 + r * 4 + q], ce[r][q], 1.0 / 512.0, 1e-6, i * 16 + r * 4 + q); + var cs2: f64 = 0; + for (0..d) |j| { + cv[j] = p[0] * x[i * wide + j] + p[1] * x[i * wide + d + j] + p[2] * x[i * wide + 2 * d + j] + p[3] * x[i * wide + 3 * d + j]; + cs2 += cv[j] * cv[j]; + } + const rinv2 = 1.0 / @sqrt(cs2 / @as(f64, @floatFromInt(d)) + c.eps); + for (0..d) |j| b.rel(normed[i * d + j], nwv[j] * cv[j] * rinv2, 1.0 / 64.0, 3e-3, i * d + j); + } + try rp.stage(name, b); + try rp.stage(try std.fmt.allocPrint(a, "{s} post", .{name}), bp); + try rp.stage(try std.fmt.allocPrint(a, "{s} comb", .{name}), bc); +} + +/// glue.hc_post in f64: bf16(rank 0 + rank 1) times the post gate plus the comb mix of the streams. +pub fn hcPostStage(rp: *Report, name: []const u8, d: usize, x: []const f64, g0: []const f32, g1: []const f32, post: []const f32, comb: []const f32, rows: usize, got: []const f64) !void { + var b = Bound{}; + for (0..rows) |i| for (0..4) |sv| for (0..d) |j| { + const branch = bf16r(@as(f64, g0[i * d + j]) + g1[i * d + j]); + var mixed: f64 = 0; + for (0..4) |q| mixed += comb[i * 16 + q * 4 + sv] * x[(i * 4 + q) * d + j]; + const at = (i * 4 + sv) * d + j; + b.rel(got[at], post[i * 4 + sv] * branch + mixed, 1.0 / 64.0, 3e-3, at); + }; + try rp.stage(name, b); +} + +/// One file of the window (stage outputs are equal bits on both sides by now). +pub const Win = struct { + a: std.mem.Allocator, + io: std.Io, + dir: []const u8, + e: Obj, + + pub fn raw(w: Win, site: ?[]const u8, key: []const u8) ![]u8 { + const o = if (site) |s| w.e.get(s).?.object else w.e; + return ref.readFile(w.a, w.io, w.dir, ref.str(o, key)); + } + + pub fn bf(w: Win, site: ?[]const u8, key: []const u8) ![]f64 { + return bfs(w.a, try w.raw(site, key)); + } + + pub fn f(w: Win, site: ?[]const u8, key: []const u8) ![]const f32 { + return f32s(try w.raw(site, key)); + } +}; + +/// The whole layer in f64, stage by stage from each stage's inputs; returns the worst fraction of tolerance. +pub fn layerF64(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, c: lay.Config, h: *const Host, tag: []const u8, e: Obj, lid: []const u8) !f64 { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const w: Win = .{ .a = a, .io = io, .dir = dir, .e = e }; + var rp: Report = .{ .a = a, .tag = tag, .id = lid }; + const rows: usize = @intCast(ref.int(e, "rows")); + const d: usize = c.d; + const kd = if (h.kda) |*v| v else return error.GlmLayerKind; + const ps: usize = kd.p_stride; + const sm = samples(rows); + // the attention site: hc_pre, the KDA block's projections, its chain, o_proj, hc_post + const x0 = try w.bf(null, "x0"); + const n0 = try w.bf("s0", "normed"); + try hcPreStage(&rp, a, "attn hc_pre", c, h.attn_hc, h.in_norm, x0, rows, n0, try w.f("s0", "post"), try w.f("s0", "comb")); + const p = try w.bf("s0", "p"); + try exl3Stage(&rp, a, "qkv", kd.qkv, n0, sm, p, ps, 0, false); + try b16Stage(&rp, a, "f_a", kd.fa, n0, d, 0, p, ps, kd.fa_off, rows); + try b16Stage(&rp, a, "g_a", kd.ga, n0, d, 0, p, ps, kd.ga_off, rows); + try b16Stage(&rp, a, "b", kd.b, n0, d, 0, p, ps, kd.b_off, rows); + try b16Stage(&rp, a, "f_b", kd.fb, p, ps, kd.fa_off, try w.bf("s0", "ka"), kd.fb.n, 0, rows); + try b16Stage(&rp, a, "g_b", kd.gb, p, ps, kd.ga_off, try w.bf("s0", "kg"), kd.gb.n, 0, rows); + try kdaStage(&rp, a, w, c, h, rows); + const kout = try w.bf("s0", "kout"); + const part0 = try w.f("s0", "part"); + const part0v = try a.alloc(f64, part0.len); + for (part0v, part0) |*o, v| o.* = v; + try exl3Stage(&rp, a, "o_proj", kd.o, kout, sm, part0v, d, 0, true); + const x1 = try w.bf("s1", "x"); + try hcPostStage(&rp, "attn hc_post", d, x0, part0, try w.f("s0", "peer"), try w.f("s0", "post"), try w.f("s0", "comb"), rows, x1); + // the ffn site: hc_pre, the MoE (router, select, routed and shared experts, combine) or the dense MLP, hc_post + const n1 = try w.bf("s1", "normed"); + try hcPreStage(&rp, a, "ffn hc_pre", c, h.ffn_hc, h.post_norm, x1, rows, n1, try w.f("s1", "post"), try w.f("s1", "comb")); + if (h.dense) _ = try swigluStage(&rp, a, w, c, h.sh, n1, rows, "mlp", "part") else try moeStage(&rp, a, w, c, h, n1, rows); + try hcPostStage(&rp, "ffn hc_post", d, x1, try w.f("s1", "part"), try w.f("s1", "peer"), try w.f("s1", "post"), try w.f("s1", "comb"), rows, try w.bf(null, "x")); + std.debug.print("RESULT {s} {s} f64 per stage (fraction of tolerance):{s}\n", .{ lid, tag, rp.line }); + return rp.worst; +} + +/// The chain (serial for decode windows, chunked at pos for a prompt chunk) from the engine's prep, f64. +fn kdaStage(rp: *Report, a: std.mem.Allocator, w: Win, c: lay.Config, h: *const Host, rows: usize) !void { + const hd: usize = h.kda.?.heads; + const per = hd * ref.DK; + const out64 = try a.alloc(f64, rows * per); + const st64 = try a.alloc(f64, per * ref.DK); + const q = f32s(try w.raw("prep", "q")); + const k = f32s(try w.raw("prep", "k")); + const v = f32s(try w.raw("prep", "v")); + const lg = f32s(try w.raw("prep", "lg")); + const beta = f32s(try w.raw("prep", "beta")); + const rec0 = try w.f(null, "rec0"); + const gate = ref.bf16s(try w.raw("s0", "kg")); + const norm = ref.bf16s(h.kda.?.norm); + if (w.e.get("prefill").?.bool) { + try ref.chunkedRef(a, q, k, v, lg, beta, rec0, gate, norm, c.eps, hd, rows, @intCast(ref.int(w.e, "pos")), out64, st64); + } else { + try ref.serialRef(a, q, k, v, lg, beta, rec0, gate, norm, c.eps, hd, rows, out64, st64); + } + const kout = try w.bf("s0", "kout"); + var b = Bound{}; + for (out64, 0..) |want, i| b.rel(kout[i], want, 1.0 / 128.0, 2e-3, i); + try rp.stage("kda out", b); + const rec = try w.f(null, "rec"); + var bs = Bound{}; + for (st64, 0..) |want, i| bs.rel(rec[i], want, 2e-2, 2e-3, i); + try rp.stage("kda state", bs); +} + +/// MoE in f64: router dots, the biased top-k and its weights, routed experts, the shared expert, the combine. +pub fn moeStage(rp: *Report, a: std.mem.Allocator, w: Win, c: lay.Config, h: *const Host, n1: []const f64, rows: usize) !void { + const d: usize = c.d; + const ne: usize = c.experts; + const slots: usize = c.top_k + 1; + const rw = try bfs(a, h.router); + const mlog = try w.f("s1", "mlog"); + var b = Bound{}; + for (0..rows) |i| for (0..ne) |ex| { + var acc: f64 = 0; + for (0..d) |j| acc += n1[i * d + j] * rw[ex * d + j]; + b.rel(mlog[i * ne + ex], acc, 1.0 / 1024.0, 1e-5, i * ne + ex); + }; + try rp.stage("router", b); + const bias = f32s(h.bias); + const picks: []const i32 = @alignCast(std.mem.bytesAsSlice(i32, try w.raw("s1", "pick"))); + const wts = try w.f("s1", "wts"); + var bw = Bound{}; + const score = try a.alloc(f64, ne); + const used = try a.alloc(bool, ne); + for (0..rows) |i| { + for (0..ne) |ex| { + score[ex] = sigmoid(mlog[i * ne + ex]); + used[ex] = false; + } + var total: f64 = 0; + for (0..c.top_k) |kk| { // the biased greedy pick, lowest id on ties; a different pick must be a near-tie + var best: f64 = -std.math.inf(f64); + var at: usize = ne; + for (0..ne) |ex| if (!used[ex] and score[ex] + bias[ex] > best) { + best = score[ex] + bias[ex]; + at = ex; + }; + const got: usize = @intCast(picks[i * slots + kk]); + if (got != at and !(got < ne and !used[got] and @abs(score[got] + bias[got] - best) < 1e-6)) { + std.debug.print("FAIL {s} {s} select f64: row {d} slot {d} picks {d}, f64 picks {d}\n", .{ rp.id, rp.tag, i, kk, got, at }); + return error.TestFailed; + } + used[got] = true; + total += score[got]; + } + for (0..c.top_k) |kk| bw.rel(wts[i * slots + kk], score[@intCast(picks[i * slots + kk])] / total * c.routed_scale, 1.0 / 1024.0, 1e-6, i * slots + kk); + if (picks[i * slots + c.top_k] != @as(i32, @intCast(ne)) or wts[i * slots + c.top_k] != 1.0) { + std.debug.print("FAIL {s} {s} select: row {d}'s last slot is not the shared expert at weight 1\n", .{ rp.id, rp.tag, i }); + return error.TestFailed; + } + } + try rp.stage("select", bw); + try expertStage(rp, a, w, c, h, n1, picks, rows); + _ = try swigluStage(rp, a, w, c, h.sh, n1, rows, "shared", "sy"); + const ey = try w.f("s1", "ey"); + const part = try w.f("s1", "part"); + var bc = Bound{}; + for (0..rows) |i| for (0..d) |j| { + var acc: f64 = 0; + for (0..slots) |kk| acc += @as(f64, wts[i * slots + kk]) * ey[(i * slots + kk) * d + j]; + bc.rel(part[i * d + j], acc, 1.0 / 1024.0, 1e-5, i * d + j); + }; + try rp.stage("combine", bc); +} + +/// A SwiGLU MLP (the shared expert, or a dense layer's): gate and up (EXL3), SwiGLU, down (EXL3, fp32 into `out`). +fn swigluStage(rp: *Report, a: std.mem.Allocator, w: Win, c: lay.Config, lin: [3]Lin, n1: []const f64, rows: usize, label: []const u8, out: []const u8) ![]const f32 { + const sm = samples(rows); + const sw: usize = lin[0].n; + const smg = try w.bf("s1", "smg"); + const smu = try w.bf("s1", "smu"); + try exl3Stage(rp, a, try std.fmt.allocPrint(a, "{s} gate", .{label}), lin[0], n1, sm, smg, sw, 0, false); + try exl3Stage(rp, a, try std.fmt.allocPrint(a, "{s} up", .{label}), lin[1], n1, sm, smu, sw, 0, false); + const sact = try w.bf("s1", "sact"); + var bs = Bound{}; + for (0..rows * sw) |i| { + const g = @min(smg[i], c.limit); + const u = std.math.clamp(smu[i], -c.limit, c.limit); + bs.rel(sact[i], bf16r(bf16r(g * sigmoid(g)) * u), 1.0 / 64.0, 3e-3, i); + } + try rp.stage(try std.fmt.allocPrint(a, "{s} swiglu", .{label}), bs); + const y = try w.f("s1", out); + const yv = try a.alloc(f64, y.len); + for (yv, y) |*o, v| o.* = v; + try exl3Stage(rp, a, try std.fmt.allocPrint(a, "{s} down", .{label}), lin[2], sact, sm, yv, c.d, 0, true); + return y; +} + +/// Routed slots of the first and last rows: gate/up, the bf16 SwiGLU the epilogue rounds, down, each slot's y. +fn expertStage(rp: *Report, a: std.mem.Allocator, w: Win, c: lay.Config, h: *const Host, n1: []const f64, picks: []const i32, rows: usize) !void { + const d: usize = c.d; + const slots: usize = c.top_k + 1; + const ex = &h.ex; + const ey = try w.f("s1", "ey"); + var b = Bound{}; + const which = [2]usize{ 0, rows - 1 }; + for (which[0..if (rows > 1) 2 else 1]) |i| for (0..c.top_k) |kk| { + const e: usize = @intCast(picks[i * slots + kk]); + const x = n1[i * d ..][0..d]; + const g = try exl3Ref(a, ex.lin(0, e), x, 1); + const u = try exl3Ref(a, ex.lin(1, e), x, 1); + const act = try a.alloc(f64, ex.width); + for (act, g.want, u.want) |*o, gv, uv| { + const gg = @min(bf16r(gv), c.limit); + const uu = std.math.clamp(bf16r(uv), -c.limit, c.limit); + o.* = bf16r(bf16r(gg * sigmoid(gg)) * uu); + } + const y = try exl3Ref(a, ex.lin(2, e), act, 1); + for (0..d) |j| { + const at = (i * slots + kk) * d + j; + b.add(ey[at], y.want[j], @abs(y.want[j]) / 128 + y.s1[j] / 128 + 1e-30, at); + } + }; + try rp.stage("routed experts", b); +} diff --git a/zig/tests/cuda/glm_layer_tests.zig b/zig/tests/cuda/glm_layer_tests.zig new file mode 100644 index 0000000000..40f717c96c --- /dev/null +++ b/zig/tests/cuda/glm_layer_tests.zig @@ -0,0 +1,247 @@ +//! glm.layer.kda (glm.layer.mlp: a dense layer): a real KDA layer at TP=2 rank 0 vs Python's layer_forward, bit-exact. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const lref = @import("glm_layer_ref.zig"); +const lay = cuda.glm_layer; +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; + +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const Obj = std.json.ObjectMap; +const id = "glm.layer.kda"; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + kdaLayer(r, dir, "layer", id) catch |e| state.fail(e, id); +} + +fn skip(lid: []const u8) void { + std.debug.print("SKIP {s}\n", .{lid}); +} + +pub fn file(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + return ref.readFile(r.a, r.gpu.io, dir, name); +} + +pub fn up(r: *Rig, dir: []const u8, name: []const u8) !u64 { + return r.dev(try file(r, dir, name)); +} + +pub fn uint(o: Obj, key: []const u8) u32 { + return @intCast(ref.int(o, key)); +} + +/// The layer's constants (manifest params). +pub fn config(p: Obj) lay.Config { + return .{ + .eps = @floatCast(ref.float(p, "eps")), + .hc_eps = @floatCast(ref.float(p, "hc_eps")), + .hc_iters = uint(p, "hc_iters"), + .lower = @floatCast(ref.float(p, "lower")), + .limit = @floatCast(ref.float(p, "limit")), + .routed_scale = @floatCast(ref.float(p, "routed_scale")), + .top_k = uint(p, "top_k"), + .experts = uint(p, "experts"), + .d = uint(p, "hidden"), + .streams = uint(p, "streams"), + }; +} + +/// The manifest section's KDA layer (MoE or dense MLP) through every window, reported as `lid`. +pub fn kdaLayer(r: *Rig, dir: []const u8, section: []const u8, lid: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) return skip(lid); + var man = try ref.Manifest.open(r.gpu.gpa, r.gpu.io, dir); + defer man.deinit(); + const o = (man.root().get(section) orelse return skip(lid)).object; + const c = config(o.get("params").?.object); + const host = try lref.Host.read(r.a, r.gpu.io, dir, o.get("weights").?.object); + const w = try weights(r, &host, c); + var k = try lay.Kernels.load(r.gpu.d); + defer k.unload(); + var worst: f64 = 0; + const cases = o.get("cases").?.object; + var it = cases.iterator(); + while (it.next()) |kv| { + const t0 = check.now(r.gpu.io); + const tag = kv.key_ptr.*; + const e = kv.value_ptr.object; + try window(r, &k, dir, c, &w, tag, e, lid); + const q = try lref.layerF64(r.a, r.gpu.io, dir, c, &host, tag, e, lid); + worst = @max(worst, q); + const dt = @as(f64, @floatFromInt(check.now(r.gpu.io) - t0)) / 1e9; + std.debug.print("RESULT {s} {s}: bit-exact, f64 worst {d:.4} of tolerance ({d:.1} s)\n", .{ lid, tag, q, dt }); + } + std.debug.print("RESULT {s}: {d} windows bit-exact, f64 worst {d:.4} of tolerance\n", .{ lid, cases.count(), worst }); + check.pass("{s}", .{lid}); +} + +pub fn linW(r: *Rig, l: lref.Lin) !exl3.Weight { + const st = exl3.strides(l.k, l.n, l.halves, true); + return .{ + .trellis = try r.dev(l.words), + .suh = try r.dev(l.suh), + .svh = try r.dev(l.svh), + .k = @intCast(l.k), + .n = @intCast(l.n), + .halves = l.halves, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; +} + +pub fn b16W(r: *Rig, b: lref.Mat) !cuda.glm_bmm.B16 { + return .{ .w = try r.dev(b.w), .n = @intCast(b.n), .k = @intCast(b.k) }; +} + +/// The MoE on the device as the engine holds it; the routed experts' trellises stay in one buffer each. +pub fn moeW(r: *Rig, h: *const lref.Host, c: lay.Config) !lay.MoeW { + const ex = &h.ex; + var tables: [3]u64 = undefined; + for (ex.proj, ex.offs, 0..) |bytes, offs, j| { + const base = try r.dev(bytes); + const ptrs = try r.a.alloc(u64, ex.count); + for (offs, 0..) |off, i| ptrs[i] = base + off; + tables[j] = try r.dev(std.mem.sliceAsBytes(ptrs)); + } + const experts: exl3e.Layer = .{ + .gate = tables[0], + .up = tables[1], + .down = tables[2], + .gate_k2 = try r.dev(std.mem.sliceAsBytes(ex.k2[0])), + .up_k2 = try r.dev(std.mem.sliceAsBytes(ex.k2[1])), + .down_k2 = try r.dev(std.mem.sliceAsBytes(ex.k2[2])), + .suh_g = try r.dev(ex.suh_g), + .suh_u = try r.dev(ex.suh_u), + .svh_g = try r.dev(ex.svh_g), + .svh_u = try r.dev(ex.svh_u), + .suh_d = try r.dev(ex.suh_d), + .svh_d = try r.dev(ex.svh_d), + .experts = ex.count, + .dims = ex.dims, + .width = ex.width, + .codebook = .mul1, + .limit = c.limit, + .act_mode = 0, + }; + return .{ + .router = .{ .w = try r.dev(h.router), .n = c.experts, .k = c.d }, + .bias = try r.dev(h.bias), + .experts = experts, + .gate = try linW(r, h.sh[0]), + .up = try linW(r, h.sh[1]), + .down = try linW(r, h.sh[2]), + }; +} + +/// Rank 0's tensors on the device as the engine holds them. +fn weights(r: *Rig, h: *const lref.Host, c: lay.Config) !lay.LayerW { + const kd = h.kda orelse return error.GlmLayerKind; + return .{ + .attn_hc = .{ .fn_w = try r.dev(h.attn_hc[0]), .base = try r.dev(h.attn_hc[1]), .scale = try r.dev(h.attn_hc[2]) }, + .ffn_hc = .{ .fn_w = try r.dev(h.ffn_hc[0]), .base = try r.dev(h.ffn_hc[1]), .scale = try r.dev(h.ffn_hc[2]) }, + .in_norm = try r.dev(h.in_norm), + .post_norm = try r.dev(h.post_norm), + .mixer = .{ .kda = .{ + .qkv = try linW(r, kd.qkv), + .fab = try b16W(r, .{ .w = try std.mem.concat(r.a, u8, &.{ kd.fa.w, kd.ga.w, kd.b.w }), .n = kd.fa.n + kd.ga.n + kd.b.n, .k = kd.fa.k }), + .fb = try b16W(r, kd.fb), + .gb = try b16W(r, kd.gb), + .o = try linW(r, kd.o), + .conv = try r.dev(kd.conv), + .a_log = try r.dev(kd.a_log), + .dt_bias = try r.dev(kd.dt_bias), + .norm = try r.dev(kd.norm), + .heads = kd.heads, + .fa_off = kd.fa_off, + .ga_off = kd.ga_off, + .b_off = kd.b_off, + .p_stride = kd.p_stride, + } }, + .ffn = if (h.dense) .{ .mlp = .{ .gate = try linW(r, h.sh[0]), .up = try linW(r, h.sh[1]), .down = try linW(r, h.sh[2]) } } else .{ .moe = try moeW(r, h, c) }, + }; +} + +pub fn same(r: *Rig, dir: []const u8, label: []const u8, ptr: u64, name: []const u8) !void { + const want = try file(r, dir, name); + try check.sameBytes(label, try r.back(ptr, want.len), want); +} + +pub const Stage = struct { []const u8, u64 }; + +/// The test's all-gather: own part in slot 0, the fixture's rank-1 partial in slot 1; then the site's stage bits. +pub const Ctx = struct { + r: *Rig, + id: []const u8, + dir: []const u8, + tag: []const u8, + e: Obj, + sc: *const lay.Scratch, + attn: []const Stage, // the mixer's stages + x: u64, + d: usize, + + pub fn gather(p: *anyopaque, s: cuda.Stream, site: lay.Site, part: u64, gath: u64, rows: usize) anyerror!void { + const c: *Ctx = @ptrCast(@alignCast(p)); + try s.synchronize(); + const n = rows * c.d * 4; + const so = c.e.get(if (site == .attn) "s0" else "s1").?.object; + const g: cuda.DeviceBuffer = .{ .d = c.r.gpu.d, .ptr = gath, .len = 2 * n, .borrowed = true }; + try g.copyFrom(0, part, n, s.handle); + try g.upload(n, try file(c.r, c.dir, ref.str(so, "peer"))); + const sc = c.sc; + // s0 holds no x (the layer's input); s1's is the mixer's hc_post + const hc = [_]Stage{ .{ "x", c.x }, .{ "normed", sc.normed }, .{ "post", sc.post }, .{ "comb", sc.comb } }; + const moe = [_]Stage{ + .{ "mlog", sc.mlog }, .{ "pick", sc.pick }, .{ "wts", sc.wts }, .{ "ey", sc.ey }, + .{ "smg", sc.smg }, .{ "smu", sc.smu }, .{ "sact", sc.sact }, .{ "sy", sc.sy }, + }; + const tail = [_]Stage{.{ "part", part }}; + for ([_][]const Stage{ &hc, if (site == .attn) c.attn else &moe, &tail }) |list| for (list) |st| { + const name = (so.get(st[0]) orelse continue).string; + const label = try std.fmt.allocPrint(c.r.a, "{s} {s} {s} {s}", .{ c.id, c.tag, @tagName(site), st[0] }); + try same(c.r, c.dir, label, st[1], name); + }; + } +}; + +fn sameAt(r: *Rig, dir: []const u8, lid: []const u8, tag: []const u8, what: []const u8, ptr: u64, name: []const u8) !void { + try same(r, dir, try std.fmt.allocPrint(r.a, "{s} {s} {s}", .{ lid, tag, what }), ptr, name); +} + +/// One window from the fixture's state: forward (+ commit for decode), every stage's bits, x and the KDA state. +fn window(r: *Rig, k: *lay.Kernels, dir: []const u8, c: lay.Config, w: *const lay.LayerW, tag: []const u8, e: Obj, lid: []const u8) !void { + const rows: usize = @intCast(ref.int(e, "rows")); + const prefill = e.get("prefill").?.bool; + var sc = try lay.Scratch.init(r.gpu.d, r.s, c, w, rows, 2); + defer sc.deinit(); + const x = try up(r, dir, ref.str(e, "x0")); + const rec0 = try file(r, dir, ref.str(e, "rec0")); + const st: lay.KdaState = .{ .conv = try up(r, dir, ref.str(e, "conv0")), .rec_in = try r.dev(rec0), .rec_out = try r.zeros(rec0.len) }; + const attn = [_]Stage{ + .{ "p", sc.p }, .{ "ka", sc.ka }, .{ "kg", sc.kg }, .{ "kout", sc.kout }, + .{ "sk", sc.kda.k }, .{ "sv", sc.kda.v }, .{ "sg", sc.kda.g }, .{ "sb", sc.kda.b }, + }; + var ctx: Ctx = .{ .r = r, .id = lid, .dir = dir, .tag = tag, .e = e, .sc = &sc, .attn = &attn, .x = x, .d = c.d }; + const g: lay.Gather = .{ .ctx = &ctx, .world = 2, .run = Ctx.gather }; + const mode: lay.Mode = if (prefill) .{ .prompt = @intCast(ref.int(e, "pos")) } else .decode; + try lay.forward(k, &r.s, c, w, .{ .kda = .{ .st = st, .mode = mode } }, &sc, x, rows, g); + const kw = &w.mixer.kda; + if (!prefill) try lay.commit(k, &r.s, kw, st, &sc, rows, rows); + try sameAt(r, dir, lid, tag, "x", x, ref.str(e, "x")); + try sameAt(r, dir, lid, tag, "rec", st.rec_out, ref.str(e, "rec")); + try sameAt(r, dir, lid, tag, "conv", st.conv, ref.str(e, "conv")); + // the oracle ran the chunk row-split too (HC_SPLIT=1, PREFILL_OVERLAP=2): own rows and both states equal unsplit + if (prefill) try check.expect(if (e.get("split_equal")) |v| v.bool else false, "{s} {s}: no split check", .{ lid, tag }); + if (e.get("keep")) |kv| { + const ko = kv.object; + const kst: lay.KdaState = .{ .conv = try up(r, dir, ref.str(e, "conv0")), .rec_in = st.rec_in, .rec_out = try r.zeros(rec0.len) }; + try lay.commit(k, &r.s, kw, kst, &sc, rows, @intCast(ref.int(ko, "keep"))); + try sameAt(r, dir, lid, tag, "keep rec", kst.rec_out, ref.str(ko, "rec")); + try sameAt(r, dir, lid, tag, "keep conv", kst.conv, ref.str(ko, "conv")); + } +} diff --git a/zig/tests/cuda/glm_mlp_tests.zig b/zig/tests/cuda/glm_mlp_tests.zig new file mode 100644 index 0000000000..4a3ed1378d --- /dev/null +++ b/zig/tests/cuda/glm_mlp_tests.zig @@ -0,0 +1,9 @@ +//! glm.layer.mlp: a real RED-SNOW dense-MLP KDA layer at TP=2 rank 0 vs the Python engine's layer_forward, bit for bit. +const glm_tests = @import("glm_tests.zig"); +const layer_tests = @import("glm_layer_tests.zig"); + +const id = "glm.layer.mlp"; + +pub fn run(r: *glm_tests.Rig, state: *glm_tests.Run, dir: []const u8) !void { + layer_tests.kdaLayer(r, dir, "mlp", id) catch |e| state.fail(e, id); +} diff --git a/zig/tests/cuda/glm_model_tests.zig b/zig/tests/cuda/glm_model_tests.zig new file mode 100644 index 0000000000..d7a177ea8d --- /dev/null +++ b/zig/tests/cuda/glm_model_tests.zig @@ -0,0 +1,103 @@ +//! tf-cuda-test glm-model: the whole RED-SNOW model at TP=2, a process a rank, vs the Python engine's greedy replies. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm"); +const check = @import("check.zig"); + +const cap = 32768; // cache slots: the reference's longest case (16,480 tokens) and the Python engine's pool buckets + +const Case = struct { name: []const u8, prompt: []u32, ids: []u32 }; + +/// `reference`: {"cases": [{"name", "prompt_ids", "ids"}], "eos_ids"}; rank 0 prints the PASS/FAIL/RESULT lines. +pub fn run(gpu: check.Gpu, rank_arg: []const u8, addr: []const u8, model: []const u8, reference: []const u8) !void { + const rank = try std.fmt.parseInt(u8, rank_arg, 10); + if (rank > 1) return error.BadRank; + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const ref = try readReference(a, gpu.io, reference); + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + const comm = try cuda.nccl_id.pair(&lib, gpu.io, rank, addr); + defer _ = lib.api.ncclCommDestroy(comm); + + const t0 = check.now(gpu.io); + var w = glm.Weights.load(gpu.gpa, gpu.io, gpu.d, model, rank) catch |e| { + if (rank == 0) std.debug.print("FAIL glm.model.load: {t}\n", .{e}); + return e; + }; + defer w.deinit(); + var e = try glm.Engine.init(gpu.gpa, gpu.d, &w, &lib, comm, cap); + defer e.deinit(); + if (rank == 0) { + std.debug.print("RESULT glm.model.load: {d:.1} GiB a rank in {d:.0} s\n", .{ @as(f64, @floatFromInt(w.pool.bytes)) / (1 << 30), secs(gpu.io, t0) }); + check.pass("glm.model.load", .{}); + } + + var first_ok = true; + var prompt_tokens: usize = 0; + var prefill_s: f64 = 0; + var decoded: usize = 0; + var decode_s: f64 = 0; + for (ref.cases) |c| { + const tp = check.now(gpu.io); + var out: std.ArrayList(u32) = .empty; + try out.append(a, try e.prefill(c.prompt)); + prefill_s += secs(gpu.io, tp); + prompt_tokens += c.prompt.len; + if (out.items[0] != c.ids[0]) first_ok = false; + const td = check.now(gpu.io); + // both ranks pick the same tokens, so both stop together: at the reply's length, an EOS or the first difference + while (out.items.len < c.ids.len and out.items[out.items.len - 1] == c.ids[out.items.len - 1] and !isEos(ref.eos, out.items[out.items.len - 1])) { + try out.append(a, try e.step(out.items[out.items.len - 1])); + } + decode_s += secs(gpu.io, td); + decoded += out.items.len - 1; + if (rank != 0) continue; + const id = try std.fmt.allocPrint(a, "glm.model.greedy.{s}", .{c.name}); + if (std.mem.eql(u32, out.items, c.ids)) { + check.pass("{s}", .{id}); + continue; + } + const n = @min(out.items.len, c.ids.len); + var at: usize = 0; + while (at < n and out.items[at] == c.ids[at]) at += 1; + if (at < n) { + std.debug.print("FAIL {s}: token {d} of {d} is {d}, the Python engine's {d}\n", .{ id, at, c.ids.len, out.items[at], c.ids[at] }); + } else { + std.debug.print("FAIL {s}: {d} tokens, the Python engine's {d}\n", .{ id, out.items.len, c.ids.len }); + } + } + if (rank != 0) return; + if (first_ok) check.pass("glm.model.first_token", .{}) else std.debug.print("FAIL glm.model.first_token: a case's first token differs\n", .{}); + std.debug.print("RESULT glm.model.prefill_tps {d:.1}\n", .{@as(f64, @floatFromInt(prompt_tokens)) / prefill_s}); + std.debug.print("RESULT glm.model.decode_tps {d:.2}\n", .{@as(f64, @floatFromInt(decoded)) / decode_s}); +} + +fn secs(io: std.Io, t0: i96) f64 { + return @as(f64, @floatFromInt(check.now(io) - t0)) / 1e9; +} + +fn isEos(eos: []const u32, tok: u32) bool { + return std.mem.indexOfScalar(u32, eos, tok) != null; +} + +fn ints(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const items = v.array.items; + const out = try a.alloc(u32, items.len); + for (items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +fn readReference(a: std.mem.Allocator, io: std.Io, path: []const u8) !struct { cases: []Case, eos: []u32 } { + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 26)); + const root = (try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{})).object; + const items = root.get("cases").?.array.items; + const cases = try a.alloc(Case, items.len); + for (items, cases) |v, *c| { + const o = v.object; + c.* = .{ .name = o.get("name").?.string, .prompt = try ints(a, o.get("prompt_ids").?), .ids = try ints(a, o.get("ids").?) }; + if (c.prompt.len == 0 or c.ids.len == 0) return error.BadReference; + } + return .{ .cases = cases, .eos = try ints(a, root.get("eos_ids").?) }; +} diff --git a/zig/tests/cuda/glm_sparse_ref.zig b/zig/tests/cuda/glm_sparse_ref.zig new file mode 100644 index 0000000000..ccdd66597f --- /dev/null +++ b/zig/tests/cuda/glm_sparse_ref.zig @@ -0,0 +1,203 @@ +//! The sparse checks' helpers: fixture access and float64 references for the DSA sparse indexer and its selection. + +const std = @import("std"); +const kref = @import("glm_kda_ref.zig"); + +pub const D = 128; // index key dims +pub const H = 32; // index heads +pub const K = 512; // pools a row keeps +pub const pool = 4; // tokens a pool + +pub const Manifest = kref.Manifest; + +pub fn readFile(a: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) ![]u8 { + return kref.readFile(a, io, dir, name); +} + +/// A float64 as the kernel's bf16 rounding (round to nearest even) sees it, widened back to f64. +pub fn bf16rn(x: f64) f64 { + const u: u32 = @bitCast(@as(f32, @floatCast(x))); + const rounded = u + 0x7FFF + ((u >> 16) & 1); + return kref.bf(@intCast(rounded >> 16)); +} + +/// float64 index LayerNorm: y = (x - mean(x)) / sqrt(var(x) + eps) * w + b over bf16 inputs, rows of ikr. +pub fn indexRef(a: std.mem.Allocator, ikr: []const u16, ikr_stride: usize, lnw: []const u16, lnb: []const u16, eps: f64, rows: usize) ![]f64 { + const y = try a.alloc(f64, rows * D); + var r: usize = 0; + while (r < rows) : (r += 1) { + var sum: f64 = 0; + for (0..D) |i| sum += kref.bf(ikr[r * ikr_stride + i]); + const mean = sum / 128.0; + var vsum: f64 = 0; + for (0..D) |i| { + const x = kref.bf(ikr[r * ikr_stride + i]); + vsum += (x - mean) * (x - mean); + } + const dv = @sqrt(vsum / 128.0 + eps); + for (0..D) |i| { + const x = kref.bf(ikr[r * ikr_stride + i]); + y[r * D + i] = (x - mean) / dv * f16v(lnw[i]) + f16v(lnb[i]); + } + } + return y; +} + +/// An fp16 value as f64 (the k_norm weights' checkpoint dtype). +pub fn f16v(u: u16) f64 { + return @as(f64, @as(f16, @bitCast(u))); +} + +/// float64 pooled keys: the gated softmax over a pool's four gate rows times the keys, bf16 rounded, summed. +pub fn poolRef(a: std.mem.Allocator, ig: []const u16, ape: []const u16, ik: []const u16, p0: usize, p1: usize) ![]f64 { + const out = try a.alloc(f64, (p1 - p0) * D); + var p = p0; + while (p < p1) : (p += 1) { + var l: [4][D]f64 = undefined; + var m: [D]f64 = undefined; + for (0..D) |i| { + var mx: f64 = 0; + for (0..4) |j| { + const g = kref.bf(ig[(4 * p + j) * D + i]); + const ap = kref.bf(ape[j * D + i]); + l[j][i] = g + ap; + if (j == 0 or l[j][i] > mx) mx = l[j][i]; + } + m[i] = mx; + } + for (0..D) |i| { + var e: [4]f64 = undefined; + for (0..4) |j| e[j] = @exp(l[j][i] - m[i]); + const s = ((e[0] + e[1]) + e[2]) + e[3]; + var t: [4]f64 = undefined; + for (0..4) |j| { + const k = kref.bf(ik[(4 * p + j) * D + i]); + t[j] = bf16rn(bf16rn(e[j] / s) * k); + } + out[(p - p0) * D + i] = ((t[0] + t[1]) + t[2]) + t[3]; + } + } + return out; +} + +/// float64 pool scores: s_p = sum_h w_h * max(scale * (q_h . k_p), 0) over bf16 inputs. +pub fn scoresRef(a: std.mem.Allocator, qi: []const u16, qi_stride: usize, wts: []const u16, w_stride: usize, pk: []const u16, rows: usize, np: usize, scale: f64, wscale: f64, pos0: i64) ![]f64 { + const out = try a.alloc(f64, rows * np); + var r: usize = 0; + while (r < rows) : (r += 1) { + const npool = @divTrunc(pos0 + @as(i64, @intCast(r)) + 1, 4); + var p: usize = 0; + while (p < np) : (p += 1) { + if (p >= npool) { + out[r * np + p] = -std.math.inf(f64); // pools past the row's last visible complete pool + continue; + } + var s: f64 = 0; + for (0..H) |h| { + var dot: f64 = 0; + for (0..D) |i| { + const q = kref.bf(qi[r * qi_stride + h * D + i]); + const k = kref.bf(pk[p * D + i]); + dot += q * k; + } + const w = kref.bf(wts[r * w_stride + h]); + s += w * wscale * @max(dot * scale, 0.0); + } + out[r * np + p] = s; + } + } + return out; +} + +/// The selection's float64-exact reference: order-key sort of the fp32 scores (ties to lower pool), k best ascending. +pub fn selectRef(a: std.mem.Allocator, scores: []const f32, np: usize, k: usize, pos: i64, rows: usize, tokens: []i32, counts: []i32) !void { + var r: usize = 0; + while (r < rows) : (r += 1) { + const idx = try a.alloc(usize, np); + defer a.free(idx); + for (0..np) |i| idx[i] = i; + const sb = scores[r * np .. (r + 1) * np]; + const S = struct { + sb: []const f32, + fn lessThan(self: @This(), x: usize, y: usize) bool { + const bx = orderKey(self.sb[x]); + const by = orderKey(self.sb[y]); + return if (bx == by) x < y else bx > by; // best first, ties to the lower pool + } + }; + std.mem.sort(usize, idx, S{ .sb = sb }, S.lessThan); + std.mem.sort(usize, idx[0..k], {}, std.sort.asc(usize)); // Python's select returns pools ascending + const q = pos + @as(i64, @intCast(r)); + const npool = @divTrunc(q + 1, 4); + for (0..k) |i| { + const p: i64 = @intCast(idx[i]); + for (0..pool) |j| tokens[r * 2051 + i * pool + j] = @intCast(p * pool + @as(i64, @intCast(j))); + } + var tail_ok: i32 = 0; + for (0..3) |j| { + const tail = npool * 4 + @as(i64, @intCast(j)); + tokens[r * 2051 + 2048 + j] = if (tail <= q) @intCast(tail) else -1; + if (tail <= q) tail_ok += 1; + } + counts[r] = if (npool > K) 2048 + tail_ok else 0; + } +} + +/// fp8 e4m3fn (finite, no inf) as f64; 0x7F/0xFF are NaN. +pub fn fp8e4m3(code: u8) f64 { + const s: f64 = if (code & 0x80 != 0) -1.0 else 1.0; + const e: u4 = @truncate(code >> 3); + const m: u3 = @truncate(code); + return if (e == 15 and m == 7) + std.math.nan(f64) + else if (e == 0) + s * @as(f64, @floatFromInt(m)) * 0.001953125 // subnormal: mantissa / 8 * 2^-6 = m * 2^-9 + else + s * (1.0 + @as(f64, @floatFromInt(m)) / 8.0) * std.math.pow(f64, 2.0, @as(f64, @floatFromInt(e)) - 7.0); +} + +/// float64 sparse attention: the selected tokens' fp8 latent rows, a plain running softmax per (row, head). +pub fn sparseRef(a: std.mem.Allocator, qa: []const u16, cache: []const u8, scales: []const f32, tokens: []const i32, counts: []const i32, heads: usize, lw: usize, scale: f64) ![]f64 { + const rows = counts.len; + const out = try a.alloc(f64, rows * heads * lw); + @memset(out, 0); // rows with count 0 stay untouched by the kernel: the test compares the rest + var r: usize = 0; + while (r < rows) : (r += 1) { + if (counts[r] == 0) continue; // the kernel leaves these rows' out untouched + var h: usize = 0; + while (h < heads) : (h += 1) { + var m: f64 = -std.math.inf(f64); + var l: f64 = 0; + const o = try a.alloc(f64, lw); + defer a.free(o); + @memset(o, 0); + for (tokens[r * 2051 ..][0..2051]) |tok| { + if (tok < 0) continue; + const t: usize = @intCast(tok); + const qrow = qa[(r * heads + h) * lw ..][0..lw]; + var dot: f64 = 0; + for (0..lw) |i| dot += kref.bf(qrow[i]) * (fp8e4m3(cache[t * lw + i]) * scales[t]); + const s = dot * scale; + const nm = @max(m, s); + const al = if (m == -std.math.inf(f64)) 0 else @exp(m - nm); + const p = @exp(s - nm); + const krow = cache[t * lw ..][0..lw]; + for (0..lw) |i| o[i] = o[i] * al + p * (fp8e4m3(krow[i]) * scales[t]); + l = l * al + p; + m = nm; + } + if (l > 0) { + for (0..lw) |i| out[(r * heads + h) * lw + i] = o[i] / l; + } else { + for (0..lw) |i| out[(r * heads + h) * lw + i] = 0; + } + } + } + return out; +} + +fn orderKey(s: f32) u32 { + const b: u32 = @bitCast(s + 0.0); + const arith: u32 = 0 -% (b >> 31); // 0 or 0xFFFFFFFF: the sign's arithmetic extension + return b ^ (arith | 0x80000000); +} diff --git a/zig/tests/cuda/glm_sparse_tests.zig b/zig/tests/cuda/glm_sparse_tests.zig new file mode 100644 index 0000000000..a871fa2bc2 --- /dev/null +++ b/zig/tests/cuda/glm_sparse_tests.zig @@ -0,0 +1,206 @@ +//! GLM-5.3-Flash's DSA sparse path on CUDA: indexer and pool scores against float64 refs and Python engine bytes. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const ref = @import("glm_kda_ref.zig"); +const sref = @import("glm_sparse_ref.zig"); +const ktests = @import("glm_kda_tests.zig"); +const sparse = cuda.glm_sparse; + +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const T = ktests.T; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + index(r, state, dir) catch |e| state.fail(e, "glm.dsa.indexer"); + scores(r, state, dir) catch |e| state.fail(e, "glm.dsa.scores"); +} + +fn skip(id: []const u8) void { + std.debug.print("SKIP {s}\n", .{id}); +} + +/// The fixture's sparse params (shared by every check). +const Params = struct { + scale: f64, + wscale: f64, + + fn read(o: std.json.ObjectMap) !Params { + const p = o.get("params").?.object; + return .{ + .scale = ref.float(p, "scale"), + .wscale = ref.float(p, "wscale"), + }; + } +}; + +/// The index update: LayerNorm'd keys and gates into the window's cache rows, then the completed pools' keys. +fn index(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + _ = state; + skip("glm.dsa.indexer"); + return; + } + var man = try sref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + const fam = man.root().get("sparse") orelse { + skip("glm.dsa.indexer"); + return; + }; + _ = try Params.read(fam.object); + var k = try sparse.Sparse.load(r.gpu.d); + defer k.unload(); + var t = T{}; + var it = fam.object.get("cases").?.object.iterator(); + while (it.next()) |kv| { + const e = kv.value_ptr.object; + const tag = kv.key_ptr.*; + const pos: i32 = @intCast(ref.int(e, "pos")); + const rows: usize = @intCast(ref.int(e, "R")); + const ikr = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "ikr")); + const igr = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "igr")); + const lnw = try sref.readFile(r.a, r.gpu.io, dir, "ln_w.bin"); + const lnb = try sref.readFile(r.a, r.gpu.io, dir, "ln_b.bin"); + const ape = try sref.readFile(r.a, r.gpu.io, dir, "ape.bin"); + const want_ik = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "ik")); + const want_ig = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "ig")); + const want_pk = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "pk")); + // the kernels run from the pre-state: rows before the window hold earlier windows' values + const ik = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "ik_pre"))); + const ig = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "ig_pre"))); + const pk = try r.dev(try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "pk_pre"))); + try k.indexWrite(&r.s, try r.dev(ikr), 160, try r.dev(igr), try r.dev(lnw), try r.dev(lnb), ik, ig, pos, 0, 1.0e-6, @intCast(rows)); + try k.poolKeys(&r.s, ik, ig, try r.dev(ape), pk, pos, 0, @intCast(rows)); + const got_ik = try r.back(ik, want_ik.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.indexer {s} ik", .{tag}), got_ik, want_ik); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.indexer {s} ig", .{tag}), try r.back(ig, want_ig.len), want_ig); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.indexer {s} pk", .{tag}), try r.back(pk, want_pk.len), want_pk); + // float64: the LayerNorm rows and the window's completed pools + const want64 = try sref.indexRef(r.a, u16s(ikr), 160, u16s(lnw), u16s(lnb), 1.0e-6, rows); + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.indexer {s} ik f64", .{tag}), got_ik[@as(usize, @intCast(pos)) * sref.D * 2 ..], want64, false, 1.0 / 128.0, 2.0e-3); + const npool_end = @divTrunc(@as(i64, pos) + @as(i64, @intCast(rows)), 4); + const want_pool = try sref.poolRef(r.a, u16s(want_ig), u16s(ape), u16s(want_ik), @intCast(@divTrunc(pos, 4)), @intCast(npool_end)); + const got_pk = try r.back(pk, want_pk.len); + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.indexer {s} pk f64", .{tag}), got_pk[@as(usize, @intCast(@divTrunc(pos, 4))) * sref.D * 2 ..], want_pool, false, 1.0 / 64.0, 2.0e-3); + } + std.debug.print("RESULT glm.dsa.indexer: f64 worst {d:.4} of tolerance\n", .{t.worst}); + check.pass("glm.dsa.indexer", .{}); +} + +/// The pool scores: every visible pool per row, one block launch per 512 rows, then the float64 reference. +fn scores(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + _ = state; + skip("glm.dsa.scores"); + return; + } + var man = try sref.Manifest.open(r.a, r.gpu.io, dir); + defer man.deinit(); + const fam = man.root().get("sparse") orelse { + skip("glm.dsa.scores"); + return; + }; + const q = try Params.read(fam.object); + var k = try sparse.Sparse.load(r.gpu.d); + defer k.unload(); + var t = T{}; + var it = fam.object.get("cases").?.object.iterator(); + while (it.next()) |kv| { + const e = kv.value_ptr.object; + const tag = kv.key_ptr.*; + const pos: i32 = @intCast(ref.int(e, "pos")); + const rows: usize = @intCast(ref.int(e, "R")); + const np_max: usize = @intCast(ref.int(e, "np_max")); + const qi = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "qi")); + const ikr = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "ikr")); + const pk = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "pk")); + const qi_dev = try r.dev(qi); + const w_base = try r.dev(ikr); // the weights are ikr's last 32 columns, per block row + const pk_dev = try r.dev(pk); + const blocks = e.get("blocks").?.array; + for (blocks.items, 0..) |bname, bi| { + const b0 = bi * 512; + const n = @min(@as(usize, 512), rows - b0); + const want = try sref.readFile(r.a, r.gpu.io, dir, bname.string); + const out = try r.zeros(want.len); + try k.scores(&r.s, qi_dev + b0 * 4096 * 2, 4096, w_base + (b0 * 160 + sref.D) * 2, 160, pk_dev, out, pos + @as(i32, @intCast(b0)), 0, @intCast(n), @intCast(np_max), @floatCast(q.scale), @floatCast(q.wscale)); + const got = try r.back(out, want.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.scores {s} b{d}", .{ tag, b0 }), got, want); + if (rows <= 128) { + const want64 = try sref.scoresRef(r.a, u16s(qi)[b0 * 4096 ..], 4096, u16s(ikr)[b0 * 160 + sref.D ..], 160, u16s(pk), n, np_max, q.scale, q.wscale, pos + @as(i32, @intCast(b0))); + try t.close(try std.fmt.allocPrint(r.a, "glm.dsa.scores {s} f64", .{tag}), got, want64, true, 1.0 / 64.0, 2.0e-3); + } + } + } + std.debug.print("RESULT glm.dsa.scores: score f64 worst {d:.4} of tolerance\n", .{t.worst}); + try selectPhase(r, fam, q, dir); + check.pass("glm.dsa.scores", .{}); +} + +/// The selection phase: scores over the row's visible bucket, the radix select, then tokens and counts. +fn selectPhase(r: *Rig, fam: std.json.Value, q: Params, dir: []const u8) !void { + var k = try sparse.Sparse.load(r.gpu.d); + defer k.unload(); + var it = fam.object.get("cases").?.object.iterator(); + while (it.next()) |kv| { + const e = kv.value_ptr.object; + const tag = kv.key_ptr.*; + const pos: i32 = @intCast(ref.int(e, "pos")); + const rows: usize = @intCast(ref.int(e, "R")); + const np_max: usize = @intCast(ref.int(e, "np_max")); + const visible = @divTrunc(pos + @as(i32, @intCast(rows)), 4) + 1; + var bucket: usize = 1; + var v: u32 = @intCast(visible - 1); + while (v != 0) : (v >>= 1) bucket <<= 1; + bucket = @max(bucket, 1024); + bucket = @min(bucket, np_max); + const qi = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "qi")); + const ikr = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "ikr")); + const pk = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "pk")); + const qi_dev = try r.dev(qi); + const w_base = try r.dev(ikr); + const pk_dev = try r.dev(pk); + const sc = try r.zeros(rows * bucket * 4); + var b0: usize = 0; + while (b0 < rows) : (b0 += 512) { + const n = @min(@as(usize, 512), rows - b0); + try k.scores(&r.s, qi_dev + b0 * 4096 * 2, 4096, w_base + (b0 * 160 + sref.D) * 2, 160, pk_dev, sc + b0 * bucket * 4, pos + @as(i32, @intCast(b0)), 0, @intCast(n), @intCast(bucket), @floatCast(q.scale), @floatCast(q.wscale)); + } + const pools = try r.zeros(rows * 512 * 8); + try k.selectRows(&r.s, sc, @intCast(bucket), pools, pos, 0, 512, @intCast(rows), true); + const tokens = try r.zeros(rows * 2051 * 4); + const counts = try r.zeros(rows * 4); + try k.tokensCounts(&r.s, pools, tokens, counts, pos, 0, @intCast(rows)); + const want_tokens = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "tokens")); + const want_counts = try sref.readFile(r.a, r.gpu.io, dir, ref.str(e, "counts")); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.scores {s} tokens", .{tag}), try r.back(tokens, want_tokens.len), want_tokens); + try check.sameBytes(try std.fmt.allocPrint(r.a, "glm.dsa.scores {s} counts", .{tag}), try r.back(counts, want_counts.len), want_counts); + // float64-exact reference: the order-key sort of the fp32 scores + if (rows <= 128) { + const got_tokens = try r.back(tokens, want_tokens.len); + const got_counts = try r.back(counts, want_counts.len); + const got_scores_raw = try r.back(sc, rows * bucket * 4); + const got_scores = try r.a.alloc(f32, rows * bucket); + @memcpy(std.mem.sliceAsBytes(got_scores), got_scores_raw); + const ref_tokens = try r.a.alloc(i32, rows * 2051); + const ref_counts = try r.a.alloc(i32, rows); + try sref.selectRef(r.a, got_scores, bucket, 512, pos, rows, ref_tokens, ref_counts); + for (ref_tokens, 0..) |want_v, i| { + const got_v = std.mem.readInt(i32, got_tokens[4 * i ..][0..4], .little); + if (got_v != want_v) return error.TestFailed; + } + for (ref_counts, 0..) |want_v, i| { + const got_v = std.mem.readInt(i32, got_counts[4 * i ..][0..4], .little); + if (got_v != want_v) return error.TestFailed; + } + } + } + std.debug.print("RESULT glm.dsa.scores: tokens exact against the order-key sort\n", .{}); +} + +fn u16s(bytes: []const u8) []const u16 { + return @alignCast(std.mem.bytesAsSlice(u16, bytes)); +} diff --git a/zig/tests/cuda/glm_tests.zig b/zig/tests/cuda/glm_tests.zig new file mode 100644 index 0000000000..7a4dbd19f4 --- /dev/null +++ b/zig/tests/cuda/glm_tests.zig @@ -0,0 +1,34 @@ +//! GLM-5.3-Flash on CUDA: each kernel family against float64 and the Python engine's bytes, then whole layers. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); + +pub const Rig = quant_ref_tests.RigT; +pub const Run = exl3_tests.Run; + +/// Every family's checks, in the order a decode token meets them; `dir` none prints SKIP for fixture checks. +pub fn run(gpu: check.Gpu, dir: []const u8) !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(17); + 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 state = Run{}; + try @import("glm_glue_tests.zig").run(&r, &state, dir); + try @import("glm_kda_tests.zig").run(&r, &state, dir); + try @import("glm_sparse_tests.zig").run(&r, &state, dir); + try @import("glm_latent_tests.zig").run(&r, &state, dir); + try @import("glm_layer_tests.zig").run(&r, &state, dir); + try @import("glm_mlp_tests.zig").run(&r, &state, dir); + try @import("glm_dsa_block_tests.zig").run(&r, &state, dir); + try @import("glm_layer_dsa_tests.zig").run(&r, &state, dir); + try @import("glm_head_tests.zig").run(&r, &state, dir); + if (state.failed) return error.TestFailed; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 7b8a089947..f3314741e4 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -22,6 +22,10 @@ const quant_ref_tests = @import("quant_ref_tests.zig"); const carveout_tests = @import("carveout_tests.zig"); const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const exl3_tp2_tests = @import("exl3_tp2_tests.zig"); +const glm_tests = @import("glm_tests.zig"); +const glm_model_tests = @import("glm_model_tests.zig"); const train_pack = @import("train_pack.zig"); const exlr8_tests = @import("exlr8_tests.zig"); const exlr8_dense_tests = @import("exlr8_dense_tests.zig"); @@ -64,6 +68,10 @@ const usage = \\ exlr8 EXLR8 projections: every codebook state, FP64, row windows, the Metal frozen digests \\ 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 + \\ exl3 [dir] EXL3 decode-row linear: float64 checks, row invariance, the Python engine's fixture (dir none skips it) + \\ exl3-tp2 [dir] EXL3 tensor parallel: rank shards vs the Python engine's, all-gathered (dir none skips them) + \\ glm [dir] GLM-5.3-Flash kernels and layers vs float64 and the Python engine (dir none skips those) + \\ glm-model MODEL REF the whole GLM-5.3-Flash at TP=2 vs the Python engine's greedy replies \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ ; @@ -123,6 +131,10 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { 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); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); + if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else "none"); + if (std.mem.eql(u8, cmd, "exl3-tp2")) return exl3_tp2_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else "none"); + if (std.mem.eql(u8, cmd, "glm")) return glm_tests.run(gpu, if (rest.len > 0) rest[0] else "none"); + if (std.mem.eql(u8, cmd, "glm-model")) return glm_model_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); if (std.mem.eql(u8, cmd, "norms")) return norms_tests.run(gpu); if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu); diff --git a/zig/tests/cuda/oracle/exl3_oracle.py b/zig/tests/cuda/oracle/exl3_oracle.py new file mode 100644 index 0000000000..17374dd307 --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle.py @@ -0,0 +1,124 @@ +"""Fixtures for `tf-cuda-test exl3`: the Python engine's EXL3 linear outputs on real RED-SNOW tensors.""" +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +# the other fixtures: siblings in this folder, which leads sys.path when the script runs +import exl3_oracle_dense +import exl3_oracle_experts +import exl3_oracle_layer +import exl3_oracle_tp2 + +ROWS = (1, 3, 16) +SEED = 1234 +SIZE_LIMIT = 8192 # prefer layers with N, K at or under this + + +def rank(t, bits: float) -> int: + """Which layers to prefer for a width: a gate_proj for 4..6 bits, a shared-expert or attention one for 7..8.""" + if bits <= 6 and "gate_proj" in t.prefix: + return 0 + if bits >= 7 and ("shared_experts" in t.prefix or "self_attn" in t.prefix): + return 0 + return 1 + + +def pick(model_dir: Path) -> list: + """One EXL3 linear per whole bits value 4..8, first by sorted name among the preferred ones.""" + from tensorfold.cuda.exl3 import format as fmt + + cp = fmt.scan(model_dir) + chosen = [] + for bits in (4.0, 5.0, 6.0, 7.0, 8.0): + cands = [t for t in cp.groups.values() + if t.bits == bits and t.in_scales == "suh" and t.out_scales == "svh"] + if not cands: + print(f"no EXL3 linear at {bits:g} bits in {model_dir}") + continue + modest = [t for t in cands if t.k <= SIZE_LIMIT and t.n <= SIZE_LIMIT] + t = min(modest or cands, key=lambda t: (rank(t, bits), t.prefix)) + marker = cp.markers.get(f"{t.prefix}.{t.codebook}") + if t.codebook in fmt.MARKERS and marker is not None and marker != fmt.MARKERS[t.codebook]: + raise ValueError(f"{t.prefix}: {t.codebook} marker is 0x{marker:08X}, want 0x{fmt.MARKERS[t.codebook]:08X}") + chosen.append((t, marker)) + return chosen + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def one(out: Path, model_dir: Path, t, marker: int | None) -> dict: + from tensorfold.cuda.exl3 import format as fmt + from tensorfold.cuda.exl3.linear import Exl3Linear + + lid = f"python.{t.bits:g}" + # "stored" keeps the checkpoint's trellis word order; the forward path is identical either way + lin = Exl3Linear.load(model_dir, t.prefix, layout="stored") + g = torch.Generator().manual_seed(SEED) + x = torch.randn((max(ROWS), lin.k), generator=g).to(torch.bfloat16) + + save(out, f"{lid}.trellis", lin.words) + save(out, f"{lid}.suh", lin.suh) + save(out, f"{lid}.svh", lin.svh) + if lin.bias is not None: + save(out, f"{lid}.bias", lin.bias) + + ys, xfiles, yfiles, sha = [], {}, {}, {} + for m in ROWS: + y = lin(x[:m].cuda()) + torch.cuda.synchronize() + save(out, f"{lid}.x{m}", x[:m]) + save(out, f"{lid}.y{m}", y) + xfiles[str(m)] = f"{lid}.x{m}.bin" + yfiles[str(m)] = f"{lid}.y{m}.bin" + sha[str(m)] = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + ys.append(y) + + rel, ref = 0.0, fmt.forward(x.float().numpy().astype(np.float64), lin.words.view(torch.int16), lin.suh, + lin.svh, t.bits, t.codebook, lin.bias) + for m, y in zip(ROWS, ys): + want = ref[:m] + rel = max(rel, np.linalg.norm(y.cpu().float().numpy().astype(np.float64) - want) + / max(np.linalg.norm(want), 1e-30)) + + entry = {"id": lid, "name": t.prefix, "bits": int(t.bits), "codebook": t.codebook, + "k": lin.k, "n": lin.n, "rows": list(ROWS), "trellis": f"{lid}.trellis.bin", + "suh": f"{lid}.suh.bin", "svh": f"{lid}.svh.bin", + "bias": None if lin.bias is None else f"{lid}.bias.bin", + "x": xfiles, "y": yfiles, "x_dtype": "bf16", "y_dtype": "bf16", "sha256": sha} + if marker is not None: + entry["marker"] = f"0x{marker:08X}" + print(f"{lid}: {t.prefix} EXL3 {t.bits:g}-bit {t.codebook} K={lin.k} N={lin.n} " + f"rel={rel:.2e} sha256={sha['16'][:16]}…") + return entry + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.partition("\n")[0]) + ap.add_argument("--model", required=True, help="the EXL3 checkpoint folder") + ap.add_argument("--out", required=True, help="the fixture folder to (over)write") + a = ap.parse_args() + out, model_dir = Path(a.out), Path(a.model) + out.mkdir(parents=True, exist_ok=True) + linears = [one(out, model_dir, t, marker) for t, marker in pick(model_dir)] + dense = exl3_oracle_dense.write(out, model_dir) + experts = exl3_oracle_experts.write(out, model_dir) + layers = exl3_oracle_layer.write(out, model_dir) + tp2 = exl3_oracle_tp2.write(out, model_dir) + manifest = {"version": 1, "linears": linears, "dense": dense, "experts": experts, "layers": layers, "tp2": tp2, + "model": str(model_dir)} + (out / "manifest.json").write_text(json.dumps(manifest, indent=1) + "\n") + print(f"{len(linears)} linears -> {out / 'manifest.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/oracle/exl3_oracle_dense.py b/zig/tests/cuda/oracle/exl3_oracle_dense.py new file mode 100644 index 0000000000..afb5b620e8 --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle_dense.py @@ -0,0 +1,67 @@ +"""The dense prompt fixture: one real projection's 300-row prompt chunk through forward.py's exl3_mm path.""" +from __future__ import annotations + +import hashlib +from pathlib import Path + +import numpy as np +import torch + +ROWS = 300 +SEED = 1235 +SIZE_LIMIT = 8192 # prefer layers with N, K at or under this +PREFER = ("self_attn.qkv_proj", "shared_experts.gate_proj", "self_attn.q_b_proj") + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.cuda.exl3 import format as fmt + from tensorfold.cuda.exl3.linear import Exl3Linear + + cp = fmt.scan(model_dir) + cands = [t for t in cp.groups.values() + if t.bits in (4.0, 5.0, 6.0, 7.0, 8.0) and t.in_scales == "suh" and t.out_scales == "svh"] + + def pref(t): + rank = next((i for i, pat in enumerate(PREFER) if pat in t.prefix), len(PREFER)) + return (rank, t.k * t.n) + + modest = [t for t in cands if t.k <= SIZE_LIMIT and t.n <= SIZE_LIMIT] + t = min(modest or cands, key=pref) + marker = cp.markers.get(f"{t.prefix}.{t.codebook}") + if t.codebook in fmt.MARKERS and marker is not None and marker != fmt.MARKERS[t.codebook]: + raise ValueError(f"{t.prefix}: {t.codebook} marker is 0x{marker:08X}, want 0x{fmt.MARKERS[t.codebook]:08X}") + lin = Exl3Linear.load(model_dir, t.prefix, layout="stored") + g = torch.Generator().manual_seed(SEED) + x = torch.randn((ROWS, lin.k), generator=g).to(torch.bfloat16) + + # forward.py's exl3_mm: the decode-row linear over 128-row chunks, a fresh output tile each chunk + y = torch.empty((ROWS, lin.n), dtype=torch.bfloat16) + o = torch.empty((min(ROWS, 128), lin.n), dtype=torch.bfloat16, device="cuda") + for r in range(0, ROWS, 128): + e = min(ROWS - r, 128) + lin(x[r:r + e].cuda(), out=o[:e]) + torch.cuda.synchronize() + y[r:r + e] = o[:e].cpu() + + lid = "dense" + save(out, f"{lid}.trellis", lin.words) + save(out, f"{lid}.suh", lin.suh) + save(out, f"{lid}.svh", lin.svh) + if lin.bias is not None: + save(out, f"{lid}.bias", lin.bias) + save(out, f"{lid}.x", x) + save(out, f"{lid}.y", y) + sha = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + entry = {"id": lid, "name": t.prefix, "bits": int(t.bits), "codebook": t.codebook, + "k": lin.k, "n": lin.n, "rows": ROWS, "trellis": f"{lid}.trellis.bin", + "suh": f"{lid}.suh.bin", "svh": f"{lid}.svh.bin", + "bias": None if lin.bias is None else f"{lid}.bias.bin", + "x": f"{lid}.x.bin", "y": f"{lid}.y.bin", "x_dtype": "bf16", "y_dtype": "bf16", "sha256": sha} + if marker is not None: + entry["marker"] = f"0x{marker:08X}" + print(f"{lid}: {t.prefix} EXL3 {t.bits:g}-bit {t.codebook} K={lin.k} N={lin.n} sha256={sha[:16]}…") + return entry diff --git a/zig/tests/cuda/oracle/exl3_oracle_experts.py b/zig/tests/cuda/oracle/exl3_oracle_experts.py new file mode 100644 index 0000000000..54acb4a9a9 --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle_experts.py @@ -0,0 +1,115 @@ +"""The routed-experts fixture: nine real experts of one MoE layer through the Python engine's generic.routed.""" +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path + +import torch + +ROWS = (1, 3, 16) +SEED = 1236 +SLOTS = 8 +LIMIT = 10.0 +POOL = 3 # experts per width + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def load_group(model_dir: Path, prefix: str) -> tuple: + """(trellis, suh, svh) of one expert projection, on the GPU (the checkpoint's stored trellis order).""" + from safetensors import safe_open + + root = Path(model_dir) + index = root / "model.safetensors.index.json" + rel = None + if index.exists(): + weight_map = json.loads(index.read_text())["weight_map"] + rel = weight_map.get(f"{prefix}.trellis") + if rel is None: + for part in sorted(root.glob("*.safetensors")): + with safe_open(str(part), framework="pt") as g: + if f"{prefix}.trellis" in set(g.keys()): + rel = part.name + break + if rel is None: + raise ValueError(f"{prefix}.trellis is in no file of {root}") + file = root / (rel or "model.safetensors") + with safe_open(str(file), framework="pt", device="cuda") as f: + return (f.get_tensor(f"{prefix}.trellis"), f.get_tensor(f"{prefix}.suh"), f.get_tensor(f"{prefix}.svh")) + + +def pick_layer(cp) -> str: + """The first layer (sorted) whose routed experts cover 4, 5 and 6 bits.""" + layers: dict[str, dict[float, list[int]]] = {} + for prefix, t in cp.groups.items(): + m = re.match(r"^(.*)\.mlp\.experts\.(\d+)\.gate_proj$", prefix) + if m is not None and t.in_scales == "suh" and t.out_scales == "svh": + layers.setdefault(m.group(1), {}).setdefault(t.bits, []).append(int(m.group(2))) + for base in sorted(layers): + if all(len(layers[base].get(b, ())) >= POOL for b in (4.0, 5.0, 6.0)): + return base + raise ValueError("no MoE layer has 4, 5 and 6-bit routed experts") + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.cuda.exl3 import experts as generic + from tensorfold.cuda.exl3 import format as fmt + + cp = fmt.scan(model_dir) + base = pick_layer(cp) + + by_bits: dict[float, list[int]] = {} + for prefix, t in cp.groups.items(): + m = re.match(rf"^{re.escape(base)}\.mlp\.experts\.(\d+)\.gate_proj$", prefix) + if m is not None: + by_bits.setdefault(t.bits, []).append(int(m.group(1))) + pool = [e for b in (4.0, 5.0, 6.0) for e in sorted(by_bits[b])[:POOL]] + + gate, up, down = [], [], [] + trellis = {"gate": [], "up": [], "down": []} + for slot, e in enumerate(pool): + for part, short, mats, names in (("gate_proj", "gate", gate, trellis["gate"]), + ("up_proj", "up", up, trellis["up"]), + ("down_proj", "down", down, trellis["down"])): + mats.append(load_group(model_dir, f"{base}.mlp.experts.{e}.{part}")) + names.append(f"experts.{short}{slot}.bin") + ex = generic.prepare(gate, up, down, "mul1", device="cuda") + + suh_files = {} + for key, t in (("suh_g", ex.suh_g), ("suh_u", ex.suh_u), ("svh_g", ex.svh_g), + ("svh_u", ex.svh_u), ("suh_d", ex.suh_d), ("svh_d", ex.svh_d)): + save(out, f"experts.{key}", t) + suh_files[key] = f"experts.{key}.bin" + for slot, mats in enumerate(zip(gate, up, down)): + for name, t in zip(("gate", "up", "down"), mats): + save(out, f"experts.{name}{slot}", t[0]) + + g = torch.Generator().manual_seed(SEED) + x = torch.randn((max(ROWS), ex.dims), generator=g).to(torch.bfloat16) + pick = torch.stack([torch.randperm(ex.count, generator=g)[:SLOTS] for _ in range(max(ROWS))]).int() + wts = (0.1 + 0.9 * torch.rand((max(ROWS), SLOTS), generator=g)).float() + + scratch = generic.Scratch(ex, max(ROWS), SLOTS) + xmap, pickmap, wtsmap, ymap, outmap, sha = {}, {}, {}, {}, {}, {} + for m in ROWS: + y = generic.routed(x[:m].cuda(), pick[:m].cuda(), wts[:m].cuda(), ex, scratch, None, m, + limit=LIMIT, act_mode=generic.ACT_BF16) + torch.cuda.synchronize() + saved = {"x": x[:m], "pick": pick[:m], "wts": wts[:m], "y": scratch.y[:m * SLOTS], "out": y} + maps = {"x": xmap, "pick": pickmap, "wts": wtsmap, "y": ymap, "out": outmap} + for key, t in saved.items(): + save(out, f"experts.{key}{m}", t) + maps[key][str(m)] = f"experts.{key}{m}.bin" + sha[str(m)] = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + + entry = {"id": "experts", "name": base, "count": ex.count, "dims": ex.dims, "width": ex.width, + "slots": SLOTS, "codebook": "mul1", "limit": LIMIT, "act_mode": generic.ACT_BF16, + "rows": list(ROWS), "gate_k2": [generic.k2_of(m[0][0]) for m in gate], + "up_k2": [generic.k2_of(m[0]) for m in up], "down_k2": [generic.k2_of(m[0]) for m in down], + "trellis": trellis, **suh_files, "x": xmap, "pick": pickmap, "wts": wtsmap, + "y": ymap, "out": outmap, "sha256": sha} + print(f"experts: {base} {len(pool)} experts (4/5/6-bit pool), rows {ROWS}, sha256={sha['16'][:16]}…") + return entry diff --git a/zig/tests/cuda/oracle/exl3_oracle_layer.py b/zig/tests/cuda/oracle/exl3_oracle_layer.py new file mode 100644 index 0000000000..469f12e0ba --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle_layer.py @@ -0,0 +1,170 @@ +"""The whole-layer fixtures: every EXL3 group of one KDA and one DSA MoE layer through the Python engine.""" +from __future__ import annotations + +import hashlib +import json +import zlib +from pathlib import Path + +import torch + +ROWS = (1, 16, 300) # the dense groups' row counts +EXPERT_ROWS = (1, 16) +SEED = 1239 # + layer index + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def sha_of(y: torch.Tensor) -> str: + return hashlib.sha256(y.cpu().contiguous().view(torch.uint8).numpy().tobytes()).hexdigest() + + +def load_tensor(model_dir: Path, name: str, device: str = "cpu") -> torch.Tensor: + from safetensors import safe_open + + root = Path(model_dir) + index = root / "model.safetensors.index.json" + if index.exists(): + rel = json.loads(index.read_text())["weight_map"][name] + else: + rel = "model.safetensors" + with safe_open(str(root / rel), framework="pt", device=device) as f: + return f.get_tensor(name) + + +def mm_rows(x: torch.Tensor, lin, m: int) -> torch.Tensor: + """lin on m rows: one call to 128, a prompt chunk beyond that as 128-row blocks (forward.py's exl3_mm).""" + if m <= 128: + y = lin(x[:m].cuda()) + torch.cuda.synchronize() + return y + y = torch.empty((m, lin.n), dtype=torch.bfloat16) + o = torch.empty((min(m, 128), lin.n), dtype=torch.bfloat16, device="cuda") + for r in range(0, m, 128): + e = min(m - r, 128) + lin(x[r:r + e].cuda(), out=o[:e]) + torch.cuda.synchronize() + y[r:r + e] = o[:e].cpu() + return y + + +def lid_of(base: str) -> str: + return base.rsplit(".", 1)[1] + + +def pick_layers(model_dir: Path) -> list[tuple[str, int]]: + """(kind, layer index) of the first KDA and the first DSA MoE layer.""" + cfg = json.loads((Path(model_dir) / "config.json").read_text())["text_config"] + keys = {"kda": "linear_attention", "dsa": "deepseek_sparse_attention"} + out = [] + for kind, key in keys.items(): + for i, (attn, mlp) in enumerate(zip(cfg["layer_types"], cfg["mlp_layer_types"])): + if attn == key and mlp == "sparse": + out.append((kind, i)) + break + else: + raise ValueError(f"no {kind} MoE layer in {model_dir}") + return out + + +def one_dense(out: Path, model_dir: Path, cp, prefix: str, base: str, hidden: torch.Tensor) -> dict: + """One non-routed group's inputs and the engine's outputs, bit for bit.""" + from tensorfold.cuda.exl3 import linear + + t = cp.groups[prefix] + lin = linear.Exl3Linear.load(model_dir, prefix, layout="stored") + gid = prefix[len(base) + 1:] + if lin.k == hidden.shape[1]: + x = hidden + else: + g = torch.Generator().manual_seed(SEED + zlib.crc32(gid.encode())) + x = torch.randn((max(ROWS), lin.k), generator=g).to(torch.bfloat16) + xmap, ymap, sha = {}, {}, {} + for m in ROWS: + y = mm_rows(x, lin, m) + save(out, f"{lid_of(base)}.{gid}.x{m}", x[:m]) + save(out, f"{lid_of(base)}.{gid}.y{m}", y) + xmap[str(m)] = f"{lid_of(base)}.{gid}.x{m}.bin" + ymap[str(m)] = f"{lid_of(base)}.{gid}.y{m}.bin" + sha[str(m)] = sha_of(y) + entry = {"id": gid, "name": prefix, "bits": int(t.bits), "codebook": t.codebook, + "k": lin.k, "n": lin.n, "rows": list(ROWS), "x": xmap, "y": ymap, "sha256": sha} + print(f" {gid}: EXL3 {t.bits:g}-bit {t.codebook} K={lin.k} N={lin.n}") + return entry + + +def expert_part(out: Path, model_dir: Path, cp, base: str, hidden: torch.Tensor, cfg: dict) -> dict: + """All experts as one grouped call with the engine's routing, bit for bit.""" + from tensorfold.cuda.exl3 import experts as generic + from tensorfold.families.glm5_next.cuda import glue + + codebook = next(t.codebook for p, t in cp.groups.items() if p.startswith(f"{base}.mlp.experts.0.")) + gate, up, down = [], [], [] + k2 = {"gate": [], "up": [], "down": []} + for e in range(cfg["n_routed_experts"]): + for part, mats in (("gate_proj", gate), ("up_proj", up), ("down_proj", down)): + mats.append(tuple(load_tensor(model_dir, f"{base}.mlp.experts.{e}.{part}.{x}", "cuda") + for x in ("trellis", "suh", "svh"))) + for mats, key in ((gate, "gate"), (up, "up"), (down, "down")): + k2[key] = [generic.k2_of(m[0]) for m in mats] + ex = generic.prepare(gate, up, down, codebook, device="cuda") + nbytes = sum(m[0].element_size() * m[0].numel() for mats in (gate, up, down) for m in mats) \ + + sum(t.element_size() * t.numel() for t in (ex.suh_g, ex.suh_u, ex.svh_g, ex.svh_u, ex.suh_d, ex.svh_d)) + print(f" experts: {ex.count} experts, {ex.dims}x{ex.width}, {nbytes / (1 << 30):.2f} GiB") + + # the routing, as the GLM engine's moe_block routes: glue.router on the bf16 gate weight, glue.select on the bias + gate_w = load_tensor(model_dir, f"{base}.mlp.gate.weight", "cuda").to(torch.bfloat16).contiguous() + bias = load_tensor(model_dir, f"{base}.mlp.gate.e_score_correction_bias", "cuda").float().contiguous() + logits = torch.empty((hidden.shape[0], ex.count), dtype=torch.float32, device="cuda") + glue.router(hidden.cuda(), gate_w, logits) + topk = cfg["num_experts_per_tok"] + pick = torch.empty((hidden.shape[0], topk + 1), dtype=torch.int32, device="cuda") + wts = torch.empty((hidden.shape[0], topk + 1), dtype=torch.float32, device="cuda") + glue.select(logits, bias, pick, wts, topk, ex.count, cfg["routed_scaling_factor"], + cfg.get("norm_topk_prob", True)) + pick, wts = pick[:, :topk].contiguous(), wts[:, :topk].contiguous() + + scratch = generic.Scratch(ex, max(EXPERT_ROWS), topk) + xmap, pmap, wmap, ymap, omap, sha = {}, {}, {}, {}, {}, {} + for m in EXPERT_ROWS: + y = generic.routed(hidden[:m].cuda(), pick[:m].cuda(), wts[:m].cuda(), ex, scratch, None, m, + limit=cfg.get("swiglu_limit", 10.0), act_mode=generic.ACT_BF16) + torch.cuda.synchronize() + for key, t in (("x", hidden[:m]), ("pick", pick[:m]), ("wts", wts[:m]), + ("y", scratch.y[:m * topk]), ("out", y)): + save(out, f"{lid_of(base)}.experts.{key}{m}", t) + xmap[str(m)] = f"{lid_of(base)}.experts.x{m}.bin" + pmap[str(m)] = f"{lid_of(base)}.experts.pick{m}.bin" + wmap[str(m)] = f"{lid_of(base)}.experts.wts{m}.bin" + ymap[str(m)] = f"{lid_of(base)}.experts.y{m}.bin" + omap[str(m)] = f"{lid_of(base)}.experts.out{m}.bin" + sha[str(m)] = sha_of(y) + return {"name": f"{base}.mlp.experts", "count": ex.count, "dims": ex.dims, "width": ex.width, + "slots": topk, "codebook": codebook, "limit": cfg.get("swiglu_limit", 10.0), + "act_mode": generic.ACT_BF16, "rows": list(EXPERT_ROWS), "gate_k2": k2["gate"], + "up_k2": k2["up"], "down_k2": k2["down"], "x": xmap, "pick": pmap, "wts": wmap, + "y": ymap, "out": omap, "sha256": sha} + + +def one_layer(out: Path, model_dir: Path, cp, cfg: dict, li: int, kind: str) -> dict: + base = f"model.language_model.layers.{li}" + print(f"{kind} layer {li}:") + g = torch.Generator().manual_seed(SEED + li) + hidden = torch.randn((max(ROWS), cfg["hidden_size"]), generator=g).to(torch.bfloat16) + groups = [one_dense(out, model_dir, cp, prefix, base, hidden) + for prefix in sorted(cp.groups) + if prefix.startswith(base + ".") and ".experts." not in prefix] + return {"id": kind, "layer": li, "name": base, "dims": cfg["hidden_size"], + "groups": groups, "experts": expert_part(out, model_dir, cp, base, hidden, cfg)} + + +def write(out: Path, model_dir: Path) -> list: + from tensorfold.cuda.exl3 import format as fmt + + cp = fmt.scan(model_dir) + cfg = json.loads((Path(model_dir) / "config.json").read_text())["text_config"] + layers = [one_layer(out, model_dir, cp, cfg, li, kind) for kind, li in pick_layers(model_dir)] + print(f"{len(layers)} layers -> {out}") + return layers diff --git a/zig/tests/cuda/oracle/exl3_oracle_tp2.py b/zig/tests/cuda/oracle/exl3_oracle_tp2.py new file mode 100644 index 0000000000..961df5f5b9 --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle_tp2.py @@ -0,0 +1,232 @@ +"""The tensor-parallel fixture: two real projections and one MoE layer's routed experts, split as the ranks are.""" +from __future__ import annotations + +import hashlib +import re +from pathlib import Path + +import torch + +import exl3_oracle_experts + +ROWS = (1, 3, 16) +SEED = 1237 +EXPERTS_SEED = 1238 +POOL = 3 # experts per width, as the routed fixture's pool +SLOTS = 8 +LIMIT = 10.0 + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def cut_dim1(raw: torch.Tensor, rank: int) -> torch.Tensor: + """A tensor's second axis in half (the rank's tile columns / scale rows), contiguous.""" + half = raw.shape[1] // 2 + return raw[:, rank * half:(rank + 1) * half].contiguous() + + +def cut_row(raw: torch.Tensor, rank: int) -> torch.Tensor: + """A tensor's leading axis in half (the rank's tile rows), contiguous.""" + half = raw.shape[0] // 2 + return raw[rank * half:(rank + 1) * half].contiguous() + + +def load_group(model_dir: Path, prefix: str) -> tuple: + """(trellis int16 as stored, suh, svh) of one linear, from the checkpoint (on the GPU).""" + from safetensors import safe_open + + root = Path(model_dir) + index = root / "model.safetensors.index.json" + rel = None + if index.exists(): + weight_map = __import__("json").loads(index.read_text())["weight_map"] + rel = weight_map.get(f"{prefix}.trellis") + if rel is None: + for part in sorted(root.glob("*.safetensors")): + with safe_open(str(part), framework="pt") as g: + if f"{prefix}.trellis" in set(g.keys()): + rel = part.name + break + if rel is None: + raise ValueError(f"{prefix}.trellis is in no file of {root}") + with safe_open(str(root / (rel or "model.safetensors")), framework="pt", device="cuda") as f: + return (f.get_tensor(f"{prefix}.trellis"), f.get_tensor(f"{prefix}.suh"), f.get_tensor(f"{prefix}.svh")) + + +def pick_out_linear(cp) -> object: + """The out-split layer: a shared_experts.gate_proj (or any gate/up one) whose cut stays on 128 boundaries.""" + cands = [t for t in cp.groups.values() + if re.search(r"\.(mlp\.shared_experts\.)?(gate|up)_proj$", t.prefix) + and t.in_scales == "suh" and t.out_scales == "svh" + and t.k % 256 == 0 and t.n % 256 == 0 and t.k <= 8192 and t.n <= 8192] + if not cands: + raise ValueError("no EXL3 gate/up projection splits on 128-boundaries at world 2") + shared = [t for t in cands if "shared_experts" in t.prefix] + return min(shared or cands, key=lambda t: (t.k * t.n, t.prefix)) + + +def pick_in_linear(cp, out) -> object: + """The in-split layer: the same layer's shared_experts.down_proj, else a down_proj, else an o_proj.""" + base = out.prefix.rsplit(".mlp.", 1)[0] if ".mlp." in out.prefix else out.prefix.rsplit(".self_attn.", 1)[0] + prefs = (f"{base}.mlp.shared_experts.down_proj", f"{base}.mlp.down_proj") + for name in prefs: + t = cp.groups.get(name) + if t is not None and t.in_scales == "suh" and t.out_scales == "svh" and t.k % 256 == 0 and t.n % 256 == 0: + return t + cands = [t for t in cp.groups.values() + if (t.prefix.endswith(".mlp.down_proj") or t.prefix.endswith(".self_attn.o_proj")) + and t.in_scales == "suh" and t.out_scales == "svh" and t.k % 256 == 0 and t.n % 256 == 0 + and t.k <= 8192 and t.n <= 8192] + if not cands: + raise ValueError("no EXL3 down/o projection splits on 128-boundaries at world 2") + return min(cands, key=lambda t: (t.k * t.n, t.prefix)) + + +def linear_entry(out: Path, model_dir: Path, t, split: str, lid: str) -> dict: + """One projection's full tensors, x rows, and per rank the engine's output on its shard.""" + from tensorfold.cuda.exl3.linear import Exl3Linear + + trellis, suh, svh = load_group(model_dir, t.prefix) + save(out, f"{lid}.trellis", trellis) + save(out, f"{lid}.suh", suh) + save(out, f"{lid}.svh", svh) + g = torch.Generator().manual_seed(SEED) + x = torch.randn((max(ROWS), t.k), generator=g).to(torch.bfloat16) + xmap, sha = {}, {} + for m in ROWS: + save(out, f"{lid}.x{m}", x[:m]) + xmap[str(m)] = f"{lid}.x{m}.bin" + + if split == "dim1": # out-split: the rank's bf16 half of the outputs + shards = {} + for rank in (0, 1): + lin = Exl3Linear.from_tensors(cut_dim1(trellis, rank), suh, cut_row(svh, rank), t.codebook, + None, layout="stored") + for m in ROWS: + y = lin(x[:m].cuda()) + torch.cuda.synchronize() + save(out, f"{lid}.shard{m}.r{rank}", y) + shards.setdefault(str(m), []).append(f"{lid}.shard{m}.r{rank}.bin") + entry = {"shards": shards} + else: # in-split: the rank's fp32 partial over the full outputs + partials, ymap = {}, {} + plist = {0: [], 1: []} + for rank in (0, 1): + lin = Exl3Linear.from_tensors(cut_row(trellis, rank), cut_row(suh, rank), svh, t.codebook, + None, layout="stored") + xr = x[:, rank * (t.k // 2):(rank + 1) * (t.k // 2)].contiguous() + for m in ROWS: + part = lin(xr[:m].cuda(), out_dtype=torch.float32) + torch.cuda.synchronize() + save(out, f"{lid}.part{m}.r{rank}", part) + partials.setdefault(str(m), []).append(f"{lid}.part{m}.r{rank}.bin") + plist[rank].append(part) + for m, p0, p1 in zip(ROWS, plist[0], plist[1]): + y = (p0 + p1).to(torch.bfloat16) # rank 0 first, fp32, then bf16 + save(out, f"{lid}.y{m}", y) + ymap[str(m)] = f"{lid}.y{m}.bin" + sha[str(m)] = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + entry = {"partials": partials, "y": ymap, "sha256": sha} + + print(f"{lid}: {t.prefix} EXL3 {t.bits:g}-bit {t.codebook} K={t.k} N={t.n} split={split} " + f"sha256={sha.get('16', '')[:16]}…") + return {"id": lid.rsplit(".", 1)[1], "name": t.prefix, "split": split, "bits": int(t.bits), + "codebook": t.codebook, "k": t.k, "n": t.n, "rows": list(ROWS), + "trellis": f"{lid}.trellis.bin", "suh": f"{lid}.suh.bin", "svh": f"{lid}.svh.bin", + "bias": None, "x": xmap, **entry} + + +def experts_entry(out: Path, model_dir: Path, cp) -> dict: + """One MoE layer's routed experts, full tensors and each rank's fp32 partials (its gate/up/down shards).""" + from tensorfold.cuda.exl3 import experts as generic + + base = exl3_oracle_experts.pick_layer(cp) + by_bits: dict[float, list[int]] = {} + for prefix, t in cp.groups.items(): + m = re.match(rf"^{re.escape(base)}\.mlp\.experts\.(\d+)\.gate_proj$", prefix) + if m is not None: + by_bits.setdefault(t.bits, []).append(int(m.group(1))) + pool = [e for b in (4.0, 5.0, 6.0) for e in sorted(by_bits[b])[:POOL]] + + gate, up, down = [], [], [] + trellis = {"gate": [], "up": [], "down": []} + for slot, e in enumerate(pool): + for part, short, mats, names in (("gate_proj", "gate", gate, trellis["gate"]), + ("up_proj", "up", up, trellis["up"]), + ("down_proj", "down", down, trellis["down"])): + mats.append(load_group(model_dir, f"{base}.mlp.experts.{e}.{part}")) + names.append(f"tp2.experts.{short}{slot}.bin") + save(out, "tp2.experts.suh_g", torch.stack([m[1] for m in gate])) + save(out, "tp2.experts.suh_u", torch.stack([m[1] for m in up])) + save(out, "tp2.experts.svh_g", torch.stack([m[2] for m in gate])) + save(out, "tp2.experts.svh_u", torch.stack([m[2] for m in up])) + save(out, "tp2.experts.suh_d", torch.stack([m[1] for m in down])) + save(out, "tp2.experts.svh_d", torch.stack([m[2] for m in down])) + for slot, mats in enumerate(zip(gate, up, down)): + for name, t in zip(("gate", "up", "down"), mats): + save(out, f"tp2.experts.{name}{slot}", t[0]) + + g = torch.Generator().manual_seed(EXPERTS_SEED) + dims, width = gate[0][0].shape[0] * 16, gate[0][0].shape[1] * 16 + x = torch.randn((max(ROWS), dims), generator=g).to(torch.bfloat16) + pick = torch.stack([torch.randperm(len(pool), generator=g)[:SLOTS] for _ in range(max(ROWS))]).int() + wts = (0.1 + 0.9 * torch.rand((max(ROWS), SLOTS), generator=g)).float() + xmap, pickmap, wtsmap = {}, {}, {} + for m in ROWS: + save(out, f"tp2.experts.x{m}", x[:m]) + save(out, f"tp2.experts.pick{m}", pick[:m]) + save(out, f"tp2.experts.wts{m}", wts[:m]) + xmap[str(m)] = f"tp2.experts.x{m}.bin" + pickmap[str(m)] = f"tp2.experts.pick{m}.bin" + wtsmap[str(m)] = f"tp2.experts.wts{m}.bin" + + partials, ymap, sha = {}, {}, {} + plist = {0: [], 1: []} + for rank in (0, 1): + def shard(mats, part): + """The rank's (trellis, suh, svh) of one expert projection: out-split for gate/up, in for down.""" + cut = cut_dim1 if part in ("gate", "up") else (lambda t, r: cut_row(t, r)) + scale_at = 2 if part in ("gate", "up") else 1 + return [(cut(t, rank), suh if scale_at == 2 else cut_row(suh, rank), + cut_row(svh, rank) if scale_at == 2 else svh) for t, suh, svh in mats] + + ex = generic.prepare(shard(gate, "gate"), shard(up, "up"), shard(down, "down"), "mul1") + scratch = generic.Scratch(ex, max(ROWS), SLOTS) + for m in ROWS: + part = generic.routed(x[:m].cuda(), pick[:m].cuda(), wts[:m].cuda(), ex, scratch, None, m, + limit=LIMIT, act_mode=generic.ACT_BF16) + torch.cuda.synchronize() + save(out, f"tp2.experts.part{m}.r{rank}", part) + partials.setdefault(str(m), []).append(f"tp2.experts.part{m}.r{rank}.bin") + plist[rank].append(part) + for m, p0, p1 in zip(ROWS, plist[0], plist[1]): + y = (p0 + p1).to(torch.bfloat16) # rank 0 first, fp32, then bf16 + save(out, f"tp2.experts.y{m}", y) + ymap[str(m)] = f"tp2.experts.y{m}.bin" + sha[str(m)] = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + + print(f"tp2 experts: {base} {len(pool)} experts (4/5/6-bit pool), width {width}, " + f"sha256={sha['16'][:16]}…") + return {"name": base, "count": len(pool), "dims": dims, "width": width, "slots": SLOTS, + "codebook": "mul1", "limit": LIMIT, "act_mode": generic.ACT_BF16, "rows": list(ROWS), + "gate_k2": [generic.k2_of(m[0]) for m in gate], + "up_k2": [generic.k2_of(m[0]) for m in up], + "down_k2": [generic.k2_of(m[0]) for m in down], + "trellis": trellis, + "suh_g": "tp2.experts.suh_g.bin", "suh_u": "tp2.experts.suh_u.bin", + "svh_g": "tp2.experts.svh_g.bin", "svh_u": "tp2.experts.svh_u.bin", + "suh_d": "tp2.experts.suh_d.bin", "svh_d": "tp2.experts.svh_d.bin", + "x": xmap, "pick": pickmap, "wts": wtsmap, "partials": partials, "y": ymap, "sha256": sha} + + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.cuda.exl3 import format as fmt + + cp = fmt.scan(model_dir) + t_out = pick_out_linear(cp) + t_in = pick_in_linear(cp, t_out) + return {"linears": [linear_entry(out, model_dir, t_out, "dim1", "tp2.linear.out"), + linear_entry(out, model_dir, t_in, "row", "tp2.linear.in")], + "experts": experts_entry(out, model_dir, cp)} diff --git a/zig/tests/cuda/oracle/glm_oracle.py b/zig/tests/cuda/oracle/glm_oracle.py new file mode 100644 index 0000000000..8a19fc52cc --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle.py @@ -0,0 +1,45 @@ +"""Fixtures for `tf-cuda-test glm`: the Python engine's GLM-5.3-Flash kernels and layers on real RED-SNOW tensors.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +# the families' fixtures: siblings in this folder, which leads sys.path when the script runs +import glm_oracle_gaps +import glm_oracle_glue +import glm_oracle_kda +import glm_oracle_latent +import glm_oracle_sparse +import glm_oracle_layer +import glm_oracle_dsa +import glm_oracle_layer_dsa +import glm_oracle_mlp +import glm_oracle_head + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", required=True, help="the checkpoint folder (the recipe's serve view)") + ap.add_argument("--out", required=True, help="the fixture folder to (over)write") + a = ap.parse_args() + out, model_dir = Path(a.out), Path(a.model) + out.mkdir(parents=True, exist_ok=True) + manifest = {"version": 1, "model": str(model_dir)} + manifest["glue"] = glm_oracle_glue.write(out, model_dir) + manifest["kda"] = glm_oracle_kda.write(out, model_dir) + manifest["sparse"] = glm_oracle_sparse.write(out, model_dir) + manifest["latent"] = glm_oracle_latent.write(out, model_dir) + manifest["gaps"] = glm_oracle_gaps.write(out, model_dir) + manifest["layer"] = glm_oracle_layer.write(out, model_dir) + manifest["dsa"] = glm_oracle_dsa.write(out, model_dir) + manifest["mlp"] = glm_oracle_mlp.write(out, model_dir) + manifest["layer_dsa"] = glm_oracle_layer_dsa.write(out, model_dir) + manifest["head"] = glm_oracle_head.write(out, model_dir) + (out / "manifest.json").write_text(json.dumps(manifest, indent=1) + "\n") + print(f"{len(manifest) - 2} families -> {out / 'manifest.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/oracle/glm_oracle_dsa.py b/zig/tests/cuda/oracle/glm_oracle_dsa.py new file mode 100644 index 0000000000..4de8c2a12f --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_dsa.py @@ -0,0 +1,155 @@ +"""glm.dsa.block fixtures: one real RED-SNOW DSA layer's dsa_block at TP=2 rank 0, fp8 latent cache, kv_b bf16.""" +from __future__ import annotations + +import dataclasses +import gc +import os + +os.environ.setdefault("TF_GLM_LATENT", "1") + +from pathlib import Path + +import torch + +LAYER = 3 # the first DSA layer (MoE) +DECODE = (1, 3, 8) # decode windows +# decode positions: dense with the indexer, a window across the dense limit, all-sparse at and past it, ~8K +DECODE_POS = (100, 2048, 2051, 2052, 8000) +PROMPT = ((1800, 512), (1536, 520)) # prompt chunks across the dense limit; the second ends in an 8-row block +FILL = 8192 # the pre-state: this many tokens written by the real block in prompt chunks +CAP = 10240 # cache rows (cap / 4 >= the 8K window's 2048-pool bucket) +DECODE_ROWS, PREFILL_ROWS = 8, 2048 # the engine's buffers (their row counts set the dense pass's chunk cap) +SEED = 5353 + + +def save(out: Path, name: str, t: torch.Tensor) -> str: + t = t.detach().contiguous().cpu() + (out / f"{name}.bin").write_bytes(t.view(torch.uint8).numpy().tobytes() if t.numel() else b"") + return f"{name}.bin" + + +def _load(model_dir: Path): + """weights.load for rank 0, cut to layers 0 .. LAYER (Config.layers patched), then just the layer kept.""" + from tensorfold.families.glm5_next.cuda import weights + + desc = weights.Config.__dict__["read"] + orig = weights.Config.read + weights.Config.read = classmethod(lambda cls, d: dataclasses.replace(orig(d), layers=LAYER + 1)) + try: + w = weights.load(model_dir, rank=0, mtp=False) + finally: + weights.Config.read = desc + w.layers = [w.layers[LAYER]] + w.meta["long_context"] = True # the recipe's 128K context: indexer caches, the indexer on every call + gc.collect() + torch.cuda.empty_cache() + return w + + +def _weights(out: Path, a, pre: str = "dbw") -> dict: + """Rank 0's DSA tensors exactly as the engine holds them (EXL3 words in its strips layout).""" + + def lin(name: str, q) -> dict: + assert q.layout == "strips" and q.bias is None, name + return {"words": save(out, f"{pre}_{name}_words", q.words), "suh": save(out, f"{pre}_{name}_suh", q.suh), + "svh": save(out, f"{pre}_{name}_svh", q.svh), "k": q.k, "n": q.n, "k2": q.k2, "codebook": q.codebook} + + def b16(name: str, q) -> dict: + assert q.weight.dtype == torch.bfloat16, name + return {"w": save(out, f"{pre}_{name}", q.weight), "n": q.n, "k": q.k} + + def plain(name: str, t: torch.Tensor, dtype: torch.dtype = torch.bfloat16) -> str: + assert t.dtype == dtype, (name, t.dtype) + return save(out, f"{pre}_{name}", t) + + ab, ix = a.absorb, a.index + assert type(ab).__name__ == "AbsorbW" and a.proj is None, "the recipe's bf16 AbsorbW and separate q_a/kv_a" + return {"q_a": lin("q_a", a.q_a), "kv_a": lin("kv_a", a.kv_a), "q_b": lin("q_b", a.q_b), "o": lin("o", a.o), + "q_norm": plain("q_norm", a.q_norm), "kv_norm": plain("kv_norm", a.kv_norm), + "wk": plain("wk", ab.wk), "wv": plain("wv", ab.wv), + "kw": b16("kw", ix.kw), "qb": b16("qb", ix.qb), "gate": plain("gate", ix.gate), + "ln_w": plain("ln_w", ix.ln_w, torch.float16), "ln_b": plain("ln_b", ix.ln_b, torch.float16), + "ape": plain("ape", ix.ape)} + + +def _block(w, st, b, R: int, pos: int) -> torch.Tensor: + """forward.dsa_block as _mixer calls it for this layer (eager: host_pos set, nch from the window).""" + from tensorfold.families.glm5_next.cuda import forward + + st.set_pos(pos) + return forward.dsa_block(w.layers[0], w, st.kc[0], st.vc[0], st.pos_dev, b, R, forward.chunks_for(st, R), + st.index[0], pos, None, st.lat_scales[0]) + + +def _caches(st) -> list[torch.Tensor]: + ik, ig, pk = st.index[0] + return [st.kc[0].view(torch.uint8), st.lat_scales[0], ik, ig, pk] + + +NAMES = ("kc", "scales", "ik", "ig", "pk") + + +def _case(out: Path, w, st, pre: list[torch.Tensor], tag: str, R: int, pos: int, prefill: bool) -> dict: + """One window from the shared pre-state: rank 0's stage bits, the fp32 partial and the caches' written rows.""" + from tensorfold.families.glm5_next.cuda import forward, sparse + + c = w.cfg + for dst, src in zip(_caches(st), pre): + dst.copy_(src) + b = forward.Buffers(w, PREFILL_ROWS if prefill else DECODE_ROWS, CAP, prefill=prefill) + g = torch.Generator().manual_seed(SEED + 7 * R + pos) + normed = torch.randn((R, c.hidden), generator=g).to(torch.bfloat16).cuda() + b.normed[:R].copy_(normed) + part = _block(w, st, b, R, pos).reshape(R, c.hidden).clone() + torch.cuda.synchronize() + s = b.lat_s + e = {"rows": R, "pos": pos, "prefill": prefill, "buf_rows": b.rows} + stages = {"normed": normed, "dqa": b.dqa[:R], "dkv": b.dkv[:R], "qr": b.qr[:R], "lat": b.lat[:R], "q": b.q[:R], + "ikr": b.ikr[:R], "igr": b.igr[:R], "qa": s.qa[:R], "ol": s.ol[:R], "vn": b.vn[:R], "part": part} + sparse_rows = pos + R - 1 >= c.dense_limit + if sparse_rows: # the selection again from the post-state pools it read (index_update ran before it) + tokens, counts = sparse.select_tokens(b.qi[:R], b.ikr[:R, c.index_dim:], st.index[0][2], pos, R, + st.index[0][2].shape[0] - 2, st.pos_dev) + stages.update(qi=b.qi[:R], tokens=tokens, counts=counts) + for key, t in stages.items(): + e[key] = save(out, f"db_{tag}_{key}", t) + # every cache row outside the window's must be untouched: the fixture keeps the written rows only + win = {"kc": (pos, pos + R), "scales": (pos, pos + R), "ik": (pos, pos + R), "ig": (pos, pos + R), + "pk": (pos // 4, (pos + R) // 4)} + for name, post, before in zip(NAMES, _caches(st), pre): + lo, hi = win[name] + assert torch.equal(post[:lo], before[:lo]) and torch.equal(post[hi:], before[hi:]), (tag, name) + e[name] = {"lo": lo, "hi": hi, "file": save(out, f"db_{tag}_{name}", post[lo:hi])} + print(f" dsa block {tag}: {R} rows at {pos}, part {part.abs().max():.3g} max, sparse {sparse_rows}") + return e + + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.families.glm5_next.cuda import forward + + out.mkdir(parents=True, exist_ok=True) + w = _load(model_dir) + c, a = w.cfg, w.layers[0].dsa + assert w.layers[0].kind == "dsa" and a.index is not None and c.dense_limit == 2051 + st = forward.State(w, CAP, DECODE_ROWS, kv="fp8") + # the true pre-state: FILL tokens through the real block in 2,048-row prompt chunks (dense, straddling, sparse) + g = torch.Generator().manual_seed(SEED) + b = forward.Buffers(w, PREFILL_ROWS, CAP, prefill=True) + for p0 in range(0, FILL, PREFILL_ROWS): + R = min(PREFILL_ROWS, FILL - p0) + b.normed[:R].copy_(torch.randn((R, c.hidden), generator=g).to(torch.bfloat16).cuda()) + _block(w, st, b, R, p0) + torch.cuda.synchronize() + del b + pre = [t.clone() for t in _caches(st)] + man = {"layer": LAYER, "params": {"eps": c.eps, "index_eps": 1e-6, "scale": c.qk_dim ** -0.5, + "index_scale": c.index_dim ** -0.5, "wscale": 1.0 / 5.656854249492381, + "dense_limit": c.dense_limit, "cap": CAP, "hidden": c.hidden, "fill": FILL}, + "weights": _weights(out, a), "pre": {n: save(out, f"db_pre_{n}", t) for n, t in zip(NAMES, pre)}, + "cases": {}} + for pos in DECODE_POS: + for R in DECODE: + man["cases"][f"d{pos}r{R}"] = _case(out, w, st, pre, f"d{pos}r{R}", R, pos, False) + for pos, R in PROMPT: + man["cases"][f"p{pos}r{R}"] = _case(out, w, st, pre, f"p{pos}r{R}", R, pos, True) + return man diff --git a/zig/tests/cuda/oracle/glm_oracle_gaps.py b/zig/tests/cuda/oracle/glm_oracle_gaps.py new file mode 100644 index 0000000000..a76c4d566e --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_gaps.py @@ -0,0 +1,144 @@ +"""GLM-5.3-Flash gap fixtures on real weights ("gaps"): 1536/512 rmsnorms, fp32-y combine, MoE per-slot y.""" +from __future__ import annotations + +import hashlib +from pathlib import Path + +import torch + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def sha(t: torch.Tensor) -> str: + return hashlib.sha256(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.cuda.exl3 import experts as generic + from tensorfold.cuda.exl3.linear import Exl3Linear + from tensorfold.families.glm5_next.cuda import exl3_generic, glue + from tensorfold.families.glm5_next.cuda.split import RankReader + + dev = "cuda" + L, EPS, LIM, SCALE, TOPK, NE, D, I = 3, 1e-5, 10.0, 2.5, 8, 288, 4096, 1024 + P = f"model.language_model.layers.{L}.mlp." + rd = RankReader(model_dir, 0) + g = torch.Generator(device="cpu").manual_seed(70707) + man: dict = {"layer": L, "eps": EPS} + + def as_int16(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.int16) if t.element_size() == 4 else t + + def lin_mm(lin: Exl3Linear, x: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """forward.py exl3_mm: the dense kernel takes 1..128 rows and is row-invariant.""" + m = x.shape[0] + if m <= 128: + return lin(x, out=out) + o = torch.empty((min(m, 128), lin.n), dtype=out.dtype, device=out.device) + for r in range(0, m, 128): + n = min(m - r, 128) + lin(x[r:r + n], out=o[:n]) + out[r:r + n].copy_(o[:n]) + return out + MP = "model.language_model." + nw1 = rd.get(MP + f"layers.{L}.self_attn.q_a_layernorm.weight").to(torch.bfloat16).to(dev) + nw2 = rd.get(MP + f"layers.{L}.self_attn.kv_a_layernorm.weight").to(torch.bfloat16).to(dev) + rows = 4 + if tuple(nw1.shape) != (1536,) or tuple(nw2.shape) != (512,): + raise ValueError(f"layernorm widths {tuple(nw1.shape)} {tuple(nw2.shape)}") + dp = torch.randn((rows, 2048), generator=g).to(torch.bfloat16).to(dev) + out1 = torch.empty((rows, 1536), dtype=torch.bfloat16, device=dev) + sums1 = torch.empty((rows, 24), dtype=torch.float32, device=dev) + glue.rmsnorm(dp[:, :1536], nw1, EPS, out1, sums1) + out2 = torch.empty((rows, 512), dtype=torch.bfloat16, device=dev) + sums2 = torch.empty((rows, 8), dtype=torch.float32, device=dev) + glue.rmsnorm(dp[:, 1536:], nw2, EPS, out2, sums2) + for name, t in (("rms_dp", dp), ("rms_nw1", nw1), ("rms_nw2", nw2), ("rms_out1", out1), + ("rms_sums1", sums1), ("rms_out2", out2), ("rms_sums2", sums2)): + save(out, name, t) + man["rms"] = {"rows": rows, "d1": 1536, "d2": 512, "eps": EPS, "dp": "rms_dp.bin", "nw1": "rms_nw1.bin", + "nw2": "rms_nw2.bin", "out1": "rms_out1.bin", "sums1": "rms_sums1.bin", + "out2": "rms_out2.bin", "sums2": "rms_sums2.bin"} + + # -- glm.glue.combine: the fp32 per-slot outputs (the EXL3 path's ey is scratch.y, fp32) ------------- + yy = torch.randn((rows * 9, D), generator=g).to(torch.float32).to(dev) + wts = torch.rand((rows, 9), generator=g).to(torch.float32).to(dev) + cout = torch.empty((rows, D), dtype=torch.float32, device=dev) + glue.combine(yy, wts, cout) + for name, t in (("c_y32", yy), ("c_wts32", wts), ("c_out32", cout)): + save(out, name, t) + man["combine"] = {"rows": rows, "d": D, "y": "c_y32.bin", "wts": "c_wts32.bin", "out": "c_out32.bin"} + + # -- the real MoE layer: 288 routed experts, the router, the shared expert, per-slot y and combine --- + moe: dict = {"layer": L, "experts": NE, "dims": D, "width": I, "codebook": "mul1", "limit": LIM, + "act_mode": 0, "slots": TOPK + 1} + gate, up, down = [], [], [] + trellis: dict[str, list[str]] = {"gate": [], "up": [], "down": []} + for e in range(NE): + for part, short, mats in (("gate_proj", "gate", gate), ("up_proj", "up", up), ("down_proj", "down", down)): + mats.append((as_int16(rd.get(P + f"experts.{e}.{part}.trellis")).to(dev), + rd.get(P + f"experts.{e}.{part}.suh").to(dev), + rd.get(P + f"experts.{e}.{part}.svh").to(dev))) + trellis[short].append(f"moe_{short}{e:03d}.bin") + save(out, f"moe_{short}{e:03d}", mats[-1][0]) + ex = generic.prepare(gate, up, down, "mul1", device=dev) + moe["gate_k2"] = [int(v) for v in ex.gate_k2.cpu()] + moe["up_k2"] = [int(v) for v in ex.up_k2.cpu()] + moe["down_k2"] = [int(v) for v in ex.down_k2.cpu()] + moe["gate"] = trellis["gate"] + moe["up"] = trellis["up"] + moe["down"] = trellis["down"] + for name, t in (("suh_g", ex.suh_g), ("suh_u", ex.suh_u), ("svh_g", ex.svh_g), ("svh_u", ex.svh_u), + ("suh_d", ex.suh_d), ("svh_d", ex.svh_d)): + save(out, f"moe_{name}", t) + moe[name] = f"moe_{name}.bin" + gate_w = rd.get(P + "gate.weight").to(torch.bfloat16).to(dev) + bias = rd.get(P + "gate.e_score_correction_bias").to(torch.float32).to(dev) + if tuple(gate_w.shape) != (NE, D) or tuple(bias.shape) != (NE,): + raise ValueError(f"router shapes {tuple(gate_w.shape)} {tuple(bias.shape)}") + save(out, "moe_gate_w", gate_w) + save(out, "moe_bias", bias) + moe["gate_w"] = "moe_gate_w.bin" + moe["bias"] = "moe_bias.bin" + s_lin = {} + for proj in ("gate_proj", "up_proj", "down_proj"): + s_lin[proj] = Exl3Linear.from_tensors( + as_int16(rd.get(P + f"shared_experts.{proj}.trellis")), + rd.get(P + f"shared_experts.{proj}.suh").to(torch.float16).to(dev), + rd.get(P + f"shared_experts.{proj}.svh").to(torch.float16).to(dev), "mul1", device=dev) + if s_lin["gate_proj"].n != I or s_lin["up_proj"].n != I or s_lin["down_proj"].k != I: + raise ValueError(f"shared expert widths {s_lin['gate_proj'].n} {s_lin['up_proj'].n} {s_lin['down_proj'].k}") + xm, pm, wm, ym, sym, partm, hy, hp = {}, {}, {}, {}, {}, {}, {}, {} + for m in (1, 3, 16, 300): + x = torch.randn((m, D), generator=g).to(torch.bfloat16).to(dev) + logits = torch.empty((m, NE), dtype=torch.float32, device=dev) + glue.router(x, gate_w, logits) + pick = torch.empty((m, TOPK + 1), dtype=torch.int32, device=dev) + wts = torch.empty((m, TOPK + 1), dtype=torch.float32, device=dev) + glue.select(logits, bias, pick, wts, TOPK, NE, SCALE, True) + scratch = generic.Scratch(ex, m, TOPK + 1) + y = exl3_generic.routed(x, pick, ex, scratch, m, LIM).clone() # [m * 9, D] fp32 + smg = torch.empty((m, I), dtype=torch.bfloat16, device=dev) + smu = torch.empty((m, I), dtype=torch.bfloat16, device=dev) + lin_mm(s_lin["gate_proj"], x, smg) + lin_mm(s_lin["up_proj"], x, smu) + sgu = torch.cat((smg, smu), dim=1) + sact = torch.empty((m, I), dtype=torch.bfloat16, device=dev) + glue.swiglu(sgu, sact, torch.empty((m,), dtype=torch.float32, device=dev), LIM) + sy = lin_mm(s_lin["down_proj"], sact, torch.empty((m, D), dtype=torch.float32, device=dev)) + y.view(m, TOPK + 1, D)[:, TOPK] = sy + part = torch.empty((m, D), dtype=torch.float32, device=dev) + glue.combine(y, wts, part) + k = str(m) + for name, t, bag in (("x", x, xm), ("pick", pick, pm), ("wts", wts, wm), ("y", y, ym), + ("sy", sy, sym), ("part", part, partm)): + save(out, f"moe_{name}{m}", t) + bag[k] = f"moe_{name}{m}.bin" + hy[k], hp[k] = sha(y), sha(part) + del scratch + moe.update({"rows": [1, 3, 16, 300], "x": xm, "pick": pm, "wts": wm, "y": ym, "sy": sym, "part": partm, + "hash": {"y": hy, "part": hp}}) + man["moe"] = moe + return man diff --git a/zig/tests/cuda/oracle/glm_oracle_glue.py b/zig/tests/cuda/oracle/glm_oracle_glue.py new file mode 100644 index 0000000000..cc08b74b2e --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_glue.py @@ -0,0 +1,198 @@ +"""GLM-5.3-Flash glue fixtures from the engine's Triton kernels on real weights, for glm_glue_tests.zig ("glue").""" +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch + + +def _load(model_dir: Path, want: dict[str, tuple]) -> dict[str, torch.Tensor]: + """The tensors named in `want` (name -> (dtype, shape)) gathered from the checkpoint's shards.""" + from safetensors import safe_open + + found: dict[str, torch.Tensor] = {} + for f in sorted(model_dir.glob("*.safetensors")): + with safe_open(f, framework="pt") as s: + for k in s.keys(): + short = k.removeprefix("model.language_model.") + if short in want and short not in found: + found[short] = s.get_tensor(k) + if len(found) == len(want): + break + missing = set(want) - set(found) + if missing: + raise KeyError(f"missing checkpoint tensors: {sorted(missing)}") + out = {} + for name, (dt, shape) in want.items(): + t = found[name] + assert tuple(t.shape) == shape, f"{name}: {tuple(t.shape)} != {shape}" + out[name] = t.to(dt).contiguous().to("cuda") + return out + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def write(out: Path, model_dir: Path) -> dict: + dev = "cuda" + L = 5 # the layer whose hyper-connection and norm weights the fixtures carry + w = _load(model_dir, { + f"layers.{L}.hc_attn_fn": (torch.float32, (24, 16384)), + f"layers.{L}.hc_attn_base": (torch.float32, (24,)), + f"layers.{L}.hc_attn_scale": (torch.float32, (3,)), + f"layers.{L}.input_layernorm.weight": (torch.bfloat16, (4096,)), + f"layers.{L}.mlp.gate.weight": (torch.bfloat16, (288, 4096)), + f"layers.{L}.mlp.gate.e_score_correction_bias": (torch.float32, (288,)), + "lm_head.weight": (torch.bfloat16, (154880, 4096)), + "embed_tokens.weight": (torch.bfloat16, (154880, 4096)), + }) + from tensorfold.families.glm5_next.cuda import glue, qmm + + EPS, HC_EPS, ITERS = 1e-6, 1e-6, 20 + LIM, SCALE, TOPK, NE, D = 6144.0, 2.5, 8, 144, 4096 + g = torch.Generator(device="cpu").manual_seed(4242) + rows = 4 + + def bf16(*shape) -> torch.Tensor: + return torch.randn(shape, generator=g).to(torch.bfloat16).to(dev) + + nw = w[f"layers.{L}.input_layernorm.weight"] + fn = w[f"layers.{L}.hc_attn_fn"] + base, scale = w[f"layers.{L}.hc_attn_base"], w[f"layers.{L}.hc_attn_scale"] + gate = w[f"layers.{L}.mlp.gate.weight"][:NE].contiguous() + bias = w[f"layers.{L}.mlp.gate.e_score_correction_bias"][:NE].contiguous() + head = w["lm_head.weight"][:4096].contiguous() + table = w["embed_tokens.weight"][:8192].contiguous() + # glm.glue.embed + ids = torch.randint(0, 8192, (rows,), generator=g, dtype=torch.int32).to(dev) + eout = torch.empty((rows, 4, D), dtype=torch.bfloat16, device=dev) + glue.embed(ids, table, D, 4, eout) + save(out, "glue_e_ids", ids) + save(out, "glue_e_tab", table) + save(out, "glue_e_out", eout) + + # glm.glue.rmsnorm + rx = bf16(rows, D) + rout = torch.empty_like(rx) + rsums = torch.empty((rows, D // 64), dtype=torch.float32, device=dev) + glue.rmsnorm(rx, nw, EPS, rout, rsums) + save(out, "glue_r_x", rx) + save(out, "glue_r_out", rout) + save(out, "glue_r_sums", rsums) + + # glm.glue.hc_pre (partial dots + finish) + hx = bf16(rows, 4 * D) + hout = torch.empty((rows, D), dtype=torch.bfloat16, device=dev) + xs = torch.empty((rows, D // 64), dtype=torch.float32, device=dev) + post = torch.empty((rows, 4), dtype=torch.float32, device=dev) + comb = torch.empty((rows, 16), dtype=torch.float32, device=dev) + part = torch.empty((rows, 16, 32), dtype=torch.float32, device=dev) + glue.hc_pre(hx, fn, base, scale, nw, hout, xs, post, comb, part, EPS, HC_EPS, ITERS) + save(out, "glue_h_x", hx) + save(out, "glue_h_out", hout) + save(out, "glue_h_xs", xs) + save(out, "glue_h_post", post) + save(out, "glue_h_comb", comb) + save(out, "glue_h_part", part) + + # glm.glue.hc_post (world partials, contiguous and held apart) + gx = bf16(rows, 4 * D) + gg = torch.randn((2, rows, D), generator=g).to(dev) + gout = torch.empty_like(gx) + glue.hc_post(gx, gout, gg, post, comb) + save(out, "glue_p_x", gx) + save(out, "glue_p_out", gout) + save(out, "glue_p_g", gg) + g0, g1 = gg[0].contiguous(), gg[1].contiguous() + pout = torch.empty_like(gx) + glue.hc_post_pair(gx, pout, g0, g1, post, comb) + save(out, "glue_p_pair_out", pout) + + # glm.glue.residual (residual_add and stream_mean) + resx = bf16(rows, D) + resout = torch.empty_like(resx) + glue.residual_add(resx, resout, gg) + save(out, "glue_res_x", resx) + save(out, "glue_res_out", resout) + smx = bf16(rows, 4 * D) + smout = torch.empty((rows, D), dtype=torch.bfloat16, device=dev) + glue.stream_mean(smx, smout) + save(out, "glue_sm_x", smx) + save(out, "glue_sm_out", smout) + + # glm.glue.swiglu + gu = bf16(rows, 2 * 6144) + sout = torch.empty((rows, 6144), dtype=torch.bfloat16, device=dev) + sxs = torch.empty((rows, 96), dtype=torch.float32, device=dev) + glue.swiglu(gu, sout, sxs, LIM) + save(out, "glue_s_gu", gu) + save(out, "glue_s_out", sout) + save(out, "glue_s_xs", sxs) + + # glm.glue.router (the sliced path: d % (8*64) == 0) + mm_x = bf16(rows, D) + lrout = torch.empty((rows, NE), dtype=torch.float32, device=dev) + glue.router(mm_x, gate, lrout) + save(out, "glue_router_x", mm_x) + save(out, "glue_router_w", gate) + save(out, "glue_router_out", lrout) + + # glm.glue.select (the router's logits, the real bias) + pick = torch.empty((rows, TOPK + 1), dtype=torch.int32, device=dev) + wts = torch.empty((rows, TOPK + 1), dtype=torch.float32, device=dev) + glue.select(lrout, bias, pick, wts, TOPK, NE, SCALE, True) + save(out, "glue_sel_logits", lrout) + save(out, "glue_sel_bias", bias) + save(out, "glue_sel_pick", pick) + save(out, "glue_sel_wts", wts) + + # glm.glue.combine (9 slots: the picked experts' outputs + the shared one) + yy = bf16(rows * (TOPK + 1), D) + cout = torch.empty((rows, D), dtype=torch.float32, device=dev) + glue.combine(yy, wts, cout) + save(out, "glue_c_y", yy) + save(out, "glue_c_wts", wts) + save(out, "glue_c_out", cout) + + # glm.bmm (the lm_head slice through qmm.matmul: the split-K shape) + bx = bf16(rows, D) + bm = qmm.B16(head, 4096, D) + by = qmm.matmul(bx, bm) + save(out, "glue_bmm_x", bx) + save(out, "glue_bmm_w", head) + save(out, "glue_bmm_out", by) + + # glm.sample.greedy (bf16 logits from the head slice, argmax with the tie rule) + lg = qmm.matmul(bx, bm).to(torch.bfloat16).contiguous() + tok = torch.argmax(lg.float(), dim=-1).to(torch.int32) + save(out, "glue_g_logits", lg) + save(out, "glue_g_tok", tok) + + entry = { + "rows": rows, "d": D, "wide": 4 * D, "copies": 4, "ne": NE, "topk": TOPK, + "eps": EPS, "hc_eps": HC_EPS, "iters": ITERS, "lim": LIM, "scale": SCALE, + "ids": "glue_e_ids.bin", "e_tab": "glue_e_tab.bin", "e_out": "glue_e_out.bin", + "r_x": "glue_r_x.bin", "r_w": "glue_nw.bin", "r_out": "glue_r_out.bin", "r_sums": "glue_r_sums.bin", + "h_x": "glue_h_x.bin", "h_out": "glue_h_out.bin", "h_xs": "glue_h_xs.bin", + "h_post": "glue_h_post.bin", "h_comb": "glue_h_comb.bin", + "p_x": "glue_p_x.bin", "p_g": "glue_p_g.bin", "p_out": "glue_p_out.bin", "p_pair_out": "glue_p_pair_out.bin", + "h_part": "glue_h_part.bin", + "res_x": "glue_res_x.bin", "res_out": "glue_res_out.bin", "sm_x": "glue_sm_x.bin", "sm_out": "glue_sm_out.bin", + "s_gu": "glue_s_gu.bin", "s_out": "glue_s_out.bin", "s_xs": "glue_s_xs.bin", + "router_x": "glue_router_x.bin", "router_w": "glue_router_w.bin", "router_out": "glue_router_out.bin", + "sel_logits": "glue_sel_logits.bin", "sel_bias": "glue_sel_bias.bin", + "sel_pick": "glue_sel_pick.bin", "sel_wts": "glue_sel_wts.bin", + "c_y": "glue_c_y.bin", "c_wts": "glue_c_wts.bin", "c_out": "glue_c_out.bin", + "bmm_x": "glue_bmm_x.bin", "bmm_w": "glue_bmm_w.bin", "bmm_out": "glue_bmm_out.bin", + "g_logits": "glue_g_logits.bin", "g_tok": "glue_g_tok.bin", + "fn": "glue_fn.bin", "base": "glue_base.bin", "hc_scale": "glue_scale.bin", "nw": "glue_nw.bin", + } + save(out, "glue_fn", fn) + save(out, "glue_base", base) + save(out, "glue_scale", scale) + save(out, "glue_nw", nw) + print(f"glue: {rows} rows, d={D}, ne={NE} -> {len(entry) - 12} tensors") + return entry diff --git a/zig/tests/cuda/oracle/glm_oracle_head.py b/zig/tests/cuda/oracle/glm_oracle_head.py new file mode 100644 index 0000000000..42e7496ab5 --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_head.py @@ -0,0 +1,86 @@ +"""glm.head fixtures: forward.compute's tail on both ranks (stream mean, final norm, head half), then sample_rows.""" +from __future__ import annotations + +import threading +from pathlib import Path + +import torch + +from glm_oracle_layer import CAP, DECODE_ROWS, PREFILL_ROWS, Comm, Hub, _load, save + +CASES = (("d1", 1, False), ("d8", 8, False), ("p300", 300, True)) # a prompt chunk's head reads its last row only +SEED = 7171 + + +def _tail(w, x: torch.Tensor, R: int, prefill: bool) -> dict: + """compute()'s lines after the layers, on the engine's own buffers.""" + from tensorfold.families.glm5_next.cuda import forward, glue + + c = w.cfg + b = forward.Buffers(w, PREFILL_ROWS if prefill else DECODE_ROWS, CAP, prefill=prefill) + b.x[:R].copy_(x) + glue.stream_mean(b.x[:R], b.hidden[:R]) + glue.rmsnorm(b.hidden[:R], w.norm, c.eps, b.fnormed[:R], b.fxs[:R]) + if prefill: + lg = forward.mm(b, b.fnormed[R - 1:R], w.head, b.fxs[R - 1:R], b.logits[:1]) + else: + lg = forward.mm(b, b.fnormed[:R], w.head, b.fxs[:R], b.logits[:R]) + torch.cuda.synchronize() + return {"hidden": b.hidden[:R].clone(), "fnormed": b.fnormed[:R].clone(), "fxs": b.fxs[:R].clone(), + "logits": lg.clone()} + + +def _sample(W, logits: list[torch.Tensor]) -> list[int]: + """decode.sample_rows greedy on both ranks at once (its all_gather of the packed candidates).""" + from tensorfold.families.glm5_next.cuda import decode + + hub, res, errs = Hub(), [None, None], [] + + def body(rank: int) -> None: + try: + W[rank].comm = Comm(hub, rank) + res[rank] = decode.sample_rows(W[rank], logits[rank], list(range(logits[rank].shape[0])), None) + except BaseException as e: # the peer must not wait at the barrier forever + errs.append(e) + hub.bar.abort() + + threads = [threading.Thread(target=body, args=(rank,)) for rank in (0, 1)] + for t in threads: + t.start() + for t in threads: + t.join() + for w in W: + w.comm = None + if errs: + raise errs[0] + assert res[0] == res[1], "the ranks chose differently" + return res[0] + + +def write(out: Path, model_dir: Path) -> dict: + W = [_load(model_dir, rank, 0) for rank in (0, 1)] + c = W[0].cfg + assert torch.equal(W[0].norm, W[1].norm) and W[0].norm.dtype == torch.bfloat16 + heads = [] + for rank, w in enumerate(W): + assert w.head.weight.dtype == torch.bfloat16 and w.head.n == c.vocab // 2 and w.vocab_offset == rank * w.head.n + heads.append({"w": save(out, f"hd_head{rank}", w.head.weight), "n": w.head.n, "k": w.head.k, + "offset": w.vocab_offset}) + man = {"params": {"eps": c.eps, "hidden": c.hidden, "streams": c.streams, "vocab": c.vocab}, + "norm": save(out, "hd_norm", W[0].norm), "heads": heads, "cases": {}} + g = torch.Generator().manual_seed(SEED) + for tag, R, prefill in CASES: + x = torch.randn((R, c.streams * c.hidden), generator=g).to(torch.bfloat16).cuda() + t = [_tail(w, x, R, prefill) for w in W] + for key in ("hidden", "fnormed", "fxs"): + assert torch.equal(t[0][key], t[1][key]), f"{tag}: the ranks' {key} differ" + tok = _sample(W, [r["logits"] for r in t]) + e = {"rows": R, "prefill": prefill, "x": save(out, f"hd_{tag}_x", x), + "tok": save(out, f"hd_{tag}_tok", torch.tensor(tok, dtype=torch.int32)), + "logits": [save(out, f"hd_{tag}_logits{rank}", r["logits"]) for rank, r in enumerate(t)]} + for key in ("hidden", "fnormed", "fxs"): + e[key] = save(out, f"hd_{tag}_{key}", t[0][key]) + wins = sum(int(v >= W[0].head.n) for v in tok) + print(f" head {tag}: {R} rows, {len(tok)} tokens ({wins} from rank 1)") + man["cases"][tag] = e + return man diff --git a/zig/tests/cuda/oracle/glm_oracle_kda.py b/zig/tests/cuda/oracle/glm_oracle_kda.py new file mode 100644 index 0000000000..dab75bd5ef --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_kda.py @@ -0,0 +1,253 @@ +# GLM-5.3-Flash's KDA oracle fixtures: the engine's rank-0 (world 2) chain, chunks and block, bit for bit. + +from __future__ import annotations + +import json +from pathlib import Path + +import torch + +import exl3_oracle_layer # save, sha_of, load_tensor + +SEED = 4477 +CHAIN_ROWS = (1, 3, 8, 64, 128) # 1/3/8 narrow (serial), 64/128 wide (three kernels, same bits) +CHUNK_ROWS = 2048 # one prompt chunk: 64 sub-chunks +CHUNK_SPLITS = (32, 96, 1024, 2016) # 32-aligned cuts, the last leaving a cut segment +BLOCK_ROWS = 300 # the block's prompt chunk (the Cut path's whole run) +CUT = 64 # the block's cut point (first segment 64 rows, then the cut segment) + + +def _attn_prefix(model_dir: Path, li: int) -> str: + """The checkpoint's attention prefix (the serve view may nest under model.language_model).""" + index = Path(model_dir) / "model.safetensors.index.json" + wm = json.loads(index.read_text())["weight_map"] + for name in wm: + if f"layers.{li}.self_attn.qkv_proj.trellis" in name: + return name[: name.index(f"layers.{li}.self_attn.qkv_proj") + len(f"layers.{li}")] + raise KeyError("qkv_proj.trellis") + + +def _save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def _pick_kda_layer(model_dir: Path) -> int: + cfg = json.loads((Path(model_dir) / "config.json").read_text())["text_config"] + for i, attn in enumerate(cfg["layer_types"]): + if attn == "linear_attention": + return i + raise ValueError(f"no KDA layer in {model_dir}") + + +def _rand(g, *shape, dtype=torch.float32, device="cuda", scale=1.0): + return (torch.randn(*shape, generator=g, dtype=torch.float32) * scale).to(dtype).to(device) + + +def _front(k, qmm, x, rows): + # The block's projection rows (kda_block's front, rank 0): p = q|k|v | f_a | g_a | b and the pair's second mats. + dev = x.device + p = torch.zeros((rows, k.out_n), dtype=torch.bfloat16, device=dev) + m = x.shape[0] + o = torch.empty((128, k.qkv.n), dtype=torch.bfloat16, device=dev) + for r in range(0, m, 128): + n = min(m - r, 128) + k.qkv(x[r:r + n], out=o[:n]) + p[r:r + n, :k.qkv.n].copy_(o[:n]) + fa = torch.empty((m, k.fa.n), dtype=torch.bfloat16, device=dev) + ga = torch.empty((m, k.ga.n), dtype=torch.bfloat16, device=dev) + bb = torch.empty((m, k.b.n), dtype=torch.bfloat16, device=dev) + qmm.matmul(x, k.fa, out=fa) + qmm.matmul(x, k.ga, out=ga) + qmm.matmul(x, k.b, out=bb) + p[:, k.fa_off:k.fa_off + k.fa.n].copy_(fa) + p[:, k.ga_off:k.ga_off + k.ga.n].copy_(ga) + p[:, k.b_off:].copy_(bb) + a = qmm.matmul(fa, k.fb, xs=qmm.group_sums(fa)) + g = qmm.matmul(ga, k.gb, xs=qmm.group_sums(ga)) + return p, a, g + + +def _store_entry(out: Path, tag: str, p, a, g, conv0, state0, o, state_out, prep, scratch=None) -> dict: + keys = {"p": p, "a": a, "g": g, "conv0": conv0, "state0": state0, "out": o, "state_out": state_out, + "q": prep[0], "k": prep[1], "v": prep[2], "lg": prep[3], "beta": prep[4]} + e = {key: f"{tag}_{key}.bin" for key in keys} + e["sha256"] = {"out": exl3_oracle_layer.sha_of(o), "state_out": exl3_oracle_layer.sha_of(state_out)} + if scratch is not None: + e["scratch"] = {} + for key in ("k", "v", "g", "b"): + e["scratch"][key] = f"{tag}_s{key}.bin" + _save(out, f"{tag}_s{key}", getattr(scratch, key)) + for key, t in keys.items(): + _save(out, f"{tag}_{key}", t) + return e + + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.families.glm5_next.cuda import kda, kda_chunked, qmm, weights + + dev = "cuda" + li = _pick_kda_layer(model_dir) + w = weights.load(model_dir, rank=0, device=dev) + k = w.layers[li].kda + c = w.cfg + H, LL, DK = k.heads, c.lin_heads // w.world, kda.DK + hidden, b_off, eps, lower = c.hidden, k.b_off, c.eps, c.lower + C = 3 * H * DK + assert (H, b_off) == (c.lin_heads // w.world, 3 * (c.lin_heads // w.world) * 128 + 256), (H, b_off) + + fam: dict = {"params": {"layer": li, "H": H, "LL": LL, "dk": DK, "hidden": hidden, "b_off": b_off, + "eps": eps, "lower": lower, "heads_full": c.lin_heads, "conv_taps": c.conv, + "world": w.world, "rank": w.rank, "qkv_n": k.qkv.n, "p_stride": k.out_n}, + "chain": {}, "chunked": {}, "block": {}} + + def inputs(rows, salt): + g = torch.Generator().manual_seed(SEED + salt) + x = _rand(g, rows, hidden, dtype=torch.bfloat16) + conv0 = _rand(g, 3, C, dtype=torch.bfloat16) + state0 = _rand(g, H, DK, DK, scale=0.3) + return x, conv0, state0 + + # -- the decode chain: narrow (serial) and wide (three kernels), the engine's own bits ------------------------- + for rows in CHAIN_ROWS: + x, conv0, state0 = inputs(rows, rows) + p, a, g = _front(k, qmm, x, rows) + state_out = torch.empty((H, DK, DK), dtype=torch.float32, device=dev) + scr = kda.KDAScratch(rows, H, dev) + o = kda.chain(p, b_off, a, g, conv0, k.conv, state0, k.a_log, k.dt_bias, k.norm, eps, lower, rows, scr, + state_out) + prep = kda_chunked.prep_reference(p, b_off, a, conv0, k.conv, k.a_log, k.dt_bias, lower, rows) + tag = f"chain_r{rows}" + e = _store_entry(out, tag, p, a, g, conv0, state0, o, state_out, prep, scratch=scr if rows == 8 else None) + _save(out, f"{tag}_x", x) + e["x"] = f"{tag}_x.bin" + if rows == 8: # one row at a time with the conv window shifted in between: the engine's own state + conv_r = conv0.clone() + st_r = state0.clone() + for i in range(rows): + o_r = torch.empty((1, H * DK), dtype=torch.bfloat16, device=dev) + so_r = torch.empty((H, DK, DK), dtype=torch.float32, device=dev) + kda.chain(p[i:i + 1], b_off, a[i:i + 1], g[i:i + 1], conv_r, k.conv, st_r, k.a_log, k.dt_bias, + k.norm, eps, lower, 1, kda.KDAScratch(1, H, dev), so_r) + st_r = so_r + if i < rows - 1: + conv_r = torch.cat([conv_r[1:], p[i:i + 1, :C]]) + _save(out, f"{tag}_rowwise_state", st_r) + e["rowwise_state"] = f"{tag}_rowwise_state.bin" + if rows == 8: # the engine's own replay of the saved rows must reach the same state + st = torch.empty_like(state_out) + kda.replay(state0, scr, rows, st) + assert torch.equal(st, state_out), "engine replay != chain state" + fam["chain"][str(rows)] = e + + # -- the prompt chunk in chunked (WY) form, plus its 32-aligned cuts ------------------------------------------- + rows = CHUNK_ROWS + x, conv0, state0 = inputs(rows, 100000) + p, a, g = _front(k, qmm, x, rows) + o = torch.empty((rows, H * DK), dtype=torch.bfloat16, device=dev) + state_out = torch.empty((H, DK, DK), dtype=torch.float32, device=dev) + kda_chunked.chain(p, b_off, a, g, conv0, k.conv, state0, k.a_log, k.dt_bias, k.norm, eps, lower, rows, o, + state_out, pos=0) + prep = kda_chunked.prep_reference(p, b_off, a, conv0, k.conv, k.a_log, k.dt_bias, lower, rows) + tag = "chunk_r2048" + e = _store_entry(out, tag, p, a, g, conv0, state0, o, state_out, prep) + splits = {} + for n in CHUNK_SPLITS: # the engine's own cuts must reproduce the whole run's bits (chunk invariance) + o1 = torch.empty((n, H * DK), dtype=torch.bfloat16, device=dev) + st1 = torch.empty((H, DK, DK), dtype=torch.float32, device=dev) + kda_chunked.chain(p[:n], b_off, a, g, conv0, k.conv, state0, k.a_log, k.dt_bias, k.norm, eps, lower, n, + o1, st1, pos=0) + conv_mid = p[n - 3:n, :C].clone() + o2 = torch.empty((rows - n, H * DK), dtype=torch.bfloat16, device=dev) + st2 = torch.empty((H, DK, DK), dtype=torch.float32, device=dev) + kda_chunked.chain(p[n:], b_off, a[n:], g[n:], conv_mid, k.conv, st1, k.a_log, k.dt_bias, k.norm, eps, + lower, rows - n, o2, st2, pos=n) + assert torch.equal(torch.cat((o1, o2)), o), f"engine cut {n} != whole run" + assert torch.equal(st2, state_out), f"engine cut {n} state" + splits[str(n)] = {"st1": f"{tag}_cut{n}_st1.bin"} + _save(out, f"{tag}_cut{n}_st1", st1) + fam["chunked"] = {**e, "rows": rows, "splits": splits} + + # the block starts at the layer's normed rows (the layernorm/HC mixing is the glue family's gate): seeded bf16 + g = torch.Generator().manual_seed(SEED + 200000) + normed = _rand(g, rows, hidden, dtype=torch.bfloat16) + p, a, g = _front(k, qmm, normed, rows) + o = torch.empty((rows, H * DK), dtype=torch.bfloat16, device=dev) + state_out = torch.empty((H, DK, DK), dtype=torch.float32, device=dev) + kda_chunked.chain(p, b_off, a, g, conv0, k.conv, state0, k.a_log, k.dt_bias, k.norm, eps, lower, rows, o, + state_out, pos=0) + prep = kda_chunked.prep_reference(p, b_off, a, conv0, k.conv, k.a_log, k.dt_bias, lower, rows) + for key, t in (("q", prep[0]), ("k", prep[1]), ("v", prep[2]), ("lg", prep[3]), ("beta", prep[4])): + _save(out, f"block_{key}", t) + part = torch.empty((rows, hidden), dtype=torch.float32, device=dev) + from tensorfold.cuda.exl3.linear import Exl3Linear + if isinstance(k.o, Exl3Linear): # the engine's exl3_mm: the kernel takes 1..128 rows + for r in range(0, rows, 128): + n = min(rows - r, 128) + k.o(o[r:r + n], out=part[r:r + n]) + else: + qmm.matmul(o, k.o, xs=qmm.group_sums(o), out=part, f32=True) + oc1 = torch.empty((CUT, H * DK), dtype=torch.bfloat16, device=dev) + stc1 = torch.empty((H, DK, DK), dtype=torch.float32, device=dev) + kda_chunked.chain(p[:CUT], b_off, a, g, conv0, k.conv, state0, k.a_log, k.dt_bias, k.norm, eps, lower, CUT, + oc1, stc1, pos=0) + conv_cut = p[CUT - 3:CUT, :C].clone() + oc2 = torch.empty((rows - CUT, H * DK), dtype=torch.bfloat16, device=dev) + stc2 = torch.empty((H, DK, DK), dtype=torch.float32, device=dev) + kda_chunked.chain(p[CUT:], b_off, a[CUT:], g[CUT:], conv_cut, k.conv, stc1, k.a_log, k.dt_bias, k.norm, eps, + lower, rows - CUT, oc2, stc2, pos=CUT) + assert torch.equal(torch.cat((oc1, oc2)), o), "engine block cut != whole run" + assert torch.equal(stc2, state_out), "engine block cut state" + pfx = _attn_prefix(model_dir, li) + trellis = exl3_oracle_layer.load_tensor(model_dir, f"{pfx}.self_attn.qkv_proj.trellis") + svh = exl3_oracle_layer.load_tensor(model_dir, f"{pfx}.self_attn.qkv_proj.svh") + from tensorfold.families.glm5_next.cuda.split import fused_rank_runs + t_runs = fused_rank_runs(trellis.shape[1], 3, w.rank, w.world) + s_runs = fused_rank_runs(svh.shape[0], 3, w.rank, w.world) + tag = "block_r300" + blk = {"rows": rows, "cut": CUT, "p": f"{tag}_p.bin", "a": f"{tag}_a.bin", "g": f"{tag}_g.bin", + "conv0": f"{tag}_conv0.bin", "state0": f"{tag}_state0.bin", "out": f"{tag}_out.bin", + "state_out": f"{tag}_state_out.bin", "part": f"{tag}_part.bin", + "sha256": {"out": exl3_oracle_layer.sha_of(o), "part": exl3_oracle_layer.sha_of(part)}, + "qkv": {"prefix": f"{pfx}.self_attn.qkv_proj", "k": k.qkv.k, "n": k.qkv.n, + "t_runs": [[r.start, r.stop] for r in t_runs], "s_runs": [[r.start, r.stop] for r in s_runs], + "codebook": c.codebook}, + "o_proj": {"prefix": f"{pfx}.self_attn.o_proj", + "exl3": op_trellis(model_dir, li, pfx) is not None, "n": k.o.n, "k": k.o.k}} + assert k.fa is not None and k.ga is not None and k.b is not None, "fused qkv with no low-rank pair" + mms = {} + for key, q in (("fa", k.fa), ("ga", k.ga), ("b", k.b), ("fb", k.fb), ("gb", k.gb)): + _save(out, f"{tag}_{key}", q.weight) + mms[key] = {"n": q.n, "k": q.k, "file": f"{tag}_{key}.bin"} + blk["mms"] = mms + for key in ("q", "k", "v", "lg", "beta"): + blk[f"prep_{key}"] = f"block_{key}.bin" + if not blk["o_proj"]["exl3"]: + _save(out, f"{tag}_o", k.o.weight) + blk["o_proj"]["w"] = f"{tag}_o.bin" + for key, t in (("p", p), ("a", a), ("g", g), ("conv0", conv0), ("state0", state0), ("out", o), + ("state_out", state_out), ("part", part)): + _save(out, f"{tag}_{key}", t) + _save(out, f"{tag}_normed", normed) + blk["normed"] = f"{tag}_normed.bin" + fam["block"] = blk + + for key, t in (("a_log", k.a_log), ("dt_bias", k.dt_bias), ("norm", k.norm), ("conv", k.conv)): + _save(out, f"kda_{key}", t) + fam["tensors"] = {"a_log": "kda_a_log.bin", "dt_bias": "kda_dt_bias.bin", "norm": "kda_norm.bin", + "conv": "kda_conv.bin"} + print(f" kda: layer {li}, H={H} heads/rank, chain rows {CHAIN_ROWS}, chunk {rows}, block {BLOCK_ROWS}") + return fam + + +def op_trellis(model_dir: Path, li: int, pfx: str): + from safetensors import safe_open + root = Path(model_dir) + index = root / "model.safetensors.index.json" + name = f"{pfx}.self_attn.o_proj.trellis" + rel = json.loads(index.read_text())["weight_map"].get(name) if index.exists() else None + if rel is None: + return None + with safe_open(str(root / rel), framework="pt", device="cpu") as f: + return f.get_tensor(name) + + diff --git a/zig/tests/cuda/oracle/glm_oracle_latent.py b/zig/tests/cuda/oracle/glm_oracle_latent.py new file mode 100644 index 0000000000..b6b6d572cc --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_latent.py @@ -0,0 +1,135 @@ +# GLM-5.3-Flash latent-cache oracle: fp8/bf16 cache write, kv_b absorb/expand, dense 512-key chunks + merge. + +from __future__ import annotations + +import os + +os.environ.setdefault("TF_GLM_LATENT", "1") + +import json +from pathlib import Path + +import torch + +import exl3_oracle_layer # save, sha_of + +SEED = 9311 +CAP = 4096 +# dense (pos, R): contexts 1, 37, 511, 512, 2051; mid-chunk starts; the 512-row prompt block over 4K. +DENSE = ((0, 1), (29, 8), (1000, 8), (2040, 8), (504, 8), (2043, 8), (0, 512), (1536, 512), (2048, 512)) +ABSORB_ROWS = (1, 8, 64, 512) +WRITE_ROWS = ((0, 1), (2040, 8), (2048, 512)) + + +def _save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def _pick_dsa_layer(model_dir: Path) -> int: + cfg = json.loads((Path(model_dir) / "config.json").read_text())["text_config"] + for i, attn in enumerate(cfg["layer_types"]): + if attn != "linear_attention": + return i + raise ValueError(f"no DSA layer in {model_dir}") + + +def _rand(g, *shape, dtype=torch.float32, device="cuda", scale=1.0): + return (torch.randn(*shape, generator=g, dtype=torch.float32) * scale).to(dtype).to(device) + + +def _seeded_cache(g, cap: int, kind, lw: int): + """A cache whose rows hold seeded values; the window's rows are garbage the kernel must overwrite.""" + cache = (torch.randn(cap, lw, generator=g, dtype=torch.float32) * 0.5).to(kind).to("cuda") + scales = torch.rand(cap, generator=g, dtype=torch.float32).to("cuda") + 0.01 + return cache, scales + + +def write(out: Path, model_dir: Path) -> dict: + out.mkdir(parents=True, exist_ok=True) + from tensorfold.families.glm5_next.cuda import latent, weights + + dev = "cuda" + li = _pick_dsa_layer(model_dir) + w = weights.load(model_dir, rank=0, device=dev) + a = w.layers[li].dsa + ab = a.absorb + H, D, LW, V = ab.heads, ab.qk_dim, ab.lw, ab.v_dim + scale = D ** -0.5 + _save(out, "wk", ab.wk) + _save(out, "wv", ab.wv) + fam: dict = {"params": {"layer": li, "H": H, "D": D, "LW": LW, "V": V, "heads": a.heads, + "scale": scale, "world": w.world, "rank": w.rank}, + "write": {}, "absorb": {}, "expand": {}, "dense": {}} + + # -- the cache write: fp8 codes with a per-token amax / 448 scale, and a plain bf16 copy --------------------- + for n, (pos, R) in enumerate(WRITE_ROWS): + g = torch.Generator().manual_seed(SEED + n) + lat = _rand(g, R, LW, dtype=torch.bfloat16, scale=0.7) + for kind in ("fp8", "bf16"): + fp8 = kind == "fp8" + cache8, scales = _seeded_cache(torch.Generator().manual_seed(SEED + 40 + n), CAP, torch.bfloat16, LW) + if fp8: + cache8 = cache8.to(torch.float8_e4m3fn) + cache_pre, scales_pre = cache8.clone(), scales.clone() + pd = torch.tensor([pos], dtype=torch.int32, device=dev) + latent.latent_write(lat, cache8, pd, scales if fp8 else None) + torch.cuda.synchronize() + tag = f"w{pos}r{R}_{kind}" + e = {"pos": pos, "R": R, "fp8": fp8, "lat": f"{tag}_lat.bin", "cache_pre": f"{tag}_cpre.bin", + "cache": f"{tag}_cache.bin"} + if fp8: + e["scales_pre"] = f"{tag}_spre.bin" + e["scales"] = f"{tag}_scales.bin" + for name, t in (("lat", lat), ("cache_pre", cache_pre), ("cache", cache8), + ("scales_pre", scales_pre), ("scales", scales)): + if name in e: + _save(out, e[name][:-4], t) + fam["write"][tag] = e + + # -- absorb and expand on the layer's kv_b (bf16 AbsorbW: the recipe's TF_GLM_KVB default) ------------------- + for n, R in enumerate(ABSORB_ROWS): + g = torch.Generator().manual_seed(SEED + 300 + n) + q = _rand(g, R, H, D, dtype=torch.bfloat16) + qa = torch.empty(R, H, LW, dtype=torch.bfloat16, device=dev) + latent.absorb_q(q, ab, qa) + ol = _rand(g, R, H, LW, dtype=torch.bfloat16, scale=0.4) + ov = torch.empty(R, H, V, dtype=torch.bfloat16, device=dev) + latent.expand_v(ol, ab, ov) + torch.cuda.synchronize() + e = {"R": R, "q": f"a{R}_q.bin", "qa": f"a{R}_qa.bin", "ol": f"a{R}_ol.bin", "out": f"a{R}_out.bin"} + for name, t in (("q", q), ("qa", qa), ("ol", ol), ("out", ov)): + _save(out, f"a{R}_{name}", t) + fam["absorb"][str(R)] = e + fam["expand"][str(R)] = e + + # -- dense chunks and the merge: every row over its own causal keys in 512-key chunks ------------------------- + for n, (pos, R) in enumerate(DENSE): + for kind in ("fp8", "bf16"): + fp8 = kind == "fp8" + tag = f"d{pos}r{R}_{kind}" + g = torch.Generator().manual_seed(SEED + 600 + n) + cache, scales = _seeded_cache(g, CAP, torch.bfloat16, LW) + if fp8: + cache = cache.to(torch.float8_e4m3fn) + qa = _rand(g, R, H, LW, dtype=torch.bfloat16) + nch = -(-(pos + R) // latent.CHUNK) + s = latent.LatentScratch(R, H, nch, dev, lw=LW) + outb = torch.empty(R, H, LW, dtype=torch.bfloat16, device=dev) + pd = torch.tensor([pos], dtype=torch.int32, device=dev) + latent.attention(qa, cache, pd, s, scale=scale, nch=nch, out=outb, + scales=scales if fp8 else None) + torch.cuda.synchronize() + e = {"pos": pos, "R": R, "nch": nch, "fp8": fp8, "heads": H, + "qa": f"{tag}_qa.bin", "cache_pre": f"{tag}_cpre.bin", "cache": f"{tag}_cache.bin", + "out": f"{tag}_out.bin"} + if fp8: + e["scales"] = f"{tag}_scales.bin" + e["scales_pre"] = f"{tag}_spre.bin" + for name, t in (("qa", qa), ("cache_pre", cache.clone()), ("cache", cache), + ("scales_pre", scales.clone()), ("scales", scales)): + if name in e: + _save(out, e[name][:-4], t) + _save(out, f"{tag}_out", outb) + e["sha256"] = {"out": exl3_oracle_layer.sha_of(outb)} + fam["dense"][tag] = e + return fam diff --git a/zig/tests/cuda/oracle/glm_oracle_layer.py b/zig/tests/cuda/oracle/glm_oracle_layer.py new file mode 100644 index 0000000000..4838a88336 --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_layer.py @@ -0,0 +1,275 @@ +"""glm.layer.kda/mlp fixtures: a real KDA layer's forward (MoE or dense MLP) at TP=2, both ranks in one process.""" +from __future__ import annotations + +import dataclasses +import gc +import threading +from pathlib import Path + +import torch + +LAYER = 4 # the first KDA layer with a MoE (0-2 carry the dense MLP, 3 is DSA) +DECODE = (1, 3, 8) # decode windows: the narrow chain, then commit(R, R) and its conv shift +KEEP = 5 # the 8-row window also commits its first 5 rows (a draft's accepted prefix) +PROMPT, PROMPT_POS = 300, 37 # one prompt chunk at an unaligned position: its first KDA sub-chunk starts 5 rows in +DECODE_POS = 512 # a decode window's position (KDA reads none; the State carries it) +DECODE_ROWS, PREFILL_ROWS, CAP = 8, 2048, 4096 # the engine's decode and prompt buffers +SEED = 6161 + + +def save(out: Path, name: str, t: torch.Tensor) -> str: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + return f"{name}.bin" + + +def _load(model_dir: Path, rank: int, layer: int = LAYER, long: bool = False): + """weights.load for one rank, cut to layers 0 .. layer, then just the layer kept; ``long``: indexer caches.""" + from tensorfold.families.glm5_next.cuda import weights + + desc = weights.Config.__dict__["read"] + orig = weights.Config.read + weights.Config.read = classmethod(lambda cls, d: dataclasses.replace(orig(d), layers=layer + 1)) + try: + w = weights.load(model_dir, rank=rank, mtp=False) + finally: + weights.Config.read = desc + w.layers = [w.layers[layer]] + if long: + w.meta["long_context"] = True + gc.collect() + torch.cuda.empty_cache() + return w + + +class Hub: + """Where the two rank threads meet: one barrier, one mailbox per rank.""" + + def __init__(self) -> None: + self.bar = threading.Barrier(2, timeout=1800) + self.box: list = [None, None] + + +class Comm: + """A two-rank communicator over threads: all_gather in rank order and exchange, synchronous on the host.""" + + def __init__(self, hub: Hub, rank: int) -> None: + self.hub, self.rank, self.snap = hub, rank, None + + def _post(self, sends) -> list: + torch.cuda.synchronize() + self.hub.box[self.rank] = [t.clone() for t in sends] + self.hub.bar.wait() + return self.hub.box + + def _done(self) -> None: + torch.cuda.synchronize() + self.hub.bar.wait() + + def all_gather(self, src: torch.Tensor, out: torch.Tensor) -> None: + box = self._post([src.reshape(-1)]) + n = src.numel() + for k in (0, 1): + out[k * n:(k + 1) * n].copy_(box[k][0]) + if self.snap is not None: + torch.cuda.synchronize() + self.snap(out[:2 * n]) + self._done() + + def exchange(self, sends, recvs, peer: int) -> None: + box = self._post(sends) + for dst, src in zip(recvs, box[peer]): + dst.copy_(src) + self._done() + + +def _run(W, R: int, pos: int, prefill: bool, split: bool, states, x0, snap0=None, keep: int | None = None) -> list[dict]: + """layer_forward (or split_layers) on both ranks at once; each rank's x, new KDA state and conv window.""" + from tensorfold.families.glm5_next.cuda import forward + from tensorfold.families.glm5_next.cuda.hcsplit import HcSplit, SplitSettings + + hub, res, errs = Hub(), [None, None], [] + + def body(rank: int) -> None: + try: + w = W[rank] + w.comm = Comm(hub, rank) + st = forward.State(w, CAP, DECODE_ROWS) + st.conv[0].copy_(states[rank][0]) + st.rec[0, 0].copy_(states[rank][1]) + st.set_pos(pos) + b = forward.Buffers(w, PREFILL_ROWS if prefill else DECODE_ROWS, CAP, prefill=prefill) + if split: # TF_GLM_HC_SPLIT=1, TF_GLM_PREFILL_OVERLAP=2 (two pieces, p2p swaps on a second stream) + b.split = HcSplit(w, b, SplitSettings(True, True, 2, 64, "p2p", True)) + b.x[:R].copy_(x0) + if rank == 0 and snap0 is not None: + w.comm.snap = lambda g: snap0(g, st, b) + if split: + forward.split_layers(w, st, b, R, None, pos) + else: + forward.layer_forward(w.layers[0], w, st, b, R, host_pos=pos) + if not prefill: + forward.commit(w, st, b, R, keep or R) + torch.cuda.synchronize() + res[rank] = {"x": b.x[:R].clone(), "rec": st.rec[st.cur[0], 0].clone(), "conv": st.conv[0].clone()} + except BaseException as e: # the peer must not wait at the barrier forever + errs.append(e) + hub.bar.abort() + + threads = [threading.Thread(target=body, args=(rank,)) for rank in (0, 1)] + for t in threads: + t.start() + for t in threads: + t.join() + if errs: + raise errs[0] + for w in W: + w.comm = None + return res + + +def _weights(out: Path, L, c, pre: str = "lkw") -> dict: + """Rank 0's tensors exactly as the engine holds them (EXL3 words in its strips layout); a KDA layer's mixer too.""" + k, m = L.kda, L.moe + + def lin(name: str, q) -> dict: + assert q.layout == "strips" and q.bias is None, name + return {"words": save(out, f"{pre}_{name}_words", q.words), "suh": save(out, f"{pre}_{name}_suh", q.suh), + "svh": save(out, f"{pre}_{name}_svh", q.svh), "k": q.k, "n": q.n, "k2": q.k2, "codebook": q.codebook} + + def b16(name: str, q) -> dict: + assert q.weight.dtype == torch.bfloat16, name + return {"w": save(out, f"{pre}_{name}", q.weight), "n": q.n, "k": q.k} + + def plain(name: str, t: torch.Tensor, dtype: torch.dtype) -> str: + assert t.dtype == dtype, (name, t.dtype) + return save(out, f"{pre}_{name}", t) + + man = {} + if L.kind == "kda": + man["kda"] = {"qkv": lin("qkv", k.qkv), "o": lin("o", k.o), "fa": b16("fa", k.fa), "ga": b16("ga", k.ga), + "b": b16("b", k.b), "fb": b16("fb", k.fb), "gb": b16("gb", k.gb), + "conv": plain("conv", k.conv, torch.bfloat16), "a_log": plain("a_log", k.a_log, torch.float32), + "dt_bias": plain("dt_bias", k.dt_bias, torch.float32), + "norm": plain("norm", k.norm, torch.bfloat16), "heads": k.heads, "b_off": k.b_off, + "fa_off": k.fa_off, "ga_off": k.ga_off, "p_stride": k.out_n} + for site, h in (("attn_hc", L.attn_hc), ("ffn_hc", L.ffn_hc)): + man[site] = {"fn": plain(f"{site}_fn", h.fn, torch.float32), "base": plain(f"{site}_base", h.base, torch.float32), + "scale": plain(f"{site}_scale", h.scale, torch.float32)} + man["in_norm"] = plain("in_norm", L.in_norm, torch.bfloat16) + man["post_norm"] = plain("post_norm", L.post_norm, torch.bfloat16) + if m is None: # a dense layer: the MLP's gate, up (the rank's columns) and down (its rows) + d = L.mlp + assert d.gu is None, "a fused gate|up pack" + man["mlp"] = {"gate": lin("mlp_gate", d.gate), "up": lin("mlp_up", d.up), "down": lin("mlp_down", d.down)} + return man + ex = m.experts + E = ex.count + assert len(ex.keep) == 3 * E and ex.gu_stride == 0 + experts = {"count": E, "dims": ex.dims, "width": ex.width, "cb": ex.cb} + for j, proj in enumerate(("gate", "up", "down")): + mats = ex.keep[j * E:(j + 1) * E] + offs, at = [], 0 + for t in mats: + offs.append(at) + at += t.numel() * t.element_size() + experts[proj] = save(out, f"{pre}_ex_{proj}", torch.cat([t.reshape(-1) for t in mats])) + experts[f"{proj}_offs"] = offs + for name in ("gate_k2", "up_k2", "down_k2"): + experts[name] = [int(v) for v in getattr(ex, name).cpu()] + for name in ("suh_g", "suh_u", "svh_g", "svh_u", "suh_d", "svh_d"): + experts[name] = plain(f"ex_{name}", getattr(ex, name), torch.float16) + s = m.shared + man["moe"] = {"router": plain("router", m.router, torch.bfloat16), "bias": plain("bias", m.bias, torch.float32), + "experts": experts, "shared": {"gate": lin("sh_gate", s.gate), "up": lin("sh_up", s.up), + "down": lin("sh_down", s.down)}} + return man + + +def _case(out: Path, W, tag: str, R: int, pos: int, prefill: bool, pre: str = "lk") -> dict: + """One window from a seeded state: rank 0's bits at every stage, rank 1's two partials, the final x and state.""" + from tensorfold.families.glm5_next.cuda import kda_chunked + + c, k = W[0].cfg, W[0].layers[0].kda + g = torch.Generator().manual_seed(SEED + 1000 * R + int(prefill)) + x0 = torch.randn((R, c.streams * c.hidden), generator=g).to(torch.bfloat16).cuda() + C = 3 * k.heads * 128 + states = [] + for _rank in (0, 1): # each rank's own heads: its own conv window and recurrent state + conv = torch.randn((3, C), generator=g).to(torch.bfloat16).cuda() + rec = (torch.randn((k.heads, 128, 128), generator=g) * 0.3).cuda() + states.append((conv, rec)) + snaps: list[dict] = [] + + def snap(gath: torch.Tensor, st, b) -> None: + site = {"gath": gath.clone(), "x": b.x[:R].clone(), "normed": b.normed[:R].clone(), + "post": b.post[:R].clone(), "comb": b.comb[:R].clone()} + if not snaps: # the attention site: the KDA block's projections, read-out and saved rows + site["p"] = (b.kproj[0, :R] if prefill else st.proj[0, :R]).clone() + site["ka"], site["kg"] = b.ka[:R].clone(), b.kg[:R].clone() + sc = b.kscratch if prefill else st.scratch[0] + site["kout"] = sc.out[:R].clone() + if not prefill: + for key in ("k", "v", "g", "b"): + site[f"s{key}"] = getattr(sc, key)[:R].clone() + elif W[0].layers[0].moe is None: # a dense MLP: gate, up and the SwiGLU, under the engine's names + for key, src in (("smg", "mg"), ("smu", "mu"), ("sact", "act")): + site[key] = getattr(b, src)[:R].clone() + else: # the MoE site: route, per-slot y (shared expert in slot 8) + for key in ("mlog", "pick", "wts", "sy", "smg", "smu", "sact"): + site[key] = getattr(b, key)[:R].clone() + site["ey"] = b.ey[:R].clone() + snaps.append(site) + + res = _run(W, R, pos, prefill, False, states, x0, snap) + assert len(snaps) == 2, len(snaps) + e = {"rows": R, "pos": pos, "prefill": prefill, "x0": save(out, f"{pre}_{tag}_x0", x0), + "conv0": save(out, f"{pre}_{tag}_conv0", states[0][0]), "rec0": save(out, f"{pre}_{tag}_rec0", states[0][1]), + "x": save(out, f"{pre}_{tag}_x", res[0]["x"]), "rec": save(out, f"{pre}_{tag}_rec", res[0]["rec"]), + "conv": save(out, f"{pre}_{tag}_conv", res[0]["conv"])} + assert torch.equal(res[0]["x"], res[1]["x"]), "the two ranks' x differ" + n = R * c.hidden + for i, site in enumerate(snaps): + gath = site.pop("gath") + e[f"s{i}"] = {"part": save(out, f"{pre}_{tag}_s{i}_part", gath[:n]), + "peer": save(out, f"{pre}_{tag}_s{i}_peer", gath[n:2 * n])} + for key, t in site.items(): + e[f"s{i}"][key] = save(out, f"{pre}_{tag}_s{i}_{key}", t) + prep = kda_chunked.prep_reference(snaps[0]["p"], k.b_off, snaps[0]["ka"], states[0][0], k.conv, k.a_log, + k.dt_bias, c.lower, R) + e["prep"] = {key: save(out, f"{pre}_{tag}_prep_{key}", t) for key, t in zip(("q", "k", "v", "lg", "beta"), prep)} + if prefill: # the same chunk row-split between the ranks: every rank's own rows and state equal + sres = _run(W, R, pos, prefill, True, states, x0) + h = R // 2 + for rank, rows in ((0, slice(0, h)), (1, slice(h, R))): + assert torch.equal(sres[rank]["x"][rows], res[rank]["x"][rows]), f"split x rank {rank}" + assert torch.equal(sres[rank]["rec"], res[rank]["rec"]), f"split rec rank {rank}" + assert torch.equal(sres[rank]["conv"], res[rank]["conv"]), f"split conv rank {rank}" + e["split_equal"] = True + elif R == max(DECODE): # commit(R, keep < R): the engine's replay of the kept rows, the conv shifted by keep + kres = _run(W, R, pos, prefill, False, states, x0, keep=KEEP) + assert torch.equal(kres[0]["x"], res[0]["x"]), "keep changed x" + e["keep"] = {"keep": KEEP, "rec": save(out, f"{pre}_{tag}_keep_rec", kres[0]["rec"]), + "conv": save(out, f"{pre}_{tag}_keep_conv", kres[0]["conv"])} + print(f" layer {tag}: {R} rows at {pos}, x {res[0]['x'].float().abs().max():.3g} max") + return e + + +def write(out: Path, model_dir: Path, layer: int = LAYER, pre: str = "lk") -> dict: + from tensorfold.families.glm5_next.cuda import kda + + W = [_load(model_dir, rank, layer) for rank in (0, 1)] + c, L = W[0].cfg, W[0].layers[0] + assert L.kind == "kda" and (L.moe is None or L.moe.shared is not None) and c.norm_topk + man = {"layer": layer, "params": {"eps": c.eps, "hc_eps": c.hc_eps, "hc_iters": c.hc_iters, "lower": c.lower, + "limit": c.limit, "routed_scale": c.routed_scale, "top_k": c.top_k, + "experts": c.experts, "hidden": c.hidden, "streams": c.streams}, + "weights": _weights(out, L, c, f"{pre}w"), "cases": {}} + chunked = kda.CHUNKED + try: + kda.CHUNKED = True # the recipe's TF_GLM_KDA_CHUNKED=1 (read per call) + for R in DECODE: + man["cases"][f"d{R}"] = _case(out, W, f"d{R}", R, DECODE_POS, False, pre) + man["cases"][f"p{PROMPT}"] = _case(out, W, f"p{PROMPT}", PROMPT, PROMPT_POS, True, pre) + finally: + kda.CHUNKED = chunked + return man diff --git a/zig/tests/cuda/oracle/glm_oracle_layer_dsa.py b/zig/tests/cuda/oracle/glm_oracle_layer_dsa.py new file mode 100644 index 0000000000..c9d6533c12 --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_layer_dsa.py @@ -0,0 +1,162 @@ +"""glm.layer.dsa fixtures: one real DSA MoE layer's whole forward at TP=2 from a real prefill's caches, fp8 latents.""" +from __future__ import annotations + +import os + +os.environ.setdefault("TF_GLM_LATENT", "1") + +import threading +from pathlib import Path + +import torch + +import glm_oracle_dsa as dsa_oracle +from glm_oracle_layer import Comm, Hub, _load, _weights, save + +LAYER = 3 # the first DSA layer (MoE) +DECODE = (1, 3, 8) # decode windows +DECODE_POS = (512, 4000) # below the dense limit (2,051) and past it +PROMPT, PROMPT_POS = 300, 1900 # one prompt chunk across the dense limit (rows 151.. are sparse) +FILL = 4096 # the pre-state: this many tokens through the real layer_forward in prompt chunks +CAP = 4096 # cache rows (cap / 4 = 1,024 pools: the 4K windows' bucket) +DECODE_ROWS, PREFILL_ROWS = 8, 2048 # the engine's decode and prompt buffers +SEED = 7373 +NAMES = dsa_oracle.NAMES + + +def _caches(st) -> list[torch.Tensor]: + ik, ig, pk = st.index[0] + return [st.kc[0].view(torch.uint8), st.lat_scales[0], ik, ig, pk] + + +def _run(W, R: int, pos: int, prefill: bool, split: bool, pres, x0s, snap0=None) -> list[dict]: + """layer_forward (or split_layers) on both ranks at once from each rank's caches; each rank's x and caches.""" + from tensorfold.families.glm5_next.cuda import forward + from tensorfold.families.glm5_next.cuda.hcsplit import HcSplit, SplitSettings + + hub, res, errs = Hub(), [None, None], [] + + def body(rank: int) -> None: + try: + w = W[rank] + w.comm = Comm(hub, rank) + st = forward.State(w, CAP, DECODE_ROWS, kv="fp8") + if pres is not None: # else the fill: State's zeroed caches + for dst, src in zip(_caches(st), pres[rank]): + dst.copy_(src) + b = forward.Buffers(w, PREFILL_ROWS if prefill else DECODE_ROWS, CAP, prefill=prefill) + if split: # TF_GLM_HC_SPLIT=1, TF_GLM_PREFILL_OVERLAP=2 (two pieces, p2p swaps on a second stream) + b.split = HcSplit(w, b, SplitSettings(True, True, 2, 64, "p2p", True)) + if rank == 0 and snap0 is not None: + w.comm.snap = lambda g: snap0(g, st, b) + for p0, x0 in zip(range(pos, pos + R * len(x0s), R), x0s): # the fill runs several chunks + n = x0.shape[0] + st.set_pos(p0) + b.x[:n].copy_(x0) + if split: + forward.split_layers(w, st, b, n, forward.chunks_for(st, n), p0) + else: + forward.layer_forward(w.layers[0], w, st, b, n, forward.chunks_for(st, n), p0) + torch.cuda.synchronize() + res[rank] = {"x": b.x[:x0s[-1].shape[0]].clone(), "caches": [t.clone() for t in _caches(st)]} + except BaseException as e: # the peer must not wait at the barrier forever + errs.append(e) + hub.bar.abort() + + threads = [threading.Thread(target=body, args=(rank,)) for rank in (0, 1)] + for t in threads: + t.start() + for t in threads: + t.join() + if errs: + raise errs[0] + for w in W: + w.comm = None + return res + + +def _case(out: Path, W, pres, tag: str, R: int, pos: int, prefill: bool) -> dict: + """One window from the pre-state: rank 0's bits at every stage, rank 1's two partials, x and the caches after.""" + from tensorfold.families.glm5_next.cuda import sparse + + c = W[0].cfg + g = torch.Generator().manual_seed(SEED + 7 * R + pos) + x0 = torch.randn((R, c.streams * c.hidden), generator=g).to(torch.bfloat16).cuda() + sparse_rows = pos + R - 1 >= c.dense_limit + snaps: list[dict] = [] + + def snap(gath: torch.Tensor, st, b) -> None: + site = {"gath": gath.clone(), "x": b.x[:R].clone(), "normed": b.normed[:R].clone(), + "post": b.post[:R].clone(), "comb": b.comb[:R].clone()} + if not snaps: # the attention site: the DSA block's stages (dsa_oracle's) + s = b.lat_s + for key, t in (("dqa", b.dqa), ("dkv", b.dkv), ("qr", b.qr), ("lat", b.lat), ("q", b.q), + ("ikr", b.ikr), ("igr", b.igr), ("qa", s.qa), ("ol", s.ol), ("vn", b.vn)): + site[key] = t[:R].clone() + if sparse_rows: # the selection again from the post-state pools it read + pk = st.index[0][2] + tokens, counts = sparse.select_tokens(b.qi[:R], b.ikr[:R, c.index_dim:], pk, pos, R, + pk.shape[0] - 2, st.pos_dev) + site.update(qi=b.qi[:R].clone(), tokens=tokens.clone(), counts=counts.clone()) + else: # the MoE site: route, per-slot y (shared expert in slot 8) + for key in ("mlog", "pick", "wts", "sy", "smg", "smu", "sact", "ey"): + site[key] = getattr(b, key)[:R].clone() + snaps.append(site) + + res = _run(W, R, pos, prefill, False, pres, [x0], snap) + assert len(snaps) == 2, len(snaps) + assert torch.equal(res[0]["x"], res[1]["x"]), "the two ranks' x differ" + e = {"rows": R, "pos": pos, "prefill": prefill, "buf_rows": PREFILL_ROWS if prefill else DECODE_ROWS, + "sparse": sparse_rows, "x0": save(out, f"ld_{tag}_x0", x0), "x": save(out, f"ld_{tag}_x", res[0]["x"])} + n = R * c.hidden + for i, site in enumerate(snaps): + gath = site.pop("gath") + e[f"s{i}"] = {"part": save(out, f"ld_{tag}_s{i}_part", gath[:n]), + "peer": save(out, f"ld_{tag}_s{i}_peer", gath[n:2 * n])} + for key, t in site.items(): + e[f"s{i}"][key] = save(out, f"ld_{tag}_s{i}_{key}", t) + # every cache row outside the window's must be untouched: the fixture keeps the written rows only + win = {"kc": (pos, pos + R), "scales": (pos, pos + R), "ik": (pos, pos + R), "ig": (pos, pos + R), + "pk": (pos // 4, (pos + R) // 4)} + for name, post, before in zip(NAMES, res[0]["caches"], pres[0]): + lo, hi = win[name] + assert torch.equal(post[:lo], before[:lo]) and torch.equal(post[hi:], before[hi:]), (tag, name) + e[name] = {"lo": lo, "hi": hi, "file": save(out, f"ld_{tag}_{name}", post[lo:hi])} + if prefill: # the same chunk row-split between the ranks: own rows of x and every cache equal + sres = _run(W, R, pos, prefill, True, pres, [x0]) + h = R // 2 + for rank, rows in ((0, slice(0, h)), (1, slice(h, R))): + assert torch.equal(sres[rank]["x"][rows], res[rank]["x"][rows]), f"split x rank {rank}" + for name, a, b in zip(NAMES, sres[rank]["caches"], res[rank]["caches"]): + assert torch.equal(a, b), f"split {name} rank {rank}" + e["split_equal"] = True + print(f" layer dsa {tag}: {R} rows at {pos}, x {res[0]['x'].float().abs().max():.3g} max, sparse {sparse_rows}") + return e + + +def write(out: Path, model_dir: Path) -> dict: + W = [_load(model_dir, rank, LAYER, long=True) for rank in (0, 1)] + c, L = W[0].cfg, W[0].layers[0] + a = L.dsa + assert L.kind == "dsa" and a.index is not None and L.moe is not None and L.moe.shared is not None + assert c.dense_limit == 2051 and c.norm_topk + # the true pre-state: FILL tokens through the real layer_forward on both ranks in 2,048-row prompt chunks + g = torch.Generator().manual_seed(SEED) + fill = [torch.randn((PREFILL_ROWS, c.streams * c.hidden), generator=g).to(torch.bfloat16).cuda() + for _ in range(FILL // PREFILL_ROWS)] + pres = [r["caches"] for r in _run(W, PREFILL_ROWS, 0, True, False, None, fill)] + del fill + man = {"layer": LAYER, + "params": {"eps": c.eps, "hc_eps": c.hc_eps, "hc_iters": c.hc_iters, "lower": c.lower, "limit": c.limit, + "routed_scale": c.routed_scale, "top_k": c.top_k, "experts": c.experts, "hidden": c.hidden, + "streams": c.streams, "index_eps": 1e-6, "scale": c.qk_dim ** -0.5, + "index_scale": c.index_dim ** -0.5, "wscale": 1.0 / 5.656854249492381, + "dense_limit": c.dense_limit, "cap": CAP, "fill": FILL}, + "weights": _weights(out, L, c, "ldw"), "pre": {n: save(out, f"ld_pre_{n}", t) for n, t in zip(NAMES, pres[0])}, + "cases": {}} + man["weights"]["dsa"] = dsa_oracle._weights(out, a, "ldw_dsa") + for pos in DECODE_POS: + for R in DECODE: + man["cases"][f"d{pos}r{R}"] = _case(out, W, pres, f"d{pos}r{R}", R, pos, False) + man["cases"][f"p{PROMPT_POS}r{PROMPT}"] = _case(out, W, pres, f"p{PROMPT_POS}r{PROMPT}", PROMPT, PROMPT_POS, True) + return man diff --git a/zig/tests/cuda/oracle/glm_oracle_mlp.py b/zig/tests/cuda/oracle/glm_oracle_mlp.py new file mode 100644 index 0000000000..dfc8f6445c --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_mlp.py @@ -0,0 +1,12 @@ +"""glm.layer.mlp fixtures: a dense-MLP KDA layer (0-2 carry the dense MLP) through glm_oracle_layer's windows.""" +from __future__ import annotations + +from pathlib import Path + +import glm_oracle_layer + +LAYER = 1 + + +def write(out: Path, model_dir: Path) -> dict: + return glm_oracle_layer.write(out, model_dir, LAYER, "lm") diff --git a/zig/tests/cuda/oracle/glm_oracle_sparse.py b/zig/tests/cuda/oracle/glm_oracle_sparse.py new file mode 100644 index 0000000000..16838954f0 --- /dev/null +++ b/zig/tests/cuda/oracle/glm_oracle_sparse.py @@ -0,0 +1,141 @@ +# GLM-5.3-Flash's DSA sparse-path oracle fixtures: indexer update, scores, top-pool select; fp8 cache, bit-exact. + +from __future__ import annotations + +import os + +os.environ.setdefault("TF_GLM_LATENT", "1") + +import json +from pathlib import Path + +import torch + +import exl3_oracle_layer # save, sha_of + +SEED = 7753 +# The windows: decode rows at the count boundaries (2,048/2,051/2,052) and long contexts, one prompt chunk. +CASES = ((2048, 8), (2052, 8), (3000, 8), (16000, 128), (40000, 8), (2016, 2048)) + + +def _save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def _pick_dsa_layer(model_dir: Path) -> int: + cfg = json.loads((Path(model_dir) / "config.json").read_text())["text_config"] + for i, attn in enumerate(cfg["layer_types"]): + if attn != "linear_attention": + return i + raise ValueError(f"no DSA layer in {model_dir}") + + +def _rand(g, *shape, dtype=torch.float32, device="cuda", scale=1.0): + return (torch.randn(*shape, generator=g, dtype=torch.float32) * scale).to(dtype).to(device) + + +def _cap(rows: int) -> int: + return max(4096, (rows + 2047) // 2048 * 2048) + + +def write(out: Path, model_dir: Path) -> dict: + out.mkdir(parents=True, exist_ok=True) + from tensorfold.families.glm5_next.cuda import glue, latent, qmm, sparse, weights + + dev = "cuda" + li = _pick_dsa_layer(model_dir) + w = weights.load(model_dir, rank=0, device=dev) + a = w.layers[li].dsa + ix, c = a.index, w.cfg + H, D, HD = c.index_heads, c.index_dim, c.index_heads * c.index_dim + assert (H, D) == (32, 128), (H, D) + scale = D ** -0.5 + wscale = 1.0 / 5.656854249492381 + _save(out, "ln_w", ix.ln_w) + _save(out, "ln_b", ix.ln_b) + _save(out, "ape", ix.ape) + + fam: dict = {"params": {"layer": li, "H": H, "D": D, "HP": 32, "hidden": c.hidden, "pool": sparse.POOL, + "topk_pools": sparse.TOPK_POOLS, + "width": sparse.TOPK_POOLS * sparse.POOL + sparse.POOL - 1, + "chunk": latent.CHUNK, "lw": latent.L, "dense_limit": c.dense_limit, + "scale": scale, "wscale": wscale, "world": w.world, "rank": w.rank}, + "ln_w": "ln_w.bin", "ln_b": "ln_b.bin", "ape": "ape.bin", "cases": {}} + + for n, (pos, R) in enumerate(CASES): + tag = f"{'p' if R > 512 else 'd'}{pos}r{R}" + g = torch.Generator().manual_seed(SEED + n) + cap = _cap(pos + R) + normed = _rand(g, R, c.hidden, dtype=torch.bfloat16) + ikr = torch.empty((R, D + H), dtype=torch.bfloat16, device=dev) + qmm.matmul(normed, ix.kw, out=ikr) + igr = torch.empty((R, D), dtype=torch.float32, device=dev) + glue.router(normed, ix.gate, igr) + torch.cuda.synchronize() # qmm.matmul runs on a side stream: index_update must see the finished k_raw + qi = _rand(g, R, HD, dtype=torch.bfloat16) + # the caches: rows before the window hold earlier windows' values (seeded), so partial pools re-pool them + ik = torch.zeros((cap, D), dtype=torch.bfloat16, device=dev) + ig = torch.zeros((cap, D), dtype=torch.bfloat16, device=dev) + if pos: + g2 = torch.Generator().manual_seed(SEED + 100 + n) + ik[:pos] = _rand(g2, pos, D, dtype=torch.bfloat16) + ig[:pos] = _rand(g2, pos, D, dtype=torch.bfloat16) + pk = torch.zeros((cap // 4 + 2, D), dtype=torch.bfloat16, device=dev) + if pos >= 4: + pk[:pos // 4] = _rand(torch.Generator().manual_seed(SEED + 200 + n), pos // 4, D, dtype=torch.bfloat16) + pos_dev = torch.tensor([pos], dtype=torch.int32, device=dev) + sparse.index_update(ikr[:, :D], igr, ix.ln_w, ix.ln_b, ix.ape, ik, ig, pk, pos_dev) + # the kernels must produce layer-normalised keys (loose guard: a racy or stale run drifts ~30x) + x = ikr[:, :D].float() + t = (x - x.mean(1, keepdim=True)) / torch.sqrt(x.var(1, unbiased=False, keepdim=True) + 1e-6) + y = (t * ix.ln_w.float() + ix.ln_b.float()).to(torch.bfloat16) + if (y - ik[pos:pos + R]).abs().max().item() > 0.1: + raise RuntimeError("index_update window rows are not LN(ikr): racy or stale kernel run") + e: dict = {"pos": pos, "R": R, "cap": cap, "np_max": pk.shape[0] - 2, "blocks": []} + rows = pos + R + for name, t in (("ikr", ikr), ("igr", igr), ("wts", ikr[:, D:]), ("qi", qi), + ("ik_pre", ik[:rows].clone()), ("ig_pre", ig[:rows].clone()), ("pk_pre", pk.clone()), + ("ik", ik[:rows].clone()), ("ig", ig[:rows].clone()), ("pk", pk.clone())): + e[name] = f"{tag}_{name}.bin" + _save(out, f"{tag}_{name}", t) + + # -- pool scores: the engine's own block loop (SELECT_ROWS rows at a time, one np_max for the window) ------- + B = min(R, sparse.SELECT_ROWS) + np_max = pk.shape[0] - 2 + scores = torch.empty((B, np_max), dtype=torch.float32, device=dev) + at = pos_dev + for b0 in range(0, R, B): + n = min(B, R - b0) + if b0: + at = pos_dev + b0 + sparse._scores[(n, (np_max + 64 * sparse.SCORE_LOOP - 1) // (64 * sparse.SCORE_LOOP))]( + qi[b0:b0 + n], ikr[b0:b0 + n, D:], D + H, pk, scores, at, n, np_max, scale, wscale, + H=H, HP=32, D=D, BP=64, L=sparse.SCORE_LOOP, num_warps=4) + e["blocks"].append(f"{tag}_sc{b0}.bin") + _save(out, f"{tag}_sc{b0}", scores[:n]) + + # -- selection on those scores: tokens and counts, the kernel path (np_max >= 512) -------------------------- + tokens, counts = sparse.select_tokens(qi, ikr[:, D:], pk, pos, R, np_max, pos_dev) + e["tokens"], e["counts"] = f"{tag}_tokens.bin", f"{tag}_counts.bin" + _save(out, f"{tag}_tokens", tokens) + _save(out, f"{tag}_counts", counts) + e["sha256"] = {"tokens": exl3_oracle_layer.sha_of(tokens), "counts": exl3_oracle_layer.sha_of(counts)} + + # -- sparse latent attention on the fp8 cache over the window's own selection ------------------------------- + lat = _rand(g, cap, latent.L, dtype=torch.bfloat16, scale=0.5) + cache = torch.empty((cap, latent.L), dtype=torch.float8_e4m3fn, device=dev) + scales = torch.empty((cap,), dtype=torch.float32, device=dev) + for w0 in range(0, cap, 2048): + pw = torch.tensor([w0], dtype=torch.int32, device=dev) + latent.latent_write(lat[w0:w0 + 2048], cache, pw, scales) + qa = _rand(g, R, a.heads, latent.L, dtype=torch.bfloat16) + out_in = _rand(g, R, a.heads, latent.L, dtype=torch.bfloat16, scale=0.3) + out_out = out_in.clone() + latent.sparse_attention(qa, cache, tokens, counts, out_out, scale, scales) + e["heads"] = a.heads + for name, t in (("qa", qa), ("cache", cache), ("scales", scales), ("out_in", out_in), ("out", out_out)): + e[name] = f"{tag}_{name}.bin" + _save(out, f"{tag}_{name}", t) + e["sha256"]["out"] = exl3_oracle_layer.sha_of(out_out) + fam["cases"][tag] = e + return fam