Repository navigation
qwen3_5 cuda: --kv-cache fp8, e4m3 attention caches with a scale a row - #886
BobClawblaw wants to merge 35 commits into
Conversation
- zig/src/platform: `sys`, imported by every module (zig/build/sys.zig wires it). On POSIX it is std.c, std.posix and std.DynLib unchanged; on Windows win_c.zig gives libc's surface over the UCRT and Win32 (integer descriptors, directory descriptors, readdir, pread/pwrite, high-resolution sleep, monotonic clock, /tmp and /proc/self/exe mapped), win_posix.zig the server's sockets over Winsock, console events as SIGINT/SIGTERM, and LoadLibraryW for DynLib. - std.c and std.DynLib calls go through `platform` (mechanical). - build.zig routes Windows targets to the CUDA build. - Server: socket closes and writes are explicit, the key file's mode check is skipped on NTFS (ACLs), descriptor limits are a no-op. - shard_edit: the shard's map is released before the rename, since Windows refuses to replace a mapped file. - Tests that need symlink privilege or mode bits skip on Windows. zig build test-cuda-host passes on x86_64-windows.
- Kernel fatbins build with nvcc and cl.exe: MSVC's conforming preprocessor is passed on Windows hosts (CUDA 13's headers need it), and explicit kernel instantiations carry the __restrict__ qualifiers their templates have (MSVC counts them in the function type; the generated code is unchanged). - Nemotron's checkpoint source keeps io_uring to Linux; elsewhere its existing thread reads run. - pull_exlr8 raises no descriptor limit on Windows. - modelopt test: the lane's slice partials are allocated. GPUs without clusters (before sm_90) reduce split K through them; the test passed none, so it wrote through a null pointer on Ampere and Ada. zig build native and install build with -Dnvcc=... -Dsm=89, and tf-cuda-test's synthetic suites pass on an RTX 4090.
- pull: where Windows refuses the snapshot's symlink, the snapshot path
becomes a hard link to its blob (CreateHardLinkW through `sys.c.link`):
no extra space, and learning already refuses linked shards. Link
failures are logged with their reason.
- HOME falls back to USERPROFILE (Windows sets no HOME) for the hub cache
and every libc getenv("HOME").
- Tests read optional fixture variables through `platform.getenv`;
std.process.Environ.getPosix does not build on Windows.
- The learned-state symlink tests skip on Windows (symlinks need
privilege); `sys.c.symlink` asks for an unprivileged one.
zig build test passes on x86_64-windows.
- docs/WINDOWS.md: requirements, build, run, how the port works and how it differs from Linux; the README points to it. - -Dstrip reaches tensorfold-native before its install step is made, so the step no longer expects a .pdb the stripped build does not write. - zig fmt on the platform files.
8080 and 8000 are taken by too many other servers; serve listens on 8888 unless --port says otherwise, and the docs and bench scripts follow.
The MLX affine 4-bit checkpoint on the native CUDA server, beside Nemotron: qwen35.cu holds the Metal engine's glue at its bf16 rounding points (gated delta rule, causal convolution, gated attention), cuda.qlinear's affine-4 kernels do the projections, sample.cu draws on the device. Windows of up to 16 rows record each row's recurrence and convolution tail, so the lane core shares rounds across streams and keeps partial copies; a window's rows equal one-row steps bit for bit. draw.zig is the host draw both backends now share. On an RTX 4090: 210 tok/s for one stream, 1250 tok/s across eight; logits agree with an fp32 reference at every position of a 36-token check.
…mory The 2B's CUDA engine now reads its widths from config.json (cuda_geometry.zig): linear value heads over fewer key heads, more query heads a key/value head, an untied head, an A_log stored as bf16. The 27B's logits agree with an fp32 reference at every position of a 36-token check, and its windows equal one-row steps bit for bit; the 2B's output is unchanged. Two Windows driver costs shaped the rest. Each kernel launch grew with the process's allocation count (1.2 ms a launch over 1,900 weight buffers), so a model's weights and scratch now live in one arena (cuda_weights.zig packs each projection through a reused staging buffer: the 27B loads in 25 s, not 200). And a process's first forwards through a buffer stall for seconds, so the engine warms every request path at load and the backend keeps a pool of stream caches that `reserve` allocates and runs a row through once the server knows its stream count. On an RTX 4090: the 27B decodes at 37 tok/s in the runner, 31 to 35 tok/s served, 210 tok/s across eight streams; the 2B stays at 210 and 1,250.
A forward's launches (the embedding to the kept rows' logits) are captured the first time a shape (rows, windows, recording) is seen and replayed after; the attention chunk count moves into the per-forward table so a graph stays valid as caches grow. Prompt chunks stay uncaptured: their GEMMs dwarf the launches, and each length would be a graph of its own. A new graph's first run pays the Windows driver's stall (about a second on a wide model), so `reserve` also captures the shared rounds' shapes, a row of each of 1 to n streams, at load. Bit for bit the same logits; on an RTX 4090 the 2B steps in 2.8 ms instead of 5.3 (served: 324 tok/s, 1,560 across eight streams) and the 27B in 23.8 ms instead of 26.9 (served: 40 tok/s, 243 across eight).
--drafter loads z-lab's DFlash2 draft model behind the lane core's Drafted wrapper: the target writes each row's residual after the drafter's five tapped layers into a ring a stream, the drafter folds those rows into a context of keys and values, and after every verify it drafts a block of seven tokens past the pending one (the target's own embedding and head at its ends, a codebook selector tracing one path through each row's best sixteen). Its projections are quantized to affine-4 on the host at load and run through the target's kernels; its glue is qwen35_draft.cu. The Windows driver stalls a buffer's first kernel and a new graph's first run for seconds, so the drafter pools its contexts and a captured forward now serves every verify shape: its launches cover a bucket of rows and every window, and the kernels read the true counts from the device. On an RTX 4090 the 27B serves 52 to 58 tok/s greedy with drafts against 40 plain, about 100 tok/s on long replies, the same tokens either way; plain decoding keeps its 24 ms step. The 2B is unchanged.
A verify window's rows may form a tree (each row's parent given): the convolution taps, the delta rule and attention follow each row's own branch, and the DFlash2 drafter fills the lane core's planned lanes (each lane its parent's rank-r choice by the codebook selector), so a lone stream verifies up to 12 lanes a round and the planner learns which shapes land. A verify no longer snapshots every row's state (1.3 GB shared, which capped the 27B at 8 rows): each stream logs its rows' decay, write and key, and keep replays the kept path onto the recurrence in one pass (and moves a tree's kept keys down), so any window width fits. One-row steps and prompts update in place as before. The lane host times the target's windows at load, which the depth rule had lacked (it had been capping chains at 3 drafts). The runner's new tree mode verifies tree windows with junk branches against step logits (bit-exact on both models).
What the port serves beyond Nemotron: Qwen3.5-2B and Qwen3.8-27B with its DFlash2 drafter, the speeds measured on an RTX 4090, the load time and why, the timing and chain-only switches, and what is not there.
The lean rule (tools/zig/lean_check.py) counts a comment that runs on to a second line; the Windows port's new files had a hundred such runs. Each is one line of at most 120 columns now, and the checker reports the same 539 problems as main.
direct_io opened files with O_DIRECT on Linux alone. On Windows the same open now asks for FILE_FLAG_NO_BUFFERING (which win_c maps O_DIRECT to), and the 4 KiB-aligned spans it already reads satisfy the sector alignment that needs; a file system that refuses it still falls back to buffered reads. The direct-read test runs on Windows too, checking the executable's bytes against buffered reads of it.
SIG.IGN is a function pointer made from 1, as std.posix has it. aarch64 requires function pointers to be aligned unless the type says align(1), so the cross-build for aarch64-windows refused it. With align(1), as std's own handler type has, tensorfold-native cross-compiles for aarch64-windows with the GB10 (sm_121) fatbins; no hardware has run it.
The Windows page opens with what the build is, what it is tested on and that it cross-compiles for aarch64-windows; the README's model table and the CUDA paragraph say Windows where they said unofficial port.
The no-GPU workflow runs zig build test on windows-latest beside ubuntu-latest, so the platform layer is built and tested on every pull request.
Every other module calls libc through `sys`; the drafter's four getenv calls were the last direct std.c calls in a file the Windows build compiles.
After a prompt the lane core hands the drafter every prompt row's features; the CUDA backend refused any request wider than its tap ring (2,176 rows), so a 27B request with --drafter and a prompt past that ended with RowsNotHeld, where the same prompt without the drafter ran. The ring keeps the newest tap_rows rows, which hold the drafter's whole context window, and absorb already reads only that window of what it is handed; the backend now refuses only rows not yet written or all gone. On an RTX 4090 the 27B with its drafter takes a 6,041-token prompt in 7.4 s, and its reply equals the one with "draft": false.
… attention cuts chunks at logical positions - rowanc held a window's rows as bits of a u32, but prompt chunks run up to 128 rows: rows from 32 on rotated at the wrong positions and attended the wrong keys. A chain row (no parents) now carries its index in its window under the top bit; tree windows (at most 16 rows) keep their bits. - tf_q35_attn_partial cut its 128-key chunks at a key's cache row, so a kept tree path whose window crossed a chunk boundary summed in other chunks than one-row steps and changed bits. Chunks are now cut at the key's logical position (the stream's length plus the ancestors before it). Checked on an RTX 5090 (Linux, sm_120), tf-qwen35-cuda over 300 ids on Qwen3.5-2B and Qwen3.8-27B: step bits unchanged, window and tree equal step byte for byte (tree differed from row 126 on the 2B and 254 on the 27B before). Prompt mode's last row against step's at 16 to 300 ids: KL at most 0.003 and the same argmax in 20 of 20 (before: 27B at 40 ids KL 2.3 with a different argmax, at 128 ids KL 1.35).
zig build test-golden and zig build fuzz-cortheon exist on every non-macOS host since the golden step moved into zig/build/golden.zig, so the golden check compiles against the platform layer on Windows: - the platform's posix gains sockaddr.in, timeval, connect and kill for a child's handle (TerminateProcess; HUP has no equivalent), setsockopt turns a timeval for SO_RCVTIMEO or SO_SNDTIMEO into Winsock's milliseconds, and a socket read that times out is WouldBlock, as std.posix reports it. - the check carries SystemRoot into the spawned server's environment, which Winsock needs to start, and drains the server's stdout after the port line on every host: the fake server logs each request there, and a full pipe (4 KiB on Windows) blocked it after sixteen cases. - on Windows the check cannot send SIGTERM or SIGHUP, so the sigterm-exit-0 outcome and the cases after a key rewrite are skipped and printed as such. On Windows the goldens give 306 of 315 equal with the same nine known differences as Linux, and the fuzz step passes its eight tests; the no-GPU workflow's Windows job runs both.
The loader tried only libcublasLt.so.13 and libcublasLt.so, so the cublaslt suite of tf-cuda-test reported LibraryUnavailable on Windows; there it opens cublasLt64_13.dll from the toolkit's bin folder.
The platform handed O_APPEND to the C runtime as _O_APPEND, whose write seeks to the end and then writes, so two threads appending at once could overwrite each other's bytes: the request log lost lines under concurrent requests, which the server goldens' request-log case caught on a slow runner. A descriptor opened with O_APPEND now writes through WriteFile at the end-of-file offset, which the kernel places whole, as POSIX does; the file keeps its write access, so ftruncate on such a descriptor still works (prompt_imprint repairs a torn index that way). A request_log test appends 200 lines from each of 8 threads and expects every line whole; it lost lines on Windows before this change.
The lifecycle/request-log outcome said only that the lines differ; it now gives both logs' line counts and the first line that differs, cut to 100 bytes each, so a mismatch seen on a CI runner can be read from its log.
A checkout with core.autocrlf, as GitHub's Windows runners make, gives the frozen request log CRLF line endings while the server writes LF, so the lifecycle/request-log case differed at its first line's end on windows-latest and nowhere else. Both logs compare with CRLF read as LF; the other goldens are parsed as JSON and never saw the endings.
…'s widest stream The chunk width came from the longest stream or window in the forward, so a 201-key row cut its keys into 128-key chunks alone and 256-key chunks beside a stream past 8192, and a window crossing 8192 recut its earlier rows; either changed the fp32 sums. Each row now takes chunk_keys(its logical keys), the width its one-row step uses, in both the partial and the merge kernels, and the partials sit at a fixed max_chunks stride. tf-qwen35-cuda takes TF_Q35_PREFILL=N (start after N cycled ids) and TF_Q35_BESIDE=N (step an N-token stream in each forward), which cover both cases: before this change, windows from 8151 keys kept 36 of 120 rows and trees 35, and a 151-key stream beside one from 8189 kept 4 to 11 of 120; after it all of them equal one-row steps on the 2B and 27B.
…fits its selector rank The integer reader accepted zero, so a conv_group_size or num_key_value_heads of 0 reached a remainder at admission; sizes and divisors are now positive (mask_token_id alone may be 0). The host staging area held 16 rows at rank 256, so a larger selector_rank overran it; it is now sized for the admitted rank. Config.parse takes the bytes, so a host test covers both.
Each module loads at the width the config's quantization table names it. Codes of 2 to 4 bits pack onto the affine-4 tiles and codes of 5 to 8 bits onto new 8-bit tiles, every code keeping its value, scale and bias, so nothing is rounded. affine8.cu's lane GEMV is lane_gemv's arithmetic a byte wide (each byte exact in bf16), with its split-K form; an odd-width embedding is widened at load and 8-bit rows have their own embed kernel. A uniform 4-bit checkpoint takes the same packs and kernels as before: the 2B's step, window, tree and prompt logits are byte-equal to the previous build's. The rope config's older "type" key is read as "rope_type" (mlx-community's OptiQ conversions write it). Checked on an RTX 4080 SUPER (sm_89): - tf-cuda-test affine8: 24 cases over 2, 3, 4, 5, 6 and 8 bits; tiles equal the host-widened codes' tiles, 8-bit products within 0.32 of the float64 bound, 1/3/9/16-row calls and split-K byte-equal. - A 2B requantized at 2, 3, 5, 6 and 8 bits (and a 6-bit embedding) gives the same logits as the same codes stored widened, in step, window, tree and prompt modes. - Qwen3.5-4B-OptiQ-4bit and Qwen3.5-9B-OptiQ-4bit: window and tree equal step over 300 ids; bench_concurrent --alone --serial at 1 and 4 streams all equal; about 130 and 78 tokens/s for one request. 8-bit modules have no prompt GEMM yet; prompt chunks run them on the decode tiles, 16 rows at a time.
The mixed-width loader widened 2 and 3-bit codes onto the 4-bit tiles and 5 and 6-bit codes onto the 8-bit ones, so a mostly 2-bit checkpoint took the 4-bit checkpoint's memory. affine8.cu's lane GEMV is now templated on the width: a lane's 16 code pairs a group sit in `bits` words, first codes in the low halves and second in the high, and bf16 (0x4300 | c) - 128 rebuilds each code exactly. The MMAs see the operands the 8-bit tiles give them, so a code's products do not depend on the width it is stored at. - AffineN (was Affine8) carries its width; AffineNKernels resolves a gemv and a split-K kernel for 2, 3, 5, 6 and 8 bits, and tf_affinen_pack packs them. - 4-bit modules keep the affine-4 tiles and prompt GEMM; uniform 4-bit checkpoints are unchanged. Embeddings still widen to 4 or 8 bits (rowBits). - TensorFold/Qwen3.8-27B-oQ2 takes 10.3 GiB and fits a 16 GB card. Checks on an RTX 4080 SUPER: - tf-cuda-test affine8: 24 cases; 2, 3, 5 and 6-bit products byte-equal to the codes widened onto the 8-bit tiles at 1, 3, 9 and 16 rows, within 0.32 of the float64 bound, split-K equal to the slice loop. - Qwen3.5-2B 4-bit: step, window, tree and prompt logits byte-equal to the build before. - The 2B requantized at 2, 3, 5, 6 and 8 bits: byte-equal to the same codes widened to 8 bits in all four modes, and step/window/tree byte-equal to the previous build, where 2 and 3-bit codes ran on the 4-bit tiles. - Qwen3.8-27B-oQ2: window, tree and prompt equal step over 300 ids; bench_concurrent --alone --serial at 1 and 2 streams all equal; about 45 to 64 tokens/s for one request, about 750 tokens/s prompt.
Prompt chunks ran modules of every width but 4 on the decode tiles, 16 rows at a time. affine8.cu now has qmm_prefill's GEMM over AffineN's tiles for 2, 3, 5, 6 and 8 bits: a 128 by 128 block on 2 by 4 warps, three stages, each weight rounded once to bf16 (s * q + b) and one fp32 chain over K. Its operands are the affine-4 prompt GEMM's, so a code of up to 4 bits gives the same bits as there. Checks on an RTX 4080 SUPER: - tf-cuda-test affine8: the prompt GEMM over 144 rows within 0.07 of a bound that allows for the bf16 weights; 2 and 3-bit byte-equal to the affine-4 prompt GEMM of the same codes, 5 and 6-bit to the 8-bit tiles'. - Qwen3.5-2B 4-bit: all four modes byte-equal to the build before. - The 2B requantized at 2, 3, 5, 6 and 8 bits: byte-equal to the same codes at 8 bits in all four modes, prompt included; decode modes unchanged. - Qwen3.8-27B-oQ2: window and tree equal step, step unchanged; bench 0 unequal. Prompts about 770 tokens/s (from 750): already bound by the bf16 tensor rate with fp32 accumulation. - Qwen3.5-9B MLX 8-bit: prompts about 1,940 tokens/s, from 1,520; bench at 1 and 4 streams 0 unequal.
…ffine_codes.cuh, one-line comments affine8.cu passed 600 lines with the prompt GEMM. The GEMM moves to its own kernel module and the code readers both use to a header; lean_check reports 528 problems, as on the base.
The weight loader's staging buffer is now Engine.init's and freed once the weights are on the card, so the server's stream budget counts that memory; the drafter's loader stages its own for its load.
--kv-cache fp8 stores attention keys and values as e4m3 rows with an fp32 scale each (the tf_kv_quant format): qwen35_kv8.cu holds rope, the attention partials and compaction for it, and the cache sizes, the startup line and the refusals follow. Windows and trees stay equal to steps bit for bit; the 9B 4-bit README perplexity is 9.267 against 9.289 with bf16 caches. On a 16 GiB card the 9B 8-bit serves a 262144-token window and the 27B oQ2 a 106496-token one.
1e149e9 to
0eebbd3
Compare
|
Thank you, @BobClawblaw, for making fp8 an explicit cache option: At 0eebbd3, the standalone write, read and compaction layouts agree in source. Combining the image and reuse changes exposes two gaps: Please make the combined vision/reuse path position- and format-aware, including step, window, tree and resume comparisons in its integration checks. Posted by Ash's AI Assistant Synthena |
…otate at the bf16 path's positions rope_rows in qwen35.cuh computes each row's position and rotates its queries and key once; tf_q35_rope and tf_q35_rope8 differ only in how they store the key and value (bf16 rows, or e4m3 rows and scales). A later change to the positions, such as an image prompt's M-RoPE planes and rope_delta, then reaches both caches. Logits are byte-equal to before in step, window, tree and prompt modes, with bf16 and fp8 caches, on the 2B and the 27B.
|
Follow-up at c9e78c2: the bf16 and fp8 rope kernels now share one body ( Checks, RTX 5090:
|
|
Second-platform check of the follow-ups: #886 at The merge. The conflicts are the ones the follow-up predicts.
Served (
Host: Windows gaps, already noted on #888 and outside this PR:
|
Stacked on #859 (then #849 and #836). Until those merge, this PR shows their commits too. The change here is the last commit,
0eebbd32, rebased onto #859's7cddf336(on #836's5e1c3fea). It was split out of #859 so that each PR makes one change.What this changes
tensorfold-native serve --kv-cache fp8stores the Qwen3.5 hybrids' attention keys and values as e4m3 rows with one fp32 scale each: amax / 448, thetf_kv_quantformat. The linear-attention state stays fp32 and the convolution tails stay bf16. A full-attention layer's cache takes 520 bytes a position and head, against bf16's 1,024. On a 16 GB card, the 9B 8-bit serves its full 262,144-token window and the 27B oQ2 serves 106,496 tokens.zig/kernels/cuda/qwen35_kv8.cuadds three kernels:tf_q35_rope8quantizes each new key and value row as it lands.tf_q35_attn_partial8runs the bf16 path's chunked attention over e4m3 rows, with chunks cut from each row's own keys, as Windows: the native CUDA server, and Qwen3.5-2B and Qwen3.8-27B on CUDA #836's8570e068does for bf16. It takes each key's dot product on its codes and then applies the key's scale, and each value's weight carries the value's scale.tf_q35k_compact8moves kept rows and their scales for a tree's accepted path.past_rows,rotate,rows_seen,seen_row) move intoqwen35.cuh.Geometry.kvBytes(cap, kv8)sizes the cache, and the startup line says "(fp8 caches)".engine_api.Open.kv8;--kv-cache bf16|fp8incli.zig, with bf16 the default;tf-qwen35-cudatakes an optional sixth argument,fp8.Which sums change: only with
--kv-cache fp8. Attention scores and values then read e4m3-rounded keys and values. The flag is off by default, and bf16 caches give the same bits as before: logits are byte-equal to #859's in step, window, tree and prompt modes on the 2B, a mixed-width 2B and the 27B.Receipt
Environment: this PR's head (
0eebbd32) on an RTX 4080 SUPER 16 GB (sm_89), Windows 11, CUDA 13.3, Zig 0.17.0,-Dsm=89. Checkpoints:lmstudio-community/Qwen3.5-9B-MLX-4bit@b455506blmstudio-community/Qwen3.5-9B-MLX-8bit@950e0459TensorFold/Qwen3.8-27B-oQ2@086dcb02mlx-community/Qwen3.5-2B-MLX-4bit@93760be4Exactness:
bench_concurrent --alone --serialwith--kv-cache fp8(levels 1 and 4, 4 cells): 48/48 equal alone and 10/10 equal to"draft": false, on the 9B 4-bit and on the 27B oQ2.tf-qwen35-cudawith fp8 caches and a 20,000-token second stream stepping beside (TF_Q35_BESIDE): step, window and tree logits equal the stream alone, on the 2B and the 27B. Also after a 3,000-token prefill with a 200-token stream beside.tf-qwen35-cudaover 300 ids with fp8 caches: window and tree equal step bit for bit, on the 2B and the 27B.z-lab/Qwen3.8-27B-DFlash2@50307d4c), on the 27B oQ2 served with--drafter --parallel 2 --context 8192 --kv-cache fp8:"draft": false.--context 262144 --parallel 2 --kv-cache fp8, given two different 120,011-token prompts together.token_shaequals its run alone, and the first prompt's drafted reply equals its"draft": falseone.Decode speed (
bench_openai, 9B 4-bit, one stream, medians over 5 reps):Prompt speed (
bench_cache --require-cold, cold 16,000-token prompt, 9B 4-bit): 9.30 s with fp8 caches, 8.89 s with bf16 (27B oQ2: 25.04 s against 23.66 s). The prompt writes and reads quantized rows, so it costs about 5% at this length.What fp8 costs:
README perplexity (9B 4-bit, 1,024 tokens): 9.267 with fp8, 9.289 with bf16.
Top-1 over 300 ids: agrees with bf16 caches on 291/300 (2B) and 295/300 (27B oQ2).
Long context (needles and facts,
--parallel 1, measured before the rebase; one stream's fp8 logits are byte-equal before and after it):Capacity on 16 GB (startup lines):
Tests:
zig build test: 48/48 steps; 747 of 761 tests pass, 14 skipped, 0 fail.--kv-cache fp8opens only a family whose options take it.Not run, and why:
--kv-cache fp8is refused.cuda_fp8.h's e4m3 types.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo, with fp8 caches as with bf16. This engine keeps no prompt cache on CUDA, so a resumed prompt is a fresh one.Limits
qwen35_kv8.fatbintoo.