Skip to content

core: add the row-split pair seam and split simd_qmm over K - #1432

Merged
ashhart merged 3 commits into
ashhart:mainfrom
akol1:qwen27/pair-seam
Oct 29, 2026
Merged

ashhart merged 3 commits into
ashhart:mainfrom
akol1:qwen27/pair-seam

Conversation

@akol1

@akol1 akol1 commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

Step 1 of 6 of the two-Mac Qwen3.8-27B series, rebased on main (7c068895). The branch has two commits: the seam (dcccfe7b), then the pick's planted control (52b11e46). Steps 2-6 stack on them (qwen27/rank-slice, qwen27/pair-lifecycle, qwen27/pair-tests, qwen27/pair-tools, qwen27/pair-docs), and each of their bodies says what it stacks on.

This commit adds the seam a pair exchanges through, and nothing above it. zig/src/core/pair_rows_layout.zig fixes the exchange window's ABI; zig/src/core/pair_rows.zig and zig/kernels/metal/core/pair_rows.metal implement the four launches: pair_stage copies this rank's own half of a frame row into the exchange slot (a bf16 copy, so it carries the frame's bits), pair_gather copies each rank's half back to its own offset in the frame row (a bf16 copy, no arithmetic), pair_argmax reduces one rank's logit span lo..lo+n to (value, id) with the lower id on a tie and a row-level non-finite flag, and pair_pick chooses over both ranks' argmaxes — a flagged row or the broken-link word draws the same 0xFFFFFFFF sentinel, so one downstream check covers both, and that sentinel is never an id. zig/src/core/pair_rows_loopback.zig is the peer for a host without a second Mac: both ranks on one GPU through exactly the seam. A broken link makes pair_gather write bf16 NaN (0x7FC0) on both halves, so a row that lost its peer never reads as a valid one.

The branch's second commit gives pair_pick's flag rule a planted control of its own: zig/tests/contract/metal_pair_controls.zig replaces ((am | ap) & PAIR_FLAG) != 0u with false and gates it at nonfinite, where both ranks' words are finite — so dropping the rule is invisible at shapes and only the flag test catches it. metal_main.zig routes the new pick_no_flag control to the pick wave, and the contract section now counts six controls. That is the same shape as argmax_no_flag, which holds the argmax's rule.

simd_qmm.chunks(n, whole) makes the pair's K split the rank's own total width, so a split matrix's K chunk tree is the tree one Mac uses over that width; every unsplit simd_qmm.call keeps passing 0 or the full width, so its bits are main's. zig/src/core/simd_qmm.zig pins the arithmetic, e.g. chunks(3168, 6240) == splits(6240) and chunks(6240, 0) == splits(6240).

glm_dsa's pair transport (zig/kernels/metal/core/pair.metal, zig/src/core/pair.zig) is untouched: this commit neither edits nor re-registers it, and no record here names its kernels.

Receipt

  • Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision: this commit (52b11e46defdfe8fc82db1266fd21838f177750d) on top of main (7c0688951584497dad8b0d5894b657d432fa917f); no MLX or PyTorch; macOS 27.0.1 on an Apple M4 Pro (Mac16,11), Metal only (applegpu_g16, no tensor units), Apple metal version 32023.921. No checkpoint: the contract gates run the launches on the GPU and compare them with float64; they do not read weights.
  • Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false: n/a at this commit — nothing serves yet, and the seam's exactness is the contract's own word-level comparison. The four pair launches pass all four gates, e.g. PASS nonfinite metal pair stage pair_stage: 5 non-finite and 40955 finite outputs match the float64 classes; gather's line is the same shape with 10 non-finite and 81910 finite; PASS nonfinite metal pair argmax pair_argmax: the stated rule holds: a winner is (value bits, id): NaN loses every comparison, a row of only NaNs reports the -inf sentinel with id 0xFFFFFFFF, and a row with any non-finite logit reports NaN bits and id | 0x80000000; PASS nonfinite metal pair pick pair_pick: the stated rule holds: the pick copies one rank's (value, id) word pair; an id with bit 31 set, and the broken-link word, draw 0xFFFFFFFF: the same sentinel, so one downstream check covers both. PASS limits: 313 registered Metal launches against M1-M5, 0 problems is the same run's register check; SUMMARY 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings; 6 controls: 6 red as planted, 0 missed; 0 limit problems; gates 31 ms, limits 127 ms (the run reads --only "metal pair", so the tally is the pair's four launches and their six controls).
  • The six planted controls all fail as planted, which is the evidence that the four launches' gates can catch a changed seam: FAIL shapes control: metal pair stage, this rank stages the peer's half: at 8x5120x5120 output 0 is 8.593750e-1 where float64 gives 5.444336e-2: inf of its bound; FAIL shapes control: metal pair gather, this rank writes its own half on both sides: at 8x10240x5120 output 0 is -7.656250e-1 where float64 gives 1.572266e-1: inf of its bound; FAIL shapes control: metal pair argmax, a tie goes to the highest id: at 8x1x124160: row 0: (2e0 186240) where its own rule gives (2e0 124160); FAIL nonfinite control: metal pair argmax, a nonfinite logit is not flagged: row 0: value 2e0 where a non-finite logit must write NaN; FAIL nonfinite control: metal pair pick, a flagged row is not turned into the sentinel: row 0 picks 429031 where its own rule gives 4294967295; FAIL shapes control: metal pair pick, a tie goes to the higher id: at 32x1x1: row 2 picks 702 where its own rule gives 302. Each is followed by CONTROL …, red at <gate>, as planted in the log.
  • A K-split change is also contract-checked where it is used: zig build test-kernel-contract's filtered --only simd run passes four kernels and zero controls, SUMMARY 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems; gates 30263 ms, limits 114 ms.
  • Decode speed before and after, from tools/bench_openai.py: n/a, no serving path at this commit; the pair's speed needs two Macs and is measured in step 5's receipts, not here.
  • Prompt speed before and after: n/a at this commit, and simd_qmm keeps every unsplit call's split, so a one-Mac prompt reduction is bit-for-bit the same tree. The series-level cold-prefill A/B on one Mac, final tree (7718b1c2, binary sha256 0ed3d1a35d8332c4…) against main (7c068895, binary sha256 415303535b9593fd…), both tensorfold-native 1.1.8 (cached_tokens == 0, 2048/8192/16384 rows, 2 repeats a side): series/main median ratios 1.00016 at 2048 rows (92.310 vs 92.295 tok/s), 1.00013 at 8192 (90.108 vs 90.097) and 1.00018 at 16384 (87.065 vs 87.049); the medians differ by at most 0.02% while one side's own two reps differ by up to 0.05%, so neither a speed-up nor a slowdown is claimed. This GPU binds the simdgroup prompt variant (applegpu_g16, no tensor units), so that A/B is the fallback path, not the M5 prompt-kernel cell.
  • Tests run, with pass, fail and skip counts: zig build test: 568/568 steps succeeded; 1279/1298 tests passed (19 skipped), exit 0. zig build test-kernel-contract: 43/43 steps succeeded, exit 0. zig build dev-tools: 215/215 steps succeeded. zig build tf-qwen27-tp-check: that step does not exist at this commit (no such step, expected: step 4 adds it). zig fmt --check build.zig zig: clean (fmt rc=0). Lean rule: lean: 425 problems, 425 on the base, 0 new against main's baseline, exit 0 (tools/zig/lean_new.py).
  • Not run, and why: the model gates (tf-qwen27-tp-check arrives in step 4, on weights this commit does not need); a two-Mac run (the serving link arrives in step 3 and the supervisor in step 5); the unfiltered kernel corpus (it stops at UNSUPPORTED metal moss image: the family runs on tensor units (M5) alone and needs a recorded rotary file, so the gates above run with --only). The series' two-Mac acceptance receipt was recorded before this rebase, and the second Mac is gone (the cluster was broken), so that receipt cannot be re-run and stays dated to its run; CUDA, HIP/RDNA, M1-M3, M5 and Windows are not in this room.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. — this commit adds the exchange kernels only; no serving or scheduling path changes, and the loopback peer drives the same lane rounds.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. — no serving path at this commit; step 3 refuses drafts and step 4's fixtures are greedy, so there is nothing drafted to compare yet.
  • No precision traded for speed. If the bits change, the description says which sums change and why. — the seam moves bf16 rows with a copy and does no arithmetic across ranks; the only intended arithmetic change is simd_qmm's K split, which stays inside the rank's own width and is pinned by host tests.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. — the four launches have records (zig/src/core/contract/metal_pair_records.zig), a contract section in docs/kernel-contract.md:146, and the SUMMARY above.
  • Prompt processing is no slower than the last release. — n/a, this commit changes no prompt path; every unsplit prompt projection keeps main's K split.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not. — Metal on an M4 Pro (applegpu_g16); the launches are also offline-compiled against the M1 and M2 limits by the contract run; no CUDA, HIP, M1-M3, M5 or Windows.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. — the parent commit has no pair records, so the contract records and the six controls fail to build or gate there; the loopback peer needs no second Mac.
  • Comments and docstrings are one line. No measurements or history in the source. — the lean rule over the diff is clean (0 new problems); the kernel comments state the rule, not a measurement.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines. — the diff names 127.0.0.1 only in the loopback peer and no user path, hostname or peer address.

