LongSpec: Long-Context Lossless Speculative Decoding with Efficient Drafting and Verification

1 Introduction

• Short-context speculative decoding methods typically train on sequences under 4k tokens, while long-context inference is increasingly important for agents and reasoning tasks.The paper contrasts lossless speculative decoding with alternatives such as quantization and model cascades, which may compromise output quality.Source: §1 · p. 1
• Some long-context methods draft with the full target model and compressed KV cache, but its computational cost can limit speed and underperform short-context methods.The draft model’s weight in low-batch settings can outweigh speculative benefits.Source: §1 · p. 2

Why cannot SoTA SD methods for short contexts

• Long-context speculative decoding faces three challenges: draft-model KV-cache growth, scarce long-sequence training data with fixed target RoPE bases, and inefficient tree attention.Conventional RoPE-base extension is not directly applicable because the draft model must match the target model’s fixed base.Source: Why cannot SoTA SD methods for short… · p. 2
• LONGSPEC pairs a constant-memory draft architecture with Anchor-Offset Indices, enabling short-sequence-trained drafts to generalize to longer inference contexts.Hybrid Tree Attention is the third innovation, addressing inefficient tree-attention computation.Source: Why cannot SoTA SD methods for short… · p. 2
• Across five long-context understanding datasets, LONGSPEC achieves up to 3.26× speedup over strong baselines using Flash Attention2.The paper also reports up to 7× speedup over common baselines using HuggingFace implementation.Source: Why cannot SoTA SD methods for short… · p. 2

2 Related Work

• Early speculative decoding methods use smaller LLMs to draft tokens, while later work improves these methods and develops specialized drafts and tree verification (Xia et al., 2023; Leviathan et al., 2023; Liu et al., 2024d; Bae et al., 2023; Liu et al., 2025a; Sun et al., 2023; Chen et al., 2024; Cai et al., 2024; Li et al., 2024c; Du et al., 2024; Huang et al., 2025a).Tree-based approaches support more adaptive and parallelizable decoding, including applications to vision-language models (Huang et al., 2025b).Source: §2 · pp. 2–3
• Other approaches use part of the target model as a draft, while retrieval-based methods use N-gram matching and pre-existing data patterns to construct drafts without additional model training (Zhang et al., 2024; Elhoushi et al., 2024; Xia et al., 2025; Fu et al., 2024; He et al., 2024; Zhao et al., 2024; Liu et al., 2025b; Shen et al., 2026).The paper uses Flash Attention to refer to FlashDecoding, implemented via flash_attn_with_kvcache.Source: §2 · pp. 2–3
• Only a few prior works focus on lossless long-context speculative decoding; TriForce, MagicDec, and QuantSpec mainly use the target model with a sparse KV cache (Sun et al., 2024; Tiwari et al., 2025).The paper identifies computationally intensive drafts as a constraint across batch sizes and contrasts its one-transformer-block draft model.Source: §2 · p. 3

3.1 Memory-Efficient Architecture

• The draft model uses sliding-window self-attention for local context and target-model KV-cache cross-attention for long-range information, keeping additional memory constant with context length.The self-attention window is set to 512; the design also shares the embedding layer and LM head with the target model.Source: §3.1 · p. 3, Figure 2 · p. 4

3.2 Effective Training Regimes

• Anchor-Offset Indices reserve positions 0–3 and assign later tokens consecutive indices from a random offset, so short-context training covers larger positions while limiting target-model loss increase to about 0.001.Flash Noisy Training simulates inference-time KV visibility constraints and increases acceptance length by 14.7% versus training without it.Source: §3.2 · pp. 4–5

3.3 Fast Tree Attention

• Tree attention’s prefix-tree masks have limited compatibility with optimized long-context kernels, so the hybrid design applies masking only to the small set of speculative tokens.Earlier approaches that apply masks to the full attention matrix rely on PyTorch eager execution and slow as sequence length increases.Source: §3.3 · p. 5
• Hybrid Tree Attention separates cached prefix KV pairs, computed with Flash Attention, from speculative-token KV pairs, computed with a custom masked Triton kernel.The two attention outputs are aggregated using a log-sum-exp trick, retaining parallel branch verification without compromising correctness.Source: §3.3 · pp. 5–6

4.1 Settings

• Evaluation spans five long-context understanding datasets and four math-reasoning datasets; drafts are pretrained, long-context trained, then fine-tuned on a long-context SFT dataset.The study compares against the target model, PLD, and MagicDec, using walltime speedup relative to vanilla autoregressive decoding and average acceptance length τ.Source: §4.1 · p. 6

