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..48c7b8d27f 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,8 @@ 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) | 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 +90,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. 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 +153,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..2a9c1afa1f 100644 --- a/build.zig +++ b/build.zig @@ -12,6 +12,8 @@ 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"); +const vision_build = @import("zig/build/vision.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -153,13 +155,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 +191,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). @@ -267,7 +270,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b counters_build.targets(b, target, optimize, metal, mods.sources, test_step); // The engines registry: the families this platform's binaries serve; the CLI reads it for models/info/pull. - 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 = mods.lanes }} }); + 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 = mods.lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } }); const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = target, @@ -572,7 +575,7 @@ pub fn distMetalServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_id m.metal.addSystemFrameworkPath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "System/Library/Frameworks" }))); m.metal.addLibraryPath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/lib" }))); m.metal.addSystemIncludePath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/include" }))); - const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = m.lanes }} }); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = m.lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, .ReleaseFast) } } }); const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = m.metal }, .{ .name = "engine_api", .module = api }, .{ .name = "tensorfold", .module = m.engine } } }); return nativeServer(b, target, .fast, m.metal, m.engine, api, engines, build_options, null); } diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md new file mode 100644 index 0000000000..52b17191c6 --- /dev/null +++ b/docs/WINDOWS.md @@ -0,0 +1,81 @@ +# 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 +``` + +- 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. +- 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.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..d33e8b5162 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -1,8 +1,8 @@ # Image input This page describes image input in the Python engine, frozen at 0.6.6. Install that engine from the -[`v0.6.6` tag](https://github.com/ashhart/TensorFold/tree/v0.6.6) to serve images. The native binary refuses every -`--vision` flag at startup, reads no vision tower, and answers image parts with HTTP 400. +[`v0.6.6` tag](https://github.com/ashhart/TensorFold/tree/v0.6.6) to serve images. The native binary reads images +on its Qwen3.5/3.8 CUDA engine ([below](#native-cuda-qwen3538)); elsewhere it refuses `--vision` at startup. Parts of it came from pull requests: GLM-5.3-Flash images by @mgoldwasser ([#101](https://github.com/ashhart/TensorFold/pull/101)), Flash Next images on CUDA by @shantanugoel ([#146](https://github.com/ashhart/TensorFold/pull/146)) from @MiaAI-Lab's patch, video by @MiaAI-Lab @@ -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: @@ -149,3 +149,14 @@ Tests cover input validation, bounded fetching, expanded prompt accounting, cach Hardware qualification is separate from these tests: each backend and chip needs real image understanding, drafted/serial equality, concurrency, chunked-prefill and memory checks before a release claim. Checkpoint metadata must describe the decoder separately from MTP: the `mlp_layer_types` list must match `num_hidden_layers` for Transformers validation. Preserve the separate MTP configuration and weights. Image input stays experimental until the hardware matrix is complete, and makes no throughput claim. + +## Native CUDA: Qwen3.5/3.8 + +`tensorfold-native serve --vision` loads the checkpoint's vision tower beside the language model on the +Qwen CUDA engine (`families/qwen3_5/cuda_vision.zig`, kernels in `qwen_vision.cu`, its Linears on cuBLASLt). The +tower and the scratch of one 16,384-patch call are allocated at startup and counted as model memory, so the +stream budget never has to make room for an image. A request's images go through the tower before its prompt; +their rows replace the `<|image_pad|>` placeholders in the prompt chunks, which rotate at M-RoPE's t, h and w +positions (interleaved sections 11, 11, 10), and the reply's rows rotate at their position plus `rope_delta`, as +transformers' `get_rope_index` lays them out. Text-only requests keep their bits. Images arrive through the native +server's image parts (decoded by the shared JPEG and PNG decoder); video parts are refused for now. 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..51c7dbdfef 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -3,6 +3,7 @@ const std = @import("std"); const exlr8_build = @import("exlr8.zig"); const golden_build = @import("golden.zig"); +const vision_build = @import("vision.zig"); /// Each .cu in zig/kernels/cuda (`src`, else `name`) with the flags its Python extension passes in `extra_cuda_cflags`. const Kernel = struct { name: []const u8, flags: []const []const u8, src: ?[]const u8 = null, arch_specific: bool = false }; @@ -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 = "qwen_vision", .flags = &.{ "-O3", "--fmad=false" } }, // ours: the Qwen3.5 vision tower's glue and attention .{ .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,13 +131,14 @@ 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 })); const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); runner.addImport("cuda", cuda); runner.addImport("core", mods.core); + runner.addImport("qwen_image", vision_build.module(b, target, optimize)); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); @@ -139,18 +149,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 } { - 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 }} }); +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 }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } }); 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 +176,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 +192,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 +207,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 +227,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), vision_build.module(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 +243,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 +265,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/build/vision.zig b/zig/build/vision.zig new file mode 100644 index 0000000000..0b78064367 --- /dev/null +++ b/zig/build/vision.zig @@ -0,0 +1,6 @@ +//! The image-input module (vision/qwen_image.zig), shared by the engine API and the CUDA test runner. +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/vision/qwen_image.zig"), .target = target, .optimize = optimize, .link_libc = true }); +} diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index edd83834c1..d19f6ee90c 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -276,7 +276,7 @@ __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) diff --git a/zig/kernels/cuda/experts_pack.cu b/zig/kernels/cuda/experts_pack.cu index 6aa7413793..b47ded555f 100644 --- a/zig/kernels/cuda/experts_pack.cu +++ b/zig/kernels/cuda/experts_pack.cu @@ -38,5 +38,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); diff --git a/zig/kernels/cuda/experts_prefill.cu b/zig/kernels/cuda/experts_prefill.cu index c1c7fbbac9..3a1df1dfee 100644 --- a/zig/kernels/cuda/experts_prefill.cu +++ b/zig/kernels/cuda/experts_prefill.cu @@ -152,6 +152,6 @@ __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) 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/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..df81a2103f --- /dev/null +++ b/zig/kernels/cuda/qwen35.cu @@ -0,0 +1,526 @@ +// 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); + } +} + +// Image rows into a prompt chunk's embeddings: row m of H (D wide) takes feature row src[m] where src[m] >= 0. +extern "C" __global__ void __launch_bounds__(256) tf_q35_splice(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ F, + const int* __restrict__ src, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const int s = src[blockIdx.x]; + if (s < 0) return; + for (int i = threadIdx.x; i < D; i += blockDim.x) H[(size_t)blockIdx.x * D + i] = F[(size_t)s * D + i]; +} + +// 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)); + } +} + +// M-RoPE's interleaved sections (mrope_section 11, 11, 10): frequency i turns by the t, h or w position. +__device__ __forceinline__ int mrope_axis(int i) { return (i % 3 == 1 && i < 33) ? 1 : ((i % 3 == 2 && i < 30) ? 2 : 0); } + +// RoPE on the first 64 of HD: pairs (i, i + 32) turned by pos[axis(i)] * theta^(-i / 32); text has one position. +__device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, const int pos[3]) { + if (d >= 64) return x[d]; + const int i = d % 32; + const float angle = (float)pos[mrope_axis(i)] * 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* __restrict__ M) { + 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]]; + // an image prompt's chunk reads each row's t, h and w positions from M ([3][MAXR]); other rows are shifted text + const int text = sg.len + rows_seen(rowanc[m]) - 1 + sg.delta; + int pos[3] = {text, text, text}; + if (M != nullptr) + for (int a = 0; a < 3; a++) pos[a] = M[a * MAXR + m]; + 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..4dffc534ce --- /dev/null +++ b/zig/kernels/cuda/qwen35.cuh @@ -0,0 +1,81 @@ +// 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) +constexpr int MAXR = 128; // a forward's rows at most (cuda_engine.zig's max_rows): an image chunk's position rows + +// 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; + int delta; // an image prompt's rotary shift: rows past it rotate at position + delta +}; + +// 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/qwen_vision.cu b/zig/kernels/cuda/qwen_vision.cu new file mode 100644 index 0000000000..7281b7a943 --- /dev/null +++ b/zig/kernels/cuda/qwen_vision.cu @@ -0,0 +1,202 @@ +// Qwen3.5 vision tower kernels (0.6.6's --vision in bf16); epilogues add the bias to cuBLASLt's fp32, round once. +#include +#include + +__device__ __forceinline__ float bf(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float rnd(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +// Patch embedding plus the resampled position grid: x = bf16(acc + b) + bf16(sum_t pos[idx_t] * w_t). +extern "C" __global__ void __launch_bounds__(256) tf_qv_embed(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + const __nv_bfloat16* __restrict__ pos, + const int* __restrict__ idx, + const float* __restrict__ wts, + __nv_bfloat16* __restrict__ out, int D) { + const long long row = blockIdx.x; + const int* ri = idx + row * 4; + const float* rw = wts + row * 4; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float h = rnd(acc[row * D + i] + bf(bias[i])); + float p = 0.0f; + for (int t = 0; t < 4; ++t) p += bf(pos[(long long)ri[t] * D + i]) * rw[t]; + out[row * D + i] = __float2bfloat16_rn(h + rnd(p)); + } +} + +// LayerNorm of bf16 rows in fp32 (mean, then the biased variance about it), bf16 out; a block a row. +extern "C" __global__ void __launch_bounds__(256) tf_qv_layernorm(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ y, int D, float eps) { + __shared__ float part[32]; + __shared__ float stat; + const long long row = blockIdx.x; + const __nv_bfloat16* xr = x + row * D; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, warps = blockDim.x >> 5; + float s = 0.0f; + for (int i = threadIdx.x; i < D; i += blockDim.x) s += bf(xr[i]); + for (int o = 16; o; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o); + if (lane == 0) part[warp] = s; + __syncthreads(); + if (threadIdx.x == 0) { + float t = 0.0f; + for (int i = 0; i < warps; ++i) t += part[i]; + stat = t / D; + } + __syncthreads(); + const float mean = stat; + float v = 0.0f; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float d = bf(xr[i]) - mean; + v += d * d; + } + for (int o = 16; o; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + __syncthreads(); + if (lane == 0) part[warp] = v; + __syncthreads(); + if (threadIdx.x == 0) { + float t = 0.0f; + for (int i = 0; i < warps; ++i) t += part[i]; + stat = rsqrtf(t / D + eps); + } + __syncthreads(); + const float r = stat; + for (int i = threadIdx.x; i < D; i += blockDim.x) + y[row * D + i] = __float2bfloat16_rn((bf(xr[i]) - mean) * r * bf(w[i]) + bf(b[i])); +} + +// qkv = bf16(acc + b) split into q, k, v [heads][rows][hd]; q and k get the axial 2-D rotary in fp32 (h, then w). +extern "C" __global__ void __launch_bounds__(128) tf_qv_qkv_rope(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + const int* __restrict__ hw, + const float* __restrict__ inv, + __nv_bfloat16* __restrict__ q, + __nv_bfloat16* __restrict__ k, + __nv_bfloat16* __restrict__ v, int rows, int heads, + int hd) { + const long long row = blockIdx.x; + const int head = blockIdx.y; + const int D = heads * hd; + const int half = hd / 2, quarter = hd / 4; + const float ph = static_cast(hw[row * 2]), pw = static_cast(hw[row * 2 + 1]); + const float* a = acc + row * 3 * D; + const long long o = (static_cast(head) * rows + row) * hd; + for (int j = threadIdx.x; j < hd; j += blockDim.x) { + const int c = head * hd + j; + // rotate_half pairs j with j +- hd/2; angle h * inv[jj] for the first quarter's frequencies, w * inv after + const int pj = j < half ? j + half : j - half; + const int pc = head * hd + pj; + const int jj = j % half; + const float ang = jj < quarter ? ph * inv[jj] : pw * inv[jj - quarter]; + const float cs = cosf(ang), sn = sinf(ang); + const float sign = j < half ? -1.0f : 1.0f; + const float qx = rnd(a[c] + bf(bias[c])), qp = rnd(a[pc] + bf(bias[pc])); + const float kx = rnd(a[D + c] + bf(bias[D + c])), kp = rnd(a[D + pc] + bf(bias[D + pc])); + q[o + j] = __float2bfloat16_rn(qx * cs + sign * qp * sn); + k[o + j] = __float2bfloat16_rn(kx * cs + sign * kp * sn); + v[o + j] = __float2bfloat16_rn(a[2 * D + c] + bf(bias[2 * D + c])); + } +} + +// Attention within one image's rows [s0, s0 + n); a thread a query row, online softmax in fp32, 64-row key tiles. +template +__device__ void qv_attention(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ out, int rows, + int heads, int s0, int n, float scale) { + constexpr int T = 64; + __shared__ float ks[T][HD]; + __shared__ float vs[T][HD]; + const int head = blockIdx.y; + const int qi = blockIdx.x * T + threadIdx.x; + const bool live = qi < n; + const long long base = static_cast(head) * rows + s0; + float qr[HD], acc[HD]; + float m = -INFINITY, l = 0.0f; + if (live) { + const __nv_bfloat16* qp = q + (base + qi) * HD; + for (int d = 0; d < HD; ++d) { + qr[d] = bf(qp[d]) * scale; + acc[d] = 0.0f; + } + } + for (int t0 = 0; t0 < n; t0 += T) { + const int count = min(T, n - t0); + __syncthreads(); + for (int i = threadIdx.x; i < count * HD; i += blockDim.x) { + const int r = i / HD, d = i % HD; + ks[r][d] = bf(k[(base + t0 + r) * HD + d]); + vs[r][d] = bf(v[(base + t0 + r) * HD + d]); + } + __syncthreads(); + if (!live) continue; + for (int r = 0; r < count; ++r) { + float s = 0.0f; +#pragma unroll + for (int d = 0; d < HD; ++d) s += qr[d] * ks[r][d]; + const float mn = fmaxf(m, s); + const float corr = __expf(m - mn); + const float p = __expf(s - mn); + l = l * corr + p; +#pragma unroll + for (int d = 0; d < HD; ++d) acc[d] = acc[d] * corr + p * vs[r][d]; + m = mn; + } + } + if (!live) return; + const float il = 1.0f / l; + __nv_bfloat16* op = out + static_cast(s0 + qi) * heads * HD + head * HD; + for (int d = 0; d < HD; ++d) op[d] = __float2bfloat16_rn(acc[d] * il); +} + +extern "C" __global__ void __launch_bounds__(64) tf_qv_attention72(const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ out, int rows, + int heads, int s0, int n, float scale) { + qv_attention<72>(q, k, v, out, rows, heads, s0, n, scale); +} + +extern "C" __global__ void __launch_bounds__(64) tf_qv_attention64(const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ out, int rows, + int heads, int s0, int n, float scale) { + qv_attention<64>(q, k, v, out, rows, heads, s0, n, scale); +} + +// y = gelu(bf16(acc + b)), the tanh form (gelu_pytorch_tanh) or the erf form (nn.GELU), computed in fp32, bf16 out. +extern "C" __global__ void __launch_bounds__(256) tf_qv_bias_gelu(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ y, int D, int tanh_form) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float x = rnd(acc[row * D + i] + bf(bias[i])); + float g; + if (tanh_form) { + const float kBeta = 0.7978845608028654f; // sqrt(2 / pi) + g = 0.5f * x * (1.0f + tanhf(kBeta * (x + 0.044715f * x * x * x))); + } else { + g = 0.5f * x * (1.0f + erff(x * 0.7071067811865476f)); + } + y[row * D + i] = __float2bfloat16_rn(g); + } +} + +// x += bf16(acc + b), the residual add in bf16 (both operands bf16, the sum rounded). +extern "C" __global__ void __launch_bounds__(256) tf_qv_bias_residual(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ x, int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float y = rnd(acc[row * D + i] + bf(bias[i])); + x[row * D + i] = __float2bfloat16_rn(bf(x[row * D + i]) + y); + } +} + +// y = bf16(acc + b): the merger's last Linear. +extern "C" __global__ void __launch_bounds__(256) tf_qv_bias_cast(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ y, int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) y[row * D + i] = __float2bfloat16_rn(acc[row * D + i] + bf(bias[i])); +} 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..9b5edacb0e 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -38,7 +38,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/engine_api.zig b/zig/src/core/engine_api.zig index 903dc4d40b..82fe572e94 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -1,5 +1,8 @@ //! The engine a server drives (``Engine``), and ``LaneHost``: the lane core served on one thread, rounds left to it. const std = @import("std"); +/// Image inputs (decode, resize, patches) for families with a vision tower; the server reaches them here. +pub const Image = lanes.Image; +pub const qwen_image = @import("qwen_image"); const lanes = @import("lanes"); const Allocator = std.mem.Allocator; @@ -88,6 +91,15 @@ pub const Request = struct { logprobs: ?u8 = null, /// Rows for each prompt token after the first, with ranks and this many best tokens, when ``Info.prompt_logprobs``. prompt_logprobs: ?u8 = null, + /// An image prompt's inputs in prompt order (Info.vision engines only; the prompt holds their placeholder rows). + images: []const lanes.Image = &.{}, + /// The prompt as the prompt cache keys it: each image's placeholder rows as ids from its content (empty: prompt). + cache_key: []const u32 = &.{}, + + /// The ids the prompt cache keys this request's states by: placeholders stand for their images. + pub fn cacheKey(r: *const Request) []const u32 { + return if (r.cache_key.len > 0) r.cache_key else r.prompt; + } }; pub const Reason = enum { stop, length, cancelled, failed }; @@ -162,6 +174,10 @@ pub const Info = struct { prompt_cache_plan: ?PromptCachePlan = null, /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, + /// The engine encodes image inputs (Request.images): it was opened with ``vision`` and its family has a tower. + vision: bool = false, + /// Its tower takes GLM-5.3-Flash's preparation and placeholders (qwen_image.glm) instead of Qwen's. + vision_glm: bool = false, /// Whether a request's caches fit beside the running streams' now (null: a free lane does); else it waits in order. admit: ?Admit = null, }; @@ -225,10 +241,20 @@ pub const Open = struct { master_port: u16 = 29551, /// --policy: the engine's policy as `key=value,...`; empty: its defaults. policy: []const u8 = "", + /// --vision: load the family's vision tower and accept image inputs. + vision: bool = false, + /// --kv-dtype: bits of the attention cache's codes (8 or 4) for families that keep a quantized cache. + kv_bits: u8 = 8, }; /// An opened engine; ``close`` stops its thread and frees its backend. -pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque }; +pub const Opened = struct { + engine: Engine, + close: *const fn (ctx: *anyopaque) void, + ctx: *anyopaque, + /// Rank 1 of two: replay rank 0's engine calls until it stops (the server then serves no HTTP); null elsewhere. + follow: ?*const fn (ctx: *anyopaque) anyerror!void = null, +}; pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0 }; 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/lane_host.zig b/zig/src/core/lane_host.zig index 4394fdac05..acb7632595 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -77,7 +77,7 @@ pub const LaneHost = struct { fn kept(ptr: *anyopaque, s: *lanes.Stream, at: u32) void { const job: *Job = @ptrCast(@alignCast(ptr)); job.reported(); // before a keep can evict the entry the pass restored - if (job.host.cache) |store| _ = store.keepBorn(job.request.prompt, born(job.request), at, s, job.request.chunks, job.host.spans(job.request)); + if (job.host.cache) |store| _ = store.keepBorn(job.request.cacheKey(), born(job.request), at, s, job.request.chunks, job.host.spans(job.request)); } /// Report a restored prefix or its failed copy; an untouched prefix remains kept. @@ -85,9 +85,10 @@ pub const LaneHost = struct { const e = job.entry orelse return; job.entry = null; const store = job.host.cache orelse return; + defer store.unpin(e); // after resumed, which reads e if (!job.started) return; job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it - if (job.stream.reuse_failed) store.resumed(e, job.request.prompt, false) else if (job.stream.cached == e.at) store.resumed(e, job.request.prompt, true); + if (job.stream.reuse_failed) store.resumed(e, job.request.cacheKey(), false) else if (job.stream.cached == e.at) store.resumed(e, job.request.cacheKey(), true); } }; @@ -285,16 +286,16 @@ pub const LaneHost = struct { h.mutex.unlock(h.io); } - /// Prefills the next queued request into a free lane; false when none waits or no lane is free. - fn admitOne(h: *LaneHost) bool { + /// The next queued request into the admitted ones (null: none waits, no lane is free, or its caches do not fit). + pub fn takeJob(h: *LaneHost) ?*Job { h.lock(); if (h.queued.items.len == 0 or h.admitted.items.len >= @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes)) { h.unlock(); - return false; + return null; } const at = admit.pick(h) orelse { // a shared cache pool's head waits for room h.unlock(); - return false; + return null; }; const job = h.queued.orderedRemove(at); job.seq = h.admissions; @@ -302,20 +303,30 @@ pub const LaneHost = struct { h.admitted.append(h.gpa, job) catch { h.unlock(); h.finish(job, .failed, "out of memory"); - return true; + return null; }; h.unlock(); + return job; + } + + /// The admitted request's stream (its kept state, drafter and grammar); false when it was dropped instead. + pub fn setup(h: *LaneHost, job: *Job) bool { const r = job.request; - if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran + if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) { // other work came first: it never ran + _ = h.cancel(job); + return false; + } var reuse: lanes.stream.Reuse = .{}; - // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass - if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.prompt, born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { - job.entry = l.entry; + if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.cacheKey(), born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { + job.entry = if (l.entry) |e| store.pin(e) else null; // kept alive until its restore reports job.kept0 = store.counts.kept; job.marks = l.marks; reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; } else |_| {}; - job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start"); + job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch { + _ = h.drop(job, "the drafter could not start"); + return false; + }; job.stream = lanes.Stream.init(h.gpa, .{ .id = "request", .prompt = r.prompt, @@ -339,18 +350,31 @@ pub const LaneHost = struct { .reuse = reuse, .grammar = if (r.structure) |st| h.constraint(st) catch |e| { job.proposer.deinit(); - return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + { + _ = h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + return false; + } } else null, + .images = r.images, .decode_spans = h.spans(r), .decide = r.decide, .untaught = r.untaught, }) catch { job.proposer.deinit(); - return h.drop(job, "out of memory"); + { + _ = h.drop(job, "out of memory"); + return false; + } }; job.started = true; - const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); - job.began = began; + job.began = std.Io.Clock.awake.now(h.io).toNanoseconds(); + return true; + } + + /// An admitted request into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. + pub fn launch(h: *LaneHost, job: *Job) bool { + const r = job.request; + const began = job.began; if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); if (lone.fits(h, job)) return lone.run(h, job, began, false); if (h.core.fills()) { @@ -424,7 +448,7 @@ pub const LaneHost = struct { return h.cancel(job); }; h.prefilled(job, began); - if (h.cache) |store| _ = if (!job.request.untaught and store.rules.lookahead == 0) store.keepBorn(job.request.prompt, born(job.request), @intCast(job.request.prompt.len), &job.stream, job.request.chunks, h.spans(job.request)) else false; + if (h.cache) |store| _ = if (!job.request.untaught and store.rules.lookahead == 0) store.keepBorn(job.request.cacheKey(), born(job.request), @intCast(job.request.prompt.len), &job.stream, job.request.chunks, h.spans(job.request)) else false; be.release(&job.stream); h.remove(job); h.finish(job, .length, ""); @@ -444,12 +468,13 @@ pub const LaneHost = struct { } /// No memory for another lane yet: the job waits at the queue's head, and admission holds until a job ends. - fn requeue(h: *LaneHost, job: *Job) bool { + pub fn requeue(h: *LaneHost, job: *Job) bool { if (job.stream.grammar) |g| g.deinit(h.gpa); // the next admission compiles its own job.stream.deinit(h.gpa); job.proposer.deinit(); job.started = false; job.fill_began = false; + if (job.entry) |e| h.cache.?.unpin(e); job.entry = null; // looked up again at its next admission h.gpa.free(job.marks); job.marks = &.{}; @@ -475,7 +500,7 @@ pub const LaneHost = struct { const alone_words = "the engine has no memory for this request's caches with nothing else running: lower --context or --parallel"; /// A job cancelled in its prompt pass, its lane already released. - fn cancel(h: *LaneHost, job: *Job) bool { + pub fn cancel(h: *LaneHost, job: *Job) bool { h.remove(job); h.finish(job, .cancelled, ""); return true; @@ -494,7 +519,7 @@ pub const LaneHost = struct { fn run(h: *LaneHost) void { while (true) { h.takeCancels(); - while (h.admitOne()) {} + while (admit.many(h)) {} h.fillOne(); status.noteLive(h); h.lock(); diff --git a/zig/src/core/lane_host_admit.zig b/zig/src/core/lane_host_admit.zig index 016ab4398a..f394a3729e 100644 --- a/zig/src/core/lane_host_admit.zig +++ b/zig/src/core/lane_host_admit.zig @@ -2,6 +2,8 @@ const api = @import("engine_api.zig"); const LaneHost = @import("lane_host.zig").LaneHost; const Job = LaneHost.Job; +const lanes = @import("lanes"); +const lone = @import("lane_host_lone.zig"); /// The queued job to admit (its index), or null to hold; the caller holds the lock and a lane is free. pub fn pick(h: *LaneHost) ?usize { @@ -21,3 +23,51 @@ pub fn pick(h: *LaneHost) ?usize { fn fits(a: api.Admit, job: *const Job) bool { return a.fits(a.ctx, @intCast(job.request.prompt.len), job.request.max_tokens); } + +/// Every queued request a lane is free for, prompts in one pass where the backend takes several; else one at a time. +pub fn many(h: *LaneHost) bool { + if (!h.core.batchesPrompts() or h.core.fills()) return one(h); + h.lock(); + const waiting = h.queued.items.len; + const free = h.info_.lanes -| h.admitted.items.len; + h.unlock(); + if (waiting < 2 or free < 2) return one(h); + var jobs: [lanes.Engine.max_together]*Job = undefined; + var n: usize = 0; + while (n < jobs.len and n < free) { + const job = h.takeJob() orelse break; + if (!h.setup(job)) continue; + if ((job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job)) { + _ = h.launch(job); + continue; + } + jobs[n] = job; + n += 1; + } + if (n == 0) return true; + if (n == 1) return h.launch(jobs[0]); + var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; + for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; + var errs: [lanes.Engine.max_together]?anyerror = undefined; + h.core.addStreams(streams[0..n], errs[0..n]); + for (jobs[0..n], errs[0..n]) |job, err| { + if (err) |e| { + _ = switch (e) { + error.Cancelled => h.cancel(job), + error.LaneMemory => h.requeue(job), + else => h.drop(job, h.words(e)), + }; + continue; + } + h.prefilled(job, job.began); + if (h.deliver(job)) h.remove(job); + } + return true; +} + +/// The next queued request admitted alone: its stream set up and launched. +fn one(h: *LaneHost) bool { + const job = h.takeJob() orelse return false; + if (!h.setup(job)) return true; + return h.launch(job); +} diff --git a/zig/src/core/lane_host_reuse_test.zig b/zig/src/core/lane_host_reuse_test.zig index 01b1fc6bec..69fca5ae1f 100644 --- a/zig/src/core/lane_host_reuse_test.zig +++ b/zig/src/core/lane_host_reuse_test.zig @@ -106,6 +106,94 @@ test "a lane host resumes a conversation from its kept prompt state and reports try std.testing.expectEqual(@as(usize, 2), store.entries.items.len); } +test "an image's placeholders are cached by its content: another image under the same ids resumes none of it" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + const Snaps = struct { + fn bytes(_: *anyopaque, at: u32) u64 { + return 64 + 4 * @as(u64, at); + } + fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!pc.Saved { + const f: *lanes.fake.Fake = @ptrCast(@alignCast(ptr)); + return f.save(@ptrCast(@alignCast(owner.?)), at); + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: pc.Saved) anyerror!void { + return error.BackendRestores; // lane-core backends restore inside their own prompt pass + } + fn drop(ptr: *anyopaque, saved: pc.Saved) void { + const f: *lanes.fake.Fake = @ptrCast(@alignCast(ptr)); + f.drop(saved); + } + }; + var store = pc.Store.init(gpa, .{ .ptr = &target, .vtable = &.{ .bytes = Snaps.bytes, .save = Snaps.save, .restore = Snaps.restore, .drop = Snaps.drop } }, .{ .min_gap = 4, .min_prompt = 0 }, 1 << 20); + defer store.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + host.cache = &store; + try host.start(); + defer host.stop(); + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + cached: ?u32 = null, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .prefilled => |c| b.cached = c, + .logprobs, .prompt_logprobs => {}, + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| b.done = f.reason, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const e = host.engine(); + // prompt ids 3..5 are one image's placeholders (token 9); cache_key holds ids derived from each image's content + const t1 = [_]u32{ 3, 1, 9, 9, 9, 9, 2, 6, 5, 3 }; + var key_a = t1; + var key_b = t1; + for (2..6) |k| { + key_a[k] = 0x8000_0000 + @as(u32, @intCast(k)); + key_b[k] = 0x9000_0000 + @as(u32, @intCast(k)); + } + var b1: Box = .{}; + defer b1.tokens.deinit(gpa); + const r1: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_a }; + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b1.wait()); + try std.testing.expectEqual(@as(usize, 1), store.entries.items.len); // image A's state at 8, past its image + // the same text and placeholder ids with image B: the kept state is A's, so B resumes none of it + var b2: Box = .{}; + defer b2.tokens.deinit(gpa); + const r2: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_b }; + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b2.wait()); + try std.testing.expectEqual(@as(?u32, 0), b2.cached); + // image A again resumes its own state + var b3: Box = .{}; + defer b3.tokens.deinit(gpa); + const r3: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_a }; + try e.submit(3, &r3, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(?u32, 8), b3.cached); +} + test "a turn's mark evicting the state it resumed from: the hit is reported first and the entry freed once" { const gpa = std.testing.allocator; var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index fe26d3f003..23a4efd892 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -1,4 +1,4 @@ -//! What the round loop asks a backend (Metal, CUDA or the fake) to compute; drawn tokens stay on the GPU behind handles. +//! What the round loop asks a backend (Metal, CUDA or the fake); drawn tokens stay on the GPU behind handles. const std = @import("std"); const Stream = @import("stream.zig").Stream; const Shape = @import("shape.zig").Shape; @@ -68,6 +68,8 @@ pub const Backend = struct { prefill: *const fn (ptr: *anyopaque, s: *Stream) anyerror!void, /// One chunk of the prompt pass, false while chunks remain; null: the pass runs whole. prefill_step: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, + /// Several streams' prompts in one pass, each as after `prefill`; false, doing nothing, when they don't fit. + prefill_many: ?*const fn (ptr: *anyopaque, ss: []const *Stream) anyerror!bool = null, /// Draw the first token from the prompt's last row at `position`; a handle to it. first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64, /// Feed one token and queue the draw of the next at `position` (a one-token round); a handle to it. @@ -137,7 +139,7 @@ pub const Backend = struct { b.vtable.release(b.ptr, s); } - /// A driver of its own starts a stream: the prompt pass and its first token, committed; a cancel or failure releases it. + /// A driver of its own starts a stream: prompt pass and first token, committed; cancel or failure releases it. pub fn opening(b: Backend, gpa: std.mem.Allocator, s: *Stream) !u32 { errdefer b.release(s); try b.prefill(s); 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/engine.zig b/zig/src/core/lanes/engine.zig index cf0a97ea4a..e4c7950781 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -83,7 +83,16 @@ pub const Engine = struct { pub const fills = open.fills; pub const fillStream = open.fillStream; - /// Take over a stream another driver decoded so far: its cache settled, `pending` and `cache_len` set, no drafts held. + pub const adoptPrefilled = open.adoptPrefilled; + pub const max_together = open.max_together; + pub const addStreams = open.addStreams; + + /// Whether the backend prefills several prompts in one pass (addStreams batches them). + pub fn batchesPrompts(e: *const Engine) bool { + return e.backend.vtable.prefill_many != null; + } + + /// Take over a stream another driver decoded: cache settled, `pending` and `cache_len` set, no drafts held. pub fn adopt(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); } 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/lanes.zig b/zig/src/core/lanes/lanes.zig index 5d50a16ef3..57d8ae9eb7 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -29,6 +29,7 @@ pub const grammar = @import("grammar.zig"); pub const Engine = engine.Engine; pub const Config = config.Config; +pub const Image = @import("stream.zig").Image; pub const Model = config.Model; pub const Stream = stream.Stream; pub const Sampling = sampling.Sampling; 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/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..c75fa81b86 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -88,3 +88,58 @@ pub fn fillStream(e: *Engine, s: *Stream, first: bool) !bool { try opened(e, s); return true; } + +/// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. +pub fn adoptPrefilled(e: *Engine, s: *Stream, token: u32) !void { + _ = e.arena.reset(.retain_capacity); + s.context.shrinkRetainingCapacity(s.prompt_len); + s.pending = null; + s.cache_len = s.prompt_len; + if (e.cfg.family_mtp and s.drafts) { + const d: u32 = @intCast(try e.rule.depth(win.who(s))); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = token }, .rows = null, .start = s.prompt_len, .position = s.prompt_len + 1, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { + s.next = held; + }; + } + if (s.logprobs != null) return error.LogprobsUnsupported; // another driver's prompt pass gives no first row + _ = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + if (s.finished) return e.release(s); + try e.live.append(e.gpa, s); +} + +/// Streams a backend prefills in one pass at most. +pub const max_together = 16; + +/// open.add for several streams, prompts in one pass when the backend allows; ``errs[i]``: stream i's error. +pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (errs) |*x| x.* = null; + const many = e.backend.vtable.prefill_many orelse return addEach(e, ss, errs); + if (ss.len < 2 or ss.len > max_together) return addEach(e, ss, errs); + _ = e.arena.reset(.retain_capacity); + for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; + const together = many(e.backend.ptr, ss) catch |err| { + for (errs) |*x| x.* = err; + return; + }; + if (!together) return addEach(e, ss, errs); + for (ss, errs) |s, *x| { + if (s.isCancelled()) { // cancelled in its pass: no first token + e.backend.release(s); + x.* = error.Cancelled; + continue; + } + opened(e, s) catch |err| { + x.* = err; + }; + } +} + +fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (ss, errs) |s, *x| add(e, s) catch |err| { + x.* = err; + }; +} diff --git a/zig/src/core/lanes/sampling.zig b/zig/src/core/lanes/sampling.zig index e9d469756c..43bddfa3a1 100644 --- a/zig/src/core/lanes/sampling.zig +++ b/zig/src/core/lanes/sampling.zig @@ -75,7 +75,7 @@ pub fn pairwiseSum(a: []const f64) f64 { return pairwiseSum(a[0..n2]) + pairwiseSum(a[n2..]); } -const Candidate = struct { value: f64, id: u64 }; +pub const Candidate = struct { value: f64, id: u64 }; fn byValueThenId(_: void, a: Candidate, b: Candidate) bool { return a.value > b.value or (a.value == b.value and a.id < b.id); @@ -88,6 +88,12 @@ pub fn choose(gpa: Allocator, values: []const f64, ids: []const u64, position: u defer gpa.free(order); for (order, values, ids) |*o, v, id| o.* = .{ .value = v, .id = id }; std.mem.sort(Candidate, order, {}, byValueThenId); + return chooseOrdered(gpa, order, position, s); +} + +/// ``choose`` over candidates already in its order (by value, then the lower id at a tie). +pub fn chooseOrdered(gpa: Allocator, order: []const Candidate, position: u64, s: Sampling) !u64 { + const width = order.len; if (s.temperature <= 0) return order[0].id; // temperature 0 is greedy: the first maximum, ties to the lower id const k = @max(1, @min(if (s.top_k != 0) @as(usize, s.top_k) else width, width)); const scaled = try gpa.alloc(f64, k); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index f028d10cc4..bc0574e1ca 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -76,7 +76,7 @@ pub const CancelCheck = struct { check: *const fn (ptr: *anyopaque) bool, }; -/// A kept prompt state the backend restores before its prompt pass (core/prompt_cache.zig), and where it keeps new ones. +/// A kept prompt state the backend restores before its prompt pass (core/prompt_cache.zig), and keeps new ones in. pub const Reuse = struct { saved: ?*anyopaque = null, // the backend's copy of the state after `at` prompt tokens (null: start at 0) at: u32 = 0, @@ -97,6 +97,22 @@ pub const Held = struct { parents: ?[]i32 = null, }; +/// An image for a vision tower: patches (fp32 [h * w, 1536], merge-block order), grid [1, h, w]; prompt placeholders. +pub const Image = struct { + patches: []const f32, + grid: [3]i64, + + /// BLAKE3 of the grid (little-endian) and the patches' bits: the image's identity for kept prompt states. + pub fn digest(im: Image) [32]u8 { + var h = std.crypto.hash.Blake3.init(.{}); + for (im.grid) |g| h.update(&std.mem.toBytes(std.mem.nativeToLittle(i64, g))); + h.update(std.mem.sliceAsBytes(im.patches)); + var out: [32]u8 = undefined; + h.final(&out); + return out; + } +}; + pub const Spec = struct { id: []const u8, prompt: []const u32, @@ -123,6 +139,7 @@ pub const Spec = struct { prompt_logprobs: ?u8 = null, // and for each prompt token after the first, with ranks (the prompt pass starts at 0) grammar: ?*grammar.Constraint = null, // structured output: the reply stays in it (the creator frees it) decode_spans: []const [2]u32 = &.{}, // prompt rows an earlier reply decoded (Request.decode_spans), for backends that keep their bits + images: []const Image = &.{}, // an image prompt's inputs, in prompt order (the backend encodes them at prefill) }; pub const Stream = struct { @@ -140,6 +157,7 @@ pub const Stream = struct { think_close: []const u32, think_end: i64, think_open: bool, + images: []const Image = &.{}, loop_guard: bool, loop_period: ?u32 = null, chunks: []const u32, @@ -202,6 +220,7 @@ pub const Stream = struct { .cancel_check = spec.cancel_check, .yield_check = spec.yield_check, .think_budget = spec.think_budget, + .images = spec.images, .think_close = spec.think_close, .think_end = spec.think_end, .think_open = spec.think_open orelse (spec.think_budget > 0 or (spec.loop_guard and spec.think_end >= 0)), @@ -447,3 +466,12 @@ test "loop guard fires, closes thinking, and latches once" { _ = try s.commit(gpa, &.{8}, &.{}); try std.testing.expect(!s.finished and s.loop_period.? == 1); } + +test "an image's digest covers its grid and every patch bit" { + const a: [3]f32 = .{ 0.1, 0.2, 0.3 }; + const b: [3]f32 = .{ 0.1, 0.2, -0.3 }; + const da = (Image{ .patches = &a, .grid = .{ 1, 2, 2 } }).digest(); + try std.testing.expect(!std.mem.eql(u8, &da, &(Image{ .patches = &b, .grid = .{ 1, 2, 2 } }).digest())); + try std.testing.expect(!std.mem.eql(u8, &da, &(Image{ .patches = &a, .grid = .{ 1, 1, 4 } }).digest())); + try std.testing.expectEqualSlices(u8, &da, &(Image{ .patches = &a, .grid = .{ 1, 2, 2 } }).digest()); +} 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.zig b/zig/src/core/prompt_cache.zig index f1465325e1..81d5cf4ed4 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -108,6 +108,8 @@ pub const Entry = struct { last: []u32, // the last prompt that kept or resumed it shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded + pins: u32 = 0, // admitted jobs that restore it and have not yet: eviction passes it over + dead: bool = false, // removed while pinned: no lookup finds it, and its last unpin frees it }; pub const Counts = struct { hits: u64 = 0, misses: u64 = 0, kept: u64 = 0, evicted: u64 = 0, refused: u64 = 0, failed: u64 = 0 }; @@ -148,6 +150,7 @@ pub const Store = struct { } fn free(s: *Store, e: *Entry) void { + s.held -= e.bytes; s.family.vtable.drop(s.family.ptr, e.saved); s.gpa.free(e.tokens); s.gpa.free(e.decode_spans); @@ -157,8 +160,19 @@ pub const Store = struct { fn remove(s: *Store, i: usize) void { const e = s.entries.orderedRemove(i); - s.held -= e.bytes; - s.free(e); + e.dead = true; + if (e.pins == 0) s.free(e); + } + + /// `e` waits for a restore: nothing frees it until its last `unpin` (removed meanwhile, no lookup finds it). + pub fn pin(_: *Store, e: *Entry) *Entry { + e.pins += 1; + return e; + } + + pub fn unpin(s: *Store, e: *Entry) void { + e.pins -= 1; + if (e.pins == 0 and e.dead) s.free(e); } /// The family's free storage kept for later saves. @@ -450,7 +464,7 @@ pub const Store = struct { if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| s.held -| bytes); // spare that cannot take this state goes before any state if (s.spare() < free_) continue; } - s.remove(s.victimBut(null, prompt, born).?); // its storage may come back spare and take this state + s.remove(s.victimBut(null, prompt, born) orelse return s.fail(at, error.RestoresPending)); // spare for this state s.counts.evicted += 1; } if (s.entries.items.len == 0) if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| bytes); @@ -473,11 +487,11 @@ pub const Store = struct { const charged = if (s.family.vtable.charged) |f| f(s.family.ptr, saved) else bytes; e.* = .{ .tokens = tokens, .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .last = last, .shared = shared }; owns_decoded = false; + s.held += charged; // free takes it back s.entries.append(s.gpa, e) catch { s.free(e); return s.fail(at, error.OutOfMemory); }; - s.held += charged; s.counts.kept += 1; if (shared) learned.learn(s, e, starts); if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| s.held); // spare storage only inside what the budget leaves @@ -534,17 +548,17 @@ pub const Store = struct { return false; } - /// The entry to free first, never `skip`: a superseded one, no branch point, oldest; else the oldest, rewind last. + /// The entry to free first, never `skip` or pinned: superseded, no branch point, oldest; else oldest, rewind last. fn victimBut(s: *const Store, skip: ?*Entry, prompt: []const u32, born: u32) ?usize { var best: ?usize = null; for (s.entries.items, 0..) |e, i| { - if (e.shared or e.branched or e == skip or s.isActiveRewind(e, prompt)) continue; // shared cuts and branch points go only as the least recently used + if (e.shared or e.branched or e == skip or e.pins > 0 or s.isActiveRewind(e, prompt)) continue; // shared cuts and branch points go only as the least recently used const now = born > e.born and prompt.len > e.tokens.len and std.mem.eql(u32, prompt[0..e.tokens.len], e.tokens); // a longer request's tokens extend it if ((now or s.extended(e)) and (best == null or e.used < s.entries.items[best.?].used)) best = i; } if (best) |i| return i; var oldest: ?usize = null; - for (s.entries.items, 0..) |e, i| if (e != skip and (oldest == null or s.before(e, s.entries.items[oldest.?], prompt))) { + for (s.entries.items, 0..) |e, i| if (e != skip and e.pins == 0 and (oldest == null or s.before(e, s.entries.items[oldest.?], prompt))) { oldest = i; }; return oldest; 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_cache_test.zig b/zig/src/core/prompt_cache_test.zig index 8452036b75..cde079b745 100644 --- a/zig/src/core/prompt_cache_test.zig +++ b/zig/src/core/prompt_cache_test.zig @@ -391,3 +391,80 @@ test "a kept state grows only into what the budget leaves, and evicts nothing to var stranger: u8 = 0; try std.testing.expect(!s.grow(&stranger, 0)); // a state the store doesn't hold } + +test "an entry pinned for a restore still to come outlives the lookups set up after it; unpinned, it goes again" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 220); // fake bytes: 100 + at + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); // 103 bytes, the oldest + const y = [_]u32{ 6, 6, 6, 6 }; + _ = f.pass(&s, &y, try s.begin(a, &y, 3, &.{}, &.{}, null, &.{})); // 103: 206 held + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const l = try s.lookup(a, &xa, 3, &.{}, &.{}, &.{}); // a job set up to resume x + const e = l.entry.?; + _ = s.pin(e); + const z = [_]u32{ 7, 7, 7, 7, 7, 7, 7, 7 }; + _ = try s.lookup(a, &z, 7, &.{}, &.{}, &.{}); // the next job's room: y goes, not x + try std.testing.expect(s.find(&xa, &.{}, &.{}) == e); + try std.testing.expect(s.find(&.{ 6, 6, 6, 6 }, &.{}, &.{}) == null); + const w: [130]u32 = @splat(8); + _ = try s.lookup(a, &w, 120, &.{}, &.{}, &.{}); // 220 bytes, and nothing but x to evict: it stays + try std.testing.expect(!s.keep(&w, 120, null, &.{}, &.{})); // and a keep with no room is refused + try std.testing.expect(s.find(&xa, &.{}, &.{}) == e); + s.unpin(e); + _ = try s.lookup(a, &w, 120, &.{}, &.{}, &.{}); + try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); + try std.testing.expectEqual(s.entries.items.len, f.live); +} + +test "two jobs pin one entry and the first restore fails: lookups miss it, and it lives until the second job unpins" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const xb = [_]u32{ 5, 5, 5, 5, 8 }; + const e = s.pin((try s.lookup(a, &xa, 3, &.{}, &.{}, &.{})).entry.?); + try std.testing.expect(s.pin((try s.lookup(a, &xb, 3, &.{}, &.{}, &.{})).entry.?) == e); + const held = s.held; + s.resumed(e, &xa, false); // the first job's restore failed: it prefilled from 0 + s.unpin(e); + try std.testing.expect(s.find(&xb, &.{}, &.{}) == null); + try std.testing.expectEqual(@as(usize, 0), s.entries.items.len); + try std.testing.expectEqual(@as(usize, 1), f.live); // the second job's restore still reads it + try std.testing.expectEqual(held, s.held); + try std.testing.expectEqual(@as(u32, 1), e.pins); + s.resumed(e, &xb, true); // no longer kept: its report changes nothing + s.unpin(e); + try std.testing.expectEqual(@as(usize, 0), f.live); + try std.testing.expectEqual(@as(u64, 0), s.held); +} + +test "forget leaves a pinned entry to its restore" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const e = s.pin((try s.lookup(a, &xa, 3, &.{}, &.{}, &.{})).entry.?); + s.forget(&xa, 3, &.{}); + try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); + try std.testing.expectEqual(@as(usize, 1), f.live); + s.unpin(e); + try std.testing.expectEqual(@as(usize, 0), f.live); +} 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/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..08b94a05d5 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -30,6 +30,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 qwen_vision: []const u8 = if (available) &Blob("fatbin_qwen_vision").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/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/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..5e27358447 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda.zig @@ -0,0 +1,20 @@ +//! 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 draft = @import("cuda_draft.zig"); +pub const vision = @import("cuda_vision.zig"); +pub const reuse = @import("cuda_reuse.zig"); + +test { + _ = config; + _ = geometry; + _ = @import("draw.zig"); + _ = @import("cuda_draft_weights.zig"); + _ = draft; + _ = vision; + _ = @import("cuda_vision_prompt.zig"); + _ = reuse; + _ = native; +} 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..7b482fa9a0 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -0,0 +1,103 @@ +//! 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, delta: 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, + /// An image prompt's rotary shift (get_rope_index's rope_delta): rows after its prompt rotate at position + delta. + delta: i32 = 0, + + 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; + self.delta = 0; + 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), .delta = self.delta }; + } + + /// 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..ea6437f0da --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -0,0 +1,551 @@ +//! 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 }; +pub const Layer = struct { + input_norm: u64, + post_norm: u64, + gate: Affine4, + up: Affine4, + down: Affine4, + 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, + splice: 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, images: ?*const Images = null }; + +/// An image prompt's rows for its prompt chunks: device feature rows, each prompt row's feature row, its positions. +pub const Images = struct { + features: u64, // bf16 [image rows, hidden] + src: []const i32, // [prompt]: a feature row, or -1 for text + pos: []const i32, // [3][prompt]: t, h and w rotary positions (get_rope_index) +}; + +/// 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: [7]cuda.Module, + draws: cuda.sample.Draws, + affine: cuda.qlinear.Affine4Kernels, + 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, + // 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, + /// During an image prompt chunk's forward: its feature rows (0: none); the chunk's src and positions are staged. + features: u64 = 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 src_at = params_at + 8 * 4; // an image chunk's rows' feature rows + const mpos_at = src_at + max_rows * 4; // and their t, h and w positions, max_rows apart (qwen35.cuh's MAXR) + const tables_len = mpos_at + 3 * max_rows * 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.features = 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 }; + 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]); + 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); + 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; + for (parts) |p| if (p.images != null and record) return error.InvalidQwenImages; + 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); + // an image prompt's chunk (one part, its rows at prompt positions len..): feature rows and t, h, w positions + const images: ?*const Images = if (record) null else parts[0].images; + if (images) |im| { + const p0 = parts[0].cache.len; + if (p0 + total > im.src.len or im.pos.len != 3 * im.src.len) return error.InvalidQwenImages; + const n = im.src.len; + @memcpy(std.mem.bytesAsSlice(i32, st[src_at..][0 .. total * 4]), im.src[p0..][0..total]); + const mpos = std.mem.bytesAsSlice(i32, st[mpos_at..][0 .. 3 * max_rows * 4]); + for (0..3) |a| @memcpy(mpos[a * max_rows ..][0..total], im.pos[a * n + p0 ..][0..total]); + } + try e.tables.uploadAsync(0, st[0..if (images != null) tables_len else src_at], e.stream.handle); // text: as before + // 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 + e.features = if (images) |im| im.features else 0; + defer e.features = 0; + 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 }); + // an image chunk's rows take the tower's features, and rotate at their staged t, h, w positions + const mpos: u64 = if (e.features != 0) e.tables.ptr + mpos_at else 0; + if (e.features != 0) try e.run(e.fns.splice, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, e.features, e.tables.ptr + src_at, @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, mpos }); + 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); + try e.proj(l.gate, e.x, e.gate, total, prompt); + try e.proj(l.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(l.down, e.act, e.r, total, prompt); + 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); + } + + /// 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..fdd6ecb0c2 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -0,0 +1,186 @@ +//! 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; + +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, + + 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); + } + + /// 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); + try is(root, "model_type", []const u8, "qwen3_5"); + const t = try object(root.get("text_config") orelse return error.BadQwenConfig); + try is(t, "model_type", []const u8, "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, "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 (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 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_native.zig b/zig/src/families/qwen3_5/cuda_native.zig new file mode 100644 index 0000000000..306d84df4e --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -0,0 +1,443 @@ +//! 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 vision = @import("cuda_vision.zig"); +const vision_prompt = @import("cuda_vision_prompt.zig"); +const reuse = @import("cuda_reuse.zig"); +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; + +/// Kept prompt states for the shared cache (cuda_reuse.zig), in pinned host memory, on the zero-anchored chunk grid. +pub const snap_bytes = reuse.bytes; +pub const snap_save = reuse.save; +pub const snap_restore = reuse.restore; +pub const snap_drop = reuse.drop; +pub const snap_host = true; +pub const snap_lookahead: u32 = 0; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, vision: 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, + /// What the prompt cache's saves copy: the engine and the streams' caches. + target: *reuse.Target, +}; + +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 + // --vision: the checkpoint's tower, its weights and one call's scratch loaded beside the model's + const tower: ?*vision.Tower = if (!o.vision) null else blk: { + const vc = try vision.Config.read(gpa, io, dir) orelse return error.QwenCheckpointHasNoVision; + if (vc.out_hidden != e.g.hidden) return error.UnsupportedQwenVision; + break :blk try vision.Tower.load(gpa, io, ctx.d, e.stream, dir, vc); + }; + errdefer if (tower) |t| t.deinit(); + const l = try gpa.create(Lanes); + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .tower = tower, .target = undefined }; + l.target = .{ .e = e, .caches = &l.caches, .seen = &l.seen }; + 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, + .target = &l.target, + }; +} + +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 }; +} + +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.QwenVideoUnsupported => "video inputs are not supported by the Qwen CUDA engine yet: send images", + error.BadImagePlaceholders => "the image placeholders do not match the prepared images", + error.BadImageGrid, error.BadImagePatches => "an image's patches do not fit the vision tower", + error.QwenServedWithoutVision => "image inputs need the server started with --vision", + error.OutOfMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + else => null, + }; +} + +/// Where a prompt pass from `at` cuts its chunk: max_rows rows on, or the request's next chunk start before that. +fn nextCut(len: usize, chunks: []const u32, at: usize) usize { + var end: usize = @min(len, at + E.max_rows); + for (chunks) |boundary| if (boundary > at) { + end = @min(end, boundary); + break; + }; + return end; +} + +/// Whether a cold pass over `len` rows (from 0) cuts a chunk at `pos`. +fn coldCut(len: usize, chunks: []const u32, pos: usize) bool { + var at: usize = 0; + while (at < pos and at < len) at = nextCut(len, chunks, at); + return at == pos; +} + +test "a pass resumes only where a cold pass cuts: every max_rows rows from 0 or from the request's chunk starts" { + try std.testing.expect(coldCut(1000, &.{}, E.max_rows * 3)); + try std.testing.expect(!coldCut(1000, &.{}, E.max_rows * 3 + 1)); + try std.testing.expect(coldCut(1000, &.{200}, 200)); + try std.testing.expect(coldCut(1000, &.{200}, 200 + E.max_rows)); + try std.testing.expect(!coldCut(1000, &.{200}, 2 * E.max_rows)); // cut at 200 instead, then 328 + try std.testing.expectEqual(@as(usize, 200), nextCut(1000, &.{200}, E.max_rows)); +} + +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, + /// --vision's tower: an image prompt's images encoded at its prefill. + tower: ?*vision.Tower = null, + target: reuse.Target, + caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, + /// An image prompt's images as its kept states name them (cuda_reuse.zig), from its prefill to its release. + seen: std.AutoHashMapUnmanaged(*lanes.Stream, []reuse.Seen) = .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); + var seen = self.seen.valueIterator(); + while (seen.next()) |v| self.gpa.free(v.*); + self.seen.deinit(self.gpa); + for (self.pool.items) |cache| self.drop(cache); + self.pool.deinit(self.gpa); + if (self.tower) |t| t.deinit(); + 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: *Draft) !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; + // an image prompt: its images through the tower first, then every chunk at its t, h, w positions + var media: ?vision_prompt.Media = null; + if (stream.images.len > 0) { + const tower = self.tower orelse return error.QwenServedWithoutVision; + media = try vision_prompt.Media.init(self.gpa, tower, ids, stream.images); + } + defer if (media) |*m| m.deinit(); + const images: ?E.Images = if (media) |m| .{ .features = m.features.ptr, .src = m.src, .pos = m.pos } else null; + const seen = try self.seenOf(ids, stream.images); + self.seen.put(self.gpa, stream, seen) catch |err| { + self.gpa.free(seen); + return err; + }; + errdefer if (self.seen.fetchRemove(stream)) |kv| self.gpa.free(kv.value); + const cache = try self.take(); + errdefer self.give(cache); + // the stream's cache is found by a save at a mark (cuda_reuse.zig) while its pass runs + try self.caches.put(self.gpa, stream, cache); + errdefer _ = self.caches.remove(stream); + // a kept state resumes only where a cold pass cuts a chunk, so the chunks after it keep a cold pass's bits + var at: usize = 0; + stream.cached = 0; + if (stream.reuse.saved) |saved| { + if (stream.reuse.at > 0 and stream.reuse.at < ids.len and coldCut(ids.len, stream.chunks, stream.reuse.at)) { + if (reuse.restoreCache(self.e, cache, saved, seen)) { + at = stream.reuse.at; + stream.cached = stream.reuse.at; + } else |_| { + try cache.reset(self.e); + stream.reuse_failed = true; + } + } else stream.reuse_failed = true; + } + while (at < ids.len) { + const end = nextCut(ids.len, stream.chunks, at); + try self.e.forward(&.{.{ .cache = cache, .rows = end - at, .images = if (images) |*im| im else null }}, ids[at..end], .last); + at = end; + // a mark the store asked for, where this pass cut anyway: the lane host copies the state here + if (end < ids.len and std.mem.indexOfScalar(u32, stream.reuse.marks, @intCast(end)) != null) { + if (stream.reuse.hook) |k| k.at(k.ptr, stream, @intCast(end)); + } + } + // the reply's rows rotate past the image prompt's positions (rope_delta) + if (media) |m| cache.delta = m.delta; + } + + 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); + if (self.seen.fetchRemove(stream)) |entry| self.gpa.free(entry.value); + } + + /// Each image's first placeholder row and digest: a kept state resumes only the same images. + fn seenOf(self: *Lanes, ids: []const u32, images: []const lanes.Image) ![]reuse.Seen { + if (images.len == 0) return self.gpa.alloc(reuse.Seen, 0); + const tower = self.tower orelse return error.QwenServedWithoutVision; + const runs = try self.gpa.alloc(vision_prompt.Run, images.len); + defer self.gpa.free(runs); + try vision_prompt.placeholders(tower.c, ids, images, runs); + const seen = try self.gpa.alloc(reuse.Seen, images.len); + for (seen, runs, images) |*s, r, im| s.* = .{ .at = @intCast(r.at), .digest = im.digest() }; + return seen; + } +}; diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig new file mode 100644 index 0000000000..41f93ee298 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -0,0 +1,288 @@ +//! Qwen3.5 CUDA prompt states for the shared cache in pinned host memory: recurrences, tails, keys, values, rope shift. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Engine = engine.Engine; +const Cache = engine.Cache; +const Geometry = geometry.Geometry; +const head_dim = geometry.head_dim; + +/// A cache's attention rows: bf16, or e4m3 with an fp32 scale each (--kv-cache fp8, where the engine has it). +pub const Format = enum(u8) { bf16, fp8 }; + +/// The engine's cache format; an engine without fp8 caches keeps bf16 rows. +pub fn formatOf(e: *const Engine) Format { + if (comptime @hasField(Engine, "kv8")) return if (e.kv8) .fp8 else .bf16; + return .bf16; +} + +/// An image a kept state's prompt holds: the row its placeholders start at, its digest (lanes.Image.digest). +pub const Seen = struct { at: u32, digest: [32]u8 }; + +/// One kept state: its bytes on the host, the position it stands at, its rotary shift, its cache format and images. +pub const Copy = struct { + host: cuda.HostBuffer, + len: u64, + at: usize, + delta: i32, + format: Format, + images: []Seen, +}; + +/// The engine, the streams' caches a save copies, and their prompts' images (cuda_native.zig's Lanes owns the maps). +pub const Target = struct { + e: *Engine, + caches: *std.AutoHashMapUnmanaged(*lanes.Stream, *Cache), + seen: *std.AutoHashMapUnmanaged(*lanes.Stream, []Seen), +}; + +/// One piece of a stream's state: its offset in the state or key/value buffer, its length. +const Piece = struct { kv: bool, off: usize, len: usize }; + +/// A state's pieces by index, any geometry: the state, then each (layer, key/value, head)'s first `at` rows. +const Layout = struct { + g: Geometry, + format: Format, + cap: usize, + at: usize, + + fn parts(l: Layout) usize { + return if (l.format == .fp8) 2 else 1; + } + + fn count(l: Layout) usize { + return 1 + l.g.attentions() * 2 * l.g.kv_heads * l.parts(); + } + + /// Piece `i`: qwen35.cu's layout (a layer's keys then values, kv_heads blocks of cap rows), or qwen35_kv8.cu's. + fn piece(l: Layout, i: usize) Piece { + if (i == 0) return .{ .kv = false, .off = 0, .len = l.g.stateBytes() }; + const j = (i - 1) / l.parts(); + const kvh = l.g.kv_heads; + const block = j % (2 * kvh); // key heads, then value heads + const layer = j / (2 * kvh); + return switch (l.format) { + .bf16 => .{ .kv = true, .off = (layer * 2 * kvh + block) * l.cap * head_dim * 2, .len = l.at * head_dim * 2 }, + .fp8 => blk: { + const base = layer * 2 * kvh * l.cap * (head_dim + 4); + break :blk if ((i - 1) % 2 == 0) + .{ .kv = true, .off = base + block * l.cap * head_dim, .len = l.at * head_dim } + else + .{ .kv = true, .off = base + 2 * kvh * l.cap * head_dim + block * l.cap * 4, .len = l.at * 4 }; + }, + }; + } + + /// The bytes a save at `at` takes. + fn total(l: Layout) u64 { + var n: u64 = 0; + for (0..l.count()) |i| n += l.piece(i).len; + return n; + } + + /// The key/value buffer's bytes at `cap` (Geometry.kvBytes). + fn kvBytes(l: Layout) usize { + return l.g.attentions() * 2 * l.g.kv_heads * l.cap * @as(usize, if (l.format == .fp8) head_dim + 4 else head_dim * 2); + } +}; + +/// Whether a request's images match a kept state's: the same digests at the same rows, for every image before `at`. +pub fn sameImages(kept: []const Seen, request: []const Seen, at: usize) bool { + var n: usize = 0; + for (request) |r| { + if (r.at >= at) continue; + if (n >= kept.len or kept[n].at != r.at or !std.mem.eql(u8, &kept[n].digest, &r.digest)) return false; + n += 1; + } + return n == kept.len; +} + +/// Host bytes a save at `at` takes. +pub fn bytes(ptr: *anyopaque, at: u32) u64 { + const t: *Target = @ptrCast(@alignCast(ptr)); + return (Layout{ .g = t.e.g, .format = formatOf(t.e), .cap = 0, .at = at }).total(); +} + +/// Copy the stream's state after `at` prompt tokens; its prompt pass stands there. +pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const t: *Target = @ptrCast(@alignCast(ptr)); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const cache = t.caches.get(s) orelse return error.UnknownStream; + return @ptrCast(try saveCache(t.e, cache, at, t.seen.get(s) orelse &.{})); +} + +/// A host copy of `cache` at `at` rows, naming the prompt's images that start before it. +pub fn saveCache(e: *Engine, cache: *const Cache, at: usize, seen: []const Seen) !*Copy { + if (cache.len != at or cache.last != null or at > cache.cap) return error.NotAtMark; + const l: Layout = .{ .g = e.g, .format = formatOf(e), .cap = cache.cap, .at = at }; + var n: usize = 0; + for (seen) |im| n += @intFromBool(im.at < at); + const copy = try e.gpa.create(Copy); + errdefer e.gpa.destroy(copy); + const images = try e.gpa.alloc(Seen, n); + errdefer e.gpa.free(images); + n = 0; + for (seen) |im| if (im.at < at) { + images[n] = im; + n += 1; + }; + copy.* = .{ .host = try cuda.HostBuffer.alloc(e.d, l.total()), .len = l.total(), .at = at, .delta = cache.delta, .format = l.format, .images = images }; + errdefer copy.host.free(); + var off: usize = 0; + for (0..l.count()) |i| { + const p = l.piece(i); + const dev = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off; + if (p.len > 0) try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off, dev, p.len, e.stream.handle), "cuMemcpyDtoHAsync"); + off += p.len; + } + try e.stream.synchronize(); + return copy; +} + +/// Whether `cache` can take the kept state: the same cache format and bytes, a position inside it, the same images. +pub fn restorable(e: *const Engine, cache: *const Cache, copy: *const Copy, seen: []const Seen) bool { + if (copy.format != formatOf(e) or copy.at >= cache.cap) return false; + const l: Layout = .{ .g = e.g, .format = copy.format, .cap = cache.cap, .at = copy.at }; + return copy.len == l.total() and copy.host.bytes.len >= copy.len and sameImages(copy.images, seen, copy.at); +} + +/// Make `cache` the kept state's: its next prompt chunk starts at the copy's position; refused before any copy. +pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque, seen: []const Seen) !void { + const copy: *Copy = @ptrCast(@alignCast(saved)); + if (copy.format != formatOf(e)) return error.KeptCacheFormat; + if (!restorable(e, cache, copy, seen)) return error.BadReuse; + const l: Layout = .{ .g = e.g, .format = copy.format, .cap = cache.cap, .at = copy.at }; + var off: usize = 0; + for (0..l.count()) |i| { + const p = l.piece(i); + const dev = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off; + if (p.len > 0) try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(dev, copy.host.bytes.ptr + off, p.len, e.stream.handle), "cuMemcpyHtoDAsync"); + off += p.len; + } + try e.stream.synchronize(); // the host copy may be dropped once this returns + cache.len = copy.at; + cache.written = copy.at; + cache.last = null; + cache.delta = copy.delta; +} + +/// The store's restore: the owner's cache made the kept state's. +pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + const t: *Target = @ptrCast(@alignCast(ptr)); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + return restoreCache(t.e, t.caches.get(s) orelse return error.UnknownStream, saved, t.seen.get(s) orelse &.{}); +} + +/// Free a kept state's host copy. +pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { + const t: *Target = @ptrCast(@alignCast(ptr)); + dropCopy(t.e, @ptrCast(@alignCast(saved))); +} + +pub fn dropCopy(e: *Engine, copy: *Copy) void { + copy.host.free(); + e.gpa.free(copy.images); + e.gpa.destroy(copy); +} + +fn testGeometry(layers: usize, attention_every: usize, kv_heads: usize) Geometry { + var g: Geometry = undefined; + g.hidden = 5120; + g.k_heads = 16; + g.v_heads = 48; + g.q_heads = 32; + g.kv_heads = kv_heads; + g.layers = layers; + @memset(&g.linear, false); + for (0..layers) |i| g.linear[i] = i % attention_every != attention_every - 1; + return g; +} + +test "a kept state holds the recurrences and tails whole and each attention head's rows it read" { + const g = testGeometry(8, 4, 4); // the 27B's pattern: three linear layers, then attention + const l: Layout = .{ .g = g, .format = .bf16, .cap = 2048, .at = 1000 }; + try std.testing.expectEqual(@as(usize, 1 + 2 * 2 * 4), l.count()); + try std.testing.expectEqual(@as(u64, g.stateBytes()), (Layout{ .g = g, .format = .bf16, .cap = 2048, .at = 0 }).total()); + try std.testing.expectEqual(@as(u64, g.stateBytes() + 2 * 2 * 4 * 1000 * head_dim * 2), l.total()); + // fp8: e4m3 rows and a scale each, 260 bytes a position and head + const l8: Layout = .{ .g = g, .format = .fp8, .cap = 2048, .at = 1000 }; + try std.testing.expectEqual(@as(usize, 1 + 2 * 2 * 2 * 4), l8.count()); + try std.testing.expectEqual(@as(u64, g.stateBytes() + 2 * 2 * 4 * 1000 * (head_dim + 4)), l8.total()); +} + +test "every geometry's pieces lie inside its allocation, apart, with no fixed limit on layers or heads" { + const a = std.testing.allocator; + // 33 attention layers of 32 key heads passed the old fixed array; the widest the engine reads is 128 layers + for ([_][3]usize{ .{ 8, 4, 4 }, .{ 99, 3, 32 }, .{ 33, 1, 32 }, .{ 128, 1, 32 }, .{ 24, 4, 2 } }) |shape| { + const g = testGeometry(shape[0], shape[1], shape[2]); + for ([_]Format{ .bf16, .fp8 }) |f| for ([_]usize{ 1, 7, 64 }) |at| { + const l: Layout = .{ .g = g, .format = f, .cap = 64, .at = at }; + const end = l.kvBytes(); + if (comptime @typeInfo(@TypeOf(Geometry.kvBytes)).@"fn".param_types.len == 3) try std.testing.expectEqual(g.kvBytes(64, f == .fp8), end); + const spans = try a.alloc([2]usize, l.count() - 1); + defer a.free(spans); + try std.testing.expectEqual(@as(usize, g.stateBytes()), l.piece(0).len); + for (spans, 1..) |*sp, i| { + const p = l.piece(i); + try std.testing.expect(p.kv and p.len > 0 and p.off + p.len <= end); + sp.* = .{ p.off, p.off + p.len }; + } + std.mem.sort([2]usize, spans, {}, struct { + fn lt(_: void, x: [2]usize, y: [2]usize) bool { + return x[0] < y[0]; + } + }.lt); + for (spans[1..], spans[0 .. spans.len - 1]) |sp, prev| try std.testing.expect(sp[0] >= prev[1]); + }; + } +} + +test "fp8 pieces follow qwen35_kv8.cu: a layer's key then value rows, then their scales" { + const g = testGeometry(4, 4, 2); // one attention layer + const l: Layout = .{ .g = g, .format = .fp8, .cap = 100, .at = 10 }; + // (key, head 1): rows at 1 * cap * 256, its scales after both blocks' rows, at 2 * 2 * cap * 256 + 1 * cap * 4 + try std.testing.expectEqual(Piece{ .kv = true, .off = 100 * 256, .len = 10 * 256 }, l.piece(1 + 2 * 1)); + try std.testing.expectEqual(Piece{ .kv = true, .off = 4 * 100 * 256 + 100 * 4, .len = 10 * 4 }, l.piece(2 + 2 * 1)); + // (value, head 0): the third block's rows and scales + try std.testing.expectEqual(Piece{ .kv = true, .off = 2 * 100 * 256, .len = 10 * 256 }, l.piece(1 + 2 * 2)); + try std.testing.expectEqual(Piece{ .kv = true, .off = 4 * 100 * 256 + 2 * 100 * 4, .len = 10 * 4 }, l.piece(2 + 2 * 2)); +} + +test "a state resumes only a request with the same images before its position" { + const d = [_][32]u8{ @splat(1), @splat(2), @splat(3) }; + const kept = [_]Seen{.{ .at = 5, .digest = d[0] }}; + try std.testing.expect(sameImages(&kept, &.{ .{ .at = 5, .digest = d[0] }, .{ .at = 40, .digest = d[2] } }, 20)); // a later image is the pass's + try std.testing.expect(!sameImages(&kept, &.{.{ .at = 5, .digest = d[1] }}, 20)); // another image under the same ids + try std.testing.expect(!sameImages(&kept, &.{.{ .at = 6, .digest = d[0] }}, 20)); + try std.testing.expect(!sameImages(&kept, &.{}, 20)); // a text prompt never resumes an image's state + try std.testing.expect(!sameImages(&.{}, &.{.{ .at = 5, .digest = d[0] }}, 20)); // nor an image prompt a text state past its image + try std.testing.expect(sameImages(&.{}, &.{.{ .at = 5, .digest = d[0] }}, 5)); // a text prefix before it +} + +test "a restore refuses another cache format, another length, a position past the cache and other images" { + var e: Engine = undefined; + e.g = testGeometry(8, 4, 4); + if (comptime @hasField(Engine, "kv8")) e.kv8 = false; + var cache: Cache = undefined; + cache.cap = 64; + const d: [32]u8 = @splat(9); + var images = [_]Seen{.{ .at = 2, .digest = d }}; + const l: Layout = .{ .g = e.g, .format = .bf16, .cap = 64, .at = 10 }; + const host = try std.testing.allocator.alignedAlloc(u8, .@"16", @intCast(l.total())); + defer std.testing.allocator.free(host); + var copy: Copy = .{ .host = .{ .d = undefined, .bytes = host }, .len = l.total(), .at = 10, .delta = 0, .format = .bf16, .images = &images }; + const same = [_]Seen{.{ .at = 2, .digest = d }}; + try std.testing.expect(restorable(&e, &cache, ©, &same)); + try std.testing.expect(!restorable(&e, &cache, ©, &.{.{ .at = 2, .digest = @splat(8) }})); + copy.format = .fp8; // a state kept from fp8 rows never lands in bf16 ones, nor the other way + try std.testing.expect(!restorable(&e, &cache, ©, &same)); + copy.format = .bf16; + copy.len -= 1; + try std.testing.expect(!restorable(&e, &cache, ©, &same)); + copy.len += 1; + cache.cap = 10; + try std.testing.expect(!restorable(&e, &cache, ©, &same)); +} diff --git a/zig/src/families/qwen3_5/cuda_vision.zig b/zig/src/families/qwen3_5/cuda_vision.zig new file mode 100644 index 0000000000..e790373d85 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_vision.zig @@ -0,0 +1,472 @@ +//! The Qwen3.5/3.8 vision tower on CUDA (transformers' Qwen3_5VisionModel, bf16); prompts: cuda_vision_prompt.zig. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const lt = cuda.cublaslt; +const DeviceBuffer = cuda.DeviceBuffer; +const prefix = "vision_tower."; + +/// What every Qwen3.5/3.8 tower shares (vision_config); its width, heads, depth and MLP are the checkpoint's. +pub const patch = 16; +pub const temporal = 2; +pub const merge = 2; +pub const channels = 3; +pub const patch_width = channels * temporal * patch * patch; // 1536 +pub const grid_side = 48; // sqrt(num_position_embeddings 2304) +pub const eps: f32 = 1e-6; +/// Blocks a tower has at most; its heads are 64 or 72 wide (qwen_vision.cu's attention instances). +pub const max_depth = 32; +/// One tower call's patches at most: a 4,096-token image's. +pub const max_patches = 16384; +const workspace_bytes = 32 << 20; + +/// What config.json says about the tower and its placeholders. +pub const Config = struct { + hidden: usize, + heads: usize, + depth: usize, + intermediate: usize, + out_hidden: usize, + image_token: u32, + video_token: u32, + vision_start: u32, + vision_end: u32, + + pub fn headDim(c: Config) usize { + return c.hidden / c.heads; + } + + /// The merger's width: four merge-ordered rows read as one. + pub fn merged(c: Config) usize { + return c.hidden * merge * merge; + } + + /// A checkpoint without a vision_config is null; a tower or M-RoPE this module does not run is refused. + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !?Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !?Config { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = parsed.value.object; + const v = (root.get("vision_config") orelse return null).object; + const want = [_]struct { []const u8, i64 }{ .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "spatial_merge_size", merge }, .{ "in_channels", channels }, .{ "num_position_embeddings", grid_side * grid_side } }; + for (want) |w| if (num(v.get(w[0])) != w[1]) return error.UnsupportedQwenVision; + const act = v.get("hidden_act") orelse return error.UnsupportedQwenVision; + if (act != .string or !std.mem.eql(u8, act.string, "gelu_pytorch_tanh")) return error.UnsupportedQwenVision; + if (v.get("deepstack_visual_indexes")) |ds| if (ds != .array or ds.array.items.len != 0) return error.UnsupportedQwenVision; + // qwen35.cu's rotate: interleaved M-RoPE over sections 11, 11 and 10 + const text = if (root.get("text_config")) |t| t.object else root; + const rope = (text.get("rope_parameters") orelse return error.UnsupportedQwenVision).object; + const inter = rope.get("mrope_interleaved") orelse return error.UnsupportedQwenVision; + const sec = rope.get("mrope_section") orelse return error.UnsupportedQwenVision; + if (inter != .bool or !inter.bool or sec != .array or sec.array.items.len != 3) return error.UnsupportedQwenVision; + for (sec.array.items, [_]i64{ 11, 11, 10 }) |s, x| if (num(s) != x) return error.UnsupportedQwenVision; + const out = num(v.get("out_hidden_size")); + const hid = num(v.get("hidden_size")); + const nh = num(v.get("num_heads")); + const dep = num(v.get("depth")); + const mlp = num(v.get("intermediate_size")); + if (out <= 0 or hid <= 0 or nh <= 0 or mlp <= 0 or dep <= 0 or dep > max_depth) return error.UnsupportedQwenVision; + if (@mod(hid, nh) != 0 or (@divExact(hid, nh) != 64 and @divExact(hid, nh) != 72)) return error.UnsupportedQwenVision; + return .{ .hidden = @intCast(hid), .heads = @intCast(nh), .depth = @intCast(dep), .intermediate = @intCast(mlp), .out_hidden = @intCast(out), .image_token = try id(root, "image_token_id"), .video_token = try id(root, "video_token_id"), .vision_start = try id(root, "vision_start_token_id"), .vision_end = try id(root, "vision_end_token_id") }; + } + + fn num(x: ?std.json.Value) i64 { + const v = x orelse return -1; + return if (v == .integer) v.integer else -1; + } + + fn id(o: std.json.ObjectMap, name: []const u8) !u32 { + const v = num(o.get(name)); + if (v < 0) return error.UnsupportedQwenVision; + return @intCast(v); + } +}; + +/// One image's patch grid (rows and columns of 16x16 patches, both even). +pub const Grid = struct { h: u32, w: u32 }; + +const Block = struct { norm1_w: u64, norm1_b: u64, qkv_w: u64, qkv_b: u64, proj_w: u64, proj_b: u64, norm2_w: u64, norm2_b: u64, fc1_w: u64, fc1_b: u64, fc2_w: u64, fc2_b: u64 }; + +const Fns = struct { + embed: cuda.Function, + layernorm: cuda.Function, + qkv_rope: cuda.Function, + attention: cuda.Function, + bias_gelu: cuda.Function, + bias_residual: cuda.Function, + bias_cast: cuda.Function, +}; + +/// The tower's weights and one call's scratch in one allocation, its kernels; `encode` turns patches into model rows. +pub const Tower = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + s: cuda.Stream, + c: Config, + mod: cuda.Module, + f: Fns, + lib: lt.Library, + arena: DeviceBuffer, + at: usize = 0, + scratch: u64 = 0, + patch_w: u64 = 0, + patch_b: u64 = 0, + pos: u64 = 0, + blocks: [max_depth]Block = undefined, + merger_norm_w: u64 = 0, + merger_norm_b: u64 = 0, + fc1_w: u64 = 0, + fc1_b: u64 = 0, + fc2_w: u64 = 0, + fc2_b: u64 = 0, + inv: u64 = 0, + + /// The checkpoint's vision_tower tensors (bf16) on stream `s`; the patch convolution's weight as [C, T, y, x] rows. + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, dir: []const u8, c: Config) !*Tower { + const t = try gpa.create(Tower); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .d = d, .s = s, .c = c, .mod = undefined, .f = undefined, .lib = undefined, .arena = undefined }; + t.lib = try lt.Library.open(); + errdefer t.lib.close(); + t.mod = try cuda.Module.load(d, cuda.kernels.qwen_vision); + errdefer t.mod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(t.f, name) = try t.mod.function(if (comptime std.mem.eql(u8, name, "attention")) "tf_qv_attention72" else "tf_qv_" ++ name); + if (c.headDim() == 64) t.f.attention = try t.mod.function("tf_qv_attention64"); + const hidden = c.hidden; + const intermediate = c.intermediate; + const merged = c.merged(); + t.arena = try DeviceBuffer.alloc(d, weightBytes(c) + scratchBytes(c, max_patches)); + errdefer t.arena.free(); + try t.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly + var ck = try core.Checkpoint.openModelPrefix(gpa, io, dir, prefix); + defer ck.close(); + { // [hidden, T, P, P, C] as stored (MLX) to [hidden, C, T, P, P] + const tn = try ck.expect(prefix ++ "patch_embed.proj.weight", .bf16, &.{ hidden, temporal, patch, patch, channels }); + const host = try gpa.alloc(u16, hidden * patch_width); + defer gpa.free(host); + const stored: []align(1) const u16 = std.mem.bytesAsSlice(u16, tn.bytes); + for (0..hidden) |o| for (0..channels) |ch| for (0..temporal) |tt| for (0..patch) |y| for (0..patch) |x| { + host[o * patch_width + ((ch * temporal + tt) * patch + y) * patch + x] = stored[(((o * temporal + tt) * patch + y) * patch + x) * channels + ch]; + }; + t.patch_w = try t.put(std.mem.sliceAsBytes(host)); + } + t.patch_b = try t.get(&ck, "patch_embed.proj.bias", &.{hidden}); + t.pos = try t.get(&ck, "pos_embed.weight", &.{ grid_side * grid_side, hidden }); + for (t.blocks[0..c.depth], 0..) |*b, i| { + const shapes = .{ + .{ "norm1_w", "norm1.weight", .{hidden} }, .{ "norm1_b", "norm1.bias", .{hidden} }, + .{ "qkv_w", "attn.qkv.weight", .{ 3 * hidden, hidden } }, .{ "qkv_b", "attn.qkv.bias", .{3 * hidden} }, + .{ "proj_w", "attn.proj.weight", .{ hidden, hidden } }, .{ "proj_b", "attn.proj.bias", .{hidden} }, + .{ "norm2_w", "norm2.weight", .{hidden} }, .{ "norm2_b", "norm2.bias", .{hidden} }, + .{ "fc1_w", "mlp.linear_fc1.weight", .{ intermediate, hidden } }, .{ "fc1_b", "mlp.linear_fc1.bias", .{intermediate} }, + .{ "fc2_w", "mlp.linear_fc2.weight", .{ hidden, intermediate } }, .{ "fc2_b", "mlp.linear_fc2.bias", .{hidden} }, + }; + inline for (shapes) |sh| { + var nb: [96]u8 = undefined; + const shape: [sh[2].len]usize = sh[2]; + @field(b, sh[0]) = try t.get(&ck, try std.fmt.bufPrint(&nb, "blocks.{d}.{s}", .{ i, sh[1] }), &shape); + } + } + t.merger_norm_w = try t.get(&ck, "merger.norm.weight", &.{hidden}); + t.merger_norm_b = try t.get(&ck, "merger.norm.bias", &.{hidden}); + t.fc1_w = try t.get(&ck, "merger.linear_fc1.weight", &.{ merged, merged }); + t.fc1_b = try t.get(&ck, "merger.linear_fc1.bias", &.{merged}); + t.fc2_w = try t.get(&ck, "merger.linear_fc2.weight", &.{ c.out_hidden, merged }); + t.fc2_b = try t.get(&ck, "merger.linear_fc2.bias", &.{c.out_hidden}); + if (ck.unused() != 0) return error.UnusedVisionTensors; + var inv: [72 / 4]f32 = undefined; + invFreq(inv[0 .. c.headDim() / 4], c.headDim()); + t.inv = try t.put(std.mem.sliceAsBytes(inv[0 .. c.headDim() / 4])); + t.scratch = try t.take(scratchBytes(c, max_patches)); + return t; + } + + pub fn deinit(t: *Tower) void { + t.arena.free(); + t.mod.unload(); + t.lib.close(); + t.gpa.destroy(t); + } + + /// The weights' bytes in the arena (each allocation 256-byte aligned). + pub fn weightBytes(c: Config) usize { + const a = aligned; + const hidden = c.hidden; + const intermediate = c.intermediate; + const merged = c.merged(); + const out_hidden = c.out_hidden; + var n: usize = a(hidden * patch_width * 2) + a(hidden * 2) + a(grid_side * grid_side * hidden * 2); + n += c.depth * (4 * a(hidden * 2) + a(3 * hidden * hidden * 2) + a(3 * hidden * 2) + a(hidden * hidden * 2) + a(hidden * 2) + a(intermediate * hidden * 2) + a(intermediate * 2) + a(hidden * intermediate * 2) + a(hidden * 2)); + n += 2 * a(hidden * 2) + a(merged * merged * 2) + a(merged * 2) + a(out_hidden * merged * 2) + a(out_hidden * 2); + return n + a(c.headDim() / 4 * 4); + } + + /// The scratch one tower call of `n` patches needs, in bytes (the patches' bf16 copy included). + pub fn scratchBytes(c: Config, n: usize) usize { + const hidden = c.hidden; + const intermediate = c.intermediate; + return n * patch_width * 2 + // patches + 2 * n * hidden * 2 + // x, h + n * @max(3 * hidden, intermediate) * 4 + // fp32 Linear out + 3 * n * hidden * 2 + // q, k, v + n * hidden * 2 + // attention out + n * @max(intermediate, hidden) * 2 + // gelu rows + n * (2 + 4) * 4 + n * 4 * 4 + // hw, idx, weights + 16 * 256 + workspace_bytes; + } + + fn take(t: *Tower, len: usize) !u64 { + const at = aligned(t.at); + if (at + len > t.arena.len) return error.ArenaExhausted; + t.at = at + len; + return t.arena.ptr + at; + } + + fn put(t: *Tower, bytes: []const u8) !u64 { + const p = try t.take(bytes.len); + try t.arena.upload(p - t.arena.ptr, bytes); + return p; + } + + fn get(t: *Tower, ck: *core.Checkpoint, name: []const u8, shape: []const usize) !u64 { + var buf: [96]u8 = undefined; + const tn = try ck.expect(try std.fmt.bufPrint(&buf, prefix ++ "{s}", .{name}), .bf16, shape); + return t.put(tn.bytes); + } + + fn upload(t: *Tower, dst: u64, bytes: []const u8) !void { + if (bytes.len == 0) return; + try t.d.check(t.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, t.s.handle), "cuMemcpyHtoDAsync"); + } + + /// Images' patches (fp32, merge order) to merged bf16 [sum(h * w) / 4, out_hidden] at `out`, each image on its own. + pub fn encode(t: *Tower, pixels: []const f32, grids: []const Grid, out: u64) !void { + var n: usize = 0; + for (grids) |g| { + if (g.h == 0 or g.w == 0 or g.h % merge != 0 or g.w % merge != 0) return error.BadImageGrid; + n += @as(usize, g.h) * g.w; + } + if (n == 0 or n > max_patches or pixels.len != n * patch_width) return error.BadImagePatches; + const gpa = t.gpa; + const c = t.c; + const out_hidden = c.out_hidden; + const hidden = c.hidden; + const heads = c.heads; + const head_dim = c.headDim(); + const intermediate = c.intermediate; + const merged = c.merged(); + // host side: bf16 patches, each patch's (row, column) and its 4 position-grid taps + const px16 = try gpa.alloc(u16, n * patch_width); + defer gpa.free(px16); + for (px16, pixels) |*o, v| o.* = bf16(v); + const hw = try gpa.alloc(i32, n * 2); + defer gpa.free(hw); + const idx = try gpa.alloc(i32, n * 4); + defer gpa.free(idx); + const wts = try gpa.alloc(f32, n * 4); + defer gpa.free(wts); + var at: usize = 0; + for (grids) |g| { + const len = @as(usize, g.h) * g.w; + positions(g, hw[at * 2 ..][0 .. len * 2], idx[at * 4 ..][0 .. len * 4], wts[at * 4 ..][0 .. len * 4]); + at += len; + } + var cur: u64 = t.scratch; + const end = t.scratch + scratchBytes(c, max_patches); + const take_ = struct { + fn f(p: *u64, len: usize) u64 { + const r = p.*; + p.* += std.mem.alignForward(usize, len, 256); + return r; + } + }.f; + const d_px = take_(&cur, n * patch_width * 2); + const x = take_(&cur, n * hidden * 2); + const h = take_(&cur, n * hidden * 2); + const acc = take_(&cur, n * @max(3 * hidden, intermediate) * 4); + const q = take_(&cur, n * hidden * 2); + const k = take_(&cur, n * hidden * 2); + const v = take_(&cur, n * hidden * 2); + const att = take_(&cur, n * hidden * 2); + const mid = take_(&cur, n * @max(intermediate, hidden) * 2); + const d_hw = take_(&cur, n * 2 * 4); + const d_idx = take_(&cur, n * 4 * 4); + const d_wts = take_(&cur, n * 4 * 4); + const ws = cur; + const ws_len = end - cur; + try t.upload(d_px, std.mem.sliceAsBytes(px16)); + try t.upload(d_hw, std.mem.sliceAsBytes(hw)); + try t.upload(d_idx, std.mem.sliceAsBytes(idx)); + try t.upload(d_wts, std.mem.sliceAsBytes(wts)); + + var mm: Gemms = .{ .lib = &t.lib }; + defer mm.deinit(); + const m4 = n / (merge * merge); + try mm.run(n, hidden, patch_width, d_px, t.patch_w, acc, ws, ws_len, t.s); + try t.go(t.f.embed, n, 1, 256, .{ acc, t.patch_b, t.pos, d_idx, d_wts, x, int(hidden) }); + const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(head_dim))); + for (t.blocks[0..c.depth]) |b| { + try t.go(t.f.layernorm, n, 1, 256, .{ x, b.norm1_w, b.norm1_b, h, int(hidden), eps }); + try mm.run(n, 3 * hidden, hidden, h, b.qkv_w, acc, ws, ws_len, t.s); + try t.go(t.f.qkv_rope, n, heads, 128, .{ acc, b.qkv_b, d_hw, t.inv, q, k, v, int(n), int(heads), int(head_dim) }); + var s0: usize = 0; + for (grids) |g| { + const len = @as(usize, g.h) * g.w; + try t.go(t.f.attention, (len + 63) / 64, heads, 64, .{ q, k, v, att, int(n), int(heads), int(s0), int(len), scale }); + s0 += len; + } + try mm.run(n, hidden, hidden, att, b.proj_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_residual, n, 1, 256, .{ acc, b.proj_b, x, int(hidden) }); + try t.go(t.f.layernorm, n, 1, 256, .{ x, b.norm2_w, b.norm2_b, h, int(hidden), eps }); + try mm.run(n, intermediate, hidden, h, b.fc1_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_gelu, n, 1, 256, .{ acc, b.fc1_b, mid, int(intermediate), int(1) }); + try mm.run(n, hidden, intermediate, mid, b.fc2_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_residual, n, 1, 256, .{ acc, b.fc2_b, x, int(hidden) }); + } + // merger: LayerNorm over 1152, four merge-ordered rows read as one of 4608 + try t.go(t.f.layernorm, n, 1, 256, .{ x, t.merger_norm_w, t.merger_norm_b, h, int(hidden), eps }); + try mm.run(m4, merged, merged, h, t.fc1_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_gelu, m4, 1, 256, .{ acc, t.fc1_b, mid, int(merged), int(0) }); + try mm.run(m4, out_hidden, merged, mid, t.fc2_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_cast, m4, 1, 256, .{ acc, t.fc2_b, out, int(out_hidden) }); + try t.s.synchronize(); // the host's patch rows and tables are freed on return + } + + fn go(t: *Tower, f: cuda.Function, gx: usize, gy: usize, block: u32, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(gx), .y = @intCast(gy) }, .block = .{ .x = block } }, t.s, &a); + } +}; + +/// cuBLASLt plans for the shapes one tower call runs, made on first use. +const Gemms = struct { + lib: *const lt.Library, + plans: [8]?struct { m: usize, n: usize, k: usize, lin: lt.Linear } = @splat(null), + + fn run(g: *Gemms, m: usize, n: usize, k: usize, x: u64, w: u64, out: u64, ws: u64, ws_len: usize, s: cuda.Stream) !void { + for (&g.plans) |*p| { + if (p.*) |*have| { + if (have.m == m and have.n == n and have.k == k) return have.lin.run(x, w, out, ws, ws_len, s.handle); + continue; + } + p.* = .{ .m = m, .n = n, .k = k, .lin = try lt.Linear.init(g.lib, m, n, k, .f32, ws_len) }; + return p.*.?.lin.run(x, w, out, ws, ws_len, s.handle); + } + return error.TooManyShapes; + } + + fn deinit(g: *Gemms) void { + for (&g.plans) |*p| if (p.*) |*have| have.lin.deinit(); + } +}; + +fn aligned(n: usize) usize { + return std.mem.alignForward(usize, n, 256); +} + +fn int(x: usize) c_int { + return @intCast(x); +} + +/// fp32 to bf16, round to nearest even (torch's .to(bfloat16)). +fn bf16(v: f32) u16 { + const b: u32 = @bitCast(v); + if (std.math.isNan(v)) return 0x7fc0; + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// The axial rotary's frequencies: 1 / 10000 ** (i / (head_dim / 2)) for even i, in fp32 as torch computes them. +pub fn invFreq(out: []f32, head_dim: usize) void { + const dim: f32 = @floatFromInt(head_dim / 2); + for (out, 0..) |*f, j| { + const e: f32 = @as(f32, @floatFromInt(2 * j)) / dim; + f.* = 1.0 / std.math.pow(f32, 10000.0, e); + } +} + +/// One image's patches in merge order: each one's (row, column) for the rotary and 4 bilinear taps into the 48x48 grid. +pub fn positions(g: Grid, hw: []i32, idx: []i32, wts: []f32) void { + const bw = g.w / merge; + const n = @as(usize, g.h) * g.w; + for (0..n) |i| { + const in_col: u32 = @intCast(i % merge); + const in_row: u32 = @intCast((i / merge) % merge); + const block_col: u32 = @intCast((i / (merge * merge)) % bw); + const block_row: u32 = @intCast(i / (merge * merge * bw)); + const r = block_row * merge + in_row; + const c = block_col * merge + in_col; + hw[i * 2] = @intCast(r); + hw[i * 2 + 1] = @intCast(c); + const ht = taps(r, g.h); + const wt = taps(c, g.w); + for (0..2) |a| for (0..2) |b| { + idx[i * 4 + a * 2 + b] = ht.i[a] * grid_side + wt.i[b]; + wts[i * 4 + a * 2 + b] = ht.w[a] * wt.w[b]; + }; + } +} + +const Taps = struct { i: [2]i32, w: [2]f32 }; + +fn taps(index: u32, size: u32) Taps { + const side_m1: f32 = grid_side - 1; + const denom: f32 = @floatFromInt(@max(size, 2) - 1); + const src = @as(f32, @floatFromInt(index)) * side_m1 / denom; + const fl = @floor(src); + var out: Taps = undefined; + for (0..2) |o| { + const raw: i32 = @as(i32, @intFromFloat(fl)) + @as(i32, @intCast(o)); + out.i[o] = std.math.clamp(raw, 0, grid_side - 1); + const dist = @abs(src - fl - @as(f32, @floatFromInt(o))); + out.w[o] = @max(1.0 - dist, 0.0); + } + return out; +} + +test "bilinear taps match align_corners resampling at the ends and middle" { + const a = taps(0, 10); + try std.testing.expectEqual(@as(i32, 0), a.i[0]); + try std.testing.expectEqual(@as(f32, 1), a.w[0]); + const b = taps(9, 10); + try std.testing.expectEqual(@as(i32, 47), b.i[0]); + try std.testing.expectEqual(@as(f32, 1), b.w[0]); +} + +test "a vision_config's tower is read; other widths, depths and M-RoPE are refused; a text-only config has none" { + const gpa = std.testing.allocator; + try std.testing.expectEqual(@as(?Config, null), try Config.parse(gpa, "{\"hidden_size\":2048}")); + const good = + \\{"vision_config":{"depth":27,"hidden_size":1152,"num_heads":16,"intermediate_size":4304,"patch_size":16, + \\"temporal_patch_size":2,"spatial_merge_size":2,"in_channels":3,"num_position_embeddings":2304, + \\"hidden_act":"gelu_pytorch_tanh","deepstack_visual_indexes":[],"out_hidden_size":5120}, + \\"text_config":{"rope_parameters":{"mrope_interleaved":true,"mrope_section":[11,11,10]}}, + \\"image_token_id":248056,"video_token_id":248057,"vision_start_token_id":248053,"vision_end_token_id":248054} + ; + const c = (try Config.parse(gpa, good)).?; + try std.testing.expectEqual(@as(usize, 5120), c.out_hidden); + try std.testing.expectEqual(@as(u32, 248056), c.image_token); + const deeper = try std.mem.replaceOwned(u8, gpa, good, "\"depth\":27", "\"depth\":33"); + defer gpa.free(deeper); + try std.testing.expectError(error.UnsupportedQwenVision, Config.parse(gpa, deeper)); + const wide = try std.mem.replaceOwned(u8, gpa, good, "\"num_heads\":16", "\"num_heads\":12"); + defer gpa.free(wide); + try std.testing.expectError(error.UnsupportedQwenVision, Config.parse(gpa, wide)); + // the 2B's tower: 24 blocks of 1024, heads of 64 + const small = try std.mem.replaceOwned(u8, gpa, good, "\"depth\":27,\"hidden_size\":1152,\"num_heads\":16,\"intermediate_size\":4304", "\"depth\":24,\"hidden_size\":1024,\"num_heads\":16,\"intermediate_size\":4096"); + defer gpa.free(small); + const c2 = (try Config.parse(gpa, small)).?; + try std.testing.expectEqual(@as(usize, 64), c2.headDim()); + try std.testing.expectEqual(@as(usize, 24), c2.depth); + const sections = try std.mem.replaceOwned(u8, gpa, good, "[11,11,10]", "[16,24,24]"); + defer gpa.free(sections); + try std.testing.expectError(error.UnsupportedQwenVision, Config.parse(gpa, sections)); +} diff --git a/zig/src/families/qwen3_5/cuda_vision_prompt.zig b/zig/src/families/qwen3_5/cuda_vision_prompt.zig new file mode 100644 index 0000000000..3aa5f3b4c4 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_vision_prompt.zig @@ -0,0 +1,175 @@ +//! An image prompt's rows for the CUDA engine: its images through the tower, its M-RoPE positions (get_rope_index). +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const vision = @import("cuda_vision.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const Config = vision.Config; +const Tower = vision.Tower; +const Grid = vision.Grid; +const merge = vision.merge; +const patch_width = vision.patch_width; +const max_patches = vision.max_patches; + +/// An image prompt's encoded rows: the tower's features on the device, each prompt row's feature row and positions. +pub const Media = struct { + gpa: std.mem.Allocator, + features: DeviceBuffer, // bf16 [image rows, out_hidden] + src: []i32, // [prompt]: feature row or -1 + pos: []i32, // [3][prompt]: t, h, w + delta: i32, + + /// `images` (the request's, in prompt order) through the tower; the prompt's positions as get_rope_index lays them. + pub fn init(gpa: std.mem.Allocator, tower: *Tower, prompt: []const u32, images: []const lanes.Image) !Media { + const c = tower.c; + const n = prompt.len; + const runs = try gpa.alloc(Run, images.len); + defer gpa.free(runs); + try placeholders(c, prompt, images, runs); + var rows: usize = 0; + for (runs) |r| rows += r.tokens(); + var m: Media = .{ .gpa = gpa, .features = try DeviceBuffer.alloc(tower.d, @max(rows, 1) * c.out_hidden * 2), .src = &.{}, .pos = &.{}, .delta = 0 }; + errdefer m.features.free(); + m.src = try gpa.alloc(i32, n); + errdefer gpa.free(m.src); + m.pos = try gpa.alloc(i32, 3 * n); + errdefer gpa.free(m.pos); + // the tower: runs of whole images within one call's patches, in prompt order + var done: usize = 0; // feature rows written + var i: usize = 0; + while (i < images.len) { + var j = i; + var patches: usize = 0; + while (j < images.len and (j == i or patches + runs[j].patches() <= max_patches)) : (j += 1) patches += runs[j].patches(); + const grids = try gpa.alloc(Grid, j - i); + defer gpa.free(grids); + for (runs[i..j], grids) |r, *g| g.* = .{ .h = r.gh, .w = r.gw }; + if (j - i == 1) { + try tower.encode(images[i].patches, grids, m.features.ptr + done * c.out_hidden * 2); + } else { + const px = try gpa.alloc(f32, patches * patch_width); + defer gpa.free(px); + var at: usize = 0; + for (images[i..j]) |img| { + @memcpy(px[at..][0..img.patches.len], img.patches); + at += img.patches.len; + } + try tower.encode(px, grids, m.features.ptr + done * c.out_hidden * 2); + } + for (runs[i..j]) |r| done += r.tokens(); + i = j; + } + m.delta = layout(n, runs, m.src, m.pos); + return m; + } + + pub fn deinit(m: *Media) void { + m.features.free(); + m.gpa.free(m.src); + m.gpa.free(m.pos); + } +}; + +/// An image's placeholder run in the prompt: its first row and its patch grid. +pub const Run = struct { + at: usize, + gh: u32, + gw: u32, + + pub fn tokens(r: Run) usize { + return @as(usize, r.gh / merge) * (r.gw / merge); + } + + pub fn patches(r: Run) usize { + return @as(usize, r.gh) * r.gw; + } +}; + +/// Each image's run of image_token rows between vision_start and vision_end, in order; videos are refused by name. +pub fn placeholders(c: Config, prompt: []const u32, images: []const lanes.Image, runs: []Run) !void { + var k: usize = 0; + var r: usize = 0; + while (r < prompt.len) : (r += 1) { + const t = prompt[r]; + if (t == c.video_token) return error.QwenVideoUnsupported; + if (t != c.image_token) continue; + if (r == 0 or prompt[r - 1] != c.vision_start) return error.BadImagePlaceholders; + if (k == images.len) return error.BadImagePlaceholders; + const img = images[k]; + if (img.grid[0] != 1) return error.QwenVideoUnsupported; + if (img.grid[1] <= 0 or img.grid[2] <= 0 or img.grid[1] > 4096 or img.grid[2] > 4096) return error.BadImageGrid; + const run: Run = .{ .at = r, .gh = @intCast(img.grid[1]), .gw = @intCast(img.grid[2]) }; + if (run.gh % merge != 0 or run.gw % merge != 0 or run.patches() > max_patches) return error.BadImageGrid; + if (img.patches.len != run.patches() * patch_width) return error.BadImagePatches; + const end = r + run.tokens(); + if (end >= prompt.len or prompt[end] != c.vision_end) return error.BadImagePlaceholders; + for (prompt[r..end]) |x| if (x != c.image_token) return error.BadImagePlaceholders; + runs[k] = run; + k += 1; + r = end; + } + if (k != images.len) return error.BadImagePlaceholders; +} + +/// The prompt's rotary positions ([3][n]) and feature rows (src, -1 for text) per get_rope_index; returns rope_delta. +pub fn layout(n: usize, runs: []const Run, src: []i32, pos: []i32) i32 { + @memset(src, -1); + var next: i32 = 0; + var cursor: usize = 0; + var row: i32 = 0; + for (runs) |img| { + const begin = img.at; + const h: usize = img.gh / merge; + const w: usize = img.gw / merge; + for (cursor..begin) |r| { + const v = next + @as(i32, @intCast(r - cursor)); + for (0..3) |a| pos[a * n + r] = v; + } + const base = next + @as(i32, @intCast(begin - cursor)); + for (0..h * w) |k| { + const r = begin + k; + pos[r] = base; + pos[n + r] = base + @as(i32, @intCast(k / w)); + pos[2 * n + r] = base + @as(i32, @intCast(k % w)); + src[r] = row; + row += 1; + } + next = base + @as(i32, @intCast(@max(h, w))); + cursor = begin + h * w; + } + for (cursor..n) |r| { + const v = next + @as(i32, @intCast(r - cursor)); + for (0..3) |a| pos[a * n + r] = v; + } + return next - @as(i32, @intCast(cursor)); +} + +const test_config: Config = .{ .hidden = 1152, .heads = 16, .depth = 27, .intermediate = 4304, .out_hidden = 8, .image_token = 50, .video_token = 51, .vision_start = 52, .vision_end = 53 }; + +test "media positions: text counts up, an image sits at one position with its grid, text resumes past its side" { + const prompt = [_]u32{ 7, 52, 50, 50, 50, 50, 50, 50, 53, 9 }; + const px: [24 * patch_width]f32 = @splat(0); + var runs: [1]Run = undefined; + try placeholders(test_config, &prompt, &.{.{ .patches = &px, .grid = .{ 1, 4, 6 } }}, &runs); + try std.testing.expectEqual(@as(usize, 2), runs[0].at); + var src: [10]i32 = undefined; + var pos: [30]i32 = undefined; + const delta = layout(prompt.len, &runs, &src, &pos); + try std.testing.expectEqual(@as(i32, -3), delta); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 2, 2, 2, 2, 2, 5, 6 }, pos[0..10]); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 2, 2, 3, 3, 3, 5, 6 }, pos[10..20]); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3, 4, 2, 3, 4, 5, 6 }, pos[20..30]); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, 0, 1, 2, 3, 4, 5, -1, -1 }, &src); +} + +test "placeholders refuse a run of the wrong length, a missing image, an extra image and a video" { + const px: [24 * patch_width]f32 = @splat(0); + const img: lanes.Image = .{ .patches = &px, .grid = .{ 1, 4, 6 } }; + var runs: [2]Run = undefined; + try std.testing.expectError(error.BadImagePlaceholders, placeholders(test_config, &.{ 52, 50, 50, 53 }, &.{img}, runs[0..1])); + try std.testing.expectError(error.BadImagePlaceholders, placeholders(test_config, &.{ 52, 50, 50, 50, 50, 50, 50, 53 }, &.{ img, img }, &runs)); + try std.testing.expectError(error.BadImagePlaceholders, placeholders(test_config, &.{ 7, 8 }, &.{img}, runs[0..1])); + try std.testing.expectError(error.QwenVideoUnsupported, placeholders(test_config, &.{ 52, 51, 53 }, &.{}, runs[0..0])); + try std.testing.expectError(error.BadImagePlaceholders, placeholders(test_config, &.{ 7, 50 }, &.{img}, runs[0..1])); +} 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..2a2b75c361 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -0,0 +1,142 @@ +//! 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). +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); +} + +/// 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) + 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])); + return most; +} + +fn upload(e: *Engine, bytes: []const u8) !u64 { + const p = try e.alloc(bytes.len); + try e.arena.upload(p - e.arena.ptr, bytes); + return p; +} + +fn put(e: *Engine, t: core.checkpoint.Tensor) !u64 { + return upload(e, 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)); +} + +fn tensor(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + var buf: [160]u8 = undefined; + return put(e, try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); +} + +/// An MLX affine-4 projection packed into qlinear's tiles from the staging buffer, which holds it as stored. +fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { + var buf: [160]u8 = undefined; + var raw: [3]u64 = undefined; + var at: usize = 0; + const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (parts, &raw) |p, *r| { + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); + try e.stage.upload(at, t.bytes); + r.* = e.stage.ptr + at; + at += t.bytes.len; + } + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try e.alloc(l.words), .s = try e.alloc(l.scales), .b = try e.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; +} + +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, ck, m, "embed_tokens", g.vocab, g.hidden) else try linear(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}), + .gate = try linear(e, ck, p, "mlp.gate_proj", g.intermediate, g.hidden), + .up = try linear(e, ck, p, "mlp.up_proj", g.intermediate, g.hidden), + .down = try linear(e, ck, p, "mlp.down_proj", g.hidden, g.intermediate), + .mixer = undefined, + }; + if (g.linear[i]) { + l.mixer = .{ .delta = .{ + .qkv = try linear(e, ck, p, "linear_attn.in_proj_qkv", conv, g.hidden), + .z = try linear(e, ck, p, "linear_attn.in_proj_z", vw, g.hidden), + .a = try linear(e, ck, p, "linear_attn.in_proj_a", g.v_heads, g.hidden), + .b = try linear(e, ck, p, "linear_attn.in_proj_b", g.v_heads, g.hidden), + .out = try linear(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, ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden), + .k = try linear(e, ck, p, "self_attn.k_proj", g.kv_heads * head_dim, g.hidden), + .v = try linear(e, ck, p, "self_attn.v_proj", g.kv_heads * head_dim, g.hidden), + .o = try linear(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..ae52d9cd9d 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,24 +1,30 @@ //! CUDA family registration, device setup, memory admission and the native lane host. + +/// The serve flags this build's engine takes past the common ones (server/cli.zig served()): two ranks. +pub const serves: []const []const u8 = &.{"--tp"}; + 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}; - -/// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. -const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ - .bytes = nemotron.native.snap_bytes, - .save = nemotron.native.snap_save, - .restore = nemotron.native.snap_restore, - .drop = nemotron.native.snap_drop, -}; +const registry = .{ nemotron.native, qwen35.native }; + +/// A family's prompt-cache copies. A comptime value per family, so the store's pointer stays valid for the process. +fn Snaps(comptime F: type) type { + return struct { + const vtable: api.prompt_cache.Snapshots.VTable = .{ .bytes = F.snap_bytes, .save = F.snap_save, .restore = F.snap_restore, .drop = F.snap_drop }; + }; +} pub const backends: []const []const u8 = &.{"cuda"}; pub const families: []const api.Family = blk: { @@ -44,7 +50,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 +158,15 @@ 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 + follow_fn: ?*const fn (*anyopaque) anyerror!void = null, // rank 1 of two: the family's replay loop + + /// Rank 1: the family replays rank 0's calls on this thread (the context was made current here at open). + fn followFamily(p: *anyopaque) anyerror!void { + const h: *Host = @ptrCast(@alignCast(p)); + if (h.gpu) |g| try g.ctx.makeCurrent(); + return h.follow_fn.?(h.family); + } fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); @@ -165,6 +180,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 +204,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(); @@ -241,6 +261,21 @@ const Host = struct { const refusals = @import("cuda_refusal.zig"); +/// A tensor-parallel rank above 0: the family opened as that rank, then its replay loop; false with ``problem`` set. +pub fn follow(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !bool { + const opened = (try open(a, gpa, io, o, problem)) orelse return false; + defer opened.close(opened.ctx); + const f = opened.follow orelse { + problem.* = "this family has no tensor-parallel ranks"; + return false; + }; + f(opened.ctx) catch |e| { + problem.* = try std.fmt.allocPrint(a, "rank {d} stopped: {s}", .{ o.rank, @errorName(e) }); + return false; + }; + return true; +} + /// The engine for `o.dir`, or null with `problem` set when `refusal` names a flag or no CUDA family reads it. pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { if (refusals.backend(o)) |why| { @@ -322,49 +357,94 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }, }; const before = try pool(a, io, &g.ctx, problem) orelse return null; - const weights = weightBytes(io, o.dir); + const two = @hasDecl(F, "two_ranks") and F.two_ranks; + if (two and o.tp != 2) { + problem.* = try std.fmt.allocPrint(a, "{s} runs on two GPUs here: serve it with --tp 2 --rank 0 on one host and --tp 2 --rank 1 --master on the other", .{F.model_type}); + return null; + } + if (!two and o.tp != 1) { + problem.* = try std.fmt.allocPrint(a, "{s} runs on one GPU: drop --tp", .{F.model_type}); + return null; + } + const weights = if (@hasDecl(F, "deviceWeightBytes")) F.deviceWeightBytes(io, o.dir) else weightBytes(io, o.dir); const held0 = cuda.usage(false).device; if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); return null; } - const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments, .slide = o.slide }) catch |e| { + var fo: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments }; + if (@hasField(F.Options, "vision")) fo.vision = o.vision; + if (@hasField(F.Options, "kv_bits")) fo.kv_bits = o.kv_bits else if (o.kv_bits != 8) { + problem.* = try std.fmt.allocPrint(a, "{s} keeps an int8 attention cache only: drop --kv-dtype", .{F.model_type}); + return null; + } + if (@hasField(F.Options, "slide")) fo.slide = o.slide; + if (@hasField(F.Options, "rank")) { + fo.tp = @intCast(o.tp); + fo.rank = @intCast(o.rank); + fo.master = if (o.master.len > 0) o.master else null; + fo.master_port = o.master_port; + } + const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, fo) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} with {s} glue kernels ({s})", .{ o.dir, kernels orelse "its own", @errorName(e) }); return null; }; 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; + }; + if (@hasDecl(F, "setStreams")) F.setStreams(loaded.ctx, streams); + if (@hasDecl(F, "setRoom")) F.setRoom(loaded.ctx, room) catch |e| { // a family that sizes each stream's caches out of one pool + problem.* = try std.fmt.allocPrint(a, "the ranks could not agree on the stream pool ({s})", .{@errorName(e)}); + 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 if (@hasDecl(F, "snap_host")) budget.hostCacheBytes(a, io, o.prompt_cache_gib) 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 "", - 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", + 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{s}", .{ toGib(cache), if (@hasDecl(F, "snap_host")) " in host memory" else "" }) 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 - store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &nemotron_snaps }, .{ .lookahead = 1, .min_prompt = 0, .warm = .waits }, cache - try nemotron.native.snap_ready(loaded.target, cache)); + // Nemotron keeps states at any row, its head one token on; a host-memory family on its zero-anchored chunk grid + const rules: pc.Rules = if (@hasDecl(F, "snap_host")) .{ .lookahead = F.snap_lookahead, .planned = true, .grid = F.prompt_rows } else .{ .lookahead = 1, .min_prompt = 0, .warm = .waits }; + const ready: u64 = if (@hasDecl(F, "snap_ready")) try F.snap_ready(loaded.target, cache) else 0; + store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &Snaps(F).vtable }, rules, cache - ready); h.store = store; } errdefer if (h.store) |store| { 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,8 +464,13 @@ 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, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; + const follow_loop: ?*const fn (*anyopaque) anyerror!void = if (@hasField(@TypeOf(loaded), "follow")) loaded.follow else null; + if (follow_loop) |f| { + h.follow_fn = f; + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h, .follow = Host.followFamily }; + } return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 40977dfa50..fe929ae849 100644 --- a/zig/src/native/cuda_bound.zig +++ b/zig/src/native/cuda_bound.zig @@ -18,6 +18,12 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.prefill_step.?(x.inner.ptr, s); } }.f else null, + .prefill_many = if (v.prefill_many != null) struct { + fn f(p: *anyopaque, ss: []const *lanes.Stream) anyerror!bool { + const x = H.bind(p); + return x.inner.vtable.prefill_many.?(x.inner.ptr, ss); + } + }.f else null, .first = struct { fn f(p: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const x = H.bind(p); 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/cuda_memory.zig b/zig/src/native/cuda_memory.zig index 563b78bba1..cf34402888 100644 --- a/zig/src/native/cuda_memory.zig +++ b/zig/src/native/cuda_memory.zig @@ -30,6 +30,14 @@ pub fn meminfo(text: []const u8) ?MemInfo { return if (t > 0 and a <= t) .{ .total = t, .available = a } else null; } +/// --prompt-cache-gib in pinned host memory (8 GiB unasked), at most half of what /proc/meminfo has available. +pub fn hostCacheBytes(a: std.mem.Allocator, io: std.Io, asked: ?f64) u64 { + const want: u64 = if (asked) |g| (if (g > 0) std.math.lossyCast(u64, g * (1 << 30)) else 0) else 8 << 30; + const text = procText(a, io, "/proc/meminfo") orelse return 0; + const m = meminfo(text) orelse return 0; + return @min(want, m.available / 2); +} + /// A /proc file's text, streamed: procfs reports size 0, and a positional read (readFileAlloc) stops there. pub fn procText(a: std.mem.Allocator, io: std.Io, path: []const u8) ?[]u8 { var file = std.Io.Dir.cwd().openFile(io, path, .{}) catch return null; 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/no_family.zig b/zig/src/native/no_family.zig new file mode 100644 index 0000000000..d81d7567ac --- /dev/null +++ b/zig/src/native/no_family.zig @@ -0,0 +1,2 @@ +//! Stands in for a CUDA family the build leaves out (-Dflashnext=false): nothing is registered. +pub const native = struct {}; diff --git a/zig/src/native/no_family_glm.zig b/zig/src/native/no_family_glm.zig new file mode 100644 index 0000000000..2dcb07a292 --- /dev/null +++ b/zig/src/native/no_family_glm.zig @@ -0,0 +1,2 @@ +//! Stands in for the GLM CUDA family when the build leaves it out (-Dglm=false): nothing is registered. +pub const native = struct {}; 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/chat.zig b/zig/src/server/chat.zig index 7c702ef432..e1e868b662 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -13,6 +13,7 @@ const ids = @import("ids.zig"); const clock = @import("clock.zig"); const Server = @import("server.zig").Server; const chunk_plan = @import("chunk_plan.zig"); +const image_key = @import("image_key.zig"); const Mailbox = @import("mailbox.zig").Mailbox; const Value = json.Value; const Cx = errors.Cx; @@ -35,6 +36,8 @@ pub const Input = struct { logprobs: ?u8 = null, /// ``/v1/completions``' ``prompt_logprobs``: rows for the prompt with this many alternatives; null: not asked. prompt_logprobs: ?u8 = null, + /// Image inputs (vision_inputs.extract), in prompt order; the rendered prompt's image pads expand to their tokens. + images: []const api.Image = &.{}, }; /// A streamed piece: content text (a string) or a delta object (reasoning or tool calls). @@ -90,9 +93,30 @@ pub const Prepared = struct { preparing: bool, }; +/// --tool-system: the server's instruction first in a tools request with no system message; others pass unchanged. +fn withToolSystem(srv: *Server, a: Allocator, input: Input) Allocator.Error!Input { + const text = srv.config.tool_system orelse return input; + if (input.tools.len == 0 or input.prompt != null or input.messages != .array) return input; + for (input.messages.array) |m| { + if (m != .object) continue; + const role = m.object.get("role") orelse continue; + if (role == .string and std.mem.eql(u8, role.string, "system")) return input; + } + const sys = try json.newObject(a); + try sys.put(a, "role", .{ .string = "system" }); + try sys.put(a, "content", .{ .string = text }); + const list = try a.alloc(Value, input.messages.array.len + 1); + list[0] = .{ .object = sys }; + @memcpy(list[1..], input.messages.array); + var out = input; + out.messages = .{ .array = list }; + return out; +} + /// Render ``input`` and run every check that can refuse it, before anything reaches the client or the engine. -pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepared { +pub fn prepare(srv: *Server, cx: *Cx, input_in: Input, gone: anytype) Failure!Prepared { const a = cx.a; + const input = try withToolSystem(srv, a, input_in); const io = srv.io; const received = nowNs(io); const f = input.fields; @@ -108,7 +132,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking; if (input.prompt != null) thinking = false; const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); - const rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + var rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + if (input.images.len > 0) rendered.history_len = 0; // image prompts keep no prefix for a later turn (as Python) if (gone.check()) return error.Cancelled; if (rendered.ids.len == 0) return cx.refuse("rendered prompt is empty"); const window: i64 = srv.info.context_window; @@ -119,7 +144,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa if (input.max_tokens != null and limit > room) return cx.fail(.context_length, "{s} {d} tokens, but the rendered prompt has {d} tokens and requests {d} reply tokens, which exceeds the context window. Reduce the prompt to at most {d} prompt tokens or request at most {d} reply tokens, including chat template and thinking tokens.", .{ errors.context_limit, window, n, limit, @max(0, window - limit), room }); limit = @min(limit, room); } - const system_len: usize = if (input.prompt != null) 0 else prompt_mod.systemPrefixLen(srv, cx, input.messages, input.tools, rendered.ids, thinking, effort); + const system_len: usize = if (input.prompt != null or input.images.len > 0) 0 else prompt_mod.systemPrefixLen(srv, cx, input.messages, input.tools, rendered.ids, thinking, effort); var shared: std.ArrayList(u32) = .empty; if (system_len > 0) for ([_]i64{ @as(i64, @intCast(system_len)) - 2048, @as(i64, @intCast(system_len)) - 512, @intCast(system_len) }) |cut| { if (cut >= 512) try shared.append(a, @intCast(cut)); @@ -137,11 +162,14 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa .history_len = @intCast(rendered.history_len), .rewind_len = @intCast(rendered.rewind_len), .shared_prefixes = shared.items, - .chunks = try chunk_plan.withCut(a, try srv.chunks.starts(a, rendered.ids), if (srv.chunks.step > 0) @intCast(@max(system_len, 1) - 1) else 0, rendered.ids.len, srv.chunks.min_chunk), // a cut before the conversation's own text: fresh sessions resume their whole harness + .images = input.images, + // a cut just before the conversation's own text: fresh sessions resume their whole harness + .chunks = try chunk_plan.withCut(a, try srv.chunks.starts(a, rendered.ids), if (srv.chunks.step > 0) @intCast(@max(system_len, 1) - 1) else 0, rendered.ids.len, srv.chunks.min_chunk), .decode_spans = try prompt_mod.replySpans(srv, cx, rendered.ids), // from the tokens alone, for every request: cached states are keyed by tokens .tools_json = if (input.tools.len > 0) try json.stringify(a, .{ .array = @constCast(input.tools) }, .{ .ascii = false }) else "", .decide = if (input.prompt == null and !background) prompt_mod.decideOf(srv, rendered.ids) else null, }; + if (input.images.len > 0) request.cache_key = try imageCacheKey(srv, a, rendered.ids, input.images); try srv.checkFeatures(cx, f, input.tools.len > 0, thinking, rendered.ids, input.tools, &request); if (thinking) { const budget_field = f.get("thinking_budget"); @@ -291,7 +319,7 @@ const Generation = struct { return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; } - /// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. + /// What closes a call the end token left open (``tool_parse.closeCall``); nothing for stop, length or ignore_eos. fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { const t = g.collected.items; if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; @@ -371,7 +399,7 @@ const Generation = struct { }; } - /// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. + /// Ends the request and waits for the engine; one ended unfinished counts as a disconnect, as on the Mac. fn cancel(g: *Generation) void { if (!g.engine_done) g.srv.engine.cancel(g.id); g.drain(); @@ -551,3 +579,17 @@ pub fn deltaOf(a: Allocator, key: []const u8, text: []const u8) Allocator.Error! try o.put(a, key, .{ .string = text }); return .{ .object = o }; } + +/// An image prompt's cache key: its placeholder rows (Qwen's image and video pads, GLM's image token) by content. +fn imageCacheKey(srv: *Server, a: std.mem.Allocator, prompt: []const u32, images: []const api.Image) ![]const u32 { + var pads: [3]u32 = undefined; + var n: usize = 0; + for ([_][]const u8{ "<|image_pad|>", "<|video_pad|>", "<|image|>" }) |name| if (srv.text.tokenId(name)) |t| { + pads[n] = t; + n += 1; + }; + const key = try image_key.cacheKey(a, prompt, images, pads[0..n]); + // no placeholder token known: the key's first id is the images', so no other prompt shares any state + if (n == 0 and key.len > 0) key[0] = (try image_key.cacheKey(a, &.{0}, images, &.{0}))[0]; + return key; +} diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 76b205eeaa..af5aa48cd7 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -49,15 +49,16 @@ pub const flags = [_]Flag{ .{ .name = "--api-key-file", .native = true }, .{ .name = "--metrics-open", .kind = .store_true, .native = true }, .{ .name = "--dashboard", .kind = .store_true, .native = true }, - .{ .name = "--vision", .kind = .store_true }, - .{ .name = "--vision-urls", .kind = .store_true }, + .{ .name = "--vision", .kind = .store_true, .native = true }, + .{ .name = "--vision-urls", .kind = .store_true, .native = true }, .{ .name = "--vision-offload", .kind = .store_true }, - .{ .name = "--vision-max-images" }, - .{ .name = "--vision-image-tokens" }, + .{ .name = "--vision-max-images", .native = true }, + .{ .name = "--vision-image-tokens", .native = true }, .{ .name = "--context", .native = true }, .{ .name = "--speed-up", .native = !cuda_build }, .{ .name = "--max-tokens", .native = true }, .{ .name = "--temperature", .native = true }, + .{ .name = "--tool-system", .native = true }, .{ .name = "--top-p", .native = true }, .{ .name = "--top-k", .native = true }, .{ .name = "--min-p", .native = true }, @@ -99,7 +100,7 @@ pub const flags = [_]Flag{ .{ .name = "--master-port", .native = served("--tp") }, // the GPU engine's policy switches as key=value pairs, repeatable .{ .name = "--policy", .kind = .append, .native = served("--policy") }, - .{ .name = "--kv-dtype", .choices = &.{ "bf16", "int8", "int4" } }, + .{ .name = "--kv-dtype", .choices = &.{ "bf16", "int8", "int4" }, .native = served("--kv-dtype"), .native_values = &.{ "int8", "int4" } }, .{ .name = "--prefill-fp8", .kind = .store_true }, .{ .name = "--no-prefill-fp8", .kind = .store_true }, .{ .name = "--precision", .choices = &.{ "checkpoint", "full" } }, @@ -122,7 +123,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 = &.{}, @@ -140,6 +141,11 @@ pub const Args = struct { learn_gib: f64 = 32, // disk for learned states, every model and build together max_tokens: i64 = 4096, temperature: ?f64 = null, + vision: bool = false, // accept image inputs (an engine whose family has a vision tower) + vision_urls: bool = false, // also public HTTPS image and video URLs + vision_max_images: i64 = 4, + vision_image_tokens: i64 = 4096, // visual tokens a request's images share + tool_system: ?[]const u8 = null, // an instruction added as the system message of tool requests that send none top_p: ?f64 = null, top_k: ?i64 = null, min_p: ?f64 = null, @@ -170,6 +176,7 @@ pub const Args = struct { master_port: u16 = 29551, policy: []const u8 = "", // every --policy joined by commas cortheon: bool = false, // Cortheon's program, research before building, and cortheon_search/cortheon_fetch run by the server + kv_dtype: []const u8 = "int8", // --kv-dtype (CUDA): the attention cache's codes, int8 or int4 }; /// A usage error's message (argparse's ``error:`` line); the caller exits 2. @@ -265,7 +272,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa } else if (is(name, "--learn-gib")) { out.learn = true; out.learn_gib = try gib(u, a, name, v); - } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { + } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--tool-system")) out.tool_system = v else if (is(name, "--vision")) out.vision = true else if (is(name, "--vision-urls")) out.vision_urls = true else if (is(name, "--vision-max-images")) out.vision_max_images = try int(u, a, name, v) else if (is(name, "--vision-image-tokens")) out.vision_image_tokens = try int(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); @@ -282,9 +289,19 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa fn cudaFlag(a: Allocator, out: *Args, name: []const u8, v: []const u8, u: *Usage) error{ Usage, OutOfMemory }!bool { if (std.mem.eql(u8, name, "--device")) { out.device = std.math.cast(u32, try int(u, a, name, v)) orelse return fail(u, a, "argument --device: a GPU ordinal from 0: '{s}'", .{v}); + } else if (std.mem.eql(u8, name, "--kv-dtype")) { + out.kv_dtype = v; } else if (std.mem.eql(u8, name, "--segments")) { const n = try int(u, a, name, v); out.segments = if (n >= 1) @intCast(@min(n, std.math.maxInt(u32))) else return fail(u, a, "argument --segments: a count from 1: '{s}'", .{v}); + } else if (std.mem.eql(u8, name, "--tp")) { + out.tp = if (std.mem.eql(u8, v, "2")) 2 else 1; + } else if (std.mem.eql(u8, name, "--rank")) { + out.rank = if (std.mem.eql(u8, v, "1")) 1 else 0; + } else if (std.mem.eql(u8, name, "--master")) { + out.master = v; + } else if (std.mem.eql(u8, name, "--master-port")) { + out.master_port = std.math.cast(u16, try int(u, a, name, v)) orelse return fail(u, a, "argument --master-port: a TCP port: '{s}'", .{v}); } else return false; return true; } @@ -362,6 +379,8 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .mtp_drafts = args.mtp_drafts, .mtp_confidence = args.mtp_confidence, .keep = args.checkpoint_slots, + .vision = args.vision, + .kv_bits = if (std.mem.eql(u8, args.kv_dtype, "int4")) 4 else 8, .tp = args.tp, .rank = args.rank, .master = args.master, 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/image_key.zig b/zig/src/server/image_key.zig new file mode 100644 index 0000000000..0819d481d4 --- /dev/null +++ b/zig/src/server/image_key.zig @@ -0,0 +1,97 @@ +//! An image prompt's cache key: placeholder rows as ids from each image's digest; restores compare whole digests. +const std = @import("std"); +const api = @import("engine_api"); + +/// Ids the tokenizer never makes: the top bit set. +const marked: u32 = 0x8000_0000; + +/// `ids` with placeholder run k (ids in `pads`) as ids from image k's digest; unpaired runs key by all the images. +pub fn cacheKey(a: std.mem.Allocator, ids: []const u32, images: []const api.Image, pads: []const u32) ![]u32 { + const key = try a.dupe(u32, ids); + errdefer a.free(key); + var runs: usize = 0; + for (ids, 0..) |t, i| { + if (isPad(t, pads) and (i == 0 or !isPad(ids[i - 1], pads))) runs += 1; + } + const seeds = try a.alloc(u64, images.len); + defer a.free(seeds); + var all = std.hash.Wyhash.init(0x51ce_7a92); + for (images, seeds) |im, *seed| { + const d = im.digest(); + all.update(&d); + seed.* = std.mem.readInt(u64, d[0..8], .little); + } + const together = all.final(); + var run: usize = 0; + var row: u64 = 0; + var h: u64 = together; + for (key, 0..) |*t, i| { + if (!isPad(ids[i], pads)) continue; + if (i == 0 or !isPad(ids[i - 1], pads)) { + h = if (runs == images.len) seeds[run] else together; + run += 1; + row = 0; + } + t.* = marked | @as(u32, @truncate(std.hash.Wyhash.hash(h, std.mem.asBytes(&row)) & 0x7fff_ffff)); + row += 1; + } + return key; +} + +fn isPad(t: u32, pads: []const u32) bool { + return std.mem.indexOfScalar(u32, pads, t) != null; +} + +test "two images under the same placeholder ids key apart; the text around them keys as itself" { + const a = std.testing.allocator; + const px_a: [4]f32 = .{ 0.1, 0.2, 0.3, 0.4 }; + const px_b: [4]f32 = .{ 0.1, 0.2, 0.3, 0.5 }; + const ids = [_]u32{ 7, 52, 9, 9, 53, 8 }; + const ka = try cacheKey(a, &ids, &.{.{ .patches = &px_a, .grid = .{ 1, 2, 2 } }}, &.{9}); + defer a.free(ka); + const kb = try cacheKey(a, &ids, &.{.{ .patches = &px_b, .grid = .{ 1, 2, 2 } }}, &.{9}); + defer a.free(kb); + const ka2 = try cacheKey(a, &ids, &.{.{ .patches = &px_a, .grid = .{ 1, 2, 2 } }}, &.{9}); + defer a.free(ka2); + try std.testing.expectEqualSlices(u32, &.{ 7, 52 }, ka[0..2]); + try std.testing.expectEqualSlices(u32, &.{ 53, 8 }, ka[4..6]); + try std.testing.expect(ka[2] >= marked and ka[2] != ka[3]); + try std.testing.expect(ka[2] != kb[2] and ka[3] != kb[3]); + try std.testing.expectEqualSlices(u32, ka, ka2); + // the same pixels in another grid are another image + const kg = try cacheKey(a, &ids, &.{.{ .patches = &px_a, .grid = .{ 1, 1, 4 } }}, &.{9}); + defer a.free(kg); + try std.testing.expect(kg[2] != ka[2]); +} + +test "a second image's change keys only its own run apart; runs that do not pair with images key by all of them" { + const a = std.testing.allocator; + const p1: [2]f32 = .{ 1, 2 }; + const p2: [2]f32 = .{ 3, 4 }; + const p3: [2]f32 = .{ 3, 5 }; + const ids = [_]u32{ 52, 9, 9, 53, 7, 52, 9, 53 }; + const k12 = try cacheKey(a, &ids, &.{ .{ .patches = &p1, .grid = .{ 1, 2, 2 } }, .{ .patches = &p2, .grid = .{ 1, 2, 2 } } }, &.{9}); + defer a.free(k12); + const k13 = try cacheKey(a, &ids, &.{ .{ .patches = &p1, .grid = .{ 1, 2, 2 } }, .{ .patches = &p3, .grid = .{ 1, 2, 2 } } }, &.{9}); + defer a.free(k13); + try std.testing.expectEqualSlices(u32, k12[0..6], k13[0..6]); + try std.testing.expect(k12[6] != k13[6]); + const one = try cacheKey(a, &ids, &.{.{ .patches = &p1, .grid = .{ 1, 2, 2 } }}, &.{9}); + defer a.free(one); + try std.testing.expect(one[1] != k12[1]); // two runs, one image: neither run keys as image 1 alone +} + +test "two different images behind one placeholder key apart" { + const a = std.testing.allocator; + const ids = [_]u32{ 7, 52, 9, 53, 8 }; + var seen: [64]u32 = undefined; + // 64 one-row images that differ in one patch value each: 64 different keys at the placeholder + for (&seen, 0..) |*k, i| { + const px: [2]f32 = .{ 0.5, @floatFromInt(i) }; + const key = try cacheKey(a, &ids, &.{.{ .patches = &px, .grid = .{ 1, 1, 1 } }}, &.{9}); + defer a.free(key); + try std.testing.expectEqualSlices(u32, &.{ 7, 52 }, key[0..2]); + k.* = key[2]; + for (seen[0..i]) |o| try std.testing.expect(o != k.*); + } +} 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/main.zig b/zig/src/server/main.zig index 4b85fca8cd..18a6d3a5f7 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -78,6 +78,10 @@ pub fn main(init: std.process.Init) !u8 { } orelse return fail(problem); defer up.text.deinit(); defer up.engine.close(up.engine.ctx); + if (up.engine.follow) |follow| { // rank 1 of two: rank 0 serves, this rank replays its engine calls + follow(up.engine.ctx) catch |e| return fail(@errorName(e)); + return 0; + } return serve.run(gpa, io, args, .{ .engine = up.engine.engine, .text = up.text.text(), diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig new file mode 100644 index 0000000000..a236d97152 --- /dev/null +++ b/zig/src/server/media_fetch.zig @@ -0,0 +1,336 @@ +//! Public HTTPS media for --vision-urls as vision/images_http.py fetches it: checked hosts, TLS, 3 redirects, limits. + +const std = @import("std"); +const builtin = @import("builtin"); +const Io = std.Io; +const net = std.Io.net; + +pub const Error = error{ OutOfMemory, Media }; + +/// The failure's words for the client (the request is refused with them). +pub const Failure = struct { text: []const u8 = "" }; + +pub const max_url_chars = 4096; +pub const max_redirects = 3; +pub const timeout_s = 10; // one download +pub const total_s = 30; // a request's downloads together + +fn fail(f: *Failure, text: []const u8) Error { + f.text = text; + return error.Media; +} + +/// Python ipaddress is_global, less multicast, reserved, 192.0.0.0/24, 168.63.129.16 and IPv6 embedding an IPv4 one. +pub fn publicIp(a: net.IpAddress) bool { + switch (a) { + .ip4 => |v| { + const b = v.bytes; + const in = struct { + fn net4(x: [4]u8, base: [4]u8, bits: u5) bool { + const xi = std.mem.readInt(u32, &x, .big); + const bi = std.mem.readInt(u32, &base, .big); + const mask: u32 = if (bits == 0) 0 else ~@as(u32, 0) << @intCast(32 - @as(u6, bits)); + return xi & mask == bi & mask; + } + }.net4; + const blocked = [_]struct { [4]u8, u5 }{ + .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, + .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, + .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, + .{ .{ 224, 0, 0, 0 }, 4 }, .{ .{ 240, 0, 0, 0 }, 4 }, + }; + for (blocked) |blk| if (in(b, blk[0], blk[1])) return false; + if (std.mem.eql(u8, &b, &.{ 168, 63, 129, 16 })) return false; + return true; + }, + .ip6 => |v| { + const b = v.bytes; + // 2000::/3 only, less 2001:db8::/32, the 2001::/23 IETF block (Teredo among it) and 6to4 2002::/16 + if (b[0] & 0xe0 != 0x20) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] == 0x0d and b[3] == 0xb8) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] < 0x02) return false; + if (b[0] == 0x20 and b[1] == 0x02) return false; + return true; + }, + } +} + +const Url = struct { host: []const u8, target: []const u8 }; + +/// grammar of images_http._url: https, a host, port 443, no credentials, fragment, backslash, whitespace or controls. +fn parseUrl(a: std.mem.Allocator, value: []const u8, f: *Failure) Error!Url { + if (value.len > max_url_chars) return fail(f, "image URL is too long or contains whitespace/control characters"); + for (value) |c| if (c <= 32 or c == 127) return fail(f, "image URL is too long or contains whitespace/control characters"); + const bad = "image URL must be HTTPS on port 443, without credentials or a fragment"; + if (!std.ascii.startsWithIgnoreCase(value, "https://")) return fail(f, bad); + if (std.mem.indexOfScalar(u8, value, '\\') != null or std.mem.indexOfScalar(u8, value, '#') != null) return fail(f, bad); + const rest = value["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const authority = rest[0..end]; + if (std.mem.indexOfScalar(u8, authority, '@') != null) return fail(f, bad); + var host = authority; + if (std.mem.lastIndexOfScalar(u8, authority, ':')) |c| { + if (authority.len > 0 and authority[0] == '[') return fail(f, bad); // a literal IPv6 host: not a name + if (!std.mem.eql(u8, authority[c + 1 ..], "443") and authority[c + 1 ..].len > 0) return fail(f, bad); + host = authority[0..c]; + } + if (host.len == 0 or std.mem.indexOfScalar(u8, host, '%') != null) return fail(f, bad); + for (host) |c| if (c >= 128) return fail(f, bad); // IDNA names: their ASCII (xn--) form + const lower = try std.ascii.allocLowerString(a, std.mem.trimEnd(u8, host, ".")); + for ([_][]const u8{ "localhost", "metadata.google.internal", "instance-data" }) |n| if (std.mem.eql(u8, lower, n)) return fail(f, "image URLs must use public internet hosts"); + // the path and query, percent-quoted as urllib.parse.quote(safe=...) leaves them + var t: std.ArrayList(u8) = .empty; + var tail = rest[end..]; + if (tail.len == 0 or tail[0] == '?') try t.append(a, '/'); + const safe_path = "/%:@!$&'()*+,;=-._~"; + var in_query = false; + for (tail) |c| { + if (c == '?' and !in_query) { + in_query = true; + try t.append(a, '?'); + continue; + } + const safe = std.ascii.isAlphanumeric(c) or std.mem.indexOfScalar(u8, safe_path, c) != null or (in_query and c == '?'); + if (safe) try t.append(a, c) else try t.print(a, "%{X:0>2}", .{c}); + } + tail = t.items; + return .{ .host = lower, .target = tail }; +} + +/// The location a redirect names, resolved against ``base`` (urljoin for absolute and absolute-path forms). +fn join(a: std.mem.Allocator, base: []const u8, location: []const u8) ![]const u8 { + if (std.ascii.startsWithIgnoreCase(location, "https://") or std.ascii.startsWithIgnoreCase(location, "http://")) return a.dupe(u8, location); + if (std.mem.startsWith(u8, location, "//")) return std.fmt.allocPrint(a, "https:{s}", .{location}); + const rest = base["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const origin = base[0 .. "https://".len + end]; + if (std.mem.startsWith(u8, location, "/")) return std.fmt.allocPrint(a, "{s}{s}", .{ origin, location }); + const path_end = std.mem.indexOfScalar(u8, rest[end..], '?') orelse rest.len - end; + const path = rest[end..][0..path_end]; + const dir = if (std.mem.lastIndexOfScalar(u8, path, '/')) |s| path[0 .. s + 1] else "/"; + return std.fmt.allocPrint(a, "{s}{s}{s}", .{ origin, dir, location }); +} + +fn remaining(io: Io, deadline: Io.Timestamp, f: *Failure) Error!i64 { + const left = Io.Clock.awake.now(io).durationTo(deadline).toMilliseconds(); + if (left <= 0) return fail(f, "image download timed out"); + return left; +} + +fn setTimeouts(handle: net.Socket.Handle, ms: i64) void { + const tv: std.posix.timeval = .{ .sec = @intCast(@divTrunc(ms, 1000)), .usec = @intCast(@mod(ms, 1000) * 1000) }; + std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.RCVTIMEO, std.mem.asBytes(&tv)) catch {}; + std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.SNDTIMEO, std.mem.asBytes(&tv)) catch {}; +} + +/// A TCP connection to the checked address, connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds connect). +fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { + // the socket calls below are Linux's; a Windows server refuses image URLs by name + return if (comptime builtin.os.tag == .windows) fail(f, "image URLs are not supported on Windows yet") else connectLinux(addr, ms, f); +} + +fn connectLinux(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { + const posix = std.posix; + const linux = std.os.linux; + const family: u32 = switch (addr) { + .ip4 => posix.AF.INET, + .ip6 => posix.AF.INET6, + }; + const rc = linux.socket(family, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (linux.errno(rc) != .SUCCESS) return fail(f, "image download failed or timed out"); + const fd: posix.fd_t = @intCast(rc); + errdefer _ = linux.close(fd); + setTimeouts(fd, ms); + var sa4: posix.sockaddr.in = undefined; + var sa6: posix.sockaddr.in6 = undefined; + const r = switch (addr) { + .ip4 => |v| blk: { + sa4 = .{ .port = std.mem.nativeToBig(u16, 443), .addr = @bitCast(v.bytes) }; + break :blk linux.connect(fd, @ptrCast(&sa4), @sizeOf(posix.sockaddr.in)); + }, + .ip6 => |v| blk: { + sa6 = .{ .port = std.mem.nativeToBig(u16, 443), .flowinfo = 0, .addr = v.bytes, .scope_id = 0 }; + break :blk linux.connect(fd, @ptrCast(&sa6), @sizeOf(posix.sockaddr.in6)); + }, + }; + if (linux.errno(r) != .SUCCESS) return fail(f, "image download failed or timed out"); + return .{ .socket = .{ .handle = fd, .address = addr } }; +} + +/// The bytes of ``url`` (and its media type, one of ``media_types``), at most ``max_bytes``, before ``deadline``. +pub fn fetch(a: std.mem.Allocator, io: Io, url_in: []const u8, max_bytes: usize, deadline_in: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!struct { data: []u8, media: []const u8 } { + var url = url_in; + const deadline = blk: { + const one = Io.Clock.awake.now(io).addDuration(.fromSeconds(timeout_s)); + break :blk if (one.nanoseconds < deadline_in.nanoseconds) one else deadline_in; + }; + var hop: usize = 0; + while (true) : (hop += 1) { + const u = try parseUrl(a, url, f); + const got = try request(a, io, u, max_bytes, deadline, media_types, f); + switch (got) { + .body => |b| return .{ .data = b.data, .media = b.media }, + .redirect => |loc| { + if (hop == max_redirects) return fail(f, "image download has too many redirects"); + url = join(a, url, loc) catch return error.OutOfMemory; + }, + } + } +} + +const Got = union(enum) { body: struct { data: []u8, media: []const u8 }, redirect: []const u8 }; + +var bundle_lock: Io.RwLock = .init; +var bundle: std.crypto.Certificate.Bundle = .empty; +var bundle_loaded = false; + +fn request(a: std.mem.Allocator, io: Io, u: Url, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!Got { + // every address the name resolves to must be public; the connection goes to the first of them + const name = net.HostName.init(u.host) catch return fail(f, "image host could not be resolved"); + var results: [32]net.HostName.LookupResult = undefined; + var q: Io.Queue(net.HostName.LookupResult) = .init(&results); + name.lookup(io, &q, .{ .port = 443 }) catch return fail(f, "image host could not be resolved"); + var first: ?net.IpAddress = null; + var count: usize = 0; + while (q.getOneUncancelable(io)) |r| switch (r) { + .address => |addr| { + if (!publicIp(addr)) return fail(f, "image URLs must resolve only to public internet addresses"); + if (first == null) first = addr; + count += 1; + }, + .canonical_name => {}, + } else |_| {} + const addr = first orelse return fail(f, "image host could not be resolved"); + _ = try remaining(io, deadline, f); + // the CA bundle, read once + { + bundle_lock.lockUncancelable(io); + defer bundle_lock.unlock(io); + if (!bundle_loaded) { + bundle.rescan(std.heap.page_allocator, io, Io.Clock.real.now(io)) catch return fail(f, "image download failed: no CA certificates on the server"); + bundle_loaded = true; + } + } + const stream = try connect(addr, try remaining(io, deadline, f), f); + defer stream.close(io); + const tls_len = std.crypto.tls.Client.min_buffer_len; + const bufs = try a.alloc(u8, 4 * tls_len); + var sr = stream.reader(io, bufs[0..tls_len]); + var sw = stream.writer(io, bufs[tls_len .. 2 * tls_len]); + var entropy: [std.crypto.tls.Client.Options.entropy_len]u8 = undefined; + io.random(&entropy); + var tls = std.crypto.tls.Client.init(&sr.interface, &sw.interface, .{ + .host = .{ .explicit = u.host }, + .ca = .{ .bundle = .{ .gpa = std.heap.page_allocator, .io = io, .lock = &bundle_lock, .bundle = &bundle } }, + .read_buffer = bufs[2 * tls_len .. 3 * tls_len], + .write_buffer = bufs[3 * tls_len ..], + .entropy = &entropy, + .realtime_now = Io.Clock.real.now(io), + .allow_truncation_attacks = true, // HTTP's own lengths and framing end the body + }) catch return fail(f, "image download failed: TLS could not be verified for the host"); + // the request + const accept = std.mem.join(a, ", ", media_types) catch return error.OutOfMemory; + tls.writer.print("GET {s} HTTP/1.1\r\nHost: {s}\r\nAccept: {s}\r\nAccept-Encoding: identity\r\nUser-Agent: TensorFold-native\r\nConnection: close\r\n\r\n", .{ u.target, u.host, accept }) catch return fail(f, "image download failed or timed out"); + tls.writer.flush() catch return fail(f, "image download failed or timed out"); + sw.interface.flush() catch return fail(f, "image download failed or timed out"); + const r = &tls.reader; + // status and headers + const status_line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + if (status_line.len < 12 or !std.mem.startsWith(u8, status_line, "HTTP/1.")) return fail(f, "image download failed"); + const status = std.fmt.parseInt(u16, status_line[9..12], 10) catch return fail(f, "image download failed"); + var location: ?[]const u8 = null; + var media: []const u8 = ""; + var length: ?usize = null; + var chunked = false; + var encoded = false; + var header_bytes: usize = 0; + while (true) { + const line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + header_bytes += line.len; + if (header_bytes > 64 * 1024) return fail(f, "image download failed"); + const l = std.mem.trimEnd(u8, line, "\r\n"); + if (l.len == 0) break; + const colon = std.mem.indexOfScalar(u8, l, ':') orelse continue; + const key = std.mem.trim(u8, l[0..colon], " \t"); + const val = std.mem.trim(u8, l[colon + 1 ..], " \t"); + if (std.ascii.eqlIgnoreCase(key, "location")) location = try a.dupe(u8, val); + if (std.ascii.eqlIgnoreCase(key, "content-type")) { + const semi = std.mem.indexOfScalar(u8, val, ';') orelse val.len; + media = try std.ascii.allocLowerString(a, std.mem.trim(u8, val[0..semi], " \t")); + } + if (std.ascii.eqlIgnoreCase(key, "content-length")) { + for (val) |c| if (!std.ascii.isDigit(c)) return fail(f, "image response exceeds the encoded byte limit"); + length = std.fmt.parseInt(usize, val, 10) catch return fail(f, "image response exceeds the encoded byte limit"); + } + if (std.ascii.eqlIgnoreCase(key, "transfer-encoding")) chunked = std.mem.indexOf(u8, try std.ascii.allocLowerString(a, val), "chunked") != null; + if (std.ascii.eqlIgnoreCase(key, "content-encoding") and !std.ascii.eqlIgnoreCase(val, "identity")) encoded = true; + } + switch (status) { + 301, 302, 303, 307, 308 => return .{ .redirect = location orelse return fail(f, "image redirect has no destination") }, + 200 => {}, + else => return fail(f, try std.fmt.allocPrint(a, "image download returned HTTP {d}", .{status})), + } + if (encoded) return fail(f, "compressed HTTP image responses are unsupported"); + for (media_types) |m| { + if (std.mem.eql(u8, m, media)) break; + } else return fail(f, if (std.mem.startsWith(u8, media_types[0], "image/")) "image URL content type must be JPEG, PNG or WebP" else try std.fmt.allocPrint(a, "media URL content type must be one of {s}", .{accept})); + if (length) |n| if (n > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + var data: std.ArrayList(u8) = .empty; + if (chunked) { + while (true) { + _ = try remaining(io, deadline, f); + const size_line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + const semi = std.mem.indexOfScalar(u8, size_line, ';') orelse size_line.len; + const hex = std.mem.trim(u8, size_line[0..semi], " \t\r\n"); + const n = std.fmt.parseInt(usize, hex, 16) catch return fail(f, "image download failed"); + if (n == 0) break; + if (data.items.len + n > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + const piece = r.take(n) catch return fail(f, "image download failed or timed out"); + try data.appendSlice(a, piece); + _ = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed"); + } + } else { + while (length == null or data.items.len < length.?) { + _ = try remaining(io, deadline, f); + const want = if (length) |n| @min(n - data.items.len, 64 * 1024) else 64 * 1024; + const piece = r.peekGreedy(1) catch |e| switch (e) { + error.EndOfStream => break, + else => return fail(f, "image download failed or timed out"), + }; + const take = @min(piece.len, want); + if (data.items.len + take > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + try data.appendSlice(a, piece[0..take]); + r.toss(take); + } + if (length) |n| if (data.items.len != n) return fail(f, "image download failed or timed out"); + } + _ = try remaining(io, deadline, f); + return .{ .body = .{ .data = data.items, .media = media } }; +} + +test "public addresses as Python's ipaddress checks them" { + const ip4 = struct { + fn of(b: [4]u8) net.IpAddress { + return .{ .ip4 = .{ .bytes = b, .port = 443 } }; + } + }.of; + try std.testing.expect(publicIp(ip4(.{ 8, 8, 8, 8 }))); + try std.testing.expect(publicIp(ip4(.{ 151, 101, 1, 69 }))); + for ([_][4]u8{ .{ 10, 1, 2, 3 }, .{ 127, 0, 0, 1 }, .{ 169, 254, 169, 254 }, .{ 172, 20, 0, 1 }, .{ 192, 168, 1, 1 }, .{ 100, 100, 0, 1 }, .{ 0, 0, 0, 0 }, .{ 224, 0, 0, 1 }, .{ 255, 255, 255, 255 }, .{ 168, 63, 129, 16 }, .{ 192, 0, 0, 9 } }) |b| + try std.testing.expect(!publicIp(ip4(b))); +} + +test "URLs as images_http._url takes them" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f: Failure = .{}; + try std.testing.expectError(error.Media, parseUrl(a, "https://example.com/a b/c.png", &f)); + const ok = try parseUrl(a, "https://example.com:443/img/c.png?x=%20&y=ä", &f); + try std.testing.expectEqualStrings("example.com", ok.host); + try std.testing.expectEqualStrings("/img/c.png?x=%20&y=%C3%A4", ok.target); + for ([_][]const u8{ "http://example.com/x.png", "https://user@example.com/x", "https://example.com:8443/x", "https://localhost/x", "https://example.com/x#frag", "https://[::1]/x" }) |bad| + try std.testing.expectError(error.Media, parseUrl(a, bad, &f)); + try std.testing.expectEqualStrings("https://e.com/b/c.png", try join(a, "https://e.com/b/a.png", "c.png")); + try std.testing.expectEqualStrings("https://e.com/z", try join(a, "https://e.com/b/a.png", "/z")); +} diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index b8782aaf96..aaf23610ae 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -127,7 +127,12 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { if (srv.config.request_log) |path| request_log.append(cx.a, path, body); var input: chat.Input = .{ .fields = undefined }; if (is_chat) { - input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); + var msgs = body.get("messages"); + if (try @import("vision_inputs.zig").extract(srv, cx, msgs)) |x| { // images: decoded, prepared, their parts as text + msgs = x.messages; + input.images = x.images; + } + input.messages = try messages.normalize(cx, msgs, "system", srv.needs_user_after_tool); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index e68363bfe7..f48b0a76dd 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -39,6 +39,7 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t error.OutOfMemory => return error.OutOfMemory, error.Template => return cx.fail(.server, "{s}", .{problem}), }; + if (!thinking) ids = try withoutEffortLine(srv, cx, ids); if (!thinking and generation and ids.len > 0) { const last = try srv.text.decode(cx.a, ids[ids.len - 1 ..]); if (std.mem.eql(u8, @import("reply_text.zig").pyStrip(last), "")) { @@ -52,6 +53,17 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t return ids; } +/// GLM-5.3's effort system line, dropped with thinking off as glm5_next/prompts.py does; its ids are its own segment. +fn withoutEffortLine(srv: *Server, cx: *Cx, ids: []u32) errors.Refused![]u32 { + const line = srv.text.encode(cx.a, "<|system|>Reasoning Effort: Max", false) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return ids, + }; + if (line.len < 2) return ids; + const at = std.mem.indexOf(u32, ids, line) orelse return ids; + return try std.mem.concat(cx.a, u32, &.{ ids[0..at], ids[at + line.len ..] }); +} + /// The template's generation prompt (thinking on or off): what a probe conversation renders past its history. fn generationPrompt(srv: *Server, cx: *Cx, thinking: bool) errors.Refused![]const u32 { const user = try json.newObject(cx.a); 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..6673db74cd 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -34,6 +34,7 @@ test { _ = @import("listener.zig"); _ = @import("hub.zig"); _ = @import("chunk_plan.zig"); + _ = @import("image_key.zig"); _ = @import("status_routes.zig"); _ = @import("memory_routes.zig"); _ = @import("process_memory.zig"); @@ -44,6 +45,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..6a80e16ce6 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 { @@ -109,6 +110,10 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .model_dir = s.model_dir, .model_ids = ids.items, .default_max_tokens = args.max_tokens, + .vision_max_images = @intCast(@max(args.vision_max_images, 1)), + .vision_urls = args.vision_urls, + .vision_image_tokens = @intCast(std.math.clamp(args.vision_image_tokens, 1, 65536)), + .tool_system = if (args.tool_system) |t| (if (std.mem.trim(u8, t, " \t\r\n").len > 0) std.mem.trim(u8, t, " \t\r\n") else null) else null, .enable_thinking = args.thinking, .reasoning_effort = args.reasoning_effort, .thinking_budget = args.thinking_budget, @@ -211,6 +216,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..2c263e6939 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"); @@ -36,6 +37,10 @@ pub const Config = struct { /// The served name first, then aliases, without repeats. model_ids: []const []const u8, default_max_tokens: i64 = 4096, + tool_system: ?[]const u8 = null, // --tool-system (trimmed; null when empty) + vision_max_images: usize = 4, // --vision-max-images + vision_urls: bool = false, // --vision-urls: public HTTPS media URLs are fetched (media_fetch.zig) + vision_image_tokens: u64 = 4096, // --vision-image-tokens: the visual tokens a request's images share enable_thinking: bool = true, reasoning_effort: ?[]const u8 = null, thinking_budget: i64 = 0, @@ -135,6 +140,7 @@ pub const Server = struct { if (srv.keepalive) |k| k.stop(); // before the engine's queue goes away if (srv.grammars) |g| g.close(); srv.metrics.deinit(); + if (srv.grammars) |g| g.close(); srv.arena.deinit(); srv.store.deinit(); srv.slide.deinit(); @@ -280,7 +286,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 +296,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/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig new file mode 100644 index 0000000000..07e87a1cae --- /dev/null +++ b/zig/src/server/vision_inputs.zig @@ -0,0 +1,336 @@ +//! A chat request's image_url and video_url parts, prepared and replaced by the text Python's frontend expands them to. + +const std = @import("std"); +const builtin = @import("builtin"); +const sys = @import("sys"); +const media_fetch = @import("media_fetch.zig"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const Server = @import("server.zig").Server; +const Value = json.Value; +const Cx = errors.Cx; +const qi = api.qwen_image; + +const vision_start = "<|vision_start|>"; +const vision_end = "<|vision_end|>"; +const max_video_bytes: usize = 16 * 1024 * 1024; +const max_videos = 2; +const max_bytes_one: usize = 10 * 1024 * 1024; +const max_bytes_all: usize = 20 * 1024 * 1024; +const max_pixels_all: u64 = 32 * 1024 * 1024; +const max_tokens_one: u64 = 4096; + +pub const Extracted = struct { messages: Value, images: []api.Image }; + +/// The image URL of a content part (``image_url``, ``image`` or ``input_image``) and its ``detail``; else null. +const Kind = enum { image, video }; + +fn mediaPart(part: Value) ?struct { kind: Kind, url: ?[]const u8, detail: []const u8 } { + if (part != .object) return null; + const ty = part.get("type"); + const is_video = (ty != null and ty.? == .string and (std.mem.eql(u8, ty.?.string, "video_url") or std.mem.eql(u8, ty.?.string, "video"))) or + part.get("video_url") != null; + if (is_video) { + var vurl: ?[]const u8 = null; + if (part.get("video_url")) |vu| switch (vu) { + .string => |x| vurl = x, + .object => if (vu.get("url")) |u| if (u == .string) { + vurl = u.string; + }, + else => {}, + }; + if (vurl == null) if (part.get("video")) |v| if (v == .string) { + vurl = v.string; + }; + return .{ .kind = .video, .url = vurl, .detail = "auto" }; + } + const is_image = (ty != null and ty.? == .string and (std.mem.eql(u8, ty.?.string, "image_url") or std.mem.eql(u8, ty.?.string, "image") or + std.mem.eql(u8, ty.?.string, "input_image"))) or part.get("image_url") != null or part.get("image") != null; + if (!is_image) return null; + var url: ?[]const u8 = null; + var detail: []const u8 = "auto"; + if (part.get("image_url")) |iu| switch (iu) { + .string => |s| url = s, + .object => { + if (iu.get("url")) |u| if (u == .string) { + url = u.string; + }; + if (iu.get("detail")) |d| if (d == .string) { + detail = d.string; + }; + }, + else => {}, + }; + if (url == null) if (part.get("image")) |im| if (im == .string) { + url = im.string; + }; + if (part.get("detail")) |d| if (d == .string) { + detail = d.string; + }; + return .{ .kind = .image, .url = url, .detail = detail }; +} + +fn hasImages(messages: Value) bool { + if (messages != .array) return false; + for (messages.array) |m| { + if (m != .object) continue; + const c = m.get("content") orelse continue; + if (c != .array) continue; + for (c.array) |p| if (mediaPart(p) != null) return true; + } + return false; +} + +const image_media = [_][]const u8{ "image/jpeg", "image/png", "image/webp" }; +const video_media = [_][]const u8{ "video/mp4", "video/webm", "video/quicktime", "video/x-matroska" }; + +/// The bytes of a ``data:`` URL, or (--vision-urls) a public HTTPS URL fetched before ``deadline`` (images.py). +fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, deadline: std.Io.Timestamp) errors.Refused![]u8 { + if (!std.mem.startsWith(u8, url, "data:")) { + if (!std.mem.startsWith(u8, url, "https://")) return cx.refuse(if (kind == .image) "images require data URLs or public HTTPS URLs" else "videos require data URLs or public HTTPS URLs"); + if (!srv.config.vision_urls) return cx.refuse("image URLs are off on this server; send the image as a data URL, or start the server with --vision-urls"); + if (url.len > media_fetch.max_url_chars) return cx.refuse("image URL is too long"); + var f: media_fetch.Failure = .{}; + const got = media_fetch.fetch(cx.a, srv.io, url, limit, deadline, if (kind == .image) &image_media else &video_media, &f) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Media => return cx.fail(.request, "{s}", .{f.text}), + }; + return got.data; + } + const comma = std.mem.indexOfScalar(u8, url, ',') orelse return cx.refuse("an image data: URL needs a comma before its bytes"); + if (!std.mem.endsWith(u8, url[0..comma], ";base64")) return cx.refuse("an image data: URL must be base64-encoded"); + const b64 = url[comma + 1 ..]; + const dec = std.base64.standard.Decoder; + const n = dec.calcSizeForSlice(b64) catch return cx.refuse("an image's base64 bytes are invalid"); + if (n > limit) return cx.refuse(if (kind == .image) "an image exceeds the 10 MiB limit (20 MiB for a request's media together)" else "a video exceeds the 16 MiB limit (20 MiB for a request's media together)"); + const out = try cx.a.alloc(u8, n); + dec.decode(out, b64) catch return cx.refuse("an image's base64 bytes are invalid"); + return out; +} + +/// A request's images prepared, its image parts as the template's image text; null without images. Needs --vision. +pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracted { + const list = messages orelse return null; + if (!hasImages(list)) return null; + if (!srv.info.vision) return cx.refuse("this server accepts text parts only; image, audio and video inputs are unsupported"); + // every media part, in prompt order, with where its text goes + const Slot = struct { kind: Kind, url: ?[]const u8, detail: []const u8, msg: usize, part: usize }; + var slots: std.ArrayList(Slot) = .empty; + for (list.array, 0..) |m, i| { + if (m != .object) continue; + const c = m.get("content") orelse continue; + if (c != .array) continue; + const role = m.get("role"); + for (c.array, 0..) |p, j| { + const im = mediaPart(p) orelse continue; + if (role != null and role.? == .string and (std.mem.eql(u8, role.?.string, "system") or std.mem.eql(u8, role.?.string, "developer"))) + return cx.refuse("a system message cannot contain images or videos"); + try slots.append(cx.a, .{ .kind = im.kind, .url = im.url, .detail = im.detail, .msg = i, .part = j }); + } + } + var n_images: usize = 0; + var n_videos: usize = 0; + for (slots.items) |sl| switch (sl.kind) { + .image => n_images += 1, + .video => n_videos += 1, + }; + if (n_images > srv.config.vision_max_images) return cx.fail(.request, "a request may carry at most {d} images; this one has {d}", .{ srv.config.vision_max_images, n_images }); + if (n_videos > max_videos) return cx.fail(.request, "a request may carry at most {d} videos; this one has {d}", .{ max_videos, n_videos }); + const budget: u64 = srv.config.vision_image_tokens; + if (n_images > 0 and budget < n_images) return cx.refuse("the image count exceeds the visual-token budget"); + const per: u64 = if (n_images > 0) @min(budget / n_images, max_tokens_one) else 0; + var images: std.ArrayList(api.Image) = .empty; + const texts = try cx.a.alloc([]const u8, slots.items.len); + var total_bytes: usize = 0; + var total_pixels: u64 = 0; + const deadline = std.Io.Clock.awake.now(srv.io).addDuration(.fromSeconds(media_fetch.total_s)); // every URL's download + for (slots.items, 0..) |sl, k| { + const url = sl.url orelse return cx.refuse(if (sl.kind == .image) "an image part needs an image_url with a url" else "a video part needs a video_url with a url"); + const bytes = try dataBytes(srv, cx, url, @min(if (sl.kind == .image) max_bytes_one else max_video_bytes, max_bytes_all - total_bytes), sl.kind, deadline); + total_bytes += bytes.len; + if (total_bytes > max_bytes_all) return cx.refuse("a request's images and videos exceed the 20 MiB limit"); + var text: std.ArrayList(u8) = .empty; + switch (sl.kind) { + .image => { + var img = qi.decode(cx.a, bytes, .{}) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.ImageTooLarge => return cx.refuse("image dimensions exceed the decoded pixel limit"), + error.NoDecoder => return cx.refuse("image inputs need the shared JPEG and PNG decoder, which this server does not have yet"), + else => return cx.refuse("image bytes are invalid or unsupported; use JPEG or PNG"), + }; + total_pixels += @as(u64, img.w) * img.h; + if (total_pixels > max_pixels_all) return cx.refuse("a request's images exceed the decoded pixel limit"); + const cap = if (std.mem.eql(u8, sl.detail, "low")) @min(per, 256) else per; + if (srv.info.vision_glm) { // GLM: its canvas within [min(16, cap), cap] tokens, <|image|> placeholders + const p = qi.glm.prepare(cx.a, img, @min(16, cap), cap) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the image does not fit the visual-token budget"), + }; + img.deinit(cx.a); + try images.append(cx.a, .{ .patches = p.patches, .grid = p.grid }); + try text.appendSlice(cx.a, "<|begin_of_image|>"); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image|>"); + try text.appendSlice(cx.a, "<|end_of_image|>"); + texts[k] = text.items; + continue; + } + const p = qi.prepare(cx.a, img, cap) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("absolute aspect ratio must be smaller than 200"), + }; + img.deinit(cx.a); + try images.append(cx.a, .{ .patches = p.patches, .grid = p.grid }); + try text.appendSlice(cx.a, vision_start); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image_pad|>"); + try text.appendSlice(cx.a, vision_end); + }, + .video => if (srv.info.vision_glm) { + // GLM: <|begin_of_video|>, image groups with their "t.t seconds", <|end_of_video|> (glm5_next) + var v = try videoGlm(srv, cx, bytes); + defer v.deinit(cx.a); + try text.appendSlice(cx.a, "<|begin_of_video|>"); + for (0..v.groups) |g| { + try images.append(cx.a, .{ .patches = try cx.a.dupe(f32, v.groupPatches(g)), .grid = .{ 1, @intCast(v.gh), @intCast(v.gw) } }); + try text.appendSlice(cx.a, "<|begin_of_image|>"); + for (0..v.tokens()) |_| try text.appendSlice(cx.a, "<|image|>"); + try text.appendSlice(cx.a, "<|end_of_image|>"); + var tb: [32]u8 = undefined; + try text.appendSlice(cx.a, qi.formatSeconds(&tb, v.times[g]) catch unreachable); + try text.appendSlice(cx.a, " seconds"); + } + try text.appendSlice(cx.a, "<|end_of_video|>"); + } else { + var v = try video(srv, cx, bytes); + defer v.deinit(cx.a); + for (0..v.groups) |g| { + try images.append(cx.a, .{ .patches = try cx.a.dupe(f32, v.groupPatches(g)), .grid = .{ 1, @intCast(v.gh), @intCast(v.gw) } }); + var tb: [32]u8 = undefined; + try text.appendSlice(cx.a, "<"); + try text.appendSlice(cx.a, qi.formatSeconds(&tb, v.times[g]) catch unreachable); + try text.appendSlice(cx.a, " seconds>"); + try text.appendSlice(cx.a, vision_start); + for (0..v.tokens()) |_| try text.appendSlice(cx.a, "<|video_pad|>"); + try text.appendSlice(cx.a, vision_end); + } + }, + } + texts[k] = text.items; + } + // the messages with each media part as its text + const out = try cx.a.alloc(Value, list.array.len); + @memcpy(out, list.array); + var k: usize = 0; + while (k < slots.items.len) { + const i = slots.items[k].msg; + const c = list.array[i].get("content").?; + const parts = try cx.a.dupe(Value, c.array); + while (k < slots.items.len and slots.items[k].msg == i) : (k += 1) { + const t = try json.newObject(cx.a); + try t.put(cx.a, "type", .{ .string = "text" }); + try t.put(cx.a, "text", .{ .string = texts[k] }); + parts[slots.items[k].part] = .{ .object = t }; + } + const copy = try json.copyObject(cx.a, list.array[i].object); + try copy.put(cx.a, "content", .{ .array = parts }); + out[i] = .{ .object = copy }; + } + return .{ .messages = .{ .array = out }, .images = images.items }; +} + +// -- videos: libtfvideo (FFmpeg), loaded when the first video arrives ------------------------------------------ + +const TfVideo = extern struct { rgb: ?[*]u8, frames: c_int, height: c_int, width: c_int, indices: ?[*]c_int, fps: f64 }; +const SizeFn = *const fn (ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int; +const DecodeFn = *const fn (data: [*]const u8, len: usize, rate: f64, min_frames: c_int, max_frames: c_int, max_seconds: f64, max_dim: c_int, size: SizeFn, ctx: ?*anyopaque, out: *TfVideo) callconv(.c) c_int; +const FreeFn = *const fn (v: *TfVideo) callconv(.c) void; + +var lib_mutex: std.Io.Mutex = .init; +var lib: ?struct { decode: DecodeFn, free: FreeFn } = null; + +fn library(io: std.Io) ?@TypeOf(lib.?) { + lib_mutex.lockUncancelable(io); + defer lib_mutex.unlock(io); + if (lib == null) { + var dl = sys.DynLib.open(if (builtin.os.tag == .windows) "tfvideo.dll" else "libtfvideo.so") catch return null; + const d = dl.lookup(DecodeFn, "tf_video_decode") orelse return null; + const f = dl.lookup(FreeFn, "tf_video_free") orelse return null; + lib = .{ .decode = d, .free = f }; // kept open for the server's life + } + return lib; +} + +fn sizeFor(ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int { + const budget: *const u64 = @ptrCast(@alignCast(ctx.?)); + const s = qi.videoSize(@intCast(frames), @intCast(height), @intCast(width), budget.*) catch return -1; + oh.* = @intCast(s[0]); + ow.* = @intCast(s[1]); + return 0; +} + +/// GLM's video canvas: the frames' content size for libtfvideo (geometry within the budget), the canvas kept. +const GlmSize = struct { budget: u64, th: u32 = 0, tw: u32 = 0 }; + +fn sizeForGlm(ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int { + const gs: *GlmSize = @ptrCast(@alignCast(ctx.?)); + const g = qi.glm.geometry(@intCast(frames), @intCast(height), @intCast(width), 16, gs.budget) catch return -1; + gs.th = g[0]; + gs.tw = g[1]; + oh.* = @intCast(g[2]); + ow.* = @intCast(g[3]); + return 0; +} + +fn videoGlm(srv: *Server, cx: *Cx, bytes: []const u8) errors.Refused!qi.glm.Video { + const l = library(srv.io) orelse return cx.refuse("video inputs need libtfvideo (FFmpeg) beside the server"); + var gs: GlmSize = .{ .budget = 16384 }; + if (std.c.getenv("TENSORFOLD_VIDEO_TOKENS")) |v| gs.budget = std.fmt.parseInt(u64, std.mem.span(v), 10) catch gs.budget; + var v: TfVideo = undefined; + const rc = l.decode(bytes.ptr, bytes.len, 2.0, 2, 256, 3600.0, 8192, sizeForGlm, @ptrCast(&gs), &v); + if (rc != 0) return cx.refuse(switch (rc) { + -2 => "the video has no video stream", + -3 => "video dimensions are missing or exceed the pixel limit", + -4 => "the video's frame count or rate is missing", + -5 => "videos are limited to 60 minutes", + -6 => "the video has fewer than two decodable frames", + -8 => "a video needs at least 2 frames and an aspect ratio under 200", + else => "video bytes are invalid or unsupported; use MP4 or WebM", + }); + defer l.free(&v); + const frames: usize = @intCast(v.frames); + const h: usize = @intCast(v.height); + const w: usize = @intCast(v.width); + const idx = try cx.a.alloc(i32, frames); + for (idx, 0..) |*x, i| x.* = v.indices.?[i]; + return qi.glm.prepareVideo(cx.a, v.rgb.?[0 .. frames * h * w * 3], frames, h, w, gs.th, gs.tw, idx, v.fps) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the video's frames do not fit the tower's patch grid"), + }; +} + +fn video(srv: *Server, cx: *Cx, bytes: []const u8) errors.Refused!qi.Video { + const l = library(srv.io) orelse return cx.refuse("video inputs need libtfvideo (FFmpeg) beside the server"); + var budget: u64 = 16384; + if (std.c.getenv("TENSORFOLD_VIDEO_TOKENS")) |v| budget = std.fmt.parseInt(u64, std.mem.span(v), 10) catch budget; + var v: TfVideo = undefined; + const rc = l.decode(bytes.ptr, bytes.len, 2.0, 4, 256, 3600.0, 8192, sizeFor, @ptrCast(&budget), &v); + if (rc != 0) return cx.refuse(switch (rc) { + -2 => "the video has no video stream", + -3 => "video dimensions are missing or exceed the pixel limit", + -4 => "the video's frame count or rate is missing", + -5 => "videos are limited to 60 minutes", + -6 => "the video has fewer than two decodable frames", + -8 => "a video needs at least 2 frames and an aspect ratio under 200", + else => "video bytes are invalid or unsupported; use MP4 or WebM", + }); + defer l.free(&v); + const frames: usize = @intCast(v.frames); + const h: usize = @intCast(v.height); + const w: usize = @intCast(v.width); + const idx = try cx.a.alloc(i32, frames); + for (idx, 0..) |*x, i| x.* = v.indices.?[i]; + return qi.prepareVideo(cx.a, v.rgb.?[0 .. frames * h * w * 3], frames, h, w, idx, v.fps) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the video's frames do not fit the tower's patch grid"), + }; +} diff --git a/zig/src/vision/glm_image.zig b/zig/src/vision/glm_image.zig new file mode 100644 index 0000000000..889d3e679b --- /dev/null +++ b/zig/src/vision/glm_image.zig @@ -0,0 +1,190 @@ +//! GLM-5.3-Flash's image and video preparation as mlx-vlm's glm5_next/processing.py: canvas, resize, norm, patches. + +const std = @import("std"); +const qi = @import("qwen_image.zig"); + +pub const PATCH = 14; +pub const MERGE = 2; +pub const TEMPORAL = 2; +pub const FACTOR = PATCH * MERGE; +const MEAN = [3]f32{ 0.48145466, 0.4578275, 0.40821073 }; +const STD = [3]f32{ 0.26862954, 0.26130258, 0.27577711 }; +const PER = 3 * TEMPORAL * PATCH * PATCH; + +pub const Error = qi.Error; + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const f = @floor(x); + const d = x - f; + if (d > 0.5) return f + 1; + if (d < 0.5) return f; + return if (@mod(f, 2) == 0) f else f + 1; +} + +fn alignUp(v: f64, f: f64) f64 { + return @ceil(v / f) * f; +} + +/// smart_resize: the canvas (height, width) of ``frames`` frames within [min_tokens, max_tokens] merged tokens. +pub fn smartResize(frames: u32, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![2]u32 { + if (frames == 0 or height == 0 or width == 0 or min_tokens == 0 or min_tokens > max_tokens) return error.UnsupportedImage; + const f: f64 = FACTOR; + const ppt: f64 = TEMPORAL * f * f; + const min_px = @as(f64, @floatFromInt(min_tokens)) * ppt; + const max_px = @as(f64, @floatFromInt(max_tokens)) * ppt; + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const n: f64 = @floatFromInt(frames); + const af = @max(@as(f64, TEMPORAL), pyRound(n / TEMPORAL) * TEMPORAL); + var ah = alignUp(h, f); + var aw = alignUp(w, f); + var px = af * ah * aw; + if (px < min_px) { + const scale = @sqrt(min_px / (n * h * w)); + ah = alignUp(@max(1, @ceil(h * scale)), f); + aw = alignUp(@max(1, @ceil(w * scale)), f); + px = af * ah * aw; + } + if (px > max_px) { // the tallest content within the budget, by bisection on its height + if (max_px < af * f * f) return error.UnsupportedImage; + var low: f64 = 1; + var high: f64 = h; + var best: [2]f64 = .{ f, f }; + while (low <= high) { + const ch = @floor((low + high) / 2); + const cw = @max(1, @floor(w * ch / h)); + const cand = [2]f64{ alignUp(ch, f), alignUp(cw, f) }; + if (af * cand[0] * cand[1] <= max_px) { + best = cand; + low = ch + 1; + } else high = ch - 1; + } + ah = best[0]; + aw = best[1]; + } + return .{ @intFromFloat(ah), @intFromFloat(aw) }; +} + +/// _resize_geometry: the canvas (th, tw) and the content (ch, cw) resized into its top-left corner. +pub fn geometry(frames: u32, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![4]u32 { + const t = try smartResize(frames, height, width, min_tokens, max_tokens); + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + var scale = @min(@as(f64, @floatFromInt(t[0])) / h, @as(f64, @floatFromInt(t[1])) / w); + const ppt: f64 = TEMPORAL * FACTOR * FACTOR; + if (@as(f64, @floatFromInt(frames)) * h * w >= @as(f64, @floatFromInt(min_tokens)) * ppt) scale = @min(1.0, scale); + const ch: u32 = @intFromFloat(@max(1, @min(@as(f64, @floatFromInt(t[0])), @floor(h * scale)))); + const cw: u32 = @intFromFloat(@max(1, @min(@as(f64, @floatFromInt(t[1])), @floor(w * scale)))); + return .{ t[0], t[1], ch, cw }; +} + +pub const Prepared = struct { + patches: []f32, // [gh * gw, 1176] + grid: [3]i64, + + pub fn tokens(p: Prepared) usize { + return @intCast(@divExact(p.grid[1] * p.grid[2], MERGE * MERGE)); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +fn norm(v: u8, c: usize) f32 { + const x = @as(f32, @floatFromInt(v)) * @as(f32, @floatCast(@as(f64, 1.0) / 255.0)); + return (x - MEAN[c]) / STD[c]; +} + +/// A frame's normalized value at canvas (y, x) (zero padding past the content's ch x cw). +fn at(frame: []const u8, cw: usize, ch: usize, y: usize, x: usize, c: usize) f32 { + return norm(if (y < ch and x < cw) frame[(y * cw + x) * 3 + c] else 0, c); +} + +/// ``im`` on its canvas within [min_tokens, max_tokens] tokens, normalized and cut into patches. +pub fn prepare(gpa: std.mem.Allocator, im: qi.Image, min_tokens: u64, max_tokens: u64) Error!Prepared { + const g = try geometry(TEMPORAL, im.h, im.w, min_tokens, max_tokens); + var r = if (g[2] == im.h and g[3] == im.w) null else try qi.resize(gpa, im, g[3], g[2]); + defer if (r) |*x| x.deinit(gpa); + const src = if (r) |x| x.rgb else im.rgb; + const gh: usize = g[0] / PATCH; + const gw: usize = g[1] / PATCH; + const out = try gpa.alloc(f32, gh * gw * PER); + var n: usize = 0; + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * PER ..][0..PER]; + for (0..3) |c| for (0..PATCH) |y| for (0..PATCH) |x| { + const v = at(src, g[3], g[2], py + y, px + x, c); + for (0..TEMPORAL) |t| dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = v; + }; + n += 1; + }; + return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; +} + +/// A video's frame groups: patches of two consecutive frames (an odd last repeats) and the first frame's seconds. +pub const Video = struct { + patches: []f32, // groups x gh x gw x 1176 + groups: usize, + gh: usize, + gw: usize, + times: []f64, + + pub fn groupPatches(v: Video, g: usize) []f32 { + const per = v.gh * v.gw * PER; + return v.patches[g * per ..][0..per]; + } + + pub fn tokens(v: Video) usize { + return v.gh * v.gw / (MERGE * MERGE); + } + + pub fn deinit(v: *Video, gpa: std.mem.Allocator) void { + gpa.free(v.patches); + gpa.free(v.times); + } +}; + +/// ``rgb``: ``frames`` RGB frames at the content size (ch x cw: geometry's), placed on the th x tw canvas. +pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, ch: usize, cw: usize, th: usize, tw: usize, indices: []const i32, fps: f64) Error!Video { + if (th % FACTOR != 0 or tw % FACTOR != 0 or ch > th or cw > tw or frames < 1) return error.UnsupportedImage; + const groups = (frames + TEMPORAL - 1) / TEMPORAL; + const gh = th / PATCH; + const gw = tw / PATCH; + const out = try gpa.alloc(f32, groups * gh * gw * PER); + errdefer gpa.free(out); + var n: usize = 0; + for (0..groups) |g| { + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * PER ..][0..PER]; + for (0..3) |c| for (0..TEMPORAL) |t| { + const f = @min(g * TEMPORAL + t, frames - 1); + const frame = rgb[f * ch * cw * 3 ..][0 .. ch * cw * 3]; + for (0..PATCH) |y| for (0..PATCH) |x| { + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = at(frame, cw, ch, py + y, px + x, c); + }; + }; + n += 1; + }; + } + const times = try gpa.alloc(f64, groups); + for (0..groups) |g| times[g] = @as(f64, @floatFromInt(indices[g * TEMPORAL])) / fps; + return .{ .patches = out, .groups = groups, .gh = gh, .gw = gw, .times = times }; +} + +test "smart_resize and geometry match the processor's" { + // 640 x 480 within 8000 tokens: aligned to 504 x 644, content unscaled + try std.testing.expectEqual([4]u32{ 504, 644, 480, 640 }, try geometry(TEMPORAL, 480, 640, 16, 8000)); + // 40 x 30 below 16 tokens: scaled up, content scaled + const g = try geometry(TEMPORAL, 30, 40, 16, 8000); + try std.testing.expectEqual(@as(u32, 112), g[0]); + try std.testing.expectEqual(@as(u32, 140), g[1]); + // 300 x 900 (tall) + const t = try geometry(TEMPORAL, 900, 300, 16, 8000); + try std.testing.expectEqual([2]u32{ 924, 308 }, .{ t[0], t[1] }); +} diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig new file mode 100644 index 0000000000..79e7c4be3a --- /dev/null +++ b/zig/src/vision/qwen_image.zig @@ -0,0 +1,361 @@ +//! Image inputs for Qwen-style vision towers as Qwen2VL's processor makes them: decode, smart_resize, resize, patches. + +const std = @import("std"); +/// GLM-5.3-Flash's preparation (its canvas, padding, normalization and patch order) on this module's decode and resize. +pub const glm = @import("glm_image.zig"); + +pub const Limits = struct { + max_dimension: u32 = 8192, + max_pixels: u64 = 16 * 1024 * 1024, +}; + +/// An RGB image, rows top to bottom, three bytes a pixel. +pub const Image = struct { + w: u32, + h: u32, + rgb: []u8, + + pub fn deinit(im: *Image, gpa: std.mem.Allocator) void { + gpa.free(im.rgb); + } +}; + +pub const Error = error{ UnsupportedImage, ImageTooLarge, OutOfMemory, BadAspectRatio, NoDecoder }; + +/// JPEG or PNG bytes to RGB: the shared runtime decoder's job (pinned libjpeg-turbo); until it lands, refused by name. +pub fn decode(gpa: std.mem.Allocator, bytes: []const u8, limits: Limits) Error!Image { + _ = .{ gpa, bytes, limits }; + return error.NoDecoder; +} + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const r = @round(x); + if (@abs(x - @trunc(x)) == 0.5) return 2.0 * @round(x / 2.0); + return r; +} + +/// transformers' smart_resize: both sides multiples of ``factor``, the area within [min_pixels, max_pixels]. +pub fn smartResize(height: u32, width: u32, factor: u32, min_pixels: u64, max_pixels: u64) Error![2]u32 { + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const f: f64 = @floatFromInt(factor); + if (@max(h, w) / @min(h, w) > 200) return error.BadAspectRatio; + var hb = pyRound(h / f) * f; + var wb = pyRound(w / f) * f; + if (hb * wb > @as(f64, @floatFromInt(max_pixels))) { + const beta = @sqrt(h * w / @as(f64, @floatFromInt(max_pixels))); + hb = @max(f, @floor(h / beta / f) * f); + wb = @max(f, @floor(w / beta / f) * f); + } else if (hb * wb < @as(f64, @floatFromInt(min_pixels))) { + const beta = @sqrt(@as(f64, @floatFromInt(min_pixels)) / (h * w)); + hb = @ceil(h * beta / f) * f; + wb = @ceil(w * beta / f) * f; + } + return .{ @intFromFloat(hb), @intFromFloat(wb) }; +} + +fn bicubic(x0: f64) f64 { + const a = -0.5; + const x = @abs(x0); + if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1; + if (x < 2.0) return (((x - 5) * x + 8) * x - 4) * a; + return 0.0; +} + +const Coeffs = struct { ksize: usize, bounds: []usize, kk: []i32, precision: u5 }; + +/// PIL's precompute_coeffs for an axis, as int16 fixed point at the most precision (< 22 bits) torch's uint8 takes. +fn coeffs(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Coeffs { + const scale = @as(f64, @floatFromInt(in_size)) / @as(f64, @floatFromInt(out_size)); + const filterscale = @max(scale, 1.0); + const support = 2.0 * filterscale; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const bounds = try gpa.alloc(usize, out_size * 2); + const kk = try gpa.alloc(i32, out_size * ksize); + const pre_all = try gpa.alloc(f64, out_size * ksize); + defer gpa.free(pre_all); + for (0..out_size) |xx| { + const pre = pre_all[xx * ksize ..][0..ksize]; + const center = (@as(f64, @floatFromInt(xx)) + 0.5) * scale; + var ww: f64 = 0; + const ss = 1.0 / filterscale; + var xmin_i: i64 = @intFromFloat(center - support + 0.5); + if (xmin_i < 0) xmin_i = 0; + var xmax_i: i64 = @intFromFloat(center + support + 0.5); + if (xmax_i > @as(i64, @intCast(in_size))) xmax_i = @intCast(in_size); + const xmin: usize = @intCast(xmin_i); + const xmax: usize = @intCast(xmax_i - xmin_i); + @memset(pre, 0); + for (0..xmax) |x| { + const wgt = bicubic((@as(f64, @floatFromInt(x + xmin)) - center + 0.5) * ss); + pre[x] = wgt; + ww += wgt; + } + for (0..xmax) |x| { + if (ww != 0.0) pre[x] /= ww; + } + bounds[xx * 2] = xmin; + bounds[xx * 2 + 1] = xmax; + } + var w_max: f64 = 0; + for (pre_all) |v| w_max = @max(w_max, @abs(v)); + var precision: u5 = 0; + while (precision < 22) : (precision += 1) { + const next: i64 = @intFromFloat(0.5 + w_max * @as(f64, @floatFromInt(@as(i64, 1) << (precision + 1)))); + if (next >= (1 << 15)) break; + } + const one: f64 = @floatFromInt(@as(i64, 1) << precision); + for (pre_all, 0..) |v, i| kk[i] = @intFromFloat(if (v < 0) -0.5 + v * one else 0.5 + v * one); + return .{ .ksize = ksize, .bounds = bounds, .kk = kk, .precision = precision }; +} + +fn clip8(v: i64, precision: u5) u8 { + const r = v >> precision; + return @intCast(std.math.clamp(r, 0, 255)); +} + +/// The antialiased bicubic resize (PIL's ImagingResample, 8 bits): horizontal pass, then vertical, rounding to bytes. +pub fn resize(gpa: std.mem.Allocator, im: Image, ow: u32, oh: u32) Error!Image { + var cur = im; + var owned = false; + defer if (owned) cur.deinit(gpa); + if (ow != im.w) { + const c = try coeffs(gpa, im.w, ow); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, ow) * im.h * 3); + for (0..im.h) |y| for (0..ow) |x| { + const xmin = c.bounds[x * 2]; + const xmax = c.bounds[x * 2 + 1]; + const k = c.kk[x * c.ksize ..]; + for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..xmax) |j| ss += @as(i64, cur.rgb[(y * cur.w + xmin + j) * 3 + ch]) * k[j]; + out[(y * ow + x) * 3 + ch] = clip8(ss, c.precision); + } + }; + cur = .{ .w = ow, .h = im.h, .rgb = out }; + owned = true; + } + if (oh != cur.h) { + const c = try coeffs(gpa, cur.h, oh); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, cur.w) * oh * 3); + for (0..oh) |y| { + const ymin = c.bounds[y * 2]; + const ymax = c.bounds[y * 2 + 1]; + const k = c.kk[y * c.ksize ..]; + for (0..cur.w) |x| for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..ymax) |j| ss += @as(i64, cur.rgb[((ymin + j) * cur.w + x) * 3 + ch]) * k[j]; + out[(y * cur.w + x) * 3 + ch] = clip8(ss, c.precision); + }; + } + if (owned) cur.deinit(gpa); + cur = .{ .w = cur.w, .h = oh, .rgb = out }; + owned = true; + } + if (!owned) return .{ .w = im.w, .h = im.h, .rgb = try gpa.dupe(u8, im.rgb) }; + owned = false; + return cur; +} + +/// One image prepared for the tower: its patches (fp32 [gh * gw, 1536], merge-block order) and grid [1, gh, gw]. +pub const Prepared = struct { + patches: []f32, + grid: [3]i64, + + pub fn tokens(p: Prepared) usize { + return @intCast(@divExact(p.grid[1] * p.grid[2], 4)); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +pub const PATCH = 16; +pub const MERGE = 2; +pub const TEMPORAL = 2; +pub const FACTOR = PATCH * MERGE; + +/// ``im`` resized within ``max_tokens`` visual tokens (min_pixels: one merged patch), normalized and cut into patches. +pub fn prepare(gpa: std.mem.Allocator, im: Image, max_tokens: u64) Error!Prepared { + const cap = max_tokens * FACTOR * FACTOR; + const size = try smartResize(im.h, im.w, FACTOR, @min(FACTOR * FACTOR, cap), cap); + var r = try resize(gpa, im, size[1], size[0]); + defer r.deinit(gpa); + const gh: usize = size[0] / PATCH; + const gw: usize = size[1] / PATCH; + const per = 3 * TEMPORAL * PATCH * PATCH; + const out = try gpa.alloc(f32, gh * gw * per); + var n: usize = 0; + var bh: usize = 0; + while (bh < gh / MERGE) : (bh += 1) { + var bw: usize = 0; + while (bw < gw / MERGE) : (bw += 1) { + for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * per ..][0..per]; + for (0..3) |c| for (0..TEMPORAL) |t| for (0..PATCH) |y| for (0..PATCH) |x| { + const v: f32 = @floatFromInt(r.rgb[((py + y) * r.w + px + x) * 3 + c]); + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = (v - 127.5) / 127.5; + }; + n += 1; + }; + } + } + return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; +} + +// -- videos --------------------------------------------------------------------------------------------------------- + +/// A video's frame size: Qwen3-VL's smart_resize, at most 768 tokens a frame group, the video within ``budget``. +pub fn videoSize(frames: u32, height0: u32, width0: u32, budget: u64) Error![2]u32 { + const factor: f64 = FACTOR; + const shortest: f64 = 128 * FACTOR * FACTOR; + const longest: f64 = @as(f64, @floatFromInt(budget)) * TEMPORAL * FACTOR * FACTOR; + const per_frame = @max(@min(768 * FACTOR * FACTOR, @floor(longest / @as(f64, @floatFromInt(@max(1, frames))))), @floor(shortest * 1.05)); + const max_pixels = per_frame * @as(f64, @floatFromInt(frames)); + if (frames < TEMPORAL) return error.UnsupportedImage; + var height: f64 = @floatFromInt(height0); + var width: f64 = @floatFromInt(width0); + if (height < factor or width < factor) { + const scale = @max(factor / height, factor / width); + height = @trunc(height * scale); + width = @trunc(width * scale); + } + if (@max(height, width) / @min(height, width) > 200) return error.BadAspectRatio; + var hb = pyRound(height / factor) * factor; + var wb = pyRound(width / factor) * factor; + const tb = pyRound(@as(f64, @floatFromInt(frames)) / TEMPORAL) * TEMPORAL; + const n: f64 = @floatFromInt(frames); + if (tb * hb * wb > max_pixels) { + const beta = @sqrt(n * height * width / max_pixels); + hb = @max(factor, @floor(height / beta / factor) * factor); + wb = @max(factor, @floor(width / beta / factor) * factor); + } else if (tb * hb * wb < shortest) { + const beta = @sqrt(shortest / (n * height * width)); + hb = @ceil(height * beta / factor) * factor; + wb = @ceil(width * beta / factor) * factor; + } + return .{ @intFromFloat(hb), @intFromFloat(wb) }; +} + +/// A video's frame groups: patches of two consecutive frames, normalized as numpy does, and their mean time. +pub const Video = struct { + patches: []f32, // groups x gh x gw x 1536 + groups: usize, + gh: usize, + gw: usize, + times: []f64, + + pub fn groupPatches(v: Video, g: usize) []f32 { + const per = v.gh * v.gw * 3 * TEMPORAL * PATCH * PATCH; + return v.patches[g * per ..][0..per]; + } + + pub fn tokens(v: Video) usize { + return v.gh * v.gw / (MERGE * MERGE); + } + + pub fn deinit(v: *Video, gpa: std.mem.Allocator) void { + gpa.free(v.patches); + gpa.free(v.times); + } +}; + +/// ``rgb``: ``frames`` RGB frames of ``h`` x ``w`` (already at the tower's size), their source ``indices`` and ``fps``. +pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, h: usize, w: usize, indices: []const i32, fps: f64) Error!Video { + if (h % FACTOR != 0 or w % FACTOR != 0 or frames < 2) return error.UnsupportedImage; + const groups = (frames + TEMPORAL - 1) / TEMPORAL; + const gh = h / PATCH; + const gw = w / PATCH; + const per = 3 * TEMPORAL * PATCH * PATCH; + const out = try gpa.alloc(f32, groups * gh * gw * per); + errdefer gpa.free(out); + const scale: f32 = @floatCast(@as(f64, 1.0) / 255.0); + var n: usize = 0; + for (0..groups) |g| { + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * per ..][0..per]; + for (0..3) |c| for (0..TEMPORAL) |t| { + const f = @min(g * TEMPORAL + t, frames - 1); + const frame = rgb[f * h * w * 3 ..]; + for (0..PATCH) |y| for (0..PATCH) |x| { + const v: f32 = @as(f32, @floatFromInt(frame[((py + y) * w + px + x) * 3 + c])) * scale; + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = (v - 0.5) / 0.5; + }; + }; + n += 1; + }; + } + const times = try gpa.alloc(f64, groups); + for (0..groups) |g| { + const a: f64 = @as(f64, @floatFromInt(indices[g * TEMPORAL])) / fps; + const b: f64 = @as(f64, @floatFromInt(indices[@min(g * TEMPORAL + TEMPORAL - 1, frames - 1)])) / fps; + times[g] = (a + b) / 2; + } + return .{ .patches = out, .groups = groups, .gh = gh, .gw = gw, .times = times }; +} + +/// Python's f"{t:.1f}": one decimal, an exact tie rounded to even. +pub fn formatSeconds(buf: []u8, t: f64) ![]const u8 { + const v = tenths(t); + return std.fmt.bufPrint(buf, "{d}.{d}", .{ @divTrunc(v, 10), @as(u64, @intCast(@mod(v, 10))) }); +} + +/// t times ten, an exact tie to even, in integers: t is m * 2^e and 10m fits 57 bits (f128 would link compiler-rt). +fn tenths(t: f64) i64 { + const bits: u64 = @bitCast(@abs(t)); + const biased: i32 = @intCast(bits >> 52); + const m: u64 = if (biased == 0) bits & ((1 << 52) - 1) else (bits & ((1 << 52) - 1)) | (1 << 52); + const e = @as(i32, @max(biased, 1)) - 1075; + if (e >= 0) return @intFromFloat(t * 10); // at least 2^52 seconds: never a frame's time + const tm = m * 10; + if (e <= -58) return 0; // t * 10 under a half: tm < 2^57 <= 2^(s - 1) + const s: u6 = @intCast(-e); + const q = tm >> s; + const rem = tm & ((@as(u64, 1) << s) - 1); + const half = @as(u64, 1) << (s - 1); + const r: i64 = @intCast(if (rem > half or (rem == half and q & 1 == 1)) q + 1 else q); + return if (t < 0) -r else r; +} + +test "tenths rounds t * 10 exactly, ties to even, as the f128 product does" { + const ref = struct { + fn f(t: f64) i64 { + const x: f128 = @as(f128, t) * 10; + const fl = @floor(x); + const frac = x - fl; + var r = if (frac > 0.5) fl + 1 else fl; + if (frac == 0.5 and @rem(fl, 2) != 0) r = fl + 1; + return @intFromFloat(r); + } + }.f; + const fixed = [_]f64{ 0, -0.0, 0.05, 0.15, 0.25, 0.35, 0.45, 1.25, 2.5, 0.95, 1.05, 3.0, 123.45, 1e-300, 5e-324, 4503599627370495.5, -0.25, -1.35 }; + for (fixed) |t| try std.testing.expectEqual(ref(t), tenths(t)); + var prng: std.Random.DefaultPrng = .init(0x7e17); + const rand = prng.random(); + for (0..200_000) |i| { + const t: f64 = switch (i % 4) { + 0 => rand.float(f64) * 4000, + 1 => @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 0, 80_000))) / 20, // halves of tenths + 2 => @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 1, 100_000))) / @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 1, 121))), + else => @bitCast(@as(u64, rand.uintLessThan(u64, 1075)) << 52 | rand.int(u64) >> 12), // any double below 2^52 + }; + try std.testing.expectEqual(ref(t), tenths(t)); + } + var buf: [32]u8 = undefined; + try std.testing.expectEqualStrings("0.2", try formatSeconds(&buf, 0.25)); + try std.testing.expectEqualStrings("1.4", try formatSeconds(&buf, 1.35)); +} + +test { + _ = glm; +} diff --git a/zig/src/vision/tf_video.c b/zig/src/vision/tf_video.c new file mode 100644 index 0000000000..9cd1290cc7 --- /dev/null +++ b/zig/src/vision/tf_video.c @@ -0,0 +1,211 @@ +// libtfvideo: a video's sampled frames as videos.py decodes them with PyAV 19.0.1 (FFmpeg 9), RGB via SWS_BICUBIC. +#include +#include +#include +#include +#include +#include +#include +#include + +typedef int (*tf_size_fn)(void* ctx, int frames, int height, int width, int* out_h, int* out_w); + +typedef struct { + uint8_t* rgb; // frames x height x width x 3 + int frames, height, width; + int* indices; // each frame's index in the source + double fps; // the source's rate +} tf_video; + +enum { + TF_OK = 0, + TF_INVALID = -1, // not a decodable container or stream + TF_NO_STREAM = -2, // no video stream + TF_DIMENSIONS = -3, // missing or over max_dim + TF_NO_COUNT = -4, // the frame count or rate is missing + TF_TOO_LONG = -5, // over max_seconds + TF_FEW_FRAMES = -6, // fewer than two decodable frames + TF_MEMORY = -7, + TF_SIZE = -8, // the size callback refused +}; + +typedef struct { const uint8_t* data; size_t len, pos; } mem_t; + +static int mem_read(void* o, uint8_t* buf, int n) { + mem_t* m = (mem_t*)o; + if (m->pos >= m->len) return AVERROR_EOF; + size_t k = m->len - m->pos < (size_t)n ? m->len - m->pos : (size_t)n; + memcpy(buf, m->data + m->pos, k); + m->pos += k; + return (int)k; +} + +static int64_t mem_seek(void* o, int64_t off, int whence) { + mem_t* m = (mem_t*)o; + if (whence == AVSEEK_SIZE) return (int64_t)m->len; + int64_t at = whence == SEEK_SET ? off : whence == SEEK_CUR ? (int64_t)m->pos + off : (int64_t)m->len + off; + if (at < 0 || at > (int64_t)m->len) return -1; + m->pos = (size_t)at; + return at; +} + +// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): the frame's own AVColorSpace passes through. +static void set_colorspace(AVFrame* f, int cs) { + if (cs == SWS_CS_ITU709) f->colorspace = AVCOL_SPC_BT709; + else if (cs == SWS_CS_FCC) f->colorspace = AVCOL_SPC_FCC; + else if (cs == SWS_CS_ITU601) f->colorspace = AVCOL_SPC_SMPTE170M; + else if (cs == SWS_CS_SMPTE240M) f->colorspace = AVCOL_SPC_SMPTE240M; + else if (cs == SWS_CS_BT2020) f->colorspace = AVCOL_SPC_BT2020_NCL; +} + +static double round_half_even(double x) { + double r = round(x); + if (fabs(x - trunc(x)) == 0.5) r = 2.0 * round(x / 2.0); + return r; +} + +void tf_video_free(tf_video* v) { + if (!v) return; + free(v->rgb); + free(v->indices); + v->rgb = NULL; + v->indices = NULL; +} + +int tf_video_decode(const uint8_t* data, size_t len, double rate, int min_frames, int max_frames, double max_seconds, + int max_dim, tf_size_fn size, void* ctx, tf_video* out) { + memset(out, 0, sizeof(*out)); + int rc = TF_INVALID; + mem_t mem = {data, len, 0}; + AVFormatContext* fmt = NULL; + AVIOContext* io = NULL; + AVCodecContext* dec = NULL; + AVPacket* pkt = NULL; + AVFrame* frame = NULL; + struct SwsContext* sws = NULL; + int* wanted = NULL; + unsigned char* iobuf = av_malloc(1 << 16); + if (!iobuf) return TF_MEMORY; + io = avio_alloc_context(iobuf, 1 << 16, 0, &mem, mem_read, NULL, mem_seek); + if (!io) { av_free(iobuf); return TF_MEMORY; } + fmt = avformat_alloc_context(); + if (!fmt) { rc = TF_MEMORY; goto done; } + fmt->pb = io; + if (avformat_open_input(&fmt, NULL, NULL, NULL) < 0) { fmt = NULL; goto done; } + if (avformat_find_stream_info(fmt, NULL) < 0) goto done; + int si = av_find_best_stream(fmt, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0); + if (si < 0) { rc = TF_NO_STREAM; goto done; } + // PyAV's streams.video[0] is the first video stream + for (unsigned i = 0; i < fmt->nb_streams; ++i) + if (fmt->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) { si = (int)i; break; } + AVStream* st = fmt->streams[si]; + int width = st->codecpar->width, height = st->codecpar->height; + if (width <= 0 || height <= 0 || width > max_dim || height > max_dim) { rc = TF_DIMENSIONS; goto done; } + AVRational r = st->avg_frame_rate; + if (r.num <= 0 || r.den <= 0) r = av_guess_frame_rate(fmt, st, NULL); + double fps = (r.num > 0 && r.den > 0) ? av_q2d(r) : 24.0; + double seconds = 0.0; + if (st->duration > 0 && st->duration != AV_NOPTS_VALUE) seconds = (double)st->duration * av_q2d(st->time_base); + else if (fmt->duration > 0 && fmt->duration != AV_NOPTS_VALUE) seconds = (double)fmt->duration / 1e6; + int64_t total = st->nb_frames > 0 ? st->nb_frames : (int64_t)round_half_even(seconds * fps); + if (!(fps > 0 && fps <= 1000) || total <= 0) { rc = TF_NO_COUNT; goto done; } + if ((double)total / fps > max_seconds) { rc = TF_TOO_LONG; goto done; } + int64_t count = (int64_t)((double)total / fps * rate); + if (count < min_frames) count = min_frames; + if (count > max_frames) count = max_frames; + if (count > total) count = total; + wanted = malloc(sizeof(int) * (size_t)count); + if (!wanted) { rc = TF_MEMORY; goto done; } + // numpy.linspace(0, total - 1, count): step = (stop - start) / (count - 1), the last exactly stop + double stop = (double)(total - 1); + for (int64_t i = 0; i < count; ++i) { + double v = count > 1 ? (double)i * (stop / (double)(count - 1)) : 0.0; + if (count > 1 && i == count - 1) v = stop; + wanted[i] = (int)round_half_even(v); + } + int oh = 0, ow = 0; + if (size(ctx, (int)count, height, width, &oh, &ow) != 0 || oh <= 0 || ow <= 0) { rc = TF_SIZE; goto done; } + const AVCodec* codec = avcodec_find_decoder(st->codecpar->codec_id); + if (!codec) goto done; + dec = avcodec_alloc_context3(codec); + if (!dec || avcodec_parameters_to_context(dec, st->codecpar) < 0) goto done; + dec->thread_count = 0; // PyAV's thread_type "AUTO" + dec->thread_type = FF_THREAD_FRAME | FF_THREAD_SLICE; + if (avcodec_open2(dec, codec, NULL) < 0) goto done; + out->rgb = malloc((size_t)count * (size_t)oh * (size_t)ow * 3); + out->indices = malloc(sizeof(int) * (size_t)count); + pkt = av_packet_alloc(); + frame = av_frame_alloc(); + if (!out->rgb || !out->indices || !pkt || !frame) { rc = TF_MEMORY; goto done; } + AVFrame* rgb = av_frame_alloc(); + if (!rgb) { rc = TF_MEMORY; goto done; } + sws = sws_alloc_context(); + if (!sws) { av_frame_free(&rgb); rc = TF_MEMORY; goto done; } + sws->threads = 0; + sws->flags = SWS_BICUBIC; + int n = 0, next = 0, got = 0, last = wanted[count - 1]; + int flushing = 0, finished = 0; + while (!finished) { + if (!flushing) { + int e = av_read_frame(fmt, pkt); + if (e < 0) { flushing = 1; avcodec_send_packet(dec, NULL); } + else { + if (pkt->stream_index == si) avcodec_send_packet(dec, pkt); + av_packet_unref(pkt); + } + } + for (;;) { + int e = avcodec_receive_frame(dec, frame); + if (e == AVERROR(EAGAIN)) { + if (flushing) finished = 1; // a flushed decoder only ends + break; + } + if (e < 0) { finished = 1; break; } + if (n > last) { finished = 1; av_frame_unref(frame); break; } + if (next < count && n == wanted[next]) { + // PyAV's _reformat: the frame's colour space and range kept; transfer and primaries unspecified + av_frame_unref(rgb); + av_frame_copy_props(rgb, frame); + set_colorspace(frame, (int)frame->colorspace); + rgb->colorspace = frame->colorspace; + rgb->color_range = frame->color_range; + set_colorspace(rgb, (int)frame->colorspace); + const enum AVColorTransferCharacteristic trc = frame->color_trc; + const enum AVColorPrimaries pri = frame->color_primaries; + frame->color_trc = rgb->color_trc = AVCOL_TRC_UNSPECIFIED; + frame->color_primaries = rgb->color_primaries = AVCOL_PRI_UNSPECIFIED; + rgb->format = AV_PIX_FMT_RGB24; + rgb->width = ow; + rgb->height = oh; + if (av_frame_get_buffer(rgb, 0) < 0) { rc = TF_MEMORY; av_frame_unref(frame); av_frame_free(&rgb); goto done; } + int e2 = sws_scale_frame(sws, rgb, frame); + frame->color_trc = trc; + frame->color_primaries = pri; + if (e2 < 0) { av_frame_unref(frame); av_frame_free(&rgb); goto done; } + uint8_t* to = out->rgb + (size_t)got * oh * ow * 3; + for (int y = 0; y < oh; ++y) memcpy(to + (size_t)y * ow * 3, rgb->data[0] + (size_t)y * rgb->linesize[0], (size_t)ow * 3); + out->indices[got++] = n; + while (next < count && wanted[next] == n) ++next; + } + ++n; + av_frame_unref(frame); + } + } + av_frame_free(&rgb); + if (got < 2) { rc = TF_FEW_FRAMES; goto done; } + out->frames = got; + out->height = oh; + out->width = ow; + out->fps = fps; + rc = TF_OK; +done: + if (rc != TF_OK) tf_video_free(out); + free(wanted); + sws_freeContext(sws); + av_frame_free(&frame); + av_packet_free(&pkt); + avcodec_free_context(&dec); + if (fmt) avformat_close_input(&fmt); + if (io) { av_freep(&io->buffer); avio_context_free(&io); } + return rc; +} 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/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..656dea33d8 --- /dev/null +++ b/zig/tests/qwen35_cuda_main.zig @@ -0,0 +1,307 @@ +//! `tf-qwen35-cuda ...`: rows to compare (TF_Q35_RESUME: from a copy). +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); + const is_vision = args.len >= 3 and std.mem.eql(u8, args[2], "vision"); + if (args.len != 5 and !(is_vision and args.len <= 7)) { + std.debug.print("usage: tf-qwen35-cuda \n" ++ + " tf-qwen35-cuda vision [min mean cosine, 0.999] [min row cosine, 0.94]\n", .{}); + return 2; + } + if (is_vision) { + const bounds: Bounds = .{ + .mean = if (args.len > 5) try std.fmt.parseFloat(f64, args[5]) else 0.999, + .row = if (args.len > 6) try std.fmt.parseFloat(f64, args[6]) else 0.94, + }; + return vision(init, a, args[1], args[3], args[4], bounds); + } + 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(); + // 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); + // TF_Q35_DELTA=d: the rows after the prefill rotate at position + d, as an image prompt's reply does + if (std.c.getenv("TF_Q35_DELTA")) |v| cache.delta = try std.fmt.parseInt(i32, std.mem.span(v), 10); + // TF_Q35_RESUME=1: the stream goes on from a kept copy of its state, restored into another (poisoned) cache + const resume_at = try envCount("TF_Q35_RESUME"); + if (resume_at > 0 and !std.mem.eql(u8, args[2], "prompt") and !std.mem.eql(u8, args[2], "mrope")) try resumeFrom(init.gpa, e, &cache); + 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 { + // mrope: chunks through an image prompt's table (text positions, no image rows): prompt's bits + const table = std.mem.eql(u8, args[2], "mrope"); + const n_ids = ids.items.len; + const src = try a.alloc(i32, n_ids); + @memset(src, -1); + const pos = try a.alloc(i32, 3 * n_ids); + // TF_Q35_SHIFT=k moves the h axis by k: a control whose logits must differ, as the table is read + const shift = try envCount("TF_Q35_SHIFT"); + for (0..3) |ax| for (0..n_ids) |i| { + pos[ax * n_ids + i] = @intCast(cache.len + i + if (ax == 1) shift else 0); + }; + var none = try cuda.DeviceBuffer.alloc(ctx.d, 256); + defer none.free(); + const im: E.Images = .{ .features = none.ptr, .src = src, .pos = pos }; + const p0 = cache.len; + var at: usize = 0; + while (at < n_ids) { + const n = @min(E.max_rows, n_ids - at); + if (table) { + // the table's rows are the prompt's from position p0: shift it so row `at` is the cache's next + const part: E.Images = .{ .features = im.features, .src = try padded(a, src, p0), .pos = try paddedPos(a, pos, p0, n_ids) }; + try e.forward(&.{.{ .cache = &cache, .rows = n, .images = &part }}, ids.items[at..][0..n], .last); + } else try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .last); + at += n; + // TF_Q35_RESUME=K in prompt and mrope: the pass goes on from a kept copy at position K + if (resume_at > 0 and cache.len == resume_at) try resumeFrom(init.gpa, e, &cache); + } + 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; +} + +/// `cache` replaced by a fresh one, poisoned, that takes a kept copy of its state (cuda_reuse.zig's save and restore). +fn resumeFrom(gpa: std.mem.Allocator, e: *E.Engine, cache: *E.Cache) !void { + const copy = try q.reuse.saveCache(e, cache, cache.len, &.{}); + defer q.reuse.dropCopy(e, copy); + var fresh = try E.Cache.init(gpa, e, cache.cap); + errdefer fresh.deinit(gpa); + try fresh.state.fill8(0x5a, e.stream.handle); + try fresh.kv.fill8(0x5a, e.stream.handle); + try q.reuse.restoreCache(e, &fresh, copy, &.{}); + if (fresh.len != cache.len or fresh.delta != cache.delta) return error.ResumeMismatch; + cache.deinit(gpa); + cache.* = fresh; +} + +/// 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); + } +}; + +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; + } +} + +/// Each image's mean and worst row cosine to 0.6.6's bf16 rows; with fp32 rows, no further from them than 0.6.6's. +const Bounds = struct { mean: f64, row: f64, slack: f64 = 0.001 }; + +/// Rows' cosine, mean and worst, of bf16 `got` against `want` (bf16, or fp32 when `want32`), and the identical rows. +const Cosines = struct { mean: f64, least: f64, same: usize }; + +fn cosines(got: []align(1) const u16, want: []const u8, want32: bool, rows: usize, width: usize) Cosines { + var sum: f64 = 0; + var least: f64 = 1; + var same: usize = 0; + for (0..rows) |r| { + var dot: f64 = 0; + var ng: f64 = 0; + var nw: f64 = 0; + var equal = !want32; + for (0..width) |i| { + const x = bf(got[r * width + i]); + const y: f64 = if (want32) @as(f64, @as(f32, @bitCast(std.mem.readInt(u32, want[(r * width + i) * 4 ..][0..4], .little)))) else bf(std.mem.readInt(u16, want[(r * width + i) * 2 ..][0..2], .little)); + if (!want32 and got[r * width + i] != std.mem.readInt(u16, want[(r * width + i) * 2 ..][0..2], .little)) equal = false; + dot += x * y; + ng += x * x; + nw += y * y; + } + const cos = dot / @sqrt(ng * nw); + sum += cos; + least = @min(least, cos); + same += @intFromBool(equal); + } + return .{ .mean = sum / @as(f64, @floatFromInt(rows)), .least = least, .same = same }; +} + +/// Each fixture's patches (tools/zig/qwen_vision_fixtures.py) through the tower: rows' cosine to 0.6.6's, bounded. +fn vision(init: std.process.Init, a: std.mem.Allocator, model: []const u8, dir: []const u8, out_path: []const u8, bounds: Bounds) !u8 { + const V = q.vision; + const io = init.io; + const cfg = try V.Config.read(a, io, model) orelse return error.QwenCheckpointHasNoVision; + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var stream = try cuda.Stream.init(ctx.d, true); + defer stream.deinit(); + const tower = try V.Tower.load(init.gpa, io, ctx.d, stream, model, cfg); + defer tower.deinit(); + const cwd = std.Io.Dir.cwd(); + const manifest = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "manifest.json" }), a, .limited(1 << 20)); + const parsed = try std.json.parseFromSliceLeaky(std.json.Value, a, manifest, .{}); + var out: std.ArrayList(u8) = .empty; + var worst: f64 = 1; + var failed: usize = 0; + for (parsed.object.get("images").?.array.items) |im| { + const o = im.object; + const grid = o.get("grid").?.array.items; + const g: V.Grid = .{ .h = @intCast(grid[0].integer), .w = @intCast(grid[1].integer) }; + const tokens: usize = @intCast(o.get("tokens").?.integer); + const px_bytes = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, o.get("pixels").?.string }), a, .limited(1 << 30)); + const want = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, o.get("features").?.string }), a, .limited(1 << 30)); + const px = try a.alloc(f32, px_bytes.len / 4); + @memcpy(std.mem.sliceAsBytes(px), px_bytes); + var features = try cuda.DeviceBuffer.alloc(ctx.d, tokens * cfg.out_hidden * 2); + defer features.free(); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + try tower.encode(px, &.{g}, features.ptr); + const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - t0)) / 1e6; + const got = try a.alloc(u8, tokens * cfg.out_hidden * 2); + try features.download(0, got); + try out.appendSlice(a, got); + if (want.len != got.len) { + std.debug.print("{s}: 0.6.6 has {d} bytes of rows, the tower {d}\n", .{ o.get("name").?.string, want.len, got.len }); + return 1; + } + const gw: []align(1) const u16 = std.mem.bytesAsSlice(u16, got); + const name = o.get("name").?.string; + const c = cosines(gw, want, false, tokens, cfg.out_hidden); + worst = @min(worst, c.least); + std.debug.print("{s}: grid {d}x{d}, {d} rows, row cosine mean {d:.5} min {d:.5}, {d} rows identical, {d:.1} ms", .{ name, g.h, g.w, tokens, c.mean, c.least, c.same, ms }); + var ok = c.mean >= bounds.mean and c.least >= bounds.row; + if (o.get("reference")) |ref| { + // both bf16 towers against 0.6.6's fp32 rows: this one may not be the further, beyond the slack + const r32 = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, ref.string }), a, .limited(1 << 31)); + if (r32.len != tokens * cfg.out_hidden * 4) { + std.debug.print("\n{s}: the fp32 reference has {d} bytes, {d} rows need {d}\n", .{ name, r32.len, tokens, tokens * cfg.out_hidden * 4 }); + return 1; + } + const ours = cosines(gw, r32, true, tokens, cfg.out_hidden); + const theirs = cosines(std.mem.bytesAsSlice(u16, want), r32, true, tokens, cfg.out_hidden); + std.debug.print("; against fp32: mean {d:.5} (0.6.6 bf16 {d:.5})", .{ ours.mean, theirs.mean }); + // with a reference, the absolute bounds give way to it: 0.6.6's own bf16 rows can sit under them + ok = ours.mean + bounds.slack >= theirs.mean; + } + std.debug.print("{s}\n", .{if (ok) "" else " FAIL"}); + failed += @intFromBool(!ok); + } + try cwd.writeFile(io, .{ .sub_path = out_path, .data = out.items }); + std.debug.print("worst row cosine {d:.5}; {d} image{s} outside the bounds (mean {d}, row {d}, fp32 slack {d})\n", .{ worst, failed, if (failed == 1) "" else "s", bounds.mean, bounds.row, bounds.slack }); + return if (failed == 0) 0 else 1; +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +/// `src` behind `p0` rows of text (-1): the table indexed by the cache's position, as a prompt from position 0 is. +fn padded(a: std.mem.Allocator, src: []const i32, p0: usize) ![]i32 { + const out = try a.alloc(i32, p0 + src.len); + @memset(out[0..p0], -1); + @memcpy(out[p0..], src); + return out; +} + +/// `pos` ([3][n]) behind `p0` rows, each axis: the positions indexed by the cache's position. +fn paddedPos(a: std.mem.Allocator, pos: []const i32, p0: usize, n: usize) ![]i32 { + const m = p0 + n; + const out = try a.alloc(i32, 3 * m); + for (0..3) |ax| { + for (0..p0) |i| out[ax * m + i] = @intCast(i); + @memcpy(out[ax * m + p0 ..][0..n], pos[ax * n ..][0..n]); + } + return out; +} 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;