I shared about this in AiOS Meetup yesterday. Adding here for archive purpose.
Draw Things is media generation software running on Apple Silicon. Media generation uses diffusion models, and the process is, in some ways, similar to LLM prefilling: both involve a large amount of parallel computation.
Thus, our initial bet was that by leveraging our know-how in diffusion engines, we could achieve amazing LLM prefill performance and deliver a nice little shock to the community enthusiastic about local AI and Macs.
Of course, only after we were halfway through this journey did we start to understand why it was hard. In the end, we shipped Local Code, which still has amazing prefill and decode performance, in addition to having the complete harness and environment built in. Sharing these learnings with the community might help us build on each other’s work and ultimately come up with better engineering ideas to overcome these limitations.
1. Prefill with SSD Streaming Is Hard
Draw Things was one of the first software to implement weight streaming (we call it JIT weights) a few years back. It served us well (until M6, which is awkward because the chip is fast and the RAM is limited).
Since prefill is generally considered a compute-bound task, it was natural to think that offloading weights to SSD should be “free”. Well, it sorta is, if you always have a very long prompt to prefill.
Let’s go over some numbers. On M5 Ultra, we usually reach ~4,800 tokens/s for prompt processing (prefill). For M5 Max, we would expect ~2,400 tokens/s if it is truly compute-bound. That also means moving the full 160 GB of quantized weights into memory for a single prefill pass. Assuming 10 GB/s of effective SSD bandwidth, we need ~38k tokens per pass just to get back to balance.
With our usual ~8k prompt processing between turns, we are looking at a 4–5× slowdown even with perfect overlap between SSD reads and computation, assuming the full weights must be loaded each time.
There is no easy way to fix this, as we are fundamentally limited by SSD speed. Sharding across multiple SSDs might help, but that is also not what normal people have set up. Adaptive prefill chunking could help alleviate this a little, although we need to be careful not to reload the same weights multiple times.
One thing we can address is finding the sweet spot where loading only the selected experts stops being beneficial and loading the full layer becomes preferable.
We tested this on DeepSeek V4 Flash, with 256 routed experts per layer. The table below compares selective expert loading against asynchronous full-layer loading, with a 16 GiB configuration and initially empty resident slots. We also counted how many distinct experts were activated at least once in each chunk, averaged across all 43 MoE layers.
Interestingly, at around 3k tokens, we have already activated roughly 75% of the experts in each layer. Selective loading and full-layer loading become almost equally fast, and beyond that, full-layer loading starts to win.
But this crossover also depends on how many experts are already resident in memory. With a warm 48 GiB configuration, at the same 3,072-token chunk size, the 40 layers that needed streaming selected roughly 190 experts per layer. However, about 90 of them were already resident, leaving only 100 to load from SSD. In that case, selective loading still wins.
So rather than using a fixed prompt-length threshold, we can make this decision based on the number of selected experts that actually need to be loaded. This is one optimization opportunity we are still exploring.
2. The Sensitivity to Numerical Accuracy (in Attention)
One thing that took me a while to realize is that both diffusion models and autoregressive models have some self-correcting mechanisms. For diffusion models, each step takes the previous prediction as input and has an opportunity to correct the trajectory. For autoregressive models, reasoning is another powerful way to correct earlier mistakes.
The difference, though, is in the attention operation.
Diffusion models often operate at shorter attention sequence lengths (unless it is a video model). These shorter sequences tend to be more forgiving of numerical accuracy losses in low-precision attention (FP8 or INT8, as in SageAttention-style implementations). Autoregressive models can run on very long sequences, and numerical errors in attention can become more consequential at longer context lengths.
Techniques we use, such as Metal Quantized Attention, therefore need to go through more scrutiny before they can be useful for LLMs. Luckily, benchmarks such as LongBench v2 are pretty effective at catching these issues.
Here is what we observed on the original 60-case subset:
Our INT8 QK + FP16 PV implementation came very close to Splash BF16, while the fully INT8 implementations showed more noticeable accuracy losses.
Of course, 60 cases are not enough to establish statistically significant differences, particularly between 39/60 and 40/60. But they are useful for identifying potential regressions that deserve further investigation.
3. When Kernel Specialization Stops Paying Off
Without even thinking about it much, for Draw Things, all of our kernels are specialized by shape. What we mean is that the matrix dimensions—or, for attention, the row and column counts—are baked into the shader as function constants.
Later, we added the ability to disable specialization by M (as in M × K × N for matrix multiplication) for MoE models (the HiDream series). Besides that, our buffer bindings are really just tensor bytes.
For LLM prefill, this is no longer profitable in many cases. Sequence lengths are irregular, and even with chunked prefill, that is still the case. So we have to use marginally less efficient kernels that read dimensions at runtime.
How much performance do we actually lose?
For larger GEMM operations and causal prefill attention, the overhead is surprisingly small, usually just 1–2%. The real challenge is with smaller shapes, particularly FP16 attention with only a few query rows, where we can see more than 25% slowdown.
With this data, we rewrote a large part of kernels involved in LLM prefill path to support runtime shapes, and only use shape specialization for chunked prefill, disabling shape specialization for tails.