HELIX
Metal prefill kernel for Mamba-2 hybrids
State-space models compute a sequential recurrence that defeats most mobile compilers. HELIX reassociates it into a chunked associative scan, so the work lands on Apple's matrix hardware instead of a serial loop, and keeps the carried state in threadgroup memory instead of spilling it to device memory on every chunk. It drops into llama.cpp behind one interception point and declines any shape it cannot serve, so it can never be a regression.
- 2.57x
- over upstream fp32, L=512
- 7.45x
- over the single-pass kernel
- 3
- shapes upstream declines
What it does
A state-space layer computes h_t = A_t h_{t-1} + B_t x_t. The dependency between steps is what makes it cheap in theory and awkward in practice: it is a scan, not a matmul, and scans do not use the matrix units that dominate a modern GPU's throughput. The chunked associative scan breaks the sequence into fixed chunks, expresses everything inside a chunk as dense matrix products, and carries a compact summary between chunks. HELIX splits that into three passes so the per-chunk work runs fully in parallel rather than one threadgroup walking chunks in sequence.
The second decision matters as much as the first. At d_state 128 and head_dim 64 the carried state is 32 KiB in fp32 — larger than the threadgroup memory budget, which is why the upstream kernel writes it back to device memory every chunk and pays two barriers for the privilege. Holding it in bf16 halves it to 16 KiB and it becomes resident for the whole walk. The precision decision and the memory-locality decision are the same decision.
Measured
| measurement | result |
|---|---|
| SSM_SCAN vs upstream fp32 simdgroup_matrix (L=512) | 2.57x |
| SSM_SCAN vs the single-pass kernel | 7.45x |
| end-to-end prefill, 2048 tokens | 1428.6 -> 1941.8 tok/s (1.36x) |
| decode | unchanged, by design |
The gap between 2.57x on the kernel and 1.36x end-to-end is Amdahl's law, and it is worth stating plainly rather than quoting only the first number. Falcon-H1 is a hybrid: the SSM scan is a minority of the computation, and ternary dequantization across the rest of the network dominates prefill.
Shapes upstream declines
llama.cpp already ships a chunked SSD Metal kernel, so "SSMs fall back to CPU on Apple Silicon" is not true and HELIX is not pitched that way. What that kernel does do is gate itself narrowly, falling back to a scalar path outside its supported shape. HELIX covers three cases it declines:
- head_dim 128 — models such as Falcon-H1 and Nemotron-H fall back entirely; HELIX tiles the channel dimension across threadgroups.
- Diagonal A — Mamba-1 style per-channel decay, rather than one scalar per head.
- Ragged tails — a final chunk shorter than the chunk size is handled by a zero-padded copy rather than a scalar tail kernel, which measured 6x slower than upstream at L=96.
Honest limits
The MPP path — the one benchmarked at 4.62x on the M5's neural accelerators — has never engaged for Falcon-H1 on any hardware. It is compiled for d_state = 128 and Falcon-H1 uses 256, so the runtime falls back silently to the fp32 path. Supporting it needs dynamic extents.
- Prefill only. Decode is bandwidth-bound and deliberately untouched, which means HELIX does not move the number a chat user actually feels.
- On an A16 iPhone the end-to-end prefill difference could not be resolved above run-to-run noise: three HELIX-on samples spanned 38.4 to 51.2 tok/s against a between-arm gap of roughly 10%.
- Gated at 64 tokens. Below that the interception declines and ggml's own kernel runs, because the chunk-parallel path costs more than it saves on short sequences.