perf(qwen27): a drafted round overlaps its cache commit with the next draft, and the drafter's attention and top-k use more of the GPU

[jkuepker]

Three changes to DFlash2 drafted decode on Metal. Tokens, draft rounds, verified rows and accepted drafts are unchanged, and the plain decode binary is byte-identical to main's.

1. A drafted round's cache and state commit runs behind the next draft (4f5f387b). accept() now splits the round's command buffer in two: the target's picks with the accepted rows' taps and head row (all the host waits for), then the keep (KV rows, recurrent state, history), which nothing reads before the next verify. The keep is committed without a wait and waited for just before that verify, and before a run returns, so the drafter's forward overlaps it. This is only correct if nothing encoded in between reads or writes what the keep writes. The drafter's forward, the taps ring and Runner.ensure from Qwen3.8-27B serves concurrent requests: shared rounds across runner slots, DFlash2 for one stream, and tf-qwen27-bench #578 don't touch those buffers, and debug asserts at the target's entry points (commitDevice, accept, absorb, advance, prefill, promptChunk, decodeChunk) check that no keep is still pending. The served path (runRound) finishes the keep at once, so the overlap is in the CLI run path only. The drafter's first command buffer also commits after six ops instead of forty, and its row copies go out as one dispatch.
2. Shared-context attention loads a tile of keys before it folds them (654c85b5). Each lane loaded one key's score and value and folded it before the next key's load started. A tile of four keys now loads together, then folds in key order with the same fp32 chain, so every output byte is unchanged.
3. Wide bf16 top-k rows are split over sixteen threadgroups (6ca24f79). One threadgroup scanned a whole row, so a 15-row draft head over 98592 columns used 15 of the GPU's cores. Rows of 4096 columns or more now scan in sixteen spans, and a second dispatch merges them with the same tournament and the same (score, lower token) order. A compile-time check ties the kernel's lane, span and max_k constants to the Zig side, so a change to one of them fails the build instead of the merge. zig/tests/bf16_topk_bench.zig, which tools/zig/bf16_topk_bench.sh builds, now dispatches the single-group kernel directly as its control.

The three can land separately. 2 and 3 also run on the served path; 1 doesn't.

Exactness

All on one Apple M5 Max, against origin/main e7c523eb (1.0.7) built the same way.

check	result
top-k kernel: 576 cases over widths 4096 to 248320 (ties, both zeros, NaN and infinities, padded strides)	entries, counts and non-finite masks byte-equal to the single-group kernel's
tools/zig/bf16_topk_bench.sh on the GPU	72 exact cases; the new path and the control both byte-equal to the CPU oracle, Row ABI and guards pass
shared-context attention: 390 masked, fully masked and unmasked cases at 1 to 2063 keys, bf16 and fp32 outputs	byte-equal to the old kernel
qwen27_gate.py check (4 prompts × 256 greedy tokens, --chunk 128 --compare-chunk 16)	tokens equal to main for all 4 prompts; prompt_state_byte_equal 4/4
drafted == plain (tf-qwen27-dflash-run --prototype q4 --compare-plain, DFlash2)	4/4; rounds 89/30/80/32, verified rows and accepted drafts equal to main
qwen27_http_gate.py (raw/chat × plain/draft)	16/16, token SHAs equal to main
concurrent (bench_concurrent.py --levels 1,4 --alone --serial, code and chat)	12/12 at 4 streams, serial 1/1, equal to main
tf-qwen27-streams (real weights), 4 streams	pass; output identical to main's
model-free device checks: tf-qwen27-synthetic, -lanes-check, -served-check, -dflash-served-check, -dflash-synthetic	exit 0; every log byte-identical to main's

The gate exits 1 on main too: story (token 68) and code (token 21) leave the Python 0.6.6 oracle at a BF16 near-tie. This branch gives exactly main's tokens on all four.

Speed

Same Mac and settings throughout, with no other model process running. Builds of origin/main e7c523eb and of this branch, run alternately (ABBA).

Drafted decode, CLI (tf-qwen27-dflash-run --prototype q4, DFlash2 4-bit, 256 tokens, 5 interleaved pairs per prompt). Tokens, rounds, verified rows and accepted drafts are equal in every run.

