Metal: queued rmsq, swq, gdn and atq kernels and a 4-row q4mv, so a decode token is one command buffer with one wait

[totaland]

Stacked on #1431, #1435, #1436 and #1437. Only the last commit is new here.

Why

With #1437, a model's Q4 products can share one Metal command buffer, but everything between them (norms, activations, recurrent and attention mixers) still runs on the CPU. So a decode step still has one wait per dependent stage. For Qwen3.5-2B (MLX 4-bit) that is 97 waits per token, and the GPU idles between short command buffers, which also makes the products themselves 2–3× slower than when they run back to back.

This PR adds the missing per-token kernels as queued intrinsics, so the whole token can be encoded into one command buffer and waited on once.

Change

New Base intrinsics in bend2/base.bend. Each has a pure Base body that defines its meaning (the interpreter and JS run it), a C twin in bend2/comp.ts that follows the body's order, and a Metal kernel that is queued like Array.q4go: barriers on both sides, and the result is ready after the next Array.q4wait. Without a Metal heap, each runs its C twin at once.

• Array.rmsq: RMS norm with weights as Array<U32> (F32 bits), plus the 64-value group sums that Array.q4mv reads. The Metal kernel sums with simd_sum trees in one threadgroup.
• Array.swq(a, o, n): SwiGLU, round(silu(a[i]) * a[n + i]), plus the group sums. One SIMD group per 64-value group.
• Array.gdn(...): one token of a gated DeltaNet mixer: conv window and activation, q/k L2 norms, decay and beta, the delta-rule rows, the gated norm and group sums. One threadgroup per value head; every phase is local to its head, so there is no cross-threadgroup sync.
• Array.atq(...): one token of gated attention over GPU-resident key and value caches: q/k RMS norm and rotation, the new k/v into the caches, scores, softmax, weighted values, sigmoid gate and group sums. One threadgroup per query head. The first head of each KV group writes the cache row; the others use the new k/v from threadgroup memory, so no head reads a row another head writes in the same dispatch. Metal handles positions below 4096 (ATQ_MAX); past that the C twin runs at once on the CPU, so the caller must wait first.
• bend_q4mv: a SIMD group now does 4 rows and reads each group's 64 activations once for all of them. Each row's arithmetic and reduction tree are unchanged from Array.q4mv sums in a fixed tree order, so NEON and one Metal SIMD group per row keep the body's bits (stacked; over comp.ts cap) #1436, so it stays bit-identical to the Bend body.
• tpl in comp.ts now reads $10 and up (/\$(\d+)/). Before, $10 was split into $1 followed by 0. No upstream intrinsic has 10 or more arguments, so this only shows with gdn and atq.
• Tests (new): tests/base/array_rmsq.bend, array_swq.bend, array_gdn.bend, array_atq.bend.

q4mv, rmsq and the plumbing are exact. swq, gdn and atq are not bit-identical on Metal: Metal's exp, its summation order and fused multiply-adds move the last bits. array_gdn.bend prints its state rounded to bfloat16 for that reason (Metal moves the last F32 bit of 5 of its 64 values).

Verification

Apple M1 Pro, macOS 26, GPU runs one at a time.

• The four new tests and array_q4mv.bend, array_q4mv_tree.bend, array_q4go.bend pass on the interpreter, C --gpu off, C --gpu 8GB (Metal) and JS against their #| lines.
• Regression. All 186 tests that upstream 60fa05dc has in tests/base and tests/io give byte-identical interpreter, C and JS output on upstream and on this branch (temp paths in JS build-error messages normalized). 120 of 186 match their #| lines on both trees. The 7 new Q4/kernel tests pass on all three lanes.
• Whole model. In a Bend inference engine (TensorFold), one decode token of Qwen3.5-2B (MLX 4-bit) is now one command buffer: 194 kernels, about 9.4 ms of GPU time, one wait. Measured with these commits on a local build that also carries File.read_into reads a file's bytes straight into an Array<U32> #1430, Metal's OUTLINE is noinline, so a large kernel stays under the AGX backend's 256 flow states #1433 and Metal: a full lane stack spills to the heap, so non-tail recursion under ! no longer faults at 1023 frames (#1393, #1394) #1439:
  • Greedy token IDs equal the CPU forward at every step over 40 and 200 tokens, and with the model state swapped out and back before every step. Max absolute logit difference: 0.234 over 40 tokens, 2.19 over 200. With the GPU off, the fused path is bit-identical to the CPU path.
  • Two sessions decoded in turn on one model each match their solo runs.
  • Four alternating pairs, previous queued path (Array.q4go/Array.q4wait queue Metal Q4 products in one command buffer and wait once on a flag (stacked; over comp.ts cap) #1437) vs this one: 66/11, 60/11, 63/11, 65/11 ms per decode step. A/A pairs: 11/11, 11/11.
• 4-row product kernel. With CPU work still between the waits it made no difference (24 vs 24 ms/token, 4 pairs). Once the GPU stayed busy for the whole token it gave 12 vs 18–19 ms/token (4/4 pairs). 2 and 8 rows were no better.

uvx ttok@0.3 < bend2/comp.ts: 68,289 (#1437) → 78,654 (+10,365). Most of it is the C twins of gdn and atq.

Limits / open

• Far over the comp.ts cap (64,000), like the rest of this stack. I'd welcome direction on whether model-specific kernels like gdn and atq belong in the compiler at all, or whether a smaller set of generic queued primitives would be preferred.
• Not bit-exact on Metal for swq, gdn and atq (see above). Greedy outputs matched in every run, but logits drift.
• Lifetime. As with q4go, every array the GPU reads or writes must stay alive until Array.q4wait; nothing in the types enforces it.
• atq past 4096 positions falls back to the CPU; that path is implemented but was not exercised (runs stopped at 205 positions).
• Tested only on an M1 Pro. No CUDA path.

[Lorenzobattistela]

Thanks @totaland, this is a lot of careful work, and the measurements are useful. This PR went up before Nico's note on #1435 to #1437, so to close the loop: we're closing it for the same reasons, plus a few we found while testing it on an M4 Pro.

• Base would carry kernels for one model family (rmsq, swq, gdn and atq, about 880 lines of base.bend across the stack), and comp.ts would be at 78,654 tokens against its permanent 64,000 cap.
• The queued ops don't always match their Bend bodies. Base says Array.q4wait(x) is x, and the checker relies on that, but on Metal, taking the wait out of array_q4go.bend prints the untouched 7.0 slots. Array.atq past position 4096 runs on the CPU right away, so right after a queued Array.swq all 256 of its outputs differ from the interpreter's. Shapes outside the fast path (Array.swq with n = 32) leave o untouched in C and on Metal, while the interpreter and JS fill it.
• Programs that use Array.q4mv or Array.q4go fail --verdict, because BendTT can't check the fold/foldb mutual recursion.
• On an M4 Pro, 2 of array_atq's Metal values differ from its expected lines, so the test gate fails there (1602/1603).

The per-token kernels fit better in a library. A package can already ship its own .c and .js effects that take an Array and run Metal kernels. We tried a small one that queues three kernels in one command buffer and waits once. As effects, the order of queueing and waiting is explicit in IO. The cost is that --verdict reports them as foreign code in any program that imports them.

The sync cost you found is real, and it isn't specific to Q4. On main, each bang is one command buffer with a blocking wait, about 0.23 ms on an M4 Pro even for an empty call. A general way to avoid that wait is worth an issue, as Nico suggested, if you want to take it on.