Windowed-MTP: Removing the Full-Context Draft-KV Tax at Million-Token Context

Background Speculative decoding accelerates autoregressive LLM decoding by having a cheap draft model propose tokens that an expensive target model verifies in one batched forward pass. Recent frontier models increasingly ship a built-in Multi-Token-Prediction (MTP/NEXTN) draft head instead of a separate small drafter, designed under the implicit assumption that this draft head is essentially free compared to the target's verification pass.

배경 스펙 디코딩(speculative decoding)은 저렴한 draft 모델이 토큰을 제안하고 비싼 target 모델이 한 번의 배치 forward로 이를 검증하는 방식으로 자기회귀 LLM 디코딩을 가속한다. 최근 프런티어 모델들은 별도 draft 모델 대신 자체 내장 Multi-Token-Prediction(MTP/NEXTN) draft head를 함께 출시하는 추세이며, 이 draft head가 target의 검증 패스에 비해 사실상 공짜라는 암묵적 가정 위에서 설계돼 있다.

Problem At million-token context, that near-free assumption breaks. The built-in MTP head re-reads the entire KV cache with full attention at every draft step, so its cost grows linearly (O(S)) with context length, and on hard, low-acceptance workloads speculative decoding can end up slower than plain dense decoding. As target models shift toward hybrid/linear attention (Mamba2, GDN) to cut their own KV cost, the verify step gets cheaper while the draft step's burden stays the same, making the problem even more pronounced.

문제 백만 토큰급 컨텍스트에서는 이 '거의 공짜' 가정이 깨진다. 내장 MTP head는 매 draft step마다 전체 KV 캐시를 full attention으로 다시 읽으므로 비용이 컨텍스트 길이에 선형(O(S))으로 증가하고, 수용률이 낮은 어려운 워크로드에서는 스펙 디코딩이 아예 dense 디코딩보다 느려지는 역전까지 발생한다. target 모델들이 자체 KV 비용을 줄이려 hybrid/linear attention(Mamba2, GDN)으로 옮겨갈수록 검증 단계는 저렴해지는데 draft 단계의 부담은 그대로 남아, 문제가 오히려 더 두드러진다.

Prior limitations Prior long-context speculative decoding work (MagicDec) diagnosed a similar KV-bandwidth bottleneck but used a separately fixed-size-KV draft; LongSpec trains a dedicated bounded-KV drafter, and SpecExtend retrieves relevant keys across models -- none of them directly patches the built-in MTP/NEXTN head that frontier models already ship. Existing SGLang windowing flags are also wired only to separately trained draft paths, leaving the built-in MTP head's O(S) tax unaddressed.

기존 한계 이전 장문맥 스펙 디코딩 연구(MagicDec)는 비슷한 KV-대역폭 병목을 진단했지만 별도의 고정 크기 KV draft를 사용했고, LongSpec은 전용 bounded-KV drafter를 학습시키며 SpecExtend는 모델 간 검색으로 관련 key를 가져온다 -- 어느 것도 프런티어 모델이 이미 내장한 MTP/NEXTN head를 직접 손보지 않는다. 기존 SGLang의 windowing 플래그도 별도로 학습된 draft 경로에만 연결돼 있어, 내장 MTP head의 O(S) 비용은 그대로 방치돼 있었다.

Goal Remove the draft-side long-context KV tax from already-shipped MTP/NEXTN heads without retraining, without touching the target's output distribution, and without adding a new kernel -- a purely inference-time, drop-in fix.

목표 재학습 없이, target의 출력 분포를 건드리지 않고, 새 커널도 추가하지 않으면서 -- 이미 출시된 MTP/NEXTN head의 장문맥 draft-side KV 비용을 순수 추론 시점 개입만으로 제거하는 것이 목표다.

Method Windowed-MTP applies a StreamingLLM-style sliding window plus a handful of sink tokens to the draft's attention only, implemented by shrinking the draft's paged-KV block table to [sink blocks] + [recent-window blocks]; the target's own full-attention verification pass is left completely untouched. Because RoPE is already baked into cached keys, dropping the intervening blocks changes nothing about the retained blocks' scores, so the change needs no new kernel, no training, and no extra parameters -- just a few lines in how the draft's KV index is built.

방법 Windowed-MTP는 StreamingLLM 스타일의 슬라이딩 윈도우와 소수의 sink 토큰을 오직 draft의 attention에만 적용한다 -- draft의 paged-KV block table을 [sink 블록] + [최근 윈도우 블록]으로 줄이는 방식으로 구현되며, target 자체의 full-attention 검증 패스는 전혀 건드리지 않는다. RoPE가 이미 캐시된 key에 반영돼 있으므로 중간 블록을 빼도 남은 블록들의 점수는 변하지 않아, 새 커널도 학습도 추가 파라미터도 필요 없이 draft의 KV 인덱스 구성 코드 몇 줄만 바꾸면 된다.