prompt	main	this PR	tokens/s	pairs faster
code, 24 tokens	5.391 ms/token	5.197	+3.7% (paired median +3.9%)	5/5
16384 tokens	6.016	5.371	+12.0% (paired median +11.9%)	5/5

At 16K main's runs spread 6.3%, with two fast runs; against main's fastest run the branch's median is still 6.9% faster.

Served drafted (native server, tools/zig/qwen27_served_bench.py, DFlash2 4-bit, 3 repeats, 2 sessions per build, alternated with main and two other branches): copy −1.1%, arithmetic +0.9%, story −0.4%, code +0.6%. tools/bench_openai.py with the drafter: fibonacci −0.2%, chat +2.1%. All are inside main's own session-to-session spread, so I'm not claiming a served change. The served code prompt is the same 24 tokens as the CLI's short case. The likely difference is commit 1, which only the CLI path uses, but I haven't measured the commits separately on this base. The served prompts are all short (14–31 tokens), so there is no served long-context number.

Cold prefill (fresh KV, --prompt-cache-gib 0, 2k/8k/32k): no change. Plain decode wasn't timed: tf-qwen27-run built from this branch is byte-identical to main's.

Host checks

• zig build test test-golden -Dcpu=apple_m1: 143/143 steps, 1290/1292 tests passed, 0 failed, 2 skipped. test-golden: 309/319 equal, the 10 known differences, 0 unexpected. All the same as main.
• python3 tools/zig/lean_check.py: 534 problems on main and 534 on this branch, 0 new.
• zig fmt --check is clean on the touched files.

Setup: Apple M5 Max (40-core GPU, 128 GB), macOS 27.0 (26A428), Xcode 27.0, Zig 0.17.0. Checkpoint TensorFold/Qwen3.8-27B-MLX-4bit @ 22d8d538; drafter z-lab/Qwen3.8-27B-DFlash2 @ 50307d4c, 4-bit. Not run: CUDA and ROCm (the changed kernels are Metal), and Apple M1–M4.

[ashhart]

Thank you, @jkuepker. Three careful speed changes to the 27B on Metal, each with its exactness argument written down, and all three hold on our machines: the in-place DeltaNet scan (#711), the decode attention without the tail kernel (#712), and the draft-side changes (#713). Your ten commits landed in 1.0.9 under your name, rebased onto 1.0.8 (only #713's first commit changed: three cb.failure() != null checks became the gpuFailed(...) helper that 1.0.8 added). Each PR's head is merged into the history with no change to the tree (7098195be, fed0e24d1, 4a8002d8d), so GitHub shows all three PRs as merged.

What we ran, with all three together against the tree they started from:

• Exactness. On an M3 Ultra and an M5 Ultra, 17 cases (9 greedy, thinking included, and 8 seeded sampled, 96 tokens each): drafted equal to "draft": false on both builds, and the same tokens across builds. On the M5 Ultra, 9 prompts' full raw token arrays were equal across builds, and each drafted reply's final state equalled its plain run's. On the M3 Ultra, a conversation with a regenerated, an edited and a deleted turn gave the same replies with the prompt cache on and off, and 2 and 4 concurrent streams gave each stream's solo reply.
• Speed, one stream, 256 tokens. M5 Ultra (four runs a build): greedy code 406.1 to 417.0 tok/s, prose 100.6 to 103.6; seeded code 294.0 to 316.9, prose 86.2 to 89.0. M3 Ultra (six runs): greedy code 150.6 to 154.7, prose 53.3 to 54.8. Prompt processing at 16k, 32k and 64k tokens on the M3 Ultra was 0.2 to 0.4% faster.
• On the release. On an M5 Ultra, the 17 cases gave the same tokens on the release candidate as on 1.0.8. In the 1.0.9 release gate, against 1.0.8, an M5 Max gave 1.0.8's tokens in all 17 cases and 12 chat turns, resumed equal to fresh, and an M3 Ultra did too while decoding prose 2.2 to 2.8% faster and code level to 1.5% faster (a fresh server's code 0.8 to 2.0% faster), with prompts 0.4 to 0.7% faster.