@ashhart

ashhart commented Oct 27, 2026

Copy link
Copy Markdown
Owner

Thanks for the pair work: the transport, the watchdog and the contract records for the four new kernels are good. The four kernels pass their gates on an M3 Ultra here. One thing has to change before this stack can land, and it starts in this PR.

The pair's tokens must equal the single Mac's. The seam computes each layer's o_proj (or GDN out_proj) and down_proj as two fp32 K-halves, one per rank, then rounds bf16(mine + peer). That is a different reduction order from the single Mac's one-pass sum. At chat-256 token 236 the single Mac's lead is 0.125 logit (one bf16 ulp at that magnitude, 24.625 vs 24.5), so it isn't a tie there. On the pair, the reordered sum closes the gap and the reply diverges. A policy that accepts this goes beyond what CONTRIBUTING and the recipe skill allow: a one-step tie policy covers only the same two contenders with both margins inside the tie.

An exact design keeps every sum on one rank. Split o_proj, out_proj and down_proj by output rows, not by K:

  1. All-gather the bf16 input activations.
  2. Have each rank compute its half of the 5,120 outputs with the full-K reduction the single Mac uses.
  3. All-gather the outputs.

No cross-rank partial sum exists then, and no tie policy is needed. It costs about two all-gathers per seam instead of one reduce, with the same weight memory per rank.

