[totaland]

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

Why

Each GPU Array.q4mv commits its own command buffer and waits for it. That costs about 0.23–0.27 ms of host time per product, whatever the size. One decode step of a 2B-parameter 4-bit model has 187 products in 97 dependent stages:

• One wait per product: 86.7 ms/token in a microbench of those 187 products.
• GPU work alone: 26–32 ms.

Change

• bend2/base.bend:
  • Array.q4go(w, x, y, r, n, cols, rows) has the same body as Array.q4mv, so the interpreter and JS run it unchanged.
  • Array.q4wait(x) is x.
  • The doc states the lifetime rule.
• bend2/comp.ts:
  • q4go: in a native build with a Metal heap, Array.q4go encodes a product of any size into one open, concurrent command buffer and returns at once. Without a heap it runs Array.q4mv.
  • q4wait: Array.q4wait commits the buffer, ends it with a barrier and a 1-thread kernel that stores a sequence number in shared memory, and spins on that number. It falls back on the buffer's status. The flag saves about 0.27 ms per wait over waitUntilCompleted.
  • A program that calls Array.q4go loads its Metal library without needing a bang.
  • BEND_Q4_TRACE=1 prints each wait's product count, GPU time and wall time.
• tests/base/array_q4go.bend (new): array_q4mv.bend's 13x128 case through the queue. On the C --gpu 8GB lane it really queues, because q4go has no size threshold.

Queued products run concurrently, so none may read another's y. Put a wait between dependent stages. The GPU reads w and x and writes y until the wait, and Bend frees a value at its last use. So wait on x, and use w and y after the wait.

Verification

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

• array_q4mv.bend, array_q4mv_tree.bend and array_q4go.bend pass on the interpreter, C --gpu off, C --gpu 8GB and JS.
• Negative controls:
  • a build whose Array.q4wait returns without waiting dies with bend: memory fault;
  • a test that dropped w before the wait got a wrong row 0 every time. The runtime had freed w and written its free-list link into word 0, which the GPU then read. The doc now states the lifetime rule.
• Per-call cost. The 4096x2048 scratch program from Array.q4mv on Metal: one kernel thread per row for products of 4M weights or more, bit-identical (stacked on #1431; over comp.ts cap) #1435, with one q4go + q4wait per call, prints the same hash (1289093127) as q4mv on CPU and Metal. Six alternating rounds: 0.25 ms/call, against 0.36 for q4mv's own command buffer. These runs are noisy.
• Microbench of one decode step (the 187 products of one token, a wait per stage, 10 tokens):
  • one wait per product: 86.7 ms/token;
  • one buffer per stage with waitUntilCompleted: 50.5–53.8 ms/token;
  • one buffer per stage with this PR's flag wait: 24.2–28.5 ms/token.
• Whole model. In a Bend inference engine (TensorFold), the CPU runs the forward pass and queues each stage's products with q4go, then waits once per stage before reading. Qwen3.5-2B (MLX 4-bit) greedy token IDs and logits equal the CPU run over 40 tokens.
  • Eight queued runs (four alternating pairs of two engine builds that differ only in argmax): 59–73 ms/step, median about 65.
  • Three CPU runs of the same builds: 89–113 ms/step.
  • That is about 15 tokens/s, against about 11 on the CPU.
• 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 Array.q4go/Array.q4wait queue Metal Q4 products in one command buffer and wait once on a flag (stacked; over comp.ts cap) #1437, the top of this stack. The only difference is the temp path in 21 JS build-error messages. 120 of 186 match their #| lines on both trees. The new q4 tests pass on all three lanes.

uvx ttok@0.3 < bend2/comp.ts: 67,215 (#1436) → 68,289 (+1,074).

Limits / open

• Over the comp.ts cap, as Array.q4mv: affine 4-bit matvec primitive (CPU 3.2-3.6x, bit-identical; over comp.ts cap) #1431 is.
• Coherence. The flag wait relies on the CPU seeing a GPU store to shared memory after a device barrier. That held in every run on an M1, but the Metal documentation does not promise it. Deleting the spin falls back to waitUntilCompleted, at about 0.27 ms per wait.
• Lifetime. The rule is easy to break: drop w early and the GPU reads freed memory. A safer API would take the whole ((w, x), y) at the wait.
• Time between waits. In the inference engine, the CPU work between waits (norms, recurrence, attention) still dominates: about 55–65 ms per step. Moving it to the GPU needs more queued kernels.
• Tested only on an M1 Pro. No CUDA path.

[nicolas-abril]

Thanks @totaland for the careful work and measurements on this stack.

We're closing #1435, #1436 and #1437 for the same reasons as #1431:

• They are stacked on Array.q4mv: affine 4-bit matvec primitive (CPU 3.2-3.6x, bit-identical; over comp.ts cap) #1431, which is closed, and each adds more to bend2/comp.ts, which is already over its token cap. The top of the stack is at 68,289 tokens, about 4,300 over the 64,000 cap.
• A matvec specific to the 4-bit format, and now its Metal kernel, NEON loops and q4go/q4wait API, belongs in user code, not in Base.

One difference from #1431: there, plain Bend on main already matched the CPU speedup. The measurements here suggest stock bangs don't yet match a dedicated kernel on Metal (the ~11.5x head and ~15 vs ~11 tokens/s). That gap is worth closing, but with a general mechanism (faster bangs for this kind of loop, or a general async GPU dispatch/wait) rather than one specific to the 4-bit format. If you want to explore that, please open an issue first so we can agree on the design.