4.2 Main Results

• At T = 0, LongSpec reaches up to 2.67× speedup on summarization and 3.26× on code completion; at T = 1, it achieves around 2.5× speedup.It outperforms the other approaches on both task types at T = 0 and maintains a substantial lead over MagicDec at T = 1.Source: §4.2 · p. 6, Table 1 · p. 7, Figure 3 · p. 7

4.3 Ablation Studies

• 3.93× faster to reach the same loss level, models with Anchor-Offset Indices also achieve higher output speed and larger acceptance length on Multi-News and RepoBench-P.The loss comparison is on real long-context data; the performance comparison is with and without the indices.Source: Figure 4 · p. 8, §4.3 · p. 8, Table 2 · p. 8
• 49.92 ms to 12.54 ms, Hybrid Tree Attention reduces target-model attention latency by approximately 75% while draft-forward and target-FFN times remain comparable.The verification-step time difference is minimal.Source: Figure 5 · p. 8, §4.3 · p. 8

4.4 Long Reasoning Acceleration

• On QwQ-32B across four math-reasoning datasets, LongSpec achieves about 2.34× higher generation speed than the baseline and averages 3.81 accepted tokens.Evaluation uses a maximum output length of 32k tokens; MagicDec's draft degrades to the target model with a relatively short prefix in long-reasoning tasks, failing to accelerate.Source: Table 3 · p. 8, §4.4 · p. 9

4.5 Throughput

• At batch size 8, LONGSPEC reaches 561.32 tokens/s, versus MagicDec's 310.58 tokens/s and Vanilla's 286.96 tokens/s on RepoBench-P.LONGSPEC maintains higher throughput than both baselines across all tested batch sizes.Source: Figure 6 · p. 9, §4.5 · p. 9

5 Conclusion

• LongSpec combines a constant-memory draft model, Anchor-Offset Indices, and Hybrid Tree Attention to address long-context decoding challenges.The draft model combines sliding-window self-attention with cache-free cross-attention; the position indices train large positional indices on short sequences, and Hybrid Tree Attention integrates tree decoding with Flash Attention.Source: §5 · p. 9
• LongSpec is designed to improve lossless speculative decoding in long-context scenarios.The framework addresses memory overhead, inadequate training for large position indices, and inefficient tree attention computation.Source: §5 · p. 9
• Experiments demonstrate LongSpec’s effectiveness on long-context understanding and real-world long reasoning tasks.The authors identify long-context-specific speculative decoding as important for efficient large-scale language-model inference.Source: §5 · p. 9

A Related Work about Lossy Speculative Decoding

• BiLD and speculative cascading use refinement or deferral strategies to seek speedups or improved cost-quality trade-offs, while accepting quality degradation or combining cascade rules with speculation.BiLD uses a small autoregressive model and occasionally invokes a larger model nonautoregressively; speculative cascading integrates cascade-style deferral rules with speculative execution.Source: A Related Work about Lossy Speculative… · p. 13
• MTAD approximates a larger model’s multi-token joint distribution with a smaller auxiliary model, accepting bounded approximation error.A trained judge module can recognize valid continuations without perfect target-model alignment, boosting acceptance rates and speed.Source: A Related Work about Lossy Speculative… · p. 13
• RSD uses a process reward model to guide target-model invocation, introducing controlled bias toward high-reward outputs to optimize cost-quality trade-offs.RAPID instead uses a retrieval-augmented-generation approach on shortened contexts as its drafter.Source: A Related Work about Lossy Speculative… · p. 13
• TokenSwift combines LLMs with partial KV cache and N-gram tables for generation up to 100k tokens, reducing computation time from hours to minutes.The passage describes it as accelerating ultra-long sequence generation.Source: A Related Work about Lossy Speculative… · p. 13

B The Intuition of Why KV Cache Can Help

• The KV cache stores context accumulated from earlier tokens and contributes to next-token prediction alongside input embeddings and model parameters.The target model uses all three components when predicting the next token.Source: B The Intuition of Why KV Cache Can Help · p. 13
• LongSpec’s draft model shares the target model’s input embeddings and KV cache, leaving model parameters as the remaining difference between their predictions.The shared cache removes discrepancies caused by differing contextual representations and aligns draft predictions more closely with the target’s.Source: B The Intuition of Why KV Cache Can Help · pp. 13–14

C Correctness for Attention Aggregation

