Windowed-MTP: Removing the Full-Context Draft-KV Tax at Million-Token Context
Alagappan Valliappan
Abstract
Frontier LLMs moved to hybrid and linear attention to curb long-context KV cost, but the built-in Multi-Token-Prediction (MTP/NEXTN) draft heads they shipped with largely kept full attention, and the KV cache is read on every draft forward. That co-design gap inverts the very optimization it was meant to deliver. One draft layer issues $\gamma$ full-context reads per decode step while a hybrid target issues one per full-attention layer, and only 9-25% of its layers are full attention, so the two read counts become comparable. At long context those reads are what set step time. Measured at $\gamma = 6$, the entire draft phase adds +92% to +138% to per-decode-step latency on top of the bare verify, and on hard, low-acceptance tasks a deep native draft turns net-negative: slower than no speculation at all. We give a per-decode-step cost model that predicts this inversion, then close the gap at the serving layer with Windowed-MTP: a StreamingLLM-style window plus attention sink on the draft's attention only, training-free, drop-in, and lossless by construction, since the draft only proposes and the target still decides the output. Across three architecture families at 1M on a single B200 in SGLang it cuts per-decode-step cost by +28% to +44% (input-invariant, widening from 261K to 1M) and turns the 7.7-11.1% of total KV held by the now-dead draft pool into reclaimable serving capacity.
Chat is not available.
Successful Page Load