[cshintov]

What this changes

On a tensor-parallel rank of the 27B (12 query heads over 2 KV heads, from #1433's rank slice), keep's gather and scatter indexed 4 heads a staging row, while Scratch sizes keep_k and keep_v at the binding's 2. A keep that accepts more than half the frame's rows wrote past both buffers. Now the two kernels index by the binding's KV heads; the 27B on one Mac computes the same indices as before.

• Issue: n/a, a bounds fix to merged code (qwen27: rank-slice the 27B family and gather the whole embedding table #1433's rank slice).
• Stacked on: main.
• atc overlap --json: atc is not installed here, so I ran git merge-tree against every open pull request touching these five files (30). This adds a conflict with 4, none with the change itself:
  • qwen27: serve Qwen3.5-MoE (Ornith-1.5-122B) on Metal: routed experts, its DFlash2 drafter and the publisher's 4-bit weights #1125, qwen3_5_moe: the experts' prompt sort as one threadgroup (stacked on #1125) #1799 and Implement GLP control vectors with a fused Metal kernel for Qwen3.8-27B #1835: the kernel registration list in zig/tests/contract/metal_main.zig (adjacent appends).
  • Qwen3.5-9B decisions in Zig engine #1871: zig/src/families/qwen27/attention.zig, the args structs next to PromptArgs (adjacent lines).

Commits:

• contract: qwen27 attention_io's pack, append and keep at 24 over 4 and a rank's 12 over 2
• qwen27: keep's gather and scatter stage the binding's KV heads a row, not 4

Where it reaches. The host keep and the device keep alike: both stage only the accepted rows (tree_commit.metal marks the rest 0xffffffff, and gather and scatter return on them). With default frames (512 rows and up) a round accepts too few rows to reach it; a smaller TF_PROMPT_ROWS frame does. Gather and scatter used the same wrong index, so the caches came out right and only the staging overran, which is why no existing check saw it.

Design.

• The two kernels take Q27KeepArgs { kept, kv_heads } in place of kept, and index row * kv_heads + head.
• Attention.keep fills it from the binding's geometry, and refuses a scratch whose KV heads differ from the binding's, since the staging belongs to the scratch.

Alternatives rejected.

• Staging sized at 4 heads always: it hides the mismatch and doubles a rank's staging.
• A compile-time KV-head define: the attention library has none. An argument matches how q27_attn_append already takes kv_heads in Q27PackArgs.
• Reading the grid's y extent in the kernel: implicit, and wrong if a launch ever pads the grid.

Checks on this head

• Head commit: df7e488
• Lean checker, python3 tools/zig/lean_check.py: 424 problems on the base, 424 on this head, 0 new.
• Tests run: zig build test-qwen27, test-qwen27-keep-host, test-qwen27-plans-host and test-kernel-contract pass; zig fmt --check build.zig zig passes.
• CI: pending.
• Not run, and why: a two-Mac rank serve with real 27B weights (I have one Mac); the device keep end to end through tree_commit's metadata kernel (the contract writes its map in the same layout, read from tree_commit.metal).

Receipt

• Environment: TensorFold ef9fa70 plus these two commits, Zig 0.17.0, macOS 26.5, Apple M4 Pro (applegpu_g16, no tensor units). No checkpoint: the change is checked by the kernel contract.
• Paths and platforms the change reaches: the qwen27 Metal keep on one Mac (unchanged indices) and on a rank of two (fixed). Run on the M4 Pro; the limits check covers M1 to M5 launch limits; not run on an M5.
• Exactness: n/a for served tokens on one Mac, where the indices are unchanged; the contract checks cache and staging rows exactly.
• Decode speed before and after: n/a, the same two launches with one more argument word.
• Prompt speed before and after: n/a, keep does not run in prompt processing.
• Licence of any new weights or drafter: n/a.

Kernel contract, only when a GPU kernel is new or changed

• Kernels registered with the contract: qwen27 attention_io's pack, append, host keep and device keep, at 24 over 4 and 12 over 2. Each keep checks the caches and the staging rows themselves (staging row i, head h at (i * kv_heads + h) * 256); the device keep keeps exactly rows / 2 + 1 rows of one chain, with the map as tree_commit writes it.
• Planted controls: an append one row late, and keep staging 4 heads a row on a rank (both kernels' index planted back). Both fail as planted.
• SUMMARY lines (tf-kernel-contract --only qwen27):

Run	Test commit (bug present)	Fix commit
Plain	8 kernels: 6 pass gates 1-4, 2 fail (rank keep, rank device keep: "staging keys row 1 head 0 is not stream 0's row 66"); 1 control red	8 kernels: 8 pass gates 1-4, 0 fail; 2 controls red as planted; 0 limit problems
MTL_SHADER_VALIDATION=1	the same 2 rank keep cases fail; 24 over 4 passes	8 pass, 0 fail; 2 controls red

The limits check covers 481 Metal launches against M1 to M5 with 0 problems. Exit status 0 on the fix commit.

• Chips I could not run: M5 (tensor units).

Checklist

• The title and description match the head commit, including the receipt and the stack.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.
• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo: not run with weights; on one Mac the indices are unchanged.
• No precision traded for speed. The bits do not change.
• Decode and prompt processing are no slower than main: not measured; the change adds one argument word to two launches.

[ashhart]

Thanks @cshintov. An independent review on current main (d8b80bd):
Verdict: LAND after the Metal kernel contract run, with one small addition. It's exact and a real bounds fix.

Applying it: both commits apply cleanly on main d8b80bd. The lean check stays at 424 at each commit. Both author
lines are @cshintov's noreply. It's Metal only: one .metal file, qwen27/attention.zig and contract tests. No CUDA
file changes.

Main's bug

• Scratch sizes keep_k and keep_v at rows × kv_heads × 256 bf16 words (attention_memory.zig:17, the last
  two sizes). Each is its own Metal buffer.
• q27_attn_gather and q27_attn_scatter are launched over (dim, kv_heads, kept). But they bounded at.y < 4 and
  indexed the staging row as (row × 4 + head) × 256 (attention_io.metal, both kernels).
• At 4 KV heads the index is right.
• At 2 KV heads the grid's y is 0..1, so only heads 0 and 1 run. But staging row r lands at r × 4 + h, which is
  twice as far as the buffer allows.
  • Any kept row r ≥ rows / 2 writes past the end of keep_k and keep_v.
  • Scatter reads back from the same wrong offsets. So when nothing else touches that memory, the caches come out
    right and only the staging overruns. That's why no check saw it.
  • An out-of-bounds GPU write into separate allocations is undefined. Main could corrupt whatever lies past those
    buffers, or lose the staged rows. Nothing was refused.

The fix is exact

• Both kernels take Q27KeepArgs { kept, kv_heads } and index (row × kv_heads + head) × 256.
• The cache side (at.y × stride + source × 256 + x) is unchanged.
• At 4 KV heads, the bound at.y < a.kv_heads and the index are identical to main's, so the same bytes are kept
  and restored. The launch grid is unchanged, and the argument grows from one word to two.
• At 2 KV heads, staging now fits its buffer exactly. Each kept row's heads are copied out and back to the same
  cache positions, so the restored caches hold the same bytes main's restored caches would have held.
• Attention.keep now refuses with AttentionKeep when the scratch's KV heads differ from the binding's. The staging
  belongs to the scratch, so a mismatch is caught before any launch. The refusal comes after the existing kept and
  cache checks.

Which shapes it matters on

Geometry.check (attention_plan.zig:49-53) admits four shapes, and config.parse reads num_key_value_heads
straight from config.json:

Shape (heads / KV heads)	Where	Main
24 / 4	the 27B on one Mac	correct; unchanged here
16 / 4	the 9B on one Mac	correct; unchanged here
12 / 2	a rank of the 27B split two ways	overran on a keep of more than rows / 2 rows; fixed
32 / 2	GQA16 on one Mac (the qualification test names Ornith-1.5-122B)	the same overrun; also fixed

The PR body names only the rank case. The 32-over-2 one-Mac shape had the same bug and gets the same fix. The body's
"the 27B on one Mac is unchanged" is right.

Tests and the contract

• The first commit adds the qwen27 attention_io contract suite: pack, append, host keep and device keep, at 24/4 and
  12/2. Every keep checks the caches and each staging row at (i × kv_heads + h) × 256 against host copies.
• Its planted control puts main's × 4 back in both kernels (the Edit now allows count = 2) on the rank shape, and
  must go red. The append-one-row-late control is kept.
• The author's table: on the test commit, the two rank keeps fail ("staging keys row 1 head 0 is not stream 0's row
  66"). On the fix commit, 8 of 8 pass, both controls are red, and there are 0 limit problems over 481 launches. So the
  new tests fail on main's kernels.

Small addition (not blocking)

Register the keep at the other two admitted shapes too: 32/2 (the one-Mac GQA16 case this also fixes) and 16/4. Today
Heads covers only q27b (24/4) and rank (12/2).

Device cells

• any Apple Silicon Mac (and an M5) | no model | tf-kernel-contract --only qwen27, also with
  MTL_SHADER_VALIDATION=1: 8 kernels pass gates 1-4, both controls red, 0 limit problems
• Mac with 32 GB+ | Qwen3.8-27B-MLX-4bit | one Mac, against main, greedy and seeded with drafts: equal token_sha256
  (the 24/4 path is unchanged)
• two Macs, TP=2 | Qwen3.8-27B-MLX-4bit | a small TF_PROMPT_ROWS frame so a round keeps more than half its rows,
  drafts on: tokens equal to the one-Mac run and to "draft": false, with shader validation clean

[ashhart]

Thanks @cshintov. The Metal kernel contract at this head (--only qwen27, M5 Ultra) passes: 8 kernels through gates 1-4, both controls red as planted, 0 limit problems. With the review's LAND, this is queued to land.

[cshintov]

Thanks @ashhart, and for catching the 32-over-2 one-Mac shape: the same overrun, the same fix. The keep registrations at 32/2 and 16/4 are ready as one test-only commit (16 of 16 kernels pass on an M4 Pro, both controls red; with main's kernels the 32/2 keep fails and 16/4 passes). I'll send it as a follow-up once this lands, so your M5 run on this head stands.

[ashhart]

Thank you @cshintov. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

[ashhart]

Thanks @cshintov! This is in TensorFold 1.2.2, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.2