• The proof reduces the matrix-level correctness claim to showing that each individual query’s output meets the requirements.Because each row of Q represents a separate query, the proof can treat its key-value inputs as already subject to the tree mask.Source: C Correctness for Attention Aggregation · p. 14
• Merged attention attends over stacked speculative and cached key-value pairs, whose logits are split into corresponding sub-logit matrices.Rows of Zspecs and Zcache give dot products for the same query against speculative and cached keys, respectively.Source: C Correctness for Attention Aggregation · p. 14
• The aggregation tracks log-sum-exp values for the separate sub-logit sets to combine partial attention outputs.The passage introduces the log-sum-exp of merged attention as part of the aggregation argument.Source: C Correctness for Attention Aggregation · p. 14

• Training uses eight A100 80GB GPUs, with ZeRO-1 for short-context 7B, 8B, and 13B targets and ZeRO-3 for long-context targets and all 33B settings.Source: D Experiments Details · p. 15
• Draft models are optimized with cross-entropy while target-model parameters remain frozen, using fused linear-and-cross-entropy computation to alleviate peak VRAM use.The fused loss computes the LM head and softmax together.Source: D Experiments Details · p. 15
• Short-context SlimPajama-6B training uses batch size 2048 and maximum learning rate 5e-4; long-context training uses batch size 256 and maximum learning rate 5e-6.Both use a cosine schedule and AdamW, and each dataset receives one training epoch.Source: D Experiments Details · p. 15
• The main computational cost is forwarding the target model to obtain its KV cache; pre-stored context-caching KV data can serve as training data and accelerate training.Source: D Experiments Details · p. 15
• Dynamic beam search halts descendant computation on low-likelihood branches instead of discarding nodes, avoiding KV-cache movement while allowing draft sequences of unequal lengths.Beam width is [4, 16, 16, 16, 16] across speculation steps.Source: D Experiments Details · p. 15

• Vanilla HF decoding runs at 14–30 tokens/s across five datasets, while switching to Flash Attention raises speed to about 50 tokens/s, more than 2.5× faster.Source: E Experimental Results of EAGLE and… · p. 15
• Token Recycling reaches an average acceptance length around 3 and provides moderate acceleration on many tasks, but consistently underperforms LongSpec.Source: E Experimental Results of EAGLE and… · p. 15
• EAGLE reaches an acceptance length around 2 and 26–40 tokens/s, a 30–50% speedup over Vanilla HF, but remains slower than Vanilla FA in every setting.Source: E Experimental Results of EAGLE and… · p. 15
• About 100 tokens/s: LongSpec with hybrid tree attention achieves this decoding speed across all models and datasets, exceeding the compared methods.Vanilla Flash Attention reaches about 50 tokens/s, while EAGLE remains below it in every setting and Token Recycling consistently underperforms LongSpec.Source: E Experimental Results of EAGLE and… · p. 15

F Performance Analysis with Varying Prefill Lengths

• Generation speed and average acceptance length remain stable across prefill ranges, with only a small speed drop in the 25k–32k range.The authors report that draft quality is unaffected by prefill length.Source: F Performance Analysis with Varying… · p. 16
• Draft time rises from 8.91 ms to 9.25 ms and target time from 25.63 ms to 30.89 ms across the shortest-to-longest contexts, while verify time increases from 6.18 ms to 6.28 ms.Source: F Performance Analysis with Varying… · p. 16

G Pseudo Code

• Flash Noisy Training pseudocode randomly selects 1–4 tokens to drop, slices the inputs and cache, then pads the output to the original query length.The pseudocode identifies the query, key-cache, and value-cache inputs and returns the output projection.Source: G Pseudo Code · pp. 16–17
• Anchor-Offset Indexing applies an offset after the first four anchor positions; for N = 128, the example offset is 16257.The pseudocode takes sequence length, maximum length, and query states as inputs and outputs query states with RoPE applied using modified indices.Source: G Pseudo Code · p. 16

• The examples are from GovReport on LongChat-7B, where blue marks draft tokens accepted by the target model; the complete answer is omitted.Source: H Case Study · p. 17
• The opioid report summary covers medical uses, risks including dependence, addiction, overdose, and diversion, and the need for proper monitoring and continued treatment research.Source: H Case Study · p. 17
• The railroad report summary describes the Railroad Retirement Board’s RRA and RUIA benefit programs and trends in railroad employment.Source: H Case Study · p. 18
• The Congressional Gold Medal summary describes the award’s purpose, legislative process, design, and use to honor individuals and groups.Source: H Case Study · p. 18
• The DHS report summary reviews FY2019 appropriations, the budget process, and debates including border security, cybersecurity, and disaster response.Source: H Case Study · p. 19