Repository navigation
Qwen3.5-9B grow-as-used K/V in Zig engine - #1896
Open
grantoverton wants to merge 9 commits into
Open
grantoverton wants to merge 9 commits into
grantoverton wants to merge 9 commits into
Conversation
This was referenced Nov 2, 2026
A DFlashDraftModel config now parses as variant v1: its taps (up to eight), block, mask token and layer kinds come from the config, with no conv or selector keys. Its graph is the plain qwen3 body (eleven tensors a layer, fc, hidden_norm, norm), each layer adding attention and MLP to its residual through a new add operator on the target's residual kernel. With no codebooks (rank 0) the selector's edge term is zero, so trees come from each depth's marginals. The target's tap planes follow the drafter's tap count (eight here), and the keep, gather and glue paths pass that count instead of five. DFlash2 parsing, its schedule and its tests are unchanged.
…nd fusion rows runnerBytes kept five raw planes a prompt row. tapRowBytes now prices a drafter's planes, and for a fusing (v1) drafter the row-major copy, fc's rows and the fused row: 18 x hidden x 2 bytes for the 9B drafter's eight taps. Fit carries it as tap_row; 0 keeps the five-plane figure, so DFlash2 and drafter-less serves plan exactly as before.
…the whole context The target's taps pass now runs the drafter's fc and hidden_norm as a v1 drafter's taps complete (Taps.setFuser), so every gather reads one plane of fused rows: a ring row and a prompt-cache state's drafter part are 8 KiB a token instead of 64 KiB. The drafter's context append starts from those rows (k and v alone). DFlash2 keeps its raw planes. The ring and the context now keep the served span (context plus the lane window) instead of the 4096-token window. Sliding layers attend their own window of it; the full-attention layer attends all of it, with its own cache and attention sized to the span.
…tap slots Commit "read and run the 9B's DFlash drafter beside DFlash2" made the tap ids and a draft config's taps eight-slot arrays, and three tf-qwen27 tools still set five: zig build dev-tools failed on dflash_graph.zig, dflash_runtime.zig and forward_diff.zig. They now name the same five taps with three empty slots (the plane count stays five).
The serving line already named a v1 drafter "DFlash", but the startup log line still said "DFlash2 drafts for one of them" for either. Config.kind() gives the name, DFlash for z-lab's v1 layout and DFlash2 for DFlash2, and both lines take it from there. Test: a v1 config is named DFlash and the stock DFlash2 config DFlash2.
Every slot took K/V rows for the whole --context on its first use, so four live 9B streams at --context 32768 held 4 GiB of K/V while each stream used a few thousand tokens. With TF_Q27_KV_GROW=<rows> a slot starts at that many rows and, when a round, a restored state or a moved slot needs more, grows to the need rounded up to the step or a quarter more than it holds, at most the context. Growth goes one attention layer at a time: a wider buffer pair, one waited blit of every old byte to its place (each K/V head from h*old*dim to h*new*dim), the old pair out of the residency set and freed, so growing holds at most one layer's old pair beside the new rows. It runs where moveSlot and reset already touch a slot (settleSlot, no pending keep, no failure). The attention kernels already take each stream's head stride from its cache, so the plan now passes the slot's held rows; arithmetic and key order depend only on positions. Unset, every slot holds the context as before, and the startup fit and slot admission still charge the full context.
With TF_Q27_KV_GROW a slot kept its high-water rows after its stream ended, and every parked stream's drafter taps ring was sized to the whole span, so one long request left its rows behind in every slot it touched. Give-back: when a stream leaves slot k, the slot goes back to its first rows (fresh small buffers, no copy) as soon as no blit reads it. The check is a non-blocking status read of the kept-state save, retried at each backend call, so no stream waits on it. Slot 0 keeps its rows while it stands for a next turn's resume; the next fresh pass there gives them back, as does any fresh pass in any slot. Rings: a parked stream's taps ring starts at the same first rows and grows by the K/V rule when taps are absorbed or a kept state is restored or moved into it (Ring.grow, Ring.hold); the old rows move to their places on the CPU after the slot's pending reads land. Drafter frames wait for their ring reads before closing, so a ring only changes between them. A context's reset gives its ring back. The growth helpers move to kv_slots.zig to keep decode_round.zig under 600 lines.
…ROW) With K/V grown as used, a fresh prompt pass gave its slot back to the first rows and then grew it step by step through the prompt: ten growths for a 65,536-token prompt, each a waited blit of every layer into new buffers. The slot is empty at that point, so it now takes the rows the prompt fills (rounded up to the step) in new buffers with no copy, and keeps them when it already holds exactly that many. Decode and restores still grow by the same rule.
grantoverton
force-pushed
the
qwen35-9b-kv-grow
branch
from
November 2, 2026 13:16
8e2463c to
f6ea8ee
Compare
Owner
|
Thanks @grantoverton. Unset is main exactly, and set, the per-head growth is right. One change, and it matters: a growth that fails partway leaves a slot's attention layers at different row counts.
Small: a growth briefly holds one layer's old pair beside its new rows, about 1/L above the full-context charge. Note it, or charge it in the fit. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacks on #1851. On the 27B engine every slot takes K/V rows for the whole
--contextthe first time it is used, so a 9B serve at--parallel 4 --context 65536holds 8 GiB of K/V for streams that mostly use a few thousand tokens. WithTF_Q27_KV_GROW=<rows>a slot starts at that many rows, grows as its stream does, and gives the rows back when its stream ends; drafter tap rings do the same. Unset, nothing changes. Replies are byte-identical with it on or off.This branch carries #1851 (and through it #1826 and #1125) until those merge, so please review the last three commits here (from "qwen27: grow a slot's K/V as its stream grows" on).
What this changes
kv_grow.zig,decode_round.zig,snapshot.zig,qwen27_served.zig). WithTF_Q27_KV_GROW=<rows>(a multiple of 16; 0 or unset is off) a slot starts at that many rows. When a round, a restored state or a moved slot needs more, it grows to the need rounded up to the step, or a quarter more than it holds, whichever is larger, at most the context. Growth goes one attention layer at a time: a wider buffer pair, one waited blit of every old byte to its place (each K/V head moves fromh * old * dimtoh * new * dim), then the old pair leaves the residency set and is freed, so growing never holds more than one layer's old pair beside the new rows. It runs only wheremoveSlotandresetalready touch a slot (aftersettleSlot, with no pending keep and no failure). The attention kernels already take each stream's head stride from its cache, so the plan now passes the slot's held rows instead of the context; arithmetic and key order depend only on positions. The startup fit and slot admission still charge the full context, so this lowers memory held, not the admission limit. Tests that fail before: the step parser, the growth rule (need rounded to the step, a quarter more, capped at the context), and the per-head moves into a wider buffer.kv_slots.zig,tap_ring.zig,dflash/runtime_model.zig,qwen27_slots.zig,qwen27_lanes.zig). When a stream leaves slot k, the slot goes back to its first rows (fresh small buffers, no copy) as soon as no blit reads it. The check is a non-blocking status read of the kept-state save, retried at each backend call, so no stream waits on it. Slot 0 keeps its rows while it stands for a next turn's resume; the next fresh pass there sizes them to its prompt (3). A parked stream's taps ring starts at the same first rows and grows by the same rule when taps are absorbed or a kept state is restored or moved into it (Ring.grow,Ring.hold); the old rows move to their places on the CPU after the slot's pending reads land. Drafter frames wait for their ring reads before closing, so a ring only changes between them. A context's reset gives its ring back. The growth helpers move tokv_slots.zigto keepdecode_round.zigunder 600 lines.kv_slots.zig,kv_grow.zig,decode_round.zig,qwen27_slots.zig). A fresh prompt pass used to give its slot back to the first rows and then grow it step by step through the prompt: ten growths for a 65,536-token prompt, each a waited blit of every layer into new buffers. The slot is empty at that point (it was just reset), so it now takes the rows the prompt fills, rounded up to the step and at least one step, in new buffers with no copy, and keeps its buffers when it already holds exactly that many. Decode and restores still grow by the rule in (1). Test that fails before: the admission rows (the prompt rounded up to the step, at least one step, at most the context; one growth for 65,536 tokens where (1) takes ten).Exactness
Target
mlx-community/Qwen3.5-9B-MLX-4bit@938d891 and drafterornith-ai/Ornith-1.5-9B-DFlash@21be344 at--drafter-bits 4, as in #1851. Before is #1851's tip as it stood,bb1028e7b; after is this branch on it,3c28eb83c, run withTF_Q27_KV_GROW=4096and unset. Greedy, thinking off,serve <target> --drafter <drafter> --parallel 4 --context 65536 --prompt-cache-gib 0. Each run sends, in order: the 12 drafted-versus-"draft": falseprompts from #1851 (6,027- and 12,026-token prompts among them),bench_openai.py,bench_concurrent.py --levels 1,4 --alone --serial, a 30,000-token exact-id prompt (the cold-prefill array cut at 30,000, 32 tokens; it grows slot 0 six times, from 4,096 to 32,000 rows; with (3) it takes 32,768 rows at once), thenbench_concurrent.py --levels 4 --aloneagain. Every finished reply's prompt length, token count and token hash is in the server log; the digest below is over all of them, sorted.30ac466ff7d830ac466ff7d830ac466ff7d8d9608ce72d81d9608ce72d81d9608ce72d8130ac466ff7d830ac466ff7d8d9608ce72d81d9608ce72d81(The M3's first before run lost its 30,000-token request to a bad client field; its 80 other replies equal the second run's, hash for hash.) The "this tip" rows ran with (3) added (the M5's build predates only a shortened doc comment). The cold-prefill runs below also give equal replies before and after on each chip, at every size.
"draft": false, 12/12 in every run, with partial acceptance (1.9 to 11.6 tokens a round).TF_Q27_KV_GROW=4096, prompt cache on): Qwen3.5-9B DFlash drafter in Zig engine #1851's probe, the 6,027-token prompt twice; the second resumes at 6,020 cached tokens with the token hash Qwen3.5-9B DFlash drafter in Zig engine #1851 records,c40a2accd7f8, on both chips, at this tip and at the previous one (1ba21dc0e). At the previous tip the restore went into a slot given back to 4,096 rows, so it grew first (restoreSlotandmoveSlotcallreservebefore the copy); at this tip the slot kept the 8,192 rows its first pass took. The replies with cached tokens in the table come from slot 0's own resume, and equal before's.TF_Q27_KV_GROW=4096,--prompt-cache-gib 0): Qwen3.5-9B DFlash drafter in Zig engine #1851's 14-step probe passes on both chips at both tips, with Qwen3.5-9B DFlash drafter in Zig engine #1851's reference hashes (a9af571d979c/a4dc1a079732on the M3,a9af571d979c/05cf379ac28eon the M5). Its drops all land before any slot outgrows 4,096 rows, so a second 14-step probe drops streams after growth: a 22k-token prompt dropped 8 s in and again at its third reply line, a 4,000-token prompt dropped 100 lines into a reply whose decode grows its slot past 4,096 rows, and the same beside a running probe (whose hash is checked too); after each, both probes equal their references. Pass on both chips at both tips. At this tip the dropped 22k-token prompt had sized its slot to 24,576 rows on both chips and each dropped reply had grown its slot to 8,192 rows (slot 1's K/V and taps ring too, on the M3). At the previous tip the M5's dropped prompt had grown its slot five times to 25,600 rows and slot 1's K/V and taps ring grew there, while the busier M3 dropped the prompt before it grew; the server logs show each growth and the resize or give-back after it.Memory
Physical footprint from
/memory(physical_footprint_bytes, andlifetime_peak_physical_footprint_bytesfor the peak), and the peak from/usr/bin/time -l, same runs as above.With it unset the after binary's footprints equal before's (M3 25.99 GB peak, M5 27.70 GB). With (3) the 30,000-token prompt takes 32,768 rows at once instead of growing to 32,000, and the footprints stay as above: M5 loaded 16.27 GB, held after 4 streams 18.10 GB, after the 30,000-token prompt 19.04 GB, peak 19.07 GB; M3 13.83, 16.01, 16.96 and 16.97 GB (unset 25.80 GB peak). The cold-prefill serve (
--parallel 1 --context 131072) peaks at 25.93 to 25.97 GB before, 24.55 GB after without (3) and 23.80 to 23.81 GB with it on the M5, and 23.96 to 23.97 GB before and 21.83 GB with (3) on the M3.Performance
Zig 0.17.0,
zig build native -Dcpu=apple_m1 -Dversion=1.0.0on each Mac (tensorfold-nativesha256 before4de52daaca0b…M3,4b736b20c17a…M5; after3558a86efc0c…M3,568826dc7c07…M5; this tip's runs4a4c3185e962…M3,2b5fe567ff4b…M5). One model process of mine at a time; both Macs were also serving other models, so single runs move by a few percent either way.Decode, same runs as Exactness (
bench_openai.py … --tokens 256 --reps 3 --temperatures 0, medians, tok/s;bench_concurrent.pyaggregate tok/s at 1 and 4 streams):M5 Ultra (runs in order; * marks a cell 10% or more below the same cell in every other run, where the Mac's other work took the GPU: grow 2's chat reps were 289.13 / 182.22 / 180.08)
bench_openaicodeM3 Ultra (runs in order; the Mac was building and testing other work during all five, so these are noisy; * as above)
bench_openaicodeA later pair on the quieter M3:
bench_openaicodeDecode is unchanged within the runs' spread on both chips. (3) doesn't touch decode; this tip's 4-stream runs came while both Macs were busier, both arms of the same binary alike (
bench_openaicode 470.52 with grow and 402.99 unset on the M5, 92.48 and 92.71 on the M3), so they aren't a decode comparison.Cold prefill: exact token arrays from #1826 (upstream's public lighthouse paragraph, tokenized once by the checkpoint's
tokenizer.json),POST /v1/completionswith the ids,max_tokens1, temperature 0,ignore_eos, server--drafter <drafter> --parallel 1 --context 131072 --prompt-cache-gib 0, one request at a time with a short unrelated request between measurements; every record hascached_tokens0 andprompt_tokensequal to N. One 2,048-token warm-up, then 3 rounds, medians, tok/s (s). Runs alternate before and after. Without (3) a 65,536-token prompt grows slot 0 ten times, from 4,096 to 78,125 rows; with it the prompt takes 65,536 rows at once.M5 Ultra (runs in order)
M3 Ultra (runs in order)
Replies are equal in every prefill run on each chip.
Without (3), cold prefill with
TF_Q27_KV_GROWset was slower: in the tables above the M3's two grow runs are 9% to 17% below both before runs at every size, and the M5's four pairs move by -13%, -2%, +5% and -16%. Each measured prompt gave its slot back and grew it again step by step, ten waited growths of all 8 attention layers for 65,536 tokens, each into new buffers. With (3) a fresh prompt takes its rows once with no copy (one resize, then one back to 4,096 rows on the next fresh pass). The same cold-prefill runs at this tip, alternated with before:M5 Ultra (runs in order; in each run the Mac's other work slowed rounds 2 and 3 alike, so medians sit below the tables above)
M3 Ultra (runs in order; other jobs held most of the M3's GPU throughout, their leased share moving from 95-120% during the grow runs to 70% during before 4, so every run is at a third to a half of the M3 tables above)
These M3 runs can't separate a few percent, so the M3 has no answer yet (see Not run). On the M5 both grow runs are level with or above both before runs at every size (-0.1% to +5.2%), so the cost is gone there. A shorter alternated check on the M5 (2 rounds of 2,048, 8,192 and 32,768 tokens; the previous tip unset and with grow, then this tip with grow; its first three servers, before the Mac's other work slowed every arm alike) showed the same: without (3) the 8,192-token prompt ran at 3,896 and 4,078 tok/s against 4,590 unset; with (3), 4,576 and 4,579. The option is off by default, so a serve that doesn't set it pays nothing.
Build and host tests
At this tip
8e2463c7b, on #1851's current tip5c5667b95: on the M5,zig build native -Dcpu=apple_m1 -Dversion=1.0.0exit 0;zig build test test-golden -Dcpu=apple_m1 --summary allexit 0, 594/594 steps, 9,227/9,285 tests passed and 58 skipped; goldens 55/55 and 322/335 equal with 0 unexpected differences;zig build dev-tools -Dcpu=apple_m1exit 0;zig build test-qwen27-memory-hostexit 0, 12/12 (its new admission test fails to build at1ba21dc0e, which has no admission rule). On the M3: native and dev-tools exit 0; test + test-golden 592/594 steps, 9,228/9,285 passed, 55 skipped and 2 failed, goldens 55/55 and 322/335 with 0 unexpected; a rerun failed one. Each failure is in the server's benchmark-client tests (benchmark_stream_test,benchmark_cli_test,benchmark_run_test), a local fake server counting one request more than its test sent, a different test each run; this branch doesn't touch those files, they pass on the M5 at this tip and on the M3 at1ba21dc0e, and the M3 was under other jobs' load at the time. Main shares them: main3ce4a0480's ownzig build test test-goldenon the same M3 failedbenchmark_run_test(9,361/9,417 passed), with no qwen27 change in it.python3 tools/zig/lean_check.py: 424 problems at5c5667b95and at each commit (0 added). The runs above at3c28eb83cran on #1851 as it stood atbb1028e7b; since then #1851 changed only its startup line (qwen27_served.zig,dflash/config.zig: the line names a v1 drafter "DFlash"), so they ran the same qwen27 code apart from that log text. At1ba21dc0e(the first two commits on5c5667b95), the M5 gave 594/594 steps, 9,227/9,285 passed and 58 skipped, goldens equal, native and dev-tools exit 0, and the M3 native, test + test-golden (594/594 steps, goldens equal; most tests cached from the same tree's full run) and dev-tools all exit 0. At the previous base (ce9c264ce, same qwen27 code but the startup line) the M3 gave the same: native, 594/594 steps, 9,228/9,283 passed, 55 skipped, goldens equal, dev-tools exit 0.Before the rebase, on the M3 at
3c28eb83c:zig build native -Dcpu=apple_m1 -Dversion=1.0.0exit 0;zig build test test-golden -Dcpu=apple_m1 --summary allexit 0, 587/587 steps, 9,077/9,129 tests passed and 52 skipped; goldens 55/55 and 322/335 equal with 0 unexpected differences.zig build dev-tools -Dcpu=apple_m1exit 0.python3 tools/zig/lean_check.py: 424 problems at the base and at each commit (0 added). The M5 built the same tip (zig build native, exit 0). No kernel is added or changed, so there is no kernel contractSUMMARYline.Not run
serve <target> --drafter <drafter> --parallel 1 --context 131072 --prompt-cache-gib 0with and withoutTF_Q27_KV_GROW=4096, alternated, the prefill arrays above. The M3 pairs here ran under other jobs' GPU load.TF_Q27_KV_GROWset:serve <Qwen3.8-27B checkpoint> --drafter <DFlash2 dir> --parallel 4 --context 65536 --prompt-cache-gib 0with and without it, then the two client commands. Unset, the 27B runs the same code as before (unit tests only ran here).atc overlap --json(ATC isn't installed here). By hand,git merge-treeof this branch against every open pull request that touches its files, counting only conflicts its three commits add over Qwen3.5-9B DFlash drafter in Zig engine #1851's: Qwen3.5-9B decisions in Zig engine #1871 inqwen27_served.zig(a sibling on Qwen3.5-9B DFlash drafter in Zig engine #1851), and qwen27: serve the 27B pair, refuse what a pair cannot keep, and die with the peer #1434 to qwen27: document two-Mac 27B serving and the row split's exactness #1437 inqwen27_lanes.zig. I'll rebase onto whichever lands first.