diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 6acdde321e..e74bb9e833 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -38,3 +38,18 @@ jobs: (cd ../lean-base && python3 tools/zig/lean_check.py > ../lean-base.txt || true) python3 tools/zig/lean_check.py > ../lean-head.txt || true python3 tools/zig/lean_new.py ../lean-base.txt ../lean-head.txt + zig-windows: + runs-on: windows-latest + steps: + - uses: actions/checkout@v4 + - uses: mlugg/setup-zig@v2 + with: + version: 0.17.0 + # The same host unit tests on Windows: the platform layer (zig/src/platform) over the UCRT, Win32 and Winsock. + - name: zig build test + run: zig build test + # The server goldens and the Cortheon fuzz harness on Windows too; the goldens skip the signal-driven cases. + - name: zig build test-golden + run: zig build test-golden + - name: zig build fuzz-cortheon + run: zig build fuzz-cortheon diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5a9c4d1014..0630be4816 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -54,9 +54,9 @@ Record pass, fail and skip counts and explain any unrun platform or feature. For served decode, the existing standard-library client needs Python 3.11 or newer: ```sh -python3 tools/bench_openai.py http://127.0.0.1:8080 local-model \ +python3 tools/bench_openai.py http://127.0.0.1:8888 local-model \ --tokens 256 --reps 3 --temperatures 0 --label candidate --output decode.json -python3 tools/bench_concurrent.py http://127.0.0.1:8080 local-model \ +python3 tools/bench_concurrent.py http://127.0.0.1:8888 local-model \ --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial --output concurrent.json ``` diff --git a/README.md b/README.md index d77b9f9f3e..78b99d0745 100644 --- a/README.md +++ b/README.md @@ -23,15 +23,17 @@ tensorfold serve "$HOME/models/nemotron-lightning" \ Use a complete local model directory or an existing Hugging Face cache. The binary in a release archive is `bin/tensorfold-native`; the [runbook](RUNBOOK.md) covers archive installation and CUDA runtime files. + +On Windows with an NVIDIA GPU, build `tensorfold-native.exe` from source as [docs/WINDOWS.md](docs/WINDOWS.md) describes. Use `tensorfold-native --help` or `-h`, and `tensorfold-native capabilities --json`, to inspect the installed binary. Model weights have separate downloads and licenses. -The API listens at `http://127.0.0.1:8080/v1` by default. +The API listens at `http://127.0.0.1:8888/v1` by default. It serves OpenAI chat completions, completions and Responses, plus Anthropic Messages. Use the model ID returned by `/v1/models`, or set one with `--name`. ```sh -curl -fsS http://127.0.0.1:8080/v1/chat/completions \ +curl -fsS http://127.0.0.1:8888/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{"model":"local-model","messages":[{"role":"user","content":"Say hello in one sentence."}],"max_tokens":128,"temperature":0}' ``` @@ -47,11 +49,11 @@ tensorfold pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit src=$(ls -d ~/.cache/huggingface/hub/models--TensorFold--NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit/snapshots/* | head -1) mkdir -p ~/models && cp -cRL "$src" ~/models/nemotron-living # always teach a copy tensorfold serve ~/models/nemotron-living --slide -curl -N http://127.0.0.1:8080/v1/slide/learn -H 'Content-Type: application/json' -d '{"text": "I like blue."}' +curl -N http://127.0.0.1:8888/v1/slide/learn -H 'Content-Type: application/json' -d '{"text": "I like blue."}' ``` Send it a document instead and it finds the facts in it and learns them one by one. Open -`http://127.0.0.1:8080/slide` to watch in 3D: the facts, the neurons each one took in the model's layers, the shard +`http://127.0.0.1:8888/slide` to watch in 3D: the facts, the neurons each one took in the model's layers, the shard files that hold them, and the lesson running now. Learning runs on Apple-silicon Macs and on an NVIDIA GB10 (DGX Spark), and a folder that learned on one serves on the @@ -71,8 +73,9 @@ The platform column names the hardware tested for each model. | Qwen3.8 Flash Next | MLX affine 6-bit, group 32 | Metal on M5 Ultra | | GLM-5.3-Flash | MLX affine 4-bit, group 64 | Metal on two M5 Ultras | | GLM-5.3 | EXLR8 `home` profile (trellis K2/K3/K4 experts) | Metal on two M5 Ultras | -| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max | -| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra | +| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max; CUDA on RTX 4090 (Windows) | +| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (Windows) | +| Qwen3.6-35B-A3B | MLX affine 4-bit, group 64, 8-bit routers; its MTP layer in `mtp-4bit.safetensors` | CUDA on RTX 4090 (Windows) | Nemotron's named checkpoint is `TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`. NVIDIA's NVFP4 checkpoint, `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4`, serves on CUDA as stored; its bf16 MTP head is quantized to 4-bit at load for drafts. @@ -88,6 +91,7 @@ Its prompt processing trails mlx_lm for now, and we are fixing it. An M3 Ultra r Bonsai, Gemma 4, Qwen3.6 and DeepSeek-V4 are still under qualification for 1.0.x. The Python 0.6.6 engine remains available at the `v0.6.6` tag for those models and other backends. On CUDA, 1.0.11 serves Nemotron on a GB10, greedy and sampled, with concurrent requests sharing each round. +On Windows ([docs/WINDOWS.md](docs/WINDOWS.md)), the native build also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees, and Qwen3.6-35B-A3B with its MTP layer drafting chains. The Linux x86_64 archive adds NVIDIA Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), tested on an RTX 3090, and Blackwell cards with compute capability 12.0 (the RTX 50 series and RTX PRO 6000 Blackwell), tested on an RTX PRO 6000; the A100 (8.0) is not supported yet. ## Exact decoding @@ -150,7 +154,7 @@ The binary's `capabilities --json` response lists its supported flags and platfo | Flag | Meaning | | --- | --- | -| `--host HOST`, `--port PORT` | Listen address and port, default `127.0.0.1:8080`. | +| `--host HOST`, `--port PORT` | Listen address and port, default `127.0.0.1:8888`. | | `--name NAME` | Model ID advertised to clients. | | `--alias NAME` | Another accepted model ID; repeat for multiple aliases. | | `--api-key KEY` | Require a bearer key; repeat for multiple keys. | diff --git a/RUNBOOK.md b/RUNBOOK.md index 408f2d91e8..0d0282fb49 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -104,14 +104,14 @@ bin/tensorfold-native serve "$HOME/models/nemotron-lightning" \ Leave the server running and check it from another terminal: ```sh -curl -fsS http://127.0.0.1:8080/health -curl -fsS http://127.0.0.1:8080/v1/models -curl -fsS http://127.0.0.1:8080/v1/chat/completions \ +curl -fsS http://127.0.0.1:8888/health +curl -fsS http://127.0.0.1:8888/v1/models +curl -fsS http://127.0.0.1:8888/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{"model":"local-model","messages":[{"role":"user","content":"Say hello in one sentence."}],"max_tokens":128,"temperature":0}' ``` -Use `http://127.0.0.1:8080/v1` as the OpenAI client base URL. +Use `http://127.0.0.1:8888/v1` as the OpenAI client base URL. For streaming, add `"stream": true` to the request. The native server also has `/v1/responses` and Anthropic `/v1/messages` routes. Serve flags and request overrides are described in [README.md](README.md#serve-flags). diff --git a/build.zig b/build.zig index 08765bd1df..01a2a30bb4 100644 --- a/build.zig +++ b/build.zig @@ -12,6 +12,7 @@ const counters_build = @import("zig/build/metal_counters.zig"); const exl3_build = @import("zig/build/exl3.zig"); const exlr8_build = @import("zig/build/exlr8.zig"); const golden_build = @import("zig/build/golden.zig"); +const sys_build = @import("zig/build/sys.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -153,13 +154,13 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&cost_run.step); switch (target.result.os.tag) { .macos => metalTargets(b, target, optimize, draft_ids, build_options, test_step), - .linux => cuda_build.targets(b, target, optimize, draft_ids, build_options), + .linux, .windows => cuda_build.targets(b, target, optimize, draft_ids, build_options), else => {}, } const xgrammar = grammar_build.option(b); dist_build.targets(b, draft_ids, build_options, release_version, xgrammar); cuda_build.hostTests(b, draft_ids, build_options, test_step, target.result.os.tag != .macos); - if (target.result.os.tag == .linux) { + if (target.result.os.tag == .linux or target.result.os.tag == .windows) { // the engine seam's tests (lane_host.zig and its siblings), which the macOS branch runs from its own api module const lanes_mod = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const api_mod = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes_mod }} }); @@ -189,6 +190,7 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&warning_run.step); b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); + sys_build.wire(b); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md new file mode 100644 index 0000000000..8bd8bf5845 --- /dev/null +++ b/docs/WINDOWS.md @@ -0,0 +1,90 @@ +# TensorFold on Windows (native) + +`tensorfold-native.exe` runs on Windows with an NVIDIA GPU through the same CUDA engine the Linux build uses. Only the +CUDA backend is built on Windows; Metal is macOS's and HIP is not built. It is tested on an RTX 4090 (x86_64); the +host code also cross-compiles for `aarch64-windows` with the GB10 (sm_121) kernels, untested on hardware so far. + +## What you need + +- An NVIDIA GPU and a current display driver (it ships `nvcuda.dll`; no CUDA toolkit is needed to run). +- To build: [Zig 0.17.0](https://ziglang.org/download/) (the version in `.zig-version`), the CUDA Toolkit + (13.x tested) for `nvcc`, and Visual Studio Build Tools with the C++ workload (`cl.exe`, nvcc's host compiler). + +## Build + +From a "x64 Native Tools" prompt (or after running `vcvars64.bat`), so `cl.exe` is on `PATH`: + +```bat +zig build native install -Dsm=89 -Dstrip=true ^ + -Dnvcc="%CUDA_PATH%\bin\nvcc.exe" +``` + +`-Dsm` takes the GPU's compute capability without the dot, comma separated for several: 86 (RTX 30), 89 (RTX 40), +120 (RTX 50). The server is `zig-out\native\bin\tensorfold-native.exe`; `zig-out\bin\tf-cuda-test.exe` runs the GPU +suites (`tf-cuda-test info`, `smoke`, `glue`, `sample`, `quant-ref`, `modelopt`, `exlr8`). + +`zig build test` runs the host tests without a GPU. + +## Run + +```bat +tensorfold-native pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit +tensorfold-native models +tensorfold-native serve --name local-model --parallel 1 --context 8192 --temperature 0 --no-thinking +``` + +The server listens on `http://127.0.0.1:8888/v1`. Ctrl+C (or closing the console) stops it cleanly. + +## Models on CUDA + +Nemotron 3.5 Lightning serves as on Linux. The Qwen3.5 hybrids have a CUDA engine of their own in this port +(`zig/src/families/qwen3_5/cuda*.zig`, kernels `zig/kernels/cuda/qwen35*.cu`), the Metal kernels' arithmetic with +the same rounding points, so a window of up to 16 rows reproduces one-token steps bit for bit: + +```bat +tensorfold-native serve --name qwen2b --context 32768 +tensorfold-native serve --name qwen27b --context 8192 --parallel 4 ^ + --drafter +tensorfold-native serve --name qwen36 --context 8192 --parallel 2 ^ + --drafter +``` + +- Qwen3.5-2B: about 340 tokens/s for one request on an RTX 4090, 1,500 across eight. +- Qwen3.8-27B (16 GiB): 40 tokens/s plain, at the card's memory bandwidth; 240 across eight requests. +- `--drafter` (4-bit only) loads z-lab's DFlash2 draft model: the target writes five layers' residuals into a ring + for it, and it drafts a tree of up to 12 lanes a round that the round loop's planner shapes from what the target + took. A lone greedy request decodes at 80 to 160 tokens/s by how predictable the text is (prose at the low end, + code at the high), sampled ones at 50 to 70, and four at once reach about 160 aggregate; replies stay identical to + plain decoding. Loading takes about 80 s: the drafter's weights are quantized to 4 bits on the host and every + shape is run once, as the driver stalls a process's first use of each kernel and buffer. +- Qwen3.6-35B-A3B (`qwen3_5_moe`, 18.6 GiB): the same layers with 256 routed experts, 8 a token, in place of the + dense MLP. Routing follows the frozen engine's `cuda/moe.py` rule, and the experts run through the grouped affine-4 + kernel's SwiGLU form, the shared expert last in the same tables. About 100 tokens/s plain for one request. With + `--drafter` naming the checkpoint's own directory, its MTP layer (`mtp-4bit.safetensors`) drafts chains of three + from each kept row's final normed state: 230 to 260 tokens/s greedy, 180 to 250 sampled, 420 aggregate across two + requests; replies stay identical to plain decoding. `TF_MTP_DRAFTS=N` sets the chain's length, and `TF_MTP_TREES=1` + drafts the planner's trees instead, which measured no faster than chains on this head. +- A verify window's rows are not snapshotted: each stream logs its rows' recurrence updates and a keep replays the + accepted path, so the 27B records 32 rows a forward where 8 fit before. +- `TF_TIMING=1` prints each verify, keep, absorb and hold's milliseconds and the drafted shape (chain or lanes); + `TF_NO_TREES=1` drafts chains only. +- No prompt cache (`--prompt-cache-gib`), `--learn` or `--slide` for these two models. + +## How the port works + +- `zig/src/platform/sys.zig` (`@import("sys")`, wired into every module by `zig/build/sys.zig`) is the operating + system as the engine calls it. On Linux and macOS it is `std.c`, `std.posix` and `std.DynLib` unchanged; on Windows + `win_c.zig` and `win_posix.zig` give the same names over the UCRT, Win32 and Winsock. +- Files are opened with every share mode, so rename-over and unlink-while-open behave as on POSIX. Checkpoint reads + skip the page cache (`FILE_FLAG_NO_BUFFERING`, as `O_DIRECT` on Linux). `/tmp` is the user's temporary folder and + `HOME` falls back to `USERPROFILE`. +- The Hugging Face cache links snapshot files to blobs. Windows allows symlinks only with Developer Mode or elevation, + so `pull` falls back to hard links (no extra space). +- Checkpoint reads use threads (io_uring is Linux's). +- The API key file is checked for type and size; Windows keeps access in ACLs, not mode bits, so the "unreadable by + other users" check is the folder's ACL rather than `chmod 600`. + +## Differences from Linux + +- No io_uring, `ulimit` or GB10 carveout, and no `SIGHUP` (the key file is still reread when it changes). +- `TF_HEAT_*` thermal pacing reads Linux's `/sys/class/thermal` and does nothing on Windows. diff --git a/docs/api.md b/docs/api.md index a0563b0048..657eca442a 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1,6 +1,6 @@ # Compatible APIs -The base URL is `http://127.0.0.1:8080/v1` with the default server settings. +The base URL is `http://127.0.0.1:8888/v1` with the default server settings. ## Native chat template override @@ -480,5 +480,5 @@ redacted thinking and other context edits return HTTP 400. Connect Claude Code using the ID from `/v1/models`: ```bash -ANTHROPIC_BASE_URL=http://127.0.0.1:8080 ANTHROPIC_API_KEY=local claude --model local-model +ANTHROPIC_BASE_URL=http://127.0.0.1:8888 ANTHROPIC_API_KEY=local claude --model local-model ``` diff --git a/docs/cache-benchmarks.md b/docs/cache-benchmarks.md index 2c9b9b6f95..e3977293f4 100644 --- a/docs/cache-benchmarks.md +++ b/docs/cache-benchmarks.md @@ -23,7 +23,7 @@ The standard-library-only cache client runs one first request, identical repeate concurrent group using the same prompt, greedy temperature, seed and reply limit. There is no warm-up: ```sh -python3 tools/bench_cache.py http://127.0.0.1:8080 MODEL \ +python3 tools/bench_cache.py http://127.0.0.1:8888 MODEL \ --tokens 64 --repeats 2 --streams 4 --require-cold --output cache.json ``` @@ -36,7 +36,7 @@ prefix retention (`--prompt-cache-gib 0`, where supported) for independent cold- For exact 2K/8K/32K/64K raw prefill, supply a JSON array of IDs from the checkpoint's own tokenizer: ```sh -python3 tools/bench_cache.py http://127.0.0.1:8080 MODEL \ +python3 tools/bench_cache.py http://127.0.0.1:8888 MODEL \ --prompt-file public-8192-tokens.json --expected-prompt-tokens 8192 \ --tokens 64 --require-cold --output cache-8192.json ``` diff --git a/docs/control.md b/docs/control.md index 4ed791242f..eb5e849364 100644 --- a/docs/control.md +++ b/docs/control.md @@ -118,7 +118,7 @@ The error from `tensorfold tui` prints that command with the venv's Python. `ric ```bash tensorfold tui tensorfold tui --profile code -tensorfold tui --url http://127.0.0.1:8080/v1 +tensorfold tui --url http://127.0.0.1:8888/v1 tensorfold tui --url https://spark.example.invalid --token-env TENSORFOLD_API_KEY tensorfold tui --demo tensorfold tui --demo --snapshot preview.svg diff --git a/docs/cortheon.md b/docs/cortheon.md index d8bee40e3e..bf6b242be1 100644 --- a/docs/cortheon.md +++ b/docs/cortheon.md @@ -100,8 +100,8 @@ tensorfold serve TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit --cor ``` ```bash -curl -s http://127.0.0.1:8080/v1/chat/completions -H 'Content-Type: application/json' \ +curl -s http://127.0.0.1:8888/v1/chat/completions -H 'Content-Type: application/json' \ -d '{"messages": [{"role": "user", "content": "What changed in the latest Zig release?"}]}' ``` -Point any OpenAI-compatible client at `http://127.0.0.1:8080/v1` as before. +Point any OpenAI-compatible client at `http://127.0.0.1:8888/v1` as before. diff --git a/docs/living-weights.md b/docs/living-weights.md index 9fd99bdb83..b7004b1e6c 100644 --- a/docs/living-weights.md +++ b/docs/living-weights.md @@ -61,12 +61,12 @@ tensorfold serve ~/models/nemotron-living --slide --slide-graph ~/models/nemotro `--slide` turns learning on. `--slide-graph` keeps the list of learned facts in a file, so it survives a restart; without it the list lives in memory, though the weights keep what they learned either way. The server listens on -`http://127.0.0.1:8080` unless you pass `--host` or `--port`. +`http://127.0.0.1:8888` unless you pass `--host` or `--port`. ### 4. Teach it a fact ```sh -curl -N http://127.0.0.1:8080/v1/slide/learn \ +curl -N http://127.0.0.1:8888/v1/slide/learn \ -H 'Content-Type: application/json' \ -d '{"text": "I like blue."}' ``` @@ -82,7 +82,7 @@ answer. A fact takes 1.5 to 3 minutes. Use the normal chat API, from any client: ```sh -curl http://127.0.0.1:8080/v1/chat/completions -H 'Content-Type: application/json' \ +curl http://127.0.0.1:8888/v1/chat/completions -H 'Content-Type: application/json' \ -d '{"messages": [{"role": "user", "content": "What colour do I like?"}]}' ``` @@ -112,7 +112,7 @@ nothing behind. An earlier fact's neuron that the new lesson would disturb is re - Several facts at once: put them in one text, as in `{"text": "I like blue. My sister is called Ana."}`. Each is learned in turn. - A file: send its contents as the text, for example - `jq -Rs '{text: .}' notes.md | curl -N http://127.0.0.1:8080/v1/slide/learn -H 'Content-Type: application/json' -d @-`. + `jq -Rs '{text: .}' notes.md | curl -N http://127.0.0.1:8888/v1/slide/learn -H 'Content-Type: application/json' -d @-`. - A web page: fetch its text first, or use omp's `/learn ` below, which does it for you. The server keeps answering chat requests while it learns: learning runs when no reply is being generated. Short, @@ -136,7 +136,7 @@ A panel shows each fact as queued, learning, learned or missed. ## See what it learned -Open `http://127.0.0.1:8080/slide` in a browser while the server runs. It draws the model in 3D and updates live: +Open `http://127.0.0.1:8888/slide` in a browser while the server runs. It draws the model in 3D and updates live: - The model as a tower of its 52 layers. Each MoE layer is a ring carrying its 64 pocket neurons, and a neuron glows in the colour of the fact that owns it. @@ -150,9 +150,9 @@ Drag to turn, scroll or pinch to zoom, and click a fact for its neurons, the que rounds. Phones get a bottom sheet. Without WebGL2, or with `/slide?flat`, the page draws a flat graph from the same data. -- `curl http://127.0.0.1:8080/v1/slide/graph` returns the same data as JSON: each fact's state, neurons (layer, index +- `curl http://127.0.0.1:8888/v1/slide/graph` returns the same data as JSON: each fact's state, neurons (layer, index and key), questions and rounds, and the model's layers, pocket and shards. -- `curl -X DELETE http://127.0.0.1:8080/v1/slide/graph` clears the graph. The weights keep what they learned. +- `curl -X DELETE http://127.0.0.1:8888/v1/slide/graph` clears the graph. The weights keep what they learned. ## Keep, move or reset what it learned diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 1e38e07605..28698ef465 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -61,7 +61,7 @@ how matrix multiplication uses a GPU. The chat fixture disables thinking. After `--name bench`, run from the repository root: ```bash -python3 tools/bench_openai.py http://127.0.0.1:8080 bench \ +python3 tools/bench_openai.py http://127.0.0.1:8888 bench \ --tokens 64 --reps 5 --temperatures 1.0,0 --output bench.json ``` diff --git a/docs/recipes/qwen3.6-moe.md b/docs/recipes/qwen3.6-moe.md index 15b756169b..28d7696b96 100644 --- a/docs/recipes/qwen3.6-moe.md +++ b/docs/recipes/qwen3.6-moe.md @@ -6,7 +6,8 @@ engine. @cshintov brought the model to Macs in [#138](https://github.com/ashhart @philip-pentatonic ([#78](https://github.com/ashhart/TensorFold/pull/78), [#84](https://github.com/ashhart/TensorFold/pull/84)) and @BHCC2025 ([#270](https://github.com/ashhart/TensorFold/pull/270)) worked on its CUDA path; see [credits](../credits.md). -The native binary has no `qwen3_5_moe` family yet, so `tensorfold pull` and `tensorfold info` refuse the checkpoint. +The native CUDA engine serves the checkpoint with its MTP layer ([docs/WINDOWS.md](../WINDOWS.md)); the Metal native +binary has no `qwen3_5_moe` family yet, so `tensorfold pull` and `tensorfold info` refuse the checkpoint there. The `qwen3_5_moe` family serves Qwen3.6-35B-A3B on Apple Silicon and on one NVIDIA GPU. Its layers are the 27B's (Gated DeltaNet and gated full attention, every fourth layer attention) with routed experts in place of the dense diff --git a/docs/recipes/qwen3.8-27b.md b/docs/recipes/qwen3.8-27b.md index d4084dcfc1..6027d7d01a 100644 --- a/docs/recipes/qwen3.8-27b.md +++ b/docs/recipes/qwen3.8-27b.md @@ -90,7 +90,7 @@ input is qualified on this checkpoint. ```bash tensorfold pull nvidia/Qwen3.8-27B-NVFP4 z-lab/Qwen3.8-27B-DFlash2 -tensorfold serve nvidia/Qwen3.8-27B-NVFP4 --host 0.0.0.0 --port 8080 +tensorfold serve nvidia/Qwen3.8-27B-NVFP4 --host 0.0.0.0 --port 8888 ``` Prompts take bf16 rows on a prompt GEMM that reads the stored bytes: each FP8 byte and each NVFP4 code times its @@ -122,7 +122,7 @@ serve the folder: ```bash python -c "from huggingface_hub import snapshot_download as d; d('turboderp/Qwen3.8-27B-exl3', revision='3.00bpw', local_dir='qwen27b-exl3-3.00bpw')" tensorfold pull z-lab/Qwen3.8-27B-DFlash2 -tensorfold serve qwen27b-exl3-3.00bpw --host 0.0.0.0 --port 8080 +tensorfold serve qwen27b-exl3-3.00bpw --host 0.0.0.0 --port 8888 ``` Verify windows use the row-invariant EXL3 linear, so drafted replies equal `"draft": false` ones; prompts use the diff --git a/docs/recipes/qwen3.8-flash-next.md b/docs/recipes/qwen3.8-flash-next.md index af65c4fccd..bed3056e40 100644 --- a/docs/recipes/qwen3.8-flash-next.md +++ b/docs/recipes/qwen3.8-flash-next.md @@ -98,7 +98,7 @@ local-inference-lab's. ship, MTP head included, on one GPU: ```bash -tensorfold serve local-inference-lab/Qwen3.8-Flash-Next-NVFP4 --host 0.0.0.0 --port 8080 +tensorfold serve local-inference-lab/Qwen3.8-Flash-Next-NVFP4 --host 0.0.0.0 --port 8888 ``` | Checkpoint (revision) | Routed experts | DeltaNet, attention, shared expert | n-gram table | @@ -169,7 +169,7 @@ Download that branch, then serve the folder: ```bash python -c "from huggingface_hub import snapshot_download as d; d('turboderp/Qwen3.8-Flash-Next-exl3', revision='3.05bpw_h5_ng5', local_dir='flashnext-exl3-3.05bpw')" python -m tensorfold.cuda.exl3.inspect flashnext-exl3-3.05bpw # bits per tensor, from the headers -tensorfold serve flashnext-exl3-3.05bpw --host 0.0.0.0 --port 8080 +tensorfold serve flashnext-exl3-3.05bpw --host 0.0.0.0 --port 8888 ``` A pack maps its n-gram table from its own file and runs on one GPU: `--tp 2` and `--ple-on-ssd` are for the MLX diff --git a/docs/speed-up-mode.md b/docs/speed-up-mode.md index 69803a8df7..f9ad574d48 100644 --- a/docs/speed-up-mode.md +++ b/docs/speed-up-mode.md @@ -119,7 +119,7 @@ Flags that change decoding, such as `--no-drafts` or `--no-thinking`, go on both ### 6. Connect a client -Point any OpenAI-compatible client at `http://RANK0:8080/v1` with the model name `flash-next`, where RANK0 is rank 0's +Point any OpenAI-compatible client at `http://RANK0:8888/v1` with the model name `flash-next`, where RANK0 is rank 0's address. Send requests only to rank 0, since rank 1 refuses its own. ## Measured speeds diff --git a/docs/vision.md b/docs/vision.md index 9536f7e7fc..486c283663 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -64,7 +64,7 @@ body = { 'max_tokens': 128, 'stream': False, } -request = Request('http://127.0.0.1:8080/v1/chat/completions', +request = Request('http://127.0.0.1:8888/v1/chat/completions', data=json.dumps(body).encode(), headers={'Content-Type': 'application/json'}) with urlopen(request) as response: diff --git a/tools/bench_concurrent.py b/tools/bench_concurrent.py index 446592072e..c1c19eb591 100644 --- a/tools/bench_concurrent.py +++ b/tools/bench_concurrent.py @@ -16,7 +16,7 @@ server with a small listen backlog resets simultaneous connections; those resets are reported, not hidden). Standard library only. - python3 tools/bench_concurrent.py http://127.0.0.1:8080 MODEL --levels 1,2,4,8 --alone --output out.json + python3 tools/bench_concurrent.py http://127.0.0.1:8888 MODEL --levels 1,2,4,8 --alone --output out.json """ import argparse diff --git a/tools/bench_openai.py b/tools/bench_openai.py index fca9bf1067..d18685df6b 100644 --- a/tools/bench_openai.py +++ b/tools/bench_openai.py @@ -4,7 +4,7 @@ CUDA engine's bench times (the first sampled token excluded). Standard library only, so it runs on a bare host. - python3 tools/bench_openai.py http://127.0.0.1:8080 MODEL --tokens 64 --reps 5 --output out.json + python3 tools/bench_openai.py http://127.0.0.1:8888 MODEL --tokens 64 --reps 5 --output out.json """ import argparse diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..210b37a8aa 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -21,6 +21,7 @@ const kernels = [_]Kernel{ .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7 .{ .name = "experts_prefill", .flags = &.{"-O3"} }, .{ .name = "experts_pack", .flags = &.{"-O3"} }, + .{ .name = "moe", .flags = &.{"-O3"} }, // cuda/moe.py: softmax top-k routing with a sigmoid-gated shared expert .{ .name = "prefill_attention", .flags = &.{ "-O3", "--fmad=false" } }, // tensorfold_prefill_attention_v1 .{ .name = "scan_rows", .flags = &.{ "-O3", "--fmad=false" } }, // nemotron_h/cuda/mamba.py .{ .name = "nemotron_ops", .flags = &.{"-O3"} }, // ours: Nemotron's serial feed, routed plan and rests @@ -32,6 +33,10 @@ const kernels = [_]Kernel{ .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) + .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: the Qwen3.5 hybrids' glue, the Metal kernels' rounding points + .{ .name = "qwen35_draft", .flags = &.{"-O3"} }, // ours: the DFlash2 drafter's glue (draft_ops.metal's rounding) + .{ .name = "qwen35_keep", .flags = &.{"-O3"} }, // ours: a verified window's kept path applied to its stream's state + .{ .name = "qwen36_mtp", .flags = &.{"-O3"} }, // ours: the Qwen3.6 MTP drafter's attention at the target's head width .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, @@ -78,7 +83,7 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return cuda; } -fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { +fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); core.addImport("tokenizer", tokenizer); @@ -90,7 +95,11 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin nemotron.addImport("lanes", lanes); nemotron.addImport("heat", heat); nemotron.addImport("nemotron_draft_ids", draft_ids); - return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat }; + const qwen35 = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3_5/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + qwen35.addImport("cuda", cuda); + qwen35.addImport("core", core); + qwen35.addImport("lanes", lanes); + return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .qwen35 = qwen35, .tokenizer = tokenizer, .heat = heat }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -122,7 +131,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu cli.addImport("core", mods.core); cli.addImport("lanes", mods.lanes); cli.addImport("nemotron", mods.nemotron); - const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron); + const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.qwen35); cli.addOptions("build_options", build_options); cli.addImport("checkpoint_cli", checkpointCli(b, target, optimize, strip, native.engines, mods.core)); b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli })); @@ -139,18 +148,23 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; + _ = nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.tokenizer, build_options, true, strip); + // tf-qwen35-cuda: Qwen3.5-2B's CUDA forward, its logits written for a reference to compare + const q35 = b.createModule(.{ .root_source_file = b.path("zig/tests/qwen35_cuda_main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); + q35.addImport("cuda", cuda); + q35.addImport("qwen35", mods.qwen35); + b.installArtifact(b.addExecutable(.{ .name = "tf-qwen35-cuda", .root_module = q35 })); } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. -fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { +fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, - .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron } }, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron }, .{ .name = "qwen35", .module = qwen35 } }, }); return .{ .api = api, .engines = mod }; } @@ -161,8 +175,8 @@ fn checkpointCli(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. } /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it. -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile { - const m = engines(b, target, optimize, cuda, lanes, nemotron); +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool, strip: ?bool) *std.Build.Step.Compile { + const m = engines(b, target, optimize, cuda, lanes, nemotron, qwen35); // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's) const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const exe = b.addExecutable(.{ .name = "tensorfold-native", .root_module = b.createModule(.{ @@ -177,6 +191,8 @@ fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b exe.root_module.strip = true; return exe; } + // before the install step, which decides then whether a .pdb comes with the executable + exe.root_module.strip = strip; const install = b.addInstallArtifact(exe, .{ .dest_dir = .{ .override = .{ .custom = "native/bin" } } }); b.step("native", "tensorfold-native with the CUDA engines into zig-out/native/bin").dependOn(&install.step); return exe; @@ -190,7 +206,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/build/check_fatbins.zig"), .target = host }) })).step); const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); - const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron); + const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron, mods.qwen35); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = host, .optimize = .debug, .link_libc = true }); @@ -210,7 +226,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); @@ -226,6 +242,8 @@ fn fatbin(b: *std.Build, nvcc: []const u8, version: std.Build.LazyPath, k: Kerne const run = b.addSystemCommand(&.{ nvcc, "-fatbin" }); run.addFileInput(version); run.addArgs(&torch_flags); + // CUDA 13's headers need MSVC's conforming preprocessor when cl.exe is the host compiler + if (b.graph.host.result.os.tag == .windows) run.addArg("-Xcompiler=/Zc:preprocessor"); run.addArgs(k.flags); const a = if (k.arch_specific) "a" else ""; var it = std.mem.tokenizeScalar(u8, sms, ','); @@ -246,7 +264,7 @@ pub fn distServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *s }; const cuda = runtime(b, target, .fast, if (prebuilt != null) &images else &.{}); const mods = family(b, target, .fast, cuda, draft_ids); - return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.tokenizer, build_options, false, null); } /// SASS targets every release fatbin holds: Ampere (sm_86), RTX 50 and RTX PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/build/sys.zig b/zig/build/sys.zig new file mode 100644 index 0000000000..1c1309af78 --- /dev/null +++ b/zig/build/sys.zig @@ -0,0 +1,40 @@ +//! `sys` (zig/src/platform/sys.zig) for every module the build compiles, by one walk of the step graph once declared. + +const std = @import("std"); + +pub fn wire(b: *std.Build) void { + var seen_steps: std.AutoHashMapUnmanaged(*std.Build.Step, void) = .empty; + var seen_mods: std.AutoHashMapUnmanaged(*std.Build.Module, void) = .empty; + var per_target: std.StringHashMapUnmanaged(*std.Build.Module) = .empty; + var w: Walk = .{ .b = b, .seen_steps = &seen_steps, .seen_mods = &seen_mods, .per_target = &per_target }; + for (b.top_level_steps.values()) |top| w.step(&top.step); +} + +const Walk = struct { + b: *std.Build, + seen_steps: *std.AutoHashMapUnmanaged(*std.Build.Step, void), + seen_mods: *std.AutoHashMapUnmanaged(*std.Build.Module, void), + per_target: *std.StringHashMapUnmanaged(*std.Build.Module), + + fn step(w: *Walk, s: *std.Build.Step) void { + if ((w.seen_steps.getOrPut(w.b.allocator, s) catch @panic("OOM")).found_existing) return; + if (s.cast(std.Build.Step.Compile)) |compile| w.module(compile.root_module, compile.root_module.resolved_target); + for (s.dependencies.items) |dep| w.step(dep); + } + + fn module(w: *Walk, m: *std.Build.Module, inherited: ?std.Build.ResolvedTarget) void { + if ((w.seen_mods.getOrPut(w.b.allocator, m) catch @panic("OOM")).found_existing) return; + const target = m.resolved_target orelse inherited orelse w.b.graph.host; + // only TensorFold's own sources see it; the sys module itself and generated modules do not + if (m.root_source_file) |root| if (root == .src_path and !m.import_table.contains("sys") and !std.mem.endsWith(u8, root.src_path.sub_path, "platform/sys.zig")) + m.addImport("sys", w.sys(target)); + for (m.import_table.values()) |dep| w.module(dep, target); + } + + fn sys(w: *Walk, target: std.Build.ResolvedTarget) *std.Build.Module { + const key = target.result.zigTriple(w.b.allocator) catch @panic("OOM"); + const slot = w.per_target.getOrPut(w.b.allocator, key) catch @panic("OOM"); + if (!slot.found_existing) slot.value_ptr.* = w.b.createModule(.{ .root_source_file = w.b.path("zig/src/platform/sys.zig"), .target = target, .link_libc = true }); + return slot.value_ptr.*; + } +}; diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index edd83834c1..3969b15cee 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -276,7 +276,11 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const } // namespace tf_experts // Decode form for groups of 64: up with relu^2 (epilogue 1), down to fp32 (epilogue 0). -#define TF_EXPERT(EPI) template __global__ void tf_experts::expert_kernel<64, 1, EPI, 2, 4>(const __nv_bfloat16*, \ - int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); +#define TF_EXPERT(EPI) template __global__ void tf_experts::expert_kernel<64, 1, EPI, 2, 4>(const __nv_bfloat16* __restrict__, \ + int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); TF_EXPERT(1) TF_EXPERT(0) + +// Pair form: gate and up blocks interleaved per group (M = 2), SwiGLU in the epilogue (2) as fp8/nvfp4's up does. +#define TF_EXPERT_PAIR(EPI) template __global__ void tf_experts::expert_kernel<64, 2, EPI, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); +TF_EXPERT_PAIR(2) diff --git a/zig/kernels/cuda/experts_pack.cu b/zig/kernels/cuda/experts_pack.cu index 6aa7413793..0ca3d77fda 100644 --- a/zig/kernels/cuda/experts_pack.cu +++ b/zig/kernels/cuda/experts_pack.cu @@ -16,10 +16,11 @@ __device__ __forceinline__ uint32_t shuffle(uint32_t w) { template // group size / 32 __global__ void __launch_bounds__(32 * 4 * H + 32) pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ scales, - const uint16_t* __restrict__ biases, uint32_t* __restrict__ out, int n, int k8, int kg, int nb) { + const uint16_t* __restrict__ biases, uint32_t* __restrict__ out, int n, int k8, int kg, int nb, int m, int mi) { constexpr int WORDS = 32 * 4 * H; // a block's weight words: NTW * H words a lane const int g = blockIdx.x, b = blockIdx.y, e = blockIdx.z, i = threadIdx.x; - uint32_t* dst = out + ((static_cast(e) * nb + b) * kg + g) * (WORDS + 32); + // block (e, b, g) of m interleaved tables, this one the mi-th: [E][nb][kg][m][WORDS + 32] + uint32_t* dst = out + ((((static_cast(e) * nb + b) * kg + g) * m) + mi) * (WORDS + 32); if (i < WORDS) { const int row = i / (4 * H), kk = i % (4 * H); // row = t * 8 + r; kk = q * H + j const int t = row / 8, r = row % 8, q = kk / H, j = kk % H; @@ -38,5 +39,5 @@ pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ sca } // namespace tf_experts_pack // Groups of 64 inputs. -template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t*, const uint16_t*, const uint16_t*, - uint32_t*, int, int, int, int); +template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t* __restrict__, const uint16_t* __restrict__, const uint16_t* __restrict__, + uint32_t* __restrict__, int, int, int, int, int, int); diff --git a/zig/kernels/cuda/experts_prefill.cu b/zig/kernels/cuda/experts_prefill.cu index c1c7fbbac9..f56c4704a8 100644 --- a/zig/kernels/cuda/experts_prefill.cu +++ b/zig/kernels/cuda/experts_prefill.cu @@ -152,6 +152,9 @@ __global__ void __launch_bounds__(WM * WN * 32) // Prefill form for groups of 64: up with relu^2 (epilogue 1), down to bf16 (epilogue 3). #define TF_PREFILL(EPI) template __global__ void tf_experts_prefill::prefill_kernel<64, 1, EPI, 2, 2, 4>( \ - const __nv_bfloat16*, int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); + const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); TF_PREFILL(1) TF_PREFILL(3) + +// Pair form for prompt rows: gate and up interleaved (M = 2) under SwiGLU (epilogue 2). +template __global__ void tf_experts_prefill::prefill_kernel<64, 2, 2, 2, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); diff --git a/zig/kernels/cuda/gdn.cu b/zig/kernels/cuda/gdn.cu index 574fb92de5..2c96e1a3ad 100644 --- a/zig/kernels/cuda/gdn.cu +++ b/zig/kernels/cuda/gdn.cu @@ -305,14 +305,14 @@ __global__ void __launch_bounds__(32 * WARPS) replay_kernel( } // namespace tf_gdn // The instantiations the Python wrappers launch (gdn_replay_cuda, dispatch_tree for bf16 keys). -#define TF_REPLAY(QK) template __global__ void tf_gdn::replay_kernel(const long long*, int, const int*, int, \ - const int*, int, float*, int, int, int); +#define TF_REPLAY(QK) template __global__ void tf_gdn::replay_kernel(const long long* __restrict__, int, const int* __restrict__, int, \ + const int* __restrict__, int, float* __restrict__, int, int, int); TF_REPLAY(__nv_bfloat16) TF_REPLAY(float) #define TF_TREE(S, R, W, C) template __global__ void tf_gdn::tree_kernel<__nv_bfloat16, S, R, W, C>( \ - const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*, const float*, const float*, const float*, \ - const long long*, const int*, const int*, int, __nv_bfloat16*, int, int, int, tf_gdn::Pending<__nv_bfloat16>, \ - float*, const long long*, bool); + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const float* __restrict__, const float* __restrict__, const float* __restrict__, \ + const long long* __restrict__, const int* __restrict__, const int* __restrict__, int, __nv_bfloat16* __restrict__, int, int, int, tf_gdn::Pending<__nv_bfloat16>, \ + float*, const long long* __restrict__, bool); TF_TREE(0, 8, 4, true) TF_TREE(1, 8, 4, false) TF_TREE(2, 8, 2, false) diff --git a/zig/kernels/cuda/lane_gemv.cu b/zig/kernels/cuda/lane_gemv.cu index 94a78a5ae8..af6f0990ac 100644 --- a/zig/kernels/cuda/lane_gemv.cu +++ b/zig/kernels/cuda/lane_gemv.cu @@ -266,7 +266,7 @@ __device__ __forceinline__ void gemv_split(const __nv_bfloat16* __restrict__ x, } // namespace tf_lane_gemv -template __global__ void tf_lane_gemv::gemv_kernel<64, 64, 8>(const __nv_bfloat16*, const float*, +template __global__ void tf_lane_gemv::gemv_kernel<64, 64, 8>(const __nv_bfloat16* __restrict__, const float* __restrict__, const __grid_constant__ tf_lane_gemv::Part, int, int, int); extern "C" __global__ void __launch_bounds__(128) tf_lane_gemv_split(const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, diff --git a/zig/kernels/cuda/moe.cu b/zig/kernels/cuda/moe.cu new file mode 100644 index 0000000000..2adb06a220 --- /dev/null +++ b/zig/kernels/cuda/moe.cu @@ -0,0 +1,113 @@ +// Softmax top-k routing with a sigmoid-gated shared expert on fp32 logits, as tensorfold/cuda/moe.py routes. + +#include +#include +#include + +namespace { + +__device__ __forceinline__ float bf(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float rbf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +constexpr int TILE = 16; // rows and experts a block covers +constexpr int KC = 256; // inputs a shared-memory chunk holds + +} // namespace + +// OUT[r][e] = fp32 x[r] . w[e] over D in index order (one fma chain a pair); grid (rows / 16, E / 16), block 256. +extern "C" __global__ void __launch_bounds__(256) tf_moe_router(const __nv_bfloat16* __restrict__ X, int x_stride, + const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, + int R, int D, int E) { + __shared__ __nv_bfloat16 xs[TILE][KC + 2]; + __shared__ __nv_bfloat16 ws[TILE][KC + 2]; + const int row0 = blockIdx.x * TILE, e0 = blockIdx.y * TILE; + const int rr = threadIdx.x / TILE, ee = threadIdx.x % TILE; + float acc = 0.0f; + for (int k0 = 0; k0 < D; k0 += KC) { + __syncthreads(); + for (int c = threadIdx.x; c < TILE * KC; c += 256) { + const int r = c / KC, k = c % KC; + const bool kin = k0 + k < D; + xs[r][k] = kin && row0 + r < R ? X[static_cast(row0 + r) * x_stride + k0 + k] : __float2bfloat16_rn(0.0f); + ws[r][k] = kin && e0 + r < E ? W[static_cast(e0 + r) * D + k0 + k] : __float2bfloat16_rn(0.0f); + } + __syncthreads(); + const int kn = min(KC, D - k0); + for (int k = 0; k < kn; ++k) acc = __fmaf_rn(bf(xs[rr][k]), bf(ws[ee][k]), acc); + } + if (row0 + rr < R && e0 + ee < E) OUT[static_cast(row0 + rr) * E + e0 + ee] = acc; +} + +namespace { + +constexpr int MAXC = 8; // experts a lane holds: up to 256 routed experts + +// The larger of two (value, index) picks: a higher value, or the lower index on a tie (values are never NaN). +__device__ __forceinline__ bool beats(float v, int i, float bv, int bi) { return v > bv || (v == bv && i < bi); } + +} // namespace + +// Row r by a warp: TOPK experts by logit (ties to the lower id), weights bf16(exp(l - top) / sum), the shared slot. +extern "C" __global__ void __launch_bounds__(32) tf_moe_topk(const float* __restrict__ L, int* __restrict__ PICK, + float* __restrict__ WT, int NE, int TOPK) { + const int64_t r = blockIdx.x; + const int lane = threadIdx.x; + const int NL = NE + 1, SLOTS = TOPK + 1; // L[r][NE], the shared gate's logit, weighs slot TOPK by bf16(sigmoid(bf16(it))) + float v[MAXC]; +#pragma unroll + for (int c = 0; c < MAXC; ++c) { + const int e = lane + 32 * c; + v[c] = e < NE ? L[r * NL + e] : -INFINITY; + } + float top = -INFINITY; +#pragma unroll + for (int c = 0; c < MAXC; ++c) top = fmaxf(top, v[c]); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) top = fmaxf(top, __shfl_xor_sync(0xffffffffu, top, off)); + float total = 0.0f; + float ex[MAXC]; + int id[MAXC]; + for (int k = 0; k < TOPK; ++k) { + float bv = -INFINITY; + int bi = 0x7fffffff; +#pragma unroll + for (int c = 0; c < MAXC; ++c) + if (lane + 32 * c < NE && beats(v[c], lane + 32 * c, bv, bi)) bv = v[c], bi = lane + 32 * c; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (beats(ov, oi, bv, bi)) bv = ov, bi = oi; + } + const float e = expf(bv - top); + total += e; + if (lane == k % 32) ex[k / 32] = e, id[k / 32] = bi; // slot k lives in lane k % 32 (TOPK <= 32) +#pragma unroll + for (int c = 0; c < MAXC; ++c) + if (lane + 32 * c == bi) v[c] = -INFINITY; + } + if (lane < TOPK) { + PICK[r * SLOTS + lane] = id[0]; + WT[r * SLOTS + lane] = rbf(ex[0] / total); + } + if (lane == 0) { + const float sg = rbf(L[r * NL + NE]); + PICK[r * SLOTS + TOPK] = NE; + WT[r * SLOTS + TOPK] = rbf(1.0f / (1.0f + expf(-sg))); + } +} + +// OUT[r][d] = bf16(sum over slots in order of fma(w[r][s], y[r][s][d])); y fp32 (decode) or bf16 (prompt rows). +extern "C" __global__ void __launch_bounds__(256) tf_moe_combine(const void* __restrict__ Y, int y_f32, const float* __restrict__ WT, + __nv_bfloat16* __restrict__ OUT, int S, int D) { + const int64_t r = blockIdx.x; + for (int d = blockIdx.y * 256 + threadIdx.x; d < D; d += gridDim.y * 256) { + float acc = 0.0f; + for (int s = 0; s < S; ++s) { + const int64_t at = (r * S + s) * D + d; + const float y = y_f32 ? static_cast(Y)[at] : bf(static_cast(Y)[at]); + acc = __fmaf_rn(WT[r * S + s], y, acc); + } + OUT[r * D + d] = __float2bfloat16_rn(acc); + } +} diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 122f404495..b2521ae4a6 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -241,5 +241,5 @@ pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric } // namespace tf_prefill_attention // Head dim 128, eight warps, eight query heads a block, eight staging slots (Nemotron's 16 heads a KV head). -template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16*, - const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, int, float); +template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, float); diff --git a/zig/kernels/cuda/qmm_group.cu b/zig/kernels/cuda/qmm_group.cu index f3f65e184f..f9b2f9f2a7 100644 --- a/zig/kernels/cuda/qmm_group.cu +++ b/zig/kernels/cuda/qmm_group.cu @@ -338,4 +338,4 @@ __global__ void __launch_bounds__(WM * WN * 32) group_kernel( // The instantiations Nemotron's windows launch on sm_121 (tile 2: rows <= 16, bf16 out). template __global__ void tf_qmm_group::group_kernel<64, 16, 64, 1, 4, 8, false, false, false, false>( - const __nv_bfloat16*, const float*, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int); + const __nv_bfloat16* __restrict__, const float* __restrict__, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int); diff --git a/zig/kernels/cuda/qmm_prefill.cu b/zig/kernels/cuda/qmm_prefill.cu index 1610b22704..3c61031a25 100644 --- a/zig/kernels/cuda/qmm_prefill.cu +++ b/zig/kernels/cuda/qmm_prefill.cu @@ -149,5 +149,5 @@ __global__ void __launch_bounds__(WM * WN * 32) prefill_kernel( // The instantiation prefill_matmul launches (tile 0: 128x128 on 2x4 warps, 3 stages), bf16 out. template __global__ void tf_qmm_prefill::prefill_kernel<64, 128, 128, 2, 4, 3, false>( - const __nv_bfloat16*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, int, int, int, + const __nv_bfloat16* __restrict__, const uint32_t* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, void* __restrict__, int, int, int, int, int, int); diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu new file mode 100644 index 0000000000..9217a19fa3 --- /dev/null +++ b/zig/kernels/cuda/qwen35.cu @@ -0,0 +1,510 @@ +// Qwen3.5's hybrids (2B, 27B) on CUDA: the Metal kernels' arithmetic and rounding; a segment is one stream's window. + +#include "qwen35.cuh" + +using namespace tf_q35; + +// P (device counts, one graph for every shape): chunks, width, rows, windows, recorded; rowpar/rowanc give the tree. +__device__ __forceinline__ bool past_rows(const int* P, int m) { return P != nullptr && m >= P[2]; } +__device__ __forceinline__ bool past_windows(const int* P, int s) { return P != nullptr && s >= P[3]; } +__device__ __forceinline__ bool recorded(const int* P) { return P != nullptr && P[4] != 0; } + +// MLX affine 4-bit token rows: out = bf16(s * q + b), D wide. +extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, + const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B, + __nv_bfloat16* __restrict__ out, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const size_t id = ids[blockIdx.x]; + for (int b = threadIdx.x; b < D / 2; b += blockDim.x) { + const uint8_t v = W[id * (D / 2) + b]; + const size_t g = id * (D / 64) + (2 * b) / 64; + const float s = bf(S[g]), z = bf(B[g]); + out[(size_t)blockIdx.x * D + 2 * b] = __float2bfloat16_rn(s * (float)(v & 15) + z); + out[(size_t)blockIdx.x * D + 2 * b + 1] = __float2bfloat16_rn(s * (float)(v >> 4) + z); + } +} + +// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)); D a multiple of 256 up to 8192; a 256-thread block a row. +extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, + const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D, + const int* __restrict__ P) { + constexpr int MAXE = 32; + __shared__ float red[8]; + if (past_rows(P, blockIdx.x)) return; + const size_t row = blockIdx.x; + const int t = threadIdx.x; + const int E = D / 256; + float hv[MAXE]; + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < MAXE; i++) { + if (i >= E) break; + const size_t at = row * D + t * E + i; + float h = bf(H[at]); + if (R != nullptr) { + const __nv_bfloat16 hb = __float2bfloat16_rn(h + bf(R[at])); + H[at] = hb; + h = bf(hb); + } + hv[i] = h; + ss = fmaf(h, h, ss); + } + ss = warp_sum(ss); + if ((t & 31) == 0) red[t >> 5] = ss; + __syncthreads(); + float total = 0.0f; + for (int i = 0; i < 8; i++) total += red[i]; + const float inv = rsqrtf(total / (float)D + eps); +#pragma unroll + for (int i = 0; i < MAXE; i++) { + if (i >= E) break; + const int d = t * E + i; + X[row * D + d] = __float2bfloat16_rn(bf(W[d]) * (hv[i] * inv)); + } +} + +// The 64-group sums of bf16 rows D wide (the decode projections' input sums). Grid (rows), a thread a group. +extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, float* __restrict__ XS, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const size_t row = blockIdx.x; + for (int g = threadIdx.x; g < D / 64; g += blockDim.x) { + const __nv_bfloat16* x = X + row * D + g * 64; + float s = 0.0f; + for (int j = 0; j < 64; j++) s += bf(x[j]); + XS[row * (D / 64) + g] = s; + } +} + +// Conv, silu, q/k norms, the linear heads' decay and write: a warp a (row, head) over q, k, v; taps from the parents. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const int* __restrict__ rowpar, int layer, + const __nv_bfloat16* __restrict__ CW, + const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, + const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, + __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ K, + __nv_bfloat16* __restrict__ V, float* __restrict__ G, __nv_bfloat16* __restrict__ BETA, + const int* __restrict__ P) { + const int lane = threadIdx.x; + const int head = blockIdx.y; + const int w = blockIdx.z; + if (past_rows(P, w)) return; + const int NK = geo.nk, NV = geo.nv, C = geo.conv(); + const Seg sg = segs[rowseg[w]]; + const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * geo.conv_rows(); + const bool isq = head < NK, isk = !isq && head < 2 * NK; + const int c0 = head * DK; // q, k and v channels are consecutive heads of DK + int src[TAPS]; // tap t reads the row TAPS - 1 - t tokens back: a forward row, or -1 - (tail row) +#pragma unroll + for (int tap = 0; tap < TAPS; tap++) src[tap] = back(rowpar, w, TAPS - 1 - tap); + constexpr int PER = DK / 32; + float vals[PER]; +#pragma unroll + for (int j = 0; j < PER; j++) { + const int c = c0 + lane * PER + j; + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; tap++) { + const int at = src[tap]; + const float xv = at < 0 ? bf(CS[(size_t)(-1 - at) * C + c]) : bf(QKV[(size_t)at * C + c]); + acc += bf(CW[c * TAPS + tap]) * xv; + } + const float conv = rbf(acc); + const float sig = rbf(1.0f / (1.0f + expf(-conv))); + vals[j] = rbf(conv * sig); + } + if (isq || isk) { + float ss = 0.0f; +#pragma unroll + for (int j = 0; j < PER; j++) ss += vals[j] * vals[j]; + ss = warp_sum(ss); + const float inv = rsqrtf(ss / (float)DK + (1e-6f / (float)DK)); + const float scale = isq ? rbf(1.0f / (float)DK) : rbf(rsqrtf((float)DK)); + __nv_bfloat16* out = isq ? Q + ((size_t)w * NK + head) * DK : K + ((size_t)w * NK + head - NK) * DK; +#pragma unroll + for (int j = 0; j < PER; j++) out[lane * PER + j] = __float2bfloat16_rn(scale * rbf(vals[j] * inv)); + } else { + const int hv = head - 2 * NK; +#pragma unroll + for (int j = 0; j < PER; j++) V[((size_t)w * NV + hv) * DK + lane * PER + j] = __float2bfloat16_rn(vals[j]); + if (lane == 0) { + const float s = rbf(bf(Ain[w * NV + hv]) + bf(DT[hv])); + const float sp = rbf(fmaxf(s, 0.0f) + logf(1.0f + expf(-fabsf(s)))); + G[w * NV + hv] = expf(-expf(ALOG[hv]) * sp); + BETA[w * NV + hv] = __float2bfloat16_rn(1.0f / (1.0f + expf(-bf(Bin[w * NV + hv])))); + } + } +} + +// A segment's conv tail: a recorded window logs each row's for keep and leaves the stream's; others update it in place. +extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, + const int* __restrict__ rowpar, int layer, const int* __restrict__ P) { + const int C = geo.conv(); + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= C || past_windows(P, blockIdx.y)) return; + const Seg sg = segs[blockIdx.y]; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * geo.conv_rows(); + if (recorded(P)) { + for (int j = 0; j < sg.rows; j++) { + __nv_bfloat16* out = log_row(geo, sg, layer, j).conv; + for (int r = 0; r < TAPS - 1; r++) { + const int at = back(rowpar, sg.row0 + j, TAPS - 2 - r); + out[(size_t)r * C + c] = at < 0 ? CS[(size_t)(-1 - at) * C + c] : QKV[(size_t)at * C + c]; + } + } + if (sg.rows != 1) return; + } + // ascending: slot r takes [r + rows], which is either a later state slot (not yet written) or a row + for (int r = 0; r < TAPS - 1; r++) { + const int at = r + sg.rows; + CS[(size_t)r * C + c] = at < TAPS - 1 ? CS[(size_t)at * C + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]; + } +} + +// Gated delta rule, each row's state its parent's: a chain keeps it in registers, a branch rebuilds it from the log. +extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, + const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, + const int* __restrict__ rowpar, int layer, __nv_bfloat16* __restrict__ y, + const int* __restrict__ P) { + const int NK = geo.nk, NV = geo.nv; + if (past_windows(P, blockIdx.z / NV)) return; + const int hv = blockIdx.z % NV; + const int hk = hv / (NV / NK); + const Seg sg = segs[blockIdx.z / NV]; + const int dk = threadIdx.x; + const int dv = blockIdx.y * 4 + threadIdx.y; + constexpr int PT = DK / 32; + const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; + const bool record = recorded(P); + float state[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + int prev = -1; // the row the registers' state follows + for (int r = 0; r < sg.rows; r++) { + const int row = sg.row0 + r; + const int parent = rowpar[row]; + if (parent != prev) { + // a branch: its parent's state from the stream's and the ancestors' logged updates, oldest first + int anc[MAXW]; + int n = 0; + for (int a = parent; a >= 0; a = rowpar[a]) anc[n++] = a; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int j = n - 1; j >= 0; j--) { + const int a = anc[j]; + const float ga = g[a * NV + hv], da = log_row(geo, sg, layer, a - sg.row0).delta[hv * DK + dv]; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * ga; + state[i] = fmaf(bf(k[((size_t)a * NK + hk) * DK + PT * dk + i]), da, state[i]); + } + } + } + prev = row; + float kc[PT], qc[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) { + kc[i] = bf(k[((size_t)row * NK + hk) * DK + PT * dk + i]); + qc[i] = bf(q[((size_t)row * NK + hk) * DK + PT * dk + i]); + } + const float vc = bf(v[((size_t)row * NV + hv) * DK + dv]); + const float gc = g[row * NV + hv]; + const float bc = bf(beta[row * NV + hv]); + float kv_mem = 0.0f; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * gc; + kv_mem += state[i] * kc[i]; + } + kv_mem = warp_sum(kv_mem); + const float delta = (vc - kv_mem) * bc; + float out = 0.0f; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = fmaf(kc[i], delta, state[i]); + out += state[i] * qc[i]; + } + out = warp_sum(out); + if (dk == 0) y[((size_t)row * NV + hv) * DK + dv] = __float2bfloat16_rn(out); + if (record) { + const LogRow l = log_row(geo, sg, layer, r); + if (dk == 0) { + l.delta[hv * DK + dv] = delta; + if (dv == 0) l.g[hv] = gc; + } + if (dv == 0 && hv % (NV / NK) == 0) +#pragma unroll + for (int i = 0; i < PT; i++) l.k[hk * DK + PT * dk + i] = k[((size_t)row * NK + hk) * DK + PT * dk + i]; + } + } + if (record && sg.rows != 1) return; +#pragma unroll + for (int i = 0; i < PT; i++) st[i] = state[i]; +} + +// out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear value head; rows nv * DK wide. Grid (1, nv, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, + const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT, + const int* __restrict__ P) { + const int lane = threadIdx.x, hv = blockIdx.y; + if (past_rows(P, blockIdx.z)) return; + const size_t m = blockIdx.z; + const size_t W = (size_t)geo.nv * DK; + constexpr int PER = DK / 32; + float yv[PER], ss = 0.0f; +#pragma unroll + for (int j = 0; j < PER; j++) { + yv[j] = bf(Y[m * W + hv * DK + lane * PER + j]); + ss += yv[j] * yv[j]; + } + ss = warp_sum(ss); + const float inv = rsqrtf(ss / (float)DK + eps); +#pragma unroll + for (int j = 0; j < PER; j++) { + const int d = lane * PER + j; + const float x = rbf(bf(NW[d]) * (yv[j] * inv)); + const float zf = bf(Z[m * W + hv * DK + d]); + OUT[m * W + hv * DK + d] = __float2bfloat16_rn(zf / (1.0f + expf(-zf)) * x); + } +} + +// Per-head RMSNorm of attention rows, out packed (rows, heads, HD); grid (rows * heads), 64 threads of 4 values. +extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps, + const int* __restrict__ P) { + __shared__ float sums[2]; + const int head = blockIdx.x % heads; + if (past_rows(P, blockIdx.x / heads)) return; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + head * hstride; + float v[4], acc = 0.0f; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = bf(x[4 * t + i]); + acc += v[i] * v[i]; + } + acc = warp_sum(acc); + if ((t & 31) == 0) sums[t >> 5] = acc; + __syncthreads(); + const float inv = rsqrtf((sums[0] + sums[1]) / (float)HD + eps); +#pragma unroll + for (int i = 0; i < 4; i++) { + const int d = 4 * t + i; + OUT[(m * heads + head) * HD + d] = __float2bfloat16_rn(bf(W[d]) * (v[i] * inv)); + } +} + +// Text RoPE on the first 64 of HD: pairs (i, i + 32) turned by position * theta^(-i / 32). +__device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, int position) { + if (d >= 64) return x[d]; + const int i = d % 32; + const float angle = (float)position * powf(10000000.0f, -(float)i / 32.0f); + float sn, cs; + sincosf(angle, &sn, &cs); + const float a = bf(x[i]), b = bf(x[i + 32]); + return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); +} + +// A row's ancestors-and-self: its tree window's row bits, or with the top bit set a chain row's index in its window. +__device__ __forceinline__ int rows_seen(unsigned anc) { return (anc >> 31) ? (int)(anc & 0x7fffffffu) + 1 : __popc(anc); } + +// A row's chunk width from its own logical keys alone, as its one-row step's: at most MAXCH chunks, 128-key multiples. +__device__ __forceinline__ int chunk_keys(int keys) { return max(128, (keys + MAXCH * 128 - 1) / (MAXCH * 128) * 128); } + +// The window row holding a row's nth seen window key. +__device__ __forceinline__ int seen_row(unsigned anc, int n) { + if (anc >> 31) return n; + for (int j = 0; j < n; j++) anc &= anc - 1; + return __ffs(anc) - 1; +} + +// Queries rotated in place, keys rotated and values copied into the row's cache slot; grid (rows), HD threads. +extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, + const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, + const int* __restrict__ P) { + const int QH = geo.qh, KVH = geo.kvh; + const int d = threadIdx.x; + if (past_rows(P, blockIdx.x)) return; + const size_t m = blockIdx.x; + const Seg sg = segs[rowseg[m]]; + const int pos = sg.len + rows_seen(rowanc[m]) - 1; + const int slot = sg.len + (int)m - sg.row0; + __shared__ __nv_bfloat16 qs[MAXQ * HD]; + for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; + __syncthreads(); + for (int h = 0; h < QH; h++) Qn[(m * QH + h) * HD + d] = rotate(qs + h * HD, d, pos); + __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; + __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + for (int h = 0; h < KVH; h++) { + const size_t dst = ((size_t)h * sg.cap + slot) * HD + d; + keys[dst] = rotate(Kn + (m * KVH + h) * HD, d, pos); + values[dst] = Vr[(m * KVH + h) * HD + d]; + } +} + +// Causal attention in chunks of chunk_keys(the row's keys), a warp a query head; partials by (row, head, chunk). +extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, + const int* __restrict__ P, float scale, + float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { + constexpr int DPL = HD / 32; + const int QH = geo.qh, KVH = geo.kvh; + const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; + const int c = blockIdx.x, h = blockIdx.y; + if (c >= MAXCH || past_rows(P, blockIdx.z)) return; + const size_t m = blockIdx.z; + const int qh = h * (QH / KVH) + g; + const Seg sg = segs[rowseg[m]]; + const unsigned anc = rowanc[m]; + // chunks cut at logical positions (a window key's: the stream's length plus its ancestors), as a step's are + const int keys = sg.len + rows_seen(anc), ck = chunk_keys(keys); + const int k0 = c * ck, k1 = min(k0 + ck, keys); + const size_t slot = (m * QH + qh) * MAXCH + c; + if (k0 >= k1) { + if (lane == 0) { + PM[slot] = -INFINITY; + PL[slot] = 0.0f; + } + return; + } + const __nv_bfloat16* K = reinterpret_cast(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; + const __nv_bfloat16* V = K + (size_t)KVH * sg.cap * HD; + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[(m * QH + qh) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + for (int t = k0; t < k1; t++) { + const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); + const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const __nv_bfloat16* kb = reinterpret_cast(&kw); + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(kb[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const uint4 vw = *reinterpret_cast(V + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const __nv_bfloat16* vb = reinterpret_cast(&vw); +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); + mx = mn; + } + if (lane == 0) { + PM[slot] = mx; + PL[slot] = l; + } +#pragma unroll + for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; +} + +// A row's chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, const float* __restrict__ PM, const float* __restrict__ PL, + const float* __restrict__ PO, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, + const int* __restrict__ P, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int QH = geo.qh; + const int lane = threadIdx.x, qh = blockIdx.x; + if (past_rows(P, blockIdx.y)) return; + const size_t m = blockIdx.y; + const int keys = segs[rowseg[m]].len + rows_seen(rowanc[m]); + const int ck = chunk_keys(keys), nch = (keys + ck - 1) / ck; + const size_t base = (m * QH + qh) * MAXCH; + float mx = -INFINITY; + for (int c = 0; c < nch; c++) + if (PL[base + c] > 0.0f) mx = fmaxf(mx, PM[base + c]); + float lsum = 0.0f, acc[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) acc[i] = 0.0f; + for (int c = 0; c < nch; c++) { + if (!(PL[base + c] > 0.0f)) continue; + const float e = __expf(PM[base + c] - mx); + lsum = fmaf(PL[base + c], e, lsum); +#pragma unroll + for (int i = 0; i < DPL; i++) acc[i] = fmaf(PO[(base + c) * HD + lane * DPL + i], e, acc[i]); + } +#pragma unroll + for (int i = 0; i < DPL; i++) OUT[(m * QH + qh) * HD + lane * DPL + i] = __float2bfloat16_rn(acc[i] / lsum); +} + +// out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows; rows qh * HD wide. +extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, + __nv_bfloat16* __restrict__ OUT, int rows, const int* __restrict__ P) { + const size_t W = (size_t)geo.qh * HD; + if (P != nullptr) rows = P[2]; + const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= (size_t)rows * W) return; + const size_t r = i / W; + const int d = (int)(i % W), head = d / HD; + const float gate = bf(QRAW[r * 2 * W + head * 2 * HD + HD + d % HD]); + const float s = rbf(1.0f / (1.0f + expf(-gate))); + OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); +} + +// A tapped layer's residual bf16(h + r) into the segment's tap ring at the row's position; a 256-thread block a row. +extern "C" __global__ void __launch_bounds__(256) tf_q35_tap(const __nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, + const Seg* __restrict__ segs, const int* __restrict__ rowseg, int tap, int D, + const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const size_t m = blockIdx.x; + const Seg sg = segs[rowseg[m]]; + if (sg.tap == 0) return; + const long long pos = sg.len + (long long)m - sg.row0; + __nv_bfloat16* out = reinterpret_cast<__nv_bfloat16*>(sg.tap) + ((size_t)(pos % sg.tap_rows) * sg.taps + tap) * D; + for (int i = threadIdx.x; i < D; i += blockDim.x) out[i] = __float2bfloat16_rn(bf(H[m * D + i]) + bf(R[m * D + i])); +} + +// out = bf16(silu(gate) * up) over `total` values. +extern "C" __global__ void tf_q35_mlp_act(const __nv_bfloat16* __restrict__ GATE, const __nv_bfloat16* __restrict__ UP, + __nv_bfloat16* __restrict__ OUT, long long total, const int* __restrict__ P, int D) { + if (P != nullptr) total = (long long)P[2] * D; + const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + const float gf = bf(GATE[i]); + OUT[i] = __float2bfloat16_rn(gf / (1.0f + expf(-gf)) * bf(UP[i])); +} + +// Each row's greedy token: the largest logit, a tie to the lower id (the host draw's rule). Grid (rows), 1024 threads. +extern "C" __global__ void __launch_bounds__(1024) tf_q35_argmax(const __nv_bfloat16* __restrict__ L, int V, unsigned* __restrict__ TOK) { + __shared__ float bv[32]; + __shared__ int bi[32]; + const __nv_bfloat16* x = L + (size_t)blockIdx.x * V; + float best = -INFINITY; + int id = 0x7fffffff; + for (int i = threadIdx.x; i < V; i += blockDim.x) { + const float v = bf(x[i]); + if (v > best) { + best = v; + id = i; + } + } +#pragma unroll + for (int m = 16; m; m >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, best, m); + const int oi = __shfl_xor_sync(0xffffffffu, id, m); + if (ov > best || (ov == best && oi < id)) { + best = ov; + id = oi; + } + } + if ((threadIdx.x & 31) == 0) { + bv[threadIdx.x >> 5] = best; + bi[threadIdx.x >> 5] = id; + } + __syncthreads(); + if (threadIdx.x == 0) { + for (int w = 1; w < 32; w++) + if (bv[w] > best || (bv[w] == best && bi[w] < id)) { + best = bv[w]; + id = bi[w]; + } + TOK[blockIdx.x] = (unsigned)id; + } +} diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh new file mode 100644 index 0000000000..5df37e308e --- /dev/null +++ b/zig/kernels/cuda/qwen35.cuh @@ -0,0 +1,79 @@ +// Shared by qwen35.cu and qwen35_keep.cu: the widths (Geo), a stream's window (Seg), a row's log record, parent walks. + +#pragma once +#include +#include + +namespace tf_q35 { + +constexpr int DK = 128; // linear-attention key and value head width +constexpr int TAPS = 4; // causal convolution taps +constexpr int HD = 256; // attention head width +constexpr int MAXQ = 32; // attention query heads at most +constexpr int MAXW = 16; // a window's rows at most (the log's rows a stream) +constexpr int MAXCH = 64; // attention chunks a row at most (cuda_engine.zig's max_chunks) + +// The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads, linear layers. +struct Geo { + int hid, nk, nv, qh, kvh, nl; + __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees + __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence + __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } + // a row's record of a linear layer: decay, strength and write a value head, key a key head, the tail after it + __device__ size_t log_bytes() const { return (size_t)nv * 4 + (size_t)nv * DK * 4 + (size_t)nk * DK * 2 + conv_rows() * 2; } +}; + +// One stream's window in this forward: its state blocks, cache length before the window, and rows. +struct Seg { + unsigned long long rec; // float recurrences, one rec() block a linear layer + unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer + unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time + unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each + unsigned long long log; // the window's rows' records (MAXW rows of nl layers of log_bytes), for keep + int cap, len, row0, rows, tap_rows, taps; +}; + +// A window row's record of one linear layer (Seg.log). +struct LogRow { + float* g; // [nv] decay + float* delta; // [nv][DK] the write's strength along dv + __nv_bfloat16* k; // [nk][DK] the key + __nv_bfloat16* conv; // [TAPS - 1][conv()] the tail after the row +}; + +__device__ __forceinline__ LogRow log_row(const Geo& geo, const Seg& sg, int layer, int row) { + char* base = reinterpret_cast(sg.log) + ((size_t)row * geo.nl + layer) * geo.log_bytes(); + LogRow l; + l.g = reinterpret_cast(base); + l.delta = l.g + geo.nv; + l.k = reinterpret_cast<__nv_bfloat16*>(l.delta + (size_t)geo.nv * DK); + l.conv = l.k + (size_t)geo.nk * DK; + return l; +} + +// The rows of a window to keep, in order (row 0 first), passed by value. +struct Path { + int n; + int rows[MAXW]; +}; + +// The row `steps` tokens before forward row `w` through rowpar: a forward row, or -1 - t for the stream's tail row t. +__device__ __forceinline__ int back(const int* __restrict__ rowpar, int w, int steps) { + int a = w; + for (int s = 1; s <= steps; s++) { + a = rowpar[a]; + if (a < 0) return -1 - (TAPS - 2 - (steps - s)); + } + return a; +} + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q35 diff --git a/zig/kernels/cuda/qwen35_draft.cu b/zig/kernels/cuda/qwen35_draft.cu new file mode 100644 index 0000000000..af2533bc51 --- /dev/null +++ b/zig/kernels/cuda/qwen35_draft.cu @@ -0,0 +1,164 @@ +// DFlash2 drafter glue on CUDA (cuda_draft.zig): a block of rows over the target's taps, draft_ops.metal's rounding. + +#include +#include + +namespace tf_q35d { + +constexpr int HD = 128; // the drafter's head width (q, k, v) + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q35d + +using namespace tf_q35d; + +// Tap rows (start + i) of a stream's ring of R rows, `width` bf16 each, into OUT rows 0..n. Grid (n), 256 threads. +extern "C" __global__ void tf_q35d_gather_taps(const __nv_bfloat16* __restrict__ RING, int R, long long start, int width, + __nv_bfloat16* __restrict__ OUT) { + const long long row = start + blockIdx.x; + const __nv_bfloat16* src = RING + (size_t)(row % R) * width; + __nv_bfloat16* dst = OUT + (size_t)blockIdx.x * width; + for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; +} + +// Rows of a window to gather, by value (a tree's kept path). +struct TapRows { + int n; + int rows[16]; +}; + +// Tap rows (start + rows.rows[i]) of a stream's ring into OUT rows 0..n: a kept path's rows in order. Grid (n). +extern "C" __global__ void tf_q35d_gather_rows(const __nv_bfloat16* __restrict__ RING, int R, long long start, TapRows rows, int width, + __nv_bfloat16* __restrict__ OUT) { + const long long row = start + rows.rows[blockIdx.x]; + const __nv_bfloat16* src = RING + (size_t)(row % R) * width; + __nv_bfloat16* dst = OUT + (size_t)blockIdx.x * width; + for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; +} + +// y[t, c] = sum over lags l of (BASE[l, c] + DYN[t, l, c / group]) * x[t - l, c] in bf16 steps; x[-1] = 0; RES last. +extern "C" __global__ void tf_q35d_dyn_conv(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ DYN, + const __nv_bfloat16* __restrict__ BASE, const __nv_bfloat16* __restrict__ RES, + __nv_bfloat16* __restrict__ OUT, int D, int group, int branch) { + const int t = blockIdx.x; + const int groups = D / group; + for (int c = threadIdx.x; c < D; c += blockDim.x) { + const int g = c / group; + float total = 0.0f; +#pragma unroll + for (int l = 0; l < 2; l++) { + if (t - l < 0) continue; + const float x = bf(X[(size_t)(t - l) * D + c]); + total = rbf(total + rbf(bf(BASE[(branch * 2 + l) * D + c]) * x)); + total = rbf(total + rbf(bf(DYN[(size_t)t * (4 * groups) + (branch * 2 + l) * groups + g]) * x)); + } + if (RES != nullptr) total = rbf(total + bf(RES[(size_t)t * D + c])); + OUT[(size_t)t * D + c] = __float2bfloat16_rn(total); + } +} + +// Per-head RMSNorm then rotate-half RoPE at position pos0 + row; OUT packed (rows, heads, HD), grid (rows * heads). +extern "C" __global__ void __launch_bounds__(32) tf_q35d_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, long long pos0, + float eps, float theta) { + const int head = blockIdx.x % heads; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + head * HD; + float v[4], acc = 0.0f; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = bf(x[4 * t + i]); + acc += v[i] * v[i]; + } + acc = warp_sum(acc); + const float inv = rsqrtf(acc / (float)HD + eps); + float n[4]; +#pragma unroll + for (int i = 0; i < 4; i++) n[i] = rbf(rbf(v[i] * inv) * bf(W[4 * t + i])); + // the pair of lane element d is d + HD / 2: this lane's partner values come from the lane HD / 8 = 16 away + const float pos = (float)(pos0 + (long long)m); + __nv_bfloat16* out = OUT + (m * heads + head) * HD; +#pragma unroll + for (int i = 0; i < 4; i++) { + const int d = 4 * t + i; + const float partner = __shfl_xor_sync(0xffffffffu, n[i], 16); + const int j = d < HD / 2 ? d : d - HD / 2; + const float angle = pos * powf(theta, -(float)j / (float)(HD / 2)); + float sn, cs; + sincosf(angle, &sn, &cs); + out[d] = __float2bfloat16_rn(d < HD / 2 ? n[i] * cs - partner * sn : n[i] * cs + partner * sn); + } +} + +// Rows of packed (kv_heads, HD) keys and values at positions pos0 + i into a layer's rings (slot pos % W). +extern "C" __global__ void tf_q35d_ring_put(const __nv_bfloat16* __restrict__ K, const __nv_bfloat16* __restrict__ V, + __nv_bfloat16* __restrict__ RK, __nv_bfloat16* __restrict__ RV, int width, long long pos0, int W) { + const int i = blockIdx.x; + const size_t slot = (size_t)((pos0 + i) % W); + for (int c = threadIdx.x; c < width; c += blockDim.x) { + RK[slot * width + c] = K[(size_t)i * width + c]; + RV[slot * width + c] = V[(size_t)i * width + c]; + } +} + +// Block attention: a row sees ring positions [max(begin, p - window + 1), pos0) and every block row; a warp a row. +extern "C" __global__ void __launch_bounds__(32) tf_q35d_attn(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KB, + const __nv_bfloat16* __restrict__ VB, const __nv_bfloat16* __restrict__ RK, + const __nv_bfloat16* __restrict__ RV, int rows, int heads, int kv_heads, int W, + long long begin, long long pos0, int window, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + const float scale = rsqrtf((float)HD); + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + const long long p = pos0 + m; + long long first = p - window + 1; + if (first < begin) first = begin; + for (long long pos = first; pos < pos0; pos++) { + const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + for (int r = 0; r < rows; r++) { + const __nv_bfloat16* k = KB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = VB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } +#pragma unroll + for (int i = 0; i < DPL; i++) OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(o[i] / l); +} diff --git a/zig/kernels/cuda/qwen35_keep.cu b/zig/kernels/cuda/qwen35_keep.cu new file mode 100644 index 0000000000..20d00fc212 --- /dev/null +++ b/zig/kernels/cuda/qwen35_keep.cu @@ -0,0 +1,58 @@ +// Keeping a verified window's path (cuda_cache.zig): the recurrences take its logged updates and its keys move down. + +#include "qwen35.cuh" + +using namespace tf_q35; + +// Each linear layer's recurrence through the path's (decay, key, write), as the forward's chain would have left it. +extern "C" __global__ void __launch_bounds__(128) tf_q35k_apply(Geo geo, Seg sg, Path path) { + const int NK = geo.nk, NV = geo.nv; + const int layer = blockIdx.z / NV, hv = blockIdx.z % NV; + const int hk = hv / (NV / NK); + const int dk = threadIdx.x; + const int dv = blockIdx.y * 4 + threadIdx.y; + constexpr int PT = DK / 32; + const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; + float state[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int j = 0; j < path.n; j++) { + const LogRow l = log_row(geo, sg, layer, path.rows[j]); + const float g = l.g[hv], d = l.delta[hv * DK + dv]; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * g; + state[i] = fmaf(bf(l.k[hk * DK + PT * dk + i]), d, state[i]); + } + } +#pragma unroll + for (int i = 0; i < PT; i++) st[i] = state[i]; +} + +// Each linear layer's convolution tail as logged after window row `last`. Grid (nl, tail values / 256), 256 threads. +extern "C" __global__ void tf_q35k_conv(Geo geo, Seg sg, int last) { + const int layer = blockIdx.x; + const size_t n = geo.conv_rows(); + const size_t i = (size_t)blockIdx.y * blockDim.x + threadIdx.x; + if (i >= n) return; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * n; + CS[i] = log_row(geo, sg, layer, last).conv[i]; +} + +// The path's keys and values move down to the slots after the cache; row j lands at or below j, never on one unread. +extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact(Geo geo, Seg sg, Path path) { + const int KVH = geo.kvh; + const int layer = blockIdx.x, h = blockIdx.y; + __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + ((size_t)layer * 2 * KVH + h) * sg.cap * HD; + __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + const int d = threadIdx.x; + for (int j = 0; j < path.n; j++) { + const int from = sg.len + path.rows[j], to = sg.len + j; + if (from != to) { + keys[(size_t)to * HD + d] = keys[(size_t)from * HD + d]; + values[(size_t)to * HD + d] = values[(size_t)from * HD + d]; + } + __syncthreads(); + } +} diff --git a/zig/kernels/cuda/qwen36_mtp.cu b/zig/kernels/cuda/qwen36_mtp.cu new file mode 100644 index 0000000000..5118a8aee7 --- /dev/null +++ b/zig/kernels/cuda/qwen36_mtp.cu @@ -0,0 +1,174 @@ +// The Qwen3.6 MTP drafter's glue on CUDA (cuda_mtp.zig): its attention at the target's head width over a ring. + +#include +#include + +namespace tf_q36m { + +constexpr int HD = 256; // the head width (q, k, v) +constexpr int ROT = 64; // rotary dims: pairs (i, i + 32), the target's partial rotary + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q36m + +using namespace tf_q36m; + +// A warp a (row, head): RMS norm, rope on [0, ROT) as pairs (i, i + 32) at pos0 + row * step; OUT (rows, heads, HD) +extern "C" __global__ void __launch_bounds__(32) tf_q36m_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, int in_head, + long long pos0, int step, float eps, float theta) { + constexpr int DPL = HD / 32; + const int head = blockIdx.x % heads; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + (size_t)head * in_head; // in_head 512 for q beside its gate, 256 for k + float v[DPL], acc = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) { + v[i] = bf(x[DPL * t + i]); + acc = fmaf(v[i], v[i], acc); + } + acc = warp_sum(acc); + const float inv = rsqrtf(acc / (float)HD + eps); + float n[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) n[i] = rbf(bf(W[DPL * t + i]) * (v[i] * inv)); + // a rotary pair's partner sits 32 elements away: DPL = 8 a lane, so 4 lanes away at the same i + const float pos = (float)(pos0 + (long long)m * step); + __nv_bfloat16* out = OUT + (m * heads + head) * HD; +#pragma unroll + for (int i = 0; i < DPL; i++) { + const int d = DPL * t + i; + const float partner = __shfl_xor_sync(0xffffffffu, n[i], 4); + float r = n[i]; + if (d < ROT) { + const int j = d < ROT / 2 ? d : d - ROT / 2; + const float angle = pos * powf(theta, -(float)j / (float)(ROT / 2)); + float sn, cs; + sincosf(angle, &sn, &cs); + r = d < ROT / 2 ? n[i] * cs - partner * sn : n[i] * cs + partner * sn; + } + out[d] = __float2bfloat16_rn(r); + } +} + +// A warp a (row, head): ring slots [begin, pos0), then block rows 0..m; OUT = softmax-weighted values x sigmoid(gate). +extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ G, int gstride, + const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, + const __nv_bfloat16* __restrict__ RK, const __nv_bfloat16* __restrict__ RV, int heads, + int kv_heads, int W, long long begin, long long pos0, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + const float scale = 1.0f / 16.0f; + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + for (long long pos = begin; pos < pos0; pos++) { + const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + for (int r = 0; r <= m; r++) { + const __nv_bfloat16* k = KB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = VB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; // the gate half +#pragma unroll + for (int i = 0; i < DPL; i++) { + const float gate = 1.0f / (1.0f + __expf(-bf(g[i]))); + OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(rbf(o[i] / l) * gate); + } +} + +// Tree rows at one depth, a warp a (row, head): ring [begin, end), the ancestor lanes ANC[m] names, then its own key. +extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn_tree(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ G, int gstride, + const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, + const __nv_bfloat16* __restrict__ LK, const __nv_bfloat16* __restrict__ LV, + const int* __restrict__ ANC, int anc_stride, const __nv_bfloat16* __restrict__ RK, + const __nv_bfloat16* __restrict__ RV, int heads, int kv_heads, int W, long long begin, + long long end, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + const float scale = 1.0f / 16.0f; + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + for (long long pos = begin; pos < end; pos++) { + const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + const int* anc = ANC + (size_t)m * anc_stride; + const int n = anc[0]; + for (int j = 0; j <= n; j++) { + const bool self = j == n; + const __nv_bfloat16* k = (self ? KB + (size_t)m * width : LK + (size_t)anc[1 + j] * width) + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = (self ? VB + (size_t)m * width : LV + (size_t)anc[1 + j] * width) + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) { + const float gate = 1.0f / (1.0f + __expf(-bf(g[i]))); + OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(rbf(o[i] / l) * gate); + } +} diff --git a/zig/kernels/cuda/scan_rows.cu b/zig/kernels/cuda/scan_rows.cu index 569ffedce8..b3ee6784cb 100644 --- a/zig/kernels/cuda/scan_rows.cu +++ b/zig/kernels/cuda/scan_rows.cu @@ -117,7 +117,7 @@ __global__ void __launch_bounds__(TPR * ROWS) scan_kernel( } // namespace tf_scan_rows -template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16*, const __nv_bfloat16*, float*, const float*, - const float*, const float*, __nv_bfloat16*, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); -template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16*, const __nv_bfloat16*, float*, const float*, - const float*, const float*, __nv_bfloat16*, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); +template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, float* __restrict__, const float* __restrict__, + const float* __restrict__, const float* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); +template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, float* __restrict__, const float* __restrict__, + const float* __restrict__, const float* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 64347ae50c..7474d2f077 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -13,6 +13,7 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "qwen4_exp", .title = "Qwen3.8 Flash Next" }, .{ .model_type = "glm5_next", .title = "GLM-5.3-Flash" }, .{ .model_type = "qwen3_5", .title = "Qwen3.5-2B or Qwen3.8-27B" }, + .{ .model_type = "qwen3_5_moe", .title = "Qwen3.6-35B-A3B" }, }; /// The family that serves `model_type` from the engine registry, or null when Zig cannot serve it (the 0.6 line may). @@ -38,7 +39,7 @@ pub fn cacheDir(a: Allocator, env: ?*const std.process.Environ.Map, override: ?[ } }.f; if (get(env, "HF_HUB_CACHE")) |c| return c; - const home = get(env, "HOME") orelse ""; + const home = get(env, "HOME") orelse get(env, "USERPROFILE") orelse ""; // Windows sets USERPROFILE, not HOME if (get(env, "HF_HOME")) |h| return std.fs.path.join(a, &.{ h, "hub" }); return std.fs.path.join(a, &.{ home, ".cache", "huggingface", "hub" }); } diff --git a/zig/src/cli/pull_exlr8.zig b/zig/src/cli/pull_exlr8.zig index 3e8c641f7e..c56d16ada2 100644 --- a/zig/src/cli/pull_exlr8.zig +++ b/zig/src/cli/pull_exlr8.zig @@ -88,6 +88,7 @@ pub fn run(a: Allocator, io: std.Io, out: *std.Io.Writer, err_out: *std.Io.Write /// Index maps, connections and scans pass the default 256 descriptors: raise this process's own limit. fn moreFiles() void { + if (@import("builtin").os.tag == .windows) return; // no descriptor limit to raise var l = std.posix.getrlimit(.NOFILE) catch return; const want: std.posix.rlim_t = if (l.max == std.posix.RLIM.INFINITY) 4096 else @min(l.max, 4096); if (l.cur >= want) return; diff --git a/zig/src/cli/pull_file.zig b/zig/src/cli/pull_file.zig index c3de2654cf..a5750e3e79 100644 --- a/zig/src/cli/pull_file.zig +++ b/zig/src/cli/pull_file.zig @@ -1,5 +1,6 @@ //! One file from the hub into the blob store: range resume, size and sha256 checks, retries, and its snapshot link. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; /// One file the hub's tree lists for the revision: LFS files carry their sha256, small files their git blob sha1. @@ -143,7 +144,21 @@ pub fn linkIntoSnapshot(a: Allocator, io: std.Io, snapshot_dir: []const u8, path if (std.fs.path.dirname(link_path)) |parent| try std.Io.Dir.cwd().createDirPath(io, parent); std.Io.Dir.cwd().deleteFile(io, link_path) catch {}; const target = try std.fs.path.join(a, &.{ "..", "..", "blobs", blob_name }); - std.Io.Dir.cwd().symLink(io, target, link_path, .{}) catch return error.SymLinkFailed; + std.Io.Dir.cwd().symLink(io, target, link_path, .{}) catch |err| { + // Windows grants symlinks only with Developer Mode or elevation: a hard link to the blob costs no space + if (@import("builtin").os.tag != .windows) { + std.log.err("cannot link {s} to its blob: {t}", .{ link_path, err }); + return error.SymLinkFailed; + } + const repo_dir = std.fs.path.dirname(std.fs.path.dirname(snapshot_dir) orelse return error.SymLinkFailed) orelse return error.SymLinkFailed; + const blob = try std.fs.path.joinZ(a, &.{ repo_dir, "blobs", blob_name }); + const linked = try a.dupeSentinel(u8, link_path, 0); + _ = platform.c.unlink(linked); // the refused symlink can leave an empty file behind + if (platform.c.link(blob, linked) != 0) { + std.log.err("cannot link {s} to its blob: {t}", .{ link_path, platform.c.errno(-1) }); + return error.SymLinkFailed; + } + }; } test "a hub tree path stays inside the snapshot: relative, without . or .. segments" { diff --git a/zig/src/cluster/admission_test.zig b/zig/src/cluster/admission_test.zig index 1b5b0f1b3b..200fe2b111 100644 --- a/zig/src/cluster/admission_test.zig +++ b/zig/src/cluster/admission_test.zig @@ -1,5 +1,6 @@ //! 32 to 128 streams: what fits at 8k and 32k context, prompt rows a round, and 32 toy streams joining while others decode. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const plan = @import("plan.zig"); const budget = @import("budget.zig"); @@ -115,7 +116,7 @@ test "32 toy streams join while others decode, prompt chunks in the same rounds: } test "the admission tables for CLUSTER.md (TF_CLUSTER_REPORT=1)" { - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; + if (platform.getenv("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/cli_test.zig b/zig/src/cluster/cli_test.zig index 267fca93d6..bc7b7f58bd 100644 --- a/zig/src/cluster/cli_test.zig +++ b/zig/src/cluster/cli_test.zig @@ -1,5 +1,6 @@ //! The cluster commands end to end on fake hosts: init, check, status, serve with in-process bring-up, and node. const std = @import("std"); +const platform = @import("sys"); const cli = @import("cli.zig"); const fixtures = @import("probe_fixtures.zig"); @@ -86,7 +87,7 @@ const studios = ; test "check plans K3 from its config and Hub file list and names the raise a low node needs; serve stops before launching" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/cost_test.zig b/zig/src/cluster/cost_test.zig index dec13181ca..47293a6d20 100644 --- a/zig/src/cluster/cost_test.zig +++ b/zig/src/cluster/cost_test.zig @@ -1,5 +1,6 @@ //! K3 and GLM-5.3 rounds on four 512 GiB Studios at 1 to 128 rows: dedup, caches, collectives, and the MLA split. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const model = @import("model.zig"); const estimate = @import("estimate.zig"); @@ -66,7 +67,7 @@ test "GLM-5.3 in FP8 reads far less a token than K3, and 128 streams at 32k cont } test "the cost tables for CLUSTER.md (TF_CLUSTER_REPORT=1)" { - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; + if (platform.getenv("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/estimate.zig b/zig/src/cluster/estimate.zig index 74917ba740..9d33e43313 100644 --- a/zig/src/cluster/estimate.zig +++ b/zig/src/cluster/estimate.zig @@ -1,5 +1,6 @@ //! A checkpoint's tensor table from its config and the Hub's file list, so a cluster can be planned before the weights land. const std = @import("std"); +const platform = @import("sys"); const checkpoint = @import("checkpoint.zig"); const model = @import("model.zig"); const roles = @import("roles.zig"); @@ -282,7 +283,7 @@ test "the K3 table matches every shard size to within its header" { } test "the real K3 file list (TF_K3_DIR): 96 files, 1.561 TB, every header under 0.01% of its file" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/model.zig b/zig/src/cluster/model.zig index 644b117264..91b37ea140 100644 --- a/zig/src/cluster/model.zig +++ b/zig/src/cluster/model.zig @@ -1,5 +1,6 @@ //! A model's shape from its config.json (text_config when nested): what the planner and the memory budgets need. const std = @import("std"); +const platform = @import("sys"); pub const Attention = enum(u8) { gqa, mla, linear }; @@ -257,7 +258,7 @@ test "GLM-5.3: every layer MLA with a 2,048-token indexer, 75 MoE layers, one MT } test "the real K3 config (TF_K3_DIR) parses to the literal" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/plan_test.zig b/zig/src/cluster/plan_test.zig index 666ecb3ba2..0514244c35 100644 --- a/zig/src/cluster/plan_test.zig +++ b/zig/src/cluster/plan_test.zig @@ -1,5 +1,6 @@ //! Kimi K3 on four 512 GiB M3 Ultras: the default layout, layouts from a config, determinism, and refusals. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const model = @import("model.zig"); const estimate = @import("estimate.zig"); @@ -34,7 +35,7 @@ pub fn k3(a: std.mem.Allocator) !checkpoint.Checkpoint { const size: u64 = if (i == 0) 2_341_216_112 else if (i < s.layers) (if (kda) 16_990_911_504 else 16_567_501_776) else if (i == s.layers) 4_697_664_072 else if (i == s.layers + 1) 92_289_328 else 802_448_352; f.* = .{ .name = try std.fmt.allocPrint(a, "model-{d:0>5}-of-{d:0>6}.safetensors", .{ i + 1, s.layers + 3 }), .size = size }; } - if (std.testing.environ.getPosix("TF_K3_DIR")) |dir| { + if (platform.getenv("TF_K3_DIR")) |dir| { const text = try std.Io.Dir.cwd().readFileAlloc(std.testing.io, try std.fs.path.join(a, &.{ dir, "files.tsv" }), a, .limited(1 << 20)); files = try estimate.manifest(a, text); } diff --git a/zig/src/cluster/status.zig b/zig/src/cluster/status.zig index ecb6dd7457..49e0395dbf 100644 --- a/zig/src/cluster/status.zig +++ b/zig/src/cluster/status.zig @@ -1,5 +1,6 @@ //! What `tensorfold cluster status` and `check` print: nodes, links, memory, the placement per rank, costs and load progress. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const topology = @import("topology.zig"); const plan_mod = @import("plan.zig"); @@ -99,7 +100,7 @@ test "the K3 status on four Studios prints every rank, fits, and the round costs try placement(&out.writer, &p, needs, fs); try costs(&out.writer, &p, needs, &s, 8192); const text = try a.dupe(u8, out.written()); - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") != null) { + if (platform.getenv("TF_CLUSTER_REPORT") != null) { std.debug.print("\n{s}\n", .{text}); const others = [_]struct { []const u8, plan_mod.Options }{ .{ "experts split by tensor parallelism", .{ .layout = .{ .tensor = 4, .expert = 1 } } }, diff --git a/zig/src/cluster/topology.zig b/zig/src/cluster/topology.zig index 6e5292c8dd..6c4d2d7564 100644 --- a/zig/src/cluster/topology.zig +++ b/zig/src/cluster/topology.zig @@ -1,5 +1,6 @@ //! The cluster's wiring from every node's ports: which port reaches which node, at what rate, and the graph's shape. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); /// One cable: node `a`'s port to node `b`'s port (indices into the node list and their port lists). @@ -134,7 +135,7 @@ test "four Macs with six cables form a mesh; three cables a chain or a star; a c } test "the lent cluster's real probes (TF_CLUSTER_PROBES): four M3 Ultras in a full mesh at 80 Gb/s" { - const root = std.testing.environ.getPosix("TF_CLUSTER_PROBES") orelse return error.SkipZigTest; + const root = platform.getenv("TF_CLUSTER_PROBES") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/core/checkpoint_metal.zig b/zig/src/core/checkpoint_metal.zig index aabd70f340..45939c6bb5 100644 --- a/zig/src/core/checkpoint_metal.zig +++ b/zig/src/core/checkpoint_metal.zig @@ -1,5 +1,6 @@ //! The shared safetensors index read into Metal buffers: one shared buffer a shard, each tensor at its data offset. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("safetensors.zig"); @@ -33,7 +34,7 @@ pub const Deferred = struct { file: usize, at: usize, bytes: usize, dtype: DType pub const Later = *const fn (name: []const u8, rank: usize) bool; /// A shard file that deferred tensors stay in, and its address-only views' reservation once one is asked for. -const File = struct { fd: std.c.fd_t, end: usize = 0, views: ?[]align(std.heap.page_size_min) u8 = null }; +const File = struct { fd: platform.c.fd_t, end: usize = 0, views: ?[]align(std.heap.page_size_min) u8 = null }; /// Kept tensors pack at this alignment; reads this large or larger run on 8 threads. pub const alignment = 256; @@ -63,14 +64,14 @@ pub const Checkpoint = struct { /// As addFileSelected; with `later`, the tensors it names stay in the file for `read` and the rest pack into one buffer. pub fn addFileSplit(self: *Checkpoint, device: mtl.Device, path: [:0]const u8, prefix: []const u8, selected: ?[]const u8, later: ?Later) !void { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) { std.log.err("cannot open {s}", .{path}); return error.OpenFailed; } var held = false; defer if (!held) { - _ = std.c.close(fd); + _ = platform.c.close(fd); }; var head: [8]u8 = undefined; try readAll(fd, &head, 0); @@ -78,7 +79,7 @@ pub const Checkpoint = struct { const header = try self.allocator.alloc(u8, header_len); defer self.allocator.free(header); try readAll(fd, header, 8); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); const data_start = 8 + header_len; const data_len: usize = @as(usize, @intCast(end)) - data_start; @@ -114,7 +115,7 @@ pub const Checkpoint = struct { } /// The kept entries packed into one buffer, the `later` ones recorded for `read`; `held` once the file stays open. - fn split(self: *Checkpoint, device: mtl.Device, fd: std.c.fd_t, held: *bool, data_start: usize, entries: st.Header, prefix: []const u8, later: Later) !void { + fn split(self: *Checkpoint, device: mtl.Device, fd: platform.c.fd_t, held: *bool, data_start: usize, entries: st.Header, prefix: []const u8, later: Later) !void { var arena = std.heap.ArenaAllocator.init(self.allocator); defer arena.deinit(); const names = try arena.allocator().alloc([]const u8, entries.count()); @@ -184,7 +185,7 @@ pub const Checkpoint = struct { /// Close the files the deferred tensors live in, once load has read them. pub fn closeFiles(self: *Checkpoint) void { for (self.files.items) |*file| if (file.fd >= 0) { - _ = std.c.close(file.fd); + _ = platform.c.close(file.fd); file.fd = -1; }; } @@ -300,16 +301,16 @@ pub const Stage = struct { } }; -fn readAll(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readAll(fd: platform.c.fd_t, dest: []u8, at: usize) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } } -const Part = struct { fd: std.c.fd_t, dest: []u8, at: usize, failed: bool = false }; +const Part = struct { fd: platform.c.fd_t, dest: []u8, at: usize, failed: bool = false }; fn readPart(part: *Part) void { readAll(part.fd, part.dest, part.at) catch { @@ -318,7 +319,7 @@ fn readPart(part: *Part) void { } /// pread `dest.len` bytes from `at` on 8 threads (page cache or disk, the copy is the load's main cost). -fn readParallel(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readParallel(fd: platform.c.fd_t, dest: []u8, at: usize) !void { const n_threads = 8; var parts: [n_threads]Part = undefined; var threads: [n_threads]?std.Thread = @splat(null); @@ -335,7 +336,7 @@ fn readParallel(fd: std.c.fd_t, dest: []u8, at: usize) !void { } /// One tensor's bytes: small ones on this thread, large ones on 8. -fn readInto(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readInto(fd: platform.c.fd_t, dest: []u8, at: usize) !void { return if (dest.len >= parallel_bytes) readParallel(fd, dest, at) else readAll(fd, dest, at); } diff --git a/zig/src/core/direct_io.zig b/zig/src/core/direct_io.zig index c06f213388..f5bba757e2 100644 --- a/zig/src/core/direct_io.zig +++ b/zig/src/core/direct_io.zig @@ -1,6 +1,7 @@ //! Positional reads that skip the page cache (O_DIRECT on Linux): 4 KiB-aligned spans into aligned buffers. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); /// Offset, length and buffer alignment a direct read needs (the logical block size of every NVMe we run on). @@ -20,22 +21,22 @@ pub fn fits(bytes: usize) usize { } pub const File = struct { - fd: std.c.fd_t, + fd: platform.c.fd_t, direct: bool, /// `path` read-only, direct when its file system allows it, else through the page cache. pub fn open(path: [:0]const u8) !File { - if (builtin.os.tag == .linux) { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }); + if (builtin.os.tag == .linux or builtin.os.tag == .windows) { + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }); if (fd >= 0) return .{ .fd = fd, .direct = true }; } - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }); if (fd < 0) return error.FileNotFound; return .{ .fd = fd, .direct = false }; } pub fn close(f: *File) void { - _ = std.c.close(f.fd); + _ = platform.c.close(f.fd); f.* = undefined; } @@ -46,11 +47,11 @@ pub const File = struct { } }; -fn sysPread(_: *anyopaque, fd: std.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { - return std.c.pread(fd, buf, len, @intCast(offset)); +fn sysPread(_: *anyopaque, fd: platform.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { + return platform.c.pread(fd, buf, len, @intCast(offset)); } -fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: *const fn (*anyopaque, std.c.fd_t, [*]u8, usize, u64) isize) ![]u8 { +fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: *const fn (*anyopaque, platform.c.fd_t, [*]u8, usize, u64) isize) ![]u8 { const lo = std.mem.alignBackward(u64, offset, alignment); const want: usize = @intCast(offset + len - lo); const whole = span(offset, len); @@ -62,8 +63,8 @@ fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: const ask = whole - at; const n = pread(ctx, f.fd, buf.ptr + at, ask, lo + @as(u64, @intCast(at))); if (n < 0) { - if (std.c.errno(n) == .INTR) continue; - std.log.err("read of {d} bytes at {d} failed: {t}", .{ ask, lo + at, std.c.errno(n) }); + if (platform.c.errno(n) == .INTR) continue; + std.log.err("read of {d} bytes at {d} failed: {t}", .{ ask, lo + at, platform.c.errno(n) }); return error.ReadFailed; } if (n == 0) return error.EndOfFile; @@ -84,12 +85,12 @@ test "span and fits" { } test "direct reads return the same bytes as buffered reads" { - if (builtin.os.tag != .linux) return error.SkipZigTest; + if (builtin.os.tag != .linux and builtin.os.tag != .windows) return error.SkipZigTest; const gpa = std.testing.allocator; var f = try File.open("/proc/self/exe"); defer f.close(); // The reference uses the page cache, so its unaligned reads do not inherit O_DIRECT. - var plain: File = .{ .fd = std.c.open("/proc/self/exe", .{ .ACCMODE = .RDONLY, .CLOEXEC = true }), .direct = false }; + var plain: File = .{ .fd = platform.c.open("/proc/self/exe", .{ .ACCMODE = .RDONLY, .CLOEXEC = true }), .direct = false }; try std.testing.expect(plain.fd >= 0); defer plain.close(); const buf = try gpa.alignedAlloc(u8, .fromByteUnits(alignment), 1 << 20); @@ -97,7 +98,7 @@ test "direct reads return the same bytes as buffered reads" { const want = try gpa.alloc(u8, 1 << 19); defer gpa.free(want); for ([_][2]usize{ .{ 0, 1 }, .{ 1, 4095 }, .{ 4095, 2 }, .{ 12345, 100000 }, .{ 8192, 1 << 19 } }) |c| { - const n = std.c.pread(plain.fd, want.ptr, c[1], @intCast(c[0])); + const n = platform.c.pread(plain.fd, want.ptr, c[1], @intCast(c[0])); try std.testing.expectEqual(@as(isize, @intCast(c[1])), n); try std.testing.expectEqualSlices(u8, want[0..c[1]], try f.read(buf, c[0], c[1])); } @@ -111,7 +112,7 @@ const Script = struct { calls: usize = 0, bad: bool = false, - fn pread(ctx: *anyopaque, _: std.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { + fn pread(ctx: *anyopaque, _: platform.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { const s: *Script = @ptrCast(@alignCast(ctx)); if (s.calls < s.offs.len) s.offs[s.calls] = offset; s.calls += 1; diff --git a/zig/src/core/exlr8_oracle_test.zig b/zig/src/core/exlr8_oracle_test.zig index 13c854dee9..6bdb2addc2 100644 --- a/zig/src/core/exlr8_oracle_test.zig +++ b/zig/src/core/exlr8_oracle_test.zig @@ -1,6 +1,7 @@ //! EXLR8 decoding against fixtures the reference loader (MIT) wrote outside the tree: layers 0, 3, 40, 78, vocab. const std = @import("std"); +const platform = @import("sys"); const codec = @import("exlr8_codec.zig"); const ex = @import("exlr8_experts.zig"); const dense = @import("exlr8_dense.zig"); @@ -169,7 +170,7 @@ const Layers = struct { }; test "expert halves of layers 3, 40 and 78: stored checksums, the reference's tiles and W bits" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -193,7 +194,7 @@ test "expert halves of layers 3, 40 and 78: stored checksums, the reference's ti h.suh(suh[0..h.rows()]); h.svh(svh[0..h.cols()]); try checkSamples(codes, h.rows(), h.cols(), suh[0..h.rows()], svh[0..h.cols()], unit.w_samples, false, &stats); - if (std.testing.environ.getPosix("TF_EXLR8_ORACLE_FULL")) |full| try wholeHalf(full, unit, codes, h, suh[0..h.rows()], svh[0..h.cols()]); + if (platform.getenv("TF_EXLR8_ORACLE_FULL")) |full| try wholeHalf(full, unit, codes, h, suh[0..h.rows()], svh[0..h.cols()]); } stats.print("expert halves"); } @@ -233,7 +234,7 @@ fn wholeHalf(dir: []const u8, unit: Unit, codes: []const u8, h: ex.Half, suh: [] } test "every half of every stored width verifies, as the reference's --verify counts them" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit(); const fx = try load(arena.allocator()); @@ -260,7 +261,7 @@ test "every half of every stored width verifies, as the reference's --verify cou } test "dense and vocabulary files: records, checksums, every format bit for bit, plain tensors as shipped" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -331,7 +332,7 @@ fn denseBits(m: dense.Matrix, unit: DenseUnit) !void { } test "the eight manifests parse to the reference's widths, and every width they ask of a present layer is stored" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -371,7 +372,7 @@ test "the eight manifests parse to the reference's widths, and every width they } test "every present release file matches its published sha256 and size" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var checked: usize = 0; for ([_][]const u8{ "experts", "dense" }) |sub| { diff --git a/zig/src/core/heat.zig b/zig/src/core/heat.zig index dc024d653d..eec72a6348 100644 --- a/zig/src/core/heat.zig +++ b/zig/src/core/heat.zig @@ -1,5 +1,6 @@ //! Opt-in wait before a prompt chunk when the hottest thermal zone is above a band. const std = @import("std"); +const platform = @import("sys"); pub const default_root = "/sys/class/thermal"; @@ -129,7 +130,7 @@ fn noteNoZone(root: []const u8) void { } fn envSpan(name: [*:0]const u8) ?[]const u8 { - const p = std.c.getenv(name) orelse return null; + const p = platform.c.getenv(name) orelse return null; const s = std.mem.span(p); if (s.len == 0) return null; return s; @@ -194,13 +195,15 @@ test "the chunk runs under the high band, and a hot reading waits until the low try std.testing.expectEqual(@as(u32, 1), stuck.sleeps); // a cancel ends the wait after the pause it came in } -extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: i32) i32; -extern "c" fn unsetenv(name: [*:0]const u8) i32; +const env = if (@import("builtin").os.tag == .windows) @import("sys").c else struct { + extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: i32) i32; + extern "c" fn unsetenv(name: [*:0]const u8) i32; +}; fn putEnv(name: [*:0]const u8, value: ?[*:0]const u8) void { if (value) |v| { - _ = setenv(name, v, 1); - } else _ = unsetenv(name); + _ = env.setenv(name, v, 1); + } else _ = env.unsetenv(name); } test "TF_HEAT_HIGH and TF_HEAT_LOW set the bands together, and TF_HEAT_ROOT the zones' folder" { diff --git a/zig/src/core/lanes/cost_cache.zig b/zig/src/core/lanes/cost_cache.zig index 1b04fe8fde..279aca3931 100644 --- a/zig/src/core/lanes/cost_cache.zig +++ b/zig/src/core/lanes/cost_cache.zig @@ -1,5 +1,6 @@ //! Timed costs are measured once per engine build, chip and model shape, then read back by later loads. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; const magic = "TFWC2\n"; @@ -53,7 +54,7 @@ fn hashPart(h: *std.crypto.hash.sha2.Sha256, part: []const u8) void { } fn path(a: Allocator, k: [32]u8, suffix: []const u8) !?[]u8 { - const home = std.c.getenv("HOME") orelse return null; + const home = platform.c.getenv("HOME") orelse return null; return try std.fmt.allocPrint(a, "{s}/.cache/tensorfold/window-costs/{s}{s}", .{ std.mem.span(home), std.fmt.bytesToHex(k, .lower), suffix }); } @@ -78,7 +79,7 @@ pub fn save(comptime T: type, a: Allocator, io: std.Io, k: [32]u8, value: T) voi const p = (path(a, k, ".bin") catch return) orelse return; defer a.free(p); std.Io.Dir.cwd().createDirPath(io, std.fs.path.dirname(p) orelse return) catch return; - const temp = std.fmt.allocPrint(a, "{s}.{d}.tmp", .{ p, std.c.getpid() }) catch return; + const temp = std.fmt.allocPrint(a, "{s}.{d}.tmp", .{ p, platform.c.getpid() }) catch return; defer a.free(temp); const bytes = std.mem.concat(a, u8, &.{ magic, std.mem.asBytes(&value) }) catch return; defer a.free(bytes); diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 3bdffe7c7a..e1a79fd2de 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -60,7 +60,9 @@ pub const Drafted = struct { m.plain_guard = d.plain_guard; m.draft_streams = d.batched; m.draft_probabilities = false; - m.head_trees = false; + m.head_trees = d.trees; + if (d.window_costs.len > 0) m.window_costs = d.window_costs; + if (d.lane_costs.len > 0) m.lane_costs = d.lane_costs; return m; } @@ -157,15 +159,17 @@ pub const Drafted = struct { x.holds.clearRetainingCapacity(); try x.firsts.resize(x.gpa, requests.len); for (requests, x.firsts.items) |r, *tok| { - if (r.lanes != null) return error.TreesNotBuilt; + if (r.lanes != null and !x.drafter.facts().trees) return error.TreesNotBuilt; const lane = x.lanes.getPtr(r.stream) orelse return error.NotDrafting; lane.filler = 0; // a new request supersedes fillers a copy round left unused var pending: u32 = undefined; if (r.rows) |rows| { - for (rows, 0..) |row, i| if (row != i) return error.TreesNotBuilt; - const n: u32 = @intCast(rows.len); + // a tree's kept path names the window's rows out of order: the features of all up to its last + var prefix = true; + for (rows, 0..) |row, i| prefix = prefix and row == i; + const n: u32 = @intCast(if (prefix) rows.len else rows[rows.len - 1] + 1); const f = try x.target.features(r.stream, x.drafter.taps(), r.start, n); - try x.absorbs.append(x.gpa, .{ .stream = r.stream, .features = f, .start = r.start, .follow = r.follow[0..rows.len] }); + try x.absorbs.append(x.gpa, .{ .stream = r.stream, .features = f, .start = r.start, .follow = r.follow[0..rows.len], .rows = if (prefix) null else rows }); pending = r.follow[rows.len - 1]; lane.rows = true; } else { @@ -187,7 +191,7 @@ pub const Drafted = struct { lane.filler = r.depth; continue; } - try x.holds.append(x.gpa, .{ .stream = r.stream, .pending = pending, .position = r.position, .depth = r.depth }); + try x.holds.append(x.gpa, .{ .stream = r.stream, .pending = pending, .position = r.position, .depth = r.depth, .lanes = r.lanes }); } if (x.absorbs.items.len > 0) try x.drafter.absorb(x.absorbs.items); if (x.holds.items.len > 0) try x.drafter.hold(x.holds.items); diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 0f8eab242a..c956522f50 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -1,6 +1,8 @@ //! A drafter outside the target family reading its tapped states; each call takes a round's streams together. const Stream = @import("stream.zig").Stream; const Features = @import("backend.zig").Features; +const Cost = @import("config.zig").Cost; +const Shape = @import("shape.zig").Shape; /// What the round loop's depth rule needs from the drafter itself (not from the target, whose own head may not exist). pub const Facts = struct { @@ -10,13 +12,16 @@ pub const Facts = struct { plain_guard: bool = true, // plain rounds compete with drafted depths (the rule may choose no drafts) batched: bool = false, // hold() runs a shared round's streams as one batch stream_bytes: usize = 0, // device bytes each opened stream holds (its caches), for memory admission + trees: bool = false, // hold() takes a tree of lanes (Hold.lanes), each lane its parent's rank-r token + window_costs: []const Cost = &.{}, // the target's window ms by rows (timed at load), for the depth rule + lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) }; -/// Target rows [start, start + features.rows) and the token after each; rows at or past `start` replace what was held. -pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32 }; +/// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb. +pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32, rows: ?[]const u32 = null }; -/// Draft `depth` tokens after `pending`, the first at `position`; asked only once row `position - 2` was absorbed. -pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32 }; +/// `depth` drafts after `pending` at `position`; with `lanes` a tree, lane i the rank-r choice after its parent. +pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32, lanes: ?*const Shape = null }; pub const Drafter = struct { ptr: *anyopaque, diff --git a/zig/src/core/lanes/grammar.zig b/zig/src/core/lanes/grammar.zig index 371c2ac4d8..2fe9ae3d87 100644 --- a/zig/src/core/lanes/grammar.zig +++ b/zig/src/core/lanes/grammar.zig @@ -1,6 +1,7 @@ //! Structured output: libtfgrammar (xgrammar, loaded at the first structured request) and a stream's grammar walk. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Allocator = std.mem.Allocator; @@ -37,7 +38,7 @@ fn load(io: std.Io, a: Allocator) ?Fns { library, }; var dl = for (names) |name| { - break std.DynLib.open(name) catch continue; + break platform.DynLib.open(name) catch continue; } else return null; var f: Fns = undefined; const info = @typeInfo(Fns).@"struct"; diff --git a/zig/src/core/lanes/learned_dirs.zig b/zig/src/core/lanes/learned_dirs.zig index 780a997a5c..cdf9648438 100644 --- a/zig/src/core/lanes/learned_dirs.zig +++ b/zig/src/core/lanes/learned_dirs.zig @@ -1,5 +1,6 @@ //! Bounded-name learned directories: both ranks plan the same identity deletions before either writes. const std = @import("std"); +const platform = @import("sys"); pub fn rootOf(dir: []const u8) []const u8 { return std.fs.path.dirname(dir) orelse "."; } @@ -13,23 +14,23 @@ fn isCurrent(keep: []const u8, id: u64) bool { } pub fn next(root: []const u8, keep: []const u8, after: ?u64) !?u64 { var buf: [1100]u8 = undefined; - const d = std.c.opendir(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0)) orelse return error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0)) orelse return error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); var out: ?u64 = null; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (name.len != 16) continue; const id = std.fmt.parseInt(u64, name, 16) catch continue; if (id == std.math.maxInt(u64)) return error.LearnedDirectoryType; - if (ent.type != std.c.DT.DIR) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.DIR) return error.LearnedDirectoryType; if (isCurrent(keep, id) or (after != null and id <= after.?)) continue; if (out == null or id < out.?) out = id; } return out; } pub fn unlink(file: [:0]const u8) !void { - const rc = std.c.unlink(file); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.unlink(file); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } fn stateKey(name: []const u8) ?u64 { @@ -52,38 +53,38 @@ fn identityName(name: []const u8) bool { } fn unlinkAt(fd: c_int, name: [*:0]const u8) !void { - const rc = std.c.unlinkat(fd, name, 0); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.unlinkat(fd, name, 0); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } /// Only the configured root may be an alias; cleanup never follows an identity or file link. pub fn cleanupParts(root: []const u8) !void { var buf: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0), .{ .ACCMODE = .RDONLY, .DIRECTORY = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0), .{ .ACCMODE = .RDONLY, .DIRECTORY = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const d = std.c.fdopendir(fd) orelse { - _ = std.c.close(fd); + const d = platform.c.fdopendir(fd) orelse { + _ = platform.c.close(fd); return error.LearnedDirectoryRead; }; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); - if (ent.type == std.c.DT.REG and std.mem.eql(u8, name, "clock.part")) { + if (ent.type == platform.c.DT.REG and std.mem.eql(u8, name, "clock.part")) { try unlinkAt(fd, @ptrCast(&ent.name)); continue; } if (!identityName(name)) continue; - if (ent.type != std.c.DT.DIR) return error.LearnedDirectoryType; - const sub_fd = std.c.openat(fd, @ptrCast(&ent.name), .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + if (ent.type != platform.c.DT.DIR) return error.LearnedDirectoryType; + const sub_fd = platform.c.openat(fd, @ptrCast(&ent.name), .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (sub_fd < 0) return error.LearnedDirectoryRead; - const sub = std.c.fdopendir(sub_fd) orelse { - _ = std.c.close(sub_fd); + const sub = platform.c.fdopendir(sub_fd) orelse { + _ = platform.c.close(sub_fd); return error.LearnedDirectoryRead; }; { - defer _ = std.c.closedir(sub); - while (std.c.readdir(sub)) |file| { - if (file.type == std.c.DT.REG and partName(std.mem.sliceTo(&file.name, 0))) try unlinkAt(sub_fd, @ptrCast(&file.name)); + defer _ = platform.c.closedir(sub); + while (platform.c.readdir(sub)) |file| { + if (file.type == platform.c.DT.REG and partName(std.mem.sliceTo(&file.name, 0))) try unlinkAt(sub_fd, @ptrCast(&file.name)); } } } @@ -91,15 +92,15 @@ pub fn cleanupParts(root: []const u8) !void { /// Payloads have no token metadata to recover a lost index entry; only exact owned, unreferenced names go. pub fn removeUnindexed(dir: [:0]const u8, metas: anytype) !void { - const fd = std.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const d = std.c.fdopendir(fd) orelse { - _ = std.c.close(fd); + const d = platform.c.fdopendir(fd) orelse { + _ = platform.c.close(fd); return error.LearnedDirectoryRead; }; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { - if (ent.type != std.c.DT.REG) continue; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { + if (ent.type != platform.c.DT.REG) continue; const name = std.mem.sliceTo(&ent.name, 0); const key = stateKey(name) orelse continue; const indexed = for (metas) |x| { @@ -110,29 +111,29 @@ pub fn removeUnindexed(dir: [:0]const u8, metas: anytype) !void { } } pub fn bytes(dir: [:0]const u8) !u64 { - const d = std.c.opendir(dir) orelse return if (std.c.errno(-1) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(dir) orelse return if (platform.c.errno(-1) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); var n: u64 = 0; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) continue; - if (ent.type != std.c.DT.REG) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.REG) return error.LearnedDirectoryType; var buf: [1400]u8 = undefined; const f = try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0); - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - _ = std.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); + _ = platform.c.close(fd); if (end < 0) return error.LearnedDirectoryRead; n = try std.math.add(u64, n, @intCast(end)); } return n; } pub fn fileBytes(file: [:0]const u8) !u64 { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.LearnedDirectoryRead; return @intCast(end); } @@ -156,16 +157,16 @@ pub fn total(root: []const u8) !u64 { pub fn used(root: []const u8, id: u64) !u64 { var buf: [1100]u8 = undefined; const f = try std.fmt.bufPrintSentinel(&buf, "{s}/{x:0>16}/used", .{ root, id }, 0); - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.close(fd); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.close(fd); var n: u64 = 0; const data = std.mem.asBytes(&n); var done: usize = 0; while (done < data.len) { - const count = std.c.read(fd, data.ptr + done, data.len - done); + const count = platform.c.read(fd, data.ptr + done, data.len - done); if (count < 0) { - if (std.c.errno(count) == .INTR) continue; + if (platform.c.errno(count) == .INTR) continue; return error.LearnedDirectoryRead; } if (count == 0) return 0; // A torn recency hint is oldest; it cannot make a learned payload valid. @@ -178,23 +179,23 @@ pub fn removeOther(root: []const u8, keep: []const u8, id: u64) !void { var buf: [1100]u8 = undefined; const dir = try path(&buf, root, id); _ = try bytes(dir); - const d = std.c.opendir(dir) orelse return if (std.c.errno(-1) == .NOENT) {} else error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + const d = platform.c.opendir(dir) orelse return if (platform.c.errno(-1) == .NOENT) {} else error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) continue; - if (ent.type != std.c.DT.REG) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.REG) return error.LearnedDirectoryType; var f: [1400]u8 = undefined; try unlink(try std.fmt.bufPrintSentinel(&f, "{s}/{s}", .{ dir, name }, 0)); } - const rc = std.c.rmdir(dir); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.rmdir(dir); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } test "acknowledged deletion refuses a directory and accepts an absent file" { var buf: [128]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-unlink-{d}", .{std.c.getpid()}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-unlink-{d}", .{platform.c.getpid()}, 0); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); try std.testing.expectError(error.LearnedUnlink, unlink(dir)); var missing: [160]u8 = undefined; try unlink(try std.fmt.bufPrintSentinel(&missing, "{s}/absent", .{dir}, 0)); diff --git a/zig/src/core/lanes/learned_disk.zig b/zig/src/core/lanes/learned_disk.zig index 9ca5fbd4c0..d6e275fae1 100644 --- a/zig/src/core/lanes/learned_disk.zig +++ b/zig/src/core/lanes/learned_disk.zig @@ -1,5 +1,6 @@ //! Reserve learned-state writes above a free-disk floor, with bounded backoff after refusals or write failures. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Count = if (builtin.os.tag == .linux) u64 else u32; const Stat = extern struct { @@ -19,14 +20,15 @@ const Stat = extern struct { extern "c" fn statvfs(path: [*:0]const u8, result: *Stat) c_int; pub fn available(dir: [:0]const u8) ?u64 { + if (builtin.os.tag == .windows) return platform.c.diskAvailable(dir); var stat: Stat = undefined; if (statvfs(dir, &stat) != 0) return null; return std.math.mul(u64, stat.frsize, stat.bavail) catch null; } fn now() ?u64 { - var t: std.c.timespec = undefined; - if (std.c.clock_gettime(std.c.CLOCK.MONOTONIC, &t) != 0 or t.sec < 0) return null; + var t: platform.c.timespec = undefined; + if (platform.c.clock_gettime(platform.c.CLOCK.MONOTONIC, &t) != 0 or t.sec < 0) return null; return @intCast(t.sec); } diff --git a/zig/src/core/learned_fault_test.zig b/zig/src/core/learned_fault_test.zig index fefd560d99..98d9c9209b 100644 --- a/zig/src/core/learned_fault_test.zig +++ b/zig/src/core/learned_fault_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const Imprint = @import("prompt_imprint.zig").Imprint; const Fake = @import("prompt_cache_fake.zig").Fake; @@ -14,14 +15,14 @@ const Disk = struct { } }; fn root(buf: []u8, suffix: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-learn-fault-{s}-{d}", .{ suffix, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-learn-fault-{s}-{d}", .{ suffix, platform.c.getpid() }); } fn put(dir: []const u8, name: []const u8, count: usize) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.Create; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(count)) != 0) return error.Create; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(count)) != 0) return error.Create; } pub fn learnOne(im: *Imprint, prompt: []const u32) !void { @@ -61,9 +62,9 @@ pub fn recallOne(im: *Imprint, prompt: []const u32) !void { fn appendIndex(dir: []const u8, bytes: []const u8) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/index", .{dir}, 0), .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/index", .{dir}, 0), .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try @import("prompt_imprint.zig").writeAll(fd, bytes); } @@ -125,9 +126,9 @@ test "invalid complete records salvage the valid prefix but unknown versions rem } else |err| try std.testing.expectEqual(error.ImprintRead, err); try std.testing.expectEqual(@as(u64, 104), try im.indexScratchBytes()); var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var actual: [44]u8 = undefined; try std.testing.expect(@import("prompt_imprint.zig").readAt(fd, &actual, 60)); try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&record), &actual); @@ -206,9 +207,9 @@ test "short future-version prefixes preserve index bytes and payloads on open an defer im.deinit(); if (prefix) try learnOne(&im, &first); var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try @import("prompt_imprint.zig").writeAll(fd, std.mem.sliceAsBytes(&future)[0..cut]); var expected: [96]u8 = undefined; const len = (if (prefix) @as(usize, 60) else 0) + cut; @@ -225,13 +226,13 @@ test "short future-version prefixes preserve index bytes and payloads on open an var actual: [96]u8 = undefined; try std.testing.expect(@import("prompt_imprint.zig").readAt(fd, actual[0..len], 0)); try std.testing.expectEqualSlices(u8, expected[0..len], actual[0..len]); - const payload = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/00000000000000af.bin", .{im.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const payload = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/00000000000000af.bin", .{im.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (payload < 0) return error.FuturePayloadRemoved; - _ = std.c.close(payload); + _ = platform.c.close(payload); if (prefix) { - const kept = std.c.open(try Fake.file(&path, im.dir, Imprint.keyOf(first[0..6])), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const kept = platform.c.open(try Fake.file(&path, im.dir, Imprint.keyOf(first[0..6])), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (kept < 0) return error.ValidPrefixPayloadRemoved; - _ = std.c.close(kept); + _ = platform.c.close(kept); } } }; @@ -272,14 +273,14 @@ test "an index replacement failure retains its metadata and removes its partial var idx_buf: [1200]u8 = undefined; const idx = try std.fmt.bufPrintSentinel(&idx_buf, "{s}/index", .{im.dir}, 0); try dirs.unlink(idx); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(idx, 0o700)); - defer _ = std.c.rmdir(idx); // the tree removal goes two levels down + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(idx, 0o700)); + defer _ = platform.c.rmdir(idx); // the tree removal goes two levels down try std.testing.expectError(error.ImprintWrite, im.remove(Imprint.keyOf(&.{ 1, 2 }))); try std.testing.expect(im.has(Imprint.keyOf(&.{ 1, 2 }))); const part = try std.fmt.bufPrintSentinel(&path, "{s}/index.part", .{im.dir}, 0); - const fd = std.c.open(part, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(part, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd >= 0) { - _ = std.c.close(fd); + _ = platform.c.close(fd); return error.PartRemains; } } @@ -315,8 +316,8 @@ test "unlink failure is observable and missing files are idempotently removed" { const r = try root(&buf, "unlink"); var path: [150]u8 = undefined; const dir = try std.fmt.bufPrintSentinel(&path, "{s}", .{r}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); try std.testing.expectError(error.LearnedUnlink, dirs.unlink(dir)); var absent: [180]u8 = undefined; try dirs.unlink(try std.fmt.bufPrintSentinel(&absent, "{s}/absent", .{dir}, 0)); diff --git a/zig/src/core/learned_index_orphan_test.zig b/zig/src/core/learned_index_orphan_test.zig index 7b361f0907..6b210bb469 100644 --- a/zig/src/core/learned_index_orphan_test.zig +++ b/zig/src/core/learned_index_orphan_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const imprint = @import("prompt_imprint.zig"); const Imprint = imprint.Imprint; const Fake = @import("prompt_cache_fake.zig").Fake; @@ -26,15 +27,15 @@ fn learn(im: *Imprint, prompt: []const u32) !void { try std.testing.expect(im.has(Imprint.keyOf(prompt[0..6]))); } fn exists(file: [:0]const u8) bool { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return false; - _ = std.c.close(fd); + _ = platform.c.close(fd); return true; } fn put(file: [:0]const u8, bytes: []const u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(file, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imprint.writeAll(fd, bytes); } @@ -42,7 +43,7 @@ test "reopen sweeps only unindexed owned payloads left by old post-tear learning const a = std.testing.allocator; for (0..3) |kind| { var buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-lost-{d}-{d}", .{ std.c.getpid(), kind }); + const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-lost-{d}-{d}", .{ platform.c.getpid(), kind }); defer @import("prompt_imprint_test.zig").rmTree(root); const first = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8 }; var later = [_]u32{ 11, 12, 13, 14, 15, 16, 17, 18 }; @@ -57,14 +58,14 @@ test "reopen sweeps only unindexed owned payloads left by old post-tear learning later[0] += 10; } var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var original: [240]u8 = undefined; try std.testing.expect(imprint.readAt(fd, &original, 0)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 60)); - if (std.c.lseek(fd, 60, std.c.SEEK.SET) != 60) return error.TestFile; + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 60)); + if (platform.c.lseek(fd, 60, platform.c.SEEK.SET) != 60) return error.TestFile; const key = Imprint.keyOf(&.{ 31, 32 }); const torn = [_]u32{ 0x494d5052, 2, @truncate(key), @truncate(key >> 32), 1, if (kind == 0) 512 else 2, 0, 100, 0 }; if (kind != 1) try imprint.writeAll(fd, std.mem.sliceAsBytes(&torn)); diff --git a/zig/src/core/learned_part_test.zig b/zig/src/core/learned_part_test.zig index c013e5de1b..a2984e345d 100644 --- a/zig/src/core/learned_part_test.zig +++ b/zig/src/core/learned_part_test.zig @@ -1,22 +1,23 @@ const std = @import("std"); +const platform = @import("sys"); const Imprint = @import("prompt_imprint.zig").Imprint; const life = @import("learned_fault_test.zig"); fn scratch(buf: []u8, tag: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-learn-part-{s}-{d}", .{ tag, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-learn-part-{s}-{d}", .{ tag, platform.c.getpid() }); } fn put(dir: []const u8, name: []const u8, count: usize) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(count)) != 0) return error.TestFile; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(count)) != 0) return error.TestFile; } fn exists(dir: []const u8, name: []const u8) bool { var buf: [1200]u8 = undefined; - const fd = std.c.open(std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0) catch return false, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0) catch return false, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return false; - _ = std.c.close(fd); + _ = platform.c.close(fd); return true; } @@ -84,6 +85,7 @@ test "orphan temporary pressure cannot evict an older valid learned state" { } test "owned-name links and directories stay untouched and foreign directories are not traversed" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "types"); @@ -95,33 +97,34 @@ test "owned-name links and directories stay untouched and foreign directories ar const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}/target", .{root}, 0); var link_buf: [1200]u8 = undefined; const link = try std.fmt.bufPrintSentinel(&link_buf, "{s}/00000000000000af.bin.part", .{im.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, link)); - defer _ = std.c.unlink(link); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, link)); + defer _ = platform.c.unlink(link); var dir_buf: [1200]u8 = undefined; const dir = try std.fmt.bufPrintSentinel(&dir_buf, "{s}/index.part", .{im.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); var foreign_buf: [1200]u8 = undefined; const foreign = try std.fmt.bufPrintSentinel(&foreign_buf, "{s}/foreign", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(foreign, 0o700)); - defer _ = std.c.rmdir(foreign); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(foreign, 0o700)); + defer _ = platform.c.rmdir(foreign); try put(foreign, "index.part", 3); var reopened = try Imprint.open(a, root, 1, 1 << 20); defer reopened.deinit(); var actual: [1200]u8 = undefined; - const n = std.c.readlink(link, &actual, actual.len); + const n = platform.c.readlink(link, &actual, actual.len); try std.testing.expectEqual(@as(isize, @intCast(target.len)), n); try std.testing.expectEqualSlices(u8, target, actual[0..@intCast(n)]); - const directory = std.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const directory = platform.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); try std.testing.expect(directory >= 0); - _ = std.c.close(directory); + _ = platform.c.close(directory); try std.testing.expect(exists(root, "target")); try std.testing.expect(exists(foreign, "index.part")); var marker_buf: [1400]u8 = undefined; - _ = std.c.unlink(try std.fmt.bufPrintSentinel(&marker_buf, "{s}/index.part", .{foreign}, 0)); + _ = platform.c.unlink(try std.fmt.bufPrintSentinel(&marker_buf, "{s}/index.part", .{foreign}, 0)); } test "an identity-directory symlink refuses open and preserves foreign payloads" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "identity-link"); @@ -132,18 +135,19 @@ test "an identity-directory symlink refuses open and preserves foreign payloads" } var target_buf: [1200]u8 = undefined; const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}/foreign", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(target, 0o700)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(target, 0o700)); try put(target, "00000000000000af.bin", 7); var link_buf: [1200]u8 = undefined; const link = try std.fmt.bufPrintSentinel(&link_buf, "{s}/0000000000000001", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, link)); - defer _ = std.c.unlink(link); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, link)); + defer _ = platform.c.unlink(link); try std.testing.expectError(error.ImprintRead, Imprint.open(a, root, 1, 1 << 20)); try std.testing.expect(exists(target, "00000000000000af.bin")); try std.testing.expect(!exists(target, "used")); } test "an explicitly configured root alias still cleans only its real identity children" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "root-target"); @@ -154,11 +158,11 @@ test "an explicitly configured root alias still cleans only its real identity ch try put(seeded.dir, "index.part", 4096); } var alias_buf: [128]u8 = undefined; - const alias = try std.fmt.bufPrintSentinel(&alias_buf, "/tmp/tf-learn-root-alias-{d}", .{std.c.getpid()}, 0); + const alias = try std.fmt.bufPrintSentinel(&alias_buf, "/tmp/tf-learn-root-alias-{d}", .{platform.c.getpid()}, 0); var target_buf: [128]u8 = undefined; const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, alias)); - defer _ = std.c.unlink(alias); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, alias)); + defer _ = platform.c.unlink(alias); var reopened = try Imprint.open(a, alias, 1, 1 << 20); defer reopened.deinit(); try std.testing.expect(!exists(reopened.dir, "index.part")); diff --git a/zig/src/core/paced_read.zig b/zig/src/core/paced_read.zig index b23dd15c4a..0b72289b1f 100644 --- a/zig/src/core/paced_read.zig +++ b/zig/src/core/paced_read.zig @@ -1,5 +1,6 @@ //! Checkpoint reads that keep out of the page cache and wait for free memory: a load never outruns the kernel's reclaim. const std = @import("std"); +const platform = @import("sys"); /// Offsets, lengths and buffers of uncached reads align to the VM page (F_NOCACHE only skips the cache for aligned reads). pub const alignment = 16384; @@ -16,17 +17,17 @@ pub fn stage(gpa: std.mem.Allocator, bytes: usize) !Stage { pub fn floor() u64 { var mem: u64 = 0; var len: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; + if (platform.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; return @max(8 << 30, mem / 100 * 3); } pub fn freeBytes() u64 { var pages: u32 = 0; var len: usize = @sizeOf(u32); - if (std.c.sysctlbyname("vm.page_free_count", &pages, &len, null, 0) != 0) return std.math.maxInt(u64); + if (platform.c.sysctlbyname("vm.page_free_count", &pages, &len, null, 0) != 0) return std.math.maxInt(u64); var size: u64 = 0; var len2: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.pagesize", &size, &len2, null, 0) != 0) size = alignment; + if (platform.c.sysctlbyname("hw.pagesize", &size, &len2, null, 0) != 0) size = alignment; return @as(u64, pages) * size; } @@ -35,13 +36,13 @@ pub fn waitForRoom(limit_ms: u32) void { const want = floor(); var waited: u32 = 0; while (waited < limit_ms and freeBytes() < want) : (waited += 2) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 2 * std.time.ns_per_ms }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 2 * std.time.ns_per_ms }; + _ = platform.c.nanosleep(&ts, null); } } /// `dest.len` bytes at `at` of `fd` (opened with F_NOCACHE) through page-aligned reads into `buf`, pausing for free memory. -pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { +pub fn read(fd: platform.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { var done: usize = 0; while (done < dest.len) { waitForRoom(10_000); @@ -52,7 +53,7 @@ pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { const span = std.mem.alignForward(usize, skip + want, alignment); var got: usize = 0; while (got < skip + want) { // a read past the file's end comes back short once the bytes asked for are in - const n = std.c.pread(fd, buf.ptr + got, span - got, @intCast(lo + got)); + const n = platform.c.pread(fd, buf.ptr + got, span - got, @intCast(lo + got)); if (n <= 0) return error.ShortRead; got += @intCast(n); } @@ -64,17 +65,17 @@ pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { test "aligned spans copy the bytes asked for, at any offset and length" { const gpa = std.testing.allocator; var name_buf: [64]u8 = undefined; - const name = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-paced-read-{d}", .{std.c.getpid()}, 0); - const fd = std.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const name = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-paced-read-{d}", .{platform.c.getpid()}, 0); + const fd = platform.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); try std.testing.expect(fd >= 0); defer { - _ = std.c.close(fd); - _ = std.c.unlink(name); + _ = platform.c.close(fd); + _ = platform.c.unlink(name); } const data = try gpa.alloc(u8, 3 * alignment + 777); defer gpa.free(data); for (data, 0..) |*b, i| b.* = @truncate(i *% 131 +% 7); - try std.testing.expectEqual(@as(isize, @intCast(data.len)), std.c.write(fd, data.ptr, data.len)); + try std.testing.expectEqual(@as(isize, @intCast(data.len)), platform.c.write(fd, data.ptr, data.len)); const buf = try stage(gpa, 2 * alignment); defer gpa.free(buf); for ([_][2]usize{ .{ 0, 1 }, .{ 5, 100 }, .{ alignment - 3, 10 }, .{ 1, 2 * alignment + 50 }, .{ 3 * alignment, 777 }, .{ 9, data.len - 9 } }) |c| { diff --git a/zig/src/core/pair.zig b/zig/src/core/pair.zig index b7d6daed68..d6fc64cd68 100644 --- a/zig/src/core/pair.zig +++ b/zig/src/core/pair.zig @@ -1,5 +1,6 @@ //! Two Macs' tensor-parallel exchanges in a command buffer over MCDMA, and rank 0's commands to rank 1. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const sources = @import("kernel_sources"); @@ -54,7 +55,7 @@ threadlocal var raised = false; /// This thread at the user-interactive QoS once: busy CPUs must not delay the pair's sends and commands. pub fn raise() void { if (raised) return; - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); raised = true; } @@ -230,16 +231,16 @@ pub const Pair = struct { /// Wait until this Mac's word at `off` reaches `value`: spins, then short sleeps (`bounded`: 60 s at most). fn reach(p: *Pair, off: usize, value: u64, bounded: bool) !void { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); var spins: usize = 0; while (@atomicLoad(u64, p.word64(off), .acquire) < value) { if (p.stop.load(.acquire)) return error.PairStopped; if (p.failed.load(.acquire)) return error.PairLinkFailed; - if (bounded and std.c.mach_absolute_time() - t0 > 60 * 24_000_000) return error.PairPeerSilent; + if (bounded and platform.c.mach_absolute_time() - t0 > 60 * 24_000_000) return error.PairPeerSilent; spins += 1; if (spins > 200_000) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } } @@ -254,8 +255,8 @@ pub const Pair = struct { if (!reached(posted, next_x) or !reached(p.queued.load(.acquire), next_x)) { idle += 1; if (idle > 1_000_000) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); continue; } diff --git a/zig/src/core/prompt_cache_fake.zig b/zig/src/core/prompt_cache_fake.zig index 6f6754d0b6..734ecafae9 100644 --- a/zig/src/core/prompt_cache_fake.zig +++ b/zig/src/core/prompt_cache_fake.zig @@ -1,5 +1,6 @@ //! A host-memory snapshot family shared by the prompt cache's tests and the learned-disk fault tests. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const imprint = @import("prompt_imprint.zig"); const Snapshots = pc.Snapshots; @@ -96,7 +97,7 @@ pub const Fake = struct { } var path: [512]u8 = undefined; - _ = std.c.unlink(file(&path, dir, key) catch return); + _ = platform.c.unlink(file(&path, dir, key) catch return); } pub fn file(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { return std.fmt.bufPrintSentinel(buf, "{s}/{x:0>16}.bin", .{ dir, key }, 0); @@ -106,17 +107,17 @@ pub const Fake = struct { f.writes += 1; if (f.fail_write) return error.WriteFailed; var path: [512]u8 = undefined; - const fd = std.c.open(try file(&path, dir, key), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try file(&path, dir, key), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.WriteFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imprint.writeAll(fd, std.mem.asBytes(@as(*State, @ptrCast(@alignCast(saved))))); } fn readFn(ptr: *anyopaque, dir: [:0]const u8, key: u64, at: u32) anyerror!Saved { const f = of(ptr); var path: [512]u8 = undefined; - const fd = std.c.open(try file(&path, dir, key), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try file(&path, dir, key), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.ReadFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const st = try f.gpa.create(State); errdefer f.gpa.destroy(st); if (!imprint.readAt(fd, std.mem.asBytes(st), 0) or st.at != at) return error.ReadFailed; diff --git a/zig/src/core/prompt_cache_learn_test.zig b/zig/src/core/prompt_cache_learn_test.zig index 73eb857f2a..26f9f3f28e 100644 --- a/zig/src/core/prompt_cache_learn_test.zig +++ b/zig/src/core/prompt_cache_learn_test.zig @@ -1,5 +1,6 @@ //! The prompt cache's learned-state tests: harness states kept on disk, the cap, the disk floor and peers' halves. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const imprint = @import("prompt_imprint.zig"); const rmTree = @import("prompt_imprint_test.zig").rmTree; @@ -14,7 +15,7 @@ test "a learned harness state outlives its store: a fresh session on a new one r defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-test-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-test-{d}", .{platform.c.getpid()}); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const harness = [_]u32{ 7, 7, 7, 7, 7, 7 }; // shared by every session: a state at 5 reads it all (lookahead 1) const key = imprint.Imprint.keyOf(&harness); @@ -50,7 +51,7 @@ test "past the learned-state cap the least recently used state is forgotten, and defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-cap-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-cap-{d}", .{platform.c.getpid()}); defer rmTree(root); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const one = [_]u32{ 7, 7, 7, 7, 7, 7 }; @@ -67,8 +68,8 @@ test "past the learned-state cap the least recently used state is forgotten, and try std.testing.expectEqual(fresh(&p2), f.pass(&s, &p2, try s.begin(a, &p2, 7, &.{5}, &.{}, null, &.{}))); try std.testing.expect(im.has(imprint.Imprint.keyOf(&two)) and !im.has(imprint.Imprint.keyOf(&one))); var path: [512]u8 = undefined; - try std.testing.expect(std.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&one))) != 0); // its file went too - _ = std.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&two))); // a file lost behind the store's back + try std.testing.expect(platform.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&one))) != 0); // its file went too + _ = platform.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&two))); // a file lost behind the store's back var g: Fake = .{ .gpa = gpa }; var t = Store.init(gpa, g.learned(), rules, 1 << 20); defer t.deinit(); @@ -87,7 +88,7 @@ test "learned states are prompt arithmetic: a request that decodes rows below on defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-spans-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-spans-{d}", .{platform.c.getpid()}); defer rmTree(root); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const harness = [_]u32{ 7, 7, 7, 7, 7, 7 }; @@ -138,7 +139,7 @@ test "Store learning preserves the disk floor, backs off writes across keys, and var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-backoff-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-backoff-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 11, 1 << 20); defer im.deinit(); @@ -180,7 +181,7 @@ test "learning reserves the existing index rewrite above the disk floor and resu var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-index-floor-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-index-floor-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 12, 1 << 20); defer im.deinit(); @@ -213,7 +214,7 @@ test "actual Store learn writes neither half on peer refusal, backs off, and rem }; const a = std.testing.allocator; var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-refuse-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-refuse-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 51, 1 << 20); defer im.deinit(); @@ -242,9 +243,9 @@ test "actual Store learn writes neither half on peer refusal, backs off, and rem try std.testing.expectEqual(@as(u64, 0), im.admission.reserved); var file_path: [512]u8 = undefined; const file = try Fake.file(&file_path, im.dir, imprint.Imprint.keyOf(prompt[0..6])); - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd >= 0) { - _ = std.c.close(fd); + _ = platform.c.close(fd); return error.OrphanedHalf; } } @@ -263,16 +264,16 @@ test "under-cap disk pressure selects the same LRU victim before either half is }; const a = std.testing.allocator; var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-victim-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-victim-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 52, 1 << 20); defer im.deinit(); try im.add(imprint.Imprint.keyOf(&.{ 1, 2 }), 1, &.{ 1, 2 }, &.{}, 16); var victim_path: [512]u8 = undefined; - const victim_fd = std.c.open(try Fake.file(&victim_path, im.dir, imprint.Imprint.keyOf(&.{ 1, 2 })), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const victim_fd = platform.c.open(try Fake.file(&victim_path, im.dir, imprint.Imprint.keyOf(&.{ 1, 2 })), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (victim_fd < 0) return error.Create; - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(victim_fd, 16)); - _ = std.c.close(victim_fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(victim_fd, 16)); + _ = platform.c.close(victim_fd); im.admission = .{ .floor = 100, .free = Disk.free, .clock = Disk.clock }; const prompt = [_]u32{ 7, 7, 7, 7, 7, 7, 9, 10 }; Disk.space = 100 + 105 + 256 + (6 * 4) + (try im.indexScratchBytes()) - 8; diff --git a/zig/src/core/prompt_imprint.zig b/zig/src/core/prompt_imprint.zig index 8fee9da5ad..024a8a9b6a 100644 --- a/zig/src/core/prompt_imprint.zig +++ b/zig/src/core/prompt_imprint.zig @@ -1,5 +1,6 @@ //! Prompt-cache states on disk (Imprint): a learned harness state outlives the server and serves later fresh sessions. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; /// One state on disk: the tokens it read (position plus lookahead), the chunk starts below it, and its bytes. @@ -74,10 +75,10 @@ pub const Imprint = struct { /// Reserve the current index's size too: LRU eviction rewrites it through a temporary file. pub fn indexScratchBytes(m: *const Imprint) !u64 { var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.ImprintRead; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.ImprintRead; + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.ImprintRead; return @intCast(end); } @@ -119,10 +120,10 @@ pub const Imprint = struct { pub fn pendingWrite(m: *const Imprint) !?u64 { var buf: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/pending", .{m.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) null else error.ImprintRead; - defer _ = std.c.close(fd); - if (std.c.lseek(fd, 0, std.c.SEEK.END) != 48) return error.ImprintRead; + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/pending", .{m.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) null else error.ImprintRead; + defer _ = platform.c.close(fd); + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != 48) return error.ImprintRead; var data: [48]u8 = undefined; if (!readAt(fd, &data, 0)) return error.ImprintRead; var digest: [32]u8 = undefined; @@ -135,19 +136,19 @@ pub const Imprint = struct { var buf: [1100]u8 = undefined; var dest: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&buf, "{s}/pending.part", .{m.dir}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var data: [48]u8 = undefined; std.mem.writeInt(u64, data[0..8], 0x494d5052494e5431, .little); std.mem.writeInt(u64, data[8..16], key, .little); std.crypto.hash.sha2.Sha256.hash(data[0..16], data[16..48], .{}); try writeAll(fd, &data); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; } - if (std.c.rename(tmp, try std.fmt.bufPrintSentinel(&dest, "{s}/pending", .{m.dir}, 0)) != 0) return error.ImprintWrite; + if (platform.c.rename(tmp, try std.fmt.bufPrintSentinel(&dest, "{s}/pending", .{m.dir}, 0)) != 0) return error.ImprintWrite; } pub fn clearPending(m: *const Imprint) !void { var buf: [1100]u8 = undefined; @@ -159,16 +160,16 @@ pub const Imprint = struct { if (m.has(key)) return; const below = cut(starts, at); var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true, .NOFOLLOW = true, .NONBLOCK = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true, .NOFOLLOW = true, .NONBLOCK = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); _ = indexSize(fd) catch return error.ImprintWrite; - var end = std.c.lseek(fd, 0, std.c.SEEK.END); + var end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.ImprintWrite; if (end != m.index_end) { try m.reload(fd); if (m.has(key)) return; - end = std.c.lseek(fd, 0, std.c.SEEK.END); + end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0 or end != m.index_end) return error.ImprintWrite; } const t = try m.gpa.dupe(u32, tokens); @@ -176,9 +177,9 @@ pub const Imprint = struct { const s = try m.gpa.dupe(u32, below); errdefer m.gpa.free(s); try m.metas.ensureUnusedCapacity(m.gpa, 1); - errdefer _ = std.c.ftruncate(fd, end); + errdefer _ = platform.c.ftruncate(fd, end); try record(fd, key, at, tokens, below, bytes); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; m.index_end = @as(u64, @intCast(end)) + HEAD * 4 + 4 * (tokens.len + below.len); m.clock += 1; m.metas.appendAssumeCapacity(.{ .key = key, .at = at, .tokens = t, .starts = s, .bytes = bytes, .used = m.clock }); @@ -190,15 +191,15 @@ pub const Imprint = struct { var path: [1100]u8 = undefined; var part: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&part, "{s}/index.part", .{m.dir}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); for (m.metas.items, 0..) |x, j| if (j != i) try record(fd, x.key, x.at, x.tokens, x.starts, x.bytes); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; } - if (std.c.rename(tmp, try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0)) != 0) return error.ImprintWrite; + if (platform.c.rename(tmp, try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0)) != 0) return error.ImprintWrite; free(m.gpa, m.metas.orderedRemove(i)); m.index_end = 0; for (m.metas.items) |x| m.index_end += HEAD * 4 + 4 * (x.tokens.len + x.starts.len); @@ -207,9 +208,9 @@ pub const Imprint = struct { /// Keep a valid prefix across torn or invalid suffixes; unknown versions and I/O errors refuse learning. fn load(m: *Imprint) !void { var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .NONBLOCK = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) {} else error.ImprintRead; - defer _ = std.c.close(fd); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .NONBLOCK = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) {} else error.ImprintRead; + defer _ = platform.c.close(fd); try m.readIndex(fd); } @@ -300,20 +301,20 @@ fn indexSize(fd: c_int) !u64 { if (@import("builtin").os.tag == .linux) { const linux = std.os.linux; var sx: linux.Statx = undefined; - if (std.c.statx(fd, "", linux.AT.EMPTY_PATH, .{ .TYPE = true, .SIZE = true }, &sx) != 0) return error.ImprintRead; - if (!sx.mask.TYPE or !sx.mask.SIZE or @as(u32, sx.mode) & std.c.S.IFMT != std.c.S.IFREG) return error.ImprintRead; + if (platform.c.statx(fd, "", linux.AT.EMPTY_PATH, .{ .TYPE = true, .SIZE = true }, &sx) != 0) return error.ImprintRead; + if (!sx.mask.TYPE or !sx.mask.SIZE or @as(u32, sx.mode) & platform.c.S.IFMT != platform.c.S.IFREG) return error.ImprintRead; if (sx.size > @as(u64, std.math.maxInt(i64))) return error.ImprintRead; return sx.size; } else { - var stat: std.c.Stat = undefined; - if (std.c.fstat(fd, &stat) != 0 or stat.mode & std.c.S.IFMT != std.c.S.IFREG or stat.size < 0) return error.ImprintRead; + var stat: platform.c.Stat = undefined; + if (platform.c.fstat(fd, &stat) != 0 or stat.mode & platform.c.S.IFMT != platform.c.S.IFREG or stat.size < 0) return error.ImprintRead; return @intCast(stat.size); } } fn repairTail(fd: c_int, whole: u64, observed: u64) !void { - if (std.c.lseek(fd, 0, std.c.SEEK.END) != observed) return error.ImprintRead; - if (std.c.ftruncate(fd, @intCast(whole)) != 0 or std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != observed) return error.ImprintRead; + if (platform.c.ftruncate(fd, @intCast(whole)) != 0 or platform.c.fsync(fd) != 0) return error.ImprintWrite; } fn ordered(starts: []const u32) bool { @@ -358,9 +359,9 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { var dirs: std.ArrayList(Dir) = .empty; defer dirs.deinit(gpa); var buf: [1100]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return 0) orelse return 0; + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return 0) orelse return 0; var held: u64 = 0; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (!identityName(name) or std.mem.endsWith(u8, keep, name)) continue; var sub: [1100]u8 = undefined; @@ -370,7 +371,7 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { held += x.bytes; dirs.append(gpa, x) catch break; } - _ = std.c.closedir(d); + _ = platform.c.closedir(d); std.mem.sort(Dir, dirs.items, {}, struct { fn older(_: void, a: Dir, b: Dir) bool { return a.used < b.used; @@ -381,7 +382,7 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { var sub: [1100]u8 = undefined; const path = std.fmt.bufPrintSentinel(&sub, "{s}/{s}", .{ root, &x.name }, 0) catch continue; _ = files(path, true); - if (std.c.rmdir(path) == 0) held -|= x.bytes; + if (platform.c.rmdir(path) == 0) held -|= x.bytes; } return held; } @@ -389,27 +390,27 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { /// The bytes of the regular files directly in `dir`, each unlinked when `remove`. fn files(dir: []const u8, remove: bool) u64 { var buf: [1100]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{dir}, 0) catch return 0) orelse return 0; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{dir}, 0) catch return 0) orelse return 0; + defer _ = platform.c.closedir(d); var n: u64 = 0; - while (std.c.readdir(d)) |ent| { - if (ent.type != std.c.DT.REG) continue; + while (platform.c.readdir(d)) |ent| { + if (ent.type != platform.c.DT.REG) continue; var p: [1400]u8 = undefined; const f = std.fmt.bufPrintSentinel(&p, "{s}/{s}", .{ dir, std.mem.sliceTo(&ent.name, 0) }, 0) catch continue; - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) continue; - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - _ = std.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); + _ = platform.c.close(fd); if (end > 0) n += @intCast(end); - if (remove) _ = std.c.unlink(f); + if (remove) _ = platform.c.unlink(f); } return n; } fn readWord(path: [:0]const u8) u64 { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return 0; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var w: u64 = 0; return if (readAt(fd, std.mem.asBytes(&w), 0)) w else 0; } @@ -417,19 +418,19 @@ fn readWord(path: [:0]const u8) u64 { fn writeWord(path: [:0]const u8, w: u64) void { var buf: [1200]u8 = undefined; const part = std.fmt.bufPrintSentinel(&buf, "{s}.part", .{path}, 0) catch return; - const fd = std.c.open(part, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(part, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return; - defer _ = std.c.close(fd); - defer _ = std.c.unlink(part); + defer _ = platform.c.close(fd); + defer _ = platform.c.unlink(part); writeAll(fd, std.mem.asBytes(&w)) catch return; - if (std.c.fsync(fd) != 0) return; - _ = std.c.rename(part, path); + if (platform.c.fsync(fd) != 0) return; + _ = platform.c.rename(part, path); } pub fn writeAll(fd: c_int, bytes: []const u8) !void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes.ptr + done, bytes.len - done); + const n = platform.c.write(fd, bytes.ptr + done, bytes.len - done); if (n <= 0) return error.ImprintWrite; done += @intCast(n); } @@ -438,7 +439,7 @@ pub fn writeAll(fd: c_int, bytes: []const u8) !void { pub fn readAt(fd: c_int, dest: []u8, at: u64) bool { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return false; done += @intCast(n); } @@ -448,13 +449,13 @@ pub fn readAt(fd: c_int, dest: []u8, at: u64) bool { /// This Mac's OS build (sysctl kern.osversion): its Metal compiler can change a kernel's bits. pub fn osBuild(buf: []u8) []const u8 { var n: usize = buf.len; - if (std.c.sysctlbyname("kern.osversion", buf.ptr, &n, null, 0) != 0) return ""; + if (platform.c.sysctlbyname("kern.osversion", buf.ptr, &n, null, 0) != 0) return ""; return std.mem.sliceTo(buf[0..n], 0); } /// Where learned states live unless --learn-dir says: $HOME/.cache/tensorfold/learned. pub fn defaultRoot(a: Allocator) ![]const u8 { - const home = std.c.getenv("HOME") orelse return error.ImprintDir; + const home = platform.c.getenv("HOME") orelse return error.ImprintDir; return std.fmt.allocPrint(a, "{s}/.cache/tensorfold/learned", .{std.mem.span(home)}); } @@ -464,11 +465,11 @@ fn makePath(gpa: Allocator, dir: [:0]const u8) !void { defer gpa.free(tmp); for (tmp[1..], 1..) |ch, i| if (ch == '/') { tmp[i] = 0; - _ = std.c.mkdir(tmp.ptr, 0o700); + _ = platform.c.mkdir(tmp.ptr, 0o700); tmp[i] = '/'; }; - const rc = std.c.mkdir(tmp.ptr, 0o700); - if (rc != 0 and std.c.errno(rc) != .EXIST) return error.ImprintDir; + const rc = platform.c.mkdir(tmp.ptr, 0o700); + if (rc != 0 and platform.c.errno(rc) != .EXIST) return error.ImprintDir; } test { @@ -479,19 +480,19 @@ test { test "a failed tail truncation keeps the bytes and refuses repair" { var buf: [128]u8 = undefined; - const path = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-index-repair-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(path); - const writer = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const path = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-index-repair-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(path); + const writer = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (writer < 0) return error.TestFile; { - defer _ = std.c.close(writer); + defer _ = platform.c.close(writer); try writeAll(writer, "record"); } - const reader = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const reader = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (reader < 0) return error.TestFile; - defer _ = std.c.close(reader); + defer _ = platform.c.close(reader); try std.testing.expectError(error.ImprintWrite, repairTail(reader, 0, 6)); - try std.testing.expectEqual(@as(i64, 6), std.c.lseek(reader, 0, std.c.SEEK.END)); + try std.testing.expectEqual(@as(i64, 6), platform.c.lseek(reader, 0, platform.c.SEEK.END)); var actual: [6]u8 = undefined; try std.testing.expect(readAt(reader, &actual, 0)); try std.testing.expectEqualStrings("record", &actual); diff --git a/zig/src/core/prompt_imprint_fault_test.zig b/zig/src/core/prompt_imprint_fault_test.zig index 418f69e0b5..a1d954bfd2 100644 --- a/zig/src/core/prompt_imprint_fault_test.zig +++ b/zig/src/core/prompt_imprint_fault_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const Imprint = @import("prompt_imprint.zig").Imprint; const imp = @import("prompt_imprint.zig"); const dirs = @import("lanes").learned_dirs; @@ -41,12 +42,12 @@ const Tracked = struct { } }; fn scratch(buf: []u8, suffix: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-imprint-fault-{s}-{d}", .{ suffix, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-imprint-fault-{s}-{d}", .{ suffix, platform.c.getpid() }); } fn word(path: [:0]const u8, n: u64) !void { - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imp.writeAll(fd, std.mem.asBytes(&n)); } test "Imprint failed open releases each allocation once after another identity scan" { @@ -58,8 +59,8 @@ test "Imprint failed open releases each allocation once after another identity s defer old.deinit(); var path: [256]u8 = undefined; const unreadable = try std.fmt.bufPrintSentinel(&path, "{s}/unreadable", .{old.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(unreadable, 0o700)); - defer _ = std.c.rmdir(unreadable); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(unreadable, 0o700)); + defer _ = platform.c.rmdir(unreadable); var tracked: Tracked = .{}; try std.testing.expectError(error.ImprintRead, Imprint.open(tracked.allocator(), root, 2, 1 << 20)); try std.testing.expectEqual(@as(usize, 0), tracked.duplicate_frees); @@ -92,8 +93,8 @@ test "failed atomic stamp replacement preserves the previous complete used word" var path: [256]u8 = undefined; try word(try std.fmt.bufPrintSentinel(&path, "{s}/used", .{m.dir}, 0), 42); const part = try std.fmt.bufPrintSentinel(&path, "{s}/used.part", .{m.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(part, 0o700)); - defer _ = std.c.rmdir(part); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(part, 0o700)); + defer _ = platform.c.rmdir(part); var reopened = try Imprint.open(a, root, 1, 1 << 20); defer reopened.deinit(); try std.testing.expectEqual(@as(u64, 42), try dirs.used(root, 1)); @@ -108,7 +109,7 @@ test "used stamp reads still refuse an actual I/O error" { var path: [256]u8 = undefined; const used = try std.fmt.bufPrintSentinel(&path, "{s}/used", .{m.dir}, 0); try dirs.unlink(used); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(used, 0o700)); - defer _ = std.c.rmdir(used); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(used, 0o700)); + defer _ = platform.c.rmdir(used); try std.testing.expectError(error.LearnedDirectoryRead, dirs.used(root, 1)); } diff --git a/zig/src/core/prompt_imprint_test.zig b/zig/src/core/prompt_imprint_test.zig index 66bdabd15e..f3e856e21b 100644 --- a/zig/src/core/prompt_imprint_test.zig +++ b/zig/src/core/prompt_imprint_test.zig @@ -1,5 +1,6 @@ //! Learned prompt states on disk: reopen, eviction order, removal and the sweep of other identities. const std = @import("std"); +const platform = @import("sys"); const imp = @import("prompt_imprint.zig"); const Imprint = imp.Imprint; @@ -7,38 +8,38 @@ const Imprint = imp.Imprint; pub fn rmTree(root: []const u8) void { var buf: [512]u8 = undefined; const r = std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return; - const d = std.c.opendir(r) orelse return; - while (std.c.readdir(d)) |ent| { + const d = platform.c.opendir(r) orelse return; + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (name[0] == '.') continue; var sub: [512]u8 = undefined; const p = std.fmt.bufPrintSentinel(&sub, "{s}/{s}", .{ root, name }, 0) catch continue; - if (std.c.unlink(p) == 0) continue; - if (std.c.opendir(p)) |s| { - while (std.c.readdir(s)) |f| { + if (platform.c.unlink(p) == 0) continue; + if (platform.c.opendir(p)) |s| { + while (platform.c.readdir(s)) |f| { const fname = std.mem.sliceTo(&f.name, 0); if (fname[0] == '.') continue; var fb: [600]u8 = undefined; - _ = std.c.unlink(std.fmt.bufPrintSentinel(&fb, "{s}/{s}", .{ p, fname }, 0) catch continue); + _ = platform.c.unlink(std.fmt.bufPrintSentinel(&fb, "{s}/{s}", .{ p, fname }, 0) catch continue); } - _ = std.c.closedir(s); + _ = platform.c.closedir(s); } - _ = std.c.rmdir(p); + _ = platform.c.rmdir(p); } - _ = std.c.closedir(d); - _ = std.c.rmdir(r); + _ = platform.c.closedir(d); + _ = platform.c.rmdir(r); } fn scratch(buf: []u8, tag: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-imprint-{s}-{d}", .{ tag, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-imprint-{s}-{d}", .{ tag, platform.c.getpid() }); } /// A file of `n` bytes at `dir`/`name`. fn fill(dir: []const u8, name: []const u8, n: usize) !void { var buf: [512]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var zeros: [256]u8 = @splat(0); var left = n; while (left > 0) : (left -= @min(left, zeros.len)) try imp.writeAll(fd, zeros[0..@min(left, zeros.len)]); @@ -46,8 +47,8 @@ fn fill(dir: []const u8, name: []const u8, n: usize) !void { fn exists(path: []const u8) bool { var buf: [512]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{path}, 0) catch return false) orelse return false; - _ = std.c.closedir(d); + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{path}, 0) catch return false) orelse return false; + _ = platform.c.closedir(d); return true; } diff --git a/zig/src/core/shard_edit.zig b/zig/src/core/shard_edit.zig index b340a40523..dff4f9dd5a 100644 --- a/zig/src/core/shard_edit.zig +++ b/zig/src/core/shard_edit.zig @@ -1,5 +1,6 @@ //! Replaced tensors written into a checkpoint's own shards, index and config: in place when sizes hold, else by rename. const std = @import("std"); +const platform = @import("sys"); const st = @import("safetensors.zig"); const Allocator = std.mem.Allocator; const Io = std.Io; @@ -200,18 +201,18 @@ fn inPlace(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacement) } var arena_state = std.heap.ArenaAllocator.init(gpa); defer arena_state.deinit(); - const fd = std.c.open(try arena_state.allocator().dupeSentinel(u8, path, 0), .{ .ACCMODE = .WRONLY }); + const fd = platform.c.open(try arena_state.allocator().dupeSentinel(u8, path, 0), .{ .ACCMODE = .WRONLY }); if (fd < 0) return error.ShardWrite; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); for (edits, places[0..edits.len]) |r, at| { var done: usize = 0; while (done < r.bytes.len) { - const n = std.c.pwrite(fd, r.bytes.ptr + done, @min(r.bytes.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pwrite(fd, r.bytes.ptr + done, @min(r.bytes.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.ShardWrite; done += @intCast(n); } } - if (std.c.fsync(fd) != 0) return error.ShardWrite; + if (platform.c.fsync(fd) != 0) return error.ShardWrite; return true; } @@ -221,7 +222,8 @@ pub fn rewrite(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacem defer arena_state.deinit(); const arena = arena_state.allocator(); var file = try st.File.open(gpa, io, path); - defer file.close(io); + var mapped = true; + defer if (mapped) file.close(io); for (edits) |r| if (!file.names.contains(r.name)) return error.MissingTensor; const raw = file.map.memory; const original = try std.json.parseFromSliceLeaky(Value, arena, raw[8..file.data], .{}); @@ -257,7 +259,11 @@ pub fn rewrite(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacem const spaces = try arena.alloc(u8, padded - text.len); @memset(spaces, ' '); try pieces.insertSlice(arena, 0, &.{ &head, text, spaces }); - try writeFile(arena, path, pieces.items); + const tmp = try writeTemp(arena, path, pieces.items); + // the pieces are written; the map goes before the rename, since Windows refuses to replace a mapped file + file.close(io); + mapped = false; + try commit(arena, tmp, path); return @as(i64, @intCast(offset)) - @as(i64, @intCast(before)); } @@ -325,23 +331,32 @@ fn writeJson(arena: Allocator, path: []const u8, v: Value) !void { /// `parts` to `path` through `path`.tmp, synced, then renamed over it; a reader sees the old file or the new one. fn writeFile(arena: Allocator, path: []const u8, parts: []const []const u8) !void { + try commit(arena, try writeTemp(arena, path, parts), path); +} + +/// `parts` written and synced to `path`.tmp, which is returned. +fn writeTemp(arena: Allocator, path: []const u8, parts: []const []const u8) ![:0]const u8 { const tmp = try std.fmt.allocPrintSentinel(arena, "{s}.tmp", .{path}, 0); - const dest = try arena.dupeSentinel(u8, path, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o644)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o644)); if (fd < 0) return error.ShardWrite; - errdefer _ = std.c.unlink(tmp); - { - defer _ = std.c.close(fd); - for (parts) |p| try put(fd, p); - if (std.c.fsync(fd) != 0) return error.ShardWrite; - } - if (std.c.rename(tmp, dest) != 0) return error.ShardWrite; + errdefer _ = platform.c.unlink(tmp); + defer _ = platform.c.close(fd); + for (parts) |p| try put(fd, p); + if (platform.c.fsync(fd) != 0) return error.ShardWrite; + return tmp; +} + +/// `tmp` renamed over `path`; a failed rename removes `tmp`. +fn commit(arena: Allocator, tmp: [:0]const u8, path: []const u8) !void { + errdefer _ = platform.c.unlink(tmp); + const dest = try arena.dupeSentinel(u8, path, 0); + if (platform.c.rename(tmp, dest) != 0) return error.ShardWrite; } fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.ShardWrite; done += @intCast(n); } @@ -518,13 +533,19 @@ test "a folder whose shard is a symlink or a hard link is refused for learning; defer tmp.cleanup(); try tmp.dir.createDirPath(io, "blobs"); for ([_][]const u8{ "config.json", index_name, "model-00001-of-00002.safetensors", "blobs/b2", "notes.txt" }) |name| try tmp.dir.writeFile(io, .{ .sub_path = name, .data = "x" }); - try tmp.dir.symLink(io, "model-00001-of-00002.safetensors", "readme-link", .{}); + // Windows grants symlinks only with Developer Mode or elevation + tmp.dir.symLink(io, "model-00001-of-00002.safetensors", "readme-link", .{}) catch |e| return if (e == error.PermissionDenied and @import("builtin").os.tag == .windows) error.SkipZigTest else e; const dir = try std.fs.path.join(arena, &.{ ".zig-cache", "tmp", &tmp.sub_path }); try std.testing.expect((try linked(arena, io, dir)) == null); try tmp.dir.symLink(io, "blobs/b2", "model-00002-of-00002.safetensors", .{}); try std.testing.expectEqualStrings("model-00002-of-00002.safetensors", (try linked(arena, io, dir)).?); try tmp.dir.deleteFile(io, "model-00002-of-00002.safetensors"); - try tmp.dir.hardLink("blobs/b2", tmp.dir, "model-00002-of-00002.safetensors", io, .{}); + if (@import("builtin").os.tag == .windows) { + // std's Io has no hard links on Windows; the platform's link (CreateHardLinkW) is what pull uses there + const old = try std.fs.path.joinZ(arena, &.{ dir, "blobs", "b2" }); + const new = try std.fs.path.joinZ(arena, &.{ dir, "model-00002-of-00002.safetensors" }); + try std.testing.expectEqual(@as(c_int, 0), platform.c.link(old, new)); + } else try tmp.dir.hardLink("blobs/b2", tmp.dir, "model-00002-of-00002.safetensors", io, .{}); try std.testing.expectEqualStrings("model-00002-of-00002.safetensors", (try linked(arena, io, dir)).?); } diff --git a/zig/src/core/snapshot_file.zig b/zig/src/core/snapshot_file.zig index 609916ed99..f56997b862 100644 --- a/zig/src/core/snapshot_file.zig +++ b/zig/src/core/snapshot_file.zig @@ -1,5 +1,6 @@ //! Learned snapshots carry a SHA-256 of their entire header and payload; exact length is required on admission. const std = @import("std"); +const platform = @import("sys"); const Hash = std.crypto.hash.sha2.Sha256; @@ -15,7 +16,7 @@ pub fn write(fd: c_int, header: []const u8, payload: []const u8) !void { pub fn read(fd: c_int, header: []const u8, payload: []u8) !void { const length = std.math.add(usize, header.len, payload.len) catch return error.SnapshotRead; const total = std.math.add(usize, length, 32) catch return error.SnapshotRead; - if (std.c.lseek(fd, 0, std.c.SEEK.END) != total) return error.SnapshotRead; + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != total) return error.SnapshotRead; if (!get(fd, payload, header.len)) return error.SnapshotRead; var stored: [32]u8 = undefined; if (!get(fd, &stored, length)) return error.SnapshotRead; @@ -27,21 +28,21 @@ pub fn read(fd: c_int, header: []const u8, payload: []u8) !void { test "payload bit flips, header changes, truncation, legacy files and extra bytes are refused" { var path: [128]u8 = undefined; - const name = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-snapshot-test-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(name); - const fd = std.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const name = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-snapshot-test-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(name); + const fd = platform.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var output: [4]u8 = undefined; try write(fd, "head", "body"); try read(fd, "head", &output); try std.testing.expectEqualStrings("body", &output); try std.testing.expectError(error.SnapshotRead, read(fd, "dead", &output)); - try std.testing.expectEqual(@as(isize, 1), std.c.pwrite(fd, "X", 1, 5)); + try std.testing.expectEqual(@as(isize, 1), platform.c.pwrite(fd, "X", 1, 5)); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 8)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 8)); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); - _ = std.c.lseek(fd, 0, std.c.SEEK.SET); + _ = platform.c.lseek(fd, 0, platform.c.SEEK.SET); try write(fd, "head", "body"); try put(fd, "x"); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); @@ -50,7 +51,7 @@ test "payload bit flips, header changes, truncation, legacy files and extra byte fn put(fd: c_int, bytes: []const u8) !void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30)); + const n = platform.c.write(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -58,7 +59,7 @@ fn put(fd: c_int, bytes: []const u8) !void { fn get(fd: c_int, bytes: []u8, at: usize) bool { var done: usize = 0; while (done < bytes.len) { - const n = std.c.pread(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return false; done += @intCast(n); } diff --git a/zig/src/core/template/methods.zig b/zig/src/core/template/methods.zig index 2ec14e09e1..ecccb1497e 100644 --- a/zig/src/core/template/methods.zig +++ b/zig/src/core/template/methods.zig @@ -1,5 +1,6 @@ //! Python methods templates call on strings, dicts and lists, the loop helpers, and the template globals. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const v = @import("value.zig"); const uni = @import("unicode.zig"); @@ -432,14 +433,22 @@ pub fn callFunc(rt: *Rt, f: v.Func, pos: []const Value, kw: []const Kw) Error!Va } const Tm = extern struct { sec: c_int, min: c_int, hour: c_int, mday: c_int, mon: c_int, year: c_int, wday: c_int, yday: c_int, isdst: c_int, gmtoff: c_long, zone: ?[*:0]const u8 }; -extern "c" fn localtime_r(t: *const std.c.time_t, out: *Tm) ?*Tm; +const localtime_r = if (builtin.os.tag == .windows) localtimeWindows else struct { + extern "c" fn localtime_r(t: *const platform.c.time_t, out: *Tm) ?*Tm; +}.localtime_r; +/// The UCRT's localtime_s fills the nine ISO fields; gmtoff and zone are POSIX's and stay unset. +extern "c" fn _localtime64_s(out: *Tm, t: *const i64) c_int; +fn localtimeWindows(t: *const platform.c.time_t, out: *Tm) ?*Tm { + out.* = std.mem.zeroes(Tm); + return if (_localtime64_s(out, t) == 0) out else null; +} extern "c" fn strftime(buf: [*]u8, max: usize, fmt: [*:0]const u8, tm: *const Tm) usize; /// datetime.now().strftime(): %f, %z and %Z are Python's own (naive time), the rest is the C library's. fn strftimeNow(rt: *Rt, fmt: []const u8) Error![]const u8 { if (!builtin.link_libc) return rt.fail("strftime_now needs libc", .{}); - var tv: std.c.timeval = undefined; - _ = std.c.gettimeofday(&tv, null); + var tv: platform.c.timeval = undefined; + _ = platform.c.gettimeofday(&tv, null); var tm: Tm = undefined; if (localtime_r(&tv.sec, &tm) == null) return rt.fail("localtime failed", .{}); var pre: std.ArrayList(u8) = .empty; diff --git a/zig/src/cuda/cublaslt.zig b/zig/src/cuda/cublaslt.zig index cc733b8537..cf5b167fc3 100644 --- a/zig/src/cuda/cublaslt.zig +++ b/zig/src/cuda/cublaslt.zig @@ -1,6 +1,7 @@ //! cuBLASLt opened at run time: descriptors, layouts, heuristics and matmul, with fp32 compute only (no TF32 or fast modes). const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Status = c_int; @@ -54,12 +55,13 @@ pub const Api = struct { }; pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, - /// The CUDA 13 soname first, then the unversioned link name. + /// The CUDA 13 soname first, then the unversioned link name; on Windows the toolkit's DLL. pub fn open() Error!Library { - for ([_][]const u8{ "libcublasLt.so.13", "libcublasLt.so" }) |path| { + const names: []const []const u8 = if (@import("builtin").os.tag == .windows) &.{"cublasLt64_13.dll"} else &.{ "libcublasLt.so.13", "libcublasLt.so" }; + for (names) |path| { return openPath(path) catch |e| switch (e) { error.LibraryUnavailable => continue, else => return e, @@ -69,7 +71,7 @@ pub const Library = struct { } pub fn openPath(path: []const u8) Error!Library { - var lib = std.DynLib.open(path) catch return error.LibraryUnavailable; + var lib = platform.DynLib.open(path) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 0a860eac92..1d644f3742 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -1,21 +1,22 @@ //! The CUDA driver opened at run time: libcuda.so.1's versioned entry points in one table, and checked calls. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Error = error{ DriverUnavailable, MissingSymbol, CudaFailed, OutOfDeviceMemory, NotReady, NotFound, Invalid }; pub const Driver = struct { - lib: std.DynLib, + lib: platform.DynLib, api: abi.Api, pub fn open() Error!Driver { - return openPath("libcuda.so.1"); + return openPath(platform.cuda_driver); } /// Resolves every field of `abi.Api` by its exact name; a missing symbol refuses the whole driver. pub fn openPath(path: []const u8) Error!Driver { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); var api: abi.Api = undefined; const info = @typeInfo(abi.Api).@"struct"; diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index c940202bea..86b1086501 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -27,6 +27,7 @@ pub const Layer = struct { experts: usize, limit: f32 = 0.0, // fp8g, nvfp4: the SwiGLU clamp, 0 for none relu2: bool = false, // nvfp4: up is one projection under relu^2 (Nemotron), not SwiGLU's gate and up + swiglu: bool = false, // affine4: up is gate and up interleaved per group ([E][NI/32][D/64][2][288]) under SwiGLU }; /// A call's epilogue: up with its format's activation (affine4 relu^2, else SwiGLU), down to fp32 or to bf16 sums. @@ -51,7 +52,7 @@ pub const Grouped = struct { const up = epi == .up; const nb = (if (up) l.width else l.dims) / 32; const w = Affine4Experts.view(l, up); - return a.run(s, up, in.x, in.stride, if (up) in.slots else 0, w.ptr, w.kg, nb, p, out, w.n, items * nb); + return a.run(s, up, up and l.swiglu, in.x, in.stride, if (up) in.slots else 0, w.ptr, w.kg, nb, p, out, w.n, items * nb); }, .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, false, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), @@ -68,7 +69,7 @@ pub const Grouped = struct { const up = epi == .up; const nb = (if (up) l.width else l.dims) / 32; const w = Affine4Experts.view(l, up); - return a.runPrompt(s, up, in.x, in.stride, if (up) in.slots else 0, w.ptr, w.kg, nb, p, out, w.n, items); + return a.runPrompt(s, up, up and l.swiglu, in.x, in.stride, if (up) in.slots else 0, w.ptr, w.kg, nb, p, out, w.n, items); }, .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, true, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), @@ -91,20 +92,25 @@ fn int(x: usize) c_int { pub const Affine4Experts = struct { up_fn: Function, down_fn: Function, + pair_fn: Function, pre_up_fn: Function, pre_down_fn: Function, + pre_pair_fn: Function, pack_fn: Function, - resident: [2]usize, // blocks the decode up and down kernels fill: per SM times SMs + resident: [3]usize, // blocks the decode up, down and pair kernels fill: per SM times SMs pub const prompt_smem: u32 = 38400; // Pre<64, 1, 2, 2, 4>: three stages of 800 uint4 + pub const pair_prompt_smem: u32 = 34816; // Pre<64, 2, 2, 2, 4>: two stages of 1088 uint4 /// The instances' names (cuobjdump -symbols of the experts, experts_prefill and experts_pack fatbins). pub const symbols = struct { pub const up = "_ZN10tf_experts13expert_kernelILi64ELi1ELi1ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; pub const down = "_ZN10tf_experts13expert_kernelILi64ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + pub const pair = "_ZN10tf_experts13expert_kernelILi64ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; pub const pre_up = "_ZN18tf_experts_prefill14prefill_kernelILi64ELi1ELi1ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; pub const pre_down = "_ZN18tf_experts_prefill14prefill_kernelILi64ELi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; - pub const pack = "_ZN15tf_experts_pack11pack_kernelILi2EEEvPKjPKtS4_Pjiiii"; + pub const pre_pair = "_ZN18tf_experts_prefill14prefill_kernelILi64ELi2ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + pub const pack = "_ZN15tf_experts_pack11pack_kernelILi2EEEvPKjPKtS4_Pjiiiiii"; }; /// From loaded experts, experts_prefill and experts_pack modules. @@ -112,12 +118,15 @@ pub const Affine4Experts = struct { var a: Affine4Experts = undefined; a.up_fn = try decode_mod.function(symbols.up); a.down_fn = try decode_mod.function(symbols.down); + a.pair_fn = try decode_mod.function(symbols.pair); a.pre_up_fn = try prompt_mod.function(symbols.pre_up); a.pre_down_fn = try prompt_mod.function(symbols.pre_down); + a.pre_pair_fn = try prompt_mod.function(symbols.pre_pair); a.pack_fn = try pack_mod.function(symbols.pack); try a.pre_up_fn.allowDynamicShared(prompt_smem); try a.pre_down_fn.allowDynamicShared(prompt_smem); - a.resident = .{ @max(1, try a.up_fn.occupancy(128, 0)) * sms, @max(1, try a.down_fn.occupancy(128, 0)) * sms }; + try a.pre_pair_fn.allowDynamicShared(pair_prompt_smem); + a.resident = .{ @max(1, try a.up_fn.occupancy(128, 0)) * sms, @max(1, try a.down_fn.occupancy(128, 0)) * sms, @max(1, try a.pair_fn.occupancy(128, 0)) * sms }; return a; } @@ -140,31 +149,38 @@ pub const Affine4Experts = struct { return a; } - /// experts.run (decode form): `up` takes token rows (relu^2, bf16 out), else pair rows (fp32 out). - pub fn run(e: *const Affine4Experts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_units: usize) !void { - const grid = @min((max_units + 3) / 4, e.resident[if (up) 0 else 1]); + /// experts.run (decode form): `up` takes token rows (relu^2, or SwiGLU over a `pair` table), else pair rows to f32. + pub fn run(e: *const Affine4Experts, s: Stream, up: bool, pair: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_units: usize) !void { + const which: usize = if (!up) 1 else if (pair) 2 else 0; + const grid = @min((max_units + 3) / 4, e.resident[which]); if (grid < 1) return; var a = args(x, x_stride, slots, w, kg, nb, p, out, n); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 }, .shared = 0 }; - try launch_.launch(if (up) e.up_fn else e.down_fn, cfg, s, &a); + try launch_.launch(if (!up) e.down_fn else if (pair) e.pair_fn else e.up_fn, cfg, s, &a); } - /// experts.prefill: 64 pairs x 128 columns a CTA; `up` relu^2, else bf16 sums (epilogue 3). - pub fn runPrompt(e: *const Affine4Experts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_items: usize) !void { + /// experts.prefill: 64 pairs x 128 columns a CTA; `up` relu^2 (or SwiGLU over a `pair` table), else bf16 sums. + pub fn runPrompt(e: *const Affine4Experts, s: Stream, up: bool, pair: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_items: usize) !void { const grid = max_items * ((nb + 3) / 4); if (grid < 1) return; var a = args(x, x_stride, slots, w, kg, nb, p, out, n); - const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 256 }, .shared = prompt_smem }; - try launch_.launch(if (up) e.pre_up_fn else e.pre_down_fn, cfg, s, &a); + const smem: u32 = if (up and pair) pair_prompt_smem else prompt_smem; + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 256 }, .shared = smem }; + try launch_.launch(if (!up) e.pre_down_fn else if (pair) e.pre_pair_fn else e.pre_up_fn, cfg, s, &a); } /// experts.make: e experts' MLX words (n, k/8), scales and biases (n, k/64) on the device into `out`'s blocks. pub fn pack(e: *const Affine4Experts, s: Stream, words: u64, scales: u64, biases: u64, out: u64, count: usize, n: usize, k: usize) !void { + return e.packInto(s, words, scales, biases, out, count, n, k, 1, 0); + } + + /// pack into the `mi`-th of `m` interleaved tables ([E][n/32][k/64][m][288]): gate then up for the pair form. + pub fn packInto(e: *const Affine4Experts, s: Stream, words: u64, scales: u64, biases: u64, out: u64, count: usize, n: usize, k: usize, m: usize, mi: usize) !void { const kg = k / 64; const nb = n / 32; var a: launch_.Args = .{}; for ([_]u64{ words, scales, biases, out }) |v| a.add(v); - for ([_]usize{ n, k / 8, kg, nb }) |v| a.add(int(v)); + for ([_]usize{ n, k / 8, kg, nb, m, mi }) |v| a.add(int(v)); try launch_.launch(e.pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(count) }, .block = .{ .x = 288 } }, s, &a); } @@ -172,6 +188,11 @@ pub const Affine4Experts = struct { pub fn bytes(count: usize, n: usize, k: usize) usize { return count * (n / 32) * (k / 64) * 288 * 4; } + + /// Bytes of e experts' gate-and-up pair tables ([E][n/32][k/64][2][288] int32). + pub fn pairBytes(count: usize, n: usize, k: usize) usize { + return 2 * bytes(count, n, k); + } }; /// Block-FP8 experts (fp8/experts.cu): the Python packing, kernels and grid choice, decode and prompt forms. diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..c3014d0ed4 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -19,6 +19,7 @@ pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").b pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack").bytes else &.{}; +pub const moe: []const u8 = if (available) &Blob("fatbin_moe").bytes else &.{}; pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_attention").bytes else &.{}; pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes else &.{}; pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; @@ -30,6 +31,10 @@ pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").b pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; +pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; +pub const qwen35_draft: []const u8 = if (available) &Blob("fatbin_qwen35_draft").bytes else &.{}; +pub const qwen35_keep: []const u8 = if (available) &Blob("fatbin_qwen35_keep").bytes else &.{}; +pub const qwen36_mtp: []const u8 = if (available) &Blob("fatbin_qwen36_mtp").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/cuda/moe.zig b/zig/src/cuda/moe.zig new file mode 100644 index 0000000000..d293ae770d --- /dev/null +++ b/zig/src/cuda/moe.zig @@ -0,0 +1,65 @@ +//! Softmax top-k routing with a sigmoid-gated shared expert (moe.cu), by the rules of tensorfold/cuda/moe.py. + +const std = @import("std"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// Routed experts a row may pick from (a lane holds eight), and the picks a row may hold (a lane a slot). +pub const max_experts = 256; +pub const max_top_k = 31; + +pub const Fns = struct { + router_fn: Function, + topk_fn: Function, + combine_fn: Function, + + /// From the loaded moe module (kernels.moe). + pub fn resolve(mod: Module) !Fns { + return .{ .router_fn = try mod.function("tf_moe_router"), .topk_fn = try mod.function("tf_moe_topk"), .combine_fn = try mod.function("tf_moe_combine") }; + } + + /// logits[r][e] = fp32 x[r] . w[e]: `rows` bf16 rows `x_stride` apart against `count` rows of `dims`, gate last. + pub fn route(f: Fns, s: Stream, x: u64, x_stride: usize, rows: usize, w: u64, logits: u64, dims: usize, count: usize) !void { + var a: launch_.Args = .{}; + a.add(x); + a.add(int(x_stride)); + a.add(w); + a.add(logits); + for ([_]usize{ rows, dims, count }) |v| a.add(int(v)); + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast((rows + 15) / 16), .y = @intCast((count + 15) / 16), .z = 1 }, .block = .{ .x = 256 } }; + try launch_.launch(f.router_fn, cfg, s, &a); + } + + /// Each row's top-k picks and bf16-rounded weights ([rows][top_k + 1]); slot top_k is the shared expert. + pub fn pick(f: Fns, s: Stream, logits: u64, picks: u64, wts: u64, rows: usize, experts: usize, top_k: usize) !void { + if (experts > max_experts or top_k > max_top_k) return error.RoutingTooWide; + var a: launch_.Args = .{}; + a.add(logits); + a.add(picks); + a.add(wts); + a.add(int(experts)); + a.add(int(top_k)); + try launch_.launch(f.topk_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 32 } }, s, &a); + } + + /// out[r] = bf16(sum of w[r][s] * y[r][s]) over the slots in order; y fp32 (decode rows) or bf16 (prompt rows). + pub fn combine(f: Fns, s: Stream, y: u64, y_f32: bool, wts: u64, out: u64, rows: usize, slots: usize, dims: usize) !void { + var a: launch_.Args = .{}; + a.add(y); + a.add(@as(c_int, @intFromBool(y_f32))); + a.add(wts); + a.add(out); + a.add(int(slots)); + a.add(int(dims)); + const blocks = @min((dims + 255) / 256, 16); + try launch_.launch(f.combine_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(blocks), .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } +}; + +fn int(x: usize) c_int { + return @intCast(x); +} diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..5bc478c6b5 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -1,6 +1,7 @@ //! NCCL opened at run time for tensor-parallel ranks: communicators and the collectives the lane engines use. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Result = c_int; @@ -30,7 +31,7 @@ pub const Api = struct { }; pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, pub fn open() Error!Library { @@ -38,7 +39,7 @@ pub const Library = struct { } pub fn openPath(path: []const u8) Error!Library { - var lib = std.DynLib.open(path) catch return error.LibraryUnavailable; + var lib = platform.DynLib.open(path) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..629a66bfe7 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -34,6 +34,7 @@ pub const qlinear = @import("qlinear.zig"); pub const norms = @import("norms.zig"); pub const embed_rows = @import("embed_rows.zig"); pub const experts = @import("experts.zig"); +pub const moe = @import("moe.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); diff --git a/zig/src/fabric/mailbox.zig b/zig/src/fabric/mailbox.zig index 3010d50af6..49b08f4b4b 100644 --- a/zig/src/fabric/mailbox.zig +++ b/zig/src/fabric/mailbox.zig @@ -1,5 +1,6 @@ //! One mcdma-rpcd mailbox (protocol 1): the mapped halves, a client's ordered call and a service's request and reply. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const layout = @import("layout.zig"); const words = @import("words.zig"); @@ -48,32 +49,32 @@ pub const Mailbox = struct { fn openShm(name: []const u8) Error!Mailbox { var path: [64]u8 = undefined; const z = std.mem.printSentinel(&path, "/mcdma-rpc.{s}", .{name}, 0) catch return error.InvalidName; - const flags: c_int = @bitCast(std.c.O{ .ACCMODE = .RDWR }); - const fd = if (builtin.os.tag == .macos) std.c.shm_open(z.ptr, flags, @as(c_uint, 0)) else std.c.shm_open(z.ptr, flags, 0); + const flags: c_int = @bitCast(platform.c.O{ .ACCMODE = .RDWR }); + const fd = if (builtin.os.tag == .macos) platform.c.shm_open(z.ptr, flags, @as(c_uint, 0)) else platform.c.shm_open(z.ptr, flags, 0); if (fd < 0) return error.NoMailbox; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); return mapFd(fd); } /// The listen end's mailbox: POSIX shared memory on a Mac, the file /dev/shm/mcdma-rpc.NAME on Linux; MCDMA_RPC_BOX_DIR names a test directory. pub fn openListen(name: []const u8) Error!Mailbox { if (!layout.validName(name)) return error.InvalidName; - const env = std.c.getenv("MCDMA_RPC_BOX_DIR"); + const env = platform.c.getenv("MCDMA_RPC_BOX_DIR"); if (env == null and builtin.os.tag == .macos) return openShm(name); const dir = if (env) |d| std.mem.span(d) else "/dev/shm"; var path: [256]u8 = undefined; const z = std.mem.printSentinel(&path, "{s}/mcdma-rpc.{s}", .{ dir, name }, 0) catch return error.InvalidName; - const fd = std.c.open(z.ptr, .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .CLOEXEC = true }); + const fd = platform.c.open(z.ptr, .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .CLOEXEC = true }); if (fd < 0) return error.NoMailbox; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); return mapFd(fd); } /// Map a mailbox the daemon created; it must be this user's and owner-only, as the daemons make them. - fn mapFd(fd: std.c.fd_t) Error!Mailbox { - var st: std.c.Stat = undefined; - if (std.c.fstat(fd, &st) != 0) return error.NoMailbox; - if (st.uid != std.c.getuid() or st.mode & 0o077 != 0) return error.UnsafeMailbox; + fn mapFd(fd: platform.c.fd_t) Error!Mailbox { + var st: platform.c.Stat = undefined; + if (platform.c.fstat(fd, &st) != 0) return error.NoMailbox; + if (st.uid != platform.c.getuid() or st.mode & 0o077 != 0) return error.UnsafeMailbox; const size: usize = @intCast(st.size); if (size < 2 * layout.segment) return error.BadSizes; const mem = std.posix.mmap(null, size, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0) catch return error.NoMailbox; @@ -210,14 +211,14 @@ pub fn allocFake(req: usize, rep: usize) ![]align(page) u8 { } test "a listen mailbox on this Mac is POSIX shared memory, as MCDMA 2.0.0's Mac listen end makes it" { - if (builtin.os.tag != .macos or std.c.getenv("MCDMA_RPC_BOX_DIR") != null) return error.SkipZigTest; + if (builtin.os.tag != .macos or platform.c.getenv("MCDMA_RPC_BOX_DIR") != null) return error.SkipZigTest; const name = "/mcdma-rpc.tf-zig-shm-test"; - const flags: c_int = @bitCast(std.c.O{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true }); - const fd = std.c.shm_open(name, flags, @as(c_uint, 0o600)); + const flags: c_int = @bitCast(platform.c.O{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true }); + const fd = platform.c.shm_open(name, flags, @as(c_uint, 0o600)); try std.testing.expect(fd >= 0); - defer _ = std.c.shm_unlink(name); - defer _ = std.c.close(fd); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 2 * layout.segment)); + defer _ = platform.c.shm_unlink(name); + defer _ = platform.c.close(fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 2 * layout.segment)); const mem = try std.posix.mmap(null, 2 * layout.segment, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); std.mem.writeInt(u64, mem[layout.sizes..][0..8], layout.segment, .little); std.mem.writeInt(u64, mem[layout.sizes + 8 ..][0..8], layout.segment, .little); diff --git a/zig/src/fabric/mcdma.zig b/zig/src/fabric/mcdma.zig index 290bcc1c61..d3292477a9 100644 --- a/zig/src/fabric/mcdma.zig +++ b/zig/src/fabric/mcdma.zig @@ -1,5 +1,6 @@ //! Registered windows over first-party MCDMA: one progress context per port, with staged sources held through flush. const std = @import("std"); +const platform = @import("sys"); const abi = @import("mcdma_abi.zig"); const rma = @import("rdma.zig"); const words = @import("words.zig"); @@ -367,7 +368,7 @@ test "an endpoint gives the library's placing thread user-interactive QoS unless for ([_]?[*:0]const u8{ null, "0", "1" }) |set| { _ = if (set) |v| setenv("MCDMA_FABRIC_QOS", v, 1) else unsetenv("MCDMA_FABRIC_QOS"); try std.testing.expectError(error.LibraryUnavailable, Endpoint.create(std.testing.allocator, "/nonexistent/libmcdma-fabric.dylib", config)); - try std.testing.expectEqualStrings(if (set) |v| std.mem.span(v) else "1", if (std.c.getenv("MCDMA_FABRIC_QOS")) |v| std.mem.span(v) else "(unset)"); + try std.testing.expectEqualStrings(if (set) |v| std.mem.span(v) else "1", if (platform.c.getenv("MCDMA_FABRIC_QOS")) |v| std.mem.span(v) else "(unset)"); } } diff --git a/zig/src/fabric/mcdma_abi.zig b/zig/src/fabric/mcdma_abi.zig index 0b4bcf062c..feeacbb75c 100644 --- a/zig/src/fabric/mcdma_abi.zig +++ b/zig/src/fabric/mcdma_abi.zig @@ -1,5 +1,6 @@ //! First-party libmcdma-fabric ABI 1, loaded only from an explicit absolute library path. const std = @import("std"); +const platform = @import("sys"); pub const Context = opaque {}; pub const Peer = opaque {}; @@ -27,7 +28,7 @@ pub const MaxLinks = *const fn () callconv(.c) c_uint; pub const ws_room: usize = 64; pub const Library = struct { - handle: std.DynLib, + handle: platform.DynLib, api: Api, write_signal: ?WriteSignal = null, connect_links: ?ConnectLinks = null, @@ -36,7 +37,7 @@ pub const Library = struct { /// No install-path or loader-search fallback; the caller supplies the inspected user-directory library. pub fn open(path: [:0]const u8) Error!Library { if (!std.fs.path.isAbsolute(path)) return error.InvalidPath; - var handle = std.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; + var handle = platform.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; errdefer handle.close(); var api: Api = undefined; inline for (@typeInfo(Api).@"struct".field_names) |name| { diff --git a/zig/src/fabric/registration.zig b/zig/src/fabric/registration.zig index 2bbaf084b4..11a18e92cc 100644 --- a/zig/src/fabric/registration.zig +++ b/zig/src/fabric/registration.zig @@ -1,13 +1,14 @@ //! A service's registration with its listen daemon ("MODE poll"), and the daemons' STATUS lines, as protocol 1 defines. const std = @import("std"); +const platform = @import("sys"); const layout = @import("layout.zig"); pub const Error = error{ InvalidName, NoDaemon, Busy, Refused, Io }; /// The registration channel: the daemon sends nothing after OK until the link drops (BYE or a closed socket). pub const Registration = struct { - in_fd: std.c.fd_t, - out_fd: std.c.fd_t, + in_fd: platform.c.fd_t, + out_fd: platform.c.fd_t, /// Register with the listen daemon at MCDMA_RPCD_SOCKET, else /tmp/mcdma-rpcd.NAME.sock. pub fn attach(name: []const u8) Error!Registration { @@ -35,21 +36,21 @@ pub const Registration = struct { /// Whether the daemon still holds the registration: anything readable means BYE or a closed socket. pub fn alive(r: *const Registration) bool { - var fds = [_]std.c.pollfd{.{ .fd = r.in_fd, .events = std.c.POLL.IN, .revents = 0 }}; - const n = std.c.poll(&fds, 1, 0); + var fds = [_]platform.c.pollfd{.{ .fd = r.in_fd, .events = platform.c.POLL.IN, .revents = 0 }}; + const n = platform.c.poll(&fds, 1, 0); return n == 0; } pub fn close(r: *Registration) void { - _ = std.c.close(r.in_fd); - if (r.out_fd != r.in_fd) _ = std.c.close(r.out_fd); + _ = platform.c.close(r.in_fd); + if (r.out_fd != r.in_fd) _ = platform.c.close(r.out_fd); r.* = undefined; } }; /// The listen daemon's socket for link `name`, honouring MCDMA_RPCD_SOCKET as mcdma-rpcd does. pub fn socketPath(buf: []u8, name: []const u8) ![:0]const u8 { - if (std.c.getenv("MCDMA_RPCD_SOCKET")) |env| { + if (platform.c.getenv("MCDMA_RPCD_SOCKET")) |env| { const s = std.mem.span(env); if (s.len > 0) return std.mem.printSentinel(buf, "{s}", .{s}, 0); } @@ -57,35 +58,35 @@ pub fn socketPath(buf: []u8, name: []const u8) ![:0]const u8 { } /// Connect to the daemon's local Unix socket; the only socket the fabric opens, and never TCP. -fn connectUnix(path: [:0]const u8) ?std.c.fd_t { - var addr: std.c.sockaddr.un = .{ .path = undefined }; +fn connectUnix(path: [:0]const u8) ?platform.c.fd_t { + var addr: platform.c.sockaddr.un = .{ .path = undefined }; if (path.len >= addr.path.len) return null; @memset(&addr.path, 0); @memcpy(addr.path[0..path.len], path); - const fd = std.c.socket(std.c.AF.UNIX, std.c.SOCK.STREAM, 0); + const fd = platform.c.socket(platform.c.AF.UNIX, platform.c.SOCK.STREAM, 0); if (fd < 0) return null; - if (std.c.connect(fd, @ptrCast(&addr), @sizeOf(std.c.sockaddr.un)) != 0) { - _ = std.c.close(fd); + if (platform.c.connect(fd, @ptrCast(&addr), @sizeOf(platform.c.sockaddr.un)) != 0) { + _ = platform.c.close(fd); return null; } return fd; } -fn writeAll(fd: std.c.fd_t, bytes: []const u8) Error!void { +fn writeAll(fd: platform.c.fd_t, bytes: []const u8) Error!void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes[done..].ptr, bytes.len - done); + const n = platform.c.write(fd, bytes[done..].ptr, bytes.len - done); if (n <= 0) return error.Io; done += @intCast(n); } } /// Read one line a byte at a time, as the daemons write them; the newline is dropped. -fn readLine(fd: std.c.fd_t, buf: []u8) Error![]const u8 { +fn readLine(fd: platform.c.fd_t, buf: []u8) Error![]const u8 { var n: usize = 0; while (n < buf.len) { var c: [1]u8 = undefined; - if (std.c.read(fd, &c, 1) != 1) return error.Io; + if (platform.c.read(fd, &c, 1) != 1) return error.Io; if (c[0] == '\n') return buf[0..n]; buf[n] = c[0]; n += 1; @@ -139,9 +140,9 @@ pub fn parsePeer(line: []const u8) ?Peer { return p; } -fn pipePair() ![2]std.c.fd_t { - var fds: [2]std.c.fd_t = undefined; - if (std.c.pipe(&fds) != 0) return error.Io; +fn pipePair() ![2]platform.c.fd_t { + var fds: [2]platform.c.fd_t = undefined; + if (platform.c.pipe(&fds) != 0) return error.Io; return fds; } @@ -149,14 +150,14 @@ test "registration answers OK, ERR busy or anything else, and BYE ends it" { for ([_][]const u8{ "OK\n", "ERR busy\n", "ERR unknown command\n" }, 0..) |answer, i| { const to_daemon = try pipePair(); const from_daemon = try pipePair(); - defer for ([_]std.c.fd_t{ to_daemon[0], to_daemon[1], from_daemon[0], from_daemon[1] }) |fd| { - _ = std.c.close(fd); + defer for ([_]platform.c.fd_t{ to_daemon[0], to_daemon[1], from_daemon[0], from_daemon[1] }) |fd| { + _ = platform.c.close(fd); }; try writeAll(from_daemon[1], answer); const r: Registration = .{ .in_fd = from_daemon[0], .out_fd = to_daemon[1] }; const got = r.handshake(); var sent: [16]u8 = undefined; - try std.testing.expectEqual(@as(isize, 10), std.c.read(to_daemon[0], &sent, sent.len)); + try std.testing.expectEqual(@as(isize, 10), platform.c.read(to_daemon[0], &sent, sent.len)); try std.testing.expectEqualStrings("MODE poll\n", sent[0..10]); switch (i) { 0 => { diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index 1da50d386d..9b9353c554 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -1,5 +1,6 @@ //! The system's verbs library by dlopen (Apple's librdma, Linux libibverbs): every entry point resolved by name. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const abi = @import("verbs_abi.zig"); @@ -52,13 +53,13 @@ const GidEx = *const fn (ctx: Ctx, port: u32, index: u32, entry: *abi.GidEntry, const InOrder = *const fn (qp: *abi.Qp, op: c_int, flags: u32) callconv(.c) c_int; pub const Verbs = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, gid_ex: ?GidEx = null, in_order: ?InOrder = null, pub fn open(path: ?[:0]const u8) Error!Verbs { - var lib = std.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; + var lib = platform.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; @@ -195,7 +196,7 @@ test "a data-device GID scan skips zero entries and prefers link-local over an I } test "verbs smoke: open the device list and close it (TF_FABRIC_VERBS=1)" { - const flag = std.c.getenv("TF_FABRIC_VERBS") orelse return error.SkipZigTest; + const flag = platform.c.getenv("TF_FABRIC_VERBS") orelse return error.SkipZigTest; if (flag[0] != '1') return error.SkipZigTest; var v = Verbs.open(null) catch |err| { std.debug.print("verbs: {s} unavailable ({s})\n", .{ default_path, @errorName(err) }); diff --git a/zig/src/fabric/verbs_abi.zig b/zig/src/fabric/verbs_abi.zig index 1941fc3157..801fa9d80a 100644 --- a/zig/src/fabric/verbs_abi.zig +++ b/zig/src/fabric/verbs_abi.zig @@ -1,5 +1,6 @@ //! The verbs ABI (rdma-core's infiniband/verbs.h, which Apple's SDK ships for librdma): the structs and constants we use. const std = @import("std"); +const platform = @import("sys"); pub const Device = extern struct { _ops: [2]?*anyopaque, @@ -22,8 +23,8 @@ pub const Qp = extern struct { qp_num: u32, state: c_int, qp_type: c_int, - mutex: std.c.pthread_mutex_t, - cond: std.c.pthread_cond_t, + mutex: platform.c.pthread_mutex_t, + cond: platform.c.pthread_cond_t, events_completed: u32, }; @@ -33,8 +34,8 @@ pub const Cq = extern struct { cq_context: ?*anyopaque, handle: u32, cqe: c_int, - mutex: std.c.pthread_mutex_t, - cond: std.c.pthread_cond_t, + mutex: platform.c.pthread_mutex_t, + cond: platform.c.pthread_cond_t, comp_events_completed: u32, async_events_completed: u32, }; @@ -59,7 +60,7 @@ pub const Context = extern struct { cmd_fd: c_int, async_fd: c_int, num_comp_vectors: c_int, - mutex: std.c.pthread_mutex_t, + mutex: platform.c.pthread_mutex_t, abi_compat: ?*anyopaque, }; diff --git a/zig/src/fabric/words.zig b/zig/src/fabric/words.zig index 53d41a0f58..f9c749fc72 100644 --- a/zig/src/fabric/words.zig +++ b/zig/src/fabric/words.zig @@ -1,5 +1,6 @@ //! Ordered mailbox words: MCDMA's libmcdma-rpc (ABI 1, rpc/mcdma_rpc.h) by dlopen, or its same rules in Zig. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const layout = @import("layout.zig"); @@ -30,20 +31,20 @@ pub fn load(w: *const u64) u64 { /// CLOCK_MONOTONIC in nanoseconds, the clock libmcdma-rpc waits on. pub fn nowNs() u64 { - var ts: std.c.timespec = undefined; - _ = std.c.clock_gettime(.MONOTONIC, &ts); + var ts: platform.c.timespec = undefined; + _ = platform.c.clock_gettime(.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); } /// Sleep about `ns` nanoseconds (the library's short step between polls). pub fn pause(ns: u64) void { - const ts: std.c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + _ = platform.c.nanosleep(&ts, null); } pub fn checkedNowNs() error{ClockUnavailable}!u64 { - var ts: std.c.timespec = undefined; - if (std.c.clock_gettime(.MONOTONIC, &ts) != 0 or ts.sec < 0 or ts.nsec < 0 or ts.nsec >= std.time.ns_per_s) return error.ClockUnavailable; + var ts: platform.c.timespec = undefined; + if (platform.c.clock_gettime(.MONOTONIC, &ts) != 0 or ts.sec < 0 or ts.nsec < 0 or ts.nsec >= std.time.ns_per_s) return error.ClockUnavailable; return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); } @@ -81,7 +82,7 @@ else /// libmcdma-rpc opened at run time; the three Metal functions are optional, as the header says. pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, abi_fn: AbiFn, wait_word: WaitFn, store_word: StoreFn, @@ -92,7 +93,7 @@ pub const Library = struct { /// Open `path`, else MCDMA_RPC_LIBRARY, else the install prefix; refuse any ABI but 1. pub fn open(path: ?[:0]const u8) Error!Library { if (path) |p| return openPath(p); - if (std.c.getenv("MCDMA_RPC_LIBRARY")) |env| { + if (platform.c.getenv("MCDMA_RPC_LIBRARY")) |env| { const p = std.mem.span(env); if (p.len > 0) return openPath(p); } @@ -104,7 +105,7 @@ pub const Library = struct { } pub fn openPath(path: [:0]const u8) Error!Library { - var lib = std.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; + var lib = platform.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; errdefer lib.close(); var l: Library = .{ .lib = lib, diff --git a/zig/src/families/flashnext/call_log.zig b/zig/src/families/flashnext/call_log.zig index 387ec89220..5afdca5da9 100644 --- a/zig/src/families/flashnext/call_log.zig +++ b/zig/src/families/flashnext/call_log.zig @@ -1,5 +1,6 @@ //! FZ_CALL_LOG: a prompt pass's calls (rows, parts, wall and GPU milliseconds), logged in one line after its first token. const std = @import("std"); +const platform = @import("sys"); const N = 16; @@ -17,7 +18,7 @@ pub const CallLog = struct { if (c.n == N) return; c.rows[c.n] = @intCast(rows); c.parts[c.n] = @intCast(parts); - c.wall[c.n] = @floatCast(@as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6); // mach ticks: 125/3 ns + c.wall[c.n] = @floatCast(@as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6); // mach ticks: 125/3 ns c.gpu[c.n] = @floatCast((g1 - g0) * 1e3); c.n += 1; } diff --git a/zig/src/families/flashnext/engine.zig b/zig/src/families/flashnext/engine.zig index c9f8481249..9ff39b2845 100644 --- a/zig/src/families/flashnext/engine.zig +++ b/zig/src/families/flashnext/engine.zig @@ -1,6 +1,7 @@ //! Flash Next served from the replay engine: one reply at a time, its prompt through prompt chunks (the MTP head's //! keys for every prompt row), then GPU-side rounds whose drafts the depth rule sets; tokens leave as rounds are read. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fz = @import("replay.zig"); const config = @import("config.zig"); @@ -151,8 +152,8 @@ pub const Engine = struct { e.copy_long = 6; e.segments = true; e.mark_taps = true; - e.pair_min = tp_settings.pairMin(if (std.c.getenv("FZ_PAIR_MIN")) |v| std.mem.span(v) else null, PAIR_MIN); // rank 1 takes rank 0's - e.call_log = std.c.getenv("FZ_CALL_LOG") != null; + e.pair_min = tp_settings.pairMin(if (platform.c.getenv("FZ_PAIR_MIN")) |v| std.mem.span(v) else null, PAIR_MIN); // rank 1 takes rank 0's + e.call_log = platform.c.getenv("FZ_CALL_LOG") != null; e.passed = null; e.snap_pool = .{}; e.handoff_ns = 0; @@ -177,8 +178,8 @@ pub const Engine = struct { r.event = try device.sharedEvent(); if (dump_dir) |d| try r.compile(d) else try r.compileChecked(); r.sel = try Select.init(r, MAXR); - r.lane_new = std.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (the recorded bits) - if (std.c.getenv("FZ_GDN")) |v| { + r.lane_new = platform.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (the recorded bits) + if (platform.c.getenv("FZ_GDN")) |v| { r.gdn_pipe = try fz.gdn_step.compile(r, false); if (v[0] == '2') r.gdn_kept = try fz.gdn_step.compile(r, true); // GPU-side rounds keep one state a layer } @@ -211,7 +212,7 @@ pub const Engine = struct { }; e.identity_hash = std.hash.Wyhash.hash(0x6964, identity); const owned_pack_dir = if (dump_dir == null) blk: { - const home = if (std.c.getenv("HOME")) |v| std.mem.span(v) else ""; + const home = if (platform.c.getenv("HOME")) |v| std.mem.span(v) else ""; const root = try pack_cache.defaultRoot(arena, home); break :blk try pack_cache.ensure(gpa, io, model_dir, root, identity, PackBuilder{}); } else null; @@ -434,7 +435,7 @@ pub const Engine = struct { e.tsel = try GSelect.init(r, 0); e.msel = try GSelect.init(r, 0); // speed-up mode: TP=2 with the peer the settings name (tp.zig); both ranks then run the same replies in lockstep - const tp_path: ?[]const u8 = speed_up orelse if (std.c.getenv("TF_FLASHNEXT_TP")) |p| std.mem.span(p) else null; + const tp_path: ?[]const u8 = speed_up orelse if (platform.c.getenv("TF_FLASHNEXT_TP")) |p| std.mem.span(p) else null; if (tp_path != null and !uniform) return error.SpeedUpNeeds6Bit32; // TP splits the 6/32 kernels only if (tp_path) |path| r.tp = try fz.Tp2.init(arena, r.device, path); if (r.tp == null) try fz.rows_load.attach( @@ -504,8 +505,8 @@ pub const Engine = struct { const pool = mtl.objc.Pool.push(); defer pool.pop(); if (r.tp) |tp| if (tp.rank == 0) { // served speed-up mode: rank 1 runs the same request (follow) - const t0 = std.c.mach_absolute_time(); - defer e.handoff_ns = @intCast((std.c.mach_absolute_time() - t0) * 125 / 3); // mach ticks are 125/3 ns on Apple silicon + const t0 = platform.c.mach_absolute_time(); + defer e.handoff_ns = @intCast((platform.c.mach_absolute_time() - t0) * 125 / 3); // mach ticks are 125/3 ns on Apple silicon try follow_mod.send(e, tp, prompt, from, marks, max_tokens, eos, depth, budget); if (!try tp.waitAck()) return error.PeerNotResumed; }; @@ -522,7 +523,7 @@ pub const Engine = struct { var calls: CallLog = .{}; while (at < prompt.len) { if (try e.agree(out.cancelled(out.ctx))) return .{ .reason = .cancelled }; - const t_call = std.c.mach_absolute_time(); + const t_call = platform.c.mach_absolute_time(); const g_call = m.gpu_seconds; if (r.tp) |tp| { // speed-up mode: the call's rows split across the two Macs, running on past up to MARKS marks const pair_taps = e.mark_taps and m.marks != null; diff --git a/zig/src/families/flashnext/follow.zig b/zig/src/families/flashnext/follow.zig index 4255faeb6e..482870c098 100644 --- a/zig/src/families/flashnext/follow.zig +++ b/zig/src/families/flashnext/follow.zig @@ -1,5 +1,6 @@ //! Speed-up mode's rank 1: rank 0's requests run here in step, and this Mac keeps its own states under rank 0's names. const std = @import("std"); +const platform = @import("sys"); const engine = @import("engine.zig"); const snapshot = @import("snapshot.zig"); const fz = @import("replay.zig"); @@ -9,7 +10,7 @@ const Engine = engine.Engine; /// Rank 0's requests as they come, until its empty request; between them a free buffer is readied inside the budget. pub fn run(e: *Engine) !void { while (try one(e, null)) |_| { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const got = e.snap_pool.ready(e.r.device, engine.NEXT_TURN, .{ .ctx = e.r.tp.?, .check = requestWaiting }, e.peer_budget -| e.peer_held -| e.snap_pool.spare()); if (got.cap > 0) std.log.info("speed-up rank 1: readied {d} of {d} MiB for the next save in {d:.1} ms", .{ got.touched >> 20, got.cap >> 20, msSince(t0) }); } @@ -21,7 +22,7 @@ fn requestWaiting(ctx: *anyopaque) bool { } fn msSince(t0: u64) f64 { - return @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6; // mach ticks are 125/3 ns on Apple silicon + return @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6; // mach ticks are 125/3 ns on Apple silicon } /// Rank 0's request for `one`: a 16-word head, then the prompt, marks, dropped states' keys and the budget. @@ -65,7 +66,7 @@ pub fn one(e: *Engine, out: ?engine.Out) !?engine.Result { const from = head[12]; if (head[15] == 0) return error.TpProtocol; // rank 0 sends its positive pair minimum with every request e.pair_min = head[15]; - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const ok = from == 0 or blk: { const st = e.peer_kept.get(engine.keyOf(prompt[0 .. from + 1])) orelse break :blk false; snapshot.restore(e, st) catch break :blk false; @@ -123,7 +124,7 @@ const Reply = struct { } fn marked(ctx: *anyopaque, at: usize) void { const r: *Reply = @ptrCast(@alignCast(ctx)); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); defer { r.kept += 1; r.keep_ns += msSince(t0) * 1e6; diff --git a/zig/src/families/flashnext/host_test.zig b/zig/src/families/flashnext/host_test.zig index bc198ce4ad..642812ce54 100644 --- a/zig/src/families/flashnext/host_test.zig +++ b/zig/src/families/flashnext/host_test.zig @@ -1,5 +1,6 @@ //! Host-only contracts exercise strict text configuration, affine six-bit packing and borrowed-view shapes. const std = @import("std"); +const platform = @import("sys"); const posix = std.posix; const family = @import("flashnext.zig"); const st = @import("../../core/safetensors.zig"); @@ -298,12 +299,12 @@ test "the prompt's checked-in text resolves and patches from another working dir const sources = @import("kernel_sources").flashnext_gen; var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); - const old_fd = std.c.open(".", .{}); // O_RDONLY via the flags struct's default - defer _ = std.c.close(old_fd); + const old_fd = platform.c.open(".", .{}); // O_RDONLY via the flags struct's default + defer _ = platform.c.close(old_fd); const tmpz = try std.fmt.allocPrintSentinel(std.testing.allocator, ".zig-cache/tmp/{s}", .{tmp.sub_path}, 0); defer std.testing.allocator.free(tmpz); - try std.testing.expect(std.c.chdir(tmpz) == 0); - defer _ = std.c.fchdir(old_fd); + try std.testing.expect(platform.c.chdir(tmpz) == 0); + defer _ = platform.c.fchdir(old_fd); for (table.entries) |e| { var text: []const u8 = ""; for (sources.sources) |s| { diff --git a/zig/src/families/flashnext/learned.zig b/zig/src/families/flashnext/learned.zig index 5d1485abdf..262f44f447 100644 --- a/zig/src/families/flashnext/learned.zig +++ b/zig/src/families/flashnext/learned.zig @@ -1,5 +1,6 @@ //! Flash Next's learned prompt states (--learn): a kept state's file on disk, and the probe keying this build's bits. const std = @import("std"); +const platform = @import("sys"); const ks = @import("kernel_sources"); const snapshot = @import("snapshot.zig"); const roles_gen = @import("roles_gen.zig"); @@ -21,24 +22,24 @@ pub fn path(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { pub fn writeFile(st: *const State, file: [:0]const u8) !void { var tmp_buf: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head: Head = .{ .magic = FILE_MAGIC, .layout = @backingInt(st.layout), .at = st.at, .bytes = st.bytes, .hist = st.hist }; try put(fd, std.mem.asBytes(&head)); try put(fd, st.buf.contents()[0..st.bytes]); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } /// `file`'s state after `at` tokens, in a buffer from the engine's pool; another position, layout or size is refused. pub fn readFile(e: *Engine, gpa: std.mem.Allocator, at: usize, file: [:0]const u8) !*State { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: Head = undefined; try get(fd, std.mem.asBytes(&head), 0); const layout: u32 = @backingInt(snapshot.layoutOf(e)); @@ -53,7 +54,7 @@ pub fn readFile(e: *Engine, gpa: std.mem.Allocator, at: usize, file: [:0]const u fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -62,7 +63,7 @@ fn put(fd: c_int, b: []const u8) !void { fn get(fd: c_int, b: []u8, at: u64) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.SnapshotRead; done += @intCast(n); } diff --git a/zig/src/families/flashnext/replay.zig b/zig/src/families/flashnext/replay.zig index 834c0a732a..7021f8c6b1 100644 --- a/zig/src/families/flashnext/replay.zig +++ b/zig/src/families/flashnext/replay.zig @@ -1,6 +1,7 @@ //! Flash Next on our Metal runtime: the Python engine's recorded kernels (a dump from tools/zig/flashnext_dump.py) //! replayed from Zig, with our own kernels for wide experts, prompt chunks, block selection and GPU-side rounds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); const roles_gen = @import("roles_gen.zig"); @@ -43,7 +44,7 @@ pub const Buf = struct { return .{ .b = x.b, .off = x.off + bytes }; } }; -pub const Entry = struct { fd: std.c.fd_t, at: usize, len: usize }; +pub const Entry = struct { fd: platform.c.fd_t, at: usize, len: usize }; pub const Variant = struct { inputs: [][]const u8, outputs: [][]const u8, meta: [][]const u8, pipe: mtl.Pipeline, file: []const u8 = "", name: []const u8 = "", text: []const u8 = "" }; pub const Site = struct { v: *Variant, grid: mtl.Size, tg: mtl.Size }; @@ -814,10 +815,10 @@ const hc_up_source = \\ ; -pub fn readAll(fd: std.c.fd_t, dest: []u8, at: usize) !void { +pub fn readAll(fd: platform.c.fd_t, dest: []u8, at: usize) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } @@ -915,7 +916,7 @@ pub const Run = struct { /// Every tensor of a safetensors file in the index, at its absolute offset. pub fn indexFile(r: *Run, path: [:0]const u8) !void { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; var head: [8]u8 = undefined; try readAll(fd, &head, 0); diff --git a/zig/src/families/flashnext/rows_load.zig b/zig/src/families/flashnext/rows_load.zig index cf8e89cfbb..1932830cd1 100644 --- a/zig/src/families/flashnext/rows_load.zig +++ b/zig/src/families/flashnext/rows_load.zig @@ -1,5 +1,6 @@ //! Rows mode without tensor units (0.6.6's pre-M5 decode): every dense lane also as MLX-layout members for rows_w. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const replay = @import("replay.zig"); pub const rows_w = @import("rows_w.zig"); @@ -14,7 +15,7 @@ const KO = 6144; // the attention and DeltaNet output projections' inputs /// Whether decode reads its dense projections on the row kernels: no tensor units and FZ_M3_LANES unset. pub fn on(device: mtl.Device) bool { - return !device.tensorUnits() and std.c.getenv("FZ_M3_LANES") == null; + return !device.tensorUnits() and platform.c.getenv("FZ_M3_LANES") == null; } /// `d` on the run's encoder at `rows` rows, unless the role is knocked out. diff --git a/zig/src/families/flashnext/tp.zig b/zig/src/families/flashnext/tp.zig index e5718126d2..3fe8c3a514 100644 --- a/zig/src/families/flashnext/tp.zig +++ b/zig/src/families/flashnext/tp.zig @@ -1,5 +1,6 @@ //! Speed-up mode for Flash Next: two Macs with the whole model each, half the work each, over MCDMA. Decode: each rank computes every routed and the shared expert for its half of the window's rows (`ownRows`) and those rows' MoE branch with one Mac's sums; after each target layer's experts the GPU posts a sequence; the host sends the branch rows (bf16) to the peer, and the next layer's combine waits for the flag the peer's message sets, then updates every row's streams (one Mac's bits). DeltaNet layers split their heads: each rank's out-projection partial (fp32), exchanged and added in rank order. Prefill: a chunk's rows split across the two Macs (replay.zig `chunkPair`); the host sends each layer's handoff into the peer's slot for that layer and signals the layer's flag, which the peer's GPU waits for. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const host_wait = @import("tp_host_wait.zig"); @@ -341,9 +342,9 @@ pub const Tp2 = struct { .jobs = try gpa.alloc(Job, JOBS), }; t.one.slice(i32, 1)[0] = 1; - t.trace = std.c.getenv("TF_TP_TRACE") != null; - t.local = std.c.getenv("TF_TP_LOCAL") != null; - t.stats = std.c.getenv("TF_TP_STATS") != null; + t.trace = platform.c.getenv("TF_TP_TRACE") != null; + t.local = platform.c.getenv("TF_TP_LOCAL") != null; + t.stats = platform.c.getenv("TF_TP_STATS") != null; // both ranks up before the first round: a word each way try rd.signal(t.peer, HELLO, 1); while (@atomicLoad(u64, t.word64(HELLO), .acquire) < 1) std.atomic.spinLoopHint(); @@ -529,11 +530,11 @@ pub const Tp2 = struct { /// Served, rank 0: rank 1's answer for the request just sent: true when it restored the same resume state. pub fn waitAck(t: *Tp2) !bool { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (true) { const v = @atomicLoad(u64, t.word64(ACK_FLAG), .acquire); if (v >> 1 >= t.req) return v >> 1 == t.req and v & 1 != 0; - if (std.c.mach_absolute_time() - t0 > 7_200_000_000) return error.TpPeerSilent; // 300 s: rank 1 may still be loading + if (platform.c.mach_absolute_time() - t0 > 7_200_000_000) return error.TpPeerSilent; // 300 s: rank 1 may still be loading std.atomic.spinLoopHint(); } } @@ -551,8 +552,8 @@ pub const Tp2 = struct { if (t.quitting.load(.acquire)) return null; spins += 1; if (spins > 100_000) { // idle: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 50_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } return .{ .head = @ptrCast(t.win.ptr + REQ), .tokens = @ptrCast(@alignCast(t.win.ptr + REQ + 64)) }; @@ -565,11 +566,11 @@ pub const Tp2 = struct { try t.rd.signal(t.peer, CTRL_FLAG, (t.ctrl << 6) | (@as(u64, depth & 31) << 1) | @intFromBool(quit)); return .{ .quit = quit, .depth = depth }; } - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (true) { const v = @atomicLoad(u64, t.word64(CTRL_FLAG), .acquire); if (v >> 6 >= t.ctrl) return .{ .quit = quit or (v >> 6 == t.ctrl and v & 1 != 0), .depth = @intCast((v >> 1) & 31) }; - if (std.c.mach_absolute_time() - t0 > 240_000_000) return error.TpPeerSilent; // 10 s + if (platform.c.mach_absolute_time() - t0 > 240_000_000) return error.TpPeerSilent; // 10 s std.atomic.spinLoopHint(); } } @@ -604,7 +605,7 @@ pub const Tp2 = struct { /// for the peer's, by the flag its message sets) or a prefill send (its writes, then its flag). While idle it /// watches the last exchange: 10 s without the peer's answer fails the link. fn service(t: *Tp2, first: u32) void { - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send const posted = t.word32(SYNC + GPU * 4); const flag = t.word64(FLAG); var seq = first; @@ -614,7 +615,7 @@ pub const Tp2 = struct { var w = @atomicLoad(u32, posted, .acquire); while (!postReached(w, seq) or !reached(t.queued.load(.acquire), seq)) { if (t.stop.load(.acquire)) return; - if (want_at > 0 and !reached(@truncate(@atomicLoad(u64, flag, .acquire)), @truncate(want)) and !t.failed.load(.acquire) and std.c.mach_absolute_time() - want_at > 240_000_000) { // 10 s of 24 MHz ticks + if (want_at > 0 and !reached(@truncate(@atomicLoad(u64, flag, .acquire)), @truncate(want)) and !t.failed.load(.acquire) and platform.c.mach_absolute_time() - want_at > 240_000_000) { // 10 s of 24 MHz ticks t.fail(seq, error.PeerSilent); t.land(want); } @@ -622,7 +623,7 @@ pub const Tp2 = struct { w = @atomicLoad(u32, posted, .acquire); } const job = t.jobs[seq % JOBS]; - const seen = if (t.stats) std.c.mach_absolute_time() else 0; + const seen = if (t.stats) platform.c.mach_absolute_time() else 0; if ((t.local and job.kind != .send) or t.failed.load(.acquire)) { // a failed link drains: the GPU never hangs if (job.kind != .send) t.land(job.value); seq +%= 1; @@ -638,11 +639,11 @@ pub const Tp2 = struct { } if (job.kind != .send) { want = x; - want_at = std.c.mach_absolute_time(); + want_at = platform.c.mach_absolute_time(); if (t.failed.load(.acquire)) t.land(x); } if (t.stats) { - const held = std.c.mach_absolute_time() - seen; // ticks: 41.67 ns each on Apple silicon + const held = platform.c.mach_absolute_time() - seen; // ticks: 41.67 ns each on Apple silicon t.held_ns += held; t.held_max = @max(t.held_max, held); t.held_kind[@backingInt(job.kind)][0] += held; @@ -705,13 +706,13 @@ test "the host service follows the GPU's posts across the 27-bit post wrap and t for (0..8) |_| { // local serving lands each exchange's flag once its post is seen x +%= 1; const seq = t.queue(.{ .kind = .exchange, .value = x }); - const ts: std.c.timespec = .{ .sec = 0, .nsec = 2_000_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 2_000_000 }; + _ = platform.c.nanosleep(&ts, null); try std.testing.expect(@as(u32, @truncate(@atomicLoad(u64, t.word64(FLAG), .acquire))) != x); // not before its post @atomicStore(u32, t.word32(SYNC + GPU * 4), (seq << 5) | 4, .release); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (@as(u32, @truncate(@atomicLoad(u64, t.word64(FLAG), .acquire))) != x) { - if (std.c.mach_absolute_time() - t0 > 24_000_000) return error.ServiceStuck; // 1 s + if (platform.c.mach_absolute_time() - t0 > 24_000_000) return error.ServiceStuck; // 1 s std.atomic.spinLoopHint(); } } @@ -729,8 +730,8 @@ test "partNext and hostWait across the 32-bit wrap: the next exchange's slot, an } const Peer = struct { // the peer's message landing a call's flag, a little later fn land(w: *u64, v: u64) void { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000_000 }; + _ = platform.c.nanosleep(&ts, null); @atomicStore(u64, w, v, .release); } }; diff --git a/zig/src/families/glm/checks.zig b/zig/src/families/glm/checks.zig index 98404ca22a..f1a92ef8c3 100644 --- a/zig/src/families/glm/checks.zig +++ b/zig/src/families/glm/checks.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's dev-time checks on a loaded engine: captures, traces, teacher forcing and knock-out profiles. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("state.zig"); const fwd = @import("forward.zig"); @@ -27,9 +28,9 @@ pub fn capture(e: *Engine, prompt: []const u32, path: []const u8) !void { fwd.head(&x, b.enc, e.sc.hidden, dump.at(x.dump_at), e.sc.picks, n); try e.finish(b.cb, b.enc); fwd.flipKda(&x); - const file = std.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; std.debug.print("captured {d} rows ({d} bytes) in {s}\n", .{ n, bytes, path }); } @@ -47,8 +48,8 @@ pub fn trace(e: *Engine, prompt: []const u32, reply: []const u32, steps: usize, @memcpy(Engine.u32s(e.prompt_ids.at(@as(usize, P) * 4), total - P), reply[0 .. total - P]); const most = @as(usize, st.max_rows) * c.hidden * 2 * (2 + 4 * @as(usize, c.run)) + @as(usize, st.max_rows) * c.vocab * 2; const dump = try e.arena.buffer(most); - const file = std.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); var x = e.ctx(); var at: u32 = 0; var calls: usize = 0; @@ -62,7 +63,7 @@ pub fn trace(e: *Engine, prompt: []const u32, reply: []const u32, steps: usize, try e.finish(b.cb, b.enc); fwd.flipKda(&x); const bytes = x.dump_at + @as(usize, n) * c.vocab * 2; - if (std.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; + if (platform.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; at += n; } std.debug.print("traced {d} calls ({d} prompt rows, {d} steps) in {s}\n", .{ calls, P, total - P, path }); @@ -106,7 +107,7 @@ pub fn profile(e: *Engine, depth: u32, reps: usize, only: bool, parts: bool) !vo const c = &e.c; const D = c.hidden; const d: u32 = if (e.w.mtp == null) 0 else @min(depth, st.max_rows - 1); - const wide = if (std.c.getenv("GLM_PROFILE_ROWS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 1 else 1; // a copy window's rows, without the MTP head + const wide = if (platform.c.getenv("GLM_PROFILE_ROWS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 1 else 1; // a copy window's rows, without the MTP head const R = if (d == 0) std.math.clamp(wide, 1, st.max_rows) else d + 1; const pool = mtl.objc.Pool.push(); defer pool.pop(); diff --git a/zig/src/families/glm/engine.zig b/zig/src/families/glm/engine.zig index e1ffe5ac41..e030c7200a 100644 --- a/zig/src/families/glm/engine.zig +++ b/zig/src/families/glm/engine.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash on Metal: one greedy reply at a time, the prompt in 16-row windows or tensor-unit chunks, then drafted rounds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -95,7 +96,7 @@ pub const Engine = struct { /// The checkpoint in `dir`, caches for `cap` tokens; GLM_LAYERS=N: the first N layers only; GLM_EP=settings: half the experts. pub fn load(gpa: std.mem.Allocator, dir: []const u8, cap: u32) !*Engine { - return loadWith(gpa, dir, cap, if (std.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false); + return loadWith(gpa, dir, cap, if (platform.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false); } /// `load` with expert parallel over the link in `ep_path` (this Mac's settings), or on one Mac when null. @@ -104,23 +105,23 @@ pub const Engine = struct { errdefer gpa.destroy(e); const pool = mtl.objc.Pool.push(); defer pool.pop(); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); e.gpa = gpa; e.ev = 0; e.residency = null; e.gpu = .{ 0, 0 }; - e.fused_route = if (std.c.getenv("GLM_ROUTE")) |v| v[0] != '0' else true; - e.draft_vocab = if (std.c.getenv("GLM_DRAFT_VOCAB")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; + e.fused_route = if (platform.c.getenv("GLM_ROUTE")) |v| v[0] != '0' else true; + e.draft_vocab = if (platform.c.getenv("GLM_DRAFT_VOCAB")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; e.committed = 0; e.ep = null; e.pr = null; e.trace_last = null; e.margins = null; e.chunk_rows = prompt_mod.max_rows; - e.cut = if (std.c.getenv("GLM_CUTS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; - e.copy_min = if (std.c.getenv("GLM_COPY")) |v| std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0, 0, 8) else 0; - e.rank_log = if (std.c.getenv("GLM_RANKS")) |v| v[0] == '1' else false; - if (std.c.getenv("GLM_CHUNK")) |v| e.chunk_rows = std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch prompt_mod.max_rows, st.max_rows + 1, prompt_mod.max_rows); + e.cut = if (platform.c.getenv("GLM_CUTS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; + e.copy_min = if (platform.c.getenv("GLM_COPY")) |v| std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0, 0, 8) else 0; + e.rank_log = if (platform.c.getenv("GLM_RANKS")) |v| v[0] == '1' else false; + if (platform.c.getenv("GLM_CHUNK")) |v| e.chunk_rows = std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch prompt_mod.max_rows, st.max_rows + 1, prompt_mod.max_rows); e.ep_arena = .init(gpa); errdefer e.ep_arena.deinit(); e.device = try mtl.Device.init(); @@ -138,17 +139,17 @@ pub const Engine = struct { e.draft_vocab -= e.draft_vocab % 4; // the head's kernel takes four rows a simdgroup const model = try load_plan.modelHash(gpa, dir, f.bytes[0..f.size]); e.model_hash = model; - if (std.c.getenv("GLM_LAYERS")) |v| try cfg.subset(&e.c, std.fmt.parseInt(u32, std.mem.span(v), 10) catch return error.BadLayerCount); - if (std.c.getenv("GLM_MTP")) |v| if (v[0] == '0') { // no MTP head: no drafts, its ~4 GB left on disk + if (platform.c.getenv("GLM_LAYERS")) |v| try cfg.subset(&e.c, std.fmt.parseInt(u32, std.mem.span(v), 10) catch return error.BadLayerCount); + if (platform.c.getenv("GLM_MTP")) |v| if (v[0] == '0') { // no MTP head: no drafts, its ~4 GB left on disk e.c.mtp = 0; }; const link: ?ep_mod.Settings = if (ep_path) |sp| blk: { const sf = try mtl.MappedFile.open(try e.ep_arena.allocator().dupeSentinel(u8, sp, 0)); defer sf.deinit(); const s = try ep_mod.readSettings(e.ep_arena.allocator(), sf.bytes[0..sf.size]); - const by_rows = if (std.c.getenv("GLM_EP_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "experts") else true; + const by_rows = if (platform.c.getenv("GLM_EP_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "experts") else true; if (by_rows) try cfg.splitRows(&e.c, s.rank, 2) else try cfg.split(&e.c, s.rank, 2); - const tp = if (std.c.getenv("GLM_TP")) |v| v[0] != '0' else true; // TP2: each Mac half the KDA heads (GLM_TP=0: off) + const tp = if (platform.c.getenv("GLM_TP")) |v| v[0] != '0' else true; // TP2: each Mac half the KDA heads (GLM_TP=0: off) if (by_rows and tp) try cfg.splitHeads(&e.c, s.rank, 2); if (e.c.tp > 1) { // TP2's head halves: the whole vocabulary (its kernel's pitch); no host logits to rank e.draft_vocab = e.c.vocab; @@ -162,7 +163,7 @@ pub const Engine = struct { gpa.destroy(e.k); } try frags.check(e.device, e.queue, gpa); // the fragment layout nax.h gives GLM's prompt and attention kernels - const chunked = if (std.c.getenv("GLM_PROMPT")) |v| v[0] != '0' else true; + const chunked = if (platform.c.getenv("GLM_PROMPT")) |v| v[0] != '0' else true; e.limit = try load_plan.loadLimit(e.device); const limit = e.limit.bytes; var plan_bytes = try load_plan.planBytes(gpa, e.device, dir, &e.c); @@ -176,7 +177,7 @@ pub const Engine = struct { plan_bytes = lean_bytes; } } - if (std.c.getenv("GLM_DRY") != null) return error.DryRun; + if (platform.c.getenv("GLM_DRY") != null) return error.DryRun; e.arena = .{ .device = e.device, .gpa = gpa }; errdefer e.arena.deinit(); const both = try st.init(&e.arena, &e.c, cap); @@ -195,9 +196,9 @@ pub const Engine = struct { return error.OverMemoryLimit; } { // free memory short of the load: a cached copy of the checkpoint holds it, the paced reads would wait (#583) - const dropped_at = std.c.mach_absolute_time(); + const dropped_at = platform.c.mach_absolute_time(); if (page_cache.dropForLoad(gpa, dir, plan_bytes + e.arena.bytes)) |d| { - if (d.dropped > 0) std.log.info("glm: dropped {d:.1} GiB of the checkpoint from the file cache, the {d:.1} GiB free memory lacked for the load ({d:.2} s)", .{ @as(f64, @floatFromInt(d.dropped)) / (1 << 30), @as(f64, @floatFromInt(d.deficit)) / (1 << 30), @as(f64, @floatFromInt(std.c.mach_absolute_time() - dropped_at)) / 24e6 }); + if (d.dropped > 0) std.log.info("glm: dropped {d:.1} GiB of the checkpoint from the file cache, the {d:.1} GiB free memory lacked for the load ({d:.2} s)", .{ @as(f64, @floatFromInt(d.dropped)) / (1 << 30), @as(f64, @floatFromInt(d.deficit)) / (1 << 30), @as(f64, @floatFromInt(platform.c.mach_absolute_time() - dropped_at)) / 24e6 }); } else |err| std.log.warn("glm: could not drop the checkpoint's cached pages ({s}); the load reads beside them", .{@errorName(err)}); } e.w = try wts.load(gpa, e.device, dir, &e.c, 16, false); @@ -216,7 +217,7 @@ pub const Engine = struct { e.ep = try ep_mod.Ep.init(gpa, e.device, s, me); } // opt-in: wiring 181 GB leaves macOS nothing to reclaim if another model shares the Mac (Flash Next runs without) - if (std.c.getenv("GLM_RESIDENCY") == null) {} else if (e.device.residencySet(e.w.buffers.items.len + e.arena.buffers.items.len)) |set| { + if (platform.c.getenv("GLM_RESIDENCY") == null) {} else if (e.device.residencySet(e.w.buffers.items.len + e.arena.buffers.items.len)) |set| { for (e.w.buffers.items) |b| set.add(b); for (e.arena.buffers.items) |b| set.add(b); set.commit(); @@ -228,13 +229,13 @@ pub const Engine = struct { if (e.residency) |set| e.keepalive_sets[0] = set; e.keepalive_target = .{ .queue = e.queue, .sets = if (e.residency != null) &e.keepalive_sets else &.{} }; // a process's first prompt pays a one-time cost: every rank of a pair pays it here at the same step (GLM_NO_WARMUP: skip) - if (e.pr != null and cap >= 128 and std.c.getenv("GLM_NO_WARMUP") == null) { + if (e.pr != null and cap >= 128 and platform.c.getenv("GLM_NO_WARMUP") == null) { var ids: [64]u32 = undefined; for (&ids, 0..) |*t, i| t.* = @intCast(1000 + i); var quiet: Quiet = .{}; _ = try e.generate(&ids, 2, &.{}, 0, .{ .ctx = &quiet, .prefilled = Quiet.prefilled, .tokens = Quiet.tokens, .cancelled = Quiet.cancelled }); } - e.load_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) / 24e6; + e.load_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) / 24e6; return e; } @@ -324,7 +325,7 @@ pub const Engine = struct { e.ev += 1; cb.signal(e.event, e.ev); cb.commit(); - e.committed = std.c.mach_absolute_time(); + e.committed = platform.c.mach_absolute_time(); if (spin) { while (@backingInt(cb.status()) < @backingInt(mtl.types.CommandBufferStatus.completed)) std.atomic.spinLoopHint(); } else cb.wait(); @@ -408,7 +409,7 @@ pub const Engine = struct { e.s.reset(); @memcpy(u32s(e.prompt_ids, prompt.len), prompt); var x = e.ctx(); - const t_start = std.c.mach_absolute_time(); + const t_start = platform.c.mach_absolute_time(); // the prompt: windows of up to 16 rows, the MTP head taking each row with the token after it var at: u32 = 0; var last_n: u32 = 1; @@ -456,7 +457,7 @@ pub const Engine = struct { at += n; } e.s.pos = P; - const t_prompt = std.c.mach_absolute_time(); + const t_prompt = platform.c.mach_absolute_time(); out.prefilled(out.ctx); var res: Result = .{ .reason = .length, .prompt_seconds = @as(f64, @floatFromInt(t_prompt - t_start)) / 24e6 }; var tok = u32s(e.sc.picks, 1)[0]; @@ -469,7 +470,7 @@ pub const Engine = struct { if (max_tokens <= 1) return res; if (d == 0 and e.ep == null and e.copy_min == 0 and e.margins == null) { try plain_rounds.run(e, &x, tok, max_tokens, eos, out, &res); - res.decode_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t_prompt)) / 24e6; + res.decode_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t_prompt)) / 24e6; return res; } var emitted: usize = 1; @@ -490,7 +491,7 @@ pub const Engine = struct { u32s(e.sc.next, 1)[0] = tok; var last_end: f64 = 0; while (true) { - const t_enc = std.c.mach_absolute_time(); + const t_enc = platform.c.mach_absolute_time(); const b = e.begin(); if (window > 0) { // the last window's kept rows (the prompt's windows kept all theirs) fwd.keepKda(&x, b.enc, window, keep); @@ -584,7 +585,7 @@ pub const Engine = struct { break; } } - res.decode_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t_prompt)) / 24e6; + res.decode_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t_prompt)) / 24e6; return res; } }; diff --git a/zig/src/families/glm/ep.zig b/zig/src/families/glm/ep.zig index 74ea506c26..1fd8a94b9f 100644 --- a/zig/src/families/glm/ep.zig +++ b/zig/src/families/glm/ep.zig @@ -1,5 +1,6 @@ //! Expert parallel over two Macs (MCDMA): half the routed experts each, a MoE call's (row, slot) outputs swapped, all combined on both. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const settings = @import("../flashnext/tp_settings.zig"); @@ -120,8 +121,8 @@ pub const Ep = struct { t.ctl = try control.Control.init(gpa, rd, CONTROL, &t.failed); errdefer t.ctl.deinit(gpa); _ = try t.ctl.hello(me); // both Macs up, running the same thing, before the first exchange - t.trace = std.c.getenv("GLM_EP_TRACE") != null; - if (std.c.getenv("GLM_EP_DELAY_US")) |v| t.delay_ns = 1000 * (std.fmt.parseInt(u64, std.mem.span(v), 10) catch 0); + t.trace = platform.c.getenv("GLM_EP_TRACE") != null; + if (platform.c.getenv("GLM_EP_DELAY_US")) |v| t.delay_ns = 1000 * (std.fmt.parseInt(u64, std.mem.span(v), 10) catch 0); t.thread = try std.Thread.spawn(.{}, service, .{t}); std.log.info("expert parallel: rank {d} of 2 holds experts {d}-{d}, connected", .{ t.rank, me.own_lo, me.own_hi - 1 }); return t; @@ -307,7 +308,7 @@ pub const Ep = struct { /// Each posted exchange's entries to the peer, then its flag; 10 s without an answer fails the link and lands every flag. fn service(t: *Ep) void { - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send const posted = t.word32(POSTED); const flag = t.word64(FLAG); var x: u64 = 1; @@ -320,26 +321,26 @@ pub const Ep = struct { t.st.watch(@atomicLoad(u64, flag, .acquire), x - 1); if (t.stop.load(.acquire)) return; if (waiting and @atomicLoad(u64, flag, .acquire) >= want) waiting = false; - if (waiting and !t.failed.load(.acquire) and std.c.mach_absolute_time() - want_at > 240_000_000) { + if (waiting and !t.failed.load(.acquire) and platform.c.mach_absolute_time() - want_at > 240_000_000) { t.fail(want, error.PeerSilent); t.land(want); waiting = false; } spins += 1; if (spins > 1_000_000) { // idle: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } - const seen = std.c.mach_absolute_time(); + const seen = platform.c.mach_absolute_time(); const n: usize = @atomicLoad(u32, t.word32(COUNT + 4 * @as(usize, @intCast(x % 2))), .acquire); const theirs: u32 = @bitCast(@as([*]const i32, @ptrCast(@alignCast(t.lists.contents() + COUNTS)))[1]); t.st.posted(x, seen, @intCast(n), theirs); t.st.watch(@atomicLoad(u64, flag, .acquire), x); if (t.trace) std.debug.print("EP rank{d} exchange {d}: {d} entries\n", .{ t.rank, x, n }); if (t.delay_ns > 0 and n * ENTRY > 1 << 20) { - const ts: std.c.timespec = .{ .sec = @intCast(t.delay_ns / std.time.ns_per_s), .nsec = @intCast(t.delay_ns % std.time.ns_per_s) }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = @intCast(t.delay_ns / std.time.ns_per_s), .nsec = @intCast(t.delay_ns % std.time.ns_per_s) }; + _ = platform.c.nanosleep(&ts, null); } if (t.failed.load(.acquire) or n * ENTRY > SLOT) { if (n * ENTRY > SLOT) t.fail(x, error.BadCount); @@ -349,7 +350,7 @@ pub const Ep = struct { t.land(x); }; want = x; - want_at = std.c.mach_absolute_time(); + want_at = platform.c.mach_absolute_time(); waiting = true; t.sent += 1; t.held_ticks += want_at - seen; @@ -394,7 +395,7 @@ pub const Ep = struct { /// The flag's new value seen now; every exchange up to `upto` whose post and landing are both known is counted. fn watch(s: *Stats, flag: u64, upto: u64) void { if (flag > s.flag_seen) { - const now = std.c.mach_absolute_time(); + const now = platform.c.mach_absolute_time(); var v = s.flag_seen + 1; while (v <= flag) : (v += 1) s.flag_t[v % 64] = now; s.flag_seen = flag; diff --git a/zig/src/families/glm/ep_control.zig b/zig/src/families/glm/ep_control.zig index 420d7411d4..4e69e3e39b 100644 --- a/zig/src/families/glm/ep_control.zig +++ b/zig/src/families/glm/ep_control.zig @@ -1,5 +1,6 @@ //! The expert-parallel pair's control protocol over a link's window: identities, requests and stop decisions, each acked. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const Rdma = fabric.rdma.Rdma; @@ -54,7 +55,7 @@ pub const Request = struct { max_tokens: usize, depth: usize, eos: []const u32, pub const Command = union(enum) { request: Request, lanes: struct { kind: u32, words: []const u32 } }; fn ticks() u64 { - return std.c.mach_absolute_time(); // 24 MHz on Apple silicon + return platform.c.mach_absolute_time(); // 24 MHz on Apple silicon } pub const Control = struct { @@ -96,8 +97,8 @@ pub const Control = struct { if (bounded and ticks() - t0 > c.limit) return error.EpPeerSilent; spins += 1; if (spins > 100_000) { // a long wait: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 50_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } } @@ -295,8 +296,8 @@ test "early cancels: rank 0 waits for rank 1's acks, nothing is overwritten or s }; var e0: ?anyerror = undefined; const th = try std.Thread.spawn(.{}, Rank0.run, .{ &p.c[0], &e0 }); - const ts: std.c.timespec = .{ .sec = 0, .nsec = 200_000_000 }; // rank 1 held back while rank 0 runs ahead - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 200_000_000 }; // rank 1 held back while rank 0 runs ahead + _ = platform.c.nanosleep(&ts, null); for (prompts, 0..) |pr, i| { const r = (try p.c[1].waitRequest()).?; try std.testing.expectEqualSlices(u32, &pr, r.prompt); diff --git a/zig/src/families/glm/learned_pair.zig b/zig/src/families/glm/learned_pair.zig index 4d15eb992f..b1c8c57321 100644 --- a/zig/src/families/glm/learned_pair.zig +++ b/zig/src/families/glm/learned_pair.zig @@ -1,5 +1,6 @@ //! A peer's disk reservation covers its snapshot half until rank 0 has committed or removed the shared index entry. const std = @import("std"); +const platform = @import("sys"); const disk = @import("lanes").learned_disk; pub const State = struct { admission: disk.Admission = .{}, @@ -98,12 +99,12 @@ test "each rank includes actual other-identity files in its independent cap befo var filename: [1400]u8 = undefined; const current = try std.fmt.bufPrintSentinel(&one, "{s}/0000000000000001", .{root}, 0); const other = try std.fmt.bufPrintSentinel(&two, "{s}/0000000000000002", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(current, 0o700)); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(other, 0o700)); - const fd = std.c.open(try std.fmt.bufPrintSentinel(&filename, "{s}/state.bin", .{other}, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(current, 0o700)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(other, 0o700)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&filename, "{s}/state.bin", .{other}, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestCreate; - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 96)); - _ = std.c.close(fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 96)); + _ = platform.c.close(fd); var rank0: State = .{ .admission = .{ .floor = 0, .free = Fake.free, .clock = Fake.clock }, .cap = 200 }; var rank1: State = .{ .admission = .{ .floor = 0, .free = Fake.free, .clock = Fake.clock }, .cap = 100 }; try std.testing.expectEqual(@as(u64, 0), rank0.need(current, 8, 2)); diff --git a/zig/src/families/glm/load_plan.zig b/zig/src/families/glm/load_plan.zig index 53a0dddbcd..94d2b98b68 100644 --- a/zig/src/families/glm/load_plan.zig +++ b/zig/src/families/glm/load_plan.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's load plan: the checkpoint's identity, the bytes its weights and caches take, the chunk height, the limit. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -137,8 +138,8 @@ pub fn limitFor(ram: u64, working_set: u64, text: ?[]const u8) error{BadLoadLimi pub fn loadLimit(device: mtl.Device) error{BadLoadLimit}!Limit { var mem: u64 = 0; var len: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; // RAM unknown: a zero limit refuses every load - const text: ?[]const u8 = if (std.c.getenv("GLM_LOAD_LIMIT_GB")) |v| std.mem.span(v) else null; + if (platform.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; // RAM unknown: a zero limit refuses every load + const text: ?[]const u8 = if (platform.c.getenv("GLM_LOAD_LIMIT_GB")) |v| std.mem.span(v) else null; const l = limitFor(mem, device.maxWorkingSet(), text) catch |err| { std.log.err("glm: GLM_LOAD_LIMIT_GB={s} is not a load limit: give a number of GB above 0, or unset it for 70% of RAM", .{text.?}); return err; diff --git a/zig/src/families/glm/mirror.zig b/zig/src/families/glm/mirror.zig index 6a3a7dbb70..4c9d862f9f 100644 --- a/zig/src/families/glm/mirror.zig +++ b/zig/src/families/glm/mirror.zig @@ -1,5 +1,6 @@ //! Speed-up mode's slot commands: rank 0 sends each slot operation before it runs it; rank 1 runs the same, in order. const std = @import("std"); +const platform = @import("sys"); const eng = @import("engine.zig"); const slots_mod = @import("slots.zig"); const st = @import("state.zig"); @@ -192,10 +193,10 @@ fn learned(sl: *Slots, k: Kind, w: []const u32) !void { } pub fn fileBytes(file: [:0]const u8) u64 { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return 0; - defer _ = std.c.close(fd); - const n = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const n = platform.c.lseek(fd, 0, platform.c.SEEK.END); return if (n > 0) @intCast(n) else 0; } diff --git a/zig/src/families/glm/names.zig b/zig/src/families/glm/names.zig index abea94eabb..3058b447e8 100644 --- a/zig/src/families/glm/names.zig +++ b/zig/src/families/glm/names.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's checkpoint index in either layout (the original or mlx-lm's), each tensor by its original name. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("../../core/safetensors.zig"); const load_fix = @import("weights_fix.zig"); @@ -45,15 +46,15 @@ pub fn read(gpa: std.mem.Allocator, arena: std.mem.Allocator, dir: []const u8, m for (files.keys(), 0..) |name, si| { const full = try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ dir, name }, 0); try shards.append(gpa, full); - const fd = std.c.open(full, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(full, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: [8]u8 = undefined; try load_fix.readAll(fd, &head, 0); const hlen = std.mem.readInt(u64, &head, .little); const header = try arena.alloc(u8, hlen); try load_fix.readAll(fd, header, 8); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); const data = 8 + hlen; const doc2 = try std.json.parseFromSliceLeaky(std.json.Value, arena, header, .{}); var eit = doc2.object.iterator(); diff --git a/zig/src/families/glm/page_cache.zig b/zig/src/families/glm/page_cache.zig index ab64f39d75..cdd9fb2ab9 100644 --- a/zig/src/families/glm/page_cache.zig +++ b/zig/src/families/glm/page_cache.zig @@ -1,5 +1,6 @@ //! A checkpoint's pages in the file cache, dropped before a load: F_NOCACHE reads keep out of it but evict nothing. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const paced_read = @import("../../core/paced_read.zig"); @@ -33,15 +34,15 @@ pub fn dropCached(gpa: std.mem.Allocator, dir: []const u8, want: u64, d: *Droppe const page: usize = std.heap.pageSize(); for (files.keys()) |name| { if (d.dropped >= want) break; - const fd = std.c.open(try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end <= 0) continue; const len: usize = @intCast(end); - const map = std.c.mmap(null, len, .{ .READ = true }, .{ .TYPE = .SHARED }, fd, 0); - if (map == std.c.MAP_FAILED) return error.MapFailed; - defer _ = std.c.munmap(@alignCast(map), len); + const map = platform.c.mmap(null, len, .{ .READ = true }, .{ .TYPE = .SHARED }, fd, 0); + if (map == platform.c.MAP_FAILED) return error.MapFailed; + defer _ = platform.c.munmap(@alignCast(map), len); const vec = try a.alloc(u8, (len + page - 1) / page); var cached: u64 = 0; if (mincore(map, len, vec.ptr) == 0) for (vec) |v| { @@ -57,31 +58,31 @@ pub fn dropCached(gpa: std.mem.Allocator, dir: []const u8, want: u64, d: *Droppe test "a checkpoint's cached pages are dropped only as far as wanted, and a shard with none is left alone" { const gpa = std.testing.allocator; var name_buf: [64]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-drop-cached-{d}", .{std.c.getpid()}, 0); - _ = std.c.mkdir(dir, 0o700); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-drop-cached-{d}", .{platform.c.getpid()}, 0); + _ = platform.c.mkdir(dir, 0o700); + defer _ = platform.c.rmdir(dir); var path_buf: [128]u8 = undefined; const index = try std.fmt.bufPrintSentinel(&path_buf, "{s}/model.safetensors.index.json", .{dir}, 0); const text = "{\"weight_map\": {\"a\": \"s.safetensors\", \"b\": \"s.safetensors\"}}"; var shard_buf: [128]u8 = undefined; const shard = try std.fmt.bufPrintSentinel(&shard_buf, "{s}/s.safetensors", .{dir}, 0); for ([_][:0]const u8{ index, shard }, [_][]const u8{ text, "" }) |file, body| { - const fd = std.c.open(file, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); try std.testing.expect(fd >= 0); - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); if (body.len > 0) { - _ = std.c.write(fd, body.ptr, body.len); + _ = platform.c.write(fd, body.ptr, body.len); continue; } const data = try gpa.alloc(u8, 4 << 20); defer gpa.free(data); @memset(data, 7); - try std.testing.expectEqual(@as(isize, @intCast(data.len)), std.c.write(fd, data.ptr, data.len)); - _ = std.c.fsync(fd); - _ = std.c.pread(fd, data.ptr, data.len, 0); // read through the file cache + try std.testing.expectEqual(@as(isize, @intCast(data.len)), platform.c.write(fd, data.ptr, data.len)); + _ = platform.c.fsync(fd); + _ = platform.c.pread(fd, data.ptr, data.len, 0); // read through the file cache } - defer _ = std.c.unlink(index); - defer _ = std.c.unlink(shard); + defer _ = platform.c.unlink(index); + defer _ = platform.c.unlink(shard); var none: Dropped = .{ .deficit = 0 }; try dropCached(gpa, dir, 0, &none); // nothing wanted: no shard touched, the pages stay cached try std.testing.expectEqual(@as(u64, 0), none.dropped); diff --git a/zig/src/families/glm/plain_rounds.zig b/zig/src/families/glm/plain_rounds.zig index c5f2114916..3289fdc513 100644 --- a/zig/src/families/glm/plain_rounds.zig +++ b/zig/src/families/glm/plain_rounds.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's plain greedy rounds in generate: the next round committed before this one's pick is read. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fwd = @import("forward.zig"); const Fence = @import("../../core/fence.zig").Fence; @@ -12,7 +13,7 @@ const Pending = struct { cb: mtl.CommandBuffer, t_enc: u64, committed: u64 }; /// One plain round at `pos`, its token from `from` (null: the ids as set), its pick into `to`; after a round it waits on the fence, not the event. fn plainRound(e: *Engine, x: *fwd.Ctx, from: ?Ref, to: Ref, pos: u32) Pending { - const t_enc = std.c.mach_absolute_time(); + const t_enc = platform.c.mach_absolute_time(); const fence = e.round_fence.?; const cb = e.queue.commandBuffer(); if (from == null and e.ev > 0) cb.waitFor(e.event, e.ev); @@ -31,7 +32,7 @@ fn plainRound(e: *Engine, x: *fwd.Ctx, from: ?Ref, to: Ref, pos: u32) Pending { e.ev += 1; b.cb.signal(e.event, e.ev); b.cb.commit(); - e.committed = std.c.mach_absolute_time(); + e.committed = platform.c.mach_absolute_time(); return .{ .cb = b.cb, .t_enc = t_enc, .committed = e.committed }; } diff --git a/zig/src/families/glm/pool.zig b/zig/src/families/glm/pool.zig index fb14809842..56f36077f0 100644 --- a/zig/src/families/glm/pool.zig +++ b/zig/src/families/glm/pool.zig @@ -1,5 +1,6 @@ //! One Mac's streams share the engine's MLA cache, a contiguous region each; kept prefixes hold blocks until needed. const std = @import("std"); +const platform = @import("sys"); const st = @import("state.zig"); const snapshot = @import("snapshot.zig"); const Engine = @import("engine.zig").Engine; @@ -41,7 +42,7 @@ pub const Pool = struct { /// Whether the slots share one pool: on one Mac (a pair's ranks mirror slot commands), unless GLM_POOL=0. pub fn on(e: *const Engine) bool { if (e.ep != null or e.followsPeer()) return false; - const v = std.c.getenv("GLM_POOL") orelse return true; + const v = platform.c.getenv("GLM_POOL") orelse return true; return !std.mem.eql(u8, std.mem.span(v), "0"); } diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index 4f73d916b0..4be9a124cf 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -1,5 +1,6 @@ //! Prompt chunks: matmuls on the M5 tensor units at the row kernels' arithmetic (core/affine_mm.zig), the rest on the row kernels. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -83,7 +84,7 @@ pub fn init(gpa: std.mem.Allocator, arena: *st.Arena, device: mtl.Device, c: *co /// The prompt's sparse MLA on the tensor units (GLM_SPARSE_NAX=0: the row kernel, its output beside its queries). fn sparseNax() bool { - return if (std.c.getenv("GLM_SPARSE_NAX")) |v| v[0] != '0' else true; + return if (platform.c.getenv("GLM_SPARSE_NAX")) |v| v[0] != '0' else true; } /// The bytes `init` takes from the arena for `rows`-row chunks. diff --git a/zig/src/families/glm/slots.zig b/zig/src/families/glm/slots.zig index 90aa566962..fc97750612 100644 --- a/zig/src/families/glm/slots.zig +++ b/zig/src/families/glm/slots.zig @@ -1,5 +1,6 @@ //! Stream slots on one Mac: a stream's caches each; prompt chunks, shared windows, keeps and drafts run on them. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("state.zig"); const fwd = @import("forward.zig"); @@ -76,7 +77,7 @@ pub const Slots = struct { .last = try e.arena.buffer(@as(usize, e.c.hidden) * 2), }; const plane = @as(usize, st.max_rows) * e.c.hidden * 2; - var sl: Slots = .{ .gpa = gpa, .e = e, .slots = slots, .log = std.c.getenv("GLM_WINDOWS") != null, .rows = try e.arena.buffer(plane), .ids = try e.arena.buffer(st.max_rows * 4), .lasts = try e.arena.buffer(plane), .picks = try e.arena.buffer(st.max_rows * 4), .ring = try e.arena.buffer(ring_len * 4) }; + var sl: Slots = .{ .gpa = gpa, .e = e, .slots = slots, .log = platform.c.getenv("GLM_WINDOWS") != null, .rows = try e.arena.buffer(plane), .ids = try e.arena.buffer(st.max_rows * 4), .lasts = try e.arena.buffer(plane), .picks = try e.arena.buffer(st.max_rows * 4), .ring = try e.arena.buffer(ring_len * 4) }; if (pool_mod.on(e)) pool_mod.setUp(&sl); return sl; } @@ -247,7 +248,7 @@ pub const Slots = struct { const e = sl.e; var segs: [st.max_rows]fwd.Seg = undefined; if (wins.len == 0 or wins.len > segs.len) return error.WindowOutOfStep; - const t_in = std.c.mach_absolute_time(); + const t_in = platform.c.mach_absolute_time(); sl.settleAll(); try sl.launch(); // the GPU runs the keeps and drafts while the window encodes const ids = Engine.u32s(e.sc.ids, st.max_rows); @@ -274,7 +275,7 @@ pub const Slots = struct { fwd.head(&x, enc, e.sc.hidden, e.sc.logits, e.sc.picks, total); sl.ahead = .{ 0, 0 }; try sl.flushBy(true); - const ret = std.c.mach_absolute_time(); + const ret = platform.c.mach_absolute_time(); if (sl.log) { const first = if (sl.ahead[1] > 0) sl.ahead[0] else e.gpu[0]; std.log.info("glm window: streams {d} rows {d} gpu {d:.2} ms gap {d:.2} ms ahead {d:.2} join {d:.3} host {d:.2} encode {d:.2} ms", .{ wins.len, total, (e.gpu[1] - e.gpu[0]) * 1e3, (first - sl.last_end) * 1e3, (sl.ahead[1] - sl.ahead[0]) * 1e3, if (sl.ahead[1] > 0) (e.gpu[0] - sl.ahead[1]) * 1e3 else 0, @as(f64, @floatFromInt(t_in -| sl.last_ret)) / 24e3, @as(f64, @floatFromInt(e.committed - t_in)) / 24e3 }); diff --git a/zig/src/families/glm/snapshot_file.zig b/zig/src/families/glm/snapshot_file.zig index 901b5d90b6..b8b058d7d3 100644 --- a/zig/src/families/glm/snapshot_file.zig +++ b/zig/src/families/glm/snapshot_file.zig @@ -1,45 +1,46 @@ //! GLM learned halves use a versioned checksum envelope; legacy, corrupt and partial files never enter a cache. const std = @import("std"); +const platform = @import("sys"); const checked = @import("../../core/snapshot_file.zig"); const magic: u32 = 0x474c4d54; pub const Head = extern struct { magic: u32 = magic, version: u32 = 1, at: u32, reserved: u32 = 0, bytes: u64 }; pub fn write(file: [:0]const u8, at: u32, payload: []const u8) !void { var tmp_buf: [1200]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head: Head = .{ .at = at, .bytes = payload.len }; try checked.write(fd, std.mem.asBytes(&head), payload); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } pub fn read(file: [:0]const u8, at: u32, payload: []u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const expected: Head = .{ .at = at, .bytes = payload.len }; var head: Head = undefined; - if (std.c.pread(fd, std.mem.asBytes(&head).ptr, @sizeOf(Head), 0) != @sizeOf(Head) or !std.mem.eql(u8, std.mem.asBytes(&head), std.mem.asBytes(&expected))) return error.SnapshotRead; + if (platform.c.pread(fd, std.mem.asBytes(&head).ptr, @sizeOf(Head), 0) != @sizeOf(Head) or !std.mem.eql(u8, std.mem.asBytes(&head), std.mem.asBytes(&expected))) return error.SnapshotRead; try checked.read(fd, std.mem.asBytes(&head), payload); } test "both GLM half files refuse a changed body, legacy header, truncation and another position" { var path: [160]u8 = undefined; - const file = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-glm-file-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(file); + const file = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-glm-file-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(file); try write(file, 128, "body"); var out: [4]u8 = undefined; try read(file, 128, &out); try std.testing.expectEqualStrings("body", &out); try std.testing.expectError(error.SnapshotRead, read(file, 129, &out)); - const fd = std.c.open(file, .{ .ACCMODE = .RDWR }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDWR }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - try std.testing.expectEqual(@as(isize, 1), std.c.pwrite(fd, "X", 1, @sizeOf(Head) + 1)); + defer _ = platform.c.close(fd); + try std.testing.expectEqual(@as(isize, 1), platform.c.pwrite(fd, "X", 1, @sizeOf(Head) + 1)); try std.testing.expectError(error.SnapshotRead, read(file, 128, &out)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 16)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 16)); try std.testing.expectError(error.SnapshotRead, read(file, 128, &out)); } diff --git a/zig/src/families/glm/weights.zig b/zig/src/families/glm/weights.zig index 5beb758f14..d85bebec5b 100644 --- a/zig/src/families/glm/weights.zig +++ b/zig/src/families/glm/weights.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's weights in the Python family's layout (stacked projections and experts, repacked mixes), read by parallel preads. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("../../core/safetensors.zig"); const cfg = @import("config.zig"); diff --git a/zig/src/families/glm/weights_fix.zig b/zig/src/families/glm/weights_fix.zig index ec0861fe6e..3d12f8b75b 100644 --- a/zig/src/families/glm/weights_fix.zig +++ b/zig/src/families/glm/weights_fix.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash weights as they load: parallel reads into place, and rewrites (mixes, router, conv, norms, columns). const std = @import("std"); +const platform = @import("sys"); const st = @import("../../core/safetensors.zig"); const paced_read = @import("../../core/paced_read.zig"); @@ -12,16 +13,16 @@ pub const Transform = struct { src: Src, dst: [*]u8, fix: Fix, extra: [3]?Src = /// One tensor's bytes read straight into their place in a weights buffer. pub const Copy = struct { shard: u32, off: u64, len: u64, dst: [*]u8 }; -pub fn readAll(fd: std.c.fd_t, dest: []u8, at: u64) !void { +pub fn readAll(fd: platform.c.fd_t, dest: []u8, at: u64) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } } -pub fn transform(gpa: std.mem.Allocator, fds: []std.c.fd_t, t: Transform) !void { +pub fn transform(gpa: std.mem.Allocator, fds: []platform.c.fd_t, t: Transform) !void { const raw = try gpa.alloc(u8, t.src.len); defer gpa.free(raw); try readAll(fds[t.src.shard], raw, t.src.off); @@ -110,7 +111,7 @@ test "fp32 values that are bf16 narrow in place; any other value is refused" { } const Pool = struct { - fds: []std.c.fd_t, + fds: []platform.c.fd_t, jobs: []const Copy, next: std.atomic.Value(usize) = .init(0), failed: std.atomic.Value(bool) = .init(false), @@ -132,15 +133,15 @@ const Pool = struct { /// Every copy on `threads` threads, uncached, paced by free memory (a large load otherwise fills the page cache too). pub fn runCopies(gpa: std.mem.Allocator, shards: []const [:0]const u8, copies: []const Copy, transforms: []const Transform, threads: usize) !void { - const fds = try gpa.alloc(std.c.fd_t, shards.len); + const fds = try gpa.alloc(platform.c.fd_t, shards.len); defer gpa.free(fds); for (shards, fds) |path, *fd| { - fd.* = std.c.open(path, .{ .ACCMODE = .RDONLY }); + fd.* = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd.* < 0) return error.OpenFailed; - _ = std.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE + _ = platform.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE } defer for (fds) |fd| { - _ = std.c.close(fd); + _ = platform.c.close(fd); }; var pool: Pool = .{ .fds = fds, .jobs = copies }; const workers = try gpa.alloc(?std.Thread, threads); @@ -153,7 +154,7 @@ pub fn runCopies(gpa: std.mem.Allocator, shards: []const [:0]const u8, copies: [ const Tx = struct { gpa: std.mem.Allocator, jobs: []const Transform, - fds: []std.c.fd_t, + fds: []platform.c.fd_t, next: std.atomic.Value(usize) = .init(0), failed: std.atomic.Value(bool) = .init(false), fn run(x: *@This()) void { diff --git a/zig/src/families/glm/weights_test.zig b/zig/src/families/glm/weights_test.zig index 87e0c7fd01..92dc8ea962 100644 --- a/zig/src/families/glm/weights_test.zig +++ b/zig/src/families/glm/weights_test.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's loader on synthetic checkpoints: one sparse shard whose header names every tensor at its width. const std = @import("std"); +const platform = @import("sys"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); const quant = @import("quant.zig"); @@ -156,10 +157,10 @@ const Fixture = struct { std.mem.writeInt(u64, &len, head.items.len, .little); try tmp.dir.writeFile(io, .{ .sub_path = "model-00001.safetensors", .data = try std.mem.concat(f.a, u8, &.{ &len, head.items }) }); const path = try std.fmt.allocPrintSentinel(f.a, "{s}/model-00001.safetensors", .{dir}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(8 + head.items.len + at)) != 0) return error.TestFile; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(8 + head.items.len + at)) != 0) return error.TestFile; } }; diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index f52e5a4a6c..ee3c374e4c 100644 --- a/zig/src/families/nemotron/backend.zig +++ b/zig/src/families/nemotron/backend.zig @@ -1,5 +1,6 @@ //! The lane core's Metal backend for Nemotron: streams' rows in one forward a round, drawn tokens in GPU slots. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const Fence = @import("../../core/fence.zig").Fence; @@ -89,7 +90,7 @@ pub const Metal = struct { .tokens = try m.device.buffer(ring * 4, opts), .prompt = try m.device.buffer(@max(o.capacity, 16) * 4, opts), .fence = try Fence.init(m.device), - .trace = std.c.getenv("TF_SLIDE_TRACE") != null, + .trace = platform.c.getenv("TF_SLIDE_TRACE") != null, }; if (o.chunk > fused_rows) b.wide = try prefill.Prefill.init(m, @min(o.chunk, o.capacity)); if (o.drafts and m.weights.mtp != null) b.head = try mtp.Head.init(m.device, m.config, &m.kernels, &m.weights, &b.pool, max_lanes, o.capacity); diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index f287c4ced4..7483f02db6 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -1,6 +1,7 @@ //! The Nemotron CUDA engine (engine.py): weights, kernels and one sequence's buffers; windows and commits. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const Config = @import("config.zig").Config; const kern = @import("cuda_kernels.zig"); @@ -130,7 +131,7 @@ pub const Engine = struct { const e = try gpa.create(Engine); errdefer gpa.destroy(e); e.* = .{ .gpa = gpa, .io = io, .ctx = ctx, .stream = undefined, .k = undefined, .w = undefined, .b = undefined, .c = undefined, .max_len = 0, .pinned = undefined, .history = undefined }; - e.trace = std.c.getenv("TF_SLIDE_TRACE") != null; + e.trace = platform.c.getenv("TF_SLIDE_TRACE") != null; e.c = try Config.read(gpa, io, model_dir); if (opts.segments < 1 or opts.segments > segs.MAX) return error.BadSegments; e.segments = opts.segments; diff --git a/zig/src/families/nemotron/cuda_native.zig b/zig/src/families/nemotron/cuda_native.zig index 64594342c4..f09ad552e6 100644 --- a/zig/src/families/nemotron/cuda_native.zig +++ b/zig/src/families/nemotron/cuda_native.zig @@ -1,5 +1,6 @@ //! Nemotron's CUDA engine, MTP head and lane backend behind the native family interface. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); @@ -134,7 +135,7 @@ fn release(p: *anyopaque) void { /// TF_CUDA_CARVEOUT=1: the display carveout KV planes go to first (cuda/carveout.zig); null when it is off. fn openCarveout(gpa: std.mem.Allocator, ctx: *const cuda.Context) !?*cuda.Carveout { const bytes = (try cuda.carveout.requested(false, env("TF_CUDA_CARVEOUT"), env("TF_CUDA_CARVEOUT_MIB"))) orelse return null; - const card: [*:0]const u8 = std.c.getenv("TF_DRM_CARD") orelse cuda.carveout.default_card; + const card: [*:0]const u8 = platform.c.getenv("TF_DRM_CARD") orelse cuda.carveout.default_card; const c = try gpa.create(cuda.Carveout); errdefer gpa.destroy(c); c.* = cuda.Carveout.open(ctx.d, card, bytes) catch |err| { @@ -151,5 +152,5 @@ fn closeCarveout(gpa: std.mem.Allocator, c: *cuda.Carveout) void { } fn env(name: [:0]const u8) ?[]const u8 { - return std.mem.span(std.c.getenv(name) orelse return null); + return std.mem.span(platform.c.getenv(name) orelse return null); } diff --git a/zig/src/families/nemotron/cuda_seq_graphs.zig b/zig/src/families/nemotron/cuda_seq_graphs.zig index 9e51e1556b..09fb6965a4 100644 --- a/zig/src/families/nemotron/cuda_seq_graphs.zig +++ b/zig/src/families/nemotron/cuda_seq_graphs.zig @@ -1,6 +1,7 @@ //! Graphs a sequence other than the engine's own captures on its own buffers, once each, and replays after. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); /// One sequence's captured graphs, by the step they run (a window's or a head level's id). @@ -27,6 +28,6 @@ pub const Stats = struct { /// Whether TF_GRAPH_LOG=1 asks for a line a request: its captures, their time, its replays and eager windows. pub fn logFromEnv() bool { - const v = std.c.getenv("TF_GRAPH_LOG") orelse return false; + const v = platform.c.getenv("TF_GRAPH_LOG") orelse return false; return std.mem.eql(u8, std.mem.span(v), "1"); } diff --git a/zig/src/families/nemotron/cuda_source.zig b/zig/src/families/nemotron/cuda_source.zig index ffc4ba2eec..f4a1f07fdd 100644 --- a/zig/src/families/nemotron/cuda_source.zig +++ b/zig/src/families/nemotron/cuda_source.zig @@ -2,6 +2,8 @@ const std = @import("std"); const linux = std.os.linux; +/// io_uring is Linux's; elsewhere every read runs on its own thread. +const has_ring = @import("builtin").os.tag == .linux; const cuda = @import("cuda"); const core = @import("core"); const kern = @import("cuda_kernels.zig"); @@ -42,7 +44,7 @@ pub const Source = struct { copied: [slot_count]?cuda.Event = @splat(null), // a page-locked slot's last copy: its next read waits for that, not the stream pending: [slot_count]?Pending = @splat(null), in_flight: usize = 0, - ring: ?linux.IoUring = null, // null where io_uring is unavailable: each read then runs on its own thread + ring: if (has_ring) ?linux.IoUring else ?void = null, // null where io_uring is unavailable: each read then runs on its own thread jobs: [slot_count]Job = undefined, asked: [slot_count]u64 = @splat(0), // the order reads were asked in: copies follow it without a ring next: u64 = 0, @@ -69,7 +71,7 @@ pub const Source = struct { s.pinned[made] = h; s.slots[made] = @alignCast(h.bytes); } - s.ring = linux.IoUring.init(slot_count, 0) catch |err| blk: { + if (has_ring) s.ring = linux.IoUring.init(slot_count, 0) catch |err| blk: { std.log.info("io_uring unavailable ({t}): checkpoint reads run on threads", .{err}); break :blk null; }; @@ -79,7 +81,7 @@ pub const Source = struct { /// Waits for reads still in flight, then frees the ring and slots and closes the files. pub fn deinit(s: *Source) void { while (s.in_flight > 0) s.reap(1) catch break; - if (s.ring) |*r| r.deinit() else s.joinAll(); + if (!has_ring) s.joinAll() else if (s.ring) |*r| r.deinit() else s.joinAll(); s.freeSlots(slot_count); for (s.files.items) |*m| m.file.close(); s.files.deinit(s.gpa); @@ -156,7 +158,8 @@ pub const Source = struct { const lo = std.mem.alignBackward(u64, p.offset, dio.alignment); var q = p; q.at = @intCast(p.offset - lo); - if (s.ring) |*r| { + if (has_ring and s.ring != null) { + const r = &s.ring.?; _ = try r.read(k, p.file.fd, .{ .buffer = s.slots[k][0..dio.span(p.offset, p.len)] }, lo); _ = try r.submit(); } else { @@ -192,6 +195,7 @@ pub const Source = struct { /// Finished reads, copied on; waits for at least `wait` of them. fn reap(s: *Source, wait: u32) !void { + if (!has_ring) return s.reapThreads(wait); const r = if (s.ring) |*ring| ring else return s.reapThreads(wait); var cqes: [slot_count]linux.io_uring_cqe = undefined; const n = try r.copy_cqes(&cqes, wait); diff --git a/zig/src/families/nemotron/forward.zig b/zig/src/families/nemotron/forward.zig index 7dc93e2de2..10751fbbf7 100644 --- a/zig/src/families/nemotron/forward.zig +++ b/zig/src/families/nemotron/forward.zig @@ -1,5 +1,6 @@ //! One Nemotron-H forward over streams' row segments, encoded as FusedDecode's kernels with their geometry. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -376,7 +377,7 @@ var split_mode: enum(u8) { unread, on, off } = .unread; /// TF_COOP_SPLIT=0 keeps one threadgroup a column tile, as the Python engine dispatches. fn splitOn() bool { if (split_mode == .unread) { - const v = std.c.getenv("TF_COOP_SPLIT"); + const v = platform.c.getenv("TF_COOP_SPLIT"); split_mode = if (v != null and v.?[0] == '0') .off else .on; } return split_mode == .on; diff --git a/zig/src/families/nemotron/gpu_round.zig b/zig/src/families/nemotron/gpu_round.zig index f8db78584b..6515ec7750 100644 --- a/zig/src/families/nemotron/gpu_round.zig +++ b/zig/src/families/nemotron/gpu_round.zig @@ -1,5 +1,6 @@ //! A lone greedy stream's rounds with no host wait: the GPU accepts, keeps rows and drafts; the host reads behind it. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const st = @import("state.zig"); @@ -349,7 +350,7 @@ pub fn run(gpa: std.mem.Allocator, b: *Metal, s: *lanes.Stream, o: Options, hook var mtp_known: usize = mtp0 + 1; var queued: u64 = 0; // rounds submitted var read: u64 = 0; // rounds whose records the host took - const log = std.c.getenv("TF_ROUND_LOG") != null; + const log = platform.c.getenv("TF_ROUND_LOG") != null; var depth = Depth{ .levels = o.depth, .most = o.depth, .fixed = o.fixed }; var windows: [ring]usize = undefined; // each queued round's verify rows var levels = depth.levels; // the levels the last queued head drafts (the next verify's rows - 1) @@ -502,7 +503,7 @@ pub fn run(gpa: std.mem.Allocator, b: *Metal, s: *lanes.Stream, o: Options, hook if (o.level_bench) { // every level live (zero bars), each chain alone in its command buffer r.o.stop = false; - if (std.c.getenv("TF_HEAD_GEOMETRY")) |g| b.head.?.geometry = std.fmt.parseInt(usize, std.mem.span(g), 10) catch 0; + if (platform.c.getenv("TF_HEAD_GEOMETRY")) |g| b.head.?.geometry = std.fmt.parseInt(usize, std.mem.span(g), 10) catch 0; var med: [2]f64 = undefined; for ([_]usize{ 1, @min(8, o.depth) }, 0..) |n, k| { var ms: [9]f64 = undefined; diff --git a/zig/src/families/nemotron/learned_prompts.zig b/zig/src/families/nemotron/learned_prompts.zig index 3f0cdfb1f6..7e7edc8818 100644 --- a/zig/src/families/nemotron/learned_prompts.zig +++ b/zig/src/families/nemotron/learned_prompts.zig @@ -1,5 +1,6 @@ //! Nemotron's learned prompt states (--learn): a kept state's file on disk, and the probe that keys this build's bits. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const snapshot = @import("snapshot.zig"); @@ -20,23 +21,23 @@ pub fn path(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { pub fn writeFile(snap: *const Snap, file: [:0]const u8) !void { var tmp_buf: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head = [HEAD]u32{ FILE_MAGIC, snap.at, @intFromBool(snap.head), @truncate(snap.bytes), @truncate(snap.bytes >> 32) }; try checked.write(fd, std.mem.sliceAsBytes(&head), snap.buf.contents()[0..snap.bytes]); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } /// `file`'s state into `snap`, whose buffer holds its bytes; a file of another position, layout or size is refused. pub fn readFile(snap: *Snap, file: [:0]const u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: [HEAD]u32 = undefined; try get(fd, std.mem.sliceAsBytes(&head), 0); const bytes = @as(u64, head[3]) | @as(u64, head[4]) << 32; @@ -47,7 +48,7 @@ pub fn readFile(snap: *Snap, file: [:0]const u8) !void { fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -56,7 +57,7 @@ fn put(fd: c_int, b: []const u8) !void { fn get(fd: c_int, b: []u8, at: u64) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.SnapshotRead; done += @intCast(n); } @@ -127,7 +128,7 @@ pub fn snapRead(ptr: *anyopaque, dir: [:0]const u8, key: u64, at: u32) anyerror! /// --learn: learned state `key`'s file removed. pub fn snapForget(_: *anyopaque, dir: [:0]const u8, key: u64) void { var buf: [1100]u8 = undefined; - _ = std.c.unlink(path(&buf, dir, key) catch return); + _ = platform.c.unlink(path(&buf, dir, key) catch return); } test "a learned state's file name is its key in hex" { diff --git a/zig/src/families/nemotron/learner.zig b/zig/src/families/nemotron/learner.zig index 0607eb90c1..69d70cb360 100644 --- a/zig/src/families/nemotron/learner.zig +++ b/zig/src/families/nemotron/learner.zig @@ -1,5 +1,6 @@ //! Sliding Weights' lessons, one learner for every backend over a host that gives the trainer, weights and files. const std = @import("std"); +const platform = @import("sys"); const adapters = @import("learn_shape.zig"); const choice = @import("choice.zig"); const router = @import("router.zig"); @@ -96,7 +97,7 @@ pub fn Learner(comptime H: type) type { pocketed: bool = false, // the last lesson put neurons in the pocket pub fn init(gpa: std.mem.Allocator, io: std.Io, b: H.Ctx) Self { - return .{ .gpa = gpa, .io = io, .b = b, .whole = std.c.getenv("TF_SLIDE_ROUTED") == null, .edits = if (has_edits) Edits.init(gpa) else {}, .method = if (has_edits) whole_lessons.method() else .pocket, .probe = if (has_edits) Probes.init(gpa) else {}, .pockets = if (has_pocket) Pockets.init(gpa) else {} }; + return .{ .gpa = gpa, .io = io, .b = b, .whole = platform.c.getenv("TF_SLIDE_ROUTED") == null, .edits = if (has_edits) Edits.init(gpa) else {}, .method = if (has_edits) whole_lessons.method() else .pocket, .probe = if (has_edits) Probes.init(gpa) else {}, .pockets = if (has_pocket) Pockets.init(gpa) else {} }; } pub fn deinit(l: *Self) void { diff --git a/zig/src/families/nemotron/living.zig b/zig/src/families/nemotron/living.zig index 1f0e8b073b..f479c17c2d 100644 --- a/zig/src/families/nemotron/living.zig +++ b/zig/src/families/nemotron/living.zig @@ -1,5 +1,6 @@ //! Saved lessons as shards beside the model's, a file a lesson (its blocks and router); both backends pass in `core`. const std = @import("std"); +const platform = @import("sys"); const Config = @import("config.zig").Config; const shape = @import("learn_shape.zig"); @@ -27,7 +28,7 @@ pub fn Io(comptime core: type) type { while (true) : (k += 1) { const path = try std.fs.path.joinZ(gpa, &.{ dir, try fileName(&buf, k) }); defer gpa.free(path); - if (std.c.unlink(path) != 0) break; + if (platform.c.unlink(path) != 0) break; } return bytes; } @@ -127,7 +128,7 @@ pub fn Io(comptime core: type) type { pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !?Saved { const old = try std.fs.path.joinZ(gpa, &.{ dir, legacy }); defer gpa.free(old); - if (std.c.access(old, std.c.F_OK) == 0) std.log.warn("slide: {s} holds lessons saved by an earlier version; they are not read", .{legacy}); + if (platform.c.access(old, platform.c.F_OK) == 0) std.log.warn("slide: {s} holds lessons saved by an earlier version; they are not read", .{legacy}); var lessons: std.ArrayList(Lesson) = .empty; errdefer { for (lessons.items) |*l| l.file.close(io); diff --git a/zig/src/families/nemotron/model.zig b/zig/src/families/nemotron/model.zig index 8ff4dae263..740846b96d 100644 --- a/zig/src/families/nemotron/model.zig +++ b/zig/src/families/nemotron/model.zig @@ -1,5 +1,6 @@ //! A loaded Nemotron-H model: config, checkpoint buffers, kernel-layout weights and compiled pipelines. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ckpt = @import("../../core/checkpoint_metal.zig"); const shards = @import("../../core/checkpoint.zig"); @@ -84,7 +85,7 @@ pub const Model = struct { if (mtp) { const path = try std.fmt.allocPrintSentinel(allocator, "{s}/mtp-4bit.safetensors", .{dir}, 0); defer allocator.free(path); - if (std.c.access(path, 0) == 0) try m.checkpoint.addFile(m.device, path, "mtp."); + if (platform.c.access(path, 0) == 0) try m.checkpoint.addFile(m.device, path, "mtp."); } const drafting = m.checkpoint.has("mtp.layers.0.eh_proj.weight"); m.draft_ids = if (drafting) try draft_ids.load(allocator, m.config.vocab) else &.{}; diff --git a/zig/src/families/nemotron/pocket.zig b/zig/src/families/nemotron/pocket.zig index ff24ac5ee9..90232ef763 100644 --- a/zig/src/families/nemotron/pocket.zig +++ b/zig/src/families/nemotron/pocket.zig @@ -1,5 +1,6 @@ //! A pocket of new neurons beside each MoE layer's shared expert, where facts are learned; nothing the model had moves. const std = @import("std"); +const platform = @import("sys"); const examples = @import("lesson_examples.zig"); const lesson_route = @import("lesson_route.zig"); const pocket_fit = @import("pocket_fit.zig"); @@ -52,7 +53,7 @@ const Refit = struct { slot: usize, neuron: usize }; /// Whether environment variable `name` is set to `value`. fn set(name: [*:0]const u8, value: []const u8) bool { - const v = std.c.getenv(name) orelse return false; + const v = platform.c.getenv(name) orelse return false; return std.mem.eql(u8, std.mem.span(v), value); } @@ -136,7 +137,7 @@ pub fn Pocket(comptime T: type) type { read: reading.Seen = .{}, // what the last read pass read, which a lesson begun again over the same keeps pub fn init(gpa: std.mem.Allocator) Self { - const asked: Mode = if (std.c.getenv("TF_SLIDE_POCKET_KEY")) |v| std.meta.stringToEnum(Mode, std.mem.span(v)) orelse .auto else .auto; + const asked: Mode = if (platform.c.getenv("TF_SLIDE_POCKET_KEY")) |v| std.meta.stringToEnum(Mode, std.mem.span(v)) orelse .auto else .auto; return .{ .gpa = gpa, .mode = asked, .strict = !set("TF_SLIDE_POCKET_LEADS", "off"), .keeping = !set("TF_SLIDE_POCKET_KEEPS", "off"), .valuing = set("TF_SLIDE_POCKET_VALUES", "on"), .silencing = set("TF_SLIDE_POCKET_SAME", "silent"), .worst = set("TF_SLIDE_POCKET_REPLAY", "worst"), .joining = set("TF_SLIDE_POCKET_GROUP", "joint") }; } @@ -360,7 +361,7 @@ pub fn Pocket(comptime T: type) type { const d = c.hidden; const n = p.slots(); const m = p.labels.items.len; - if (std.c.getenv("TF_SLIDE_POCKET_DUMP")) |dir| pocket_dump.write(p, std.mem.span(dir), p.learned_before, d) catch |e| std.log.warn("slide pocket: the read rows were not written: {t}", .{e}); + if (platform.c.getenv("TF_SLIDE_POCKET_DUMP")) |dir| pocket_dump.write(p, std.mem.span(dir), p.learned_before, d) catch |e| std.log.warn("slide pocket: the read rows were not written: {t}", .{e}); const kinds = try p.gpa.alloc(Kind, 3 * m); defer p.gpa.free(kinds); for (p.labels.items, 0..) |l, i| for (0..3) |k| { diff --git a/zig/src/families/nemotron/pocket_dump.zig b/zig/src/families/nemotron/pocket_dump.zig index 1653589631..b1f5d152a6 100644 --- a/zig/src/families/nemotron/pocket_dump.zig +++ b/zig/src/families/nemotron/pocket_dump.zig @@ -1,11 +1,12 @@ //! A lesson's read rows written out for offline study (TF_SLIDE_POCKET_DUMP=dir): inputs, labels and examples. const std = @import("std"); +const platform = @import("sys"); /// `bytes` written to `path`. fn put(path: [:0]const u8, bytes: []const u8) !void { - const f = std.c.fopen(path, "wb") orelse return error.DumpOpen; - defer _ = std.c.fclose(f); - if (std.c.fwrite(bytes.ptr, 1, bytes.len, f) != bytes.len) return error.DumpWrite; + const f = platform.c.fopen(path, "wb") orelse return error.DumpOpen; + defer _ = platform.c.fclose(f); + if (platform.c.fwrite(bytes.ptr, 1, bytes.len, f) != bytes.len) return error.DumpWrite; } /// Pocket `p`'s read rows under `dir` for `tag`: bf16 [read, slots, D], then their labels and examples as text. diff --git a/zig/src/families/nemotron/whole.zig b/zig/src/families/nemotron/whole.zig index 4fa4a9fc01..b8f6fc8a99 100644 --- a/zig/src/families/nemotron/whole.zig +++ b/zig/src/families/nemotron/whole.zig @@ -1,5 +1,6 @@ //! Whole-mode lessons: pocket neurons (default), a weight edit (TF_SLIDE_EDIT) or a probe (TF_SLIDE_PROBE). const std = @import("std"); +const platform = @import("sys"); const examples = @import("lesson_examples.zig"); const adapters = @import("learn_shape.zig"); @@ -10,8 +11,8 @@ pub const Method = enum { pocket, edit, probe }; /// The method the environment picks for whole lessons. pub fn method() Method { - if (std.c.getenv("TF_SLIDE_PROBE") != null) return .probe; - if (std.c.getenv("TF_SLIDE_EDIT") != null) return .edit; + if (platform.c.getenv("TF_SLIDE_PROBE") != null) return .probe; + if (platform.c.getenv("TF_SLIDE_EDIT") != null) return .edit; return .pocket; } diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 53771c8eda..428f764854 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -1,5 +1,6 @@ //! A native text model keeps immutable checkpoint bytes and prepared affine layouts under one owner. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const core = @import("core"); const cfg = @import("config.zig"); @@ -71,7 +72,7 @@ pub const Model = struct { m.frame = try Frame.init(m.device, c, rows); errdefer m.frame.deinit(); m.ane = null; - const columns = if (std.c.getenv("TF_ANE_COLUMNS")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else ane_columns; + const columns = if (platform.c.getenv("TF_ANE_COLUMNS")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else ane_columns; const share = columns > 0 and m.device.tensorUnits() and mtl.ane.available() and m.weights.uniformMlp(); const fits = if (share_fits) |f| !share or f(ane_mlp.Share.bytes(c, columns)) else true; if (!fits) std.log.info("qwen27: the Neural Engine's prompt share ({d} MiB of fp16 weights) would take this Mac past its load limit; prompts run on the GPU alone", .{ane_mlp.Share.bytes(c, columns) >> 20}); diff --git a/zig/src/families/qwen3_5/backend.zig b/zig/src/families/qwen3_5/backend.zig index 55fc09eecb..ef369077a8 100644 --- a/zig/src/families/qwen3_5/backend.zig +++ b/zig/src/families/qwen3_5/backend.zig @@ -163,106 +163,11 @@ pub const Metal = struct { } }; -/// A row's allowed-token bits (structured output): one bit a logit. -const mask_words = (c.vocab + 31) / 32; +const draws = @import("draw.zig"); +const mask_words = draws.mask_words; +pub const value = draws.value; +pub const draw = draws.draw; -pub fn value(word: u16) f64 { - return @as(f32, @bitCast(@as(u32, word) << 16)); -} - -/// The unchanged top-k candidates for the shared fp64 keyed sampler; a `mask` keeps its allowed ones (-inf's draw). -pub fn draw(gpa: std.mem.Allocator, logits: []const u16, settings: ?lanes.Sampling, position: u64, mask: []const u32) !u32 { - var best: ?usize = null; - var allowed: usize = 0; - for (logits, 0..) |word, id| { - const x = value(word); - if (!std.math.isFinite(x)) return error.NonfiniteQwenLogits; - if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; - allowed += 1; - if (best == null or x > value(logits[best.?])) best = id; - } - if (best == null) return error.NoAllowedQwenToken; - const sampling = settings orelse return @intCast(best.?); - if (sampling.temperature <= 0) return @intCast(best.?); - const k = if (sampling.top_k == 0) allowed else @min(allowed, sampling.top_k); - const values = try gpa.alloc(f64, k); - defer gpa.free(values); - const ids = try gpa.alloc(u64, k); - defer gpa.free(ids); - var filled: usize = 0; - for (logits, 0..) |word, id| { - if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; - const x = value(word); - if (k == allowed) { - values[filled] = x; - ids[filled] = id; - filled += 1; - continue; - } - if (filled == k and x <= values[k - 1]) continue; - var at = @min(filled, k - 1); - while (at > 0 and x > values[at - 1]) : (at -= 1) { - values[at] = values[at - 1]; - ids[at] = ids[at - 1]; - } - values[at] = x; - ids[at] = id; - filled = @min(filled + 1, k); - } - return @intCast(try lanes.sampling.choose(gpa, values, ids, position, sampling)); -} - -test "candidate selection preserves full-vocabulary keyed draws and tie order" { - const gpa = std.testing.allocator; - const words = try gpa.alloc(u16, c.vocab); - defer gpa.free(words); - const values = try gpa.alloc(f64, c.vocab); - defer gpa.free(values); - const ids = try gpa.alloc(u64, c.vocab); - defer gpa.free(ids); - for (words, values, ids, 0..) |*w, *v, *id, i| { - const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; - w.* = @intCast(@as(u32, @bitCast(x)) >> 16); - v.* = value(w.*); - id.* = i; - } - for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { - for ([_]u64{ 0, 128, 65536 }) |position| { - const s: lanes.Sampling = .{ .seed = 1234, .top_k = k, .temperature = 0.7, .top_p = 0.95, .min_p = 0.1 }; - const want = try lanes.sampling.choose(gpa, values, ids, position, s); - try std.testing.expectEqual(want, try draw(gpa, words, s, position, &.{})); - } - } - try std.testing.expectEqual(try draw(gpa, words, null, 0, &.{}), try draw(gpa, words, .{ .seed = 7, .temperature = 0 }, 99, &.{})); - words[0] = 0x7fc0; - try std.testing.expectError(error.NonfiniteQwenLogits, draw(gpa, words, null, 0, &.{})); -} - -test "a masked draw is the full-vocabulary draw with the other tokens at -inf" { - const gpa = std.testing.allocator; - const words = try gpa.alloc(u16, c.vocab); - defer gpa.free(words); - const values = try gpa.alloc(f64, c.vocab); - defer gpa.free(values); - const ids = try gpa.alloc(u64, c.vocab); - defer gpa.free(ids); - var mask: [mask_words]u32 = @splat(0); - for (words, values, ids, 0..) |*w, *v, *id, i| { - const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; - w.* = @intCast(@as(u32, @bitCast(x)) >> 16); - const on = i % 7 == 3 or i % 1000 == 1; - if (on) mask[i / 32] |= @as(u32, 1) << @intCast(i % 32); - v.* = if (on) value(w.*) else -std.math.inf(f64); - id.* = i; - } - for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { - for ([_]u64{ 0, 128, 65536 }) |position| { - const s: lanes.Sampling = .{ .seed = 99, .top_k = k, .temperature = 0.8, .top_p = 0.9, .min_p = 0.05 }; - const want = try lanes.sampling.choose(gpa, values, ids, position, s); - try std.testing.expectEqual(want, try draw(gpa, words, s, position, &mask)); - } - } - try std.testing.expectEqual(try lanes.sampling.choose(gpa, values, ids, 0, .{ .seed = 1, .temperature = 0 }), try draw(gpa, words, null, 0, &mask)); - @memset(&mask, 0); - try std.testing.expectError(error.NoAllowedQwenToken, draw(gpa, words, null, 0, &mask)); +test { + _ = draws; } diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig new file mode 100644 index 0000000000..653029762c --- /dev/null +++ b/zig/src/families/qwen3_5/cuda.zig @@ -0,0 +1,17 @@ +//! The Qwen3.5 hybrids on CUDA (2B, Qwen3.8-27B): the engine and its native family entry; Metal's 2B is qwen3_5.zig. +pub const config = @import("config.zig"); +pub const geometry = @import("cuda_geometry.zig"); +pub const engine = @import("cuda_engine.zig"); +pub const native = @import("cuda_native.zig"); +pub const native_moe = @import("cuda_native_moe.zig"); +pub const mtp = @import("cuda_mtp.zig"); +pub const mtp_check = @import("cuda_mtp_check.zig"); +pub const draft = @import("cuda_draft.zig"); + +test { + _ = config; + _ = geometry; + _ = @import("draw.zig"); + _ = @import("cuda_draft_weights.zig"); + _ = draft; +} diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig new file mode 100644 index 0000000000..b4376c9490 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -0,0 +1,100 @@ +//! A stream's state on CUDA: recurrences, conv tails, keys and values, the drafter's tap ring, the last window's log. +const std = @import("std"); +const cuda = @import("cuda"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Engine = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows a drafter's tap ring holds a stream: its window and a prompt chunk more. +pub const tap_rows = 2048 + engine.max_rows; + +/// qwen35.cuh's Seg and Path. +pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; +pub const Path = extern struct { n: c_int, rows: [engine.most_window]c_int }; + +pub const Cache = struct { + state: DeviceBuffer, + kv: DeviceBuffer, + /// A drafter's tapped rows (`Engine.taps`), a ring of tap_rows: the row at position p in slot p % tap_rows. + tap: ?DeviceBuffer = null, + /// The newest verify window's rows' records (qwen35.cuh's LogRow), most_window rows a linear layer. + log: DeviceBuffer, + cap: usize, + len: usize = 0, + /// Positions the newest forward wrote up to: a tap ring holds them until the next forward, whatever `keep` lands. + written: usize = 0, + logits: []u16, + /// The newest verify window of more than one row, until its path is kept: where it started, its rows. + last: ?struct { start: usize, rows: usize } = null, + + pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { + var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); + errdefer state.free(); + try state.fill8(0, null); + var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); + errdefer kv.free(); + try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) + var log = try DeviceBuffer.alloc(e.d, e.g.logBytes(engine.most_window)); + errdefer log.free(); + try log.fill8(0, null); + var tap: ?DeviceBuffer = null; + if (e.n_taps > 0) { + tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); + try tap.?.fill8(0, null); + } + errdefer if (tap) |*t| t.free(); + return .{ .state = state, .kv = kv, .tap = tap, .log = log, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + } + + /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). + pub fn reset(self: *Cache, e: *const Engine) !void { + self.len = 0; + self.written = 0; + self.last = null; + try self.state.fill8(0, e.stream.handle); + } + + pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { + self.state.free(); + self.kv.free(); + self.log.free(); + if (self.tap) |*t| t.free(); + gpa.free(self.logits); + } + + /// The stream's window of `rows` rows from forward row `row0`, after `len` cached positions. + pub fn seg(self: *const Cache, e: *const Engine, len: usize, row0: usize, rows: usize) Seg { + const g = e.g; + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; + } + + /// Keeps `path` (window rows, ascending from 0): the state as that chain would leave it, a tree's keys moved down. + pub fn keep(self: *Cache, e: *Engine, path: []const u32) !void { + const last = self.last orelse return error.NothingToKeep; + if (path.len == 0 or path.len > last.rows or path[0] != 0) return error.NothingToKeep; + var p: Path = .{ .n = @intCast(path.len), .rows = @splat(0) }; + var prefix = true; + for (path, 0..) |r, i| { + if (r >= last.rows or (i > 0 and r <= path[i - 1])) return error.NothingToKeep; + p.rows[i] = @intCast(r); + prefix = prefix and r == i; + } + self.last = null; + self.len = last.start + path.len; + const g = e.g; + const sg = self.seg(e, last.start, 0, last.rows); + try e.run(e.keeps.apply, .{ 1, geometry.dk / 4, g.v_heads * g.deltas() }, .{ 32, 4, 1 }, .{ e.geo, sg, p }); + try e.run(e.keeps.conv, .{ g.deltas(), (g.convBytes() / 2 + 255) / 256, 1 }, .{ 256, 1, 1 }, .{ e.geo, sg, @as(c_int, @intCast(path[path.len - 1])) }); + if (!prefix and g.attentions() > 0) try e.run(e.keeps.compact, .{ g.attentions(), g.kv_heads, 1 }, .{ geometry.head_dim, 1, 1 }, .{ e.geo, sg, p }); + } + + /// The newest window kept whole (the round loop keeps nothing when every row landed), before the next forward. + pub fn settle(self: *Cache, e: *Engine) !void { + const last = self.last orelse return; + var all: [engine.most_window]u32 = undefined; + for (all[0..last.rows], 0..) |*r, i| r.* = @intCast(i); + try self.keep(e, all[0..last.rows]); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig new file mode 100644 index 0000000000..7f0e4e61be --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -0,0 +1,584 @@ +//! The DFlash2 drafter on CUDA behind core/lanes/drafter.zig: a tap ring a stream, 7 drafts as a chain or a tree. +const std = @import("std"); +const platform = @import("sys"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const weights = @import("cuda_draft_weights.zig"); + +const E = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; +const dr = lanes.drafter; +const Stream = lanes.Stream; + +/// Rows an absorb runs at once (the target's prompt chunks are at most this). +const max_rows = engine.max_rows; +/// The logprob row's words (sample.cu's LP_WORDS) and its best-token capacity. +const lp_words = 3 + 2 * 20; +const lp_top = 20; + +pub const Config = struct { + hidden: usize, + intermediate: usize, + layers: usize, + vocab: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + block: usize, + mask: u32, + window: usize, + group: usize, + rank: usize, + topk: usize, + eps: f32, + theta: f32, + taps: [8]u32, + n_taps: usize, + + fn qWidth(c: Config) usize { + return c.heads * c.head_dim; + } + fn kvWidth(c: Config) usize { + return c.kv_heads * c.head_dim; + } + fn dynamic(c: Config) usize { + return 4 * (c.hidden / c.group); + } + /// The ring's rows: the window and a block more. + fn ringRows(c: Config) usize { + return c.window + 16; + } + fn ringBytes(c: Config) usize { + return c.layers * 2 * c.ringRows() * c.kvWidth() * 2; + } + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, target: engine.Geometry) !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 << 20)); + defer gpa.free(bytes); + return parse(gpa, bytes, target); + } + + fn parse(gpa: std.mem.Allocator, bytes: []const u8, target: engine.Geometry) !Config { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadDraftConfig; + const o = doc.value.object; + const d = (o.get("dflash_config") orelse return error.BadDraftConfig); + if (d != .object) return error.BadDraftConfig; + var c: Config = .{ + .hidden = try int(o, "hidden_size"), + .intermediate = try int(o, "intermediate_size"), + .layers = try int(o, "num_hidden_layers"), + .vocab = try int(o, "vocab_size"), + .heads = try int(o, "num_attention_heads"), + .kv_heads = try int(o, "num_key_value_heads"), + .head_dim = try int(o, "head_dim"), + .block = try int(d.object, "block_size"), + .mask = @intCast(try whole(d.object, "mask_token_id")), + .window = try int(o, "sliding_window"), + .group = try int(d.object, "conv_group_size"), + .rank = try int(d.object, "selector_rank"), + .topk = try int(d.object, "selector_top_k"), + .eps = @floatCast(try num(o, "rms_norm_eps")), + .theta = @floatCast(try num((o.get("rope_parameters") orelse return error.BadDraftConfig).object, "rope_theta")), + .taps = undefined, + .n_taps = 0, + }; + const taps = d.object.get("target_layer_ids") orelse return error.BadDraftConfig; + if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > c.taps.len) return error.BadDraftConfig; + for (taps.array.items, 0..) |t, i| { + if (t != .integer or t.integer < 0 or t.integer >= target.layers) return error.BadDraftConfig; + c.taps[i] = @intCast(t.integer); + } + c.n_taps = taps.array.items.len; + // what the kernels fix: the head width, two-tap convolutions, a block of 8, the selector's top-k within a row + if (c.head_dim != 128 or try int(d.object, "conv_kernel_size") != 2 or c.block != 8 or c.topk > lp_top or c.mask >= c.vocab) return error.UnsupportedDraftConfig; + if (c.hidden != target.hidden or c.vocab != target.vocab or c.hidden % 256 != 0 or c.hidden % c.group != 0 or c.heads % c.kv_heads != 0 or c.layers > 8) return error.UnsupportedDraftConfig; + if (c.window < c.block or c.window > engine.tap_rows - max_rows) return error.UnsupportedDraftConfig; + return c; + } +}; + +/// A size, count or divisor: positive, so a zero can't reach a remainder or an empty allocation. +fn int(o: std.json.ObjectMap, key: []const u8) !usize { + const v = try whole(o, key); + if (v == 0) return error.BadDraftConfig; + return v; +} + +/// A token id: any whole number. +fn whole(o: std.json.ObjectMap, key: []const u8) !usize { + const v = o.get(key) orelse return error.BadDraftConfig; + if (v != .integer or v.integer < 0) return error.BadDraftConfig; + return @intCast(v.integer); +} + +fn num(o: std.json.ObjectMap, key: []const u8) !f64 { + return switch (o.get(key) orelse return error.BadDraftConfig) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadDraftConfig, + }; +} + +const Fns = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, dyn_conv: cuda.Function, norm_rope: cuda.Function, ring_put: cuda.Function, attn: cuda.Function }; + +/// qwen35_draft.cu's TapRows. +const TapRows = extern struct { n: c_int, rows: [engine.most_window]c_int }; + +/// A stream's context: its layers' key and value rings over [begin, end), and the drafts it holds (a chain or tree). +const Ctx = struct { + ring: DeviceBuffer, + begin: u64 = 0, + end: u64 = 0, + held: [engine.most_window]u32 = undefined, + held_n: usize = 0, +}; + +pub const Draft = struct { + gpa: std.mem.Allocator, + io: std.Io, + e: *E, + c: Config, + w: weights.Weights, + mod: cuda.Module, + fns: Fns, + scratch: DeviceBuffer, + at: usize = 0, + // absorb rows (max_rows) and block rows (block), all bf16 unless said + taps_in: u64, + th: u64, + thn: u64, + kabs: u64, + kn: u64, + vabs: u64, + xs: u64, // 64-group sums for the projections (f32) + h: u64, + h2: u64, + x: u64, + xin: u64, + dyn: u64, + qraw: u64, + q: u64, + kraw: u64, + kb: u64, + vb: u64, + att: u64, + o: u64, + gate: u64, + up: u64, + act: u64, + logits: u64, + words: u64, // u32 logprob words a row + picks: u64, // u32 zeros: the logprob rows' pick (unused) + hp: u64, + ids: u64, // u32 the block's ids + staged: cuda.HostBuffer, // pinned: the ids up, the words and projected rows down + streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, + /// Contexts between streams, allocated and run once (`reserve`), as the driver stalls a buffer's first kernel. + pool: std.ArrayList(*Ctx) = .empty, + /// One level of a hold (a hold drafts every level at once; timed at load, over the block's levels). + step_ms: f64 = 1.0, + /// The target's window ms by rows (cuda_native.zig times them), for the round loop's prices. + costs: [engine.most_window]lanes.config.Cost = undefined, + n_costs: usize = 0, + + const words_at = 64; + const hp_at = words_at + 16 * lp_words * 4; + /// The host staging area: the words then 16 rows of `rank` selector projections (the block's mask rows fit). + fn stagedLen(rank: usize) usize { + return hp_at + 16 * rank * 2; + } + + pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8) !*Draft { + const c = try Config.read(gpa, io, dir, e.g); + const d = try gpa.create(Draft); + errdefer gpa.destroy(d); + d.* = .{ .gpa = gpa, .io = io, .e = e, .c = c, .w = undefined, .mod = undefined, .fns = undefined, .scratch = undefined, .taps_in = 0, .th = 0, .thn = 0, .kabs = 0, .kn = 0, .vabs = 0, .xs = 0, .h = 0, .h2 = 0, .x = 0, .xin = 0, .dyn = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .o = 0, .gate = 0, .up = 0, .act = 0, .logits = 0, .words = 0, .picks = 0, .hp = 0, .ids = 0, .staged = undefined }; + d.mod = try cuda.Module.load(e.d, cuda.kernels.qwen35_draft); + errdefer d.mod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(d.fns, name) = try d.mod.function("tf_q35d_" ++ name); + d.w = try weights.load(gpa, io, e, dir, .{ .hidden = c.hidden, .intermediate = c.intermediate, .layers = c.layers, .vocab = c.vocab, .q_width = c.qWidth(), .kv_width = c.kvWidth(), .head_dim = c.head_dim, .dynamic = c.dynamic(), .rank = c.rank, .taps = c.n_taps }); + errdefer d.w.deinit(gpa); + d.scratch = try DeviceBuffer.alloc(e.d, d.scratchBytes()); + errdefer d.scratch.free(); + try d.scratch.fill8(0, null); + d.at = 0; + _ = try d.layout(true); + d.staged = try cuda.HostBuffer.alloc(e.d, stagedLen(c.rank)); + errdefer d.staged.free(); + try d.picksZero(); + // the first context, and a hold timed on it (the quickest of several: the first ones pay the driver's stalls) + const ctx = try d.create(); + errdefer d.drop(ctx); + try d.absorb(ctx, d.taps_in, 0, 1, null); + var quickest: f64 = std.math.inf(f64); + for (0..8) |_| { + ctx.end = 1; + const t0 = std.Io.Clock.awake.now(io); + try d.hold(ctx, 0, 1, c.block - 1, null); + quickest = @min(quickest, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6); + } + d.step_ms = quickest / @as(f64, @floatFromInt(c.block - 1)); + ctx.* = .{ .ring = ctx.ring }; + try d.pool.append(gpa, ctx); + return d; + } + + /// The drafter's taps (the target's layers it reads). + pub fn taps(d: *const Draft) []const u32 { + return d.c.taps[0..d.c.n_taps]; + } + + /// The target's window of `rows` rows takes `ms` (told by cuda_native.zig before the facts are read). + pub fn setCost(d: *Draft, rows: u32, ms: f64) void { + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a hold level {d:.2})\n", .{ rows, ms, d.step_ms }); + if (d.n_costs < d.costs.len) { + d.costs[d.n_costs] = .{ .width = rows, .ms = ms }; + d.n_costs += 1; + } + } + + pub fn deinit(d: *Draft) void { + var it = d.streams.valueIterator(); + while (it.next()) |ctx| d.drop(ctx.*); + d.streams.deinit(d.gpa); + for (d.pool.items) |ctx| d.drop(ctx); + d.pool.deinit(d.gpa); + d.staged.free(); + d.scratch.free(); + d.w.deinit(d.gpa); + d.mod.unload(); + d.gpa.destroy(d); + } + + fn scratchBytes(d: *Draft) usize { + d.at = 0; + return d.layout(false) catch unreachable; + } + + /// The scratch rows, measured (nothing assigned) or assigned from the scratch buffer. + fn layout(d: *Draft, assign: bool) !usize { + const c = d.c; + const B = c.block; + const R = max_rows; + const sizes = .{ + .{ "taps_in", R * c.n_taps * c.hidden * 2 }, .{ "th", R * c.hidden * 2 }, .{ "thn", R * c.hidden * 2 }, .{ "kabs", R * c.kvWidth() * 2 }, + .{ "kn", R * c.kvWidth() * 2 }, .{ "vabs", R * c.kvWidth() * 2 }, .{ "xs", R * (c.n_taps * c.hidden / 64) * 4 }, .{ "h", B * c.hidden * 2 }, + .{ "h2", B * c.hidden * 2 }, .{ "x", B * c.hidden * 2 }, .{ "xin", B * c.hidden * 2 }, .{ "dyn", B * c.dynamic() * 2 }, + .{ "qraw", B * c.qWidth() * 2 }, .{ "q", B * c.qWidth() * 2 }, .{ "kraw", B * c.kvWidth() * 2 }, .{ "kb", B * c.kvWidth() * 2 }, + .{ "vb", B * c.kvWidth() * 2 }, .{ "att", B * c.qWidth() * 2 }, .{ "o", B * c.hidden * 2 }, .{ "gate", B * c.intermediate * 2 }, + .{ "up", B * c.intermediate * 2 }, .{ "act", B * c.intermediate * 2 }, .{ "logits", B * c.vocab * 2 }, .{ "words", B * lp_words * 4 }, + .{ "picks", B * 4 }, .{ "hp", B * c.rank * 2 }, .{ "ids", B * 4 }, + }; + var total: usize = 0; + inline for (sizes) |s| { + const at = engine.aligned(total); + total = at + s[1]; + if (assign) @field(d, s[0]) = d.scratch.ptr + at; + } + return total; + } + + fn picksZero(d: *Draft) !void { + const zero: [64]u8 = @splat(0); + try d.e.d.check(d.e.d.api.cuMemcpyHtoD_v2(d.picks, &zero, 64), "cuMemcpyHtoD"); + } + + fn drop(d: *Draft, ctx: *Ctx) void { + ctx.ring.free(); + d.gpa.destroy(ctx); + } + + /// A new context, its ring touched once (an absorb, then a draft), so no request pays the driver's stall. + fn create(d: *Draft) !*Ctx { + const ctx = try d.gpa.create(Ctx); + errdefer d.gpa.destroy(ctx); + ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.c.ringBytes()) }; + errdefer ctx.ring.free(); + try ctx.ring.fill8(0, null); + try d.absorb(ctx, d.taps_in, 0, 1, null); // any tap rows do: the drafter's own zeroed staging + try d.hold(ctx, 0, 1, d.c.block - 1, null); + ctx.* = .{ .ring = ctx.ring }; + return ctx; + } + + /// `n` contexts ready in the pool, for the streams the server admits. + pub fn reserve(d: *Draft, n: usize) !void { + try d.pool.ensureTotalCapacity(d.gpa, n); + while (d.pool.items.len < n) d.pool.appendAssumeCapacity(try d.create()); + } + + fn ringK(d: *Draft, ctx: *Ctx, layer: usize) u64 { + return ctx.ring.ptr + layer * 2 * d.c.ringRows() * d.c.kvWidth() * 2; + } + + fn ringV(d: *Draft, ctx: *Ctx, layer: usize) u64 { + return d.ringK(ctx, layer) + d.c.ringRows() * d.c.kvWidth() * 2; + } + + /// out (rows, n) = x (rows, k) through `l`, as the target's decode tiles or prompt GEMM, with the drafter's sums. + fn proj(d: *Draft, l: cuda.qlinear.Affine4, x: u64, out: u64, rows: usize) !void { + const e = d.e; + if (rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, d.xs, @as(c_int, @intCast(l.k)), @as(u64, 0) }); + try e.affine.decode(e.stream, x, d.xs, l, out, rows); + } + + fn norm(d: *Draft, h: u64, w: u64, out: u64, rows: usize) !void { + const e = d.e; + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ h, @as(u64, 0), w, out, d.c.eps, @as(c_int, @intCast(d.c.hidden)), @as(u64, 0) }); + } + + fn dynConv(d: *Draft, x: u64, dyn: u64, base: u64, res: u64, out: u64, rows: usize, branch: usize) !void { + try d.e.run(d.fns.dyn_conv, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ x, dyn, base, res, out, @as(c_int, @intCast(d.c.hidden)), @as(c_int, @intCast(d.c.group)), @as(c_int, @intCast(branch)) }); + } + + fn normRope(d: *Draft, x: u64, w: u64, out: u64, rows: usize, heads: usize, pos0: u64) !void { + try d.e.run(d.fns.norm_rope, .{ rows * heads, 1, 1 }, .{ 32, 1, 1 }, .{ x, w, out, @as(c_int, @intCast(heads)), @as(c_int, @intCast(heads * d.c.head_dim)), @as(c_longlong, @intCast(pos0)), d.c.eps, d.c.theta }); + } + + /// Target rows [start, start + n) from the tap ring into the context; with `rows`, ring rows start + rows[i]. + fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32) !void { + const c = d.c; + var start = start_; + var n = n_; + if (rows_) |rs| if (rs.len != n or n > engine.most_window) return error.BadDraftRows; + if (start < ctx.end) ctx.end = start; + if (start > ctx.end) ctx.begin = start; // a gap (a prompt restored from a cache): the context starts here + if (n > c.window) { + start += n - c.window; + n = c.window; + ctx.begin = start; + } + ctx.end = start; + const e = d.e; + while (n > 0) { + const rows = @min(n, max_rows); + if (rows_) |rs| { + var tr: TapRows = .{ .n = @intCast(rows), .rows = @splat(0) }; + for (rs, 0..) |r, i| tr.rows[i] = @intCast(r); + try e.run(d.fns.gather_rows, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), tr, @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); + } else try e.run(d.fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); + try d.proj(d.w.fc, d.taps_in, d.th, rows); + try d.norm(d.th, d.w.hidden_norm, d.thn, rows); + for (d.w.layers, 0..) |l, i| { + try d.proj(l.k, d.thn, d.kabs, rows); + try d.normRope(d.kabs, l.k_norm, d.kn, rows, c.kv_heads, start); + try d.proj(l.v, d.thn, d.vabs, rows); + try e.run(d.fns.ring_put, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kn, d.vabs, d.ringK(ctx, i), d.ringV(ctx, i), @as(c_int, @intCast(c.kvWidth())), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.ringRows())) }); + } + start += rows; + n -= rows; + ctx.end = start; + } + if (ctx.end > c.window and ctx.begin < ctx.end - c.window) ctx.begin = ctx.end - c.window; + } + + /// The block [pending, mask x 7] at pos0.. over the context: `depth` drafts as a chain, or the `lanes` tree. + fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { + const c = d.c; + const e = d.e; + const B = c.block; + if (ctx.end != pos0) return error.DraftContextBehind; + const ids = std.mem.bytesAsSlice(u32, d.staged.bytes[0 .. B * 4]); + ids[0] = pending; + @memset(ids[1..B], c.mask); + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(d.ids, ids.ptr, B * 4, e.stream.handle), "cuMemcpyHtoDAsync"); + try e.run(e.fns.embed, .{ B, 1, 1 }, .{ 256, 1, 1 }, .{ d.ids, e.embed_w, e.embed_s, e.embed_b, d.h, @as(c_int, @intCast(c.hidden)), @as(u64, 0) }); + for (d.w.layers, 0..) |l, i| { + try d.norm(d.h, l.input_norm, d.x, B); + try d.proj(l.attn_conv, d.x, d.dyn, B); + try d.dynConv(d.x, d.dyn, l.attn_base, 0, d.xin, B, 0); + try d.proj(l.q, d.xin, d.qraw, B); + try d.proj(l.k, d.xin, d.kraw, B); + try d.proj(l.v, d.xin, d.vb, B); + try d.normRope(d.qraw, l.q_norm, d.q, B, c.heads, pos0); + try d.normRope(d.kraw, l.k_norm, d.kb, B, c.kv_heads, pos0); + try e.run(d.fns.attn, .{ B, c.heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.kb, d.vb, d.ringK(ctx, i), d.ringV(ctx, i), @as(c_int, @intCast(B)), @as(c_int, @intCast(c.heads)), @as(c_int, @intCast(c.kv_heads)), @as(c_int, @intCast(c.ringRows())), @as(c_longlong, @intCast(ctx.begin)), @as(c_longlong, @intCast(pos0)), @as(c_int, @intCast(c.window)), d.att }); + try d.proj(l.o, d.att, d.o, B); + try d.dynConv(d.o, d.dyn, l.attn_base, d.h, d.h2, B, 1); + try d.norm(d.h2, l.post_norm, d.x, B); + try d.proj(l.mlp_conv, d.x, d.dyn, B); + try d.dynConv(d.x, d.dyn, l.mlp_base, 0, d.xin, B, 0); + try d.proj(l.gate, d.xin, d.gate, B); + try d.proj(l.up, d.xin, d.up, B); + try e.run(e.fns.mlp_act, .{ (B * c.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ d.gate, d.up, d.act, @as(c_longlong, @intCast(B * c.intermediate)), @as(u64, 0), @as(c_int, @intCast(c.intermediate)) }); + try d.proj(l.down, d.act, d.o, B); + try d.dynConv(d.o, d.dyn, l.mlp_base, d.h2, d.h, B, 1); + } + // the mask rows' logits through the target's head, their best tokens and their selector projections + const rows = B - 1; + try d.norm(d.h + c.hidden * 2, d.w.final_norm, d.x, rows); + try d.proj(e.head, d.x, d.logits, rows); + try e.draws.logprobRows(e.stream, d.logits, c.vocab, d.picks, @intCast(c.topk), false, d.words, rows); + try d.proj(d.w.selector, d.x, d.hp, rows); + const st = d.staged.bytes; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + words_at, d.words, rows * lp_words * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + hp_at, d.hp, rows * c.rank * 2, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + if (lanes_) |l| { + try d.selectLanes(ctx, pending, rows, l); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ l.parents.len, l.levels() }); + } else { + d.select(ctx, pending, rows); + ctx.held_n = @min(depth, rows); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] chain {d}\n", .{ctx.held_n}); + } + } + + /// Row `r`'s `rank`th candidate after `pred`: unary logit + pred . projected . successor, ties to the earlier one. + fn choice(d: *const Draft, r: usize, pred: u32, rank: usize) u32 { + const c = d.c; + const w = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][r * lp_words * 4 ..][0 .. lp_words * 4]); + const hp = std.mem.bytesAsSlice(u16, d.staged.bytes[hp_at..][r * c.rank * 2 ..][0 .. c.rank * 2]); + var scores: [lp_top]f32 = undefined; + for (0..c.topk) |j| { + const cand = w[2 + j]; + const unary: f32 = @bitCast(w[2 + lp_top + j]); + var edge: f32 = 0; + for (0..c.rank) |k| edge += bf16(d.w.predecessor[pred * c.rank + k]) * bf16(hp[k]) * bf16(d.w.successor[cand * c.rank + k]); + scores[j] = unary + edge; + } + var taken: [lp_top]bool = @splat(false); + var pick: usize = 0; + for (0..rank + 1) |_| { + var best: f32 = -std.math.inf(f32); + pick = 0; + for (0..c.topk) |j| if (!taken[j] and scores[j] > best) { + best = scores[j]; + pick = j; + }; + taken[pick] = true; + } + return w[2 + pick]; + } + + /// The path through the rows' candidates: each row's best after the one before. + fn select(d: *Draft, ctx: *Ctx, pending: u32, rows: usize) void { + var pred = pending; + for (0..rows) |r| { + ctx.held[r] = d.choice(r, pred, 0); + pred = ctx.held[r]; + } + } + + /// The tree's lanes in order: lane i the rank-r candidate at its depth's row after its parent lane's token. + fn selectLanes(d: *Draft, ctx: *Ctx, pending: u32, rows: usize, l: *const lanes.shape.Shape) !void { + if (l.parents.len > ctx.held.len) return error.TooManyLanes; + for (l.parents, l.depths, l.ranks, 0..) |par, depth, rank, i| { + if (depth >= rows or rank >= d.c.topk or par >= @as(i32, @intCast(i))) return error.BadLanes; + const pred = if (par < 0) pending else ctx.held[@intCast(par)]; + ctx.held[i] = d.choice(depth, pred, rank); + } + ctx.held_n = l.parents.len; + } + + // the Drafter contract + + pub fn drafter(d: *Draft) dr.Drafter { + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(p: *anyopaque) *Draft { + return @ptrCast(@alignCast(p)); + } + + fn tapsFn(p: *anyopaque) []const u32 { + const d = cast(p); + return d.c.taps[0..d.c.n_taps]; + } + + fn factsFn(p: *anyopaque) dr.Facts { + const d = cast(p); + // TF_NO_TREES=1: chains only (for comparing) + return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes(), .trees = platform.c.getenv("TF_NO_TREES") == null, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; + } + + fn openFn(p: *anyopaque, s: *Stream) anyerror!void { + const d = cast(p); + if (d.streams.contains(s)) return error.StreamOpen; + const ctx = d.pool.pop() orelse try d.create(); + errdefer d.pool.append(d.gpa, ctx) catch d.drop(ctx); + ctx.* = .{ .ring = ctx.ring }; + try d.streams.put(d.gpa, s, ctx); + } + + fn absorbFn(p: *anyopaque, items: []const dr.Absorb) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("absorb", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; + // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) + if (it.rows) |rs| { + try d.absorb(ctx, it.features.buffer, it.features.offset, rs.len, rs); + } else try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows, null); + } + } + + fn holdFn(p: *anyopaque, items: []const dr.Hold) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("hold", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + try d.hold(ctx, it.pending, it.position - 1, it.depth, it.lanes); + } + } + + fn heldFn(p: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const d = cast(p); + for (streams, out) |s, o| { + const ctx = d.streams.get(s) orelse return error.UnknownStream; + if (o.len > ctx.held_n) return error.TooFewDrafts; + @memcpy(o, ctx.held[0..o.len]); + } + } + + fn releaseFn(p: *anyopaque, s: *Stream) void { + const d = cast(p); + if (d.streams.fetchRemove(s)) |kv| d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + } +}; + +/// Acceptance by depth until a stream has its own (the reference's chat text: measured, see the commit). +const prior = [_]f64{ 0.75, 0.65, 0.58, 0.52, 0.47, 0.43, 0.4 }; + +fn bf16(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// TF_TIMING=1: each call's wall time on stderr. +pub fn timing(what: []const u8, io: std.Io, t0: std.Io.Timestamp) void { + if (platform.c.getenv("TF_TIMING") == null) return; + const now = std.Io.Clock.awake.now(io); + const ns: f64 = @floatFromInt(t0.durationTo(now).toNanoseconds()); + const at: f64 = @floatFromInt(@mod(now.toNanoseconds(), 1_000_000_000_000)); + std.debug.print("[timing] {s} {d:.2} ms at {d:.1}\n", .{ what, ns / 1e6, at / 1e6 }); +} + +test "a drafter config with a zero divisor is refused, and any admitted rank's selector rows fit the staging area" { + const gpa = std.testing.allocator; + const target: engine.Geometry = .{ .hidden = 5120, .intermediate = 17408, .layers = 64, .vocab = 248320, .k_heads = 16, .v_heads = 48, .q_heads = 24, .kv_heads = 4, .eps = 1e-6, .tied = false, .context = 262144, .linear = @splat(false) }; + const text = + \\{"dflash_config":{"block_size":8,"conv_group_size":16,"conv_kernel_size":2,"mask_token_id":248070,"selector_rank":256, + \\ "selector_top_k":16,"target_layer_ids":[5,19,33,47,61]},"head_dim":128,"hidden_size":5120,"intermediate_size":17408, + \\ "num_attention_heads":32,"num_hidden_layers":5,"num_key_value_heads":8,"rms_norm_eps":1e-06, + \\ "rope_parameters":{"rope_theta":10000000},"sliding_window":2048,"vocab_size":248320} + ; + try std.testing.expectEqual(@as(usize, 256), (try Config.parse(gpa, text, target)).rank); + for ([_][2][]const u8{ .{ "\"conv_group_size\":16", "\"conv_group_size\":0" }, .{ "\"num_key_value_heads\":8", "\"num_key_value_heads\":0" } }) |r| { + const bad = try std.mem.replaceOwned(u8, gpa, text, r[0], r[1]); + defer gpa.free(bad); + try std.testing.expectError(error.BadDraftConfig, Config.parse(gpa, bad, target)); + } + const wide = try std.mem.replaceOwned(u8, gpa, text, "\"selector_rank\":256", "\"selector_rank\":1024"); + defer gpa.free(wide); + const c = try Config.parse(gpa, wide, target); + try std.testing.expect(Draft.stagedLen(c.rank) >= Draft.hp_at + (c.block - 1) * c.rank * 2); +} diff --git a/zig/src/families/qwen3_5/cuda_draft_weights.zig b/zig/src/families/qwen3_5/cuda_draft_weights.zig new file mode 100644 index 0000000000..6d614f7982 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_draft_weights.zig @@ -0,0 +1,221 @@ +//! The DFlash2 drafter's weights onto the device: projections quantized to affine-4 on the host, codebooks kept there. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const engine = @import("cuda_engine.zig"); + +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const Engine = engine.Engine; +const aligned = engine.aligned; + +pub const Layer = struct { + input_norm: u64, + post_norm: u64, + q: Affine4, + k: Affine4, + v: Affine4, + o: Affine4, + q_norm: u64, + k_norm: u64, + gate: Affine4, + up: Affine4, + down: Affine4, + attn_conv: Affine4, // kernel_projection: hidden -> 4 * groups dynamic coefficients + attn_base: u64, // base_kernel (2 branches, 2 lags, hidden) bf16 + mlp_conv: Affine4, + mlp_base: u64, +}; + +pub const Weights = struct { + fc: Affine4, // the taps' concatenation -> hidden + hidden_norm: u64, + final_norm: u64, + selector: Affine4, // hidden_projection: hidden -> rank + predecessor: []u16, // codebooks (vocab, rank) bf16, on the host + successor: []u16, + layers: []Layer, + arena: DeviceBuffer, + arena_at: usize = 0, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + w.arena.free(); + gpa.free(w.predecessor); + gpa.free(w.successor); + gpa.free(w.layers); + } + + fn alloc(w: *Weights, bytes: usize) !u64 { + const at = aligned(w.arena_at); + if (at + bytes > w.arena.len) return error.ArenaExhausted; + w.arena_at = at + bytes; + return w.arena.ptr + at; + } +}; + +/// The drafter's widths the loader needs (cuda_draft.zig's Config). +pub const Shape = struct { hidden: usize, intermediate: usize, layers: usize, vocab: usize, q_width: usize, kv_width: usize, head_dim: usize, dynamic: usize, rank: usize, taps: usize }; + +fn packedBytes(n: usize, k: usize) usize { + const l = Affine4.layout(n, k); + return aligned(l.words) + 2 * aligned(l.scales); +} + +/// Device bytes of everything `load` packs. +pub fn total(s: Shape) usize { + var t = packedBytes(s.hidden, s.taps * s.hidden) + packedBytes(s.rank, s.hidden) + 2 * aligned(s.hidden * 2); + t += s.layers * (2 * aligned(s.hidden * 2) + packedBytes(s.q_width, s.hidden) + 2 * packedBytes(s.kv_width, s.hidden) + packedBytes(s.hidden, s.q_width) + + 2 * aligned(s.head_dim * 2) + 2 * packedBytes(s.intermediate, s.hidden) + packedBytes(s.hidden, s.intermediate) + 2 * packedBytes(s.dynamic, s.hidden) + 2 * aligned(4 * s.hidden * 2)); + return t; +} + +fn bf16(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +fn toBf16(x: f32) u16 { + // round to nearest even, as the device does + const u: u32 = @bitCast(x); + const rounding: u32 = 0x7fff + ((u >> 16) & 1); + return @intCast((u +% rounding) >> 16); +} + +/// MLX affine-4, group 64: scale bf16((max - min) / 15), bias bf16(min), nibbles round((w - bias) / scale), low first. +fn quantize(gpa: std.mem.Allocator, w: []align(1) const u16, n: usize, k: usize) !struct { words: []u32, scales: []u16, biases: []u16 } { + const groups = k / 64; + const words = try gpa.alloc(u32, n * k / 8); + errdefer gpa.free(words); + const scales = try gpa.alloc(u16, n * groups); + errdefer gpa.free(scales); + const biases = try gpa.alloc(u16, n * groups); + errdefer gpa.free(biases); + for (0..n) |r| { + for (0..groups) |g| { + const row = w[r * k + g * 64 ..][0..64]; + var lo: f32 = std.math.inf(f32); + var hi: f32 = -std.math.inf(f32); + for (row) |h| { + const x = bf16(h); + lo = @min(lo, x); + hi = @max(hi, x); + } + const scale = toBf16((hi - lo) / 15.0); + const bias = toBf16(lo); + scales[r * groups + g] = scale; + biases[r * groups + g] = bias; + const s = bf16(scale); + const b = bf16(bias); + const inv = if (s == 0) 0.0 else 1.0 / s; + for (0..8) |j| { + var word: u32 = 0; + for (0..8) |i| { + const q = std.math.clamp(@round((bf16(row[j * 8 + i]) - b) * inv), 0.0, 15.0); + word |= @as(u32, @intFromFloat(q)) << @intCast(4 * i); + } + words[(r * k + g * 64) / 8 + j] = word; + } + } + } + return .{ .words = words, .scales = scales, .biases = biases }; +} + +/// A bf16 (n, k) tensor quantized and packed through the target's staging buffer into the drafter's arena. +fn linear(gpa: std.mem.Allocator, e: *Engine, w: *Weights, ck: *core.Checkpoint, name: []const u8, n: usize, k: usize) !Affine4 { + const t = try ck.expect(name, .bf16, &.{ n, k }); + const q = try quantize(gpa, std.mem.bytesAsSlice(u16, t.bytes), n, k); + defer { + gpa.free(q.words); + gpa.free(q.scales); + gpa.free(q.biases); + } + const parts = [_][]const u8{ std.mem.sliceAsBytes(q.words), std.mem.sliceAsBytes(q.scales), std.mem.sliceAsBytes(q.biases) }; + var raw: [3]u64 = undefined; + var at: usize = 0; + for (parts, &raw) |p, *r| { + if (at + p.len > e.stage.len) return error.StageTooSmall; + try e.stage.upload(at, p); + r.* = e.stage.ptr + at; + at += p.len; + } + const l = Affine4.layout(n, k); + const out: Affine4 = .{ .w = try w.alloc(l.words), .s = try w.alloc(l.scales), .b = try w.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + try e.affine.pack(e.stream, raw[0], raw[1], raw[2], out); + try e.stream.synchronize(); + return out; +} + +fn vector(w: *Weights, ck: *core.Checkpoint, name: []const u8, shape: []const usize) !u64 { + const t = try ck.expect(name, .bf16, shape); + const p = try w.alloc(t.bytes.len); + try w.arena.upload(p - w.arena.ptr, t.bytes); + return p; +} + +fn hostBf16(gpa: std.mem.Allocator, ck: *core.Checkpoint, name: []const u8, shape: []const usize) ![]u16 { + const t = try ck.expect(name, .bf16, shape); + const out = try gpa.alloc(u16, t.bytes.len / 2); + @memcpy(out, std.mem.bytesAsSlice(u16, t.bytes)); + return out; +} + +pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *Engine, dir: []const u8, s: Shape) !Weights { + var ck = try core.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + var w: Weights = .{ .fc = undefined, .hidden_norm = 0, .final_norm = 0, .selector = undefined, .predecessor = &.{}, .successor = &.{}, .layers = &.{}, .arena = try DeviceBuffer.alloc(e.d, total(s)) }; + errdefer w.arena.free(); + try w.arena.fill8(0, null); + w.layers = try gpa.alloc(Layer, s.layers); + errdefer gpa.free(w.layers); + w.fc = try linear(gpa, e, &w, &ck, "fc.weight", s.hidden, s.taps * s.hidden); + w.hidden_norm = try vector(&w, &ck, "hidden_norm.weight", &.{s.hidden}); + w.final_norm = try vector(&w, &ck, "norm.weight", &.{s.hidden}); + w.selector = try linear(gpa, e, &w, &ck, "candidate_selector.hidden_projection.weight", s.rank, s.hidden); + w.predecessor = try hostBf16(gpa, &ck, "candidate_selector.predecessor_codebook", &.{ s.vocab, s.rank }); + errdefer gpa.free(w.predecessor); + w.successor = try hostBf16(gpa, &ck, "candidate_selector.successor_codebook", &.{ s.vocab, s.rank }); + errdefer gpa.free(w.successor); + for (w.layers, 0..) |*l, i| { + var buf: [96]u8 = undefined; + const Name = struct { + buf: *[96]u8, + i: usize, + fn of(self: @This(), suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(self.buf, "layers.{d}.{s}", .{ self.i, suffix }); + } + }; + const nm: Name = .{ .buf = &buf, .i = i }; + l.* = .{ + .input_norm = try vector(&w, &ck, try nm.of("input_layernorm.weight"), &.{s.hidden}), + .post_norm = try vector(&w, &ck, try nm.of("post_attention_layernorm.weight"), &.{s.hidden}), + .q = try linear(gpa, e, &w, &ck, try nm.of("self_attn.q_proj.weight"), s.q_width, s.hidden), + .k = try linear(gpa, e, &w, &ck, try nm.of("self_attn.k_proj.weight"), s.kv_width, s.hidden), + .v = try linear(gpa, e, &w, &ck, try nm.of("self_attn.v_proj.weight"), s.kv_width, s.hidden), + .o = try linear(gpa, e, &w, &ck, try nm.of("self_attn.o_proj.weight"), s.hidden, s.q_width), + .q_norm = try vector(&w, &ck, try nm.of("self_attn.q_norm.weight"), &.{s.head_dim}), + .k_norm = try vector(&w, &ck, try nm.of("self_attn.k_norm.weight"), &.{s.head_dim}), + .gate = try linear(gpa, e, &w, &ck, try nm.of("mlp.gate_proj.weight"), s.intermediate, s.hidden), + .up = try linear(gpa, e, &w, &ck, try nm.of("mlp.up_proj.weight"), s.intermediate, s.hidden), + .down = try linear(gpa, e, &w, &ck, try nm.of("mlp.down_proj.weight"), s.hidden, s.intermediate), + .attn_conv = try linear(gpa, e, &w, &ck, try nm.of("attention_conv.kernel_projection.weight"), s.dynamic, s.hidden), + .attn_base = try vector(&w, &ck, try nm.of("attention_conv.base_kernel"), &.{ 2, 2, s.hidden }), + .mlp_conv = try linear(gpa, e, &w, &ck, try nm.of("mlp_conv.kernel_projection.weight"), s.dynamic, s.hidden), + .mlp_base = try vector(&w, &ck, try nm.of("mlp_conv.base_kernel"), &.{ 2, 2, s.hidden }), + }; + } + return w; +} + +test "bf16 rounding of a quantized group's scale and bias" { + try std.testing.expectEqual(@as(u16, 0x3f80), toBf16(1.0)); + try std.testing.expectEqual(@as(f32, 1.0), bf16(0x3f80)); + const w: [64]u16 = @splat(0x3f80); // a flat group: scale 0, bias 1, every nibble 0 + const q = try quantize(std.testing.allocator, &w, 1, 64); + defer { + std.testing.allocator.free(q.words); + std.testing.allocator.free(q.scales); + std.testing.allocator.free(q.biases); + } + try std.testing.expectEqual(@as(u16, 0), q.scales[0]); + try std.testing.expectEqual(@as(u16, 0x3f80), q.biases[0]); + try std.testing.expectEqual(@as(u32, 0), q.words[0]); +} diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig new file mode 100644 index 0000000000..8cca782f8a --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -0,0 +1,583 @@ +//! Qwen3.5's hybrids on CUDA: qlinear projections, qwen35.cu glue, one captured forward over every stream's window. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const geometry = @import("cuda_geometry.zig"); +const weights = @import("cuda_weights.zig"); +const caches = @import("cuda_cache.zig"); + +pub const Geometry = geometry.Geometry; +pub const Cache = caches.Cache; +pub const tap_rows = caches.tap_rows; +const Seg = caches.Seg; +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const head_dim = geometry.head_dim; + +/// Rows of one forward at most: a prompt chunk, or a shared round's windows. +pub const max_rows = 128; +/// Streams' windows in one forward at most. +pub const max_segs = 32; +/// Rows of a verify forward at most: each row's logits are kept and drawn. +pub const most_verify = 32; +/// A window's rows at most (exact: one-row steps' bits; qwen35.cuh's MAXW, the rows a stream's log holds). +pub const most_window = 16; +/// Partial attention chunks a row at most (qwen35.cuh's MAXCH): each row's chunk width grows with its own keys. +const max_chunks = 64; + +/// Device allocations are 256-byte aligned, as the driver's are. +pub fn aligned(n: usize) usize { + return std.mem.alignForward(usize, n, 256); +} + +/// qwen35.cuh's Geo. +pub const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int, nl: c_int }; + +pub const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; +pub const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; +/// A dense SwiGLU MLP, or routed experts (the shared one last in their tables) behind a router of bf16 rows. +pub const Mlp = union(enum) { dense: struct { gate: Affine4, up: Affine4, down: Affine4 }, moe: struct { router: u64, experts: cuda.experts.Layer } }; +pub const Layer = struct { + input_norm: u64, + post_norm: u64, + mlp: Mlp, + mixer: union(enum) { delta: Delta, attention: Attention }, +}; + +const Fns = struct { + embed: cuda.Function, + norm: cuda.Function, + sums: cuda.Function, + gdn_pre: cuda.Function, + conv_commit: cuda.Function, + gdn_chain: cuda.Function, + gdn_post: cuda.Function, + head_norm: cuda.Function, + rope: cuda.Function, + attn_partial: cuda.Function, + attn_merge: cuda.Function, + attn_gate: cuda.Function, + mlp_act: cuda.Function, + argmax: cuda.Function, + tap: cuda.Function, +}; + +/// qwen35_keep.cu's kernels (Cache.keep). +pub const Keeps = struct { apply: cuda.Function, conv: cuda.Function, compact: cuda.Function }; + +/// A window of one stream in a forward: `parents`, for a tree, names each row's parent row (-1: the first row). +pub const Part = struct { cache: *Cache, rows: usize, parents: ?[]const i32 = null }; + +/// What a forward keeps: a prompt chunk its last row's logits, a verify window every row's draw or logits and log. +pub const Keep = union(enum) { last, draws: []const Draw, logits }; + +/// A kept row's draw on the device: sample.cu's keyed rule at `position`, or greedy without one. +pub const Draw = struct { rule: ?cuda.sample.Rule, position: u64 }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + g: Geometry, + geo: Geo, + stream: cuda.Stream, + mods: [11]cuda.Module, + draws: cuda.sample.Draws, + affine: cuda.qlinear.Affine4Kernels, + experts: cuda.experts.Affine4Experts, + plan: cuda.grouped.Router, + moe: cuda.moe.Fns, + fns: Fns, + keeps: Keeps, + /// Every weight and scratch buffer in one allocation: WDDM's launch cost grows with the allocation count. + arena: DeviceBuffer, + arena_at: usize = 0, + stage: DeviceBuffer, // the load's staging of each projection as stored, packed from there + /// Rows a verify forward takes at most, and a window's rows at most. + verify_rows: usize, + window_rows: usize, + // weights + embed_w: u64, + embed_s: u64, + embed_b: u64, + head: Affine4, + final_norm: u64, + layers: []Layer, + // scratch, max_rows rows each + h: u64, + x: u64, + r: u64, + xs: u64, + qkv: u64, + z: u64, + a: u64, + b: u64, + q: u64, + k: u64, + v: u64, + g_: u64, + beta: u64, + y: u64, + mix: u64, + qraw: u64, + kraw: u64, + vraw: u64, + qn: u64, + kn: u64, + gate: u64, + up: u64, + act: u64, + pm: u64, + pl: u64, + po: u64, + logits: u64, + // routed experts' scratch (qwen3_5_moe): logits, picks and weights a row, the plan, the pairs' activations and sums + rlogits: u64 = 0, + picks: u64 = 0, + wts: u64 = 0, + members: u64 = 0, + items: u64 = 0, + counts: u64 = 0, + rank: u64 = 0, + hist: u64 = 0, + eact: u64 = 0, + ey: u64 = 0, + // per-forward tables (ids, segments, rows' segment, parent and ancestors, the draws' rules) staged in pinned memory + tables: DeviceBuffer, + staged: cuda.HostBuffer, + host_logits: cuda.HostBuffer, + /// The kept rows' device draws, when forward was given them. + tokens: [most_verify]u32 = undefined, + /// The captured forwards by shape (`Shape.key`): the launches from the embedding to the head. + graphs: std.AutoHashMapUnmanaged(u32, cuda.graph.Exec) = .empty, + /// A drafter's tapped layers (`setTaps`): each row's residual after them lands in its stream's tap ring. + taps: [8]u32 = undefined, + n_taps: usize = 0, + + const ids_at = 0; + const segs_at = max_rows * 4; + const rowseg_at = segs_at + max_segs * @sizeOf(Seg); + const rowpar_at = rowseg_at + max_rows * 4; + const rowanc_at = rowpar_at + max_rows * 4; + const rules_at = rowanc_at + max_rows * 4; + const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); + const tokens_at = meta_at + most_verify * 4; + const params_at = tokens_at + most_verify * 4; // two unused words, the rows, the windows, record + const tables_len = params_at + 8 * 4; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const g = try Geometry.read(gpa, io, dir); + const d = ctx.d; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = undefined; + e.gpa = gpa; + e.d = d; + e.g = g; + e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads), .nl = @intCast(g.deltas()) }; + e.verify_rows = most_verify; + e.window_rows = most_window; + e.arena_at = 0; + e.graphs = .empty; + e.n_taps = 0; + e.stream = try cuda.Stream.init(d, true); + errdefer e.stream.deinit(); + e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); + errdefer e.arena.free(); + try e.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly + e.stage = try DeviceBuffer.alloc(d, weights.stageBytes(g)); + errdefer e.stage.free(); + const kk = cuda.kernels; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample, kk.qwen35_keep, kk.experts, kk.experts_prefill, kk.experts_pack, kk.moe }; + var loaded: usize = 0; + errdefer for (e.mods[0..loaded]) |*m| m.unload(); + for (images, 0..) |img, i| { + e.mods[i] = try cuda.Module.load(d, img); + loaded += 1; + } + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + const gb10 = try ctx.attribute(.compute_capability_major) == 12 and try ctx.attribute(.compute_capability_minor) == 1; + e.affine = try cuda.qlinear.Affine4Kernels.resolve(d, e.mods[0], e.mods[1], e.mods[2], e.mods[3], sms, gb10); + errdefer e.affine.deinit(); + e.draws = try cuda.sample.Draws.resolve(e.mods[5]); + e.experts = try cuda.experts.Affine4Experts.resolve(e.mods[7], e.mods[8], e.mods[9], sms); + e.plan = try cuda.grouped.Router.resolve(e.mods[7]); + e.moe = try cuda.moe.Fns.resolve(e.mods[10]); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(e.fns, name) = try e.mods[4].function("tf_q35_" ++ name); + inline for (@typeInfo(Keeps).@"struct".field_names) |name| @field(e.keeps, name) = try e.mods[6].function("tf_q35k_" ++ name); + e.layers = try gpa.alloc(Layer, g.layers); + errdefer gpa.free(e.layers); + var ck = try core.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + try weights.load(e, &ck); + _ = try e.scratch(.allocate); + e.tables = try DeviceBuffer.alloc(d, tables_len); + errdefer e.tables.free(); + e.staged = try cuda.HostBuffer.alloc(d, tables_len); + errdefer e.staged.free(); + e.host_logits = try cuda.HostBuffer.alloc(d, most_verify * g.vocab * 2); + errdefer e.host_logits.free(); + try e.warm(gpa); + return e; + } + + /// Every path a request takes, run once now: WDDM stalls a process's first forwards for seconds (a 20 s launch). + fn warm(e: *Engine, gpa: std.mem.Allocator) !void { + var warm_cache = try Cache.init(gpa, e, 64); + defer warm_cache.deinit(gpa); + const ids: [32]u32 = @splat(0); + const greedy: Draw = .{ .rule = null, .position = 33 }; + try e.forward(&.{.{ .cache = &warm_cache, .rows = ids.len }}, &ids, .last); // a prompt chunk: the prompt GEMMs + try e.forward(&.{.{ .cache = &warm_cache, .rows = 1 }}, ids[0..1], .{ .draws = &.{greedy} }); + try e.forward(&.{.{ .cache = &warm_cache, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = .{ .seed = 1, .inv_t = 1.0 / 0.7, .top_p = 0.95, .near = 20.0, .min_log = -std.math.inf(f32), .top_k = 20 }, .position = 34 }} }); + // a tree window and a branch of it kept: the keep kernels, and the forward's branch path + try e.forward(&.{.{ .cache = &warm_cache, .rows = 3, .parents = &.{ -1, 0, 0 } }}, ids[0..3], .{ .draws = &.{ greedy, greedy, greedy } }); + try warm_cache.keep(e, &.{ 0, 2 }); + try e.stream.synchronize(); + } + + pub fn deinit(e: *Engine) void { + var it = e.graphs.valueIterator(); + while (it.next()) |x| x.deinit(); + e.graphs.deinit(e.gpa); + e.host_logits.free(); + e.staged.free(); + e.tables.free(); + e.stage.free(); + e.arena.free(); + e.gpa.free(e.layers); + e.affine.deinit(); + for (&e.mods) |*m| m.unload(); + e.stream.deinit(); + e.gpa.destroy(e); + } + + /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values, taps, log. + pub fn streamBytes(e: *const Engine, cap: usize) usize { + return e.g.stateBytes() + e.g.kvBytes(cap) + e.tapBytes() + e.g.logBytes(most_window); + } + + /// A stream's tap ring: tap_rows rows of the tapped layers' residuals. + pub fn tapBytes(e: *const Engine) usize { + return tap_rows * e.n_taps * e.g.hidden * 2; + } + + /// The layers a drafter taps, told before the streams' caches exist; each forward then fills their rings. + pub fn setTaps(e: *Engine, taps: []const u32) !void { + if (taps.len > e.taps.len) return error.TooManyTaps; + for (taps) |t| if (t >= e.g.layers) return error.BadTap; + @memcpy(e.taps[0..taps.len], taps); + e.n_taps = taps.len; + } + + /// The next `bytes` of the arena, 256-byte aligned. + pub fn alloc(e: *Engine, bytes: usize) !u64 { + const at = aligned(e.arena_at); + if (at + bytes > e.arena.len) return error.ArenaExhausted; + e.arena_at = at + bytes; + return e.arena.ptr + at; + } + + /// The scratch rows' bytes (`measure`: summed, nothing taken), or the rows themselves from the arena. + fn scratch(e: *Engine, mode: enum { measure, allocate }) !usize { + var total: usize = 0; + const Take = struct { + e: *Engine, + mode: @TypeOf(mode), + total: *usize, + fn alloc(t: @This(), bytes: usize) !u64 { + t.total.* += aligned(bytes); + return if (t.mode == .measure) 0 else try t.e.alloc(bytes); + } + }; + const t = Take{ .e = e, .mode = mode, .total = &total }; + const g = e.g; + const R = max_rows; + const wide = @max(g.valueWidth(), g.attnWidth(), g.hidden); // y and mix: either mixer's rows + inline for (.{ "h", "x", "r" }) |name| @field(e, name) = try t.alloc(R * g.hidden * 2); + inline for (.{ "y", "mix", "z", "v" }) |name| @field(e, name) = try t.alloc(R * wide * 2); + inline for (.{ "q", "k" }) |name| @field(e, name) = try t.alloc(R * g.k_heads * geometry.dk * 2); + e.qkv = try t.alloc(R * g.convWidth() * 2); + inline for (.{ "gate", "up", "act" }) |name| @field(e, name) = try t.alloc(R * g.intermediate * 2); + e.xs = try t.alloc(R * (@max(g.intermediate, wide) / 64) * 4); + inline for (.{ "a", "b", "beta" }) |name| @field(e, name) = try t.alloc(R * g.v_heads * 2); + e.g_ = try t.alloc(R * g.v_heads * 4); + e.qraw = try t.alloc(R * 2 * g.attnWidth() * 2); + e.qn = try t.alloc(R * g.attnWidth() * 2); + inline for (.{ "kraw", "vraw", "kn" }) |name| @field(e, name) = try t.alloc(R * g.kv_heads * head_dim * 2); + e.pm = try t.alloc(R * g.q_heads * max_chunks * 4); + e.pl = try t.alloc(R * g.q_heads * max_chunks * 4); + e.po = try t.alloc(R * g.q_heads * max_chunks * head_dim * 4); + e.logits = try t.alloc(most_verify * g.vocab * 2); + if (g.moe()) { + const pairs = R * g.slots(); + const count = g.experts + 1; + e.rlogits = try t.alloc(R * count * 4); + e.picks = try t.alloc(pairs * 4); + e.wts = try t.alloc(pairs * 4); + e.members = try t.alloc(pairs * 4); + e.items = try t.alloc(3 * cuda.grouped.maxItems(pairs, count, 16) * 4); + e.counts = try t.alloc(2 * 4); + e.rank = try t.alloc(pairs * 4); + e.hist = try t.alloc(((pairs + 1023) / 1024) * count * 4); + e.eact = try t.alloc(pairs * g.intermediate * 2); + e.ey = try t.alloc(pairs * g.hidden * 4); + } + return total; + } + + pub fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + const cfg: cuda.Config = .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = @intCast(block[0]), .y = @intCast(block[1]), .z = @intCast(block[2]) } }; + try cuda.launch.launch(f, cfg, e.stream, &a); + } + + /// out (rows, n) = x (rows, k) through `l`: 16-row decode tiles (a row's bits batch-independent) or a prompt GEMM. + pub fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { + if (prompt and rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)), e.tables.ptr + params_at }); + var at: usize = 0; + while (at < rows) : (at += 16) { + const n = @min(16, rows - at); + try e.affine.decode(e.stream, x + at * l.k * 2, e.xs + at * (l.k / 64) * 4, l, out + at * l.n * 2, n); + } + } + + fn norm(e: *Engine, residual: bool, w: u64, rows: usize) !void { + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); + } + + /// A chain row's rowanc: its index in its window, as a prompt chunk's rows outnumber the bits. + const chain_row: u32 = 1 << 31; + + /// The row tables of `parts`: each row's segment, parent row (-1: a window's first) and ancestors-and-self bits. + fn rowTables(st: []u8, parts: []const Part) !void { + const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. max_rows * 4]); + const rowpar = std.mem.bytesAsSlice(c_int, st[rowpar_at..][0 .. max_rows * 4]); + const rowanc = std.mem.bytesAsSlice(u32, st[rowanc_at..][0 .. max_rows * 4]); + var row0: usize = 0; + for (parts, 0..) |p, i| { + if (p.parents) |ps| if (ps.len != p.rows or p.rows > most_window or ps[0] != -1) return error.InvalidQwenParents; + for (0..p.rows) |r| { + const parent: i32 = if (p.parents) |ps| ps[r] else if (r == 0) -1 else @intCast(r - 1); + if (parent >= @as(i32, @intCast(r)) or parent < -1) return error.InvalidQwenParents; + rowseg[row0 + r] = @intCast(i); + rowpar[row0 + r] = if (parent < 0) -1 else @intCast(row0 + @as(usize, @intCast(parent))); + rowanc[row0 + r] = if (p.parents == null) chain_row | @as(u32, @intCast(r)) else (@as(u32, 1) << @intCast(r)) | (if (parent < 0) 0 else rowanc[row0 + @as(usize, @intCast(parent))]); + } + row0 += p.rows; + } + } + + /// One forward over `parts`' windows, `ids` in order; `keep` names the kept rows; an unkept verify is kept first. + pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { + const g = e.g; + var total: usize = 0; + for (parts) |p| { + try p.cache.settle(e); + if (p.rows == 0 or p.cache.len + p.rows > p.cache.cap) return error.QwenContextFull; + total += p.rows; + } + if (parts.len == 0 or parts.len > max_segs or total != ids.len or total > max_rows) return error.InvalidQwenSegments; + const record = keep != .last; + if (record and total > e.verify_rows) return error.InvalidQwenSegments; + if (!record and parts.len != 1) return error.InvalidQwenSegments; + if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; + for (ids) |id| if (id >= g.vocab) return error.InvalidQwenToken; + for (parts) |p| if (p.parents != null and !record) return error.InvalidQwenParents; + const heads: usize = if (record) total else 1; + // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps + const prompt = !record; + // the tables + const st = e.staged.bytes; + @memcpy(std.mem.bytesAsSlice(u32, st[ids_at..][0 .. total * 4]), ids); + const segs = std.mem.bytesAsSlice(Seg, st[segs_at..][0 .. parts.len * @sizeOf(Seg)]); + try rowTables(st, parts); + var row0: usize = 0; + for (parts, 0..) |p, i| { + segs[i] = p.cache.seg(e, p.cache.len, row0, p.rows); + row0 += p.rows; + } + if (keep == .draws) { + const rules = std.mem.bytesAsSlice(cuda.sample.Rule, st[rules_at..][0 .. total * @sizeOf(cuda.sample.Rule)]); + const meta = std.mem.bytesAsSlice(c_int, st[meta_at..][0 .. total * 4]); + for (keep.draws, rules, meta) |p, *r, *m| { + if (p.rule) |x| r.* = x; + m.* = @intCast(p.position -% 1); // sample.cu keys row r at meta + r + 1 + } + } + // words 0 and 1 are reserved (each row's attention chunks follow its own keys: qwen35.cu's chunk_keys) + const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..32]); + params[0] = 0; + params[1] = 0; + params[2] = @intCast(total); + params[3] = @intCast(parts.len); + params[4] = @intFromBool(record); + try e.tables.uploadAsync(0, st, e.stream.handle); + // a captured forward covers a row bucket and every window; kernels read the true counts; a lone step is exact + const lone = total == 1 and parts.len == 1; + const bucket: usize = if (lone) 1 else if (total <= 16) 16 else most_verify; + const windows: usize = if (lone) 1 else max_segs; + const shape: Shape = .{ .total = @intCast(if (prompt) total else bucket), .nseg = @intCast(if (prompt) parts.len else windows), .record = record, .prompt = prompt, .taps = e.n_taps > 0 }; + if (prompt) { + // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length's graph would pay the stall + try e.body(shape); + } else { + const gop = try e.graphs.getOrPut(e.gpa, shape.key()); + if (!gop.found_existing) { + errdefer _ = e.graphs.remove(shape.key()); + // captured, nothing runs: the replay below is this shape's first run + try cuda.graph.beginCapture(e.stream, .thread_local); + e.body(shape) catch |err| { + var broken = cuda.graph.endCapture(e.stream) catch return err; + broken.deinit(); + return err; + }; + var captured = try cuda.graph.endCapture(e.stream); + defer captured.deinit(); + gop.value_ptr.* = try captured.instantiate(); + } + try gop.value_ptr.launchOn(e.stream); + } + for (parts) |p| { + p.cache.last = if (record and p.rows > 1) .{ .start = p.cache.len, .rows = p.rows } else null; + p.cache.len += p.rows; + p.cache.written = p.cache.len; + } + switch (keep) { + .draws => |pk| { + const toks = e.tables.ptr + tokens_at; + for (pk, 0..) |p, i| { + const row = e.logits + i * g.vocab * 2; + if (p.rule != null) { + try e.draws.draw(e.stream, row, g.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); + } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, @intCast(g.vocab)), toks + i * 4 }); + } + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); + }, + .logits, .last => { + const out = e.host_logits.bytes[0 .. heads * g.vocab * 2]; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); + }, + } + } + + /// A forward's shape, which fixes its launches: its tables and the caches' lengths are read on the device. + const Shape = struct { + total: u8, + nseg: u8, + record: bool, + prompt: bool, + taps: bool, + + fn key(s: Shape) u32 { + return @as(u32, s.total) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17 | @as(u32, @intFromBool(s.taps)) << 18; + } + }; + + /// The launches of a forward of `shape`, from the embedding to the kept rows' logits (captured once a shape). + fn body(e: *Engine, shape: Shape) !void { + const g = e.g; + const total: usize = shape.total; + const nseg: usize = shape.nseg; + const record = shape.record; + const prompt = shape.prompt; + const heads: usize = if (record) total else 1; + const t_ids = e.tables.ptr + ids_at; + const t_segs = e.tables.ptr + segs_at; + const t_rows = e.tables.ptr + rowseg_at; + const t_par = e.tables.ptr + rowpar_at; + const t_anc = e.tables.ptr + rowanc_at; + const t_params = e.tables.ptr + params_at; + const geo = e.geo; + const group = g.q_heads / g.kv_heads; + + try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)), t_params }); + var di: usize = 0; + var ai: c_int = 0; + for (e.layers, 0..) |l, i| { + try e.norm(i != 0, l.input_norm, total); + switch (l.mixer) { + .delta => |m| { + const layer: c_int = @intCast(di); + try e.proj(m.qkv, e.x, e.qkv, total, prompt); + try e.proj(m.z, e.x, e.z, total, prompt); + try e.proj(m.a, e.x, e.a, total, prompt); + try e.proj(m.b, e.x, e.b, total, prompt); + try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, t_par, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta, t_params }); + try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, t_par, layer, t_params }); + try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, t_par, layer, e.y, t_params }); + try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix, t_params }); + try e.proj(m.out, e.mix, e.r, total, prompt); + di += 1; + }, + .attention => |m| { + const aw = g.attnWidth(); + try e.proj(m.q, e.x, e.qraw, total, prompt); + try e.proj(m.k, e.x, e.kraw, total, prompt); + try e.proj(m.v, e.x, e.vraw, total, prompt); + try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps, t_params }); + try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, t_anc, ai, t_params }); + try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, t_anc, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_segs, t_rows, t_anc, t_params, e.y }); + try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); + try e.proj(m.o, e.mix, e.r, total, prompt); + ai += 1; + }, + } + try e.norm(true, l.post_norm, total); + switch (l.mlp) { + .dense => |m| { + try e.proj(m.gate, e.x, e.gate, total, prompt); + try e.proj(m.up, e.x, e.up, total, prompt); + const n: c_longlong = @intCast(total * g.intermediate); + try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n, t_params, @as(c_int, @intCast(g.intermediate)) }); + try e.proj(m.down, e.act, e.r, total, prompt); + }, + .moe => |m| try e.routedRows(m.router, m.experts, e.x, e.r, total, prompt and total > most_window), + } + if (shape.taps) for (e.taps[0..e.n_taps], 0..) |t, ti| if (t == i) { + try e.run(e.fns.tap, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, e.r, t_segs, t_rows, @as(c_int, @intCast(ti)), @as(c_int, @intCast(g.hidden)), t_params }); + }; + } + try e.norm(true, e.final_norm, total); + // the head over the kept rows + try e.proj(e.head, e.x + (total - heads) * g.hidden * 2, e.logits, heads, false); + } + + /// Routed experts over `total` normed rows at x into out; `wide` rows take the prefill kernels, else decode's. + pub fn routedRows(e: *Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u64, total: usize, wide: bool) !void { + const g = e.g; + const slots = g.slots(); + const count = g.experts + 1; + const pairs = total * slots; + const tile: usize = if (wide) 64 else 16; + const items = cuda.grouped.maxItems(pairs, count, tile); + const plan: cuda.grouped.Plan = .{ .members = e.members, .items = e.items, .counts = e.counts, .rank = e.rank, .hist = e.hist }; + const grouped: cuda.experts.Grouped = .{ .affine4 = &e.experts }; + try e.moe.route(e.stream, x, g.hidden, total, router, e.rlogits, g.hidden, count); + try e.moe.pick(e.stream, e.rlogits, e.picks, e.wts, total, g.experts, g.top_k); + try e.plan.route(e.stream, e.picks, pairs, count, tile, plan); + const rows_in: cuda.experts.Rows = .{ .x = x, .stride = g.hidden, .slots = slots }; + const act_in: cuda.experts.Rows = .{ .x = e.eact, .stride = g.intermediate }; + if (wide) { + try grouped.prompt(e.stream, .up, rows_in, layer, plan, e.eact, items, -1); + try grouped.prompt(e.stream, .down_bf16, act_in, layer, plan, e.ey, items, -1); + } else { + try grouped.decode(e.stream, .up, rows_in, layer, plan, e.eact, items, -1); + try grouped.decode(e.stream, .down_f32, act_in, layer, plan, e.ey, items, -1); + } + try e.moe.combine(e.stream, e.ey, !wide, e.wts, out, total, slots, g.hidden); + } + + /// Row `i`'s logits from the newest forward that kept logits. + pub fn rowLogits(e: *const Engine, i: usize) []const u16 { + return e.host_logits.slice(u16)[i * e.g.vocab ..][0..e.g.vocab]; + } +}; diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig new file mode 100644 index 0000000000..3640559862 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -0,0 +1,235 @@ +//! A Qwen3.5 hybrid's widths (2B, 27B) as the CUDA engine reads them from config.json; what qwen35.cu fixes is checked. +const std = @import("std"); + +pub const max_layers = 128; +/// qwen35.cu's fixed widths. +pub const dk = 128; +pub const taps = 4; +pub const head_dim = 256; +pub const max_query_heads = 32; +/// The rotary base the kernels fix (parse checks the config says so). +pub const theta: f32 = 10000000; + +pub const Geometry = struct { + hidden: usize, + intermediate: usize, + layers: usize, + vocab: usize, + k_heads: usize, // linear-attention key heads + v_heads: usize, // linear-attention value heads + q_heads: usize, // attention query heads + kv_heads: usize, // attention key/value heads + eps: f32, + tied: bool, + context: usize, + linear: [max_layers]bool, + /// Routed experts (qwen3_5_moe): their count and the picks a row makes; 0: a dense MLP of `intermediate`. + experts: usize = 0, + top_k: usize = 0, + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Geometry { + 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 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + /// Whether the MLPs are routed experts, whose width is `intermediate` (the shared expert's too). + pub fn moe(g: Geometry) bool { + return g.experts != 0; + } + /// A row's expert slots: its picks, then the shared expert. + pub fn slots(g: Geometry) usize { + return g.top_k + 1; + } + /// Value-head width of the gated delta rule's rows (and of z and the out projection's input). + pub fn valueWidth(g: Geometry) usize { + return g.v_heads * dk; + } + + /// Channels the causal convolution sees: q, k and v. + pub fn convWidth(g: Geometry) usize { + return (2 * g.k_heads + g.v_heads) * dk; + } + + /// Attention output width (q_heads of head_dim), the o projection's input. + pub fn attnWidth(g: Geometry) usize { + return g.q_heads * head_dim; + } + + pub fn deltas(g: Geometry) usize { + var n: usize = 0; + for (g.linear[0..g.layers]) |l| n += @intFromBool(l); + return n; + } + + pub fn attentions(g: Geometry) usize { + return g.layers - g.deltas(); + } + + /// Bytes of one linear layer's recurrence (fp32) and convolution tail (bf16). + pub fn recBytes(g: Geometry) usize { + return g.v_heads * dk * dk * 4; + } + + pub fn convBytes(g: Geometry) usize { + return (taps - 1) * g.convWidth() * 2; + } + + /// A stream's recurrences and tails. + pub fn stateBytes(g: Geometry) usize { + return g.deltas() * (g.recBytes() + g.convBytes()); + } + + /// A row's record of a linear layer (qwen35.cuh log_bytes): decay and write a value head, key a key head, the tail. + pub fn logRowBytes(g: Geometry) usize { + return g.v_heads * 4 + g.v_heads * dk * 4 + g.k_heads * dk * 2 + g.convBytes(); + } + + /// A stream's log of a window's `rows` rows. + pub fn logBytes(g: Geometry, rows: usize) usize { + return rows * g.deltas() * g.logRowBytes(); + } + + /// A stream's keys and values at a `cap`-token window. + pub fn kvBytes(g: Geometry, cap: usize) usize { + return g.attentions() * 2 * g.kv_heads * cap * head_dim * 2; + } +}; + +fn object(v: std.json.Value) !std.json.ObjectMap { + return if (v == .object) v.object else error.BadQwenConfig; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) !usize { + const v = o.get(key) orelse return error.BadQwenConfig; + if (v != .integer or v.integer < 1) return error.BadQwenConfig; + return @intCast(v.integer); +} + +fn number(o: std.json.ObjectMap, key: []const u8) !f64 { + return switch (o.get(key) orelse return error.BadQwenConfig) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadQwenConfig, + }; +} + +fn is(o: std.json.ObjectMap, key: []const u8, comptime T: type, want: T) !void { + const v = o.get(key) orelse return error.UnsupportedQwenConfig; + const ok = switch (T) { + bool => v == .bool and v.bool == want, + []const u8 => v == .string and std.mem.eql(u8, v.string, want), + usize => v == .integer and v.integer == want, + else => @compileError("type"), + }; + if (!ok) return error.UnsupportedQwenConfig; +} + +pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer doc.deinit(); + const root = try object(doc.value); + const moe = blk: { + const v = root.get("model_type") orelse return error.UnsupportedQwenConfig; + if (v != .string) return error.UnsupportedQwenConfig; + if (std.mem.eql(u8, v.string, "qwen3_5")) break :blk false; + if (std.mem.eql(u8, v.string, "qwen3_5_moe")) break :blk true; + return error.UnsupportedQwenConfig; + }; + const t = try object(root.get("text_config") orelse return error.BadQwenConfig); + try is(t, "model_type", []const u8, if (moe) "qwen3_5_moe_text" else "qwen3_5_text"); + try is(t, "attention_bias", bool, false); + try is(t, "attn_output_gate", bool, true); + try is(t, "hidden_act", []const u8, "silu"); + try is(t, "head_dim", usize, head_dim); + try is(t, "linear_key_head_dim", usize, dk); + try is(t, "linear_value_head_dim", usize, dk); + try is(t, "linear_conv_kernel_dim", usize, taps); + const rope = try object(t.get("rope_parameters") orelse return error.UnsupportedQwenRotary); + try is(rope, "rope_type", []const u8, "default"); + if (try number(rope, "rope_theta") != 10000000 or try number(rope, "partial_rotary_factor") != 0.25) return error.UnsupportedQwenRotary; + const q = try object(root.get("quantization") orelse return error.UnsupportedQwenQuantization); + if (try integer(q, "bits") != 4 or try integer(q, "group_size") != 64) return error.UnsupportedQwenQuantization; + if (q.get("mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "affine")) return error.UnsupportedQwenQuantization; + var g: Geometry = .{ + .hidden = try integer(t, "hidden_size"), + .intermediate = try integer(t, if (moe) "moe_intermediate_size" else "intermediate_size"), + .layers = try integer(t, "num_hidden_layers"), + .vocab = try integer(t, "vocab_size"), + .k_heads = try integer(t, "linear_num_key_heads"), + .v_heads = try integer(t, "linear_num_value_heads"), + .q_heads = try integer(t, "num_attention_heads"), + .kv_heads = try integer(t, "num_key_value_heads"), + .eps = @floatCast(try number(t, "rms_norm_eps")), + .tied = blk: { + const v = t.get("tie_word_embeddings") orelse root.get("tie_word_embeddings") orelse break :blk false; + if (v != .bool) return error.BadQwenConfig; + break :blk v.bool; + }, + .context = try integer(t, "max_position_embeddings"), + .linear = @splat(false), + }; + if (g.layers > max_layers or g.hidden % 256 != 0 or g.hidden > 8192 or g.intermediate % 64 != 0) return error.UnsupportedQwenGeometry; + if (moe) { + g.experts = try integer(t, "num_experts"); + g.top_k = try integer(t, "num_experts_per_tok"); + // the shared expert runs in the routed experts' table, so it needs their width + if (try integer(t, "shared_expert_intermediate_size") != g.intermediate) return error.UnsupportedQwenGeometry; + if (g.experts > 256 or g.top_k > 31 or g.top_k >= g.experts) return error.UnsupportedQwenGeometry; + } + if (g.v_heads % g.k_heads != 0 or g.q_heads % g.kv_heads != 0 or g.q_heads > max_query_heads) return error.UnsupportedQwenGeometry; + const types = t.get("layer_types") orelse return error.UnsupportedQwenGeometry; + if (types != .array or types.array.items.len != g.layers) return error.UnsupportedQwenGeometry; + for (types.array.items, 0..) |kind, i| { + if (kind != .string) return error.UnsupportedQwenGeometry; + if (std.mem.eql(u8, kind.string, "linear_attention")) g.linear[i] = true else if (!std.mem.eql(u8, kind.string, "full_attention")) return error.UnsupportedQwenGeometry; + } + return g; +} + +test "the 35B-A3B's routed experts: 256 of width 512, 8 a row, the shared one the same width" { + const text = + \\{"model_type":"qwen3_5_moe","quantization":{"group_size":64,"bits":4,"mode":"affine"},"tie_word_embeddings":false, + \\ "text_config":{"model_type":"qwen3_5_moe_text","attention_bias":false,"attn_output_gate":true,"hidden_act":"silu", + \\ "head_dim":256,"linear_key_head_dim":128,"linear_value_head_dim":128,"linear_conv_kernel_dim":4, + \\ "hidden_size":2048,"moe_intermediate_size":512,"shared_expert_intermediate_size":512,"num_experts":256, + \\ "num_experts_per_tok":8,"num_hidden_layers":4,"vocab_size":248320,"linear_num_key_heads":16, + \\ "linear_num_value_heads":32,"num_attention_heads":16,"num_key_value_heads":2,"rms_norm_eps":1e-06, + \\ "max_position_embeddings":262144, + \\ "layer_types":["linear_attention","linear_attention","linear_attention","full_attention"], + \\ "rope_parameters":{"rope_type":"default","rope_theta":10000000,"partial_rotary_factor":0.25}}} + ; + const g = try parse(std.testing.allocator, text); + try std.testing.expect(g.moe()); + try std.testing.expectEqual(@as(usize, 256), g.experts); + try std.testing.expectEqual(@as(usize, 9), g.slots()); + try std.testing.expectEqual(@as(usize, 512), g.intermediate); + try std.testing.expectEqual(@as(usize, 8192), g.convWidth()); + const odd = try std.mem.replaceOwned(u8, std.testing.allocator, text, "\"shared_expert_intermediate_size\":512", "\"shared_expert_intermediate_size\":1024"); + defer std.testing.allocator.free(odd); + try std.testing.expectError(error.UnsupportedQwenGeometry, parse(std.testing.allocator, odd)); +} + +test "the 27B's widths: 48 value heads over 16 key heads, 24 query heads over 4" { + const text = + \\{"model_type":"qwen3_5","quantization":{"group_size":64,"bits":4,"mode":"affine"},"tie_word_embeddings":false, + \\ "text_config":{"model_type":"qwen3_5_text","attention_bias":false,"attn_output_gate":true,"hidden_act":"silu", + \\ "head_dim":256,"linear_key_head_dim":128,"linear_value_head_dim":128,"linear_conv_kernel_dim":4, + \\ "hidden_size":5120,"intermediate_size":17408,"num_hidden_layers":4,"vocab_size":248320,"linear_num_key_heads":16, + \\ "linear_num_value_heads":48,"num_attention_heads":24,"num_key_value_heads":4,"rms_norm_eps":1e-06, + \\ "max_position_embeddings":262144,"tie_word_embeddings":false, + \\ "layer_types":["linear_attention","linear_attention","linear_attention","full_attention"], + \\ "rope_parameters":{"rope_type":"default","rope_theta":10000000,"partial_rotary_factor":0.25}}} + ; + const g = try parse(std.testing.allocator, text); + try std.testing.expectEqual(@as(usize, 3), g.deltas()); + try std.testing.expectEqual(@as(usize, 10240), g.convWidth()); + try std.testing.expectEqual(@as(usize, 6144), g.attnWidth()); + try std.testing.expectEqual(@as(usize, 48 * 128 * 128 * 4), g.recBytes()); + try std.testing.expect(!g.tied); + const wide = try std.mem.replaceOwned(u8, std.testing.allocator, text, "\"head_dim\":256", "\"head_dim\":128"); + defer std.testing.allocator.free(wide); + try std.testing.expectError(error.UnsupportedQwenConfig, parse(std.testing.allocator, wide)); +} diff --git a/zig/src/families/qwen3_5/cuda_mtp.zig b/zig/src/families/qwen3_5/cuda_mtp.zig new file mode 100644 index 0000000000..fddb55f77a --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_mtp.zig @@ -0,0 +1,530 @@ +//! The Qwen3.6 MTP layer as the drafter on CUDA: chains or trees from a kept row's final normed state and next token. +const std = @import("std"); +const platform = @import("sys"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const geometry = @import("cuda_geometry.zig"); +const weights = @import("cuda_mtp_weights.zig"); +const timing = @import("cuda_draft.zig").timing; + +const E = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; +const dr = lanes.drafter; +const Stream = lanes.Stream; +const head_dim = geometry.head_dim; +const max_rows = engine.max_rows; +const most_window = engine.most_window; + +/// Drafts a round at most (the frozen engine's DEPTH); TF_MTP_DRAFTS=N sets it, up to a window's rows less one. +const default_depth = 3; +/// A tree's lanes at most (the lane planner's widest shape within a window), and the ranks a row's candidates cover. +const max_lanes = most_window; +const ranks = lanes.shape.ranks; +/// The logprob row's words (sample.cu's LP_WORDS): the ids of a row's best tokens sit at words 2 onward. +const lp_words = 3 + 2 * 20; + +const Fns = struct { norm_rope: cuda.Function, attn: cuda.Function, attn_tree: cuda.Function }; +/// qwen35_draft.cu's ring kernels, shared with the DFlash2 drafter. +const Ring = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, ring_put: cuda.Function }; +const TapRows = extern struct { n: c_int, rows: [most_window]c_int }; + +/// A stream's context: its ring of keys and values over [begin, end), the newest row's state, and its held drafts. +pub const Ctx = struct { + ring: DeviceBuffer, + last: DeviceBuffer, // the newest absorbed row's normed output (bf16 hidden): the first draft row's state + begin: u64 = 0, + end: u64 = 0, + cands: [ranks]u32 = @splat(0), // the newest absorbed row's best tokens by rank (rank 0: its chain draft) + held: [max_lanes]u32 = undefined, + held_n: usize = 0, +}; + +/// How a layer pass attends: absorbed or chained rows over the ring and the block's earlier rows, into the ring too. +const Pass = union(enum) { ring, tree: u64 }; + +pub const Mtp = struct { + gpa: std.mem.Allocator, + io: std.Io, + e: *E, + w: weights.Weights, + mod: cuda.Module, + dmod: cuda.Module, + fns: Fns, + ring_fns: Ring, + scratch: DeviceBuffer, + at: usize = 0, + capacity: usize, // ring rows: every position of the stream's window + depth: u32, + trees: bool, // TF_MTP_TREES=1: the planner's trees instead of chains (measured a wash on this head) + tap: [1]u32, + // rows of max_rows, bf16 unless said + taps_in: u64, + hin: u64, + hn: u64, + emb: u64, + en: u64, + h: u64, + r: u64, + x: u64, + qraw: u64, + q: u64, + kraw: u64, + kb: u64, + vb: u64, + att: u64, + xs: u64, // f32 group sums for the projections + logits: u64, // max_lanes rows + words: u64, // u32 logprob words a lane row + picks: u64, // u32 zeros: the logprob rows' pick (unused) + lane_state: u64, // each tree lane's normed output + lane_k: u64, // each tree lane's key and value rows + lane_v: u64, + anc: u64, // int32 ancestor table a row: a count, then lane indices + ids: u64, // u32 the rows' tokens + tok: u64, // u32 the drawn token + staged: cuda.HostBuffer, // pinned: the ids and ancestors up, the token and words down + streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, + /// Contexts between streams, allocated and run once (`reserve`), as the driver stalls a buffer's first kernel. + pool: std.ArrayList(*Ctx) = .empty, + /// One draft level's cost (timed at load): a chained row, or a tree's rows at one depth. + step_ms: f64 = 1.0, + /// The target's window ms by rows (cuda_native.zig times them), for the round loop's prices. + costs: [most_window]lanes.config.Cost = undefined, + n_costs: usize = 0, + + const anc_stride = 1 + max_lanes; + const ids_at = 0; + const tok_at = max_rows * 4; + const words_at = tok_at + 64; + const anc_at = words_at + max_lanes * lp_words * 4; + const staged_len = anc_at + max_lanes * anc_stride * 4; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8, capacity: usize) !*Mtp { + if (!e.g.moe()) return error.NotRouted; + const d = try gpa.create(Mtp); + errdefer gpa.destroy(d); + d.* = .{ .gpa = gpa, .io = io, .e = e, .w = undefined, .mod = undefined, .dmod = undefined, .fns = undefined, .ring_fns = undefined, .scratch = undefined, .capacity = capacity, .depth = default_depth, .trees = platform.c.getenv("TF_MTP_TREES") != null, .tap = .{@intCast(e.g.layers - 1)}, .taps_in = 0, .hin = 0, .hn = 0, .emb = 0, .en = 0, .h = 0, .r = 0, .x = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .xs = 0, .logits = 0, .words = 0, .picks = 0, .lane_state = 0, .lane_k = 0, .lane_v = 0, .anc = 0, .ids = 0, .tok = 0, .staged = undefined }; + if (platform.c.getenv("TF_MTP_DRAFTS")) |v| d.depth = std.fmt.parseInt(u32, std.mem.span(v), 10) catch default_depth; + d.depth = @min(@max(d.depth, 1), most_window - 1); + d.mod = try cuda.Module.load(e.d, cuda.kernels.qwen36_mtp); + errdefer d.mod.unload(); + d.dmod = try cuda.Module.load(e.d, cuda.kernels.qwen35_draft); + errdefer d.dmod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(d.fns, name) = try d.mod.function("tf_q36m_" ++ name); + inline for (@typeInfo(Ring).@"struct".field_names) |name| @field(d.ring_fns, name) = try d.dmod.function("tf_q35d_" ++ name); + d.w = try weights.load(gpa, io, e, dir); + errdefer d.w.deinit(); + d.scratch = try DeviceBuffer.alloc(e.d, d.scratchBytes()); + errdefer d.scratch.free(); + try d.scratch.fill8(0, null); + d.at = 0; + _ = try d.layout(true); + d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); + errdefer d.staged.free(); + // the first context, and a chain timed on it: the quickest of sixteen, as the driver's stalls last a dozen + const ctx = try d.create(); + errdefer d.drop(ctx); + var quickest: f64 = std.math.inf(f64); + for (0..16) |_| { + try d.absorb(ctx, d.taps_in, 0, 1, null, &.{0}); + const t0 = std.Io.Clock.awake.now(io); + try d.hold(ctx, 0, 1, d.depth, null); + const ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6; + quickest = @min(quickest, ms); + } + d.step_ms = quickest / @as(f64, @floatFromInt(d.depth)); + d.resetCtx(ctx); + try d.pool.append(gpa, ctx); + return d; + } + + /// The drafter's tap: the target's last layer, whose residual the final norm then reads. + pub fn taps(d: *const Mtp) []const u32 { + return &d.tap; + } + + pub fn setCost(d: *Mtp, rows: u32, ms: f64) void { + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a draft level {d:.2})\n", .{ rows, ms, d.step_ms }); + if (d.n_costs < d.costs.len) { + d.costs[d.n_costs] = .{ .width = rows, .ms = ms }; + d.n_costs += 1; + } + } + + pub fn deinit(d: *Mtp) void { + var it = d.streams.valueIterator(); + while (it.next()) |ctx| d.drop(ctx.*); + d.streams.deinit(d.gpa); + for (d.pool.items) |ctx| d.drop(ctx); + d.pool.deinit(d.gpa); + d.staged.free(); + d.scratch.free(); + d.w.deinit(); + d.dmod.unload(); + d.mod.unload(); + d.gpa.destroy(d); + } + + pub fn kvWidth(d: *const Mtp) usize { + return d.e.g.kv_heads * head_dim; + } + + fn ringBytes(d: *const Mtp) usize { + return 2 * d.capacity * d.kvWidth() * 2; + } + + fn scratchBytes(d: *Mtp) usize { + d.at = 0; + return d.layout(false) catch unreachable; + } + + /// The scratch rows, measured (nothing assigned) or assigned from the scratch buffer. + fn layout(d: *Mtp, assign: bool) !usize { + const g = d.e.g; + const R = max_rows; + const L = max_lanes; + const aw = g.attnWidth(); + const kw = d.kvWidth(); + const sizes = .{ + .{ "taps_in", R * g.hidden * 2 }, .{ "hin", R * g.hidden * 2 }, .{ "hn", R * g.hidden * 2 }, .{ "emb", R * g.hidden * 2 }, + .{ "en", R * g.hidden * 2 }, .{ "h", R * g.hidden * 2 }, .{ "r", R * g.hidden * 2 }, .{ "x", R * g.hidden * 2 }, + .{ "qraw", R * 2 * aw * 2 }, .{ "q", R * aw * 2 }, .{ "kraw", R * kw * 2 }, .{ "kb", R * kw * 2 }, + .{ "vb", R * kw * 2 }, .{ "att", R * aw * 2 }, .{ "xs", R * (@max(g.hidden, aw) / 64) * 4 }, .{ "logits", L * g.vocab * 2 }, + .{ "words", L * lp_words * 4 }, .{ "picks", L * 4 }, .{ "lane_state", L * g.hidden * 2 }, .{ "lane_k", L * kw * 2 }, + .{ "lane_v", L * kw * 2 }, .{ "anc", L * anc_stride * 4 }, .{ "ids", R * 4 }, .{ "tok", 64 }, + }; + var total: usize = 0; + inline for (sizes) |s| { + const at = engine.aligned(total); + total = at + s[1]; + if (assign) @field(d, s[0]) = d.scratch.ptr + at; + } + return total; + } + + pub fn drop(d: *Mtp, ctx: *Ctx) void { + ctx.last.free(); + ctx.ring.free(); + d.gpa.destroy(ctx); + } + + fn resetCtx(d: *Mtp, ctx: *Ctx) void { + _ = d; + ctx.* = .{ .ring = ctx.ring, .last = ctx.last }; + } + + /// A new context, touched once (an absorb, then a chain), so no request pays the driver's stall. + pub fn create(d: *Mtp) !*Ctx { + const ctx = try d.gpa.create(Ctx); + errdefer d.gpa.destroy(ctx); + ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.ringBytes()), .last = undefined }; + errdefer ctx.ring.free(); + ctx.last = try DeviceBuffer.alloc(d.e.d, d.e.g.hidden * 2); + errdefer ctx.last.free(); + try ctx.ring.fill8(0, null); + try ctx.last.fill8(0, null); + try d.absorb(ctx, d.taps_in, 0, 1, null, &.{0}); // any tap rows do: the drafter's own zeroed staging + try d.hold(ctx, 0, 1, d.depth, null); + d.resetCtx(ctx); + return ctx; + } + + /// `n` contexts ready in the pool, for the streams the server admits. + pub fn reserve(d: *Mtp, n: usize) !void { + try d.pool.ensureTotalCapacity(d.gpa, n); + while (d.pool.items.len < n) d.pool.appendAssumeCapacity(try d.create()); + } + + pub fn ringK(d: *Mtp, ctx: *Ctx) u64 { + _ = d; + return ctx.ring.ptr; + } + + pub fn ringV(d: *Mtp, ctx: *Ctx) u64 { + return ctx.ring.ptr + d.capacity * d.kvWidth() * 2; + } + + /// out (rows, n) = x (rows, k) through `l`, as the target's decode tiles or prompt GEMM, with the drafter's sums. + fn proj(d: *Mtp, l: cuda.qlinear.Affine4, x: u64, out: u64, rows: usize) !void { + const e = d.e; + if (rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, d.xs, @as(c_int, @intCast(l.k)), @as(u64, 0) }); + try e.affine.decode(e.stream, x, d.xs, l, out, rows); + } + + /// out = norm(h + r) with w; h takes the sum (r 0: no residual). + fn norm(d: *Mtp, h: u64, r: u64, w: u64, out: u64, rows: usize) !void { + const e = d.e; + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ h, r, w, out, e.g.eps, @as(c_int, @intCast(e.g.hidden)), @as(u64, 0) }); + } + + /// The layer over `rows` rows at slots pos0..: d.ids their tokens, d.hin their states; x gets their normed outputs. + pub fn layer(d: *Mtp, ctx: *Ctx, rows: usize, pos0: u64, pass: Pass) !void { + const e = d.e; + const g = e.g; + const w = &d.w; + const aw = g.attnWidth(); + const kw = d.kvWidth(); + const p0: c_longlong = @intCast(pos0); + // a tree pass's rows are sibling lanes at one depth, so they share pos0; chained or absorbed rows step by one + const step: c_int = @intFromBool(pass == .ring); + try e.run(e.fns.embed, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.ids, e.embed_w, e.embed_s, e.embed_b, d.emb, @as(c_int, @intCast(g.hidden)), @as(u64, 0) }); + try d.norm(d.emb, 0, w.norm_e, d.en, rows); + try d.norm(d.hin, 0, w.norm_h, d.hn, rows); + try d.proj(w.fc_e, d.en, d.h, rows); + try d.proj(w.fc_h, d.hn, d.r, rows); + try d.norm(d.h, d.r, w.input_norm, d.x, rows); + try d.proj(w.q, d.x, d.qraw, rows); + try d.proj(w.k, d.x, d.kraw, rows); + try d.proj(w.v, d.x, d.vb, rows); + try e.run(d.fns.norm_rope, .{ rows * g.q_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.qraw, w.q_norm, d.q, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), p0, step, g.eps, geometry.theta }); + try e.run(d.fns.norm_rope, .{ rows * g.kv_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.kraw, w.k_norm, d.kb, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(kw)), @as(c_int, head_dim), p0, step, g.eps, geometry.theta }); + const heads: c_int = @intCast(g.q_heads); + const kvh: c_int = @intCast(g.kv_heads); + const cap: c_int = @intCast(d.capacity); + const begin: c_longlong = @intCast(ctx.begin); + switch (pass) { + .ring => try e.run(d.fns.attn, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), heads, kvh, cap, begin, p0, d.att }), + .tree => |anc| try e.run(d.fns.attn_tree, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.lane_k, d.lane_v, anc, @as(c_int, anc_stride), d.ringK(ctx), d.ringV(ctx), heads, kvh, cap, begin, @as(c_longlong, @intCast(ctx.end)), d.att }), + } + try d.proj(w.o, d.att, d.r, rows); + try d.norm(d.h, d.r, w.post_norm, d.x, rows); + try e.routedRows(w.mlp.moe.router, w.mlp.moe.experts, d.x, d.r, rows, rows > most_window); + try d.norm(d.h, d.r, w.norm, d.x, rows); + if (pass == .ring) try e.run(d.ring_fns.ring_put, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), @as(c_int, @intCast(kw)), p0, cap }); + } + + /// Rows [first, first + n) of x through the target's head: their logits and best tokens' words, row 0's pick. + fn draw(d: *Mtp, first: usize, n: usize, cands: bool) !void { + const e = d.e; + try d.proj(e.head, d.x + first * e.g.hidden * 2, d.logits, n); + try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ d.logits, @as(c_int, @intCast(e.g.vocab)), d.tok }); + if (cands and d.trees) try e.draws.logprobRows(e.stream, d.logits, e.g.vocab, d.picks, ranks, false, d.words, n); + } + + /// The drawn token, and the first `n` rows' candidates by rank, after a draw. + fn fetch(d: *Mtp, n: usize, cands: [][ranks]u32) !u32 { + const e = d.e; + const st = d.staged.bytes; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + tok_at, d.tok, 4, e.stream.handle), "cuMemcpyDtoHAsync"); + if (d.trees and n > 0) try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + words_at, d.words, n * lp_words * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + for (cands[0..n], 0..) |*c, r| { + const w = std.mem.bytesAsSlice(u32, st[words_at + r * lp_words * 4 ..][0 .. lp_words * 4]); + for (c, 0..) |*t, j| t.* = if (d.trees) w[2 + j] else std.mem.bytesToValue(u32, st[tok_at..][0..4]); + } + return std.mem.bytesToValue(u32, st[tok_at..][0..4]); + } + + pub fn putIds(d: *Mtp, ids: []const u32) !void { + const e = d.e; + try e.stream.synchronize(); // the previous upload from this staging area has run: an absorb queues one a chunk + @memcpy(std.mem.bytesAsSlice(u32, d.staged.bytes[ids_at..][0 .. ids.len * 4]), ids); + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(d.ids, d.staged.bytes.ptr + ids_at, ids.len * 4, e.stream.handle), "cuMemcpyHtoDAsync"); + } + + /// Target rows [start, start + n) from the tap ring with the token after each; `rows`, when given, picks them. + pub fn absorb(d: *Mtp, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32, follow: []const u32) !void { + const e = d.e; + const g = e.g; + var start = start_; + var n = n_; + var fi: usize = 0; + if (follow.len < n) return error.BadDraftRows; + if (rows_) |rs| if (rs.len != n or n > most_window) return error.BadDraftRows; + if (start < ctx.end) ctx.end = start; + if (start > ctx.end) ctx.begin = start; // a gap (a prompt restored from a cache): the context starts here + const held = engine.tap_rows - max_rows; // the target's ring holds this many rows behind its newest + if (n > held) { + fi = n - held; + start += n - held; + n = held; + ctx.begin = start; + } + ctx.end = start; + while (n > 0) { + const rows = @min(n, max_rows); + if (rows_) |rs| { + var tr: TapRows = .{ .n = @intCast(rows), .rows = @splat(0) }; + for (rs, 0..) |r, i| tr.rows[i] = @intCast(r); + try e.run(d.ring_fns.gather_rows, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), tr, @as(c_int, @intCast(g.hidden)), d.taps_in }); + } else try e.run(d.ring_fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(g.hidden)), d.taps_in }); + try d.norm(d.taps_in, 0, e.final_norm, d.hin, rows); // the target's final norm: the rows the head read + try d.putIds(follow[fi..][0..rows]); + try d.layer(ctx, rows, start, .ring); + start += rows; + n -= rows; + fi += rows; + ctx.end = start; + // the newest row: its state kept for the drafts, its best tokens drawn + try ctx.last.copyFrom(0, d.x + (rows - 1) * g.hidden * 2, g.hidden * 2, e.stream.handle); + try d.draw(rows - 1, 1, true); + } + var cands: [1][ranks]u32 = undefined; + ctx.cands[0] = try d.fetch(1, &cands); + if (d.trees) ctx.cands = cands[0]; + } + + /// `depth` drafts after `pending` at pos0 (the pending token's slot): a chain, or the planner's tree of lanes. + fn hold(d: *Mtp, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { + _ = pending; // the newest absorbed row's follow token, whose candidates ctx.cands are + if (ctx.end != pos0) return error.DraftContextBehind; + if (lanes_) |l| return d.holdTree(ctx, pos0, l); + const e = d.e; + const g = e.g; + ctx.held[0] = ctx.cands[0]; + ctx.held_n = 1; + var k: usize = 1; + while (k < depth and k < ctx.held.len) : (k += 1) { + // row k: (the previous row's normed output, the draft before) at slot pos0 + k - 1 + const src = if (k == 1) ctx.last.ptr else d.x; + try d.scratch.copyFrom(d.hin - d.scratch.ptr, src, g.hidden * 2, e.stream.handle); + try d.putIds(&.{ctx.held[k - 1]}); + try d.layer(ctx, 1, pos0 + k - 1, .ring); + try d.draw(0, 1, false); + var none: [0][ranks]u32 = undefined; + ctx.held[k] = try d.fetch(0, &none); + ctx.held_n = k + 1; + } + } + + /// The tree's lanes, lane i the rank-r candidate of its parent's row; a depth's rows run as one pass. + pub fn holdTree(d: *Mtp, ctx: *Ctx, pos0: u64, l: *const lanes.shape.Shape) !void { + const e = d.e; + const g = e.g; + const n = l.parents.len; + if (n > max_lanes) return error.TooManyLanes; + var cands: [max_lanes][ranks]u32 = undefined; // each lane's row's candidates, once its row has run + var internal: [max_lanes]bool = @splat(false); + for (l.parents, l.depths, l.ranks, 0..) |par, dd, rank, i| { + if (rank >= ranks or par >= @as(i32, @intCast(i)) or dd >= most_window) return error.BadLanes; + if (par >= 0) internal[@intCast(par)] = true else ctx.held[i] = ctx.cands[rank]; + } + // internal lanes by depth: each pass's rows are the lanes at that depth whose rows later lanes read + var done: [max_lanes]bool = @splat(false); + var depth: u8 = 0; + while (depth < most_window) : (depth += 1) { + var batch: [max_lanes]usize = undefined; + var rows: usize = 0; + for (l.depths, 0..) |dd, i| if (dd == depth and internal[i]) { + batch[rows] = i; + rows += 1; + }; + if (rows == 0) { + var any = false; + for (l.depths, 0..) |dd, i| any = any or (dd > depth and internal[i]); + if (!any) break; + continue; + } + // each row's inputs: its parent's output state, its own token; its ancestors as lane indices + const st = d.staged.bytes; + const anc = std.mem.bytesAsSlice(i32, st[anc_at..][0 .. rows * anc_stride * 4]); + var toks: [max_lanes]u32 = undefined; + for (batch[0..rows], 0..) |i, m| { + const par = l.parents[i]; + const src = if (par < 0) ctx.last.ptr else d.lane_state + @as(u64, @intCast(par)) * g.hidden * 2; + try d.scratch.copyFrom(d.hin - d.scratch.ptr + m * g.hidden * 2, src, g.hidden * 2, e.stream.handle); + toks[m] = ctx.held[i]; + var count: usize = 0; + var a = par; + while (a >= 0) : (a = l.parents[@intCast(a)]) { + anc[m * anc_stride + 1 + count] = a; + count += 1; + } + // root first, as a chain's ring holds them: a lane's sums then equal its path's chained rows + std.mem.reverse(i32, anc[m * anc_stride + 1 ..][0..count]); + anc[m * anc_stride] = @intCast(count); + } + try d.putIds(toks[0..rows]); + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(d.anc, st.ptr + anc_at, rows * anc_stride * 4, e.stream.handle), "cuMemcpyHtoDAsync"); + try d.layer(ctx, rows, pos0 + depth, .{ .tree = d.anc }); + const kw = d.kvWidth(); + for (batch[0..rows], 0..) |i, m| { + try d.scratch.copyFrom(d.lane_state - d.scratch.ptr + i * g.hidden * 2, d.x + m * g.hidden * 2, g.hidden * 2, e.stream.handle); + try d.scratch.copyFrom(d.lane_k - d.scratch.ptr + i * kw * 2, d.kb + m * kw * 2, kw * 2, e.stream.handle); + try d.scratch.copyFrom(d.lane_v - d.scratch.ptr + i * kw * 2, d.vb + m * kw * 2, kw * 2, e.stream.handle); + done[i] = true; + } + try d.draw(0, rows, true); + var got: [max_lanes][ranks]u32 = undefined; + _ = try d.fetch(rows, got[0..rows]); + for (batch[0..rows], 0..) |i, m| cands[i] = got[m]; + // the lanes whose parent's row just ran take their tokens now + for (l.parents, l.ranks, 0..) |par, rank, i| if (par >= 0 and done[@intCast(par)]) { + ctx.held[i] = cands[@intCast(par)][rank]; + }; + } + ctx.held_n = n; + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ n, l.levels() }); + } + + // the Drafter contract + + pub fn drafter(d: *Mtp) dr.Drafter { + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(p: *anyopaque) *Mtp { + return @ptrCast(@alignCast(p)); + } + + fn tapsFn(p: *anyopaque) []const u32 { + return cast(p).taps(); + } + + fn factsFn(p: *anyopaque) dr.Facts { + const d = cast(p); + return .{ .depth = d.depth, .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.ringBytes() + d.e.g.hidden * 2, .trees = d.trees, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; + } + + fn openFn(p: *anyopaque, s: *Stream) anyerror!void { + const d = cast(p); + if (d.streams.contains(s)) return error.StreamOpen; + const ctx = d.pool.pop() orelse try d.create(); + errdefer d.pool.append(d.gpa, ctx) catch d.drop(ctx); + d.resetCtx(ctx); + try d.streams.put(d.gpa, s, ctx); + } + + fn absorbFn(p: *anyopaque, items: []const dr.Absorb) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("mtp absorb", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; + // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) + if (it.rows) |rs| { + try d.absorb(ctx, it.features.buffer, it.features.offset, rs.len, rs, it.follow); + } else try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows, null, it.follow); + } + } + + fn holdFn(p: *anyopaque, items: []const dr.Hold) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("mtp hold", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + try d.hold(ctx, it.pending, it.position - 1, it.depth, it.lanes); + } + } + + fn heldFn(p: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const d = cast(p); + for (streams, out) |s, o| { + const ctx = d.streams.get(s) orelse return error.UnknownStream; + if (o.len > ctx.held_n) return error.TooFewDrafts; + @memcpy(o, ctx.held[0..o.len]); + } + } + + fn releaseFn(p: *anyopaque, s: *Stream) void { + const d = cast(p); + if (d.streams.fetchRemove(s)) |kv| d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + } +}; + +/// Acceptance by depth until a stream has its own (the frozen engine's chains kept about three of four). +const prior = [_]f64{ 0.7, 0.55, 0.45, 0.4, 0.35, 0.3, 0.3 }; diff --git a/zig/src/families/qwen3_5/cuda_mtp_check.zig b/zig/src/families/qwen3_5/cuda_mtp_check.zig new file mode 100644 index 0000000000..1eda62b862 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_mtp_check.zig @@ -0,0 +1,133 @@ +//! GPU regressions for the MTP drafter: absorbs in chunks, a tree's lanes against their paths chained, a poisoned ring. +const std = @import("std"); +const lanes = @import("lanes"); +const Mtp = @import("cuda_mtp.zig").Mtp; +const max_rows = @import("cuda_engine.zig").max_rows; + +/// Siblings 0 and 1 at depth 0, six siblings 3 to 8 at depth 1, lane 9 at depth 2 under 3 and 0; each has a child. +const parents = [_]i32{ -1, -1, -1, 0, 0, 0, 0, 1, 1, 3, 4, 5, 6, 7, 8, 9 }; +const ranks = [_]u8{ 0, 1, 2, 0, 1, 2, 3, 0, 1, 0, 0, 0, 0, 0, 0, 0 }; +const depths = [_]u8{ 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3 }; +const internal = [_]usize{ 0, 1, 3, 4, 5, 6, 7, 8, 9 }; + +/// A lane's rows from a pass: its normed output, its key and its value. +const Rows = struct { state: []u8, k: []u8, v: []u8 }; + +/// Absorbs `ids` from the target's tap ring (`tap`, filled by forwards over them), then checks the tree twice. +pub fn run(gpa: std.mem.Allocator, d: *Mtp, tap: u64, ids: []const u32) !bool { + const e = d.e; + const g = e.g; + d.trees = true; + const ctx = try d.create(); + defer d.drop(ctx); + try d.absorb(ctx, tap, 0, ids.len - 1, null, ids[1..]); + const pos0 = ctx.end; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + + // one absorb over several prompt chunks against the same rows absorbed a chunk a call + const apart = try d.create(); + defer d.drop(apart); + var from: usize = 0; + while (from < pos0) : (from += max_rows) { + const n = @min(max_rows, pos0 - from); + try d.absorb(apart, tap, from, n, null, ids[from + 1 ..][0..n]); + } + const ring = 2 * d.capacity * d.kvWidth() * 2; + const whole = try a.alloc(u8, ring + g.hidden * 2); + const parts = try a.alloc(u8, ring + g.hidden * 2); + try download(d, whole[0..ring], ctx.ring.ptr); + try download(d, whole[ring..], ctx.last.ptr); + try download(d, parts[0..ring], apart.ring.ptr); + try download(d, parts[ring..], apart.last.ptr); + var ok = std.mem.eql(u8, whole, parts) and std.mem.eql(u32, &ctx.cands, &apart.cands); + std.debug.print("{s} {d} rows absorbed at once vs a chunk of {d} a call: ring, state and candidates\n", .{ if (ok) "PASS" else "FAIL", pos0, max_rows }); + var p = parents; + var r = ranks; + var dp = depths; + const shape: lanes.shape.Shape = .{ .parents = &p, .ranks = &r, .depths = &dp, .expected = 0 }; + const hidden = g.hidden * 2; + const kw = d.kvWidth() * 2; + + try d.holdTree(ctx, pos0, &shape); + const tree = try lanesRows(a, d, hidden, kw); + const held = ctx.held; + + // the ring at and past the committed end poisoned: a tree pass reads none of it, so nothing may change + const junk = try a.alloc(u8, 16 * kw); + @memset(junk, 0x7f); + for (0..16) |i| { + const slot = (pos0 + i) % d.capacity; + try upload(d, d.ringK(ctx) + slot * kw, junk[0..kw]); + try upload(d, d.ringV(ctx) + slot * kw, junk[0..kw]); + } + try d.holdTree(ctx, pos0, &shape); + const poisoned = try lanesRows(a, d, hidden, kw); + ok = report("poisoned ring suffix, tree rows", poisoned, tree, &internal) and ok; + if (!std.mem.eql(u32, &ctx.held, &held)) { + std.debug.print("FAIL poisoned ring suffix: the lanes' tokens changed\n", .{}); + ok = false; + } + + // each internal lane's path chained from the newest absorbed row, one row a slot, against the tree's rows + var chain: Rows = .{ .state = try a.alloc(u8, tree.state.len), .k = try a.alloc(u8, tree.k.len), .v = try a.alloc(u8, tree.v.len) }; + for (internal) |lane| { + var path: [4]usize = undefined; + var n: usize = 0; + var at: i32 = @intCast(lane); + while (at >= 0) : (at = parents[@intCast(at)]) { + path[n] = @intCast(at); + n += 1; + } + std.mem.reverse(usize, path[0..n]); + for (path[0..n], 0..) |step, k| { + const src = if (k == 0) ctx.last.ptr else d.x; + try d.scratch.copyFrom(d.hin - d.scratch.ptr, src, hidden, e.stream.handle); + try d.putIds(&.{held[step]}); + try d.layer(ctx, 1, pos0 + k, .ring); + } + try download(d, chain.state[lane * hidden ..][0..hidden], d.x); + try download(d, chain.k[lane * kw ..][0..kw], d.kb); + try download(d, chain.v[lane * kw ..][0..kw], d.vb); + } + ok = report("six siblings a depth, tree rows vs each lane's path chained", tree, chain, &internal) and ok; + return ok; +} + +fn lanesRows(a: std.mem.Allocator, d: *Mtp, hidden: usize, kw: usize) !Rows { + const n = parents.len; + const out: Rows = .{ .state = try a.alloc(u8, n * hidden), .k = try a.alloc(u8, n * kw), .v = try a.alloc(u8, n * kw) }; + try download(d, out.state, d.lane_state); + try download(d, out.k, d.lane_k); + try download(d, out.v, d.lane_v); + return out; +} + +/// PASS when every internal lane's state, key and value bytes are equal; else the first lane that differs. +fn report(what: []const u8, got: Rows, want: Rows, lanes_: []const usize) bool { + const hidden = got.state.len / parents.len; + const kw = got.k.len / parents.len; + for (lanes_) |l| { + const same = std.mem.eql(u8, got.state[l * hidden ..][0..hidden], want.state[l * hidden ..][0..hidden]) and + std.mem.eql(u8, got.k[l * kw ..][0..kw], want.k[l * kw ..][0..kw]) and std.mem.eql(u8, got.v[l * kw ..][0..kw], want.v[l * kw ..][0..kw]); + if (!same) { + std.debug.print("FAIL {s}: lane {d} (depth {d}) differs\n", .{ what, l, depths[l] }); + return false; + } + } + std.debug.print("PASS {s}: {d} lanes byte-equal\n", .{ what, lanes_.len }); + return true; +} + +fn upload(d: *Mtp, dst: u64, bytes: []const u8) !void { + const e = d.e; + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, e.stream.handle), "cuMemcpyHtoDAsync"); + try e.stream.synchronize(); +} + +fn download(d: *Mtp, out: []u8, src: u64) !void { + const e = d.e; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, src, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); +} diff --git a/zig/src/families/qwen3_5/cuda_mtp_weights.zig b/zig/src/families/qwen3_5/cuda_mtp_weights.zig new file mode 100644 index 0000000000..09b9770951 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_mtp_weights.zig @@ -0,0 +1,80 @@ +//! The Qwen3.6 MTP layer's weights (mtp-4bit.safetensors beside the checkpoint) onto the device, their own arena. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const engine = @import("cuda_engine.zig"); +const geometry = @import("cuda_geometry.zig"); +const weights = @import("cuda_weights.zig"); + +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const Engine = engine.Engine; +const aligned = engine.aligned; +const head_dim = geometry.head_dim; + +pub const file = "mtp-4bit.safetensors"; + +pub const Weights = struct { + arena: DeviceBuffer, + arena_at: usize = 0, + fc_e: Affine4, // fc's halves: the embedding's inputs, then the hidden state's + fc_h: Affine4, + norm_e: u64, + norm_h: u64, + input_norm: u64, + post_norm: u64, + q_norm: u64, + k_norm: u64, + norm: u64, + q: Affine4, + k: Affine4, + v: Affine4, + o: Affine4, + mlp: engine.Mlp, + + pub fn alloc(w: *Weights, bytes: usize) !u64 { + const at = aligned(w.arena_at); + if (at + bytes > w.arena.len) return error.ArenaExhausted; + w.arena_at = at + bytes; + return w.arena.ptr + at; + } + + pub fn deinit(w: *Weights) void { + w.arena.free(); + } +}; + +/// Device bytes of everything `load` packs. +pub fn total(g: geometry.Geometry) usize { + var t = 2 * weights.packedBytes(g.hidden, g.hidden) + 7 * aligned(@max(g.hidden, head_dim) * 2); + t += weights.packedBytes(2 * g.attnWidth(), g.hidden) + 2 * weights.packedBytes(g.kv_heads * head_dim, g.hidden) + weights.packedBytes(g.hidden, g.attnWidth()); + t += weights.routedBytes(g.experts, g.intermediate, g.hidden); + return t; +} + +pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *Engine, dir: []const u8) !Weights { + const g = e.g; + var ck: core.Checkpoint = .{ .gpa = gpa, .io = io }; + defer ck.close(); + try ck.add(dir, file); + var w: Weights = .{ .arena = try DeviceBuffer.alloc(e.d, total(g)), .fc_e = undefined, .fc_h = undefined, .norm_e = 0, .norm_h = 0, .input_norm = 0, .post_norm = 0, .q_norm = 0, .k_norm = 0, .norm = 0, .q = undefined, .k = undefined, .v = undefined, .o = undefined, .mlp = undefined }; + errdefer w.arena.free(); + try w.arena.fill8(0, null); + const m = "language_model.mtp."; + const p = "language_model.mtp.layers.0."; + w.fc_e = try weights.linearSlice(e, &w, &ck, m, "fc", g.hidden, 2 * g.hidden, 0, g.hidden); + w.fc_h = try weights.linearSlice(e, &w, &ck, m, "fc", g.hidden, 2 * g.hidden, g.hidden, g.hidden); + w.norm_e = try weights.tensor(&w, &ck, "{s}pre_fc_norm_embedding.weight", .{m}, .bf16, &.{g.hidden}); + w.norm_h = try weights.tensor(&w, &ck, "{s}pre_fc_norm_hidden.weight", .{m}, .bf16, &.{g.hidden}); + w.norm = try weights.tensor(&w, &ck, "{s}norm.weight", .{m}, .bf16, &.{g.hidden}); + w.input_norm = try weights.tensor(&w, &ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{g.hidden}); + w.post_norm = try weights.tensor(&w, &ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{g.hidden}); + w.q_norm = try weights.tensor(&w, &ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{head_dim}); + w.k_norm = try weights.tensor(&w, &ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{head_dim}); + w.q = try weights.linear(e, &w, &ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden); + w.k = try weights.linear(e, &w, &ck, p, "self_attn.k_proj", g.kv_heads * head_dim, g.hidden); + w.v = try weights.linear(e, &w, &ck, p, "self_attn.v_proj", g.kv_heads * head_dim, g.hidden); + w.o = try weights.linear(e, &w, &ck, p, "self_attn.o_proj", g.hidden, g.attnWidth()); + w.mlp = try weights.routed(e, &w, &ck, p); + return w; +} diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig new file mode 100644 index 0000000000..3f463e6c76 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -0,0 +1,361 @@ +//! The Qwen3.5 hybrids' CUDA lane backend (native/cuda.zig's registry): 16-row windows, every stream in one forward. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const c = @import("config.zig"); +const E = @import("cuda_engine.zig"); +const draws = @import("draw.zig"); +const Draft = @import("cuda_draft.zig").Draft; +const Mtp = @import("cuda_mtp.zig").Mtp; +const be = lanes.backend; + +pub const model_type = "qwen3_5"; +pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const default_context: i64 = 32768; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = E.max_rows; +/// --drafter loads the DFlash2 draft model (openDrafter); the checkpoints have no head of their own. +pub const takes_drafter = true; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; + +/// What the native server drives: the lane backend, the facts its round loop reads, and how to free it. +pub const Loaded = struct { + backend: be.Backend, + facts: lanes.Model, + rows: u32, + /// Device bytes each admitted stream allocates (recurrences, convolution tails, keys and values). + stream_bytes: usize, + free_streams: u32 = 0, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + /// Called with the admitted stream count: that many caches allocated and touched now, as memory commits on touch. + reserve: *const fn (*anyopaque, u32) anyerror!void, +}; + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.QwenHasNoCapturedKernels; + const e = try E.Engine.init(gpa, io, ctx, dir); + errdefer e.deinit(); + if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary + const l = try gpa.create(Lanes); + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context }; + return .{ + .backend = l.backend(), + .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, + .rows = @intCast(e.window_rows), + .stream_bytes = e.streamBytes(o.context), + .ctx = l, + .deinit = release, + .reserve = reserve, + }; +} + +fn reserve(p: *anyopaque, streams: u32) anyerror!void { + const l: *Lanes = @ptrCast(@alignCast(p)); + try l.pool.ensureTotalCapacity(l.gpa, streams); + while (l.pool.items.len < streams) { + const cache = try l.gpa.create(E.Cache); + errdefer l.gpa.destroy(cache); + cache.* = try E.Cache.init(l.gpa, l.e, l.capacity); + errdefer cache.deinit(l.gpa); + // a row through it: kernels have to have referenced a buffer once before the driver's stall is behind it + try l.e.forward(&.{.{ .cache = cache, .rows = 1 }}, &.{0}, .last); + try cache.reset(l.e); + l.pool.appendAssumeCapacity(cache); + } + // the shared rounds' shapes, a row of 1 to `streams` streams, captured now: a new graph's first run stalls too + var parts: [E.max_segs]E.Part = undefined; + var ids: [E.max_segs]u32 = @splat(0); + var picks: [E.max_segs]E.Draw = undefined; + for (0..@min(streams, E.max_segs)) |n| { + parts[n] = .{ .cache = l.pool.items[n], .rows = 1 }; + picks[n] = .{ .rule = null, .position = 1 }; + try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); + for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); + } + // with a drafter: a lone stream's windows of the pending token and its drafts, chains and trees + if (l.e.n_taps > 0) for (2..l.e.window_rows + 1) |rows| { + for (picks[0..rows]) |*pk| pk.* = .{ .rule = null, .position = 1 }; + try l.e.forward(&.{.{ .cache = l.pool.items[0], .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); + try l.pool.items[0].reset(l.e); + }; +} + +/// The DFlash2 drafter in `dir` over engine `p`, 4-bit only, and the target's windows timed for the round loop. +pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { + if (bits != 4) return error.DrafterBitsUnsupported; + const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; + const d = try Draft.init(gpa, io, l.e, dir); + errdefer d.deinit(); + try l.e.setTaps(d.taps()); + try l.timeWindows(d); + return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter }; +} + +/// The MTP layer beside the checkpoint in `dir` as the drafter over engine `p` (qwen3_5_moe), its ring the window. +pub fn openMtp(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { + if (bits != 4) return error.DrafterBitsUnsupported; + const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; + const d = try Mtp.init(gpa, io, l.e, dir, l.capacity); + errdefer d.deinit(); + try l.e.setTaps(d.taps()); + try l.timeWindows(d); + return .{ .drafter = d.drafter(), .ctx = d, .close = closeMtp, .reserve = reserveMtp }; +} + +fn reserveMtp(p: *anyopaque, streams: u32) anyerror!void { + const d: *Mtp = @ptrCast(@alignCast(p)); + try d.reserve(streams); +} + +fn closeMtp(p: *anyopaque) void { + const d: *Mtp = @ptrCast(@alignCast(p)); + d.deinit(); +} + +fn reserveDrafter(p: *anyopaque, streams: u32) anyerror!void { + const d: *Draft = @ptrCast(@alignCast(p)); + try d.reserve(streams); +} + +fn closeDrafter(p: *anyopaque) void { + const d: *Draft = @ptrCast(@alignCast(p)); + d.deinit(); +} + +/// A request this engine refuses, in words; null: none of its own. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.InvalidQwenPrompt => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.OutOfMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + else => null, + }; +} + +fn release(p: *anyopaque) void { + const l: *Lanes = @ptrCast(@alignCast(p)); + l.deinit(); +} + +/// The device draw of a stream's rule at `position`: sample.cu's keyed rule, greedy at temperature 0 or without one. +fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { + const s = sampling orelse return .{ .rule = null, .position = position }; + if (s.temperature <= 0) return .{ .rule = null, .position = position }; + return .{ .rule = cuda.sample.Rule.from(lanes.gpu_rule.rule(s)), .position = position }; +} + +const Lanes = struct { + gpa: std.mem.Allocator, + io: std.Io, + e: *E.Engine, + capacity: usize, + caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, + /// Caches at the full window, kept between streams: allocated and touched once (`reserve`), reset on reuse. + pool: std.ArrayList(*E.Cache) = .empty, + + fn deinit(self: *Lanes) void { + var it = self.caches.valueIterator(); + while (it.next()) |cache| self.drop(cache.*); + self.caches.deinit(self.gpa); + for (self.pool.items) |cache| self.drop(cache); + self.pool.deinit(self.gpa); + self.e.deinit(); + self.gpa.destroy(self); + } + + fn drop(self: *Lanes, cache: *E.Cache) void { + cache.deinit(self.gpa); + self.gpa.destroy(cache); + } + + /// A cache for a new stream: a pooled one reset, else one more (a window beyond the reserve). + fn take(self: *Lanes) !*E.Cache { + if (self.pool.pop()) |cache| { + try cache.reset(self.e); + return cache; + } + const cache = try self.gpa.create(E.Cache); + errdefer self.gpa.destroy(cache); + cache.* = try E.Cache.init(self.gpa, self.e, self.capacity); + return cache; + } + + fn give(self: *Lanes, cache: *E.Cache) void { + self.pool.append(self.gpa, cache) catch self.drop(cache); + } + + /// Each width's forward timed on its own cache, on its third run: graph captured, stall paid, a keep in front. + fn timeWindows(self: *Lanes, d: anytype) !void { + var cache = try E.Cache.init(self.gpa, self.e, 64); + defer cache.deinit(self.gpa); + const ids: [E.most_window]u32 = @splat(0); + var picks: [E.most_window]E.Draw = undefined; + for (&picks) |*pk| pk.* = .{ .rule = null, .position = 1 }; + for (1..self.e.window_rows + 1) |rows| { + var ms: f64 = 0; + for (0..3) |run| { + try cache.reset(self.e); + const t0 = std.Io.Clock.awake.now(self.io); + try self.e.forward(&.{.{ .cache = &cache, .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); + if (run == 2) ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(self.io)).toNanoseconds())) / 1e6; + } + d.setCost(@intCast(rows), ms); + } + } + + fn backend(self: *Lanes) be.Backend { + return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .prepare_features = prepareFeaturesFn, .features = featuresFn, .release = releaseFn } }; + } + + fn cast(ptr: *anyopaque) *Lanes { + return @ptrCast(@alignCast(ptr)); + } + + fn cacheOf(self: *Lanes, stream: *lanes.Stream) !*E.Cache { + return self.caches.get(stream) orelse error.UnknownQwenStream; + } + + fn prefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void { + const self = cast(ptr); + const ids = stream.prompt(); + if (ids.len == 0 or ids.len + stream.max_new > self.capacity or self.caches.contains(stream)) return error.InvalidQwenPrompt; + const cache = try self.take(); + errdefer self.give(cache); + var at: usize = 0; + while (at < ids.len) { + var end: usize = @min(ids.len, at + E.max_rows); + for (stream.chunks) |boundary| if (boundary > at) { + end = @min(end, boundary); + break; + }; + try self.e.forward(&.{.{ .cache = cache, .rows = end - at }}, ids[at..end], .last); + at = end; + } + try self.caches.put(self.gpa, stream, cache); + } + + fn firstFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64) !u64 { + return firstMaskedFn(ptr, stream, position, &.{}); + } + + fn firstMaskedFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64, mask: []const u32) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len) return error.InvalidQwenPosition; + return try draws.draw(self.gpa, cache.logits, stream.sampling, position, mask); + } + + fn queueFn(ptr: *anyopaque, stream: *lanes.Stream, feed: be.Feed, position: u64) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len + 1) return error.InvalidQwenPosition; + const token: u32 = switch (feed) { + .value => |v| v, + .handle => |h| @intCast(h), + }; + try self.e.forward(&.{.{ .cache = cache, .rows = 1 }}, &.{token}, .{ .draws = &.{pick(stream.sampling, position)} }); + cache.last = null; + return self.e.tokens[0]; + } + + fn readFn(_: *anyopaque, handle: u64) !u32 { + if (handle >= c.vocab) return error.InvalidQwenToken; + return @intCast(handle); + } + + /// Each window (a pending token and drafts) checked, then as many as fit a forward; a tree's rows keyed by depth. + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { + const self = cast(ptr); + const t0 = std.Io.Clock.awake.now(self.io); + defer @import("cuda_draft.zig").timing("verify", self.io, t0); + if (windows.len != out.len) return error.InvalidQwenWindows; + for (windows, out) |w, o| { + const rows = w.rows(); + if (w.held != 0 or rows > self.e.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; + if (w.masks.len != 0 and w.masks.len != rows * draws.mask_words) return error.InvalidQwenMasks; + if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; + const cache = try self.cacheOf(w.stream); + var depth: [E.most_window]u64 = undefined; + for (w.positions, 0..) |p, r| { + const parent: i32 = if (w.parents) |ps| ps[r] else if (r == 0) -1 else @intCast(r - 1); + if (r == 0 and parent != -1) return error.UnsupportedQwenWindow; + if (r > 0 and (parent < 0 or parent >= @as(i32, @intCast(r)))) return error.UnsupportedQwenWindow; + depth[r] = if (parent < 0) 0 else depth[@intCast(parent)] + 1; + if (p != cache.len + depth[r] + 1) return error.InvalidQwenPosition; + } + } + var at: usize = 0; + while (at < windows.len) { + var parts: [E.max_segs]E.Part = undefined; + var ids: [E.max_rows]u32 = undefined; + var picks: [E.max_rows]E.Draw = undefined; + var n: usize = 0; + var total: usize = 0; + var masked = false; + while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= self.e.verify_rows) : (n += 1) { + const w = windows[at + n]; + parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows(), .parents = w.parents }; + ids[total] = w.pending; + @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); + for (w.positions, 0..) |p, r| picks[total + r] = pick(w.stream.sampling, p); + masked = masked or w.masks.len != 0; + total += w.rows(); + } + // structured output draws on the host, over each row's logits with its mask applied + try self.e.forward(parts[0..n], ids[0..total], if (masked) .logits else .{ .draws = picks[0..total] }); + var row: usize = 0; + for (windows[at..][0..n], out[at..][0..n]) |w, o| { + for (o.sampled, w.positions, 0..) |*t, p, r| { + const mask = if (w.masks.len > 0) w.masks[r * draws.mask_words ..][0..draws.mask_words] else &.{}; + t.* = if (masked) try draws.draw(self.gpa, self.e.rowLogits(row + r), w.stream.sampling, p, mask) else self.e.tokens[row + r]; + } + @memcpy(o.drafts, w.tokens); + row += w.rows(); + } + at += n; + } + } + + /// The accepted path of each window (a chain's prefix or a tree's branch): the state as that chain leaves it. + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { + const self = cast(ptr); + const t0 = std.Io.Clock.awake.now(self.io); + defer @import("cuda_draft.zig").timing("keep", self.io, t0); + for (windows, paths) |w, path| { + const cache = try self.cacheOf(w.stream); + if (cache.last == null) continue; // a one-row window (or one kept already): nothing to roll back + try cache.keep(self.e, path); + } + } + + fn draftFn(_: *anyopaque, requests: []const be.DraftRequest) !void { + for (requests) |r| if (r.depth != 0) return error.QwenCheckpointHasNoDraftHead; + } + + /// A drafter's tapped layers, before any stream's caches exist: each forward fills their rings after. + fn prepareFeaturesFn(ptr: *anyopaque, taps: []const u32) !void { + const self = cast(ptr); + if (self.caches.count() != 0 or self.pool.items.len != 0) return error.TapsAfterStreams; + try self.e.setTaps(taps); + } + + /// Rows [start, start + count) of a stream's tap ring (slot position % tap_rows), the newest tap_rows rows kept. + fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { + const self = cast(ptr); + if (taps.len != self.e.n_taps) return error.TapsMismatch; + const cache = try self.cacheOf(stream); + const ring = cache.tap orelse return error.NoTaps; + // a long prompt asks for more rows than the ring keeps; the drafter reads only its window of the newest + const held_from = cache.written -| E.tap_rows; + if (start + count > cache.written or start + count <= held_from) return error.RowsNotHeld; + return .{ .buffer = ring.ptr, .offset = start, .rows = count, .row_bytes = @intCast(self.e.n_taps * self.e.g.hidden * 2), .space = .device, .dtype = .bf16 }; + } + + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { + const self = cast(ptr); + if (self.caches.fetchRemove(stream)) |entry| self.give(entry.value); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_native_moe.zig b/zig/src/families/qwen3_5/cuda_native_moe.zig new file mode 100644 index 0000000000..16d71b8d2b --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_native_moe.zig @@ -0,0 +1,14 @@ +//! The qwen3_5_moe model type (Qwen3.6-35B-A3B) on the same CUDA engine: routed experts in place of the dense MLP. +const native = @import("cuda_native.zig"); + +pub const model_type = "qwen3_5_moe"; +pub const formats = native.formats; +pub const default_context = native.default_context; +pub const max_segments = native.max_segments; +pub const prompt_rows = native.prompt_rows; +/// --drafter names the checkpoint's own directory: its MTP layer (mtp-4bit.safetensors) drafts chains. +pub const takes_drafter = true; +pub const Options = native.Options; +pub const open = native.open; +pub const openDrafter = native.openMtp; +pub const explain = native.explain; diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig new file mode 100644 index 0000000000..35af7ef55e --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -0,0 +1,253 @@ +//! The checkpoint's weights onto the device: affine-4 projections packed as qlinear tiles, out of the engine's arena. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Geometry = geometry.Geometry; +const Engine = engine.Engine; +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const aligned = engine.aligned; +const head_dim = geometry.head_dim; + +/// The packed tiles' bytes of an n-by-k affine-4 projection (qlinear's words, scales and biases). +pub fn packedBytes(n: usize, k: usize) usize { + const l = Affine4.layout(n, k); + return aligned(l.words) + 2 * aligned(l.scales); +} + +/// The checkpoint's own bytes of an n-by-k affine-4 projection: words, then its scales and biases. +fn rawBytes(n: usize, k: usize) usize { + return n * k / 8 * 4 + 2 * (n * k / 64 * 2); +} + +/// A layer's routed experts' device bytes: the pair table, the down table and the router rows. +pub fn routedBytes(experts: usize, intermediate: usize, hidden: usize) usize { + const ex = cuda.experts.Affine4Experts; + return aligned(ex.pairBytes(experts + 1, intermediate, hidden)) + aligned(ex.bytes(experts + 1, hidden, intermediate)) + aligned((experts + 1) * hidden * 2); +} + +/// What `load` takes from the arena: every packed projection and vector, each allocation aligned. +pub fn total(g: Geometry) usize { + var t: usize = 0; + t += aligned(g.vocab * g.hidden / 8 * 4) + 2 * aligned(g.vocab * g.hidden / 64 * 2); // the embedding as stored + t += packedBytes(g.vocab, g.hidden); // the head, tied or not + t += aligned(g.hidden * 2); + for (g.linear[0..g.layers]) |lin| { + t += 2 * aligned(g.hidden * 2); + if (g.moe()) t += routedBytes(g.experts, g.intermediate, g.hidden) else t += 2 * packedBytes(g.intermediate, g.hidden) + packedBytes(g.hidden, g.intermediate); + if (lin) { + t += packedBytes(g.convWidth(), g.hidden) + packedBytes(g.valueWidth(), g.hidden) + 2 * packedBytes(g.v_heads, g.hidden) + packedBytes(g.hidden, g.valueWidth()); + t += aligned(g.convWidth() * geometry.taps * 2) + aligned(geometry.dk * 2) + aligned(g.v_heads * 4) + aligned(g.v_heads * 2); + } else { + t += packedBytes(2 * g.attnWidth(), g.hidden) + 2 * packedBytes(g.kv_heads * head_dim, g.hidden) + packedBytes(g.hidden, g.attnWidth()); + t += 2 * aligned(head_dim * 2); + } + } + return t; +} + +/// The largest projection's bytes as stored: the staging buffer the load packs each one from. +pub fn stageBytes(g: Geometry) usize { + var most: usize = 0; + for ([_][2]usize{ .{ g.vocab, g.hidden }, .{ g.intermediate, g.hidden }, .{ g.hidden, g.intermediate }, .{ g.convWidth(), g.hidden }, .{ g.valueWidth(), g.hidden }, .{ 2 * g.attnWidth(), g.hidden } }) |nk| most = @max(most, rawBytes(nk[0], nk[1])); + if (g.moe()) most = @max(most, @max(rawBytes(g.experts * g.intermediate, g.hidden), rawBytes(g.experts * g.hidden, g.intermediate))); + return most; +} + +/// `into` is the engine or a drafter's weights: anything with `alloc(bytes)` out of its `arena`. +pub fn upload(into: anytype, bytes: []const u8) !u64 { + const p = try into.alloc(bytes.len); + try into.arena.upload(p - into.arena.ptr, bytes); + return p; +} + +fn put(into: anytype, t: core.checkpoint.Tensor) !u64 { + return upload(into, t.bytes); +} + +/// A vector the kernels read as f32, stored as f32 (the 2B) or bf16 (the 27B's A_log). +fn f32vec(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, n: usize) !u64 { + var buf: [160]u8 = undefined; + const t = try ck.get(try std.fmt.bufPrint(&buf, fmt, args)); + if (t.is(.f32, &.{n})) return put(e, t); + if (!t.is(.bf16, &.{n})) return error.UnexpectedTensor; + const words = try e.gpa.alloc(u32, n); + defer e.gpa.free(words); + for (words, std.mem.bytesAsSlice(u16, t.bytes)) |*w, h| w.* = @as(u32, h) << 16; + return upload(e, std.mem.sliceAsBytes(words)); +} + +pub fn tensor(into: anytype, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + var buf: [160]u8 = undefined; + return put(into, try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); +} + +/// Raw words, scales and biases staged in order; each part's device pointer. +fn stage(e: *Engine, parts: [3][]const u8) ![3]u64 { + var raw: [3]u64 = undefined; + var at: usize = 0; + for (parts, &raw) |p, *r| { + if (at + p.len > e.stage.len) return error.StageTooSmall; + try e.stage.upload(at, p); + r.* = e.stage.ptr + at; + at += p.len; + } + return raw; +} + +fn packStaged(e: *Engine, into: anytype, raw: [3]u64, n: usize, k: usize) !Affine4 { + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try into.alloc(l.words), .s = try into.alloc(l.scales), .b = try into.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + try e.affine.pack(e.stream, raw[0], raw[1], raw[2], q); + try e.stream.synchronize(); // the next projection's staging overwrites this one's + return q; +} + +/// An MLX affine-4 projection packed into qlinear's tiles from the staging buffer, which holds it as stored. +pub fn linear(e: *Engine, into: anytype, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { + var buf: [160]u8 = undefined; + var parts: [3][]const u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (specs, &parts) |p, *part| part.* = (try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] })).bytes; + return packStaged(e, into, try stage(e, parts), n, k); +} + +/// Inputs [from, from + len) of an n-by-k projection as their own packed projection (the MTP's fc halves). +pub fn linearSlice(e: *Engine, into: anytype, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize, from: usize, len: usize) !Affine4 { + var buf: [160]u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize, usize }{ .{ "weight", .u32, k / 8, 4 }, .{ "scales", .bf16, k / 64, 2 }, .{ "biases", .bf16, k / 64, 2 } }; + var parts: [3][]u8 = undefined; + var taken: usize = 0; + errdefer for (parts[0..taken]) |p| e.gpa.free(p); + for (specs, &parts) |p, *part| { + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); + const per = k / p[2]; // inputs a column holds + const cols = len / per; + part.* = try e.gpa.alloc(u8, n * cols * p[3]); + taken += 1; + for (0..n) |r| @memcpy(part.*[r * cols * p[3] ..][0 .. cols * p[3]], t.bytes[(r * p[2] + from / per) * p[3] ..][0 .. cols * p[3]]); + } + defer for (parts) |p| e.gpa.free(p); + return packStaged(e, into, try stage(e, .{ parts[0], parts[1], parts[2] }), n, len); +} + +fn bf16f(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// f32 to bf16, round to nearest even (finite values). +fn f32bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @intCast((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +/// 8-bit affine rows (four codes a word, a scale and bias a group of 64) dequantized to bf16: the router rows. +fn dequant8(dst: []u16, words: []align(1) const u32, scales: []align(1) const u16, biases: []align(1) const u16, k: usize) void { + const groups = k / 64; + for (0..dst.len / k) |r| { + for (0..k) |i| { + const code: f32 = @floatFromInt((words[r * (k / 4) + i / 4] >> @intCast(8 * (i % 4))) & 0xff); + dst[r * k + i] = f32bf16(code * bf16f(scales[r * groups + i / 64]) + bf16f(biases[r * groups + i / 64])); + } + } +} + +/// The router's rows as bf16 [experts + 1][hidden]: the routed experts' gate, then the shared expert's gate row. +fn routerRows(e: *Engine, into: anytype, ck: *core.Checkpoint, prefix: []const u8) !u64 { + const g = e.g; + const rows = try e.gpa.alloc(u16, (g.experts + 1) * g.hidden); + defer e.gpa.free(rows); + var buf: [160]u8 = undefined; + const names = [_]struct { []const u8, usize, usize }{ .{ "mlp.gate", 0, g.experts }, .{ "mlp.shared_expert_gate", g.experts, 1 } }; + for (names) |n| { + const w = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, n[0] }), .u32, &.{ n[2], g.hidden / 4 }); + const s = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.scales", .{ prefix, n[0] }), .bf16, &.{ n[2], g.hidden / 64 }); + const b = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.biases", .{ prefix, n[0] }), .bf16, &.{ n[2], g.hidden / 64 }); + dequant8(rows[n[1] * g.hidden ..][0 .. n[2] * g.hidden], std.mem.bytesAsSlice(u32, w.bytes), std.mem.bytesAsSlice(u16, s.bytes), std.mem.bytesAsSlice(u16, b.bytes), g.hidden); + } + return upload(into, std.mem.sliceAsBytes(rows)); +} + +/// `count` stacked experts' projection (one at count 0) staged as stored and packed into table `out` as the mi-th of m. +fn expertTable(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, count: usize, n: usize, k: usize, out: u64, m: usize, mi: usize) !void { + var buf: [160]u8 = undefined; + var parts: [3][]const u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (specs, &parts) |p, *part| { + const stacked: []const usize = &.{ count, n, p[2] }; + const one: []const usize = &.{ n, p[2] }; + part.* = (try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], if (count == 0) one else stacked)).bytes; + } + const raw = try stage(e, parts); + try e.experts.packInto(e.stream, raw[0], raw[1], raw[2], out, @max(count, 1), n, k, m, mi); + try e.stream.synchronize(); // the next table's staging overwrites this one's +} + +/// A layer's routed experts: gate and up interleaved as one pair table, down apart, the shared expert last in each. +pub fn routed(e: *Engine, into: anytype, ck: *core.Checkpoint, p: []const u8) !engine.Mlp { + const g = e.g; + const ex = cuda.experts.Affine4Experts; + const ni = g.intermediate; + const pair = try into.alloc(ex.pairBytes(g.experts + 1, ni, g.hidden)); + const down = try into.alloc(ex.bytes(g.experts + 1, g.hidden, ni)); + try expertTable(e, ck, p, "mlp.switch_mlp.gate_proj", g.experts, ni, g.hidden, pair, 2, 0); + try expertTable(e, ck, p, "mlp.switch_mlp.up_proj", g.experts, ni, g.hidden, pair, 2, 1); + try expertTable(e, ck, p, "mlp.switch_mlp.down_proj", g.experts, g.hidden, ni, down, 1, 0); + const shared_pair = pair + ex.pairBytes(g.experts, ni, g.hidden); + try expertTable(e, ck, p, "mlp.shared_expert.gate_proj", 0, ni, g.hidden, shared_pair, 2, 0); + try expertTable(e, ck, p, "mlp.shared_expert.up_proj", 0, ni, g.hidden, shared_pair, 2, 1); + try expertTable(e, ck, p, "mlp.shared_expert.down_proj", 0, g.hidden, ni, down + ex.bytes(g.experts, g.hidden, ni), 1, 0); + const layer: cuda.experts.Layer = .{ .format = .affine4, .up = pair, .down = down, .width = ni, .dims = g.hidden, .experts = g.experts + 1, .swiglu = true }; + return .{ .moe = .{ .router = try routerRows(e, into, ck, p), .experts = layer } }; +} + +pub fn load(e: *Engine, ck: *core.Checkpoint) !void { + const g = e.g; + const m = "language_model.model."; + e.embed_w = try tensor(e, ck, "{s}embed_tokens.weight", .{m}, .u32, &.{ g.vocab, g.hidden / 8 }); + e.embed_s = try tensor(e, ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 }); + e.embed_b = try tensor(e, ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 }); + e.head = if (g.tied) try linear(e, e, ck, m, "embed_tokens", g.vocab, g.hidden) else try linear(e, e, ck, "language_model.", "lm_head", g.vocab, g.hidden); + e.final_norm = try tensor(e, ck, "{s}norm.weight", .{m}, .bf16, &.{g.hidden}); + const conv = g.convWidth(); + const vw = g.valueWidth(); + for (e.layers, 0..) |*l, i| { + var pbuf: [96]u8 = undefined; + const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); + l.* = .{ + .input_norm = try tensor(e, ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{g.hidden}), + .post_norm = try tensor(e, ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{g.hidden}), + .mlp = if (g.moe()) try routed(e, e, ck, p) else .{ .dense = .{ + .gate = try linear(e, e, ck, p, "mlp.gate_proj", g.intermediate, g.hidden), + .up = try linear(e, e, ck, p, "mlp.up_proj", g.intermediate, g.hidden), + .down = try linear(e, e, ck, p, "mlp.down_proj", g.hidden, g.intermediate), + } }, + .mixer = undefined, + }; + if (g.linear[i]) { + l.mixer = .{ .delta = .{ + .qkv = try linear(e, e, ck, p, "linear_attn.in_proj_qkv", conv, g.hidden), + .z = try linear(e, e, ck, p, "linear_attn.in_proj_z", vw, g.hidden), + .a = try linear(e, e, ck, p, "linear_attn.in_proj_a", g.v_heads, g.hidden), + .b = try linear(e, e, ck, p, "linear_attn.in_proj_b", g.v_heads, g.hidden), + .out = try linear(e, e, ck, p, "linear_attn.out_proj", g.hidden, vw), + .conv = try tensor(e, ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ conv, geometry.taps, 1 }), + .norm = try tensor(e, ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{geometry.dk}), + .a_log = try f32vec(e, ck, "{s}linear_attn.A_log", .{p}, g.v_heads), + .dt_bias = try tensor(e, ck, "{s}linear_attn.dt_bias", .{p}, .bf16, &.{g.v_heads}), + } }; + } else { + l.mixer = .{ .attention = .{ + .q = try linear(e, e, ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden), + .k = try linear(e, e, ck, p, "self_attn.k_proj", g.kv_heads * head_dim, g.hidden), + .v = try linear(e, e, ck, p, "self_attn.v_proj", g.kv_heads * head_dim, g.hidden), + .o = try linear(e, e, ck, p, "self_attn.o_proj", g.hidden, g.attnWidth()), + .q_norm = try tensor(e, ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{head_dim}), + .k_norm = try tensor(e, ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{head_dim}), + } }; + } + } +} diff --git a/zig/src/families/qwen3_5/draw.zig b/zig/src/families/qwen3_5/draw.zig new file mode 100644 index 0000000000..c9acd623bf --- /dev/null +++ b/zig/src/families/qwen3_5/draw.zig @@ -0,0 +1,108 @@ +//! The host draw over a row of bf16 logits, shared by the Metal and CUDA backends: greedy, keyed sampling or a mask. +const std = @import("std"); +const lanes = @import("lanes"); +const c = @import("config.zig"); + +/// A row's allowed-token bits (structured output): one bit a logit. +pub const mask_words = (c.vocab + 31) / 32; + +pub fn value(word: u16) f64 { + return @as(f32, @bitCast(@as(u32, word) << 16)); +} + +/// The unchanged top-k candidates for the shared fp64 keyed sampler; a `mask` keeps its allowed ones (-inf's draw). +pub fn draw(gpa: std.mem.Allocator, logits: []const u16, settings: ?lanes.Sampling, position: u64, mask: []const u32) !u32 { + var best: ?usize = null; + var allowed: usize = 0; + for (logits, 0..) |word, id| { + const x = value(word); + if (!std.math.isFinite(x)) return error.NonfiniteQwenLogits; + if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; + allowed += 1; + if (best == null or x > value(logits[best.?])) best = id; + } + if (best == null) return error.NoAllowedQwenToken; + const sampling = settings orelse return @intCast(best.?); + if (sampling.temperature <= 0) return @intCast(best.?); + const k = if (sampling.top_k == 0) allowed else @min(allowed, sampling.top_k); + const values = try gpa.alloc(f64, k); + defer gpa.free(values); + const ids = try gpa.alloc(u64, k); + defer gpa.free(ids); + var filled: usize = 0; + for (logits, 0..) |word, id| { + if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; + const x = value(word); + if (k == allowed) { + values[filled] = x; + ids[filled] = id; + filled += 1; + continue; + } + if (filled == k and x <= values[k - 1]) continue; + var at = @min(filled, k - 1); + while (at > 0 and x > values[at - 1]) : (at -= 1) { + values[at] = values[at - 1]; + ids[at] = ids[at - 1]; + } + values[at] = x; + ids[at] = id; + filled = @min(filled + 1, k); + } + return @intCast(try lanes.sampling.choose(gpa, values, ids, position, sampling)); +} + +test "candidate selection preserves full-vocabulary keyed draws and tie order" { + const gpa = std.testing.allocator; + const words = try gpa.alloc(u16, c.vocab); + defer gpa.free(words); + const values = try gpa.alloc(f64, c.vocab); + defer gpa.free(values); + const ids = try gpa.alloc(u64, c.vocab); + defer gpa.free(ids); + for (words, values, ids, 0..) |*w, *v, *id, i| { + const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; + w.* = @intCast(@as(u32, @bitCast(x)) >> 16); + v.* = value(w.*); + id.* = i; + } + for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { + for ([_]u64{ 0, 128, 65536 }) |position| { + const s: lanes.Sampling = .{ .seed = 1234, .top_k = k, .temperature = 0.7, .top_p = 0.95, .min_p = 0.1 }; + const want = try lanes.sampling.choose(gpa, values, ids, position, s); + try std.testing.expectEqual(want, try draw(gpa, words, s, position, &.{})); + } + } + try std.testing.expectEqual(try draw(gpa, words, null, 0, &.{}), try draw(gpa, words, .{ .seed = 7, .temperature = 0 }, 99, &.{})); + words[0] = 0x7fc0; + try std.testing.expectError(error.NonfiniteQwenLogits, draw(gpa, words, null, 0, &.{})); +} + +test "a masked draw is the full-vocabulary draw with the other tokens at -inf" { + const gpa = std.testing.allocator; + const words = try gpa.alloc(u16, c.vocab); + defer gpa.free(words); + const values = try gpa.alloc(f64, c.vocab); + defer gpa.free(values); + const ids = try gpa.alloc(u64, c.vocab); + defer gpa.free(ids); + var mask: [mask_words]u32 = @splat(0); + for (words, values, ids, 0..) |*w, *v, *id, i| { + const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; + w.* = @intCast(@as(u32, @bitCast(x)) >> 16); + const on = i % 7 == 3 or i % 1000 == 1; + if (on) mask[i / 32] |= @as(u32, 1) << @intCast(i % 32); + v.* = if (on) value(w.*) else -std.math.inf(f64); + id.* = i; + } + for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { + for ([_]u64{ 0, 128, 65536 }) |position| { + const s: lanes.Sampling = .{ .seed = 99, .top_k = k, .temperature = 0.8, .top_p = 0.9, .min_p = 0.05 }; + const want = try lanes.sampling.choose(gpa, values, ids, position, s); + try std.testing.expectEqual(want, try draw(gpa, words, s, position, &mask)); + } + } + try std.testing.expectEqual(try lanes.sampling.choose(gpa, values, ids, 0, .{ .seed = 1, .temperature = 0 }), try draw(gpa, words, null, 0, &mask)); + @memset(&mask, 0); + try std.testing.expectError(error.NoAllowedQwenToken, draw(gpa, words, null, 0, &mask)); +} diff --git a/zig/src/hip/driver.zig b/zig/src/hip/driver.zig index 0982c96660..7fa1ef9cca 100644 --- a/zig/src/hip/driver.zig +++ b/zig/src/hip/driver.zig @@ -1,5 +1,6 @@ //! The first HIP admission boundary, with signatures from ROCm's hip_runtime_api.h; no CUDA symbol aliases or GPU kernels. const std = @import("std"); +const platform = @import("sys"); pub const Error = error{ DriverUnavailable, MissingSymbol, HipFailed, InvalidDeviceCount }; pub const Api = struct { @@ -8,11 +9,11 @@ pub const Api = struct { }; pub const Driver = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, pub fn openPath(path: []const u8) Error!Driver { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); const api = try @import("symbols.zig").resolve(Api, &lib); try initialize(api); diff --git a/zig/src/hip/runtime.zig b/zig/src/hip/runtime.zig index 8f5527bd84..2583cf6bbd 100644 --- a/zig/src/hip/runtime.zig +++ b/zig/src/hip/runtime.zig @@ -1,10 +1,11 @@ //! Complete model-free HIP runtime admission; partial symbol tables never escape. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Error = error{ DriverUnavailable, MissingSymbol, HipFailed, Invalid }; pub const Runtime = struct { - lib: std.DynLib, + lib: platform.DynLib, api: abi.Api, pub fn open() Error!Runtime { @@ -12,7 +13,7 @@ pub const Runtime = struct { } pub fn openPath(path: []const u8) Error!Runtime { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); const api = try @import("symbols.zig").resolve(abi.Api, &lib); try check(api.hipInit(0)); diff --git a/zig/src/hip/symbols.zig b/zig/src/hip/symbols.zig index e334f716cb..54c2c14bc6 100644 --- a/zig/src/hip/symbols.zig +++ b/zig/src/hip/symbols.zig @@ -1,7 +1,8 @@ //! One symbol-resolution path for admission fixtures and the full runtime. const std = @import("std"); +const platform = @import("sys"); -pub fn resolve(comptime Api: type, lib: *std.DynLib) error{MissingSymbol}!Api { +pub fn resolve(comptime Api: type, lib: *platform.DynLib) error{MissingSymbol}!Api { var api: Api = undefined; const info = @typeInfo(Api).@"struct"; inline for (info.field_names, info.field_types) |name, T| { diff --git a/zig/src/learned_warning_test.zig b/zig/src/learned_warning_test.zig index cb9a9346b1..382ac2659e 100644 --- a/zig/src/learned_warning_test.zig +++ b/zig/src/learned_warning_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const pc = @import("core/prompt_cache.zig"); const Imprint = @import("core/prompt_imprint.zig").Imprint; const Fake = @import("core/prompt_cache_fake.zig").Fake; @@ -56,7 +57,7 @@ fn expectedWarnings(n: usize) !void { pub fn main() !void { const a = std.heap.page_allocator; var buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-warning-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-warning-{d}", .{platform.c.getpid()}); defer @import("core/prompt_imprint_test.zig").rmTree(root); { var im = try Imprint.open(a, root, 1, 1 << 20); diff --git a/zig/src/metal/clock.zig b/zig/src/metal/clock.zig index 6461c8e227..9e0572c597 100644 --- a/zig/src/metal/clock.zig +++ b/zig/src/metal/clock.zig @@ -1,13 +1,14 @@ //! Host time in seconds on the clock Metal's GPUStartTime and GPUEndTime use (mach_absolute_time). const std = @import("std"); +const platform = @import("sys"); var scale: f64 = 0; pub fn seconds() f64 { if (scale == 0) { - var info: std.c.mach_timebase_info_data = undefined; - _ = std.c.mach_timebase_info(&info); + var info: platform.c.mach_timebase_info_data = undefined; + _ = platform.c.mach_timebase_info(&info); scale = @as(f64, @floatFromInt(info.numer)) / @as(f64, @floatFromInt(info.denom)) / 1e9; } - return @as(f64, @floatFromInt(std.c.mach_absolute_time())) * scale; + return @as(f64, @floatFromInt(platform.c.mach_absolute_time())) * scale; } diff --git a/zig/src/metal/device.zig b/zig/src/metal/device.zig index 69d67e108c..353cd6d8ce 100644 --- a/zig/src/metal/device.zig +++ b/zig/src/metal/device.zig @@ -1,5 +1,6 @@ //! The GPU device: queues, buffers (zero-copy over mmap too), events and residency sets. const std = @import("std"); +const platform = @import("sys"); const objc = @import("objc.zig"); const types = @import("types.zig"); const command = @import("command.zig"); @@ -139,10 +140,10 @@ pub const MappedFile = struct { size: usize, pub fn open(path: [*:0]const u8) !MappedFile { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end <= 0) return error.EmptyFile; const size: usize = @intCast(end); const page = std.heap.pageSize(); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..8ce4e0bebb 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,16 +1,19 @@ //! CUDA family registration, device setup, memory admission and the native lane host. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const qwen35 = @import("qwen35"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); const nemotron_slide = @import("nemotron_slide.zig"); +const drafted = @import("cuda_drafted.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{nemotron.native}; +const registry = .{ nemotron.native, qwen35.native, qwen35.native_moe }; /// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ @@ -44,7 +47,7 @@ fn chipClass(a: Allocator, capability: u32) ?[]const u8 { } fn getenv(name: [:0]const u8) ?[]const u8 { - return std.mem.span(std.c.getenv(name) orelse return null); + return std.mem.span(platform.c.getenv(name) orelse return null); } /// A variable's whole number: null when unset or empty, an error for anything else that is not one. @@ -152,6 +155,7 @@ const Host = struct { slide: ?*anyopaque = null, // that learner's adapter, freed by slide_free before the family slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into + served: drafted.Served = undefined, // the family's backend, wrapped with --drafter's draft model when given fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); @@ -165,6 +169,7 @@ const Host = struct { } if (h.slide) |s| h.slide_free.?(h.gpa, s); h.gpa.free(h.dir); + h.served.close(h.gpa); h.release(h.family); if (h.gpu) |g| { g.ctx.deinit(); @@ -188,6 +193,10 @@ const Host = struct { fn serve(h: *Host, io: std.Io, facts: lanes.Model, rows: u32, info: api.Info, explain: ?api.Explain) !void { h.cfg = try lanes.Config.init(h.gpa, facts, rows, rows - 1); errdefer h.cfg.deinit(h.gpa); + if (facts.head_trees) { // a drafter's trees: the planner picks a lone stream's lanes, as deep as it drafts + h.cfg.head_lanes = rows - 1; + h.cfg.head_depth = @max(1, facts.drafts); + } h.clock = .{ .io = io }; h.core = lanes.Engine.init(h.gpa, &h.cfg, h.backend(), h.clock.clock()); errdefer h.core.deinit(); @@ -333,27 +342,43 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return null; }; defer if (!opened) loaded.deinit(loaded.ctx); + var served = try drafted.wrap(F, a, gpa, io, loaded, o, problem) orelse return null; + defer if (!opened) served.close(gpa); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - const room = after.room(model) -| loaded.learn_bytes; + // what a family's Loaded may leave out: no learner, no lone driver, no prompt-cache copies + const L = @TypeOf(loaded); + const learn_bytes: usize = if (@hasField(L, "learn_bytes")) loaded.learn_bytes else 0; + const room = after.room(model) -| learn_bytes; const free = loaded.free_streams; - const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { + const streams = budget.admit(room, served.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { - error.NoStream => try std.fmt.allocPrint(a, "the CUDA memory budget fits no stream: one at a {d}-token window takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; lower --context, or free device memory", .{ window, toGib(loaded.stream_bytes), toGib(room), toGib(model), toGib(after.reserve) }), - error.TooMany => try std.fmt.allocPrint(a, "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, and the CUDA memory budget leaves {d:.2} GiB: serve --parallel {d}, or lower --context", .{ o.lanes, toGib(loaded.stream_bytes * (o.lanes -| free)), window, toGib(room), free + room / loaded.stream_bytes }), + error.NoStream => try std.fmt.allocPrint(a, "the CUDA memory budget fits no stream: one at a {d}-token window takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; lower --context, or free device memory", .{ window, toGib(served.stream_bytes), toGib(room), toGib(model), toGib(after.reserve) }), + error.TooMany => try std.fmt.allocPrint(a, "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, and the CUDA memory budget leaves {d:.2} GiB: serve --parallel {d}, or lower --context", .{ o.lanes, toGib(served.stream_bytes * (o.lanes -| free)), window, toGib(room), free + room / served.stream_bytes }), + }; + return null; + }; + // a family that pre-allocates its streams' buffers does so now that their count is known + if (@hasField(L, "reserve")) { + loaded.reserve(loaded.ctx, streams) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot reserve {d} streams' caches at a {d}-token window ({s}): serve fewer with --parallel, or lower --context", .{ streams, window, @errorName(e) }); + return null; }; + } + served.reserve(streams) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot reserve {d} streams' caches at a {d}-token window ({s}): serve fewer with --parallel, or lower --context", .{ streams, window, @errorName(e) }); return null; }; const h = try gpa.create(Host); errdefer gpa.destroy(h); - h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); + h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = served.backend, .served = served, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; + const cache = if (!@hasField(L, "target")) 0 else budget.cacheBytes(o.prompt_cache_gib, room, served.stream_bytes * (streams -| free)); h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(served.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (@hasField(L, "target") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet @@ -364,7 +389,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { + if (@hasField(L, "learns") and loaded.learns != null) { const Slide = nemotron_slide.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); @@ -384,7 +409,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, served.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = served.backend.vtable.first_row != null, .prompt_logprobs = served.backend.vtable.prompt_rows, .structures = served.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } diff --git a/zig/src/native/cuda_drafted.zig b/zig/src/native/cuda_drafted.zig new file mode 100644 index 0000000000..d259af0993 --- /dev/null +++ b/zig/src/native/cuda_drafted.zig @@ -0,0 +1,57 @@ +//! --drafter on the CUDA host: a family's draft model behind the lane core's Drafted wrapper over its backend. +const std = @import("std"); +const api = @import("engine_api"); +const lanes = @import("lanes"); +const Allocator = std.mem.Allocator; + +/// What the host serves: the family's backend as is, or wrapped with a drafter in front. +pub const Served = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + stream_bytes: usize, + held: ?*Held = null, + + /// The drafter's own per-stream buffers for `streams`, once the server knows the count. + pub fn reserve(s: *Served, streams: u32) !void { + if (s.held) |h| try h.reserve(h.ctx, streams); + } + + pub fn close(s: *Served, gpa: Allocator) void { + if (s.held) |h| { + h.drafted.deinit(); + h.close(h.ctx); + gpa.destroy(h); + } + s.held = null; + } +}; + +const Held = struct { drafted: lanes.drafted.Drafted, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void }; + +/// `loaded`'s backend, wrapped with the drafter `o.drafter` names when it does; null with `problem` set on refusal. +pub fn wrap(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, loaded: anytype, o: api.Open, problem: *[]const u8) !?Served { + var s: Served = .{ .backend = loaded.backend, .facts = loaded.facts, .stream_bytes = loaded.stream_bytes }; + const dir = o.drafter orelse return s; + if (!o.drafts) return s; + if (!@hasDecl(F, "openDrafter")) { + problem.* = try std.fmt.allocPrint(a, "--drafter: {s} checkpoints load no draft model on CUDA", .{F.model_type}); + return null; + } + const opened = F.openDrafter(gpa, io, loaded.ctx, dir, o.drafter_bits) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load the drafter {s} ({s})", .{ dir, @errorName(e) }); + return null; + }; + const h = try gpa.create(Held); + errdefer gpa.destroy(h); + h.* = .{ .drafted = undefined, .ctx = opened.ctx, .close = opened.close, .reserve = opened.reserve }; + h.drafted = lanes.drafted.Drafted.init(gpa, loaded.backend, opened.drafter) catch |e| { + opened.close(opened.ctx); + problem.* = try std.fmt.allocPrint(a, "the drafter {s} cannot attach to this model ({s})", .{ dir, @errorName(e) }); + return null; + }; + s.backend = h.drafted.backend(); + s.facts = h.drafted.facts(loaded.facts); + s.stream_bytes = h.drafted.streamBytes(loaded.stream_bytes); + s.held = h; + return s; +} diff --git a/zig/src/native/flashnext_host.zig b/zig/src/native/flashnext_host.zig index cdba9b907b..539964a682 100644 --- a/zig/src/native/flashnext_host.zig +++ b/zig/src/native/flashnext_host.zig @@ -1,6 +1,7 @@ //! Flash Next behind the native server: one reply at a time on the replay engine (prompt chunks, then GPU-side //! rounds), greedy only; the engine's thread owns the GPU, requests wait in arrival order (foreground first). const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -420,7 +421,7 @@ const Snaps = struct { /// --learn: learned state `key`'s file removed. fn forget(_: *anyopaque, dir: [:0]const u8, key: u64) void { var buf: [1100]u8 = undefined; - _ = std.c.unlink(learned.path(&buf, dir, key) catch return); + _ = platform.c.unlink(learned.path(&buf, dir, key) catch return); } }; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 39fe531408..8370f45ed7 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -1,5 +1,6 @@ //! The Metal engines: lane core (Nemotron, Qwen3.5-2B, GLM, Qwen3.8-27B), Flash Next replay. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -260,9 +261,9 @@ fn shardStamps(h: *std.hash.Wyhash, a: Allocator, io: std.Io, dir: []const u8) ! /// A safetensors file's header (its tensors' names, shapes and offsets), in `a`. fn safetensorsHeader(a: Allocator, dir: []const u8, name: []const u8) ![]const u8 { const path = try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0); - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.FileNotFound; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var n: u64 = 0; if (!api.prompt_imprint.readAt(fd, std.mem.asBytes(&n), 0) or n > 1 << 24) return error.BadHeader; const head = try a.alloc(u8, @intCast(n)); @@ -305,7 +306,7 @@ fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ /// Flash Next uses the dump named by TF_FLASHNEXT_DUMP, or the checked-in kernels when that variable is unset. fn openFlashNext(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { - const dump: ?[]const u8 = if (std.c.getenv("TF_FLASHNEXT_DUMP")) |d| std.mem.span(d) else null; + const dump: ?[]const u8 = if (platform.c.getenv("TF_FLASHNEXT_DUMP")) |d| std.mem.span(d) else null; const native = modelContext(a, io, o.dir); const window: i64 = o.context orelse native; if (window < 0 or (native > 0 and window > native)) { @@ -373,7 +374,7 @@ fn openGlmDsa(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ } var ram: u64 = 0; var len: usize = 8; - if (std.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; + if (platform.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; const pool = mtl.objc.Pool.push(); defer pool.pop(); const h = glm_dsa.open(gpa, io, o.dir, @intCast(window), settings, @max(1, @min(o.lanes, 16)), ram / 100 * 66, o.prompt_cache_gib) catch |e| { @@ -469,10 +470,10 @@ test "a shard rewritten in place, same size and bytes, gives the learned states try std.testing.expectEqual(first.final(), notes.final()); std.Io.sleep(io, .fromMilliseconds(5), .awake) catch {}; const path = try std.fmt.allocPrintSentinel(a, "{s}/model-00001-of-00001.safetensors", .{dir}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY }, @as(platform.c.mode_t, 0)); try std.testing.expect(fd >= 0); - try std.testing.expectEqual(@as(isize, 7), std.c.pwrite(fd, "weights", 7, 0)); // as Sliding Weights writes a tensor of the same size - _ = std.c.close(fd); + try std.testing.expectEqual(@as(isize, 7), platform.c.pwrite(fd, "weights", 7, 0)); // as Sliding Weights writes a tensor of the same size + _ = platform.c.close(fd); var moved = std.hash.Wyhash.init(0); try shardStamps(&moved, a, io, dir); try std.testing.expect(first.final() != moved.final()); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index e09a508a8c..e7da14ef17 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -1,5 +1,6 @@ //! The 27B's lane core backend: a stream a runner slot, shared rounds across slots, DFlash2 trees for slot 0. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -125,7 +126,7 @@ pub const Metal = struct { const drafting = b.draft != null; const shared = b.runner.slots > 1; // a lone stream drafts DFlash2 trees, as in 1.0.5: copies land less often there (TF_QWEN27_COPIES=1 opts in) - const lone_copies = if (std.c.getenv("TF_QWEN27_COPIES")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; + const lone_copies = if (platform.c.getenv("TF_QWEN27_COPIES")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; return .{ .exact_width = if (drafting or shared) max_rows else 1, .mtp = drafting, .speculate = drafting, .speculate_early = false, .drafts = b.nodes, .draft_probabilities = true, .window_costs = b.costs[0..b.timed], .hidden_rows = shared, .streams_exact = true, .batch_rows = if (shared) batch_rows else max_rows, .max_streams = b.runner.slots, .lone_copies = lone_copies }; } diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 09da59e53f..514cecdc8d 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -1,5 +1,6 @@ //! Qwen3.8-27B on the lane host: a stream a runner slot in shared rounds, DFlash2 drafts for slot 0's stream. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const mtl = @import("metal"); const api = @import("engine_api"); @@ -289,7 +290,7 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df h.host.decoded_rows = true; // a reply's rows prefill with the bits decoding gave them, with the cache on or off // a lone greedy request takes the serial loop on tensor units, where it is faster (TF_QWEN27_LANES overrides) const units = model.device.tensorUnits(); - const serial = if (std.c.getenv("TF_QWEN27_LANES")) |v| !std.mem.eql(u8, std.mem.span(v), "1") else units; + const serial = if (platform.c.getenv("TF_QWEN27_LANES")) |v| !std.mem.eql(u8, std.mem.span(v), "1") else units; if (h.draft != null and serial) h.host.lone = .{ .ctx = h, .run = Host.lone }; h.warm = .{ .queue = model.queue }; // the ready slots are wired before the cache's budget reads the footprint, and the later ones count as taken diff --git a/zig/src/platform/sys.zig b/zig/src/platform/sys.zig new file mode 100644 index 0000000000..204eeb3d4a --- /dev/null +++ b/zig/src/platform/sys.zig @@ -0,0 +1,51 @@ +//! The OS as TensorFold calls it, imported as `sys`: `c` is std.c or win_c.zig, `DynLib` is std.DynLib or LoadLibraryW. + +const std = @import("std"); +const builtin = @import("builtin"); + +pub const windows = builtin.os.tag == .windows; + +pub const c = if (windows) @import("win_c.zig") else std.c; + +/// std.posix, or its sockets, signals and limits over Winsock and console events. +pub const posix = if (windows) @import("win_posix.zig") else std.posix; + +pub const DynLib = if (windows) WinDynLib else std.DynLib; + +/// An environment variable's value, through libc (std.process.Environ.getPosix does not build on Windows). +pub fn getenv(name: [*:0]const u8) ?[:0]const u8 { + return std.mem.span(c.getenv(name) orelse return null); +} + +/// The CUDA driver's file name: nvcuda.dll comes with the display driver on Windows. +pub const cuda_driver = if (windows) "nvcuda.dll" else "libcuda.so.1"; + +extern "kernel32" fn LoadLibraryW(name: [*:0]const u16) callconv(.winapi) ?*anyopaque; +extern "kernel32" fn GetProcAddress(module: *anyopaque, name: [*:0]const u8) callconv(.winapi) ?*anyopaque; +extern "kernel32" fn FreeLibrary(module: *anyopaque) callconv(.winapi) c_int; + +const WinDynLib = struct { + module: *anyopaque, + + pub const Error = error{ FileNotFound, NameTooLong, InvalidWtf8 }; + + pub fn open(path: []const u8) Error!WinDynLib { + var buf: [1024:0]u16 = undefined; + const n = std.unicode.wtf8ToWtf16Le(buf[0 .. buf.len - 1], path) catch return error.InvalidWtf8; + buf[n] = 0; + return .{ .module = LoadLibraryW(buf[0..n :0].ptr) orelse return error.FileNotFound }; + } + + pub fn openZ(path: [*:0]const u8) Error!WinDynLib { + return open(std.mem.span(path)); + } + + pub fn close(self: *WinDynLib) void { + _ = FreeLibrary(self.module); + } + + pub fn lookup(self: *WinDynLib, comptime T: type, name: [:0]const u8) ?T { + const p = GetProcAddress(self.module, name.ptr) orelse return null; + return @ptrCast(@alignCast(p)); + } +}; diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig new file mode 100644 index 0000000000..bfb41928c5 --- /dev/null +++ b/zig/src/platform/win_c.zig @@ -0,0 +1,563 @@ +//! libc's surface on Windows (`sys.c`): descriptors over the UCRT, CreateFileW with every share mode, -1 and errno. + +const std = @import("std"); + +pub const fd_t = c_int; +pub const mode_t = c_uint; +pub const off_t = i64; +pub const pid_t = c_int; +pub const time_t = i64; +pub const E = std.c.E; +pub const errno = std.c.errno; +pub const timespec = extern struct { sec: time_t, nsec: c_long }; + +pub const SEEK = struct { + pub const SET: c_int = 0; + pub const CUR: c_int = 1; + pub const END: c_int = 2; +}; +pub const F_OK: c_int = 0; +pub const R_OK: c_int = 4; +pub const W_OK: c_int = 2; +pub const X_OK: c_int = 0; // Windows has no execute bit; existence is the answer +pub const AT = struct { + pub const FDCWD: fd_t = -100; + pub const REMOVEDIR: u32 = 0x200; +}; +pub const DT = struct { + pub const UNKNOWN: u8 = 0; + pub const DIR: u8 = 4; + pub const REG: u8 = 8; + pub const LNK: u8 = 10; +}; +pub const S = struct { + pub const IFMT: mode_t = 0o170000; + pub const IFDIR: mode_t = 0o040000; + pub const IFREG: mode_t = 0o100000; + pub const IFLNK: mode_t = 0o120000; +}; +pub const RLIM = struct { + pub const INFINITY: u64 = std.math.maxInt(u64); +}; +pub const timeval = extern struct { sec: time_t, usec: c_long }; +pub const CLOCK = struct { + pub const REALTIME: c_int = 0; + pub const MONOTONIC: c_int = 1; +}; + +/// Linux's open flags by name; DIRECT asks for unbuffered reads (FILE_FLAG_NO_BUFFERING), the rest map or are moot. +pub const O = packed struct(u32) { + ACCMODE: enum(u2) { RDONLY = 0, WRONLY = 1, RDWR = 2 } = .RDONLY, + CREAT: bool = false, + EXCL: bool = false, + NOCTTY: bool = false, + TRUNC: bool = false, + APPEND: bool = false, + NONBLOCK: bool = false, + DSYNC: bool = false, + ASYNC: bool = false, + DIRECT: bool = false, + LARGEFILE: bool = false, + DIRECTORY: bool = false, + NOFOLLOW: bool = false, + NOATIME: bool = false, + CLOEXEC: bool = false, + SYNC: bool = false, + PATH: bool = false, + TMPFILE: bool = false, + _: u13 = 0, +}; + +// ---- Win32 and UCRT entry points ---- + +const HANDLE = *anyopaque; +const BOOL = c_int; +const invalid_handle: isize = -1; + +const OVERLAPPED = extern struct { internal: usize = 0, internal_high: usize = 0, offset: u32, offset_high: u32, event: ?HANDLE = null }; +const FILETIME = extern struct { low: u32, high: u32 }; +const FindData = extern struct { + attributes: u32, + created: FILETIME, + accessed: FILETIME, + written: FILETIME, + size_high: u32, + size_low: u32, + reserved0: u32, + reserved1: u32, + name: [260]u16, + alternate: [14]u16, +}; + +extern "kernel32" fn CreateFileW(name: [*:0]const u16, access: u32, share: u32, security: ?*anyopaque, disposition: u32, flags: u32, template: ?HANDLE) callconv(.winapi) isize; +extern "kernel32" fn ReadFile(h: HANDLE, buf: [*]u8, len: u32, got: ?*u32, ov: ?*OVERLAPPED) callconv(.winapi) BOOL; +extern "kernel32" fn WriteFile(h: HANDLE, buf: [*]const u8, len: u32, put: ?*u32, ov: ?*OVERLAPPED) callconv(.winapi) BOOL; +extern "kernel32" fn GetLastError() callconv(.winapi) u32; +extern "kernel32" fn GetFileAttributesW(name: [*:0]const u16) callconv(.winapi) u32; +extern "kernel32" fn MoveFileExW(from: [*:0]const u16, to: [*:0]const u16, flags: u32) callconv(.winapi) BOOL; +extern "kernel32" fn FindFirstFileW(pattern: [*:0]const u16, data: *FindData) callconv(.winapi) isize; +extern "kernel32" fn FindNextFileW(h: HANDLE, data: *FindData) callconv(.winapi) BOOL; +extern "kernel32" fn FindClose(h: HANDLE) callconv(.winapi) BOOL; +extern "kernel32" fn GetDiskFreeSpaceExW(dir: [*:0]const u16, avail: ?*u64, total: ?*u64, free: ?*u64) callconv(.winapi) BOOL; +extern "kernel32" fn GetTempPathW(len: u32, buf: [*]u16) callconv(.winapi) u32; +extern "kernel32" fn GetModuleFileNameW(module: ?HANDLE, buf: [*]u16, len: u32) callconv(.winapi) u32; +extern "kernel32" fn CreateHardLinkW(new: [*:0]const u16, existing: [*:0]const u16, security: ?*anyopaque) callconv(.winapi) BOOL; +extern "kernel32" fn CreateSymbolicLinkW(link: [*:0]const u16, target: [*:0]const u16, flags: u32) callconv(.winapi) u8; +extern "kernel32" fn CloseHandle(h: HANDLE) callconv(.winapi) BOOL; + +extern "c" fn _errno() *c_int; +extern "c" fn _open_osfhandle(h: isize, flags: c_int) c_int; +extern "c" fn _get_osfhandle(fd: c_int) isize; +extern "c" fn _close(fd: c_int) c_int; +extern "c" fn _read(fd: c_int, buf: [*]u8, len: c_uint) c_int; +extern "c" fn _write(fd: c_int, buf: [*]const u8, len: c_uint) c_int; +extern "c" fn _lseeki64(fd: c_int, off: i64, whence: c_int) i64; +extern "c" fn _commit(fd: c_int) c_int; +extern "c" fn _chsize_s(fd: c_int, size: i64) c_int; +extern "c" fn _wunlink(path: [*:0]const u16) c_int; +extern "c" fn _wrmdir(path: [*:0]const u16) c_int; +extern "c" fn _wmkdir(path: [*:0]const u16) c_int; +extern "c" fn _waccess(path: [*:0]const u16, mode: c_int) c_int; +extern "c" fn _isatty(fd: c_int) c_int; +extern "c" fn _fstat64(fd: c_int, st: *Stat64) c_int; +const ucrt = struct { + extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; +}; + +/// libc's getenv, except that HOME, which Windows does not set, falls back to USERPROFILE. +pub fn getenv(name: [*:0]const u8) ?[*:0]u8 { + if (ucrt.getenv(name)) |v| return v; + if (std.mem.eql(u8, std.mem.span(name), "HOME")) return ucrt.getenv("USERPROFILE"); + return null; +} +pub extern "c" fn fopen(path: [*:0]const u8, mode: [*:0]const u8) ?*std.c.FILE; +pub extern "c" fn fwrite(ptr: [*]const u8, size: usize, n: usize, f: *std.c.FILE) usize; +pub extern "c" fn fclose(f: *std.c.FILE) c_int; +pub extern "c" fn free(ptr: ?*anyopaque) void; +pub extern "c" fn malloc(n: usize) ?*anyopaque; + +/// The UCRT's struct _stat64. +const Stat64 = extern struct { dev: u32, ino: u16, mode: u16, nlink: i16, uid: i16, gid: i16, rdev: u32, size: i64, atime: i64, mtime: i64, ctime: i64 }; + +pub fn fail(e: E) c_int { + _errno().* = @backingInt(e); + return -1; +} + +/// The errno a failed Win32 call means. +fn lastError() E { + return switch (GetLastError()) { + 2, 3, 15, 123, 161 => .NOENT, // file, path, drive not found; bad name + 5, 32, 33 => .ACCES, // access denied, sharing and lock violations + 80, 183 => .EXIST, + 112 => .NOSPC, + 6 => .BADF, + 145 => .NOTEMPTY, + 267 => .NOTDIR, + 87 => .INVAL, + else => .IO, + }; +} + +// ---- paths ---- + +const Wide = [32767:0]u16; + +/// UTF-8 path as a wide path in `buf`, null with errno if it does not fit; /tmp and /proc/self/exe map as in Cygwin. +fn wide(buf: *Wide, path: []const u8) ?[*:0]const u16 { + var at: usize = 0; + var rest = path; + if (std.mem.eql(u8, path, "/proc/self/exe")) { + const n = GetModuleFileNameW(null, buf, buf.len - 1); + if (n == 0 or n >= buf.len - 1) { + _ = fail(.NAMETOOLONG); + return null; + } + buf[n] = 0; + return buf[0..n :0].ptr; + } + if (std.mem.startsWith(u8, path, "/tmp") and (path.len == 4 or path[4] == '/')) { + at = GetTempPathW(buf.len - 1, buf); // ends in a separator + if (at == 0 or at >= buf.len - 1) { + _ = fail(.NAMETOOLONG); + return null; + } + if (path.len == 4) at -= 1; + rest = if (path.len > 5) path[5..] else ""; + } + const n = std.unicode.wtf8ToWtf16Le(buf[at .. buf.len - 1], rest) catch { + _ = fail(.NAMETOOLONG); + return null; + }; + buf[at + n] = 0; + return buf[0 .. at + n :0].ptr; +} + +/// Directory descriptors: open(O.DIRECTORY) remembers the path; openat and unlinkat join onto it. +const dir_base: fd_t = 0x4000_0000; +var dirs_lock: std.atomic.Mutex = .unlocked; +var dir_paths: [256]?[:0]u8 = @splat(null); + +fn lockDirs() void { + while (!dirs_lock.tryLock()) std.atomic.spinLoopHint(); +} + +fn dirPath(fd: fd_t) ?[:0]const u8 { + if (fd < dir_base or fd >= dir_base + dir_paths.len) return null; + lockDirs(); + defer dirs_lock.unlock(); + return dir_paths[@intCast(fd - dir_base)]; +} + +fn openDir(path: []const u8) c_int { + var buf: Wide = undefined; + const w = wide(&buf, path) orelse return -1; + const attrs = GetFileAttributesW(w); + if (attrs == 0xFFFF_FFFF) return fail(lastError()); + if (attrs & 0x10 == 0) return fail(.NOTDIR); + const copy = std.heap.c_allocator.dupeSentinel(u8, path, 0) catch return fail(.NOMEM); + lockDirs(); + defer dirs_lock.unlock(); + for (&dir_paths, 0..) |*slot, i| if (slot.* == null) { + slot.* = copy; + return dir_base + @as(fd_t, @intCast(i)); + }; + std.heap.c_allocator.free(copy); + return fail(.MFILE); +} + +/// `name` under the directory `fd` (or the working directory, or absolute) in `buf`. +fn joinAt(buf: []u8, fd: fd_t, name: [*:0]const u8) ?[:0]const u8 { + const leaf = std.mem.span(name); + if (fd == AT.FDCWD or std.fs.path.isAbsolute(leaf)) return std.fmt.bufPrintSentinel(buf, "{s}", .{leaf}, 0) catch null; + const base = dirPath(fd) orelse { + _ = fail(.BADF); + return null; + }; + return std.fmt.bufPrintSentinel(buf, "{s}\\{s}", .{ base, leaf }, 0) catch { + _ = fail(.NAMETOOLONG); + return null; + }; +} + +// ---- descriptors ---- + +/// Variadic as libc's is, so callers pass a mode or not; the mode is moot on NTFS without ACL work. +pub fn open(path: [*:0]const u8, flags: O, ...) callconv(.c) c_int { + return openFile(path, flags); +} + +fn openFile(path: [*:0]const u8, flags: O) c_int { + const p = std.mem.span(path); + if (flags.DIRECTORY) return openDir(p); + var buf: Wide = undefined; + const w = wide(&buf, p) orelse return -1; + const generic_read: u32 = 0x8000_0000; + const generic_write: u32 = 0x4000_0000; + const rights: u32 = switch (flags.ACCMODE) { + .RDONLY => generic_read, + .WRONLY => generic_write, + .RDWR => generic_read | generic_write, + }; + const disposition: u32 = if (flags.CREAT and flags.EXCL) 1 // CREATE_NEW + else if (flags.CREAT and flags.TRUNC) 2 // CREATE_ALWAYS + else if (flags.CREAT) 4 // OPEN_ALWAYS + else if (flags.TRUNC) 5 // TRUNCATE_EXISTING + else 3; // OPEN_EXISTING + var attrs: u32 = 0x80; // FILE_ATTRIBUTE_NORMAL + if (flags.DIRECT) attrs |= 0x2000_0000; // FILE_FLAG_NO_BUFFERING + const h = CreateFileW(w, rights, 0x7, null, disposition, attrs, null); // share read, write and delete + if (h == invalid_handle) { + // a directory opened without O.DIRECTORY, as POSIX allows for reading its entries + if (GetLastError() == 5 and flags.ACCMODE == .RDONLY and GetFileAttributesW(w) & 0x10 != 0 and GetFileAttributesW(w) != 0xFFFF_FFFF) return openDir(p); + return fail(lastError()); + } + var crt: c_int = 0x8000; // _O_BINARY + crt |= switch (flags.ACCMODE) { + .RDONLY => 0, + .WRONLY => 1, + .RDWR => 2, + }; + if (flags.APPEND) crt |= 0x8; + const fd = _open_osfhandle(h, crt); + if (fd < 0) { + _ = CloseHandle(@ptrFromInt(@as(usize, @bitCast(h)))); + return fail(.MFILE); + } + if (fd < append_fds.len) @atomicStore(u8, &append_fds[@intCast(fd)], @intFromBool(flags.APPEND), .release); + return fd; +} + +pub fn openat(dir: fd_t, path: [*:0]const u8, flags: O, ...) callconv(.c) c_int { + var buf: [4096]u8 = undefined; + const p = joinAt(&buf, dir, path) orelse return -1; + return openFile(p.ptr, flags); +} + +/// Descriptors opened with O_APPEND, whose writes go to the end of the file whole, as POSIX appends do. +var append_fds: [8192]u8 = @splat(0); + +pub fn close(fd: fd_t) c_int { + if (fd >= 0 and fd < append_fds.len) @atomicStore(u8, &append_fds[@intCast(fd)], 0, .release); + if (fd >= dir_base and fd < dir_base + dir_paths.len) { + lockDirs(); + defer dirs_lock.unlock(); + const slot = &dir_paths[@intCast(fd - dir_base)]; + const p = slot.* orelse return fail(.BADF); + std.heap.c_allocator.free(p); + slot.* = null; + return 0; + } + return _close(fd); +} + +const chunk: usize = 1 << 30; + +pub fn read(fd: fd_t, buf: [*]u8, len: usize) isize { + return _read(fd, buf, @intCast(@min(len, chunk))); +} + +pub fn write(fd: fd_t, buf: [*]const u8, len: usize) isize { + if (fd >= 0 and fd < append_fds.len and @atomicLoad(u8, &append_fds[@intCast(fd)], .acquire) != 0) { + // WriteFile at the end-of-file offset lands whole; the CRT's seek-then-write lets writers interleave + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = 0xFFFF_FFFF, .offset_high = 0xFFFF_FFFF }; + var put: u32 = 0; + if (WriteFile(h, buf, @intCast(@min(len, chunk)), &put, &ov) == 0) return fail(lastError()); + return put; + } + return _write(fd, buf, @intCast(@min(len, chunk))); +} + +fn osHandle(fd: fd_t) ?HANDLE { + const h = _get_osfhandle(fd); + if (h == invalid_handle or h == -2) { + _ = fail(.BADF); + return null; + } + return @ptrFromInt(@as(usize, @bitCast(h))); +} + +/// Positional read by an OVERLAPPED offset; end of file reads 0 as pread does. +pub fn pread(fd: fd_t, buf: [*]u8, len: usize, offset: off_t) isize { + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = @truncate(@as(u64, @bitCast(offset))), .offset_high = @truncate(@as(u64, @bitCast(offset)) >> 32) }; + var got: u32 = 0; + if (ReadFile(h, buf, @intCast(@min(len, chunk)), &got, &ov) == 0) { + if (GetLastError() == 38) return 0; // ERROR_HANDLE_EOF + return fail(lastError()); + } + return got; +} + +pub fn pwrite(fd: fd_t, buf: [*]const u8, len: usize, offset: off_t) isize { + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = @truncate(@as(u64, @bitCast(offset))), .offset_high = @truncate(@as(u64, @bitCast(offset)) >> 32) }; + var put: u32 = 0; + if (WriteFile(h, buf, @intCast(@min(len, chunk)), &put, &ov) == 0) return fail(lastError()); + return put; +} + +pub fn lseek(fd: fd_t, offset: off_t, whence: c_int) off_t { + return _lseeki64(fd, offset, whence); +} + +pub fn fsync(fd: fd_t) c_int { + if (dirPath(fd) != null) return 0; // directory entries are durable once NTFS has them + return _commit(fd); +} + +pub fn ftruncate(fd: fd_t, size: off_t) c_int { + const rc = _chsize_s(fd, size); + if (rc != 0) return fail(@fromBackingInt(@intCast(rc))); + return 0; +} + +pub fn isatty(fd: fd_t) c_int { + return _isatty(fd); +} + +pub const Stat = struct { + dev: u64, + ino: u64, + mode: mode_t, + size: i64, + mtim: timespec, + + pub fn mtime(s: Stat) timespec { + return s.mtim; + } +}; + +pub fn fstat(fd: fd_t, out: *Stat) c_int { + if (dirPath(fd) != null) { + out.* = .{ .dev = 0, .ino = 0, .mode = S.IFDIR | 0o700, .size = 0, .mtim = .{ .sec = 0, .nsec = 0 } }; + return 0; + } + var st: Stat64 = undefined; + if (_fstat64(fd, &st) != 0) return -1; + out.* = .{ .dev = st.dev, .ino = st.ino, .mode = st.mode, .size = st.size, .mtim = .{ .sec = st.mtime, .nsec = 0 } }; + return 0; +} + +// ---- names ---- + +pub fn unlink(path: [*:0]const u8) c_int { + var buf: Wide = undefined; + return _wunlink(wide(&buf, std.mem.span(path)) orelse return -1); +} + +pub fn rmdir(path: [*:0]const u8) c_int { + var buf: Wide = undefined; + return _wrmdir(wide(&buf, std.mem.span(path)) orelse return -1); +} + +pub fn unlinkat(dir: fd_t, path: [*:0]const u8, flags: u32) c_int { + var buf: [4096]u8 = undefined; + const p = joinAt(&buf, dir, path) orelse return -1; + return if (flags & AT.REMOVEDIR != 0) rmdir(p.ptr) else unlink(p.ptr); +} + +pub fn mkdir(path: [*:0]const u8, mode: mode_t) c_int { + _ = mode; + var buf: Wide = undefined; + return _wmkdir(wide(&buf, std.mem.span(path)) orelse return -1); +} + +/// POSIX rename: replaces an existing target. +pub fn rename(from: [*:0]const u8, to: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wa = wide(&a, std.mem.span(from)) orelse return -1; + const wb = wide(&b, std.mem.span(to)) orelse return -1; + if (MoveFileExW(wa, wb, 0x1 | 0x8) == 0) return fail(lastError()); // REPLACE_EXISTING | WRITE_THROUGH + return 0; +} + +/// POSIX link: `new` names the same file as `old` (same volume). +pub fn link(old: [*:0]const u8, new: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wa = wide(&a, std.mem.span(old)) orelse return -1; + const wb = wide(&b, std.mem.span(new)) orelse return -1; + if (CreateHardLinkW(wb, wa, null) == 0) return fail(lastError()); + return 0; +} + +pub fn access(path: [*:0]const u8, mode: c_int) c_int { + var buf: Wide = undefined; + return _waccess(wide(&buf, std.mem.span(path)) orelse return -1, mode & 6); +} + +// ---- directories ---- + +pub const dirent = extern struct { + type: u8, + name: [1024]u8, +}; + +pub const DIR = struct { + find: ?HANDLE, + data: FindData, + first: bool, + fd: ?fd_t, + ent: dirent, +}; + +pub fn opendir(path: [*:0]const u8) ?*DIR { + var pattern: [4096]u8 = undefined; + const p = std.fmt.bufPrint(&pattern, "{s}\\*", .{std.mem.span(path)}) catch { + _ = fail(.NAMETOOLONG); + return null; + }; + var buf: Wide = undefined; + const w = wide(&buf, p) orelse return null; + const d = std.heap.c_allocator.create(DIR) catch { + _ = fail(.NOMEM); + return null; + }; + const h = FindFirstFileW(w, &d.data); + if (h == invalid_handle) { + const e = lastError(); + std.heap.c_allocator.destroy(d); + _ = fail(e); + return null; + } + d.* = .{ .find = @ptrFromInt(@as(usize, @bitCast(h))), .data = d.data, .first = true, .fd = null, .ent = undefined }; + return d; +} + +/// The stream owns `fd` from here: closedir releases it. +pub fn fdopendir(fd: fd_t) ?*DIR { + const path = dirPath(fd) orelse { + _ = fail(.BADF); + return null; + }; + const d = opendir(path.ptr) orelse return null; + d.fd = fd; + return d; +} + +pub fn readdir(d: *DIR) ?*dirent { + if (d.first) { + d.first = false; + } else if (FindNextFileW(d.find.?, &d.data) == 0) { + return null; + } + const name = std.mem.sliceTo(&d.data.name, 0); + const n = std.unicode.wtf16LeToWtf8(d.ent.name[0 .. d.ent.name.len - 1], name); + d.ent.name[n] = 0; + d.ent.type = if (d.data.attributes & 0x400 != 0) DT.LNK else if (d.data.attributes & 0x10 != 0) DT.DIR else DT.REG; + return &d.ent; +} + +pub fn closedir(d: *DIR) c_int { + if (d.find) |h| _ = FindClose(h); + if (d.fd) |fd| _ = close(fd); + std.heap.c_allocator.destroy(d); + return 0; +} + +// ---- process, environment, time: win_clock.zig ---- + +const clock = @import("win_clock.zig"); +pub const getpid = clock.getpid; +pub const getuid = clock.getuid; +pub const setenv = clock.setenv; +pub const unsetenv = clock.unsetenv; +pub const nanosleep = clock.nanosleep; +pub const clock_gettime = clock.clock_gettime; +pub const gettimeofday = clock.gettimeofday; + +/// Bytes free to this user under `dir` (statvfs's bavail * frsize). +pub fn diskAvailable(dir: [*:0]const u8) ?u64 { + var buf: Wide = undefined; + const w = wide(&buf, std.mem.span(dir)) orelse return null; + var avail: u64 = 0; + if (GetDiskFreeSpaceExW(w, &avail, null, null) == 0) return null; + return avail; +} + +/// NTFS keeps access in ACLs; mode bits are not stored, so a mode change succeeds and changes nothing. +pub fn fchmodat(dir: fd_t, path: [*:0]const u8, mode: mode_t, flags: u32) c_int { + _ = .{ dir, path, mode, flags }; + return 0; +} + +/// Symbolic links need Developer Mode or elevation on Windows; this asks for an unprivileged one. +pub fn symlink(target: [*:0]const u8, path: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wt = wide(&a, std.mem.span(target)) orelse return -1; + const wp = wide(&b, std.mem.span(path)) orelse return -1; + const attrs = GetFileAttributesW(wt); + const dir: u32 = if (attrs != 0xFFFF_FFFF and attrs & 0x10 != 0) 0x1 else 0; + if (CreateSymbolicLinkW(wp, wt, dir | 0x2) == 0) return fail(lastError()); // ALLOW_UNPRIVILEGED_CREATE + return 0; +} + +/// Reading a link's target needs its reparse data; nothing on Windows asks for it yet. +pub fn readlink(path: [*:0]const u8, buf: [*]u8, len: usize) isize { + _ = .{ path, buf, len }; + return fail(.NOSYS); +} diff --git a/zig/src/platform/win_clock.zig b/zig/src/platform/win_clock.zig new file mode 100644 index 0000000000..0c0c67ab4e --- /dev/null +++ b/zig/src/platform/win_clock.zig @@ -0,0 +1,81 @@ +//! Process, environment and clock calls of the Windows libc surface, re-exported by win_c.zig: UCRT and kernel timers. + +const std = @import("std"); +const c = @import("win_c.zig"); + +const HANDLE = *anyopaque; +const BOOL = c_int; +const FILETIME = extern struct { low: u32, high: u32 }; + +extern "kernel32" fn QueryPerformanceCounter(count: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn QueryPerformanceFrequency(freq: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn GetSystemTimePreciseAsFileTime(t: *FILETIME) callconv(.winapi) void; +extern "kernel32" fn CreateWaitableTimerExW(security: ?*anyopaque, name: ?[*:0]const u16, flags: u32, access: u32) callconv(.winapi) ?HANDLE; +extern "kernel32" fn SetWaitableTimer(h: HANDLE, due: *const i64, period: i32, routine: ?*anyopaque, arg: ?*anyopaque, resume_: BOOL) callconv(.winapi) BOOL; +extern "kernel32" fn WaitForSingleObject(h: HANDLE, ms: u32) callconv(.winapi) u32; +extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; +extern "c" fn _getpid() c_int; +extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; + +pub fn getpid() c.pid_t { + return _getpid(); +} + +pub fn getuid() u32 { + return 0; +} + +pub fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int { + if (overwrite == 0 and c.getenv(name) != null) return 0; + if (_putenv_s(name, value) != 0) return c.fail(.INVAL); + return 0; +} + +pub fn unsetenv(name: [*:0]const u8) c_int { + if (_putenv_s(name, "") != 0) return c.fail(.INVAL); + return 0; +} + +threadlocal var sleep_timer: ?HANDLE = null; + +/// High-resolution waitable timer per thread, so microsecond waits are not rounded to the 15.6 ms tick. +pub fn nanosleep(req: *const c.timespec, rem: ?*c.timespec) c_int { + if (rem) |r| r.* = .{ .sec = 0, .nsec = 0 }; + const ns: i64 = req.sec * std.time.ns_per_s + req.nsec; + if (ns <= 0) return 0; + if (sleep_timer == null) sleep_timer = CreateWaitableTimerExW(null, null, 0x2, 0x1F0003); // HIGH_RESOLUTION, TIMER_ALL_ACCESS + if (sleep_timer) |t| { + const due: i64 = -@divFloor(ns + 99, 100); // relative, in 100 ns units + if (SetWaitableTimer(t, &due, 0, null, null, 0) != 0) { + _ = WaitForSingleObject(t, 0xFFFF_FFFF); + return 0; + } + } + Sleep(@intCast(@max(1, @divFloor(ns + std.time.ns_per_ms - 1, std.time.ns_per_ms)))); + return 0; +} + +pub fn clock_gettime(clock: c_int, out: *c.timespec) c_int { + if (clock == c.CLOCK.MONOTONIC) { + var count: i64 = 0; + var freq: i64 = 0; + _ = QueryPerformanceCounter(&count); + _ = QueryPerformanceFrequency(&freq); + out.* = .{ .sec = @divFloor(count, freq), .nsec = @intCast(@divFloor(@mod(count, freq) * std.time.ns_per_s, freq)) }; + return 0; + } + var ft: FILETIME = undefined; + GetSystemTimePreciseAsFileTime(&ft); + const since_1601: i64 = @intCast((@as(u64, ft.high) << 32) | ft.low); + const unix = since_1601 - 116_444_736_000_000_000; + out.* = .{ .sec = @divFloor(unix, 10_000_000), .nsec = @intCast(@mod(unix, 10_000_000) * 100) }; + return 0; +} + +pub fn gettimeofday(tv: *c.timeval, tz: ?*anyopaque) c_int { + _ = tz; + var ts: c.timespec = undefined; + _ = clock_gettime(c.CLOCK.REALTIME, &ts); + tv.* = .{ .sec = ts.sec, .usec = @intCast(@divFloor(ts.nsec, 1000)) }; + return 0; +} diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig new file mode 100644 index 0000000000..4088f3214b --- /dev/null +++ b/zig/src/platform/win_posix.zig @@ -0,0 +1,308 @@ +//! std.posix as the server calls it, on Windows: Winsock sockets, WSAPoll, console events as SIGINT and SIGTERM. + +const std = @import("std"); +const c = @import("win_c.zig"); + +pub const E = c.E; +pub const errno = c.errno; +pub const fd_t = c.fd_t; +pub const AT = c.AT; +pub const S = c.S; + +// ---- sockets ---- + +pub const socket_t = isize; +pub const socklen_t = c_int; +pub const sockaddr = extern struct { + family: u16, + data: [14]u8, + + pub const in = extern struct { family: u16, port: u16, addr: u32, zero: [8]u8 = @splat(0) }; +}; +pub const timeval = c.timeval; + +pub const AF = struct { + pub const UNIX: u32 = 1; + pub const INET: u32 = 2; + pub const INET6: u32 = 23; +}; +pub const SOCK = struct { + pub const STREAM: u32 = 1; +}; +pub const SOL = struct { + pub const SOCKET: i32 = 0xffff; +}; +pub const SO = struct { + /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, the opposite of POSIX's. + pub const REUSEADDR: u32 = @bitCast(@as(i32, -5)); + pub const RCVTIMEO: u32 = 0x1006; + pub const SNDTIMEO: u32 = 0x1005; +}; +pub const SHUT = struct { + pub const RD: i32 = 0; + pub const WR: i32 = 1; + pub const RDWR: i32 = 2; +}; +pub const MSG = struct { + pub const PEEK: u32 = 0x2; + /// Winsock has no per-call nonblocking flag; every caller polls first, so the receive never waits. + pub const DONTWAIT: u32 = 0; +}; +pub const POLL = struct { + pub const IN: i16 = 0x0300; // POLLRDNORM | POLLRDBAND + /// WSAPoll refuses POLLPRI; out-of-band data is not a signal the server uses. + pub const PRI: i16 = 0; + pub const OUT: i16 = 0x0010; + pub const ERR: i16 = 0x0001; + pub const HUP: i16 = 0x0002; + pub const NVAL: i16 = 0x0004; +}; +pub const pollfd = extern struct { fd: socket_t, events: i16, revents: i16 }; + +const WSAData = extern struct { version: u16, high: u16, rest: [400]u8 }; +const ws = struct { + extern "ws2_32" fn WSAStartup(version: u16, data: *WSAData) callconv(.winapi) c_int; + extern "ws2_32" fn WSAGetLastError() callconv(.winapi) c_int; + extern "ws2_32" fn WSAPoll(fds: [*]pollfd, n: u32, timeout: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn socket(af: c_int, kind: c_int, proto: c_int) callconv(.winapi) isize; + extern "ws2_32" fn bind(s: isize, addr: *const anyopaque, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn listen(s: isize, backlog: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn connect(s: isize, addr: *const anyopaque, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn accept(s: isize, addr: ?*anyopaque, len: ?*c_int) callconv(.winapi) isize; + extern "ws2_32" fn getsockname(s: isize, addr: *anyopaque, len: *c_int) callconv(.winapi) c_int; + extern "ws2_32" fn setsockopt(s: isize, level: c_int, opt: c_int, val: [*]const u8, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn shutdown(s: isize, how: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn closesocket(s: isize) callconv(.winapi) c_int; + extern "ws2_32" fn send(s: isize, buf: [*]const u8, len: c_int, flags: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn recv(s: isize, buf: [*]u8, len: c_int, flags: c_int) callconv(.winapi) c_int; +}; +extern "c" fn _errno() *c_int; + +var started = std.atomic.Value(u8).init(0); + +/// WSAStartup once per process, before the first socket call. +fn start() void { + if (started.load(.acquire) == 2) return; + if (started.cmpxchgStrong(0, 1, .acq_rel, .acquire) == null) { + var data: WSAData = undefined; + _ = ws.WSAStartup(0x0202, &data); + started.store(2, .release); + return; + } + while (started.load(.acquire) != 2) std.atomic.spinLoopHint(); +} + +/// -1 with errno from the thread's Winsock error. +fn failed() c_int { + const e: E = switch (ws.WSAGetLastError()) { + 10004 => .INTR, + 10013 => .ACCES, + 10014 => .FAULT, + 10022 => .INVAL, + 10024 => .MFILE, + 10035 => .AGAIN, + 10038 => .NOTSOCK, + 10048 => .ADDRINUSE, + 10049 => .ADDRNOTAVAIL, + 10050 => .NETDOWN, + 10053 => .CONNABORTED, + 10054 => .CONNRESET, + 10055 => .NOBUFS, + 10057 => .NOTCONN, + 10058 => .PIPE, + 10060 => .TIMEDOUT, + 10061 => .CONNREFUSED, + else => .IO, + }; + _errno().* = @backingInt(e); + return -1; +} + +fn len32(n: usize) c_int { + return @intCast(@min(n, std.math.maxInt(c_int))); +} + +pub const system = struct { + pub fn socket(family: u32, kind: u32, proto: u32) isize { + start(); + const fd = ws.socket(@intCast(family), @intCast(kind), @intCast(proto)); + return if (fd == -1) failed() else fd; + } + /// `len` as std.Io.Threaded's addressToPosix gives it (its socklen_t is unsigned on Windows). + pub fn bind(fd: socket_t, addr: *const anyopaque, len: anytype) c_int { + return if (ws.bind(fd, addr, @intCast(len)) != 0) failed() else 0; + } + pub fn connect(fd: socket_t, addr: *const anyopaque, len: anytype) c_int { + return if (ws.connect(fd, addr, @intCast(len)) != 0) failed() else 0; + } + pub fn listen(fd: socket_t, backlog: u31) c_int { + return if (ws.listen(fd, backlog) != 0) failed() else 0; + } + pub fn accept(fd: socket_t, addr: ?*anyopaque, len: ?*socklen_t) isize { + const s = ws.accept(fd, addr, len); + return if (s == -1) failed() else s; + } + pub fn getsockname(fd: socket_t, addr: *anyopaque, len: *socklen_t) c_int { + return if (ws.getsockname(fd, addr, len) != 0) failed() else 0; + } + pub fn shutdown(fd: socket_t, how: i32) c_int { + return if (ws.shutdown(fd, how) != 0) failed() else 0; + } + /// Closes a socket; files close through `sys.c.close`. + pub fn close(fd: socket_t) c_int { + return if (ws.closesocket(fd) != 0) failed() else 0; + } + /// Sends on a socket. + pub fn write(fd: socket_t, buf: [*]const u8, len: usize) isize { + const n = ws.send(fd, buf, len32(len), 0); + return if (n < 0) failed() else n; + } + pub fn recvfrom(fd: socket_t, buf: [*]u8, len: usize, flags: u32, addr: ?*anyopaque, alen: ?*socklen_t) isize { + _ = addr; + _ = alen; + const n = ws.recv(fd, buf, len32(len), @intCast(flags)); + return if (n < 0) failed() else n; + } +}; + +pub fn setsockopt(fd: socket_t, level: i32, opt: u32, value: []const u8) error{SetSockOptFailed}!void { + if (opt == SO.RCVTIMEO or opt == SO.SNDTIMEO) { + // POSIX gives a timeval; Winsock takes the milliseconds as a DWORD + const tv: *align(1) const timeval = @ptrCast(value.ptr); + const ms: u32 = @intCast(@as(i64, tv.sec) * 1000 + @divTrunc(@as(i64, tv.usec), 1000)); + if (ws.setsockopt(fd, level, @bitCast(opt), @ptrCast(&ms), @sizeOf(u32)) != 0) return error.SetSockOptFailed; + return; + } + if (ws.setsockopt(fd, level, @bitCast(opt), value.ptr, @intCast(value.len)) != 0) return error.SetSockOptFailed; +} + +/// WSAPoll, with POSIX's rule that an entry whose fd is negative is skipped (WSAPoll would fail the call). +pub fn poll(fds: []pollfd, timeout_ms: i32) error{PollFailed}!usize { + start(); + var live: [64]pollfd = undefined; + var where: [64]usize = undefined; + var n: usize = 0; + for (fds, 0..) |*f, i| { + f.revents = 0; + if (f.fd < 0) continue; + if (n == live.len) return error.PollFailed; + live[n] = f.*; + where[n] = i; + n += 1; + } + if (n == 0) { + if (timeout_ms > 0) Sleep(@intCast(timeout_ms)); + return 0; + } + const rc = ws.WSAPoll(&live, @intCast(n), timeout_ms); + if (rc < 0) return error.PollFailed; + for (live[0..n], where[0..n]) |f, i| fds[i].revents = f.revents; + return @intCast(rc); +} + +extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; + +/// Receives from a socket; 0 is the peer's orderly close, and a receive timeout (SO.RCVTIMEO) is WouldBlock. +pub fn read(fd: socket_t, buf: []u8) error{ ConnectionResetByPeer, WouldBlock }!usize { + const n = ws.recv(fd, buf.ptr, len32(buf.len), 0); + if (n < 0) return switch (ws.WSAGetLastError()) { + 10035, 10060 => error.WouldBlock, + else => error.ConnectionResetByPeer, + }; + return @intCast(n); +} + +// ---- processes ---- + +extern "kernel32" fn TerminateProcess(process: *anyopaque, code: c_uint) callconv(.winapi) c_int; + +/// kill for a spawned child's handle: TERM ends it with exit code 1; HUP and the rest have no Windows equivalent. +pub fn kill(process: std.os.windows.HANDLE, sig: SIG) error{ Unsupported, Unexpected }!void { + if (sig != .TERM) return error.Unsupported; + if (TerminateProcess(process, 1) == 0) return error.Unexpected; +} + +// ---- signals: console control events ---- + +pub const SIG = enum(c_int) { + HUP = 1, + INT = 2, + PIPE = 13, + TERM = 15, + _, + + pub const IGN: ?Handler = @ptrFromInt(1); + pub const DFL: ?Handler = null; +}; +pub const Handler = *align(1) const fn (SIG) callconv(.c) void; // align(1): IGN is address 1, as std.posix has it +pub const sigset_t = u32; +pub const Sigaction = struct { + handler: extern union { handler: ?Handler, sigaction: ?*const anyopaque }, + mask: sigset_t, + flags: u32, +}; + +pub fn sigemptyset() sigset_t { + return 0; +} + +var on_int: ?Handler = null; +var on_term: ?Handler = null; +var console_hooked = std.atomic.Value(bool).init(false); + +extern "kernel32" fn SetConsoleCtrlHandler(routine: ?*const fn (u32) callconv(.winapi) c_int, add: c_int) callconv(.winapi) c_int; + +/// Ctrl+C and Ctrl+Break are SIGINT; closing the console, logoff and shutdown are SIGTERM. +fn consoleEvent(kind: u32) callconv(.winapi) c_int { + const h = switch (kind) { + 0, 1 => on_int, + else => on_term, + } orelse return 0; + if (h == SIG.IGN.?) return 1; + h(if (kind <= 1) .INT else .TERM); + return 1; +} + +/// SIGINT and SIGTERM through console events; Windows has no SIGHUP or SIGPIPE (a closed socket fails the send). +pub fn sigaction(sig: SIG, act: ?*const Sigaction, old: ?*Sigaction) void { + if (old) |o| o.* = .{ .handler = .{ .handler = SIG.DFL }, .mask = 0, .flags = 0 }; + const a = act orelse return; + switch (sig) { + .INT => on_int = a.handler.handler, + .TERM => on_term = a.handler.handler, + else => return, + } + if (!console_hooked.swap(true, .acq_rel)) _ = SetConsoleCtrlHandler(consoleEvent, 1); +} + +// ---- descriptor limits: Windows has none to raise ---- + +pub const rlim_t = u64; +pub const rlimit = struct { cur: rlim_t, max: rlim_t }; +pub const rlimit_resource = enum { NOFILE }; + +pub fn getrlimit(_: rlimit_resource) error{Unsupported}!rlimit { + return error.Unsupported; +} + +pub fn setrlimit(_: rlimit_resource, _: rlimit) error{Unsupported}!void { + return error.Unsupported; +} + +/// A connected stream pair over loopback TCP (Windows has no socketpair); for tests that drive a connection. +pub fn socketpair(pair: *[2]socket_t) error{SocketPair}!void { + start(); + const l = ws.socket(@intCast(AF.INET), @intCast(SOCK.STREAM), 0); + if (l == -1) return error.SocketPair; + defer _ = ws.closesocket(l); + var addr: extern struct { family: u16 = 2, port: u16 = 0, addr: [4]u8 = .{ 127, 0, 0, 1 }, zero: [8]u8 = @splat(0) } = .{}; + var len: c_int = @sizeOf(@TypeOf(addr)); + if (ws.bind(l, &addr, len) != 0 or ws.listen(l, 1) != 0 or ws.getsockname(l, &addr, &len) != 0) return error.SocketPair; + const a = ws.socket(@intCast(AF.INET), @intCast(SOCK.STREAM), 0); + if (a == -1) return error.SocketPair; + errdefer _ = ws.closesocket(a); + if (ws.connect(a, &addr, len) != 0) return error.SocketPair; + const b = ws.accept(l, null, null); + if (b == -1) return error.SocketPair; + pair.* = .{ a, b }; +} diff --git a/zig/src/server/auth.zig b/zig/src/server/auth.zig index 5457651dd9..b2274e2dc0 100644 --- a/zig/src/server/auth.zig +++ b/zig/src/server/auth.zig @@ -1,6 +1,7 @@ //! API keys held as SHA-256 digests with labels; a restricted key file reread on change or SIGHUP. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const builtin = @import("builtin"); const log = @import("log.zig"); const Allocator = std.mem.Allocator; @@ -84,10 +85,10 @@ pub const Store = struct { const path = s.path.?; const pathz = try s.gpa.dupeSentinel(u8, path, 0); defer s.gpa.free(pathz); - const fd = posix.openatZ(posix.AT.FDCWD, pathz, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true, .NONBLOCK = true }, 0) catch return error.Unsafe; - defer _ = posix.system.close(fd); + const fd = openKeys(pathz) catch return error.Unsafe; + defer _ = platform.c.close(fd); const st = meta(fd, null) orelse return error.Unsafe; - if ((st.mode & posix.S.IFMT) != posix.S.IFREG or st.mode & 0o044 != 0) { + if ((st.mode & platform.c.S.IFMT) != platform.c.S.IFREG or st.mode & 0o044 != 0) { problem.* = "API key file must be a regular file unreadable by other users; use chmod 600"; return error.Invalid; } @@ -101,7 +102,7 @@ pub const Store = struct { const text = try a.alloc(u8, (1 << 20) + 1); var got: usize = 0; while (got < text.len) { - const n = posix.read(fd, text[got..]) catch return error.Unsafe; + const n = readKeys(fd, text[got..]) catch return error.Unsafe; if (n == 0) break; got += n; } @@ -290,7 +291,7 @@ test "auth cleanup on every initialization allocation failure" { try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "keys", .data = "ops: sk-file\nsk-unlabelled\n" }); const path = try std.fmt.allocPrintSentinel(gpa, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); defer gpa.free(path); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); const Case = struct { fn run(a: Allocator, file: ?[]const u8) !void { var problem: []const u8 = ""; @@ -304,6 +305,19 @@ test "auth cleanup on every initialization allocation failure" { try std.testing.checkAllAllocationFailures(gpa, Case.run, .{@as(?[]const u8, path)}); } +/// The key file opened for reading, never through a link. +fn openKeys(pathz: [*:0]const u8) !platform.c.fd_t { + if (builtin.os.tag != .windows) return posix.openatZ(posix.AT.FDCWD, pathz, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true, .NONBLOCK = true }, 0); + const fd = platform.c.open(pathz, .{ .ACCMODE = .RDONLY }); + return if (fd < 0) error.OpenFailed else fd; +} + +fn readKeys(fd: platform.c.fd_t, buf: []u8) !usize { + if (builtin.os.tag != .windows) return posix.read(fd, buf); + const n = platform.c.read(fd, buf.ptr, buf.len); + return if (n < 0) error.ReadFailed else @intCast(n); +} + /// A file's type and permission bits, size, identity and mtime: statx on Linux (its libc has no stat), stat elsewhere. const Meta = struct { mode: u32, size: u64, dev: u64, ino: u64, mtime_ns: i128 }; @@ -313,11 +327,23 @@ fn meta(fd: posix.fd_t, path: ?[*:0]const u8) ?Meta { const linux = std.os.linux; var sx: linux.Statx = undefined; const want: linux.STATX = .{ .TYPE = true, .MODE = true, .INO = true, .SIZE = true, .MTIME = true }; - const rc = if (path) |p| std.c.statx(posix.AT.FDCWD, p, 0, want, &sx) else std.c.statx(fd, "", 0x1000, want, &sx); // AT_EMPTY_PATH + const rc = if (path) |p| platform.c.statx(posix.AT.FDCWD, p, 0, want, &sx) else platform.c.statx(fd, "", 0x1000, want, &sx); // AT_EMPTY_PATH if (rc != 0) return null; const dev = (@as(u64, sx.dev_major) << 32) | sx.dev_minor; return .{ .mode = sx.mode, .size = sx.size, .dev = dev, .ino = sx.ino, .mtime_ns = @as(i128, sx.mtime.sec) * std.time.ns_per_s + sx.mtime.nsec }; } + if (builtin.os.tag == .windows) { + // NTFS keeps access in ACLs, not mode bits: the type and size are checked, and the profile's ACL is the guard + var st: platform.c.Stat = undefined; + const f = if (path) |p| platform.c.open(p, .{ .ACCMODE = .RDONLY }) else fd; + if (f < 0) return null; + defer if (path != null) { + _ = platform.c.close(f); + }; + if (platform.c.fstat(f, &st) != 0) return null; + const t = st.mtime(); + return .{ .mode = (st.mode & platform.c.S.IFMT) | 0o600, .size = @intCast(st.size), .dev = st.dev, .ino = st.ino, .mtime_ns = @as(i128, t.nsec) + @as(i128, t.sec) * std.time.ns_per_s }; + } var st: posix.Stat = undefined; const rc = if (path) |p| posix.system.stat(p, &st) else posix.system.fstat(fd, &st); if (posix.errno(rc) != .SUCCESS) return null; @@ -326,6 +352,7 @@ fn meta(fd: posix.fd_t, path: ?[*:0]const u8) ?Meta { } test "the key file: refused while others can read it, read at 0600, reread after an edit" { + if (builtin.os.tag == .windows) return error.SkipZigTest; // refusal by mode bits; NTFS keeps access in ACLs var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); @@ -334,11 +361,11 @@ test "the key file: refused while others can read it, read at 0600, reread after defer tmp.cleanup(); try tmp.dir.writeFile(io, .{ .sub_path = "keys", .data = "ops: sk-one\n" }); const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); var problem: []const u8 = ""; try std.testing.expectError(error.KeyFile, Store.init(a, &.{}, "", path, false, &problem)); try std.testing.expect(std.mem.indexOf(u8, problem, "chmod 600") != null); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); var s = try Store.init(std.testing.allocator, &.{"sk-cli"}, "sk-env", path, false, &problem); defer s.deinit(); try std.testing.expectEqualStrings("cli-1", s.match(a, "Bearer sk-cli", null).?); @@ -354,6 +381,7 @@ test "the key file: refused while others can read it, read at 0600, reread after } test "a request's key label outlives a reload of the key file" { + if (builtin.os.tag == .windows) return error.SkipZigTest; // refusal by mode bits; NTFS keeps access in ACLs const gpa = std.testing.allocator; const io = std.testing.io; var arena: std.heap.ArenaAllocator = .init(gpa); @@ -363,7 +391,7 @@ test "a request's key label outlives a reload of the key file" { defer tmp.cleanup(); try tmp.dir.writeFile(io, .{ .sub_path = "keys", .data = "ops: sk-one\n" }); const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); var problem: []const u8 = ""; var s = try Store.init(gpa, &.{}, "", path, false, &problem); defer if (s.file_arena) |*keys| keys.deinit(); @@ -373,7 +401,7 @@ test "a request's key label outlives a reload of the key file" { try std.testing.expectEqualStrings("ci", s.match(a, null, "sk-two").?); // another request rereads the file try std.testing.expectEqualStrings("ops", label); const kept = s.match(a, null, "sk-two").?; - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); s.checked_ns = null; try std.testing.expect(s.match(a, "Bearer sk-one", null) == null); // a refused reload frees the file's keys too try std.testing.expectEqualStrings("ci", kept); diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 76b205eeaa..3ad95bf2fc 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -122,7 +122,7 @@ pub const env = [_][]const u8{ "TENSORFOLD_API_KEY", "TENSORFOLD_NO_LIVE", "TENS pub const Args = struct { model: []const u8 = "", host: []const u8 = "127.0.0.1", - port: u16 = 8080, + port: u16 = 8888, name: []const u8 = "", chat_template: ?[]const u8 = null, alias: []const []const u8 = &.{}, diff --git a/zig/src/server/http_conn.zig b/zig/src/server/http_conn.zig index f29d9f8684..536e8649d7 100644 --- a/zig/src/server/http_conn.zig +++ b/zig/src/server/http_conn.zig @@ -1,6 +1,7 @@ //! One HTTP/1.1 connection as Python's BaseHTTPRequestHandler serves it: bounded heads, keep-alive, error pages. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const Allocator = std.mem.Allocator; const log = @import("log.zig"); diff --git a/zig/src/server/hub.zig b/zig/src/server/hub.zig index 5a155e6a1b..4e0efb1a5a 100644 --- a/zig/src/server/hub.zig +++ b/zig/src/server/hub.zig @@ -21,7 +21,7 @@ fn isDir(io: std.Io, path: []const u8) bool { /// The checkpoint directory for ``model``, or null with ``problem`` set (no downloads: the cache must hold it). pub fn resolve(a: Allocator, io: std.Io, environ: ?*const std.process.Environ.Map, model: []const u8, problem: *[]const u8) !?[]const u8 { - const home = if (environ) |e| e.get("HOME") orelse "" else ""; + const home = if (environ) |e| e.get("HOME") orelse e.get("USERPROFILE") orelse "" else ""; // Windows sets USERPROFILE const expanded = if (std.mem.startsWith(u8, model, "~/")) try std.fs.path.join(a, &.{ home, model[2..] }) else model; if (isDir(io, expanded)) return expanded; if (!isRepoId(model)) { diff --git a/zig/src/server/late_system_test.zig b/zig/src/server/late_system_test.zig index e661ae2fea..7f4622dcd1 100644 --- a/zig/src/server/late_system_test.zig +++ b/zig/src/server/late_system_test.zig @@ -1,5 +1,6 @@ //! A later system message is re-roled to user on Messages when the template cannot keep it, matching chat. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const server = @import("root.zig"); const json = server.json; @@ -132,10 +133,10 @@ const Capture = struct { const Reply = struct { status: u16, body: []u8 }; fn exchange(srv: *server.Server, path: []const u8, body: []const u8) !Reply { - var pair: [2]std.c.fd_t = undefined; - if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; - defer _ = std.c.close(pair[0]); - defer _ = std.c.close(pair[1]); + var pair: [2]platform.posix.socket_t = undefined; + if (@import("builtin").os.tag == .windows) try platform.posix.socketpair(&pair) else if (platform.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer _ = platform.posix.system.close(pair[0]); + defer _ = platform.posix.system.close(pair[1]); const head = try std.fmt.allocPrint(std.testing.allocator, "POST {s} HTTP/1.1\r\nHost: t\r\nContent-Length: {d}\r\n\r\n", .{ path, body.len }); defer std.testing.allocator.free(head); try writeAll(pair[0], head); @@ -157,23 +158,23 @@ fn exchange(srv: *server.Server, path: []const u8, body: []const u8) !Reply { return .{ .status = status, .body = payload }; } -fn writeAll(fd: std.c.fd_t, data: []const u8) !void { +fn writeAll(fd: platform.posix.socket_t, data: []const u8) !void { var sent: usize = 0; while (sent < data.len) { - const n = std.c.write(fd, data[sent..].ptr, data.len - sent); + const n = platform.posix.system.write(fd, data[sent..].ptr, data.len - sent); if (n <= 0) return error.Closed; sent += @intCast(n); } } -fn readReady(a: std.mem.Allocator, fd: std.c.fd_t) ![]u8 { +fn readReady(a: std.mem.Allocator, fd: platform.posix.socket_t) ![]u8 { var out: std.ArrayList(u8) = .empty; errdefer out.deinit(a); var tmp: [4096]u8 = undefined; while (true) { - var pollfd = [_]std.posix.pollfd{.{ .fd = fd, .events = std.posix.POLL.IN, .revents = 0 }}; - if (try std.posix.poll(&pollfd, 200) == 0) break; - const n = std.posix.read(fd, &tmp) catch |e| return e; + var pollfd = [_]platform.posix.pollfd{.{ .fd = fd, .events = platform.posix.POLL.IN, .revents = 0 }}; + if (try platform.posix.poll(&pollfd, 200) == 0) break; + const n = platform.posix.read(fd, &tmp) catch |e| return e; if (n == 0) break; try out.appendSlice(a, tmp[0..n]); } diff --git a/zig/src/server/listener.zig b/zig/src/server/listener.zig index 8eb7a27df6..a6492e737d 100644 --- a/zig/src/server/listener.zig +++ b/zig/src/server/listener.zig @@ -1,6 +1,7 @@ //! The listening socket: SO_REUSEADDR only (never REUSEPORT, so a second server on the port fails), backlog 128. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const net = std.Io.net; const log = @import("log.zig"); const Threaded = std.Io.Threaded; diff --git a/zig/src/server/live.zig b/zig/src/server/live.zig index cb0cbaf640..3d6f8b482e 100644 --- a/zig/src/server/live.zig +++ b/zig/src/server/live.zig @@ -1,5 +1,6 @@ //! The live status line (connections, decode and prefill tok/s), redrawn in a terminal, and /health's copy of it. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const log = @import("log.zig"); @@ -94,7 +95,7 @@ fn columns() usize { if (os == .linux) { if (std.os.linux.ioctl(1, std.os.linux.T.IOCGWINSZ, @intFromPtr(&ws)) == 0 and ws.col > 0) return ws.col; } else if (os.isDarwin()) { - if (std.c.ioctl(1, @bitCast(@as(u32, 0x40087468)), &ws) == 0 and ws.col > 0) return ws.col; + if (platform.c.ioctl(1, @bitCast(@as(u32, 0x40087468)), &ws) == 0 and ws.col > 0) return ws.col; } return 100; } diff --git a/zig/src/server/log.zig b/zig/src/server/log.zig index 40430985a2..629bf90bee 100644 --- a/zig/src/server/log.zig +++ b/zig/src/server/log.zig @@ -1,5 +1,6 @@ //! Server output: ``[tensorfold]`` lines on stdout, and the live status line they clear first. const std = @import("std"); +const platform = @import("sys"); const Conn = @import("http_conn.zig").Conn; const afterChars = @import("reply_text.zig").afterChars; @@ -23,8 +24,8 @@ const clear = "\r\x1b[2K"; fn writeOut(bytes: []const u8) void { var off: usize = 0; while (off < bytes.len) { - const rc = std.posix.system.write(1, bytes[off..].ptr, bytes.len - off); - if (std.posix.errno(rc) != .SUCCESS or rc <= 0) return; + const rc = platform.c.write(1, bytes[off..].ptr, bytes.len - off); + if (platform.c.errno(rc) != .SUCCESS or rc <= 0) return; off += @intCast(rc); } } diff --git a/zig/src/server/request_log.zig b/zig/src/server/request_log.zig index ea4c8b7997..6d6513113e 100644 --- a/zig/src/server/request_log.zig +++ b/zig/src/server/request_log.zig @@ -1,5 +1,6 @@ //! TENSORFOLD_REQUEST_LOG: each chat and completion body appended as one JSON line, images redacted, for replays. const std = @import("std"); +const platform = @import("sys"); const json = @import("json.zig"); const Value = json.Value; const Allocator = std.mem.Allocator; @@ -33,8 +34,51 @@ pub fn append(a: Allocator, path: []const u8, body: Value) void { if (body.strField("priority")) |p| if (std.mem.eql(u8, p, "background")) return; const line = json.stringify(a, redact(a, body) catch return, .{}) catch return; const pathz = a.dupeSentinel(u8, path, 0) catch return; - const fd = std.posix.openatZ(std.posix.AT.FDCWD, pathz, .{ .ACCMODE = .WRONLY, .APPEND = true, .CREAT = true }, 0o644) catch return; - defer _ = std.posix.system.close(fd); + const fd = platform.c.open(pathz, .{ .ACCMODE = .WRONLY, .APPEND = true, .CREAT = true }, @as(platform.c.mode_t, 0o644)); + if (fd < 0) return; + defer _ = platform.c.close(fd); const framed = std.mem.concat(a, u8, &.{ line, "\n" }) catch return; - _ = std.posix.system.write(fd, framed.ptr, framed.len); + _ = platform.c.write(fd, framed.ptr, framed.len); +} + +test "appends from several threads keep every line whole" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const cwd = try std.process.currentPathAlloc(std.testing.io, a); + defer a.free(cwd); + const path = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}/log.jsonl", .{ cwd, tmp.sub_path }); + defer a.free(path); + const threads = 8; + const lines = 200; + const head = "{\"line\": \""; + const Worker = struct { + fn run(p: []const u8, k: u8) void { + var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator); + defer arena.deinit(); + const text = arena.allocator().alloc(u8, 300) catch return; + @memset(text, 'a' + k); + for (0..lines) |_| { + const o = json.newObject(arena.allocator()) catch return; + o.put(arena.allocator(), "line", .{ .string = text }) catch return; + append(arena.allocator(), p, .{ .object = o }); + } + } + }; + var workers: [threads]std.Thread = undefined; + for (&workers, 0..) |*w, k| w.* = try std.Thread.spawn(.{}, Worker.run, .{ path, @as(u8, @intCast(k)) }); + for (workers) |w| w.join(); + const got = try std.Io.Dir.cwd().readFileAlloc(std.testing.io, path, a, .limited(1 << 24)); + defer a.free(got); + var counts: [threads]usize = @splat(0); + var it = std.mem.splitScalar(u8, got, '\n'); + while (it.next()) |line| { + if (line.len == 0) continue; + try std.testing.expectEqual(head.len + 300 + 2, line.len); + const k = line[head.len] - 'a'; + try std.testing.expect(k < threads); + for (line[head.len..][0..300]) |c| try std.testing.expectEqual('a' + k, c); + counts[k] += 1; + } + for (counts) |n| try std.testing.expectEqual(@as(usize, lines), n); } diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 7630f5217a..9cfa2dc1d4 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -44,6 +44,7 @@ test { _ = @import("logprobs_test.zig"); _ = @import("late_system_test.zig"); _ = @import("compact_test.zig"); + _ = @import("request_log.zig"); _ = @import("completion_prompt_test.zig"); _ = @import("server_init_test.zig"); _ = @import("decisions.zig"); diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index da27cb6b0e..46a8415f3b 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -1,5 +1,6 @@ //! A server's life: keys, the listening socket, the startup lines, the live line, and a clean stop on SIGTERM or SIGINT. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const auth = @import("auth.zig"); @@ -10,7 +11,7 @@ const listener_mod = @import("listener.zig"); const model_text = @import("model_text.zig"); const server_mod = @import("server.zig"); const Allocator = std.mem.Allocator; -const posix = std.posix; +const posix = platform.posix; var stop_requested: std.atomic.Value(bool) = .init(false); @@ -33,7 +34,7 @@ pub fn stopping() bool { } /// The hard limit macOS reads as unlimited. -const rlim_infinity: u64 = std.c.RLIM.INFINITY; +const rlim_infinity: u64 = platform.c.RLIM.INFINITY; /// The targets to try, in order: the hard limit when finite, else the fallbacks 65536 then 10240, first that works; never an unlimited soft limit, since a loop over every descriptor would do absurd work. fn fillTargets(max: u64, out: *[2]u64) []const u64 { @@ -211,6 +212,7 @@ test "the targets are the hard limit, or the fallbacks under an unlimited one" { } test "the raise lifts a low soft limit, and a second raise does nothing" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // no descriptor limits to raise const before = try posix.getrlimit(.NOFILE); defer posix.setrlimit(.NOFILE, before) catch {}; if (@as(u64, @intCast(before.max)) <= 256) return error.SkipZigTest; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 011734ee9a..4d76d98105 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -1,5 +1,6 @@ //! The server: one engine and tokenizer behind the HTTP routes, one thread a connection. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const errors = @import("errors.zig"); @@ -280,7 +281,7 @@ pub const Server = struct { pub fn serve(srv: *Server, listener: listener_mod.Listener, stop: *const std.atomic.Value(bool)) void { while (listener.accept(stop)) |c| { const t = std.Thread.spawn(.{ .stack_size = 8 << 20 }, connection, .{ srv, c }) catch { - _ = std.posix.system.close(c.fd); + _ = platform.posix.system.close(c.fd); continue; }; t.detach(); @@ -290,7 +291,7 @@ pub const Server = struct { fn connection(srv: *Server, accepted: listener_mod.Listener.Accepted) void { _ = srv.open_connections.fetchAdd(1, .acq_rel); defer _ = srv.open_connections.fetchSub(1, .acq_rel); - defer _ = std.posix.system.close(accepted.fd); + defer _ = platform.posix.system.close(accepted.fd); var peer_buf: [64]u8 = undefined; const peer = listener_mod.peerText(&peer_buf, accepted.peer); var conn = http_conn.Conn.init(srv.gpa, accepted.fd, peer) catch return; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aec6b9ec1b..c83cdf5123 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -15,6 +15,7 @@ const resume_tests = @import("resume_tests.zig"); const prompt_cuts = @import("prompt_cuts.zig"); const chunk_costs = @import("chunk_costs.zig"); const grouped_tests = @import("grouped_tests.zig"); +const moe_tests = @import("moe_tests.zig"); const fp8_tests = @import("fp8_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); @@ -58,6 +59,7 @@ const usage = \\ prompt-cuts MODEL IDS_FILE a prompt keeping states inside chunks and at their ends (background passes too), and resumed from each, equals one uncut pass \\ recapture MODEL graphs captured again while another sequence is bound: both keep their sampling rules \\ grouped-plan the shared expert plan against experts.route's bytes (oracle/grouped_plan.py) + \\ moe softmax top-k routing and the SwiGLU pair experts against host references (synthetic) \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes @@ -141,6 +143,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "prompt-cuts")) return prompt_cuts.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "recapture")) return recapture_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "grouped-plan")) return grouped_tests.plan(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "moe")) return moe_tests.run(gpu); if (std.mem.eql(u8, cmd, "carveout")) { const mib = if (rest.len > 0) try std.fmt.parseInt(usize, rest[0], 10) else cuda.carveout.default_bytes >> 20; return carveout_tests.run(gpu, if (rest.len > 1) rest[1] else cuda.carveout.default_card, mib); diff --git a/zig/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig index faa25fdfdd..cf4cf9fc62 100644 --- a/zig/tests/cuda/modelopt_tests.zig +++ b/zig/tests/cuda/modelopt_tests.zig @@ -187,22 +187,31 @@ fn lanes(gpu: Gpu, rng: std.Random, p: *const modelopt.Packer, lin: cuda.qlinear defer y.free(); var y2 = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 2); defer y2.free(); - try lin.decode(s, q, .{ .x = dx.ptr }, y.ptr, 16); + // slice partials for the lane's split K: GPUs without clusters (before sm_90) reduce through them + const lw = switch (q) { + .fp8, .bf16 => |v| v, + else => unreachable, + }; + var part_bytes: usize = 4; + for ([_]usize{ 16, 128, m }) |r| part_bytes = @max(part_bytes, lin.lane.?.partBytes(r, lw)); + var part = try cuda.DeviceBuffer.alloc(gpu.d, part_bytes); + defer part.free(); + try lin.decode(s, q, .{ .x = dx.ptr, .part = part.ptr }, y.ptr, 16); try s.synchronize(); const wide = try check.download(gpu, y); defer a.free(wide); for (1..16) |rows| { try y2.fill8(0xff, s.handle); - try lin.decode(s, q, .{ .x = dx.ptr }, y2.ptr, rows); + try lin.decode(s, q, .{ .x = dx.ptr, .part = part.ptr }, y2.ptr, rows); try s.synchronize(); try same(gpu, "decode rows at a narrower width", y2, wide[0 .. rows * n * 2]); } var worst = [2]f64{ 0, 0 }; - try lin.prompt(s, q, .{ .x = dx.ptr }, y.ptr, m); + try lin.prompt(s, q, .{ .x = dx.ptr, .part = part.ptr }, y.ptr, m); try s.synchronize(); const prompt = try check.download(gpu, y); defer a.free(prompt); - try lin.prompt(s, q, .{ .x = dx.ptr }, y2.ptr, 128); + try lin.prompt(s, q, .{ .x = dx.ptr, .part = part.ptr }, y2.ptr, 128); try s.synchronize(); try same(gpu, "prompt rows in a chunk of 128", y2, prompt[0 .. 128 * n * 2]); for (0..m) |r| for (0..n) |c| { diff --git a/zig/tests/cuda/moe_tests.zig b/zig/tests/cuda/moe_tests.zig new file mode 100644 index 0000000000..2bfd33a415 --- /dev/null +++ b/zig/tests/cuda/moe_tests.zig @@ -0,0 +1,309 @@ +//! Softmax top-k routing (cuda/moe.zig) and the SwiGLU pair form of the affine-4 experts against host references. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +fn bf16f(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +fn f32bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @intCast((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +fn rbf(x: f32) f32 { + return bf16f(f32bf16(x)); +} + +fn randBf16(r: std.Random, scale: f32) u16 { + return f32bf16((r.float(f32) * 2.0 - 1.0) * scale); +} + +fn upload(gpu: Gpu, bytes: []const u8) !cuda.DeviceBuffer { + return cuda.DeviceBuffer.fromHost(gpu.d, bytes); +} + +/// The router's fp32 logits: one fma chain a (row, expert) in index order, as tf_moe_router sums. +fn router(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { + const a = gpu.gpa; + const rows = 5; + const dims = 2048; + const count = 257; + const x = try a.alloc(u16, rows * dims); + defer a.free(x); + const w = try a.alloc(u16, count * dims); + defer a.free(w); + for (x) |*v| v.* = randBf16(r, 1.0); + for (w) |*v| v.* = randBf16(r, 0.1); + const want = try a.alloc(f32, rows * count); + defer a.free(want); + for (0..rows) |i| for (0..count) |e| { + var acc: f32 = 0; + for (0..dims) |k| acc = @mulAdd(f32, bf16f(x[i * dims + k]), bf16f(w[e * dims + k]), acc); + want[i * count + e] = acc; + }; + var dx = try upload(gpu, std.mem.sliceAsBytes(x)); + defer dx.free(); + var dw = try upload(gpu, std.mem.sliceAsBytes(w)); + defer dw.free(); + var dl = try cuda.DeviceBuffer.alloc(gpu.d, rows * count * 4); + defer dl.free(); + try f.route(s, dx.ptr, dims, rows, dw.ptr, dl.ptr, dims, count); + try s.synchronize(); + const got = try check.download(gpu, dl); + defer a.free(got); + try check.sameBytes("router logits", got, std.mem.sliceAsBytes(want)); + check.pass("moe router: {d} rows x {d} experts over {d} inputs equal the host's fma chains", .{ rows, count, dims }); +} + +/// Each row's picks and weights by cuda/moe.py's rule: top-k by logit, ties to the lower id, exp(l - top) / sum. +fn topk(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { + const a = gpu.gpa; + const rows = 7; + const experts = 256; + const top_k = 8; + const slots = top_k + 1; + const logits = try a.alloc(f32, rows * (experts + 1)); + defer a.free(logits); + for (logits) |*v| v.* = (r.float(f32) * 2.0 - 1.0) * 4.0; + // a planted tie, a row of equal logits and a NaN that never wins + logits[1 * (experts + 1) + 3] = logits[1 * (experts + 1) + 200]; + for (0..experts) |e| logits[2 * (experts + 1) + e] = 0.5; + logits[3 * (experts + 1) + 7] = std.math.nan(f32); + const want_ids = try a.alloc(i32, rows * slots); + defer a.free(want_ids); + const want_w = try a.alloc(f32, rows * slots); + defer a.free(want_w); + for (0..rows) |i| { + const row = logits[i * (experts + 1) ..][0..experts]; + var v: [experts]f32 = undefined; + for (&v, row) |*d, l| d.* = if (l != l) -std.math.inf(f32) else l; + var top = -std.math.inf(f32); + for (v) |l| top = @max(top, l); + var total: f32 = 0; + var ex: [top_k]f32 = undefined; + for (0..top_k) |k| { + var bi: usize = 0; + var bv = -std.math.inf(f32); + for (v, 0..) |l, e| if (l > bv) { + bv = l; + bi = e; + }; + ex[k] = @exp(bv - top); + total += ex[k]; + want_ids[i * slots + k] = @intCast(bi); + v[bi] = -std.math.inf(f32); + } + for (0..top_k) |k| want_w[i * slots + k] = rbf(ex[k] / total); + want_ids[i * slots + top_k] = experts; + const sg = rbf(logits[i * (experts + 1) + experts]); + want_w[i * slots + top_k] = rbf(1.0 / (1.0 + @exp(-sg))); + } + var dl = try upload(gpu, std.mem.sliceAsBytes(logits)); + defer dl.free(); + var dp = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer dp.free(); + var dw = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer dw.free(); + try f.pick(s, dl.ptr, dp.ptr, dw.ptr, rows, experts, top_k); + try s.synchronize(); + const got_ids = try check.download(gpu, dp); + defer a.free(got_ids); + try check.sameBytes("top-k ids", got_ids, std.mem.sliceAsBytes(want_ids)); + const got_w = try check.download(gpu, dw); + defer a.free(got_w); + var off: usize = 0; + for (std.mem.bytesAsSlice(f32, got_w), want_w) |g, wv| { + // exp on the device and the host may differ in the last fp32 place: the bf16 weight then moves one step at most + if (g != wv and @abs(g - wv) > @abs(wv) * (1.0 / 128.0)) off += 1; + } + try check.expect(off == 0, "top-k weights: {d} of {d} differ by more than a bf16 step", .{ off, want_w.len }); + check.pass("moe top-k: {d} rows' picks of {d} over {d} experts equal the host's (a tie, a flat row and a NaN included), weights within a bf16 step", .{ rows, top_k, experts }); +} + +/// The weighted sum of a row's slots in slot order, fma in fp32, rounded once to bf16. +fn combine(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { + const a = gpu.gpa; + const rows = 3; + const slots = 4; + const dims = 2048; + const y = try a.alloc(f32, rows * slots * dims); + defer a.free(y); + const w = try a.alloc(f32, rows * slots); + defer a.free(w); + for (y) |*v| v.* = (r.float(f32) * 2.0 - 1.0) * 3.0; + for (w) |*v| v.* = rbf(r.float(f32)); + const want = try a.alloc(u16, rows * dims); + defer a.free(want); + for (0..rows) |i| for (0..dims) |d| { + var acc: f32 = 0; + for (0..slots) |k| acc = @mulAdd(f32, w[i * slots + k], y[(i * slots + k) * dims + d], acc); + want[i * dims + d] = f32bf16(acc); + }; + var dy = try upload(gpu, std.mem.sliceAsBytes(y)); + defer dy.free(); + var dw = try upload(gpu, std.mem.sliceAsBytes(w)); + defer dw.free(); + var dout = try cuda.DeviceBuffer.alloc(gpu.d, rows * dims * 2); + defer dout.free(); + try f.combine(s, dy.ptr, true, dw.ptr, dout.ptr, rows, slots, dims); + try s.synchronize(); + const got = try check.download(gpu, dout); + defer a.free(got); + try check.sameBytes("combined rows", got, std.mem.sliceAsBytes(want)); + check.pass("moe combine: {d} rows of {d} slots equal the host's fma chains rounded to bf16", .{ rows, slots }); +} + +/// Random affine-4 experts (words, scales, biases) as the checkpoint stores them. +const Table = struct { words: []u32, scales: []u16, biases: []u16, count: usize, n: usize, k: usize }; + +fn table(a: std.mem.Allocator, r: std.Random, count: usize, n: usize, k: usize) !Table { + const t: Table = .{ .words = try a.alloc(u32, count * n * k / 8), .scales = try a.alloc(u16, count * n * k / 64), .biases = try a.alloc(u16, count * n * k / 64), .count = count, .n = n, .k = k }; + for (t.words) |*w| w.* = r.int(u32); + for (t.scales) |*v| v.* = f32bf16(0.02 + r.float(f32) * 0.03); + for (t.biases) |*v| v.* = randBf16(r, 0.2); + return t; +} + +/// Weight (e, j, i) dequantized: scale * code + bias. +fn weight(t: Table, e: usize, j: usize, i: usize) f64 { + const row = (e * t.n + j); + const code: f64 = @floatFromInt((t.words[row * (t.k / 8) + i / 8] >> @intCast(4 * (i % 8))) & 15); + return code * bf16f(t.scales[row * (t.k / 64) + i / 64]) + bf16f(t.biases[row * (t.k / 64) + i / 64]); +} + +fn silu(x: f32) f32 { + return x / (1.0 + @exp(-x)); +} + +/// The pair form: gate and up interleaved per group, SwiGLU with the kernel's bf16 rounding points, down to fp32 sums. +fn pairExperts(gpu: Gpu, s: cuda.Stream, r: std.Random) !void { + const a = gpu.gpa; + var mods: [3]cuda.Module = undefined; + mods[0] = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mods[0].unload(); + mods[1] = try cuda.Module.load(gpu.d, cuda.kernels.experts_prefill); + defer mods[1].unload(); + mods[2] = try cuda.Module.load(gpu.d, cuda.kernels.experts_pack); + defer mods[2].unload(); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + const ex = try cuda.experts.Affine4Experts.resolve(mods[0], mods[1], mods[2], sms); + const plan_fns = try cuda.grouped.Router.resolve(mods[0]); + const count = 6; + const ni = 64; + const dims = 128; + const rows = 3; + const slots = 4; + const pairs = rows * slots; + const gate = try table(a, r, count, ni, dims); + defer { + a.free(gate.words); + a.free(gate.scales); + a.free(gate.biases); + } + const up = try table(a, r, count, ni, dims); + defer { + a.free(up.words); + a.free(up.scales); + a.free(up.biases); + } + const down = try table(a, r, count, dims, ni); + defer { + a.free(down.words); + a.free(down.scales); + a.free(down.biases); + } + const x = try a.alloc(u16, rows * dims); + defer a.free(x); + for (x) |*v| v.* = randBf16(r, 1.0); + const picks = [pairs]i32{ 0, 2, 5, 5, 1, 2, 3, 5, 4, 4, 0, 5 }; + // the host reference + const want = try a.alloc(f64, pairs * dims); + defer a.free(want); + const act = try a.alloc(f32, ni); + defer a.free(act); + for (0..pairs) |p| { + const e: usize = @intCast(picks[p]); + const xr = x[(p / slots) * dims ..][0..dims]; + for (0..ni) |j| { + var g: f64 = 0; + var u: f64 = 0; + for (0..dims) |i| { + g += weight(gate, e, j, i) * bf16f(xr[i]); + u += weight(up, e, j, i) * bf16f(xr[i]); + } + const gv = rbf(@floatCast(g)); + const uv = rbf(@floatCast(u)); + act[j] = rbf(rbf(silu(gv)) * uv); + } + for (0..dims) |d| { + var y: f64 = 0; + for (0..ni) |j| y += weight(down, e, d, j) * act[j]; + want[p * dims + d] = y; + } + } + // the device: packed tables, the plan, up (pair) then down + var dpair = try cuda.DeviceBuffer.alloc(gpu.d, cuda.experts.Affine4Experts.pairBytes(count, ni, dims)); + defer dpair.free(); + var ddown = try cuda.DeviceBuffer.alloc(gpu.d, cuda.experts.Affine4Experts.bytes(count, dims, ni)); + defer ddown.free(); + for ([_]struct { Table, u64, usize, usize }{ .{ gate, dpair.ptr, 2, 0 }, .{ up, dpair.ptr, 2, 1 }, .{ down, ddown.ptr, 1, 0 } }) |item| { + const t = item[0]; + var w = try upload(gpu, std.mem.sliceAsBytes(t.words)); + defer w.free(); + var sc = try upload(gpu, std.mem.sliceAsBytes(t.scales)); + defer sc.free(); + var bi = try upload(gpu, std.mem.sliceAsBytes(t.biases)); + defer bi.free(); + try ex.packInto(s, w.ptr, sc.ptr, bi.ptr, item[1], t.count, t.n, t.k, item[2], item[3]); + try s.synchronize(); + } + var dx = try upload(gpu, std.mem.sliceAsBytes(x)); + defer dx.free(); + var dpicks = try upload(gpu, std.mem.sliceAsBytes(&picks)); + defer dpicks.free(); + var members = try cuda.DeviceBuffer.alloc(gpu.d, pairs * 4); + defer members.free(); + var items = try cuda.DeviceBuffer.alloc(gpu.d, cuda.grouped.maxItems(pairs, count, 16) * 12); + defer items.free(); + var counts = try cuda.DeviceBuffer.alloc(gpu.d, 8); + defer counts.free(); + var dact = try cuda.DeviceBuffer.alloc(gpu.d, pairs * ni * 2); + defer dact.free(); + var dy = try cuda.DeviceBuffer.alloc(gpu.d, pairs * dims * 4); + defer dy.free(); + const plan: cuda.grouped.Plan = .{ .members = members.ptr, .items = items.ptr, .counts = counts.ptr, .rank = 0, .hist = 0 }; + try plan_fns.route(s, dpicks.ptr, pairs, count, 16, plan); + const layer: cuda.experts.Layer = .{ .format = .affine4, .up = dpair.ptr, .down = ddown.ptr, .width = ni, .dims = dims, .experts = count, .swiglu = true }; + const grouped: cuda.experts.Grouped = .{ .affine4 = &ex }; + const items_n = cuda.grouped.maxItems(pairs, count, 16); + try grouped.decode(s, .up, .{ .x = dx.ptr, .stride = dims, .slots = slots }, layer, plan, dact.ptr, items_n, -1); + try grouped.decode(s, .down_f32, .{ .x = dact.ptr, .stride = ni }, layer, plan, dy.ptr, items_n, -1); + try s.synchronize(); + const got = try check.download(gpu, dy); + defer a.free(got); + var worst: f64 = 0; + var scale: f64 = 0; + for (std.mem.bytesAsSlice(f32, got), want) |g, wv| { + worst = @max(worst, @abs(@as(f64, g) - wv)); + scale = @max(scale, @abs(wv)); + } + try check.expect(worst <= scale * 2e-3, "pair experts: worst error {e} against a largest sum of {e}", .{ worst, scale }); + check.pass("moe experts: {d} pairs over {d} SwiGLU experts (gate and up interleaved) within {e} of the fp64 reference (largest sum {e})", .{ pairs, count, worst, scale }); +} + +pub fn run(gpu: Gpu) !void { + var prng = std.Random.DefaultPrng.init(0x5eed); + const r = prng.random(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.moe); + defer mod.unload(); + const f = try cuda.moe.Fns.resolve(mod); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + try router(gpu, f, s, r); + try topk(gpu, f, s, r); + try combine(gpu, f, s, r); + try pairExperts(gpu, s, r); +} diff --git a/zig/tests/cuda/pack_compare.zig b/zig/tests/cuda/pack_compare.zig index b1cc379a36..10f9041c65 100644 --- a/zig/tests/cuda/pack_compare.zig +++ b/zig/tests/cuda/pack_compare.zig @@ -1,5 +1,6 @@ //! tf-train-pack's checks on either backend: packed passes against one a pass, and passes after kept prefixes. const std = @import("std"); +const platform = @import("sys"); /// What one sequence's pass left: its loss, whether it came back, its pushes and its keys at the pocket's layers. const Seen = struct { loss: f32, recalled: bool, push: []f32, keys: []u16 }; @@ -169,7 +170,7 @@ pub fn For(comptime train: type) type { /// TF_PACK_DUMP: the packed target pass's losses and pushes written there, for another build's to compare. fn dump(gpa: std.mem.Allocator, io: std.Io, seen: []const Seen) !void { - const path = std.c.getenv("TF_PACK_DUMP") orelse return; + const path = platform.c.getenv("TF_PACK_DUMP") orelse return; var bytes: std.ArrayList(u8) = .empty; for (seen) |x| { try bytes.appendSlice(gpa, std.mem.asBytes(&x.loss)); diff --git a/zig/tests/cuda/train_pack.zig b/zig/tests/cuda/train_pack.zig index 2ab91949b9..d721c19887 100644 --- a/zig/tests/cuda/train_pack.zig +++ b/zig/tests/cuda/train_pack.zig @@ -1,5 +1,6 @@ //! Training passes holding several sequences against one a pass, and from kept prefixes against whole, on CUDA. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const nemotron = @import("nemotron"); const check_mod = @import("check.zig"); @@ -17,7 +18,7 @@ const open = 2; pub fn run(gpu: check_mod.Gpu, model: []const u8, path: []const u8, count: usize) !void { const gpa = gpu.gpa; const io = gpu.io; - const kernels = std.c.getenv("TENSORFOLD_CUDA_KERNELS"); + const kernels = platform.c.getenv("TENSORFOLD_CUDA_KERNELS"); const e = try nemotron.Engine.init(gpa, io, gpu.ctx, model, if (kernels) |k| std.mem.span(k) else null, .{ .context = 8192, .mtp = false, .graphs = false, .sampling = null, .segments = 1 }); defer e.deinit(); var living: ?nemotron.sites.Sites = null; diff --git a/zig/tests/fabric_library.zig b/zig/tests/fabric_library.zig index 04141e546a..2e09f6aefb 100644 --- a/zig/tests/fabric_library.zig +++ b/zig/tests/fabric_library.zig @@ -1,5 +1,6 @@ //! The fabric bound to MCDMA's own libmcdma-rpc (built from MCDMA's rpc/ by tools/zig/test_fabric.sh, MCDMA_RPC_LIBRARY). const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const fabric = @import("fabric"); const layout = fabric.layout; @@ -9,7 +10,7 @@ const mailbox = fabric.mailbox; const gpa = std.testing.allocator; fn open() !words.Library { - if (std.c.getenv("MCDMA_RPC_LIBRARY") == null) return error.SkipZigTest; + if (platform.c.getenv("MCDMA_RPC_LIBRARY") == null) return error.SkipZigTest; return words.Library.open(null); } diff --git a/zig/tests/flashnext_bench.zig b/zig/tests/flashnext_bench.zig index 7a0c73d528..d15901c68a 100644 --- a/zig/tests/flashnext_bench.zig +++ b/zig/tests/flashnext_bench.zig @@ -1,5 +1,6 @@ //! FZ_DBENCH: Flash Next dense and DeltaNet kernels at 1-16 rows against the recorded ones, bit for bit, then timed. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const fz = tf.flashnext_replay; @@ -238,7 +239,7 @@ pub fn run(r: *Run, m: *Model, prompt: []const u32, gpa: std.mem.Allocator, aren }; const pfs = [_]usize{ 1, 2 }; - const bits = std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_DBENCH").?), "bits"); + const bits = std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_DBENCH").?), "bits"); // bits: core lane against the recorded kernel, every layer, one Mac and speed-up layouts, every read-ahead depth const keep = try gpa.alloc(u8, fz.MAXR * fz.VOCAB * 2); defer gpa.free(keep); diff --git a/zig/tests/flashnext_reuse.zig b/zig/tests/flashnext_reuse.zig index eeeda53247..c9c66a3ca1 100644 --- a/zig/tests/flashnext_reuse.zig +++ b/zig/tests/flashnext_reuse.zig @@ -1,5 +1,6 @@ //! Flash Next prompt reuse against fresh passes: agent-like turns, every depth, copy drafts off and on, state bytes. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const api = @import("engine_api"); @@ -121,9 +122,9 @@ pub fn main(init: std.process.Init) !void { } const pool = mtl.objc.Pool.push(); defer pool.pop(); - const n_out: usize = if (std.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 64; + const n_out: usize = if (platform.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 64; var depths: std.ArrayList(?usize) = .empty; - var dit = std.mem.tokenizeScalar(u8, if (std.c.getenv("FZ_DEPTHS")) |v| std.mem.span(v) else "1,2,3,5,0", ','); + var dit = std.mem.tokenizeScalar(u8, if (platform.c.getenv("FZ_DEPTHS")) |v| std.mem.span(v) else "1,2,3,5,0", ','); while (dit.next()) |d| { const k = try std.fmt.parseInt(usize, d, 10); try depths.append(a, if (k == 0) null else k); diff --git a/zig/tests/flashnext_run.zig b/zig/tests/flashnext_run.zig index 87efe89c96..5c1e1c3234 100644 --- a/zig/tests/flashnext_run.zig +++ b/zig/tests/flashnext_run.zig @@ -2,6 +2,7 @@ //! recorded variant, every weight the pack's or the checkpoint's. Checks one-row greedy tokens and every row of the //! Python engine's drafted windows (with rollback), then times each window size. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); const tf = @import("tensorfold"); @@ -46,13 +47,13 @@ const jsonInt = fz.jsonInt; fn armName(seg_ab: bool, arm: usize) []const u8 { if (seg_ab) return if (arm == 1) "segments on " else "segments off"; - if (std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "fz")) return if (arm == 1) "fz kernels" else "recorded "; + if (platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "fz")) return if (arm == 1) "fz kernels" else "recorded "; return if (arm == 1) "copy on " else "copy off"; } /// FZ_PROMPT_N: only the prompt's first n tokens (all when unset). fn promptCut() !usize { - const v = std.c.getenv("FZ_PROMPT_N") orelse return std.math.maxInt(usize); + const v = platform.c.getenv("FZ_PROMPT_N") orelse return std.math.maxInt(usize); return std.fmt.parseInt(usize, std.mem.span(v), 10); } @@ -124,7 +125,7 @@ pub fn main(init: std.process.Init) !void { defer arena_state.deinit(); const arena = arena_state.allocator(); const device = try mtl.Device.init(); - if (std.c.getenv("FZ_COMPILE_CHECK") != null) { // every prompt-kernel source and the 6-bit file on this Metal compiler + if (platform.c.getenv("FZ_COMPILE_CHECK") != null) { // every prompt-kernel source and the 6-bit file on this Metal compiler var failed: usize = 0; var files: std.ArrayList(ks.File) = .empty; try files.appendSlice(gpa, &ks.prefill); @@ -149,8 +150,8 @@ pub fn main(init: std.process.Init) !void { } std.process.exit(if (failed == 0) 0 else 1); } - if (std.c.getenv("FZ_OPORDER") != null) return @import("flashnext_bench.zig").opOrder(device); // no model loaded - if (std.c.getenv("FZ_MMA_PEAK") != null) { // the tensor units' rate on register fragments: 16x32x16 ops a second + if (platform.c.getenv("FZ_OPORDER") != null) return @import("flashnext_bench.zig").opOrder(device); // no model loaded + if (platform.c.getenv("FZ_MMA_PEAK") != null) { // the tensor units' rate on register fragments: 16x32x16 ops a second const alib = try mtl.Library.fromSource(device, try fz.frags.source(device, arena, ks.flashnext_attn), mtl.CompileOptions.mlx()); const pipe = try mtl.Pipeline.init(device, alib, "tf_mma_peak", false); const queue = try device.queue(); @@ -174,9 +175,9 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_ENGINE") != null) { // the served engine: its own load, warm-up and reply, on FZ_REF's prompt + if (platform.c.getenv("FZ_ENGINE") != null) { // the served engine: its own load, warm-up and reply, on FZ_REF's prompt const fx = @import("tensorfold").flashnext_engine; - const rp = std.mem.span(std.c.getenv("FZ_REF") orelse return error.NoRef); + const rp = std.mem.span(platform.c.getenv("FZ_REF") orelse return error.NoRef); const rf = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}", .{rp}, 0)); const rj = try std.json.parseFromSliceLeaky(std.json.Value, arena, rf.bytes[0..rf.size], .{}); const all_items = rj.object.get("prompt").?.array.items; @@ -184,9 +185,9 @@ pub fn main(init: std.process.Init) !void { const toks = try arena.alloc(u32, items.len); for (items, 0..) |x, i| toks[i] = @intCast(x.integer); const e = try fx.Engine.load(gpa, init.io, args[1], args[2]); - if (std.c.getenv("FZ_COPY_MIN")) |v| e.copy_min = try std.fmt.parseInt(u32, std.mem.span(v), 10); - if (std.c.getenv("FZ_COPY_LONG")) |v| e.copy_long = try std.fmt.parseInt(u32, std.mem.span(v), 10); - if (std.c.getenv("FZ_ENGINE_WARM") != null) try e.warm(); + if (platform.c.getenv("FZ_COPY_MIN")) |v| e.copy_min = try std.fmt.parseInt(u32, std.mem.span(v), 10); + if (platform.c.getenv("FZ_COPY_LONG")) |v| e.copy_long = try std.fmt.parseInt(u32, std.mem.span(v), 10); + if (platform.c.getenv("FZ_ENGINE_WARM") != null) try e.warm(); const Show = struct { got: std.ArrayList(u32) = .empty, a: std.mem.Allocator, @@ -200,11 +201,11 @@ pub fn main(init: std.process.Init) !void { return false; } }; - const n_out: usize = if (std.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 256; + const n_out: usize = if (platform.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 256; var replies: [2][]u32 = undefined; - const seg_ab = std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "seg"); + const seg_ab = platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "seg"); var depths: std.ArrayList(?usize) = .empty; // FZ_DEPTHS=3,5,7: each depth's arms in turn (0: the depth rule, p: plain) - if (std.c.getenv("FZ_DEPTHS")) |v| { + if (platform.c.getenv("FZ_DEPTHS")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |d| { if (std.mem.eql(u8, d, "p")) { @@ -214,9 +215,9 @@ pub fn main(init: std.process.Init) !void { const n = try std.fmt.parseInt(usize, d, 10); try depths.append(arena, if (n == 0) null else n); } - } else try depths.append(arena, if (std.c.getenv("FZ_DEPTH")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else null); + } else try depths.append(arena, if (platform.c.getenv("FZ_DEPTH")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else null); // FZ_AB=fz (with FZ_LANE=1 FZ_GDN=2): recorded dense and DeltaNet kernels, then core lane, fz_gdn, kept states - const fz_ab = std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "fz"); + const fz_ab = platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "fz"); const fz_on = .{ e.r.lane_new, e.r.gdn_pipe, e.r.gdn_kept }; for (depths.items) |depth| for (0..2) |arm| { // copy drafts off, then on (FZ_AB=seg: staggered prompt segments off, then on) const plain = depth != null and depth.? == 0; // one token a round: one arm (copies need drafts) @@ -256,7 +257,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("{s}: reply hash {x:0>16}\n", .{ name, std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(sh.got.items)) }); if (arm == 1) for (res.copy_by_len, 0..) |cl, n| if (cl[0] > 0) std.debug.print(" match {d}: {d} rounds, {d:.2} landed\n", .{ n, cl[0], @as(f64, @floatFromInt(cl[1])) / @as(f64, @floatFromInt(cl[0])) }); }; - if (std.c.getenv("FZ_NO_REF") != null) return; // the reply hashes are the check (two Macs against one) + if (platform.c.getenv("FZ_NO_REF") != null) return; // the reply hashes are the check (two Macs against one) // the reference: the same prompt chunks, then one row a step e.hostMode(); const mm = e.m; @@ -287,28 +288,28 @@ pub fn main(init: std.process.Init) !void { // the replay below reads every expert, embedding and table at 6/32: other widths are the engine's (Engine.load) if (!try fz.engine_widths.uniformDir(arena, init.io, args[1])) return error.HarnessNeeds6Bit32; var r = Run{ .arena = arena, .device = device, .queue = try device.queue() }; - r.fused_xsum = std.c.getenv("FZ_FUSED_XSUM") != null; - r.serial = std.c.getenv("FZ_SERIAL") != null; - r.xnew = std.c.getenv("FZ_XNEW") != null; - r.split = std.c.getenv("FZ_SPLIT") != null; - r.gdn_step = std.c.getenv("FZ_GDN_STEP") != null; - r.hc_mma = std.c.getenv("FZ_HC_MMA") != null; - r.hc_up = std.c.getenv("FZ_HCCHECK") != null or if (std.c.getenv("FZ_HC_UP")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else false; + r.fused_xsum = platform.c.getenv("FZ_FUSED_XSUM") != null; + r.serial = platform.c.getenv("FZ_SERIAL") != null; + r.xnew = platform.c.getenv("FZ_XNEW") != null; + r.split = platform.c.getenv("FZ_SPLIT") != null; + r.gdn_step = platform.c.getenv("FZ_GDN_STEP") != null; + r.hc_mma = platform.c.getenv("FZ_HC_MMA") != null; + r.hc_up = platform.c.getenv("FZ_HCCHECK") != null or if (platform.c.getenv("FZ_HC_UP")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else false; r.event = try device.sharedEvent(); - r.dense = r.xnew and std.c.getenv("FZ_DENSE") != null; - r.dense_target = r.dense and std.c.getenv("FZ_DENSE_TARGET") != null; - r.xpack = r.xnew and std.c.getenv("FZ_XPACK") != null; - r.grouped = r.xnew and !r.xpack and std.c.getenv("FZ_GROUPED") != null; - r.xfused = r.xnew and !r.xpack and !r.grouped and std.c.getenv("FZ_XFUSED") != null; - r.xsx = r.xnew and std.c.getenv("FZ_XSX") != null; - r.copy = std.c.getenv("FZ_COPY") != null; - r.prefetch = !r.serial and std.c.getenv("FZ_PREFETCH") != null; + r.dense = r.xnew and platform.c.getenv("FZ_DENSE") != null; + r.dense_target = r.dense and platform.c.getenv("FZ_DENSE_TARGET") != null; + r.xpack = r.xnew and platform.c.getenv("FZ_XPACK") != null; + r.grouped = r.xnew and !r.xpack and platform.c.getenv("FZ_GROUPED") != null; + r.xfused = r.xnew and !r.xpack and !r.grouped and platform.c.getenv("FZ_XFUSED") != null; + r.xsx = r.xnew and platform.c.getenv("FZ_XSX") != null; + r.copy = platform.c.getenv("FZ_COPY") != null; + r.prefetch = !r.serial and platform.c.getenv("FZ_PREFETCH") != null; const t0 = mtl.clock.seconds(); try r.compile(args[2]); r.sel = try Select.init(&r, MAXR); - if (std.c.getenv("FZ_CATCH_CHECK") != null) return catchCheck(&r, arena); - r.lane_new = r.xnew and std.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (same bits) - if (r.xnew and std.c.getenv("FZ_GDN") != null) r.gdn_pipe = try fz.gdn_step.compile(&r, false); + if (platform.c.getenv("FZ_CATCH_CHECK") != null) return catchCheck(&r, arena); + r.lane_new = r.xnew and platform.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (same bits) + if (r.xnew and platform.c.getenv("FZ_GDN") != null) r.gdn_pipe = try fz.gdn_step.compile(&r, false); const t1 = mtl.clock.seconds(); const index_file = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors.index.json", .{args[1]}, 0)); @@ -320,8 +321,8 @@ pub fn main(init: std.process.Init) !void { while (fit.next()) |name| try r.indexFile(try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ args[1], name.* }, 0)); try r.indexFile(try std.fmt.allocPrintSentinel(arena, "{s}/pack.safetensors", .{args[2]}, 0)); - const own_prompt = std.c.getenv("FZ_REF") != null; // another prompt: passes against Python's records are skipped - const ref_path = if (std.c.getenv("FZ_REF")) |p| try std.fmt.allocPrintSentinel(arena, "{s}", .{std.mem.span(p)}, 0) else try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{args[2]}, 0); + const own_prompt = platform.c.getenv("FZ_REF") != null; // another prompt: passes against Python's records are skipped + const ref_path = if (platform.c.getenv("FZ_REF")) |p| try std.fmt.allocPrintSentinel(arena, "{s}", .{std.mem.span(p)}, 0) else try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{args[2]}, 0); const ref_file = try mtl.MappedFile.open(ref_path); const ref = try std.json.parseFromSliceLeaky(std.json.Value, arena, ref_file.bytes[0..ref_file.size], .{}); const ple_ref = ref.object.get("ple").?.object; @@ -502,15 +503,15 @@ pub fn main(init: std.process.Init) !void { const t2 = mtl.clock.seconds(); std.debug.print("compiled in {d:.2} s, loaded {d:.1} GB in {d:.1} s\n", .{ t1 - t0, @as(f64, @floatFromInt(r.loaded)) / 1e9, t2 - t1 }); - if (std.c.getenv("FZ_DBENCH") != null) { // dense classes and DeltaNet timed by rows (flashnext_bench.zig) + if (platform.c.getenv("FZ_DBENCH") != null) { // dense classes and DeltaNet timed by rows (flashnext_bench.zig) var toks: [MAXR]u32 = undefined; for (0..MAXR) |i| toks[i] = @intCast(ref.object.get("prompt").?.array.items[i].integer); return @import("flashnext_bench.zig").run(&r, m, &toks, gpa, arena); } - if (std.c.getenv("FZ_PROFILE") != null) { - if (std.c.getenv("TF_FLASHNEXT_TP")) |path| r.tp = try fz.Tp2.init(arena, r.device, std.mem.span(path)); + if (platform.c.getenv("FZ_PROFILE") != null) { + if (platform.c.getenv("TF_FLASHNEXT_TP")) |path| r.tp = try fz.Tp2.init(arena, r.device, std.mem.span(path)); var hcskips: std.ArrayList(u32) = .empty; // FZ_HCSKIP=1,2,4: the profile again under each hc knock-out - if (std.c.getenv("FZ_HCSKIP")) |v| { + if (platform.c.getenv("FZ_HCSKIP")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try hcskips.append(arena, try std.fmt.parseInt(u32, x, 10)); } else try hcskips.append(arena, 0); @@ -520,7 +521,7 @@ pub fn main(init: std.process.Init) !void { var pk: [MAXR]u32 = undefined; m.reset(); for (0..3) |_| try m.window(toks[0..1], &pk); - if (std.c.getenv("FZ_PROFILE_AB") != null) { // in-model A/B, interleaved: recorded, core lane, fz_gdn, both + if (platform.c.getenv("FZ_PROFILE_AB") != null) { // in-model A/B, interleaved: recorded, core lane, fz_gdn, both const gpipe = try fz.gdn_step.compile(&r, false); const arms = [_][2]bool{ .{ false, false }, .{ true, false }, .{ false, true }, .{ true, true } }; for ([_]usize{ 1, 2, 4, 6, 8, 16 }) |rows| { @@ -537,12 +538,12 @@ pub fn main(init: std.process.Init) !void { return; } var ups: std.ArrayList(bool) = .empty; // FZ_HC_UP=0,1: the profile with the recorded up projection, then fz_hc_up - if (std.c.getenv("FZ_HC_UP")) |v| { + if (platform.c.getenv("FZ_HC_UP")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try ups.append(arena, !std.mem.eql(u8, x, "0")); } else try ups.append(arena, false); var prow: std.ArrayList(usize) = .empty; // FZ_PROFILE_ROWS=2,3: the windows' widths (default 1, 4, 8) - if (std.c.getenv("FZ_PROFILE_ROWS")) |v| { + if (platform.c.getenv("FZ_PROFILE_ROWS")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try prow.append(arena, try std.fmt.parseInt(usize, x, 10)); } else try prow.appendSlice(arena, &.{ 1, 4, 8 }); @@ -551,7 +552,7 @@ pub fn main(init: std.process.Init) !void { r.hcskip = hs; var base: f64 = 0; for (names, 0..) |name, c| { - const only = std.c.getenv("FZ_PROFILE_ONLY") != null; // every class but this one knocked out + const only = platform.c.getenv("FZ_PROFILE_ONLY") != null; // every class but this one knocked out r.skip = if (c == 0) (if (only) 0x1ff else 0) else if (only) 0x1ff & ~(@as(u32, 1) << @intCast(c - 1)) else @as(u32, 1) << @intCast(c - 1); m.gpu_seconds = 0; for (0..20) |_| try m.window(toks[0..rows], &pk); @@ -573,7 +574,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_DUAL") != null) { // two lane groups on two queues: one window, two in a row, two at once + if (platform.c.getenv("FZ_DUAL") != null) { // two lane groups on two queues: one window, two in a row, two at once const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -648,7 +649,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_HCCHECK") != null) { // fz_hc_up against the recorded up projection: every row's logits and streams, bit for bit + if (platform.c.getenv("FZ_HCCHECK") != null) { // fz_hc_up against the recorded up projection: every row's logits and streams, bit for bit const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -687,7 +688,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("hc check: {s}\n", .{if (bad == 0) "bit-identical" else "DIFFERENT"}); return; } - if (std.c.getenv("FZ_GCHECK") != null) { // grouped or fused experts against fz_xgu/fz_xdown: every row's logits, bit for bit + if (platform.c.getenv("FZ_GCHECK") != null) { // grouped or fused experts against fz_xgu/fz_xdown: every row's logits, bit for bit const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -714,12 +715,12 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_XTIME") != null) { // the experts' share of a window: chain rows against one token repeated + if (platform.c.getenv("FZ_XTIME") != null) { // the experts' share of a window: chain rows against one token repeated const want0 = ref.object.get("tokens").?.array.items; var pk: [MAXR]u32 = undefined; m.reset(); for (0..3) |_| try m.window(&.{@intCast(want0[0].integer)}, &pk); - if (std.c.getenv("FZ_GPARTS") != null) for ([_]usize{ 2, 8 }) |rows| { + if (platform.c.getenv("FZ_GPARTS") != null) for ([_]usize{ 2, 8 }) |rows| { var toks: [MAXR]u32 = undefined; for (0..rows) |i| toks[i] = @intCast(want0[i].integer); for ([_]u32{ 0, 1, 4, 6 }) |gs| { @@ -747,21 +748,21 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_PREFILL") != null) { // the prompt path against the exact 8-row windows, then one-row decode + if (platform.c.getenv("FZ_PREFILL") != null) { // the prompt path against the exact 8-row windows, then one-row decode const all_items = ref.object.get("prompt").?.array.items; const pr_items = all_items[0..@min(all_items.len, try promptCut())]; // FZ_PROMPT_N: the prompt's first n tokens const toks = try gpa.alloc(u32, pr_items.len); defer gpa.free(toks); for (pr_items, 0..) |x, i| toks[i] = @intCast(x.integer); - if (std.c.getenv("FZ_SATTN_CHECK") != null) { // prompt attention: the decode's kernels vs the tensor units + if (platform.c.getenv("FZ_SATTN_CHECK") != null) { // prompt attention: the decode's kernels vs the tensor units var pr = try Prompt.init(&r, args[2], r.xnew_header); - if (std.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); - const ab0: []const u8 = if (std.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; + if (platform.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); + const ab0: []const u8 = if (platform.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; const seg2 = std.mem.eql(u8, ab0, "seg2"); // seg2: one chunk at a time against staggered segments (core/segments.zig): FZ_SEGS=2 (default) the // engine's rule (two segments of at least FZ_SEG_MIN rows), 1 one segment, 3-4 that many of up to a chunk - const nseg: usize = if (std.c.getenv("FZ_SEGS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 2; - const seg_min: usize = if (std.c.getenv("FZ_SEG_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 1; + const nseg: usize = if (platform.c.getenv("FZ_SEGS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 2; + const seg_min: usize = if (platform.c.getenv("FZ_SEG_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 1; if (seg2 and (nseg == 0 or nseg > tf.segments.MAX)) return error.Segments; var extra: [3]Prompt = undefined; var ps: [4]*Prompt = undefined; @@ -773,7 +774,7 @@ pub fn main(init: std.process.Init) !void { var outs: [2][48]u32 = undefined; const logits: [2][]u16 = .{ try gpa.alloc(u16, VOCAB), try gpa.alloc(u16, VOCAB) }; var best: [2]f64 = .{ 1e9, 1e9 }; - const ab: []const u8 = if (std.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; + const ab: []const u8 = if (platform.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; for (0..2) |arm| { // arm 0 the old kernels, arm 1 the new: FZ_AB attn (default), scan or tiles if (std.mem.eql(u8, ab, "scan")) pr.scan4 = arm == 1 else if (std.mem.eql(u8, ab, "tiles")) pr.tall_tiles = arm == 1 else if (std.mem.eql(u8, ab, "gu")) pr.fused_gu = arm == 1 else if (std.mem.eql(u8, ab, "bm32")) pr.expert_bm = if (arm == 1) 32 else 0 else if (std.mem.eql(u8, ab, "scan8")) pr.scan8 = arm == 1 else if (std.mem.eql(u8, ab, "seg2")) {} else pr.fast_attn = arm == 1; for (0..3) |run| { @@ -830,7 +831,7 @@ pub fn main(init: std.process.Init) !void { const a0 = mtl.clock.seconds(); var at: usize = 0; var last_n: usize = 1; - const ref_upto: usize = if (std.c.getenv("FZ_PF_REF")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else toks.len; + const ref_upto: usize = if (platform.c.getenv("FZ_PF_REF")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else toks.len; while (at < @min(ref_upto, toks.len)) { const n: usize = @min(MAXR, toks.len - at); try m.window(toks[at .. at + n], &pk); @@ -865,7 +866,7 @@ pub fn main(init: std.process.Init) !void { } // the prompt path var pr = try Prompt.init(&r, args[2], r.xnew_header); - if (std.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); + if (platform.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); m.reset(); const b0 = mtl.clock.seconds(); var first: u32 = 0; @@ -927,7 +928,7 @@ pub fn main(init: std.process.Init) !void { pr.skip = 0; return; } - if (std.c.getenv("FZ_QMM6") != null) { // 6-bit tensor-unit projections for prompt chunks: checked and timed + if (platform.c.getenv("FZ_QMM6") != null) { // 6-bit tensor-unit projections for prompt chunks: checked and timed const lib = try mtl.Library.fromSource(device, try fz.frags.source(device, arena, ks.flashnext_qmm6), mtl.CompileOptions.mlx()); const qmm_pipe = try mtl.Pipeline.init(device, lib, "tf_qmm6_t_nax", false); const off_pipe = try mtl.Pipeline.init(device, lib, "tf_expert_offsets6", false); @@ -1061,7 +1062,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_OVERLAP") != null) { // how many experts a window's rows share, layer by layer + if (platform.c.getenv("FZ_OVERLAP") != null) { // how many experts a window's rows share, layer by layer const pr = ref.object.get("prompt").?.array.items; const want0 = ref.object.get("tokens").?.array.items; r.ar = .{ .b = try r.buffer(64) }; @@ -1148,10 +1149,10 @@ pub fn main(init: std.process.Init) !void { var pick: [MAXR]u32 = undefined; // FZ_LONG: prompts of any length through the prompt chunks (reference and GPU-side rounds alike) - const long_mode = std.c.getenv("FZ_LONG") != null; - if (std.c.getenv("FZ_PRE1") != null) try m.window(&.{1000}, &pick); // a decode window before the prompt chunks + const long_mode = platform.c.getenv("FZ_LONG") != null; + if (platform.c.getenv("FZ_PRE1") != null) try m.window(&.{1000}, &pick); // a decode window before the prompt chunks var pr_long: ?Prompt = if (long_mode) try Prompt.init(&r, args[2], r.xnew_header) else null; - if (pr_long) |*pr| if (std.c.getenv("FZ_STEP")) |v| { + if (pr_long) |*pr| if (platform.c.getenv("FZ_STEP")) |v| { pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); }; const ptoks = try gpa.alloc(u32, prompt.len); @@ -1190,20 +1191,20 @@ pub fn main(init: std.process.Init) !void { const steps: f64 = @floatFromInt(want.len - 1); std.debug.print("one-row reference starts {any}\n", .{got.items[0..@min(24, got.items.len)]}); std.debug.print("one row: {d}/{d} tokens equal to Python's; {d:.1} tok/s ({d:.2} ms a step, GPU {d:.2} ms)\n", .{ same, want.len, steps / (t4 - t3), (t4 - t3) * 1e3 / steps, m.gpu_seconds * 1e3 / steps }); - if (std.c.getenv("FZ_ONLY1") != null) { + if (platform.c.getenv("FZ_ONLY1") != null) { if (same < want.len) std.debug.print("first difference at token {d}: got {d}, Python {d}\n", .{ same, got.items[same], want[same].integer }); return; } // 2. the Python engine's drafted windows, every row's pick, with rollback m.reset(); - if (!long_mode and std.c.getenv("FZ_ONLY17") == null) for (prompt) |tok| { + if (!long_mode and platform.c.getenv("FZ_ONLY17") == null) for (prompt) |tok| { try m.window(&.{@intCast(tok.integer)}, &pick); m.keepRows(&.{@intCast(tok.integer)}, 1); }; var bad: usize = 0; var total_rows: usize = 0; - const skip_mid = long_mode or std.c.getenv("FZ_ONLY17") != null; // the one-row reference, then GPU-side rounds only + const skip_mid = long_mode or platform.c.getenv("FZ_ONLY17") != null; // the one-row reference, then GPU-side rounds only const rounds = if (own_prompt or skip_mid) &[_]std.json.Value{} else ref.object.get("rounds").?.array.items; for (rounds, 0..) |round, ri| { const o = round.object; @@ -1310,7 +1311,7 @@ pub fn main(init: std.process.Init) !void { const adapt_story = [_][3]usize{ .{ 3, 3, 0 }, .{ 0, 7, 9 } }; const adapt_own = [_][3]usize{ .{ 3, 3, 0 }, .{ 6, 6, 0 }, .{ 0, 7, 9 } }; const adapt: []const [3]usize = if (own_prompt) &adapt_own else &adapt_story; - const copy_min: usize = if (std.c.getenv("FZ_COPY_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 3; + const copy_min: usize = if (platform.c.getenv("FZ_COPY_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 3; var hist: std.ArrayList(u32) = .empty; for (prompt) |x| try hist.append(gpa, @intCast(x.integer)); const n_prompt = hist.items.len; @@ -1441,7 +1442,7 @@ pub fn main(init: std.process.Init) !void { // N+1 while round N runs and reads the emitted tokens from a ring. Depth 0: the depth rule picks each round's // drafts from the rounds read back so far (the window two rounds behind the one being encoded). var depths7: std.ArrayList(usize) = .empty; - if (std.c.getenv("FZ_D7")) |v| { // e.g. FZ_D7=6,0 + if (platform.c.getenv("FZ_D7")) |v| { // e.g. FZ_D7=6,0 var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |d| try depths7.append(gpa, try std.fmt.parseInt(usize, d, 10)); } else try depths7.appendSlice(gpa, &.{ 3, 6, 0 }); diff --git a/zig/tests/flashnext_think.zig b/zig/tests/flashnext_think.zig index 283d217b36..8c7bbe0e0b 100644 --- a/zig/tests/flashnext_think.zig +++ b/zig/tests/flashnext_think.zig @@ -1,5 +1,6 @@ //! tf-flashnext-think [CASES] [SEED]: Flash Next's verdict kernels under thinking budgets against the lane core. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -13,8 +14,8 @@ pub fn main(init: std.process.Init) !void { defer dev.deinit(); var g = try tf.flashnext_think_check.Rig.init(dev); defer g.deinit(); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const sum = try tf.flashnext_think_check.run(init.gpa, &g, cases, seed); - const s = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e9; // mach ticks are 125/3 ns on Apple silicon + const s = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e9; // mach ticks are 125/3 ns on Apple silicon std.debug.print("{s}: {d} cases (seed {d}), plain and drafted == the lane core in each; closed at the budget {d}, think end first {d}; {d:.1} s\n", .{ std.mem.span(dev.name()), sum.cases, seed, sum.closed, sum.ended_first, s }); } diff --git a/zig/tests/glm_attn_probe.zig b/zig/tests/glm_attn_probe.zig index c6bc8ac6e6..e86a001d61 100644 --- a/zig/tests/glm_attn_probe.zig +++ b/zig/tests/glm_attn_probe.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's dense latent attention (forward.attendDense) on given inputs, for a bit check against a reference. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -57,8 +58,8 @@ pub fn main(init: std.process.Init) !void { try out.appendSlice(arena, bufs[3].contents()[0 .. 64 * n * 2]); try out.appendSlice(arena, bufs[4].contents()[0..qb]); } - const file = std.c.fopen(try std.fmt.allocPrintSentinel(arena, "{s}", .{args[2]}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(out.items.ptr, 1, out.items.len, file) != out.items.len) return error.WriteFailed; + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(arena, "{s}", .{args[2]}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(out.items.ptr, 1, out.items.len, file) != out.items.len) return error.WriteFailed; std.debug.print("{d} cases\n", .{cases}); } diff --git a/zig/tests/glm_dsa_prompt_check.zig b/zig/tests/glm_dsa_prompt_check.zig index e88d50ebac..9141bf7fcf 100644 --- a/zig/tests/glm_dsa_prompt_check.zig +++ b/zig/tests/glm_dsa_prompt_check.zig @@ -1,5 +1,6 @@ //! A prompt window's experts (every expert once, on the tensor units) against the decode path's 16-row windows, by row. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const g = tf.glm_dsa; @@ -53,8 +54,8 @@ pub fn main(init: std.process.Init) !void { for (q.x.slice(f32, @as(usize, rows) * hid)) |*v| v.* = prng.random().floatNorm(f32); const dec = try device.buffer(@as(usize, rows) * hid * 4, mtl.ResourceOptions.shared); const out = try device.buffer(@as(usize, rows) * hid * 4, mtl.ResourceOptions.shared); - if (std.c.getenv("GLM_ATTN") != null) return attention(&K, &L, &s, &q, queue, device, &c, rows, w.full); - if (std.c.getenv("GLM_NARROW") != null) return narrow(&K, &L, &q, queue, device, &c); + if (platform.c.getenv("GLM_ATTN") != null) return attention(&K, &L, &s, &q, queue, device, &c, rows, w.full); + if (platform.c.getenv("GLM_NARROW") != null) return narrow(&K, &L, &q, queue, device, &c); // the decode path: route and experts 16 rows at a time (each window's rows copied into the 16-row scratch) var dec_ms: f64 = 0; var k: u32 = 0; diff --git a/zig/tests/glm_dsa_run.zig b/zig/tests/glm_dsa_run.zig index 6c918d5676..94132983a1 100644 --- a/zig/tests/glm_dsa_run.zig +++ b/zig/tests/glm_dsa_run.zig @@ -1,5 +1,6 @@ //! GLM-5.3 on the pair without the server: rank 0 runs a prompt and its greedy reply, rank 1 follows. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const tok = @import("tokenizer"); const g = tf.glm_dsa; @@ -13,7 +14,7 @@ const usage = ; fn env(name: [:0]const u8, default: u64) u64 { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(u64, std.mem.span(v), 10) catch default; } @@ -30,9 +31,9 @@ pub fn main(init: std.process.Init) !void { } var ram: u64 = 0; var len: usize = 8; - if (std.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; + if (platform.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; const limit = env("GLM_LIMIT_GB", ram * 7 / 10 / 1_000_000_000) * 1_000_000_000; - const layers: ?u32 = if (std.c.getenv("GLM_LAYERS") != null) @intCast(env("GLM_LAYERS", 78)) else null; + const layers: ?u32 = if (platform.c.getenv("GLM_LAYERS") != null) @intCast(env("GLM_LAYERS", 78)) else null; const depth: u32 = @intCast(env("GLM_DEPTH", 0)); const e = try g.engine.Engine.load(a, io, argv[1], argv[2], layers, depth > 0, @intCast(env("GLM_CONTEXT", 8192)), 1, limit); defer e.deinit(); @@ -53,7 +54,7 @@ pub fn main(init: std.process.Init) !void { } fn once(e: *g.engine.Engine, io: std.Io, t: *const tok.Tokenizer, prompt: []const u32, max: u64, depth: u32) !void { - const forced = std.c.getenv("GLM_FORCED") != null; + const forced = platform.c.getenv("GLM_FORCED") != null; var agree: usize = 0; const t0 = std.Io.Timestamp.now(io, .awake); try g.follow.begin(e, 0, prompt); diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig index f2ceb725a0..e91e0657db 100644 --- a/zig/tests/glm_rows_speed.zig +++ b/zig/tests/glm_rows_speed.zig @@ -1,5 +1,6 @@ //! tf-glm-rows-speed: GLM-5.3-Flash's decode products at the checkpoints' widths and shapes, GPU time and GB/s. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const sources = @import("kernel_sources"); @@ -555,13 +556,13 @@ pub fn main() !void { 42, 3 * affine.bytes(N, D, .{ .bits = 2, .group = 64 }) * TOPK + sh8[0].bytes + sh8[1].bytes, ); - if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + if (platform.c.getenv("GLM_SPEED_CANDIDATES") != null) { const lib = try cand.candidateLib(device); defer lib.deinit(); try cand.moeCandidates(&b, lib, p, gu2, dn2, sh8); } } - if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + if (platform.c.getenv("GLM_SPEED_CANDIDATES") != null) { const lib = try cand.candidateLib(device); defer lib.deinit(); std.debug.print("candidates at 8 bits (bits checked == glm_rows):\n", .{}); @@ -571,7 +572,7 @@ pub fn main() !void { try cand.candidates(&b, lib, p, m, s.name); } } - if (std.c.getenv("GLM_SPEED_PROMPT") != null) try speed_prompt.prompt(&b, k, p); + if (platform.c.getenv("GLM_SPEED_PROMPT") != null) try speed_prompt.prompt(&b, k, p); std.debug.print("uniform 4/64:\n", .{}); for (uni.rows[0..uni.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); uni.total("uniform 4/64"); diff --git a/zig/tests/glm_run.zig b/zig/tests/glm_run.zig index 13194d6d0f..97ad4105a8 100644 --- a/zig/tests/glm_run.zig +++ b/zig/tests/glm_run.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash replies on the native engine at each draft depth: hashes, first differences from depth 0 and the reference, speeds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const glm = tf.glm; @@ -53,7 +54,7 @@ const Collect = struct { }; fn env(name: [:0]const u8, default: []const u8) []const u8 { - return if (std.c.getenv(name)) |v| std.mem.span(v) else default; + return if (platform.c.getenv(name)) |v| std.mem.span(v) else default; } fn hash(t: []const u32) u64 { @@ -124,13 +125,13 @@ pub fn main(init: std.process.Init) !void { const e = try glm.engine.Engine.load(gpa, args[1], cap); defer e.deinit(); std.debug.print("loaded in {d:.1} s: {d} of {d} layers, routed experts {d}-{d}, {d:.1} GB of weights, MTP head {s}, {d:.2} GB of caches and buffers, prompt chunks of {d} rows\n", .{ e.load_seconds, e.c.run, e.c.layers, e.c.own[0], e.c.own[1] - 1, @as(f64, @floatFromInt(e.w.bytes)) / 1e9, if (e.hasMtp()) "yes" else "no", @as(f64, @floatFromInt(e.arena.bytes)) / 1e9, if (e.pr) |p| p.rows else 0 }); - const strict = std.c.getenv("GLM_REF_STRICT") != null; + const strict = platform.c.getenv("GLM_REF_STRICT") != null; const eos = e.c.eos[0..e.c.eos_n]; - if (std.c.getenv("GLM_CAPTURE")) |path| { // the first prompt's first window, sublayer by sublayer (glm_ref.py --capture) + if (platform.c.getenv("GLM_CAPTURE")) |path| { // the first prompt's first window, sublayer by sublayer (glm_ref.py --capture) const first = doc.object.get("prompts").?.array.items[0]; try e.capture(try ints(arena, first.object.get("ids").?), std.mem.span(path)); } - if (std.c.getenv("GLM_TRACE")) |spec| { // NAME:STEPS:PATH, glm_ref.py --trace NAME STEPS PATH's twin + if (platform.c.getenv("GLM_TRACE")) |spec| { // NAME:STEPS:PATH, glm_ref.py --trace NAME STEPS PATH's twin var parts = std.mem.splitScalar(u8, std.mem.span(spec), ':'); const want = parts.next() orelse return error.BadTrace; const steps = try std.fmt.parseInt(usize, parts.next() orelse return error.BadTrace, 10); @@ -139,18 +140,18 @@ pub fn main(init: std.process.Init) !void { try e.trace(try ints(arena, p.object.get("ids").?), try ints(arena, p.object.get("expect").?), steps, out_path); }; } - if (std.c.getenv("GLM_FORCED") != null) for (doc.object.get("prompts").?.array.items) |p| { + if (platform.c.getenv("GLM_FORCED") != null) for (doc.object.get("prompts").?.array.items) |p| { const want = try ints(arena, p.object.get("expect") orelse continue); const r = try e.forced(try ints(arena, p.object.get("ids").?), want); std.debug.print("{s} forced: {d} of {d} picks equal the reference's", .{ p.object.get("name").?.string, r.same, want.len }); if (r.first) |i| std.debug.print(", first differs at {d}\n", .{i}) else std.debug.print("\n", .{}); }; var failures: usize = 0; - const vs: ?std.json.Value = if (std.c.getenv("GLM_VS")) |path| blk: { + const vs: ?std.json.Value = if (platform.c.getenv("GLM_VS")) |path| blk: { const f = try mtl.MappedFile.open(path); break :blk try std.json.parseFromSliceLeaky(std.json.Value, arena, f.bytes[0..f.size], .{}); } else null; - const trace_prefix: ?[*:0]const u8 = std.c.getenv("GLM_TRACE_LAST"); + const trace_prefix: ?[*:0]const u8 = platform.c.getenv("GLM_TRACE_LAST"); const trace_bytes = (2 + 4 * @as(usize, e.c.run)) * e.c.hidden * 2; if (trace_prefix != null) e.trace_last = try e.arena.buffer(trace_bytes); var saved: std.ArrayList(u8) = .empty; @@ -168,10 +169,10 @@ pub fn main(init: std.process.Init) !void { if (trace_prefix != null and run == 0 and d == depths.items[0]) col.trace_from = e.trace_last.?.addr()[0..trace_bytes]; var margins: std.ArrayList(f32) = .empty; defer margins.deinit(gpa); - const want_margins = run == 0 and d == depths.items[0] and (std.c.getenv("GLM_MARGINS") != null or std.c.getenv("GLM_MARGINS_VS") != null); + const want_margins = run == 0 and d == depths.items[0] and (platform.c.getenv("GLM_MARGINS") != null or platform.c.getenv("GLM_MARGINS_VS") != null); e.margins = if (want_margins) &margins else null; defer e.margins = null; - const want_logits = run == 0 and d == depths.items[0] and (std.c.getenv("GLM_LOGITS") != null or std.c.getenv("GLM_LOGITS_VS") != null); + const want_logits = run == 0 and d == depths.items[0] and (platform.c.getenv("GLM_LOGITS") != null or platform.c.getenv("GLM_LOGITS_VS") != null); if (want_logits) col.logits_from = @as([*]const u16, @ptrCast(@alignCast(e.sc.logits.addr())))[0..e.c.vocab]; const r = try e.generate(ids, max, eos, d, .{ .ctx = &col, .prefilled = Collect.prefilled, .tokens = Collect.tokens, .cancelled = Collect.cancelled }); const toks = col.toks.items; @@ -182,7 +183,7 @@ pub fn main(init: std.process.Init) !void { const rn: f64 = @floatFromInt(r.rounds); std.debug.print(" a round: {d:.2} ms wall, {d:.2} ms GPU, {d:.2} ms encoding, {d:.2} ms GPU idle between rounds; {d:.2} tokens\n", .{ r.decode_seconds * 1e3 / rn, r.gpu_seconds * 1e3 / rn, r.encode_seconds * 1e3 / rn, r.gap_seconds * 1e3 / @max(rn - 1, 1), @as(f64, @floatFromInt(toks.len -| 1)) / rn }); if (r.copy_rounds > 0) std.debug.print(" copy drafts: {d} rounds, {d} drafts kept ({d:.2} a copy round)\n", .{ r.copy_rounds, r.copy_accepted, @as(f64, @floatFromInt(r.copy_accepted)) / @as(f64, @floatFromInt(r.copy_rounds)) }); - if (std.c.getenv("GLM_RANKS") != null) { + if (platform.c.getenv("GLM_RANKS") != null) { std.debug.print(" draft ranks (kept: all / miss at 2nd / 3rd / later):", .{}); for (r.draft_ranks, 0..) |row, m| if (row[0] + row[1] + row[2] + row[3] > 0) std.debug.print(" m{d} {d}/{d}/{d}/{d}", .{ m, row[0], row[1], row[2], row[3] }); std.debug.print("\n", .{}); @@ -192,17 +193,17 @@ pub fn main(init: std.process.Init) !void { failures += 1; std.debug.print(" DIFFERS from depth {d} at token {d}\n", .{ depths.items[0], at }); } else std.debug.print(" equal to depth {d}\n", .{depths.items[0]}); - if (want_margins) if (std.c.getenv("GLM_MARGINS")) |prefix| { + if (want_margins) if (platform.c.getenv("GLM_MARGINS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.margins", .{ prefix, name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(std.mem.sliceAsBytes(margins.items).ptr, 4, margins.items.len, file) != margins.items.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(std.mem.sliceAsBytes(margins.items).ptr, 4, margins.items.len, file) != margins.items.len) return error.WriteFailed; }; if (vs) |other| if (other.object.get(name)) |theirs| { const want = try ints(arena, theirs); if (firstDiff(want, toks)) |at| { std.debug.print(" vs GLM_VS: first token {s}, first difference at token {d} of {d}\n", .{ if (at == 0) "DIFFERS" else "equal", at, @min(want.len, toks.len) }); - if (want_margins) if (std.c.getenv("GLM_MARGINS_VS")) |prefix| { + if (want_margins) if (platform.c.getenv("GLM_MARGINS_VS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.margins", .{ prefix, name }, 0); if (mtl.MappedFile.open(path)) |f| { defer f.deinit(); @@ -216,18 +217,18 @@ pub fn main(init: std.process.Init) !void { }; if (col.trace.len > 0) { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.trace", .{ std.mem.span(trace_prefix.?), name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(col.trace.ptr, 1, col.trace.len, file) != col.trace.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(col.trace.ptr, 1, col.trace.len, file) != col.trace.len) return error.WriteFailed; } if (want_logits and col.logits.len == 0) return error.NoLogits; - if (want_logits) if (std.c.getenv("GLM_LOGITS")) |prefix| { + if (want_logits) if (platform.c.getenv("GLM_LOGITS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(std.mem.sliceAsBytes(col.logits).ptr, 2, col.logits.len, file) != col.logits.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(std.mem.sliceAsBytes(col.logits).ptr, 2, col.logits.len, file) != col.logits.len) return error.WriteFailed; }; - if (want_logits) if (std.c.getenv("GLM_LOGITS_VS")) |prefix| try compareLogits(gpa, col.logits, try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0)); + if (want_logits) if (platform.c.getenv("GLM_LOGITS_VS")) |prefix| try compareLogits(gpa, col.logits, try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0)); if (run == 0 and d == depths.items[0]) { if (saved.items.len > 1) try saved.append(arena, ','); try saved.print(arena, "\"{s}\":[", .{name}); @@ -241,7 +242,7 @@ pub fn main(init: std.process.Init) !void { } else std.debug.print(" reference: {d} of {d} tokens equal\n", .{ @min(want.len, toks.len), want.len }); } }; - if (std.c.getenv("GLM_CANCEL_TEST") != null) { // a reply cancelled after its first prompt window, then a fresh one + if (platform.c.getenv("GLM_CANCEL_TEST") != null) { // a reply cancelled after its first prompt window, then a fresh one var cut: Collect = .{ .gpa = gpa, .cancel_at = 1 }; defer cut.toks.deinit(gpa); const rc = try e.generate(ids, max, eos, depths.items[0], .{ .ctx = &cut, .prefilled = Collect.prefilled, .tokens = Collect.tokens, .cancelled = Collect.cancelled }); @@ -253,23 +254,23 @@ pub fn main(init: std.process.Init) !void { std.debug.print(" cancel test: cancelled reply ended {t} after {d} tokens; the next fresh reply {s} the plain one\n", .{ rc.reason, cut.toks.items.len, if (same) "equals" else "DIFFERS from" }); } } - if (std.c.getenv("GLM_PROFILE")) |spec| { + if (platform.c.getenv("GLM_PROFILE")) |spec| { const reps = try std.fmt.parseInt(usize, env("GLM_PROFILE_REPS", "12"), 10); var pit = std.mem.tokenizeScalar(u8, std.mem.span(spec), ','); - const only = std.c.getenv("GLM_PROFILE_ONLY") != null; // each class alone instead of each left out - const parts = std.c.getenv("GLM_PROFILE_PARTS") != null; // by launch instead of by class + const only = platform.c.getenv("GLM_PROFILE_ONLY") != null; // each class alone instead of each left out + const parts = platform.c.getenv("GLM_PROFILE_PARTS") != null; // by launch instead of by class while (pit.next()) |dv| try e.profile(try std.fmt.parseInt(u32, dv, 10), reps, only, parts); } - if (std.c.getenv("GLM_MM_CHECK") != null) try e.checkMatmul(try ints(arena, doc.object.get("prompts").?.array.items[0].object.get("ids").?)); - if (std.c.getenv("GLM_PROMPT_PROFILE")) |v| { // each prompt's first chunk by class (GLM_PROFILE_ONLY: each alone) + if (platform.c.getenv("GLM_MM_CHECK") != null) try e.checkMatmul(try ints(arena, doc.object.get("prompts").?.array.items[0].object.get("ids").?)); + if (platform.c.getenv("GLM_PROMPT_PROFILE")) |v| { // each prompt's first chunk by class (GLM_PROFILE_ONLY: each alone) const reps = try std.fmt.parseInt(usize, std.mem.span(v), 10); - for (doc.object.get("prompts").?.array.items) |p| try e.profilePrompt(try ints(arena, p.object.get("ids").?), reps, std.c.getenv("GLM_PROFILE_ONLY") != null, std.c.getenv("GLM_PROFILE_PARTS") != null); + for (doc.object.get("prompts").?.array.items) |p| try e.profilePrompt(try ints(arena, p.object.get("ids").?), reps, platform.c.getenv("GLM_PROFILE_ONLY") != null, platform.c.getenv("GLM_PROFILE_PARTS") != null); } try saved.append(arena, '}'); - if (std.c.getenv("GLM_OUT")) |path| { - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(saved.items.ptr, 1, saved.items.len, file) != saved.items.len) return error.WriteFailed; + if (platform.c.getenv("GLM_OUT")) |path| { + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(saved.items.ptr, 1, saved.items.len, file) != saved.items.len) return error.WriteFailed; } if (failures > 0) { std.debug.print("{d} replies differ from the plain ones{s}\n", .{ failures, if (strict) " or the references" else "" }); diff --git a/zig/tests/linktest.zig b/zig/tests/linktest.zig index dfdad16efe..bf251fd1ec 100644 --- a/zig/tests/linktest.zig +++ b/zig/tests/linktest.zig @@ -1,5 +1,6 @@ //! Link qualification over the system's verbs; `run` swaps connection details on stdin and stdout and opens no socket. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const suite = @import("linktest_suite.zig"); const probe = @import("linktest_probe.zig"); @@ -30,7 +31,7 @@ fn readLine(buf: []u8) ![]const u8 { var n: usize = 0; while (n < buf.len) { var c: [1]u8 = undefined; - if (std.c.read(0, &c, 1) != 1) return error.EndOfInput; + if (platform.c.read(0, &c, 1) != 1) return error.EndOfInput; if (c[0] == '\n') return buf[0..n]; buf[n] = c[0]; n += 1; @@ -61,7 +62,7 @@ fn parseHello(line: []const u8) ![2]vqp.Info { /// Exit when stdin closes: the run ends with the session that started it. fn tether() void { var c: [1]u8 = undefined; - while (std.c.read(0, &c, 1) == 1) {} + while (platform.c.read(0, &c, 1) == 1) {} out("R stdin_closed", .{}); std.process.exit(4); } @@ -79,7 +80,7 @@ fn mapped(len: usize) ![]align(vl.page) u8 { /// Two queue pairs on one device: collectives on one, raw measurements on the other, so their receives never mix. fn run(gpa: std.mem.Allocator, v: *const Verbs, a: Args) !void { _ = try std.Thread.spawn(.{}, watchdog, .{a.deadline_s}); - out("PID {d}", .{std.c.getpid()}); + out("PID {d}", .{platform.c.getpid()}); var peer: vl.Peer = .{ .ep = try vqp.Endpoint.open(v, a.dev, a.kind, 4095, a.gid) }; defer peer.ep.close(); var raw = try peer.ep.sibling(a.kind, 4095); @@ -138,7 +139,7 @@ pub fn main(init: std.process.Init) !void { if (!std.mem.eql(u8, args[1], "run")) std.process.exit(2); const a = try parse(args); run(init.gpa, &v, a) catch |err| { - out("R error what={s} errno={d}", .{ @errorName(err), std.c._errno().* }); + out("R error what={s} errno={d}", .{ @errorName(err), platform.c._errno().* }); std.process.exit(1); }; } diff --git a/zig/tests/linktest_probe.zig b/zig/tests/linktest_probe.zig index ce3402075a..a5c6af2541 100644 --- a/zig/tests/linktest_probe.zig +++ b/zig/tests/linktest_probe.zig @@ -1,5 +1,6 @@ //! What a verbs device reports and what it actually allows: registration sizes and counts, queue pair types and counts. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const abi = fabric.verbs_abi; const Verbs = fabric.verbs.Verbs; @@ -8,7 +9,7 @@ const out = @import("linktest_suite.zig").out; const mib = 1 << 20; fn errno() c_int { - return std.c._errno().*; + return platform.c._errno().*; } pub fn device(v: *const Verbs, dev: []const u8) !void { diff --git a/zig/tests/linktest_suite.zig b/zig/tests/linktest_suite.zig index 37b998b9f7..f356e6e519 100644 --- a/zig/tests/linktest_suite.zig +++ b/zig/tests/linktest_suite.zig @@ -1,5 +1,6 @@ //! The link measurements over Thunderbolt's two-sided UC: opcodes, ping-pong, bandwidth, integrity, barrier, all-reduce. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const collective = fabric.collective; const vqp = fabric.verbs_qp; @@ -53,7 +54,7 @@ pub const Ctx = struct { pub fn out(comptime fmt: []const u8, args: anytype) void { var buf: [4096]u8 = undefined; const line = std.fmt.bufPrint(&buf, fmt ++ "\n", args) catch return; - _ = std.c.write(1, line.ptr, line.len); + _ = platform.c.write(1, line.ptr, line.len); } fn us(ns: u64) f64 { diff --git a/zig/tests/mm_bench.zig b/zig/tests/mm_bench.zig index 1ec779f916..76f6402c12 100644 --- a/zig/tests/mm_bench.zig +++ b/zig/tests/mm_bench.zig @@ -1,5 +1,6 @@ //! Prompt matmul kernels compiled from .metal files at run time: chained speed at one shape, bits against the first, error against fp64. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -17,7 +18,7 @@ const Strides = extern struct { x_row: i32 = 0, y_row: i32 = 0, sums_row: i32 = const Kernel = struct { path: []const u8, bm: u32 = 64, bn: u32 = 64, threads: u32 = 128, entry: []const u8 = "tf_affine_mm", lib: mtl.Library = undefined, pipe: mtl.Pipeline = undefined, sums_pipe: mtl.Pipeline = undefined, sums: mtl.Buffer = undefined, y: mtl.Buffer = undefined, ms: [64]f64 = undefined, sums_ms: f64 = 0 }; fn envInt(name: [:0]const u8, default: usize) usize { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; } @@ -55,7 +56,7 @@ pub fn main(init: std.process.Init) !void { const trials = @min(envInt("MM_TRIALS", 5), 64); const copies = envInt("MM_COPIES", 4); const check_rows = @min(envInt("MM_CHECK_ROWS", 16), M); - const out_f32 = std.c.getenv("MM_F32") != null; + const out_f32 = platform.c.getenv("MM_F32") != null; const ob: usize = if (out_f32) 4 else 2; const kernels = try arena.alloc(Kernel, args.len - 4); for (kernels, args[4..]) |*k, s| k.* = try parseKernel(s); diff --git a/zig/tests/moe_bench.zig b/zig/tests/moe_bench.zig index 6e381ff8d5..f611ef2575 100644 --- a/zig/tests/moe_bench.zig +++ b/zig/tests/moe_bench.zig @@ -1,5 +1,6 @@ //! Routed-expert prompt kernels from .metal files compiled at run time: a chunk's skewed top-k routing, chained speed, bits against the first. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -16,12 +17,12 @@ const Args = extern struct { rows: i32, n: i32, k: i32, experts: i32 }; const Kernel = struct { path: []const u8, bm: u32 = 64, entry: []const u8 = "", lib: mtl.Library = undefined, pipe: mtl.Pipeline = undefined, y: mtl.Buffer = undefined, ms: [64]f64 = undefined }; fn envInt(name: [:0]const u8, default: usize) usize { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; } fn envFloat(name: [:0]const u8, default: f64) f64 { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseFloat(f64, std.mem.span(v)) catch default; } @@ -37,7 +38,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("usage: {s}\n", .{usage}); std.process.exit(2); } - const down = if (std.c.getenv("MOE_MODE")) |v| std.mem.eql(u8, std.mem.span(v), "down") else false; + const down = if (platform.c.getenv("MOE_MODE")) |v| std.mem.eql(u8, std.mem.span(v), "down") else false; const T = try std.fmt.parseInt(u32, args[1], 10); const E: u32 = @intCast(envInt("MOE_EXPERTS", 288)); const topk: u32 = @intCast(envInt("MOE_TOPK", 8)); @@ -65,7 +66,7 @@ pub fn main(init: std.process.Init) !void { total += w.*; } const pick = try arena.alloc(u32, n); - const exact = std.c.getenv("MOE_EXACT") != null; // pairs dealt round-robin: every expert the same row count + const exact = platform.c.getenv("MOE_EXACT") != null; // pairs dealt round-robin: every expert the same row count if (exact) for (pick, 0..) |*q, i| { q.* = @intCast((i % topk + (i / topk) * topk) % E); }; diff --git a/zig/tests/nemotron_prefill_check.zig b/zig/tests/nemotron_prefill_check.zig index 5bca853939..f1b91302aa 100644 --- a/zig/tests/nemotron_prefill_check.zig +++ b/zig/tests/nemotron_prefill_check.zig @@ -1,5 +1,6 @@ //! One prompt chunk on the prefill kernels against tools/zig/prefill_dump.py's rows: embedding, every layer, final norm. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -17,7 +18,7 @@ pub fn main(init: std.process.Init) !void { var it = std.mem.tokenizeScalar(u8, args[2], ','); while (it.next()) |t| try list.append(arena, try std.fmt.parseInt(u32, t, 10)); // TF_PREFILL_PREFIX=N: the first N ids fill the caches as a chunk first; the rest is the checked chunk - const before: usize = if (std.c.getenv("TF_PREFILL_PREFIX")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 0; + const before: usize = if (platform.c.getenv("TF_PREFILL_PREFIX")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 0; const all = list.items; const ids = all[before..]; const n = ids.len; diff --git a/zig/tests/pocket_fit_check.zig b/zig/tests/pocket_fit_check.zig index 7f7bfaf2f5..0fe514684a 100644 --- a/zig/tests/pocket_fit_check.zig +++ b/zig/tests/pocket_fit_check.zig @@ -1,5 +1,6 @@ //! A dumped pocket lesson's read rows fitted as the learner fits them, timed; each fit's numbers and weights written. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const pocket_fit = tf.nemotron.pocket_fit; @@ -57,7 +58,7 @@ pub fn main(init: std.process.Init) !void { const us = try gpa.alloc(f32, jobs.items.len * d); const fits = try gpa.alloc(Fit, jobs.items.len); // TF_FIT_THREADS: that many threads take the fits in turn (default: a thread a fit, as the learner runs them) - const want = if (std.c.getenv("TF_FIT_THREADS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else jobs.items.len; + const want = if (platform.c.getenv("TF_FIT_THREADS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else jobs.items.len; const threads = try gpa.alloc(std.Thread, @min(@max(want, 1), jobs.items.len)); var next = std.atomic.Value(usize).init(0); const work: Work = .{ .gpa = gpa, .rows = rows, .kinds = all_kinds, .m = m, .slots = slots, .d = d, .jobs = jobs.items, .us = us, .fits = fits, .next = &next }; diff --git a/zig/tests/qwen27/dflash_graph.zig b/zig/tests/qwen27/dflash_graph.zig index bbd9db4865..e382276faf 100644 --- a/zig/tests/qwen27/dflash_graph.zig +++ b/zig/tests/qwen27/dflash_graph.zig @@ -1,5 +1,6 @@ //! Generated draft tensors exercise the real callback backend without checkpoint files. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const q = tf.qwen27; @@ -47,19 +48,19 @@ pub fn createQ4(a: std.mem.Allocator, target: *q.model.Model) !*Model { } try header.writer.writeByte('}'); var path_buf: [96]u8 = undefined; - const path = try std.fmt.bufPrintSentinel(&path_buf, "/tmp/tf-q27-q4-draft-{d}.safetensors", .{std.c.getpid()}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const path = try std.fmt.bufPrintSentinel(&path_buf, "/tmp/tf-q27-q4-draft-{d}.safetensors", .{platform.c.getpid()}, 0); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.Create; - defer _ = std.c.unlink(path); + defer _ = platform.c.unlink(path); var head: [8]u8 = undefined; std.mem.writeInt(u64, &head, header.written().len, .little); for ([_][]const u8{ &head, header.written(), std.mem.sliceAsBytes(data.items) }) |part| { - if (std.c.write(fd, part.ptr, part.len) != @as(isize, @intCast(part.len))) { - _ = std.c.close(fd); + if (platform.c.write(fd, part.ptr, part.len) != @as(isize, @intCast(part.len))) { + _ = platform.c.close(fd); return error.ShortWrite; } } - _ = std.c.close(fd); + _ = platform.c.close(fd); return Model.open(a, target, config, &.{path}, .prepared_q4_reference); } diff --git a/zig/tests/qwen27/dflash_run.zig b/zig/tests/qwen27/dflash_run.zig index 06660eed7e..2d75b9d491 100644 --- a/zig/tests/qwen27/dflash_run.zig +++ b/zig/tests/qwen27/dflash_run.zig @@ -1,5 +1,6 @@ //! The draft prototype on real weights: target tokens and final state against ordinary native generation. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("qwen27"); const profile = @import("core").gpu_profile; @@ -71,7 +72,7 @@ pub fn main(init: std.process.Init) !void { try generation.prefill(prompt, chunk); const prompt_ns = start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds(); var trace = profile.Trace{ .queue = target.queue }; - const profiled = std.c.getenv("QWEN27_PROFILE") != null; + const profiled = platform.c.getenv("QWEN27_PROFILE") != null; if (profiled) { runner.profile(&trace); trace.enabled = true; diff --git a/zig/tests/qwen27/dflash_served.zig b/zig/tests/qwen27/dflash_served.zig index ca9b3c2bfb..6955e546f0 100644 --- a/zig/tests/qwen27/dflash_served.zig +++ b/zig/tests/qwen27/dflash_served.zig @@ -1,5 +1,6 @@ //! Generated native draft and plain requests retain CLI tokens and final state through the served host. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("tensorfold").qwen27; const api = @import("engine_api"); @@ -73,7 +74,7 @@ const Box = struct { }; /// Drafted requests on the lane core (its own round counts), by the host's rule: TF_QWEN27_LANES, else no tensor units. fn lanesMode(h: *host.Host) bool { - if (std.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); + if (platform.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); return !h.model.device.tensorUnits(); } fn stop(_: *anyopaque, ids: []const u32) bool { diff --git a/zig/tests/qwen27/reuse.zig b/zig/tests/qwen27/reuse.zig index 46ce1ed90d..e41b5d8c04 100644 --- a/zig/tests/qwen27/reuse.zig +++ b/zig/tests/qwen27/reuse.zig @@ -1,5 +1,6 @@ //! Qwen27 prompt reuse through the served host against fresh prompt passes: each prompt's state and reply, bit for bit. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("tensorfold").qwen27; const api = @import("engine_api"); @@ -90,7 +91,7 @@ const Run = struct { /// Drafted requests on the lane core, by the served host's rule: TF_QWEN27_LANES, else a chip without tensor units. fn lanesMode(h: *host.Host) bool { - if (std.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); + if (platform.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); return !h.model.device.tensorUnits(); } diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig new file mode 100644 index 0000000000..847240302e --- /dev/null +++ b/zig/tests/qwen35_cuda_main.zig @@ -0,0 +1,189 @@ +//! `tf-qwen35-cuda `: bf16 logit rows for a reference, or picks. +const std = @import("std"); +const cuda = @import("cuda"); +const q = @import("qwen35"); +const E = q.engine; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const args = try init.minimal.args.toSlice(a); + if (args.len != 5) { + std.debug.print("usage: tf-qwen35-cuda \n", .{}); + return 2; + } + var ids: std.ArrayList(u32) = .empty; + var it = std.mem.tokenizeAny(u8, args[3], ", "); + while (it.next()) |t| try ids.append(a, try std.fmt.parseInt(u32, t, 10)); + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + const e = try E.Engine.init(init.gpa, init.io, &ctx, args[1]); + defer e.deinit(); + if (std.mem.eql(u8, args[2], "mtp")) return mtpCheck(init.gpa, init.io, e, args[1], ids.items); + // TF_Q35_PREFILL=N starts the stream after N cycled ids; TF_Q35_BESIDE=N steps an N-token stream beside it + const prefill = try envCount("TF_Q35_PREFILL"); + const beside = try envCount("TF_Q35_BESIDE"); + var cache = try E.Cache.init(init.gpa, e, 4096 + prefill); + defer cache.deinit(init.gpa); + try fill(e, &cache, ids.items, prefill); + var other = try E.Cache.init(init.gpa, e, 4096 + beside); + defer other.deinit(init.gpa); + try fill(e, &other, ids.items, beside); + var f: Fwd = .{ .e = e, .other = if (beside > 0) &other else null, .ids = ids.items }; + var out: std.ArrayList(u8) = .empty; + if (std.mem.eql(u8, args[2], "step")) { + var first: i64 = 0; + var rest: i64 = 0; + for (ids.items, 0..) |id, i| { + const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + try f.run(.{ .cache = &cache, .rows = 1 }, &.{id}); + const dt: i64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); + if (i == 0) first = dt else rest += dt; + @memcpy(cache.logits, e.rowLogits(0)); + try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); + } + const n: f64 = @floatFromInt(ids.items.len - 1); + std.debug.print("{d} steps: the first {d:.1} ms, then {d:.1} ms a step ({d:.1} tok/s)\n", .{ ids.items.len, @as(f64, @floatFromInt(first)) / 1e6, @as(f64, @floatFromInt(rest)) / 1e6 / n, n * 1e9 / @as(f64, @floatFromInt(rest)) }); + } else if (std.mem.eql(u8, args[2], "window")) { + // windows of 1, 2, 3, ... rows: every row's logits, which must equal step's bit for bit + var at: usize = 0; + var w: usize = 1; + while (at < ids.items.len) : (w = w % e.window_rows + 1) { + const n = @min(w, ids.items.len - at); + try f.run(.{ .cache = &cache, .rows = n }, ids.items[at..][0..n]); + for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); + at += n; + } + } else if (std.mem.eql(u8, args[2], "tree")) { + // tree windows of growing widths, junk branches beside the kept true one, whose logits must equal step's bits + var at: usize = 0; + var w: usize = 2; + while (at < ids.items.len) : (w = w % e.window_rows + 1) { + var toks: [E.most_window]u32 = undefined; + var parents: [E.most_window]i32 = undefined; + var path: [E.most_window]u32 = undefined; + toks[0] = ids.items[at]; + parents[0] = -1; + path[0] = 0; + var n: usize = 1; + var kept: usize = 1; + var last: i32 = 0; + var depth: usize = 1; + while (n + 2 <= w and at + depth < ids.items.len) : (depth += 1) { + toks[n] = (ids.items[at + depth] + 7) % @as(u32, @intCast(e.g.vocab)); + parents[n] = last; + toks[n + 1] = ids.items[at + depth]; + parents[n + 1] = last; + path[kept] = @intCast(n + 1); + last = @intCast(n + 1); + n += 2; + kept += 1; + } + if (n < w and n >= 3) { + toks[n] = 1; + parents[n] = 1; + n += 1; + } + try f.run(.{ .cache = &cache, .rows = n, .parents = parents[0..n] }, toks[0..n]); + for (path[0..kept]) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); + if (n > 1) try cache.keep(e, path[0..kept]); + at += kept; + } + } else if (std.mem.eql(u8, args[2], "probe")) { + // the first id through layers 0..N (N = the out argument): that layer's routing and rows, for a reference + const layer = try std.fmt.parseInt(usize, args[4], 10); + const all = e.layers; + e.layers = all[0 .. layer + 1]; + e.graphs.clearRetainingCapacity(); // the warm-up captured every layer; this forward captures the cut + try e.forward(&.{.{ .cache = &cache, .rows = 1 }}, ids.items[0..1], .logits); + e.layers = all; + const g = e.g; + if (!g.moe()) return error.NotRouted; + const count = g.experts + 1; + const logits_f = try a.alloc(f32, count); + const picks = try a.alloc(i32, g.slots()); + const wts = try a.alloc(f32, g.slots()); + var xrow: [8]u16 = undefined; + var rrow: [8]u16 = undefined; + try e.arena.download(e.rlogits - e.arena.ptr, std.mem.sliceAsBytes(logits_f)); + try e.arena.download(e.picks - e.arena.ptr, std.mem.sliceAsBytes(picks)); + try e.arena.download(e.wts - e.arena.ptr, std.mem.sliceAsBytes(wts)); + try e.arena.download(e.x - e.arena.ptr, std.mem.sliceAsBytes(&xrow)); + try e.arena.download(e.r - e.arena.ptr, std.mem.sliceAsBytes(&rrow)); + std.debug.print("probe layer {d} x[:8]", .{layer}); + for (xrow) |h| std.debug.print(" {d:.5}", .{@as(f32, @bitCast(@as(u32, h) << 16))}); + std.debug.print("\nprobe top logits", .{}); + for (0..10) |_| { + var best: usize = 0; + for (logits_f[0 .. count - 1], 0..) |v, i| if (v > logits_f[best]) { + best = i; + }; + std.debug.print(" ({d}, {d:.5})", .{ best, logits_f[best] }); + logits_f[best] = -std.math.inf(f32); + } + std.debug.print(" shared gate {d:.5}\nprobe picks", .{logits_f[count - 1]}); + for (picks) |pk| std.debug.print(" {d}", .{pk}); + std.debug.print(" weights", .{}); + for (wts) |w| std.debug.print(" {d:.5}", .{w}); + std.debug.print("\nprobe r[:8]", .{}); + for (rrow) |h| std.debug.print(" {d:.5}", .{@as(f32, @bitCast(@as(u32, h) << 16))}); + std.debug.print("\n", .{}); + return 0; + } else { + var at: usize = 0; + while (at < ids.items.len) { + const n = @min(E.max_rows, ids.items.len - at); + try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .last); + at += n; + } + try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); + } + try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[4], .data = out.items }); + return 0; +} + +/// One forward keeping logits: the stream's part first (its rows' logits come first), then the second stream's step. +const Fwd = struct { + e: *E.Engine, + other: ?*E.Cache, + ids: []const u32, + n: usize = 0, + + fn run(f: *Fwd, part: E.Part, toks: []const u32) !void { + const o = f.other orelse return f.e.forward(&.{part}, toks, .logits); + var all: [E.most_window + 1]u32 = undefined; + @memcpy(all[0..toks.len], toks); + all[toks.len] = f.ids[f.n % f.ids.len]; + f.n += 1; + try f.e.forward(&.{ part, .{ .cache = o, .rows = 1 } }, all[0 .. toks.len + 1], .logits); + } +}; + +/// The MTP drafter's tree regressions (cuda_mtp_check.zig) after a prompt over the ids fills the target's tap ring. +fn mtpCheck(gpa: std.mem.Allocator, io: std.Io, e: *E.Engine, dir: []const u8, ids: []const u32) !u8 { + const d = try q.mtp.Mtp.init(gpa, io, e, dir, 4096); + defer d.deinit(); + try e.setTaps(d.taps()); + var cache = try E.Cache.init(gpa, e, 4096); + defer cache.deinit(gpa); + try fill(e, &cache, ids, ids.len); + return if (try q.mtp_check.run(gpa, d, cache.tap.?.ptr, ids)) 0 else 1; +} + +fn envCount(name: [:0]const u8) !usize { + const v = std.c.getenv(name) orelse return 0; + return std.fmt.parseInt(usize, std.mem.span(v), 10); +} + +/// `n` positions of the ids cycled, through prompt forwards. +fn fill(e: *E.Engine, cache: *E.Cache, ids: []const u32, n: usize) !void { + var buf: [E.max_rows]u32 = undefined; + var at: usize = 0; + while (at < n) { + const k = @min(E.max_rows, n - at); + for (buf[0..k], at..) |*t, i| t.* = ids[i % ids.len]; + try e.forward(&.{.{ .cache = cache, .rows = k }}, buf[0..k], .last); + at += k; + } +} diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index 184c401968..73d607350f 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -1,6 +1,8 @@ //! Frozen server goldens for `zig build test-golden`, with the ten differences that are still known. const std = @import("std"); -const posix = std.posix; +const builtin = @import("builtin"); +const platform = @import("sys"); +const posix = platform.posix; const Io = std.Io; const wire = @import("wire.zig"); const Reply = wire.Reply; @@ -21,6 +23,38 @@ const known_diffs = [_]Known{ .{ .group = "anthropic", .name = "tool-use-stream" }, }; +/// Why two request logs differ: their line counts and the first line that does, each cut to 100 bytes. +fn logDiff(a: std.mem.Allocator, want: []const u8, got: []const u8) ![]const []const u8 { + var out: std.ArrayList([]const u8) = .empty; + var wl = std.mem.splitScalar(u8, want, '\n'); + var gl = std.mem.splitScalar(u8, got, '\n'); + var n: usize = 0; + while (true) : (n += 1) { + const w = wl.next(); + const g = gl.next(); + if (w == null and g == null) break; + if (w != null and g != null and std.mem.eql(u8, w.?, g.?)) continue; + const ws = w orelse ""; + const gs = g orelse ""; + try out.append(a, try std.fmt.allocPrint(a, "line {d}: golden [{s}] zig [{s}]", .{ n + 1, ws[0..@min(ws.len, 100)], gs[0..@min(gs.len, 100)] })); + break; + } + try out.append(a, try std.fmt.allocPrint(a, "TENSORFOLD_REQUEST_LOG: golden {d} lines, zig {d} lines", .{ std.mem.count(u8, want, "\n"), std.mem.count(u8, got, "\n") })); + return out.items; +} + +/// Reads a spawned server's stdout to its end, so its log lines never fill the pipe (4 KiB on Windows). +fn drainStdout(io: Io, out: Io.File) void { + var buf: [4096]u8 = undefined; + var reader = out.reader(io, &buf); + while (true) { + _ = reader.interface.discardShort(buf.len) catch return; + } +} + +/// The parent's SystemRoot, which a Windows child needs for Winsock; null elsewhere. +var system_root: ?[]const u8 = null; + fn isKnown(group: []const u8, name: []const u8) bool { for (known_diffs) |k| { if (std.mem.eql(u8, k.group, group) and std.mem.eql(u8, k.name, name)) return true; @@ -354,6 +388,7 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con } else try argv.append(a, arg); } var map = std.process.Environ.Map.init(a); + if (system_root) |root| try map.put("SystemRoot", root); // Winsock refuses to start without it try map.put("TENSORFOLD_NO_LIVE", "1"); try map.put("TF_FAKE_CONTEXT", context_tokens); var it = env.iterator(); @@ -379,7 +414,7 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con try fw.interface.writeAll(content); try fw.interface.flush(); const kpz = try std.fmt.allocPrintSentinel(a, "{s}", .{kp}, 0); - _ = std.c.fchmodat(posix.AT.FDCWD, kpz, 0o600, 0); + _ = platform.c.fchmodat(posix.AT.FDCWD, kpz, 0o600, 0); } } const child = try a.create(std.process.Child); @@ -397,12 +432,16 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con const line = reader.interface.takeDelimiterExclusive('\n') catch return error.NoPortLine; if (std.mem.startsWith(u8, line, "PORT ")) { const port = try std.fmt.parseInt(u16, std.mem.trim(u8, line[5..], " \r"), 10); + // the server's log lines keep coming on stdout: drained, or a full pipe blocks the server + const drain = try std.Thread.spawn(.{}, drainStdout, .{ io, child.stdout.? }); + drain.detach(); return .{ .child = child, .port = port, .log_path = log_path, .key_path = key_path }; } } } pub fn main(init: std.process.Init) !void { + if (builtin.os.tag == .windows) system_root = init.environ_map.get("SystemRoot"); const gpa = init.gpa; const io = init.io; const args = try init.minimal.args.toSlice(init.arena.allocator()); @@ -411,7 +450,7 @@ pub fn main(init: std.process.Init) !void { const cases_path = args[2]; const golden_dir = args[3]; const fixtures = args[4]; - const tmp_dir = try std.fmt.allocPrint(init.arena.allocator(), ".zig-tmp-golden-{d}", .{std.c.getpid()}); + const tmp_dir = try std.fmt.allocPrint(init.arena.allocator(), ".zig-tmp-golden-{d}", .{platform.c.getpid()}); Io.Dir.cwd().createDir(io, tmp_dir, .default_dir) catch {}; defer Io.Dir.cwd().deleteTree(io, tmp_dir) catch {}; @@ -435,15 +474,26 @@ pub fn main(init: std.process.Init) !void { if (obj.get("cases")) |cases| { for (cases.array.items) |case| try runCase(init.arena.allocator(), io, golden_dir, server.?.port, case); } + var after_signal = false; // Windows: a step's cases after a signal it cannot send are skipped, and said so for (obj.get("then").?.array.items) |then| { const t = then.object; const kind = t.get("kind").?.string; const s = &server.?; if (std.mem.eql(u8, kind, "cases")) { + if (after_signal) { + std.debug.print("SKIPPED {d} cases after a signal Windows cannot send\n", .{t.get("cases").?.array.items.len}); + continue; + } for (t.get("cases").?.array.items) |case| try runCase(init.arena.allocator(), io, golden_dir, s.port, case); } else if (std.mem.eql(u8, kind, "sleep")) { std.Io.sleep(io, .fromMilliseconds(@intFromFloat(t.get("seconds").?.float * 1000.0)), .awake) catch {}; } else if (std.mem.eql(u8, kind, "sigterm")) { + if (builtin.os.tag == .windows) { + std.debug.print("SKIPPED lifecycle/sigterm-exit-0: Windows has no SIGTERM to send, the server is ended\n", .{}); + posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; + _ = s.child.wait(io) catch {}; + continue; + } posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; const term = s.child.wait(io) catch return error.WaitFailed; try outcomes.append(init.arena.allocator(), .{ @@ -454,13 +504,14 @@ pub fn main(init: std.process.Init) !void { }); } else if (std.mem.eql(u8, kind, "compare_log")) { const golden_log = try std.fmt.allocPrint(init.arena.allocator(), "{s}/{s}", .{ golden_dir, t.get("golden").?.string }); - const want = Io.Dir.cwd().readFileAlloc(io, golden_log, init.arena.allocator(), .limited(1 << 28)) catch ""; - const got = Io.Dir.cwd().readFileAlloc(io, s.log_path.?, init.arena.allocator(), .limited(1 << 28)) catch ""; + // a checkout with autocrlf gives the frozen log CRLF endings, the server writes LF: both compare as LF + const want = try std.mem.replaceOwned(u8, init.arena.allocator(), Io.Dir.cwd().readFileAlloc(io, golden_log, init.arena.allocator(), .limited(1 << 28)) catch "", "\r\n", "\n"); + const got = try std.mem.replaceOwned(u8, init.arena.allocator(), Io.Dir.cwd().readFileAlloc(io, s.log_path.?, init.arena.allocator(), .limited(1 << 28)) catch "", "\r\n", "\n"); try outcomes.append(init.arena.allocator(), .{ .group = "lifecycle", .name = "request-log", .verdict = if (std.mem.eql(u8, want, got)) .equal else .differ, - .why = if (std.mem.eql(u8, want, got)) &.{} else &.{"TENSORFOLD_REQUEST_LOG lines differ from the frozen ones"}, + .why = if (std.mem.eql(u8, want, got)) &.{} else try logDiff(init.arena.allocator(), want, got), }); } else if (std.mem.eql(u8, kind, "rewrite_key")) { var file = try Io.Dir.cwd().createFile(io, s.key_path.?, .{}); @@ -470,8 +521,10 @@ pub fn main(init: std.process.Init) !void { try fw.interface.writeAll(t.get("content").?.string); try fw.interface.flush(); posix.kill(s.child.id.?, posix.SIG.HUP) catch {}; // the server rereads the key file, as the reference does + after_signal = builtin.os.tag == .windows; } else if (std.mem.eql(u8, kind, "sighup")) { posix.kill(s.child.id.?, posix.SIG.HUP) catch {}; + after_signal = builtin.os.tag == .windows; } else if (std.mem.eql(u8, kind, "terminate")) { posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; _ = s.child.wait(io) catch {}; diff --git a/zig/tests/train_check.zig b/zig/tests/train_check.zig index 08601733c3..77acc8e7c1 100644 --- a/zig/tests/train_check.zig +++ b/zig/tests/train_check.zig @@ -1,5 +1,6 @@ //! The learning step's gradients against the engine's own forward, moving each layer along its own, on the real model. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const nm = tf.nemotron; @@ -51,8 +52,8 @@ pub fn main(init: std.process.Init) !void { const with_grad = (try t.step(&ids, start, .grad, turn)).loss; const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0.toNanoseconds())) / 1e6; std.debug.print("loss {d:.5} (with gradients {d:.5}), a step with gradients {d:.1} ms, GPU {d:.1} ms\n", .{ plain, with_grad, ms, b.last_ms }); - if (std.c.getenv("TF_TRAIN_TIME_ONLY") != null) return; - if (std.c.getenv("TF_TRAIN_PROFILE") != null) { + if (platform.c.getenv("TF_TRAIN_TIME_ONLY") != null) return; + if (platform.c.getenv("TF_TRAIN_PROFILE") != null) { var prof: nm.encoder.Profiler = .{ .queue = m.queue, .kernels = &m.kernels }; t.profile = &prof; _ = try t.step(&ids, start, .grad, turn); diff --git a/zig/tests/weight_read_check.zig b/zig/tests/weight_read_check.zig index d5acf5e7a6..1089d93666 100644 --- a/zig/tests/weight_read_check.zig +++ b/zig/tests/weight_read_check.zig @@ -1,5 +1,6 @@ //! Check the actual native loaders with a synthetic file and bounded I/O faults. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -63,7 +64,7 @@ pub fn main(init: std.process.Init) !void { } else { try replay.indexFile(path); const entry = try replay.entry(if (grouped) "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight" else "probe"); - defer _ = std.c.close(entry.fd); + defer _ = platform.c.close(entry.fd); const result = if (grouped) replay.group(0, 2, "weight") else replay.load("probe"); const buffer = result catch |err| { if (!eof or err != error.ShortRead) return err;