Prefetch-steered routing speeds up large ai model decoding by managing cache use
Prefetch-Steered Routing: Deterministic Cache-Aware Expert Selection for Offloaded Mixture-of-Experts Decoding
Artificial Intelligence
Summary
Large AI models called Mixture-of-Experts (MoE) need to use parts stored in separate memory, which slows down processing when switching between them. The authors developed a way to predict and prepare the needed parts in a fixed, logical cache to reduce copying delays and make the model’s responses more consistent. Their method speeds up decoding by over twice as fast compared to traditional approaches while keeping the quality of results similar. They tested this on several big models and found it works better especially when limited memory is available.
What this means in practice
- •For machine learning engineers: Reduce latency when running big Mixture-of-Experts language models on limited GPU memory using prefetch-steered routing to improve throughput and consistency.
- •For cloud gpu operators: Optimize resource usage and speed for GPU hosts running offloaded MoE models by implementing deterministic cache management to control memory transfers.
Authors
Siming Huang
Abstract
When a Mixture-of-Experts (MoE) model does not fit in GPU memory, batch-1 decoding is dominated by copying expert weights from host memory. Cache-aware routing reduces these copies by biasing the router toward resident experts, but combined with asynchronous prefetching it makes the model's outputs depend on transfer timing, so a harness that does not simulate that timing no longer evaluates the deployed model. We propose prefetch-steered routing (PSR): routing conditions on a deterministic logical cache, into which at most n experts predicted by the next layer's router are inserted per layer and token; physical copies affect only latency. On OLMoE-1B-7B with 50% and 25% of experts resident on an RTX 4090, PSR decodes 2.2-2.3x faster than exact routing with LRU and 25-30% faster than the Cache-Prior operating point, with mean benchmark differences of +0.1 and +0.5 points (95% lower bounds -0.8 and -0.3). In greedy generation its lead over Cache-Prior depends on the workload and reverses on HumanEval at 50% residency. On Moonlight-16B-A3B, PSR is 1.9x faster than exact routing with LRU at 37.5% residency and, at the latency of Cache-Prior, significantly better on three of eight benchmarks. On Qwen3.6-35B-A3B, where copies are cheap relative to computation, PSR keeps its advantage over Cache-Prior on the likelihood benchmarks but not its speed. Distillation-trained gates did not help. Speculative decoding with block-level expert selection was slower than PSR with a 1B autoregressive drafter on OLMoE, but with a block drafter on Qwen3.6 it was 2.1-2.9x faster than any autoregressive configuration with at most its GPU memory on the 64 prompts we timed.