Blog post8 min read

Porting FlashPrefill V2 to Apple silicon for faster prefill

A visual explanation of block-sparse prefill attention, followed by what happened when we built our own FlashPrefill V2 kernel for Apple silicon in MLX, Apple's machine-learning framework.

Long prompts make LLM inference expensive before the model generates a single token. During this prefill phase, every token looks back at the tokens before it. Double the prompt length and the attention matrix contains roughly four times as many token pairs.

FlashPrefill V2, by Fan et al. at WeChat, starts from a useful observation: most of that matrix is quiet. Attention concentrates around the local diagonal and a limited number of distant regions. If a cheap preview can find those regions, the full attention calculation only has to run there.

The authors implemented FlashPrefill V2 on an NVIDIA H20, while we were curious about how the same idea would behave on the much smaller GPU in an Apple M-series chip. We expect edge devices to matter more over time, and Apple's M-series processors present an interesting constraint: high memory bandwidth, but relatively little compute compared with discrete consumer GPUs. That imbalance makes avoiding unnecessary work especially valuable.

The short version: summarize each key block, use those summaries to decide which blocks deserve full attention, and keep a cheap correction for the blocks that do not.

Part I · How it works

Attention is dense. Its useful structure is not.

An attention matrix places queries on one axis and keys on the other. Each cell asks how much one token should use another token when forming its next representation. Causal language models only look backward, which gives the matrix its triangular shape.

Meaning appears as bright points

“Foolishness” reaches back to “worst” and “wisdom”. Sparse attention must preserve links like these, not merely keep a fixed local window.

GPUs compute the matrix in tiles

A useful block and a nearly empty block normally cost the same dense matrix multiplication. That makes quiet regions expensive.

Compress keys before the heavy work

Average the keys in a block, then compare every query against that one pooled key. This toy block falls from 36 dot products to six.

A cheap map of where the attention lives

One pooled score per query row per block, local stays exact. The map already knows where the light is, before any heavy math. At runtime this map is all the kernel gets to see.

The algorithm starts here

Everything before this was dense attention and pooling. The blocks worth exact math are those above a threshold, plus the local diagonal. More context, more sparsity.

Compute the peaks, summarize the rest

The selected blocks are recomputed exactly. Every other block keeps its pooled row from the map, so the softmax does not treat skipped context as empty.

The algorithm therefore has three pieces: build the pooled index, run full attention over selected blocks, and apply the pooled correction. One number to keep straight: the toy above keeps blocks above 20% of the row’s hottest block (alpha = 0.2) so that a 36-token example has something left to prune; the kernel runs at alpha = 0.1 end to end. The per-layer figure below sets alpha per prompt length so that its synthetic prompt skips the same share of blocks the real runs did. The interesting systems question is whether those savings survive a real implementation. We ported FlashPrefill V2 to Apple silicon by writing our own Metal kernel in MLX and have released the implementation.

Part II · What happened on MLX

Porting the idea to a much smaller GPU

FlashPrefill V2 is designed around NVIDIA Hopper. Our implementation targets Metal through MLX, where the hardware budget and execution model are different. A Hopper thread block can work on a 128 × 64 score tile with asynchronous, double-buffered loads. On our MLX path, the production tile is 32 × 16 and loading is sequential. The algorithm transfers; the kernel architecture does not transfer unchanged.

Hardware
Apple M4 Pro, 16-core GPU, 24 GB, for both the kernel benchmarks and the end-to-end runs
Software
MLX 0.32.2, custom Metal kernel
Model shape
Qwen3-4B: 32 query heads, 8 KV heads, head dimension 128
Kernel benchmark
BF16, clustered synthetic Q/K/V, alpha per length matched to the end-to-end block sparsity, best of 5 runs after warm-up
End to end
mlx-lm Qwen3-4B-Thinking-2507-8bit, the four NIAH prompts in the repository, mean of 3 runs, M4 Pro

The crossover starts around 4k tokens

At short lengths, building the index and correcting pruned blocks cost about as much as they save. At 2k tokens the sparse path is effectively tied with stock MLX attention. By 4k, 38% of the blocks are skipped and the balance changes. At 16k, the complete attention path is 2.9× faster; at 32k, 4.3×; at 64k, 5.6×.