Key idea Speculative decoding is exact by construction -- the target always decides which proposed tokens get accepted, so a cheaper, less-informed drafter can only change which tokens are proposed, never which are accepted, meaning any approximation on the draft side is automatically safe. It is like a junior reporter drafting copy for a strict fact-checking editor-in-chief: the reporter can skim only the last few paragraphs of a long archive instead of rereading everything, because the editor will catch and correct anything wrong regardless. The shortcut only changes how many drafts survive review, never what gets published.

핵심 아이디어 스펙 디코딩은 구조적으로 정확하다 -- target이 항상 어떤 제안 토큰을 받아들일지 결정하므로, draft가 더 저렴하고 부정확해져도 바뀌는 것은 무엇이 제안되는가일 뿐 무엇이 채택되는가는 절대 바뀌지 않는다. 이는 깐깐한 편집장에게 검수를 맡기는 신입 기자와 같다 -- 신입 기자가 방대한 과거 기사 전체를 다시 읽는 대신 최근 몇 단락만 훑고 초안을 써도, 편집장이 어차피 틀린 부분을 다 잡아내므로 결과물의 정확성은 그대로다. 지름길은 얼마나 많은 초안이 검수를 통과하느냐만 바꿀 뿐, 무엇이 최종 게재되느냐는 절대 바꾸지 않는다.

Validation Losslessness is checked directly via an all-cell greedy output diff between native and windowed decoding, confirming outputs match up to pre-existing bf16 floating-point nondeterminism across every tested (model x input) combination. Acceptance preservation is checked via three separate probes -- a window-size dose-response sweep, a per-position conditional-acceptance decomposition, and an in-run decision-invariance check -- run across three production-scale architecture families (Qwen GDN-MoE 35B/122B, Mamba2-hybrid Nemotron-3-120B) at true 1-million-token context inside a production serving framework (SGLang), not a toy single-request harness.

검증 무손실성은 native와 windowed 디코딩 간 all-cell greedy 출력 diff로 직접 확인했으며, 테스트한 모든 (모델 x 입력) 조합에서 기존 bf16 부동소수점 비결정성 이상의 차이는 없었다. 수용률 보존은 window 크기 dose-response 스윕, 위치별 조건부 수용률 분해, in-run 결정-불변성 점검이라는 세 가지 독립적 프로브로 확인했으며, 실제 프로덕션 서빙 프레임워크(SGLang) 안에서 세 가지 프로덕션급 아키텍처 계열(Qwen GDN-MoE 35B/122B, Mamba2-hybrid Nemotron-3-120B)을 대상으로 진짜 백만 토큰 컨텍스트에서 검증했다.

Results At 1M context, windowing cuts the per-decode-step cost over the shipped native MTP draft by 28% to 44% -- an input-invariant margin that widens further with context length -- while the windowed draft's top-1 token proposal matches the unwindowed draft's proposal 86-94% of the time. End-to-end, Windowed-MTP is 11% to 53% faster than native MTP and 1.58-2.55x faster than no speculation at all on a true 1M-token retrieval benchmark, and reclaims the now-unread 7.7-11% of total KV as a compact ring buffer at zero acceptance or quality cost.

결과 1M 컨텍스트에서 windowing은 출시된 native MTP draft 대비 decode step당 비용을 28%에서 44% 절감하며 -- 이 폭은 입력에 무관하고 컨텍스트가 길어질수록 더 커진다 -- windowed draft의 top-1 제안 토큰은 native draft의 제안과 86-94%의 확률로 일치한다. end-to-end로는 진짜 1M 토큰 retrieval 벤치마크에서 native MTP 대비 11-53% 빠르고 스펙 디코딩을 아예 쓰지 않는 것보다 1.58-2.55배 빠르며, 더 이상 읽히지 않는 전체 KV의 7.7-11%를 수용률이나 품질 손실 없이 compact ring buffer로 회수한다.

Comparison Against the native full-context MTP head the target already ships -- the paper's deliberately chosen, tightly controlled baseline that isolates the windowing effect with weights, verifier, and tree shape held fixed -- Windowed-MTP wins every tested cell at every draft depth. Against a separately trained drafter like EAGLE-3, the paper argues, without directly testing, that the same windowing trick should apply just as well, since an EAGLE-style head is itself a transformer layer that reads the full KV cache and would pay an identical long-context tax.

