Prefix Sliding for efficient test-time scaling
Niklas Muennighoff, Zhengyang Wang, Zeyi Chen, Weijia Shi, Binyuan Hui, John Yang, Dapeng Jiang, Mika Senghaas, Fares Obeid, Johannes Hagemann, Sami Jaghouar, Ludwig Schmidt, Percy Liang, Jason Wei, Andrew Y. Ng, Luke Zettlemoyer, Yejin Choi, Mike Lewis
cs.CL, cs.AI, cs.LG
2026-08-27
On Qwen3-1.7B, Prefix Sliding keeps the prompt prefix plus a tail window. No training: ~3x faster. An 8192 window hits 35.8 AIME25 vs 34.2 full attention.
The usual way to spend extra test-time compute is to let a language model think longer. Full attention makes that expensive: every new token attends to the entire reasoning trace, so cost grows linearly with length. Hard problems that need tens of thousands of tokens hit memory and latency walls. Long traces also distract the model, poison the context, loop, and bury earlier conclusions.
A toy example makes the waste obvious. Once ((42 + 84) × 4) - 5 has finished the addition, the scratch work for 42 + 84 is dead weight; only the intermediate result matters. Attention maps on a Qwen3-1.7B AIME25 trace agree: the prefix (system prompt, the problem, the <think> delimiter) and roughly the last 1,000 tokens soak up most of the mass. The middle of the trace barely gets looked at.
Prefix Sliding keeps two slices of the sequence and drops the rest.
A 100-token prefix plus a 4,096-token window caps memory at 4,196 tokens. After that, each new token costs the same whether the model has already generated a million tokens or a billion.
Position encodings default to Continue PE: IDs keep climbing, so RoPE-applied KV entries stay valid. Reset PE reassigns positions and recomputes those entries. The two are close on AIME25, so the cheaper option won.
For RL, they use truncated backprop. A sliding window's theoretical receptive field is W×L layers; in practice it is about 1.5×W. On a 100k-token rollout with W=2,048, the trainer receives the last 8,192 tokens and applies the RL loss only to the final 2,048. A 7B async RL check (window 8,192, 4× context) matched full attention with full backprop on AIME24. The FlashAttention kernel filters at two levels: elementwise masks inside tiles that straddle the allowed region, and skipped tiles that sit entirely outside the prefix or the window.
Main numbers are Qwen3-1.7B, avg@64, temperature 0.6. Training-free, window 4,096: AIME25 33.9, GPQA 37.0, MATH500 91.5, against full attention at 34.2 / 37.6 / 91.7. Window 8,192 actually beats full attention on AIME25 (35.8) and GPQA (38.0). Throughput at 128k tokens: 5,224 tok/s with a 4,096 window versus 448 for full attention. At 32k the gap is 5,479 vs 1,477, about 3.7×. The abstract's "3× faster" is wall-clock thinking time in Figure 1: more tokens in the same budget, not better tokens.
On one H100 the custom kernel roughly matches a vanilla sliding-window kernel and settles near 5,000 tok/s after warmup. Full attention keeps slowing down.
Under a matched 8,192 memory budget, RL with Prefix Sliding rolls past 100k tokens while full attention is stuck at 8k, and reward is higher. Pure sliding windows forget the problem. Last-k and summary restarts reprocess tokens, sawtooth the memory curve, and add extra hyperparameters. On AIME25 with a 262,144 max generation and a 4,096 local window, Prefix Sliding has the best accuracy-vs-time curve, and it adds only one knob: window size.
This is an attention mask on an existing reasoner. No retraining, constant marginal cost once the window is full. Contest math and long-horizon agents can try 4,096 or 8,192 first. On the training side, you no longer have to discard overlong generations.
It is not a new architecture. Linear models such as Mamba and RWKV need a retrain; Prefix Sliding wraps current Transformer weights. StreamingLLM keeps about four sink tokens. A tool spec can sit unused for thousands of steps and still be required later, so the whole prefix has to stay pinned. H2O keeps heavy hitters by past attention (backward-looking); Prefix Sliding is forward-looking. Combining them is plausible. The paper could not compare fairly because H2O does not plug into vLLM or FlashAttention.
Do not expect a win on short tasks. HealthBench averages 2,086 tokens; with a 2,048 window the model almost never slides, so speed matches full attention.
LiveCodeBench needs a window of at least 16,384 to match full attention. The model drafts a function, then thinks in comments for thousands of tokens; by the time it writes code again, the function header has slid out. That evaluation is training-free. RL might teach it to stop rambling in comments. This paper does not show that.
Tool output and multi-turn chats are real failure modes. A dumped webpage can overflow the window, so the model cannot even see the full page, and it may flush earlier instructions. Whether later user turns should be nailed into the prefix or left to slide off is left open. Very long prefixes also keep their full prefill KV; this method does not shrink that.
Headline numbers are 1.7B. Truncated backprop was only checked briefly at 7B. There is no matched-compute comparison to RNNs, SSMs, or hybrid sliding-window models. Last-k and summary hyperparameters in Figure 9 were swept in a smaller-context setup and may not be optimal at 262k generation length.