Skip to content

core: add the pair seam and register its Metal kernels with the kernel contract - #1432

Open
akol1 wants to merge 1 commit into
ashhart:mainfrom
akol1:qwen27/pair-seam
Open

akol1 wants to merge 1 commit into
ashhart:mainfrom
akol1:qwen27/pair-seam

Conversation

@akol1

@akol1 akol1 commented Oct 27, 2026

Copy link
Copy Markdown
Contributor

What this changes

Adds the seam that carries one decode round's partials between two Macs and registers its four Metal kernels with the kernel contract: pair_layout fixes which rows of each weight a rank owns, pair_loopback stands a peer in on one Mac, pair.metal writes the frame and completes the partial on the far rank, and a float store variant of simd_qmm lets a rank write its exact part. No serving path changes here; the qwen27 wiring follows in the next pull request.

Refs #1236 and #808. #1236 tracks a mixed Metal/CUDA group over RDMA; this is two Metal Macs over a socket, it leaves prefill unsplit, and the serving pull request states what it refuses. atc overlap --json against 8c36359d finds 172 open pull requests sharing at least one file with this stack; the six it could merge-test are three clean (#1336, #1295, #1429) and three conflicting only inside the two append-only registry tables (#1377, #1391, #1394, lfm2, another contributor), so whoever lands second rebases. The rest read unknown in this partial clone.

Receipt

  • Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision: 8c36359d plus this commit (137cf631); 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: this commit and its gates move bytes, they do not read weights.
  • Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false: n/a, no tokens are served by this commit. The seam's reference is float64, and tf-kernel-contract --only 'metal pair ' passes all four kernels' gates: pair add_rms 3 shapes within the bound (worst 0.139 of it) and 3 edges refused, 10248 non-finite and 71672 finite outputs in their float64 classes, five runs byte-identical over 327680 bytes, and the sum stays fp32 at k 6144 where a bf16 sum would miss by 672.2x; pair sum 2 shapes (worst 0.000) and 2 edges refused, five runs byte-identical over 81920 bytes; pair argmax holds its rule (NaN loses every comparison, an all-NaN row reports the -inf sentinel with id 0xFFFFFFFF), five runs byte-identical over 64 bytes; pair pick copies one rank's (value, id) word pair and reads an id of 0xFFFFFFFF as a broken link, five runs byte-identical over 128 bytes. SUMMARY 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems, and the same run's static table reports 265 registered Metal launches against M1-M5, 0 problems. Two caveats: --only filters subjects, so the 0 controls is a scope artefact and the two CONTROL red as planted lines in the log are the limits table's own planted controls; and this is a scoped diagnostic, not the full qualification receipt (docs/recipes/SKILL.md), which this GPU cannot produce — see below.
  • Decode speed before and after, from tools/bench_openai.py: n/a, no serving path to measure; the serving pull request carries the pair's numbers.
  • Prompt speed before and after: n/a, no prompt path changes, and the pair leaves prefill on both ranks by design.
  • Tests run, with pass, fail and skip counts:zig build test test-golden at this tree: 7116 passed, 45 skipped, 0 failed of 7161, with 535 of 535 build steps; the 45 skips are the image-codec and pixel cases that need libjpeg-turbo 3.1.4.1 or libwebp 1.6.0, plus SKIP qwen27-tp-check: no --model DIR; the filtered tf-kernel-contract run quoted above is separate from that count and exits 0 with SUMMARY 4 kernels: 4 pass gates 1-4, 0 fail, 0 known findings.
  • Not run, and why: the full tf-kernel-contract corpus stops at UNSUPPORTED gemma4 31B rounded dense: tensor units required; Gemma coverage unchecked and exits error: TensorUnitsRequired, so Gemma 4 31B needs a tensor-unit Mac; the rotary probes behind it want --rotary-accuracy from tf-rotary-sweep, which this draft does not claim; CUDA, HIP/RDNA, M1-M3, M5 and Windows are not in this room, so the CI zig-windows job and a reviewer's CUDA box have to carry them; and there is no maintainer-agreed gate for a new tensor-parallel family — zig/gate.json has no entries at 8c36359d and neither Tensor parallel across Metal and CUDA ranks (Mac + NVIDIA): tracker #1236 nor Two-Mac (--speed-up) for Qwen3.8-27B, and a mid-size model that fits a pair of 128 GB Macs #808 names a threshold, so this draft asks for that decision rather than inventing one.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. — no serving path changes here; the seam runs inside the round the lanes already execute, and pairLoopback is a test double for the far rank.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. — not reachable at this commit; the pair refuses drafts and keeps no prompt cache, and the serving pull request records why.
  • No precision traded for speed. If the bits change, the description says which sums change and why. — the frame is fp32 and the only arithmetic change is simd_qmm's TF_OUT float store, which exists so a rank writes its exact part; every bf16 store and every single-Mac sum is untouched.
  • 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 carry a spec, a record and a probe each, the metal pair limits route is registered, and the float store is a spec variant; Metal is the only backend in this change and its filtered SUMMARY is quoted above.
  • Prompt processing is no slower than the last release. — n/a, no prompt path changes; the serving pull request makes the same statement with its reason.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not. — ran Metal on an M4 Pro; no CUDA, no HIP, no M1-M3/M5, no Windows (see Not run).
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. — the four probes test kernels that do not exist on the parent commit, so red-before is a missing symbol rather than a failing test; the pair gates do skip cleanly when the GPU or the contract binary is absent, which is how CI runs them.
  • Comments and docstrings are one line. No measurements or history in the source. — the lean check reports 428 problems, 428 on the base, 0 new.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines. — nothing in the diff names a machine, a user path or an address; the loopback peer is 127.0.0.1, and every later script in the stack reads its host and paths from arguments or the environment.

The seam carries one round's partials between two Macs: pair.metal writes a
frame, the host relay owns the exchange instead of spinning on device
counters (which only works while both ranks share a launch group), and
pair_sum completes the partial on the far rank. The payload is f32, the
seam's accumulation, and simd_qmm grows the matching store. pair.zig splits
into pair_layout and pair_loopback so each module stays under the 600-line
rule.

Nothing registered the seam before: three of the four kernels already ran
under the older ABI on the shipped glm_dsa speed-up path, but there was no
record and no spec, so no gate covered the kernels this stack changes. Add
the four launch records, the seam spec with its four probes, the
`metal pair ` route in the limits table, and the fp32-out simd_qmm variant,
so the transport and family changes that follow have a SUMMARY to post.

The glm_dsa call sites and the two Pipeline.load test call sites are forced
by the shared signatures; the family changes themselves come later.
@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.

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