비교 target이 이미 출시한 native full-context MTP head 대비 -- 논문이 의도적으로 선택한, weight·verifier·tree 구조를 고정한 채 windowing 효과만 분리하는 엄격한 baseline -- Windowed-MTP는 테스트한 모든 draft depth에서 승리한다. EAGLE-3처럼 별도로 학습된 drafter와는 직접 비교하지 않았지만, 논문은 EAGLE 스타일 head 역시 전체 KV 캐시를 읽는 transformer layer이므로 동일한 장문맥 세금을 물 것이고 같은 windowing 기법이 똑같이 적용될 것이라고 주장한다.

Significance It quantifies and fixes a specific, previously invisible failure mode in a technique -- built-in MTP speculative decoding -- that frontier serving stacks already ship by default. At short context the tax is genuinely negligible, so the problem simply did not exist until recent models started supporting million-token context windows. The fix also generalizes cleanly across three quite different architecture families (MoE-heavy, Mamba2-hybrid) and composes with orthogonal techniques the field is simultaneously adopting, like fp8 KV quantization and hybrid/linear-attention targets, without losing its relative benefit.

의의 프런티어 서빙 스택이 이미 기본값으로 탑재한 기법(내장 MTP 스펙 디코딩)에 존재하던, 지금까지 보이지 않았던 특정 실패 모드를 정량화하고 고쳤다는 점이 의미가 크다 -- 짧은 컨텍스트에서는 이 세금이 정말로 무시할 만했기 때문에, 최근 모델들이 백만 토큰급 컨텍스트를 지원하기 전까지는 문제 자체가 존재하지 않았다. 서로 꽤 다른 세 아키텍처 계열(MoE 위주, Mamba2-hybrid)에 걸쳐 깔끔하게 일반화되고, fp8 KV 양자화나 hybrid/linear-attention target처럼 업계가 동시에 채택 중인 다른 기법들과도 상대적 이득을 잃지 않고 결합된다.

Limitations The paper's own Discussion section states that the cost-side gain scales with how windowable the draft's per-step cost share is -- architectures whose draft attention is a smaller fraction of total draft cost see a smaller, still-positive speedup, e.g. 28% for the Mamba2-hybrid Nemotron versus 30-44% for the Qwen MoE drafts. All headline numbers use bf16 KV at B=1/TP1 (batch/TP2 sweeps are reported as supporting results), and (reviewer judgment) the technique's benefit necessarily vanishes at short context, where the window already spans the whole sequence and Windowed-MTP degenerates to native MTP.

한계 논문 스스로의 Discussion 절은 비용 절감 폭이 draft의 step당 비용 중 windowable한 비중에 비례한다고 명시한다 -- draft attention이 전체 draft 비용에서 차지하는 비중이 작은 아키텍처는 절감폭도 작다(여전히 양수이긴 하지만), 예를 들어 Mamba2-hybrid Nemotron은 28%인 반면 Qwen MoE draft는 30-44%다. 모든 헤드라인 수치는 bf16 KV, B=1/TP1 기준이며(batch/TP2 스윕은 보조 결과로 보고), (리뷰어 판단) 이 기법의 이득은 짧은 컨텍스트에서는 필연적으로 사라진다 -- 윈도우가 이미 시퀀스 전체를 덮어 Windowed-MTP가 native MTP로 그대로 수렴하기 때문이다.

Future work The paper's own Discussion section names concrete next steps: applying the identical windowing trick to trained EAGLE-3 drafters and to the emerging class of block-diffusion drafters (which the authors flag as facing an even larger tax, since their draft-KV grows per layer with context); building an online adaptive-depth controller now that windowing makes shallower depths often optimal; and a kernel-level implementation of prefill-side pruning to also cut time-to-first-token, not just per-token decode latency.

향후 과제 논문 스스로의 Discussion 절은 구체적인 다음 단계를 명시한다: 동일한 windowing 기법을 학습된 EAGLE-3 drafter와, (저자들이 draft-KV가 레이어마다 컨텍스트에 비례해 늘어나 세금이 더 클 것이라 지적하는) 신흥 block-diffusion drafter에도 적용하는 것, windowing 덕분에 얕은 depth가 자주 최적이 되는 상황을 활용할 온라인 적응형 depth 컨트롤러 구축, 그리고 decode 지연뿐 아니라 첫 토큰까지의 시간(TTFT)도 줄이기 위한 prefill-side pruning의 커널 수준 구현이다.

Released resources Official code and reproduction scripts (an SGLang patch plus run/config scripts and seeded inputs) are verified public on GitHub. The paper itself is on arXiv under a CC BY 4.0 license.

자원 공개 공식 코드와 재현 스크립트(SGLang 패치, 실행/설정 스크립트, seed가 고정된 입력)는 GitHub에 공개돼 있음을 확인했다. 논문 자체는 CC BY 4.0 라이선스로 arXiv에 공개돼 있다.