Repository navigation
Qwen3.5-9B decisions in Zig engine - #1871
Open
grantoverton wants to merge 14 commits into
Open
grantoverton wants to merge 14 commits into
grantoverton wants to merge 14 commits into
Conversation
grantoverton
force-pushed
the
qwen35-9b-decisions
branch
from
November 1, 2026 22:53
11f9220 to
a410915
Compare
This was referenced Nov 1, 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.
Known values (causal or block_bidirectional full attention, causal linear attention, the slot_lm readout, 1 to 256 slots) are read; anything else is refused by name, and overrides that name no readout are refused rather than served causally. A checkpoint with neither file is a plain chat model.
The prompt kernels gain a BLOCK instantiation (q27_attn_prompt_block_8, q27_attn_prompt_sg_block_8) whose rows read keys [0, end): a block attends both ways within itself and to the cache before it (end may run past the pass for layer-major chunks). The causal kernels keep their names and compile from their old expressions; encode() passes end 0 and takes them. A bidirectional pass that cannot take the prompt kernel is refused rather than run causally. tf-qwen27-attention-block checks GQA 4 against a float64 reference: 8 cases, each mask far from the other's reference, and encode() equal to encodeTo(end 0) byte for byte.
A checkpoint whose readout is slot_lm serves decisions only:
- block_forward runs a block-bidirectional block in one prompt pass, or
layer-major in prompt-row chunks when it is longer (every chunk's keys
before any row reads them; GDN chunks stay causal and in order);
- the decider runs the trunk once, saves its DeltaNet state, and runs each
question's suffix from it; logits are the LM head rows of the answer
letters (dequantized as MLX does) dotted with the final norm's row;
- decisions go through the lane host's score queue (Engine.decide, set only
for a slot_lm checkpoint), so concurrent decisions take turns on the
engine thread;
- the server renders the qev-labels-v2 prompt through the chat template
(thinking off), splits at the shared prefix, and answers /v1/decisions
and /v1/systemone in oMLX's shapes; other models keep the 0.6.6 and
llama-server routes, so a chat checkpoint answers neither.
Answers round to 6 places from the exact binary value, as Python's round
does (Zig's {d:.6} rounds the shortest decimal: 0.1234565 gave 0.123457,
Python gives 0.123456); checked against Python on 200,000 values.
…one plan encodeBlocks takes a plan of chain streams and an end for each: stream s's rows read every key before ends[s]. The K/V writes and the query packing run once for the plan; then the one-stream prompt kernel runs once per stream over that stream's packed tiles, cache pointers and strides and rows of the output. tf-qwen27-attention-block gains 4 multi-stream sets (2 and 3 streams, with blocks across 512-key chunks, a 1-row stream and a 3-row stream): every stream's bytes equal its one-stream encodeTo pass. It ran red first (Unimplemented).
A family's Decider may set run_many: the engine thread hands it every decision queued when it drains the score queue (two or more), with an error slot per decision (Refused keeps its words, the rest fail and log). Decisions it handled are skipped by the one-at-a-time loop, and done is still set only in the final pass under the lock, so an asker's stack Ask is never touched after it returns. lane_host_score_test: three queued decisions (one refused) reach run_many in one call, each with its own answer or refusal. It ran red first (0 batches).
…irectional passes block_forward.runMany runs several slots' blocks in one pass, each stream reading the keys before its own end (encodeBlocks). The decider uses it two ways: - one decision: its trunk on slot 0, copied to slot i for question i (moveSlot), then every question in one pass instead of one pass each, when the questions fit one prompt pass's rows (128 before M5); - run_many: queued decisions in groups that fill the slots, their trunks in shared passes, then all their questions in shared passes of at most the chunk's rows. Decisions too large to batch run alone. A decisions-only server keeps 8 streams (KV made on first use). TF_DECIDE_BATCH=0 turns batching off. Each question's bytes equal its one-question pass's.
Engine.decide waited with no cancel or timeout. A Decision now carries timeout_ms, a client check (gone) and a stop flag (core/decision.zig, which also keeps engine_api.zig within 600 lines). While its decision is still queued, the asker removes it from the score queue once the time is up or the client has left, so the engine never sees it; once the engine has taken it, the asker sets stop and the qwen27 decider returns Cancelled at its next pass boundary (before the trunk, after it, and between questions; a batch skips a stopped decision), leaving its slot empty. Either way decide returns Failed with stopped saying why. The server gives each decision 60 s (503 when the time runs out) and drops one whose client closed its connection (peerGone), with nothing to send. lane_host_score_test: a queued decision past its timeout, and one whose client is gone, leave the queue unanswered; a running one past its timeout stops at its next pass and the next decision is answered. Before the change the test binary waited until it was killed.
…ass, bit for bit With TF_DECIDE_CHECK set, the decider recomputes every question from an empty slot (the trunk, then that question alone: no saved DeltaNet state, no moved slot, no shared pass) after answering, and the decision fails with CheckFailed unless every label logit is bit-equal. It logs one line a decision with the count.
grantoverton
force-pushed
the
qwen35-9b-decisions
branch
from
November 2, 2026 13:16
a410915 to
8f54eb7
Compare
Owner
|
Thanks @grantoverton. Chat checkpoints are unchanged: the causal prompt kernels keep main's expressions and names, the block forms are separate instantiations, and routes reach main's handlers without a decider. The mask is the model's (each block row reads keys [0, end), both ways within the block, with GDN causal and a bad end refused), and the readout dots each question's last final-norm row in f64 with the dequantized LM-head rows of the labels. Three small changes:
Small: the block kernels have no contract record, and |
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.
Part of #503; stacks on #1851. This serves
baselabs/elephant-1, a Qwen3.5-9B decision model, on the 27B engine: its full-attention layers read each block both ways (modeling_overrides: block_bidirectional full attention, causal linear attention) and it answers through a slot_lm readout (readout.json: the LM head rows of the answer letters at each question's last row). With this branch such a checkpoint answers/v1/decisionsand/v1/systemonein oMLX's shapes, and only those. A chat checkpoint (the 9B and the 27B among them) serves as before: its causal prompt kernels are main's, compiled from the same expressions, and both routes keep main's handlers.This is a fresh take on #1189, #1191 and #1192 that follows your 26 Oct review: decisions use main's score queue (
lane_host_score.zig), the qev shapes exist only for a slot_lm checkpoint, the wait has a time limit and notices a client that leaves, and every answer below was checked bit for bit against a fresh pass. The JointHead of #1191 is gone; this model's weights are public.This branch carries #1851 (and through it #1826 and #1125) until those merge, so please review the last eight commits here (from "qwen27: read a decision checkpoint's modeling_overrides and readout.json" on).
What this changes
modeling_overridesandreadout.json(families/qwen27/readout.zig). Known values (causal or block_bidirectional full attention, causal linear attention, the slot_lm readout, 1 to 256 slots) are read; anything else is refused by name, and overrides that name no readout are refused rather than served causally. A checkpoint with neither file is a chat model. Tests that fail before: the readout cases (known values, each unknown value refused, overrides without a readout refused).attention_prompt.metal,attention_prompt_sg.metal,attention.zig). The prompt kernels gain aBLOCKinstantiation (q27_attn_prompt_block_8,q27_attn_prompt_sg_block_8) whose rows read keys [0, end): a block attends both ways within itself and to the cache before it. The causal kernels keep their names and compile from their old expressions, andencode()(end 0) takes them. A bidirectional pass that can't take the prompt kernel is refused rather than run causally.tf-qwen27-attention-blockchecks GQA 4 against a float64 reference.block_forward.zig,native/qwen27_decide.zig,qwen27_served.zig,server/qev*.zig,lane_host_score.zig). A block runs in one prompt pass, or layer-major in prompt-row chunks when it is longer (every chunk's keys before any row reads them; GDN chunks stay causal and in order). The decider runs the trunk once, saves its DeltaNet state and runs each question from it; a question's logits are the answer letters' LM head rows (dequantized as MLX does) dotted with the final norm's row. Decisions go through main's score queue as a newEngine.decideentry, set only for a slot_lm checkpoint, so concurrent decisions take turns on the engine thread. The server renders the qev-labels-v2 prompt through the chat template (thinking off), splits at the shared prefix, and rounds answers to 6 places from the exact binary value, as Python'srounddoes. Tests that fail before: the decide queue test (answers on the engine thread, refusals keep their words, a stopped host fails) andqev_test(a decisions body, a System One body, refusals, rounding).attention.zig).encodeBlockstakes a plan of streams and an end for each; the K/V writes and query packing run once, then the one-stream block kernel runs once per stream.tf-qwen27-attention-blockgains 4 multi-stream sets. It ran red first (Unimplemented).run_manytogether (lane_host_score.zig). The engine thread hands a family'srun_manyevery decision queued when it drains the queue, with an error slot each;doneis still set only in the final pass under the lock. Test that fails before: three queued decisions (one refused) reachrun_manyin one call with their own answers.block_forward.runMany,qwen27_decide.zig). One decision: its trunk on slot 0, copied to slot i for question i, then every question in one pass when they fit one prompt pass's rows (128 before M5).run_many: queued decisions in groups that fill the slots, trunks then questions in shared passes of at most a pass's rows.TF_DECIDE_BATCH=0turns batching off.core/decision.zig,lane_host_score.zig,qwen27_decide.zig,qev.zig). ADecisioncarriestimeout_ms, a client check and a stop flag (in their own file, which also keepsengine_api.zigwithin 600 lines). While a decision is still queued, its asker removes it from the queue once the time is up or the client has gone, so the engine never sees it; once the engine has taken it, the asker setsstopand waits, and the decider returnsCancelledat its next pass boundary (before the trunk, after it, between questions; a batch skips a stopped decision), leaving its slot empty.decidethen returnsFailedwithstoppedsaying why. The server gives a decision 60 s (503 after that) and drops one whose client closed its connection (peerGone, as embeddings does), sending nothing. Tests that fail before: a queued decision past its timeout, and one whose client is gone, leave the queue unanswered; a running one past its timeout stops at its next pass, and the next decision is answered. Before the change the test binary waited until it was killed (296 s, undertimeout 300).TF_DECIDE_CHECKcompares each answer with a fresh pass, bit for bit (qwen27_decide.zig). With it set, the decider recomputes every question from an empty slot (the trunk, then that question alone: no saved state, no moved slot, no shared pass) after answering, and the decision fails withCheckFailedunless every label logit is bit-equal; one log line a decision gives the count.Exactness
Checkpoint: a mixed 4/5-bit requant of
baselabs/elephant-1@b3f8a93 (MIT): affine 4-bit, group 64, with 53 modules at 5 bits; not published (model-00001-of-00002.safetensorssha256313c27bd7ce3c7b4…,model-00002-of-00002.safetensors7eaa8d15118b7114…). Serverserve <checkpoint> --context 8192 --prompt-cache-gib 0 --no-drafts. The reference answers are oMLX's, recorded from a running system on the same checkpoint: 116 decision requests (216 questions: 115 decisions bodies and one System One body) and 966 more distinct bodies (1,833 questions). They're recorded traffic, so I can't attach them; the per-question records (probabilities, both engines' top two, margins) are available on request. Every check ran at11f922019(tensorfold-nativesha2562f7e7036f27be1f6…built on the M3,55375deada067762…on the M5). That was before the branch moved onto #1851's tip5c5667b95on main 9e8c368. The move adds #1826's pre-load--draftercheck ahead of this branch'sreadout.jsonread, and #1851's renamed startup line. Main's lane changes are opt-in, with defaults unchanged. At the tipa4109159cthe host suite ran again (below).TF_DECIDE_CHECK=1): all 966 bodies answered one at a time and then 4 at once, each decision recomputed question by question from an empty slot. M5: 1,932 decisions checked, every label logit bit-equal, 0CheckFailed. M3: the same 1,932, every label logit bit-equal, 0CheckFailed.mlx-community/Qwen3.5-9B-MLX-4bit@938d891 with the DFlash drafter of Qwen3.5-9B DFlash drafter in Zig engine #1851 on this tip, M3: drafted equals"draft": falseon Qwen3.5-9B DFlash drafter in Zig engine #1851's 12 prompts, the 14-step cancellation passes with Qwen3.5-9B DFlash drafter in Zig engine #1851's hashes (a9af571d979c/a4dc1a079732), and the drafter's tokens a round equal Qwen3.5-9B DFlash drafter in Zig engine #1851's on 11 of the 12 prompts in each of four runs (on the 6,027-token prompt they vary from run to run on both: 10.67 and 11.13 on Qwen3.5-9B DFlash drafter in Zig engine #1851, 10.67 to 11.64 here, with the same tokens)./v1/decisionsand/v1/systemonethere give main's answers (400 for a body without questions, 404 for System One), byte for byte as Qwen3.5-9B DFlash drafter in Zig engine #1851's.tf-qwen27-attention-block, GQA 4): the 8 one-stream cases (1 to 128 rows, block ends past 512-key chunk edges and layer-major ends) are within 0.00407 of float64 on the M5 (tensor-unit kernel) and 0.00390 on the M3 (simdgroup kernel), each mask far from the other's reference, andencode()equalsencodeTo(end 0)byte for byte; in the 4 multi-stream sets every stream's bytes (10 streams) equal its one-stream pass on both chips.Fidelity: agreement with oMLX
The 966 bodies (1,833 questions) against oMLX's recorded answers. A question agrees when both name the same top label; |dp| is the largest difference in any label's probability.
Every difference is a near-tie where this branch picked oMLX's runner-up and oMLX's top two were within 0.01: on the M5, four
fitspredicates (oMLX gaps 0.0066 to 0.0083) and onepickchoice (0.4857 / 0.4835, a gap of 0.0023); on the M3, sevenfitspredicates (gaps 0.0013 to 0.0083). The logit of oMLX's top label differs with an sd of 0.0426 on both chips. On the 116 requests: 216/216 top labels equal on the M5, 215/216 on the M3 (onefitsat oMLX 0.5007 / 0.4993).Performance
Zig 0.17.0. Latency is the 116 requests replayed one at a time and 4 at once (each request's wall clock, p50 and p95, and the replay's wall time). Main has no before-run (it doesn't serve this model), so the comparison is oMLX on the same M5 and checkpoint, recorded on 31 Oct; I didn't rerun it, since that oMLX serves other work on that Mac. Both Macs also served other models during these runs, and the two M5 rounds differ for that reason.
On the M3 this branch takes 836.4 ms p50 and 1,995.3 ms p95 one at a time (104.73 s), and 4,101.0 / 6,913.8 ms 4 at once (115.63 s). On the M5 the replay takes 1.87× oMLX's wall time one at a time in round 1 (2.28× in round 2), and 1.73× (2.21×) 4 at once. Decisions are exact here first; speed is a follow-up (the 5-bit members of this checkpoint run on the generic row kernels, as on #1826). Peak footprint 13.81 GB on the M5 and 10.86 GB on the M3 (
--context 8192, 8 decision streams).Build and host tests (M3, at
a4109159c):zig build native -Dcpu=apple_m1 -Dversion=1.0.0exit 0 (tensorfold-nativesha25602584c7d78294ae9…);zig build test test-golden -Dcpu=apple_m1 --summary allfrom an empty cache exit 0, 594/594 steps, 9,274/9,329 tests passed and 55 skipped (#1851's tip: 9,230/9,285); goldens 55/55 and 322/335 equal, 0 unexpected differences;zig build dev-tools -Dcpu=apple_m1exit 0;python3 tools/zig/lean_check.py424 problems at each of the eight commits, main's 424. The M5 built the same tip (zig build native tf-qwen27-run, exit 0,tensorfold-nativesha256bfcc4802f85159b2…). Main's kernel contract registry has no entry forq27_attn_prompt(qwen27's only one isq27_take), so there is noSUMMARYline for the new block instantiations; registering them is under Not run.Not run
0x36/elephant-1-9B-MLX-4bitor anmlx_lm.convertofbaselabs/elephant-1@b3f8a93 at 4 bits:tensorfold-native serve <dir> --context 8192 --prompt-cache-gib 0 --no-drafts, then the same replays.q27_attn_promptand its block instantiations: a test type inzig/tests/contract/and launches inregistry.zig, then theSUMMARYlines on M3 and M5.--context 8192.max_rowspiece at a time between chat rounds, as you suggested for lanes: a queued decision takes a free slot beside the live streams, as a new stream's prompt pass would #1192, only matters for such a checkpoint and is left until one exists.atc overlap --json(ATC isn't installed here). By hand, withgit merge-treeagainst every open pull request that touches these files: Implement GLP control vectors with a fused Metal kernel for Qwen3.8-27B #1835 conflicts inqwen27_served.zig; qwen3_5_moe: the experts' prompt sort as one threadgroup (stacked on #1125) #1799 and the 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 pair series add conflicts inattention_prompt.metalandattention.zigbeyond the ones they have with Qwen3.5-9B support in Zig engine #1826; Metal: the dense Qwen3.5 loader takes its shape from the config #1817 touches Qwen3.5-9B support in Zig engine #1826'ssupported()only. I'll rebase onto whichever lands first.