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
MPP is opt-in, and should be. bf16 operands fail llama.cpp's own SSM_SCAN op test at its 2e-7 fp32 bar — 7 of 13 cases, against 13/13 for the fp32 path — so shipping it as the default would mean failing the host project's test suite. GGML_HELIX_MPP=1 takes it, and the bridge names the kernel it chose so the choice is never inferred from throughput.
Neural accelerators, on a shape that used to decline
Pass C is compiled once per supported d_state, because matmul2d_descriptor needs the reduction extent as a constant expression and specialising on it is what reaches 65.68 TFLOP/s. Until recently only 128 was compiled, so Falcon-H1 at 256 fell back to the fp32 path — silently, which is the danger: the fallback is safe but invisible, and throughput alone cannot tell you which ran. Both extents are now compiled.
| fp32 | MPP | ||
|---|---|---|---|
| SSM_SCAN alone | 3.733 ms | 2.250 ms | 1.66x |
| end-to-end prefill, pp2048 | 1922.8 t/s | 2078.6 t/s | 1.08x |
Threadgroup memory did not change: the staged tiles are sized by chunk and head-dim slab, not by d_state, so 256 costs the same ~17 KiB as 128 and stays well inside the 32 KiB limit.
The bf16 operands cost nothing measurable in quality. Wikitext-2 over 40 chunks puts the MPP path at 10.8018 against 10.8016 for fp32 — a delta of 0.001 standard errors — which is the interesting part, because that same kernel fails llama.cpp's op-level test at 2e-7. The op tolerance is a numerical bound, not a model-quality threshold.
- 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.