GLM-5.3's two-Mac path. This rewrites the shipped glm_dsa transport: the device spin-wait becomes a host relay that reopens the encoder at every exchange, and several Pipeline signatures change. The arithmetic looks unchanged, but please include a GLM two-Mac A/B against main: identical token shas, and decode within noise. A host round trip per exchange, two per layer, may slow it.

Checked here: the branch merges cleanly into main, zig build dev-tools builds, fmt is clean, lean has no new findings, and the pair kernel contract passes 4/4. Nothing needing the checkpoint or two Macs was run.

@akol1

akol1 commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor Author

Thanks for the review. The blocking item is addressed: the pair's tokens now equal the single Mac's, byte for byte, on every fixture in the gate.

Exactness. Stage-3 pair gate with --dump, four fixtures: chat31 full mode at 96 and 256 steps (chunk 128), prompt-n4hi full and mlp modes at 8 steps (chunk 1). All fields green: pair_logits_byte_equal, pair_logits_compared, ranks_token_equal, reference_token_equal, shard1_logits_byte_equal; first_difference and bad_pick_at null; seam_gave_up_rank0/1 both 0. At the prompt merge the top-2 margin (ids 13962, 8160; 25.625 vs 25, a 0.625 lead) is byte-identical across single, rank 0 and merged. A stage-1 cross-tree check matches the pre-stack receipt for the 256-token chat stream: equal single_mac_tokens element for element, equal single_mac_prompt_state (f46a4012...746b68), equal single_mac_resident_bytes (19,232,096,256), and true shard1_tokens_equal/shard1_prompt_state_equal.

Stage 1 against main's own binary. Rather than compare loadShard(ranks=1) with load only on the new code, I built upstream/main at b50d4e23 in a separate worktree (zig build native: 57/57 steps) and replayed the same 31-id chat prompt to it as a raw id array (/v1/completions with an integer prompt, return_token_ids: true, max_tokens: 256), so no template sat between the two runs. Main's reply: sampling: greedy, 256 ids, first [13962, 2500, 11296, 56727, 1633, 37539, 264, 21966], last [799, 1324, 303, 279], token_sha256 a06fab348fbe653ec676860c774c20c30a3bc37c57465018b4a9e4b5430b8f52. That digest is exactly the comma-join hash of the stack's stage-1 single_mac_tokens, so main's binary and the stack's single-Mac path agree token for token.

Mechanism. The old seam computed o_proj, GDN out_proj and down_proj as two fp32 K-halves and summed them, which cannot be bit-equal. It is now split by output rows: each rank reads the whole input K, computes its rows with full K, and writes bf16 at its row offset (2,560 of 5,120 hidden rows; 8,704 of 17,408 for down_proj). pair_stage copies a rank's half of the relay frame; pair_gather joins own and peer rows at row offsets with no arithmetic between them. They replace pair_sum/pair_part. A broken link still writes bf16 NaN.

Exactness also needed a reduction-tree fix: simd_qmm chose its K-chunk tree from the output width, so halving a projection changed fp32 reduction order where the count straddled a threshold (zba 6,240 to 3,168; qkv 10,240 to 5,120). The new chunks(n, whole) keeps splits(n) when no split rank is present, and split callers pass the checkpoint's whole width. Unsplit runs are unchanged.

Receipts. tf-kernel-contract --only "metal pair" --seed 7: 5 kernels pass gates 1-4, 0 fail. --only pair --rotary-accuracy: 16 kernels pass, 30/30 controls red as planted. --only 'simd_qmm q4': 2 kernels pass. zig build test-kernel-contract: 18/18 steps, 101/101 tests. Clean-cache zig build test: 563/563 steps, 7,766/7,813 tests passed, 47 skipped, exit 0. Lean rule: 0 new problems against base (428 on both). zig build dev-tools: 214/214 steps, exit 0. zig fmt --check clean.

One honesty note: the two-Mac A/B on GLM is not reproducible on 2x64 GiB, so this PR's transport coverage rests on the qwen27 pair receipts rather than that A/B. There is no GLM two-Mac receipt, and there will not be one.

