[totaland]

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

Why

Array.q4mv's body adds each row's terms in index order: one dependent F32 chain per row. The CPU loop can only go wide across rows, and the Metal kernel from #1435 needs one thread per row, which reads the weights strided. A faster kernel has to add in a different order. To keep the same bits on every lane, the Bend body has to define that order.

Change

The order (T32), for row i and group g (64 weights, code words k = 0..7, nibbles j = 0..7):

1. Lane j: L_j = 0; for k in 0..7: L_j = L_j + q(k,j) * x[64g+8k+j]. Each mul and add is rounded; there is no FMA.
2. qx = ((L0+L4)+(L2+L6)) + ((L1+L5)+(L3+L7)).
3. t_g = scale_g*qx + bias_g*xsum_g.
4. 32 partials: p[l] = 0.0 + t_l + t_{l+32} + ..., with groups in order.
5. Fold p[l] += p[l+o] for o = 16, 8, 4, 2, 1, then round p[0] to BF16 (nearest even).

• bend2/base.bend: the body and its doc follow T32.

• bend2/comp.ts:

  • a scalar row (q4_row) and a NEON loop that runs 4 rows at once (q4_neon4, under __ARM_NEON);
  • the Metal kernel uses one SIMD group per row, 8 rows per 256-thread group. Lane l owns groups l, l+32, …, and step 5 is simd_shuffle_down;
  • one bounds check is shared by the NEON and GPU paths.

• tests/base/array_q4mv_tree.bend (new):

  • a 5x2176 product: 34 groups, so the partials wrap;
  • rows 1..4 alone;
  • lane and fold cancellations at ±2^25 and ±2^24.

  The expected bits come from an independent C model of the order.

This changes results. After BF16 rounding, 1 of 6144, 1 of 2048 and 58 of 248320 rows differ from #1431's order on real weights. array_q4mv.bend's expected lines do not change.

Verification

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

• array_q4mv.bend and array_q4mv_tree.bend pass on the interpreter, C --gpu off, C --gpu 8GB and JS.

• The new test catches the defects it should:

  • Array.q4mv: affine 4-bit matvec primitive (CPU 3.2-3.6x, bit-identical; over comp.ts cap) #1431's order fails it on all three lanes.
  • A NEON lane-pairing mutation fails only the C lane.
  • The C model's mutants (wrong lane meet, 16 partials, sequential fold, reversed fold) all differ.

• Kernel microbench. The Metal tree kernel, the NEON tree loop and the C reference are bit-identical at these shapes:

  • 13x128, with a 3-row offset tail;
  • 64x2048, 192x2048, 2048x2048 and 7760x2048;
  • 6144x2048, 2048x6144 and 248320x2048.

  All 18 AB/BA runs exited bit-exact.

• Whole model. In a Bend inference engine (TensorFold), Qwen3.5-2B (MLX 4-bit) gives bit-identical full logits on the CPU (NEON) and with --gpu 8GB (248,320 of 248,320). Greedy token IDs match the old order.

• 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 prints the same hash on CPU and Metal (1289093127). Six alternating rounds, milliseconds per call (400 calls, single-threaded host):

  	--gpu off	--gpu 8GB
  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 (index order)	2.36	0.30
  this PR	1.09	0.21

• Kernel microbench timings (best of 10 reps, 3 AB/BA pairs, against 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's kernels):

  shape	Metal	CPU, 1 thread
  6144x2048	0.55 → 0.29 ms	1.06–1.10x
  2048x6144	0.74–0.91 → 0.22–0.35 ms	1.09–1.17x
  248320x2048	5.44 → 3.64 ms	1.10x

  On Metal each lane now reads 32 contiguous bytes, and the head reaches 78 GB/s.

The CPU gain differs between the two measurements: 2.2x in the 400-call program, but 1.06–1.17x in the microbench. I have not explained the difference. Treat the CPU gain as somewhere in that range. In the inference engine, CPU decode improved by about 4% (127 → 121.5 ms/step, 10 pairs), which is inside the noise.

• 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: 66,466 (#1435) → 67,215 (+749).

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.
• This changes the documented order of a Base primitive, so the bits differ from Array.q4mv: affine 4-bit matvec primitive (CPU 3.2-3.6x, bit-identical; over comp.ts cap) #1431.
• The NEON path is arm64 only; other hosts use the scalar row.
• 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.