Saved in:
Bibliographic Details
Main Authors: He, Liang, Wen, Jingbo, Zhan, Qishi, Chen, Yixiong, Cui, Kangning, Lan, Qizhen, Wang, Xilu
Format: Preprint
Published: 2026
Subjects:
Online Access:https://arxiv.org/abs/2606.00144
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910273087995904
author He, Liang
Wen, Jingbo
Zhan, Qishi
Chen, Yixiong
Cui, Kangning
Lan, Qizhen
Wang, Xilu
author_facet He, Liang
Wen, Jingbo
Zhan, Qishi
Chen, Yixiong
Cui, Kangning
Lan, Qizhen
Wang, Xilu
contents Speculative decoding speeds up autoregressive decoding by using a drafter to propose multiple tokens that a verifier validates in parallel. In resource-constrained deployments, the drafter uses a sparse KV cache to limit peak GPU memory and end-to-end latency under a fixed KV budget, while the verifier keeps a full KV cache. Mid-to-long context inference (4K--16K context length) is common in real applications. However, naive sparse/full speculative decoding suffers from the sparse/full mismatch as context length grows, causing the acceptance rate to drop quickly. We propose BudgetDraft, a multi-view sparse training method for sparse drafting in mid-to-long inference. The drafter is exposed to multiple sampled KV budgets during training and learns to align each sparse view with one shared full-cache teacher target. BudgetDraft combines an acceptance-aware loss on a full-cache branch with a multi-view loss on a sparse-cache branch, producing a single budget-robust drafter that recovers acceptance across sparsity levels without extra inference-time components. Experimental results on PG-19, LongBench, and LWM show that BudgetDraft achieves up to 6.55x, 4.46x, and 2.10x end-to-end speedup vs AR at 4K, 8K, and 16K context lengths, while keeping the inference pipeline memory-friendly.
format Preprint
id arxiv_https___arxiv_org_abs_2606_00144
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding
He, Liang
Wen, Jingbo
Zhan, Qishi
Chen, Yixiong
Cui, Kangning
Lan, Qizhen
Wang, Xilu
Machine Learning
Artificial Intelligence
Speculative decoding speeds up autoregressive decoding by using a drafter to propose multiple tokens that a verifier validates in parallel. In resource-constrained deployments, the drafter uses a sparse KV cache to limit peak GPU memory and end-to-end latency under a fixed KV budget, while the verifier keeps a full KV cache. Mid-to-long context inference (4K--16K context length) is common in real applications. However, naive sparse/full speculative decoding suffers from the sparse/full mismatch as context length grows, causing the acceptance rate to drop quickly. We propose BudgetDraft, a multi-view sparse training method for sparse drafting in mid-to-long inference. The drafter is exposed to multiple sampled KV budgets during training and learns to align each sparse view with one shared full-cache teacher target. BudgetDraft combines an acceptance-aware loss on a full-cache branch with a multi-view loss on a sparse-cache branch, producing a single budget-robust drafter that recovers acceptance across sparsity levels without extra inference-time components. Experimental results on PG-19, LongBench, and LWM show that BudgetDraft achieves up to 6.55x, 4.46x, and 2.10x end-to-end speedup vs AR at 4K, 8K, and 16K context lengths, while keeping the inference pipeline memory-friendly.
title BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2606.00144