REMIT starts in shadow mode: two weeks, on your own servers, with your own data, zero actions taken.See the pilot
Products

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.

Diagram of the chunked associative scan: four chunks processed in parallel, each with a lower-triangular intra-chunk mask, and a scan line carrying the state summary between them.
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

measurementresult
SSM_SCAN vs upstream fp32 simdgroup_matrix (L=512)2.57x
SSM_SCAN vs the single-pass kernel7.45x
end-to-end prefill, 2048 tokens1428.6 -> 1941.8 tok/s (1.36x)
decodeunchanged, by design
M5 Max, GPU-timer based, min of N.

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.

fp32MPP
SSM_SCAN alone3.733 ms2.250 ms1.66x
end-to-end prefill, pp20481922.8 t/s2078.6 t/s1.08x
M5 Max, Falcon-H1-7B at d_state=256.

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.