Next step. The refreshed work is 16 commits on upstream/main (b50d4e23, head 1b3f0ed7), 58 files, +4,156/−531, on fork branch akol1/TensorFold:qwen27/stack-main. Main moved by two test-only commits after the first push of that branch, so I rebased it onto the new tip and re-ran every receipt above there; the 16th commit is the simd_qmm call-site arity fix in the affine reference tool that zig build dev-tools found. The six open PR branches are unchanged and will not be force-pushed. Would you prefer to (a) review the refreshed stack from that branch (compare: main...akol1:qwen27/stack-main), or (b) retarget/refresh the PRs on your side? We would rather not rewrite history that has already been reviewed. Since this comment, main advanced one more commit (9ed939bb, the offline kernel-contract limit, #1493); the same stack rebases onto it without conflicts, and the gates I re-ran there pass (build 63/63, dev-tools 214/214, host contract 18/18 with 120/120 tests (the newer base adds 19 contract tests), the three focused filters 5/5, 16/16 with 30/30 controls red as planted, 2/2, and all four stage-3 receipts green). I have not republished that rebase; say the word if you would rather review it on 9ed939bb than on b50d4e23.

@ashhart

ashhart commented Oct 27, 2026

Copy link
Copy Markdown
Owner

Thanks @akol1. Splitting by output rows with full K on each rank is the exact design we hoped for, and checking the single-Mac path against main's own binary is a useful receipt.

The PR branches haven't changed, though. #1432's head still sums fp32 K-halves across ranks (pair.metal:92, simd_qmm.zig:41), so on this PR the blocking item is still open. A single 58-file branch is too much to review as one change, so we'd rather not take (a). Please push the refreshed commits onto the existing branches instead, one step per PR, rebased on current main, each saying what it stacks on. A force-push is fine here: the review stays in the comments, and re-reviewing each step keeps the split you made.

On GLM-5.3: since there won't be a two-Mac GLM receipt, please leave glm_dsa's transport as it is on main. That means the changes to glm_dsa/engine.zig, model.zig and prompt.zig, and the Pipeline signatures they need. Use the new seam for qwen27 alone. A shipped path can't change without an A/B against main.

Once #1432 carries the refreshed code, we'll review the row split and the simd_qmm K-chunk change. Unsplit callers must stay byte-identical. We'll also run the pair contract, and later tf-qwen27-tp-check --stage 3 once the PR that adds it is refreshed.

@akol1 akol1 changed the title core: add the pair seam and register its Metal kernels with the kernel contract core: add the row-split pair seam and split simd_qmm over K Oct 27, 2026
@akol1

akol1 commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor Author

Refreshed and force-pushed, one step per branch, all rebased on current main (5897e0bc):

step branch head stacks on
1 qwen27/pair-seam 0b79f9c9 main alone
2 qwen27/rank-slice 7d57101a step 1
3 qwen27/pair-lifecycle 277929ca steps 1-2
4 qwen27/pair-tests 0c2df0ec steps 1-3
5 qwen27/pair-tools 0b4f1f4e steps 1-4
6 qwen27/pair-docs 0c500376 steps 1-5

Each commit message carries the same "stacked on ..." line, so the stack is readable from git log. The 58-file branch is not part of this: qwen27/stack-main is untouched, and its content now lives in the six steps. Every body is rewritten against main (5897e0bc), and every receipt in them was produced on the refreshed commits.

GLM. glm_dsa/engine.zig, model.zig and prompt.zig are back to main, and the Pipeline signatures they need are unchanged: the seam is qwen27-only. There is no GLM two-Mac receipt in this series and none is promised.

Honesty on what was and was not run. The pair in #1435 runs both ranks through pair.Loopback, a host-memcpy peer on one Mac, so the bodies quote no two-Mac number and claim no cross-Mac transport. The kernel gate is tf-kernel-contract --only "metal pair" --seed 7: SUMMARY 5 kernels: 5 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems, and PASS limits: 281 registered Metal launches against M1-M5, 0 problems. The pair table's PASS nonfinite/PASS repeat/PASS shapes lines are in #1432's body. The 4/4 control result is your M3 Ultra report, and the 30 controls: 30 red line is the rotary table, not this run; I am not attributing either to the pair.

Unsplit callers. chunks(n, whole) returns splits(n) for whole = 0, every unsplit call site passes the literal 0 (zig/src/families/qwen27/dflash/runtime_callbacks.zig:32 and the simd_qmm self-check, with its scalar twins), and zig/src/core/simd_qmm.zig:185-186 pins chunks(3168, 6240) == splits(6240) and chunks(6240, 0) == splits(6240). A ranked caller passes the whole matrix's width, so a rank that owns rows keeps the chunk tree one Mac uses; a split 10,240-wide qkv no longer takes 8 chunks where one Mac takes 16. zig build test-kernel-contract is 18/18 steps and 120/120 tests on the refreshed heads.

--stage 3 is on the refreshed #1435 (qwen27/pair-tests, 0c2df0ec): tf-qwen27-tp-check --model DIR --prompt-file F --steps N --chunk C --stage 3 --dump. Stage 3 runs stage 1 (a rank-local shard against the single-Mac load: greedy ids, logit bytes, prompt-state digest) and stage 2 (two rank-local shards, every layer seam through the loopback peer); a token difference returns error.PairTokensDifferFromSingleMac with no exemption. On the recorded 31-id chat prompt at 256 greedy tokens the pair's token_sha256 is a06fab348fbe653ec676860c774c20c30a3bc37c57465018b4a9e4b5430b8f52, which is also what main's own tensorfold-native produced from 5897e0bc for the same prompt.

The row split and the simd_qmm K-chunk change are step 1 and step 2 review material as you asked; happy to answer on any step.

@ashhart

ashhart commented Oct 27, 2026

Copy link
Copy Markdown
Owner

Thanks @akol1 for moving the seam to an output-row split and for leaving glm_dsa's transport as it is on main. pair.zig, pair.metal and the Pipeline signatures are untouched. We merged this into main locally. fmt is clean and the lean count is unchanged. test-kernel-contract, test-kernel-seed-host, compile-gemma4-contract-host and check-qwen27-speed-host pass. Every unsplit simd_qmm.call passes 0 or the full width, so its bits are main's.

  1. pair_add_rms (zig/kernels/metal/core/pair_rows.metal:32-60, zig/src/core/pair_rows.zig:297-309) still adds the two ranks' fp32 parts (pair_rows.metal:48). Nothing in the stack calls it, and it is the cross-rank sum the row split exists to avoid. Please remove it, with its record (metal_pair_records.zig:7), its probe and the docs line.
  2. The receipt shows 0 controls. Please plant controls for the four remaining kernels and show each one red. For example: a gather that writes the peer's half at its own offset, an argmax or pick tie that keeps the higher id, and a stage that reads the wrong half.
  3. One Mac fails a row that holds any NaN or ±Inf (tree_round.metal:23,45). pair_argmax skips NaN and lets +Inf win (pair_rows.metal:111-112), so the pair returns a token where one Mac returns an error. Please carry the row's non-finite flag through pair_pick to the sentinel, and add a control for it.
  4. The branch no longer merges. docs/kernel-contract.md, zig/src/core/contract.zig and zig/src/core/contract/registry.zig conflict with main's latest contract landings. Please rebase.

Next we will run the pair contract and the simd_qmm contract with its controls on an M5 Ultra.

@akol1
akol1 force-pushed the qwen27/pair-seam branch 3 times, most recently from 05d69b8 to 55065d5 Compare October 28, 2026 04:09
@akol1

akol1 commented Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Refreshed and force-pushed: qwen27/pair-seam is now 55065d55, rebased onto current main (db3fd5ab). The body is rewritten around this head. glm_dsa's transport is untouched — the 20-file diff touches no glm, dsa, mcdma, fabric or transport file, and no record here names its kernels.

  1. pair_add_rms is removed: kernel, record, probe and docs line. The seam is pair_stage and pair_gather only, both bf16 copies, and nothing in it adds two ranks' fp32 parts. Every unsplit simd_qmm.call still passes 0 or the full width.

  2. Each of the four kernels has a planted control, and each is red as planted:

    • FAIL shapes control: metal pair stage, this rank stages the peer's half: at 8x5120x5120 output 0 is 8.593750e-1 where float64 gives 5.444336e-2: inf of its bound
    • FAIL shapes control: metal pair gather, this rank writes its own half on both sides: at 8x10240x5120 output 0 is -7.656250e-1 where float64 gives 1.572266e-1: inf of its bound
    • FAIL shapes control: metal pair argmax, a tie goes to the highest id: at 8x1x124160: row 0: (2e0 186240) where its own rule gives (2e0 124160)
    • FAIL shapes control: metal pair pick, a tie goes to the higher id: at 32x1x1: row 2 picks 702 where its own rule gives 302

    Each is followed by CONTROL …, red at <gate>, as planted. The filtered pair run on this head is SUMMARY 15 kernels: 15 pass gates 1-4, 0 fail, 0 known findings; 35 controls: 35 red as planted, 0 missed; 0 limit problems; gates 316 ms, limits 135 ms.

  3. The row's non-finite flag reaches the sentinel. pair_argmax reduces a rank's logit span to (value, id) plus a row-level non-finite flag, and pair_pick fails closed when either rank's row is flagged — a flagged row or the broken-link word outranks any value, so the pair never returns a token where one Mac returns an error. The gate says so: PASS nonfinite metal pair argmax pair_argmax: the stated rule holds: a winner is (value bits, id): NaN loses every comparison, a row of only NaNs reports the -inf sentinel with id 0xFFFFFFFF, and a row with any non-finite logit reports NaN bits and id | 0x80000000. The control for that is FAIL nonfinite control: metal pair argmax, a nonfinite logit is not flagged: row 0: value 2e0 where a non-finite logit must write NaN.

  4. Rebased onto db3fd5ab: docs/kernel-contract.md, contract.zig and contract/registry.zig are main's, with no conflict.

Gates on 55065d55: zig build test 557/557 steps succeeded; 5026/5055 tests passed (29 skipped); test-kernel-contract 43/43; dev-tools 215/215; zig fmt --check build.zig zig clean; --only simd SUMMARY 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems; gates 32923 ms, limits 120 ms; lean 428 problems, 428 on the base, 0 new against main's 429. tf-qwen27-tp-check does not exist at this commit; step 4 adds it.

@ashhart

ashhart commented Oct 28, 2026

Copy link
Copy Markdown
Owner

Thanks @akol1. All four points are in on 55065d55, and the branch merges cleanly into current main:

  1. pair_add_rms is gone, with its record, probe and docs line. Nothing in the seam adds two ranks' parts.
  2. Each of the four kernels has a planted control, and your run shows each one red at its gate.
  3. pair_argmax flags a row holding any NaN or ±Inf, and pair_pick turns the flag into the sentinel, so a pair fails where one Mac fails.
  4. Other families' simd_qmm and simd_qmm_bits calls pass 0, and on one Mac qwen27's gn is the matrix's own rows, so one Mac keeps main's chunk tree.

One small addition, not blocking: a pair_pick control that ignores bit 31 (the flag test replaced by false), so the pick's own flag rule has a control that goes red, as argmax_no_flag does for the argmax.

Next is our run of the whole series on an M5 Ultra, at #1437's head merged onto main:

  • the kernel contract: --only "metal pair" with its five controls, --only simd with its controls, and the full Metal summary;
  • tf-qwen27-tp-check --stage 3 --dump on TensorFold/Qwen3.8-27B-MLX-4bit: full and mlp, --chunk 1 and 128, --decode-prompt at 128, and the 512-id prompt;
  • stage 1 on the 512-id prompt at --chunk 512 and at --chunk 128, which should give the same prompt-state digest. One Mac serves 512-row prompt rounds on an M5 and the pair serves 128-row rounds, so this is the check that the cap leaves one Mac's bits alone;
  • one Mac on its own against main: equal token hashes from bench_concurrent.py --levels 1,4 --alone --serial, drafted equal to "draft": false, and decode and cold prefill within noise.

If those pass, the series can land in order, starting with this PR.

@ashhart

ashhart commented Oct 28, 2026

Copy link
Copy Markdown
Owner

Thanks @akol1. Here is our run of the whole series on an M5 Ultra: #1437's head (660f7a35) against main db3fd5ab, with TensorFold/Qwen3.8-27B-MLX-4bit.

The pair checks all pass:

  • Kernel contract --only "metal pair": 4 kernels pass gates 1-4, 5 controls red as planted, no limit problems. --only simd: 4 kernels pass. The unfiltered Metal run stops before its first gate on main and on the series alike, because this Mac has no rotary accuracy receipt (UnrecordedRotaryCompiler). So there is no full summary yet.
  • tf-qwen27-tp-check --stage 3: full and mlp at --chunk 1 and 128, --decode-prompt at 128, and the 512-id prompt. Both ranks' tokens equal the reference, and the pair's logits are byte-equal every time.
  • Stage 1 on the 512-id prompt: --chunk 512 and --chunk 128 give the same prompt-state digest, and it equals one Mac's.

One Mac against main, alternating base, new, new, base:

  • Tokens: all 20 prompt rows equal main's. Every stream in the 1 to 8 stream sweep equals its run alone and its serial run, and chat resume equals fresh with thinking on and off.
  • Decode: the served sweep is within 0.1% at 1, 2, 4 and 8 streams. The prompt cells can't resolve decode here, because main's own two runs differ by up to 4% at 8k.
  • Cold prefill is slower at every size. At 2k, 8k and 16k every series run is below every main run:
prompt main (tok/s) series (tok/s) median
2k 1,648-1,656 1,638-1,647 -0.6%
8k 1,611-1,616 1,602-1,607 -0.5%
16k 1,555-1,564 1,550-1,554 -0.5%
32k 1,446-1,467 1,455-1,459 -0.4%
64k 1,292-1,303 1,290-1,296 -0.4%

One Mac should run main's program at main's speed, so this is the one thing left before the series lands. The two Metal kernels whose one-Mac code changed are the first place I'd look. q27_attn_prompt_merge now indexes with args.kv_heads, and q27_attn_merge with the heads buffer, where main has the compile-time PG * 4 and G * 4. A function constant, or a separate instance, would give one Mac main's code there. If that isn't it, a prefill A/B at each PR's head should show which commit costs it. We'll rerun the prefill cells on the fix.

@ashhart

ashhart commented Oct 28, 2026

Copy link
Copy Markdown
Owner

Thanks @akol1. q27_attn_prompt_merge_single and q27_attn_merge_single with the 24 heads fixed at compile time are what we hoped for, and the pair_pick flag control and the 4-bit qualification are in too. We've queued the same M5 Ultra cells on the new head (dccc7ead): one Mac against main for prefill, decode, the sweep and chat resume, then the pair contracts and the stage 1 and 3 checks. If one Mac's prefill is back to main's, the series lands in order.

@akol1
akol1 force-pushed the qwen27/pair-seam branch 2 times, most recently from c152fc5 to cf6f0cb Compare October 28, 2026 17:00
@akol1

akol1 commented Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Heads-up before you spend more cells: main moved to dd0a9cc5 while our receipts ran, and four of the six branches went CONFLICTING on it. The series is rebased onto dd0a9cc5; heads now are

The two _single kernels and the if (split) dispatch are identical to what dccc7ead carried — git diff dccc7ead 380ae404 -- zig/kernels is empty, and the only attention.zig delta is dd0a9cc5's own attention_memory.zig extraction. So cells pinned to dccc7ead still measure the fix, but the run that decides the landing should be on 380ae404 (base dd0a9cc5). The full reply, with the AIR diff and fresh receipts on the new heads, follows.

@akol1

akol1 commented Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Thanks — the two kernels are back to main's code on one Mac, and all six branches are rebased onto current main (dd0a9cc5). Heads: cf6f0cb4 (this PR), bb00f2b0, c7dcbce2, c7114a1a, d2e1e62b, 05bb79a9.

The fix is the separate instance, and it is in the second branch's kernel commit (bb00f2b0). q27_attn_merge_single (zig/kernels/metal/qwen27/attention_mpp.metal:330) is q27_attn_merge minus the heads argument, and q27_attn_prompt_merge_single (attention_prompt.metal:175) writes PG * 4 where the split prompt merge reads PG * args.kv_heads (:170). encode chooses a.io[if (split) 4 else 10] (attention.zig:405) and a.prompt.?[if (split) 1 else 2] (:361), split is the caller's r.reduce != null (decode_round.zig:449), so a one-Mac pass binds main's instance and never sets the heads uniform (:408 runs only when split). A function constant was not an option: there is no MTLFunctionConstantValues anywhere under zig/src or zig/kernels, so this tree has no mechanism for one.

They are main's kernels, not lookalikes. I compiled dd0a9cc5's tree and the series tip's tree with the flags zig build gives these sources (metal -x metal -std=metal4.1 -fmetal-math-mode=safe -fmetal-math-fp32-functions=fast -c), read both all.airs with metal-objdump, and compared the one-Mac kernel each side would bind. This PR's head does not touch either source file — git diff dd0a9cc5 cf6f0cb4 -- zig/kernels/metal/qwen27/attention_mpp.metal zig/kernels/metal/qwen27/attention_prompt.metal is empty, and the _single variants are added in bb00f2b0 — so the main side of each comparison is dd0a9cc5's own function:

  • q27_attn_merge (main) vs q27_attn_merge_single (series): 329 IR lines each, 0 differing.
  • q27_attn_prompt_merge vs q27_attn_prompt_merge_single: 125 IR lines each, exactly one line differing — dereferenceable(28) against dereferenceable(32) on the %struct.Q27PromptArgs addrspace(2)* parameter. Those four bytes are the int kv_heads member the split kernel reads; the single kernel never loads it, and its body is main's text with PG * 4 in the store. So the only thing one Mac lost is a uniform load.

Correction to my last note, before you spend on it: I wrote that git diff dccc7ead against the series tip -- zig/kernels is empty. That sentence was wrong — over that range zig/kernels is 67 files, +756/−613, because main moved under the series. What is byte-identical between the two trees is the set that carries this change: attention_mpp.metal, attention_prompt.metal, zig/kernels/metal/core/pair.metal, zig/src/core/pair.zig, zig/src/families/qwen27/pair_rows.zig, and the two if (split) dispatch lines in attention.zig (that file's only delta in the range is main's extraction of the scratch layout into attention_memory.zig, +3/−10). So the cells you queued on dccc7ead are measuring this fix. The final tree's own kernel delta against dd0a9cc5 is four files, +221/−2: pair_rows.metal (+129) and sources.zig (+2) from this PR, and the two attention files from step 2.

A/B on one Mac, in the shape of your prefill table. Both builds served the same checkpoint and returned the same token_sha per length (36790ecd55c2 at 2k and 8k, 8d1ede4f889e at 16k, both builds), alternating order, one warm-up request per process, two repetitions per length; the A/B loop refuses any recorded request whose reply does not report prompt_tokens_details.cached_tokens == 0:

prompt main (tok/s) series (tok/s) ratio
2k 89.2 90.7 1.0174
8k 88.6 88.8 1.0032
16k 86.2 85.5 0.9912

On this Mac a one-Mac pass binds the fixed prompt instance directly: with split false, promptRows sets a.prompt.?[2] (attention.zig:361), which is q27_attn_prompt_merge_single. None of the four serve logs from the A/B carries attention.zig:229's "prompt rows take the decode path" line, which prints only when the simdgroup prompt kernel here asks for fewer than 256 threads a group (:228; tensorUnits() is generation() >= 17 on zig/src/metal/device.zig:68-70 and this M4 reports applegpu_g16, so the simdgroup variant is the one it builds). The merge half of the prompt path is therefore the kernel the AIR check above covers; the chunk kernel differs by variant on an M5 (tensor ops) versus here (simdgroup matrices), so this table prices the merge half and your queued M5 cells price the rest. The three ratios are 1.017, 1.003 and 0.991 while the spread between two repetitions on this Mac is wider than that (90.8 against 87.6 at 2k), so nothing here separates the trees — it cannot reproduce your M5 Ultra's cold-prefill cell, but it is the same before/after comparison and it is clean.

Fresh receipts on this head (cf6f0cb4, rebased): zig build test 568/568 steps succeeded; 6207/6241 tests passed (34 skipped), exit 0; zig build dev-tools 215/215; zig build test-kernel-contract 43/43; the pair filter SUMMARY 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings; 6 controls: 6 red as planted, 0 missed; 0 limit problems; gates 41 ms, limits 163 ms and the simd filter SUMMARY 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems; gates 74490 ms, limits 380 ms; lean 425 problems, 425 on the base, 0 new; tf-qwen27-tp-check does not exist at this commit (it arrives in step 4, c7114a1a), so no receipt from it here.

Your pair_pick control is in, in this PR's second commit: metal_pair_controls.zig replaces ((am | ap) & PAIR_FLAG) != 0u with false, and it is red as planted — the sixth control in that summary line.

@akol1

akol1 commented Oct 29, 2026

Copy link
Copy Markdown
Contributor Author

Rebased onto the current main (7c068895) and force-pushed; the head is now 52b11e46defdfe8fc82db1266fd21838f177750d. The series order is unchanged, and these are re-run numbers from the rebased head, not carried over.

zig build test: 568/568 steps succeeded; 1279/1298 tests passed (19 skipped); zig build test-kernel-contract: 43/43; zig build dev-tools: 215/215; zig fmt --check build.zig zig clean; filtered pair run gates 31 ms, limits 127 ms, filtered simd run gates 30263 ms, limits 114 ms; lean rule 425 problems, 425 on the base, 0 new.

The series-level cold-prefill A/B on one Mac, re-run at the rebased tip against main (7c068895, cached_tokens == 0, 2048/8192/16384 rows, 2 repeats a side): series/main median ratios 1.00016 at 2048 rows (92.310 vs 92.295 tok/s), 1.00013 at 8192 (90.108 vs 90.097) and 1.00018 at 16384 (87.065 vs 87.049), while one side's own two repeats differ by up to 0.05%, so neither a speed-up nor a slowdown is claimed.

On your M5 cells: this rebase did not touch either _single kernel — git diff 55865b01 7718b1c2 -- zig/kernels/metal/qwen27/attention_mpp.metal zig/kernels/metal/qwen27/attention_prompt.metal is empty — so a one-Mac pass still binds the fixed instance, and the cells you queued measure the same merge kernels on these heads.

This room is one Mac now: the second Mac the pair receipts were recorded on is gone (the cluster was broken), so a fresh two-Mac pair_accept.py run is not possible here. The bodies keep that run labelled as a dated receipt from before the rebase and put the single-Mac loopback gate (tf-qwen27-tp-check) and the host tests next to it. No other question from the earlier comments is still open.

akol1 added 2 commits October 29, 2026 23:46
Step 1 of 6 of the two-Mac 27B series; stacked on main (1b9a987) alone,
and steps 2-6 stack on this commit.

The seam exchanges whole output rows instead of K halves: rank r owns
2,560 of the 5,120 output columns (8,704 of 17,408 for the down
projection), reduces over the full K on its own rank, and `pair_gather`
copies whole rows in bf16 without doing any arithmetic, so no exchange
rounds or sums a partial product across ranks. A broken link surfaces as
bf16 NaN `0x7FC0`.

The seam's four kernels (`pair_stage`, `pair_gather`, `pair_argmax`,
`pair_pick`) move bf16 frames and compare values only, so they carry no
sums and no accumulator's error. One Mac fails a row holding any NaN or
±Inf, so `pair_argmax` reports such a row as NaN bits with bit 31 set in
its id word, and `pair_pick` turns that flag into the id sentinel: both
Macs then refuse a non-finite round where one Mac would, through the same
downstream check as a broken link. Five planted controls (a stage on the
peer's half, a gather writing its own half twice, an argmax and a pick
tie to the higher id, an argmax that does not flag) each go red on the
gate named in `docs/kernel-contract.md`; the next commit adds the sixth,
for `pair_pick`'s flag rule itself.

That transport is a new module, `pair_rows.zig` and `pair_rows.metal`,
over `pair_rows_layout.zig`'s slots and the relay protocol unchanged.
`pair.metal` and `pair.zig` stay for GLM-5.3-Flash, whose transport this
series does not touch: a shipped path cannot change without an A/B
against main.

`simd_qmm` gains a K-chunk width: `chunks(n, whole)` lets a rank that
still holds the whole output width keep the split one Mac derives from
that width, instead of one derived from the rank's narrower slice.
The pick's sentinel rule was the one pair rule without a planted
control. The argmax control proves the row's non-finite flag is written,
but nothing proved `pair_pick` acts on it: a pick that dropped the
`((am | ap) & PAIR_FLAG) != 0u` test would still pass every gate.

The sixth control replaces that test with `false`, planted only at the
nonfinite fill -- at shapes both ranks' words are finite, so dropping
the rule is invisible there and the control has to be red at the gate
that plants it. The control is red as planted, and the pair docs now
count six.
@ashhart
ashhart merged commit 3f0a715 into ashhart:main Oct 29, 2026
1 of 9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants