Repository navigation
Qwen3.5-9B prompt rows in one product in Zig engine - #1897
Open
grantoverton wants to merge 11 commits into
Open
grantoverton wants to merge 11 commits into
grantoverton wants to merge 11 commits into
Conversation
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.
…product A mixed-width fused stack (an oQ4e checkpoint's 4/5-bit q|k|v or qkv|z) ran each lane member one row a launch, so a 2,048-row prompt pass read that member's weights 2,048 times. Prompt rows now go through one lane product into a stage buffer, and a copy kernel places them at the stack's stride; decoded rows keep the per-row launch. The lane product is the call quantOne already makes for a single matrix of that width, and quant_rows_test's 300-row pass through a 4/5-bit stack still equals its rows run one at a time. The GPU profile also reports its per-shape dense totals.
…h (opt-in TF_Q27_FIVE_BYTES) An oQ4e checkpoint's 5-bit g64 matrices (on the 9B: DeltaNet out_proj in all 24 linear layers, some down, z, a/b and q/k/o) run on tf_lane, a 1-16 row kernel: a 2,048-row prompt pass is 128 launches a matrix, each unpacking the 5-bit stream and summing its input groups again. With TF_Q27_FIVE_BYTES=1 on chips with tensor units, those matrices store their codes one byte a weight ([n][k], scales and biases as prepareLanes lays them) and run tf_byte: tf_lane's arithmetic on the same codes (the uint8 tensor op, the same scale and bias FMAs, its K split and reduction order, input sums in tf_lane's order from one sums pass), with every 16-row block in one dispatch. Unset, nothing changes. The byte copy holds 8 bits a weight instead of 5. Test: on a chip with tensor units, a 300-row prompt pass through byte codes equals the lane kernel's rows one at a time, and byte codes' decoded rows equal the lane kernel's prompt pass.
grantoverton
force-pushed
the
qwen35-9b-prompt-rows
branch
from
November 2, 2026 13:16
f4e7d6b to
37d8b13
Compare
Owner
|
Thanks @grantoverton. The bit-identity claim against
Two small points: say how the start-up fit counts the byte copy (+0.23 GB on the 9B, +0.57 GB on the 27B). And replaced scratch buffers are released mid-pass, which is safe because qwen27's command buffers retain bound resources, but please confirm that on #1826's base. |
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 #1896 (grow-as-used K/V, branch
qwen35-9b-kv-grow). On an oQ4e checkpoint (4-bit in groups of 64, with 5-bit matrices and mixed-width fused stacks) the 27B engine's prompt pass runs two kinds of projection one row a launch: a mixed stack's lane member, and every 5-bit g64 matrix ontf_lane, a 1-16 row kernel. An M5 prompt pass is up to 2,048 rows, so those matrices are read hundreds of times a pass. This puts every prompt row of a lane member through one product, and, withTF_Q27_FIVE_BYTES=1on chips with tensor units, runs 5-bit matrices from byte codes with every 16-row block in one dispatch. Outputs are bit-identical: both keeptf_lane's arithmetic, and the replies below equal before's on both chips.Please review the last two commits here (from "qwen27: a mixed stack's lane member takes all its prompt rows in one product" on); the rest is #1896 (grow-as-used K/V) and its stack.
What this changes
quant_gpu.zig,core/gpu_profile.zig,core/contract/qwen27_records.zig).quantRunslaunched each lane member once a row. Prompt rows (k.promptand more than one row) now go through onequantOnecall into a stage buffer, the call it already makes for a single matrix of that width, and a copy kernel (q27_stack) places them at the stack's stride. Decoded rows keep the per-row launch. The GPU profile also prints its per-shape dense totals.quant_rows_test's 300-row pass through a 4/5-bit stack, which equals its rows decoded one at a time, now runs this path on every chip. The copy kernel's launch is registered with the kernel contract (gate 5).projection.zig,quant_gpu.zig,gpu_weights.zig,quant_rows_test.zig,qwen27_records.zig). WithTF_Q27_FIVE_BYTES=1on a chip with tensor units, a 5-bit g64 matrix stores its codes one byte a weight ([n][k], scales and biases laid out asprepareLaneslays them) and runstf_byte:tf_lane's arithmetic on the same codes (the uint8 tensor op, the same scale and bias FMAs, its K split and reduction order, input sums intf_lane's order from onetf_byte_sumspass), every 16-row block in one dispatch. Decoded rows use it too. The byte copy holds 8 bits a weight instead of 5. Unset, or on a chip without tensor units, nothing changes. Test (tensor units): a 300-row prompt pass through byte codes equals the lane kernel's rows decoded one at a time, and byte codes' decoded rows equal the lane kernel's prompt pass, byte for byte. Both kernels' launches are registered for gate 5.Exactness
Target: a local oQ4e quantization of
ornith-ai/Ornith-1.5-9B(not published; the 9B shape #1826 serves): 4-bit in groups of 64, with 52 matrices at 5 bits (DeltaNetout_projin all 24 linear layers,in_proj_ain 10,in_proj_bin 9,in_proj_zin 3,down_projin 4 andk_projin 2), so both changed paths load. Drafterornith-ai/Ornith-1.5-9B-DFlash@21be344 at--drafter-bits 4. Before is #1896's tip; after is this tip, unset and withTF_Q27_FIVE_BYTES=1. Greedy, thinking off. The 4-stream runs (--parallel 4 --context 65536 --prompt-cache-gib 0) send #1896's sequence: the 12 drafted-versus-plain prompts,bench_openai.py,bench_concurrent.py --levels 1,4 --alone --serial, a 30,000-token exact-id prompt andbench_concurrent.py --levels 4 --alone; the digest is over every reply's prompt length, token count and token hash, sorted.6944e5041f256944e5041f25TF_Q27_FIVE_BYTES=16944e5041f25TF_Q27_FIVE_BYTES=16944e5041f256944e5041f259dab160716a49dab160716a49dab160716a49dab160716a4TF_Q27_FIVE_BYTES=16944e5041f25Every cold-prefill run below also gives equal replies on its chip (25 each, digest
6ffc6af68772)."draft": false, 12/12 in every run; concurrent equals solo, every alone and serial cell equal.quantOne's own call for that width, and the byte kernel's 16-row blocks each run the same group loop. The host tests above check this at 300 rows (two full blocks plus a partial one at every 128- and 16-row edge); the 2,048-row M5 chunks and 128-row M3 chunks of the 30,000- and 65,536-token prompts give before's replies.tf_lane's K split, intf_lane's order.The 27B. Both changed paths serve any qwen27 checkpoint with mixed stacks or 5-bit matrices, so the 27B was run before and after: a local oQ4e quantization of Qwen3.8-27B (not published; the layout #1822 describes: 4-bit in groups of 64, 5-bit
out_projon all 48 DeltaNet layers,in_proj_aandin_proj_bon 23,in_proj_zon one, a 5-bit embedding and a 6-bitlm_head),serve <27B> --parallel 1 --context 131072 --prompt-cache-gib 0, no drafter: the 12 prompts above,bench_openai.py, and cold prefill of the same prompts as text.bench_openaicode / chat832fc648f137832fc648f137TF_Q27_FIVE_BYTES=1832fc648f137832fc648f137(Prefill here is the text of the same paragraph, 2 rounds, tok/s; the 27B's tokenizer gives the same token counts.)
Performance
Zig 0.17.0,
zig build native -Dcpu=apple_m1 -Dversion=1.0.0on each Mac. One model process of mine at a time; both Macs were also serving other models, and the M5's other work moved single runs by up to 30% (each table names the spread), so runs alternate before and after and the tables give every run.Cold prefill: the exact token arrays from #1826 (upstream's public lighthouse paragraph, tokenized by Qwen3.5-9B's
tokenizer.json, sent as ids),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.M5 Ultra, 9B oQ4e (runs in order; the M5's other work moved whole runs by up to 30%: before 2 is 22% below before 1 at 2,048)
TF_Q27_FIVE_BYTES=1TF_Q27_FIVE_BYTES=1TF_Q27_FIVE_BYTES=1TF_Q27_FIVE_BYTES=1TF_Q27_FIVE_BYTES=1TF_Q27_FIVE_BYTES=1In every cycle after is above the before run beside it, and the byte path is above after. Against the before run next to it, after gains 15% to 21% (cycle 1), 3% to 16% (2), 1% to 11% (4), 18% to 20% (5) and 19% to 29% (6); the byte path adds 10% to 58% over the after run next to it. The M5's absolute level drifts by 30% between cycles, so compare within a cycle.
M3 Ultra, 9B oQ4e (runs in order; the byte path needs tensor units and does not run here, and an M3 prompt pass is at most 128 rows)
On the M3, after equals before within the runs' spread.
Decode (
bench_openai.py … --tokens 256 --reps 3 --temperatures 0, medians;bench_concurrent.pyaggregates), from the 4-stream runs:M5 Ultra (tok/s)
bench_openaicodeTF_Q27_FIVE_BYTES=1TF_Q27_FIVE_BYTES=1Decoded rows don't take the stack path, and the after runs decode as before does. Decoded rows do take byte codes when they are on.
M3 Ultra (tok/s)
bench_openaicodeMemory: with
TF_Q27_FIVE_BYTES=1the 9B loads at 18.45 GB instead of 18.33 GB (M5,/memoryafter load) and peaks at 28.35 GB against 27.91 to 28.11 GB unset (4-stream runs); the 27B peaks at 44.95 GB instead of 44.37 GB; unset, the footprint equals before's.Build and host tests
At the tip
f4e7d6b68: this series rebased onto #1896's current tip8e2463c7b, which adds a third commit (a fresh pass takes its prompt's K/V rows at once) that runs only withTF_Q27_KV_GROWset; every run above has it unset, so they ran this code. 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,231/9,289 tests passed and 58 skipped; goldens 55/55 and 322/335 equal, 0 unexpected differences;zig build dev-tools -Dcpu=apple_m1exit 0. On the M3: native and dev-tools exit 0; test + test-golden 590/594 steps, 9,226/9,289 passed, 57 skipped and 6 failed, goldens 55/55 and 322/335 with 0 unexpected. All six failures are in the server's benchmark-client tests (benchmark_cli_test,benchmark_stream_test,benchmark_upload_test,benchmark_run_test), each a local fake server counting one request more than its test sent; this series doesn't touch those files, they pass on the M5 at the same tip and on the M3 ataa20a833a, and the M3 was under other jobs' load at the time (#1896's tip shows the same tests failing, a different one or two each run).python3 tools/zig/lean_check.py: 424 problems at5c5667b95and at each commit (0 added). Before this rebase, ataa20a833a: on the M5, native exit 0, test + test-golden exit 0 (594/594 steps; 1,281/1,305 tests rerun, 24 skipped, the rest cached from the same tree's full run at2af064e9d: 9,229/9,287 passed, 58 skipped), goldens 55/55 and 322/335 with 0 unexpected, dev-tools exit 0; on the M3, native, test + test-golden (594/594 steps, goldens 55/55 and 322/335 with 0 unexpected; most tests cached from the same tree's full run) and dev-tools all exit 0. At the previous base (ce9c264ce) the M3 gave: native exit 0, 594/594 steps, 9,230/9,287 passed, 57 skipped (the byte-code test skips there), goldens equal, dev-tools exit 0. 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.Before the rebase, at the same qwen27 code:
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,079/9,133 tests passed and 54 skipped (the byte-code test skips without tensor units); goldens 55/55 and 322/335 equal, 0 unexpected differences;zig build dev-tools -Dcpu=apple_m1exit 0.zig build test test-golden -Dcpu=apple_m1 --summary allexit 0, 587/587 steps, 9,078/9,133 passed and 55 skipped (the byte-code test runs here); goldens 55/55 and 322/335 equal, 0 unexpected differences.python3 tools/zig/lean_check.py: 424 problems at the base and at each commit (0 added), before and after the rebase.zig build test): the three new launch records pass.zig build kernel-contract -Dcpu=apple_m1was not run (see Not run).Not run
serve <27B oQ4e> --parallel 1 --context 131072 --prompt-cache-gib 0, the 12 prompts,bench_openai.pyand the cold-prefill text, before and after, alternated.zig build kernel-contract -Dcpu=apple_m1on either Mac, so there is noSUMMARYline yet.tf_byte(a replacement fortf_lanethat must equal its bits) andq27_stack: only their gate-5 launch records are added.zig build kernel-contract -Dcpu=apple_m1with atf_bytetest type namingtf_laneas the kernel it replaces.TF_Q27_FIVE_BYTES=1as the default on tensor units: it costs 3 bits a weight on the 5-bit matrices and was only measured on the 9B.tf-qwen27-widths <checkpoint>lists the mixed stacks and 5-bit matrices any checkpoint has.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 this series adds over Qwen3.5-9B DFlash drafter in Zig engine #1851's: the one-lineall_metalchange inregistry.zigagainst EXL3 Flash Next on CUDA: one rank, offloaded experts, sm_120 gate (1/8) #1441 to EXL3 Flash Next: draft lane guard for offloaded experts (8/8) #1446, KV pool: host tier for the core pool (5/8) #1501, EXL3 Flash Next: host-resident K/V on the tier (6/8) #1522, flashnext cuda: a shared round's attention in one launch a kernel #1739, Implement GLP control vectors with a fused Metal kernel for Qwen3.8-27B #1835, Ling: grouped expert routing on Metal, checked against a float64 reference #1836, Metal: the 27B verifies 32 rows a pass on pre-M5 Macs #1841, CUDA: Qwen4 and tensor-parallel loaders drop the checkpoint's pages after every rank is ready #1887, tools: bench_image_session plays one image chat session and records cached tokens and first-token time #1889, Kolibri-1 on CUDA: prompts in the shared KV store with TF_KV_SHARE (off by default) #1891 and DeepSeek-V4.1: --prompt-cache-gib bounds the whole prompt cache, the same on every rank #1892 (each adds records there);qwen27_served.zigagainst Qwen3.5-9B decisions in Zig engine #1871 (from Qwen3.5-9B grow-as-used K/V in Zig engine #1896); andprojection.zigandqwen27_lanes.zigagainst 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. I'll rebase onto whichever lands first.