diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 597ff6b522..5be651004e 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -41,3 +41,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 c46dcad158..757dba31ff 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -55,9 +55,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 cdfeca38f8..aade55795e 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}' ``` @@ -48,12 +50,12 @@ src=$(ls -d ~/.cache/huggingface/hub/models--TensorFold--NVIDIA-Nemotron-3.5-Lig mkdir -p ~/models && cp -cRL "$src" ~/models/nemotron-living # always teach a copy (umask 077; openssl rand -hex 24 > ~/models/slide-key) # --slide starts only with a key tensorfold serve ~/models/nemotron-living --slide --api-key-file ~/models/slide-key -curl -N http://127.0.0.1:8080/v1/slide/learn -H "Authorization: Bearer $(cat ~/models/slide-key)" \ +curl -N http://127.0.0.1:8888/v1/slide/learn -H "Authorization: Bearer $(cat ~/models/slide-key)" \ -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 @@ -76,8 +78,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; images with `--vision` ([image input](docs/vision.md#qwen38-27b)) | 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; images with `--vision` ([image input](docs/vision.md#qwen38-27b)) | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (Windows) | | LFM2.5-24B-A2B | MLX affine 4-, 6- or 8-bit, group 64, tied embeddings | Metal on M4 Pro | Nemotron's named checkpoint is `TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`. @@ -110,6 +112,7 @@ 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 the native server runs Nemotron 3.5 Lightning on a GB10, an RTX 3090 and an RTX PRO 6000, and Nemotron 3.5 Super VL with images on one GB10 or two, greedy and sampled, with concurrent requests sharing each round. GLM-5.3-Flash also serves across two CUDA ranks, and Liquid AI's d1-3B answers `/v1/systemone` and `/v1/decisions`. +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.0 (the A100) and 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; FP8 and NVFP4 checkpoints need compute capability 9.0 or newer. ## Exact decoding @@ -177,7 +180,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 d837f9423b..0fbda42425 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 380a38b3de..5582b36da8 100644 --- a/build.zig +++ b/build.zig @@ -18,6 +18,7 @@ const server_build = @import("zig/build/server.zig"); const golden_build = @import("zig/build/golden.zig"); const contract_build = @import("zig/build/kernel_contract.zig"); const video_build = @import("zig/build/video.zig"); +const sys_build = @import("zig/build/sys.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -163,14 +164,14 @@ pub fn build(b: *std.Build) void { merge_build.targets(b, target, optimize); // tf-merge-check and test-merge, host only on every OS 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); xpu_build.hostTests(b, test_step); - 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 }} }); @@ -204,6 +205,7 @@ pub fn build(b: *std.Build) void { b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); contract_build.hostTests(b); + 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). @@ -557,7 +559,8 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = helpers })).step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = engine })).step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = lanes })).step); - test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fabric })).step); + // the fake fabric maps its windows with mmap: the Macs' RDMA, not Windows + if (target.result.os.tag != .windows) test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fabric })).step); const dashboard = b.createModule(.{ .root_source_file = b.path("zig/src/server/dashboard_test.zig"), .target = target, .optimize = optimize, .link_libc = true }); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = dashboard })).step); const kimi_test = b.step("test-k3", "Kimi K3's host-side unit tests (no GPU work)"); diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md new file mode 100644 index 0000000000..400b2eed70 --- /dev/null +++ b/docs/WINDOWS.md @@ -0,0 +1,99 @@ +# 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`, `affine8`, `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. + +### Mixed-width checkpoints + +The Qwen3.5 engine also loads MLX checkpoints whose modules carry their own widths (the config's `quantization` +table names a width per module, as mlx-lm's mixed recipes and the OptiQ and oQ conversions write), 2 to 8 bits in +groups of 64. Nothing is rounded: a module of 2 or 3 bits is packed onto the 4-bit tiles and one of 5 or 6 bits onto +8-bit tiles, each code keeping its value, scale and bias, and the 8-bit tiles run `affine8.cu`'s lane GEMV, the +4-bit one's arithmetic a byte wide. An embedding of an odd width is widened the same way at load. Windows still +reproduce one-token steps bit for bit, and a uniform 4-bit checkpoint loads and runs exactly as before. + +- Tested on an RTX 4080 SUPER (16 GB): `mlx-community/Qwen3.5-4B-OptiQ-4bit` (4 and 8 bits, a tied 8-bit + embedding) at about 130 tokens/s for one request and 400 across four, and `mlx-community/Qwen3.5-9B-OptiQ-4bit` + (an untied 8-bit head) at about 78 and 260; replies alone, at four streams and with `"draft": false` equal. +- Widened modules take their served width's memory: a 2-bit module the 4-bit size, a 5-bit one the 8-bit size. So a + mostly 2-bit checkpoint such as `TensorFold/Qwen3.8-27B-oQ2` needs about as much memory as the 4-bit 27B. +- 8-bit modules have no prompt GEMM yet: prompt chunks run them on the decode tiles, 16 rows at a time. +- `tf-cuda-test affine8` checks every width's tiles against its codes widened on the host, the 8-bit products + against float64, and every row count and split-K path byte for byte. + +## 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 85194d2150..ecfe474bad 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 @@ -586,5 +586,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/community-benchmarks.md b/docs/community-benchmarks.md index 9603239a13..d8df042745 100644 --- a/docs/community-benchmarks.md +++ b/docs/community-benchmarks.md @@ -30,7 +30,7 @@ tensorfold benchmark TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit - # Attach to a server you started first, in another terminal. Give the public checkpoint ID it serves, # and its --name as --model-id. tensorfold serve example-org/public-checkpoint --name local-model -tensorfold benchmark example-org/public-checkpoint --server http://127.0.0.1:8080 --model-id local-model --backend mlx +tensorfold benchmark example-org/public-checkpoint --server http://127.0.0.1:8888 --model-id local-model --backend mlx ``` Attached-server results do not claim the client's hardware is the server's hardware. They carry unknown hardware, versions and checkpoint revisions and remain unranked. An attached run records the backend, model and context you declare, unchecked against the server. The runner does not clear or restart an attached server's caches. 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 3c2871e4db..4b9359bcde 100644 --- a/docs/living-weights.md +++ b/docs/living-weights.md @@ -69,13 +69,13 @@ there is no key, when the key file is missing, empty or readable by other users space in it. Every request then sends the key, as the commands below do; [API keys](api.md#api-keys) has the rest. `--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 +memory, though the weights keep what they learned either way. The server listens on `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 "Authorization: Bearer $(cat ~/models/slide-key)" \ -H 'Content-Type: application/json' \ -d '{"text": "I like blue."}' @@ -93,7 +93,7 @@ few seconds on a Mac. Use the normal chat API, from any client: ```sh -curl http://127.0.0.1:8080/v1/chat/completions -H "Authorization: Bearer $(cat ~/models/slide-key)" \ +curl http://127.0.0.1:8888/v1/chat/completions -H "Authorization: Bearer $(cat ~/models/slide-key)" \ -H 'Content-Type: application/json' -d '{"messages": [{"role": "user", "content": "What colour do I like?"}]}' ``` @@ -125,7 +125,7 @@ nothing behind. An earlier fact's neuron that the new lesson would disturb is re - A file: send its contents as the text, for example: ```sh - jq -Rs '{text: .}' notes.md | curl -N http://127.0.0.1:8080/v1/slide/learn \ + jq -Rs '{text: .}' notes.md | curl -N http://127.0.0.1:8888/v1/slide/learn \ -H "Authorization: Bearer $(cat ~/models/slide-key)" -H 'Content-Type: application/json' -d @- ``` @@ -153,7 +153,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. @@ -167,7 +167,7 @@ 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. The page asks for the server's key once and keeps it for that browser tab. -- `curl -H "Authorization: Bearer $(cat ~/models/slide-key)" http://127.0.0.1:8080/v1/slide/graph` +- `curl -H "Authorization: Bearer $(cat ~/models/slide-key)" 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. - The same with `-X DELETE` clears the graph. The weights keep what they learned. To take one fact out of the weights, @@ -176,7 +176,7 @@ data. The page asks for the server's key once and keeps it for that browser tab. ## Delete one fact ```sh -curl -N -X DELETE http://127.0.0.1:8080/v1/slide/fact/2 \ +curl -N -X DELETE http://127.0.0.1:8888/v1/slide/fact/2 \ -H "Authorization: Bearer $(cat ~/models/slide-key)" \ -H 'Content-Type: application/json' -d '{"reason": "my sister changed her name"}' ``` diff --git a/docs/recipes/README.md b/docs/recipes/README.md index f3e82f73b4..9a9bb5b918 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -63,7 +63,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/glm-sampling-target-rows.md b/docs/recipes/glm-sampling-target-rows.md index 7e2e61da45..f69d9d3bef 100644 --- a/docs/recipes/glm-sampling-target-rows.md +++ b/docs/recipes/glm-sampling-target-rows.md @@ -26,8 +26,8 @@ Start the newly built native server using the existing admitted configuration an Replace `local-model` and the URL below with the values exposed by that server. ```sh -python3 tools/zig/check_glm_sampling.py http://127.0.0.1:8080 local-model --output glm-sampling-receipt.json -python3 tools/bench_concurrent.py http://127.0.0.1:8080 local-model --tokens 64 --reps 2 --temperatures 0,1 --levels 1,4 --alone --serial --output glm-concurrent-receipt.json +python3 tools/zig/check_glm_sampling.py http://127.0.0.1:8888 local-model --output glm-sampling-receipt.json +python3 tools/bench_concurrent.py http://127.0.0.1:8888 local-model --tokens 64 --reps 2 --temperatures 0,1 --levels 1,4 --alone --serial --output glm-concurrent-receipt.json ``` Require positive restored-token counts, native token hashes, drafted telemetry greater than zero, equality between 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 81c7b3610c..3bb86e0934 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -65,7 +65,7 @@ body = { 'max_tokens': 128, 'stream': False, } -request = Request('http://127.0.0.1:8080/v1/chat/completions', +request = Request('http://127.0.0.1:8888/v1/chat/completions', data=json.dumps(body).encode(), headers={'Content-Type': 'application/json'}) with urlopen(request) as response: diff --git a/tools/bench_concurrent.py b/tools/bench_concurrent.py index 583c7a644f..55c5be6be2 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 f29a5a323c..624bbb0631 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 684d6c3cc9..59d2bd6358 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -25,6 +25,7 @@ const kernels = [_]Kernel{ .{ .name = "scan_rows", .flags = &.{ "-O3", "--fmad=false" } }, // nemotron_h/cuda/mamba.py .{ .name = "nemotron_ops", .flags = &.{"-O3"} }, // ours: Nemotron's serial feed, routed plan and rests .{ .name = "affine4_pack", .flags = &.{"-O3"} }, // ours: the MLX affine-4 repack into qlinear's tiles + .{ .name = "affine8", .flags = &.{"-O3"} }, // ours: MLX affine 2 to 8 bits into 4- and 8-bit tiles, the 8-bit lane_gemv .{ .name = "fp8_experts", .flags = &.{"-O3"} }, // fp8/experts.cu's device code, tensorfold_fp8_experts_v6 .{ .name = "exlr8", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 lane projections, the Metal kernels' order .{ .name = "exlr8_dense", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 dense matrices by row, the Metal order @@ -33,6 +34,9 @@ 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 = "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 }, @@ -88,7 +92,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, radio: *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, radio: *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); @@ -102,7 +106,11 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin nemotron.addImport("nemotron_draft_ids", draft_ids); const radio = radioModule(b, target, optimize, cuda, core, lanes); nemotron.addImport("radio", radio); - return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat, .radio = radio }; + 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, .radio = radio }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -134,7 +142,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu cli.addImport("core", mods.core); cli.addImport("lanes", mods.lanes); cli.addImport("nemotron", mods.nemotron); - const native = engines(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat); + const native = engines(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, 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 })); @@ -158,7 +166,12 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu b.installArtifact(runner_exe); const controls = if (nvcc) |tool| fatbin(b, tool, version.?, .{ .name = "contract_controls", .flags = &.{"-O3"} }, sms) else null; @import("kernel_contract.zig").cuda(b, runner_exe, controls); - nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, build_options, true).root_module.strip = strip; + _ = nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, 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 GLM-5.3-Flash CUDA family over the runtime, the core, the lane core, canon's shares and core/heat.zig. @@ -184,7 +197,7 @@ fn radioModule(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu } /// 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, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, heat: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { +fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, heat: *std.Build.Module, qwen35: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), @@ -199,6 +212,7 @@ fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti .{ .name = "lfm2", .module = lfm2Module(b, target, optimize, cuda, core) }, .{ .name = "core", .module = core }, .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) }, + .{ .name = "qwen35", .module = qwen35 }, }, }); return .{ .api = api, .engines = mod }; @@ -210,8 +224,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, heat: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile { - const m = engines(b, target, optimize, cuda, core, lanes, nemotron, heat); +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, heat: *std.Build.Module, 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, core, lanes, nemotron, heat, 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(.{ @@ -227,6 +241,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; @@ -240,7 +256,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.core, mods.lanes, mods.nemotron, mods.heat); + const native_modules = engines(b, host, .debug, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, 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 }); @@ -268,7 +284,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, lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, 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, lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); // the GLM-5.3-Flash CUDA family's host side const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm })).step); @@ -294,6 +310,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); + // with cl.exe as host compiler, CUDA 13's headers need MSVC's conforming preprocessor, and M_SQRT2 its math define + if (b.graph.host.result.os.tag == .windows) run.addArgs(&.{ "-Xcompiler=/Zc:preprocessor", "-D_USE_MATH_DEFINES" }); run.addArgs(k.flags); var it = std.mem.tokenizeScalar(u8, sms, ','); while (it.next()) |sm| { @@ -317,7 +335,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.heat, mods.tokenizer, build_options, false); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.qwen35, mods.tokenizer, build_options, false, null); } /// SASS targets every release fatbin holds: Ampere (sm_80, sm_86), RTX 50/PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/build/sys.zig b/zig/build/sys.zig new file mode 100644 index 0000000000..1c1309af78 --- /dev/null +++ b/zig/build/sys.zig @@ -0,0 +1,40 @@ +//! `sys` (zig/src/platform/sys.zig) for every module the build compiles, by one walk of the step graph once declared. + +const std = @import("std"); + +pub fn wire(b: *std.Build) void { + var seen_steps: std.AutoHashMapUnmanaged(*std.Build.Step, void) = .empty; + var seen_mods: std.AutoHashMapUnmanaged(*std.Build.Module, void) = .empty; + var per_target: std.StringHashMapUnmanaged(*std.Build.Module) = .empty; + var w: Walk = .{ .b = b, .seen_steps = &seen_steps, .seen_mods = &seen_mods, .per_target = &per_target }; + for (b.top_level_steps.values()) |top| w.step(&top.step); +} + +const Walk = struct { + b: *std.Build, + seen_steps: *std.AutoHashMapUnmanaged(*std.Build.Step, void), + seen_mods: *std.AutoHashMapUnmanaged(*std.Build.Module, void), + per_target: *std.StringHashMapUnmanaged(*std.Build.Module), + + fn step(w: *Walk, s: *std.Build.Step) void { + if ((w.seen_steps.getOrPut(w.b.allocator, s) catch @panic("OOM")).found_existing) return; + if (s.cast(std.Build.Step.Compile)) |compile| w.module(compile.root_module, compile.root_module.resolved_target); + for (s.dependencies.items) |dep| w.step(dep); + } + + fn module(w: *Walk, m: *std.Build.Module, inherited: ?std.Build.ResolvedTarget) void { + if ((w.seen_mods.getOrPut(w.b.allocator, m) catch @panic("OOM")).found_existing) return; + const target = m.resolved_target orelse inherited orelse w.b.graph.host; + // only TensorFold's own sources see it; the sys module itself and generated modules do not + if (m.root_source_file) |root| if (root == .src_path and !m.import_table.contains("sys") and !std.mem.endsWith(u8, root.src_path.sub_path, "platform/sys.zig")) + m.addImport("sys", w.sys(target)); + for (m.import_table.values()) |dep| w.module(dep, target); + } + + fn sys(w: *Walk, target: std.Build.ResolvedTarget) *std.Build.Module { + const key = target.result.zigTriple(w.b.allocator) catch @panic("OOM"); + const slot = w.per_target.getOrPut(w.b.allocator, key) catch @panic("OOM"); + if (!slot.found_existing) slot.value_ptr.* = w.b.createModule(.{ .root_source_file = w.b.path("zig/src/platform/sys.zig"), .target = target, .link_libc = true }); + return slot.value_ptr.*; + } +}; diff --git a/zig/kernels/cuda/affine8.cu b/zig/kernels/cuda/affine8.cu new file mode 100644 index 0000000000..8598cdb43e --- /dev/null +++ b/zig/kernels/cuda/affine8.cu @@ -0,0 +1,348 @@ +// Ours: MLX affine codes of 2-4 bits on lane_gemv's 4-bit tiles, 5-8 bits on 8-bit tiles; no code is rounded. + +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace tf_affine8 { + +using namespace qmm_frag; + +// Code i of an MLX row packed `bits` wide, low code first, in little-endian words (a code may cross a word). +__device__ __forceinline__ uint32_t code_at(const uint32_t* __restrict__ row, long long i, int bits) { + const long long bit = i * bits; + const int sh = static_cast(bit & 31); + const long long wd = bit >> 5; + uint32_t v = row[wd] >> sh; + if (sh + bits > 32) v |= row[wd + 1] << (32 - sh); + return v & ((1u << bits) - 1u); +} + +// Bytes at s and s + 8 of w as a bf16 pair (the first low): exact, as every byte is a bf16 integer. +__device__ __forceinline__ uint32_t byte_pair(uint32_t w, int s) { + const float a = __uint_as_float(0x4B000000u | ((w >> s) & 0xFFu)) - 8388608.0f; + const float b = __uint_as_float(0x4B000000u | ((w >> (s + 8)) & 0xFFu)) - 8388608.0f; + const __nv_bfloat162 h = __floats2bfloat162_rn(a, b); + return *reinterpret_cast(&h); +} + +struct Part { + const uint32_t* w; + const __nv_bfloat16* scales; + const __nv_bfloat16* biases; + void* out; + int n, npad, sk, tiles; +}; + +// lane_gemv's tile with a byte a weight: 16 rows, BN columns over four warps, one 64-group a stage. +template +struct Tile { + static constexpr int THREADS = 128; + static constexpr int NT = BN / 4 / 8; + static constexpr int ROW = GS * 2; + static constexpr int CHUNKS = ROW / 16; + static constexpr int X = 16 * ROW; + static constexpr int W = BN * GS; + static constexpr int S = BN * 2; + static constexpr int XS = 16 * 4; + static constexpr int STAGE = X + W + 2 * S + XS; + static constexpr int SMEM = STAGES * STAGE; +}; + +// lane_gemv's gemv_items on 8-bit tiles [tile][group][8][32][GS / 16]: acc = x_sum * b + (d * s + acc) a group. +template +__device__ __forceinline__ void items(const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const Part& P, int M, int K, + int ldx, int count, Map map, float (&acc)[Tile::NT][4], After after) { + using T = Tile; + constexpr int NT = T::NT; + constexpr int KW = GS / 16; // words a lane a group: one per k16 step + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, wn = tid >> 5; + const int KG = K / GS; + auto stage = [&](int s) { return buf + s * T::STAGE; }; + const int rows = min(16, M); + constexpr int TILE_BYTES = 64 * GS; + constexpr int XC = (16 * T::CHUNKS + T::THREADS - 1) / T::THREADS, WC = (T::W / 16 + T::THREADS - 1) / T::THREADS; + constexpr int SC = (2 * T::S / 16 + T::THREADS - 1) / T::THREADS; + const __nv_bfloat16* xsrc[XC]; + size_t woff[WC]; + int soff[SC], xdst[XC], wdst[WC], sdst[SC]; + bool xok[XC], wok[WC], sok[SC], sbias[SC]; +#pragma unroll + for (int j = 0; j < XC; ++j) { + const int c = tid + j * T::THREADS, r = c / T::CHUNKS, ch = c % T::CHUNKS; + xok[j] = c < 16 * T::CHUNKS && r < rows; + xsrc[j] = x + static_cast(min(r, rows - 1)) * ldx + ch * 8; + xdst[j] = r * T::ROW + swz(r, ch) * 16; + } +#pragma unroll + for (int j = 0; j < WC; ++j) { + const int c = tid + j * T::THREADS, t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + wok[j] = c < T::W / 16; + woff[j] = static_cast(t) * KG * TILE_BYTES + off * 16; + wdst[j] = T::X + c * 16; + } +#pragma unroll + for (int j = 0; j < SC; ++j) { + const int c = tid + j * T::THREADS, which = c / (T::S / 16), off = c % (T::S / 16); + sok[j] = c < 2 * (T::S / 16); + sbias[j] = which; + soff[j] = off * 8; + sdst[j] = T::X + T::W + which * T::S + off * 16; + } + const bool xsok = tid < rows; + const float* xssrc = xs + static_cast(min(tid, rows - 1)) * KG; + auto load_x = [&](int s, int i) { + const int g = map(i).x; + unsigned char* p = stage(s); +#pragma unroll + for (int j = 0; j < XC; ++j) + if (xok[j]) cp16(p + xdst[j], xsrc[j] + g * GS); + if (xsok) cp4(p + T::X + T::W + 2 * T::S + tid * 4, xssrc + g); + }; + auto load_w = [&](int s, int i) { + const int2 at = map(i); + const int g = at.x, n0 = at.y; + unsigned char* p = stage(s); + const unsigned char* wb = reinterpret_cast(P.w) + static_cast(n0 / 64) * KG * TILE_BYTES + (n0 % 64) * GS; +#pragma unroll + for (int j = 0; j < WC; ++j) + if (wok[j]) cp16(p + wdst[j], wb + woff[j] + static_cast(g) * TILE_BYTES); +#pragma unroll + for (int j = 0; j < SC; ++j) + if (sok[j]) cp16(p + sdst[j], (sbias[j] ? P.biases : P.scales) + n0 + soff[j] + static_cast(g) * P.npad); + }; + + for (int c = tid; c < 16 * T::CHUNKS; c += T::THREADS) + if (c / T::CHUNKS >= rows) +#pragma unroll + for (int s = 0; s < STAGES; ++s) + *reinterpret_cast(stage(s) + c / T::CHUNKS * T::ROW + c % T::CHUNKS * 16) = uint4{}; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < count) load_w(s, s); + commit(); + } + grid_wait(); +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < count) load_x(s, s); + commit(); + } + grid_launch(); + for (int it = 0; it < count; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < count) { + load_x(next % STAGES, next); + load_w(next % STAGES, next); + } + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W); + const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S); + uint32_t words[NT][KW]; +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int v = 0; v < KW; ++v) words[j][v] = pw[((wn * NT + j) * 32 + lane) * KW + v]; + float d[NT][4]; +#pragma unroll + for (int kt = 0; kt < GS / 16; ++kt) { + uint32_t a[4]; + const int r = (lane & 7) + ((lane >> 3) & 1) * 8, ch = kt * 2 + (lane >> 4); + ldmatrix4(a, p + r * T::ROW + swz(r, ch) * 16); +#pragma unroll + for (int j = 0; j < NT; ++j) { + const uint32_t b0 = byte_pair(words[j][kt], 0); + const uint32_t b1 = byte_pair(words[j][kt], 16); + if (kt == 0) mma0(d[j], a, b0, b1); + else mma(d[j], a, b0, b1); + } + } + const int row = lane >> 2; + const float xv[2] = {px[row], px[row + 8]}; +#pragma unroll + for (int j = 0; j < NT; ++j) { + const int col = wn * (BN / 4) + j * 8 + (lane & 3) * 2; + const __nv_bfloat162 s2 = *reinterpret_cast(ps + col); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + BN + col); + const float sv[2] = {__low2float(s2), __high2float(s2)}; + const float bv[2] = {__low2float(b2), __high2float(b2)}; +#pragma unroll + for (int e = 0; e < 4; ++e) + acc[j][e] = __fmaf_rn(xv[e >> 1], bv[e & 1], __fmaf_rn(d[j][e], sv[e & 1], acc[j][e])); + } + after(it); + } +} + +// Rows below M of the BN columns at n0 as bf16, from fp32 pairs got(j, h, row, col) in the accumulator's layout. +template +__device__ __forceinline__ void store(const Part& P, int M, int n0, Got got) { + const int lane = threadIdx.x & 31, wn = threadIdx.x >> 5, row = lane >> 2, N = P.n; +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = n0 + wn * (BN / 4) + j * 8 + (lane & 3) * 2, r = row + h * 8; + if (r >= M) continue; + const float2 v = got(j, h, r, col - n0); + auto* dst = reinterpret_cast<__nv_bfloat16*>(P.out) + static_cast(r) * N + col; + if (col + 1 < N && (N & 1) == 0) { + *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v.x, v.y); + } else { + if (col < N) dst[0] = __float2bfloat16_rn(v.x); + if (col + 1 < N) dst[1] = __float2bfloat16_rn(v.y); + } + } +} + +// lane_gemv's gemv_kernel: CTAs loop over column tiles, each tile's K slices summed in slice order. +template +__device__ __forceinline__ void gemv(const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const Part& P, int M, int K, int ldx) { + constexpr int NT = Tile::NT; + const int KG = K / GS, per = KG / P.sk; + const int mine = P.tiles > (int)blockIdx.x ? (P.tiles - 1 - (int)blockIdx.x) / (int)gridDim.x + 1 : 0; + auto tile_of_item = [&](int i) { return (int)blockIdx.x + (i / KG) * (int)gridDim.x; }; + float acc[NT][4], tot[NT][4]; +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[j][e] = tot[j][e] = 0.0f; + items(x, xs, P, M, K, ldx, mine * KG, [&](int i) { return make_int2(i % KG, tile_of_item(i) * BN); }, acc, [&](int it) { + const int g = it % KG; + if ((g + 1) % per != 0) return; + const bool first = g + 1 == per; +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) { + tot[j][e] = first ? acc[j][e] : tot[j][e] + acc[j][e]; + acc[j][e] = 0.0f; + } + if (g != KG - 1) return; + store(P, M, tile_of_item(it) * BN, [&](int j, int h, int, int) { return make_float2(tot[j][2 * h], tot[j][2 * h + 1]); }); + }); +} + +// lane_gemv's gemv_split: a CTA a tile's K slice, the tile's last CTA summing the slices in their order from the first. +template +__device__ __forceinline__ void split(const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const Part& P, int M, int K, int ldx, + float* __restrict__ work, unsigned* __restrict__ tickets) { + constexpr int NT = Tile::NT; + __shared__ unsigned ticket; + const int tid = threadIdx.x, lane = tid & 31, wn = tid >> 5; + const int per = K / GS / P.sk; + const int tile = (int)blockIdx.x / P.sk, slice = (int)blockIdx.x % P.sk, g0 = slice * per; + float acc[NT][4]; +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[j][e] = 0.0f; + items(x, xs, P, M, K, ldx, per, [&](int i) { return make_int2(g0 + i, tile * BN); }, acc, [](int) {}); + const int row = lane >> 2; + float* part = work + (static_cast(tile) * P.sk + slice) * 16 * BN; +#pragma unroll + for (int j = 0; j < NT; ++j) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int c = wn * (BN / 4) + j * 8 + (lane & 3) * 2, r = row + h * 8; + if (r >= M) continue; + part[r * BN + c] = acc[j][2 * h]; + part[r * BN + c + 1] = acc[j][2 * h + 1]; + } + __threadfence(); + __syncthreads(); + if (tid == 0) ticket = atomicAdd(tickets + tile, 1u); + __syncthreads(); + if (ticket != static_cast(P.sk - 1)) return; + __threadfence(); + const float* base = work + static_cast(tile) * P.sk * 16 * BN; + store(P, M, tile * BN, [&](int, int, int r, int c) { + float t0 = __ldcg(base + r * BN + c), t1 = __ldcg(base + r * BN + c + 1); + for (int s = 1; s < P.sk; ++s) { + t0 = t0 + __ldcg(base + (s * 16 + r) * BN + c); + t1 = t1 + __ldcg(base + (s * 16 + r) * BN + c + 1); + } + return make_float2(t0, t1); + }); + if (tid == 0) tickets[tile] = 0; +} + +} // namespace tf_affine8 + +// Shared bytes: 6 stages for gemv, 4 for split (cuda/qlinear.zig's affine8 constants). +static_assert(tf_affine8::Tile<64, 64, 6>::SMEM == 38784, "gemv8 smem"); +static_assert(tf_affine8::Tile<64, 64, 4>::SMEM == 25856, "split8 smem"); + +extern "C" __global__ void __launch_bounds__(128) tf_affine8_gemv(const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, + const __grid_constant__ tf_affine8::Part P, int M, int K, int ldx) { + tf_affine8::gemv<64, 64, 6>(x, xs, P, M, K, ldx); +} + +extern "C" __global__ void __launch_bounds__(128) tf_affine8_split(const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, + const __grid_constant__ tf_affine8::Part P, int M, int K, int ldx, + float* __restrict__ work, unsigned* __restrict__ tickets) { + tf_affine8::split<64, 64, 4>(x, xs, P, M, K, ldx, work, tickets); +} + +// MLX words of 2 to 4 bits -> lane_gemv's 4-bit tiles [npad/64][kg][8][32][2], tf_pack_dense's layout. +extern "C" __global__ void tf_affine8_pack4(const uint32_t* __restrict__ w, int bits, int n, int kw, int kg, uint32_t* __restrict__ out, + long long total) { + const long long idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) return; + const int v = idx & 1, lane = (idx >> 1) & 31, j = (idx >> 6) & 7; + const long long tg = idx >> 9; + const int g = static_cast(tg % kg); + const long long t = tg / kg; + const long long row = t * 64 + j * 8 + (lane >> 2); + const int c = lane & 3; + const int offs[8] = {0, 8, 16, 24, 1, 9, 17, 25}; + uint32_t word = 0; + if (row < n) { + const uint32_t* src = w + row * kw; +#pragma unroll + for (int p = 0; p < 8; ++p) word |= tf_affine8::code_at(src, g * 64 + 32 * v + 2 * c + offs[p], bits) << (4 * p); + } + out[idx] = word; +} + +// MLX words of 5 to 8 bits -> 8-bit tiles [npad/64][kg][8][32][4], word kt holding inputs k0+2c, +1, +8, +9. +extern "C" __global__ void tf_affine8_pack8(const uint32_t* __restrict__ w, int bits, int n, int kw, int kg, uint32_t* __restrict__ out, + long long total) { + const long long idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) return; + const int kt = idx & 3, lane = (idx >> 2) & 31, j = (idx >> 7) & 7; + const long long tg = idx >> 10; + const int g = static_cast(tg % kg); + const long long t = tg / kg; + const long long row = t * 64 + j * 8 + (lane >> 2); + const int k0 = g * 64 + 16 * kt + 2 * (lane & 3); + const int offs[4] = {0, 1, 8, 9}; + uint32_t word = 0; + if (row < n) { + const uint32_t* src = w + row * kw; +#pragma unroll + for (int p = 0; p < 4; ++p) word |= tf_affine8::code_at(src, k0 + offs[p], bits) << (8 * p); + } + out[idx] = word; +} + +// MLX rows of `bits` codes -> MLX rows of `wide` (4 or 8) bit codes, the same values: total output words. +extern "C" __global__ void tf_affine8_widen(const uint32_t* __restrict__ w, int bits, int k, uint32_t* __restrict__ out, int wide, + long long total) { + const long long idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) return; + const int per = 32 / wide; // codes a word out + const long long row_words = static_cast(k) / per; + const long long row = idx / row_words, at = idx % row_words; + const uint32_t* src = w + row * (static_cast(k) * bits / 32); + uint32_t word = 0; + for (int c = 0; c < per; ++c) word |= tf_affine8::code_at(src, at * per + c, bits) << (c * wide); + out[idx] = word; +} 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 88e32e7be9..8db3479278 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -244,14 +244,14 @@ vision_kernel(const __nv_bfloat16* __restrict__ qkv, __nv_bfloat16* __restrict__ } // 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); // Head dim 64, eight warps, four query heads a block, eight staging slots (LFM2's 4 heads a KV head, causal or trees). -template __global__ void tf_prefill_attention::tree_kernel<64, 8, 4, 8>(const __nv_bfloat16*, - const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, const int*, int, int, int, int, int, float); +template __global__ void tf_prefill_attention::tree_kernel<64, 8, 4, 8>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, const int* __restrict__, int, int, int, int, int, float); // Head dim 72, eight warps of one head, eight staging slots (SigLIP2's tower, 72 dims padded to 80 inside). -template __global__ void tf_prefill_attention::full_kernel<72, 8, 8>(const __nv_bfloat16*, __nv_bfloat16*, int, int, +template __global__ void tf_prefill_attention::full_kernel<72, 8, 8>(const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, float); // Head dim 80, eight warps, two blocks an SM (RADIO's tower, 16 heads of 80): full_kernel<80>'s bits, 1.35x faster. -template __global__ void tf_prefill_attention::vision_kernel<80, 8>(const __nv_bfloat16*, __nv_bfloat16*, int, int, +template __global__ void tf_prefill_attention::vision_kernel<80, 8>(const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, 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..d69759b8dc --- /dev/null +++ b/zig/kernels/cuda/qwen35.cu @@ -0,0 +1,522 @@ +// 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); + } +} + +// MLX affine 8-bit token rows (a mixed checkpoint's, or widened to 8 at load): out = bf16(s * q + b), D wide. +extern "C" __global__ void __launch_bounds__(256) tf_q35_embed8(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 d = threadIdx.x; d < D; d += blockDim.x) { + const size_t g = id * (D / 64) + d / 64; + out[(size_t)blockIdx.x * D + d] = __float2bfloat16_rn(bf(S[g]) * (float)W[id * D + d] + bf(B[g])); + } +} + +// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)); D a multiple of 256 up to 8192; a 256-thread block a row. +extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, + const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D, + const int* __restrict__ P) { + constexpr int MAXE = 32; + __shared__ float red[8]; + if (past_rows(P, blockIdx.x)) return; + const size_t row = blockIdx.x; + const int t = threadIdx.x; + const int E = D / 256; + float hv[MAXE]; + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < MAXE; i++) { + if (i >= E) break; + const size_t at = row * D + t * E + i; + float h = bf(H[at]); + if (R != nullptr) { + const __nv_bfloat16 hb = __float2bfloat16_rn(h + bf(R[at])); + H[at] = hb; + h = bf(hb); + } + hv[i] = h; + ss = fmaf(h, h, ss); + } + ss = warp_sum(ss); + if ((t & 31) == 0) red[t >> 5] = ss; + __syncthreads(); + float total = 0.0f; + for (int i = 0; i < 8; i++) total += red[i]; + const float inv = rsqrtf(total / (float)D + eps); +#pragma unroll + for (int i = 0; i < MAXE; i++) { + if (i >= E) break; + const int d = t * E + i; + X[row * D + d] = __float2bfloat16_rn(bf(W[d]) * (hv[i] * inv)); + } +} + +// The 64-group sums of bf16 rows D wide (the decode projections' input sums). Grid (rows), a thread a group. +extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, float* __restrict__ XS, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const size_t row = blockIdx.x; + for (int g = threadIdx.x; g < D / 64; g += blockDim.x) { + const __nv_bfloat16* x = X + row * D + g * 64; + float s = 0.0f; + for (int j = 0; j < 64; j++) s += bf(x[j]); + XS[row * (D / 64) + g] = s; + } +} + +// Conv, silu, q/k norms, the linear heads' decay and write: a warp a (row, head) over q, k, v; taps from the parents. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const int* __restrict__ rowpar, int layer, + const __nv_bfloat16* __restrict__ CW, + const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, + const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, + __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ K, + __nv_bfloat16* __restrict__ V, float* __restrict__ G, __nv_bfloat16* __restrict__ BETA, + const int* __restrict__ P) { + const int lane = threadIdx.x; + const int head = blockIdx.y; + const int w = blockIdx.z; + if (past_rows(P, w)) return; + const int NK = geo.nk, NV = geo.nv, C = geo.conv(); + const Seg sg = segs[rowseg[w]]; + const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * geo.conv_rows(); + const bool isq = head < NK, isk = !isq && head < 2 * NK; + const int c0 = head * DK; // q, k and v channels are consecutive heads of DK + int src[TAPS]; // tap t reads the row TAPS - 1 - t tokens back: a forward row, or -1 - (tail row) +#pragma unroll + for (int tap = 0; tap < TAPS; tap++) src[tap] = back(rowpar, w, TAPS - 1 - tap); + constexpr int PER = DK / 32; + float vals[PER]; +#pragma unroll + for (int j = 0; j < PER; j++) { + const int c = c0 + lane * PER + j; + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; tap++) { + const int at = src[tap]; + const float xv = at < 0 ? bf(CS[(size_t)(-1 - at) * C + c]) : bf(QKV[(size_t)at * C + c]); + acc += bf(CW[c * TAPS + tap]) * xv; + } + const float conv = rbf(acc); + const float sig = rbf(1.0f / (1.0f + expf(-conv))); + vals[j] = rbf(conv * sig); + } + if (isq || isk) { + float ss = 0.0f; +#pragma unroll + for (int j = 0; j < PER; j++) ss += vals[j] * vals[j]; + ss = warp_sum(ss); + const float inv = rsqrtf(ss / (float)DK + (1e-6f / (float)DK)); + const float scale = isq ? rbf(1.0f / (float)DK) : rbf(rsqrtf((float)DK)); + __nv_bfloat16* out = isq ? Q + ((size_t)w * NK + head) * DK : K + ((size_t)w * NK + head - NK) * DK; +#pragma unroll + for (int j = 0; j < PER; j++) out[lane * PER + j] = __float2bfloat16_rn(scale * rbf(vals[j] * inv)); + } else { + const int hv = head - 2 * NK; +#pragma unroll + for (int j = 0; j < PER; j++) V[((size_t)w * NV + hv) * DK + lane * PER + j] = __float2bfloat16_rn(vals[j]); + if (lane == 0) { + const float s = rbf(bf(Ain[w * NV + hv]) + bf(DT[hv])); + const float sp = rbf(fmaxf(s, 0.0f) + logf(1.0f + expf(-fabsf(s)))); + G[w * NV + hv] = expf(-expf(ALOG[hv]) * sp); + BETA[w * NV + hv] = __float2bfloat16_rn(1.0f / (1.0f + expf(-bf(Bin[w * NV + hv])))); + } + } +} + +// A segment's conv tail: a recorded window logs each row's for keep and leaves the stream's; others update it in place. +extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, + const int* __restrict__ rowpar, int layer, const int* __restrict__ P) { + const int C = geo.conv(); + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= C || past_windows(P, blockIdx.y)) return; + const Seg sg = segs[blockIdx.y]; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * geo.conv_rows(); + if (recorded(P)) { + for (int j = 0; j < sg.rows; j++) { + __nv_bfloat16* out = log_row(geo, sg, layer, j).conv; + for (int r = 0; r < TAPS - 1; r++) { + const int at = back(rowpar, sg.row0 + j, TAPS - 2 - r); + out[(size_t)r * C + c] = at < 0 ? CS[(size_t)(-1 - at) * C + c] : QKV[(size_t)at * C + c]; + } + } + if (sg.rows != 1) return; + } + // ascending: slot r takes [r + rows], which is either a later state slot (not yet written) or a row + for (int r = 0; r < TAPS - 1; r++) { + const int at = r + sg.rows; + CS[(size_t)r * C + c] = at < TAPS - 1 ? CS[(size_t)at * C + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]; + } +} + +// Gated delta rule, each row's state its parent's: a chain keeps it in registers, a branch rebuilds it from the log. +extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, + const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, + const int* __restrict__ rowpar, int layer, __nv_bfloat16* __restrict__ y, + const int* __restrict__ P) { + const int NK = geo.nk, NV = geo.nv; + if (past_windows(P, blockIdx.z / NV)) return; + const int hv = blockIdx.z % NV; + const int hk = hv / (NV / NK); + const Seg sg = segs[blockIdx.z / NV]; + const int dk = threadIdx.x; + const int dv = blockIdx.y * 4 + threadIdx.y; + constexpr int PT = DK / 32; + const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; + const bool record = recorded(P); + float state[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + int prev = -1; // the row the registers' state follows + for (int r = 0; r < sg.rows; r++) { + const int row = sg.row0 + r; + const int parent = rowpar[row]; + if (parent != prev) { + // a branch: its parent's state from the stream's and the ancestors' logged updates, oldest first + int anc[MAXW]; + int n = 0; + for (int a = parent; a >= 0; a = rowpar[a]) anc[n++] = a; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int j = n - 1; j >= 0; j--) { + const int a = anc[j]; + const float ga = g[a * NV + hv], da = log_row(geo, sg, layer, a - sg.row0).delta[hv * DK + dv]; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * ga; + state[i] = fmaf(bf(k[((size_t)a * NK + hk) * DK + PT * dk + i]), da, state[i]); + } + } + } + prev = row; + float kc[PT], qc[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) { + kc[i] = bf(k[((size_t)row * NK + hk) * DK + PT * dk + i]); + qc[i] = bf(q[((size_t)row * NK + hk) * DK + PT * dk + i]); + } + const float vc = bf(v[((size_t)row * NV + hv) * DK + dv]); + const float gc = g[row * NV + hv]; + const float bc = bf(beta[row * NV + hv]); + float kv_mem = 0.0f; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * gc; + kv_mem += state[i] * kc[i]; + } + kv_mem = warp_sum(kv_mem); + const float delta = (vc - kv_mem) * bc; + float out = 0.0f; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = fmaf(kc[i], delta, state[i]); + out += state[i] * qc[i]; + } + out = warp_sum(out); + if (dk == 0) y[((size_t)row * NV + hv) * DK + dv] = __float2bfloat16_rn(out); + if (record) { + const LogRow l = log_row(geo, sg, layer, r); + if (dk == 0) { + l.delta[hv * DK + dv] = delta; + if (dv == 0) l.g[hv] = gc; + } + if (dv == 0 && hv % (NV / NK) == 0) +#pragma unroll + for (int i = 0; i < PT; i++) l.k[hk * DK + PT * dk + i] = k[((size_t)row * NK + hk) * DK + PT * dk + i]; + } + } + if (record && sg.rows != 1) return; +#pragma unroll + for (int i = 0; i < PT; i++) st[i] = state[i]; +} + +// out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear value head; rows nv * DK wide. Grid (1, nv, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, + const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT, + const int* __restrict__ P) { + const int lane = threadIdx.x, hv = blockIdx.y; + if (past_rows(P, blockIdx.z)) return; + const size_t m = blockIdx.z; + const size_t W = (size_t)geo.nv * DK; + constexpr int PER = DK / 32; + float yv[PER], ss = 0.0f; +#pragma unroll + for (int j = 0; j < PER; j++) { + yv[j] = bf(Y[m * W + hv * DK + lane * PER + j]); + ss += yv[j] * yv[j]; + } + ss = warp_sum(ss); + const float inv = rsqrtf(ss / (float)DK + eps); +#pragma unroll + for (int j = 0; j < PER; j++) { + const int d = lane * PER + j; + const float x = rbf(bf(NW[d]) * (yv[j] * inv)); + const float zf = bf(Z[m * W + hv * DK + d]); + OUT[m * W + hv * DK + d] = __float2bfloat16_rn(zf / (1.0f + expf(-zf)) * x); + } +} + +// Per-head RMSNorm of attention rows, out packed (rows, heads, HD); grid (rows * heads), 64 threads of 4 values. +extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps, + const int* __restrict__ P) { + __shared__ float sums[2]; + const int head = blockIdx.x % heads; + if (past_rows(P, blockIdx.x / heads)) return; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + head * hstride; + float v[4], acc = 0.0f; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = bf(x[4 * t + i]); + acc += v[i] * v[i]; + } + acc = warp_sum(acc); + if ((t & 31) == 0) sums[t >> 5] = acc; + __syncthreads(); + const float inv = rsqrtf((sums[0] + sums[1]) / (float)HD + eps); +#pragma unroll + for (int i = 0; i < 4; i++) { + const int d = 4 * t + i; + OUT[(m * heads + head) * HD + d] = __float2bfloat16_rn(bf(W[d]) * (v[i] * inv)); + } +} + +// Text RoPE on the first 64 of HD: pairs (i, i + 32) turned by position * theta^(-i / 32). +__device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, int position) { + if (d >= 64) return x[d]; + const int i = d % 32; + const float angle = (float)position * powf(10000000.0f, -(float)i / 32.0f); + float sn, cs; + sincosf(angle, &sn, &cs); + const float a = bf(x[i]), b = bf(x[i + 32]); + return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); +} + +// A row's ancestors-and-self: its tree window's row bits, or with the top bit set a chain row's index in its window. +__device__ __forceinline__ int rows_seen(unsigned anc) { return (anc >> 31) ? (int)(anc & 0x7fffffffu) + 1 : __popc(anc); } + +// A row's chunk width from its own logical keys alone, as its one-row step's: at most MAXCH chunks, 128-key multiples. +__device__ __forceinline__ int chunk_keys(int keys) { return max(128, (keys + MAXCH * 128 - 1) / (MAXCH * 128) * 128); } + +// The window row holding a row's nth seen window key. +__device__ __forceinline__ int seen_row(unsigned anc, int n) { + if (anc >> 31) return n; + for (int j = 0; j < n; j++) anc &= anc - 1; + return __ffs(anc) - 1; +} + +// Queries rotated in place, keys rotated and values copied into the row's cache slot; grid (rows), HD threads. +extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, + const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, + const int* __restrict__ P) { + const int QH = geo.qh, KVH = geo.kvh; + const int d = threadIdx.x; + if (past_rows(P, blockIdx.x)) return; + const size_t m = blockIdx.x; + const Seg sg = segs[rowseg[m]]; + const int pos = sg.len + rows_seen(rowanc[m]) - 1; + const int slot = sg.len + (int)m - sg.row0; + __shared__ __nv_bfloat16 qs[MAXQ * HD]; + for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; + __syncthreads(); + for (int h = 0; h < QH; h++) Qn[(m * QH + h) * HD + d] = rotate(qs + h * HD, d, pos); + __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; + __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + for (int h = 0; h < KVH; h++) { + const size_t dst = ((size_t)h * sg.cap + slot) * HD + d; + keys[dst] = rotate(Kn + (m * KVH + h) * HD, d, pos); + values[dst] = Vr[(m * KVH + h) * HD + d]; + } +} + +// Causal attention in chunks of chunk_keys(the row's keys), a warp a query head; partials by (row, head, chunk). +extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, + const int* __restrict__ P, float scale, + float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { + constexpr int DPL = HD / 32; + const int QH = geo.qh, KVH = geo.kvh; + const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; + const int c = blockIdx.x, h = blockIdx.y; + if (c >= MAXCH || past_rows(P, blockIdx.z)) return; + const size_t m = blockIdx.z; + const int qh = h * (QH / KVH) + g; + const Seg sg = segs[rowseg[m]]; + const unsigned anc = rowanc[m]; + // chunks cut at logical positions (a window key's: the stream's length plus its ancestors), as a step's are + const int keys = sg.len + rows_seen(anc), ck = chunk_keys(keys); + const int k0 = c * ck, k1 = min(k0 + ck, keys); + const size_t slot = (m * QH + qh) * MAXCH + c; + if (k0 >= k1) { + if (lane == 0) { + PM[slot] = -INFINITY; + PL[slot] = 0.0f; + } + return; + } + const __nv_bfloat16* K = reinterpret_cast(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; + const __nv_bfloat16* V = K + (size_t)KVH * sg.cap * HD; + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[(m * QH + qh) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + for (int t = k0; t < k1; t++) { + const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); + const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const __nv_bfloat16* kb = reinterpret_cast(&kw); + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(kb[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const uint4 vw = *reinterpret_cast(V + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const __nv_bfloat16* vb = reinterpret_cast(&vw); +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); + mx = mn; + } + if (lane == 0) { + PM[slot] = mx; + PL[slot] = l; + } +#pragma unroll + for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; +} + +// A row's chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, const float* __restrict__ PM, const float* __restrict__ PL, + const float* __restrict__ PO, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, + const int* __restrict__ P, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int QH = geo.qh; + const int lane = threadIdx.x, qh = blockIdx.x; + if (past_rows(P, blockIdx.y)) return; + const size_t m = blockIdx.y; + const int keys = segs[rowseg[m]].len + rows_seen(rowanc[m]); + const int ck = chunk_keys(keys), nch = (keys + ck - 1) / ck; + const size_t base = (m * QH + qh) * MAXCH; + float mx = -INFINITY; + for (int c = 0; c < nch; c++) + if (PL[base + c] > 0.0f) mx = fmaxf(mx, PM[base + c]); + float lsum = 0.0f, acc[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) acc[i] = 0.0f; + for (int c = 0; c < nch; c++) { + if (!(PL[base + c] > 0.0f)) continue; + const float e = __expf(PM[base + c] - mx); + lsum = fmaf(PL[base + c], e, lsum); +#pragma unroll + for (int i = 0; i < DPL; i++) acc[i] = fmaf(PO[(base + c) * HD + lane * DPL + i], e, acc[i]); + } +#pragma unroll + for (int i = 0; i < DPL; i++) OUT[(m * QH + qh) * HD + lane * DPL + i] = __float2bfloat16_rn(acc[i] / lsum); +} + +// out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows; rows qh * HD wide. +extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, + __nv_bfloat16* __restrict__ OUT, int rows, const int* __restrict__ P) { + const size_t W = (size_t)geo.qh * HD; + if (P != nullptr) rows = P[2]; + const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= (size_t)rows * W) return; + const size_t r = i / W; + const int d = (int)(i % W), head = d / HD; + const float gate = bf(QRAW[r * 2 * W + head * 2 * HD + HD + d % HD]); + const float s = rbf(1.0f / (1.0f + expf(-gate))); + OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); +} + +// A tapped layer's residual bf16(h + r) into the segment's tap ring at the row's position; a 256-thread block a row. +extern "C" __global__ void __launch_bounds__(256) tf_q35_tap(const __nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, + const Seg* __restrict__ segs, const int* __restrict__ rowseg, int tap, int D, + const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const size_t m = blockIdx.x; + const Seg sg = segs[rowseg[m]]; + if (sg.tap == 0) return; + const long long pos = sg.len + (long long)m - sg.row0; + __nv_bfloat16* out = reinterpret_cast<__nv_bfloat16*>(sg.tap) + ((size_t)(pos % sg.tap_rows) * sg.taps + tap) * D; + for (int i = threadIdx.x; i < D; i += blockDim.x) out[i] = __float2bfloat16_rn(bf(H[m * D + i]) + bf(R[m * D + i])); +} + +// out = bf16(silu(gate) * up) over `total` values. +extern "C" __global__ void tf_q35_mlp_act(const __nv_bfloat16* __restrict__ GATE, const __nv_bfloat16* __restrict__ UP, + __nv_bfloat16* __restrict__ OUT, long long total, const int* __restrict__ P, int D) { + if (P != nullptr) total = (long long)P[2] * D; + const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + const float gf = bf(GATE[i]); + OUT[i] = __float2bfloat16_rn(gf / (1.0f + expf(-gf)) * bf(UP[i])); +} + +// Each row's greedy token: the largest logit, a tie to the lower id (the host draw's rule). Grid (rows), 1024 threads. +extern "C" __global__ void __launch_bounds__(1024) tf_q35_argmax(const __nv_bfloat16* __restrict__ L, int V, unsigned* __restrict__ TOK) { + __shared__ float bv[32]; + __shared__ int bi[32]; + const __nv_bfloat16* x = L + (size_t)blockIdx.x * V; + float best = -INFINITY; + int id = 0x7fffffff; + for (int i = threadIdx.x; i < V; i += blockDim.x) { + const float v = bf(x[i]); + if (v > best) { + best = v; + id = i; + } + } +#pragma unroll + for (int m = 16; m; m >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, best, m); + const int oi = __shfl_xor_sync(0xffffffffu, id, m); + if (ov > best || (ov == best && oi < id)) { + best = ov; + id = oi; + } + } + if ((threadIdx.x & 31) == 0) { + bv[threadIdx.x >> 5] = best; + bi[threadIdx.x >> 5] = id; + } + __syncthreads(); + if (threadIdx.x == 0) { + for (int w = 1; w < 32; w++) + if (bv[w] > best || (bv[w] == best && bi[w] < id)) { + best = bv[w]; + id = bi[w]; + } + TOK[blockIdx.x] = (unsigned)id; + } +} diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh new file mode 100644 index 0000000000..5df37e308e --- /dev/null +++ b/zig/kernels/cuda/qwen35.cuh @@ -0,0 +1,79 @@ +// Shared by qwen35.cu and qwen35_keep.cu: the widths (Geo), a stream's window (Seg), a row's log record, parent walks. + +#pragma once +#include +#include + +namespace tf_q35 { + +constexpr int DK = 128; // linear-attention key and value head width +constexpr int TAPS = 4; // causal convolution taps +constexpr int HD = 256; // attention head width +constexpr int MAXQ = 32; // attention query heads at most +constexpr int MAXW = 16; // a window's rows at most (the log's rows a stream) +constexpr int MAXCH = 64; // attention chunks a row at most (cuda_engine.zig's max_chunks) + +// The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads, linear layers. +struct Geo { + int hid, nk, nv, qh, kvh, nl; + __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees + __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence + __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } + // a row's record of a linear layer: decay, strength and write a value head, key a key head, the tail after it + __device__ size_t log_bytes() const { return (size_t)nv * 4 + (size_t)nv * DK * 4 + (size_t)nk * DK * 2 + conv_rows() * 2; } +}; + +// One stream's window in this forward: its state blocks, cache length before the window, and rows. +struct Seg { + unsigned long long rec; // float recurrences, one rec() block a linear layer + unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer + unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time + unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each + unsigned long long log; // the window's rows' records (MAXW rows of nl layers of log_bytes), for keep + int cap, len, row0, rows, tap_rows, taps; +}; + +// A window row's record of one linear layer (Seg.log). +struct LogRow { + float* g; // [nv] decay + float* delta; // [nv][DK] the write's strength along dv + __nv_bfloat16* k; // [nk][DK] the key + __nv_bfloat16* conv; // [TAPS - 1][conv()] the tail after the row +}; + +__device__ __forceinline__ LogRow log_row(const Geo& geo, const Seg& sg, int layer, int row) { + char* base = reinterpret_cast(sg.log) + ((size_t)row * geo.nl + layer) * geo.log_bytes(); + LogRow l; + l.g = reinterpret_cast(base); + l.delta = l.g + geo.nv; + l.k = reinterpret_cast<__nv_bfloat16*>(l.delta + (size_t)geo.nv * DK); + l.conv = l.k + (size_t)geo.nk * DK; + return l; +} + +// The rows of a window to keep, in order (row 0 first), passed by value. +struct Path { + int n; + int rows[MAXW]; +}; + +// The row `steps` tokens before forward row `w` through rowpar: a forward row, or -1 - t for the stream's tail row t. +__device__ __forceinline__ int back(const int* __restrict__ rowpar, int w, int steps) { + int a = w; + for (int s = 1; s <= steps; s++) { + a = rowpar[a]; + if (a < 0) return -1 - (TAPS - 2 - (steps - s)); + } + return a; +} + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q35 diff --git a/zig/kernels/cuda/qwen35_draft.cu b/zig/kernels/cuda/qwen35_draft.cu new file mode 100644 index 0000000000..af2533bc51 --- /dev/null +++ b/zig/kernels/cuda/qwen35_draft.cu @@ -0,0 +1,164 @@ +// DFlash2 drafter glue on CUDA (cuda_draft.zig): a block of rows over the target's taps, draft_ops.metal's rounding. + +#include +#include + +namespace tf_q35d { + +constexpr int HD = 128; // the drafter's head width (q, k, v) + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q35d + +using namespace tf_q35d; + +// Tap rows (start + i) of a stream's ring of R rows, `width` bf16 each, into OUT rows 0..n. Grid (n), 256 threads. +extern "C" __global__ void tf_q35d_gather_taps(const __nv_bfloat16* __restrict__ RING, int R, long long start, int width, + __nv_bfloat16* __restrict__ OUT) { + const long long row = start + blockIdx.x; + const __nv_bfloat16* src = RING + (size_t)(row % R) * width; + __nv_bfloat16* dst = OUT + (size_t)blockIdx.x * width; + for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; +} + +// Rows of a window to gather, by value (a tree's kept path). +struct TapRows { + int n; + int rows[16]; +}; + +// Tap rows (start + rows.rows[i]) of a stream's ring into OUT rows 0..n: a kept path's rows in order. Grid (n). +extern "C" __global__ void tf_q35d_gather_rows(const __nv_bfloat16* __restrict__ RING, int R, long long start, TapRows rows, int width, + __nv_bfloat16* __restrict__ OUT) { + const long long row = start + rows.rows[blockIdx.x]; + const __nv_bfloat16* src = RING + (size_t)(row % R) * width; + __nv_bfloat16* dst = OUT + (size_t)blockIdx.x * width; + for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; +} + +// y[t, c] = sum over lags l of (BASE[l, c] + DYN[t, l, c / group]) * x[t - l, c] in bf16 steps; x[-1] = 0; RES last. +extern "C" __global__ void tf_q35d_dyn_conv(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ DYN, + const __nv_bfloat16* __restrict__ BASE, const __nv_bfloat16* __restrict__ RES, + __nv_bfloat16* __restrict__ OUT, int D, int group, int branch) { + const int t = blockIdx.x; + const int groups = D / group; + for (int c = threadIdx.x; c < D; c += blockDim.x) { + const int g = c / group; + float total = 0.0f; +#pragma unroll + for (int l = 0; l < 2; l++) { + if (t - l < 0) continue; + const float x = bf(X[(size_t)(t - l) * D + c]); + total = rbf(total + rbf(bf(BASE[(branch * 2 + l) * D + c]) * x)); + total = rbf(total + rbf(bf(DYN[(size_t)t * (4 * groups) + (branch * 2 + l) * groups + g]) * x)); + } + if (RES != nullptr) total = rbf(total + bf(RES[(size_t)t * D + c])); + OUT[(size_t)t * D + c] = __float2bfloat16_rn(total); + } +} + +// Per-head RMSNorm then rotate-half RoPE at position pos0 + row; OUT packed (rows, heads, HD), grid (rows * heads). +extern "C" __global__ void __launch_bounds__(32) tf_q35d_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, long long pos0, + float eps, float theta) { + const int head = blockIdx.x % heads; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + head * HD; + float v[4], acc = 0.0f; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = bf(x[4 * t + i]); + acc += v[i] * v[i]; + } + acc = warp_sum(acc); + const float inv = rsqrtf(acc / (float)HD + eps); + float n[4]; +#pragma unroll + for (int i = 0; i < 4; i++) n[i] = rbf(rbf(v[i] * inv) * bf(W[4 * t + i])); + // the pair of lane element d is d + HD / 2: this lane's partner values come from the lane HD / 8 = 16 away + const float pos = (float)(pos0 + (long long)m); + __nv_bfloat16* out = OUT + (m * heads + head) * HD; +#pragma unroll + for (int i = 0; i < 4; i++) { + const int d = 4 * t + i; + const float partner = __shfl_xor_sync(0xffffffffu, n[i], 16); + const int j = d < HD / 2 ? d : d - HD / 2; + const float angle = pos * powf(theta, -(float)j / (float)(HD / 2)); + float sn, cs; + sincosf(angle, &sn, &cs); + out[d] = __float2bfloat16_rn(d < HD / 2 ? n[i] * cs - partner * sn : n[i] * cs + partner * sn); + } +} + +// Rows of packed (kv_heads, HD) keys and values at positions pos0 + i into a layer's rings (slot pos % W). +extern "C" __global__ void tf_q35d_ring_put(const __nv_bfloat16* __restrict__ K, const __nv_bfloat16* __restrict__ V, + __nv_bfloat16* __restrict__ RK, __nv_bfloat16* __restrict__ RV, int width, long long pos0, int W) { + const int i = blockIdx.x; + const size_t slot = (size_t)((pos0 + i) % W); + for (int c = threadIdx.x; c < width; c += blockDim.x) { + RK[slot * width + c] = K[(size_t)i * width + c]; + RV[slot * width + c] = V[(size_t)i * width + c]; + } +} + +// Block attention: a row sees ring positions [max(begin, p - window + 1), pos0) and every block row; a warp a row. +extern "C" __global__ void __launch_bounds__(32) tf_q35d_attn(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KB, + const __nv_bfloat16* __restrict__ VB, const __nv_bfloat16* __restrict__ RK, + const __nv_bfloat16* __restrict__ RV, int rows, int heads, int kv_heads, int W, + long long begin, long long pos0, int window, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + const float scale = rsqrtf((float)HD); + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + const long long p = pos0 + m; + long long first = p - window + 1; + if (first < begin) first = begin; + for (long long pos = first; pos < pos0; pos++) { + const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + for (int r = 0; r < rows; r++) { + const __nv_bfloat16* k = KB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = VB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } +#pragma unroll + for (int i = 0; i < DPL; i++) OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(o[i] / l); +} diff --git a/zig/kernels/cuda/qwen35_keep.cu b/zig/kernels/cuda/qwen35_keep.cu new file mode 100644 index 0000000000..20d00fc212 --- /dev/null +++ b/zig/kernels/cuda/qwen35_keep.cu @@ -0,0 +1,58 @@ +// Keeping a verified window's path (cuda_cache.zig): the recurrences take its logged updates and its keys move down. + +#include "qwen35.cuh" + +using namespace tf_q35; + +// Each linear layer's recurrence through the path's (decay, key, write), as the forward's chain would have left it. +extern "C" __global__ void __launch_bounds__(128) tf_q35k_apply(Geo geo, Seg sg, Path path) { + const int NK = geo.nk, NV = geo.nv; + const int layer = blockIdx.z / NV, hv = blockIdx.z % NV; + const int hk = hv / (NV / NK); + const int dk = threadIdx.x; + const int dv = blockIdx.y * 4 + threadIdx.y; + constexpr int PT = DK / 32; + const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; + float state[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int j = 0; j < path.n; j++) { + const LogRow l = log_row(geo, sg, layer, path.rows[j]); + const float g = l.g[hv], d = l.delta[hv * DK + dv]; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * g; + state[i] = fmaf(bf(l.k[hk * DK + PT * dk + i]), d, state[i]); + } + } +#pragma unroll + for (int i = 0; i < PT; i++) st[i] = state[i]; +} + +// Each linear layer's convolution tail as logged after window row `last`. Grid (nl, tail values / 256), 256 threads. +extern "C" __global__ void tf_q35k_conv(Geo geo, Seg sg, int last) { + const int layer = blockIdx.x; + const size_t n = geo.conv_rows(); + const size_t i = (size_t)blockIdx.y * blockDim.x + threadIdx.x; + if (i >= n) return; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * n; + CS[i] = log_row(geo, sg, layer, last).conv[i]; +} + +// The path's keys and values move down to the slots after the cache; row j lands at or below j, never on one unread. +extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact(Geo geo, Seg sg, Path path) { + const int KVH = geo.kvh; + const int layer = blockIdx.x, h = blockIdx.y; + __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + ((size_t)layer * 2 * KVH + h) * sg.cap * HD; + __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + const int d = threadIdx.x; + for (int j = 0; j < path.n; j++) { + const int from = sg.len + path.rows[j], to = sg.len + j; + if (from != to) { + keys[(size_t)to * HD + d] = keys[(size_t)from * HD + d]; + values[(size_t)to * HD + d] = values[(size_t)from * HD + d]; + } + __syncthreads(); + } +} diff --git a/zig/kernels/cuda/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 99d6835a49..100b81ecf9 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -40,7 +40,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/eaq_container.zig b/zig/src/core/eaq_container.zig index 068faeefc1..914e86e2b9 100644 --- a/zig/src/core/eaq_container.zig +++ b/zig/src/core/eaq_container.zig @@ -291,7 +291,7 @@ fn resolve(a: std.mem.Allocator, io: Io, dir: []const u8, file: []const u8, root }; for (roots) |r| { const root = r orelse continue; - if (real.len > root.len and std.mem.startsWith(u8, real, root) and real[root.len] == '/') return real; + if (real.len > root.len and std.mem.startsWith(u8, real, root) and std.fs.path.isSep(real[root.len])) return real; } Why.set(why, "{s}", .{file}); return error.ShardOutside; diff --git a/zig/src/core/eaq_container_test.zig b/zig/src/core/eaq_container_test.zig index 05e3576aa3..4099e203c5 100644 --- a/zig/src/core/eaq_container_test.zig +++ b/zig/src/core/eaq_container_test.zig @@ -222,7 +222,8 @@ test "a shard linked outside the pack's folder is refused; a snapshot's links in defer pack.close(io); try fixture.write(gpa, io, pack, .none, .{}); try std.Io.Dir.rename(pack, "model-00002.safetensors", tmp.dir, "outside.safetensors", io); - try pack.symLink(io, "../outside.safetensors", "model-00002.safetensors", .{}); + pack.symLink(io, "../outside.safetensors", "model-00002.safetensors", .{}) catch |e| + return if (e == error.PermissionDenied and @import("builtin").os.tag == .windows) error.SkipZigTest else e; // symlinks need privilege on Windows var why: eaq.Why = .{}; try std.testing.expectError(error.ShardOutside, eaq.open(gpa, io, try std.fmt.bufPrint(&path, "{s}/pack", .{root}), .{ .why = &why })); try std.testing.expectEqualStrings("model-00002.safetensors", why.text()); 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_rows_test.zig b/zig/src/core/lane_host_rows_test.zig index 3a6cb1e6ed..937bb50faf 100644 --- a/zig/src/core/lane_host_rows_test.zig +++ b/zig/src/core/lane_host_rows_test.zig @@ -1,5 +1,6 @@ //! A lane host's kept states on disk: written while idle, finished on a stop, read back by the next host, asked about. const std = @import("std"); +const platform = @import("sys"); const lanes = @import("lanes"); const api = @import("engine_api.zig"); const pc = @import("prompt_cache.zig"); @@ -114,7 +115,7 @@ fn prompt(n: usize) [700]u32 { test "a host writes its kept states while idle and on its stop; the next host resumes from disk, its reply unchanged" { var buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&buf, "/tmp/tf-host-rows-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&buf, "/tmp/tf-host-rows-{d}", .{platform.c.getpid()}); defer rmTree(root); const S = struct { var first: [706]u32 = undefined; // turn one's prompt and its reply diff --git a/zig/src/core/lanes/cost_cache.zig b/zig/src/core/lanes/cost_cache.zig index 1b04fe8fde..279aca3931 100644 --- a/zig/src/core/lanes/cost_cache.zig +++ b/zig/src/core/lanes/cost_cache.zig @@ -1,5 +1,6 @@ //! Timed costs are measured once per engine build, chip and model shape, then read back by later loads. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; const magic = "TFWC2\n"; @@ -53,7 +54,7 @@ fn hashPart(h: *std.crypto.hash.sha2.Sha256, part: []const u8) void { } fn path(a: Allocator, k: [32]u8, suffix: []const u8) !?[]u8 { - const home = std.c.getenv("HOME") orelse return null; + const home = platform.c.getenv("HOME") orelse return null; return try std.fmt.allocPrint(a, "{s}/.cache/tensorfold/window-costs/{s}{s}", .{ std.mem.span(home), std.fmt.bytesToHex(k, .lower), suffix }); } @@ -78,7 +79,7 @@ pub fn save(comptime T: type, a: Allocator, io: std.Io, k: [32]u8, value: T) voi const p = (path(a, k, ".bin") catch return) orelse return; defer a.free(p); std.Io.Dir.cwd().createDirPath(io, std.fs.path.dirname(p) orelse return) catch return; - const temp = std.fmt.allocPrint(a, "{s}.{d}.tmp", .{ p, std.c.getpid() }) catch return; + const temp = std.fmt.allocPrint(a, "{s}.{d}.tmp", .{ p, platform.c.getpid() }) catch return; defer a.free(temp); const bytes = std.mem.concat(a, u8, &.{ magic, std.mem.asBytes(&value) }) catch return; defer a.free(bytes); diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 3bdffe7c7a..e1a79fd2de 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -60,7 +60,9 @@ pub const Drafted = struct { m.plain_guard = d.plain_guard; m.draft_streams = d.batched; m.draft_probabilities = false; - m.head_trees = false; + m.head_trees = d.trees; + if (d.window_costs.len > 0) m.window_costs = d.window_costs; + if (d.lane_costs.len > 0) m.lane_costs = d.lane_costs; return m; } @@ -157,15 +159,17 @@ pub const Drafted = struct { x.holds.clearRetainingCapacity(); try x.firsts.resize(x.gpa, requests.len); for (requests, x.firsts.items) |r, *tok| { - if (r.lanes != null) return error.TreesNotBuilt; + if (r.lanes != null and !x.drafter.facts().trees) return error.TreesNotBuilt; const lane = x.lanes.getPtr(r.stream) orelse return error.NotDrafting; lane.filler = 0; // a new request supersedes fillers a copy round left unused var pending: u32 = undefined; if (r.rows) |rows| { - for (rows, 0..) |row, i| if (row != i) return error.TreesNotBuilt; - const n: u32 = @intCast(rows.len); + // a tree's kept path names the window's rows out of order: the features of all up to its last + var prefix = true; + for (rows, 0..) |row, i| prefix = prefix and row == i; + const n: u32 = @intCast(if (prefix) rows.len else rows[rows.len - 1] + 1); const f = try x.target.features(r.stream, x.drafter.taps(), r.start, n); - try x.absorbs.append(x.gpa, .{ .stream = r.stream, .features = f, .start = r.start, .follow = r.follow[0..rows.len] }); + try x.absorbs.append(x.gpa, .{ .stream = r.stream, .features = f, .start = r.start, .follow = r.follow[0..rows.len], .rows = if (prefix) null else rows }); pending = r.follow[rows.len - 1]; lane.rows = true; } else { @@ -187,7 +191,7 @@ pub const Drafted = struct { lane.filler = r.depth; continue; } - try x.holds.append(x.gpa, .{ .stream = r.stream, .pending = pending, .position = r.position, .depth = r.depth }); + try x.holds.append(x.gpa, .{ .stream = r.stream, .pending = pending, .position = r.position, .depth = r.depth, .lanes = r.lanes }); } if (x.absorbs.items.len > 0) try x.drafter.absorb(x.absorbs.items); if (x.holds.items.len > 0) try x.drafter.hold(x.holds.items); diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 0f8eab242a..c956522f50 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -1,6 +1,8 @@ //! A drafter outside the target family reading its tapped states; each call takes a round's streams together. const Stream = @import("stream.zig").Stream; const Features = @import("backend.zig").Features; +const Cost = @import("config.zig").Cost; +const Shape = @import("shape.zig").Shape; /// What the round loop's depth rule needs from the drafter itself (not from the target, whose own head may not exist). pub const Facts = struct { @@ -10,13 +12,16 @@ pub const Facts = struct { plain_guard: bool = true, // plain rounds compete with drafted depths (the rule may choose no drafts) batched: bool = false, // hold() runs a shared round's streams as one batch stream_bytes: usize = 0, // device bytes each opened stream holds (its caches), for memory admission + trees: bool = false, // hold() takes a tree of lanes (Hold.lanes), each lane its parent's rank-r token + window_costs: []const Cost = &.{}, // the target's window ms by rows (timed at load), for the depth rule + lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) }; -/// Target rows [start, start + features.rows) and the token after each; rows at or past `start` replace what was held. -pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32 }; +/// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb. +pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32, rows: ?[]const u32 = null }; -/// Draft `depth` tokens after `pending`, the first at `position`; asked only once row `position - 2` was absorbed. -pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32 }; +/// `depth` drafts after `pending` at `position`; with `lanes` a tree, lane i the rank-r choice after its parent. +pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32, lanes: ?*const Shape = null }; pub const Drafter = struct { ptr: *anyopaque, diff --git a/zig/src/core/lanes/grammar.zig b/zig/src/core/lanes/grammar.zig index 371c2ac4d8..2fe9ae3d87 100644 --- a/zig/src/core/lanes/grammar.zig +++ b/zig/src/core/lanes/grammar.zig @@ -1,6 +1,7 @@ //! Structured output: libtfgrammar (xgrammar, loaded at the first structured request) and a stream's grammar walk. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Allocator = std.mem.Allocator; @@ -37,7 +38,7 @@ fn load(io: std.Io, a: Allocator) ?Fns { library, }; var dl = for (names) |name| { - break std.DynLib.open(name) catch continue; + break platform.DynLib.open(name) catch continue; } else return null; var f: Fns = undefined; const info = @typeInfo(Fns).@"struct"; diff --git a/zig/src/core/lanes/learned_dirs.zig b/zig/src/core/lanes/learned_dirs.zig index 780a997a5c..cdf9648438 100644 --- a/zig/src/core/lanes/learned_dirs.zig +++ b/zig/src/core/lanes/learned_dirs.zig @@ -1,5 +1,6 @@ //! Bounded-name learned directories: both ranks plan the same identity deletions before either writes. const std = @import("std"); +const platform = @import("sys"); pub fn rootOf(dir: []const u8) []const u8 { return std.fs.path.dirname(dir) orelse "."; } @@ -13,23 +14,23 @@ fn isCurrent(keep: []const u8, id: u64) bool { } pub fn next(root: []const u8, keep: []const u8, after: ?u64) !?u64 { var buf: [1100]u8 = undefined; - const d = std.c.opendir(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0)) orelse return error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0)) orelse return error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); var out: ?u64 = null; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (name.len != 16) continue; const id = std.fmt.parseInt(u64, name, 16) catch continue; if (id == std.math.maxInt(u64)) return error.LearnedDirectoryType; - if (ent.type != std.c.DT.DIR) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.DIR) return error.LearnedDirectoryType; if (isCurrent(keep, id) or (after != null and id <= after.?)) continue; if (out == null or id < out.?) out = id; } return out; } pub fn unlink(file: [:0]const u8) !void { - const rc = std.c.unlink(file); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.unlink(file); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } fn stateKey(name: []const u8) ?u64 { @@ -52,38 +53,38 @@ fn identityName(name: []const u8) bool { } fn unlinkAt(fd: c_int, name: [*:0]const u8) !void { - const rc = std.c.unlinkat(fd, name, 0); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.unlinkat(fd, name, 0); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } /// Only the configured root may be an alias; cleanup never follows an identity or file link. pub fn cleanupParts(root: []const u8) !void { var buf: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0), .{ .ACCMODE = .RDONLY, .DIRECTORY = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0), .{ .ACCMODE = .RDONLY, .DIRECTORY = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const d = std.c.fdopendir(fd) orelse { - _ = std.c.close(fd); + const d = platform.c.fdopendir(fd) orelse { + _ = platform.c.close(fd); return error.LearnedDirectoryRead; }; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); - if (ent.type == std.c.DT.REG and std.mem.eql(u8, name, "clock.part")) { + if (ent.type == platform.c.DT.REG and std.mem.eql(u8, name, "clock.part")) { try unlinkAt(fd, @ptrCast(&ent.name)); continue; } if (!identityName(name)) continue; - if (ent.type != std.c.DT.DIR) return error.LearnedDirectoryType; - const sub_fd = std.c.openat(fd, @ptrCast(&ent.name), .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + if (ent.type != platform.c.DT.DIR) return error.LearnedDirectoryType; + const sub_fd = platform.c.openat(fd, @ptrCast(&ent.name), .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (sub_fd < 0) return error.LearnedDirectoryRead; - const sub = std.c.fdopendir(sub_fd) orelse { - _ = std.c.close(sub_fd); + const sub = platform.c.fdopendir(sub_fd) orelse { + _ = platform.c.close(sub_fd); return error.LearnedDirectoryRead; }; { - defer _ = std.c.closedir(sub); - while (std.c.readdir(sub)) |file| { - if (file.type == std.c.DT.REG and partName(std.mem.sliceTo(&file.name, 0))) try unlinkAt(sub_fd, @ptrCast(&file.name)); + defer _ = platform.c.closedir(sub); + while (platform.c.readdir(sub)) |file| { + if (file.type == platform.c.DT.REG and partName(std.mem.sliceTo(&file.name, 0))) try unlinkAt(sub_fd, @ptrCast(&file.name)); } } } @@ -91,15 +92,15 @@ pub fn cleanupParts(root: []const u8) !void { /// Payloads have no token metadata to recover a lost index entry; only exact owned, unreferenced names go. pub fn removeUnindexed(dir: [:0]const u8, metas: anytype) !void { - const fd = std.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const d = std.c.fdopendir(fd) orelse { - _ = std.c.close(fd); + const d = platform.c.fdopendir(fd) orelse { + _ = platform.c.close(fd); return error.LearnedDirectoryRead; }; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { - if (ent.type != std.c.DT.REG) continue; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { + if (ent.type != platform.c.DT.REG) continue; const name = std.mem.sliceTo(&ent.name, 0); const key = stateKey(name) orelse continue; const indexed = for (metas) |x| { @@ -110,29 +111,29 @@ pub fn removeUnindexed(dir: [:0]const u8, metas: anytype) !void { } } pub fn bytes(dir: [:0]const u8) !u64 { - const d = std.c.opendir(dir) orelse return if (std.c.errno(-1) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(dir) orelse return if (platform.c.errno(-1) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); var n: u64 = 0; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) continue; - if (ent.type != std.c.DT.REG) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.REG) return error.LearnedDirectoryType; var buf: [1400]u8 = undefined; const f = try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0); - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - _ = std.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); + _ = platform.c.close(fd); if (end < 0) return error.LearnedDirectoryRead; n = try std.math.add(u64, n, @intCast(end)); } return n; } pub fn fileBytes(file: [:0]const u8) !u64 { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.LearnedDirectoryRead; return @intCast(end); } @@ -156,16 +157,16 @@ pub fn total(root: []const u8) !u64 { pub fn used(root: []const u8, id: u64) !u64 { var buf: [1100]u8 = undefined; const f = try std.fmt.bufPrintSentinel(&buf, "{s}/{x:0>16}/used", .{ root, id }, 0); - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.close(fd); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.close(fd); var n: u64 = 0; const data = std.mem.asBytes(&n); var done: usize = 0; while (done < data.len) { - const count = std.c.read(fd, data.ptr + done, data.len - done); + const count = platform.c.read(fd, data.ptr + done, data.len - done); if (count < 0) { - if (std.c.errno(count) == .INTR) continue; + if (platform.c.errno(count) == .INTR) continue; return error.LearnedDirectoryRead; } if (count == 0) return 0; // A torn recency hint is oldest; it cannot make a learned payload valid. @@ -178,23 +179,23 @@ pub fn removeOther(root: []const u8, keep: []const u8, id: u64) !void { var buf: [1100]u8 = undefined; const dir = try path(&buf, root, id); _ = try bytes(dir); - const d = std.c.opendir(dir) orelse return if (std.c.errno(-1) == .NOENT) {} else error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + const d = platform.c.opendir(dir) orelse return if (platform.c.errno(-1) == .NOENT) {} else error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) continue; - if (ent.type != std.c.DT.REG) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.REG) return error.LearnedDirectoryType; var f: [1400]u8 = undefined; try unlink(try std.fmt.bufPrintSentinel(&f, "{s}/{s}", .{ dir, name }, 0)); } - const rc = std.c.rmdir(dir); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.rmdir(dir); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } test "acknowledged deletion refuses a directory and accepts an absent file" { var buf: [128]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-unlink-{d}", .{std.c.getpid()}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-unlink-{d}", .{platform.c.getpid()}, 0); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); try std.testing.expectError(error.LearnedUnlink, unlink(dir)); var missing: [160]u8 = undefined; try unlink(try std.fmt.bufPrintSentinel(&missing, "{s}/absent", .{dir}, 0)); diff --git a/zig/src/core/lanes/learned_disk.zig b/zig/src/core/lanes/learned_disk.zig index 9ca5fbd4c0..d6e275fae1 100644 --- a/zig/src/core/lanes/learned_disk.zig +++ b/zig/src/core/lanes/learned_disk.zig @@ -1,5 +1,6 @@ //! Reserve learned-state writes above a free-disk floor, with bounded backoff after refusals or write failures. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Count = if (builtin.os.tag == .linux) u64 else u32; const Stat = extern struct { @@ -19,14 +20,15 @@ const Stat = extern struct { extern "c" fn statvfs(path: [*:0]const u8, result: *Stat) c_int; pub fn available(dir: [:0]const u8) ?u64 { + if (builtin.os.tag == .windows) return platform.c.diskAvailable(dir); var stat: Stat = undefined; if (statvfs(dir, &stat) != 0) return null; return std.math.mul(u64, stat.frsize, stat.bavail) catch null; } fn now() ?u64 { - var t: std.c.timespec = undefined; - if (std.c.clock_gettime(std.c.CLOCK.MONOTONIC, &t) != 0 or t.sec < 0) return null; + var t: platform.c.timespec = undefined; + if (platform.c.clock_gettime(platform.c.CLOCK.MONOTONIC, &t) != 0 or t.sec < 0) return null; return @intCast(t.sec); } diff --git a/zig/src/core/lanes/recall_store.zig b/zig/src/core/lanes/recall_store.zig index 857075b5cc..9628906433 100644 --- a/zig/src/core/lanes/recall_store.zig +++ b/zig/src/core/lanes/recall_store.zig @@ -1,13 +1,14 @@ //! Recall's record on disk: one tokenizer's ids in chunk files of a 0700 directory, the oldest dropped past the cap. const std = @import("std"); +const platform = @import("sys"); const recall = @import("recall.zig"); const rr = recall.record; const Admission = @import("learned_disk.zig").Admission; const Allocator = std.mem.Allocator; -extern "c" fn flock(fd: c_int, operation: c_int) c_int; -const lock_ex = 2; -const lock_nb = 4; +const flock = platform.c.flock; +const lock_ex = platform.c.LOCK.EX; +const lock_nb = platform.c.LOCK.NB; const magic = "TFRECALL"; const version: u32 = 1; @@ -122,7 +123,7 @@ pub const Store = struct { } if (s.writes) s.takeLock(); errdefer if (s.lock_fd >= 0) { - _ = std.c.close(s.lock_fd); + _ = platform.c.close(s.lock_fd); }; try s.load(); return s; @@ -130,7 +131,7 @@ pub const Store = struct { pub fn close(s: *Store) void { s.closeFile(); - if (s.lock_fd >= 0) _ = std.c.close(s.lock_fd); + if (s.lock_fd >= 0) _ = platform.c.close(s.lock_fd); s.shared.record.deinit(); s.gpa.free(s.name); s.gpa.free(s.dir); @@ -139,8 +140,8 @@ pub const Store = struct { fn closeFile(s: *Store) void { if (s.file < 0) return; - _ = std.c.fsync(s.file); - _ = std.c.close(s.file); + _ = platform.c.fsync(s.file); + _ = platform.c.close(s.file); s.file = -1; } @@ -189,10 +190,10 @@ pub const Store = struct { fn takeLock(s: *Store) void { var buf: [1100]u8 = undefined; const path = std.fmt.bufPrintSentinel(&buf, "{s}/lock", .{s.dir}, 0) catch return s.readOnly(); - const fd = std.c.open(path, .{ .ACCMODE = .RDWR, .CREAT = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDWR, .CREAT = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return s.readOnly(); if (flock(fd, lock_ex | lock_nb) != 0) { - _ = std.c.close(fd); + _ = platform.c.close(fd); return s.readOnly(); } s.lock_fd = fd; @@ -226,10 +227,10 @@ pub const Store = struct { fn readChunk(s: *Store, seq: u64, pending: *std.ArrayList(u32), ids: []u32) !usize { var buf: [1100]u8 = undefined; const path = try s.chunkPath(&buf, seq); - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.RecallRead; - defer _ = std.c.close(fd); - const size = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const size = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (size < head_bytes) return error.RecallRead; var head: [head_bytes]u8 = undefined; if (!readAt(fd, &head, 0) or !std.mem.eql(u8, &head, &s.header())) return error.RecallRead; @@ -260,13 +261,13 @@ pub const Store = struct { var buf: [1100]u8 = undefined; const path = s.chunkPath(&buf, seq) catch continue; if (k >= chunks or seq != s.seq0 + k) { - _ = std.c.unlink(path); + _ = platform.c.unlink(path); } else if (k + 1 == chunks) { const keep = head_bytes + 4 * (@as(u64, r.end) - @as(u64, k) * rr.chunk_tokens); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd >= 0) { - _ = std.c.ftruncate(fd, @intCast(keep)); - _ = std.c.close(fd); + _ = platform.c.ftruncate(fd, @intCast(keep)); + _ = platform.c.close(fd); } } } @@ -278,9 +279,9 @@ pub const Store = struct { /// The chunk files' numbers, ascending. fn listChunks(s: *Store, out: *std.ArrayList(u64)) !void { - const d = std.c.opendir(s.dir) orelse return error.RecallRead; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + const d = platform.c.opendir(s.dir) orelse return error.RecallRead; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (name.len != 20 or !std.mem.endsWith(u8, name, ".ids")) continue; const seq = std.fmt.parseInt(u64, name[0..16], 16) catch continue; @@ -296,7 +297,7 @@ pub const Store = struct { if (s.writes) { var buf: [1100]u8 = undefined; if (s.chunkPath(&buf, s.seq0 + r.first - 1)) |path| { - _ = std.c.unlink(path); + _ = platform.c.unlink(path); } else |_| {} } return true; @@ -394,16 +395,16 @@ pub const Store = struct { s.closeFile(); var buf: [1100]u8 = undefined; const path = s.chunkPath(&buf, s.seq0 + k) catch return false; - var flags: std.c.O = .{ .ACCMODE = .WRONLY, .APPEND = true, .NOFOLLOW = true }; - var fd = std.c.open(path, flags, @as(std.c.mode_t, 0)); + var flags: platform.c.O = .{ .ACCMODE = .WRONLY, .APPEND = true, .NOFOLLOW = true }; + var fd = platform.c.open(path, flags, @as(platform.c.mode_t, 0)); if (fd < 0) { flags.CREAT = true; flags.EXCL = true; - fd = std.c.open(path, flags, @as(std.c.mode_t, 0o600)); + fd = platform.c.open(path, flags, @as(platform.c.mode_t, 0o600)); if (fd < 0) return false; const head = s.header(); writeAll(fd, &head) catch { - _ = std.c.close(fd); + _ = platform.c.close(fd); return false; }; } @@ -437,8 +438,8 @@ pub const Store = struct { for (seqs.items) |seq| { var buf: [1100]u8 = undefined; const path = s.chunkPath(&buf, seq) catch continue; - const rc = std.c.unlink(path); - const e = std.c.errno(rc); + const rc = platform.c.unlink(path); + const e = platform.c.errno(rc); if (rc == 0 or e == .NOENT or s.problem.len > 0) continue; s.note("{s} could not be removed ({s})", .{ path, @tagName(e) }); } @@ -494,18 +495,18 @@ fn makePath(dir: [:0]const u8) !void { buf[dir.len] = 0; for (1..dir.len) |i| if (buf[i] == '/') { buf[i] = 0; - _ = std.c.mkdir(@ptrCast(&buf), 0o700); + _ = platform.c.mkdir(@ptrCast(&buf), 0o700); buf[i] = '/'; }; - const rc = std.c.mkdir(dir, 0o700); - if (rc != 0 and std.c.errno(rc) != .EXIST) return error.RecallDir; + const rc = platform.c.mkdir(dir, 0o700); + if (rc != 0 and platform.c.errno(rc) != .EXIST) return error.RecallDir; } fn readSmall(a: Allocator, path: [:0]const u8) ![]u8 { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) error.FileNotFound else error.RecallRead; - defer _ = std.c.close(fd); - const size = std.c.lseek(fd, 0, std.c.SEEK.END); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) error.FileNotFound else error.RecallRead; + defer _ = platform.c.close(fd); + const size = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (size < 0 or size > 4096) return error.RecallRead; const out = try a.alloc(u8, @intCast(size)); if (!readAt(fd, out, 0)) return error.RecallRead; @@ -516,21 +517,21 @@ fn readSmall(a: Allocator, path: [:0]const u8) ![]u8 { fn writeSmall(path: [:0]const u8, bytes: []const u8) !void { var buf: [1200]u8 = undefined; const part = try std.fmt.bufPrintSentinel(&buf, "{s}.part", .{path}, 0); - const flags: std.c.O = .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }; - const fd = std.c.open(part, flags, @as(std.c.mode_t, 0o600)); + const flags: platform.c.O = .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }; + const fd = platform.c.open(part, flags, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.RecallWrite; { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try writeAll(fd, bytes); - if (std.c.fsync(fd) != 0) return error.RecallWrite; + if (platform.c.fsync(fd) != 0) return error.RecallWrite; } - if (std.c.rename(part, path) != 0) return error.RecallWrite; + if (platform.c.rename(part, path) != 0) return error.RecallWrite; } 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.RecallWrite; done += @intCast(n); } @@ -539,7 +540,7 @@ fn writeAll(fd: c_int, bytes: []const u8) !void { 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); } @@ -548,7 +549,7 @@ fn readAt(fd: c_int, dest: []u8, at: u64) bool { /// Where Recall keeps records unless --recall-dir names one: $HOME/.cache/tensorfold/recall/. pub fn defaultDir(a: Allocator, identity: Identity) ![]const u8 { - const home = std.c.getenv("HOME") orelse return error.RecallDir; + const home = platform.c.getenv("HOME") orelse return error.RecallDir; const short = identity.short(); return std.fmt.allocPrint(a, "{s}/.cache/tensorfold/recall/{s}", .{ std.mem.span(home), &short }); } diff --git a/zig/src/core/lanes/recall_store_test.zig b/zig/src/core/lanes/recall_store_test.zig index 408a60ef0c..5fc1ae4ea5 100644 --- a/zig/src/core/lanes/recall_store_test.zig +++ b/zig/src/core/lanes/recall_store_test.zig @@ -1,5 +1,6 @@ //! The store on disk: reloads, refusal by name, torn tails, bad ids, the cap, dedup, read-only, emptying. const std = @import("std"); +const platform = @import("sys"); const store = @import("recall_store.zig"); const rr = @import("recall_record.zig"); const Store = store.Store; @@ -120,11 +121,11 @@ test "a torn tail and an id past the vocabulary are cut at reload, and writing g } var buf: [1100]u8 = undefined; const path = try std.fmt.bufPrintSentinel(&buf, "{s}/{x:0>16}.ids", .{ d.path, @as(u64, 0) }, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(platform.c.mode_t, 0)); try std.testing.expect(fd >= 0); const torn = [_]u32{ 7, 8, 9 }; // a segment whose mark never landed - _ = std.c.write(fd, std.mem.sliceAsBytes(&torn).ptr, 12); - _ = std.c.close(fd); + _ = platform.c.write(fd, std.mem.sliceAsBytes(&torn).ptr, 12); + _ = platform.c.close(fd); const b = try ids(80, 7); defer gpa.free(b); { @@ -138,9 +139,9 @@ test "a torn tail and an id past the vocabulary are cut at reload, and writing g try std.testing.expectEqual(@as(u32, 101 + 81), s.shared.record.held()); try std.testing.expect(s.shared.record.longest(b[0..40]).len > 0); const bad = [_]u32{ 1, 2, vocab + 3, rr.mark }; - const fd2 = std.c.open(path, .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(std.c.mode_t, 0)); - _ = std.c.write(fd2, std.mem.sliceAsBytes(&bad).ptr, 16); - _ = std.c.close(fd2); + const fd2 = platform.c.open(path, .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(platform.c.mode_t, 0)); + _ = platform.c.write(fd2, std.mem.sliceAsBytes(&bad).ptr, 16); + _ = platform.c.close(fd2); const again = try open(&d, identity(5, "m"), .{ .read_only = true }); defer again.close(); try std.testing.expectEqual(@as(u32, 182), again.shared.record.held()); @@ -161,7 +162,7 @@ test "past the cap the oldest chunk goes, file and all" { try std.testing.expectEqual(@as(u32, 0), s.shared.record.longest(first[0..20]).len); var buf: [1100]u8 = undefined; const oldest = try std.fmt.bufPrintSentinel(&buf, "{s}/{x:0>16}.ids", .{ d.path, @as(u64, 0) }, 0); - try std.testing.expect(std.c.access(oldest, 0) != 0); + try std.testing.expect(platform.c.access(oldest, 0) != 0); try std.testing.expectEqual(@as(u32, rr.max_match), s.shared.record.longest(bulk[0 .. bulk.len - 1]).len); } @@ -208,17 +209,17 @@ test "an empty that cannot remove a chunk file says which, keeps nothing in memo s.readPrompt(a, 0); var buf: [1100]u8 = undefined; const stuck = try std.fmt.bufPrintSentinel(&buf, "{s}/{x:0>16}.ids", .{ d.path, @as(u64, 0xff) }, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(stuck, 0o700)); // a directory: no unlink removes it + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(stuck, 0o700)); // a directory: no unlink removes it var inner: [1200]u8 = undefined; const keep = try std.fmt.bufPrintSentinel(&inner, "{s}/x", .{stuck}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(keep, 0o700)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(keep, 0o700)); try std.testing.expectError(error.Unlinked, s.empty()); try std.testing.expect(std.mem.indexOf(u8, s.why(), "00000000000000ff.ids") != null); try std.testing.expect(s.shared.record.held() == 0 and s.added == 0 and s.held == 0 and s.stats().left); s.readPrompt(a, 0); try std.testing.expectEqual(@as(u64, 0), s.added); // nothing goes in while a file is left - _ = std.c.rmdir(keep); - _ = std.c.rmdir(stuck); + _ = platform.c.rmdir(keep); + _ = platform.c.rmdir(stuck); try s.empty(); try std.testing.expect(!s.stats().left); s.readPrompt(a, 0); diff --git a/zig/src/core/lanes/recorder.zig b/zig/src/core/lanes/recorder.zig index e62eda60fd..c57a1ae250 100644 --- a/zig/src/core/lanes/recorder.zig +++ b/zig/src/core/lanes/recorder.zig @@ -1,5 +1,6 @@ //! Drafter training data from serving: each committed row's token, position and tapped features, files per stream. const std = @import("std"); +const platform = @import("sys"); const Io = std.Io; const Allocator = std.mem.Allocator; const be = @import("backend.zig"); @@ -44,7 +45,7 @@ pub const Recorder = struct { defer root.close(io); var k: u32 = 0; while (true) : (k += 1) { - const name = try std.fmt.allocPrint(gpa, "{d}-{d}", .{ std.c.getpid(), k }); + const name = try std.fmt.allocPrint(gpa, "{d}-{d}", .{ platform.c.getpid(), k }); errdefer gpa.free(name); root.createDir(io, name, .default_dir) catch |e| switch (e) { error.PathAlreadyExists => { @@ -122,7 +123,7 @@ pub const Recorder = struct { for (exts, &files) |ext, *f| { const path = try std.fmt.allocPrint(r.gpa, "{s}{s}", .{ base, ext }); defer r.gpa.free(path); - f.* = try r.space.createFile(r.io, path, .{ .exclusive = true }); + f.* = try r.createNew(path); made += 1; } r.streams.putAssumeCapacity(s, .{ .base = base, .files = files }); @@ -194,9 +195,17 @@ pub const Recorder = struct { try r.create(try std.fmt.allocPrint(a, "{s}.json", .{o.base}), text.written()); } + /// A new file at `path`, never through a link there: Windows' exclusive create follows a dangling one. + fn createNew(r: *Recorder, path: []const u8) !Io.File { + if (comptime @import("builtin").os.tag == .windows) { + if (r.space.statFile(r.io, path, .{ .follow_symlinks = false })) |_| return error.PathAlreadyExists else |_| {} + } + return r.space.createFile(r.io, path, .{ .exclusive = true }); + } + /// A new file in the recorder's folder with these bytes; one already there is an error, never overwritten. fn create(r: *Recorder, path: []const u8, bytes: []const u8) !void { - const f = try r.space.createFile(r.io, path, .{ .exclusive = true }); + const f = try r.createNew(path); defer f.close(r.io); try f.writeStreamingAll(r.io, bytes); } diff --git a/zig/src/core/lanes/recorder_test.zig b/zig/src/core/lanes/recorder_test.zig index ab3ac33079..47d8ab8d5d 100644 --- a/zig/src/core/lanes/recorder_test.zig +++ b/zig/src/core/lanes/recorder_test.zig @@ -1,5 +1,6 @@ //! The recorder in front of the fake target: replies unchanged, and each stream's files hold its committed rows. const std = @import("std"); +const platform = @import("sys"); const Config = @import("config.zig").Config; const Model = @import("config.zig").Model; const Cost = @import("config.zig").Cost; @@ -69,7 +70,7 @@ fn run(cases: []const Case, mode: Mode, record: bool) ![][]u32 { var recorder = try rec.Recorder.init(gpa, io, dir, &.{ 30, 31 }); recorder.floor = 0; defer recorder.deinit(); - if (mode.planted) |p| try recorder.space.symLink(io, p[0], p[1], .{}); + if (mode.planted) |p| recorder.space.symLink(io, p[0], p[1], .{}) catch |e| return if (e == error.PermissionDenied and @import("builtin").os.tag == .windows) error.SkipZigTest else e; // symlinks need privilege on Windows const inner = if (mode.chunks > 0) target.stepped() else target.backend(); var recorded = try rec.Recorded.init(inner, &recorder); var clock: fake.FixedClock = .{}; @@ -193,10 +194,11 @@ test "each recorder makes its own new folder: two in one folder, and a link wher defer away.cleanup(); var buf: [64]u8 = undefined; const dir = try std.fmt.bufPrintSentinel(&buf, ".zig-cache/tmp/{s}", .{tmp.sub_path}, 0); - const first = try std.fmt.allocPrint(gpa, "{d}-0", .{std.c.getpid()}); + const first = try std.fmt.allocPrint(gpa, "{d}-0", .{platform.c.getpid()}); defer gpa.free(first); var abuf: [64]u8 = undefined; - try tmp.dir.symLink(io, try std.fmt.bufPrint(&abuf, "../{s}", .{away.sub_path}), first, .{ .is_directory = true }); + tmp.dir.symLink(io, try std.fmt.bufPrint(&abuf, "../{s}", .{away.sub_path}), first, .{ .is_directory = true }) catch |e| + return if (e == error.PermissionDenied and @import("builtin").os.tag == .windows) error.SkipZigTest else e; // symlinks need privilege on Windows var one = try rec.Recorder.init(gpa, io, dir, &.{1}); defer one.deinit(); var two = try rec.Recorder.init(gpa, io, dir, &.{1}); @@ -219,13 +221,13 @@ test "a recorder that finds every name taken gives up once, owning its name" { const dir = try std.fmt.bufPrintSentinel(&buf, ".zig-cache/tmp/{s}", .{tmp.sub_path}, 0); for (0..3) |k| { var nbuf: [32]u8 = undefined; - try tmp.dir.createDir(io, try std.fmt.bufPrint(&nbuf, "{d}-{d}", .{ std.c.getpid(), k }), .default_dir); + try tmp.dir.createDir(io, try std.fmt.bufPrint(&nbuf, "{d}-{d}", .{ platform.c.getpid(), k }), .default_dir); } try std.testing.expectError(error.PathAlreadyExists, rec.Recorder.initTries(gpa, io, dir, &.{1}, 3)); var r = try rec.Recorder.initTries(gpa, io, dir, &.{1}, 4); defer r.deinit(); var nbuf: [32]u8 = undefined; - try std.testing.expectEqualStrings(try std.fmt.bufPrint(&nbuf, "{d}-3", .{std.c.getpid()}), r.name); + try std.testing.expectEqualStrings(try std.fmt.bufPrint(&nbuf, "{d}-3", .{platform.c.getpid()}), r.name); } test "a stream's files are only ever new ones: links planted in the recorder's folder are not written through" { 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/originals.zig b/zig/src/core/originals.zig index 395ddbf0cb..fe17d2d9b1 100644 --- a/zig/src/core/originals.zig +++ b/zig/src/core/originals.zig @@ -1,5 +1,6 @@ //! A model folder's original files, cloned once before --slide first teaches it, and put back by `slide rollback`. const std = @import("std"); +const platform = @import("sys"); const disk = @import("lanes").learned_disk; const files = @import("originals_files.zig"); const Allocator = std.mem.Allocator; @@ -50,10 +51,10 @@ pub const Restored = struct { /// A process's hold on a model folder, gone when it is released or the process ends. pub const Hold = struct { - fd: std.c.fd_t, + fd: platform.c.fd_t, pub fn release(h: Hold) void { - _ = std.c.close(h.fd); + _ = platform.c.close(h.fd); } }; @@ -61,11 +62,11 @@ pub const Hold = struct { pub fn hold(a: Allocator, dir: []const u8, exclusive: bool) error{ Busy, Unlockable, OutOfMemory }!Hold { const fd = files.openFolder(try a.dupeSentinel(u8, dir, 0)); if (fd < 0) return error.Unlockable; - const kind: c_int = if (exclusive) std.c.LOCK.EX else std.c.LOCK.SH; - const rc = std.c.flock(fd, kind | std.c.LOCK.NB); + const kind: c_int = if (exclusive) platform.c.LOCK.EX else platform.c.LOCK.SH; + const rc = platform.c.flock(fd, kind | platform.c.LOCK.NB); if (rc == 0) return .{ .fd = fd }; - const busy = std.c.errno(rc) == .AGAIN; - _ = std.c.close(fd); + const busy = platform.c.errno(rc) == .AGAIN; + _ = platform.c.close(fd); return if (busy) error.Busy else error.Unlockable; } @@ -124,7 +125,7 @@ fn keepFiles(io: Io, dir: []const u8, graph: ?[]const u8, o: Options, a: Allocat } const part = try std.fs.path.join(a, &.{ dir, partial_name }); Io.Dir.cwd().deleteTree(io, part) catch {}; - if (std.c.mkdir(try a.dupeSentinel(u8, part, 0), 0o755) != 0) return error.SnapshotFolder; + if (platform.c.mkdir(try a.dupeSentinel(u8, part, 0), 0o755) != 0) return error.SnapshotFolder; errdefer Io.Dir.cwd().deleteTree(io, part) catch {}; const names = try learnables(a, io, dir); if (names.len == 0) { @@ -216,11 +217,11 @@ fn rollbackFiles(io: Io, dir: []const u8, o: Options, a: Allocator, why: *[]cons } // nothing is put back before every clone is made and the copies fit; a stop at any point leaves no temp behind errdefer for (backs) |b| { - _ = std.c.unlink(b.tmp); + _ = platform.c.unlink(b.tmp); }; var largest: u64 = 0; for (backs) |*b| { - _ = std.c.unlink(b.tmp); // an interrupted rollback's + _ = platform.c.unlink(b.tmp); // an interrupted rollback's b.cloned = try cloneIf(o, b.src, b.tmp); if (b.cloned) out.cloned += 1 else largest = @max(largest, b.size); } @@ -380,7 +381,7 @@ fn removeAdded(a: Allocator, io: Io, dir: []const u8, entries: []const Entry) ![ if ((try d.statFile(io, e.name, .{ .follow_symlinks = false })).kind == .directory) continue; try gone.append(a, try a.dupe(u8, e.name)); } - for (gone.items) |name| if (std.c.unlink(try join(a, dir, name)) != 0) return error.RemoveFailed; + for (gone.items) |name| if (platform.c.unlink(try join(a, dir, name)) != 0) return error.RemoveFailed; return gone.items; } @@ -426,7 +427,7 @@ fn removeGraphs(a: Allocator, io: Io, snap: []const u8) ![]const []const u8 { const text = Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 26)) catch continue; const v = std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}) catch continue; if (v != .object or v.object.get("facts") == null or v.object.get("next") == null) continue; - if (std.c.unlink(try a.dupeSentinel(u8, path, 0)) != 0) return error.RemoveFailed; + if (platform.c.unlink(try a.dupeSentinel(u8, path, 0)) != 0) return error.RemoveFailed; try gone.append(a, path); } return gone.items; diff --git a/zig/src/core/originals_files.zig b/zig/src/core/originals_files.zig index 5773526041..e6e2e231aa 100644 --- a/zig/src/core/originals_files.zig +++ b/zig/src/core/originals_files.zig @@ -1,16 +1,17 @@ //! The file work the snapshot and the rollback share: clones, hashed copies, whole writes and synced folders. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Allocator = std.mem.Allocator; const Sha256 = std.crypto.hash.sha2.Sha256; const chunk = 4 << 20; // bytes a read or write moves at a time -const new_flags: std.c.O = .{ .ACCMODE = .WRONLY, .CREAT = true, .EXCL = true, .CLOEXEC = true }; -const file_mode: std.c.mode_t = 0o644; +const new_flags: platform.c.O = .{ .ACCMODE = .WRONLY, .CREAT = true, .EXCL = true, .CLOEXEC = true }; +const file_mode: platform.c.mode_t = 0o644; /// The folder at `path` opened to read (a descriptor to lock or sync), negative when it cannot be. -pub fn openFolder(path: [:0]const u8) std.c.fd_t { - return std.c.open(path, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .CLOEXEC = true }, @as(std.c.mode_t, 0)); +pub fn openFolder(path: [:0]const u8) platform.c.fd_t { + return platform.c.open(path, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .CLOEXEC = true }, @as(platform.c.mode_t, 0)); } extern "c" fn clonefile(src: [*:0]const u8, dst: [*:0]const u8, flags: u32) c_int; @@ -23,22 +24,22 @@ pub fn clone(src: [:0]const u8, dst: [:0]const u8) !bool { .macos => { const rc = clonefile(src, dst, clone_nofollow); if (rc == 0) return true; - return switch (std.c.errno(rc)) { + return switch (platform.c.errno(rc)) { .OPNOTSUPP, .XDEV, .INVAL => false, else => error.CloneFailed, }; }, .linux => { - const in = std.c.open(src, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + const in = platform.c.open(src, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, @as(platform.c.mode_t, 0)); if (in < 0) return error.ReadFailed; - defer _ = std.c.close(in); - const out = std.c.open(dst, new_flags, file_mode); + defer _ = platform.c.close(in); + const out = platform.c.open(dst, new_flags, file_mode); if (out < 0) return error.WriteFailed; - const rc = std.c.ioctl(out, ficlone, in); - const e = std.c.errno(rc); - _ = std.c.close(out); + const rc = platform.c.ioctl(out, ficlone, in); + const e = platform.c.errno(rc); + _ = platform.c.close(out); if (rc == 0) return true; - _ = std.c.unlink(dst); + _ = platform.c.unlink(dst); return switch (e) { .OPNOTSUPP, .XDEV, .INVAL, .NOTTY, .NOSYS, .BADF => false, else => error.CloneFailed, @@ -50,9 +51,9 @@ pub fn clone(src: [:0]const u8, dst: [:0]const u8) !bool { /// SHA-256 of the file at `path`. pub fn hashFile(path: [:0]const u8) ![32]u8 { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.ReadFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const buf = try std.heap.page_allocator.alloc(u8, chunk); defer std.heap.page_allocator.free(buf); var h = Sha256.init(.{}); @@ -62,12 +63,12 @@ pub fn hashFile(path: [:0]const u8) ![32]u8 { /// `src` copied to the new file `dst`, synced; the SHA-256 of the bytes read from `src`. pub fn copyHashed(src: [:0]const u8, dst: [:0]const u8) ![32]u8 { - const in = std.c.open(src, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + const in = platform.c.open(src, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, @as(platform.c.mode_t, 0)); if (in < 0) return error.ReadFailed; - defer _ = std.c.close(in); - const out = std.c.open(dst, new_flags, file_mode); + defer _ = platform.c.close(in); + const out = platform.c.open(dst, new_flags, file_mode); if (out < 0) return error.WriteFailed; - defer _ = std.c.close(out); + defer _ = platform.c.close(out); const buf = try std.heap.page_allocator.alloc(u8, chunk); defer std.heap.page_allocator.free(buf); var h = Sha256.init(.{}); @@ -75,25 +76,25 @@ pub fn copyHashed(src: [:0]const u8, dst: [:0]const u8) ![32]u8 { h.update(got); try writeAll(out, got); } - if (std.c.fsync(out) != 0) return error.WriteFailed; + if (platform.c.fsync(out) != 0) return error.WriteFailed; return h.finalResult(); } /// The next bytes of `fd` in `buf`, or null at its end. -fn readSome(fd: std.c.fd_t, buf: []u8) !?[]u8 { +fn readSome(fd: platform.c.fd_t, buf: []u8) !?[]u8 { while (true) { - const n = std.c.read(fd, buf.ptr, buf.len); - if (n < 0 and std.c.errno(n) == .INTR) continue; + const n = platform.c.read(fd, buf.ptr, buf.len); + if (n < 0 and platform.c.errno(n) == .INTR) continue; if (n < 0) return error.ReadFailed; return if (n == 0) null else buf[0..@intCast(n)]; } } -fn writeAll(fd: std.c.fd_t, bytes: []const u8) !void { +fn writeAll(fd: platform.c.fd_t, bytes: []const u8) !void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes.ptr + done, bytes.len - done); - if (n < 0 and std.c.errno(n) == .INTR) continue; + const n = platform.c.write(fd, bytes.ptr + done, bytes.len - done); + if (n < 0 and platform.c.errno(n) == .INTR) continue; if (n <= 0) return error.WriteFailed; done += @intCast(n); } @@ -102,24 +103,24 @@ fn writeAll(fd: std.c.fd_t, bytes: []const u8) !void { /// `bytes` to `path` through `path`.tmp, synced, then renamed over it. pub fn writeWhole(a: Allocator, path: []const u8, bytes: []const u8) !void { const tmp = try std.fmt.allocPrintSentinel(a, "{s}.tmp", .{path}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .CLOEXEC = true }, file_mode); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .CLOEXEC = true }, file_mode); if (fd < 0) return error.WriteFailed; { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try writeAll(fd, bytes); - if (std.c.fsync(fd) != 0) return error.WriteFailed; + if (platform.c.fsync(fd) != 0) return error.WriteFailed; } try rename(tmp, try a.dupeSentinel(u8, path, 0)); } pub fn rename(from: [:0]const u8, to: [:0]const u8) !void { - if (std.c.rename(from, to) != 0) return error.RenameFailed; + if (platform.c.rename(from, to) != 0) return error.RenameFailed; } /// The folder's entries made durable (best effort: a file system that cannot sync a folder still has them). pub fn syncDir(a: Allocator, dir: []const u8) void { const fd = openFolder(a.dupeSentinel(u8, dir, 0) catch return); if (fd < 0) return; - _ = std.c.fsync(fd); - _ = std.c.close(fd); + _ = platform.c.fsync(fd); + _ = platform.c.close(fd); } diff --git a/zig/src/core/originals_test.zig b/zig/src/core/originals_test.zig index 588eb350f2..4cd2ebd822 100644 --- a/zig/src/core/originals_test.zig +++ b/zig/src/core/originals_test.zig @@ -1,5 +1,6 @@ //! The snapshot --slide keeps before a folder's first lesson, the loader passing it by, and the rollback that uses it. const std = @import("std"); +const platform = @import("sys"); const originals = @import("originals.zig"); const Allocator = std.mem.Allocator; const Io = std.Io; @@ -64,12 +65,12 @@ fn model(a: Allocator, io: Io, tmp: std.testing.TmpDir) ![]const u8 { /// A lesson's writes as the learner makes them: a shard rewritten, one written in place, the index, config.json, more. fn teach(a: Allocator, io: Io, dir: []const u8) !void { try replace(a, io, dir, shard1, try tensor(a, "layers.0.up.weight", &.{ 1, 2, 3, 4, 5, 6, 0.5, 0, 0 })); - const fd = std.c.open(try path(a, dir, shard2), .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try path(a, dir, shard2), .{ .ACCMODE = .WRONLY }, @as(platform.c.mode_t, 0)); try std.testing.expect(fd >= 0); - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const moved = [_]f32{99}; const end = (try read(a, io, try path(a, dir, shard2))).len - 4; - try std.testing.expectEqual(@as(isize, 4), std.c.pwrite(fd, std.mem.asBytes(&moved), 4, @intCast(end))); + try std.testing.expectEqual(@as(isize, 4), platform.c.pwrite(fd, std.mem.asBytes(&moved), 4, @intCast(end))); try replace(a, io, dir, "model.safetensors.index.json", index_taught); const config = "{\"model_type\": \"nemotron_h\", \"moe_shared_expert_intermediate_size\": 3776, " ++ "\"tensorfold_learned_neurons\": 64}"; @@ -105,13 +106,13 @@ fn has(text: []const u8, part: []const u8) bool { /// The first byte of `p` flipped, in place. fn flip(p: [:0]const u8) void { - const fd = std.c.open(p, .{ .ACCMODE = .RDWR }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(p, .{ .ACCMODE = .RDWR }, @as(platform.c.mode_t, 0)); if (fd < 0) return; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var b: [1]u8 = undefined; - if (std.c.pread(fd, &b, 1, 0) != 1) return; + if (platform.c.pread(fd, &b, 1, 0) != 1) return; b[0] ^= 1; - _ = std.c.pwrite(fd, &b, 1, 0); + _ = platform.c.pwrite(fd, &b, 1, 0); } fn short(_: [:0]const u8) ?u64 { diff --git a/zig/src/core/paced_read.zig b/zig/src/core/paced_read.zig index fc53af3023..2f249efd52 100644 --- a/zig/src/core/paced_read.zig +++ b/zig/src/core/paced_read.zig @@ -1,5 +1,6 @@ //! Checkpoint reads that skip the page cache and pace on the memory the kernel can give: a load never outruns 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; @@ -59,7 +60,7 @@ fn pageBytes() u64 { fn sysctl(comptime name: [:0]const u8) ?u64 { var v: u64 = 0; var len: usize = @sizeOf(u64); - if (std.c.sysctlbyname(name, &v, &len, null, 0) != 0) return null; + if (platform.c.sysctlbyname(name, &v, &len, null, 0) != 0) return null; return switch (len) { 8 => v, 4 => @as(u32, @truncate(v)), @@ -97,13 +98,13 @@ const Mac = struct { } fn pause(_: Mac, ms: u32) void { - const ts: std.c.timespec = .{ .sec = 0, .nsec = @intCast(@as(u64, ms) * std.time.ns_per_ms) }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = @intCast(@as(u64, ms) * std.time.ns_per_ms) }; + _ = platform.c.nanosleep(&ts, null); } }; /// `dest.len` bytes at `at` of `fd` (F_NOCACHE) through page-aligned reads into `buf`, each once there is room for it. -pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage, short: *Memory) !void { +pub fn read(fd: platform.c.fd_t, dest: []u8, at: u64, buf: Stage, short: *Memory) !void { const want = floor(); var done: usize = 0; while (done < dest.len) { @@ -115,7 +116,7 @@ pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage, short: *Memory) !vo const span = std.mem.alignForward(usize, skip + take, alignment); var got: usize = 0; while (got < skip + take) { // 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); } @@ -172,17 +173,17 @@ test "this Mac's counters: free and reclaimable pages within its RAM" { 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); var short: Memory = undefined; 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/picture/jpeg.zig b/zig/src/core/picture/jpeg.zig index 4cca1b6c48..331e289f85 100644 --- a/zig/src/core/picture/jpeg.zig +++ b/zig/src/core/picture/jpeg.zig @@ -1,6 +1,7 @@ //! JPEG through one pinned libjpeg-turbo, loaded at the first JPEG and run at Pillow's settings: 0.6.6's pixels. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Allocator = std.mem.Allocator; const picture = @import("picture.zig"); @@ -16,7 +17,7 @@ const needs = "JPEG images need libjpeg-turbo " ++ pinned ++ " (libturbojpeg)"; pub const missing = needs ++ ", which this server could not load; send PNG, or install it and name the file in " ++ env; pub const wrong_version = needs ++ "; the one this server found is another release, whose pixels can differ"; -const library = if (builtin.os.tag == .macos) "libturbojpeg.0.dylib" else "libturbojpeg.so.0"; +const library = if (builtin.os.tag == .macos) "libturbojpeg.0.dylib" else if (builtin.os.tag == .windows) "turbojpeg.dll" else "libturbojpeg.so.0"; // TurboJPEG 3's enums: an init type, pixel formats, parameters and error codes const init_decompress = 1; @@ -60,7 +61,7 @@ fn load(io: std.Io, a: Allocator) error{ NoJpegLibrary, JpegLibraryVersion }!Fns } fn find(io: std.Io, a: Allocator) Loaded { - if (std.c.getenv(env)) |named| return open(io, a, std.mem.span(named)) orelse .missing; + if (platform.c.getenv(env)) |named| return open(io, a, std.mem.span(named)) orelse .missing; var seen_wrong = false; if (std.process.executableDirPathAlloc(io, a)) |exe| { defer a.free(exe); @@ -84,7 +85,7 @@ fn find(io: std.Io, a: Allocator) Loaded { /// One candidate: null when it does not load or lacks TurboJPEG 3; `.wrong` when its version string is not the pin. fn open(io: std.Io, a: Allocator, path: []const u8) ?Loaded { - var dl = std.DynLib.open(path) catch return null; + var dl = platform.DynLib.open(path) catch return null; var f: Fns = undefined; const info = @typeInfo(Fns).@"struct"; inline for (info.field_names, info.field_types) |name, T| { @@ -103,9 +104,12 @@ fn open(io: std.Io, a: Allocator, path: []const u8) ?Loaded { /// The version string libjpeg-turbo builds into every binary ("libjpeg-turbo version X (build Y)") equals the pin. fn versionIs(io: std.Io, a: Allocator, symbol: *const anyopaque) bool { - var info: DlInfo = undefined; - if (dladdr(symbol, &info) == 0) return false; - const file = std.mem.span(info.fname orelse return false); + var path: [1100]u8 = undefined; + const file = if (comptime builtin.os.tag == .windows) platform.libraryOf(symbol, &path) orelse return false else blk: { + var info: DlInfo = undefined; + if (dladdr(symbol, &info) == 0) return false; + break :blk std.mem.span(info.fname orelse return false); + }; const bytes = std.Io.Dir.cwd().readFileAlloc(io, file, a, .limited(64 << 20)) catch return false; defer a.free(bytes); return versionOf(bytes) != null and std.mem.eql(u8, versionOf(bytes).?, pinned); diff --git a/zig/src/core/picture/picture_test.zig b/zig/src/core/picture/picture_test.zig index f8346d7f9c..e1fc0136ce 100644 --- a/zig/src/core/picture/picture_test.zig +++ b/zig/src/core/picture/picture_test.zig @@ -36,7 +36,7 @@ fn payload(arena: std.mem.Allocator, url: []const u8) !?[]u8 { /// JPEG cases need the pinned library: required when TENSORFOLD_LIBTURBOJPEG names one, skipped when none loads. fn jpegReady(a: std.mem.Allocator) !bool { if (picture.jpeg.available(std.testing.io, a)) return true; - if (std.testing.environ.getPosix(picture.jpeg.env) != null) return error.PinnedLibraryDidNotLoad; + if (@import("sys").getenv(picture.jpeg.env) != null) return error.PinnedLibraryDidNotLoad; std.debug.print("libjpeg-turbo {s} not found: JPEG pixel cases skipped\n", .{picture.jpeg.pinned}); return false; } diff --git a/zig/src/core/prompt_cache_fake.zig b/zig/src/core/prompt_cache_fake.zig index 6f6754d0b6..734ecafae9 100644 --- a/zig/src/core/prompt_cache_fake.zig +++ b/zig/src/core/prompt_cache_fake.zig @@ -1,5 +1,6 @@ //! A host-memory snapshot family shared by the prompt cache's tests and the learned-disk fault tests. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const imprint = @import("prompt_imprint.zig"); const Snapshots = pc.Snapshots; @@ -96,7 +97,7 @@ pub const Fake = struct { } var path: [512]u8 = undefined; - _ = std.c.unlink(file(&path, dir, key) catch return); + _ = platform.c.unlink(file(&path, dir, key) catch return); } pub fn file(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { return std.fmt.bufPrintSentinel(buf, "{s}/{x:0>16}.bin", .{ dir, key }, 0); @@ -106,17 +107,17 @@ pub const Fake = struct { f.writes += 1; if (f.fail_write) return error.WriteFailed; var path: [512]u8 = undefined; - const fd = std.c.open(try file(&path, dir, key), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try file(&path, dir, key), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.WriteFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imprint.writeAll(fd, std.mem.asBytes(@as(*State, @ptrCast(@alignCast(saved))))); } fn readFn(ptr: *anyopaque, dir: [:0]const u8, key: u64, at: u32) anyerror!Saved { const f = of(ptr); var path: [512]u8 = undefined; - const fd = std.c.open(try file(&path, dir, key), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try file(&path, dir, key), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.ReadFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const st = try f.gpa.create(State); errdefer f.gpa.destroy(st); if (!imprint.readAt(fd, std.mem.asBytes(st), 0) or st.at != at) return error.ReadFailed; diff --git a/zig/src/core/prompt_cache_learn_test.zig b/zig/src/core/prompt_cache_learn_test.zig index 73eb857f2a..26f9f3f28e 100644 --- a/zig/src/core/prompt_cache_learn_test.zig +++ b/zig/src/core/prompt_cache_learn_test.zig @@ -1,5 +1,6 @@ //! The prompt cache's learned-state tests: harness states kept on disk, the cap, the disk floor and peers' halves. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const imprint = @import("prompt_imprint.zig"); const rmTree = @import("prompt_imprint_test.zig").rmTree; @@ -14,7 +15,7 @@ test "a learned harness state outlives its store: a fresh session on a new one r defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-test-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-test-{d}", .{platform.c.getpid()}); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const harness = [_]u32{ 7, 7, 7, 7, 7, 7 }; // shared by every session: a state at 5 reads it all (lookahead 1) const key = imprint.Imprint.keyOf(&harness); @@ -50,7 +51,7 @@ test "past the learned-state cap the least recently used state is forgotten, and defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-cap-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-cap-{d}", .{platform.c.getpid()}); defer rmTree(root); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const one = [_]u32{ 7, 7, 7, 7, 7, 7 }; @@ -67,8 +68,8 @@ test "past the learned-state cap the least recently used state is forgotten, and try std.testing.expectEqual(fresh(&p2), f.pass(&s, &p2, try s.begin(a, &p2, 7, &.{5}, &.{}, null, &.{}))); try std.testing.expect(im.has(imprint.Imprint.keyOf(&two)) and !im.has(imprint.Imprint.keyOf(&one))); var path: [512]u8 = undefined; - try std.testing.expect(std.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&one))) != 0); // its file went too - _ = std.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&two))); // a file lost behind the store's back + try std.testing.expect(platform.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&one))) != 0); // its file went too + _ = platform.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&two))); // a file lost behind the store's back var g: Fake = .{ .gpa = gpa }; var t = Store.init(gpa, g.learned(), rules, 1 << 20); defer t.deinit(); @@ -87,7 +88,7 @@ test "learned states are prompt arithmetic: a request that decodes rows below on defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-spans-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-spans-{d}", .{platform.c.getpid()}); defer rmTree(root); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const harness = [_]u32{ 7, 7, 7, 7, 7, 7 }; @@ -138,7 +139,7 @@ test "Store learning preserves the disk floor, backs off writes across keys, and var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-backoff-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-backoff-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 11, 1 << 20); defer im.deinit(); @@ -180,7 +181,7 @@ test "learning reserves the existing index rewrite above the disk floor and resu var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-index-floor-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-index-floor-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 12, 1 << 20); defer im.deinit(); @@ -213,7 +214,7 @@ test "actual Store learn writes neither half on peer refusal, backs off, and rem }; const a = std.testing.allocator; var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-refuse-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-refuse-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 51, 1 << 20); defer im.deinit(); @@ -242,9 +243,9 @@ test "actual Store learn writes neither half on peer refusal, backs off, and rem try std.testing.expectEqual(@as(u64, 0), im.admission.reserved); var file_path: [512]u8 = undefined; const file = try Fake.file(&file_path, im.dir, imprint.Imprint.keyOf(prompt[0..6])); - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd >= 0) { - _ = std.c.close(fd); + _ = platform.c.close(fd); return error.OrphanedHalf; } } @@ -263,16 +264,16 @@ test "under-cap disk pressure selects the same LRU victim before either half is }; const a = std.testing.allocator; var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-victim-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-victim-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 52, 1 << 20); defer im.deinit(); try im.add(imprint.Imprint.keyOf(&.{ 1, 2 }), 1, &.{ 1, 2 }, &.{}, 16); var victim_path: [512]u8 = undefined; - const victim_fd = std.c.open(try Fake.file(&victim_path, im.dir, imprint.Imprint.keyOf(&.{ 1, 2 })), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const victim_fd = platform.c.open(try Fake.file(&victim_path, im.dir, imprint.Imprint.keyOf(&.{ 1, 2 })), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (victim_fd < 0) return error.Create; - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(victim_fd, 16)); - _ = std.c.close(victim_fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(victim_fd, 16)); + _ = platform.c.close(victim_fd); im.admission = .{ .floor = 100, .free = Disk.free, .clock = Disk.clock }; const prompt = [_]u32{ 7, 7, 7, 7, 7, 7, 9, 10 }; Disk.space = 100 + 105 + 256 + (6 * 4) + (try im.indexScratchBytes()) - 8; diff --git a/zig/src/core/prompt_imprint.zig b/zig/src/core/prompt_imprint.zig index 8acf83fe21..d3f4b96564 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; } /// A learned state's file under `imprint.dir`, read whole into a count (the single-rank reclaim). @@ -491,19 +492,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/rows_file.zig b/zig/src/core/rows_file.zig index 4087f96a94..7b98d03d78 100644 --- a/zig/src/core/rows_file.zig +++ b/zig/src/core/rows_file.zig @@ -1,5 +1,6 @@ //! Row block and state files: payload, a state's tokens, then a trailer; read past the page cache and checked first. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Name = @import("rows_chain.zig").Name; @@ -51,9 +52,9 @@ pub const Writer = struct { var w: Writer = .{ .fd = -1 }; const t = try path(&w.tmp, dir, stem, ext, true); w.tmp_len = t.len; - w.fd = std.c.open(t, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + w.fd = platform.c.open(t, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true, .CLOEXEC = true }, @as(platform.c.mode_t, 0o600)); if (w.fd < 0) return error.RowsWrite; - if (builtin.os.tag == .macos) _ = std.c.fcntl(w.fd, std.c.F.NOCACHE, @as(c_int, 1)); + if (builtin.os.tag == .macos) _ = platform.c.fcntl(w.fd, platform.c.F.NOCACHE, @as(c_int, 1)); return w; } @@ -81,18 +82,18 @@ pub const Writer = struct { } trailer.meta_hash = mh.final(); try put(w.fd, std.mem.asBytes(&trailer)); - if (sync and std.c.fsync(w.fd) != 0) return error.RowsWrite; + if (sync and platform.c.fsync(w.fd) != 0) return error.RowsWrite; if (builtin.os.tag == .linux) { - _ = std.c.fdatasync(w.fd); + _ = platform.c.fdatasync(w.fd); _ = posix_fadvise(w.fd, 0, 0, 4); // POSIX_FADV_DONTNEED: written pages leave the cache } - _ = std.c.close(w.fd); + _ = platform.c.close(w.fd); w.fd = -1; const t_path = w.tmpPath(); var dest: [1200]u8 = undefined; const d = try std.fmt.bufPrintSentinel(&dest, "{s}", .{t_path[0 .. t_path.len - ".part".len]}, 0); - if (std.c.rename(t_path, d) != 0) { - _ = std.c.unlink(t_path); + if (platform.c.rename(t_path, d) != 0) { + _ = platform.c.unlink(t_path); return error.RowsWrite; } return trailer.payload + trailer.meta + @sizeOf(Trailer); @@ -100,18 +101,18 @@ pub const Writer = struct { /// A write given up: the `.part` file goes. pub fn abort(w: *Writer) void { - if (w.fd >= 0) _ = std.c.close(w.fd); + if (w.fd >= 0) _ = platform.c.close(w.fd); w.fd = -1; - _ = std.c.unlink(w.tmpPath()); + _ = platform.c.unlink(w.tmpPath()); } }; 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) { - if (n < 0 and std.c.errno(n) == .INTR) continue; + if (n < 0 and platform.c.errno(n) == .INTR) continue; return error.RowsWrite; } done += @intCast(n); @@ -121,8 +122,8 @@ fn put(fd: c_int, bytes: []const u8) !void { fn get(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, @min(dest.len - done, 1 << 30), @intCast(at + done)); - if (n < 0 and std.c.errno(n) == .INTR) continue; + const n = platform.c.pread(fd, dest.ptr + done, @min(dest.len - done, 1 << 30), @intCast(at + done)); + if (n < 0 and platform.c.errno(n) == .INTR) continue; if (n <= 0) return false; done += @intCast(n); } @@ -135,7 +136,7 @@ pub const Opened = struct { trailer: Trailer, pub fn close(o: *Opened) void { - _ = std.c.close(o.fd); + _ = platform.c.close(o.fd); o.fd = -1; } }; @@ -143,19 +144,19 @@ pub const Opened = struct { /// `file` opened past the page cache, its trailer read and checked against its size, kind and (when given) its name. pub fn open(file: [:0]const u8, kind: Kind, name: ?Name) !Opened { var fd: c_int = -1; - if (builtin.os.tag == .linux) fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .NOFOLLOW = true, .DIRECT = true }, @as(std.c.mode_t, 0)); - if (fd < 0) fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + if (builtin.os.tag == .linux) fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .NOFOLLOW = true, .DIRECT = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.RowsMissing; - errdefer _ = std.c.close(fd); - if (builtin.os.tag == .macos) _ = std.c.fcntl(fd, std.c.F.NOCACHE, @as(c_int, 1)); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + errdefer _ = platform.c.close(fd); + if (builtin.os.tag == .macos) _ = platform.c.fcntl(fd, platform.c.F.NOCACHE, @as(c_int, 1)); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < @sizeOf(Trailer)) return error.RowsRead; var t: Trailer = undefined; var raw: [@sizeOf(Trailer) + 2 * 4096]u8 align(4096) = undefined; // O_DIRECT: an aligned read around the trailer const at: u64 = @as(u64, @intCast(end)) - @sizeOf(Trailer); const lo = std.mem.alignBackward(u64, at, 4096); const span = std.mem.alignForward(u64, @as(u64, @intCast(end)) - lo, 4096); - const n = std.c.pread(fd, &raw, @intCast(span), @intCast(lo)); + const n = platform.c.pread(fd, &raw, @intCast(span), @intCast(lo)); if (n < @as(isize, @intCast(@as(u64, @intCast(end)) - lo))) return error.RowsRead; @memcpy(std.mem.asBytes(&t), raw[@intCast(at - lo)..][0..@sizeOf(Trailer)]); if (t.magic != MAGIC or t.version != VERSION or t.kind != @backingInt(kind)) return error.RowsRead; @@ -172,8 +173,8 @@ pub fn payload(o: *const Opened, into: Stage) ![]u8 { var done: u64 = 0; while (done < n) { // past the payload the read may run into the metadata or come back short at the file's end const from = std.mem.alignBackward(u64, done, 4096); // a short O_DIRECT read resumes on an aligned boundary - const r = std.c.pread(o.fd, into.ptr + from, @intCast(@min(span - from, 1 << 30)), @intCast(from)); - if (r < 0 and std.c.errno(r) == .INTR) continue; + const r = platform.c.pread(o.fd, into.ptr + from, @intCast(@min(span - from, 1 << 30)), @intCast(from)); + if (r < 0 and platform.c.errno(r) == .INTR) continue; if (r <= 0) return error.RowsRead; const next = from + @as(u64, @intCast(r)); if (next <= done) return error.RowsRead; @@ -206,8 +207,8 @@ fn readUnaligned(fd: c_int, dest: []u8, at: u64) bool { const span = std.mem.alignForward(usize, skip + want, 4096); var got: usize = 0; while (got < skip + want) { - const n = std.c.pread(fd, bounce[got..].ptr, span - got, @intCast(lo + got)); - if (n < 0 and std.c.errno(n) == .INTR) continue; + const n = platform.c.pread(fd, bounce[got..].ptr, span - got, @intCast(lo + got)); + if (n < 0 and platform.c.errno(n) == .INTR) continue; if (n <= 0) return false; got += @intCast(n); } @@ -225,9 +226,9 @@ pub fn stage(a: std.mem.Allocator, bytes: u64) !Stage { test "a block file reads back byte for byte, and a flipped byte, a cut tail or another name is refused" { const a = std.testing.allocator; var dir_buf: [64]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&dir_buf, "/tmp/tf-rows-file-{d}", .{std.c.getpid()}, 0); - _ = std.c.mkdir(dir, 0o700); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&dir_buf, "/tmp/tf-rows-file-{d}", .{platform.c.getpid()}, 0); + _ = platform.c.mkdir(dir, 0o700); + defer _ = platform.c.rmdir(dir); const body = try a.alloc(u8, 70_000); defer a.free(body); for (body, 0..) |*b, i| b.* = @truncate(i *% 31 +% 7); @@ -239,7 +240,7 @@ test "a block file reads back byte for byte, and a flipped byte, a cut tail or a try std.testing.expectEqual(@as(u64, body.len + 96), bytes); var buf: [128]u8 = undefined; const file = try path(&buf, dir, "b", ".blk", false); - defer _ = std.c.unlink(file); + defer _ = platform.c.unlink(file); const s = try stage(a, body.len); defer a.free(s); { @@ -249,32 +250,32 @@ test "a block file reads back byte for byte, and a flipped byte, a cut tail or a } try std.testing.expectError(error.RowsRead, open(file, .block, @splat(6))); try std.testing.expectError(error.RowsRead, open(file, .state, null)); - 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)); try std.testing.expect(fd >= 0); - defer _ = std.c.close(fd); - try std.testing.expectEqual(@as(isize, 1), std.c.pwrite(fd, "X", 1, 12345)); + defer _ = platform.c.close(fd); + try std.testing.expectEqual(@as(isize, 1), platform.c.pwrite(fd, "X", 1, 12345)); { var o = try open(file, .block, name); defer o.close(); try std.testing.expectError(error.RowsHash, payload(&o, s)); } - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 5000)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 5000)); try std.testing.expectError(error.RowsRead, open(file, .block, name)); } test "a state file's metadata reads back checked, and an aborted write leaves nothing" { const a = std.testing.allocator; var dir_buf: [64]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&dir_buf, "/tmp/tf-rows-meta-{d}", .{std.c.getpid()}, 0); - _ = std.c.mkdir(dir, 0o700); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&dir_buf, "/tmp/tf-rows-meta-{d}", .{platform.c.getpid()}, 0); + _ = platform.c.mkdir(dir, 0o700); + defer _ = platform.c.rmdir(dir); const tokens = [_]u32{ 1, 2, 3, 4, 5 }; var w = try Writer.open(dir, "s", ".st"); try w.append("own bytes"); _ = try w.finish(.{ .kind = @backingInt(Kind.state), .payload = 0, .payload_hash = 0, .name = @splat(1), .at = 4 }, &.{std.mem.sliceAsBytes(&tokens)}, true); var buf: [128]u8 = undefined; const file = try path(&buf, dir, "s", ".st", false); - defer _ = std.c.unlink(file); + defer _ = platform.c.unlink(file); var o = try open(file, .state, null); defer o.close(); try std.testing.expectEqual(@as(u32, 4), o.trailer.at); diff --git a/zig/src/core/rows_store.zig b/zig/src/core/rows_store.zig index ec84675869..1434501d4c 100644 --- a/zig/src/core/rows_store.zig +++ b/zig/src/core/rows_store.zig @@ -1,5 +1,6 @@ //! Conversation states on disk (--learn-conversations): rows in blocks states share, each state's own parts in a file. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const pc = @import("prompt_cache.zig"); const modes = @import("cache_modes.zig"); @@ -89,7 +90,7 @@ pub const Rows = struct { try r.scan(); r.settle(); r.stamp(); - r.verify = if (std.c.getenv("TF_IMPRINT_VERIFY")) |v| v[0] == '1' else false; + r.verify = if (platform.c.getenv("TF_IMPRINT_VERIFY")) |v| v[0] == '1' else false; return r; } @@ -307,11 +308,11 @@ pub const Rows = struct { const today = dayNow(); var buf: [1100]u8 = undefined; const p = std.fmt.bufPrintSentinel(&buf, "{s}/written", .{r.root}, 0) catch return .{ .day = today, .bytes = std.math.maxInt(u64) }; - const fd = std.c.open(p, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(p, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return .{ .day = today, .bytes = 0 }; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var d: Day = undefined; - if (std.c.pread(fd, std.mem.asBytes(&d), @sizeOf(Day), 0) != @sizeOf(Day) or d.day != today) return .{ .day = today, .bytes = 0 }; + if (platform.c.pread(fd, std.mem.asBytes(&d), @sizeOf(Day), 0) != @sizeOf(Day) or d.day != today) return .{ .day = today, .bytes = 0 }; return d; } @@ -323,10 +324,10 @@ pub const Rows = struct { d.bytes +|= bytes; var buf: [1100]u8 = undefined; const p = std.fmt.bufPrintSentinel(&buf, "{s}/written", .{r.root}, 0) catch return; - const fd = std.c.open(p, .{ .ACCMODE = .WRONLY, .CREAT = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(p, .{ .ACCMODE = .WRONLY, .CREAT = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return; - defer _ = std.c.close(fd); - _ = std.c.pwrite(fd, std.mem.asBytes(&d), @sizeOf(Day), 0); + defer _ = platform.c.close(fd); + _ = platform.c.pwrite(fd, std.mem.asBytes(&d), @sizeOf(Day), 0); } /// Every state and block of this identity forgotten (DELETE /v1/imprint); the queue too. The bytes freed. @@ -348,15 +349,15 @@ pub const Rows = struct { /// The files under the store's directory: states read back and checked by name, blocks by size; leftovers go. fn scan(r: *Rows) !void { - const d = std.c.opendir(r.dir) orelse return error.RowsDir; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { - if (ent.type != std.c.DT.REG) continue; + const d = platform.c.opendir(r.dir) orelse return error.RowsDir; + 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); var buf: [1200]u8 = undefined; const p = std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ r.dir, name }, 0) catch continue; if (std.mem.endsWith(u8, name, ".part")) { - _ = std.c.unlink(p); + _ = platform.c.unlink(p); continue; } const n = parseName(name) orelse continue; @@ -365,7 +366,7 @@ pub const Rows = struct { try r.blocks.put(r.gpa, n, .{ .bytes = bytes }); } else if (std.mem.endsWith(u8, name, ".st")) { r.readState(p) catch { - _ = std.c.unlink(p); + _ = platform.c.unlink(p); }; } } @@ -448,8 +449,8 @@ pub fn withAt(a: Allocator, starts: []const u32, at: u32, planned: bool) ![]u32 } fn dayNow() u64 { - var t: std.c.timespec = undefined; - if (std.c.clock_gettime(std.c.CLOCK.REALTIME, &t) != 0 or t.sec < 0) return 0; + var t: platform.c.timespec = undefined; + if (platform.c.clock_gettime(platform.c.CLOCK.REALTIME, &t) != 0 or t.sec < 0) return 0; return @as(u64, @intCast(t.sec)) / 86400; } @@ -465,28 +466,28 @@ fn parseName(name: []const u8) ?chain.Stem { pub fn unlinkName(dir: []const u8, n: anytype, ext: []const u8) void { var buf: [1200]u8 = undefined; const p = file.path(&buf, dir, &chain.hex(n), ext, false) catch return; - _ = std.c.unlink(p); + _ = platform.c.unlink(p); } fn readWord(p: [:0]const u8) u64 { - const fd = std.c.open(p, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(p, .{ .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 (std.c.pread(fd, std.mem.asBytes(&w), 8, 0) == 8) w else 0; + return if (platform.c.pread(fd, std.mem.asBytes(&w), 8, 0) == 8) w else 0; } fn writeWord(p: [:0]const u8, w: u64) void { var buf: [1200]u8 = undefined; const part = std.fmt.bufPrintSentinel(&buf, "{s}.part", .{p}, 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); - if (std.c.write(fd, std.mem.asBytes(&w), 8) != 8) { - _ = std.c.unlink(part); + defer _ = platform.c.close(fd); + if (platform.c.write(fd, std.mem.asBytes(&w), 8) != 8) { + _ = platform.c.unlink(part); return; } - _ = std.c.rename(part, p); + _ = platform.c.rename(part, p); } /// `dir` and its parents, 0700. @@ -495,11 +496,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.RowsDir; + const rc = platform.c.mkdir(tmp.ptr, 0o700); + if (rc != 0 and platform.c.errno(rc) != .EXIST) return error.RowsDir; } test { diff --git a/zig/src/core/rows_test.zig b/zig/src/core/rows_test.zig index f1b4a31af8..a5426ffdaf 100644 --- a/zig/src/core/rows_test.zig +++ b/zig/src/core/rows_test.zig @@ -1,5 +1,6 @@ //! The rows store's tests: states written behind and read back equal, shared blocks, checks, the cap, budget and floor. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const rows = @import("rows_store.zig"); const Fake = @import("rows_fake.zig").Fake; @@ -10,7 +11,7 @@ const Allocator = std.mem.Allocator; const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; fn root(buf: []u8, tag: []const u8) []const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-rows-{s}-{d}", .{ tag, std.c.getpid() }) catch unreachable; + return std.fmt.bufPrint(buf, "/tmp/tf-rows-{s}-{d}", .{ tag, platform.c.getpid() }) catch unreachable; } fn options(dir: []const u8) rows.Options { @@ -126,10 +127,10 @@ test "a block that fails its check is forgotten with every state naming it, and var it = r.blocks.keyIterator(); var path: [1200]u8 = undefined; const file = try rows.file.path(&path, r.dir, &rows.chain.hex(it.next().?.*), ".blk", false); - 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)); try std.testing.expect(fd >= 0); - defer _ = std.c.close(fd); - try std.testing.expectEqual(@as(isize, 1), std.c.pwrite(fd, "X", 1, 100)); + defer _ = platform.c.close(fd); + try std.testing.expectEqual(@as(isize, 1), platform.c.pwrite(fd, "X", 1, 100)); } var f: Fake = .{ .gpa = gpa }; const r = try rows.Rows.open(gpa, f.snapshots(), options(dir)); @@ -373,10 +374,10 @@ test "a learn dir that is a regular file is refused, and open frees what it made const path = root(&buf, "file"); var zbuf: [129]u8 = undefined; const z = try std.fmt.bufPrintSentinel(&zbuf, "{s}", .{path}, 0); - const fd = std.c.open(z, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(z, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); try std.testing.expect(fd >= 0); - _ = std.c.close(fd); - defer _ = std.c.unlink(z); + _ = platform.c.close(fd); + defer _ = platform.c.unlink(z); var f: Fake = .{ .gpa = gpa }; try std.testing.expectError(error.RowsDir, rows.Rows.open(gpa, f.snapshots(), options(path))); } diff --git a/zig/src/core/rows_write.zig b/zig/src/core/rows_write.zig index 96d3bf81da..8e37173e4c 100644 --- a/zig/src/core/rows_write.zig +++ b/zig/src/core/rows_write.zig @@ -1,5 +1,6 @@ //! Kept states written behind: a block (or a block's worth of a state's own file) a step, only while the engine idles. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const image = @import("rows_image.zig"); const chain = @import("rows_chain.zig"); @@ -86,8 +87,8 @@ pub fn drain(r: *Rows, deadline_ns: i128) usize { } pub fn nowNs() i128 { - var t: std.c.timespec = undefined; - _ = std.c.clock_gettime(std.c.CLOCK.MONOTONIC, &t); + var t: platform.c.timespec = undefined; + _ = platform.c.clock_gettime(platform.c.CLOCK.MONOTONIC, &t); return @as(i128, t.sec) * std.time.ns_per_s + t.nsec; } 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/core/tfvideo.zig b/zig/src/core/tfvideo.zig index 8d62ea36bc..739c14282d 100644 --- a/zig/src/core/tfvideo.zig +++ b/zig/src/core/tfvideo.zig @@ -30,6 +30,8 @@ const Fns = struct { /// Beside the server (zig-out/native/lib, where `zig build video` installs it), then the loader's own search path. const names: []const [:0]const u8 = if (builtin.os.tag == .macos) &.{ "@executable_path/../lib/libtfvideo.dylib", "libtfvideo.dylib" } +else if (builtin.os.tag == .windows) + &.{"tfvideo.dll"} else &.{"libtfvideo.so"}; @@ -44,7 +46,7 @@ fn load() ?Fns { if (lib_failed) return null; lib_failed = true; var dl = for (names) |name| { - break std.DynLib.open(name) catch continue; + break @import("sys").DynLib.open(name) catch continue; } else return null; var f: Fns = undefined; const info = @typeInfo(Fns).@"struct"; diff --git a/zig/src/cuda/cublaslt.zig b/zig/src/cuda/cublaslt.zig index 1e0ba48ae4..618b4ebbb8 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; @@ -60,12 +61,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, @@ -75,7 +77,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 72e33edc92..e81d71fff2 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 18542b8cb4..efab27542c 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -26,6 +26,7 @@ pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; pub const lane_gemv: []const u8 = if (available) &Blob("fatbin_lane_gemv").bytes else &.{}; pub const affine4_pack: []const u8 = if (available) &Blob("fatbin_affine4_pack").bytes else &.{}; +pub const affine8: []const u8 = if (available) &Blob("fatbin_affine8").bytes else &.{}; pub const fp8_experts: []const u8 = if (available) &Blob("fatbin_fp8_experts").bytes else &.{}; pub const exlr8: []const u8 = if (available) &Blob("fatbin_exlr8").bytes else &.{}; pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").bytes else &.{}; @@ -33,6 +34,9 @@ pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts pub const rank_sum: []const u8 = if (available) &Blob("fatbin_rank_sum").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 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 b0d527ce75..8d0877fe6a 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; @@ -32,7 +33,7 @@ pub const Api = struct { }; pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, pub fn open() Error!Library { @@ -40,7 +41,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/qlinear.zig b/zig/src/cuda/qlinear.zig index dbb86bb8d0..839d38b3e9 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -30,6 +30,26 @@ pub const Affine4 = struct { } }; +/// MLX affine 8-bit tiled for affine8.cu's lane_gemv: a byte a weight, Affine4's padding, scales and biases. +pub const Affine8 = struct { + w: u64, + s: u64, + b: u64, + n: usize, + k: usize, + npad: usize, + + pub fn layout(n: usize, k: usize) struct { npad: usize, words: usize, scales: usize } { + const npad = (n + 127) / 128 * 128; + return .{ .npad = npad, .words = npad * k, .scales = k / 64 * npad * 2 }; + } +}; + +/// The width an MLX code of `bits` is served at: 2 to 4 bits on the 4-bit tiles, 5 to 8 on the 8-bit ones. +pub fn servedBits(bits: u8) u8 { + return if (bits <= 4) 4 else 8; +} + /// One projection's weights in any format this module serves; every one but affine4 is a lane matmul mode. pub const Weight = union(enum) { affine4: Affine4, @@ -349,6 +369,114 @@ pub const Affine4Kernels = struct { } }; +pub const gemv8_smem: u32 = 38784; // affine8.cu's Tile<64, 64, 6>::SMEM +pub const split8_smem: u32 = 25856; // Tile<64, 64, 4>::SMEM + +/// MLX affine codes of 5 to 8 bits on 8-bit tiles with lane_gemv's arithmetic (2 to 4 bits use Affine4Kernels). +pub const Affine8Kernels = struct { + gemv_fn: Function, + split_fn: Function, + pack4_fn: Function, // MLX codes of 2 to 4 bits into the affine-4 tiles + pack8_fn: Function, // MLX codes of 5 to 8 bits into the 8-bit tiles + widen_fn: Function, // MLX rows into MLX rows of 4 or 8-bit codes (token embeddings) + scales_fn: Function, // affine4_pack.cu's tf_transpose_pad16 + gemv_blocks: usize, + split: ?Split, + pdl: bool, + + /// From loaded affine8 and affine4_pack modules; GB10 decodes without the split form, as Affine4Kernels does. + pub fn resolve(d: *const Driver, mod: Module, pack_mod: Module, sms: usize, gb10: bool) !Affine8Kernels { + var a: Affine8Kernels = undefined; + a.gemv_fn = try mod.function("tf_affine8_gemv"); + a.split_fn = try mod.function("tf_affine8_split"); + a.pack4_fn = try mod.function("tf_affine8_pack4"); + a.pack8_fn = try mod.function("tf_affine8_pack8"); + a.widen_fn = try mod.function("tf_affine8_widen"); + a.scales_fn = try pack_mod.function(symbols.pack_scales); + try a.gemv_fn.allowDynamicShared(gemv8_smem); + try a.split_fn.allowDynamicShared(split8_smem); + a.gemv_blocks = @max(1, try a.gemv_fn.occupancy(128, gemv8_smem)) * sms; + a.pdl = gb10; + a.split = if (gb10) null else try Split.init(d); + return a; + } + + pub fn deinit(a: *Affine8Kernels) void { + if (a.split) |*sp| sp.deinit(); + } + + /// MLX words (n, k * bits / 32) on the device into the tiles at `w` of the width `servedBits(bits)` serves. + pub fn packWords(a: *const Affine8Kernels, s: Stream, words: u64, bits: u8, n: usize, k: usize, w: u64) !void { + if (k % 64 != 0 or k * bits % 32 != 0) return error.UnsupportedShape; + const kg = k / 64; + const npad = (n + 127) / 128 * 128; + const wide = servedBits(bits); + const total: u64 = @as(u64, npad / 64) * kg * (if (wide == 4) @as(u64, 512) else 1024); + var g: launch_.Args = .{}; + g.add(words); + for ([_]usize{ bits, n, k * bits / 32, kg }) |v| g.add(int(v)); + g.add(w); + g.add(@as(c_longlong, @intCast(total))); + const f = if (wide == 4) a.pack4_fn else a.pack8_fn; + try launch_.launch(f, .{ .grid = .{ .x = @intCast((total + 255) / 256) }, .block = .{ .x = 256 } }, s, &g); + } + + /// (n, k/64) scales or biases into the tiles' (k/64, npad) layout. + pub fn packTable(a: *const Affine8Kernels, s: Stream, src: u64, n: usize, k: usize, npad: usize, out: u64) !void { + const kg = k / 64; + var t: launch_.Args = .{}; + t.add(src); + for ([_]usize{ n, kg, npad }) |v| t.add(int(v)); + t.add(out); + const cells = @as(u64, kg) * npad; + try launch_.launch(a.scales_fn, .{ .grid = .{ .x = @intCast((cells + 255) / 256) }, .block = .{ .x = 256 } }, s, &t); + } + + /// MLX rows (n, k) of `bits` codes into MLX rows of `wide`-bit codes at `out`, every code's value kept. + pub fn widen(a: *const Affine8Kernels, s: Stream, words: u64, bits: u8, n: usize, k: usize, out: u64, wide: u8) !void { + if (wide != 4 and wide != 8) return error.UnsupportedWidth; + const total: u64 = @as(u64, n) * k * wide / 32; + var g: launch_.Args = .{}; + g.add(words); + g.add(int(bits)); + g.add(int(k)); + g.add(out); + g.add(int(wide)); + g.add(@as(c_longlong, @intCast(total))); + try launch_.launch(a.widen_fn, .{ .grid = .{ .x = @intCast((total + 255) / 256) }, .block = .{ .x = 256 } }, s, &g); + } + + /// x (rows, k) bf16 with group sums xs -> out (rows, n) bf16; a row's bits the same at every row count up to 16. + pub fn decode(a: *const Affine8Kernels, s: Stream, x: u64, xs: u64, q: Affine8, out: u64, rows: usize) !void { + if (rows > 16) return error.WindowTooWide; + const sk = splitK(q.n, q.k); + const tiles = (q.n + 63) / 64; + const part: GemvPart = .{ .w = q.w, .scales = q.s, .biases = q.b, .out = out, .n = int(q.n), .npad = int(q.npad), .sk = int(sk), .tiles = int(tiles) }; + var g: launch_.Args = .{}; + g.add(x); + g.add(xs); + g.add(part); + for ([_]usize{ rows, q.k, q.k }) |v| g.add(int(v)); + if (sk > 1) if (a.split) |sp| { + if (tiles > Split.tiles_max or tiles * sk * 16 * 64 * 4 > sp.work.len) return error.SplitTooWide; + g.add(sp.work.ptr); + g.add(sp.tickets.ptr); + return launch_.launch(a.split_fn, .{ .grid = .{ .x = u(tiles * sk) }, .block = .{ .x = 128 }, .shared = split8_smem }, s, &g); + }; + const cfg: launch_.Config = .{ .grid = .{ .x = u(@min(tiles, a.gemv_blocks)) }, .block = .{ .x = 128 }, .shared = gemv8_smem, .pdl = a.pdl }; + try launch_.launch(a.gemv_fn, cfg, s, &g); + } +}; + +test "the 8-bit tiles hold a byte a weight over Affine4's padding" { + const l = Affine8.layout(10304, 2688); + try std.testing.expectEqual(@as(usize, 10368), l.npad); + try std.testing.expectEqual(@as(usize, 10368 * 2688), l.words); + try std.testing.expectEqual(Affine4.layout(10304, 2688).scales, l.scales); + for ([_]u8{ 2, 3, 4 }) |b| try std.testing.expectEqual(@as(u8, 4), servedBits(b)); + for ([_]u8{ 5, 6, 8 }) |b| try std.testing.expectEqual(@as(u8, 8), servedBits(b)); +} + test "split_k follows the Python shapes" { try std.testing.expectEqual(@as(usize, 2), splitK(10304, 2688)); try std.testing.expectEqual(@as(usize, 8), splitK(2688, 4096)); diff --git a/zig/src/cuda/tp_watch.zig b/zig/src/cuda/tp_watch.zig index f912ba0ec8..b0b127bc8b 100644 --- a/zig/src/cuda/tp_watch.zig +++ b/zig/src/cuda/tp_watch.zig @@ -1,6 +1,7 @@ //! Ranks' liveness lines: a TCP line from each rank to rank 0, open before NCCL, carrying ready and error records. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = @import("sys").posix; /// One record on a line: ready, or an error and its name. const Record = extern struct { kind: u32, rank: u32, text: [56]u8 = @splat(0) }; @@ -24,7 +25,7 @@ pub fn ended() void { /// NCCL_COMM_ID's host (IPv4) at its port plus one. fn address() !posix.sockaddr.in { - const id = std.mem.span(std.c.getenv("NCCL_COMM_ID") orelse return error.NoCommId); + const id = std.mem.span(platform.c.getenv("NCCL_COMM_ID") orelse return error.NoCommId); const colon = std.mem.lastIndexOfScalar(u8, id, ':') orelse return error.NoCommId; const port = (std.fmt.parseInt(u16, id[colon + 1 ..], 10) catch return error.NoCommId) +% 1; const ip = std.Io.net.IpAddress.parse(id[0..colon], port) catch return error.NoCommId; @@ -127,7 +128,7 @@ pub fn fail(name: []const u8) void { /// Ends this process at once: a stream may hold kernels that wait on the gone rank, which a normal exit would join. fn stop() noreturn { - std.c._exit(1); + platform.c._exit(1); } /// A line's records; it closing or carrying an error stops this rank unless it ends on purpose. diff --git a/zig/src/fabric/fake.zig b/zig/src/fabric/fake.zig index 0b050f8b2f..607e0051af 100644 --- a/zig/src/fabric/fake.zig +++ b/zig/src/fabric/fake.zig @@ -53,7 +53,8 @@ pub const Cluster = struct { .prng = .init(0x7f4a), }; for (c.windows, c.access, c.endpoints, 0..) |*w, *a, *e, i| { - w.* = try std.posix.mmap(null, window_bytes, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + w.* = try std.heap.page_allocator.alignedAlloc(u8, .fromByteUnits(rdma.page), window_bytes); + @memset(w.*, 0); // as an anonymous map comes a.* = .{}; e.* = .{ .cluster = c, .me = @intCast(i) }; } @@ -67,7 +68,7 @@ pub const Cluster = struct { } pub fn deinit(c: *Cluster) void { - for (c.windows) |w| std.posix.munmap(w); + for (c.windows) |w| std.heap.page_allocator.free(w); for (c.queues) |*q| { for (q.items) |op| if (op.kind == .write) c.gpa.free(op.bytes); q.deinit(c.gpa); 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..c619ecb752 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(platform.c.CLOCK.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(platform.c.CLOCK.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 fa9e9a02d2..72ab29e3c7 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 lanes = @import("lanes"); const fz = @import("replay.zig"); @@ -156,8 +157,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; @@ -182,8 +183,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 } @@ -216,7 +217,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; @@ -439,7 +440,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) |t| try r.held.appendSlice(arena, &.{ t.wbuf, t.one, t.pick_tmp }); // every pair round uses them @@ -511,8 +512,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; }; @@ -529,7 +530,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 562f2dc905..42c476da87 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 7b24f4701d..9e1ab16f2a 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); } @@ -936,7 +937,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 0c4b1a8f1d..fa5b098993 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 50c0285c40..4c0834b0e9 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 b29b9d0219..21294d5a00 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"); @@ -103,7 +104,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, null); + return loadWith(gpa, dir, cap, if (platform.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false, null); } /// `load` with expert parallel over `ep_path`'s link (one Mac when null) and `vision`'s image tower and rows. @@ -112,25 +113,25 @@ 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.vision = 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; - e.sampling = try fwd.envSampling(std.c.getenv); - 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; + e.sampling = try fwd.envSampling(platform.c.getenv); + 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(); @@ -149,17 +150,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; @@ -173,7 +174,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; // --vision: the image tower and its rows count against the limit beside the weights and caches @@ -189,7 +190,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); @@ -209,9 +210,9 @@ pub const Engine = struct { return error.OverMemoryLimit; } { // short of free memory: drop the checkpoint's own cached pages before the kernel takes other files' (#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 + tower)) |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); @@ -233,13 +234,13 @@ pub const Engine = struct { errdefer if (e.vision) |v| v.deinit(); residency.wire(e); // GLM_RESIDENCY's opt-in set, then the idle keepalive's target // 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; } @@ -330,7 +331,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(); @@ -385,7 +386,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; @@ -432,7 +433,7 @@ pub const Engine = struct { } x.draws = null; // the decode loop re-arms the sampled draw each round; the MTP head drafts greedily 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]; @@ -445,7 +446,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; @@ -467,7 +468,7 @@ pub const Engine = struct { var last_end: f64 = 0; var round_draws: fwd.Draws = .{}; // the current decode round's rules, rebuilt each round 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); @@ -566,7 +567,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 b30e02c6e1..e732a14778 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; @@ -56,7 +57,7 @@ pub const Command = union(enum) { request: Request, lanes: struct { kind: u32, w fn ticks() u64 { if (@import("builtin").is_test and @import("builtin").os.tag != .macos) return @intCast(@divTrunc(std.Io.Clock.awake.now(std.testing.io).toNanoseconds() * 24, 1000)); - 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 { @@ -98,8 +99,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(); } } @@ -297,8 +298,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 9e214c3b13..02d9d3be7d 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 333cea1de6..7b7bcd9d73 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 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); @@ -35,7 +36,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 156a4bda34..504b7de04f 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; @@ -42,7 +43,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 b325424805..b08e362f22 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 f9efec49ce..29c3a9a93a 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"); @@ -78,7 +79,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; } @@ -252,7 +253,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); @@ -280,7 +281,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 3a22c38b98..5d39ab59af 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), @@ -138,15 +139,15 @@ const Pool = struct { /// Every copy on `threads` threads, uncached, paced by the memory the kernel can give (else a load fills the cache). 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); @@ -164,7 +165,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/lfm2/image_test.zig b/zig/src/families/lfm2/image_test.zig index 635e7944bc..94ba22bca9 100644 --- a/zig/src/families/lfm2/image_test.zig +++ b/zig/src/families/lfm2/image_test.zig @@ -9,7 +9,7 @@ const jpeg = core.picture.jpeg; /// JPEG cases need the pinned libjpeg-turbo: required when TENSORFOLD_LIBTURBOJPEG names one, skipped when none loads. fn jpegReady() !bool { if (jpeg.available(std.testing.io, std.testing.allocator)) return true; - if (std.testing.environ.getPosix(jpeg.env) != null) return error.PinnedLibraryDidNotLoad; + if (@import("sys").getenv(jpeg.env) != null) return error.PinnedLibraryDidNotLoad; std.debug.print("libjpeg-turbo {s} not found: LFM2's JPEG cases skipped\n", .{jpeg.pinned}); return false; } diff --git a/zig/src/families/lfm2_moe/backend.zig b/zig/src/families/lfm2_moe/backend.zig index 01b742716d..c90e2753a1 100644 --- a/zig/src/families/lfm2_moe/backend.zig +++ b/zig/src/families/lfm2_moe/backend.zig @@ -1,5 +1,6 @@ //! LFM2 conv state, windows and keyed draws behind the existing lane engine's backend interface. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const c = @import("config.zig"); @@ -69,11 +70,11 @@ pub const Metal = struct { for (&self.costs, 1..) |*cost, w| { var best: f64 = std.math.floatMax(f64); for (0..3) |_| { - var t0: std.c.timespec = undefined; - var t1: std.c.timespec = undefined; - _ = std.c.clock_gettime(.MONOTONIC, &t0); + var t0: platform.c.timespec = undefined; + var t1: platform.c.timespec = undefined; + _ = platform.c.clock_gettime(platform.c.CLOCK.MONOTONIC, &t0); try fwd.run(self.model, &self.scratch, &.{.{ .cache = &cache, .rows = w }}, ids[0..w], true, .all); - _ = std.c.clock_gettime(.MONOTONIC, &t1); + _ = platform.c.clock_gettime(platform.c.CLOCK.MONOTONIC, &t1); cache.len = base; const ms = @as(f64, @floatFromInt(t1.sec - t0.sec)) * 1e3 + @as(f64, @floatFromInt(t1.nsec - t0.nsec)) / 1e6; best = @min(best, ms); diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index 7bc9e5de6e..2fe9440070 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 99e6b17816..9ad312137e 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"); @@ -140,7 +141,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); try latent.place(&e.c); if (e.c.rank > 0) return error.RankServesExperts; // ranks past 0 run cuda_worker.serve, not an engine diff --git a/zig/src/families/nemotron/cuda_identity.zig b/zig/src/families/nemotron/cuda_identity.zig index c080aaee87..10570b816d 100644 --- a/zig/src/families/nemotron/cuda_identity.zig +++ b/zig/src/families/nemotron/cuda_identity.zig @@ -370,7 +370,8 @@ test "a Hugging Face snapshot's shard is checked against its blob name: the righ const snap = try tmp.dir.createDirPathOpen(io, try std.fmt.bufPrint(&buf, "{s}/snapshots/r1", .{repo}), .{}); defer snap.close(io); try snap.writeFile(io, .{ .sub_path = "config.json", .data = "{\"hidden_size\": 8}" }); - try snap.symLink(io, try std.fmt.bufPrint(&buf, "../../blobs/{s}", .{&hexname}), "model.safetensors", .{}); + snap.symLink(io, try std.fmt.bufPrint(&buf, "../../blobs/{s}", .{&hexname}), "model.safetensors", .{}) catch |e| + return if (e == error.PermissionDenied and @import("builtin").os.tag == .windows) error.SkipZigTest else e; // symlinks need privilege on Windows } const good = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}/good/snapshots/r1", .{tmp.sub_path}); defer gpa.free(good); diff --git a/zig/src/families/nemotron/cuda_native.zig b/zig/src/families/nemotron/cuda_native.zig index 1f2544dd8d..0a9c3b63a7 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"); @@ -154,7 +155,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| { @@ -171,5 +172,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 3f859410b5..94b3e3e8da 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"); @@ -436,7 +437,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 4e8a084e0f..f6ced585d6 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 f11d9e1b94..69cf472cd3 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"); @@ -100,7 +101,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 562391207d..2dcb203a2e 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"); @@ -53,7 +54,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); } @@ -138,7 +139,7 @@ pub fn Pocket(comptime T: type) type { forgot: forgetting.Kept = .{}, // what the last delete changed, until a save writes it or it is put back 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") }; } @@ -363,7 +364,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 9583f9713b..5219d059b7 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"); const forgetting = @import("pocket_forget.zig"); @@ -11,8 +12,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/nemotron/xpu_engine.zig b/zig/src/families/nemotron/xpu_engine.zig index c6151738dc..8a1e6359af 100644 --- a/zig/src/families/nemotron/xpu_engine.zig +++ b/zig/src/families/nemotron/xpu_engine.zig @@ -1,6 +1,7 @@ //! Nemotron-H serial decode on an Intel GPU as one engine value: device checkpoint, one sequence. const std = @import("std"); +const platform = @import("sys"); const xpu = @import("xpu"); const cfgm = @import("xpu_config.zig"); const widths = @import("config.zig"); @@ -21,8 +22,8 @@ pub const Options = struct { }; 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(platform.c.CLOCK.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); } diff --git a/zig/src/families/nemotron/xpu_mtp.zig b/zig/src/families/nemotron/xpu_mtp.zig index 4e447f7a9f..e22a5533c9 100644 --- a/zig/src/families/nemotron/xpu_mtp.zig +++ b/zig/src/families/nemotron/xpu_mtp.zig @@ -1,5 +1,6 @@ //! Nemotron-H's own MTP head as a speculative drafter: eh_proj, attention (own KV), MoE, norm, lm_head; equals greedy. const std = @import("std"); +const platform = @import("sys"); const rt = @import("xpu").rt; const ld = @import("xpu").loader; const model = @import("xpu_model.zig"); @@ -23,8 +24,8 @@ fn run(k: *rt.Kernel, groups: [3]u32, a: anytype) !void { pub const Stats = struct { windows: u64 = 0, plain: u64 = 0, drafted: u64 = 0, accepted: u64 = 0, tokens: u64 = 0, win_ns: u64 = 0, mtp_ns: u64 = 0, steps: u64 = 0, absorbs: u64 = 0 }; 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(platform.c.CLOCK.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); } diff --git a/zig/src/families/nemotron/xpu_win.zig b/zig/src/families/nemotron/xpu_win.zig index c5dcd24e5b..84e728b93d 100644 --- a/zig/src/families/nemotron/xpu_win.zig +++ b/zig/src/families/nemotron/xpu_win.zig @@ -1,5 +1,6 @@ //! Nemotron-H multi-row windows: 1..16 tokens in one pass, rows bit-identical to the one-token path; else commit(c). const std = @import("std"); +const platform = @import("sys"); const rt = @import("xpu").rt; const model = @import("xpu_model.zig"); @@ -30,8 +31,8 @@ pub var prof_on: bool = false; var prof: [48]struct { h: ?*const anyopaque = null, ns: u64 = 0, calls: u64 = 0 } = @splat(.{}); 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(platform.c.CLOCK.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); } 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 8fd6c179e0..d4a811fc34 100644 --- a/zig/src/families/qwen3_5/backend.zig +++ b/zig/src/families/qwen3_5/backend.zig @@ -164,106 +164,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/backend/hip/engine/hip_costs.zig b/zig/src/families/qwen3_5/backend/hip/engine/hip_costs.zig index 0fd6318a3c..afd9cd9271 100644 --- a/zig/src/families/qwen3_5/backend/hip/engine/hip_costs.zig +++ b/zig/src/families/qwen3_5/backend/hip/engine/hip_costs.zig @@ -1,6 +1,7 @@ //! The HIP lanes backend's own costs, timed at load on scratch streams: what the depth rule weighs drafts against. const std = @import("std"); +const platform = @import("sys"); const hip = @import("hip"); const lanes = @import("lanes"); const Engine = @import("engine.zig").Engine; @@ -29,8 +30,8 @@ pub const Costs = struct { }; 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(platform.c.CLOCK.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); } diff --git a/zig/src/families/qwen3_5/backend/hip/engine/round.zig b/zig/src/families/qwen3_5/backend/hip/engine/round.zig index 8af857901b..363d8e5f3c 100644 --- a/zig/src/families/qwen3_5/backend/hip/engine/round.zig +++ b/zig/src/families/qwen3_5/backend/hip/engine/round.zig @@ -1,6 +1,7 @@ //! A lane round on the engine: its device plan, its forward (replayed, captured or eager), its draws and its keep. const std = @import("std"); +const platform = @import("sys"); const hip = @import("hip"); const state = @import("../forward/state.zig"); const win = @import("../forward/window.zig"); @@ -125,8 +126,8 @@ pub fn choose(e: *Engine, rows: []const Rows, forced: ?Pick) !Pick { } 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(platform.c.CLOCK.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); } diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig new file mode 100644 index 0000000000..47158adda9 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda.zig @@ -0,0 +1,14 @@ +//! 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"); + +test { + _ = config; + _ = geometry; + _ = @import("draw.zig"); + _ = @import("cuda_draft_weights.zig"); + _ = draft; +} diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig new file mode 100644 index 0000000000..b4376c9490 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -0,0 +1,100 @@ +//! A stream's state on CUDA: recurrences, conv tails, keys and values, the drafter's tap ring, the last window's log. +const std = @import("std"); +const cuda = @import("cuda"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Engine = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows a drafter's tap ring holds a stream: its window and a prompt chunk more. +pub const tap_rows = 2048 + engine.max_rows; + +/// qwen35.cuh's Seg and Path. +pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; +pub const Path = extern struct { n: c_int, rows: [engine.most_window]c_int }; + +pub const Cache = struct { + state: DeviceBuffer, + kv: DeviceBuffer, + /// A drafter's tapped rows (`Engine.taps`), a ring of tap_rows: the row at position p in slot p % tap_rows. + tap: ?DeviceBuffer = null, + /// The newest verify window's rows' records (qwen35.cuh's LogRow), most_window rows a linear layer. + log: DeviceBuffer, + cap: usize, + len: usize = 0, + /// Positions the newest forward wrote up to: a tap ring holds them until the next forward, whatever `keep` lands. + written: usize = 0, + logits: []u16, + /// The newest verify window of more than one row, until its path is kept: where it started, its rows. + last: ?struct { start: usize, rows: usize } = null, + + pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { + var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); + errdefer state.free(); + try state.fill8(0, null); + var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); + errdefer kv.free(); + try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) + var log = try DeviceBuffer.alloc(e.d, e.g.logBytes(engine.most_window)); + errdefer log.free(); + try log.fill8(0, null); + var tap: ?DeviceBuffer = null; + if (e.n_taps > 0) { + tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); + try tap.?.fill8(0, null); + } + errdefer if (tap) |*t| t.free(); + return .{ .state = state, .kv = kv, .tap = tap, .log = log, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + } + + /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). + pub fn reset(self: *Cache, e: *const Engine) !void { + self.len = 0; + self.written = 0; + self.last = null; + try self.state.fill8(0, e.stream.handle); + } + + pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { + self.state.free(); + self.kv.free(); + self.log.free(); + if (self.tap) |*t| t.free(); + gpa.free(self.logits); + } + + /// The stream's window of `rows` rows from forward row `row0`, after `len` cached positions. + pub fn seg(self: *const Cache, e: *const Engine, len: usize, row0: usize, rows: usize) Seg { + const g = e.g; + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; + } + + /// Keeps `path` (window rows, ascending from 0): the state as that chain would leave it, a tree's keys moved down. + pub fn keep(self: *Cache, e: *Engine, path: []const u32) !void { + const last = self.last orelse return error.NothingToKeep; + if (path.len == 0 or path.len > last.rows or path[0] != 0) return error.NothingToKeep; + var p: Path = .{ .n = @intCast(path.len), .rows = @splat(0) }; + var prefix = true; + for (path, 0..) |r, i| { + if (r >= last.rows or (i > 0 and r <= path[i - 1])) return error.NothingToKeep; + p.rows[i] = @intCast(r); + prefix = prefix and r == i; + } + self.last = null; + self.len = last.start + path.len; + const g = e.g; + const sg = self.seg(e, last.start, 0, last.rows); + try e.run(e.keeps.apply, .{ 1, geometry.dk / 4, g.v_heads * g.deltas() }, .{ 32, 4, 1 }, .{ e.geo, sg, p }); + try e.run(e.keeps.conv, .{ g.deltas(), (g.convBytes() / 2 + 255) / 256, 1 }, .{ 256, 1, 1 }, .{ e.geo, sg, @as(c_int, @intCast(path[path.len - 1])) }); + if (!prefix and g.attentions() > 0) try e.run(e.keeps.compact, .{ g.attentions(), g.kv_heads, 1 }, .{ geometry.head_dim, 1, 1 }, .{ e.geo, sg, p }); + } + + /// The newest window kept whole (the round loop keeps nothing when every row landed), before the next forward. + pub fn settle(self: *Cache, e: *Engine) !void { + const last = self.last orelse return; + var all: [engine.most_window]u32 = undefined; + for (all[0..last.rows], 0..) |*r, i| r.* = @intCast(i); + try self.keep(e, all[0..last.rows]); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig new file mode 100644 index 0000000000..10862f6164 --- /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.embed(d.ids, d.h, B, 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 e.projSums(e.head, d.x, d.xs, d.logits, rows, rows > 16, 0); + 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..5cf32e6349 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -0,0 +1,566 @@ +//! 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 Affine8 = cuda.qlinear.Affine8; +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 }; + +/// A projection on the tiles of the width it is served at: 4 bits (codes of 2 to 4) or 8 (codes of 5 to 8). +pub const Proj = union(enum) { + a4: Affine4, + a8: Affine8, + + pub fn outputs(p: Proj) usize { + return switch (p) { + inline else => |q| q.n, + }; + } + + pub fn inputs(p: Proj) usize { + return switch (p) { + inline else => |q| q.k, + }; + } +}; + +pub const Delta = struct { qkv: Proj, z: Proj, a: Proj, b: Proj, out: Proj, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; +pub const Attention = struct { q: Proj, k: Proj, v: Proj, o: Proj, q_norm: u64, k_norm: u64 }; +pub const Layer = struct { + input_norm: u64, + post_norm: u64, + gate: Proj, + up: Proj, + down: Proj, + mixer: union(enum) { delta: Delta, attention: Attention }, +}; + +const Fns = struct { + embed: cuda.Function, + embed8: cuda.Function, + norm: cuda.Function, + sums: cuda.Function, + gdn_pre: cuda.Function, + conv_commit: cuda.Function, + gdn_chain: cuda.Function, + gdn_post: cuda.Function, + head_norm: cuda.Function, + rope: cuda.Function, + attn_partial: cuda.Function, + attn_merge: cuda.Function, + attn_gate: cuda.Function, + mlp_act: cuda.Function, + argmax: cuda.Function, + tap: cuda.Function, +}; + +/// qwen35_keep.cu's kernels (Cache.keep). +pub const Keeps = struct { apply: cuda.Function, conv: cuda.Function, compact: cuda.Function }; + +/// A window of one stream in a forward: `parents`, for a tree, names each row's parent row (-1: the first row). +pub const Part = struct { cache: *Cache, rows: usize, parents: ?[]const i32 = null }; + +/// What a forward keeps: a prompt chunk its last row's logits, a verify window every row's draw or logits and log. +pub const Keep = union(enum) { last, draws: []const Draw, logits }; + +/// A kept row's draw on the device: sample.cu's keyed rule at `position`, or greedy without one. +pub const Draw = struct { rule: ?cuda.sample.Rule, position: u64 }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + g: Geometry, + geo: Geo, + stream: cuda.Stream, + mods: [8]cuda.Module, + draws: cuda.sample.Draws, + affine: cuda.qlinear.Affine4Kernels, + wide: cuda.qlinear.Affine8Kernels, // 8-bit tiles, and the packs of every width but 4 + 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, + embed_bits: u8, // the embedding's rows' width: 4 or 8 + head: Proj, + 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, + + const ids_at = 0; + const segs_at = max_rows * 4; + const rowseg_at = segs_at + max_segs * @sizeOf(Seg); + const rowpar_at = rowseg_at + max_rows * 4; + const rowanc_at = rowpar_at + max_rows * 4; + const rules_at = rowanc_at + max_rows * 4; + const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); + const tokens_at = meta_at + most_verify * 4; + const params_at = tokens_at + most_verify * 4; // two unused words, the rows, the windows, record + const tables_len = params_at + 8 * 4; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const g = try Geometry.read(gpa, io, dir); + const d = ctx.d; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = undefined; + e.gpa = gpa; + e.d = d; + e.g = g; + e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads), .nl = @intCast(g.deltas()) }; + e.verify_rows = most_verify; + e.window_rows = most_window; + e.arena_at = 0; + e.graphs = .empty; + e.n_taps = 0; + e.stream = try cuda.Stream.init(d, true); + errdefer e.stream.deinit(); + e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); + errdefer e.arena.free(); + try e.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly + e.stage = try DeviceBuffer.alloc(d, weights.stageBytes(g)); + errdefer e.stage.free(); + const kk = cuda.kernels; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample, kk.qwen35_keep, kk.affine8 }; + 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.wide = try cuda.qlinear.Affine8Kernels.resolve(d, e.mods[7], e.mods[3], sms, gb10); + errdefer e.wide.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(); + e.wide.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: Proj, x: u64, out: u64, rows: usize, prompt: bool) !void { + return e.projSums(l, x, e.xs, out, rows, prompt, e.tables.ptr + params_at); + } + + /// `proj` with group sums in `xs` and row counts at `params`; 8-bit tiles run prompts 16 rows at a time. + pub fn projSums(e: *Engine, l: Proj, x: u64, xs: u64, out: u64, rows: usize, prompt: bool, params: u64) !void { + const n_out = l.outputs(); + const k = l.inputs(); + if (prompt and rows > 16) if (l == .a4) return e.affine.prompt(e.stream, x, l.a4, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, xs, @as(c_int, @intCast(k)), params }); + var at: usize = 0; + while (at < rows) : (at += 16) { + const n = @min(16, rows - at); + const xa = x + at * k * 2; + const sa = xs + at * (k / 64) * 4; + const oa = out + at * n_out * 2; + switch (l) { + .a4 => |q| try e.affine.decode(e.stream, xa, sa, q, oa, n), + .a8 => |q| try e.wide.decode(e.stream, xa, sa, q, oa, n), + } + } + } + + /// Token rows of the embedding into `out` (rows, hidden) bf16, at the width its words are served. + pub fn embed(e: *Engine, ids: u64, out: u64, rows: usize, params: u64) !void { + const f = if (e.embed_bits == 8) e.fns.embed8 else e.fns.embed; + try e.run(f, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ ids, e.embed_w, e.embed_s, e.embed_b, out, @as(c_int, @intCast(e.g.hidden)), params }); + } + + fn norm(e: *Engine, residual: bool, w: u64, rows: usize) !void { + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); + } + + /// A chain row's rowanc: its index in its window, as a prompt chunk's rows outnumber the bits. + const chain_row: u32 = 1 << 31; + + /// The row tables of `parts`: each row's segment, parent row (-1: a window's first) and ancestors-and-self bits. + fn rowTables(st: []u8, parts: []const Part) !void { + const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. max_rows * 4]); + const rowpar = std.mem.bytesAsSlice(c_int, st[rowpar_at..][0 .. max_rows * 4]); + const rowanc = std.mem.bytesAsSlice(u32, st[rowanc_at..][0 .. max_rows * 4]); + var row0: usize = 0; + for (parts, 0..) |p, i| { + if (p.parents) |ps| if (ps.len != p.rows or p.rows > most_window or ps[0] != -1) return error.InvalidQwenParents; + for (0..p.rows) |r| { + const parent: i32 = if (p.parents) |ps| ps[r] else if (r == 0) -1 else @intCast(r - 1); + if (parent >= @as(i32, @intCast(r)) or parent < -1) return error.InvalidQwenParents; + rowseg[row0 + r] = @intCast(i); + rowpar[row0 + r] = if (parent < 0) -1 else @intCast(row0 + @as(usize, @intCast(parent))); + rowanc[row0 + r] = if (p.parents == null) chain_row | @as(u32, @intCast(r)) else (@as(u32, 1) << @intCast(r)) | (if (parent < 0) 0 else rowanc[row0 + @as(usize, @intCast(parent))]); + } + row0 += p.rows; + } + } + + /// One forward over `parts`' windows, `ids` in order; `keep` names the kept rows; an unkept verify is kept first. + pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { + const g = e.g; + var total: usize = 0; + for (parts) |p| { + try p.cache.settle(e); + if (p.rows == 0 or p.cache.len + p.rows > p.cache.cap) return error.QwenContextFull; + total += p.rows; + } + if (parts.len == 0 or parts.len > max_segs or total != ids.len or total > max_rows) return error.InvalidQwenSegments; + const record = keep != .last; + if (record and total > e.verify_rows) return error.InvalidQwenSegments; + if (!record and parts.len != 1) return error.InvalidQwenSegments; + if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; + for (ids) |id| if (id >= g.vocab) return error.InvalidQwenToken; + for (parts) |p| if (p.parents != null and !record) return error.InvalidQwenParents; + const heads: usize = if (record) total else 1; + // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps + const prompt = !record; + // the tables + const st = e.staged.bytes; + @memcpy(std.mem.bytesAsSlice(u32, st[ids_at..][0 .. total * 4]), ids); + const segs = std.mem.bytesAsSlice(Seg, st[segs_at..][0 .. parts.len * @sizeOf(Seg)]); + try rowTables(st, parts); + var row0: usize = 0; + for (parts, 0..) |p, i| { + segs[i] = p.cache.seg(e, p.cache.len, row0, p.rows); + row0 += p.rows; + } + if (keep == .draws) { + const rules = std.mem.bytesAsSlice(cuda.sample.Rule, st[rules_at..][0 .. total * @sizeOf(cuda.sample.Rule)]); + const meta = std.mem.bytesAsSlice(c_int, st[meta_at..][0 .. total * 4]); + for (keep.draws, rules, meta) |p, *r, *m| { + if (p.rule) |x| r.* = x; + m.* = @intCast(p.position -% 1); // sample.cu keys row r at meta + r + 1 + } + } + // words 0 and 1 are reserved (each row's attention chunks follow its own keys: qwen35.cu's chunk_keys) + const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..32]); + params[0] = 0; + params[1] = 0; + params[2] = @intCast(total); + params[3] = @intCast(parts.len); + params[4] = @intFromBool(record); + try e.tables.uploadAsync(0, st, e.stream.handle); + // a captured forward covers a row bucket and every window; kernels read the true counts; a lone step is exact + const lone = total == 1 and parts.len == 1; + const bucket: usize = if (lone) 1 else if (total <= 16) 16 else most_verify; + const windows: usize = if (lone) 1 else max_segs; + const shape: Shape = .{ .total = @intCast(if (prompt) total else bucket), .nseg = @intCast(if (prompt) parts.len else windows), .record = record, .prompt = prompt, .taps = e.n_taps > 0 }; + if (prompt) { + // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length's graph would pay the stall + try e.body(shape); + } else { + const gop = try e.graphs.getOrPut(e.gpa, shape.key()); + if (!gop.found_existing) { + errdefer _ = e.graphs.remove(shape.key()); + // captured, nothing runs: the replay below is this shape's first run + try cuda.graph.beginCapture(e.stream, .thread_local); + e.body(shape) catch |err| { + var broken = cuda.graph.endCapture(e.stream) catch return err; + broken.deinit(); + return err; + }; + var captured = try cuda.graph.endCapture(e.stream); + defer captured.deinit(); + gop.value_ptr.* = try captured.instantiate(); + } + try gop.value_ptr.launchOn(e.stream); + } + for (parts) |p| { + p.cache.last = if (record and p.rows > 1) .{ .start = p.cache.len, .rows = p.rows } else null; + p.cache.len += p.rows; + p.cache.written = p.cache.len; + } + switch (keep) { + .draws => |pk| { + const toks = e.tables.ptr + tokens_at; + for (pk, 0..) |p, i| { + const row = e.logits + i * g.vocab * 2; + if (p.rule != null) { + try e.draws.draw(e.stream, row, g.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); + } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, @intCast(g.vocab)), toks + i * 4 }); + } + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); + }, + .logits, .last => { + const out = e.host_logits.bytes[0 .. heads * g.vocab * 2]; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); + }, + } + } + + /// A forward's shape, which fixes its launches: its tables and the caches' lengths are read on the device. + const Shape = struct { + total: u8, + nseg: u8, + record: bool, + prompt: bool, + taps: bool, + + fn key(s: Shape) u32 { + return @as(u32, s.total) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17 | @as(u32, @intFromBool(s.taps)) << 18; + } + }; + + /// The launches of a forward of `shape`, from the embedding to the kept rows' logits (captured once a shape). + fn body(e: *Engine, shape: Shape) !void { + const g = e.g; + const total: usize = shape.total; + const nseg: usize = shape.nseg; + const record = shape.record; + const prompt = shape.prompt; + const heads: usize = if (record) total else 1; + const t_ids = e.tables.ptr + ids_at; + const t_segs = e.tables.ptr + segs_at; + const t_rows = e.tables.ptr + rowseg_at; + const t_par = e.tables.ptr + rowpar_at; + const t_anc = e.tables.ptr + rowanc_at; + const t_params = e.tables.ptr + params_at; + const geo = e.geo; + const group = g.q_heads / g.kv_heads; + + try e.embed(t_ids, e.h, total, t_params); + var di: usize = 0; + var ai: c_int = 0; + for (e.layers, 0..) |l, i| { + try e.norm(i != 0, l.input_norm, total); + switch (l.mixer) { + .delta => |m| { + const layer: c_int = @intCast(di); + try e.proj(m.qkv, e.x, e.qkv, total, prompt); + try e.proj(m.z, e.x, e.z, total, prompt); + try e.proj(m.a, e.x, e.a, total, prompt); + try e.proj(m.b, e.x, e.b, total, prompt); + try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, t_par, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta, t_params }); + try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, t_par, layer, t_params }); + try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, t_par, layer, e.y, t_params }); + try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix, t_params }); + try e.proj(m.out, e.mix, e.r, total, prompt); + di += 1; + }, + .attention => |m| { + const aw = g.attnWidth(); + try e.proj(m.q, e.x, e.qraw, total, prompt); + try e.proj(m.k, e.x, e.kraw, total, prompt); + try e.proj(m.v, e.x, e.vraw, total, prompt); + try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps, t_params }); + try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, t_anc, ai, t_params }); + try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, t_anc, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_segs, t_rows, t_anc, t_params, e.y }); + try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); + try e.proj(m.o, e.mix, e.r, total, prompt); + ai += 1; + }, + } + try e.norm(true, l.post_norm, total); + 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..ef98544306 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -0,0 +1,285 @@ +//! 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"); +const core = @import("core"); + +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; + +/// A layer's projections, each its own MLX width in a mixed checkpoint. +pub const Slot = enum { gate, up, down, qkv, z, a, b, out, q, k, v, o }; + +/// Each slot's module under its layer (`layers.{i}.` before it), as the config's quantization table names it. +pub const slot_modules = std.enums.EnumArray(Slot, []const u8).init(.{ + .gate = "mlp.gate_proj", + .up = "mlp.up_proj", + .down = "mlp.down_proj", + .qkv = "linear_attn.in_proj_qkv", + .z = "linear_attn.in_proj_z", + .a = "linear_attn.in_proj_a", + .b = "linear_attn.in_proj_b", + .out = "linear_attn.out_proj", + .q = "self_attn.q_proj", + .k = "self_attn.k_proj", + .v = "self_attn.v_proj", + .o = "self_attn.o_proj", +}); + +/// The MLX code width of every quantized module (2 to 8 bits, groups of 64): a uniform checkpoint's are all one. +pub const Widths = struct { + embed: u8 = 4, + head: u8 = 4, + proj: [max_layers][std.enums.values(Slot).len]u8 = @splat(@splat(4)), + + pub fn of(w: *const Widths, layer: usize, slot: Slot) u8 { + return w.proj[layer][@backingInt(slot)]; + } + + /// Whether any module is not 4 bits: the loader then packs that module from its own width. + pub fn mixed(w: *const Widths, layers: usize) bool { + if (w.embed != 4 or w.head != 4) return true; + for (w.proj[0..layers]) |row| for (row) |b| if (b != 4) return true; + return false; + } +}; + +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, + bits: Widths = .{}, + + 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, if (rope.get("rope_type") != null) "rope_type" else "type", []const u8, "default"); // "type": the older key + 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); + 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; + } + g.bits = try widths(gpa, root, q, g); + return g; +} + +/// Every module's width from the quantization table, global then per module: MLX affine, groups of 64. +fn widths(gpa: std.mem.Allocator, root: std.json.ObjectMap, q: std.json.ObjectMap, g: Geometry) !Widths { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const mlx = core.quant.mlx; + const table = (mlx.detect(arena.allocator(), .{ .root = root, .quantization = q }) catch return error.UnsupportedQwenQuantization) orelse + return error.UnsupportedQwenQuantization; + const Look = struct { + fn bits(t: mlx.Config, key: []const u8) !u8 { + const w = t.width(key) catch return error.UnsupportedQwenQuantization; + if (w.group != 64) return error.UnsupportedQwenQuantization; + return w.bits; + } + }; + var out: Widths = .{ .embed = try Look.bits(table, "embed_tokens"), .head = undefined }; + out.head = if (g.tied) out.embed else try Look.bits(table, "lm_head"); + var buf: [96]u8 = undefined; + for (0..g.layers) |i| for (std.enums.values(Slot)) |slot| { + const used = switch (slot) { + .gate, .up, .down => true, + .qkv, .z, .a, .b, .out => g.linear[i], + .q, .k, .v, .o => !g.linear[i], + }; + if (used) out.proj[i][@backingInt(slot)] = try Look.bits(table, try std.fmt.bufPrint(&buf, "layers.{d}.{s}", .{ i, slot_modules.get(slot) })); + }; + return out; +} + +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); + try std.testing.expect(!g.bits.mixed(g.layers)); // uniform 4-bit: every module the global width + 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)); +} + +test "a mixed checkpoint's widths by module, under any spelling; other groups and modes refused" { + const text = + \\{"model_type":"qwen3_5","quantization":{"group_size":64,"bits":2,"mode":"affine", + \\ "language_model.model.embed_tokens":{"group_size":64,"bits":8}, + \\ "language_model.model.layers.0.linear_attn.in_proj_a":{"group_size":64,"bits":5}, + \\ "model.layers.1.self_attn.v_proj":{"group_size":64,"bits":6}, + \\ "language_model.model.layers.1.mlp.down_proj":{"group_size":64,"bits":3}}, + \\ "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":2048,"intermediate_size":6144,"num_hidden_layers":2,"vocab_size":248320,"linear_num_key_heads":16, + \\ "linear_num_value_heads":16,"num_attention_heads":8,"num_key_value_heads":2,"rms_norm_eps":1e-06, + \\ "max_position_embeddings":262144,"tie_word_embeddings":true, + \\ "layer_types":["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(u8, 8), g.bits.embed); + try std.testing.expectEqual(@as(u8, 8), g.bits.head); // tied: the embedding's width + try std.testing.expectEqual(@as(u8, 5), g.bits.of(0, .a)); + try std.testing.expectEqual(@as(u8, 2), g.bits.of(0, .qkv)); + try std.testing.expectEqual(@as(u8, 6), g.bits.of(1, .v)); + try std.testing.expectEqual(@as(u8, 3), g.bits.of(1, .down)); + try std.testing.expectEqual(@as(u8, 2), g.bits.of(1, .gate)); + try std.testing.expect(g.bits.mixed(g.layers)); + const g32 = try std.mem.replaceOwned(u8, std.testing.allocator, text, "\"group_size\":64,\"bits\":3", "\"group_size\":32,\"bits\":3"); + defer std.testing.allocator.free(g32); + try std.testing.expectError(error.UnsupportedQwenQuantization, parse(std.testing.allocator, g32)); + const gptq = try std.mem.replaceOwned(u8, std.testing.allocator, text, "\"mode\":\"affine\"", "\"mode\":\"gptq\""); + defer std.testing.allocator.free(gptq); + try std.testing.expectError(error.UnsupportedQwenQuantization, parse(std.testing.allocator, gptq)); +} 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..8773dffc71 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -0,0 +1,338 @@ +//! 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 be = lanes.backend; + +pub const model_type = "qwen3_5"; +pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const default_context: i64 = 32768; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = E.max_rows; +/// --drafter loads the DFlash2 draft model (openDrafter); the checkpoints have no head of their own. +pub const takes_drafter = true; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; + +/// What the native server drives: the lane backend, the facts its round loop reads, and how to free it. +pub const Loaded = struct { + backend: be.Backend, + facts: lanes.Model, + rows: u32, + /// Device bytes each admitted stream allocates (recurrences, convolution tails, keys and values). + stream_bytes: usize, + free_streams: u32 = 0, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + /// Called with the admitted stream count: that many caches allocated and touched now, as memory commits on touch. + reserve: *const fn (*anyopaque, u32) anyerror!void, +}; + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.QwenHasNoCapturedKernels; + const e = try E.Engine.init(gpa, io, ctx, dir); + errdefer e.deinit(); + if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary + const l = try gpa.create(Lanes); + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context }; + return .{ + .backend = l.backend(), + .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, + .rows = @intCast(e.window_rows), + .stream_bytes = e.streamBytes(o.context), + .ctx = l, + .deinit = release, + .reserve = reserve, + }; +} + +fn reserve(p: *anyopaque, streams: u32) anyerror!void { + const l: *Lanes = @ptrCast(@alignCast(p)); + try l.pool.ensureTotalCapacity(l.gpa, streams); + while (l.pool.items.len < streams) { + const cache = try l.gpa.create(E.Cache); + errdefer l.gpa.destroy(cache); + cache.* = try E.Cache.init(l.gpa, l.e, l.capacity); + errdefer cache.deinit(l.gpa); + // a row through it: kernels have to have referenced a buffer once before the driver's stall is behind it + try l.e.forward(&.{.{ .cache = cache, .rows = 1 }}, &.{0}, .last); + try cache.reset(l.e); + l.pool.appendAssumeCapacity(cache); + } + // the shared rounds' shapes, a row of 1 to `streams` streams, captured now: a new graph's first run stalls too + var parts: [E.max_segs]E.Part = undefined; + var ids: [E.max_segs]u32 = @splat(0); + var picks: [E.max_segs]E.Draw = undefined; + for (0..@min(streams, E.max_segs)) |n| { + parts[n] = .{ .cache = l.pool.items[n], .rows = 1 }; + picks[n] = .{ .rule = null, .position = 1 }; + try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); + for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); + } + // with a drafter: a lone stream's windows of the pending token and its drafts, chains and trees + if (l.e.n_taps > 0) for (2..l.e.window_rows + 1) |rows| { + for (picks[0..rows]) |*pk| pk.* = .{ .rule = null, .position = 1 }; + try l.e.forward(&.{.{ .cache = l.pool.items[0], .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); + try l.pool.items[0].reset(l.e); + }; +} + +/// The DFlash2 drafter in `dir` over engine `p`, 4-bit only, and the target's windows timed for the round loop. +pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { + if (bits != 4) return error.DrafterBitsUnsupported; + const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; + const d = try Draft.init(gpa, io, l.e, dir); + errdefer d.deinit(); + try l.e.setTaps(d.taps()); + try l.timeWindows(d); + return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter }; +} + +fn reserveDrafter(p: *anyopaque, streams: u32) anyerror!void { + const d: *Draft = @ptrCast(@alignCast(p)); + try d.reserve(streams); +} + +fn closeDrafter(p: *anyopaque) void { + const d: *Draft = @ptrCast(@alignCast(p)); + d.deinit(); +} + +/// A request this engine refuses, in words; null: none of its own. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.InvalidQwenPrompt => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.OutOfMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + else => null, + }; +} + +fn release(p: *anyopaque) void { + const l: *Lanes = @ptrCast(@alignCast(p)); + l.deinit(); +} + +/// The device draw of a stream's rule at `position`: sample.cu's keyed rule, greedy at temperature 0 or without one. +fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { + const s = sampling orelse return .{ .rule = null, .position = position }; + if (s.temperature <= 0) return .{ .rule = null, .position = position }; + return .{ .rule = cuda.sample.Rule.from(lanes.gpu_rule.rule(s)), .position = position }; +} + +const Lanes = struct { + gpa: std.mem.Allocator, + io: std.Io, + e: *E.Engine, + capacity: usize, + caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, + /// Caches at the full window, kept between streams: allocated and touched once (`reserve`), reset on reuse. + pool: std.ArrayList(*E.Cache) = .empty, + + fn deinit(self: *Lanes) void { + var it = self.caches.valueIterator(); + while (it.next()) |cache| self.drop(cache.*); + self.caches.deinit(self.gpa); + for (self.pool.items) |cache| self.drop(cache); + self.pool.deinit(self.gpa); + self.e.deinit(); + self.gpa.destroy(self); + } + + fn drop(self: *Lanes, cache: *E.Cache) void { + cache.deinit(self.gpa); + self.gpa.destroy(cache); + } + + /// A cache for a new stream: a pooled one reset, else one more (a window beyond the reserve). + fn take(self: *Lanes) !*E.Cache { + if (self.pool.pop()) |cache| { + try cache.reset(self.e); + return cache; + } + const cache = try self.gpa.create(E.Cache); + errdefer self.gpa.destroy(cache); + cache.* = try E.Cache.init(self.gpa, self.e, self.capacity); + return cache; + } + + fn give(self: *Lanes, cache: *E.Cache) void { + self.pool.append(self.gpa, cache) catch self.drop(cache); + } + + /// Each width's forward timed on its own cache, on its third run: graph captured, stall paid, a keep in front. + fn timeWindows(self: *Lanes, d: *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; + const cache = try self.take(); + errdefer self.give(cache); + var at: usize = 0; + while (at < ids.len) { + var end: usize = @min(ids.len, at + E.max_rows); + for (stream.chunks) |boundary| if (boundary > at) { + end = @min(end, boundary); + break; + }; + try self.e.forward(&.{.{ .cache = cache, .rows = end - at }}, ids[at..end], .last); + at = end; + } + try self.caches.put(self.gpa, stream, cache); + } + + fn firstFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64) !u64 { + return firstMaskedFn(ptr, stream, position, &.{}); + } + + fn firstMaskedFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64, mask: []const u32) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len) return error.InvalidQwenPosition; + return try draws.draw(self.gpa, cache.logits, stream.sampling, position, mask); + } + + fn queueFn(ptr: *anyopaque, stream: *lanes.Stream, feed: be.Feed, position: u64) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len + 1) return error.InvalidQwenPosition; + const token: u32 = switch (feed) { + .value => |v| v, + .handle => |h| @intCast(h), + }; + try self.e.forward(&.{.{ .cache = cache, .rows = 1 }}, &.{token}, .{ .draws = &.{pick(stream.sampling, position)} }); + cache.last = null; + return self.e.tokens[0]; + } + + fn readFn(_: *anyopaque, handle: u64) !u32 { + if (handle >= c.vocab) return error.InvalidQwenToken; + return @intCast(handle); + } + + /// Each window (a pending token and drafts) checked, then as many as fit a forward; a tree's rows keyed by depth. + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { + const self = cast(ptr); + const t0 = std.Io.Clock.awake.now(self.io); + defer @import("cuda_draft.zig").timing("verify", self.io, t0); + if (windows.len != out.len) return error.InvalidQwenWindows; + for (windows, out) |w, o| { + const rows = w.rows(); + if (w.held != 0 or rows > self.e.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; + if (w.masks.len != 0 and w.masks.len != rows * draws.mask_words) return error.InvalidQwenMasks; + if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; + const cache = try self.cacheOf(w.stream); + var depth: [E.most_window]u64 = undefined; + for (w.positions, 0..) |p, r| { + const parent: i32 = if (w.parents) |ps| ps[r] else if (r == 0) -1 else @intCast(r - 1); + if (r == 0 and parent != -1) return error.UnsupportedQwenWindow; + if (r > 0 and (parent < 0 or parent >= @as(i32, @intCast(r)))) return error.UnsupportedQwenWindow; + depth[r] = if (parent < 0) 0 else depth[@intCast(parent)] + 1; + if (p != cache.len + depth[r] + 1) return error.InvalidQwenPosition; + } + } + var at: usize = 0; + while (at < windows.len) { + var parts: [E.max_segs]E.Part = undefined; + var ids: [E.max_rows]u32 = undefined; + var picks: [E.max_rows]E.Draw = undefined; + var n: usize = 0; + var total: usize = 0; + var masked = false; + while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= self.e.verify_rows) : (n += 1) { + const w = windows[at + n]; + parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows(), .parents = w.parents }; + ids[total] = w.pending; + @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); + for (w.positions, 0..) |p, r| picks[total + r] = pick(w.stream.sampling, p); + masked = masked or w.masks.len != 0; + total += w.rows(); + } + // structured output draws on the host, over each row's logits with its mask applied + try self.e.forward(parts[0..n], ids[0..total], if (masked) .logits else .{ .draws = picks[0..total] }); + var row: usize = 0; + for (windows[at..][0..n], out[at..][0..n]) |w, o| { + for (o.sampled, w.positions, 0..) |*t, p, r| { + const mask = if (w.masks.len > 0) w.masks[r * draws.mask_words ..][0..draws.mask_words] else &.{}; + t.* = if (masked) try draws.draw(self.gpa, self.e.rowLogits(row + r), w.stream.sampling, p, mask) else self.e.tokens[row + r]; + } + @memcpy(o.drafts, w.tokens); + row += w.rows(); + } + at += n; + } + } + + /// The accepted path of each window (a chain's prefix or a tree's branch): the state as that chain leaves it. + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { + const self = cast(ptr); + const t0 = std.Io.Clock.awake.now(self.io); + defer @import("cuda_draft.zig").timing("keep", self.io, t0); + for (windows, paths) |w, path| { + const cache = try self.cacheOf(w.stream); + if (cache.last == null) continue; // a one-row window (or one kept already): nothing to roll back + try cache.keep(self.e, path); + } + } + + fn draftFn(_: *anyopaque, requests: []const be.DraftRequest) !void { + for (requests) |r| if (r.depth != 0) return error.QwenCheckpointHasNoDraftHead; + } + + /// A drafter's tapped layers, before any stream's caches exist: each forward fills their rings after. + fn prepareFeaturesFn(ptr: *anyopaque, taps: []const u32) !void { + const self = cast(ptr); + if (self.caches.count() != 0 or self.pool.items.len != 0) return error.TapsAfterStreams; + try self.e.setTaps(taps); + } + + /// Rows [start, start + count) of a stream's tap ring (slot position % tap_rows), the newest tap_rows rows kept. + fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { + const self = cast(ptr); + if (taps.len != self.e.n_taps) return error.TapsMismatch; + const cache = try self.cacheOf(stream); + const ring = cache.tap orelse return error.NoTaps; + // a long prompt asks for more rows than the ring keeps; the drafter reads only its window of the newest + const held_from = cache.written -| E.tap_rows; + if (start + count > cache.written or start + count <= held_from) return error.RowsNotHeld; + return .{ .buffer = ring.ptr, .offset = start, .rows = count, .row_bytes = @intCast(self.e.n_taps * self.e.g.hidden * 2), .space = .device, .dtype = .bf16 }; + } + + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { + const self = cast(ptr); + if (self.caches.fetchRemove(stream)) |entry| self.give(entry.value); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig new file mode 100644 index 0000000000..8c3d9bf29d --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -0,0 +1,233 @@ +//! The checkpoint's weights onto the device as qlinear tiles in the engine's arena: 2-4 bits affine-4, 5-8 bits 8-bit. +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 Slot = geometry.Slot; +const Engine = engine.Engine; +const Proj = engine.Proj; +const Affine4 = cuda.qlinear.Affine4; +const Affine8 = cuda.qlinear.Affine8; +const servedBits = cuda.qlinear.servedBits; +const DeviceBuffer = cuda.DeviceBuffer; +const aligned = engine.aligned; +const head_dim = geometry.head_dim; + +/// The packed tiles' bytes of an n-by-k projection of `bits` codes (qlinear's words, scales and biases). +fn packedBytes(n: usize, k: usize, bits: u8) usize { + if (servedBits(bits) == 4) { + const l = Affine4.layout(n, k); + return aligned(l.words) + 2 * aligned(l.scales); + } + const l = Affine8.layout(n, k); + return aligned(l.words) + 2 * aligned(l.scales); +} + +/// The checkpoint's own bytes of an n-by-k projection of `bits` codes: words, then its scales and biases. +fn rawBytes(n: usize, k: usize, bits: u8) usize { + return n * k * bits / 8 + 2 * (n * k / 64 * 2); +} + +/// A slot's output and input widths. +fn shape(g: Geometry, slot: Slot) [2]usize { + return switch (slot) { + .gate, .up => .{ g.intermediate, g.hidden }, + .down => .{ g.hidden, g.intermediate }, + .qkv => .{ g.convWidth(), g.hidden }, + .z => .{ g.valueWidth(), g.hidden }, + .a, .b => .{ g.v_heads, g.hidden }, + .out => .{ g.hidden, g.valueWidth() }, + .q => .{ 2 * g.attnWidth(), g.hidden }, + .k, .v => .{ g.kv_heads * head_dim, g.hidden }, + .o => .{ g.hidden, g.attnWidth() }, + }; +} + +/// The slots layer `i` loads: the MLP's and its mixer's. +fn slots(g: Geometry, i: usize) []const Slot { + return if (g.linear[i]) &.{ .gate, .up, .down, .qkv, .z, .a, .b, .out } else &.{ .gate, .up, .down, .q, .k, .v, .o }; +} + +/// The token embedding as the embed kernels read it: MLX rows at the width it is served (4 or 8 bits). +fn embedBytes(g: Geometry) usize { + return aligned(g.vocab * g.hidden * servedBits(g.bits.embed) / 8) + 2 * aligned(g.vocab * g.hidden / 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 = embedBytes(g); + t += packedBytes(g.vocab, g.hidden, g.bits.head); // the head, tied or not + t += aligned(g.hidden * 2); + for (0..g.layers) |i| { + t += 2 * aligned(g.hidden * 2); + for (slots(g, i)) |slot| { + const nk = shape(g, slot); + t += packedBytes(nk[0], nk[1], g.bits.of(i, slot)); + } + if (g.linear[i]) { + t += aligned(g.convWidth() * geometry.taps * 2) + aligned(geometry.dk * 2) + aligned(g.v_heads * 4) + aligned(g.v_heads * 2); + } else { + t += 2 * aligned(head_dim * 2); + } + } + return t; +} + +/// The largest projection's bytes as stored: the staging buffer the load packs each one from (and widens through). +pub fn stageBytes(g: Geometry) usize { + var most = rawBytes(g.vocab, g.hidden, g.bits.head); + if (servedBits(g.bits.embed) != g.bits.embed) most = @max(most, rawBytes(g.vocab, g.hidden, g.bits.embed)); + for (0..g.layers) |i| for (slots(g, i)) |slot| { + const nk = shape(g, slot); + most = @max(most, rawBytes(nk[0], nk[1], g.bits.of(i, slot))); + }; + 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 projection packed into qlinear's tiles from staging; 4-bit takes the affine-4 pack, others by width. +fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize, bits: u8) !Proj { + 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 * bits / 32 }, .{ "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 out: Proj = if (servedBits(bits) == 4) blk: { + 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 }; + if (bits == 4) { + try e.affine.pack(e.stream, raw[0], raw[1], raw[2], q); + } else { + try e.wide.packWords(e.stream, raw[0], bits, n, k, q.w); + try e.wide.packTable(e.stream, raw[1], n, k, q.npad, q.s); + try e.wide.packTable(e.stream, raw[2], n, k, q.npad, q.b); + } + break :blk .{ .a4 = q }; + } else blk: { + const l = Affine8.layout(n, k); + const q: Affine8 = .{ .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.wide.packWords(e.stream, raw[0], bits, n, k, q.w); + try e.wide.packTable(e.stream, raw[1], n, k, q.npad, q.s); + try e.wide.packTable(e.stream, raw[2], n, k, q.npad, q.b); + break :blk .{ .a8 = q }; + }; + try e.stream.synchronize(); // the next projection's staging overwrites this one's + return out; +} + +/// The token embedding's words for the embed kernels: as stored at 4 or 8 bits, else widened through staging. +fn embedWords(e: *Engine, ck: *core.Checkpoint, m: []const u8) !u64 { + const g = e.g; + const bits = g.bits.embed; + const wide = servedBits(bits); + var buf: [160]u8 = undefined; + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}embed_tokens.weight", .{m}), .u32, &.{ g.vocab, g.hidden * bits / 32 }); + if (bits == wide) return put(e, t); + try e.stage.upload(0, t.bytes); + const out = try e.alloc(g.vocab * g.hidden * wide / 8); + try e.wide.widen(e.stream, e.stage.ptr, bits, g.vocab, g.hidden, out, wide); + try e.stream.synchronize(); + return out; +} + +pub fn load(e: *Engine, ck: *core.Checkpoint) !void { + const g = e.g; + const m = "language_model.model."; + e.embed_bits = servedBits(g.bits.embed); + e.embed_w = try embedWords(e, ck, m); + 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, g.bits.head) else try linear(e, ck, "language_model.", "lm_head", g.vocab, g.hidden, g.bits.head); + e.final_norm = try tensor(e, ck, "{s}norm.weight", .{m}, .bf16, &.{g.hidden}); + const conv = g.convWidth(); + for (e.layers, 0..) |*l, i| { + var pbuf: [96]u8 = undefined; + const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); + var ps: std.enums.EnumArray(Slot, Proj) = .initUndefined(); + for (slots(g, i)) |slot| { + const nk = shape(g, slot); + ps.set(slot, try linear(e, ck, p, geometry.slot_modules.get(slot), nk[0], nk[1], g.bits.of(i, slot))); + } + 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 = ps.get(.gate), + .up = ps.get(.up), + .down = ps.get(.down), + .mixer = undefined, + }; + if (g.linear[i]) { + l.mixer = .{ .delta = .{ + .qkv = ps.get(.qkv), + .z = ps.get(.z), + .a = ps.get(.a), + .b = ps.get(.b), + .out = ps.get(.out), + .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 = ps.get(.q), + .k = ps.get(.k), + .v = ps.get(.v), + .o = ps.get(.o), + .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}), + } }; + } + } +} + +test "a uniform 4-bit checkpoint keeps its sizes; wider modules take a byte a weight" { + var g: Geometry = undefined; + g = .{ .hidden = 2048, .intermediate = 6144, .layers = 2, .vocab = 248320, .k_heads = 16, .v_heads = 16, .q_heads = 8, .kv_heads = 2, .eps = 1e-6, .tied = true, .context = 4096, .linear = @splat(false) }; + g.linear[0] = true; + const uniform = total(g); + try std.testing.expectEqual(rawBytes(g.vocab, g.hidden, 4), stageBytes(g)); + g.bits.proj[1][@backingInt(Slot.down)] = 8; + const nk = shape(g, .down); + try std.testing.expectEqual(uniform - packedBytes(nk[0], nk[1], 4) + packedBytes(nk[0], nk[1], 8), total(g)); + g.bits.proj[1][@backingInt(Slot.down)] = 3; // widened onto the 4-bit tiles: the same device bytes + try std.testing.expectEqual(uniform, total(g)); + g.bits.embed = 6; + g.bits.head = 6; + try std.testing.expect(total(g) > uniform); + try std.testing.expectEqual(rawBytes(g.vocab, g.hidden, 6), stageBytes(g)); +} 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/families/qwen3_5/weights/real_test.zig b/zig/src/families/qwen3_5/weights/real_test.zig index c86fe1bc0b..3c5cf3d574 100644 --- a/zig/src/families/qwen3_5/weights/real_test.zig +++ b/zig/src/families/qwen3_5/weights/real_test.zig @@ -49,7 +49,7 @@ pub fn run(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Summary { } test "the real checkpoint (TF_QWEN_DIR) resolves every tensor with consistent shapes" { - const dir = std.testing.environ.getPosix("TF_QWEN_DIR") orelse return error.SkipZigTest; + const dir = @import("sys").getenv("TF_QWEN_DIR") orelse return error.SkipZigTest; const sum = try run(std.testing.allocator, std.testing.io, dir); std.debug.print("qwen35 index: {d} layers ({d} full attention, {d} MoE), tied head {}, MTP {}, {d} MiB of projections and MLPs\n", .{ sum.layers, sum.full, sum.moe, sum.tied, sum.mtp, sum.bytes >> 20 }); try std.testing.expect(sum.layers > 0); diff --git a/zig/src/hip/comm/link.zig b/zig/src/hip/comm/link.zig index c58048f4cc..d00d5dc73e 100644 --- a/zig/src/hip/comm/link.zig +++ b/zig/src/hip/comm/link.zig @@ -1,7 +1,7 @@ //! The ranks' TCP link: rank 0 listens, the others connect once, every rank's hello is checked, then rank 0 commands. const std = @import("std"); -const posix = std.posix; +const posix = @import("sys").posix; const rccl = @import("rccl.zig"); const agree_ = @import("../agree.zig"); const Policy = @import("../policy.zig").Policy; diff --git a/zig/src/hip/comm/loopback.zig b/zig/src/hip/comm/loopback.zig index 45d0a3dfc8..bc504ffd7f 100644 --- a/zig/src/hip/comm/loopback.zig +++ b/zig/src/hip/comm/loopback.zig @@ -1,7 +1,7 @@ //! Groups of threads joined over 127.0.0.1 for host tests: each rank holds the collective.Link a served rank has. const std = @import("std"); -const posix = std.posix; +const posix = @import("sys").posix; const collective = @import("../collective.zig"); const link = @import("link.zig"); const group = @import("group.zig"); diff --git a/zig/src/hip/comm/rccl.zig b/zig/src/hip/comm/rccl.zig index 52794f458f..396e974a44 100644 --- a/zig/src/hip/comm/rccl.zig +++ b/zig/src/hip/comm/rccl.zig @@ -33,7 +33,7 @@ const paths = [_][]const u8{ "librccl.so.1", "librccl.so", "/opt/rocm/lib/librcc const sum_op: c_int = 0; pub const Rccl = struct { - lib: std.DynLib, + lib: @import("sys").DynLib, api: Api, /// `first` (the policy's `rccl_lib`) is tried before the usual places. @@ -49,7 +49,7 @@ pub const Rccl = struct { n += 1; } for (candidates[0..n]) |path| { - var lib = std.DynLib.open(path) catch continue; + var lib = @import("sys").DynLib.open(path) catch continue; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; 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 690665c336..5c1ca39834 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 254a8799a9..1c7951d2e6 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,11 +1,13 @@ //! CUDA family registration, device setup, memory admission and the native lane host. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); const lfm2 = @import("lfm2.zig"); const glm = @import("glm_cuda"); +const qwen35 = @import("qwen35"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); @@ -13,9 +15,10 @@ const snaps = @import("cuda_snaps.zig"); const nemotron_slide = @import("nemotron_slide.zig"); const cuda_disk = @import("cuda_disk.zig"); const vision_offer = @import("cuda_vision_offer.zig"); +const drafted = @import("cuda_drafted.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; @@ -112,6 +115,7 @@ const Host = struct { slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into recording: ?*record.Recording = null, // --record: the recorder in front of the family's backend + served: drafted.Served = undefined, // the family's backend, wrapped with --drafter's draft model when given fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); @@ -128,6 +132,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(); @@ -151,6 +156,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(); @@ -370,20 +379,24 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O problem.* = why; return null; }; + 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; - var room = after.room(model) -| loaded.learn_bytes; + const L = @TypeOf(loaded); // a family's Loaded may leave out a learner and a lone driver + const learn_bytes: usize = if (@hasField(L, "learn_bytes")) loaded.learn_bytes else 0; + var room = after.room(model) -| learn_bytes; var low: []const u8 = ""; // names the rank whose room bounds the streams, when it is not this one if (@hasField(@TypeOf(loaded), "peer_room")) if (loaded.peer_room) |pr| if (pr.bytes < room) { room = pr.bytes; low = try std.fmt.allocPrint(a, " on rank {d}, the smallest room of the ranks", .{pr.rank}); }; const free = loaded.free_streams; - const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { + 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, no_stream, .{ window, - toGib(loaded.stream_bytes), + toGib(served.stream_bytes), toGib(room), toGib(model), toGib(after.reserve), @@ -391,22 +404,26 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }), error.TooMany => try std.fmt.allocPrint(a, too_many, .{ o.lanes, - toGib(loaded.stream_bytes * (o.lanes -| free)), + toGib(served.stream_bytes * (o.lanes -| free)), window, toGib(room), low, - free + room / loaded.stream_bytes, + free + room / served.stream_bytes, }), }; 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 held = 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 held = served.stream_bytes * (streams -| free); const cache = if (@hasDecl(F, "snap_save")) budget.cacheBytes(o.prompt_cache_gib, room, held) else 0; h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(served.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); @@ -446,7 +463,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.free(h.dir); }; if (o.record) |dir| { - const r = try record.open(gpa, io, dir, taps, loaded.backend); + const r = try record.open(gpa, io, dir, taps, served.backend); h.recording = r; h.inner = r.recorded.backend(); h.lone = null; // every token through the lane rounds, where the recorder sees it @@ -454,15 +471,15 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O } errdefer if (h.recording) |r| r.free(gpa); // 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, .{ + 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 and @hasDecl(F, "warm_quiet"), - .logprobs = loaded.backend.vtable.first_row != null, - .prompt_logprobs = loaded.backend.vtable.prompt_rows, - .structures = loaded.backend.vtable.first_masked != null, + .logprobs = served.backend.vtable.first_row != null, + .prompt_logprobs = served.backend.vtable.prompt_rows, + .structures = served.backend.vtable.first_masked != null, .vision = if (@hasField(@TypeOf(loaded), "vision")) vision_offer.offer(loaded.vision) else null, }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; diff --git a/zig/src/native/cuda_drafted.zig b/zig/src/native/cuda_drafted.zig new file mode 100644 index 0000000000..df4f0ddeb6 --- /dev/null +++ b/zig/src/native/cuda_drafted.zig @@ -0,0 +1,61 @@ +//! --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 family's own per-stream buffers, when its Loaded pre-allocates them. + family: ?struct { ctx: *anyopaque, reserve: *const fn (*anyopaque, u32) anyerror!void } = null, + + /// The family's and the drafter's per-stream buffers for `streams`, once the server knows the count. + pub fn reserve(s: *Served, streams: u32) !void { + if (s.family) |f| try f.reserve(f.ctx, streams); + 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 }; + if (@hasField(@TypeOf(loaded), "reserve")) s.family = .{ .ctx = loaded.ctx, .reserve = loaded.reserve }; + const dir = o.drafter orelse return s; + if (!o.drafts) return s; + if (!@hasDecl(F, "openDrafter")) { + problem.* = try std.fmt.allocPrint(a, "--drafter: {s} checkpoints load no draft model on CUDA", .{F.model_type}); + return null; + } + const opened = F.openDrafter(gpa, io, loaded.ctx, dir, o.drafter_bits) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load the drafter {s} ({s})", .{ dir, @errorName(e) }); + return null; + }; + const h = try gpa.create(Held); + errdefer gpa.destroy(h); + h.* = .{ .drafted = undefined, .ctx = opened.ctx, .close = opened.close, .reserve = opened.reserve }; + h.drafted = lanes.drafted.Drafted.init(gpa, loaded.backend, opened.drafter) catch |e| { + opened.close(opened.ctx); + problem.* = try std.fmt.allocPrint(a, "the drafter {s} cannot attach to this model ({s})", .{ dir, @errorName(e) }); + return null; + }; + s.backend = h.drafted.backend(); + s.facts = h.drafted.facts(loaded.facts); + s.stream_bytes = h.drafted.streamBytes(loaded.stream_bytes); + s.held = h; + return s; +} diff --git a/zig/src/native/flashnext_host.zig b/zig/src/native/flashnext_host.zig index 8281b19df4..5162a8c3bd 100644 --- a/zig/src/native/flashnext_host.zig +++ b/zig/src/native/flashnext_host.zig @@ -1,5 +1,6 @@ //! Flash Next behind the native server: one reply at a time, sampled on the GPU on request (speed-up mode greedy only). const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 0fdb9cc5f8..ccc53f61b8 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, LFM2-MoE, 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"); @@ -326,9 +327,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)); @@ -378,7 +379,7 @@ fn gibBytes(g: f64) u64 { /// 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)) { @@ -453,7 +454,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| { @@ -564,10 +565,10 @@ test "a shard rewritten in place, same size and bytes, gives the learned states try std.testing.expectEqual(first.final(), notes.final()); std.Io.sleep(io, .fromMilliseconds(5), .awake) catch {}; const path = try std.fmt.allocPrintSentinel(a, "{s}/model-00001-of-00001.safetensors", .{dir}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY }, @as(platform.c.mode_t, 0)); try std.testing.expect(fd >= 0); - try std.testing.expectEqual(@as(isize, 7), std.c.pwrite(fd, "weights", 7, 0)); // as Sliding Weights writes a tensor of the same size - _ = std.c.close(fd); + try std.testing.expectEqual(@as(isize, 7), platform.c.pwrite(fd, "weights", 7, 0)); // as Sliding Weights writes a tensor of the same size + _ = platform.c.close(fd); var moved = std.hash.Wyhash.init(0); try shardStamps(&moved, a, io, dir); try std.testing.expect(first.final() != moved.final()); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index e1f0d34134..e8c6eba8d3 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"); @@ -126,7 +127,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 168d9304c8..44c043b79f 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"); @@ -325,7 +326,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..6e42e9ee75 --- /dev/null +++ b/zig/src/platform/sys.zig @@ -0,0 +1,63 @@ +//! 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"; + +/// The file a loaded library's `symbol` comes from, as dladdr's dli_fname, in `buf` on Windows. +pub fn libraryOf(symbol: *const anyopaque, buf: []u8) ?[]const u8 { + if (!windows) @compileError("libraryOf: dladdr gives the path elsewhere"); + var module: ?*anyopaque = null; + // GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | UNCHANGED_REFCOUNT: the module holding the address, not pinned + if (GetModuleHandleExA(0x4 | 0x2, @ptrCast(symbol), &module) == 0) return null; + const n = GetModuleFileNameA(module, buf.ptr, @intCast(@min(buf.len, std.math.maxInt(u32)))); + return if (n == 0 or n >= buf.len) null else buf[0..n]; +} + +extern "kernel32" fn GetModuleHandleExA(flags: u32, name: ?[*]const u8, module: *?*anyopaque) callconv(.winapi) c_int; +extern "kernel32" fn GetModuleFileNameA(module: ?*anyopaque, name: [*]u8, size: u32) callconv(.winapi) u32; +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..f2ac507732 --- /dev/null +++ b/zig/src/platform/win_c.zig @@ -0,0 +1,569 @@ +//! 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 _exit(code: c_int) noreturn; +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 }; + +const lock = @import("win_lock.zig"); +pub const LOCK = lock.LOCK; +pub const flock = lock.flock; + +pub fn fail(e: E) c_int { + _errno().* = @backingInt(e); + return -1; +} + +/// The errno a failed Win32 call means. +pub 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. +pub 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(); +} + +pub 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; + lock.release(fd - dir_base); + 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_lock.zig b/zig/src/platform/win_lock.zig new file mode 100644 index 0000000000..aa77426264 --- /dev/null +++ b/zig/src/platform/win_lock.zig @@ -0,0 +1,72 @@ +//! flock on Windows: LockFileEx over a file's first byte; a directory's lock is on `.tensorfold-lock` inside it. + +const std = @import("std"); +const c = @import("win_c.zig"); + +pub const LOCK = struct { + pub const SH: c_int = 1; + pub const EX: c_int = 2; + pub const NB: c_int = 4; + pub const UN: c_int = 8; +}; + +const HANDLE = *anyopaque; +const invalid: isize = -1; +const OVERLAPPED = extern struct { internal: usize = 0, internal_high: usize = 0, offset: u32 = 0, offset_high: u32 = 0, event: ?HANDLE = null }; +extern "kernel32" fn LockFileEx(h: HANDLE, flags: u32, reserved: u32, low: u32, high: u32, ov: *OVERLAPPED) callconv(.winapi) c_int; +extern "kernel32" fn UnlockFileEx(h: HANDLE, reserved: u32, low: u32, high: u32, ov: *OVERLAPPED) callconv(.winapi) c_int; +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 CloseHandle(h: HANDLE) callconv(.winapi) c_int; +extern "kernel32" fn GetLastError() callconv(.winapi) u32; +extern "c" fn _get_osfhandle(fd: c_int) isize; + +/// Each directory descriptor's lock file handle (by its slot), closed with the descriptor. +var dir_locks: [256]isize = @splat(invalid); +var locks_mutex: std.atomic.Mutex = .unlocked; + +/// flock's whole-file lock: shared or exclusive, blocking unless NB; a lock held against it is AGAIN. +pub fn flock(fd: c.fd_t, op: c_int) c_int { + const h: isize = if (c.dirPath(fd)) |dir| dirLock(fd - c.dir_base, dir) orelse return -1 else _get_osfhandle(fd); + if (h == invalid) return c.fail(.BADF); + const handle: HANDLE = @ptrFromInt(@as(usize, @bitCast(h))); + var ov: OVERLAPPED = .{}; + if (op & LOCK.UN != 0) return if (UnlockFileEx(handle, 0, 1, 0, &ov) != 0) 0 else c.fail(c.lastError()); + const flags: u32 = (if (op & LOCK.EX != 0) @as(u32, 2) else 0) | (if (op & LOCK.NB != 0) @as(u32, 1) else 0); + if (LockFileEx(handle, flags, 0, 1, 0, &ov) != 0) return 0; + return c.fail(if (GetLastError() == 33) .AGAIN else c.lastError()); // ERROR_LOCK_VIOLATION +} + +/// The lock file handle of directory slot `slot`, opened on first use. +fn dirLock(slot: c.fd_t, dir: []const u8) ?isize { + while (!locks_mutex.tryLock()) std.atomic.spinLoopHint(); + defer locks_mutex.unlock(); + const held = &dir_locks[@intCast(slot)]; + if (held.* != invalid) return held.*; + var path: [1100]u8 = undefined; + var wide: [1100:0]u16 = undefined; + const p = std.fmt.bufPrint(&path, "{s}\\.tensorfold-lock", .{dir}) catch { + _ = c.fail(.NAMETOOLONG); + return null; + }; + const n = std.unicode.wtf8ToWtf16Le(&wide, p) catch { + _ = c.fail(.INVAL); + return null; + }; + wide[n] = 0; + const h = CreateFileW(wide[0..n :0].ptr, 0x8000_0000 | 0x4000_0000, 0x7, null, 4, 0x82, null); // OPEN_ALWAYS, hidden + if (h == invalid) { + _ = c.fail(c.lastError()); + return null; + } + held.* = h; + return h; +} + +/// Directory slot `slot` closed: its lock file handle goes, and the locks with it. +pub fn release(slot: c.fd_t) void { + while (!locks_mutex.tryLock()) std.atomic.spinLoopHint(); + defer locks_mutex.unlock(); + const held = &dir_locks[@intCast(slot)]; + if (held.* != invalid) _ = CloseHandle(@ptrFromInt(@as(usize, @bitCast(held.*)))); + held.* = invalid; +} diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig new file mode 100644 index 0000000000..e9f017ea13 --- /dev/null +++ b/zig/src/platform/win_posix.zig @@ -0,0 +1,329 @@ +//! 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 = AF.INET, 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 CLOEXEC: u32 = 0; // Winsock handles are not inherited by children unless asked +}; +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 KEEPALIVE: u32 = 0x0008; +}; +pub const IPPROTO = struct { + pub const TCP: i32 = 6; +}; +pub const TCP = struct { + pub const NODELAY: u32 = 1; + pub const KEEPIDLE: u32 = 3; // TCP_KEEPALIVE: seconds idle before the first probe + pub const KEEPINTVL: u32 = 17; + pub const KEEPCNT: u32 = 16; +}; +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 read(fd: socket_t, buf: [*]u8, len: usize) isize { + const n = ws.recv(fd, buf, len32(len), 0); + return if (n < 0) failed() else n; + } + pub fn poll(fds: [*]pollfd, n: usize, timeout_ms: i32) isize { + const got = posixPoll(fds[0..n], timeout_ms) catch return -1; + return @intCast(got); + } + 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). +const posixPoll = poll; + +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 976c78da26..cc40f9d49a 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; } @@ -297,7 +298,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 = ""; @@ -311,6 +312,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 }; @@ -320,11 +334,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; @@ -333,6 +359,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(); @@ -341,11 +368,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).?); @@ -361,6 +388,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); @@ -370,7 +398,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(); @@ -380,7 +408,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/benchmark_cli_test.zig b/zig/src/server/benchmark_cli_test.zig index a5aa227fa0..be20a7585a 100644 --- a/zig/src/server/benchmark_cli_test.zig +++ b/zig/src/server/benchmark_cli_test.zig @@ -1,5 +1,6 @@ //! The command end to end against fakes: attached run, consent, login, publish, cancel and the checks. const std = @import("std"); +const platform = @import("sys"); const json = @import("json"); const suite = @import("benchmark_suite.zig"); const receipt = @import("benchmark_receipt.zig"); @@ -495,9 +496,9 @@ test "the hidden login prompt reads with echo off and a guarded terminal, then r const f = try world(a, false, "202 Accepted"); defer f.stop(); var fds: [2]std.posix.fd_t = undefined; - if (std.c.pipe(&fds) != 0) return error.SkipZigTest; + if (platform.c.pipe(&fds) != 0) return error.SkipZigTest; defer for (fds) |fd| { - _ = std.c.close(fd); + _ = platform.c.close(fd); }; var buf: [256]u8 = undefined; var piped = (std.Io.File{ .handle = fds[0], .flags = .{ .nonblocking = false } }).readerStreaming(std.testing.io, &buf); @@ -515,7 +516,7 @@ test "the hidden login prompt reads with echo off and a guarded terminal, then r t.quiet = if (std.posix.tcgetattr(t.term)) |now| !now.lflag.ECHO else |_| false; if (t.quiet) t.guarded = stop.guardedTerminal() else fake.nap(2); } - _ = std.c.write(t.pipe, token ++ "\n", token.len + 1); + _ = platform.c.write(t.pipe, token ++ "\n", token.len + 1); } }; var typist: Typist = .{ .term = pty.term, .pipe = fds[1] }; diff --git a/zig/src/server/benchmark_fake.zig b/zig/src/server/benchmark_fake.zig index bb653fe539..17198d3bc5 100644 --- a/zig/src/server/benchmark_fake.zig +++ b/zig/src/server/benchmark_fake.zig @@ -1,5 +1,6 @@ //! A loopback HTTP fake for the benchmark tests: scripted replies, pauses between writes, recorded requests. const std = @import("std"); +const platform = @import("sys"); const posix = std.posix; const Allocator = std.mem.Allocator; @@ -197,23 +198,23 @@ pub const Pty = struct { const master = posix_openpt(@bitCast(posix.O{ .ACCMODE = .RDWR, .NOCTTY = true })); if (master < 0) return null; const name = if (grantpt(master) == 0 and unlockpt(master) == 0) ptsname(master) else null; - const term = if (name) |n| std.c.open(n, .{ .ACCMODE = .RDWR, .NOCTTY = true }) else -1; + const term = if (name) |n| platform.c.open(n, .{ .ACCMODE = .RDWR, .NOCTTY = true }) else -1; if (term < 0) { - _ = std.c.close(master); + _ = platform.c.close(master); return null; } return .{ .master = master, .term = term }; } pub fn close(p: Pty) void { - _ = std.c.close(p.term); - _ = std.c.close(p.master); + _ = platform.c.close(p.term); + _ = platform.c.close(p.master); } }; pub fn nap(ms: u32) void { var ts: posix.timespec = .{ .sec = @intCast(ms / 1000), .nsec = @intCast(@as(u64, ms % 1000) * std.time.ns_per_ms) }; - _ = std.c.nanosleep(&ts, null); + _ = platform.c.nanosleep(&ts, null); } /// A response head with ``Connection: close`` and no length, as the native server sends SSE. diff --git a/zig/src/server/benchmark_server.zig b/zig/src/server/benchmark_server.zig index 6462429255..7fadb7a3d6 100644 --- a/zig/src/server/benchmark_server.zig +++ b/zig/src/server/benchmark_server.zig @@ -1,5 +1,6 @@ //! The benchmark's own serve child: lock, flags, environment, port, readiness and a stop that always reaps. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const json = @import("json"); const cli = @import("cli.zig"); @@ -215,7 +216,7 @@ pub const Managed = struct { stop_flag.releaseChild(m.pid); var status: c_int = 0; while (true) { - const rc = std.c.waitpid(m.pid, &status, 0); + const rc = platform.c.waitpid(m.pid, &status, 0); if (rc >= 0 or posix.errno(rc) != .INTR) break; } m.reaped = true; @@ -224,7 +225,7 @@ pub const Managed = struct { }; const wait = struct { - extern "c" fn waitid(idtype: c_uint, id: c_uint, info: *std.c.siginfo_t, options: c_int) c_int; + extern "c" fn waitid(idtype: c_uint, id: c_uint, info: *platform.c.siginfo_t, options: c_int) c_int; const p_pid = 1; const exited: c_int = if (builtin.os.tag == .linux) std.os.linux.W.EXITED else 0x4; const no_wait: c_int = if (builtin.os.tag == .linux) std.os.linux.W.NOWAIT else 0x20; @@ -232,9 +233,9 @@ const wait = struct { /// True once ``pid`` has exited, or is no child of this process, without reaping it. pub fn ended(pid: posix.pid_t) bool { - var info: std.c.siginfo_t = undefined; + var info: platform.c.siginfo_t = undefined; @memset(std.mem.asBytes(&info), 0); - const rc = wait.waitid(wait.p_pid, @intCast(pid), &info, wait.exited | wait.no_wait | std.c.W.NOHANG); + const rc = wait.waitid(wait.p_pid, @intCast(pid), &info, wait.exited | wait.no_wait | platform.c.W.NOHANG); if (rc < 0) return posix.errno(rc) == .CHILD; const got = if (builtin.os.tag == .linux) info.fields.common.first.piduid.pid else info.pid; return got == pid; diff --git a/zig/src/server/benchmark_server_test.zig b/zig/src/server/benchmark_server_test.zig index ec2f5bae6f..7bffdb94ed 100644 --- a/zig/src/server/benchmark_server_test.zig +++ b/zig/src/server/benchmark_server_test.zig @@ -1,5 +1,6 @@ //! The serve child's life with a shell stand-in: ready, exit, timeout, stubborn child, lock, Ctrl-C. const std = @import("std"); +const platform = @import("sys"); const cli = @import("cli.zig"); const server = @import("benchmark_server.zig"); const stop = @import("benchmark_stop.zig"); @@ -326,7 +327,7 @@ test "a second Ctrl-C kills the running serve child before the benchmark exits" var waited: u32 = 0; while (!gone(m.pid) and waited < 2000) : (waited += 10) { var status: c_int = 0; - _ = std.c.waitpid(m.pid, &status, std.c.W.NOHANG); + _ = platform.c.waitpid(m.pid, &status, platform.c.W.NOHANG); fake.nap(10); } try std.testing.expect(waited < 2000); @@ -356,7 +357,7 @@ test "a lock path that is a symlink to an existing file, or a FIFO, is refused a try std.testing.expectEqualStrings("Could not create the benchmark lock; check the TensorFold directory permissions", problem); } -extern "c" fn mkfifo(path: [*:0]const u8, mode: std.c.mode_t) c_int; +extern "c" fn mkfifo(path: [*:0]const u8, mode: platform.c.mode_t) c_int; const Racer = struct { bench: []const u8, @@ -449,7 +450,7 @@ test "an exited child is seen without being reaped, and once stopped it is relea var waited: u32 = 0; while (!server.ended(quick.id.?) and waited < 2000) : (waited += 10) fake.nap(10); var status: c_int = 0; - try std.testing.expectEqual(quick.id.?, std.c.waitpid(quick.id.?, &status, std.c.W.NOHANG)); + try std.testing.expectEqual(quick.id.?, platform.c.waitpid(quick.id.?, &status, platform.c.W.NOHANG)); const routes = try healthy(a, "{\"ok\": true}", ".tensorfold-benchmark"); const f = try fake.Fake.start(std.testing.allocator, &routes); defer f.stop(); @@ -461,10 +462,10 @@ test "an exited child is seen without being reaped, and once stopped it is relea try std.testing.expectEqual(@as(posix.pid_t, 0), stop.guardedChild()); const other = try std.process.spawn(io, .{ .argv = &.{ "/bin/sh", "-c", "exec sleep 30" }, .stdin = .ignore, .stdout = .ignore, .stderr = .ignore }); defer { - _ = std.c.kill(other.id.?, .KILL); - _ = std.c.waitpid(other.id.?, &status, 0); + _ = platform.c.kill(other.id.?, .KILL); + _ = platform.c.waitpid(other.id.?, &status, 0); } stop.abandon(); fake.nap(100); - try std.testing.expectEqual(@as(posix.pid_t, 0), std.c.waitpid(other.id.?, &status, std.c.W.NOHANG)); + try std.testing.expectEqual(@as(posix.pid_t, 0), platform.c.waitpid(other.id.?, &status, platform.c.W.NOHANG)); } diff --git a/zig/src/server/benchmark_stop.zig b/zig/src/server/benchmark_stop.zig index 334517c061..9968e18c5d 100644 --- a/zig/src/server/benchmark_stop.zig +++ b/zig/src/server/benchmark_stop.zig @@ -1,5 +1,6 @@ //! Ctrl-C, SIGTERM and SIGHUP during a benchmark: one flag the waits read; a second Ctrl-C exits at once. const std = @import("std"); +const platform = @import("sys"); const posix = std.posix; var stored: std.atomic.Value(u32) = .init(0); @@ -17,7 +18,7 @@ fn onSignal(sig: posix.SIG) callconv(.c) void { abandon(); const default: posix.Sigaction = .{ .handler = .{ .handler = posix.SIG.DFL }, .mask = posix.sigemptyset(), .flags = 0 }; posix.sigaction(.INT, &default, null); - _ = std.c.raise(.INT); + _ = platform.c.raise(.INT); }; } @@ -87,7 +88,7 @@ pub fn guardedTerminal() posix.fd_t { /// What a second Ctrl-C does before the process dies: kill the guarded child and restore the guarded terminal. pub fn abandon() void { const pid = child.swap(0, .acq_rel); - if (pid > 0) _ = std.c.kill(pid, .KILL); + if (pid > 0) _ = platform.c.kill(pid, .KILL); const fd = tty_fd.swap(-1, .acq_rel); - if (fd >= 0) _ = std.c.tcsetattr(fd, .NOW, &tty_saved); + if (fd >= 0) _ = platform.c.tcsetattr(fd, .NOW, &tty_saved); } diff --git a/zig/src/server/benchmark_stop_test.zig b/zig/src/server/benchmark_stop_test.zig index bc7b833d61..b851177a47 100644 --- a/zig/src/server/benchmark_stop_test.zig +++ b/zig/src/server/benchmark_stop_test.zig @@ -1,5 +1,6 @@ //! The stop flag through its real handler, in forked copies of the test process, with children the tests reap. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const fake = @import("benchmark_fake.zig"); const posix = std.posix; @@ -37,22 +38,22 @@ const Sleeper = struct { var waited: u32 = 0; while (s.status == null and waited <= ms) : (waited += 10) { var status: c_int = 0; - if (std.c.waitpid(s.pid, &status, std.c.W.NOHANG) == s.pid) s.status = @bitCast(status) else if (waited < ms) fake.nap(10); + if (platform.c.waitpid(s.pid, &status, platform.c.W.NOHANG) == s.pid) s.status = @bitCast(status) else if (waited < ms) fake.nap(10); } return s.status; } fn end(s: *Sleeper) void { if (s.ended(0) != null) return; - _ = std.c.kill(s.pid, .KILL); + _ = platform.c.kill(s.pid, .KILL); var status: c_int = 0; - _ = std.c.waitpid(s.pid, &status, 0); + _ = platform.c.waitpid(s.pid, &status, 0); } }; fn killedBy(status: ?u32, sig: posix.SIG) bool { const s = status orelse return false; - return std.c.W.IFSIGNALED(s) and std.c.W.TERMSIG(s) == sig; + return platform.c.W.IFSIGNALED(s) and platform.c.W.TERMSIG(s) == sig; } test "a second Ctrl-C kills the guarded child, and a released child is left alone" { @@ -72,17 +73,17 @@ test "a second Ctrl-C kills the guarded child, and a released child is left alon /// Runs ``body`` in a forked copy of this process; its wait status, or null when it ran past 5 s and was killed. fn forked(body: *const fn () u8) !?u32 { - const pid = std.c.fork(); + const pid = platform.c.fork(); if (pid < 0) return error.SkipZigTest; - if (pid == 0) std.c._exit(body()); + if (pid == 0) platform.c._exit(body()); var status: c_int = 0; var waited: u32 = 0; while (waited < 5000) : (waited += 10) { - if (std.c.waitpid(pid, &status, std.c.W.NOHANG) == pid) return @bitCast(status); + if (platform.c.waitpid(pid, &status, platform.c.W.NOHANG) == pid) return @bitCast(status); fake.nap(10); } - _ = std.c.kill(pid, .KILL); - _ = std.c.waitpid(pid, &status, 0); + _ = platform.c.kill(pid, .KILL); + _ = platform.c.waitpid(pid, &status, 0); return null; } @@ -94,24 +95,24 @@ const InFork = struct { stop.releaseTerminal(); stop.install(); stop.guardChild(target); - _ = std.c.raise(.INT); + _ = platform.c.raise(.INT); if (stop.requested() != .INT) return 2; - _ = std.c.raise(.INT); + _ = platform.c.raise(.INT); return 3; } fn firstStopWins() u8 { stop.reset(); stop.install(); - _ = std.c.raise(.HUP); - _ = std.c.raise(.TERM); + _ = platform.c.raise(.HUP); + _ = platform.c.raise(.TERM); if (stop.requested() != .HUP) return 4; stop.reset(); - _ = std.c.raise(.TERM); - _ = std.c.raise(.HUP); + _ = platform.c.raise(.TERM); + _ = platform.c.raise(.HUP); if (stop.requested() != .TERM) return 5; stop.reset(); - _ = std.c.raise(.INT); + _ = platform.c.raise(.INT); if (stop.requested() != .INT) return 6; stop.restore(); return 0; @@ -130,8 +131,8 @@ test "through the real handler, a second Ctrl-C kills the guarded child and ends test "through the real handler, HUP, TERM and a first INT are recorded and the first one wins" { if (builtin.os.tag != .macos and builtin.os.tag != .linux) return error.SkipZigTest; const status = (try forked(InFork.firstStopWins)) orelse return error.TestUnexpectedResult; - try std.testing.expect(std.c.W.IFEXITED(status)); - try std.testing.expectEqual(@as(u8, 0), std.c.W.EXITSTATUS(status)); + try std.testing.expect(platform.c.W.IFEXITED(status)); + try std.testing.expectEqual(@as(u8, 0), platform.c.W.EXITSTATUS(status)); } test "a second Ctrl-C at the hidden prompt turns echo back on" { diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 460acd4599..3c4b20195c 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -131,7 +131,7 @@ pub const env = [_][]const u8{ "TENSORFOLD_API_KEY", "TENSORFOLD_NO_LIVE", "TENS pub const Args = struct { model: []const u8 = "", host: []const u8 = "127.0.0.1", - port: u16 = 8080, + port: u16 = 8888, name: []const u8 = "", chat_template: ?[]const u8 = null, alias: []const []const u8 = &.{}, diff --git a/zig/src/server/health_test.zig b/zig/src/server/health_test.zig index 272933be59..285abd56f6 100644 --- a/zig/src/server/health_test.zig +++ b/zig/src/server/health_test.zig @@ -8,7 +8,7 @@ const listener = @import("listener.zig"); const model_text = @import("model_text.zig"); const server = @import("server.zig"); const status_routes = @import("status_routes.zig"); -const posix = std.posix; +const posix = @import("sys").posix; const Threaded = std.Io.Threaded; const Allocator = std.mem.Allocator; 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 0c6aa1b7ab..cac8221617 100644 --- a/zig/src/server/hub.zig +++ b/zig/src/server/hub.zig @@ -21,7 +21,7 @@ pub fn isDir(io: std.Io, path: []const u8) bool { /// The checkpoint directory for ``model``, or null with ``problem`` set (no downloads: the cache must hold it). pub fn resolve(a: Allocator, io: std.Io, environ: ?*const std.process.Environ.Map, model: []const u8, problem: *[]const u8) !?[]const u8 { - const home = if (environ) |e| e.get("HOME") orelse "" else ""; + const home = if (environ) |e| e.get("HOME") orelse e.get("USERPROFILE") orelse "" else ""; // Windows sets USERPROFILE const expanded = if (std.mem.startsWith(u8, model, "~/")) try std.fs.path.join(a, &.{ home, model[2..] }) else model; if (isDir(io, expanded)) return expanded; if (!isRepoId(model)) { diff --git a/zig/src/server/late_system_test.zig b/zig/src/server/late_system_test.zig index a603bf552e..85dfc6b161 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; @@ -141,10 +142,10 @@ fn exchange(srv: *server.Server, path: []const u8, body: []const u8) !Reply { } fn exchangeAs(srv: *server.Server, method: []const u8, 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 request_line = "{s} {s} HTTP/1.1\r\nHost: t\r\nContent-Length: {d}\r\n\r\n"; const head = try std.fmt.allocPrint(std.testing.allocator, request_line, .{ method, path, body.len }); defer std.testing.allocator.free(head); @@ -167,23 +168,23 @@ fn exchangeAs(srv: *server.Server, method: []const u8, path: []const u8, body: [ 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 8a1f7df742..659662bfbd 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"); 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 4803950807..dec9a4b6b5 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -68,7 +68,13 @@ pub fn main(init: std.process.Init) !u8 { try std.Io.File.stdout().writeStreamingAll(io, out.written()); return 0; } - if (argv.len >= 2 and std.mem.eql(u8, argv[1], "benchmark")) return benchmark.main(init, argv[2..], engines.capabilities(a)); + if (argv.len >= 2 and std.mem.eql(u8, argv[1], "benchmark")) { + // the benchmark runs its serve child under POSIX signals and waits; Windows refuses it by name + if (comptime @import("builtin").os.tag == .windows) { + std.debug.print("tensorfold benchmark is not supported on Windows yet: run tensorfold serve and a client against it\n", .{}); + return 2; + } else return benchmark.main(init, argv[2..], engines.capabilities(a)); + } if (argv.len < 2 or !std.mem.eql(u8, argv[1], "serve")) { std.debug.print("usage: tensorfold-native capabilities --json | models | info MODEL | pull REPO[@REVISION] | slide rollback FOLDER | serve MODEL [flags] | benchmark MODEL [flags]\n", .{}); return 2; diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig index a906709a93..9b149741dd 100644 --- a/zig/src/server/media_fetch.zig +++ b/zig/src/server/media_fetch.zig @@ -253,6 +253,11 @@ fn readInto(a: std.mem.Allocator, io: Io, r: *Io.Reader, data: *std.ArrayList(u8 /// A TCP connection to the checked address on port 443, its connect bounded by `ms` (std's connect has no timeout). fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { + // the socket calls below are libc's on Linux and macOS; a Windows server refuses image URLs by name + return if (comptime @import("builtin").os.tag == .windows) fail(f, "image URLs are not supported on Windows yet") else connectPosix(addr, ms, f); +} + +fn connectPosix(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { const c = std.c; const failed = "image download failed or timed out"; var storage: c.sockaddr.storage = undefined; diff --git a/zig/src/server/request_log.zig b/zig/src/server/request_log.zig index 6a27e2cae3..1afc09f677 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"); const Value = json.Value; const Allocator = std.mem.Allocator; @@ -33,10 +34,11 @@ 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 }, 0o600) 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, 0o600)); + 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 "the request log is created owner-only and keeps image URLs out" { @@ -54,9 +56,53 @@ test "the request log is created owner-only and keeps image URLs out" { const body = try json.newObject(a); try body.put(a, "messages", .{ .array = try a.dupe(Value, &.{.{ .object = part }}) }); append(a, path, .{ .object = body }); - const st = try std.Io.Dir.cwd().statFile(std.testing.io, path, .{}); // Linux's libc has no stat for std.c - try std.testing.expectEqual(@as(u32, 0o600), @as(u32, @intCast(st.permissions.toMode())) & 0o777); + if (comptime @import("builtin").os.tag != .windows) { // Windows files have no mode bits + const st = try std.Io.Dir.cwd().statFile(std.testing.io, path, .{}); // Linux's libc has no stat for std.c + try std.testing.expectEqual(@as(u32, 0o600), @as(u32, @intCast(st.permissions.toMode())) & 0o777); + } const text = try std.Io.Dir.cwd().readFileAlloc(std.testing.io, path, a, .limited(1 << 16)); try std.testing.expect(std.mem.indexOf(u8, text, "") != null); try std.testing.expect(std.mem.indexOf(u8, text, "example.test") == null); } + +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 381c864ad0..efd1f7a234 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -51,6 +51,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"); @@ -85,17 +86,20 @@ test { _ = @import("cortheon/cortheon.zig"); _ = @import("cortheon/cortheon_test.zig"); _ = @import("recall_routes.zig"); - _ = @import("benchmark_suite.zig"); - _ = @import("benchmark_check.zig"); - _ = @import("benchmark_receipt_test.zig"); - _ = @import("benchmark_stop_test.zig"); - _ = @import("benchmark_sse_test.zig"); - _ = @import("benchmark_stream_test.zig"); - _ = @import("benchmark_server_test.zig"); - _ = @import("benchmark_facts_test.zig"); - _ = @import("benchmark_upload.zig"); - _ = @import("benchmark_upload_test.zig"); - _ = @import("benchmark_args.zig"); - _ = @import("benchmark_cli_test.zig"); - _ = @import("benchmark_run_test.zig"); + // the benchmark command's serve child runs under POSIX signals and waits; Windows refuses the command + if (comptime @import("builtin").os.tag != .windows) { + _ = @import("benchmark_suite.zig"); + _ = @import("benchmark_check.zig"); + _ = @import("benchmark_receipt_test.zig"); + _ = @import("benchmark_stop_test.zig"); + _ = @import("benchmark_sse_test.zig"); + _ = @import("benchmark_stream_test.zig"); + _ = @import("benchmark_server_test.zig"); + _ = @import("benchmark_facts_test.zig"); + _ = @import("benchmark_upload.zig"); + _ = @import("benchmark_upload_test.zig"); + _ = @import("benchmark_args.zig"); + _ = @import("benchmark_cli_test.zig"); + _ = @import("benchmark_run_test.zig"); + } } diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 2fde4ce7fd..d96b96a436 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"); const auth = @import("auth.zig"); @@ -11,7 +12,7 @@ const model_text = @import("model_text.zig"); const server_mod = @import("server.zig"); const recall_routes = @import("recall_routes.zig"); const Allocator = std.mem.Allocator; -const posix = std.posix; +const posix = platform.posix; var stop_requested: std.atomic.Value(bool) = .init(false); @@ -34,7 +35,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 { @@ -225,6 +226,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 5406581774..6797c084a7 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"); const errors = @import("errors.zig"); @@ -301,7 +302,7 @@ pub const Server = struct { _ = srv.open_connections.fetchAdd(1, .acq_rel); const t = std.Thread.spawn(.{ .stack_size = 8 << 20 }, connection, .{ srv, c }) catch { _ = srv.open_connections.fetchSub(1, .acq_rel); - _ = std.posix.system.close(c.fd); + _ = platform.posix.system.close(c.fd); continue; }; t.detach(); @@ -310,7 +311,7 @@ pub const Server = struct { fn connection(srv: *Server, accepted: listener_mod.Listener.Accepted) void { 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/slide_key.zig b/zig/src/server/slide_key.zig index 9047783865..763f85771b 100644 --- a/zig/src/server/slide_key.zig +++ b/zig/src/server/slide_key.zig @@ -1,5 +1,6 @@ //! The --slide key gate: a server that rewrites its weights starts only with a key its own key store accepts. const std = @import("std"); +const platform = @import("sys"); const auth = @import("auth.zig"); const cli = @import("cli.zig"); const Allocator = std.mem.Allocator; @@ -59,10 +60,12 @@ test "--slide refuses before any load a key file or a key the server's store wou const missing = (try slideKeyless(a, keyed, "")) orelse return error.TestExpectedRefusal; try std.testing.expect(has(missing, file)); try tmp.dir.writeFile(io, .{ .sub_path = "keys", .data = "ops: sk-file\n" }); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(std.posix.AT.FDCWD, file, 0o644, 0)); - const open = (try slideKeyless(a, keyed, "")) orelse return error.TestExpectedRefusal; - try std.testing.expect(has(open, file) and has(open, "chmod 600")); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(std.posix.AT.FDCWD, file, 0o600, 0)); + if (comptime @import("builtin").os.tag != .windows) { // Windows files have no mode bits for others to read through + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(platform.c.AT.FDCWD, file, 0o644, 0)); + const open = (try slideKeyless(a, keyed, "")) orelse return error.TestExpectedRefusal; + try std.testing.expect(has(open, file) and has(open, "chmod 600")); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(platform.c.AT.FDCWD, file, 0o600, 0)); + } try std.testing.expect(try slideKeyless(a, keyed, "") == null); try tmp.dir.writeFile(io, .{ .sub_path = "keys", .data = "" }); const empty = (try slideKeyless(a, keyed, "")) orelse return error.TestExpectedRefusal; diff --git a/zig/src/xpu/loader.zig b/zig/src/xpu/loader.zig index 5a2fb9e342..4d408065b6 100644 --- a/zig/src/xpu/loader.zig +++ b/zig/src/xpu/loader.zig @@ -3,9 +3,15 @@ const std = @import("std"); const rt = @import("rt.zig"); -extern "c" fn open(path: [*:0]const u8, flags: c_int, ...) c_int; -extern "c" fn pread(fd: c_int, buf: [*]u8, n: usize, off: i64) isize; -extern "c" fn close(fd: c_int) c_int; +const platform = @import("sys"); +const pread = platform.c.pread; +const close = platform.c.close; + +/// `path` opened to read (the only mode the loader asks for, as flags 0). +fn open(path: [*:0]const u8, flags: c_int) c_int { + _ = flags; + return platform.c.open(path, .{}, @as(platform.c.mode_t, 0)); +} const chunk_bytes: usize = 64 << 20; diff --git a/zig/src/xpu/rt.zig b/zig/src/xpu/rt.zig index 169ec2fcc9..3d458b4cda 100644 --- a/zig/src/xpu/rt.zig +++ b/zig/src/xpu/rt.zig @@ -1,6 +1,7 @@ //! Prototype-style runtime view: one Runtime over a Context and its Stream, with plain buffers, modules and kernels. const std = @import("std"); +const platform = @import("sys"); const ze = @import("ze"); const xpu = @import("context.zig"); const Context = @import("context.zig").Context; @@ -12,8 +13,8 @@ const ProfEntry = struct { name: ?[*:0]const u8 = null, ns: u64 = 0, calls: u64 var prof_tab: [160]ProfEntry = @splat(.{}); fn profNow() u64 { - var ts: std.c.timespec = undefined; - _ = std.c.clock_gettime(.MONOTONIC, &ts); + var ts: platform.c.timespec = undefined; + _ = platform.c.clock_gettime(platform.c.CLOCK.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); } @@ -169,7 +170,7 @@ pub const Runtime = struct { live_list = s.list; live_sync = c.d.api.zeCommandListHostSynchronize; _ = atexit(&atExit); // exit handlers run last-in first-out: after the loader is open, so ours runs before it tears down - if (std.c.getenv("XPU_PROF") != null) prof_on = true; + if (platform.c.getenv("XPU_PROF") != null) prof_on = true; return .{ .drv = c.d, .device = c.gpu.device, .ctx = c.gpu.context, .list = s.list }; } diff --git a/zig/tests/cuda/affine8_tests.zig b/zig/tests/cuda/affine8_tests.zig new file mode 100644 index 0000000000..a0a74c1aa7 --- /dev/null +++ b/zig/tests/cuda/affine8_tests.zig @@ -0,0 +1,184 @@ +//! MLX affine widths on qlinear's tiles: tiles equal host-widened codes, products in float64 bounds, bits row-stable. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref = @import("quant_ref_tests.zig"); +const Gpu = check.Gpu; + +const qlinear = cuda.qlinear; +const Rig = quant_ref.RigT; +const bf = quant_ref.helpers.bfOf; +const bf16 = quant_ref.helpers.bf16Of; + +/// `codes` packed `bits` wide, low code first, into little-endian words (MLX's layout; a code may cross a word). +fn packCodes(bits: u8, codes: []const u8, words: []u32) void { + @memset(words, 0); + for (codes, 0..) |code, i| { + const bit = i * bits; + const v: u64 = @as(u64, code) << @intCast(bit & 31); + words[bit >> 5] |= @truncate(v); + if ((bit & 31) + bits > 32) words[(bit >> 5) + 1] |= @truncate(v >> 32); + } +} + +const Matrix = struct { codes: []u8, scales: []u16, biases: []u16, n: usize, k: usize, bits: u8 }; + +/// Random codes of `bits`, and per-64 scales and biases as a converter writes them (bias near -scale * half range). +fn matrix(r: *Rig, n: usize, k: usize, bits: u8) !Matrix { + const codes = try r.a.alloc(u8, n * k); + for (codes) |*c| c.* = @intCast(r.rng.uintLessThan(u16, @as(u16, 1) << @intCast(bits))); + const scales = try r.a.alloc(u16, n * k / 64); + const biases = try r.a.alloc(u16, n * k / 64); + const half: f32 = @floatFromInt(@as(u16, 1) << @intCast(bits - 1)); + for (scales, biases) |*s, *b| { + const sv = (r.rng.float(f32) * 0.9 + 0.1) * 0.02 / half; + s.* = bf16(sv); + b.* = bf16(-bf32(s.*) * half * (0.8 + 0.4 * r.rng.float(f32))); + } + return .{ .codes = codes, .scales = scales, .biases = biases, .n = n, .k = k, .bits = bits }; +} + +fn bf32(x: u16) f32 { + return @floatCast(bf(x)); +} + +/// The matrix's MLX words at `bits` (its own width, or a wider one its codes fit). +fn mlxWords(r: *Rig, m: Matrix, bits: u8) ![]u32 { + const out = try r.a.alloc(u32, m.n * m.k * bits / 32); + packCodes(bits, m.codes, out); + return out; +} + +/// y = x (s q + b)^T in float64, and its bound: a bf16 ulp of |y| and 2^-12 of sum |x| (|s q| + |b|). +fn reference(r: *Rig, x: []const u16, rows: usize, m: Matrix) !struct { want: []f64, tol: []f64 } { + const want = try r.a.alloc(f64, rows * m.n); + const tol = try r.a.alloc(f64, rows * m.n); + for (0..rows) |row| for (0..m.n) |j| { + var sum: f64 = 0; + var abs: f64 = 0; + for (0..m.k) |i| { + const g = j * (m.k / 64) + i / 64; + const s = bf(m.scales[g]); + const b = bf(m.biases[g]); + const q: f64 = @floatFromInt(m.codes[j * m.k + i]); + const xv = bf(x[row * m.k + i]); + sum += xv * (s * q + b); + abs += @abs(xv) * (@abs(s * q) + @abs(b)); + } + want[row * m.n + j] = sum; + tol[row * m.n + j] = @abs(sum) * std.math.pow(f64, 2, -7) + abs * std.math.pow(f64, 2, -12) + 1e-30; + }; + return .{ .want = want, .tol = tol }; +} + +/// tf_q35_sums's 64-group sums: fp32, in order. +fn sums(r: *Rig, x: []const u16, rows: usize, k: usize) ![]f32 { + const out = try r.a.alloc(f32, rows * k / 64); + for (0..rows) |row| for (0..k / 64) |g| { + var s: f32 = 0; + for (0..64) |j| s += bf32(x[row * k + g * 64 + j]); + out[row * (k / 64) + g] = s; + }; + return out; +} + +const Kernels = struct { a4: *const qlinear.Affine4Kernels, a8: *const qlinear.Affine8Kernels }; + +/// One width at one shape: tiles against the host-widened codes' tiles, products against float64, rows and paths. +fn width(r: *Rig, k_: Kernels, n: usize, k: usize, bits: u8) !void { + var name: [128]u8 = undefined; + const m = try matrix(r, n, k, bits); + const served = qlinear.servedBits(bits); + const own = try r.dev(std.mem.sliceAsBytes(try mlxWords(r, m, bits))); + const wide = try r.dev(std.mem.sliceAsBytes(try mlxWords(r, m, served))); + const ds = try r.dev(std.mem.sliceAsBytes(m.scales)); + const db = try r.dev(std.mem.sliceAsBytes(m.biases)); + const l4 = qlinear.Affine4.layout(n, k); + const l8 = qlinear.Affine8.layout(n, k); + const words_len = if (served == 4) l4.words else l8.words; + const npad = l4.npad; + const t_own = try r.zeros(words_len); + const t_wide = try r.zeros(words_len); + const ts = try r.zeros(l4.scales); + const tb = try r.zeros(l4.scales); + try k_.a8.packWords(r.s, own, bits, n, k, t_own); + if (served == 4) { + // the affine-4 pack of the codes widened to 4 bits: the tiles a 4-bit checkpoint of these codes gets + try k_.a4.pack(r.s, wide, ds, db, .{ .w = t_wide, .s = ts, .b = tb, .n = n, .k = k, .npad = npad }); + } else { + try k_.a8.packWords(r.s, wide, 8, n, k, t_wide); + try k_.a8.packTable(r.s, ds, n, k, npad, ts); + try k_.a8.packTable(r.s, db, n, k, npad, tb); + } + try check.sameBytes(try std.fmt.bufPrint(&name, "{d}-bit codes [{d}, {d}]: tiles equal the {d}-bit widened codes' tiles", .{ bits, n, k, served }), try r.back(t_own, words_len), try r.back(t_wide, words_len)); + // the embedding's widen: rows of `bits` codes into rows of the served width, word for word + const widened = try r.zeros(n * k * served / 8); + try k_.a8.widen(r.s, own, bits, n, k, widened, served); + try check.sameBytes(try std.fmt.bufPrint(&name, "{d}-bit rows [{d}, {d}] widened to {d} bits", .{ bits, n, k, served }), try r.back(widened, n * k * served / 8), try r.back(wide, n * k * served / 8)); + // products: 16 rows against float64, fewer rows byte-equal to the 16-row call's first rows + const most = 16; + const x = try r.bfs(most * k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dxs = try r.dev(std.mem.sliceAsBytes(try sums(r, x, most, k))); + const ref = try reference(r, x, most, m); + const counts = [_]usize{ 1, 3, 9, most }; + var outs: [counts.len]u64 = undefined; + for (counts, &outs) |rows, *o| { + o.* = try r.zeros(rows * n * 2); + if (served == 4) { + try k_.a4.decode(r.s, dx, dxs, .{ .w = t_own, .s = ts, .b = tb, .n = n, .k = k, .npad = npad }, o.*, rows); + } else { + try k_.a8.decode(r.s, dx, dxs, .{ .w = t_own, .s = ts, .b = tb, .n = n, .k = k, .npad = npad }, o.*, rows); + } + } + try r.close(try std.fmt.bufPrint(&name, "{d}-bit [{d}, {d}] decode, {d} rows", .{ bits, n, k, most }), outs[counts.len - 1], false, ref.want, ref.tol); + for (counts[0 .. counts.len - 1], outs[0 .. counts.len - 1]) |rows, o| { + const a = try r.back(outs[counts.len - 1], rows * n * 2); + const b = try r.back(o, rows * n * 2); + try check.sameBytes(try std.fmt.bufPrint(&name, "{d}-bit [{d}, {d}]: {d} rows equal the 16-row call's", .{ bits, n, k, rows }), b, a); + } + // the 8-bit tiles' two paths: split-K CTAs and lane_gemv's loop over slices give one row the same bits + if (served == 8) if (k_.a8.split != null and qlinear.splitK(n, k) > 1) { + var plain = k_.a8.*; + plain.split = null; + const o = try r.zeros(most * n * 2); + try plain.decode(r.s, dx, dxs, .{ .w = t_own, .s = ts, .b = tb, .n = n, .k = k, .npad = npad }, o, most); + try check.sameBytes(try std.fmt.bufPrint(&name, "{d}-bit [{d}, {d}]: split-K equals lane_gemv's slices", .{ bits, n, k }), try r.back(o, most * n * 2), try r.back(outs[counts.len - 1], most * n * 2)); + }; +} + +pub fn run(gpu: Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(17); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + const kk = cuda.kernels; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.affine8 }; + var mods: [images.len]cuda.Module = undefined; + var loaded: usize = 0; + defer for (mods[0..loaded]) |*md| md.unload(); + for (images, &mods) |img, *md| { + md.* = try cuda.Module.load(gpu.d, img); + loaded += 1; + } + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + var a4 = try qlinear.Affine4Kernels.resolve(gpu.d, mods[0], mods[1], mods[2], mods[3], sms, false); + defer a4.deinit(); + var a8 = try qlinear.Affine8Kernels.resolve(gpu.d, mods[4], mods[3], sms, false); + defer a8.deinit(); + const k_: Kernels = .{ .a4 = &a4, .a8 = &a8 }; + // shapes: one K slice and several (split-K), rows short of a tile, a 64-row tile past the 128 padding + const shapes = [_][2]usize{ .{ 320, 2560 }, .{ 2560, 512 }, .{ 200, 1024 }, .{ 48, 2048 } }; + var cases: usize = 0; + for ([_]u8{ 2, 3, 4, 5, 6, 8 }) |bits| for (shapes) |nk| { + try width(&r, k_, nk[0], nk[1], bits); + cases += 1; + }; + check.pass("affine8: {d} cases over 2, 3, 4, 5, 6 and 8 bits: tiles equal the widened codes', products within {d:.3} of the float64 bound, every row count and path byte-equal", .{ cases, r.worst }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..418cfce5dd 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -21,6 +21,7 @@ const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); const glm_ref_tests = @import("glm_ref_tests.zig"); +const affine8_tests = @import("affine8_tests.zig"); const carveout_tests = @import("carveout_tests.zig"); const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); @@ -81,6 +82,7 @@ const usage = \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) + \\ affine8 MLX affine 2 to 8-bit tiles and the 8-bit lane_gemv: widened codes, float64, rows (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) \\ nvfp4 NVFP4 projections against recorded lane-matmul and prompt-GEMM fixture bytes \\ modelopt ModelOpt repack, FP8 and BF16 lane rows and relu^2 NVFP4 experts against fp64 @@ -157,6 +159,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-ref")) return glm_ref_tests.run(gpu); + if (std.mem.eql(u8, cmd, "affine8")) return affine8_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); diff --git a/zig/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig index 1e9f5162f2..427e387fa2 100644 --- a/zig/tests/cuda/modelopt_tests.zig +++ b/zig/tests/cuda/modelopt_tests.zig @@ -193,22 +193,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 59be964440..761207534f 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; @@ -56,7 +57,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 { @@ -127,13 +128,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); @@ -142,18 +143,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; @@ -171,10 +172,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; @@ -185,7 +186,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", .{}); @@ -195,17 +196,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(); @@ -219,18 +220,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}); @@ -244,7 +245,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 }); @@ -256,23 +257,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..e0bb3ac025 --- /dev/null +++ b/zig/tests/qwen35_cuda_main.zig @@ -0,0 +1,137 @@ +//! `tf-qwen35-cuda `: the ids' logits as bf16 rows for a reference. +const std = @import("std"); +const cuda = @import("cuda"); +const q = @import("qwen35"); +const E = q.engine; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const args = try init.minimal.args.toSlice(a); + if (args.len != 5) { + std.debug.print("usage: tf-qwen35-cuda \n", .{}); + return 2; + } + var ids: std.ArrayList(u32) = .empty; + var it = std.mem.tokenizeAny(u8, args[3], ", "); + while (it.next()) |t| try ids.append(a, try std.fmt.parseInt(u32, t, 10)); + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + const e = try E.Engine.init(init.gpa, init.io, &ctx, args[1]); + defer e.deinit(); + // TF_Q35_PREFILL=N starts the stream after N cycled ids; TF_Q35_BESIDE=N steps an N-token stream beside it + const prefill = try envCount("TF_Q35_PREFILL"); + const beside = try envCount("TF_Q35_BESIDE"); + var cache = try E.Cache.init(init.gpa, e, 4096 + prefill); + defer cache.deinit(init.gpa); + try fill(e, &cache, ids.items, prefill); + var other = try E.Cache.init(init.gpa, e, 4096 + beside); + defer other.deinit(init.gpa); + try fill(e, &other, ids.items, beside); + var f: Fwd = .{ .e = e, .other = if (beside > 0) &other else null, .ids = ids.items }; + var out: std.ArrayList(u8) = .empty; + if (std.mem.eql(u8, args[2], "step")) { + var first: i64 = 0; + var rest: i64 = 0; + for (ids.items, 0..) |id, i| { + const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + try f.run(.{ .cache = &cache, .rows = 1 }, &.{id}); + const dt: i64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); + if (i == 0) first = dt else rest += dt; + @memcpy(cache.logits, e.rowLogits(0)); + try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); + } + const n: f64 = @floatFromInt(ids.items.len - 1); + std.debug.print("{d} steps: the first {d:.1} ms, then {d:.1} ms a step ({d:.1} tok/s)\n", .{ ids.items.len, @as(f64, @floatFromInt(first)) / 1e6, @as(f64, @floatFromInt(rest)) / 1e6 / n, n * 1e9 / @as(f64, @floatFromInt(rest)) }); + } else if (std.mem.eql(u8, args[2], "window")) { + // windows of 1, 2, 3, ... rows: every row's logits, which must equal step's bit for bit + var at: usize = 0; + var w: usize = 1; + while (at < ids.items.len) : (w = w % e.window_rows + 1) { + const n = @min(w, ids.items.len - at); + try f.run(.{ .cache = &cache, .rows = n }, ids.items[at..][0..n]); + for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); + at += n; + } + } else if (std.mem.eql(u8, args[2], "tree")) { + // tree windows of growing widths, junk branches beside the kept true one, whose logits must equal step's bits + var at: usize = 0; + var w: usize = 2; + while (at < ids.items.len) : (w = w % e.window_rows + 1) { + var toks: [E.most_window]u32 = undefined; + var parents: [E.most_window]i32 = undefined; + var path: [E.most_window]u32 = undefined; + toks[0] = ids.items[at]; + parents[0] = -1; + path[0] = 0; + var n: usize = 1; + var kept: usize = 1; + var last: i32 = 0; + var depth: usize = 1; + while (n + 2 <= w and at + depth < ids.items.len) : (depth += 1) { + toks[n] = (ids.items[at + depth] + 7) % @as(u32, @intCast(e.g.vocab)); + parents[n] = last; + toks[n + 1] = ids.items[at + depth]; + parents[n + 1] = last; + path[kept] = @intCast(n + 1); + last = @intCast(n + 1); + n += 2; + kept += 1; + } + if (n < w and n >= 3) { + toks[n] = 1; + parents[n] = 1; + n += 1; + } + try f.run(.{ .cache = &cache, .rows = n, .parents = parents[0..n] }, toks[0..n]); + for (path[0..kept]) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); + if (n > 1) try cache.keep(e, path[0..kept]); + at += kept; + } + } else { + var at: usize = 0; + while (at < ids.items.len) { + const n = @min(E.max_rows, ids.items.len - at); + try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .last); + at += n; + } + try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); + } + try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[4], .data = out.items }); + return 0; +} + +/// One forward keeping logits: the stream's part first (its rows' logits come first), then the second stream's step. +const Fwd = struct { + e: *E.Engine, + other: ?*E.Cache, + ids: []const u32, + n: usize = 0, + + fn run(f: *Fwd, part: E.Part, toks: []const u32) !void { + const o = f.other orelse return f.e.forward(&.{part}, toks, .logits); + var all: [E.most_window + 1]u32 = undefined; + @memcpy(all[0..toks.len], toks); + all[toks.len] = f.ids[f.n % f.ids.len]; + f.n += 1; + try f.e.forward(&.{ part, .{ .cache = o, .rows = 1 } }, all[0 .. toks.len + 1], .logits); + } +}; + +fn envCount(name: [:0]const u8) !usize { + const v = std.c.getenv(name) orelse return 0; + return std.fmt.parseInt(usize, std.mem.span(v), 10); +} + +/// `n` positions of the ids cycled, through prompt forwards. +fn fill(e: *E.Engine, cache: *E.Cache, ids: []const u32, n: usize) !void { + var buf: [E.max_rows]u32 = undefined; + var at: usize = 0; + while (at < n) { + const k = @min(E.max_rows, n - at); + for (buf[0..k], at..) |*t, i| t.* = ids[i % ids.len]; + try e.forward(&.{.{ .cache = cache, .rows = k }}, buf[0..k], .last); + at += k; + } +} diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index 184c401968..73d607350f 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -1,6 +1,8 @@ //! Frozen server goldens for `zig build test-golden`, with the ten differences that are still known. const std = @import("std"); -const posix = std.posix; +const builtin = @import("builtin"); +const platform = @import("sys"); +const posix = platform.posix; const Io = std.Io; const wire = @import("wire.zig"); const Reply = wire.Reply; @@ -21,6 +23,38 @@ const known_diffs = [_]Known{ .{ .group = "anthropic", .name = "tool-use-stream" }, }; +/// Why two request logs differ: their line counts and the first line that does, each cut to 100 bytes. +fn logDiff(a: std.mem.Allocator, want: []const u8, got: []const u8) ![]const []const u8 { + var out: std.ArrayList([]const u8) = .empty; + var wl = std.mem.splitScalar(u8, want, '\n'); + var gl = std.mem.splitScalar(u8, got, '\n'); + var n: usize = 0; + while (true) : (n += 1) { + const w = wl.next(); + const g = gl.next(); + if (w == null and g == null) break; + if (w != null and g != null and std.mem.eql(u8, w.?, g.?)) continue; + const ws = w orelse ""; + const gs = g orelse ""; + try out.append(a, try std.fmt.allocPrint(a, "line {d}: golden [{s}] zig [{s}]", .{ n + 1, ws[0..@min(ws.len, 100)], gs[0..@min(gs.len, 100)] })); + break; + } + try out.append(a, try std.fmt.allocPrint(a, "TENSORFOLD_REQUEST_LOG: golden {d} lines, zig {d} lines", .{ std.mem.count(u8, want, "\n"), std.mem.count(u8, got, "\n") })); + return out.items; +} + +/// Reads a spawned server's stdout to its end, so its log lines never fill the pipe (4 KiB on Windows). +fn drainStdout(io: Io, out: Io.File) void { + var buf: [4096]u8 = undefined; + var reader = out.reader(io, &buf); + while (true) { + _ = reader.interface.discardShort(buf.len) catch return; + } +} + +/// The parent's SystemRoot, which a Windows child needs for Winsock; null elsewhere. +var system_root: ?[]const u8 = null; + fn isKnown(group: []const u8, name: []const u8) bool { for (known_diffs) |k| { if (std.mem.eql(u8, k.group, group) and std.mem.eql(u8, k.name, name)) return true; @@ -354,6 +388,7 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con } else try argv.append(a, arg); } var map = std.process.Environ.Map.init(a); + if (system_root) |root| try map.put("SystemRoot", root); // Winsock refuses to start without it try map.put("TENSORFOLD_NO_LIVE", "1"); try map.put("TF_FAKE_CONTEXT", context_tokens); var it = env.iterator(); @@ -379,7 +414,7 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con try fw.interface.writeAll(content); try fw.interface.flush(); const kpz = try std.fmt.allocPrintSentinel(a, "{s}", .{kp}, 0); - _ = std.c.fchmodat(posix.AT.FDCWD, kpz, 0o600, 0); + _ = platform.c.fchmodat(posix.AT.FDCWD, kpz, 0o600, 0); } } const child = try a.create(std.process.Child); @@ -397,12 +432,16 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con const line = reader.interface.takeDelimiterExclusive('\n') catch return error.NoPortLine; if (std.mem.startsWith(u8, line, "PORT ")) { const port = try std.fmt.parseInt(u16, std.mem.trim(u8, line[5..], " \r"), 10); + // the server's log lines keep coming on stdout: drained, or a full pipe blocks the server + const drain = try std.Thread.spawn(.{}, drainStdout, .{ io, child.stdout.? }); + drain.detach(); return .{ .child = child, .port = port, .log_path = log_path, .key_path = key_path }; } } } pub fn main(init: std.process.Init) !void { + if (builtin.os.tag == .windows) system_root = init.environ_map.get("SystemRoot"); const gpa = init.gpa; const io = init.io; const args = try init.minimal.args.toSlice(init.arena.allocator()); @@ -411,7 +450,7 @@ pub fn main(init: std.process.Init) !void { const cases_path = args[2]; const golden_dir = args[3]; const fixtures = args[4]; - const tmp_dir = try std.fmt.allocPrint(init.arena.allocator(), ".zig-tmp-golden-{d}", .{std.c.getpid()}); + const tmp_dir = try std.fmt.allocPrint(init.arena.allocator(), ".zig-tmp-golden-{d}", .{platform.c.getpid()}); Io.Dir.cwd().createDir(io, tmp_dir, .default_dir) catch {}; defer Io.Dir.cwd().deleteTree(io, tmp_dir) catch {}; @@ -435,15 +474,26 @@ pub fn main(init: std.process.Init) !void { if (obj.get("cases")) |cases| { for (cases.array.items) |case| try runCase(init.arena.allocator(), io, golden_dir, server.?.port, case); } + var after_signal = false; // Windows: a step's cases after a signal it cannot send are skipped, and said so for (obj.get("then").?.array.items) |then| { const t = then.object; const kind = t.get("kind").?.string; const s = &server.?; if (std.mem.eql(u8, kind, "cases")) { + if (after_signal) { + std.debug.print("SKIPPED {d} cases after a signal Windows cannot send\n", .{t.get("cases").?.array.items.len}); + continue; + } for (t.get("cases").?.array.items) |case| try runCase(init.arena.allocator(), io, golden_dir, s.port, case); } else if (std.mem.eql(u8, kind, "sleep")) { std.Io.sleep(io, .fromMilliseconds(@intFromFloat(t.get("seconds").?.float * 1000.0)), .awake) catch {}; } else if (std.mem.eql(u8, kind, "sigterm")) { + if (builtin.os.tag == .windows) { + std.debug.print("SKIPPED lifecycle/sigterm-exit-0: Windows has no SIGTERM to send, the server is ended\n", .{}); + posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; + _ = s.child.wait(io) catch {}; + continue; + } posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; const term = s.child.wait(io) catch return error.WaitFailed; try outcomes.append(init.arena.allocator(), .{ @@ -454,13 +504,14 @@ pub fn main(init: std.process.Init) !void { }); } else if (std.mem.eql(u8, kind, "compare_log")) { const golden_log = try std.fmt.allocPrint(init.arena.allocator(), "{s}/{s}", .{ golden_dir, t.get("golden").?.string }); - const want = Io.Dir.cwd().readFileAlloc(io, golden_log, init.arena.allocator(), .limited(1 << 28)) catch ""; - const got = Io.Dir.cwd().readFileAlloc(io, s.log_path.?, init.arena.allocator(), .limited(1 << 28)) catch ""; + // a checkout with autocrlf gives the frozen log CRLF endings, the server writes LF: both compare as LF + const want = try std.mem.replaceOwned(u8, init.arena.allocator(), Io.Dir.cwd().readFileAlloc(io, golden_log, init.arena.allocator(), .limited(1 << 28)) catch "", "\r\n", "\n"); + const got = try std.mem.replaceOwned(u8, init.arena.allocator(), Io.Dir.cwd().readFileAlloc(io, s.log_path.?, init.arena.allocator(), .limited(1 << 28)) catch "", "\r\n", "\n"); try outcomes.append(init.arena.allocator(), .{ .group = "lifecycle", .name = "request-log", .verdict = if (std.mem.eql(u8, want, got)) .equal else .differ, - .why = if (std.mem.eql(u8, want, got)) &.{} else &.{"TENSORFOLD_REQUEST_LOG lines differ from the frozen ones"}, + .why = if (std.mem.eql(u8, want, got)) &.{} else try logDiff(init.arena.allocator(), want, got), }); } else if (std.mem.eql(u8, kind, "rewrite_key")) { var file = try Io.Dir.cwd().createFile(io, s.key_path.?, .{}); @@ -470,8 +521,10 @@ pub fn main(init: std.process.Init) !void { try fw.interface.writeAll(t.get("content").?.string); try fw.interface.flush(); posix.kill(s.child.id.?, posix.SIG.HUP) catch {}; // the server rereads the key file, as the reference does + after_signal = builtin.os.tag == .windows; } else if (std.mem.eql(u8, kind, "sighup")) { posix.kill(s.child.id.?, posix.SIG.HUP) catch {}; + after_signal = builtin.os.tag == .windows; } else if (std.mem.eql(u8, kind, "terminate")) { posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; _ = s.child.wait(io) catch {}; diff --git a/zig/tests/train_check.zig b/zig/tests/train_check.zig index 08601733c3..77acc8e7c1 100644 --- a/zig/tests/train_check.zig +++ b/zig/tests/train_check.zig @@ -1,5 +1,6 @@ //! The learning step's gradients against the engine's own forward, moving each layer along its own, on the real model. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const nm = tf.nemotron; @@ -51,8 +52,8 @@ pub fn main(init: std.process.Init) !void { const with_grad = (try t.step(&ids, start, .grad, turn)).loss; const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0.toNanoseconds())) / 1e6; std.debug.print("loss {d:.5} (with gradients {d:.5}), a step with gradients {d:.1} ms, GPU {d:.1} ms\n", .{ plain, with_grad, ms, b.last_ms }); - if (std.c.getenv("TF_TRAIN_TIME_ONLY") != null) return; - if (std.c.getenv("TF_TRAIN_PROFILE") != null) { + if (platform.c.getenv("TF_TRAIN_TIME_ONLY") != null) return; + if (platform.c.getenv("TF_TRAIN_PROFILE") != null) { var prof: nm.encoder.Profiler = .{ .queue = m.queue, .kernels = &m.kernels }; t.profile = &prof; _ = try t.step(&ids, start, .grad, turn); diff --git a/zig/tests/weight_read_check.zig b/zig/tests/weight_read_check.zig index d5acf5e7a6..1089d93666 100644 --- a/zig/tests/weight_read_check.zig +++ b/zig/tests/weight_read_check.zig @@ -1,5 +1,6 @@ //! Check the actual native loaders with a synthetic file and bounded I/O faults. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -63,7 +64,7 @@ pub fn main(init: std.process.Init) !void { } else { try replay.indexFile(path); const entry = try replay.entry(if (grouped) "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight" else "probe"); - defer _ = std.c.close(entry.fd); + defer _ = platform.c.close(entry.fd); const result = if (grouped) replay.group(0, 2, "weight") else replay.load("probe"); const buffer = result catch |err| { if (!eof or err != error.ShortRead) return err;