Attention time per layer as the prompt grows
0 2,000 4,000 6,000 8,000 ms per layer 7 ms 1.0× 2k 16% sparsity 27 ms 1.4× 4k 38% sparsity 104 ms 1.9× 8k 56% sparsity 413 ms 2.9× 16k 73% sparsity 1,671 ms 4.3× 32k 83% sparsity 6,766 ms 5.6× 64k 89% sparsity stock MLX attention FPV2 main loop, kept blocks only FPV2 mean correction FPV2 index stage
One attention layer, Qwen3-4B shapes, alpha set per length so the synthetic prompt skips the same share of blocks the NIAH prompts did end to end, measured on an M4 Pro. The stock kernel grows with the square of the prompt; FPV2 grows with the kept blocks.

At a 16k-token prompt with alpha set to zero, so that nothing is pruned, the sparse kernel takes 425 ms against 413 ms for stock MLX. The gain therefore comes from sparsity, not a faster baseline kernel.

At 16k, only 28% of visible blocks run in full

The next figure is not an illustration. It is the selection map produced by one KV head during a 16k-token prefill. Coloured blocks run full attention, shaded by how many of the eight KV heads keep them; the grey sinks, local window and dense tail are kept by rule. Dark blocks are represented by pooled rows, so most of the older context is pruned.

Where the kept blocks are
0 0 4k 4k 8k 8k 12k 12k 16k 16k key position (128-token blocks) query position (32-position tiles) sinks: the first 256 tokens, always kept local window and diagonal, always kept pruned: represented by one pooled row each last two tiles run dense One 16k-token prefill, one KV head alpha = 0.1, 512 query tiles × 128 key blocks kept by score, attended in full shade: 1 to 8 of 8 KV heads keep it kept by rule: sinks, local window, dense tail pruned, folded in as one pooled row not visible (future tokens) 28% of visible blocks kept 128× cheaper per pruned block, in FLOPs Kernel time, measured 95% walking kept blocks 1% the correction loop existing 4% the pooled-row passes
Each row is a 32-position query tile choosing which 128-token key blocks to attend. Pruned blocks are three quarters of the triangle but cost one matmul pass per sixteen.

A 2.9× attention path does not mean 2.9× prefill

FPV2 changes attention and nothing else. QKV projections, output projections, normalization, and the MLP still take the same time. At 16k tokens, attention represents about 40% of Qwen3-4B's prefill FLOPs. Even if attention became free, the theoretical end-to-end speedup would be only 1.7×.

Why a faster attention kernel is a smaller prefill speedup
1k 4k 16k 64k 256k 0% 50% 100% everything else projections, MLP, norms attention the only part FPV2 touches 25% 40% 57% 84% prompt length share of prefill FLOPs 1× 2× 3× 4× 5× 6× 7× 8k 16k 32k 64k 128k 256k 1.9× 2.9× 4.3× 5.6× 1.12× 1.36× 1.78× 2.58× attention kernel alone, M4 Pro whole prefill measured, M4 Pro ceiling (if attention was instant) prompt length speedup over stock
Qwen3-4B. At 16k, attention is 40% of the FLOPs, so a 2.9× attention kernel can at most give 1.7× end to end. Kernel and whole-prefill series are both measured on an M4 Pro. Past 64k the ceiling climbs steeply, which is the regime the paper reports.

The measured whole-prefill gains on an M4 Pro were 1.12× at 8k, 1.36× at 16k, 1.78× at 32k, and 2.58× at 64k; the 8k to 64k runs saw 56%, 73%, 84%, and 89% block sparsity, respectively. In wall-clock terms a 32k-token prefill falls from 105 s to 59 s, and a 64k prefill from 347 s to 134 s. That gap between kernel and application speedup is expected, not a failure of the sparse kernel. It is also why long context matters: attention's quadratic cost becomes a larger share of total work as the prompt grows. For this model it rises from 25% at 8k to an estimated 84% at 128k, where the theoretical ceiling passes 6×.

For small GPUs the result is encouraging: on an M4 Pro, a chip with a fraction of an H20’s compute, block-sparse prefill takes a 64k-token prompt from nearly six minutes to just over two, and the algorithm needed no change to get there, only the kernel did. Our measurements focus on implementation performance; the FlashPrefill V2 paper reports accuracy across various benchmarks.

← All posts sevren.ai