Files
colibri-strix/c/tests
Code Arranger 3716e4006a Metal backend (Apple Silicon): batched experts + fused attention on GPU, unified-memory zero-copy, gated behind COLI_METAL — 2.06 tok/s M5 Max (#72, #87, #103)
* docs: Metal expert-matmul backend design (Apple Silicon)

Empirically-validated design for a batched MoE expert-matmul Metal backend.
Microbenchmarks (scratchpad) establish: runtime-compiled Metal needs no Xcode;
V3 (float4 + threadgroup reduction) kernel is correct and fast; synchronous
per-matmul dispatch loses to CPU due to ~150us Metal launch latency, so the win
is batched full-layer dispatch (854us/layer, 707 GFLOP/s) reading expert slabs
zero-copy from unified memory.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: backend infrastructure + kernel-correctness test (M1)

Add backend_metal.{h,mm} — an opt-in Apple-GPU backend built with METAL=1 on
macOS. Runtime-compiled shader (no Xcode needed), zero-copy over unified memory.
Implements coli_metal_matmul (general quantized GEMV, f32/int8/int4/int2) via a
threadgroup-reduction + float4 kernel; batched moe_block is stubbed (returns 0 ->
CPU fallback) for M2. tests/test_backend_metal.mm validates all formats and edge
shapes (odd S, non-mult-4 dims) against a CPU reference (nerr ~2e-6). Makefile
gains a METAL=1 Darwin branch and a metal-test target. Default build unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: batched moe_block + zero-copy slab registry (M2 backend)

Implement coli_metal_moe_block: gate/up/silu/down for a whole expert block in ONE
command buffer, with GPU memory barriers between stages and BINDLESS gpuAddress
pointers so each expert is read zero-copy from its own RAM slab (exceeds Metal's
~31 buffer-binding limit). coli_metal_register/unregister wrap page-aligned slabs
via newBufferWithBytesNoCopy and resolve interior pointers to GPU addresses.
Per-row ragged expert routing supported; CPU does the final weighted scatter-add.
test_backend_metal validates decode + ragged blocks vs a CPU reference (nerr ~2e-6).
Still gated off in glm.c until the moe() wiring lands.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: wire batched moe_block into glm.c, token-exact (M2 integration)

moe() now dispatches each routed-expert block through the GPU in one command
buffer when COLI_METAL=1, reading expert weights zero-copy from page-aligned
RAM slabs (registered in expert_load). Falls back to CPU per-block on any
unresolved slab or GPU fault. Default build byte-identical (all #ifdef COLI_METAL).

Fixes a heap-corruption crash: expert_load registers slabs from parallel OpenMP
threads, so the slab registry is now mutex-guarded (buffer creation stays outside
the lock). Added command-buffer error checking (fall back to CPU on GPU fault)
and a COLI_METAL_DEBUG one-shot trace.

Validated token-exact vs the CPU path (greedy): identical 12-token output;
expert-matmul time 29.9s -> 21.1s with pinned experts still on CPU.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: instrument moe_block (GPU/CPU split, wall-vs-kernel time)

Add diagnostics printed on the PROFILO line under COLI_METAL: GPU vs CPU-fallback
block counts, experts-on-GPU, and a per-block time split (setup / gpu-wall /
kernel / scatter). Reveals that with a warm cache all experts run on the GPU
(0 fallback) and expert-matmul drops ~1.3x vs CPU, but ~62% of GPU wall-time is
idle/scheduling latency (3.1s kernel of 8.3s wall over 396 sporadic submits) —
the GPU powers down between blocks because attention runs on the CPU per layer.
Points the next optimization at keeping the GPU hot (offload attention).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: Metal backend measured results + next levers

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: Phase 2 fused decode attention plan + absorption-core validated

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: fused decode attention on GPU, token-exact (Phase 2)

coli_metal_attn_decode runs a full S=1 decode attention layer in ONE command
buffer: q_a -> rmsnorm -> q_b -> RoPE ; kv_a -> latent rmsnorm@pos + krot RoPE@pos
(cache write) ; MLA absorption core (qabs/score/softmax/clat/ctx) ; o_proj. The
absorption-core kernels were validated in isolation (nerr ~1e-6) before wiring.
Projection matmuls reuse the mm_gemv kernel; attention weights are uploaded+cached
(serial path, no lock); Lc/Rc caches are page-aligned + registered in kv_alloc for
zero-copy GPU read/write. GLM-5.2 dims compiled in; falls back to CPU for S>1
(prefill/MTP verify), st0!=0, active DSA selection (context>topk), or mismatched
dims. DSA index-key write stays on CPU so future selection still works.

Validated token-exact vs CPU (identical greedy output); attention time 16.5s ->
10.5s (~1.57x), end-to-end 0.20 -> 0.28 tok/s.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: Phase 2 fused attention complete + known limits

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: attention coverage/latency instrumentation + honest results

Add per-layer fused-attention counters (METAL-ATTN line): GPU layer count, gpu-wall
and true kernel time. Measurement (DRAFT=0, all-S=1 decode) shows the fused attention
triggers on all decode layers but is submit-latency-bound: gpu-wall 3.70s vs kernel
0.63s (83% idle latency over 546 sporadic command buffers). Attention time is neutral
vs CPU; the earlier MTP-on "16.5->10.5" was run-to-run variance. Design doc corrected
with the honest result: both offloads are gated by Metal's ~5ms cold-GPU submit
latency; reducing submit count (fuse attention+experts per layer) is the real lever.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: fused attention handles S<=4 (covers MTP verify forwards)

Extend coli_metal_attn_decode from S=1 to S<=4: the core kernels (qabs/score/
smax/clat/ctx) gain a query-row dimension with per-row causal masking (query s
attends keys [0, pos_base+s]); rmsnorm/rope/copy became row-aware; projections
run S rows via mm_gemv. This covers the default MTP config (draft=3 -> S=4 verify
forwards), which previously fell back to CPU attention entirely.

Token-exact vs CPU (identical greedy output, MTP on). Perf is inconclusive at
short context: still submit-latency-bound (attn gpu-wall 5.5s vs kernel 0.9s) and
the measurement is dominated by disk-streaming variance (+/-15s between runs).
Next: measure with a fully-warm cache to isolate compute, then reduce submit count.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: clean warm A/B shows real ~1.4x (experts+S<=4 attention), token-exact

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: interleave attention q/kv paths, 7->4 barriers (iter 2)

The q-path (q_a->rmsnorm->q_b->rope) and kv-path (kv_a->copy->rmsnorm+rope) are
independent until the absorption core, but were serialized by memory barriers.
Interleave them into 4 barrier-separated stages so the GPU overlaps independent
dispatches. Token-exact; attention gpu-wall 3.04s -> 2.73s (~10%).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: zero-copy attention weights + fuse shared expert into GPU block (iter 2)

Dense QT weights/scales now allocate page-aligned + registered (qalloc) under
METAL, so the fused attention reads q_a/q_b/kv_a/kv_b/o zero-copy instead of
uploading ~6 GB of duplicates (RSS -3 GB, upload copies gone). bind_gemv resolves
registered pointers (buffer,offset) with a pre-check guard.

Phase E's shared expert (identical shapes to a routed expert: gate/up [I,D],
down [D,I], same int4 container) is appended to the first Metal moe_block as an
extra expert with rw=1.0 over all S rows — removes 3 CPU matmuls per layer and
fills the same GPU submit. CPU Phase E still runs on any fallback.

Zero-copy validated token-exact: 35.1s -> 29.7s (0.34 tok/s) warm.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: iteration 2 findings

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: iter 2 final ~1.56x + iter 3 plan (disk/GPU overlap)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: overlap disk loads with GPU compute inside the layer (iter 3)

Split each MoE block into two GPU submits: the RESIDENT experts (pin/LRU hits,
plus the fused shared expert) are encoded and committed BEFORE the missed
experts' OMP pread loop, so the GPU computes while the disk reads; the missed
subset follows in a second (sync) submit once loaded. New two-phase backend API
(coli_metal_moe_block_begin/end) with handle-owned scratch so the async submit
cannot collide with the sync path's static buffers; moe_submit/moe_finish are
shared by both. Per-subset CPU fallback preserved (resident and missed fall back
independently on unresolved slab or GPU fault).

Token-exact. Warm 96GB: expert-matmul 8.96 -> 4.92s (resident compute now hidden
inside the disk window; expert idle latency ~5.7s -> ~0.9s), total 28.97s
(0.35 tok/s) vs CPU 50.2s = ~1.73x.

Note: 'make glm METAL=1' after a default build does NOT rebuild (target looks
up-to-date) — touch glm.c or clean when switching build flavors.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: iter 3 disk/GPU overlap results (~1.73x)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: keep-alive spinner experiment (env-gated) + latency decomposition

COLI_METAL_SPIN=1 keeps trivial GPU work in flight on a separate queue to probe
whether inter-submit idle is clock ramp-down; thread is detached (a joinable
global thread std::terminate'd the process at exit). First contended A/B was
inconclusive but showed the spinner does NOT collapse attention wall per-call
(~16ms both ways), so ramp-down is not the whole story. METAL-ATTN now decomposes
latency: cpu-sched (commit->kernelStart) vs gpu-sched (kernelStart->GPUStart) vs
kernel execution, to pinpoint where the ~13ms/call goes. Default behavior
unchanged (spinner off).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: standalone regression tests for fused decode attention

run_attn builds full-size fake GLM-5.2 attention weights (int4, page-aligned,
registered), replicates glm.c's absorb-branch math exactly on the CPU (q_a ->
rmsnorm -> q_b -> rope; kv_a -> latent rmsnorm + krot rope -> cache; per-head
qabs/score/softmax/clat/ctx; o_proj), and checks coli_metal_attn_decode against
it at S=1/3/4 and pos_base 0/12/37 — including the Lc/Rc cache write-back, which
end-to-end runs cannot isolate. All pass (nerr ~5e-6, cache ~1.4e-5). The whole
Metal path (gemv, moe_block, fused attention) is now testable without the model.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: route large row-batch matmul_qt GEMMs to the GPU (prefill)

matmul_qt now dispatches to a new coli_metal_gemm when S >= COLI_METAL_GEMM_MIN
(default 16), the weight is int8/int4 and registered (all dense QT allocs are,
via qalloc), and we're not inside an OpenMP region (mirrors the CUDA guard).
Decode-sized matmuls stay on the CPU where NEON wins vs submit latency; prefill's
big GEMMs (kv_b reconstruction at S=Tk, o_proj, dense MLP, step_all's S x vocab
logits) amortize it — microbench showed ~6x over the CPU idot at S=16.
Standalone test: registered int4 GEMM S=64 vs cpu_ref (nerr 2.9e-6).
Machine busy again; end-to-end token-exactness + threshold sweep pending idle.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* README: document the experimental Metal backend (Apple Silicon)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: 1.5-2.1x faster moe_gemv (simdgroup-per-row + 8-value loads)

Replace one-threadgroup-per-output-row (128 threads reducing via threadgroup
memory) with one SIMDGROUP per output row, 4 rows per threadgroup, and uchar4
loads (8 nibbles / 8 int8 per lane-iteration). Removes the threadgroup barrier
+ shared-memory reduction entirely (simd_sum only) and doubles load width.
Engine-like block-shape microbench (pure GPU time): S=4 block 2548->1739us,
S=1 block 934->437us, big block 4582->3414us — 358-389 GB/s vs 182-264.
Row-bound guard added (NT) since the grid rounds up to 4 rows/TG.
All backend tests pass (moe_block nerr 2.4e-6, attention unchanged).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: mm_gemv simdgroup-per-row + 8-value loads (attention projections, prefill GEMM)

Apply the moe_gemv V2 transformation to the general quantized GEMV: one simdgroup
per output element (4/threadgroup), 8-value loads for i8/i4/f32, no threadgroup
reduction. Same measured 1.5-2.1x class of win; serves the fused-attention
projections (q_a/q_b/kv_a/o), coli_metal_gemm (prefill), and coli_metal_matmul.
All three dispatch sites updated (NT row-bound guard, grid ceil(NT/4) x 128).
Full test suite green, incl. non-mult-of-8 tail paths (2050x6146) and all fmts.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: experimental COLI_MMAP=1 — experts as zero-copy views into mmap'd files

Lazily mmap each safetensors file (PROT_READ, MAP_SHARED, mutex-guarded — expert
loads are OMP-parallel), register the mapping with Metal, and make expert_load a
pointer assignment into the map: no pread, no slab, no copy; the OS page cache is
the cache. Alignment guards fall back to the slab path. Default OFF.

First validation (machine at load 66 + 46GB swap): token-exact, RSS 58 -> 10.5 GB
as designed, but GPU wall exploded (~130 MB/s effective) — the GPU demand-faults
file-backed pages, catastrophic when memory pressure evicts them. Needs an
idle-machine A/B to judge fairly (llama.cpp's identical technique relies on pages
staying resident); possible fixes if slow even idle: CPU pre-touch of missed
experts' pages before the GPU submit, or madvise/mlock windows.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: CPU pre-touch for COLI_MMAP expert pages (fix GPU demand-faulting)

In mmap mode, fault the missed expert's pages in on the CPU inside expert_load
(madvise WILLNEED for async readahead + a page-stride touch): this is pread's I/O
without the copy and without the slab, it runs inside the existing OMP loop that
overlaps with the resident-experts GPU submit (iter 3), and it guarantees the GPU
only ever reads resident pages — GPU demand-faulting of file-backed pages
measured catastrophic (~130 MB/s). Read-only addition: outputs unchanged from the
validated mmap run; perf pending the idle-machine A/B.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: idle-machine suite results (~1.33x same-session; mmap negative result)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: COLI_METAL_UNTRACKED experiment (negative result, default off)

Env-gated MTLResourceHazardTrackingModeUntracked on registered wraps + scratch to
test whether cross-CB hazard tracking causes the ~10ms/CB gpu-sched delay. Idle
A/B: no effect (gpu-sched 3.9 vs 3.4s, noise), token-exact. Together with the
spinner negative, this pins the attention CB delay as inherent scheduler/wake
overhead on an empty pipeline — removable only by eliminating the CB boundary,
which CPU-side routing at ~58% hit-rate forces. Metal side is at its floor:
kernel 3.5s+0.8s (near BW ceiling), sched ~3.2s, disk ~15s dominant (10 tok).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: loop conclusion — best config DIRECT=1+COLI_METAL=1, 0.42 tok/s (~1.4x)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: refactor attention into encode_attention()+resolve_attn() (layer-CB prep)

Behavior-preserving: attn_decode is now a thin wrapper; all attention tests
byte-identical. Prepares embedding the chain in a full-layer command buffer.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* metal: full decode layer in ONE command buffer (token-exact)

coli_metal_layer_decode runs the whole layer prelude on the GPU in a single
submit: in_ln rmsnorm -> fused attention -> residual add -> post_ln rmsnorm ->
shared expert (gate/up/silu/down) -> router (f32 simdgroup matvec + sigmoid) ->
exact phase-A top-K selection (greedy argmax over sigmoid+bias with CPU tie
order, --topp truncation, norm_topk, routed_scale) in a serial-per-row kernel.
The CPU's per-layer work shrinks to: read 8 expert IDs, resolve/load, expert CBs
(disk/GPU overlap unchanged), scatter. moe() consumes the precomputed routing
(g_pre_*: skips phase A, keeps eusage/eheat/ereq counters for the learning
cache) and adds the GPU shared-expert output instead of computing phase E.
ld() tensors (norms/router/bias) now allocate registered so the GPU reads them
zero-copy. DSA index keys still computed on CPU from the in_ln-normed x (new
inrm output). Every missing condition falls back to the full CPU layer.

Validated token-exact vs CPU (identical greedy output, MTP on). Profile:
"altro" 3.8s -> 0.53s (12 tok); 0.42 tok/s despite disk-variance headwind.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: Phase 3 full-layer CB results — 0.43 tok/s record, token-exact

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* gitignore: Metal build artifacts, venv, bench datasets

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* remove internal design docs before PR

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-13 08:57:10 +02:00
..