diff --git a/RUNBOOK.md b/RUNBOOK.md index 67d8ecdc2a..b604f2b308 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -168,6 +168,12 @@ The server logs `heat_wait_s`, a cancel ends the wait, and the reply's tokens do On three GB10s serving GLM-5.3-Flash NVFP4 at TP=3, a 1,035,295-token prompt without the bands kept one host's hottest zone at 94 C or more for most of the prompt (97.7 C at the top); `TF_HEAT_HIGH=92` and `TF_HEAT_LOW=88` held all three at 93.0 C or under, and the first token came 10% later. +`TF_GRAPH_CACHE=0` (or `off`, any case) runs every step eagerly in any family that installs the keyed CUDA graph cache, +without a rebuild. Set it alike on every tensor-parallel rank: an off rank never joins the cache's per-miss vote, so a +rank left on would wait at its first miss, and only the off rank logs the switch. Ranks that share a cache compare +`Settings.digest` of their effective settings at start, before building it, and refuse to start when they differ. +A tensor-parallel family passes its cache settings in the TP config; the bootstrap then compares them and refuses to +start on every rank, naming TF_GRAPH_CACHE among the settings that must match. Archives whose names contain `host-only` are CPU verification artifacts and cannot serve CUDA inference. ## Flash Next and paired Metal serving diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index db27613ce9..3706d0d4fe 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -183,6 +183,12 @@ the one-shot all-gather protocol of [b12x](https://github.com/local-inference-la sequence parity and a timeout that fails the transport. Copyright 2026 Luke Alonso and the b12x contributors, Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships no NOTICE file. +The Zig tensor-parallel fabric's RoCE path (`zig/src/fabric/tp/mailbox.cu`, `roce.zig`, `roce_verbs.zig`) adapts +the one-shot all-gather protocol of [b12x](https://github.com/local-inference-lab/b12x)'s RoCEnante +(`b12x/comm/roce/`, PRs #295 and #315): a device-resident epoch, sequence flags after each payload, two slots by +sequence parity and a timeout that fails the transport. Copyright 2026 Luke Alonso and the b12x contributors, +Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships no NOTICE file. + ## Vendored code and weights In the Python engine, frozen at 0.6.6, `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified diff --git a/build.zig b/build.zig index 8361b06746..e04ab811ac 100644 --- a/build.zig +++ b/build.zig @@ -19,6 +19,9 @@ const golden_build = @import("zig/build/golden.zig"); const contract_build = @import("zig/build/kernel_contract.zig"); const video_build = @import("zig/build/video.zig"); const sys_build = @import("zig/build/sys.zig"); +const tp_build = @import("zig/build/tp.zig"); +const dsv41_build = @import("zig/build/dsv41.zig"); +const sessions_build = @import("zig/build/sessions.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -202,6 +205,10 @@ pub fn build(b: *std.Build) void { b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); contract_build.hostTests(b, gemma4_contract_host); + tp_build.add(b, target, optimize, test_step); + sessions_build.add(b, test_step); + // DeepSeek-V4.1's host code reads POSIX and Linux APIs (pthreads, clock_gettime, io_uring) that Windows lacks + if (b.graph.host.result.os.tag != .windows) dsv41_build.add(b, test_step); sys_build.wire(b); @import("zig/build/checkpoint_config.zig").wire(b); } diff --git a/docs/recipes/README.md b/docs/recipes/README.md index e8094405d7..70299c896e 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -18,6 +18,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | Gemma 4 31B | [Native Metal, affine 4-bit](gemma-4-native.md) | | Gemma 4 12B | [Native Metal, affine 4-bit, MTP assistant](gemma-4-12b-native.md) | | DeepSeek-V4-Flash | [MLX on a 256 GB Mac, DSpark and MTP drafts](deepseek-v4-flash.md) | +| DeepSeek-V4.1-Flash | [CUDA (Zig): tools and checks](deepseek-v4.1-cuda.md) | | Qwen3.6-35B-A3B | [MLX with MTP drafts, one-GPU CUDA](qwen3.6-moe.md) | | Step-5-Preview | [Native Metal on one 256 GB Mac: 2-bit experts, MTP drafts, images](step-5-preview.md) | diff --git a/docs/recipes/deepseek-v4.1-cuda.md b/docs/recipes/deepseek-v4.1-cuda.md new file mode 100644 index 0000000000..54b3fad765 --- /dev/null +++ b/docs/recipes/deepseek-v4.1-cuda.md @@ -0,0 +1,37 @@ +# DeepSeek-V4.1-Flash on CUDA (Zig): tools and checks + +The `deepseek_v41` family runs on the native Zig engine with its own CUDA kernels (`zig/kernels/cuda/deepseek_v41/`). +Its development tools are Python and live outside this repository, in jayleaton/deepseek-v41-tensorfold-spark, +pinned to commit c1cb91840fb8: +[`tools/tensorfold-dsv41/`](https://github.com/jayleaton/deepseek-v41-tensorfold-spark/tree/c1cb91840fb8ff5db0bd27418fc7cc766a222cbc/tools/tensorfold-dsv41). +That folder's README lists each tool's dependencies (the TensorFold Python engine, PyTorch, Triton, a GPU where +noted) and `MAP.tsv` names the TensorFold path each file came from; its `check.py` runs the tools' own checks without +a GPU. + +Nothing in `zig build` or `zig build test` runs Python. The checks below read fixtures committed here; references, +captures, AOT compiles, SASS identity and GPU measurements are separate receipts made with the tools on real hardware. + +## SASS identity of the kernel copies + +The family's kernel copies (line 1 of each names its source) must compile to the reference extensions' SASS. With +Docker and `nvcr.io/nvidia/pytorch:26.07-py3`, from a TensorFold checkout: + +``` +OUT=/scratch/sass SM=121 TF_DSV41_KERNELS=$PWD/zig/kernels/cuda/deepseek_v41 TF_DSV41_REF_SRC=/path/to/reference \ + /path/to/tools/tensorfold-dsv41/kernels/sass_check.sh [name ...] +``` + +`TF_DSV41_REF_SRC` (and optionally `TF_DSV41_REF_SRC_R1` / `_X3GM2` with `TF_DSV41_REF_REV*`) names the reference +repository and revision; the check stops before building without them. +`-Dnvcc=zig/kernels/cuda/deepseek_v41/nvcc-docker.sh` builds the images themselves through the same container. + +## Stored fixtures checked natively + +Paths under `zig/src/families/deepseek_v41/`. Host suites (`zig build test-dsv41`, part of `zig build test`; no GPU): +- `fixtures/widths-exl3-2p8.txt` and `fixtures/config.json`: a 2.8 bpw pack's expert widths and release config + (`host_test.zig`). +- `fixtures/lanes-ref.bin`: the run-time lane layout (`prepare.zig`). +- `fixtures/engram-hash-ref.bin` and `fixtures/engram-rows-ref.txt`: Engram's hasher and rows (`engram_host.zig`, + `engram_rows_test.zig`). +- Not stored: the device kernel fixtures (`cuda_tests/oracle.py`, `cuda_tests/dsv41_capture.py` in the tools) + that `tf-dsv41-test all DIR` and `tf-dsv41-test replay DIR` compare on a GPU. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 5fadb17f06..b18db4380b 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -105,6 +105,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("glm_cuda", glm_cuda); const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }); b.installArtifact(runner_exe); + b.installArtifact(@import("dsv41.zig").runner(b, target, optimize, cuda, strip)); // tf-tp-watch-check: the ranks' lines with no GPU or NCCL; a failing or dying rank stops every rank b.installArtifact(b.addExecutable(.{ .name = "tf-tp-watch-check", .root_module = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/tp_watch_check.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{.{ .name = "cuda", .module = cuda }} }) })); flashnextChecks(b, target, optimize, cuda, mods.core, mods.lanes, mods.qwen_vision); diff --git a/zig/build/dsv41.zig b/zig/build/dsv41.zig new file mode 100644 index 0000000000..350306a089 --- /dev/null +++ b/zig/build/dsv41.zig @@ -0,0 +1,68 @@ +//! DeepSeek-V4.1's build: the family's host side (zig/src/families/deepseek_v41) and its host tests on any OS. + +const std = @import("std"); + +/// The backend-neutral core (the safetensors reader the pack headers use) for the family's host side. +fn core(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *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 m = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("tokenizer", tokenizer); + return m; +} + +/// The CUDA runtime without kernel images: the family's launch bindings and device loads compile on any host. +fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const options = b.addOptions(); + options.addOption(bool, "with_kernels", false); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + cuda.addOptions("kernel_options", options); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); + return cuda; +} + +/// The family's kernel bindings and loader (`dsv41_kernels`, kernels.zig) over `cuda`, without fatbins (host-only). +pub fn kernels(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module { + const options = b.addOptions(); + options.addOption(bool, "with_kernels", false); + options.addOption(usize, "dsv41_images", 0); // also keeps this options file apart from the runtime's + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/kernels.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addOptions("dsv41_kernel_options", options); + // the optional images (keyed sampler, two-tile linears) are embedded as empty files + const empty = b.addWriteFiles().add("empty.fatbin", ""); + for ([_][]const u8{ "sampling", "lin2" }) |n| m.addAnonymousImport(b.fmt("dsv41_fatbin_dsv41_{s}", .{n}), .{ .root_source_file = empty }); + return m; +} + +/// The family's host side (config, pack headers, a rank's weight plan, kernel bindings) as the `deepseek_v41` module. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/deepseek_v41.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const cuda = runtime(b, target, optimize); + m.addImport("core", core(b, target, optimize)); + m.addImport("cuda", cuda); + m.addImport("dsv41_kernels", kernels(b, target, optimize, cuda)); + return m; +} + +/// `tf-dsv41-test`, the device test runner (zig/tests/cuda/dsv41_main.zig), over `cuda` and host-only family kernels. +pub fn runner(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, strip: bool) *std.Build.Step.Compile { + const m = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/dsv41_main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); + m.addImport("cuda", cuda); + m.addImport("dsv41_kernels", kernels(b, target, optimize, cuda)); + return b.addExecutable(.{ .name = "tf-dsv41-test", .root_module = m }); +} + +pub fn add(b: *std.Build, test_step: *std.Build.Step) void { + const host = b.graph.host; + const step = b.step("test-dsv41", "DeepSeek-V4.1's host-side unit tests (no GPU work)"); + const family = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug) })); + const bindings = b.addRunArtifact(b.addTest(.{ .root_module = kernels(b, host, .debug, runtime(b, host, .debug)) })); + for ([_]*std.Build.Step.Run{ family, bindings }) |run| { + step.dependOn(&run.step); + test_step.dependOn(&run.step); + } + // the device runner compiles on any host; running it needs a GPU and every family image + const compiled = &runner(b, host, .debug, runtime(b, host, .debug), false).step; + step.dependOn(compiled); + test_step.dependOn(compiled); +} diff --git a/zig/build/sessions.zig b/zig/build/sessions.zig new file mode 100644 index 0000000000..5dd2331845 --- /dev/null +++ b/zig/build/sessions.zig @@ -0,0 +1,20 @@ +//! Session memory: portable bookkeeping tests everywhere, NVMe tier tests on Linux. + +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/sessions/sessions.zig"), .target = target, .optimize = optimize, .link_libc = true }); +} + +pub fn add(b: *std.Build, test_step: *std.Build.Step) void { + const host = b.graph.host; + const tests = if (host.result.os.tag == .linux) module(b, host, .debug) else b.createModule(.{ + .root_source_file = b.path("zig/src/sessions/portable_test.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + }); + const run = b.addRunArtifact(b.addTest(.{ .root_module = tests })); + b.step("test-sessions", "Session memory's host tests: pool, prefix trie, RAM tier, NVMe tier on Linux (no GPU)").dependOn(&run.step); + test_step.dependOn(&run.step); +} diff --git a/zig/build/tp.zig b/zig/build/tp.zig new file mode 100644 index 0000000000..e07230abbc --- /dev/null +++ b/zig/build/tp.zig @@ -0,0 +1,88 @@ +//! Tensor parallelism's build: Linux host tests, and on Linux the `tf-tp-test` GPU runner with the mailbox kernel. + +const std = @import("std"); + +/// The CUDA runtime without the engine's kernels (tp only needs the driver, memory, streams and NCCL). +fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const options = b.addOptions(); + options.addOption(bool, "with_kernels", false); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + cuda.addOptions("kernel_options", options); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); + return cuda; +} + +/// The `tp` module over `cuda` (an engine passes its own runtime module); verbs (zig/src/fabric/verbs.zig) for RoCE. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/tp.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("verbs", b.createModule(.{ .root_source_file = b.path("zig/src/fabric/verbs.zig"), .target = target, .optimize = optimize, .link_libc = true })); + return m; +} + +/// mailbox.cu as a fatbin: sm_120 (RTX PRO 6000), sm_121 (GB10), and compute_90 PTX for anything else. +fn mailboxFatbin(b: *std.Build, nvcc: []const u8) std.Build.LazyPath { + const run = b.addSystemCommand(&.{ nvcc, "-fatbin", "-O3", "-std=c++17", "-gencode=arch=compute_120,code=sm_120", "-gencode=arch=compute_121,code=sm_121", "-gencode=arch=compute_90,code=compute_90", "-o" }); + const out = run.addOutputFileArg("tp_mailbox.fatbin"); + run.addFileArg(b.path("zig/src/fabric/tp/mailbox.cu")); + return out; +} + +/// `pub const mailbox: []const u8`, the fatbin's bytes or empty. +fn kernelModule(b: *std.Build, image: ?std.Build.LazyPath) *std.Build.Module { + const wf = b.addWriteFiles(); + if (image) |file| { + _ = wf.addCopyFile(file, "tp_mailbox.fatbin"); + // the driver takes 8-byte aligned images (Module.load refuses others); @embedFile alone promises nothing + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "const bytes align(16) = @embedFile(\"tp_mailbox.fatbin\").*;\npub const mailbox: []const u8 = &bytes;\n") }); + } + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "pub const mailbox: []const u8 = \"\";\n") }); +} + +pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, test_step: *std.Build.Step) void { + const host = b.graph.host; + const cli_tests = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_cli.zig"), .target = host, .optimize = optimize }) })); + const cli_step = b.step("test-tp-cli", "TP runner command validation (no GPU)"); + cli_step.dependOn(&cli_tests.step); + test_step.dependOn(&cli_tests.step); + if (host.result.os.tag == .linux) { + const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); + const tp_step = b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)"); + tp_step.dependOn(&tests.step); + tp_step.dependOn(cli_step); + test_step.dependOn(&tests.step); + } + if (target.result.os.tag != .linux) return; + const nvcc = b.option([]const u8, "tp-nvcc", "nvcc that builds the TP mailbox kernel"); + const prebuilt = b.option([]const u8, "tp-fatbin", "absolute path of a prebuilt tp_mailbox.fatbin to embed"); + const image: ?std.Build.LazyPath = if (prebuilt) |p| b.graph.cwdRelativePath(p) else if (nvcc) |tool| mailboxFatbin(b, tool) else null; + const cuda = runtime(b, target, optimize); + const runner = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + runner.addImport("cuda", cuda); + runner.addImport("tp", module(b, target, optimize, cuda)); + runner.addImport("tp_kernel", kernelModule(b, image)); + const exe = b.addExecutable(.{ .name = "tf-tp-test", .root_module = runner }); + if (host.result.os.tag == .linux and target.query.isNative()) { + const cli = @import("../src/fabric/tp/gpu_cli.zig"); + for (cli.invalid) |args| { + const check = b.addRunArtifact(exe); + check.addArgs(args); + check.setEnvironmentVariable("TF_TP_WORLD", "invalid"); + check.expectExitCode(2); + check.expectStdErrEqual(cli.usage); + cli_step.dependOn(&check.step); + test_step.dependOn(&check.step); + } + if (image == null) { + const kernel = b.addRunArtifact(exe); + kernel.addArg("kernel"); + kernel.setEnvironmentVariable("TF_TP_WORLD", "invalid"); + kernel.expectExitCode(1); + kernel.expectStdErrEqual("mailbox kernel 0 B\n"); + cli_step.dependOn(&kernel.step); + test_step.dependOn(&kernel.step); + } + } + b.installArtifact(exe); + b.step("tf-tp-test", "TP=2 on two GPUs: bootstrap, every collective bit-exact, graph capture, fail-fast, latency").dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/kernels/cuda/deepseek_v41/attn_cuda.cu b/zig/kernels/cuda/deepseek_v41/attn_cuda.cu new file mode 100644 index 0000000000..ec43c81ecb --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/attn_cuda.cu @@ -0,0 +1,541 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/csa2/attn_cuda.cu (blob f4be064c7368): sass_check.sh +#define DSV41_ATTN_NO_TORCH + + + +// CSA2 decode attention (chunks + merge) in ONE launch, TF_DSV41_ATTN_CUDA=1: the Triton pair's ops, same bits +#include +#include +#include +#include + +namespace dsv41_attn { + +constexpr int D = 512, CH = 128, NCOMP = 4, NCH = 5, BMQ = 16, WINDOW = 128; // tiles of 32 keys +constexpr int NOPE = 448, VB = 576, SB = 8, NT = 256; +constexpr int SM_Q = 0; // q A fragments [32 steps][32 lanes] uint4 16384 +constexpr int SM_KV = 16384; // 128 rows x 576 B, 8-byte pieces swizzled by key & 7 73728 +constexpr int SM_SC = SM_KV + CH * VB; // 128 x 8 scale bytes 1024 +constexpr int SM_P = SM_SC + CH * SB; // P A fragments [4 tiles][2 steps][32 lanes] uint4 4096 +constexpr int SM_TMAX = SM_P + 4096; // [4 tiles][4 n8][16 rows] fp32 1024 +constexpr int SM_WSUM = SM_TMAX + 1024; // same shape: the warps' row sums 1024 +constexpr int SM_ROW = SM_WSUM + 1024; // [4 tiles][16 rows] nm, alpha, active 768 +constexpr int SM_PHYS = SM_ROW + 768; // 128 physical rows (int32, -1: masked) 512 +constexpr int SM_MISC = SM_PHYS + 512; // final m, l [16] each + the ticket flag 256 +constexpr int SMEM = SM_MISC + 256; // 98,816 B (sm_121: <= 101,376 a block) +static_assert(SMEM <= 101376, "shared memory"); +constexpr unsigned FULL = 0xffffffffu; + +struct Args { + const uint16_t* q; // [R, H, 512] bf16 (RoPE'd) + const uint8_t* cv; const uint8_t* csc; long long cvs, css; // compressed rows (values, scales), row strides + const int* tok; long long ts; const int* cnt; // selection [R, ts] int32 (-1 pad), counts [R] + const uint8_t* sv; const uint8_t* ssc; long long svs, sss; // SWA rings + const int* lo; const int* hi; // [R] window starts; window ends (HAS_HI) or null + const void* pos; // int32 [1] (POS + r) or int64 [R] (ROWS) + const long long* sl; // ROWS: the rows' slots [R] + const int* pt; long long pts; int psh; // page table of the compressed rows (psh 0: contiguous) + const float* sink; const float* cs; long long csst; // [H] fp32, RoPE table fp32 [max_pos, 64] + float* po; float* pm; float* pl; // partials (attn.Scratch layout) + uint16_t* out; // [R, H, 512] bf16 + int* ticket; // [R, H / 16] (self-cleaning) + int R, H, ring, has_comp; + int qrope; // split_kernel: q arrives un-rotated, RoPE'd as it loads +}; + +__device__ __forceinline__ float ex2(float x) { + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float div_full(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float texp(float x) { return ex2(__fmul_rn(x, __int_as_float(0x3FB8AA3B))); } +__device__ __forceinline__ float pow2b(uint32_t s) { return __int_as_float(static_cast(s) << 23); } + +__device__ __forceinline__ void mma(float (&d)[4], const uint4& a, uint32_t b0, uint32_t b1) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b0), "r"(b1)); +} + +// two e4m3 bytes (low byte = the lower k) with their own scales -> bf16x2 (exact: e4m3 x 2^e is a bf16) +__device__ __forceinline__ uint32_t deq2(uint32_t two, float s0, float s1) { + uint32_t h2; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(h2) : "h"(static_cast(two))); + const __half2 h = *reinterpret_cast(&h2); + const __nv_bfloat162 b = __floats2bfloat162_rn(__fmul_rn(__low2float(h), s0), __fmul_rn(__high2float(h), s1)); + return *reinterpret_cast(&b); +} + +// a byte offset x of row `key` in the staged rows: 8-byte pieces swizzled within groups of 8 (bank-conflict free) +__device__ __forceinline__ int swz(int key, int x) { + const int ch = x >> 3; + return key * VB + ((((ch & ~7) | ((ch & 7) ^ (key & 7)))) << 3) + (x & 7); +} + +// Triton's dot-operand k order: step s (16 k) of a 32-k block takes k = klo(s), klo + 1 and khi(s), khi + 1 +template __device__ __forceinline__ int klo(int s, int t) { + return KW == 4 ? 32 * (s >> 1) + 4 * t + 2 * (s & 1) : 16 * s + 2 * t; +} +template __device__ __forceinline__ int khi(int s, int t) { return klo(s, t) + (KW == 4 ? 16 : 8); } +// where element k (< 32) of a row sits in a 2-step A fragment: step, quad lane, high half (+16 / +8) +template __device__ __forceinline__ void frag_of(int k, int& step, int& tq, int& hi) { + if (KW == 4) { step = (k >> 1) & 1; hi = (k >> 4) & 1; tq = (k & 15) >> 2; } + else { step = k >> 4; hi = (k >> 3) & 1; tq = (k & 7) >> 1; } +} + +// the bf16 pair (k, k + 1) of staged row `key` (S's B operand: k along the row) +__device__ __forceinline__ uint32_t pair_k(const uint8_t* sm, int key, int k) { + if (k < NOPE) { + const uint32_t two = *reinterpret_cast(sm + SM_KV + swz(key, k)); + const float s = pow2b(sm[SM_SC + key * SB + (k >> 6)]); + return deq2(two, s, s); + } + return *reinterpret_cast(sm + SM_KV + swz(key, NOPE + 2 * (k - NOPE))); +} + +// value (key, n) of the staged rows as a bf16 bit pattern / its e4m3 byte (PV's B operand: k along the keys) +__device__ __forceinline__ uint32_t pair_n(const uint8_t* sm, int k0, int k1, int n) { + if (n < NOPE) { + const uint32_t two = sm[SM_KV + swz(k0, n)] | (static_cast(sm[SM_KV + swz(k1, n)]) << 8); + return deq2(two, pow2b(sm[SM_SC + k0 * SB + (n >> 6)]), pow2b(sm[SM_SC + k1 * SB + (n >> 6)])); + } + const int x = NOPE + 2 * (n - NOPE); + return *reinterpret_cast(sm + SM_KV + swz(k0, x)) | + (static_cast(*reinterpret_cast(sm + SM_KV + swz(k1, x))) << 16); +} + +__device__ __forceinline__ void cp8(uint32_t dst, const void* src, bool ok) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8, %2;" :: "r"(dst), "l"(src), "r"(ok ? 8 : 0)); +} + +template __device__ __forceinline__ int row_pos(const Args& a, int r) { + return ROWS ? static_cast(static_cast(a.pos)[r]) : static_cast(a.pos)[0] + r; +} + +// how many chunks of row r have keys (the merge's ticket target; the others are Triton's empty partial) +__device__ __forceinline__ bool chunk_works(const Args& a, int r, int c) { + return c == NCOMP || (a.has_comp && c * CH < a.cnt[r]); +} + +template +__device__ void merge(const Args& a, uint8_t* sm, int r, int hb) { + const int tid = threadIdx.x, H = a.H, R = a.R; + float* ca = reinterpret_cast(sm + SM_TMAX); // [5][16] a, [5][16] b, [16] af, [16] den + float* cb = ca + NCH * BMQ; + float* fa = cb + NCH * BMQ; + float* fd = fa + BMQ; + if (tid < BMQ) { // a head's scalar chain (Triton's _merge, unrolled) + const int h = hb * BMQ + tid; + float m = -INFINITY, l = 0.f; + for (int c = 0; c < NCH; ++c) { + const bool w = chunk_works(a, r, c); + const long long base = (static_cast(c) * R + r) * H + h; + const float cm = w ? __ldcg(a.pm + base) : -INFINITY, cl = w ? __ldcg(a.pl + base) : 0.f; + const bool act = cl > 0.f; + float aa, bb; + if (c == 0) { + const float nm = act ? fmaxf(cm, -INFINITY) : -INFINITY; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + aa = 0.f; + l = __fmaf_rn(bb, cl, 0.f); + m = nm; + } else { + const float nm = act ? fmaxf(m, cm) : m; + aa = act ? (m == -INFINITY ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + l = __fmaf_rn(l, aa, __fmul_rn(bb, cl)); + m = nm; + } + ca[c * BMQ + tid] = aa; + cb[c * BMQ + tid] = bb; + } + const float sink = a.sink[h]; + const float top = fmaxf(m, sink); + const float af = m == -INFINITY ? 0.f : texp(__fsub_rn(m, top)); + const float e = sink == -INFINITY ? 0.f : texp(__fsub_rn(sink, top)); + fa[tid] = af; + fd[tid] = __fmaf_rn(l, af, e); + } + __syncthreads(); + const int at = row_pos(a, r); + const float* csr = a.cs + static_cast(at) * a.csst; + for (int u = tid; u < BMQ * (D / 4); u += NT) { // float4 units: head hl, columns d .. d + 3 + const int hl = u / (D / 4), d = (u % (D / 4)) * 4, h = hb * BMQ + hl; + float o[4]; + for (int c = 0; c < NCH; ++c) { + float4 v = make_float4(0.f, 0.f, 0.f, 0.f); + if (chunk_works(a, r, c)) + v = __ldcg(reinterpret_cast(a.po + ((static_cast(c) * R + r) * H + h) * D + + d)); + const float co[4] = {v.x, v.y, v.z, v.w}; + const float aa = ca[c * BMQ + hl], bb = cb[c * BMQ + hl]; + for (int i = 0; i < 4; ++i) + o[i] = c == 0 ? __fmaf_rn(bb, co[i], 0.f) : __fmaf_rn(o[i], aa, __fmul_rn(co[i], bb)); + } + float x[4]; + for (int i = 0; i < 4; ++i) x[i] = div_full(__fmul_rn(o[i], fa[hl]), fd[hl]); + uint16_t ob[4]; + for (int i = 0; i < 4; i += 2) { // pair (d + i, d + i + 1): GPT-J, inverse + const int p = (d + i) >> 1; + const float c = p >= 224 ? csr[p - 224] : 1.f, s = p >= 224 ? csr[32 + p - 224] : 0.f; + const float ne = __fmaf_rn(x[i], c, __fmul_rn(x[i + 1], s)); + const float no = __fmaf_rn(x[i + 1], c, -__fmul_rn(x[i], s)); + const __nv_bfloat16 b0 = __float2bfloat16_rn(ne), b1 = __float2bfloat16_rn(no); + ob[i] = *reinterpret_cast(&b0); + ob[i + 1] = *reinterpret_cast(&b1); + } + uint2 w; + w.x = ob[0] | (static_cast(ob[1]) << 16); + w.y = ob[2] | (static_cast(ob[3]) << 16); + *reinterpret_cast(a.out + (static_cast(r) * H + h) * D + d) = w; + } +} + +template +__global__ void __launch_bounds__(NT, 1) attn_kernel(Args a) { + extern __shared__ __align__(16) uint8_t sm[]; + const int c = blockIdx.x, hb = blockIdx.y, r = blockIdx.z; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, t = lane & 3; + const int H = a.H, R = a.R; + if (!chunk_works(a, r, c)) return; + int* phys = reinterpret_cast(sm + SM_PHYS); + const long long sl = ROWS ? a.sl[r] : 0; + if (tid < CH) { // the chunk's 128 keys -> physical rows (-1: masked) + int row = -1; + if (c < NCOMP) { + const int idx = c * CH + tid; + if (idx < a.cnt[r]) { + const int tk = a.tok[static_cast(r) * a.ts + idx]; + if (tk >= 0) { + row = tk; + if (a.psh > 0) { + const int* ptr = a.pt + (ROWS ? sl * a.pts : 0); + row = (ptr[tk >> a.psh] << a.psh) + (tk & ((1 << a.psh) - 1)); + } + } + } + } else { + const int hi = HAS_HI ? a.hi[r] : row_pos(a, r); + const int lo = max(a.lo[r], hi - (WINDOW - 1)); + const int p = hi - (WINDOW - 1) + tid; + if (p >= lo && p >= 0) row = (p % a.ring) + (ROWS ? static_cast(sl) * a.ring : 0); + } + phys[tid] = row; + } + { // q -> A fragments (Triton's k order) + const uint16_t* qb = a.q + (static_cast(r) * H + hb * BMQ) * D; + for (int i = tid; i < 32 * 32; i += NT) { + const int s = i >> 5, ln = i & 31, gg = ln >> 2, tt = ln & 3; + const int k0 = klo(s, tt), k1 = khi(s, tt); + uint4 f; + f.x = *reinterpret_cast(qb + gg * D + k0); + f.y = *reinterpret_cast(qb + (gg + 8) * D + k0); + f.z = *reinterpret_cast(qb + gg * D + k1); + f.w = *reinterpret_cast(qb + (gg + 8) * D + k1); + reinterpret_cast(sm + SM_Q)[i] = f; + } + } + __syncthreads(); + { // the rows: 72 value pieces + 1 scale piece a key + const bool comp = c < NCOMP; + const uint8_t* V = comp ? a.cv : a.sv; + const uint8_t* S = comp ? a.csc : a.ssc; + const long long vs = comp ? a.cvs : a.svs, ss = comp ? a.css : a.sss; + const uint32_t base = static_cast(__cvta_generic_to_shared(sm)); + for (int i = tid; i < CH * 73; i += NT) { + const int key = i / 73, pc = i % 73, row = phys[key]; + const bool ok = row >= 0; + const long long rr = ok ? row : 0; + if (pc < 72) cp8(base + SM_KV + swz(key, pc * 8), V + rr * vs + pc * 8, ok); + else cp8(base + SM_SC + key * SB, S + rr * ss, ok); + } + asm volatile("cp.async.commit_group;\ncp.async.wait_group 0;" ::: "memory"); + } + __syncthreads(); + const float NINF = -INFINITY; + const int n8 = warp & 3, ta = warp >> 2; // this warp: S of tiles ta and ta + 2, block n8 + float sv[2][4] = {{0.f, 0.f, 0.f, 0.f}, {0.f, 0.f, 0.f, 0.f}}; + { + const int kA = 32 * ta + 8 * n8 + g, kB = kA + 64; + const uint4* qf = reinterpret_cast(sm + SM_Q); +#pragma unroll 4 + for (int s = 0; s < 32; ++s) { + const uint4 af = qf[s * 32 + lane]; + const int k0 = klo(s, t), k1 = khi(s, t); + mma(sv[0], af, pair_k(sm, kA, k0), pair_k(sm, kA, k1)); + mma(sv[1], af, pair_k(sm, kB, k0), pair_k(sm, kB, k1)); + } + } + float* tmax = reinterpret_cast(sm + SM_TMAX); + float* wsum = reinterpret_cast(sm + SM_WSUM); + float* rowv = reinterpret_cast(sm + SM_ROW); // [tile][row]: nm, alpha, active at 0 / 64 / 128 + bool val[2][2]; + for (int j = 0; j < 2; ++j) { + const int tile = ta + 2 * j, kc = 32 * tile + 8 * n8 + 2 * t; + val[j][0] = phys[kc] >= 0; + val[j][1] = phys[kc + 1] >= 0; + for (int e = 0; e < 4; ++e) + sv[j][e] = val[j][e & 1] ? __fmul_rn(sv[j][e], __int_as_float(0x3D3504F3)) : NINF; + float m0 = fmaxf(sv[j][0], sv[j][1]), m1 = fmaxf(sv[j][2], sv[j][3]); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 2)); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 1)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 2)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 1)); + if (t == 0) { + tmax[(tile * 4 + n8) * 16 + g] = m0; + tmax[(tile * 4 + n8) * 16 + g + 8] = m1; + } + } + __syncthreads(); + if (tid < BMQ) { // the running max over the tiles, a row a thread + float m = NINF; + for (int tile = 0; tile < 4; ++tile) { + const float* tm4 = tmax + tile * 64 + tid; + const float tm = fmaxf(fmaxf(tm4[0], tm4[16]), fmaxf(tm4[32], tm4[48])); + const bool act = tm != NINF; + const float nm = act ? fmaxf(m, tm) : m; + const float al = act ? (m == NINF ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + rowv[tile * 16 + tid] = nm; + rowv[64 + tile * 16 + tid] = al; + rowv[128 + tile * 16 + tid] = act ? 1.f : 0.f; + m = nm; + } + reinterpret_cast(sm + SM_MISC)[tid] = m; + } + __syncthreads(); + uint32_t* pf = reinterpret_cast(sm + SM_P); // [tile][step][lane][4] + for (int j = 0; j < 2; ++j) { // p, the thread / quad sums, bf16(p) -> fragments + const int tile = ta + 2 * j; + float p[4]; + for (int e = 0; e < 4; ++e) { + const int row = g + (e >> 1) * 8; + const bool act = rowv[128 + tile * 16 + row] != 0.f; + p[e] = (val[j][e & 1] && act) ? texp(__fsub_rn(sv[j][e], rowv[tile * 16 + row])) : 0.f; + } + float a0 = __fadd_rn(p[0], p[1]), a1 = __fadd_rn(p[2], p[3]); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 2)); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 1)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 2)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 1)); + if (t == 0) { + wsum[(tile * 4 + n8) * 16 + g] = a0; + wsum[(tile * 4 + n8) * 16 + g + 8] = a1; + } + int step, tq, hi; + frag_of(8 * n8 + 2 * t, step, tq, hi); + const __nv_bfloat162 lo2 = __floats2bfloat162_rn(p[0], p[1]), hi2 = __floats2bfloat162_rn(p[2], p[3]); + uint32_t* f = pf + ((tile * 2 + step) * 32 + g * 4 + tq) * 4; + f[2 * hi] = *reinterpret_cast(&lo2); // reg 0 / 2: row g; 1 / 3: row g + 8 + f[2 * hi + 1] = *reinterpret_cast(&hi2); + } + __syncthreads(); + if (tid < BMQ) { // l: fma(l, alpha, (W0 + W2) + (W1 + W3)) + float l = 0.f; + for (int tile = 0; tile < 4; ++tile) { + const float* w4 = wsum + tile * 64 + tid; + const float sum = __fadd_rn(__fadd_rn(w4[0], w4[32]), __fadd_rn(w4[16], w4[48])); + l = __fmaf_rn(l, rowv[64 + tile * 16 + tid], sum); + } + reinterpret_cast(sm + SM_MISC)[16 + tid] = l; + } + float acc[8][4]; + for (int j = 0; j < 8; ++j) + for (int e = 0; e < 4; ++e) acc[j][e] = 0.f; + for (int tile = 0; tile < 4; ++tile) { // o = mma(P, kv, rn(o * alpha)), 2 k steps a tile + const float al0 = rowv[64 + tile * 16 + g], al1 = rowv[64 + tile * 16 + g + 8]; + for (int j = 0; j < 8; ++j) { + acc[j][0] = __fmul_rn(acc[j][0], al0); + acc[j][1] = __fmul_rn(acc[j][1], al0); + acc[j][2] = __fmul_rn(acc[j][2], al1); + acc[j][3] = __fmul_rn(acc[j][3], al1); + } + for (int step = 0; step < 2; ++step) { + const uint4 af = reinterpret_cast(pf)[(tile * 2 + step) * 32 + lane]; + const int k0 = 32 * tile + klo(step, t), k1 = 32 * tile + khi(step, t); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int n = 64 * j + 8 * warp + g; + mma(acc[j], af, pair_n(sm, k0, k0 + 1, n), pair_n(sm, k1, k1 + 1, n)); + } + } + } + const long long pbase = (static_cast(c) * R + r) * H + hb * BMQ; + for (int j = 0; j < 8; ++j) { + const int n = 64 * j + 8 * warp + 2 * t; + *reinterpret_cast(a.po + (pbase + g) * D + n) = make_float2(acc[j][0], acc[j][1]); + *reinterpret_cast(a.po + (pbase + g + 8) * D + n) = make_float2(acc[j][2], acc[j][3]); + } + if (tid < BMQ) { + a.pm[pbase + tid] = reinterpret_cast(sm + SM_MISC)[tid]; + a.pl[pbase + tid] = reinterpret_cast(sm + SM_MISC)[16 + tid]; + } + int* flag = reinterpret_cast(sm + SM_MISC) + 32; + __threadfence(); + __syncthreads(); + if (tid == 0) { + int need = 0; + for (int cc = 0; cc < NCH; ++cc) need += chunk_works(a, r, cc) ? 1 : 0; + *flag = atomicAdd(a.ticket + r * (H / BMQ) + hb, 1) == need - 1; + } + __syncthreads(); + if (!*flag) return; + __threadfence(); + if (tid == 0) atomicExch(a.ticket + r * (H / BMQ) + hb, 0); // self-cleaning: the next launch / replay + merge(a, sm, r, hb); +} + +template +cudaError_t launch_one(const Args& a, cudaStream_t stream) { + static bool init = false; + if (!init) { + const cudaError_t e = cudaFuncSetAttribute(attn_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM); + if (e != cudaSuccess) return e; + init = true; + } + attn_kernel<<>>(a); + return cudaGetLastError(); +} + +#include "attn_split.cuh" + +template +cudaError_t launch_split(const Args& a, cudaStream_t stream) { + using G = Split

; + static bool init = false; + if (!init) { + const cudaError_t e = cudaFuncSetAttribute(split_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, G::SMEM); + if (e != cudaSuccess) return e; + init = true; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(NCH * P, a.H / BMQ, a.R); + cfg.blockDim = dim3(NT); + cfg.dynamicSmemBytes = G::SMEM; + cfg.stream = stream; + cudaLaunchAttribute at[1]; + at[0].id = cudaLaunchAttributeClusterDimension; + at[0].val.clusterDim.x = P; + at[0].val.clusterDim.y = 1; + at[0].val.clusterDim.z = 1; + cfg.attrs = at; + cfg.numAttrs = P > 1 ? 1 : 0; + return cudaLaunchKernelEx(&cfg, split_kernel, a); +} + +template inline cudaError_t dispatch_split(const Args& a, int kw, bool rows, bool has_hi, cudaStream_t s) { + switch ((kw == 4 ? 4 : 0) + (rows ? 2 : 0) + (has_hi ? 1 : 0)) { + case 0: return launch_split<2, false, false, P>(a, s); + case 1: return launch_split<2, false, true, P>(a, s); + case 2: return launch_split<2, true, false, P>(a, s); + case 3: return launch_split<2, true, true, P>(a, s); + case 4: return launch_split<4, false, false, P>(a, s); + case 5: return launch_split<4, false, true, P>(a, s); + case 6: return launch_split<4, true, false, P>(a, s); + default: return launch_split<4, true, true, P>(a, s); + } +} + +// split: 0 attn_kernel (the default launch), 1 / 4 split_kernel with that many CTAs a chunk (the same bits) +inline cudaError_t dispatch(const Args& a, int kw, bool rows, bool has_hi, cudaStream_t stream, int split = 0) { + if (a.R < 1 || a.H % BMQ || a.ring < WINDOW || (a.ring & (a.ring - 1))) return cudaErrorInvalidValue; + if (split == 1) return dispatch_split<1>(a, kw, rows, has_hi, stream); + if (split == 4) return dispatch_split<4>(a, kw, rows, has_hi, stream); + if (split != 0) return cudaErrorInvalidValue; + switch ((kw == 4 ? 4 : 0) + (rows ? 2 : 0) + (has_hi ? 1 : 0)) { + case 0: return launch_one<2, false, false>(a, stream); + case 1: return launch_one<2, false, true>(a, stream); + case 2: return launch_one<2, true, false>(a, stream); + case 3: return launch_one<2, true, true>(a, stream); + case 4: return launch_one<4, false, false>(a, stream); + case 5: return launch_one<4, false, true>(a, stream); + case 6: return launch_one<4, true, false>(a, stream); + default: return launch_one<4, true, true>(a, stream); + } +} + +} // namespace dsv41_attn + +#ifndef DSV41_ATTN_NO_TORCH // the compile test builds the kernels alone (nvcc, no torch headers) +#include +#include +#include + +namespace { +template T* tp(const at::Tensor& t) { return t.numel() ? static_cast(t.data_ptr()) : nullptr; } +} + +void dsv41_attn_run_cuda(const at::Tensor& q, const at::Tensor& cv, const at::Tensor& csc, const at::Tensor& tok, + const at::Tensor& cnt, const at::Tensor& sv, const at::Tensor& ssc, const at::Tensor& lo, + const at::Tensor& hi, const at::Tensor& pos, const at::Tensor& sl, const at::Tensor& pt, + const at::Tensor& sink, const at::Tensor& cs, at::Tensor& po, at::Tensor& pm, + at::Tensor& pl, at::Tensor& out, at::Tensor& ticket, int64_t ring, int64_t psh, + int64_t pts, int64_t kw, int64_t split, int64_t qrope) { + using namespace dsv41_attn; + Args a{}; + a.q = tp(q); + a.has_comp = cv.numel() ? 1 : 0; + a.cv = tp(cv); a.csc = tp(csc); + a.cvs = cv.numel() ? cv.stride(0) : 0; a.css = csc.numel() ? csc.stride(0) : 0; + a.tok = tp(tok); a.ts = tok.numel() ? tok.stride(0) : 0; a.cnt = tp(cnt); + a.sv = tp(sv); a.ssc = tp(ssc); a.svs = sv.stride(0); a.sss = ssc.stride(0); + a.lo = tp(lo); a.hi = tp(hi); + a.pos = pos.data_ptr(); + a.sl = tp(sl); + a.pt = tp(pt); a.pts = pts; a.psh = static_cast(psh); + a.sink = tp(sink); a.cs = tp(cs); a.csst = cs.stride(0); + a.po = tp(po); a.pm = tp(pm); a.pl = tp(pl); + a.out = tp(out); a.ticket = tp(ticket); + a.R = static_cast(q.size(0)); a.H = static_cast(q.size(1)); a.ring = static_cast(ring); + a.qrope = qrope && split > 0 ? 1 : 0; + TORCH_CHECK(!qrope || split > 0, "attn_cuda: q's RoPE folds into the split kernel only"); + C10_CUDA_CHECK(dispatch(a, static_cast(kw), sl.numel() > 0, hi.numel() > 0, + at::cuda::getCurrentCUDAStream(), static_cast(split))); +} + +std::vector dsv41_attn_info_cuda(int64_t device) { // GB10's limits for the two launches + int sms = 0, optin = 0, cl = 0, blocks = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, static_cast(device)); + cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, static_cast(device)); + cudaDeviceGetAttribute(&cl, cudaDevAttrClusterLaunch, static_cast(device)); + cudaFuncSetAttribute(dsv41_attn::attn_kernel<4, true, false>, cudaFuncAttributeMaxDynamicSharedMemorySize, + dsv41_attn::SMEM); + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, dsv41_attn::attn_kernel<4, true, false>, dsv41_attn::NT, + dsv41_attn::SMEM); + return {sms, optin, cl, blocks, dsv41_attn::SMEM}; +} +#endif + +// The instantiations dispatch() launches: attn_kernel. +template __global__ void dsv41_attn::attn_kernel<2, false, false>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<2, false, true>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<2, true, false>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<2, true, true>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<4, false, false>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<4, false, true>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<4, true, false>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<4, true, true>(dsv41_attn::Args); +// dispatch_split's split_kernel at P 1 and 4 (a cluster of 4) +template __global__ void dsv41_attn::split_kernel<2, false, false, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, false, true, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, true, false, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, true, true, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, false, false, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, false, true, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, true, false, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, true, true, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, false, false, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, false, true, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, true, false, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, true, true, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, false, false, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, false, true, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, true, false, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, true, true, 4>(dsv41_attn::Args); diff --git a/zig/kernels/cuda/deepseek_v41/attn_split.cuh b/zig/kernels/cuda/deepseek_v41/attn_split.cuh new file mode 100644 index 0000000000..cd70b1594a --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/attn_split.cuh @@ -0,0 +1,538 @@ +// attn_cuda.cu's split_kernel (TF_DSV41_ATTN_SPLIT): its helpers and the cluster kernel, inside namespace dsv41_attn +#pragma once +// ---- split_kernel (TF_DSV41_ATTN_SPLIT=P, 1 / 4): a chunk over a CLUSTER of P CTAs, attn_kernel's bits at any P --- +constexpr int MAXP = 4; // the ticket's ranks a (row, head tile) + +__host__ __device__ constexpr int col2byte(int n) { return n <= NOPE ? n : NOPE + 2 * (n - NOPE); } + +template struct Split { + static constexpr int TPC = 4 / P; // S tiles a CTA + static constexpr int KS = 32 * TPC; // S keys a CTA + static constexpr int W = D / P; // PV / merge columns a CTA + static constexpr int J = 8 / P; // n8 blocks a warp in PV + static constexpr int NPW = TPC == 4 ? 2 : 1; // S chains a warp ((tile, n8) pairs: 4 TPC over 8 warps) + static constexpr int SLB = P == 1 ? VB : col2byte(D) - col2byte(D - W); // slice pitch (the widest rank's) + static constexpr int SM_Q = 0; // q A fragments 16384 + static constexpr int SM_KP = 16384; // PV slice: 128 keys x SLB (P = 1: the whole rows) + static constexpr int SM_KS = P == 1 ? SM_KP : SM_KP + CH * SLB; // S rows: KS keys x 576 (P = 1: the slice) + static constexpr int SM_SC = SM_KS + KS * VB; // 128 keys x 8 scale bytes 1024 + static constexpr int SM_KB = SM_SC + CH * SB; // P > 1: S's B operands as bf16 pairs [KS][256] words + static constexpr int SM_P = SM_KB + (P == 1 ? 0 : KS * D * 2); // P A fragments [4][2][32] uint4 4096 + static constexpr int SM_TMAX = SM_P + 4096; // [4 tiles][4 n8][16 rows] fp32 (merge: ca, cb, ...) + static constexpr int SM_WSUM = SM_TMAX + 1024; + static constexpr int SM_ROW = SM_WSUM + 1024; // [4 tiles][16 rows] nm, alpha, active + static constexpr int SM_PHYS = SM_ROW + 768; // 128 physical rows (-1: masked) + static constexpr int SM_MISC = SM_PHYS + 512; // final m, l [16] each + the ticket flag + static constexpr int SMEM = SM_MISC + 256; // P = 1: 98,816 B; P = 4: 100,864 B + static constexpr int SM_VB = SM_KB; // P > 1, after S: PV's B pairs [64 key pairs][W] words + static_assert(P == 1 || (CH / 2) * W * 4 <= KS * D * 2, "PV's pairs fit where S's were"); + static_assert(SMEM <= 101376, "shared memory"); + static_assert(SLB % 64 == 0 && W % 128 == 0, "16-byte swizzle groups; whole n8 blocks a warp"); +}; + +__device__ __forceinline__ int fsw(int key) { return ((key >> 1) ^ (key >> 3)) & 3; } +__device__ __forceinline__ int swz16(int key, int x, int pitch) { + const int c = x >> 4; + return key * pitch + (((c & ~3) | ((c & 3) ^ fsw(key))) << 4) + (x & 15); +} + +// S's B pair (k, k + 1) of the S row `key` (local to the S rows; `ck`: its chunk key, for the scale) +template +__device__ __forceinline__ uint32_t spair_k(const uint8_t* sm, int key, int ck, int k) { + using G = Split

; + if (k < NOPE) { + const uint32_t two = *reinterpret_cast(sm + G::SM_KS + swz16(key, k, VB)); + const float s = pow2b(sm[G::SM_SC + ck * SB + (k >> 6)]); + return deq2(two, s, s); + } + return *reinterpret_cast(sm + G::SM_KS + swz16(key, NOPE + 2 * (k - NOPE), VB)); +} + +// P > 1: S's B operands dequantized once into bf16 pairs, word k / 2 of S row `key` at w ^ bsw(key): conflict-free +template __device__ __forceinline__ int bsw(int key) { + return (KW == 4 ? (key & 1) : ((key & 1) << 2)) | (((key >> 1) & 3) << 3); +} +template +__device__ __forceinline__ uint32_t sbf_k(const uint8_t* sm, int key, int k) { + return reinterpret_cast(sm + Split

::SM_KB)[key * (D / 2) + ((k >> 1) ^ bsw(key))]; +} + +// PV's B pair: value (k0, n), (k1, n) of the slice (byte x0: the rank's first byte) +template +__device__ __forceinline__ uint32_t spair_n(const uint8_t* sm, int k0, int k1, int n, int x0) { + using G = Split

; + if (n < NOPE) { + const int x = n - x0; + const uint32_t two = sm[G::SM_KP + swz16(k0, x, G::SLB)] | + (static_cast(sm[G::SM_KP + swz16(k1, x, G::SLB)]) << 8); + return deq2(two, pow2b(sm[G::SM_SC + k0 * SB + (n >> 6)]), pow2b(sm[G::SM_SC + k1 * SB + (n >> 6)])); + } + const int x = col2byte(n) - x0; + return *reinterpret_cast(sm + G::SM_KP + swz16(k0, x, G::SLB)) | + (static_cast(*reinterpret_cast(sm + G::SM_KP + swz16(k1, x, G::SLB))) << 16); +} + +// P > 1: PV's B pairs dequantized once, word (key pair kp, column x) at x ^ vsw(kp): conflict-free +template __device__ __forceinline__ int vsw(int kp) { return (KW == 4 ? (kp >> 1) & 3 : kp & 3) << 3; } + +// TF_DSV41_ATTN_ROPE: q's rope pair j in place, (rn(e c) - rn(o s), rn(o c) + rn(e s)), no contraction, then bf16 +__device__ __forceinline__ uint32_t rope_word(uint32_t w, const float* csr, int j) { + const float e = __uint_as_float(w << 16), o = __uint_as_float(w & 0xffff0000u); + const float c = csr[j], sn = csr[32 + j]; + const __nv_bfloat162 b = __floats2bfloat162_rn(__fsub_rn(__fmul_rn(e, c), __fmul_rn(o, sn)), + __fadd_rn(__fmul_rn(o, c), __fmul_rn(e, sn))); + return *reinterpret_cast(&b); +} + +__device__ __forceinline__ void cp16(uint32_t dst, const void* src, bool ok) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" :: "r"(dst), "l"(src), "r"(ok ? 16 : 0)); +} +__device__ __forceinline__ void cl_arrive() { asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); } +__device__ __forceinline__ void cl_wait() { asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); } +__device__ __forceinline__ uint32_t cl_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} +// store word v at shared address `p` of cluster rank q (DSMEM) +__device__ __forceinline__ void st_peer(void* p, uint32_t q, uint32_t v) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)), b; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(b) : "r"(a), "r"(q)); + asm volatile("st.shared::cluster.u32 [%0], %1;" :: "r"(b), "r"(v) : "memory"); +} +// word i of shared array `p` in cluster rank q (DSMEM) +__device__ __forceinline__ uint32_t ld_peer(const void* p, uint32_t q) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)), b, v; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(b) : "r"(a), "r"(q)); + asm volatile("ld.shared::cluster.u32 %0, [%1];" : "=r"(v) : "r"(b) : "memory"); + return v; +} + +// stage row bytes [b0, b0 + 16 n) at key slot `key` (16-byte pieces, or 8-byte halves without v16); masked: zeros +__device__ __forceinline__ void stage_piece(uint32_t base, int region, int key, int pitch, int pc, const uint8_t* src, + bool ok, bool v16) { + const uint32_t dst = base + region + swz16(key, pc * 16, pitch); + if (v16) { + cp16(dst, src, ok); + } else { + cp8(dst, src, ok); + cp8(dst + 8, src + 8, ok); + } +} + +// the merge of columns [c0, c0 + NC) of (row r, head tile hb): attn_kernel's merge per element, loads first +template +__device__ void merge_cols(const Args& a, uint8_t* sm, int smt, int r, int hb, int c0) { + const int tid = threadIdx.x, H = a.H, R = a.R; + float* ca = reinterpret_cast(sm + smt); // [5][16] a, [5][16] b, [16] af, [16] den + float* cb = ca + NCH * BMQ; + float* fa = cb + NCH * BMQ; + float* fd = fa + BMQ; + bool cw[NCH]; +#pragma unroll + for (int c = 0; c < NCH; ++c) cw[c] = chunk_works(a, r, c); + constexpr int UB = BMQ * (NC / 4) / NT; // float4 units a thread (8 at P = 1, 2 at P = 4) + static_assert(UB * NT == BMQ * (NC / 4), "whole units"); + float4 pv[UB][NCH]; +#pragma unroll + for (int ub = 0; ub < UB; ++ub) { // the partials in flight first + const int u = tid + ub * NT, hl = u / (NC / 4), d = c0 + (u % (NC / 4)) * 4, h = hb * BMQ + hl; +#pragma unroll + for (int c = 0; c < NCH; ++c) + pv[ub][c] = cw[c] ? __ldcg(reinterpret_cast( + a.po + ((static_cast(c) * R + r) * H + h) * D + d)) + : make_float4(0.f, 0.f, 0.f, 0.f); + } + if (tid < BMQ) { // a head's scalar chain (Triton's _merge, unrolled) + const int h = hb * BMQ + tid; + float pmv[NCH], plv[NCH]; +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const long long base = (static_cast(c) * R + r) * H + h; + pmv[c] = cw[c] ? __ldcg(a.pm + base) : -INFINITY; + plv[c] = cw[c] ? __ldcg(a.pl + base) : 0.f; + } + const float sink = a.sink[h]; + float m = -INFINITY, l = 0.f; +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const float cm = pmv[c], cl = plv[c]; + const bool act = cl > 0.f; + float aa, bb; + if (c == 0) { + const float nm = act ? fmaxf(cm, -INFINITY) : -INFINITY; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + aa = 0.f; + l = __fmaf_rn(bb, cl, 0.f); + m = nm; + } else { + const float nm = act ? fmaxf(m, cm) : m; + aa = act ? (m == -INFINITY ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + l = __fmaf_rn(l, aa, __fmul_rn(bb, cl)); + m = nm; + } + ca[c * BMQ + tid] = aa; + cb[c * BMQ + tid] = bb; + } + const float top = fmaxf(m, sink); + const float af = m == -INFINITY ? 0.f : texp(__fsub_rn(m, top)); + const float e = sink == -INFINITY ? 0.f : texp(__fsub_rn(sink, top)); + fa[tid] = af; + fd[tid] = __fmaf_rn(l, af, e); + } + __syncthreads(); + const float* csr = a.cs + static_cast(row_pos(a, r)) * a.csst; +#pragma unroll + for (int ub = 0; ub < UB; ++ub) { + const int u = tid + ub * NT, hl = u / (NC / 4), d = c0 + (u % (NC / 4)) * 4, h = hb * BMQ + hl; + float o[4]; +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const float4 v = pv[ub][c]; + const float co[4] = {v.x, v.y, v.z, v.w}; + const float aa = ca[c * BMQ + hl], bb = cb[c * BMQ + hl]; +#pragma unroll + for (int i = 0; i < 4; ++i) + o[i] = c == 0 ? __fmaf_rn(bb, co[i], 0.f) : __fmaf_rn(o[i], aa, __fmul_rn(co[i], bb)); + } + float x[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) x[i] = div_full(__fmul_rn(o[i], fa[hl]), fd[hl]); + uint16_t ob[4]; +#pragma unroll + for (int i = 0; i < 4; i += 2) { // pair (d + i, d + i + 1): GPT-J, inverse + const int p = (d + i) >> 1; + const float c = p >= 224 ? csr[p - 224] : 1.f, s = p >= 224 ? csr[32 + p - 224] : 0.f; + const float ne = __fmaf_rn(x[i], c, __fmul_rn(x[i + 1], s)); + const float no = __fmaf_rn(x[i + 1], c, -__fmul_rn(x[i], s)); + const __nv_bfloat16 b0 = __float2bfloat16_rn(ne), b1 = __float2bfloat16_rn(no); + ob[i] = *reinterpret_cast(&b0); + ob[i + 1] = *reinterpret_cast(&b1); + } + uint2 w; + w.x = ob[0] | (static_cast(ob[1]) << 16); + w.y = ob[2] | (static_cast(ob[3]) << 16); + *reinterpret_cast(a.out + (static_cast(r) * H + h) * D + d) = w; + } +} + +template +__global__ void __launch_bounds__(NT, 1) split_kernel(Args a) { + using G = Split

; + extern __shared__ __align__(16) uint8_t sm[]; + const int k = P == 1 ? 0 : static_cast(cl_rank()); + const int c = blockIdx.x / P, hb = blockIdx.y, r = blockIdx.z; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, t = lane & 3; + const int H = a.H, R = a.R; + if (!chunk_works(a, r, c)) return; // the whole cluster: no barrier is pending + uint4 qv[32 * 32 / NT]; // q's A fragments: every load in flight at once + { + const uint16_t* qb = a.q + (static_cast(r) * H + hb * BMQ) * D; +#pragma unroll + for (int u = 0; u < 32 * 32 / NT; ++u) { + const int i = tid + u * NT, s = i >> 5, ln = i & 31, gg = ln >> 2, tt = ln & 3; + const int k0 = klo(s, tt), k1 = khi(s, tt); + qv[u].x = *reinterpret_cast(qb + gg * D + k0); + qv[u].y = *reinterpret_cast(qb + (gg + 8) * D + k0); + qv[u].z = *reinterpret_cast(qb + gg * D + k1); + qv[u].w = *reinterpret_cast(qb + (gg + 8) * D + k1); + } + if (a.qrope) { // the rope dims (k >= 448): their pair's rotation + const float* csr = a.cs + static_cast(row_pos(a, r)) * a.csst; +#pragma unroll + for (int u = 0; u < 32 * 32 / NT; ++u) { + const int i = tid + u * NT, s = i >> 5, tt = (i & 31) & 3; + const int k0 = klo(s, tt), k1 = khi(s, tt); + if (k0 >= NOPE) { + qv[u].x = rope_word(qv[u].x, csr, (k0 - NOPE) >> 1); + qv[u].y = rope_word(qv[u].y, csr, (k0 - NOPE) >> 1); + } + if (k1 >= NOPE) { + qv[u].z = rope_word(qv[u].z, csr, (k1 - NOPE) >> 1); + qv[u].w = rope_word(qv[u].w, csr, (k1 - NOPE) >> 1); + } + } + } + } + int* phys = reinterpret_cast(sm + G::SM_PHYS); + const long long sl = ROWS ? a.sl[r] : 0; + if (tid < CH) { // attn_kernel's row map + int row = -1; + if (c < NCOMP) { + const int idx = c * CH + tid; + if (idx < a.cnt[r]) { + const int tk = a.tok[static_cast(r) * a.ts + idx]; + if (tk >= 0) { + row = tk; + if (a.psh > 0) { + const int* ptr = a.pt + (ROWS ? sl * a.pts : 0); + row = (ptr[tk >> a.psh] << a.psh) + (tk & ((1 << a.psh) - 1)); + } + } + } + } else { + const int hi = HAS_HI ? a.hi[r] : row_pos(a, r); + const int lo = max(a.lo[r], hi - (WINDOW - 1)); + const int p = hi - (WINDOW - 1) + tid; + if (p >= lo && p >= 0) row = (p % a.ring) + (ROWS ? static_cast(sl) * a.ring : 0); + } + phys[tid] = row; + } + __syncthreads(); + { // the rows in flight before q's fragments + const bool comp = c < NCOMP; + const uint8_t* V = comp ? a.cv : a.sv; + const uint8_t* S = comp ? a.csc : a.ssc; + const long long vs = comp ? a.cvs : a.svs, ss = comp ? a.css : a.sss; + const bool v16 = ((reinterpret_cast(V) | static_cast(vs)) & 15) == 0; + const uint32_t base = static_cast(__cvta_generic_to_shared(sm)); + const int x0 = col2byte(G::W * k), nsl = (col2byte(G::W * k + G::W) - x0) / 16; // the slice's pieces a key + for (int i = tid; i < CH * nsl; i += NT) { + const int key = i / nsl, pc = i % nsl, row = phys[key]; + const long long rr = row >= 0 ? row : 0; + stage_piece(base, G::SM_KP, key, G::SLB, pc, V + rr * vs + x0 + pc * 16, row >= 0, v16); + } + if (P > 1) { // the S tiles' whole rows + for (int i = tid; i < G::KS * (VB / 16); i += NT) { + const int key = i / (VB / 16), pc = i % (VB / 16), row = phys[G::KS * k + key]; + const long long rr = row >= 0 ? row : 0; + stage_piece(base, G::SM_KS, key, VB, pc, V + rr * vs + pc * 16, row >= 0, v16); + } + } + for (int key = tid; key < CH; key += NT) { + const int row = phys[key]; + cp8(base + G::SM_SC + key * SB, S + (row >= 0 ? row : 0) * ss, row >= 0); + } + asm volatile("cp.async.commit_group;" ::: "memory"); + } +#pragma unroll + for (int u = 0; u < 32 * 32 / NT; ++u) reinterpret_cast(sm + G::SM_Q)[tid + u * NT] = qv[u]; + asm volatile("cp.async.wait_group 0;" ::: "memory"); + __syncthreads(); + if (P > 1) { // S's operands -> bf16 pairs, 4 k a step (deq2) + uint32_t* kb = reinterpret_cast(sm + G::SM_KB); + for (int i = tid; i < G::KS * (D / 4); i += NT) { + const int key = i / (D / 4), k4 = (i % (D / 4)) * 4, w = k4 >> 1; + uint32_t w0, w1; + if (k4 < NOPE) { + const uint32_t b = *reinterpret_cast(sm + G::SM_KS + swz16(key, k4, VB)); + const float sc = pow2b(sm[G::SM_SC + (key + G::KS * k) * SB + (k4 >> 6)]); + w0 = deq2(b & 0xffffu, sc, sc); + w1 = deq2(b >> 16, sc, sc); + } else { + const uint2 b = *reinterpret_cast(sm + G::SM_KS + swz16(key, NOPE + 2 * (k4 - NOPE), VB)); + w0 = b.x; + w1 = b.y; + } + kb[key * (D / 2) + (w ^ bsw(key))] = w0; + kb[key * (D / 2) + ((w + 1) ^ bsw(key))] = w1; + } + __syncthreads(); + } + const float NINF = -INFINITY; + float* tmax = reinterpret_cast(sm + G::SM_TMAX); + float* wsum = reinterpret_cast(sm + G::SM_WSUM); + float* rowv = reinterpret_cast(sm + G::SM_ROW); + // S: chain pair w (+ 8): local tile w >> 2, block n8 = w & 3 (P = 1: attn_kernel's tiles ta and ta + 2) + const int n8 = warp & 3; + const bool sw = warp < 4 * G::TPC; // this warp has S chains + float sv[G::NPW][4]; + bool val[G::NPW][2]; + int tile_of[G::NPW]; +#pragma unroll + for (int j = 0; j < G::NPW; ++j) { + tile_of[j] = G::TPC * k + (warp >> 2) + 2 * j; + for (int e = 0; e < 4; ++e) sv[j][e] = 0.f; + } + if (sw) { + const uint4* qf = reinterpret_cast(sm + G::SM_Q); + int lk[G::NPW]; +#pragma unroll + for (int j = 0; j < G::NPW; ++j) lk[j] = 32 * (tile_of[j] - G::TPC * k) + 8 * n8 + g; +#pragma unroll 4 + for (int s = 0; s < 32; ++s) { + const uint4 af = qf[s * 32 + lane]; + const int k0 = klo(s, t), k1 = khi(s, t); +#pragma unroll + for (int j = 0; j < G::NPW; ++j) { + if (P == 1) { + mma(sv[j], af, spair_k

(sm, lk[j], lk[j], k0), spair_k

(sm, lk[j], lk[j], k1)); + } else { + mma(sv[j], af, sbf_k(sm, lk[j], k0), sbf_k(sm, lk[j], k1)); + } + } + } +#pragma unroll + for (int j = 0; j < G::NPW; ++j) { + const int tile = tile_of[j], kc = 32 * tile + 8 * n8 + 2 * t; + val[j][0] = phys[kc] >= 0; + val[j][1] = phys[kc + 1] >= 0; + for (int e = 0; e < 4; ++e) + sv[j][e] = val[j][e & 1] ? __fmul_rn(sv[j][e], __int_as_float(0x3D3504F3)) : NINF; + float m0 = fmaxf(sv[j][0], sv[j][1]), m1 = fmaxf(sv[j][2], sv[j][3]); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 2)); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 1)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 2)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 1)); + if (t == 0) { + tmax[(tile * 4 + n8) * 16 + g] = m0; + tmax[(tile * 4 + n8) * 16 + g + 8] = m1; + } + } + } + if (P > 1) { // this rank's tile maxima pushed to the others + __syncthreads(); + for (int i = tid; i < 64 * G::TPC * (P - 1); i += NT) { + const int e = 64 * G::TPC * k + i % (64 * G::TPC), q = (k + 1 + i / (64 * G::TPC)) % P; + st_peer(tmax + e, q, __float_as_uint(tmax[e])); + } + cl_arrive(); + const int x0 = col2byte(G::W * k); // meanwhile: PV's B pairs, the fp8 slice -> bf16, + uint32_t* vb = reinterpret_cast(sm + G::SM_VB); // 4 columns a step (deq2 per column) + for (int i = tid; i < (CH / 2) * (G::W / 4); i += NT) { + const int kp = i / (G::W / 4), x = (i % (G::W / 4)) * 4, n = G::W * k + x, k0 = 2 * kp, k1 = k0 + 1; + uint32_t o[4]; + if (n < NOPE) { + const uint32_t b0 = *reinterpret_cast(sm + G::SM_KP + swz16(k0, n - x0, G::SLB)); + const uint32_t b1 = *reinterpret_cast(sm + G::SM_KP + swz16(k1, n - x0, G::SLB)); + const float s0 = pow2b(sm[G::SM_SC + k0 * SB + (n >> 6)]), s1 = pow2b(sm[G::SM_SC + k1 * SB + (n >> 6)]); +#pragma unroll + for (int e = 0; e < 4; ++e) o[e] = deq2(((b0 >> (8 * e)) & 0xffu) | (((b1 >> (8 * e)) & 0xffu) << 8), s0, s1); + } else { + const uint2 a0 = *reinterpret_cast(sm + G::SM_KP + swz16(k0, col2byte(n) - x0, G::SLB)); + const uint2 a1 = *reinterpret_cast(sm + G::SM_KP + swz16(k1, col2byte(n) - x0, G::SLB)); + o[0] = (a0.x & 0xffffu) | (a1.x << 16); + o[1] = (a0.x >> 16) | (a1.x & 0xffff0000u); + o[2] = (a0.y & 0xffffu) | (a1.y << 16); + o[3] = (a0.y >> 16) | (a1.y & 0xffff0000u); + } +#pragma unroll + for (int e = 0; e < 4; ++e) vb[kp * G::W + ((x + e) ^ vsw(kp))] = o[e]; + } + cl_wait(); + } + __syncthreads(); + if (tid < BMQ) { // the running max over the tiles, a row a thread + float m = NINF; + for (int tile = 0; tile < 4; ++tile) { + const float* tm4 = tmax + tile * 64 + tid; + const float tm = fmaxf(fmaxf(tm4[0], tm4[16]), fmaxf(tm4[32], tm4[48])); + const bool act = tm != NINF; + const float nm = act ? fmaxf(m, tm) : m; + const float al = act ? (m == NINF ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + rowv[tile * 16 + tid] = nm; + rowv[64 + tile * 16 + tid] = al; + rowv[128 + tile * 16 + tid] = act ? 1.f : 0.f; + m = nm; + } + reinterpret_cast(sm + G::SM_MISC)[tid] = m; + } + __syncthreads(); + uint32_t* pf = reinterpret_cast(sm + G::SM_P); // [tile][step][lane][4] + if (sw) { +#pragma unroll + for (int j = 0; j < G::NPW; ++j) { // p, the quad sums, bf16(p) -> fragments + const int tile = tile_of[j]; + float p[4]; + for (int e = 0; e < 4; ++e) { + const int row = g + (e >> 1) * 8; + const bool act = rowv[128 + tile * 16 + row] != 0.f; + p[e] = (val[j][e & 1] && act) ? texp(__fsub_rn(sv[j][e], rowv[tile * 16 + row])) : 0.f; + } + float a0 = __fadd_rn(p[0], p[1]), a1 = __fadd_rn(p[2], p[3]); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 2)); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 1)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 2)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 1)); + if (t == 0) { + wsum[(tile * 4 + n8) * 16 + g] = a0; + wsum[(tile * 4 + n8) * 16 + g + 8] = a1; + } + int step, tq, hi; + frag_of(8 * n8 + 2 * t, step, tq, hi); + const __nv_bfloat162 lo2 = __floats2bfloat162_rn(p[0], p[1]), hi2 = __floats2bfloat162_rn(p[2], p[3]); + uint32_t* f = pf + ((tile * 2 + step) * 32 + g * 4 + tq) * 4; + f[2 * hi] = *reinterpret_cast(&lo2); + f[2 * hi + 1] = *reinterpret_cast(&hi2); + } + } + if (P > 1) { // this rank's row sums and fragments pushed + __syncthreads(); + constexpr int NW = G::TPC * (64 + 256); // words of this rank's tiles + for (int i = tid; i < NW * (P - 1); i += NT) { + const int e = i % NW, q = (k + 1 + i / NW) % P; + uint32_t* w = e < 64 * G::TPC ? reinterpret_cast(wsum) + 64 * G::TPC * k + e + : pf + 256 * G::TPC * k + (e - 64 * G::TPC); + st_peer(w, q, *w); + } + cl_arrive(); // the last cluster access: no barrier before exit + cl_wait(); + } + __syncthreads(); + if (tid < BMQ) { // l: fma(l, alpha, (W0 + W2) + (W1 + W3)) + float l = 0.f; + for (int tile = 0; tile < 4; ++tile) { + const float* w4 = wsum + tile * 64 + tid; + const float sum = __fadd_rn(__fadd_rn(w4[0], w4[32]), __fadd_rn(w4[16], w4[48])); + l = __fmaf_rn(l, rowv[64 + tile * 16 + tid], sum); + } + reinterpret_cast(sm + G::SM_MISC)[16 + tid] = l; + } + const int x0 = col2byte(G::W * k); + float acc[G::J][4]; +#pragma unroll + for (int j = 0; j < G::J; ++j) + for (int e = 0; e < 4; ++e) acc[j][e] = 0.f; + for (int tile = 0; tile < 4; ++tile) { // o = mma(P, kv, rn(o * alpha)), 2 k steps a tile + const float al0 = rowv[64 + tile * 16 + g], al1 = rowv[64 + tile * 16 + g + 8]; +#pragma unroll + for (int j = 0; j < G::J; ++j) { + acc[j][0] = __fmul_rn(acc[j][0], al0); + acc[j][1] = __fmul_rn(acc[j][1], al0); + acc[j][2] = __fmul_rn(acc[j][2], al1); + acc[j][3] = __fmul_rn(acc[j][3], al1); + } + for (int step = 0; step < 2; ++step) { + const uint4 af = reinterpret_cast(pf)[(tile * 2 + step) * 32 + lane]; + const int k0 = 32 * tile + klo(step, t), k1 = 32 * tile + khi(step, t); +#pragma unroll + for (int j = 0; j < G::J; ++j) { + const int n = 64 * (G::J * k + j) + 8 * warp + g; + if (P == 1) { + mma(acc[j], af, spair_n

(sm, k0, k0 + 1, n, x0), spair_n

(sm, k1, k1 + 1, n, x0)); + } else { + const uint32_t* vb = reinterpret_cast(sm + G::SM_VB); + const int x = n - G::W * k; + mma(acc[j], af, vb[(k0 >> 1) * G::W + (x ^ vsw(k0 >> 1))], + vb[(k1 >> 1) * G::W + (x ^ vsw(k1 >> 1))]); + } + } + } + } + const long long pbase = (static_cast(c) * R + r) * H + hb * BMQ; +#pragma unroll + for (int j = 0; j < G::J; ++j) { + const int n = 64 * (G::J * k + j) + 8 * warp + 2 * t; + *reinterpret_cast(a.po + (pbase + g) * D + n) = make_float2(acc[j][0], acc[j][1]); + *reinterpret_cast(a.po + (pbase + g + 8) * D + n) = make_float2(acc[j][2], acc[j][3]); + } + if (tid < BMQ) { // every rank: the same (m, l), so each merge sees it + a.pm[pbase + tid] = reinterpret_cast(sm + G::SM_MISC)[tid]; + a.pl[pbase + tid] = reinterpret_cast(sm + G::SM_MISC)[16 + tid]; + } + int* flag = reinterpret_cast(sm + G::SM_MISC) + 32; + int* tk = a.ticket + (r * (H / BMQ) + hb) * MAXP + k; + __threadfence(); + __syncthreads(); + if (tid == 0) { + int need = 0; + for (int cc = 0; cc < NCH; ++cc) need += chunk_works(a, r, cc) ? 1 : 0; + *flag = atomicAdd(tk, 1) == need - 1; + } + __syncthreads(); + if (!*flag) return; + __threadfence(); + if (tid == 0) atomicExch(tk, 0); // self-cleaning: the next launch / replay + merge_cols(a, sm, G::SM_TMAX, r, hb, G::W * k); +} diff --git a/zig/kernels/cuda/deepseek_v41/decode.cuh b/zig/kernels/cuda/deepseek_v41/decode.cuh new file mode 100644 index 0000000000..2a663ae1c9 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/decode.cuh @@ -0,0 +1,166 @@ +// Copy of src/tensorfold/cuda/exl3/decode.cuh (blob 1acca8f0457c): sass_check.sh + +// After ExLlamaV3, MIT, Copyright (c) 2025 Turboderp. + +// EXL3 tiles: lane L decodes a tile's values 8L..8L+7 straight into its two mma.m16n8k16 B fragments, bit for bit. + +#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 (stream_ends). +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; +} + +// This lane's words of a tile (tile: the tile's first 32-bit word; global or shared memory). +template +__device__ __forceinline__ void load_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] = tile[(word + i) % tile_words()]; +} + +// 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); + } +} + +// One codebook value. +template +__device__ __forceinline__ half decode1(uint32_t s) { + const uint32_t v = decode2(s, s); + return __ushort_as_half(static_cast(v & 0xffffu)); +} + +// 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]); +} + +// Row and column in the 16x16 tile (row = k, column = n) of this lane's value j (0..7). +__device__ __forceinline__ int value_row(int lane, int j) { return 2 * (lane & 3) + (j & 1) + 8 * ((j >> 1) & 1); } +__device__ __forceinline__ int value_col(int lane, int j) { return (lane >> 2) + 8 * (j >> 2); } + +// 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/deepseek_v41/exl3_experts.cu b/zig/kernels/cuda/deepseek_v41/exl3_experts.cu new file mode 100644 index 0000000000..a0959e216b --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/exl3_experts.cu @@ -0,0 +1,283 @@ +// Copy of src/tensorfold/cuda/exl3/experts.cu (blob 8ce66ca83bce): sass_check.sh + +// EXL3 routed experts, a width per expert: fixed-order splits, slots and butterflies, no atomics, so bit-exact. + +#include +#include +#include // what the cut torch includes provided: the fixed-width integers + +#include "experts_grouped.cuh" + +namespace tf_exl3_experts { + +constexpr float HAD_SCALE = 0.08838834764831845f; // 1 / sqrt(128) + +// 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; + +__global__ void __launch_bounds__(GROUP_THREADS) group_kernel(const int* __restrict__ pick, int* __restrict__ uids, + int* __restrict__ ucount, int* __restrict__ members, + int R, int slots, int E, int maxm) { + extern __shared__ int sh_pick[]; + __shared__ int warp_tot[GROUP_THREADS / 32]; + const int n = R * slots; + for (int i = threadIdx.x; i < n; i += GROUP_THREADS) sh_pick[i] = pick[i]; + __syncthreads(); + int cnt[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 c = 0; + if (e < E) + for (int i = 0; i < n; ++i) c += sh_pick[i] == e; + cnt[q] = c; + used += c > 0; + } + // exclusive scan of `used` over threads + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + 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; +#pragma unroll + 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) ucount[0] = s; + } + __syncthreads(); + int place = warp_tot[warp] + inc - used; +#pragma unroll + for (int q = 0; q < GROUP_PER_THREAD; ++q) { + if (cnt[q] == 0) continue; + const int e = threadIdx.x * GROUP_PER_THREAD + q; + uids[place] = e; + int j = 0; + for (int i = 0; i < n && j < maxm; ++i) + if (sh_pick[i] == e) members[place * maxm + j++] = (i / slots) * 32 + (i % slots); + for (; j < maxm; ++j) members[place * maxm + j] = -1; + ++place; + } +} + +// Walsh-Hadamard of 128 values, 4 a lane, fixed butterfly order (strides 1, 2 in registers, 4..64 across lanes). +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + 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) { + float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +template __device__ __forceinline__ float to_f(T v); +template <> __device__ __forceinline__ float to_f<__nv_bfloat16>(__nv_bfloat16 v) { return __bfloat162float(v); } +template <> __device__ __forceinline__ float to_f(half v) { return __half2float(v); } + +// Program (member row, 128-block of K, matrix): Xh = fp16((x * suh) @ H) for gate and up of every routed slot. +template +__global__ void rot_in_kernel(const TIN* __restrict__ x, int x_stride, const int* __restrict__ pick, + const half* __restrict__ suh0, const half* __restrict__ suh1, half* __restrict__ out0, + half* __restrict__ out1, int K, int slots, int E) { + const int p = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int row = p / slots; + const int e = pick[p]; + if (e < 0 || e >= E) return; + const int lane = threadIdx.x; + const half* suh = (mat ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + const TIN* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = to_f(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? out1 : out0) + (size_t)p * 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 (row, 128-block of width): ordered splits, rotate, * svh, SwiGLU (0: bf16 roundings, 1: fp32), act -> Xd. +__global__ void gateup_epilogue_kernel(const float* __restrict__ Z, const int* __restrict__ pick, + const half* __restrict__ svh_g, const half* __restrict__ svh_u, + const half* __restrict__ suh_d, half* __restrict__ xd, int P, int N, int SK, + int E, float limit, int act_mode) { + const int p = blockIdx.x, blk = blockIdx.y; + const int e = pick[p]; + if (e < 0 || e >= E) 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 < SK; ++s) { + sg += Z[((size_t)(0 * SK + s) * P + p) * N + n + j]; + su += Z[((size_t)(1 * SK + s) * P + p) * 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 (act_mode == 0) { + float gg = fminf(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = fminf(fmaxf(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), -limit), + limit); + act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + } else { + float gg = fminf(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j]), limit); + float uu = fminf(fmaxf(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j]), -limit), limit); + act = gg / (1.f + expf(-gg)) * uu; + } + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * 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 model width): Y = (splits summed in order) @ H * svh_d, fp32. +__global__ void down_epilogue_kernel(const float* __restrict__ Z, const int* __restrict__ pick, + const half* __restrict__ svh_d, float* __restrict__ y, int P, int D, int SK, + int E) { + const int p = blockIdx.x, blk = blockIdx.y; + const int e = pick[p]; + if (e < 0 || e >= E) return; + const int lane = threadIdx.x; + const int n = blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int k = 0; k < SK; ++k) s += Z[((size_t)k * P + p) * D + n + j]; + v[j] = s; + } + fwht128(v, lane); + float* o = y + (size_t)p * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); +} + +// out[r][d] = sum over slots in order of wts[r][k] * y[r * slots + k][d] (fp32, fma chain from 0). +__global__ void combine_kernel(const float* __restrict__ y, const float* __restrict__ wts, float* __restrict__ out, + int D, int slots) { + const int r = blockIdx.x; + const int d = blockIdx.y * blockDim.x + threadIdx.x; + if (d >= D) return; + float acc = 0.f; + for (int k = 0; k < slots; ++k) acc = fmaf(wts[r * slots + k], y[((size_t)r * slots + k) * D + d], acc); + out[(size_t)r * D + d] = acc; +} + +// down_epilogue_kernel then combine_kernel in one launch, the same arithmetic in the same order (the same bits). +__device__ __forceinline__ void store_out(float* o, float v) { *o = v; } +__device__ __forceinline__ void store_out(__nv_bfloat16* o, float v) { *o = __float2bfloat16_rn(v); } + +// OutT bf16: the fp32 sum rounded to nearest even as it is stored (same bits as fp32 out then to bfloat16) +template +__global__ void down_combine_kernel(const float* __restrict__ Z, const int* __restrict__ pick, + const half* __restrict__ svh_d, float* __restrict__ y, + const float* __restrict__ wts, OutT* __restrict__ out, int P, int D, int SK, + int E, int slots) { + __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 * slots + k; + const int e = pick[p]; + float o[4]; + if (e >= 0 && e < E) { + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int q = 0; q < SK; ++q) s += Z[((size_t)q * P + p) * 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(svh_d[(size_t)e * D + n + j]); + y[(size_t)p * D + n + j] = o[j]; + } + } else { +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = y[(size_t)p * 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 < slots; ++q) { + const float w = wts[r * 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) store_out(out + (size_t)r * D + n + j, acc[j]); +} + +} // namespace tf_exl3_experts + +// upstream grouped_kernel at mul1 for every (nt, warps, pf) x TF_RANGES, and dequant_kernel<2, K2> (tests) +#define TF_EXL3X_INST(NT, W, PF, LO, HI) template __global__ void tf_exl3x::grouped_kernel<2, NT, W, PF, LO, HI>( \ + const half*, const half*, const int64_t*, const int64_t*, const int*, const int*, const int*, const int*, \ + const int*, float*, int, int, int, int, int, int); +TF_EXL3X_INST(8, 4, 1, 8, 8) +TF_EXL3X_INST(8, 4, 1, 2, 10) +TF_EXL3X_INST(8, 4, 1, 2, 16) +TF_EXL3X_INST(8, 4, 2, 8, 8) +TF_EXL3X_INST(8, 4, 2, 2, 10) +TF_EXL3X_INST(8, 4, 2, 2, 16) +TF_EXL3X_INST(4, 4, 2, 8, 8) +TF_EXL3X_INST(4, 4, 2, 2, 10) +TF_EXL3X_INST(4, 4, 2, 2, 16) +template __global__ void tf_exl3x::dequant_kernel<2, 2>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 3>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 4>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 5>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 6>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 7>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 8>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 9>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 10>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 11>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 12>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 13>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 14>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 15>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 16>(const uint32_t*, half*, int, int); +// down_combine_kernel at both outputs the binding takes (fp32, bf16: TF_DSV41_MOE_BF16) +template __global__ void tf_exl3_experts::down_combine_kernel(const float*, const int*, const half*, float*, \ + const float*, float*, int, int, int, int, int); +template __global__ void tf_exl3_experts::down_combine_kernel<__nv_bfloat16>(const float*, const int*, const half*, float*, \ + const float*, __nv_bfloat16*, int, int, int, int, int); +// rot_in_kernel for both inputs the binding takes +template __global__ void tf_exl3_experts::rot_in_kernel<__nv_bfloat16>(const __nv_bfloat16*, int, const int*, const half*, const half*, half*, \ + half*, int, int, int); +template __global__ void tf_exl3_experts::rot_in_kernel(const half*, int, const int*, const half*, const half*, half*, \ + half*, int, int, int); diff --git a/zig/kernels/cuda/deepseek_v41/experts_grouped.cuh b/zig/kernels/cuda/deepseek_v41/experts_grouped.cuh new file mode 100644 index 0000000000..ab7591ac2e --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/experts_grouped.cuh @@ -0,0 +1,311 @@ +// Copy of src/tensorfold/cuda/exl3/experts_grouped.cuh (blob 57be4c027785): sass_check.sh + +// After ExLlamaV3, MIT, Copyright (c) 2025 Turboderp. + +// Grouped EXL3 expert GEMV: rows stay independent, K ranges fixed by shape, warps summed in order. + +#pragma once + +#include +#include +#include // what the cut torch includes provided: the fixed-width integers + +namespace tf_exl3x { + +// Two codebook values (CB 0 3inst, 1 mcg, 2 mul1) as a half2, bit-identical to ExLlamaV3's decode_3inst_2. +template +__device__ __forceinline__ uint32_t cb_pair(uint32_t s0, uint32_t s1) { + if constexpr (CB == 2) { + const uint32_t x0 = s0 * 0x83DCD12Du, x1 = s1 * 0x83DCD12Du; + const uint32_t sum0 = __dp4a(x0, 0x01010101u, 0x6400u); + const uint32_t sum1 = __dp4a(x1, 0x01010101u, 0x6400u); + const uint32_t hv = __byte_perm(sum0, sum1, 0x5410); + half2 h = *reinterpret_cast(&hv); + half2 r = __hfma2(h, __half2half2(__ushort_as_half(0x1eee)), __half2half2(__ushort_as_half(0xc931))); + return *reinterpret_cast(&r); + } else { + uint32_t x0, x1; + if constexpr (CB == 1) { + x0 = s0 * 0xCBAC1FEDu; + x1 = s1 * 0xCBAC1FEDu; + } else { + x0 = s0 * 89226354u + 64248484u; + x1 = s1 * 89226354u + 64248484u; + } + x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u; + x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u; + uint32_t lo = __byte_perm(x0, x1, 0x5410); + uint32_t hi = __byte_perm(x0, x1, 0x7632); + half2 r = __hadd2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + return *reinterpret_cast(&r); + } +} + +// K2 half-bits a value: a tile is 4 * K2 words; a lane's eight windows fall in NG runs of GV within two words. +template +struct Fmt { + static constexpr int TW = 4 * K2; + static constexpr int LW = (TW + 31) / 32; + // windows sharing one 64-bit merge (tested for every K2) + static constexpr int GV = (K2 >= 13) ? 2 : ((K2 == 7 || (K2 >= 9 && K2 <= 12) || K2 == 16) ? 4 : 8); + static constexpr int NG = 8 / GV; + __host__ __device__ static constexpr int end(int p) { return (p >> 1) * K2 + ((p & 1) ? K2 : (K2 >> 1)); } + // right shift of window j of a run (run starts at an even position) relative to the run's last window + __host__ __device__ static constexpr int off(int j) { return end(GV - 1) - end(j); } +}; + +template +struct LaneMap { + int hi[Fmt::NG], lo[Fmt::NG], sh[Fmt::NG]; + __device__ __forceinline__ explicit LaneMap(int lane) { + constexpr int TW = Fmt::TW, GV = Fmt::GV; +#pragma unroll + for (int g = 0; g < Fmt::NG; ++g) { + const int last_end = Fmt::end(8 * lane + g * GV + GV - 1) + 128 * K2; + const int hr = (last_end - 1) >> 5; + hi[g] = hr % TW; + lo[g] = (hr + TW - 1) % TW; + sh[g] = (hr + 1) * 32 - last_end; + } + } +}; + +template +__device__ __forceinline__ uint32_t fetch(const uint32_t (&w)[LW], int idx) { + if constexpr (LW == 1) { + return __shfl_sync(0xffffffffu, w[0], idx); + } else { + const uint32_t a = __shfl_sync(0xffffffffu, w[0], idx & 31); + const uint32_t b = __shfl_sync(0xffffffffu, w[1], idx & 31); + return idx < 32 ? a : b; + } +} + +// This lane's eight values of a tile as the B fragments of its two n8 halves. +template +__device__ __forceinline__ void decode_tile(const uint32_t (&w)[Fmt::LW], const LaneMap& m, int lane, + uint32_t (&b0)[2], uint32_t (&b1)[2]) { + uint32_t st[8]; + if constexpr (K2 == 8) { + // 4 bits: lane L's windows are exactly words L-1 and L (as the 4-bit EXL3 decode) + const uint32_t p = __shfl_sync(0xffffffffu, w[0], (lane + 31) & 31); + const uint32_t s = __funnelshift_r(w[0], p, 20); + st[0] = (s >> 8) & 0xffffu; + st[1] = (s >> 4) & 0xffffu; + st[2] = s & 0xffffu; + st[3] = w[0] >> 16; + st[4] = (w[0] >> 12) & 0xffffu; + st[5] = (w[0] >> 8) & 0xffffu; + st[6] = (w[0] >> 4) & 0xffffu; + st[7] = w[0] & 0xffffu; + } else { + constexpr int GV = Fmt::GV, NG = Fmt::NG; +#pragma unroll + for (int g = 0; g < NG; ++g) { + const uint32_t whi = fetch::LW>(w, m.hi[g]); + const uint32_t wlo = fetch::LW>(w, m.lo[g]); + const uint64_t mm = ((((uint64_t)wlo) << 32) | whi) >> m.sh[g]; +#pragma unroll + for (int j = 0; j < GV; ++j) st[g * GV + j] = (uint32_t)(mm >> Fmt::off(j)) & 0xffffu; + } + } + b0[0] = cb_pair(st[0], st[1]); + b0[1] = cb_pair(st[2], st[3]); + b1[0] = cb_pair(st[4], st[5]); + b1[1] = cb_pair(st[6], st[7]); +} + +__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])); +} + +__device__ __forceinline__ uint32_t load_pair(const half* x, bool ok) { + return ok ? *reinterpret_cast(x) : 0u; +} + +template +__device__ __forceinline__ void load_words(uint32_t (&dst)[Fmt::LW], const uint32_t* p, int lane) { + constexpr int TW = Fmt::TW; +#pragma unroll + for (int l = 0; l < Fmt::LW; ++l) { + if constexpr ((TW % 32) == 0) + dst[l] = __ldg(p + l * 32); + else + dst[l] = (l * 32 + lane < TW) ? __ldg(p + l * 32) : 0u; + } +} + +// One warp's k tiles [kt0, kt0 + nkt) of an expert matrix into acc, PF tiles in flight. +template +__device__ __forceinline__ void warp_tiles(const uint32_t* __restrict__ 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]) { + constexpr int TW = Fmt::TW, LW = Fmt::LW; + const LaneMap map(lane); + const size_t kstride = (size_t)NTILES * TW; + const uint32_t* tp = T + ((size_t)kt0 * NTILES + nt0) * TW + lane; + + uint32_t pf[PF][NT][LW]; +#pragma unroll + for (int d = 0; d < PF; ++d) + if (d < nkt) +#pragma unroll + for (int i = 0; i < NT; ++i) load_words(pf[d][i], tp + d * kstride + i * TW, lane); + + for (int ib = 0; ib < nkt; ib += PF) { +#pragma unroll + for (int d = 0; d < PF; ++d) { + const int it = ib + d; + if (it < nkt) { + uint32_t w[NT][LW]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int l = 0; l < LW; ++l) w[i][l] = pf[d][i][l]; + if (it + PF < nkt) +#pragma unroll + for (int i = 0; i < NT; ++i) + load_words(pf[d][i], tp + (size_t)(it + PF) * kstride + i * TW, lane); + const int k = (kt0 + it) * 16; + uint32_t a[4] = {load_pair(x0 + k, ok0), load_pair(x1 + k, ok1), load_pair(x0 + k + 8, ok0), + load_pair(x1 + k + 8, ok1)}; +#pragma unroll + for (int i = 0; i < NT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile(w[i], map, lane, b0, b1); + mma16816(acc[i][0], a, b0); + mma16816(acc[i][1], a, b1); + } + } + } + } +} + +// The K2 values an instance covering [LO, HI] compiles (half-bits 2..16). +__host__ __device__ constexpr bool k2_supported(int k2) { + return k2 >= 2 && k2 <= 16; +} + +// Program (expert u, n block, split, member tile): up to 16 members x W_q over the split's K range; warps in order. +template +__global__ void __launch_bounds__(W * 32) grouped_kernel( + const half* __restrict__ X0, const half* __restrict__ X1, const int64_t* __restrict__ TP0, + const int64_t* __restrict__ TP1, const int* __restrict__ K2_0, const int* __restrict__ K2_1, + const int* __restrict__ uids, const int* __restrict__ ucount, const int* __restrict__ members, + float* __restrict__ Z, int K, int N, int P, int SK, int maxm, int slots) { + const int u = blockIdx.x; + if (u >= ucount[0]) return; + const int MT = (maxm + 15) / 16; + const int mtile = blockIdx.z % MT; + const int split = (blockIdx.z / MT) % SK; + const int mat = blockIdx.z / MT / SK; + const half* X = mat ? X1 : X0; + const int e = uids[u]; + const uint32_t* T = reinterpret_cast(mat ? TP1[e] : TP0[e]); + const int k2 = mat ? K2_1[e] : K2_0[e]; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int KT = K >> 4, NTILES = N >> 4; + + __shared__ int rows_sh[16]; + if (threadIdx.x < 16) { + const int m = mtile * 16 + threadIdx.x; + const int code = m < maxm ? members[u * maxm + m] : -1; + rows_sh[threadIdx.x] = code >= 0 ? (code >> 5) * slots + (code & 31) : -1; + } + __syncthreads(); + if (rows_sh[0] < 0) return; // members come first, so this tile is empty + const int r0 = rows_sh[g], r1 = rows_sh[g + 8]; + const half* x0 = X + (size_t)(r0 < 0 ? 0 : r0) * K + 2 * t; + const half* x1 = X + (size_t)(r1 < 0 ? 0 : r1) * K + 2 * t; + + const int per_split = KT / SK, per_warp = per_split / W; + const int kt0 = split * per_split + warp * per_warp; + const int nt0 = blockIdx.y * NT; + + 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_: \ + if constexpr (K2_ >= LO && K2_ <= HI) \ + warp_tiles(T, NTILES, kt0, per_warp, nt0, x0, x1, r0 >= 0, r1 >= 0, lane, acc); \ + else \ + __trap(); \ + 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(); + } + + // warps' partial sums through shared memory, added in warp order + __shared__ float red[W][16][NT * 16]; +#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]; + Z[(((size_t)mat * SK + split) * P + r) * N + nt0 * 16 + col] = s; + } +} + +// W_q [K, N] fp16 of one matrix through the same lane decode (tests; not on the forward path). +template +__global__ void dequant_kernel(const uint32_t* __restrict__ T, half* __restrict__ out, int K, int N) { + const int kt = blockIdx.x, nt = blockIdx.y, lane = threadIdx.x; + const int NTILES = N >> 4; + uint32_t w[Fmt::LW]; + load_words(w, T + ((size_t)kt * NTILES + nt) * Fmt::TW + lane, lane); + const LaneMap map(lane); + uint32_t b0[2], b1[2]; + decode_tile(w, map, lane, b0, b1); + const int g = lane >> 2, t = lane & 3; + uint32_t v[4] = {b0[0], b0[1], b1[0], b1[1]}; +#pragma unroll + for (int q = 0; q < 4; ++q) { + half2 h = *reinterpret_cast(&v[q]); + const int col = nt * 16 + g + 8 * (q >> 1); + const int row = kt * 16 + 2 * t + 8 * (q & 1); + out[(size_t)row * N + col] = __low2half(h); + out[(size_t)(row + 1) * N + col] = __high2half(h); + } +} + +} // namespace tf_exl3x diff --git a/zig/kernels/cuda/deepseek_v41/glue.cu b/zig/kernels/cuda/deepseek_v41/glue.cu new file mode 100644 index 0000000000..ad589080e5 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/glue.cu @@ -0,0 +1,403 @@ +// DeepSeek-V4.1 window glue: torch's CUDA arithmetic per element (fp32 opmath, bf16 / f16 rn casts, --fmad=false). + +#include +#include +#include + +namespace dsv41_glue { + +__device__ __forceinline__ float bf(uint16_t b) { return __uint_as_float(static_cast(b) << 16); } +__device__ __forceinline__ uint16_t to_bf(float x) { return __bfloat16_as_ushort(__float2bfloat16(x)); } + +// rows [n, D] = bf16_rn(f32 embed[ids - lo]) from the rank's slice [V, D], +0 off it (NaN -> 0x7FFF) +__global__ void __launch_bounds__(256) embed_send_kernel(const long long* __restrict__ ids, int n, + const uint16_t* __restrict__ embed, long long vocab_lo, int V, + int D, uint16_t* __restrict__ rows) { + const long long total = (long long)n * D; + for (long long e = (long long)blockIdx.x * blockDim.x + threadIdx.x; e < total; e += (long long)gridDim.x * blockDim.x) { + const long long i = e / D, d = e - i * D; + const long long local = ids[i] - vocab_lo; + rows[e] = (local >= 0 && local < V) ? to_bf(bf(embed[local * D + d])) : (uint16_t)0; + } +} + +// recv [W, n, D] bf16 summed in rank order (fp32 adds) -> out [n, 4 D] bf16, the sum repeated x4 +__global__ void __launch_bounds__(256) embed_sum_kernel(const uint16_t* __restrict__ recv, int W, int n, int D, + uint16_t* __restrict__ out) { + const long long total = (long long)n * D, plane = total; + for (long long e = (long long)blockIdx.x * blockDim.x + threadIdx.x; e < total; e += (long long)gridDim.x * blockDim.x) { + float s = bf(recv[e]); + for (int w = 1; w < W; ++w) s = __fadd_rn(s, bf(recv[w * plane + e])); + const uint16_t b = to_bf(s); + const long long i = e / D, d = e - i * D; + uint16_t* o = out + i * 4 * D + d; +#pragma unroll + for (int r = 0; r < 4; ++r) o[r * D] = b; + } +} + +// y = bf16_rn(x) (the MoE partial) +__global__ void __launch_bounds__(256) cast_bf16_kernel(const float* __restrict__ x, uint16_t* __restrict__ y, + long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + y[i] = to_bf(x[i]); +} + +// y = f32(f16_rn(x)) +__global__ void __launch_bounds__(256) kit_weights_kernel(const float* __restrict__ x, float* __restrict__ y, + long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + y[i] = __half2float(__float2half(x[i])); +} + +// x = f32(bf16_rn(x)) in place +__global__ void __launch_bounds__(256) kit_logits_kernel(float* __restrict__ x, long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + x[i] = bf(to_bf(x[i])); +} + +// pos[0] = start; pos64[i] = start + i; lo[i] = 0 +__global__ void __launch_bounds__(256) positions_kernel(int start, int n, int* __restrict__ pos, + long long* __restrict__ pos64, int* __restrict__ lo) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i == 0) pos[0] = start; + if (i < n) { + pos64[i] = (long long)start + i; + lo[i] = 0; + } +} + +// a graphed window's positions: `start` from the window's staged statics, so one captured graph serves every window +__global__ void __launch_bounds__(256) positions_dev_kernel(const int* __restrict__ start, int n, int* __restrict__ pos, + long long* __restrict__ pos64, int* __restrict__ lo) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int s = *start; + if (i == 0) pos[0] = s; + if (i < n) { + pos64[i] = (long long)s + i; + lo[i] = 0; + } +} + +// recv [W, n, c] -> out [n, W c] (bf16 words, any 2-byte type) +__global__ void __launch_bounds__(256) gather_cols_kernel(const uint16_t* __restrict__ recv, int W, int n, int c, + uint16_t* __restrict__ out) { + const long long total = (long long)W * n * c; + for (long long e = (long long)blockIdx.x * blockDim.x + threadIdx.x; e < total; e += (long long)gridDim.x * blockDim.x) { + const long long i = e / ((long long)W * c), r = e - i * W * c, w = r / c, j = r - w * c; + out[e] = recv[(w * n + i) * c + j]; + } +} + +// compressor carry [cols] fp32 <- bf16 row `row` of src (row stride ld); the row from the device when `accepted` is set +__global__ void __launch_bounds__(256) carry_kernel(const uint16_t* __restrict__ src, long long ld, + const int* __restrict__ accepted, int row, int cols, + float* __restrict__ carry) { + const int r = accepted ? *accepted : row; + for (int j = blockIdx.x * blockDim.x + threadIdx.x; j < cols; j += gridDim.x * blockDim.x) + carry[j] = bf(src[(long long)r * ld + j]); +} + +// bf16 -> fp32 (exact: the bits shifted into the high half), n elements +__global__ void __launch_bounds__(256) widen_f32_kernel(const uint16_t* __restrict__ x, float* __restrict__ y, + long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + y[i] = bf(x[i]); +} + +// DSpark candidates: a CTA per fp32 logit row [rows, V] (stride ld); the best k by (value desc, column asc), exact. +constexpr int DS_NT = 1024, DS_KMAX = 128; + +__device__ __forceinline__ unsigned long long ds_key(float x, int col) { // unique in a row; -0.0 == +0.0 + const unsigned int b = __float_as_uint(x + 0.0f); // -0.0 + 0.0 = +0.0 + const unsigned int mono = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + return (static_cast(mono) << 32) | (0xFFFFFFFFull - static_cast(col)); +} + +__global__ void __launch_bounds__(DS_NT) ds_cands_kernel(const float* __restrict__ lg, long long ld, int V, int k, + int id0, float* __restrict__ out) { + __shared__ int hist[256]; + __shared__ unsigned long long s_prefix, s_mask; + __shared__ int s_need, s_done, s_n; + __shared__ unsigned long long key_sh[DS_KMAX]; + const int tid = threadIdx.x; + const float* row = lg + (long long)blockIdx.x * ld; + float* o = out + (long long)blockIdx.x * 2 * k; // k values, then id0 + col as int bits; padding -inf, -1 + const int c = V < k ? V : k; + if (tid == 0) { s_prefix = 0; s_mask = 0; s_need = c; s_done = 0; } + __syncthreads(); + for (int shift = 56; shift >= 0 && c > 0; shift -= 8) { // radix select of the k-th key, 8 bits a pass + if (tid < 256) hist[tid] = 0; + __syncthreads(); + const unsigned long long prefix = s_prefix, mask = s_mask; + for (int i = tid; i < V; i += DS_NT) { + const unsigned long long u = ds_key(row[i], i); + if ((u & mask) == prefix) atomicAdd(&hist[(u >> shift) & 255], 1); + } + __syncthreads(); + if (tid == 0) { + int need = s_need, d = 255; + for (; d > 0 && hist[d] < need; --d) need -= hist[d]; + s_need = need; + s_prefix = prefix | ((unsigned long long)d << shift); + s_mask = mask | (255ull << shift); + s_done = hist[d] == need; // the digit's keys: all winners + } + __syncthreads(); + if (s_done) break; + } + const unsigned long long prefix = s_prefix, mask = s_mask; + if (tid == 0) s_n = 0; + __syncthreads(); + for (int i = tid; i < V && c > 0; i += DS_NT) { + const unsigned long long u = ds_key(row[i], i); + if ((u & mask) >= prefix) { + const int at = atomicAdd(&s_n, 1); + if (at < DS_KMAX) key_sh[at] = u; + } + } + __syncthreads(); + int m = 1; + while (m < c) m <<= 1; + for (int i = c + tid; i < m; i += DS_NT) key_sh[i] = 0; // sorts last + __syncthreads(); + for (int size = 2; size <= m; size <<= 1) { // bitonic sort, descending + for (int stride = size >> 1; stride > 0; stride >>= 1) { + for (int i = tid; i < m; i += DS_NT) { + const int j = i ^ stride; + if (j > i) { + const bool desc = (i & size) == 0; + const unsigned long long a = key_sh[i], b = key_sh[j]; + if (desc ? a < b : a > b) { key_sh[i] = b; key_sh[j] = a; } + } + } + __syncthreads(); + } + } + for (int i = tid; i < k; i += DS_NT) { + if (i < c) { + const int col = (int)(0xFFFFFFFFull - (key_sh[i] & 0xFFFFFFFFull)); + o[i] = row[col]; + o[k + i] = __int_as_float(id0 + col); + } else { + o[i] = -__int_as_float(0x7F800000); + o[k + i] = __int_as_float(-1); + } + } +} + +// pairs[r] = (f64 vals[r], f64 (cols[r] + id0)), a rank's best; one thread a row +__global__ void __launch_bounds__(256) pick_pack_kernel(const float* __restrict__ vals, const long long* __restrict__ cols, + int n, long long id0, double* __restrict__ pairs) { + const int r = blockIdx.x * blockDim.x + threadIdx.x; + if (r >= n) return; + pairs[2 * r] = static_cast(vals[r]); + pairs[2 * r + 1] = static_cast(cols[r] + id0); +} + +// out[r]: best id over ranks (tie / NaN: lower rank), then accepted, bonus, next position; one block, n <= 1024 +__global__ void __launch_bounds__(1024) pick_merge_kernel(const double* __restrict__ all, int W, int n, + const long long* __restrict__ ids, long long start, + long long* __restrict__ out) { + const int r = threadIdx.x; + if (r < n) { + int bk = 0; + for (int k = 1; k < W; k++) + if (all[2 * (k * n + r)] > all[2 * (bk * n + r)]) bk = k; + out[r] = static_cast(all[2 * (bk * n + r) + 1]); + } + __syncthreads(); + if (r == 0) { + int acc = 0; + if (ids != nullptr) + while (acc + 1 < n && ids[acc + 1] == out[acc]) acc++; + out[n] = acc; + out[n + 1] = out[acc]; + out[n + 2] = start + acc + 1; + } +} + +// one block: DSpark's one-slot statics from the device pick, (anchor, start) = (bonus, next position); rest as staged +__global__ void __launch_bounds__(256) ds_stage_kernel(const long long* __restrict__ pick, int nw, long long* __restrict__ s64, + int* __restrict__ s32, int n, int W, int ring, long long valid, + int noise, int o_pos, int o_tok, int o_cnt, int o_lo, int o_hi, + int o_anchor) { + const long long A = pick[nw + 1], P = pick[nw + 2]; + const long long lo_w = P > W ? P - W : 0; // p -| window + const long long first = lo_w > valid ? lo_w : valid; + const int count = P > first ? (int)(P - first) : 0; + const int t = threadIdx.x; + for (int i = t; i < n; i += blockDim.x) { + s64[i] = i == 0 ? A : (long long)noise; // ids (offset 0) + s64[o_pos + i] = P + i; + s32[o_cnt + i] = count; + s32[o_lo + i] = (int)P; + s32[o_hi + i] = (int)(P + n - 1); + } + for (int e = t; e < n * W; e += blockDim.x) { + const int j = e % W; + s32[o_tok + e] = j < count ? (int)((first + j) % ring) : -1; + } + if (t == 0) { + s32[0] = (int)P; // start (offset 0) + s32[o_anchor] = (int)A; + } +} + +// one thread a slot (segs [k][3]: first row, rows, start): pick_merge's chain walk -> acc [k][3] +__global__ void __launch_bounds__(32) ds_accept_rows_kernel(const long long* __restrict__ picks, const long long* __restrict__ ids, + const long long* __restrict__ segs, int k, long long* __restrict__ acc) { + const int s = threadIdx.x; + if (s >= k) return; + const long long a = segs[3 * s], n = segs[3 * s + 1], start = segs[3 * s + 2]; + long long c = 0; + while (c + 1 < n && ids[a + c + 1] == picks[a + c]) c++; + acc[3 * s] = c; + acc[3 * s + 1] = picks[a + c]; + acc[3 * s + 2] = start + c + 1; +} + +// one block: a pass's members (mem [g][3]: acc index, slot, valid), each one's n rows from its (bonus, next position) +__global__ void __launch_bounds__(256) ds_stage_rows_kernel(const long long* __restrict__ acc, const long long* __restrict__ mem, + int g, long long* __restrict__ s64, int* __restrict__ s32, int n, + int W, int ring, int noise, int o_ids, int o_pos, int o_tok, + int o_cnt, int o_lo, int o_hi) { + const int t = threadIdx.x; + for (int m = 0; m < g; m++) { + const long long j = mem[3 * m], slot = mem[3 * m + 1], valid = mem[3 * m + 2]; + const long long A = acc[3 * j + 1], P = acc[3 * j + 2]; + const long long lo_w = P > W ? P - W : 0; + const long long first = lo_w > valid ? lo_w : valid; + const int count = P > first ? (int)(P - first) : 0; + const int a = m * n; + for (int i = t; i < n; i += blockDim.x) { + s64[o_ids + a + i] = i == 0 ? A : (long long)noise; + s64[o_pos + a + i] = P + i; + s32[o_cnt + a + i] = count; + s32[o_lo + a + i] = (int)P; + s32[o_hi + a + i] = (int)(P + n - 1); + } + for (int e = t; e < n * W; e += blockDim.x) { + const int jj = e % W; + s32[o_tok + a * W + e] = jj < count ? (int)(slot * ring + (first + jj) % ring) : -1; + } + } +} + +__device__ __forceinline__ unsigned long long vs_mix(unsigned long long x) { + x ^= x >> 30; + x *= 0xBF58476D1CE4E5B9ull; + x ^= x >> 27; + x *= 0x94D049BB133111EBull; + return x ^ (x >> 31); +} + +__device__ __forceinline__ double vs_uniform(unsigned long long seed, unsigned long long pos, unsigned long long id) { + unsigned long long x = vs_mix(seed + 0x9E3779B97F4A7C15ull); + x = vs_mix(x ^ (pos * 0xD1B54A32D192ED03ull)); + x = vs_mix(x ^ id); + return (double)(x >> 11) * 0x1p-53 + 0x1p-54; +} + +__device__ __forceinline__ float vs_value(unsigned long long k) { + const unsigned int mono = (unsigned int)(k >> 32); + return __uint_as_float((mono & 0x80000000u) ? (mono & 0x7FFFFFFFu) : ~mono); +} + +// numpy's float64 pairwise sum (leaves of 8..128, halves cut to x8); depth a template, so no runtime recursion +__device__ __forceinline__ double vs_leaf(const double* a, int n) { + if (n < 8) { + double r = -0.0; + for (int i = 0; i < n; i++) r += a[i]; + return r; + } + double r[8]; + for (int j = 0; j < 8; j++) r[j] = a[j]; + int i = 8; + for (; i < n - (n % 8); i += 8) + for (int j = 0; j < 8; j++) r[j] += a[i + j]; + double res = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7])); + for (; i < n; i++) res += a[i]; + return res; +} + +template __device__ __forceinline__ double vs_pairwise(const double* a, int n) { + if (n <= 128) return vs_leaf(a, n); + int n2 = n / 2; + n2 -= n2 % 8; + return vs_pairwise(a, n2) + vs_pairwise(a + n2, n - n2); +} + +template <> __device__ __forceinline__ double vs_pairwise<0>(const double* a, int n) { return vs_leaf(a, n); } + +constexpr int VS_NT = 256, VS_MAX = 2048; + +// one block a row: speculative keyed pick over g [W][n][2 k] (values, ids as int bits); the host's pick decides +__global__ void __launch_bounds__(VS_NT) vs_choose_kernel(const float* __restrict__ g, int W, int n, int k, int count, + unsigned long long seed, long long pos0, int top_k, double temp, + double top_p, double min_p, long long* __restrict__ out) { + __shared__ unsigned long long keys[VS_MAX]; + __shared__ double work[1024]; + const int r = blockIdx.x, t = threadIdx.x, C = W * k; + if (C > VS_MAX) { if (t == 0) out[r] = -1; return; } + for (int i = t; i < C; i += VS_NT) { + const int w = i / k, j = i - w * k; + const float* row = g + ((long long)w * n + r) * 2 * k; + const unsigned int bits = __float_as_uint(row[j] + 0.0f); + const unsigned int mono = (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); + const unsigned int id = __float_as_uint(row[k + j]); + keys[i] = ((unsigned long long)mono << 32) | (0xFFFFFFFFull - id); + } + int m = 1; + while (m < C) m <<= 1; + for (int i = C + t; i < m; i += VS_NT) keys[i] = 0; + __syncthreads(); + for (int size = 2; size <= m; size <<= 1) + for (int stride = size >> 1; stride > 0; stride >>= 1) { + for (int i = t; i < m; i += VS_NT) { + const int j = i ^ stride; + if (j > i) { + const bool desc = (i & size) == 0; + const unsigned long long a = keys[i], b = keys[j]; + if (desc ? a < b : a > b) { keys[i] = b; keys[j] = a; } + } + } + __syncthreads(); + } + if (t != 0) return; + const long long id0 = (long long)(0xFFFFFFFFull - (keys[0] & 0xFFFFFFFFull)); + if (!(temp > 0.0)) { out[r] = id0; return; } // greedy: the first merged candidate + if (top_k == 0 && top_p > 0.0 && top_p < 1.0) { out[r] = -1; return; } // a nucleus row: the host's statistics + const int width = (count > 0 && count < C) ? count : C; // count 0: all W k + int kk = top_k != 0 ? (top_k < width ? top_k : width) : width; + if (kk < 1) kk = 1; + if (kk > 1024) kk = 1024; + const double tt = temp > 1e-6 ? temp : 1e-6; + const double s0 = (double)vs_value(keys[0]) / tt; + int keep = kk; + if (top_p > 0.0 && top_p < 1.0) { + for (int j = 0; j < kk; j++) work[j] = exp((double)vs_value(keys[j]) / tt - s0); + const double total = vs_pairwise<3>(work, kk); // kk <= 1024 + double cum = 0.0; + int below = 0; + for (int j = 0; j < kk; j++) { + cum += work[j] / total; + if (cum < top_p) below++; + } + keep = below + 1; + } + const double floor_ = s0 + (min_p > 0.0 ? log(min_p) : -INFINITY); + const int lim = keep < kk ? keep : kk; + int best = 0; + double best_score = -INFINITY; + for (int j = 0; j < lim; j++) { + const double x = (double)vs_value(keys[j]) / tt; + if (min_p > 0.0 && x < floor_) continue; + const unsigned long long id = 0xFFFFFFFFull - (keys[j] & 0xFFFFFFFFull); + const double score = x - log(-log(vs_uniform(seed, (unsigned long long)(pos0 + r), id))); + if (j == 0 || score > best_score) { best = j; best_score = score; } + } + out[r] = (long long)(0xFFFFFFFFull - (keys[best] & 0xFFFFFFFFull)); +} + +} // namespace dsv41_glue diff --git a/zig/kernels/cuda/deepseek_v41/linear.cu b/zig/kernels/cuda/deepseek_v41/linear.cu new file mode 100644 index 0000000000..b30ad5e4fa --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/linear.cu @@ -0,0 +1,363 @@ +// Copy of src/tensorfold/cuda/exl3/linear.cu (blob 83a8dd081226): sass_check.sh + +// EXL3 linear, 1-128 rows: a row's bits depend only on it (mma keeps rows apart, fixed K ranges, fixed-order sums). + +#include +#include +#include // what the cut torch includes provided: the fixed-width integers +#include + +#include "decode.cuh" + +using namespace tf_exl3; + +namespace tf_exl3_linear { + +enum DType : int { F16 = 0, BF16 = 1, F32 = 2 }; + +__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)); + if (bias) v[j] += __half2float(__ldg(bias + col + j)); + } +} + +// A lane's words of one k step: at 1 and 2 bits the warp loads the step and each lane takes its words by shuffle. +template +__host__ __device__ constexpr bool step_shuffled() { + return K2 == 2 || K2 == 4; +} + +template +__host__ __device__ constexpr int step_regs() { + return step_shuffled() ? tile_words() / 4 : 8 * lane_words(); +} + +template +__device__ __forceinline__ void load_step(const uint32_t* step, int lane, uint32_t (&raw)[step_regs()]) { + constexpr int TW = tile_words(), LW = lane_words(); + if constexpr (step_shuffled()) { +#pragma unroll + for (int c = 0; c < step_regs(); ++c) raw[c] = __ldg(step + c * 32 + lane); + } else { + 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); + } +} + +// Tile j's lane words from a loaded step; prev: the lane's first window starts in the word before its own. +template +__device__ __forceinline__ void step_lane_words(const uint32_t (&raw)[step_regs()], int j, int lane, bool prev, + uint32_t (&w)[lane_words()]) { + constexpr int TW = tile_words(), LW = lane_words(); + if constexpr (step_shuffled()) { + static_assert(LW == 2, "a lane's windows span two words at 1 and 2 bits"); + constexpr int LPW = 8 / K2; // lanes whose windows end in the same word + const uint32_t r = raw[j * TW / 32]; + const int base = (j * TW) % 32; + const uint32_t own = __shfl_sync(0xffffffffu, r, base + lane / LPW); + const uint32_t before = __shfl_sync(0xffffffffu, r, base + (lane / LPW + TW - 1) % TW); + w[0] = prev ? before : own; + w[1] = prev ? own : 0u; // a window inside one word never reads w[1] + } else { +#pragma unroll + for (int q = 0; q < LW; ++q) w[q] = raw[j * LW + q]; + } +} + +__global__ void __launch_bounds__(128) rot_in_kernel(const void* __restrict__ x, int x_dtype, + const half* __restrict__ suh, half* __restrict__ xh, int K) { + const int blk = blockIdx.x * 4 + (threadIdx.x >> 5), row = blockIdx.y, lane = threadIdx.x & 31; + if (blk * 128 >= K) return; + const int k = blk * 128 + 4 * lane; + float v[4], s[4]; + load4(x, x_dtype, (size_t)row * K + k, v); + load4(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(xh, F16, (size_t)row * K + k, v); +} + +template +__global__ void __launch_bounds__(WK * 32) linear_kernel( + const half* __restrict__ xh, const uint32_t* __restrict__ T, long long stride_k, long long stride_nb, + const half* __restrict__ svh, const half* __restrict__ bias, void* __restrict__ y, int y_dtype, + float* __restrict__ Z, int* __restrict__ counters, int M, int K, int N, int SK, const int* __restrict__ skip) { + if (skip != nullptr && *skip != 0) return; // a device-side skip (every program: no counter moves) + constexpr int TW = tile_words(); + constexpr int LW = lane_words(); + extern __shared__ __align__(16) float red[]; // WK * RH * 128 floats + __shared__ int last; + const int RH = min(M, 8); // rows of red a warp + + const int nb = blockIdx.x, split = blockIdx.y, NB = gridDim.x; + 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 = T + nb * stride_nb; + const int col0 = nb * 128; + bool prev = false; + if constexpr (step_shuffled()) { + int word, offset; + lane_start(lane, word, offset); + prev = word != lane / (8 / K2); + } + + for (int m0 = 0, pass = 0; m0 < M; m0 += 16, ++pass) { + const int R = min(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 = xh + (size_t)r0 * K; + const half* x1 = xh + (size_t)r1 * K; + const uint32_t* tile = tiles + (size_t)kt0 * stride_k; + // up to 6 bits the next k step's words are loaded while this one is decoded; 7 and 8 bits load per tile + constexpr bool PF = K2 <= 12; + constexpr int SR = step_regs(); + uint32_t cur[PF ? SR : 1], nxt[PF ? SR : 1]; + if constexpr (PF) load_step(tile, lane, cur); +#pragma unroll 1 + for (int i = 0; i < per_warp; ++i) { + const int kt = kt0 + i; + uint32_t a[4]; + a[0] = __ldg(reinterpret_cast(x0 + kt * 16 + 2 * t)); + a[1] = __ldg(reinterpret_cast(x1 + kt * 16 + 2 * t)); + a[2] = __ldg(reinterpret_cast(x0 + kt * 16 + 2 * t + 8)); + a[3] = __ldg(reinterpret_cast(x1 + kt * 16 + 2 * t + 8)); + if constexpr (PF) { + if (i + 1 < per_warp) load_step(tile + (size_t)(i + 1) * 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, lane, prev, w); + else ldg_lane_words(tile + (size_t)i * stride_k + j * TW, lane, w); + uint32_t b0[2], b1[2]; + decode_lane(w, lane, b0, b1); + mma16816(acc[j][0], a, b0); + mma16816(acc[j][1], a, 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 = min(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, svh, bias, col0 + 4 * lane); + store4(y, 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(Z + ((size_t)split * M + m0 + rlo + r) * N + col0 + c) = s; + } + } + } + if (SK > 1) { + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) last = atomicAdd(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(Z + at)); + for (int q = 1; q < SK; ++q) { + const float4 u = __ldcg(reinterpret_cast(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, svh, bias, col0 + 4 * lane); + store4(y, y_dtype, (size_t)(m0 + r) * N + col0 + 4 * lane, v); + } + if (threadIdx.x == 0) counters[pass * NB + nb] = 0; // every program of the block has arrived + } + } + __syncthreads(); // red is reused in the next pass + } +} + +// W_q [K, N] fp16 from the trellis words (tile (kt, nt) at kt * stride_k + (nt / 8) * stride_nb): one warp per tile. +template +__global__ void __launch_bounds__(32) unpack_kernel(const uint32_t* __restrict__ T, half* __restrict__ W, int N, + int64_t stride_k, int64_t stride_nb) { + const int kt = blockIdx.y, nt = blockIdx.x, lane = threadIdx.x; + uint32_t w[lane_words()]; + ldg_lane_words(T + kt * stride_k + (nt >> 3) * stride_nb + (nt & 7) * tile_words(), lane, w); + uint32_t b[2][2]; + decode_lane(w, lane, b[0], b[1]); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const uint32_t v = b[j >> 2][(j >> 1) & 1]; + const unsigned short h = (j & 1) ? (unsigned short)(v >> 16) : (unsigned short)(v & 0xffffu); + W[(size_t)(kt * 16 + value_row(lane, j)) * N + nt * 16 + value_col(lane, j)] = __ushort_as_half(h); + } +} + +} // namespace tf_exl3_linear + +// Launched at mul1 (TF_EXL3_WIDTHS(X, 2) + K2 3, 5, 7): linear_kernel at WK 2 / 4 / 8, unpack_kernel. +template __global__ void tf_exl3_linear::linear_kernel<2, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<2, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<2, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<2, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<3, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<3, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<3, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<3, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<4, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<4, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<4, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<4, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<5, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<5, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<5, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<5, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<6, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<6, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<6, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<6, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<7, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<7, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<7, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<7, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<8, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<8, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<8, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<8, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<10, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<10, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<10, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<10, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<12, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<12, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<12, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<12, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<14, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<14, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<14, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<14, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<16, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<16, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<16, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<16, 2>(const uint32_t*, half*, int, int64_t, int64_t); diff --git a/zig/kernels/cuda/deepseek_v41/mhc_cuda.cu b/zig/kernels/cuda/deepseek_v41/mhc_cuda.cu new file mode 100644 index 0000000000..92414ed276 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/mhc_cuda.cu @@ -0,0 +1,539 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/mhc_cuda.cu (blob caa24ba4bb72): sass_check.sh +#define DSV41_MHC_NO_TORCH + + + +// TF_DSV41_MHC_CUDA=1: one mHC boundary (post + site + finish + next RMSNorm) as ONE launch for decode rows (R <= 16). + +// Triton's arithmetic operation for operation (same bits); the last CTA's finish reads every partial, so order is moot. + +#include +#include + +namespace dsv41_mhc { + +constexpr int D = 5120; +constexpr int NB = 40; // column blocks: part of the arithmetic +constexpr int CB = D / NB; // 128 columns a block: one FMA chain +constexpr int G8 = CB / 8; // 16-byte groups a block row +constexpr int NMIX = 24; +constexpr int MAXR = 16; +constexpr int NT = 128; // threads a CTA +constexpr int PB = CB + 8; // bf16 pitch of a staged row (272 B: conflict-free 16-byte reads) +constexpr int PF = CB + 4; // fp32 pitch (fp32 fn) +constexpr int NE = 4 * NB; // partial entries a row, (stream, block) order +constexpr int PW = 32; // floats an entry +constexpr int NQ = 7; // float4 groups of an entry the finish reads: [0, 28) +constexpr int STAGE4 = 640; // float4 a finish stage (10 KiB; 2 stages) +constexpr int ROWB = 128; // dynamic shared memory a row (the finish's row sums; nsys reads R from it) +static_assert(D % (NB * 8) == 0 && CB == 128, "the chains are 128 columns"); + +struct Args { + const uint16_t* x; long long xs; // streams [R, 4 D] bf16, row stride xs + uint16_t* xout; // new streams [R, 4 D] (POST) + const void* g; long long gr; int world; int gbf16; // gathered partials [world, R, D], rank stride gr + const float* post; const float* comb; const float* pre; // [R, 4], [R, 16] (comb[i, j] at 4 i + j), [R, 4] + const void* fn; const float* base; const float* scale; // [24, 4 D] bf16 / fp32, [24], [3] + float* part; uint16_t* c; // [R, 4, NB, 32] fp32, the collapsed row [R, D] bf16 + uint16_t* tap; long long ts; // [R, D] row stride ts (or null) + const void* nw; int nwbf16; // norm weight [D] bf16 / fp32 + uint16_t* out; // the next sublayer's input [R, D] bf16 + float* opre; float* opost; float* ocomb; // the next site's coefficients (MIX) + int* cnt; // [0] the ticket, [1] the tail's items, [2] exits (self-cleaning) + int R; + float eps, hc_eps, post_alpha; int iters; + long long spin; // TAIL: SM cycles a CTA waits for the last arrival (0: it leaves) + int defer; // TAIL: the coefficient items run in coef_kernel (a side stream) +}; + +__device__ __forceinline__ uint4 ld_stream(const void* 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; +} +__device__ __forceinline__ uint4 ld_plain(const void* p) { return *reinterpret_cast(p); } +__device__ __forceinline__ void widen8(const uint4 v, float (&f)[8]) { // 8 bf16, column ascending -> fp32 + f[0] = __uint_as_float(v.x << 16); f[1] = __uint_as_float(v.x & 0xffff0000u); + f[2] = __uint_as_float(v.y << 16); f[3] = __uint_as_float(v.y & 0xffff0000u); + f[4] = __uint_as_float(v.z << 16); f[5] = __uint_as_float(v.z & 0xffff0000u); + f[6] = __uint_as_float(v.w << 16); f[7] = __uint_as_float(v.w & 0xffff0000u); +} +__device__ __forceinline__ uint32_t bfbits(float f) { // cvt.rn.bf16.f32: Triton's .to(tl.bfloat16) + unsigned short h; + asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f)); + return h; +} +__device__ __forceinline__ float bfr(float f) { return __uint_as_float(bfbits(f) << 16); } +__device__ __forceinline__ uint4 pack8(const float (&f)[8]) { // values already bf16: exact + return make_uint4(bfbits(f[0]) | (bfbits(f[1]) << 16), bfbits(f[2]) | (bfbits(f[3]) << 16), + bfbits(f[4]) | (bfbits(f[5]) << 16), bfbits(f[6]) | (bfbits(f[7]) << 16)); +} +__device__ __forceinline__ void load8f(const void* p, bool bf, float (&f)[8]) { // 8 bf16 or 8 fp32, read-only + if (bf) { + widen8(ld_stream(p), f); + } else { + const float4 a = __ldg(reinterpret_cast(p)), b = __ldg(reinterpret_cast(p) + 1); + f[0] = a.x; f[1] = a.y; f[2] = a.z; f[3] = a.w; f[4] = b.x; f[5] = b.y; f[6] = b.z; f[7] = b.w; + } +} +__device__ __forceinline__ float ex2a(float x) { // tl.exp's ex2.approx.f32 (not .ftz) + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float divf(float a, float b) { // Triton's fp32 '/': div.full.f32 + float y; + asm("div.full.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ void cp16(void* s, const void* g) { + const uint32_t a = static_cast(__cvta_generic_to_shared(s)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" :: "r"(a), "l"(g) : "memory"); +} +__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;" ::: "memory"); } +template __device__ __forceinline__ void cp_wait() { + asm volatile("cp.async.wait_group %0;" :: "n"(N) : "memory"); +} + +// X'_i of 8 columns: the post expression, one rounding an operation +__device__ __forceinline__ void post1(const Args& a, int r, int i, const float (&br)[8], const float (&x)[4][8], + float (&y)[8]) { + const float p = __ldg(a.post + r * 4 + i); + const float c0 = __ldg(a.comb + r * 16 + i), c1 = __ldg(a.comb + r * 16 + 4 + i); + const float c2 = __ldg(a.comb + r * 16 + 8 + i), c3 = __ldg(a.comb + r * 16 + 12 + i); +#pragma unroll + for (int k = 0; k < 8; ++k) { + const float m = __fadd_rn(__fadd_rn(__fadd_rn(__fmul_rn(c0, x[0][k]), __fmul_rn(c1, x[1][k])), + __fmul_rn(c2, x[2][k])), __fmul_rn(c3, x[3][k])); + y[k] = bfr(__fadd_rn(__fmul_rn(p, br[k]), m)); + } +} + +template struct Smem { + static constexpr int MAIN = 2 * MAXR * PB * 2 + (FN32 ? NMIX * PF * 4 : NMIX * PB * 2); + static constexpr int BYTES = MAIN > 2 * STAGE4 * 16 ? MAIN : 2 * STAGE4 * 16; +}; + +// phase B: NS chains a thread at once (chain c = t + s NT -> row c / K, kind c % K; kind 25 = the collapsed row) +template +__device__ __forceinline__ void chains(const Args& a, const uint16_t* xs, const uint16_t* cs, const void* fns, int K, + int nch, int b, int jp, int t) { + const uint16_t* src[NS]; + const void* wr[NS]; + int row[NS], kind[NS]; + float acc[NS]; +#pragma unroll + for (int s = 0; s < NS; ++s) { + const int cc = min(t + s * NT, nch - 1); + row[s] = cc / K; + kind[s] = MIX ? cc % K : 25; + src[s] = (kind[s] == 25 ? cs : xs) + row[s] * PB; + wr[s] = kind[s] < NMIX ? (FN32 ? static_cast(static_cast(fns) + kind[s] * PF) + : static_cast(static_cast(fns) + kind[s] * PB)) + : static_cast(src[s]); + acc[s] = 0.0f; + } +#pragma unroll + for (int k8 = 0; k8 < G8; ++k8) { +#pragma unroll + for (int s = 0; s < NS; ++s) { + float v[8], w[8]; + widen8(*reinterpret_cast(src[s] + k8 * 8), v); + if (FN32 && kind[s] < NMIX) { + const float4 w0 = reinterpret_cast(wr[s])[2 * k8]; + const float4 w1 = reinterpret_cast(wr[s])[2 * k8 + 1]; + w[0] = w0.x; w[1] = w0.y; w[2] = w0.z; w[3] = w0.w; w[4] = w1.x; w[5] = w1.y; w[6] = w1.z; w[7] = w1.w; + } else { + widen8(reinterpret_cast(wr[s])[k8], w); + } +#pragma unroll + for (int i = 0; i < 8; ++i) acc[s] = __fmaf_rn(v[i], w[i], acc[s]); + } + } +#pragma unroll + for (int s = 0; s < NS; ++s) + if (t + s * NT < nch) a.part[((row[s] * 4 + jp) * NB + b) * PW + kind[s]] = acc[s]; +} + +// row r's coefficients on 16 lanes (comb (i, j) at lane bits 3..2, 1..0): _finish_k's ops and butterfly order +__device__ __forceinline__ void coefs_of(const Args& a, const float* mx, int r, bool store, int e) { + const float rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(mx[24], 4.0f * D), a.eps))); + const float LOG2E = __int_as_float(0x3FB8AA3B), NLOG2E = __int_as_float(0xBFB8AA3B); + if (e < 4) { + const float pl = __fadd_rn(__fmul_rn(__fmul_rn(mx[e], rinv), __ldg(a.scale + 0)), __ldg(a.base + e)); + const float ql = __fadd_rn(__fmul_rn(__fmul_rn(mx[4 + e], rinv), __ldg(a.scale + 1)), __ldg(a.base + 4 + e)); + const float pre = __fadd_rn(a.hc_eps, divf(1.0f, __fadd_rn(ex2a(__fmul_rn(pl, NLOG2E)), 1.0f))); + const float post = __fmul_rn(a.post_alpha, divf(1.0f, __fadd_rn(ex2a(__fmul_rn(ql, NLOG2E)), 1.0f))); + if (store) { a.opre[r * 4 + e] = pre; a.opost[r * 4 + e] = post; } + } + const float cl = __fadd_rn(__fmul_rn(__fmul_rn(mx[8 + e], rinv), __ldg(a.scale + 2)), __ldg(a.base + 8 + e)); + float mxv = fmaxf(cl, __shfl_xor_sync(0xffffffffu, cl, 2)); + mxv = fmaxf(mxv, __shfl_xor_sync(0xffffffffu, mxv, 1)); + const float ce = ex2a(__fmul_rn(__fsub_rn(cl, mxv), LOG2E)); + float s = __fadd_rn(ce, __shfl_xor_sync(0xffffffffu, ce, 2)); + s = __fadd_rn(s, __shfl_xor_sync(0xffffffffu, s, 1)); + float cm = __fadd_rn(a.hc_eps, divf(ce, s)); + s = __fadd_rn(cm, __shfl_xor_sync(0xffffffffu, cm, 8)); + s = __fadd_rn(s, __shfl_xor_sync(0xffffffffu, s, 4)); + cm = divf(cm, __fadd_rn(a.hc_eps, s)); + for (int it = 1; it < a.iters; ++it) { + s = __fadd_rn(cm, __shfl_xor_sync(0xffffffffu, cm, 2)); + s = __fadd_rn(s, __shfl_xor_sync(0xffffffffu, s, 1)); + cm = divf(cm, __fadd_rn(a.hc_eps, s)); + s = __fadd_rn(cm, __shfl_xor_sync(0xffffffffu, cm, 8)); + s = __fadd_rn(s, __shfl_xor_sync(0xffffffffu, s, 4)); + cm = divf(cm, __fadd_rn(a.hc_eps, s)); + } + if (store) a.ocomb[r * 16 + e] = cm; +} +__device__ __forceinline__ void coefs(const Args& a, const float* rowv, int r, bool store, int e) { + coefs_of(a, rowv + r * (ROWB / 4), r, store, e); +} + +__host__ __device__ constexpr int chunk_entries(int R) { // partial entries a finish stage holds a row + return R <= 1 ? 80 : R <= 2 ? 40 : R <= 4 ? 20 : R <= 9 ? 10 : 5; +} + +// the last CTA: row sums (F1), then the coefficients (F3) beside the normed input (F2) +template +__device__ void finish(const Args& a, float* rowv, float4* st, int t) { + const int R = a.R; + if (MIX) { // F1: 7 float4 sums a row, a thread each; partials staged through 2 cp.async stages + const int ch = chunk_entries(R), nck = NE / ch, per = R * ch * NQ; + const int ri = t / NQ, qi = t % NQ; + float4 acc = make_float4(0.f, 0.f, 0.f, 0.f); + float ssc = 0.0f; + auto issue = [&](int c) { + float4* dst = st + (c & 1) * STAGE4; + for (int i = t; i < per; i += NT) { + const int r = i / (ch * NQ), rem = i % (ch * NQ), e = rem / NQ, q = rem % NQ; + cp16(dst + i, a.part + (static_cast(r) * NE + c * ch + e) * PW + q * 4); + } + cp_commit(); + }; + issue(0); + if (nck > 1) issue(1); + for (int c = 0; c < nck; ++c) { + if (c + 1 < nck) cp_wait<1>(); else cp_wait<0>(); + __syncthreads(); + if (t < R * NQ) { + const float4* src = st + (c & 1) * STAGE4 + ri * ch * NQ + qi; +#pragma unroll 5 + for (int e = 0; e < ch; ++e) { + const float4 v = src[e * NQ]; + acc.x = __fadd_rn(acc.x, v.x); + if (qi < 6) { + acc.y = __fadd_rn(acc.y, v.y); acc.z = __fadd_rn(acc.z, v.z); acc.w = __fadd_rn(acc.w, v.w); + } else if (c * ch + e < NB) { + ssc = __fadd_rn(ssc, v.y); // [25]: stream 0's blocks only + } + } + } + __syncthreads(); + if (c + 2 < nck) issue(c + 2); + } + if (t < R * NQ) { + float* rv = rowv + ri * (ROWB / 4); + if (qi < 6) { + rv[4 * qi] = acc.x; rv[4 * qi + 1] = acc.y; rv[4 * qi + 2] = acc.z; rv[4 * qi + 3] = acc.w; + } else { + rv[24] = acc.x; + rv[25] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ssc, 1.0f * D), a.eps))); + } + } + } else if (t < R) { // final: the collapsed row's square sum alone + const float* p = a.part + static_cast(t) * NE * PW + 25; + float v[NB]; +#pragma unroll + for (int e = 0; e < NB; ++e) v[e] = __ldcg(p + e * PW); + float ssc = 0.0f; +#pragma unroll + for (int e = 0; e < NB; ++e) ssc = __fadd_rn(ssc, v[e]); + rowv[t * (ROWB / 4) + 25] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ssc, 1.0f * D), a.eps))); + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int f3w = MIX ? min(4, (R + 1) / 2) : 0; // warps with coefficient rows (2 rows a warp a pass) + if (warp < f3w) { +#pragma unroll 1 + for (int p = 0; p < 2 && warp * 2 + 8 * p < R; ++p) { + const int r = warp * 2 + (lane >> 4) + 8 * p; + coefs(a, rowv, min(r, R - 1), r < R, lane & 15); + } + } + const bool split = f3w < 4; // F2 beside F3 on the other warps, else after it + if (split && warp < f3w) return; + const int wid = split ? t - f3w * 32 : t, nwk = split ? NT - f3w * 32 : NT; + const bool nwbf = a.nwbf16 != 0; +#pragma unroll 4 + for (int i = wid; i < R * (D / 8); i += nwk) { + const int r = i / (D / 8), d0 = (i % (D / 8)) * 8; + const float rn = rowv[r * (ROWB / 4) + 25]; + float c[8], w[8], o[8]; + widen8(__ldcg(reinterpret_cast(a.c + static_cast(r) * D + d0)), c); + load8f(static_cast(a.nw) + d0 * (nwbf ? 2 : 4), nwbf, w); +#pragma unroll + for (int k = 0; k < 8; ++k) o[k] = __fmul_rn(__fmul_rn(c[k], rn), w[k]); + *reinterpret_cast(a.out + static_cast(r) * D + d0) = pack8(o); + } +} + +#include "mhc_tail.cuh" + +// DEFER: a row's coefficient item on its own CTA, on a side stream that joins before the next boundary +__global__ void __launch_bounds__(NT) coef_kernel(const Args a) { + __shared__ __align__(16) float4 st[NE * NQ]; + __shared__ float mx[32]; + coef_item(a, blockIdx.x, st, mx, threadIdx.x); +} + +// POST: post from gathered partials; COLL 1 / 2: c = stream 0 / pre-mix; MIX: next site; ALL4: a CTA writes 4 streams +template +__global__ void __launch_bounds__(NT, 4) boundary_kernel(const Args a) { + extern __shared__ float4 dyn[]; // R x ROWB: the finish's row sums + __shared__ __align__(16) unsigned char sm[Smem::BYTES]; + __shared__ int last; + const int b = blockIdx.x, j = ALL4 ? 0 : blockIdx.y, t = threadIdx.x, R = a.R; + const bool lead = j == 0; + uint16_t* xs = reinterpret_cast(sm); // stream j's new values: the chains' input + uint16_t* cs = xs + MAXR * PB; // the collapsed row (lead) + void* fns = cs + MAXR * PB; // this CTA's fn columns + constexpr int FV = FN32 ? 6 : 3; // 16-byte fn vectors a thread + uint4 fr[FV]; + if (MIX) { // in flight through phase A +#pragma unroll + for (int i = 0; i < FV; ++i) { + const int v = t + i * NT, m = v / (FN32 ? 32 : 16), q = v % (FN32 ? 32 : 16); + fr[i] = ld_stream(static_cast(a.fn) + + (static_cast(m) * 4 * D + j * D + b * CB) * (FN32 ? 4 : 2) + q * 16); + } + } + // ---- phase A: the post, the collapsed row, the taps; 8 columns a thread a step, every row ---- + for (int w = t; w < R * G8; w += NT) { + const int r = w / G8, q = w % G8, col = b * CB + q * 8; + const uint16_t* xr = a.x + static_cast(r) * a.xs + col; + float x[4][8], vj[8]; + if (POST || lead) { +#pragma unroll + for (int s = 0; s < 4; ++s) widen8(ld_plain(xr + s * D), x[s]); + } + if (POST) { + float br[8], gv[8]; + const char* gp = static_cast(a.g); + const int gsz = a.gbf16 ? 2 : 4; + load8f(gp + (static_cast(r) * D + col) * gsz, a.gbf16, br); + for (int rk = 1; rk < a.world; ++rk) { + load8f(gp + (rk * a.gr + static_cast(r) * D + col) * gsz, a.gbf16, gv); +#pragma unroll + for (int k = 0; k < 8; ++k) br[k] = __fadd_rn(br[k], gv[k]); + } +#pragma unroll + for (int k = 0; k < 8; ++k) br[k] = bfr(br[k]); + uint16_t* xo = a.xout + static_cast(r) * 4 * D + col; + if (lead) { // every stream's X' (the collapse, the taps) + float y[4][8]; +#pragma unroll + for (int i = 0; i < 4; ++i) post1(a, r, i, br, x, y[i]); +#pragma unroll + for (int i = 0; i < 4; ++i) + if (ALL4 || i == 0) *reinterpret_cast(xo + i * D) = pack8(y[i]); +#pragma unroll + for (int s = 0; s < 4; ++s) +#pragma unroll + for (int k = 0; k < 8; ++k) x[s][k] = y[s][k]; +#pragma unroll + for (int k = 0; k < 8; ++k) vj[k] = y[0][k]; + } else { + post1(a, r, j, br, x, vj); + *reinterpret_cast(xo + j * D) = pack8(vj); + } + } else if (lead) { +#pragma unroll + for (int k = 0; k < 8; ++k) vj[k] = x[0][k]; + } else { + widen8(ld_plain(xr + j * D), vj); + } + if (MIX) *reinterpret_cast(xs + r * PB + q * 8) = pack8(vj); + if (lead) { + float cv[8]; + if (COLL == 2) { + const float q0 = __ldg(a.pre + r * 4), q1 = __ldg(a.pre + r * 4 + 1); + const float q2 = __ldg(a.pre + r * 4 + 2), q3 = __ldg(a.pre + r * 4 + 3); +#pragma unroll + for (int k = 0; k < 8; ++k) + cv[k] = bfr(__fadd_rn(__fadd_rn(__fadd_rn(__fmul_rn(q0, x[0][k]), __fmul_rn(q1, x[1][k])), + __fmul_rn(q2, x[2][k])), __fmul_rn(q3, x[3][k]))); + } else { +#pragma unroll + for (int k = 0; k < 8; ++k) cv[k] = x[0][k]; + } + const uint4 cp = pack8(cv); + *reinterpret_cast(a.c + static_cast(r) * D + col) = cp; + *reinterpret_cast(cs + r * PB + q * 8) = cp; + if (POST && a.tap != nullptr) { + float tp[8]; +#pragma unroll + for (int k = 0; k < 8; ++k) + tp[k] = bfr(__fmul_rn(__fadd_rn(__fadd_rn(__fadd_rn(x[0][k], x[1][k]), x[2][k]), x[3][k]), 0.25f)); + *reinterpret_cast(a.tap + static_cast(r) * a.ts + col) = pack8(tp); + } + } + } + if (MIX) { +#pragma unroll + for (int i = 0; i < FV; ++i) { + const int v = t + i * NT, m = v / (FN32 ? 32 : 16), q = v % (FN32 ? 32 : 16); + reinterpret_cast(fns)[m * (FN32 ? PF / 4 : PB / 8) + q] = fr[i]; + } + } + __syncthreads(); + // ---- phase B: the FMA chains ---- + const int K = (MIX ? 25 : 0) + (lead ? 1 : 0); + const int nch = R * K; + switch ((nch + NT - 1) / NT) { + case 1: chains<1, FN32, MIX>(a, xs, cs, fns, K, nch, b, j, t); break; + case 2: chains<2, FN32, MIX>(a, xs, cs, fns, K, nch, b, j, t); break; + case 3: chains<3, FN32, MIX>(a, xs, cs, fns, K, nch, b, j, t); break; + case 4: chains<4, FN32, MIX>(a, xs, cs, fns, K, nch, b, j, t); break; + default: break; + } + // ---- the ticket: the last CTA runs the finish ---- + __threadfence(); + __syncthreads(); + if (TAIL) { + tail(a, reinterpret_cast(dyn), reinterpret_cast(sm), t); + return; + } + if (t == 0) last = atomicAdd(a.cnt, 1) == static_cast(gridDim.x * gridDim.y) - 1; + __syncthreads(); + if (!last) return; + __threadfence(); + if (t == 0) atomicExch(a.cnt, 0); // self-cleaning: the next launch / replay starts at 0 + finish(a, reinterpret_cast(dyn), reinterpret_cast(sm), t); +} + +template +cudaError_t launch_one(const Args& a, cudaStream_t stream, bool tail = false) { + const dim3 grid(NB, ALL4 ? 1 : 4); + if (tail) boundary_kernel<<>>(a); + else boundary_kernel<<>>(a); + return cudaGetLastError(); +} + +// mode: 0 boundary (POST, pre-mix collapse, MIX), 1 site from stream 0, 2 site with a carried pre-mix, 3 final +inline cudaError_t coef_launch(const Args& a, cudaStream_t stream) { + if (a.R < 1 || a.R > MAXR) return cudaErrorInvalidValue; + coef_kernel<<>>(a); + return cudaGetLastError(); +} + +// tail: the parallel tail instead of the last CTA's finish (the same bits) +inline cudaError_t dispatch(const Args& a, int mode, bool fn32, cudaStream_t stream, bool tail = false) { + if (a.R < 1 || a.R > MAXR || a.world < 1 || a.world > 8) return cudaErrorInvalidValue; + switch (mode * 2 + (fn32 ? 1 : 0)) { + case 0: return launch_one(a, stream, tail); + case 1: return launch_one(a, stream, tail); + case 2: return launch_one(a, stream, tail); + case 3: return launch_one(a, stream, tail); + case 4: return launch_one(a, stream, tail); + case 5: return launch_one(a, stream, tail); + case 6: case 7: return launch_one(a, stream, tail); + default: return cudaErrorInvalidValue; + } +} + +} // namespace dsv41_mhc + +#ifndef DSV41_MHC_NO_TORCH // the compile test builds the kernels alone (nvcc, no torch headers) +#include +#include +#include + +namespace { +template T* ptr(const at::Tensor& t) { return t.numel() ? static_cast(t.data_ptr()) : nullptr; } +} + +void dsv41_mhc_run_cuda(const at::Tensor& x, at::Tensor& xout, const at::Tensor& g, const at::Tensor& post, + const at::Tensor& comb, const at::Tensor& pre, const at::Tensor& fn, const at::Tensor& base, + const at::Tensor& scale, at::Tensor& part, at::Tensor& c, at::Tensor& tap, + const at::Tensor& nw, at::Tensor& out, at::Tensor& opre, at::Tensor& opost, + at::Tensor& ocomb, at::Tensor& cnt, int64_t R, int64_t mode, double eps, double hc_eps, + double post_alpha, int64_t iters, int64_t spin, int64_t defer) { + using namespace dsv41_mhc; + Args a{}; + a.x = ptr(x); a.xs = x.stride(0); + a.xout = ptr(xout); + a.g = g.numel() ? g.data_ptr() : nullptr; a.gr = g.numel() ? g.stride(0) : 0; + a.world = g.numel() ? static_cast(g.size(0)) : 1; + a.gbf16 = g.scalar_type() == at::kBFloat16 ? 1 : 0; + a.post = ptr(post); a.comb = ptr(comb); a.pre = ptr(pre); + a.fn = fn.numel() ? fn.data_ptr() : nullptr; a.base = ptr(base); a.scale = ptr(scale); + a.part = ptr(part); a.c = ptr(c); + a.tap = ptr(tap); a.ts = tap.numel() ? tap.stride(0) : 0; + a.nw = nw.data_ptr(); a.nwbf16 = nw.scalar_type() == at::kBFloat16 ? 1 : 0; + a.out = ptr(out); + a.opre = ptr(opre); a.opost = ptr(opost); a.ocomb = ptr(ocomb); + a.cnt = ptr(cnt); + a.R = static_cast(R); + a.eps = static_cast(eps); a.hc_eps = static_cast(hc_eps); + a.post_alpha = static_cast(post_alpha); a.iters = static_cast(iters); + a.spin = spin; // < 0: the last CTA's finish; >= 0: the parallel tail + a.defer = spin >= 0 && defer ? 1 : 0; + const bool fn32 = fn.numel() && fn.scalar_type() == at::kFloat; + C10_CUDA_CHECK(dispatch(a, static_cast(mode), fn32, at::cuda::getCurrentCUDAStream(), spin >= 0)); +} + +void dsv41_mhc_coef_cuda(const at::Tensor& part, const at::Tensor& base, const at::Tensor& scale, at::Tensor& opre, + at::Tensor& opost, at::Tensor& ocomb, int64_t R, double eps, double hc_eps, double post_alpha, + int64_t iters) { + using namespace dsv41_mhc; + Args a{}; + a.part = ptr(part); a.base = ptr(base); a.scale = ptr(scale); + a.opre = ptr(opre); a.opost = ptr(opost); a.ocomb = ptr(ocomb); + a.R = static_cast(R); + a.eps = static_cast(eps); a.hc_eps = static_cast(hc_eps); + a.post_alpha = static_cast(post_alpha); a.iters = static_cast(iters); + C10_CUDA_CHECK(coef_launch(a, at::cuda::getCurrentCUDAStream())); +} + +std::vector dsv41_mhc_info_cuda(int64_t device) { // what the GPU offers for cross-CTA reductions + int cl = 0, sms = 0, optin = 0, maxcl = 0; + cudaDeviceGetAttribute(&cl, cudaDevAttrClusterLaunch, static_cast(device)); + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, static_cast(device)); + cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, static_cast(device)); + if (cl) { + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(dsv41_mhc::NB, 4); + cfg.blockDim = dim3(dsv41_mhc::NT); + cfg.dynamicSmemBytes = dsv41_mhc::MAXR * dsv41_mhc::ROWB; + cudaFuncSetAttribute(dsv41_mhc::boundary_kernel, + cudaFuncAttributeNonPortableClusterSizeAllowed, 1); + if (cudaOccupancyMaxPotentialClusterSize(&maxcl, dsv41_mhc::boundary_kernel, + &cfg) != cudaSuccess) { + maxcl = -1; + cudaGetLastError(); + } + } + int blocks = 0; + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, dsv41_mhc::boundary_kernel, + dsv41_mhc::NT, dsv41_mhc::MAXR * dsv41_mhc::ROWB); + return {cl, maxcl, sms, optin, blocks}; +} +#endif + +// The instantiations dispatch() launches: modes 0-2 at bf16 / fp32 fn, mode 3 at bf16; tail off and on. +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); diff --git a/zig/kernels/cuda/deepseek_v41/mhc_tail.cuh b/zig/kernels/cuda/deepseek_v41/mhc_tail.cuh new file mode 100644 index 0000000000..e3f236de73 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/mhc_tail.cuh @@ -0,0 +1,94 @@ +// mhc_cuda.cu's parallel tail (TF_DSV41_MHC_TAIL=1): coef/norm items and tail(), inside namespace dsv41_mhc +#pragma once +// ---- parallel tail (TF_DSV41_MHC_TAIL=1): live CTAs claim finish()'s items from a counter, same chains, same bits +constexpr int F2C = NT * 8; // columns a normed-input item +constexpr int F2N = D / F2C; // items a row +static_assert(D % F2C == 0 && NE * NQ * 16 <= 2 * STAGE4 * 16, "tail items"); + +__device__ __forceinline__ int ld_acquire(const int* p) { + int v; + asm volatile("ld.acquire.gpu.b32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ float rn_of(float ssc, float eps) { + return __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ssc, 1.0f * D), eps))); +} + +// row r's 25 sums over its 160 entries and the square sum over stream 0's 40 (finish()'s chains), then coefs() +__device__ void coef_item(const Args& a, int r, float4* st, float* mx, int t) { + for (int i = t; i < NE * NQ; i += NT) + cp16(st + i, a.part + (static_cast(r) * NE + i / NQ) * PW + (i % NQ) * 4); + cp_commit(); + cp_wait<0>(); + __syncthreads(); + if (t < 26) { + const float* sf = reinterpret_cast(st); + const int n = t == 25 ? NB : NE; + float acc = 0.0f; +#pragma unroll 8 + for (int e = 0; e < n; ++e) acc = __fadd_rn(acc, sf[e * (NQ * 4) + t]); + mx[t] = t == 25 ? rn_of(acc, a.eps) : acc; + } + __syncthreads(); + if (t < 32) coefs_of(a, mx, r, t < 16, t & 15); +} + +// normed input of row r, columns [F2C q, F2C q + F2C): bf16((c * rn) * w) +__device__ void norm_item(const Args& a, int r, int q, float* buf, int t) { + const int d0 = q * F2C + t * 8; + const bool nwbf = a.nwbf16 != 0; + float c[8], w[8], o[8]; + widen8(__ldcg(reinterpret_cast(a.c + static_cast(r) * D + d0)), c); + load8f(static_cast(a.nw) + d0 * (nwbf ? 2 : 4), nwbf, w); + if (t < NB) buf[t] = __ldcg(a.part + (static_cast(r) * NE + t) * PW + 25); + __syncthreads(); + if (t == 0) { + float ssc = 0.0f; +#pragma unroll 8 + for (int e = 0; e < NB; ++e) ssc = __fadd_rn(ssc, buf[e]); + buf[NB] = rn_of(ssc, a.eps); + } + __syncthreads(); + const float rn = buf[NB]; +#pragma unroll + for (int k = 0; k < 8; ++k) o[k] = __fmul_rn(__fmul_rn(c[k], rn), w[k]); + *reinterpret_cast(a.out + static_cast(r) * D + d0) = pack8(o); +} + +template +__device__ void tail(const Args& a, float* rowv, float4* st, int t) { + __shared__ int live, item; + __shared__ float nbuf[NB + 1]; // a normed-input item's 40 square sums + rn + const int total = static_cast(gridDim.x * gridDim.y), R = a.R; + if (t == 0) { + int n = atomicAdd(a.cnt, 1) + 1; + if (n < total && a.spin > 0) { + const long long t0 = clock64(); + while ((n = ld_acquire(a.cnt)) < total && clock64() - t0 < a.spin) __nanosleep(32); + } + live = n >= total; + } + __syncthreads(); + if (live) { + __threadfence(); + const int nc = MIX && !a.defer ? R : 0, nitems = nc + R * F2N; + for (;;) { + if (t == 0) item = atomicAdd(a.cnt + 1, 1); + __syncthreads(); + const int it = item; + if (it >= nitems) break; + if (it < nc) { + coef_item(a, it, st, rowv, t); + } else { + const int f = it - nc; + norm_item(a, f / F2N, f % F2N, nbuf, t); + } + __syncthreads(); + } + } + if (t == 0 && atomicAdd(a.cnt + 2, 1) == total - 1) { // the last to leave: every counter back at 0 + atomicExch(a.cnt, 0); + atomicExch(a.cnt + 1, 0); + atomicExch(a.cnt + 2, 0); + } +} diff --git a/zig/kernels/cuda/deepseek_v41/nvcc-docker.sh b/zig/kernels/cuda/deepseek_v41/nvcc-docker.sh new file mode 100755 index 0000000000..e71d12832e --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/nvcc-docker.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# nvcc from the NGC PyTorch image for hosts without CUDA (`zig build -Dnvcc=`); paths mount 1:1. +set -euo pipefail +IMAGE="${TF_ZIG_IMAGE:-nvcr.io/nvidia/pytorch:26.07-py3}" +# the container's memory and memory+swap cap: the Docker daemon runs it outside the caller's own limits +MEM="${TF_ZIG_DOCKER_MEMORY:-8g}" +# the mount of a directory: its repository's root (a source's headers sit beside it), else the directory (a cache) +root() { git -C "$1" rev-parse --show-toplevel 2>/dev/null || printf '%s\n' "$1"; } +roots=("$PWD" "$(root "$PWD")") +for a in "$@"; do + [ "${a:0:1}" = / ] || continue + d="$(dirname "$a")" + mkdir -p "$d" + roots+=("$(root "$d")") +done +mapfile -t roots < <(printf '%s\n' "${roots[@]}" | sort -u) +mounts=() +for r in "${roots[@]}"; do mounts+=(-v "$r:$r"); done +docker run --rm -i --network none --memory "$MEM" --memory-swap "$MEM" -u "$(id -u):$(id -g)" "${mounts[@]}" -w "$PWD" \ + --entrypoint /usr/local/cuda/bin/nvcc "$IMAGE" "$@" +# depfilter FILE CWD ROOT...: keep the -MF file's dependencies under a mounted root (image headers drop out) +depfilter() { + local path=$1 cwd=$2 text target d full keep=() out=() reals k deps=() + shift 2 + for r in "$@"; do keep+=("$(realpath -m -- "$r")/"); done + text=$(sed -e ':a' -e '/\\$/{N;s/\\\n/ /;ba' -e '}' "$path") + target=${text%%:*} + read -r -d "" -a deps <<< "${text#*:}" || true + [ "${#deps[@]}" -gt 0 ] && for d in "${deps[@]}"; do + if [ "${d:0:1}" = / ]; then full=$d; else full=$cwd/$d; fi + reals=$(realpath -m -- "$full") + for k in "${keep[@]}"; do [[ "$reals" == "$k"* ]] && { out+=("$d"); break; }; done + done + { printf '%s: ' "$target"; local first=1; for d in "${out[@]}"; do [ $first = 1 ] && printf '%s' "$d" || printf ' \\\n %s' "$d"; first=0; done; printf '\n'; } > "$path" +} +dep="" +prev="" +for a in "$@"; do + [ "$prev" = "-MF" ] && dep="$a" + prev="$a" +done +if [ -n "$dep" ] && [ -f "$dep" ]; then + depfilter "$dep" "$PWD" "${roots[@]}" +fi diff --git a/zig/kernels/cuda/deepseek_v41/pointwise.cu b/zig/kernels/cuda/deepseek_v41/pointwise.cu new file mode 100644 index 0000000000..eefb387097 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/pointwise.cu @@ -0,0 +1,40 @@ +// Torch pointwise ops of DeepSeek-V4.1's fast prefill MoE, bit-exact: fp32 opmath, NaN-propagating clamps. + +// swiglu: act = silu(min(g, limit)) * clamp(u, -limit, limit); add: out += b; add_bf16: y = bf16_rn(a + b). + +// Built with --fmad=false: torch runs each op as its own kernel with fp32 rounding, so nothing may contract. + +#include +#include + +namespace dsv41_pointwise { + +// torch's clamp_max / clamp: a NaN stays NaN +__device__ __forceinline__ float clamp_max(float v, float hi) { return isnan(v) ? v : fminf(v, hi); } +__device__ __forceinline__ float clamp(float v, float lo, float hi) { return isnan(v) ? v : fminf(fmaxf(v, lo), hi); } +// torch's silu (ActivationSiluKernel.cu): x / (1 + exp(-x)) in fp32 +__device__ __forceinline__ float silu(float x) { return x / (1.0f + expf(-x)); } + +// g [R, ldg], u [R, ldu], act [R, lda]; n columns a row +__global__ void __launch_bounds__(256) swiglu_kernel(const float* __restrict__ g, long long ldg, + const float* __restrict__ u, long long ldu, + float* __restrict__ act, long long lda, int n, float limit) { + const int r = blockIdx.y; + for (int c = blockIdx.x * blockDim.x + threadIdx.x; c < n; c += gridDim.x * blockDim.x) { + const float a = silu(clamp_max(g[r * ldg + c], limit)); + act[r * lda + c] = a * clamp(u[r * ldu + c], -limit, limit); + } +} + +__global__ void __launch_bounds__(256) add_kernel(float* __restrict__ out, const float* __restrict__ b, long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + out[i] = out[i] + b[i]; +} + +__global__ void __launch_bounds__(256) add_bf16_kernel(const float* __restrict__ a, const float* __restrict__ b, + __nv_bfloat16* __restrict__ y, long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + y[i] = __float2bfloat16_rn(a[i] + b[i]); +} + +} // namespace dsv41_pointwise diff --git a/zig/kernels/cuda/deepseek_v41/proto/attn_fastmerge.cuh b/zig/kernels/cuda/deepseek_v41/proto/attn_fastmerge.cuh new file mode 100644 index 0000000000..7409ec1763 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/proto/attn_fastmerge.cuh @@ -0,0 +1,499 @@ +// attn_cuda.cu with a latency-hiding merge (all chunk loads issued before the chains); same bits as attn_cuda.cu. +#ifndef DSV41_ATTN_NO_TORCH +#define DSV41_ATTN_NO_TORCH +#endif +// Copy of src/tensorfold/families/deepseek_v41/cuda/csa2/attn_cuda.cu (blob 0f3234101360): sass_check.sh + + + +// CSA2 decode attention core in ONE launch: CTA (chunk of 5, 16-head tile, row), 8 warps, ops = Triton's sm_121 PTX +#include +#include +#include +#include + +namespace dsv41_attn_fm { + +constexpr int D = 512, CH = 128, NCOMP = 4, NCH = 5, BMQ = 16, WINDOW = 128; // tiles of 32 keys +constexpr int NOPE = 448, VB = 576, SB = 8, NT = 256; +constexpr int SM_Q = 0; // q A fragments [32 steps][32 lanes] uint4 16384 +constexpr int SM_KV = 16384; // 128 rows x 576 B, 8-byte pieces swizzled by key & 7 73728 +constexpr int SM_SC = SM_KV + CH * VB; // 128 x 8 scale bytes 1024 +constexpr int SM_P = SM_SC + CH * SB; // P A fragments [4 tiles][2 steps][32 lanes] uint4 4096 +constexpr int SM_TMAX = SM_P + 4096; // [4 tiles][4 n8][16 rows] fp32 1024 +constexpr int SM_WSUM = SM_TMAX + 1024; // same shape: the warps' row sums 1024 +constexpr int SM_ROW = SM_WSUM + 1024; // [4 tiles][16 rows] nm, alpha, active 768 +constexpr int SM_PHYS = SM_ROW + 768; // 128 physical rows (int32, -1: masked) 512 +constexpr int SM_MISC = SM_PHYS + 512; // final m, l [16] each + the ticket flag 256 +constexpr int SMEM = SM_MISC + 256; // 98,816 B (sm_121: <= 101,376 a block) +static_assert(SMEM <= 101376, "shared memory"); +constexpr unsigned FULL = 0xffffffffu; + +struct Args { + const uint16_t* q; // [R, H, 512] bf16 (RoPE'd) + const uint8_t* cv; const uint8_t* csc; long long cvs, css; // compressed rows (values, scales), row strides + const int* tok; long long ts; const int* cnt; // selection [R, ts] int32 (-1 pad), counts [R] + const uint8_t* sv; const uint8_t* ssc; long long svs, sss; // SWA rings + const int* lo; const int* hi; // [R] window starts; window ends (HAS_HI) or null + const void* pos; // int32 [1] (POS + r) or int64 [R] (ROWS) + const long long* sl; // ROWS: the rows' slots [R] + const int* pt; long long pts; int psh; // page table of the compressed rows (psh 0: contiguous) + const float* sink; const float* cs; long long csst; // [H] fp32, RoPE table fp32 [max_pos, 64] + float* po; float* pm; float* pl; // partials (attn.Scratch layout) + uint16_t* out; // [R, H, 512] bf16 + int* ticket; // [R, H / 16] (self-cleaning) + int R, H, ring, has_comp; +}; + +__device__ __forceinline__ float ex2(float x) { + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float div_full(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float texp(float x) { return ex2(__fmul_rn(x, __int_as_float(0x3FB8AA3B))); } +__device__ __forceinline__ float pow2b(uint32_t s) { return __int_as_float(static_cast(s) << 23); } + +__device__ __forceinline__ void mma(float (&d)[4], const uint4& a, uint32_t b0, uint32_t b1) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b0), "r"(b1)); +} + +// two e4m3 bytes (low byte = the lower k) with their own scales -> bf16x2 (exact: e4m3 x 2^e is a bf16) +__device__ __forceinline__ uint32_t deq2(uint32_t two, float s0, float s1) { + uint32_t h2; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(h2) : "h"(static_cast(two))); + const __half2 h = *reinterpret_cast(&h2); + const __nv_bfloat162 b = __floats2bfloat162_rn(__fmul_rn(__low2float(h), s0), __fmul_rn(__high2float(h), s1)); + return *reinterpret_cast(&b); +} + +// a byte offset x of row `key` in the staged rows: 8-byte pieces swizzled within groups of 8 (bank-conflict free) +__device__ __forceinline__ int swz(int key, int x) { + const int ch = x >> 3; + return key * VB + ((((ch & ~7) | ((ch & 7) ^ (key & 7)))) << 3) + (x & 7); +} + +// Triton's dot-operand k order: step s (16 k) of a 32-k block takes k = klo(s), klo + 1 and khi(s), khi + 1 +template __device__ __forceinline__ int klo(int s, int t) { + return KW == 4 ? 32 * (s >> 1) + 4 * t + 2 * (s & 1) : 16 * s + 2 * t; +} +template __device__ __forceinline__ int khi(int s, int t) { return klo(s, t) + (KW == 4 ? 16 : 8); } +// where element k (< 32) of a row sits in a 2-step A fragment: step, quad lane, high half (+16 / +8) +template __device__ __forceinline__ void frag_of(int k, int& step, int& tq, int& hi) { + if (KW == 4) { step = (k >> 1) & 1; hi = (k >> 4) & 1; tq = (k & 15) >> 2; } + else { step = k >> 4; hi = (k >> 3) & 1; tq = (k & 7) >> 1; } +} + +// the bf16 pair (k, k + 1) of staged row `key` (S's B operand: k along the row) +__device__ __forceinline__ uint32_t pair_k(const uint8_t* sm, int key, int k) { + if (k < NOPE) { + const uint32_t two = *reinterpret_cast(sm + SM_KV + swz(key, k)); + const float s = pow2b(sm[SM_SC + key * SB + (k >> 6)]); + return deq2(two, s, s); + } + return *reinterpret_cast(sm + SM_KV + swz(key, NOPE + 2 * (k - NOPE))); +} + +// value (key, n) of the staged rows as a bf16 bit pattern / its e4m3 byte (PV's B operand: k along the keys) +__device__ __forceinline__ uint32_t pair_n(const uint8_t* sm, int k0, int k1, int n) { + if (n < NOPE) { + const uint32_t two = sm[SM_KV + swz(k0, n)] | (static_cast(sm[SM_KV + swz(k1, n)]) << 8); + return deq2(two, pow2b(sm[SM_SC + k0 * SB + (n >> 6)]), pow2b(sm[SM_SC + k1 * SB + (n >> 6)])); + } + const int x = NOPE + 2 * (n - NOPE); + return *reinterpret_cast(sm + SM_KV + swz(k0, x)) | + (static_cast(*reinterpret_cast(sm + SM_KV + swz(k1, x))) << 16); +} + +__device__ __forceinline__ void cp8(uint32_t dst, const void* src, bool ok) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8, %2;" :: "r"(dst), "l"(src), "r"(ok ? 8 : 0)); +} + +template __device__ __forceinline__ int row_pos(const Args& a, int r) { + return ROWS ? static_cast(static_cast(a.pos)[r]) : static_cast(a.pos)[0] + r; +} + +// how many chunks of row r have keys (the merge's ticket target; the others are Triton's empty partial) +__device__ __forceinline__ bool chunk_works(const Args& a, int r, int c) { + return c == NCOMP || (a.has_comp && c * CH < a.cnt[r]); +} + +template +__device__ void merge(const Args& a, uint8_t* sm, int r, int hb) { + const int tid = threadIdx.x, H = a.H, R = a.R; + float* ca = reinterpret_cast(sm + SM_TMAX); // [5][16] a, [5][16] b, [16] af, [16] den + float* cb = ca + NCH * BMQ; + float* fa = cb + NCH * BMQ; + float* fd = fa + BMQ; + if (tid < BMQ) { // a head's scalar chain (the merge, unrolled) + const int h = hb * BMQ + tid; + float m = -INFINITY, l = 0.f; + float pmv[NCH], plv[NCH]; // every chunk's (m, l) in flight at once +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const bool w = chunk_works(a, r, c); + const long long base = (static_cast(c) * R + r) * H + h; + pmv[c] = w ? __ldcg(a.pm + base) : -INFINITY; + plv[c] = w ? __ldcg(a.pl + base) : 0.f; + } +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const float cm = pmv[c], cl = plv[c]; + const bool act = cl > 0.f; + float aa, bb; + if (c == 0) { + const float nm = act ? fmaxf(cm, -INFINITY) : -INFINITY; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + aa = 0.f; + l = __fmaf_rn(bb, cl, 0.f); + m = nm; + } else { + const float nm = act ? fmaxf(m, cm) : m; + aa = act ? (m == -INFINITY ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + l = __fmaf_rn(l, aa, __fmul_rn(bb, cl)); + m = nm; + } + ca[c * BMQ + tid] = aa; + cb[c * BMQ + tid] = bb; + } + const float sink = a.sink[h]; + const float top = fmaxf(m, sink); + const float af = m == -INFINITY ? 0.f : texp(__fsub_rn(m, top)); + const float e = sink == -INFINITY ? 0.f : texp(__fsub_rn(sink, top)); + fa[tid] = af; + fd[tid] = __fmaf_rn(l, af, e); + } + __syncthreads(); + const int at = row_pos(a, r); + const float* csr = a.cs + static_cast(at) * a.csst; + constexpr int UB = BMQ * (D / 4) / NT; // 8 units a thread, every load issued first + bool cw[NCH]; +#pragma unroll + for (int c = 0; c < NCH; ++c) cw[c] = chunk_works(a, r, c); + float4 pv[UB][NCH]; +#pragma unroll + for (int ub = 0; ub < UB; ++ub) { + const int u = tid + ub * NT, hl = u / (D / 4), d = (u % (D / 4)) * 4, h = hb * BMQ + hl; +#pragma unroll + for (int c = 0; c < NCH; ++c) + pv[ub][c] = cw[c] ? __ldcg(reinterpret_cast(a.po + ((static_cast(c) * R + r) * H + h) * D + d)) + : make_float4(0.f, 0.f, 0.f, 0.f); + } +#pragma unroll + for (int ub = 0; ub < UB; ++ub) { // float4 units: head hl, columns d .. d + 3 + const int u = tid + ub * NT, hl = u / (D / 4), d = (u % (D / 4)) * 4, h = hb * BMQ + hl; + float o[4]; +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const float4 v = pv[ub][c]; + const float co[4] = {v.x, v.y, v.z, v.w}; + const float aa = ca[c * BMQ + hl], bb = cb[c * BMQ + hl]; + for (int i = 0; i < 4; ++i) + o[i] = c == 0 ? __fmaf_rn(bb, co[i], 0.f) : __fmaf_rn(o[i], aa, __fmul_rn(co[i], bb)); + } + float x[4]; + for (int i = 0; i < 4; ++i) x[i] = div_full(__fmul_rn(o[i], fa[hl]), fd[hl]); + uint16_t ob[4]; + for (int i = 0; i < 4; i += 2) { // pair (d + i, d + i + 1): GPT-J, inverse + const int p = (d + i) >> 1; + const float c = p >= 224 ? csr[p - 224] : 1.f, s = p >= 224 ? csr[32 + p - 224] : 0.f; + const float ne = __fmaf_rn(x[i], c, __fmul_rn(x[i + 1], s)); + const float no = __fmaf_rn(x[i + 1], c, -__fmul_rn(x[i], s)); + const __nv_bfloat16 b0 = __float2bfloat16_rn(ne), b1 = __float2bfloat16_rn(no); + ob[i] = *reinterpret_cast(&b0); + ob[i + 1] = *reinterpret_cast(&b1); + } + uint2 w; + w.x = ob[0] | (static_cast(ob[1]) << 16); + w.y = ob[2] | (static_cast(ob[3]) << 16); + *reinterpret_cast(a.out + (static_cast(r) * H + h) * D + d) = w; + } +} + +template +__global__ void __launch_bounds__(NT, 1) attn_kernel(Args a) { + extern __shared__ __align__(16) uint8_t sm[]; + const int c = blockIdx.x, hb = blockIdx.y, r = blockIdx.z; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, t = lane & 3; + const int H = a.H, R = a.R; + if (!chunk_works(a, r, c)) return; + int* phys = reinterpret_cast(sm + SM_PHYS); + const long long sl = ROWS ? a.sl[r] : 0; + if (tid < CH) { // the chunk's 128 keys -> physical rows (-1: masked) + int row = -1; + if (c < NCOMP) { + const int idx = c * CH + tid; + if (idx < a.cnt[r]) { + const int tk = a.tok[static_cast(r) * a.ts + idx]; + if (tk >= 0) { + row = tk; + if (a.psh > 0) { + const int* ptr = a.pt + (ROWS ? sl * a.pts : 0); + row = (ptr[tk >> a.psh] << a.psh) + (tk & ((1 << a.psh) - 1)); + } + } + } + } else { + const int hi = HAS_HI ? a.hi[r] : row_pos(a, r); + const int lo = max(a.lo[r], hi - (WINDOW - 1)); + const int p = hi - (WINDOW - 1) + tid; + if (p >= lo && p >= 0) row = (p % a.ring) + (ROWS ? static_cast(sl) * a.ring : 0); + } + phys[tid] = row; + } + { // q -> A fragments (Triton's k order) + const uint16_t* qb = a.q + (static_cast(r) * H + hb * BMQ) * D; + for (int i = tid; i < 32 * 32; i += NT) { + const int s = i >> 5, ln = i & 31, gg = ln >> 2, tt = ln & 3; + const int k0 = klo(s, tt), k1 = khi(s, tt); + uint4 f; + f.x = *reinterpret_cast(qb + gg * D + k0); + f.y = *reinterpret_cast(qb + (gg + 8) * D + k0); + f.z = *reinterpret_cast(qb + gg * D + k1); + f.w = *reinterpret_cast(qb + (gg + 8) * D + k1); + reinterpret_cast(sm + SM_Q)[i] = f; + } + } + __syncthreads(); + { // the rows: 72 value pieces + 1 scale piece a key + const bool comp = c < NCOMP; + const uint8_t* V = comp ? a.cv : a.sv; + const uint8_t* S = comp ? a.csc : a.ssc; + const long long vs = comp ? a.cvs : a.svs, ss = comp ? a.css : a.sss; + const uint32_t base = static_cast(__cvta_generic_to_shared(sm)); + for (int i = tid; i < CH * 73; i += NT) { + const int key = i / 73, pc = i % 73, row = phys[key]; + const bool ok = row >= 0; + const long long rr = ok ? row : 0; + if (pc < 72) cp8(base + SM_KV + swz(key, pc * 8), V + rr * vs + pc * 8, ok); + else cp8(base + SM_SC + key * SB, S + rr * ss, ok); + } + asm volatile("cp.async.commit_group;\ncp.async.wait_group 0;" ::: "memory"); + } + __syncthreads(); + const float NINF = -INFINITY; + const int n8 = warp & 3, ta = warp >> 2; // this warp: S of tiles ta and ta + 2, block n8 + float sv[2][4] = {{0.f, 0.f, 0.f, 0.f}, {0.f, 0.f, 0.f, 0.f}}; + { + const int kA = 32 * ta + 8 * n8 + g, kB = kA + 64; + const uint4* qf = reinterpret_cast(sm + SM_Q); +#pragma unroll 4 + for (int s = 0; s < 32; ++s) { + const uint4 af = qf[s * 32 + lane]; + const int k0 = klo(s, t), k1 = khi(s, t); + mma(sv[0], af, pair_k(sm, kA, k0), pair_k(sm, kA, k1)); + mma(sv[1], af, pair_k(sm, kB, k0), pair_k(sm, kB, k1)); + } + } + float* tmax = reinterpret_cast(sm + SM_TMAX); + float* wsum = reinterpret_cast(sm + SM_WSUM); + float* rowv = reinterpret_cast(sm + SM_ROW); // [tile][row]: nm, alpha, active at 0 / 64 / 128 + bool val[2][2]; + for (int j = 0; j < 2; ++j) { + const int tile = ta + 2 * j, kc = 32 * tile + 8 * n8 + 2 * t; + val[j][0] = phys[kc] >= 0; + val[j][1] = phys[kc + 1] >= 0; + for (int e = 0; e < 4; ++e) + sv[j][e] = val[j][e & 1] ? __fmul_rn(sv[j][e], __int_as_float(0x3D3504F3)) : NINF; + float m0 = fmaxf(sv[j][0], sv[j][1]), m1 = fmaxf(sv[j][2], sv[j][3]); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 2)); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 1)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 2)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 1)); + if (t == 0) { + tmax[(tile * 4 + n8) * 16 + g] = m0; + tmax[(tile * 4 + n8) * 16 + g + 8] = m1; + } + } + __syncthreads(); + if (tid < BMQ) { // the running max over the tiles, a row a thread + float m = NINF; + for (int tile = 0; tile < 4; ++tile) { + const float* tm4 = tmax + tile * 64 + tid; + const float tm = fmaxf(fmaxf(tm4[0], tm4[16]), fmaxf(tm4[32], tm4[48])); + const bool act = tm != NINF; + const float nm = act ? fmaxf(m, tm) : m; + const float al = act ? (m == NINF ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + rowv[tile * 16 + tid] = nm; + rowv[64 + tile * 16 + tid] = al; + rowv[128 + tile * 16 + tid] = act ? 1.f : 0.f; + m = nm; + } + reinterpret_cast(sm + SM_MISC)[tid] = m; + } + __syncthreads(); + uint32_t* pf = reinterpret_cast(sm + SM_P); // [tile][step][lane][4] + for (int j = 0; j < 2; ++j) { // p, the thread / quad sums, bf16(p) -> fragments + const int tile = ta + 2 * j; + float p[4]; + for (int e = 0; e < 4; ++e) { + const int row = g + (e >> 1) * 8; + const bool act = rowv[128 + tile * 16 + row] != 0.f; + p[e] = (val[j][e & 1] && act) ? texp(__fsub_rn(sv[j][e], rowv[tile * 16 + row])) : 0.f; + } + float a0 = __fadd_rn(p[0], p[1]), a1 = __fadd_rn(p[2], p[3]); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 2)); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 1)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 2)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 1)); + if (t == 0) { + wsum[(tile * 4 + n8) * 16 + g] = a0; + wsum[(tile * 4 + n8) * 16 + g + 8] = a1; + } + int step, tq, hi; + frag_of(8 * n8 + 2 * t, step, tq, hi); + const __nv_bfloat162 lo2 = __floats2bfloat162_rn(p[0], p[1]), hi2 = __floats2bfloat162_rn(p[2], p[3]); + uint32_t* f = pf + ((tile * 2 + step) * 32 + g * 4 + tq) * 4; + f[2 * hi] = *reinterpret_cast(&lo2); // reg 0 / 2: row g; 1 / 3: row g + 8 + f[2 * hi + 1] = *reinterpret_cast(&hi2); + } + __syncthreads(); + if (tid < BMQ) { // l: fma(l, alpha, (W0 + W2) + (W1 + W3)) + float l = 0.f; + for (int tile = 0; tile < 4; ++tile) { + const float* w4 = wsum + tile * 64 + tid; + const float sum = __fadd_rn(__fadd_rn(w4[0], w4[32]), __fadd_rn(w4[16], w4[48])); + l = __fmaf_rn(l, rowv[64 + tile * 16 + tid], sum); + } + reinterpret_cast(sm + SM_MISC)[16 + tid] = l; + } + float acc[8][4]; + for (int j = 0; j < 8; ++j) + for (int e = 0; e < 4; ++e) acc[j][e] = 0.f; + for (int tile = 0; tile < 4; ++tile) { // o = mma(P, kv, rn(o * alpha)), 2 k steps a tile + const float al0 = rowv[64 + tile * 16 + g], al1 = rowv[64 + tile * 16 + g + 8]; + for (int j = 0; j < 8; ++j) { + acc[j][0] = __fmul_rn(acc[j][0], al0); + acc[j][1] = __fmul_rn(acc[j][1], al0); + acc[j][2] = __fmul_rn(acc[j][2], al1); + acc[j][3] = __fmul_rn(acc[j][3], al1); + } + for (int step = 0; step < 2; ++step) { + const uint4 af = reinterpret_cast(pf)[(tile * 2 + step) * 32 + lane]; + const int k0 = 32 * tile + klo(step, t), k1 = 32 * tile + khi(step, t); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int n = 64 * j + 8 * warp + g; + mma(acc[j], af, pair_n(sm, k0, k0 + 1, n), pair_n(sm, k1, k1 + 1, n)); + } + } + } + const long long pbase = (static_cast(c) * R + r) * H + hb * BMQ; + for (int j = 0; j < 8; ++j) { + const int n = 64 * j + 8 * warp + 2 * t; + *reinterpret_cast(a.po + (pbase + g) * D + n) = make_float2(acc[j][0], acc[j][1]); + *reinterpret_cast(a.po + (pbase + g + 8) * D + n) = make_float2(acc[j][2], acc[j][3]); + } + if (tid < BMQ) { + a.pm[pbase + tid] = reinterpret_cast(sm + SM_MISC)[tid]; + a.pl[pbase + tid] = reinterpret_cast(sm + SM_MISC)[16 + tid]; + } + int* flag = reinterpret_cast(sm + SM_MISC) + 32; + __threadfence(); + __syncthreads(); + if (tid == 0) { + int need = 0; + for (int cc = 0; cc < NCH; ++cc) need += chunk_works(a, r, cc) ? 1 : 0; + *flag = atomicAdd(a.ticket + r * (H / BMQ) + hb, 1) == need - 1; + } + __syncthreads(); + if (!*flag) return; + __threadfence(); + if (tid == 0) atomicExch(a.ticket + r * (H / BMQ) + hb, 0); // self-cleaning: the next launch / replay + merge(a, sm, r, hb); +} + +template +cudaError_t launch_one(const Args& a, cudaStream_t stream) { + static bool init = false; + if (!init) { + const cudaError_t e = cudaFuncSetAttribute(attn_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM); + if (e != cudaSuccess) return e; + init = true; + } + attn_kernel<<>>(a); + return cudaGetLastError(); +} + +inline cudaError_t dispatch(const Args& a, int kw, bool rows, bool has_hi, cudaStream_t stream) { + if (a.R < 1 || a.H % BMQ || a.ring < WINDOW || (a.ring & (a.ring - 1))) return cudaErrorInvalidValue; + switch ((kw == 4 ? 4 : 0) + (rows ? 2 : 0) + (has_hi ? 1 : 0)) { + case 0: return launch_one<2, false, false>(a, stream); + case 1: return launch_one<2, false, true>(a, stream); + case 2: return launch_one<2, true, false>(a, stream); + case 3: return launch_one<2, true, true>(a, stream); + case 4: return launch_one<4, false, false>(a, stream); + case 5: return launch_one<4, false, true>(a, stream); + case 6: return launch_one<4, true, false>(a, stream); + default: return launch_one<4, true, true>(a, stream); + } +} + +} // namespace dsv41_attn_fm + +#ifndef DSV41_ATTN_NO_TORCH // the compile test builds the kernels alone (nvcc, no torch headers) +#include +#include +#include + +namespace { +template T* tp(const at::Tensor& t) { return t.numel() ? static_cast(t.data_ptr()) : nullptr; } +} + +void dsv41_attn_run_cuda(const at::Tensor& q, const at::Tensor& cv, const at::Tensor& csc, const at::Tensor& tok, + const at::Tensor& cnt, const at::Tensor& sv, const at::Tensor& ssc, const at::Tensor& lo, + const at::Tensor& hi, const at::Tensor& pos, const at::Tensor& sl, const at::Tensor& pt, + const at::Tensor& sink, const at::Tensor& cs, at::Tensor& po, at::Tensor& pm, + at::Tensor& pl, at::Tensor& out, at::Tensor& ticket, int64_t ring, int64_t psh, + int64_t pts, int64_t kw) { + using namespace dsv41_attn; + Args a{}; + a.q = tp(q); + a.has_comp = cv.numel() ? 1 : 0; + a.cv = tp(cv); a.csc = tp(csc); + a.cvs = cv.numel() ? cv.stride(0) : 0; a.css = csc.numel() ? csc.stride(0) : 0; + a.tok = tp(tok); a.ts = tok.numel() ? tok.stride(0) : 0; a.cnt = tp(cnt); + a.sv = tp(sv); a.ssc = tp(ssc); a.svs = sv.stride(0); a.sss = ssc.stride(0); + a.lo = tp(lo); a.hi = tp(hi); + a.pos = pos.data_ptr(); + a.sl = tp(sl); + a.pt = tp(pt); a.pts = pts; a.psh = static_cast(psh); + a.sink = tp(sink); a.cs = tp(cs); a.csst = cs.stride(0); + a.po = tp(po); a.pm = tp(pm); a.pl = tp(pl); + a.out = tp(out); a.ticket = tp(ticket); + a.R = static_cast(q.size(0)); a.H = static_cast(q.size(1)); a.ring = static_cast(ring); + C10_CUDA_CHECK(dispatch(a, static_cast(kw), sl.numel() > 0, hi.numel() > 0, + at::cuda::getCurrentCUDAStream())); +} + +std::vector dsv41_attn_info_cuda(int64_t device) { // the device limits of the two launches + int sms = 0, optin = 0, cl = 0, blocks = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, static_cast(device)); + cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, static_cast(device)); + cudaDeviceGetAttribute(&cl, cudaDevAttrClusterLaunch, static_cast(device)); + cudaFuncSetAttribute(dsv41_attn_fm::attn_kernel<4, true, false>, cudaFuncAttributeMaxDynamicSharedMemorySize, + dsv41_attn_fm::SMEM); + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, dsv41_attn_fm::attn_kernel<4, true, false>, dsv41_attn_fm::NT, + dsv41_attn_fm::SMEM); + return {sms, optin, cl, blocks, dsv41_attn_fm::SMEM}; +} +#endif + +// The instantiations dispatch() launches: attn_kernel. +template __global__ void dsv41_attn_fm::attn_kernel<2, false, false>(dsv41_attn_fm::Args); +template __global__ void dsv41_attn_fm::attn_kernel<2, false, true>(dsv41_attn_fm::Args); +template __global__ void dsv41_attn_fm::attn_kernel<2, true, false>(dsv41_attn_fm::Args); +template __global__ void dsv41_attn_fm::attn_kernel<2, true, true>(dsv41_attn_fm::Args); +template __global__ void dsv41_attn_fm::attn_kernel<4, false, false>(dsv41_attn_fm::Args); +template __global__ void dsv41_attn_fm::attn_kernel<4, false, true>(dsv41_attn_fm::Args); +template __global__ void dsv41_attn_fm::attn_kernel<4, true, false>(dsv41_attn_fm::Args); +template __global__ void dsv41_attn_fm::attn_kernel<4, true, true>(dsv41_attn_fm::Args); diff --git a/zig/kernels/cuda/deepseek_v41/proto/attn_phases.cuh b/zig/kernels/cuda/deepseek_v41/proto/attn_phases.cuh new file mode 100644 index 0000000000..c205a8a4d6 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/proto/attn_phases.cuh @@ -0,0 +1,489 @@ +// Timing-only copy of attn_cuda.cu with globaltimer stamps STAMP(0..7) at phase boundaries per CTA. +#ifndef DSV41_ATTN_NO_TORCH +#define DSV41_ATTN_NO_TORCH +#endif +// Copy of src/tensorfold/families/deepseek_v41/cuda/csa2/attn_cuda.cu (blob 0f3234101360): sass_check.sh + + + +// CSA2 decode attention core in ONE launch: CTA (chunk of 5, 16-head tile, row), 8 warps, ops = Triton's sm_121 PTX +#include +#include +#include +#include + +namespace dsv41_attn_ph { +__device__ unsigned long long g_stamp[256][10]; +__device__ __forceinline__ unsigned long long gtime() { unsigned long long t; asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); return t; } +#define STAMP(k) do { __syncthreads(); if (threadIdx.x == 0) g_stamp[(blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x][k] = gtime(); } while (0) + +constexpr int D = 512, CH = 128, NCOMP = 4, NCH = 5, BMQ = 16, WINDOW = 128; // tiles of 32 keys +constexpr int NOPE = 448, VB = 576, SB = 8, NT = 256; +constexpr int SM_Q = 0; // q A fragments [32 steps][32 lanes] uint4 16384 +constexpr int SM_KV = 16384; // 128 rows x 576 B, 8-byte pieces swizzled by key & 7 73728 +constexpr int SM_SC = SM_KV + CH * VB; // 128 x 8 scale bytes 1024 +constexpr int SM_P = SM_SC + CH * SB; // P A fragments [4 tiles][2 steps][32 lanes] uint4 4096 +constexpr int SM_TMAX = SM_P + 4096; // [4 tiles][4 n8][16 rows] fp32 1024 +constexpr int SM_WSUM = SM_TMAX + 1024; // same shape: the warps' row sums 1024 +constexpr int SM_ROW = SM_WSUM + 1024; // [4 tiles][16 rows] nm, alpha, active 768 +constexpr int SM_PHYS = SM_ROW + 768; // 128 physical rows (int32, -1: masked) 512 +constexpr int SM_MISC = SM_PHYS + 512; // final m, l [16] each + the ticket flag 256 +constexpr int SMEM = SM_MISC + 256; // 98,816 B (sm_121: <= 101,376 a block) +static_assert(SMEM <= 101376, "shared memory"); +constexpr unsigned FULL = 0xffffffffu; + +struct Args { + const uint16_t* q; // [R, H, 512] bf16 (RoPE'd) + const uint8_t* cv; const uint8_t* csc; long long cvs, css; // compressed rows (values, scales), row strides + const int* tok; long long ts; const int* cnt; // selection [R, ts] int32 (-1 pad), counts [R] + const uint8_t* sv; const uint8_t* ssc; long long svs, sss; // SWA rings + const int* lo; const int* hi; // [R] window starts; window ends (HAS_HI) or null + const void* pos; // int32 [1] (POS + r) or int64 [R] (ROWS) + const long long* sl; // ROWS: the rows' slots [R] + const int* pt; long long pts; int psh; // page table of the compressed rows (psh 0: contiguous) + const float* sink; const float* cs; long long csst; // [H] fp32, RoPE table fp32 [max_pos, 64] + float* po; float* pm; float* pl; // partials (attn.Scratch layout) + uint16_t* out; // [R, H, 512] bf16 + int* ticket; // [R, H / 16] (self-cleaning) + int R, H, ring, has_comp; +}; + +__device__ __forceinline__ float ex2(float x) { + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float div_full(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float texp(float x) { return ex2(__fmul_rn(x, __int_as_float(0x3FB8AA3B))); } +__device__ __forceinline__ float pow2b(uint32_t s) { return __int_as_float(static_cast(s) << 23); } + +__device__ __forceinline__ void mma(float (&d)[4], const uint4& a, uint32_t b0, uint32_t b1) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b0), "r"(b1)); +} + +// two e4m3 bytes (low byte = the lower k) with their own scales -> bf16x2 (exact: e4m3 x 2^e is a bf16) +__device__ __forceinline__ uint32_t deq2(uint32_t two, float s0, float s1) { + uint32_t h2; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(h2) : "h"(static_cast(two))); + const __half2 h = *reinterpret_cast(&h2); + const __nv_bfloat162 b = __floats2bfloat162_rn(__fmul_rn(__low2float(h), s0), __fmul_rn(__high2float(h), s1)); + return *reinterpret_cast(&b); +} + +// a byte offset x of row `key` in the staged rows: 8-byte pieces swizzled within groups of 8 (bank-conflict free) +__device__ __forceinline__ int swz(int key, int x) { + const int ch = x >> 3; + return key * VB + ((((ch & ~7) | ((ch & 7) ^ (key & 7)))) << 3) + (x & 7); +} + +// Triton's dot-operand k order: step s (16 k) of a 32-k block takes k = klo(s), klo + 1 and khi(s), khi + 1 +template __device__ __forceinline__ int klo(int s, int t) { + return KW == 4 ? 32 * (s >> 1) + 4 * t + 2 * (s & 1) : 16 * s + 2 * t; +} +template __device__ __forceinline__ int khi(int s, int t) { return klo(s, t) + (KW == 4 ? 16 : 8); } +// where element k (< 32) of a row sits in a 2-step A fragment: step, quad lane, high half (+16 / +8) +template __device__ __forceinline__ void frag_of(int k, int& step, int& tq, int& hi) { + if (KW == 4) { step = (k >> 1) & 1; hi = (k >> 4) & 1; tq = (k & 15) >> 2; } + else { step = k >> 4; hi = (k >> 3) & 1; tq = (k & 7) >> 1; } +} + +// the bf16 pair (k, k + 1) of staged row `key` (S's B operand: k along the row) +__device__ __forceinline__ uint32_t pair_k(const uint8_t* sm, int key, int k) { + if (k < NOPE) { + const uint32_t two = *reinterpret_cast(sm + SM_KV + swz(key, k)); + const float s = pow2b(sm[SM_SC + key * SB + (k >> 6)]); + return deq2(two, s, s); + } + return *reinterpret_cast(sm + SM_KV + swz(key, NOPE + 2 * (k - NOPE))); +} + +// value (key, n) of the staged rows as a bf16 bit pattern / its e4m3 byte (PV's B operand: k along the keys) +__device__ __forceinline__ uint32_t pair_n(const uint8_t* sm, int k0, int k1, int n) { + if (n < NOPE) { + const uint32_t two = sm[SM_KV + swz(k0, n)] | (static_cast(sm[SM_KV + swz(k1, n)]) << 8); + return deq2(two, pow2b(sm[SM_SC + k0 * SB + (n >> 6)]), pow2b(sm[SM_SC + k1 * SB + (n >> 6)])); + } + const int x = NOPE + 2 * (n - NOPE); + return *reinterpret_cast(sm + SM_KV + swz(k0, x)) | + (static_cast(*reinterpret_cast(sm + SM_KV + swz(k1, x))) << 16); +} + +__device__ __forceinline__ void cp8(uint32_t dst, const void* src, bool ok) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8, %2;" :: "r"(dst), "l"(src), "r"(ok ? 8 : 0)); +} + +template __device__ __forceinline__ int row_pos(const Args& a, int r) { + return ROWS ? static_cast(static_cast(a.pos)[r]) : static_cast(a.pos)[0] + r; +} + +// how many chunks of row r have keys (the merge's ticket target; the others are Triton's empty partial) +__device__ __forceinline__ bool chunk_works(const Args& a, int r, int c) { + return c == NCOMP || (a.has_comp && c * CH < a.cnt[r]); +} + +template +__device__ void merge(const Args& a, uint8_t* sm, int r, int hb) { + const int tid = threadIdx.x, H = a.H, R = a.R; + float* ca = reinterpret_cast(sm + SM_TMAX); // [5][16] a, [5][16] b, [16] af, [16] den + float* cb = ca + NCH * BMQ; + float* fa = cb + NCH * BMQ; + float* fd = fa + BMQ; + if (tid < BMQ) { // a head's scalar chain (the merge, unrolled) + const int h = hb * BMQ + tid; + float m = -INFINITY, l = 0.f; + for (int c = 0; c < NCH; ++c) { + const bool w = chunk_works(a, r, c); + const long long base = (static_cast(c) * R + r) * H + h; + const float cm = w ? __ldcg(a.pm + base) : -INFINITY, cl = w ? __ldcg(a.pl + base) : 0.f; + const bool act = cl > 0.f; + float aa, bb; + if (c == 0) { + const float nm = act ? fmaxf(cm, -INFINITY) : -INFINITY; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + aa = 0.f; + l = __fmaf_rn(bb, cl, 0.f); + m = nm; + } else { + const float nm = act ? fmaxf(m, cm) : m; + aa = act ? (m == -INFINITY ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + l = __fmaf_rn(l, aa, __fmul_rn(bb, cl)); + m = nm; + } + ca[c * BMQ + tid] = aa; + cb[c * BMQ + tid] = bb; + } + const float sink = a.sink[h]; + const float top = fmaxf(m, sink); + const float af = m == -INFINITY ? 0.f : texp(__fsub_rn(m, top)); + const float e = sink == -INFINITY ? 0.f : texp(__fsub_rn(sink, top)); + fa[tid] = af; + fd[tid] = __fmaf_rn(l, af, e); + } + __syncthreads(); + const int at = row_pos(a, r); + const float* csr = a.cs + static_cast(at) * a.csst; + for (int u = tid; u < BMQ * (D / 4); u += NT) { // float4 units: head hl, columns d .. d + 3 + const int hl = u / (D / 4), d = (u % (D / 4)) * 4, h = hb * BMQ + hl; + float o[4]; + for (int c = 0; c < NCH; ++c) { + float4 v = make_float4(0.f, 0.f, 0.f, 0.f); + if (chunk_works(a, r, c)) + v = __ldcg(reinterpret_cast(a.po + ((static_cast(c) * R + r) * H + h) * D + + d)); + const float co[4] = {v.x, v.y, v.z, v.w}; + const float aa = ca[c * BMQ + hl], bb = cb[c * BMQ + hl]; + for (int i = 0; i < 4; ++i) + o[i] = c == 0 ? __fmaf_rn(bb, co[i], 0.f) : __fmaf_rn(o[i], aa, __fmul_rn(co[i], bb)); + } + float x[4]; + for (int i = 0; i < 4; ++i) x[i] = div_full(__fmul_rn(o[i], fa[hl]), fd[hl]); + uint16_t ob[4]; + for (int i = 0; i < 4; i += 2) { // pair (d + i, d + i + 1): GPT-J, inverse + const int p = (d + i) >> 1; + const float c = p >= 224 ? csr[p - 224] : 1.f, s = p >= 224 ? csr[32 + p - 224] : 0.f; + const float ne = __fmaf_rn(x[i], c, __fmul_rn(x[i + 1], s)); + const float no = __fmaf_rn(x[i + 1], c, -__fmul_rn(x[i], s)); + const __nv_bfloat16 b0 = __float2bfloat16_rn(ne), b1 = __float2bfloat16_rn(no); + ob[i] = *reinterpret_cast(&b0); + ob[i + 1] = *reinterpret_cast(&b1); + } + uint2 w; + w.x = ob[0] | (static_cast(ob[1]) << 16); + w.y = ob[2] | (static_cast(ob[3]) << 16); + *reinterpret_cast(a.out + (static_cast(r) * H + h) * D + d) = w; + } +} + +template +__global__ void __launch_bounds__(NT, 1) attn_kernel(Args a) { + extern __shared__ __align__(16) uint8_t sm[]; + const int c = blockIdx.x, hb = blockIdx.y, r = blockIdx.z; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, t = lane & 3; + const int H = a.H, R = a.R; + if (!chunk_works(a, r, c)) return; + STAMP(0); + int* phys = reinterpret_cast(sm + SM_PHYS); + const long long sl = ROWS ? a.sl[r] : 0; + if (tid < CH) { // the chunk's 128 keys -> physical rows (-1: masked) + int row = -1; + if (c < NCOMP) { + const int idx = c * CH + tid; + if (idx < a.cnt[r]) { + const int tk = a.tok[static_cast(r) * a.ts + idx]; + if (tk >= 0) { + row = tk; + if (a.psh > 0) { + const int* ptr = a.pt + (ROWS ? sl * a.pts : 0); + row = (ptr[tk >> a.psh] << a.psh) + (tk & ((1 << a.psh) - 1)); + } + } + } + } else { + const int hi = HAS_HI ? a.hi[r] : row_pos(a, r); + const int lo = max(a.lo[r], hi - (WINDOW - 1)); + const int p = hi - (WINDOW - 1) + tid; + if (p >= lo && p >= 0) row = (p % a.ring) + (ROWS ? static_cast(sl) * a.ring : 0); + } + phys[tid] = row; + } + { // q -> A fragments (Triton's k order) + const uint16_t* qb = a.q + (static_cast(r) * H + hb * BMQ) * D; + for (int i = tid; i < 32 * 32; i += NT) { + const int s = i >> 5, ln = i & 31, gg = ln >> 2, tt = ln & 3; + const int k0 = klo(s, tt), k1 = khi(s, tt); + uint4 f; + f.x = *reinterpret_cast(qb + gg * D + k0); + f.y = *reinterpret_cast(qb + (gg + 8) * D + k0); + f.z = *reinterpret_cast(qb + gg * D + k1); + f.w = *reinterpret_cast(qb + (gg + 8) * D + k1); + reinterpret_cast(sm + SM_Q)[i] = f; + } + } + STAMP(1); + { // the rows: 72 value pieces + 1 scale piece a key + const bool comp = c < NCOMP; + const uint8_t* V = comp ? a.cv : a.sv; + const uint8_t* S = comp ? a.csc : a.ssc; + const long long vs = comp ? a.cvs : a.svs, ss = comp ? a.css : a.sss; + const uint32_t base = static_cast(__cvta_generic_to_shared(sm)); + for (int i = tid; i < CH * 73; i += NT) { + const int key = i / 73, pc = i % 73, row = phys[key]; + const bool ok = row >= 0; + const long long rr = ok ? row : 0; + if (pc < 72) cp8(base + SM_KV + swz(key, pc * 8), V + rr * vs + pc * 8, ok); + else cp8(base + SM_SC + key * SB, S + rr * ss, ok); + } + asm volatile("cp.async.commit_group;\ncp.async.wait_group 0;" ::: "memory"); + } + STAMP(2); + const float NINF = -INFINITY; + const int n8 = warp & 3, ta = warp >> 2; // this warp: S of tiles ta and ta + 2, block n8 + float sv[2][4] = {{0.f, 0.f, 0.f, 0.f}, {0.f, 0.f, 0.f, 0.f}}; + { + const int kA = 32 * ta + 8 * n8 + g, kB = kA + 64; + const uint4* qf = reinterpret_cast(sm + SM_Q); +#pragma unroll 4 + for (int s = 0; s < 32; ++s) { + const uint4 af = qf[s * 32 + lane]; + const int k0 = klo(s, t), k1 = khi(s, t); + mma(sv[0], af, pair_k(sm, kA, k0), pair_k(sm, kA, k1)); + mma(sv[1], af, pair_k(sm, kB, k0), pair_k(sm, kB, k1)); + } + } + STAMP(3); + float* tmax = reinterpret_cast(sm + SM_TMAX); + float* wsum = reinterpret_cast(sm + SM_WSUM); + float* rowv = reinterpret_cast(sm + SM_ROW); // [tile][row]: nm, alpha, active at 0 / 64 / 128 + bool val[2][2]; + for (int j = 0; j < 2; ++j) { + const int tile = ta + 2 * j, kc = 32 * tile + 8 * n8 + 2 * t; + val[j][0] = phys[kc] >= 0; + val[j][1] = phys[kc + 1] >= 0; + for (int e = 0; e < 4; ++e) + sv[j][e] = val[j][e & 1] ? __fmul_rn(sv[j][e], __int_as_float(0x3D3504F3)) : NINF; + float m0 = fmaxf(sv[j][0], sv[j][1]), m1 = fmaxf(sv[j][2], sv[j][3]); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 2)); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 1)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 2)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 1)); + if (t == 0) { + tmax[(tile * 4 + n8) * 16 + g] = m0; + tmax[(tile * 4 + n8) * 16 + g + 8] = m1; + } + } + __syncthreads(); + if (tid < BMQ) { // the running max over the tiles, a row a thread + float m = NINF; + for (int tile = 0; tile < 4; ++tile) { + const float* tm4 = tmax + tile * 64 + tid; + const float tm = fmaxf(fmaxf(tm4[0], tm4[16]), fmaxf(tm4[32], tm4[48])); + const bool act = tm != NINF; + const float nm = act ? fmaxf(m, tm) : m; + const float al = act ? (m == NINF ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + rowv[tile * 16 + tid] = nm; + rowv[64 + tile * 16 + tid] = al; + rowv[128 + tile * 16 + tid] = act ? 1.f : 0.f; + m = nm; + } + reinterpret_cast(sm + SM_MISC)[tid] = m; + } + __syncthreads(); + uint32_t* pf = reinterpret_cast(sm + SM_P); // [tile][step][lane][4] + for (int j = 0; j < 2; ++j) { // p, the thread / quad sums, bf16(p) -> fragments + const int tile = ta + 2 * j; + float p[4]; + for (int e = 0; e < 4; ++e) { + const int row = g + (e >> 1) * 8; + const bool act = rowv[128 + tile * 16 + row] != 0.f; + p[e] = (val[j][e & 1] && act) ? texp(__fsub_rn(sv[j][e], rowv[tile * 16 + row])) : 0.f; + } + float a0 = __fadd_rn(p[0], p[1]), a1 = __fadd_rn(p[2], p[3]); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 2)); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 1)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 2)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 1)); + if (t == 0) { + wsum[(tile * 4 + n8) * 16 + g] = a0; + wsum[(tile * 4 + n8) * 16 + g + 8] = a1; + } + int step, tq, hi; + frag_of(8 * n8 + 2 * t, step, tq, hi); + const __nv_bfloat162 lo2 = __floats2bfloat162_rn(p[0], p[1]), hi2 = __floats2bfloat162_rn(p[2], p[3]); + uint32_t* f = pf + ((tile * 2 + step) * 32 + g * 4 + tq) * 4; + f[2 * hi] = *reinterpret_cast(&lo2); // reg 0 / 2: row g; 1 / 3: row g + 8 + f[2 * hi + 1] = *reinterpret_cast(&hi2); + } + __syncthreads(); + if (tid < BMQ) { // l: fma(l, alpha, (W0 + W2) + (W1 + W3)) + float l = 0.f; + for (int tile = 0; tile < 4; ++tile) { + const float* w4 = wsum + tile * 64 + tid; + const float sum = __fadd_rn(__fadd_rn(w4[0], w4[32]), __fadd_rn(w4[16], w4[48])); + l = __fmaf_rn(l, rowv[64 + tile * 16 + tid], sum); + } + reinterpret_cast(sm + SM_MISC)[16 + tid] = l; + } + STAMP(4); + float acc[8][4]; + for (int j = 0; j < 8; ++j) + for (int e = 0; e < 4; ++e) acc[j][e] = 0.f; + for (int tile = 0; tile < 4; ++tile) { // o = mma(P, kv, rn(o * alpha)), 2 k steps a tile + const float al0 = rowv[64 + tile * 16 + g], al1 = rowv[64 + tile * 16 + g + 8]; + for (int j = 0; j < 8; ++j) { + acc[j][0] = __fmul_rn(acc[j][0], al0); + acc[j][1] = __fmul_rn(acc[j][1], al0); + acc[j][2] = __fmul_rn(acc[j][2], al1); + acc[j][3] = __fmul_rn(acc[j][3], al1); + } + for (int step = 0; step < 2; ++step) { + const uint4 af = reinterpret_cast(pf)[(tile * 2 + step) * 32 + lane]; + const int k0 = 32 * tile + klo(step, t), k1 = 32 * tile + khi(step, t); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int n = 64 * j + 8 * warp + g; + mma(acc[j], af, pair_n(sm, k0, k0 + 1, n), pair_n(sm, k1, k1 + 1, n)); + } + } + } + STAMP(5); + const long long pbase = (static_cast(c) * R + r) * H + hb * BMQ; + for (int j = 0; j < 8; ++j) { + const int n = 64 * j + 8 * warp + 2 * t; + *reinterpret_cast(a.po + (pbase + g) * D + n) = make_float2(acc[j][0], acc[j][1]); + *reinterpret_cast(a.po + (pbase + g + 8) * D + n) = make_float2(acc[j][2], acc[j][3]); + } + if (tid < BMQ) { + a.pm[pbase + tid] = reinterpret_cast(sm + SM_MISC)[tid]; + a.pl[pbase + tid] = reinterpret_cast(sm + SM_MISC)[16 + tid]; + } + int* flag = reinterpret_cast(sm + SM_MISC) + 32; + __threadfence(); + __syncthreads(); + if (tid == 0) { + int need = 0; + for (int cc = 0; cc < NCH; ++cc) need += chunk_works(a, r, cc) ? 1 : 0; + *flag = atomicAdd(a.ticket + r * (H / BMQ) + hb, 1) == need - 1; + } + __syncthreads(); + STAMP(6); + if (!*flag) return; + __threadfence(); + if (tid == 0) atomicExch(a.ticket + r * (H / BMQ) + hb, 0); // self-cleaning: the next launch / replay + merge(a, sm, r, hb); + STAMP(7); +} + +template +cudaError_t launch_one(const Args& a, cudaStream_t stream) { + static bool init = false; + if (!init) { + const cudaError_t e = cudaFuncSetAttribute(attn_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM); + if (e != cudaSuccess) return e; + init = true; + } + attn_kernel<<>>(a); + return cudaGetLastError(); +} + +inline cudaError_t dispatch(const Args& a, int kw, bool rows, bool has_hi, cudaStream_t stream) { + if (a.R < 1 || a.H % BMQ || a.ring < WINDOW || (a.ring & (a.ring - 1))) return cudaErrorInvalidValue; + switch ((kw == 4 ? 4 : 0) + (rows ? 2 : 0) + (has_hi ? 1 : 0)) { + case 0: return launch_one<2, false, false>(a, stream); + case 1: return launch_one<2, false, true>(a, stream); + case 2: return launch_one<2, true, false>(a, stream); + case 3: return launch_one<2, true, true>(a, stream); + case 4: return launch_one<4, false, false>(a, stream); + case 5: return launch_one<4, false, true>(a, stream); + case 6: return launch_one<4, true, false>(a, stream); + default: return launch_one<4, true, true>(a, stream); + } +} + +} // namespace dsv41_attn_ph + +#ifndef DSV41_ATTN_NO_TORCH // the compile test builds the kernels alone (nvcc, no torch headers) +#include +#include +#include + +namespace { +template T* tp(const at::Tensor& t) { return t.numel() ? static_cast(t.data_ptr()) : nullptr; } +} + +void dsv41_attn_run_cuda(const at::Tensor& q, const at::Tensor& cv, const at::Tensor& csc, const at::Tensor& tok, + const at::Tensor& cnt, const at::Tensor& sv, const at::Tensor& ssc, const at::Tensor& lo, + const at::Tensor& hi, const at::Tensor& pos, const at::Tensor& sl, const at::Tensor& pt, + const at::Tensor& sink, const at::Tensor& cs, at::Tensor& po, at::Tensor& pm, + at::Tensor& pl, at::Tensor& out, at::Tensor& ticket, int64_t ring, int64_t psh, + int64_t pts, int64_t kw) { + using namespace dsv41_attn; + Args a{}; + a.q = tp(q); + a.has_comp = cv.numel() ? 1 : 0; + a.cv = tp(cv); a.csc = tp(csc); + a.cvs = cv.numel() ? cv.stride(0) : 0; a.css = csc.numel() ? csc.stride(0) : 0; + a.tok = tp(tok); a.ts = tok.numel() ? tok.stride(0) : 0; a.cnt = tp(cnt); + a.sv = tp(sv); a.ssc = tp(ssc); a.svs = sv.stride(0); a.sss = ssc.stride(0); + a.lo = tp(lo); a.hi = tp(hi); + a.pos = pos.data_ptr(); + a.sl = tp(sl); + a.pt = tp(pt); a.pts = pts; a.psh = static_cast(psh); + a.sink = tp(sink); a.cs = tp(cs); a.csst = cs.stride(0); + a.po = tp(po); a.pm = tp(pm); a.pl = tp(pl); + a.out = tp(out); a.ticket = tp(ticket); + a.R = static_cast(q.size(0)); a.H = static_cast(q.size(1)); a.ring = static_cast(ring); + C10_CUDA_CHECK(dispatch(a, static_cast(kw), sl.numel() > 0, hi.numel() > 0, + at::cuda::getCurrentCUDAStream())); +} + +std::vector dsv41_attn_info_cuda(int64_t device) { // the device limits of the two launches + int sms = 0, optin = 0, cl = 0, blocks = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, static_cast(device)); + cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, static_cast(device)); + cudaDeviceGetAttribute(&cl, cudaDevAttrClusterLaunch, static_cast(device)); + cudaFuncSetAttribute(dsv41_attn_ph::attn_kernel<4, true, false>, cudaFuncAttributeMaxDynamicSharedMemorySize, + dsv41_attn_ph::SMEM); + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, dsv41_attn_ph::attn_kernel<4, true, false>, dsv41_attn_ph::NT, + dsv41_attn_ph::SMEM); + return {sms, optin, cl, blocks, dsv41_attn_ph::SMEM}; +} +#endif + +// The instantiations dispatch() launches: attn_kernel. +template __global__ void dsv41_attn_ph::attn_kernel<2, false, false>(dsv41_attn_ph::Args); +template __global__ void dsv41_attn_ph::attn_kernel<2, false, true>(dsv41_attn_ph::Args); +template __global__ void dsv41_attn_ph::attn_kernel<2, true, false>(dsv41_attn_ph::Args); +template __global__ void dsv41_attn_ph::attn_kernel<2, true, true>(dsv41_attn_ph::Args); +template __global__ void dsv41_attn_ph::attn_kernel<4, false, false>(dsv41_attn_ph::Args); +template __global__ void dsv41_attn_ph::attn_kernel<4, false, true>(dsv41_attn_ph::Args); +template __global__ void dsv41_attn_ph::attn_kernel<4, true, false>(dsv41_attn_ph::Args); +template __global__ void dsv41_attn_ph::attn_kernel<4, true, true>(dsv41_attn_ph::Args); diff --git a/zig/kernels/cuda/deepseek_v41/proto/gb10emu.cuh b/zig/kernels/cuda/deepseek_v41/proto/gb10emu.cuh new file mode 100644 index 0000000000..59c05b20bd --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/proto/gb10emu.cuh @@ -0,0 +1,132 @@ +// Emulates a 48-SM, bandwidth-limited GB10-class part on a bigger sm_120 GPU: blocker CTAs park or hog the other SMs. + +// Call gb10::eager() before any CUDA call: with lazy loading a first launch waits on the parked blockers (deadlock). +#pragma once +#include +#include +#include +#include +#include + +#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \ + fprintf(stderr, "%s:%d %s: %s\n", __FILE__, __LINE__, #x, cudaGetErrorString(e_)); exit(1); } } while (0) + +namespace gb10 { + +constexpr int FREE = 48; // the GB10's SMs +constexpr int BLOCK_SMEM = 99 * 1024 - 64; // one blocker an SM (sm_120: 99 KB a CTA with the static flag) + +__device__ __forceinline__ unsigned smid() { unsigned s; asm volatile("mov.u32 %0, %%smid;" : "=r"(s)); return s; } + +// hogs: the first `hogs` blocked SMs stream `buf` (16-byte loads, 512 threads); the others spin +__global__ void __launch_bounds__(512) blocker(volatile int* release, int* ready, const int4* __restrict__ buf, + size_t n16, int hogs, unsigned long long* sink) { + extern __shared__ int4 pad[]; + const unsigned s = smid(); + if (s < FREE) return; + if (threadIdx.x == 0) atomicAdd(ready, 1); + const bool hog = (int)(s - FREE) < hogs; + int4 acc = make_int4(0, 0, 0, 0); + size_t i = ((size_t)(s - FREE) * 512 + threadIdx.x) * 8; + const size_t stride = (size_t)hogs * 512 * 8; + __shared__ int stop; + for (;;) { // one thread polls the host flag, every 64 sweeps + if (threadIdx.x == 0) stop = *release; + __syncthreads(); + if (stop) break; + __syncthreads(); + for (int k = 0; k < 64; ++k) { + if (hog) { +#pragma unroll + for (int u = 0; u < 8; ++u) { + const int4 v = __ldcs(buf + ((i + u) & (n16 - 1))); // n16: a power of two + acc.x ^= v.x; acc.y ^= v.y; acc.z ^= v.z; acc.w ^= v.w; + } + i = (i + stride) & (n16 - 1); + } else { + __nanosleep(2000); + } + } + } + if (acc.x == 0x7fffffff) sink[0] = (unsigned long long)(acc.y ^ acc.z ^ acc.w ^ pad[0].x); +} + +// a plain stream read of n16 16-byte words on the kernel's SMs: the bandwidth the kernel under test sees +__global__ void __launch_bounds__(512) stream_read(const int4* __restrict__ p, size_t n16, unsigned long long* sink) { + int4 acc = make_int4(0, 0, 0, 0); + for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; i < n16; i += (size_t)gridDim.x * blockDim.x) { + const int4 v = __ldcs(p + i); + acc.x ^= v.x; acc.y ^= v.y; acc.z ^= v.z; acc.w ^= v.w; + } + if (acc.x == 0x7fffffff) sink[1] = (unsigned long long)(acc.y ^ acc.z ^ acc.w); +} + +inline void eager() { setenv("CUDA_MODULE_LOADING", "EAGER", 1); } + +enum class Mode { Full, Sms, Gb10 }; + +// begin() parks the blockers, end() releases them; wrap a timed region: emu.begin(); time(kernel); emu.end(); +struct Emu { + Mode mode = Mode::Full; + int sms = 0, hogs = 0; + cudaStream_t side{}; + int* release_h = nullptr; int* release_d = nullptr; int* ready = nullptr; + int4* buf = nullptr; size_t n16 = 0; unsigned long long* sink = nullptr; + + void init(Mode m, size_t hog_bytes = (size_t)4 << 30) { + mode = m; + cudaDeviceProp p{}; CK(cudaGetDeviceProperties(&p, 0)); sms = p.multiProcessorCount; + CK(cudaStreamCreateWithFlags(&side, cudaStreamNonBlocking)); + CK(cudaHostAlloc(&release_h, sizeof(int), cudaHostAllocMapped)); + CK(cudaHostGetDevicePointer(&release_d, release_h, 0)); + CK(cudaMalloc(&ready, sizeof(int))); + CK(cudaMalloc(&sink, 16)); + if (m == Mode::Gb10) { + n16 = hog_bytes / 16; CK(cudaMalloc(&buf, n16 * 16)); CK(cudaMemset(buf, 1, n16 * 16)); + } + CK(cudaFuncSetAttribute(blocker, cudaFuncAttributeMaxDynamicSharedMemorySize, BLOCK_SMEM)); + } + int grid_sms() const { return mode == Mode::Full ? sms : FREE; } // what the kernel under test may size for + void begin() { + if (mode == Mode::Full) return; + *release_h = 0; CK(cudaMemset(ready, 0, sizeof(int))); CK(cudaDeviceSynchronize()); + blocker<<>>(release_h ? release_d : nullptr, ready, buf, n16, + mode == Mode::Gb10 ? hogs : 0, sink); + CK(cudaGetLastError()); + int r = 0; + for (int spin = 0; r < sms - FREE && spin < 200000; ++spin) CK(cudaMemcpy(&r, ready, 4, cudaMemcpyDeviceToHost)); + if (r < sms - FREE) { fprintf(stderr, "gb10emu: only %d of %d blockers resident\n", r, sms - FREE); exit(1); } + } + void end() { + if (mode == Mode::Full) return; + *release_h = 1; __sync_synchronize(); CK(cudaDeviceSynchronize()); + } + // GB/s of a plain 16-byte stream read of `bytes` on the kernel's SMs (inside begin / end) + double stream_gbps(const void* p, size_t bytes, cudaStream_t st) { + cudaEvent_t a, b; CK(cudaEventCreate(&a)); CK(cudaEventCreate(&b)); + const int grid = grid_sms() * 4; + stream_read<<>>((const int4*)p, bytes / 16, sink); + CK(cudaEventRecord(a, st)); + stream_read<<>>((const int4*)p, bytes / 16, sink); + CK(cudaEventRecord(b, st)); CK(cudaEventSynchronize(b)); + float ms = 0; CK(cudaEventElapsedTime(&ms, a, b)); + CK(cudaEventDestroy(a)); CK(cudaEventDestroy(b)); + return bytes / (ms * 1e6); + } + // Gb10 mode: the hog count that leaves the kernel's SMs closest to `target` GB/s of stream read + void calibrate(const void* p, size_t bytes, cudaStream_t st, double target) { + if (mode != Mode::Gb10) return; + int best = 0; double bestd = 1e30; + for (int h = 0; h <= sms - FREE; h += 4) { + hogs = h; begin(); const double g = stream_gbps(p, bytes, st); end(); + printf(" calib hogs %3d -> stream %.1f GB/s\n", h, g); + if (fabs(g - target) < bestd) { bestd = fabs(g - target); best = h; } + if (g < target * 0.9) break; + } + hogs = best; + begin(); const double g = stream_gbps(p, bytes, st); end(); + printf("gb10emu: %d hog SMs, the 48 SMs stream %.1f GB/s (target %.0f)\n", hogs, g, target); + } +}; + +} // namespace gb10 diff --git a/zig/kernels/cuda/deepseek_v41/proto/glue_bench.cu b/zig/kernels/cuda/deepseek_v41/proto/glue_bench.cu new file mode 100644 index 0000000000..28d142de74 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/proto/glue_bench.cu @@ -0,0 +1,186 @@ +// Latency-bound decode kernels (CSA2 attention core, mHC boundary) at 1-16 rows on 48 emulated SMs, cold L2 a launch. +#include "gb10emu.cuh" +#include "../attn_cuda.cu" +#include "../mhc_cuda.cu" +#include "attn_phases.cuh" +#include "attn_fastmerge.cuh" +#include +#include +#include +#include + +namespace { + +__global__ void fill_bytes(uint8_t* p, size_t n, uint32_t seed, uint8_t mask) { + for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (size_t)gridDim.x * blockDim.x) { + uint32_t x = (uint32_t)i * 0x9E3779B1u ^ seed; + x ^= x >> 16; x *= 0x7feb352du; x ^= x >> 15; + p[i] = (uint8_t)(x & mask); + } +} +__global__ void fill_f32(float* p, size_t n, uint32_t seed, float lo, float hi) { + for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (size_t)gridDim.x * blockDim.x) { + uint32_t x = (uint32_t)i * 0x9E3779B1u ^ seed; + x ^= x >> 16; x *= 0x7feb352du; x ^= x >> 15; + p[i] = lo + (hi - lo) * (x & 0xffffff) / 16777216.f; + } +} +__global__ void flush(int4* p, size_t n) { + for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (size_t)gridDim.x * blockDim.x) + p[i] = make_int4((int)i, 0, 0, 0); +} + +template T* dalloc(size_t n) { T* p; CK(cudaMalloc(&p, n * sizeof(T) + 256)); CK(cudaMemset(p, 0, n * sizeof(T) + 256)); return p; } +uint8_t* rbytes(size_t n, uint32_t seed, uint8_t mask) { auto* p = dalloc(n); fill_bytes<<<512, 256>>>(p, n, seed, mask); return p; } +float* rf32(size_t n, uint32_t seed, float lo, float hi) { auto* p = dalloc(n); fill_f32<<<512, 256>>>(p, n, seed, lo, hi); return p; } + +// times fn() `reps` times, each after a cold-L2 write; returns the median in us +template +double time_us(gb10::Emu& emu, cudaStream_t st, int reps, int4* fl, size_t fln, F fn) { + cudaEvent_t a, b; CK(cudaEventCreate(&a)); CK(cudaEventCreate(&b)); + std::vector v; + fn(); // warm (first-launch attributes) + CK(cudaStreamSynchronize(st)); + emu.begin(); + for (int i = 0; i < reps; ++i) { + flush<<>>(fl, fln); + CK(cudaEventRecord(a, st)); + fn(); + CK(cudaEventRecord(b, st)); + CK(cudaEventSynchronize(b)); + float ms; CK(cudaEventElapsedTime(&ms, a, b)); + v.push_back(ms * 1e3); + } + emu.end(); + CK(cudaGetLastError()); + std::sort(v.begin(), v.end()); + return v[v.size() / 2]; +} + +} // namespace + +int main(int argc, char** argv) { // glue_bench [--mode full|sms] [--reps 40] + gb10::eager(); + std::string mode = "sms"; int reps = 40; + for (int i = 1; i + 1 < argc; i += 2) { + const std::string k = argv[i]; + if (k == "--mode") mode = argv[i + 1]; else if (k == "--reps") reps = atoi(argv[i + 1]); + } + gb10::Emu emu; emu.init(mode == "full" ? gb10::Mode::Full : gb10::Mode::Sms); + cudaStream_t st; CK(cudaStreamCreateWithFlags(&st, cudaStreamNonBlocking)); + const size_t fln = (256u << 20) / 16; int4* fl = dalloc(fln); + constexpr int RMAX = 16, H = 32, NCOMP_ROWS = 16384, RING = 128, POS = 32768; + using namespace dsv41_attn; + // ---- attention: 512 selected compressed rows + the 128-row window a row, 32 heads (one rank) ---- + uint8_t* cv = rbytes((size_t)NCOMP_ROWS * VB, 1, 0x3b); + uint8_t* csc = dalloc((size_t)NCOMP_ROWS * SB); CK(cudaMemset(csc, 127, (size_t)NCOMP_ROWS * SB)); + uint8_t* sv = rbytes((size_t)RMAX * RING * VB, 2, 0x3b); + uint8_t* ssc = dalloc((size_t)RMAX * RING * SB); CK(cudaMemset(ssc, 127, (size_t)RMAX * RING * SB)); + std::vector tok((size_t)RMAX * 512), cnt(RMAX, 512), lo(RMAX, 0); + for (size_t i = 0; i < tok.size(); ++i) tok[i] = (int)((i * 2654435761u) % NCOMP_ROWS); + int *dtok = dalloc(tok.size()), *dcnt = dalloc(RMAX), *dlo = dalloc(RMAX), *dpos = dalloc(1); + CK(cudaMemcpy(dtok, tok.data(), tok.size() * 4, cudaMemcpyHostToDevice)); + CK(cudaMemcpy(dcnt, cnt.data(), RMAX * 4, cudaMemcpyHostToDevice)); + CK(cudaMemcpy(dlo, lo.data(), RMAX * 4, cudaMemcpyHostToDevice)); + const int pos = POS; CK(cudaMemcpy(dpos, &pos, 4, cudaMemcpyHostToDevice)); + uint8_t* q = rbytes((size_t)RMAX * H * D * 2, 3, 0x3b); + float* cs = rf32((size_t)(POS + RMAX + 8) * 64, 4, -1.f, 1.f); + float* sink = rf32(H, 5, -1.f, 1.f); + float* po = dalloc((size_t)NCH * RMAX * H * D); + float *pm = dalloc((size_t)NCH * RMAX * H), *pl = dalloc((size_t)NCH * RMAX * H); + uint16_t* out = dalloc((size_t)RMAX * H * D); + int* ticket = dalloc(RMAX * H / BMQ); + // ---- mHC boundary (mode 0: post of 2 ranks' partials, collapse, 24 mixes, finish, norm) ---- + constexpr int DM = dsv41_mhc::D, NBM = dsv41_mhc::NB; + uint8_t* x = rbytes((size_t)RMAX * 4 * DM * 2, 6, 0x3b); + uint16_t* xout = dalloc((size_t)RMAX * 4 * DM); + uint8_t* gp = rbytes((size_t)2 * RMAX * DM * 2, 7, 0x3b); + float *post = rf32(RMAX * 4, 8, 0.1f, 1.f), *comb = rf32(RMAX * 16, 9, 0.f, 0.5f), *pre = rf32(RMAX * 4, 10, 0.f, 1.f); + uint8_t* fn = rbytes((size_t)24 * 4 * DM * 2, 11, 0x3b); + float *base = rf32(24, 12, -0.1f, 0.1f), *scale = rf32(3, 13, 0.5f, 1.f); + float* part = dalloc((size_t)RMAX * 4 * NBM * 32); + uint16_t *cc = dalloc((size_t)RMAX * DM), *mout = dalloc((size_t)RMAX * DM); + uint8_t* nw = rbytes((size_t)DM * 2, 14, 0x3b); + float *opre = dalloc(RMAX * 4), *opost = dalloc(RMAX * 4), *ocomb = dalloc(RMAX * 16); + int* mcnt = dalloc(1); + CK(cudaDeviceSynchronize()); + + // a CTA's time against its key count: chunk 0 holds `kc` selected keys (chunks 1-3 exit at once), the window kc + if (getenv("ATTN_KEYS")) { + for (int kc : {128, 64, 32, 16}) { + std::vector c1(RMAX, kc), l1(RMAX); + for (int r = 0; r < RMAX; ++r) l1[r] = POS + r - (kc - 1); + CK(cudaMemcpy(dcnt, c1.data(), RMAX * 4, cudaMemcpyHostToDevice)); + CK(cudaMemcpy(dlo, l1.data(), RMAX * 4, cudaMemcpyHostToDevice)); + Args a{}; + a.q = reinterpret_cast(q); a.cv = cv; a.csc = csc; a.cvs = VB; a.css = SB; + a.tok = dtok; a.ts = 512; a.cnt = dcnt; a.sv = sv; a.ssc = ssc; a.svs = VB; a.sss = SB; + a.lo = dlo; a.pos = dpos; a.sink = sink; a.cs = cs; a.csst = 64; a.po = po; a.pm = pm; a.pl = pl; + a.out = out; a.ticket = ticket; a.R = 1; a.H = H; a.ring = RING; a.has_comp = 1; + const double t = time_us(emu, st, reps, fl, fln, [&] { CK(dispatch(a, 4, false, false, st)); }); + printf("KEYS %3d a chunk (2 chunks x 2 head tiles working): %.1f us\n", kc, t); + } + CK(cudaMemcpy(dcnt, cnt.data(), RMAX * 4, cudaMemcpyHostToDevice)); + CK(cudaMemcpy(dlo, lo.data(), RMAX * 4, cudaMemcpyHostToDevice)); + } + if (getenv("ATTN_PHASES")) { // per-phase times of the R = 1 launch's CTAs + dsv41_attn_ph::Args a{}; + a.q = reinterpret_cast(q); a.cv = cv; a.csc = csc; a.cvs = VB; a.css = SB; + a.tok = dtok; a.ts = 512; a.cnt = dcnt; a.sv = sv; a.ssc = ssc; a.svs = VB; a.sss = SB; + a.lo = dlo; a.pos = dpos; a.sink = sink; a.cs = cs; a.csst = 64; a.po = po; a.pm = pm; a.pl = pl; + a.out = out; a.ticket = ticket; a.R = 1; a.H = H; a.ring = RING; a.has_comp = 1; + const double t = time_us(emu, st, reps, fl, fln, [&] { CK(dsv41_attn_ph::dispatch(a, 4, false, false, st)); }); + unsigned long long hs[256][10]; + CK(cudaMemcpyFromSymbol(hs, dsv41_attn_ph::g_stamp, sizeof(hs))); + unsigned long long t0 = ~0ull; + for (int i = 0; i < 10; ++i) t0 = std::min(t0, hs[i][0]); + printf("PHASES launch %.1f us; per CTA (us from the first CTA's start): start rows+q keys S softmax PV ticket merge\n", t); + for (int i = 0; i < 10; ++i) { + printf("PHASES cta %d:", i); + for (int k = 0; k < 8; ++k) printf(" %6.2f", hs[i][k] ? (hs[i][k] - t0) / 1e3 : -1.0); + printf("\n"); + } + } + { // fast merge: same bits, time at 1-16 rows + std::vector ref((size_t)RMAX * H * D), got((size_t)RMAX * H * D); + for (int R : {1, 4, 16}) { + Args a{}; + a.q = reinterpret_cast(q); a.cv = cv; a.csc = csc; a.cvs = VB; a.css = SB; + a.tok = dtok; a.ts = 512; a.cnt = dcnt; a.sv = sv; a.ssc = ssc; a.svs = VB; a.sss = SB; + a.lo = dlo; a.pos = dpos; a.sink = sink; a.cs = cs; a.csst = 64; a.po = po; a.pm = pm; a.pl = pl; + a.out = out; a.ticket = ticket; a.R = R; a.H = H; a.ring = RING; a.has_comp = 1; + CK(cudaMemset(out, 0, (size_t)RMAX * H * D * 2)); + const double t0 = time_us(emu, st, reps, fl, fln, [&] { CK(dispatch(a, 4, false, false, st)); }); + CK(cudaMemcpy(ref.data(), out, ref.size() * 2, cudaMemcpyDeviceToHost)); + dsv41_attn_fm::Args b{}; + static_assert(sizeof(b) == sizeof(a), "same Args"); + memcpy(&b, &a, sizeof(a)); + CK(cudaMemset(out, 0, (size_t)RMAX * H * D * 2)); + const double t1 = time_us(emu, st, reps, fl, fln, [&] { CK(dsv41_attn_fm::dispatch(b, 4, false, false, st)); }); + CK(cudaMemcpy(got.data(), out, got.size() * 2, cudaMemcpyDeviceToHost)); + size_t diff = 0, nz = 0; + for (size_t i = 0; i < (size_t)R * H * D; ++i) { diff += ref[i] != got[i]; nz += ref[i] != 0; } + printf("FASTMERGE R %2d: attn_cuda %.1f us, fast merge %.1f us, %zu of %zu outputs differ (%zu nonzero)\n", + R, t0, t1, diff, (size_t)R * H * D, nz); + } + } + printf("mode %s: per launch, median of %d cold launches (us)\n rows | attn_kernel | mhc boundary\n", mode.c_str(), reps); + for (int R : {1, 2, 4, 8, 16}) { + Args a{}; + a.q = reinterpret_cast(q); a.cv = cv; a.csc = csc; a.cvs = VB; a.css = SB; + a.tok = dtok; a.ts = 512; a.cnt = dcnt; a.sv = sv; a.ssc = ssc; a.svs = VB; a.sss = SB; + a.lo = dlo; a.hi = nullptr; a.pos = dpos; a.sl = nullptr; a.pt = nullptr; a.pts = 0; a.psh = 0; + a.sink = sink; a.cs = cs; a.csst = 64; a.po = po; a.pm = pm; a.pl = pl; a.out = out; a.ticket = ticket; + a.R = R; a.H = H; a.ring = RING; a.has_comp = 1; + const double ta = time_us(emu, st, reps, fl, fln, [&] { CK(dispatch(a, 4, false, false, st)); }); + dsv41_mhc::Args m{}; + m.x = reinterpret_cast(x); m.xs = 4 * DM; m.xout = xout; m.g = gp; m.gr = (long long)R * DM; + m.world = 2; m.gbf16 = 1; m.post = post; m.comb = comb; m.pre = pre; m.fn = fn; m.base = base; m.scale = scale; + m.part = part; m.c = cc; m.tap = nullptr; m.ts = 0; m.nw = nw; m.nwbf16 = 1; m.out = mout; + m.opre = opre; m.opost = opost; m.ocomb = ocomb; m.cnt = mcnt; m.R = R; + m.eps = 1e-20f; m.hc_eps = 1e-6f; m.post_alpha = 2.f; m.iters = 20; + const double tm = time_us(emu, st, reps, fl, fln, [&] { CK(dsv41_mhc::dispatch(m, 0, false, st)); }); + printf("RESULT %4d | %9.1f | %9.1f\n", R, ta, tm); + } + return 0; +} diff --git a/zig/kernels/cuda/deepseek_v41/proto/mma_rate.cu b/zig/kernels/cuda/deepseek_v41/proto/mma_rate.cu new file mode 100644 index 0000000000..69977cd09e --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/proto/mma_rate.cu @@ -0,0 +1,80 @@ +// mma.sync issue rates (FLOP/SM/clock) on sm_120/sm_121, 8 accumulator chains a warp; usage: mma_rate [warps_per_sm] +#include "gb10emu.cuh" +#include + +namespace { + +constexpr int ITERS = 4096, CH = 8; + +template +__global__ void __launch_bounds__(512) mma_loop(float* out, uint32_t seed) { + uint32_t a[4], b[2]; +#pragma unroll + for (int i = 0; i < 4; ++i) a[i] = 0x3c003c00u ^ (seed * (threadIdx.x + i)); + b[0] = a[0] ^ 1u; b[1] = a[1] ^ 3u; + float d[CH][4] = {}; + uint32_t h[CH][2] = {}; + for (int it = 0; it < ITERS; ++it) { +#pragma unroll + for (int c = 0; c < CH; ++c) { + if constexpr (V == 0) + 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[c][0]), "+f"(d[c][1]), "+f"(d[c][2]), "+f"(d[c][3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + else if constexpr (V == 1) + 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[c][0]), "+f"(d[c][1]), "+f"(d[c][2]), "+f"(d[c][3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + else if constexpr (V == 2) + asm volatile("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%0,%1};\n" + : "+r"(h[c][0]), "+r"(h[c][1]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + else if constexpr (V == 3) + asm volatile("mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f32.e4m3.e4m3.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(d[c][0]), "+f"(d[c][1]), "+f"(d[c][2]), "+f"(d[c][3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + else + asm volatile("mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(d[c][0]), "+f"(d[c][1]), "+f"(d[c][2]), "+f"(d[c][3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + } + } + float s = 0; +#pragma unroll + for (int c = 0; c < CH; ++c) s += d[c][0] + d[c][3] + __half2float(__ushort_as_half((unsigned short)h[c][0])); + if (s == 1234.5f) out[threadIdx.x] = s; +} + +template +void run(const char* name, gb10::Emu& emu, int warps, double flop_per_mma, float* out) { + cudaStream_t st; CK(cudaStreamCreateWithFlags(&st, cudaStreamNonBlocking)); + cudaEvent_t a, b; CK(cudaEventCreate(&a)); CK(cudaEventCreate(&b)); + const int grid = emu.grid_sms(); + mma_loop<<>>(out, 7); // warm + emu.begin(); + CK(cudaEventRecord(a, st)); + mma_loop<<>>(out, 7); + CK(cudaEventRecord(b, st)); CK(cudaEventSynchronize(b)); + emu.end(); + float ms; CK(cudaEventElapsedTime(&ms, a, b)); + int clk; CK(cudaDeviceGetAttribute(&clk, cudaDevAttrClockRate, 0)); + const double flop = (double)grid * warps * ITERS * CH * flop_per_mma; + const double per = flop / (ms * 1e-3) / grid / (clk * 1e3); + printf("RATE %-28s %2d warps/SM: %8.1f TFLOP/s on %d SMs, %6.0f FLOP/SM/clk (at %d MHz) -> GB10 48 SMs @2.5 GHz %.1f TFLOP/s\n", + name, warps, flop / (ms * 1e9), grid, per, clk / 1000, per * 48 * 2.5e9 / 1e12); +} + +} // namespace + +int main(int argc, char** argv) { + gb10::eager(); + const int warps = argc > 1 ? atoi(argv[1]) : 8; + gb10::Emu emu; emu.init(gb10::Mode::Sms); + float* out; CK(cudaMalloc(&out, 4096)); + run<0>("f16 x f16 -> f32 k16", emu, warps, 2.0 * 16 * 8 * 16, out); + run<1>("bf16 x bf16 -> f32 k16", emu, warps, 2.0 * 16 * 8 * 16, out); + run<2>("f16 x f16 -> f16 k16", emu, warps, 2.0 * 16 * 8 * 16, out); + run<3>("e4m3 kind::f8f6f4 -> f32 k32", emu, warps, 2.0 * 16 * 8 * 32, out); + run<4>("e4m3 -> f32 k32", emu, warps, 2.0 * 16 * 8 * 32, out); + return 0; +} diff --git a/zig/kernels/cuda/deepseek_v41/proto/router_glue_check.cu b/zig/kernels/cuda/deepseek_v41/proto/router_glue_check.cu new file mode 100644 index 0000000000..facf2620f7 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/proto/router_glue_check.cu @@ -0,0 +1,162 @@ +// Oracle and timing for router_glue.cu against the pinned upstream code; usage: router_glue_check [--reps N]. + +// Keep fmad on: the oracle reproduces upstream's FMA tree bit for bit (SASS-audited on sm_120/sm_121). +#include "../router_glue.cu" +#include "../exl3_experts.cu" +#include +#include +#include +#include +#include + +#define CK(call) do { const cudaError_t err = (call); if (err != cudaSuccess) { \ + std::fprintf(stderr, "%s:%d: %s: %s\n", __FILE__, __LINE__, #call, cudaGetErrorString(err)); \ + std::exit(2); } } while (0) + +template struct Device { + T* p; + size_t n; + explicit Device(size_t count) : n(count) { CK(cudaMalloc(&p, n * sizeof(T))); } + ~Device() { cudaFree(p); } + void set(const std::vector& v) { CK(cudaMemcpy(p, v.data(), n * sizeof(T), cudaMemcpyHostToDevice)); } + void sentinel() { CK(cudaMemset(p, 0xa5, n * sizeof(T))); } + std::vector get() const { + std::vector v(n); CK(cudaMemcpy(v.data(), p, n * sizeof(T), cudaMemcpyDeviceToHost)); return v; + } +}; + +template void same(const Device& ref, const Device& got, const char* what, int R, int mode, int edges) { + const auto a = ref.get(), b = got.get(); + if (std::memcmp(a.data(), b.data(), a.size() * sizeof(T)) == 0) return; + for (size_t i = 0; i < a.size(); ++i) + if (std::memcmp(a.data() + i, b.data() + i, sizeof(T)) != 0) { + std::fprintf(stderr, "FAIL %s R=%d mode=%d edges=%d element=%zu\n", what, R, mode, edges, i); + std::exit(1); + } +} + +uint32_t mix(uint32_t v) { v ^= v >> 16; v *= 0x7feb352du; v ^= v >> 15; v *= 0x846ca68bu; return v ^ (v >> 16); } + +template float time_us(int reps, F launch) { + cudaEvent_t a, b; CK(cudaEventCreate(&a)); CK(cudaEventCreate(&b)); + launch(); CK(cudaDeviceSynchronize()); + CK(cudaEventRecord(a)); + for (int q = 0; q < reps; ++q) launch(); + CK(cudaEventRecord(b)); CK(cudaEventSynchronize(b)); + float ms; CK(cudaEventElapsedTime(&ms, a, b)); + CK(cudaEventDestroy(a)); CK(cudaEventDestroy(b)); + return ms * 1000.0f / reps; +} + +template float time_graph_us(int reps, F launch) { + cudaStream_t stream; CK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + cudaGraph_t graph; cudaGraphExec_t exec; + CK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + launch(stream); + CK(cudaStreamEndCapture(stream, &graph)); + CK(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); + const float us = time_us(reps, [&] { CK(cudaGraphLaunch(exec, nullptr)); }); + CK(cudaGraphExecDestroy(exec)); CK(cudaGraphDestroy(graph)); CK(cudaStreamDestroy(stream)); + return us; +} + +void check_case(int R, int mode, bool edges, int reps) { + constexpr int K = 5120; + const int slots = mode == 4 ? 1 : mode == 6 ? 4 : (mode == 5 || mode == 7) ? 7 : 9; + const int E = mode == 4 ? 1 : mode == 5 ? 512 : mode == 6 ? 17 : mode == 7 ? 385 : 257; + const int maxm = mode == 2 ? 3 : R; + const int stride = K + 8, P = R * slots; + const size_t out_n = (size_t)P * K + 32; + Device x((size_t)R * stride), suh0((size_t)E * K), suh1((size_t)E * K); + Device pick(P), ids0(E + 32), ids1(E + 32), ids2(E + 32), count0(1), count1(1), count2(1); + Device mem0((size_t)E * maxm + 32), mem1((size_t)E * maxm + 32), mem2((size_t)E * maxm + 32); + Device out00(out_n), out01(out_n), out10(out_n), out11(out_n); + std::vector xv(x.n), s0(suh0.n), s1(suh1.n); + // bf16/fp16 edge bits (signed zeros, subnormals, maxima, infinities, NaNs) plus finite values of both signs + const uint16_t bf_edge[] = {0, 0x8000, 1, 0x8001, 0x007f, 0x0080, 0x3f80, 0xbf80, + 0x7f7f, 0xff7f, 0x7f80, 0xff80, 0x7fc1, 0xffc2}; + const uint16_t hf_edge[] = {0, 0x8000, 1, 0x8001, 0x03ff, 0x0400, 0x3c00, 0xbc00, + 0x7bff, 0xfbff, 0x7c00, 0xfc00, 0x7e01, 0xfe02}; + for (size_t i = 0; i < xv.size(); ++i) { + const auto h = mix((uint32_t)i + 77); + xv[i] = edges ? bf_edge[h % 14] : (uint16_t)((h & 0x8000u) | (0x3e00u + h % 512)); + } + for (size_t i = 0; i < s0.size(); ++i) { + const auto h = mix((uint32_t)i + 41), j = mix((uint32_t)i + 43); + s0[i] = edges ? hf_edge[h % 14] : (uint16_t)((h & 0x8000u) | (0x3400u + h % 2048)); + s1[i] = edges ? hf_edge[j % 14] : (uint16_t)((j & 0x8000u) | (0x3400u + j % 2048)); + } + std::vector pv(P); + for (int i = 0; i < P; ++i) { + const int slot = i % slots; + if (mode == 0 || mode == 5 || mode == 6 || mode == 7) + pv[i] = slot == slots - 1 ? E - 1 : (i * 13 + slot * 7) % (E - 1); + else if (mode == 1) pv[i] = i % 5 == 0 ? E : i % 7 == 0 ? -1 : (i * 11) % E; + else if (mode == 2) pv[i] = i % 3; // duplicates and truncation to maxm=3 + else if (mode == 3) pv[i] = i % 2 ? E : -1; // no valid experts + else pv[i] = 0; + } + x.set(xv); suh0.set(s0); suh1.set(s1); pick.set(pv); + ids0.sentinel(); ids1.sentinel(); ids2.sentinel(); mem0.sentinel(); mem1.sentinel(); mem2.sentinel(); + out00.sentinel(); out01.sentinel(); out10.sentinel(); out11.sentinel(); + const auto reference_group = [&](cudaStream_t stream = nullptr) { + tf_exl3_experts::group_kernel<<<1, 1024, P * sizeof(int), stream>>>(pick.p, ids0.p, count0.p, mem0.p, R, slots, E, maxm); + CK(cudaGetLastError()); + }; + const auto reference = [&](cudaStream_t stream = nullptr) { + reference_group(stream); + tf_exl3_experts::rot_in_kernel<__nv_bfloat16><<>>( + reinterpret_cast(x.p), stride, pick.p, + reinterpret_cast(suh0.p), reinterpret_cast(suh1.p), + reinterpret_cast(out00.p), reinterpret_cast(out01.p), K, slots, E); + CK(cudaGetLastError()); + }; + const auto group = [&](cudaStream_t stream = nullptr) { + router_group<<<1, 1024, 0, stream>>>(pick.p, ids2.p, count2.p, mem2.p, R, slots, E, maxm); + CK(cudaGetLastError()); + }; + const auto fused = [&](cudaStream_t stream = nullptr) { + router_group_rot<<<1 + (P * (K / 128) * 2 + 31) / 32, 1024, 0, stream>>>(x.p, stride, pick.p, + reinterpret_cast(suh0.p), reinterpret_cast(suh1.p), + reinterpret_cast(out10.p), reinterpret_cast(out11.p), + ids1.p, count1.p, mem1.p, R, K, slots, E, maxm); + CK(cudaGetLastError()); + }; + reference(); group(); fused(); CK(cudaDeviceSynchronize()); + same(count0, count1, "fused count", R, mode, edges); same(count0, count2, "group count", R, mode, edges); + same(ids0, ids1, "fused ids", R, mode, edges); same(ids0, ids2, "group ids", R, mode, edges); + same(mem0, mem1, "fused members", R, mode, edges); same(mem0, mem2, "group members", R, mode, edges); + same(out00, out10, "gate rotation", R, mode, edges); same(out01, out11, "up rotation", R, mode, edges); + // Repeated launches require no ticket reset; check both integer and fp bits. + for (int repeat = 0; repeat < 3; ++repeat) fused(); + CK(cudaDeviceSynchronize()); + same(mem0, mem1, "replayed members", R, mode, edges); + same(out00, out10, "replayed gate", R, mode, edges); + if ((mode == 0 || mode == 6 || mode == 7) && !edges) { + const float old_group = time_us(reps, reference_group), new_group = time_us(reps, group); + const float old_pair = time_us(reps, reference), new_pair = time_us(reps, fused); + std::printf("PASS R=%d slots=%d E=%d group_us %.3f -> %.3f pair_us %.3f -> %.3f (warm eager)\n", + R, slots, E, old_group, new_group, old_pair, new_pair); + const float graph_old = time_graph_us(reps, reference), graph_new = time_graph_us(reps, fused); + same(count0, count1, "graphed count", R, mode, edges); + same(ids0, ids1, "graphed ids", R, mode, edges); + same(mem0, mem1, "graphed members", R, mode, edges); + same(out00, out10, "graphed gate", R, mode, edges); + same(out01, out11, "graphed up", R, mode, edges); + std::printf("PASS R=%d slots=%d E=%d graph_pair_us %.3f -> %.3f nodes 2 -> 1 (warm graph replay)\n", + R, slots, E, graph_old, graph_new); + } else std::printf("PASS R=%d slots=%d E=%d maxm=%d mode=%d edges=%d\n", R, slots, E, maxm, mode, edges); +} + +int main(int argc, char** argv) { + int reps = 100; + if (argc == 3 && std::strcmp(argv[1], "--reps") == 0) reps = std::max(1, std::atoi(argv[2])); + else if (argc != 1) { std::fprintf(stderr, "usage: %s [--reps N]\n", argv[0]); return 2; } + int count = 0; CK(cudaGetDeviceCount(&count)); + if (count == 0) { std::fprintf(stderr, "No CUDA GPU: oracle not run\n"); return 2; } + for (int R : {17, 20, 24, 32, 48, 64}) + for (int mode = 0; mode < 8; ++mode) + for (bool edges : {false, true}) check_case(R, mode, edges, reps); + std::printf("96/96 cases bit-exact, replay and sentinel guards included\n"); + return 0; +} diff --git a/zig/kernels/cuda/deepseek_v41/proto/x3gm_ablate.cuh b/zig/kernels/cuda/deepseek_v41/proto/x3gm_ablate.cuh new file mode 100644 index 0000000000..d3a111ae23 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/proto/x3gm_ablate.cuh @@ -0,0 +1,180 @@ +// Timing ablations of x3gm.cu's gm_kernel: ABL 0 = the kernel, 1-7 each drop one part; ABL > 0 outputs are junk. +#pragma once +namespace dsv41_x3gm { +template +__global__ void __launch_bounds__(Cfg::THREADS, + Cfg::MINB) gm_ablate(const Args a) { + using C = Cfg; + constexpr int THREADS = C::THREADS, BM = C::BM, LDA = C::LDA, LDE = C::LDE, ER = C::ER, CPR = C::CPR, + TW = C::TW, W = C::W; + extern __shared__ __align__(16) unsigned char smem[]; + float* ep = reinterpret_cast(smem); // [MATS][ER][LDE] after the K loop (aliases the ring) + __shared__ int rows_sh[BM]; + __shared__ int item_sh; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int mat = warp / C::WPM, slice = warp % C::WPM; // columns slice 16 MTL .. + 16 MTL of the item + const int xmat = XM == 1 ? 0 : mat; + const int g = lane >> 2, t = lane & 3; + const int K = a.K, N = a.N, KT = K >> 4, NTILES = N >> 4, S = KT / KS, NBLK = N >> 7; + const size_t se = (size_t)KT * NTILES * TW; // words of one expert's matrix + const int total = a.npass[0] * NBLK; + const int lr = (lane & 7) + ((lane >> 4) << 3), lc = (lane >> 3) & 1; + const LaneMap map(lane); + + int next = blockIdx.x; + for (;;) { + __syncthreads(); // the last item is done with rows_sh and shared memory + if (threadIdx.x == 0) item_sh = a.ticket ? atomicAdd(a.ticket, 1) : next; + next += gridDim.x; + __syncthreads(); + const int item = item_sh; + if (item >= total) break; + const int tt = item / NBLK, nb = item % NBLK; + const int e = a.pe[tt], off = a.poff[tt], cnt = a.pcnt[tt]; + const uint32_t* te0 = a.tp0 ? reinterpret_cast(a.tp0[e]) : a.t0 + (size_t)e * se; + const uint32_t* te1 = a.tp1 ? reinterpret_cast(a.tp1[e]) : a.t1 + (size_t)e * se; + for (int i = threadIdx.x; i < BM; i += THREADS) rows_sh[i] = i < cnt ? a.order[off + i] : -1; + __syncthreads(); + const int cnt16 = (cnt + 15) & ~15; // rows past it are never read (warp-uniform skip) + + // stage s: member rows (zero past the count), then the item's words of k tiles KS s .. KS s + KS - 1, block nb + auto load_stage = [&](int s) { + unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + half* ra = reinterpret_cast(st); + const int k0 = s * KS * 16; + for (int i = threadIdx.x; i < XM * BM * CPR && !(ABL == 4 && s >= NSA); i += THREADS) { + const int m = i / (BM * CPR), r = (i / CPR) % BM, c = i % CPR; + if (r < cnt16) { + const int row = rows_sh[r]; + const half* src = (m ? a.x1 : a.x0) + (size_t)(row < 0 ? 0 : row) * K + k0 + c * 8; + cp_async16(ra + ((size_t)m * BM + r) * LDA + swz(r, c) * 8, src, row >= 0); + } + } + uint32_t* rw = reinterpret_cast(st + C::A_BYTES); + constexpr int WCH = MATS * KS * NB * K2; // 16-byte chunks: K2 a tile + if (ABL == 3 && s >= NSA) return; + for (int i = threadIdx.x; i < WCH; i += THREADS) { + const int q = i % K2, n = (i / K2) % NB, kk = (i / (K2 * NB)) % KS, m = i / (K2 * NB * KS); + const uint32_t* src = (m ? te1 : te0) + ((size_t)(s * KS + kk) * NTILES + nb * NB + n) * TW + q * 4; + cp_async16(rw + ((m * KS + kk) * NB + n) * TW + q * 4, src, true); + } + }; + + float acc[MTL][NG][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) +#pragma unroll + for (int n = 0; n < NG; ++n) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[l][n][c] = 0.f; + +#pragma unroll + for (int s = 0; s < NSA - 1; ++s) { + if (s < S) load_stage(s); + cp_commit(); + } + cp_wait(); + __syncthreads(); + for (int s = 0; s < S; ++s) { + if (s + NSA - 1 < S) load_stage(s + NSA - 1); + cp_commit(); + const unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + const half* A = reinterpret_cast(st) + (size_t)xmat * BM * LDA; + const uint32_t* Wd = reinterpret_cast(st + C::A_BYTES) + (size_t)mat * KS * NB * TW + + slice * MTL * TW; +#pragma unroll + for (int kk = 0; kk < KS; ++kk) { + uint32_t af[MTL][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + if constexpr (ABL == 1) { + const uint32_t* tw = Wd + (kk * NB + l) * TW; + af[l][0] = tw[lane % TW] & 0x3bff3bffu; af[l][1] = af[l][0] ^ 0x1u; + af[l][2] = af[l][0] ^ 0x2u; af[l][3] = af[l][0] ^ 0x4u; + } else { + decode_a(Wd + (kk * NB + l) * TW, map, af[l]); + } + } + const int chunk = swz(lr, kk * 2 + lc); +#pragma unroll + for (int np = 0; np < NG / 2; ++np) { + if (16 * np < cnt) { // warp-uniform + uint32_t b[4]; + ldsm4(b, A + (16 * np + lr) * LDA + chunk * 8); + const uint32_t b01[2] = {b[0], b[1]}, b23[2] = {b[2], b[3]}; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + if constexpr (ABL == 2) { + acc[l][2 * np][0] += __uint_as_float((af[l][0] ^ b01[0]) & 0x3f7fffffu); + acc[l][2 * np + 1][0] += __uint_as_float((af[l][1] ^ b23[1]) & 0x3f7fffffu); + } else { + mma16816(acc[l][2 * np], af[l], b01); + mma16816(acc[l][2 * np + 1], af[l], b23); + } + } + } + } + } + cp_wait(); + __syncthreads(); + } + + // epilogue: ER members a round, accumulators -> ep[mat][member][column], then one warp a member row + if constexpr (ABL == 5) { // no epilogue: one store keeps the chains live + if (acc[0][0][0] == 1234.5f) a.y[threadIdx.x] = acc[MTL - 1][NG - 1][3]; + continue; + } +#pragma unroll + for (int rd = 0; rd < C::ROUNDS; ++rd) { + if (rd) __syncthreads(); + if (ER * rd < cnt) { +#pragma unroll + for (int n = 0; n < NGR; ++n) { + const int ng = rd * NGR + n; + float* E = ep + ((size_t)mat * ER + 8 * n + 2 * t) * LDE + slice * 16 * MTL + g; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + E[l * 16] = acc[l][ng][0]; + E[LDE + l * 16] = acc[l][ng][1]; + E[l * 16 + 8] = acc[l][ng][2]; + E[LDE + l * 16 + 8] = acc[l][ng][3]; + } + } + } + __syncthreads(); + for (int r = warp; r < ER && ER * rd + r < cnt; r += W) { + const int row = rows_sh[ER * rd + r]; + const int c0 = 4 * lane; + const size_t col = (size_t)e * N + nb * 128 + c0; + float v[4]; + *reinterpret_cast(v) = *reinterpret_cast(ep + (size_t)r * LDE + c0); + if constexpr (ABL != 7) tf_exl3::fwht128(v, lane); + if constexpr (MATS == 2) { + float u[4]; + *reinterpret_cast(u) = *reinterpret_cast(ep + ((size_t)ER + r) * LDE + c0); + if constexpr (ABL != 7) tf_exl3::fwht128(u, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { // upstream gateup_epilogue, ACT_F32 + const float gg = fminf(v[j] * HAD * (ABL == 6 ? 1.f : __half2float(a.sv0[col + j])), a.limit); + const float uu = fminf(fmaxf(u[j] * HAD * (ABL == 6 ? 1.f : __half2float(a.sv1[col + j])), -a.limit), a.limit); + const float act = gg / (1.f + expf(-gg)) * uu; + v[j] = act * (ABL == 6 ? 1.f : __half2float(a.sd[col + j])); + } + if constexpr (ABL != 7) tf_exl3::fwht128(v, lane); + half2* o = reinterpret_cast(a.xd + (size_t)row * N + nb * 128 + c0); + o[0] = __halves2half2(__float2half_rn(v[0] * HAD), __float2half_rn(v[1] * HAD)); + o[1] = __halves2half2(__float2half_rn(v[2] * HAD), __float2half_rn(v[3] * HAD)); + } else { + float4 o; // upstream down_epilogue + o.x = v[0] * HAD * (ABL == 6 ? 1.f : __half2float(a.sv0[col + 0])); + o.y = v[1] * HAD * (ABL == 6 ? 1.f : __half2float(a.sv0[col + 1])); + o.z = v[2] * HAD * (ABL == 6 ? 1.f : __half2float(a.sv0[col + 2])); + o.w = v[3] * HAD * (ABL == 6 ? 1.f : __half2float(a.sv0[col + 3])); + *reinterpret_cast(a.y + (size_t)row * N + nb * 128 + c0) = o; + } + } + } + } +} + +} // namespace dsv41_x3gm diff --git a/zig/kernels/cuda/deepseek_v41/proto/x3gm_bench.cu b/zig/kernels/cuda/deepseek_v41/proto/x3gm_bench.cu new file mode 100644 index 0000000000..00e6325fe5 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/proto/x3gm_bench.cu @@ -0,0 +1,218 @@ +// Times x3gm rotation, gate/up and down at one DeepSeek-V4.1 layer's per-rank shapes under gb10emu.cuh modes. +#include "gb10emu.cuh" +#include "../x3gm.cu" +#include "x3gm_ablate.cuh" +#include +#include +#include +#include +#include +#include + +using namespace dsv41_x3gm; + +namespace { + +constexpr int E = 384, S = 6, D = 5120, I = 1152; + +__global__ void fill_words(uint32_t* p, size_t n, uint32_t seed) { + for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (size_t)gridDim.x * blockDim.x) { + uint32_t x = (uint32_t)i * 0x9E3779B1u ^ seed; + x ^= x >> 16; x *= 0x7feb352du; x ^= x >> 15; x *= 0x846ca68bu; x ^= x >> 16; + p[i] = x; + } +} +__global__ void fill_half(half* p, size_t n, uint32_t seed, float lo, float hi) { + for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (size_t)gridDim.x * blockDim.x) { + uint32_t x = (uint32_t)i * 0x9E3779B1u ^ seed; + x ^= x >> 16; x *= 0x7feb352du; x ^= x >> 15; + p[i] = __float2half(lo + (hi - lo) * (x & 0xffffff) / 16777216.f); + } +} + +using Kern = void (*)(const Args); +struct Launch { Kern k; int threads; size_t smem; int bm; }; + +template +Launch make() { + using C = Cfg; + Kern k = gm_kernel; + CK(cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)C::SMEM)); + CK(cudaFuncSetAttribute(k, cudaFuncAttributePreferredSharedMemoryCarveout, 100)); + return {k, C::THREADS, C::SMEM, C::BM}; +} +template +Launch make_abl() { + using C = Cfg; + Kern k = gm_ablate; + CK(cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)C::SMEM)); + CK(cudaFuncSetAttribute(k, cudaFuncAttributePreferredSharedMemoryCarveout, 100)); + return {k, C::THREADS, C::SMEM, C::BM}; +} +// ablations (x3gm_ablate.cuh) at the default configurations, one rotated input, 2 and 3 bits +#define ABL_K(A, K2) if (abl == A && k2 == K2) return down ? make_abl() : make_abl(); +#define ABL_A(A) ABL_K(A, 4) ABL_K(A, 6) +Launch abl_launch(int abl, int k2, bool down) { + ABL_A(1) ABL_A(2) ABL_A(3) ABL_A(4) ABL_A(5) ABL_A(6) ABL_A(7) + fprintf(stderr, "no ablation %d at k2 %d\n", abl, k2); exit(2); +} + +// default instances (2,048 rows: gate/up 0 with one rotated input, else 1; down 0) + alternatives, K2 4-10 = 2-5 bits +#define GU_CASE(XM, K2, CFG, ID) if (xm == XM && k2 == K2 && id == ID) return make<2, XM, K2, CFG>(); +#define DN_CASE(K2, CFG, ID) if (k2 == K2 && id == ID) return make<1, 1, K2, CFG>(); +#define GU_W(K2) GU_CASE(1, K2, GM_GU0, 0) GU_CASE(1, K2, GM_GU1, 1) GU_CASE(1, K2, GM_GU5, 5) GU_CASE(1, K2, GM_GU2, 2) GU_CASE(1, K2, GM_GU3, 3) \ + GU_CASE(2, K2, GM_GU0, 0) GU_CASE(2, K2, GM_GU1, 1) GU_CASE(2, K2, GM_GU5, 5) +#define DN_W(K2) DN_CASE(K2, GM_DN0, 0) DN_CASE(K2, GM_DN1, 1) DN_CASE(K2, GM_DN2, 2) DN_CASE(K2, GM_DN3, 3) +Launch gu_launch(int xm, int k2, int id) { + GU_W(4) GU_W(6) GU_W(8) GU_W(10) + fprintf(stderr, "no gate/up instance xm %d k2 %d cfg %d\n", xm, k2, id); exit(2); +} +Launch dn_launch(int k2, int id) { + DN_W(4) DN_W(6) DN_W(8) DN_W(10) + fprintf(stderr, "no down instance k2 %d cfg %d\n", k2, id); exit(2); +} + +struct Plan { std::vector pick, order, pe, poff, pcnt; int npass = 0; }; + +// routing: each row picks 6 distinct experts; zipf > 0 skews the experts' popularity (rank^-zipf) +Plan route(int rows, double zipf, int bm, uint32_t seed) { + Plan p; std::mt19937 rng(seed); + std::vector w(E); + for (int e = 0; e < E; ++e) w[e] = zipf > 0 ? std::pow(e + 1.0, -zipf) : 1.0; + std::shuffle(w.begin(), w.end(), rng); + std::discrete_distribution pickd(w.begin(), w.end()); + p.pick.resize((size_t)rows * S); + for (int r = 0; r < rows; ++r) { + int got = 0; + while (got < S) { + const int e = pickd(rng); + bool dup = false; + for (int j = 0; j < got; ++j) dup |= p.pick[(size_t)r * S + j] == e; + if (!dup) p.pick[(size_t)r * S + got++] = e; + } + } + const int P = rows * S; + p.order.resize(P); std::iota(p.order.begin(), p.order.end(), 0); + std::stable_sort(p.order.begin(), p.order.end(), [&](int a, int b) { return p.pick[a] < p.pick[b]; }); + std::vector cnt(E, 0); + for (int x : p.pick) ++cnt[x]; + int off = 0; + for (int e = 0; e < E; ++e) { + for (int j = 0; j < cnt[e]; j += bm) { + p.pe.push_back(e); p.poff.push_back(off + j); p.pcnt.push_back(std::min(bm, cnt[e] - j)); + } + off += cnt[e]; + } + p.npass = (int)p.pe.size(); + return p; +} + +template T* dev(const std::vector& v) { + T* d; CK(cudaMalloc(&d, v.size() * sizeof(T) + 16)); + CK(cudaMemcpy(d, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + return d; +} + +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } + +} // namespace + +int main(int argc, char** argv) { // flags: --k2 --rows --mode --gu --dn --xm --zipf --reps --target --abl + int k2 = 4, rows = 2048, gu = 0, dn = 0, xm = 1, reps = 5, abl = 0; + double zipf = 0, target = 240; + std::string mode = "full"; + for (int i = 1; i + 1 < argc; i += 2) { + const std::string a = argv[i]; + if (a == "--k2") k2 = atoi(argv[i + 1]); + else if (a == "--rows") rows = atoi(argv[i + 1]); + else if (a == "--gu") gu = atoi(argv[i + 1]); + else if (a == "--dn") dn = atoi(argv[i + 1]); + else if (a == "--xm") xm = atoi(argv[i + 1]); + else if (a == "--abl") abl = atoi(argv[i + 1]); + else if (a == "--reps") reps = atoi(argv[i + 1]); + else if (a == "--zipf") zipf = atof(argv[i + 1]); + else if (a == "--target") target = atof(argv[i + 1]); + else if (a == "--mode") mode = argv[i + 1]; + else { fprintf(stderr, "unknown %s\n", a.c_str()); return 2; } + } + gb10::eager(); + gb10::Emu emu; + emu.init(mode == "gb10" ? gb10::Mode::Gb10 : mode == "sms" ? gb10::Mode::Sms : gb10::Mode::Full); + cudaStream_t st; CK(cudaStreamCreateWithFlags(&st, cudaStreamNonBlocking)); + + const Launch G = abl ? abl_launch(abl, k2, false) : gu_launch(xm, k2, gu); + const Launch Dn = abl ? abl_launch(abl, k2, true) : dn_launch(k2, dn); + const Plan pl = route(rows, zipf, G.bm, 1234); + const Plan pd = G.bm == Dn.bm ? pl : route(rows, zipf, Dn.bm, 1234); + const int P = rows * S; + const size_t mw = (size_t)E * (D / 16) * (I / 16) * 4 * k2; // words of one stacked matrix + uint32_t *tg, *tu, *td; + CK(cudaMalloc(&tg, mw * 4)); CK(cudaMalloc(&tu, mw * 4)); CK(cudaMalloc(&td, mw * 4)); + fill_words<<<1024, 256>>>(tg, mw, 1); fill_words<<<1024, 256>>>(tu, mw, 2); fill_words<<<1024, 256>>>(td, mw, 3); + half *x, *suh_g, *suh_u, *svg, *svu, *sud, *svd, *xg, *xu, *xd; float* y; + CK(cudaMalloc(&x, (size_t)rows * D * 2)); + CK(cudaMalloc(&suh_g, (size_t)E * D * 2)); CK(cudaMalloc(&suh_u, (size_t)E * D * 2)); + CK(cudaMalloc(&svg, (size_t)E * I * 2)); CK(cudaMalloc(&svu, (size_t)E * I * 2)); + CK(cudaMalloc(&sud, (size_t)E * I * 2)); CK(cudaMalloc(&svd, (size_t)E * D * 2)); + CK(cudaMalloc(&xg, (size_t)P * D * 2)); CK(cudaMalloc(&xu, (size_t)P * D * 2)); CK(cudaMalloc(&xd, (size_t)P * I * 2)); + CK(cudaMalloc(&y, (size_t)P * D * 4)); + fill_half<<<1024, 256>>>(x, (size_t)rows * D, 4, -1.f, 1.f); + for (half* p : {suh_g, suh_u}) fill_half<<<1024, 256>>>(p, (size_t)E * D, 5, -1.f, 1.f); + for (half* p : {svg, svu, sud}) fill_half<<<1024, 256>>>(p, (size_t)E * I, 6, 0.5f, 1.5f); + fill_half<<<1024, 256>>>(svd, (size_t)E * D, 7, 0.5f, 1.5f); + int *pick = dev(pl.pick), *order = dev(pl.order), *pe = dev(pl.pe), *poff = dev(pl.poff), *pcnt = dev(pl.pcnt); + int *orderd = dev(pd.order), *ped = dev(pd.pe), *poffd = dev(pd.poff), *pcntd = dev(pd.pcnt); + int *np = dev(std::vector{pl.npass}), *npd = dev(std::vector{pd.npass}); + int* ticket; CK(cudaMalloc(&ticket, 8)); + CK(cudaDeviceSynchronize()); + + int per_g = 0, per_d = 0; + CK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_g, G.k, G.threads, G.smem)); + CK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_d, Dn.k, Dn.threads, Dn.smem)); + emu.calibrate(tg, mw * 4, st, target); + const int gg = emu.grid_sms() * per_g, gd = emu.grid_sms() * per_d; + + Args ag{}; ag.x0 = xg; ag.x1 = xm == 2 ? xu : nullptr; ag.t0 = tg; ag.t1 = tu; ag.order = order; ag.pe = pe; + ag.poff = poff; ag.pcnt = pcnt; ag.npass = np; ag.ticket = ticket; ag.sv0 = svg; ag.sv1 = svu; ag.sd = sud; + ag.xd = xd; ag.K = D; ag.N = I; ag.limit = 10.f; + Args ad{}; ad.x0 = xd; ad.t0 = td; ad.order = orderd; ad.pe = ped; ad.poff = poffd; ad.pcnt = pcntd; ad.npass = npd; + ad.ticket = ticket + 1; ad.sv0 = svd; ad.y = y; ad.K = I; ad.N = D; + const int rot_items = P * (D / 128); + + cudaEvent_t ev[4]; + for (auto& e : ev) CK(cudaEventCreate(&e)); + std::vector t_rot, t_gu, t_dn; + for (int r = 0; r < reps + 1; ++r) { + CK(cudaMemsetAsync(ticket, 0, 8, st)); + emu.begin(); + const bool dbg = getenv("X3GM_DEBUG") != nullptr; + auto step = [&](const char* w) { if (dbg) { CK(cudaStreamSynchronize(st)); fprintf(stderr, " %s done\n", w); } }; + step("begin"); + CK(cudaEventRecord(ev[0], st)); + if (xm == 2) rot_kernel<<<(rot_items + 7) / 8, 256, 0, st>>>(x, D, pick, suh_g, suh_u, xg, xu, D, S, rot_items); + else rot_kernel<<<(rot_items + 7) / 8, 256, 0, st>>>(x, D, pick, suh_g, nullptr, xg, nullptr, D, S, rot_items); + step("rot"); + CK(cudaEventRecord(ev[1], st)); + G.k<<>>(ag); + step("gateup"); + CK(cudaEventRecord(ev[2], st)); + Dn.k<<>>(ad); + CK(cudaEventRecord(ev[3], st)); + CK(cudaGetLastError()); + CK(cudaEventSynchronize(ev[3])); + emu.end(); + float a, b, c; + CK(cudaEventElapsedTime(&a, ev[0], ev[1])); CK(cudaEventElapsedTime(&b, ev[1], ev[2])); CK(cudaEventElapsedTime(&c, ev[2], ev[3])); + if (r) { t_rot.push_back(a); t_gu.push_back(b); t_dn.push_back(c); } + } + const double bits = k2 / 2.0, wbytes = 3.0 * E * D * I * bits / 8; + const double flop = 2.0 * P * 3.0 * D * I; + const double tr = median(t_rot), tg_ = median(t_gu), td_ = median(t_dn), tot = tr + tg_ + td_; + cudaDeviceProp prop{}; CK(cudaGetDeviceProperties(&prop, 0)); + int clk = 0; CK(cudaDeviceGetAttribute(&clk, cudaDevAttrClockRate, 0)); + printf("RESULT abl=%d mode=%s k2=%d rows=%d xm=%d gu=%d dn=%d zipf=%.2f passes=%d/%d grid=%d/%d (per_sm %d/%d) sms=%d clk=%dMHz hogs=%d\n", + abl, mode.c_str(), k2, rows, xm, gu, dn, zipf, pl.npass, pd.npass, gg, gd, per_g, per_d, emu.grid_sms(), clk / 1000, emu.hogs); + printf("RESULT ms rot %.3f gateup %.3f down %.3f total %.3f | weights %.3f GB -> %.1f GB/s | %.1f TFLOP/s\n", + tr, tg_, td_, tot, wbytes / 1e9, wbytes / (tot * 1e6), flop / (tot * 1e9)); + return 0; +} diff --git a/zig/kernels/cuda/deepseek_v41/ptxas.sh b/zig/kernels/cuda/deepseek_v41/ptxas.sh new file mode 100755 index 0000000000..a90929d20d --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/ptxas.sh @@ -0,0 +1,27 @@ +#!/bin/bash +# Compiles DeepSeek-V4.1 kernel copies to fatbins in the NGC PyTorch image with -Xptxas -v; reports spills / stack. +set -euo pipefail +# usage: OUT=

SM=121,120 bash ptxas.sh name[:flags] ... (flags comma separated as in cuda.zig, e.g. "x3ld:-O3") +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT="${OUT:?set OUT}" +SM="${SM:-121}" +IMAGE="${TF_ZIG_IMAGE:-nvcr.io/nvidia/pytorch:26.07-py3}" +mkdir -p "$OUT" +TORCH="-D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -std=c++20" +cmds="" +for spec in "$@"; do + name="${spec%%:*}"; flags="${spec#*:}"; [ "$flags" = "$spec" ] && flags="-O3" + gen="" + for sm in ${SM//,/ }; do gen="$gen -gencode=arch=compute_${sm},code=sm_${sm}"; done + cmds="$cmds nvcc -fatbin $TORCH ${flags//,/ } $gen -Xptxas -v -o /out/$name.fatbin /k/$name.cu > /out/$name.ptxas 2>&1 && echo built $name || { echo FAILED $name; tail -20 /out/$name.ptxas; } ;" +done +docker run --rm --network none --memory "${TF_ZIG_DOCKER_MEMORY:-12g}" --memory-swap "${TF_ZIG_DOCKER_MEMORY:-12g}" -u "$(id -u):$(id -g)" -v "$HERE:/k:ro" -v "$OUT:/out" --entrypoint bash "$IMAGE" -c "$cmds" +for spec in "$@"; do + name="${spec%%:*}" + f="$OUT/$name.ptxas" + [ -f "$f" ] || continue + fns=$(grep -c "Compiling entry function" "$f" || true) + spill=$(grep -E "bytes spill (stores|loads)" "$f" | grep -vc " 0 bytes spill stores, 0 bytes spill loads" || true) + stack=$(grep -E "bytes stack frame" "$f" | grep -vc " 0 bytes stack frame" || true) + echo "$name: $fns entry functions; with spills: $spill; with a stack frame: $stack" +done diff --git a/zig/kernels/cuda/deepseek_v41/router_gemv.cu b/zig/kernels/cuda/deepseek_v41/router_gemv.cu new file mode 100644 index 0000000000..60db87747a --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/router_gemv.cu @@ -0,0 +1,521 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/router_gemv.cu (blob 5841ced53bf5): sass_check.sh +#define DSV41_RG_NO_TORCH + + + +// DeepSeek-V4.1 MoE router, one GEMV + top-K launch: each logit is one fp32 FMA chain + xor butterfly, any tiling +#include +#include +#include + +namespace dsv41_rg { + +constexpr int NV = 20; // 16-byte vectors a lane a gate row: D = 32 x 8 x NV = 5,120 +constexpr int D = 32 * 8 * NV; +constexpr int JC = 2; // vectors a lane a chunk +constexpr int NC = NV / JC; // chunks of x a row tile +constexpr int CK = 32 * 8 * JC; // k a chunk (512) +constexpr int EPL = 16; // experts a lane in the selection: E <= 512 +constexpr int KMAX = 8; // top-k at most +static_assert(NV % JC == 0, "whole chunks"); + +__device__ __forceinline__ uint4 ldg_stream(const uint16_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; +} + +// 8 bf16 (k ascending: the low half of each word first) -> fp32, exactly +__device__ __forceinline__ void widen8(const uint4 v, float (&f)[8]) { + f[0] = __uint_as_float(v.x << 16); f[1] = __uint_as_float(v.x & 0xffff0000u); + f[2] = __uint_as_float(v.y << 16); f[3] = __uint_as_float(v.y & 0xffff0000u); + f[4] = __uint_as_float(v.z << 16); f[5] = __uint_as_float(v.z & 0xffff0000u); + f[6] = __uint_as_float(v.w << 16); f[7] = __uint_as_float(v.w & 0xffff0000u); +} + +// the reference router's safe softplus, operation for operation (no contraction: the intrinsics are never fused) +__device__ __forceinline__ float softplus(float z) { + const float e = expf(fminf(z, 20.0f)); + if (z > 20.0f) return z; + if (z < -10.0f) return __fsub_rn(e, __fmul_rn(0.5f, __fmul_rn(e, e))); + return logf(__fadd_rn(1.0f, e)); +} + +// TF_DSV41_RG_PRUNE: the top-p prune folded into the selection, its operations in order on the row's K weights +struct Prune { + float min_w, topp, scale; + int min_k, orig, drop, dup, on; +}; + +__device__ __forceinline__ void prune_row(const Prune& P, int K, float (&w)[KMAX], int (&p)[KMAX]) { + float total = 0.0f, f[KMAX], excl[KMAX]; + int rank[KMAX]; + bool keep[KMAX]; +#pragma unroll + for (int i = 0; i < KMAX; ++i) if (i < K) total = __fadd_rn(total, w[i]); +#pragma unroll + for (int i = 0; i < KMAX; ++i) if (i < K) f[i] = __fdiv_rn(w[i], total); +#pragma unroll + for (int i = 0; i < KMAX; ++i) { + rank[i] = 0; +#pragma unroll + for (int j = 0; j < KMAX; ++j) + if (i < K && j < K) rank[i] += (f[j] > f[i]) || (f[j] == f[i] && j < i); + } + float acc = 0.0f; +#pragma unroll + for (int r = 0; r < KMAX; ++r) +#pragma unroll + for (int i = 0; i < KMAX; ++i) + if (r < K && i < K && rank[i] == r) { excl[i] = acc; acc = __fadd_rn(acc, f[i]); } + int nk = 0, top = 0; + float ks = 0.0f; +#pragma unroll + for (int i = 0; i < KMAX; ++i) { + if (i >= K) break; + keep[i] = ((excl[i] < P.topp) && (f[i] >= P.min_w)) || (rank[i] < P.min_k); + nk += keep[i]; + if (rank[i] == 0) top = p[i]; + ks = __fadd_rn(ks, keep[i] ? w[i] : 0.0f); + } + if (nk >= K) return; +#pragma unroll + for (int i = 0; i < KMAX; ++i) { + if (i >= K) break; + w[i] = keep[i] ? (P.orig ? w[i] : __fmul_rn(P.scale, __fdiv_rn(w[i], ks))) : 0.0f; + if (!keep[i]) p[i] = P.dup ? top : P.drop; + } +} + +// Prune pass of one selected row (lane 0); noinline so its arrays never share a live range with select_row's +__device__ __noinline__ void prune_pass(const Prune& P, int K, int kit, int* pick, float* wts, int* psh) { + float wv[KMAX]; + int ip[KMAX]; +#pragma unroll + for (int t = 0; t < KMAX; ++t) { + wv[t] = t < K ? wts[t] : 0.0f; + ip[t] = t < K ? pick[t] : 0; + } + prune_row(P, K, wv, ip); +#pragma unroll + for (int t = 0; t < KMAX; ++t) { + if (t >= K) break; + pick[t] = ip[t]; + wts[t] = kit ? __half2float(__float2half_rn(wv[t])) : wv[t]; + psh[t] = ip[t]; + } +} + +// One row's selection by one warp (lane l holds experts l + 32 i). +__device__ __forceinline__ void select_row(const float* lg, const float* __restrict__ bias, int E, int K, int slots, + float scale, int kit, int* pick, float* wts, int* psh, int lane, + const Prune& P) { + float s[EPL], ch[EPL]; +#pragma unroll + for (int i = 0; i < EPL; ++i) { + const int e = lane + 32 * i; + if (e < E) { + s[i] = __fsqrt_rn(softplus(__ldcg(lg + e))); + ch[i] = __fadd_rn(s[i], bias[e]); + } else { + s[i] = 0.0f; + ch[i] = -INFINITY; + } + } + float sp[KMAX]; + int ip[KMAX]; + float total = 0.0f; +#pragma unroll + for (int t = 0; t < KMAX; ++t) { + if (t >= K) break; + float bv = -INFINITY; + int bi = 0x7fffffff; +#pragma unroll + for (int i = 0; i < EPL; ++i) // ids ascend with i: strict > keeps the lower id on a tie + if (ch[i] > bv) { bv = ch[i]; bi = lane + 32 * i; } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { // (score desc, id asc): a total order, every lane ends equal + const float ov = __shfl_xor_sync(0xffffffffu, bv, o); + const int oi = __shfl_xor_sync(0xffffffffu, bi, o); + if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } + } + float mine = 0.0f; +#pragma unroll + for (int i = 0; i < EPL; ++i) + if (lane + 32 * i == bi) { mine = s[i]; ch[i] = -INFINITY; } + sp[t] = __shfl_sync(0xffffffffu, mine, bi & 31); + ip[t] = bi; + total = __fadd_rn(total, sp[t]); + } + if (lane == 0) { +#pragma unroll + for (int t = 0; t < KMAX; ++t) { + if (t >= K) break; + float w = __fmul_rn(scale, __fdiv_rn(sp[t], total)); + if (kit && !P.on) w = __half2float(__float2half_rn(w)); // pick.kit_weights (with prune: after it) + pick[t] = ip[t]; + wts[t] = w; + psh[t] = ip[t]; + } + if (slots > K) { pick[K] = E; wts[K] = 1.0f; psh[K] = E; } + } +} + +struct Args { + const uint16_t* x; // [R, x_stride] bf16, rows 16-byte aligned + long long xs; + const uint16_t* w; // [E, D] bf16 + const float* bias; // [E] + int* pick; // [R, slots] + float* wts; // [R, slots] + float* lg; // [R, E] fp32 (scratch, or the caller's logits) + int* cnt; // [>= grid.y] int32, 0 between launches + int R, E, K, slots, select; + float scale; + int kit; // the routed weights rounded to fp16 (pick.kit_weights, folded) + int* gid; // group != 0: upstream's group kernel folded (one row tile only): distinct ids (< GE) + int* gcnt; // in id order, their count, members [., maxm] row * 32 + slot in row order, -1 after + int* gmem; + int GE, maxm, group; + Prune prune; // the top-p prune rule folded into the selection (on = 0: off) +}; + +constexpr int GE_MAX = 512; // table entries the folded grouping takes +constexpr int PICKS_MAX = 16 * (KMAX + 1); +constexpr int TAIL_INTS = PICKS_MAX; // the tail's picks (both kernels) +constexpr int GROUP_INTS = TAIL_INTS + 2 * GE_MAX + PICKS_MAX + 8; // + the grouping's (the narrow kernel) + +// Upstream's group_kernel (exl3/experts.cu) on the R <= 16 rows of one row tile, by the selecting CTA; integers only +template +__device__ __forceinline__ void group_rows(const Args& a, const int* psh, int* sh, int tid) { + constexpr int T = NW * 32; + int* gc = sh; // the main loop's shared memory, reused + int* gpl = gc + GE_MAX; + int* cpl = gpl + GE_MAX; + int* wtot = cpl + PICKS_MAX; + const int n = a.R * a.slots, GE = a.GE, lane = tid & 31, warp = tid >> 5; + for (int e = tid; e < GE; e += T) gc[e] = 0; + __syncthreads(); + for (int i = tid; i < n; i += T) { + const int e = psh[i]; + if (e >= 0 && e < GE) atomicAdd(gc + e, 1); + } + __syncthreads(); + const int per = (GE + T - 1) / T, lo = min(GE, tid * per), hi = min(GE, lo + per); + int used = 0; + for (int e = lo; e < hi; ++e) used += gc[e] > 0; + int inc = used; // inclusive scan in the warp, then over warps +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int v = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) inc += v; + } + if (lane == 31) wtot[warp] = inc; + __syncthreads(); + int base = 0, total = 0; +#pragma unroll + for (int w = 0; w < NW; ++w) { + base += w < warp ? wtot[w] : 0; + total += wtot[w]; + } + int place = base + inc - used; + for (int e = lo; e < hi; ++e) + if (gc[e] > 0) { + a.gid[place] = e; + gpl[e] = place; + cpl[place] = gc[e]; + ++place; + } + if (tid == 0) a.gcnt[0] = total; + __syncthreads(); + for (int i = tid; i < n; i += T) { // rank of pick i among its expert's picks: row order + const int e = psh[i]; + if (e < 0 || e >= GE) continue; + int rank = 0; + for (int q = 0; q < i; ++q) rank += psh[q] == e; + if (rank < a.maxm) a.gmem[(size_t)gpl[e] * a.maxm + rank] = (i / a.slots) * 32 + i % a.slots; + } + for (int q = tid; q < total * a.maxm; q += T) + if (q % a.maxm >= cpl[q / a.maxm]) a.gmem[q] = -1; +} + +// Fused tail: the row tile's last CTA selects its rows, groups them (GROUP), resets the counter; sh: TAIL_INTS ints +template +__device__ __forceinline__ void tail(const Args& a, int r0, int nr, int tid, int* sh) { + __shared__ int last_sh; + int* psh = sh; + const int lane = tid & 31, warp = tid >> 5; + __threadfence(); // this CTA's logits visible device-wide before its arrival + __syncthreads(); // (also: every thread is past the main loop's shared memory) + if (tid == 0) last_sh = atomicAdd(a.cnt + blockIdx.y, 1) == (int)gridDim.x - 1; + __syncthreads(); + if (!last_sh) return; + __threadfence(); + for (int r = warp; r < nr; r += NW) + select_row(a.lg + (size_t)(r0 + r) * a.E, a.bias, a.E, a.K, a.slots, a.scale, a.kit, + a.pick + (size_t)(r0 + r) * a.slots, a.wts + (size_t)(r0 + r) * a.slots, psh + r * a.slots, lane, + a.prune); + if (a.prune.on && lane == 0) // the rule, then the kit rounding (each lane 0: its own rows) + for (int r = warp; r < nr; r += NW) + prune_pass(a.prune, a.K, a.kit, a.pick + (size_t)(r0 + r) * a.slots, a.wts + (size_t)(r0 + r) * a.slots, + psh + r * a.slots); + if (GROUP && a.group) { + __syncthreads(); + group_rows(a, psh, sh + TAIL_INTS, tid); + } + if (tid == 0) a.cnt[blockIdx.y] = 0; // self-cleaning +} + +// RT rows a tile (1 .. 16), NW warps a CTA, EW experts a warp. +template +__global__ void __launch_bounds__(NW * 32) gemv_kernel(const Args a) { + constexpr int PD = 8 / EW; // vectors in flight a lane, per expert + constexpr int XU = RT * CK / 8; // uint4 of x a chunk + constexpr int XT = (XU + NW * 32 - 1) / (NW * 32); + static_assert(PD < NV, "ring shorter than a row"); + __shared__ __align__(16) uint16_t xsh[2][RT][CK]; + + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int r0 = blockIdx.y * RT; + const int nr = min(RT, a.R - r0); + const int e0 = (blockIdx.x * NW + warp) * EW; + + auto wload = [&](int j, uint4 (&dst)[EW]) { +#pragma unroll + for (int e = 0; e < EW; ++e) { + if (e0 + e < a.E) dst[e] = ldg_stream(a.w + (size_t)(e0 + e) * D + (size_t)(j * 32 + lane) * 8); + else dst[e] = make_uint4(0u, 0u, 0u, 0u); + } + }; + auto xload = [&](int c, uint4 (&xr)[XT]) { +#pragma unroll + for (int q = 0; q < XT; ++q) { + const int u = tid + NW * 32 * q, row = u / (CK / 8), col = (u % (CK / 8)) * 8; + xr[q] = make_uint4(0u, 0u, 0u, 0u); + if (u < XU && row < nr) + xr[q] = __ldg(reinterpret_cast(a.x + (size_t)(r0 + row) * a.xs + c * CK + col)); + } + }; + auto xstore = [&](int buf, const uint4 (&xr)[XT]) { +#pragma unroll + for (int q = 0; q < XT; ++q) { + const int u = tid + NW * 32 * q, row = u / (CK / 8), col = (u % (CK / 8)) * 8; + if (u < XU) *reinterpret_cast(&xsh[buf][row][col]) = xr[q]; + } + }; + + float acc[EW][RT]; +#pragma unroll + for (int e = 0; e < EW; ++e) +#pragma unroll + for (int r = 0; r < RT; ++r) acc[e][r] = 0.0f; + uint4 ring[PD][EW]; +#pragma unroll + for (int j = 0; j < PD; ++j) wload(j, ring[j]); + uint4 xr[XT]; + xload(0, xr); + xstore(0, xr); + __syncthreads(); + +#pragma unroll + for (int j = 0; j < NV; ++j) { + const int c = j / JC, jj = j % JC; + if (jj == 0 && c + 1 < NC) xload(c + 1, xr); // next chunk's x: L2, consumed at the chunk's end + float wf[EW][8]; +#pragma unroll + for (int e = 0; e < EW; ++e) widen8(ring[j % PD][e], wf[e]); + if (j + PD < NV) wload(j + PD, ring[j % PD]); // the slot refilled PD steps ahead +#pragma unroll + for (int r = 0; r < RT; ++r) { + float xf[8]; + widen8(*reinterpret_cast(&xsh[c & 1][r][(jj * 32 + lane) * 8]), xf); +#pragma unroll + for (int e = 0; e < EW; ++e) +#pragma unroll + for (int i = 0; i < 8; ++i) acc[e][r] = __fmaf_rn(xf[i], wf[e][i], acc[e][r]); + } + if (jj == JC - 1) { + if (c + 1 < NC) xstore((c + 1) & 1, xr); // the buffer chunk c - 1 read (all past its barrier) + __syncthreads(); + } + } + +#pragma unroll + for (int e = 0; e < EW; ++e) +#pragma unroll + for (int r = 0; r < RT; ++r) { + float v = acc[e][r]; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, o)); + if (lane == r && r < nr && e0 + e < a.E) a.lg[(size_t)(r0 + r) * a.E + e0 + e] = v; + } + static_assert(sizeof(xsh) >= TAIL_INTS * sizeof(int), "the tail's picks fit the staged x"); + if (a.select) tail(a, r0, nr, tid, reinterpret_cast(&xsh[0][0][0])); +} + +#include "router_narrow.cuh" + +template +cudaError_t launch_narrow(const Args& a, cudaStream_t stream) { + static const bool carved = [] { // the most shared memory a SM (8 / NW CTAs resident at once); a + cudaFuncSetAttribute(narrow_kernel, // preference only: an error here is cleared, never fatal + cudaFuncAttributePreferredSharedMemoryCarveout, (int)cudaSharedmemCarveoutMaxShared); + (void)cudaGetLastError(); + return true; + }(); + (void)carved; + const dim3 grid((unsigned)((a.E + NW - 1) / NW), (unsigned)((a.R + RT - 1) / RT)); + narrow_kernel<<>>(a); + return cudaGetLastError(); +} + +template +cudaError_t by_rt_narrow(const Args& a, int rt, cudaStream_t stream) { + switch (rt) { + case 1: return launch_narrow<1, NW>(a, stream); + case 2: return launch_narrow<2, NW>(a, stream); + case 4: return launch_narrow<4, NW>(a, stream); + case 8: return launch_narrow<8, NW>(a, stream); + case 16: return launch_narrow<16, NW>(a, stream); + default: return cudaErrorInvalidValue; + } +} + +template +cudaError_t launch(const Args& a, cudaStream_t stream) { + const dim3 grid((unsigned)((a.E + NW * EW - 1) / (NW * EW)), (unsigned)((a.R + RT - 1) / RT)); + gemv_kernel<<>>(a); + return cudaGetLastError(); +} + +template +cudaError_t by_rt(const Args& a, int rt, cudaStream_t stream) { + switch (rt) { + case 1: return launch<1, NW, EW>(a, stream); + case 2: return launch<2, NW, EW>(a, stream); + case 4: return launch<4, NW, EW>(a, stream); + case 8: return launch<8, NW, EW>(a, stream); + case 16: return launch<16, NW, EW>(a, stream); + default: return cudaErrorInvalidValue; + } +} + +// rt rows a tile (1, 2, 4, 8, 16); (nw, ew) in {4, 8} x {1, 2}, or ew == 0: narrow, nw in {1, 2, 4, 8} +inline cudaError_t dispatch(const Args& a, int rt, int nw, int ew, cudaStream_t stream) { + if (ew == 0) { + switch (nw) { + case 1: return by_rt_narrow<1>(a, rt, stream); + case 2: return by_rt_narrow<2>(a, rt, stream); + case 4: return by_rt_narrow<4>(a, rt, stream); + case 8: return by_rt_narrow<8>(a, rt, stream); + default: return cudaErrorInvalidValue; + } + } + if (nw == 8 && ew == 1) return by_rt<8, 1>(a, rt, stream); + if (nw == 4 && ew == 1) return by_rt<4, 1>(a, rt, stream); + if (nw == 8 && ew == 2) return by_rt<8, 2>(a, rt, stream); + if (nw == 4 && ew == 2) return by_rt<4, 2>(a, rt, stream); + return cudaErrorInvalidValue; +} + +} // namespace dsv41_rg + +#ifndef DSV41_RG_NO_TORCH // the compile test builds the kernels alone (nvcc, no torch headers) +#include +#include +#include + +void dsv41_rg_route_cuda(const at::Tensor& X, const at::Tensor& W, const at::Tensor& B, at::Tensor& pick, + at::Tensor& wts, at::Tensor& LG, at::Tensor& cnt, int64_t K, int64_t slots, double scale, + int64_t rt, int64_t nw, int64_t ew, bool select, bool kit, at::Tensor& gid, at::Tensor& gcnt, + at::Tensor& gmem, int64_t GE, bool group, const std::vector& prune) { + using namespace dsv41_rg; + TORCH_CHECK(prune.empty() || prune.size() == 7, "router_gemv: prune = [min_w, topp, scale, min_k, orig, drop, dup]"); + const int64_t R = X.size(0), E = W.size(0); + TORCH_CHECK(X.size(1) == D && W.size(1) == D, "router_gemv: D must be ", D); + TORCH_CHECK(X.stride(1) == 1 && X.stride(0) % 8 == 0 && reinterpret_cast(X.data_ptr()) % 16 == 0, + "router_gemv: x rows must be unit-stride and 16-byte aligned"); + TORCH_CHECK(E >= 1 && E <= 32 * EPL && K >= 1 && K <= KMAX && K <= E && slots >= K && slots <= K + 1, + "router_gemv: E <= ", 32 * EPL, ", 1 <= K <= ", KMAX, ", slots K or K + 1"); + TORCH_CHECK(LG.numel() >= R * E, "router_gemv: logits buffer too small"); + TORCH_CHECK(!select || cnt.numel() >= (R + rt - 1) / rt, "router_gemv: counters too few for ", R, " rows"); + const bool wide = (nw == 4 || nw == 8) && (ew == 1 || ew == 2); + const bool narrow = ew == 0 && (nw == 1 || nw == 2 || nw == 4 || nw == 8); + TORCH_CHECK((rt == 1 || rt == 2 || rt == 4 || rt == 8 || rt == 16) && (wide || narrow), + "router_gemv: no instance (rt, nw, ew) = (", rt, ", ", nw, ", ", ew, ")"); + TORCH_CHECK(!group || (select && narrow && R <= rt && GE >= 1 && GE <= GE_MAX && gid.numel() >= std::min(R * slots, GE) + && gcnt.numel() >= 1 && gmem.dim() == 2 && gmem.size(0) >= std::min(R * slots, GE) + && gmem.size(1) >= 1), + "router_gemv: the folded grouping takes one row tile, GE <= ", GE_MAX, " and ids / members sized"); + if (R == 0) return; + Args a; + a.x = reinterpret_cast(X.data_ptr()); + a.xs = X.stride(0); + a.w = reinterpret_cast(W.data_ptr()); + a.bias = B.data_ptr(); + a.pick = pick.data_ptr(); + a.wts = wts.data_ptr(); + a.lg = LG.data_ptr(); + a.cnt = cnt.data_ptr(); + a.R = (int)R; a.E = (int)E; a.K = (int)K; a.slots = (int)slots; a.select = select ? 1 : 0; + a.scale = (float)scale; + a.kit = kit ? 1 : 0; + a.group = group ? 1 : 0; + a.gid = group ? gid.data_ptr() : nullptr; + a.gcnt = group ? gcnt.data_ptr() : nullptr; + a.gmem = group ? gmem.data_ptr() : nullptr; + a.GE = (int)GE; + a.maxm = group ? (int)gmem.size(1) : 0; // upstream's Scratch.window: members [maxu, R] + a.prune = Prune{}; + if (!prune.empty()) { // fp32 thresholds: the rounding of Triton's float arguments + a.prune.min_w = (float)prune[0]; a.prune.topp = (float)prune[1]; a.prune.scale = (float)prune[2]; + a.prune.min_k = (int)prune[3]; a.prune.orig = (int)prune[4]; a.prune.drop = (int)prune[5]; + a.prune.dup = (int)prune[6]; a.prune.on = 1; + } + C10_CUDA_CHECK(dispatch(a, (int)rt, (int)nw, (int)ew, at::cuda::getCurrentCUDAStream())); +} +#endif + +// The instantiations dispatch() launches: gemv_kernel and narrow_kernel. +template __global__ void dsv41_rg::gemv_kernel<1, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<1, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<1, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<1, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<1, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<1, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<1, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<1, 8>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<2, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<2, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<2, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<2, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<2, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<2, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<2, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<2, 8>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<4, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<4, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<4, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<4, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<4, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<4, 8>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<8, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<8, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<8, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<8, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<8, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<8, 8>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<16, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<16, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<16, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<16, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<16, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<16, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<16, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<16, 8>(const dsv41_rg::Args); diff --git a/zig/kernels/cuda/deepseek_v41/router_glue.cu b/zig/kernels/cuda/deepseek_v41/router_glue.cu new file mode 100644 index 0000000000..5e39b67d95 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/router_glue.cu @@ -0,0 +1,115 @@ +// Decode router glue, arithmetic of exl3/experts.cu (blob 8ce66ca83bce); -O3, fmad on, ftz off; R <= 64, K == 5120. +#include +#include +#include +#include + +namespace dsv41_router_glue { +constexpr int THREADS = 1024; +constexpr int WARPS = THREADS / 32; +constexpr int MAX_E = 512; +constexpr int MAX_PICKS = 64 * 9; +constexpr float HAD_SCALE = 0.08838834764831845f; + +// Verbatim arithmetic from tf_exl3_experts::fwht128 at the source above. +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + 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) { + float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +// Integer-equal to upstream group_kernel (histogram, id scan, earlier equal picks; duplicates past maxm truncated). +__device__ __forceinline__ void group_rows(const int* pick, int* uids, int* ucount, + int* members, int R, int slots, int E, int maxm) { + __shared__ int picks[MAX_PICKS], count[MAX_E], place[MAX_E], warp_tot[WARPS]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int n = R * slots; + if (tid < n) picks[tid] = pick[tid]; + if (tid < E) count[tid] = 0; + __syncthreads(); + if (tid < n) { + const int e = picks[tid]; + if (e >= 0 && e < E) atomicAdd(count + e, 1); + } + __syncthreads(); + const int used = tid < E && count[tid] > 0; + int inc = used; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int v = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) inc += v; + } + if (lane == 31) warp_tot[warp] = inc; + __syncthreads(); + if (warp == 0) { + const int v = warp_tot[lane]; + int s = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int x = __shfl_up_sync(0xffffffffu, s, o); + if (lane >= o) s += x; + } + warp_tot[lane] = s - v; + if (lane == 31) ucount[0] = s; + } + __syncthreads(); + if (used) { + const int p = warp_tot[warp] + inc - used; + uids[p] = tid; + place[tid] = p; + } + __syncthreads(); + if (tid < n) { + const int e = picks[tid]; + if (e >= 0 && e < E) { + int rank = 0; + for (int q = 0; q < tid; ++q) rank += picks[q] == e; + if (rank < maxm) members[(size_t)place[e] * maxm + rank] = (tid / slots) * 32 + tid % slots; + } + } + if (used) + for (int j = count[tid]; j < maxm; ++j) members[(size_t)place[tid] * maxm + j] = -1; +} +} // namespace dsv41_router_glue + +// One CTA, block 1024, no dynamic shared memory. Same argument order as group. +extern "C" __global__ void __launch_bounds__(1024) router_group( + const int* pick, int* uids, int* ucount, int* members, int R, int slots, int E, int maxm) { + dsv41_router_glue::group_rows(pick, uids, ucount, members, R, slots, E, maxm); +} + +// grid.x = 1 + ceil(R slots (K / 128) 2 / 32), block 1024: CTA 0 groups, other warps rotate one (pick, block, mat). +extern "C" __global__ void __launch_bounds__(1024) router_group_rot( + const uint16_t* x, int x_stride, const int* pick, const half* suh0, const half* suh1, + half* out0, half* out1, int* uids, int* ucount, int* members, + int R, int K, int slots, int E, int maxm) { + using namespace dsv41_router_glue; + if (blockIdx.x == 0) { + group_rows(pick, uids, ucount, members, R, slots, E, maxm); + return; + } + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int tile = (blockIdx.x - 1) * WARPS + warp; + const int P = R * slots, blocks = K / 128; + if (tile >= P * blocks * 2) return; // uniform within each warp + const int p = tile % P, blk = (tile / P) % blocks, mat = tile / (P * blocks); + const int row = p / slots, e = pick[p]; + if (e < 0 || e >= E) return; // invalid picks leave outputs untouched + const half* suh = (mat ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = reinterpret_cast(x) + + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? out1 : out0) + (size_t)p * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} diff --git a/zig/kernels/cuda/deepseek_v41/router_narrow.cuh b/zig/kernels/cuda/deepseek_v41/router_narrow.cuh new file mode 100644 index 0000000000..2b38f41edb --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/router_narrow.cuh @@ -0,0 +1,96 @@ +// router_gemv.cu's narrow_kernel and its cp.async helpers, inside namespace dsv41_rg +#pragma once +// x through L1 (ld.global.nc), volatile so it stays in the step loop (a plain __ldg was hoisted, spilling registers) +__device__ __forceinline__ uint4 ldg_keep(const uint16_t* p) { + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p)); + return v; +} + +__device__ __forceinline__ void cp_async16(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) : "memory"); +} +__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;" ::: "memory"); } +template +__device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;" ::"n"(N) : "memory"); } +// n a constant once the step loop is unrolled: the switch folds to one wait_group +__device__ __forceinline__ void cp_wait_n(int n) { + switch (n) { +#define DSV41_RG_W(k) case k: cp_wait(); break; + DSV41_RG_W(0) DSV41_RG_W(1) DSV41_RG_W(2) DSV41_RG_W(3) DSV41_RG_W(4) DSV41_RG_W(5) DSV41_RG_W(6) DSV41_RG_W(7) + DSV41_RG_W(8) DSV41_RG_W(9) DSV41_RG_W(10) DSV41_RG_W(11) DSV41_RG_W(12) DSV41_RG_W(13) DSV41_RG_W(14) + DSV41_RG_W(15) DSV41_RG_W(16) DSV41_RG_W(17) DSV41_RG_W(18) DSV41_RG_W(19) +#undef DSV41_RG_W + default: cp_wait<0>(); break; + } +} + +// Narrow kernel (TF_DSV41_RG_NARROW = NW, rows <= 16): gemv_kernel's chains and bits; gate via cp.async, x via L1 +template +__global__ void __launch_bounds__(NW * 32) narrow_kernel(const Args a) { + constexpr int PD = NW <= 4 && RT <= 4 ? NV : NV / 2; // gate vectors in flight a lane + constexpr int XP = RT <= 4 ? 2 : 0; // steps of x in flight (0: loaded at use, L1) + constexpr int TS = (GROUP_INTS * 4 + NW * 512 - 1) / (NW * 512); // slots the tail needs + __shared__ __align__(16) uint4 wsh[NW][PD > TS ? PD : TS][32]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int r0 = blockIdx.y * RT; + const int nr = min(RT, a.R - r0); + const int e = blockIdx.x * NW + warp; + const bool live = e < a.E; + const uint16_t* wrow = a.w + (size_t)(live ? e : 0) * D + (size_t)lane * 8; + + auto wload = [&](int j) { // a dead warp's slots are never read: committed empty + if (live) cp_async16(&wsh[warp][j % PD][lane], wrow + (size_t)j * 256); + cp_commit(); + }; + auto xload = [&](int j, uint4 (&dst)[RT]) { +#pragma unroll + for (int r = 0; r < RT; ++r) + dst[r] = (live && r < nr) + ? ldg_keep(a.x + (size_t)(r0 + r) * a.xs + (size_t)(j * 32 + lane) * 8) + : make_uint4(0u, 0u, 0u, 0u); + }; + +#pragma unroll + for (int j = 0; j < PD; ++j) wload(j); + uint4 xr[XP > 0 ? XP : 1][RT]; +#pragma unroll + for (int j = 0; j < XP; ++j) xload(j, xr[j]); + float acc[RT]; +#pragma unroll + for (int r = 0; r < RT; ++r) acc[r] = 0.0f; + +#pragma unroll + for (int j = 0; j < NV; ++j) { + // groups committed before step j: min(NV, PD + j); group j done = at most that - j - 1 pending + cp_wait_n(min(NV - j - 1, PD - 1)); + float wf[8]; + widen8(wsh[warp][j % PD][lane], wf); + uint4 xc[RT]; + if constexpr (XP == 0) { + xload(j, xc); + } else { +#pragma unroll + for (int r = 0; r < RT; ++r) xc[r] = xr[j % XP][r]; + if (j + XP < NV) xload(j + XP, xr[j % XP]); + } +#pragma unroll + for (int r = 0; r < RT; ++r) { + float xf[8]; + widen8(xc[r], xf); +#pragma unroll + for (int i = 0; i < 8; ++i) acc[r] = __fmaf_rn(xf[i], wf[i], acc[r]); + } + if (j + PD < NV) wload(j + PD); // the slot just read (its value is in wf: used above) + } + +#pragma unroll + for (int r = 0; r < RT; ++r) { + float v = acc[r]; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, o)); + if (lane == r && r < nr && live) a.lg[(size_t)(r0 + r) * a.E + e] = v; + } + if (a.select) tail(a, r0, nr, tid, reinterpret_cast(&wsh[0][0][0])); +} diff --git a/zig/kernels/cuda/deepseek_v41/topk_cuda.cu b/zig/kernels/cuda/deepseek_v41/topk_cuda.cu new file mode 100644 index 0000000000..f633f78e22 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/topk_cuda.cu @@ -0,0 +1,296 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/csa2/topk_cuda.cu (blob 839ee8a85f0e): sass_check.sh +#define DSV41_ATTN_NO_TORCH + + + +// Indexer decode top-k (TF_DSV41_ATTN_CUDA=1): top K of unique keys (score image, then lower position), exact. + +// Radix select of the K-th 32-bit image over a cluster of <= 8 CTAs a row (DSMEM histograms), then scan-order pick. + +// Image: s == 0 -> 0, NaN -> 0x7FC00000, negative -> bits ^ 0x7FFFFFFF, then ^ 0x80000000 (0 marks missing). + +#include +#include +#include + +namespace dsv41_topk { +namespace cg = cooperative_groups; + +constexpr int NT = 512, NW = NT / 32, MAXCL = 8, MAXEPT = 31; +constexpr unsigned FULL = 0xffffffffu; + +struct Job { + const float* s; long long ss; // scores [R, n] (unit stride) + const int* cand; long long cs; // mode 3: ascending candidate blocks [R, cols] (-1 last) + int* out; long long os; // [R, k] + int* cnt; // [R] visible counts or null + int nk, k, mode, ept; // entries, K, 0 dense / 2 blocks / 3 candidate positions, entries a thread +}; + +struct Args { + Job job[2]; + const void* pos; int pos64; // int32 [1] (POS + r) or int64 [R] + int ratio, bs; +}; + +__device__ __forceinline__ uint32_t image(float s) { + int b = __float_as_int(s); + if (s == 0.f) b = 0; + if (s != s) b = 0x7FC00000; + if (b < 0) b ^= 0x7FFFFFFF; + return static_cast(b) ^ 0x80000000u; +} + +// exclusive scan of one int a thread over the CTA (512 threads); returns the prefix, *total the sum +__device__ int block_scan(int v, int* ws, int* total) { + const int lane = threadIdx.x & 31, w = threadIdx.x >> 5; + int x = v; + for (int o = 1; o < 32; o <<= 1) { + const int y = __shfl_up_sync(FULL, x, o); + if (lane >= o) x += y; + } + if (lane == 31) ws[w] = x; + __syncthreads(); + if (w == 0) { + int z = lane < NW ? ws[lane] : 0; + for (int o = 1; o < 32; o <<= 1) { + const int y = __shfl_up_sync(FULL, z, o); + if (lane >= o) z += y; + } + if (lane < NW) ws[NW + lane] = z; + } + __syncthreads(); + const int pre = (w ? ws[NW + w - 1] : 0) + x - v; + *total = ws[2 * NW - 1]; + __syncthreads(); + return pre; +} + +__global__ void __launch_bounds__(NT, 1) topk_kernel(Args a) { + cg::cluster_group cl = cg::this_cluster(); + const Job J = blockIdx.z == 0 ? a.job[0] : a.job[1]; // (no dynamic indexing: no local memory) + const int CL = static_cast(cl.num_blocks()), rank = static_cast(cl.block_rank()); + const int r = blockIdx.y, tid = threadIdx.x, lane = tid & 31; + extern __shared__ uint32_t img[]; // SL = 512 x ept images (0: missing) + __shared__ int hist[2][256]; + __shared__ int tot[256]; + __shared__ int ws[2 * NW]; + __shared__ int ex[2]; // (above, equal) of this CTA, read by the cluster + __shared__ int dec[4]; + const int SL = NT * J.ept, lo = rank * SL, n = max(0, min(SL, J.nk - lo)); + const int q = a.pos64 ? static_cast(static_cast(a.pos)[r]) + : static_cast(a.pos)[0] + r; + const int nvis = (q + 1) / a.ratio; + if (rank == 0 && tid == 0 && J.cnt) J.cnt[r] = 0; // before the first cluster barrier (release) + const float* srow = J.s + static_cast(r) * J.ss; + const int* crow = J.cand ? J.cand + static_cast(r) * J.cs : nullptr; + for (int i = tid; i < SL; i += NT) { // images, coalesced + const int e = lo + i; + uint32_t u = 0; + if (i < n) { + if (J.mode == 2) { // block e: max over its visible members, newest +inf + float m = 0.f; // tl.max (maxnumf) over the members, masked -inf + for (int j = 0; j < a.bs; ++j) { + const int p = e * a.bs + j; + const float v = p < nvis ? srow[p] : -INFINITY; + m = j ? fmaxf(m, v) : v; + } + if (e == (nvis - 1) / a.bs) m = INFINITY; + if (e * a.bs < nvis) u = image(m); + } else if (J.mode == 3) { + if (crow[e / a.bs] >= 0) u = image(srow[e]); + } else { + u = image(srow[e]); + } + } + img[i] = u; + } + __syncthreads(); + uint32_t prefix = 0, pmask = 0; + int above = 0, quota_mode = 0; // 0: every match selected; 1: equal images capped + for (int p = 0; p < 4; ++p) { + int* h = hist[p & 1]; + if (tid < 256) h[tid] = 0; + __syncthreads(); + const int shift = 24 - 8 * p; + for (int i0 = 0; i0 < SL; i0 += NT) { + const uint32_t v = img[i0 + tid]; + const bool m = v != 0 && (v & pmask) == prefix; + const unsigned act = __ballot_sync(FULL, m); + if (m) { + const int d = (v >> shift) & 255; + const unsigned peers = __match_any_sync(act, d); + if (lane == __ffs(peers) - 1) atomicAdd(&h[d], __popc(peers)); + } + } + cl.sync(); + if (tid < 256) { + int v[MAXCL]; +#pragma unroll + for (int x = 0; x < MAXCL; ++x) v[x] = x < CL ? cl.map_shared_rank(h, x)[tid] : 0; + int sum = 0; +#pragma unroll + for (int x = 0; x < MAXCL; ++x) sum += v[x]; + tot[tid] = sum; + } + __syncthreads(); + if (tid < 32) { // suffix sums over 256 digits, 8 a lane; the digit + int loc[8], run = 0; + for (int x = 7; x >= 0; --x) { run += tot[lane * 8 + x]; loc[x] = run; } + int suf = run; // inclusive suffix over lanes >= this one + for (int o = 1; o < 32; o <<= 1) { + const int y = __shfl_down_sync(FULL, suf, o); + if (lane + o < 32) suf += y; + } + const int after = suf - run; // lanes above this one + int best = -1, bge = 0; + for (int x = 7; x >= 0; --x) + if (best < 0 && above + after + loc[x] >= J.k) { best = lane * 8 + x; bge = after + loc[x]; } + const unsigned has = __ballot_sync(FULL, best >= 0); + const int src = has ? 31 - __clz(has) : 0; + const int dg = __shfl_sync(FULL, best, src), ge = __shfl_sync(FULL, bge, src); + const int all = __shfl_sync(FULL, suf, 0); + if (lane == 0) { dec[0] = dg; dec[1] = ge; dec[2] = all; } + } + __syncthreads(); + const int dg = dec[0], ge = dec[1], all = dec[2]; + if (p == 0 && all <= J.k) break; // every valid entry (prefix / mask stay 0) + const int hd = tot[dg]; + above += ge - hd; + prefix |= static_cast(dg) << shift; + pmask |= 255u << shift; + if (above + hd == J.k) break; // the prefix's entries complete the top K + if (p == 3) quota_mode = 1; // ties at the K-th image: the first ones win + } + // classify: 2 above (selected), 1 equal (selected up to the quota, in scan order), 0 out + const int e0 = tid * J.ept; + int gt = 0, eq = 0; + for (int x = 0; x < J.ept; ++x) { + const uint32_t v = img[e0 + x]; + if (v == 0) continue; + const uint32_t mv = v & pmask; + gt += mv > prefix; + eq += mv == prefix; + } + int tgt, teq; + const int pre_eq = block_scan(eq, ws, &teq); + block_scan(gt, ws, &tgt); + if (tid == 0) { ex[0] = tgt; ex[1] = teq; } + cl.sync(); + int base = 0, eq_before = 0, total = 0, quota = teq; + { + const int need = J.k - above; // equal images to take (quota mode) + for (int x = 0; x < CL; ++x) { + const int* rx = cl.map_shared_rank(ex, x); + const int xg = rx[0], xe = rx[1]; + const int xq = quota_mode ? max(0, min(xe, need - eq_before)) : xe; + if (x < rank) base += xg + xq; + if (x == rank) quota = xq; + eq_before += xe; + total += xg + xq; + } + } + cl.sync(); // no DSMEM access after this: CTAs may exit + int sel = 0; + { + int er = pre_eq; + for (int x = 0; x < J.ept; ++x) { + const uint32_t v = img[e0 + x]; + if (v == 0) continue; + const uint32_t mv = v & pmask; + if (mv > prefix) ++sel; + else if (mv == prefix) { sel += er < quota; ++er; } + } + } + int tsel; + const int off = base + block_scan(sel, ws, &tsel); + int* orow = J.out + static_cast(r) * J.os; + int vis = 0; + { + int er = pre_eq, o = off; + for (int x = 0; x < J.ept; ++x) { + const uint32_t v = img[e0 + x]; + if (v == 0) continue; + const uint32_t mv = v & pmask; + bool s = mv > prefix; + if (!s && mv == prefix) { s = er < quota; ++er; } + if (!s) continue; + const int e = lo + e0 + x; + const int pp = J.mode == 3 ? crow[e / a.bs] * a.bs + e % a.bs : e; + orow[o++] = pp; + vis += pp < nvis; + } + } + if (J.cnt) { + int tv; + block_scan(vis, ws, &tv); + if (tid == 0 && tv) atomicAdd(J.cnt + r, tv); + } + if (rank == 0) + for (int i = total + tid; i < J.k; i += NT) orow[i] = -1; +} + +// Shared memory a CTA for ``ept`` entries a thread (the images; the rest is static). +inline size_t smem(int ept) { return static_cast(NT) * ept * 4; } + +inline cudaError_t launch(const Args& a, int jobs, int R, int CL, int ept, cudaStream_t stream) { + if (R < 1 || jobs < 1 || jobs > 2 || CL < 1 || CL > MAXCL || ept < 1 || ept > MAXEPT) + return cudaErrorInvalidValue; + static bool init = false; + if (!init) { + const cudaError_t e = cudaFuncSetAttribute(topk_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem(MAXEPT))); + if (e != cudaSuccess) return e; + init = true; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(CL, R, jobs); + cfg.blockDim = dim3(NT); + cfg.dynamicSmemBytes = smem(ept); + cfg.stream = stream; + cudaLaunchAttribute at[1]; + at[0].id = cudaLaunchAttributeClusterDimension; + at[0].val.clusterDim.x = CL; + at[0].val.clusterDim.y = 1; + at[0].val.clusterDim.z = 1; + cfg.attrs = at; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, topk_kernel, a); +} + +} // namespace dsv41_topk + +#ifndef DSV41_ATTN_NO_TORCH +#include +#include +#include + +namespace { +dsv41_topk::Job job_of(const at::Tensor& s, const at::Tensor& cand, at::Tensor& out, at::Tensor& cnt, int64_t nk, + int64_t k, int64_t mode, int64_t ept) { + dsv41_topk::Job j{}; + j.s = static_cast(s.data_ptr()); j.ss = s.stride(0); + j.cand = cand.numel() ? static_cast(cand.data_ptr()) : nullptr; + j.cs = cand.numel() ? cand.stride(0) : 0; + j.out = static_cast(out.data_ptr()); j.os = out.stride(0); + j.cnt = cnt.numel() ? static_cast(cnt.data_ptr()) : nullptr; + j.nk = static_cast(nk); j.k = static_cast(k); j.mode = static_cast(mode); + j.ept = static_cast(ept); + return j; +} +} + +void dsv41_topk_run_cuda(const at::Tensor& s0, const at::Tensor& c0, at::Tensor& o0, at::Tensor& n0, int64_t nk0, + int64_t k0, int64_t m0, int64_t e0, const at::Tensor& s1, const at::Tensor& c1, + at::Tensor& o1, at::Tensor& n1, int64_t nk1, int64_t k1, int64_t m1, int64_t e1, + int64_t jobs, const at::Tensor& pos, int64_t ratio, int64_t bs, int64_t R, int64_t CL) { + dsv41_topk::Args a{}; + a.job[0] = job_of(s0, c0, o0, n0, nk0, k0, m0, e0); + a.job[1] = jobs > 1 ? job_of(s1, c1, o1, n1, nk1, k1, m1, e1) : a.job[0]; + a.pos = pos.data_ptr(); a.pos64 = pos.scalar_type() == at::kLong ? 1 : 0; + a.ratio = static_cast(ratio); a.bs = static_cast(bs); + const int ept = static_cast(std::max(e0, jobs > 1 ? e1 : e0)); + C10_CUDA_CHECK(dsv41_topk::launch(a, static_cast(jobs), static_cast(R), static_cast(CL), ept, + at::cuda::getCurrentCUDAStream())); +} +#endif diff --git a/zig/kernels/cuda/deepseek_v41/topk_keys.cu b/zig/kernels/cuda/deepseek_v41/topk_keys.cu new file mode 100644 index 0000000000..1bf475ce08 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/topk_keys.cu @@ -0,0 +1,94 @@ +// Row top-k of bf16-valued fp32 logits on the device: value descending, ties by lower column (torch.topk of keys). + +// keys = (mono(bits(x + 0.0)) & -65536) | (65535 - col): unique in a row (C <= 65,536), so any exact select matches. + +// One CTA a row: 4-pass 8-bit radix select of the k-th key, gather k winners, bitonic sort; cols int64, vals fp32. + +#include + +namespace dsv41_topk_keys { + +constexpr int NT = 1024, KMAX = 1024, COLS = 1 << 16; + +__device__ __forceinline__ uint32_t ukey(float x, int col) { + int b = __float_as_int(x + 0.0f); // -0.0 + 0.0 = +0.0 + const int mono = b >= 0 ? b : (b ^ 0x7FFFFFFF); + const int key = (mono & -COLS) | (COLS - 1 - col); + return static_cast(key) ^ 0x80000000u; // signed order as unsigned +} + +__global__ void __launch_bounds__(NT) topk_kernel(const float* __restrict__ lg, long long ld, int C, int k, + float* __restrict__ vals, long long* __restrict__ cols) { + __shared__ int hist[256]; + __shared__ uint32_t s_prefix, s_mask; + __shared__ int s_need, s_n; + __shared__ uint32_t key_sh[KMAX]; + __shared__ int col_sh[KMAX]; + const int tid = threadIdx.x; + const float* row = lg + (long long)blockIdx.x * ld; + + // radix select: the k-th largest key, 8 bits a pass from the top + if (tid == 0) { s_prefix = 0; s_mask = 0; s_need = k; } + __syncthreads(); + for (int shift = 24; shift >= 0; shift -= 8) { + if (tid < 256) hist[tid] = 0; + __syncthreads(); + const uint32_t prefix = s_prefix, mask = s_mask; + for (int c = tid; c < C; c += NT) { + const uint32_t u = ukey(row[c], c); + if ((u & mask) == prefix) atomicAdd(&hist[(u >> shift) & 255], 1); + } + __syncthreads(); + if (tid == 0) { // the digit holding the need-th largest, from the top + int need = s_need, d = 255; + for (; d > 0 && hist[d] < need; --d) need -= hist[d]; + s_need = need; + s_prefix = prefix | ((uint32_t)d << shift); + s_mask = mask | (255u << shift); + } + __syncthreads(); + } + const uint32_t kth = s_prefix; // the k-th largest key itself (keys are unique) + + // the winners (> kth, and kth): exactly k, in any order + if (tid == 0) s_n = 0; + __syncthreads(); + for (int c = tid; c < C; c += NT) { + const uint32_t u = ukey(row[c], c); + if (u >= kth) { + const int at = atomicAdd(&s_n, 1); + key_sh[at] = u; + col_sh[at] = c; + } + } + __syncthreads(); + int n = 1; + while (n < k) n <<= 1; + for (int i = k + tid; i < n; i += NT) { key_sh[i] = 0; col_sh[i] = -1; } // pads sort last (keys >= 1) + __syncthreads(); + + // bitonic sort, descending + for (int size = 2; size <= n; size <<= 1) { + for (int stride = size >> 1; stride > 0; stride >>= 1) { + for (int i = tid; i < n; i += NT) { + const int j = i ^ stride; + if (j > i) { + const bool desc = (i & size) == 0; + const uint32_t a = key_sh[i], b = key_sh[j]; + if (desc ? a < b : a > b) { + key_sh[i] = b; key_sh[j] = a; + const int t = col_sh[i]; col_sh[i] = col_sh[j]; col_sh[j] = t; + } + } + } + __syncthreads(); + } + } + for (int i = tid; i < k; i += NT) { + const int c = col_sh[i]; + cols[(long long)blockIdx.x * k + i] = c; + vals[(long long)blockIdx.x * k + i] = row[c]; + } +} + +} // namespace dsv41_topk_keys diff --git a/zig/kernels/cuda/deepseek_v41/x3gm.cu b/zig/kernels/cuda/deepseek_v41/x3gm.cu new file mode 100644 index 0000000000..3798f4e298 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/x3gm.cu @@ -0,0 +1,477 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/x3gm.cu (blob 7aed2b027afc): sass_check.sh + + + +// x3gm: routed EXL3 prefill experts (gate/up, down); each output is one mma chain over ascending k from +0, any tiling + +#include +#include + +#include "decode.cuh" +#include "experts_grouped.cuh" + +namespace dsv41_x3gm { + +using tf_exl3x::cb_pair; +using tf_exl3x::Fmt; +using tf_exl3x::LaneMap; +using tf_exl3x::mma16816; + +constexpr float HAD = 0.08838834764831845f; // 1 / sqrt(128) +constexpr int NB = 8; // column tiles an item: 128 columns, one Hadamard block +constexpr int MUL1 = 2; + +__device__ __forceinline__ void cp_async16(void* smem, const void* gmem, bool ok) { + const unsigned s = (unsigned)__cvta_generic_to_shared(smem); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(s), "l"(gmem), "r"(ok ? 16 : 0)); +} +__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } +__device__ __forceinline__ void ldsm4(uint32_t (&a)[4], const void* p) { + const unsigned s = (unsigned)__cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(s)); +} + +// the 16-byte chunk c of member row r in a stage (CPR chunks a row): 8 rows of an ldmatrix phase in 8 bank groups +template +__device__ __forceinline__ int swz(int r, int c) { + static_assert(CPR == 8 || CPR == 4, "rows of 64 or 32 halves"); + return CPR == 8 ? (c ^ (r & 7)) : (c ^ ((r >> 1) & 3)); +} + +// A 16x16 W_q tile from shared words as the A fragment of W^T: upstream decode_tile, registers reordered +template +__device__ __forceinline__ void decode_a(const uint32_t* tile, const LaneMap& m, uint32_t (&a)[4]) { + constexpr int GV = Fmt::GV, NG = Fmt::NG; + uint32_t st[8]; +#pragma unroll + for (int g = 0; g < NG; ++g) { + const uint32_t whi = tile[m.hi[g]], wlo = tile[m.lo[g]]; + const uint64_t mm = ((((uint64_t)wlo) << 32) | whi) >> m.sh[g]; +#pragma unroll + for (int j = 0; j < GV; ++j) st[g * GV + j] = (uint32_t)(mm >> Fmt::off(j)) & 0xffffu; + } + a[0] = cb_pair(st[0], st[1]); // b0[0]: k 2t, 2t+1 of column g -> A row g, k 2t.. + a[1] = cb_pair(st[4], st[5]); // b1[0]: k 2t, 2t+1 of column g + 8 -> A row g + 8, k 2t.. + a[2] = cb_pair(st[2], st[3]); // b0[1]: k 2t+8, 2t+9 of column g -> A row g, k 2t+8.. + a[3] = cb_pair(st[6], st[7]); // b1[1] -> A row g + 8, k 2t+8.. +} + +template +struct Cfg { + static constexpr int BM = NG * 8, WPM = NB / MTL, W = MATS * WPM, THREADS = W * 32; + static constexpr int CPR = KS * 2, LDA = KS * 16, TW = 4 * K2; + static constexpr int LDE = 132, ER = NGR * 8, ROUNDS = NG / NGR; + static constexpr size_t A_BYTES = (size_t)XM * BM * LDA * sizeof(half); + static constexpr size_t W_BYTES = (size_t)MATS * KS * NB * TW * sizeof(uint32_t); + static constexpr size_t STAGE = A_BYTES + W_BYTES; + static constexpr size_t RING = (size_t)NSA * STAGE; + static constexpr size_t EPI = (size_t)MATS * ER * LDE * sizeof(float); + static constexpr size_t SMEM = RING > EPI ? RING : EPI; + // CTAs an SM the registers are capped for: 2 when two rings fit, <= 8 warps, <= 64 accumulators a lane; else 1 + static constexpr int MINB = 2 * (SMEM + 1536) <= 102400 && THREADS <= 256 && MTL * NG <= 16 ? 2 : 1; + static_assert(NG % 2 == 0 && NG % NGR == 0 && NGR % 2 == 0 && NB % MTL == 0, "tiling"); + static_assert(A_BYTES % 16 == 0 && (TW * 4) % 16 == 0, "16-byte stages"); + static_assert(XM == 1 || XM == MATS, "one shared input or one a matrix"); +}; + +struct Args { + const half* x0; // [P, K] rows of matrix 0 (gate / down input) + const half* x1; // [P, K] rows of matrix 1 (up) when XM == 2 + const uint32_t* t0; // [E, K/16, N/16, 4 K2] trellis words, stacked (when tp0 is null) + const uint32_t* t1; + const int64_t* tp0; // [E] expert e's words [K/16, N/16, 4 K2] (a ragged stack), or null + const int64_t* tp1; + const int* k2e; // [E] each expert's K2 (gm2_kernel; the launch's width table) + const int* order; // pair indices by expert (stable: rows ascending) + const int* pe; // pass -> expert + const int* poff; // pass -> first index into order + const int* pcnt; // pass -> members (<= BM) + const int* npass; // [1] passes in all + int* ticket; // [1] zeroed, or null (static stride) + const half* sv0; // gate/up: svh_g [E, N]; down: svh_d [E, N] + const half* sv1; // svh_u [E, N] + const half* sd; // suh_d [E, N] + half* xd; // gate/up out [P, N] + float* y; // down out [P, N] + int K, N; + float limit; +}; + +template +__global__ void __launch_bounds__(Cfg::THREADS, + Cfg::MINB) gm_kernel(const Args a) { + using C = Cfg; + constexpr int THREADS = C::THREADS, BM = C::BM, LDA = C::LDA, LDE = C::LDE, ER = C::ER, CPR = C::CPR, + TW = C::TW, W = C::W; + extern __shared__ __align__(16) unsigned char smem[]; + float* ep = reinterpret_cast(smem); // [MATS][ER][LDE] after the K loop (aliases the ring) + __shared__ int rows_sh[BM]; + __shared__ int item_sh; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int mat = warp / C::WPM, slice = warp % C::WPM; // columns slice 16 MTL .. + 16 MTL of the item + const int xmat = XM == 1 ? 0 : mat; + const int g = lane >> 2, t = lane & 3; + const int K = a.K, N = a.N, KT = K >> 4, NTILES = N >> 4, S = KT / KS, NBLK = N >> 7; + const size_t se = (size_t)KT * NTILES * TW; // words of one expert's matrix + const int total = a.npass[0] * NBLK; + const int lr = (lane & 7) + ((lane >> 4) << 3), lc = (lane >> 3) & 1; + const LaneMap map(lane); + + int next = blockIdx.x; + for (;;) { + __syncthreads(); // the last item is done with rows_sh and shared memory + if (threadIdx.x == 0) item_sh = a.ticket ? atomicAdd(a.ticket, 1) : next; + next += gridDim.x; + __syncthreads(); + const int item = item_sh; + if (item >= total) break; + const int tt = item / NBLK, nb = item % NBLK; + const int e = a.pe[tt], off = a.poff[tt], cnt = a.pcnt[tt]; + const uint32_t* te0 = a.tp0 ? reinterpret_cast(a.tp0[e]) : a.t0 + (size_t)e * se; + const uint32_t* te1 = a.tp1 ? reinterpret_cast(a.tp1[e]) : a.t1 + (size_t)e * se; + for (int i = threadIdx.x; i < BM; i += THREADS) rows_sh[i] = i < cnt ? a.order[off + i] : -1; + __syncthreads(); + const int cnt16 = (cnt + 15) & ~15; // rows past it are never read (warp-uniform skip) + + // stage s: member rows' k [16 KS s, + 16 KS) (zero past the count), then the item's k tile words, block nb + auto load_stage = [&](int s) { + unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + half* ra = reinterpret_cast(st); + const int k0 = s * KS * 16; + for (int i = threadIdx.x; i < XM * BM * CPR; i += THREADS) { + const int m = i / (BM * CPR), r = (i / CPR) % BM, c = i % CPR; + if (r < cnt16) { + const int row = rows_sh[r]; + const half* src = (m ? a.x1 : a.x0) + (size_t)(row < 0 ? 0 : row) * K + k0 + c * 8; + cp_async16(ra + ((size_t)m * BM + r) * LDA + swz(r, c) * 8, src, row >= 0); + } + } + uint32_t* rw = reinterpret_cast(st + C::A_BYTES); + constexpr int WCH = MATS * KS * NB * K2; // 16-byte chunks: K2 a tile + for (int i = threadIdx.x; i < WCH; i += THREADS) { + const int q = i % K2, n = (i / K2) % NB, kk = (i / (K2 * NB)) % KS, m = i / (K2 * NB * KS); + const uint32_t* src = (m ? te1 : te0) + ((size_t)(s * KS + kk) * NTILES + nb * NB + n) * TW + q * 4; + cp_async16(rw + ((m * KS + kk) * NB + n) * TW + q * 4, src, true); + } + }; + + float acc[MTL][NG][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) +#pragma unroll + for (int n = 0; n < NG; ++n) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[l][n][c] = 0.f; + +#pragma unroll + for (int s = 0; s < NSA - 1; ++s) { + if (s < S) load_stage(s); + cp_commit(); + } + cp_wait(); + __syncthreads(); + for (int s = 0; s < S; ++s) { + if (s + NSA - 1 < S) load_stage(s + NSA - 1); + cp_commit(); + const unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + const half* A = reinterpret_cast(st) + (size_t)xmat * BM * LDA; + const uint32_t* Wd = reinterpret_cast(st + C::A_BYTES) + (size_t)mat * KS * NB * TW + + slice * MTL * TW; +#pragma unroll + for (int kk = 0; kk < KS; ++kk) { + uint32_t af[MTL][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) decode_a(Wd + (kk * NB + l) * TW, map, af[l]); + const int chunk = swz(lr, kk * 2 + lc); +#pragma unroll + for (int np = 0; np < NG / 2; ++np) { + if (16 * np < cnt) { // warp-uniform + uint32_t b[4]; + ldsm4(b, A + (16 * np + lr) * LDA + chunk * 8); + const uint32_t b01[2] = {b[0], b[1]}, b23[2] = {b[2], b[3]}; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + mma16816(acc[l][2 * np], af[l], b01); + mma16816(acc[l][2 * np + 1], af[l], b23); + } + } + } + } + cp_wait(); + __syncthreads(); + } + + // epilogue: ER members a round, accumulators -> ep[mat][member][column], then one warp a member row +#pragma unroll + for (int rd = 0; rd < C::ROUNDS; ++rd) { + if (rd) __syncthreads(); + if (ER * rd < cnt) { +#pragma unroll + for (int n = 0; n < NGR; ++n) { + const int ng = rd * NGR + n; + float* E = ep + ((size_t)mat * ER + 8 * n + 2 * t) * LDE + slice * 16 * MTL + g; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + E[l * 16] = acc[l][ng][0]; + E[LDE + l * 16] = acc[l][ng][1]; + E[l * 16 + 8] = acc[l][ng][2]; + E[LDE + l * 16 + 8] = acc[l][ng][3]; + } + } + } + __syncthreads(); + for (int r = warp; r < ER && ER * rd + r < cnt; r += W) { + const int row = rows_sh[ER * rd + r]; + const int c0 = 4 * lane; + const size_t col = (size_t)e * N + nb * 128 + c0; + float v[4]; + *reinterpret_cast(v) = *reinterpret_cast(ep + (size_t)r * LDE + c0); + tf_exl3::fwht128(v, lane); + if constexpr (MATS == 2) { + float u[4]; + *reinterpret_cast(u) = *reinterpret_cast(ep + ((size_t)ER + r) * LDE + c0); + tf_exl3::fwht128(u, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { // upstream gateup_epilogue, ACT_F32 + const float gg = fminf(v[j] * HAD * __half2float(a.sv0[col + j]), a.limit); + const float uu = fminf(fmaxf(u[j] * HAD * __half2float(a.sv1[col + j]), -a.limit), a.limit); + const float act = gg / (1.f + expf(-gg)) * uu; + v[j] = act * __half2float(a.sd[col + j]); + } + tf_exl3::fwht128(v, lane); + half2* o = reinterpret_cast(a.xd + (size_t)row * N + nb * 128 + c0); + o[0] = __halves2half2(__float2half_rn(v[0] * HAD), __float2half_rn(v[1] * HAD)); + o[1] = __halves2half2(__float2half_rn(v[2] * HAD), __float2half_rn(v[3] * HAD)); + } else { + float4 o; // upstream down_epilogue + o.x = v[0] * HAD * __half2float(a.sv0[col + 0]); + o.y = v[1] * HAD * __half2float(a.sv0[col + 1]); + o.z = v[2] * HAD * __half2float(a.sv0[col + 2]); + o.w = v[3] * HAD * __half2float(a.sv0[col + 3]); + *reinterpret_cast(a.y + (size_t)row * N + nb * 128 + c0) = o; + } + } + } + } +} + +#include "x3gm2.cuh" + +// Rotated routed inputs: out_m[p] = fp16((x[p / slots] * suh_m[e]) H / sqrt 128), upstream rot_in's arithmetic +template __device__ __forceinline__ float to_f(TIN v); +template <> __device__ __forceinline__ float to_f<__nv_bfloat16>(__nv_bfloat16 v) { return __bfloat162float(v); } +template <> __device__ __forceinline__ float to_f(half v) { return __half2float(v); } + +template +__global__ void __launch_bounds__(256) rot_kernel(const TIN* __restrict__ x, int x_stride, const int* __restrict__ pick, + const half* __restrict__ suh0, const half* __restrict__ suh1, + half* __restrict__ out0, half* __restrict__ out1, int K, int slots, + int items) { + const int item = blockIdx.x * 8 + (threadIdx.x >> 5); + if (item >= items) return; // warp-uniform + const int KB = K >> 7, p = item / KB, blk = item % KB, lane = threadIdx.x & 31; + const int e = pick[p]; + const TIN* xr = x + (size_t)(p / slots) * x_stride + blk * 128 + 4 * lane; + float xv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) xv[j] = to_f(xr[j]); +#pragma unroll + for (int m = 0; m < MATS; ++m) { + const half* suh = (m ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = xv[j] * __half2float(suh[j]); + tf_exl3::fwht128(v, lane); + half* o = (m ? out1 : out0) + (size_t)p * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD); + } +} + +// Tests: W_q [K, N] fp16 of one trellis through decode_a (words staged in shared memory as the gm kernels stage them). +template +__global__ void __launch_bounds__(32) dequant_kernel(const uint32_t* __restrict__ T, half* __restrict__ out, int N) { + constexpr int TW = 4 * K2; + __shared__ __align__(16) uint32_t tile[TW]; + const int kt = blockIdx.x, nt = blockIdx.y, lane = threadIdx.x, NTILES = N >> 4; + for (int i = lane; i < TW; i += 32) tile[i] = T[((size_t)kt * NTILES + nt) * TW + i]; + __syncwarp(); + const LaneMap map(lane); + uint32_t af[4]; + decode_a(tile, map, af); + const int g = lane >> 2, t = lane & 3; +#pragma unroll + for (int q = 0; q < 4; ++q) { // A row (q & 1 ? g + 8 : g) = W column; k 2t (+8) + const half2 h = *reinterpret_cast(&af[q]); + const int col = nt * 16 + g + 8 * (q & 1), row = kt * 16 + 2 * t + 8 * (q >> 1); + out[(size_t)row * N + col] = __low2half(h); + out[(size_t)(row + 1) * N + col] = __high2half(h); + } +} + + +// Tile table (MTL, NG, KS, NSA, NGR), mirrored by the host tile lists (a test parses these lines); 2.. = 4K tiles +#define GM_GU0 2, 8, 4, 3, 4 +#define GM_GU1 2, 8, 2, 4, 4 +#define GM_GU2 1, 16, 2, 4, 4 +#define GM_GU3 1, 16, 4, 3, 4 +#define GM_GU4 2, 16, 2, 3, 4 +#define GM_GU5 2, 8, 2, 6, 4 +#define GM_DN0 2, 8, 4, 4, 8 +#define GM_DN1 1, 8, 4, 4, 8 +#define GM_DN2 1, 16, 4, 4, 8 +#define GM_DN3 1, 16, 2, 6, 8 +#define GM_DN4 2, 16, 4, 3, 16 +#define GM_NGU 6 +#define GM_NDN 5 + + +} // namespace dsv41_x3gm + +#ifndef DSV41_X3GM_NO_INSTANCES +// Instances the bindings launch: gate/up (XM 1 shared input, 2 own) GM_GU0-5, down GM_DN0-4, K2 3-8 and 10 +#define DSV41_X3GM_GU(XM, K2, CFG) template __global__ void dsv41_x3gm::gm_kernel<2, XM, K2, CFG>(const dsv41_x3gm::Args); +#define DSV41_X3GM_DN(K2, CFG) template __global__ void dsv41_x3gm::gm_kernel<1, 1, K2, CFG>(const dsv41_x3gm::Args); +// v2 (gm2_kernel: every width in one launch): gate/up GM_GU0 / GM_GU1 at either input mode, down GM_DN0 +template __global__ void dsv41_x3gm::gm2_kernel<2, 1, GM_GU0>(const dsv41_x3gm::Args); +template __global__ void dsv41_x3gm::gm2_kernel<2, 1, GM_GU1>(const dsv41_x3gm::Args); +template __global__ void dsv41_x3gm::gm2_kernel<2, 2, GM_GU0>(const dsv41_x3gm::Args); +template __global__ void dsv41_x3gm::gm2_kernel<2, 2, GM_GU1>(const dsv41_x3gm::Args); +template __global__ void dsv41_x3gm::gm2_kernel<1, 1, GM_DN0>(const dsv41_x3gm::Args); +DSV41_X3GM_GU(1, 3, GM_GU0) +DSV41_X3GM_GU(1, 3, GM_GU1) +DSV41_X3GM_GU(1, 3, GM_GU2) +DSV41_X3GM_GU(1, 3, GM_GU3) +DSV41_X3GM_GU(1, 3, GM_GU4) +DSV41_X3GM_GU(1, 3, GM_GU5) +DSV41_X3GM_GU(2, 3, GM_GU0) +DSV41_X3GM_GU(2, 3, GM_GU1) +DSV41_X3GM_GU(2, 3, GM_GU2) +DSV41_X3GM_GU(2, 3, GM_GU3) +DSV41_X3GM_GU(2, 3, GM_GU4) +DSV41_X3GM_GU(2, 3, GM_GU5) +DSV41_X3GM_DN(3, GM_DN0) +DSV41_X3GM_DN(3, GM_DN1) +DSV41_X3GM_DN(3, GM_DN2) +DSV41_X3GM_DN(3, GM_DN3) +DSV41_X3GM_DN(3, GM_DN4) +DSV41_X3GM_GU(1, 4, GM_GU0) +DSV41_X3GM_GU(1, 4, GM_GU1) +DSV41_X3GM_GU(1, 4, GM_GU2) +DSV41_X3GM_GU(1, 4, GM_GU3) +DSV41_X3GM_GU(1, 4, GM_GU4) +DSV41_X3GM_GU(1, 4, GM_GU5) +DSV41_X3GM_GU(2, 4, GM_GU0) +DSV41_X3GM_GU(2, 4, GM_GU1) +DSV41_X3GM_GU(2, 4, GM_GU2) +DSV41_X3GM_GU(2, 4, GM_GU3) +DSV41_X3GM_GU(2, 4, GM_GU4) +DSV41_X3GM_GU(2, 4, GM_GU5) +DSV41_X3GM_DN(4, GM_DN0) +DSV41_X3GM_DN(4, GM_DN1) +DSV41_X3GM_DN(4, GM_DN2) +DSV41_X3GM_DN(4, GM_DN3) +DSV41_X3GM_DN(4, GM_DN4) +DSV41_X3GM_GU(1, 5, GM_GU0) +DSV41_X3GM_GU(1, 5, GM_GU1) +DSV41_X3GM_GU(1, 5, GM_GU2) +DSV41_X3GM_GU(1, 5, GM_GU3) +DSV41_X3GM_GU(1, 5, GM_GU4) +DSV41_X3GM_GU(1, 5, GM_GU5) +DSV41_X3GM_GU(2, 5, GM_GU0) +DSV41_X3GM_GU(2, 5, GM_GU1) +DSV41_X3GM_GU(2, 5, GM_GU2) +DSV41_X3GM_GU(2, 5, GM_GU3) +DSV41_X3GM_GU(2, 5, GM_GU4) +DSV41_X3GM_GU(2, 5, GM_GU5) +DSV41_X3GM_DN(5, GM_DN0) +DSV41_X3GM_DN(5, GM_DN1) +DSV41_X3GM_DN(5, GM_DN2) +DSV41_X3GM_DN(5, GM_DN3) +DSV41_X3GM_DN(5, GM_DN4) +DSV41_X3GM_GU(1, 6, GM_GU0) +DSV41_X3GM_GU(1, 6, GM_GU1) +DSV41_X3GM_GU(1, 6, GM_GU2) +DSV41_X3GM_GU(1, 6, GM_GU3) +DSV41_X3GM_GU(1, 6, GM_GU4) +DSV41_X3GM_GU(1, 6, GM_GU5) +DSV41_X3GM_GU(2, 6, GM_GU0) +DSV41_X3GM_GU(2, 6, GM_GU1) +DSV41_X3GM_GU(2, 6, GM_GU2) +DSV41_X3GM_GU(2, 6, GM_GU3) +DSV41_X3GM_GU(2, 6, GM_GU4) +DSV41_X3GM_GU(2, 6, GM_GU5) +DSV41_X3GM_DN(6, GM_DN0) +DSV41_X3GM_DN(6, GM_DN1) +DSV41_X3GM_DN(6, GM_DN2) +DSV41_X3GM_DN(6, GM_DN3) +DSV41_X3GM_DN(6, GM_DN4) +DSV41_X3GM_GU(1, 7, GM_GU0) +DSV41_X3GM_GU(1, 7, GM_GU1) +DSV41_X3GM_GU(1, 7, GM_GU2) +DSV41_X3GM_GU(1, 7, GM_GU3) +DSV41_X3GM_GU(1, 7, GM_GU4) +DSV41_X3GM_GU(1, 7, GM_GU5) +DSV41_X3GM_GU(2, 7, GM_GU0) +DSV41_X3GM_GU(2, 7, GM_GU1) +DSV41_X3GM_GU(2, 7, GM_GU2) +DSV41_X3GM_GU(2, 7, GM_GU3) +DSV41_X3GM_GU(2, 7, GM_GU4) +DSV41_X3GM_GU(2, 7, GM_GU5) +DSV41_X3GM_DN(7, GM_DN0) +DSV41_X3GM_DN(7, GM_DN1) +DSV41_X3GM_DN(7, GM_DN2) +DSV41_X3GM_DN(7, GM_DN3) +DSV41_X3GM_DN(7, GM_DN4) +DSV41_X3GM_GU(1, 8, GM_GU0) +DSV41_X3GM_GU(1, 8, GM_GU1) +DSV41_X3GM_GU(1, 8, GM_GU2) +DSV41_X3GM_GU(1, 8, GM_GU3) +DSV41_X3GM_GU(1, 8, GM_GU4) +DSV41_X3GM_GU(1, 8, GM_GU5) +DSV41_X3GM_GU(2, 8, GM_GU0) +DSV41_X3GM_GU(2, 8, GM_GU1) +DSV41_X3GM_GU(2, 8, GM_GU2) +DSV41_X3GM_GU(2, 8, GM_GU3) +DSV41_X3GM_GU(2, 8, GM_GU4) +DSV41_X3GM_GU(2, 8, GM_GU5) +DSV41_X3GM_DN(8, GM_DN0) +DSV41_X3GM_DN(8, GM_DN1) +DSV41_X3GM_DN(8, GM_DN2) +DSV41_X3GM_DN(8, GM_DN3) +DSV41_X3GM_DN(8, GM_DN4) +DSV41_X3GM_GU(1, 10, GM_GU0) +DSV41_X3GM_GU(1, 10, GM_GU1) +DSV41_X3GM_GU(1, 10, GM_GU2) +DSV41_X3GM_GU(1, 10, GM_GU3) +DSV41_X3GM_GU(1, 10, GM_GU4) +DSV41_X3GM_GU(1, 10, GM_GU5) +DSV41_X3GM_GU(2, 10, GM_GU0) +DSV41_X3GM_GU(2, 10, GM_GU1) +DSV41_X3GM_GU(2, 10, GM_GU2) +DSV41_X3GM_GU(2, 10, GM_GU3) +DSV41_X3GM_GU(2, 10, GM_GU4) +DSV41_X3GM_GU(2, 10, GM_GU5) +DSV41_X3GM_DN(10, GM_DN0) +DSV41_X3GM_DN(10, GM_DN1) +DSV41_X3GM_DN(10, GM_DN2) +DSV41_X3GM_DN(10, GM_DN3) +DSV41_X3GM_DN(10, GM_DN4) +#define DSV41_X3GM_ROT(T, M) template __global__ void dsv41_x3gm::rot_kernel(const T*, int, const int*, const half*, \ + const half*, half*, half*, int, int, int); +DSV41_X3GM_ROT(__nv_bfloat16, 1) +DSV41_X3GM_ROT(__nv_bfloat16, 2) +DSV41_X3GM_ROT(half, 1) +DSV41_X3GM_ROT(half, 2) +template __global__ void dsv41_x3gm::dequant_kernel<3>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<4>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<5>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<6>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<7>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<8>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<10>(const uint32_t*, half*, int); +#endif // DSV41_X3GM_NO_INSTANCES diff --git a/zig/kernels/cuda/deepseek_v41/x3gm2.cuh b/zig/kernels/cuda/deepseek_v41/x3gm2.cuh new file mode 100644 index 0000000000..526092524d --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/x3gm2.cuh @@ -0,0 +1,256 @@ +// x3gm.cu's v2 kernel (TF_DSV41_GM_V2): gm2_item, kloop2 and gm2_kernel, inside namespace dsv41_x3gm +#pragma once +// x3gm v2 (TF_DSV41_GM_V2): gm_kernel's items, chains and bits with less issue work (member groups, width per item) + +__device__ __forceinline__ void ldsm2(uint32_t (&a)[2], const void* p) { + const unsigned s = (unsigned)__cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" : "=r"(a[0]), "=r"(a[1]) : "r"(s)); +} + +// The K loop of one item with N8 active n8 member groups (acc[l][0 .. N8 - 1]; the rest stay 0.0, never stored) +template +__device__ __forceinline__ void kloop2(float (&acc)[MTL][NG][4], const unsigned char* smem, int S, int xmat, int mat, + int slice, int lr, int lc, const LaneMap& map, LoadFn&& load_stage) { + using C = Cfg; + constexpr int BM = C::BM, LDA = C::LDA, CPR = C::CPR, TW = C::TW, PAIRS = N8 / 2; +#pragma unroll + for (int s = 0; s < NSA - 1; ++s) { + if (s < S) load_stage(s); + cp_commit(); + } + cp_wait(); + __syncthreads(); + for (int s = 0; s < S; ++s) { + if (s + NSA - 1 < S) load_stage(s + NSA - 1); + cp_commit(); + const unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + const half* A = reinterpret_cast(st) + (size_t)xmat * BM * LDA; + const uint32_t* Wd = reinterpret_cast(st + C::A_BYTES) + (size_t)mat * KS * NB * TW + + slice * MTL * TW; +#pragma unroll + for (int kk = 0; kk < KS; ++kk) { + uint32_t af[MTL][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) decode_a(Wd + (kk * NB + l) * TW, map, af[l]); + const int chunk = swz(lr, kk * 2 + lc); +#pragma unroll + for (int np = 0; np < PAIRS; ++np) { + uint32_t b[4]; + ldsm4(b, A + (16 * np + lr) * LDA + chunk * 8); + const uint32_t b01[2] = {b[0], b[1]}, b23[2] = {b[2], b[3]}; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + mma16816(acc[l][2 * np], af[l], b01); + mma16816(acc[l][2 * np + 1], af[l], b23); + } + } + if constexpr ((N8 & 1) != 0) { // members 16 PAIRS .. + 7 (x2: lanes 0-15's addresses) + uint32_t b[2]; + ldsm2(b, A + (16 * PAIRS + lr) * LDA + chunk * 8); +#pragma unroll + for (int l = 0; l < MTL; ++l) mma16816(acc[l][2 * PAIRS], af[l], b); + } + } + cp_wait(); + __syncthreads(); + } +} + +// One item (pass x 128-column block) at width K2, rows_sh filled by the caller; otherwise gm_kernel's body +template +__device__ __forceinline__ void gm2_item(const Args& a, unsigned char* smem, const int* rows_sh, int e, int cnt, + int nb) { + using C = Cfg; + constexpr int THREADS = C::THREADS, BM = C::BM, LDA = C::LDA, LDE = C::LDE, ER = C::ER, CPR = C::CPR, + TW = C::TW, W = C::W; + constexpr int ACH = XM * BM * CPR, WCH = MATS * KS * NB * K2; // 16-byte chunks a stage: member rows, words + constexpr int APT = (ACH + THREADS - 1) / THREADS, WPT = (WCH + THREADS - 1) / THREADS; + float* ep = reinterpret_cast(smem); // [MATS][ER][LDE] after the K loop (aliases the ring) + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int mat = warp / C::WPM, slice = warp % C::WPM; + const int xmat = XM == 1 ? 0 : mat; + const int g = lane >> 2, t = lane & 3; + const int K = a.K, N = a.N, KT = K >> 4, NTILES = N >> 4, S = KT / KS; + const size_t se = (size_t)KT * NTILES * TW; + const int lr = (lane & 7) + ((lane >> 4) << 3), lc = (lane >> 3) & 1; + const LaneMap map(lane); + const uint32_t* te0 = a.tp0 ? reinterpret_cast(a.tp0[e]) : a.t0 + (size_t)e * se; + const uint32_t* te1 = a.tp1 ? reinterpret_cast(a.tp1[e]) : a.t1 + (size_t)e * se; + const int cnt16 = (cnt + 15) & ~15; // rows past it are never read + + // this thread's chunks: row chunk element offset at k 0 (-1: none, -2: zero fill), word chunk offset at stage 0 + int a_src[APT], w_src[WPT]; +#pragma unroll + for (int j = 0; j < APT; ++j) { + const int i = threadIdx.x + j * THREADS, r = (i / CPR) % BM, c = i % CPR; + int v = -1; + if (i < ACH && r < cnt16) v = rows_sh[r] < 0 ? -2 : rows_sh[r] * K + c * 8; + a_src[j] = v; + } +#pragma unroll + for (int j = 0; j < WPT; ++j) { + const int i = threadIdx.x + j * THREADS, q = i % K2, n = (i / K2) % NB, kk = (i / (K2 * NB)) % KS; + w_src[j] = (kk * NTILES + nb * NB + n) * TW + q * 4; + } + auto load_stage = [&](int s) { + unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + half* ra = reinterpret_cast(st); + uint32_t* rw = reinterpret_cast(st + C::A_BYTES); + const int k0 = s * KS * 16; +#pragma unroll + for (int j = 0; j < APT; ++j) { + const int i = threadIdx.x + j * THREADS, v = a_src[j]; + if (v != -1) { + const int m = i / (BM * CPR), r = (i / CPR) % BM, c = i % CPR; + cp_async16(ra + ((size_t)m * BM + r) * LDA + swz(r, c) * 8, (m ? a.x1 : a.x0) + (v < 0 ? 0 : v + k0), + v >= 0); + } + } + const size_t ws = (size_t)s * KS * NTILES * TW; +#pragma unroll + for (int j = 0; j < WPT; ++j) { + const int i = threadIdx.x + j * THREADS; + if (WCH % THREADS == 0 || i < WCH) { + const int q = i % K2, n = (i / K2) % NB, kk = (i / (K2 * NB)) % KS, m = i / (K2 * NB * KS); + cp_async16(rw + ((m * KS + kk) * NB + n) * TW + q * 4, (m ? te1 : te0) + ws + w_src[j], true); + } + } + }; + + float acc[MTL][NG][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) +#pragma unroll + for (int n = 0; n < NG; ++n) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[l][n][c] = 0.f; + + static_assert(NG <= 16 && NG % 2 == 0, "n8 groups"); + const int n8 = (cnt + 7) >> 3; // >= 1 (a pass has members) +#define GM2_K(G) kloop2(acc, smem, S, xmat, mat, slice, lr, lc, map, load_stage) + switch (n8 <= 8 ? n8 : (n8 + 1) & ~1) { // 1..8, then even counts (NG 16) + case 1: GM2_K(1); break; + case 2: GM2_K(2); break; + case 3: GM2_K(3); break; + case 4: GM2_K(4); break; + case 5: if constexpr (NG >= 6) GM2_K(5); break; + case 6: if constexpr (NG >= 6) GM2_K(6); break; + case 7: if constexpr (NG >= 8) GM2_K(7); break; + case 8: if constexpr (NG >= 8) GM2_K(8); break; + case 10: if constexpr (NG >= 10) GM2_K(10); break; + case 12: if constexpr (NG >= 12) GM2_K(12); break; + case 14: if constexpr (NG >= 14) GM2_K(14); break; + default: if constexpr (NG >= 16) GM2_K(16); break; + } +#undef GM2_K + + // epilogue: gm_kernel's +#pragma unroll + for (int rd = 0; rd < C::ROUNDS; ++rd) { + if (rd) __syncthreads(); + if (ER * rd < cnt) { +#pragma unroll + for (int n = 0; n < NGR; ++n) { + const int ng = rd * NGR + n; + float* E = ep + ((size_t)mat * ER + 8 * n + 2 * t) * LDE + slice * 16 * MTL + g; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + E[l * 16] = acc[l][ng][0]; + E[LDE + l * 16] = acc[l][ng][1]; + E[l * 16 + 8] = acc[l][ng][2]; + E[LDE + l * 16 + 8] = acc[l][ng][3]; + } + } + } + __syncthreads(); + for (int r = warp; r < ER && ER * rd + r < cnt; r += W) { + const int row = rows_sh[ER * rd + r]; + const int c0 = 4 * lane; + const size_t col = (size_t)e * N + nb * 128 + c0; + float v[4]; + *reinterpret_cast(v) = *reinterpret_cast(ep + (size_t)r * LDE + c0); + tf_exl3::fwht128(v, lane); + if constexpr (MATS == 2) { + float u[4]; + *reinterpret_cast(u) = *reinterpret_cast(ep + ((size_t)ER + r) * LDE + c0); + tf_exl3::fwht128(u, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { // upstream gateup_epilogue, ACT_F32 + const float gg = fminf(v[j] * HAD * __half2float(a.sv0[col + j]), a.limit); + const float uu = fminf(fmaxf(u[j] * HAD * __half2float(a.sv1[col + j]), -a.limit), a.limit); + const float act = gg / (1.f + expf(-gg)) * uu; + v[j] = act * __half2float(a.sd[col + j]); + } + tf_exl3::fwht128(v, lane); + half2* o = reinterpret_cast(a.xd + (size_t)row * N + nb * 128 + c0); + o[0] = __halves2half2(__float2half_rn(v[0] * HAD), __float2half_rn(v[1] * HAD)); + o[1] = __halves2half2(__float2half_rn(v[2] * HAD), __float2half_rn(v[3] * HAD)); + } else { + float4 o; // upstream down_epilogue + o.x = v[0] * HAD * __half2float(a.sv0[col + 0]); + o.y = v[1] * HAD * __half2float(a.sv0[col + 1]); + o.z = v[2] * HAD * __half2float(a.sv0[col + 2]); + o.w = v[3] * HAD * __half2float(a.sv0[col + 3]); + *reinterpret_cast(a.y + (size_t)row * N + nb * 128 + c0) = o; + } + } + } +} + +// The widths a v2 launch dispatches and, per width, the ring depth (cut where it would cost a second CTA an SM) +#define GM2_WIDTHS(X) X(3) X(4) X(5) X(6) X(7) X(8) X(10) +constexpr int GM2_NWIDTHS = 7; +template +struct Ring2 { + static constexpr bool fits2(int n) { + return 2 * (Cfg::STAGE * (size_t)n + 1536) <= 102400; + } + static constexpr int cut(int n) { return n <= 2 || fits2(n) ? n : cut(n - 1); } + static constexpr int value = Cfg::MINB == 2 ? cut(NSA) : NSA; +}; +template +struct Smem2 { +#define GM2_SM(K) Cfg::value, NGR>::SMEM, + static constexpr size_t list[GM2_NWIDTHS] = {GM2_WIDTHS(GM2_SM)}; +#undef GM2_SM + static constexpr size_t max_from(int i) { + return i == GM2_NWIDTHS ? 0 : (list[i] > max_from(i + 1) ? list[i] : max_from(i + 1)); + } + static constexpr size_t SMEM = max_from(0); + static constexpr int THREADS = Cfg::THREADS; + static constexpr int MINB = 2 * (SMEM + 1536) <= 102400 && THREADS <= 256 && MTL * NG <= 16 ? 2 : 1; +}; + +// One v2 launch over the passes of every width (uniform stack: one width; ragged: each expert its own, a.k2e) +template +__global__ void __launch_bounds__(Smem2::THREADS, + Smem2::MINB) gm2_kernel(const Args a) { + constexpr int BM = NG * 8, THREADS = Smem2::THREADS; + extern __shared__ __align__(16) unsigned char smem[]; + __shared__ int rows_sh[BM]; + __shared__ int item_sh; + const int NBLK = a.N >> 7, total = a.npass[0] * NBLK; + int next = blockIdx.x; + for (;;) { + __syncthreads(); // the last item is done with rows_sh and shared memory + if (threadIdx.x == 0) item_sh = a.ticket ? atomicAdd(a.ticket, 1) : next; + next += gridDim.x; + __syncthreads(); + const int item = item_sh; + if (item >= total) break; + const int tt = item / NBLK, nb = item % NBLK; + const int e = a.pe[tt], off = a.poff[tt], cnt = a.pcnt[tt]; + for (int i = threadIdx.x; i < BM; i += THREADS) rows_sh[i] = i < cnt ? a.order[off + i] : -1; + __syncthreads(); + switch (a.k2e[e]) { // block-uniform +#define GM2_CASE(K) \ + case K: \ + gm2_item::value, NGR>(a, smem, rows_sh, e, \ + cnt, nb); \ + break; + GM2_WIDTHS(GM2_CASE) +#undef GM2_CASE + default: break; // never: the host checks every width is in K2S + } + } +} diff --git a/zig/kernels/cuda/deepseek_v41/x3gm_plan.cu b/zig/kernels/cuda/deepseek_v41/x3gm_plan.cu new file mode 100644 index 0000000000..ee7355a061 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/x3gm_plan.cu @@ -0,0 +1,104 @@ +// x3gm's pass tables on the device in one launch: same integers as the torch ops of x3gm.plan. + +// order = stable argsort(pick); cnt/off/npe/pend per expert; pass t -> expert e, poff, pcnt; npass = pend[E - 1]. + +// One CTA (E < 512): histogram, scan, stable chunked scatter (per-warp counts + match_any); integers only, so exact. + +#include + +namespace dsv41_x3gm_plan { + +constexpr int NT = 512, NW = NT / 32, EMAX = 512; // E <= 511: the routers take <= 512 experts (GE_MAX) + +__global__ void __launch_bounds__(NT) plan_kernel(const int* __restrict__ pick, int P, int E, int bm, + int* __restrict__ order, int* __restrict__ pe, + int* __restrict__ poff, int* __restrict__ pcnt, + int* __restrict__ npass) { + __shared__ int cnt[EMAX + 1]; // counts of values 0..E, then (scatter) the moving bases + __shared__ int pend[EMAX]; // inclusive scan of passes an expert + __shared__ int wc[NW][EMAX + 1]; // a chunk's per-warp counts of each value (scatter) + __shared__ int scan_tmp[NT]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, V = E + 1; + + for (int v = tid; v < V; v += NT) cnt[v] = 0; + __syncthreads(); + for (int i = tid; i < P; i += NT) atomicAdd(&cnt[pick[i]], 1); // counts: order-free + __syncthreads(); + + // exclusive scans: bases of values 0..E and pend over the first E; V <= NT * 2, two elements a thread + auto block_scan = [&](int x) { // inclusive scan of one value a thread + scan_tmp[tid] = x; + __syncthreads(); + for (int o = 1; o < NT; o <<= 1) { + const int y = tid >= o ? scan_tmp[tid - o] : 0; + __syncthreads(); + scan_tmp[tid] += y; + __syncthreads(); + } + const int r = scan_tmp[tid]; + __syncthreads(); + return r; + }; + int c0 = 2 * tid < V ? cnt[2 * tid] : 0, c1 = 2 * tid + 1 < V ? cnt[2 * tid + 1] : 0; + const int b_end = block_scan(c0 + c1); // inclusive over pairs of values + const int n0 = 2 * tid < E ? (c0 + bm - 1) / bm : 0, n1 = 2 * tid + 1 < E ? (c1 + bm - 1) / bm : 0; + const int p_end = block_scan(n0 + n1); + __syncthreads(); + if (2 * tid < V) cnt[2 * tid] = b_end - c0 - c1; // base of value 2 tid + if (2 * tid + 1 < V) cnt[2 * tid + 1] = b_end - c1; + if (2 * tid < E) pend[2 * tid] = p_end - n1; + if (2 * tid + 1 < E) pend[2 * tid + 1] = p_end; + __syncthreads(); + + // pass tables (counts recomputed from the bases: base[v + 1] - base[v], base[E + 1] = P) + const int tmax = P / bm + E + 1; + const int npe_last = E > 0 ? pend[E - 1] : 0; + for (int t = tid; t < tmax; t += NT) { + int lo = 0, hi = E; // searchsorted(pend, t, right): first pend > t + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (pend[mid] <= t) lo = mid + 1; else hi = mid; + } + const int e = lo < E - 1 ? lo : E - 1; + const int ce = (e + 1 < V ? cnt[e + 1] : P) - cnt[e]; + const int npe_e = (ce + bm - 1) / bm; + const long long j = (long long)t - (pend[e] - npe_e); + long long o = (long long)cnt[e] + j * bm; + const long long cap = P - 1 > 0 ? P - 1 : 0; + if (o > cap) o = cap; + long long c = (long long)ce - j * bm; + c = c < 0 ? 0 : (c > bm ? bm : c); + pe[t] = e; + poff[t] = (int)o; + pcnt[t] = (int)c; + } + if (tid == 0) npass[0] = npe_last; + __syncthreads(); + + // stable scatter, NT pairs a chunk in index order + for (int base = 0; base < P; base += NT) { + for (int i = tid; i < NW * V; i += NT) (&wc[0][0])[(i / V) * (EMAX + 1) + i % V] = 0; + __syncthreads(); + const int i = base + tid; + const bool live = i < P; + const int v = live ? pick[i] : -1; + const unsigned same = __match_any_sync(0xffffffffu, v); + const int before = __popc(same & ((1u << lane) - 1u)); // equal values in earlier lanes of the warp + if (live && before == 0) wc[warp][v] = __popc(same); // the warp's count of v, by its first lane + __syncthreads(); + if (live) { + int slot = cnt[v] + before; + for (int w = 0; w < warp; ++w) slot += wc[w][v]; + order[slot] = i; + } + __syncthreads(); + for (int u = tid; u < V; u += NT) { // the bases move past this chunk's pairs + int s = 0; + for (int w = 0; w < NW; ++w) s += wc[w][u]; + cnt[u] += s; + } + __syncthreads(); + } +} + +} // namespace dsv41_x3gm_plan diff --git a/zig/kernels/cuda/deepseek_v41/x3ld.cu b/zig/kernels/cuda/deepseek_v41/x3ld.cu new file mode 100644 index 0000000000..15bdfb4fcc --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/x3ld.cu @@ -0,0 +1,249 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/x3ld.cu (blob 5cd25f2cdf3c): sass_check.sh + + + +// DeepSeek-V4.1-Flash routed experts (TF_DSV41_EXPERT_LOADS=1): grouped EXL3 expert GEMV with a faster load path. + +// Same bits as grouped_kernel at any row count; data movement changes (prologue, ring, PDL). + +// PROBE 3 (timing only, wrong results): no decode and no mma, the load path alone. + +#include "experts_grouped.cuh" + +namespace dsv41_x3ld { + +using tf_exl3x::decode_tile; +using tf_exl3x::Fmt; +using tf_exl3x::LaneMap; +using tf_exl3x::mma16816; + +constexpr int W = 4; // warps a CTA: upstream's default_config, which fixes the K ranges + +__device__ __forceinline__ uint4 ldg_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; +} + +// load_pair's value (ok ? *(const uint32_t*)x : 0), the load pinned where it stands +__device__ __forceinline__ uint32_t ldg_pair(const half* x, bool ok) { + uint32_t v = 0u; + if (ok) asm volatile("ld.global.nc.u32 %0, [%1];" : "=r"(v) : "l"(x)); + return v; +} + +__device__ __forceinline__ void pdl_wait() { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif +} + +template +struct Step { + static constexpr int TW = Fmt::TW; // words a tile + static constexpr int LW = Fmt::LW; // words a lane a tile (load_words) + static constexpr int WORDS = NT * TW; // a warp's k step, contiguous + static constexpr int NV = (WORDS + 127) / 128; // 16-byte loads a lane a step +}; + +// One warp's k tiles [kt0, kt0 + nkt) into acc (warp_tiles' arithmetic, loads PD ahead); all CTA threads call it. +template +__device__ __forceinline__ void ld_tiles(const uint32_t* __restrict__ T, int NTILES, int kt0, int nkt, int nt0, + const half* __restrict__ X, int K, const int* rows_sh, uint32_t* stage, + int lane, float (&acc)[NT][2][4]) { + using S = Step; + const int g = lane >> 2, t = lane & 3; + const size_t kstride = (size_t)NTILES * S::TW; + const uint32_t* tp = T + ((size_t)kt0 * NTILES + nt0) * S::TW; + + uint4 ring[PD][S::NV]; + auto issue_w = [&](int d, int s) { + const uint32_t* src = tp + (size_t)s * kstride; +#pragma unroll + for (int v = 0; v < S::NV; ++v) { + const int off = 128 * v + 4 * lane; + if (off < S::WORDS) ring[d][v] = ldg_v4(src + off); + } + }; +#pragma unroll + for (int d = 0; d < PD; ++d) issue_w(d, d); // PD <= nkt (host) + + pdl_wait(); // X (and Z) belong to the kernels before + __syncthreads(); // rows_sh + const int r0 = rows_sh[g], r1 = rows_sh[g + 8]; + const half* x0 = X + (size_t)(r0 < 0 ? 0 : r0) * K + 2 * t; + const half* x1 = X + (size_t)(r1 < 0 ? 0 : r1) * K + 2 * t; + uint32_t a_ring[PD][4]; + auto issue_a = [&](int d, int s) { + const int k = (kt0 + s) * 16; + a_ring[d][0] = ldg_pair(x0 + k, r0 >= 0); + a_ring[d][1] = ldg_pair(x1 + k, r1 >= 0); + a_ring[d][2] = ldg_pair(x0 + k + 8, r0 >= 0); + a_ring[d][3] = ldg_pair(x1 + k + 8, r1 >= 0); + }; +#pragma unroll + for (int d = 0; d < PD; ++d) issue_a(d, d); + + const LaneMap map(lane); + for (int kb = 0; kb < nkt; kb += PD) { // nkt % PD == 0 (host) +#pragma unroll + for (int d = 0; d < PD; ++d) { + const int s = kb + d; + const bool more = s + PD < nkt; + __syncwarp(); // every lane has read step s - 1 back +#pragma unroll + for (int v = 0; v < S::NV; ++v) { + const int off = 128 * v + 4 * lane; + if (off < S::WORDS) *reinterpret_cast(stage + off) = ring[d][v]; + } + if (more) issue_w(d, s + PD); // the slot's registers are free again + __syncwarp(); + uint32_t w[NT][S::LW]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int l = 0; l < S::LW; ++l) + w[i][l] = (S::TW % 32 == 0 || l * 32 + lane < S::TW) ? stage[i * S::TW + l * 32 + lane] : 0u; + uint32_t a[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) a[j] = a_ring[d][j]; + if (more) issue_a(d, s + PD); +#pragma unroll + for (int i = 0; i < NT; ++i) { + if constexpr (PROBE == 3) { + uint32_t x = a[0] ^ a[1] ^ a[2] ^ a[3]; +#pragma unroll + for (int l = 0; l < S::LW; ++l) x ^= w[i][l]; + acc[i][0][0] = __uint_as_float(__float_as_uint(acc[i][0][0]) ^ x); + } else { + uint32_t b0[2], b1[2]; + decode_tile(w[i], map, lane, b0, b1); + mma16816(acc[i][0], a, b0); + mma16816(acc[i][1], a, b1); + } + } + } + } + __syncwarp(); // the staging area becomes red[warp] again +} + +// Program (u, n block, (mat * SK + split) * MT + member tile): grouped_kernel's work item. +template +__global__ void __launch_bounds__(W * 32, NT == 8 ? 3 : 4) ld_kernel( + const half* __restrict__ X0, const half* __restrict__ X1, const int64_t* __restrict__ TP0, + const int64_t* __restrict__ TP1, const int* __restrict__ K2_0, const int* __restrict__ K2_1, + const int* __restrict__ uids, const int* __restrict__ ucount, const int* __restrict__ members, + float* __restrict__ Z, int K, int N, int P, int SK, int maxm, int slots) { + static_assert(PD >= 1 && PD <= 4, "PD 1-4"); + const int u = blockIdx.x; + const int MT = (maxm + 15) / 16; + const int mtile = blockIdx.z % MT; + const int split = (blockIdx.z / MT) % SK; + const int mat = blockIdx.z / MT / SK; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + + // round trip 1: none of these depends on another (uids / members are sized by the grid: always in bounds) + const int cnt = __ldcg(ucount); + const int e = __ldcg(uids + u); + const int m = mtile * 16 + threadIdx.x; + const int code = threadIdx.x < 16 && m < maxm ? __ldcg(members + u * maxm + m) : -1; + const int first = __ldcg(members + u * maxm + mtile * 16); + if (u >= cnt || first < 0) return; // grouped_kernel's exits (CTA-uniform) + + __shared__ int rows_sh[16]; + __shared__ __align__(16) float red[W][16][NT * 16]; + if (threadIdx.x < 16) rows_sh[threadIdx.x] = code >= 0 ? (code >> 5) * slots + (code & 31) : -1; + + // round trip 2: the expert's trellis and width + const uint32_t* T = reinterpret_cast(mat ? TP1[e] : TP0[e]); + const int k2 = mat ? K2_1[e] : K2_0[e]; + const half* X = mat ? X1 : X0; + const int KT = K >> 4, NTILES = N >> 4; + const int per_split = KT / SK, per_warp = per_split / W; + const int kt0 = split * per_split + warp * per_warp; + const int nt0 = blockIdx.y * NT; + uint32_t* stage = reinterpret_cast(&red[warp][0][0]); + + 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 DSV41_X3LD_CASE(K2_) \ + case K2_: \ + if constexpr (K2_ >= LO && K2_ <= HI) \ + ld_tiles(T, NTILES, kt0, per_warp, nt0, X, K, rows_sh, stage, lane, acc); \ + else \ + __trap(); \ + break; + DSV41_X3LD_CASE(2) + DSV41_X3LD_CASE(3) + DSV41_X3LD_CASE(4) + DSV41_X3LD_CASE(5) + DSV41_X3LD_CASE(6) + DSV41_X3LD_CASE(7) + DSV41_X3LD_CASE(8) + DSV41_X3LD_CASE(9) + DSV41_X3LD_CASE(10) + DSV41_X3LD_CASE(11) + DSV41_X3LD_CASE(12) + DSV41_X3LD_CASE(13) + DSV41_X3LD_CASE(14) + DSV41_X3LD_CASE(15) + DSV41_X3LD_CASE(16) +#undef DSV41_X3LD_CASE + default: + __trap(); + } + + // grouped_kernel's reduction, verbatim: the warps' partial sums through shared memory, added in warp order + const int g = lane >> 2, t = lane & 3; +#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]; + Z[(((size_t)mat * SK + split) * P + r) * N + nt0 * 16 + col] = s; + } +} + +} // namespace dsv41_x3ld + +// The instantiations dsv41_x3ld_grouped_cuda dispatches: mul1, (nt, pd) x probe x the (lo, hi) ranges. +#define DSV41_X3LD_INST(NT, PD, LO, HI, PROBE) template __global__ void dsv41_x3ld::ld_kernel<2, NT, PD, LO, HI, PROBE>( \ + const half*, const half*, const int64_t*, const int64_t*, const int*, const int*, const int*, const int*, \ + const int*, float*, int, int, int, int, int, int); +DSV41_X3LD_INST(8, 1, 8, 8, 0) +DSV41_X3LD_INST(8, 1, 2, 10, 0) +DSV41_X3LD_INST(8, 1, 2, 12, 0) +DSV41_X3LD_INST(8, 2, 8, 8, 0) +DSV41_X3LD_INST(8, 2, 2, 10, 0) +DSV41_X3LD_INST(8, 2, 2, 12, 0) +DSV41_X3LD_INST(4, 2, 8, 8, 0) +DSV41_X3LD_INST(4, 2, 2, 10, 0) +DSV41_X3LD_INST(4, 2, 2, 12, 0) +DSV41_X3LD_INST(8, 1, 8, 8, 3) +DSV41_X3LD_INST(8, 1, 2, 10, 3) +DSV41_X3LD_INST(8, 1, 2, 12, 3) +DSV41_X3LD_INST(8, 2, 8, 8, 3) +DSV41_X3LD_INST(8, 2, 2, 10, 3) +DSV41_X3LD_INST(8, 2, 2, 12, 3) diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index b784479964..fe21a66fc4 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -39,6 +39,8 @@ pub const DeviceAttribute = enum(c_int) { compute_capability_minor = 76, max_shared_memory_per_multiprocessor = 81, max_shared_memory_per_block_optin = 97, + virtual_memory_management_supported = 102, + gpu_direct_rdma_with_cuda_vmm_supported = 110, }; pub const FunctionAttribute = enum(c_int) { @@ -242,3 +244,10 @@ comptime { std.debug.assert(@sizeOf(KernelNodeParams) == 72 and @offsetOf(KernelNodeParams, "params") == 40); std.debug.assert(@sizeOf(ExecUpdateResultInfo) == 24); } + +test "CUDA device attribute matches the RDMA VMM capability ABI" { + // CUDA's CUdevice_attribute assigns 110 to RDMA VMM support; 126 is the memory sync domain count. + try @import("std").testing.expectEqual(@as(c_int, 110), @backingInt( + DeviceAttribute.gpu_direct_rdma_with_cuda_vmm_supported, + )); +} diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig new file mode 100644 index 0000000000..de3da5718c --- /dev/null +++ b/zig/src/cuda/arena.zig @@ -0,0 +1,405 @@ +//! A device sub-allocator under DeviceBuffer.alloc: buffers carved from a few large chunks, not one cuMemAlloc each. + +const std = @import("std"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const memory = @import("memory.zig"); + +pub const Error = error{ DriverUnavailable, MissingSymbol, Unsupported, CudaFailed, OutOfMemory }; + +/// Where chunks come from: `chunk` maps `len` bytes (a multiple of `granularity`), `release` gives a chunk back. +pub const Backend = struct { + ctx: *anyopaque, + granularity: usize, + chunk: *const fn (ctx: *anyopaque, len: usize) ?u64, + release: *const fn (ctx: *anyopaque, base: u64, len: usize) void, + name: []const u8, +}; + +pub const Options = struct { + small_max: usize = 4 << 20, // requests up to this many bytes share chunks + small_chunk: usize = 64 << 20, // a shared chunk's bytes, rounded up to the granularity + alignment: usize = 512, // every range starts at a multiple of this +}; + +pub const Stats = struct { small: u64 = 0, large: u64 = 0, reused: u64 = 0, fallback: u64 = 0, chunk_bytes: u64 = 0 }; + +/// Small requests bump 512-byte aligned through shared chunks and are reused by size; large ones get their own chunk. +pub const Arena = struct { + gpa: std.mem.Allocator, + be: Backend, + o: Options, + mutex: std.atomic.Mutex = .unlocked, + small: std.ArrayList([2]u64) = .empty, // shared chunks (base, len) + cursor: u64 = 0, // the current shared chunk's next free byte + limit: u64 = 0, + free_small: std.AutoHashMapUnmanaged(u64, std.ArrayList(u64)) = .empty, // freed small ranges by rounded size + live: std.AutoHashMapUnmanaged(u64, u64) = .empty, // live small ranges: address to rounded size + large: std.AutoHashMapUnmanaged(u64, u64) = .empty, // large chunks: address to mapped size + dead: [2]u32 = .{ 0, 0 }, // removals from `live` and `large` since each was last rehashed + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, be: Backend, o: Options) Arena { + return .{ .gpa = gpa, .be = be, .o = o }; + } + + pub fn deinit(a: *Arena) void { + var it = a.large.iterator(); + while (it.next()) |e| a.be.release(a.be.ctx, e.key_ptr.*, e.value_ptr.*); + for (a.small.items) |c| a.be.release(a.be.ctx, c[0], c[1]); + var f = a.free_small.valueIterator(); + while (f.next()) |l| l.deinit(a.gpa); + a.free_small.deinit(a.gpa); + a.live.deinit(a.gpa); + a.large.deinit(a.gpa); + a.small.deinit(a.gpa); + a.* = undefined; + } + + fn lock(a: *Arena) void { + while (!a.mutex.tryLock()) std.Thread.yield() catch {}; + } + + fn roundUp(x: usize, m: usize) usize { + return std.mem.alignForward(usize, x, m); + } + + /// An address for `len` bytes, or null (the caller then allocates as before). + pub fn alloc(a: *Arena, len: usize) ?u64 { + if (len == 0) return null; + a.lock(); + defer a.mutex.unlock(); + return a.allocLocked(len) catch null; + } + + fn allocLocked(a: *Arena, len: usize) !?u64 { + if (len > a.o.small_max) { + const n = roundUp(len, a.be.granularity); + const p = a.be.chunk(a.be.ctx, n) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, p, n); + try a.large.put(a.gpa, p, n); + a.stats.large += 1; + a.stats.chunk_bytes += n; + return p; + } + const n: u64 = roundUp(len, a.o.alignment); + try a.live.ensureUnusedCapacity(a.gpa, 1); // before a range is taken: a failure leaves free list and cursor alone + if (a.free_small.getPtr(n)) |l| if (l.pop()) |p| { + a.live.putAssumeCapacityNoClobber(p, n); + a.stats.reused += 1; + return p; + }; + if (a.limit - a.cursor < n) { + const cl = roundUp(@max(a.o.small_chunk, n), a.be.granularity); + const base = a.be.chunk(a.be.ctx, cl) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, base, cl); + try a.small.append(a.gpa, .{ base, cl }); + a.cursor = base; + a.limit = base + cl; + a.stats.chunk_bytes += cl; + } + const p = a.cursor; + a.live.putAssumeCapacityNoClobber(p, n); + a.cursor += n; + a.stats.small += 1; + return p; + } + + /// True when `p` was the arena's (and is now free); false: not ours, the caller frees it as before. + pub fn free(a: *Arena, p: u64) bool { + a.lock(); + defer a.mutex.unlock(); + if (a.large.fetchRemove(p)) |e| { + removed(&a.large, &a.dead[1]); + a.be.release(a.be.ctx, p, e.value); + return true; + } + const e = a.live.fetchRemove(p) orelse return false; + removed(&a.live, &a.dead[0]); + const g = a.free_small.getOrPut(a.gpa, e.value) catch return true; // leaked range, still ours + if (!g.found_existing) g.value_ptr.* = .empty; + g.value_ptr.append(a.gpa, p) catch {}; + return true; + } + + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + fn removed(map: *std.AutoHashMapUnmanaged(u64, u64), dead: *u32) void { + dead.* += 1; + if (@as(u64, dead.*) * 10 < map.capacity()) return; + map.rehash(std.hash_map.AutoContext(u64){}); + dead.* = 0; + } + + // -- memory.zig's hook ---------------------------------------------------------------------------------------- + + fn hookAlloc(ctx: *anyopaque, len: usize) ?u64 { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.alloc(len); + } + + fn hookFree(ctx: *anyopaque, p: u64) bool { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.free(p); + } + + /// Every later DeviceBuffer.alloc comes from the arena, which must outlive its buffers; only addresses change. + pub fn install(a: *Arena) void { + memory.setHook(.{ .ctx = a, .alloc = hookAlloc, .free = hookFree }); + } +}; + +// -- backends --------------------------------------------------------------------------------------------------------- + +/// One cuMemAlloc a chunk (the driver's own 2 MiB-aligned allocations). +pub const Plain = struct { + d: *const Driver, + + pub fn backend(p: *Plain) Backend { + return .{ .ctx = p, .granularity = 2 << 20, .chunk = chunk, .release = release, .name = "plain" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const p: *Plain = @ptrCast(@alignCast(ctx)); + var q: abi.DevicePtr = 0; + if (p.d.api.cuMemAlloc_v2(&q, len) != abi.success) return null; + return q; + } + + fn release(ctx: *anyopaque, base: u64, _: usize) void { + const p: *Plain = @ptrCast(@alignCast(ctx)); + _ = p.d.api.cuMemFree_v2(base); + } +}; + +/// CUDA's virtual memory management: one reserved virtual range, chunks of the recommended granularity mapped in order. +pub const Vmm = struct { + lib: std.DynLib, + device: abi.Device, + prop: Prop, + granularity: usize, + va: u64 = 0, + va_len: usize = 0, + next: u64 = 0, + api: Api, + + const Api = struct { + reserve: *const fn (*u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + addressFree: *const fn (u64, usize) callconv(.c) abi.Result, + create: *const fn (*u64, usize, *const Prop, c_ulonglong) callconv(.c) abi.Result, + release: *const fn (u64) callconv(.c) abi.Result, + map: *const fn (u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + unmap: *const fn (u64, usize) callconv(.c) abi.Result, + setAccess: *const fn (u64, usize, *const AccessDesc, usize) callconv(.c) abi.Result, + granularity: *const fn (*usize, *const Prop, c_int) callconv(.c) abi.Result, + }; + + pub const Prop = extern struct { // CUmemAllocationProp + type: c_int = 1, // CU_MEM_ALLOCATION_TYPE_PINNED + handle_types: c_int = 0, + loc_type: c_int = 1, // CU_MEM_LOCATION_TYPE_DEVICE + loc_id: c_int = 0, + win32: ?*anyopaque = null, + compression: u8 = 0, + rdma: u8 = 0, + usage: u16 = 0, + reserved: [4]u8 = .{ 0, 0, 0, 0 }, + }; + const AccessDesc = extern struct { loc_type: c_int = 1, loc_id: c_int, flags: c_int = 3 }; // CUmemAccessDesc, read-write + + comptime { + std.debug.assert(@sizeOf(Prop) == 32 and @offsetOf(Prop, "compression") == 24 and @sizeOf(AccessDesc) == 12); + } + + /// `va_bytes` of virtual addresses reserved (cheap: no memory behind them until a chunk maps). + pub fn open(d: *const Driver, device: abi.Device, va_bytes: usize) Error!Vmm { + var vmm_ok: c_int = 0; + if (d.api.cuDeviceGetAttribute(&vmm_ok, .virtual_memory_management_supported, device) != abi.success or vmm_ok == 0) return error.Unsupported; // VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + var api: Api = undefined; + const names = .{ .{ "reserve", "cuMemAddressReserve" }, .{ "addressFree", "cuMemAddressFree" }, .{ "create", "cuMemCreate" }, .{ "release", "cuMemRelease" }, .{ "map", "cuMemMap" }, .{ "unmap", "cuMemUnmap" }, .{ "setAccess", "cuMemSetAccess" }, .{ "granularity", "cuMemGetAllocationGranularity" } }; + inline for (names) |n| @field(api, n[0]) = lib.lookup(@FieldType(Api, n[0]), n[1]) orelse return error.MissingSymbol; + var prop: Prop = .{ .loc_id = device }; + var rdma: c_int = 0; // GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED: keep the memory registrable as cuMemAlloc's is + if (d.api.cuDeviceGetAttribute(&rdma, .gpu_direct_rdma_with_cuda_vmm_supported, device) == abi.success and rdma != 0) prop.rdma = 1; + var g: usize = 0; + if (api.granularity(&g, &prop, 1) != abi.success or g == 0) return error.Unsupported; // RECOMMENDED + g = @max(g, 2 << 20); + const len = std.mem.alignForward(usize, va_bytes, g); + var va: u64 = 0; + if (api.reserve(&va, len, g, 0, 0) != abi.success) return error.CudaFailed; + return .{ .lib = lib, .device = device, .prop = prop, .granularity = g, .va = va, .va_len = len, .next = va, .api = api }; + } + + pub fn close(v: *Vmm) void { + _ = v.api.addressFree(v.va, v.va_len); + v.lib.close(); + } + + pub fn backend(v: *Vmm) Backend { + return .{ .ctx = v, .granularity = v.granularity, .chunk = chunk, .release = release, .name = "vmm" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + if (v.va + v.va_len - v.next < len) return null; + var h: u64 = 0; + if (v.api.create(&h, len, &v.prop, 0) != abi.success) return null; + defer _ = v.api.release(h); // the mapping holds the memory until it is unmapped + const at = v.next; + if (v.api.map(at, len, 0, h, 0) != abi.success) return null; + const desc: AccessDesc = .{ .loc_id = v.device }; + if (v.api.setAccess(at, len, &desc, 1) != abi.success) { + _ = v.api.unmap(at, len); + return null; + } + v.next += len; + return at; + } + + fn release(ctx: *anyopaque, base: u64, len: usize) void { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + _ = v.api.unmap(base, len); // the virtual range is not reused: the reservation is sized for the process + } +}; + +// -- tests (a fake backend: addresses only) --------------------------------------------------------------------------- + +const Fake = struct { + next: u64 = 1 << 40, + chunks: u32 = 0, + released: u32 = 0, + fail: bool = false, + + fn backend(f: *Fake, g: usize) Backend { + return .{ .ctx = f, .granularity = g, .chunk = chunk, .release = release, .name = "fake" }; + } + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail) return null; + const p = f.next; + f.next += len; + f.chunks += 1; + return p; + } + fn release(ctx: *anyopaque, _: u64, _: usize) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.released += 1; + } +}; + +test "arena: small requests pack 512-aligned into shared chunks, are reused by size, large ones get their own" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + const p1 = a.alloc(20480).?; // a norm weight + const p2 = a.alloc(983040).?; // an mHC fn16 + const p3 = a.alloc(100).?; + try t.expectEqual(p1 + 20480, p2); + try t.expectEqual(p2 + 983040, p3); + for ([_]u64{ p1, p2, p3 }) |p| try t.expectEqual(@as(u64, 0), p % 512); + try t.expectEqual(@as(u32, 1), f.chunks); + // a freed range comes back for the same rounded size + try t.expect(a.free(p3)); + try t.expectEqual(p3, a.alloc(17).?); + try t.expectEqual(@as(u64, 1), a.stats.reused); + // three more 1 MiB ranges fit the 4 MiB chunk, the fourth opens a new shared chunk + for (0..3) |_| _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 1), f.chunks); + _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 2), f.chunks); + // large: its own chunk, rounded to the granularity, released when freed + const big = a.alloc(5 << 20).?; + try t.expectEqual(@as(u32, 3), f.chunks); + try t.expectEqual(@as(u64, 6 << 20), a.large.get(big).?); + try t.expect(a.free(big)); + try t.expectEqual(@as(u32, 1), f.released); + // not ours + try t.expect(!a.free(12345)); + // a backend that cannot map: null, the caller allocates as before + f.fail = true; + try t.expectEqual(@as(?u64, null), a.alloc(64 << 20)); + try t.expectEqual(@as(u64, 1), a.stats.fallback); +} + +test "arena: under DeviceBuffer.alloc through memory.setHook" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{}); + defer a.deinit(); + a.install(); + defer memory.setHook(null); + const h = memory.currentHook().?; + const p = h.alloc(h.ctx, 4096).?; + try t.expect(h.free(h.ctx, p)); + try t.expect(!h.free(h.ctx, 7)); +} + +/// Slots of `m` that hold neither a key nor a tombstone. +fn freeSlots(m: *const std.AutoHashMapUnmanaged(u64, u64)) usize { + const meta = m.metadata orelse return 0; + var n: usize = 0; + for (meta[0..m.capacity()]) |x| n += @intFromBool(x.isFree()); + return n; +} + +test "arena: a long run of allocations and frees leaves both address maps free slots" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + var held: std.ArrayList(u64) = .empty; + defer held.deinit(t.allocator); + var rng = std.Random.DefaultPrng.init(3); + for (0..20000) |_| { + // a serving process's buffers come and go (prompt workspaces, grown buffers): small and large alike + if (held.items.len > 0 and rng.random().boolean()) { + const k = rng.random().uintLessThan(usize, held.items.len); + try t.expect(a.free(held.swapRemove(k))); + } else { + const big = rng.random().uintLessThan(u32, 4) == 0; + try held.append(t.allocator, a.alloc(if (big) 3 << 20 else 512 * (1 + rng.random().uintLessThan(u64, 64))).?); + } + for ([_]*const std.AutoHashMapUnmanaged(u64, u64){ &a.live, &a.large }) |m| { + if (m.capacity() == 0) continue; + try t.expect(freeSlots(m) * 10 >= m.capacity()); + } + } +} + +test "arena: a reuse whose address-map insertion fails keeps its range free, and the retry gets it" { + const t = std.testing; + var f: Fake = .{}; + var fa = t.FailingAllocator.init(t.allocator, .{}); + { + var a = Arena.init(fa.allocator(), f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); // testing.allocator underneath reports any leaked host bytes + const p = a.alloc(4096).?; + try t.expect(a.free(p)); + // other sizes take back the room the free gave: inserting the reused range must grow the map + while (a.live.available > 0) _ = a.alloc(512).?; + const live = a.live.count(); + const cursor = a.cursor; + const chunks = f.chunks; + fa.fail_index = fa.alloc_index; // the next allocation fails, and every one after it until healed + try t.expectEqual(@as(?u64, null), a.alloc(4096)); + try t.expectEqual(live, a.live.count()); + try t.expectEqual(cursor, a.cursor); + try t.expectEqual(chunks, f.chunks); + try t.expectEqualSlices(u64, &.{p}, a.free_small.get(4096).?.items); + try t.expectEqual(@as(u64, 0), a.stats.reused); + fa.fail_index = std.math.maxInt(usize); + try t.expectEqual(p, a.alloc(4096).?); + try t.expectEqual(@as(u64, 1), a.stats.reused); + try t.expectEqual(@as(usize, 0), a.free_small.get(4096).?.items.len); + try t.expect(a.free(p)); + } + try t.expectEqual(f.chunks, f.released); // every chunk went back to the backend +} diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 7c7fb7f068..38a9b930d7 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -9,6 +9,10 @@ pub const Error = error{ CopySourceLimit, StageBusy, StageUnprepared, UnownedCap pub const Driver = struct { lib: platform.DynLib, api: abi.Api, + /// cuFuncGetName (CUDA 12.3+; optional: an older driver opens without it): a failed launch names its kernel + func_name: ?FuncGetName = null, + + pub const FuncGetName = *const fn (*?[*:0]const u8, abi.Function) callconv(.c) abi.Result; pub fn open() Error!Driver { return openPath(platform.cuda_driver); @@ -26,7 +30,7 @@ pub const Driver = struct { return error.MissingSymbol; }; } - const d: Driver = .{ .lib = lib, .api = api }; + const d: Driver = .{ .lib = lib, .api = api, .func_name = lib.lookup(FuncGetName, "cuFuncGetName") }; try d.check(api.cuInit(0), "cuInit"); return d; } @@ -43,18 +47,33 @@ pub const Driver = struct { pub fn check(self: *const Driver, res: abi.Result, what: []const u8) Error!void { if (res == abi.success) return; if (res == abi.error_not_ready) return error.NotReady; - std.log.err("{s}: {s} ({d}) {s}", .{ what, self.errorName(res), res, self.errorText(res) }); + std.log.err("{s}: {s} ({d}) {s}{s}", .{ what, self.errorName(res), res, self.errorText(res), shutdownNote(res) }); + return errorOf(res); + } + + /// The error a failed call returns (`check` without its log line, for callers that log more). + pub fn errorOf(res: abi.Result) Error { return switch (res) { 2 => error.OutOfDeviceMemory, abi.error_not_found => error.NotFound, + abi.error_not_ready => error.NotReady, else => error.CudaFailed, }; } + /// The driver's name for the code, else knownName's: during the exit's teardown cuGetErrorName fails too. pub fn errorName(self: *const Driver, res: abi.Result) []const u8 { var s: ?[*:0]const u8 = null; - if (self.api.cuGetErrorName(res, &s) != abi.success) return "CUDA_ERROR_UNKNOWN_CODE"; - return if (s) |p| std.mem.span(p) else "CUDA_ERROR_UNKNOWN_CODE"; + if (self.api.cuGetErrorName(res, &s) == abi.success) if (s) |p| return std.mem.span(p); + return knownName(res) orelse "CUDA_ERROR_UNKNOWN_CODE"; + } + + /// A kernel's symbol name (cuFuncGetName), or "?" when the driver cannot say. + pub fn functionName(self: *const Driver, f: abi.Function) []const u8 { + const get = self.func_name orelse return "?"; + var s: ?[*:0]const u8 = null; + if (get(&s, f) != abi.success) return "?"; + return if (s) |p| std.mem.span(p) else "?"; } pub fn errorText(self: *const Driver, res: abi.Result) []const u8 { @@ -96,3 +115,43 @@ pub const Driver = struct { return std.mem.sliceTo(buf, 0); } }; + +/// CUresult names (cuda.h) for when cuGetErrorName cannot answer. +pub fn knownName(res: abi.Result) ?[]const u8 { + return switch (res) { + 1 => "CUDA_ERROR_INVALID_VALUE", + 2 => "CUDA_ERROR_OUT_OF_MEMORY", + 3 => "CUDA_ERROR_NOT_INITIALIZED", + 4 => "CUDA_ERROR_DEINITIALIZED", + 101 => "CUDA_ERROR_INVALID_DEVICE", + 200 => "CUDA_ERROR_INVALID_IMAGE", + 201 => "CUDA_ERROR_INVALID_CONTEXT", + 209 => "CUDA_ERROR_NO_BINARY_FOR_GPU", + 400 => "CUDA_ERROR_INVALID_HANDLE", + 500 => "CUDA_ERROR_NOT_FOUND", + 600 => "CUDA_ERROR_NOT_READY", + 700 => "CUDA_ERROR_ILLEGAL_ADDRESS", + 701 => "CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES", + 702 => "CUDA_ERROR_LAUNCH_TIMEOUT", + 709 => "CUDA_ERROR_CONTEXT_IS_DESTROYED", + 710 => "CUDA_ERROR_ASSERT", + 714 => "CUDA_ERROR_HARDWARE_STACK_ERROR", + 715 => "CUDA_ERROR_ILLEGAL_INSTRUCTION", + 716 => "CUDA_ERROR_MISALIGNED_ADDRESS", + 719 => "CUDA_ERROR_LAUNCH_FAILED", + 999 => "CUDA_ERROR_UNKNOWN", + else => null, + }; +} + +/// CUDA_ERROR_DEINITIALIZED: the process is exiting (a stop signal, then exit) while a thread still calls the driver. +pub fn shutdownNote(res: abi.Result) []const u8 { + return if (res == 4) " (the CUDA driver is shutting down: the process is exiting)" else ""; +} + +test "CUresult names without the driver: DEINITIALIZED is 4, unknown codes stay unknown" { + try std.testing.expectEqualStrings("CUDA_ERROR_DEINITIALIZED", knownName(4).?); + try std.testing.expectEqualStrings("CUDA_ERROR_ILLEGAL_ADDRESS", knownName(700).?); + try std.testing.expect(knownName(12345) == null); + try std.testing.expect(shutdownNote(4).len > 0 and shutdownNote(1).len == 0); +} diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig new file mode 100644 index 0000000000..1916b22b57 --- /dev/null +++ b/zig/src/cuda/graph_cache.zig @@ -0,0 +1,448 @@ +//! Keyed CUDA graphs for any family: one captured executable per key, launched again when the key comes back. + +const std = @import("std"); +const graph = @import("graph.zig"); +const Driver = @import("driver.zig").Driver; +const Stream = @import("stream.zig").Stream; +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel + +/// One step's launches on `stream`; it must not synchronize, allocate device memory or read the device on the host. +pub const Body = struct { + ctx: *anyopaque, + run: *const fn (ctx: *anyopaque, stream: Stream) anyerror!void, +}; + +/// Every rank's verdict AND-ed on a miss, never a hit; an off cache never votes, so ranks match `Settings.digest`. +pub const Agree = struct { + ctx: *anyopaque, + all: *const fn (ctx: *anyopaque, ok: bool) anyerror!bool, +}; + +/// Capture a body into an executable graph, launch it, free it: a fake engine runs the policy in host tests. +pub const Engine = struct { + ctx: *anyopaque, + capture: *const fn (ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque, + launch: *const fn (ctx: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void, + free: *const fn (ctx: *anyopaque, exec: *anyopaque) void, +}; + +pub const Settings = struct { + on: bool = true, // off: every step runs eagerly + max: u32 = 48, // graphs held at most + min_keep: u32 = 8, // under the memory floor the cache evicts down to this many, then stops capturing + hold: bool = false, // under the floor a miss runs eagerly with no eviction: held graphs cost next to no memory + + /// These settings under the operator's TF_GRAPH_CACHE value: `0` or `off` runs every step eagerly; else unchanged. + pub fn operator(s: Settings, value: ?[]const u8) Settings { + var out = s; + const v = value orelse return out; + if (std.mem.eql(u8, v, "0") or std.ascii.eqlIgnoreCase(v, "off")) out.on = false; + return out; + } + + /// What a cache runs with: the operator's value, min_keep under max (0 with hold), and every off cache one value. + pub fn effective(s: Settings, value: ?[]const u8) Settings { + var out = s.operator(value); + if (!out.on or out.max == 0) return .{ .on = false, .max = 0, .min_keep = 0, .hold = false }; + out.min_keep = if (out.hold) 0 else @min(out.min_keep, out.max); + return out; + } + + /// `effective` under this process's TF_GRAPH_CACHE. + pub fn fromEnv(s: Settings) Settings { + return s.effective(if (std.c.getenv("TF_GRAPH_CACHE")) |v| std.mem.span(v) else null); + } + + /// The effective settings as bytes for ranks to compare at start (a family's bootstrap agreement), never per step. + pub fn digest(s: Settings) [10]u8 { + const e = s.effective(null); + var out: [10]u8 = undefined; + out[0] = @intFromBool(e.on); + out[1] = @intFromBool(e.hold); + std.mem.writeInt(u32, out[2..6], e.max, .little); + std.mem.writeInt(u32, out[6..10], e.min_keep, .little); + return out; + } +}; + +pub const Outcome = enum { replayed, captured, eager }; + +pub fn Cache(comptime Key: type) type { + return struct { + const Self = @This(); + const Entry = struct { exec: *anyopaque, used: u64 }; + + gpa: std.mem.Allocator, + engine: Engine, + settings: Settings, + agree: ?Agree = null, + room: ?*const fn () bool = null, // false when this rank is below its memory floor; null: always room + entries: std.AutoArrayHashMapUnmanaged(Key, Entry) = .empty, + cap: u32, + tick: u64 = 0, + fingerprint: ?u64 = null, + failed: bool = false, // a capture failed on some rank: eager from here on + stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, + + /// The cache with `settings.fromEnv()`: every rank sharing an `agree` must hold an equal `Settings.digest`. + pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { + const s = settings.fromEnv(); + if (!quiet and settings.on and settings.max > 0 and !s.on) std.log.info("graph cache: off (TF_GRAPH_CACHE), every step eager", .{}); + return .{ .gpa = gpa, .engine = engine, .settings = s, .cap = s.max }; + } + + pub fn deinit(c: *Self) void { + c.dropAll(); + c.entries.deinit(c.gpa); + } + + pub fn count(c: *const Self) usize { + return c.entries.count(); + } + + pub fn has(c: *const Self, key: Key) bool { + return c.entries.contains(key); + } + + /// Every graph freed (a fingerprint change, a shutdown). + pub fn dropAll(c: *Self) void { + for (c.entries.values()) |e| c.engine.free(c.engine.ctx, e.exec); + c.stats.dropped += c.entries.count(); + c.entries.clearRetainingCapacity(); + } + + fn agreed(c: *Self, ok: bool) !bool { + const a = c.agree orelse return ok; + return a.all(a.ctx, ok); + } + + fn evictOldest(c: *Self) void { + var at: ?usize = null; + for (c.entries.values(), 0..) |e, i| { + if (at == null or e.used < c.entries.values()[at.?].used) at = i; + } + const i = at orelse return; + c.engine.free(c.engine.ctx, c.entries.values()[i].exec); + c.entries.swapRemoveAt(i); + c.stats.evicted += 1; + } + + fn eager(c: *Self, stream: Stream, body: Body) !Outcome { + try body.run(body.ctx, stream); + c.stats.eager += 1; + return .eager; + } + + /// One step: the key's graph, captured first on a miss; a new `fingerprint` (same step on all ranks) drops all. + pub fn run(c: *Self, key: Key, stream: Stream, fingerprint: u64, body: Body) !Outcome { + if (!c.settings.on or c.failed or c.settings.max == 0) return c.eager(stream, body); + if (c.fingerprint != null and c.fingerprint.? != fingerprint) c.dropAll(); + c.fingerprint = fingerprint; + c.tick += 1; + if (c.entries.getPtr(key)) |e| { + e.used = c.tick; + try c.engine.launch(c.engine.ctx, e.exec, stream); + c.stats.replayed += 1; + return .replayed; + } + // a miss: room on every rank (a host map that cannot grow is none), else shed a quarter and lower the cap + const reserved = if (c.entries.ensureUnusedCapacity(c.gpa, 1)) |_| true else |_| false; + const room_here = reserved and (if (c.room) |f| f() else true); + if (!try c.agreed(room_here)) { + if (c.settings.hold) { + if (!quiet and c.stats.floor_hits == 0) std.log.warn("graph cache: under the memory floor on {s} rank: this key eager, {d} graphs held (hold)", .{ if (room_here) "another" else "this", c.entries.count() }); + c.stats.floor_hits += 1; + return c.eager(stream, body); + } + if (c.entries.count() > c.settings.min_keep) { + const drop = @max(1, c.entries.count() / 4); + for (0..drop) |_| if (c.entries.count() > c.settings.min_keep) c.evictOldest(); + } + const was = c.cap; + c.cap = @intCast(@max(c.entries.count(), c.settings.min_keep)); // min_keep <= max (init) + c.stats.floor_hits += 1; + if (!quiet and c.cap != was) std.log.warn("graph cache: under the memory floor on {s} rank: cap {d} -> {d} ({d} held, {d} evicted, {d} eager so far)", .{ if (room_here) "another" else "this", was, c.cap, c.entries.count(), c.stats.evicted, c.stats.eager + 1 }); + return c.eager(stream, body); + } + if (c.cap < c.settings.max) { + if (!quiet) std.log.info("graph cache: room again on every rank: cap {d} -> {d}", .{ c.cap, c.settings.max }); + c.cap = c.settings.max; + } + while (c.entries.count() >= c.cap) c.evictOldest(); + const exec = c.engine.capture(c.engine.ctx, stream, body) catch |err| blk: { + // the only trace of why graphs went off (the step then runs eagerly, which may fail the same way) + if (!quiet) std.log.warn("graph capture failed: {s} (key {any}, {d} graphs held)", .{ @errorName(err), key, c.entries.count() }); + break :blk null; + }; + var owned = exec; // freed here until the map holds it, so a failed agreement cannot leak it + defer if (owned) |x| c.engine.free(c.engine.ctx, x); + if (!try c.agreed(exec != null)) { + c.failed = true; + c.dropAll(); + return c.eager(stream, body); + } + c.entries.putAssumeCapacityNoClobber(key, .{ .exec = exec.?, .used = c.tick }); + owned = null; + try c.engine.launch(c.engine.ctx, exec.?, stream); + c.stats.captured += 1; + return .captured; + } + }; +} + +/// The driver's engine: thread-local stream capture (other threads' CUDA calls stay out), instantiate, upload, launch. +pub const CudaEngine = struct { + d: *const Driver, + + pub fn engine(e: *CudaEngine) Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + + const Held = struct { exec: graph.Exec }; + + fn capture(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const e: *CudaEngine = @ptrCast(@alignCast(ctx)); + try graph.beginCapture(stream, .thread_local); + body.run(body.ctx, stream) catch |err| { + // end the capture so the stream is usable again; the partial graph is dropped + if (graph.endCapture(stream)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return err; + }; + var g = try graph.endCapture(stream); + defer g.deinit(); + const h = try std.heap.c_allocator.create(Held); + errdefer std.heap.c_allocator.destroy(h); + h.exec = try g.instantiate(); + errdefer h.exec.deinit(); + try h.exec.upload(stream); + _ = e; + return h; + } + + fn launch(_: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void { + const h: *Held = @ptrCast(@alignCast(exec)); + try h.exec.launchOn(stream); + } + + fn free(_: *anyopaque, exec: *anyopaque) void { + const h: *Held = @ptrCast(@alignCast(exec)); + h.exec.deinit(); + std.heap.c_allocator.destroy(h); + } +}; + +/// FNV-1a over the addresses a body captures (buffer bases, sizes): the fingerprint `run` compares, on one rank only. +pub fn fingerprintOf(words: []const u64) u64 { + var h: u64 = 0xcbf29ce484222325; + for (words) |w| { + var x = w; + for (0..8) |_| { + h = (h ^ (x & 0xff)) *% 0x100000001b3; + x >>= 8; + } + } + return h; +} + +// -- host tests: the policy with a fake engine -------------------------------------------------------------------- + +const Fake = struct { + captures: u32 = 0, + launches: u32 = 0, + live: u32 = 0, + fail_capture: bool = false, + slots: [64]u8 = undefined, + + fn engine(f: *Fake) Engine { + return .{ .ctx = f, .capture = cap, .launch = lau, .free = fre }; + } + fn cap(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail_capture) return error.CaptureFailed; + _ = stream; + _ = body; // a capture records the body; nothing runs + f.captures += 1; + f.live += 1; + return &f.slots[f.captures % f.slots.len]; + } + fn lau(ctx: *anyopaque, _: *anyopaque, _: Stream) anyerror!void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.launches += 1; + } + fn fre(ctx: *anyopaque, _: *anyopaque) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.live -= 1; + } +}; + +const Ran = struct { + n: u32 = 0, + fn body(r: *Ran) Body { + return .{ .ctx = r, .run = go }; + } + fn go(ctx: *anyopaque, _: Stream) anyerror!void { + const r: *Ran = @ptrCast(@alignCast(ctx)); + r.n += 1; + } +}; + +const no_stream: Stream = .{ .d = undefined, .handle = null }; + +test "graph cache: capture on first use, replay after, LRU at the cap" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 2 }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.replayed, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.captured, try c.run(2, no_stream, 7, r.body())); + _ = try c.run(1, no_stream, 7, r.body()); // 1 is now the most recent + try std.testing.expectEqual(Outcome.captured, try c.run(3, no_stream, 7, r.body())); // evicts 2 + try std.testing.expect(c.has(1) and c.has(3) and !c.has(2)); + try std.testing.expectEqual(@as(u32, 0), r.n); // the body never ran eagerly + try std.testing.expectEqual(@as(u32, 2), f.live); + try std.testing.expectEqual(@as(u32, 5), f.launches); +} + +test "graph cache: a new fingerprint drops every graph" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + defer c.deinit(); + _ = try c.run(1, no_stream, 7, r.body()); + _ = try c.run(2, no_stream, 7, r.body()); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 8, r.body())); + try std.testing.expectEqual(@as(usize, 1), c.count()); + try std.testing.expectEqual(@as(u32, 1), f.live); +} + +const Peer = struct { + theirs: []const bool, // the other rank's verdicts, in call order + at: usize = 0, + fn agree(p: *Peer) Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) anyerror!bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const t = p.theirs[p.at]; + p.at += 1; + return ok and t; + } +}; + +test "graph cache: a capture failing on another rank turns graphs off here too" { + var f: Fake = .{}; + var r: Ran = .{}; + // per miss: room, then the capture's verdict; the peer's capture fails on the second key + var peer: Peer = .{ .theirs = &.{ true, true, true, false } }; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + c.agree = peer.agree(); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(2, no_stream, 7, r.body())); + try std.testing.expect(c.failed); + try std.testing.expectEqual(@as(u32, 0), f.live); // this rank's good graphs are freed as well + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 2), r.n); +} + +var room_flag = true; +fn roomFn() bool { + return room_flag; +} + +test "graph cache: below the memory floor the least recent quarter goes, the cap follows, the step runs eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 6), c.count()); + try std.testing.expectEqual(@as(u32, 6), c.cap); + try std.testing.expect(!c.has(0) and !c.has(1) and c.has(7)); + room_flag = true; + // room again: the cap is back to max, the miss captures without evicting + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 7), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + try std.testing.expect(c.has(7) and c.has(2)); + try std.testing.expectEqual(@as(u64, 1), c.stats.floor_hits); +} + +test "graph cache: a dip under the floor does not leave later keys churning once memory is back" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 16, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..12) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; // a long prompt's workspace held + _ = try c.run(200, no_stream, 7, r.body()); + room_flag = true; // freed: the decode after it meets 12 new keys (a new context bucket), then replays them + const before = c.stats.evicted; + for (0..12) |k| _ = try c.run(@intCast(300 + k), no_stream, 7, r.body()); + for (0..12) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(300 + k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 16), c.cap); + // only what the cap of 16 forces: 9 left after the dip + 12 new = 21 -> 5 evictions, none during the replays + try std.testing.expectEqual(before + 5, c.stats.evicted); +} + +test "graph cache: hold - below the floor a miss runs eagerly, nothing is evicted, the held keys replay" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2, .hold = true }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(101, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 8), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + for (0..8) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u64, 0), c.stats.evicted); + try std.testing.expectEqual(@as(u64, 2), c.stats.floor_hits); + room_flag = true; + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 9), c.count()); +} + +test "graph cache: off runs every step eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .on = false }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 0), f.captures); +} + +test "graph cache: TF_GRAPH_CACHE=0 or off turns a family's cache off without a rebuild; anything else leaves it" { + const on: Settings = .{}; + for ([_][]const u8{ "0", "off", "OFF" }) |v| try std.testing.expect(!on.operator(v).on); + for ([_]?[]const u8{ null, "1", "on", "" }) |v| try std.testing.expect(on.operator(v).on); + try std.testing.expect(!(Settings{ .on = false }).operator("1").on); // the operator only turns a cache off + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), on.operator("off")); + defer c.deinit(); + for (0..3) |_| try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 0), f.captures); +} + +test "graph cache: fingerprint" { + try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); +} + +test { + _ = @import("graph_cache_regression_test.zig"); + _ = @import("graph_cache_ranks_test.zig"); +} diff --git a/zig/src/cuda/graph_cache_ranks_test.zig b/zig/src/cuda/graph_cache_ranks_test.zig new file mode 100644 index 0000000000..18b8a87e9d --- /dev/null +++ b/zig/src/cuda/graph_cache_ranks_test.zig @@ -0,0 +1,136 @@ +//! Two in-process ranks with fake engines: equal effective settings vote on misses only; a mismatch is a stall. + +const std = @import("std"); +const cache = @import("graph_cache.zig"); +const Stream = @import("stream.zig").Stream; +const no_stream: Stream = .{ .d = undefined, .handle = null }; +const equal = std.testing.expectEqual; +const Settings = cache.Settings; + +const Engine = struct { + captures: usize = 0, + live: usize = 0, + handles: [8]u8 = undefined, + + fn engine(e: *Engine) cache.Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + fn capture(p: *anyopaque, _: Stream, _: cache.Body) !*anyopaque { + const e: *Engine = @ptrCast(@alignCast(p)); + e.captures += 1; + e.live += 1; + return &e.handles[e.captures % e.handles.len]; + } + fn launch(_: *anyopaque, _: *anyopaque, _: Stream) !void {} + fn free(p: *anyopaque, _: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.live -= 1; + } + fn body(e: *Engine) cache.Body { + return .{ .ctx = e, .run = run }; + } + fn run(_: *anyopaque, _: Stream) !void {} +}; + +/// The two ranks' votes in call order; a rank waits at most `wait_ns` for the other's vote of the same round. +const Board = struct { + votes: [2][16]bool = undefined, + cast: [2]std.atomic.Value(u32) = .{ .init(0), .init(0) }, + wait_ns: u64, + + fn vote(b: *Board, rank: u1, ok: bool) !bool { + const r = b.cast[rank].load(.monotonic); + if (r == b.votes[rank].len) return error.TooManyVotes; + b.votes[rank][r] = ok; + b.cast[rank].store(r + 1, .release); + const other = 1 - rank; + const deadline = now() + b.wait_ns; + while (b.cast[other].load(.acquire) <= r) { + if (now() > deadline) return error.PeerNeverVoted; + std.Thread.yield() catch {}; + } + return ok and b.votes[other][r]; + } + + fn now() u64 { + return @intCast(std.Io.Clock.awake.now(std.testing.io).toNanoseconds()); + } +}; + +const Rank = struct { + board: *Board, + id: u1, + engine: Engine = .{}, + outcomes: [8]cache.Outcome = undefined, + err: ?anyerror = null, + + fn all(ctx: *anyopaque, ok: bool) anyerror!bool { + const r: *Rank = @ptrCast(@alignCast(ctx)); + return r.board.vote(r.id, ok); + } + + /// One rank's steps over `keys` with `settings` exactly as its process resolved them. + fn steps(r: *Rank, settings: Settings, ks: []const u32) void { + var c = cache.Cache(u32).init(std.testing.allocator, r.engine.engine(), settings); + defer c.deinit(); + c.agree = .{ .ctx = r, .all = all }; + for (ks, 0..) |k, i| r.outcomes[i] = c.run(k, no_stream, 7, r.engine.body()) catch |e| { + r.err = e; + return; + }; + } +}; + +const keys = [_]u32{ 1, 2, 1, 2, 1 }; + +test "graph cache ranks: equal enabled settings capture together and vote on each miss, never on a hit" { + var board: Board = .{ .wait_ns = 5 * std.time.ns_per_s }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + const s = (Settings{ .max = 4 }).effective(null); + const t = try std.Thread.spawn(.{}, Rank.steps, .{ &ranks[1], s, &keys }); + ranks[0].steps(s, &keys); + t.join(); + const want = [_]cache.Outcome{ .captured, .captured, .replayed, .replayed, .replayed }; + for (&ranks) |*r| { + try equal(@as(?anyerror, null), r.err); + try std.testing.expectEqualSlices(cache.Outcome, &want, r.outcomes[0..keys.len]); + try equal(@as(u32, 4), board.cast[r.id].load(.acquire)); // two misses: room and capture each, hits none + try equal(@as(usize, 0), r.engine.live); + } +} + +test "graph cache ranks: TF_GRAPH_CACHE 0, off and OFF are one effective setting, and an off cache never votes" { + const d: Settings = .{}; + const off = d.effective("0").digest(); + try std.testing.expectEqualSlices(u8, &off, &d.effective("off").digest()); + try std.testing.expectEqualSlices(u8, &off, &d.effective("OFF").digest()); + try std.testing.expectEqualSlices(u8, &off, &(Settings{ .max = 0, .hold = true }).digest()); + try std.testing.expect(!std.mem.eql(u8, &off, &d.effective(null).digest())); + const held: Settings = .{ .hold = true, .min_keep = 2 }; + try std.testing.expectEqualSlices(u8, &held.digest(), &(Settings{ .hold = true, .min_keep = 30 }).digest()); + try std.testing.expectEqualSlices(u8, &d.effective("1").digest(), &d.effective(null).digest()); + var board: Board = .{ .wait_ns = 0 }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + ranks[0].steps(d.effective("0"), &keys); + ranks[1].steps(d.effective("OFF"), &keys); + for (&ranks) |*r| { + try equal(@as(?anyerror, null), r.err); + for (r.outcomes[0..keys.len]) |o| try equal(cache.Outcome.eager, o); + try equal(@as(u32, 0), board.cast[r.id].load(.acquire)); + try equal(@as(usize, 0), r.engine.captures); + } +} + +test "graph cache ranks: one rank switched off leaves the other waiting at its first miss; digests show it first" { + const d: Settings = .{}; + try std.testing.expect(!std.mem.eql(u8, &d.effective(null).digest(), &d.effective("off").digest())); + var board: Board = .{ .wait_ns = 50 * std.time.ns_per_ms }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + ranks[1].steps(d.effective("off"), &keys); + try equal(@as(?anyerror, null), ranks[1].err); + try equal(@as(u32, 0), board.cast[1].load(.acquire)); // the off rank ran every step eagerly without a vote + ranks[0].steps(d.effective(null), &keys); + try equal(@as(?anyerror, error.PeerNeverVoted), ranks[0].err); // in production this wait has no deadline + try equal(@as(u32, 1), board.cast[0].load(.acquire)); + try equal(@as(usize, 0), ranks[0].engine.live); +} diff --git a/zig/src/cuda/graph_cache_regression_test.zig b/zig/src/cuda/graph_cache_regression_test.zig new file mode 100644 index 0000000000..7f2da1063a --- /dev/null +++ b/zig/src/cuda/graph_cache_regression_test.zig @@ -0,0 +1,181 @@ +//! Failure ownership, rank verdict counts and bounded-cap regressions for the graph cache. + +const std = @import("std"); +const cache = @import("graph_cache.zig"); +const Stream = @import("stream.zig").Stream; +const no_stream: Stream = .{ .d = undefined, .handle = null }; +const expect = std.testing.expect; +const equal = std.testing.expectEqual; + +const Engine = struct { + captures: usize = 0, + live: usize = 0, + launches: usize = 0, + bodies: usize = 0, + handles: [32]u8 = undefined, + + fn engine(e: *Engine) cache.Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + fn capture(p: *anyopaque, _: Stream, _: cache.Body) !*anyopaque { + const e: *Engine = @ptrCast(@alignCast(p)); + e.captures += 1; + e.live += 1; + return &e.handles[e.captures % e.handles.len]; + } + fn launch(p: *anyopaque, _: *anyopaque, _: Stream) !void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.launches += 1; + } + fn free(p: *anyopaque, _: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(p)); + std.debug.assert(e.live > 0); + e.live -= 1; + } + fn body(e: *Engine) cache.Body { + return .{ .ctx = e, .run = run }; + } + fn run(p: *anyopaque, _: Stream) !void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.bodies += 1; + } +}; + +const Peer = struct { + verdicts: []const bool, + at: usize = 0, + throw_at: ?usize = null, + fn agree(p: *Peer) cache.Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) !bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const at = p.at; + p.at += 1; + if (p.throw_at == at) return error.CollectiveFailed; + if (at >= p.verdicts.len) return error.UnexpectedCollective; + return ok and p.verdicts[at]; + } +}; + +test "graph cache regression: capture agreement throw frees the not-yet-inserted executable" { + var e: Engine = .{}; + // per miss: room, then the capture's verdict + var p: Peer = .{ .verdicts = &.{ true, true }, .throw_at = 1 }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try std.testing.expectError(error.CollectiveFailed, c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), e.captures); + try equal(@as(usize, 0), e.live); + try equal(@as(usize, 0), c.count()); + try equal(@as(usize, 2), p.at); +} + +test "graph cache regression: local map OOM declines before any capture and can recover" { + var e: Engine = .{}; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var c = cache.Cache(u32).init(failing.allocator(), e.engine(), .{}); + defer c.deinit(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try expect(failing.has_induced_failure); + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 0), e.live); + try expect(!c.failed); + failing.fail_index = std.math.maxInt(usize); + try equal(cache.Outcome.captured, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), e.live); +} + +test "graph cache regression: peer map OOM makes a locally prepared rank eager before capture" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{false} }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), p.at); + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 1), e.bodies); +} + +test "graph cache regression: local OOM still contributes the room verdict" { + var e: Engine = .{}; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var p: Peer = .{ .verdicts = &.{true} }; + var c = cache.Cache(u32).init(failing.allocator(), e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), p.at); + try equal(@as(usize, 0), e.captures); +} + +test "graph cache regression: hits make no collective; a fingerprint change shared by every rank recaptures" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{ true, true, true, true, true, true } }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + _ = try c.run(1, no_stream, 7, e.body()); + _ = try c.run(2, no_stream, 7, e.body()); + try equal(@as(usize, 4), p.at); + for (0..4) |_| try equal(cache.Outcome.replayed, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 4), p.at); + // a new layout generation on every rank: both graphs go, the next miss agrees twice as any other + try equal(cache.Outcome.captured, try c.run(1, no_stream, 8, e.body())); + try equal(@as(usize, 6), p.at); + try equal(@as(u64, 2), c.stats.dropped); + try equal(@as(usize, 1), c.count()); + try equal(@as(usize, 1), e.live); +} + +var room = true; +fn available() bool { + return room; +} + +test "graph cache regression: zero maximum remains eager across pressure recovery, with no collective" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{} }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 0 }); + defer c.deinit(); + c.room = available; + c.agree = p.agree(); + for ([_]bool{ false, true, false, true }) |ok| { + room = ok; + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 0), c.cap); + try equal(@as(usize, 0), c.count()); + } + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 0), p.at); +} + +test "graph cache regression: default floor cannot enlarge a two-entry maximum" { + var e: Engine = .{}; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 2 }); + defer c.deinit(); + c.room = available; + try equal(@as(u32, 2), c.settings.min_keep); + room = false; + _ = try c.run(1, no_stream, 7, e.body()); + room = true; + for (0..8) |key| { + _ = try c.run(@intCast(key), no_stream, 7, e.body()); + try expect(c.count() <= 2 and c.cap <= 2 and e.live <= 2); + } +} + +test "graph cache regression: zero keep floor recovers after reducing the cap to zero" { + var e: Engine = .{}; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 2, .min_keep = 0 }); + defer c.deinit(); + c.room = available; + room = false; + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 0), c.cap); + room = true; + try equal(cache.Outcome.captured, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 2), c.cap); +} diff --git a/zig/src/cuda/launch.zig b/zig/src/cuda/launch.zig index 56bba5272a..47cbf720b6 100644 --- a/zig/src/cuda/launch.zig +++ b/zig/src/cuda/launch.zig @@ -4,7 +4,8 @@ const std = @import("std"); const abi = @import("abi.zig"); const Function = @import("module.zig").Function; const Stream = @import("stream.zig").Stream; -const Error = @import("driver.zig").Error; +const driver = @import("driver.zig"); +const Error = driver.Error; pub const Dim3 = abi.Dim3; @@ -88,7 +89,7 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void const g = cfg.grid; const b = cfg.block; if (cfg.plain()) { - return d.check(d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); + return launched(f, cfg, d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); } var attrs: [4]abi.LaunchAttribute = undefined; const n = fillAttributes(cfg, &attrs); @@ -104,7 +105,37 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void .attrs = &attrs, .num_attrs = @intCast(n), }; - try d.check(d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); + try launched(f, cfg, d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); +} + +/// A launch's result; a failure names the kernel and its configuration with the code. +fn launched(f: Function, cfg: Config, res: abi.Result, what: []const u8) Error!void { + if (res == abi.success) return; + if (res == abi.error_not_ready) return error.NotReady; // as Driver.check: not a failure to log + const d = f.d; + var buf: [160]u8 = undefined; + std.log.err("{s}: {s} ({d}) {s}; {s}{s}", .{ what, d.errorName(res), res, d.errorText(res), describe(&buf, d.functionName(f.handle), cfg), driver.shutdownNote(res) }); + return driver.Driver.errorOf(res); +} + +/// "kernel , grid XxYxZ, block XxYxZ, shared N B" (+ cluster / PDL / cooperative when set). +pub fn describe(buf: []u8, name: []const u8, cfg: Config) []const u8 { + const g = cfg.grid; + const b = cfg.block; + var w = std.Io.Writer.fixed(buf); + w.print("kernel {s}, grid {d}x{d}x{d}, block {d}x{d}x{d}, shared {d} B", .{ name, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared }) catch return w.buffered(); + if (cfg.cluster) |c| w.print(", cluster {d}x{d}x{d}", .{ c.x, c.y, c.z }) catch return w.buffered(); + if (cfg.pdl) w.writeAll(", pdl") catch return w.buffered(); + if (cfg.cooperative) w.writeAll(", cooperative") catch return w.buffered(); + return w.buffered(); +} + +test "a failed launch's description: the kernel, its grid, block and shared bytes, bounded" { + var buf: [160]u8 = undefined; + const cfg: Config = .{ .grid = .{ .x = 70000, .y = 2, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 }, .shared = 49152, .pdl = true }; + try std.testing.expectEqualStrings("kernel _stream_pf, grid 70000x2x1, block 128x1x1, shared 49152 B, pdl", describe(&buf, "_stream_pf", cfg)); + var small: [24]u8 = undefined; + try std.testing.expect(describe(&small, "a_very_long_kernel_symbol_name", cfg).len <= small.len); } fn attribute(id: abi.LaunchAttributeId) abi.LaunchAttribute { diff --git a/zig/src/cuda/memory.zig b/zig/src/cuda/memory.zig index 28c5b6244f..0f8f407806 100644 --- a/zig/src/cuda/memory.zig +++ b/zig/src/cuda/memory.zig @@ -100,18 +100,39 @@ fn lockCounts() void { while (!counts_mutex.tryLock()) std.Thread.yield() catch {}; } +/// A sub-allocator under DeviceBuffer (arena.zig): `alloc` gives an address or null for cuMemAlloc, `free` owns it. +pub const Hook = struct { + ctx: *anyopaque, + alloc: *const fn (ctx: *anyopaque, len: usize) ?abi.DevicePtr, + free: *const fn (ctx: *anyopaque, p: abi.DevicePtr) bool, // true when the address was the hook's +}; + +var alloc_hook: ?Hook = null; + +/// Installs, or with null removes, the sub-allocator: once at boot, before the buffers it should hold. +pub fn setHook(h: ?Hook) void { + alloc_hook = h; +} + +pub fn currentHook() ?Hook { + return alloc_hook; +} + pub const DeviceBuffer = struct { d: *const Driver, ptr: abi.DevicePtr, len: usize, borrowed: bool = false, // a span of memory another value owns (a carveout's): free leaves it alone - /// `len` bytes, 256-byte aligned by the driver; zero bytes allocate nothing and hold address 0. + /// `len` bytes, 256-byte aligned (512 from an installed arena); zero bytes allocate nothing and hold address 0. pub fn alloc(d: *const Driver, len: usize) Error!DeviceBuffer { try reserveAllocation(len); errdefer finishAllocation(null, len); var p: abi.DevicePtr = 0; - if (len > 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + if (len > 0) { + if (alloc_hook) |h| p = h.alloc(h.ctx, len) orelse 0; + if (p == 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + } finishAllocation(&device_bytes, len); return .{ .d = d, .ptr = p, .len = len }; } @@ -129,7 +150,10 @@ pub const DeviceBuffer = struct { pub fn free(self: *DeviceBuffer) void { if (!self.borrowed) { - if (self.ptr != 0) _ = self.d.api.cuMemFree_v2(self.ptr); + if (self.ptr != 0) { + const mine = if (alloc_hook) |h| h.free(h.ctx, self.ptr) else false; + if (!mine) _ = self.d.api.cuMemFree_v2(self.ptr); + } freed(&device_bytes, self.len); } self.* = undefined; @@ -266,6 +290,28 @@ test "freeing a borrowed span leaves the driver and the counts alone" { try std.testing.expectEqual(before.device, usage(false).device); } +test "a buffer the hook owns is freed by the hook, never by the driver" { + const Owner = struct { + var frees: u32 = 0; + fn alloc(_: *anyopaque, _: usize) ?abi.DevicePtr { + return null; + } + fn free(_: *anyopaque, p: abi.DevicePtr) bool { + frees += 1; + return p == 0x2000; + } + }; + var ctx: u8 = 0; + setHook(.{ .ctx = &ctx, .alloc = Owner.alloc, .free = Owner.free }); + defer setHook(null); + held(&device_bytes, 4096); + var b: DeviceBuffer = .{ .d = undefined, .ptr = 0x2000, .len = 4096 }; + b.free(); + var c: DeviceBuffer = .{ .d = undefined, .ptr = 0x3000, .len = 64, .borrowed = true }; + c.free(); + try std.testing.expectEqual(@as(u32, 1), Owner.frees); +} + test "peak reset preserves an allocation between its read and store" { const before = usage(true); held(&device_bytes, 100); diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index 420b2d460a..66fc50c3f1 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -39,6 +39,8 @@ pub const Config = extern struct { pub const Error = error{ LibraryUnavailable, MissingSymbol, NcclFailed }; +pub const in_progress: Result = 7; // ncclInProgress: a non-blocking communicator's call is still running + const R = Result; const D = abi.DevicePtr; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index bd3ffa6845..0a52db2a17 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -54,6 +54,8 @@ pub const tp_group = @import("tp_group.zig"); pub const tp_watch = @import("tp_watch.zig"); pub const ck = @import("ck.zig"); pub const modelopt = @import("modelopt.zig"); +pub const graph_cache = @import("graph_cache.zig"); +pub const arena = @import("arena.zig"); pub const tree_attention = @import("tree_attention.zig"); pub const block_attention = @import("block_attention.zig"); pub const window_attention = @import("window_attention.zig"); @@ -100,6 +102,8 @@ test { _ = @import("modelopt.zig"); _ = @import("ck.zig"); _ = @import("tp_watch.zig"); + _ = graph_cache; + _ = arena; _ = @import("shared_test.zig"); _ = @import("tree_attention.zig"); _ = @import("block_attention.zig"); diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md new file mode 100644 index 0000000000..7ad00ef7d5 --- /dev/null +++ b/zig/src/fabric/tp/README.md @@ -0,0 +1,72 @@ +# Tensor parallelism across processes + +`zig/src/fabric/tp` runs tensor parallelism with one GPU per process, either on one host or on two hosts over RoCE. +`TF_COMM_BACKEND` picks the transport for small exchanges: `nccl` (the default), `mailbox` (two GPUs on one host) +or `roce` (two hosts). NCCL carries everything larger, so it must work on its own as well. + +## Building `tf-tp-test` + +The `mailbox` and `roce` backends need the mailbox kernel (`mailbox.cu`) built into the binary: + +```sh +zig build tf-tp-test -Doptimize=ReleaseSafe -Dtp-nvcc=/usr/local/cuda/bin/nvcc +# or embed a fatbin built elsewhere: +zig build tf-tp-test -Doptimize=ReleaseSafe -Dtp-fatbin=/abs/path/tp_mailbox.fatbin +``` + +A build without `-Dtp-nvcc` or `-Dtp-fatbin` still runs `nccl` and `roce-host`, which needs no GPU. +In that build the `mailbox` and `roce` suites fail at load with `NoKernel` (an engine falls back to NCCL instead). +The `INFO` line `tf-tp-test` prints at start shows the embedded kernel's size; `0 B` means none was built in. + +## Choosing the RDMA devices + +| Variable | Used by | Needed when | +| --- | --- | --- | +| `TF_TP_ROCE_HCA` | the RoCE wire | the host has more than one active RoCE port that is not cabled to the peer | +| `NCCL_IB_HCA` | NCCL | the same as above; NCCL otherwise picks its own devices | +| `NCCL_SOCKET_IFNAME` | NCCL bootstrap | the host has several IP interfaces and NCCL picks one the peer cannot reach | +| `TF_TP_MASTER` | the TCP bootstrap | always on two hosts: rank 0's address as the other rank reaches it | + +`TF_TP_ROCE_HCA` is `name[:gid],name[:gid]`: up to two devices, and optionally each one's RoCE v2 GID index. +Without it the rank picks the devices itself, at most `TF_TP_ROCE_HCAS` (default 2). +A host cabled to more than one neighbour should always set it explicitly. +The two hosts may list their devices in a different order and use different names. +Each device pairs with the peer device on the same IPv4 /24. + +`NCCL_IB_HCA` should name the same devices without the `:gid` suffix. +In NCCL, `:N` after a device name selects port `N`, not a GID index. + +`run_hosts.sh` sets `NCCL_IB_HCA` from `TF_TP_ROCE_HCA` when only that is set. +When neither is set, it first runs `tf-tp-test hca` on both ranks: the RoCE pick, made with the peer over the bootstrap. +The same rule picks the wire's devices, and each rank prints its own as `NCCL_IB_HCA=name,...`. +`hca` needs no GPU, and it exits 1 when no port shares a /24 with the peer, so no host falls back to its first ports. +The runner logs each rank's pick and stops with that reason when `hca` fails. +Both hosts must use the same mode: if one names its ports and the other names none, only the latter runs `hca`. +It finds no peer and stops after `TF_TP_PLAN_TIMEOUT` (default 300 s) with that reason. + +Two DGX Spark nodes, with the second ConnectX-7 port of each cabled to the other (both PCIe functions of that +port), and the first port cabled to a third node: + +```sh +# on both hosts; rank 0's link address is 192.0.2.1 +export TF_TP_ROCE_HCA=rocep1s0f1,roceP2p1s0f1 +export NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1 +export NCCL_SOCKET_IFNAME=enp1s0f1np1 +zig/src/fabric/tp/run_hosts.sh 0 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank0 # host A +zig/src/fabric/tp/run_hosts.sh 1 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank1 # host B +``` + +`ibv_devices` lists the RDMA devices; `ibv_devinfo -d NAME` shows a port's state and link layer. +`/sys/class/infiniband/NAME/ports/1/gids/N` and `gid_attrs/ndevs/N` show each GID and its netdev. +`ulimit -l` must allow the RoCE region to be pinned (`unlimited` is simplest). + +## RoCE limits and the failed-transport check + +No RDMA write is larger than 1 MiB, so on `TF_COMM_BACKEND=roce` a `TF_TP_MAILBOX_MAX_KB` above 1024 fails at boot. +`TF_TP_ROCE_CHUNKED_KB` (default 0, off) sends variable-length gathers past the slot over RoCE in slot-sized ops. +It needs a kernel built from this tree's `mailbox.cu`. +At most 2 MiB is posted and not yet completed at any time. +A rank whose transport fails moves its queue pairs to the error state and keeps its region until the process exits. + +`run_transport.sh MASTER_ADDR PEER_SSH OUTDIR` checks that failure path on two hosts with `tf-tp-test failfast V +transport` and stops at the first kernel log line or HCA error counter that grows; run it with no arguments for help. diff --git a/zig/src/fabric/tp/bootstrap.zig b/zig/src/fabric/tp/bootstrap.zig new file mode 100644 index 0000000000..e1a6911efb --- /dev/null +++ b/zig/src/fabric/tp/bootstrap.zig @@ -0,0 +1,311 @@ +//! Rank discovery over TCP (a star on rank 0): small host blobs all-gathered, settings agreed, sockets kept for fate. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; + +pub const max_ranks = 64; +/// Largest blob a rank contributes to one `allGather`. +pub const max_blob = 64 * 1024; + +const magic: u32 = 0x31425054; // "TPB1" +const version: u32 = 2; + +pub const Error = error{ Protocol, Mismatch, TooLarge } || sock.Error; + +/// `role`: 0 the control connection (bootstrap, then the fate channel), 1 the plan link. +const Hello = extern struct { magic: u32, version: u32, world: u32, rank: u32, role: u32 }; + +pub const Bootstrap = struct { + rank: u32, + world: u32, + /// Rank 0: the socket of rank i at index i (index 0 unused); other ranks: rank 0's socket at index 0. + fds: [max_ranks]sock.Fd = @splat(-1), + /// The plan link's sockets, laid out as `fds`. + plan_fds: [max_ranks]sock.Fd = @splat(-1), + /// Each call's limit for the peers' bytes (the connect timeout). + timeout_ns: u64, + + fn deadline(b: *const Bootstrap) u64 { + return sock.nowNs() + b.timeout_ns; + } + + /// Rank 0 listens on `link_port` and accepts every other rank; the others connect to `master` (connect timeout). + pub fn open(cfg: Config) Error!Bootstrap { + if (cfg.world > max_ranks or cfg.rank >= cfg.world) return error.Protocol; + var b: Bootstrap = .{ .rank = cfg.rank, .world = cfg.world, .timeout_ns = cfg.connect_timeout_ns }; + errdefer b.close(); + if (cfg.world == 1) return b; + if (cfg.rank == 0) { + const l = try sock.listen(cfg.link_port); + defer sock.close(l); + try b.acceptAll(l); + } else { + const until = b.deadline(); + for ([_]u32{ 0, 1 }) |role| { + const fd = try sock.connect(cfg.master, cfg.link_port, until); + if (role == 0) b.fds[0] = fd else b.plan_fds[0] = fd; + const h: Hello = .{ .magic = magic, .version = version, .world = cfg.world, .rank = cfg.rank, .role = role }; + try sock.sendAll(fd, std.mem.asBytes(&h)); + } + } + return b; + } + + /// Rank 0 over an already listening socket (tests: an ephemeral port). + pub fn acceptOn(l: sock.Fd, world: u32, timeout_ns: u64) Error!Bootstrap { + var b: Bootstrap = .{ .rank = 0, .world = world, .timeout_ns = timeout_ns }; + errdefer b.close(); + try b.acceptAll(l); + return b; + } + + fn acceptAll(b: *Bootstrap, l: sock.Fd) Error!void { + // every other rank's control and plan connections, in any order + var joined: u32 = 0; + const until = b.deadline(); + while (joined < 2 * (b.world - 1)) { + const fd = try sock.accept(l, until); + var h: Hello = undefined; + sock.recvExact(fd, std.mem.asBytes(&h), until) catch { + sock.close(fd); + continue; + }; + if (h.magic != magic or h.version != version) { + log.warn("tp bootstrap: a connection that is not a rank (magic {x}); ignored", .{h.magic}); + sock.close(fd); + continue; + } + const table = if (h.role == 0) &b.fds else &b.plan_fds; + if (h.world != b.world or h.rank == 0 or h.rank >= b.world or h.role > 1 or table[h.rank] >= 0) { + log.warn("tp bootstrap: rank {d} of {d} joined a world of {d}", .{ h.rank, h.world, b.world }); + sock.close(fd); + return error.Mismatch; + } + table[h.rank] = fd; + joined += 1; + } + } + + /// Takes the control sockets (the fate channel keeps them); `close` then leaves them open. + pub fn release(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.fds; + b.fds = @splat(-1); + return out; + } + + /// Takes the plan link's sockets. + pub fn releasePlan(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.plan_fds; + b.plan_fds = @splat(-1); + return out; + } + + pub fn close(b: *Bootstrap) void { + for ([_]*[max_ranks]sock.Fd{ &b.fds, &b.plan_fds }) |table| for (table) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } + + /// Every rank's `mine` in rank order into `out`, returned as slices of `out`; every rank calls it equally often. + pub fn allGather(b: *Bootstrap, mine: []const u8, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + if (mine.len > max_blob) return error.TooLarge; + std.debug.assert(blobs.len >= b.world); + if (b.world == 1) { + if (out.len < mine.len) return error.TooLarge; + @memcpy(out[0..mine.len], mine); + blobs[0] = out[0..mine.len]; + return blobs[0..1]; + } + if (b.rank != 0) { + try sendFrame(b.fds[0], mine); + return b.readAll(b.fds[0], out, blobs); + } + // rank 0: everyone's blob, then the whole table to everyone + var used: usize = 0; + for (0..b.world) |i| { + const len: u32 = if (i == 0) @intCast(mine.len) else blk: { + var l: u32 = undefined; + try sock.recvExact(b.fds[i], std.mem.asBytes(&l), b.deadline()); + break :blk l; + }; + if (len > max_blob or used + len > out.len) return error.TooLarge; + if (i == 0) @memcpy(out[used..][0..len], mine) else try sock.recvExact(b.fds[i], out[used..][0..len], b.deadline()); + blobs[i] = out[used..][0..len]; + used += len; + } + for (1..b.world) |i| { + for (blobs[0..b.world]) |blob| try sendFrame(b.fds[i], blob); + } + return blobs[0..b.world]; + } + + fn readAll(b: *Bootstrap, fd: sock.Fd, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + var used: usize = 0; + for (0..b.world) |i| { + var l: u32 = undefined; + try sock.recvExact(fd, std.mem.asBytes(&l), b.deadline()); + if (l > max_blob or used + l > out.len) return error.TooLarge; + try sock.recvExact(fd, out[used..][0..l], b.deadline()); + blobs[i] = out[used..][0..l]; + used += l; + } + return blobs[0..b.world]; + } + + /// Rank 0's `bytes` on every rank (into `out` on the others; `bytes` is ignored there). + pub fn broadcast(b: *Bootstrap, bytes: []const u8, out: []u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [4096]u8 = undefined; + if (out.len * b.world > scratch.len) return error.TooLarge; + const all = try b.allGather(if (b.rank == 0) bytes else &.{}, &scratch, &blobs); + if (all[0].len != out.len) return error.Protocol; + @memcpy(out, all[0]); + } + + /// Every rank's `bytes` must be equal (settings that change bytes on the wire); `what` names them in the log. + pub fn agree(b: *Bootstrap, bytes: []const u8, what: []const u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [8192]u8 = undefined; + const all = try b.allGather(bytes, &scratch, &blobs); + for (all, 0..) |blob, i| if (!std.mem.eql(u8, blob, all[0])) { + log.warn("tp bootstrap: rank {d}'s {s} differs from rank 0's; both ranks refuse to start", .{ i, what }); + return error.Mismatch; + }; + } +}; + +fn sendFrame(fd: sock.Fd, bytes: []const u8) sock.Error!void { + const l: u32 = @intCast(bytes.len); + try sock.sendAll(fd, std.mem.asBytes(&l)); + try sock.sendAll(fd, bytes); +} + +/// An ephemeral listening port for tests. +pub fn testListener() !struct { fd: sock.Fd, port: u16 } { + const l = try sock.listen(0); + var addr: std.posix.sockaddr.in = undefined; + var len: std.posix.socklen_t = @sizeOf(std.posix.sockaddr.in); + _ = std.c.getsockname(l, @ptrCast(&addr), &len); + return .{ .fd = l, .port = std.mem.bigToNative(u16, addr.port) }; +} + +fn testRank(cfg: Config, out: *[3][32]u8, err: *?anyerror) void { + var b = Bootstrap.open(cfg) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var id: [32]u8 = undefined; + b.broadcast(&@as([32]u8, @splat(0xAB)), &id) catch |e| { + err.* = e; + return; + }; + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const mine: [3]u8 = @splat(@intCast(cfg.rank)); + const all = b.allGather(mine[0 .. cfg.rank + 1], &buf, &blobs) catch |e| { + err.* = e; + return; + }; + for (all, 0..) |blob, i| @memcpy(out[i][0..blob.len], blob); + b.agree("same", "settings") catch |e| { + err.* = e; + }; +} + +test "three ranks gather in rank order and agree" { + const l = try testListener(); + defer sock.close(l.fd); + var outs: [3][3][32]u8 = undefined; + var errs: [3]?anyerror = @splat(null); + var threads: [2]std.Thread = undefined; + for (0..2) |i| threads[i] = try std.Thread.spawn(.{}, testRank, .{ Config{ .world = 3, .rank = @intCast(i + 1), .link_port = l.port, .connect_timeout_ns = 5 * std.time.ns_per_s }, &outs[i + 1], &errs[i + 1] }); + var b = try Bootstrap.acceptOn(l.fd, 3, 5 * std.time.ns_per_s); + defer b.close(); + var id: [32]u8 = undefined; + try b.broadcast(&@as([32]u8, @splat(0xAB)), &id); + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const all = try b.allGather(&.{0}, &buf, &blobs); + try b.agree("same", "settings"); + for (threads) |t| t.join(); + for (errs) |e| try std.testing.expect(e == null); + try std.testing.expectEqual(@as(usize, 3), all.len); + try std.testing.expectEqualSlices(u8, &.{ 2, 2, 2 }, all[2]); + try std.testing.expectEqualSlices(u8, &.{ 1, 1 }, outs[2][1][0..2]); +} + +test "settings that differ stop every rank" { + const l = try testListener(); + defer sock.close(l.fd); + const T = struct { + fn go(port: u16, err: *?anyerror) void { + var b = Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = port, .connect_timeout_ns = 5 * std.time.ns_per_s }) catch |e| { + err.* = e; + return; + }; + defer b.close(); + b.agree("knobs-b", "settings") catch |e| { + err.* = e; + }; + } + }; + var err: ?anyerror = null; + const t = try std.Thread.spawn(.{}, T.go, .{ l.port, &err }); + var b = try Bootstrap.acceptOn(l.fd, 2, 5 * std.time.ns_per_s); + defer b.close(); + try std.testing.expectError(error.Mismatch, b.agree("knobs-a", "settings")); + t.join(); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), err); +} + +/// Rank 1 agrees on `theirs.digest()` against rank 0's `ours.digest()`; returns both ranks' results. +fn agreeDigests(ours: Config, theirs: Config) !struct { ?anyerror, ?anyerror } { + const l = try testListener(); + defer sock.close(l.fd); + const T = struct { + fn go(port: u16, cfg: Config, err: *?anyerror) void { + var c = cfg; + c.world = 2; + c.rank = 1; + c.link_port = port; + c.connect_timeout_ns = 5 * std.time.ns_per_s; + var b = Bootstrap.open(c) catch |e| { + err.* = e; + return; + }; + defer b.close(); + b.agree(&c.digest(), "settings") catch |e| { + err.* = e; + }; + } + }; + var err1: ?anyerror = null; + const t = try std.Thread.spawn(.{}, T.go, .{ l.port, theirs, &err1 }); + var b = try Bootstrap.acceptOn(l.fd, 2, 5 * std.time.ns_per_s); + defer b.close(); + var err0: ?anyerror = null; + b.agree(&ours.digest(), "settings") catch |e| { + err0 = e; + }; + t.join(); + return .{ err0, err1 }; +} + +test "a graph cache switched off on one rank stops both at start; 0 and OFF on the two ranks start" { + const on: Config = .{ .graph_cache = .{} }; + const off: Config = .{ .graph_cache = .{}, .graph_cache_env = "off" }; + const r = try agreeDigests(on, off); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), r[0]); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), r[1]); + const none = try agreeDigests(.{}, .{ .graph_cache_env = "off" }); // a family without the cache is not refused + try std.testing.expectEqual(@as(?anyerror, null), none[0]); + try std.testing.expectEqual(@as(?anyerror, null), none[1]); + const same = try agreeDigests(.{ .graph_cache = .{}, .graph_cache_env = "0" }, .{ .graph_cache = .{}, .graph_cache_env = "OFF" }); + try std.testing.expectEqual(@as(?anyerror, null), same[0]); + try std.testing.expectEqual(@as(?anyerror, null), same[1]); + const family = try agreeDigests(.{ .graph_cache = .{ .max = 8 } }, .{ .graph_cache = .{ .max = 16 } }); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), family[0]); +} diff --git a/zig/src/fabric/tp/collective.zig b/zig/src/fabric/tp/collective.zig new file mode 100644 index 0000000000..d2334dc74a --- /dev/null +++ b/zig/src/fabric/tp/collective.zig @@ -0,0 +1,136 @@ +//! The tensor-parallel communicator every backend implements: stream-ordered, graph-capturable, rank-order sums. +const std = @import("std"); +const cuda = @import("cuda"); + +pub const DevicePtr = cuda.abi.DevicePtr; +pub const Stream = cuda.abi.Stream; + +pub const Error = error{ Aborted, Unsupported, Invalid, PeerFailed, Timeout, BackendFailed }; + +/// Element types the collectives move (`reduce` takes the float ones and the 32/64-bit integers). +pub const DType = enum(u8) { + u8, + i32, + i64, + f16, + bf16, + f32, + f64, + + pub fn size(t: DType) usize { + return switch (t) { + .u8 => 1, + .f16, .bf16 => 2, + .i32, .f32 => 4, + .i64, .f64 => 8, + }; + } +}; + +pub const Op = enum(u8) { sum, max, min }; + +/// Which implementation is behind a `Collective` (logs, bench labels). +pub const Kind = enum { nccl, mailbox, hybrid, roce, host }; + +pub const Collective = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + kind: Kind, + rank: *const fn (ptr: *anyopaque) u32, + world: *const fn (ptr: *anyopaque) u32, + /// recv[count] <- the element-wise reduction of every rank's send[count], added in rank order. + all_reduce: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void, + /// recv[world * count] <- every rank's send[count], rank 0's first (NCCL's layout). + all_gather: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, stream: Stream) Error!void, + /// Two ranks trade: this rank's send[count] lands in `peer`'s recv[count] and the peer's in ours (one group). + exchange: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + send: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + recv: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + /// Host: returns once every rank has called it and this rank's earlier collectives have finished. + barrier: *const fn (ptr: *anyopaque) Error!void, + /// Host: a failure the transport recorded since the last check (a peer that timed out, an NCCL async error). + check: *const fn (ptr: *anyopaque) Error!void, + /// Any thread: end pending and later work at once; safe while another thread is inside a call. + abort: *const fn (ptr: *anyopaque) void, + /// recv[world][stride]: rank r's first lens[r] bytes (device int32 [world]) land at recv + r * stride. + all_gather_v: ?*const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void = null, + /// Whether `all_gather_v` moves only the lengths for shards of `bytes` here (rank-local). Null: never. + var_fits: ?*const fn (ptr: *anyopaque, bytes: usize) bool = null, + /// Host: whether every rank's transport takes device lengths, agreed once per communicator. Null: false. + var_agreed: ?*const fn (ptr: *anyopaque) Error!bool = null, + }; + + pub fn kind(c: Collective) Kind { + return c.vtable.kind; + } + + pub fn rank(c: Collective) u32 { + return c.vtable.rank(c.ptr); + } + + pub fn world(c: Collective) u32 { + return c.vtable.world(c.ptr); + } + + pub fn allReduce(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void { + return c.vtable.all_reduce(c.ptr, src, dst, count, t, op, stream); + } + + pub fn allGather(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: Stream) Error!void { + return c.vtable.all_gather(c.ptr, src, dst, count, t, stream); + } + + /// `allGather` of `stride` elements a rank where only rank r's first lens[r] bytes matter (`VTable.all_gather_v`). + pub fn allGatherV(c: Collective, src: DevicePtr, dst: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void { + if (c.vtable.all_gather_v) |f| return f(c.ptr, src, dst, stride, lens, t, stream); + return c.vtable.all_gather(c.ptr, src, dst, stride, t, stream); + } + + /// Whether `allGatherV` of `bytes`-byte strides moves only the lengths here (callers agree first: `varAgreed`). + pub fn varFits(c: Collective, bytes: usize) bool { + const f = c.vtable.var_fits orelse return false; + return f(c.ptr, bytes); + } + + /// Every rank's transport takes device lengths: one agreement a communicator, called alike outside any graph. + pub fn varAgreed(c: Collective) Error!bool { + const f = c.vtable.var_agreed orelse return false; + return f(c.ptr); + } + + pub fn exchange(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.exchange(c.ptr, src, dst, count, t, peer, stream); + } + + pub fn send(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.send(c.ptr, buf, count, t, peer, stream); + } + + pub fn recv(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.recv(c.ptr, buf, count, t, peer, stream); + } + + pub fn barrier(c: Collective) Error!void { + return c.vtable.barrier(c.ptr); + } + + pub fn check(c: Collective) Error!void { + return c.vtable.check(c.ptr); + } + + pub fn abort(c: Collective) void { + c.vtable.abort(c.ptr); + } +}; + +/// Refuses a peer that is this rank or out of range. +pub fn checkPeer(rank: u32, world: u32, peer: u32) Error!void { + if (peer == rank or peer >= world) return error.Invalid; +} + +test "dtype sizes" { + try std.testing.expectEqual(@as(usize, 2), DType.bf16.size()); + try std.testing.expectEqual(@as(usize, 8), DType.i64.size()); +} diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig new file mode 100644 index 0000000000..6f96091c1c --- /dev/null +++ b/zig/src/fabric/tp/config.zig @@ -0,0 +1,288 @@ +//! A rank's TP settings from the environment: rank, master, port, backend, fail-fast knobs. +const std = @import("std"); +const log = @import("log.zig"); +const graph_cache = @import("cuda").graph_cache; + +pub const Backend = enum { + /// NCCL for everything (the default). + nccl, + /// Exchanges of at most `mailbox_max_bytes` a rank through the one-shot mailbox kernel, NCCL above that. + mailbox, + /// The same split between two hosts: the small exchanges over RoCE (doorbell + RDMA proxy), NCCL above. + roce, +}; + +/// What `roce` does when its setup or load-time probe fails on any rank (TF_TP_ROCE_FALLBACK). +pub const Fallback = enum { nccl, @"error" }; + +pub const Error = error{BadSetting}; + +/// RoCE's per-write cap (roce_guard.max_write): on the roce backend a slot past it fails boot, never clamped. +pub const max_rdma_write: usize = @import("roce_guard.zig").max_write; + +/// TF_TP_PLAN_PIN: the CPU the plan threads pin to at their first plan (auto: the highest cpu_capacity). +pub const PlanPin = union(enum) { auto, cpu: u32 }; + +pub const Config = struct { + world: u32 = 2, + rank: u32 = 0, + /// Rank 0's address as the other ranks reach it. + master: []const u8 = "localhost", + /// The rendezvous port (TF_TP_PORT); the link listens on `link_port`. + port: u16 = 29500, + /// Rank 0 listens here: TF_TP_PLAN_PORT, else `port` + 7. + link_port: u16 = 29507, + /// How long rank 0 waits for the others and they keep retrying the connect (TF_TP_PLAN_TIMEOUT). + connect_timeout_ns: u64 = 300 * std.time.ns_per_s, + /// CUDA device ordinal this rank uses. + device: u32 = 0, + backend: Backend = .nccl, + mailbox_max_bytes: usize = 256 * 1024, + /// TF_TP_ROCE_CHUNKED_KB: `allGatherV` strides up to this go RoCE in ops of at most a slot (0: off). + roce_chunked_bytes: usize = 0, + /// A mailbox wait for the peer gives up after this long and poisons the transport (fail-stop). + mailbox_timeout_ns: u64 = 120 * std.time.ns_per_s, + failfast: bool = true, + /// After a failure: time for the HTTP threads to answer before exit 70. + grace_ns: u64 = std.time.ns_per_s, + /// Rank 1's memory report period on the fate channel (also how fast it notices rank 0 hanging up). + report_ns: u64 = std.time.ns_per_s / 2, + nccl_lib: []const u8 = "libnccl.so.2", + /// `name[:gid],...` RDMA devices for `roce` (null: active RoCE v2 ports on an IPv4 /24 the peer has, roce_pick). + roce_hca: ?[]const u8 = null, + /// Stripe over at most this many HCAs (auto detection only). + roce_hcas: u8 = 2, + roce_tc: u8 = 0, + /// TF_TP_ROCE_FAST as the RoCE kernel's `opts` word: bit 0 on, bits 8-15 pollers, bits 16-31 stagger ns (local). + roce_fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest RoCE grid, a power of two 1..16; launch shape only, the bytes are the same. + roce_blocks: u32 = 16, + /// Pin the proxy thread (it busy-spins one core between exchanges). + roce_cpu: ?u32 = null, + /// `shm` runs the RoCE protocol on one host without RDMA (tests); `verbs` is the real wire. + roce_wire: enum { verbs, shm } = .verbs, + roce_fallback: Fallback = .nccl, + /// TF_TP_PLAN_LINK=rdma: the follower polls the plan socket this long before blocking (TF_TP_PLAN_SPIN_US). + plan_spin_ns: u64 = 0, + plan_pin: ?PlanPin = null, + /// The keyed graph cache settings a family builds with, set before `Session.open` (null: the family has none). + graph_cache: ?graph_cache.Settings = null, + /// TF_GRAPH_CACHE as read (`Cache.init` applies the same value); only its effect on `graph_cache` is compared. + graph_cache_env: ?[]const u8 = null, + + /// Reads every setting from `get` (the environment in `fromEnv`); unset names keep the defaults. + pub fn parse(get: *const fn ([]const u8) ?[]const u8) Error!Config { + var c: Config = .{}; + c.world = try int(u32, get, "TF_TP_WORLD", c.world, 1, 64); + c.rank = try int(u32, get, "TF_TP_RANK", c.rank, 0, c.world - 1); + if (get("TF_TP_MASTER")) |m| c.master = m; + c.port = try int(u16, get, "TF_TP_PORT", c.port, 1, 65535 - 8); + c.link_port = try int(u16, get, "TF_TP_PLAN_PORT", c.port + 7, 1, 65535); + c.connect_timeout_ns = try seconds(get, "TF_TP_PLAN_TIMEOUT", 300, 0.1, 86400); + c.device = try int(u32, get, "TF_TP_DEVICE", c.device, 0, 63); + const backend = "TF_COMM_BACKEND"; + if (get(backend)) |b| c.backend = std.meta.stringToEnum(Backend, std.mem.trim(u8, b, " ")) orelse { + log.warn("{s}={s}: expected nccl, mailbox or roce", .{ backend, b }); + return error.BadSetting; + }; + c.mailbox_max_bytes = 1024 * try int(usize, get, "TF_TP_MAILBOX_MAX_KB", 256, 4, 16 * 1024); + if (c.backend == .roce and c.mailbox_max_bytes > max_rdma_write) { + log.warn("TF_TP_MAILBOX_MAX_KB={d}: a RoCE write is capped at {d} KiB; refusing to start", .{ c.mailbox_max_bytes / 1024, max_rdma_write / 1024 }); + return error.BadSetting; + } + c.roce_chunked_bytes = 1024 * try int(usize, get, "TF_TP_ROCE_CHUNKED_KB", 0, 0, 64 * 1024); + c.mailbox_timeout_ns = try seconds(get, "TF_TP_MAILBOX_TIMEOUT_S", 120, 0.001, 86400); + if (get("TF_TP_FAILFAST")) |v| c.failfast = !(std.mem.eql(u8, v, "0") or std.mem.eql(u8, v, "off")); + c.grace_ns = try seconds(get, "TF_TP_FAILFAST_GRACE_S", 1.0, 0, 60); + c.report_ns = try seconds(get, "TF_TP_PEER_MEM_S", 0.5, 0.01, 60); + if (get("TF_NCCL_LIB")) |p| c.nccl_lib = p; + c.roce_hca = get("TF_TP_ROCE_HCA"); + c.roce_hcas = try int(u8, get, "TF_TP_ROCE_HCAS", 2, 1, 2); + c.roce_tc = try int(u8, get, "TF_TP_ROCE_TC", try int(u8, get, "NCCL_IB_TC", 0, 0, 255), 0, 255); + if (try int(u8, get, "TF_TP_ROCE_FAST", 0, 0, 1) == 1) { + const pollers = try int(u32, get, "TF_TP_ROCE_POLLERS", 1, 1, 8); + const stagger = try int(u32, get, "TF_TP_ROCE_STAGGER_NS", 150, 0, 65535); + c.roce_fast = 1 | (pollers << 8) | (stagger << 16); + } + c.roce_blocks = try int(u32, get, "TF_TP_ROCE_BLOCKS", 16, 1, 16); + if (c.roce_blocks & (c.roce_blocks - 1) != 0) { + log.warn("TF_TP_ROCE_BLOCKS={d}: expected a power of two", .{c.roce_blocks}); + return error.BadSetting; + } + if (get("TF_TP_ROCE_CPU")) |_| c.roce_cpu = try int(u32, get, "TF_TP_ROCE_CPU", 0, 0, 4095); + if (get("TF_TP_ROCE_WIRE")) |w| c.roce_wire = std.meta.stringToEnum(@TypeOf(c.roce_wire), w) orelse { + log.warn("TF_TP_ROCE_WIRE={s}: expected verbs or shm", .{w}); + return error.BadSetting; + }; + const fallback = "TF_TP_ROCE_FALLBACK"; + if (get(fallback)) |f| c.roce_fallback = std.meta.stringToEnum(Fallback, f) orelse { + log.warn("{s}={s}: expected nccl or error", .{ fallback, f }); + return error.BadSetting; + }; + if (get("TF_TP_PLAN_LINK")) |l| { + const mode = std.mem.trim(u8, l, " "); + if (std.mem.eql(u8, mode, "rdma")) { + c.plan_spin_ns = 1000 * try int(u64, get, "TF_TP_PLAN_SPIN_US", 5000, 0, 10_000_000); + } else if (!std.mem.eql(u8, mode, "tcp") and !std.mem.eql(u8, mode, "nccl")) { + log.warn("TF_TP_PLAN_LINK={s}: expected nccl, tcp or rdma", .{l}); + return error.BadSetting; + } + } + if (get("TF_TP_PLAN_PIN")) |raw| { + const v = std.mem.trim(u8, raw, " "); + if (std.mem.eql(u8, v, "auto")) { + c.plan_pin = .auto; + } else if (!std.mem.eql(u8, v, "-1") and !std.mem.eql(u8, v, "off")) { + c.plan_pin = .{ .cpu = try int(u32, get, "TF_TP_PLAN_PIN", 0, 0, 4095) }; + } + } + c.graph_cache_env = get("TF_GRAPH_CACHE"); + return c; + } + + pub fn fromEnv() Error!Config { + return parse(&envGet); + } + + /// The settings both ranks must share (everything but rank, device and timeouts), as bytes to compare. + pub fn digest(c: Config) [32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + h.update(std.mem.asBytes(&c.world)); + h.update(std.mem.asBytes(&c.backend)); + h.update(std.mem.asBytes(&c.mailbox_max_bytes)); + h.update(std.mem.asBytes(&c.roce_chunked_bytes)); + h.update(std.mem.asBytes(&c.failfast)); + h.update(std.mem.asBytes(&c.roce_wire)); + h.update(std.mem.asBytes(&c.roce_fallback)); + if (c.graph_cache) |g| h.update(&g.effective(c.graph_cache_env).digest()); + var out: [32]u8 = undefined; + h.final(&out); + return out; + } +}; + +fn envGet(name: []const u8) ?[]const u8 { + var buf: [64]u8 = undefined; + if (name.len >= buf.len) return null; + @memcpy(buf[0..name.len], name); + buf[name.len] = 0; + const v = std.c.getenv(buf[0..name.len :0]) orelse return null; + const s = std.mem.span(v); + return if (s.len == 0) null else s; +} + +fn int(comptime T: type, get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: T, lo: T, hi: T) Error!T { + const raw = get(name) orelse return default; + const v = std.fmt.parseInt(T, std.mem.trim(u8, raw, " "), 10) catch { + log.warn("{s}={s}: expected an integer", .{ name, raw }); + return error.BadSetting; + }; + if (v < lo or v > hi) { + log.warn("{s}={d}: expected {d}..{d}", .{ name, v, lo, hi }); + return error.BadSetting; + } + return v; +} + +fn seconds(get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: f64, lo: f64, hi: f64) Error!u64 { + const v = if (get(name)) |raw| std.fmt.parseFloat(f64, std.mem.trim(u8, raw, " ")) catch { + log.warn("{s}={s}: expected seconds", .{ name, raw }); + return error.BadSetting; + } else default; + if (!(v >= lo and v <= hi)) { + log.warn("{s}={d}: expected {d}..{d} s", .{ name, v, lo, hi }); + return error.BadSetting; + } + return @intFromFloat(v * std.time.ns_per_s); +} + +const Fake = struct { + var pairs: []const [2][]const u8 = &.{}; + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "defaults: the plan link 7 ports past the rendezvous, NCCL, fail-fast on" { + Fake.pairs = &.{ .{ "TF_TP_RANK", "1" }, .{ "TF_TP_PORT", "29600" } }; + const c = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u32, 1), c.rank); + try std.testing.expectEqual(@as(u16, 29607), c.link_port); + try std.testing.expectEqual(Backend.nccl, c.backend); + try std.testing.expectEqual(std.time.ns_per_s, c.grace_ns); + try std.testing.expectEqual(std.time.ns_per_s / 2, c.report_ns); +} + +test "bad settings are refused and the digest ignores the rank" { + Fake.pairs = &.{.{ "TF_TP_RANK", "2" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{.{ "TF_COMM_BACKEND", "rdma" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_ROCE_HCA", "rocep1s0f1,roceP2p1s0f1" }, .{ "NCCL_IB_TC", "106" } }; + const r = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, r.backend); + try std.testing.expectEqual(@as(u8, 106), r.roce_tc); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", r.roce_hca.?); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "64" }, .{ "TF_TP_FAILFAST_GRACE_S", "0.25" } }; + const a = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(usize, 64 * 1024), a.mailbox_max_bytes); + try std.testing.expectEqual(std.time.ns_per_s / 4, a.grace_ns); + var b = a; + b.rank = 1; + try std.testing.expectEqualSlices(u8, &a.digest(), &b.digest()); + b.mailbox_max_bytes = 1024; + try std.testing.expect(!std.mem.eql(u8, &a.digest(), &b.digest())); +} + +test "the plan link's spin and pin, the RoCE grid cap" { + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" }, .{ "TF_TP_ROCE_FALLBACK", "error" }, .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_PIN", "auto" } }; + const p = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, p.backend); + try std.testing.expectEqual(@as(usize, 1024 * 1024), p.mailbox_max_bytes); + try std.testing.expectEqual(Fallback.@"error", p.roce_fallback); + try std.testing.expectEqual(@as(u64, 5000 * 1000), p.plan_spin_ns); + try std.testing.expectEqual(PlanPin.auto, p.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "tcp" }, .{ "TF_TP_PLAN_PIN", "7" } }; // tcp blocks at once; an explicit CPU + const o = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 0), o.plan_spin_ns); + try std.testing.expectEqual(PlanPin{ .cpu = 7 }, o.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_SPIN_US", "200" }, .{ "TF_TP_PLAN_PIN", "-1" } }; + const s = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 200 * 1000), s.plan_spin_ns); + try std.testing.expect(s.plan_pin == null); + Fake.pairs = &.{.{ "TF_TP_PLAN_LINK", "udp" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; // the RoCE grid cap: 16 unset, powers of two only + try std.testing.expectEqual(@as(u32, 16), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "4" }}; + try std.testing.expectEqual(@as(u32, 4), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "12" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; +} + +test "the digest holds the graph cache's effective settings: 0, off and OFF agree, on and off differ" { + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "0" }}; + var zero = try Config.parse(&Fake.get); + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "OFF" }}; + var upper = try Config.parse(&Fake.get); + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "1" }}; + var one = try Config.parse(&Fake.get); + Fake.pairs = &.{}; + var unset = try Config.parse(&Fake.get); + try std.testing.expectEqualSlices(u8, &zero.digest(), &unset.digest()); // no cache installed: not compared + for ([_]*Config{ &zero, &upper, &one, &unset }) |c| c.graph_cache = .{}; + var max0 = unset; + max0.graph_cache = .{ .max = 0 }; + try std.testing.expectEqualSlices(u8, &zero.digest(), &upper.digest()); + try std.testing.expectEqualSlices(u8, &zero.digest(), &max0.digest()); + try std.testing.expectEqualSlices(u8, &one.digest(), &unset.digest()); + try std.testing.expect(!std.mem.eql(u8, &zero.digest(), &unset.digest())); + var family = unset; + family.graph_cache = .{ .max = 8, .hold = true }; + try std.testing.expect(!std.mem.eql(u8, &family.digest(), &unset.digest())); + var family_off = upper; + family_off.graph_cache = family.graph_cache; + try std.testing.expectEqualSlices(u8, &zero.digest(), &family_off.digest()); +} diff --git a/zig/src/fabric/tp/fate.zig b/zig/src/fabric/tp/fate.zig new file mode 100644 index 0000000000..5084c4648b --- /dev/null +++ b/zig/src/fabric/tp/fate.zig @@ -0,0 +1,261 @@ +//! Fail fast on every rank: a failure or a vanished rank ends all ranks with exit 70 about a second later. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; +const posix = std.posix; +/// Linux POLLRDHUP: the peer shut its side down (std has no name for it). +const pollrdhup: i16 = 0x2000; + +pub const exit_code: u8 = 70; +const magic: u32 = 0x46463431; +const kind_mem: u32 = 1; +const kind_fail: u32 = 2; +/// rank 0's clean stop: the hang-up that follows is no failure (sent before rank 0 closes its sockets) +const kind_stop: u32 = 3; +const Frame = extern struct { magic: u32, kind: u32, value: i64 }; + +pub const max_ranks = 64; +pub const max_hooks = 8; + +/// A callback with its context (an abort of a collective, rank 0's request failing); `reason` lives until exit. +pub const Hook = struct { ctx: *anyopaque, call: *const fn (ctx: *anyopaque, reason: []const u8) void }; + +pub const Fate = struct { + rank: u32, + world: u32, + fds: [max_ranks]sock.Fd, + grace_ns: u64, + report_ns: u64, + /// This rank's usable memory in bytes for the MEM report (default: MemAvailable). + usable: *const fn () u64 = memAvailable, + /// The process exit; tests replace it (it may return there). + exit_fn: *const fn (u8) void = exitNow, + hooks: [max_hooks]Hook = undefined, + n_hooks: std.atomic.Value(u32) = .init(0), + dead: std.atomic.Value(bool) = .init(false), + closing: std.atomic.Value(bool) = .init(false), + /// Set once the exit was asked for (tests read it; the process is gone in serving). + exited: std.atomic.Value(bool) = .init(false), + reason_buf: [512]u8 = undefined, + reason_len: usize = 0, + send_lock: std.atomic.Value(bool) = .init(false), + /// Rank 0: each rank's last report (bytes) and its monotonic time. + peer_mem: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + peer_mem_at: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + wake: [2]sock.Fd = .{ -1, -1 }, + thread: ?std.Thread = null, + + /// The channel over the bootstrap's sockets (`Bootstrap.release`); `start` runs its thread. + pub fn init(cfg: Config, fds: [max_ranks]sock.Fd) Fate { + return .{ .rank = cfg.rank, .world = cfg.world, .fds = fds, .grace_ns = cfg.grace_ns, .report_ns = cfg.report_ns }; + } + + pub fn start(f: *Fate) !void { + var p: [2]c_int = undefined; + if (std.c.pipe2(&p, .{ .CLOEXEC = true }) != 0) return error.SocketFailed; + f.wake = p; + f.thread = if (f.rank == 0) try std.Thread.spawn(.{}, watch0, .{f}) else try std.Thread.spawn(.{}, watchK, .{f}); + } + + /// Runs `hook` when the ranks go down; collectives register their abort here. + pub fn onFail(f: *Fate, hook: Hook) void { + const i = f.n_hooks.load(.acquire); + std.debug.assert(i < max_hooks); + f.hooks[i] = hook; + f.n_hooks.store(i + 1, .release); + } + + /// This rank failed (`where`: what it was doing): every rank goes down. Returns only in tests. + pub fn fatal(f: *Fate, where: []const u8, err: anyerror) void { + var buf: [256]u8 = undefined; + const msg = std.fmt.bufPrint(&buf, "rank {d} failed in {s}: {t}", .{ f.rank, where, err }) catch "rank failed"; + f.die(msg, true); + } + + /// This rank failed while booting: every rank goes down, and the thread is joined before the caller frees this. + pub fn failBoot(f: *Fate, where: []const u8, err: anyerror) void { + f.fatal(where, err); + f.close(); + } + + /// A normal stop is coming (rank 0 before it tells the others; the others when told): later hang-ups are fine. + pub fn expectStop(f: *Fate) void { + f.closing.store(true, .release); + } + + /// Rank 0's clean stop of every rank: a STOP frame to each ahead of its close, so followers end quietly. + pub fn announceStop(f: *Fate) void { + if (f.closing.swap(true, .acq_rel)) return; + if (f.rank != 0) return; + const h: Frame = .{ .magic = magic, .kind = kind_stop, .value = 0 }; + f.lock(); + defer f.unlock(); + for (f.fds[1..f.world]) |fd| if (fd >= 0) { + _ = sock.trySend(fd, std.mem.asBytes(&h)); + }; + } + + /// A normal stop: the thread ends and the sockets close. + pub fn close(f: *Fate) void { + f.closing.store(true, .release); + if (f.wake[1] >= 0) _ = std.c.write(f.wake[1], "x", 1); + if (f.thread) |t| t.join(); + f.thread = null; + for (&f.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + for (f.wake) |fd| if (fd >= 0) sock.close(fd); + f.wake = .{ -1, -1 }; + } + + /// Rank 0: rank `r`'s last reported usable bytes and the age of the report, or null before its first. + pub fn peerMem(f: *const Fate, r: u32) ?struct { bytes: u64, age_ns: u64 } { + const at = f.peer_mem_at[r].load(.acquire); + if (at == 0) return null; + return .{ .bytes = f.peer_mem[r].load(.acquire), .age_ns = sock.nowNs() -| at }; + } + + pub fn reason(f: *const Fate) []const u8 { + return f.reason_buf[0..f.reason_len]; + } + + fn lock(f: *Fate) void { + while (f.send_lock.swap(true, .acquire)) std.atomic.spinLoopHint(); + } + + fn unlock(f: *Fate) void { + f.send_lock.store(false, .release); + } + + fn die(f: *Fate, why: []const u8, own: bool) void { + if (f.dead.swap(true, .acq_rel)) { + // the first caller exits the process; a second (the fate thread or a round thread) waits for it + while (!f.exited.load(.acquire)) sock.sleepNs(10 * std.time.ns_per_ms); + return; + } + const n = @min(why.len, f.reason_buf.len); + @memcpy(f.reason_buf[0..n], why[0..n]); + f.reason_len = n; + const r = f.reason(); + log.warn("[tensorfold] tp fail-fast: rank {d}: {s}; exiting ({d}) in {d} ms so every rank restarts", .{ f.rank, r, exit_code, f.grace_ns / std.time.ns_per_ms }); + // aborts first: kernels waiting on the dead peer end now; a hook that blocks cannot hold the exit + if (std.Thread.spawn(.{}, runHooks, .{ f, r })) |t| t.detach() else |_| runHooks(f, r); + if (!f.closing.load(.acquire)) { + const head: Frame = .{ .magic = magic, .kind = kind_fail, .value = @intCast(r.len) }; + f.lock(); + for (f.fds[0..f.world]) |fd| if (fd >= 0) { + if (own or f.rank == 0) { + _ = sock.trySend(fd, std.mem.asBytes(&head)); + _ = sock.trySend(fd, r); + } + sock.hangup(fd); + }; + f.unlock(); + } + sock.sleepNs(f.grace_ns); + f.exited.store(true, .release); + f.exit_fn(exit_code); + } + + fn runHooks(f: *Fate, r: []const u8) void { + const n = f.n_hooks.load(.acquire); + for (f.hooks[0..n]) |h| h.call(h.ctx, r); + } + + /// Reads one frame from `fd`: a MEM report, a FAIL (the ranks go down), or the end (the peer is gone). + fn readFrame(f: *Fate, peer: u32, fd: sock.Fd) bool { + var h: Frame = undefined; + const deadline = sock.nowNs() + std.time.ns_per_s; + sock.recvExact(fd, std.mem.asBytes(&h), deadline) catch { + if (f.closing.load(.acquire)) return false; + var buf: [128]u8 = undefined; + f.die(std.fmt.bufPrint(&buf, "rank {d} closed the fate channel (it failed, exited or was killed)", .{peer}) catch "peer gone", false); + return false; + }; + if (h.magic != magic) { + f.die("a frame that is not the fate channel's", false); + return false; + } + if (h.kind == kind_stop) { + // rank 0 stops every rank: its hang-up is next, and no failure + f.closing.store(true, .release); + return false; + } + if (h.kind == kind_mem) { + f.peer_mem[peer].store(@intCast(@max(h.value, 0)), .release); + f.peer_mem_at[peer].store(sock.nowNs(), .release); + return true; + } + var text: [512]u8 = undefined; + const len: usize = @intCast(std.math.clamp(h.value, 0, text.len)); + sock.recvExact(fd, text[0..len], deadline) catch {}; + f.die(if (len > 0) text[0..len] else "a peer failed", false); + return false; + } + + fn watch0(f: *Fate) void { + var p: [max_ranks + 1]std.c.pollfd = undefined; + var who: [max_ranks + 1]u32 = undefined; + var n: usize = 0; + for (f.fds[1..f.world], 1..) |fd, r| if (fd >= 0) { + p[n] = .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }; + who[n] = @intCast(r); + n += 1; + }; + p[n] = .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }; + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + if (std.c.poll(&p, @intCast(n + 1), -1) <= 0) continue; + for (p[0..n], who[0..n]) |*q, r| if (q.revents != 0) { + if (!f.readFrame(r, q.fd)) return; + }; + } + } + + fn watchK(f: *Fate) void { + const fd = f.fds[0]; + var p = [_]std.c.pollfd{ + .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }, + .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }, + }; + const ms: c_int = @intCast(@max(f.report_ns / std.time.ns_per_ms, 1)); + f.report(fd); + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + const ready = std.c.poll(&p, 2, ms); + if (ready == 0) { + f.report(fd); + continue; + } + if (ready > 0 and p[0].revents != 0 and !f.readFrame(0, fd)) return; + } + } + + /// One MEM frame, never blocking (a full buffer drops it: the next comes one period later). + fn report(f: *Fate, fd: sock.Fd) void { + const h: Frame = .{ .magic = magic, .kind = kind_mem, .value = @intCast(f.usable()) }; + f.lock(); + defer f.unlock(); + _ = sock.trySend(fd, std.mem.asBytes(&h)); + } +}; + +fn exitNow(code: u8) void { + std.c._exit(code); +} + +/// MemAvailable from /proc/meminfo in bytes (0 when unreadable). +pub fn memAvailable() u64 { + const fd = std.c.open("/proc/meminfo", .{ .ACCMODE = .RDONLY }); + if (fd < 0) return 0; + defer _ = std.c.close(fd); + var buf: [4096]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return 0; + const text = buf[0..@intCast(n)]; + const key = "MemAvailable:"; + const at = std.mem.indexOf(u8, text, key) orelse return 0; + var it = std.mem.tokenizeAny(u8, text[at + key.len ..], " \n"); + const kb = std.fmt.parseInt(u64, it.next() orelse return 0, 10) catch return 0; + return kb * 1024; +} diff --git a/zig/src/fabric/tp/fate_test.zig b/zig/src/fabric/tp/fate_test.zig new file mode 100644 index 0000000000..728f4bfdf3 --- /dev/null +++ b/zig/src/fabric/tp/fate_test.zig @@ -0,0 +1,241 @@ +//! The fate channel between two ranks over real TCP: a failure or a vanished rank ends both, a stop ends neither. +const std = @import("std"); +const sock = @import("sock.zig"); +const fate = @import("fate.zig"); +const bootstrap = @import("bootstrap.zig"); +const Config = @import("config.zig").Config; +const Fate = fate.Fate; + +var exits: std.atomic.Value(u32) = .init(0); +var hooks: std.atomic.Value(u32) = .init(0); + +fn recordExit(code: u8) void { + std.debug.assert(code == fate.exit_code); + _ = exits.fetchAdd(1, .acq_rel); +} + +fn hook(_: *anyopaque, reason: []const u8) void { + std.debug.assert(reason.len > 0); + _ = hooks.fetchAdd(1, .acq_rel); +} + +const grace = 50 * std.time.ns_per_ms; + +/// Two connected ranks' fate channels (rank 0 first), started, with recorded exits and hooks. +fn pair(fates: *[2]Fate) !void { + return pairAt(.{ &fates[0], &fates[1] }); +} + +fn pairAt(fates: [2]*Fate) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var fds0: [fate.max_ranks]sock.Fd = @splat(-1); + var fds1: [fate.max_ranks]sock.Fd = @splat(-1); + fds0[1] = b; + fds1[0] = a; + for (fates, [2][fate.max_ranks]sock.Fd{ fds0, fds1 }, 0..) |f, fds, r| { + f.* = Fate.init(.{ .world = 2, .rank = @intCast(r), .grace_ns = grace, .report_ns = 20 * std.time.ns_per_ms }, fds); + f.exit_fn = recordExit; + f.usable = fakeUsable; + f.onFail(.{ .ctx = f, .call = hook }); + } + for (fates) |f| try f.start(); +} + +fn fakeUsable() u64 { + return 123 << 20; +} + +fn waitExits(want: u32) !u64 { + const t0 = sock.nowNs(); + while (exits.load(.acquire) < want) { + if (sock.nowNs() - t0 > 5 * std.time.ns_per_s) return error.Timeout; + sock.sleepNs(std.time.ns_per_ms); + } + return sock.nowNs() - t0; +} + +fn reset() void { + exits.store(0, .release); + hooks.store(0, .release); +} + +test "rank 1's failure ends both ranks with its reason" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + const mem = fates[0].peerMem(1) orelse return error.NoReport; + try std.testing.expectEqual(@as(u64, 123 << 20), mem.bytes); + fates[1].fatal("prefill", error.OutOfMemory); + const took = try waitExits(2); + try std.testing.expect(took < grace + 500 * std.time.ns_per_ms); + try std.testing.expectEqual(@as(u32, 2), hooks.load(.acquire)); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 failed in prefill: OutOfMemory") != null); +} + +test "rank 0's failure ends rank 1" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + fates[0].fatal("window", error.Unexpected); + _ = try waitExits(2); + try std.testing.expect(std.mem.indexOf(u8, fates[1].reason(), "rank 0 failed in window") != null); +} + +test "a rank that vanishes without a word ends the other" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + // rank 1 killed: its socket closes with no FAIL frame (its own fate thread ends quietly first) + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); +} + +test "a normal stop ends nobody" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + fates[0].expectStop(); + fates[1].expectStop(); + fates[1].close(); + fates[0].close(); + sock.sleepNs(2 * grace); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); +} + +test "a rank failing at boot ends both with its reason, and its channel is down before its owner frees it" { + // a follower's boot failed after its session came up; its Fate is freed while the fate thread could still read it + reset(); + // rank 1's Fate where a Model keeps it: on the heap, its pages unmapped when freed (a stale read faults) + const pa = std.heap.page_allocator; + var f0: Fate = undefined; + const f1 = try pa.create(Fate); + try pairAt(.{ &f0, f1 }); + defer f0.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + f1.failBoot("boot", error.RopeTooShort); + try std.testing.expect(f1.thread == null); + pa.destroy(f1); + _ = try waitExits(2); + // several report periods with rank 1's Fate gone: its thread would have reported from freed memory + sock.sleepNs(5 * 20 * std.time.ns_per_ms); + try std.testing.expect(std.mem.indexOf(u8, f0.reason(), "rank 1 failed in boot: RopeTooShort") != null); +} + +test "rank 0's clean stop ends every rank with no fail-fast exit, whatever order the ranks close in" { + // rank 0 ended its run and closed; the follower, still tearing down, must not see the hang-up as a failure + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + sock.sleepNs(60 * std.time.ns_per_ms); + fates[0].announceStop(); + fates[0].close(); // rank 0 gone before rank 1 even started its own stop + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + fates[1].close(); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} + +// -- a server's drain on two ranks: the plan link and the fate channel in the order a family's stop uses them -- + +const planlink = @import("planlink.zig"); + +/// Plan words a family's rank 0 sends: a round's operation, and a stop word that ends a follower's loop. +const op_round: i64 = 1; +const op_stop: i64 = 4; + +fn linkPair(links: *[2]planlink.PlanLink) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [planlink.max_ranks]sock.Fd = @splat(-1); + var f1: [planlink.max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + links[0] = planlink.PlanLink.init(0, 2, f0); + links[1] = planlink.PlanLink.init(1, 2, f1); +} + +/// Rank 1's loop: each round as it comes, until the stop word, or until the link ends. +const Follower = struct { + link: *planlink.PlanLink, + rounds: u32 = 0, + ended: enum { running, stopped, peer_closed, failed } = .running, + + fn run(f: *Follower) void { + var buf: [16]i64 = undefined; + while (true) { + const msg = f.link.recv(&buf) catch |e| { + f.ended = if (e == error.PeerClosed) .peer_closed else .failed; + return; + }; + if (msg.len == 0) continue; + if (msg[0] == op_stop) { + f.ended = .stopped; + return; + } + if (msg[0] == op_round) f.rounds += 1; + } + } +}; + +test "a drain on rank 0, then its clean stop at a round boundary: rank 1 ends its loop there and nobody exits 70" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + var follower: Follower = .{ .link = &links[1] }; + const t = try std.Thread.spawn(.{}, Follower.run, .{&follower}); + // serving, then draining (rank 0 refuses new work, rounds go on; the fate channel is untouched), past a grace + var sent: u32 = 0; + for (0..30) |_| { + try links[0].send(&.{ op_round, sent }); + sent += 1; + sock.sleepNs(5 * std.time.ns_per_ms); + } + try std.testing.expect(!fates[0].closing.load(.acquire)); + // the stop's order: the stop word first, so the follower ends its loop before the fate STOP and the hang-up + try links[0].send(&.{op_stop}); + fates[0].announceStop(); + links[0].close(); + fates[0].close(); + t.join(); + // rank 1's loop returned: its model closes and the process exits 0 + fates[1].close(); + links[1].close(); + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + try std.testing.expectEqual(.stopped, follower.ended); + try std.testing.expectEqual(sent, follower.rounds); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} + +test "a drain on rank 0 does not disarm fail-fast: rank 1 killed mid-drain still ends rank 0 (exit 70)" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + defer links[0].close(); + for (0..5) |i| try links[0].send(&.{ op_round, @intCast(i) }); + // rank 0 is draining (no stop announced); rank 1 is killed: its sockets close with no word + links[1].close(); + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); + try std.testing.expectEqual(@as(u32, 1), exits.load(.acquire)); +} diff --git a/zig/src/fabric/tp/gpu_bench.zig b/zig/src/fabric/tp/gpu_bench.zig new file mode 100644 index 0000000000..99831e6d77 --- /dev/null +++ b/zig/src/fabric/tp/gpu_bench.zig @@ -0,0 +1,82 @@ +//! Latency and bandwidth of each backend's collectives, eager and graphed; RESULT lines are JSON, busbw as nccl-tests. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Ctx = @import("gpu_checks.zig").Ctx; +const Collective = tp.Collective; + +pub const Kind = enum { all_reduce, all_gather }; + +fn issue(c: *Ctx, comm: Collective, kind: Kind, bytes: usize) !void { + const count = bytes / 2; + switch (kind) { + .all_reduce => try comm.allReduce(c.in.ptr, c.out.ptr, count, .bf16, .sum, c.stream.handle), + .all_gather => try comm.allGather(c.in.ptr, c.out.ptr, count, .bf16, c.stream.handle), + } +} + +/// GPU microseconds a call, `iters` calls issued back to back on the stream. +fn eager(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, iters: u32) !f64 { + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + for (0..10) |_| try issue(c, comm, kind, bytes); + try c.stream.synchronize(); + try comm.barrier(); + try e0.record(c.stream); + for (0..iters) |_| try issue(c, comm, kind, bytes); + try e1.record(c.stream); + try e1.synchronize(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(iters)); +} + +/// GPU microseconds a call inside one graph of `per_graph` calls, replayed `replays` times. +fn graphed(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, per_graph: u32, replays: u32) !f64 { + try cuda.graph.beginCapture(c.stream, .thread_local); + const err: ?tp.collective.Error = blk: { + for (0..per_graph) |_| issue(c, comm, kind, bytes) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (err) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + try exec.upload(c.stream); + try exec.launchOn(c.stream); + try c.stream.synchronize(); + try comm.barrier(); + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + try e0.record(c.stream); + for (0..replays) |_| try exec.launchOn(c.stream); + try e1.record(c.stream); + try e1.synchronize(); + try comm.check(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(per_graph * replays)); +} + +/// Every power of two from `lo` to `hi` bytes a rank, both kinds, eager and graphed. +pub fn sweep(c: *Ctx, comm: Collective, label: []const u8, lo: usize, hi: usize) !void { + for ([_]Kind{ .all_reduce, .all_gather }) |kind| { + var bytes = lo; + while (bytes <= hi) : (bytes *= 2) { + const per_graph: u32 = if (bytes >= 1 << 20) 20 else 100; + const e = try eager(c, comm, kind, bytes, if (bytes >= 1 << 20) 50 else 200); + const g = try graphed(c, comm, kind, bytes, per_graph, 5); + const moved: f64 = @floatFromInt(if (kind == .all_gather) 2 * bytes else bytes); + const factor: f64 = if (kind == .all_gather) 0.5 else 1.0; + const busbw = moved / (g * 1e-6) * factor / 1e9; + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"{t}\",\"bytes\":{d},\"eager_us\":{d:.2},\"graph_us\":{d:.2},\"busbw_GBps\":{d:.2}}}\n", .{ label, kind, bytes, e, g, busbw }); + } + } +} + +/// One decode window's exchanges: `n` all-gathers of one row's bf16 partials (`row_bytes` a rank) in one graph. +pub fn decodeWindow(c: *Ctx, comm: Collective, label: []const u8, n: u32, row_bytes: usize) !void { + const per = try graphed(c, comm, .all_gather, row_bytes, n, 20); + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"decode_window\",\"exchanges\":{d},\"bytes\":{d},\"per_exchange_us\":{d:.2},\"window_us\":{d:.1}}}\n", .{ label, n, row_bytes, per, per * @as(f64, @floatFromInt(n)) }); +} diff --git a/zig/src/fabric/tp/gpu_checks.zig b/zig/src/fabric/tp/gpu_checks.zig new file mode 100644 index 0000000000..5dc4ed6d4c --- /dev/null +++ b/zig/src/fabric/tp/gpu_checks.zig @@ -0,0 +1,188 @@ +//! Every collective on real GPUs against the host reference, eager and replayed from a graph. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Collective = tp.Collective; +const DType = tp.DType; +const Op = tp.Op; +const reference = tp.reference; + +pub const Ctx = struct { + d: *const cuda.Driver, + gpa: std.mem.Allocator, + stream: cuda.Stream, + rank: u32, + /// Device buffers sized for the largest case: input, output (world x), a spare. + in: cuda.DeviceBuffer, + out: cuda.DeviceBuffer, + /// Host scratch: this rank's input, the peer's, the expected and the downloaded output. + mine: []u8, + peer: []u8, + want: []u8, + got: []u8, + failures: u32 = 0, + passes: u32 = 0, + + pub fn init(d: *const cuda.Driver, gpa: std.mem.Allocator, rank: u32, max_bytes: usize) !Ctx { + return .{ + .d = d, + .gpa = gpa, + .stream = try cuda.Stream.init(d, true), + .rank = rank, + .in = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .out = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .mine = try gpa.alloc(u8, max_bytes), + .peer = try gpa.alloc(u8, max_bytes), + .want = try gpa.alloc(u8, 2 * max_bytes), + .got = try gpa.alloc(u8, 2 * max_bytes), + }; + } + + pub fn deinit(c: *Ctx) void { + c.stream.deinit(); + c.in.free(); + c.out.free(); + for ([_][]u8{ c.mine, c.peer, c.want, c.got }) |b| c.gpa.free(b); + } + + fn verdict(c: *Ctx, ok: bool, comptime fmt: []const u8, args: anytype) void { + if (ok) c.passes += 1 else c.failures += 1; + std.debug.print("{s} " ++ fmt ++ "\n", .{if (ok) "PASS" else "FAIL"} ++ args); + } + + /// Both ranks' inputs for case `seed`: this rank's into `mine`, the peer's into `peer` (n bytes each). + fn inputs(c: *Ctx, n: usize, t: DType, seed: u64) void { + reference.fill(c.mine[0..n], t, c.rank, seed); + reference.fill(c.peer[0..n], t, 1 - c.rank, seed); + } + + fn rankInput(c: *Ctx, r: u32, n: usize) []const u8 { + return if (r == c.rank) c.mine[0..n] else c.peer[0..n]; + } + + fn compare(c: *Ctx, comm: Collective, label: []const u8, what: []const u8, t: DType, n: usize, out_off: usize, want: []const u8) !void { + try c.stream.synchronize(); + comm.check() catch |e| { + c.verdict(false, "{s} {s} {t} {d} B: transport error {t}", .{ label, what, t, n, e }); + return; + }; + try c.out.download(out_off, c.got[0..want.len]); + const ok = std.mem.eql(u8, c.got[0..want.len], want); + var first: usize = 0; + if (!ok) first = std.mem.indexOfDiff(u8, c.got[0..want.len], want) orelse 0; + c.verdict(ok, "{s} {s} {t} {d} B a rank{s}", .{ label, what, t, n, if (ok) "" else " (first difference at a byte shown below)" }); + if (!ok) std.debug.print(" byte {d}: got {x} want {x}\n", .{ first, c.got[first], want[first] }); + } + + /// all-reduce (out of place, in place, misaligned by `skew` bytes), all-gather (out of place, in place), exchange. + pub fn collectives(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const reduce_types = [_]DType{ .f32, .bf16, .f16, .i32 }; + for (reduce_types, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + const count = n / t.size(); + c.inputs(n, t, seed + ti); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, t, .sum); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum", t, n, 0, c.want[0..n]); + // in place, at an offset of one element (not 16-byte aligned for 2-byte types) + const skew = t.size(); + try c.out.uploadAsync(skew, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.out.ptr + skew, c.out.ptr + skew, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum in place, skewed", t, n, skew, c.want[0..n]); + } + for ([_]Op{ .max, .min }) |op| { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed + 11); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .f32, op); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, n / 4, .f32, op, c.stream.handle); + try c.compare(comm, label, if (op == .max) "all-reduce max" else "all-reduce min", .f32, n, 0, c.want[0..n]); + } + for ([_]DType{ .u8, .bf16, .f32 }, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + c.inputs(n, t, seed + 20 + ti); + @memcpy(c.want[0..n], c.rankInput(0, n)); + @memcpy(c.want[n..][0..n], c.rankInput(1, n)); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allGather(c.in.ptr, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather", t, n, 0, c.want[0 .. 2 * n]); + try c.out.fill8(0, c.stream.handle); + try c.out.uploadAsync(c.rank * n, c.mine[0..n], c.stream.handle); + try comm.allGather(c.out.ptr + c.rank * n, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather in place", t, n, 0, c.want[0 .. 2 * n]); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.exchange(c.in.ptr, c.out.ptr, n / t.size(), t, 1 - c.rank, c.stream.handle); + try c.compare(comm, label, "exchange", t, n, 0, c.peer[0..n]); + } + } + + /// send then recv one way, then the other way (NCCL behind every backend). + pub fn pointToPoint(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + for (0..2) |dir| { + const sender: u32 = @intCast(dir); + if (c.rank == sender) { + try comm.send(c.in.ptr, n / 4, .f32, 1 - c.rank, c.stream.handle); + try c.stream.synchronize(); + } else { + try comm.recv(c.out.ptr, n / 4, .f32, sender, c.stream.handle); + try c.compare(comm, label, "send/recv", .f32, n, 0, c.peer[0..n]); + } + } + } + + /// A captured chain replayed with fresh inputs, eager calls between replays (the mailbox epoch spans both). + pub fn graphs(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, replays: u32) !void { + const n = n_bytes - n_bytes % 16; + const half = n / 2; // bf16 elements + try cuda.graph.beginCapture(c.stream, .thread_local); + // x1 = sum(x0) [bf16], x2 = gather(x1), x3 = sum(x2) [f32 view], x4 = exchange(x3) + const x0 = c.in.ptr; + const x1 = c.out.ptr; + const x2 = c.in.ptr + n; + const x3 = c.out.ptr + n; + const x4 = c.out.ptr + 3 * n; + const ops: ?tp.collective.Error = blk: { + comm.allReduce(x0, x1, half, .bf16, .sum, c.stream.handle) catch |e| break :blk e; + comm.allGather(x1, x2, half, .bf16, c.stream.handle) catch |e| break :blk e; + comm.allReduce(x2, x3, n / 2, .f32, .sum, c.stream.handle) catch |e| break :blk e; + comm.exchange(x3, x4, n / 2, .f32, 1 - c.rank, c.stream.handle) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (ops) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + const nodes = try g.nodeCount(); + const want1 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want1); + const want3 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want3); + for (0..replays) |k| { + c.inputs(n, .bf16, 1000 + k); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + reference.reduce(want1[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + @memcpy(want1[n..][0..n], want1[0..n]); // the gather of equal sums: twice the same + reference.reduce(want3[0 .. 2 * n], &.{ want1[0 .. 2 * n], want1[0 .. 2 * n] }, .f32, .sum); + try exec.launchOn(c.stream); + try c.compare(comm, label, "graph replay: sum, gather, sum, exchange", .bf16, n, n, want3[0 .. 2 * n]); + try c.compare(comm, label, "graph replay: exchange of the sum", .f32, 2 * n, 3 * n, want3[0 .. 2 * n]); + // an eager call between replays + try c.collectivesQuick(comm, label, k); + } + std.debug.print(" ({s}: {d} graph nodes, {d} replays)\n", .{ label, nodes, replays }); + } + + fn collectivesQuick(c: *Ctx, comm: Collective, label: []const u8, k: usize) !void { + const n = 4096; + c.inputs(n, .bf16, 5000 + k); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + try c.in.uploadAsync(5 * n, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr + 5 * n, c.out.ptr + 5 * n, n / 2, .bf16, .sum, c.stream.handle); + try c.compare(comm, label, "eager between replays", .bf16, n, 5 * n, c.want[0..n]); + } +}; diff --git a/zig/src/fabric/tp/gpu_cli.zig b/zig/src/fabric/tp/gpu_cli.zig new file mode 100644 index 0000000000..dc576fc03b --- /dev/null +++ b/zig/src/fabric/tp/gpu_cli.zig @@ -0,0 +1,66 @@ +//! Validate the runner's command before opening a driver, bootstrap or fate channel. +const std = @import("std"); + +pub const usage = "usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel\n"; +pub const Mode = enum { fatal, kill, transport }; +pub const Failure = struct { victim: u32, mode: Mode, n: u32 = 50, hold: ?u32 = null }; +pub const Command = union(enum) { suite, failfast: Failure, @"roce-host", hca, kernel }; + +pub fn parse(args: []const []const u8) error{InvalidArguments}!Command { + if (args.len == 0) return error.InvalidArguments; + const tag = std.meta.stringToEnum(std.meta.Tag(Command), args[0]) orelse return error.InvalidArguments; + switch (tag) { + .suite, .@"roce-host", .hca, .kernel => { + if (args.len != 1) return error.InvalidArguments; + return switch (tag) { + .suite => .suite, + .@"roce-host" => .@"roce-host", + .hca => .hca, + .kernel => .kernel, + else => unreachable, + }; + }, + .failfast => { + if (args.len < 3 or args.len > 5) return error.InvalidArguments; + const mode = std.meta.stringToEnum(Mode, args[2]) orelse return error.InvalidArguments; + if (mode != .transport and args.len != 3) return error.InvalidArguments; + return .{ .failfast = .{ + .victim = try number(args[1]), + .mode = mode, + .n = if (args.len > 3) try number(args[3]) else 50, + .hold = if (args.len > 4) try number(args[4]) else null, + } }; + }, + } +} + +fn number(s: []const u8) error{InvalidArguments}!u32 { + return std.fmt.parseInt(u32, s, 10) catch error.InvalidArguments; +} + +// Also executed against the Linux entry point with an invalid environment, before any CUDA work. +pub const invalid = [_][]const []const u8{ + &.{}, &.{"UNKNOWN"}, &.{ "suite", "extra" }, &.{ "kernel", "extra" }, + &.{ "hca", "extra" }, &.{ "roce-host", "extra" }, &.{"failfast"}, &.{ "failfast", "0" }, + &.{ "failfast", "0", "typo" }, &.{ "failfast", "bad", "fatal" }, &.{ "failfast", "-1", "kill" }, &.{ "failfast", "4294967296", "fatal" }, + &.{ "failfast", "0", "fatal", "50" }, &.{ "failfast", "0", "kill", "50" }, &.{ "failfast", "0", "transport", "bad" }, &.{ "failfast", "0", "transport", "50", "bad" }, + &.{ "failfast", "0", "transport", "50", "0", "extra" }, +}; + +test "invalid commands are refused before resources open" { + for (invalid) |args| try std.testing.expectError(error.InvalidArguments, parse(args)); +} + +test "documented commands and failfast defaults" { + inline for (.{ "suite", "roce-host", "hca", "kernel" }) |name| { + try std.testing.expectEqualStrings(name, @tagName(try parse(&.{name}))); + } + inline for (.{ Mode.fatal, Mode.kill, Mode.transport }) |mode| { + const f = (try parse(&.{ "failfast", "2", @tagName(mode) })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 2, .mode = mode }, f); + } + const explicit = (try parse(&.{ "failfast", "63", "transport", "0", "0" })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 63, .mode = .transport, .n = 0, .hold = 0 }, explicit); + const n = (try parse(&.{ "failfast", "1", "transport", "123" })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 1, .mode = .transport, .n = 123 }, n); +} diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig new file mode 100644 index 0000000000..1243ab794e --- /dev/null +++ b/zig/src/fabric/tp/gpu_main.zig @@ -0,0 +1,155 @@ +//! `tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel`: one rank of TP=2. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const tp_kernel = @import("tp_kernel"); +const checks = @import("gpu_checks.zig"); +const bench = @import("gpu_bench.zig"); +const cli = @import("gpu_cli.zig"); + +const max_bytes = (4 << 20) + 4096; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + const command = cli.parse(args[1..]) catch { + std.debug.print("{s}", .{cli.usage}); + return 2; + }; + // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first + if (command == .kernel) { + std.debug.print("mailbox kernel {d} B\n", .{tp_kernel.mailbox.len}); + return if (tp_kernel.mailbox.len > 0) 0 else 1; + } + const cfg = try tp.Config.fromEnv(); + if (command == .@"roce-host") return @import("roce_host.zig").run(cfg, init.gpa); + if (command == .hca) return hca(cfg, init.io); + if (command == .failfast and (!cfg.failfast or command.failfast.hold != null) and !hostOnlyAllowed(cfg)) { + std.debug.print("REFUSED: failfast with TF_TP_FAILFAST=0 or a hold is host-only: TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire\n", .{}); + return 2; + } + if (cfg.backend != .nccl and tp_kernel.mailbox.len == 0) { + std.debug.print("FAIL TF_COMM_BACKEND={t} needs the mailbox kernel: rebuild with -Dtp-nvcc or -Dtp-fatbin\n", .{cfg.backend}); + return 2; + } + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, @intCast(cfg.device)); + defer ctx.deinit(); + var name_buf: [128]u8 = undefined; + const t0 = tp.sock.nowNs(); + var s: tp.Session = undefined; + try s.open(&driver, cfg, tp_kernel.mailbox); + const up_ms = @as(f64, @floatFromInt(tp.sock.nowNs() - t0)) / 1e6; + std.debug.print("INFO rank {d}/{d} on {s} (sm_{d}), backend {t}, mailbox kernel {d} B, bootstrap + NCCL init {d:.0} ms\n", .{ cfg.rank, cfg.world, try ctx.name(&name_buf), try ctx.capability(), s.comm().kind(), tp_kernel.mailbox.len, up_ms }); + return switch (command) { + .suite => suite(&s, &driver, init.gpa, cfg), + .failfast => |f| if (f.mode == .transport) + @import("gpu_transport.zig").run(&s, &driver, init.gpa, cfg, f.victim, f.n, f.hold orelse 0) + else + failfast(&s, &driver, init.gpa, cfg, f.victim, @tagName(f.mode)), + else => unreachable, + }; +} + +/// `hca` (host only, no GPU): both ranks' RoCE pick over the bootstrap; this rank's `NCCL_IB_HCA=...` line on stdout. +fn hca(cfg: tp.Config, io: std.Io) !u8 { + var b = tp.Bootstrap.open(cfg) catch |e| { + const s = @as(f64, @floatFromInt(cfg.connect_timeout_ns)) / std.time.ns_per_s; + switch (e) { + // a peer that never runs `hca` (it names its ports, or is not up) leaves this rank waiting until here + error.Timeout => std.debug.print("FAIL hca rank {d}: no peer joined within TF_TP_PLAN_TIMEOUT ({d:.1} s): is its runner up, and leaving TF_TP_ROCE_HCA and NCCL_IB_HCA unset as this one does?\n", .{ cfg.rank, s }), + error.AddressInUse => std.debug.print("FAIL hca rank {d}: port {d} (TF_TP_PORT + 7) is in use on this host\n", .{ cfg.rank, cfg.link_port }), + error.Closed => std.debug.print("FAIL hca rank {d}: the peer closed the bootstrap connection\n", .{cfg.rank}), + else => std.debug.print("FAIL hca rank {d}: bootstrap: {t}\n", .{ cfg.rank, e }), + } + return 1; + }; + defer b.close(); + var specs: [tp.roce_verbs.max_hcas]tp.roce_verbs.HcaSpec = undefined; + const pick = tp.roce_pick.select(cfg, &b, &specs) catch |e| { + std.debug.print("FAIL hca rank {d}: {t} while picking the ports with the peer\n", .{ cfg.rank, e }); + return 1; + }; + var buf: [tp.roce_pick.list_bytes + 16]u8 = undefined; + const line = tp.roce_pick.hcaLine(pick, &specs, &buf) catch |e| { + const why = switch (e) { + error.LibraryUnavailable => "could not load libibverbs.so.1 (LibraryUnavailable); check the RDMA userspace installation", + error.MissingSymbol => "libibverbs.so.1 lacks required symbols (MissingSymbol); check the RDMA userspace installation", + error.NoRoceDevice => "no active Ethernet RDMA port with an IPv4 RoCE v2 GID on this host: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA", + error.NoSharedSubnet => "no RoCE port here shares an IPv4 /24 with the peer's: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA", + else => @errorName(e), + }; + std.debug.print("FAIL hca rank {d}: {s}\n", .{ cfg.rank, why }); + return 1; + }; + var names: [tp.roce_pick.list_bytes]u8 = undefined; + const how = if (pick.explicit) "TF_TP_ROCE_HCA" else "on a /24 the peer shares"; + std.debug.print("INFO hca rank {d}: {s} ({s})\n", .{ cfg.rank, try tp.roce_pick.joinNames(specs[0..pick.n], true, &names), how }); + try std.Io.File.stdout().writeStreamingAll(io, line); + return 0; +} + +fn suite(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, max_bytes); + defer c.deinit(); + var backends: [2]struct { comm: tp.Collective, label: []const u8 } = undefined; + var n_backends: usize = 1; + backends[0] = .{ .comm = s.nccl.iface(), .label = "nccl" }; + if (s.comm().kind() != .nccl) { + backends[1] = .{ .comm = s.comm(), .label = if (s.roce != null) "roce" else "mailbox" }; + n_backends = 2; + } + const sizes = [_]usize{ 6, 1000, 4096, 65536, 256 * 1024, 256 * 1024 + 16, 1 << 20, 4 << 20 }; + for (backends[0..n_backends]) |b| { + for (sizes, 0..) |n, i| try c.collectives(b.comm, b.label, n, 100 * i); + try c.pointToPoint(b.comm, b.label, 64 * 1024, 7); + for ([_]usize{ 4096, 65536, 1 << 20 }) |n| try c.graphs(b.comm, b.label, n, 4); + try b.comm.barrier(); + } + for (backends[0..n_backends]) |b| { + try bench.sweep(&c, b.comm, b.label, 4096, 4 << 20); + try bench.decodeWindow(&c, b.comm, b.label, 86, 8192); + } + // the stop: rank 0 marks it expected first, then both meet and close + if (s.fate) |*f| f.expectStop(); + try s.comm().barrier(); + std.debug.print("SUMMARY rank {d}: {d} passed, {d} failed\n", .{ cfg.rank, c.passes, c.failures }); + s.close(); + return if (c.failures == 0) 0 else 1; +} + +/// The failed-transport check's host-only forms (fail-fast off, a hold): TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire. +fn hostOnlyAllowed(cfg: tp.Config) bool { + const flag = std.mem.eql(u8, std.mem.span(std.c.getenv("TF_TP_ALLOW_HOST_ONLY") orelse ""), "1"); + return flag and (cfg.backend != .roce or cfg.roce_wire == .shm); +} + +pub fn realNow() f64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.REALTIME, &ts); + return @as(f64, @floatFromInt(ts.sec)) + @as(f64, @floatFromInt(ts.nsec)) / 1e9; +} + +/// Exchanges until the victim fails; the survivor is left waiting inside one, which only the fate channel ends. +fn failfast(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, mode: []const u8) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, 1 << 20); + defer c.deinit(); + const comm = s.comm(); + var k: u32 = 0; + while (true) : (k += 1) { + if (k == 50 and cfg.rank == victim) { + std.debug.print("INJECT {s} rank {d} at {d:.6}\n", .{ mode, cfg.rank, realNow() }); + if (std.mem.eql(u8, mode, "kill")) _ = std.c.raise(.KILL); + s.fatal("window", error.Injected); + return 3; // not reached: the process exits 70 + } + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + c.stream.synchronize() catch |e| { + std.debug.print("rank {d}: stream error {t} after the peer failed\n", .{ cfg.rank, e }); + }; + comm.check() catch |e| { + std.debug.print("rank {d}: transport reports {t} at exchange {d}; waiting for the fate channel\n", .{ cfg.rank, e, k }); + while (true) tp.sock.sleepNs(std.time.ns_per_s); + }; + } +} diff --git a/zig/src/fabric/tp/gpu_transport.zig b/zig/src/fabric/tp/gpu_transport.zig new file mode 100644 index 0000000000..49cd625844 --- /dev/null +++ b/zig/src/fabric/tp/gpu_transport.zig @@ -0,0 +1,71 @@ +//! `tf-tp-test failfast V transport [N [HOLD_S]]`: the failed-transport path on real HCAs, injected on an idle link. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const checks = @import("gpu_checks.zig"); +const realNow = @import("gpu_main.zig").realNow; + +/// N clean sums, an NCCL meet, the victim's `Roce.inject` once every link is idle, then its close (fail-fast: exit 70). +pub fn run(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, n: u32, hold_s: u32) !u8 { + const r = s.roce orelse { + std.debug.print("failfast transport needs TF_COMM_BACKEND=roce (the session runs {t})\n", .{s.comm().kind()}); + return 2; + }; + var c = try checks.Ctx.init(d, gpa, cfg.rank, 4096); + defer c.deinit(); + const comm = s.comm(); + for (0..n) |_| { + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + try c.stream.synchronize(); + try comm.check(); + } + try s.nccl.iface().barrier(); + std.debug.print("rank {d}: {d} clean exchanges; both ranks met over NCCL\n", .{ cfg.rank, n }); + // the inject only on an idle link: every rank's posts completed, else no rank injects + const t_idle = tp.sock.nowNs(); + while (!r.linkIdle() and tp.sock.nowNs() - t_idle < 2 * std.time.ns_per_s) tp.sock.sleepNs(std.time.ns_per_ms); + var line: [384]u8 = undefined; + std.debug.print("rank {d}: link {s} before the inject: {s}\n", .{ cfg.rank, if (r.linkIdle()) "idle" else "BUSY", r.describeLink(&line) }); + if (try idleRanks(s, &c, r.linkIdle()) != cfg.world) { + std.debug.print("REFUSED rank {d}: a link was not idle; no inject, a clean stop\n", .{cfg.rank}); + if (s.fate) |*f| f.expectStop(); + try s.nccl.iface().barrier(); + s.close(); + return 2; + } + if (cfg.rank == victim) { + std.debug.print("INJECT transport rank {d} at {d:.6}\n", .{ cfg.rank, realNow() }); + r.inject(); + comm.check() catch |e| std.debug.print("rank {d}: transport reports {t}\n", .{ cfg.rank, e }); + } else if (s.fate != null) { + while (true) tp.sock.sleepNs(std.time.ns_per_s); // the fate channel ends this rank (exit 70) + } + const kept0 = tp.roce.Roce.kept.load(.monotonic); + s.close(); // fail-fast on: exit 70 from here + const kept = tp.roce.Roce.kept.load(.monotonic) - kept0; + const want: u32 = @intFromBool(cfg.rank == victim); + const live = cfg.rank != victim or mapped(r.region); + std.debug.print("rank {d}: closed; regions kept {d} (want {d}){s}\n", .{ cfg.rank, kept, want, if (cfg.rank != victim) "" else if (live) ", region still mapped" else ", region NOT mapped" }); + if (hold_s > 0) { + std.debug.print("HOLD rank {d} pid {d} for {d} s\n", .{ cfg.rank, std.c.getpid(), hold_s }); + tp.sock.sleepNs(@as(u64, hold_s) * std.time.ns_per_s); + } + const ok = kept == want and live; + std.debug.print("{s} rank {d}\n", .{ if (ok) "PASS" else "FAIL", cfg.rank }); + return if (ok) 0 else 1; +} + +fn mapped(region: []const u8) bool { + var vec: [1]u8 = undefined; + return std.os.linux.errno(std.os.linux.mincore(@constCast(region.ptr), 4096, &vec)) == .SUCCESS; +} + +/// The ranks whose link was idle, summed over NCCL (one f32 a rank). +fn idleRanks(s: *tp.Session, c: *checks.Ctx, idle: bool) !u32 { + var v: [1]f32 = .{@floatFromInt(@intFromBool(idle))}; + try c.in.uploadAsync(0, std.mem.asBytes(&v), c.stream.handle); + try s.nccl.iface().allReduce(c.in.ptr, c.out.ptr, 1, .f32, .sum, c.stream.handle); + try c.stream.synchronize(); + try c.out.download(0, std.mem.asBytes(&v)); + return @intFromFloat(v[0]); +} diff --git a/zig/src/fabric/tp/host.zig b/zig/src/fabric/tp/host.zig new file mode 100644 index 0000000000..051daea537 --- /dev/null +++ b/zig/src/fabric/tp/host.zig @@ -0,0 +1,293 @@ +//! In-process ranks over host memory: the `Collective` contract without a GPU (desyncs refused, abort wakes waiters). +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub const max_ranks = 8; + +const Call = enum(u8) { none, all_reduce, all_gather, all_gather_v, exchange, barrier }; + +const Post = struct { call: Call = .none, send: DevicePtr = 0, recv: DevicePtr = 0, count: usize = 0, t: DType = .u8, op: Op = .sum, peer: u32 = 0, lens: DevicePtr = 0 }; + +/// One-way message slot from rank `src` to rank `dst` (send / recv). +const Slot = struct { buf: std.atomic.Value(DevicePtr) = .init(0), bytes: usize = 0, taken: std.atomic.Value(bool) = .init(false) }; + +pub const Group = struct { + world: u32, + /// A wait longer than this is `error.Timeout` (a peer that never calls). + timeout_ns: u64 = 10 * std.time.ns_per_s, + posts: [max_ranks]Post = @splat(.{}), + arrived: std.atomic.Value(u32) = .init(0), + generation: std.atomic.Value(u32) = .init(0), + aborted: std.atomic.Value(bool) = .init(false), + /// Ranks (a bit each) whose transport refuses device lengths: their `all_gather_v` moves whole strides (tests). + var_refuse: u32 = 0, + slots: [max_ranks][max_ranks]Slot = @splat(@splat(.{})), + + pub fn init(world: u32) Group { + std.debug.assert(world >= 1 and world <= max_ranks); + return .{ .world = world }; + } + + pub fn rank(g: *Group, r: u32) Rank { + std.debug.assert(r < g.world); + return .{ .group = g, .me = r }; + } + + /// Sense-reversing barrier over every rank; `error.Aborted` once any rank aborts. + fn meet(g: *Group) Error!void { + const gen = g.generation.load(.acquire); + if (g.arrived.fetchAdd(1, .acq_rel) + 1 == g.world) { + g.arrived.store(0, .release); + _ = g.generation.fetchAdd(1, .acq_rel); + return if (g.aborted.load(.acquire)) error.Aborted else {}; + } + const t0 = now(); + while (g.generation.load(.acquire) == gen) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + if (g.aborted.load(.acquire)) return error.Aborted; + } +}; + +fn now() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +fn bytesAt(p: DevicePtr, n: usize) []u8 { + if (n == 0) return &.{}; + return @as([*]u8, @ptrFromInt(p))[0..n]; +} + +pub const Rank = struct { + group: *Group, + me: u32, + /// `varAgreed`'s answer, once this communicator agreed. + var_agreed: ?bool = null, + + pub fn iface(self: *Rank) collective.Collective { + return .{ .ptr = self, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .host, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + + fn self_(ptr: *anyopaque) *Rank { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).group.world; + } + + /// Posts this rank's call, waits for all, and refuses calls that do not match rank 0's. + fn enter(self: *Rank, p: Post) Error!void { + const g = self.group; + if (g.aborted.load(.acquire)) return error.Aborted; + g.posts[self.me] = p; + try g.meet(); + const r0 = g.posts[0]; + if (r0.call != p.call or r0.count != p.count or r0.t != p.t or r0.op != p.op) { + abort(self); + return error.Invalid; + } + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_reduce, .send = s, .recv = r, .count = count, .t = t, .op = op }); + const n = count * t.size(); + var ins: [max_ranks][]const u8 = undefined; + for (0..g.world) |i| ins[i] = bytesAt(g.posts[i].send, n); + // the output may alias this rank's input: reduce into scratch first, then meet before writing + var scratch: [1 << 16]u8 = undefined; + if (n > scratch.len) return error.Unsupported; + if (n > 0) reference.reduce(scratch[0..n], ins[0..g.world], t, op); + try g.meet(); + @memcpy(bytesAt(r, n), scratch[0..n]); + try g.meet(); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather, .send = s, .recv = r, .count = count, .t = t }); + const n = count * t.size(); + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| @memcpy(out[i * n ..][0..n], bytesAt(g.posts[i].send, n)); + try g.meet(); + } + + fn lensAt(p: DevicePtr, world: u32) []const i32 { + return @as([*]const i32, @ptrFromInt(p))[0..world]; + } + + /// Copies only each rank's lens[r] bytes (`var_refuse` ranks take whole strides); unequal lengths are refused. + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather_v, .send = s, .recv = r, .count = stride, .t = t, .lens = lens }); + const n = stride * t.size(); + const mine = lensAt(lens, g.world); + for (0..g.world) |i| if (!std.mem.eql(i32, mine, lensAt(g.posts[i].lens, g.world))) { + abort(self); + return error.Invalid; + }; + const whole = g.var_refuse & (@as(u32, 1) << @intCast(self.me)) != 0; + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| { + if (mine[i] < 0 or mine[i] > n) { + abort(self); + return error.Invalid; + } + const len: usize = if (whole) n else @intCast(mine[i]); + @memcpy(out[i * n ..][0..len], bytesAt(g.posts[i].send, len)); + } + try g.meet(); + } + + fn varFits(ptr: *anyopaque, bytes: usize) bool { + const self = self_(ptr); + return bytes > 0 and self.group.var_refuse & (@as(u32, 1) << @intCast(self.me)) == 0; + } + + /// One int all-gather of `varFits(16)` the first time, then the kept answer. + fn varAgreed(ptr: *anyopaque) Error!bool { + const self = self_(ptr); + if (self.var_agreed) |v| return v; + const mine: i32 = @intFromBool(varFits(ptr, 16)); + var all: [max_ranks]i32 = undefined; + try allGather(ptr, @intFromPtr(&mine), @intFromPtr(&all), 1, .i32, null); + self.var_agreed = std.mem.min(i32, all[0..self.group.world]) == 1; + return self.var_agreed.?; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + try self.enter(.{ .call = .exchange, .send = s, .recv = r, .count = count, .t = t, .peer = peer }); + if (g.posts[peer].peer != self.me) { + abort(self); + return error.Invalid; + } + const n = count * t.size(); + @memcpy(bytesAt(r, n), bytesAt(g.posts[peer].send, n)); + try g.meet(); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[self.me][peer]; + slot.bytes = count * t.size(); + slot.taken.store(false, .release); + slot.buf.store(if (buf == 0) 1 else buf, .release); + try waitFor(g, &slot.taken, true); + slot.buf.store(0, .release); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[peer][self.me]; + const t0 = now(); + var src = slot.buf.load(.acquire); + while (src == 0) : (src = slot.buf.load(.acquire)) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + const n = count * t.size(); + if (slot.bytes != n) { + abort(self); + return error.Invalid; + } + @memcpy(bytesAt(buf, n), bytesAt(src, n)); + slot.taken.store(true, .release); + } + + fn waitFor(g: *Group, flag: *std.atomic.Value(bool), want: bool) Error!void { + const t0 = now(); + while (flag.load(.acquire) != want) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + } + + fn barrier(ptr: *anyopaque) Error!void { + const self = self_(ptr); + try self.enter(.{ .call = .barrier }); + try self.group.meet(); + } + + fn check(ptr: *anyopaque) Error!void { + if (self_(ptr).group.aborted.load(.acquire)) return error.Aborted; + } + + fn abort(ptr: *anyopaque) void { + self_(ptr).group.aborted.store(true, .release); + } +}; + +/// Runs `body(rank_collective, rank, ctx)` on `world` threads at once and returns the first error any rank hit. +pub fn run(world: u32, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + return runWith(world, .{}, ctx, body); +} + +pub const Options = struct { + /// `Group.var_refuse`: ranks whose transport takes no device lengths. + var_refuse: u32 = 0, +}; + +pub fn runWith(world: u32, opts: Options, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + var group = Group.init(world); + group.var_refuse = opts.var_refuse; + var ranks: [max_ranks]Rank = undefined; + var errs: [max_ranks]?anyerror = @splat(null); + var threads: [max_ranks]std.Thread = undefined; + const Wrap = struct { + fn go(r: *Rank, c: @TypeOf(ctx), e: *?anyerror) void { + body(r.iface(), r.me, c) catch |err| { + e.* = err; + r.group.aborted.store(true, .release); + }; + } + }; + for (0..world) |i| { + ranks[i] = group.rank(@intCast(i)); + threads[i] = try std.Thread.spawn(.{}, Wrap.go, .{ &ranks[i], ctx, &errs[i] }); + } + for (threads[0..world]) |t| t.join(); + for (errs[0..world]) |e| if (e) |err| return err; +} diff --git a/zig/src/fabric/tp/host_test.zig b/zig/src/fabric/tp/host_test.zig new file mode 100644 index 0000000000..7de32008af --- /dev/null +++ b/zig/src/fabric/tp/host_test.zig @@ -0,0 +1,154 @@ +//! The `Collective` contract on the in-process ranks: every call's bits against the reference, desyncs and aborts. +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const host = @import("host.zig"); +const Collective = collective.Collective; +const DType = collective.DType; + +fn addr(b: []u8) collective.DevicePtr { + return @intFromPtr(b.ptr); +} + +const n = 96; + +const Calls = struct { + fn go(c: Collective, rank: u32, _: void) !void { + const world = c.world(); + inline for (.{ DType.f32, DType.bf16, DType.i32, DType.f16 }) |t| { + const bytes = n * t.size(); + var mine: [n * 8]u8 = undefined; + var out: [n * 8 * host.max_ranks]u8 = undefined; + reference.fill(mine[0..bytes], t, rank, 7); + // all-reduce against the reference sum of every rank's input, in rank order + var ins: [host.max_ranks][n * 8]u8 = undefined; + var views: [host.max_ranks][]const u8 = undefined; + for (0..world) |r| { + reference.fill(ins[r][0..bytes], t, @intCast(r), 7); + views[r] = ins[r][0..bytes]; + } + var want: [n * 8]u8 = undefined; + reference.reduce(want[0..bytes], views[0..world], t, .sum); + try c.allReduce(addr(&mine), addr(&out), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], out[0..bytes]); + // in place + try c.allReduce(addr(&mine), addr(&mine), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], mine[0..bytes]); + // all-gather: rank order + reference.fill(mine[0..bytes], t, rank, 7); + try c.allGather(addr(&mine), addr(&out), n, t, null); + for (0..world) |r| try std.testing.expectEqualSlices(u8, views[r], out[r * bytes ..][0..bytes]); + } + try c.barrier(); + } +}; + +test "every call on two and three in-process ranks equals the reference" { + try host.run(2, {}, Calls.go); + try host.run(3, {}, Calls.go); +} + +const Pairwise = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var mine: [64]u8 = @splat(@intCast(rank + 1)); + var got: [64]u8 = undefined; + const peer = 1 - rank; + try c.exchange(addr(&mine), addr(&got), 64, .u8, peer, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(@intCast(peer + 1))), &got); + // send / recv, ordered so neither waits forever + if (rank == 0) { + try c.send(addr(&mine), 16, .f32, 1, null); + } else { + var buf: [64]u8 = undefined; + try c.recv(addr(&buf), 16, .f32, 0, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(1)), &buf); + } + try std.testing.expectError(error.Invalid, c.exchange(addr(&mine), addr(&got), 64, .u8, rank, null)); + } +}; + +test "exchange, send and recv between two ranks" { + try host.run(2, {}, Pairwise.go); +} + +const Desync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + // the ranks disagree on the count: both see it instead of hanging + const r = c.allGather(addr(&a), addr(&b), if (rank == 0) 16 else 32, .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "a desynchronised call is refused on every rank" { + try host.run(2, {}, Desync.go); +} + +const Aborts = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [16]u8 = @splat(0); + var b: [32]u8 = undefined; + if (rank == 1) { + // rank 1 never calls: it aborts after a while, which must free rank 0's waiting call + @import("sock.zig").sleepNs(50 * std.time.ns_per_ms); + c.abort(); + return; + } + try std.testing.expectError(error.Aborted, c.allGather(addr(&a), addr(&b), 16, .u8, null)); + try std.testing.expectError(error.Aborted, c.check()); + try std.testing.expectError(error.Aborted, c.barrier()); + } +}; + +test "abort frees a rank waiting in a collective" { + try host.run(2, {}, Aborts.go); +} + +const VarGather = struct { + fn go(c: Collective, rank: u32, refuse: u32) !void { + const world = c.world(); + const stride = 200; + const whole = refuse & (@as(u32, 1) << @intCast(rank)) != 0; + try std.testing.expectEqual(!whole, c.varFits(stride)); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); // kept: no second collective + var ins: [host.max_ranks][stride]u8 = undefined; + for (0..world) |r| reference.fill(&ins[r], .u8, @intCast(r), 11); + // lengths of 0, odd, 16 and the whole stride, in bytes; .f32 strides count elements + const cases = [_][3]i32{ .{ 0, 17, 200 }, .{ 200, 0, 3 }, .{ 16, 16, 16 }, .{ 199, 200, 1 } }; + inline for (.{ DType.u8, DType.f32 }) |t| for (cases) |lens| { + var out: [stride * host.max_ranks]u8 = @splat(0xEE); + try c.allGatherV(addr(&ins[rank]), addr(&out), stride / t.size(), @intFromPtr(&lens), t, null); + for (0..world) |r| { + const len: usize = @intCast(lens[r]); + try std.testing.expectEqualSlices(u8, ins[r][0..len], out[r * stride ..][0..len]); + // past a length: the transport's choice (untouched here, or the whole stride where lengths are refused) + const past = out[r * stride + len .. (r + 1) * stride]; + if (whole) try std.testing.expectEqualSlices(u8, ins[r][len..], past) else for (past) |x| try std.testing.expectEqual(@as(u8, 0xEE), x); + } + }; + try c.barrier(); + } +}; + +test "a variable-length all-gather moves each rank's length, and the ranks agree on device lengths once" { + for ([_]u32{ 2, 3 }) |w| { + try host.runWith(w, .{}, @as(u32, 0), VarGather.go); + try host.runWith(w, .{ .var_refuse = 2 }, @as(u32, 2), VarGather.go); + } +} + +const VarDesync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + const lens = [2]i32{ 8, if (rank == 0) 8 else 16 }; + const r = c.allGatherV(addr(&a), addr(&b), 64, @intFromPtr(&lens), .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "lengths that differ between ranks are refused on every rank" { + try host.run(2, {}, VarDesync.go); +} diff --git a/zig/src/fabric/tp/hybrid.zig b/zig/src/fabric/tp/hybrid.zig new file mode 100644 index 0000000000..e5f6d72c31 --- /dev/null +++ b/zig/src/fabric/tp/hybrid.zig @@ -0,0 +1,183 @@ +//! Small exchanges through the mailbox (one host) or RoCE (two hosts), the rest through NCCL; ranks always agree. +const std = @import("std"); +const collective = @import("collective.zig"); +const Mailbox = @import("mailbox.zig").Mailbox; +const Roce = @import("roce.zig").Roce; +const Nccl = @import("nccl.zig").Nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +/// The one-kernel path for small shards: both expose the same calls. +pub const Small = union(enum) { + mailbox: *Mailbox, + roce: *Roce, + + fn fits(s: Small, n: usize) bool { + return switch (s) { + inline else => |p| p.fits(n), + }; + } + + fn sum(s: Small, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.sum(src, dst, count, t, stream), + }; + } + + fn gather(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.allGatherBytes(src, dst, n, stream), + }; + } + + fn exchange(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.exchangeBytes(src, dst, n, stream), + }; + } + + fn check(s: Small) Error!void { + return switch (s) { + inline else => |p| p.check(), + }; + } + + fn abort(s: Small) void { + switch (s) { + inline else => |p| p.abort(), + } + } +}; + +/// Whether a variable-length gather of `n`-byte strides one op cannot take goes RoCE in slot-sized ops (needs the ABI). +pub fn chunkRoute(chunked_bytes: usize, n: usize, single: bool, abi: bool) bool { + return abi and !single and n > 0 and n <= chunked_bytes; +} + +pub const Hybrid = struct { + small: Small, + nccl: *Nccl, + /// Largest shard a rank sends through the small path. + max_bytes: usize, + /// Largest stride a variable-length gather takes on RoCE in slot-sized ops (TF_TP_ROCE_CHUNKED_KB; 0: none). + chunked_bytes: usize = 0, + /// Every rank's `varFits(16)`, agreed at `Session.open` (the bootstrap's all-gather): `varAgreed`'s answer. + var_agreed: bool = false, + + pub fn iface(h: *Hybrid) collective.Collective { + return .{ .ptr = h, .vtable = if (h.small == .roce) &roce_vtable else &mailbox_vtable }; + } + + fn vtable(comptime kind: collective.Kind) collective.Collective.VTable { + return .{ + .kind = kind, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + } + + const mailbox_vtable = vtable(.hybrid); + const roce_vtable = vtable(.roce); + + fn self_(ptr: *anyopaque) *Hybrid { + return @ptrCast(@alignCast(ptr)); + } + + fn routed(h: *const Hybrid, n: usize) bool { + return n <= h.max_bytes and h.small.fits(n); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.size; + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (op == .sum and (t == .f32 or t == .bf16) and h.routed(count * t.size())) return h.small.sum(s, r, count, t, stream); + return Nccl.allReduce(h.nccl, s, r, count, t, op, stream); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (h.routed(count * t.size())) return h.small.gather(s, r, count * t.size(), stream); + return Nccl.allGather(h.nccl, s, r, count, t, stream); + } + + /// RoCE moves each rank's device length; the mailbox and NCCL move whole strides (the same bytes where lens says). + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + const n = stride * t.size(); + if (singleFits(h, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + if (varFits(ptr, n)) return h.small.roce.allGatherVChunked(s, r, n, lens, stream); + return allGather(ptr, s, r, stride, t, stream); + } + + fn singleFits(h: *const Hybrid, n: usize) bool { + return switch (h.small) { + .roce => |p| n > 0 and h.routed(n) and p.d != null, + .mailbox => false, + }; + } + + /// RoCE with its device side, in one op or in slot-sized ops (the route is the size's alone: every rank alike). + pub fn varFits(ptr: *anyopaque, n: usize) bool { + const h = self_(ptr); + const single = singleFits(h, n); + return single or switch (h.small) { + .roce => |p| chunkRoute(h.chunked_bytes, n, single, p.d != null and p.chunk_abi), + .mailbox => false, + }; + } + + fn varAgreed(ptr: *anyopaque) Error!bool { + return self_(ptr).var_agreed; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const h = self_(ptr); + try collective.checkPeer(h.nccl.me, h.nccl.size, peer); + if (h.routed(count * t.size())) return h.small.exchange(s, r, count * t.size(), stream); + return Nccl.exchange(h.nccl, s, r, count, t, peer, stream); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).send(buf, count, t, peer, stream); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).recv(buf, count, t, peer, stream); + } + + fn barrier(ptr: *anyopaque) Error!void { + return Nccl.iface(self_(ptr).nccl).barrier(); + } + + fn check(ptr: *anyopaque) Error!void { + const h = self_(ptr); + try h.small.check(); + try Nccl.iface(h.nccl).check(); + } + + fn abort(ptr: *anyopaque) void { + const h = self_(ptr); + h.small.abort(); + Nccl.iface(h.nccl).abort(); + } +}; diff --git a/zig/src/fabric/tp/log.zig b/zig/src/fabric/tp/log.zig new file mode 100644 index 0000000000..4e3234ab49 --- /dev/null +++ b/zig/src/fabric/tp/log.zig @@ -0,0 +1,15 @@ +//! TP's log lines, silent under test: a test's stderr is the build runner's channel. +const std = @import("std"); +const quiet = @import("builtin").is_test; + +pub fn err(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.err(format, args); +} + +pub fn warn(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.warn(format, args); +} + +pub fn info(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.info(format, args); +} diff --git a/zig/src/fabric/tp/mailbox.cu b/zig/src/fabric/tp/mailbox.cu new file mode 100644 index 0000000000..843b4f1af2 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.cu @@ -0,0 +1,564 @@ +// One-shot two-rank exchange through pinned host memory both GPUs map (the TP mailbox; mailbox.zig is its host side). + +#include +#include + +// Protocol adapted from b12x's RoCEnante (Apache-2.0): a device epoch, slots by seq parity, a seq flag after the data. +namespace tftp { + +constexpr int kMaxBlocks = 16; +constexpr int kFlagStride = 32; // words: one 128-byte line per flag +// ctrl words: [0] abort (any rank, any host thread); per rank r at 16 + 8 r: failed, seq, seen, block, why, waited us +constexpr int kCtrlAbort = 0; +constexpr int kCtrlRank = 16; + +enum Mode : uint32_t { kGather = 0, kSumF32 = 1, kSumBf16 = 2, kExchange = 3 }; + +struct MailArgs { + const uint8_t *in; // this rank's shard (device) + uint8_t *out; // gather [2 n], exchange / sums [n] + uint8_t *peer_slots; // the peer's inbox data, as this GPU maps it + uint32_t *peer_flags; // the peer's inbox flags + const uint8_t *my_slots; // this rank's inbox data + const uint32_t *my_flags; // this rank's inbox flags + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison + uint32_t *ctrl; // shared control words + uint64_t nbytes; // bytes a rank + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; +}; + +__device__ __forceinline__ uint32_t ld_relaxed_gpu(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_acquire_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 ld_relaxed_sys_v4(const void *p) { + uint4 v; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u8(const uint8_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u8 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void st_relaxed_sys(uint32_t *p, uint32_t v) { + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_sys(uint32_t *p, uint32_t v) { + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_gpu(uint32_t *p, uint32_t v) { + asm volatile("st.release.gpu.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_gpu(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_sys(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.sys.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_cas_gpu(uint32_t *p, uint32_t cmp, uint32_t v) { + uint32_t old; + asm volatile("atom.relaxed.gpu.global.cas.b32 %0, [%1], %2, %3;" : "=r"(old) : "l"(p), "r"(cmp), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ void fence_sc_sys() { asm volatile("fence.sc.sys;" ::: "memory"); } + +__device__ __forceinline__ uint64_t globaltimer() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// dst[i] <- src[i] for i in [0, n), grid-stride, 16-byte vectors when aligned; SysSrc: system-scope loads. +template +__device__ __forceinline__ void copy_bytes(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 v = SysSrc ? ld_relaxed_sys_v4(src + (i << 4)) : reinterpret_cast(src)[i]; + reinterpret_cast(dst)[i] = v; + } + } else { + for (uint64_t i = index; i < n; i += stride) { + dst[i] = static_cast(SysSrc ? ld_relaxed_sys_u8(src + i) : src[i]); + } + } +} + +// TF_TP_ROCE_FAST: copy_bytes with two grid strides' loads issued before their stores (same bytes, one round trip). +template +__device__ __forceinline__ void copy_bytes2(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) != 0u) { + copy_bytes(dst, src, n, index, stride); + return; + } + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += 2u * stride) { + const uint64_t j = i + stride; + uint4 v0, v1 = make_uint4(0u, 0u, 0u, 0u); + if (SysSrc) { + v0 = ld_relaxed_sys_v4(src + (i << 4)); + if (j < packs) v1 = ld_relaxed_sys_v4(src + (j << 4)); + } else { + v0 = reinterpret_cast(src)[i]; + if (j < packs) v1 = reinterpret_cast(src)[j]; + } + reinterpret_cast(dst)[i] = v0; + if (j < packs) reinterpret_cast(dst)[j] = v1; + } +} + +template +__device__ __forceinline__ void copy_any(bool fast, uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + if (fast) { + copy_bytes2(dst, src, n, index, stride); + } else { + copy_bytes(dst, src, n, index, stride); + } +} + +__device__ __forceinline__ uint32_t add_bf16x2(uint32_t x, uint32_t y) { + const float2 a = __bfloat1622float2(*reinterpret_cast(&x)); + const float2 b = __bfloat1622float2(*reinterpret_cast(&y)); + const __nv_bfloat162 r = __floats2bfloat162_rn(a.x + b.x, a.y + b.y); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t add_f32(uint32_t x, uint32_t y) { + return __float_as_uint(__uint_as_float(x) + __uint_as_float(y)); +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u16(const uint8_t *p) { + uint16_t v; + asm volatile("ld.relaxed.sys.global.u16 %0, [%1];" : "=h"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint16_t add_bf16(uint16_t x, uint16_t y) { + const __nv_bfloat16 r = __float2bfloat16_rn(__bfloat162float(__ushort_as_bfloat16(x)) + + __bfloat162float(__ushort_as_bfloat16(y))); + return __bfloat16_as_ushort(r); +} + +// out <- mine + peer's, rank 0's first; IEEE addition commutes, so both ranks get the same bits. +template +__device__ __forceinline__ void sum_bytes(uint8_t *out, const uint8_t *mine, const uint8_t *peer, uint64_t n, + uint32_t rank, uint64_t index, uint64_t stride) { + const uint64_t al = reinterpret_cast(out) | reinterpret_cast(mine) | + reinterpret_cast(peer) | n; + if ((al & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 m = reinterpret_cast(mine)[i]; + const uint4 p = ld_relaxed_sys_v4(peer + (i << 4)); + const uint4 a = rank == 0 ? m : p; + const uint4 b = rank == 0 ? p : m; + uint4 r; + r.x = Bf16 ? add_bf16x2(a.x, b.x) : add_f32(a.x, b.x); + r.y = Bf16 ? add_bf16x2(a.y, b.y) : add_f32(a.y, b.y); + r.z = Bf16 ? add_bf16x2(a.z, b.z) : add_f32(a.z, b.z); + r.w = Bf16 ? add_bf16x2(a.w, b.w) : add_f32(a.w, b.w); + reinterpret_cast(out)[i] = r; + } + } else if (Bf16) { + for (uint64_t i = index; i < (n >> 1); i += stride) { + const uint16_t m = reinterpret_cast(mine)[i]; + const uint16_t p = static_cast(ld_relaxed_sys_u16(peer + (i << 1))); + reinterpret_cast(out)[i] = rank == 0 ? add_bf16(m, p) : add_bf16(p, m); + } + } else { + for (uint64_t i = index; i < (n >> 2); i += stride) { + const uint32_t m = reinterpret_cast(mine)[i]; + const uint32_t p = ld_relaxed_sys(reinterpret_cast(peer) + i); + reinterpret_cast(out)[i] = rank == 0 ? add_f32(m, p) : add_f32(p, m); + } + } +} + +// Thread 0 of a block: spin until the peer block's flag holds seq; false on abort (why 1) or timeout (why 2). +__device__ __noinline__ bool wait_flag(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, uint64_t timeout_ns, + uint32_t *seen, uint32_t *why, uint32_t *waited_us) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *seen = v; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + return *seen == seq; + } + } + } +} + +// TF_TP_ROCE_FAST: wait_flag for one of several pollers of a flag; `done` (shared) ends the others once one saw it. +__device__ __noinline__ bool wait_flag_shared(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, + uint64_t timeout_ns, uint32_t *seen, uint32_t *why, uint32_t *waited_us, + volatile uint32_t *done) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *done = 1u; + *seen = v; + return true; + } + if (*done != 0u) { + *seen = seq; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + if (*seen == seq) { + *done = 1u; + return true; + } + return *done != 0u; + } + } + } +} + +// One failure record a rank (the first failing wait's): its fields, then its valid word, each fenced at system scope. +__device__ __forceinline__ void fail_record(uint32_t *poison, uint32_t *rec, uint32_t seq, uint32_t seen, uint32_t who, + uint32_t why, uint32_t waited) { + if (atom_cas_gpu(poison, 0u, 1u) == 0u) { + st_relaxed_sys(rec + 1, seq); + st_relaxed_sys(rec + 2, seen); + st_relaxed_sys(rec + 3, who); + st_relaxed_sys(rec + 4, why); + st_relaxed_sys(rec + 5, waited); + fence_sc_sys(); + st_relaxed_sys(rec, 1u); + fence_sc_sys(); + } +} + +} // namespace tftp + +using namespace tftp; + +extern "C" __global__ void __launch_bounds__(512) tp_mailbox(MailArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + // a poisoned transport does nothing (the host raises on its next check); poison and epoch in one round trip + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const uint64_t flag_at = (slot * kMaxBlocks + blockIdx.x) * kFlagStride; + + // 1. this block's share of the shard into the peer's inbox, then its flag after it (system-scope release) + copy_bytes(a.peer_slots + slot * a.slot_bytes, a.in, a.nbytes, index, stride); + __syncthreads(); + if (tid == 0) { + fence_sc_sys(); + st_release_sys(a.peer_flags + flag_at, seq); + } + // the local shard while the peer's writes travel + if (a.mode == kGather) { + copy_bytes(a.out + static_cast(a.rank) * a.nbytes, a.in, a.nbytes, index, stride); + } + + // 2. the peer's block of the same index (it wrote exactly the bytes this block reads) + if (tid == 0) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.my_flags + flag_at, seq, a.ctrl, a.timeout_ns, &seen, &why, &waited)) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kCtrlRank + 8u * a.rank, seq, seen, blockIdx.x, why, waited); + } + } + __syncthreads(); + + // 3. copy out, gated by this block's own wait + if (s_bad == 0u) { + const uint8_t *inbox = a.my_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_bytes(a.out + static_cast(1u - a.rank) * a.nbytes, inbox, a.nbytes, index, + stride); + break; + case kExchange: + copy_bytes(a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the last block publishes the epoch, unless a wait of this launch failed (later launches stay no-ops) + __syncthreads(); + if (tid == 0) { + // the counter is back at 0 for the next launch (stream order), whatever its grid + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + } + } +} + +// The RoCE variant (roce.zig): blocks stage into the send slot, the last rings the doorbell, a host proxy RDMA-writes. +namespace tftp { + +constexpr int kRoceHcas = 2; +// ctrl words shared with roce.zig's proxy: [0] doorbell seq, [4 + slot] padded bytes, [8] completed, [16..21] failure +constexpr int kRoceSeq = 0; +constexpr int kRoceSlotBytes = 4; +constexpr int kRoceCompleted = 8; +constexpr int kRoceFailed = 16; +constexpr int kRoceAbort = 12; // wait_flag reads ctrl[kCtrlAbort]: the RoCE abort word is passed as that ctrl + +struct RoceArgs { + const uint8_t *in; + uint8_t *out; + uint8_t *send_slots; // this rank's send slots (pinned host, read by the NIC) + const uint8_t *recv_slots; // this rank's receive slots (written by the peer's NIC) + const uint32_t *flags; // this rank's flags [2][kRoceHcas], written by the peer's NIC + uint32_t *ctrl; + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison, [3] stage counter + uint64_t nbytes; + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; + uint32_t n_hca; + uint32_t opts; // TF_TP_ROCE_FAST: bit 0 fast, bits 8-15 pollers a flag, bits 16-31 start stagger ns; 0: classic + // gather only (nullptr: every shard is nbytes): lens[r] = rank r's bytes, padded to 16 on the wire + const int32_t *lens; + uint64_t chunk; // chunked gathers: this op's offset into each shard (at most a slot of each from there; else 0) +}; + +// Rank r's bytes this op: ``nbytes``, or lens[r] capped at nbytes. +__device__ __forceinline__ uint64_t roce_shard(const RoceArgs &a, uint32_t r) { + if (a.lens == nullptr) { + return a.nbytes; + } + const int32_t v = __ldg(a.lens + r); + const uint64_t n = v > 0 ? static_cast(v) : 0u; + return n < a.nbytes ? n : a.nbytes; +} + +// The bytes a copy of an n-byte shard moves: n rounded up to 16 when that stays inside a 16-aligned stride, else n. +__device__ __forceinline__ uint64_t roce_span(const RoceArgs &a, uint64_t n) { + if (a.lens == nullptr || (a.nbytes & 15u) != 0u) { + return n; + } + return (n + 15u) & ~uint64_t(15u); +} + +// A shard's bytes past `off`, at most a slot (chunked gathers; off 0 and a shard within the slot: the shard). +__device__ __forceinline__ uint64_t roce_part(const RoceArgs &a, uint64_t shard, uint64_t off) { + const uint64_t rest = shard > off ? shard - off : 0u; + return rest < a.slot_bytes ? rest : a.slot_bytes; +} + +} // namespace tftp + +// Present in images whose RoceArgs carry ``chunk`` (roce.zig refuses TF_TP_ROCE_CHUNKED_KB without it); never launched. +extern "C" __global__ void tp_roce_chunk_abi() {} + +extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + __shared__ uint32_t s_done[kRoceHcas]; // fast: some poller of HCA h saw its flag + const bool fast = (a.opts & 1u) != 0u; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + if (tid < static_cast(kRoceHcas)) { + s_done[tid] = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const bool var = a.mode == kGather && a.lens != nullptr; + const uint64_t off = var ? a.chunk : 0u; + const uint64_t mine = var ? roce_part(a, roce_shard(a, a.rank), off) : a.nbytes; + const uint64_t theirs = var ? roce_part(a, roce_shard(a, 1u - a.rank), off) : a.nbytes; + const uint64_t padded = var ? (mine < 16u ? 16u : (mine + 15u) & ~uint64_t(15u)) : (a.nbytes + 15u) & ~uint64_t(15u); + if (off != 0u && mine == 0u && theirs == 0u) { + return; // nothing past this offset on either rank: no op, the epoch stays (both ranks skip alike) + } + if (padded > a.slot_bytes) { // never staged past a slot: a refused wait (why 3), the transport poisoned + if (tid == 0 && blockIdx.x == 0) { + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, 0u, 0u, 3u, 0u); + } + return; + } + + // 1. stage this block's share into the send slot; the last block to finish rings the doorbell + copy_any(fast, a.send_slots + slot * a.slot_bytes, a.in + off, var ? roce_span(a, mine) : a.nbytes, index, + stride); + __syncthreads(); + if (tid == 0 && fast) { + // lean doorbell: release RMW per block arrival, the last acquires, then a release store (no SC system fences) + const bool last = gdim == 1u || atom_add_acq_rel_sys(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } else if (tid == 0) { + fence_sc_sys(); + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + fence_sc_sys(); // every block's staging (seen through the counter) before the doorbell + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } + if (a.mode == kGather) { + copy_any(fast, a.out + static_cast(a.rank) * a.nbytes + off, a.in + off, + var ? roce_span(a, mine) : a.nbytes, index, stride); + } + + // 2. every HCA's flag for this slot; fast: poller k of HCA h is thread 32 k + h, started k x stagger ns late + const uint32_t pollers = (a.opts >> 8) & 0xffu; + const uint32_t stagger_ns = a.opts >> 16; + if (fast) { + const uint32_t lane = tid & 31u; + const uint32_t k = tid >> 5; + if (lane < a.n_hca && k < (pollers == 0u ? 1u : pollers)) { + if (k > 0u && stagger_ns != 0u) { + __nanosleep(k * stagger_ns); + } + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag_shared(a.flags + (slot * kRoceHcas + lane) * kFlagStride, seq, a.ctrl + kRoceAbort, + a.timeout_ns, &seen, &why, &waited, &s_done[lane])) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, lane, why, waited); + } + } + } else if (tid == 0) { + for (uint32_t h = 0; h < a.n_hca; h++) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.flags + (slot * kRoceHcas + h) * kFlagStride, seq, a.ctrl + kRoceAbort, a.timeout_ns, &seen, + &why, &waited)) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, h, why, waited); + break; + } + } + } + __syncthreads(); + + // 3. copy out of the receive slot the peer's NIC wrote + if (s_bad == 0u) { + const uint8_t *inbox = a.recv_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_any(fast, a.out + static_cast(1u - a.rank) * a.nbytes + off, inbox, + var ? roce_span(a, theirs) : a.nbytes, index, stride); + break; + case kExchange: + copy_any(fast, a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the epoch, as tp_mailbox; the host also sees the last completed seq + __syncthreads(); + if (tid == 0) { + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + st_relaxed_sys(a.ctrl + kRoceCompleted, seq); + } + } +} diff --git a/zig/src/fabric/tp/mailbox.zig b/zig/src/fabric/tp/mailbox.zig new file mode 100644 index 0000000000..902bb79393 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.zig @@ -0,0 +1,215 @@ +//! The mailbox backend: two-rank exchanges in one kernel through pinned host memory both GPUs map (host side). +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const flag_bytes = 2 * max_blocks * 128; +const ctrl_bytes = 4096; +const ctrl_abort = 0; +const ctrl_rank = 16; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `MailArgs`, field for field. +const MailArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + peer_slots: DevicePtr, + peer_flags: DevicePtr, + my_slots: DevicePtr, + my_flags: DevicePtr, + state: DevicePtr, + ctrl: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, +}; + +/// Driver calls the runtime's table does not carry (page-locking memory this process mapped itself). +pub const HostRegister = struct { + lib: std.DynLib, + register: *const fn (?*anyopaque, usize, c_uint) callconv(.c) cuda.abi.Result, + unregister: *const fn (?*anyopaque) callconv(.c) cuda.abi.Result, + + pub fn open() !HostRegister { + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + return .{ + .register = lib.lookup(@FieldType(HostRegister, "register"), "cuMemHostRegister_v2") orelse return error.MissingSymbol, + .unregister = lib.lookup(@FieldType(HostRegister, "unregister"), "cuMemHostUnregister") orelse return error.MissingSymbol, + .lib = lib, + }; + } +}; + +/// Grid for `n` bytes a rank: one block per 16 KiB, a power of two, at most `max_blocks` (both ranks pick the same). +pub fn blocksFor(n: usize) u32 { + const want = std.math.divCeil(usize, @max(n, 1), 16 * 1024) catch 1; + return @intCast(@min(std.math.ceilPowerOfTwoAssert(usize, want), max_blocks)); +} + +pub const Mailbox = struct { + d: *const cuda.Driver, + me: u32, + slot_bytes: usize, + timeout_ns: u64, + region: []align(std.heap.page_size_min) u8, + dev: DevicePtr, + hr: HostRegister, + state: cuda.DeviceBuffer, + module: cuda.Module, + func: cuda.Function, + aborted: std.atomic.Value(bool) = .init(false), + + /// Both ranks call it after the bootstrap (two ranks, one host); `image` is mailbox.cu's fatbin. + pub fn init(d: *const cuda.Driver, b: *Bootstrap, image: []const u8, max_bytes: usize, timeout_ns: u64) !Mailbox { + if (b.world != 2) return error.Unsupported; + if (image.len == 0) return error.NoKernel; + const slot = std.mem.alignForward(usize, max_bytes, 4096); + const len = ctrl_bytes + 2 * (flag_bytes + 2 * slot); + // rank 0 creates and sizes the file before anyone learns its name, the others open it, then rank 0 unlinks it + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + const path = name[0..name_len :0]; + if (b.rank != 0) { + fd = std.c.open(path, .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + try b.agree("mapped", "mailbox file"); + const region = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + errdefer posix.munmap(region); + try b.agree("open", "mailbox mapping"); + if (b.rank == 0) _ = std.c.unlink(path); + var hr = try HostRegister.open(); + errdefer hr.lib.close(); + try d.check(hr.register(region.ptr, len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + errdefer _ = hr.unregister(region.ptr); + var dev: DevicePtr = 0; + try d.check(d.api.cuMemHostGetDevicePointer_v2(&dev, region.ptr, 0), "cuMemHostGetDevicePointer"); + var state = try cuda.DeviceBuffer.alloc(d, 16); + errdefer state.free(); + try state.fill8(0, null); + var module = try cuda.Module.load(d, image); + errdefer module.unload(); + const func = try module.function("tp_mailbox"); + try b.agree("ready", "mailbox setup"); + return .{ .d = d, .me = b.rank, .slot_bytes = slot, .timeout_ns = timeout_ns, .region = region, .dev = dev, .hr = hr, .state = state, .module = module, .func = func }; + } + + pub fn deinit(m: *Mailbox) void { + m.module.unload(); + m.state.free(); + _ = m.hr.unregister(m.region.ptr); + m.hr.lib.close(); + posix.munmap(m.region); + } + + fn inbox(m: *const Mailbox, r: u32) usize { + return ctrl_bytes + r * (flag_bytes + 2 * m.slot_bytes); + } + + fn ctrlWord(m: *const Mailbox, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(m.region[4 * i ..].ptr)); + } + + /// Whether `n` bytes a rank fit one slot. + pub fn fits(m: *const Mailbox, n: usize) bool { + return n <= m.slot_bytes; + } + + fn launch(m: *Mailbox, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, stream: collective.Stream) Error!void { + if (m.aborted.load(.acquire)) return error.Aborted; + if (n > m.slot_bytes) return error.Unsupported; + if (n == 0) return; + const peer = 1 - m.me; + var a: MailArgs = .{ + .in = in, + .out = out, + .peer_slots = m.dev + m.inbox(peer) + flag_bytes, + .peer_flags = m.dev + m.inbox(peer), + .my_slots = m.dev + m.inbox(m.me) + flag_bytes, + .my_flags = m.dev + m.inbox(m.me), + .state = m.state.ptr, + .ctrl = m.dev, + .nbytes = n, + .slot_bytes = m.slot_bytes, + .timeout_ns = m.timeout_ns, + .rank = m.me, + .mode = @backingInt(mode), + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const r = m.d.api.cuLaunchKernel(m.func.handle, blocksFor(n), 1, 1, threads, 1, 1, 0, stream, ¶ms, null); + m.d.check(r, "cuLaunchKernel tp_mailbox") catch return error.BackendFailed; + } + + pub fn allGatherBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.gather, s, r, n, stream); + } + + pub fn exchangeBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.exchange, s, r, n, stream); + } + + /// Sums only, f32 or bf16 (what the model's partials are); the rest is `error.Unsupported`. + pub fn sum(m: *Mailbox, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return m.launch(mode, s, r, count * t.size(), stream); + } + + /// Ends pending and later exchanges on both ranks (the abort word is shared); this rank's launches stop. + pub fn abort(m: *Mailbox) void { + m.aborted.store(true, .release); + m.ctrlWord(ctrl_abort).store(1, .release); + } + + /// A wait that timed out or was aborted on either rank. + pub fn check(m: *Mailbox) Error!void { + for (0..2) |r| { + const rec = ctrl_rank + 8 * r; + if (m.ctrlWord(rec).load(.acquire) != 0) { + log.warn("tp mailbox: rank {d}'s wait for seq {d} failed ({s} after {d} us; block {d}, flag held {d})", .{ r, m.ctrlWord(rec + 1).load(.acquire), if (m.ctrlWord(rec + 4).load(.acquire) == 1) "aborted" else "timed out", m.ctrlWord(rec + 5).load(.acquire), m.ctrlWord(rec + 3).load(.acquire), m.ctrlWord(rec + 2).load(.acquire) }); + return error.PeerFailed; + } + } + if (m.aborted.load(.acquire) or m.ctrlWord(ctrl_abort).load(.acquire) != 0) return error.Aborted; + } +}; + +test "grids grow with the shard and stay powers of two" { + try std.testing.expectEqual(@as(u32, 1), blocksFor(4096)); + try std.testing.expectEqual(@as(u32, 1), blocksFor(16 * 1024)); + try std.testing.expectEqual(@as(u32, 2), blocksFor(16 * 1024 + 1)); + try std.testing.expectEqual(@as(u32, 4), blocksFor(64 * 1024)); + try std.testing.expectEqual(@as(u32, 16), blocksFor(4 << 20)); + try std.testing.expectEqual(@as(usize, 96), @sizeOf(MailArgs)); +} diff --git a/zig/src/fabric/tp/nccl.zig b/zig/src/fabric/tp/nccl.zig new file mode 100644 index 0000000000..c9369369ef --- /dev/null +++ b/zig/src/fabric/tp/nccl.zig @@ -0,0 +1,164 @@ +//! The NCCL backend: one communicator a rank, every collective one enqueue on the caller's stream. +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const nccl = cuda.nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub fn dataType(t: DType) nccl.DataType { + return switch (t) { + .u8 => .u8, + .i32 => .i32, + .i64 => .i64, + .f16 => .f16, + .bf16 => .bf16, + .f32 => .f32, + .f64 => .f64, + }; +} + +pub fn redOp(op: Op) nccl.RedOp { + return switch (op) { + .sum => .sum, + .max => .max, + .min => .min, + }; +} + +pub const Nccl = struct { + lib: *const nccl.Library, + d: *const cuda.Driver, + comm: nccl.Comm, + me: u32, + size: u32, + aborted: std.atomic.Value(bool) = .init(false), + /// The barrier's one-float buffers and its own stream (host path only). + scratch: cuda.DeviceBuffer, + stream: cuda.Stream, + + /// Joins the communicator `id` names (every rank calls it with the same id; it returns once all have). + pub fn init(lib: *const nccl.Library, d: *const cuda.Driver, id: nccl.UniqueId, rank: u32, world: u32) !Nccl { + var comm: nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(world), id, @intCast(rank)), "ncclCommInitRank"); + errdefer _ = lib.api.ncclCommDestroy(comm); + var scratch = try cuda.DeviceBuffer.alloc(d, 4 * (world + 1)); + errdefer scratch.free(); + try scratch.fill8(0, null); + return .{ .lib = lib, .d = d, .comm = comm, .me = rank, .size = world, .scratch = scratch, .stream = try cuda.Stream.init(d, true) }; + } + + pub fn deinit(n: *Nccl) void { + _ = n.stream.synchronize() catch {}; + n.stream.deinit(); + n.scratch.free(); + if (!n.aborted.load(.acquire)) _ = n.lib.api.ncclCommDestroy(n.comm); + } + + pub fn iface(n: *Nccl) collective.Collective { + return .{ .ptr = n, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .nccl, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + }; + + fn self_(ptr: *anyopaque) *Nccl { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).size; + } + + fn call(n: *Nccl, r: nccl.Result, what: []const u8) Error!void { + if (r == 0) return; + if (n.aborted.load(.acquire)) return error.Aborted; + n.lib.check(r, what) catch return error.BackendFailed; + } + + fn live(n: *Nccl) Error!void { + if (n.aborted.load(.acquire)) return error.Aborted; + } + + /// NCCL adds in its ring's order: the rank-order sum `all_reduce` promises holds at two ranks, not past them. + pub fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllReduce(s, r, count, dataType(t), redOp(op), n.comm, stream), "ncclAllReduce"); + } + + pub fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(s, r, count, dataType(t), n.comm, stream), "ncclAllGather"); + } + + pub fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclGroupStart(), "ncclGroupStart"); + const a = n.lib.api.ncclSend(s, count, dataType(t), @intCast(peer), n.comm, stream); + const b = n.lib.api.ncclRecv(r, count, dataType(t), @intCast(peer), n.comm, stream); + const e = n.lib.api.ncclGroupEnd(); + try n.call(a, "ncclSend"); + try n.call(b, "ncclRecv"); + try n.call(e, "ncclGroupEnd"); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclSend(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclSend"); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclRecv(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclRecv"); + } + + /// One float all-gathered on the barrier's own stream, then a sync of it (host path, never captured). + fn barrier(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(n.scratch.ptr, n.scratch.ptr + 4, 1, .f32, n.comm, n.stream.handle), "ncclAllGather"); + n.stream.synchronize() catch return if (n.aborted.load(.acquire)) error.Aborted else error.BackendFailed; + } + + fn check(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + var r: nccl.Result = 0; + try n.call(n.lib.api.ncclCommGetAsyncError(n.comm, &r), "ncclCommGetAsyncError"); + if (r != 0 and r != nccl.in_progress) { + n.lib.check(r, "NCCL async error") catch {}; + return error.PeerFailed; + } + } + + fn abort(ptr: *anyopaque) void { + const n = self_(ptr); + if (n.aborted.swap(true, .acq_rel)) return; + _ = n.lib.api.ncclCommAbort(n.comm); + } +}; diff --git a/zig/src/fabric/tp/planlink.zig b/zig/src/fabric/tp/planlink.zig new file mode 100644 index 0000000000..44dbb61a0c --- /dev/null +++ b/zig/src/fabric/tp/planlink.zig @@ -0,0 +1,467 @@ +//! The round plan from rank 0 to the other ranks over its own TCP connection: a 4-byte length, then int64 words. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const builtin = @import("builtin"); +const config = @import("config.zig"); +const linux = std.os.linux; + +pub const max_ranks = 64; + +/// The CPU a plan thread pinned itself to (TF_TP_PLAN_PIN), -1 until one did; an engine's CPU isolation keeps off it. +pub var pinned_cpu: std.atomic.Value(i32) = .init(-1); +const more: u32 = 0x8000; // MSG_MORE: the length and the ints leave as one segment + +pub const Error = error{ PeerClosed, TooLarge, NotRankZero, NotFollower }; + +/// A batch frame's first word (a plan's op is >= 0): [batch, n, len 1, plan 1, ..., len n, plan n]. +pub const batch: i64 = std.math.minInt(i64); +/// Words rank 0 holds at most (more flushes the held plans first), and the most a batch frame carries. +pub const hold_words: usize = 1024; +pub const batch_words: usize = 2 * hold_words + 2; + +/// Largest plan either side accepts: 64M words (512 MiB), so a garbled length cannot become a 4 GiB allocation. +pub const max_frame_words: usize = 64 << 20; + +comptime { + std.debug.assert(builtin.cpu.arch.endian() == .little); // the ints go out as they are in memory +} + +pub const PlanLink = struct { + rank: u32, + world: u32, + /// Rank 0: rank i's socket at index i; the others: rank 0's at index 0 (the bootstrap's `releasePlan`). + fds: [max_ranks]sock.Fd, + sent: u64 = 0, + received: u64 = 0, + /// TF_TP_PLAN_LINK=rdma (config.plan_spin_ns): the follower polls the socket this long before it blocks. + spin_ns: u64 = 0, + /// TF_TP_PLAN_PIN: the sending / receiving thread pins itself at its first plan; `avoid`: the RoCE proxy's CPU. + pin: ?config.PlanPin = null, + avoid: ?u32 = null, + pinned: bool = false, + /// rank 0: plans held for the next send (`hold`), each as [len, plan...]; `held_n` of them + held: [hold_words]i64 = undefined, + held_len: usize = 0, + held_n: usize = 0, + /// the follower: a batch frame's plans not handed out yet, as [len, plan...] from `inbox_at` + inbox: [batch_words]i64 = undefined, + inbox_len: usize = 0, + inbox_at: usize = 0, + batches: u64 = 0, + + pub fn init(rank: u32, world: u32, fds: [max_ranks]sock.Fd) PlanLink { + return .{ .rank = rank, .world = world, .fds = fds }; + } + + /// Rank 0: one plan to every other rank (host only); held plans go first, in the same frame when they fit. + pub fn send(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (ints.len > max_frame_words) return error.TooLarge; + if (p.held_n > 0) { + if (2 + p.held_len + 1 + ints.len <= batch_words) { + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n + 1); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + frame[2 + p.held_len] = @intCast(ints.len); + @memcpy(frame[3 + p.held_len ..][0..ints.len], ints); + const n = 3 + p.held_len + ints.len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + try p.flush(); + } + return p.sendFrame(ints); + } + + /// Rank 0: a plan the followers run before the next sent one, carried in that plan's frame (`send` or `flush`). + pub fn hold(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (1 + ints.len > hold_words) { + try p.flush(); + return p.sendFrame(ints); + } + if (p.held_len + 1 + ints.len > hold_words) try p.flush(); + p.held[p.held_len] = @intCast(ints.len); + @memcpy(p.held[p.held_len + 1 ..][0..ints.len], ints); + p.held_len += 1 + ints.len; + p.held_n += 1; + } + + /// Rank 0: the held plans now, as one batch frame (nothing held: nothing sent). + pub fn flush(p: *PlanLink) Error!void { + if (p.held_n == 0) return; + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + const n = 2 + p.held_len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + + fn sendFrame(p: *PlanLink, ints: []const i64) Error!void { + p.pinOnce(); + const bytes = std.mem.sliceAsBytes(ints); + const len: u32 = @intCast(bytes.len); + for (p.fds[1..p.world]) |fd| { + if (std.c.send(fd, &len, 4, std.posix.MSG.NOSIGNAL | more) != 4) return error.PeerClosed; + sock.sendAll(fd, bytes) catch return error.PeerClosed; + } + p.sent += 1; + } + + /// The other ranks: the next plan into `out`; `error.PeerClosed` once rank 0 is gone, `error.TooLarge` skips one. + pub fn recv(p: *PlanLink, out: []i64) Error![]i64 { + if (p.inboxNext()) |m| { + if (m.len > out.len) return error.TooLarge; + @memcpy(out[0..m.len], m); + return out[0..m.len]; + } + const n = try p.header(); + if (n > out.len) { + try p.drop(n); + return error.TooLarge; + } + const got = try p.body(out[0..n]); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recv(out); + } + return got; + } + + /// The next plan into `buf`, grown to the frame; the buffer is kept, so steady plans allocate nothing. + pub fn recvGrow(p: *PlanLink, gpa: std.mem.Allocator, buf: *std.ArrayList(i64)) (Error || std.mem.Allocator.Error)![]i64 { + if (p.inboxNext()) |m| { + try buf.resize(gpa, m.len); + @memcpy(buf.items, m); + return buf.items; + } + const n = try p.header(); + if (n > max_frame_words) { + try p.drop(n); + return error.TooLarge; + } + try buf.resize(gpa, n); + const got = try p.body(buf.items); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recvGrow(gpa, buf); + } + return got; + } + + /// The follower: a batch frame's plans into the inbox (checked whole first: a garbled batch ends the link). + fn unpack(p: *PlanLink, frame: []const i64) Error!void { + if (frame.len < 2 or frame.len > batch_words or frame[1] < 1) return error.PeerClosed; + var at: usize = 2; + var k: i64 = 0; + while (k < frame[1]) : (k += 1) { + if (at >= frame.len or frame[at] < 0) return error.PeerClosed; + at += 1 + @as(usize, @intCast(frame[at])); + if (at > frame.len) return error.PeerClosed; + } + if (at != frame.len) return error.PeerClosed; + @memcpy(p.inbox[0 .. frame.len - 2], frame[2..]); + p.inbox_len = frame.len - 2; + p.inbox_at = 0; + p.batches += 1; + } + + /// The follower: the next plan of the last batch frame, if any (a slice of the inbox, valid until the next recv). + fn inboxNext(p: *PlanLink) ?[]const i64 { + if (p.inbox_at >= p.inbox_len) return null; + const n: usize = @intCast(p.inbox[p.inbox_at]); + const m = p.inbox[p.inbox_at + 1 ..][0..n]; + p.inbox_at += 1 + n; + return m; + } + + /// The next frame's length in words. + fn header(p: *PlanLink) Error!usize { + if (p.rank == 0) return error.NotFollower; + p.pinOnce(); + if (p.spin_ns > 0) p.spin(); + var len: u32 = undefined; + sock.recvExact(p.fds[0], std.mem.asBytes(&len), null) catch return error.PeerClosed; + if (len % 8 != 0) return error.PeerClosed; // not a frame boundary: the stream is lost + return len / 8; + } + + fn body(p: *PlanLink, out: []i64) Error![]i64 { + sock.recvExact(p.fds[0], std.mem.sliceAsBytes(out), null) catch return error.PeerClosed; + p.received += 1; + return out; + } + + /// Reads and drops `n` words. + fn drop(p: *PlanLink, n: usize) Error!void { + var left: usize = n * 8; + var scratch: [4096]u8 = undefined; + while (left > 0) : (left -= @min(left, scratch.len)) sock.recvExact(p.fds[0], scratch[0..@min(left, scratch.len)], null) catch return error.PeerClosed; + } + + /// Polls rank 0's socket until it is readable (or hung up) or `spin_ns` passed; the read that follows blocks. + fn spin(p: *PlanLink) void { + if (comptime builtin.os.tag != .linux) return; + var fd = [1]linux.pollfd{.{ .fd = p.fds[0], .events = linux.POLL.IN, .revents = 0 }}; + const end = sock.nowNs() + p.spin_ns; + while (true) { + const rc = linux.poll(&fd, 1, 0); + if (linux.errno(rc) == .SUCCESS and rc > 0) return; + if (sock.nowNs() >= end) return; + std.atomic.spinLoopHint(); + } + } + + /// Pins the calling thread once: an explicit CPU, or auto's highest cpu_capacity. + fn pinOnce(p: *PlanLink) void { + if (p.pinned) return; + p.pinned = true; + if (comptime builtin.os.tag != .linux) return; + const want = p.pin orelse return; + const cpu = switch (want) { + .cpu => |c| c, + .auto => bestCpu(p.avoid) orelse { + log.info("[tensorfold] tp: TF_TP_PLAN_PIN=auto: no cpu_capacity differs among the allowed CPUs; the plan thread stays unpinned", .{}); + return; + }, + }; + pinTo(cpu) catch |e| { + log.warn("[tensorfold] tp: TF_TP_PLAN_PIN: CPU {d} refused ({t}); the plan thread stays unpinned", .{ cpu, e }); + return; + }; + pinned_cpu.store(@intCast(cpu), .release); + log.info("[tensorfold] tp: plan link rank {d} pinned to CPU {d}", .{ p.rank, cpu }); + } + + pub fn close(p: *PlanLink) void { + for (&p.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } +}; + +/// The calling thread onto `cpu` alone. +pub fn pinTo(cpu: u32) !void { + if (cpu >= linux.CPU_SETSIZE) return error.BadCpu; + var set: linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + try linux.sched_setaffinity(0, &set); +} + +/// Auto pin: the allowed CPU (not `avoid`) with the highest cpu_capacity, lowest on ties; null when all are equal. +pub fn bestCpu(avoid: ?u32) ?u32 { + var set: linux.cpu_set_t = undefined; + if (linux.errno(linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set)) != .SUCCESS) return null; + var best: ?u32 = null; + var best_cap: u32 = 0; + var low_cap: u32 = std.math.maxInt(u32); + for (set, 0..) |word, wi| { + var bits = word; + while (bits != 0) : (bits &= bits - 1) { + const cpu: u32 = @intCast(wi * @bitSizeOf(usize) + @ctz(bits)); + if (avoid == cpu) continue; + const c = capacity(cpu) orelse continue; + low_cap = @min(low_cap, c); + if (best == null or c > best_cap) { + best = cpu; + best_cap = c; + } + } + } + return if (best != null and best_cap > low_cap) best else null; +} + +fn capacity(cpu: u32) ?u32 { + var pb: [64]u8 = undefined; + const path = std.fmt.bufPrintSentinel(&pb, "/sys/devices/system/cpu/cpu{d}/cpu_capacity", .{cpu}, 0) catch return null; + const rc = linux.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) return null; + const fd: i32 = @intCast(rc); + defer _ = linux.close(fd); + var b: [16]u8 = undefined; + const n = linux.read(fd, &b, b.len); + if (linux.errno(n) != .SUCCESS) return null; + return std.fmt.parseInt(u32, std.mem.trim(u8, b[0..n], " \n"), 10) catch null; +} + +test "plans arrive whole and in order; a closed rank 0 ends the follower" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @as(i64, @intCast(i)) * -3 + 0x7FFF_0000_0000; + try lead.send(&.{ 0x504C414E, 2, 1 }); + try lead.send(&big); + try lead.send(&.{}); + var buf: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 2, 1 }, try follow.recv(&buf)); + try std.testing.expectEqualSlices(i64, &big, try follow.recv(&buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&buf)).len); + try std.testing.expectError(error.NotRankZero, follow.send(&.{1})); + var small: [2]i64 = undefined; + try lead.send(&.{ 1, 2, 3 }); + try std.testing.expectError(error.TooLarge, follow.recv(&small)); + lead.close(); + try std.testing.expectError(error.PeerClosed, follow.recv(&buf)); +} + +test "an admission-sized plan (1.2M words) arrives whole through the growable path" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + const gpa = std.testing.allocator; + const prompt = try gpa.alloc(i64, 1_200_000); + defer gpa.free(prompt); + for (prompt, 0..) |*v, i| v.* = @intCast((i *% 2654435761) % 163840); + // the socket buffer holds far less than 9.6 MB: rank 0 sends from its own thread, as it would + const Send = struct { + fn go(link: *PlanLink, ints: []const i64) void { + link.send(&.{ 0x504C414E, 3, 20 }) catch {}; + link.send(ints) catch {}; + link.send(&.{ 0x504C414E, 3, 21 }) catch {}; + } + }; + const t = try std.Thread.spawn(.{}, Send.go, .{ &lead, prompt }); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(gpa); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 20 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, prompt, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 21 }, try follow.recvGrow(gpa, &buf)); + t.join(); + try std.testing.expect(buf.capacity >= prompt.len); // kept for the next admission +} + +test "TF_TP_PLAN_LINK=rdma: a spinning follower gets every plan, late or early; TF_TP_PLAN_PIN pins the thread" { + if (builtin.os.tag != .linux) return error.SkipZigTest; + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + follow.spin_ns = 2 * std.time.ns_per_ms; + // the CPU this thread may run on first: the pin keeps it allowed, so the test runner is not starved + var set: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set); + var first: u32 = 0; + for (set, 0..) |w, i| if (w != 0) { + first = @intCast(i * @bitSizeOf(usize) + @ctz(w)); + break; + }; + const Follower = struct { + fn run(f: *PlanLink, cpu: u32, out: *[3]i64, ok: *bool) void { + f.pin = .{ .cpu = cpu }; + var buf: [8]i64 = undefined; + const x = (f.recv(&buf) catch return)[0]; // arrives within the spin + const y = (f.recv(&buf) catch return)[0]; // arrives after it: the blocking read + var now: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &now); + var count: usize = 0; + for (now) |w| count += @popCount(w); + out.* = .{ x, y, @intCast(count) }; + ok.* = f.pinned and (now[cpu / @bitSizeOf(usize)] >> @intCast(cpu % @bitSizeOf(usize))) & 1 == 1; + } + }; + var out: [3]i64 = .{ 0, 0, 0 }; + var ok = false; + const t = try std.Thread.spawn(.{}, Follower.run, .{ &follow, first, &out, &ok }); + try lead.send(&.{11}); + var ts: linux.timespec = .{ .sec = 0, .nsec = 20 * std.time.ns_per_ms }; + _ = linux.nanosleep(&ts, null); + try lead.send(&.{22}); + t.join(); + try std.testing.expectEqual([3]i64{ 11, 22, 1 }, out); + try std.testing.expect(ok); + if (bestCpu(null)) |c| try std.testing.expect((set[c / @bitSizeOf(usize)] >> @intCast(c % @bitSizeOf(usize))) & 1 == 1); +} + +test "held plans leave with the next plan, in order, as one frame; a large plan after them flushes them first" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("127.0.0.1", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + const lead = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(lead); + lead.* = PlanLink.init(0, 2, f0); + const follow = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(follow); + follow.* = PlanLink.init(1, 2, f1); + defer follow.close(); + // a round: three keeps held, the window sent: one frame + try lead.hold(&.{ 22, 0, 3 }); + try lead.hold(&.{ 22, 1, 0 }); + try lead.hold(&.{ 22, 2, 5 }); + try lead.send(&.{ 21, 4, 9, 9, 9 }); + try std.testing.expectEqual(@as(u64, 1), lead.sent); + // a held plan, then one too large for a batch: two frames, the held one first + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @intCast(i); + try lead.hold(&.{ 22, 3, 1 }); + try lead.send(&big); + try std.testing.expectEqual(@as(u64, 3), lead.sent); + // flush alone, an empty plan held, a plain send after + try lead.hold(&.{}); + try lead.flush(); + try lead.flush(); // nothing held: nothing sent + try lead.send(&.{7}); + try std.testing.expectEqual(@as(u64, 5), lead.sent); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(std.testing.allocator); + const gpa = std.testing.allocator; + try std.testing.expectEqualSlices(i64, &.{ 22, 0, 3 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 1, 0 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 2, 5 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 21, 4, 9, 9, 9 }, try follow.recvGrow(gpa, &buf)); + var out: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 22, 3, 1 }, try follow.recv(&out)); + try std.testing.expectEqualSlices(i64, &big, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&out)).len); + try std.testing.expectEqualSlices(i64, &.{7}, try follow.recv(&out)); + try std.testing.expectEqual(@as(u64, 3), follow.batches); + try std.testing.expectError(error.NotRankZero, follow.hold(&.{1})); + lead.close(); +} diff --git a/zig/src/fabric/tp/reference.zig b/zig/src/fabric/tp/reference.zig new file mode 100644 index 0000000000..b98966fcc8 --- /dev/null +++ b/zig/src/fabric/tp/reference.zig @@ -0,0 +1,121 @@ +//! Host references for the collectives: reductions in rank order, bf16 / f16 added in fp32 and rounded once. +const std = @import("std"); +const collective = @import("collective.zig"); +const DType = collective.DType; +const Op = collective.Op; + +/// bf16 bits to fp32 (exact). +pub fn bf16ToF32(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// fp32 to bf16 bits, round to nearest even; NaN stays a quiet NaN. +pub fn f32ToBf16(f: f32) u16 { + const u: u32 = @bitCast(f); + if (std.math.isNan(f)) return @intCast((u >> 16) | 0x40); + const round = 0x7FFF + ((u >> 16) & 1); + return @intCast((u +% round) >> 16); +} + +fn combine(comptime T: type, a: T, b: T, op: Op) T { + return switch (op) { + .sum => if (@typeInfo(T) == .int) a +% b else a + b, + .max => @max(a, b), + .min => @min(a, b), + }; +} + +fn reduceTyped(comptime T: type, acc: []u8, src: []const u8, op: Op) void { + const n = acc.len / @sizeOf(T); + for (0..n) |i| { + const a = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const b = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), src[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const r: T = combine(T, @bitCast(a), @bitCast(b), op); + std.mem.writeInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], @bitCast(r), .little); + } +} + +fn reduceHalf(acc: []u8, src: []const u8, op: Op, bf: bool) void { + const n = acc.len / 2; + for (0..n) |i| { + const a = std.mem.readInt(u16, acc[2 * i ..][0..2], .little); + const b = std.mem.readInt(u16, src[2 * i ..][0..2], .little); + const fa: f32 = if (bf) bf16ToF32(a) else @floatCast(@as(f16, @bitCast(a))); + const fb: f32 = if (bf) bf16ToF32(b) else @floatCast(@as(f16, @bitCast(b))); + const r = combine(f32, fa, fb, op); + const out: u16 = if (bf) f32ToBf16(r) else @bitCast(@as(f16, @floatCast(r))); + std.mem.writeInt(u16, acc[2 * i ..][0..2], out, .little); + } +} + +/// acc <- acc (op) src, element-wise; `acc` holds the lower ranks' result, `src` the next rank's input. +pub fn accumulate(acc: []u8, src: []const u8, t: DType, op: Op) void { + std.debug.assert(acc.len == src.len and acc.len % t.size() == 0); + switch (t) { + .u8 => reduceTyped(u8, acc, src, op), + .i32 => reduceTyped(i32, acc, src, op), + .i64 => reduceTyped(i64, acc, src, op), + .f32 => reduceTyped(f32, acc, src, op), + .f64 => reduceTyped(f64, acc, src, op), + .bf16 => reduceHalf(acc, src, op, true), + .f16 => reduceHalf(acc, src, op, false), + } +} + +/// out <- inputs[0] (op) inputs[1] (op) ..., in rank order. +pub fn reduce(out: []u8, inputs: []const []const u8, t: DType, op: Op) void { + @memcpy(out, inputs[0]); + for (inputs[1..]) |in| accumulate(out, in, t, op); +} + +/// Deterministic test data for `rank`'s input of `bytes`: finite values of type `t` that differ per rank and position. +pub fn fill(out: []u8, t: DType, rank: u32, seed: u64) void { + var rng = std.Random.DefaultPrng.init(seed ^ (@as(u64, rank) *% 0x9E3779B97F4A7C15)); + const r = rng.random(); + const n = out.len / t.size(); + for (0..n) |i| { + const x = r.float(f32) * 8.0 - 4.0; + switch (t) { + .u8 => out[i] = r.int(u8), + .i32 => std.mem.writeInt(i32, out[4 * i ..][0..4], r.int(i32), .little), + .i64 => std.mem.writeInt(i64, out[8 * i ..][0..8], r.int(i64), .little), + .f32 => std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(x), .little), + .f64 => std.mem.writeInt(u64, out[8 * i ..][0..8], @bitCast(@as(f64, x) * 1.000000123), .little), + .bf16 => std.mem.writeInt(u16, out[2 * i ..][0..2], f32ToBf16(x), .little), + .f16 => std.mem.writeInt(u16, out[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(x))), .little), + } + } +} + +test "bf16 rounding is nearest even" { + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.0)); + // 1 + 2^-8 is exactly halfway between bf16 1.0 and 1.0078125: rounds to the even one (1.0) + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.00390625)); + try std.testing.expectEqual(@as(u16, 0x3F82), f32ToBf16(1.01171875)); + try std.testing.expect(std.math.isNan(bf16ToF32(f32ToBf16(std.math.nan(f32))))); +} + +test "rank-order reduction of bf16 adds in fp32 and rounds once" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(u16, a[0..2], f32ToBf16(1.0), .little); + std.mem.writeInt(u16, b[0..2], f32ToBf16(0.00390625), .little); + std.mem.writeInt(u16, a[2..4], f32ToBf16(-2.5), .little); + std.mem.writeInt(u16, b[2..4], f32ToBf16(3.0), .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .bf16, .sum); + try std.testing.expectEqual(f32ToBf16(1.0), std.mem.readInt(u16, out[0..2], .little)); + try std.testing.expectEqual(f32ToBf16(0.5), std.mem.readInt(u16, out[2..4], .little)); + reduce(&out, &.{ &a, &b }, .bf16, .max); + try std.testing.expectEqual(f32ToBf16(3.0), std.mem.readInt(u16, out[2..4], .little)); +} + +test "integer sums wrap" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(i32, &a, std.math.maxInt(i32), .little); + std.mem.writeInt(i32, &b, 1, .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .i32, .sum); + try std.testing.expectEqual(@as(i32, std.math.minInt(i32)), std.mem.readInt(i32, &out, .little)); +} diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig new file mode 100644 index 0000000000..c9bb72a2c3 --- /dev/null +++ b/zig/src/fabric/tp/roce.zig @@ -0,0 +1,525 @@ +//! The RoCE backend between two hosts: the tp_roce kernel stages and rings a doorbell, a proxy thread moves the slot. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const rv = @import("roce_verbs.zig"); +const shm = @import("roce_shm.zig"); +const guard = @import("roce_guard.zig"); +const HostRegister = @import("mailbox.zig").HostRegister; +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const page = std.heap.page_size_min; +const flag_stride = 128; +// region: ctrl [4 KiB], flags [2 slots][2 HCAs] at 128 B, send [2][slot], recv [2][slot] (mailbox.cu's tp_roce) +const ctrl_off = 0; +const flag_off = 4096; +const send_off = 8192; +// ctrl words (mailbox.cu): doorbell, padded bytes per slot, completed, abort, proxy failure, a failed wait's record +const w_seq = 0; +const w_slot_bytes = 4; +const w_completed = 8; +const w_abort = 12; +const w_proxy_failed = 13; +const w_failed = 16; +/// Model graphs leave sub-millisecond gaps: the proxy stays hot this many idle polls, then naps 20 us at a time. +const idle_spins: u64 = 20_000_000; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `RoceArgs`, field for field. +const RoceArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + send_slots: DevicePtr, + recv_slots: DevicePtr, + flags: DevicePtr, + ctrl: DevicePtr, + state: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, + n_hca: u32, + /// TF_TP_ROCE_FAST (config.roce_fast; 0: the classic path) + opts: u32 = 0, + /// gather only: device int32 [2] of each rank's bytes (0: every shard is `nbytes`) + lens: DevicePtr = 0, + /// Chunked gathers: this op's byte offset into each rank's shard (it moves at most a slot from there; 0 otherwise). + chunk: u64 = 0, +}; + +pub const WireKind = enum { + /// RDMA writes over the HCAs. + verbs, + /// The proxy copies into the peer's region in shared memory: the same protocol on one host without RDMA (tests). + shm, +}; + +pub const Options = struct { + wire: WireKind = .verbs, + hcas: []const rv.HcaSpec = &.{}, + traffic_class: u8 = 0, + cpu: ?u32 = null, + max_bytes: usize = 256 * 1024, + timeout_ns: u64 = 120 * std.time.ns_per_s, + /// TF_TP_ROCE_FAST: tp_roce's `opts` word (config.roce_fast) + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid (config.roce_blocks) + blocks: u32 = 16, + /// The shm wire's stripes, as HCAs: its posts take the verbs wire's split and bounds. + shm_stripes: u32 = 1, + /// Tests: sees every shm post's stripes before the copy. + tap: ?*shm.PostTap = null, +}; + +const Wire = union(WireKind) { + verbs: rv.VerbsWire, + shm: shm.ShmWire, + + fn post(w: *Wire, slot: u32, seq: u32, padded: u32) !void { + switch (w.*) { + .verbs => |*v| try v.post(slot, seq, padded), + .shm => |*s| try s.post(slot, seq, padded), + } + } + + fn reap(w: *Wire) !void { + switch (w.*) { + .verbs => |*v| try v.reap(), + .shm => {}, + } + } + + fn hcas(w: *const Wire) u32 { + return switch (w.*) { + .verbs => |v| v.n_hca, + .shm => |s| s.stripes, + }; + } +}; + +pub const Roce = struct { + me: u32, + slot_bytes: usize, + timeout_ns: u64, + /// TF_TP_ROCE_FAST: the kernel's `opts` word + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid + max_grid: u32 = 16, + /// This rank's region (for shm: half of the mapping in `mapping`). + region: []align(page) u8, + mapping: []align(page) u8, + wire: Wire, + cpu: ?u32, + // device side (absent in host-only use: tests and `tf-tp-test roce-host`) + d: ?*const cuda.Driver = null, + hr: ?HostRegister = null, + dev: DevicePtr = 0, + state: ?cuda.DeviceBuffer = null, + module: ?cuda.Module = null, + func: ?cuda.Function = null, + /// The kernel image takes `RoceArgs.chunk` (it exports tp_roce_chunk_abi): TF_TP_ROCE_CHUNKED_KB needs it. + chunk_abi: bool = false, + // the proxy thread + thread: ?std.Thread = null, + running: std.atomic.Value(bool) = .init(false), + last_seq: u32 = 0, + err_buf: [160]u8 = undefined, + err_len: usize = 0, + // host emulation of the kernel's side (tests) + host_epoch: u32 = 0, + aborted: std.atomic.Value(bool) = .init(false), + + pub fn regionBytes(slot: usize) usize { + return send_off + 4 * slot; + } + + /// Both ranks call it after the bootstrap; any rank failing makes all return `error.RoceUnavailable` (collective). + pub fn init(d: ?*const cuda.Driver, b: *Bootstrap, image: []const u8, o: Options) !*Roce { + if (b.world != 2) return error.Unsupported; + if (o.max_bytes > guard.max_write) return error.WriteTooLarge; // a slot no post could fill in one write + const slot = std.mem.alignForward(usize, o.max_bytes, 4096); + const len = regionBytes(slot); + const r = try std.heap.page_allocator.create(Roce); + errdefer std.heap.page_allocator.destroy(r); + r.* = .{ .me = b.rank, .slot_bytes = slot, .timeout_ns = o.timeout_ns, .fast = o.fast, .max_grid = o.blocks, .region = undefined, .mapping = undefined, .wire = undefined, .cpu = o.cpu }; + // 1. the regions and the wire, each rank alone; every rank learns whether all did + const ok = r.openWire(b, o, len) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, ok, "RoCE setup")) { + if (ok) r.closeWire(); + return error.RoceUnavailable; + } + errdefer r.closeWire(); + // 2. the verbs connection records, through the bootstrap + const connected = r.connectWire(b) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} connect failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, connected, "RoCE connect")) return error.RoceUnavailable; + // 3. the device side + if (d) |drv| { + const dev_ok = r.openDevice(drv, image) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} device setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, dev_ok, "RoCE device setup")) { + if (dev_ok) r.closeDevice(); + return error.RoceUnavailable; + } + } + r.running.store(true, .release); + r.thread = try std.Thread.spawn(.{}, proxyLoop, .{r}); + try b.agree("proxy", "RoCE proxy start"); + return r; + } + + fn agreeAll(b: *Bootstrap, ok: bool, _: []const u8) !bool { + var blobs: [2][]const u8 = undefined; + var buf: [2]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(ok)}, &buf, &blobs); + return all[0][0] == 1 and all[1][0] == 1; + } + + fn openWire(r: *Roce, b: *Bootstrap, o: Options, len: usize) !bool { + switch (o.wire) { + .verbs => { + r.mapping = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + r.region = r.mapping; + errdefer posix.munmap(r.mapping); + try rv.prepareRegion(r.mapping); + r.wire = .{ .verbs = try rv.VerbsWire.open(o.hcas, r.region, .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes }, o.traffic_class) }; + }, + .shm => { + r.mapping = try shm.sharedPair(b, len); + r.region = @alignCast(r.mapping[r.me * len ..][0..len]); + const layout: guard.Layout = .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes, .len = len }; + r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .layout = layout, .local = r.region, .stripes = o.shm_stripes, .tap = o.tap } }; + }, + } + return true; + } + + fn connectWire(r: *Roce, b: *Bootstrap) !bool { + switch (r.wire) { + .shm => return true, + .verbs => |*v| { + const mine = v.record(); + var blobs: [2][]const u8 = undefined; + var buf: [2 * @sizeOf(rv.Record)]u8 = undefined; + const all = try b.allGather(std.mem.asBytes(&mine), &buf, &blobs); + var peer: rv.Record = undefined; + if (all[1 - r.me].len != @sizeOf(rv.Record)) return error.BadRecord; + @memcpy(std.mem.asBytes(&peer), all[1 - r.me]); + try v.connect(peer); + var line: [512]u8 = undefined; + log.info("[tensorfold] tp roce: rank {d} registered {s}", .{ r.me, v.describe(&line) }); + return true; + }, + } + } + + fn openDevice(r: *Roce, d: *const cuda.Driver, image: []const u8) !bool { + if (image.len == 0) return error.NoKernel; + r.d = d; + r.hr = try HostRegister.open(); + try d.check(r.hr.?.register(r.region.ptr, r.region.len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + try d.check(d.api.cuMemHostGetDevicePointer_v2(&r.dev, r.region.ptr, 0), "cuMemHostGetDevicePointer"); + r.state = try cuda.DeviceBuffer.alloc(d, 16); + try r.state.?.fill8(0, null); + r.module = try cuda.Module.load(d, image); + r.func = try r.module.?.function("tp_roce"); + r.chunk_abi = if (r.module.?.function("tp_roce_chunk_abi")) |_| true else |_| false; + return true; + } + + fn closeDevice(r: *Roce) void { + if (r.module) |*m| m.unload(); + if (r.state) |*s| s.free(); + if (r.hr) |*h| { + _ = h.unregister(r.region.ptr); + h.lib.close(); + } + r.module = null; + r.state = null; + r.hr = null; + } + + fn closeWire(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| v.close(), + .shm => {}, + } + posix.munmap(r.mapping); + } + + /// Regions a failed transport left registered and mapped (the process exit is their only teardown). + pub var kept: std.atomic.Value(u32) = .init(0); + + /// HCAs and GPU quiet, then unregister and unmap; a failed transport's region is kept until exit. + pub fn deinit(r: *Roce) void { + const was_failed = r.failed(); + r.word(w_abort).store(1, .release); + r.running.store(false, .release); + if (r.thread) |t| t.join(); + const quiet = r.quiesce(); + if (was_failed or !quiet) { + const why = if (was_failed) "failed" else "did not quiesce"; + log.warn("[tensorfold] tp roce: rank {d}: the transport {s}; region 0x{x} ({d} B) kept registered until exit", .{ r.me, why, @intFromPtr(r.region.ptr), r.region.len }); + _ = kept.fetchAdd(1, .monotonic); + return; + } + r.closeDevice(); + r.closeWire(); + std.heap.page_allocator.destroy(r); + } + + /// The QPs to error with every posted request reaped, and the GPU's work on the region done. + fn quiesce(r: *Roce) bool { + const wire = switch (r.wire) { + .verbs => |*v| v.quiesce(rv.quiesce_ns), + .shm => true, + }; + const gpu = if (r.d) |d| d.api.cuCtxSynchronize() == cuda.abi.success else true; + var line: [384]u8 = undefined; + log.info("[tensorfold] tp roce: rank {d} quiesced ({s}): {s}; GPU sync {s}", .{ r.me, if (wire) "drained" else "NOT drained", r.describeLink(&line), if (gpu) "done" else "FAILED" }); + return wire and gpu; + } + + /// Each HCA's QP state, posts and completions, and what a drain reaped (tf-tp-test's idle check; shm: none). + pub fn describeLink(r: *const Roce, buf: []u8) []const u8 { + return switch (r.wire) { + .verbs => |*v| v.report(buf), + .shm => "shm wire", + }; + } + + fn word(r: *const Roce, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(r.region[ctrl_off + 4 * i ..].ptr)); + } + + fn fail(r: *Roce, comptime fmt: []const u8, args: anytype) void { + const s = std.fmt.bufPrint(&r.err_buf, fmt, args) catch "proxy failure"; + r.err_len = s.len; + log.err("[tensorfold] tp roce proxy: {s}", .{s}); + r.word(w_proxy_failed).store(1, .release); + r.word(w_abort).store(1, .release); // waits on this rank end now instead of at the timeout + r.halt(); + } + + /// The HCAs stop: nothing of this rank's goes out and none of the peer's writes lands while the ranks go down. + fn halt(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| { + const ok = v.halt(); + var line: [384]u8 = undefined; + log.warn("[tensorfold] tp roce: rank {d} halted{s}: {s}", .{ r.me, if (ok) "" else " (a QP refused ERR)", v.report(&line) }); + }, + .shm => {}, + } + } + + /// Every post of this rank's completed (shm: nothing is ever in flight). + pub fn linkIdle(r: *const Roce) bool { + return switch (r.wire) { + .verbs => |*v| v.idle(), + .shm => true, + }; + } + + /// tf-tp-test's `failfast RANK transport`: the proxy's failure path (`fail`), here with no op in flight. + pub fn inject(r: *Roce) void { + r.fail("failure injected by tf-tp-test (no RDMA fault)", .{}); + } + + /// The transport failed or was aborted: a device may still write to its region. + pub fn failed(r: *const Roce) bool { + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return true; + return r.word(w_proxy_failed).load(.acquire) != 0 or r.word(w_failed).load(.acquire) != 0; + } + + /// The proxy loop: spin on the doorbell, post each new sequence (at most two pending), reap while idle, then nap. + fn proxyLoop(r: *Roce) void { + if (r.cpu) |cpu| pin(cpu); + var idle: u64 = 0; + while (r.running.load(.monotonic)) { + if (r.word(w_abort).load(.monotonic) != 0) return; + const seq = r.word(w_seq).load(.acquire); + if (seq == r.last_seq) { + idle += 1; + if (idle % 64 == 0) r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + if (idle >= idle_spins) sock.sleepNs(20 * std.time.ns_per_us) else std.atomic.spinLoopHint(); + continue; + } + idle = 0; + var pending = seq -% r.last_seq; + if (pending > 2) return r.fail("doorbell skipped {d} ops (posted {d}, now {d})", .{ pending, r.last_seq, seq }); + var s = r.last_seq +% 1; + while (pending > 0) : ({ + s +%= 1; + pending -= 1; + }) { + const padded = r.word(w_slot_bytes + (s & 1)).load(.acquire); + if (padded == 0 or padded % 16 != 0 or padded > r.slot_bytes) return r.fail("op {d} of {d} bytes", .{ s, padded }); + r.wire.post(s & 1, s, padded) catch |e| return r.fail("post of op {d}: {t}", .{ s, e }); + r.last_seq = s; + } + r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + } + } + + pub fn fits(r: *const Roce, n: usize) bool { + return n <= r.slot_bytes; + } + + fn launch(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + return r.launchAt(mode, in, out, n, lens, 0, stream); + } + + fn launchAt(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, chunk: u64, stream: collective.Stream) Error!void { + if (r.aborted.load(.acquire)) return error.Aborted; + if (n > r.slot_bytes and (lens == 0 or !r.chunk_abi)) return error.Unsupported; + if (n == 0) return; + const d = r.d orelse return error.Unsupported; + var a: RoceArgs = .{ + .in = in, + .out = out, + .send_slots = r.dev + send_off, + .recv_slots = r.dev + send_off + 2 * r.slot_bytes, + .flags = r.dev + flag_off, + .ctrl = r.dev + ctrl_off, + .state = r.state.?.ptr, + .nbytes = n, + .slot_bytes = r.slot_bytes, + .timeout_ns = r.timeout_ns, + .rank = r.me, + .mode = @backingInt(mode), + .n_hca = r.wire.hcas(), + .opts = r.fast, + .lens = lens, + .chunk = chunk, + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const blocks = @min(@import("mailbox.zig").blocksFor(@min(n, r.slot_bytes)), r.max_grid); + d.check(d.api.cuLaunchKernel(r.func.?.handle, blocks, 1, 1, threads, 1, 1, 0, stream, ¶ms, null), "cuLaunchKernel tp_roce") catch return error.BackendFailed; + } + + pub fn allGatherBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.gather, s, o, n, 0, stream); + } + + /// `allGatherBytes` where rank r posts only lens[r] bytes (device int32 [2], padded to 16, at least 16). + pub fn allGatherV(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (n > r.slot_bytes) return error.Unsupported; // one op: past a slot is allGatherVChunked's + if (lens == 0) return error.Invalid; + return r.launch(.gather, s, o, n, lens, stream); + } + + /// `allGatherV` past the slot: op c moves at most a slot of each shard from c x slot (empty on both ranks: no op). + pub fn allGatherVChunked(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (lens == 0) return error.Invalid; + if (!r.chunk_abi) return error.Unsupported; + var off: usize = 0; + while (off < n) : (off += r.slot_bytes) try r.launchAt(.gather, s, o, n, lens, off, stream); + } + + pub fn exchangeBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.exchange, s, o, n, 0, stream); + } + + pub fn sum(r: *Roce, s: DevicePtr, o: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return r.launch(mode, s, o, count * t.size(), 0, stream); + } + + /// Ends this rank's waits and the proxy (the peer learns through the fate channel and aborts its own). + pub fn abort(r: *Roce) void { + r.aborted.store(true, .release); + r.word(w_abort).store(1, .release); + r.halt(); + } + + pub fn check(r: *Roce) Error!void { + if (r.word(w_failed).load(.acquire) != 0) { + log.warn("tp roce: rank {d}'s wait for seq {d} failed ({s} after {d} us; HCA {d}, flag held {d}); doorbell {d}, posted {d}, completed {d}, abort {d}, proxy failed {d}, flags {d} {d}", .{ r.me, r.word(w_failed + 1).load(.acquire), if (r.word(w_failed + 4).load(.acquire) == 1) "aborted" else "timed out", r.word(w_failed + 5).load(.acquire), r.word(w_failed + 3).load(.acquire), r.word(w_failed + 2).load(.acquire), r.word(w_seq).load(.acquire), r.last_seq, r.word(w_completed).load(.acquire), r.word(w_abort).load(.acquire), r.word(w_proxy_failed).load(.acquire), r.flagWord(0), r.flagWord(1) }); + return error.PeerFailed; + } + if (r.word(w_proxy_failed).load(.acquire) != 0) return error.BackendFailed; + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return error.Aborted; + } + + fn flagWord(r: *const Roce, slot: usize) u32 { + const p: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + slot * 2 * flag_stride ..].ptr)); + return p.load(.acquire); + } + + pub fn proxyError(r: *const Roce) []const u8 { + return r.err_buf[0..r.err_len]; + } + + /// The kernel's side of one op done by the host (tests, `roce-host`): a gather through the same slots and flags. + pub fn hostGather(r: *Roce, in: []const u8, out: []u8) Error!void { + return r.hostGatherV(in, out, null); + } + + /// `hostGather` with variable lengths (`allGatherV`): rank k's first lens[k] bytes land at out + k x n. + pub fn hostGatherV(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32) Error!void { + if (in.len > r.slot_bytes) return error.Invalid; + return r.hostGatherAt(in, out, lens, 0); + } + + /// `allGatherVChunked`'s ops done by the host (the kernel's side of each, as `hostGatherV`). + pub fn hostGatherVChunked(r: *Roce, in: []const u8, out: []u8, lens: [2]u32) Error!void { + var off: usize = 0; + while (off < in.len) : (off += r.slot_bytes) try r.hostGatherAt(in, out, lens, off); + } + + /// One op from shard offset `off`: each rank's bytes past it, at most a slot; an op empty on both ranks is skipped. + fn hostGatherAt(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32, off: usize) Error!void { + const n = in.len; + if (n == 0 or out.len < 2 * n) return error.Invalid; + if (lens) |l| if (l[0] > n or l[1] > n) return error.Invalid; + const full: [2]usize = if (lens) |l| .{ l[0], l[1] } else .{ n, n }; + const mine: usize = @min(full[r.me] -| off, r.slot_bytes); + const theirs: usize = @min(full[1 - r.me] -| off, r.slot_bytes); + if (off > 0 and mine == 0 and theirs == 0) return; + const seq = r.host_epoch +% 1; + const slot: usize = seq & 1; + @memcpy(r.region[send_off + slot * r.slot_bytes ..][0..mine], in[off..][0..mine]); + r.word(w_slot_bytes + slot).store(@intCast(@max(16, std.mem.alignForward(usize, mine, 16))), .release); + r.word(w_seq).store(seq, .release); + @memcpy(out[r.me * n + off ..][0..mine], in[off..][0..mine]); + const t0 = sock.nowNs(); + for (0..r.wire.hcas()) |h| { + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + (slot * 2 + h) * flag_stride ..].ptr)); + while (flag.load(.acquire) != seq) { + if (r.word(w_abort).load(.monotonic) != 0) return error.Aborted; + if (sock.nowNs() - t0 > r.timeout_ns) return error.Timeout; + std.atomic.spinLoopHint(); + } + } + @memcpy(out[(1 - r.me) * n + off ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); + r.host_epoch = seq; + } +}; + +fn pin(cpu: u32) void { + var set: std.os.linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + std.os.linux.sched_setaffinity(0, &set) catch log.warn("tp roce: could not pin the proxy to CPU {d}", .{cpu}); +} diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig new file mode 100644 index 0000000000..c316448ef8 --- /dev/null +++ b/zig/src/fabric/tp/roce_fake.zig @@ -0,0 +1,143 @@ +//! Fake verbs for host tests of VerbsWire: each QP's posted writes counted from the requests, completions that lag. +const std = @import("std"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const rv = @import("roce_verbs.zig"); + +pub const slot: u64 = 1 << 20; +pub const send_off = 8192; + +/// What the HCAs hold (posted, not completed); every count comes from the work requests themselves. +pub const Fake = struct { + pub var region: []const u8 = &.{}; + /// The wire's slot (`Hw.wire`), every data write's bound with the send slots. + pub var slot_bytes: u64 = slot; + pub var lag: u32 = 1; + pub var polls: u64 = 0; + pub var queue: [2][512]u64 = undefined; + pub var head: [2]u32 = .{ 0, 0 }; + pub var len: [2]u32 = .{ 0, 0 }; + pub var inflight: u64 = 0; + pub var max_inflight: u64 = 0; + pub var max_write: u64 = 0; + pub var outside: u64 = 0; + pub var bytes: u64 = 0; + pub var posts: u64 = 0; + + pub fn reset(l: u32) void { + lag = l; + polls = 0; + head = .{ 0, 0 }; + len = .{ 0, 0 }; + inflight = 0; + max_inflight = 0; + max_write = 0; + outside = 0; + bytes = 0; + posts = 0; + } + + pub fn postSend(qp: *abi.Qp, wr: *abi.SendWr, _: *?*abi.SendWr) callconv(.c) c_int { + const q = qp.qp_num; + posts += 1; + var sum: u64 = 0; + var at: ?*abi.SendWr = wr; + while (at) |w| : (at = w.next) { + if (w.send_flags & abi.send_inline == 0) { + const sge = w.sg_list[0]; + const lo = @intFromPtr(region.ptr); + if (sge.addr < lo + send_off or sge.addr + sge.length > lo + send_off + 2 * slot_bytes) outside += 1; + max_write = @max(max_write, sge.length); + sum += sge.length; + } + if (w.send_flags & abi.send_signaled != 0) { + queue[q][(head[q] + len[q]) % 512] = sum; + len[q] += 1; + inflight += sum; + bytes += sum; + max_inflight = @max(max_inflight, inflight); + sum = 0; + } + } + return 0; + } + + /// One completion every `lag` polls of a QP's CQ: the HCA finishes writes far slower than they are posted. + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { + const q = cq.handle; + polls += 1; + if (polls % lag != 0 or len[q] == 0) return 0; + inflight -= queue[q][head[q]]; + head[q] = (head[q] + 1) % 512; + len[q] -= 1; + wc[0] = std.mem.zeroes(abi.Wc); + wc[0].status = abi.wc_success; + return 1; + } +}; + +/// QPs that complete nothing on their own: a QP moved to error flushes what it holds, as an HCA reports it. +pub const Halts = struct { + pub var halted: [2]bool = .{ false, false }; + pub var flushed_before_halt: u32 = 0; + pub var flushes: u32 = 0; + + pub fn reset() void { + halted = .{ false, false }; + flushed_before_halt = 0; + flushes = 0; + } + + pub fn modifyQp(qp: *abi.Qp, attr: *abi.QpAttr, mask: c_int) callconv(.c) c_int { + if (attr.qp_state == abi.qps_err and mask == abi.mask.state) halted[qp.qp_num] = true; + return 0; + } + + /// ibv_query_qp as the device answers it: ERR once halted, else RTS. + pub fn queryQp(qp: *abi.Qp, attr: *abi.QpAttr, _: c_int, _: *abi.QpInitAttr) callconv(.c) c_int { + attr.qp_state = if (halted[qp.qp_num]) abi.qps_err else abi.qps_rts; + return 0; + } + + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { + const q = cq.handle; + if (Fake.len[q] == 0 or !halted[q]) return 0; + if (!halted[0] or !halted[1]) flushed_before_halt += 1; + Fake.inflight -= Fake.queue[q][Fake.head[q]]; + Fake.head[q] = (Fake.head[q] + 1) % 512; + Fake.len[q] -= 1; + flushes += 1; + wc[0] = std.mem.zeroes(abi.Wc); + wc[0].status = 5; // IBV_WC_WR_FLUSH_ERR + return 1; + } + + /// Points a wire's QPs at these handlers. + pub fn install(hw: *Hw, w: *rv.VerbsWire) void { + w.v.api.ibv_modify_qp = modifyQp; + w.v.query_qp = queryQp; + for (&hw.ctx) |*c| c.ops.poll_cq = pollCq; + } +}; + +/// Two fake HCAs behind one VerbsWire, connected to a peer whose region has this rank's layout. +pub const Hw = struct { + ctx: [2]abi.Context = undefined, + qp: [2]abi.Qp = undefined, + cq: [2]abi.Cq = undefined, + mr: [2]abi.Mr = undefined, + + pub fn wire(hw: *Hw, region: []u8) rv.VerbsWire { + const s = Fake.slot_bytes; + var w: rv.VerbsWire = .{ .v = undefined, .traffic_class = 0, .recv_off = send_off + 2 * s, .flag_off = 4096, .send_off = send_off, .slot_bytes = s, .region = region }; + for (0..2) |i| { + hw.ctx[i].ops = .{ ._before_poll = @splat(null), .poll_cq = Fake.pollCq, ._before_post = @splat(null), .post_send = Fake.postSend, .post_recv = undefined, ._after = @splat(null) }; + hw.qp[i].qp_num = @intCast(i); + hw.cq[i].handle = @intCast(i); + hw.mr[i].lkey = 7; + w.hcas[i] = .{ .ctx = &hw.ctx[i], .pd = undefined, .mr = &hw.mr[i], .cq = &hw.cq[i], .qp = &hw.qp[i], .gid_index = 0, .gid = @splat(0), .mtu = 0, .lid = 0 }; + } + w.n_hca = 2; + return w; + } +}; diff --git a/zig/src/fabric/tp/roce_guard.zig b/zig/src/fabric/tp/roce_guard.zig new file mode 100644 index 0000000000..8f85247afd --- /dev/null +++ b/zig/src/fabric/tp/roce_guard.zig @@ -0,0 +1,73 @@ +//! Every RoCE post's bounds before the verb: each HCA's stripe inside its slot and the registered region. +const std = @import("std"); + +pub const max_hcas = 2; +pub const flag_stride = 128; + +pub const Error = error{ OutOfRegion, WriteTooLarge }; + +/// The largest RDMA write a post may make, whatever the slot (a 9.6 MB post faulted a DGX Spark's HCA); never raised. +pub const max_write: u32 = 1 << 20; + +/// Both ranks' region layout (roce.zig's): byte offsets and the registered length every post must stay inside. +pub const Layout = struct { send: u64, recv: u64, flags: u64, slot: u64, len: u64 }; + +/// One HCA's share of a post: its data stripe in this rank's send slot and the peer's receive slot, and its flag word. +pub const Stripe = struct { local: u64, remote: u64, len: u32, flag: u64 }; + +/// Slot `slot`'s first `padded` bytes striped over `n_hca` HCAs; error.OutOfRegion before any verb past a slot or MR. +pub fn stripes(l: Layout, n_hca: u32, slot: u32, padded: u32, out: *[max_hcas]Stripe) Error!void { + if (n_hca == 0 or n_hca > max_hcas or slot > 1) return error.OutOfRegion; + if (padded > max_write) return error.WriteTooLarge; + if (padded == 0 or padded % 16 != 0 or padded > l.slot) return error.OutOfRegion; + const packs = padded / 16; + var offset: u64 = 0; + for (out[0..n_hca], 0..) |*o, i| { + var stripe = packs / n_hca; + if (i < packs % n_hca) stripe += 1; + o.* = .{ .local = l.send + slot * l.slot + offset, .remote = l.recv + slot * l.slot + offset, .len = stripe * 16, .flag = l.flags + (@as(u64, slot) * max_hcas + i) * flag_stride }; + if (!within(o.local, o.len, l.send + slot * l.slot, l.slot) or !within(o.remote, o.len, l.recv + slot * l.slot, l.slot)) return error.OutOfRegion; + if (!within(o.flag, 4, l.flags, l.send - l.flags) or !within(o.local, o.len, 0, l.len) or !within(o.remote, o.len, 0, l.len)) return error.OutOfRegion; + offset += o.len; + } +} + +fn within(at: u64, len: u64, lo: u64, size: u64) bool { + return at >= lo and at + len <= lo + size; +} + +/// roce.zig's layout for a 64 KiB slot: ctrl 4 KiB, flags 4 KiB, send [2][slot], recv [2][slot]. +fn testLayout() Layout { + const slot: u64 = 64 << 10; + return .{ .send = 8192, .recv = 8192 + 2 * slot, .flags = 4096, .slot = slot, .len = 8192 + 4 * slot }; +} + +test "a post striped over two HCAs: the stripes tile the slot's first bytes, each flag on its own word" { + const l = testLayout(); + var s: [max_hcas]Stripe = undefined; + try stripes(l, 2, 1, 48, &s); + try std.testing.expectEqual(Stripe{ .local = l.send + l.slot, .remote = l.recv + l.slot, .len = 32, .flag = l.flags + 2 * flag_stride }, s[0]); + try std.testing.expectEqual(Stripe{ .local = l.send + l.slot + 32, .remote = l.recv + l.slot + 32, .len = 16, .flag = l.flags + 3 * flag_stride }, s[1]); + try stripes(l, 2, 1, @intCast(l.slot), &s); + try std.testing.expectEqual(l.recv + 2 * l.slot, s[1].remote + s[1].len); +} + +test "a post past its slot, a third slot, an unpadded or empty post, or no HCA is refused" { + const l = testLayout(); + var s: [max_hcas]Stripe = undefined; + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, @intCast(l.slot + 16), &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 2, 16, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 24, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 0, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 0, 0, 16, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 3, 0, 16, &s)); + var big = l; + big.slot = 10 << 20; + big.recv = big.send + 2 * big.slot; + big.len = big.send + 4 * big.slot; + try std.testing.expectError(error.WriteTooLarge, stripes(big, 2, 0, max_write + 16, &s)); + try stripes(big, 2, 0, max_write, &s); + var short = l; + short.len -= 16; + try std.testing.expectError(error.OutOfRegion, stripes(short, 2, 1, @intCast(l.slot), &s)); +} diff --git a/zig/src/fabric/tp/roce_host.zig b/zig/src/fabric/tp/roce_host.zig new file mode 100644 index 0000000000..1dec07e9ba --- /dev/null +++ b/zig/src/fabric/tp/roce_host.zig @@ -0,0 +1,54 @@ +//! `tf-tp-test roce-host`: the RoCE wire and proxy between two hosts without a GPU, checked bit for bit, then timed. +const std = @import("std"); +const tp = @import("tp"); +const roce = tp.roce; +const rv = tp.roce_verbs; + +const sizes = [_]usize{ 16, 4096, 8192, 16384, 65536, 262144 }; + +pub fn run(cfg: tp.Config, gpa: std.mem.Allocator) !u8 { + var b = try tp.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var n: usize = 0; + if (cfg.roce_wire == .verbs) n = try tp.roce_pick.resolve(cfg, &b, &specs); + const r = try roce.Roce.init(null, &b, "", .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = 256 * 1024, .timeout_ns = 10 * std.time.ns_per_s }); + defer r.deinit(); + const in = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(in); + const peer = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(peer); + const out = try gpa.alloc(u8, 512 * 1024); + defer gpa.free(out); + // 1. every size from 1 byte to the slot, both slots, bit for bit + var checked: u32 = 0; + for (0..1000) |k| { + const len = 1 + (k * 7919) % (256 * 1024); + tp.reference.fill(in[0..len], .u8, cfg.rank, k); + tp.reference.fill(peer[0..len], .u8, 1 - cfg.rank, k); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + const mine_at = cfg.rank * len; + const peer_at = (1 - cfg.rank) * len; + if (!std.mem.eql(u8, out[mine_at..][0..len], in[0..len]) or !std.mem.eql(u8, out[peer_at..][0..len], peer[0..len])) { + std.debug.print("FAIL roce-host gather {d} of {d} bytes\n", .{ k, len }); + return 1; + } + checked += 1; + } + try r.check(); + std.debug.print("PASS roce-host: {d} gathers of 1 B to 256 KiB bit-exact over {s}\n", .{ checked, @tagName(cfg.roce_wire) }); + // 2. one exchange's host-to-host time by size: median and p90 of 2,000 back-to-back ops + var times: [2000]u64 = undefined; + for (sizes) |len| { + for (×) |*t| { + const t0 = tp.sock.nowNs(); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + t.* = tp.sock.nowNs() - t0; + } + std.mem.sort(u64, ×, {}, std.sort.asc(u64)); + if (cfg.rank == 0) std.debug.print("RESULT {{\"backend\":\"roce-host\",\"wire\":\"{s}\",\"bytes\":{d},\"p50_us\":{d:.2},\"p90_us\":{d:.2}}}\n", .{ @tagName(cfg.roce_wire), len, @as(f64, @floatFromInt(times[1000])) / 1e3, @as(f64, @floatFromInt(times[1800])) / 1e3 }); + } + try r.check(); + try b.agree("done", "roce-host end"); + return 0; +} diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig new file mode 100644 index 0000000000..fd7be53262 --- /dev/null +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -0,0 +1,110 @@ +//! VerbsWire on fake verbs whose completions lag: what it holds in flight and what a failed transport leaves behind. +const std = @import("std"); +const rv = @import("roce_verbs.zig"); +const fake = @import("roce_fake.zig"); +const Roce = @import("roce.zig").Roce; +const Fake = fake.Fake; +const Halts = fake.Halts; +const slot = fake.slot; + +test "chunked posts: the wire's own in-flight count and peak match the HCAs', and reach the 2 MiB cap" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = slot; + Fake.reset(3); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + for (1..9) |seq| { + try w.post(@intCast(seq & 1), @intCast(seq), slot); + try std.testing.expectEqual(Fake.inflight, w.inflight()); + try std.testing.expect(w.inflight() <= rv.max_inflight); + } + while (w.inflight() > 0) try w.reap(); + try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + try std.testing.expectEqual(Fake.max_inflight, w.peak.load(.monotonic)); + try std.testing.expectEqual(@as(u64, 2 << 20), w.peak.load(.monotonic)); +} + +test "a failed transport: both QPs go to error before the first flush and every posted write is reaped" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(slot)); + defer std.testing.allocator.free(region); + @memset(region, 0); + Fake.region = region; + Fake.slot_bytes = slot; + Fake.reset(1); + Halts.reset(); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + Halts.install(&hw, &w); + // two 1 MiB ops posted, neither completed: 2 MiB in flight at the HCAs, the link busy + try std.testing.expect(w.idle()); + try w.post(1, 1, slot); + try w.post(0, 2, slot); + try std.testing.expectEqual(@as(u64, 2 * slot), w.inflight()); + try std.testing.expect(!w.idle()); + try std.testing.expect(w.quiesce(std.time.ns_per_s)); + try std.testing.expect(Halts.halted[0] and Halts.halted[1]); + try std.testing.expectEqual(@as(u32, 0), Halts.flushed_before_halt); + try std.testing.expectEqual(@as(u32, 4), Halts.flushes); + try std.testing.expectEqual(@as(u64, 0), w.inflight()); + try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + // the line tf-tp-test's transport run reads: the device's states and every reap, all of them flushes + var line: [256]u8 = undefined; + try std.testing.expect(w.idle()); + try std.testing.expectEqualStrings("HCA 0 QP ERR; posted 2, completed 2; HCA 1 QP ERR; posted 2, completed 2; reaped 4 (flushed 4), still posted 0", w.report(&line)); + // an aborted transport reports itself failed: its deinit keeps the region + Halts.reset(); + var r: Roce = .{ .me = 0, .slot_bytes = slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; + try std.testing.expect(!r.failed()); + r.abort(); + try std.testing.expect(Halts.halted[0] and Halts.halted[1]); + try std.testing.expect(r.failed()); +} + +test "a QP whose inline capacity cannot take the 4-byte flag post is refused" { + try std.testing.expect(!rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 3 })); + try std.testing.expect(rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 4 })); +} + +/// /proc/self/smaps' VmFlags of the mapping that contains `at`. +fn vmFlags(at: usize, buf: []u8) ![]const u8 { + const fd = std.c.open("/proc/self/smaps", .{}); + if (fd < 0) return error.NoSmaps; + defer _ = std.c.close(fd); + var n: usize = 0; + while (n < buf.len) { + const got = std.c.read(fd, buf[n..].ptr, buf.len - n); + if (got <= 0) break; + n += @intCast(got); + } + var lines = std.mem.splitScalar(u8, buf[0..n], '\n'); + var inside = false; + while (lines.next()) |l| { + const sp = std.mem.indexOfScalar(u8, l, ' ') orelse l.len; + if (std.mem.indexOfScalar(u8, l[0..sp], '-')) |dash| { + const lo = std.fmt.parseInt(usize, l[0..dash], 16) catch continue; + const hi = std.fmt.parseInt(usize, l[dash + 1 .. sp], 16) catch continue; + inside = lo <= at and at < hi; + } + if (inside and std.mem.startsWith(u8, l, "VmFlags:")) return l; + } + return error.NotFound; +} + +test "the region is DONTFORK and NOHUGEPAGE before registration" { + if (@import("builtin").os.tag != .linux) return error.SkipZigTest; + const len = Roce.regionBytes(slot); + const m = try std.posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + defer std.posix.munmap(m); + const buf = try std.testing.allocator.alloc(u8, 4 << 20); + defer std.testing.allocator.free(buf); + // smaps: dc = MADV_DONTFORK, nh = MADV_NOHUGEPAGE + const before = try vmFlags(@intFromPtr(m.ptr), buf); + try std.testing.expect(std.mem.indexOf(u8, before, " dc") == null and std.mem.indexOf(u8, before, " nh") == null); + try rv.prepareRegion(m); + // An interior page stays inside the VMA even when a neighbouring mapping merges with it. + const flags = try vmFlags(@intFromPtr(m.ptr) + 4096, buf); + try std.testing.expect(std.mem.indexOf(u8, flags, " dc") != null); + try std.testing.expect(std.mem.indexOf(u8, flags, " nh") != null); +} diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig new file mode 100644 index 0000000000..d1ed7a6529 --- /dev/null +++ b/zig/src/fabric/tp/roce_pick.zig @@ -0,0 +1,448 @@ +//! Which RoCE devices a rank stripes over: TF_TP_ROCE_HCA, else the usable ports on an IPv4 /24 the peer also has. +const std = @import("std"); +const log = @import("log.zig"); +const rv = @import("roce_verbs.zig"); +const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; +const Verbs = @import("verbs").Verbs; +const OpenFn = *const fn (?[:0]const u8) @import("verbs").Error!Verbs; + +/// Ports a rank lists to its peers (a DGX Spark's ConnectX-7 shows four RDMA devices). +pub const max_ports = 8; +const blob_bytes = 1 + 4 * max_ports; + +/// How many devices, and how: on a /24 the peer shares (`shared`), named by TF_TP_ROCE_HCA (`explicit`), or neither. +pub const Pick = struct { + n: usize, + shared: bool, + explicit: bool = false, + verbs_open_error: ?@import("verbs").Error = null, +}; + +fn usable(p: rv.Port) bool { + return p.active and p.ethernet and p.ipv4 != null; +} + +fn subnet(a: [4]u8) u32 { + return @as(u32, a[0]) << 16 | @as(u32, a[1]) << 8 | a[2]; +} + +fn countIn(list: []const [4]u8, net: u32) usize { + var k: usize = 0; + for (list) |a| k += @intFromBool(subnet(a) == net); + return k; +} + +/// Usable ports on a /24 the peers list, by subnet, each /24 as often as both sides have it; else the first usable. +pub fn choose(ports: []const rv.Port, peer: []const [4]u8, out: *[rv.max_hcas]rv.HcaSpec, limit: usize) Pick { + const cap = @min(limit, rv.max_hcas); + var order: [max_ports]usize = undefined; + var m: usize = 0; + for (ports[0..@min(ports.len, max_ports)], 0..) |p, i| if (usable(p)) { + order[m] = i; + m += 1; + }; + var n: usize = 0; + for (order[0..m]) |i| n += @intFromBool(countIn(peer, subnet(ports[i].ipv4.?)) > 0); + if (n == 0) { + for (order[0..@min(m, cap)], 0..) |i, k| out[k] = ports[i].spec; + return .{ .n = @min(m, cap), .shared = false }; + } + // stable insertion sort by subnet (at most max_ports entries) + for (1..m) |a| { + var j = a; + while (j > 0 and subnet(ports[order[j - 1]].ipv4.?) > subnet(ports[order[j]].ipv4.?)) : (j -= 1) { + std.mem.swap(usize, &order[j - 1], &order[j]); + } + } + n = 0; + var taken: [max_ports]u32 = undefined; + for (order[0..m]) |i| { + if (n == cap) break; + const net = subnet(ports[i].ipv4.?); + var mine: usize = 0; + for (taken[0..n]) |t| mine += @intFromBool(t == net); + if (mine >= countIn(peer, net)) continue; + taken[n] = net; + out[n] = ports[i].spec; + n += 1; + } + return .{ .n = n, .shared = true }; +} + +/// `[count][count x IPv4]` of the usable ports, the blob each rank sends. +fn encode(ports: []const rv.Port, out: *[blob_bytes]u8) []const u8 { + var n: u8 = 0; + for (ports[0..@min(ports.len, max_ports)]) |p| if (usable(p)) { + out[1 + 4 * @as(usize, n) ..][0..4].* = p.ipv4.?; + n += 1; + }; + out[0] = n; + return out[0 .. 1 + 4 * @as(usize, n)]; +} + +/// A peer's blob into `out`; a malformed one counts as no ports. +fn decode(blob: []const u8, out: [][4]u8) usize { + if (blob.len == 0 or blob[0] > max_ports or blob.len != 1 + 4 * @as(usize, blob[0]) or out.len < blob[0]) return 0; + for (0..blob[0]) |k| out[k] = blob[1 + 4 * k ..][0..4].*; + return blob[0]; +} + +/// The devices for the verbs wire. A collective: every rank calls it, including ranks that set TF_TP_ROCE_HCA. +pub fn resolve(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !usize { + return (try select(cfg, b, out)).n; +} + +/// `resolve` with how the devices were chosen, over this host's ports (a collective, as `resolve`). +pub fn select(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !Pick { + return selectWithOpen(cfg, b, out, Verbs.open); +} + +fn selectWithOpen(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec, open: OpenFn) !Pick { + var ports: [max_ports]rv.Port = undefined; + var n_ports: usize = 0; + var open_error: ?@import("verbs").Error = null; + if (open(null)) |opened| { + var v = opened; + defer v.close(); + n_ports = rv.scanPorts(&v, &ports); + } else |e| { + open_error = e; + } + var pick = try selectFrom(cfg, b, ports[0..n_ports], out); + pick.verbs_open_error = open_error; + return pick; +} + +/// `select` over given ports (the scan's, or a test's fixture ports). +pub fn selectFrom(cfg: Config, b: *bootstrap.Bootstrap, ports: []const rv.Port, out: *[rv.max_hcas]rv.HcaSpec) !Pick { + var mine: [blob_bytes]u8 = undefined; + var buf: [bootstrap.max_ranks * blob_bytes]u8 = undefined; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + const all = try b.allGather(encode(ports, &mine), &buf, &blobs); + if (cfg.roce_hca) |text| return .{ .n = try rv.parseSpecs(text, out), .shared = false, .explicit = true }; + var peer: [bootstrap.max_ranks * max_ports][4]u8 = undefined; + var n_peer: usize = 0; + for (all, 0..) |blob, r| if (r != b.rank) { + n_peer += decode(blob, peer[n_peer..]); + }; + const pick = choose(ports, peer[0..n_peer], out, cfg.roce_hcas); + if (!pick.shared and pick.n > 0) log.warn("tp roce: no RoCE port shares an IPv4 /24 with the peer's; using the first {d} active (set TF_TP_ROCE_HCA)", .{pick.n}); + return pick; +} + +/// Device names joined by commas, with each `:gid` when `gid` (TF_TP_ROCE_HCA's form) or without (NCCL_IB_HCA's). +pub fn joinNames(specs: []const rv.HcaSpec, gid: bool, buf: []u8) error{WriteFailed}![]const u8 { + var w: std.Io.Writer = .fixed(buf); + for (specs, 0..) |s, i| { + try w.print("{s}{s}", .{ if (i == 0) "" else ",", s.name[0..s.name_len] }); + if (gid) if (s.gid) |g| try w.print(":{d}", .{g}); + } + return w.buffered(); +} + +/// Room for `joinNames` of every device a rank stripes over (a name is under 64 bytes, a GID index at most 3 digits). +pub const list_bytes = rv.max_hcas * (64 + 4); + +pub const HcaError = error{ + /// no active Ethernet port with an IPv4 RoCE v2 GID on this host + NoRoceDevice, + /// this host's ports share no IPv4 /24 with the peer's: NCCL would not reach the peer on them either + NoSharedSubnet, +} || @import("verbs").Error; + +/// `tf-tp-test hca`'s line for run_hosts.sh: `NCCL_IB_HCA=` the picked devices; never the first local ports alone. +pub fn hcaLine(pick: Pick, specs: []const rv.HcaSpec, buf: *[list_bytes + 16]u8) HcaError![]const u8 { + if (pick.n == 0) { + if (pick.verbs_open_error) |e| return e; + return error.NoRoceDevice; + } + if (!pick.shared and !pick.explicit) return error.NoSharedSubnet; + const prefix = "NCCL_IB_HCA="; + @memcpy(buf[0..prefix.len], prefix); + const joined = joinNames(specs[0..pick.n], false, buf[prefix.len .. buf.len - 1]) catch unreachable; // list_bytes holds them + buf[prefix.len + joined.len] = '\n'; + return buf[0 .. prefix.len + joined.len + 1]; +} + +fn port(name: []const u8, active: bool, ipv4: ?[4]u8) rv.Port { + var p: rv.Port = .{ .active = active, .ethernet = true, .ipv4 = ipv4 }; + p.spec.name_len = @intCast(name.len); + @memcpy(p.spec.name[0..name.len], name); + p.spec.gid = 3; + return p; +} + +fn names(specs: []const rv.HcaSpec, buf: []u8) []const u8 { + return joinNames(specs, false, buf) catch unreachable; +} + +// A host with two neighbours: port f0 cabled to one Spark (both functions on one /24), port f1 to the other. +const a_ports = [_]rv.Port{ + port("rocep1s0f0", true, .{ 203, 0, 113, 1 }), + port("rocep1s0f1", true, .{ 192, 0, 2, 1 }), + port("roceP2p1s0f0", true, .{ 203, 0, 113, 65 }), + port("roceP2p1s0f1", true, .{ 198, 51, 100, 1 }), +}; +// Its peer for this run: only port f1 is cabled, f0 is down. +const b_ports = [_]rv.Port{ + port("rocep1s0f0", false, null), + port("rocep1s0f1", true, .{ 192, 0, 2, 2 }), + port("roceP2p1s0f0", false, null), + port("roceP2p1s0f1", true, .{ 198, 51, 100, 2 }), +}; + +fn listed(ports: []const rv.Port, out: *[max_ports][4]u8) []const [4]u8 { + var blob: [blob_bytes]u8 = undefined; + return out[0..decode(encode(ports, &blob), out)]; +} + +test "a two-neighbour host picks the rails its peer shares, not its first two active ports" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [128]u8 = undefined; + // the old rule (no peer list): the first two active ports, one of them cabled to the other neighbour + const old = choose(&a_ports, &.{}, &specs, 2); + try std.testing.expect(!old.shared); + try std.testing.expectEqualStrings("rocep1s0f0,rocep1s0f1", names(specs[0..old.n], &buf)); + var peer_buf: [max_ports][4]u8 = undefined; + const b_addrs = listed(&b_ports, &peer_buf); + try std.testing.expectEqual(@as(usize, 0), countIn(b_addrs, subnet(.{ 203, 0, 113, 1 }))); + // with the peer's addresses: the two functions of the cabled port, on each side + const a = choose(&a_ports, b_addrs, &specs, 2); + try std.testing.expect(a.shared); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", names(specs[0..a.n], &buf)); + var a_buf: [max_ports][4]u8 = undefined; + const b = choose(&b_ports, listed(&a_ports, &a_buf), &specs, 2); + try std.testing.expectEqual(@as(usize, 2), b.n); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", names(specs[0..b.n], &buf)); + // one rail asked for: the lower /24 on both sides + try std.testing.expectEqualStrings("rocep1s0f1", names(specs[0..choose(&a_ports, b_addrs, &specs, 1).n], &buf)); +} + +test "two rails on one /24 count once per peer port; no shared /24 keeps the list order" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [128]u8 = undefined; + // the neighbour cabled to port f0 lists one address on that /24: one rail each, equal counts + const one = [_][4]u8{.{ 203, 0, 113, 2 }}; + const p = choose(&a_ports, &one, &specs, 2); + try std.testing.expectEqual(@as(usize, 1), p.n); + try std.testing.expectEqualStrings("rocep1s0f0", names(specs[0..p.n], &buf)); + const two = [_][4]u8{ .{ 203, 0, 113, 2 }, .{ 203, 0, 113, 66 } }; + try std.testing.expectEqualStrings("rocep1s0f0,roceP2p1s0f0", names(specs[0..choose(&a_ports, &two, &specs, 2).n], &buf)); + // inactive and non-Ethernet ports never count; a peer on other subnets falls back to the first usable ports + var mixed = b_ports; + mixed[1].ethernet = false; + const other = [_][4]u8{.{ 198, 51, 100, 9 }}; + const q = choose(&mixed, &.{.{ 203, 0, 113, 9 }}, &specs, 2); + try std.testing.expect(!q.shared); + try std.testing.expectEqualStrings("roceP2p1s0f1", names(specs[0..q.n], &buf)); + try std.testing.expectEqualStrings("roceP2p1s0f1", names(specs[0..choose(&mixed, &other, &specs, 2).n], &buf)); + // a malformed blob is no ports; the detected GID index travels with the pick + var out: [max_ports][4]u8 = undefined; + try std.testing.expectEqual(@as(usize, 0), decode(&.{ 2, 1, 2, 3, 4 }, &out)); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); +} + +test "hca's line: the picked names for NCCL; no device or no shared /24 is refused, an explicit list is taken" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var line: [list_bytes + 16]u8 = undefined; + var peer_buf: [max_ports][4]u8 = undefined; + const shared = choose(&a_ports, listed(&b_ports, &peer_buf), &specs, 2); + try std.testing.expectEqualStrings("NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1\n", try hcaLine(shared, &specs, &line)); + var buf: [list_bytes]u8 = undefined; + try std.testing.expectEqualStrings("rocep1s0f1:3,roceP2p1s0f1:3", try joinNames(specs[0..shared.n], true, &buf)); + // the old fallback (first active ports, no /24 in common) is what NCCL failed on: refused, never printed + try std.testing.expectError(error.NoSharedSubnet, hcaLine(choose(&a_ports, &.{.{ 10, 9, 9, 9 }}, &specs, 2), &specs, &line)); + try std.testing.expectError(error.NoRoceDevice, hcaLine(choose(&.{}, &.{}, &specs, 2), &specs, &line)); + const n = try rv.parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1\n", try hcaLine(.{ .n = n, .shared = false, .explicit = true }, &specs, &line)); +} + +/// One RDMA device for `rv.scanPorts` through a stub verbs table: port 1 up or down, Ethernet, mlx5's GIDs 0-3. +const FakeDev = struct { name: [:0]const u8, up: bool, ipv4: [4]u8 = @splat(0) }; + +const ScanFixture = struct { + const vabi = @import("verbs").abi; + var host: []const FakeDev = &.{}; + var list: [max_ports:null]?*vabi.Device = undefined; + + fn handle(i: usize) usize { + return 0x10000 * (i + 1); + } + fn index(p: *const anyopaque) usize { + return @intFromPtr(p) / 0x10000 - 1; + } + fn getList(num: *c_int) callconv(.c) ?[*:null]?*vabi.Device { + for (host, 0..) |_, i| list[i] = @ptrFromInt(handle(i)); + list[host.len] = null; + num.* = @intCast(host.len); + return &list; + } + fn freeList(_: [*:null]?*vabi.Device) callconv(.c) void {} + fn devName(d: *vabi.Device) callconv(.c) ?[*:0]const u8 { + return host[index(d)].name.ptr; + } + fn openDev(d: *vabi.Device) callconv(.c) ?*vabi.Context { + return @ptrFromInt(@intFromPtr(d)); + } + fn closeDev(_: *vabi.Context) callconv(.c) c_int { + return 0; + } + fn queryPort(c: *vabi.Context, _: u8, a: *vabi.PortAttr) callconv(.c) c_int { + a.* = std.mem.zeroes(vabi.PortAttr); + a.state = if (host[index(c)].up) vabi.port_active else 1; + a.link_layer = 2; // Ethernet + a.gid_tbl_len = 4; + return 0; + } + fn queryGid(c: *vabi.Context, _: u8, i: c_int, g: *vabi.Gid) callconv(.c) c_int { + g.* = .{ .raw = @splat(0) }; + if (i < 2) { + g.raw[0] = 0xfe; + g.raw[1] = 0x80; + g.raw[15] = @intCast(index(c) + 1); + } else { + g.raw[10] = 0xff; + g.raw[11] = 0xff; + g.raw[12..16].* = host[index(c)].ipv4; + } + return 0; + } + fn gidEx(_: *vabi.Context, _: u32, i: u32, e: *vabi.GidEntry, _: u32, _: usize) callconv(.c) c_int { + e.* = std.mem.zeroes(vabi.GidEntry); + e.gid_type = if (i % 2 == 1) 2 else 1; // odd indices RoCE v2 + return 0; + } + + /// `rv.scanPorts` over `devs` (names that no real /sys/class/infiniband has, so sysfs never answers for them). + fn scan(devs: []const FakeDev, out: *[max_ports]rv.Port) []const rv.Port { + host = devs; + var v: Verbs = undefined; + v.api.ibv_get_device_list = getList; + v.api.ibv_free_device_list = freeList; + v.api.ibv_get_device_name = devName; + v.api.ibv_open_device = openDev; + v.api.ibv_close_device = closeDev; + v.api.ibv_query_port = queryPort; + v.api.ibv_query_gid = queryGid; + v.gid_ex = gidEx; + return out[0..rv.scanPorts(&v, out)]; + } +}; + +/// #636's reported two GB10 hosts, modelled (no capture): device order and picks reported; addresses, states assumed. +const hosts = [2][4]FakeDev{ .{ + .{ .name = "fx_rocep1s0f0", .up = true, .ipv4 = .{ 192, 0, 2, 1 } }, + .{ .name = "fx_rocep1s0f1", .up = false }, + .{ .name = "fx_roceP2p1s0f0", .up = true, .ipv4 = .{ 198, 51, 100, 1 } }, + .{ .name = "fx_roceP2p1s0f1", .up = false }, +}, .{ + .{ .name = "fx_rocep1s0f0", .up = false }, + .{ .name = "fx_rocep1s0f1", .up = true, .ipv4 = .{ 192, 0, 2, 2 } }, + .{ .name = "fx_roceP2p1s0f0", .up = false }, + .{ .name = "fx_roceP2p1s0f1", .up = true, .ipv4 = .{ 198, 51, 100, 2 } }, +} }; + +test "the port scan reads a host's devices through verbs: state, Ethernet, the IPv4 RoCE v2 GID index" { + var out: [max_ports]rv.Port = undefined; + const a = ScanFixture.scan(&hosts[0], &out); + try std.testing.expectEqual(@as(usize, 4), a.len); + try std.testing.expectEqualStrings("fx_rocep1s0f0", a[0].spec.name[0..a[0].spec.name_len]); + try std.testing.expect(a[0].active and a[0].ethernet and !a[1].active); + try std.testing.expectEqual(@as(?[4]u8, .{ 192, 0, 2, 1 }), a[0].ipv4); + try std.testing.expectEqual(@as(?u8, 3), a[0].spec.gid); // the v2 IPv4 GID, not v1's index 2 + try std.testing.expectEqual(@as(?[4]u8, null), a[1].ipv4); +} + +/// One rank's `selectFrom` over a real bootstrap (rank 0 on `l`), its hca line or error into `out`. +fn hcaRank(cfg: Config, l: ?sock.Fd, ports: []const rv.Port, open: ?OpenFn, out: *[list_bytes + 16]u8, line: *HcaError![]const u8, err: *?anyerror) void { + var b = (if (l) |fd| bootstrap.Bootstrap.acceptOn(fd, 2, 5 * std.time.ns_per_s) else bootstrap.Bootstrap.open(cfg)) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + const pick = (if (open) |f| selectWithOpen(cfg, &b, &specs, f) else selectFrom(cfg, &b, ports, &specs)) catch |e| { + err.* = e; + return; + }; + line.* = hcaLine(pick, &specs, out); +} + +fn hcaPair(cfgs: [2]Config, ports: [2][]const rv.Port, lines: *[2]HcaError![]const u8, bufs: *[2][list_bytes + 16]u8) !void { + return hcaPairWithOpen(cfgs, ports, .{ null, null }, lines, bufs); +} + +fn hcaPairWithOpen(cfgs: [2]Config, ports: [2][]const rv.Port, open: [2]?OpenFn, lines: *[2]HcaError![]const u8, bufs: *[2][list_bytes + 16]u8) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + var errs: [2]?anyerror = @splat(null); + var c1 = cfgs[1]; + c1.link_port = l.port; + c1.connect_timeout_ns = 5 * std.time.ns_per_s; + const t = try std.Thread.spawn(.{}, hcaRank, .{ c1, null, ports[1], open[1], &bufs[1], &lines[1], &errs[1] }); + hcaRank(cfgs[0], l.fd, ports[0], open[0], &bufs[0], &lines[0], &errs[0]); + t.join(); + for (errs) |e| if (e) |x| return x; +} + +fn missingLibrary(_: ?[:0]const u8) @import("verbs").Error!Verbs { + return error.LibraryUnavailable; +} + +fn missingSymbol(_: ?[:0]const u8) @import("verbs").Error!Verbs { + return error.MissingSymbol; +} + +test "hca keeps a verbs load failure while completing the peer's port exchange" { + var lines: [2]HcaError![]const u8 = undefined; + var bufs: [2][list_bytes + 16]u8 = undefined; + var scanned: [max_ports]rv.Port = undefined; + const peer = ScanFixture.scan(&hosts[1], &scanned); + const cfgs: [2]Config = .{ .{ .rank = 0 }, .{ .rank = 1 } }; + try hcaPairWithOpen(cfgs, .{ &.{}, peer }, .{ missingLibrary, null }, &lines, &bufs); + try std.testing.expectError(error.LibraryUnavailable, lines[0]); + try std.testing.expectError(error.NoSharedSubnet, lines[1]); +} + +test "hca distinguishes missing verbs symbols from an empty successful scan" { + const cfg: Config = .{ .world = 1 }; + var b = try bootstrap.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [list_bytes + 16]u8 = undefined; + const pick = try selectWithOpen(cfg, &b, &specs, missingSymbol); + try std.testing.expectError(error.MissingSymbol, hcaLine(pick, &specs, &buf)); + try std.testing.expectError(error.NoRoceDevice, hcaLine(try selectFrom(cfg, &b, &.{}, &specs), &specs, &buf)); +} + +test "hca keeps an explicit list when the verbs scan cannot open its library" { + const cfg: Config = .{ .world = 1, .roce_hca = "fx_rocep1s0f0:3" }; + var b = try bootstrap.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [list_bytes + 16]u8 = undefined; + const pick = try selectWithOpen(cfg, &b, &specs, missingLibrary); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0\n", try hcaLine(pick, &specs, &buf)); +} + +test "hca over a real bootstrap: both ranks print their peer-linked ports; an explicit rank keeps its list; no ports fail both" { + var lines: [2]HcaError![]const u8 = undefined; + var bufs: [2][list_bytes + 16]u8 = undefined; + var scanned: [2][max_ports]rv.Port = undefined; + const ports = [2][]const rv.Port{ ScanFixture.scan(&hosts[0], &scanned[0]), ScanFixture.scan(&hosts[1], &scanned[1]) }; + const auto: [2]Config = .{ .{ .rank = 0 }, .{ .rank = 1 } }; + try hcaPair(auto, ports, &lines, &bufs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0,fx_roceP2p1s0f0\n", try lines[0]); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f1,fx_roceP2p1s0f1\n", try lines[1]); + // TF_TP_ROCE_HCA on one rank: its list as given (no GID for NCCL); it still sends its ports to the other + const one: [2]Config = .{ .{ .rank = 0, .roce_hca = "fx_rocep1s0f0:3" }, .{ .rank = 1 } }; + try hcaPair(one, ports, &lines, &bufs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0\n", try lines[0]); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f1,fx_roceP2p1s0f1\n", try lines[1]); + // a host without usable ports: it has none, its peer shares no /24 with it, and both refuse + try hcaPair(auto, .{ &.{}, ports[1] }, &lines, &bufs); + try std.testing.expectError(error.NoRoceDevice, lines[0]); + try std.testing.expectError(error.NoSharedSubnet, lines[1]); +} diff --git a/zig/src/fabric/tp/roce_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig new file mode 100644 index 0000000000..a284293000 --- /dev/null +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -0,0 +1,120 @@ +//! VerbsWire.post on fake verbs: a post that would leave its slot or region is refused before any verb. +const std = @import("std"); +const fake = @import("roce_fake.zig"); +const Roce = @import("roce.zig").Roce; +const Config = @import("config.zig").Config; +const Fake = fake.Fake; + +test "a post past its slot is refused before any verb; a full slot posts inside the region" { + const small: u64 = 64 << 10; + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(small)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = small; + defer Fake.slot_bytes = fake.slot; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.OutOfRegion, w.post(1, 1, small + 16)); + try std.testing.expectError(error.OutOfRegion, w.post(2, 1, 16)); + try std.testing.expectEqual(@as(u64, 0), Fake.posts); + try w.post(1, 1, small); + try std.testing.expectEqual(@as(u64, 2), Fake.posts); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); +} + +test "no RDMA write past 1 MiB, whatever the slot: a 9.57 MB and a 1 MiB + 16 post refused in a 10 MiB slot" { + const big: u64 = 10 << 20; + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(big)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = big; + defer Fake.slot_bytes = fake.slot; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.WriteTooLarge, w.post(0, 1, 9_568_256)); + try std.testing.expectError(error.WriteTooLarge, w.post(1, 1, (1 << 20) + 16)); + try std.testing.expectEqual(@as(u64, 0), Fake.posts); + try w.post(0, 1, 1 << 20); + try std.testing.expectEqual(@as(u64, 1 << 20), Fake.max_write * 2); +} + +const Env = struct { + var pairs: []const [2][]const u8 = &.{}; + + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "boot refuses a RoCE slot past 1 MiB (TF_TP_MAILBOX_MAX_KB), never clamping it" { + for ([_][]const u8{ "1028", "2048", "16384" }) |kb| { + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", kb } }; + try std.testing.expectError(error.BadSetting, Config.parse(&Env.get)); + } + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" } }; + try std.testing.expectEqual(@as(usize, 1 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); + // the cap is RoCE's: the same-host mailbox keeps its larger slots + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "4096" } }; + try std.testing.expectEqual(@as(usize, 4 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); +} + +test "chunked strides posted back to back: at most 2 MiB in flight over both HCAs and no write past 1 MiB" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = fake.slot; + // gathers of R x 512 x 584 B strides and 584 B rows up to a 128 MiB union, each in <= 1 MiB ops + var strides: [16]u64 = undefined; + var n: usize = 0; + for ([_]u64{ 1, 2, 3, 4, 6, 16, 32, 64 }) |r| { + strides[n] = r * 512 * 584; + n += 1; + } + for ([_]u64{ 1795, 1796, (128 << 20) / (3 * 584) }) |m| { + strides[n] = m * 584; + n += 1; + } + for ([_]u32{ 1, 3, 17 }) |lag| { + Fake.reset(lag); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + var seq: u32 = 0; + var want: u64 = 0; + for (strides[0..n]) |stride| { + var off: u64 = 0; + while (off < stride) : (off += fake.slot) { + seq += 1; + const padded: u32 = @intCast(std.mem.alignForward(u64, @min(fake.slot, stride - off), 16)); + try w.post(seq & 1, seq, padded); + want += padded; + } + } + while (Fake.inflight > 0) try w.reap(); + try std.testing.expectEqual(want, Fake.bytes); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); + // the cap's numbers, not roce_verbs.zig's constants: 1 MiB a write, 2 MiB in flight + try std.testing.expect(Fake.max_write <= 1 << 20); + try std.testing.expect(Fake.max_inflight <= 2 << 20); + // the bound is reached, not avoided: posts outran completions + try std.testing.expect(Fake.max_inflight > 1 << 20); + } +} + +test "Roce.abort (the fail-fast hook) moves both QPs to error: no send goes out and no peer write lands" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + @memset(region, 0); + Fake.region = region; + Fake.slot_bytes = fake.slot; + Fake.reset(1); + fake.Halts.reset(); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + fake.Halts.install(&hw, &w); + var r: Roce = .{ .me = 0, .slot_bytes = fake.slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; + r.abort(); + try std.testing.expect(fake.Halts.halted[0] and fake.Halts.halted[1]); +} diff --git a/zig/src/fabric/tp/roce_shm.zig b/zig/src/fabric/tp/roce_shm.zig new file mode 100644 index 0000000000..3924be3e96 --- /dev/null +++ b/zig/src/fabric/tp/roce_shm.zig @@ -0,0 +1,63 @@ +//! The RoCE protocol's shm wire: both ranks' regions in one shared mapping, posts through the guard's stripes. +const std = @import("std"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const guard = @import("roce_guard.zig"); +const posix = std.posix; +const page = std.heap.page_size_min; + +/// A test's view of each post: slot, padded bytes and the stripes the wire is about to move. +pub const PostTap = struct { + ctx: *anyopaque, + seen: *const fn (ctx: *anyopaque, l: guard.Layout, slot: u32, padded: u32, s: []const guard.Stripe) void, +}; + +/// The peer's region is mapped here; a post copies each stripe, then stores its flag with release order. +pub const ShmWire = struct { + peer: []u8, + layout: guard.Layout, + local: []u8, + /// Stripes a post makes, as HCAs would: the verbs wire's split, flags and bounds on one host. + stripes: u32 = 1, + tap: ?*PostTap = null, + + pub fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) guard.Error!void { + var plan: [guard.max_hcas]guard.Stripe = undefined; + try guard.stripes(w.layout, w.stripes, slot, padded, &plan); + if (w.tap) |t| t.seen(t.ctx, w.layout, slot, padded, plan[0..w.stripes]); + for (plan[0..w.stripes]) |s| { + @memcpy(w.peer[s.remote..][0..s.len], w.local[s.local..][0..s.len]); + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[s.flag..].ptr)); + flag.store(seq, .release); + } + } +}; + +/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). +pub fn sharedPair(b: *Bootstrap, len: usize) ![]align(page) u8 { + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-roce-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(2 * len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + if (b.rank != 0) { + fd = std.c.open(name[0..name_len :0], .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + const m = try posix.mmap(null, 2 * len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + try b.agree("mapped", "RoCE shm wire"); + if (b.rank == 0) _ = std.c.unlink(name[0..name_len :0]); + return m; +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig new file mode 100644 index 0000000000..ec6bc4d65e --- /dev/null +++ b/zig/src/fabric/tp/roce_test.zig @@ -0,0 +1,173 @@ +//! The RoCE protocol end to end on host threads: two ranks, TCP bootstrap, proxy, shm or verbs (TF_TP_ROCE_TEST_HCA). +const std = @import("std"); +const roce = @import("roce.zig"); +const rv = @import("roce_verbs.zig"); +const shm = @import("roce_shm.zig"); +const guard = @import("roce_guard.zig"); +const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); +const reference = @import("reference.zig"); +const Bootstrap = bootstrap.Bootstrap; + +const Run = struct { + port: u16, + wire: roce.WireKind, + devs: [2][]const u8 = .{ "", "" }, + ops: u32, + timeout_ns: u64 = 10 * std.time.ns_per_s, + /// Rank 1 stops after this many ops (rank 0 must then fail its next wait, by timeout or abort). + stop_after: ?u32 = null, + abort_after_ms: ?u64 = null, + errs: [2]?anyerror = .{ null, null }, + done: [2]u32 = .{ 0, 0 }, + max_bytes: usize = 16 * 1024, + /// Variable lengths (`Roce.allGatherV`'s contract): each op's lens[r] bytes of rank r's n, 0 and n included. + var_lens: bool = false, + /// Strides up to four slots through `hostGatherVChunked` (variable lengths, slot-sized ops). + chunked: bool = false, + /// The shm wire's stripes (2: the verbs wire's split over two HCAs). + stripes: u32 = 1, + /// Each rank's shm posts as the tap saw them: how many, and how many left their slot or the region. + taps: [2]Tap = .{ .{}, .{} }, +}; + +/// Counts an shm wire's posts and any stripe outside its slot, its flag words or the region. +const Tap = struct { + posts: u32 = 0, + outside: u32 = 0, + tap: shm.PostTap = undefined, + + fn seen(ctx: *anyopaque, l: guard.Layout, slot: u32, padded: u32, s: []const guard.Stripe) void { + const t: *Tap = @ptrCast(@alignCast(ctx)); + t.posts += 1; + var sum: u64 = 0; + for (s) |x| { + sum += x.len; + const in_slot = x.local >= l.send + slot * l.slot and x.local + x.len <= l.send + (slot + 1) * l.slot; + const in_peer = x.remote >= l.recv + slot * l.slot and x.remote + x.len <= l.recv + (slot + 1) * l.slot; + if (!in_slot or !in_peer or x.remote + x.len > l.len or x.flag < l.flags or x.flag + 4 > l.send) t.outside += 1; + } + if (sum != padded) t.outside += 1; + } +}; + +fn rank(run: *Run, me: u32, lfd: sock.Fd) void { + run.errs[me] = body(run, me, lfd); +} + +fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { + var b = (if (me == 0) Bootstrap.acceptOn(lfd, 2, 5 * std.time.ns_per_s) else Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = run.port, .connect_timeout_ns = 5 * std.time.ns_per_s })) catch |e| return e; + defer b.close(); + var spec: [1]rv.HcaSpec = .{.{}}; + if (run.wire == .verbs) _ = rv.parseSpecs(run.devs[me], @ptrCast(&spec)) catch |e| return e; + const mine = &run.taps[me]; + mine.tap = .{ .ctx = mine, .seen = Tap.seen }; + const r = roce.Roce.init(null, &b, "", .{ .wire = run.wire, .hcas = if (run.wire == .verbs) &spec else &.{}, .max_bytes = run.max_bytes, .timeout_ns = run.timeout_ns, .shm_stripes = run.stripes, .tap = &mine.tap }) catch |e| return e; + defer r.deinit(); + var abort_thread: ?std.Thread = null; + defer if (abort_thread) |t| t.join(); + var in: [64 * 1024]u8 = undefined; + var peer: [64 * 1024]u8 = undefined; + var out: [128 * 1024]u8 = undefined; + for (0..run.ops) |k| { + if (run.stop_after) |s| if (k == s) { + // rank 1 skips op s; rank 0 waits in it until its timeout, or until the abort armed now + if (me == 1) { + sock.sleepNs(300 * std.time.ns_per_ms); + return null; + } + if (run.abort_after_ms) |ms| abort_thread = std.Thread.spawn(.{}, abortLater, .{ r, ms }) catch null; + }; + // sizes from 1 byte to the slot, odd ones included, both slots and wrap-arounds + const n = 1 + (k * 7919) % (if (run.chunked) 4 * run.max_bytes else run.max_bytes); + reference.fill(in[0..n], .u8, me, k); + reference.fill(peer[0..n], .u8, 1 - me, k); + const lens: [2]u32 = if (k % 5 == 0) .{ 0, @intCast(n) } else .{ @intCast((k * 31) % (n + 1)), @intCast((k * 17 + 5) % (n + 1)) }; + const l: [2]usize = if (run.var_lens or run.chunked) .{ lens[0], lens[1] } else .{ n, n }; + if (run.chunked) { + r.hostGatherVChunked(in[0..n], out[0 .. 2 * n], lens) catch |e| return e; + } else r.hostGatherV(in[0..n], out[0 .. 2 * n], if (run.var_lens) lens else null) catch |e| return e; + const r0 = if (me == 0) in[0..n] else peer[0..n]; + const r1 = if (me == 1) in[0..n] else peer[0..n]; + if (!std.mem.eql(u8, out[0..l[0]], r0[0..l[0]]) or !std.mem.eql(u8, out[n..][0..l[1]], r1[0..l[1]])) return error.WrongBytes; + run.done[me] += 1; + } + r.check() catch |e| return e; + // both ranks finish before either tears down (the peer may still read a flag this rank's wire writes) + b.agree("done", "test end") catch |e| return e; + return null; +} + +fn abortLater(r: *roce.Roce, ms: u64) void { + sock.sleepNs(ms * std.time.ns_per_ms); + r.abort(); +} + +fn twoRanks(run: *Run) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + run.port = l.port; + const t1 = try std.Thread.spawn(.{}, rank, .{ run, 1, l.fd }); + rank(run, 0, l.fd); + t1.join(); +} + +test "the RoCE protocol over the shm wire: 2,000 gathers of every size, bit for bit" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "the RoCE protocol with variable lengths over the shm wire: each rank posts its own padded count, empty shards included" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "the shm wire striped as over two HCAs: every post inside its slot and the region, bytes and lengths exact" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true, .stripes = 2 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); + for (run.taps) |t| { + try std.testing.expectEqual(@as(u32, 2000), t.posts); + try std.testing.expectEqual(@as(u32, 0), t.outside); + } +} + +test "chunked gathers over the shm wire: strides up to four slots in slot-sized ops, tails empty on both ranks skipped" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .chunked = true, .stripes = 2 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); + for (run.taps) |t| { + try std.testing.expect(t.posts > 2000); + try std.testing.expectEqual(@as(u32, 0), t.outside); + } +} + +test "a peer that stops makes the waiting rank time out, and abort ends a wait at once" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .timeout_ns = 100 * std.time.ns_per_ms }; + try twoRanks(&run); + try std.testing.expectEqual(@as(?anyerror, error.Timeout), run.errs[0]); + try std.testing.expectEqual(@as(u32, 20), run.done[0]); + var run2: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .abort_after_ms = 50 }; + const kept = roce.Roce.kept.load(.monotonic); + try twoRanks(&run2); + // the aborted rank keeps its region until exit, the other frees it + try std.testing.expectEqual(kept + 1, roce.Roce.kept.load(.monotonic)); + try std.testing.expectEqual(@as(?anyerror, error.Aborted), run2.errs[0]); + try std.testing.expectEqual(@as(?anyerror, null), run2.errs[1]); +} + +test "the RoCE protocol over verbs (TF_TP_ROCE_TEST_HCA=dev0,dev1: soft-RoCE or real HCAs)" { + const raw = std.c.getenv("TF_TP_ROCE_TEST_HCA") orelse return error.SkipZigTest; + var it = std.mem.splitScalar(u8, std.mem.span(raw), ','); + var run: Run = .{ .port = 0, .wire = .verbs, .ops = 2000 }; + run.devs = .{ it.next() orelse return error.SkipZigTest, it.next() orelse return error.SkipZigTest }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig new file mode 100644 index 0000000000..6917d292be --- /dev/null +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -0,0 +1,455 @@ +//! The RoCE wire over verbs: one RC queue pair a HCA to the peer, a send slot striped over the HCAs with seq flags. +const std = @import("std"); +const log = @import("log.zig"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const Verbs = verbs.Verbs; +const guard = @import("roce_guard.zig"); +const sock = @import("sock.zig"); + +pub const max_hcas = guard.max_hcas; +/// Work requests a queue pair holds; at most a quarter outstanding (two requests a stripe). +const send_depth = 256; +const link_ethernet = 2; +const record_magic: u32 = 0x5446_5244 + 2; // "TFRD", bumped for each change of this record's layout + +pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord } || guard.Error; + +/// Bytes posted and not yet completed, every HCA together: a post waits for completions past it; never raised. +pub const max_inflight: u64 = 2 * @as(u64, guard.max_write); +/// A post that waits on completions this long fails the proxy (a peer or HCA that stopped answering). +const inflight_wait_ns: u64 = 10 * std.time.ns_per_s; +const ring_len = send_depth / 4; +/// The flag post's bytes: a u32 on the stack under lkey 0, legal only while the QP sends it inline. +pub const flag_bytes: u32 = @sizeOf(u32); +/// The longest `quiesce` waits for the HCAs to flush what was posted. +pub const quiesce_ns: u64 = 5 * std.time.ns_per_s; + +/// Before ibv_reg_mr: no fork shares or copies the region (DONTFORK), its pages stay 4 KiB (NOHUGEPAGE); else no start. +pub fn prepareRegion(region: []align(std.heap.page_size_min) u8) Error!void { + if (comptime @import("builtin").os.tag != .linux) return; // the verbs wire runs on Linux only + const advice = [_]u32{ std.os.linux.MADV.DONTFORK, std.os.linux.MADV.NOHUGEPAGE }; + for (advice) |a| std.posix.madvise(region.ptr, region.len, a) catch |e| { + log.warn("tp roce: madvise({d}) on the {d} B region failed: {t}; refusing to start", .{ a, region.len, e }); + return error.VerbsFailed; + }; +} + +/// Whether a QP's actual inline capacity takes the flag post. +pub fn inlineFits(cap: abi.QpCap) bool { + return cap.max_inline_data >= flag_bytes; +} + +/// One HCA as configured: its device name and RoCE v2 GID index (null: detected). +pub const HcaSpec = struct { name: [64]u8 = @splat(0), name_len: u8 = 0, gid: ?u8 = null }; + +/// What a rank tells its peer to connect (both ranks build it the same way; sent through the bootstrap). +pub const Record = extern struct { + magic: u32 = record_magic, + n_hca: u32 = 0, + region: u64 = 0, + /// The region's slot and registered length: a peer whose layout differs is refused at connect. + slot_bytes: u64 = 0, + region_len: u64 = 0, + rkey: [max_hcas]u32 = @splat(0), + lid: [max_hcas]u32 = @splat(0), + mtu: [max_hcas]u32 = @splat(0), + qpn: [max_hcas]u32 = @splat(0), + gid: [max_hcas][16]u8 = @splat(@splat(0)), +}; + +const Hca = struct { + ctx: *abi.Context, + pd: *abi.Pd, + mr: *abi.Mr, + cq: *abi.Cq, + qp: *abi.Qp, + gid_index: u8, + gid: [16]u8, + mtu: u32, + lid: u16, + outstanding: u32 = 0, + /// Each outstanding post's stripe bytes, oldest at `head` (completions arrive in order). + ring: [ring_len]u32 = undefined, + head: u32 = 0, + inflight: u64 = 0, + /// Posts made and their completions reaped (the proxy writes them, others read): equal = nothing in flight. + posts: std.atomic.Value(u64) = .init(0), + done: std.atomic.Value(u64) = .init(0), +}; + +/// Parses `name[:gid],name[:gid]` (TF_TP_ROCE_HCA). +pub fn parseSpecs(text: []const u8, out: *[max_hcas]HcaSpec) !usize { + var n: usize = 0; + var it = std.mem.tokenizeScalar(u8, text, ','); + while (it.next()) |item| { + if (n == max_hcas) return error.TooManyHcas; + var parts = std.mem.splitScalar(u8, std.mem.trim(u8, item, " "), ':'); + const name = parts.next().?; + if (name.len == 0 or name.len >= 64) return error.BadHcaName; + out[n] = .{ .name_len = @intCast(name.len) }; + @memcpy(out[n].name[0..name.len], name); + if (parts.next()) |g| out[n].gid = try std.fmt.parseInt(u8, g, 10); + n += 1; + } + return n; +} + +/// RoCE v2 by sysfs (the kernel's own word for each GID); null when sysfs cannot say. +fn sysfsRoceV2(name: []const u8, index: u32) ?bool { + var path: [256]u8 = undefined; + const p = std.fmt.bufPrintSentinel(&path, "/sys/class/infiniband/{s}/ports/1/gid_attrs/types/{d}", .{ name, index }, 0) catch return null; + const fd = std.c.open(p, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return null; + defer _ = std.c.close(fd); + var buf: [32]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return false; + return std.mem.startsWith(u8, buf[0..@intCast(n)], "RoCE v2"); +} + +/// An IPv4-mapped GID (::ffff:a.b.c.d): the address the HCA's netdev carries. +pub fn ipv4Mapped(g: [16]u8) bool { + return std.mem.allEqual(u8, g[0..10], 0) and g[10] == 0xff and g[11] == 0xff; +} + +/// The lowest RoCE v2 GID index that is IPv4-mapped (GID indices move on reboot: never hard-coded). +fn detectGid(v: *const Verbs, ctx: *abi.Context, name: []const u8, table: c_int) ?u8 { + var i: u32 = 0; + while (i < @min(@as(u32, @intCast(@max(table, 0))), 255)) : (i += 1) { + const g = v.gid(ctx, 1, i) catch continue; + if (!ipv4Mapped(g.raw)) continue; + const v2 = sysfsRoceV2(name, i) orelse (if (v.gidType(ctx, 1, i)) |t| t == 2 else true); + if (v2) return @intCast(i); + } + return null; +} + +/// A device's port 1 as `roce_pick.choose` sees it: up, Ethernet, and its IPv4 RoCE v2 GID (index and address). +pub const Port = struct { spec: HcaSpec = .{}, active: bool = false, ethernet: bool = false, ipv4: ?[4]u8 = null }; + +/// Port 1 of every device that opens, in the device list's order, at most `out.len`. +pub fn scanPorts(v: *const Verbs, out: []Port) usize { + var count: c_int = 0; + const list = v.api.ibv_get_device_list(&count) orelse return 0; + defer v.api.ibv_free_device_list(list); + var n: usize = 0; + var i: usize = 0; + while (list[i]) |dev| : (i += 1) { + if (n == out.len) break; + const name = std.mem.span(v.api.ibv_get_device_name(dev) orelse continue); + const ctx = v.api.ibv_open_device(dev) orelse continue; + defer _ = v.api.ibv_close_device(ctx); + const port = v.port(ctx, 1) catch continue; + var p: Port = .{ .active = port.state == abi.port_active, .ethernet = port.link_layer == link_ethernet }; + p.spec.name_len = @intCast(@min(name.len, 63)); + @memcpy(p.spec.name[0..p.spec.name_len], name[0..p.spec.name_len]); + if (p.active and p.ethernet) if (detectGid(v, ctx, name, port.gid_tbl_len)) |gi| { + if (v.gid(ctx, 1, gi)) |g| { + p.spec.gid = gi; + p.ipv4 = g.raw[12..16].*; + } else |_| {} + }; + out[n] = p; + n += 1; + } + return n; +} + +pub const VerbsWire = struct { + v: Verbs, + hcas: [max_hcas]Hca = undefined, + n_hca: u32 = 0, + traffic_class: u8, + peer_region: u64 = 0, + peer_rkey: [max_hcas]u32 = @splat(0), + /// Offsets in both ranks' regions (the same layout): receive slots and flags. + recv_off: usize, + flag_off: usize, + send_off: usize, + slot_bytes: usize, + region: []u8, + posted: u64 = 0, + /// The most stripe bytes in flight at once (posted, not completed), for tests and reports. + peak: std.atomic.Value(u64) = .init(0), + completed: u64 = 0, + /// Completions `quiesce` reaped, and those of them flushed (not sent: the QP was in error). + reaped: u64 = 0, + flushed: u64 = 0, + + /// Opens the HCAs, registers `region` on each and creates the queue pairs (INIT); `record` then goes to the peer. + pub fn open(specs: []const HcaSpec, region: []u8, layout: struct { send: usize, recv: usize, flags: usize, slot: usize }, traffic_class: u8) !VerbsWire { + if (specs.len == 0 or specs.len > max_hcas) return error.NoHca; + var w: VerbsWire = .{ .v = try Verbs.open(null), .traffic_class = traffic_class, .recv_off = layout.recv, .flag_off = layout.flags, .send_off = layout.send, .slot_bytes = layout.slot, .region = region }; + errdefer w.close(); + for (specs) |s| { + try w.openHca(s); + } + return w; + } + + fn openHca(w: *VerbsWire, s: HcaSpec) !void { + const api = w.v.api; + const name = s.name[0..s.name_len]; + const ctx = try w.v.openDevice(name); + errdefer _ = api.ibv_close_device(ctx); + const port = try w.v.port(ctx, 1); + if (port.state != abi.port_active) { + log.warn("tp roce: {s} port 1 is not active", .{name}); + return error.NoHca; + } + const gi = s.gid orelse detectGid(&w.v, ctx, name, port.gid_tbl_len) orelse { + log.warn("tp roce: {s} has no IPv4 RoCE v2 GID", .{name}); + return error.NoHca; + }; + const gid = try w.v.gid(ctx, 1, gi); + const pd = api.ibv_alloc_pd(ctx) orelse return error.VerbsFailed; + errdefer _ = api.ibv_dealloc_pd(pd); + const mr = api.ibv_reg_mr(pd, w.region.ptr, w.region.len, abi.access_local_write | abi.access_remote_write) orelse { + log.warn("tp roce: ibv_reg_mr on {s} failed (memlock limit? ulimit -l unlimited)", .{name}); + return error.VerbsFailed; + }; + errdefer _ = api.ibv_dereg_mr(mr); + const cq = api.ibv_create_cq(ctx, send_depth, null, null, 0) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_cq(cq); + var init: abi.QpInitAttr = .{ .send_cq = cq, .recv_cq = cq, .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 }, .qp_type = abi.qpt_rc }; + const qp = api.ibv_create_qp(pd, &init) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_qp(qp); + if (!inlineFits(init.cap)) { + log.warn("tp roce: {s}'s QP sends {d} B inline, the flag post needs {d}; refusing to start", .{ name, init.cap.max_inline_data, flag_bytes }); + return error.VerbsFailed; + } + var a: abi.QpAttr = .{ .qp_state = abi.qps_init, .port_num = 1, .qp_access_flags = abi.access_remote_write }; + if (api.ibv_modify_qp(qp, &a, abi.mask.state | abi.mask.pkey_index | abi.mask.port | abi.mask.access_flags) != 0) return error.VerbsFailed; + w.hcas[w.n_hca] = .{ .ctx = ctx, .pd = pd, .mr = mr, .cq = cq, .qp = qp, .gid_index = gi, .gid = gid.raw, .mtu = @intCast(port.active_mtu), .lid = port.lid }; + w.n_hca += 1; + log.info("[tensorfold] tp roce: {s} GID {d} ({d}.{d}.{d}.{d}), MTU enum {d}", .{ name, gi, gid.raw[12], gid.raw[13], gid.raw[14], gid.raw[15], port.active_mtu }); + } + + pub fn close(w: *VerbsWire) void { + const api = w.v.api; + for (w.hcas[0..w.n_hca]) |h| { + _ = api.ibv_destroy_qp(h.qp); + _ = api.ibv_destroy_cq(h.cq); + _ = api.ibv_dereg_mr(h.mr); + _ = api.ibv_dealloc_pd(h.pd); + _ = api.ibv_close_device(h.ctx); + } + w.n_hca = 0; + w.v.close(); + } + + /// The region and each HCA's keys, this side's and the peer's (a protection error is checked against them). + pub fn describe(w: *const VerbsWire, buf: []u8) []const u8 { + var out: std.Io.Writer = .fixed(buf); + out.print("region 0x{x} len {d}", .{ @intFromPtr(w.region.ptr), w.region.len }) catch return out.buffered(); + for (w.hcas[0..w.n_hca], 0..) |h, i| { + out.print("; HCA {d} lkey 0x{x} rkey 0x{x}, peer region 0x{x} rkey 0x{x}", .{ i, h.mr.lkey, h.mr.rkey, w.peer_region, w.peer_rkey[i] }) catch break; + } + return out.buffered(); + } + + /// The QPs to error: the HCAs stop sending and refuse the peer's writes, and what was posted completes as flushed. + pub fn halt(w: *VerbsWire) bool { + var ok = true; + for (w.hcas[0..w.n_hca]) |h| { + var a: abi.QpAttr = .{ .qp_state = abi.qps_err }; + if (w.v.api.ibv_modify_qp(h.qp, &a, abi.mask.state) != 0) ok = false; + } + return ok; + } + + /// `halt`, then every posted request's completion reaped, flushed or not: no request of the region left at an HCA. + pub fn quiesce(w: *VerbsWire, timeout_ns: u64) bool { + if (!w.halt()) return false; + const t0 = sock.nowNs(); + for (w.hcas[0..w.n_hca]) |*h| while (h.outstanding > 0) { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return false; + for (wc[0..@intCast(n)]) |c| w.flushed += @intFromBool(c.status != abi.wc_success); + w.reaped += @intCast(n); + if (n > 0) w.popped(h, @intCast(n)) else if (sock.nowNs() - t0 > timeout_ns) return false; + }; + return true; + } + + /// Every HCA's posts completed: no request of the region at an HCA (readable beside the proxy). + pub fn idle(w: *const VerbsWire) bool { + for (w.hcas[0..w.n_hca]) |*h| if (h.posts.load(.acquire) != h.done.load(.acquire)) return false; + return true; + } + + /// Each QP's state as its device reports it (ibv_query_qp), then what `quiesce` reaped and what is still posted. + pub fn report(w: *const VerbsWire, buf: []u8) []const u8 { + const names = [_][]const u8{ "RESET", "INIT", "RTR", "RTS", "SQD", "SQE", "ERR" }; + var out: std.Io.Writer = .fixed(buf); + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + var a: abi.QpAttr = .{}; + var init: abi.QpInitAttr = undefined; + const q = w.v.query_qp orelse { + out.print("HCA {d} QP state unknown (no ibv_query_qp); ", .{i}) catch break; + continue; + }; + if (q(h.qp, &a, abi.mask.state, &init) != 0) { + out.print("HCA {d} QP query failed; ", .{i}) catch break; + } else if (a.qp_state >= 0 and a.qp_state < names.len) { + out.print("HCA {d} QP {s}; ", .{ i, names[@intCast(a.qp_state)] }) catch break; + } else out.print("HCA {d} QP state {d}; ", .{ i, a.qp_state }) catch break; + out.print("posted {d}, completed {d}; ", .{ h.posts.load(.acquire), h.done.load(.acquire) }) catch break; + } + var held: u64 = 0; + for (w.hcas[0..w.n_hca]) |h| held += h.outstanding; + out.print("reaped {d} (flushed {d}), still posted {d}", .{ w.reaped, w.flushed, held }) catch {}; + return out.buffered(); + } + + pub fn record(w: *const VerbsWire) Record { + var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr), .slot_bytes = w.slot_bytes, .region_len = w.region.len }; + for (w.hcas[0..w.n_hca], 0..) |h, i| { + r.rkey[i] = h.mr.rkey; + r.lid[i] = h.lid; + r.mtu[i] = h.mtu; + r.qpn[i] = h.qp.qp_num; + r.gid[i] = h.gid; + } + return r; + } + + /// The peer HCA for each of ours: the same IPv4 /24 (functions pair by subnet, names may differ), else position. + pub fn pairing(mine: Record, peer: Record) ![max_hcas]u32 { + if (peer.magic != record_magic) return error.BadRecord; + if (peer.n_hca != mine.n_hca or peer.slot_bytes != mine.slot_bytes or peer.region_len != mine.region_len) return error.BadRecord; + var out: [max_hcas]u32 = undefined; + for (0..mine.n_hca) |i| { + out[i] = @intCast(i); + var hits: u32 = 0; + for (0..peer.n_hca) |j| if (std.mem.eql(u8, mine.gid[i][12..15], peer.gid[j][12..15])) { + out[i] = @intCast(j); + hits += 1; + }; + if (hits > 1) out[i] = @intCast(i); // one subnet for both: keep the order + } + return out; + } + + /// INIT to RTR to RTS against the peer's record. + pub fn connect(w: *VerbsWire, peer: Record) !void { + const pairs = try pairing(w.record(), peer); + const api = w.v.api; + const m = abi.mask; + w.peer_region = peer.region; + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + const j = pairs[i]; + w.peer_rkey[i] = peer.rkey[j]; + var rtr: abi.QpAttr = .{ .qp_state = abi.qps_rtr, .path_mtu = @intCast(@min(peer.mtu[j], h.mtu)), .dest_qp_num = peer.qpn[j], .rq_psn = 0, .max_dest_rd_atomic = 1, .min_rnr_timer = 12 }; + rtr.ah_attr = .{ .grh = .{ .dgid = .{ .raw = peer.gid[j] }, .sgid_index = h.gid_index, .hop_limit = 64, .traffic_class = w.traffic_class }, .dlid = @intCast(peer.lid[j]), .is_global = 1, .port_num = 1 }; + if (api.ibv_modify_qp(h.qp, &rtr, m.state | m.av | m.path_mtu | m.dest_qpn | m.rq_psn | m.max_dest_rd_atomic | m.min_rnr_timer) != 0) return error.VerbsFailed; + var rts: abi.QpAttr = .{ .qp_state = abi.qps_rts, .timeout = 14, .retry_cnt = 7, .rnr_retry = 7, .sq_psn = 0, .max_rd_atomic = 1 }; + if (api.ibv_modify_qp(h.qp, &rts, m.state | m.timeout | m.retry_cnt | m.rnr_retry | m.sq_psn | m.max_qp_rd_atomic) != 0) return error.VerbsFailed; + } + } + + /// Slot `slot`'s first `padded` bytes to the peer, striped over the HCAs, each stripe's `seq` flag on the same QP. + pub fn post(w: *VerbsWire, slot: u32, seq: u32, padded: u32) !void { + var plan: [max_hcas]guard.Stripe = undefined; + try guard.stripes(w.regionLayout(), w.n_hca, slot, padded, &plan); + while (w.inflight() + padded > max_inflight) try w.reapWait(); + var seq_copy = seq; + for (w.hcas[0..w.n_hca], plan[0..w.n_hca], 0..) |*h, s, i| { + while (h.outstanding >= ring_len) try w.reapHca(h); + var flag_sge: abi.Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; + var flag_wr: abi.SendWr = .{ .wr_id = seq, .sg_list = @ptrCast(&flag_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = abi.send_signaled | abi.send_inline, .remote_addr = w.peer_region + s.flag, .rkey = w.peer_rkey[i] }; + var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + s.local, .length = s.len, .lkey = h.mr.lkey }; + var data_wr: abi.SendWr = .{ .wr_id = seq, .next = &flag_wr, .sg_list = @ptrCast(&data_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = 0, .remote_addr = w.peer_region + s.remote, .rkey = w.peer_rkey[i] }; + var bad: ?*abi.SendWr = null; + const first: *abi.SendWr = if (s.len > 0) &data_wr else &flag_wr; + if (h.ctx.ops.post_send(h.qp, first, &bad) != 0) return error.PostFailed; + h.ring[(h.head + h.outstanding) % ring_len] = s.len; + h.outstanding += 1; + h.inflight += s.len; + _ = h.posts.fetchAdd(1, .release); + } + if (w.inflight() > w.peak.load(.monotonic)) w.peak.store(w.inflight(), .monotonic); + w.posted += 1; + } + + /// This rank's layout; the MR covers `region` exactly and the peer's record matched it at connect. + pub fn regionLayout(w: *const VerbsWire) guard.Layout { + return .{ .send = w.send_off, .recv = w.recv_off, .flags = w.flag_off, .slot = w.slot_bytes, .len = w.region.len }; + } + + fn reapHca(w: *VerbsWire, h: *Hca) !void { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return error.Completion; + for (wc[0..@intCast(n)]) |c| if (c.status != abi.wc_success) { + log.err("tp roce: RDMA write of seq {d} failed: {s} (vendor 0x{x})", .{ c.wr_id, w.v.statusText(c.status), c.vendor_err }); + return error.Completion; + }; + w.popped(h, @intCast(n)); + } + + fn popped(w: *VerbsWire, h: *Hca, n: u32) void { + for (0..n) |_| { + h.inflight -= h.ring[h.head]; + h.head = (h.head + 1) % ring_len; + } + h.outstanding -= n; + w.completed += n; + _ = h.done.fetchAdd(n, .release); + } + + /// Bytes posted and not yet completed over every HCA. + pub fn inflight(w: *const VerbsWire) u64 { + var sum: u64 = 0; + for (w.hcas[0..w.n_hca]) |h| sum += h.inflight; + return sum; + } + + /// Reaps until some post completes (error.Completion when none does within `inflight_wait_ns`). + fn reapWait(w: *VerbsWire) !void { + const before = w.completed; + const t0 = sock.nowNs(); + while (w.completed == before) { + try w.reap(); + if (w.completed == before and sock.nowNs() - t0 > inflight_wait_ns) return error.Completion; + std.atomic.spinLoopHint(); + } + } + + /// Reap every HCA's completions (after each op, and in the proxy's idle loop). + pub fn reap(w: *VerbsWire) !void { + for (w.hcas[0..w.n_hca]) |*h| try w.reapHca(h); + } +}; + +test "HCA lists parse and subnets pair the functions" { + var specs: [max_hcas]HcaSpec = undefined; + const n = try parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqual(@as(usize, 2), n); + try std.testing.expectEqualStrings("rocep1s0f1", specs[0].name[0..specs[0].name_len]); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); + try std.testing.expectEqual(@as(?u8, null), specs[1].gid); + var a: Record = .{ .n_hca = 2 }; + var b: Record = .{ .n_hca = 2 }; + const sub100 = [16]u8{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, 192, 168, 100, 10 }; + var sub101 = sub100; + sub101[14] = 101; + a.gid = .{ sub100, sub101 }; + b.gid = .{ sub101, sub100 }; // the peer lists its functions the other way round + b.gid[0][15] = 11; + b.gid[1][15] = 11; + const p = try VerbsWire.pairing(a, b); + try std.testing.expectEqual(@as(u32, 1), p[0]); + try std.testing.expectEqual(@as(u32, 0), p[1]); + try std.testing.expect(ipv4Mapped(sub100)); + b.n_hca = 1; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); + b.n_hca = 2; + b.slot_bytes = 1 << 20; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); + b.slot_bytes = 0; + b.region_len = 4096; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); +} diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh new file mode 100644 index 0000000000..75f0c7900b --- /dev/null +++ b/zig/src/fabric/tp/run_hosts.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# One rank of the two-host TP=2 gate over RoCE: `run_hosts.sh RANK RANK0_ADDR ./tf-tp-test OUTDIR` on both hosts. +set -u +RANK=${1:?rank}; MASTER=${2:?master address}; BIN=${3:?tf-tp-test}; OUT=${4:?outdir} +mkdir -p "$OUT" +[ -n "${NCCL_SOCKET_IFNAME:-}" ] && export NCCL_SOCKET_IFNAME # else NCCL picks the interface +# NCCL uses the RoCE wire's devices unless NCCL_IB_HCA is set: the names without `:gid` (NCCL reads `:N` as a port) +[ -n "${TF_TP_ROCE_HCA:-}" ] && [ -z "${NCCL_IB_HCA:-}" ] && NCCL_IB_HCA=$(echo "$TF_TP_ROCE_HCA" | sed 's/:[0-9]*//g; s/ //g') +[ -n "${TF_TP_ROCE_HCA:-}" ] && export TF_TP_ROCE_HCA +"$BIN" kernel > /dev/null 2>&1 || { echo "error: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } +ulimit -l unlimited 2> /dev/null || echo "warning: memlock is $(ulimit -l) KiB: ibv_reg_mr of the RoCE region may fail" >&2 +# neither named: NCCL takes the ports the RoCE picker chose with the peer (`tf-tp-test hca`, both ranks), not its own +if [ -z "${TF_TP_ROCE_HCA:-}${NCCL_IB_HCA:-}" ]; then + # hca gives up after TF_TP_PLAN_TIMEOUT (default 300 s) with its reason; the cap here only catches a hang past that + CAP=$(awk 'BEGIN { t = ARGV[1] + 0; if (t <= 0) t = 300; printf "%d", int(t) + 31 }' "${TF_TP_PLAN_TIMEOUT:-300}") + echo "hca: rank $RANK picking the ports with the peer (up to $CAP s)" + PICK=$(TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=29700 timeout "$CAP" "$BIN" hca 2> "$OUT/hca.log") + rc=$? + if [ "$rc" -eq 124 ]; then + echo "error: $BIN hca still running after $CAP s: killed ($OUT/hca.log)" >&2 + exit 3 + fi + if [ "$rc" -ne 0 ]; then + echo "error: $BIN hca exited $rc: $(grep -h '^FAIL\|^error' "$OUT/hca.log" | tail -n 1 | grep . || tail -n 1 "$OUT/hca.log") ($OUT/hca.log)" >&2 + exit 3 + fi + if ! echo "$PICK" | grep -Eqx 'NCCL_IB_HCA=[A-Za-z0-9_.-]+(,[A-Za-z0-9_.-]+)*'; then + echo "error: $BIN hca printed '$PICK', not NCCL_IB_HCA=name,... ($OUT/hca.log)" >&2 + exit 3 + fi + NCCL_IB_HCA=${PICK#NCCL_IB_HCA=} + echo "hca: rank $RANK NCCL_IB_HCA=$NCCL_IB_HCA ($(grep -h '^INFO hca' "$OUT/hca.log"))" +fi +[ -n "${NCCL_IB_HCA:-}" ] && export NCCL_IB_HCA +{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; env | grep -E '^(TF_TP_ROCE_HCA|NCCL_IB_HCA|NCCL_SOCKET_IFNAME)='; } > "$OUT/host.txt" 2>&1 +PORT=29700 +step() { # NAME BACKEND ARGS...: this host's rank; log OUT/NAME.log ending with an EXIT line + local name=$1 backend=$2; shift 2 + PORT=$((PORT + 10)) + ( TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend TF_TP_ROCE_FALLBACK=error \ + timeout 900 "$BIN" "$@"; echo "EXIT rank $RANK rc $? at $(date +%s.%N)" ) > "$OUT/$name.log" 2>&1 + echo "$name: $(grep -h '^SUMMARY\|^PASS roce-host\|^INJECT\|^EXIT' "$OUT/$name.log" | tr '\n' ' ')" +} +step roce-host roce roce-host +step suite-roce roce suite +step suite-nccl nccl suite +for backend in roce nccl; do + for mode in fatal kill; do + for victim in 0 1; do step "ff-$backend-$mode-$victim" "$backend" failfast "$victim" "$mode"; done + done +done diff --git a/zig/src/fabric/tp/run_pair.sh b/zig/src/fabric/tp/run_pair.sh new file mode 100644 index 0000000000..1cf73c6261 --- /dev/null +++ b/zig/src/fabric/tp/run_pair.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# Two ranks of tf-tp-test on one two-GPU host: `[SUITES=..] [FAILFAST=..] run_pair.sh BIN OUTDIR`. +set -u +BIN=${1:?bin}; OUT=${2:?outdir} +mkdir -p "$OUT" +nvidia-smi > "$OUT/smi.txt" 2>&1 +nvidia-smi topo -m > "$OUT/topo.txt" 2>&1 +ldconfig -p | grep -E "libnccl|libcuda" > "$OUT/libs.txt" 2>&1 +PORT=29500 +pair() { # NAME BACKEND ARGS...: both ranks, logs OUT/NAME.rR.log ending with an EXIT line + local name=$1 backend=$2 r; shift 2 + PORT=$((PORT + 10)) + for r in 0 1; do + ( TF_TP_RANK=$r TF_TP_DEVICE=$r TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend NCCL_DEBUG=${NCCL_DEBUG:-WARN} \ + timeout 600 "$BIN" "$@"; echo "EXIT rank $r rc $? at $(date +%s.%N)" ) > "$OUT/$name.r$r.log" 2>&1 & + done + wait +} +# ffsum OUT NAME VICTIM: a fail-fast pair's verdict line (exit codes, exit times after the injection) +ffsum() { + [ -f "$1/$2.r0.log" ] && [ -f "$1/$2.r1.log" ] || { echo "ffsum: $1/$2.r0.log or .r1.log missing" >&2; return 1; } + awk -v name="$2" -v victim="$3" ' + FNR == 1 { r = (FILENAME ~ /\.r1\.log$/) ? 1 : 0 } + !(r in ex) && match($0, /EXIT rank [0-9] rc [0-9]+ at [0-9.]+/) { + split(substr($0, RSTART, RLENGTH), f, " "); rc[r] = f[5]; at[r] = f[7]; ex[r] = 1 + } + r == victim && !inj && match($0, /INJECT [A-Za-z0-9_]+ rank [0-9]+ at [0-9.]+/) { + split(substr($0, RSTART, RLENGTH), g, " "); t0 = g[6]; inj = 1 + } + END { + if (!inj || !(0 in ex) || !(1 in ex)) { print name ": FAIL (no inject / exit line)"; exit } + for (i = 0; i < 2; i++) { want[i] = (i == victim && name ~ /kill/) ? 137 : 70; dt[i] = at[i] - t0 } + ok = rc[0] == want[0] && rc[1] == want[1] && dt[1 - victim] < 3.0 + printf "%s: %s rc [%d, %d] (want [%d, %d]), exit after injection: rank0 %.3f s, rank1 %.3f s\n", name, + ok ? "PASS" : "FAIL", rc[0], rc[1], want[0], want[1], dt[0], dt[1] + }' "$1/$2.r0.log" "$1/$2.r1.log" +} +for backend in ${SUITES:-mailbox nccl}; do + pair "suite-$backend" "$backend" suite +done +: > "$OUT/failfast.txt" +for backend in ${FAILFAST:-nccl mailbox}; do + for mode in fatal kill; do + for victim in 0 1; do + name=ff-$backend-$mode-$victim + pair "$name" "$backend" failfast "$victim" "$mode" + ffsum "$OUT" "$name" "$victim" >> "$OUT/failfast.txt" + done + done +done +grep -h "^SUMMARY\|^FAIL" "$OUT"/suite*.log +cat "$OUT/failfast.txt" diff --git a/zig/src/fabric/tp/run_transport.sh b/zig/src/fabric/tp/run_transport.sh new file mode 100755 index 0000000000..68ebcddd9a --- /dev/null +++ b/zig/src/fabric/tp/run_transport.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# The failed-transport check on two hosts (tf-tp-test failfast V transport); `run_transport.sh` alone prints how. +set -u +usage() { + cat << 'TXT' +run_transport.sh MASTER_ADDR PEER_SSH OUTDIR (on rank 0's host, as root; rank 1 runs over ssh) + Two hosts with this script and tf-tp-test (built with the mailbox kernel) at the same path on both. Before the steps, + on both hosts: no GPU process and memlock unlimited. Steps: T0 (suite), T1 (victim 1), T3 (victim 0), fail-fast on; + tf-tp-test injects only when both links are idle. After each step, on both hosts: the kernel log since the step + began and every HCA *err* counter against the step's start. A matching line, a counter that grew, an unexpected rc + or a failure line prints STOP: and ends the run (exit 1), with no later step. +run_transport.sh t2 OUTDIR (one host: TF_TP_ALLOW_HOST_ONLY=1) + T2 (fail-fast off, a hold) runs both ranks on this host over the shm wire; tf-tp-test refuses it on an RDMA wire. +TXT +} +[ $# -ge 2 ] || { usage; exit 2; } +KIT=$(cd "$(dirname "$0")" && pwd) +BIN=$KIT/tf-tp-test +KLOG='EXT4-fs|smmu|iommu|mlx5|AER|nvme|Xid' +BAD='RDMA write of seq|local protection|NOT drained|GPU sync FAILED|REFUSED|refused ERR|panic' +"$BIN" kernel > /dev/null 2>&1 || { echo "STOP: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } + +if [ "${1:-}" = t2 ]; then + OUT=${2:?outdir} + [ "${TF_TP_ALLOW_HOST_ONLY:-}" = 1 ] || { echo "STOP: T2 is host-only: set TF_TP_ALLOW_HOST_ONLY=1" >&2; exit 2; } + mkdir -p "$OUT" + for r in 0 1; do + TF_TP_RANK=$r TF_TP_MASTER=127.0.0.1 TF_TP_PORT=29990 TF_COMM_BACKEND=roce TF_TP_ROCE_WIRE=shm TF_TP_FAILFAST=0 \ + timeout 300 "$BIN" failfast 1 transport 200 5 > "$OUT/T2-r$r.log" 2>&1 & + done + wait + grep -h '^PASS\|^FAIL\|REFUSED\|regions kept' "$OUT"/T2-r*.log + exit 0 +fi + +MASTER=${1:?master address}; PEER=${2:?peer ssh target}; OUT=${3:?outdir} +mkdir -p "$OUT" +stop() { echo "STOP: $*" | tee -a "$OUT/verdict.txt"; exit 1; } + +precheck() { + [ -z "$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2> /dev/null)" ] || { echo "a GPU process runs"; return 1; } + [ "$(bash -c 'ulimit -l unlimited 2> /dev/null; ulimit -l')" = unlimited ] || { echo "memlock not unlimited"; return 1; } + echo "ok $(hostname) $(uname -r)" +} +counters() { + find -L /sys/class/infiniband/ -maxdepth 5 -path '*ports*counters*' -name '*err*' 2> /dev/null | sort | while read -r f; do + echo "$f $(cat "$f" 2> /dev/null)" + done +} +klog() { # kernel lines since epoch $1 matching $2 (journalctl must answer: an unreadable log is a STOP) + journalctl -k --since "@$1" --no-pager -o short-iso 2> /dev/null | grep -E "$2" + return "${PIPESTATUS[0]}" +} +probe() { # this host's tf-tp-test QPs and MRs once a second until it exits + local pid="" t + for _ in $(seq 200); do pid=$(pgrep -nx tf-tp-test) && break; sleep 0.1; done + [ -z "$pid" ] && return + for t in $(seq 180); do + kill -0 "$pid" 2> /dev/null || break + echo "== t=$t pid $pid" + rdma resource show qp pid "$pid" 2>&1 + rdma resource show mr pid "$pid" 2>&1 + sleep 1 + done +} +on() { # NODE CMD ARGS...: a command or one of the functions above, here (local) or on the peer + local node=$1 fn=$2 def + shift 2 + if [ "$node" = local ]; then "$fn" "$@"; return; fi + def=$(declare -f "$fn" 2> /dev/null) + ssh -o BatchMode=yes "$PEER" "${def:+$def +}$fn $(printf '%q ' "$@")" +} +grew() { # BEFORE AFTER: counters that went up + awk 'NR == FNR { b[$1] = $2; next } ($1 in b) && $2 > b[$1] { print $1 " " b[$1] " -> " $2 }' "$1" "$2" +} + +for n in local peer; do + v=$(on $n precheck) || stop "$n: $v" + echo "precheck $n: $v" | tee -a "$OUT/verdict.txt" +done +on peer test -x "$BIN" || stop "peer: no $BIN (the kit must sit at the same path on both hosts)" + +PORT=29900 +step() { # NAME WANT_RC -- ARGS... + local name=$1 want=$2 n rc rc0 rc1 hits g + shift 3 + PORT=$((PORT + 10)) + declare -A t + for n in local peer; do + t[$n]=$(on $n date +%s) + on $n counters > "$OUT/$name.$n.err0" + done + on local probe > "$OUT/$name.local.rdma.txt" 2>&1 & + on peer probe > "$OUT/$name.peer.rdma.txt" 2>&1 & + local envs=(TF_TP_MASTER="$MASTER" TF_TP_PORT=$PORT TF_COMM_BACKEND=roce TF_TP_ROCE_FALLBACK=error TF_TP_MAILBOX_MAX_KB=1024 + TF_TP_ROCE_FAST=1 TF_TP_FAILFAST_GRACE_S=15) + ssh -o BatchMode=yes "$PEER" "ulimit -l unlimited && env $(printf '%q ' "${envs[@]}") TF_TP_RANK=1 timeout 300 $(printf '%q ' "$BIN" "$@"); echo \"EXIT rank 1 rc \$? at \$(date +%s.%N)\"" > "$OUT/$name-r1.log" 2>&1 & + ( ulimit -l unlimited && env "${envs[@]}" TF_TP_RANK=0 timeout 300 "$BIN" "$@"; echo "EXIT rank 0 rc $? at $(date +%s.%N)" ) > "$OUT/$name-r0.log" 2>&1 + wait + rc0=$(grep -oP '^EXIT rank 0 rc \K[0-9]+' "$OUT/$name-r0.log") + rc1=$(grep -oP '^EXIT rank 1 rc \K[0-9]+' "$OUT/$name-r1.log") + echo "$name: rank 0 rc ${rc0:-?}, rank 1 rc ${rc1:-?} (want $want)" | tee -a "$OUT/verdict.txt" + for n in local peer; do + on $n counters > "$OUT/$name.$n.err1" + g=$(grew "$OUT/$name.$n.err0" "$OUT/$name.$n.err1") + [ -n "$g" ] && stop "$name $n: HCA error counters grew: $g" + hits=$(on $n klog "${t[$n]}" "$KLOG"); rc=$? + [ "$rc" -ne 0 ] && stop "$name $n: the kernel log is unreadable (journalctl rc $rc)" + [ -n "$hits" ] && { echo "$hits" > "$OUT/$name.$n.klog.txt"; stop "$name $n: kernel lines: $hits"; } + done + hits=$(grep -hE "$BAD" "$OUT/$name-r0.log" "$OUT/$name-r1.log") + [ -n "$hits" ] && stop "$name: $hits" + [ "$rc0" = "$want" ] && [ "$rc1" = "$want" ] || stop "$name: rc $rc0 / $rc1, want $want" + grep -h 'halted\|quiesced\|idle before\|regions kept' "$OUT/$name"-r*.log | sed "s/^/ /" +} +step T0 0 -- suite +step T1 70 -- failfast 1 transport 200 +step T3 70 -- failfast 0 transport 200 +echo "PASS transport check (T0, T1, T3)" | tee -a "$OUT/verdict.txt" diff --git a/zig/src/fabric/tp/sock.zig b/zig/src/fabric/tp/sock.zig new file mode 100644 index 0000000000..b382884034 --- /dev/null +++ b/zig/src/fabric/tp/sock.zig @@ -0,0 +1,169 @@ +//! TCP for the bootstrap and the fate channel: listen, accept and connect with deadlines, whole sends and receives. +const std = @import("std"); +const c = std.c; +const posix = std.posix; + +pub const Fd = c.fd_t; + +pub const Error = error{ SocketFailed, AddressUnresolved, AddressInUse, Timeout, Closed, Interrupted }; + +pub fn nowNs() u64 { + var ts: c.timespec = undefined; + _ = c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +pub fn sleepNs(ns: u64) void { + var ts: c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + while (c.nanosleep(&ts, &ts) != 0) {} +} + +/// TCP_NODELAY (frames are tiny and latency matters) and keepalive (a vanished host fails a read in ~90 s, not never). +fn tune(fd: Fd) void { + const one: c_int = 1; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, &one, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.KEEPALIVE, &one, @sizeOf(c_int)); + const idle: c_int = 30; + const intvl: c_int = 10; + const cnt: c_int = 6; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPIDLE, &idle, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPINTVL, &intvl, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPCNT, &cnt, @sizeOf(c_int)); +} + +/// Rank 0's listening socket on every IPv4 address at `port`. +pub fn listen(port: u16) Error!Fd { + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + errdefer _ = c.close(fd); + const one: c_int = 1; + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, &one, @sizeOf(c_int)); + var addr: posix.sockaddr.in = .{ .port = std.mem.nativeToBig(u16, port), .addr = 0 }; + if (c.bind(fd, @ptrCast(&addr), @sizeOf(posix.sockaddr.in)) != 0) return error.AddressInUse; + if (c.listen(fd, 64) != 0) return error.SocketFailed; + return fd; +} + +/// One connection before `deadline_ns` (monotonic). +pub fn accept(fd: Fd, deadline_ns: u64) Error!Fd { + while (true) { + try waitReadable(fd, deadline_ns); + const conn = c.accept(fd, null, null); + if (conn >= 0) { + tune(conn); + return conn; + } + } +} + +/// Connects to `host:port`, retrying every 50 ms (the master may not listen yet) until `deadline_ns`. +pub fn connect(host: []const u8, port: u16, deadline_ns: u64) Error!Fd { + var name: [256]u8 = undefined; + if (host.len >= name.len) return error.AddressUnresolved; + @memcpy(name[0..host.len], host); + name[host.len] = 0; + var service: [8]u8 = undefined; + const svc = std.fmt.bufPrintSentinel(&service, "{d}", .{port}, 0) catch unreachable; + const hints: c.addrinfo = .{ .flags = .{}, .family = posix.AF.INET, .socktype = posix.SOCK.STREAM, .protocol = 0, .addrlen = 0, .addr = null, .canonname = null, .next = null }; + while (true) { + var res: ?*c.addrinfo = null; + if (@backingInt(c.getaddrinfo(name[0..host.len :0], svc, &hints, &res)) == 0) { + defer c.freeaddrinfo(res.?); + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + if (c.connect(fd, res.?.addr.?, res.?.addrlen) == 0) { + tune(fd); + return fd; + } + _ = c.close(fd); + } + if (nowNs() >= deadline_ns) return error.Timeout; + sleepNs(50 * std.time.ns_per_ms); + } +} + +fn waitReadable(fd: Fd, deadline_ns: u64) Error!void { + while (true) { + const t = nowNs(); + if (t >= deadline_ns) return error.Timeout; + const ms: c_int = @intCast(@min((deadline_ns - t) / std.time.ns_per_ms + 1, 1000)); + var p = [_]c.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const n = c.poll(&p, 1, ms); + if (n > 0) return; + } +} + +/// Every byte of `bytes`, or `error.Closed`. +pub fn sendAll(fd: Fd, bytes: []const u8) Error!void { + var off: usize = 0; + while (off < bytes.len) { + const n = c.send(fd, bytes[off..].ptr, bytes.len - off, posix.MSG.NOSIGNAL); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// Fills `out` before `deadline_ns` (null: no deadline), or `error.Closed` at the end of the stream. +pub fn recvExact(fd: Fd, out: []u8, deadline_ns: ?u64) Error!void { + var off: usize = 0; + while (off < out.len) { + if (deadline_ns) |d| try waitReadable(fd, d); + const n = c.recv(fd, out[off..].ptr, out.len - off, 0); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// A non-blocking send of the whole of `bytes` if the socket buffer has room (false: dropped, or the peer is gone). +pub fn trySend(fd: Fd, bytes: []const u8) bool { + const n = c.send(fd, bytes.ptr, bytes.len, posix.MSG.NOSIGNAL | posix.MSG.DONTWAIT); + return n == @as(isize, @intCast(bytes.len)); +} + +pub fn hangup(fd: Fd) void { + _ = c.shutdown(fd, posix.SHUT.RDWR); +} + +pub fn close(fd: Fd) void { + _ = c.close(fd); +} + +/// This host's identity (the kernel's boot id, else the host name): ranks with equal ids share memory. +pub fn hostId(out: *[64]u8) []const u8 { + @memset(out, 0); + const f = c.open("/proc/sys/kernel/random/boot_id", .{ .ACCMODE = .RDONLY }); + if (f >= 0) { + defer _ = c.close(f); + const n = c.read(f, out, out.len); + if (n > 0) return std.mem.trim(u8, out[0..@intCast(n)], " \n"); + } + if (c.gethostname(out, out.len) != 0) return ""; + return std.mem.sliceTo(out, 0); +} + +test "a listener and a client trade bytes and see the hang-up" { + const l = try listen(0); + defer close(l); + var addr: posix.sockaddr.in = undefined; + var len: posix.socklen_t = @sizeOf(posix.sockaddr.in); + try std.testing.expectEqual(@as(c_int, 0), c.getsockname(l, @ptrCast(&addr), &len)); + const port = std.mem.bigToNative(u16, addr.port); + const deadline = nowNs() + 5 * std.time.ns_per_s; + const a = try connect("localhost", port, deadline); + defer close(a); + const b = try accept(l, deadline); + defer close(b); + try sendAll(a, "hello"); + var got: [5]u8 = undefined; + try recvExact(b, &got, deadline); + try std.testing.expectEqualStrings("hello", &got); + try std.testing.expectError(error.Timeout, recvExact(b, &got, nowNs() + 20 * std.time.ns_per_ms)); + hangup(a); + try std.testing.expectError(error.Closed, recvExact(b, &got, deadline)); +} diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig new file mode 100644 index 0000000000..2c9eeb85ec --- /dev/null +++ b/zig/src/fabric/tp/tp.zig @@ -0,0 +1,217 @@ +//! Tensor parallelism across processes, one GPU each: bootstrap, collectives, fail-fast; `Session.open` starts a rank. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); + +pub const collective = @import("collective.zig"); +pub const reference = @import("reference.zig"); +pub const config = @import("config.zig"); +pub const sock = @import("sock.zig"); +pub const bootstrap = @import("bootstrap.zig"); +pub const fate = @import("fate.zig"); +pub const host = @import("host.zig"); +pub const nccl = @import("nccl.zig"); +pub const mailbox = @import("mailbox.zig"); +pub const hybrid = @import("hybrid.zig"); +pub const roce = @import("roce.zig"); +pub const roce_verbs = @import("roce_verbs.zig"); +pub const roce_pick = @import("roce_pick.zig"); +pub const planlink = @import("planlink.zig"); +pub const Verbs = @import("verbs").Verbs; + +pub const Collective = collective.Collective; +pub const DType = collective.DType; +pub const Op = collective.Op; +pub const Config = config.Config; +pub const Bootstrap = bootstrap.Bootstrap; +pub const Fate = fate.Fate; + +pub const Session = struct { + cfg: Config, + lib: cuda.nccl.Library, + nccl: nccl.Nccl, + mail: ?mailbox.Mailbox = null, + roce: ?*roce.Roce = null, + route: hybrid.Hybrid = undefined, + /// The round plan's link (rank 0 sends, the others receive), on its own TCP connection. + plan: planlink.PlanLink = undefined, + fate: ?Fate = null, + + /// Brings this rank up in place; `kernel_image` is mailbox.cu's fatbin, empty sends those backends to NCCL. + pub fn open(s: *Session, d: *const cuda.Driver, cfg: Config, kernel_image: []const u8) !void { + var b = try Bootstrap.open(cfg); + defer b.close(); + try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_ROCE_CHUNKED_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK, TF_GRAPH_CACHE)"); + var lib = try cuda.nccl.Library.openPath(cfg.nccl_lib); + errdefer lib.close(); + const v = try lib.version(); + try b.agree(std.mem.asBytes(&v), "NCCL version"); + var id: cuda.nccl.UniqueId = undefined; + if (cfg.rank == 0) try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try b.broadcast(&id.internal, &id.internal); + s.* = .{ .cfg = cfg, .lib = lib, .nccl = undefined }; + s.nccl = try nccl.Nccl.init(&s.lib, d, id, cfg.rank, cfg.world); + errdefer s.nccl.deinit(); + switch (cfg.backend) { + .nccl => {}, + .mailbox => if (try mailboxPossible(&b, kernel_image)) { + s.mail = try mailbox.Mailbox.init(d, &b, kernel_image, cfg.mailbox_max_bytes, cfg.mailbox_timeout_ns); + s.route = .{ .small = .{ .mailbox = &s.mail.? }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + } else if (cfg.rank == 0) { + log.warn("[tensorfold] tp: TF_COMM_BACKEND=mailbox needs two ranks on one host and the kernel; every rank uses NCCL", .{}); + }, + .roce => try s.openRoce(d, &b, kernel_image), + } + try s.agreeVar(&b); + s.plan = planlink.PlanLink.init(cfg.rank, cfg.world, b.releasePlan()); + s.plan.spin_ns = cfg.plan_spin_ns; + s.plan.pin = cfg.plan_pin; + s.plan.avoid = cfg.roce_cpu; + if (cfg.failfast) { + s.fate = Fate.init(cfg, b.release()); + s.fate.?.onFail(.{ .ctx = s, .call = abortHook }); + try s.fate.?.start(); + } + log.info("[tensorfold] tp: rank {d} of {d} up ({t}, NCCL {d}, fail-fast {s}, roce fast {s}: {d} poller(s), stagger {d} ns)", .{ cfg.rank, cfg.world, s.comm().kind(), v, if (cfg.failfast) "on" else "off", if (cfg.roce_fast & 1 != 0) "on" else "off", (cfg.roce_fast >> 8) & 0xff, cfg.roce_fast >> 16 }); + } + + /// RoCE for the small exchanges, probed against NCCL at load; on failure every rank obeys TF_TP_ROCE_FALLBACK. + fn openRoce(s: *Session, d: *const cuda.Driver, b: *Bootstrap, image: []const u8) !void { + const cfg = s.cfg; + var specs: [roce_verbs.max_hcas]roce_verbs.HcaSpec = undefined; + var n_specs: usize = 0; + if (cfg.roce_wire == .verbs) n_specs = try roce_pick.resolve(cfg, b, &specs); + const opts: roce.Options = .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n_specs], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = cfg.mailbox_max_bytes, .timeout_ns = cfg.mailbox_timeout_ns, .fast = cfg.roce_fast, .blocks = cfg.roce_blocks }; + const r = roce.Roce.init(d, b, image, opts) catch |e| return s.roceFallback(e); + s.roce = r; + if (cfg.roce_chunked_bytes > 0 and !r.chunk_abi) { + log.warn("TF_TP_ROCE_CHUNKED_KB={d}: the RoCE kernel image has no chunked ABI (tp_roce_chunk_abi); refusing to start", .{cfg.roce_chunked_bytes / 1024}); + r.deinit(); + s.roce = null; + return error.ChunkAbiMissing; + } + s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes, .chunked_bytes = cfg.roce_chunked_bytes }; + const good = probe(s, d) catch false; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(good)}, &buf, &blobs); + for (all) |x| if (x[0] != 1) { + r.deinit(); + s.roce = null; + return s.roceFallback(error.ProbeMismatch); + }; + } + + /// `Collective.varAgreed`: whether every rank's transport moves device lengths, agreed once a session. + fn agreeVar(s: *Session, b: *Bootstrap) !void { + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(s.comm().varFits(16))}, &buf, &blobs); + var ok = true; + for (all) |x| ok = ok and x[0] == 1; + if (s.mail != null or s.roce != null) s.route.var_agreed = ok; + } + + fn roceFallback(s: *Session, e: anyerror) !void { + if (s.cfg.roce_fallback == .@"error") return e; + if (s.cfg.rank == 0) log.warn("[tensorfold] tp: RoCE unavailable ({t}); every rank uses NCCL (TF_TP_ROCE_FALLBACK=error refuses instead)", .{e}); + } + + /// One 4 KiB all-gather and one bf16 sum through RoCE and through NCCL: the bytes must be equal. + fn probe(s: *Session, d: *const cuda.Driver) !bool { + const n = 4096; + // input [0, n), gathers [n, 3n) and [3n, 5n), sums [5n, 6n) and [6n, 7n) + var a = try cuda.DeviceBuffer.alloc(d, 7 * n); + defer a.free(); + var stream = try cuda.Stream.init(d, true); + defer stream.deinit(); + var bytes: [n]u8 = undefined; + reference.fill(&bytes, .bf16, s.cfg.rank, 0x7072_6f62); + try a.uploadAsync(0, &bytes, stream.handle); + const fast = s.route.iface(); + const slow = s.nccl.iface(); + try fast.allGather(a.ptr, a.ptr + n, n, .u8, stream.handle); + try slow.allGather(a.ptr, a.ptr + 3 * n, n, .u8, stream.handle); + try fast.allReduce(a.ptr, a.ptr + 5 * n, n / 2, .bf16, .sum, stream.handle); + try slow.allReduce(a.ptr, a.ptr + 6 * n, n / 2, .bf16, .sum, stream.handle); + try stream.synchronize(); + try fast.check(); + var got: [6 * n]u8 = undefined; + try a.download(n, &got); + const gathers = std.mem.eql(u8, got[0 .. 2 * n], got[2 * n .. 4 * n]); + const sums = std.mem.eql(u8, got[4 * n .. 5 * n], got[5 * n .. 6 * n]); + if (!gathers or !sums) log.warn("[tensorfold] tp roce probe: gather {s}, sum {s} against NCCL", .{ if (gathers) "equal" else "DIFFERS", if (sums) "equal" else "DIFFERS" }); + return gathers and sums; + } + + fn mailboxPossible(b: *Bootstrap, image: []const u8) !bool { + if (b.world != 2) return false; + var id_buf: [64]u8 = undefined; + const id = sock.hostId(&id_buf); + var mine: [65]u8 = @splat(0); + mine[0] = @intFromBool(image.len > 0); + @memcpy(mine[1..][0..id.len], id); + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [2 * 65]u8 = undefined; + const all = try b.allGather(&mine, &buf, &blobs); + return all[0][0] == 1 and std.mem.eql(u8, all[0], all[1]); + } + + fn abortHook(ctx: *anyopaque, _: []const u8) void { + const s: *Session = @ptrCast(@alignCast(ctx)); + s.comm().abort(); + } + + /// The communicator the model uses: NCCL, or the hybrid route over the mailbox / RoCE. + pub fn comm(s: *Session) Collective { + if (s.mail != null or s.roce != null) return s.route.iface(); + return s.nccl.iface(); + } + + /// This rank failed inside a round: every rank goes down (exit 70). Without fail-fast it only aborts. + pub fn fatal(s: *Session, where: []const u8, err: anyerror) void { + if (s.fate) |*f| return f.fatal(where, err); + s.comm().abort(); + } + + /// A clean stop announced by rank 0 before it closes: followers end with no fail-fast exit. + pub fn announceStop(s: *Session) void { + if (s.fate) |*f| f.announceStop(); + } + + /// This rank failed after `open`, before serving: the other ranks learn why and go down (fail-fast). + pub fn bootFailed(s: *Session, err: anyerror) void { + if (s.fate) |*f| f.failBoot("boot", err); + s.close(); + } + + /// A normal stop: every rank calls it (rank 0 first marks the stop as expected with `fate.expectStop`). + pub fn close(s: *Session) void { + // a failed transport: every rank goes down (the peers stop posting) before this one lets go of anything + if (s.roce) |r| if (r.failed()) if (s.fate) |*f| f.fatal("transport", error.BackendFailed); + if (s.fate) |*f| f.close(); + s.plan.close(); + if (s.roce) |r| r.deinit(); + if (s.mail) |*m| m.deinit(); + s.nccl.deinit(); + s.lib.close(); + } +}; + +test { + _ = collective; + _ = reference; + _ = config; + _ = sock; + _ = bootstrap; + _ = mailbox; + _ = roce_verbs; + _ = roce_pick; + _ = planlink; + _ = @import("roce_guard.zig"); + _ = @import("roce_test.zig"); + _ = @import("roce_post_test.zig"); + _ = @import("roce_inflight_test.zig"); + _ = @import("host_test.zig"); + _ = @import("fate_test.zig"); +} diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index 9b9353c554..3c32714d30 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -2,7 +2,7 @@ const std = @import("std"); const platform = @import("sys"); const builtin = @import("builtin"); -const abi = @import("verbs_abi.zig"); +pub const abi = @import("verbs_abi.zig"); pub const Error = error{ LibraryUnavailable, MissingSymbol, NoDeviceList, NoSuchDevice, VerbsFailed, BadDeviceList, NoUsableGid }; pub const max_links = 8; @@ -51,12 +51,15 @@ pub const Api = struct { /// Exported by rdma-core 34+ but not in every header; used only when present. const GidEx = *const fn (ctx: Ctx, port: u32, index: u32, entry: *abi.GidEntry, flags: u32, size: usize) callconv(.c) c_int; const InOrder = *const fn (qp: *abi.Qp, op: c_int, flags: u32) callconv(.c) c_int; +/// ibv_query_qp: the QP's state as the device holds it (not every library exports it; used only when present). +pub const QueryQp = *const fn (qp: *abi.Qp, attr: *abi.QpAttr, mask: c_int, init: *abi.QpInitAttr) callconv(.c) c_int; pub const Verbs = struct { lib: platform.DynLib, api: Api, gid_ex: ?GidEx = null, in_order: ?InOrder = null, + query_qp: ?QueryQp = null, pub fn open(path: ?[:0]const u8) Error!Verbs { var lib = platform.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; @@ -69,6 +72,7 @@ pub const Verbs = struct { var v: Verbs = .{ .lib = lib, .api = api }; v.gid_ex = v.lib.lookup(GidEx, "_ibv_query_gid_ex"); v.in_order = v.lib.lookup(InOrder, "ibv_query_qp_data_in_order"); + v.query_qp = v.lib.lookup(QueryQp, "ibv_query_qp"); return v; } diff --git a/zig/src/families/deepseek_v41/block.zig b/zig/src/families/deepseek_v41/block.zig new file mode 100644 index 0000000000..cee25c90e4 --- /dev/null +++ b/zig/src/families/deepseek_v41/block.zig @@ -0,0 +1,520 @@ +//! A decode window as calls: every launch with its arguments and role-named tensors ("s." / "w." / "L." lifetimes). + +const std = @import("std"); +const dk = @import("dsv41_kernels"); +const calls = @import("block_calls.zig"); +const config = @import("config.zig"); +const knob_types = @import("knob_types.zig"); +const moe = @import("block_moe.zig"); +const block_prefetch = @import("block_prefetch.zig"); +const block_attn = @import("block_attn.zig"); +const block_index = @import("block_index.zig"); +const Config = config.Config; +const Arg = calls.Arg; +const Named = calls.Named; +const Dt = calls.Dt; +const dense = dk.dense; + +/// The widths the loader picked: K2 of each dense group and each routed layer's expert K2 range (gate/up, down). +pub const Widths = struct { + dense: std.StringHashMapUnmanaged(u32) = .empty, + experts: std.AutoHashMapUnmanaged(u32, [2][2]u32) = .empty, + /// each routed layer's expert K2s present (bit k2 set), gate/up and down: x3gm runs one launch a width present + gm: std.AutoHashMapUnmanaged(u32, [2]u32) = .empty, + + pub fn k2(w: *const Widths, name: []const u8) !u32 { + return w.dense.get(name) orelse error.NoWidth; + } +}; + +/// Run-time sizes fixed at load (scratch and pools), and the L2 prefetch settings. +pub const Options = struct { + /// TF_DSV41_IMAGES=native: this prefill segment holds image positions, so Engram fuses with a keep vector + image_keep: bool = false, + /// TF_DSV41_BIAS_VL: the segment's image rows (MoE call apart, gate.bias_vl); 0: none; < 0: the buffer plan's + image_rows: i64 = 0, + /// TP ranks (a boot constant): every per-rank width below is the model's / world; the plan checked divisibility + world: u32 = 2, + /// KV positions a slot holds (pool rows = limit / ratio + 1) + limit: i64 = 4096, + /// rope table rows (cos / sin pairs of 32) + rope_rows: i64 = 6144, + /// rows the routed experts' scratch holds: the backbone's and a DSpark block's + expert_rows: i64 = 1024, + dspark_rows: i64 = 64, + /// TF_DSV41_BRANCHES: the attention sublayer's indexer / compressor chains tagged for the side stream + branches: bool = false, + branch_rows: i64 = 64, + /// TF_DSV41_MHC_DEFER (with `r1`): its coefficient launch on the deferred stream, joined by the next mHC call + mhc_defer: bool = false, + /// TF_DSV41_MHC_DEFER_AT=exchange: deferred launches held to the next exchange; _ROWS: only up to that many rows + mhc_defer_hold: bool = false, + /// TF_DSV41_COEF_LATE (with MHC_DEFER): a deferred coef launch moves to just before its sublayer's exchange + coef_late: bool = false, + mhc_defer_hold_rows: i64 = 0, + /// TF_DSV41_MHC_PFDEC (0 = off): mixing sites of windows of at least this many rows on mhc_pf (block_wide.zig) + mhc_pf_rows: i64 = 0, + /// TF_DSV41_L2PF_MB (bytes), _CHUNK_KB, _PACE_GBPS, _PACE_CTAS; the GPU's SMs (a segments launch's grid bound) + l2pf_budget: u64 = 12 << 20, + l2pf_chunk: i64 = 32768, + l2pf_pace_gbps: f64 = 150, + l2pf_pace_ctas: i64 = 2, + sms: i64 = 188, + /// TF_DSV41_EXPERT_TOPP (lossy). null: no pruning; the router's tail rounds kit weights and groups experts itself + expert_topp: ?f64 = 0.85, + /// after each block, its streams with the last post applied out of place, to glue step "trace" (never served) + trace: bool = false, + /// DSpark's taps: boundaries entering target layers write the streams' mean into "w.taps" bf16 [n, 3 x hidden] + taps: bool = false, + /// rows a prefill segment holds at most (forward.pchunk): sizes the SWA staging ring (block_prefill.zig) + prefill_rows: i64 = 2048, + /// TF_DSV41_PF_4K: prefill segments up to 4,096 rows, router in 2,048-row launches; off: 2,048 (longer refused) + pf4k: bool = false, + /// TF_DSV41_PF_OVERLAP: a prefill layer's two exchanges in this many row pieces on the side stream (branches, TP=2) + pf_overlap: i64 = 0, + /// TF_DSV41_PF_OVERLAP_SITE (with pf_overlap): all exchange pieces on the side stream, the next site in row pieces + pf_overlap_site: bool = false, + /// TF_DSV41_MHC_SITE_ROWS (0 off; 256-4,096, multiple of 16): a prefill boundary site in row pieces on main + mhc_site_rows: i64 = 0, + /// TF_DSV41_PF_TBO: CED encoder segments two at a time, the second a layer behind on the side stream (`branches`) + pf_tbo: bool = false, + /// TF_DSV41_STREAM_RB (0 off; 1 / 2 / 4): the long-context stream top-k as the row-blocked twin `_stream_rb` + stream_rb: i64 = 0, + /// TF_DSV41_INDEX_BOUND: decode index scores and top-k as row-bounded twins that stop at each row's visible end + index_bound: bool = false, + /// decode knob set (same bits): ATTN_SPLIT, ATTN_ROPE, MHC_TAIL, MHC_DEFER, RG_PRUNE, PRUNE_KIT, MOE_BF16, RMS_FOLD + r1: bool = false, + /// TF_DSV41_ROUTER_GROUP_ROT: wide grouping and expert input rotation in one launch (default off). + router_group_rot: bool = false, + /// TF_DSV41_X3LD_EPI: decode x3ld + epilogues as two fused launches (ticket words "s.ex.epi" / "s.dx.epi") + x3ld_epi: bool = false, + /// KV-only RMS feeding SWA store; separately gated, default off. + kv_norm_store: bool = false, + /// `r1` off but a twin's bindings that always take its trailing arguments, at their off values (same launches) + r1_sig: bool = false, + /// MHC_TAIL_SPIN (SM cycles a CTA waits for the last arrival) + mhc_spin: i64 = 8000, + /// the KV families in the paged pool (kv.DevicePool) instead of the slot's contiguous buffers; null: contiguous + pool: ?Pool = null, + /// TF_DSV41_PF_DENSE=fused: prefill dense EXL3 GEMMs on the pfdense kernel; null: W_q unpack + `_gemm` (same bits) + pfd: ?*const knob_types.PfDense = null, + /// TF_DSV41_GM_V2 (gm2pf.zig): prefill's routed experts on x3gm v2; off: v1. The same bits either way. + gm_v2: knob_types.GmMode = .off, + /// TF_DSV41_GM_GU2 (with gm_v2 `one`): a 4,096-row segment's gate / up at the 128-member tile GU2 + gm_gu2: bool = false, + /// TF_DSV41_INDEX_BUDGET_MIB (bytes): indexer scores + sort keys a launch; past it, row blocks (same bits) + index_budget: i64 = 256 << 20, + /// a row-mode program (batch.zig, rowmode.zig): row blocks score the context bucket's keys; no streamed top-k + rows: bool = false, + + /// TF_DSV41_MHC_DEFER_AT(_ROWS): a program of `rows` rows issues its deferred launches at the next exchange. + pub fn holdsDeferred(o: Options, rows: i64) bool { + return o.mhc_defer_hold and (o.mhc_defer_hold_rows == 0 or rows <= o.mhc_defer_hold_rows); + } +}; + +/// The paged pool as a window's calls see it: comp rows, index keys, page table and split table (TF_DSV41_KV_SPLIT). +pub const Pool = struct { + /// pages each family tensor holds (comp: split's local pages + null + discard) + comp_pages: i64, + ik_pages: i64, + /// pages a slot's table holds + pts: i64, + split: bool = false, + /// TF_DSV41_KV_SPLIT_EARLY: an index layer's exchange right after its selection (side stream), beside the q group + early: bool = false, + /// split's union cap (TF_DSV41_KV_SPLIT_UNION_MIB): the prefill receive buffer's bytes + union_cap: i64 = 128 << 20, + pub const page: i64 = 256; + pub const row_bytes: i64 = 584; + pub const ik_bytes: i64 = 132; + + /// log2 of a family's rows a page + pub fn shift(ratio: i64) i64 { + return std.math.log2_int(u64, @intCast(@divExact(page, ratio))); + } +}; + +pub const Emitter = struct { + a: std.mem.Allocator, + cfg: *const Config, + w: *const Widths, + o: Options, + /// rows this window and the position of its first + n: i64, + start: i64, + out: std.ArrayList(calls.Call) = .empty, + /// which of the two stream buffers holds the streams, and which coefficient set the last boundary wrote + cur: u1 = 0, + coef: u1 = 0, + begin: calls.Begin = .window, + /// the current layer's compression ratio (picks its rope table) + layer_ratio: u32 = 0, + /// TF_DSV41_BRANCHES: the stream the next calls go to, and the marks the next call carries + br_side: bool = false, + br_fork: bool = false, + br_join: bool = false, + br_mark: u8 = 0, + br_wait: u8 = 0, + /// TF_DSV41_MHC_DEFER: the next call is the coefficient launch / joins it; a coefficient launch not joined yet + df_side: bool = false, + df_join: bool = false, + df_pending: bool = false, + /// TF_DSV41_COEF_LATE: the index of a deferred coefficient launch not joined yet (moved to the next exchange) + df_at: ?usize = null, + /// a multi-segment prefill: the current segment's first row in the run's per-row roles; 0 elsewhere + row0: i64 = 0, + /// TF_DSV41_PF_TBO (block_prefill.emitTbo): the call indexes where each layer's attention part and MoE part start + marks: ?*std.ArrayList(usize) = null, + /// its first row in the run's narrow per-row roles, a multiple of 4 (16-byte aligned slices); 0 elsewhere + cnt0: i64 = 0, + + pub fn fmt(e: *Emitter, comptime f: []const u8, args: anytype) ![]const u8 { + return std.fmt.allocPrint(e.a, f, args); + } + + /// A contiguous tensor of `role`. + pub fn t(e: *Emitter, role: calls.Role, dt: Dt, shape: []const i64) !Arg { + const s = try e.a.dupe(i64, shape); + return .{ .t = .{ .role = role, .dt = dt, .shape = s, .stride = try calls.contiguous(e.a, s) } }; + } + + /// A view of `role`: its own strides and a byte offset. + pub fn view(e: *Emitter, role: calls.Role, dt: Dt, shape: []const i64, stride: []const i64, offset: i64) !Arg { + return .{ .t = .{ .role = role, .dt = dt, .shape = try e.a.dupe(i64, shape), .stride = try e.a.dupe(i64, stride), .offset = offset } }; + } + + pub fn buf(e: *Emitter, comptime f: []const u8, args: anytype, dt: Dt, shape: []const i64) !Arg { + return e.t(.{ .buf = try e.fmt(f, args) }, dt, shape); + } + + pub fn weight(e: *Emitter, comptime f: []const u8, args: anytype, dt: Dt, shape: []const i64) !Arg { + return e.t(.{ .weight = try e.fmt(f, args) }, dt, shape); + } + + pub fn empty(e: *Emitter, dt: Dt) !Arg { + return e.t(.empty, dt, &.{0}); + } + + pub fn ints(e: *Emitter, xs: []const i64) !Arg { + const l = try e.a.alloc(Arg, xs.len); + for (xs, l) |x, *y| y.* = .{ .i = x }; + return .{ .list = l }; + } + + fn push(e: *Emitter, c: calls.Call) !void { + var x = c; + x.begin = e.begin; + e.begin = .none; + x.side = e.br_side; + x.fork = e.br_fork; + x.join = e.br_join; + x.mark = e.br_mark; + x.wait = e.br_wait; + e.br_fork = false; + e.br_join = false; + e.br_mark = 0; + e.br_wait = 0; + x.defer_side = e.df_side; + x.defer_join = e.df_join; + e.df_side = false; + e.df_join = false; + if (e.o.coef_late and e.marks == null) try e.lateCoef(&x); + try e.out.append(e.a, x); + } + + /// TF_DSV41_COEF_LATE: a deferred coef launch moves to just before the next exchange of its fork region; same data. + fn lateCoef(e: *Emitter, x: *const calls.Call) !void { + if (x.defer_join) e.df_at = null; + if (x.defer_side) { + e.df_at = e.out.items.len; + return; + } + const at = e.df_at orelse return; + if (!(x.glue and std.mem.startsWith(u8, x.name, "glue.exchange"))) return; + e.df_at = null; + const c = e.out.orderedRemove(at); + std.debug.assert(c.defer_side and c.begin == .none); + try e.out.append(e.a, c); + } + + /// Appends positional arguments to the last emitted call (`r1`'s trailing arguments). + pub fn appendLast(e: *Emitter, args: []const Arg) !void { + const last = &e.out.items[e.out.items.len - 1]; + const more = try e.a.alloc(calls.Named, last.args.len + args.len); + @memcpy(more[0..last.args.len], last.args); + for (more[last.args.len..], args) |*m, x| m.* = .{ .arg = x }; + last.args = more; + } + + pub fn ext(e: *Emitter, name: []const u8, args: []const Arg) !void { + const l = try e.a.alloc(Named, args.len); + for (args, l) |x, *y| y.* = .{ .arg = x }; + try e.push(.{ .triton = false, .name = name, .args = l }); + } + + /// A glue step ("glue."): the roles it reads / writes and its scalars, for the forward to run. + pub fn glue(e: *Emitter, name: []const u8, args: []const Arg) !void { + const l = try e.a.alloc(Named, args.len); + for (args, l) |x, *y| y.* = .{ .arg = x }; + try e.push(.{ .triton = false, .glue = true, .name = try e.fmt("glue.{s}", .{name}), .args = l }); + } + + pub fn triton(e: *Emitter, name: []const u8, grid: [3]i64, args: []const Named) !void { + try e.push(.{ .triton = true, .name = name, .grid = grid, .args = try e.a.dupe(Named, args) }); + } + + /// A config value's float as the JSON number (f64), not the f32 we keep. + pub fn dec(x: f32) f64 { + var b: [64]u8 = undefined; + const s = std.fmt.bufPrint(&b, "{e}", .{x}) catch return x; + return std.fmt.parseFloat(f64, s) catch x; + } + + /// Positions after this window: the last row's + 1. + pub fn end(e: *const Emitter) i64 { + return e.start + e.n; + } + + pub fn streams(e: *Emitter) !Arg { + return e.buf("w.x{d}", .{e.cur}, .bf16, &.{ e.n, 4 * @as(i64, e.cfg.hidden) }); + } + + pub fn spare(e: *Emitter) !Arg { + return e.buf("w.x{d}", .{1 - e.cur}, .bf16, &.{ e.n, 4 * @as(i64, e.cfg.hidden) }); + } + + // dense EXL3 projections --------------------------------------------------------------------------------------- + + /// A trellis's run-time words (int32 [N/128, K/16, 8, 4 K2]): strips "s..T" or dense3 lanes "s..lanes". + pub fn strips(e: *Emitter, w: []const u8, k: i64, n: i64) !Arg { + return e.words(w, k, n, "T"); + } + + fn words(e: *Emitter, w: []const u8, k: i64, n: i64, form: []const u8) !Arg { + const k2: i64 = try e.w.k2(w); + return e.t(.{ .buf = try e.fmt("s.{s}.{s}", .{ w, form }) }, .i32, &.{ @divExact(n, 128), @divExact(k, 16), 8, 4 * k2 }); + } + + /// Bytes of a group's trellis (= its strips). + pub fn trellisBytes(e: *Emitter, w: []const u8, k: i64, n: i64) !u64 { + const k2: u64 = try e.w.k2(w); + return @as(u64, @intCast(@divExact(k, 16) * @divExact(n, 16))) * 8 * k2 * 2; + } + + /// One projection of a fused group: weight prefix, input, output (a whole buffer; `col` its first column), K x N. + pub const Proj = struct { w: []const u8, x: Arg, y: Arg, col: usize = 0, k: usize, n: usize }; + + /// `r1`'s RMS_FOLD: rot_in norms its rows itself (rms(x, w, eps), rmsnorm's order) before the rotation. + pub const Norm = struct { w: Arg, eps: f64 }; + + /// x3seg rot_in, then dense3 a width: the group's projections in one rot_in launch, one linear launch a K2. + pub fn fused(e: *Emitter, tag: []const u8, ps: []const Proj) !void { + return e.fusedNorm(tag, ps, null); + } + + pub fn fusedNorm(e: *Emitter, tag: []const u8, ps: []const Proj, norm: ?Norm) !void { + const m: usize = @intCast(e.n); + const segs = try e.a.alloc(dense.Seg, ps.len); + for (ps, segs, 0..) |p, *s, i| { + const pl = dense.plan(p.k, p.n); + s.* = .{ .k = p.k, .n = p.n, .k2 = try e.w.k2(p.w), .sk = pl[0], .wk = pl[1], .out = i, .col = p.col }; + } + const lay = try dense.layout(segs, m); + const xs = try e.a.alloc(Arg, ps.len); + const suh = try e.a.alloc(Arg, ps.len); + const offs = try e.a.alloc(i64, ps.len); + for (ps, xs, suh, offs, 0..) |p, *x, *s, *o, i| { + x.* = p.x; + s.* = try e.weight("{s}.suh", .{p.w}, .f16, &.{@intCast(p.k)}); + o.* = @intCast(lay.xh_off[i]); + } + const xh = try e.buf("L.{s}.xh", .{tag}, .f16, &.{@intCast(lay.xh_total)}); + if (norm) |nm| { + const ws = try e.a.alloc(Arg, ps.len); + @memset(ws, nm.w); + try e.ext("tf_dsv41_x3seg_v3.rot_in", &.{ .{ .list = xs }, .{ .list = suh }, xh, try e.ints(offs), .{ .b = false }, .{ .list = ws }, .{ .f = nm.eps } }); + } else try e.ext("tf_dsv41_x3seg_v3.rot_in", &.{ .{ .list = xs }, .{ .list = suh }, xh, try e.ints(offs), .{ .b = false } }); + const z: Arg = if (lay.z_total > 0) try e.buf("L.{s}.z", .{tag}, .f32, &.{@intCast(lay.z_total)}) else .none; + const counters = try e.buf("L.{s}.cnt", .{tag}, .i32, &.{@intCast(lay.c_total)}); + for (lay.launches[0..lay.n_launches]) |l| { + const T = try e.a.alloc(Arg, l.count); + const svh = try e.a.alloc(Arg, l.count); + const ys = try e.a.alloc(Arg, l.count); + var meta: std.ArrayList(i64) = .empty; + for (l.segs[0..l.count], 0..) |si, j| { + const s = segs[si]; + const p = ps[si]; + T[j] = try e.words(p.w, @intCast(p.k), @intCast(p.n), "lanes"); + svh[j] = try e.weight("{s}.svh", .{p.w}, .f16, &.{@intCast(p.n)}); + ys[j] = p.y; + const st = dense.strides(s.k, s.k2); + try meta.appendSlice(e.a, &.{ @intCast(s.k), @intCast(s.n), @intCast(s.sk), @intCast(s.wk), st[0], st[1], @intCast(lay.xh_off[si]), @intCast(lay.z_off[si]), @intCast(lay.c_off[si]), @intCast(s.col), @intCast(l.p0[j]) }); + } + try e.ext("tf_dsv41_dense3_v1.linear", &.{ + xh, .{ .list = T }, .{ .list = svh }, .{ .list = ys }, + z, counters, try e.ints(meta.items), .{ .i = @intCast(l.programs) }, + .{ .i = l.k2 }, .{ .i = @intCast(l.wk) }, .{ .i = if (l.wk == 8) 1 else 3 }, .{ .b = false }, + }); + } + } + + /// Upstream's rot_in + linear: one projection, its trellis as strips unless one 128-column block. + pub fn upstream(e: *Emitter, w: []const u8, x: Arg, k: i64, n: i64, y: Arg) !void { + const k2: i64 = try e.w.k2(w); + const xh = try e.buf("L.{s}.xh", .{w}, .f16, &.{ e.n, k }); + try e.ext("tensorfold_exl3_linear_v4.rot_in", &.{ x, try e.weight("{s}.suh", .{w}, .f16, &.{k}), xh }); + const pl = dense.plan(@intCast(k), @intCast(n)); + const st = dense.strides(@intCast(k), @intCast(k2)); + const T = if (n == 128) + try e.view(.{ .weight = try e.fmt("{s}.trellis", .{w}) }, .i32, &.{ 1, @divExact(k, 16), 8, 4 * k2 }, &.{ 32 * k2, 32 * k2, 4 * k2, 1 }, 0) + else + try e.strips(w, k, n); + if (pl[0] != 1) return error.Unsupported; // no upstream linear on the decode path splits K + try e.ext("tensorfold_exl3_linear_v4.linear", &.{ + xh, T, .{ .i = st[0] }, .{ .i = st[1] }, try e.weight("{s}.svh", .{w}, .f16, &.{n}), .none, y, .none, + try e.buf("L.{s}.cnt", .{w}, .i32, &.{8 * @divExact(n, 128)}), .{ .i = k2 }, .{ .i = 2 }, .{ .i = @intCast(pl[0]) }, .{ .i = @intCast(pl[1]) }, + }); + } + + // L2 prefetch (TF_DSV41_L2PF): table rows filled by l2pf.zig at boot, issued on the runner's side stream ------ + + /// What a site prefetches, in read order: a weight, a projection's trellis words, or an expert entry's trellis. + pub const L2Src = struct { kind: enum { weight, words, expert }, name: []const u8, index: u32 = 0, bytes: u64 }; + pub const L2Site = enum { o, f, x }; + + pub const linSrcs = block_prefetch.linSrcs; + pub const l2pfSources = block_prefetch.l2pfSources; + pub const srcSizes = block_prefetch.srcSizes; + pub const take = block_prefetch.take; + pub const prefetchO = block_prefetch.prefetchO; + pub const paced = block_prefetch.paced; + pub const prefetchF = block_prefetch.prefetchF; + pub const prefetchX = block_prefetch.prefetchX; + + const Group = struct { ws: []const []const u8, ks: []const i64, ns: []const i64 }; + + /// The x group (fused_proj): wq_a, wkv, then a KV source's compressor projections. + pub fn xGroup(e: *Emitter, L: u32) !Group { + const D: i64 = e.cfg.hidden; + const ratio = e.cfg.compressRatio(L); + const comp: usize = if (e.cfg.isKvSource(L)) (if (ratio == 2) 2 else 1) else 0; + const ws = try e.a.alloc([]const u8, 2 + comp); + const ns = try e.a.alloc(i64, 2 + comp); + const names = [_][]const u8{ "wq_a", "wkv", "comp_wkv", "comp_wgate" }; + for (ws, ns, 0..) |*w, *n, i| { + w.* = try e.fmt("L{d}.attn.{s}", .{ L, names[i] }); + n.* = if (i == 0) e.cfg.q_lora else e.cfg.head_dim; + } + const ks = try e.a.alloc(i64, ws.len); + @memset(ks, D); + return .{ .ws = ws, .ks = ks, .ns = ns }; + } + + /// The q group: wq_b, and an indexer's ix_wq_b. + pub fn qGroup(e: *Emitter, L: u32) !Group { + const ix = switch (e.cfg.mode(L)) { + .full, .reindex => true, + else => false, + }; + const ws = try e.a.alloc([]const u8, if (ix) 2 else 1); + const ns = try e.a.alloc(i64, ws.len); + ws[0] = try e.fmt("L{d}.attn.wq_b", .{L}); + ns[0] = @as(i64, e.cfg.heads / e.o.world) * e.cfg.head_dim; + if (ix) { + ws[1] = try e.fmt("L{d}.attn.ix_wq_b", .{L}); + ns[1] = @as(i64, e.cfg.index_heads) * e.cfg.index_dim; + } + const ks = try e.a.alloc(i64, ws.len); + @memset(ks, e.cfg.q_lora); + return .{ .ws = ws, .ks = ks, .ns = ns }; + } + + pub const rope = block_attn.rope; + pub const ropeTable = block_attn.ropeTable; + pub const pos = block_attn.pos; + pub const kvStore = block_attn.kvStore; + + /// A family's paging: the slot's table, rows-a-page shift and row strides (pool mode); none for ring / contiguous. + pub const Paging = struct { pt: Arg, psh: i64, vs: i64, ss: i64 }; + + pub const paging = block_attn.paging; + pub const pool = block_attn.pool; + pub const indexKeys = block_attn.indexKeys; + + /// Split KV: source `src`'s selected rows from their owners (dense up to 16 rows, else a union), read unpaged. + pub const Received = struct { cv: Arg, cs: Arg, tok: Arg }; + + pub const received = block_attn.received; + pub const exchange = block_attn.exchange; + pub const selection = block_index.selection; + pub const ticketLen = block_index.ticketLen; + pub const attnSplit = block_index.attnSplit; + pub const rms = block_attn.rms; + pub const kvNormStoreOn = block_attn.kvNormStoreOn; + pub const normStore = block_attn.normStore; + pub const swaStore = block_attn.swaStore; + pub const branching = block_attn.branching; + pub const fork = block_attn.fork; + pub const onMain = block_attn.onMain; + pub const join = block_attn.join; + pub const attention = block_attn.attention; + pub const index = block_index.index; + pub const scoreRows = block_index.scoreRows; + pub const blocked = block_index.blocked; + pub const blockPos = block_index.blockPos; + pub const topk = block_index.topk; + pub const topkOf = block_index.topkOf; + + // the window ------------------------------------------------------------------------------------------------- + + /// Every launch of a decode window of `n` rows from `start` over `layers` (ascending), and the head when `head`. + pub fn window(e: *Emitter, layers: []const u32, head: bool) !void { + // the window's positions and streams (this rank's embedding rows, their all-sum: the first exchange) + try e.glue("positions", &.{ try e.pos(), try e.buf("w.pos64", .{}, .i64, &.{e.n}), try e.buf("w.lo", .{}, .i32, &.{e.n}) }); + try e.glue("embed", &.{ try e.streams(), try e.buf("w.recv.moe", .{}, .bf16, &.{ e.o.world, e.n, e.cfg.hidden }) }); + for (layers, 0..) |L, i| { + if (i > 0) e.begin = .layer; + e.layer_ratio = e.cfg.compressRatio(L); + if (e.cfg.isEngram(L)) { + // an Engram block after the first: the last exchange's post alone, Engram, then a site2 collapse + if (i > 0) try moe.postOnly(e, null); + try moe.engram(e, L); + try moe.mhc(e, L, "hc_attn", if (i == 0) .site1 else .site2); + } else try moe.mhc(e, L, "hc_attn", if (i == 0) .site1 else .boundary); + try e.attention(L); + try e.prefetchF(L); + const D: i64 = e.cfg.hidden; + try e.glue("exchange", &.{ try e.buf("L.part", .{}, .bf16, &.{ e.n, D }), try e.buf("w.recv.attn", .{}, .bf16, &.{ e.o.world, e.n, D }) }); + try moe.mhc(e, L, "hc_ffn", .boundary); + try moe.moe(e, L); + if (i + 1 < layers.len) try e.prefetchX(layers[i + 1]); + // `r1` (MOE_BF16): down_combine stores the partial as bf16, the exchange takes it as is + if (e.o.r1) + try e.glue("exchange", &.{ try e.buf("L.moe", .{}, .bf16, &.{ e.n, D }), try e.buf("w.recv.moe", .{}, .bf16, &.{ e.o.world, e.n, D }) }) + else + try e.glue("exchange_f32", &.{ try e.buf("L.moe", .{}, .f32, &.{ e.n, D }), try e.buf("w.recv.moe", .{}, .bf16, &.{ e.o.world, e.n, D }) }); + if (e.o.trace) { + const y = try e.buf("w.trace", .{}, .bf16, &.{ e.n, 4 * D }); + try moe.postOnly(e, y); + try e.glue("trace", &.{ y, .{ .i = L } }); + } + } + if (head) try moe.finish(e); + } +}; + +/// The window's calls (arena-owned by `a`). +pub fn emit(a: std.mem.Allocator, cfg: *const Config, w: *const Widths, o: Options, layers: []const u32, n: i64, start: i64, head: bool) ![]calls.Call { + var e: Emitter = .{ .a = a, .cfg = cfg, .w = w, .o = o, .n = n, .start = start }; + try e.window(layers, head); + if (o.holdsDeferred(n)) calls.holdDeferred(e.out.items); + return e.out.items; +} + +/// exl3 expert_loads.k2_range: the instance covering a layer's half-bit widths [lo, hi]. +pub fn k2Range(lo: u32, hi: u32) ![2]i64 { + const ranges = [_][2]u32{ .{ 8, 8 }, .{ 2, 10 }, .{ 2, 12 } }; + for (ranges) |r| if (r[0] <= lo and hi <= r[1]) return .{ r[0], r[1] }; + return error.NoInstance; +} diff --git a/zig/src/families/deepseek_v41/block_attn.zig b/zig/src/families/deepseek_v41/block_attn.zig new file mode 100644 index 0000000000..24b828879f --- /dev/null +++ b/zig/src/families/deepseek_v41/block_attn.zig @@ -0,0 +1,344 @@ +//! A decode window's attention: rope, the KV stores and pools, the compressed exchange, SWA and CSA2 launches. + +const calls = @import("block_calls.zig"); +const block = @import("block.zig"); +const Arg = calls.Arg; +const Pool = block.Pool; +const Emitter = block.Emitter; +const dec = block.Emitter.dec; +const Proj = block.Emitter.Proj; +const Norm = block.Emitter.Norm; +const Paging = block.Emitter.Paging; +const Received = block.Emitter.Received; +const attnSplit = block.Emitter.attnSplit; + +// attention ---------------------------------------------------------------------------------------------------- + +pub fn rope(e: *Emitter, x: Arg, heads: i64, dim: i64, nope: i64) !void { + const n = e.n; + const xv = try e.view(x.t.role, .bf16, &.{ n, heads, dim }, &.{ heads * dim, dim, 1 }, 0); + try e.triton("_rope", .{ n, @divExact(heads, 4), 1 }, &.{ + .{ .name = "X", .arg = xv }, .{ .name = "x_rs", .arg = .{ .i = heads * dim } }, + .{ .name = "x_hs", .arg = .{ .i = dim } }, .{ .name = "OUT", .arg = xv }, + .{ .name = "o_rs", .arg = .{ .i = heads * dim } }, .{ .name = "o_hs", .arg = .{ .i = dim } }, + .{ .name = "CS", .arg = try e.ropeTable() }, .{ .name = "cs_rs", .arg = .{ .i = 64 } }, + .{ .name = "POS", .arg = try e.buf("w.pos64", .{}, .i64, &.{n}) }, .{ .name = "H", .arg = .{ .i = heads } }, + .{ .name = "NOPE", .arg = .{ .i = nope } }, .{ .name = "NB", .arg = .{ .i = 64 } }, + .{ .name = "HALF", .arg = .{ .i = 32 } }, .{ .name = "BH", .arg = .{ .i = 4 } }, + .{ .name = "INV", .arg = .{ .b = false } }, .{ .name = "COPY", .arg = .{ .b = false } }, + .{ .name = "PDL", .arg = .{ .b = false } }, + }); +} + +/// The rope table (cos / sin) a layer uses: the compressed layers' (compress_theta) or the SWA layers'. +pub fn ropeTable(e: *Emitter) !Arg { + return e.buf("s.rope.{s}", .{if (e.layer_ratio > 0) "comp" else "main"}, .f32, &.{ e.o.rope_rows, 64 }); +} + +pub fn pos(e: *Emitter) !Arg { + return e.buf("w.pos", .{}, .i32, &.{1}); +} + +/// The KV store (_kv_store): the SWA ring or a compressed family (split KV: only the rows this rank owns). +pub fn kvStore(e: *Emitter, lat: Arg, src: u32, ratio: i64) !void { + const rows = try e.pool(src, ratio); + const pg = try e.paging(ratio, true); + try e.triton("_kv_store", .{ e.n, 1, 1 }, &.{ + .{ .name = "LAT", .arg = lat }, .{ .name = "l_stride", .arg = .{ .i = e.cfg.head_dim } }, + .{ .name = "CS", .arg = try e.ropeTable() }, .{ .name = "cs_stride", .arg = .{ .i = 64 } }, + .{ .name = "V", .arg = rows[0] }, .{ .name = "S", .arg = rows[1] }, + .{ .name = "v_stride", .arg = .{ .i = pg.vs } }, .{ .name = "s_stride", .arg = .{ .i = pg.ss } }, + .{ .name = "POS", .arg = try e.pos() }, .{ .name = "RATIO", .arg = .{ .i = ratio } }, + .{ .name = "RING", .arg = .{ .i = if (ratio == 0) 2 * @as(i64, e.cfg.window) else 1 } }, .{ .name = "PT", .arg = pg.pt }, + .{ .name = "PSH", .arg = .{ .i = pg.psh } }, .{ .name = "SL", .arg = .none }, + .{ .name = "PTS", .arg = .{ .i = 0 } }, .{ .name = "ROWS", .arg = .{ .b = false } }, + .{ .name = "PDL", .arg = .{ .b = false } }, + }); +} + +pub fn paging(e: *Emitter, ratio: i64, store: bool) !Paging { + const p = e.o.pool orelse return .{ .pt = .none, .psh = 0, .vs = 576, .ss = 8 }; + if (ratio == 0) return .{ .pt = .none, .psh = 0, .vs = 576, .ss = 8 }; + const tab: []const u8 = if (store and p.split) "s.kv.ct" else "s.kv.pt"; + return .{ .pt = try e.buf("{s}", .{tab}, .i32, &.{p.pts}), .psh = Pool.shift(ratio), .vs = Pool.row_bytes, .ss = Pool.row_bytes }; +} + +/// A layer's KV rows (576 B) and scales (8 B): SWA ring, contiguous compressed rows, or the pool's 584 B records. +pub fn pool(e: *Emitter, L: u32, ratio: i64) ![2]Arg { + if (ratio > 0) if (e.o.pool) |p| { + const rows = p.comp_pages * @divExact(Pool.page, ratio); + const role: calls.Role = .{ .buf = try e.fmt("s.kv.comp.L{d}", .{L}) }; + return .{ try e.view(role, .u8, &.{ rows, 576 }, &.{ Pool.row_bytes, 1 }, 0), try e.view(role, .u8, &.{ rows, 8 }, &.{ Pool.row_bytes, 1 }, 576) }; + }; + const rows: i64 = if (ratio == 0) 2 * @as(i64, e.cfg.window) else @divFloor(e.o.limit, ratio) + 1; + const kind = if (ratio == 0) "swa" else "comp"; + return .{ try e.buf("s.L{d}.{s}.v", .{ L, kind }, .u8, &.{ rows, 576 }), try e.buf("s.L{d}.{s}.s", .{ L, kind }, .u8, &.{ rows, 8 }) }; +} + +pub fn indexKeys(e: *Emitter, L: u32) !Arg { + const r: i64 = e.cfg.compressRatio(L); + if (e.o.pool) |p| return e.buf("s.kv.ik.L{d}", .{L}, .u8, &.{ p.ik_pages * @divExact(Pool.page, r), Pool.ik_bytes }); + return e.buf("s.L{d}.ik", .{L}, .u8, &.{ @divFloor(e.o.limit, r) + 1, @as(i64, e.cfg.index_dim) + 4 }); +} + +/// The rows and tokens of the last exchange (a reuse layer: the same selection and kv source). +pub fn received(e: *Emitter, src: u32, ratio: i64) !Received { + _ = src; + _ = ratio; + const p = e.o.pool.?; + const n = e.n; + const K: i64 = e.cfg.index_topk; + const W: i64 = e.o.world; + const sent: i64 = if (n <= 16) n * K else @divFloor(p.union_cap, (W + 1) * Pool.row_bytes); + const recv: calls.Role = .{ .buf = "w.kx.recv" }; + return .{ + .cv = try e.view(recv, .u8, &.{ W * sent, 576 }, &.{ Pool.row_bytes, 1 }, 0), + .cs = try e.view(recv, .u8, &.{ W * sent, 8 }, &.{ Pool.row_bytes, 1 }, 576), + .tok = try e.buf("w.kx.tok", .{}, .i32, &.{ n, K }), + }; +} + +pub fn exchange(e: *Emitter, src: u32, ratio: i64, sel: Arg) !Received { + const p = e.o.pool.?; + const n = e.n; + const K: i64 = e.cfg.index_topk; + const W: i64 = e.o.world; + // row windows exchange dense at any rows, as the reference does + const small = n <= 16 or e.o.rows; + // dense: every rank sends R x K rows; union: (W + 1) x M rows under the cap + const sent: i64 = if (small) n * K else @divFloor(p.union_cap, (W + 1) * Pool.row_bytes); + const recv = try e.buf("w.kx.recv", .{}, .u8, &.{ W * sent, Pool.row_bytes }); + const tok = try e.buf("w.kx.tok", .{}, .i32, &.{ n, K }); + const comp = try e.buf("s.kv.comp.L{d}", .{src}, .u8, &.{ p.comp_pages * @divExact(Pool.page, ratio), Pool.row_bytes }); + try e.glue(if (small) "kx_dense" else "kx_union", &.{ + sel, comp, + try e.buf("s.kv.ct", .{}, .i32, &.{p.pts}), try e.buf("w.kx.send", .{}, .u8, &.{ sent, Pool.row_bytes }), + recv, tok, + .{ .i = Pool.shift(ratio) }, + }); + return .{ + .cv = try e.view(recv.t.role, .u8, &.{ W * sent, 576 }, &.{ Pool.row_bytes, 1 }, 0), + .cs = try e.view(recv.t.role, .u8, &.{ W * sent, 8 }, &.{ Pool.row_bytes, 1 }, 576), + .tok = tok, + }; +} + +/// Triton `_rms` (rmsnorm.rms: a weight, bf16 in and out): `r1`'s kv_norm beside the SWA store. +pub fn rms(e: *Emitter, x: Arg, w: Arg, out: Arg, k: i64) !void { + try e.triton("_rms", .{ e.n, 1, 1 }, &.{ + .{ .name = "X", .arg = x }, .{ .name = "x_rs", .arg = .{ .i = k } }, + .{ .name = "W", .arg = w }, .{ .name = "OUT", .arg = out }, + .{ .name = "o_rs", .arg = .{ .i = k } }, .{ .name = "K", .arg = .{ .i = k } }, + .{ .name = "inv_k", .arg = .{ .f = 1.0 / @as(f64, @floatFromInt(k)) } }, .{ .name = "eps", .arg = .{ .f = dec(e.cfg.eps) } }, + .{ .name = "BK", .arg = .{ .i = 512 } }, .{ .name = "HAS_W", .arg = .{ .b = true } }, + .{ .name = "NARROW", .arg = .{ .b = true } }, .{ .name = "PDL", .arg = .{ .b = false } }, + }); +} + +/// Only the 512-wide KV producer belongs to this fusion; other RMS norms keep their calls. +pub fn kvNormStoreOn(e: *const Emitter) bool { + return e.o.kv_norm_store and e.cfg.head_dim == 512 and e.n >= 1 and e.n <= 64; +} + +pub fn normStore(e: *Emitter, x: Arg, w: Arg, v: Arg, s: Arg, p: Arg, sl: Arg, ring: i64, rows: bool) !void { + try e.ext("tf_dsv41_kv_glue_v1.norm_store", &.{ + x, w, try e.ropeTable(), v, s, p, sl, + .{ .i = 512 }, .{ .i = 64 }, .{ .i = 576 }, .{ .i = 8 }, .{ .f = 1.0 / 512.0 }, .{ .f = dec(e.cfg.eps) }, .{ .i = ring }, + .{ .b = rows }, + }); +} + +/// The SWA store of the window's kv rows (`r1`: kv_norm runs here, RMS_FOLD's `Normed.tensor()`). +pub fn swaStore(e: *Emitter, L: u32, kv: Arg, kvn: Arg) !void { + if (e.o.r1 and e.kvNormStoreOn()) { + const rows = try e.pool(L, 0); + try e.normStore(kv, try e.weight("L{d}.attn.kv_norm", .{L}, .f32, &.{e.cfg.head_dim}), rows[0], rows[1], try e.pos(), .none, 2 * @as(i64, e.cfg.window), false); + } else { + if (e.o.r1) try e.rms(kv, try e.weight("L{d}.attn.kv_norm", .{L}, .f32, &.{e.cfg.head_dim}), kvn, e.cfg.head_dim); + try e.kvStore(kvn, L, 0); + } +} + +/// branches.plan's marks (`fork`, `onMain`, `join`); nothing unless TF_DSV41_BRANCHES branches this window +pub fn branching(e: *const Emitter) bool { + return e.o.branches and e.n <= e.o.branch_rows; +} + +pub fn fork(e: *Emitter) void { + if (!e.branching()) return; + e.br_fork = true; + e.br_side = true; +} + +pub fn onMain(e: *Emitter) void { + e.br_side = false; +} + +pub fn join(e: *Emitter) void { + if (!e.branching()) return; + e.br_join = true; + e.br_side = false; +} + +pub fn attention(e: *Emitter, L: u32) !void { + const n = e.n; + const D: i64 = e.cfg.hidden; + const md = e.cfg.mode(L); + const ratio: i64 = e.cfg.compressRatio(L); + const kv_src = e.cfg.kvSource(L); + const H: i64 = e.cfg.heads / e.o.world; + const hd: i64 = e.cfg.head_dim; + const IH: i64 = e.cfg.index_heads; + const ID: i64 = e.cfg.index_dim; + const x = try e.buf("w.out", .{}, .bf16, &.{ n, D }); + const indexer = md == .full or md == .reindex; + const own_kv = md == .full; + // the indexer's head weights (a bf16 GEMV over the normed input) + const iw = try e.buf("L.ix_w", .{}, .bf16, &.{ n, IH }); + // index layer: fork, side indexer chain, main q group, join, attn; other layers: side SWA store, main q group + if (indexer) e.fork(); + if (indexer) try e.triton("_plain", .{ n, IH, 1 }, &.{ + .{ .name = "X", .arg = x }, .{ .name = "x_rs", .arg = .{ .i = D } }, + .{ .name = "W", .arg = try e.weight("L{d}.attn.ix_wp", .{L}, .f32, &.{ IH, D }) }, .{ .name = "w_rs", .arg = .{ .i = D } }, + .{ .name = "OUT", .arg = iw }, .{ .name = "K", .arg = .{ .i = D } }, + .{ .name = "N", .arg = .{ .i = IH } }, .{ .name = "BK", .arg = .{ .i = 256 } }, + .{ .name = "PDL", .arg = .{ .b = false } }, .{ .name = "BF16", .arg = .{ .b = true } }, + }); + e.onMain(); + // x group: wq_a, wkv (+ the compressor's on a KV source) + const xg = try e.xGroup(L); + const qa = try e.buf("L.qa", .{}, .bf16, &.{ n, e.cfg.q_lora }); + const kv = try e.buf("L.kv", .{}, .bf16, &.{ n, hd }); + const cw: i64 = if (ratio == 2) 2 * hd else hd; + // a KV source's compressor projection outlives its layer: keep() carries the accepted row (ratio 2) + const comp = if (ratio == 2) try e.buf("w.L{d}.comp", .{L}, .bf16, &.{ n, cw }) else try e.buf("L.comp", .{}, .bf16, &.{ n, cw }); + var ps: [4]Proj = undefined; + for (xg.ws, 0..) |w, i| ps[i] = .{ .w = w, .x = x, .k = @intCast(D), .n = @intCast(xg.ns[i]), .y = switch (i) { + 0 => qa, + 1 => kv, + else => comp, + }, .col = if (i == 3) @intCast(hd) else 0 }; + try e.fused("x", ps[0..xg.ws.len]); + const qn = try e.buf("L.qn", .{}, .bf16, &.{ n, e.cfg.q_lora }); + const kvn = try e.buf("L.kvn", .{}, .bf16, &.{ n, hd }); + // `r1` (RMS_FOLD): no launch here; q_norm runs in its consumers' rot_in, kv_norm beside the SWA store + const qnorm: ?Norm = if (e.o.r1) .{ .w = try e.weight("L{d}.attn.q_norm", .{L}, .f32, &.{e.cfg.q_lora}), .eps = dec(e.cfg.eps) } else null; + if (!e.o.r1) try e.triton("_rms2", .{ n, 2, 1 }, &.{ + .{ .name = "XA", .arg = qa }, .{ .name = "xa_rs", .arg = .{ .i = e.cfg.q_lora } }, + .{ .name = "WA", .arg = try e.weight("L{d}.attn.q_norm", .{L}, .f32, &.{e.cfg.q_lora}) }, .{ .name = "OA", .arg = qn }, + .{ .name = "oa_rs", .arg = .{ .i = e.cfg.q_lora } }, .{ .name = "KA", .arg = .{ .i = e.cfg.q_lora } }, + .{ .name = "inv_ka", .arg = .{ .f = 1.0 / @as(f64, @floatFromInt(e.cfg.q_lora)) } }, .{ .name = "XB", .arg = kv }, + .{ .name = "xb_rs", .arg = .{ .i = hd } }, .{ .name = "WB", .arg = try e.weight("L{d}.attn.kv_norm", .{L}, .f32, &.{hd}) }, + .{ .name = "OB", .arg = kvn }, .{ .name = "ob_rs", .arg = .{ .i = hd } }, + .{ .name = "KB", .arg = .{ .i = hd } }, .{ .name = "inv_kb", .arg = .{ .f = 1.0 / @as(f64, @floatFromInt(hd)) } }, + .{ .name = "eps", .arg = .{ .f = dec(e.cfg.eps) } }, .{ .name = "BK", .arg = .{ .i = 512 } }, + .{ .name = "NARROW", .arg = .{ .b = true } }, .{ .name = "PDL", .arg = .{ .b = false } }, + }); + if (!indexer) { + e.fork(); + try e.swaStore(L, kv, kvn); // `r1` (RMS_FOLD): its kv_norm with it, on the side + e.onMain(); + } + if (indexer) e.fork(); + if (own_kv) { + // the compressor: pooled rows, normed, into the pool; the indexer's keys from the same rows + const lat = try e.buf("L.lat", .{}, .bf16, &.{ n, hd }); + try e.triton("_pool_norm", .{ n, 1, 1 }, &.{ + .{ .name = "BUF", .arg = comp }, .{ .name = "b_stride", .arg = .{ .i = cw } }, + .{ .name = "W", .arg = try e.weight("L{d}.attn.comp_norm", .{L}, .f32, &.{hd}) }, .{ .name = "LAT", .arg = lat }, + .{ .name = "POS", .arg = try e.pos() }, .{ .name = "n", .arg = .{ .i = n } }, + .{ .name = "EPS", .arg = .{ .f = dec(e.cfg.eps) } }, .{ .name = "RATIO", .arg = .{ .i = ratio } }, + .{ .name = "SL", .arg = .none }, .{ .name = "PREV", .arg = .none }, + .{ .name = "OFF", .arg = .{ .i = if (ratio == 2) 1 else 0 } }, .{ .name = "ROWS", .arg = .{ .b = false } }, + .{ .name = "PDL", .arg = .{ .b = false } }, .{ .name = "CARRY", .arg = if (ratio == 2) try e.buf("s.L{d}.carry", .{L}, .f32, &.{ 1, cw }) else .none }, + .{ .name = "c_stride", .arg = .{ .i = if (ratio == 2) cw else 0 } }, .{ .name = "SPLIT", .arg = .{ .b = ratio == 2 } }, + }); + try e.kvStore(lat, L, ratio); + const kp = try e.buf("L.ix_k", .{}, .bf16, &.{ n, ID }); + try e.upstream(try e.fmt("L{d}.attn.ix_wk", .{L}), lat, hd, ID, kp); + const kpg = try e.paging(ratio, false); + try e.triton("_index_k", .{ n, 1, 1 }, &.{ + .{ .name = "KP", .arg = kp }, .{ .name = "k_stride", .arg = .{ .i = ID } }, + .{ .name = "W", .arg = try e.weight("L{d}.attn.ix_knorm", .{L}, .f32, &.{ID}) }, .{ .name = "CS", .arg = try e.ropeTable() }, + .{ .name = "cs_stride", .arg = .{ .i = 64 } }, .{ .name = "IK", .arg = try e.indexKeys(L) }, + .{ .name = "POS", .arg = try e.pos() }, .{ .name = "EPS", .arg = .{ .f = dec(e.cfg.eps) } }, + .{ .name = "RATIO", .arg = .{ .i = ratio } }, .{ .name = "PT", .arg = kpg.pt }, + .{ .name = "PSH", .arg = .{ .i = kpg.psh } }, .{ .name = "KFP8", .arg = .{ .b = true } }, + .{ .name = "SL", .arg = .none }, .{ .name = "PTS", .arg = .{ .i = 0 } }, + .{ .name = "ROWS", .arg = .{ .b = false } }, .{ .name = "PDL", .arg = .{ .b = false } }, + }); + } + if (indexer) try e.index(L, if (e.o.r1) qa else qn, iw, qnorm); + const early: ?Received = if (indexer and e.o.pool != null and e.o.pool.?.split and e.o.pool.?.early and (n <= 16 or e.o.rows)) + try e.exchange(kv_src.?, e.cfg.compressRatio(kv_src.?), (try e.selection(L))[0]) + else + null; + e.onMain(); + // q group: wq_b + const q = try e.buf("L.q", .{}, .bf16, &.{ n, H * hd }); + try e.fusedNorm("q", &.{.{ .w = try e.fmt("L{d}.attn.wq_b", .{L}), .x = if (e.o.r1) qa else qn, .y = q, .k = e.cfg.q_lora, .n = @intCast(H * hd) }}, qnorm); + try e.prefetchO(L); + if (!e.o.r1) try e.rope(q, H, hd, hd - 64); // `r1` (ATTN_ROPE): the attention kernel rotates q as it loads it + if (indexer) try e.swaStore(L, kv, kvn); + e.join(); // the attention core reads both streams' values + // attention over the SWA ring and (a compressed layer) the source's pool at the source's selection + const o = try e.buf("L.o", .{}, .bf16, &.{ n, H, hd }); + const swa = try e.pool(L, 0); + var cv = try e.empty(.u8); + var cs = try e.empty(.u8); + var tok = try e.empty(.i32); + var cnt = try e.empty(.i32); + var apg: Paging = .{ .pt = try e.empty(.i32), .psh = 0, .vs = 0, .ss = 0 }; + if (ratio > 0) { + const src_ratio = e.cfg.compressRatio(kv_src.?); + const sel = try e.selection(e.cfg.indexSource(L).?); + cnt = sel[1]; + if (e.o.pool != null and e.o.pool.?.split) { + // split: an index layer's selection brings its rows over the exchange; its reuse layers read them too + const got = early orelse if (indexer) try e.exchange(kv_src.?, src_ratio, sel[0]) else try e.received(kv_src.?, src_ratio); + cv = got.cv; + cs = got.cs; + tok = got.tok; + } else { + const p = try e.pool(kv_src.?, src_ratio); + cv = p[0]; + cs = p[1]; + tok = sel[0]; + if (e.o.pool != null) { + const pg = try e.paging(src_ratio, false); + apg = .{ .pt = pg.pt, .psh = pg.psh, .vs = 0, .ss = 0 }; + } + } + } + const rows_scr = n * 5 * H; + try e.ext("tf_dsv41_attn_cuda_v1.attn", &.{ + try e.view(q.t.role, .bf16, &.{ n, H, hd }, &.{ H * hd, hd, 1 }, 0), cv, cs, tok, cnt, swa[0], swa[1], + try e.buf("w.lo", .{}, .i32, &.{n}), try e.empty(.i32), try e.pos(), try e.empty(.i64), apg.pt, try e.weight("L{d}.attn.sink", .{L}, .f32, &.{H}), try e.ropeTable(), + try e.buf("s.attn.po", .{}, .f32, &.{rows_scr * 512}), try e.buf("s.attn.pm", .{}, .f32, &.{rows_scr}), try e.buf("s.attn.pl", .{}, .f32, &.{rows_scr}), o, try e.buf("s.attn.ticket", .{}, .i32, &.{e.ticketLen()}), .{ .i = 2 * @as(i64, e.cfg.window) }, .{ .i = apg.psh }, + .{ .i = 0 }, .{ .i = 4 }, + }); + if (e.o.r1) { + // ATTN_SPLIT 4 (a cluster of 4 a chunk up to ATTN_SPLIT_ROWS = 2 rows, 1 above) and ATTN_ROPE + const last = &e.out.items[e.out.items.len - 1]; + const more = try e.a.alloc(calls.Named, last.args.len + 2); + @memcpy(more[0..last.args.len], last.args); + more[last.args.len] = .{ .arg = .{ .i = attnSplit(n) } }; + more[last.args.len + 1] = .{ .arg = .{ .i = 1 } }; + last.args = more; + } else if (e.o.r1_sig) try e.appendLast(&.{ .{ .i = 0 }, .{ .i = 0 } }); + // o: wo_a's groups over the heads' columns, then wo_b (this rank's partial) + const g: usize = e.cfg.o_groups / e.o.world; + const gk: i64 = @divExact(H * hd, @as(i64, @intCast(g))); + const oa = try e.buf("L.oa", .{}, .bf16, &.{ n, @as(i64, @intCast(g)) * e.cfg.o_lora }); + var po: [8]Proj = undefined; + for (po[0..g], 0..) |*p, i| p.* = .{ + .w = try e.fmt("L{d}.attn.wo_a.{d}", .{ L, i }), + .x = try e.view(o.t.role, .bf16, &.{ n, gk }, &.{ H * hd, 1 }, @as(i64, @intCast(i)) * gk * 2), + .y = oa, + .col = i * e.cfg.o_lora, + .k = @intCast(gk), + .n = e.cfg.o_lora, + }; + try e.fused("o", po[0..g]); + try e.fused("wo_b", &.{.{ .w = try e.fmt("L{d}.attn.wo_b", .{L}), .x = oa, .y = try e.buf("L.part", .{}, .bf16, &.{ n, D }), .k = g * e.cfg.o_lora, .n = @intCast(D) }}); +} diff --git a/zig/src/families/deepseek_v41/block_calls.zig b/zig/src/families/deepseek_v41/block_calls.zig new file mode 100644 index 0000000000..99765cf940 --- /dev/null +++ b/zig/src/families/deepseek_v41/block_calls.zig @@ -0,0 +1,257 @@ +//! A block's launches as the forward issues them, tensors named by role, and the check of each against a captured op. + +const std = @import("std"); +const Value = std.json.Value; + +pub const Dt = enum { bf16, f16, f32, f64, i8, u8, i16, i32, i64, bool }; + +/// Where a tensor's bytes live: one of our loader's weights, a named buffer of the forward, or nothing (numel 0). +pub const Role = union(enum) { + weight: []const u8, + /// scope by prefix: "s." across windows (slot state, tables), "L." its layer, anything else the window + buf: []const u8, + empty, +}; + +pub const Tensor = struct { + role: Role, + dt: Dt, + shape: []const i64, + stride: []const i64, + /// bytes from the role's first byte + offset: i64 = 0, +}; + +pub const Arg = union(enum) { + t: Tensor, + i: i64, + f: f64, + b: bool, + none, + list: []const Arg, + /// the capture's table contents are not compared (an L2 prefetch list of addresses) + opaque_table: Tensor, +}; + +pub const Named = struct { name: []const u8 = "", arg: Arg }; + +/// Where a call starts a scope: the window's or a layer's roles begin fresh there. +pub const Begin = enum { none, window, layer }; + +pub const Call = struct { + triton: bool, + glue: bool = false, // the forward's own step (positions, exchanges, Engram rows): no captured launch + begin: Begin = .none, + /// the compiled kernel's name (Triton) or "." + name: []const u8, + grid: [3]i64 = .{ 1, 1, 1 }, + args: []const Named, + side: bool = false, // TF_DSV41_BRANCHES: issued on the side stream (same launch; checks compare launches) + fork: bool = false, // the side stream first waits for every main-stream launch so far + join: bool = false, // the main stream first waits for the side stream + mark: u8 = 0, // TF_DSV41_PF_OVERLAP_SITE: on a side call, record side mark m after it + wait: u8 = 0, // on a main call, wait for mark m only (a `join` still waits for all) + defer_side: bool = false, // TF_DSV41_MHC_DEFER: an mHC coefficient launch on the deferred stream + defer_join: bool = false, // the main stream waits for it first (the coefficients' reader) +}; + +/// TF_DSV41_MHC_DEFER_AT=exchange: move each deferred launch to just before the next exchange or its join, in-region. +pub fn holdDeferred(cs: []Call) void { + var i: usize = 0; + while (i < cs.len) : (i += 1) { + const c = cs[i]; + if (!c.defer_side or c.defer_join or c.side or c.fork or c.join or c.begin != .none) continue; + var j = i + 1; + while (j < cs.len and !cs[j].defer_join and !cs[j].defer_side and !isExchange(cs[j])) : (j += 1) {} + if (j == cs.len or cs[j].defer_side) continue; // no exchange or join after it: left where it was + std.mem.copyForwards(Call, cs[i .. j - 1], cs[i + 1 .. j]); + cs[j - 1] = c; + i = j - 1; + } +} + +/// A collective exchange step of the forward (glue.exchange / glue.exchange_f32). +pub fn isExchange(c: Call) bool { + return c.glue and std.mem.startsWith(u8, c.name, "glue.exchange"); +} + +/// The calls a capture records (the glue steps left out), in order; a dropped step's scope start moves to the next. +pub fn launches(a: std.mem.Allocator, cs: []const Call) ![]const Call { + var out: std.ArrayList(Call) = .empty; + var pending: Begin = .none; + for (cs) |c| { + if (c.glue) { + if (@backingInt(c.begin) > @backingInt(pending)) pending = c.begin; + continue; + } + var x = c; + if (pending == .window or (pending == .layer and x.begin == .none)) x.begin = pending; + pending = .none; + try out.append(a, x); + } + return out.items; +} + +/// Contiguous strides of `shape` (elements). +pub fn contiguous(a: std.mem.Allocator, shape: []const i64) ![]const i64 { + const s = try a.alloc(i64, shape.len); + var acc: i64 = 1; + var i = shape.len; + while (i > 0) { + i -= 1; + s[i] = acc; + acc *= @max(shape[i], 1); + } + return s; +} + +fn dtName(dt: Dt) []const u8 { + return switch (dt) { + .bf16 => "bfloat16", + .f16 => "float16", + .f32 => "float32", + .f64 => "float64", + .i8 => "int8", + .u8 => "uint8", + .i16 => "int16", + .i32 => "int32", + .i64 => "int64", + .bool => "bool", + }; +} + +/// Role -> captured storage key within its scope, learned as the check goes; another key for a role is a mismatch. +pub const Roles = struct { + a: std.mem.Allocator, + keys: std.StringHashMapUnmanaged([]const u8) = .empty, + /// storage size a role maps to (the GPU run sizes its buffer from it) + sizes: std.StringHashMapUnmanaged(u64) = .empty, + + /// Forgets the roles whose scope ends at `b`: "L." at a layer, "L." and "w." at a window. + pub fn begin(r: *Roles, b: Begin) void { + switch (b) { + .none => {}, + .layer => r.endScope("w."), + .window => r.endScope(""), + } + } + + /// Forgets every role but the persistent ("s.") ones and those starting with `prefix_keep`. + pub fn endScope(r: *Roles, prefix_keep: []const u8) void { + var it = r.keys.iterator(); + var drop: std.ArrayList([]const u8) = .empty; + defer drop.deinit(r.a); + while (it.next()) |e| { + const k = e.key_ptr.*; + if (std.mem.startsWith(u8, k, "s.")) continue; + if (prefix_keep.len > 0 and std.mem.startsWith(u8, k, prefix_keep)) continue; + drop.append(r.a, k) catch {}; + } + for (drop.items) |k| _ = r.keys.remove(k); + } +}; + +pub const Mismatch = struct { what: []const u8 }; + +/// The first difference between `c` and the captured op `op` (null: they agree); learns the roles' storage keys. +pub fn check(a: std.mem.Allocator, c: *const Call, op: std.json.ObjectMap, roles: *Roles) !?[]const u8 { + const name = op.get("name").?.string; + if (!std.mem.eql(u8, name, c.name)) return try std.fmt.allocPrint(a, "name {s} vs captured {s}", .{ c.name, name }); + if (c.triton) { + const g = op.get("grid").?.array.items; + for (c.grid, 0..) |x, i| if (g[i].integer != x) return try std.fmt.allocPrint(a, "grid {any} vs captured [{d},{d},{d}]", .{ c.grid, g[0].integer, g[1].integer, g[2].integer }); + } + const got = op.get("args").?.array.items; + if (got.len != c.args.len) return try std.fmt.allocPrint(a, "{d} args vs captured {d}", .{ c.args.len, got.len }); + var keys: std.StringHashMapUnmanaged(i64) = .empty; // this op's role -> buffer id (aliasing within the op) + const bufs = op.get("buffers").?.array.items; + for (c.args, got, 0..) |want, cap, i| { + if (c.triton) { + const cn = cap.object.get("name").?.string; + if (!std.mem.eql(u8, cn, want.name)) return try std.fmt.allocPrint(a, "arg {d} named {s} vs captured {s}", .{ i, want.name, cn }); + } + if (try argCheck(a, want.arg, cap, roles, &keys, bufs)) |why| return try std.fmt.allocPrint(a, "arg {d} ({s}): {s}", .{ i, if (want.name.len > 0) want.name else "-", why }); + } + return null; +} + +fn argCheck(a: std.mem.Allocator, want: Arg, cap: Value, roles: *Roles, keys: *std.StringHashMapUnmanaged(i64), bufs: []const Value) !?[]const u8 { + const o = cap.object; + const t = o.get("t").?.string; + switch (want) { + .none => return if (std.mem.eql(u8, t, "none")) null else try std.fmt.allocPrint(a, "none vs {s}", .{t}), + .b => |v| { + if (!std.mem.eql(u8, t, "bool")) return try std.fmt.allocPrint(a, "bool vs {s}", .{t}); + return if (o.get("v").?.bool == v) null else try std.fmt.allocPrint(a, "{} vs {}", .{ v, o.get("v").?.bool }); + }, + .i => |v| { + if (!std.mem.eql(u8, t, "int")) return try std.fmt.allocPrint(a, "int {d} vs {s}", .{ v, t }); + return if (o.get("v").?.integer == v) null else try std.fmt.allocPrint(a, "{d} vs {d}", .{ v, o.get("v").?.integer }); + }, + .f => |v| { + if (!std.mem.eql(u8, t, "float")) return try std.fmt.allocPrint(a, "float vs {s}", .{t}); + const bits = try std.fmt.parseInt(u64, o.get("f64").?.string[2..], 16); + return if (bits == @as(u64, @bitCast(v))) null else try std.fmt.allocPrint(a, "{e} vs {e}", .{ v, @as(f64, @bitCast(bits)) }); + }, + .list => |items| { + if (!std.mem.eql(u8, t, "list")) return try std.fmt.allocPrint(a, "list vs {s}", .{t}); + const ci = o.get("items").?.array.items; + if (ci.len != items.len) return try std.fmt.allocPrint(a, "list of {d} vs {d}", .{ items.len, ci.len }); + for (items, ci, 0..) |x, y, j| if (try argCheck(a, x, y, roles, keys, bufs)) |why| return try std.fmt.allocPrint(a, "[{d}] {s}", .{ j, why }); + return null; + }, + .t, .opaque_table => |w| { + if (!std.mem.eql(u8, t, "tensor")) return try std.fmt.allocPrint(a, "tensor vs {s}", .{t}); + if (!std.mem.eql(u8, o.get("dtype").?.string, dtName(w.dt))) return try std.fmt.allocPrint(a, "dtype {s} vs {s}", .{ dtName(w.dt), o.get("dtype").?.string }); + const shape = o.get("shape").?.array.items; + if (!intsEq(w.shape, shape)) return try std.fmt.allocPrint(a, "shape {any} vs captured", .{w.shape}); + var numel: i64 = 1; + for (w.shape) |d| numel *= d; + if (numel == 0) return null; // an empty tensor: only its dtype and shape count + if (w.shape.len > 0 and !intsEq(w.stride, o.get("stride").?.array.items)) return try std.fmt.allocPrint(a, "stride {any} vs captured", .{w.stride}); + switch (w.role) { + .empty => return try std.fmt.allocPrint(a, "empty role for a tensor of {d} elements", .{numel}), + .weight => |name| { + const cw = o.get("weight") orelse return try std.fmt.allocPrint(a, "weight {s} vs a captured buffer", .{name}); + if (!std.mem.eql(u8, cw.string, name)) return try std.fmt.allocPrint(a, "weight {s} vs {s}", .{ name, cw.string }); + return if (o.get("offset").?.integer == w.offset) null else try std.fmt.allocPrint(a, "offset {d} vs {d}", .{ w.offset, o.get("offset").?.integer }); + }, + .buf => |role| { + if (o.get("weight")) |cw| return try std.fmt.allocPrint(a, "buffer {s} vs captured weight {s}", .{ role, cw.string }); + if (o.get("offset").?.integer != w.offset) return try std.fmt.allocPrint(a, "{s} offset {d} vs {d}", .{ role, w.offset, o.get("offset").?.integer }); + const id = o.get("buf").?.integer; + // within the op: the same role is the same buffer and two roles are two buffers + if (keys.get(role)) |prev| { + if (prev != id) return try std.fmt.allocPrint(a, "{s} is two captured buffers", .{role}); + } else { + var it = keys.iterator(); + while (it.next()) |e| if (e.value_ptr.* == id) return try std.fmt.allocPrint(a, "{s} and {s} are one captured buffer", .{ role, e.key_ptr.* }); + try keys.put(a, role, id); + } + const key = for (bufs) |b| { + if (b.object.get("id").?.integer == id) break b.object.get("key").?.string; + } else return "no buffer entry"; + const gop = try roles.keys.getOrPut(roles.a, role); + if (gop.found_existing) { + if (!std.mem.eql(u8, gop.value_ptr.*, key)) return try std.fmt.allocPrint(a, "{s} moved to another captured storage", .{role}); + } else { + gop.key_ptr.* = try roles.a.dupe(u8, role); + gop.value_ptr.* = try roles.a.dupe(u8, key); + } + const size: u64 = @intCast(for (bufs) |b| { + if (b.object.get("id").?.integer == id) break b.object.get("nbytes").?.integer; + } else 0); + const sg = try roles.sizes.getOrPut(roles.a, gop.key_ptr.*); + if (!sg.found_existing or sg.value_ptr.* < size) sg.value_ptr.* = size; + return null; + }, + } + }, + } +} + +fn intsEq(w: []const i64, c: []const Value) bool { + if (w.len != c.len) return false; + for (w, c) |x, y| if (y.integer != x) return false; + return true; +} diff --git a/zig/src/families/deepseek_v41/block_index.zig b/zig/src/families/deepseek_v41/block_index.zig new file mode 100644 index 0000000000..3e5019c282 --- /dev/null +++ b/zig/src/families/deepseek_v41/block_index.zig @@ -0,0 +1,224 @@ +//! A decode window's indexer: index scores (whole, row-blocked), the selections, attn_cuda's top-k. + +const std = @import("std"); +const calls = @import("block_calls.zig"); +const ft = @import("forward_types.zig"); +const block = @import("block.zig"); +const Arg = calls.Arg; +const Emitter = block.Emitter; +const Norm = block.Emitter.Norm; + +/// attn_cuda's MAX_EPT: past 8 CTAs x 512 threads x 31 entries a selection runs the Triton dtopk instead. +const topk_max_ept: i64 = 31; + +/// topk.plan: (cluster CTAs, entries a thread) for `nk` entries. +fn topkPlan(nk: i64, cl_in: ?i64) [2]i64 { + const cl = cl_in orelse @max(1, @min(8, std.math.divCeil(i64, nk, 4096) catch 1)); + var ept = @max(1, std.math.divCeil(i64, nk, cl * 512) catch 1); + ept += 1 - @mod(ept, 2); + return .{ cl, ept }; +} + +/// The window's index selection of layer `src` (reuse layers read the source's). +pub fn selection(e: *Emitter, src: u32) ![2]Arg { + const k: i64 = e.cfg.index_topk; + return .{ + try e.view(.{ .buf = try e.fmt("w.sel{d}", .{src}) }, .i32, &.{ e.n, k }, &.{ k, 1 }, e.row0 * k * 4), + // the counts' 4-byte rows: a run's segment from a row a multiple of 4 (cnt0), so the slice is 16-aligned + try e.view(.{ .buf = try e.fmt("w.cnt{d}", .{src}) }, .i32, &.{e.n}, &.{1}, e.cnt0 * 4), + }; +} + +/// attn_cuda's ticket words: a (row, head tile)'s ranks, 4 with the split kernel (`r1` or `r1_sig`) +pub fn ticketLen(e: *const Emitter) i64 { + return if (e.o.r1 or e.o.r1_sig) 1024 else 256; +} + +/// attn_cuda.split: CTAs a chunk with `r1` (ATTN_SPLIT 4 up to ATTN_SPLIT_ROWS = 2 rows, 1 above) +pub fn attnSplit(n: i64) i64 { + return if (n <= 2) 4 else 1; +} + +/// The indexer: query projection and rope, scores over the keys, top-k (and a candidate source's blocks). +pub fn index(e: *Emitter, L: u32, qn: Arg, iw: Arg, qnorm: ?Norm) !void { + const n = e.n; + const IH: i64 = e.cfg.index_heads; + const ID: i64 = e.cfg.index_dim; + const qi = try e.buf("L.qi", .{}, .bf16, &.{ n, IH * ID }); + try e.fusedNorm("qi", &.{.{ .w = try e.fmt("L{d}.attn.ix_wq_b", .{L}), .x = qn, .y = qi, .k = e.cfg.q_lora, .n = @intCast(IH * ID) }}, qnorm); + try e.rope(qi, IH, ID, ID - 64); + const reindex = e.cfg.mode(L) == .reindex; + const src = if (reindex) e.cfg.kvSource(L).? else L; + const ratio: i64 = e.cfg.compressRatio(src); + const cbs: i64 = e.cfg.candidate_block_size; + const nk: i64 = if (reindex) @as(i64, e.cfg.candidate_blocks) * cbs else @divFloor(e.end(), ratio); + const cand = try e.buf("w.cand", .{}, .i32, &.{ n, e.cfg.candidate_blocks }); + // the materialised scores + sort keys within the index budget, else row blocks through `_keys` + if (reindex) { + if (n * nk * 16 > e.o.index_budget) return error.Unsupported; // index.reindex in row blocks: > 256 rows at 64 MiB + } else { + const nvis = @max(nk, 1); + if (!e.o.rows and n > 16 and nvis >= ft.stream_min and !e.cfg.isCandidateSource(L) and @popCount(e.cfg.index_topk) == 1) return error.Unsupported; + if (n * nvis * 12 > e.o.index_budget) return e.blocked(L, qi, iw, src, ratio, nvis); + } + const sc = try e.buf("L.scores", .{}, .f32, &.{ n, nk }); + try e.scoreRows(qi, iw, src, ratio, 0, n, try e.pos(), nk, sc, if (reindex) cand else null); + try e.topkOf(L, sc, cand, nk, ratio, e.o.index_bound); +} + +/// index.scores over rows [a, a + rows) into `out` [rows, nk]: dense, or over candidate blocks `cand` (Reindex). +pub fn scoreRows(e: *Emitter, qi: Arg, iw: Arg, src: u32, ratio: i64, a: i64, rows: i64, p: Arg, nk: i64, out: Arg, cand: ?Arg) !void { + const IH: i64 = e.cfg.index_heads; + const ID: i64 = e.cfg.index_dim; + const cbs: i64 = e.cfg.candidate_block_size; + const spg = try e.paging(ratio, false); + try e.triton(if (e.o.index_bound and cand == null) "_scores_b" else "_scores", .{ rows, std.math.divCeil(i64, nk, 64) catch unreachable, 1 }, &.{ + .{ .name = "QI", .arg = try e.view(qi.t.role, .bf16, &.{ rows, IH, ID }, &.{ IH * ID, ID, 1 }, a * IH * ID * 2) }, + .{ .name = "W", .arg = try e.view(iw.t.role, .bf16, &.{ rows, IH }, &.{ IH, 1 }, a * IH * 2) }, + .{ .name = "w_stride", .arg = .{ .i = IH } }, + .{ .name = "IK", .arg = try e.indexKeys(src) }, + .{ .name = "OUT", .arg = out }, + .{ .name = "POS", .arg = p }, + .{ .name = "KEYS", .arg = cand orelse out }, + .{ .name = "k_stride", .arg = .{ .i = if (cand != null) e.cfg.candidate_blocks else nk } }, + .{ .name = "NK", .arg = .{ .i = nk } }, + .{ .name = "o_stride", .arg = .{ .i = nk } }, + .{ .name = "RATIO", .arg = .{ .i = ratio } }, + .{ .name = "H", .arg = .{ .i = IH } }, + .{ .name = "D", .arg = .{ .i = ID } }, + .{ .name = "BP", .arg = .{ .i = 64 } }, + .{ .name = "WS", .arg = .{ .f = 1.0 / @sqrt(@as(f64, @floatFromInt(IH))) } }, + .{ .name = "SCALE", .arg = .{ .f = 1.0 / @sqrt(@as(f64, @floatFromInt(ID))) } }, + .{ .name = "GATHER", .arg = .{ .b = cand != null } }, + .{ .name = "PT", .arg = spg.pt }, + .{ .name = "PSH", .arg = .{ .i = spg.psh } }, + .{ .name = "KFP8", .arg = .{ .b = true } }, + .{ .name = "SL", .arg = .none }, + .{ .name = "PTS", .arg = .{ .i = 0 } }, + .{ .name = "ROWS", .arg = .{ .b = false } }, + .{ .name = "CBS", .arg = .{ .i = if (cand != null) cbs else 0 } }, + }); +} + +/// Row-blocked selection past the index budget: per block `_scores`, `_keys`, top positions; then counts. +pub fn blocked(e: *Emitter, L: u32, qi: Arg, iw: Arg, src: u32, ratio: i64, nvis: i64) !void { + const n = e.n; + const k: i64 = e.cfg.index_topk; + const cb: i64 = e.cfg.candidate_blocks; + const cbs: i64 = e.cfg.candidate_block_size; + const sel = try e.selection(L); + const cand_src = e.cfg.isCandidateSource(L); + const cand = try e.buf("w.cand", .{}, .i32, &.{ n, cb }); + const st: i64 = if (n > 16) @max(16, @divFloor(e.o.index_budget, nvis * 12)) else n; + var a: i64 = 0; + var j: usize = 0; + while (a < n) : ({ + a += st; + j += 1; + }) { + const rows = @min(st, n - a); + const p = if (st >= n) try e.pos() else try e.blockPos(j, a, rows); + const nk = if (st >= n or e.o.rows) nvis else @max(@divFloor(e.start + a + rows, ratio), 1); + const sc = try e.buf("L.scores", .{}, .f32, &.{ rows, nk }); + try e.scoreRows(qi, iw, src, ratio, a, rows, p, nk, sc, null); + const keys = try e.buf("L.ix.keys", .{}, .i64, &.{ rows, nk }); + try e.triton("_keys", .{ rows, std.math.divCeil(i64, nk, 1024) catch unreachable, 1 }, &.{ + .{ .name = "S", .arg = sc }, .{ .name = "s_stride", .arg = .{ .i = nk } }, + .{ .name = "K", .arg = keys }, .{ .name = "k_stride", .arg = .{ .i = nk } }, + .{ .name = "POSN", .arg = sc }, .{ .name = "p_stride", .arg = .{ .i = nk } }, + .{ .name = "NK", .arg = .{ .i = nk } }, .{ .name = "BLOCK", .arg = .{ .i = 1024 } }, + .{ .name = "GATHER", .arg = .{ .b = false } }, + }); + try e.glue("topk", &.{ keys, try e.view(sel[0].t.role, .i32, &.{ rows, k }, &.{ k, 1 }, a * k * 4), .{ .i = k } }); + if (!cand_src) continue; + const nb = std.math.divCeil(i64, nk, cbs) catch unreachable; + const bk = try e.buf("L.ix.bkeys", .{}, .i64, &.{ rows, nb }); + try e.triton("_block_keys", .{ rows, std.math.divCeil(i64, nb, 256) catch unreachable, 1 }, &.{ + .{ .name = "S", .arg = sc }, .{ .name = "s_stride", .arg = .{ .i = nk } }, + .{ .name = "K", .arg = bk }, .{ .name = "k_stride", .arg = .{ .i = nb } }, + .{ .name = "POS", .arg = p }, .{ .name = "NB", .arg = .{ .i = nb } }, + .{ .name = "RATIO", .arg = .{ .i = ratio } }, .{ .name = "BS", .arg = .{ .i = cbs } }, + .{ .name = "TB", .arg = .{ .i = 256 } }, .{ .name = "ROWS", .arg = .{ .b = false } }, + }); + try e.glue("topk", &.{ bk, try e.view(cand.t.role, .i32, &.{ rows, cb }, &.{ cb, 1 }, a * cb * 4), .{ .i = cb } }); + } + try e.glue("counts", &.{ sel[0], sel[1], try e.buf("w.pos64", .{}, .i64, &.{n}), .{ .i = ratio } }); +} + +/// Win.sub's first position (int32 [1]) of row block [a, a + rows); row mode reads the block's table rows instead. +pub fn blockPos(e: *Emitter, j: usize, a: i64, rows: i64) !Arg { + const p = try e.buf("L.ix.pos{d}", .{j}, .i32, &.{1}); + try e.glue("block_pos", &.{ p, .{ .i = e.start + a }, .{ .i = a }, .{ .i = rows } }); + return p; +} + +/// attn_cuda topk over `scores` [n, nk] of index layer L (up to two jobs a launch); candidate blocks into `cand`. +pub fn topk(e: *Emitter, L: u32, scores: Arg, cand: Arg, nk: i64, ratio: i64) !void { + return e.topkOf(L, scores, cand, nk, ratio, false); +} + +/// `topk`; `bound` (TF_DSV41_INDEX_BOUND): dense select and candidate blocks as the row-bounded twins, same bytes. +pub fn topkOf(e: *Emitter, L: u32, scores: Arg, cand: Arg, nk: i64, ratio: i64, bound: bool) !void { + const n = e.n; + const dtopk = if (bound) "_dtopk_b" else "_dtopk"; + const topk_ext = if (bound and e.cfg.mode(L) != .reindex) "tf_dsv41_topk_b_v1.topk" else "tf_dsv41_attn_cuda_v1.topk"; + const reindex = e.cfg.mode(L) == .reindex; + const cbs: i64 = e.cfg.candidate_block_size; + const sel = try e.selection(L); + const k: i64 = e.cfg.index_topk; + const pl = topkPlan(nk, null); + if (pl[1] > topk_max_ept) { + // past attn_cuda's plan limit (8 x 512 x 31 entries): Triton _dtopk with counts, then candidate blocks + if (reindex) return error.Unsupported; + try e.triton(dtopk, .{ n, 1, 1 }, &.{ + .{ .name = "S", .arg = scores }, .{ .name = "s_stride", .arg = .{ .i = scores.t.stride[0] } }, + .{ .name = "P", .arg = scores }, .{ .name = "p_stride", .arg = .{ .i = 0 } }, + .{ .name = "POS", .arg = try e.pos() }, .{ .name = "NK", .arg = .{ .i = nk } }, + .{ .name = "OUT", .arg = sel[0] }, .{ .name = "o_stride", .arg = .{ .i = k } }, + .{ .name = "CNT", .arg = sel[1] }, .{ .name = "RATIO", .arg = .{ .i = ratio } }, + .{ .name = "K", .arg = .{ .i = k } }, .{ .name = "MODE", .arg = .{ .i = 0 } }, + .{ .name = "ROWS", .arg = .{ .b = false } }, .{ .name = "T", .arg = .{ .i = 4096 } }, + .{ .name = "BS", .arg = .{ .i = 8 } }, .{ .name = "RB", .arg = .{ .i = 8 } }, + .{ .name = "NPASS", .arg = .{ .i = 8 } }, .{ .name = "SORT", .arg = .{ .b = false } }, + .{ .name = "COUNT", .arg = .{ .b = true } }, .{ .name = "KP", .arg = .{ .i = @intCast(std.math.ceilPowerOfTwo(u64, @intCast(k)) catch unreachable) } }, + }); + if (!e.cfg.isCandidateSource(L)) return; + const nb = std.math.divCeil(i64, nk, cbs) catch unreachable; + const bp = topkPlan(nb, null); + if (bp[1] > topk_max_ept) { + // past 8 x 512 x 31 blocks too (positions past ~1,015,808): dtopk.blocks (_dtopk mode 2, no counts) + const cb: i64 = e.cfg.candidate_blocks; + return e.triton(dtopk, .{ n, 1, 1 }, &.{ + .{ .name = "S", .arg = scores }, .{ .name = "s_stride", .arg = .{ .i = scores.t.stride[0] } }, + .{ .name = "P", .arg = scores }, .{ .name = "p_stride", .arg = .{ .i = 0 } }, + .{ .name = "POS", .arg = try e.pos() }, .{ .name = "NK", .arg = .{ .i = nb } }, + .{ .name = "OUT", .arg = cand }, .{ .name = "o_stride", .arg = .{ .i = cand.t.stride[0] } }, + .{ .name = "CNT", .arg = cand }, .{ .name = "RATIO", .arg = .{ .i = ratio } }, + .{ .name = "K", .arg = .{ .i = cb } }, .{ .name = "MODE", .arg = .{ .i = 2 } }, + .{ .name = "ROWS", .arg = .{ .b = false } }, .{ .name = "T", .arg = .{ .i = 4096 } }, + .{ .name = "BS", .arg = .{ .i = cbs } }, .{ .name = "RB", .arg = .{ .i = 8 } }, + .{ .name = "NPASS", .arg = .{ .i = 8 } }, .{ .name = "SORT", .arg = .{ .b = false } }, + .{ .name = "COUNT", .arg = .{ .b = false } }, .{ .name = "KP", .arg = .{ .i = @intCast(std.math.ceilPowerOfTwo(u64, @intCast(cb)) catch unreachable) } }, + }); + } + const job: [8]Arg = .{ scores, try e.empty(.i32), cand, try e.view(cand.t.role, .i32, &.{0}, &.{e.cfg.candidate_blocks}, 0), .{ .i = nb }, .{ .i = e.cfg.candidate_blocks }, .{ .i = 2 }, .{ .i = bp[1] } }; + var one: std.ArrayList(Arg) = .empty; + try one.appendSlice(e.a, &job); + try one.appendSlice(e.a, &job); + try one.appendSlice(e.a, &.{ .{ .i = 1 }, try e.pos(), .{ .i = ratio }, .{ .i = cbs }, .{ .i = n }, .{ .i = bp[0] } }); + return e.ext(topk_ext, one.items); + } + var jobs: [2][8]Arg = undefined; + jobs[0] = .{ scores, if (reindex) cand else try e.empty(.i32), sel[0], sel[1], .{ .i = nk }, .{ .i = k }, .{ .i = if (reindex) 3 else 0 }, .{ .i = pl[1] } }; + const blocks = e.cfg.isCandidateSource(L); + jobs[1] = jobs[0]; + if (blocks) { + const nb = std.math.divCeil(i64, nk, cbs) catch unreachable; + jobs[1] = .{ scores, try e.empty(.i32), cand, try e.view(cand.t.role, .i32, &.{0}, &.{e.cfg.candidate_blocks}, 0), .{ .i = nb }, .{ .i = e.cfg.candidate_blocks }, .{ .i = 2 }, .{ .i = topkPlan(nb, pl[0])[1] } }; + } + var args: std.ArrayList(Arg) = .empty; + try args.appendSlice(e.a, &jobs[0]); + try args.appendSlice(e.a, &jobs[1]); + try args.appendSlice(e.a, &.{ .{ .i = if (blocks) 2 else 1 }, try e.pos(), .{ .i = ratio }, .{ .i = if (blocks or reindex) cbs else 1 }, .{ .i = n }, .{ .i = pl[0] } }); + try e.ext(topk_ext, args.items); +} diff --git a/zig/src/families/deepseek_v41/block_moe.zig b/zig/src/families/deepseek_v41/block_moe.zig new file mode 100644 index 0000000000..34168fa012 --- /dev/null +++ b/zig/src/families/deepseek_v41/block_moe.zig @@ -0,0 +1,319 @@ +//! The decode window's mHC sites, MoE (router, prune, group, x3ld, epilogues), Engram and head, as calls. + +const std = @import("std"); +const calls = @import("block_calls.zig"); +const block = @import("block.zig"); +const E = block.Emitter; +const Arg = calls.Arg; +const wide = @import("block_wide.zig"); + +/// mhc_cuda.run modes: first site (no post), boundary (post, collapse 2, mix), final norm (post, collapse 2, no mix). +pub const Mode = enum(i64) { boundary = 0, site1 = 1, site2 = 2, final = 3 }; + +fn coefs(e: *E, set: u1) ![3]Arg { + return .{ + try e.buf("w.c{d}.pre", .{set}, .f32, &.{ e.n, 4 }), + try e.buf("w.c{d}.post", .{set}, .f32, &.{ e.n, 4 }), + try e.buf("w.c{d}.comb", .{set}, .f32, &.{ e.n, 16 }), + }; +} + +/// One mHC site of layer L (`which`: hc_attn / hc_ffn), or the final norm (L null). +pub fn mhc(e: *E, layer: ?u32, which: []const u8, mode: Mode) !void { + return mhcInto(e, layer, which, mode, false); +} + +/// `mhc`, writing the streams in place when `in_place` (a run without the second streams buffer: CED's replay). +pub fn mhcInto(e: *E, layer: ?u32, which: []const u8, mode: Mode, in_place: bool) !void { + // past 16 rows or with TF_DSV41_MHC_PFDEC at the knob's rows: the wide path (block_wide.zig) + if (e.n > wide.cuda_rows or wide.pfOn(e, mode)) return wide.decode(e, layer, which, mode, in_place); + const n = e.n; + const D: i64 = e.cfg.hidden; + const post = mode == .boundary or mode == .final; + const mix = mode != .final; + const f0 = try e.empty(.f32); + const b0 = try e.empty(.bf16); + const x = try e.streams(); + const prev = try coefs(e, e.coef); + const next = try coefs(e, 1 - e.coef); + var nw: Arg = undefined; + var hc = [3]Arg{ b0, f0, f0 }; + if (layer) |L| { + hc = .{ + try e.buf("s.L{d}.{s}.fn16", .{ L, which }, .bf16, &.{ 6 * @as(i64, e.cfg.hc_mult), 4 * D }), + try e.weight("L{d}.{s}.base", .{ L, which }, .f32, &.{6 * @as(i64, e.cfg.hc_mult)}), + try e.weight("L{d}.{s}.scale", .{ L, which }, .f32, &.{3}), + }; + nw = try e.weight("L{d}.{s}", .{ L, if (std.mem.eql(u8, which, "hc_attn")) "attn_norm" else "ffn_norm" }, .f32, &.{D}); + } else nw = try e.weight("norm", .{}, .f32, &.{D}); + // DSpark's tap: the mean of the new streams at a boundary entering a target layer (not the window's first site) + var tap: Arg = b0; + if (e.o.taps and mode == .boundary and std.mem.eql(u8, which, "hc_attn")) if (layer) |L| { + const t = e.cfg.dspark_targets.items(); + if (std.mem.indexOfScalar(u32, t, L)) |j| { + const w: i64 = @as(i64, @intCast(t.len)) * D; + tap = try e.view(.{ .buf = "w.taps" }, .bf16, &.{ n, D }, &.{ w, 1 }, @as(i64, @intCast(j)) * D * 2); + } + }; + const xout: Arg = switch (mode) { + .site1, .site2 => b0, + .final => x, + .boundary => if (in_place) x else try e.spare(), + }; + awaitCoefs(e); + try e.ext("tf_dsv41_mhc_cuda_v1.run", &.{ + x, xout, + // the exchange that ran last: the attention's partials at the FFN site, the MoE's at the next attention / final + if (post) try e.buf("w.recv.{s}", .{if (std.mem.eql(u8, which, "hc_ffn")) "attn" else "moe"}, .bf16, &.{ e.o.world, n, D }) else b0, if (post) prev[1] else f0, + if (post) prev[2] else f0, if (post or mode == .site2) prev[0] else f0, + hc[0], hc[1], + hc[2], try e.buf("L.mhc.part", .{}, .f32, &.{ n, 4, 40, 32 }), + try e.buf("L.mhc.c", .{}, .bf16, &.{ n, D }), tap, + nw, try e.buf("w.out", .{}, .bf16, &.{ n, D }), + if (mix) next[0] else f0, if (mix) next[1] else f0, + if (mix) next[2] else f0, try e.buf("w.mhc.cnt", .{}, .i32, &.{4}), + .{ .i = n }, .{ .i = @backingInt(mode) }, + .{ .f = E.dec(e.cfg.eps) }, .{ .f = E.dec(e.cfg.hc_eps) }, + .{ .f = E.dec(e.cfg.hc_post_alpha) }, .{ .i = e.cfg.hc_sinkhorn_iters }, + }); + if (e.o.r1) { + // MHC_TAIL (parallel finish) and MHC_DEFER (a mixing site's coefficients as their own launch, joined later) + const last = &e.out.items[e.out.items.len - 1]; + const more = try e.a.alloc(calls.Named, last.args.len + 2); + @memcpy(more[0..last.args.len], last.args); + more[last.args.len] = .{ .arg = .{ .i = e.o.mhc_spin } }; + more[last.args.len + 1] = .{ .arg = .{ .i = @intFromBool(mix) } }; + last.args = more; + if (mix and e.o.mhc_defer) { + e.df_side = true; // MHC_DEFER: on the deferred stream, the next mHC call joins it + e.df_pending = true; + } + if (mix) try e.ext("tf_dsv41_mhc_cuda_v1.coef", &.{ + try e.buf("L.mhc.part", .{}, .f32, &.{ n, 4, 40, 32 }), hc[1], + hc[2], next[0], + next[1], next[2], + .{ .i = n }, .{ .f = E.dec(e.cfg.eps) }, + .{ .f = E.dec(e.cfg.hc_eps) }, .{ .f = E.dec(e.cfg.hc_post_alpha) }, + .{ .i = e.cfg.hc_sinkhorn_iters }, + }); + } else if (e.o.r1_sig) try e.appendLast(&.{ .{ .i = -1 }, .{ .i = 0 } }); // the twin's binding, `r1` off: no tail + if (mix) e.coef = 1 - e.coef; + if (mode == .boundary and !in_place) e.cur = 1 - e.cur; +} + +/// mhc.await_coefs: an mHC call reads the coefficients a deferred launch wrote: the main stream joins it first. +pub fn awaitCoefs(e: *E) void { + if (!e.df_pending) return; + e.df_pending = false; + e.df_join = true; +} + +/// mHC post only: the last exchange's post applied to the streams in place, before a non-first Engram block. +pub fn postOnly(e: *E, out: ?calls.Arg) !void { + const n = e.n; + const D: i64 = e.cfg.hidden; + const x = try e.streams(); + const xout = out orelse x; + const prev = try coefs(e, e.coef); + awaitCoefs(e); + try e.triton("_site", .{ std.math.divCeil(i64, n, 16) catch unreachable, 40, 1 }, &.{ + .{ .name = "X", .arg = x }, .{ .name = "x_stride", .arg = .{ .i = 4 * D } }, + .{ .name = "XOUT", .arg = xout }, .{ .name = "G", .arg = try e.buf("w.recv.moe", .{}, .bf16, &.{ e.o.world, n, D }) }, + .{ .name = "g_rank", .arg = .{ .i = n * D } }, .{ .name = "POST", .arg = prev[1] }, + .{ .name = "COMB", .arg = prev[2] }, .{ .name = "PRE", .arg = x }, + .{ .name = "FN", .arg = x }, .{ .name = "PART", .arg = x }, + .{ .name = "C", .arg = x }, .{ .name = "TAP", .arg = x }, + .{ .name = "tap_stride", .arg = .{ .i = 0 } }, .{ .name = "R", .arg = .{ .i = n } }, + .{ .name = "D", .arg = .{ .i = D } }, .{ .name = "NB", .arg = .{ .i = 40 } }, + .{ .name = "BM", .arg = .{ .i = 16 } }, .{ .name = "BK", .arg = .{ .i = 16 } }, + .{ .name = "WORLD", .arg = .{ .i = e.o.world } }, .{ .name = "POST_ON", .arg = .{ .b = true } }, + .{ .name = "COLLAPSE", .arg = .{ .i = 0 } }, .{ .name = "MIX", .arg = .{ .b = false } }, + .{ .name = "TAP_ON", .arg = .{ .b = false } }, .{ .name = "SPLIT", .arg = .{ .b = false } }, + }); +} + +/// Router GEMV, top-p prune, group, then routed + shared experts: rot_in, x3ld gate/up, epilogue, x3ld down, combine. +pub fn moe(e: *E, L: u32) !void { + const n = e.n; + const D: i64 = e.cfg.hidden; + const backbone = e.cfg.isBackbone(L); + // the router folds kit rounding + grouping into its tail when no prune follows; with `r1` it folds the prune too + const fold = !backbone or e.o.expert_topp == null or e.o.r1; + const ex = e.cfg.expertsOf(L); + const slots: i64 = ex.topk + 1; + const Et: i64 = ex.count + 1; + const I: i64 = e.cfg.expert_width / e.o.world; + const P = n * slots; + const rows: i64 = if (backbone) e.o.expert_rows else e.o.dspark_rows; + const sc: []const u8 = if (backbone) "s.ex" else "s.dx"; + const x = try e.buf("w.out", .{}, .bf16, &.{ n, D }); + const pick = try e.buf("L.pick", .{}, .i32, &.{ n, slots }); + const wts = try e.buf("L.wts", .{}, .f32, &.{ n, slots }); + const uids = try e.buf("{s}.ids", .{sc}, .i32, &.{P}); + const ucount = try e.buf("{s}.count", .{sc}, .i32, &.{1}); + const members = try e.buf("{s}.members", .{sc}, .i32, &.{ P, n }); + // router_gemv.config / row_tile: the narrow kernel's (2, 0) to 16 rows, then the prefill config (block_wide.zig) + const rc = wide.routerCfg(n); + // the route groups in its tail only on the narrow kernel (<= 16 rows) and GE <= 512; else upstream's group follows + const grouped = fold and rc[2] == 0 and Et <= 512; + const scale: f64 = E.dec(e.cfg.routed_scale); + try e.ext("tf_dsv41_router_gemv_v2.route", &.{ + x, try e.weight("L{d}.moe.gate", .{L}, .bf16, &.{ ex.count, D }), + try e.weight("L{d}.moe.bias", .{L}, .f32, &.{ex.count}), pick, + wts, try e.buf("L.logits", .{}, .f32, &.{ n, ex.count }), + try e.buf("s.router.cnt", .{}, .i32, &.{128}), .{ .i = ex.topk }, + .{ .i = slots }, .{ .f = scale }, + .{ .i = rc[0] }, .{ .i = rc[1] }, + .{ .i = rc[2] }, .{ .b = true }, + .{ .b = fold }, if (!grouped) try e.empty(.i32) else uids, + if (!grouped) try e.empty(.i32) else ucount, if (!grouped) try e.t(.empty, .i32, &.{ 0, 1 }) else members, + .{ .i = if (!grouped) 0 else Et }, .{ .b = grouped }, + }); + if (e.o.r1) { + // prune fold args [min_w, topp, scale, min_k, orig, drop, dup] (empty for DSpark); drop = table size + const pa: []const Arg = if (backbone and e.o.expert_topp != null) &.{ + .{ .f = 0.0 }, .{ .f = e.o.expert_topp.? }, .{ .f = scale }, .{ .i = 3 }, .{ .i = 0 }, .{ .i = Et }, .{ .i = 0 }, + } else &.{}; + const last = &e.out.items[e.out.items.len - 1]; + const more = try e.a.alloc(calls.Named, last.args.len + 1); + @memcpy(more[0..last.args.len], last.args); + more[last.args.len] = .{ .arg = .{ .list = try e.a.dupe(Arg, pa) } }; + last.args = more; + } else if (e.o.r1_sig) try e.appendLast(&.{.{ .list = &.{} }}); // the twin's route always takes the prune list + if (!fold) { + try e.triton("_prune", .{ std.math.divCeil(i64, n, 16) catch unreachable, 1, 1 }, &.{ + .{ .name = "PICK", .arg = pick }, .{ .name = "WTS", .arg = wts }, + .{ .name = "R", .arg = .{ .i = n } }, .{ .name = "stride", .arg = .{ .i = slots } }, + .{ .name = "MIN_W", .arg = .{ .f = 0.0 } }, .{ .name = "TOPP", .arg = .{ .f = e.o.expert_topp.? } }, + .{ .name = "SCALE", .arg = .{ .f = scale } }, .{ .name = "DROP", .arg = .{ .i = Et } }, + .{ .name = "MIN_K", .arg = .{ .i = 3 } }, .{ .name = "K", .arg = .{ .i = ex.topk } }, + .{ .name = "KP", .arg = .{ .i = 8 } }, .{ .name = "DUP", .arg = .{ .b = false } }, + .{ .name = "BR", .arg = .{ .i = 16 } }, .{ .name = "ORIG", .arg = .{ .b = false } }, + }); + if (e.o.r1_sig) { // the twin's `_prune` takes KIT (`r1`'s PRUNE_KIT) as its last constexpr + const last = &e.out.items[e.out.items.len - 1]; + const more = try e.a.alloc(calls.Named, last.args.len + 1); + @memcpy(more[0..last.args.len], last.args); + more[last.args.len] = .{ .name = "KIT", .arg = .{ .b = false } }; + last.args = more; + } + try e.glue("kit_weights", &.{ wts, try e.buf("L.wts16", .{}, .f32, &.{ n, slots }) }); + } + const group_rot = !grouped and e.o.router_group_rot and n <= 64 and slots <= 9 and Et <= 512 and D == 5120; + if (!grouped and !group_rot) try e.ext("tensorfold_exl3_experts_v1.group", &.{ pick, uids, ucount, members, .{ .i = n }, .{ .i = slots }, .{ .i = Et } }); + const pre = try e.fmt("s.L{d}.ex", .{L}); + const xg = try e.buf("{s}.xg", .{sc}, .f16, &.{ rows * slots, D }); + const xu = try e.buf("{s}.xu", .{sc}, .f16, &.{ rows * slots, D }); + const xd = try e.buf("{s}.xd", .{sc}, .f16, &.{ rows * slots, I }); + const z = try e.buf("{s}.z", .{sc}, .f32, &.{rows * slots * @max(4 * 2 * I, D)}); + try e.ext(if (group_rot) "tensorfold_exl3_experts_v1.group_rot_in" else "tensorfold_exl3_experts_v1.rot_in", &.{ + x, .{ .i = D }, pick, try e.buf("{s}.suh_g", .{pre}, .f16, &.{ Et, D }), try e.buf("{s}.suh_u", .{pre}, .f16, &.{ Et, D }), + xg, xu, .{ .i = n }, .{ .i = D }, .{ .i = slots }, + .{ .i = Et }, + }); + if (group_rot) try e.appendLast(&.{ uids, ucount, members }); + const r = e.w.experts.get(L) orelse return error.NoWidth; + const gu = try block.k2Range(r[0][0], r[0][1]); + const dn = try block.k2Range(r[1][0], r[1][1]); + const svh_g = try e.buf("{s}.svh_g", .{pre}, .f16, &.{ Et, I }); + const svh_u = try e.buf("{s}.svh_u", .{pre}, .f16, &.{ Et, I }); + const suh_d = try e.buf("{s}.suh_d", .{pre}, .f16, &.{ Et, I }); + const tp_g = try e.buf("{s}.tp_g", .{pre}, .i64, &.{Et}); + const tp_u = try e.buf("{s}.tp_u", .{pre}, .i64, &.{Et}); + const k2_g = try e.buf("{s}.k2_g", .{pre}, .i32, &.{Et}); + const k2_u = try e.buf("{s}.k2_u", .{pre}, .i32, &.{Et}); + const tp_d = try e.buf("{s}.tp_d", .{pre}, .i64, &.{Et}); + const k2_d = try e.buf("{s}.k2_d", .{pre}, .i32, &.{Et}); + const svh_d = try e.buf("{s}.svh_d", .{pre}, .f16, &.{ Et, D }); + const y = try e.buf("{s}.y", .{sc}, .f32, &.{ rows * slots, D }); + // after a prune, kit_weights' fp16 rounding (a pointwise op the forward runs) gives new weights + const cw = if (!fold) try e.buf("L.wts16", .{}, .f32, &.{ n, slots }) else wts; + const out = try e.buf("L.moe", .{}, if (e.o.r1 and backbone) .bf16 else .f32, &.{ n, D }); + const limit: Arg = .{ .f = E.dec(e.cfg.swiglu_limit) }; + if (e.o.x3ld_epi) { + // TF_DSV41_X3LD_EPI: the same x3ld work with the epilogues in its tail behind a last-arrival ticket; same bits + const mt = std.math.divCeil(i64, n, 16) catch unreachable; + const ticket = try e.buf("{s}.epi", .{sc}, .i32, &.{@max(P * mt * @divExact(I, 128), n * @divExact(D, 128))}); + try e.ext("tf_dsv41_x3ld_epi_v1.gateup", &.{ + xg, xu, tp_g, tp_u, k2_g, k2_u, uids, ucount, + members, z, .{ .i = 2 }, .{ .i = D }, .{ .i = I }, .{ .i = P }, .{ .i = 4 }, .{ .i = slots }, + .{ .i = 2 }, .{ .i = 8 }, .{ .i = 1 }, .{ .i = gu[0] }, .{ .i = gu[1] }, pick, svh_g, svh_u, + suh_d, xd, .{ .i = Et }, limit, .{ .i = 1 }, ticket, + }); + try e.ext("tf_dsv41_x3ld_epi_v1.down", &.{ + xd, xd, tp_d, tp_d, k2_d, k2_d, uids, ucount, + members, z, .{ .i = 1 }, .{ .i = I }, .{ .i = D }, .{ .i = P }, .{ .i = 1 }, .{ .i = slots }, + .{ .i = 2 }, .{ .i = 8 }, .{ .i = 1 }, .{ .i = dn[0] }, .{ .i = dn[1] }, pick, svh_d, y, + cw, out, .{ .i = Et }, ticket, + }); + return; + } + try e.ext("tf_dsv41_x3ld_v1.grouped", &.{ + xg, xu, + tp_g, tp_u, + k2_g, k2_u, + uids, ucount, + members, z, + .{ .i = 2 }, .{ .i = D }, + .{ .i = I }, .{ .i = P }, + .{ .i = 4 }, .{ .i = slots }, + .{ .i = 2 }, .{ .i = 8 }, + .{ .i = 1 }, .{ .i = 0 }, + .{ .i = gu[0] }, .{ .i = gu[1] }, + .{ .b = false }, + }); + try e.ext("tensorfold_exl3_experts_v1.gateup_epilogue", &.{ + z, pick, + svh_g, svh_u, + suh_d, xd, + .{ .i = n }, .{ .i = P }, + .{ .i = I }, .{ .i = 4 }, + .{ .i = slots }, .{ .i = Et }, + limit, .{ .i = 1 }, + }); + try e.ext("tf_dsv41_x3ld_v1.grouped", &.{ + xd, xd, tp_d, tp_d, k2_d, k2_d, uids, ucount, + members, z, .{ .i = 1 }, .{ .i = I }, .{ .i = D }, .{ .i = P }, .{ .i = 1 }, .{ .i = slots }, + .{ .i = 2 }, .{ .i = 8 }, .{ .i = 1 }, .{ .i = 0 }, .{ .i = dn[0] }, .{ .i = dn[1] }, .{ .b = false }, + }); + try e.ext("tensorfold_exl3_experts_v1.down_combine", &.{ + z, pick, + svh_d, y, + cw, out, + .{ .i = n }, .{ .i = P }, + .{ .i = D }, .{ .i = 1 }, + .{ .i = slots }, .{ .i = Et }, + }); +} + +/// Engram (layers 1 and 14), before the layer's first mHC site: wkv over the hashed rows, gated fusion into streams. +pub fn engram(e: *E, L: u32) !void { + const n = e.n; + const D: i64 = e.cfg.hidden; + const K: i64 = @as(i64, e.cfg.hashCols()) * e.cfg.engram_head_dim; + const N: i64 = (@as(i64, e.cfg.engram_max_ngram) + 1) * D; + const kv = try e.buf("L.engram.kv", .{}, .bf16, &.{ n, N }); + try e.glue("engram_rows", &.{ try e.buf("L.engram.rows", .{}, .bf16, &.{ n, K }), .{ .i = L } }); + try e.upstream(try e.fmt("L{d}.engram.wkv", .{L}), try e.buf("L.engram.rows", .{}, .bf16, &.{ n, K }), K, N, kv); + const qk = try e.weight("L{d}.engram.qk", .{L}, .f32, &.{ e.cfg.hc_mult, D }); + try e.triton("_fuse_dec", .{ n, 1, 1 }, &.{ + .{ .name = "X", .arg = try e.streams() }, .{ .name = "x_stride", .arg = .{ .i = 4 * D } }, + .{ .name = "KV", .arg = kv }, .{ .name = "kv_stride", .arg = .{ .i = N } }, + .{ .name = "QK", .arg = qk }, .{ .name = "KEEP", .arg = qk }, + .{ .name = "GATE", .arg = qk }, .{ .name = "eps", .arg = .{ .f = E.dec(e.cfg.eps) } }, + .{ .name = "clamp", .arg = .{ .f = E.dec(e.cfg.engram_gate_clamp) } }, .{ .name = "sqrt_d", .arg = .{ .f = @sqrt(@as(f64, @floatFromInt(D))) } }, + .{ .name = "D", .arg = .{ .i = D } }, .{ .name = "BK", .arg = .{ .i = 64 } }, + .{ .name = "UB", .arg = .{ .i = 1024 } }, .{ .name = "HAS_KEEP", .arg = .{ .b = false } }, + .{ .name = "HAS_GATE", .arg = .{ .b = false } }, + }); +} + +/// The final norm (mHC final) and this rank's head columns (fp32 logits). +pub fn finish(e: *E) !void { + try mhc(e, null, "", .final); + const D: i64 = e.cfg.hidden; + const V: i64 = e.cfg.vocab / e.o.world; + const logits = try e.buf("w.logits", .{}, .f32, &.{ e.n, V }); + try e.upstream("head", try e.buf("w.out", .{}, .bf16, &.{ e.n, D }), D, V, logits); + try e.glue("kit_logits", &.{logits}); +} diff --git a/zig/src/families/deepseek_v41/block_prefetch.zig b/zig/src/families/deepseek_v41/block_prefetch.zig new file mode 100644 index 0000000000..47a9d4bcd2 --- /dev/null +++ b/zig/src/families/deepseek_v41/block_prefetch.zig @@ -0,0 +1,110 @@ +//! A decode window's L2 prefetch (TF_DSV41_L2PF): each site's sources, the table rows, the side-stream launches. + +const std = @import("std"); +const block = @import("block.zig"); +const Emitter = block.Emitter; +const L2Src = block.Emitter.L2Src; +const L2Site = block.Emitter.L2Site; + +/// l2pf._lin_tensors: every projection's suh and svh, then their trellis words. +pub fn linSrcs(e: *Emitter, out: *std.ArrayList(L2Src), ws: []const []const u8, ks: []const i64, ns: []const i64) !void { + for (ws, ks, ns) |w, k, n| try out.appendSlice(e.a, &.{ + .{ .kind = .weight, .name = try e.fmt("{s}.suh", .{w}), .bytes = @intCast(2 * k) }, + .{ .kind = .weight, .name = try e.fmt("{s}.svh", .{w}), .bytes = @intCast(2 * n) }, + }); + for (ws, ks, ns) |w, k, n| try out.append(e.a, .{ .kind = .words, .name = w, .bytes = try e.trellisBytes(w, k, n) }); +} + +/// Site `site`'s sources: "o" and "f" of block `L`, "x" of the block before `L` (its table is "s.L.pf.x"). +pub fn l2pfSources(e: *Emitter, L: u32, site: L2Site) ![]const L2Src { + var out: std.ArrayList(L2Src) = .empty; + switch (site) { + .o => { + const D: i64 = e.cfg.hidden; + const g = e.cfg.o_groups / e.o.world; + const ki: i64 = @divExact(@as(i64, e.cfg.heads) * e.cfg.head_dim, @as(i64, e.cfg.o_groups)); + const ws = try e.a.alloc([]const u8, g); + const ks = try e.a.alloc(i64, g); + const ns = try e.a.alloc(i64, g); + for (ws, ks, ns, 0..) |*w, *k, *n, i| { + w.* = try e.fmt("L{d}.attn.wo_a.{d}", .{ L, i }); + k.* = ki; + n.* = e.cfg.o_lora; + } + try e.linSrcs(&out, ws, ks, ns); + try e.linSrcs(&out, &.{try e.fmt("L{d}.attn.wo_b", .{L})}, &.{@as(i64, g) * e.cfg.o_lora}, &.{D}); + }, + .f => { + const D: i64 = e.cfg.hidden; + const I: i64 = e.cfg.expert_width / e.o.world; + const ex = e.cfg.expertsOf(L); + const sh = try e.fmt("L{d}.moe.shared.0.w1", .{L}); + const k2: u64 = try e.w.k2(sh); + const shared: u64 = @as(u64, @intCast(@divExact(D, 16) * @divExact(I, 16))) * 8 * k2 * 2; + try out.appendSlice(e.a, &.{ + .{ .kind = .weight, .name = try e.fmt("L{d}.moe.gate", .{L}), .bytes = @as(u64, ex.count) * @as(u64, @intCast(D)) * 2 }, + .{ .kind = .expert, .name = try e.fmt("s.L{d}.ex.tp_g", .{L}), .index = ex.count, .bytes = shared }, + .{ .kind = .expert, .name = try e.fmt("s.L{d}.ex.tp_u", .{L}), .index = ex.count, .bytes = shared }, + }); + }, + .x => { + const xg = try e.xGroup(L); + try e.linSrcs(&out, xg.ws, xg.ks, xg.ns); + const qg = try e.qGroup(L); + try e.linSrcs(&out, qg.ws, qg.ks, qg.ns); + }, + } + return out.items; +} + +pub fn srcSizes(e: *Emitter, srcs: []const L2Src) ![]const u64 { + const out = try e.a.alloc(u64, srcs.len); + for (srcs, out) |x, *y| y.* = x.bytes; + return out; +} + +/// l2pf.take: spans of `sizes` in order within the budget, the last one's head. +pub fn take(e: *const Emitter, sizes: []const u64) i64 { + var left = e.o.l2pf_budget; + var rows: i64 = 0; + for (sizes) |s| { + if (left < 16) break; + const n = @min(left, s) / 16 * 16; + if (n == 0) continue; + rows += 1; + left -= n; + } + return rows; +} + +/// Site "o" (after wq_b): wo_a's groups, then wo_b; l2pf segments. +pub fn prefetchO(e: *Emitter, L: u32) !void { + const sizes = try e.srcSizes(try e.l2pfSources(L, .o)); + var total: u64 = 0; + var left = e.o.l2pf_budget; + for (sizes) |s| { + if (left < 16) break; + const n = @min(left, s) / 16 * 16; + total += n; + left -= n; + } + const pieces = std.math.divCeil(u64, total, @intCast(e.o.l2pf_chunk)) catch 1; + const grid: i64 = @max(1, @min(e.o.sms, @as(i64, @intCast(std.math.divCeil(u64, pieces, 128) catch 1)))); + const rows = e.take(sizes); + try e.ext("tf_dsv41_l2pf_v1.segments", &.{ .{ .opaque_table = (try e.buf("s.L{d}.pf.o", .{L}, .i64, &.{ rows, 2 })).t }, .{ .i = grid }, .{ .i = 128 }, .{ .i = e.o.l2pf_chunk }, .{ .b = true } }); +} + +pub fn paced(e: *Emitter, role: []const u8, rows: i64) !void { + const ns: i64 = @max(1, @as(i64, @intFromFloat(@round(@as(f64, @floatFromInt(e.o.l2pf_chunk)) / e.o.l2pf_pace_gbps)))); + try e.ext("tf_dsv41_l2pace_v1.paced", &.{ .{ .opaque_table = (try e.buf("{s}", .{role}, .i64, &.{ rows, 2 })).t }, .{ .i = e.o.l2pf_pace_ctas }, .{ .i = e.o.l2pf_chunk }, .{ .i = ns }, .{ .i = 0 } }); +} + +/// Site "f" (after the attention sublayer): the router's gate, then the shared expert's gate and up trellises. +pub fn prefetchF(e: *Emitter, L: u32) !void { + try e.paced(try e.fmt("s.L{d}.pf.f", .{L}), e.take(try e.srcSizes(try e.l2pfSources(L, .f)))); +} + +/// Site "x" of the block before `nxt` (after its MoE): `nxt`'s x group, then its q group. +pub fn prefetchX(e: *Emitter, nxt: u32) !void { + try e.paced(try e.fmt("s.L{d}.pf.x", .{nxt}), e.take(try e.srcSizes(try e.l2pfSources(nxt, .x)))); +} diff --git a/zig/src/families/deepseek_v41/block_wide.zig b/zig/src/families/deepseek_v41/block_wide.zig new file mode 100644 index 0000000000..db83dcdd8f --- /dev/null +++ b/zig/src/families/deepseek_v41/block_wide.zig @@ -0,0 +1,167 @@ +//! Decode windows of more than 16 rows (TF_DSV41_ROWS_CAP past 16): mHC sites and router config past mhc_cuda's 16. + +const std = @import("std"); +const calls = @import("block_calls.zig"); +const block = @import("block.zig"); +const moe = @import("block_moe.zig"); +const E = block.Emitter; +const Arg = calls.Arg; + +/// mhc_cuda / mhc_dec's MAX_ROWS: windows of more rows take this file's path. +pub const cuda_rows: i64 = 16; +/// TF_DSV41_MHC_SPLIT_ROWS (32): split launches and the second streams buffer up to here. +pub const split_rows: i64 = 32; +/// attn_cuda.MAX_ROWS / dtopk.MAX_ROWS / rowtab's largest bucket: the widest decode window. +pub const max_rows: i64 = 64; + +pub const Mode = enum { boundary, site1, site2, final }; + +fn coefs(e: *E, set: u1) ![3]Arg { + return .{ + try e.buf("w.c{d}.pre", .{set}, .f32, &.{ e.n, 4 }), + try e.buf("w.c{d}.post", .{set}, .f32, &.{ e.n, 4 }), + try e.buf("w.c{d}.comb", .{set}, .f32, &.{ e.n, 16 }), + }; +} + +/// One mHC site of layer L (`which`: hc_attn / hc_ffn) or the final norm (L null) for a window of more than 16 rows. +pub fn mhc(e: *E, layer: ?u32, which: []const u8, mode: Mode, in_place: bool, norm: ?[]const u8, tap: ?Arg) !void { + const n = e.n; + const pf = pfOn(e, mode); + if ((n <= cuda_rows and !pf) or n > max_rows) return error.WideRows; + // PFDEC's deferred coefficients (MHC_DEFER): this site reads them (a no-op without a pending launch) + moe.awaitCoefs(e); + const D: i64 = e.cfg.hidden; + const hm: i64 = e.cfg.hc_mult; + const post = mode == .boundary or mode == .final; + const mix = mode != .final; + const x = try e.streams(); + const prev = try coefs(e, e.coef); + const next = try coefs(e, 1 - e.coef); + const part = try e.buf("L.mhc.part", .{}, .f32, &.{ n, 4, 40, 32 }); + const c = try e.buf("L.mhc.c", .{}, .bf16, &.{ n, D }); + // the forward's second streams buffer: a boundary of at most 32 rows writes it; a final writes the streams + const spare = mode == .boundary and !in_place and n <= split_rows; + const xout: Arg = if (spare) try e.spare() else x; + const g: ?Arg = if (post) try e.buf("w.recv.{s}", .{if (std.mem.eql(u8, which, "hc_ffn")) "attn" else "moe"}, .bf16, &.{ e.o.world, n, D }) else null; + const collapse: i64 = if (mode == .site1) 1 else 2; + const pre: ?Arg = if (collapse == 2) prev[0] else null; + const fnw: ?Arg = if (mix) try e.buf("s.L{d}.{s}.fn16", .{ layer.?, which }, .bf16, &.{ 6 * hm, hm * D }) else null; + if (pf) { + const L = layer.?; + try e.ext("tf_dsv41_mhc_pf_v1.run", &.{ + x, if (post) xout else try e.empty(.bf16), + g orelse try e.empty(.bf16), if (post) prev[1] else try e.empty(.f32), + if (post) prev[2] else try e.empty(.f32), pre orelse try e.empty(.f32), + fnw.?, part, + c, tap orelse try e.empty(.bf16), + .{ .i = n }, .{ .i = @backingInt(mode) }, + }); + try finishK(e, null, part, c, try normOf(e, layer, which, norm)); + if (e.o.mhc_defer) { + e.df_side = true; // MHC_DEFER: on the deferred stream, the next mHC call joins it + e.df_pending = true; + } + try e.ext("tf_dsv41_mhc_cuda_v1.coef", &.{ + part, try e.weight("L{d}.{s}.base", .{ L, which }, .f32, &.{6 * hm}), + try e.weight("L{d}.{s}.scale", .{ L, which }, .f32, &.{3}), next[0], + next[1], next[2], + .{ .i = n }, .{ .f = E.dec(e.cfg.eps) }, + .{ .f = E.dec(e.cfg.hc_eps) }, .{ .f = E.dec(e.cfg.hc_post_alpha) }, + .{ .i = e.cfg.hc_sinkhorn_iters }, + }); + e.coef = 1 - e.coef; + if (spare) e.cur = 1 - e.cur; + return; + } + const split = mix and n <= split_rows and !(post and !spare); + try e.triton("_site", .{ std.math.divCeil(i64, n, 16) catch unreachable, 40, if (split) 4 else 1 }, &.{ + .{ .name = "X", .arg = x }, .{ .name = "x_stride", .arg = .{ .i = hm * D } }, + .{ .name = "XOUT", .arg = if (post) xout else x }, .{ .name = "G", .arg = g orelse x }, + .{ .name = "g_rank", .arg = .{ .i = if (post) n * D else 0 } }, .{ .name = "POST", .arg = if (post) prev[1] else x }, + .{ .name = "COMB", .arg = if (post) prev[2] else x }, .{ .name = "PRE", .arg = pre orelse x }, + .{ .name = "FN", .arg = fnw orelse x }, .{ .name = "PART", .arg = part }, + .{ .name = "C", .arg = c }, .{ .name = "TAP", .arg = tap orelse x }, + .{ .name = "tap_stride", .arg = .{ .i = if (tap) |t| t.t.stride[0] else 0 } }, .{ .name = "R", .arg = .{ .i = n } }, + .{ .name = "D", .arg = .{ .i = D } }, .{ .name = "NB", .arg = .{ .i = 40 } }, + .{ .name = "BM", .arg = .{ .i = 16 } }, .{ .name = "BK", .arg = .{ .i = 16 } }, + .{ .name = "WORLD", .arg = .{ .i = if (post) e.o.world else 1 } }, .{ .name = "POST_ON", .arg = .{ .b = post } }, + .{ .name = "COLLAPSE", .arg = .{ .i = collapse } }, .{ .name = "MIX", .arg = .{ .b = mix } }, + .{ .name = "TAP_ON", .arg = .{ .b = tap != null } }, .{ .name = "SPLIT", .arg = .{ .b = split } }, + }); + try finishK(e, if (mix) .{ layer.?, which, next } else null, part, c, try normOf(e, layer, which, norm)); + if (mix) e.coef = 1 - e.coef; + if (spare) e.cur = 1 - e.cur; +} + +/// TF_DSV41_MHC_PFDEC: this window's mixing sites take mhc_pf + `coef_kernel` (`mhc`'s `pf` branch). +pub fn pfOn(e: *const E, mode: anytype) bool { + const m = e.o.mhc_pf_rows; + return m > 0 and e.n >= m and e.n <= max_rows and @backingInt(mode) != @backingInt(Mode.final); +} + +/// A site's (or the final's) norm weight: the layer's attn / ffn norm, the model's (or `norm`'s) at the final. +fn normOf(e: *E, layer: ?u32, which: []const u8, norm: ?[]const u8) !Arg { + const D: i64 = e.cfg.hidden; + if (layer) |L| return e.weight("L{d}.{s}", .{ L, if (std.mem.eql(u8, which, "hc_attn")) "attn_norm" else "ffn_norm" }, .f32, &.{D}); + return e.weight("{s}", .{norm orelse "norm"}, .f32, &.{D}); +} + +/// The mHC finish kernel: with `coef` also the next site's coefficients, else the normed input alone. +fn finishK(e: *E, coef: ?struct { u32, []const u8, [3]Arg }, part: Arg, c: Arg, nw: Arg) !void { + const n = e.n; + const D: i64 = e.cfg.hidden; + const hm: i64 = e.cfg.hc_mult; + const mix = coef != null; + try e.triton("_finish_k", .{ n, 1, 1 }, &.{ + .{ .name = "PART", .arg = part }, + .{ .name = "BASE", .arg = if (coef) |k| try e.weight("L{d}.{s}.base", .{ k[0], k[1] }, .f32, &.{6 * hm}) else part }, + .{ .name = "SCALE", .arg = if (coef) |k| try e.weight("L{d}.{s}.scale", .{ k[0], k[1] }, .f32, &.{3}) else part }, + .{ .name = "PRE", .arg = if (coef) |k| k[2][0] else part }, + .{ .name = "POST", .arg = if (coef) |k| k[2][1] else part }, + .{ .name = "COMB", .arg = if (coef) |k| k[2][2] else part }, + .{ .name = "C", .arg = c }, + .{ .name = "NW", .arg = nw }, + .{ .name = "OUT", .arg = try e.buf("w.out", .{}, .bf16, &.{ n, D }) }, + .{ .name = "eps", .arg = .{ .f = E.dec(e.cfg.eps) } }, + .{ .name = "hc_eps", .arg = .{ .f = E.dec(e.cfg.hc_eps) } }, + .{ .name = "post_alpha", .arg = .{ .f = E.dec(e.cfg.hc_post_alpha) } }, + .{ .name = "D", .arg = .{ .i = D } }, + .{ .name = "NB", .arg = .{ .i = 40 } }, + .{ .name = "ITERS", .arg = .{ .i = e.cfg.hc_sinkhorn_iters } }, + .{ .name = "COEF", .arg = .{ .b = mix } }, + .{ .name = "BLOCK", .arg = .{ .i = @intCast(std.math.gcd(@as(u64, 1024), @as(u64, @intCast(D)))) } }, + .{ .name = "UNROLL", .arg = .{ .b = true } }, + .{ .name = "NORM_FIRST", .arg = .{ .b = true } }, + }); +} + +/// block_moe.mhcInto's call for a window of more than 16 rows (`mode`: block_moe.Mode), with its DSpark tap. +pub fn decode(e: *E, layer: ?u32, which: []const u8, mode: anytype, in_place: bool) !void { + const m = std.meta.stringToEnum(Mode, @tagName(mode)).?; + var tap: ?Arg = null; + if (e.o.taps and m == .boundary and std.mem.eql(u8, which, "hc_attn")) if (layer) |L| { + const t = e.cfg.dspark_targets.items(); + if (std.mem.indexOfScalar(u32, t, L)) |j| { + const D: i64 = e.cfg.hidden; + const w: i64 = @as(i64, @intCast(t.len)) * D; + tap = try e.view(.{ .buf = "w.taps" }, .bf16, &.{ e.n, D }, &.{ w, 1 }, @as(i64, @intCast(j)) * D * 2); + } + }; + return mhc(e, layer, which, m, in_place, null, tap); +} + +/// router_gemv config at `n` rows: (row tile, warps, experts a warp); past 16 rows the prefill config. +pub fn routerCfg(n: i64) [3]i64 { + if (n > cuda_rows) return .{ 16, 8, 2 }; + const t = std.math.ceilPowerOfTwo(u64, @intCast(n)) catch unreachable; + return .{ @intCast(t), 2, 0 }; +} + +test "wide rows: router config as router_gemv.config / row_tile" { + try std.testing.expectEqual([3]i64{ 1, 2, 0 }, routerCfg(1)); + try std.testing.expectEqual([3]i64{ 8, 2, 0 }, routerCfg(5)); + try std.testing.expectEqual([3]i64{ 16, 2, 0 }, routerCfg(16)); + try std.testing.expectEqual([3]i64{ 16, 8, 2 }, routerCfg(17)); + try std.testing.expectEqual([3]i64{ 16, 8, 2 }, routerCfg(64)); +} diff --git a/zig/src/families/deepseek_v41/buffers.zig b/zig/src/families/deepseek_v41/buffers.zig new file mode 100644 index 0000000000..9d547dea86 --- /dev/null +++ b/zig/src/families/deepseek_v41/buffers.zig @@ -0,0 +1,89 @@ +//! The buffer plan: every role a window's calls name, sized by its largest view, over every row bucket. + +const std = @import("std"); +const calls = @import("block_calls.zig"); + +pub const Scope = enum { persistent, window, layer }; // "s." own allocation; "w."/"L." one arena + +pub fn scopeOf(role: []const u8) Scope { + if (std.mem.startsWith(u8, role, "s.")) return .persistent; + if (std.mem.startsWith(u8, role, "w.")) return .window; + return .layer; +} + +pub fn dtSize(dt: calls.Dt) u64 { + return switch (dt) { + .bool, .i8, .u8 => 1, + .bf16, .f16, .i16 => 2, + .f32, .i32 => 4, + .f64, .i64 => 8, + }; +} + +/// Bytes from the role's first byte that a view reaches (0: an empty tensor). +pub fn extent(t: calls.Tensor) u64 { + var last: i64 = 0; + for (t.shape, t.stride) |n, s| { + if (n == 0) return 0; + last += (n - 1) * s; + } + return @as(u64, @intCast(t.offset)) + (@as(u64, @intCast(last)) + 1) * dtSize(t.dt); +} + +pub const Plan = struct { + a: std.mem.Allocator, + sizes: std.StringArrayHashMapUnmanaged(u64) = .empty, + /// the weights the programs read by name (forms_inplace.zig leaves those in place) + weights: std.StringArrayHashMapUnmanaged(void) = .empty, + + fn note(p: *Plan, x: calls.Arg) !void { + switch (x) { + .list => |items| for (items) |y| try p.note(y), + .t, .opaque_table => |t| switch (t.role) { + .buf => |r| { + const n = extent(t); + if (n == 0) return; + const g = try p.sizes.getOrPut(p.a, r); + if (!g.found_existing) { + g.key_ptr.* = try p.a.dupe(u8, r); + g.value_ptr.* = 0; + } + g.value_ptr.* = @max(g.value_ptr.*, n); + }, + .weight => |w| { + const g = try p.weights.getOrPut(p.a, w); + if (!g.found_existing) g.key_ptr.* = try p.a.dupe(u8, w); + }, + else => {}, + }, + else => {}, + } + } + + /// Takes every role of `cs` into the plan (call once a row bucket). + pub fn add(p: *Plan, cs: []const calls.Call) !void { + for (cs) |c| for (c.args) |x| try p.note(x.arg); + } + + pub const Totals = struct { roles: [3]usize = @splat(0), bytes: [3]u64 = @splat(0) }; + + /// Bytes by scope, each role rounded up to 256 (its place in the arena or its own allocation). + pub fn totals(p: *const Plan) Totals { + var t: Totals = .{}; + for (p.sizes.keys(), p.sizes.values()) |k, v| { + const s = @backingInt(scopeOf(k)); + t.roles[s] += 1; + t.bytes[s] += std.mem.alignForward(u64, v, 256); + } + return t; + } +}; + +const testing = std.testing; + +test "extent: a view's reach from the role's first byte" { + const t: calls.Tensor = .{ .role = .{ .buf = "w.x" }, .dt = .bf16, .shape = &.{ 3, 4096 }, .stride = &.{ 16384, 1 }, .offset = 8192 }; + try testing.expectEqual(@as(u64, 8192 + (2 * 16384 + 4095 + 1) * 2), extent(t)); + const e: calls.Tensor = .{ .role = .{ .buf = "w.x" }, .dt = .f32, .shape = &.{0}, .stride = &.{1} }; + try testing.expectEqual(@as(u64, 0), extent(e)); +} diff --git a/zig/src/families/deepseek_v41/config.zig b/zig/src/families/deepseek_v41/config.zig new file mode 100644 index 0000000000..85e1401978 --- /dev/null +++ b/zig/src/families/deepseek_v41/config.zig @@ -0,0 +1,340 @@ +//! DeepSeek-V4.1-Flash text config (config.json -> text_config) and the per-layer CSA2 topology. + +const std = @import("std"); + +/// Fixed-capacity id lists of the config (layers, sources, Engram tables); the release's longest has 43 entries. +pub fn List(comptime T: type, comptime cap: usize) type { + return struct { + const Self = @This(); + buf: [cap]T = undefined, + len: usize = 0, + + pub fn of(values: []const T) Self { + var s: Self = .{}; + for (values) |v| s.push(v) catch unreachable; + return s; + } + + pub fn push(s: *Self, v: T) !void { + if (s.len == cap) return error.ConfigListTooLong; + s.buf[s.len] = v; + s.len += 1; + } + + pub fn items(s: *const Self) []const T { + return s.buf[0..s.len]; + } + + pub fn contains(s: *const Self, v: T) bool { + return std.mem.indexOfScalar(T, s.items(), v) != null; + } + }; +} + +pub const Ids = List(u32, 64); + +pub const Rope = struct { + theta: f64 = 10000.0, + compress_theta: f64 = 160000.0, + factor: f64 = 16.0, + original_max_positions: u64 = 65536, + beta_fast: f64 = 32.0, + beta_slow: f64 = 1.0, + max_positions: u64 = 1048576, + yarn: bool = true, +}; + +/// The attention a layer runs: sliding window only, KV source + indexer, own indexer on a source's keys, or reuse. +pub const Mode = enum { swa, full, reindex, reuse }; + +pub const Config = struct { + vocab: u32 = 129280, + hidden: u32 = 5120, + layers: u32 = 40, + heads: u32 = 64, + head_dim: u32 = 512, + rope_dim: u32 = 64, + q_lora: u32 = 1280, + o_lora: u32 = 1024, + o_groups: u32 = 8, + eps: f32 = 1e-20, + window: u32 = 128, + experts: u32 = 384, + shared_experts: u32 = 1, + topk: u32 = 6, + expert_width: u32 = 2304, + routed_scale: f32 = 1.5, + norm_topk: bool = true, + swiglu_limit: f32 = 10.0, + compress_ratios: Ids = Ids.of(&.{ 0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0 }), + kv_sources: Ids = Ids.of(&.{ 2, 8, 14, 20 }), + index_sources: Ids = Ids.of(&.{ 2, 8, 14, 20, 24, 28, 32, 36 }), + index_heads: u32 = 32, + index_dim: u32 = 128, + index_topk: u32 = 512, + candidate_source: i64 = 20, + candidate_blocks: u32 = 2048, + candidate_block_size: u32 = 8, + hc_mult: u32 = 4, + hc_sinkhorn_iters: u32 = 20, + hc_eps: f32 = 1e-6, + hc_post_alpha: f32 = 2.0, + engram_layers: Ids = Ids.of(&.{ 1, 14 }), + engram_rows: List(u64, 8) = List(u64, 8).of(&.{ 384006168, 384016682 }), + engram_max_ngram: u32 = 4, + engram_vocab: u64 = 16000000, + engram_heads: u32 = 8, + engram_head_dim: u32 = 256, + engram_pad: u32 = 2, + engram_compressed_vocab: u32 = 99092, + engram_gate_clamp: f32 = 1e-6, + mtp_layers: u32 = 3, + dspark_block: u32 = 5, + dspark_noise_token: u32 = 128799, + dspark_targets: Ids = Ids.of(&.{ 37, 38, 39 }), + dspark_markov_rank: u32 = 256, + dspark_experts: u32 = 128, + dspark_topk: u32 = 3, + rope: Rope = .{}, + bos: u32 = 0, + eos: u32 = 1, + + /// config.json in `dir`. + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + /// Checkpoint JSON (text_config, else top level); absent keys keep release values, unknown routing is refused. + pub fn parse(gpa: std.mem.Allocator, json: []const u8) !Config { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, json, .{}); + defer parsed.deinit(); + if (parsed.value != .object) return error.BadConfig; + const top = parsed.value.object; + const text = if (top.get("text_config")) |t| (if (t == .object) t.object else return error.BadConfig) else top; + var c: Config = .{}; + const ints = .{ + .{ "vocab", "vocab_size" }, .{ "hidden", "hidden_size" }, + .{ "layers", "num_hidden_layers" }, .{ "heads", "num_attention_heads" }, + .{ "head_dim", "head_dim" }, .{ "rope_dim", "qk_rope_head_dim" }, + .{ "q_lora", "q_lora_rank" }, .{ "o_lora", "o_lora_rank" }, + .{ "o_groups", "o_groups" }, .{ "window", "sliding_window" }, + .{ "experts", "n_routed_experts" }, .{ "shared_experts", "n_shared_experts" }, + .{ "topk", "num_experts_per_tok" }, .{ "expert_width", "moe_intermediate_size" }, + .{ "index_heads", "index_n_heads" }, .{ "index_dim", "index_head_dim" }, + .{ "index_topk", "index_topk" }, .{ "candidate_source", "candidate_source_layer_id" }, + .{ "candidate_blocks", "candidate_topk_blocks" }, .{ "candidate_block_size", "candidate_block_size" }, + .{ "hc_mult", "hc_mult" }, .{ "hc_sinkhorn_iters", "hc_sinkhorn_iters" }, + .{ "engram_max_ngram", "engram_max_ngram_size" }, .{ "engram_vocab", "engram_vocab_size" }, + .{ "engram_heads", "engram_n_heads" }, .{ "engram_head_dim", "engram_head_dim" }, + .{ "engram_pad", "engram_pad_token_id" }, .{ "engram_compressed_vocab", "engram_compressed_vocab_size" }, + .{ "mtp_layers", "num_nextn_predict_layers" }, .{ "dspark_block", "dspark_block_size" }, + .{ "dspark_noise_token", "dspark_noise_token_id" }, .{ "dspark_markov_rank", "dspark_markov_rank" }, + .{ "dspark_experts", "dspark_n_routed_experts" }, .{ "dspark_topk", "dspark_num_experts_per_tok" }, + }; + inline for (ints) |p| if (text.get(p[1])) |v| { + @field(c, p[0]) = try intAs(@TypeOf(@field(c, p[0])), v); + }; + const floats = .{ + .{ "eps", "rms_norm_eps" }, .{ "routed_scale", "routed_scaling_factor" }, .{ "swiglu_limit", "swiglu_limit" }, + .{ "hc_eps", "hc_eps" }, .{ "hc_post_alpha", "hc_post_alpha" }, .{ "engram_gate_clamp", "engram_gate_clamp" }, + }; + inline for (floats) |p| if (text.get(p[1])) |v| { + @field(c, p[0]) = @floatCast(try float(v)); + }; + if (text.get("norm_topk_prob")) |v| c.norm_topk = if (v == .bool) v.bool else return error.BadConfig; + const lists = .{ + .{ "compress_ratios", "compress_ratios" }, .{ "kv_sources", "kv_source_layer_ids" }, + .{ "index_sources", "index_source_layer_ids" }, .{ "engram_layers", "engram_layer_ids" }, + .{ "dspark_targets", "dspark_target_layer_ids" }, .{ "engram_rows", "engram_num_embeddings" }, + }; + inline for (lists) |p| if (text.get(p[1])) |v| { + @field(c, p[0]) = try list(@TypeOf(@field(c, p[0])), v); + }; + // routing the kernels implement: sqrt(softplus) scores, bias-corrected top-k (vLLM noaux_tc) + if (text.get("scoring_func")) |v| if (v != .string or !std.mem.eql(u8, v.string, "sqrtsoftplus")) return error.UnsupportedScoring; + if (text.get("topk_method")) |v| if (v != .string or !std.mem.eql(u8, v.string, "noaux_tc")) return error.UnsupportedScoring; + c.rope = try rope(text); + if (top.get("bos_token_id")) |v| c.bos = try intAs(u32, v); + if (top.get("eos_token_id")) |v| c.eos = try intAs(u32, v); + try c.validate(); + return c; + } + + /// Shapes the kernels and the loader assume; a release that breaks one is refused before any allocation. + pub fn validate(c: *const Config) !void { + if (c.hidden % 128 != 0 or c.expert_width % 128 != 0 or c.vocab % 128 != 0) return error.BadConfig; + if (c.heads % c.o_groups != 0 or c.rope_dim > c.head_dim or c.topk > c.experts) return error.BadConfig; + if (c.dspark_topk > c.dspark_experts and c.dspark_experts != 0) return error.BadConfig; + if (c.compress_ratios.len < c.layers) return error.BadConfig; + for (c.compress_ratios.items()) |r| if (r > 2) return error.BadConfig; + for (c.kv_sources.items()) |s| if (s >= c.layers or c.compressRatio(s) == 0) return error.BadConfig; + for (c.index_sources.items()) |s| if (s >= c.layers or c.compressRatio(s) == 0) return error.BadConfig; + if (c.engram_layers.len != c.engram_rows.len) return error.BadConfig; + // every compressed layer needs a KV source and an index source at or before it + var l: u32 = 0; + while (l < c.layers) : (l += 1) if (c.compressRatio(l) > 0 and (c.kvSource(l) == null or c.indexSource(l) == null)) return error.BadConfig; + } + + /// Layers the forward runs: the backbone, then the DSpark blocks (checkpoint mtp.*). + pub fn allLayers(c: *const Config) u32 { + return c.layers + c.mtp_layers; + } + + pub fn compressRatio(c: *const Config, layer: u32) u32 { + return if (layer < c.compress_ratios.len) c.compress_ratios.buf[layer] else 0; + } + + pub fn isBackbone(c: *const Config, layer: u32) bool { + return layer < c.layers; + } + + /// The latest source at or before `layer` (null for sliding-window layers). + fn latest(c: *const Config, sources: *const Ids, layer: u32) ?u32 { + if (c.compressRatio(layer) == 0) return null; + var best: ?u32 = null; + for (sources.items()) |s| if (s <= layer and (best == null or s > best.?)) { + best = s; + }; + return best; + } + + pub fn kvSource(c: *const Config, layer: u32) ?u32 { + return c.latest(&c.kv_sources, layer); + } + + pub fn indexSource(c: *const Config, layer: u32) ?u32 { + return c.latest(&c.index_sources, layer); + } + + pub fn isKvSource(c: *const Config, layer: u32) bool { + return c.isBackbone(layer) and c.kv_sources.contains(layer); + } + + pub fn isIndexSource(c: *const Config, layer: u32) bool { + return c.isBackbone(layer) and c.index_sources.contains(layer); + } + + pub fn isCandidateSource(c: *const Config, layer: u32) bool { + return c.candidate_source == layer and c.candidate_blocks > 0; + } + + pub fn usesCandidates(c: *const Config, layer: u32) bool { + return c.isIndexSource(layer) and c.candidate_blocks > 0 and c.candidate_source >= 0 and c.candidate_source < layer; + } + + pub fn mode(c: *const Config, layer: u32) Mode { + if (c.compressRatio(layer) == 0) return .swa; + if (c.isKvSource(layer)) return .full; + if (c.isIndexSource(layer)) return .reindex; + return .reuse; + } + + pub fn isEngram(c: *const Config, layer: u32) bool { + return c.engram_layers.contains(layer); + } + + /// (routed experts, experts a token) of a backbone layer or a DSpark block. + pub fn expertsOf(c: *const Config, layer: u32) struct { count: u32, topk: u32 } { + if (c.isBackbone(layer)) return .{ .count = c.experts, .topk = c.topk }; + return .{ .count = if (c.dspark_experts != 0) c.dspark_experts else c.experts, .topk = if (c.dspark_topk != 0) c.dspark_topk else c.topk }; + } + + /// Engram's hash columns: (max n-gram - 1) x heads. + pub fn hashCols(c: *const Config) u32 { + return (c.engram_max_ngram - 1) * c.engram_heads; + } +}; + +fn intAs(comptime T: type, v: std.json.Value) !T { + return switch (v) { + .integer => |i| std.math.cast(T, i) orelse error.BadConfig, + .float => |f| if (@floor(f) == f) std.math.cast(T, @as(i64, @intFromFloat(f))) orelse error.BadConfig else error.BadConfig, + else => error.BadConfig, + }; +} + +fn float(v: std.json.Value) !f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadConfig, + }; +} + +fn list(comptime L: type, v: std.json.Value) !L { + if (v != .array) return error.BadConfig; + var out: L = .{}; + const T = @TypeOf(out.buf[0]); + for (v.array.items) |x| try out.push(try intAs(T, x)); + return out; +} + +/// rope_theta, compress_rope_theta and rope_scaling (or rope_parameters), as RopeConfig.from_dict. +fn rope(text: std.json.ObjectMap) !Rope { + var r: Rope = .{}; + if (text.get("rope_theta")) |v| r.theta = try float(v); + if (text.get("compress_rope_theta")) |v| r.compress_theta = try float(v); + if (text.get("max_position_embeddings")) |v| r.max_positions = try intAs(u64, v); + const scaling = text.get("rope_scaling") orelse text.get("rope_parameters"); + r.factor = 1.0; + r.original_max_positions = r.max_positions; + r.yarn = false; + if (scaling) |s| if (s == .object) { + const o = s.object; + if (o.get("factor")) |v| r.factor = try float(v); + if (o.get("original_max_position_embeddings")) |v| r.original_max_positions = try intAs(u64, v); + if (o.get("beta_fast")) |v| r.beta_fast = try float(v); + if (o.get("beta_slow")) |v| r.beta_slow = try float(v); + const kind = o.get("rope_type") orelse o.get("type"); + r.yarn = if (kind) |k| (k == .string and !std.mem.eql(u8, k.string, "default")) else false; + }; + return r; +} + +/// A small config with every V4.1 mechanism; split widths are 256 so TP=2 keeps 128-value blocks. +pub fn tiny() Config { + return .{ + .vocab = 256, + .hidden = 256, + .layers = 8, + .heads = 4, + .head_dim = 64, + .rope_dim = 8, + .q_lora = 32, + .o_lora = 128, + .o_groups = 2, + .eps = 1e-6, + .window = 4, + .experts = 8, + .topk = 2, + .expert_width = 256, + .compress_ratios = Ids.of(&.{ 0, 0, 2, 2, 1, 1, 1, 1, 0, 0 }), + .kv_sources = Ids.of(&.{ 2, 4 }), + .index_sources = Ids.of(&.{ 2, 4, 6 }), + .index_heads = 2, + .index_dim = 16, + .index_topk = 3, + .candidate_source = 4, + .candidate_blocks = 2, + .candidate_block_size = 2, + .engram_layers = Ids.of(&.{1}), + .engram_rows = List(u64, 8).of(&.{0}), + .engram_max_ngram = 3, + .engram_vocab = 31, + .engram_heads = 2, + .engram_head_dim = 8, + .engram_compressed_vocab = 50, + .mtp_layers = 2, + .dspark_block = 3, + .dspark_noise_token = 96, + .dspark_targets = Ids.of(&.{ 5, 6, 7 }), + .dspark_markov_rank = 8, + .dspark_experts = 4, + .dspark_topk = 2, + .rope = .{ .factor = 4.0, .original_max_positions = 64, .max_positions = 256 }, + }; +} diff --git a/zig/src/families/deepseek_v41/deepseek_v41.zig b/zig/src/families/deepseek_v41/deepseek_v41.zig new file mode 100644 index 0000000000..27c07a2e13 --- /dev/null +++ b/zig/src/families/deepseek_v41/deepseek_v41.zig @@ -0,0 +1,38 @@ +//! DeepSeek-V4.1-Flash on the Zig engine: release config, EXL3 pack headers, a rank's TP weight plan (host side). + +const std = @import("std"); + +pub const config = @import("config.zig"); +pub const Config = config.Config; +pub const exl3 = @import("exl3.zig"); +pub const names = @import("names.zig"); +pub const pack = @import("pack.zig"); +pub const Pack = pack.Pack; +pub const plan = @import("plan.zig"); +pub const Plan = plan.Plan; +pub const stage = @import("stage.zig"); +pub const named = @import("named.zig"); +pub const prepare = @import("prepare.zig"); +pub const engram_host = @import("engram_host.zig"); + +test { + _ = config; + _ = exl3; + _ = names; + _ = pack; + _ = plan; + _ = stage; + _ = named; + _ = prepare; + _ = engram_host; + // the device-side weight load, compiled on the host (no GPU work) + inline for (.{ @import("load.zig"), @import("forms.zig"), @import("forms_inplace.zig"), @import("host_stage.zig"), @import("triton_capture.zig"), @import("ex_stacks.zig"), @import("buffers.zig"), @import("block_calls.zig") }) |t| std.testing.refAllDecls(t); + // a block's emitter (attention, index, MoE, read-ahead, windows past 16 rows) over the launch bindings + const block = @import("block.zig"); + inline for (.{ block, block.Emitter, block.Pool, @import("block_attn.zig"), @import("block_index.zig"), @import("block_moe.zig"), @import("block_prefetch.zig"), @import("block_wide.zig"), @import("knob_types.zig"), @import("forward_types.zig") }) |t| std.testing.refAllDecls(t); + _ = @import("host_test.zig"); + _ = @import("engram_rows_test.zig"); // Engram's shard rows + _ = @import("phases.zig"); // the phase timers + // prepared weight folders (TF_DSV41_PREPARED) + _ = @import("prepared.zig"); +} diff --git a/zig/src/families/deepseek_v41/engram_host.zig b/zig/src/families/deepseek_v41/engram_host.zig new file mode 100644 index 0000000000..cb40006758 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_host.zig @@ -0,0 +1,307 @@ +//! Engram's host side: table layout, n-gram hashes of a window's tokens, a rank's heads, and the synthetic index rows. + +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel + +pub const dead_id: i64 = -1; +pub const max_layers = 4; +pub const max_ngram = 8; +pub const max_cols = 64; + +/// Where a window's table rows come from (local NVMe shards or a remote service); every source gives the same bytes. +pub const Source = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// Starts fetching Engram layer index `li`'s rows `idx` (this rank's heads' global rows); returns at once. + issue: *const fn (ptr: *anyopaque, li: usize, idx: []const i64) anyerror!void, + /// Layer index `li`'s rows `idx` as bf16 bits [idx.len, dim] (waits for any still in flight). + rows: *const fn (ptr: *anyopaque, li: usize, idx: []const i64, dim: usize, out: []u16) anyerror!void, + }; +}; + +pub const Host = struct { + layer_ids: []const u32, + max_ngram: u32, + heads: u32, + /// token id -> compressed id + token_map: []const u32, + pad_id: i64, + mult: [max_layers][max_ngram]i64 = undefined, + primes: [max_layers][max_cols]i64 = undefined, + offsets: [max_layers][max_cols]i64 = undefined, + + pub fn cols(h: *const Host) u32 { + return (h.max_ngram - 1) * h.heads; + } + + /// Prime bucket sizes: for each layer and order the next unused primes above `vocab_size - 1` (reference order). + pub fn layout(h: *Host, vocab_size: u64) void { + var seen: [max_layers * max_cols]i64 = undefined; + var nseen: usize = 0; + for (0..h.layer_ids.len) |li| { + var col: usize = 0; + for (0..h.max_ngram - 1) |_| { + var cur: i64 = @intCast(vocab_size - 1); + for (0..h.heads) |_| { + cur += 1; + while (!isPrime(@intCast(cur)) or std.mem.indexOfScalar(i64, seen[0..nseen], cur) != null) cur += 1; + seen[nseen] = cur; + nseen += 1; + h.primes[li][col] = cur; + col += 1; + } + } + var acc: i64 = 0; + for (0..col) |c| { + h.offsets[li][c] = acc; + acc += h.primes[li][c]; + } + } + } + + /// Table rows [t][layers][cols] of `ids` after `lookback` (oldest first; short: sequence start); past map: dead. + pub fn hash(h: *const Host, ids: []const u32, lookback: []const u32, out: []i64) void { + const nb = lookback.len; + const L = h.layer_ids.len; + const C = h.cols(); + std.debug.assert(out.len == ids.len * L * C); + for (0..ids.len) |t| for (0..L) |li| { + var blocked = false; + var rolling: i64 = 0; + for (0..h.max_ngram) |s| { + const j = @as(i64, @intCast(t + nb)) - @as(i64, @intCast(s)); + const before = j < 0; + const v: i64 = if (before) h.pad_id else h.src(if (j < @as(i64, @intCast(nb))) lookback[@intCast(j)] else ids[@as(usize, @intCast(j)) - nb]); + blocked = blocked or before or v == dead_id; + const value = if (blocked) h.pad_id else v; + rolling ^= value *% h.mult[li][s]; + if (s == 0) continue; + for (0..h.heads) |hh| { + const c = (s - 1) * h.heads + hh; + out[(t * L + li) * C + c] = @mod(rolling, h.primes[li][c]) + h.offsets[li][c]; + } + } + }; + } + + fn src(h: *const Host, id: u32) i64 { + return if (id >= h.token_map.len) dead_id else h.token_map[id]; + } + + /// (first head, heads) rank `rank` of `world` owns: complete heads, TP-major (head_shard). + pub fn headShard(h: *const Host, rank: u32, world: u32) [2]u32 { + const n = h.cols(); + const part = (n + world - 1) / world; + const h0 = @min(rank * part, n); + return .{ h0, @min(h0 + part, n) - h0 }; + } + + /// This rank's heads' global rows of Engram layer `layer` for a window's `ids` after `lookback`: idx [n][heads]. + pub fn shardIndex(h: *const Host, a: std.mem.Allocator, ids: []const u32, lookback: []const u32, layer: u32, rank: u32, world: u32) ![]i64 { + const li = std.mem.indexOfScalar(u32, h.layer_ids, layer) orelse return error.NotEngram; + const L = h.layer_ids.len; + const C = h.cols(); + const sh = h.headShard(rank, world); + const hashes = try a.alloc(i64, ids.len * L * C); + defer a.free(hashes); + h.hash(ids, lookback, hashes); + const idx = try a.alloc(i64, ids.len * sh[1]); + for (0..ids.len) |t| @memcpy(idx[t * sh[1] ..][0..sh[1]], hashes[(t * L + li) * C + sh[0] ..][0..sh[1]]); + return idx; + } + + /// Engram layer `layer`'s window rows, this rank's heads, bf16 [ids.len, heads x dim]: from `source` or index rows. + pub fn windowRows(h: *const Host, a: std.mem.Allocator, source: ?Source, ids: []const u32, lookback: []const u32, layer: u32, rank: u32, world: u32, dim: usize, out: []u16) !void { + const idx = try h.shardIndex(a, ids, lookback, layer, rank, world); + defer a.free(idx); + const s = source orelse return indexRows(idx, dim, out); + const li = std.mem.indexOfScalar(u32, h.layer_ids, layer).?; + return s.vtable.rows(s.ptr, li, idx, dim, out); + } + + /// Starts reading a coming window's rows on every Engram layer; a wrong guess costs a read; errors only logged. + pub fn prefetch(h: *const Host, a: std.mem.Allocator, source: ?Source, ids: []const u32, lookback: []const u32, rank: u32, world: u32) void { + const s = source orelse return; + for (h.layer_ids, 0..) |layer, li| { + const idx = h.shardIndex(a, ids, lookback, layer, rank, world) catch return; + defer a.free(idx); + s.vtable.issue(s.ptr, li, idx) catch |e| if (!quiet) std.log.scoped(.dsv41).warn("engram: prefetch of layer {d} failed ({t})", .{ layer, e }); + } + } + + /// Loads recipe/dsv41_engram_host.py's file (the map borrowed from `bytes`, which must outlive the host). + pub fn load(bytes: []const u8, heads: u32, pad_token: u32, vocab_size: u64, layer_buf: []u32) !Host { + if (bytes.len < 24 or !std.mem.eql(u8, bytes[0..8], "DSV41EH1")) return error.BadEngramHost; + const L = std.mem.readInt(u32, bytes[8..12], .little); + const M = std.mem.readInt(u32, bytes[12..16], .little); + const V = std.mem.readInt(u32, bytes[16..20], .little); + if (L > max_layers or M > max_ngram or (M - 1) * heads > max_cols or L > layer_buf.len) return error.BadEngramHost; + var at: usize = 24; + for (0..L) |i| layer_buf[i] = std.mem.readInt(u32, bytes[at + 4 * i ..][0..4], .little); + at += 4 * L; + var h: Host = .{ .layer_ids = layer_buf[0..L], .max_ngram = M, .heads = heads, .token_map = &.{}, .pad_id = 0 }; + for (0..L) |li| for (0..M) |s| { + h.mult[li][s] = std.mem.readInt(i64, bytes[at..][0..8], .little); + at += 8; + }; + if (bytes.len - at < 4 * @as(usize, V)) return error.BadEngramHost; + h.token_map = @alignCast(std.mem.bytesAsSlice(u32, bytes[at..][0 .. 4 * V])); + h.pad_id = h.token_map[pad_token]; + h.layout(vocab_size); + return h; + } +}; + +/// The committed ids before `start` (last `keep_n`): the tail, then kept leading rows of `pending`; null: unreachable. +pub fn lookbackAt(tail: []const u32, pos: u64, pending: ?[]const u32, start: u64, keep_n: usize, buf: []u32) ?[]const u32 { + if (start < pos) return null; + const k: usize = @intCast(start - pos); + const extra: []const u32 = if (pending) |p| (if (k <= p.len) p[0..k] else return null) else if (k == 0) &.{} else return null; + if (tail.len + extra.len > buf.len) return null; + @memcpy(buf[0..tail.len], tail); + @memcpy(buf[tail.len..][0..extra.len], extra); + const m = tail.len + extra.len; + const n = @min(m, keep_n); + return buf[m - n .. m]; +} + +/// A slot's Engram tail after keeping `ids`: the last `keep_n` of tail + ids, into `tail`; returns its new length. +pub fn tailAfter(tail: []u32, len: usize, keep_n: usize, ids: []const u32) usize { + var all: [2 * max_ngram + 64]u32 = undefined; + const old = tail[0..len]; + const take = ids[ids.len -| keep_n..]; + @memcpy(all[0..old.len], old); + @memcpy(all[old.len..][0..take.len], take); + const m = old.len + take.len; + const k = @min(m, keep_n); + @memcpy(tail[0..k], all[m - k .. m]); + return k; +} + +test "an Engram warm's lookback at a draft before its window's keep is the window's tail after it" { + // TF_DSV41_ENGRAM_WARM: warmEngram asks at pos + kept, pending window in flight; stageEngram reads after the keep + var prng = std.Random.DefaultPrng.init(56); + const rnd = prng.random(); + for (0..500) |_| { + const keep_n = 3; + var tail: [max_ngram]u32 = undefined; + const len = rnd.uintAtMost(usize, keep_n); + for (tail[0..len]) |*t| t.* = rnd.int(u32); + var pend: [6]u32 = undefined; + const pn = 1 + rnd.uintLessThan(usize, pend.len); + for (pend[0..pn]) |*t| t.* = rnd.int(u32); + const pos: u64 = 100 + rnd.uintAtMost(u64, 50); + const kept = rnd.uintLessThan(usize, pn); // drafts kept: rows 0..kept, the next window starts past them + const start = pos + kept + 1; + var buf: [2 * max_ngram + 64]u32 = undefined; + const before = lookbackAt(tail[0..len], pos, pend[0..pn], start, keep_n, &buf).?; + var after = tail; + const alen = tailAfter(&after, len, keep_n, pend[0 .. kept + 1]); + try std.testing.expectEqualSlices(u32, after[0..alen], before); + // a keep already done (the commit before the draft): no pending window, the same tail + var buf2: [2 * max_ngram + 64]u32 = undefined; + try std.testing.expectEqualSlices(u32, after[0..alen], lookbackAt(after[0..alen], start, null, start, keep_n, &buf2).?); + } +} + +/// Small primes, then Miller-Rabin with bases 2, 7, 61 (exact below 4.7e9). +pub fn isPrime(n: u64) bool { + if (n < 2) return false; + for ([_]u64{ 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37 }) |p| if (n % p == 0) return n == p; + var d = n - 1; + var r: u32 = 0; + while (d % 2 == 0) : (r += 1) d /= 2; + for ([_]u64{ 2, 7, 61 }) |a| { + var x = powMod(a, d, n); + if (x == 1 or x == n - 1) continue; + var ok = false; + var i: u32 = 1; + while (i < r) : (i += 1) { + x = @intCast(@as(u128, x) * x % n); + if (x == n - 1) { + ok = true; + break; + } + } + if (!ok) return false; + } + return true; +} + +fn powMod(b0: u64, e0: u64, m: u64) u64 { + var b: u128 = b0 % m; + var e = e0; + var r: u128 = 1; + while (e > 0) : (e >>= 1) { + if (e & 1 == 1) r = r * b % m; + b = b * b % m; + } + return @intCast(r); +} + +/// Synthetic rows: row i, value j = (((i * 2654435761 + j * 40503) % 511) - 255) / 1024, exact in bf16, into `out`. +pub fn indexRows(rows: []const i64, dim: usize, out: []u16) void { + for (rows, 0..) |i, r| for (0..dim) |j| { + const v = @mod(i *% 2654435761 +% @as(i64, @intCast(j)) * 40503, 511) - 255; + const f: f32 = @as(f32, @floatFromInt(v)) / 1024.0; + out[r * dim + j] = @truncate(@as(u32, @bitCast(f)) >> 16); + }; +} + +const testing = std.testing; + +test "the n-gram hashes: the reference NgramHasher on the release config and tokenizer (fixtures/engram-hash-ref.bin)" { + const a = testing.allocator; + const fx: []const u8 = @embedFile("fixtures/engram-hash-ref.bin"); + const L = std.mem.readInt(u32, fx[8..12], .little); + const M = std.mem.readInt(u32, fx[12..16], .little); + const V = std.mem.readInt(u32, fx[16..20], .little); + var at: usize = 24 + 4 * L + 8 * L * M; + // a full-size map holding only the entries the fixture's ids use + const map = try a.alloc(u32, V); + defer a.free(map); + @memset(map, 0); + const used = std.mem.readInt(u32, fx[at..][0..4], .little); + at += 4; + for (0..used) |_| { + map[std.mem.readInt(u32, fx[at..][0..4], .little)] = std.mem.readInt(u32, fx[at + 4 ..][0..4], .little); + at += 8; + } + var file: std.ArrayList(u8) = .empty; + defer file.deinit(a); + try file.appendSlice(a, fx[0 .. 24 + 4 * L + 8 * L * M]); + try file.appendSlice(a, std.mem.sliceAsBytes(map)); + const bytes = try a.alignedAlloc(u8, .@"4", file.items.len); + defer a.free(bytes); + @memcpy(bytes, file.items); + var ids_buf: [max_layers]u32 = undefined; + const h = try Host.load(bytes, 8, 2, 16_000_000, &ids_buf); + // the reference layout: the first and last primes, the last head's offset + try testing.expectEqual(@as(i64, 16_000_057), h.primes[0][0]); + try testing.expectEqual(@as(i64, 16_000_889), h.primes[1][23]); + try testing.expectEqual(@as(i64, 368_005_705), h.offsets[0][23]); + const cases = std.mem.readInt(u32, fx[at..][0..4], .little); + at += 4; + for (0..cases) |_| { + const n = std.mem.readInt(u32, fx[at..][0..4], .little); + const nb = std.mem.readInt(u32, fx[at + 4 ..][0..4], .little); + at += 8; + const ids_c = try a.alloc(u32, n); + defer a.free(ids_c); + for (ids_c, 0..) |*x, i| x.* = std.mem.readInt(u32, fx[at + 4 * i ..][0..4], .little); + at += 4 * n; + const lb_c = try a.alloc(u32, nb); + defer a.free(lb_c); + for (lb_c, 0..) |*x, i| x.* = std.mem.readInt(u32, fx[at + 4 * i ..][0..4], .little); + at += 4 * nb; + const got = try a.alloc(i64, n * L * h.cols()); + defer a.free(got); + h.hash(ids_c, lb_c, got); + for (got, 0..) |g, i| try testing.expectEqual(std.mem.readInt(i64, fx[at + 8 * i ..][0..8], .little), g); + at += 8 * got.len; + } + try testing.expectEqual(fx.len, at); + try testing.expectEqual([2]u32{ 12, 12 }, h.headShard(1, 2)); +} diff --git a/zig/src/families/deepseek_v41/engram_io.zig b/zig/src/families/deepseek_v41/engram_io.zig new file mode 100644 index 0000000000..60be56ea5a --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_io.zig @@ -0,0 +1,534 @@ +//! Engram shard reads: the batches, LRU tables and reader paths (io_uring, native AIO, a thread pool, serial). + +const std = @import("std"); +const linux = std.os.linux; +const engram_rows = @import("engram_rows.zig"); +const Rows = engram_rows.Rows; +const Header = engram_rows.Header; +const header_bytes = engram_rows.header_bytes; +const sector = engram_rows.sector; +const row_bytes = engram_rows.row_bytes; +const merge_gap = engram_rows.merge_gap; +const page = engram_rows.page; + +/// One extent of a batch: file bytes [off, off + len) into the batch buffer at `at`, its rows `first .. + count`. +const Extent = struct { off: u64, len: u64, at: usize, first: u32, count: u32 }; + +pub const Batch = struct { + table: u32, + /// sorted unique global rows + rows: []u64, + ext: []Extent, + buf: []align(page) u8, + left: u32, + failed: bool = false, + in_use: bool = true, + /// a `rows` call took at least half its rows (a bulk batch is read through once read); `taken`: this call's count + read: bool = false, + taken: u32 = 0, + unread: u8 = 0, // `rows` calls that settled it unread: a TF_DSV41_PREFETCH_AHEAD batch waits through one + + /// The record of `rows[k]` (k found by the caller). + pub fn record(b: *const Batch, k: usize, lo: u64) []const u8 { + var a: usize = 0; + var z: usize = b.ext.len; + while (z - a > 1) { + const m = (a + z) / 2; + if (b.ext[m].first <= k) a = m else z = m; + } + const e = b.ext[a]; + const at = e.at + @as(usize, @intCast(header_bytes + (b.rows[k] - lo) * row_bytes - e.off)); + return b.buf[at..][0..row_bytes]; + } +}; + +const Entry = union(enum) { ready: u32, pending: u32 }; +const none: u32 = std.math.maxInt(u32); + +/// One layer's shard and record cache. +pub const Table = struct { + fd: linux.fd_t, + direct: bool, + h: Header, + size: u64, + map: std.AutoHashMapUnmanaged(u64, Entry) = .empty, + /// cap records, their ids, and the LRU links (head: most recent) + recs: []u8, + ids: []u64, + prev: []u32, + next: []u32, + used: u32 = 0, + head: u32 = none, + tail: u32 = none, + /// keys removed since the map was last rehashed: each left a tombstone that a missing key's probe walks past + dead: u32 = 0, + + pub fn unlink(t: *Table, s: u32) void { + if (t.prev[s] != none) t.next[t.prev[s]] = t.next[s] else t.head = t.next[s]; + if (t.next[s] != none) t.prev[t.next[s]] = t.prev[s] else t.tail = t.prev[s]; + } + + pub fn front(t: *Table, s: u32) void { + t.prev[s] = none; + t.next[s] = t.head; + if (t.head != none) t.prev[t.head] = s; + t.head = s; + if (t.tail == none) t.tail = s; + } + + /// A slot for a new record: an unused one, else the least recent (its id leaves the map). + pub fn take(t: *Table) u32 { + if (t.used < t.ids.len) { + t.used += 1; + return t.used - 1; + } + const s = t.tail; + t.unlink(s); + t.drop(t.ids[s]); + return s; + } + + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + pub fn drop(t: *Table, x: u64) void { + if (!t.map.remove(x)) return; + t.dead += 1; + if (@as(u64, t.dead) * 10 < t.map.capacity()) return; + t.map.rehash(std.hash_map.AutoContext(u64){}); + t.dead = 0; + } +}; + +/// One extent for a pool thread: file bytes [off, off + buf.len) into `buf` (the batch's own buffer), batch `bi`. +const Job = struct { bi: u32, fd: linux.fd_t, size: u64, direct: bool, buf: []u8, off: u64 }; + +/// The reader threads without io_uring; `mu` guards the queue, every batch's `left` / `failed`, and the batch list. +pub const Pool = struct { + mu: std.c.pthread_mutex_t = std.c.PTHREAD_MUTEX_INITIALIZER, + work: std.c.pthread_cond_t = .{}, + done: std.c.pthread_cond_t = .{}, + jobs: std.ArrayList(Job) = .empty, + head: usize = 0, + stop: bool = false, + threads: []std.Thread = &.{}, + + pub fn lock(p: *Pool) void { + _ = std.c.pthread_mutex_lock(&p.mu); + } + + pub fn unlock(p: *Pool) void { + _ = std.c.pthread_mutex_unlock(&p.mu); + } +}; + +/// Linux native AIO (TF_DSV41_ENGRAM_AIO=1): one io_submit a batch, reaped by the waiter; works without io_uring. +pub const Aio = struct { + ctx: u64 = 0, + in_flight: u32 = 0, + poll: bool = false, // TF_DSV41_ENGRAM_POLL=1: a wait spins on io_getevents up to `poll_ns`, then sleeps + + pub const poll_ns: u64 = 20 * std.time.ns_per_ms; + + pub const depth: u32 = 1024; + + const Iocb = extern struct { + data: u64, + key: u32 = 0, + rw_flags: i32 = 0, + opcode: u16 = 0, // IOCB_CMD_PREAD + reqprio: i16 = 0, + fildes: u32, + buf: u64, + nbytes: u64, + offset: i64, + reserved2: u64 = 0, + flags: u32 = 0, + resfd: u32 = 0, + }; + + const Event = extern struct { data: u64, obj: u64, res: i64, res2: i64 }; + + pub fn init() !Aio { + var a: Aio = .{ .poll = pollEnv() }; + const rc = linux.syscall2(.io_setup, depth, @intFromPtr(&a.ctx)); + if (linux.errno(rc) != .SUCCESS) return error.AioRefused; + return a; + } + + pub fn deinit(a: *Aio) void { + _ = linux.syscall1(.io_destroy, a.ctx); + } +}; + +/// TF_DSV41_ENGRAM_POLL: 0 (default) | 1 (with TF_DSV41_ENGRAM_AIO=1). +pub fn pollEnv() bool { + const v = std.c.getenv("TF_DSV41_ENGRAM_POLL") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// TF_DSV41_ENGRAM_AIO: 0 (default) | 1. +pub fn aioEnv() bool { + const v = std.c.getenv("TF_DSV41_ENGRAM_AIO") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// TF_DSV41_ENGRAM_QD: the reader threads without io_uring (default 64; 0: none, reads serial). +pub fn poolThreads() u32 { + const v = std.c.getenv("TF_DSV41_ENGRAM_QD") orelse return 64; + return std.fmt.parseInt(u32, std.mem.span(v), 10) catch 64; +} + +/// Native AIO in place of io_uring or the pool (tests, the probe): every read async through one context. +pub fn startAio(r: *Rows) !void { + if (r.ring) |*g| g.deinit(); + r.ring = null; + r.stopPool(); + if (r.aio == null) r.aio = try Aio.init(); +} + +/// The reader threads (no io_uring): `n` workers waiting for extents. +pub fn startPool(r: *Rows, n: u32) !void { + if (r.pool != null or r.ring != null or n == 0) return error.EngramPool; + const p = try r.gpa.create(Pool); + errdefer r.gpa.destroy(p); + p.* = .{}; + p.threads = try r.gpa.alloc(std.Thread, n); + var made: usize = 0; + errdefer { + p.lock(); + p.stop = true; + p.unlock(); + _ = std.c.pthread_cond_broadcast(&p.work); + for (p.threads[0..made]) |t| t.join(); + r.gpa.free(p.threads); + } + r.pool = p; + errdefer r.pool = null; + for (p.threads) |*t| { + t.* = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, worker, .{r}); + made += 1; + } +} + +/// The reader threads stopped (their queue drained); reads are serial afterwards. +pub fn stopPool(r: *Rows) void { + const p = r.pool orelse return; + p.lock(); + p.stop = true; + p.unlock(); + _ = std.c.pthread_cond_broadcast(&p.work); + for (p.threads) |t| t.join(); + r.gpa.free(p.threads); + p.jobs.deinit(r.gpa); + r.gpa.destroy(p); + r.pool = null; +} + +/// A reader thread: extents off the queue until `stop` (the queue drained first). +pub fn worker(r: *Rows) void { + const p = r.pool.?; + while (true) { + p.lock(); + while (p.head == p.jobs.items.len and !p.stop) _ = std.c.pthread_cond_wait(&p.work, &p.mu); + if (p.head == p.jobs.items.len) { + p.unlock(); + return; + } + const j = p.jobs.items[p.head]; + p.head += 1; + if (p.head == p.jobs.items.len) { + p.jobs.clearRetainingCapacity(); + p.head = 0; + } + p.unlock(); + const ok = if (readAt(j.fd, j.size, j.direct, j.buf, j.off)) true else |_| false; + p.lock(); + const b = &r.batches.items[j.bi]; + if (!ok) b.failed = true; + b.left -= 1; + const last = b.left == 0; + p.unlock(); + if (last) _ = std.c.pthread_cond_broadcast(&p.done); + } +} + +pub fn openTable(gpa: std.mem.Allocator, path: [:0]const u8, cap: u32) !Table { + var direct = true; + var rc = linux.openat(linux.AT.FDCWD, path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }, 0); + if (linux.errno(rc) != .SUCCESS) { + direct = false; + rc = linux.openat(linux.AT.FDCWD, path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) { + std.log.scoped(.dsv41).err("engram: cannot open {s}", .{path}); + return error.FileNotFound; + } + } + const fd: linux.fd_t = @intCast(rc); + errdefer _ = linux.close(fd); + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true }, &st)) != .SUCCESS) return error.StatFailed; + const hb = try gpa.alignedAlloc(u8, .fromByteUnits(page), page); + defer gpa.free(hb); + try preadAll(fd, hb, 0); + if (!direct) _ = linux.fadvise(fd, 0, page, linux.POSIX_FADV.DONTNEED); + const h = try Header.parse(hb[0..48], st.size); + const t: Table = .{ + .fd = fd, + .direct = direct, + .h = h, + .size = st.size, + .recs = try gpa.alloc(u8, cap * row_bytes), + .ids = try gpa.alloc(u64, cap), + .prev = try gpa.alloc(u32, cap), + .next = try gpa.alloc(u32, cap), + }; + return t; +} + +pub fn closeTable(r: *Rows, t: *Table) void { + _ = linux.close(t.fd); + t.map.deinit(r.gpa); + r.gpa.free(t.recs); + r.gpa.free(t.ids); + r.gpa.free(t.prev); + r.gpa.free(t.next); +} + +/// A batch for sorted unique `rows` of table `ti`, its extents submitted; `urgent`: runs before queued read-ahead. +pub fn start(r: *Rows, ti: u32, rows_in: []const u64, urgent: bool) !u32 { + const t = &r.tables[ti]; + const gpa = r.gpa; + var ext: std.ArrayList(Extent) = .empty; + errdefer ext.deinit(gpa); + var total: usize = 0; + for (rows_in, 0..) |x, k| { + const s = header_bytes + (x - t.h.lo) * row_bytes; + const s0 = s / sector * sector; + const s1 = (s + row_bytes + sector - 1) / sector * sector; + if (ext.items.len > 0) { + const e = &ext.items[ext.items.len - 1]; + if (s0 <= e.off + e.len + merge_gap) { + const end = @max(e.off + e.len, s1); + total += @intCast(end - e.off - e.len); + e.len = end - e.off; + e.count += 1; + continue; + } + } + total = std.mem.alignForward(usize, total, page); + try ext.append(gpa, .{ .off = s0, .len = s1 - s0, .at = total, .first = @intCast(k), .count = 1 }); + total += @intCast(s1 - s0); + } + // extents keep their buffer offsets page aligned (O_DIRECT) + const buf = try gpa.alignedAlloc(u8, .fromByteUnits(page), std.mem.alignForward(usize, total, page)); + errdefer gpa.free(buf); + const owned = try gpa.dupe(u64, rows_in); + errdefer gpa.free(owned); + const b: Batch = .{ .table = ti, .rows = owned, .ext = try ext.toOwnedSlice(gpa), .buf = buf, .left = 0 }; + if (r.pool) |p| { + // the batch's slot and its extents queued under the pool's lock (a worker settles batches.items[bi]) + try t.map.ensureUnusedCapacity(gpa, @intCast(owned.len)); + p.lock(); + const queued = r.queue(p, t, b, urgent); + p.unlock(); + const bi = try queued; + _ = std.c.pthread_cond_broadcast(&p.work); + for (owned) |x| t.map.putAssumeCapacity(x, .{ .pending = bi }); + return bi; + } + const bi = try r.slotFor(b); + for (owned) |x| try t.map.put(gpa, x, .{ .pending = bi }); + const bp = &r.batches.items[bi]; + r.stats.reads += bp.ext.len; + for (bp.ext) |e| r.stats.bytes += e.len; + if (r.ring) |*g| { + for (bp.ext, 0..) |e, k| { + const sqe = g.read(@as(u64, bi) << 32 | k, t.fd, .{ .buffer = bp.buf[e.at..][0..@intCast(e.len)] }, e.off) catch blk: { + // a full submission queue: submit what is queued, take finished reads, then the slot is free + _ = try g.submit(); + try r.reap(0); + break :blk try g.read(@as(u64, bi) << 32 | k, t.fd, .{ .buffer = bp.buf[e.at..][0..@intCast(e.len)] }, e.off); + }; + _ = sqe; + bp.left += 1; + } + _ = try g.submit(); + } else if (r.aio != null) { + try r.aioSubmit(bi, t); + } else { + for (bp.ext) |e| try r.readExtent(t, bp.buf[e.at..][0..@intCast(e.len)], e.off); + } + return bi; +} + +pub fn slotFor(r: *Rows, b: Batch) !u32 { + for (r.batches.items, 0..) |*x, i| if (!x.in_use) { + x.* = b; + return @intCast(i); + }; + try r.batches.append(r.gpa, b); + return @intCast(r.batches.items.len - 1); +} + +/// A batch `b` in a free slot, its extents queued on `p` (locked): at the end, or `urgent` at the head. +pub fn queue(r: *Rows, p: *Pool, t: *const Table, b: Batch, urgent: bool) !u32 { + try p.jobs.ensureUnusedCapacity(r.gpa, b.ext.len); + const bi = try r.slotFor(b); + const bp = &r.batches.items[bi]; + r.stats.reads += bp.ext.len; + const n = bp.ext.len; + const at = if (urgent) p.head else p.jobs.items.len; + if (urgent) { + // one shift of the queued jobs past the head, then this batch's in their place + const old = p.jobs.items.len; + p.jobs.items.len += n; + std.mem.copyBackwards(Job, p.jobs.items[at + n ..], p.jobs.items[at..old]); + } else p.jobs.items.len += n; + for (bp.ext, p.jobs.items[at..][0..n]) |e, *j| { + r.stats.bytes += e.len; + j.* = .{ .bi = bi, .fd = t.fd, .size = t.size, .direct = t.direct, .buf = bp.buf[e.at..][0..@intCast(e.len)], .off = e.off }; + } + bp.left = @intCast(bp.ext.len); + return bi; +} + +pub fn readExtent(r: *Rows, t: *const Table, buf: []u8, off: u64) !void { + _ = r; + return readAt(t.fd, t.size, t.direct, buf, off); +} + +/// Batch `bi`'s extents on the AIO context (all of them, reaping finished reads while the context is full). +pub fn aioSubmit(r: *Rows, bi: u32, t: *const Table) !void { + const a = &r.aio.?; + const bp = &r.batches.items[bi]; + var cbs: [64]Aio.Iocb = undefined; + var ptrs: [64]*Aio.Iocb = undefined; + var k: usize = 0; + while (k < bp.ext.len) { + if (a.in_flight == Aio.depth) try r.reap(1); + const n = @min(bp.ext.len - k, cbs.len, Aio.depth - a.in_flight); + for (cbs[0..n], ptrs[0..n], bp.ext[k..][0..n], 0..) |*cb, *pp, e, j| { + cb.* = .{ .data = @as(u64, bi) << 32 | (k + j), .fildes = @intCast(t.fd), .buf = @intFromPtr(bp.buf[e.at..].ptr), .nbytes = e.len, .offset = @intCast(e.off) }; + pp.* = cb; + } + const rc = linux.syscall3(.io_submit, a.ctx, n, @intFromPtr(&ptrs)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .AGAIN, .INTR => { + try r.reap(if (a.in_flight > 0) 1 else 0); + continue; + }, + else => return error.EngramReadFailed, + } + bp.left += @intCast(rc); + a.in_flight += @intCast(rc); + k += rc; + } +} + +/// Takes finished AIO reads (waiting for at least `at_least`). +pub fn aioReap(r: *Rows, at_least: u32) !void { + const a = &r.aio.?; + var evs: [64]Aio.Event = undefined; + var need = @min(at_least, a.in_flight); + const zero: linux.timespec = .{ .sec = 0, .nsec = 0 }; + const t0 = if (a.poll and need > 0) nowNs() else 0; + var polling = a.poll and need > 0; + while (true) { + // polling: whatever has landed, at once; else the kernel's wait for `need` + if (polling and nowNs() - t0 > Aio.poll_ns) polling = false; + const rc = if (polling) + linux.syscall5(.io_getevents, a.ctx, 0, evs.len, @intFromPtr(&evs), @intFromPtr(&zero)) + else + linux.syscall5(.io_getevents, a.ctx, need, evs.len, @intFromPtr(&evs), 0); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.EngramReadFailed, + } + for (evs[0..rc]) |ev| { + const b = &r.batches.items[@intCast(ev.data >> 32)]; + const e = b.ext[@intCast(ev.data & 0xFFFF_FFFF)]; + const t = &r.tables[b.table]; + const want = @min(e.len, t.size - e.off); + if (ev.res < 0 or @as(u64, @intCast(ev.res)) < want) b.failed = true; + if (!t.direct) _ = linux.fadvise(t.fd, @intCast(e.off), @intCast(e.len), linux.POSIX_FADV.DONTNEED); + b.left -= 1; + a.in_flight -= 1; + } + need -|= @intCast(rc); + if (need == 0) return; + if (polling and rc == 0) std.atomic.spinLoopHint(); + } +} + +/// Takes finished reads off the ring (waiting for at least `at_least`). +pub fn reap(r: *Rows, at_least: u32) !void { + if (r.aio != null) return r.aioReap(at_least); + const g = &(r.ring orelse return); + var cqes: [64]linux.io_uring_cqe = undefined; + var need = at_least; + while (true) { + const n = try g.copy_cqes(&cqes, need); + for (cqes[0..n]) |c| { + const b = &r.batches.items[@intCast(c.user_data >> 32)]; + const e = b.ext[@intCast(c.user_data & 0xFFFF_FFFF)]; + const t = &r.tables[b.table]; + const want = @min(e.len, t.size - e.off); + if (c.res < 0 or @as(u64, @intCast(c.res)) < want) b.failed = true; + if (!t.direct) _ = linux.fadvise(t.fd, @intCast(e.off), @intCast(e.len), linux.POSIX_FADV.DONTNEED); + b.left -= 1; + } + need -|= n; + if (need == 0) return; + } +} + +/// Blocks until batch `bi`'s reads are done. +pub fn wait(r: *Rows, bi: u32) !void { + if (r.pool) |p| { + p.lock(); + defer p.unlock(); + while (r.batches.items[bi].left > 0) _ = std.c.pthread_cond_wait(&p.done, &p.mu); + if (r.batches.items[bi].failed) return error.EngramReadFailed; + return; + } + while (r.batches.items[bi].left > 0) try r.reap(1); + if (r.batches.items[bi].failed) return error.EngramReadFailed; +} + +/// An extent's bytes [off, off + buf.len) of a shard of `size` bytes (stops at its end), pages dropped if buffered. +pub fn readAt(fd: linux.fd_t, size: u64, direct: bool, buf: []u8, off: u64) !void { + const want = @min(buf.len, size - off); + var done: usize = 0; + while (done < want) { + const rc = linux.pread(fd, buf[done..].ptr, buf.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.EngramReadFailed, + } + if (rc == 0) return error.EngramShortRead; + done += rc; + } + if (!direct) _ = linux.fadvise(fd, @intCast(off), @intCast(buf.len), linux.POSIX_FADV.DONTNEED); +} + +pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { + var done: usize = 0; + while (done < out.len) { + const rc = linux.pread(fd, out[done..].ptr, out.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.EngramReadFailed, + } + if (rc == 0) return error.EngramShortRead; + done += rc; + } +} + +pub fn nowNs() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} diff --git a/zig/src/families/deepseek_v41/engram_rows.zig b/zig/src/families/deepseek_v41/engram_rows.zig new file mode 100644 index 0000000000..7854df9ec0 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_rows.zig @@ -0,0 +1,306 @@ +//! Engram table rows from packed per-rank shards on local NVMe, through io_uring, AIO or reader threads, LRU-cached. + +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel +const linux = std.os.linux; +const eh = @import("engram_host.zig"); + +pub const magic: u64 = 0x31344E4531565344; // "DSV41EN1"; file `/engram-l-rof.bin` +pub const header_bytes: u64 = 4096; // 6 little-endian u64: magic, layer, lo, hi, total rows, row bytes +pub const sector: u64 = 512; +pub const row_bytes: u64 = 264; // 256 e4m3 values and 8 UE8M0 scales (one a 32) +pub const values = 256; +pub const merge_gap: u64 = 4096; // holes up to this merged into one extent +pub const page = 4096; +const engram_io = @import("engram_io.zig"); + +pub const Header = struct { + layer: u64, + lo: u64, + hi: u64, + total: u64, + row_bytes: u64, + + /// The header's six words, checked against the file's size. + pub fn parse(b: *const [48]u8, size: u64) !Header { + var w: [6]u64 = undefined; + for (&w, 0..) |*x, i| x.* = std.mem.readInt(u64, b[8 * i ..][0..8], .little); + if (w[0] != magic) return error.NotEngramShard; + const h: Header = .{ .layer = w[1], .lo = w[2], .hi = w[3], .total = w[4], .row_bytes = w[5] }; + if (h.row_bytes != row_bytes or h.hi < h.lo or size != header_bytes + (h.hi - h.lo) * h.row_bytes) return error.BadEngramShard; + return h; + } +}; + +/// torch's float8_e4m3fn -> float32 (c10 fp8e4m3fn_to_fp32_value, bit for bit). +pub fn e4m3(b: u8) f32 { + const w: u32 = @as(u32, b) << 24; + const sign = w & 0x8000_0000; + const nonsign = w & 0x7FFF_FFFF; + var shift: u32 = @clz(nonsign); + shift = if (shift > 4) shift - 4 else 0; + const inf_nan: u32 = @bitCast((@as(i32, @bitCast(nonsign +% 0x0100_0000)) >> 8) & 0x7F80_0000); + const zero: u32 = @bitCast(@as(i32, @bitCast(nonsign -% 1)) >> 31); + const r = sign | ((((nonsign << @intCast(shift)) >> 4) +% ((0x78 -% shift) << 23) | inf_nan) & ~zero); + return @bitCast(r); +} + +/// float32 -> bfloat16 as torch's CPU vector path: round to nearest even, NaN -> 0xFFFF (no packed table holds a NaN). +pub fn bf16(x: f32) u16 { + if (std.math.isNan(x)) return 0xFFFF; + const u: u32 = @bitCast(x); + return @intCast((u +% (0x7FFF + ((u >> 16) & 1))) >> 16); +} + +/// The table [scale << 8 | byte] -> bf16 bits of e4m3(byte) x 2^(scale - 127) (scale 0: x 0). +pub fn fillTable(t: *[65536]u16) void { + for (0..256) |s| { + const mul: f32 = @bitCast(@as(u32, @intCast(s)) << 23); + for (0..256) |b| t[s << 8 | b] = bf16(e4m3(@intCast(b)) * mul); + } +} + +const Batch = engram_io.Batch; +const Table = engram_io.Table; +const Pool = engram_io.Pool; +const Aio = engram_io.Aio; +pub const pollEnv = engram_io.pollEnv; +pub const aioEnv = engram_io.aioEnv; +pub const poolThreads = engram_io.poolThreads; +const nowNs = engram_io.nowNs; + +pub const Stats = struct { hits: u64 = 0, waited: u64 = 0, misses: u64 = 0, issued: u64 = 0, reads: u64 = 0, bytes: u64 = 0, wait_ns: u64 = 0 }; + +pub const Rows = struct { + gpa: std.mem.Allocator, + tables: []Table, + ring: ?linux.IoUring, + /// no io_uring: the reader threads (null: serial reads on the engine thread) + pool: ?*Pool = null, + /// no io_uring, TF_DSV41_ENGRAM_AIO=1: native AIO (in place of the pool) + aio: ?Aio = null, + batches: std.ArrayList(Batch) = .empty, + lut: *[65536]u16, + cap: u32, + stats: Stats = .{}, + /// scratch reused by `rows` (sorted unique ids, their records) + uniq: std.ArrayList(u64) = .empty, + recs: std.ArrayList([*]const u8) = .empty, + want: std.ArrayList(u32) = .empty, + + /// Opens rank `rank` of `world`'s shard of every Engram layer in `layer_ids` under `dir`, `cap` records a layer. + pub fn open(gpa: std.mem.Allocator, dir: []const u8, rank: u32, world: u32, layer_ids: []const u32, cap: u32) !*Rows { + if (cap < 4096) return error.EngramCacheTooSmall; + const r = try gpa.create(Rows); + errdefer gpa.destroy(r); + r.* = .{ .gpa = gpa, .tables = try gpa.alloc(Table, layer_ids.len), .ring = null, .lut = try gpa.create([65536]u16), .cap = cap }; + fillTable(r.lut); + var made: usize = 0; + errdefer { + for (r.tables[0..made]) |*t| r.closeTable(t); + gpa.free(r.tables); + gpa.destroy(r.lut); + } + for (layer_ids, r.tables) |l, *t| { + var nb: [4096]u8 = undefined; + const p = try std.fmt.bufPrint(nb[0 .. nb.len - 1], "{s}/engram-l{d}-r{d}of{d}.bin", .{ dir, l, rank, world }); + nb[p.len] = 0; + const path = nb[0..p.len :0]; + t.* = try openTable(gpa, path, cap); + made += 1; + if (t.h.layer != l) return error.BadEngramShard; + } + // TF_DSV41_ENGRAM_URING=0: the reader threads even where io_uring works (a container that allows it, A/B) + const uring_off = if (std.c.getenv("TF_DSV41_ENGRAM_URING")) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; + r.ring = (if (uring_off) error.PermissionDenied else linux.IoUring.init(256, 0)) catch |e| blk: { + if (aioEnv()) { + if (Aio.init()) |a| { + r.aio = a; + std.log.scoped(.dsv41).info("engram: no io_uring ({t}); rows read by native AIO (TF_DSV41_ENGRAM_AIO, {d} in flight), waits {s}", .{ e, Aio.depth, if (a.poll) "polled (TF_DSV41_ENGRAM_POLL)" else "asleep in the kernel" }); + break :blk null; + } else |ae| if (!quiet) std.log.scoped(.dsv41).warn("engram: TF_DSV41_ENGRAM_AIO=1 but io_setup failed ({t}): the reader threads", .{ae}); + } + const n = poolThreads(); + if (n > 0) { + try r.startPool(n); + std.log.scoped(.dsv41).info("engram: no io_uring ({t}); rows read by {d} threads (TF_DSV41_ENGRAM_QD)", .{ e, n }); + } else if (!quiet) std.log.scoped(.dsv41).warn("engram: no io_uring ({t}) and TF_DSV41_ENGRAM_QD=0; rows are read when a window needs them", .{e}); + break :blk null; + }; + return r; + } + + pub const startAio = engram_io.startAio; + pub const startPool = engram_io.startPool; + pub const stopPool = engram_io.stopPool; + const worker = engram_io.worker; + const openTable = engram_io.openTable; + const closeTable = engram_io.closeTable; + + pub fn close(r: *Rows) void { + // every read in flight lands before its buffer goes + for (r.batches.items, 0..) |*b, i| if (b.in_use) { + r.wait(@intCast(i)) catch {}; + b.read = true; + }; + r.settle(null); + r.stopPool(); + if (r.ring) |*g| g.deinit(); + if (r.aio) |*a| a.deinit(); + for (r.tables) |*t| r.closeTable(t); + r.gpa.free(r.tables); + r.batches.deinit(r.gpa); + r.uniq.deinit(r.gpa); + r.recs.deinit(r.gpa); + r.want.deinit(r.gpa); + r.gpa.destroy(r.lut); + r.gpa.destroy(r); + } + + /// engram_host's interface to these rows. + pub fn source(r: *Rows) eh.Source { + return .{ .ptr = r, .vtable = &.{ .issue = issueFn, .rows = rowsFn } }; + } + + fn issueFn(p: *anyopaque, li: usize, idx: []const i64) anyerror!void { + const r: *Rows = @ptrCast(@alignCast(p)); + return r.issue(li, idx); + } + + fn rowsFn(p: *anyopaque, li: usize, idx: []const i64, dim: usize, out: []u16) anyerror!void { + const r: *Rows = @ptrCast(@alignCast(p)); + return r.rows(li, idx, dim, out); + } + + /// Sorted unique `idx` into r.uniq (checked against the shard's range). + fn unique(r: *Rows, t: *const Table, idx: []const i64) !void { + r.uniq.clearRetainingCapacity(); + try r.uniq.ensureTotalCapacity(r.gpa, idx.len); + for (idx) |i| { + if (i < 0 or @as(u64, @intCast(i)) < t.h.lo or @as(u64, @intCast(i)) >= t.h.hi) return error.EngramRowOutsideShard; + r.uniq.appendAssumeCapacity(@intCast(i)); + } + std.mem.sort(u64, r.uniq.items, {}, std.sort.asc(u64)); + var m: usize = 0; + for (r.uniq.items) |x| if (m == 0 or r.uniq.items[m - 1] != x) { + r.uniq.items[m] = x; + m += 1; + }; + r.uniq.shrinkRetainingCapacity(m); + } + + /// Starts reading layer `li`'s rows `idx` that are neither cached nor in flight; returns at once. + pub fn issue(r: *Rows, li: usize, idx: []const i64) !void { + const t = &r.tables[li]; + try r.unique(t, idx); + var m: usize = 0; + for (r.uniq.items) |x| if (!t.map.contains(x)) { + r.uniq.items[m] = x; + m += 1; + }; + if (m == 0) return; + r.stats.issued += m; + _ = try r.start(@intCast(li), r.uniq.items[0..m], false); + } + const start = engram_io.start; + pub const slotFor = engram_io.slotFor; + pub const queue = engram_io.queue; + pub const readExtent = engram_io.readExtent; + pub const aioSubmit = engram_io.aioSubmit; + pub const aioReap = engram_io.aioReap; + pub const reap = engram_io.reap; + const wait = engram_io.wait; + + /// Finished batches into the caches (bulk read through); `li`: the table whose `rows` call settles (null: none). + fn settle(r: *Rows, li: ?usize) void { + // the pool's workers write `left` / `failed`: settled under its lock (the batches' list stays put) + if (r.pool) |p| p.lock(); + defer if (r.pool) |p| p.unlock(); + for (r.batches.items) |*b| { + if (!b.in_use or b.left > 0) continue; + const t = &r.tables[b.table]; + const keep = !b.failed and b.rows.len <= r.cap / 4; + // a settle on another table never drops it: segments read layer 1 then 14, so the next 14 batch sees both + if (!keep and !b.failed and !b.read and (b.unread < 1 or li != b.table)) { + if (li == b.table) b.unread += 1; + continue; + } + for (b.rows, 0..) |x, k| { + if (!keep) { + t.drop(x); + continue; + } + const s = t.take(); + @memcpy(t.recs[s * row_bytes ..][0..row_bytes], b.record(k, t.h.lo)); + t.ids[s] = x; + t.front(s); + // the entry is this batch's (a row is in one batch at most; evictions only remove ready entries) + t.map.getPtr(x).?.* = .{ .ready = s }; + } + r.gpa.free(b.rows); + r.gpa.free(b.ext); + r.gpa.free(b.buf); + b.in_use = false; + } + } + + /// Layer `li`'s rows `idx` (global row ids this rank owns, any order, duplicates allowed) as bf16 [idx.len, dim]. + pub fn rows(r: *Rows, li: usize, idx: []const i64, dim: usize, out: []u16) !void { + if (dim > values or dim % 32 != 0 or out.len != idx.len * dim) return error.BadEngramRows; + const t = &r.tables[li]; + try r.unique(t, idx); + const u = r.uniq.items; + try r.recs.resize(r.gpa, u.len); + r.want.clearRetainingCapacity(); + // cached records now; in-flight ones after their batches; the rest read as one batch + var miss: std.ArrayList(u64) = .empty; + defer miss.deinit(r.gpa); + for (u, 0..) |x, k| { + if (t.map.get(x)) |e| switch (e) { + .ready => |s| { + r.recs.items[k] = t.recs[s * row_bytes ..].ptr; + t.unlink(s); + t.front(s); + r.stats.hits += 1; + }, + .pending => |bi| { + r.batches.items[bi].taken += 1; + if (std.mem.indexOfScalar(u32, r.want.items, bi) == null) try r.want.append(r.gpa, bi); + r.stats.waited += 1; + }, + } else try miss.append(r.gpa, x); + } + if (miss.items.len > 0) { + r.stats.misses += miss.items.len; + const bi = try r.start(@intCast(li), miss.items, true); + r.batches.items[bi].read = true; + try r.want.append(r.gpa, bi); + } + const t0 = nowNs(); + for (r.want.items) |bi| { + try r.wait(bi); + const b = &r.batches.items[bi]; + if (2 * @as(usize, b.taken) >= b.rows.len) b.read = true; + b.taken = 0; + } + r.stats.wait_ns += nowNs() - t0; + for (u, 0..) |x, k| switch (t.map.get(x).?) { + .ready => {}, + .pending => |bi| { + const b = &r.batches.items[bi]; + r.recs.items[k] = b.record(std.sort.binarySearch(u64, b.rows, x, orderU64).?, t.h.lo).ptr; + }, + }; + // every position: its record's values through the table (scale byte of each 32 values, then the e4m3 byte) + for (idx, 0..) |i, p| { + const rec = r.recs.items[std.sort.binarySearch(u64, u, @as(u64, @intCast(i)), orderU64).?]; + const o = out[p * dim ..][0..dim]; + for (o, 0..) |*v, j| v.* = r.lut[@as(usize, rec[values + j / 32]) << 8 | rec[j]]; + } + r.settle(li); + } +}; + +fn orderU64(a: u64, b: u64) std.math.Order { + return std.math.order(a, b); +} diff --git a/zig/src/families/deepseek_v41/engram_rows_test.zig b/zig/src/families/deepseek_v41/engram_rows_test.zig new file mode 100644 index 0000000000..b750d6143a --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_rows_test.zig @@ -0,0 +1,383 @@ +//! engram_rows.zig against reference Engram host fixtures: bf16 table, packed shard via sync, threads and io_uring. + +const std = @import("std"); +const testing = std.testing; +const er = @import("engram_rows.zig"); +const eh = @import("engram_host.zig"); + +const ref = @embedFile("fixtures/engram-rows-ref.txt"); + +/// engram_fixture.record: byte j of global row r. +fn recordByte(r: u64, j: u64) u8 { + const h = (r * 2654435761 + j * 40503) % (1 << 32); + return if (j < 256) @truncate(h >> 7) else @intCast(100 + (h >> 9) % 40); +} + +const Fx = struct { lo: u64, rows: u64, total: u64, sets: [3]struct { idx: []i64, sha: [32]u8 } }; + +fn parseFx(a: std.mem.Allocator, lut_sha: *[32]u8) !Fx { + var fx: Fx = undefined; + var set: usize = 0; + var it = std.mem.tokenizeScalar(u8, ref, '\n'); + while (it.next()) |line| { + var w = std.mem.tokenizeScalar(u8, line, ' '); + const kind = w.next().?; + if (std.mem.eql(u8, kind, "lut")) { + _ = try std.fmt.hexToBytes(lut_sha, w.next().?); + } else if (std.mem.eql(u8, kind, "shard")) { + fx.lo = try std.fmt.parseInt(u64, w.next().?, 10); + fx.rows = try std.fmt.parseInt(u64, w.next().?, 10); + fx.total = try std.fmt.parseInt(u64, w.next().?, 10); + } else { + var ids: std.ArrayList(i64) = .empty; + var c = std.mem.tokenizeScalar(u8, w.next().?, ','); + while (c.next()) |x| try ids.append(a, try std.fmt.parseInt(i64, x, 10)); + fx.sets[set].idx = try ids.toOwnedSlice(a); + _ = try std.fmt.hexToBytes(&fx.sets[set].sha, w.next().?); + set += 1; + } + } + return fx; +} + +/// Writes the fixture's shard (the packer's layout) as engram-l1-r0of2.bin under `dir`. +fn writeShard(a: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, fx: Fx) !void { + return writeShardAs(a, io, dir, fx, 1); +} + +fn writeShardAs(a: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, fx: Fx, layer: u64) !void { + const n: usize = @intCast(er.header_bytes + fx.rows * er.row_bytes); + const b = try a.alloc(u8, n); + defer a.free(b); + @memset(b, 0); + for ([_]u64{ er.magic, layer, fx.lo, fx.lo + fx.rows, fx.total, er.row_bytes }, 0..) |v, i| std.mem.writeInt(u64, b[8 * i ..][0..8], v, .little); + for (0..fx.rows) |k| for (0..er.row_bytes) |j| { + b[er.header_bytes + k * er.row_bytes + j] = recordByte(fx.lo + k, j); + }; + var nb: [32]u8 = undefined; + try dir.writeFile(io, .{ .sub_path = try std.fmt.bufPrint(&nb, "engram-l{d}-r0of2.bin", .{layer}), .data = b }); +} + +fn sha(out: []const u16) [32]u8 { + var d: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(std.mem.sliceAsBytes(out), &d, .{}); + return d; +} + +test "engram rows: the bf16 table equals torch's dequant(...).to(bfloat16)" { + var want: [32]u8 = undefined; + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + _ = try parseFx(arena.allocator(), &want); + const t = try testing.allocator.create([65536]u16); + defer testing.allocator.destroy(t); + er.fillTable(t); + try testing.expectEqual(want, sha(t)); + // spot checks: 1.0 x 2^0, -0.0, the smallest subnormal x 2^-6, scale byte 0 (x 0) + try testing.expectEqual(@as(u16, 0x3F80), t[127 << 8 | 0x38]); + try testing.expectEqual(@as(u16, 0x8000), t[127 << 8 | 0x80]); + try testing.expectEqual(@as(u16, 0x3800), t[121 << 8 | 0x01]); // 2^-9 x 2^-6 + try testing.expectEqual(@as(u16, 0x0000), t[0 << 8 | 0x38]); +} + +/// How a test reads: serially on the calling thread, by the reader threads (no io_uring), or by io_uring. +const Mode = enum { sync, pool, uring, aio, aio_poll }; +const modes = [_]Mode{ .sync, .pool, .uring, .aio, .aio_poll }; + +fn openRows(a: std.mem.Allocator, io: std.Io, tmp: *testing.TmpDir, mode: Mode) !*er.Rows { + const root = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(root); + const r = try er.Rows.open(a, root, 0, 2, &.{1}, 4096); + if (mode == .aio or mode == .aio_poll) { + try r.startAio(); // TF_DSV41_ENGRAM_AIO's reader: native AIO, no threads + r.aio.?.poll = mode == .aio_poll; // + TF_DSV41_ENGRAM_POLL: waits polled + return r; + } + if (r.aio) |*x| { // TF_DSV41_ENGRAM_AIO=1 in the environment: the other modes without it + _ = x; + r.aio.?.deinit(); + r.aio = null; + } + if (mode != .uring) if (r.ring) |*g| { + g.deinit(); + r.ring = null; + }; + if (mode == .pool and r.pool == null) try r.startPool(8); + if (mode == .sync) r.stopPool(); // no io_uring here: open started the pool + return r; +} + +test "engram rows: a packed shard read as the reference ShardRows.rows(...).to(bfloat16), synchronously, by the reader threads and by io_uring" { + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShard(a, io, tmp.dir, fx); + for (modes) |mode| { + const r = try openRows(a, io, &tmp, mode); + defer r.close(); + try testing.expectEqual(fx.lo, r.tables[0].h.lo); + for (fx.sets) |s| { + const out = try a.alloc(u16, s.idx.len * 256); + defer a.free(out); + try r.rows(0, s.idx, 256, out); + try testing.expectEqual(s.sha, sha(out)); + // again: every record from the cache now, the same bits + const hits = r.stats.hits; + try r.rows(0, s.idx, 256, out); + try testing.expectEqual(s.sha, sha(out)); + try testing.expect(r.stats.hits > hits); + } + try testing.expectError(error.EngramRowOutsideShard, r.rows(0, &.{@intCast(fx.lo - 1)}, 256, try arena.allocator().alloc(u16, 256))); + } +} + +test "engram rows: issued ahead, then waited for; bulk reads past the cache; eviction keeps the bits" { + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShard(a, io, tmp.dir, fx); + for (modes) |mode| { + const r = try openRows(a, io, &tmp, mode); + defer r.close(); + const s = fx.sets[1]; + try r.issue(0, s.idx); + try testing.expectEqual(@as(u64, s.idx.len), r.stats.issued); + try r.issue(0, s.idx); // nothing new: every row is in flight or cached + try testing.expectEqual(@as(u64, s.idx.len), r.stats.issued); + const out = try a.alloc(u16, s.idx.len * 256); + defer a.free(out); + try r.rows(0, s.idx, 256, out); + try testing.expectEqual(s.sha, sha(out)); + try testing.expectEqual(@as(u64, 0), r.stats.misses); + // every row of the shard once (a prompt's bulk read: past cap / 4, read through), then the sets again + const all = try a.alloc(i64, fx.rows); + defer a.free(all); + for (all, 0..) |*x, k| x.* = @intCast(fx.lo + fx.rows - 1 - k); + const big = try a.alloc(u16, all.len * 256); + defer a.free(big); + try r.rows(0, all, 256, big); + for (all, 0..) |x, k| for (0..256) |j| { + const row: u64 = @intCast(x); + const v = recordByte(row, j); + const sc = recordByte(row, 256 + j / 32); + try testing.expectEqual(r.lut[@as(usize, sc) << 8 | v], big[k * 256 + j]); + }; + for (fx.sets) |q| { + const o = try a.alloc(u16, q.idx.len * 256); + defer a.free(o); + try r.rows(0, q.idx, 256, o); + try testing.expectEqual(q.sha, sha(o)); + } + try testing.expect(r.tables[0].used <= 4096); + } +} + +test "engram rows: a prompt segment's bulk reads issued ahead wait through one unrelated read, then serve it" { + // TF_DSV41_PREFETCH_AHEAD: segment k + 1's read issued before k reads; k must not drop it, k + 1 finds every row + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShard(a, io, tmp.dir, fx); + const all = try a.alloc(i64, fx.rows); + defer a.free(all); + for (all, 0..) |*x, k| x.* = @intCast(fx.lo + k); + try testing.expect(all.len > 4096 / 4); // a bulk batch: read through, not cached + const big = try a.alloc(u16, all.len * 256); + defer a.free(big); + for (modes) |mode| { + const r = try openRows(a, io, &tmp, mode); + defer r.close(); + try r.issue(0, all); + const s = fx.sets[0]; + const o = try a.alloc(u16, s.idx.len * 256); + defer a.free(o); + try r.rows(0, s.idx, 256, o); // segment k: waits for the rows it shares with the batch, the batch stays + try testing.expectEqual(s.sha, sha(o)); + const reads = r.stats.reads; + try r.rows(0, all, 256, big); // segment k + 1: every row from the batch issued ahead + try testing.expectEqual(@as(u64, 0), r.stats.misses); + try testing.expectEqual(reads, r.stats.reads); + for (all, 0..) |x, k| for (0..256) |j| { + const row: u64 = @intCast(x); + try testing.expectEqual(r.lut[@as(usize, recordByte(row, 256 + j / 32)) << 8 | recordByte(row, j)], big[k * 256 + j]); + }; + // a wrong guess: issued, then two unrelated reads of its table, then it is gone (read again on demand) + try r.issue(0, all); + try r.rows(0, s.idx, 256, o); + try r.rows(0, s.idx, 256, o); + const misses = r.stats.misses; + try r.rows(0, all, 256, big); + try testing.expect(r.stats.misses > misses); + for (r.batches.items) |b| try testing.expect(!b.in_use); + } +} + +test "engram rows: a prompt segment ahead on two tables survives the segment's layer-1 then layer-14 reads" { + // both layers' batches of k + 1 issued, k reads 1 then 14, k + 1 reads 1 (settle) then 14: no miss on either table + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShardAs(a, io, tmp.dir, fx, 1); + try writeShardAs(a, io, tmp.dir, fx, 14); + const all = try a.alloc(i64, fx.rows); + defer a.free(all); + for (all, 0..) |*x, k| x.* = @intCast(fx.lo + k); + const big = try a.alloc(u16, all.len * 256); + defer a.free(big); + const s = fx.sets[0]; + const o = try a.alloc(u16, s.idx.len * 256); + defer a.free(o); + for (modes) |mode| { + const root = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(root); + const r = try er.Rows.open(a, root, 0, 2, &.{ 1, 14 }, 4096); + defer r.close(); + if (mode != .uring) if (r.ring) |*g| { + g.deinit(); + r.ring = null; + }; + if (mode == .pool and r.pool == null) try r.startPool(8); + if (mode == .sync) r.stopPool(); + try r.issue(0, all); + try r.issue(1, all); + for (0..2) |li| { + try r.rows(li, s.idx, 256, o); // segment k: a few shared rows of each table (waits for the batch) + try testing.expectEqual(s.sha, sha(o)); + } + const reads = r.stats.reads; + for (0..2) |li| { + try r.rows(li, all, 256, big); // segment k + 1: layer 1, then layer 14, all from the batches issued ahead + try testing.expectEqual(@as(u64, 0), r.stats.misses); + try testing.expectEqual(reads, r.stats.reads); + for (all, 0..) |x, k| for (0..256) |j| { + const row: u64 = @intCast(x); + try testing.expectEqual(r.lut[@as(usize, recordByte(row, 256 + j / 32)) << 8 | recordByte(row, j)], big[k * 256 + j]); + }; + } + for (r.batches.items) |b| try testing.expect(!b.in_use); + } +} + +test "engram rows: the reader threads with many batches in flight (issued, missed, waited for out of order)" { + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShard(a, io, tmp.dir, fx); + const r = try openRows(a, io, &tmp, .pool); + defer r.close(); + var rng = std.Random.DefaultPrng.init(4101); + const rand = rng.random(); + const sets = try a.alloc([]i64, 48); + defer { + for (sets) |q| a.free(q); + a.free(sets); + } + for (sets) |*q| { + q.* = try a.alloc(i64, 1 + rand.uintLessThan(usize, 300)); + for (q.*) |*x| x.* = @intCast(fx.lo + rand.uintLessThan(u64, fx.rows)); + } + // every set issued ahead, then read in reverse (later batches waited for first), each read checked value by value + for (sets) |q| try r.issue(0, q); + var k = sets.len; + while (k > 0) { + k -= 1; + const q = sets[k]; + const out = try a.alloc(u16, q.len * 256); + defer a.free(out); + try r.rows(0, q, 256, out); + for (q, 0..) |x, i| for (0..256) |j| { + const row: u64 = @intCast(x); + try testing.expectEqual(r.lut[@as(usize, recordByte(row, 256 + j / 32)) << 8 | recordByte(row, j)], out[i * 256 + j]); + }; + } + try testing.expect(r.stats.issued > 0); +} + +test "engram rows: shard headers" { + var b: [48]u8 = undefined; + for ([_]u64{ er.magic, 14, 10, 20, 100, er.row_bytes }, 0..) |v, i| std.mem.writeInt(u64, b[8 * i ..][0..8], v, .little); + const h = try er.Header.parse(&b, er.header_bytes + 10 * er.row_bytes); + try testing.expectEqual(@as(u64, 14), h.layer); + try testing.expectError(error.BadEngramShard, er.Header.parse(&b, er.header_bytes + 9 * er.row_bytes)); + b[0] ^= 1; + try testing.expectError(error.NotEngramShard, er.Header.parse(&b, er.header_bytes + 10 * er.row_bytes)); +} + +test "engram lookback: the tail, then the kept rows of the window in flight" { + var buf: [16]u32 = undefined; + // nothing in flight: the next window starts at the slot's position + try testing.expectEqualSlices(u32, &.{ 7, 8, 9 }, eh.lookbackAt(&.{ 7, 8, 9 }, 50, null, 50, 3, &buf).?); + try testing.expect(eh.lookbackAt(&.{ 7, 8, 9 }, 50, null, 51, 3, &buf) == null); + // a window [a, d1, d2] at 50 not kept yet; the next starts at 52: a and d1 were kept + try testing.expectEqualSlices(u32, &.{ 9, 100, 101 }, eh.lookbackAt(&.{ 7, 8, 9 }, 50, &.{ 100, 101, 102 }, 52, 3, &buf).?); + try testing.expectEqualSlices(u32, &.{ 100, 101, 102 }, eh.lookbackAt(&.{ 7, 8, 9 }, 50, &.{ 100, 101, 102 }, 53, 3, &buf).?); + try testing.expect(eh.lookbackAt(&.{ 7, 8, 9 }, 50, &.{ 100, 101, 102 }, 54, 3, &buf) == null); + // the sequence start: fewer than keep_n + try testing.expectEqualSlices(u32, &.{5}, eh.lookbackAt(&.{}, 0, &.{ 5, 6 }, 1, 3, &buf).?); +} + +/// Slots of the record map that hold neither a key nor a tombstone: a lookup that misses probes until it finds one. +fn freeSlots(r: *const er.Rows) usize { + const m = &r.tables[0].map; + const meta = m.metadata orelse return 0; + var n: usize = 0; + for (meta[0..m.capacity()]) |x| n += @intFromBool(x.isFree()); + return n; +} + +test "engram rows: a long run of bulk reads and evictions leaves the record map free slots (no probe of the whole map)" { + // bulk reads past cap / 4 (inserted, then dropped) and evictions, as a long uncached prefill does + const a = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const fx: Fx = .{ .lo = 1000, .rows = 20000, .total = 40000, .sets = undefined }; + try writeShard(a, io, tmp.dir, fx); + const root = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(root); + const r = try er.Rows.open(a, root, 0, 2, &.{1}, 4096); + defer r.close(); + const ids = try a.alloc(i64, 1500); + defer a.free(ids); + const out = try a.alloc(u16, ids.len * 256); + defer a.free(out); + var rng = std.Random.DefaultPrng.init(7); + // the worst sample: free slots against the map's slots (it grows meanwhile) + var worst: [2]usize = .{ 1, 1 }; + for (0..300) |i| { + // a prompt segment's bulk read (past cap / 4: read through, its keys dropped), then a decode window's few + const n: usize = if (i % 2 == 0) ids.len else 64; + for (ids[0..n]) |*x| x.* = @intCast(fx.lo + rng.random().uintLessThan(u64, fx.rows)); + try r.rows(0, ids[0..n], 256, out[0 .. n * 256]); + const all = r.tables[0].map.capacity(); + const free = freeSlots(r); + if (free * worst[1] < worst[0] * all) worst = .{ free, all }; + } + errdefer std.debug.print("record map: {d} free of {d} slots at the worst\n", .{ worst[0], worst[1] }); + try testing.expect(worst[0] * 10 >= worst[1]); +} diff --git a/zig/src/families/deepseek_v41/ex_stacks.zig b/zig/src/families/deepseek_v41/ex_stacks.zig new file mode 100644 index 0000000000..c04a0b8135 --- /dev/null +++ b/zig/src/families/deepseek_v41/ex_stacks.zig @@ -0,0 +1,139 @@ +//! TF_DSV41_EX_STACKS=1: routed and shared expert suh / svh uploaded back to back, so the forms' stacks are views. + +const std = @import("std"); +const buffers = @import("buffers.zig"); + +const log = std.log.scoped(.dsv41); + +pub fn enabled() bool { + const v = std.c.getenv("TF_DSV41_EX_STACKS") orelse return false; + return std.mem.eql(u8, std.mem.trim(u8, std.mem.span(v), " "), "1"); +} + +/// The shared expert's tensor for a routed suh / svh ("L3.moe.w1.suh" -> "L3.moe.shared.0.w1.suh"), else null. +pub fn sharedOf(buf: []u8, name: []const u8) ?[]const u8 { + const at = std.mem.indexOf(u8, name, ".moe.w") orelse return null; + const rest = name[at + ".moe.".len ..]; + if (rest.len != "w1.suh".len or rest[1] < '1' or rest[1] > '3') return null; + if (!std.mem.endsWith(u8, rest, ".suh") and !std.mem.endsWith(u8, rest, ".svh")) return null; + return std.fmt.bufPrint(buf, "{s}.moe.shared.0.{s}", .{ name[0..at], rest }) catch null; +} + +/// The routed weight starting a stack role ("s.L3.ex.suh_g" -> "L3.moe.w1.suh"; g w1, u w3, d w2), else null. +pub fn stackWeight(buf: []u8, role: []const u8) ?[]const u8 { + if (!std.mem.startsWith(u8, role, "s.L")) return null; + const ex = std.mem.indexOf(u8, role, ".ex.") orelse return null; + const what = role[ex + 4 ..]; + if (what.len != 5 or what[3] != '_') return null; + const kind = what[0..3]; + if (!std.mem.eql(u8, kind, "suh") and !std.mem.eql(u8, kind, "svh")) return null; + const w: []const u8 = switch (what[4]) { + 'g' => "w1", + 'u' => "w3", + 'd' => "w2", + else => return null, + }; + return std.fmt.bufPrint(buf, "{s}.moe.{s}.{s}", .{ role[2..ex], w, kind }) catch null; +} + +/// Index pairs (routed, shared) of `names` to upload together: both present and the routed one's length 16-aligned. +pub fn pairs(gpa: std.mem.Allocator, names: []const []const u8, lens: []const usize) ![][2]usize { + var at: std.StringHashMapUnmanaged(usize) = .empty; + defer at.deinit(gpa); + for (names, 0..) |n, i| try at.put(gpa, n, i); + var out: std.ArrayList([2]usize) = .empty; + var buf: [192]u8 = undefined; + for (names, 0..) |n, i| { + const s = sharedOf(&buf, n) orelse continue; + const j = at.get(s) orelse continue; + if (lens[i] % 16 != 0) continue; // the shared half must start where a kernel may read it as a tensor + try out.append(gpa, .{ i, j }); + } + return out.toOwnedSlice(gpa); +} + +pub const Bound = struct { roles: u32 = 0, bytes: u64 = 0 }; + +/// Binds `plan`'s stacks to their routed weight where the shared half follows (r: run.Runner, weights: load.Weights). +pub fn bind(r: anytype, weights: anytype, plan: *const buffers.Plan) !Bound { + var out: Bound = .{}; + var nb: [192]u8 = undefined; + var sb: [192]u8 = undefined; + for (plan.sizes.keys(), plan.sizes.values()) |role, size| { + const wn = stackWeight(&nb, role) orelse continue; + if (r.addressOf(role) != null) continue; + const t = weights.get(wn) orelse continue; + const sn = sharedOf(&sb, wn) orelse continue; + const s = weights.get(sn) orelse continue; + if (s.ptr != t.ptr + t.len or size > t.len + s.len) continue; + try r.external(role, t.ptr); + out.roles += 1; + out.bytes += t.len + s.len; + } + log.info("expert stacks (TF_DSV41_EX_STACKS): {d} stacks read in place, {d:.2} GiB not allocated", .{ out.roles, @as(f64, @floatFromInt(out.bytes)) / (1 << 30) }); + return out; +} + +const testing = std.testing; + +const Span = struct { ptr: u64, len: usize }; +const Weights = struct { + map: std.StringHashMapUnmanaged(Span) = .empty, + fn get(w: *const Weights, name: []const u8) ?Span { + return w.map.get(name); + } +}; +const Roles = struct { + gpa: std.mem.Allocator, + addrs: std.StringHashMapUnmanaged(u64) = .empty, + fn addressOf(r: *const Roles, role: []const u8) ?u64 { + return r.addrs.get(role); + } + fn external(r: *Roles, role: []const u8, addr: u64) !void { + try r.addrs.put(r.gpa, role, addr); + } +}; + +test "expert stacks: names, pairs, and a stack bound only where its shared half follows the routed one" { + const gpa = testing.allocator; + var b: [192]u8 = undefined; + try testing.expectEqualStrings("L3.moe.shared.0.w1.suh", sharedOf(&b, "L3.moe.w1.suh").?); + try testing.expectEqualStrings("L40.moe.shared.0.w2.svh", sharedOf(&b, "L40.moe.w2.svh").?); + try testing.expectEqual(@as(?[]const u8, null), sharedOf(&b, "L3.moe.w1.trellis")); + try testing.expectEqual(@as(?[]const u8, null), sharedOf(&b, "L3.moe.shared.0.w1.suh")); + try testing.expectEqualStrings("L3.moe.w1.suh", stackWeight(&b, "s.L3.ex.suh_g").?); + try testing.expectEqualStrings("L3.moe.w3.svh", stackWeight(&b, "s.L3.ex.svh_u").?); + try testing.expectEqualStrings("L3.moe.w2.suh", stackWeight(&b, "s.L3.ex.suh_d").?); + try testing.expectEqual(@as(?[]const u8, null), stackWeight(&b, "s.L3.ex.tp_g")); + try testing.expectEqual(@as(?[]const u8, null), stackWeight(&b, "s.L3.ex.k2_d")); + // the builder's order: a pair is found wherever its halves sit; an odd routed length is left alone + const names = [_][]const u8{ "L3.moe.shared.0.w1.suh", "L3.attn.wq_b.trellis", "L3.moe.w1.suh", "L3.moe.w2.svh", "L3.moe.shared.0.w2.svh" }; + const p = try pairs(gpa, &names, &.{ 64, 8, 4096, 4090, 64 }); + defer gpa.free(p); + try testing.expectEqual(@as(usize, 1), p.len); + try testing.expectEqual([2]usize{ 2, 0 }, p[0]); + // binding: suh_g's halves adjacent (bound), suh_u's apart (its own buffer), a stack larger than both (its own) + var w: Weights = .{}; + defer w.map.deinit(gpa); + try w.map.put(gpa, "L3.moe.w1.suh", .{ .ptr = 0x1000, .len = 4096 }); + try w.map.put(gpa, "L3.moe.shared.0.w1.suh", .{ .ptr = 0x2000, .len = 64 }); + try w.map.put(gpa, "L3.moe.w3.suh", .{ .ptr = 0x9000, .len = 4096 }); + try w.map.put(gpa, "L3.moe.shared.0.w3.suh", .{ .ptr = 0x20000, .len = 64 }); + try w.map.put(gpa, "L3.moe.w2.suh", .{ .ptr = 0x30000, .len = 4096 }); + try w.map.put(gpa, "L3.moe.shared.0.w2.suh", .{ .ptr = 0x31000, .len = 64 }); + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + var plan: buffers.Plan = .{ .a = arena.allocator() }; + try plan.sizes.put(arena.allocator(), "s.L3.ex.suh_g", 4096 + 64); + try plan.sizes.put(arena.allocator(), "s.L3.ex.suh_u", 4096 + 64); + try plan.sizes.put(arena.allocator(), "s.L3.ex.suh_d", 4096 + 128); + try plan.sizes.put(arena.allocator(), "s.L3.ex.tp_g", 64); + var r: Roles = .{ .gpa = gpa }; + defer r.addrs.deinit(gpa); + const got = try bind(&r, &w, &plan); + try testing.expectEqual(@as(u32, 1), got.roles); + try testing.expectEqual(@as(?u64, 0x1000), r.addressOf("s.L3.ex.suh_g")); + try testing.expectEqual(@as(?u64, null), r.addressOf("s.L3.ex.suh_u")); + try testing.expectEqual(@as(?u64, null), r.addressOf("s.L3.ex.suh_d")); + try testing.expectEqual(@as(?u64, null), r.addressOf("s.L3.ex.tp_g")); +} diff --git a/zig/src/families/deepseek_v41/exl3.zig b/zig/src/families/deepseek_v41/exl3.zig new file mode 100644 index 0000000000..68050611a4 --- /dev/null +++ b/zig/src/families/deepseek_v41/exl3.zig @@ -0,0 +1,254 @@ +//! EXL3 groups as the pack stores them: trellis geometry, codebooks, a rank's TP slice as byte spans, expert layouts. + +const std = @import("std"); + +/// The codebook a group's trellis decodes with: the `.mul1` / `.mcg` marker's value, else the 3-instruction default. +pub const Codebook = enum(u8) { + inst3, + mul1, + mcg, + + pub const marker_mul1: u32 = 0x83DCD12D; + pub const marker_mcg: u32 = 0xCBAC1FED; + + /// A marker tensor's value (the low 32 bits of its first element); unknown values are refused. + pub fn fromMarker(v: u32) !Codebook { + return switch (v) { + marker_mul1 => .mul1, + marker_mcg => .mcg, + else => error.UnknownCodebookMarker, + }; + } +}; + +/// How a rank holds a group under TP (every boundary a multiple of 128, so Hadamard blocks stay whole). +pub const Split = enum { + /// all of it, on every rank + whole, + /// this rank's output columns: N tiles and `svh` + col, + /// this rank's input rows: K tiles and `suh` (a partial sum the ranks exchange) + row, +}; + +/// One EXL3 matrix y = x @ W, W [K, N]: trellis int16 [K/16, N/16, words], suh fp16 [K], svh fp16 [N]. +pub const Group = struct { + k_tiles: u32, + n_tiles: u32, + /// int16 words a 16x16 tile: 16 x bits (the trellis's last dimension) + words: u32, + + /// From the trellis entry's shape; a tile must hold whole K2 bytes and 1 to 8 bits a value. + pub fn fromShape(shape: []const usize) !Group { + if (shape.len != 3) return error.BadTrellis; + const w = shape[2]; + if (w == 0 or w % 8 != 0 or w > 128 or shape[0] == 0 or shape[1] == 0) return error.BadTrellis; + return .{ .k_tiles = @intCast(shape[0]), .n_tiles = @intCast(shape[1]), .words = @intCast(w) }; + } + + pub fn k(g: Group) u32 { + return 16 * g.k_tiles; + } + + pub fn n(g: Group) u32 { + return 16 * g.n_tiles; + } + + /// Half-bits a value (the kernels' K2): 4 for 2 bits, 16 for 8 bits. + pub fn k2(g: Group) u32 { + return g.words / 8; + } + + /// Bits a value, in halves exact (k2 / 2). + pub fn bits(g: Group) f32 { + return @as(f32, @floatFromInt(g.words)) / 16.0; + } + + pub fn tileBytes(g: Group) u64 { + return @as(u64, g.words) * 2; + } + + pub fn trellisBytes(g: Group) u64 { + return @as(u64, g.k_tiles) * g.n_tiles * g.tileBytes(); + } + + /// The rank's part of the group; a split that does not fall on 128-value blocks is refused. + pub fn part(g: Group, split: Split, rank: u32, world: u32) !Part { + std.debug.assert(rank < world); + var p: Part = .{ .group = g, .trellis = .{ .offset = 0, .len = g.trellisBytes() }, .suh = .{ .offset = 0, .len = 2 * @as(u64, g.k()) }, .svh = .{ .offset = 0, .len = 2 * @as(u64, g.n()) } }; + switch (split) { + .whole => {}, + .row => { + if (g.k() % (world * 128) != 0) return error.SplitNotOnBlocks; + const per = g.k_tiles / world; + p.group.k_tiles = per; + const row = @as(u64, g.n_tiles) * g.tileBytes(); + p.trellis = .{ .offset = rank * per * row, .len = per * row }; + p.suh = .{ .offset = 2 * 16 * @as(u64, rank * per), .len = 2 * 16 * @as(u64, per) }; + }, + .col => { + if (g.n() % (world * 128) != 0) return error.SplitNotOnBlocks; + const per = g.n_tiles / world; + p.group.n_tiles = per; + // one K-tile row at a time: this rank's N tiles are a run inside each row + p.trellis = .{ .offset = rank * per * g.tileBytes(), .len = per * g.tileBytes(), .stride = @as(u64, g.n_tiles) * g.tileBytes(), .count = g.k_tiles }; + p.svh = .{ .offset = 2 * 16 * @as(u64, rank * per), .len = 2 * 16 * @as(u64, per) }; + }, + } + return p; + } +}; + +/// Bytes of a stored tensor to read: `count` runs of `len` bytes, `stride` apart, from `offset` into the tensor. +pub const Span = struct { + offset: u64, + len: u64, + stride: u64 = 0, + count: u64 = 1, + + pub fn bytes(s: Span) u64 { + return s.len * s.count; + } + + /// The `i`th run's start. + pub fn at(s: Span, i: u64) u64 { + return s.offset + i * s.stride; + } +}; + +/// A rank's slice of a group: its geometry and where its trellis, suh and svh bytes sit in the stored tensors. +pub const Part = struct { + group: Group, + trellis: Span, + suh: Span, + svh: Span, + + pub fn bytes(p: Part) u64 { + return p.trellis.bytes() + p.suh.bytes() + p.svh.bytes(); + } +}; + +/// Routed experts of one projection on a rank: one width (X3Stack) or a width per expert (X3Ragged, E2). +pub const Layout = struct { + kind: enum { stack, ragged }, + k_tiles: u32, + n_tiles: u32, + /// each expert's int16 words a tile + words: []u16, + /// each expert's first int16 element in the projection's trellis buffer (back to back, as `ragged()` packs them) + offsets: []u64, + /// int16 elements of every expert's trellis + elements: u64, + + /// The layout of the rank's expert parts; parts must share K and N (the kernels' tables index one shape). + pub fn of(gpa: std.mem.Allocator, parts: []const Part) !Layout { + if (parts.len == 0) return error.NoExperts; + const words = try gpa.alloc(u16, parts.len); + errdefer gpa.free(words); + const offsets = try gpa.alloc(u64, parts.len); + errdefer gpa.free(offsets); + const kt = parts[0].group.k_tiles; + const nt = parts[0].group.n_tiles; + var at: u64 = 0; + var same = true; + for (parts, words, offsets) |p, *w, *o| { + if (p.group.k_tiles != kt or p.group.n_tiles != nt) return error.ExpertShapesDiffer; + w.* = @intCast(p.group.words); + o.* = at; + at += @as(u64, kt) * nt * p.group.words; + same = same and p.group.words == parts[0].group.words; + } + return .{ .kind = if (same) .stack else .ragged, .k_tiles = kt, .n_tiles = nt, .words = words, .offsets = offsets, .elements = at }; + } + + pub fn deinit(l: *Layout, gpa: std.mem.Allocator) void { + gpa.free(l.words); + gpa.free(l.offsets); + l.* = undefined; + } + + pub fn count(l: *const Layout) usize { + return l.words.len; + } + + pub fn trellisBytes(l: *const Layout) u64 { + return 2 * l.elements; + } + + /// Byte offset of expert `e`'s trellis in the projection buffer: 16-byte aligned (a tile is 16 K2 bytes). + pub fn byteOffset(l: *const Layout, e: usize) u64 { + return 2 * l.offsets[e]; + } + + /// Experts at each K2 (index = K2, 1..16): the widths x3gm launches once each. + pub fn histogram(l: *const Layout) [17]u32 { + var h: [17]u32 = @splat(0); + for (l.words) |w| h[w / 8] += 1; + return h; + } + + /// Distinct widths: x3gm's launches for this projection. + pub fn widths(l: *const Layout) u32 { + var n: u32 = 0; + for (l.histogram()) |c| n += @intFromBool(c > 0); + return n; + } +}; + +const testing = std.testing; + +test "group geometry and K2" { + const g = try Group.fromShape(&.{ 320, 144, 32 }); + try testing.expectEqual(@as(u32, 5120), g.k()); + try testing.expectEqual(@as(u32, 2304), g.n()); + try testing.expectEqual(@as(u32, 4), g.k2()); + try testing.expectEqual(@as(f32, 2.0), g.bits()); + try testing.expectError(error.BadTrellis, Group.fromShape(&.{ 320, 144, 30 })); + try testing.expectError(error.BadTrellis, Group.fromShape(&.{ 320, 144 })); + try testing.expectEqual(Codebook.mcg, try Codebook.fromMarker(0xCBAC1FED)); + try testing.expectError(error.UnknownCodebookMarker, Codebook.fromMarker(7)); +} + +test "TP parts: rows contiguous, columns a run a K-tile row" { + const w2 = try Group.fromShape(&.{ 144, 320, 48 }); // a 3-bit routed down projection: 2304 -> 5120 + const r1 = try w2.part(.row, 1, 2); + try testing.expectEqual(@as(u32, 72), r1.group.k_tiles); + try testing.expectEqual(@as(u64, 72 * 320 * 96), r1.trellis.len); + try testing.expectEqual(r1.trellis.len, r1.trellis.offset); + try testing.expectEqual(Span{ .offset = 2 * 1152, .len = 2 * 1152 }, r1.suh); + try testing.expectEqual(@as(u64, 2 * 5120), r1.svh.len); + const w1 = try Group.fromShape(&.{ 320, 144, 32 }); + const c1 = try w1.part(.col, 1, 2); + try testing.expectEqual(@as(u32, 72), c1.group.n_tiles); + try testing.expectEqual(@as(u64, 320), c1.trellis.count); + try testing.expectEqual(@as(u64, 144 * 64), c1.trellis.stride); + try testing.expectEqual(@as(u64, 72 * 64), c1.trellis.offset); + try testing.expectEqual(c1.group.trellisBytes(), c1.trellis.bytes()); + // the 8-bit indexer key projection (512 -> 128) cannot split its 128 outputs over two ranks + const wk = try Group.fromShape(&.{ 32, 8, 128 }); + try testing.expectError(error.SplitNotOnBlocks, wk.part(.col, 0, 2)); + try testing.expectEqual(wk.trellisBytes(), (try wk.part(.whole, 1, 2)).trellis.bytes()); +} + +test "expert layouts: one width stacks, mixed widths go ragged with aligned offsets" { + const a = testing.allocator; + const two = try (try Group.fromShape(&.{ 144, 320, 32 })).part(.row, 0, 2); + const three = try (try Group.fromShape(&.{ 144, 320, 48 })).part(.row, 0, 2); + var stack = try Layout.of(a, &.{ two, two, two }); + defer stack.deinit(a); + try testing.expect(stack.kind == .stack); + try testing.expectEqual(@as(u32, 1), stack.widths()); + // a mixed-bit pack's layer 2 down projection on rank 0: 383 experts at 2 bits and one at 3 bits + var parts: [384]Part = @splat(two); + parts[200] = three; + var rag = try Layout.of(a, &parts); + defer rag.deinit(a); + try testing.expect(rag.kind == .ragged); + try testing.expectEqual(@as(u64, 566968320), rag.trellisBytes()); + try testing.expectEqual(@as(u32, 383), rag.histogram()[4]); + try testing.expectEqual(@as(u32, 1), rag.histogram()[6]); + try testing.expectEqual(@as(u64, 0), rag.byteOffset(201) % 16); + try testing.expectEqual(rag.byteOffset(200) + three.trellis.bytes(), rag.byteOffset(201)); + const odd = try (try Group.fromShape(&.{ 144, 160, 32 })).part(.row, 0, 2); + try testing.expectError(error.ExpertShapesDiffer, Layout.of(a, &.{ two, odd })); +} diff --git a/zig/src/families/deepseek_v41/fixtures/config.json b/zig/src/families/deepseek_v41/fixtures/config.json new file mode 100644 index 0000000000..eeec94126b --- /dev/null +++ b/zig/src/families/deepseek_v41/fixtures/config.json @@ -0,0 +1,184 @@ +{ + "architectures": [ + "DeepseekV41ForCausalLM" + ], + "model_type": "deepseek_v41", + "dtype": "bfloat16", + "transformers_version": "5.6.0", + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 2, + "image_token_id": 129264, + "quantization_config": { + "quant_method": "exl3", + "version": "1.4.2", + "avg_bits": 2.9, + "head_bits": 6, + "calibration": { + "rows": 250, + "cols": 2048 + }, + "out_scales": "always", + "codebook": "mul1", + "original_quantization_config": { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [ + 32, + 32 + ], + "scale_fmt": "ue8m0", + "expert_dtype": "fp4" + }, + "mtp_bits": 4 + }, + "text_config": { + "model_type": "deepseek_v41_text", + "vocab_size": 129280, + "hidden_size": 5120, + "moe_intermediate_size": 2304, + "num_hidden_layers": 40, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "qk_rope_head_dim": 64, + "q_lora_rank": 1280, + "o_lora_rank": 1024, + "o_groups": 8, + "hidden_act": "silu", + "swiglu_limit": 10.0, + "rms_norm_eps": 1e-20, + "attention_bias": false, + "attention_dropout": 0.0, + "initializer_range": 0.02, + "use_cache": true, + "tie_word_embeddings": false, + "max_position_embeddings": 1048576, + "rope_theta": 10000, + "rope_scaling": { + "rope_type": "yarn", + "factor": 16, + "beta_fast": 32, + "beta_slow": 1, + "original_max_position_embeddings": 65536 + }, + "n_routed_experts": 384, + "n_shared_experts": 1, + "num_experts_per_tok": 6, + "scoring_func": "sqrtsoftplus", + "topk_method": "noaux_tc", + "norm_topk_prob": true, + "routed_scaling_factor": 1.5, + "sliding_window": 128, + "compress_ratios": [ + 0, + 0, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0 + ], + "compress_rope_theta": 160000, + "kv_source_layer_ids": [ + 2, + 8, + 14, + 20 + ], + "index_source_layer_ids": [ + 2, + 8, + 14, + 20, + 24, + 28, + 32, + 36 + ], + "index_n_heads": 32, + "index_head_dim": 128, + "index_topk": 512, + "candidate_source_layer_id": 20, + "candidate_topk_blocks": 2048, + "candidate_block_size": 8, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "hc_eps": 1e-06, + "engram_layer_ids": [ + 1, + 14 + ], + "engram_num_embeddings": [ + 384006168, + 384016682 + ], + "engram_max_ngram_size": 4, + "engram_vocab_size": 16000000, + "engram_n_heads": 8, + "engram_head_dim": 256, + "engram_pad_token_id": 2, + "engram_compressed_vocab_size": 99092, + "num_nextn_predict_layers": 3, + "dspark_block_size": 5, + "dspark_noise_token_id": 128799, + "dspark_target_layer_ids": [ + 37, + 38, + 39 + ], + "dspark_markov_rank": 256, + "dspark_n_routed_experts": 128, + "dspark_num_experts_per_tok": 3 + }, + "vision_config": { + "model_type": "deepseek_v41_vision", + "num_hidden_layers": 32, + "hidden_size": 1024, + "num_attention_heads": 16, + "intermediate_size": 2816, + "patch_size": 14, + "rope_theta": 10000, + "downsample_ratio": 3, + "max_image_tokens": 1024, + "min_pixels": 295936, + "max_wh_ratio": null + }, + "engram_table_dir": null +} \ No newline at end of file diff --git a/zig/src/families/deepseek_v41/fixtures/engram-hash-ref.bin b/zig/src/families/deepseek_v41/fixtures/engram-hash-ref.bin new file mode 100644 index 0000000000..2e1b20b38c Binary files /dev/null and b/zig/src/families/deepseek_v41/fixtures/engram-hash-ref.bin differ diff --git a/zig/src/families/deepseek_v41/fixtures/engram-rows-ref.txt b/zig/src/families/deepseek_v41/fixtures/engram-rows-ref.txt new file mode 100644 index 0000000000..d9ff54b326 --- /dev/null +++ b/zig/src/families/deepseek_v41/fixtures/engram-rows-ref.txt @@ -0,0 +1,5 @@ +lut ac4091794caec6f553a226c3857da2756d0b29c4c7b24e02ce61274a3df93185 +shard 1000 3000 9000 +rows 1000,1001,1001,3999,1007,2500,1007 82b5fc83bc4eb5195c2d36925f57ef8af91663ed60af344641924535b8dc25b4 +rows 1000,1977,2954,3931,1908,2885,3862,1839,2816,3793,1770,2747,3724,1701,2678,3655,1632,2609,3586,1563,2540,3517,1494,2471,3448,1425,2402,3379,1356,2333,3310,1287,2264,3241,1218,2195,3172,1149,2126,3103,1080,2057,3034,1011,1988,2965,3942,1919,2896,3873,1850,2827,3804,1781,2758,3735,1712,2689,3666,1643,2620,3597,1574,2551,3528,1505,2482,3459,1436,2413,3390,1367,2344,3321,1298,2275,3252,1229,2206,3183,1160,2137,3114,1091,2068,3045,1022,1999,2976,3953,1930,2907,3884,1861,2838,3815 6a0439b98aad93e03a744b515286f8041b7567bd09e15a2f1e358e00b4950199 +rows 1259,1258,1257,1256,1255,1254,1253,1252,1251,1250,1249,1248,1247,1246,1245,1244,1243,1242,1241,1240,1239,1238,1237,1236,1235,1234,1233,1232,1231,1230,1229,1228,1227,1226,1225,1224,1223,1222,1221,1220,1219,1218,1217,1216,1215,1214,1213,1212,1211,1210,1209,1208,1207,1206,1205,1204,1203,1202,1201,1200 f8ef91955c85ec1ba1caa46f4116ddc94cea8204c72f97670292cbb9d8ce82b5 diff --git a/zig/src/families/deepseek_v41/fixtures/lanes-ref.bin b/zig/src/families/deepseek_v41/fixtures/lanes-ref.bin new file mode 100644 index 0000000000..231acd0531 Binary files /dev/null and b/zig/src/families/deepseek_v41/fixtures/lanes-ref.bin differ diff --git a/zig/src/families/deepseek_v41/fixtures/widths-exl3-2p8.txt b/zig/src/families/deepseek_v41/fixtures/widths-exl3-2p8.txt new file mode 100644 index 0000000000..317ae6ed73 --- /dev/null +++ b/zig/src/families/deepseek_v41/fixtures/widths-exl3-2p8.txt @@ -0,0 +1,44 @@ +# expert widths of a 2.8 bpw EXL3 pack, world 2 rank 0: layer w1 w2 w3 shared_k2(w1,w2,w3) rank_bytes +0 4:384 4:384 4:384 10,10,10 1709752320 +1 4:384 4:384 4:384 10,10,10 1709752320 +2 4:384 4:383,6:1 4:384 10,10,10 1710489600 +3 4:384 4:384 4:384 10,10,10 1709752320 +4 4:384 4:383,6:1 4:384 10,10,10 1710489600 +5 4:384 4:384 4:384 10,10,10 1709752320 +6 4:384 4:384 4:384 8,8,8 1707540480 +7 4:384 4:382,6:2 4:384 10,10,10 1711226880 +8 4:384 4:376,6:8 4:384 12,12,12 1717862400 +9 4:382,6:2 4:369,6:14,8:1 4:382,6:2 12,12,12 1726709760 +10 4:376,6:8 4:360,6:23,8:1 4:376,6:8 12,12,12 1742192640 +11 4:382,6:2 4:364,6:20 4:382,6:2 12,12,12 1729658880 +12 4:382,6:2 4:373,6:10,8:1 4:382,6:2 12,12,12 1723760640 +13 4:384 4:380,6:4 4:384 12,12,12 1714913280 +14 4:384 4:384 4:384 10,10,10 1709752320 +15 4:384 4:384 4:384 10,10,10 1709752320 +16 4:384 4:384 4:384 10,10,10 1709752320 +17 4:384 4:384 4:384 8,8,8 1707540480 +18 4:384 4:384 4:384 6,6,6 1705328640 +19 4:71,6:222,8:90,10:1 4:34,6:126,8:217,10:7 4:71,6:222,8:90,10:1 12,12,12 2737520640 +20 4:41,6:162,8:166,10:15 4:18,6:69,8:221,10:76 4:41,6:162,8:166,10:15 12,12,12 3051601920 +21 4:18,6:104,8:216,10:46 4:5,6:38,8:182,10:159 4:18,6:104,8:216,10:46 12,12,12 3353886720 +22 4:13,6:69,8:226,10:76 4:3,6:25,8:145,10:211 4:13,6:69,8:226,10:76 12,12,12 3515351040 +23 4:24,6:55,8:220,10:85 4:21,6:31,8:139,10:193 4:24,6:55,8:220,10:85 12,12,12 3472588800 +24 4:14,6:117,8:221,10:32 4:6,6:34,8:213,10:131 4:14,6:117,8:221,10:32 12,12,12 3306700800 +25 4:28,6:175,8:167,10:14 4:5,6:66,8:237,10:76 4:28,6:175,8:167,10:14 12,12,12 3090677760 +26 4:31,6:239,8:108,10:6 4:8,6:116,8:228,10:32 4:31,6:239,8:108,10:6 12,12,12 2901934080 +27 4:54,6:199,8:126,10:5 4:20,6:113,8:219,10:32 4:54,6:199,8:126,10:5 12,12,12 2876129280 +28 4:109,6:213,8:59,10:3 4:39,6:178,8:146,10:21 4:109,6:213,8:59,10:3 12,12,12 2606284800 +29 4:116,6:219,8:47,10:2 4:47,6:189,8:136,10:12 4:116,6:219,8:47,10:2 12,12,12 2548776960 +30 4:164,6:192,8:28 4:76,6:217,8:88,10:3 4:164,6:192,8:28 12,12,12 2374041600 +31 4:180,6:177,8:26,10:1 4:84,6:213,8:84,10:3 4:180,6:177,8:26,10:1 12,12,12 2341601280 +32 4:188,6:172,8:24 4:84,6:205,8:92,10:3 4:188,6:172,8:24 12,12,12 2329804800 +33 4:187,6:175,8:22 4:92,6:217,8:72,10:3 4:187,6:175,8:22 12,12,12 2307686400 +34 4:208,6:157,8:19 4:122,6:207,8:53,10:2 4:208,6:157,8:19 12,12,12 2234695680 +35 4:65,6:186,8:121,10:12 4:36,6:97,8:207,10:44 4:65,6:186,8:121,10:12 12,12,12 2870231040 +36 6:147,8:205,10:32 6:73,8:207,10:104 6:147,8:205,10:32 12,12,12 3263938560 +37 6:91,8:222,10:71 6:47,8:158,10:179 6:91,8:222,10:71 12,12,12 3478487040 +38 6:50,8:129,10:205 6:32,8:61,10:291 6:50,8:129,10:205 12,12,12 3830169600 +39 6:48,8:103,10:233 6:38,8:42,10:304 6:48,8:103,10:233 12,12,12 3879567360 +40 8:128 8:128 8:128 8,8,8 1141309440 +41 8:128 8:128 8:128 8,8,8 1141309440 +42 8:128 8:128 8:128 8,8,8 1141309440 diff --git a/zig/src/families/deepseek_v41/forms.zig b/zig/src/families/deepseek_v41/forms.zig new file mode 100644 index 0000000000..ad5250d048 --- /dev/null +++ b/zig/src/families/deepseek_v41/forms.zig @@ -0,0 +1,93 @@ +//! The run-time weight forms built at load from the loaded weights: strips, lanes, bf16 mHC, expert tables. + +const std = @import("std"); +const cuda = @import("cuda"); +const load = @import("load.zig"); +const plan = @import("plan.zig"); +const prepare = @import("prepare.zig"); +const named = @import("named.zig"); + +pub const Source = struct { + d: *const cuda.Driver, + weights: *const load.Weights, + plan: *const plan.Plan, + + fn download(s: *const Source, a: std.mem.Allocator, name: []const u8) ![]u8 { + const t = s.weights.get(name) orelse return error.MissingWeight; + const host = try a.alloc(u8, t.len); + try cuda.DeviceBuffer.download(.{ .d = s.d, .ptr = t.ptr, .len = t.len }, 0, host); + return host; + } + + /// (K tiles, N tiles, int16 words a tile) of a group's trellis. + fn trellis(s: *const Source, w: []const u8) ![3]usize { + var b: [256]u8 = undefined; + const t = s.weights.get(try std.fmt.bufPrint(&b, "{s}.trellis", .{w})) orelse return error.MissingWeight; + if (t.rank != 3) return error.NotTrellis; + return .{ t.shape[0], t.shape[1], t.shape[2] }; + } +}; + +/// A persistent role's initial bytes when it is a weight form (null: state, scratch, tables the forward fills). +pub fn build(a: std.mem.Allocator, s: *const Source, role: []const u8) !?[]const u8 { + if (!std.mem.startsWith(u8, role, "s.")) return null; + const lanes = std.mem.endsWith(u8, role, ".lanes"); + if (lanes or std.mem.endsWith(u8, role, ".T")) { + const w = role[2 .. role.len - (if (lanes) ".lanes".len else ".T".len)]; + const tr = try s.download(a, try std.fmt.allocPrint(a, "{s}.trellis", .{w})); + const g = try s.trellis(w); + const st = try prepare.strips(a, tr, g[0], g[1], g[2]); + return if (lanes) try prepare.lanes(a, st, g[0], g[2] / 8) else st; + } + if (std.mem.endsWith(u8, role, ".fn16")) { + const src = try s.download(a, try std.fmt.allocPrint(a, "{s}.fn", .{role[2 .. role.len - 5]})); + const out = try a.alloc(u8, src.len / 2); + for (0..src.len / 4) |i| std.mem.writeInt(u16, out[2 * i ..][0..2], named.bf16Round(@bitCast(std.mem.readInt(u32, src[4 * i ..][0..4], .little))), .little); + return out; + } + if (!std.mem.startsWith(u8, role, "s.L")) return null; + const ex = std.mem.indexOf(u8, role, ".ex.") orelse return null; + const layer = role[2..ex]; + const what = role[ex + 4 ..]; + const proj: usize = switch (what[what.len - 1]) { + 'g' => 0, + 'u' => 2, + 'd' => 1, + else => return null, + }; + const projs = [_][]const u8{ "w1", "w2", "w3" }; + const w = projs[proj]; + if (std.mem.startsWith(u8, what, "suh") or std.mem.startsWith(u8, what, "svh")) { + const kind = what[0..3]; + const routed = try s.download(a, try std.fmt.allocPrint(a, "{s}.moe.{s}.{s}", .{ layer, w, kind })); + const shared = try s.download(a, try std.fmt.allocPrint(a, "{s}.moe.shared.0.{s}.{s}", .{ layer, w, kind })); + return try std.mem.concat(a, u8, &.{ routed, shared }); + } + const L = try std.fmt.parseInt(u32, layer[1..], 10); + const lp = for (s.plan.layers) |*l| { + if (l.index == L) break l; + } else return error.LayerNotPlanned; + const lay = &lp.experts[proj].layout; + const E = lay.words.len; + const shared = try std.fmt.allocPrint(a, "{s}.moe.shared.0.{s}", .{ layer, w }); + if (std.mem.startsWith(u8, what, "k2")) { + const out = try a.alloc(u8, 4 * (E + 1)); + for (lay.words, 0..) |words, i| std.mem.writeInt(i32, out[4 * i ..][0..4], @intCast(words / 8), .little); + std.mem.writeInt(i32, out[4 * E ..][0..4], @intCast((try s.trellis(shared))[2] / 8), .little); + return out; + } + if (std.mem.startsWith(u8, what, "tp")) { + // trellis address = ragged base + offset, shared expert last; a uniform stack (".trellis"): same offsets + const data = s.weights.get(try std.fmt.allocPrint(a, "{s}.moe.{s}.data", .{ layer, w })) orelse uniform: { + const t = s.weights.get(try std.fmt.allocPrint(a, "{s}.moe.{s}.trellis", .{ layer, w })) orelse return error.MissingWeight; + if (t.len != lay.trellisBytes()) return error.ExpertStackSize; + break :uniform t; + }; + const sh = s.weights.get(try std.fmt.allocPrint(a, "{s}.trellis", .{shared})) orelse return error.MissingWeight; + const out = try a.alloc(u8, 8 * (E + 1)); + for (0..E) |e| std.mem.writeInt(u64, out[8 * e ..][0..8], data.ptr + lay.byteOffset(e), .little); + std.mem.writeInt(u64, out[8 * E ..][0..8], sh.ptr, .little); + return out; + } + return null; +} diff --git a/zig/src/families/deepseek_v41/forms_inplace.zig b/zig/src/families/deepseek_v41/forms_inplace.zig new file mode 100644 index 0000000000..22b4454bbc --- /dev/null +++ b/zig/src/families/deepseek_v41/forms_inplace.zig @@ -0,0 +1,127 @@ +//! TF_DSV41_FORMS_INPLACE=1: the dense forms (s..lanes / .T) written over their trellis, as dense3.to_lanes does. + +const std = @import("std"); +const buffers = @import("buffers.zig"); +const prepare = @import("prepare.zig"); + +const log = std.log.scoped(.dsv41); + +pub fn enabled() bool { + const v = std.c.getenv("TF_DSV41_FORMS_INPLACE") orelse return false; + return std.mem.eql(u8, std.mem.trim(u8, std.mem.span(v), " "), "1"); +} + +pub const Bound = struct { roles: u32 = 0, bytes: u64 = 0 }; + +/// The source weight of a dense form role (`` of "s..lanes" / "s..T"), null for any other role. +pub fn sourceOf(role: []const u8) ?[]const u8 { + if (!std.mem.startsWith(u8, role, "s.")) return null; + for ([_][]const u8{ ".lanes", ".T" }) |suf| if (std.mem.endsWith(u8, role, suf)) return role[2 .. role.len - suf.len]; + return null; +} + +/// Binds `plans[0]`'s forms (r: run.Runner, weights: load.Weights) unless read, a source's two, or too large. +pub fn bind(gpa: std.mem.Allocator, r: anytype, weights: anytype, plans: []const *const buffers.Plan) !Bound { + // sources with more than one form role: each would overwrite the trellis the other is built from + var forms: std.StringHashMapUnmanaged(u32) = .empty; + defer forms.deinit(gpa); + for (plans[0].sizes.keys()) |role| if (sourceOf(role)) |w| { + const g = try forms.getOrPut(gpa, w); + g.value_ptr.* = if (g.found_existing) g.value_ptr.* + 1 else 1; + }; + var out: Bound = .{}; + var name: [192]u8 = undefined; + for (plans[0].sizes.keys(), plans[0].sizes.values()) |role, size| { + const w = sourceOf(role) orelse continue; + if (forms.get(w).? != 1 or r.addressOf(role) != null) continue; + const tr = std.fmt.bufPrint(&name, "{s}.trellis", .{w}) catch continue; + const t = weights.get(tr) orelse continue; + if (size > t.len) continue; + const read = for (plans) |p| { + if (p.weights.contains(tr)) break true; + } else false; + if (read) continue; + try r.external(role, t.ptr); + out.roles += 1; + out.bytes += t.len; + } + log.info("forms in place (TF_DSV41_FORMS_INPLACE): {d} dense forms over their trellis, {d:.2} GiB not allocated", .{ out.roles, @as(f64, @floatFromInt(out.bytes)) / (1 << 30) }); + return out; +} + +const testing = std.testing; + +/// The weights' and runner's parts `bind` uses (load.Weights.get, run.Runner.addressOf / external). +const Span = struct { ptr: u64, len: usize }; +const Weights = struct { + map: std.StringHashMapUnmanaged(Span) = .empty, + fn get(w: *const Weights, name: []const u8) ?Span { + return w.map.get(name); + } +}; +const Roles = struct { + gpa: std.mem.Allocator, + addrs: std.StringHashMapUnmanaged(u64) = .empty, + fn addressOf(r: *const Roles, role: []const u8) ?u64 { + return r.addrs.get(role); + } + fn external(r: *Roles, role: []const u8, addr: u64) !void { + try r.addrs.put(r.gpa, role, addr); + } +}; + +test "forms in place: each dense form bound to its own trellis; read, shared, larger or other roles keep their own" { + const gpa = testing.allocator; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var w: Weights = .{}; + defer w.map.deinit(gpa); + const ws = [_]struct { []const u8, u64, usize }{ + .{ "L0.attn.wq_b.trellis", 0x10000, 4096 }, + .{ "head.trellis", 0x20000, 1024 }, + .{ "L1.attn.wo_b.trellis", 0x30000, 512 }, + .{ "L2.attn.wq_a.trellis", 0x40000, 256 }, + .{ "L3.attn.wkv.trellis", 0x50000, 256 }, + }; + for (ws) |x| try w.map.put(gpa, x[0], .{ .ptr = x[1], .len = x[2] }); + var p: buffers.Plan = .{ .a = a }; + for ([_]struct { []const u8, u64 }{ + .{ "s.L0.attn.wq_b.lanes", 4096 }, // bound + .{ "s.head.T", 1000 }, // bound (smaller than its trellis) + .{ "s.L1.attn.wo_b.lanes", 1024 }, // larger than its trellis: own buffer + .{ "s.L2.attn.wq_a.lanes", 256 }, // its trellis is read by a 4K program: own buffer + .{ "s.L3.attn.wkv.lanes", 256 }, // two forms of one source: both their own + .{ "s.L3.attn.wkv.T", 256 }, + .{ "s.L4.attn.wq_b.lanes", 64 }, // no trellis on this rank: own buffer + .{ "s.L0.hc_attn.fn16", 64 }, // not a dense form + .{ "w.x", 64 }, + }) |x| try p.sizes.put(a, x[0], x[1]); + var k4: buffers.Plan = .{ .a = a }; + try k4.weights.put(a, "L2.attn.wq_a.trellis", {}); + var r: Roles = .{ .gpa = gpa }; + defer r.addrs.deinit(gpa); + const got = try bind(gpa, &r, &w, &.{ &p, &k4 }); + try testing.expectEqual(@as(u32, 2), got.roles); + try testing.expectEqual(@as(u64, 4096 + 1024), got.bytes); + try testing.expectEqual(@as(?u64, 0x10000), r.addressOf("s.L0.attn.wq_b.lanes")); + try testing.expectEqual(@as(?u64, 0x20000), r.addressOf("s.head.T")); + for ([_][]const u8{ "s.L1.attn.wo_b.lanes", "s.L2.attn.wq_a.lanes", "s.L3.attn.wkv.lanes", "s.L3.attn.wkv.T", "s.L4.attn.wq_b.lanes", "s.L0.hc_attn.fn16", "w.x" }) |role| + try testing.expectEqual(@as(?u64, null), r.addressOf(role)); + try testing.expectEqualStrings("L0.attn.wq_b", sourceOf("s.L0.attn.wq_b.lanes").?); + try testing.expectEqual(@as(?[]const u8, null), sourceOf("L.x.lanes")); +} + +test "forms in place: strips and lanes are the trellis's own length, so a form fits over it" { + const gpa = testing.allocator; + // kt 2 K tiles, nt 8 N tiles, 32 int16 words a tile (K2 4): one strip, one lanes group pair + const src = try gpa.alloc(u8, 2 * 8 * 32 * 2); + defer gpa.free(src); + for (src, 0..) |*b, i| b.* = @truncate(i *% 37 +% 11); + const st = try prepare.strips(gpa, src, 2, 8, 32); + defer gpa.free(st); + const ln = try prepare.lanes(gpa, st, 2, 4); + defer gpa.free(ln); + try testing.expectEqual(src.len, st.len); + try testing.expectEqual(src.len, ln.len); +} diff --git a/zig/src/families/deepseek_v41/forward_types.zig b/zig/src/families/deepseek_v41/forward_types.zig new file mode 100644 index 0000000000..d77adf3c09 --- /dev/null +++ b/zig/src/families/deepseek_v41/forward_types.zig @@ -0,0 +1,71 @@ +//! Plain types the forward shares with its later layers (CED stash, multi-segment runs, gate items, round graph). +const std = @import("std"); +const eh = @import("engram_host.zig"); + +/// the decoder's rows over a prompt's end, its last token (the verify row) included +pub const replay_rows: u32 = 128; +/// the stash rows a slot needs (the prefilled part of the tail) +pub const keep: u32 = replay_rows - 1; +/// the stash ring's rows (position p at row p % ring): at least `keep`, a power of two +pub const ring: u32 = 128; + +/// A slot's stash bookkeeping (rows on the device, "s.ced.*" row p % ring): positions [max(from, hi - keep), hi), ids. +pub const Stash = struct { + from: u64 = 0, + hi: u64 = 0, + ids: [ring]u32 = @splat(0), + + /// A segment at `start`; a gap starts the stash over there. + pub fn extend(s: *Stash, start: u64, ids: []const u32) void { + if (s.hi != start) s.from = start; + for (ids, 0..) |t, i| s.ids[@intCast((start + i) % ring)] = t; + s.hi = start + ids.len; + } + + /// The first position held. + pub fn lo(s: *const Stash) u64 { + return @max(s.from, s.hi -| keep); + } + + /// The replay check: the stash ends at `end` and holds [r0, end). + pub fn covers(s: *const Stash, r0: u64, end: u64) bool { + return s.hi == end and s.lo() <= r0 and end - r0 <= keep; + } + + /// The ids of positions [r0, end) (`covers` first). + pub fn tail(s: *const Stash, r0: u64, end: u64, out: []u32) void { + for (out[0..@intCast(end - r0)], 0..) |*o, i| o.* = s.ids[@intCast((r0 + i) % ring)]; + } +}; + +/// One segment of a multi-segment run: a slot's prompt ids from its position `start`. +pub const MultiSeg = struct { slot: u32, start: u64, ids: []const u32 }; + +/// An Engram tail (a slot's last ids before the run). +pub const Tail = struct { + ids: [eh.max_ngram]u32 = undefined, + len: usize = 0, + + pub fn items(t: *const Tail) []const u32 { + return t.ids[0..t.len]; + } +}; + +/// The run in flight: its segments, each one's first row and Engram tail at the start, every segment's ids in order. +pub const Multi = struct { segs: []const MultiSeg, row0: []const usize, tails: []const Tail, ids: []const u32 }; + +/// One slot's decoder replay in a batched run: the tail's first position and rows, its first row in the run. +pub const Replayed = struct { slot: u32, at: u64, rows: u32, row0: u32 }; + +/// A window's part on one slot: its ids and the slot's Engram tail (lookback) before them. +pub const Item = struct { ids: []const u32, tail: []const u32 }; + +/// TF_DSV41_ROUND_GRAPH (round_graph.zig): the settings read once. +pub const RoundGraph = struct { + split: bool = false, + gate: bool = false, + /// layers in the head graph + head: u32 = 1, +}; + +pub const stream_min: i64 = 4096; // TF_DSV41_INDEX_STREAM_MIN (block.zig: a decode window past it would stream) diff --git a/zig/src/families/deepseek_v41/host_stage.zig b/zig/src/families/deepseek_v41/host_stage.zig new file mode 100644 index 0000000000..baf0bdac81 --- /dev/null +++ b/zig/src/families/deepseek_v41/host_stage.zig @@ -0,0 +1,41 @@ +//! Pinned host staging a prompt segment rewrites without draining the stream (waits only on its last copy's event). + +const std = @import("std"); +const cuda = @import("cuda"); + +/// forward.CHUNK: windows of more rows stage through pinned memory without a stream sync +pub const chunk: usize = 128; + +pub const Stage = struct { + buf: ?cuda.HostBuffer = null, + ev: ?cuda.Event = null, + /// a copy out of `buf` was issued and `ev` recorded after it + pending: bool = false, + + /// `bytes` of the buffer to write: the copy that last read it has finished; grown when too small. + pub fn take(s: *Stage, d: *const cuda.Driver, bytes: usize) ![]u8 { + if (s.pending) { + try s.ev.?.synchronize(); + s.pending = false; + } + if (s.buf == null or s.buf.?.bytes.len < bytes) { + if (s.buf) |*b| b.free(); + s.buf = null; + s.buf = try cuda.HostBuffer.alloc(d, @max(bytes, 4096)); + } + return s.buf.?.bytes[0..bytes]; + } + + /// The copy out of the buffer is on `stream`: the next `take` waits for it alone. + pub fn issued(s: *Stage, d: *const cuda.Driver, stream: cuda.Stream) !void { + if (s.ev == null) s.ev = try cuda.Event.init(d, false); + try s.ev.?.record(stream); + s.pending = true; + } + + pub fn deinit(s: *Stage) void { + if (s.ev) |*e| e.deinit(); + if (s.buf) |*b| b.free(); + s.* = .{}; + } +}; diff --git a/zig/src/families/deepseek_v41/host_test.zig b/zig/src/families/deepseek_v41/host_test.zig new file mode 100644 index 0000000000..1212ed3518 --- /dev/null +++ b/zig/src/families/deepseek_v41/host_test.zig @@ -0,0 +1,330 @@ +//! Host tests of the loader side: release config, synthetic pack, 2.8 bpw widths, header dumps (TF_DSV41_HEADERS). + +const std = @import("std"); +const testing = std.testing; +const config = @import("config.zig"); +const Config = config.Config; +const exl3 = @import("exl3.zig"); +const Pack = @import("pack.zig").Pack; +const plan = @import("plan.zig"); +const named = @import("named.zig"); + +test "the release config.json: topology, routing and rope" { + const c = try Config.parse(testing.allocator, @embedFile("fixtures/config.json")); + try testing.expectEqual(@as(u32, 129280), c.vocab); + try testing.expectEqual(@as(u32, 43), c.allLayers()); + try testing.expectEqual(@as(f32, 1e-20), c.eps); + try testing.expect(c.rope.yarn and c.rope.factor == 16 and c.rope.original_max_positions == 65536); + try testing.expectEqual(@as(f64, 160000), c.rope.compress_theta); + try testing.expectEqual(config.Mode.swa, c.mode(1)); + try testing.expectEqual(config.Mode.full, c.mode(2)); + try testing.expectEqual(config.Mode.reuse, c.mode(3)); + try testing.expectEqual(config.Mode.full, c.mode(20)); + try testing.expectEqual(config.Mode.reindex, c.mode(24)); + try testing.expectEqual(config.Mode.swa, c.mode(41)); + try testing.expectEqual(@as(?u32, 14), c.kvSource(19)); + try testing.expectEqual(@as(?u32, 20), c.kvSource(39)); + try testing.expectEqual(@as(?u32, 36), c.indexSource(39)); + try testing.expect(c.isCandidateSource(20) and c.usesCandidates(24) and !c.usesCandidates(20)); + try testing.expectEqual(@as(u32, 128), c.expertsOf(40).count); + try testing.expectEqual(@as(u32, 3), c.expertsOf(40).topk); + try testing.expectEqual(@as(u32, 24), c.hashCols()); + try testing.expect(c.isEngram(14) and !c.isEngram(2)); + try testing.expectEqual(@as(u32, 1), c.eos); + // the release's values equal the Config defaults + const d: Config = .{}; + try testing.expectEqualSlices(u32, d.compress_ratios.items(), c.compress_ratios.items()); + try testing.expectEqualSlices(u64, d.engram_rows.items(), c.engram_rows.items()); + try testing.expectError(error.UnsupportedScoring, Config.parse(testing.allocator, "{\"scoring_func\": \"sigmoid\"}")); + try testing.expectError(error.BadConfig, Config.parse(testing.allocator, "{\"compress_ratios\": [0, 3]}")); + try config.tiny().validate(); +} + +/// Header JSON of a pack for `c`, named as the release packs are (their index, pattern by pattern), header only. +fn tinyPack(a: std.mem.Allocator, c: *const Config, extra: []const u8, skip: []const u8) ![]u8 { + var n: u64 = 0; + return tinyPackLen(a, c, extra, skip, &n); +} + +/// `tinyPack`, and the data bytes its header describes. +fn tinyPackLen(a: std.mem.Allocator, c: *const Config, extra: []const u8, skip: []const u8, data_len: *u64) ![]u8 { + var w: std.Io.Writer.Allocating = .init(a); + errdefer w.deinit(); + var at: u64 = 0; + var first = true; + const T = struct { + fn put(wr: *std.Io.Writer, pos: *u64, f: *bool, sk: []const u8, name: []const u8, dtype: []const u8, shape: []const u64) !void { + if (sk.len > 0 and std.mem.eql(u8, name, sk)) return; + var n: u64 = if (std.mem.eql(u8, dtype, "F32") or std.mem.eql(u8, dtype, "I32")) 4 else 2; + for (shape) |d| n *= d; + try wr.print("{s}\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ if (f.*) "" else ",", name, dtype }); + for (shape, 0..) |d, i| try wr.print("{s}{d}", .{ if (i == 0) "" else ",", d }); + try wr.print("],\"data_offsets\":[{d},{d}]}}", .{ pos.*, pos.* + n }); + pos.* += n; + f.* = false; + } + fn group(wr: *std.Io.Writer, pos: *u64, f: *bool, sk: []const u8, p: []const u8, k: u64, n: u64, words: u64) !void { + var b: [160]u8 = undefined; + try put(wr, pos, f, sk, try std.fmt.bufPrint(&b, "{s}.trellis", .{p}), "I16", &.{ k / 16, n / 16, words }); + try put(wr, pos, f, sk, try std.fmt.bufPrint(&b, "{s}.suh", .{p}), "F16", &.{k}); + try put(wr, pos, f, sk, try std.fmt.bufPrint(&b, "{s}.svh", .{p}), "F16", &.{n}); + try put(wr, pos, f, sk, try std.fmt.bufPrint(&b, "{s}.mul1", .{p}), "I32", &.{}); + } + }; + const wr = &w.writer; + try wr.writeAll("{"); + const D: u64 = c.hidden; + var nb: [160]u8 = undefined; + var l: u32 = 0; + while (l < c.allLayers()) : (l += 1) { + var pb: [32]u8 = undefined; + const p = if (l < c.layers) try std.fmt.bufPrint(&pb, "layers.{d}", .{l}) else try std.fmt.bufPrint(&pb, "mtp.{d}", .{l - c.layers}); + const E: u64 = if (l < c.layers) c.experts else c.dspark_experts; + inline for (.{ "attn_norm.weight", "ffn_norm.weight" }) |s| try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}." ++ s, .{p}), "BF16", &.{D}); + inline for (.{ "attn", "ffn" }) |s| { + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.hc_" ++ s ++ "_fn", .{p}), "F32", &.{ 24, 4 * D }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.hc_" ++ s ++ "_base", .{p}), "F32", &.{24}); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.hc_" ++ s ++ "_scale", .{p}), "F32", &.{3}); + } + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.q_norm.weight", .{p}), "BF16", &.{c.q_lora}); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.kv_norm.weight", .{p}), "BF16", &.{c.head_dim}); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.attn_sink", .{p}), "F32", &.{c.heads}); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wq_a", .{p}), D, c.q_lora, 64); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wkv", .{p}), D, c.head_dim, 80); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wq_b", .{p}), c.q_lora, c.heads * c.head_dim, 64); + var g: u32 = 0; + while (g < c.o_groups) : (g += 1) try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wo_a.slice.{d}", .{ p, g }), c.heads / c.o_groups * c.head_dim, c.o_lora, 64); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wo_b", .{p}), c.o_groups * c.o_lora, D, 64); + if (l < c.layers and c.kv_sources.contains(l)) { + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.wkv", .{p}), D, c.head_dim, 80); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.wkv.weight", .{p}), "BF16", &.{ c.head_dim, D }); + if (c.compressRatio(l) == 2) { + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.wgate", .{p}), D, c.head_dim, 80); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.wgate.weight", .{p}), "BF16", &.{ c.head_dim, D }); + } + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.norm.weight", .{p}), "BF16", &.{c.head_dim}); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.wk", .{p}), c.head_dim, c.index_dim, 128); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.wk.weight", .{p}), "BF16", &.{ c.index_dim, c.head_dim }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.k_norm.weight", .{p}), "BF16", &.{c.index_dim}); + } + if (l < c.layers and c.index_sources.contains(l)) { + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.wq_b", .{p}), c.q_lora, c.index_heads * c.index_dim, 80); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.weights_proj.weight", .{p}), "BF16", &.{ c.index_heads, D }); + } + if (c.engram_layers.contains(l)) { + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.engram.wkv", .{p}), c.hashCols() * c.engram_head_dim, (c.engram_max_ngram + 1) * D, 64); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.engram.q_weight", .{p}), "BF16", &.{ c.hc_mult, D }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.engram.k_weight", .{p}), "BF16", &.{ c.hc_mult, D }); + } + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.gate.weight", .{p}), "BF16", &.{ E, D }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.gate.bias", .{p}), "F32", &.{E}); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.gate.bias_vl", .{p}), "F32", &.{E}); + inline for (.{ "w1", "w2", "w3" }) |m| { + const kn: [2]u64 = if (comptime std.mem.eql(u8, m, "w2")) .{ c.expert_width, D } else .{ D, c.expert_width }; + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.shared_experts." ++ m, .{p}), kn[0], kn[1], 80); + var e: u64 = 0; + while (e < E) : (e += 1) { + // block 2: expert 1's down projection at 3 bits, expert 3's gate at 3 bits (a gate/up width mismatch) + const words: u64 = if (l == 2 and ((e == 1 and comptime std.mem.eql(u8, m, "w2")) or (e == 3 and comptime std.mem.eql(u8, m, "w1")))) 48 else 32; + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.experts.{d}." ++ m, .{ p, e }), kn[0], kn[1], words); + } + } + } + try T.put(wr, &at, &first, skip, "embed.weight", "BF16", &.{ c.vocab, D }); + try T.put(wr, &at, &first, skip, "norm.weight", "BF16", &.{D}); + try T.group(wr, &at, &first, skip, "head", D, c.vocab, 96); + const last = c.mtp_layers - 1; + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "mtp.{d}.main_proj", .{last}), 3 * D, D, 64); + inline for (.{ "main_norm", "norm" }) |s| try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "mtp.{d}." ++ s ++ ".weight", .{last}), "BF16", &.{D}); + inline for (.{ "embed", "head" }) |s| try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "mtp.{d}.markov_head." ++ s ++ ".weight", .{last}), "BF16", &.{ c.vocab, c.dspark_markov_rank }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "mtp.{d}.confidence_head.proj.weight", .{last}), "F32", &.{ 1, D + c.dspark_markov_rank }); + try T.put(wr, &at, &first, skip, "vision.norm.weight", "BF16", &.{D}); + try T.put(wr, &at, &first, skip, "image_start", "BF16", &.{D}); + if (extra.len > 0) try T.put(wr, &at, &first, skip, extra, "BF16", &.{4}); + try wr.writeAll("}"); + data_len.* = at; + return w.toOwnedSlice(); +} + +fn tinyPlan(a: std.mem.Allocator, c: *const Config, extra: []const u8, skip: []const u8, rank: u32) !plan.Plan { + return tinyPlanOf(a, c, c, extra, skip, rank); +} + +/// A pack written for `stored`, planned with `c`. +fn tinyPlanOf(a: std.mem.Allocator, c: *const Config, stored: *const Config, extra: []const u8, skip: []const u8, rank: u32) !plan.Plan { + const json = try tinyPack(a, stored, extra, skip); + defer a.free(json); + var p = Pack.init(a); + defer p.deinit(); + try p.addHeader("model.safetensors", json, std.math.maxInt(u32)); + return plan.build(a, c, &p, .{ .rank = rank, .world = 2 }); +} + +test "a synthetic pack: both ranks plan every block, the parts tile the full tensors, nothing is left" { + const a = testing.allocator; + const c = config.tiny(); + var r0 = try tinyPlan(a, &c, "", "", 0); + defer r0.deinit(); + var r1 = try tinyPlan(a, &c, "", "", 1); + defer r1.deinit(); + try testing.expectEqual(@as(usize, 0), r0.leftovers.len); + try testing.expectEqual(@as(usize, 10), r0.layers.len); + const l2 = &r0.layers[2]; + try testing.expectEqual(config.Mode.full, l2.mode); + try testing.expect(l2.experts[1].layout.kind == .ragged and l2.experts[0].layout.kind == .ragged and l2.experts[2].layout.kind == .stack); + try testing.expectEqual(@as(u32, 1), l2.gate_up_mismatch); + try testing.expectEqual(@as(u32, 2), l2.experts[1].layout.widths()); + try testing.expect(r0.layers[3].experts[1].layout.kind == .stack); + try testing.expectEqual(@as(usize, 4), r0.layers[8].experts[0].layout.count()); + // every split part on the two ranks adds up to the stored tensor + for (r0.layers, r1.layers) |x, y| for (x.groups, y.groups) |g0, g1| { + if (g0.split == .whole) { + try testing.expectEqual(g0.full.trellisBytes(), g0.part.trellis.bytes()); + } else try testing.expectEqual(g0.full.trellisBytes(), g0.part.trellis.bytes() + g1.part.trellis.bytes()); + }; + try testing.expectEqual(r0.head.?.full.trellisBytes(), r0.head.?.part.trellis.bytes() + r1.head.?.part.trellis.bytes()); + try testing.expectEqualStrings("mtp.1.main_proj", r0.dspark.?.main_proj.prefix); + try testing.expect(r0.dspark.?.confidence != null); + try testing.expectEqual(r0.bytes(), r1.bytes()); + // the embedding: half the vocabulary rows on each rank + try testing.expectEqual(@as(u64, 128 * 256 * 2), r0.top[0].bytes); +} + +test "a synthetic pack: unknown names are leftovers, missing ones and wrong shapes are refused" { + const a = testing.allocator; + const c = config.tiny(); + var p = try tinyPlan(a, &c, "layers.3.attn.mystery", "", 0); + defer p.deinit(); + try testing.expectEqual(@as(usize, 1), p.leftovers.len); + try testing.expectEqualStrings("layers.3.attn.mystery", p.leftovers[0]); + try testing.expectError(error.MissingTensor, tinyPlan(a, &c, "", "layers.5.ffn.experts.6.w3.svh", 0)); + try testing.expectError(error.MissingTensor, tinyPlan(a, &c, "", "layers.4.attn.indexer.k_norm.weight", 1)); + var wide = c; + wide.q_lora = 48; + try testing.expectError(error.UnexpectedShape, tinyPlanOf(a, &wide, &c, "", "", 0)); +} + +test "a synthetic pack on disk: every planned tensor named in the reference loader's form" { + const a = testing.allocator; + const io = testing.io; + const c = config.tiny(); + var n: u64 = 0; + const json = try tinyPackLen(a, &c, "", "", &n); + defer a.free(json); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + { + var f = try tmp.dir.createFile(io, "model.safetensors", .{}); + defer f.close(io); + var len: [8]u8 = undefined; + std.mem.writeInt(u64, &len, json.len, .little); + try f.writePositionalAll(io, &len, 0); + try f.writePositionalAll(io, json, 8); + try f.setLength(io, 8 + json.len + n); + } + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var p = try Pack.open(a, io, dir); + defer p.deinit(); + var pl = try plan.build(a, &c, &p, .{ .rank = 1, .world = 2 }); + defer pl.deinit(); + var b: named.Builder = .{ .gpa = a, .io = io, .pack = &p }; + defer b.deinit(); + try b.all(&pl, c.o_groups); + try testing.expectEqual(named.Kind.i16, b.find("L2.moe.w2.data").?.kind); + try testing.expect(b.find("L3.moe.w2.trellis") != null and b.find("L3.moe.w2.data") == null); + try testing.expectEqualSlices(usize, &.{ 8, 8, 16, 32 }, b.find("L3.moe.w2.trellis").?.shape[0..4]); + const qk = b.find("L1.engram.qk").?; + try testing.expect(qk.kind == .f32 and qk.bytes.len == 4 * c.hc_mult * c.hidden); + try testing.expectEqual(named.Kind.bf16, b.find("L0.moe.gate").?.kind); + try testing.expectEqual(named.Kind.f32, b.find("L2.attn.ix_wp").?.kind); + // rank 1 of 2: half the heads' sinks, half the vocabulary rows, wo_a's one local group + try testing.expectEqual(@as(usize, 4 * c.heads / 2), b.find("L4.attn.sink").?.bytes.len); + try testing.expectEqual(@as(usize, 2 * c.vocab / 2 * c.hidden), b.find("embed").?.bytes.len); + try testing.expect(b.find("L4.attn.wo_a.0.svh") != null and b.find("L4.attn.wo_a.1.svh") == null); + try testing.expect(b.find("head.trellis") != null and b.find("norm") != null); +} + +/// One fixture histogram ("4:383,6:1") as rank-0 parts of a TP=2 projection with full shape `kn`. +fn parts(a: std.mem.Allocator, text: []const u8, kn: [2]u32, split: exl3.Split) ![]exl3.Part { + var out: std.ArrayList(exl3.Part) = .empty; + errdefer out.deinit(a); + var it = std.mem.tokenizeScalar(u8, text, ','); + while (it.next()) |item| { + const colon = std.mem.indexOfScalar(u8, item, ':') orelse return error.BadFixture; + const k2 = try std.fmt.parseInt(u32, item[0..colon], 10); + const n = try std.fmt.parseInt(u32, item[colon + 1 ..], 10); + const g = try exl3.Group.fromShape(&.{ kn[0] / 16, kn[1] / 16, 8 * k2 }); + try out.appendNTimes(a, try g.part(split, 0, 2), n); + } + return out.toOwnedSlice(a); +} + +test "the 2.8 bpw EXL3 pack: rank 0's routed + shared trellis bytes of every block match the reference loader's listing" { + const a = testing.allocator; + const c: Config = .{}; + var lines = std.mem.tokenizeScalar(u8, @embedFile("fixtures/widths-exl3-2p8.txt"), '\n'); + var blocks: u32 = 0; + var ragged: u32 = 0; + while (lines.next()) |line| { + if (line[0] == '#') continue; + var f = std.mem.tokenizeScalar(u8, line, ' '); + const layer = try std.fmt.parseInt(u32, f.next().?, 10); + var total: u64 = 0; + const projs = [_]struct { kn: [2]u32, split: exl3.Split }{ .{ .kn = .{ c.hidden, c.expert_width }, .split = .col }, .{ .kn = .{ c.expert_width, c.hidden }, .split = .row }, .{ .kn = .{ c.hidden, c.expert_width }, .split = .col } }; + for (projs) |pj| { + const ps = try parts(a, f.next().?, pj.kn, pj.split); + defer a.free(ps); + try testing.expectEqual(@as(usize, c.expertsOf(layer).count), ps.len); + var lay = try exl3.Layout.of(a, ps); + defer lay.deinit(a); + ragged += @intFromBool(lay.kind == .ragged); + total += lay.trellisBytes(); + } + var sh = std.mem.tokenizeScalar(u8, f.next().?, ','); + for (projs) |pj| { + const g = try exl3.Group.fromShape(&.{ pj.kn[0] / 16, pj.kn[1] / 16, 8 * try std.fmt.parseInt(u32, sh.next().?, 10) }); + total += (try g.part(pj.split, 0, 2)).trellis.bytes(); + } + try testing.expectEqual(try std.fmt.parseInt(u64, f.next().?, 10), total); + blocks += 1; + } + try testing.expectEqual(@as(u32, 43), blocks); + try testing.expect(ragged > 40); +} + +test "real pack headers (TF_DSV41_HEADERS: model.safetensors.index.json and .hdr dumps): planned blocks" { + const dir = testing.environ.getPosix("TF_DSV41_HEADERS") orelse return error.SkipZigTest; + const a = testing.allocator; + const io = testing.io; + var p = Pack.init(a); + defer p.deinit(); + const cwd = std.Io.Dir.cwd(); + const index = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + defer a.free(index); + const text = try cwd.readFileAlloc(io, index, a, .limited(1 << 28)); + defer a.free(text); + try p.useIndex(text); + var d = try cwd.openDir(io, dir, .{ .iterate = true }); + defer d.close(io); + var it = d.iterate(); + var blocks: std.ArrayList(u32) = .empty; + defer blocks.deinit(a); + while (try it.next(io)) |e| { + if (!std.mem.endsWith(u8, e.name, ".hdr")) continue; + const bytes = try d.readFileAlloc(io, e.name, a, .limited(1 << 28)); + defer a.free(bytes); + const shard = e.name[0 .. e.name.len - 4]; + try p.addHeaderDump(shard, bytes); + if (std.mem.startsWith(u8, shard, "layer-")) try blocks.append(a, try std.fmt.parseInt(u32, shard[6..8], 10)); + } + const c: Config = .{}; + for ([_]u32{ 0, 1 }) |rank| { + var pl = try plan.build(a, &c, &p, .{ .rank = rank, .world = 2, .blocks = blocks.items }); + defer pl.deinit(); + for (pl.layers) |*l| { + const h = l.experts[1].layout.histogram(); + std.debug.print("rank {d} {s} {t}: down {t} K2 4/6/8 = {d}/{d}/{d}, gate/up mismatches {d}, experts {d} B, block {d} B\n", .{ rank, l.prefix, l.mode, l.experts[1].layout.kind, h[4], h[6], h[8], l.gate_up_mismatch, l.expertBytes(), l.bytes() }); + } + } +} diff --git a/zig/src/families/deepseek_v41/kernels.zig b/zig/src/families/deepseek_v41/kernels.zig new file mode 100644 index 0000000000..38f446cfb8 --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels.zig @@ -0,0 +1,156 @@ +//! DeepSeek-V4.1's CUDA kernels (zig/kernels/cuda/deepseek_v41) as embedded fatbins, and their launches. + +const std = @import("std"); +const cuda = @import("cuda"); +const options = @import("dsv41_kernel_options"); + +pub const exl3 = @import("kernels_exl3.zig"); +pub const dense = @import("kernels_dense.zig"); +pub const kv_glue = @import("kernels_kv_glue.zig"); +pub const replay = @import("kernels_replay.zig"); +pub const ops = @import("kernels_ops.zig"); +/// sampling.cu (its own module, loaded by the keyed sampler) +pub const sampling = @import("sampling_kernels.zig"); + +fn Blob(comptime import_name: []const u8) type { + return struct { + pub const bytes align(16) = @embedFile(import_name).*; + }; +} + +/// False in host-only builds (no nvcc, no prebuilt fatbins); every image is then empty. +pub const available = options.with_kernels; + +/// The fatbins, by the names zig/build/cuda.zig gives them (dsv41_). +pub const images = struct { + pub const exl3_experts: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_exl3_experts").bytes else &.{}; + pub const x3ld: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_x3ld").bytes else &.{}; + pub const x3pf: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_x3pf").bytes else &.{}; + pub const x3gm: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_x3gm").bytes else &.{}; + pub const linear: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_linear").bytes else &.{}; + pub const x3seg: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_x3seg").bytes else &.{}; + pub const dense3: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_dense3").bytes else &.{}; + pub const attn: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_attn").bytes else &.{}; + pub const topk: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_topk").bytes else &.{}; + pub const mhc: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_mhc").bytes else &.{}; + pub const mhc_pf: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_mhc_pf").bytes else &.{}; + pub const router_gemv: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_router_gemv").bytes else &.{}; + pub const pfdense_k8: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_pfdense_k8").bytes else &.{}; + pub const pfdense_k10: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_pfdense_k10").bytes else &.{}; + pub const pfdense_k12: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_pfdense_k12").bytes else &.{}; + pub const pfdense_k16: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_pfdense_k16").bytes else &.{}; + pub const engram_gate: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_engram_gate").bytes else &.{}; + pub const l2pace: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_l2pace").bytes else &.{}; + pub const l2pf: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_l2pf").bytes else &.{}; + pub const x3gm_plan: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_x3gm_plan").bytes else &.{}; + pub const topk_keys: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_topk_keys").bytes else &.{}; + pub const pointwise: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_pointwise").bytes else &.{}; + pub const kvsplit: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_kvsplit").bytes else &.{}; + pub const glue: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_glue").bytes else &.{}; + pub const x3gm3: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_x3gm3").bytes else &.{}; + pub const pfglue: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_pfglue").bytes else &.{}; + pub const kv_glue: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_kv_glue").bytes else &.{}; + pub const router_glue: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_router_glue").bytes else &.{}; + pub const x3ld_epi: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_x3ld_epi").bytes else &.{}; + pub const topk_b: []const u8 = if (available) &Blob("dsv41_fatbin_dsv41_topk_b").bytes else &.{}; +}; + +/// lin2.cu (ours, TF_DSV41_LIN2): optional (zig/build/cuda.zig), an empty image when the build has none. +const Lin2Blob = struct { + pub const bytes align(16) = @embedFile("dsv41_fatbin_dsv41_lin2").*; +}; +pub const lin2_image: []const u8 = &Lin2Blob.bytes; + +fn env(name: [:0]const u8) ?[]const u8 { + const v = std.c.getenv(name) orelse return null; + return std.mem.span(v); +} + +pub const Kernels = struct { + mods: [30]cuda.Module, + kv_glue: kv_glue.Functions, + exl3: exl3.Functions, + dense: dense.Functions, + ops: ops.Functions, + sms: usize, + + /// Loads every module and resolves every instance the launches can use. + pub fn load(ctx: *const cuda.Context) !Kernels { + if (!available) return error.BuiltWithoutKernels; + const d = ctx.d; + var k: Kernels = undefined; + const list = [_][]const u8{ images.exl3_experts, images.x3ld, images.x3pf, images.x3gm, images.linear, images.x3seg, images.dense3, images.attn, images.topk, images.mhc, images.mhc_pf, images.router_gemv, images.pfdense_k8, images.pfdense_k10, images.pfdense_k12, images.pfdense_k16, images.engram_gate, images.l2pace, images.l2pf, images.x3gm_plan, images.topk_keys, images.pointwise, images.kvsplit, images.glue, images.x3gm3, images.pfglue, images.router_glue, images.kv_glue, images.x3ld_epi, images.topk_b }; + comptime std.debug.assert(list.len == @typeInfo(@FieldType(Kernels, "mods")).array.len); + var loaded: usize = 0; + errdefer for (k.mods[0..loaded]) |*m| m.unload(); + for (list, 0..) |img, i| { + k.mods[i] = try cuda.Module.load(d, img); + loaded += 1; + } + // the per-block shared memory opt-in (sm_120 / sm_121: 101,376 B): instances above it are refused one by one + const optin: u32 = @intCast(try ctx.attribute(.max_shared_memory_per_block_optin)); + k.kv_glue = try kv_glue.Functions.resolve(k.mods[27]); + k.exl3 = try exl3.Functions.resolve(k.mods[0], k.mods[1], k.mods[2], k.mods[3], k.mods[19], k.mods[24], k.mods[26], k.mods[28], optin); + k.dense = try dense.Functions.resolve(k.mods[4], k.mods[5], k.mods[6]); + const m = k.mods; + k.ops = try ops.Functions.resolve(m[7], m[8], m[9], m[10], m[11], .{ m[12], m[13], m[14], m[15] }, m[16], m[17], m[18], m[20], m[21], m[22], m[23], m[25], m[29], optin); + k.sms = @intCast(try ctx.attribute(.multiprocessor_count)); + try k.loadLin2(d, blk: { + var lm = try dense.lin2ModeWk(env("TF_DSV41_LIN2"), env("TF_DSV41_LIN2_PF"), env("TF_DSV41_LIN2_WK")); + lm.cluster = try dense.lin2Cluster(env("TF_DSV41_LIN2_CLUSTER")); + break :blk lm; + }); + return k; + } + + /// TF_DSV41_LIN2: lin2.cu's module and instances (dense.Ops); refused when the build has no lin2 fatbin. + pub fn loadLin2(k: *Kernels, d: *const cuda.Driver, mode: dense.Lin2Mode) !void { + if (!mode.lanes and !mode.linear and !mode.cluster) return; + if (lin2_image.len == 0) return error.Lin2KernelNotBuilt; + var m = try cuda.Module.load(d, lin2_image); + errdefer m.unload(); + k.dense.lin2 = try dense.Lin2.resolve(m, mode); + std.log.scoped(.dsv41).info("lin2: 17-32-row {s} launches as two-tile variants (one pass over the weights){s}{s}", .{ if (mode.lanes and mode.linear) "dense3 and linear.cu" else if (mode.lanes) "dense3" else "linear.cu", if (mode.linear and !mode.pf) ", linear2 without the k-step prefetch" else "", if (mode.lanes and !mode.wk8) ", lanes2 at 4-warp CTAs only" else if (mode.lanes and !mode.wk4) ", lanes2 at 8-warp CTAs only" else "" }); + if (mode.cluster) std.log.scoped(.dsv41).info("lin2: dense3's split-K launches of up to 32 rows reduce their partials in a cluster's shared memory (TF_DSV41_LIN2_CLUSTER)", .{}); + } + + pub fn deinit(k: *Kernels) void { + if (k.dense.lin2) |*l| l.module.unload(); + for (&k.mods) |*m| m.unload(); + } + + /// The routed-expert launches on one stream. + pub fn experts(k: *const Kernels, s: cuda.Stream) exl3.Ops { + return .{ .f = &k.exl3, .s = s, .sms = k.sms }; + } + + /// Attention, top-k, mHC, router, pfdense, Engram gate and L2 launches on one stream. + pub fn others(k: *const Kernels, s: cuda.Stream) ops.Ops { + return .{ .f = &k.ops, .s = s }; + } + + pub fn kvStore(k: *const Kernels, s: cuda.Stream) kv_glue.Ops { + return .{ .f = &k.kv_glue, .stream = s }; + } + + /// The dense-linear launches on one stream. + pub fn linears(k: *const Kernels, s: cuda.Stream) dense.Ops { + return .{ .f = &k.dense, .s = s }; + } +}; + +test { + std.testing.refAllDecls(kv_glue); + std.testing.refAllDecls(kv_glue.Ops); + std.testing.refAllDecls(exl3); + std.testing.refAllDecls(exl3.Ops); + std.testing.refAllDecls(exl3.Functions); + std.testing.refAllDecls(dense); + std.testing.refAllDecls(dense.Ops); + std.testing.refAllDecls(dense.Functions); + std.testing.refAllDecls(Kernels); + std.testing.refAllDecls(replay); + std.testing.refAllDecls(ops); + std.testing.refAllDecls(ops.Ops); + std.testing.refAllDecls(ops.Functions); +} diff --git a/zig/src/families/deepseek_v41/kernels_dense.zig b/zig/src/families/deepseek_v41/kernels_dense.zig new file mode 100644 index 0000000000..b31bec7ebc --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_dense.zig @@ -0,0 +1,555 @@ +//! DeepSeek-V4.1's dense EXL3 linears (mul1): linear.cu, x3seg segment tables and dense3's "lanes" layout. + +const std = @import("std"); +const cuda = @import("cuda"); +const lin2 = @import("kernels_lin2.zig"); + +pub const DType = enum(c_int) { f16 = 0, bf16 = 1, f32 = 2 }; + +pub fn dtypeSize(t: DType) usize { + return switch (t) { + .f16, .bf16 => 2, + .f32 => 4, + }; +} + +// -- Plans and tables (pure: host tests cover them) ------------------------------------------------------------- + +/// linear.plan: (K splits, warps a program) for a K x N layer, the shape's alone (every row count one reduction). +pub fn plan(k: usize, n: usize) [2]usize { + const blocks = 192; + const min_tiles = 8; + 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, wk }; +} + +/// Words between k tiles and between 128-column blocks of a matrix in the "strips" layout (both dense layouts). +pub fn strides(k: usize, k2: u32) [2]i64 { + const tw: i64 = 4 * @as(i64, k2); + return .{ 8 * tw, @as(i64, @intCast(k / 16)) * 8 * tw }; +} + +pub const max_seg = 8; // x3seg.cu MAXSEG +pub const rows_max = 128; // fused_proj.ROWS / dense3.ROWS: rows a launch +pub const x3seg_k2s = [_]u32{ 4, 6, 8, 10, 12, 16 }; +pub const dense3_k2s = [_]u32{ 4, 6, 8, 10, 12 }; + +/// fused_proj.Seg: one projection as the kernel sees it, and where its output goes (buffer `out`, column `col`). +pub const Seg = struct { + k: usize, + n: usize, + k2: u32, + sk: usize, + wk: usize, + out: usize = 0, + col: usize = 0, + + pub fn programs(s: Seg) usize { + return (s.n / 128) * s.sk; + } +}; + +pub const Launch = struct { + k2: u32, + wk: usize, + segs: [max_seg]usize = undefined, // indices into the call's segments + count: usize = 0, + programs: usize = 0, + p0: [max_seg]usize = undefined, // first program of each segment in this launch +}; + +/// fused_proj.Layout: scratch offsets (fp16 rotated inputs, Z floats, counter ints) and the launches at `m` rows. +pub const Layout = struct { + m: usize, + xh_off: [64]usize = undefined, + z_off: [64]usize = undefined, + c_off: [64]usize = undefined, + xh_total: usize = 0, + z_total: usize = 0, + c_total: usize = 0, + launches: [16]Launch = undefined, + n_launches: usize = 0, +}; + +pub const Error = error{ Shape, Unsupported, TooMany }; + +/// fused_proj.layout: one launch a width in first-appearance order, up to 8 segments; WK = the segments' largest. +pub fn layout(segs: []const Seg, m: usize) Error!Layout { + if (m < 1 or m > rows_max) return error.Shape; + if (segs.len > 64) return error.TooMany; + var lay: Layout = .{ .m = m }; + for (segs, 0..) |s, i| { + if (s.k % 128 != 0 or s.n % 128 != 0 or s.sk < 1 or (s.wk != 4 and s.wk != 8) or (s.k / 16) % (s.sk * s.wk) != 0) return error.Shape; + lay.xh_off[i] = lay.xh_total; + lay.z_off[i] = lay.z_total; + lay.c_off[i] = lay.c_total; + lay.xh_total += m * s.k; + if (s.sk > 1) lay.z_total += s.sk * m * s.n; + lay.c_total += 8 * (s.n / 128); + } + var seen: [64]u32 = undefined; + var nseen: usize = 0; + for (segs) |s| { + if (std.mem.indexOfScalar(u32, seen[0..nseen], s.k2) == null) { + seen[nseen] = s.k2; + nseen += 1; + } + } + for (seen[0..nseen]) |k2| { + var l: Launch = .{ .k2 = k2, .wk = 0 }; + for (segs, 0..) |s, i| { + if (s.k2 != k2) continue; + if (l.count == max_seg) { + if (lay.n_launches == lay.launches.len) return error.TooMany; + lay.launches[lay.n_launches] = l; + lay.n_launches += 1; + l = .{ .k2 = k2, .wk = 0 }; + } + l.segs[l.count] = i; + l.p0[l.count] = l.programs; + l.count += 1; + l.programs += s.programs(); + l.wk = @max(l.wk, s.wk); + } + if (lay.n_launches == lay.launches.len) return error.TooMany; + lay.launches[lay.n_launches] = l; + lay.n_launches += 1; + } + return lay; +} + +/// dense3.group_steps: k steps a load group (whole 16-byte chunks a lane). +pub fn groupSteps(k2: u32) usize { + return if (k2 == 16) 1 else 2; +} + +/// x3seg.cu rot table entry (by value); R1c: with `nw` set the input is rmsnorm-normed in the kernel (inv_k, eps). +pub const RotSeg = extern struct { x: u64, suh: u64, xh: u64, ld: i64, K: c_int, x_dtype: c_int, nw: u64 = 0, nw_dtype: c_int = 0, inv_k: f32 = 0, eps: f32 = 0 }; +pub const RotTable = extern struct { s: [max_seg]RotSeg, n: c_int }; +pub const LinSeg = extern struct { + xh: u64, + T: u64, + stride_k: i64, + stride_nb: i64, + svh: u64, + y: u64, + Z: u64, + counters: u64, + K: c_int, + N: c_int, + SK: c_int, + wk: c_int, + p0: c_int, + ld: c_int, +}; +pub const LinTable = extern struct { s: [max_seg]LinSeg, n: c_int }; + +comptime { + std.debug.assert(@sizeOf(RotSeg) == 64 and @offsetOf(RotSeg, "nw") == 40 and @sizeOf(RotTable) == 520); + std.debug.assert(@sizeOf(LinSeg) == 88 and @sizeOf(LinTable) == 712); +} + +/// One segment's buffers for a linear launch: matrix words, svh, output `y` (row stride `ld` elements). +pub const SegBufs = struct { T: u64, svh: u64, y: u64, ld: usize }; + +/// One segment's rotation input: rows at x + row * ld, dtype, suh [K]; `nw`: folded q_norm weight (0 for none). +pub const RotIn = struct { x: u64, ld: usize, dtype: DType, suh: u64, nw: u64 = 0, nw_dtype: DType = .bf16 }; + +/// fp32(1 / K): the binding's `static_cast(1.0 / K)`. +pub fn invK(k: usize) f32 { + return @floatCast(1.0 / @as(f64, @floatFromInt(k))); +} + +/// fused_proj._block's rot_in tables: one launch per 8 segments (`first` = 0, 8, ...), rows at their xh offsets. +pub fn rotTable(lay: Layout, segs: []const Seg, ins: []const RotIn, first: usize, xh: u64, eps: f32) RotTable { + var t = std.mem.zeroes(RotTable); + const end = @min(segs.len, first + max_seg); + t.n = @intCast(end - first); + for (first..end, 0..) |i, j| { + t.s[j] = .{ .x = ins[i].x, .suh = ins[i].suh, .xh = xh + 2 * lay.xh_off[i], .ld = @intCast(ins[i].ld), .K = @intCast(segs[i].k), .x_dtype = @backingInt(ins[i].dtype) }; + if (ins[i].nw != 0) { + t.s[j].nw = ins[i].nw; + t.s[j].nw_dtype = @backingInt(ins[i].nw_dtype); + t.s[j].inv_k = invK(segs[i].k); + t.s[j].eps = eps; + } + } + return t; +} + +/// dsv41_x3seg_linear_cuda / dsv41_dense3_linear_cuda's table of one launch. +pub fn linTable(lay: Layout, l: Launch, segs: []const Seg, bufs: []const SegBufs, y_dtype: DType, xh: u64, z: u64, counters: u64) LinTable { + var t = std.mem.zeroes(LinTable); + t.n = @intCast(l.count); + const esz = dtypeSize(y_dtype); + for (l.segs[0..l.count], 0..) |i, j| { + const s = segs[i]; + const st = strides(s.k, s.k2); + t.s[j] = .{ + .xh = xh + 2 * lay.xh_off[i], + .T = bufs[i].T, + .stride_k = st[0], + .stride_nb = st[1], + .svh = bufs[i].svh, + .y = bufs[i].y + s.col * esz, + .Z = if (s.sk > 1) z + 4 * lay.z_off[i] else 0, + .counters = counters + 4 * lay.c_off[i], + .K = @intCast(s.k), + .N = @intCast(s.n), + .SK = @intCast(s.sk), + .wk = @intCast(s.wk), + .p0 = @intCast(l.p0[j]), + .ld = @intCast(bufs[i].ld), + }; + } + return t; +} + +// -- Symbols ---------------------------------------------------------------------------------------------------- + +pub const linear_k2s = [_]u32{ 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16 }; + +fn linearSymbol(comptime k2: u32, comptime wk: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN14tf_exl3_linear13linear_kernelILi{d}ELi2ELi{d}EEEvPK6__halfPKjxxS3_S3_PviPfPiiiiiPKi", .{ k2, wk }); +} +fn unpackSymbol(comptime k2: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN14tf_exl3_linear13unpack_kernelILi{d}ELi2EEEvPKjP6__halfill", .{k2}); +} +fn segSymbol(comptime k2: u32, comptime wk: u32, comptime mb: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN11dsv41_x3seg17seg_linear_kernelILi{d}ELi2ELi{d}ELi{d}EEEvNS_8LinTableEii", .{ k2, wk, mb }); +} +fn lanesSymbol(comptime k2: u32, comptime wk: u32, comptime mb: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN12dsv41_dense319lanes_linear_kernelILi{d}ELi2ELi{d}ELi{d}EEEvN11dsv41_x3seg8LinTableEii", .{ k2, wk, mb }); +} +fn lanesUnpackSymbol(comptime k2: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN12dsv41_dense319lanes_unpack_kernelILi{d}ELi2EEEvPKjP6__halfill", .{k2}); +} + +pub const lanesc_cfgs = lin2.lanesc_cfgs; +pub const Lin2Mode = lin2.Lin2Mode; +pub const lin2Cluster = lin2.lin2Cluster; +pub const clusterSplits = lin2.clusterSplits; +pub const lin2Mode = lin2.lin2Mode; +pub const lin2ModeWk = lin2.lin2ModeWk; +pub const twoTile = lin2.twoTile; +pub const Lin2 = lin2.Lin2; + +/// x3seg's (WK, MINB) instances (TF_SEG_WIDTH: WK 8 takes any minb) and dense3's (TF_D3_WIDTH). +pub const seg_cfgs = [_][2]u32{ .{ 4, 3 }, .{ 4, 2 }, .{ 8, 1 } }; +pub const lanes_cfgs = [_][2]u32{ .{ 4, 3 }, .{ 4, 4 }, .{ 8, 1 }, .{ 8, 2 } }; + +pub const Functions = struct { + rot_in: cuda.Function, + linear: [linear_k2s.len][3]cuda.Function, // WK 2, 4, 8 + unpack: [linear_k2s.len]cuda.Function, + seg_rot_in: cuda.Function, + seg: [x3seg_k2s.len][seg_cfgs.len]cuda.Function, + lanes: [dense3_k2s.len][lanes_cfgs.len]cuda.Function, + lanes_unpack: [dense3_k2s.len]cuda.Function, + to_lanes: cuda.Function, + to_strips: cuda.Function, + /// TF_DSV41_LIN2: the two-tile variants, set by Kernels.load when the knob asks for them + lin2: ?Lin2 = null, + + pub fn resolve(linear: cuda.Module, x3seg: cuda.Module, dense3: cuda.Module) !Functions { + var f: Functions = undefined; + f.lin2 = null; + f.rot_in = try linear.function("_ZN14tf_exl3_linear13rot_in_kernelEPKviPK6__halfPS2_i"); + inline for (linear_k2s, 0..) |k2, i| { + inline for (.{ 2, 4, 8 }, 0..) |wk, j| f.linear[i][j] = try linear.function(linearSymbol(k2, wk)); + f.unpack[i] = try linear.function(unpackSymbol(k2)); + } + f.seg_rot_in = try x3seg.function("_ZN11dsv41_x3seg17seg_rot_in_kernelENS_8RotTableE"); + inline for (x3seg_k2s, 0..) |k2, i| inline for (seg_cfgs, 0..) |c, j| { + f.seg[i][j] = try x3seg.function(segSymbol(k2, c[0], c[1])); + }; + inline for (dense3_k2s, 0..) |k2, i| { + inline for (lanes_cfgs, 0..) |c, j| f.lanes[i][j] = try dense3.function(lanesSymbol(k2, c[0], c[1])); + f.lanes_unpack[i] = try dense3.function(lanesUnpackSymbol(k2)); + } + f.to_lanes = try dense3.function("_ZN12dsv41_dense315to_lanes_kernelEPKjPjxii"); + f.to_strips = try dense3.function("_ZN12dsv41_dense316to_strips_kernelEPKjPjxii"); + return f; + } +}; + +fn dim(x: usize) u32 { + return @intCast(x); +} + +fn int(x: usize) c_int { + return @intCast(x); +} + +fn go(f: cuda.Function, s: cuda.Stream, grid: [3]usize, block: usize, shared: usize, pdl: bool, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = .{ .x = dim(grid[0]), .y = dim(grid[1]), .z = dim(grid[2]) }, .block = .{ .x = dim(block) }, .shared = dim(shared), .pdl = pdl }, s, a); +} + +/// WK * min(M, 8) * 128 floats: the linear kernels' reduction buffer (at most 32 KiB: no opt-in needed). +fn redSmem(wk: usize, m: usize) usize { + return wk * @min(m, 8) * 128 * 4; +} + +pub const Ops = struct { + f: *const Functions, + s: cuda.Stream, + + /// exl3_rot_in_cuda: xh [M, K] fp16 = (x * suh) H / sqrt 128 by 128-blocks. + pub fn rotIn(o: Ops, x: u64, x_dtype: DType, suh: u64, xh: u64, M: usize, K: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(@backingInt(x_dtype)); + a.add(suh); + a.add(xh); + a.add(int(K)); + try go(o.f.rot_in, o.s, .{ (K / 128 + 3) / 4, M, 1 }, 128, 0, false, &a); + } + + /// exl3_linear_cuda (mul1): y [M, N] from xh [M, K]; `z` fp32 [SK, M, N] when SK > 1; `skip` a device flag or 0. + pub fn linear(o: Ops, xh: u64, T: u64, stride_k: i64, stride_nb: i64, svh: u64, bias: u64, y: u64, y_dtype: DType, z: u64, counters: u64, M: usize, K: usize, N: usize, k2: u32, SK: usize, WK: usize, skip: u64) !void { + const ki = std.mem.indexOfScalar(u32, &linear_k2s, k2) orelse return error.Unsupported; + const wi: usize = switch (WK) { + 2 => 0, + 4 => 1, + 8 => 2, + else => return error.Unsupported, + }; + if (SK == 0 or (K / 16) % (SK * WK) != 0) return error.Shape; + if (SK > 1 and z == 0) return error.Shape; + if (M < 1 or M > rows_max) return error.Shape; + var a: cuda.Args = .{}; + for ([_]u64{ xh, T }) |v| a.add(v); + a.add(stride_k); + a.add(stride_nb); + for ([_]u64{ svh, bias, y }) |v| a.add(v); + a.add(@backingInt(y_dtype)); + for ([_]u64{ z, counters }) |v| a.add(v); + for ([_]usize{ M, K, N, SK }) |v| a.add(int(v)); + a.add(skip); + var f = o.f.linear[ki][wi]; + if (o.f.lin2) |l2| if (l2.mode.linear and twoTile(M)) { + f = l2.linear[ki][wi]; // TF_DSV41_LIN2: the same programs, one pass over the weights + }; + try go(f, o.s, .{ N / 128, SK, 1 }, WK * 32, redSmem(WK, M), false, &a); + } + + /// exl3_unpack_cuda (mul1): W_q [K, N] fp16. + pub fn unpack(o: Ops, T: u64, W: u64, K: usize, N: usize, stride_k: i64, stride_nb: i64, k2: u32) !void { + const ki = std.mem.indexOfScalar(u32, &linear_k2s, k2) orelse return error.Unsupported; + var a: cuda.Args = .{}; + a.add(T); + a.add(W); + a.add(int(N)); + a.add(stride_k); + a.add(stride_nb); + try go(o.f.unpack[ki], o.s, .{ N / 16, K / 16, 1 }, 32, 0, false, &a); + } + + /// dsv41_x3seg_rot_in_cuda: up to 8 (input, suh) pairs into one xh (each at its xh offset), one launch. + pub fn segRotIn(o: Ops, tab: RotTable, M: usize, pdl: bool) !void { + if (tab.n < 1 or tab.n > max_seg) return error.Shape; + for (tab.s[0..@intCast(tab.n)]) |s| { + if (s.nw != 0 and s.x_dtype != @backingInt(DType.bf16)) return error.Shape; // the fold takes bf16 rows + } + var kmax: usize = 0; + for (tab.s[0..@intCast(tab.n)]) |s| kmax = @max(kmax, @as(usize, @intCast(s.K))); + var a: cuda.Args = .{}; + a.add(tab); + try go(o.f.seg_rot_in, o.s, .{ (kmax / 128 + 3) / 4, M, @intCast(tab.n) }, 128, 0, pdl, &a); + } + + /// dsv41_x3seg_linear_cuda: one launch of a layout (words in the strips layout). + pub fn segLinear(o: Ops, tab: LinTable, y_dtype: DType, M: usize, l: Launch, minb: u32, pdl: bool) !void { + try o.segLinearRaw(tab, y_dtype, M, l.programs, l.k2, l.wk, minb, pdl); + } + + /// segLinear with the binding's own (programs, K2, WK, minb): WK 8 takes any minb. + pub fn segLinearRaw(o: Ops, tab: LinTable, y_dtype: DType, M: usize, programs: usize, k2: u32, wk: usize, minb: u32, pdl: bool) !void { + const ki = std.mem.indexOfScalar(u32, &x3seg_k2s, k2) orelse return error.Unsupported; + const ci: usize = if (wk == 8) 2 else if (wk == 4 and minb == 3) 0 else if (wk == 4 and minb == 2) 1 else return error.Unsupported; + try o.tableLaunch(o.f.seg[ki][ci], tab, y_dtype, M, programs, wk, pdl); + } + + /// dsv41_dense3_linear_cuda: the same programs on words in the lanes layout. + pub fn lanesLinear(o: Ops, tab: LinTable, y_dtype: DType, M: usize, l: Launch, minb: u32, pdl: bool) !void { + try o.lanesLinearRaw(tab, y_dtype, M, l.programs, l.k2, l.wk, minb, pdl); + } + + /// lanesLinear with the binding's own arguments; each segment's k steps must split into whole load groups. + pub fn lanesLinearRaw(o: Ops, tab: LinTable, y_dtype: DType, M: usize, programs: usize, k2: u32, wk: usize, minb: u32, pdl: bool) !void { + const ki = std.mem.indexOfScalar(u32, &dense3_k2s, k2) orelse return error.Unsupported; + const want = [2]u32{ @intCast(wk), minb }; + var ci: ?usize = null; + for (lanes_cfgs, 0..) |c, j| if (std.mem.eql(u32, &c, &want)) { + ci = j; + }; + for (tab.s[0..@intCast(tab.n)]) |s| { + const per_warp: usize = @intCast(@divTrunc(@divTrunc(@divTrunc(s.K, 16), s.SK), s.wk)); + if (per_warp % groupSteps(k2) != 0) return error.Shape; + } + const cfg = ci orelse return error.Unsupported; + var f = o.f.lanes[ki][cfg]; + const width_ok = if (o.f.lin2) |l2| (if (wk == 8) l2.mode.wk8 else l2.mode.wk4) else false; + // TF_DSV41_LIN2_CLUSTER: a column block's split programs in one cluster, partials reduced in shared memory + if (o.f.lin2) |l2| if (l2.mode.cluster) if (clusterSplits(tab, M)) |sk| { + const two = twoTile(M); + if (!two or (l2.mode.lanes and width_ok)) { + const j: usize = (if (two) @as(usize, 2) else 0) + @intFromBool(wk == 8); + const t: usize = if (two) 2 else 1; + if (tab.n < 1 or tab.n > max_seg) return error.Shape; + var a: cuda.Args = .{}; + a.add(tab); + a.add(@backingInt(y_dtype)); + a.add(int(M)); + const shared = redSmem(wk, M) + 16 * t * 128 * 4; + return cuda.launch.launch(l2.lanesc[ki][j], .{ .grid = .{ .x = dim(programs), .y = 1, .z = 1 }, .block = .{ .x = dim(wk * 32) }, .shared = dim(shared), .cluster = .{ .x = sk, .y = 1, .z = 1 }, .pdl = pdl }, o.s, &a); + } + }; + if (o.f.lin2) |l2| if (l2.mode.lanes and twoTile(M) and width_ok) { + f = l2.lanes[ki][if (wk == 8) 1 else 0]; // TF_DSV41_LIN2: the same programs, one pass over the weights + }; + try o.tableLaunch(f, tab, y_dtype, M, programs, wk, pdl); + } + + fn tableLaunch(o: Ops, f: cuda.Function, tab: LinTable, y_dtype: DType, M: usize, programs: usize, wk: usize, pdl: bool) !void { + if (M < 1 or M > rows_max or tab.n < 1 or tab.n > max_seg) return error.Shape; + var a: cuda.Args = .{}; + a.add(tab); + a.add(@backingInt(y_dtype)); + a.add(int(M)); + try go(f, o.s, .{ programs, 1, 1 }, wk * 32, redSmem(wk, M), pdl, &a); + } + + /// dsv41_dense3_relayout_cuda: int32 words of one matrix ([N/128, K/16, 8, 4 K2]) to the lanes layout or back. + pub fn relayout(o: Ops, src: u64, dst: u64, words: usize, k2: u32, K: usize, lanes: bool) !void { + var a: cuda.Args = .{}; + a.add(src); + a.add(dst); + a.add(@as(i64, @intCast(words))); + a.add(@as(c_int, @intCast(k2))); + a.add(int(K / 16)); + try go(if (lanes) o.f.to_lanes else o.f.to_strips, o.s, .{ (words + 255) / 256, 1, 1 }, 256, 0, false, &a); + } + + /// dsv41_dense3_unpack_cuda: W_q [K, N] fp16 from lanes-layout words. + pub fn lanesUnpack(o: Ops, T: u64, W: u64, K: usize, N: usize, stride_k: i64, stride_nb: i64, k2: u32) !void { + const ki = std.mem.indexOfScalar(u32, &dense3_k2s, k2) orelse return error.Unsupported; + var a: cuda.Args = .{}; + a.add(T); + a.add(W); + a.add(int(N)); + a.add(stride_k); + a.add(stride_nb); + try go(o.f.lanes_unpack[ki], o.s, .{ N / 16, K / 16, 1 }, 32, 0, false, &a); + } +}; + +test "linear.plan as upstream" { + // wide layers: 8 warps, no split; narrow ones split K until 8 k tiles a warp would not hold + try std.testing.expectEqual([2]usize{ 1, 8 }, plan(1024, 128 * 64)); + try std.testing.expectEqual([2]usize{ 8, 4 }, plan(4096, 1024)); + try std.testing.expectEqual([2]usize{ 1, 4 }, plan(512, 4096)); +} + +test "rot_in tables: one launch a group of 8, xh offsets from the layout" { + var segs: [10]Seg = undefined; + var ins: [10]RotIn = undefined; + for (&segs, &ins, 0..) |*s, *r, i| { + s.* = .{ .k = 1024, .n = 256, .k2 = 10, .sk = 1, .wk = 4 }; + r.* = .{ .x = 0x1000 * (i + 1), .ld = 1024, .dtype = .bf16, .suh = 0x100 * (i + 1) }; + } + const lay = try layout(&segs, 2); + const t0 = rotTable(lay, &segs, &ins, 0, 0x10000, 0); + const t1 = rotTable(lay, &segs, &ins, 8, 0x10000, 0); + try std.testing.expectEqual(@as(c_int, 8), t0.n); + try std.testing.expectEqual(@as(c_int, 2), t1.n); + try std.testing.expectEqual(@as(u64, 0x10000 + 2 * 8 * 2 * 1024), t1.s[0].xh); + try std.testing.expectEqual(@as(c_int, @backingInt(DType.bf16)), t1.s[1].x_dtype); +} + +test "fused_proj.layout: offsets, one launch a width, 8 segments a launch" { + const segs = [_]Seg{ + .{ .k = 4096, .n = 1024, .k2 = 10, .sk = 4, .wk = 4 }, + .{ .k = 4096, .n = 512, .k2 = 12, .sk = 8, .wk = 4, .out = 1 }, + .{ .k = 4096, .n = 256, .k2 = 10, .sk = 8, .wk = 8, .out = 0, .col = 1024 }, + }; + const lay = try layout(&segs, 3); + try std.testing.expectEqual(@as(usize, 3 * 4096 * 3), lay.xh_total); + try std.testing.expectEqual(@as(usize, 4 * 3 * 1024 + 8 * 3 * 512 + 8 * 3 * 256), lay.z_total); + try std.testing.expectEqual(@as(usize, 8 * (8 + 4 + 2)), lay.c_total); + try std.testing.expectEqual(@as(usize, 2), lay.n_launches); + const l0 = lay.launches[0]; + try std.testing.expectEqual(@as(u32, 10), l0.k2); + try std.testing.expectEqual(@as(usize, 8), l0.wk); + try std.testing.expectEqualSlices(usize, &.{ 0, 2 }, l0.segs[0..l0.count]); + try std.testing.expectEqualSlices(usize, &.{ 0, 32 }, l0.p0[0..l0.count]); + try std.testing.expectEqual(@as(usize, 32 + 16), l0.programs); + try std.testing.expectError(error.Shape, layout(&segs, 129)); +} + +test "TF_DSV41_LIN2 reader: unset / 0 / off none, 1 both, lanes / linear one, _PF 1 / 0, anything else refused" { + const t = std.testing; + for ([_]?[]const u8{ null, "", "0", "off" }) |v| { + const m = try lin2Mode(v, null); + try t.expect(!m.lanes and !m.linear and m.pf); + } + try t.expectEqual(Lin2Mode{ .lanes = true, .linear = true }, try lin2Mode("1", null)); + try t.expectEqual(Lin2Mode{ .lanes = true }, try lin2Mode(" lanes ", "1")); + try t.expectEqual(Lin2Mode{ .linear = true, .pf = false }, try lin2Mode("linear", "0")); + for ([_][]const u8{ "2", "head", "lanes,linear" }) |v| try t.expectError(error.BadLin2, lin2Mode(v, null)); + try t.expectError(error.BadLin2, lin2Mode("1", "on")); + // TF_DSV41_LIN2_WK: lanes2 at one CTA width only + const four = try lin2ModeWk("lanes", null, "4"); + try t.expect(four.lanes and four.wk4 and !four.wk8); + const eight = try lin2ModeWk("1", "0", " 8 "); + try t.expect(eight.lanes and eight.linear and !eight.pf and !eight.wk4 and eight.wk8); + const both = try lin2ModeWk("1", null, null); + try t.expect(both.wk4 and both.wk8); + for ([_][]const u8{ "2", "48", "all" }) |w| try t.expectError(error.BadLin2, lin2ModeWk("1", null, w)); +} + +test "TF_DSV41_LIN2: only the original's two-pass launches (17-32 rows) take the variant" { + const t = std.testing; + for (1..17) |m| try t.expect(!twoTile(m)); + for (17..33) |m| try t.expect(twoTile(m)); + for (33..129) |m| try t.expect(!twoTile(m)); +} + +test "TF_DSV41_LIN2_CLUSTER: the reader, and the launches the cluster variant takes" { + const t = std.testing; + for ([_]?[]const u8{ null, "", "0" }) |v| try t.expect(!try lin2Cluster(v)); + try t.expect(try lin2Cluster(" 1 ")); + for ([_][]const u8{ "on", "2", "yes" }) |v| try t.expectError(error.BadLin2, lin2Cluster(v)); + var tab = std.mem.zeroes(LinTable); + tab.n = 2; + tab.s[0] = .{ .xh = 0, .T = 0, .stride_k = 0, .stride_nb = 0, .svh = 0, .y = 0, .Z = 0, .counters = 0, .K = 4096, .N = 512, .SK = 4, .wk = 4, .p0 = 0, .ld = 512 }; + tab.s[1] = tab.s[0]; + tab.s[1].p0 = 16; // 4 column blocks x 4 splits + try t.expectEqual(@as(?u32, 4), clusterSplits(tab, 1)); + try t.expectEqual(@as(?u32, 4), clusterSplits(tab, 32)); + try t.expectEqual(@as(?u32, null), clusterSplits(tab, 33)); // past two row tiles + tab.s[1].SK = 2; // a launch whose segments split differently + try t.expectEqual(@as(?u32, null), clusterSplits(tab, 8)); + tab.s[1].SK = 4; + tab.s[1].p0 = 18; // a segment not starting on a cluster boundary + try t.expectEqual(@as(?u32, null), clusterSplits(tab, 8)); + tab.s[1].p0 = 16; + tab.s[0].SK = 1; + tab.s[1].SK = 1; // no split: no partials to reduce + try t.expectEqual(@as(?u32, null), clusterSplits(tab, 8)); + for (&tab.s) |*x| x.SK = 16; // past the portable cluster size + try t.expectEqual(@as(?u32, null), clusterSplits(tab, 8)); +} diff --git a/zig/src/families/deepseek_v41/kernels_exl3.zig b/zig/src/families/deepseek_v41/kernels_exl3.zig new file mode 100644 index 0000000000..bb25a603b8 --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_exl3.zig @@ -0,0 +1,534 @@ +//! DeepSeek-V4.1's routed EXL3 expert kernels: x3ld, x3pf, x3gm and the grouped kernel with its helpers. + +const std = @import("std"); +const cuda = @import("cuda"); +const gmk = @import("kernels_x3gm.zig"); + +// -- Instances (the copies' footers hold the matching explicit instantiations) ---------------------------------- + +/// The (lo, hi) K2 ranges x3ld / x3pf are built for, tried in order (expert_loads.RANGES / expert_prefill.RANGES). +pub const ranges = [_][2]u32{ .{ 8, 8 }, .{ 2, 10 }, .{ 2, 12 } }; +/// Upstream's TF_RANGES: (8, 8), (2, 10), else (2, 16). +pub const upstream_ranges = [_][2]u32{ .{ 8, 8 }, .{ 2, 10 }, .{ 2, 16 } }; + +/// dispatch_ranges: the first instance whose range holds [lo, hi], or null (the binding's "unsupported"). +pub fn rangeIndex(lo: u32, hi: u32) ?usize { + if (lo == 8 and hi == 8) return 0; + if (lo >= 2 and hi <= 10) return 1; + if (lo >= 2 and hi <= 12) return 2; + return null; +} + +/// grouped_launch's TF_RANGES (always an instance: 2..16 catches the rest). +pub fn upstreamRangeIndex(lo: u32, hi: u32) usize { + if (lo == 8 and hi == 8) return 0; + if (lo >= 2 and hi <= 10) return 1; + return 2; +} + +/// x3ld's (nt, pd) settings (expert_loads.CFGS); probe 3 (timing only) has the first two. +pub const ld_cfgs = [_][2]u32{ .{ 8, 1 }, .{ 8, 2 }, .{ 4, 2 } }; +/// x3pf's (nt, mtl) settings. +pub const pf_cfgs = [_][2]u32{ .{ 4, 4 }, .{ 8, 2 } }; +/// upstream grouped_kernel's (nt, warps, pf) settings. +pub const grouped_cfgs = [_][3]u32{ .{ 8, 4, 1 }, .{ 8, 4, 2 }, .{ 4, 4, 2 } }; + +fn cfgIndex(comptime n: usize, table: []const [n]u32, want: [n]u32) ?usize { + for (table, 0..) |c, i| if (std.mem.eql(u32, &c, &want)) return i; + return null; +} + +const tail16 = "PK6__halfS3_PKlS5_PKiS7_S7_S7_S7_Pfiiiiii"; // (X0, X1, TP0, TP1, K2_0, K2_1, uids, ucount, members, Z, 6 ints) + +pub fn ldSymbol(comptime nt: u32, comptime pd: u32, comptime lo: u32, comptime hi: u32, comptime probe: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_x3ld9ld_kernelILi2ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}EEEv" ++ tail16, .{ nt, pd, lo, hi, probe }); +} + +/// TF_DSV41_X3LD_EPI (x3ld_epi.cu): x3ld (nt 8) with the gate/up (`gu`) or down_combine (`dn` fp32, `dnb` bf16) tail. +pub const EpiKind = enum { gu, dn, dnb }; +pub const epi_pds = [_]u32{ 1, 2 }; + +pub fn epiSymbol(comptime kind: EpiKind, comptime pd: u32, comptime lo: u32, comptime hi: u32) [:0]const u8 { + return std.fmt.comptimePrint("dsv41_x3ld_epi_{s}_{d}_{d}_{d}", .{ @tagName(kind), pd, lo, hi }); +} + +pub fn pfSymbol(comptime nt: u32, comptime mtl: u32, comptime lo: u32, comptime hi: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_x3pf9pf_kernelILi2ELi{d}ELi{d}ELi{d}ELi{d}EEEv" ++ tail16, .{ nt, mtl, lo, hi }); +} + +pub fn groupedSymbol(comptime nt: u32, comptime w: u32, comptime pf: u32, comptime lo: u32, comptime hi: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN8tf_exl3x14grouped_kernelILi2ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}EEEv" ++ tail16, .{ nt, w, pf, lo, hi }); +} + +pub fn dequantSymbol(comptime k2: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN8tf_exl3x14dequant_kernelILi2ELi{d}EEEvPKjP6__halfii", .{k2}); +} + +const up_sym = struct { + const group = "_ZN15tf_exl3_experts12group_kernelEPKiPiS2_S2_iiii"; + const rot_in_bf16 = "_ZN15tf_exl3_experts13rot_in_kernelI13__nv_bfloat16EEvPKT_iPKiPK6__halfS9_PS7_SA_iii"; + const rot_in_f16 = "_ZN15tf_exl3_experts13rot_in_kernelI6__halfEEvPKT_iPKiPKS1_S8_PS1_S9_iii"; + const gateup_epilogue = "_ZN15tf_exl3_experts22gateup_epilogue_kernelEPKfPKiPK6__halfS6_S6_PS4_iiiifi"; + const down_epilogue = "_ZN15tf_exl3_experts20down_epilogue_kernelEPKfPKiPK6__halfPfiiii"; + const combine = "_ZN15tf_exl3_experts14combine_kernelEPKfS1_Pfii"; + // templated on the output (R1: fp32, or bf16 rounded as stored = fp32 then .to(bf16)) + const down_combine_f32 = "_ZN15tf_exl3_experts19down_combine_kernelIfEEvPKfPKiPK6__halfPfS2_PT_iiiii"; + const down_combine_bf16 = "_ZN15tf_exl3_experts19down_combine_kernelI13__nv_bfloat16EEvPKfPKiPK6__halfPfS3_PT_iiiii"; +}; + +pub const gm_k2s = gmk.gm_k2s; +pub const gu_tiles = gmk.gu_tiles; +pub const dn_tiles = gmk.dn_tiles; +pub const gm_smem_max = gmk.gm_smem_max; +pub const gmK2Index = gmk.gmK2Index; +pub const GmShape = gmk.GmShape; +pub const v2_gu = gmk.v2_gu; +pub const gu2_cfg = gmk.gu2_cfg; +pub const v2_dn = gmk.v2_dn; +pub const gm2Shape = gmk.gm2Shape; +pub const V3 = gmk.V3; +pub const v3_variants = gmk.v3_variants; +pub const gm3Symbol = gmk.gm3Symbol; +pub const v3Of = gmk.v3Of; +pub const gmV3 = gmk.gmV3; +pub const v3_default = gmk.v3_default; +pub const gmTuned2 = gmk.gmTuned2; +pub const gm2WidthsOk = gmk.gm2WidthsOk; +pub const gm2GuSymbol = gmk.gm2GuSymbol; +pub const gm2DnSymbol = gmk.gm2DnSymbol; +pub const gmSmemBytes = gmk.gmSmemBytes; +pub const GmKind = gmk.GmKind; +pub const plan_max_experts = gmk.plan_max_experts; +pub const planPasses = gmk.planPasses; +pub const gmFits = gmk.gmFits; +pub const gmMembers = gmk.gmMembers; +pub const TuneRow = gmk.TuneRow; +pub const tune_default = gmk.tune_default; +pub const gmTuned = gmk.gmTuned; +pub const gmGuSymbol = gmk.gmGuSymbol; +pub const gmDnSymbol = gmk.gmDnSymbol; +pub const gmDequantSymbol = gmk.gmDequantSymbol; +pub const GmArgs = gmk.GmArgs; +pub const gmCheck = gmk.gmCheck; + +/// TF_DSV41_PF_TBO_GM_SMS (under TF_DSV41_PF_TBO): the SMs x3gm's persistent grid covers at most; 0: every SM +pub var gm_sms_cap: usize = 0; + +const gm_sym = struct { + const rot_bf16_1 = "_ZN10dsv41_x3gm10rot_kernelI13__nv_bfloat16Li1EEEvPKT_iPKiPK6__halfS9_PS7_SA_iii"; + const rot_bf16_2 = "_ZN10dsv41_x3gm10rot_kernelI13__nv_bfloat16Li2EEEvPKT_iPKiPK6__halfS9_PS7_SA_iii"; + const rot_f16_1 = "_ZN10dsv41_x3gm10rot_kernelI6__halfLi1EEEvPKT_iPKiPKS1_S8_PS1_S9_iii"; + const rot_f16_2 = "_ZN10dsv41_x3gm10rot_kernelI6__halfLi2EEEvPKT_iPKiPKS1_S8_PS1_S9_iii"; +}; + +// -- Loaded functions ------------------------------------------------------------------------------------------- + +pub const Gm = struct { f: cuda.Function, threads: u32, smem: u32, per_sm: u32 }; + +pub const Functions = struct { + ld: [2][ld_cfgs.len][ranges.len]?cuda.Function, // [probe 0 / 3][cfg][range] + pf: [pf_cfgs.len][ranges.len]cuda.Function, + grouped: [grouped_cfgs.len][upstream_ranges.len]cuda.Function, + dequant: [15]cuda.Function, // K2 2..16 + group: cuda.Function, + group_rot: cuda.Function, + rot_in: [2]cuda.Function, // bf16, fp16 input + gateup_epilogue: cuda.Function, + down_epilogue: cuda.Function, + combine: cuda.Function, + down_combine: [2]cuda.Function, // fp32, bf16 out + gm_gu: [gm_k2s.len][2][gu_tiles.len]Gm, // [k2][xm - 1][cfg] + gm_dn: [gm_k2s.len][dn_tiles.len]Gm, + gm_rot: [2][2]cuda.Function, // [bf16, fp16][mats - 1] + gm_dequant: [gm_k2s.len]cuda.Function, + gm_plan: cuda.Function, + gm2_gu: [2][v2_gu.len]Gm, // [xm - 1][cfg] + gm2_dn: Gm, + /// gm2_kernel at GU2 (128-member gate / up, two inputs): x3gm3.cu's instance (TF_DSV41_GM_GU2) + gm2_gu2: Gm, + gm3: [v3_variants.len]Gm, + epi: [3][epi_pds.len][ranges.len]cuda.Function, // [EpiKind][pd - 1][range] + + /// Resolves every instance (a wrong name fails here, not on a rare path) and opts x3gm into its shared memory. + pub fn resolve(experts: cuda.Module, x3ld: cuda.Module, x3pf: cuda.Module, x3gm: cuda.Module, x3gm_plan: cuda.Module, x3gm3: cuda.Module, router_glue: cuda.Module, x3ld_epi: cuda.Module, optin: u32) !Functions { + var f: Functions = undefined; + inline for (ld_cfgs, 0..) |c, ci| inline for (ranges, 0..) |r, ri| { + f.ld[0][ci][ri] = try x3ld.function(ldSymbol(c[0], c[1], r[0], r[1], 0)); + f.ld[1][ci][ri] = if (ci < 2) try x3ld.function(ldSymbol(c[0], c[1], r[0], r[1], 3)) else null; + }; + inline for (pf_cfgs, 0..) |c, ci| inline for (ranges, 0..) |r, ri| { + f.pf[ci][ri] = try x3pf.function(pfSymbol(c[0], c[1], r[0], r[1])); + }; + inline for (.{ EpiKind.gu, EpiKind.dn, EpiKind.dnb }, 0..) |kind, ki| inline for (epi_pds, 0..) |pd, pi| inline for (ranges, 0..) |r, ri| { + f.epi[ki][pi][ri] = try x3ld_epi.function(epiSymbol(kind, pd, r[0], r[1])); + }; + inline for (grouped_cfgs, 0..) |c, ci| inline for (upstream_ranges, 0..) |r, ri| { + f.grouped[ci][ri] = try experts.function(groupedSymbol(c[0], c[1], c[2], r[0], r[1])); + }; + inline for (0..15) |i| f.dequant[i] = try experts.function(dequantSymbol(i + 2)); + f.group = try experts.function(up_sym.group); + f.group_rot = try router_glue.function("router_group_rot"); + f.rot_in = .{ try experts.function(up_sym.rot_in_bf16), try experts.function(up_sym.rot_in_f16) }; + f.gateup_epilogue = try experts.function(up_sym.gateup_epilogue); + f.down_epilogue = try experts.function(up_sym.down_epilogue); + f.combine = try experts.function(up_sym.combine); + f.down_combine = .{ try experts.function(up_sym.down_combine_f32), try experts.function(up_sym.down_combine_bf16) }; + inline for (gm_k2s, 0..) |k2, ki| { + inline for (.{ 1, 2 }) |xm| inline for (0..gu_tiles.len) |c| { + f.gm_gu[ki][xm - 1][c] = try gmResolve(x3gm, gmGuSymbol(xm, k2, c), GmShape.of(2, xm, k2, gu_tiles[c]), optin); + }; + inline for (0..dn_tiles.len) |c| f.gm_dn[ki][c] = try gmResolve(x3gm, gmDnSymbol(k2, c), GmShape.of(1, 1, k2, dn_tiles[c]), optin); + f.gm_dequant[ki] = try x3gm.function(gmDequantSymbol(k2)); + } + inline for (.{ 1, 2 }) |xm| inline for (v2_gu, 0..) |c, ci| { + f.gm2_gu[xm - 1][ci] = try gmResolve(x3gm, gm2GuSymbol(xm, c), gm2Shape(2, xm, gu_tiles[c]), optin); + }; + f.gm2_dn = try gmResolve(x3gm, gm2DnSymbol(0), gm2Shape(1, 1, dn_tiles[0]), optin); + f.gm2_gu2 = try gmResolve(x3gm3, gm2GuSymbol(2, gu2_cfg), gm2Shape(2, 2, gu_tiles[gu2_cfg]), optin); + inline for (v3_variants, 0..) |v, i| f.gm3[i] = try gmResolve(x3gm3, gm3Symbol(v), gm2Shape(v.mats, v.xm, v.t), optin); + f.gm_plan = try x3gm_plan.function("_ZN15dsv41_x3gm_plan11plan_kernelEPKiiiiPiS2_S2_S2_S2_"); + f.gm_rot = .{ .{ try x3gm.function(gm_sym.rot_bf16_1), try x3gm.function(gm_sym.rot_bf16_2) }, .{ try x3gm.function(gm_sym.rot_f16_1), try x3gm.function(gm_sym.rot_f16_2) } }; + return f; + } +}; + +/// Opts `f` into `bytes` of dynamic shared memory if it fits the opt-in limit; false refuses the instance. +pub fn fitShared(f: cuda.Function, name: []const u8, bytes: u32, optin: u32) !bool { + const static: u32 = @intCast(try f.attribute(.shared_size_bytes)); + if (static + bytes > optin) { + std.log.info("dsv41 kernel refused on this GPU ({d} + {d} B shared > {d}): {s}", .{ static, bytes, optin, name }); + return false; + } + if (bytes > 0) try f.allowDynamicShared(bytes); // set whatever the size + return true; +} + +/// x3gm.cu launch<>'s first call: the shared memory attribute, then the CTAs an SM it allows (0: does not fit). +fn gmResolve(m: cuda.Module, name: [:0]const u8, shape: GmShape, optin: u32) !Gm { + const f = try m.function(name); + const ok = try fitShared(f, name, shape.smem, optin); + return .{ .f = f, .threads = shape.threads, .smem = shape.smem, .per_sm = if (ok) try f.occupancy(shape.threads, shape.smem) else 0 }; +} + +// -- Launches --------------------------------------------------------------------------------------------------- + +pub const Input = enum { bf16, f16 }; + +pub fn int(x: usize) c_int { + return @intCast(x); +} + +fn dim(x: usize) u32 { + return @intCast(x); +} + +pub fn go(f: cuda.Function, s: cuda.Stream, grid: [3]usize, block: usize, shared: usize, pdl: bool, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = .{ .x = dim(grid[0]), .y = dim(grid[1]), .z = dim(grid[2]) }, .block = .{ .x = dim(block) }, .shared = dim(shared), .pdl = pdl }, s, a); +} + +/// The arguments every grouped decode kernel (upstream's grouped_kernel, x3ld, x3pf) takes, plus the grid's sizes. +pub const Grouped = struct { + x0: u64, // fp16 [P, K] rotated inputs of matrix 0 (and of matrix 1 in x1) + x1: u64, + tp0: u64, // int64 [E]: each expert's trellis words + tp1: u64, + k2_0: u64, // int32 [E]: each expert's K2 + k2_1: u64, + uids: u64, // int32 [nexp]: distinct experts (group_kernel) + ucount: u64, // int32 [1] + members: u64, // int32 [nexp, maxm]: row * 32 + slot codes, -1 after the last + z: u64, // fp32 [mats, SK, P, N] + mats: usize, + K: usize, + N: usize, + P: usize, + SK: usize, + slots: usize, + maxm: usize, // members.size(1) + nexp: usize, // uids.size(0) + lo: u32, // the widths' range over the table + hi: u32, + + fn add(g: Grouped, a: *cuda.Args) void { + for ([_]u64{ g.x0, g.x1, g.tp0, g.tp1, g.k2_0, g.k2_1, g.uids, g.ucount, g.members, g.z }) |v| a.add(v); + for ([_]usize{ g.K, g.N, g.P, g.SK, g.maxm, g.slots }) |v| a.add(int(v)); + } +}; + +pub const Error = error{ Unsupported, Shape }; + +/// dsv41_x3ld_grouped_cuda's checks: mul1, even splits, the ring depth a divisor of a warp's k steps. +pub fn x3ldCheck(g: Grouped, nt: u32, pd: u32, probe: u32) Error!struct { cfg: usize, range: usize } { + const cfg = cfgIndex(2, &ld_cfgs, .{ nt, pd }) orelse return error.Unsupported; + if (probe != 0 and (probe != 3 or cfg >= 2)) return error.Unsupported; + const range = rangeIndex(g.lo, g.hi) orelse return error.Unsupported; + if (g.SK == 0 or g.K % (16 * g.SK * 4) != 0 or g.N % (16 * nt) != 0) return error.Shape; + const per_warp = g.K / (16 * g.SK * 4); + if (per_warp < pd or per_warp % pd != 0) return error.Shape; + return .{ .cfg = cfg, .range = range }; +} + +/// dsv41_x3pf_grouped_cuda's checks. +pub fn x3pfCheck(g: Grouped, nt: u32, mtl: u32) Error!struct { cfg: usize, range: usize } { + const cfg = cfgIndex(2, &pf_cfgs, .{ nt, mtl }) orelse return error.Unsupported; + const range = rangeIndex(g.lo, g.hi) orelse return error.Unsupported; + if (g.SK == 0 or g.K % (16 * g.SK * 4) != 0 or g.N % (16 * nt) != 0) return error.Shape; + return .{ .cfg = cfg, .range = range }; +} + +/// x3ld_epi.cu's Epi (by value): the epilogue's operands, 0 where a kind does not use one. +pub const Epi = extern struct { + pick: u64 = 0, + sv0: u64 = 0, // gate/up: svh_g; down: svh_d + sv1: u64 = 0, // gate/up: svh_u + sd: u64 = 0, // gate/up: suh_d + xd: u64 = 0, // gate/up: Xd + y: u64 = 0, // down: y + wts: u64 = 0, // down: the combine's weights + out: u64 = 0, // down: L.moe (fp32 or bf16) + ticket: u64 = 0, // int32, zero between launches + E: c_int = 0, + limit: f32 = 0, + act_mode: c_int = 0, +}; + +comptime { + std.debug.assert(@sizeOf(Epi) == 88 and @offsetOf(Epi, "E") == 72); +} + +/// The ticket words a fused launch counts on (gate/up: expert x member tile x 128 columns; down: row x 128 columns). +pub fn epiTicketWords(g: Grouped, kind: EpiKind) usize { + return if (kind == .gu) g.nexp * ((g.maxm + 15) / 16) * (g.N / 128) else (g.P / g.slots) * (g.N / 128); +} + +/// TF_DSV41_X3LD_EPI's checks: x3ld's at nt 8 and pd 1 / 2, plus the fused tails' shapes (down: one split). +pub fn x3ldEpiCheck(g: Grouped, kind: EpiKind, pd: u32) Error!struct { pd: usize, range: usize } { + const c = try x3ldCheck(g, 8, pd, 0); + const pi = std.mem.indexOfScalar(u32, &epi_pds, pd) orelse return error.Unsupported; + if (g.N % 128 != 0 or g.slots == 0 or g.slots > 32 or g.P % g.slots != 0) return error.Shape; + switch (kind) { + .gu => if (g.mats != 2) return error.Shape, + .dn, .dnb => if (g.mats != 1 or g.SK != 1) return error.Shape, + } + return .{ .pd = pi, .range = c.range }; +} + +pub const Ops = struct { + f: *const Functions, + s: cuda.Stream, + sms: usize, + + /// x3ld (expert_loads.grouped): Z of every (mat, split, live member row, column), upstream's bits. + pub fn x3ld(o: Ops, g: Grouped, nt: u32, pd: u32, probe: u32, pdl: bool) !void { + const c = try x3ldCheck(g, nt, pd, probe); + const mt = (g.maxm + 15) / 16; + var a: cuda.Args = .{}; + g.add(&a); + try go(o.f.ld[if (probe == 3) 1 else 0][c.cfg][c.range].?, o.s, .{ g.nexp, g.N / (16 * nt), g.mats * g.SK * mt }, 128, 0, pdl, &a); + } + + /// TF_DSV41_X3LD_EPI: x3ld (nt 8) with the gateup_epilogue (`gu`) or down_combine (`dn` / `dnb`) in its tail. + pub fn x3ldEpi(o: Ops, g: Grouped, kind: EpiKind, pd: u32, epi: Epi) !void { + const c = try x3ldEpiCheck(g, kind, pd); + if (epi.ticket == 0 or epi.pick == 0) return error.Shape; + const mt = (g.maxm + 15) / 16; + var a: cuda.Args = .{}; + g.add(&a); + a.add(epi); + try go(o.f.epi[@backingInt(kind)][c.pd][c.range], o.s, .{ g.nexp, g.N / 128, g.mats * g.SK * mt }, 128, 0, false, &a); + } + + /// x3pf (expert_prefill.grouped): the same Z, MTL member tiles a program. + pub fn x3pf(o: Ops, g: Grouped, nt: u32, mtl: u32) !void { + const c = try x3pfCheck(g, nt, mtl); + const mg = (g.maxm + 16 * mtl - 1) / (16 * mtl); + var a: cuda.Args = .{}; + g.add(&a); + try go(o.f.pf[c.cfg][c.range], o.s, .{ g.nexp, g.N / (16 * nt), g.mats * g.SK * mg }, 128, 0, false, &a); + } + + /// Upstream's grouped kernel (exl3x_grouped_cuda at cb 2): the reference x3ld and x3pf keep the bits of. + pub fn grouped(o: Ops, g: Grouped, nt: u32, warps: u32, pf: u32) !void { + const cfg = cfgIndex(3, &grouped_cfgs, .{ nt, warps, pf }) orelse return error.Unsupported; + if (g.SK == 0 or g.K % (16 * g.SK * warps) != 0 or g.N % (16 * nt) != 0) return error.Shape; + const mt = (g.maxm + 15) / 16; + var a: cuda.Args = .{}; + g.add(&a); + try go(o.f.grouped[cfg][upstreamRangeIndex(g.lo, g.hi)], o.s, .{ g.nexp, g.N / (16 * nt), g.mats * g.SK * mt }, warps * 32, 0, false, &a); + } + + /// exl3x_group_cuda: distinct experts of `pick` [R, slots] in id order and their members, one 1,024-thread block. + pub fn group(o: Ops, pick: u64, uids: u64, ucount: u64, members: u64, R: usize, slots: usize, E: usize, maxm: usize) !void { + if (E > 1024 * 4 or slots > 32) return error.Shape; + const smem = R * slots * 4; + if (smem > 48 * 1024) return error.Shape; // the binding never opts in to more + var a: cuda.Args = .{}; + for ([_]u64{ pick, uids, ucount, members }) |v| a.add(v); + for ([_]usize{ R, slots, E, maxm }) |v| a.add(int(v)); + try go(o.f.group, o.s, .{ 1, 1, 1 }, 1024, smem, false, &a); + } + + /// exl3x_rot_in_cuda: Xh = fp16((x * suh) H / sqrt 128) for gate and up of every routed slot. + pub fn rotIn(o: Ops, input: Input, x: u64, x_stride: usize, pick: u64, suh0: u64, suh1: u64, out0: u64, out1: u64, rows: usize, K: usize, slots: usize, E: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(int(x_stride)); + for ([_]u64{ pick, suh0, suh1, out0, out1 }) |v| a.add(v); + for ([_]usize{ K, slots, E }) |v| a.add(int(v)); + try go(o.f.rot_in[@backingInt(input)], o.s, .{ rows * slots, K / 128, 2 }, 32, 0, false, &a); + } + + /// Independent grouping/rotation CTAs share a launch, with no inter-CTA handoff. Invalid picks leave Xh untouched. + pub fn groupRotIn(o: Ops, x: u64, x_stride: usize, pick: u64, suh0: u64, suh1: u64, out0: u64, out1: u64, uids: u64, ucount: u64, members: u64, rows: usize, K: usize, slots: usize, E: usize, maxm: usize) !void { + if (rows < 1 or rows > 64 or K != 5120 or slots < 1 or slots > 9 or E < 1 or E > 512 or x_stride < K or maxm < 1) return error.Shape; + var a: cuda.Args = .{}; + a.add(x); + a.add(int(x_stride)); + for ([_]u64{ pick, suh0, suh1, out0, out1, uids, ucount, members }) |v| a.add(v); + for ([_]usize{ rows, K, slots, E, maxm }) |v| a.add(int(v)); + const tiles = rows * slots * (K / 128) * 2; + try go(o.f.group_rot, o.s, .{ 1 + (tiles + 31) / 32, 1, 1 }, 1024, 0, false, &a); + } + + /// exl3x_gateup_epilogue_cuda: splits summed, rotated, * svh, SwiGLU (act_mode 0 bf16 roundings, 1 fp32), Xd. + pub fn gateupEpilogue(o: Ops, z: u64, pick: u64, svh_g: u64, svh_u: u64, suh_d: u64, xd: u64, rows: usize, P: usize, N: usize, SK: usize, slots: usize, E: usize, limit: f32, act_mode: c_int) !void { + var a: cuda.Args = .{}; + for ([_]u64{ z, pick, svh_g, svh_u, suh_d, xd }) |v| a.add(v); + for ([_]usize{ P, N, SK, E }) |v| a.add(int(v)); + a.add(limit); + a.add(act_mode); + try go(o.f.gateup_epilogue, o.s, .{ rows * slots, N / 128, 1 }, 32, 0, false, &a); + } + + /// exl3x_down_epilogue_cuda: Y = (splits summed) H * svh_d, fp32. + pub fn downEpilogue(o: Ops, z: u64, pick: u64, svh_d: u64, y: u64, rows: usize, P: usize, D: usize, SK: usize, slots: usize, E: usize) !void { + var a: cuda.Args = .{}; + for ([_]u64{ z, pick, svh_d, y }) |v| a.add(v); + for ([_]usize{ P, D, SK, E }) |v| a.add(int(v)); + try go(o.f.down_epilogue, o.s, .{ rows * slots, D / 128, 1 }, 32, 0, false, &a); + } + + /// exl3x_combine_cuda: out[r] = fma chain over slots in order of wts[r][k] * y[r * slots + k]. + pub fn combine(o: Ops, y: u64, wts: u64, out: u64, rows: usize, D: usize, slots: usize) !void { + var a: cuda.Args = .{}; + for ([_]u64{ y, wts, out }) |v| a.add(v); + for ([_]usize{ D, slots }) |v| a.add(int(v)); + try go(o.f.combine, o.s, .{ rows, (D + 255) / 256, 1 }, 256, 0, false, &a); + } + + /// exl3x_down_combine_cuda: down_epilogue then combine in one launch, same bits; `out_bf16`: bf16 out (RNE). + pub fn downCombine(o: Ops, z: u64, pick: u64, svh_d: u64, y: u64, wts: u64, out: u64, rows: usize, P: usize, D: usize, SK: usize, slots: usize, E: usize, out_bf16: bool) !void { + if (slots > 32) return error.Shape; + var a: cuda.Args = .{}; + for ([_]u64{ z, pick, svh_d, y, wts, out }) |v| a.add(v); + for ([_]usize{ P, D, SK, E, slots }) |v| a.add(int(v)); + try go(o.f.down_combine[@intFromBool(out_bf16)], o.s, .{ rows, D / 128, 1 }, 32 * slots, 0, false, &a); + } + + /// exl3x_dequant_cuda at cb 2: W_q [K, N] fp16 through the grouped kernels' lane decode (tests). + pub fn dequant(o: Ops, t: u64, out: u64, K: usize, N: usize, k2: u32) !void { + if (k2 < 2 or k2 > 16) return error.Unsupported; + var a: cuda.Args = .{}; + a.add(t); + a.add(out); + a.add(int(K)); + a.add(int(N)); + try go(o.f.dequant[k2 - 2], o.s, .{ K / 16, N / 16, 1 }, 32, 0, false, &a); + } + + pub const gmRot = gmk.gmRot; + pub const gmGateup = gmk.gmGateup; + pub const gmDown = gmk.gmDown; + pub const gm2Gateup = gmk.gm2Gateup; + pub const gm2Down = gmk.gm2Down; + pub const gm3Gateup = gmk.gm3Gateup; + pub const gm3Down = gmk.gm3Down; + pub const gmPlan = gmk.gmPlan; + pub const gmDequant = gmk.gmDequant; + + pub fn gmLaunch(o: Ops, g: Gm, args: GmArgs, ticket: bool) !void { + if (g.per_sm == 0) return error.Unsupported; // "the configuration does not fit an SM" + var b = args; + if (!ticket) b.ticket = 0; + var a: cuda.Args = .{}; + a.add(b); + // TF_DSV41_PF_TBO_GM_SMS: persistent CTAs on at most that many SMs; which CTA runs an item changes no bit + const sms = if (gm_sms_cap > 0) @min(o.sms, gm_sms_cap) else o.sms; + try go(g.f, o.s, .{ sms * g.per_sm, 1, 1 }, g.threads, g.smem, false, &a); + } +}; + +test "ranges as dispatch_ranges and TF_RANGES pick them" { + try std.testing.expectEqual(@as(?usize, 0), rangeIndex(8, 8)); + try std.testing.expectEqual(@as(?usize, 1), rangeIndex(4, 10)); + try std.testing.expectEqual(@as(?usize, 1), rangeIndex(8, 10)); + try std.testing.expectEqual(@as(?usize, 2), rangeIndex(4, 12)); + try std.testing.expectEqual(@as(?usize, null), rangeIndex(4, 16)); + try std.testing.expectEqual(@as(?usize, null), rangeIndex(1, 4)); + try std.testing.expectEqual(@as(usize, 2), upstreamRangeIndex(4, 16)); +} + +test "x3ld checks mirror the binding's" { + const g: Grouped = .{ .x0 = 0, .x1 = 0, .tp0 = 0, .tp1 = 0, .k2_0 = 0, .k2_1 = 0, .uids = 0, .ucount = 0, .members = 0, .z = 0, .mats = 2, .K = 4096, .N = 1024, .P = 64, .SK = 1, .slots = 8, .maxm = 64, .nexp = 64, .lo = 4, .hi = 12 }; + const c = try x3ldCheck(g, 8, 2, 0); + try std.testing.expectEqual(@as(usize, 1), c.cfg); + try std.testing.expectEqual(@as(usize, 2), c.range); + try std.testing.expectError(error.Unsupported, x3ldCheck(g, 4, 2, 3)); + try std.testing.expectError(error.Unsupported, x3ldCheck(g, 4, 1, 0)); + var bad = g; + bad.K = 4096 + 16; // not 16 x SK x 4 warps + try std.testing.expectError(error.Shape, x3ldCheck(bad, 8, 1, 0)); + bad = g; + bad.K = 16 * 4 * 3; // three k steps a warp: not a multiple of pd 2 + try std.testing.expectError(error.Shape, x3ldCheck(bad, 8, 2, 0)); +} + +test "x3gm v2 shapes and tuning" { + // gate/up cfg 0, one input: K2 10's ring cut from 3 to 2 stages keeps two CTAs an SM; the widest width sets SMEM + const s = gm2Shape(2, 1, gu_tiles[0]); + try std.testing.expectEqual(@as(u32, 256), s.threads); + try std.testing.expect(2 * (@as(usize, s.smem) + 1536) <= 102400); + try std.testing.expectEqual([2]usize{ 0, 0 }, gmTuned2(null, null, true)); + try std.testing.expectEqual([2]usize{ 1, 0 }, gmTuned2(3, 2, false)); + try std.testing.expect(gm2WidthsOk(&.{ 3, 10, 6 })); + try std.testing.expect(!gm2WidthsOk(&.{ 2, 6 })); +} + +test "x3gm v2 at GU2 (TF_DSV41_GM_GU2, x3gm3.cu's instance): 512 threads, every width's ring in sm_121's opt-in memory" { + const s = gm2Shape(2, 2, gu_tiles[gu2_cfg]); + try std.testing.expectEqual(@as(u32, 512), s.threads); + try std.testing.expect(s.smem > 0 and s.smem + 4 * 8 * gu_tiles[gu2_cfg][1] + 4 <= gm_smem_max); + try std.testing.expectEqualStrings("_ZN10dsv41_x3gm10gm2_kernelILi2ELi2ELi1ELi16ELi2ELi4ELi4EEEvNS_4ArgsE", gm2GuSymbol(2, gu2_cfg)); +} + +test "x3gm shapes, fits and tuning as the reference tuner" { + // G6 gate/up cfg 0 at K2 6, one input: 3 x (64 x 64 x 2 + 2 x 4 x 8 x 24 x 4) = 49,152 + 4,608 ... (smem_bytes) + const s = GmShape.of(2, 1, 6, gu_tiles[0]); + try std.testing.expectEqual(@as(u32, 256), s.threads); + try std.testing.expectEqual(@as(u32, 3 * (64 * 64 * 2 + 2 * 4 * 8 * 24 * 4)), s.smem); + try std.testing.expect(gmFits(.gu, 0, 10, 1)); + try std.testing.expect(!gmFits(.gu, 3, 10, 2)); + try std.testing.expectEqual(@as(u32, 128), gmMembers(.dn, 2)); + // 2,048 rows x 6 slots over 384 experts (32 members): G6's (0 / 1, 0) + try std.testing.expectEqual([2]usize{ 0, 0 }, gmTuned(&tune_default, null, null, 2048, 6, 384, true, 6, 6)); + try std.testing.expectEqual([2]usize{ 1, 0 }, gmTuned(&tune_default, null, null, 2048, 6, 384, false, 6, 6)); + // 4,096 rows (64 members): the 4K entry, cfg 3 / 2 when it fits the width + try std.testing.expectEqual([2]usize{ 3, 2 }, gmTuned(&tune_default, null, null, 4096, 6, 384, true, 6, 6)); + try std.testing.expectEqual([2]usize{ 2, 2 }, gmTuned(&tune_default, null, null, 4096, 6, 384, false, 6, 6)); + try std.testing.expectEqual([2]usize{ 5, 4 }, gmTuned(&tune_default, 5, 4, 4096, 6, 384, false, 6, 6)); +} + +test "x3ld_epi checks" { + const g: Grouped = .{ .x0 = 0, .x1 = 0, .tp0 = 0, .tp1 = 0, .k2_0 = 0, .k2_1 = 0, .uids = 0, .ucount = 0, .members = 0, .z = 0, .mats = 2, .K = 5120, .N = 1152, .P = 28, .SK = 4, .slots = 7, .maxm = 4, .nexp = 28, .lo = 4, .hi = 10 }; + _ = try x3ldEpiCheck(g, .gu, 1); + try std.testing.expectError(error.Shape, x3ldEpiCheck(g, .dn, 1)); + var d = g; + d.mats = 1; + d.K = 1152; + d.N = 5120; + d.SK = 1; + _ = try x3ldEpiCheck(d, .dnb, 1); + try std.testing.expectEqual(@as(usize, 4 * 40), epiTicketWords(d, .dn)); + try std.testing.expectEqual(@as(usize, 28 * 1 * 9), epiTicketWords(g, .gu)); + try std.testing.expectEqualStrings("dsv41_x3ld_epi_dnb_2_2_10", epiSymbol(.dnb, 2, 2, 10)); +} diff --git a/zig/src/families/deepseek_v41/kernels_glue_ops.zig b/zig/src/families/deepseek_v41/kernels_glue_ops.zig new file mode 100644 index 0000000000..76d1a1f6af --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_glue_ops.zig @@ -0,0 +1,332 @@ +//! DeepSeek-V4.1 glue launches: the window's and the prefill segment's torch glue as kernels_ops.Ops methods. + +const std = @import("std"); +const cuda = @import("cuda"); +const ops = @import("kernels_ops.zig"); + +const Ops = ops.Ops; +const glue = ops.glue; +const pfglue = ops.pfglue; + +pub fn glueGo(o: Ops, kind: glue.Kind, n: u64, args: *cuda.Args) !void { + try o.go(o.f.glue[@backingInt(kind)], .{ glue.grid(n), 1, 1 }, 256, 0, null, args); +} + +/// Embedding send: rows [n, D] bf16 = embed[ids - vocab_lo] when in this rank's slice [0, V), else +0. +pub fn embedSend(o: Ops, ids: u64, n: usize, embed: u64, vocab_lo: i64, V: usize, D: usize, rows: u64) !void { + if (n == 0) return; + var a: cuda.Args = .{}; + a.add(ids); + a.add(@as(c_int, @intCast(n))); + a.add(embed); + a.add(vocab_lo); + a.add(@as(c_int, @intCast(V))); + a.add(@as(c_int, @intCast(D))); + a.add(rows); + try o.glueGo(.embed_send, @as(u64, n) * D, &a); +} + +/// Embedding sum: recv [W, n, D] bf16 -> out [n, 4 D] bf16 (fp32 adds in rank order, one rounding). +pub fn embedSum(o: Ops, recv: u64, W: usize, n: usize, D: usize, out: u64) !void { + if (n == 0) return; + if (W < 1) return error.Shape; + var a: cuda.Args = .{}; + a.add(recv); + for ([_]usize{ W, n, D }) |v| a.add(@as(c_int, @intCast(v))); + a.add(out); + try o.glueGo(.embed_sum, @as(u64, n) * D, &a); +} + +/// y bf16 = x.to(bf16) over n elements (the MoE partial into the exchange). +pub fn castBf16(o: Ops, x: u64, y: u64, n: usize) !void { + if (n == 0) return; + var a: cuda.Args = .{}; + a.add(x); + a.add(y); + a.add(@as(i64, @intCast(n))); + try o.glueGo(.cast_bf16, n, &a); +} + +/// pick.kit_weights: y = x.to(fp16).to(fp32) over n elements (y may be x). +pub fn kitWeights(o: Ops, x: u64, y: u64, n: usize) !void { + if (n == 0) return; + var a: cuda.Args = .{}; + a.add(x); + a.add(y); + a.add(@as(i64, @intCast(n))); + try o.glueGo(.kit_weights, n, &a); +} + +/// pick.kit_logits in place: x = x.to(bf16).to(fp32) over n elements. +pub fn kitLogits(o: Ops, x: u64, n: usize) !void { + if (n == 0) return; + var a: cuda.Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(n))); + try o.glueGo(.kit_logits, n, &a); +} + +/// Win.make's tensors: pos int32 [1] = start, pos64 int64 [n] = start + i, lo int32 [n] = 0. +pub fn positions(o: Ops, start: i32, n: usize, pos: u64, pos64: u64, lo: u64) !void { + var a: cuda.Args = .{}; + a.add(start); + a.add(@as(c_int, @intCast(n))); + a.add(pos); + a.add(pos64); + a.add(lo); + try o.go(o.f.glue[@backingInt(glue.Kind.positions)], .{ (@max(n, 1) + 255) / 256, 1, 1 }, 256, 0, null, &a); +} + +/// DSpark's candidates (glue.cu ds_cands): each row's best k as packed [rows][2 k] fp32 (values, then ids). +pub fn dsCandidates(o: Ops, lg: u64, ld: usize, rows: usize, V: usize, k: usize, id0: u32, out: u64) !void { + const f = o.f.ds_cands orelse return error.Unsupported; + if (k < 1 or k > glue.ds_cands_max or V < 1) return error.Shape; + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(lg); + a.add(@as(i64, @intCast(ld))); + a.add(@as(i32, @intCast(V))); + a.add(@as(i32, @intCast(k))); + a.add(@as(i32, @intCast(id0))); + a.add(out); + try o.go(f, .{ rows, 1, 1 }, 1024, 0, null, &a); +} + +/// DSpark's one-slot statics from the device pick `pick` into "w.ds.i64" / "w.ds.i32" at offsets `off`. +pub fn dsStage(o: Ops, pick: u64, nw: usize, s64: u64, s32: u64, n: usize, W: usize, ring: usize, valid: u64, noise: u32, off: [6]i64) !void { + const f = o.f.ds_stage orelse return error.Unsupported; + var a: cuda.Args = .{}; + a.add(pick); + a.add(@as(i32, @intCast(nw))); + a.add(s64); + a.add(s32); + a.add(@as(i32, @intCast(n))); + a.add(@as(i32, @intCast(W))); + a.add(@as(i32, @intCast(ring))); + a.add(@as(i64, @intCast(valid))); + a.add(@as(i32, @intCast(noise))); + for (off) |x| a.add(@as(i32, @intCast(x))); + try o.go(f, .{ 1, 1, 1 }, 256, 0, null, &a); +} + +/// A row window's slots' (accepted, bonus, next position) into `acc` [k][3] from `picks`, `ids` and `segs`. +pub fn dsAcceptRows(o: Ops, picks: u64, ids: u64, segs: u64, k: usize, acc: u64) !void { + const f = o.f.ds_accept_rows orelse return error.Unsupported; + if (k < 1 or k > 32) return error.Shape; + var a: cuda.Args = .{}; + a.add(picks); + a.add(ids); + a.add(segs); + a.add(@as(i32, @intCast(k))); + a.add(acc); + try o.go(f, .{ 1, 1, 1 }, 32, 0, null, &a); +} + +/// A several-slot pass's position-dependent statics for `g` members `mem` [g][3], from acc's bonus and position. +pub fn dsStageRows(o: Ops, acc: u64, mem: u64, g: usize, s64: u64, s32: u64, n: usize, W: usize, ring: usize, noise: u32, off: [6]i64) !void { + const f = o.f.ds_stage_rows orelse return error.Unsupported; + var a: cuda.Args = .{}; + a.add(acc); + a.add(mem); + a.add(@as(i32, @intCast(g))); + a.add(s64); + a.add(s32); + a.add(@as(i32, @intCast(n))); + a.add(@as(i32, @intCast(W))); + a.add(@as(i32, @intCast(ring))); + a.add(@as(i32, @intCast(noise))); + for (off) |x| a.add(@as(i32, @intCast(x))); + try o.go(f, .{ 1, 1, 1 }, 256, 0, null, &a); +} + +/// vsample.choose over gathered candidates `g` for `n` rows: tokens into `out` int64 [n] (-1: a nucleus row). +pub fn vsChoose(o: Ops, g: u64, W: usize, n: usize, k: usize, count: usize, seed: u64, pos0: u64, top_k: u32, temp: f64, top_p: f64, min_p: f64, out: u64) !void { + const f = o.f.vs_choose orelse return error.Unsupported; + if (W * k > glue.vs_max or n == 0) return error.Shape; + var a: cuda.Args = .{}; + a.add(g); + a.add(@as(i32, @intCast(W))); + a.add(@as(i32, @intCast(n))); + a.add(@as(i32, @intCast(k))); + a.add(@as(i32, @intCast(count))); + a.add(seed); + a.add(@as(i64, @bitCast(pos0))); // the kernel draws at pos0 + r in int64: a pos0 below 0 (a row-shifted start) wraps back + a.add(@as(i32, @intCast(top_k))); + a.add(temp); + a.add(top_p); + a.add(min_p); + a.add(out); + try o.go(f, .{ n, 1, 1 }, 256, 0, null, &a); +} + +/// bf16 [n] -> fp32 [n], exact (the DSpark pass's q before _rope). +pub fn widenF32(o: Ops, src: u64, dst: u64, n: usize) !void { + if (n == 0) return; + var a: cuda.Args = .{}; + a.add(src); + a.add(dst); + a.add(@as(i64, @intCast(n))); + try o.glueGo(.widen_f32, n, &a); +} + +pub fn pfGo(o: Ops, kind: pfglue.Kind, n: u64, args: *cuda.Args) !void { + try o.go(o.f.pfglue[@backingInt(kind)], .{ glue.grid(n), 1, 1 }, 256, 0, null, args); +} + +/// Each of R rows' `count` largest int64 keys as ascending positions, -1 padded: out int32 [R, count]. +pub fn topPositions(o: Ops, keys: u64, ld: usize, R: usize, n: usize, count: usize, out: u64) !void { + if (count < 1 or count > pfglue.max_count) return error.Shape; + if (R == 0) return; + var a: cuda.Args = .{}; + a.add(keys); + a.add(@as(i64, @intCast(ld))); + a.add(@as(c_int, @intCast(n))); + a.add(@as(c_int, @intCast(count))); + a.add(out); + try o.go(o.f.pfglue[@backingInt(pfglue.Kind.top_positions)], .{ R, 1, 1 }, 1024, 0, null, &a); +} + +/// index.candidate_keys: blocks int32 [R, nb] (row stride ld) -> out int32 [R, nb * bs]. +pub fn candKeys(o: Ops, blocks: u64, ld: usize, R: usize, nb: usize, bs: usize, out: u64) !void { + if (R * nb == 0) return; + var a: cuda.Args = .{}; + a.add(blocks); + a.add(@as(i64, @intCast(ld))); + for ([_]usize{ R, nb, bs }) |v| a.add(@as(c_int, @intCast(v))); + a.add(out); + try o.pfGo(.cand_keys, R * nb * bs, &a); +} + +/// out int32 [R] = #(sel >= 0 and sel < (pos + 1) / ratio), sel int32 [R, K], pos int64 [R]. +pub fn visibleCounts(o: Ops, sel: u64, K: usize, pos: u64, ratio: u32, R: usize, out: u64) !void { + if (R == 0) return; + if (ratio < 1) return error.Shape; + var a: cuda.Args = .{}; + a.add(sel); + a.add(@as(c_int, @intCast(K))); + a.add(pos); + a.add(@as(c_int, @intCast(ratio))); + a.add(@as(c_int, @intCast(R))); + a.add(out); + try o.go(o.f.pfglue[@backingInt(pfglue.Kind.counts)], .{ (R + 7) / 8, 1, 1 }, 256, 0, null, &a); +} + +/// Routed picks: pk int32 [n, routed] = pick[:, :routed], w6 fp32 = f32(f16(wts[:, :routed])) (shared slot after). +pub fn gmPicks(o: Ops, pick: u64, wts: u64, ld: usize, routed: usize, n: usize, pk: u64, w6: u64) !void { + if (n * routed == 0) return; + var a: cuda.Args = .{}; + a.add(pick); + a.add(wts); + for ([_]usize{ ld, routed, n }) |v| a.add(@as(c_int, @intCast(v))); + a.add(pk); + a.add(w6); + try o.pfGo(.gm_picks, n * routed, &a); +} + +/// One width's picks: out[p] = k2tab[pick[p]] == k2 ? pick[p] : E (the input of gmPlan). +pub fn widthMask(o: Ops, pick: u64, P: usize, k2tab: u64, E: usize, k2: u32, out: u64) !void { + if (P == 0) return; + var a: cuda.Args = .{}; + a.add(pick); + a.add(@as(i64, @intCast(P))); + a.add(k2tab); + a.add(@as(c_int, @intCast(E))); + a.add(@as(c_int, @intCast(k2))); + a.add(out); + try o.pfGo(.width_mask, P, &a); +} + +/// The indexer head weights' casts: y fp32 = f32(bf16(f32(x fp64))), n elements (glue step ix_w). +pub fn f64Bf16F32(o: Ops, x: u64, y: u64, n: usize) !void { + if (n == 0) return; + var a: cuda.Args = .{}; + a.add(x); + a.add(y); + a.add(@as(i64, @intCast(n))); + try o.pfGo(.f64_bf16_f32, n, &a); +} + +/// blocks._projection's fp32: out [n, hd] = f32(kv), or (gate != 0, ratio 2) [n, 2 hd] = [f32(kv) | f32(gate)]. +pub fn widenCat(o: Ops, kv: u64, gate: u64, n: usize, hd: usize, out: u64) !void { + if (n == 0) return; + var a: cuda.Args = .{}; + a.add(kv); + a.add(gate); + a.add(@as(c_int, @intCast(n))); + a.add(@as(c_int, @intCast(hd))); + a.add(out); + try o.pfGo(.widen_cat, n * hd * @as(usize, if (gate != 0) 2 else 1), &a); +} + +/// Stage rows of one tensor: dst row p % dst_ring = src row p % src_ring for p in [lo, hi). +pub fn ringCopy(o: Ops, src: u64, src_ring: usize, dst: u64, dst_ring: usize, row_bytes: usize, lo: u64, hi: u64) !void { + if (hi <= lo) return; + var a: cuda.Args = .{}; + a.add(src); + a.add(@as(c_int, @intCast(src_ring))); + a.add(dst); + a.add(@as(c_int, @intCast(dst_ring))); + a.add(@as(i64, @intCast(row_bytes))); + a.add(@as(i64, @intCast(lo))); + a.add(@as(i64, @intCast(hi))); + try o.pfGo(.ring_copy, (hi - lo) * (row_bytes / 4 + 1), &a); +} + +/// positions with `start` read from the device (int32 at `start_dev`): a graphed window's form, the same values. +pub fn positionsDev(o: Ops, start_dev: u64, n: usize, pos: u64, pos64: u64, lo: u64) !void { + var a: cuda.Args = .{}; + a.add(start_dev); + a.add(@as(c_int, @intCast(n))); + a.add(pos); + a.add(pos64); + a.add(lo); + try o.go(o.f.glue[@backingInt(glue.Kind.positions_dev)], .{ (@max(n, 1) + 255) / 256, 1, 1 }, 256, 0, null, &a); +} + +/// forward.gather_cols after the all-gather: recv [W, n, c] -> out [n, W c] (2-byte elements, rank order). +pub fn gatherCols(o: Ops, recv: u64, W: usize, n: usize, c: usize, out: u64) !void { + if (n == 0 or c == 0) return; + var a: cuda.Args = .{}; + a.add(recv); + for ([_]usize{ W, n, c }) |v| a.add(@as(c_int, @intCast(v))); + a.add(out); + try o.glueGo(.gather_cols, @as(u64, W) * n * c, &a); +} + +/// glue.cu pick_pack: f64 (value, column + id0) pairs [n][2] from topk_keys' k-1 vals fp32 [n] / cols int64 [n]. +pub fn pickPack(o: Ops, vals: u64, cols: u64, n: usize, id0: u64, pairs: u64) !void { + const f = o.f.pick_pack orelse return error.NoPickKernel; + var a: cuda.Args = .{}; + a.add(vals); + a.add(cols); + a.add(@as(c_int, @intCast(n))); + a.add(@as(i64, @intCast(id0))); + a.add(pairs); + try o.go(f, .{ glue.grid(n), 1, 1 }, 256, 0, null, &a); +} + +/// glue.cu pick_merge: out int64 [n + 3] = each row's pick over the ranks, then accepted rows, bonus, next start. +pub fn pickMerge(o: Ops, all: u64, W: usize, n: usize, ids: u64, start: i64, out: u64) !void { + const f = o.f.pick_merge orelse return error.NoPickKernel; + if (n == 0 or n > glue.pick_max_rows) return error.Shape; + var a: cuda.Args = .{}; + a.add(all); + a.add(@as(c_int, @intCast(W))); + a.add(@as(c_int, @intCast(n))); + a.add(ids); + a.add(start); + a.add(out); + try o.go(f, .{ 1, 1, 1 }, @intCast(@max(32, std.mem.alignForward(usize, n, 32))), 0, null, &a); +} + +/// The compressor's carry: carry fp32 [cols] = bf16 row of src; row *accepted (device int32) if `accepted` != 0. +pub fn carry(o: Ops, src: u64, ld: usize, accepted: u64, row: usize, cols: usize, out: u64) !void { + var a: cuda.Args = .{}; + a.add(src); + a.add(@as(i64, @intCast(ld))); + a.add(accepted); + a.add(@as(c_int, @intCast(row))); + a.add(@as(c_int, @intCast(cols))); + a.add(out); + try o.glueGo(.carry, cols, &a); +} diff --git a/zig/src/families/deepseek_v41/kernels_kv_glue.zig b/zig/src/families/deepseek_v41/kernels_kv_glue.zig new file mode 100644 index 0000000000..47b71369f8 --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_kv_glue.zig @@ -0,0 +1,46 @@ +//! KV-path RMS + SWA RoPE/FP8 store. This owns only the norm feeding the SWA producer. +const std = @import("std"); +const cuda = @import("cuda"); + +/// Mirrors kv_glue.cu; norm_out is an optional contiguous bf16 debug output, not used in production. +pub const Args = extern struct { + x: u64 = 0, + w: u64 = 0, + cs: u64 = 0, + v: u64 = 0, + s: u64 = 0, + pos: u64 = 0, + sl: u64 = 0, + norm_out: u64 = 0, + x_stride: i64 = 512, + cs_stride: i64 = 64, + v_stride: i64 = 576, + s_stride: i64 = 8, + inv_k: f32 = 1.0 / 512.0, + eps: f32 = 1e-20, + ring: c_int = 256, + rows: c_int = 0, +}; +comptime { + std.debug.assert(@sizeOf(Args) == 112 and @offsetOf(Args, "inv_k") == 96 and @offsetOf(Args, "ring") == 104); +} + +pub const Functions = struct { + norm_store: cuda.Function, + pub fn resolve(m: cuda.Module) !Functions { + return .{ .norm_store = try m.function("kv_norm_store") }; + } +}; + +pub const Ops = struct { + f: *const Functions, + stream: cuda.Stream, + pub fn normStore(o: Ops, a: Args, n: usize) !void { + if (n == 0 or n > 64 or a.x_stride < 512 or a.cs_stride < 64 or a.v_stride < 576 or a.s_stride < 8 or + a.ring <= 0 or (a.ring & (a.ring - 1)) != 0 or a.inv_k != 1.0 / 512.0 or + (a.rows != 0 and a.rows != 1) or (a.rows == 1 and a.sl == 0)) return error.Shape; + var args: cuda.Args = .{}; + args.add(a); + try cuda.launch.launch(o.f.norm_store, .{ .grid = .{ .x = @intCast(n) }, .block = .{ .x = 128 } }, o.stream, &args); + } +}; diff --git a/zig/src/families/deepseek_v41/kernels_lin2.zig b/zig/src/families/deepseek_v41/kernels_lin2.zig new file mode 100644 index 0000000000..db68e18af7 --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_lin2.zig @@ -0,0 +1,106 @@ +//! TF_DSV41_LIN2 (lin2.cu): the dense linears' one-pass 17-32-row variants, their symbols, mode readers and functions. + +const std = @import("std"); +const cuda = @import("cuda"); +const dense = @import("kernels_dense.zig"); + +const LinTable = dense.LinTable; +const dense3_k2s = dense.dense3_k2s; +const linear_k2s = dense.linear_k2s; + +// TF_DSV41_LIN2 (ours, lin2.cu): the 17-32-row programs of the dense linears in one pass over the weights +fn lanes2Symbol(comptime k2: u32, comptime wk: u32, comptime mb: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_lin213lanes2_kernelILi{d}ELi2ELi{d}ELi{d}EEEvN11dsv41_x3seg8LinTableEii", .{ k2, wk, mb }); +} +fn lanescSymbol(comptime k2: u32, comptime wk: u32, comptime mb: u32, comptime t: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_lin213lanesc_kernelILi{d}ELi2ELi{d}ELi{d}ELi{d}EEEvN11dsv41_x3seg8LinTableEii", .{ k2, wk, mb, t }); +} +/// lanesc's instances in lin2.cu's order (WK, MINB, T): a 4-warp and an 8-warp CTA at T 1 and at T 2 +pub const lanesc_cfgs = [_][3]u32{ .{ 4, 3, 1 }, .{ 8, 1, 1 }, .{ 4, 2, 2 }, .{ 8, 1, 2 } }; +fn linear2Symbol(comptime k2: u32, comptime wk: u32, comptime pf: bool) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_lin214linear2_kernelILi{d}ELi2ELi{d}ELb{d}EEEvPK6__halfPKjxxS3_S3_PviPfPiiiiiPKi", .{ k2, wk, @intFromBool(pf) }); +} + +/// TF_DSV41_LIN2[_PF|_WK|_CLUSTER]: lin2.cu's one-pass 17-32-row variants of the dense linears; bad values refused. +pub const Lin2Mode = struct { lanes: bool = false, linear: bool = false, pf: bool = true, wk4: bool = true, wk8: bool = true, cluster: bool = false }; + +/// TF_DSV41_LIN2_CLUSTER's value: unset / empty / 0 off, 1 on, anything else refused. +pub fn lin2Cluster(v: ?[]const u8) !bool { + const t = std.mem.trim(u8, v orelse "", " \t\r\n"); + if (t.len == 0 or std.mem.eql(u8, t, "0")) return false; + if (std.mem.eql(u8, t, "1")) return true; + return error.BadLin2; +} + +/// The split count (2-8) shared by every segment of a dense3 launch of at most 32 rows, else null. +pub fn clusterSplits(tab: LinTable, m: usize) ?u32 { + if (m < 1 or m > 32 or tab.n < 1) return null; + const sk = tab.s[0].SK; + if (sk < 2 or sk > 8) return null; + for (tab.s[0..@intCast(tab.n)]) |x| if (x.SK != sk or @mod(x.p0, sk) != 0) return null; + return @intCast(sk); +} + +pub fn lin2Mode(v: ?[]const u8, pf: ?[]const u8) !Lin2Mode { + return lin2ModeWk(v, pf, null); +} + +pub fn lin2ModeWk(v: ?[]const u8, pf: ?[]const u8, wk: ?[]const u8) !Lin2Mode { + var m: Lin2Mode = .{}; + const t = std.mem.trim(u8, v orelse "", " \t\r\n"); + if (t.len == 0 or std.mem.eql(u8, t, "0") or std.ascii.eqlIgnoreCase(t, "off")) { + m = .{}; + } else if (std.mem.eql(u8, t, "1") or std.ascii.eqlIgnoreCase(t, "on")) { + m = .{ .lanes = true, .linear = true }; + } else if (std.ascii.eqlIgnoreCase(t, "lanes")) { + m = .{ .lanes = true }; + } else if (std.ascii.eqlIgnoreCase(t, "linear")) { + m = .{ .linear = true }; + } else return error.BadLin2; + const p = std.mem.trim(u8, pf orelse "", " \t\r\n"); + if (p.len == 0 or std.mem.eql(u8, p, "1")) { + m.pf = true; + } else if (std.mem.eql(u8, p, "0")) { + m.pf = false; + } else return error.BadLin2; + const w = std.mem.trim(u8, wk orelse "", " \t\r\n"); + if (w.len == 0) { + m.wk4 = true; + m.wk8 = true; + } else if (std.mem.eql(u8, w, "4")) { + m.wk8 = false; + } else if (std.mem.eql(u8, w, "8")) { + m.wk4 = false; + } else return error.BadLin2; + return m; +} + +/// A launch of `m` rows the two-tile variant takes: the original's two passes (rows 17-32). +pub fn twoTile(m: usize) bool { + return m > 16 and m <= 32; +} + +/// lin2.cu's instances: lanes2 at dense3's widths (4- and 8-warp CTAs); linear2 with or without k-step prefetch. +pub const Lin2 = struct { + module: cuda.Module, + mode: Lin2Mode, + lanes: [dense3_k2s.len][2]cuda.Function, // WK 4, 8 + linear: [linear_k2s.len][3]cuda.Function, // WK 2, 4, 8 + /// TF_DSV41_LIN2_CLUSTER: lanesc at lanesc_cfgs (resolved only with the knob: an older lin2 fatbin still loads) + lanesc: [dense3_k2s.len][lanesc_cfgs.len]cuda.Function = undefined, + + pub fn resolve(module: cuda.Module, mode: Lin2Mode) !Lin2 { + var l: Lin2 = .{ .module = module, .mode = mode, .lanes = undefined, .linear = undefined }; + inline for (dense3_k2s, 0..) |k2, i| { + l.lanes[i][0] = try module.function(lanes2Symbol(k2, 4, 2)); + l.lanes[i][1] = try module.function(lanes2Symbol(k2, 8, 1)); + } + inline for (linear_k2s, 0..) |k2, i| inline for (.{ 2, 4, 8 }, 0..) |wk, j| { + l.linear[i][j] = try module.function(if (mode.pf) linear2Symbol(k2, wk, true) else linear2Symbol(k2, wk, false)); + }; + if (mode.cluster) inline for (dense3_k2s, 0..) |k2, i| inline for (lanesc_cfgs, 0..) |c, j| { + l.lanesc[i][j] = try module.function(lanescSymbol(k2, c[0], c[1], c[2])); + }; + return l; + } +}; diff --git a/zig/src/families/deepseek_v41/kernels_ops.zig b/zig/src/families/deepseek_v41/kernels_ops.zig new file mode 100644 index 0000000000..b2dff0b03a --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_ops.zig @@ -0,0 +1,590 @@ +//! DeepSeek-V4.1's other CUDA kernels (attention, top-k, mHC, router, prefill GEMM, glue): args and launches. + +const std = @import("std"); +const cuda = @import("cuda"); +const shapes = @import("kernels_shapes.zig"); +const glue_ops = @import("kernels_glue_ops.zig"); +const fitShared = @import("kernels_exl3.zig").fitShared; + +fn dim(x: usize) u32 { + return @intCast(x); +} + +pub const Error = shapes.Error; +pub const attn = shapes.attn; +pub const topk = shapes.topk; +pub const mhc = shapes.mhc; +pub const mhc_pf = shapes.mhc_pf; +pub const rg = shapes.rg; +pub const pfd = shapes.pfd; + +// -- Loaded functions and launches ------------------------------------------------------------------------------ + +pub const Functions = struct { + attn: [8]cuda.Function, + attn_split: [2][8]cuda.Function, // [P 1, P 4][instance] + topk: cuda.Function, + /// topk_b.cu's row-bounded twin (tf_dsv41_topk_b_v1.topk) + topk_b: cuda.Function, + mhc: [2][mhc.instances.len]cuda.Function, // [tail][instance] + mhc_coef: cuda.Function, + mhc_pf: [mhc_pf.instances.len]cuda.Function, + gemv: [rg.rts.len][rg.wide.len]cuda.Function, + narrow: [rg.rts.len][rg.narrow.len]cuda.Function, + pfd: [pfd.k2s.len][2][pfd.cfgs.len]?cuda.Function, // [k2][lanes][cfg]; K2 16 has no lanes kernels + pfd_dequant: [pfd.k2s.len][2]?cuda.Function, + gate_copy: cuda.Function, + paced: cuda.Function, + segments: [2]cuda.Function, // bulk false / true + topk_keys: cuda.Function, + swiglu: cuda.Function, + add: cuda.Function, + add_bf16: cuda.Function, + kv_dense: cuda.Function, + /// kvsplit's compact pack (TF_DSV41_KV_SPLIT_COMPACT); null in a fatbin built before it + kv_pack: ?cuda.Function = null, + /// kvsplit's dedup steps (TF_DSV41_KV_SPLIT_DEDUP: mark, flags, pack, tok); null in a fatbin built before them + kv_dedup: ?[4]cuda.Function = null, + kv_hist: ?[2]cuda.Function = null, + kv_gather: cuda.Function, + glue: [10]cuda.Function, // glue.syms order + /// glue.cu's DSpark candidates (TF_DSV41_DRAFT_GRAPHS); null in a fatbin built before it (the host's candidates) + ds_cands: ?cuda.Function = null, + /// glue.cu's device pick (TF_DSV41_GREEDY_GPU); null in a fatbin built before it (the pairs merged on the host) + pick_pack: ?cuda.Function = null, + pick_merge: ?cuda.Function = null, + /// glue.cu's DSpark statics from the device pick (the speculative pass's device start); optional, as ds_cands + ds_stage: ?cuda.Function = null, + /// glue.cu's row-window accept and several-slot statics (the slot passes' device start); optional + ds_accept_rows: ?cuda.Function = null, + ds_stage_rows: ?cuda.Function = null, + /// glue.cu's device keyed choice (vsample.choose; the sampled rounds' device start); optional + vs_choose: ?cuda.Function = null, + pfglue: [pfglue.syms.len]cuda.Function, // pfglue.syms order + + /// `pf`: the four pfdense modules (K2 8, 10, 12, 16). + pub fn resolve(m_attn: cuda.Module, m_topk: cuda.Module, m_mhc: cuda.Module, m_mhc_pf: cuda.Module, m_rg: cuda.Module, pf: [4]cuda.Module, m_gate: cuda.Module, m_pace: cuda.Module, m_l2pf: cuda.Module, m_keys: cuda.Module, m_pw: cuda.Module, m_kv: cuda.Module, m_glue: cuda.Module, m_pfglue: cuda.Module, m_topk_b: cuda.Module, optin: u32) !Functions { + var f: Functions = undefined; + inline for (0..8) |i| f.attn[i] = try m_attn.function(attn.symbol(if (i >= 4) 4 else 2, (i & 2) != 0, (i & 1) != 0)); + inline for (f.attn, 0..) |fun, i| { + if (!try fitShared(fun, attn.symbol(if (i >= 4) 4 else 2, (i & 2) != 0, (i & 1) != 0), attn.SMEM, optin)) return error.SharedMemoryTooSmall; + } + inline for (.{ 1, 4 }, 0..) |p, pi| inline for (0..8) |i| { + const name = attn.splitSymbol(if (i >= 4) 4 else 2, (i & 2) != 0, (i & 1) != 0, p); + f.attn_split[pi][i] = try m_attn.function(name); + if (!try fitShared(f.attn_split[pi][i], name, attn.splitSmem(p), optin)) return error.SharedMemoryTooSmall; + }; + f.topk = try m_topk.function(topk.symbol); + if (!try fitShared(f.topk, topk.symbol, topk.smem(topk.MAXEPT), optin)) return error.SharedMemoryTooSmall; + f.topk_b = try m_topk_b.function(topk.symbol_b); + if (!try fitShared(f.topk_b, topk.symbol_b, topk.smem(topk.MAXEPT), optin)) return error.SharedMemoryTooSmall; + inline for (.{ false, true }, 0..) |tail, ti| inline for (0..mhc.instances.len) |i| { + f.mhc[ti][i] = try m_mhc.function(mhc.symbol(i, tail)); + }; + f.mhc_coef = try m_mhc.function(mhc.coef_symbol); + inline for (0..mhc_pf.instances.len) |i| { + f.mhc_pf[i] = try m_mhc_pf.function(mhc_pf.symbol(i)); + if (!try fitShared(f.mhc_pf[i], mhc_pf.symbol(i), mhc_pf.smem(mhc_pf.instances[i].fn32), optin)) return error.SharedMemoryTooSmall; + } + inline for (rg.rts, 0..) |rt, i| { + inline for (rg.wide, 0..) |c, j| f.gemv[i][j] = try m_rg.function(rg.gemvSymbol(rt, c[0], c[1])); + inline for (rg.narrow, 0..) |nw, j| f.narrow[i][j] = try m_rg.function(rg.narrowSymbol(rt, nw)); + } + inline for (pfd.k2s, 0..) |k2, ki| inline for (.{ false, true }, 0..) |lanes, li| { + inline for (0..pfd.cfgs.len) |c| { + if (lanes and k2 == 16) { + f.pfd[ki][li][c] = null; + } else { + const fun = try pf[ki].function(pfd.symbol(k2, lanes, c)); + f.pfd[ki][li][c] = if (try fitShared(fun, pfd.symbol(k2, lanes, c), pfd.smem(c, k2), optin)) fun else null; + } + } + f.pfd_dequant[ki][li] = if (lanes and k2 == 16) null else try pf[ki].function(pfd.dequantSymbol(k2, lanes)); + }; + f.gate_copy = try m_gate.function("_ZN17dsv41_engram_gate16gate_copy_kernelEPxiiPKhxPhxxi"); + f.paced = try m_pace.function("_ZN12dsv41_l2pace12paced_kernelEPKlijmm"); + f.segments = .{ try m_l2pf.function("_ZN6tfl2pf15segments_kernelILb0EEEvPKlij"), try m_l2pf.function("_ZN6tfl2pf15segments_kernelILb1EEEvPKlij") }; + f.topk_keys = try m_keys.function("_ZN15dsv41_topk_keys11topk_kernelEPKfxiiPfPx"); + f.swiglu = try m_pw.function("_ZN15dsv41_pointwise13swiglu_kernelEPKfxS1_xPfxif"); + f.add = try m_pw.function("_ZN15dsv41_pointwise10add_kernelEPfPKfx"); + f.add_bf16 = try m_pw.function("_ZN15dsv41_pointwise15add_bf16_kernelEPKfS1_P13__nv_bfloat16x"); + f.kv_dense = try m_kv.function("_ZN13dsv41_kvsplit12dense_kernelENS_9DenseArgsE"); + f.kv_gather = try m_kv.function("_ZN13dsv41_kvsplit13gather_kernelEPKhjPKjjPh"); + f.kv_pack = m_kv.function("_ZN13dsv41_kvsplit11pack_kernelENS_8PackArgsE") catch null; + f.kv_dedup = kvsplit.dedupFns(m_kv) catch null; + f.kv_hist = kvsplit.histFns(m_kv) catch null; + inline for (glue.syms, 0..) |name, i| f.glue[i] = try m_glue.function(name); + f.ds_cands = m_glue.function(glue.ds_cands_sym) catch null; + f.pick_pack = m_glue.function(glue.pick_pack_sym) catch null; + f.pick_merge = m_glue.function(glue.pick_merge_sym) catch null; + f.ds_stage = m_glue.function(glue.ds_stage_sym) catch null; + f.ds_accept_rows = m_glue.function(glue.ds_accept_rows_sym) catch null; + f.ds_stage_rows = m_glue.function(glue.ds_stage_rows_sym) catch null; + f.vs_choose = m_glue.function(glue.vs_choose_sym) catch null; + inline for (pfglue.syms, 0..) |name, i| f.pfglue[i] = try m_pfglue.function(name); + return f; + } +}; + +/// kvsplit.cu: split KV's exchange copies (kv/split.zig's DenseArgs, the same layout; HostKernels the reference). +pub const kvsplit = struct { + pub const DenseArgs = extern struct { + sel: u64, + rows: u32, + k: u32, + table: u64, + pts: u32, + rslot: u64, + psh: u32, + base: u64, + row_bytes: u32, + world: u32, + send: u64, + tok: u64, + }; + + /// kv/split.zig's PackArgs (DenseArgs' fields at the same offsets, rank in psh's padding, lens after tok). + pub const PackArgs = extern struct { + sel: u64, + rows: u32, + k: u32, + table: u64, + pts: u32, + rslot: u64, + psh: u32, + rank: u32, + base: u64, + row_bytes: u32, + world: u32, + send: u64, + tok: u64, + lens: u64, + }; + + /// kv/split.zig's DedupArgs: PackArgs, then the map and the scratch. + pub const DedupArgs = extern struct { + p: PackArgs, + map: u64, + first: u64, + place: u64, + bcnt: u64, + }; + + const dedup_syms = [4][:0]const u8{ + "_ZN13dsv41_kvsplit17dedup_mark_kernelENS_9DedupArgsE", + "_ZN13dsv41_kvsplit18dedup_flags_kernelENS_9DedupArgsE", + "_ZN13dsv41_kvsplit17dedup_pack_kernelENS_9DedupArgsE", + "_ZN13dsv41_kvsplit16dedup_tok_kernelENS_9DedupArgsE", + }; + + fn dedupFns(m: cuda.Module) ![4]cuda.Function { + var out: [4]cuda.Function = undefined; + for (dedup_syms, &out) |name, *f| f.* = try m.function(name); + return out; + } + + /// kv/split.zig's HistArgs (the history union's gather and tokens). + pub const HistArgs = extern struct { + sel: u64, + rows: u32, + k: u32, + table: u64, + pages: u32, + psh: u32, + rank: u32, + world: u32, + m: u32, + row_bytes: u32, + base: u64, + send: u64, + tok: u64, + }; + + const hist_syms = [2][:0]const u8{ "_ZN13dsv41_kvsplit18hist_gather_kernelENS_8HistArgsE", "_ZN13dsv41_kvsplit15hist_tok_kernelENS_8HistArgsE" }; + + fn histFns(m: cuda.Module) ![2]cuda.Function { + return .{ try m.function(hist_syms[0]), try m.function(hist_syms[1]) }; + } + + comptime { + std.debug.assert(@sizeOf(DedupArgs) == 120 and @offsetOf(DedupArgs, "map") == 88 and @offsetOf(DedupArgs, "bcnt") == 112); + std.debug.assert(@sizeOf(HistArgs) == 72 and @offsetOf(HistArgs, "m") == 40 and @offsetOf(HistArgs, "base") == 48); + std.debug.assert(@sizeOf(DenseArgs) == 80 and @offsetOf(DenseArgs, "rslot") == 32 and @offsetOf(DenseArgs, "send") == 64); + std.debug.assert(@sizeOf(PackArgs) == 88 and @offsetOf(PackArgs, "rank") == 44 and @offsetOf(PackArgs, "lens") == 80); + } + + fn grid(words: u64) usize { + return @intCast(@max(1, @min((words + 255) / 256, 4096))); + } +}; + +/// glue.cu: the window's torch glue, torch's arithmetic bit for bit. +pub const glue = struct { + pub const syms = [_][:0]const u8{ + "_ZN10dsv41_glue17embed_send_kernelEPKxiPKtxiiPt", + "_ZN10dsv41_glue16embed_sum_kernelEPKtiiiPt", + "_ZN10dsv41_glue16cast_bf16_kernelEPKfPtx", + "_ZN10dsv41_glue18kit_weights_kernelEPKfPfx", + "_ZN10dsv41_glue17kit_logits_kernelEPfx", + "_ZN10dsv41_glue16positions_kernelEiiPiPxS0_", + "_ZN10dsv41_glue18gather_cols_kernelEPKtiiiPt", + "_ZN10dsv41_glue12carry_kernelEPKtxPKiiiPf", + "_ZN10dsv41_glue20positions_dev_kernelEPKiiPiPxS2_", + "_ZN10dsv41_glue16widen_f32_kernelEPKtPfx", + }; + pub const Kind = enum(u4) { embed_send, embed_sum, cast_bf16, kit_weights, kit_logits, positions, gather_cols, carry, positions_dev, widen_f32 }; + /// DSpark's candidates (optional: resolved apart, so an older fatbin still loads) + pub const ds_cands_sym: [:0]const u8 = "_ZN10dsv41_glue15ds_cands_kernelEPKfxiiiPf"; + /// the most candidates a rank (the kernel's shared sort) + pub const ds_cands_max = 128; + /// DSpark's statics from the device pick (optional, as ds_cands) + pub const ds_stage_sym: [:0]const u8 = "_ZN10dsv41_glue15ds_stage_kernelEPKxiPxPiiiixiiiiiii"; + pub const ds_accept_rows_sym: [:0]const u8 = "_ZN10dsv41_glue21ds_accept_rows_kernelEPKxS1_S1_iPx"; + pub const ds_stage_rows_sym: [:0]const u8 = "_ZN10dsv41_glue20ds_stage_rows_kernelEPKxS1_iPxPiiiiiiiiiii"; + pub const vs_choose_sym: [:0]const u8 = "_ZN10dsv41_glue16vs_choose_kernelEPKfiiiiyxidddPx"; + /// the candidates vs_choose merges at most (W k) + pub const vs_max = 2048; + /// the device pick (optional, as ds_cands) + pub const pick_pack_sym: [:0]const u8 = "_ZN10dsv41_glue16pick_pack_kernelEPKfPKxixPd"; + pub const pick_merge_sym: [:0]const u8 = "_ZN10dsv41_glue17pick_merge_kernelEPKdiiPKxxPx"; + /// the most rows pick_merge takes (one block) + pub const pick_max_rows = 1024; + + pub fn grid(n: u64) usize { + return @intCast(@max(1, @min((n + 255) / 256, 4096))); + } +}; + +/// prefill_glue.cu: the prefill segment's torch glue (block_prefill.zig's glue steps), torch's integers and casts. +pub const pfglue = struct { + pub const syms = [_][:0]const u8{ + "_ZN12dsv41_pfglue20top_positions_kernelEPKxxiiPi", + "_ZN12dsv41_pfglue16cand_keys_kernelEPKixiiiPi", + "_ZN12dsv41_pfglue13counts_kernelEPKiiPKxiiPi", + "_ZN12dsv41_pfglue15gm_picks_kernelEPKiPKfiiiPiPf", + "_ZN12dsv41_pfglue17width_mask_kernelEPKixS1_iiPi", + "_ZN12dsv41_pfglue19f64_bf16_f32_kernelEPKdPfx", + "_ZN12dsv41_pfglue16widen_cat_kernelEPKtS1_iiPf", + "_ZN12dsv41_pfglue16ring_copy_kernelEPKhiPhixxx", + }; + pub const Kind = enum(u3) { top_positions, cand_keys, counts, gm_picks, width_mask, f64_bf16_f32, widen_cat, ring_copy }; + /// index.top_positions' largest `count` (TOPK 512, TOPK_BLOCKS 2,048) + pub const max_count = 2048; +}; + +/// topk_keys.cu: pick.top's bounds (C columns a row within pick.COLS, k within the shared sort). +pub const topk_keys = struct { + pub const max_cols = 1 << 16; + pub const max_k = 1024; +}; + +pub const Ops = struct { + f: *const Functions, + s: cuda.Stream, + + pub fn go(o: Ops, f: cuda.Function, grid: [3]usize, block: usize, shared: u32, cluster: ?cuda.Dim3, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = .{ .x = dim(grid[0]), .y = dim(grid[1]), .z = dim(grid[2]) }, .block = .{ .x = dim(block) }, .shared = shared, .cluster = cluster }, o.s, a); + } + + /// dsv41_attn_run_cuda: `rows` slots table given, `has_hi` window ends given, `kw` Triton k width, `split` R1 CTAs. + pub fn attention(o: Ops, a: attn.Args, kw: u32, rows: bool, has_hi: bool, split: u32) !void { + const i = try attn.instance(a, kw, rows, has_hi); + if (a.qrope != 0 and split == 0) return error.Unsupported; // "q's RoPE folds into the split kernel only" + var args: cuda.Args = .{}; + args.add(a); + const hb: usize = @intCast(@divTrunc(a.H, attn.BMQ)); + const R: usize = @intCast(a.R); + switch (split) { + 0 => try o.go(o.f.attn[i], .{ attn.NCH, hb, R }, attn.NT, attn.SMEM, null, &args), + 1 => try o.go(o.f.attn_split[0][i], .{ attn.NCH, hb, R }, attn.NT, attn.splitSmem(1), null, &args), + 4 => try o.go(o.f.attn_split[1][i], .{ attn.NCH * 4, hb, R }, attn.NT, attn.splitSmem(4), .{ .x = 4 }, &args), + else => return error.Unsupported, + } + } + + /// dsv41_topk_run_cuda: `jobs` 1-2 selections, R rows, CL CTAs a cluster, ept entries a thread. + pub fn topK(o: Ops, a: topk.Args, jobs: usize, R: usize, CL: usize, ept: usize) !void { + return o.topKOf(o.f.topk, a, jobs, R, CL, ept); + } + + /// topk_b.cu (TF_DSV41_INDEX_BOUND): topK's launch of the row-bounded twin (the same args and bytes). + pub fn topKB(o: Ops, a: topk.Args, jobs: usize, R: usize, CL: usize, ept: usize) !void { + return o.topKOf(o.f.topk_b, a, jobs, R, CL, ept); + } + + fn topKOf(o: Ops, f: cuda.Function, a: topk.Args, jobs: usize, R: usize, CL: usize, ept: usize) !void { + if (R < 1 or jobs < 1 or jobs > 2 or CL < 1 or CL > topk.MAXCL or ept < 1 or ept > topk.MAXEPT) return error.Shape; + var args: cuda.Args = .{}; + args.add(a); + try o.go(f, .{ CL, R, jobs }, topk.NT, topk.smem(ept), .{ .x = dim(CL) }, &args); + } + + /// dsv41_mhc_run_cuda (`fn32`: fn is fp32); `a.spin` >= 0 R1's parallel tail, `a.defer` 1 leaves coefs to mhcCoef. + pub fn mhcBoundary(o: Ops, a: mhc.Args, mode: u32, fn32: bool) !void { + const i = try mhc.instance(a, mode, fn32); + var b = a; + if (b.spin < 0) b.@"defer" = 0; // as the binding: defer only with the tail + var args: cuda.Args = .{}; + args.add(b); + try o.go(o.f.mhc[@intFromBool(b.spin >= 0)][i], .{ mhc.NB, if (mhc.instances[i].all4) 1 else 4, 1 }, mhc.NT, @intCast(@as(u32, @intCast(a.R)) * mhc.ROWB), null, &args); + } + + /// dsv41_mhc_coef_cuda (R1): the next site's coefficients of a deferred boundary, a CTA a row. + pub fn mhcCoef(o: Ops, a: mhc.Args) !void { + if (a.R < 1 or a.R > mhc.COEF_ROWS) return error.Shape; + var args: cuda.Args = .{}; + args.add(a); + try o.go(o.f.mhc_coef, .{ @intCast(a.R), 1, 1 }, mhc.NT, 0, null, &args); + } + + /// dsv41_mhc_pf_run_cuda: `tap` = a tap buffer is given. + pub fn mhcSite(o: Ops, a: mhc_pf.Args, mode: u32, fn32: bool, tap: bool) !void { + const i = try mhc_pf.instance(a, mode, fn32, tap); + var args: cuda.Args = .{}; + args.add(a); + const R: usize = @intCast(a.R); + try o.go(o.f.mhc_pf[i], .{ mhc_pf.NB, (R + mhc_pf.BM - 1) / mhc_pf.BM, 1 }, mhc_pf.NT, mhc_pf.smem(mhc_pf.instances[i].fn32), null, &args); + } + + /// dsv41_rg_route_cuda: (rt, nw, ew) a wide instance, or ew 0 the narrow kernel (default carveout). + pub fn route(o: Ops, a: rg.Args, rt: u32, nw: u32, ew: u32, lg_len: usize, cnt_len: usize) !void { + try rg.check(a, rt, nw, ew, lg_len, cnt_len); + if (a.R == 0) return; + const ri = std.mem.indexOfScalar(u32, &rg.rts, rt).?; + const E: usize = @intCast(a.E); + const R: usize = @intCast(a.R); + var args: cuda.Args = .{}; + args.add(a); + if (ew == 0) { + const ni = std.mem.indexOfScalar(u32, &rg.narrow, nw).?; + return o.go(o.f.narrow[ri][ni], .{ (E + nw - 1) / nw, (R + rt - 1) / rt, 1 }, nw * 32, 0, null, &args); + } + var wi: usize = 0; + for (rg.wide, 0..) |c, j| if (c[0] == nw and c[1] == ew) { + wi = j; + }; + try o.go(o.f.gemv[ri][wi], .{ (E + nw * ew - 1) / (nw * ew), (R + rt - 1) / rt, 1 }, nw * 32, 0, null, &args); + } + + /// pfdense.cpp gemm: out [M, N] = xh [M, K] (rotated, fp16) through the trellis decoded in the kernel. + pub fn pfdGemm(o: Ops, a: pfd.Args, k2: u32, lanes: bool, cfg: usize, words: usize) !void { + try pfd.check(a, k2, lanes, cfg, words); + const ki = std.mem.indexOfScalar(u32, &pfd.k2s, k2).?; + const f = o.f.pfd[ki][@intFromBool(lanes)][cfg] orelse return error.Unsupported; + const bm = pfd.cfgs[cfg][0]; + const M: usize = @intCast(a.M); + const N: usize = @intCast(a.N); + var args: cuda.Args = .{}; + args.add(a); + try o.go(f, .{ ((M + bm - 1) / bm) * (N >> 7), 1, 1 }, pfd.threads(cfg), pfd.smem(cfg, k2), null, &args); + } + + /// pfdense.cpp dequant: W_q [K, N] fp16 through the kernel's decode (tests). + pub fn pfdDequant(o: Ops, T: u64, stride_k: i64, stride_nb: i64, out: u64, K: usize, N: usize, k2: u32, lanes: bool) !void { + const ki = std.mem.indexOfScalar(u32, &pfd.k2s, k2) orelse return error.Unsupported; + const f = o.f.pfd_dequant[ki][@intFromBool(lanes)] orelse return error.Unsupported; + var args: cuda.Args = .{}; + args.add(T); + args.add(stride_k); + args.add(stride_nb); + args.add(out); + args.add(@as(c_int, @intCast(N))); + try o.go(f, .{ K / 32, N / 128, 1 }, 128, 0, null, &args); + } + + /// dsv41_engram_gate_copy: wait (bounded) for the host's ready word, then copy the selected buffer. + pub fn engramGate(o: Ops, ctl: u64, ready: c_int, err_slot: c_int, base: u64, stride: i64, dst: u64, nbytes: i64, timeout_ns: i64, vec: bool) !void { + const units: i64 = if (vec) (nbytes + 15) >> 4 else nbytes; + const blocks: usize = @intCast(std.math.clamp(@divTrunc(units + 255, 256), 1, 64)); + var args: cuda.Args = .{}; + args.add(ctl); + args.add(ready); + args.add(err_slot); + args.add(base); + args.add(stride); + args.add(dst); + args.add(nbytes); + args.add(timeout_ns); + args.add(@as(c_int, @intFromBool(vec))); + try o.go(o.f.gate_copy, .{ blocks, 1, 1 }, 256, 0, null, &args); + } + + /// dsv41_l2pace launch_paced: (address, bytes) table prefetched to L2 at a paced rate on `ctas` warps. + pub fn l2Paced(o: Ops, table: u64, n: usize, ctas: usize, chunk: u32, ns_per_piece: u64, delay_ns: u64) !void { + var args: cuda.Args = .{}; + args.add(table); + args.add(@as(c_int, @intCast(n))); + args.add(chunk); + args.add(ns_per_piece); + args.add(delay_ns); + try o.go(o.f.paced, .{ ctas, 1, 1 }, 32, 0, null, &args); + } + + /// pick.top (topk_keys.cu): each row's best k of bf16-valued fp32 logits, value desc then column asc. + pub fn topKeys(o: Ops, lg: u64, ld: usize, rows: usize, C: usize, k: usize, vals: u64, cols: u64) !void { + if (rows < 1 or C < 1 or C > topk_keys.max_cols or k < 1 or k > @min(C, topk_keys.max_k)) return error.Shape; + var args: cuda.Args = .{}; + args.add(lg); + args.add(@as(i64, @intCast(ld))); + args.add(@as(c_int, @intCast(C))); + args.add(@as(c_int, @intCast(k))); + args.add(vals); + args.add(cols); + try o.go(o.f.topk_keys, .{ rows, 1, 1 }, 1024, 0, null, &args); + } + + /// The shared expert's pointwise part: act = silu(min(g, limit)) * clamp(u, +-limit) (fp32, strides in elements). + pub fn swiglu(o: Ops, g: u64, ldg: usize, u: u64, ldu: usize, act: u64, lda: usize, rows: usize, n: usize, limit: f32) !void { + var args: cuda.Args = .{}; + args.add(g); + args.add(@as(i64, @intCast(ldg))); + args.add(u); + args.add(@as(i64, @intCast(ldu))); + args.add(act); + args.add(@as(i64, @intCast(lda))); + args.add(@as(c_int, @intCast(n))); + args.add(limit); + try o.go(o.f.swiglu, .{ @min((n + 255) / 256, 64), rows, 1 }, 256, 0, null, &args); + } + + /// out += b over n fp32 elements (prefill_moe.forward: the routed partial plus the shared expert's). + pub fn addInto(o: Ops, out: u64, b: u64, n: usize) !void { + var args: cuda.Args = .{}; + args.add(out); + args.add(b); + args.add(@as(i64, @intCast(n))); + try o.go(o.f.add, .{ @min((n + 255) / 256, 4096), 1, 1 }, 256, 0, null, &args); + } + + /// y = bf16_rn(a + b) over n elements (torch.add(out, shared, out=bf16): TF_DSV41_PF_COPIES). + pub fn addBf16(o: Ops, a: u64, b: u64, y: u64, n: usize) !void { + var args: cuda.Args = .{}; + args.add(a); + args.add(b); + args.add(y); + args.add(@as(i64, @intCast(n))); + try o.go(o.f.add_bf16, .{ @min((n + 255) / 256, 4096), 1, 1 }, 256, 0, null, &args); + } + + /// kvsplit dense: a split family's selected rows packed for the all-gather (`a` @bitCast from kv/split.zig). + pub fn kvDense(o: Ops, a: kvsplit.DenseArgs) !void { + if (a.row_bytes == 0 or a.row_bytes % 8 != 0 or a.k == 0 or a.world == 0 or a.psh > 30) return error.Shape; + const words = @as(u64, a.rows) * a.k * (a.row_bytes / 8); + if (words == 0) return; + var args: cuda.Args = .{}; + args.add(a); + try o.go(o.f.kv_dense, .{ kvsplit.grid(words), 1, 1 }, 256, 0, null, &args); + } + + /// kvsplit pack (compact): this rank's owned selected rows first, lens[W] each owner's bytes; block 256. + pub fn kvPack(o: Ops, a: kvsplit.PackArgs) !void { + const f = o.f.kv_pack orelse return error.Unsupported; + if (a.row_bytes == 0 or a.row_bytes % 8 != 0 or a.k == 0 or a.world == 0 or a.world > 8 or a.psh > 30) return error.Shape; + const n = @as(u64, a.rows) * a.k; + if (n == 0) return; + var args: cuda.Args = .{}; + args.add(a); + try o.go(f, .{ @intCast((n + 255) / 256), 1, 1 }, 256, 0, null, &args); + } + + /// kvsplit dedup: the pack over each slot's unique rows (mark, flags, pack, tok: four launches of 256 threads). + pub fn kvDedup(o: Ops, d: kvsplit.DedupArgs) !void { + const fs = o.f.kv_dedup orelse return error.Unsupported; + const a = d.p; + if (a.row_bytes == 0 or a.row_bytes % 8 != 0 or a.k == 0 or a.world == 0 or a.world > 8 or a.psh > 30) return error.Shape; + const n = @as(u64, a.rows) * a.k; + if (n == 0) return; + const blocks: usize = @intCast((n + 255) / 256); + for (fs, [4]usize{ kvsplit.grid(n), blocks, blocks, kvsplit.grid(n) }) |f, g| { + var args: cuda.Args = .{}; + args.add(d); + try o.go(f, .{ g, 1, 1 }, 256, 0, null, &args); + } + } + + /// kvsplit hist_gather + hist_tok: the history union's M rows of this rank and the selection's tokens. + pub fn kvHistory(o: Ops, a: kvsplit.HistArgs) !void { + const fs = o.f.kv_hist orelse return error.Unsupported; + if (a.row_bytes == 0 or a.row_bytes % 8 != 0 or a.k == 0 or a.world == 0 or a.world > 8 or a.psh > 30) return error.Shape; + const words = @as(u64, a.m) * (a.row_bytes / 8); + const n = @as(u64, a.rows) * a.k; + for (fs, [2]u64{ words, n }) |f, w| { + if (w == 0) continue; + var args: cuda.Args = .{}; + args.add(a); + try o.go(f, .{ kvsplit.grid(w), 1, 1 }, 256, 0, null, &args); + } + } + + /// kvsplit gather: send[i] = base row phys[i] (uint32 [n]). + pub fn kvGather(o: Ops, base: u64, row_bytes: u32, phys: u64, n: u32, send: u64) !void { + if (row_bytes == 0 or row_bytes % 8 != 0) return error.Shape; + const words = @as(u64, n) * (row_bytes / 8); + if (words == 0) return; + var args: cuda.Args = .{}; + args.add(base); + args.add(row_bytes); + args.add(phys); + args.add(n); + args.add(send); + try o.go(o.f.kv_gather, .{ kvsplit.grid(words), 1, 1 }, 256, 0, null, &args); + } + + pub const glueGo = glue_ops.glueGo; + pub const embedSend = glue_ops.embedSend; + pub const embedSum = glue_ops.embedSum; + pub const castBf16 = glue_ops.castBf16; + pub const kitWeights = glue_ops.kitWeights; + pub const kitLogits = glue_ops.kitLogits; + pub const positions = glue_ops.positions; + pub const dsCandidates = glue_ops.dsCandidates; + pub const dsStage = glue_ops.dsStage; + pub const dsAcceptRows = glue_ops.dsAcceptRows; + pub const dsStageRows = glue_ops.dsStageRows; + pub const vsChoose = glue_ops.vsChoose; + pub const widenF32 = glue_ops.widenF32; + pub const pfGo = glue_ops.pfGo; + pub const topPositions = glue_ops.topPositions; + pub const candKeys = glue_ops.candKeys; + pub const visibleCounts = glue_ops.visibleCounts; + pub const gmPicks = glue_ops.gmPicks; + pub const widthMask = glue_ops.widthMask; + pub const f64Bf16F32 = glue_ops.f64Bf16F32; + pub const widenCat = glue_ops.widenCat; + pub const ringCopy = glue_ops.ringCopy; + pub const positionsDev = glue_ops.positionsDev; + pub const gatherCols = glue_ops.gatherCols; + pub const pickPack = glue_ops.pickPack; + pub const pickMerge = glue_ops.pickMerge; + pub const carry = glue_ops.carry; + + /// tfl2pf launch_segments: (address, bytes) table prefetched to L2 (bulk: cp.async.bulk.prefetch). + pub fn l2Segments(o: Ops, table: u64, n: usize, grid: usize, threads: usize, chunk: u32, bulk: bool) !void { + var args: cuda.Args = .{}; + args.add(table); + args.add(@as(c_int, @intCast(n))); + args.add(chunk); + try o.go(o.f.segments[@intFromBool(bulk)], .{ grid, 1, 1 }, threads, 0, null, &args); + } +}; + +test "pfdense shared memory and heuristic as the reference" { + // cfg 0 (128, 2, 4, 2, 3) at K2 10: 3 x (128 x 32 x 2 + 2 x 8 x 40 x 4) + 2 x 8192 = 3 x 10752 + 16384 + try std.testing.expectEqual(@as(u32, 3 * 10752 + 16384), pfd.smem(0, 10)); + try std.testing.expectEqual(@as(u32, 256), pfd.threads(0)); + // a 2,048-row [5120 -> 1024] call on 48 SMs: 16 x 8 = 128 CTAs of BM 128 + const h = pfd.heuristic(5120, 1024, 10, 2048, 48); + try std.testing.expectEqual(@as(usize, 8), h[1]); +} + +test "mhc instances follow dispatch" { + var a: mhc.Args = .{ .R = 3 }; + try std.testing.expectEqual(@as(usize, 6), try mhc.instance(a, 3, true)); + try std.testing.expectEqual(@as(usize, 3), try mhc.instance(a, 1, true)); + a.R = 17; + try std.testing.expectError(error.Shape, mhc.instance(a, 0, false)); + const p: mhc_pf.Args = .{ .R = 40 }; + try std.testing.expectEqual(@as(usize, 5), try mhc_pf.instance(p, 1, true, false)); + try std.testing.expectError(error.Shape, mhc_pf.instance(p, 1, false, true)); + try std.testing.expectEqual(@as(u32, 47872), mhc_pf.smem(false)); +} diff --git a/zig/src/families/deepseek_v41/kernels_replay.zig b/zig/src/families/deepseek_v41/kernels_replay.zig new file mode 100644 index 0000000000..ffeb282502 --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_replay.zig @@ -0,0 +1,232 @@ +//! Replays captured extension calls through the Zig launches, so block replay compares every .cu kernel's outputs. + +const std = @import("std"); +const cuda = @import("cuda"); +const kernels = @import("kernels.zig"); +const exl3 = @import("kernels_exl3.zig"); +const dense = @import("kernels_dense.zig"); +const ops = @import("kernels_ops.zig"); +const replay_experts = @import("kernels_replay_experts.zig"); +const replay_ops = @import("kernels_replay_ops.zig"); + +const Kernels = kernels.Kernels; + +pub const DType = enum { f16, bf16, f32, f64, i8, u8, i16, i32, i64, bool }; +pub const Tensor = struct { ptr: u64, dtype: DType, shape: []const i64, stride: []const i64 }; // element strides, ptr already at storage_offset +pub const Arg = union(enum) { tensor: Tensor, int: i64, float: f64, boolean: bool, none, list: []const Arg }; +pub const Kw = struct { name: []const u8, value: Arg }; + +/// Device bytes a call may use beyond its arguments (cleared per call by the caller). +pub const Scratch = struct { + ptr: u64, + len: usize, + used: usize = 0, + + pub fn take(s: *Scratch, n: usize) !u64 { + const at = std.mem.alignForward(usize, s.used, 256); + if (at + n > s.len) return error.ScratchTooSmall; + s.used = at + n; + return s.ptr + at; + } +}; + +pub const Error = error{ MissingArgument, WrongArgument, Unsupported }; + +/// The call's arguments by pybind position, or by keyword when passed so (`names` lists the binding's parameters). +pub const Call = struct { + args: []const Arg, + kwargs: []const Kw, + names: []const []const u8, + + fn get(c: Call, i: usize) Error!Arg { + if (i < c.args.len) return c.args[i]; + for (c.kwargs) |kw| if (std.mem.eql(u8, kw.name, c.names[i])) return kw.value; + return error.MissingArgument; + } + + pub fn opt(c: Call, i: usize) Error!?Arg { + const a = c.get(i) catch |e| switch (e) { + error.MissingArgument => return null, + else => return e, + }; + return if (a == .none) null else a; + } + + pub fn tensor(c: Call, i: usize) Error!Tensor { + return switch (try c.get(i)) { + .tensor => |t| t, + else => error.WrongArgument, + }; + } + + /// An optional tensor's address (None or an empty tensor: 0, as the bindings pass nullptr). + pub fn ptrOrNull(c: Call, i: usize) Error!u64 { + const a = (try c.opt(i)) orelse return 0; + return switch (a) { + .tensor => |t| if (numel(t) == 0) 0 else t.ptr, + else => error.WrongArgument, + }; + } + + pub fn ptr(c: Call, i: usize) Error!u64 { + return (try c.tensor(i)).ptr; + } + + pub fn int(c: Call, i: usize) Error!i64 { + return switch (try c.get(i)) { + .int => |v| v, + .boolean => |b| @intFromBool(b), + else => error.WrongArgument, + }; + } + + pub fn usz(c: Call, i: usize) Error!usize { + const v = try c.int(i); + if (v < 0) return error.WrongArgument; + return @intCast(v); + } + + pub fn u(c: Call, i: usize) Error!u32 { + return @intCast(try c.usz(i)); + } + + pub fn float(c: Call, i: usize) Error!f64 { + return switch (try c.get(i)) { + .float => |v| v, + .int => |v| @floatFromInt(v), + else => error.WrongArgument, + }; + } + + pub fn boolean(c: Call, i: usize, default: bool) Error!bool { + const a = (try c.opt(i)) orelse return default; + return switch (a) { + .boolean => |b| b, + .int => |v| v != 0, + else => error.WrongArgument, + }; + } + + pub fn list(c: Call, i: usize) Error![]const Arg { + return switch (try c.get(i)) { + .list => |l| l, + else => error.WrongArgument, + }; + } +}; + +pub fn numel(t: Tensor) i64 { + var n: i64 = 1; + for (t.shape) |d| n *= d; + return n; +} + +pub fn dim(t: Tensor, i: usize) Error!usize { + if (i >= t.shape.len) return error.WrongArgument; + return @intCast(t.shape[i]); +} + +pub fn ioType(t: DType) Error!dense.DType { + return switch (t) { + .f16 => .f16, + .bf16 => .bf16, + .f32 => .f32, + else => error.WrongArgument, + }; +} + +pub fn input(t: DType) Error!exl3.Input { + return switch (t) { + .bf16 => .bf16, + .f16 => .f16, + else => error.WrongArgument, + }; +} + +pub fn elem(t: DType) usize { + return switch (t) { + .i8, .u8, .bool => 1, + .f16, .bf16, .i16 => 2, + .f32, .i32 => 4, + .f64, .i64 => 8, + }; +} + +pub fn is(a: []const u8, b: []const u8) bool { + return std.mem.eql(u8, a, b); +} + +/// One captured extension call re-issued on `s`; false when there is no binding for it (reported as skipped). +pub fn call(k: *const Kernels, s: cuda.Stream, ext: []const u8, func: []const u8, args: []const Arg, kwargs: []const Kw, scratch: *Scratch) !bool { + _ = scratch; + if (is(ext, "tf_dsv41_kv_glue_v1") and is(func, "norm_store")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "x", "w", "cs", "v", "s", "pos", "sl", "x_stride", "cs_stride", "v_stride", "s_stride", "inv_k", "eps", "ring", "rows" } }; + const x = try c.tensor(0); + const rows = try c.boolean(14, false); + if (x.dtype != .bf16 or try dim(x, 1) != 512 or (try c.tensor(1)).dtype != .f32 or + (try c.tensor(2)).dtype != .f32 or (try c.tensor(5)).dtype != (if (rows) DType.i64 else DType.i32) or + (rows and (try c.tensor(6)).dtype != .i64)) return error.WrongArgument; + try k.kvStore(s).normStore(.{ + .x = x.ptr, + .w = try c.ptr(1), + .cs = try c.ptr(2), + .v = try c.ptr(3), + .s = try c.ptr(4), + .pos = try c.ptr(5), + .sl = try c.ptrOrNull(6), + .x_stride = try c.int(7), + .cs_stride = try c.int(8), + .v_stride = try c.int(9), + .s_stride = try c.int(10), + .inv_k = @floatCast(try c.float(11)), + .eps = @floatCast(try c.float(12)), + .ring = @intCast(try c.int(13)), + .rows = @intFromBool(rows), + }, try dim(x, 0)); + return true; + } + if (is(ext, "tensorfold_exl3_linear_v4")) return linearExt(k.linears(s), func, args, kwargs); + if (is(ext, "tf_dsv41_x3seg_v3")) return segExt(k.linears(s), func, args, kwargs, false); + if (is(ext, "tf_dsv41_dense3_v1")) return segExt(k.linears(s), func, args, kwargs, true); + if (is(ext, "tensorfold_exl3_experts_v1")) return expertsExt(k.experts(s), func, args, kwargs); + if (is(ext, "tf_dsv41_x3ld_v1") or is(ext, "tf_dsv41_x3pf_v1")) return groupedExt(k.experts(s), ext, func, args, kwargs); + if (is(ext, "tf_dsv41_x3gm_v1")) return x3gmExt(k.experts(s), func, args, kwargs); + if (is(ext, "tf_dsv41_x3ld_epi_v1")) return x3ldEpiExt(k.experts(s), func, args, kwargs); + if (is(ext, "tf_dsv41_attn_cuda_v1")) return attnExt(k.others(s), func, args, kwargs); + if (is(ext, "tf_dsv41_topk_b_v1") and is(func, "topk")) return topkCall(k.others(s), args, kwargs, true); + if (is(ext, "tf_dsv41_mhc_cuda_v1") and is(func, "run")) return mhcRun(k.others(s), args, kwargs); + if (is(ext, "tf_dsv41_mhc_cuda_v1") and is(func, "coef")) return mhcCoef(k.others(s), args, kwargs); + if (is(ext, "tf_dsv41_mhc_pf_v1") and is(func, "run")) return mhcPfRun(k.others(s), args, kwargs); + if (is(ext, "tf_dsv41_router_gemv_v2") and is(func, "route")) return route(k.others(s), args, kwargs); + if (is(ext, "tf_dsv41_pfdense_v1")) return pfdenseExt(k.others(s), func, args, kwargs); + if (is(ext, "tf_dsv41_l2pace_v1") and is(func, "paced")) return paced(k.others(s), args, kwargs); + if (is(ext, "tf_dsv41_l2pf_v1") and is(func, "segments")) return segments(k.others(s), args, kwargs); + return false; // engram_gate.gate_copy reads host-mapped words a capture cannot relocate: not replayed +} + +const linearExt = replay_experts.linearExt; +const segExt = replay_experts.segExt; +const expertsExt = replay_experts.expertsExt; +const groupedExt = replay_experts.groupedExt; +const x3gmExt = replay_experts.x3gmExt; +const x3ldEpiExt = replay_experts.x3ldEpiExt; +const attnExt = replay_ops.attnExt; +const topkCall = replay_ops.topkCall; +const mhcRun = replay_ops.mhcRun; +const mhcCoef = replay_ops.mhcCoef; +const mhcPfRun = replay_ops.mhcPfRun; +const route = replay_ops.route; +const pfdenseExt = replay_ops.pfdenseExt; +const paced = replay_ops.paced; +const segments = replay_ops.segments; + +test "arguments by position and keyword" { + const shape = [_]i64{ 2, 3 }; + const t: Tensor = .{ .ptr = 64, .dtype = .bf16, .shape = &shape, .stride = &.{ 3, 1 } }; + const c: Call = .{ .args = &.{ .{ .tensor = t }, .{ .int = 7 } }, .kwargs = &.{.{ .name = "pdl", .value = .{ .boolean = true } }}, .names = &.{ "x", "n", "z", "pdl" } }; + try std.testing.expectEqual(@as(u64, 64), try c.ptr(0)); + try std.testing.expectEqual(@as(usize, 7), try c.usz(1)); + try std.testing.expectEqual(@as(u64, 0), try c.ptrOrNull(2)); + try std.testing.expect(try c.boolean(3, false)); + try std.testing.expectError(error.WrongArgument, c.tensor(1)); +} diff --git a/zig/src/families/deepseek_v41/kernels_replay_experts.zig b/zig/src/families/deepseek_v41/kernels_replay_experts.zig new file mode 100644 index 0000000000..0d470ebcfd --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_replay_experts.zig @@ -0,0 +1,325 @@ +//! Replays the EXL3 linear and routed-expert extension calls (linear, x3seg, dense3, experts, x3ld/x3pf, x3gm). + +const std = @import("std"); +const exl3 = @import("kernels_exl3.zig"); +const dense = @import("kernels_dense.zig"); +const replay = @import("kernels_replay.zig"); + +const Call = replay.Call; +const Arg = replay.Arg; +const Kw = replay.Kw; +const numel = replay.numel; +const dim = replay.dim; +const ioType = replay.ioType; +const input = replay.input; +const elem = replay.elem; +const is = replay.is; + +// tensorfold_exl3_linear_v4 (linear.cpp) +pub fn linearExt(o: dense.Ops, func: []const u8, args: []const Arg, kwargs: []const Kw) !bool { + if (is(func, "rot_in")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "x", "suh", "xh" } }; + const x = try c.tensor(0); + try o.rotIn(x.ptr, try ioType(x.dtype), try c.ptr(1), try c.ptr(2), try dim(x, 0), try dim(x, 1)); + return true; + } + if (is(func, "linear") or is(func, "linear_skip")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "xh", "T", "stride_k", "stride_nb", "svh", "bias", "y", "Z", "counters", "K2", "cb", "SK", "WK", "skip" } }; + if (try c.int(10) != 2) return false; // mul1 only is built (every DeepSeek-V4.1 matrix) + const xh = try c.tensor(0); + const y = try c.tensor(6); + const skip: u64 = if (is(func, "linear_skip")) try c.ptrOrNull(13) else 0; + try o.linear(xh.ptr, try c.ptr(1), try c.int(2), try c.int(3), try c.ptr(4), try c.ptrOrNull(5), y.ptr, try ioType(y.dtype), try c.ptrOrNull(7), try c.ptr(8), try dim(xh, 0), try dim(xh, 1), try dim(y, 1), try c.u(9), try c.usz(11), try c.usz(12), skip); + return true; + } + if (is(func, "unpack")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "T", "W", "stride_k", "stride_nb", "K2", "cb" } }; + if (try c.int(5) != 2) return false; + const w = try c.tensor(1); + try o.unpack(try c.ptr(0), w.ptr, try dim(w, 0), try dim(w, 1), try c.int(2), try c.int(3), try c.u(4)); + return true; + } + return false; +} + +// tf_dsv41_x3seg_v3 / tf_dsv41_dense3_v1 (x3seg.cpp, dense3.cpp): the tables built as dsv41_*_linear_cuda builds them +pub fn segExt(o: dense.Ops, func: []const u8, args: []const Arg, kwargs: []const Kw, lanes: bool) !bool { + if (is(func, "rot_in") and !lanes) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "xs", "suh", "xh", "xh_off", "pdl", "nw", "eps" } }; + const nws: []const Arg = if (try c.opt(5)) |v| switch (v) { + .list => |l| l, + else => return error.WrongArgument, + } else &.{}; + const eps: f32 = if ((try c.opt(6)) != null) @floatCast(try c.float(6)) else 0; + const xs = try c.list(0); + const suh = try c.list(1); + const offs = try c.list(3); + const xh = try c.ptr(2); + if (xs.len < 1 or xs.len > dense.max_seg or suh.len != xs.len or offs.len != xs.len) return error.WrongArgument; + var tab = std.mem.zeroes(dense.RotTable); + tab.n = @intCast(xs.len); + var M: usize = 0; + for (xs, suh, offs, 0..) |xa, sa, oa, i| { + if (xa != .tensor or sa != .tensor or oa != .int) return error.WrongArgument; + const x = xa.tensor; + if (i == 0) M = try dim(x, 0); + tab.s[i] = .{ .x = x.ptr, .suh = sa.tensor.ptr, .xh = xh + 2 * @as(u64, @intCast(oa.int)), .ld = x.stride[0], .K = @intCast(try dim(x, 1)), .x_dtype = @backingInt(try ioType(x.dtype)) }; + if (i < nws.len) { // R1c: the folded norm (an empty weight: none) + const w = switch (nws[i]) { + .tensor => |t| t, + else => return error.WrongArgument, + }; + if (numel(w) != 0) { + tab.s[i].nw = w.ptr; + tab.s[i].nw_dtype = @backingInt(try ioType(w.dtype)); + tab.s[i].inv_k = dense.invK(try dim(x, 1)); + tab.s[i].eps = eps; + } + } + } + try o.segRotIn(tab, M, try c.boolean(4, false)); + return true; + } + if (is(func, "linear")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "xh", "T", "svh", "ys", "Z", "counters", "meta", "programs", "K2", "WK", "minb", "pdl" } }; + const T = try c.list(1); + const svh = try c.list(2); + const ys = try c.list(3); + const meta = try c.list(6); + const n = T.len; + if (n < 1 or n > dense.max_seg or svh.len != n or ys.len != n or meta.len != 11 * n) return error.WrongArgument; + const xh = try c.ptr(0); + const z = try c.ptrOrNull(4); + const counters = try c.ptr(5); + if (ys[0] != .tensor) return error.WrongArgument; + const y0 = ys[0].tensor; + const M = try dim(y0, 0); + const y_dtype = try ioType(y0.dtype); + var tab = std.mem.zeroes(dense.LinTable); + tab.n = @intCast(n); + for (0..n) |i| { + var m: [11]i64 = undefined; + for (&m, meta[11 * i ..][0..11]) |*v, a| v.* = switch (a) { + .int => |x| x, + else => return error.WrongArgument, + }; + if (T[i] != .tensor or svh[i] != .tensor or ys[i] != .tensor) return error.WrongArgument; + const y = ys[i].tensor; + tab.s[i] = .{ + .xh = xh + 2 * @as(u64, @intCast(m[6])), + .T = T[i].tensor.ptr, + .stride_k = m[4], + .stride_nb = m[5], + .svh = svh[i].tensor.ptr, + .y = y.ptr + @as(u64, @intCast(m[9])) * elem(y.dtype), + .Z = if (m[2] > 1) z + 4 * @as(u64, @intCast(m[7])) else 0, + .counters = counters + 4 * @as(u64, @intCast(m[8])), + .K = @intCast(m[0]), + .N = @intCast(m[1]), + .SK = @intCast(m[2]), + .wk = @intCast(m[3]), + .p0 = @intCast(m[10]), + .ld = @intCast(y.stride[0]), + }; + } + const programs = try c.usz(7); + const k2 = try c.u(8); + const wk = try c.usz(9); + const minb = try c.u(10); + const pdl = try c.boolean(11, false); + if (lanes) try o.lanesLinearRaw(tab, y_dtype, M, programs, k2, wk, minb, pdl) else try o.segLinearRaw(tab, y_dtype, M, programs, k2, wk, minb, pdl); + return true; + } + if (is(func, "relayout") and lanes) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "src", "dst", "K2", "K", "lanes" } }; + const src = try c.tensor(0); + try o.relayout(src.ptr, try c.ptr(1), @intCast(numel(src)), try c.u(2), try c.usz(3), try c.boolean(4, true)); + return true; + } + if (is(func, "unpack") and lanes) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "T", "W", "stride_k", "stride_nb", "K2" } }; + const w = try c.tensor(1); + try o.lanesUnpack(try c.ptr(0), w.ptr, try dim(w, 0), try dim(w, 1), try c.int(2), try c.int(3), try c.u(4)); + return true; + } + return false; +} + +const grouped_names = [_][]const u8{ "X0", "X1", "TP0", "TP1", "B0", "B1", "uids", "ucount", "members", "Z", "mats", "K", "N", "P", "SK", "slots", "cb" }; + +fn groupedArgs(c: Call, lo_i: usize) !exl3.Grouped { + const uids = try c.tensor(6); + const members = try c.tensor(8); + return .{ + .x0 = try c.ptr(0), + .x1 = try c.ptr(1), + .tp0 = try c.ptr(2), + .tp1 = try c.ptr(3), + .k2_0 = try c.ptr(4), + .k2_1 = try c.ptr(5), + .uids = uids.ptr, + .ucount = try c.ptr(7), + .members = members.ptr, + .z = try c.ptr(9), + .mats = try c.usz(10), + .K = try c.usz(11), + .N = try c.usz(12), + .P = try c.usz(13), + .SK = try c.usz(14), + .slots = try c.usz(15), + .maxm = try dim(members, 1), + .nexp = try dim(uids, 0), + .lo = try c.u(lo_i), + .hi = try c.u(lo_i + 1), + }; +} + +// tensorfold_exl3_experts_v1 (experts.cpp) +pub fn expertsExt(o: exl3.Ops, func: []const u8, args: []const Arg, kwargs: []const Kw) !bool { + if (is(func, "group_rot_in")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "x", "x_stride", "pick", "suh0", "suh1", "out0", "out1", "rows", "K", "slots", "E", "uids", "ucount", "members" } }; + const x = try c.tensor(0); + if (x.dtype != .bf16) return error.WrongArgument; + const members = try c.tensor(13); + try o.groupRotIn(x.ptr, try c.usz(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), try c.ptr(5), try c.ptr(6), try c.ptr(11), try c.ptr(12), members.ptr, try c.usz(7), try c.usz(8), try c.usz(9), try c.usz(10), try dim(members, 1)); + return true; + } + if (is(func, "grouped")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &(grouped_names ++ [_][]const u8{ "nt", "warps", "pf", "lo", "hi" }) }; + if (try c.int(16) != 2) return false; // cb 2 (mul1) instances only + try o.grouped(try groupedArgs(c, 20), try c.u(17), try c.u(18), try c.u(19)); + return true; + } + if (is(func, "group")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "pick", "uids", "ucount", "members", "R", "slots", "E" } }; + const members = try c.tensor(3); + try o.group(try c.ptr(0), try c.ptr(1), try c.ptr(2), members.ptr, try c.usz(4), try c.usz(5), try c.usz(6), try dim(members, 1)); + return true; + } + if (is(func, "rot_in")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "x", "x_stride", "pick", "suh0", "suh1", "out0", "out1", "rows", "K", "slots", "E" } }; + const x = try c.tensor(0); + try o.rotIn(try input(x.dtype), x.ptr, try c.usz(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), try c.ptr(5), try c.ptr(6), try c.usz(7), try c.usz(8), try c.usz(9), try c.usz(10)); + return true; + } + if (is(func, "gateup_epilogue")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "Z", "pick", "svh_g", "svh_u", "suh_d", "xd", "rows", "P", "N", "SK", "slots", "E", "limit", "act_mode" } }; + try o.gateupEpilogue(try c.ptr(0), try c.ptr(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), try c.ptr(5), try c.usz(6), try c.usz(7), try c.usz(8), try c.usz(9), try c.usz(10), try c.usz(11), @floatCast(try c.float(12)), @intCast(try c.int(13))); + return true; + } + if (is(func, "down_epilogue")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "Z", "pick", "svh_d", "y", "rows", "P", "D", "SK", "slots", "E" } }; + try o.downEpilogue(try c.ptr(0), try c.ptr(1), try c.ptr(2), try c.ptr(3), try c.usz(4), try c.usz(5), try c.usz(6), try c.usz(7), try c.usz(8), try c.usz(9)); + return true; + } + if (is(func, "combine")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "y", "wts", "out", "rows", "D", "slots" } }; + try o.combine(try c.ptr(0), try c.ptr(1), try c.ptr(2), try c.usz(3), try c.usz(4), try c.usz(5)); + return true; + } + if (is(func, "down_combine")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "Z", "pick", "svh_d", "y", "wts", "out", "rows", "P", "D", "SK", "slots", "E" } }; + const out = try c.tensor(5); + if (out.dtype != .f32 and out.dtype != .bf16) return error.WrongArgument; + try o.downCombine(try c.ptr(0), try c.ptr(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), out.ptr, try c.usz(6), try c.usz(7), try c.usz(8), try c.usz(9), try c.usz(10), try c.usz(11), out.dtype == .bf16); + return true; + } + if (is(func, "dequant")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "T", "out", "k2", "cb" } }; + if (try c.int(3) != 2) return false; + const t = try c.tensor(0); + try o.dequant(t.ptr, try c.ptr(1), 16 * try dim(t, 0), 16 * try dim(t, 1), try c.u(2)); + return true; + } + return false; +} + +// tf_dsv41_x3ld_v1.grouped (..., cb, nt, pd, probe, lo, hi, pdl) / tf_dsv41_x3pf_v1.grouped (..., cb, nt, mtl, lo, hi) +pub fn groupedExt(o: exl3.Ops, ext: []const u8, func: []const u8, args: []const Arg, kwargs: []const Kw) !bool { + if (!is(func, "grouped")) return false; + if (is(ext, "tf_dsv41_x3ld_v1")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &(grouped_names ++ [_][]const u8{ "nt", "pd", "probe", "lo", "hi", "pdl" }) }; + if (try c.int(16) != 2) return error.Unsupported; // the binding refuses it too + try o.x3ld(try groupedArgs(c, 20), try c.u(17), try c.u(18), try c.u(19), try c.boolean(22, false)); + return true; + } + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &(grouped_names ++ [_][]const u8{ "nt", "mtl", "lo", "hi" }) }; + if (try c.int(16) != 2) return error.Unsupported; + try o.x3pf(try groupedArgs(c, 19), try c.u(17), try c.u(18)); + return true; +} + +// tf_dsv41_x3ld_epi_v1 (TF_DSV41_X3LD_EPI): x3ld.grouped's arguments to hi, then the epilogue's, then the ticket words +pub fn x3ldEpiExt(o: exl3.Ops, func: []const u8, args: []const Arg, kwargs: []const Kw) !bool { + const base = grouped_names ++ [_][]const u8{ "nt", "pd", "lo", "hi", "pick" }; + if (is(func, "gateup")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &(base ++ [_][]const u8{ "svh_g", "svh_u", "suh_d", "xd", "E", "limit", "act_mode", "ticket" }) }; + if (try c.int(16) != 2 or try c.int(17) != 8) return error.Unsupported; + try o.x3ldEpi(try groupedArgs(c, 19), .gu, try c.u(18), .{ + .pick = try c.ptr(21), + .sv0 = try c.ptr(22), + .sv1 = try c.ptr(23), + .sd = try c.ptr(24), + .xd = try c.ptr(25), + .E = @intCast(try c.int(26)), + .limit = @floatCast(try c.float(27)), + .act_mode = @intCast(try c.int(28)), + .ticket = try c.ptr(29), + }); + return true; + } + if (is(func, "down")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &(base ++ [_][]const u8{ "svh_d", "y", "wts", "out", "E", "ticket" }) }; + if (try c.int(16) != 2 or try c.int(17) != 8) return error.Unsupported; + const out = try c.tensor(25); + if (out.dtype != .f32 and out.dtype != .bf16) return error.WrongArgument; + try o.x3ldEpi(try groupedArgs(c, 19), if (out.dtype == .bf16) .dnb else .dn, try c.u(18), .{ + .pick = try c.ptr(21), + .sv0 = try c.ptr(22), + .y = try c.ptr(23), + .wts = try c.ptr(24), + .out = out.ptr, + .E = @intCast(try c.int(26)), + .ticket = try c.ptr(27), + }); + return true; + } + return false; +} + +// tf_dsv41_x3gm_v1 (x3gm.cpp) +pub fn x3gmExt(o: exl3.Ops, func: []const u8, args: []const Arg, kwargs: []const Kw) !bool { + if (is(func, "gateup")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "xg", "xu", "tg", "tu", "order", "pe", "poff", "pcnt", "npass", "ticket", "svh_g", "svh_u", "suh_d", "xd", "K", "N", "k2", "shx", "cfg", "use_ticket", "limit", "tables" } }; + try o.gmGateup(try c.ptr(0), try c.ptr(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), try c.ptr(5), try c.ptr(6), try c.ptr(7), try c.ptr(8), try c.ptr(9), try c.ptr(10), try c.ptr(11), try c.ptr(12), try c.ptr(13), try c.usz(14), try c.usz(15), try c.u(16), try c.boolean(17, false), try c.usz(18), try c.boolean(19, true), @floatCast(try c.float(20)), try c.boolean(21, false)); + return true; + } + if (is(func, "down")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "xd", "td", "order", "pe", "poff", "pcnt", "npass", "ticket", "svh_d", "y", "K", "N", "k2", "cfg", "use_ticket", "tables" } }; + try o.gmDown(try c.ptr(0), try c.ptr(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), try c.ptr(5), try c.ptr(6), try c.ptr(7), try c.ptr(8), try c.ptr(9), try c.usz(10), try c.usz(11), try c.u(12), try c.usz(13), try c.boolean(14, true), try c.boolean(15, false)); + return true; + } + if (is(func, "gateup2")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "xg", "xu", "tg", "tu", "k2e", "order", "pe", "poff", "pcnt", "npass", "ticket", "svh_g", "svh_u", "suh_d", "xd", "K", "N", "shx", "cfg", "use_ticket", "limit", "tables" } }; + try o.gm2Gateup(try c.ptr(0), try c.ptr(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), try c.ptr(5), try c.ptr(6), try c.ptr(7), try c.ptr(8), try c.ptr(9), try c.ptr(10), try c.ptr(11), try c.ptr(12), try c.ptr(13), try c.ptr(14), try c.usz(15), try c.usz(16), try c.boolean(17, false), try c.usz(18), try c.boolean(19, true), @floatCast(try c.float(20)), try c.boolean(21, false)); + return true; + } + if (is(func, "down2")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "xd", "td", "k2e", "order", "pe", "poff", "pcnt", "npass", "ticket", "svh_d", "y", "K", "N", "cfg", "use_ticket", "tables" } }; + try o.gm2Down(try c.ptr(0), try c.ptr(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), try c.ptr(5), try c.ptr(6), try c.ptr(7), try c.ptr(8), try c.ptr(9), try c.ptr(10), try c.usz(11), try c.usz(12), try c.usz(13), try c.boolean(14, true), try c.boolean(15, false)); + return true; + } + if (is(func, "rot")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "x", "x_stride", "pick", "suh0", "suh1", "out0", "out1", "K", "slots", "pairs", "mats" } }; + const x = try c.tensor(0); + try o.gmRot(try input(x.dtype), x.ptr, try c.usz(1), try c.ptr(2), try c.ptr(3), try c.ptr(4), try c.ptr(5), try c.ptr(6), try c.usz(7), try c.usz(8), try c.usz(9), try c.usz(10)); + return true; + } + if (is(func, "dequant")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "t", "out", "k2" } }; + const out = try c.tensor(1); + try o.gmDequant(try c.ptr(0), out.ptr, try dim(out, 0), try dim(out, 1), try c.u(2)); + return true; + } + return false; +} diff --git a/zig/src/families/deepseek_v41/kernels_replay_ops.zig b/zig/src/families/deepseek_v41/kernels_replay_ops.zig new file mode 100644 index 0000000000..8be2b1bdf1 --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_replay_ops.zig @@ -0,0 +1,308 @@ +//! Replays the other kernels' extension calls: attention, top-k, mHC, router, prefill dense GEMM and L2 prefetch. + +const ops = @import("kernels_ops.zig"); +const replay = @import("kernels_replay.zig"); + +const Call = replay.Call; +const Tensor = replay.Tensor; +const Arg = replay.Arg; +const Kw = replay.Kw; +const numel = replay.numel; +const dim = replay.dim; +const is = replay.is; + +/// The bindings' `tp(t)`: an empty tensor is nullptr. +fn tp(t: Tensor) u64 { + return if (numel(t) == 0) 0 else t.ptr; +} + +fn stride0(t: Tensor) i64 { + return if (numel(t) == 0 or t.stride.len == 0) 0 else t.stride[0]; +} + +// tf_dsv41_attn_cuda_v1 (csa2/attn_cuda.cpp): attn and topk as dsv41_attn_run_cuda / dsv41_topk_run_cuda fill Args +pub fn attnExt(o: ops.Ops, func: []const u8, args: []const Arg, kwargs: []const Kw) !bool { + if (is(func, "attn")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "q", "cv", "csc", "tok", "cnt", "sv", "ssc", "lo", "hi", "pos", "sl", "pt", "sink", "cs", "po", "pm", "pl", "out", "ticket", "ring", "psh", "pts", "kw", "split", "qrope" } }; + const q = try c.tensor(0); + const cv = try c.tensor(1); + const csc = try c.tensor(2); + const tok = try c.tensor(3); + const sv = try c.tensor(5); + const ssc = try c.tensor(6); + const hi = try c.tensor(8); + const sl = try c.tensor(10); + const cs = try c.tensor(13); + const a: ops.attn.Args = .{ + .q = tp(q), + .has_comp = @intFromBool(numel(cv) != 0), + .cv = tp(cv), + .csc = tp(csc), + .cvs = stride0(cv), + .css = stride0(csc), + .tok = tp(tok), + .ts = stride0(tok), + .cnt = tp(try c.tensor(4)), + .sv = tp(sv), + .ssc = tp(ssc), + .svs = sv.stride[0], + .sss = ssc.stride[0], + .lo = tp(try c.tensor(7)), + .hi = tp(hi), + .pos = (try c.tensor(9)).ptr, + .sl = tp(sl), + .pt = tp(try c.tensor(11)), + .pts = try c.int(21), + .psh = @intCast(try c.int(20)), + .sink = tp(try c.tensor(12)), + .cs = tp(cs), + .csst = cs.stride[0], + .po = tp(try c.tensor(14)), + .pm = tp(try c.tensor(15)), + .pl = tp(try c.tensor(16)), + .out = tp(try c.tensor(17)), + .ticket = tp(try c.tensor(18)), + .R = @intCast(try dim(q, 0)), + .H = @intCast(try dim(q, 1)), + .ring = @intCast(try c.int(19)), + }; + const split: u32 = if (try c.opt(23)) |v| switch (v) { + .int => |x| @intCast(x), + else => return error.WrongArgument, + } else 0; // captures before R1 have no split: today's launch + var b = a; + const qrope: i64 = if (try c.opt(24)) |v| switch (v) { + .int => |n| n, + else => return error.WrongArgument, + } else 0; + b.qrope = @intFromBool(qrope != 0 and split > 0); + try o.attention(b, try c.u(22), numel(sl) > 0, numel(hi) > 0, split); + return true; + } + if (is(func, "topk")) return topkCall(o, args, kwargs, false); + return false; +} + +/// tf_dsv41_attn_cuda_v1.topk, or with `bounded` tf_dsv41_topk_b_v1.topk (TF_DSV41_INDEX_BOUND, same arguments) +pub fn topkCall(o: ops.Ops, args: []const Arg, kwargs: []const Kw, bounded: bool) !bool { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "s0", "c0", "o0", "n0", "nk0", "k0", "m0", "e0", "s1", "c1", "o1", "n1", "nk1", "k1", "m1", "e1", "jobs", "pos", "ratio", "bs", "R", "CL" } }; + const jobs = try c.usz(16); + const pos = try c.tensor(17); + var a: ops.topk.Args = .{ .job = undefined, .pos = pos.ptr, .pos64 = @intFromBool(pos.dtype == .i64), .ratio = @intCast(try c.int(18)), .bs = @intCast(try c.int(19)) }; + a.job[0] = try topkJob(c, 0); + a.job[1] = if (jobs > 1) try topkJob(c, 8) else a.job[0]; + const e0 = try c.usz(7); + const ept = if (jobs > 1) @max(e0, try c.usz(15)) else e0; + if (bounded) try o.topKB(a, jobs, try c.usz(20), try c.usz(21), ept) else try o.topK(a, jobs, try c.usz(20), try c.usz(21), ept); + return true; +} + +fn topkJob(c: Call, at: usize) !ops.topk.Job { + const sc = try c.tensor(at); + const cand = try c.tensor(at + 1); + const out = try c.tensor(at + 2); + return .{ .s = sc.ptr, .ss = sc.stride[0], .cand = tp(cand), .cs = stride0(cand), .out = out.ptr, .os = out.stride[0], .cnt = tp(try c.tensor(at + 3)), .nk = @intCast(try c.int(at + 4)), .k = @intCast(try c.int(at + 5)), .mode = @intCast(try c.int(at + 6)), .ept = @intCast(try c.int(at + 7)) }; +} + +// tf_dsv41_mhc_cuda_v1.run (mhc_cuda.cpp / dsv41_mhc_run_cuda) +pub fn mhcRun(o: ops.Ops, args: []const Arg, kwargs: []const Kw) !bool { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "x", "xout", "g", "post", "comb", "pre", "fn", "base", "scale", "part", "c", "tap", "nw", "out", "opre", "opost", "ocomb", "cnt", "R", "mode", "eps", "hc_eps", "post_alpha", "iters", "spin", "defer" } }; + const x = try c.tensor(0); + const g = try c.tensor(2); + const f = try c.tensor(6); + const tap = try c.tensor(11); + const nw = try c.tensor(12); + const a: ops.mhc.Args = .{ + .x = tp(x), + .xs = x.stride[0], + .xout = tp(try c.tensor(1)), + .g = tp(g), + .gr = stride0(g), + .world = if (numel(g) != 0) @intCast(try dim(g, 0)) else 1, + .gbf16 = @intFromBool(g.dtype == .bf16), + .post = tp(try c.tensor(3)), + .comb = tp(try c.tensor(4)), + .pre = tp(try c.tensor(5)), + .@"fn" = tp(f), + .base = tp(try c.tensor(7)), + .scale = tp(try c.tensor(8)), + .part = tp(try c.tensor(9)), + .c = tp(try c.tensor(10)), + .tap = tp(tap), + .ts = stride0(tap), + .nw = nw.ptr, + .nwbf16 = @intFromBool(nw.dtype == .bf16), + .out = tp(try c.tensor(13)), + .opre = tp(try c.tensor(14)), + .opost = tp(try c.tensor(15)), + .ocomb = tp(try c.tensor(16)), + .cnt = tp(try c.tensor(17)), + .R = @intCast(try c.int(18)), + .eps = @floatCast(try c.float(20)), + .hc_eps = @floatCast(try c.float(21)), + .post_alpha = @floatCast(try c.float(22)), + .iters = @intCast(try c.int(23)), + }; + // R1: spin < 0 (or no spin: a capture before R1) is today's finish; defer only with the tail + var b = a; + b.spin = if (try c.opt(24)) |v| switch (v) { + .int => |n| n, + else => return error.WrongArgument, + } else -1; + b.@"defer" = @intFromBool(b.spin >= 0 and (if (try c.opt(25)) |v| switch (v) { + .int => |n| n != 0, + .boolean => |t| t, + else => return error.WrongArgument, + } else false)); + try o.mhcBoundary(b, try c.u(19), numel(f) != 0 and f.dtype == .f32); + return true; +} + +// tf_dsv41_mhc_cuda_v1.coef (R1: dsv41_mhc_coef_cuda): a deferred boundary's coefficients +pub fn mhcCoef(o: ops.Ops, args: []const Arg, kwargs: []const Kw) !bool { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "part", "base", "scale", "opre", "opost", "ocomb", "R", "eps", "hc_eps", "post_alpha", "iters" } }; + const a: ops.mhc.Args = .{ + .part = try c.ptr(0), + .base = try c.ptr(1), + .scale = try c.ptr(2), + .opre = try c.ptr(3), + .opost = try c.ptr(4), + .ocomb = try c.ptr(5), + .R = @intCast(try c.int(6)), + .eps = @floatCast(try c.float(7)), + .hc_eps = @floatCast(try c.float(8)), + .post_alpha = @floatCast(try c.float(9)), + .iters = @intCast(try c.int(10)), + }; + try o.mhcCoef(a); + return true; +} + +// tf_dsv41_mhc_pf_v1.run (mhc_pf.cpp / dsv41_mhc_pf_run_cuda) +pub fn mhcPfRun(o: ops.Ops, args: []const Arg, kwargs: []const Kw) !bool { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "x", "xout", "g", "post", "comb", "pre", "fn", "part", "c", "tap", "R", "mode" } }; + const x = try c.tensor(0); + const g = try c.tensor(2); + const f = try c.tensor(6); + const tap = try c.tensor(9); + const a: ops.mhc_pf.Args = .{ + .x = tp(x), + .xs = x.stride[0], + .xout = tp(try c.tensor(1)), + .g = tp(g), + .gr = stride0(g), + .world = if (numel(g) != 0) @intCast(try dim(g, 0)) else 1, + .gbf16 = @intFromBool(g.dtype == .bf16), + .post = tp(try c.tensor(3)), + .comb = tp(try c.tensor(4)), + .pre = tp(try c.tensor(5)), + .@"fn" = f.ptr, + .part = tp(try c.tensor(7)), + .c = tp(try c.tensor(8)), + .tap = tp(tap), + .ts = stride0(tap), + .R = @intCast(try c.int(10)), + }; + try o.mhcSite(a, try c.u(11), f.dtype == .f32, numel(tap) > 0); + return true; +} + +/// route's `prune` (R1c): a list of 7 floats, absent or empty = off. +fn pruneArg(c: Call, i: usize) !ops.rg.Prune { + const l = (try c.opt(i)) orelse return .{}; + const items = switch (l) { + .list => |x| x, + else => return error.WrongArgument, + }; + var v: [7]f64 = undefined; + if (items.len != 0 and items.len != 7) return error.WrongArgument; + for (items, 0..) |it, j| v[j] = switch (it) { + .float => |f| f, + .int => |n| @floatFromInt(n), + else => return error.WrongArgument, + }; + return ops.rg.pruneOf(v[0..items.len]); +} + +// tf_dsv41_router_gemv_v2.route (router_gemv.cpp / dsv41_rg_route_cuda) +pub fn route(o: ops.Ops, args: []const Arg, kwargs: []const Kw) !bool { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "X", "W", "B", "pick", "wts", "LG", "cnt", "K", "slots", "scale", "rt", "nw", "ew", "select", "kit", "gid", "gcnt", "gmem", "GE", "group", "prune" } }; + const x = try c.tensor(0); + const w = try c.tensor(1); + const lg = try c.tensor(5); + const cnt = try c.tensor(6); + const group = try c.boolean(19, false); + const gmem = try c.tensor(17); + if (try dim(x, 1) != ops.rg.D or try dim(w, 1) != ops.rg.D) return error.WrongArgument; + const a: ops.rg.Args = .{ + .x = x.ptr, + .xs = x.stride[0], + .w = w.ptr, + .bias = try c.ptr(2), + .pick = try c.ptr(3), + .wts = try c.ptr(4), + .lg = lg.ptr, + .cnt = cnt.ptr, + .R = @intCast(try dim(x, 0)), + .E = @intCast(try dim(w, 0)), + .K = @intCast(try c.int(7)), + .slots = @intCast(try c.int(8)), + .select = @intFromBool(try c.boolean(13, false)), + .scale = @floatCast(try c.float(9)), + .kit = @intFromBool(try c.boolean(14, false)), + .group = @intFromBool(group), + .gid = if (group) try c.ptr(15) else 0, + .gcnt = if (group) try c.ptr(16) else 0, + .gmem = if (group) gmem.ptr else 0, + .GE = @intCast(try c.int(18)), + .maxm = if (group) @intCast(try dim(gmem, 1)) else 0, + .prune = try pruneArg(c, 20), + }; + try o.route(a, try c.u(10), try c.u(11), try c.u(12), @intCast(numel(lg)), @intCast(numel(cnt))); + return true; +} + +// tf_dsv41_pfdense_v1 (pfdense.cpp) +pub fn pfdenseExt(o: ops.Ops, func: []const u8, args: []const Arg, kwargs: []const Kw) !bool { + if (is(func, "gemm")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "xh", "T", "stride_k", "stride_nb", "svh", "bias", "out", "K2", "lanes", "cfg", "group" } }; + const xh = try c.tensor(0); + const T = try c.tensor(1); + const out = try c.tensor(6); + const ot: ops.pfd.OutType = switch (out.dtype) { + .f32 => .f32, + .bf16 => .bf16, + .f16 => .f16, + else => return error.WrongArgument, + }; + const a: ops.pfd.Args = .{ .xh = xh.ptr, .T = T.ptr, .stride_k = try c.int(2), .stride_nb = try c.int(3), .svh = try c.ptr(4), .bias = try c.ptrOrNull(5), .out = out.ptr, .o_stride = out.stride[0], .M = @intCast(try dim(xh, 0)), .K = @intCast(try dim(xh, 1)), .N = @intCast(try dim(out, 1)), .group = @intCast(try c.int(10)), .out_type = @backingInt(ot) }; + try o.pfdGemm(a, try c.u(7), try c.boolean(8, false), try c.usz(9), @intCast(numel(T))); + return true; + } + if (is(func, "dequant")) { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "T", "W", "stride_k", "stride_nb", "K2", "lanes" } }; + const w = try c.tensor(1); + try o.pfdDequant(try c.ptr(0), try c.int(2), try c.int(3), w.ptr, try dim(w, 0), try dim(w, 1), try c.u(4), try c.boolean(5, false)); + return true; + } + return false; +} + +// tf_dsv41_l2pace_v1.paced / tf_dsv41_l2pf_v1.segments: L2 prefetch only (no outputs; replayed for timing parity) +pub fn paced(o: ops.Ops, args: []const Arg, kwargs: []const Kw) !bool { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "table", "ctas", "chunk", "ns_per_piece", "delay_ns" } }; + const t = try c.tensor(0); + const n = try dim(t, 0); + if (n == 0) return true; + try o.l2Paced(t.ptr, n, try c.usz(1), try c.u(2), @intCast(try c.int(3)), @intCast(try c.int(4))); + return true; +} + +pub fn segments(o: ops.Ops, args: []const Arg, kwargs: []const Kw) !bool { + const c: Call = .{ .args = args, .kwargs = kwargs, .names = &.{ "table", "grid", "threads", "chunk", "bulk" } }; + const t = try c.tensor(0); + const n = try dim(t, 0); + if (n == 0) return true; + try o.l2Segments(t.ptr, n, try c.usz(1), try c.usz(2), try c.u(3), try c.boolean(4, false)); + return true; +} diff --git a/zig/src/families/deepseek_v41/kernels_shapes.zig b/zig/src/families/deepseek_v41/kernels_shapes.zig new file mode 100644 index 0000000000..6c7bd6c135 --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_shapes.zig @@ -0,0 +1,407 @@ +//! DeepSeek-V4.1 kernel shapes: attention, top-k, mHC, router and prefill GEMM args, symbols and checks. + +const std = @import("std"); + +pub const Error = error{ Shape, Unsupported }; + +// -- attn_cuda.cu: dsv41_attn ----------------------------------------------------------------------------------- + +pub const attn = struct { + pub const NCH = 5; + pub const BMQ = 16; + pub const WINDOW = 128; + pub const NT = 256; + pub const SMEM: u32 = 98816; + + pub const Args = extern struct { + q: u64 = 0, + cv: u64 = 0, + csc: u64 = 0, + cvs: i64 = 0, + css: i64 = 0, + tok: u64 = 0, + ts: i64 = 0, + cnt: u64 = 0, + sv: u64 = 0, + ssc: u64 = 0, + svs: i64 = 0, + sss: i64 = 0, + lo: u64 = 0, + hi: u64 = 0, + pos: u64 = 0, + sl: u64 = 0, + pt: u64 = 0, + pts: i64 = 0, + psh: c_int = 0, + sink: u64 = 0, + cs: u64 = 0, + csst: i64 = 0, + po: u64 = 0, + pm: u64 = 0, + pl: u64 = 0, + out: u64 = 0, + ticket: u64 = 0, + R: c_int = 0, + H: c_int = 0, + ring: c_int = 0, + has_comp: c_int = 0, + qrope: c_int = 0, // R1c (split_kernel only): q arrives un-rotated and is RoPE'd as it loads + }; + + comptime { + std.debug.assert(@offsetOf(Args, "psh") == 144 and @offsetOf(Args, "sink") == 152 and @offsetOf(Args, "qrope") == 232 and @sizeOf(Args) == 240); + } + + pub fn symbol(comptime kw: u32, comptime rows: bool, comptime hi: bool) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_attn11attn_kernelILi{d}ELb{d}ELb{d}EEEvNS_4ArgsE", .{ kw, @intFromBool(rows), @intFromBool(hi) }); + } + + /// R1 split_kernel: P CTAs a chunk, a cluster of P (P 4); the same bits as attn_kernel. + pub fn splitSymbol(comptime kw: u32, comptime rows: bool, comptime hi: bool, comptime p: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_attn12split_kernelILi{d}ELb{d}ELb{d}ELi{d}EEEvNS_4ArgsE", .{ kw, @intFromBool(rows), @intFromBool(hi), p }); + } + + /// Split

::SMEM: P 1 the same as attn_kernel's, P 4 the S rows' B pairs beside the PV slice. + pub fn splitSmem(p: u32) u32 { + return if (p == 4) 100864 else SMEM; + } + + /// The ticket's ranks a (row, head tile) once any launch splits: the ticket buffer must hold R H / 16 * MAXP. + pub const MAXP = 4; + + /// attn_cuda.split(R) (TF_DSV41_ATTN_SPLIT on): P 4 up to `split_rows` rows (2), P 1 above. + pub fn splitFor(rows: usize, split_rows: usize) u32 { + return if (rows <= split_rows) 4 else 1; + } + + /// dispatch(): the instance index (kw 4 | rows | has_hi), after its checks. + pub fn instance(a: Args, kw: u32, rows: bool, has_hi: bool) Error!usize { + if (a.R < 1 or @rem(a.H, BMQ) != 0 or a.ring < WINDOW or (a.ring & (a.ring - 1)) != 0) return error.Shape; + return @as(usize, if (kw == 4) 4 else 0) + @as(usize, if (rows) 2 else 0) + @intFromBool(has_hi); + } +}; + +// topk_cuda.cu: dsv41_topk +pub const topk = struct { + pub const NT = 512; + pub const MAXCL = 8; + pub const MAXEPT = 31; + + pub const Job = extern struct { s: u64 = 0, ss: i64 = 0, cand: u64 = 0, cs: i64 = 0, out: u64 = 0, os: i64 = 0, cnt: u64 = 0, nk: c_int = 0, k: c_int = 0, mode: c_int = 0, ept: c_int = 0 }; + pub const Args = extern struct { job: [2]Job, pos: u64 = 0, pos64: c_int = 0, ratio: c_int = 0, bs: c_int = 0 }; + + comptime { + std.debug.assert(@sizeOf(Job) == 72 and @offsetOf(Args, "pos") == 144 and @sizeOf(Args) == 168); + } + + pub const symbol = "_ZN10dsv41_topk11topk_kernelENS_4ArgsE"; + /// topk_b.cu (ours, TF_DSV41_INDEX_BOUND): topk_kernel with each row stopped at its own visible end, the same Args + pub const symbol_b = "_ZN10dsv41_topk13topk_b_kernelENS_4ArgsE"; + + pub fn smem(ept: usize) u32 { + return @intCast(NT * ept * 4); + } +}; + +// -- mhc_cuda.cu: dsv41_mhc (decode boundary, <= 16 rows) ------------------------------------------------------- + +pub const mhc = struct { + pub const NB = 40; + pub const NT = 128; + pub const MAXR = 16; + pub const ROWB = 128; + /// coef_kernel is a CTA a row: TF_DSV41_MHC_PFDEC launches it past MAXR, up to a decode window's 64 rows + pub const COEF_ROWS = 64; + + pub const Args = extern struct { + x: u64 = 0, + xs: i64 = 0, + xout: u64 = 0, + g: u64 = 0, + gr: i64 = 0, + world: c_int = 1, + gbf16: c_int = 0, + post: u64 = 0, + comb: u64 = 0, + pre: u64 = 0, + @"fn": u64 = 0, + base: u64 = 0, + scale: u64 = 0, + part: u64 = 0, + c: u64 = 0, + tap: u64 = 0, + ts: i64 = 0, + nw: u64 = 0, + nwbf16: c_int = 0, + out: u64 = 0, + opre: u64 = 0, + opost: u64 = 0, + ocomb: u64 = 0, + cnt: u64 = 0, + R: c_int = 0, + eps: f32 = 0, + hc_eps: f32 = 0, + post_alpha: f32 = 0, + iters: c_int = 0, + spin: i64 = -1, // R1 tail: SM cycles a CTA waits for the last arrival (the binding's -1: today's finish) + @"defer": c_int = 0, // R1 tail: the coefficient items run in coef_kernel on a side stream + }; + + comptime { + std.debug.assert(@offsetOf(Args, "post") == 48 and @offsetOf(Args, "nwbf16") == 136 and @offsetOf(Args, "R") == 184 and @offsetOf(Args, "spin") == 208 and @sizeOf(Args) == 224); + } + + pub const coef_symbol = "_ZN9dsv41_mhc11coef_kernelENS_4ArgsE"; + + /// (POST, COLL, MIX, ALL4, FN32) of dispatch's cases, by mode * 2 + fn32 (mode 3 ignores fn32). + pub const instances = [_]struct { post: bool, coll: u32, mix: bool, all4: bool, fn32: bool }{ + .{ .post = true, .coll = 2, .mix = true, .all4 = false, .fn32 = false }, + .{ .post = true, .coll = 2, .mix = true, .all4 = false, .fn32 = true }, + .{ .post = false, .coll = 1, .mix = true, .all4 = false, .fn32 = false }, + .{ .post = false, .coll = 1, .mix = true, .all4 = false, .fn32 = true }, + .{ .post = false, .coll = 2, .mix = true, .all4 = false, .fn32 = false }, + .{ .post = false, .coll = 2, .mix = true, .all4 = false, .fn32 = true }, + .{ .post = true, .coll = 2, .mix = false, .all4 = true, .fn32 = false }, + }; + + /// boundary_kernel (TAIL: R1's parallel tail, default false: still mangled) + pub fn symbol(comptime i: usize, comptime tail: bool) [:0]const u8 { + const s = instances[i]; + return std.fmt.comptimePrint("_ZN9dsv41_mhc15boundary_kernelILb{d}ELi{d}ELb{d}ELb{d}ELb{d}ELb{d}EEEvNS_4ArgsE", .{ @intFromBool(s.post), s.coll, @intFromBool(s.mix), @intFromBool(s.all4), @intFromBool(s.fn32), @intFromBool(tail) }); + } + + /// dispatch(): mode 0 boundary, 1 site from stream 0, 2 site with a carried pre-mix, 3 final. + pub fn instance(a: Args, mode: u32, fn32: bool) Error!usize { + if (a.R < 1 or a.R > MAXR or a.world < 1 or a.world > 8) return error.Shape; + if (mode == 3) return 6; + if (mode > 2) return error.Unsupported; + return mode * 2 + @intFromBool(fn32); + } +}; + +// mhc_pf.cu: dsv41_mhc_pf (prefill site, 16 rows a CTA) +pub const mhc_pf = struct { + pub const NB = 40; + pub const NT = 128; + pub const BM = 16; + pub const MAXW = 8; + + pub const Args = extern struct { + x: u64 = 0, + xs: i64 = 0, + xout: u64 = 0, + g: u64 = 0, + gr: i64 = 0, + world: c_int = 1, + gbf16: c_int = 0, + post: u64 = 0, + comb: u64 = 0, + pre: u64 = 0, + @"fn": u64 = 0, + part: u64 = 0, + c: u64 = 0, + tap: u64 = 0, + ts: i64 = 0, + R: c_int = 0, + }; + + comptime { + std.debug.assert(@offsetOf(Args, "R") == 112 and @sizeOf(Args) == 120); + } + + /// Lay::BYTES: staged streams, the collapsed rows and fn. + pub fn smem(fn32: bool) u32 { + const pb = 128 + 8; + const pf = 128 + 4; + const xs = 4 * BM * pb * 2; + const cs = BM * pb * 2; + const fnb: u32 = if (fn32) 4 * 24 * pf * 4 else 4 * 24 * pb * 2; + return xs + cs + fnb; + } + + /// (POST, COLL, FN32, TAP) by dispatch's case (mode * 4 + fn32 * 2 + tap). + pub const instances = [_]struct { case: u32, post: bool, coll: u32, fn32: bool, tap: bool }{ + .{ .case = 0, .post = true, .coll = 2, .fn32 = false, .tap = false }, + .{ .case = 1, .post = true, .coll = 2, .fn32 = false, .tap = true }, + .{ .case = 2, .post = true, .coll = 2, .fn32 = true, .tap = false }, + .{ .case = 3, .post = true, .coll = 2, .fn32 = true, .tap = true }, + .{ .case = 4, .post = false, .coll = 1, .fn32 = false, .tap = false }, + .{ .case = 6, .post = false, .coll = 1, .fn32 = true, .tap = false }, + .{ .case = 8, .post = false, .coll = 2, .fn32 = false, .tap = false }, + .{ .case = 10, .post = false, .coll = 2, .fn32 = true, .tap = false }, + }; + + pub fn symbol(comptime i: usize) [:0]const u8 { + const s = instances[i]; + return std.fmt.comptimePrint("_ZN12dsv41_mhc_pf11site_kernelILb{d}ELi{d}ELb{d}ELb{d}EEEvNS_4ArgsE", .{ @intFromBool(s.post), s.coll, @intFromBool(s.fn32), @intFromBool(s.tap) }); + } + + pub fn instance(a: Args, mode: u32, fn32: bool, tap: bool) Error!usize { + if (a.R < 1 or a.world < 1 or a.world > MAXW or (tap and mode != 0)) return error.Shape; + const case = mode * 4 + @as(u32, if (fn32) 2 else 0) + @intFromBool(tap); + for (instances, 0..) |s, i| if (s.case == case) return i; + return error.Unsupported; + } +}; + +// -- router_gemv.cu: dsv41_rg ----------------------------------------------------------------------------------- + +pub const rg = struct { + pub const D = 5120; + pub const EPL = 16; + pub const KMAX = 8; + pub const GE_MAX = 512; + pub const rts = [_]u32{ 1, 2, 4, 8, 16 }; + pub const wide = [_][2]u32{ .{ 8, 1 }, .{ 4, 1 }, .{ 8, 2 }, .{ 4, 2 } }; + pub const narrow = [_]u32{ 1, 2, 4, 8 }; + + pub const Args = extern struct { + x: u64 = 0, + xs: i64 = 0, + w: u64 = 0, + bias: u64 = 0, + pick: u64 = 0, + wts: u64 = 0, + lg: u64 = 0, + cnt: u64 = 0, + R: c_int = 0, + E: c_int = 0, + K: c_int = 0, + slots: c_int = 0, + select: c_int = 0, + scale: f32 = 0, + kit: c_int = 0, + gid: u64 = 0, + gcnt: u64 = 0, + gmem: u64 = 0, + GE: c_int = 0, + maxm: c_int = 0, + group: c_int = 0, + prune: Prune = .{}, // R1c: the prune rule folded into the selection (on 0: off) + }; + + /// router_gemv.cu Prune: [min_w, topp, scale, min_k, orig, drop, dup] (prune.fold_args) in fp32 / int, on = 1. + pub const Prune = extern struct { min_w: f32 = 0, topp: f32 = 0, scale: f32 = 0, min_k: c_int = 0, orig: c_int = 0, drop: c_int = 0, dup: c_int = 0, on: c_int = 0 }; + + /// The binding's prune vector (7 doubles, empty: off) as the kernel's struct (fp32 thresholds, as it rounds them). + pub fn pruneOf(v: []const f64) error{Shape}!Prune { + if (v.len == 0) return .{}; + if (v.len != 7) return error.Shape; + return .{ .min_w = @floatCast(v[0]), .topp = @floatCast(v[1]), .scale = @floatCast(v[2]), .min_k = @intFromFloat(v[3]), .orig = @intFromFloat(v[4]), .drop = @intFromFloat(v[5]), .dup = @intFromFloat(v[6]), .on = 1 }; + } + + comptime { + std.debug.assert(@sizeOf(Prune) == 32); + std.debug.assert(@offsetOf(Args, "gid") == 96 and @offsetOf(Args, "group") == 128 and @offsetOf(Args, "prune") == 132 and @sizeOf(Args) == 168); + } + + pub fn gemvSymbol(comptime rt: u32, comptime nw: u32, comptime ew: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN8dsv41_rg11gemv_kernelILi{d}ELi{d}ELi{d}EEEvNS_4ArgsE", .{ rt, nw, ew }); + } + + pub fn narrowSymbol(comptime rt: u32, comptime nw: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN8dsv41_rg13narrow_kernelILi{d}ELi{d}EEEvNS_4ArgsE", .{ rt, nw }); + } + + /// dsv41_rg_route_cuda's checks (row stride / alignment of x and the buffer sizes are the caller's). + pub fn check(a: Args, rt: u32, nw: u32, ew: u32, lg_len: usize, cnt_len: usize) Error!void { + const E: usize = @intCast(a.E); + const R: usize = @intCast(a.R); + if (E < 1 or E > 32 * EPL or a.K < 1 or a.K > KMAX or a.K > a.E or a.slots < a.K or a.slots > a.K + 1) return error.Shape; + if (lg_len < R * E) return error.Shape; + if (a.select != 0 and cnt_len < (R + rt - 1) / rt) return error.Shape; + const is_wide = (nw == 4 or nw == 8) and (ew == 1 or ew == 2); + const is_narrow = ew == 0 and (nw == 1 or nw == 2 or nw == 4 or nw == 8); + if (std.mem.indexOfScalar(u32, &rts, rt) == null or !(is_wide or is_narrow)) return error.Unsupported; + if (a.group != 0 and (a.select == 0 or !is_narrow or R > rt or a.GE < 1 or a.GE > GE_MAX)) return error.Shape; + } +}; + +// -- pfdense.cuh: dsv41_pfd (fused prefill dense GEMM, one fatbin a width) -------------------------------------- + +pub const pfd = struct { + pub const k2s = [_]u32{ 8, 10, 12, 16 }; + /// DSV41_PFD_CFGS: (BM, WM, WN, KS, NST) by id. + pub const cfgs = [_][5]u32{ .{ 128, 2, 4, 2, 3 }, .{ 128, 2, 4, 2, 4 }, .{ 128, 4, 2, 2, 3 }, .{ 64, 2, 2, 2, 4 }, .{ 64, 1, 4, 2, 4 }, .{ 128, 2, 4, 2, 2 }, .{ 32, 1, 4, 2, 4 }, .{ 256, 4, 2, 2, 3 } }; + pub const OutType = enum(c_int) { f16 = 0, bf16 = 1, f32 = 2 }; + + pub const Args = extern struct { + xh: u64 = 0, + T: u64 = 0, + stride_k: i64 = 0, + stride_nb: i64 = 0, + svh: u64 = 0, + bias: u64 = 0, + out: u64 = 0, + o_stride: i64 = 0, + M: c_int = 0, + K: c_int = 0, + N: c_int = 0, + group: c_int = 0, + out_type: c_int = 0, + }; + + comptime { + std.debug.assert(@offsetOf(Args, "M") == 64 and @sizeOf(Args) == 88); + } + + pub fn threads(cfg: usize) u32 { + return cfgs[cfg][1] * cfgs[cfg][2] * 32; + } + + /// Lay::SMEM (the layout does not depend on LANES). + pub fn smem(cfg: usize, k2: u32) u32 { + const bm, const wm, _, const ks, const nst = cfgs[cfg]; + const a_bytes = bm * (16 * ks) * 2; + const w_bytes = ks * 8 * (4 * k2) * 4; + const ring = nst * (a_bytes + w_bytes); + const main = ring + 2 * (ks * 8 * 32 * 16); + const epi = (bm / wm) * 128 * 2 * 2; + return @max(main, epi); + } + + pub fn symbol(comptime k2: u32, comptime lanes: bool, comptime cfg: usize) [:0]const u8 { + const c = cfgs[cfg]; + return std.fmt.comptimePrint("_ZN9dsv41_pfd10pfd_kernelILi{d}ELb{d}ENS_3CfgILi{d}ELi{d}ELi{d}ELi{d}ELi{d}EEEEEvNS_4ArgsE", .{ k2, @intFromBool(lanes), c[0], c[1], c[2], c[3], c[4] }); + } + + pub fn dequantSymbol(comptime k2: u32, comptime lanes: bool) [:0]const u8 { + return std.fmt.comptimePrint("_ZN9dsv41_pfd18pfd_dequant_kernelILi{d}ELb{d}EEEvPKjxxP6__halfi", .{ k2, @intFromBool(lanes) }); + } + + /// pfdense.cpp gemm's checks on the shapes and strides (dtype / alignment are the caller's). + pub fn check(a: Args, k2: u32, lanes: bool, cfg: usize, words: usize) Error!void { + if (std.mem.indexOfScalar(u32, &k2s, k2) == null or (lanes and k2 == 16)) return error.Unsupported; + if (cfg >= cfgs.len) return error.Unsupported; + const ks: i64 = cfgs[cfg][3]; + const K: i64 = a.K; + const N: i64 = a.N; + if (a.M < 1 or @rem(K, 128) != 0 or K < 128 or @rem(K, 16 * ks) != 0) return error.Shape; + if (@rem(N, 128) != 0 or N < 128 or @rem(a.o_stride, 2) != 0 or a.o_stride < N) return error.Shape; + const tw: i64 = 4 * @as(i64, k2); + if (@rem(a.stride_k, 4) != 0 or @rem(a.stride_nb, 4) != 0 or a.stride_k < 8 * tw or a.stride_nb < 8 * tw) return error.Shape; + const last = (@divTrunc(K, 16) - 1) * a.stride_k + (@divTrunc(N, 128) - 1) * a.stride_nb + 8 * tw; + if (last > @as(i64, @intCast(words))) return error.Shape; + if (a.group < 1 or a.group > 1024) return error.Shape; + } + + /// pfdense.heuristic (no sweep table): the largest row tile whose last wave is not mostly empty. + pub fn heuristic(k: usize, n: usize, k2: u32, m: usize, sms: usize) [2]usize { + const nb = n / 128; + var best: ?usize = null; + var best_eff: f64 = -1.0; + for ([_]usize{ 0, 3, 6 }) |cfg| { + if (k % (16 * cfgs[cfg][3]) != 0) continue; + const bm = cfgs[cfg][0]; + const blocks: usize = if (2 * (smem(cfg, k2) + 1024) <= 102400) 2 else 1; + const ctas = ((m + bm - 1) / bm) * nb; + const slots = sms * blocks; + const eff = @as(f64, @floatFromInt(ctas)) / @as(f64, @floatFromInt(((ctas + slots - 1) / slots) * slots)); + if (eff > best_eff + 0.10) { + best = cfg; + best_eff = eff; + } + } + return .{ best orelse 0, 8 }; + } +}; diff --git a/zig/src/families/deepseek_v41/kernels_x3gm.zig b/zig/src/families/deepseek_v41/kernels_x3gm.zig new file mode 100644 index 0000000000..0578fc17ad --- /dev/null +++ b/zig/src/families/deepseek_v41/kernels_x3gm.zig @@ -0,0 +1,354 @@ +//! DeepSeek-V4.1's x3gm routed-expert GEMM: width and tile tables, v2/v3 shapes and tuning, args and launches. + +const std = @import("std"); +const cuda = @import("cuda"); +const exl3 = @import("kernels_exl3.zig"); + +const Ops = exl3.Ops; +const Input = exl3.Input; +const Error = exl3.Error; +const int = exl3.int; +const go = exl3.go; + +/// x3gm's widths (x3gm.K2S) and tile tables (MTL, NG, KS, NSA, NGR) by configuration id (x3gm.cu GM_GU* / GM_DN*). +pub const gm_k2s = [_]u32{ 3, 4, 5, 6, 7, 8, 10 }; +pub const gu_tiles = [_][5]u32{ .{ 2, 8, 4, 3, 4 }, .{ 2, 8, 2, 4, 4 }, .{ 1, 16, 2, 4, 4 }, .{ 1, 16, 4, 3, 4 }, .{ 2, 16, 2, 3, 4 }, .{ 2, 8, 2, 6, 4 } }; +pub const dn_tiles = [_][5]u32{ .{ 2, 8, 4, 4, 8 }, .{ 1, 8, 4, 4, 8 }, .{ 1, 16, 4, 4, 8 }, .{ 1, 16, 2, 6, 8 }, .{ 2, 16, 4, 3, 16 } }; +/// sm_121's opt-in dynamic shared memory a CTA (x3gm.SMEM_MAX). +pub const gm_smem_max: usize = 101376; + +pub fn gmK2Index(k2: u32) ?usize { + return std.mem.indexOfScalar(u32, &gm_k2s, k2); +} + +/// x3gm.cu Cfg: threads and dynamic shared memory of one instance. +pub const GmShape = struct { + threads: u32, + smem: u32, + + pub fn of(mats: u32, xm: u32, k2: u32, t: [5]u32) GmShape { + const mtl, const ng, const ks, const nsa, const ngr = t; + const bm = ng * 8; + const w = mats * (8 / mtl); + const a_bytes: usize = @as(usize, xm) * bm * (ks * 16) * 2; + const w_bytes: usize = @as(usize, mats) * ks * 8 * (4 * k2) * 4; + const ring = nsa * (a_bytes + w_bytes); + const epi: usize = @as(usize, mats) * (ngr * 8) * 132 * 4; + return .{ .threads = w * 32, .smem = @intCast(@max(ring, epi)) }; + } +}; + +/// x3gm v2 (gm2_kernel): the tiles it is built at (V2_GU / V2_DN). +pub const v2_gu = [_]usize{ 0, 1 }; +/// gate / up's 128-member tile (gu_tiles[2], NG 16) that x3gm3.cu builds gm2_kernel at, two inputs only +pub const gu2_cfg: usize = 2; +pub const v2_dn = [_]usize{0}; + +/// x3gm.cu Smem2: the widest width's dynamic shared memory; rings cut (Ring2) to keep 2 CTAs an SM. +pub fn gm2Shape(mats: u32, xm: u32, t: [5]u32) GmShape { + const mtl, const ng, const ks, const nsa, const ngr = t; + const four = GmShape.of(mats, xm, 4, t); + const two_ctas = 2 * (@as(usize, four.smem) + 1536) <= 102400 and four.threads <= 256 and mtl * ng <= 16; // Cfg::MINB == 2 + var smem: u32 = 0; + for (gm_k2s) |k2| { + const stage: usize = @as(usize, xm) * (ng * 8) * (ks * 16) * 2 + @as(usize, mats) * ks * 8 * (4 * k2) * 4; + var ring: u32 = nsa; + if (two_ctas) { + while (ring > 2 and 2 * (stage * ring + 1536) > 102400) ring -= 1; + } + smem = @max(smem, GmShape.of(mats, xm, k2, .{ mtl, ng, ks, ring, ngr }).smem); + } + return .{ .threads = four.threads, .smem = smem }; +} + +/// x3gm v3 (x3gm3.cu's gm3_kernel): v2 plus next-item read-ahead and LEAN reads; variants in its instance order. +pub const V3 = struct { mats: u32, xm: u32, t: [5]u32, lean: u32 }; +pub const v3_variants = [_]V3{ + .{ .mats = 2, .xm = 2, .t = .{ 2, 8, 2, 4, 4 }, .lean = 0 }, // v2's GU1 + claim-ahead + .{ .mats = 2, .xm = 2, .t = .{ 2, 8, 2, 4, 4 }, .lean = 1 }, + .{ .mats = 2, .xm = 2, .t = .{ 1, 8, 2, 4, 4 }, .lean = 1 }, // 16-warp CTAs + .{ .mats = 2, .xm = 2, .t = .{ 1, 8, 4, 3, 4 }, .lean = 1 }, + .{ .mats = 2, .xm = 1, .t = .{ 2, 8, 4, 3, 4 }, .lean = 1 }, // v2's GU0 (one input) + claim-ahead + LEAN + .{ .mats = 2, .xm = 1, .t = .{ 1, 8, 4, 3, 4 }, .lean = 1 }, + .{ .mats = 1, .xm = 1, .t = .{ 2, 8, 4, 4, 8 }, .lean = 0 }, // v2's DN0 + claim-ahead + .{ .mats = 1, .xm = 1, .t = .{ 2, 8, 4, 4, 8 }, .lean = 1 }, + .{ .mats = 1, .xm = 1, .t = .{ 1, 8, 4, 4, 8 }, .lean = 0 }, // 8-warp CTAs: 16 warps an SM + .{ .mats = 1, .xm = 1, .t = .{ 1, 8, 4, 4, 8 }, .lean = 1 }, + .{ .mats = 1, .xm = 1, .t = .{ 1, 8, 2, 6, 8 }, .lean = 1 }, +}; + +pub fn gm3Symbol(comptime v: V3) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_x3gm10gm3_kernelILi{d}ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}EEEvNS_4ArgsE", .{ v.mats, v.xm, v.t[0], v.t[1], v.t[2], v.t[3], v.t[4], v.lean }); +} + +/// v3's variants for one projection: gate / up at one or two inputs (xm), or down. +pub fn v3Of(mats: u32, xm: u32, out: []usize) []usize { + var n: usize = 0; + for (v3_variants, 0..) |v, i| if (v.mats == mats and v.xm == xm) { + out[n] = i; + n += 1; + }; + return out[0..n]; +} + +/// TF_DSV41_GM_V3=1: prefill's x3gm through v3 (v3_default's variants); unset or 0: v2 (default; v3 is bit-identical). +pub fn gmV3() bool { + const v = std.c.getenv("TF_DSV41_GM_V3") orelse return false; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// v3's default variant of a projection (from per-shape timing). +pub const v3_default: struct { gu: [2]usize, dn: usize } = .{ .gu = .{ 4, 1 }, .dn = 9 }; + +/// x3gm.tuned2: (gate/up cfg, down cfg) of a v2 block: forced ids where built, else 0 / 1 by rotated inputs; down 0. +pub fn gmTuned2(force_gu: ?usize, force_dn: ?usize, shx: bool) [2]usize { + const gu = if (force_gu) |f| (if (std.mem.indexOfScalar(usize, &v2_gu, f) != null) f else null) else null; + const dn = if (force_dn) |f| (if (std.mem.indexOfScalar(usize, &v2_dn, f) != null) f else null) else null; + return .{ gu orelse (if (shx) 0 else 1), dn orelse 0 }; +} + +/// The host check before a v2 launch: every expert's width is one gm2_kernel dispatches. +pub fn gm2WidthsOk(widths: []const u32) bool { + for (widths) |w| if (gmK2Index(w) == null) return false; + return true; +} + +pub fn gm2GuSymbol(comptime xm: u32, comptime cfg: usize) [:0]const u8 { + const t = gu_tiles[cfg]; + return std.fmt.comptimePrint("_ZN10dsv41_x3gm10gm2_kernelILi2ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}EEEvNS_4ArgsE", .{ xm, t[0], t[1], t[2], t[3], t[4] }); +} + +pub fn gm2DnSymbol(comptime cfg: usize) [:0]const u8 { + const t = dn_tiles[cfg]; + return std.fmt.comptimePrint("_ZN10dsv41_x3gm10gm2_kernelILi1ELi1ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}EEEvNS_4ArgsE", .{ t[0], t[1], t[2], t[3], t[4] }); +} + +/// x3gm.smem_bytes: Cfg::SMEM plus the static rows / item words. +pub fn gmSmemBytes(t: [5]u32, mats: u32, xm: u32, k2: u32) usize { + return @as(usize, GmShape.of(mats, xm, k2, t).smem) + 4 * 8 * t[1] + 4; +} + +pub const GmKind = enum { gu, dn }; + +/// x3gm_plan.cu takes E < 512 experts. +pub const plan_max_experts = 512; + +/// x3gm.plan's table length: P / bm + E + 1 passes at most. +pub fn planPasses(pairs: usize, experts: usize, bm: usize) usize { + return pairs / bm + experts + 1; +} + +/// x3gm.fits: whether a configuration builds a CTA at this width (gate/up: `xm` rotated inputs) on sm_121. +pub fn gmFits(kind: GmKind, cfg: usize, k2: u32, xm: u32) bool { + return switch (kind) { + .gu => cfg < gu_tiles.len and gmSmemBytes(gu_tiles[cfg], 2, xm, k2) <= gm_smem_max, + .dn => cfg < dn_tiles.len and gmSmemBytes(dn_tiles[cfg], 1, 1, k2) <= gm_smem_max, + }; +} + +/// x3gm.bm_of: members a pass of a configuration. +pub fn gmMembers(kind: GmKind, cfg: usize) u32 { + return 8 * (if (kind == .gu) gu_tiles[cfg] else dn_tiles[cfg])[1]; +} + +/// x3gm.TUNE row: from `members` an expert on average, gate/up `gu` (null: by shx; else one / two inputs), down `dn`. +pub const TuneRow = struct { members: f64, gu: ?[2]u32, dn: u32 }; +pub const tune_default = [_]TuneRow{ .{ .members = 0, .gu = null, .dn = 0 }, .{ .members = 48, .gu = .{ 3, 2 }, .dn = 2 } }; + +/// x3gm.tuned: (gate/up cfg, down cfg) of a block of `rows` rows; `force_gu` / `force_dn` are TF_DSV41_GM_CFG's ids. +pub fn gmTuned(table: []const TuneRow, force_gu: ?usize, force_dn: ?usize, rows: usize, slots: usize, experts: usize, shx: bool, k2gu: u32, k2d: u32) [2]usize { + const m = @as(f64, @floatFromInt(rows * slots)) / @as(f64, @floatFromInt(@max(1, experts))); + var ent = table[0]; + for (table) |r| { + if (m > r.members or r.members == 0) ent = r; + } + var gu: usize = if (ent.gu) |p| (if (shx) p[0] else p[1]) else (if (shx) 0 else 1); + if (force_gu) |f| { + gu = f; + } else if (!gmFits(.gu, gu, k2gu, if (shx) 1 else 2)) { + gu = if (shx) 0 else 1; + } + var dn: usize = ent.dn; + if (force_dn) |f| { + dn = f; + } else if (!gmFits(.dn, dn, k2d, 1)) { + dn = 0; + } + return .{ gu, dn }; +} + +pub fn gmGuSymbol(comptime xm: u32, comptime k2: u32, comptime cfg: usize) [:0]const u8 { + const t = gu_tiles[cfg]; + return std.fmt.comptimePrint("_ZN10dsv41_x3gm9gm_kernelILi2ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}EEEvNS_4ArgsE", .{ xm, k2, t[0], t[1], t[2], t[3], t[4] }); +} + +pub fn gmDnSymbol(comptime k2: u32, comptime cfg: usize) [:0]const u8 { + const t = dn_tiles[cfg]; + return std.fmt.comptimePrint("_ZN10dsv41_x3gm9gm_kernelILi1ELi1ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}EEEvNS_4ArgsE", .{ k2, t[0], t[1], t[2], t[3], t[4] }); +} + +pub fn gmDequantSymbol(comptime k2: u32) [:0]const u8 { + return std.fmt.comptimePrint("_ZN10dsv41_x3gm14dequant_kernelILi{d}EEEvPKjP6__halfi", .{k2}); +} + +/// x3gm.cu's Args, passed by value. +pub const GmArgs = extern struct { + x0: u64 = 0, + x1: u64 = 0, + t0: u64 = 0, + t1: u64 = 0, + tp0: u64 = 0, + tp1: u64 = 0, + k2e: u64 = 0, // int32 [E] each expert's K2 (gm2_kernel; gm_kernel ignores it) + order: u64 = 0, + pe: u64 = 0, + poff: u64 = 0, + pcnt: u64 = 0, + npass: u64 = 0, + ticket: u64 = 0, + sv0: u64 = 0, + sv1: u64 = 0, + sd: u64 = 0, + xd: u64 = 0, + y: u64 = 0, + K: c_int = 0, + N: c_int = 0, + limit: f32 = 0, +}; + +comptime { + std.debug.assert(@sizeOf(GmArgs) == 160 and @offsetOf(GmArgs, "k2e") == 48 and @offsetOf(GmArgs, "K") == 144 and @offsetOf(GmArgs, "limit") == 152); +} + +/// x3gm's per-launch checks (launch<>): K in whole stages, N in 128-column Hadamard blocks. +pub fn gmCheck(kind: GmKind, cfg: usize, K: usize, N: usize) Error!void { + const ks = (if (kind == .gu) gu_tiles[cfg] else dn_tiles[cfg])[2]; + if ((K / 16) % ks != 0 or N % 128 != 0) return error.Shape; +} + +/// dsv41_x3gm_rot_cuda: out_m[p] = fp16((x[p / slots] * suh_m[e]) H / sqrt 128) for `mats` matrices. +pub fn gmRot(o: Ops, input: Input, x: u64, x_stride: usize, pick: u64, suh0: u64, suh1: u64, out0: u64, out1: u64, K: usize, slots: usize, pairs: usize, mats: usize) !void { + if (mats != 1 and mats != 2) return error.Unsupported; + const items = pairs * (K / 128); + var a: cuda.Args = .{}; + a.add(x); + a.add(int(x_stride)); + for ([_]u64{ pick, suh0, suh1, out0, out1 }) |v| a.add(v); + for ([_]usize{ K, slots, items }) |v| a.add(int(v)); + try go(o.f.gm_rot[@backingInt(input)][mats - 1], o.s, .{ (items + 7) / 8, 1, 1 }, 256, 0, false, &a); +} + +/// dsv41_x3gm_gateup_cuda: Xd of every pass's members; `tables`: tg / tu are int64 pointer tables (ragged). +pub fn gmGateup(o: Ops, xg: u64, xu: u64, tg: u64, tu: u64, order: u64, pe: u64, poff: u64, pcnt: u64, npass: u64, ticket: u64, svh_g: u64, svh_u: u64, suh_d: u64, xd: u64, K: usize, N: usize, k2: u32, shx: bool, cfg: usize, use_ticket: bool, limit: f32, tables: bool) !void { + const ki = gmK2Index(k2) orelse return error.Unsupported; + if (cfg >= gu_tiles.len) return error.Unsupported; + try gmCheck(.gu, cfg, K, N); + var a: GmArgs = .{ .x0 = xg, .x1 = xu, .order = order, .pe = pe, .poff = poff, .pcnt = pcnt, .npass = npass, .ticket = ticket, .sv0 = svh_g, .sv1 = svh_u, .sd = suh_d, .xd = xd, .K = int(K), .N = int(N), .limit = limit }; + if (tables) { + a.tp0 = tg; + a.tp1 = tu; + } else { + a.t0 = tg; + a.t1 = tu; + } + try o.gmLaunch(o.f.gm_gu[ki][if (shx) 0 else 1][cfg], a, use_ticket); +} + +/// dsv41_x3gm_down_cuda: Y (fp32) of every pass's members. +pub fn gmDown(o: Ops, xd: u64, td: u64, order: u64, pe: u64, poff: u64, pcnt: u64, npass: u64, ticket: u64, svh_d: u64, y: u64, K: usize, N: usize, k2: u32, cfg: usize, use_ticket: bool, tables: bool) !void { + const ki = gmK2Index(k2) orelse return error.Unsupported; + if (cfg >= dn_tiles.len) return error.Unsupported; + try gmCheck(.dn, cfg, K, N); + var a: GmArgs = .{ .x0 = xd, .x1 = xd, .order = order, .pe = pe, .poff = poff, .pcnt = pcnt, .npass = npass, .ticket = ticket, .sv0 = svh_d, .y = y, .K = int(K), .N = int(N) }; + if (tables) { + a.tp0 = td; + a.tp1 = td; + } else { + a.t0 = td; + a.t1 = td; + } + try o.gmLaunch(o.f.gm_dn[ki][cfg], a, use_ticket); +} + +/// dsv41_x3gm_gateup2_cuda (x3gm v2): Xd of every pass of every width in one launch; `tables`: ragged pointers. +pub fn gm2Gateup(o: Ops, xg: u64, xu: u64, tg: u64, tu: u64, k2e: u64, order: u64, pe: u64, poff: u64, pcnt: u64, npass: u64, ticket: u64, svh_g: u64, svh_u: u64, suh_d: u64, xd: u64, K: usize, N: usize, shx: bool, cfg: usize, use_ticket: bool, limit: f32, tables: bool) !void { + const gu2 = cfg == gu2_cfg and !shx; + const ci = std.mem.indexOfScalar(usize, &v2_gu, cfg) orelse if (gu2) 0 else return error.Unsupported; + try gmCheck(.gu, cfg, K, N); + var a: GmArgs = .{ .x0 = xg, .x1 = xu, .k2e = k2e, .order = order, .pe = pe, .poff = poff, .pcnt = pcnt, .npass = npass, .ticket = ticket, .sv0 = svh_g, .sv1 = svh_u, .sd = suh_d, .xd = xd, .K = int(K), .N = int(N), .limit = limit }; + if (tables) { + a.tp0 = tg; + a.tp1 = tu; + } else { + a.t0 = tg; + a.t1 = tu; + } + try o.gmLaunch(if (gu2) o.f.gm2_gu2 else o.f.gm2_gu[if (shx) 0 else 1][ci], a, use_ticket); +} + +/// dsv41_x3gm_down2_cuda (x3gm v2): Y of every pass of every width in one launch; `k2e` the experts' down widths. +pub fn gm2Down(o: Ops, xd: u64, td: u64, k2e: u64, order: u64, pe: u64, poff: u64, pcnt: u64, npass: u64, ticket: u64, svh_d: u64, y: u64, K: usize, N: usize, cfg: usize, use_ticket: bool, tables: bool) !void { + if (std.mem.indexOfScalar(usize, &v2_dn, cfg) == null) return error.Unsupported; + try gmCheck(.dn, cfg, K, N); + var a: GmArgs = .{ .x0 = xd, .x1 = xd, .k2e = k2e, .order = order, .pe = pe, .poff = poff, .pcnt = pcnt, .npass = npass, .ticket = ticket, .sv0 = svh_d, .y = y, .K = int(K), .N = int(N) }; + if (tables) { + a.tp0 = td; + a.tp1 = td; + } else { + a.t0 = td; + a.t1 = td; + } + try o.gmLaunch(o.f.gm2_dn, a, use_ticket); +} + +/// x3gm v3 gate/up (gm3_kernel), variant `v` of v3_variants (mats 2, its xm): gm2Gateup's arguments and bits. +pub fn gm3Gateup(o: Ops, v: usize, xg: u64, xu: u64, tg: u64, tu: u64, k2e: u64, order: u64, pe: u64, poff: u64, pcnt: u64, npass: u64, ticket: u64, svh_g: u64, svh_u: u64, suh_d: u64, xd: u64, K: usize, N: usize, shx: bool, use_ticket: bool, limit: f32, tables: bool) !void { + const vv = v3_variants[v]; + if (vv.mats != 2 or vv.xm != @as(u32, if (shx) 1 else 2)) return error.Unsupported; + if ((K / 16) % vv.t[2] != 0 or N % 128 != 0) return error.Shape; + var a: GmArgs = .{ .x0 = xg, .x1 = xu, .k2e = k2e, .order = order, .pe = pe, .poff = poff, .pcnt = pcnt, .npass = npass, .ticket = ticket, .sv0 = svh_g, .sv1 = svh_u, .sd = suh_d, .xd = xd, .K = int(K), .N = int(N), .limit = limit }; + if (tables) { + a.tp0 = tg; + a.tp1 = tu; + } else { + a.t0 = tg; + a.t1 = tu; + } + try o.gmLaunch(o.f.gm3[v], a, use_ticket); +} + +/// x3gm v3 down (gm3_kernel), variant `v` (mats 1): gm2Down's arguments and bits. +pub fn gm3Down(o: Ops, v: usize, xd: u64, td: u64, k2e: u64, order: u64, pe: u64, poff: u64, pcnt: u64, npass: u64, ticket: u64, svh_d: u64, y: u64, K: usize, N: usize, use_ticket: bool, tables: bool) !void { + const vv = v3_variants[v]; + if (vv.mats != 1) return error.Unsupported; + if ((K / 16) % vv.t[2] != 0 or N % 128 != 0) return error.Shape; + var a: GmArgs = .{ .x0 = xd, .x1 = xd, .k2e = k2e, .order = order, .pe = pe, .poff = poff, .pcnt = pcnt, .npass = npass, .ticket = ticket, .sv0 = svh_d, .y = y, .K = int(K), .N = int(N) }; + if (tables) { + a.tp0 = td; + a.tp1 = td; + } else { + a.t0 = td; + a.t1 = td; + } + try o.gmLaunch(o.f.gm3[v], a, use_ticket); +} + +/// x3gm.plan on the device (x3gm_plan.cu): order [pairs], pe / poff / pcnt [planPasses], npass [1] in `bm` members. +pub fn gmPlan(o: Ops, pick: u64, pairs: usize, E: usize, bm: usize, order: u64, pe: u64, poff: u64, pcnt: u64, npass: u64) !void { + if (E < 1 or E >= plan_max_experts or bm < 1 or pairs < 1) return error.Shape; + var a: cuda.Args = .{}; + a.add(pick); + for ([_]usize{ pairs, E, bm }) |v| a.add(int(v)); + for ([_]u64{ order, pe, poff, pcnt, npass }) |v| a.add(v); + try go(o.f.gm_plan, o.s, .{ 1, 1, 1 }, 512, 0, false, &a); +} + +/// dsv41_x3gm_dequant_cuda: W_q [K, N] fp16 through the gm kernels' decode_a (tests). +pub fn gmDequant(o: Ops, t: u64, out: u64, K: usize, N: usize, k2: u32) !void { + const ki = gmK2Index(k2) orelse return error.Unsupported; + var a: cuda.Args = .{}; + a.add(t); + a.add(out); + a.add(int(N)); + try go(o.f.gm_dequant[ki], o.s, .{ K / 16, N / 16, 1 }, 32, 0, false, &a); +} diff --git a/zig/src/families/deepseek_v41/knob_types.zig b/zig/src/families/deepseek_v41/knob_types.zig new file mode 100644 index 0000000000..e1cb6d4eaa --- /dev/null +++ b/zig/src/families/deepseek_v41/knob_types.zig @@ -0,0 +1,78 @@ +//! Option types the block and forward options carry: the fused dense prefill GEMM's choices, the x3gm v2 mode. +const std = @import("std"); +const dk = @import("dsv41_kernels"); + +const pfd = dk.ops.pfd; + +/// TF_DSV41_GM_V2 (gm2pf.zig): prefill's routed experts on x3gm v2; off: v1. The same bits either way. +pub const GmMode = enum { off, one, split }; + +/// The SMs the dense prefill heuristic prices waves on, on every GPU (the pick never changes a bit). +pub const pfd_sms: usize = 48; + +/// The fused dense prefill GEMM's choices (pfdense.pick): an override, the JSON table, else the heuristic. +pub const PfDense = struct { + entries: []const Entry = &.{}, + /// TF_DSV41_PF_DENSE_CFG: (id, group) for every shape + override: ?[2]u32 = null, + + pub const Entry = struct { k: u32, n: u32, cfg: u32, group: u32 = 8, small: ?u32 = null, small_rows: u32 = 0 }; + + /// pfdense.why_not's shape and width rules (the codebook is mul1 for every dense group of our packs). + pub fn eligible(k: i64, n: i64, k2: u32, lanes: bool) bool { + const widths: []const u32 = if (lanes) &.{ 8, 10, 12 } else &.{ 8, 10, 12, 16 }; + if (std.mem.indexOfScalar(u32, widths, k2) == null) return false; + return k > 0 and n > 0 and @rem(k, 128) == 0 and @rem(n, 128) == 0; + } + + /// (cfg, group) for an [m, k] @ [k, n] call: override, table, else heuristic (also when a stage does not divide K). + pub fn pick(p: *const PfDense, k: usize, n: usize, k2: u32, m: usize) [2]u32 { + var got: ?[2]u32 = p.override; + if (got == null) for (p.entries) |e| { + if (e.k != k or e.n != n) continue; + const c = if (e.small != null and m <= e.small_rows) e.small.? else e.cfg; + got = .{ c, e.group }; + break; + }; + if (got) |g| if (g[0] < pfd.cfgs.len and k % (16 * pfd.cfgs[g[0]][3]) == 0) return g; + const h = pfd.heuristic(k, n, k2, m, pfd_sms); + return .{ @intCast(h[0]), @intCast(h[1]) }; + } + + /// pfdense._table_file: `{"K,N": {"cfg": id, "group": g, "small": id, "small_rows": r}}`. + pub fn parseTable(a: std.mem.Allocator, json: []const u8) ![]const Entry { + const parsed = std.json.parseFromSliceLeaky(std.json.Value, a, json, .{}) catch return error.BadTable; + if (parsed != .object) return error.BadTable; + const out = try a.alloc(Entry, parsed.object.count()); + var it = parsed.object.iterator(); + var i: usize = 0; + while (it.next()) |kv| : (i += 1) { + var ks = std.mem.tokenizeScalar(u8, kv.key_ptr.*, ','); + const k = std.fmt.parseInt(u32, std.mem.trim(u8, ks.next() orelse return error.BadTable, " "), 10) catch return error.BadTable; + const n = std.fmt.parseInt(u32, std.mem.trim(u8, ks.next() orelse return error.BadTable, " "), 10) catch return error.BadTable; + if (kv.value_ptr.* != .object) return error.BadTable; + const o = kv.value_ptr.object; + out[i] = .{ .k = k, .n = n, .cfg = try field(o, "cfg") orelse return error.BadTable }; + if (try field(o, "group")) |g| out[i].group = g; + out[i].small = try field(o, "small"); + out[i].small_rows = (try field(o, "small_rows")) orelse 0; + } + return out; + } + + fn field(o: std.json.ObjectMap, name: []const u8) !?u32 { + const v = o.get(name) orelse return null; + return switch (v) { + .integer => |x| std.math.cast(u32, x) orelse error.BadTable, + else => error.BadTable, + }; + } + + pub fn parseOverride(raw: []const u8) ![2]u32 { + var it = std.mem.tokenizeScalar(u8, raw, ','); + const id = std.fmt.parseInt(u32, std.mem.trim(u8, it.next() orelse return error.BadKnob, " "), 10) catch return error.BadKnob; + const g: u32 = if (it.next()) |x| std.fmt.parseInt(u32, std.mem.trim(u8, x, " "), 10) catch return error.BadKnob else 8; + if (it.next() != null or id >= pfd.cfgs.len or g < 1 or g > 1024) return error.BadKnob; + return .{ id, g }; + } +}; diff --git a/zig/src/families/deepseek_v41/load.zig b/zig/src/families/deepseek_v41/load.zig new file mode 100644 index 0000000000..881e3529cc --- /dev/null +++ b/zig/src/families/deepseek_v41/load.zig @@ -0,0 +1,94 @@ +//! A rank's named weights on the GPU: one device buffer a tensor (bound by name), digested on the way up. + +const std = @import("std"); +const cuda = @import("cuda"); +const named = @import("named.zig"); +const ex_stacks = @import("ex_stacks.zig"); + +pub const Tensor = struct { + ptr: u64, + len: usize, + kind: named.Kind, + sha256: [32]u8, + shape: [4]usize = @splat(1), + rank: u8 = 0, + + pub fn dims(t: *const Tensor) []const usize { + return t.shape[0..t.rank]; + } +}; + +pub const Weights = struct { + gpa: std.mem.Allocator, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + map: std.StringHashMapUnmanaged(Tensor) = .empty, + bytes: usize = 0, + /// each tensor's SHA-256 taken on the way up (the weights check reads it); off, `sha256` is zeros (served boot) + digest: bool = true, + + pub fn deinit(w: *Weights) void { + for (w.bufs.items) |*b| b.free(); + w.bufs.deinit(w.gpa); + var it = w.map.keyIterator(); + while (it.next()) |k| w.gpa.free(k.*); + w.map.deinit(w.gpa); + w.* = undefined; + } + + /// Uploads every tensor `b` holds (its host images can be freed afterwards); a name already on the GPU is refused. + pub fn upload(w: *Weights, d: *const cuda.Driver, b: *const named.Builder) !void { + // TF_DSV41_EX_STACKS: a routed suh / svh and its shared expert's in one buffer (ex_stacks.zig) + var paired: std.AutoHashMapUnmanaged(usize, void) = .empty; + defer paired.deinit(w.gpa); + if (ex_stacks.enabled()) { + const names = try w.gpa.alloc([]const u8, b.out.items.len); + defer w.gpa.free(names); + const lens = try w.gpa.alloc(usize, b.out.items.len); + defer w.gpa.free(lens); + for (b.out.items, names, lens) |n, *x, *l| { + x.* = n.name; + l.* = n.bytes.len; + } + const ps = try ex_stacks.pairs(w.gpa, names, lens); + defer w.gpa.free(ps); + for (ps) |p| { + const r = b.out.items[p[0]]; + const s = b.out.items[p[1]]; + if (w.map.contains(r.name) or w.map.contains(s.name)) return error.DuplicateWeight; + var buf = try cuda.DeviceBuffer.alloc(d, r.bytes.len + s.bytes.len); + errdefer buf.free(); + try cuda.DeviceBuffer.upload(buf, 0, r.bytes); + try cuda.DeviceBuffer.upload(buf, r.bytes.len, s.bytes); + try w.bufs.append(w.gpa, buf); + try w.put(r, buf.ptr); + try w.put(s, buf.ptr + r.bytes.len); + try paired.put(w.gpa, p[0], {}); + try paired.put(w.gpa, p[1], {}); + } + } + for (b.out.items, 0..) |n, i| { + if (paired.contains(i)) continue; + if (w.map.contains(n.name)) return error.DuplicateWeight; + var sha: [32]u8 = @splat(0); + if (w.digest) std.crypto.hash.sha2.Sha256.hash(n.bytes, &sha, .{}); + var buf = try cuda.DeviceBuffer.fromHost(d, n.bytes); + errdefer buf.free(); + try w.bufs.append(w.gpa, buf); + try w.map.put(w.gpa, try w.gpa.dupe(u8, n.name), .{ .ptr = buf.ptr, .len = n.bytes.len, .kind = n.kind, .sha256 = sha, .shape = n.shape, .rank = n.rank }); + w.bytes += n.bytes.len; + } + // pageable uploads can still be in flight when cuMemcpyHtoD returns; the weights are whole before any launch + try d.check(d.api.cuCtxSynchronize(), "cuCtxSynchronize"); + } + + fn put(w: *Weights, n: named.Named, ptr: u64) !void { + var sha: [32]u8 = @splat(0); + if (w.digest) std.crypto.hash.sha2.Sha256.hash(n.bytes, &sha, .{}); + try w.map.put(w.gpa, try w.gpa.dupe(u8, n.name), .{ .ptr = ptr, .len = n.bytes.len, .kind = n.kind, .sha256 = sha, .shape = n.shape, .rank = n.rank }); + w.bytes += n.bytes.len; + } + + pub fn get(w: *const Weights, name: []const u8) ?Tensor { + return w.map.get(name); + } +}; diff --git a/zig/src/families/deepseek_v41/named.zig b/zig/src/families/deepseek_v41/named.zig new file mode 100644 index 0000000000..d439d03508 --- /dev/null +++ b/zig/src/families/deepseek_v41/named.zig @@ -0,0 +1,261 @@ +//! A rank's weights as named host tensors in the reference loader's forms, named as the reference capture names them. + +const std = @import("std"); +const plan = @import("plan.zig"); +const names = @import("names.zig"); +const stage = @import("stage.zig"); +const Pack = @import("pack.zig").Pack; +const DType = @import("pack.zig").DType; +const Io = std.Io; + +pub const Kind = enum { f32, bf16, f16, i16 }; + +/// One tensor the forward holds: its name (`L2.attn.wq_a.trellis`, `embed`), element kind, shape and bytes. +pub const Named = struct { + name: []const u8, + kind: Kind, + shape: [4]usize, + rank: u8, + bytes: []align(16) u8, +}; + +/// What a native tensor becomes (Loader.native and its callers). +const Out = enum { + /// fp32 values (`native(name)`: bf16 or fp32 storage widened) + f32, + /// bf16 values (`native(name, bf16, bf16_param=True)`: the router gate) + bf16, + /// fp32 holding bf16-rounded values (`native(name, bf16_param=True)`: the indexer's weights_proj) + f32_via_bf16, + /// Engram's q / k: combined into `engram.qk` = bf16(q) * bf16(k) in fp32 + qk, +}; + +const Rule = struct { name: []const u8, out: Out }; + +/// Native suffixes (names.zig blockNatives) -> the capture's field names and forms. +const rules = std.StaticStringMap(Rule).initComptime(.{ + .{ "attn_norm.weight", Rule{ .name = "attn_norm", .out = .f32 } }, + .{ "ffn_norm.weight", Rule{ .name = "ffn_norm", .out = .f32 } }, + .{ "hc_attn_fn", Rule{ .name = "hc_attn.fn", .out = .f32 } }, + .{ "hc_attn_base", Rule{ .name = "hc_attn.base", .out = .f32 } }, + .{ "hc_attn_scale", Rule{ .name = "hc_attn.scale", .out = .f32 } }, + .{ "hc_ffn_fn", Rule{ .name = "hc_ffn.fn", .out = .f32 } }, + .{ "hc_ffn_base", Rule{ .name = "hc_ffn.base", .out = .f32 } }, + .{ "hc_ffn_scale", Rule{ .name = "hc_ffn.scale", .out = .f32 } }, + .{ "attn.q_norm.weight", Rule{ .name = "attn.q_norm", .out = .f32 } }, + .{ "attn.kv_norm.weight", Rule{ .name = "attn.kv_norm", .out = .f32 } }, + .{ "attn.attn_sink", Rule{ .name = "attn.sink", .out = .f32 } }, + .{ "attn.compressor.norm.weight", Rule{ .name = "attn.comp_norm", .out = .f32 } }, + .{ "attn.indexer.k_norm.weight", Rule{ .name = "attn.ix_knorm", .out = .f32 } }, + .{ "attn.indexer.weights_proj.weight", Rule{ .name = "attn.ix_wp", .out = .f32_via_bf16 } }, + .{ "ffn.gate.weight", Rule{ .name = "moe.gate", .out = .bf16 } }, + .{ "ffn.gate.bias", Rule{ .name = "moe.bias", .out = .f32 } }, + .{ "engram.q_weight", Rule{ .name = "engram.qk", .out = .qk } }, + .{ "engram.k_weight", Rule{ .name = "engram.qk", .out = .qk } }, +}); + +/// The capture's name of a block group (AttnW / MoEW / EngramW fields). +fn groupField(g: plan.GroupPlan, gl: u32, buf: []u8) ![]const u8 { + return switch (g.proj) { + .wq_a => "attn.wq_a", + .wkv => "attn.wkv", + .wq_b => "attn.wq_b", + .wo_a => std.fmt.bufPrint(buf, "attn.wo_a.{d}", .{g.slice % gl}), + .wo_b => "attn.wo_b", + .comp_wkv => "attn.comp_wkv", + .comp_wgate => "attn.comp_wgate", + .ix_wq_b => "attn.ix_wq_b", + .ix_wk => "attn.ix_wk", + .shared_w1 => "moe.shared.0.w1", + .shared_w2 => "moe.shared.0.w2", + .shared_w3 => "moe.shared.0.w3", + .engram_wkv => "engram.wkv", + }; +} + +/// fp32 -> bf16, round to nearest even (torch's .to(bfloat16)); NaN stays a quiet NaN. +pub fn bf16Round(x: f32) u16 { + const b: u32 = @bitCast(x); + if (std.math.isNan(x)) return @intCast((b >> 16) | 0x40); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +pub fn bf16Widen(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// One stored element as fp32 (bf16, fp16 and fp32 storage). +fn load(dtype: DType, src: []const u8, i: usize) !f32 { + return switch (dtype) { + .f32 => @bitCast(std.mem.readInt(u32, src[4 * i ..][0..4], .little)), + .bf16 => bf16Widen(std.mem.readInt(u16, src[2 * i ..][0..2], .little)), + .f16 => @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, src[2 * i ..][0..2], .little)))), + else => error.UnexpectedTensor, + }; +} + +pub const Builder = struct { + gpa: std.mem.Allocator, + io: Io, + pack: *const Pack, + out: std.ArrayList(Named) = .empty, + + pub fn deinit(b: *Builder) void { + for (b.out.items) |n| { + b.gpa.free(n.name); + b.gpa.free(n.bytes); + } + b.out.deinit(b.gpa); + } + + fn put(b: *Builder, name: []const u8, kind: Kind, shape: []const usize, bytes: []align(16) u8) !void { + var n: Named = .{ .name = try b.gpa.dupe(u8, name), .kind = kind, .shape = @splat(1), .rank = @intCast(shape.len), .bytes = bytes }; + @memcpy(n.shape[0..shape.len], shape); + try b.out.append(b.gpa, n); + } + + fn copy(b: *Builder, src: []const u8) ![]align(16) u8 { + const d = try b.gpa.alignedAlloc(u8, .@"16", src.len); + @memcpy(d, src); + return d; + } + + /// A group's part as `.trellis`, `.suh`, `.svh`. + fn group(b: *Builder, base: []const u8, g: plan.GroupPlan) !void { + var img = try stage.group(b.gpa, b.io, b.pack, g); + defer img.deinit(b.gpa); + const p = g.part.group; + var nb: [160]u8 = undefined; + try b.put(try std.fmt.bufPrint(&nb, "{s}.trellis", .{base}), .i16, &.{ p.k_tiles, p.n_tiles, p.words }, try b.copy(img.trellis)); + try b.put(try std.fmt.bufPrint(&nb, "{s}.suh", .{base}), .f16, &.{p.k()}, try b.copy(img.suh)); + try b.put(try std.fmt.bufPrint(&nb, "{s}.svh", .{base}), .f16, &.{p.n()}, try b.copy(img.svh)); + } + + /// A routed projection: `.trellis` [E, kt, nt, words] (stack) or `.data` (ragged), `.suh` [E, K], `.svh` [E, N]. + fn experts(b: *Builder, base: []const u8, x: *const plan.ExpertsPlan) !void { + var img = try stage.experts(b.gpa, b.io, b.pack, x); + defer img.deinit(b.gpa); + const l = &x.layout; + const e = l.count(); + var nb: [160]u8 = undefined; + if (l.kind == .stack) { + try b.put(try std.fmt.bufPrint(&nb, "{s}.trellis", .{base}), .i16, &.{ e, l.k_tiles, l.n_tiles, l.words[0] }, try b.copy(img.trellis)); + } else try b.put(try std.fmt.bufPrint(&nb, "{s}.data", .{base}), .i16, &.{@intCast(l.elements)}, try b.copy(img.trellis)); + try b.put(try std.fmt.bufPrint(&nb, "{s}.suh", .{base}), .f16, &.{ e, 16 * @as(usize, l.k_tiles) }, try b.copy(img.suh)); + try b.put(try std.fmt.bufPrint(&nb, "{s}.svh", .{base}), .f16, &.{ e, 16 * @as(usize, l.n_tiles) }, try b.copy(img.svh)); + } + + /// The rank's slice of a native tensor, as stored. + fn raw(b: *Builder, v: plan.NativePlan, rank: u32, world: u32) ![]u8 { + const whole = v.slice == .whole; + const len = if (whole) v.info.nbytes else v.info.nbytes / world; + const bytes = try b.gpa.alloc(u8, @intCast(len)); + errdefer b.gpa.free(bytes); + try b.pack.read(b.io, v.info, .{ .offset = if (whole) 0 else rank * len, .len = len }, bytes); + return bytes; + } + + /// A native tensor in its forward form; Engram's q waits for its k (`pending`). + fn native(b: *Builder, name: []const u8, rule: Rule, v: plan.NativePlan, rank: u32, world: u32, pending: *?[]f32) !void { + const bytes = try b.raw(v, rank, world); + defer b.gpa.free(bytes); + const count = bytes.len / v.info.dtype.size(); + var shape: [4]usize = @splat(1); + const dims = v.info.dims(); + @memcpy(shape[0..dims.len], dims); + if (v.slice != .whole) shape[0] /= world; + switch (rule.out) { + .bf16 => { + const out = try b.gpa.alignedAlloc(u8, .@"16", 2 * count); + for (0..count) |i| std.mem.writeInt(u16, out[2 * i ..][0..2], bf16Round(try load(v.info.dtype, bytes, i)), .little); + try b.put(name, .bf16, shape[0..dims.len], out); + }, + .f32, .f32_via_bf16 => { + const out = try b.gpa.alignedAlloc(u8, .@"16", 4 * count); + for (0..count) |i| { + const x = try load(v.info.dtype, bytes, i); + const y = if (rule.out == .f32) x else bf16Widen(bf16Round(x)); + std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(y), .little); + } + try b.put(name, .f32, shape[0..dims.len], out); + }, + .qk => { + const vals = try b.gpa.alloc(f32, count); + for (vals, 0..) |*x, i| x.* = bf16Widen(bf16Round(try load(v.info.dtype, bytes, i))); + if (pending.*) |q| { + defer b.gpa.free(q); + defer b.gpa.free(vals); + if (q.len != vals.len) return error.UnexpectedTensor; + const out = try b.gpa.alignedAlloc(u8, .@"16", 4 * count); + for (q, vals, 0..) |x, y, i| std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(x * y), .little); + pending.* = null; + try b.put(name, .f32, shape[0..dims.len], out); + } else pending.* = vals; + }, + } + } + + pub fn layer(b: *Builder, l: *const plan.LayerPlan, o_groups: u32, rank: u32, world: u32) !void { + var nb: [160]u8 = undefined; + var fb: [32]u8 = undefined; + const gl = o_groups / world; + for (l.groups) |g| try b.group(try std.fmt.bufPrint(&nb, "L{d}.{s}", .{ l.index, try groupField(g, gl, &fb) }), g); + var pending: ?[]f32 = null; + defer if (pending) |q| b.gpa.free(q); + for (l.natives) |v| { + const suffix = v.name[l.prefix.len + 1 ..]; + const rule = rules.get(suffix) orelse return error.UnnamedNative; + try b.native(try std.fmt.bufPrint(&nb, "L{d}.{s}", .{ l.index, rule.name }), rule, v, rank, world, &pending); + } + if (pending != null) return error.EngramHalf; + for (&l.experts, [_][]const u8{ "w1", "w2", "w3" }) |*x, w| try b.experts(try std.fmt.bufPrint(&nb, "L{d}.moe.{s}", .{ l.index, w }), x); + } + + /// The vocabulary rows (bf16), the final norm (fp32) and the head's part. + pub fn top(b: *Builder, p: *const plan.Plan) !void { + var none: ?[]f32 = null; + for (p.top) |v| { + const is_embed = std.mem.eql(u8, v.name, "embed.weight"); + const rule: Rule = if (is_embed) .{ .name = "embed", .out = .bf16 } else .{ .name = "norm", .out = .f32 }; + try b.native(rule.name, rule, v, p.rank, p.world, &none); + } + if (p.head) |h| try b.group("head", h); + } + + /// DSpark's heads: main_proj (whole), main_norm and norm fp32, Markov embed / head bf16, confidence head fp32. + pub fn dspark(b: *Builder, p: *const plan.Plan) !void { + const d = p.dspark orelse return; + var none: ?[]f32 = null; + try b.group("dspark.main_proj", d.main_proj); + const out_names = [_][]const u8{ "dspark.main_norm", "dspark.norm", "dspark.markov.w1", "dspark.markov.w2" }; + const outs = [_]Out{ .f32, .f32, .bf16, .bf16 }; + for (d.natives, out_names, outs) |v, nm, o| try b.native(nm, .{ .name = nm, .out = o }, v, p.rank, p.world, &none); + if (d.confidence) |c| try b.native("dspark.conf", .{ .name = "dspark.conf", .out = .f32 }, c, p.rank, p.world, &none); + } + + /// Every planned block, then the model-level tensors (a whole-model plan only). + pub fn all(b: *Builder, p: *const plan.Plan, o_groups: u32) !void { + for (p.layers) |*l| try b.layer(l, o_groups, p.rank, p.world); + try b.top(p); + try b.dspark(p); + } + + pub fn find(b: *const Builder, name: []const u8) ?*const Named { + for (b.out.items) |*n| if (std.mem.eql(u8, n.name, name)) return n; + return null; + } +}; + +const testing = std.testing; + +test "the loader's conversions: bf16 round to nearest even, widening, fp16 storage" { + try testing.expectEqual(@as(u16, 0x3f80), bf16Round(1.0)); + // 1 + 2^-8 is a tie between 1.0 and 1 + 2^-7: even keeps 1.0; 1 + 3 * 2^-8 rounds up + try testing.expectEqual(@as(u16, 0x3f80), bf16Round(1.00390625)); + try testing.expectEqual(@as(u16, 0x3f82), bf16Round(1.01171875)); + try testing.expectEqual(@as(f32, -2.5), bf16Widen(bf16Round(-2.5))); + var h: [2]u8 = undefined; + std.mem.writeInt(u16, &h, @bitCast(@as(f16, 0.333251953125)), .little); + try testing.expectEqual(@as(f32, 0.333251953125), try load(.f16, &h, 0)); +} diff --git a/zig/src/families/deepseek_v41/names.zig b/zig/src/families/deepseek_v41/names.zig new file mode 100644 index 0000000000..16da179483 --- /dev/null +++ b/zig/src/families/deepseek_v41/names.zig @@ -0,0 +1,236 @@ +//! The pack's tensor names, each one's TP split and expected shape, plus the names the text engine leaves unread. + +const std = @import("std"); +const Config = @import("config.zig").Config; +const List = @import("config.zig").List; +const Split = @import("exl3.zig").Split; +const DType = @import("core").safetensors.DType; + +/// "layers.L" for the backbone, "mtp.i" for DSpark block i (layer = layers + i). +pub fn blockPrefix(c: *const Config, layer: u32, buf: []u8) ![]const u8 { + return if (c.isBackbone(layer)) std.fmt.bufPrint(buf, "layers.{d}", .{layer}) else std.fmt.bufPrint(buf, "mtp.{d}", .{layer - c.layers}); +} + +/// An EXL3 projection of a block, by the forward's name for it. +pub const Proj = enum { wq_a, wkv, wq_b, wo_a, wo_b, comp_wkv, comp_wgate, ix_wq_b, ix_wk, shared_w1, shared_w2, shared_w3, engram_wkv }; + +/// A block's EXL3 group: its name under the block prefix, how ranks split it, and its K x N (0: not checked). +pub const GroupSpec = struct { + proj: Proj, + suffix: []const u8, + /// wo_a: the global group index this slice is + slice: u32 = 0, + split: Split, + k: u32 = 0, + n: u32 = 0, +}; + +pub const Groups = List(GroupSpec, 32); + +/// The block's EXL3 groups this rank loads, in the loader's order (wo_a: only this rank's groups). +pub fn blockGroups(c: *const Config, layer: u32, rank: u32, world: u32, has_shared: bool) !Groups { + var g: Groups = .{}; + const per_group_k = c.heads / c.o_groups * c.head_dim; + try g.push(.{ .proj = .wq_a, .suffix = "attn.wq_a", .split = .whole, .k = c.hidden, .n = c.q_lora }); + try g.push(.{ .proj = .wkv, .suffix = "attn.wkv", .split = .whole, .k = c.hidden, .n = c.head_dim }); + try g.push(.{ .proj = .wq_b, .suffix = "attn.wq_b", .split = .col, .k = c.q_lora, .n = c.heads * c.head_dim }); + const gl = c.o_groups / world; + var s = rank * gl; + while (s < (rank + 1) * gl) : (s += 1) try g.push(.{ .proj = .wo_a, .suffix = "attn.wo_a.slice", .slice = s, .split = .whole, .k = per_group_k, .n = c.o_lora }); + try g.push(.{ .proj = .wo_b, .suffix = "attn.wo_b", .split = .row, .k = c.o_groups * c.o_lora, .n = c.hidden }); + if (c.isBackbone(layer) and c.isKvSource(layer)) { + try g.push(.{ .proj = .comp_wkv, .suffix = "attn.compressor.wkv", .split = .whole, .k = c.hidden, .n = c.head_dim }); + if (c.compressRatio(layer) > 1) try g.push(.{ .proj = .comp_wgate, .suffix = "attn.compressor.wgate", .split = .whole, .k = c.hidden, .n = c.head_dim }); + try g.push(.{ .proj = .ix_wk, .suffix = "attn.indexer.wk", .split = .whole, .k = c.head_dim, .n = c.index_dim }); + } + if (c.isBackbone(layer) and c.isIndexSource(layer)) + try g.push(.{ .proj = .ix_wq_b, .suffix = "attn.indexer.wq_b", .split = .whole, .k = c.q_lora, .n = c.index_heads * c.index_dim }); + if (has_shared) { + try g.push(.{ .proj = .shared_w1, .suffix = "ffn.shared_experts.w1", .split = .col, .k = c.hidden, .n = c.expert_width }); + try g.push(.{ .proj = .shared_w2, .suffix = "ffn.shared_experts.w2", .split = .row, .k = c.expert_width, .n = c.hidden }); + try g.push(.{ .proj = .shared_w3, .suffix = "ffn.shared_experts.w3", .split = .col, .k = c.hidden, .n = c.expert_width }); + } + if (c.isEngram(layer)) try g.push(.{ .proj = .engram_wkv, .suffix = "engram.wkv", .split = .whole, .k = c.hashCols() * c.engram_head_dim, .n = (c.engram_max_ngram + 1) * c.hidden }); + return g; +} + +/// The group's full name under `block`. +pub fn groupName(spec: GroupSpec, block: []const u8, buf: []u8) ![]const u8 { + return if (spec.proj == .wo_a) std.fmt.bufPrint(buf, "{s}.{s}.{d}", .{ block, spec.suffix, spec.slice }) else std.fmt.bufPrint(buf, "{s}.{s}", .{ block, spec.suffix }); +} + +/// Routed experts' projections: gate (w1) and up (w3) split by columns, down (w2) by rows. +pub const ExpertProj = enum { + w1, + w2, + w3, + + pub fn split(p: ExpertProj) Split { + return if (p == .w2) .row else .col; + } + + /// (K, N) of one expert's full matrix. + pub fn kn(p: ExpertProj, c: *const Config) [2]u32 { + return if (p == .w2) .{ c.expert_width, c.hidden } else .{ c.hidden, c.expert_width }; + } +}; + +pub fn expertName(block: []const u8, e: u32, p: ExpertProj, buf: []u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.ffn.experts.{d}.{t}", .{ block, e, p }); +} + +/// Which dtypes a native tensor may be stored in; the loader converts to what the forward holds. +pub const Store = enum { + /// norms: bf16 in the pack (fp32 in the forward), fp32 accepted + norm, + /// fp32 only (sinks, hc_*, router bias) + f32, + /// fp16 or bf16, rounded to bf16 once (router gate, indexer weights_proj, Markov heads) + half, + + pub fn accepts(s: Store, d: DType) bool { + return switch (s) { + .norm => d == .bf16 or d == .f32, + .f32 => d == .f32, + .half => d == .bf16 or d == .f16, + }; + } +}; + +/// How a rank holds a native tensor. +pub const Slice = enum { whole, heads, vocab_rows }; + +pub const NativeSpec = struct { + suffix: []const u8, + store: Store, + shape: [2]u32, + rank: u8, + slice: Slice = .whole, + + fn v(suffix: []const u8, store: Store, n: u32) NativeSpec { + return .{ .suffix = suffix, .store = store, .shape = .{ n, 0 }, .rank = 1 }; + } + + fn m(suffix: []const u8, store: Store, r: u32, cols: u32) NativeSpec { + return .{ .suffix = suffix, .store = store, .shape = .{ r, cols }, .rank = 2 }; + } + + pub fn dims(s: *const NativeSpec) []const u32 { + return s.shape[0..s.rank]; + } +}; + +pub const Natives = List(NativeSpec, 24); + +/// A block's native tensors: norms, mHC mixers, the sink, compressor / indexer pieces, the router, Engram's q.k. +pub fn blockNatives(c: *const Config, layer: u32) !Natives { + var n: Natives = .{}; + const hc = c.hc_mult; + const mix = (2 + hc) * hc; // pre (hc), post (hc) and the hc x hc residual mix + try n.push(.v("attn_norm.weight", .norm, c.hidden)); + try n.push(.v("ffn_norm.weight", .norm, c.hidden)); + inline for (.{ "attn", "ffn" }) |w| { + try n.push(.m("hc_" ++ w ++ "_fn", .f32, mix, hc * c.hidden)); + try n.push(.v("hc_" ++ w ++ "_base", .f32, mix)); + try n.push(.v("hc_" ++ w ++ "_scale", .f32, 3)); + } + try n.push(.v("attn.q_norm.weight", .norm, c.q_lora)); + try n.push(.v("attn.kv_norm.weight", .norm, c.head_dim)); + var sink = NativeSpec.v("attn.attn_sink", .f32, c.heads); + sink.slice = .heads; + try n.push(sink); + if (c.isBackbone(layer) and c.isKvSource(layer)) { + try n.push(.v("attn.compressor.norm.weight", .norm, c.head_dim)); + try n.push(.v("attn.indexer.k_norm.weight", .norm, c.index_dim)); + } + if (c.isBackbone(layer) and c.isIndexSource(layer)) try n.push(.m("attn.indexer.weights_proj.weight", .half, c.index_heads, c.hidden)); + const ex = c.expertsOf(layer); + try n.push(.m("ffn.gate.weight", .half, ex.count, c.hidden)); + try n.push(.v("ffn.gate.bias", .f32, ex.count)); + if (c.isEngram(layer)) { + try n.push(.m("engram.q_weight", .norm, hc, c.hidden)); + try n.push(.m("engram.k_weight", .norm, hc, c.hidden)); + } + return n; +} + +/// Model-level tensors outside the blocks: this rank's vocabulary rows, the final norm (the head is an EXL3 group). +pub fn topNatives(c: *const Config) !Natives { + var n: Natives = .{}; + var embed = NativeSpec.m("embed.weight", .half, c.vocab, c.hidden); + embed.slice = .vocab_rows; + try n.push(embed); + try n.push(.v("norm.weight", .norm, c.hidden)); + return n; +} + +pub const head = GroupSpec{ .proj = .wq_a, .suffix = "head", .split = .col }; + +/// DSpark's heads beside its blocks (the latest mtp.i that has each). +pub fn dsparkNatives(c: *const Config) !Natives { + var n: Natives = .{}; + try n.push(.v("main_norm.weight", .norm, c.hidden)); + try n.push(.v("norm.weight", .norm, c.hidden)); + try n.push(.m("markov_head.embed.weight", .half, c.vocab, c.dspark_markov_rank)); + try n.push(.m("markov_head.head.weight", .half, c.vocab, c.dspark_markov_rank)); + return n; +} + +/// The optional confidence head: fp32 values [D + r], any shape with that many. +pub const dspark_confidence = "confidence_head.proj.weight"; +pub const dspark_main_proj = "main_proj"; + +/// Names the text engine never reads: vision tower, image markers, image routing bias, bf16 copies of EXL3 groups. +pub fn ignored(name: []const u8, has: anytype) bool { + for ([_][]const u8{ "vision.", "aligner.", "image_" }) |p| if (std.mem.startsWith(u8, name, p)) return true; + if (std.mem.endsWith(u8, name, ".ffn.gate.bias_vl")) return true; + if (std.mem.endsWith(u8, name, ".weight")) { + var buf: [256]u8 = undefined; + const base = name[0 .. name.len - ".weight".len]; + const sib = std.fmt.bufPrint(&buf, "{s}.trellis", .{base}) catch return false; + return has.has(sib); + } + return false; +} + +const testing = std.testing; + +test "the release's block 2: a ratio-2 KV and index source with its compressor gate" { + const c: Config = .{}; + var buf: [64]u8 = undefined; + try testing.expectEqualStrings("mtp.1", try blockPrefix(&c, 41, &buf)); + const g = try blockGroups(&c, 2, 1, 2, true); + var names: [32][]const u8 = undefined; + var nbuf: [32][96]u8 = undefined; + for (g.items(), 0..) |s, i| names[i] = try groupName(s, "layers.2", &nbuf[i]); + // wo_a: rank 1 of 2 holds groups 4..7 + try testing.expectEqualStrings("layers.2.attn.wo_a.slice.4", names[3]); + try testing.expectEqual(@as(usize, 3 + 4 + 1 + 3 + 1 + 3), g.len); + try testing.expectEqual(@as(u32, 4096), g.items()[3].k); + try testing.expectEqual(@as(u32, 8192), g.items()[7].k); + // layer 20: ratio 1, no compressor gate; layer 24: reindex (own wq_b only); layer 30: reuse; DSpark blocks: neither + const g20 = try blockGroups(&c, 20, 0, 2, true); + for (g20.items()) |s| try testing.expect(s.proj != .comp_wgate); + try testing.expectEqual(@as(usize, 3 + 4 + 1 + 1), (try blockGroups(&c, 24, 0, 2, false)).len); + try testing.expectEqual(@as(usize, 3 + 4 + 1), (try blockGroups(&c, 30, 0, 2, false)).len); + try testing.expectEqual(@as(usize, 3 + 4 + 1 + 3), (try blockGroups(&c, 41, 0, 2, true)).len); + try testing.expect((try blockGroups(&c, 14, 0, 2, true)).items()[15].proj == .engram_wkv); + try testing.expectEqualStrings("layers.3.ffn.experts.17.w2", try expertName("layers.3", 17, .w2, &buf)); +} + +test "natives and the ignore list" { + const c: Config = .{}; + const n = try blockNatives(&c, 41); + for (n.items()) |s| if (std.mem.eql(u8, s.suffix, "ffn.gate.weight")) try testing.expectEqual(@as(u32, 128), s.shape[0]); + try testing.expectEqual(@as(u32, 24), (try blockNatives(&c, 0)).items()[2].shape[0]); + const Has = struct { + pub fn has(_: @This(), name: []const u8) bool { + return std.mem.eql(u8, name, "layers.2.attn.indexer.wk.trellis"); + } + }; + try testing.expect(ignored("layers.2.attn.indexer.wk.weight", Has{})); + try testing.expect(!ignored("layers.2.attn.kv_norm.weight", Has{})); + try testing.expect(ignored("vision.blocks.3.mlp.w1.weight", Has{})); + try testing.expect(ignored("layers.9.ffn.gate.bias_vl", Has{})); + try testing.expect(Store.half.accepts(.f16) and !Store.f32.accepts(.bf16)); +} diff --git a/zig/src/families/deepseek_v41/pack.zig b/zig/src/families/deepseek_v41/pack.zig new file mode 100644 index 0000000000..131978bcd3 --- /dev/null +++ b/zig/src/families/deepseek_v41/pack.zig @@ -0,0 +1,306 @@ +//! A DeepSeek-V4.1 pack folder by its safetensors headers alone, read with pread and no mmap. + +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel +const log = std.log.scoped(.dsv41); +const core = @import("core"); +const exl3 = @import("exl3.zig"); +const st = core.safetensors; +const Io = std.Io; + +pub const DType = st.DType; + +/// One stored tensor: its shard, dtype, shape and where its bytes start in that file. +pub const Info = struct { + file: u32, + dtype: DType, + rank: u8, + shape: [st.max_rank]usize, + /// absolute byte offset of the first byte in the shard file + start: u64, + nbytes: u64, + + pub fn dims(i: *const Info) []const usize { + return i.shape[0..i.rank]; + } + + pub fn is(i: *const Info, dtype: DType, shape: []const usize) bool { + return i.dtype == dtype and std.mem.eql(usize, i.dims(), shape); + } +}; + +pub const Pack = struct { + arena: std.heap.ArenaAllocator, + /// the folder, else "" for header-only packs (tests, header dumps) + dir: []const u8 = "", + /// shards the index lists that the folder does not hold + missing: u32 = 0, + files: std.ArrayList([]const u8) = .empty, + names: std.StringArrayHashMapUnmanaged(Info) = .empty, + /// the index's weight map (name -> shard); when set, names another shard owns are skipped + owner: std.StringHashMapUnmanaged([]const u8) = .empty, + + pub fn init(gpa: std.mem.Allocator) Pack { + return .{ .arena = std.heap.ArenaAllocator.init(gpa) }; + } + + pub fn deinit(p: *Pack) void { + p.arena.deinit(); + p.* = undefined; + } + + /// Every shard's header in `dir`: the index's weight map if present, else every *.safetensors file. + pub fn open(gpa: std.mem.Allocator, io: Io, dir: []const u8) !Pack { + var p = init(gpa); + errdefer p.deinit(); + const a = p.arena.allocator(); + p.dir = try a.dupe(u8, dir); + const cwd = Io.Dir.cwd(); + var shards: std.ArrayList([]const u8) = .empty; + const index = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + if (cwd.readFileAlloc(io, index, a, .limited(1 << 28))) |text| { + try p.useIndex(text); + var it = p.owner.valueIterator(); + while (it.next()) |f| { + for (shards.items) |s| { + if (std.mem.eql(u8, s, f.*)) break; + } else try shards.append(a, f.*); + } + } else |err| switch (err) { + error.FileNotFound => { + var d = try cwd.openDir(io, dir, .{ .iterate = true }); + defer d.close(io); + var it = d.iterate(); + while (try it.next(io)) |e| if (std.mem.endsWith(u8, e.name, ".safetensors")) try shards.append(a, try a.dupe(u8, e.name)); + }, + else => return err, + } + if (shards.items.len == 0) return error.NoSafetensors; + std.mem.sort([]const u8, shards.items, {}, lessStr); + for (shards.items) |name| { + const path = try std.fs.path.join(a, &.{ dir, name }); + // a shard the index lists but the folder lacks: its tensors stay unknown (a plan needing one fails) + const file = cwd.openFile(io, path, .{}) catch |e| switch (e) { + error.FileNotFound => { + p.missing += 1; + continue; + }, + else => return e, + }; + defer file.close(io); + var head: [8]u8 = undefined; + if (try file.readPositionalAll(io, &head, 0) != 8) return error.BadSafetensors; + const n = std.mem.readInt(u64, &head, .little); + const len = try file.length(io); + if (n > len - 8 or n > 1 << 30) return error.BadSafetensors; + const json = try a.alloc(u8, @intCast(n)); + if (try file.readPositionalAll(io, json, 8) != json.len) return error.BadSafetensors; + try p.addHeader(name, json, len - 8 - n); + } + return p; + } + + /// model.safetensors.index.json's weight map, before any header is added. + pub fn useIndex(p: *Pack, json: []const u8) !void { + const a = p.arena.allocator(); + const parsed = try std.json.parseFromSliceLeaky(std.json.Value, a, json, .{}); + if (parsed != .object) return error.BadIndex; + const map = (parsed.object.get("weight_map") orelse return error.BadIndex); + if (map != .object) return error.BadIndex; + try p.owner.ensureTotalCapacity(a, @intCast(map.object.count())); + var it = map.object.iterator(); + while (it.next()) |e| { + if (e.value_ptr.* != .string) return error.BadIndex; + p.owner.putAssumeCapacity(e.key_ptr.*, e.value_ptr.string); + } + } + + /// A shard's header JSON (`data_len` bytes follow; maxInt(u64) for dumps); a name already indexed is refused. + pub fn addHeader(p: *Pack, file_name: []const u8, json: []const u8, data_len: u64) !void { + const a = p.arena.allocator(); + const id: u32 = @intCast(p.files.items.len); + try p.files.append(a, try a.dupe(u8, file_name)); + const header = try st.parseHeader(a, json, @intCast(data_len)); + try p.names.ensureUnusedCapacity(a, header.count()); + var it = header.iterator(); + while (it.next()) |e| { + const v = e.value_ptr.*; + if (p.owner.count() > 0) { + const f = p.owner.get(e.key_ptr.*) orelse continue; + if (!std.mem.eql(u8, f, file_name)) continue; + } + const gop = p.names.getOrPutAssumeCapacity(e.key_ptr.*); + if (gop.found_existing) return error.DuplicateTensor; + gop.value_ptr.* = .{ .file = id, .dtype = v.dtype, .rank = v.rank, .shape = v.shape, .start = 8 + json.len + v.begin, .nbytes = v.end - v.begin }; + } + } + + /// A header dump (the 8-byte length, then the JSON; nothing after) as shard `file_name`. + pub fn addHeaderDump(p: *Pack, file_name: []const u8, bytes: []const u8) !void { + if (bytes.len < 8) return error.BadSafetensors; + const n = std.mem.readInt(u64, bytes[0..8], .little); + if (n > bytes.len - 8) return error.BadSafetensors; + try p.addHeader(file_name, bytes[8..][0..@intCast(n)], std.math.maxInt(u64)); + } + + pub fn get(p: *const Pack, name: []const u8) ?Info { + return p.names.get(name); + } + + pub fn has(p: *const Pack, name: []const u8) bool { + return p.names.contains(name); + } + + /// A required tensor; a missing name is logged and refused. + pub fn need(p: *const Pack, name: []const u8) !Info { + return p.get(name) orelse { + if (!quiet) log.warn("pack has no tensor {s}", .{name}); + return error.MissingTensor; + }; + } + + /// The EXL3 group at `prefix` (its trellis entry), checked int16 [K/16, N/16, words] with fp16 suh [K] and svh [N]. + pub fn group(p: *const Pack, prefix: []const u8, buf: []u8) !exl3.Group { + const t = try p.need(try std.fmt.bufPrint(buf, "{s}.trellis", .{prefix})); + if (t.dtype != .i16) return error.BadTrellis; + const g = try exl3.Group.fromShape(t.dims()); + const suh = try p.need(try std.fmt.bufPrint(buf, "{s}.suh", .{prefix})); + const svh = try p.need(try std.fmt.bufPrint(buf, "{s}.svh", .{prefix})); + if (!suh.is(.f16, &.{g.k()}) or !svh.is(.f16, &.{g.n()})) { + if (!quiet) log.warn("{s}: suh/svh do not match trellis K={d} N={d}", .{ prefix, g.k(), g.n() }); + return error.BadTrellis; + } + return g; + } + + /// The group's marker tensor name (mul1 or mcg) if it has one; its value picks the codebook at load time. + pub fn marker(p: *const Pack, prefix: []const u8, buf: []u8) !?[]const u8 { + for ([_][]const u8{ "mul1", "mcg" }) |m| { + const name = try std.fmt.bufPrint(buf, "{s}.{s}", .{ prefix, m }); + if (p.names.getKey(name)) |k| return k; + } + return null; + } + + /// `span` of tensor `info`'s bytes into `out` (span.bytes() long), run by run. + pub fn read(p: *const Pack, io: Io, info: Info, span: exl3.Span, out: []u8) !void { + if (out.len != span.bytes()) return error.BadSpan; + if (span.count > 0 and span.at(span.count - 1) + span.len > info.nbytes) return error.BadSpan; + if (p.dir.len == 0) return error.HeaderOnlyPack; + var path_buf: [Io.Dir.max_path_bytes]u8 = undefined; + const path = try std.fmt.bufPrint(&path_buf, "{s}/{s}", .{ p.dir, p.files.items[info.file] }); + const file = try Io.Dir.cwd().openFile(io, path, .{}); + defer file.close(io); + if (span.count <= 1) { + if (try file.readPositionalAll(io, out, info.start + span.offset) != out.len) return error.ShortRead; + dropCache(file, info.start + span.offset, out.len); + return; + } + // a strided part (a column split): one read covering every run, then the runs gathered + const first = span.at(0); + const cover: usize = @intCast(span.at(span.count - 1) + span.len - first); + const tmp = try std.heap.page_allocator.alloc(u8, cover); + defer std.heap.page_allocator.free(tmp); + if (try file.readPositionalAll(io, tmp, info.start + first) != cover) return error.ShortRead; + dropCache(file, info.start + first, cover); + const len: usize = @intCast(span.len); + const stride: usize = @intCast(span.stride); + for (0..@intCast(span.count)) |i| @memcpy(out[i * len ..][0..len], tmp[i * stride ..][0..len]); + } + + /// Drops the page cache of bytes just read (POSIX_FADV_DONTNEED): on unified memory it competes with the device. + fn dropCache(file: Io.File, offset: u64, len: usize) void { + if (@import("builtin").os.tag != .linux) return; + const linux = std.os.linux; + _ = linux.fadvise(file.handle, @intCast(offset), @intCast(len), linux.POSIX_FADV.DONTNEED); + } + + /// The codebook of the group at `prefix`: its marker's value, else the 3-instruction default. + pub fn codebook(p: *const Pack, io: Io, prefix: []const u8) !exl3.Codebook { + var buf: [256]u8 = undefined; + const name = (try p.marker(prefix, &buf)) orelse return .inst3; + const info = p.get(name).?; + if (info.nbytes < 4) return error.BadMarker; + var v: [4]u8 = undefined; + try p.read(io, info, .{ .offset = 0, .len = 4 }, &v); + return exl3.Codebook.fromMarker(std.mem.readInt(u32, &v, .little)); + } +}; + +fn lessStr(_: void, x: []const u8, y: []const u8) bool { + return std.mem.order(u8, x, y) == .lt; +} + +const testing = std.testing; + +/// One tensor of an in-memory image. +pub const Spec = struct { name: []const u8, dtype: []const u8, shape: []const usize, fill: u8 = 0 }; + +/// A safetensors image in memory: tensors in order, each filled with `fill`. +pub fn image(a: std.mem.Allocator, specs: []const Spec) ![]u8 { + var json: std.Io.Writer.Allocating = .init(a); + defer json.deinit(); + try json.writer.writeAll("{\"__metadata__\":{\"format\":\"pt\"}"); + var at: usize = 0; + for (specs) |s| { + var n = (DType.parse(s.dtype) orelse return error.UnsupportedDType).size(); + for (s.shape) |d| n *= d; + try json.writer.print(",\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ s.name, s.dtype }); + for (s.shape, 0..) |d, i| try json.writer.print("{s}{d}", .{ if (i == 0) "" else ",", d }); + try json.writer.print("],\"data_offsets\":[{d},{d}]}}", .{ at, at + n }); + at += n; + } + try json.writer.writeAll("}"); + const head = json.written(); + const out = try a.alloc(u8, 8 + head.len + at); + std.mem.writeInt(u64, out[0..8], head.len, .little); + @memcpy(out[8..][0..head.len], head); + var pos = 8 + head.len; + for (specs) |s| { + var n = DType.parse(s.dtype).?.size(); + for (s.shape) |d| n *= d; + @memset(out[pos..][0..n], s.fill); + pos += n; + } + return out; +} + +test "a folder by its headers: index order, groups checked, strided reads, codebook markers" { + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const a = testing.allocator; + const one = try image(a, &.{ + .{ .name = "x.trellis", .dtype = "I16", .shape = &.{ 2, 16, 16 }, .fill = 0 }, + .{ .name = "x.suh", .dtype = "F16", .shape = &.{32} }, + .{ .name = "x.svh", .dtype = "F16", .shape = &.{256} }, + .{ .name = "x.mcg", .dtype = "I32", .shape = &.{}, .fill = 0 }, + }); + defer a.free(one); + // tile (kt, nt) holds bytes of value 16 kt + nt, so a column part's runs are checkable + const hl = std.mem.readInt(u64, one[0..8], .little); + const data = one[8 + hl ..]; + for (0..2) |kt| for (0..16) |nt| @memset(data[(kt * 16 + nt) * 32 ..][0..32], @intCast(16 * kt + nt)); + std.mem.writeInt(u32, data[2 * 16 * 32 + 64 + 512 ..][0..4], exl3.Codebook.marker_mcg, .little); + const two = try image(a, &.{.{ .name = "y.weight", .dtype = "BF16", .shape = &.{ 3, 2 }, .fill = 7 }}); + defer a.free(two); + try tmp.dir.writeFile(io, .{ .sub_path = "b.safetensors", .data = one }); + try tmp.dir.writeFile(io, .{ .sub_path = "a.safetensors", .data = two }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var p = try Pack.open(a, io, dir); + defer p.deinit(); + try testing.expectEqualStrings("a.safetensors", p.files.items[0]); + var buf: [128]u8 = undefined; + const g = try p.group("x", &buf); + try testing.expectEqual(@as(u32, 256), g.n()); + const part = try g.part(.col, 1, 2); + const out = try a.alloc(u8, @intCast(part.trellis.bytes())); + defer a.free(out); + try p.read(io, p.get("x.trellis").?, part.trellis, out); + for (0..2) |kt| for (0..8) |j| try testing.expectEqual(@as(u8, @intCast(16 * kt + 8 + j)), out[(kt * 8 + j) * 32]); + try testing.expectEqual(exl3.Codebook.mcg, try p.codebook(io, "x")); + try testing.expectEqual(exl3.Codebook.inst3, try p.codebook(io, "y")); + try testing.expectError(error.MissingTensor, p.group("y", &buf)); + try testing.expectError(error.BadSpan, p.read(io, p.get("y.weight").?, .{ .offset = 8, .len = 8 }, out[0..8])); +} diff --git a/zig/src/families/deepseek_v41/phases.zig b/zig/src/families/deepseek_v41/phases.zig new file mode 100644 index 0000000000..bfd86f1caf --- /dev/null +++ b/zig/src/families/deepseek_v41/phases.zig @@ -0,0 +1,206 @@ +//! TF_DSV41_PHASES (0 | 1 | sync, _OUT=): where every millisecond of a serving round goes, by named phase. + +const std = @import("std"); + +pub const Name = enum { + prefill, + draft, + @"draft.pass", + window, + forward, + @"graph.stage", + @"graph.replay", + @"graph.capture", + @"engram.wait", + prefetch, + candidates, + commit, + /// the speculative pass's launch after a window (TF_DSV41_SPEC_DRAFT; its collect is `draft.pass`) + spec, + /// host time between top-level phases (the lanes' own work, token events, plan link) + host, + /// inside `draft.pass` (TF_DSV41_DRAFT_GRAPHS): the host waiting for the device pass and its copy + @"draft.wait", + /// inside `draft.pass`: the host's chain (Markov bias, choice, confidences) over the gathered candidates + @"draft.chain", + /// inside `draft.pass`, the host step (DRAFT_GRAPHS off): the logits' download, the host candidates, their gather + @"draft.host", +}; + +/// A gap between top-level phases this long or longer is idle (no live stream), not a round's host time. +pub const idle_ns: u64 = 50 * std.time.ns_per_ms; + +const count = std.enums.values(Name).len; + +pub const Mode = enum { off, wall, sync }; + +/// The stream the `sync` mode waits for. +pub const Sync = struct { ctx: *anyopaque, run: *const fn (ctx: *anyopaque) anyerror!void }; + +/// Counters beside the phases (an object without "ms", so phase filters skip it): written into the report. +pub const Extra = struct { ctx: *anyopaque, write: *const fn (ctx: *anyopaque, w: *std.Io.Writer) anyerror!void }; + +pub const Phases = struct { + mode: Mode, + sync: ?Sync = null, + extra: ?Extra = null, + ns: [count]u64 = @splat(0), + n: [count]u64 = @splat(0), + rounds: u64 = 0, + tokens: u64 = 0, + t0: u64, + path: ?[]const u8 = null, + reported: u64 = 0, + /// phases open now (a top-level phase starts at depth 0), and when the last top-level phase ended (0: none yet) + depth: u32 = 0, + last_end: u64 = 0, + + pub fn parseMode(v: []const u8) !Mode { + if (std.mem.eql(u8, v, "0") or v.len == 0) return .off; + if (std.mem.eql(u8, v, "1")) return .wall; + if (std.mem.eql(u8, v, "sync")) return .sync; + return error.BadPhasesMode; + } + + /// From TF_DSV41_PHASES / TF_DSV41_PHASES_OUT; null when off. + pub fn fromEnv(gpa: std.mem.Allocator) !?*Phases { + const mode = try parseMode(if (std.c.getenv("TF_DSV41_PHASES")) |v| std.mem.span(v) else "0"); + if (mode == .off) return null; + const p = try gpa.create(Phases); + p.* = .{ .mode = mode, .t0 = nowNs(), .path = if (std.c.getenv("TF_DSV41_PHASES_OUT")) |v| std.mem.span(v) else null }; + return p; + } + + pub fn reset(p: *Phases) void { + p.ns = @splat(0); + p.n = @splat(0); + p.rounds = 0; + p.tokens = 0; + p.t0 = nowNs(); + p.depth = 0; + p.last_end = 0; + } + + fn settle(p: *Phases) void { + if (p.mode != .sync) return; + const s = p.sync orelse return; + s.run(s.ctx) catch {}; + } + + pub fn add(p: *Phases, name: Name, ns: u64) void { + p.ns[@backingInt(name)] += ns; + p.n[@backingInt(name)] += 1; + } + + /// The report: {name: {ms, n, ms_avg}} plus wall_ms, mode, rounds, tokens. + pub fn write(p: *const Phases, w: *std.Io.Writer) !void { + const wall = @as(f64, @floatFromInt(nowNs() - p.t0)) / 1e6; + try w.print("{{\"mode\": \"{s}\", \"wall_ms\": {d:.2}, \"rounds\": {d}, \"tokens\": {d}", .{ if (p.mode == .sync) "sync" else "1", wall, p.rounds, p.tokens }); + for (0..count) |i| { + if (p.n[i] == 0) continue; + const ms = @as(f64, @floatFromInt(p.ns[i])) / 1e6; + try w.print(", \"{s}\": {{\"ms\": {d:.2}, \"n\": {d}, \"ms_avg\": {d:.3}}}", .{ @tagName(@as(Name, @fromBackingInt(@intCast(i)))), ms, p.n[i], ms / @as(f64, @floatFromInt(p.n[i])) }); + } + if (p.extra) |x| try x.write(x.ctx, w); + try w.writeAll("}\n"); + } + + /// The report into TF_DSV41_PHASES_OUT (when set). + pub fn flush(p: *Phases, gpa: std.mem.Allocator, io: std.Io) void { + const path = p.path orelse return; + var w: std.Io.Writer.Allocating = .init(gpa); + defer w.deinit(); + p.write(&w.writer) catch return; + std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = w.written() }) catch |e| std.log.err("phases: {s} not written ({t})", .{ path, e }); + p.reported = p.rounds; + } + + /// A verify window's round finished with `kept` tokens committed; the report is rewritten every 64 rounds. + pub fn round(p: *Phases, gpa: std.mem.Allocator, io: std.Io, kept: usize) void { + p.rounds += 1; + p.tokens += kept; + if (p.rounds - p.reported >= 64) p.flush(gpa, io); + } +}; + +/// One phase's timer: `const t = phases.start(p, .window); defer t.stop();` (p null: nothing). +pub const Timer = struct { + p: ?*Phases, + name: Name, + t0: u64, + + pub fn stop(t: Timer) void { + const p = t.p orelse return; + p.settle(); + const now = nowNs(); + p.add(t.name, now -| t.t0); + p.depth -|= 1; + if (p.depth == 0) p.last_end = now; + } +}; + +pub fn start(p: ?*Phases, name: Name) Timer { + const x = p orelse return .{ .p = null, .name = name, .t0 = 0 }; + if (x.depth == 0 and x.last_end != 0) { + const gap = nowNs() -| x.last_end; + if (gap < idle_ns) x.add(.host, gap); + } + x.depth += 1; + x.settle(); + return .{ .p = x, .name = name, .t0 = nowNs() }; +} + +pub fn nowNs() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +test "phases: the reference engine's names and report shape; sync mode syncs at both ends" { + const Count = struct { + n: u32 = 0, + fn run(ctx: *anyopaque) anyerror!void { + const c: *@This() = @ptrCast(@alignCast(ctx)); + c.n += 1; + } + }; + var c: Count = .{}; + var p: Phases = .{ .mode = .sync, .t0 = nowNs(), .sync = .{ .ctx = &c, .run = Count.run } }; + { + const t = start(&p, .@"draft.pass"); + defer t.stop(); + } + const off = start(null, .window); + off.stop(); + try std.testing.expectEqual(@as(u32, 2), c.n); + { + // a second top-level phase: one `host` gap before it; the nested phase opens none + const t = start(&p, .window); + defer t.stop(); + const n = start(&p, .forward); + n.stop(); + } + try std.testing.expectEqual(@as(u64, 1), p.n[@backingInt(Name.host)]); + p.last_end = nowNs() - idle_ns; // an idle gap: not counted + start(&p, .commit).stop(); + try std.testing.expectEqual(@as(u64, 1), p.n[@backingInt(Name.host)]); + p.add(.@"engram.wait", 1_500_000); + p.add(.@"engram.wait", 500_000); + p.rounds = 2; + p.tokens = 7; + var w: std.Io.Writer.Allocating = .init(std.testing.allocator); + defer w.deinit(); + try p.write(&w.writer); + const parsed = try std.json.parseFromSlice(std.json.Value, std.testing.allocator, w.written(), .{}); + defer parsed.deinit(); + const o = parsed.value.object; + try std.testing.expectEqualStrings("sync", o.get("mode").?.string); + try std.testing.expectEqual(@as(i64, 7), o.get("tokens").?.integer); + const ew = o.get("engram.wait").?.object; + try std.testing.expectEqual(@as(i64, 2), ew.get("n").?.integer); + try std.testing.expectEqual(@as(f64, 2.0), ew.get("ms").?.float); + try std.testing.expect(o.get("draft.pass") != null and o.get("window") != null); + try std.testing.expect(o.get("host") != null); // the gap between the two top-level phases above + try std.testing.expectEqual(Mode.off, try Phases.parseMode("0")); + try std.testing.expectError(error.BadPhasesMode, Phases.parseMode("nvtx")); +} diff --git a/zig/src/families/deepseek_v41/plan.zig b/zig/src/families/deepseek_v41/plan.zig new file mode 100644 index 0000000000..e730a4ee2d --- /dev/null +++ b/zig/src/families/deepseek_v41/plan.zig @@ -0,0 +1,295 @@ +//! A rank's weights planned from the pack's headers alone: group TP parts, routed layouts, byte totals, coverage. + +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel +const log = std.log.scoped(.dsv41); +const Config = @import("config.zig").Config; +const exl3 = @import("exl3.zig"); +const names = @import("names.zig"); +const Pack = @import("pack.zig").Pack; +const Info = @import("pack.zig").Info; + +/// One EXL3 group of the rank: its stored name and the rank's part. +pub const GroupPlan = struct { + proj: names.Proj, + prefix: []const u8, + split: exl3.Split, + full: exl3.Group, + part: exl3.Part, + /// marker tensor (mul1 / mcg) whose value picks the codebook at load; null: the 3-instruction default + marker: ?[]const u8, + /// wo_a: the global group index this slice is + slice: u32 = 0, +}; + +pub const NativePlan = struct { + name: []const u8, + info: Info, + slice: names.Slice, + /// bytes of the rank's slice as stored + bytes: u64, +}; + +/// One routed projection: each expert's group part, their layout on the device, and the stored prefixes. +pub const ExpertsPlan = struct { + proj: names.ExpertProj, + parts: []exl3.Part, + prefixes: [][]const u8, + layout: exl3.Layout, +}; + +pub const LayerPlan = struct { + index: u32, + prefix: []const u8, + mode: @import("config.zig").Mode, + groups: []GroupPlan, + natives: []NativePlan, + experts: [3]ExpertsPlan, + /// experts whose gate (w1) and up (w3) widths differ: x3gm cannot fuse them, the block goes to x3pf / grouped + gate_up_mismatch: u32, + + /// Routed trellis bytes plus the shared expert's: the quant repo's `rank_bytes`. + pub fn expertBytes(l: *const LayerPlan) u64 { + var n: u64 = 0; + for (l.experts) |e| n += e.layout.trellisBytes(); + for (l.groups) |g| switch (g.proj) { + .shared_w1, .shared_w2, .shared_w3 => n += g.part.trellis.bytes(), + else => {}, + }; + return n; + } + + /// Everything the rank holds for the block as stored: trellises, suh/svh, natives. + pub fn bytes(l: *const LayerPlan) u64 { + var n: u64 = 0; + for (l.groups) |g| n += g.part.bytes(); + for (l.natives) |v| n += v.bytes; + for (l.experts) |e| { + n += e.layout.trellisBytes(); + for (e.parts) |p| n += p.suh.bytes() + p.svh.bytes(); + } + return n; + } +}; + +pub const Dspark = struct { + main_proj: GroupPlan, + natives: []NativePlan, + confidence: ?NativePlan, +}; + +pub const Plan = struct { + arena: std.heap.ArenaAllocator, + rank: u32, + world: u32, + layers: []LayerPlan, + top: []NativePlan, + /// null when only some blocks were planned + head: ?GroupPlan, + dspark: ?Dspark, + /// stored names neither planned nor on the ignore list (the loader refuses a pack with any) + leftovers: [][]const u8, + + pub fn deinit(p: *Plan) void { + p.arena.deinit(); + p.* = undefined; + } + + pub fn bytes(p: *const Plan) u64 { + var n: u64 = if (p.head) |h| h.part.bytes() else 0; + for (p.top) |v| n += v.bytes; + for (p.layers) |*l| n += l.bytes(); + if (p.dspark) |d| { + n += d.main_proj.part.bytes(); + for (d.natives) |v| n += v.bytes; + if (d.confidence) |c| n += c.bytes; + } + return n; + } +}; + +pub const Options = struct { + rank: u32, + world: u32, + /// plan only these blocks (single-layer runs, header dumps): no top tensors, head, DSpark or coverage; null: all + blocks: ?[]const u32 = null, + /// with `blocks`: also the vocabulary rows, the final norm and the head (a window's finish) + top: bool = false, + /// with `blocks` and `top`: also DSpark's heads (main_proj, its norms, the Markov and confidence heads) + dspark: bool = false, +}; + +const Builder = struct { + a: std.mem.Allocator, + c: *const Config, + pack: *const Pack, + o: Options, + used: std.StringHashMapUnmanaged(void) = .empty, + buf: [256]u8 = undefined, + + fn mark(b: *Builder, name: []const u8) !void { + try b.used.put(b.a, b.pack.names.getKey(name) orelse return error.MissingTensor, {}); + } + + fn dupe(b: *Builder, s: []const u8) ![]const u8 { + return b.a.dupe(u8, s); + } + + fn group(b: *Builder, proj: names.Proj, prefix: []const u8, split: exl3.Split, k: u32, n: u32) !GroupPlan { + const g = try b.pack.group(prefix, &b.buf); + if ((k != 0 and g.k() != k) or (n != 0 and g.n() != n)) { + if (!quiet) log.warn("{s}: K={d} N={d}, expected K={d} N={d}", .{ prefix, g.k(), g.n(), k, n }); + return error.UnexpectedShape; + } + inline for (.{ "trellis", "suh", "svh" }) |part| try b.mark(try std.fmt.bufPrint(&b.buf, "{s}." ++ part, .{prefix})); + const marker = try b.pack.marker(prefix, &b.buf); + if (marker) |m| try b.mark(m); + return .{ .proj = proj, .prefix = try b.dupe(prefix), .split = split, .full = g, .part = try g.part(split, b.o.rank, b.o.world), .marker = if (marker) |m| try b.dupe(m) else null }; + } + + fn native(b: *Builder, name: []const u8, spec: names.NativeSpec) !NativePlan { + const info = try b.pack.need(name); + const want = spec.dims(); + var ok = spec.store.accepts(info.dtype) and info.rank == want.len; + if (ok) for (want, info.dims()) |w, d| { + ok = ok and w == d; + }; + if (!ok) { + if (!quiet) log.warn("{s}: {t} {any}, expected {t} storage of {any}", .{ name, info.dtype, info.dims(), spec.store, want }); + return error.UnexpectedTensor; + } + try b.mark(name); + const bytes = switch (spec.slice) { + .whole => info.nbytes, + .heads, .vocab_rows => blk: { + if (info.shape[0] % b.o.world != 0) return error.SplitNotEven; + break :blk info.nbytes / b.o.world; + }, + }; + return .{ .name = try b.dupe(name), .info = info, .slice = spec.slice, .bytes = bytes }; + } + + fn experts(b: *Builder, block: []const u8, index: u32, proj: names.ExpertProj) !ExpertsPlan { + const count = b.c.expertsOf(index).count; + const parts = try b.a.alloc(exl3.Part, count); + const prefixes = try b.a.alloc([]const u8, count); + const kn = proj.kn(b.c); + var codebook_marker: ?bool = null; + for (parts, prefixes, 0..) |*p, *pre, e| { + var nb: [128]u8 = undefined; + const name = try names.expertName(block, @intCast(e), proj, &nb); + const g = try b.group(.wq_a, name, proj.split(), kn[0], kn[1]); + // the codebook is one a projection; here the marker's presence, its value at load + const has = g.marker != null; + if (codebook_marker) |m| if (m != has) return error.CodebooksDiffer; + codebook_marker = has; + p.* = g.part; + pre.* = g.prefix; + } + return .{ .proj = proj, .parts = parts, .prefixes = prefixes, .layout = try exl3.Layout.of(b.a, parts) }; + } + + fn layer(b: *Builder, index: u32) !LayerPlan { + var pb: [32]u8 = undefined; + const block = try b.dupe(try names.blockPrefix(b.c, index, &pb)); + const has_shared = b.pack.has(try std.fmt.bufPrint(&b.buf, "{s}.ffn.shared_experts.w1.trellis", .{block})); + const specs = try names.blockGroups(b.c, index, b.o.rank, b.o.world, has_shared); + const groups = try b.a.alloc(GroupPlan, specs.len); + for (specs.items(), groups) |s, *g| { + var nb: [128]u8 = undefined; + g.* = try b.group(s.proj, try names.groupName(s, block, &nb), s.split, s.k, s.n); + g.slice = s.slice; + } + const nspecs = try names.blockNatives(b.c, index); + const natives = try b.a.alloc(NativePlan, nspecs.len); + for (nspecs.items(), natives) |s, *v| { + var nb: [128]u8 = undefined; + v.* = try b.native(try std.fmt.bufPrint(&nb, "{s}.{s}", .{ block, s.suffix }), s); + } + // the other ranks' wo_a groups are theirs to read: mark them so coverage holds on every rank + var g: u32 = 0; + while (g < b.c.o_groups) : (g += 1) { + var nb: [128]u8 = undefined; + const pre = try std.fmt.bufPrint(&nb, "{s}.attn.wo_a.slice.{d}", .{ block, g }); + inline for (.{ "trellis", "suh", "svh" }) |part| try b.mark(try std.fmt.bufPrint(&b.buf, "{s}." ++ part, .{pre})); + if (try b.pack.marker(pre, &b.buf)) |m| try b.mark(m); + } + var out: LayerPlan = .{ .index = index, .prefix = block, .mode = b.c.mode(index), .groups = groups, .natives = natives, .experts = undefined, .gate_up_mismatch = 0 }; + inline for (.{ .w1, .w2, .w3 }, 0..) |p, i| out.experts[i] = try b.experts(block, index, p); + for (out.experts[0].layout.words, out.experts[2].layout.words) |w1, w3| out.gate_up_mismatch += @intFromBool(w1 != w3); + return out; + } + + /// The latest mtp.i holding `suffix` (as a tensor, a group or a .weight). + fn findDspark(b: *Builder, suffix: []const u8, buf: []u8) !?[]const u8 { + var i = b.c.mtp_layers; + while (i > 0) { + i -= 1; + const base = try std.fmt.bufPrint(buf, "mtp.{d}.{s}", .{ i, suffix }); + for ([_][]const u8{ "", ".trellis", ".weight" }) |ext| { + var nb: [160]u8 = undefined; + if (b.pack.has(try std.fmt.bufPrint(&nb, "{s}{s}", .{ base, ext }))) return base; + } + } + return null; + } + + fn dspark(b: *Builder) !?Dspark { + if (b.c.mtp_layers == 0) return null; + var fb: [160]u8 = undefined; + const mp = (try b.findDspark(names.dspark_main_proj, &fb)) orelse return null; + const main_proj = try b.group(.wq_a, mp, .whole, 3 * b.c.hidden, b.c.hidden); + const specs = try names.dsparkNatives(b.c); + const natives = try b.a.alloc(NativePlan, specs.len); + for (specs.items(), natives) |s, *v| { + const stem = s.suffix[0 .. s.suffix.len - ".weight".len]; + const base = (try b.findDspark(stem, &fb)) orelse return error.MissingTensor; + var nb: [160]u8 = undefined; + v.* = try b.native(try std.fmt.bufPrint(&nb, "{s}.weight", .{base}), s); + } + var confidence: ?NativePlan = null; + if (try b.findDspark("confidence_head.proj", &fb)) |base| { + var nb: [160]u8 = undefined; + const name = try std.fmt.bufPrint(&nb, "{s}.weight", .{base}); + const info = try b.pack.need(name); + var n: u64 = 1; + for (info.dims()) |d| n *= d; + if (info.dtype != .f32 and info.dtype != .bf16 or n != b.c.hidden + b.c.dspark_markov_rank) return error.UnexpectedTensor; + try b.mark(name); + confidence = .{ .name = try b.dupe(name), .info = info, .slice = .whole, .bytes = info.nbytes }; + } + return .{ .main_proj = main_proj, .natives = natives, .confidence = confidence }; + } +}; + +/// The rank's plan; heads, groups and the vocabulary must split over `world` (blocks of 128). +pub fn build(gpa: std.mem.Allocator, c: *const Config, pack: *const Pack, o: Options) !Plan { + if (o.world == 0 or o.rank >= o.world or c.heads % o.world != 0 or c.o_groups % o.world != 0 or c.vocab % (128 * o.world) != 0) return error.BadWorld; + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + var b: Builder = .{ .a = arena.allocator(), .c = c, .pack = pack, .o = o }; + const full = o.blocks == null; + const layers = try b.a.alloc(LayerPlan, if (o.blocks) |bl| bl.len else c.allLayers()); + for (layers, 0..) |*l, i| { + const index: u32 = if (o.blocks) |bl| bl[i] else @intCast(i); + if (index >= c.allLayers()) return error.LayerOutOfBounds; + l.* = try b.layer(index); + } + var top: []NativePlan = &.{}; + var head: ?GroupPlan = null; + var ds: ?Dspark = null; + if (full or o.top) { + const tspecs = try names.topNatives(c); + top = try b.a.alloc(NativePlan, tspecs.len); + for (tspecs.items(), top) |s, *v| v.* = try b.native(s.suffix, s); + head = try b.group(.wq_a, names.head.suffix, .col, c.hidden, c.vocab); + if (full or o.dspark) ds = try b.dspark(); + } + // coverage: every stored name was planned or is knowingly left (vision, image routing, bf16 copies of groups) + var left: std.ArrayList([]const u8) = .empty; + if (full) for (pack.names.keys()) |k| { + if (b.used.contains(k) or names.ignored(k, pack)) continue; + try left.append(b.a, try b.dupe(k)); + }; + return .{ .arena = arena, .rank = o.rank, .world = o.world, .layers = layers, .top = top, .head = head, .dspark = ds, .leftovers = left.items }; +} diff --git a/zig/src/families/deepseek_v41/prepare.zig b/zig/src/families/deepseek_v41/prepare.zig new file mode 100644 index 0000000000..0869e34320 --- /dev/null +++ b/zig/src/families/deepseek_v41/prepare.zig @@ -0,0 +1,144 @@ +//! The run-time weight forms built at load from named host tensors: linear strips, routed + shared experts' tables. + +const std = @import("std"); + +/// Trellis int16 [kt, nt, words] -> int32 [nt/8, kt, 8, words/2]: each 128-column block's tiles in K order. +pub fn strips(gpa: std.mem.Allocator, src: []const u8, kt: usize, nt: usize, words: usize) ![]align(16) u8 { + if (nt % 8 != 0 or src.len != kt * nt * words * 2) return error.BadTrellis; + const run = 8 * words * 2; // 8 tiles of one K row inside one 128-column block + const out = try gpa.alignedAlloc(u8, .@"16", src.len); + const blocks = nt / 8; + for (0..blocks * kt) |i| { + const nb = i / kt; + const k = i % kt; + @memcpy(out[i * run ..][0..run], src[(k * nt + nb * 8) * words * 2 ..][0..run]); + } + return out; +} + +/// Strips words -> dense3's lanes layout (dense3.cu to_lanes_kernel): per strip, lane L's bits of each G-step group. +pub fn lanes(gpa: std.mem.Allocator, strip_words: []const u8, kt: usize, k2: usize) ![]align(16) u8 { + const G: usize = if (k2 == 16) 1 else 2; + const nb = k2 / 2; // bytes of a lane's run of one tile (4 K2 bits) + const strip = 4 * kt * 32 * k2; // bytes a strip + if (k2 < 4 or k2 % 2 != 0 or kt % G != 0 or strip_words.len % strip != 0) return error.BadStrips; + const out = try gpa.alignedAlloc(u8, .@"16", strip_words.len); + const tile = 16 * k2; // bytes a tile (4 K2 words) + for (0..strip_words.len / strip) |s| { + const src = strip_words[s * strip ..][0..strip]; + const dst = out[s * strip ..][0..strip]; + for (0..kt / G) |gi| { + const g0 = gi * G * 8 * tile; // the group's first byte (src and dst) + for (0..32) |L| for (0..8 * G) |tt| for (0..nb) |j| { + const sb = g0 + tt * tile + L * nb + j; // stream byte in the source + const lb = tt * nb + j; // stream byte in the lane's run + const w = lb / 4; // the lane's word + const db = g0 + 4 * (((w / 4) * 32 + L) * 4 + w % 4) + (lb % 4); // stream byte in the destination + dst[4 * (db / 4) + 3 - db % 4] = src[4 * (sb / 4) + 3 - sb % 4]; + }; + } + } + return out; +} + +/// One expert of a projection: its trellis's device address and K2, its scales as stored (fp16). +pub const Expert = struct { trellis: u64, k2: i32, suh: []const u8, svh: []const u8 }; + +/// A projection's tables: int64 addresses, int32 K2s, suh [E, K] and svh [E, N] fp16. +pub const Tables = struct { + ptrs: []i64, + k2s: []i32, + suh: []align(16) u8, + svh: []align(16) u8, + k2_lo: i32, + k2_hi: i32, + + pub fn deinit(t: *Tables, gpa: std.mem.Allocator) void { + gpa.free(t.ptrs); + gpa.free(t.k2s); + gpa.free(t.suh); + gpa.free(t.svh); + t.* = undefined; + } +}; + +/// The tables of `experts` in order, the shared expert last; every suh / svh must have the projection's K and N. +pub fn tables(gpa: std.mem.Allocator, experts: []const Expert) !Tables { + if (experts.len == 0) return error.NoExperts; + const kb = experts[0].suh.len; + const nb = experts[0].svh.len; + var t: Tables = .{ + .ptrs = try gpa.alloc(i64, experts.len), + .k2s = try gpa.alloc(i32, experts.len), + .suh = try gpa.alignedAlloc(u8, .@"16", experts.len * kb), + .svh = try gpa.alignedAlloc(u8, .@"16", experts.len * nb), + .k2_lo = std.math.maxInt(i32), + .k2_hi = 0, + }; + errdefer t.deinit(gpa); + for (experts, 0..) |e, i| { + if (e.suh.len != kb or e.svh.len != nb) return error.ExpertShapesDiffer; + t.ptrs[i] = @bitCast(e.trellis); + t.k2s[i] = e.k2; + @memcpy(t.suh[i * kb ..][0..kb], e.suh); + @memcpy(t.svh[i * nb ..][0..nb], e.svh); + t.k2_lo = @min(t.k2_lo, e.k2); + t.k2_hi = @max(t.k2_hi, e.k2); + } + return t; +} + +const testing = std.testing; + +test "strips: each 128-column block's tiles in K order" { + const a = testing.allocator; + const kt = 3; + const nt = 16; + const words = 2; // one int32 a tile, to keep the check readable + var src: [kt * nt * words * 2]u8 = undefined; + // tile (k, n) holds the byte 16 k + n + for (0..kt) |k| for (0..nt) |n| @memset(src[(k * nt + n) * words * 2 ..][0 .. words * 2], @intCast(16 * k + n)); + const out = try strips(a, &src, kt, nt, words); + defer a.free(out); + // out[nb][k][j] = tile (k, 8 nb + j) + for (0..2) |nb| for (0..kt) |k| for (0..8) |j| { + try testing.expectEqual(@as(u8, @intCast(16 * k + 8 * nb + j)), out[((nb * kt + k) * 8 + j) * words * 2]); + }; + try testing.expectError(error.BadTrellis, strips(a, src[0..8], kt, nt, words)); +} + +test "lanes: the numpy reference's layout (to_lanes) at every dense3 width" { + // fixtures/lanes-ref.bin: for K2 4, 6, 8, 10, 12, 16: u32 K2, u32 NB, u32 KT, strips words, lanes words + const a = testing.allocator; + var rest: []const u8 = @embedFile("fixtures/lanes-ref.bin"); + var cases: usize = 0; + while (rest.len > 0) : (cases += 1) { + const k2 = std.mem.readInt(u32, rest[0..4], .little); + const nb = std.mem.readInt(u32, rest[4..8], .little); + const kt = std.mem.readInt(u32, rest[8..12], .little); + const bytes = 4 * nb * kt * 32 * k2; + const src = rest[12..][0..bytes]; + const want = rest[12 + bytes ..][0..bytes]; + const got = try lanes(a, src, kt, k2); + defer a.free(got); + try testing.expectEqualSlices(u8, want, got); + rest = rest[12 + 2 * bytes ..]; + } + try testing.expectEqual(@as(usize, 6), cases); +} + +test "expert tables: addresses, K2s and stacked scales in order, the width range" { + const a = testing.allocator; + const s1 = [_]u8{ 1, 2, 3, 4 }; + const s2 = [_]u8{ 5, 6, 7, 8 }; + const v1 = [_]u8{ 9, 9 }; + const v2 = [_]u8{ 7, 7 }; + var t = try tables(a, &.{ .{ .trellis = 0x7f00_0000_1000, .k2 = 4, .suh = &s1, .svh = &v1 }, .{ .trellis = 0x7f00_0002_0000, .k2 = 10, .suh = &s2, .svh = &v2 } }); + defer t.deinit(a); + try testing.expectEqualSlices(i64, &.{ 0x7f00_0000_1000, 0x7f00_0002_0000 }, t.ptrs); + try testing.expectEqualSlices(i32, &.{ 4, 10 }, t.k2s); + try testing.expectEqualSlices(u8, &.{ 1, 2, 3, 4, 5, 6, 7, 8 }, t.suh); + try testing.expectEqualSlices(u8, &.{ 9, 9, 7, 7 }, t.svh); + try testing.expect(t.k2_lo == 4 and t.k2_hi == 10); + try testing.expectError(error.ExpertShapesDiffer, tables(a, &.{ .{ .trellis = 1, .k2 = 4, .suh = &s1, .svh = &v1 }, .{ .trellis = 2, .k2 = 4, .suh = s2[0..2], .svh = &v2 } })); +} diff --git a/zig/src/families/deepseek_v41/prepared.zig b/zig/src/families/deepseek_v41/prepared.zig new file mode 100644 index 0000000000..4f0e0f911f --- /dev/null +++ b/zig/src/families/deepseek_v41/prepared.zig @@ -0,0 +1,281 @@ +//! Prepared weight folders (TF_DSV41_PREPARED, _WRITE, _VERIFY): built layer images cached on disk for a fast boot. +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel +const named = @import("named.zig"); +const pack_mod = @import("pack.zig"); +const Io = std.Io; +const Allocator = std.mem.Allocator; + +const magic = "TFDSV41P"; +const version: u32 = 1; +const page: u64 = 4096; + +/// What the images depend on besides the pack's bytes. +pub const Key = struct { + rank: u32, + world: u32, + o_groups: u32, + blocks: []const u32, + dspark: bool, +}; + +/// The builders' source, so a change to how images are made never reads an old folder. +const sources = [_][]const u8{ @embedFile("named.zig"), @embedFile("plan.zig"), @embedFile("stage.zig"), @embedFile("pack.zig"), @embedFile("exl3.zig") }; + +pub const Stats = struct { hits: u32 = 0, misses: u32 = 0, written: u32 = 0, read_bytes: u64 = 0, read_ns: u64 = 0, write_ns: u64 = 0 }; + +pub const Cache = struct { + gpa: Allocator, + io: Io, + /// ``

//rank`` + root: []u8, + write: bool, + verify: bool, + stats: Stats = .{}, + + /// The cache TF_DSV41_PREPARED names, or null when unset. + pub fn fromEnv(gpa: Allocator, io: Io, p: *const pack_mod.Pack, key: Key) !?Cache { + const dir = std.mem.span(std.c.getenv("TF_DSV41_PREPARED") orelse return null); + if (std.mem.trim(u8, dir, " ").len == 0) return null; + const flag = struct { + fn on(name: [*:0]const u8) bool { + const v = std.c.getenv(name) orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); + } + }.on; + return try open(gpa, io, dir, try packDigest(io, p), key, flag("TF_DSV41_PREPARED_WRITE"), flag("TF_DSV41_PREPARED_VERIFY")); + } + + pub fn open(gpa: Allocator, io: Io, dir: []const u8, pack_digest: [32]u8, key: Key, write: bool, verify: bool) !Cache { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + h.update(magic); + h.update(std.mem.asBytes(&version)); + for (sources) |s| h.update(s); + h.update(&pack_digest); + h.update(std.mem.asBytes(&[_]u32{ key.rank, key.world, key.o_groups, @intFromBool(key.dspark), @intCast(key.blocks.len) })); + h.update(std.mem.sliceAsBytes(key.blocks)); + var d: [32]u8 = undefined; + h.final(&d); + const hex = std.fmt.bytesToHex(d[0..8].*, .lower); + const root = try std.fmt.allocPrint(gpa, "{s}/{s}/rank{d}", .{ dir, &hex, key.rank }); + errdefer gpa.free(root); + if (write) try Io.Dir.cwd().createDirPath(io, root); + return .{ .gpa = gpa, .io = io, .root = root, .write = write, .verify = verify }; + } + + pub fn deinit(c: *Cache) void { + c.gpa.free(c.root); + } + + fn path(c: *const Cache, buf: []u8, tag: []const u8, suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}/{s}.bin{s}", .{ c.root, tag, suffix }); + } + + /// Asks the kernel to read `.bin` ahead (POSIX_FADV_WILLNEED); a missing file is ignored, nothing waits. + pub fn willNeed(c: *const Cache, tag: []const u8) void { + var pbuf: [std.fs.max_path_bytes]u8 = undefined; + const p = c.path(pbuf[0 .. pbuf.len - 1], tag, "") catch return; + pbuf[p.len] = 0; + const linux = std.os.linux; + const rc = linux.openat(linux.AT.FDCWD, pbuf[0..p.len :0], .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) return; + const fd: linux.fd_t = @intCast(rc); + defer _ = linux.close(fd); + _ = linux.fadvise(fd, 0, 0, linux.POSIX_FADV.WILLNEED); + } + + /// Fills `b` from ``.bin``: true when read, false when there is no usable file (`b` is then untouched). + pub fn get(c: *Cache, b: *named.Builder, tag: []const u8) !bool { + var pbuf: [std.fs.max_path_bytes]u8 = undefined; + const t0 = Io.Clock.awake.now(c.io).toNanoseconds(); + const file = Io.Dir.cwd().openFile(c.io, try c.path(&pbuf, tag, ""), .{}) catch { + c.stats.misses += 1; + return false; + }; + defer file.close(c.io); + const got = readInto(c.gpa, c.io, file, b, c.verify) catch |e| { + if (!quiet) std.log.warn("dsv41 prepared: {s} unusable ({t}); building it from the pack", .{ pbuf[0 .. c.root.len + tag.len + 5], e }); + c.stats.misses += 1; + return false; + }; + c.stats.hits += 1; + c.stats.read_bytes += got; + c.stats.read_ns += @intCast(Io.Clock.awake.now(c.io).toNanoseconds() - t0); + return true; + } + + /// Writes ``.bin`` from `b` when writing is on and the file is missing. + pub fn put(c: *Cache, b: *const named.Builder, tag: []const u8) !void { + if (!c.write) return; + var pbuf: [std.fs.max_path_bytes]u8 = undefined; + var tbuf: [std.fs.max_path_bytes]u8 = undefined; + const final_path = try c.path(&pbuf, tag, ""); + const tmp_path = try c.path(&tbuf, tag, ".tmp"); + const t0 = Io.Clock.awake.now(c.io).toNanoseconds(); + const dir = Io.Dir.cwd(); + { + const file = try dir.createFile(c.io, tmp_path, .{ .truncate = true }); + defer file.close(c.io); + try writeOut(c.gpa, c.io, file, b); + } + try dir.rename(tmp_path, dir, final_path, c.io); + c.stats.written += 1; + c.stats.write_ns += @intCast(Io.Clock.awake.now(c.io).toNanoseconds() - t0); + } +}; + +/// The pack's shards as they are on disk: each file's name, size and modification time. +pub fn packDigest(io: Io, p: *const pack_mod.Pack) ![32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + var buf: [std.fs.max_path_bytes]u8 = undefined; + for (p.files.items) |f| { + h.update(f); + const st = Io.Dir.cwd().statFile(io, try std.fmt.bufPrint(&buf, "{s}/{s}", .{ p.dir, f }), .{}) catch |e| switch (e) { + error.FileNotFound => { + h.update("missing"); + continue; + }, + else => return e, + }; + h.update(std.mem.asBytes(&[_]i128{ @intCast(st.size), st.mtime.toNanoseconds() })); + } + var d: [32]u8 = undefined; + h.final(&d); + return d; +} + +const Entry = extern struct { + kind: u8, + rank: u8, + name_len: u16, + pad: u32 = 0, + shape: [4]u64, + offset: u64, + len: u64, + sha256: [32]u8, +}; + +fn writeOut(gpa: Allocator, io: Io, file: Io.File, b: *const named.Builder) !void { + var head: std.ArrayList(u8) = .empty; + defer head.deinit(gpa); + try head.appendSlice(gpa, magic); + try head.appendSlice(gpa, std.mem.asBytes(&version)); + try head.appendSlice(gpa, std.mem.asBytes(&@as(u32, @intCast(b.out.items.len)))); + var size: u64 = 16; + for (b.out.items) |n| size += @sizeOf(Entry) + n.name.len; + var at = std.mem.alignForward(u64, size, page); + for (b.out.items) |n| { + var e: Entry = .{ .kind = @backingInt(n.kind), .rank = n.rank, .name_len = @intCast(n.name.len), .shape = undefined, .offset = at, .len = n.bytes.len, .sha256 = undefined }; + for (&e.shape, n.shape) |*d, s| d.* = s; + std.crypto.hash.sha2.Sha256.hash(n.bytes, &e.sha256, .{}); + try head.appendSlice(gpa, std.mem.asBytes(&e)); + try head.appendSlice(gpa, n.name); + at = std.mem.alignForward(u64, at + n.bytes.len, page); + } + try file.writePositionalAll(io, head.items, 0); + var cursor = std.mem.alignForward(u64, size, page); + for (b.out.items) |n| { + try file.writePositionalAll(io, n.bytes, cursor); + cursor = std.mem.alignForward(u64, cursor + n.bytes.len, page); + } + try file.setLength(io, cursor); +} + +/// Reads every tensor of a file into `b` (each into its own buffer); returns payload bytes. On error `b` gets nothing. +fn readInto(gpa: Allocator, io: Io, file: Io.File, b: *named.Builder, verify: bool) !u64 { + var fixed: [16]u8 = undefined; + if (try file.readPositionalAll(io, &fixed, 0) != 16) return error.Truncated; + if (!std.mem.eql(u8, fixed[0..8], magic)) return error.BadMagic; + if (std.mem.readInt(u32, fixed[8..12], .little) != version) return error.BadVersion; + const count = std.mem.readInt(u32, fixed[12..16], .little); + const total = try file.length(io); + var got: std.ArrayList(named.Named) = .empty; + errdefer { + for (got.items) |n| { + gpa.free(n.name); + gpa.free(n.bytes); + } + got.deinit(gpa); + } + var off: u64 = 16; + var bytes: u64 = 0; + for (0..count) |_| { + var e: Entry = undefined; + if (try file.readPositionalAll(io, std.mem.asBytes(&e), off) != @sizeOf(Entry)) return error.Truncated; + off += @sizeOf(Entry); + if (e.offset + e.len > total or e.rank > 4 or e.kind > @backingInt(named.Kind.i16)) return error.BadEntry; + const name = try gpa.alloc(u8, e.name_len); + errdefer gpa.free(name); + if (try file.readPositionalAll(io, name, off) != name.len) return error.Truncated; + off += e.name_len; + const data = try gpa.alignedAlloc(u8, .@"16", e.len); + errdefer gpa.free(data); + if (try file.readPositionalAll(io, data, e.offset) != data.len) return error.Truncated; + if (verify) { + var sha: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(data, &sha, .{}); + if (!std.mem.eql(u8, &sha, &e.sha256)) return error.DigestMismatch; + } + var n: named.Named = .{ .name = name, .kind = @fromBackingInt(@intCast(e.kind)), .shape = undefined, .rank = e.rank, .bytes = data }; + for (&n.shape, e.shape) |*d, s| d.* = @intCast(s); + try got.append(gpa, n); + bytes += e.len; + } + try b.out.appendSlice(b.gpa, got.items); + got.deinit(gpa); + return bytes; +} + +test "a layer's images round-trip through its prepared file, and another key misses" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}", .{&tmp.sub_path}); + defer gpa.free(root); + var p = pack_mod.Pack.init(gpa); + defer p.deinit(); + const key: Key = .{ .rank = 1, .world = 2, .o_groups = 8, .blocks = &.{ 0, 1, 2 }, .dspark = true }; + var c = try Cache.open(gpa, io, root, try packDigest(io, &p), key, true, true); + defer c.deinit(); + var b: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer b.deinit(); + const sizes = [_]usize{ 4096, 5000, 16, 1 }; + for (sizes, 0..) |n, i| { + const bytes = try gpa.alignedAlloc(u8, .@"16", n); + for (bytes, 0..) |*x, j| x.* = @truncate(j *% 31 +% i); + try b.out.append(gpa, .{ .name = try std.fmt.allocPrint(gpa, "L1.t{d}.trellis", .{i}), .kind = @fromBackingInt(@intCast(i % 4)), .shape = .{ n, 1, 1, 1 }, .rank = @intCast(1 + i % 4), .bytes = bytes }); + } + var miss: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer miss.deinit(); + try std.testing.expect(!try c.get(&miss, "L1")); + try c.put(&b, "L1"); + var back: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer back.deinit(); + try std.testing.expect(try c.get(&back, "L1")); + try std.testing.expectEqual(b.out.items.len, back.out.items.len); + for (b.out.items, back.out.items) |x, y| { + try std.testing.expectEqualStrings(x.name, y.name); + try std.testing.expectEqual(x.kind, y.kind); + try std.testing.expectEqual(x.rank, y.rank); + try std.testing.expectEqualSlices(usize, &x.shape, &y.shape); + try std.testing.expectEqualSlices(u8, x.bytes, y.bytes); + try std.testing.expect(std.mem.isAligned(@intFromPtr(y.bytes.ptr), 16)); + } + try std.testing.expectEqual(@as(u32, 1), c.stats.hits); + // another rank (or blocks, pack, builder source) is another folder + var other = try Cache.open(gpa, io, root, try packDigest(io, &p), .{ .rank = 0, .world = 2, .o_groups = 8, .blocks = &.{ 0, 1, 2 }, .dspark = true }, false, false); + defer other.deinit(); + var none: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer none.deinit(); + try std.testing.expect(!try other.get(&none, "L1")); + // a damaged payload is refused with VERIFY and the builder stays empty + var pbuf: [std.fs.max_path_bytes]u8 = undefined; + const f = try Io.Dir.cwd().openFile(io, try c.path(&pbuf, "L1", ""), .{ .mode = .read_write }); + try f.writePositionalAll(io, "x", page); + f.close(io); + var bad: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer bad.deinit(); + try std.testing.expect(!try c.get(&bad, "L1")); + try std.testing.expectEqual(@as(usize, 0), bad.out.items.len); +} diff --git a/zig/src/families/deepseek_v41/sampling_kernels.zig b/zig/src/families/deepseek_v41/sampling_kernels.zig new file mode 100644 index 0000000000..2bdbc76406 --- /dev/null +++ b/zig/src/families/deepseek_v41/sampling_kernels.zig @@ -0,0 +1,118 @@ +//! sampling.cu's launches in their own optional fatbin (`dsv41_sampling`): nucleus statistics and packed rows. + +const std = @import("std"); +const cuda = @import("cuda"); + +const Blob = struct { + pub const bytes align(16) = @embedFile("dsv41_fatbin_dsv41_sampling").*; +}; +const image: []const u8 = &Blob.bytes; + +/// The build embedded the kernel (-Dnvcc, or -Dfatbins with dsv41_sampling.fatbin). +pub const available = image.len > 0; + +pub const Functions = struct { + module: cuda.Module, + stats_fn: cuda.Function, + pack_fn: cuda.Function, + /// the split statistics (every SM; null in a fatbin built before them: `stats` then runs one CTA a row) + split: ?Split = null, + /// why `split` is null: the symbol that did not resolve, and the driver's error + split_missing: ?struct { symbol: [:0]const u8, err: anyerror } = null, + + pub const Split = struct { max: cuda.Function, sum: cuda.Function, fin: cuda.Function }; + + /// Columns a CTA of the split statistics covers (64,640 a rank: 32 chunks a row). + pub const chunk: usize = 2048; + + pub fn load(d: *const cuda.Driver) !Functions { + if (!available) return error.SamplingKernelNotBuilt; + var m = try cuda.Module.load(d, image); + errdefer m.unload(); + var f: Functions = .{ + .module = m, + .stats_fn = try m.function("_ZN14dsv41_sampling12stats_kernelEPKfxidPd"), + .pack_fn = try m.function("_ZN14dsv41_sampling11pack_kernelEPKfPKxiiPf"), + }; + var fs: [3]cuda.Function = undefined; + for (split_syms, &fs) |sym, *x| x.* = m.function(sym) catch |e| { + f.split_missing = .{ .symbol = sym, .err = e }; + return f; + }; + f.split = .{ .max = fs[0], .sum = fs[1], .fin = fs[2] }; + return f; + } + + const split_syms = [3][:0]const u8{ "dsv41_sampling_stats_max", "dsv41_sampling_stats_sum", "dsv41_sampling_stats_fin" }; + + /// Bytes of device scratch `statsSplit` needs for `rows` rows of C columns. + pub fn splitScratch(rows: usize, C: usize) usize { + const chunks = std.math.divCeil(usize, C, chunk) catch unreachable; + return std.mem.alignForward(usize, 4 * rows * chunks, 8) + 8 * rows * chunks; + } + + /// `stats` over every SM (sum in chunk order); `work` holds `splitScratch(rows, C)` bytes; error.NoSplit without. + pub fn statsSplit(f: *const Functions, s: cuda.Stream, lg: u64, ld: usize, rows: usize, C: usize, t: f64, work: u64, out: u64) !void { + const sp = f.split orelse return error.NoSplit; + if (rows < 1 or C < 1 or !(t > 0)) return error.Shape; + const chunks = std.math.divCeil(usize, C, chunk) catch unreachable; + const pmax = work; + const psum = work + std.mem.alignForward(usize, 4 * rows * chunks, 8); + const grid: cuda.launch.Dim3 = .{ .x = @intCast(chunks), .y = @intCast(rows), .z = 1 }; + const block: cuda.launch.Dim3 = .{ .x = 256, .y = 1, .z = 1 }; + { + var args: cuda.Args = .{}; + args.add(lg); + args.add(@as(i64, @intCast(ld))); + args.add(@as(c_int, @intCast(C))); + args.add(@as(c_int, @intCast(chunk))); + args.add(pmax); + try cuda.launch.launch(sp.max, .{ .grid = grid, .block = block }, s, &args); + } + { + var args: cuda.Args = .{}; + args.add(lg); + args.add(@as(i64, @intCast(ld))); + args.add(@as(c_int, @intCast(C))); + args.add(@as(c_int, @intCast(chunk))); + args.add(t); + args.add(pmax); + args.add(psum); + try cuda.launch.launch(sp.sum, .{ .grid = grid, .block = block }, s, &args); + } + var args: cuda.Args = .{}; + args.add(pmax); + args.add(psum); + args.add(@as(c_int, @intCast(chunks))); + args.add(out); + try cuda.launch.launch(sp.fin, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 32, .y = 1, .z = 1 } }, s, &args); + } + + pub fn unload(f: *Functions) void { + f.module.unload(); + } + + /// nucleus.row_stats: float64 [rows, 2] at `out`, each row's (max, sum exp(v / t - max / t)) over its C columns. + pub fn stats(f: *const Functions, s: cuda.Stream, lg: u64, ld: usize, rows: usize, C: usize, t: f64, out: u64) !void { + if (rows < 1 or C < 1 or !(t > 0)) return error.Shape; + var args: cuda.Args = .{}; + args.add(lg); + args.add(@as(i64, @intCast(ld))); + args.add(@as(c_int, @intCast(C))); + args.add(t); + args.add(out); + try cuda.launch.launch(f.stats_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 1024, .y = 1, .z = 1 } }, s, &args); + } + + /// cand_gather's packed rows: fp32 [rows, 2k] at `out` from topk_keys' vals fp32 / cols int64 [rows, k]. + pub fn pack(f: *const Functions, s: cuda.Stream, vals: u64, cols: u64, rows: usize, k: usize, id0: u32, out: u64) !void { + if (rows < 1 or k < 1) return error.Shape; + var args: cuda.Args = .{}; + args.add(vals); + args.add(cols); + args.add(@as(c_int, @intCast(k))); + args.add(@as(c_int, @intCast(id0))); + args.add(out); + try cuda.launch.launch(f.pack_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &args); + } +}; diff --git a/zig/src/families/deepseek_v41/stage.zig b/zig/src/families/deepseek_v41/stage.zig new file mode 100644 index 0000000000..726618d9b1 --- /dev/null +++ b/zig/src/families/deepseek_v41/stage.zig @@ -0,0 +1,148 @@ +//! A rank's EXL3 bytes read into host images the device buffers copy as is: group parts and ragged routed experts. + +const std = @import("std"); +const exl3 = @import("exl3.zig"); +const plan = @import("plan.zig"); +const Pack = @import("pack.zig").Pack; +const Io = std.Io; + +/// One group's part: trellis int16 [kt, nt, words], suh fp16 [K part], svh fp16 [N part]. +pub const GroupImage = struct { + trellis: []u8, + suh: []u8, + svh: []u8, + + pub fn deinit(g: *GroupImage, gpa: std.mem.Allocator) void { + gpa.free(g.trellis); + gpa.free(g.suh); + gpa.free(g.svh); + g.* = undefined; + } +}; + +pub fn group(gpa: std.mem.Allocator, io: Io, pack: *const Pack, g: plan.GroupPlan) !GroupImage { + var out: GroupImage = .{ .trellis = &.{}, .suh = &.{}, .svh = &.{} }; + errdefer out.deinit(gpa); + var buf: [256]u8 = undefined; + out.trellis = try gpa.alloc(u8, @intCast(g.part.trellis.bytes())); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.trellis", .{g.prefix})), g.part.trellis, out.trellis); + out.suh = try gpa.alloc(u8, @intCast(g.part.suh.bytes())); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.suh", .{g.prefix})), g.part.suh, out.suh); + out.svh = try gpa.alloc(u8, @intCast(g.part.svh.bytes())); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.svh", .{g.prefix})), g.part.svh, out.svh); + return out; +} + +/// A routed projection: experts' trellises back to back, suh [E, K part] and svh [E, N part] stacked, K2 per expert. +pub const ExpertsImage = struct { + trellis: []align(16) u8, + suh: []u8, + svh: []u8, + /// the kernels' int32 K2 table + k2: []i32, + /// byte offset of each expert's trellis in `trellis` (the int64 address table is the device base plus these) + offsets: []u64, + + pub fn deinit(e: *ExpertsImage, gpa: std.mem.Allocator) void { + gpa.free(e.trellis); + gpa.free(e.suh); + gpa.free(e.svh); + gpa.free(e.k2); + gpa.free(e.offsets); + e.* = undefined; + } +}; + +pub fn experts(gpa: std.mem.Allocator, io: Io, pack: *const Pack, x: *const plan.ExpertsPlan) !ExpertsImage { + const n = x.layout.count(); + const kpart = 2 * @as(u64, x.layout.k_tiles) * 16; + const npart = 2 * @as(u64, x.layout.n_tiles) * 16; + var out: ExpertsImage = .{ .trellis = &.{}, .suh = &.{}, .svh = &.{}, .k2 = &.{}, .offsets = &.{} }; + errdefer out.deinit(gpa); + out.trellis = try gpa.alignedAlloc(u8, .@"16", @intCast(x.layout.trellisBytes())); + out.suh = try gpa.alloc(u8, @intCast(n * kpart)); + out.svh = try gpa.alloc(u8, @intCast(n * npart)); + out.k2 = try gpa.alloc(i32, n); + out.offsets = try gpa.alloc(u64, n); + var buf: [256]u8 = undefined; + for (x.parts, x.prefixes, 0..) |p, pre, e| { + const at = x.layout.byteOffset(e); + out.offsets[e] = at; + out.k2[e] = @intCast(p.group.k2()); + const t = out.trellis[@intCast(at)..][0..@intCast(p.trellis.bytes())]; + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.trellis", .{pre})), p.trellis, t); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.suh", .{pre})), p.suh, out.suh[@intCast(e * kpart)..][0..@intCast(kpart)]); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.svh", .{pre})), p.svh, out.svh[@intCast(e * npart)..][0..@intCast(npart)]); + } + return out; +} + +const testing = std.testing; +const image = @import("pack.zig").image; + +test "staged images: a column part's runs, ragged experts back to back, stacked scales" { + const a = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + // three experts of a gate projection (K 32, N 256): expert 1 at 3 bits, the others at 2 + const words = [_]usize{ 32, 48, 32 }; + var specs: [12]@import("pack.zig").Spec = undefined; + var nb: [12][48]u8 = undefined; + const shapes = [3][3]usize{ .{ 2, 16, words[0] }, .{ 2, 16, words[1] }, .{ 2, 16, words[2] } }; + for (0..3) |e| { + const base = e * 4; + specs[base] = .{ .name = try std.fmt.bufPrint(&nb[base], "layers.0.ffn.experts.{d}.w1.trellis", .{e}), .dtype = "I16", .shape = &shapes[e] }; + specs[base + 1] = .{ .name = try std.fmt.bufPrint(&nb[base + 1], "layers.0.ffn.experts.{d}.w1.suh", .{e}), .dtype = "F16", .shape = &.{32} }; + specs[base + 2] = .{ .name = try std.fmt.bufPrint(&nb[base + 2], "layers.0.ffn.experts.{d}.w1.svh", .{e}), .dtype = "F16", .shape = &.{256} }; + specs[base + 3] = .{ .name = try std.fmt.bufPrint(&nb[base + 3], "layers.0.ffn.experts.{d}.w1.mul1", .{e}), .dtype = "I32", .shape = &.{} }; + } + const file = try image(a, &specs); + defer a.free(file); + // every byte of expert e's tensors: trellis tile (kt, nt) = 64 e + 16 kt + nt; svh element j's low byte = j + const hl = std.mem.readInt(u64, file[0..8], .little); + var pos: usize = 8 + hl; + for (0..3) |e| { + const tile = 2 * words[e]; + for (0..2) |kt| for (0..16) |nt| @memset(file[pos + (kt * 16 + nt) * tile ..][0..tile], @intCast(64 * e + 16 * kt + nt)); + pos += 2 * 16 * tile; + @memset(file[pos..][0..64], @intCast(200 + e)); + pos += 64; + for (0..256) |j| file[pos + 2 * j] = @intCast(j); + pos += 512 + 4; + } + try tmp.dir.writeFile(io, .{ .sub_path = "x.safetensors", .data = file }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var pack = try Pack.open(a, io, dir); + defer pack.deinit(); + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var parts: [3]exl3.Part = undefined; + var prefixes: [3][]const u8 = undefined; + for (0..3) |e| { + var buf: [128]u8 = undefined; + prefixes[e] = try std.fmt.bufPrint(&nb[e], "layers.0.ffn.experts.{d}.w1", .{e}); + parts[e] = try (try pack.group(prefixes[e], &buf)).part(.col, 1, 2); + } + const x: plan.ExpertsPlan = .{ .proj = .w1, .parts = &parts, .prefixes = &prefixes, .layout = try exl3.Layout.of(arena.allocator(), &parts) }; + var img = try experts(a, io, &pack, &x); + defer img.deinit(a); + try testing.expectEqualSlices(i32, &.{ 4, 6, 4 }, img.k2); + try testing.expectEqualSlices(u64, &.{ 0, 2 * 8 * 64, 2 * 8 * 64 + 2 * 8 * 96 }, img.offsets); + // rank 1's N tiles are 8..15 of each K-tile row + for (0..3) |e| for (0..2) |kt| for (0..8) |j| { + const tile = 2 * words[e]; + try testing.expectEqual(@as(u8, @intCast(64 * e + 16 * kt + 8 + j)), img.trellis[@intCast(img.offsets[e] + (kt * 8 + j) * tile)]); + }; + // suh whole (a column split keeps K), svh this rank's 128 columns + try testing.expectEqual(@as(u8, 201), img.suh[64]); + try testing.expectEqual(@as(u8, 128), img.svh[0]); + try testing.expectEqual(@as(u8, 128), img.svh[256 + 0]); + try testing.expectEqual(@as(usize, 3 * 256), img.svh.len); + var buf: [128]u8 = undefined; + const gp: plan.GroupPlan = .{ .proj = .wq_a, .prefix = prefixes[1], .split = .row, .full = try pack.group(prefixes[1], &buf), .part = try (try pack.group(prefixes[1], &buf)).part(.whole, 0, 2), .marker = null }; + var gi = try group(a, io, &pack, gp); + defer gi.deinit(a); + try testing.expectEqual(@as(u8, 64 + 16 + 15), gi.trellis[(16 + 15) * 96]); +} diff --git a/zig/src/families/deepseek_v41/triton_capture.zig b/zig/src/families/deepseek_v41/triton_capture.zig new file mode 100644 index 0000000000..05ac448ecf --- /dev/null +++ b/zig/src/families/deepseek_v41/triton_capture.zig @@ -0,0 +1,121 @@ +//! The reference capture's Triton cubins (aot/aot.json, by hash) loaded one by one; launches packed as Triton does. + +const std = @import("std"); +const cuda = @import("cuda"); +const Value = std.json.Value; + +pub const Param = struct { name: []const u8, ty: []const u8 }; + +pub const Kernel = struct { + kernel: cuda.triton.Kernel, + params: []Param, + global_scratch: u32, + global_align: u32, +}; + +pub const Set = struct { + arena: std.heap.ArenaAllocator, + kernels: std.StringHashMapUnmanaged(Kernel) = .empty, + /// hashes whose cubin did not load, with the reason + failed: std.ArrayList([2][]const u8) = .empty, + + pub fn deinit(s: *Set) void { + var it = s.kernels.valueIterator(); + while (it.next()) |k| k.kernel.unload(); + s.arena.deinit(); + s.* = undefined; + } + + /// Every kernel `dir`/aot.json lists; one that does not load is recorded in `failed`, the rest stay usable. + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, device: cuda.abi.Device, dir: []const u8) !Set { + var s: Set = .{ .arena = std.heap.ArenaAllocator.init(gpa) }; + errdefer s.deinit(); + const a = s.arena.allocator(); + const text = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "aot.json" }), a, .limited(1 << 26)); + const root = try std.json.parseFromSliceLeaky(Value, a, text, .{}); + for (root.object.get("kernels").?.array.items) |kv| { + const k = kv.object; + const hash = k.get("hash").?.string; + const name = k.get("name").?.string; + const file = try std.fmt.allocPrint(a, "{s}/cubins/{s}.cubin", .{ dir, hash }); + const cubin = std.Io.Dir.cwd().readFileAllocOptions(io, file, gpa, .limited(1 << 26), .@"16", null) catch |e| { + try s.failed.append(a, .{ hash, @errorName(e) }); + continue; + }; + defer gpa.free(cubin); + const meta: cuda.triton.Meta = .{ + .name = name, + .num_warps = @intCast(k.get("num_warps").?.integer), + .num_ctas = @intCast(k.get("num_ctas").?.integer), + .shared = @intCast(k.get("shared").?.integer), + .global_scratch_size = @intCast(k.get("global_scratch").?.integer), + .global_scratch_align = @intCast(k.get("global_align").?.integer), + .profile_scratch_size = @intCast(k.get("profile_scratch").?.integer), + .launch_pdl = k.get("pdl").?.bool, + }; + const kernel = cuda.triton.Kernel.load(d, device, cubin, meta, try a.dupeSentinel(u8, name, 0)) catch |e| { + try s.failed.append(a, .{ hash, @errorName(e) }); + continue; + }; + const ps = k.get("params").?.array.items; + const params = try a.alloc(Param, ps.len); + for (ps, params) |p, *out| out.* = .{ .name = p.object.get("name").?.string, .ty = p.object.get("type").?.string }; + try s.kernels.put(a, hash, .{ .kernel = kernel, .params = params, .global_scratch = meta.global_scratch_size, .global_align = meta.global_scratch_align }); + } + return s; + } +}; + +/// A captured scalar's bits as the parameter's C type (Triton's ty_to_cpp: i1 is int8, floats by their bit patterns). +fn scalar(args: *cuda.Args, ty: []const u8, v: Value) !void { + const T = enum { i1, i8, i16, i32, i64, u8, u16, u32, u64, fp32, fp64 }; + const t = std.meta.stringToEnum(T, ty) orelse return error.UnsupportedScalar; + const o = v.object; + const kind = o.get("t").?.string; + if (t == .fp32 or t == .fp64) { + if (!std.mem.eql(u8, kind, "float")) return error.WrongScalar; + const bits = o.get(if (t == .fp32) "f32" else "f64").?.string; + const x = try std.fmt.parseInt(u64, bits[2..], 16); + if (t == .fp32) args.add(@as(u32, @intCast(x))) else args.add(x); + return; + } + const n: i64 = if (std.mem.eql(u8, kind, "bool")) @intFromBool(o.get("v").?.bool) else if (std.mem.eql(u8, kind, "int")) o.get("v").?.integer else return error.WrongScalar; + switch (t) { + .i1, .i8 => args.add(@as(i8, @intCast(n))), + .u8 => args.add(@as(u8, @intCast(n))), + .i16 => args.add(@as(i16, @intCast(n))), + .u16 => args.add(@as(u16, @intCast(n))), + .i32 => args.add(@as(i32, @intCast(n))), + .u32 => args.add(@as(u32, @intCast(n))), + .i64 => args.add(n), + .u64 => args.add(@as(u64, @bitCast(n))), + .fp32, .fp64 => unreachable, + } +} + +/// The address of a tensor argument (`addr`), the int it carries, or null for None. +pub const Resolve = *const fn (ctx: *anyopaque, arg: Value) anyerror!?u64; + +/// Packs `op`'s named arguments in runtime order and launches it on the captured grid (global scratch: `scratch`). +pub fn launch(k: *const Kernel, op: std.json.ObjectMap, stream: cuda.Stream, ctx: *anyopaque, resolve: Resolve, scratch: u64) !void { + var args: cuda.Args = .{}; + const given = op.get("args").?.array.items; + for (k.params) |p| { + const arg = for (given) |g| { + if (std.mem.eql(u8, g.object.get("name").?.string, p.name)) break g; + } else return error.MissingArgument; + if (p.ty[0] == '*') { + args.add(@as(u64, (try resolve(ctx, arg)) orelse 0)); + } else try scalar(&args, p.ty, arg); + } + const g = op.get("grid").?.array.items; + const grid: cuda.Dim3 = .{ .x = @intCast(g[0].integer), .y = @intCast(g[1].integer), .z = @intCast(g[2].integer) }; + try k.kernel.launchOn(grid, stream, &args, .{ .global = scratch }, &.{}); +} + +/// Global scratch bytes a launch of `op` needs (Triton: blocks x CTAs x the kernel's size). +pub fn scratchBytes(k: *const Kernel, op: std.json.ObjectMap) usize { + if (k.global_scratch == 0) return 0; + const g = op.get("grid").?.array.items; + return k.kernel.globalScratchBytes(.{ .x = @intCast(g[0].integer), .y = @intCast(g[1].integer), .z = @intCast(g[2].integer) }); +} diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index 3f0f7f72ee..bf6085416a 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -105,7 +105,7 @@ const Wire = struct { payload: Value = .null, fn out(w: *Wire) openai.Out { - return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole } }; + return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole, .keepalive = ping } }; } fn open(ctx: *anyopaque, prompt_len: usize) error{Closed}!void { @@ -130,6 +130,12 @@ const Wire = struct { w.status = status; w.payload = payload; } + + /// Anthropic's own keepalive event. + fn ping(ctx: *anyopaque) error{Closed}!void { + const w: *Wire = @ptrCast(@alignCast(ctx)); + return w.reply.send("ping", try w.reply.obj("ping")); + } }; const Call = struct { id: []const u8 = "", name: []const u8 = "", arguments: []const u8 = "", sent: usize = 0, done: bool = false }; diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 47ec679be7..8edf497c10 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -262,6 +262,7 @@ const Generation = struct { streaming_done: bool = false, first_ns: ?i96 = null, last_ns: ?i96 = null, + wrote_ns: ?i96 = null, reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so engine_done: bool = false, consumed: usize = 0, @@ -350,6 +351,7 @@ const Generation = struct { g.cancel(); return error.Cancelled; }; + g.wrote_ns = nowNs(g.srv.io); } /// End the request and wait for the engine; unfinished requests count as disconnects. @@ -395,6 +397,15 @@ const Generation = struct { g.cancel(); return error.Cancelled; } + if (chunk == null and @hasDecl(@TypeOf(gone), "beat")) { + const now = nowNs(g.srv.io); + const since = g.wrote_ns orelse now; + if (g.wrote_ns == null) g.wrote_ns = now; + if (gone.beat(now - since) catch { + g.cancel(); + return error.Cancelled; + }) g.wrote_ns = now; + } const tokens = chunk orelse continue; if (g.reason == null) try g.commit(tokens); } diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig new file mode 100644 index 0000000000..71b9514baf --- /dev/null +++ b/zig/src/server/drain_test.zig @@ -0,0 +1,379 @@ +//! A stop's drain over HTTP: serve.run with a real SIGTERM, a slow stub engine, bytes as tokens. +const std = @import("std"); +const api = @import("engine_api"); +const json = @import("json"); +const model_text = @import("model_text.zig"); +const serve = @import("serve.zig"); +const cli = @import("cli.zig"); +const routes = @import("routes.zig"); +const Allocator = std.mem.Allocator; +const testing = std.testing; + +/// Stops mid-reply: "@short" ends after 12 tokens, "@long" when halted or cancelled, "@hold" waits for `release`. +pub const Slow = struct { + gpa: Allocator, + io: std.Io, + halted: std.atomic.Value(bool) = .init(false), + reason: []const u8 = "", + active: std.atomic.Value(u32) = .init(0), + halts: u32 = 0, + closes: u32 = 0, + /// halts before the first close (the halt must come first) + halts_at_close: u32 = 0, + cancelled: std.atomic.Value(u64) = .init(0), + /// "@hold" replies stay silent until this is set + release: std.atomic.Value(bool) = .init(false), + + fn engine(s: *Slow) api.Engine { + return .{ .ctx = s, .vtable = &.{ .info = info, .submit = submit, .cancel = cancel, .status = status, .memory = memory } }; + } + + fn self(ctx: *anyopaque) *Slow { + return @ptrCast(@alignCast(ctx)); + } + + fn info(_: *anyopaque) api.Info { + return .{ .context_window = 4096, .name = "slow" }; + } + + fn cancel(ctx: *anyopaque, id: api.Id) void { + self(ctx).cancelled.store(id, .release); + } + + fn status(ctx: *anyopaque, out: *api.Status, _: []u32) void { + out.* = .{ .running = self(ctx).active.load(.acquire) }; + } + + fn memory(_: *anyopaque, _: bool) ?api.Memory { + return null; + } + + fn submit(ctx: *anyopaque, id: api.Id, r: *const api.Request, sink: api.Sink) api.SubmitError!void { + const s = self(ctx); + _ = s.active.fetchAdd(1, .acq_rel); + const t = std.Thread.spawn(.{}, run, .{ s, id, r, sink }) catch { + _ = s.active.fetchSub(1, .acq_rel); + return error.Busy; + }; + t.detach(); + } + + fn run(s: *Slow, id: api.Id, r: *const api.Request, sink: api.Sink) void { + defer _ = s.active.fetchSub(1, .acq_rel); + var arena = std.heap.ArenaAllocator.init(s.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = a.alloc(u8, r.prompt.len) catch return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = "out of memory" } }); + for (r.prompt, prompt) |t, *b| b.* = @truncate(t); + const long = std.mem.indexOf(u8, prompt, "@long") != null; + const word: []const u32 = if (long) &.{ ' ', 'l', 'a' } else &.{ ' ', 'o', 'k' }; + const hang = std.mem.indexOf(u8, prompt, "@hang") != null; + const wait = std.mem.indexOf(u8, prompt, "@wait") != null; + const hold = std.mem.indexOf(u8, prompt, "@hold") != null; + // a long prefill: nothing for 800 ms ("@wait"), or nothing until stopped ("@hang") + var quiet: u32 = 0; + while (hang or (wait and quiet < 80) or (hold and !s.release.load(.acquire))) : (quiet += 1) { + if (s.halted.load(.acquire)) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = s.reason } }); + if (s.cancelled.load(.acquire) == id) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .cancelled } }); + std.Io.sleep(s.io, .fromMilliseconds(10), .awake) catch {}; + } + sink.event(sink.ctx, id, &.{ .prefilled = 0 }); + var n: usize = 0; + while (true) : (n += 1) { + if (s.halted.load(.acquire)) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = s.reason } }); + if (s.cancelled.load(.acquire) == id) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .cancelled } }); + if (!long and n == 12) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .stop } }); + sink.event(sink.ctx, id, &.{ .tokens = word }); + std.Io.sleep(s.io, .fromMilliseconds(15), .awake) catch {}; + } + } + + fn halt(ctx: *anyopaque, reason: []const u8) void { + const s = self(ctx); + s.reason = reason; + s.halts += 1; + s.halted.store(true, .release); + // a round boundary: every reply has ended before halt returns + while (s.active.load(.acquire) > 0) std.Io.sleep(s.io, .fromMilliseconds(1), .awake) catch {}; + } + + fn close(ctx: *anyopaque) void { + const s = self(ctx); + if (s.closes == 0) s.halts_at_close = s.halts; + s.closes += 1; + } +}; + +/// Bytes as tokens; a prompt renders as its messages' JSON, so the engine sees the client's marker. +pub const TestText = struct { + fn text(t: *@This()) model_text.Text { + return .{ .ctx = t, .vtable = &.{ .encode = encode, .decode = decode, .token_id = tokenId, .token_string = tokenString, .vocab_size = vocabSize, .eos_ids = eosIds, .render = render, .template_source = templateSource } }; + } + + fn encode(_: *anyopaque, a: Allocator, input: []const u8, _: bool) model_text.Error![]u32 { + const ids = try a.alloc(u32, input.len); + for (input, ids) |byte, *id| id.* = byte; + return ids; + } + + fn decode(_: *anyopaque, a: Allocator, ids: []const u32) Allocator.Error![]u8 { + const decoded = try a.alloc(u8, ids.len); + for (ids, decoded) |id, *byte| byte.* = @intCast(id); + return decoded; + } + + fn tokenId(_: *anyopaque, _: []const u8) ?u32 { + return null; + } + + fn tokenString(_: *anyopaque, a: Allocator, id: u32) Allocator.Error![]u8 { + return std.fmt.allocPrint(a, "{d}", .{id}); + } + + fn vocabSize(_: *anyopaque) u32 { + return 256; + } + + fn eosIds(_: *anyopaque) []const u32 { + return &.{}; + } + + fn render(_: *anyopaque, a: Allocator, messages: json.Value, _: model_text.RenderOptions, _: *[]const u8) model_text.Error![]u8 { + return json.stringify(a, messages, .{}) catch error.OutOfMemory; + } + + fn templateSource(_: *anyopaque) []const u8 { + return ""; + } +}; + +pub var listen_port: std.atomic.Value(u16) = .init(0); + +fn onListen(port: u16) void { + listen_port.store(port, .release); +} + +pub const Reply = struct { status: u16, head: []const u8, body: []const u8 }; + +/// One request on its own connection (closed by us after it), the whole response read. +pub fn request(a: Allocator, io: std.Io, port: u16, method: []const u8, path: []const u8, body: []const u8) !Reply { + const addr: std.Io.net.IpAddress = .{ .ip4 = .loopback(port) }; + const stream = try addr.connect(io, .{ .mode = .stream }); + defer stream.close(io); + const head = try std.fmt.allocPrint(a, "{s} {s} HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\nContent-Length: {d}\r\nConnection: close\r\n\r\n{s}", .{ method, path, body.len, body }); + var wbuf: [4096]u8 = undefined; + var w = std.Io.net.Stream.Writer.init(stream, io, &wbuf); + try w.interface.writeAll(head); + try w.interface.flush(); + var rbuf: [16384]u8 = undefined; + var r = std.Io.net.Stream.Reader.init(stream, io, &rbuf); + const text = try r.interface.allocRemaining(a, .unlimited); + const split = std.mem.indexOf(u8, text, "\r\n\r\n") orelse return error.Head; + var payload: []const u8 = text[split + 4 ..]; + if (std.mem.indexOf(u8, text[0..split], "chunked") != null) payload = try unchunk(a, payload); + return .{ .status = try std.fmt.parseInt(u16, text[9..12], 10), .head = text[0..split], .body = payload }; +} + +fn unchunk(a: Allocator, s: []const u8) ![]const u8 { + var out: std.ArrayList(u8) = .empty; + var i: usize = 0; + while (i < s.len) { + const eol = std.mem.indexOfPos(u8, s, i, "\r\n") orelse break; + const n = try std.fmt.parseInt(usize, std.mem.trim(u8, s[i..eol], " "), 16); + if (n == 0) break; + try out.appendSlice(a, s[eol + 2 .. eol + 2 + n]); + i = eol + 2 + n + 2; + } + return out.items; +} + +pub const Background = struct { + a: Allocator, + io: std.Io, + port: u16, + body: []const u8, + reply: ?Reply = null, + failed: bool = false, + + fn run(b: *Background) void { + b.reply = request(b.a, b.io, b.port, "POST", "/v1/chat/completions", b.body) catch { + b.failed = true; + return; + }; + } +}; + +pub fn chatBody(a: Allocator, marker: []const u8, stream: bool) ![]const u8 { + return std.fmt.allocPrint(a, "{{\"messages\":[{{\"role\":\"user\",\"content\":\"go {s}\"}}],\"stream\":{},\"max_tokens\":4000}}", .{ marker, stream }); +} + +pub const Served = struct { + code: u8 = 255, + fn run(x: *Served, gpa: Allocator, io: std.Io, setup: serve.Setup) void { + x.code = serve.run(gpa, io, .{ .host = "127.0.0.1", .port = 0 }, setup); + } +}; + +/// serve.run on its own thread with `drain_s`, the engine `slow`; the port once it listens. +pub fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, tt: *TestText, x: *Served) !std.Thread { + listen_port.store(0, .release); + const t = try std.Thread.spawn(.{}, Served.run, .{ x, gpa, io, serve.Setup{ + .engine = slow.engine(), + .text = tt.text(), + .served = "slow", + .environ = env, + .on_listen = onListen, + .stop = .{ .ctx = slow, .halt = Slow.halt, .close = Slow.close }, + } }); + while (listen_port.load(.acquire) == 0) std.Io.sleep(io, .fromMilliseconds(2), .awake) catch {}; + return t; +} + +pub fn waitActive(io: std.Io, slow: *Slow, n: u32) void { + while (slow.active.load(.acquire) < n) std.Io.sleep(io, .fromMilliseconds(2), .awake) catch {}; +} + +test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finishes; at the deadline the rest end \"server restarting\" at a round boundary, then the engine closes" { + const gpa = testing.allocator; + const io = testing.io; + var tt: TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + + // 1. the reply in progress finishes within the drain: no halt, a whole reply; new work refused meanwhile + { + try env.put("TF_DRAIN_S", "10"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + const port = listen_port.load(.acquire); + // held through the checks below: a reply that ended first would end the drain and close the listener + var bg: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@hold", true) }; + const bt = try std.Thread.spawn(.{}, Background.run, .{&bg}); + errdefer slow.release.store(true, .release); + waitActive(io, &slow, 1); + try std.posix.raise(.TERM); + // the drain has begun once a new request is refused + var refused: Reply = undefined; + while (true) { + refused = try request(a, io, port, "POST", "/v1/chat/completions", try chatBody(a, "@short", false)); + if (refused.status == 503) break; + try testing.expectEqual(@as(u16, 200), refused.status); // only before the signal lands + } + try testing.expect(std.mem.indexOf(u8, refused.head, "Retry-After: 30") != null); + try testing.expect(std.mem.indexOf(u8, refused.body, "\"code\": \"server_restarting\"") != null); + const anth = try request(a, io, port, "POST", "/v1/messages", "{\"model\":\"x\",\"max_tokens\":5,\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}]}"); + try testing.expectEqual(@as(u16, 503), anth.status); + try testing.expect(std.mem.indexOf(u8, anth.body, "overloaded_error") != null); + const health = try request(a, io, port, "GET", "/health", ""); + try testing.expectEqual(@as(u16, 503), health.status); + try testing.expect(std.mem.indexOf(u8, health.head, "Retry-After: 30") != null); + try testing.expect(std.mem.indexOf(u8, health.body, "\"status\": \"draining\"") != null); + slow.release.store(true, .release); + bt.join(); + t.join(); + try testing.expect(!bg.failed); + try testing.expectEqual(@as(u16, 200), bg.reply.?.status); + try testing.expect(std.mem.indexOf(u8, bg.reply.?.body, "\"finish_reason\": \"stop\"") != null); + try testing.expect(std.mem.indexOf(u8, bg.reply.?.body, "restarting") == null); + try testing.expect(std.mem.endsWith(u8, std.mem.trimEnd(u8, bg.reply.?.body, "\r\n"), "data: [DONE]")); + try testing.expectEqual(@as(u8, 0), x.code); + try testing.expectEqual(@as(u32, 0), slow.halts); + try testing.expectEqual(@as(u32, 1), slow.closes); + } + // 2. replies that outlast the drain end "server restarting"; the engine is halted, then closed once + { + try env.put("TF_DRAIN_S", "0.3"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + const port = listen_port.load(.acquire); + var streamed: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@long", true) }; + var whole: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@long", false) }; + const st = try std.Thread.spawn(.{}, Background.run, .{&streamed}); + const wt = try std.Thread.spawn(.{}, Background.run, .{&whole}); + waitActive(io, &slow, 2); + try std.posix.raise(.TERM); + st.join(); + wt.join(); + t.join(); + try testing.expectEqual(@as(u8, 0), x.code); + try testing.expectEqual(@as(u32, 1), slow.halts); + try testing.expectEqual(@as(u32, 1), slow.closes); + try testing.expectEqual(@as(u32, 1), slow.halts_at_close); + try testing.expectEqualStrings(routes.restarting, slow.reason); + const sb = streamed.reply.?.body; + try testing.expectEqual(@as(u16, 200), streamed.reply.?.status); + try testing.expect(std.mem.indexOf(u8, sb, "\"type\": \"server_error\"") != null); + try testing.expect(std.mem.indexOf(u8, sb, routes.restarting) != null); + try testing.expect(std.mem.indexOf(u8, sb, "data: [DONE]") != null); + try testing.expect(whole.reply.?.status >= 500); + try testing.expect(std.mem.indexOf(u8, whole.reply.?.body, routes.restarting) != null); + } +} + +/// The race seam's state: the stop delivered between a request's count and its drain check. +const Gap = struct { + var io: std.Io = undefined; + var port: u16 = 0; + var fired: std.atomic.Value(bool) = .init(false); + var saw_draining: bool = false; + var drain_waited: bool = false; + var served: *Served = undefined; + + fn hook() void { + if (fired.swap(true, .acq_rel)) return; + std.posix.raise(.TERM) catch return; + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + // the drain has begun (it reads the count after it sets the flag): /health says so + var tries: u32 = 0; + while (tries < 500) : (tries += 1) { + const h = request(arena.allocator(), io, port, "GET", "/health", "") catch break; + if (h.status == 503 and std.mem.indexOf(u8, h.body, "draining") != null) { + saw_draining = true; + break; + } + std.Io.sleep(io, .fromMilliseconds(5), .awake) catch {}; + } + // this request is counted: the drain may not finish while it is between its count and its check + std.Io.sleep(io, .fromMilliseconds(300), .awake) catch {}; + drain_waited = served.code == 255; + } +}; + +test "a stop that lands between a request's count and its drain check: the request is refused (503) and the drain waited for it, never lost" { + const gpa = testing.allocator; + const io = testing.io; + var tt: TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + try env.put("TF_DRAIN_S", "10"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + Gap.io = io; + Gap.port = listen_port.load(.acquire); + Gap.served = &x; + Gap.fired.store(false, .release); + routes.counted_hook = Gap.hook; + defer routes.counted_hook = null; + const r = try request(a, io, Gap.port, "POST", "/v1/chat/completions", try chatBody(a, "@short", false)); + t.join(); + try testing.expect(Gap.fired.load(.acquire)); + try testing.expect(Gap.saw_draining); + try testing.expect(Gap.drain_waited); + try testing.expectEqual(@as(u16, 503), r.status); + try testing.expect(std.mem.indexOf(u8, r.body, "server_restarting") != null); + try testing.expectEqual(@as(u32, 0), slow.closes -| 1); // closed once, after + try testing.expectEqual(@as(u32, 0), slow.active.load(.acquire)); // nothing reached the engine + try testing.expectEqual(@as(u8, 0), x.code); +} diff --git a/zig/src/server/keepalive_test.zig b/zig/src/server/keepalive_test.zig new file mode 100644 index 0000000000..43435c1d04 --- /dev/null +++ b/zig/src/server/keepalive_test.zig @@ -0,0 +1,75 @@ +//! A silent stream keeps its client: keepalives on both routes before the first token; a failed one ends the request. +const std = @import("std"); +const server = @import("server.zig"); +const d = @import("drain_test.zig"); +const testing = std.testing; + +/// A raw request whose response is read until the connection ends. +fn messagesBody(a: std.mem.Allocator, marker: []const u8) ![]const u8 { + return std.fmt.allocPrint(a, "{{\"model\":\"x\",\"max_tokens\":50,\"stream\":true,\"messages\":[{{\"role\":\"user\",\"content\":\"go {s}\"}}]}}", .{marker}); +} + +/// Sends `body` streamed and hangs up once the response head arrived and the request reached `slow`. +fn leave(a: std.mem.Allocator, io: std.Io, port: u16, body: []const u8, slow: *d.Slow) !void { + const addr: std.Io.net.IpAddress = .{ .ip4 = .loopback(port) }; + const stream = try addr.connect(io, .{ .mode = .stream }); + defer stream.close(io); + const head = try std.fmt.allocPrint(a, "POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\nContent-Length: {d}\r\n\r\n{s}", .{ body.len, body }); + var wbuf: [4096]u8 = undefined; + var w = std.Io.net.Stream.Writer.init(stream, io, &wbuf); + try w.interface.writeAll(head); + try w.interface.flush(); + var rbuf: [64]u8 = undefined; + var r = std.Io.net.Stream.Reader.init(stream, io, &rbuf); + _ = try r.interface.peek(1); + // the head goes out before the engine has the request: a hang-up before that leaves nothing to cancel + d.waitActive(io, slow, 1); +} + +test "a silent engine's stream gets keepalives before its first token (chat: SSE comments, messages: ping events)" { + const period = comptime std.meta.fieldInfo(server.Config, .sse_keepalive_s).attrs.defaultValue(f64).?; + try testing.expect(period > 0 and period <= 15); + const gpa = testing.allocator; + const io = testing.io; + var tt: d.TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + try env.put("TF_DRAIN_S", "5"); + try env.put("TF_SSE_KEEPALIVE_S", "0.1"); + var slow: d.Slow = .{ .gpa = gpa, .io = io }; + var x: d.Served = .{}; + const t = try d.start(gpa, io, &env, &slow, &tt, &x); + const port = d.listen_port.load(.acquire); + + const chat = try d.request(a, io, port, "POST", "/v1/chat/completions", try d.chatBody(a, "@wait", true)); + try testing.expectEqual(@as(u16, 200), chat.status); + const first_token = std.mem.indexOf(u8, chat.body, "\" ok\"") orelse { + std.debug.print("{s}\n", .{chat.body[0..@min(chat.body.len, 1500)]}); + return error.NoToken; + }; + try testing.expect(std.mem.count(u8, chat.body[0..first_token], ": keepalive\n\n") >= 3); + try testing.expect(std.mem.indexOf(u8, chat.body, "data: [DONE]") != null); + + const msg = try d.request(a, io, port, "POST", "/v1/messages", try messagesBody(a, "@wait")); + try testing.expectEqual(@as(u16, 200), msg.status); + const start = std.mem.indexOf(u8, msg.body, "event: message_start") orelse return error.NoStart; + const ping = std.mem.indexOf(u8, msg.body, "event: ping\ndata: {\"type\": \"ping\"}") orelse return error.NoPing; + const delta = std.mem.indexOf(u8, msg.body, "event: content_block_delta") orelse return error.NoDelta; + try testing.expect(start < ping and ping < delta); + + // a client that leaves a silent stream: the request is cancelled (the disconnect poll or a failed keepalive) + var idle: u32 = 0; + while (slow.active.load(.acquire) > 0 and idle < 300) : (idle += 1) std.Io.sleep(io, .fromMilliseconds(10), .awake) catch {}; + try leave(a, io, port, try d.chatBody(a, "@hang", true), &slow); + var waited: u32 = 0; + while (slow.cancelled.load(.acquire) == 0 and waited < 300) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(10), .awake) catch {}; + try testing.expect(slow.cancelled.load(.acquire) != 0); + + try std.posix.raise(.TERM); + t.join(); + try testing.expectEqual(@as(u8, 0), x.code); +} diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index f1cceb9ed4..32c40047f5 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -123,7 +123,8 @@ pub fn main(init: std.process.Init) !u8 { defer up.text.deinit(); var record: ?*@import("engine_api").recall_store.Store = null; defer if (record) |r| r.close(); // after the engine below has stopped reading it - defer up.engine.close(up.engine.ctx); + var closer: Closer = .{ .opened = up.engine }; + defer closer.closeOnce(); const served = hub.servedName(args.name, args.model, dir); problem = ""; record = recall_open.open(gpa, io, a, args, up.text, served, &problem) catch |e| { @@ -131,6 +132,7 @@ pub fn main(init: std.process.Init) !u8 { }; if (record) |r| if (!up.engine.engine.attachRecall(r)) return fail(not_lanes); return serve.run(gpa, io, args, .{ + .stop = .{ .ctx = &closer, .close = Closer.close }, .engine = up.engine.engine, .text = up.text.text(), .served = served, @@ -144,6 +146,22 @@ pub fn main(init: std.process.Init) !u8 { }); } +/// Closes once: serve.run closes the engine before freeing the server, main's defer covers paths that end before it. +const Closer = struct { + opened: engines.Opened, + closed: bool = false, + + fn closeOnce(c: *Closer) void { + if (c.closed) return; + c.closed = true; + c.opened.close(c.opened.ctx); + } + + fn close(ctx: *anyopaque) void { + closeOnce(@ptrCast(@alignCast(ctx))); + } +}; + /// TF_SLIDE_FORGET_FAIL in a -Dslide-test-hooks build: the fact whose check a delete fails, so its put-back is seen. fn forgetFail(env: *const std.process.Environ.Map) u32 { const v = env.get("TF_SLIDE_FORGET_FAIL") orelse return 0; diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index 45cf437364..1cdee2eda8 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -36,6 +36,8 @@ pub const Out = struct { event: *const fn (ctx: *anyopaque, payload: ?Value) error{Closed}!void, /// A whole JSON reply. reply: *const fn (ctx: *anyopaque, status: u16, payload: Value) void, + /// Bytes that keep a silent stream alive (an SSE comment, or the API's ping); null: none. + keepalive: ?*const fn (ctx: *anyopaque) error{Closed}!void = null, }; }; @@ -56,7 +58,7 @@ const ClientOut = struct { a: Allocator, fn out(c: *ClientOut) Out { - return .{ .ctx = c, .vt = &.{ .open = open, .event = event, .reply = reply } }; + return .{ .ctx = c, .vt = &.{ .open = open, .event = event, .reply = reply, .keepalive = keepalive } }; } fn open(ctx: *anyopaque, _: usize) error{Closed}!void { @@ -73,14 +75,31 @@ const ClientOut = struct { const c: *ClientOut = @ptrCast(@alignCast(ctx)); routes.sendValue(c.conn, c.a, status, payload); } + + fn keepalive(ctx: *anyopaque) error{Closed}!void { + const c: *ClientOut = @ptrCast(@alignCast(ctx)); + return sse.comment(c.conn); + } }; /// Whether the client has gone, read between rounds (Python's socket_cancellation). pub const Gone = struct { conn: *Conn, + /// The stream a keepalive goes to once it has been silent `every_ns` (0: never) + out: ?Out = null, + every_ns: i96 = 0, pub fn check(g: Gone) bool { return g.conn.peerGone(); } + + /// A keepalive if the stream has been silent long enough; a failed write means the client left. + pub fn beat(g: Gone, silent_ns: i96) error{Closed}!bool { + const o = g.out orelse return false; + const send = o.vt.keepalive orelse return false; + if (g.every_ns <= 0 or silent_ns < g.every_ns) return false; + try send(o.ctx); + return true; + } }; /// ``error_body``: OpenAI's error object, with param and code where clients key on them. @@ -514,8 +533,11 @@ const Run = struct { defer if (!handed) chat.release(r.srv, prepared.preparing); var sink_state: StreamSink = .{ .run = r, .tools = tools, .prompt_len = prepared.prompt_len }; handed = true; + var live = gone; // keepalives start once begin has opened the stream, after the engine admits the request + live.out = r.out; + live.every_ns = @intFromFloat(r.srv.config.sse_keepalive_s * std.time.ns_per_s); const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call, .begin = StreamSink.begin }; - var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, gone) else chat.generate(r.srv, &cx, prepared, sink, gone)) catch |e| { + var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, gone) else chat.generate(r.srv, &cx, prepared, sink, live)) catch |e| { if (!sink_state.opened) return r.unsent(&cx, e, field); switch (e) { error.Cancelled => return, diff --git a/zig/src/server/responses.zig b/zig/src/server/responses.zig index e84a74d0a9..0c72b48678 100644 --- a/zig/src/server/responses.zig +++ b/zig/src/server/responses.zig @@ -203,7 +203,7 @@ const Wire = struct { payload: Value = .null, fn out(w: *Wire) openai.Out { - return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole } }; + return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole, .keepalive = keepalive } }; } fn open(ctx: *anyopaque, _: usize) error{Closed}!void { @@ -223,6 +223,11 @@ const Wire = struct { w.status = status; w.payload = payload; } + + fn keepalive(ctx: *anyopaque) error{Closed}!void { + const w: *Wire = @ptrCast(@alignCast(ctx)); + return sse.comment(w.reply.conn); + } }; /// A Response built from chat-completion deltas, in order: reasoning, text, calls; streamed as typed events. diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index c7b31b5004..35484bffbc 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -56,6 +56,8 @@ test { _ = @import("stream_preflight_test.zig"); _ = @import("stream_dsml_test.zig"); _ = @import("logprobs_test.zig"); + _ = @import("drain_test.zig"); + _ = @import("keepalive_test.zig"); _ = @import("late_system_test.zig"); _ = @import("compact_test.zig"); _ = @import("request_log.zig"); diff --git a/zig/src/server/routes.zig b/zig/src/server/routes.zig index 5cf7f19eb2..8eee7e84ce 100644 --- a/zig/src/server/routes.zig +++ b/zig/src/server/routes.zig @@ -1,6 +1,7 @@ //! Method and path dispatch, as ``do_GET``, ``do_POST`` and ``do_DELETE`` route them. const speech = @import("speech.zig"); const std = @import("std"); +const api = @import("engine_api"); const json = @import("json"); const http_conn = @import("http_conn.zig"); const http_body = @import("http_body.zig"); @@ -77,19 +78,61 @@ fn discardBody(conn: *Conn, a: Allocator) void { fn post(srv: *Server, conn: *Conn, a: Allocator) void { const route = auth.routePath(conn.path); - if (std.mem.eql(u8, route, "/v1/audio/speech")) return speech.post(srv, conn, a) catch {}; - if (std.mem.eql(u8, route, "/v1/slide/learn")) return slide.learn(srv, conn, a); - if (std.mem.eql(u8, route, "/v1/imprint/resume")) return imprint.post(srv, conn, a, .where); - if (std.mem.eql(u8, route, "/v1/imprint/load")) return imprint.post(srv, conn, a, .load); - if (std.mem.endsWith(u8, route, "/decisions")) return decisions.post(srv, conn, a); - if (std.mem.endsWith(u8, route, "/systemone")) return systemone.post(srv, conn, a); - if (anthropic.route(conn.path)) return anthropic.post(srv, conn, a); - if (responses.route(route)) |rid| if (rid.len == 0) return responses.post(srv, conn, a); for (tokens.paths) |r| if (std.mem.eql(u8, route, r)) return tokens.post(srv, conn, a, std.mem.endsWith(u8, route, "/detokenize")); + const is_anthropic = anthropic.route(conn.path); + const is_responses = if (responses.route(route)) |rid| rid.len == 0 else false; + const is_decisions = std.mem.endsWith(u8, route, "/decisions"); + const is_systemone = std.mem.endsWith(u8, route, "/systemone"); + const is_slide = std.mem.eql(u8, route, "/v1/slide/learn"); + const is_speech = std.mem.eql(u8, route, "/v1/audio/speech"); + const imprint_op: ?api.ImprintOp = if (std.mem.eql(u8, route, "/v1/imprint/resume")) .where else if (std.mem.eql(u8, route, "/v1/imprint/load")) .load else null; const is_chat = std.mem.endsWith(u8, route, "/chat/completions"); - if (!is_chat and !std.mem.endsWith(u8, route, "/completions")) { + if (!is_slide and !is_speech and imprint_op == null and !is_anthropic and !is_responses and !is_decisions and !is_systemone and !is_chat and !std.mem.endsWith(u8, route, "/completions")) { discardBody(conn, a); return unknown(conn, a); } + // counted before the drain flag is read: a stop that sees no request in progress sees every later one refused + _ = srv.generating.fetchAdd(1, .seq_cst); + defer _ = srv.generating.fetchSub(1, .seq_cst); + if (@import("builtin").is_test) if (counted_hook) |h| h(); + if (srv.draining.load(.seq_cst)) { + discardBody(conn, a); + return draining(conn, is_anthropic); + } + if (is_speech) return speech.post(srv, conn, a) catch {}; + if (is_slide) return slide.learn(srv, conn, a); + if (imprint_op) |op| return imprint.post(srv, conn, a, op); + if (is_decisions) return decisions.post(srv, conn, a); + if (is_systemone) return systemone.post(srv, conn, a); + if (is_anthropic) return anthropic.post(srv, conn, a); + if (is_responses) return responses.post(srv, conn, a); openai.post(srv, conn, a, is_chat); } + +/// Tests only: the gap between a request's count and its drain check, where a stop must still see it counted. +pub var counted_hook: ?*const fn () void = null; + +/// Seconds a refused client is told to wait while the server drains for a restart (Retry-After). +pub const retry_after_s = 30; +pub const restarting = "server restarting: retry shortly"; + +/// A request that arrives while a stop drains: 503, Retry-After, the API's own error shape; the connection closes. +pub fn draining(conn: *Conn, anthropic_shape: bool) void { + const body = if (anthropic_shape) + "{\"type\": \"error\", \"error\": {\"type\": \"overloaded_error\", \"message\": \"" ++ restarting ++ "\"}}" + else + "{\"error\": {\"message\": \"" ++ restarting ++ "\", \"type\": \"server_error\", \"param\": null, \"code\": \"server_restarting\"}}"; + sendDraining(conn, body); +} + +/// 503 with Retry-After and Connection: close, `body` as JSON. +pub fn sendDraining(conn: *Conn, body: []const u8) void { + conn.close = true; + conn.startResponse(503, null) catch return; + conn.addHeader("Content-Type", "application/json") catch return; + conn.addHeader("Retry-After", std.fmt.comptimePrint("{d}", .{retry_after_s})) catch return; + var len: [24]u8 = undefined; + conn.addHeader("Content-Length", std.fmt.bufPrint(&len, "{d}", .{body.len}) catch unreachable) catch return; + conn.addHeader("Connection", "close") catch return; + conn.finish(body) catch {}; +} diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 1342f2df4a..c97c9b207f 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -10,16 +10,29 @@ const live = @import("live.zig"); const listener_mod = @import("listener.zig"); const model_text = @import("model_text.zig"); const server_mod = @import("server.zig"); +const routes = @import("routes.zig"); const recall_routes = @import("recall_routes.zig"); const Allocator = std.mem.Allocator; const posix = platform.posix; var stop_requested: std.atomic.Value(bool) = .init(false); +/// The signal that asked for the stop (0: none yet), for the shutdown line +var stop_signal: std.atomic.Value(i32) = .init(0); +/// Stop signals so far: a second one cuts the drain short +var stop_count: std.atomic.Value(u32) = .init(0); -fn onStop(_: posix.SIG) callconv(.c) void { +fn onStop(sig: posix.SIG) callconv(.c) void { + stop_signal.store(@intCast(@backingInt(sig)), .release); + _ = stop_count.fetchAdd(1, .acq_rel); stop_requested.store(true, .release); } +fn signalName(sig: i32) []const u8 { + if (sig == @backingInt(posix.SIG.TERM)) return "SIGTERM"; + if (sig == @backingInt(posix.SIG.INT)) return "SIGINT"; + return "a stop signal"; +} + /// SIGTERM and SIGINT stop the server cleanly; SIGHUP rereads the key file; SIGPIPE is ignored, as Python ignores it. pub fn installSignals(keys: bool) void { const stop: posix.Sigaction = .{ .handler = .{ .handler = onStop }, .mask = posix.sigemptyset(), .flags = 0 }; @@ -81,6 +94,15 @@ pub const Setup = struct { on_listen: ?*const fn (port: u16) void = null, /// Recall's record, attached to the engine already (--recall); null: off. recall: ?*api.recall_store.Store = null, + /// Run by `run` so the engine stops before the server is freed; null: the caller closes it after `run`. + stop: ?EngineStop = null, +}; + +/// `halt` (optional) ends the rounds at a boundary at the drain's deadline, so `close` never meets a running round. +pub const EngineStop = struct { + ctx: *anyopaque, + halt: ?*const fn (ctx: *anyopaque, reason: []const u8) void = null, + close: *const fn (ctx: *anyopaque) void, }; fn env(s: Setup, name: []const u8) ?[]const u8 { @@ -90,6 +112,9 @@ fn env(s: Setup, name: []const u8) ?[]const u8 { /// Serves until a stop signal; the process exit status. pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { + stop_requested.store(false, .release); + stop_count.store(0, .release); + stop_signal.store(0, .release); var arena: std.heap.ArenaAllocator = .init(gpa); defer arena.deinit(); const a = arena.allocator(); @@ -112,6 +137,14 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { std.debug.print("tensorfold: TENSORFOLD_SEED_SALT={s}: an integer\n", .{t}); return 1; } else 0; + const drain_s = seconds(env(s, "TF_DRAIN_S"), 20) orelse { + std.debug.print("tensorfold: TF_DRAIN_S: seconds, a number >= 0\n", .{}); + return 1; + }; + const keepalive_s = seconds(env(s, "TF_SSE_KEEPALIVE_S"), 15) orelse { + std.debug.print("tensorfold: TF_SSE_KEEPALIVE_S: seconds, a number >= 0\n", .{}); + return 1; + }; const config: server_mod.Config = .{ .served_name = s.served, .model_dir = s.model_dir, @@ -126,6 +159,7 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .use_drafts = !args.no_drafts, .vision_urls = args.vision_urls, .seed_salt = salt, + .sse_keepalive_s = keepalive_s, .request_log = env(s, "TENSORFOLD_REQUEST_LOG"), .dashboard = args.dashboard, .compact_at = if (args.compact_auto) .{ .auto = {} } else if (args.compact_fraction) |f| .{ .fraction = f } else null, @@ -170,6 +204,8 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { s.served, args.host, port, shownSampling(a, s.sampling), if (args.no_drafts or srv.info.plain_only) "off" else "on", if (window > 0) std.fmt.bufPrint(&window_text, "{d}", .{window}) catch "?" else "unlimited", loaded, }); + if (greedy(s.sampling)) + log.line("WARNING: default sampling is greedy (temperature <= 0) - a client that omits `temperature` decodes greedily, which on a long agentic turn can loop in reasoning and return empty content", .{}); if (args.thinking and std.mem.indexOf(u8, s.text.templateSource(), "enable_thinking") != null) log.line("thinking on (the chat template's default): replies reason in reasoning_content before the answer in content, and max_tokens counts both. --no-thinking turns it off; a request can send chat_template_kwargs {{\"enable_thinking\": false}}", .{}); if (s.recall) |r| recall_routes.line(r); @@ -178,17 +214,57 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { var ticker: live.Ticker = .{ .engine = s.engine }; const drawing = live.wanted(io, if (env(s, "TENSORFOLD_NO_LIVE")) |v| std.mem.eql(u8, v, "1") else false); if (drawing) ticker.start(); - const accept = std.Thread.spawn(.{}, server_mod.Server.serve, .{ srv, lis, &stop_requested }) catch return 1; + // the listener stays open through the drain: new requests get 503 + Retry-After, /health "draining" + var accept_stop: std.atomic.Value(bool) = .init(false); + const accept = std.Thread.spawn(.{}, server_mod.Server.serve, .{ srv, lis, &accept_stop }) catch return 1; + while (!stop_requested.load(.acquire)) std.Io.sleep(io, .fromMilliseconds(50), .awake) catch {}; + drain(srv, io, s.stop, drain_s); + accept_stop.store(true, .release); accept.join(); lis.close(); if (drawing) ticker.finish(); + s.engine.drain(); // kept states waiting for disk are written (under the engine's deadline) before the engine closes + // the engine stops while nothing runs on it: no round meets a CUDA driver torn down by the exit + if (s.stop) |st| st.close(st.ctx); var waited: u32 = 0; while (srv.open_connections.load(.acquire) > 0 and waited < 40) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(50), .awake) catch {}; - s.engine.drain(); // kept states waiting for disk are written (under the engine's deadline) before the process ends - if (srv.open_connections.load(.acquire) > 0) std.process.exit(0); // replies still open: end without freeing what they read + if (srv.open_connections.load(.acquire) > 0) { + log.line("shutting down: exiting with {d} idle connection(s) still open", .{srv.open_connections.load(.acquire)}); + std.process.exit(0); // connections still open: end without freeing what they read + } + log.line("shutting down: stopped", .{}); return 0; } +/// Work in progress gets `drain_s` to finish (new work 503s) before the rest is halted; a second signal cuts it short. +fn drain(srv: *server_mod.Server, io: std.Io, stop: ?EngineStop, drain_s: f64) void { + srv.draining.store(true, .seq_cst); + const signals = stop_count.load(.acquire); + log.line("shutting down ({s}): draining: new requests get 503; {d} request(s) in progress may finish within {d:.1} s", .{ signalName(stop_signal.load(.acquire)), srv.generating.load(.seq_cst), drain_s }); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + const deadline = t0 + @as(i96, @intFromFloat(drain_s * 1e9)); + while (srv.generating.load(.seq_cst) > 0 and std.Io.Clock.awake.now(io).toNanoseconds() < deadline and stop_count.load(.acquire) == signals) + std.Io.sleep(io, .fromMilliseconds(20), .awake) catch {}; + const left = srv.generating.load(.seq_cst); + if (left == 0) { + log.line("shutting down: drained", .{}); + return; + } + const cut = if (stop_count.load(.acquire) != signals) "a second stop signal" else "the drain deadline (TF_DRAIN_S)"; + log.line("shutting down: {s}: {d} request(s) still in progress end \"{s}\" at the next round boundary", .{ cut, left, routes.restarting }); + if (stop) |st| if (st.halt) |h| h(st.ctx, routes.restarting); + var waited: u32 = 0; + while (srv.generating.load(.seq_cst) > 0 and waited < 100) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(20), .awake) catch {}; +} + +/// A seconds setting (TF_DRAIN_S, TF_SSE_KEEPALIVE_S): unset or blank keeps the default; null: not a number >= 0. +fn seconds(raw: ?[]const u8, default: f64) ?f64 { + const t = std.mem.trim(u8, raw orelse "", " \t"); + if (t.len == 0) return default; + const x = std.fmt.parseFloat(f64, t) catch return null; + return if (x >= 0) x else null; +} + /// ``--host`` as an address: an IP literal, else a name the resolver knows (localhost first). fn resolve(io: std.Io, host: []const u8, port: u16) ?std.Io.net.IpAddress { if (std.Io.net.IpAddress.parse(host, port)) |ip| return ip else |_| {} @@ -196,6 +272,17 @@ fn resolve(io: std.Io, host: []const u8, port: u16) ?std.Io.net.IpAddress { return std.Io.net.IpAddress.resolve(io, host, port) catch null; } +/// Whether requests that name no temperature decode greedily (the startup warning). +fn greedy(sampling: ?json.Value) bool { + const v = sampling orelse return true; + const t = v.get("temperature") orelse return true; + return switch (t) { + .float => |f| f <= 0, + .int => |i| (std.fmt.parseFloat(f64, i) catch 0) <= 0, + else => true, + }; +} + /// The serving line's sampling: greedy, or each default as ``name value``. fn shownSampling(a: Allocator, sampling: ?json.Value) []const u8 { const v = sampling orelse return "greedy"; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index cd95098680..2c2ecfa185 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -69,6 +69,8 @@ pub const Config = struct { /// --learn-conversations: /v1/imprint answers; without it each of its routes refuses and names the flag. saved_conversations: bool = false, timeouts: http_conn.Timeouts = .{}, + /// TF_SSE_KEEPALIVE_S: a silent stream's keepalive period before its first token (serve.zig; 0: none) + sse_keepalive_s: f64 = 15, /// The open /health names these after its other keys; empty fields are left out. identity: Identity = .{}, /// The process environment GET /v1/benchmark/manifest reads (rank count, revision, the hardware probe); null: none. @@ -97,6 +99,10 @@ pub const Server = struct { next_id: std.atomic.Value(u64) = .init(1), /// Connections open now; a stop waits for them before freeing what they read. open_connections: std.atomic.Value(u32) = .init(0), + /// A stop is draining (serve.zig): new generation requests get 503 with Retry-After, /health answers "draining" + draining: std.atomic.Value(bool) = .init(false), + /// Generation requests in progress (routes.post), what a drain waits for; counted before `draining` is read + generating: std.atomic.Value(u32) = .init(0), /// The idle keepalive on the engine's queue, armed by --keep-warm; null when off or not Metal. keepalive: ?*api.keepalive.Keepalive = null, preparing: std.atomic.Value(i64) = .init(0), @@ -179,7 +185,8 @@ pub const Server = struct { const options = try json.newObject(cx.a); const defaults = if (srv.config.default_sampling) |d| try fields.parseNumbers(cx, d) else null; if (defaults) |d| try merge(options, cx, d); - try merge(options, cx, try fields.parseNumbers(cx, f)); + const asked = try fields.parseNumbers(cx, f); + try merge(options, cx, asked); const temp = number(options.get("temperature")) orelse 0; if (temp <= 0) return null; if (srv.info.greedy_only) return cx.refuse(if ((if (defaults) |d| number(d.get("temperature")) else null) orelse 0 > 0) diff --git a/zig/src/server/sse.zig b/zig/src/server/sse.zig index 0942a7e145..2438c41eee 100644 --- a/zig/src/server/sse.zig +++ b/zig/src/server/sse.zig @@ -19,6 +19,11 @@ pub fn data(conn: *Conn, a: Allocator, payload: json.Value) error{Closed}!void { return conn.writeAll(std.mem.concat(a, u8, &.{ "data: ", text, "\n\n" }) catch return error.Closed); } +/// An SSE comment line, which clients skip: it keeps a stream alive while nothing else is sent. +pub fn comment(conn: *Conn) error{Closed}!void { + return conn.writeAll(": keepalive\n\n"); +} + pub fn done(conn: *Conn) error{Closed}!void { return conn.writeAll("data: [DONE]\n\n"); } diff --git a/zig/src/server/status_routes.zig b/zig/src/server/status_routes.zig index 8235bbb8f9..3db0c6d335 100644 --- a/zig/src/server/status_routes.zig +++ b/zig/src/server/status_routes.zig @@ -12,6 +12,8 @@ const Allocator = std.mem.Allocator; /// With keys on, health names nothing about the model. pub fn health(srv: *Server, conn: *Conn, a: Allocator) !void { + // draining: the engine is not asked, it may be stopping + if (srv.draining.load(.seq_cst)) return routes.sendDraining(conn, try std.fmt.allocPrint(a, "{{\"status\": \"draining\", \"requests_in_progress\": {d}}}", .{srv.generating.load(.seq_cst)})); var status: api.Status = .{}; srv.engine.status(&status, &.{}); if (!status.healthy) return conn.sendJson(503, "{\"status\": \"unavailable\"}"); diff --git a/zig/src/sessions/disk.zig b/zig/src/sessions/disk.zig new file mode 100644 index 0000000000..f5ab58f06f --- /dev/null +++ b/zig/src/sessions/disk.zig @@ -0,0 +1,587 @@ +//! The NVMe tier: a file an entry under //rank/, written via .tmp + rename, streamed via the lanes. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const lanes_mod = @import("lanes.zig"); +const Lanes = lanes_mod.Lanes; +const Staging = lanes_mod.Staging; +const Op = lanes_mod.Op; +const up = lanes_mod.up; +const format = @import("format.zig"); +const Layout = format.Layout; +const Head = format.Head; +const Seg = format.Seg; +const PageStore = @import("pagestore.zig").PageStore; +const pool_mod = @import("pool.zig"); +const Family = pool_mod.Family; +const Digest = @import("prefix.zig").Digest; + +pub const none = std.math.maxInt(u32); + +pub const Options = struct { + root: []const u8, + compat: Digest, + rank: u32 = 0, + world: u32 = 1, + /// bytes of files the tier keeps (least recently used files without dependents go first) + budget: u64 = 64 << 30, + /// entries shorter than this are dropped, not parked + min_tokens: u64 = 1024, + /// O_DIRECT where the filesystem takes it (buffered with the page cache dropped otherwise) + direct: bool = true, + /// staging buffers and their size: the host transient of any park or restore + staging: u32 = 16, + chunk: u32 = 8 << 20, + lanes: u32 = 4, + /// delta files a chain may stack before a park writes the entry whole + max_chain: u32 = 8, +}; + +/// A piece of bounded state (rings, carries, taps): small, host bytes. +pub const Blob = struct { + name: []const u8, + bytes: []const u8, +}; + +/// What a park writes: the entry, its tokens from first_page on, and this rank's pages per family, in logical order. +pub const WriteSrc = struct { + key: Digest, + /// the file holding pages [0, first_page) (a delta), or null (a whole entry: first_page 0) + base: ?Digest = null, + tag: u32, + kind: u32 = 0, + page: u32, + pos: u64, + first_page: u32 = 0, + npages: u32, + /// tokens [first_page x page, pos) + ids: []const i32, + store: PageStore, + fence: u64 = 0, + families: []const Family, + /// per family: family-tensor pages of this rank's logical pages [first_page, npages), in logical order + pages: []const []const u32, + blobs: []const Blob = &.{}, +}; + +/// One file of a restore: read logical pages [lo, hi) of it into the slot. +pub const Extent = struct { + key: Digest, + first_page: u32, + lo: u32, + hi: u32, +}; + +/// Where a restore puts rows: per family, the family-tensor page of each of the slot's logical pages. +pub const ReadSink = struct { + store: PageStore, + families: []const Family, + page: u32, + world: u32, + rank: u32, + /// per family: family-tensor pages for logical pages [0, npages) + pages: []const []const u32, +}; + +/// The bounded state a restore returns (owned). +pub const Restored = struct { + blobs: []Blob, + + pub fn deinit(r: *Restored, gpa: std.mem.Allocator) void { + for (r.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(r.blobs); + } +}; + +/// The index record of one file. +pub const Rec = struct { + key: Digest, + base: u32 = none, + tag: u32, + kind: u32, + pos: u64, + first_page: u32, + npages: u32, + size: u64, + /// files stacked on this one + refs: u32 = 0, + /// files under it, this one included + depth: u32 = 1, + older: u32 = none, + newer: u32 = none, + live: bool = true, +}; + +pub const Stats = struct { + writes: u64 = 0, + reads: u64 = 0, + bytes_written: u64 = 0, + bytes_read: u64 = 0, + bad: u64 = 0, + evicted: u64 = 0, +}; + +pub const Disk = struct { + gpa: std.mem.Allocator, + io: Io, + opts: Options, + dir: Io.Dir, + staging: Staging, + lanes: *Lanes, + recs: std.ArrayList(Rec) = .empty, + free_recs: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// least recently used first + oldest: u32 = none, + newest: u32 = none, + used: u64 = 0, + stats: Stats = .{}, + + pub fn open(gpa: std.mem.Allocator, io: Io, opts: Options) !*Disk { + const d = try gpa.create(Disk); + errdefer gpa.destroy(d); + var name_buf: [64]u8 = undefined; + const sub = try std.fmt.bufPrint(&name_buf, "{x}/rank{d}", .{ opts.compat[0..8], opts.rank }); + var root = try Io.Dir.cwd().createDirPathOpen(io, opts.root, .{}); + defer root.close(io); + const dir = try root.createDirPathOpen(io, sub, .{ .open_options = .{ .iterate = true } }); + errdefer dir.close(io); + d.* = .{ .gpa = gpa, .io = io, .opts = opts, .dir = dir, .staging = try Staging.init(gpa, io, @max(2, opts.staging), up(opts.chunk)), .lanes = undefined }; + errdefer d.staging.deinit(); + d.lanes = try Lanes.start(gpa, io, &d.staging, @max(1, opts.lanes)); + return d; + } + + pub fn close(d: *Disk) void { + d.lanes.shutdown(); + d.staging.deinit(); + d.dir.close(d.io); + d.recs.deinit(d.gpa); + d.free_recs.deinit(d.gpa); + d.by_key.deinit(d.gpa); + const gpa = d.gpa; + gpa.destroy(d); + } + + pub fn accepts(d: *const Disk, pos: u64) bool { + return pos >= d.opts.min_tokens; + } + + pub fn get(d: *Disk, key: Digest) ?*Rec { + const i = d.by_key.get(key) orelse return null; + return &d.recs.items[i]; + } + + fn fileName(key: Digest, ext: []const u8, buf: *[48]u8) [:0]const u8 { + const n = (std.fmt.bufPrint(buf[0..47], "{x}.{s}", .{ &key, ext }) catch unreachable).len; + buf[n] = 0; + return buf[0..n :0]; + } + + fn openFile(d: *Disk, name: [:0]const u8, write: bool) !struct { fd: linux.fd_t, direct: bool } { + var flags: linux.O = .{ .ACCMODE = if (write) .WRONLY else .RDONLY, .CLOEXEC = true }; + if (write) { + flags.CREAT = true; + flags.TRUNC = true; + } + if (d.opts.direct) { + var f2 = flags; + f2.DIRECT = true; + const rc = linux.openat(d.dir.handle, name, f2, 0o644); + if (linux.errno(rc) == .SUCCESS) return .{ .fd = @intCast(rc), .direct = true }; + } + const rc = linux.openat(d.dir.handle, name, flags, 0o644); + if (linux.errno(rc) != .SUCCESS) return if (linux.errno(rc) == .NOENT) error.FileNotFound else error.OpenFailed; + return .{ .fd = @intCast(rc), .direct = false }; + } + + /// The layout a WriteSrc gets: segments, chunk sizes (whole pages a chunk for pool segments). + pub fn plan(d: *Disk, src: *const WriteSrc) !Layout { + var segs: std.ArrayList(Seg) = .empty; + defer segs.deinit(d.gpa); + const cap = d.opts.chunk; + try segs.append(d.gpa, .{ .name = format.segName("ids"), .kind = .ids, .family = 0, .nbytes = src.ids.len * 4, .chunk = cap }); + for (src.families, 0..) |f, i| { + const pb: u32 = @intCast(f.pageBytes(src.page)); + const k = @max(1, cap / pb); + if (@as(u64, k) * pb > d.staging.size) return error.PageTooLarge; + try segs.append(d.gpa, .{ .name = format.segName(f.name), .kind = .pool, .family = @intCast(i), .nbytes = @as(u64, src.pages[i].len) * pb, .chunk = k * pb }); + } + for (src.blobs) |b| try segs.append(d.gpa, .{ .name = format.segName(b.name), .kind = .blob, .family = 0, .nbytes = b.bytes.len, .chunk = cap }); + const head: Head = .{ .key = src.key, .base = src.base orelse @splat(0), .compat = d.opts.compat, .tag = src.tag, .kind = src.kind, .page = src.page, .world = d.opts.world, .rank = d.opts.rank, .first_page = src.first_page, .npages = src.npages, .nseg = 0, .ndigest = 0, .pos = src.pos }; + return Layout.build(d.gpa, head, segs.items); + } + + /// Writes an entry's file streamed (the pump copies pages to staging, lanes hash and write); returns its bytes. + pub fn writeFile(d: *Disk, src: *const WriteSrc) !u64 { + var lay = try d.plan(src); + defer lay.deinit(d.gpa); + var tb: [48]u8 = undefined; + var fb: [48]u8 = undefined; + const tmp = fileName(src.key, "tmp", &tb); + const f = try d.openFile(tmp, true); + var ok = false; + defer { + _ = linux.close(f.fd); + if (!ok) _ = linux.unlinkat(d.dir.handle, tmp, 0); + } + try src.store.waitFence(src.fence); + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + for (lay.segs) |*s| { + for (0..s.nchunks) |ci| { + if (op.failed()) |e| return e; + const c = s.chunkAt(@intCast(ci)); + const sb = try d.staging.get(); + errdefer d.staging.put(sb); + const buf = d.staging.buf(sb)[0..c.data]; + const at = @as(u64, ci) * s.chunk; + switch (s.kind) { + .ids => @memcpy(buf, std.mem.sliceAsBytes(src.ids)[at..][0..c.data]), + .blob => @memcpy(buf, blobOf(src, s).bytes[at..][0..c.data]), + .pool => { + const pb: u32 = @intCast(src.families[s.family].pageBytes(src.page)); + const p0: usize = @intCast(at / pb); + try src.store.read(s.family, src.pages[s.family][p0..][0 .. c.data / pb], buf); + }, + _ => unreachable, + } + d.lanes.submit(.{ .kind = .write, .fd = f.fd, .staging = sb, .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + ci], .op = &op }); + } + } + try op.wait(); + const hb = try d.staging.get(); + defer d.staging.put(hb); + if (lay.head_span > d.staging.size) return error.HeaderTooLarge; + lay.encode(d.staging.buf(hb)); + try lanes_mod.pwriteAll(f.fd, d.staging.buf(hb)[0..lay.head_span], 0); + if (linux.errno(linux.fdatasync(f.fd)) != .SUCCESS) return error.SyncFailed; + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + if (linux.errno(linux.renameat(d.dir.handle, tmp, d.dir.handle, fileName(src.key, "tfs", &fb))) != .SUCCESS) return error.RenameFailed; + ok = true; + return lay.head.total; + } + + fn blobOf(src: *const WriteSrc, s: *const Seg) Blob { + for (src.blobs) |b| if (std.mem.eql(u8, b.name, s.nameOf())) return b; + unreachable; + } + + /// A file's header, checked. + pub fn readHead(d: *Disk, fd: linux.fd_t) !Layout { + const hb = try d.staging.get(); + defer d.staging.put(hb); + const buf = d.staging.buf(hb); + try lanes_mod.preadAll(fd, buf[0..4096], 0); + const hn = std.mem.readInt(u32, buf[12..16], .little); + if (up(hn) > buf.len) return error.BadHeader; + if (hn > 4096) try lanes_mod.preadAll(fd, buf[4096..up(hn)], 4096); + return Layout.decode(d.gpa, buf[0..up(hn)]); + } + + /// Reads pages [lo, hi) of one restore file into the sink, each chunk checked first; blobs too when want_blobs. + pub fn readFile(d: *Disk, ext: Extent, sink: *const ReadSink, want_blobs: bool) !?Restored { + var nb: [48]u8 = undefined; + const f = try d.openFile(fileName(ext.key, "tfs", &nb), false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + if (!std.mem.eql(u8, &lay.head.key, &ext.key) or lay.head.first_page != ext.first_page or lay.head.page != sink.page) return error.BadHeader; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true }, &st)) != .SUCCESS or st.size < lay.head.total) return error.Truncated; + var out: std.ArrayList(Blob) = .empty; + errdefer { + for (out.items) |b| { + d.gpa.free(b.name); + d.gpa.free(b.bytes); + } + out.deinit(d.gpa); + } + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + // (segment, chunk) jobs; the pump keeps at most staging - 1 in flight (one buffer for the header path) + var jobs: std.ArrayList(struct { seg: u32, chunk: u32 }) = .empty; + defer jobs.deinit(d.gpa); + for (lay.segs, 0..) |*s, si| switch (s.kind) { + .pool => { + if (s.family >= sink.families.len) return error.BadHeader; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + if (s.chunk % pb != 0 or !std.mem.eql(u8, s.nameOf(), fam.name)) return error.BadHeader; + const need = ownedIn(fam, sink, ext.first_page, ext.hi) - ownedIn(fam, sink, ext.first_page, ext.lo); + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + if (s.nbytes < (@as(u64, skip) + need) * pb) return error.BadHeader; + const kpp = s.chunk / pb; + const c0 = skip / kpp; + const c1 = if (need == 0) c0 else (skip + need + kpp - 1) / kpp; + for (c0..c1) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + .blob => if (want_blobs) { + const buf = try d.gpa.alloc(u8, s.nbytes); + errdefer d.gpa.free(buf); + const name = try d.gpa.dupe(u8, s.nameOf()); + errdefer d.gpa.free(name); + try out.append(d.gpa, .{ .name = name, .bytes = buf }); + for (0..s.nchunks) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + else => {}, + }; + const window = d.staging.count - 1; + var sent: usize = 0; + var seen: usize = 0; + var bytes: u64 = 0; + while (seen < jobs.items.len) { + while (sent < jobs.items.len and sent - seen < window and op.failed() == null) : (sent += 1) { + const jb = jobs.items[sent]; + const s = &lay.segs[jb.seg]; + const c = s.chunkAt(jb.chunk); + d.lanes.submit(.{ .kind = .read, .fd = f.fd, .staging = try d.staging.get(), .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + jb.chunk], .op = &op, .tag = @intCast(sent) }); + bytes += c.span; + } + const j = op.next() orelse break; + seen += 1; + defer d.staging.put(j.staging); + if (op.failed() != null) continue; + const jb = jobs.items[j.tag]; + const s = &lay.segs[jb.seg]; + const data = d.staging.buf(j.staging)[0..j.data]; + if (s.kind == .blob) { + for (out.items) |b| if (std.mem.eql(u8, b.name, s.nameOf())) @memcpy(@constCast(b.bytes)[@as(u64, jb.chunk) * s.chunk ..][0..data.len], data); + continue; + } + try d.drain(s, jb.chunk, data, ext, sink); + } + try op.wait(); + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + d.stats.bytes_read += bytes; + if (!want_blobs) { + out.deinit(d.gpa); + return null; + } + return .{ .blobs = try out.toOwnedSlice(d.gpa) }; + } + + /// A checked pool chunk's pages that fall in [lo, hi) into the sink's pages. + fn drain(d: *Disk, s: *const Seg, chunk: u32, data: []const u8, ext: Extent, sink: *const ReadSink) !void { + _ = d; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + const kpp = s.chunk / pb; + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + const end = ownedIn(fam, sink, ext.first_page, ext.hi); + const a = @max(chunk * kpp, skip); + const b = @min(chunk * kpp + @as(u32, @intCast(data.len / pb)), end); + if (a >= b) return; + var list: [512]u32 = undefined; + var n: u32 = 0; + var j = a; + while (j < b) : (j += 1) { + list[n] = sink.pages[s.family][logicalOf(fam, sink, ext.first_page, j)]; + n += 1; + if (n == list.len or j + 1 == b) { + const first = j + 1 - n; + try sink.store.write(s.family, list[0..n], data[(first - chunk * kpp) * pb ..][0 .. n * pb]); + n = 0; + } + } + } + + /// Records a written file in the index (round thread, at the plan's point); its base gains a dependent. + pub fn register(d: *Disk, src: *const WriteSrc, size: u64) !u32 { + const id = try d.add(src.key, src.base, src.tag, src.kind, src.pos, src.first_page, src.npages, size); + d.stats.writes += 1; + d.stats.bytes_written += size; + return id; + } + + fn add(d: *Disk, key: Digest, base: ?Digest, tag: u32, kind: u32, pos: u64, first_page: u32, npages: u32, size: u64) !u32 { + var r: Rec = .{ .key = key, .tag = tag, .kind = kind, .pos = pos, .first_page = first_page, .npages = npages, .size = size }; + if (base) |b| { + const bi = d.by_key.get(b) orelse return error.NoBase; + r.base = bi; + r.depth = d.recs.items[bi].depth + 1; + } + const id: u32 = if (d.free_recs.pop()) |i| i else blk: { + try d.recs.append(d.gpa, undefined); + break :blk @intCast(d.recs.items.len - 1); + }; + try d.by_key.put(d.gpa, key, id); + if (r.base != none) d.recs.items[r.base].refs += 1; + d.recs.items[id] = r; + d.link(id); + d.used += size; + return id; + } + + /// A valid file found on restart: its header, size and verified ids (tokens [first_page x page, pos), owned). + pub const Found = struct { + head: Head, + size: u64, + mtime: i64, + ids: []i32, + + fn older(_: void, a: Found, b: Found) bool { + return a.mtime < b.mtime or (a.mtime == b.mtime and a.head.first_page < b.head.first_page); + } + }; + + /// Indexes the directory on restart: drops .tmp and bad files, checks ids, registers bases before their deltas. + pub fn scan(d: *Disk, out: *std.ArrayList(Found)) !void { + var it = d.dir.iterate(); + var nb: [48]u8 = undefined; + while (try it.next(d.io)) |ent| { + if (ent.name.len >= nb.len) continue; + @memcpy(nb[0..ent.name.len], ent.name); + nb[ent.name.len] = 0; + const name = nb[0..ent.name.len :0]; + if (std.mem.endsWith(u8, name, ".tmp")) { + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + } + if (!std.mem.endsWith(u8, name, ".tfs")) continue; + const f = d.found(name) catch { + d.stats.bad += 1; + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + }; + try out.append(d.gpa, f); + } + std.mem.sort(Found, out.items, {}, Found.older); + var i: usize = 0; + while (i < out.items.len) { + const f = out.items[i]; + const base: ?Digest = if (f.head.hasBase()) f.head.base else null; + if (base != null and d.by_key.get(base.?) == null) { + d.stats.bad += 1; + var kb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(f.head.key, "tfs", &kb), 0); + d.gpa.free(f.ids); + _ = out.orderedRemove(i); + continue; + } + _ = try d.add(f.head.key, base, f.head.tag, f.head.kind, f.head.pos, f.head.first_page, f.head.npages, f.size); + i += 1; + } + } + + fn found(d: *Disk, name: [:0]const u8) !Found { + const f = try d.openFile(name, false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + const h = lay.head; + var kb: [48]u8 = undefined; + if (!std.mem.eql(u8, fileName(h.key, "tfs", &kb), name)) return error.BadHeader; + if (!std.mem.eql(u8, &h.compat, &d.opts.compat) or h.rank != d.opts.rank or h.world != d.opts.world) return error.Foreign; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true, .MTIME = true }, &st)) != .SUCCESS or st.size < h.total) return error.Truncated; + const seg = lay.find(.ids, "ids") orelse return error.BadHeader; + if (seg.nbytes % 4 != 0 or seg.nbytes / 4 != h.pos - @as(u64, h.first_page) * h.page) return error.BadHeader; + const ids = try d.gpa.alloc(i32, seg.nbytes / 4); + errdefer d.gpa.free(ids); + const sb = try d.staging.get(); + defer d.staging.put(sb); + for (0..seg.nchunks) |ci| { + const c = seg.chunkAt(@intCast(ci)); + const buf = d.staging.buf(sb); + try lanes_mod.preadAll(f.fd, buf[0..c.span], c.off); + var dg: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(buf[0..c.data], &dg, .{}); + if (!std.mem.eql(u8, &dg, &lay.digests[seg.first_digest + ci])) return error.ChecksumMismatch; + @memcpy(std.mem.sliceAsBytes(ids)[@as(usize, ci) * seg.chunk ..][0..c.data], buf[0..c.data]); + } + return .{ .head = h, .size = st.size, .mtime = st.mtime.sec * std.time.ns_per_s + st.mtime.nsec, .ids = ids }; + } + + fn link(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + r.older = d.newest; + r.newer = none; + if (d.newest != none) d.recs.items[d.newest].newer = id else d.oldest = id; + d.newest = id; + } + + fn unlink(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + if (r.older != none) d.recs.items[r.older].newer = r.newer else d.oldest = r.newer; + if (r.newer != none) d.recs.items[r.newer].older = r.older else d.newest = r.older; + } + + /// Marks a file (and its bases) used now: O(chain). + pub fn touch(d: *Disk, key: Digest) void { + var id = d.by_key.get(key) orelse return; + while (id != none) : (id = d.recs.items[id].base) { + d.unlink(id); + d.link(id); + } + // bases end up newest: walk order keeps the entry itself older than its bases, so trims reach dependents first + } + + /// The files a restore of key reads, base first, with the logical pages each gives. + pub fn extents(d: *Disk, key: Digest, out: *std.ArrayList(Extent)) !void { + out.clearRetainingCapacity(); + var id = d.by_key.get(key) orelse return error.FileNotFound; + var hi = d.recs.items[id].npages; + while (id != none) { + const r = &d.recs.items[id]; + try out.append(d.gpa, .{ .key = r.key, .first_page = r.first_page, .lo = r.first_page, .hi = hi }); + hi = r.first_page; + id = r.base; + } + std.mem.reverse(Extent, out.items); + } + + /// Deletes a file without dependents (error when one stacks on it). + pub fn remove(d: *Disk, key: Digest) !void { + const id = d.by_key.get(key) orelse return; + const r = &d.recs.items[id]; + if (r.refs > 0) return error.HasDependents; + if (r.base != none) d.recs.items[r.base].refs -= 1; + d.unlink(id); + _ = d.by_key.remove(key); + d.used -= r.size; + r.live = false; + var nb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(key, "tfs", &nb), 0); + try d.free_recs.append(d.gpa, id); + } + + /// Deletes least recently used files without dependents while over budget; their keys go to out. + pub fn trim(d: *Disk, out: *std.ArrayList(Digest)) !void { + var id = d.oldest; + while (d.used > d.opts.budget and id != none) { + const next = d.recs.items[id].newer; + if (d.recs.items[id].refs == 0) { + const key = d.recs.items[id].key; + try d.remove(key); + try out.append(d.gpa, key); + d.stats.evicted += 1; + id = d.oldest; // a removed file may free its base + continue; + } + id = next; + } + } +}; + +/// Owned (this rank's) logical pages of family f in [first, k): the file's index of logical page k. +pub fn ownedIn(f: Family, sink: *const ReadSink, first: u32, k: u32) u32 { + if (!(f.split and sink.world > 1)) return k - first; + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, k, sink.world, &by); + return by[sink.rank]; +} + +/// The logical page of a family's j-th stored page of a file that starts at first. +pub fn logicalOf(f: Family, sink: *const ReadSink, first: u32, j: u32) u32 { + if (!(f.split and sink.world > 1)) return first + j; + const w = sink.world; + const r0 = first + (sink.rank + w - first % w) % w; // the first owned logical page at or after first + return r0 + j * w; +} diff --git a/zig/src/sessions/disk_test.zig b/zig/src/sessions/disk_test.zig new file mode 100644 index 0000000000..6ac0525e45 --- /dev/null +++ b/zig/src/sessions/disk_test.zig @@ -0,0 +1,125 @@ +//! The NVMe tier on real files: whole and delta files round-trip bit for bit, split ranks, damage. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk = @import("disk.zig"); + +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Fills a page of every family with bytes that name (family, logical page, byte). +fn fill(h: *HostPool, p: *const Pool, logical: u32, pg: u32, salt: u8) void { + for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pg)) continue; + const v = h.pageOf(@intCast(fi), p.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 131 +% logical *% 7 +% fi *% 29 +% salt); + } +} + +fn familyPages(gpa: std.mem.Allocator, p: *const Pool, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const out = try gpa.alloc([]u32, fams.len); + for (fams, 0..) |f, fi| { + var l: std.ArrayList(u32) = .empty; + for (pages[first..], first..) |pg, k| { + if (f.split and p.world > 1 and owned_only and k % p.world != p.rank) continue; + try l.append(gpa, p.familyPage(f, pg)); + } + out[fi] = try l.toOwnedSlice(gpa); + } + return out; +} + +fn freePages(gpa: std.mem.Allocator, x: [][]u32) void { + for (x) |l| gpa.free(l); + gpa.free(x); +} + +fn roundTrip(world: u32, rank: u32) !void { + const gpa = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(root); + var p = try Pool.init(gpa, .{ .families = &fams, .page = 256 }, 64 * 256, world, rank); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + // small chunks so pool segments span many chunks; few staging buffers so the pump waits on the lanes + const d = try disk.Disk.open(gpa, io, .{ .root = root, .compat = @splat(5), .rank = rank, .world = world, .chunk = 3 * 74_752, .staging = 3, .lanes = 3, .min_tokens = 1 }); + defer d.close(); + const a = try p.newSlot(32 * 256); + const pos: u64 = 9 * 256 + 100; + try a.ensure(pos); + for (a.mapped(), 0..) |pg, k| fill(&h, &p, @intCast(k), pg, 1); + const ids = try gpa.alloc(i32, pos); + defer gpa.free(ids); + for (ids, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + var blob: [5000]u8 = undefined; + for (&blob, 0..) |*x, i| x.* = @truncate(9 - i % 5); + const fp = try familyPages(gpa, &p, a.mapped(), 0, true); + defer freePages(gpa, fp); + const src: disk.WriteSrc = .{ .key = @splat(1), .tag = 4, .page = 256, .pos = pos, .npages = a.len, .ids = ids, .store = h.store(), .families = &fams, .pages = fp, .blobs = &.{.{ .name = "rings", .bytes = &blob }} }; + const size = try d.writeFile(&src); + _ = try d.register(&src, size); + // a newer turn stacks a delta on it: pages [9, 14) only (page 9 was partial in the base) + const pos2: u64 = 13 * 256 + 7; + try a.ensure(pos2); + for (a.mapped()[9..], 9..) |pg, k| fill(&h, &p, @intCast(k), pg, 2); + const ids2 = try gpa.alloc(i32, pos2); + defer gpa.free(ids2); + for (ids2, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + const fp2 = try familyPages(gpa, &p, a.mapped(), 9, true); + defer freePages(gpa, fp2); + const src2: disk.WriteSrc = .{ .key = @splat(2), .base = @splat(1), .tag = 4, .page = 256, .pos = pos2, .first_page = 9, .npages = a.len, .ids = ids2[9 * 256 ..], .store = h.store(), .families = &fams, .pages = fp2, .blobs = &.{.{ .name = "rings", .bytes = blob[0..77] }} }; + const size2 = try d.writeFile(&src2); + _ = try d.register(&src2, size2); + try testing.expect(size2 < size); + try testing.expectError(error.HasDependents, d.remove(@splat(1))); + // restore the delta onto other pages through its chain + const b = try p.newSlot(32 * 256); + try b.ensure(pos2); + var ext: std.ArrayList(disk.Extent) = .empty; + defer ext.deinit(gpa); + try d.extents(@splat(2), &ext); + try testing.expectEqual(@as(usize, 2), ext.items.len); + try testing.expectEqual(@as(u32, 9), ext.items[0].hi); + const tp = try familyPages(gpa, &p, b.mapped(), 0, false); + defer freePages(gpa, tp); + const sink: disk.ReadSink = .{ .store = h.store(), .families = &fams, .page = 256, .world = world, .rank = rank, .pages = tp }; + try testing.expect(try d.readFile(ext.items[0], &sink, false) == null); + var got = (try d.readFile(ext.items[1], &sink, true)).?; + defer got.deinit(gpa); + try testing.expectEqualSlices(u8, blob[0..77], got.blobs[0].bytes); + for (a.mapped(), b.mapped()) |pa, pb| for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pa)) continue; + try testing.expectEqualSlices(u8, h.pageOf(@intCast(fi), p.familyPage(f, pa)), h.pageOf(@intCast(fi), p.familyPage(f, pb))); + }; + // damage one chunk: the read refuses it before or as its rows land + var nb: [48]u8 = undefined; + const key2: [16]u8 = @splat(2); + const name = try std.fmt.bufPrint(&nb, "{x}.tfs", .{&key2}); + const f = try d.dir.openFile(io, name, .{ .mode = .read_write }); + var one = [_]u8{0}; + _ = try f.readPositional(io, &.{&one}, 3 * 4096 + 100); + one[0] ^= 0xff; + try f.writePositionalAll(io, &one, 3 * 4096 + 100); + f.close(io); + try testing.expectError(error.ChecksumMismatch, d.readFile(ext.items[1], &sink, true)); +} + +test "whole and delta files round-trip bit for bit (replicated)" { + try roundTrip(1, 0); +} + +test "split ranks write and restore their own pages (rank 0 and 1 of 2)" { + try roundTrip(2, 0); + try roundTrip(2, 1); +} diff --git a/zig/src/sessions/format.zig b/zig/src/sessions/format.zig new file mode 100644 index 0000000000..2ac391d156 --- /dev/null +++ b/zig/src/sessions/format.zig @@ -0,0 +1,197 @@ +//! A session file (format 2): a header, then segments of 4 KiB-aligned chunks, each with its own SHA-256. + +const std = @import("std"); +const Sha256 = std.crypto.hash.sha2.Sha256; +const up = @import("lanes.zig").up; + +pub const magic = "TFSESS02".*; +pub const version: u32 = 2; + +pub const SegKind = enum(u32) { ids = 1, pool = 2, blob = 3, _ }; + +pub const Head = extern struct { + magic: [8]u8 = magic, + version: u32 = version, + /// bytes of Head + segments + digests (before the pad) + head_bytes: u32 = 0, + /// SHA-256 of those bytes with this field zero + head_sha: [32]u8 = @splat(0), + key: [16]u8, + /// the entry whose files hold logical pages [0, first_page); zero: none + base: [16]u8 = @splat(0), + /// what makes these bytes mean the same thing to a build (family, knobs, layout, split) + compat: [16]u8, + tag: u32, + kind: u32, + page: u32, + world: u32, + rank: u32, + first_page: u32, + npages: u32, + nseg: u32, + ndigest: u32, + _pad: u32 = 0, + pos: u64, + total: u64 = 0, + + pub fn hasBase(h: *const Head) bool { + return !std.mem.allEqual(u8, &h.base, 0); + } +}; + +pub const Seg = extern struct { + name: [40]u8, + kind: SegKind, + /// pool: the family's index in the layout + family: u32, + nbytes: u64, + offset: u64 = 0, + /// bytes a chunk (the last may be shorter) and file bytes a chunk (4 KiB multiple) + chunk: u32, + stride: u32 = 0, + first_digest: u32 = 0, + nchunks: u32 = 0, + + pub fn nameOf(s: *const Seg) []const u8 { + return std.mem.sliceTo(&s.name, 0); + } + + /// Chunk i's data bytes and file span. + pub fn chunkAt(s: *const Seg, i: u32) struct { off: u64, data: u32, span: u32 } { + const data: u32 = @intCast(@min(@as(u64, s.chunk), s.nbytes - @as(u64, i) * s.chunk)); + return .{ .off = s.offset + @as(u64, i) * s.stride, .data = data, .span = @intCast(up(data)) }; + } +}; + +pub const Digest32 = [32]u8; + +pub fn segName(name: []const u8) [40]u8 { + var n: [40]u8 = @splat(0); + @memcpy(n[0..@min(name.len, 39)], name[0..@min(name.len, 39)]); + return n; +} + +/// A file's layout: offsets, chunks and the digest slots, before or after its chunks are written. +pub const Layout = struct { + head: Head, + segs: []Seg, + digests: []Digest32, + /// file bytes of the padded header + head_span: u64, + + /// Places segments (each with name, kind, family, nbytes, chunk) after the header. + pub fn build(gpa: std.mem.Allocator, head: Head, segs_in: []const Seg) !Layout { + const segs = try gpa.dupe(Seg, segs_in); + errdefer gpa.free(segs); + var nd: u32 = 0; + for (segs) |*s| { + if (s.chunk == 0) return error.BadSegment; + s.nchunks = @intCast((s.nbytes + s.chunk - 1) / s.chunk); + s.stride = @intCast(up(s.chunk)); + s.first_digest = nd; + nd += s.nchunks; + } + var h = head; + h.nseg = @intCast(segs.len); + h.ndigest = nd; + h.head_bytes = @intCast(@sizeOf(Head) + segs.len * @sizeOf(Seg) + nd * @sizeOf(Digest32)); + const span = up(h.head_bytes); + var off = span; + for (segs) |*s| { + s.offset = off; + if (s.nchunks > 0) off += @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk); + } + h.total = off; + const d = try gpa.alloc(Digest32, nd); + @memset(d, @splat(0)); + return .{ .head = h, .segs = segs, .digests = d, .head_span = span }; + } + + pub fn deinit(l: *Layout, gpa: std.mem.Allocator) void { + gpa.free(l.segs); + gpa.free(l.digests); + l.* = undefined; + } + + /// The header's bytes with its checksum into out (len >= head_span; the pad is zeroed). + pub fn encode(l: *Layout, out: []u8) void { + @memset(out[0..l.head_span], 0); + var h = l.head; + h.head_sha = @splat(0); + const hb = std.mem.asBytes(&h); + const sb = std.mem.sliceAsBytes(l.segs); + const db = std.mem.sliceAsBytes(l.digests); + @memcpy(out[0..hb.len], hb); + @memcpy(out[hb.len..][0..sb.len], sb); + @memcpy(out[hb.len + sb.len ..][0..db.len], db); + var sha: [32]u8 = undefined; + Sha256.hash(out[0..h.head_bytes], &sha, .{}); + @memcpy(out[@offsetOf(Head, "head_sha")..][0..32], &sha); + l.head.head_sha = sha; + } + + /// The header read back from bytes (at least head_bytes of them), every field and offset checked. + pub fn decode(gpa: std.mem.Allocator, bytes: []const u8) !Layout { + if (bytes.len < @sizeOf(Head)) return error.BadHeader; + var h: Head = undefined; + @memcpy(std.mem.asBytes(&h), bytes[0..@sizeOf(Head)]); + if (!std.mem.eql(u8, &h.magic, &magic) or h.version != version) return error.BadHeader; + const want = @sizeOf(Head) + @as(u64, h.nseg) * @sizeOf(Seg) + @as(u64, h.ndigest) * @sizeOf(Digest32); + if (h.head_bytes != want or bytes.len < want) return error.BadHeader; + var check: [32]u8 = undefined; + var sha = Sha256.init(.{}); + sha.update(bytes[0..@offsetOf(Head, "head_sha")]); + sha.update(&@as([32]u8, @splat(0))); + sha.update(bytes[@offsetOf(Head, "head_sha") + 32 .. want]); + sha.final(&check); + if (!std.mem.eql(u8, &check, &h.head_sha)) return error.ChecksumMismatch; + const segs = try gpa.alloc(Seg, h.nseg); + errdefer gpa.free(segs); + @memcpy(std.mem.sliceAsBytes(segs), bytes[@sizeOf(Head)..][0 .. segs.len * @sizeOf(Seg)]); + const d = try gpa.alloc(Digest32, h.ndigest); + errdefer gpa.free(d); + @memcpy(std.mem.sliceAsBytes(d), bytes[@sizeOf(Head) + segs.len * @sizeOf(Seg) ..][0 .. d.len * 32]); + var nd: u64 = 0; + for (segs) |s| { + if (s.chunk == 0 or s.stride != up(s.chunk) or s.first_digest != nd or s.offset % 4096 != 0) return error.BadHeader; + if (s.nchunks != (s.nbytes + s.chunk - 1) / s.chunk) return error.BadHeader; + if (s.nchunks > 0 and s.offset + @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk) > h.total) return error.BadHeader; + nd += s.nchunks; + } + if (nd != h.ndigest) return error.BadHeader; + return .{ .head = h, .segs = segs, .digests = d, .head_span = up(h.head_bytes) }; + } + + pub fn find(l: *const Layout, kind: SegKind, name: []const u8) ?*const Seg { + for (l.segs) |*s| if (s.kind == kind and std.mem.eql(u8, s.nameOf(), name)) return s; + return null; + } +}; + +test "a layout encodes, decodes, and refuses a damaged header" { + const gpa = std.testing.allocator; + const head: Head = .{ .key = @splat(1), .compat = @splat(2), .tag = 3, .kind = 0, .page = 256, .world = 2, .rank = 1, .first_page = 4, .npages = 9, .nseg = 0, .ndigest = 0, .pos = 2200 }; + var l = try Layout.build(gpa, head, &.{ + .{ .name = segName("ids"), .kind = .ids, .family = 0, .nbytes = 4 * (2200 - 1024), .chunk = 1 << 20 }, + .{ .name = segName("comp.20"), .kind = .pool, .family = 0, .nbytes = 3 * 149_504, .chunk = 2 * 149_504 }, + .{ .name = segName("taps"), .kind = .blob, .family = 0, .nbytes = 0, .chunk = 4096 }, + }); + defer l.deinit(gpa); + try std.testing.expectEqual(@as(u32, 1 + 2 + 0), l.head.ndigest); + try std.testing.expectEqual(@as(u64, 4096), l.segs[0].offset); + try std.testing.expectEqual(@as(u64, 4096 + up(4 * 1176)), l.segs[1].offset); + const c1 = l.segs[1].chunkAt(1); + try std.testing.expectEqual(@as(u32, 149_504), c1.data); + try std.testing.expectEqual(l.segs[1].offset + up(2 * 149_504), c1.off); + try std.testing.expectEqual(l.segs[2].offset, l.head.total); + l.digests[1][5] = 9; + var buf: [8192]u8 = undefined; + l.encode(&buf); + var back = try Layout.decode(gpa, &buf); + defer back.deinit(gpa); + try std.testing.expectEqual(l.head.total, back.head.total); + try std.testing.expectEqual(@as(u8, 9), back.digests[1][5]); + try std.testing.expectEqualStrings("comp.20", back.find(.pool, "comp.20").?.nameOf()); + buf[200] ^= 0x40; + try std.testing.expectError(error.ChecksumMismatch, Layout.decode(gpa, &buf)); +} diff --git a/zig/src/sessions/freeset.zig b/zig/src/sessions/freeset.zig new file mode 100644 index 0000000000..b8cce323af --- /dev/null +++ b/zig/src/sessions/freeset.zig @@ -0,0 +1,114 @@ +//! A set of free page indices, lowest first in O(log64 n): a bit a page, then a bit a non-empty word a level up. + +const std = @import("std"); + +const max_levels = 6; // 64^6 pages: more than any pool + +pub const FreeSet = struct { + /// levels[0] holds a bit a page, levels[k] a bit a non-empty word of levels[k - 1] + levels: [max_levels][]u64 = @splat(&.{}), + depth: u8 = 0, + n: u32 = 0, + free: u32 = 0, + + /// All of 0 .. n - 1 free. + pub fn init(gpa: std.mem.Allocator, n: u32) !FreeSet { + var s: FreeSet = .{ .n = n }; + errdefer s.deinit(gpa); + var bits: usize = @max(n, 1); + while (true) { + const words = (bits + 63) / 64; + s.levels[s.depth] = try gpa.alloc(u64, words); + @memset(s.levels[s.depth], 0); + s.depth += 1; + if (words == 1) break; + bits = words; + } + for (0..n) |i| s.setBit(@intCast(i)); + s.free = n; + return s; + } + + pub fn deinit(s: *FreeSet, gpa: std.mem.Allocator) void { + for (s.levels[0..s.depth]) |l| gpa.free(l); + s.* = .{}; + } + + pub fn count(s: *const FreeSet) u32 { + return s.free; + } + + pub fn contains(s: *const FreeSet, i: u32) bool { + return s.levels[0][i / 64] & (@as(u64, 1) << @intCast(i % 64)) != 0; + } + + /// The lowest free index, taken; null when none is free. + pub fn takeLowest(s: *FreeSet) ?u32 { + if (s.free == 0) return null; + var idx: usize = 0; + var lv: usize = s.depth; + while (lv > 0) { + lv -= 1; + const w = s.levels[lv][idx]; + idx = idx * 64 + @ctz(w); + } + const i: u32 = @intCast(idx); + s.clearBit(i); + s.free -= 1; + return i; + } + + /// Takes a given free index (error when it is not free). + pub fn take(s: *FreeSet, i: u32) error{NotFree}!void { + if (i >= s.n or !s.contains(i)) return error.NotFree; + s.clearBit(i); + s.free -= 1; + } + + /// Gives an index back (error when it is already free: a double free). + pub fn give(s: *FreeSet, i: u32) error{DoubleFree}!void { + if (i >= s.n or s.contains(i)) return error.DoubleFree; + s.setBit(i); + s.free += 1; + } + + fn setBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + const was = l[idx / 64]; + l[idx / 64] = was | (@as(u64, 1) << @intCast(idx % 64)); + if (was != 0) return; // the word was non-empty: the levels above already say so + idx /= 64; + } + } + + fn clearBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + l[idx / 64] &= ~(@as(u64, 1) << @intCast(idx % 64)); + if (l[idx / 64] != 0) return; // the word still has a free bit: the levels above stay set + idx /= 64; + } + } +}; + +test "lowest first, give back, double free refused, across levels" { + const gpa = std.testing.allocator; + var s = try FreeSet.init(gpa, 70_000); // three levels + defer s.deinit(gpa); + try std.testing.expectEqual(@as(u8, 3), s.depth); + for (0..5000) |i| try std.testing.expectEqual(@as(?u32, @intCast(i)), s.takeLowest()); + try s.give(4097); + try s.give(17); + try std.testing.expectError(error.DoubleFree, s.give(17)); + try std.testing.expectEqual(@as(?u32, 17), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 4097), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 5000), s.takeLowest()); + try s.take(69_999); + try std.testing.expectError(error.NotFree, s.take(69_999)); + try std.testing.expectEqual(@as(u32, 70_000 - 5001 - 1), s.count()); + while (s.takeLowest()) |_| {} + try std.testing.expectEqual(@as(u32, 0), s.count()); + try s.give(69_999); + try std.testing.expectEqual(@as(?u32, 69_999), s.takeLowest()); +} diff --git a/zig/src/sessions/jobs.zig b/zig/src/sessions/jobs.zig new file mode 100644 index 0000000000..5b445f20e7 --- /dev/null +++ b/zig/src/sessions/jobs.zig @@ -0,0 +1,150 @@ +//! Parks and restores off the round: a tier thread runs each job's file I/O until the plan settles it (Store.settle). + +const std = @import("std"); +const Io = std.Io; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; + +pub const Kind = enum { park, restore }; +pub const State = enum(u8) { queued, running, done, failed }; + +pub const Job = struct { + kind: Kind, + /// the store's entry + entry: u32, + state: std.atomic.Value(State) = .init(.queued), + err: ?anyerror = null, + started_ns: i96 = 0, + elapsed_ns: u64 = 0, + /// park: the file to write (owned buffers below) and its size once written + src: disk_mod.WriteSrc = undefined, + size: u64 = 0, + /// restore: files base first, the sink, the bounded state read + extents: []disk_mod.Extent = &.{}, + sink: disk_mod.ReadSink = undefined, + restored: ?disk_mod.Restored = null, + /// restore: the slot it fills + slot: u32 = 0, + /// pool pages the job holds until settled + held: []u32 = &.{}, + /// owned buffers + ids: []i32 = &.{}, + lists: [][]u32 = &.{}, + blobs: []disk_mod.Blob = &.{}, + + pub fn done(j: *const Job) bool { + const s = j.state.load(.acquire); + return s == .done or s == .failed; + } + + pub fn deinit(j: *Job, gpa: std.mem.Allocator) void { + gpa.free(j.held); + gpa.free(j.ids); + for (j.lists) |l| gpa.free(l); + gpa.free(j.lists); + for (j.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(j.blobs); + gpa.free(j.extents); + if (j.restored) |*r| r.deinit(gpa); + gpa.destroy(j); + } + + fn run(j: *Job, d: *Disk, io: Io) void { + j.state.store(.running, .release); + const t0 = Io.Clock.awake.now(io).toNanoseconds(); + j.err = switch (j.kind) { + .park => if (d.writeFile(&j.src)) |n| blk: { + j.size = n; + break :blk null; + } else |e| e, + .restore => j.restore(d), + }; + j.elapsed_ns = @intCast(Io.Clock.awake.now(io).toNanoseconds() - t0); + j.state.store(if (j.err == null) .done else .failed, .release); + } + + fn restore(j: *Job, d: *Disk) ?anyerror { + for (j.extents, 0..) |e, i| { + const last = i + 1 == j.extents.len; + const r = d.readFile(e, &j.sink, last) catch |err| return err; + if (last) j.restored = r; + } + j.sink.store.publish() catch |err| return err; + return null; + } +}; + +/// The tier thread and its queue (null thread: jobs run inline). +pub const Runner = struct { + gpa: std.mem.Allocator, + io: Io, + disk: *Disk, + thread: ?std.Thread = null, + queue: std.ArrayList(*Job) = .empty, + head: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + + pub fn init(gpa: std.mem.Allocator, io: Io, d: *Disk) Runner { + return .{ .gpa = gpa, .io = io, .disk = d }; + } + + /// Starts the tier thread: jobs then run off the caller's thread. + pub fn startThread(r: *Runner) !void { + if (r.thread != null) return; + r.thread = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, loop, .{r}); + } + + pub fn deinit(r: *Runner) void { + if (r.thread) |t| { + r.mutex.lockUncancelable(r.io); + r.stop = true; + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + t.join(); + } + r.queue.deinit(r.gpa); + } + + pub fn submit(r: *Runner, j: *Job) !void { + if (r.thread == null) return j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (r.head == r.queue.items.len) { + r.queue.clearRetainingCapacity(); + r.head = 0; + } + try r.queue.append(r.gpa, j); + r.cond.signal(r.io); + } + + /// Waits for a job (the plan settles it next). + pub fn wait(r: *Runner, j: *Job) void { + while (!j.done()) { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (!j.done()) r.cond.waitTimeout(r.io, &r.mutex, .{ .duration = .{ .raw = .fromMilliseconds(5), .clock = .awake } }) catch {}; + } + } + + fn loop(r: *Runner) void { + while (true) { + const j = blk: { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + while (r.head == r.queue.items.len and !r.stop) r.cond.waitUncancelable(r.io, &r.mutex); + if (r.head == r.queue.items.len) return; + r.head += 1; + break :blk r.queue.items[r.head - 1]; + }; + j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + } + } +}; diff --git a/zig/src/sessions/lanes.zig b/zig/src/sessions/lanes.zig new file mode 100644 index 0000000000..2b005aca9a --- /dev/null +++ b/zig/src/sessions/lanes.zig @@ -0,0 +1,337 @@ +//! The NVMe tier's lane threads, which hash and write (or read and check) 4 KiB chunks, and their staging buffers. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const Sha256 = std.crypto.hash.sha2.Sha256; + +pub const block = 4096; + +pub fn up(n: u64) u64 { + return (n + block - 1) / block * block; +} + +/// Bounded, aligned staging buffers (allocated once; `release` gives the memory back while idle). +pub const Staging = struct { + io: Io, + size: usize, + count: u32, + bufs: [][]align(block) u8, + free_list: []u32, + free_n: u32, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + gpa: std.mem.Allocator, + + pub fn init(gpa: std.mem.Allocator, io: Io, count: u32, size: usize) !Staging { + if (count == 0 or size % block != 0) return error.BadStaging; + const bufs = try gpa.alloc([]align(block) u8, count); + errdefer gpa.free(bufs); + const fl = try gpa.alloc(u32, count); + errdefer gpa.free(fl); + for (bufs, 0..) |*b, i| { + b.* = &.{}; + fl[i] = @intCast(i); + } + return .{ .io = io, .size = size, .count = count, .bufs = bufs, .free_list = fl, .free_n = count, .gpa = gpa }; + } + + pub fn deinit(s: *Staging) void { + s.release(); + s.gpa.free(s.bufs); + s.gpa.free(s.free_list); + s.* = undefined; + } + + /// A free buffer's index, waiting for one; its memory is mapped on first use. + pub fn get(s: *Staging) !u32 { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + while (s.free_n == 0) s.cond.waitUncancelable(s.io, &s.mutex); + s.free_n -= 1; + const i = s.free_list[s.free_n]; + if (s.bufs[i].len == 0) s.bufs[i] = try s.gpa.alignedAlloc(u8, .fromByteUnits(block), s.size); + return i; + } + + pub fn put(s: *Staging, i: u32) void { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + s.free_list[s.free_n] = i; + s.free_n += 1; + s.cond.signal(s.io); + } + + pub fn buf(s: *Staging, i: u32) []align(block) u8 { + return s.bufs[i]; + } + + /// Unmaps every buffer (all must be free): the memory goes back between operations. + pub fn release(s: *Staging) void { + std.debug.assert(s.free_n == s.count); + for (s.bufs) |*b| if (b.len != 0) { + s.gpa.free(b.*); + b.* = &.{}; + }; + } + + pub fn bytes(s: *const Staging) u64 { + return @as(u64, s.count) * s.size; + } +}; + +pub const Kind = enum { write, read }; + +/// One chunk's I/O: buf[0 .. span] at file offset off (span: the chunk padded to 4 KiB; data: the chunk's own bytes). +pub const Job = struct { + kind: Kind, + fd: linux.fd_t, + staging: u32, + data: u32, + span: u32, + off: u64, + /// write: the chunk's digest goes here; read: the digest it must have + digest: *[32]u8, + op: *Op, + /// the pump's tag for a finished read + tag: u32 = 0, +}; + +/// A park's or restore's chunks in flight: how many are pending, the first error, and finished reads for the pump. +pub const Op = struct { + io: Io, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + pending: u32 = 0, + err: ?anyerror = null, + /// finished reads (their jobs), drained by the pump + done: std.ArrayList(Job) = .empty, + gpa: std.mem.Allocator, + /// write jobs hand their staging back themselves + staging: *Staging, + + pub fn init(gpa: std.mem.Allocator, io: Io, staging: *Staging) !Op { + var o: Op = .{ .io = io, .gpa = gpa, .staging = staging }; + try o.done.ensureTotalCapacity(gpa, staging.count); + return o; + } + + pub fn deinit(o: *Op) void { + o.done.deinit(o.gpa); + } + + fn finish(o: *Op, j: Job, err: ?anyerror) void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + o.pending -= 1; + if (err) |e| { + if (o.err == null) o.err = e; + } + if (j.kind == .read) o.done.appendAssumeCapacity(j) else o.staging.put(j.staging); + o.cond.broadcast(o.io); + } + + /// The next finished read, waiting for one; null when none is pending. + pub fn next(o: *Op) ?Job { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.done.items.len == 0 and o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + return o.done.pop(); + } + + /// Waits for every pending job; returns the first error. + pub fn wait(o: *Op) !void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + if (o.err) |e| return e; + } + + pub fn failed(o: *Op) ?anyerror { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + return o.err; + } +}; + +/// The lane threads and their queue. +pub const Lanes = struct { + io: Io, + gpa: std.mem.Allocator, + staging: *Staging, + threads: []std.Thread, + /// queued jobs: a ring as long as the staging set (every queued job holds a buffer) + ring: []Job, + head: usize = 0, + len: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + /// seconds the lanes spent hashing and in I/O (summed over lanes) + hash_ns: std.atomic.Value(u64) = .init(0), + io_ns: std.atomic.Value(u64) = .init(0), + + pub fn start(gpa: std.mem.Allocator, io: Io, staging: *Staging, n: u32) !*Lanes { + const l = try gpa.create(Lanes); + errdefer gpa.destroy(l); + const ring = try gpa.alloc(Job, staging.count); + errdefer gpa.free(ring); + l.* = .{ .io = io, .gpa = gpa, .staging = staging, .threads = try gpa.alloc(std.Thread, n), .ring = ring }; + errdefer gpa.free(l.threads); + var made: usize = 0; + errdefer { + l.halt(); + for (l.threads[0..made]) |t| t.join(); + } + for (l.threads) |*t| { + t.* = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, run, .{l}); + made += 1; + } + return l; + } + + fn halt(l: *Lanes) void { + l.mutex.lockUncancelable(l.io); + l.stop = true; + l.cond.broadcast(l.io); + l.mutex.unlock(l.io); + } + + pub fn shutdown(l: *Lanes) void { + l.halt(); + for (l.threads) |t| t.join(); + l.gpa.free(l.ring); + l.gpa.free(l.threads); + l.gpa.destroy(l); + } + + /// Queues a job (each holds a staging buffer, so the ring never overflows). + pub fn submit(l: *Lanes, j: Job) void { + { + j.op.mutex.lockUncancelable(l.io); + defer j.op.mutex.unlock(l.io); + j.op.pending += 1; + } + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + std.debug.assert(l.len < l.ring.len); + l.ring[(l.head + l.len) % l.ring.len] = j; + l.len += 1; + l.cond.signal(l.io); + } + + fn run(l: *Lanes) void { + while (true) { + const j = blk: { + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + while (l.len == 0 and !l.stop) l.cond.waitUncancelable(l.io, &l.mutex); + if (l.len == 0) return; + const j = l.ring[l.head]; + l.head = (l.head + 1) % l.ring.len; + l.len -= 1; + break :blk j; + }; + j.op.finish(j, l.work(j)); + } + } + + fn work(l: *Lanes, j: Job) ?anyerror { + const b = l.staging.buf(j.staging); + var t0 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + switch (j.kind) { + .write => { + Sha256.hash(b[0..j.data], j.digest, .{}); + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.hash_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + t0 = t1; + @memset(b[j.data..j.span], 0); + pwriteAll(j.fd, b[0..j.span], j.off) catch |e| return e; + _ = l.io_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t0), .monotonic); + }, + .read => { + preadAll(j.fd, b[0..j.span], j.off) catch |e| return e; + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.io_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + var d: [32]u8 = undefined; + Sha256.hash(b[0..j.data], &d, .{}); + _ = l.hash_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t1), .monotonic); + if (!std.mem.eql(u8, &d, j.digest)) return error.ChecksumMismatch; + }, + } + return null; + } +}; + +pub fn pwriteAll(fd: linux.fd_t, data: []const u8, off: u64) !void { + var done: usize = 0; + while (done < data.len) { + const rc = linux.pwrite(fd, data[done..].ptr, data.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + .NOSPC => return error.NoSpaceLeft, + else => return error.WriteFailed, + } + if (rc == 0) return error.WriteFailed; + done += rc; + } +} + +pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { + var done: usize = 0; + while (done < out.len) { + const rc = linux.pread(fd, out[done..].ptr, out.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.ReadFailed, + } + if (rc == 0) return error.ShortRead; + done += rc; + } +} + +test "lanes write, read back and check chunks through bounded staging" { + if (@import("builtin").os.tag != .linux) return error.SkipZigTest; + const gpa = std.testing.allocator; + const io = std.testing.io; + var st = try Staging.init(gpa, io, 3, 2 * block); + defer st.deinit(); + const l = try Lanes.start(gpa, io, &st, 2); + defer l.shutdown(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const f = try tmp.dir.createFile(io, "lanes.bin", .{ .read = true }); + defer f.close(io); + const fd = f.handle; + var digests: [7][32]u8 = undefined; + var op = try Op.init(gpa, io, &st); + defer op.deinit(); + for (0..7) |i| { + const s = try st.get(); + const n: u32 = if (i == 6) 1000 else 2 * block - 7; + for (st.buf(s)[0..n], 0..) |*b, k| b.* = @truncate(k * 31 + i); + l.submit(.{ .kind = .write, .fd = fd, .staging = s, .data = n, .span = @intCast(up(n)), .off = i * 2 * block, .digest = &digests[i], .op = &op }); + } + try op.wait(); + var rop = try Op.init(gpa, io, &st); + defer rop.deinit(); + var seen: u32 = 0; + var sent: u32 = 0; + while (seen < 7) { + while (sent < 7 and sent - seen < 3) : (sent += 1) { + const n: u32 = if (sent == 6) 1000 else 2 * block - 7; + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = n, .span = @intCast(up(n)), .off = sent * 2 * block, .digest = &digests[sent], .op = &rop, .tag = sent }); + } + const j = rop.next().?; + for (st.buf(j.staging)[0..j.data], 0..) |b, k| try std.testing.expectEqual(@as(u8, @truncate(k * 31 + j.tag)), b); + st.put(j.staging); + seen += 1; + } + try rop.wait(); + digests[2][0] ^= 1; // a damaged chunk is refused + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = 2 * block - 7, .span = 2 * block, .off = 2 * 2 * block, .digest = &digests[2], .op = &rop }); + st.put(rop.next().?.staging); + try std.testing.expectError(error.ChecksumMismatch, rop.wait()); +} diff --git a/zig/src/sessions/pagestore.zig b/zig/src/sessions/pagestore.zig new file mode 100644 index 0000000000..0e4792699b --- /dev/null +++ b/zig/src/sessions/pagestore.zig @@ -0,0 +1,129 @@ +//! The bytes behind the pool's pages: a family's tensor of pages, read and written a page run at a time. + +const std = @import("std"); +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; + +pub const PageStore = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// family fam's pages, back to back into dst (pages.len x a page's bytes) + read: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void, + /// src (pages.len x a page's bytes) into family fam's pages + write: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void, + /// one page's rows of family fam into another page (copy on write of a RAM resume's partial page) + copyPage: *const fn (ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void, + /// round thread: a point after every row the compute stream has written so far + fence: *const fn (ptr: *anyopaque) anyerror!u64, + /// before reads: wait until the fence's rows are visible + waitFence: *const fn (ptr: *anyopaque, token: u64) anyerror!void, + /// after writes: make them visible to the compute stream + publish: *const fn (ptr: *anyopaque) anyerror!void, + }; + + pub fn read(s: PageStore, fam: u32, pages: []const u32, dst: []u8) !void { + return s.vtable.read(s.ptr, fam, pages, dst); + } + pub fn write(s: PageStore, fam: u32, pages: []const u32, src: []const u8) !void { + return s.vtable.write(s.ptr, fam, pages, src); + } + pub fn copyPage(s: PageStore, fam: u32, src: u32, dst: u32) !void { + return s.vtable.copyPage(s.ptr, fam, src, dst); + } + pub fn fence(s: PageStore) !u64 { + return s.vtable.fence(s.ptr); + } + pub fn waitFence(s: PageStore, token: u64) !void { + return s.vtable.waitFence(s.ptr, token); + } + pub fn publish(s: PageStore) !void { + return s.vtable.publish(s.ptr); + } +}; + +/// Every family's pages in host memory: [familyTensorPages x page bytes] a family, zeroed (the null page stays zero). +pub const HostPool = struct { + gpa: std.mem.Allocator, + page: u32, + families: []const Family, + tensors: [][]u8, + + pub fn init(gpa: std.mem.Allocator, p: *const Pool) !HostPool { + const t = try gpa.alloc([]u8, p.layout.families.len); + errdefer gpa.free(t); + var made: usize = 0; + errdefer for (t[0..made]) |x| gpa.free(x); + for (p.layout.families, 0..) |f, i| { + t[i] = try gpa.alloc(u8, @as(usize, p.familyTensorPages(f)) * f.pageBytes(p.page)); + @memset(t[i], 0); + made += 1; + } + return .{ .gpa = gpa, .page = p.page, .families = p.layout.families, .tensors = t }; + } + + pub fn deinit(h: *HostPool) void { + for (h.tensors) |x| h.gpa.free(x); + h.gpa.free(h.tensors); + h.* = undefined; + } + + pub fn store(h: *HostPool) PageStore { + return .{ .ptr = h, .vtable = &vtable }; + } + + /// One page of family fam (a view). + pub fn pageOf(h: *HostPool, fam: u32, pg: u32) []u8 { + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + return h.tensors[fam][pg * n ..][0..n]; + } + + const vtable: PageStore.VTable = .{ .read = read, .write = write, .copyPage = copyPage, .fence = fence, .waitFence = waitFence, .publish = publish }; + + fn read(ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (dst.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(dst[i * n ..][0..n], h.pageOf(fam, pg)); + } + + fn write(ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (src.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(h.pageOf(fam, pg), src[i * n ..][0..n]); + } + + fn copyPage(ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + @memcpy(h.pageOf(fam, dst), h.pageOf(fam, src)); + } + + fn fence(_: *anyopaque) anyerror!u64 { + return 0; + } + fn waitFence(_: *anyopaque, _: u64) anyerror!void {} + fn publish(_: *anyopaque) anyerror!void {} +}; + +test "host pool pages round-trip through the interface" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 1, .row_bytes = 8 }, .{ .name = "ik", .ratio = 2, .row_bytes = 4 } }; + var p = try Pool.init(gpa, .{ .families = &fams, .page = 16 }, 64, 1, 0); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + const s = h.store(); + var src: [2 * 128]u8 = undefined; + for (&src, 0..) |*b, i| b.* = @truncate(i * 7 + 3); + try s.write(0, &.{ 3, 1 }, &src); + var dst: [2 * 128]u8 = undefined; + try s.read(0, &.{ 3, 1 }, &dst); + try std.testing.expectEqualSlices(u8, &src, &dst); + try s.copyPage(0, 3, 0); + try std.testing.expectEqualSlices(u8, src[0..128], h.pageOf(0, 0)); + try std.testing.expect(std.mem.allEqual(u8, h.pageOf(0, p.null_page), 0)); + try std.testing.expectError(error.BadLength, s.read(1, &.{0}, &dst)); +} diff --git a/zig/src/sessions/pool.zig b/zig/src/sessions/pool.zig new file mode 100644 index 0000000000..85ee2c48c8 --- /dev/null +++ b/zig/src/sessions/pool.zig @@ -0,0 +1,484 @@ +//! The paged KV pool's host book: lowest-free pages per residue mod W, holder counts, slot page tables, quotas. + +const std = @import("std"); +const FreeSet = @import("freeset.zig").FreeSet; + +pub const Error = error{ PoolExhausted, PastCapacity, SlotNotEmpty, DoubleFree, BadPool, ResidueMismatch } || std.mem.Allocator.Error; + +/// A cache family: rows of `row_bytes` every `ratio` tokens; `split` families shard by logical page over the ranks. +pub const Family = struct { + name: []const u8, + ratio: u32, + row_bytes: u32, + split: bool = false, + + pub fn pageRows(f: Family, page: u32) u32 { + return page / f.ratio; + } + + /// Bytes of one page of this family (a whole page's rows). + pub fn pageBytes(f: Family, page: u32) u64 { + return @as(u64, f.pageRows(page)) * f.row_bytes; + } +}; + +/// The pool's families and page size: what a token costs and how pages are laid out. +pub const Layout = struct { + families: []const Family, + /// tokens a page (a power of two; 256 for every family we serve) + page: u32 = 256, + + /// A rank's bytes a token: the split families count 1 / W. + pub fn bytesPerToken(l: Layout, world: u32) f64 { + var b: f64 = 0; + for (l.families) |f| b += @as(f64, @floatFromInt(f.row_bytes)) / @as(f64, @floatFromInt(f.ratio)) / @as(f64, @floatFromInt(if (f.split and world > 1) world else 1)); + return b; + } + + /// A rank's average bytes a page (split families 1 / W): what the residency budget counts. + pub fn pageBytes(l: Layout, world: u32) u64 { + var b: u64 = 0; + for (l.families) |f| b += f.pageBytes(l.page) / (if (f.split and world > 1) world else 1); + return b; + } +}; + +pub fn pagesFor(tokens: u64, page: u32) u32 { + return @intCast((tokens + page - 1) / page); +} + +/// Logical pages [k0, k1) counted by residue mod `w` (O(w)). +pub fn span(k0: u32, k1: u32, w: u32, out: []u32) void { + for (out[0..w], 0..) |*o, r| { + const rr: u32 = @intCast(r); + const first = k0 + (rr + w - k0 % w) % w; + o.* = if (first >= k1) 0 else (k1 - 1 - first) / w + 1; + } +} + +pub const max_world = 8; + +/// A file whose logical pages [0, pages) match a slot's or entry's rows: a later park writes only past it. +pub const Origin = struct { + key: [16]u8, + pages: u32, +}; + +/// One slot's page table: logical page i -> physical page; the device tables follow through the dirty range. +pub const Slot = struct { + pool: *Pool, + id: u32, + capacity: u64, + max_pages: u32, + /// physical page of each mapped logical page; len is the mapped count + pages: []u32, + len: u32 = 0, + /// pages the admission reserved for this slot's request (null: no admission control) + quota: ?u32 = null, + /// pages taken past the quota, from pages nobody had reserved + over: u32 = 0, + /// table entries changed since the device side last synced them: [lo, hi) + dirty_lo: u32 = std.math.maxInt(u32), + dirty_hi: u32 = 0, + /// the file the leading pages came from (truncation shortens it, release clears it) + origin: ?Origin = null, + + pub fn mapped(s: *const Slot) []const u32 { + return s.pages[0..s.len]; + } + + pub fn tokens(s: *const Slot) u64 { + return @as(u64, s.len) * s.pool.page; + } + + fn mark(s: *Slot, lo: u32, hi: u32) void { + s.dirty_lo = @min(s.dirty_lo, lo); + s.dirty_hi = @max(s.dirty_hi, hi); + } + + /// The table range to upload since the last call (empty when nothing changed): entries past len are the null page. + pub fn takeDirty(s: *Slot) ?struct { lo: u32, hi: u32 } { + if (s.dirty_lo >= s.dirty_hi) return null; + defer { + s.dirty_lo = std.math.maxInt(u32); + s.dirty_hi = 0; + } + return .{ .lo = s.dirty_lo, .hi = s.dirty_hi }; + } + + /// The device table's value at logical page i (the page, or the null page past len). + pub fn tableAt(s: *const Slot, i: u32) u32 { + return if (i < s.len) s.pages[i] else s.pool.null_page; + } + + /// The split families' device table at logical page i: the local page this rank stores, else the discard page. + pub fn localTableAt(s: *const Slot, i: u32) u32 { + return s.pool.localPage(s.tableAt(i)); + } + + fn reservedSpan(s: *const Slot, out: []u32) void { + const q = s.quota orelse 0; + span(s.len, @max(s.len, q), s.pool.world, out); + } + + /// Maps every page holding positions < end (a no-op when they are: the common case). + pub fn ensure(s: *Slot, end: u64) Error!void { + const p = s.pool; + const need = pagesFor(end, p.page); + if (need <= s.len) return; + if (need > s.max_pages) return error.PastCapacity; + const w = p.world; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + if (s.quota) |q| if (need > q) { + // past the reservation: only pages nobody else reserved, residue by residue + var extra: [max_world]u32 = undefined; + span(@max(q, s.len), need, w, &extra); + for (0..w) |r| if (extra[r] > p.free[r].count() -| (p.reserved[r] - before[r])) return error.PoolExhausted; + s.over += need - @max(q, s.len); + }; + var want: [max_world]u32 = undefined; + span(s.len, need, w, &want); + for (0..w) |r| if (want[r] > p.free[r].count()) return error.PoolExhausted; + const lo = s.len; + for (lo..need) |k| { + const pg = p.takeAt(@intCast(k)) catch unreachable; // counted above + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[k] = pg; + } + s.len = need; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(lo, need); + } + + /// Maps shared pages as the first pages of an empty slot (a RAM resume: read only, the slot writes past them). + pub fn adopt(s: *Slot, shared: []const u32) Error!void { + if (s.len != 0) return error.SlotNotEmpty; + if (shared.len > s.max_pages) return error.PastCapacity; + const p = s.pool; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + for (shared, 0..) |pg, k| if (pg >= p.npages or pg % p.world != k % p.world) return error.ResidueMismatch; + for (shared, 0..) |pg, k| { + p.hold(pg, 1, 1); + s.pages[k] = pg; + } + s.len = @intCast(shared.len); + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(0, s.len); + } + + /// Appends one page taken for logical page len (the copy-on-write page of a RAM resume). + pub fn push(s: *Slot, pg: u32) Error!void { + if (s.len >= s.max_pages) return error.PastCapacity; + const p = s.pool; + if (p.world > 1 and pg % p.world != s.len % p.world) return error.ResidueMismatch; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[s.len] = pg; + s.len += 1; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(s.len - 1, s.len); + } + + /// Lets go of the pages past position keep (rows nothing reads again); returns how many. + pub fn truncate(s: *Slot, keep: u64) Error!u32 { + const p = s.pool; + // positions >= keep get written again: a page holding one no longer equals the origin file's + if (s.origin) |*o| o.pages = @min(o.pages, @as(u32, @intCast(keep / p.page))); + const k = pagesFor(keep, p.page); + if (k >= s.len) return 0; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + const n = s.len - k; + for (s.pages[k..s.len]) |pg| try p.release(pg, 1); + const old = s.len; + s.len = k; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(k, old); + return n; + } + + /// Sets the admission's reservation (null: none). + pub fn reserve(s: *Slot, quota: ?u32) void { + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + s.quota = quota; + s.pool.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + s.pool.addReserved(&after); + } + + /// Frees the slot's pages and its reservation. + pub fn releaseAll(s: *Slot) Error!u32 { + s.reserve(null); + s.over = 0; + s.origin = null; + return s.truncate(0); + } +}; + +/// The pool's book: free sets by residue, holders a page, the slots and their reservations. +pub const Pool = struct { + gpa: std.mem.Allocator, + layout: Layout, + npages: u32, + page: u32, + /// log2 of page tokens + shift: u5, + /// ranks the split families are sharded over (1: every rank stores every row) + world: u32, + rank: u32, + /// the page unmapped table entries point at (all zeros: a stray read reads zeros) + null_page: u32, + /// a split family's local page count (npages / W), its local null page and its discard page (other ranks' rows) + local_pages: u32, + local_null: u32, + discard: u32, + free: [max_world]FreeSet = @splat(.{}), + /// pages each residue's slots reserved and have not mapped yet + reserved: [max_world]u32 = @splat(0), + /// holders a page (slots + entries + in-flight jobs) and how many are slots + refs: []u16, + slot_refs: []u16, + /// pages held by something and by no slot (session entries alone): what parking can free + entry_only: u32 = 0, + slots: std.ArrayList(*Slot) = .empty, + + pub fn init(gpa: std.mem.Allocator, layout: Layout, tokens: u64, world: u32, rank: u32) Error!Pool { + if (layout.page == 0 or layout.page & (layout.page - 1) != 0) return error.BadPool; + if (tokens % layout.page != 0 or tokens == 0) return error.BadPool; + if (world == 0 or world > max_world or rank >= world) return error.BadPool; + const npages: u32 = @intCast(tokens / layout.page); + if (npages % world != 0) return error.BadPool; + for (layout.families) |f| if (f.ratio == 0 or layout.page % f.ratio != 0) return error.BadPool; + var p: Pool = .{ + .gpa = gpa, + .layout = layout, + .npages = npages, + .page = layout.page, + .shift = @intCast(std.math.log2_int(u32, layout.page)), + .world = world, + .rank = rank, + .null_page = npages, + .local_pages = npages / world, + .local_null = npages / world, + .discard = npages / world + 1, + .refs = &.{}, + .slot_refs = &.{}, + }; + errdefer p.deinit(); + for (0..world) |r| p.free[r] = try FreeSet.init(gpa, npages / world); + p.refs = try gpa.alloc(u16, npages); + @memset(p.refs, 0); + p.slot_refs = try gpa.alloc(u16, npages); + @memset(p.slot_refs, 0); + return p; + } + + pub fn deinit(p: *Pool) void { + for (p.slots.items) |s| { + p.gpa.free(s.pages); + p.gpa.destroy(s); + } + p.slots.deinit(p.gpa); + for (0..p.world) |r| p.free[r].deinit(p.gpa); + p.gpa.free(p.refs); + p.gpa.free(p.slot_refs); + p.* = undefined; + } + + pub fn newSlot(p: *Pool, capacity: u64) Error!*Slot { + const s = try p.gpa.create(Slot); + errdefer p.gpa.destroy(s); + const mp = pagesFor(capacity, p.page); + s.* = .{ .pool = p, .id = @intCast(p.slots.items.len), .capacity = capacity, .max_pages = mp, .pages = try p.gpa.alloc(u32, mp) }; + errdefer p.gpa.free(s.pages); + try p.slots.append(p.gpa, s); + return s; + } + + /// A free physical page for logical page k (its residue when split), held once by the caller. + pub fn takeAt(p: *Pool, k: u32) Error!u32 { + const r = k % p.world; + const local = p.free[r].takeLowest() orelse return error.PoolExhausted; + const pg = local * p.world + r; + p.hold(pg, 1, 0); // the caller's hold (an entry-like holder until a slot maps it) + return pg; + } + + /// Holders change on page pg: dref total, dslot of them slots; keeps entry_only and gives the page back at zero. + fn hold(p: *Pool, pg: u32, dref: i32, dslot: i32) void { + const r0 = p.refs[pg]; + const s0 = p.slot_refs[pg]; + const r1: u16 = @intCast(@as(i32, r0) + dref); + const s1: u16 = @intCast(@as(i32, s0) + dslot); + p.refs[pg] = r1; + p.slot_refs[pg] = s1; + const was = r0 > 0 and s0 == 0; + const now = r1 > 0 and s1 == 0; + if (was and !now) p.entry_only -= 1; + if (!was and now) p.entry_only += 1; + } + + /// One holder lets go (slot: from a slot's table); the page returns to its free set when nobody holds it. + fn release(p: *Pool, pg: u32, slot: u1) Error!void { + if (p.refs[pg] == 0) return error.DoubleFree; + p.hold(pg, -1, -@as(i32, slot)); + if (p.refs[pg] == 0) p.free[pg % p.world].give(pg / p.world) catch return error.DoubleFree; + } + + /// One more holder (a session entry or an in-flight job) of each page. + pub fn share(p: *Pool, pages: []const u32) void { + for (pages) |pg| p.hold(pg, 1, 0); + } + + /// Those holders let go. + pub fn drop(p: *Pool, pages: []const u32) Error!void { + for (pages) |pg| try p.release(pg, 0); + } + + pub fn holders(p: *const Pool, pg: u32) u16 { + return p.refs[pg]; + } + + fn unreserve(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] -= by[r]; + } + + fn addReserved(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] += by[r]; + } + + /// Free pages no admission reserved (split: W x the scarcer residue, so any logical run fits). + pub fn available(p: *const Pool) u32 { + var low: u32 = std.math.maxInt(u32); + for (0..p.world) |r| low = @min(low, p.free[r].count() -| p.reserved[r]); + return low * p.world; + } + + pub fn freePages(p: *const Pool) u32 { + var n: u32 = 0; + for (0..p.world) |r| n += p.free[r].count(); + return n; + } + + /// Pages an admission needs for a request (prompt + reply + slack, at most the slot's capacity). + pub fn needPages(p: *const Pool, prompt: u64, max_tokens: u64, capacity: u64, slack: u64) u32 { + const total = prompt + max_tokens + slack; + return pagesFor(if (capacity == 0) total else @min(total, capacity), p.page); + } + + /// This rank stores the split families' rows of page pg. + pub fn owns(p: *const Pool, pg: u32) bool { + return pg != p.null_page and pg % p.world == p.rank; + } + + /// A split family's local page of pg: its own when owned, the local null for the null page, else the discard page. + pub fn localPage(p: *const Pool, pg: u32) u32 { + if (pg == p.null_page) return p.local_null; + return if (pg % p.world == p.rank) pg / p.world else p.discard; + } + + /// The page a family's tensor indexes for physical page pg (split families: localPage). + pub fn familyPage(p: *const Pool, f: Family, pg: u32) u32 { + return if (f.split and p.world > 1) p.localPage(pg) else pg; + } + + /// Pages a family's device tensor holds: npages + null, or (split) local pages + local null + discard page. + pub fn familyTensorPages(p: *const Pool, f: Family) u32 { + return if (f.split and p.world > 1) p.local_pages + 2 else p.npages + 1; + } + + /// Bytes this rank stores of family f for logical pages [first, first + n) of a slot or entry. + pub fn familyBytes(p: *const Pool, f: Family, first: u32, n: u32) u64 { + if (!(f.split and p.world > 1)) return @as(u64, n) * f.pageBytes(p.page); + var by: [max_world]u32 = undefined; + span(first, first + n, p.world, &by); + return @as(u64, by[p.rank]) * f.pageBytes(p.page); + } + + /// Bytes of every family's device tensor on this rank. + pub fn tensorBytes(p: *const Pool) u64 { + var b: u64 = 0; + for (p.layout.families) |f| b += @as(u64, p.familyTensorPages(f)) * f.pageBytes(p.page); + return b; + } + + /// Pages held by session entries alone and their bytes on this rank (the residency budget's measure): O(1). + pub fn entryOnlyBytes(p: *const Pool) u64 { + return @as(u64, p.entry_only) * p.layout.pageBytes(p.world); + } +}; + +test "residues, holders, quotas and the device tables" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 2, .row_bytes = 584, .split = true }, .{ .name = "index_k", .ratio = 2, .row_bytes = 132 } }; + var p = try Pool.init(gpa, .{ .families = &fams }, 16 * 256, 2, 1); + defer p.deinit(); + try std.testing.expectEqual(@as(u32, 8 + 2), p.familyTensorPages(fams[0])); + try std.testing.expectEqual(@as(u32, 16 + 1), p.familyTensorPages(fams[1])); + try std.testing.expectApproxEqAbs(@as(f64, 146 + 66), p.layout.bytesPerToken(2), 1e-9); + const a = try p.newSlot(8 * 256); + const b = try p.newSlot(8 * 256); + try a.ensure(3 * 256 + 1); // 4 pages: residues 0,1,0,1 + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 2, 3 }, a.mapped()); + try b.ensure(256); + try std.testing.expectEqualSlices(u32, &.{4}, b.mapped()); + try std.testing.expectEqual(p.discard, a.localTableAt(0)); // page 0 is rank 0's + try std.testing.expectEqual(@as(u32, 0), a.localTableAt(1)); // page 1 -> rank 1's local page 0 + try std.testing.expectEqual(p.local_null, a.localTableAt(7)); + const d = a.takeDirty().?; + try std.testing.expectEqual(@as(u32, 0), d.lo); + try std.testing.expectEqual(@as(u32, 4), d.hi); + try std.testing.expect(a.takeDirty() == null); + // an entry shares a's pages; a lets go; the pages stay held by the entry alone + p.share(a.mapped()[0..3]); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + _ = try a.releaseAll(); + try std.testing.expectEqual(@as(u32, 3), p.entry_only); + try std.testing.expectEqual(@as(u32, 16 - 1 - 3), p.freePages()); + try p.drop(&.{ 0, 1, 2 }); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + try std.testing.expectError(error.DoubleFree, p.drop(&.{0})); + // quotas: b reserves 6 pages; available counts W x the scarcer residue + b.reserve(6); + var by: [2]u32 = undefined; + span(1, 6, 2, &by); + try std.testing.expectEqual(@as(u32, 2), by[0]); + try std.testing.expectEqual(@as(u32, 3), by[1]); + try std.testing.expectEqual(@as(u32, 2 * @min(7 - 2, 8 - 3)), p.available()); + try b.ensure(6 * 256); + try std.testing.expectEqual(@as(u32, 0), p.reserved[0] + p.reserved[1]); + for (b.mapped(), 0..) |pg, k| try std.testing.expectEqual(k % 2, pg % 2); + // RAM resume: adopt full pages of the right residues, refuse the wrong ones + const c = try p.newSlot(8 * 256); + try std.testing.expectError(error.ResidueMismatch, c.adopt(&.{b.pages[1]})); + p.share(b.mapped()[0..2]); + try c.adopt(b.mapped()[0..2]); + try c.push(try p.takeAt(2)); + try std.testing.expectEqual(@as(u16, 3), p.holders(b.pages[0])); + try std.testing.expectEqual(@as(u32, 0), c.pages[2] % 2); +} + +test "admission without a capacity still reserves prompt, reply and slack pages" { + var p = try Pool.init(std.testing.allocator, .{ .families = &.{}, .page = 64 }, 1024, 1, 0); + defer p.deinit(); + try std.testing.expectEqual(@as(u32, 3), p.needPages(65, 64, 0, 1)); + try std.testing.expectEqual(@as(u32, 1), p.needPages(65, 64, 64, 1)); + try std.testing.expectEqual(@as(u32, 0), p.needPages(0, 0, 0, 0)); +} diff --git a/zig/src/sessions/portable_test.zig b/zig/src/sessions/portable_test.zig new file mode 100644 index 0000000000..2e70055c68 --- /dev/null +++ b/zig/src/sessions/portable_test.zig @@ -0,0 +1,9 @@ +//! Portable session bookkeeping tests; the Linux NVMe tier is tested separately. +test { + _ = @import("freeset.zig"); + _ = @import("pool.zig"); + _ = @import("pagestore.zig"); + _ = @import("prefix.zig"); + _ = @import("prefix_fail_test.zig"); + _ = @import("format.zig"); +} diff --git a/zig/src/sessions/prefix.zig b/zig/src/sessions/prefix.zig new file mode 100644 index 0000000000..80c5838c60 --- /dev/null +++ b/zig/src/sessions/prefix.zig @@ -0,0 +1,512 @@ +//! Session entries by prefix: a trie over chained page digests, one node per distinct run of full pages a tag. + +const std = @import("std"); + +pub const Digest = [16]u8; +pub const none = std.math.maxInt(u32); +const Blake = std.crypto.hash.blake2.Blake2b128; + +/// Chained digests of the full pages of ids under a tag: out[j] determines tag and ids[0 .. page (j + 1)]. +pub fn chain(tag: u32, ids: []const i32, page: u32, out: []Digest) void { + var prev = root(tag); + for (out, 0..) |*o, j| { + o.* = step(prev, ids[j * page ..][0..page]); + prev = o.*; + } +} + +/// The digest of a tag's empty prefix. +pub fn root(tag: u32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-root"); + h.update(std.mem.asBytes(&tag)); + return h.finalResult(); +} + +/// The next page's digest. +pub fn step(prev: Digest, page_ids: []const i32) Digest { + var h = Blake.init(.{}); + h.update(&prev); + h.update(std.mem.sliceAsBytes(page_ids)); + return h.finalResult(); +} + +/// An entry's identity: its node's digest (tag and full pages), its length and tail. +pub fn entryKey(node: Digest, pos: u64, tail: []const i32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-entry"); + h.update(&node); + h.update(std.mem.asBytes(&pos)); + h.update(std.mem.sliceAsBytes(tail)); + return h.finalResult(); +} + +pub const Tier = enum { ram, any }; + +pub fn Index(comptime Extra: type) type { + return struct { + const Self = @This(); + + pub const Node = struct { + digest: Digest, + parent: u32, + depth: u32, + first_child: u32 = none, + next_sib: u32 = none, + prev_sib: u32 = none, + first_entry: u32 = none, + /// entries in the subtree (this node's included): all, and those in RAM + sub_all: u32 = 0, + sub_ram: u32 = 0, + /// RAM entries at this node itself + own_ram: u32 = 0, + /// the tokens of the page leading into this node, while sub_ram > 0 + tokens: ?[]i32 = null, + }; + + pub const Entry = struct { + key: Digest, + tag: u32, + pos: u64, + node: u32, + tail: []i32, + ram: bool, + next_at: u32 = none, + prev_at: u32 = none, + live: bool = true, + extra: Extra, + }; + + /// An entry inserted, or the one already there with the same key. + pub const Placed = struct { id: u32, new: bool }; + + gpa: std.mem.Allocator, + page: u32, + nodes: std.ArrayList(Node) = .empty, + free_nodes: std.ArrayList(u32) = .empty, + by_digest: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + dead: u32 = 0, // nodes dropped from `by_digest` since it was last rehashed + entries: std.ArrayList(Entry) = .empty, + free_entries: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// page-sized token blocks, reused + blocks: std.ArrayList([]i32) = .empty, + + pub fn init(gpa: std.mem.Allocator, page: u32) Self { + return .{ .gpa = gpa, .page = page }; + } + + pub fn deinit(x: *Self) void { + for (x.nodes.items) |n| if (n.tokens) |t| x.gpa.free(t); + for (x.blocks.items) |b| x.gpa.free(b); + for (x.entries.items) |e| if (e.live) x.gpa.free(e.tail); + x.nodes.deinit(x.gpa); + x.free_nodes.deinit(x.gpa); + x.by_digest.deinit(x.gpa); + x.entries.deinit(x.gpa); + x.free_entries.deinit(x.gpa); + x.by_key.deinit(x.gpa); + x.blocks.deinit(x.gpa); + x.* = undefined; + } + + pub fn get(x: *Self, e: u32) *Entry { + return &x.entries.items[e]; + } + + pub fn node(x: *Self, n: u32) *Node { + return &x.nodes.items[n]; + } + + pub fn lookup(x: *const Self, key: Digest) ?u32 { + return x.by_key.get(key); + } + + /// Reserves first, so a failure changes nothing; free_nodes keeps room for every slot (prune never allocates). + fn newNode(x: *Self, digest: Digest, parent: u32, depth: u32) !u32 { + try x.by_digest.ensureUnusedCapacity(x.gpa, 1); + if (x.free_nodes.items.len == 0) { + try x.nodes.ensureUnusedCapacity(x.gpa, 1); + try x.free_nodes.ensureTotalCapacity(x.gpa, x.nodes.items.len + 1); + } + const id: u32 = x.free_nodes.pop() orelse blk: { + x.nodes.appendAssumeCapacity(undefined); + break :blk @intCast(x.nodes.items.len - 1); + }; + x.nodes.items[id] = .{ .digest = digest, .parent = parent, .depth = depth }; + x.by_digest.putAssumeCapacityNoClobber(digest, id); + if (parent != none) { + const p = &x.nodes.items[parent]; + x.nodes.items[id].next_sib = p.first_child; + if (p.first_child != none) x.nodes.items[p.first_child].prev_sib = id; + p.first_child = id; + } + return id; + } + + fn takeBlock(x: *Self) ![]i32 { + if (x.blocks.pop()) |b| return b; + return x.gpa.alloc(i32, x.page); + } + + /// Adds an entry for ids (chain: its full pages' digests) or returns the one with its key; an error adds none. + pub fn insert(x: *Self, tag: u32, ids: []const i32, digests: []const Digest, ram: bool, extra: Extra) !Placed { + const full: u32 = @intCast(ids.len / x.page); + std.debug.assert(digests.len >= full); + const r = root(tag); + var n = x.by_digest.get(r) orelse try x.newNode(r, none, 0); + errdefer x.prune(n); // the path's new nodes, which no entry holds yet + for (0..full) |j| { + const d = digests[j]; + n = x.by_digest.get(d) orelse try x.newNode(d, n, @intCast(j + 1)); + } + const r2 = try x.place(tag, n, ids.len, ids[full * x.page ..], extra); + if (r2.new and ram) x.setRam(r2.id, true, ids) catch |err| { + _ = x.drop(r2.id); // the errdefer above prunes its path + return err; + }; + return r2; + } + + /// Adds a disk entry continuing from node anchor (none: the tag's root); suffix: tokens [depth x page, pos). + pub fn insertFrom(x: *Self, tag: u32, anchor: u32, suffix: []const i32, extra: Extra) !Placed { + var n = if (anchor != none) anchor else blk: { + const r = root(tag); + break :blk x.by_digest.get(r) orelse try x.newNode(r, none, 0); + }; + errdefer x.prune(n); // the path's new nodes, which no entry holds yet + const d0 = x.nodes.items[n].depth; + const full: u32 = @intCast(suffix.len / x.page); + for (0..full) |j| { + const d = step(x.nodes.items[n].digest, suffix[j * x.page ..][0..x.page]); + n = x.by_digest.get(d) orelse try x.newNode(d, n, d0 + @as(u32, @intCast(j)) + 1); + } + return x.place(tag, n, @as(u64, d0) * x.page + suffix.len, suffix[full * x.page ..], extra); + } + + /// The node at depth on an entry's path (its first `depth` full pages). + pub fn nodeAt(x: *Self, id: u32, depth: u32) ?u32 { + var n = x.entries.items[id].node; + while (n != none and x.nodes.items[n].depth > depth) n = x.nodes.items[n].parent; + return if (n != none and x.nodes.items[n].depth == depth) n else null; + } + + fn place(x: *Self, tag: u32, n: u32, pos: u64, tail_ids: []const i32, extra: Extra) !Placed { + const key = entryKey(x.nodes.items[n].digest, pos, tail_ids); + if (x.by_key.get(key)) |e| return .{ .id = e, .new = false }; + try x.by_key.ensureUnusedCapacity(x.gpa, 1); + if (x.free_entries.items.len == 0) { + try x.entries.ensureUnusedCapacity(x.gpa, 1); + try x.free_entries.ensureTotalCapacity(x.gpa, x.entries.items.len + 1); + } + const tail = try x.gpa.dupe(i32, tail_ids); // the last step that can fail + const id: u32 = x.free_entries.pop() orelse blk: { + x.entries.appendAssumeCapacity(undefined); + break :blk @intCast(x.entries.items.len - 1); + }; + x.by_key.putAssumeCapacityNoClobber(key, id); + const nd = &x.nodes.items[n]; + x.entries.items[id] = .{ .key = key, .tag = tag, .pos = pos, .node = n, .tail = tail, .ram = false, .next_at = nd.first_entry, .extra = extra }; + if (nd.first_entry != none) x.entries.items[nd.first_entry].prev_at = id; + nd.first_entry = id; + x.addPath(n, 1, 0); + return .{ .id = id, .new = true }; + } + + fn addPath(x: *Self, from: u32, d_all: i32, d_ram: i32) void { + var n = from; + while (n != none) { + const nd = &x.nodes.items[n]; + nd.sub_all = @intCast(@as(i64, nd.sub_all) + d_all); + nd.sub_ram = @intCast(@as(i64, nd.sub_ram) + d_ram); + n = nd.parent; + } + } + + /// Moves an entry in or out of RAM; going in needs its ids (to keep its path's tokens). + pub fn setRam(x: *Self, id: u32, ram: bool, ids: ?[]const i32) !void { + const e = &x.entries.items[id]; + if (e.ram == ram) return; + if (ram) { + const all = ids orelse return error.IdsNeeded; + errdefer x.dropTokens(e.node); // the blocks this call took: tokens stay only where sub_ram > 0 + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.tokens != null) continue; + const b = try x.takeBlock(); + @memcpy(b, all[(nd.depth - 1) * x.page ..][0..x.page]); + nd.tokens = b; + } + } + e.ram = ram; + x.nodes.items[e.node].own_ram = if (ram) x.nodes.items[e.node].own_ram + 1 else x.nodes.items[e.node].own_ram - 1; + x.addPath(e.node, 0, if (ram) 1 else -1); + if (!ram) x.dropTokens(e.node); + } + + fn dropTokens(x: *Self, from: u32) void { + var n = from; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_ram > 0) return; + if (nd.tokens) |t| { + x.blocks.append(x.gpa, t) catch x.gpa.free(t); + nd.tokens = null; + } + } + } + + /// Removes an entry and the nodes no entry needs any more. + pub fn remove(x: *Self, id: u32) void { + x.prune(x.drop(id)); + } + + /// Unlinks an entry and frees its slot; returns its node, whose path may now hold no entry. + fn drop(x: *Self, id: u32) u32 { + x.setRam(id, false, null) catch unreachable; + const e = &x.entries.items[id]; + const n = e.node; + if (e.prev_at != none) x.entries.items[e.prev_at].next_at = e.next_at else x.nodes.items[n].first_entry = e.next_at; + if (e.next_at != none) x.entries.items[e.next_at].prev_at = e.prev_at; + _ = x.by_key.remove(e.key); + x.gpa.free(e.tail); + e.live = false; + x.free_entries.appendAssumeCapacity(id); // place keeps room for every slot + x.addPath(n, -1, 0); + return n; + } + + /// Frees the nodes from n up that no entry holds (sub_all 0). + fn prune(x: *Self, n: u32) void { + var m = n; + while (m != none and x.nodes.items[m].sub_all == 0) { + const nd = x.nodes.items[m]; + if (nd.parent != none) { + if (nd.prev_sib != none) x.nodes.items[nd.prev_sib].next_sib = nd.next_sib else x.nodes.items[nd.parent].first_child = nd.next_sib; + if (nd.next_sib != none) x.nodes.items[nd.next_sib].prev_sib = nd.prev_sib; + } + if (x.by_digest.remove(nd.digest)) x.dropped(); + x.free_nodes.appendAssumeCapacity(m); // newNode keeps room for every slot + m = nd.parent; + } + } + + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + fn dropped(x: *Self) void { + x.dead += 1; + if (@as(u64, x.dead) * 10 < x.by_digest.capacity()) return; + x.by_digest.rehash(std.hash_map.AutoContext(Digest){}); + x.dead = 0; + } + + /// The entry's tail matches the prompt there and the entry is shorter (its full pages matched by digest). + fn tailOf(e: *const Entry, prompt: []const i32) bool { + const at: usize = @intCast(e.pos - e.tail.len); + return e.pos < prompt.len and std.mem.eql(i32, e.tail, prompt[at..][0..e.tail.len]); + } + + /// The longest entry of tag whose ids strictly prefix prompt (RAM first on a tie); digests: the prompt's chain. + pub fn find(x: *Self, tag: u32, prompt: []const i32, digests: []const Digest, tier: Tier) ?u32 { + const full: u32 = @intCast(prompt.len / x.page); + const r = x.by_digest.get(root(tag)) orelse return null; + // the deepest node on the prompt's path: each prefix of an entry's path has a node: binary search + var lo: u32 = 0; + var hi: u32 = full; + while (lo < hi) { + const mid = lo + (hi - lo + 1) / 2; + if (x.by_digest.contains(digests[mid - 1])) lo = mid else hi = mid - 1; + } + var n: u32 = if (lo == 0) r else x.by_digest.get(digests[lo - 1]).?; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_all == 0 or (tier == .ram and nd.own_ram == 0)) continue; + var best: u32 = none; + var e = nd.first_entry; + while (e != none) : (e = x.entries.items[e].next_at) { + const en = &x.entries.items[e]; + if (tier == .ram and !en.ram) continue; + if (!tailOf(en, prompt)) continue; + if (best == none or en.pos > x.entries.items[best].pos or (en.pos == x.entries.items[best].pos and en.ram)) best = e; + } + if (best != none) return best; + } + return null; + } + + /// A newer entry of id's chain holds every position it does (tier .ram: in RAM), so id need not be written. + pub fn covered(x: *Self, id: u32, tier: Tier) bool { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == id or (tier == .ram and !o.ram)) continue; + if (o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) return true; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + const n = if (tier == .ram) ch.sub_ram else ch.sub_all; + if (n == 0) continue; + if (e.tail.len == 0) return true; + if (ch.tokens) |t| if (std.mem.startsWith(i32, t, e.tail)) return true; + } + return false; + } + + /// How many RAM entries extend a RAM entry (longer, same tag, its ids their prefix): branch points have more. + pub fn extensions(x: *Self, id: u32) u32 { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var n: u32 = 0; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f != id and o.ram and o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) n += 1; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + if (ch.sub_ram == 0) continue; + if (std.mem.startsWith(i32, ch.tokens.?, e.tail)) n += ch.sub_ram; + } + return n; + } + + /// A RAM entry no other RAM entry extends: a chat's newest state. + pub fn isLeaf(x: *Self, id: u32) bool { + return x.entries.items[id].ram and !x.covered(id, .ram); + } + + /// A RAM leaf's older turns (the RAM entries it extends) into out; returns the chat's recency (newest `used`). + pub fn members(x: *Self, leaf: u32, out: *std.ArrayList(u32), used: *const fn (*const Extra) u64) !u64 { + const l = &x.entries.items[leaf]; + var recency = used(&l.extra); + var below: u32 = none; // the path's node under n: its tokens are the page an entry at n must prefix + var n = l.node; + while (n != none) : ({ + below = n; + n = x.nodes.items[n].parent; + }) { + const nd = &x.nodes.items[n]; + if (nd.own_ram == 0) continue; + const page_ids: []const i32 = if (below == none) l.tail else x.nodes.items[below].tokens.?; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == leaf or !o.ram or o.pos >= l.pos) continue; + if (!std.mem.startsWith(i32, page_ids, o.tail)) continue; + try out.append(x.gpa, f); + recency = @max(recency, used(&o.extra)); + } + } + return recency; + } + + /// A RAM entry's ids, rebuilt from its path's pages and tail, into out (len = pos). + pub fn idsOf(x: *Self, id: u32, out: []i32) !void { + const e = &x.entries.items[id]; + if (out.len != e.pos) return error.BadLength; + const at = e.pos - e.tail.len; + @memcpy(out[at..], e.tail); + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + const t = nd.tokens orelse return error.NotInRam; + @memcpy(out[(nd.depth - 1) * x.page ..][0..x.page], t); + } + } + }; +} + +const T = Index(u64); + +fn usedOf(e: *const u64) u64 { + return e.*; +} + +fn seq(gpa: std.mem.Allocator, n: usize, salt: i32) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @as(i32, @intCast(i % 997)) * 3 + salt; + return v; +} + +fn add(x: *T, tag: u32, ids: []const i32, ram: bool, used: u64) !u32 { + var d: [64]Digest = undefined; + chain(tag, ids, x.page, d[0 .. ids.len / x.page]); + return (try x.insert(tag, ids, d[0 .. ids.len / x.page], ram, used)).id; +} + +test "longest strict prefix, ties, tags and coverage" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + const chat = try seq(gpa, 200, 0); + defer gpa.free(chat); + const t1 = try add(&x, 7, chat[0..40], true, 1); // a prompt's replay point + const t2 = try add(&x, 7, chat[0..72], true, 2); // a turn's end + const t3 = try add(&x, 7, chat[0..128], false, 3); // on disk only, page aligned + _ = try add(&x, 8, chat[0..100], true, 4); // another tag + var d: [16]Digest = undefined; + chain(7, chat[0..40], 16, d[0..2]); + const again = try x.insert(7, chat[0..40], d[0..2], true, 9); + try std.testing.expect(again.id == t1 and !again.new); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t3), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..150], d[0..9], .ram)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..128], d[0..8], .any)); // strict: 128 is not a prefix of 128 + chain(7, chat[0..72], 16, d[0..4]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, chat[0..72], d[0..4], .any)); + var other = try gpa.dupe(i32, chat[0..150]); + defer gpa.free(other); + other[50] = -1; // diverges inside t2's tail page: t1 still prefixes it + chain(7, other, 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, other, d[0..9], .any)); + // t3 extends t2 on disk only: its page's tokens are gone, so a tail there is undecided (the store marks covered_by) + try std.testing.expect(x.covered(t1, .ram) and !x.covered(t2, .any) and !x.covered(t2, .ram) and !x.covered(t3, .any)); + const t4 = try add(&x, 7, chat[0..64], false, 5); // page aligned: any deeper entry covers it + try std.testing.expect(x.covered(t4, .any) and x.covered(t4, .ram)); + x.remove(t4); + try std.testing.expect(x.isLeaf(t2) and !x.isLeaf(t1)); + try std.testing.expectEqual(@as(u32, 1), x.extensions(t1)); + try std.testing.expectEqual(@as(u32, 0), x.extensions(t2)); + var m: std.ArrayList(u32) = .empty; + defer m.deinit(gpa); + try std.testing.expectEqual(@as(u64, 2), try x.members(t2, &m, usedOf)); + try std.testing.expectEqualSlices(u32, &.{t1}, m.items); + const ids = try gpa.alloc(i32, 72); + defer gpa.free(ids); + try x.idsOf(t2, ids); + try std.testing.expectEqualSlices(i32, chat[0..72], ids); + // going to disk drops the path's tokens once no RAM entry needs them; removal prunes the nodes + try x.setRam(t2, false, null); + try x.setRam(t1, false, null); + try std.testing.expectError(error.NotInRam, x.idsOf(t2, ids)); + x.remove(t3); + x.remove(t2); + x.remove(t1); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, null), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(u32, 1 + 6), x.by_digest.count()); // tag 8's root + its 6 full pages +} + +test "a long run of turns saved and dropped leaves the page map free slots" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + var live: std.ArrayList(u32) = .empty; + defer live.deinit(gpa); + var rng = std.Random.DefaultPrng.init(5); + for (0..3000) |i| { + // a chat's turn (its own pages), the oldest dropped once eight are kept: nodes come and go + const ids = try seq(gpa, 16 * 60, @intCast(i)); + defer gpa.free(ids); + try live.append(gpa, try add(&x, 7, ids[0 .. 16 * (1 + rng.random().uintLessThan(usize, 60))], true, i)); + if (live.items.len > 8) x.remove(live.orderedRemove(0)); + const meta = x.by_digest.metadata orelse continue; + var free: usize = 0; + for (meta[0..x.by_digest.capacity()]) |m| free += @intFromBool(m.isFree()); + try std.testing.expect(free * 10 >= x.by_digest.capacity()); + } +} diff --git a/zig/src/sessions/prefix_fail_test.zig b/zig/src/sessions/prefix_fail_test.zig new file mode 100644 index 0000000000..dd0d8ec489 --- /dev/null +++ b/zig/src/sessions/prefix_fail_test.zig @@ -0,0 +1,127 @@ +//! prefix.Index under allocation failure: an insert that fails changes nothing, and the same insert then succeeds. + +const std = @import("std"); +const prefix = @import("prefix.zig"); +const testing = std.testing; +const T = prefix.Index(u64); +const Digest = prefix.Digest; +const page = 16; + +/// What a failed insert must leave as it was; spare token blocks may stay in the reuse pool. +const Shape = struct { keys: u32, digests: u32, ram: u32, held: u32, sub_all: u32, sub_ram: u32 }; + +fn shape(x: *T, tag: u32) Shape { + var s: Shape = .{ .keys = x.by_key.count(), .digests = x.by_digest.count(), .ram = 0, .held = 0, .sub_all = 0, .sub_ram = 0 }; + var e = x.by_key.valueIterator(); + while (e.next()) |id| s.ram += @intFromBool(x.get(id.*).ram); + var n = x.by_digest.valueIterator(); + while (n.next()) |id| s.held += @intFromBool(x.node(id.*).tokens != null); + if (x.by_digest.get(prefix.root(tag))) |r| { + s.sub_all = x.node(r).sub_all; + s.sub_ram = x.node(r).sub_ram; + } + return s; +} + +/// Every slot is live in its map or on its free list, and a node keeps its page's tokens exactly while sub_ram > 0. +fn consistent(x: *T) !void { + try testing.expectEqual(x.entries.items.len - x.free_entries.items.len, x.by_key.count()); + try testing.expectEqual(x.nodes.items.len - x.free_nodes.items.len, x.by_digest.count()); + var n = x.by_digest.valueIterator(); + while (n.next()) |id| { + const nd = x.node(id.*); + if (nd.depth > 0) try testing.expectEqual(nd.sub_ram > 0, nd.tokens != null); + } +} + +fn seq(gpa: std.mem.Allocator, n: usize, salt: i32) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @as(i32, @intCast(i % 997)) * 3 + salt; + return v; +} + +fn add(x: *T, tag: u32, ids: []const i32, ram: bool, used: u64) !u32 { + var d: [64]Digest = undefined; + prefix.chain(tag, ids, page, d[0 .. ids.len / page]); + return (try x.insert(tag, ids, d[0 .. ids.len / page], ram, used)).id; +} + +/// The (k + 1)-th allocation from now fails, and every one after it until `heal`. +fn failAfter(fa: *testing.FailingAllocator, k: usize) void { + fa.fail_index = fa.alloc_index + k; +} + +fn heal(fa: *testing.FailingAllocator) void { + fa.fail_index = std.math.maxInt(usize); +} + +test "a RAM insert failing at any allocation leaves the index as it was, and the same insert then succeeds" { + const chat = try seq(testing.allocator, 160, 0); + defer testing.allocator.free(chat); + var d: [16]Digest = undefined; + prefix.chain(7, chat[0..150], page, d[0..9]); + var k: usize = 0; + while (k < 256) : (k += 1) { + var fa = testing.FailingAllocator.init(testing.allocator, .{}); + var x = T.init(fa.allocator(), page); + defer x.deinit(); + // what the insert meets: entries sharing its first pages, a disk entry, another tag, and free slots to reuse + _ = try add(&x, 7, chat[0..40], true, 1); + _ = try add(&x, 7, chat[0..72], false, 2); + _ = try add(&x, 8, chat[0..33], true, 3); + x.remove(try add(&x, 9, chat[0..48], true, 4)); + try consistent(&x); + const before = shape(&x, 7); + failAfter(&fa, k); + const got = x.insert(7, chat[0..150], d[0..9], true, 9) catch |err| { + try testing.expectEqual(error.OutOfMemory, err); + try consistent(&x); + try testing.expectEqual(before, shape(&x, 7)); + heal(&fa); + const again = try x.insert(7, chat[0..150], d[0..9], true, 9); + try testing.expect(again.new); + try consistent(&x); + var out: [150]i32 = undefined; + try x.idsOf(again.id, &out); + try testing.expectEqualSlices(i32, chat[0..150], &out); + var d10: [10]Digest = undefined; + prefix.chain(7, chat[0..160], page, &d10); + try testing.expectEqual(@as(?u32, again.id), x.find(7, chat[0..160], &d10, .ram)); + x.remove(again.id); + try consistent(&x); + try testing.expectEqual(before, shape(&x, 7)); + continue; + }; + // no allocation left to fail: every failure point before this one was injected and rolled back + try testing.expect(got.new and k > 0); + return; + } + return error.TestUnexpectedResult; +} + +test "a disk insert into an empty index failing at any allocation leaves it empty, and the same insert then succeeds" { + const chat = try seq(testing.allocator, 90, 0); + defer testing.allocator.free(chat); + var k: usize = 0; + while (k < 256) : (k += 1) { + var fa = testing.FailingAllocator.init(testing.allocator, .{}); + var x = T.init(fa.allocator(), page); + defer x.deinit(); + failAfter(&fa, k); + const got = x.insertFrom(7, prefix.none, chat, 5) catch |err| { + try testing.expectEqual(error.OutOfMemory, err); + try consistent(&x); + try testing.expectEqual(@as(u32, 0), x.by_digest.count()); // the tag's root and pages made before the failure + try testing.expectEqual(@as(u32, 0), x.by_key.count()); + heal(&fa); + const again = try x.insertFrom(7, prefix.none, chat, 5); + try testing.expect(again.new); + try testing.expectEqual(@as(u32, 1 + 5), x.by_digest.count()); + try consistent(&x); + continue; + }; + try testing.expect(got.new and k > 0); + return; + } + return error.TestUnexpectedResult; +} diff --git a/zig/src/sessions/sessions.zig b/zig/src/sessions/sessions.zig new file mode 100644 index 0000000000..d13bef6029 --- /dev/null +++ b/zig/src/sessions/sessions.zig @@ -0,0 +1,34 @@ +//! Session memory, family-neutral: the KV pool book, page bytes, prefix entries, the RAM tier and the NVMe tier. + +pub const freeset = @import("freeset.zig"); +pub const pool = @import("pool.zig"); +pub const Pool = pool.Pool; +pub const Slot = pool.Slot; +pub const Family = pool.Family; +pub const Layout = pool.Layout; +pub const pagestore = @import("pagestore.zig"); +pub const PageStore = pagestore.PageStore; +pub const HostPool = pagestore.HostPool; +pub const prefix = @import("prefix.zig"); +pub const lanes = @import("lanes.zig"); +pub const format = @import("format.zig"); +pub const disk = @import("disk.zig"); +pub const Disk = disk.Disk; +pub const jobs = @import("jobs.zig"); +pub const store = @import("store.zig"); +pub const Store = store.Store; + +test { + _ = freeset; + _ = pool; + _ = pagestore; + _ = prefix; + _ = lanes; + _ = format; + _ = disk; + _ = jobs; + _ = store; + _ = @import("prefix_fail_test.zig"); + _ = @import("disk_test.zig"); + _ = @import("store_test.zig"); +} diff --git a/zig/src/sessions/store.zig b/zig/src/sessions/store.zig new file mode 100644 index 0000000000..0bd128dd84 --- /dev/null +++ b/zig/src/sessions/store.zig @@ -0,0 +1,570 @@ +//! The session store: RAM entries holding pool pages and bounded state, a residency budget, and the NVMe tier. + +const std = @import("std"); +const Io = std.Io; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const PageStore = @import("pagestore.zig").PageStore; +const prefix = @import("prefix.zig"); +const Digest = prefix.Digest; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; +const jobs = @import("jobs.zig"); +const Job = jobs.Job; + +pub const none = prefix.none; + +/// A slot's bounded state (rings, carries, lookback, taps) as the family snapshots it: opaque to the store. +pub const Bounded = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + nbytes: *const fn (ptr: *anyopaque) u64, + /// host copies of its pieces for a park (owned by the caller: name and bytes allocated with gpa) + blobs: *const fn (ptr: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob, + release: *const fn (ptr: *anyopaque) void, + }; +}; + +pub const Policy = struct { + /// bounded state the RAM tier keeps (least recently used entries go past it) + ram_bytes: u64 = 512 << 20, + /// Idle chats and entry-only pool bytes: null is unlimited; zero retains no idle residency. + chats: ?u32 = null, + resident_bytes: ?u64 = null, + /// park an evicted chat's older turns too (default: drop them, the leaf's file covers them) + park_prefixes: bool = false, + /// drop, never write, an entry a newer entry of its chain covers + skip_covered: bool = true, + /// write a newer turn as a delta on the file its leading pages came from + delta: bool = true, + /// parks and restores on the tier thread (settled by the plan) + async_io: bool = false, +}; + +pub const Extra = struct { + /// RAM: the pool pages it holds, logical [0, npages) + pages: []u32 = &.{}, + bounded: ?Bounded = null, + /// a file with this key is registered on the NVMe tier + file: bool = false, + used: u64 = 0, + origin: ?pool_mod.Origin = null, + /// a newer entry of its chain was parked from RAM: writing this one again is redundant + covered_by: bool = false, + kind: u32 = 0, + /// an unsettled park of it + job: ?*Job = null, + older: u32 = none, + newer: u32 = none, +}; + +const Index = prefix.Index(Extra); + +pub const Hit = struct { id: u32, ram: bool, pos: u64 }; + +pub const Stats = struct { + saves: u64 = 0, + dups: u64 = 0, + hits_ram: u64 = 0, + hits_disk: u64 = 0, + parks: u64 = 0, + delta_parks: u64 = 0, + drops: u64 = 0, + covered_drops: u64 = 0, + chat_evictions: u64 = 0, + park_ns: u64 = 0, + restore_ns: u64 = 0, +}; + +pub const Store = struct { + gpa: std.mem.Allocator, + io: Io, + pool: *Pool, + pages: PageStore, + index: Index, + disk: ?*Disk, + /// heap-held: the tier thread keeps its address + runner: ?*jobs.Runner = null, + policy: Policy, + clock: u64 = 0, + bounded_total: u64 = 0, + /// RAM entries, least recently used first + oldest: u32 = none, + newest: u32 = none, + ram_count: u32 = 0, + pending: std.ArrayList(*Job) = .empty, + scratch: std.ArrayList(Digest) = .empty, + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, io: Io, pool: *Pool, pages: PageStore, d: ?*Disk, policy: Policy) !Store { + var s: Store = .{ .gpa = gpa, .io = io, .pool = pool, .pages = pages, .index = Index.init(gpa, pool.page), .disk = d, .policy = policy }; + if (d) |dd| { + const r = try gpa.create(jobs.Runner); + r.* = jobs.Runner.init(gpa, io, dd); + s.runner = r; + if (policy.async_io) r.startThread() catch |e| { + gpa.destroy(r); + return e; + }; + } + return s; + } + + pub fn deinit(s: *Store) void { + if (s.runner) |r| { + for (s.pending.items) |j| r.wait(j); + r.deinit(); + s.gpa.destroy(r); + } + for (s.pending.items) |j| j.deinit(s.gpa); + s.pending.deinit(s.gpa); + for (s.index.entries.items) |*e| if (e.live) { + if (e.extra.bounded) |b| b.vtable.release(b.ptr); + s.pool.drop(e.extra.pages) catch {}; + s.gpa.free(e.extra.pages); + }; + s.index.deinit(); + s.scratch.deinit(s.gpa); + s.* = undefined; + } + + fn digests(s: *Store, tag: u32, ids: []const i32) ![]const Digest { + const n = ids.len / s.pool.page; + try s.scratch.resize(s.gpa, n); + prefix.chain(tag, ids, s.pool.page, s.scratch.items); + return s.scratch.items; + } + + pub fn entry(s: *Store, id: u32) *Index.Entry { + return s.index.get(id); + } + + /// The longest stored entry (RAM or NVMe) whose ids strictly prefix prompt, with the same tag. + pub fn find(s: *Store, tag: u32, prompt: []const i32) !?Hit { + const d = try s.digests(tag, prompt); + const id = s.index.find(tag, prompt, d, .any) orelse return null; + const e = s.index.get(id); + return .{ .id = id, .ram = e.ram, .pos = e.pos }; + } + + fn link(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + x.older = s.newest; + x.newer = none; + if (s.newest != none) s.index.get(s.newest).extra.newer = id else s.oldest = id; + s.newest = id; + } + + fn unlink(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + if (x.older != none) s.index.get(x.older).extra.newer = x.newer else s.oldest = x.newer; + if (x.newer != none) s.index.get(x.newer).extra.older = x.older else s.newest = x.older; + } + + fn touch(s: *Store, id: u32) void { + s.clock += 1; + const e = s.index.get(id); + e.extra.used = s.clock; + if (e.ram) { + s.unlink(id); + s.link(id); + } + } + + /// An entry for a slot's state at len(ids), sharing its pages; null for a duplicate (bounded is released then). + pub fn save(s: *Store, slot: *Slot, bounded: Bounded, ids: []const i32, tag: u32, kind: u32) !?u32 { + if (ids.len == 0) return error.EmptyEntry; + const need = pool_mod.pagesFor(ids.len, s.pool.page); + if (slot.len < need) return error.SlotTooShort; + const d = try s.digests(tag, ids); + const ins = try s.index.insert(tag, ids, d, false, .{ .kind = kind }); + errdefer if (ins.new) s.index.remove(ins.id); + const e = s.index.get(ins.id); + if (e.ram) { + bounded.vtable.release(bounded.ptr); + s.touch(ins.id); + s.stats.dups += 1; + return null; + } + const pages = try s.gpa.dupe(u32, slot.mapped()[0..need]); + errdefer s.gpa.free(pages); + try s.index.setRam(ins.id, true, ids); + s.pool.share(pages); + const x = &s.index.get(ins.id).extra; + x.pages = pages; + x.bounded = bounded; + x.kind = kind; + if (slot.origin) |o| x.origin = .{ .key = o.key, .pages = @min(o.pages, @as(u32, @intCast(ids.len / s.pool.page))) }; + s.bounded_total += bounded.vtable.nbytes(bounded.ptr); + s.ram_count += 1; + s.link(ins.id); + s.touch(ins.id); + s.stats.saves += 1; + try s.trimRam(); + _ = try s.enforce(); + return ins.id; + } + + /// A RAM entry into an empty slot: full pages shared, the partial last page copied; returns its bounded state. + pub fn restoreRam(s: *Store, id: u32, slot: *Slot) !Bounded { + const e = s.index.get(id); + if (!e.ram) return error.NotInRam; + const full: u32 = @intCast(e.pos / s.pool.page); + try slot.adopt(e.extra.pages[0..full]); + errdefer _ = slot.releaseAll() catch {}; + if (e.pos % s.pool.page != 0) { + const src = e.extra.pages[full]; + const dst = try s.pool.takeAt(full); + for (s.pool.layout.families, 0..) |f, fi| { + if (f.split and !s.pool.owns(src)) continue; + try s.pages.copyPage(@intCast(fi), s.pool.familyPage(f, src), s.pool.familyPage(f, dst)); + } + try slot.push(dst); + } + if (e.extra.origin) |o| slot.origin = .{ .key = o.key, .pages = @min(o.pages, full) }; + s.touch(id); + s.stats.hits_ram += 1; + return e.extra.bounded.?; + } + + /// Starts an NVMe resume into an empty slot (pages mapped now, filled by the job); settle it before the slot runs. + pub fn beginRestore(s: *Store, id: u32, slot: *Slot) !*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.extra.file) return error.NotOnDisk; + if (slot.len != 0) return error.SlotNotEmpty; + try slot.ensure(e.pos); + const j = try s.gpa.create(Job); + j.* = .{ .kind = .restore, .entry = id, .slot = slot.id }; + errdefer j.deinit(s.gpa); + var ext: std.ArrayList(disk_mod.Extent) = .empty; + errdefer ext.deinit(s.gpa); + try d.extents(e.key, &ext); + j.extents = try ext.toOwnedSlice(s.gpa); + const np = pool_mod.pagesFor(e.pos, s.pool.page); + j.lists = try s.familyLists(slot.mapped()[0..np], 0, false); + j.sink = .{ .store = s.pages, .families = s.pool.layout.families, .page = s.pool.page, .world = s.pool.world, .rank = s.pool.rank, .pages = j.lists }; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Per family: family-tensor pages of logical [first, len(pages)) (owned_only: this rank's of split families). + fn familyLists(s: *Store, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const fams = s.pool.layout.families; + const out = try s.gpa.alloc([]u32, fams.len); + var made: usize = 0; + errdefer { + for (out[0..made]) |l| s.gpa.free(l); + s.gpa.free(out); + } + for (fams, 0..) |f, fi| { + const w = s.pool.world; + const sharded = f.split and w > 1 and owned_only; + var n: usize = pages.len - first; + if (sharded) { + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, @intCast(pages.len), w, &by); + n = by[s.pool.rank]; + } + const l = try s.gpa.alloc(u32, n); + var i: usize = 0; + var k: usize = first; + if (sharded) k = first + (s.pool.rank + w - first % w) % w; + while (k < pages.len) : (k += if (sharded) w else 1) { + l[i] = s.pool.familyPage(f, pages[k]); + i += 1; + } + out[fi] = l; + made += 1; + } + return out; + } + + /// Starts a park of a RAM entry (null: nothing to write, the entry is already settled out of RAM). + pub fn beginPark(s: *Store, id: u32) !?*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.ram or e.extra.job != null) return error.NotParkable; + if (e.extra.file and d.get(e.key) != null) { + try s.leaveRam(id); + return null; + } + const page = s.pool.page; + const npages = pool_mod.pagesFor(e.pos, page); + var base: ?Digest = null; + var first: u32 = 0; + if (s.policy.delta) if (e.extra.origin) |o| if (o.pages > 0) { + var r = d.get(o.key); + while (r) |rec| : (r = if (rec.base == disk_mod.none) null else &d.recs.items[rec.base]) { + if (rec.first_page < o.pages) break; + } + if (r) |rec| if (rec.depth < d.opts.max_chain) { + base = rec.key; + first = o.pages; + }; + }; + const j = try s.gpa.create(Job); + j.* = .{ .kind = .park, .entry = id }; + errdefer j.deinit(s.gpa); + j.ids = try s.gpa.alloc(i32, e.pos); + try s.index.idsOf(id, j.ids); + j.lists = try s.familyLists(e.extra.pages, first, true); + j.blobs = try e.extra.bounded.?.vtable.blobs(e.extra.bounded.?.ptr, s.gpa); + j.src = .{ .key = e.key, .base = base, .tag = e.tag, .kind = e.extra.kind, .page = page, .pos = e.pos, .first_page = first, .npages = npages, .ids = j.ids[first * page ..], .store = s.pages, .fence = try s.pages.fence(), .families = s.pool.layout.families, .pages = j.lists, .blobs = j.blobs }; + if (base) |b| d.get(b).?.refs += 1; // held while the job runs: a trim cannot take the base + e.extra.job = j; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Settles a finished job at the plan's point: a park registers its file, a restore returns the bounded state. + pub fn settle(s: *Store, j: *Job) !?disk_mod.Restored { + s.runner.?.wait(j); + defer { + for (s.pending.items, 0..) |p, i| if (p == j) { + _ = s.pending.orderedRemove(i); + break; + }; + j.deinit(s.gpa); + } + const d = s.disk.?; + switch (j.kind) { + .park => { + if (j.src.base) |b| d.get(b).?.refs -= 1; + const e = s.index.get(j.entry); + e.extra.job = null; + if (j.err) |err| return err; // the entry stays in RAM; the .tmp is gone + _ = try d.register(&j.src, j.size); + e.extra.file = true; + s.stats.parks += 1; + if (j.src.base != null) s.stats.delta_parks += 1; + s.stats.park_ns += j.elapsed_ns; + if (s.policy.skip_covered) { + var m: std.ArrayList(u32) = .empty; + defer m.deinit(s.gpa); + _ = try s.index.members(j.entry, &m, usedOf); + for (m.items) |o| s.index.get(o).extra.covered_by = true; + } + try s.leaveRam(j.entry); + try s.trimDisk(); + return null; + }, + .restore => { + if (j.err) |err| { + // the slot as before the restore: empty, its reservation kept + _ = try s.pool.slots.items[j.slot].truncate(0); + return err; + } + const e = s.index.get(j.entry); + d.touch(e.key); + s.touch(j.entry); + s.pool.slots.items[j.slot].origin = .{ .key = e.key, .pages = @intCast(e.pos / s.pool.page) }; + s.stats.hits_disk += 1; + s.stats.restore_ns += j.elapsed_ns; + const r = j.restored; + j.restored = null; + return r; + }, + } + } + + /// Settles every pending job in submission order (one rank, or the plan's "settle all"). + pub fn settleAll(s: *Store) !void { + while (s.pending.items.len > 0) { + if (try s.settle(s.pending.items[0])) |r| { + var rr = r; + rr.deinit(s.gpa); + } + } + } + + fn usedOf(x: *const Extra) u64 { + return x.used; + } + + /// A RAM entry lets go of its pages and bounded state; it stays indexed if a file holds it. + fn leaveRam(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (!e.ram) return; + try s.pool.drop(e.extra.pages); + s.gpa.free(e.extra.pages); + e.extra.pages = &.{}; + if (e.extra.bounded) |b| { + s.bounded_total -= b.vtable.nbytes(b.ptr); + b.vtable.release(b.ptr); + } + e.extra.bounded = null; + s.unlink(id); + s.ram_count -= 1; + try s.index.setRam(id, false, null); + if (!e.extra.file) s.index.remove(id); + } + + /// An entry whose file could not be read back: the file deleted, the entry gone unless RAM holds it (every rank). + pub fn forget(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (e.extra.job != null) return error.Busy; + if (e.extra.file) if (s.disk) |d| { + d.stats.bad += 1; + // a delta's base stays while a newer file stacks on it (an opt-in delta mode) + d.remove(e.key) catch |err| return if (err == error.HasDependents) {} else err; + e.extra.file = false; + }; + if (!e.ram) s.index.remove(id); + } + + /// Drops an entry from RAM without writing it. + pub fn drop(s: *Store, id: u32) !void { + try s.leaveRam(id); + s.stats.drops += 1; + } + + /// Lets go of a RAM entry: dropped when covered, parked if the tier takes it, else dropped; returns a started park. + pub fn evict(s: *Store, id: u32) !?*Job { + const e = s.index.get(id); + if (e.extra.job != null) return null; + if (s.policy.skip_covered and (e.extra.covered_by or s.index.covered(id, .any))) { + try s.drop(id); + s.stats.covered_drops += 1; + return null; + } + if (s.disk) |d| if (d.accepts(e.pos)) { + const j = try s.beginPark(id) orelse return null; + if (!s.policy.async_io) _ = try s.settle(j); + return if (s.policy.async_io) j else null; + }; + try s.drop(id); + return null; + } + + /// Evicts least recently used RAM entries past the bounded-state budget. + fn trimRam(s: *Store) !void { + var id = s.oldest; + while (s.bounded_total > s.policy.ram_bytes and s.ram_count > 1 and id != none) { + const next = s.index.get(id).extra.newer; + if (s.index.get(id).extra.job == null) _ = try s.evict(id); + id = next; + } + } + + /// Files the NVMe budget deleted leave the index (a RAM entry with that key stays, unwritten again). + fn trimDisk(s: *Store) !void { + const d = s.disk orelse return; + var gone: std.ArrayList(Digest) = .empty; + defer gone.deinit(s.gpa); + try d.trim(&gone); + for (gone.items) |k| if (s.index.lookup(k)) |id| { + s.index.get(id).extra.file = false; + if (!s.index.get(id).ram) s.index.remove(id); + }; + } + + const Chat = struct { leaf: u32, recency: u64 }; + + fn chatLess(_: void, a: Chat, b: Chat) bool { + return a.recency < b.recency or (a.recency == b.recency and a.leaf < b.leaf); + } + + /// A chat is live while a slot maps its leaf's first page. + fn idleLeaf(s: *Store, id: u32) bool { + const e = s.index.get(id); + return e.extra.job == null and s.index.isLeaf(id) and (e.extra.pages.len == 0 or s.pool.slot_refs[e.extra.pages[0]] == 0); + } + + /// Evicts LRU idle chats past the residency budget (leaf parked, unshared older turns dropped); returns how many. + pub fn enforce(s: *Store) !u32 { + const pol = s.policy; + if (pol.chats == null and pol.resident_bytes == null) return 0; + var chats: std.ArrayList(Chat) = .empty; + defer chats.deinit(s.gpa); + var members: std.ArrayList(u32) = .empty; + defer members.deinit(s.gpa); + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) { + if (!s.idleLeaf(id)) continue; + members.clearRetainingCapacity(); + try chats.append(s.gpa, .{ .leaf = id, .recency = try s.index.members(id, &members, usedOf) }); + } + std.mem.sort(Chat, chats.items, {}, chatLess); + var idle: u32 = @intCast(chats.items.len); + var gone: u32 = 0; + for (chats.items) |c| { + const over_n = if (pol.chats) |limit| idle > limit else false; + const over_b = if (pol.resident_bytes) |limit| s.pool.entryOnlyBytes() > limit else false; + if (!over_n and !over_b) break; + if (!s.index.get(c.leaf).ram or !s.idleLeaf(c.leaf)) continue; + members.clearRetainingCapacity(); + _ = try s.index.members(c.leaf, &members, usedOf); + _ = try s.evict(c.leaf); + // deepest first: a member more RAM entries extend than this chat's leaf is a branch point, as is all above + const own: u32 = if (s.index.get(c.leaf).ram) 1 else 0; + for (members.items) |m| { + if (s.index.get(m).extra.job != null) continue; + if (s.index.extensions(m) > own) break; + if (pol.park_prefixes) _ = try s.evict(m) else try s.drop(m); + } + idle -= 1; + gone += 1; + s.stats.chat_evictions += 1; + } + return gone; + } + + /// RAM entries to let go of (coldest first, never keep) so need pages are free: [] if they are, null if impossible. + pub fn planSpill(s: *Store, need: u32, keep: []const u32, out: *std.ArrayList(u32)) !bool { + out.clearRetainingCapacity(); + var avail = s.pool.available(); + if (avail >= need) return true; + var left: std.AutoHashMapUnmanaged(u32, u16) = .empty; + defer left.deinit(s.gpa); + var id = s.oldest; + while (id != none and avail < need) : (id = s.index.get(id).extra.newer) { + if (std.mem.indexOfScalar(u32, keep, id) != null or s.index.get(id).extra.job != null) continue; + try out.append(s.gpa, id); + for (s.index.get(id).extra.pages) |pg| { + const g = try left.getOrPut(s.gpa, pg); + if (!g.found_existing) g.value_ptr.* = s.pool.refs[pg]; + g.value_ptr.* -= 1; + if (g.value_ptr.* == 0) avail += 1; // split: a page of the scarcer residue may count for less; the check below decides + } + } + return avail >= need; + } + + /// Rebuilds NVMe entries on restart (ranks then intersect them), deltas on their bases' paths; returns how many. + pub fn reconcile(s: *Store) !u32 { + const d = s.disk orelse return 0; + var found: std.ArrayList(Disk.Found) = .empty; + defer { + for (found.items) |f| s.gpa.free(f.ids); + found.deinit(s.gpa); + } + try d.scan(&found); + for (found.items) |f| { + var anchor: u32 = none; + if (f.head.hasBase()) { + const b = s.index.lookup(f.head.base) orelse return error.NoBase; + anchor = s.index.nodeAt(b, f.head.first_page) orelse return error.BadChain; + } + const r = try s.index.insertFrom(f.head.tag, anchor, f.ids, .{ .file = true, .kind = f.head.kind }); + s.clock += 1; + s.index.get(r.id).extra.used = s.clock; + } + return @intCast(found.items.len); + } + + /// Idle resident chats and the pool bytes held by entries alone. + pub fn resident(s: *Store) struct { chats: u32, bytes: u64 } { + var n: u32 = 0; + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) if (s.idleLeaf(id)) { + n += 1; + }; + return .{ .chats = n, .bytes = s.pool.entryOnlyBytes() }; + } +}; diff --git a/zig/src/sessions/store_test.zig b/zig/src/sessions/store_test.zig new file mode 100644 index 0000000000..97d018d8a0 --- /dev/null +++ b/zig/src/sessions/store_test.zig @@ -0,0 +1,241 @@ +//! The session store on a host pool and a real NVMe tier: chats, evictions, resumes bit for bit, deltas, two ranks. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk_mod = @import("disk.zig"); +const store_mod = @import("store.zig"); +const Store = store_mod.Store; +const Bounded = store_mod.Bounded; + +const page = 64; +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Bounded state in host memory: one blob naming its chat and length. +const HostBounded = struct { + gpa: std.mem.Allocator, + bytes: []u8, + + fn make(gpa: std.mem.Allocator, chat: u32, pos: u64) !Bounded { + const h = try gpa.create(HostBounded); + h.* = .{ .gpa = gpa, .bytes = try gpa.alloc(u8, 3000) }; + for (h.bytes, 0..) |*b, i| b.* = @truncate(i +% chat *% 17 +% pos); + return .{ .ptr = h, .vtable = &vt }; + } + const vt: Bounded.VTable = .{ .nbytes = nbytes, .blobs = blobs, .release = release }; + fn nbytes(p: *anyopaque) u64 { + const h: *HostBounded = @ptrCast(@alignCast(p)); + return h.bytes.len; + } + fn blobs(p: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob { + const h: *HostBounded = @ptrCast(@alignCast(p)); + const out = try gpa.alloc(disk_mod.Blob, 1); + out[0] = .{ .name = try gpa.dupe(u8, "rings"), .bytes = try gpa.dupe(u8, h.bytes) }; + return out; + } + fn release(p: *anyopaque) void { + const h: *HostBounded = @ptrCast(@alignCast(p)); + h.gpa.free(h.bytes); + h.gpa.destroy(h); + } +}; + +/// One rank: its pool, host pages and store, plus what each chat's slot wrote (the reference rows by logical page). +const Rank = struct { + pool: Pool, + host: HostPool, + disk: *disk_mod.Disk, + store: Store, + + fn init(r: *Rank, gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, policy: store_mod.Policy) !void { + r.pool = try Pool.init(gpa, .{ .families = &fams, .page = page }, 96 * page, world, rank); + r.host = try HostPool.init(gpa, &r.pool); + r.disk = try disk_mod.Disk.open(gpa, testing.io, .{ .root = root, .compat = @splat(7), .rank = rank, .world = world, .chunk = 64 << 10, .staging = 4, .lanes = 2, .min_tokens = 1 }); + r.store = try Store.init(gpa, testing.io, &r.pool, r.host.store(), r.disk, policy); + } + + fn deinit(r: *Rank) void { + r.store.deinit(); + r.disk.close(); + r.host.deinit(); + r.pool.deinit(); + } + + /// The rows a chat's slot writes at a logical page: a function of (chat, logical page, family) alone. + fn write(r: *Rank, s: *const Slot, chat: u32, from_page: u32) void { + for (s.mapped()[from_page..], from_page..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi); + }; + } + + /// Every stored row of the slot's pages below pos equals what write() put there. + fn check(r: *Rank, s: *const Slot, chat: u32, pos: u64) !void { + const n = pool_mod.pagesFor(pos, page); + for (s.mapped()[0..n], 0..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + const rows: usize = if (k + 1 == n and pos % page != 0) (pos % page) / f.ratio * f.row_bytes else v.len; + for (v[0..rows], 0..) |b, i| try testing.expectEqual(@as(u8, @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi)), b); + }; + } +}; + +fn chatIds(gpa: std.mem.Allocator, chat: u32, n: usize) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @intCast(chat * 100_000 + i); + return v; +} + +/// Prefill or resume a chat to len(ids) in slot s, save it, free the slot, enforce residency; returns the resume point. +fn turn(gpa: std.mem.Allocator, r: *Rank, s: *Slot, chat: u32, ids: []const i32) ![]const u8 { + var tier: []const u8 = "fresh"; + var from: u32 = 0; + if (try r.store.find(1, ids)) |hit| { + if (hit.ram) { + const b = try r.store.restoreRam(hit.id, s); + try testing.expectEqual(@as(u64, 3000), b.vtable.nbytes(b.ptr)); + tier = "ram"; + } else { + const j = try r.store.beginRestore(hit.id, s); + var got = (try r.store.settle(j)).?; + defer got.deinit(gpa); + try testing.expectEqual(@as(u8, @truncate(chat *% 17 +% hit.pos)), got.blobs[0].bytes[0]); + tier = "disk"; + } + try r.check(s, chat, hit.pos); + from = @intCast(hit.pos / page); + } + try s.ensure(ids.len); + r.write(s, chat, from); + _ = try r.store.save(s, try HostBounded.make(gpa, chat, ids.len), ids, 1, 0); + _ = try s.releaseAll(); + _ = try r.store.enforce(); + try r.store.settleAll(); // the plan's settle point (async parks finish here) + return tier; +} + +/// The scenario on one rank; returns a trace of its decisions (stats and page ids) to compare ranks. +fn scenario(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, async_io: bool, trace: *std.ArrayList(u64)) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1, .async_io = async_io }); + defer r.deinit(); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + const b = try chatIds(gpa, 2, 20 * page); + defer gpa.free(b); + const st = &r.store.stats; + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 1, a[0 .. 5 * page + 10])); + try testing.expectEqualStrings("ram", try turn(gpa, &r, s, 1, a[0 .. 8 * page + 3])); + try testing.expectEqual(@as(u64, 0), st.parks); + // a second idle chat: chat 1 (least recent) parked whole, its older turn dropped (not written) + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 2, b[0 .. 6 * page])); + try testing.expectEqual(@as(u64, 1), st.parks); + try testing.expectEqual(@as(u64, 1), st.drops); + try testing.expectEqual(@as(u32, 1), r.store.resident().chats); + const whole = r.disk.used; + // chat 1 resumes from NVMe bit for bit; chat 2 goes + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 12 * page + 40])); + try testing.expectEqual(@as(u64, 2), st.parks); + const before = r.disk.used; + // chat 2 resumes; chat 1 goes again, written as a delta on its first file (pages 8..12 of 13) + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 2, b[0 .. 9 * page])); + try testing.expectEqual(@as(u64, 3), st.parks); + try testing.expectEqual(@as(u64, 1), st.delta_parks); + // the delta holds 5 of the entry's 13 pages: less than chat 1's first whole file of 9 pages + try testing.expect(r.disk.used - before < whole); + // chat 1's newest turn restores through its two files, bit for bit + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 15 * page])); + try testing.expectEqual(@as(u64, 3), st.hits_disk); + try trace.appendSlice(gpa, &.{ st.parks, st.delta_parks, st.drops, st.covered_drops, st.hits_ram, st.hits_disk, r.pool.freePages(), r.pool.entry_only }); + var id = r.store.oldest; + while (id != store_mod.none) : (id = r.store.entry(id).extra.newer) for (r.store.entry(id).extra.pages) |pg| try trace.append(gpa, pg); +} + +/// After the scenario: a new process on the same tier finds every file and resumes chat 1 through its delta chain. +fn restart(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1 }); + defer r.deinit(); + try testing.expectEqual(@as(u32, 4), try r.store.reconcile()); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 16 * page])); + try testing.expectEqual(@as(u64, 1), r.store.stats.hits_disk); +} + +fn tmpRoot(gpa: std.mem.Allocator, tmp: *std.testing.TmpDir, sub: []const u8) ![:0]u8 { + try tmp.dir.createDirPath(testing.io, sub); + return tmp.dir.realPathFileAlloc(testing.io, sub, gpa); +} + +test "chats over turns: residency to NVMe, bit-identical resumes, deltas, older turns dropped (sync and async)" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const r0 = try tmpRoot(gpa, &tmp, "sync"); + defer gpa.free(r0); + try scenario(gpa, r0, 1, 0, false, &t0); + try restart(gpa, r0, 1, 0); + const r1 = try tmpRoot(gpa, &tmp, "async"); + defer gpa.free(r1); + try scenario(gpa, r1, 1, 0, true, &t1); + try testing.expectEqualSlices(u64, t0.items, t1.items); // the tier thread changes when bytes move, not what is decided +} + +test "two split ranks make the same decisions on the same calls and each restores its own rows" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const root = try tmpRoot(gpa, &tmp, "tier"); + defer gpa.free(root); + try scenario(gpa, root, 2, 0, false, &t0); + try scenario(gpa, root, 2, 1, true, &t1); + try restart(gpa, root, 2, 0); + try restart(gpa, root, 2, 1); + try testing.expectEqualSlices(u64, t0.items, t1.items); +} + +test "zero chat or byte budgets evict idle residency while null remains unlimited" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try tmpRoot(gpa, &tmp, "zero-budget"); + defer gpa.free(root); + var r: Rank = undefined; + try r.init(gpa, root, 1, 0, .{}); + defer r.deinit(); + const slot = try r.pool.newSlot(64 * page); + const ids = try chatIds(gpa, 1, 5 * page); + defer gpa.free(ids); + _ = try turn(gpa, &r, slot, 1, ids); + try testing.expectEqual(@as(u32, 1), r.store.resident().chats); + r.store.policy.chats = 0; + try testing.expectEqual(@as(u32, 1), try r.store.enforce()); + try testing.expectEqual(@as(u32, 0), r.store.resident().chats); + r.store.policy.chats = null; + _ = try turn(gpa, &r, slot, 1, ids); + try testing.expect(r.store.resident().bytes > 0); + r.store.policy.resident_bytes = 0; + try testing.expectEqual(@as(u32, 1), try r.store.enforce()); + try testing.expectEqual(@as(u64, 0), r.store.resident().bytes); +} diff --git a/zig/tests/cuda/deepseek_v41/experts.zig b/zig/tests/cuda/deepseek_v41/experts.zig new file mode 100644 index 0000000000..112079c569 --- /dev/null +++ b/zig/tests/cuda/deepseek_v41/experts.zig @@ -0,0 +1,471 @@ +//! Bit checks of the routed EXL3 expert kernels (decode chain, x3ld / x3pf settings, x3gm) against the oracle. + +const std = @import("std"); +const cuda = @import("cuda"); +const dsv41 = @import("dsv41_kernels"); +const check = @import("../check.zig"); +const Fixture = @import("../fixture.zig").Fixture; +const experts_cases = @import("experts_cases.zig"); + +const exl3 = dsv41.exl3; +const Gpu = check.Gpu; + +pub const Kind = struct { + buf: [32]u8 = undefined, + len: usize = 0, + pub fn slice(k: *const Kind) []const u8 { + return k.buf[0..k.len]; + } +}; + +/// The case's "kind" param (which check reads it). +pub fn caseKind(gpu: Gpu, dir: []const u8) !Kind { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + const s = try fx.string("kind"); + var k: Kind = .{}; + if (s.len > k.buf.len) return error.WrongType; + @memcpy(k.buf[0..s.len], s); + k.len = s.len; + return k; +} + +/// Device copies of a fixture's arrays, freed together. +pub const Bufs = struct { + gpu: Gpu, + fx: *const Fixture, + list: std.ArrayList(cuda.DeviceBuffer) = .empty, + + pub fn deinit(b: *Bufs) void { + for (b.list.items) |*x| x.free(); + b.list.deinit(b.gpu.gpa); + } + + pub fn up(b: *Bufs, name: []const u8) !cuda.DeviceBuffer { + const host = try b.fx.bytes(name); + defer b.gpu.gpa.free(host); + const buf = try cuda.DeviceBuffer.fromHost(b.gpu.d, host); + try b.list.append(b.gpu.gpa, buf); + return buf; + } + + pub fn zeros(b: *Bufs, len: usize, byte: u8) !cuda.DeviceBuffer { + const buf = try cuda.DeviceBuffer.alloc(b.gpu.d, @max(len, 4)); + try buf.fill8(byte, null); + try b.list.append(b.gpu.gpa, buf); + return buf; + } + + /// int64 [E] pointer table: `base` + the fixture's byte offsets `name`. + fn table(b: *Bufs, base: cuda.DeviceBuffer, name: []const u8) !cuda.DeviceBuffer { + const host = try b.fx.bytes(name); + defer b.gpu.gpa.free(host); + const offs = std.mem.bytesAsSlice(u64, host); + const ptrs = try b.gpu.gpa.alloc(u64, offs.len); + defer b.gpu.gpa.free(ptrs); + for (offs, ptrs) |o, *p| p.* = base.ptr + o; + const buf = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(ptrs)); + try b.list.append(b.gpu.gpa, buf); + return buf; + } +}; + +pub fn size(fx: Fixture, name: []const u8) !usize { + return @intCast(try fx.int(name)); +} + +/// `got` (the first `want.len` bytes of the device buffer) against the fixture's array `name`. +pub fn same(gpu: Gpu, fx: Fixture, buf: cuda.DeviceBuffer, name: []const u8, what: []const u8) !void { + const want = try fx.bytes(name); + defer gpu.gpa.free(want); + const got = try gpu.gpa.alloc(u8, want.len); + defer gpu.gpa.free(got); + try buf.download(0, got); + try check.sameBytes(what, got, want); +} + +fn has(fx: Fixture, name: []const u8) bool { + const arrays = fx.parsed.value.object.get("arrays") orelse return false; + return arrays.object.get(name) != null; +} + +/// The decode chain of one routed layer at R rows, every stage against the reference output of the same stage. +pub fn decodeChain(gpu: Gpu, k: *const dsv41.Kernels, dir: []const u8) !void { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var b: Bufs = .{ .gpu = gpu, .fx = &fx }; + defer b.deinit(); + var stream = try cuda.Stream.init(gpu.d, false); // blocking: ordered after the uploads and memsets on the default stream + defer stream.deinit(); + const o = k.experts(stream); + + const R = try size(fx, "R"); + const slots = try size(fx, "slots"); + const E = try size(fx, "E"); + const D = try size(fx, "D"); + const I = try size(fx, "I"); + const P = R * slots; + const maxu = @min(P, E); + const input: exl3.Input = if (std.mem.eql(u8, try fx.string("input"), "bf16")) .bf16 else .f16; + const limit: f32 = if ((try fx.int("limit_inf")) != 0) std.math.inf(f32) else @floatCast(try fx.float("limit")); + const act_mode: c_int = @intCast(try fx.int("act_mode")); + + const x = try b.up("x"); + const pick = try b.up("pick"); + const wts = try b.up("wts"); + const words = try b.up("words"); + const tg = try b.table(words, "off_g"); + const tu = try b.table(words, "off_u"); + const td = try b.table(words, "off_d"); + const k2g = try b.up("k2_g"); + const k2u = try b.up("k2_u"); + const k2d = try b.up("k2_d"); + const suh_g = try b.up("suh_g"); + const suh_u = try b.up("suh_u"); + const svh_g = try b.up("svh_g"); + const svh_u = try b.up("svh_u"); + const suh_d = try b.up("suh_d"); + const svh_d = try b.up("svh_d"); + + const uids = try b.zeros(maxu * 4, 0); + const ucount = try b.zeros(4, 0); + const members = try b.zeros(maxu * R * 4, 0xff); + try o.group(pick.ptr, uids.ptr, ucount.ptr, members.ptr, R, slots, E, R); + try stream.synchronize(); + try same(gpu, fx, uids, "uids", "group uids"); + try same(gpu, fx, ucount, "ucount", "group ucount"); + try same(gpu, fx, members, "members", "group members"); + + const xg = try b.zeros(P * D * 2, 0); + const xu = try b.zeros(P * D * 2, 0); + try o.rotIn(input, x.ptr, D, pick.ptr, suh_g.ptr, suh_u.ptr, xg.ptr, xu.ptr, R, D, slots, E); + try stream.synchronize(); + try same(gpu, fx, xg, "xg", "rot_in gate"); + try same(gpu, fx, xu, "xu", "rot_in up"); + + // Exercise the actual Zig ABI against the same fixture, including untouched invalid-pick rows. + if (input == .bf16 and R <= 64 and D == 5120 and slots <= 9 and E <= 512) { + const fg = try b.zeros(P * D * 2, 0); + const fu = try b.zeros(P * D * 2, 0); + const fi = try b.zeros(maxu * 4, 0); + const fc = try b.zeros(4, 0); + const fm = try b.zeros(maxu * R * 4, 0xff); + try o.groupRotIn(x.ptr, D, pick.ptr, suh_g.ptr, suh_u.ptr, fg.ptr, fu.ptr, fi.ptr, fc.ptr, fm.ptr, R, D, slots, E, R); + try stream.synchronize(); + try same(gpu, fx, fi, "uids", "group_rot ids"); + try same(gpu, fx, fc, "ucount", "group_rot count"); + try same(gpu, fx, fm, "members", "group_rot members"); + try same(gpu, fx, fg, "xg", "group_rot gate"); + try same(gpu, fx, fu, "xu", "group_rot up"); + } + + const lo_gu: u32 = @intCast(try fx.int("lo_gu")); + const hi_gu: u32 = @intCast(try fx.int("hi_gu")); + const lo_d: u32 = @intCast(try fx.int("lo_d")); + const hi_d: u32 = @intCast(try fx.int("hi_d")); + const sk_gu = try size(fx, "sk_gu"); + const sk_d = try size(fx, "sk_d"); + const gu: exl3.Grouped = .{ .x0 = xg.ptr, .x1 = xu.ptr, .tp0 = tg.ptr, .tp1 = tu.ptr, .k2_0 = k2g.ptr, .k2_1 = k2u.ptr, .uids = uids.ptr, .ucount = ucount.ptr, .members = members.ptr, .z = 0, .mats = 2, .K = D, .N = I, .P = P, .SK = sk_gu, .slots = slots, .maxm = R, .nexp = maxu, .lo = lo_gu, .hi = hi_gu }; + const z = try b.zeros(@max(2 * sk_gu * P * I, sk_d * P * D) * 4, 0); + const zlen_gu = 2 * sk_gu * P * I * 4; + const zlen_d = sk_d * P * D * 4; + + // upstream's grouped kernel at the shape's default setting: the Z the epilogue reads + var g = gu; + g.z = z.ptr; + const ucfg_gu = [3]u32{ @intCast(try fx.int("nt_gu")), @intCast(try fx.int("w_gu")), @intCast(try fx.int("pf_gu")) }; + try zeroThen(z, zlen_gu); + try o.grouped(g, ucfg_gu[0], ucfg_gu[1], ucfg_gu[2]); + try stream.synchronize(); + try same(gpu, fx, z, "z_gu", "upstream grouped gate/up Z"); + var n_ld: usize = 0; + var n_pf: usize = 0; + try settings(gpu, fx, o, stream, z, g, zlen_gu, "gu", &n_ld, &n_pf); + + // the gate/up epilogue on upstream's Z (re-run: the settings above overwrote it with the same bits) + try zeroThen(z, zlen_gu); + try o.grouped(g, ucfg_gu[0], ucfg_gu[1], ucfg_gu[2]); + const xd = try b.zeros(P * I * 2, 0); + try o.gateupEpilogue(z.ptr, pick.ptr, svh_g.ptr, svh_u.ptr, suh_d.ptr, xd.ptr, R, P, I, sk_gu, slots, E, limit, act_mode); + try stream.synchronize(); + try same(gpu, fx, xd, "xd", "gate/up epilogue Xd"); + + var dn: exl3.Grouped = .{ .x0 = xd.ptr, .x1 = xd.ptr, .tp0 = td.ptr, .tp1 = td.ptr, .k2_0 = k2d.ptr, .k2_1 = k2d.ptr, .uids = uids.ptr, .ucount = ucount.ptr, .members = members.ptr, .z = z.ptr, .mats = 1, .K = I, .N = D, .P = P, .SK = sk_d, .slots = slots, .maxm = R, .nexp = maxu, .lo = lo_d, .hi = hi_d }; + const ucfg_d = [3]u32{ @intCast(try fx.int("nt_d")), @intCast(try fx.int("w_d")), @intCast(try fx.int("pf_d")) }; + try zeroThen(z, zlen_d); + try o.grouped(dn, ucfg_d[0], ucfg_d[1], ucfg_d[2]); + try stream.synchronize(); + try same(gpu, fx, z, "z_d", "upstream grouped down Z"); + try settings(gpu, fx, o, stream, z, dn, zlen_d, "d", &n_ld, &n_pf); + + try zeroThen(z, zlen_d); + try o.grouped(dn, ucfg_d[0], ucfg_d[1], ucfg_d[2]); + const y = try b.zeros(P * D * 4, 0); + try o.downEpilogue(z.ptr, pick.ptr, svh_d.ptr, y.ptr, R, P, D, sk_d, slots, E); + try stream.synchronize(); + try same(gpu, fx, y, "y", "down epilogue Y"); + const out = try b.zeros(R * D * 4, 0); + try o.combine(y.ptr, wts.ptr, out.ptr, R, D, slots); + try stream.synchronize(); + try same(gpu, fx, out, "out", "combine"); + try y.fill8(0, null); + try out.fill8(0, null); + try o.downCombine(z.ptr, pick.ptr, svh_d.ptr, y.ptr, wts.ptr, out.ptr, R, P, D, sk_d, slots, E, false); + try stream.synchronize(); + try same(gpu, fx, out, "out", "down_combine"); + dn.z = 0; + check.pass("BITEXACT experts {s}: R {d} x {d} slots, E {d}, D {d}, I {d}, K2 gu [{d},{d}] d [{d},{d}]: group, rot_in, grouped, {d} x3ld and {d} x3pf settings, epilogues, combine", .{ std.fs.path.basename(dir), R, slots, E, D, I, lo_gu, hi_gu, lo_d, hi_d, n_ld, n_pf }); +} + +/// The reference zeroes the whole Z before each launch: rows a launch does not write stay 0 on both sides. +fn zeroThen(z: cuda.DeviceBuffer, len: usize) !void { + _ = len; + try z.fill8(0, null); +} + +/// Every x3ld setting (and PDL) and every x3pf setting that fits this matrix: Z against the reference of that setting. +fn settings(gpu: Gpu, fx: Fixture, o: exl3.Ops, stream: cuda.Stream, z: cuda.DeviceBuffer, g: exl3.Grouped, zlen: usize, tag: []const u8, n_ld: *usize, n_pf: *usize) !void { + var name_buf: [64]u8 = undefined; + for (exl3.ld_cfgs) |c| for ([_]u32{ 0, 3 }) |probe| { + _ = exl3.x3ldCheck(g, c[0], c[1], probe) catch continue; + if (probe == 3) continue; // timing only: wrong results by design + for ([_]bool{ false, true }) |pdl| { + const name = try std.fmt.bufPrint(&name_buf, "z_{s}_ld_{d}_{d}", .{ tag, c[0], c[1] }); + try zeroThen(z, zlen); + try o.x3ld(g, c[0], c[1], 0, pdl); + try stream.synchronize(); + if (!has(fx, name)) return error.FixtureMissingSetting; + try same(gpu, fx, z, name, name); + n_ld.* += 1; + } + }; + for (exl3.pf_cfgs) |c| { + _ = exl3.x3pfCheck(g, c[0], c[1]) catch continue; + const name = try std.fmt.bufPrint(&name_buf, "z_{s}_pf_{d}_{d}", .{ tag, c[0], c[1] }); + if (!has(fx, name)) continue; // the oracle runs x3pf at prefill row counts only + try zeroThen(z, zlen); + try o.x3pf(g, c[0], c[1]); + try stream.synchronize(); + try same(gpu, fx, z, name, name); + n_pf.* += 1; + } +} + +/// x3gm over one row block (rot, gate/up and down launches per width, combine): every oracle configuration checked. +pub fn x3gm(gpu: Gpu, k: *const dsv41.Kernels, dir: []const u8) !void { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var b: Bufs = .{ .gpu = gpu, .fx = &fx }; + defer b.deinit(); + var stream = try cuda.Stream.init(gpu.d, false); // blocking: ordered after the uploads and memsets on the default stream + defer stream.deinit(); + const o = k.experts(stream); + + const R = try size(fx, "R"); + const slots = try size(fx, "slots"); + const D = try size(fx, "D"); + const I = try size(fx, "I"); + const P = R * slots; + const shx = (try fx.int("shx")) != 0; + const mats_x: usize = if (shx) 1 else 2; + const ragged = (try fx.int("ragged")) != 0; + const limit: f32 = @floatCast(try fx.float("limit")); + const input: exl3.Input = if (std.mem.eql(u8, try fx.string("input"), "bf16")) .bf16 else .f16; + + const x = try b.up("x"); + const pick = try b.up("pick"); + const wts = try b.up("wts"); + const words = try b.up("words"); + const suh_g = try b.up("suh_g"); + const suh_u = try b.up("suh_u"); + const svh_g = try b.up("svh_g"); + const svh_u = try b.up("svh_u"); + const suh_d = try b.up("suh_d"); + const svh_d = try b.up("svh_d"); + // the stacked form: base + 0 / base + gate bytes / ...; the ragged form: pointer tables + const tg = if (ragged) try b.table(words, "off_g") else words; + const tu = if (ragged) try b.table(words, "off_u") else words; + const td = if (ragged) try b.table(words, "off_d") else words; + const base_u: u64 = if (ragged) 0 else @intCast(try fx.int("base_u")); + const base_d: u64 = if (ragged) 0 else @intCast(try fx.int("base_d")); + + const xs = try b.zeros(mats_x * P * D * 2, 0); + const xg = xs.ptr; + const xu = xs.ptr + (mats_x - 1) * P * D * 2; + try o.gmRot(input, x.ptr, D, pick.ptr, suh_g.ptr, suh_u.ptr, xg, xu, D, slots, P, mats_x); + try stream.synchronize(); + try same(gpu, fx, xs, "xs", "x3gm rot"); + + const xd = try b.zeros(P * I * 2, 0); + const y = try b.zeros(P * D * 4, 0); + const ticket = try b.zeros(4, 0); + const out = try b.zeros(R * D * 4, 0); + var name_buf: [96]u8 = undefined; + const runs = try fx.int("runs"); + var launches: usize = 0; + for (0..@intCast(runs)) |r| { + // run r: a configuration per (kind, width) run__ and its plan run___ + const tk = (try fx.int(try std.fmt.bufPrint(&name_buf, "run{d}_ticket", .{r}))) != 0; + try xd.fill8(0, null); + try y.fill8(0, null); + for (exl3.gm_k2s) |k2| { + const pl = (try plan(&b, fx, r, "gu", k2)) orelse continue; + const cg: usize = @intCast(try fx.int(try std.fmt.bufPrint(&name_buf, "run{d}_gu_{d}", .{ r, k2 }))); + try nativePlan(&b, fx, o, stream, r, "gu", k2, P, cg); + try ticket.fill8(0, null); + try o.gmGateup(xg, xu, tg.ptr, tu.ptr + base_u, pl[0], pl[1], pl[2], pl[3], pl[4], ticket.ptr, svh_g.ptr, svh_u.ptr, suh_d.ptr, xd.ptr, D, I, k2, shx, cg, tk, limit, ragged); + launches += 1; + } + try stream.synchronize(); + try same(gpu, fx, xd, try std.fmt.bufPrint(&name_buf, "run{d}_xd", .{r}), "x3gm gate/up Xd"); + for (exl3.gm_k2s) |k2| { + const pl = (try plan(&b, fx, r, "dn", k2)) orelse continue; + const cd: usize = @intCast(try fx.int(try std.fmt.bufPrint(&name_buf, "run{d}_dn_{d}", .{ r, k2 }))); + try nativePlan(&b, fx, o, stream, r, "dn", k2, P, cd); + try ticket.fill8(0, null); + try o.gmDown(xd.ptr, td.ptr + base_d, pl[0], pl[1], pl[2], pl[3], pl[4], ticket.ptr, svh_d.ptr, y.ptr, I, D, k2, cd, tk, ragged); + launches += 1; + } + try stream.synchronize(); + try same(gpu, fx, y, try std.fmt.bufPrint(&name_buf, "run{d}_y", .{r}), "x3gm down Y"); + try o.combine(y.ptr, wts.ptr, out.ptr, R, D, slots); + try stream.synchronize(); + try same(gpu, fx, out, try std.fmt.bufPrint(&name_buf, "run{d}_out", .{r}), "x3gm combine"); + } + var v2_launches: usize = 0; + if (has(fx, "v2_xd") or has(fx, "v2_t1_xd")) { + // x3gm v2: one plan over every pair, one gate/up and one down launch over every width + const k2g = try b.up("v2_k2g"); + const k2d = try b.up("v2_k2d"); + const pl = [_]u64{ (try b.up("v2_order")).ptr, (try b.up("v2_pe")).ptr, (try b.up("v2_poff")).ptr, (try b.up("v2_pcnt")).ptr, (try b.up("v2_npass")).ptr }; + const cg: usize = @intCast(try fx.int("v2_gu")); + const cd: usize = @intCast(try fx.int("v2_dn")); + for ([_]bool{ true, false }) |tk| { + try xd.fill8(0, null); + try y.fill8(0, null); + try ticket.fill8(0, null); + try o.gm2Gateup(xg, xu, tg.ptr, tu.ptr + base_u, k2g.ptr, pl[0], pl[1], pl[2], pl[3], pl[4], ticket.ptr, svh_g.ptr, svh_u.ptr, suh_d.ptr, xd.ptr, D, I, shx, cg, tk, limit, ragged); + try ticket.fill8(0, null); + try o.gm2Down(xd.ptr, td.ptr + base_d, k2d.ptr, pl[0], pl[1], pl[2], pl[3], pl[4], ticket.ptr, svh_d.ptr, y.ptr, I, D, cd, tk, ragged); + try o.combine(y.ptr, wts.ptr, out.ptr, R, D, slots); + try stream.synchronize(); + const t: u8 = @intFromBool(tk); + try same(gpu, fx, xd, try std.fmt.bufPrint(&name_buf, "v2_t{d}_xd", .{t}), "x3gm v2 gate/up Xd"); + try same(gpu, fx, y, try std.fmt.bufPrint(&name_buf, "v2_t{d}_y", .{t}), "x3gm v2 down Y"); + try same(gpu, fx, out, try std.fmt.bufPrint(&name_buf, "v2_t{d}_out", .{t}), "x3gm v2 combine"); + v2_launches += 2; + } + } + if (v2_launches == 0) { + // No v2 arrays in the fixture: gm2 against gm_kernel outputs (run 0, ticket on), same bits by v2's contract. + const E = try size(fx, "E"); + const T = exl3.planPasses(P, E, 64); + const order = try b.zeros(P * 4, 0); + const pe = try b.zeros(T * 4, 0); + const poff = try b.zeros(T * 4, 0); + const pcnt = try b.zeros(T * 4, 0); + const npass = try b.zeros(4, 0); + try o.gmPlan(pick.ptr, P, E, 64, order.ptr, pe.ptr, poff.ptr, pcnt.ptr, npass.ptr); + const k2g = try b.up("k2_g"); + const k2d = try b.up("k2_d"); + const cfg = exl3.gmTuned2(null, null, shx); + for ([_]bool{ true, false }) |tk| { + try xd.fill8(0, null); + try y.fill8(0, null); + try ticket.fill8(0, null); + try o.gm2Gateup(xg, xu, tg.ptr, tu.ptr + base_u, k2g.ptr, order.ptr, pe.ptr, poff.ptr, pcnt.ptr, npass.ptr, ticket.ptr, svh_g.ptr, svh_u.ptr, suh_d.ptr, xd.ptr, D, I, shx, cfg[0], tk, limit, ragged); + try ticket.fill8(0, null); + try o.gm2Down(xd.ptr, td.ptr + base_d, k2d.ptr, order.ptr, pe.ptr, poff.ptr, pcnt.ptr, npass.ptr, ticket.ptr, svh_d.ptr, y.ptr, I, D, cfg[1], tk, ragged); + try o.combine(y.ptr, wts.ptr, out.ptr, R, D, slots); + try stream.synchronize(); + try same(gpu, fx, xd, "run0_xd", "x3gm v2 Xd vs gm_kernel"); + try same(gpu, fx, y, "run0_y", "x3gm v2 Y vs gm_kernel"); + try same(gpu, fx, out, "run0_out", "x3gm v2 combine vs gm_kernel"); + v2_launches += 2; + } + } + // x3gm v3 (gm3_kernel): every variant against gm_kernel outputs (run 0), both ticket modes, one plan (bm 64). + var v3_launches: usize = 0; + { + const E = try size(fx, "E"); + const T = exl3.planPasses(P, E, 64); + const order = try b.zeros(P * 4, 0); + const pe = try b.zeros(T * 4, 0); + const poff = try b.zeros(T * 4, 0); + const pcnt = try b.zeros(T * 4, 0); + const npass = try b.zeros(4, 0); + try o.gmPlan(pick.ptr, P, E, 64, order.ptr, pe.ptr, poff.ptr, pcnt.ptr, npass.ptr); + const k2g = try b.up("k2_g"); + const k2d = try b.up("k2_d"); + const xd_want = try b.up("run0_xd"); + var vb: [exl3.v3_variants.len]usize = undefined; + for (exl3.v3Of(2, @intCast(mats_x), &vb)) |v| for ([_]bool{ true, false }) |tk| { + try xd.fill8(0, null); + try ticket.fill8(0, null); + try o.gm3Gateup(v, xg, xu, tg.ptr, tu.ptr + base_u, k2g.ptr, order.ptr, pe.ptr, poff.ptr, pcnt.ptr, npass.ptr, ticket.ptr, svh_g.ptr, svh_u.ptr, suh_d.ptr, xd.ptr, D, I, shx, tk, limit, ragged); + try stream.synchronize(); + try same(gpu, fx, xd, "run0_xd", try std.fmt.bufPrint(&name_buf, "x3gm v3 gate/up variant {d} Xd vs gm_kernel", .{v})); + v3_launches += 1; + }; + for (exl3.v3Of(1, 1, &vb)) |v| for ([_]bool{ true, false }) |tk| { + try y.fill8(0, null); + try ticket.fill8(0, null); + try o.gm3Down(v, xd_want.ptr, td.ptr + base_d, k2d.ptr, order.ptr, pe.ptr, poff.ptr, pcnt.ptr, npass.ptr, ticket.ptr, svh_d.ptr, y.ptr, I, D, tk, ragged); + try stream.synchronize(); + try same(gpu, fx, y, "run0_y", try std.fmt.bufPrint(&name_buf, "x3gm v3 down variant {d} Y vs gm_kernel", .{v})); + v3_launches += 1; + }; + } + if (has(fx, "dequant_in")) { + const t = try b.up("dequant_in"); + const dk: usize = @intCast(try fx.int("dequant_k")); + const dn: usize = @intCast(try fx.int("dequant_n")); + const dk2: u32 = @intCast(try fx.int("dequant_k2")); + const dq = try b.zeros(dk * dn * 2, 0); + try o.gmDequant(t.ptr, dq.ptr, dk, dn, dk2); + try stream.synchronize(); + try same(gpu, fx, dq, "dequant_gm", "x3gm dequant"); + try dq.fill8(0, null); + try o.dequant(t.ptr, dq.ptr, dk, dn, dk2); + try stream.synchronize(); + try same(gpu, fx, dq, "dequant_up", "upstream dequant"); + } + check.pass("BITEXACT x3gm {s}: R {d} x {d} slots, D {d}, I {d}, {s}, shx {}: rot, {d} runs ({d} launches), v2 {d} launches, v3 {d} launches, combine", .{ std.fs.path.basename(dir), R, slots, D, I, if (ragged) "ragged" else "stacked", shx, runs, launches, v2_launches, v3_launches }); +} + +pub const pointwiseCases = experts_cases.pointwiseCases; +pub const glueCases = experts_cases.glueCases; +pub const topkCases = experts_cases.topkCases; +pub const planCases = experts_cases.planCases; + +/// x3gm_plan.cu on the plan's input picks against torch's x3gm.plan (the arrays the launches then use). +fn nativePlan(b: *Bufs, fx: Fixture, o: exl3.Ops, stream: cuda.Stream, r: usize, kind: []const u8, k2: u32, P: usize, cfg: usize) !void { + var name_buf: [64]u8 = undefined; + const pk = try b.up(try std.fmt.bufPrint(&name_buf, "run{d}_{s}_{d}_pk", .{ r, kind, k2 })); + const E = try size(fx, "E"); + const bm: usize = exl3.gmMembers(if (kind[0] == 'g') .gu else .dn, cfg); + const T = exl3.planPasses(P, E, bm); + const order = try b.zeros(P * 4, 0xee); + const pe = try b.zeros(T * 4, 0xee); + const poff = try b.zeros(T * 4, 0xee); + const pcnt = try b.zeros(T * 4, 0xee); + const npass = try b.zeros(4, 0xee); + try o.gmPlan(pk.ptr, P, E, bm, order.ptr, pe.ptr, poff.ptr, pcnt.ptr, npass.ptr); + try stream.synchronize(); + for ([_][]const u8{ "order", "pe", "poff", "pcnt", "npass" }, [_]cuda.DeviceBuffer{ order, pe, poff, pcnt, npass }) |part, buf| { + const name = try std.fmt.bufPrint(&name_buf, "run{d}_{s}_{d}_{s}", .{ r, kind, k2, part }); + try same(b.gpu, fx, buf, name, "x3gm_plan.cu vs x3gm.plan"); + } +} + +/// The plan arrays of (run, kind, width), uploaded; null when that width has no launch in the run. +fn plan(b: *Bufs, fx: Fixture, r: usize, kind: []const u8, k2: u32) !?[5]u64 { + var out: [5]u64 = undefined; + var name_buf: [64]u8 = undefined; + for ([_][]const u8{ "order", "pe", "poff", "pcnt", "npass" }, 0..) |part, i| { + const name = try std.fmt.bufPrint(&name_buf, "run{d}_{s}_{d}_{s}", .{ r, kind, k2, part }); + if (!has(fx, name)) return null; + out[i] = (try b.up(name)).ptr; + } + return out; +} diff --git a/zig/tests/cuda/deepseek_v41/experts_cases.zig b/zig/tests/cuda/deepseek_v41/experts_cases.zig new file mode 100644 index 0000000000..6251342d1c --- /dev/null +++ b/zig/tests/cuda/deepseek_v41/experts_cases.zig @@ -0,0 +1,197 @@ +//! Bit checks of the pointwise, glue, top-k keys and x3gm plan kernels against torch. + +const std = @import("std"); +const cuda = @import("cuda"); +const dsv41 = @import("dsv41_kernels"); +const check = @import("../check.zig"); +const fixture = @import("../fixture.zig"); +const Fixture = fixture.Fixture; +const experts = @import("experts.zig"); + +const exl3 = dsv41.exl3; +const Gpu = check.Gpu; +const Bufs = experts.Bufs; +const size = experts.size; +const same = experts.same; + +/// pointwise.cu against torch: the shared expert's SwiGLU with the limit, and the partial adds (fp32, into bf16). +pub fn pointwiseCases(gpu: Gpu, k: *const dsv41.Kernels, dir: []const u8) !void { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var b: Bufs = .{ .gpu = gpu, .fx = &fx }; + defer b.deinit(); + var stream = try cuda.Stream.init(gpu.d, false); // blocking: ordered after the uploads and memsets on the default stream + defer stream.deinit(); + const o = k.others(stream); + const R = try size(fx, "R"); + const n = try size(fx, "n"); + const D = try size(fx, "D"); + const g = try b.up("g"); + const u = try b.up("u"); + const act = try b.zeros(R * n * 4, 0xee); + try o.swiglu(g.ptr, n, u.ptr, n, act.ptr, n, R, n, @floatCast(try fx.float("limit"))); + try stream.synchronize(); + try same(gpu, fx, act, "act", "swiglu vs torch clamp_ / silu / mul_"); + const a = try b.up("a"); + const bb = try b.up("b"); + const y16 = try b.zeros(R * D * 2, 0xee); + try o.addBf16(a.ptr, bb.ptr, y16.ptr, R * D); + try o.addInto(a.ptr, bb.ptr, R * D); + try stream.synchronize(); + try same(gpu, fx, a, "sum32", "out += shared"); + try same(gpu, fx, y16, "sum16", "torch.add(out=bf16)"); + check.pass("BITEXACT pointwise.cu == torch: swiglu with limit ({d} x {d}, NaN / inf / -0.0), adds into fp32 and bf16", .{ R, n }); +} + +/// glue.cu against torch: the window's embedding send / sum, the casts, positions, gather_cols and the carry. +pub fn glueCases(gpu: Gpu, k: *const dsv41.Kernels, dir: []const u8) !void { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var b: Bufs = .{ .gpu = gpu, .fx = &fx }; + defer b.deinit(); + var stream = try cuda.Stream.init(gpu.d, false); // blocking: ordered after the uploads on the default stream + defer stream.deinit(); + const o = k.others(stream); + const V = try size(fx, "V"); + const D = try size(fx, "D"); + const W = try size(fx, "W"); + const n = try size(fx, "n"); + const c = try size(fx, "c"); + + const embed = try b.up("embed"); + const ids = try b.up("ids"); + const send = try b.zeros(n * D * 2, 0xee); + try o.embedSend(ids.ptr, n, embed.ptr, try fx.int("lo"), V, D, send.ptr); + const recv = try b.up("recv"); + const summed = try b.zeros(n * 4 * D * 2, 0xee); + try o.embedSum(recv.ptr, W, n, D, summed.ptr); + const x = try b.up("x"); + const cast = try b.zeros(n * D * 2, 0xee); + try o.castBf16(x.ptr, cast.ptr, n * D); + const kw = try b.zeros(n * D * 4, 0xee); + try o.kitWeights(x.ptr, kw.ptr, n * D); + const kl = try b.up("x"); + try o.kitLogits(kl.ptr, n * D); + const gc = try b.up("gc"); + const cols = try b.zeros(n * W * c * 2, 0xee); + try o.gatherCols(gc.ptr, W, n, c, cols.ptr); + const src = try b.up("src"); + const row: usize = @intCast(try fx.int("row")); + const carry = try b.zeros(1024 * 4, 0xee); + try o.carry(src.ptr, @intCast(try fx.int("ld")), 0, row, 1024, carry.ptr); + // the device-held row (graph-safe form): the same bits + const acc = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&@as(i32, @intCast(row)))); + try b.list.append(gpu.gpa, acc); + const carry2 = try b.zeros(1024 * 4, 0xee); + try o.carry(src.ptr, @intCast(try fx.int("ld")), acc.ptr, 0, 1024, carry2.ptr); + const start: i32 = @intCast(try fx.int("start")); + const pos = try b.zeros(4, 0xee); + const pos64 = try b.zeros(n * 8, 0xee); + const lo = try b.zeros(n * 4, 0xee); + try o.positions(start, n, pos.ptr, pos64.ptr, lo.ptr); + // the graphed window's form: start read from the device, the same values + const start_dev = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&start)); + try b.list.append(gpu.gpa, start_dev); + const pos_d = try b.zeros(4, 0xee); + const pos64_d = try b.zeros(n * 8, 0xee); + const lo_d = try b.zeros(n * 4, 0xee); + try o.positionsDev(start_dev.ptr, n, pos_d.ptr, pos64_d.ptr, lo_d.ptr); + // widen_f32: the recv rows (bf16, NaN / inf / -0 among them) to fp32, against the bits shifted up on the host + const wide = try b.zeros(W * n * D * 4, 0xee); + try o.widenF32(recv.ptr, wide.ptr, W * n * D); + try stream.synchronize(); + { + const recv_host = try fx.bytes("recv"); + defer gpu.gpa.free(recv_host); + const want = try gpu.gpa.alloc(u32, recv_host.len / 2); + defer gpu.gpa.free(want); + for (want, 0..) |*w, i| w.* = @as(u32, std.mem.readInt(u16, recv_host[2 * i ..][0..2], .little)) << 16; + const got = try gpu.gpa.alloc(u8, 4 * want.len); + defer gpu.gpa.free(got); + try wide.download(0, got); + try check.sameBytes("glue widen_f32", got, std.mem.sliceAsBytes(want)); + } + + try same(gpu, fx, send, "send", "glue embed_send"); + try same(gpu, fx, summed, "summed", "glue embed_sum"); + try same(gpu, fx, cast, "cast", "glue cast_bf16"); + try same(gpu, fx, kw, "kw", "glue kit_weights"); + try same(gpu, fx, kl, "kl", "glue kit_logits"); + try same(gpu, fx, cols, "cols", "glue gather_cols"); + try same(gpu, fx, carry, "carry", "glue carry (host row)"); + try same(gpu, fx, carry2, "carry", "glue carry (device row)"); + // positions: no torch op to record; Win.make's values + var want_pos: [1]i32 = .{start}; + const want64 = try gpu.gpa.alloc(i64, n); + defer gpu.gpa.free(want64); + const want_lo = try gpu.gpa.alloc(i32, n); + defer gpu.gpa.free(want_lo); + for (want64, want_lo, 0..) |*p, *l, i| { + p.* = start + @as(i64, @intCast(i)); + l.* = 0; + } + for ([_]struct { cuda.DeviceBuffer, []const u8 }{ .{ pos, std.mem.sliceAsBytes(&want_pos) }, .{ pos64, std.mem.sliceAsBytes(want64) }, .{ lo, std.mem.sliceAsBytes(want_lo) }, .{ pos_d, std.mem.sliceAsBytes(&want_pos) }, .{ pos64_d, std.mem.sliceAsBytes(want64) }, .{ lo_d, std.mem.sliceAsBytes(want_lo) } }) |pw| { + const got = try gpu.gpa.alloc(u8, pw[1].len); + defer gpu.gpa.free(got); + try pw[0].download(0, got); + try check.sameBytes("glue positions", got, pw[1]); + } + check.pass("BITEXACT glue.cu == torch: embed send / sum (W {d}), cast_bf16, kit_weights, kit_logits, gather_cols, carry, positions (host and device start), widen_f32", .{W}); +} + +/// topk_keys.cu against pick.top (torch.topk on pick.keys) at k 1 / 5 / 64 / 1,024 over 64,640 columns. +pub fn topkCases(gpu: Gpu, k: *const dsv41.Kernels, dir: []const u8) !void { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var b: Bufs = .{ .gpu = gpu, .fx = &fx }; + defer b.deinit(); + var stream = try cuda.Stream.init(gpu.d, false); // blocking: ordered after the uploads and memsets on the default stream + defer stream.deinit(); + const o = k.others(stream); + const R = try size(fx, "R"); + const C = try size(fx, "C"); + const lg = try b.up("lg"); + var name_buf: [32]u8 = undefined; + for (0..try size(fx, "ks")) |i| { + const kk = try size(fx, try std.fmt.bufPrint(&name_buf, "k{d}", .{i})); + const vals = try b.zeros(R * kk * 4, 0xee); + const cols = try b.zeros(R * kk * 8, 0xee); + try o.topKeys(lg.ptr, C, R, C, kk, vals.ptr, cols.ptr); + try stream.synchronize(); + try same(gpu, fx, cols, try std.fmt.bufPrint(&name_buf, "cols{d}", .{i}), "topk_keys cols"); + try same(gpu, fx, vals, try std.fmt.bufPrint(&name_buf, "vals{d}", .{i}), "topk_keys vals"); + } + check.pass("BITEXACT topk_keys.cu == pick.top on {d} rows x {d} columns, 4 k", .{ R, C }); +} + +/// x3gm_plan.cu against torch's x3gm.plan on full-sized picks (the "plan" case). +pub fn planCases(gpu: Gpu, k: *const dsv41.Kernels, dir: []const u8) !void { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var b: Bufs = .{ .gpu = gpu, .fx = &fx }; + defer b.deinit(); + var stream = try cuda.Stream.init(gpu.d, false); // blocking: ordered after the uploads and memsets on the default stream + defer stream.deinit(); + const o = k.experts(stream); + var name_buf: [64]u8 = undefined; + const n = try size(fx, "cases"); + for (0..n) |i| { + const P = try size(fx, try std.fmt.bufPrint(&name_buf, "c{d}_P", .{i})); + const E = try size(fx, try std.fmt.bufPrint(&name_buf, "c{d}_E", .{i})); + const bm = try size(fx, try std.fmt.bufPrint(&name_buf, "c{d}_bm", .{i})); + const T = exl3.planPasses(P, E, bm); + const pk = try b.up(try std.fmt.bufPrint(&name_buf, "c{d}_pk", .{i})); + const order = try b.zeros(P * 4, 0xee); + const pe = try b.zeros(T * 4, 0xee); + const poff = try b.zeros(T * 4, 0xee); + const pcnt = try b.zeros(T * 4, 0xee); + const npass = try b.zeros(4, 0xee); + try o.gmPlan(pk.ptr, P, E, bm, order.ptr, pe.ptr, poff.ptr, pcnt.ptr, npass.ptr); + try stream.synchronize(); + for ([_][]const u8{ "order", "pe", "poff", "pcnt", "npass" }, [_]cuda.DeviceBuffer{ order, pe, poff, pcnt, npass }) |part, buf| { + const name = try std.fmt.bufPrint(&name_buf, "c{d}_{s}", .{ i, part }); + try same(gpu, fx, buf, name, name); + } + } + check.pass("BITEXACT x3gm_plan.cu == torch x3gm.plan on {d} cases (up to 24,576 pairs, 384 experts)", .{n}); +} diff --git a/zig/tests/cuda/deepseek_v41/replay.zig b/zig/tests/cuda/deepseek_v41/replay.zig new file mode 100644 index 0000000000..a91abaf707 --- /dev/null +++ b/zig/tests/cuda/deepseek_v41/replay.zig @@ -0,0 +1,211 @@ +//! Replays one recorded extension call (dsv41_capture.py) on its recorded bytes; every viewed storage bit-compared. + +const std = @import("std"); +const cuda = @import("cuda"); +const dsv41 = @import("dsv41_kernels"); +const check = @import("../check.zig"); +const Fixture = @import("../fixture.zig").Fixture; + +const replay = dsv41.replay; +const Gpu = check.Gpu; +const Value = std.json.Value; + +pub const Outcome = enum { exact, skipped }; + +const Storage = struct { name: []const u8, buf: ?cuda.DeviceBuffer, len: usize }; + +const Ctx = struct { + gpa: std.mem.Allocator, + a: std.mem.Allocator, // the case's arena: argument trees + storages: []Storage, + extra: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn ptrOf(c: *Ctx, name: []const u8) !u64 { + for (c.storages) |s| if (std.mem.eql(u8, s.name, name)) return if (s.buf) |b| b.ptr else 0; + return error.UnknownStorage; + } + + fn dtype(s: []const u8) !replay.DType { + const map = .{ .{ "float16", .f16 }, .{ "bfloat16", .bf16 }, .{ "float32", .f32 }, .{ "float64", .f64 }, .{ "int8", .i8 }, .{ "uint8", .u8 }, .{ "int16", .i16 }, .{ "int32", .i32 }, .{ "int64", .i64 }, .{ "bool", .bool } }; + inline for (map) |m| if (std.mem.eql(u8, s, m[0])) return m[1]; + return error.UnknownDType; + } + + fn ints(c: *Ctx, v: Value) ![]const i64 { + const items = v.array.items; + const out = try c.a.alloc(i64, items.len); + for (items, out) |x, *o| o.* = x.integer; + return out; + } + + /// dsv41_capture.py's argument node -> a replay Arg. + fn arg(c: *Ctx, v: Value) !replay.Arg { + const o = v.object; + if (o.get("t")) |t| { + const shape = try c.ints(o.get("shape").?); + var numel: i64 = 1; + for (shape) |d| numel *= d; + // an empty tensor's pointer can precede its storage: pass null, as launchers expect for numel 0 + const raw = o.get("off").?.integer; + const base = try c.ptrOf(t.string); + const ptr: u64 = if (base == 0 or raw < 0 or numel == 0) 0 else base + @as(u64, @intCast(raw)); + return .{ .tensor = .{ .ptr = ptr, .dtype = try dtype(o.get("dtype").?.string), .shape = shape, .stride = try c.ints(o.get("stride").?) } }; + } + if (o.get("i")) |x| return .{ .int = x.integer }; + if (o.get("f")) |x| return .{ .float = switch (x) { + .float => |f| f, + .integer => |n| @floatFromInt(n), + else => return error.WrongType, + } }; + if (o.get("b")) |x| return .{ .boolean = x.bool }; + if (o.get("n") != null) return .none; + if (o.get("l")) |l| { + const items = l.array.items; + const out = try c.a.alloc(replay.Arg, items.len); + for (items, out) |x, *y| y.* = try c.arg(x); + return .{ .list = out }; + } + return error.WrongType; + } + + /// A zeroed device buffer of `n` bytes owned by the case. + fn zeros(c: *Ctx, gpu: Gpu, n: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(gpu.d, @max(n, 16)); + try b.fill8(0, null); + try c.extra.append(c.gpa, b); + return b.ptr; + } +}; + +fn upload(gpu: Gpu, fx: Fixture, storages: []Storage) !void { + var name_buf: [64]u8 = undefined; + for (storages) |s| { + if (s.buf) |b| { + const host = try fx.bytes(try std.fmt.bufPrint(&name_buf, "{s}_before", .{s.name})); + defer gpu.gpa.free(host); + try b.upload(0, host); + } + } +} + +/// Every storage against the recorded after-call bytes, except `skip` (a ticket / counter replaced by a variant). +fn compare(gpu: Gpu, fx: Fixture, storages: []Storage, what: []const u8, skip: ?[]const u8) !void { + var name_buf: [64]u8 = undefined; + for (storages) |s| { + const b = s.buf orelse continue; + if (skip) |k| if (std.mem.eql(u8, k, s.name)) continue; + const want = try fx.bytes(try std.fmt.bufPrint(&name_buf, "{s}_after", .{s.name})); + defer gpu.gpa.free(want); + const got = try gpu.gpa.alloc(u8, want.len); + defer gpu.gpa.free(got); + try b.download(0, got); + const label = try std.fmt.allocPrint(gpu.gpa, "{s} storage {s}", .{ what, s.name }); + defer gpu.gpa.free(label); + try check.sameBytes(label, got, want); + } +} + +fn storageOf(v: Value) ?[]const u8 { + const t = v.object.get("t") orelse return null; + return t.string; +} + +pub fn run(gpu: Gpu, k: *const dsv41.Kernels, dir: []const u8) !Outcome { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const call = fx.parsed.value.object.get("call") orelse return error.MissingField; + const ext = call.object.get("ext").?.string; + const func = call.object.get("func").?.string; + + // the storages, sized by their recorded bytes + const n: usize = @intCast(try fx.int("storages")); + const storages = try a.alloc(Storage, n); + var made: usize = 0; + defer for (storages[0..made]) |*s| if (s.buf) |*b| b.free(); + const arrays = fx.parsed.value.object.get("arrays").?.object; + for (storages, 0..) |*s, i| { + s.name = try std.fmt.allocPrint(a, "s{d}", .{i}); + const ent = arrays.get(try std.fmt.allocPrint(a, "s{d}_before", .{i})).?.object; + s.len = @intCast(ent.get("shape").?.array.items[0].integer); + s.buf = if (s.len == 0) null else try cuda.DeviceBuffer.alloc(gpu.d, s.len); + made += 1; + } + var c: Ctx = .{ .gpa = gpu.gpa, .a = a, .storages = storages }; + defer { + for (c.extra.items) |*b| b.free(); + c.extra.deinit(gpu.gpa); + } + const jargs = call.object.get("args").?.array.items; + const args = try a.alloc(replay.Arg, jargs.len); + for (jargs, args) |v, *o| o.* = try c.arg(v); + var kws: std.ArrayList(replay.Kw) = .empty; + var it = call.object.get("kwargs").?.object.iterator(); + while (it.next()) |e| try kws.append(a, .{ .name = e.key_ptr.*, .value = try c.arg(e.value_ptr.*) }); + + var stream = try cuda.Stream.init(gpu.d, false); // blocking: ordered after the uploads on the default stream + defer stream.deinit(); + const scratch_buf = try cuda.DeviceBuffer.alloc(gpu.d, 64 << 20); + var scratch_owned = scratch_buf; + defer scratch_owned.free(); + var scratch: replay.Scratch = .{ .ptr = scratch_buf.ptr, .len = scratch_buf.len }; + + try upload(gpu, fx, storages); + if (!try replay.call(k, stream, ext, func, args, kws.items, &scratch)) return .skipped; + try stream.synchronize(); + try compare(gpu, fx, storages, func, null); + + var variants: usize = 0; + // R1 variants of a pre-R1 call (the twin's kernel, today's launch): the same bits by contract + if (std.mem.eql(u8, ext, "tf_dsv41_attn_cuda_v1") and std.mem.eql(u8, func, "attn") and args.len == 23) { + const ticket_name = storageOf(jargs[18]).?; + const q = args[0].tensor; + const hb: usize = @intCast(@divTrunc(q.shape[1], 16)); + for ([_]i64{ 1, 4 }) |p| { + const v = try a.alloc(replay.Arg, 24); + @memcpy(v[0..23], args); + var t = args[18].tensor; + t.ptr = try c.zeros(gpu, @as(usize, @intCast(q.shape[0])) * hb * dsv41.ops.attn.MAXP * 4); + v[18] = .{ .tensor = t }; + v[23] = .{ .int = p }; + try upload(gpu, fx, storages); + if (!try replay.call(k, stream, ext, func, v, kws.items, &scratch)) return error.VariantNotBound; + try stream.synchronize(); + try compare(gpu, fx, storages, if (p == 1) "attn split 1" else "attn split 4", ticket_name); + variants += 1; + } + } + if (std.mem.eql(u8, ext, "tf_dsv41_mhc_cuda_v1") and std.mem.eql(u8, func, "run") and args.len == 24) { + const cnt_name = storageOf(jargs[17]).?; + const mode = args[19].int; + const opre = args[14].tensor; + var opre_n: i64 = 1; + for (opre.shape) |d| opre_n *= d; + // the deferred form only where the boundary computes the next site's coefficients (MIX: modes 0-2) + const forms: []const bool = if (mode <= 2 and opre_n > 0) &.{ false, true } else &.{false}; + for (forms) |deferred| { + const v = try a.alloc(replay.Arg, 26); + @memcpy(v[0..24], args); + var t = args[17].tensor; + t.ptr = try c.zeros(gpu, 16); // the tail's 3 self-resetting ints + v[17] = .{ .tensor = t }; + v[24] = .{ .int = 0 }; // spin 0: the parallel tail, no waiting + v[25] = .{ .int = @intFromBool(deferred) }; + try upload(gpu, fx, storages); + if (!try replay.call(k, stream, ext, func, v, kws.items, &scratch)) return error.VariantNotBound; + if (deferred) { + // coef_kernel: (part, base, scale, opre, opost, ocomb, R, eps, hc_eps, post_alpha, iters) + const cv = [_]replay.Arg{ args[9], args[7], args[8], args[14], args[15], args[16], args[18], args[20], args[21], args[22], args[23] }; + if (!try replay.call(k, stream, ext, "coef", &cv, &.{}, &scratch)) return error.VariantNotBound; + } + try stream.synchronize(); + try compare(gpu, fx, storages, if (deferred) "mhc tail + defer" else "mhc tail", cnt_name); + variants += 1; + } + } + const test_name = if (call.object.get("test")) |t| t.string else ""; + check.pass("BITEXACT replay {s}.{s} ({d} storages, {d} R1 variants) from {s}", .{ ext, func, n, variants, test_name }); + return .exact; +} diff --git a/zig/tests/cuda/dsv41_main.zig b/zig/tests/cuda/dsv41_main.zig new file mode 100644 index 0000000000..6a31d9bc7f --- /dev/null +++ b/zig/tests/cuda/dsv41_main.zig @@ -0,0 +1,119 @@ +//! GPU test runner for DeepSeek-V4.1's kernels: `tf-dsv41-test [fixture dir]`, bit-exact against oracle.py. + +const std = @import("std"); +const cuda = @import("cuda"); +const dsv41 = @import("dsv41_kernels"); +const check = @import("check.zig"); +const experts = @import("deepseek_v41/experts.zig"); +const replay_case = @import("deepseek_v41/replay.zig"); + +const usage = + \\usage: tf-dsv41-test [fixture dir] + \\ symbols load every DeepSeek-V4.1 fatbin and resolve every instance + \\ experts decode chain: group, rot_in, x3ld (every setting) / upstream grouped / x3pf, epilogues, combine + \\ x3gm fast prefill: rot, gate/up and down at every configuration (uniform and ragged), combine + \\ replay one recorded Python extension call (dsv41_capture.py) re-issued and compared, storage by storage + \\ all every case directory under by its manifest's "kind" + \\ +; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("{s}", .{usage}); + return 2; + } + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + const gpu: check.Gpu = .{ .d = &driver, .ctx = &ctx, .gpa = init.gpa, .io = init.io }; + run(gpu, args[1], args[2..]) catch |e| { + std.debug.print("FAIL {s}: {t}\n", .{ args[1], e }); + return 1; + }; + return 0; +} + +fn arg(rest: []const [:0]const u8, i: usize) ![]const u8 { + if (i >= rest.len) { + std.debug.print("{s}", .{usage}); + return error.MissingArgument; + } + return rest[i]; +} + +fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { + var k = try dsv41.Kernels.load(gpu.ctx); + defer k.deinit(); + if (std.mem.eql(u8, cmd, "symbols")) { + check.pass("dsv41 symbols: every instance resolved ({d} SMs)", .{k.sms}); + return; + } + if (std.mem.eql(u8, cmd, "experts")) return experts.decodeChain(gpu, &k, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "x3gm")) return experts.x3gm(gpu, &k, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "replay")) { + if (try replay_case.run(gpu, &k, try arg(rest, 0)) == .skipped) std.debug.print("SKIP no binding\n", .{}); + return; + } + if (std.mem.eql(u8, cmd, "all")) return all(gpu, &k, try arg(rest, 0)); + std.debug.print("{s}", .{usage}); + return error.UnknownCommand; +} + +/// Every case under `root` (one directory a case, its manifest's "kind" naming the check); fails if any fails. +fn all(gpu: check.Gpu, k: *const dsv41.Kernels, root: []const u8) !void { + var dir = try std.Io.Dir.cwd().openDir(gpu.io, root, .{ .iterate = true }); + defer dir.close(gpu.io); + var names: std.ArrayList([]u8) = .empty; + defer { + for (names.items) |n| gpu.gpa.free(n); + names.deinit(gpu.gpa); + } + var it = dir.iterate(); + while (try it.next(gpu.io)) |e| { + if (e.kind == .directory) try names.append(gpu.gpa, try gpu.gpa.dupe(u8, e.name)); + } + std.mem.sort([]u8, names.items, {}, struct { + fn lt(_: void, a: []u8, b: []u8) bool { + return std.mem.lessThan(u8, a, b); + } + }.lt); + var failed: usize = 0; + var skipped: usize = 0; + for (names.items) |n| { + const path = try std.fs.path.join(gpu.gpa, &.{ root, n }); + defer gpu.gpa.free(path); + const kind = experts.caseKind(gpu, path) catch |e| { + std.debug.print("FAIL {s}: no manifest ({t})\n", .{ n, e }); + failed += 1; + continue; + }; + const r = if (std.mem.eql(u8, kind.slice(), "experts")) + experts.decodeChain(gpu, k, path) + else if (std.mem.eql(u8, kind.slice(), "x3gm")) + experts.x3gm(gpu, k, path) + else if (std.mem.eql(u8, kind.slice(), "plan")) + experts.planCases(gpu, k, path) + else if (std.mem.eql(u8, kind.slice(), "topk")) + experts.topkCases(gpu, k, path) + else if (std.mem.eql(u8, kind.slice(), "pointwise")) + experts.pointwiseCases(gpu, k, path) + else if (std.mem.eql(u8, kind.slice(), "glue")) + experts.glueCases(gpu, k, path) + else if (std.mem.eql(u8, kind.slice(), "replay")) blk: { + const o = replay_case.run(gpu, k, path) catch |e| break :blk e; + if (o == .skipped) { + std.debug.print("SKIP {s}: no Zig binding\n", .{n}); + skipped += 1; + } + break :blk {}; + } else error.UnknownKind; + r catch |e| { + std.debug.print("FAIL {s}: {t}\n", .{ n, e }); + failed += 1; + }; + } + std.debug.print("RESULT {d} cases, {d} failed, {d} skipped (no binding)\n", .{ names.items.len, failed, skipped }); + if (failed > 0) return error.TestFailed; +}