PaperlayerView exact revision on arXiv

Source-linked AI summary

ASPIRE: Asynchronous Batched Self-Speculative Decoding for Long-Context LLM Inference

Amir Ziashahabi, Hossein Entezari Zarch, Lei Gao, Murali Annavaram, Salman Avestimehr

arXiv:2609.17943v1cs.LGcs.CLcs.DC

TL;DR

Long-context decoding is limited by memory-bound attention and synchronized speculative batching that cannot adapt to heterogeneous, changing draft lengths. ASPIRE combines mixed non-synchronized execution, per-request cost-aware scheduling, and intra-draft context refresh. Across three models and five workloads, it achieves 1.70-4.58× speedup over autoregressive decoding and improves average speedup by approximately 27% over prior self-speculative baselines.

  • Problem

    Long-context inference repeatedly reads request-specific KV caches, while synchronized batched speculative decoding uses one schedule despite request-dependent and dynamically changing optimal draft lengths.

  • Method

    ASPIRE combines unified mixed forwarding, per-request online scheduling, and a single-layer full-attention refresh that updates sparse context during drafting.

  • Results

    Across three models and five reasoning and long-context workloads, ASPIRE achieves 1.70-4.58× speedup over AutoRegressive decoding and improves average speedup by approximately 27% over prior self-speculative baselines.

  • Takeaways & Limitations

    ASPIRE enables batched requests to follow different speculation trajectories while improving throughput over autoregressive and prior self-speculative decoding baselines.

  • Takeaways & Limitations

    Integrating mixed forward with a general continuous-batching engine, particularly under chunked prefill, remains future work.

Abstract

from arXiv · show

Long-context LLM inference is bottlenecked by attention, whose repeated KV-cache reads make decoding memory-bound. Self-speculative decoding alleviates this by drafting tokens with sparse attention and verifying them with full attention, but existing batched methods remain synchronized: all requests in a batch share a single draft-verify schedule, even though the optimal draft length varies widely across requests and changes dynamically within each request. We propose ASPIRE, a non-synchronized batched self-speculative decoding framework built on three components. First, a unified mixed forward allows drafting and verifying requests to coexist in the same batched forward pass, removing the need for global draft-verify phases. Second, a lightweight online speculation scheduler uses per-request acceptance-rate estimates and a batch-aware cost model to let each request independently choose when to verify. Third, an intra-draft refresh layer performs full attention at a single designated layer during drafting, updating the sparse context at every draft step to reduce staleness during drafting. Across three models and five reasoning and long-context benchmarks, ASPIRE achieves $1.70$-$4.58\times$ speedup in decoding throughput over autoregressive baselines and improves average speedup by approximately $27\%$ over the strongest prior self-speculative baselines.

1 Introduction

Long-context decoding is memory-bound because attention repeatedly reads request-specific KV caches, while synchronized speculative batching cannot accommodate heterogeneous and changing draft lengths. ASPIRE addresses this with non-synchronized mixed execution, per-request scheduling, and lightweight intra-draft context refresh.

  • 1 Introduction: Different requests have heterogeneous optimal draft lengths, and each request’s optimal length can change during generation, making shared schedules inefficient.Synchronized batching can under-draft some requests and over-draft others.
  • 1 Introduction: ASPIRE lets requests draft or verify independently within one mixed batched forward, removing synchronized batch-wide draft-verify phases.Requests can use partial KV context for drafting or full KV context for verification while preserving shared batched execution.
  • 1 Introduction: The online scheduler uses per-request acceptance estimates and a batch-aware draft/verify cost ratio to choose each request’s action on every forward pass.This enables each request to follow its own speculation trajectory.
  • 1 Introduction: The intra-draft refresh mechanism updates sparse KV context during drafting using a lightweight full-attention refresh layer, improving draft quality and enabling longer drafts.It exploits attention correlation across layers and consecutive tokens with minimal overhead.
  • 1 Introduction: 1.70-4.58× speedup over autoregressive decoding was achieved across three models and five workloads, with best overall performance against prior speculative baselines.The reported result covers reasoning and long-context workloads.

2 Background

Transformer decoding repeatedly reads each request’s full KV cache, making long-context attention memory-bound; sparse-context drafting reduces this cost while full-context verification preserves exactness. Existing batched speculative systems synchronize requests under a shared schedule despite differing verification needs.

  • 2 Background: Long-context decoding is memory-bound because every generated token reads the full prior KV cache for its request.MLP computation benefits from batching, but attention remains tied to independent per-request contexts.
  • 2 Background: Self-speculative decoding drafts tokens with sparse-context attention and verifies the speculative prefix with full-context attention.The full KV-cache access is amortized across multiple accepted drafted tokens during verification.
  • 2 Background: Synchronized batched speculative decoding makes all requests draft and verify together under a common schedule, forcing heterogeneous requests to advance in lockstep.This can mismatch each request’s optimal verification depth.

3 Motivating Observations

Measurements show that draft-length behavior varies both across requests and across verification rounds within a request, motivating per-request step-level scheduling. Predictor experiments further favor recent cross-layer or late-layer signals over stale temporal reuse, motivating ASPIRE’s refresh design.

  • 3 Motivating Observations: Accepted draft lengths range from 2.56 to 20.00 tokens across 256 LongBench samples, with mean 11.51, showing that one shared draft length is suboptimal.The study used 20-token drafts on 16k-20k-token prefixes.
  • 3 Motivating Observations: Accepted lengths fluctuate substantially across verification rounds within individual requests, so fixed request-specific lengths are also insufficient.The observed variation motivates per-request, step-level scheduling rather than synchronized batch-wide decisions.
  • 3 Motivating Observations: Temporal attention-map reuse becomes stale quickly as token lag grows, whereas nearby cross-layer signals provide more stable sparse-context predictions.The cross-layer comparison is diagnostic because fixed offsets would require multiple layer-specific refresh signals.
  • 3 Motivating Observations: A single late-layer refresh offers the best quality-overhead tradeoff: recall improves through refresh layer 2 and then changes little as more full-attention layers are added.Deeper refresh points strengthen prediction but increase full-attention work at every draft step.
  • 3 Motivating Observations: ASPIRE therefore uses a late-layer source from the previous token, combining low temporal staleness with a strong cross-layer signal at one full-attention layer per draft step.The design achieves comparable sparse-context recall to same-token refresh with less overhead.

4 ASPIRE

ASPIRE enables independently scheduled drafting and verification within one batched forward pass, while refreshing sparse contexts during drafting. Its scheduler combines acceptance estimates with batch-aware cost modeling to choose per-request draft lengths.

  • 4 ASPIRE: ASPIRE combines unified mixed forwarding, per-request scheduling, and intra-draft context refresh to improve speculative utilization in batched long-context inference.Requests may draft or verify concurrently, while a designated full-attention layer updates sparse contexts during drafting.
  • 4.1 Unified Mixed Forward: Different requests can draft or verify within the same forward step, eliminating global draft-verify phases while preserving shared batched execution.Verifying requests use full KV context, whereas drafting requests use selected KV pages.
  • 4.2 Speculation Scheduler: Each request independently targets a draft length using online acceptance-rate estimates and a draft-to-verify cost ratio, then verifies at that target or a hard cap.The scheduler applies its policy at every forward step rather than jointly optimizing the full future trajectory.
  • 4.2 Speculation Scheduler: The batch-aware timing model represents MLP cost by total forward tokens and attention cost by total KV-cache length read.This reflects compute-bound batched MLP execution and memory-bound attention under the current mixed-batch composition.
  • 4.3 Refresh Layer for Intra-Draft Context Refresh: A designated refresh layer performs full attention during drafting, converts attention signals into page scores, and selects the sparse context for the next step.The selected context adapts as generation moves farther from the last verified token, addressing staleness during drafting.
  • 4.3 Refresh Layer for Intra-Draft Context Refresh: The refresh mechanism retains the most recent pages and fills the remaining sparse-context budget with highest-scoring pages.The budget is determined from the request page count, sparsity rate, and a minimum floor.

5 Experiments

Experiments evaluate ASPIRE across three models and five reasoning and long-context workloads, comparing it with autoregressive and self-speculative baselines. ASPIRE’s gains are largest for long contexts, while ablations show benefits from mixed execution, adaptive scheduling, and refresh.

  • 5.1 Experimental Setup: Experiments use Qwen3-1.7B, Qwen3-8B, and DS-Llama-8B on AIME 2025, CodeElo, and LongBench-v2 settings spanning 30K–40K and 80K–100K contexts.Throughput is measured over decoding forward passes on a single NVIDIA H100 NVL 96GB GPU unless otherwise stated.
  • 5.2 Overall Throughput Results: Up to 4.58× speedup over AutoRegressive is achieved on long-context workloads, with gains increasing as context grows from 16K to 80K.The corresponding maxima are 4.58×, 2.75×, and 2.66× on Qwen3-1.7B, Qwen3-8B, and DS-Llama-8B.
  • 5.2 Overall Throughput Results: ASPIRE delivers the best overall average speedup on Qwen3-8B and DS-Llama-8B, at 2.08× and 2.14× across all five workloads.On Qwen3-1.7B, ASPIRE also has the best average performance at 3.34×, although the strongest variant varies by workload.
  • 5.3 In-depth Analysis: ASPIRE improves over synchronized fixed-length drafting increasingly as context grows, reaching 2.75× versus 2.22× on Qwen3-8B’s 80K–100K workload.The same pattern appears on DS-Llama-8B, where ASPIRE rises from 2.21× to 2.66× relative to ASPIRE-Fixed.
  • 5.3 In-depth Analysis: The batch-aware scheduler outperforms simple acceptance-based feedback when verification depends jointly on changing acceptance behavior and batch-dependent cost.Across models, ASPIRE averages 3.34×, 2.08×, and 2.14×, versus 3.26×, 1.94×, and 1.93× for ASPIRE-FSM.

6 Related Work

Related work reduces long-context inference cost through sparse attention, speculative decoding, and self-speculative decoding. ASPIRE differs by combining non-synchronized intra-batch scheduling with mixed target-model execution, while integration with general continuous batching remains future work.

  • Sparse attention for long-context inference: Sparse-context methods reduce attention cost, but fixed or prompt-derived retention strategies cannot adapt to changing token-level needs during generation.
  • Speculative decoding with auxiliary draft models: Auxiliary-draft speculative decoding uses a separate model or predictor to propose tokens before target-model verification, a pipeline distinct from ASPIRE’s unified target-model forward.
  • Self-speculative decoding: Self-speculative methods use the target model with sparse drafting and full-attention verification, motivating ASPIRE’s long-context design.
  • Production systems: ASPIRE addresses a different synchronization boundary by allowing draft and verify roles to coexist in one batched target-model invocation.Integrating this mixed forward with a general continuous-batching engine, particularly under chunked prefill, remains future work.

7 Conclusion

ASPIRE is a dynamic, non-synchronized batched self-speculative decoding framework combining mixed execution, per-request scheduling, and intra-draft context refresh. Across three models and five workloads, it improves throughput over autoregressive and prior self-speculative decoding.

  • 7 Conclusion: ASPIRE achieves 1.70–4.58× decoding speedup over AutoRegressive and approximately 27% higher average speedup than the strongest prior self-speculative baselines.
  • 7 Conclusion: ASPIRE combines unified mixed forward, per-request online scheduling, and intra-draft refresh so batch requests can follow different speculation trajectories while their sparse contexts stay updated.

B.1 Draft-Length Variability on Expanded AIME

Expanded AIME experiments show that accepted draft lengths vary broadly across requests and fluctuate across verification rounds. This variability appears across all three evaluated architectures on 90 reasoning problems.

  • B.1 Draft-Length Variability on Expanded AIME: The analysis covers all 90 AIME 2024–2026 problems, with requests sorted by completion time in the verification-round heatmaps.
  • B.1 Draft-Length Variability on Expanded AIME: Accepted draft length fluctuates between verification rounds within individual requests, alongside the broad differences in per-request means.Both forms of variation occur across all three architectures at magnitudes comparable to the LongBench measurements.

B.2 Output Quality

ASPIRE preserves output quality while accelerating decoding: its greedy LongBench-v1 results closely match autoregressive decoding, and its scheduling remains near a hindsight oracle.

  • ASPIRE is exact by construction because sparse attention proposes tokens, while full-attention verification determines accepted sequences.The implementation check targets errors in page selection or the acceptance test that could affect outputs without changing measured speedup.
  • ASPIRE’s LongBench-v1 macro average differs from greedy autoregressive decoding by −0.20 points across 1,460 examples.Four of eight tasks agree within 0.05 points, while the largest single deviation is 0.5 points.
  • ASPIRE’s throughput remains within 4.7–20.2% of a hindsight scheduling oracle across the evaluated models and workloads.The oracle replays realized accept/reject traces using the same mixed-forward backend, sampling configuration, and batch size.

B.4 Tensor Parallelism and Recalibration

Under tensor parallelism, ASPIRE retains strong throughput, but its timing model must be recalibrated for each deployment configuration because coefficients do not transfer across TP degrees.

  • ASPIRE reaches 2.16× speedup on AIME25 and 2.07× on LB[16k-18k] at TP=2, exceeding the corresponding single-GPU figures.Table 4 evaluates Qwen3-8B with the largest batch fitting at each tensor-parallel degree.
  • The timing model’s functional form transfers across TP degrees, but its coefficients do not, requiring calibration for each deployment configuration.Fitted models predict step time within 5–11% MAPE, whereas reusing TP=1 coefficients at TP=2 raises error to 35–50% MAPE.
  • Stale TP=1 coefficients reduce LB[16k-18k] speedup from 2.07× to 1.51×, while recalibration restores the full result without affecting correctness.The stale coefficients remain benign on AIME25, where speedup is 2.21×.

B.5 Robustness of Refresh and Sparse Selection

ASPIRE’s refresh-layer and sparse-selection settings are robust across models and workloads, with interior refresh positions performing similarly and modest sensitivity to page size and estimator constants.

  • Refresh-layer position: Refresh acceptance is low at layer 0, high and flat from N/4 through N−2, and lower again at the final layer across all six model–workload combinations.Layer 0 achieves 38.7–76.4% acceptance, while N−1 is consistently worse than N−2.
  • Refresh-layer overhead: Refresh-layer overhead is 1.03%–1.06% of Qwen3-8B cycle time at 16k–32k context and falls at or below measurement noise on the other two models.The acceptance improvement offsets this cost by allowing longer drafts before staleness forces verification.
  • Page size: Sparse-selection page size changes throughput by under 9% across a 64× granularity range, with page size 16 near the peak.At a fixed 512-token budget, coarse pages devote more budget to recent context and leave fewer slots for score-selected pages.
  • Acceptance-estimator constants: Neither acceptance-estimator sweep collapses performance, and neither the smoothing weight nor prior produces a monotone throughput trend.Their settings affect estimate convergence, which interacts with run-to-run generation variation.

C.2 Long-Context Tasks

On long-context workloads, intra-draft refresh mitigates sparse-context staleness as draft length grows, substantially improving acceptance and supporting longer speculation.

  • The long-context evaluation covers detailed statistics for LB[16k-18k], LB-v2[30k-40k], and LB-v2[80k-100k].These results are reported across the corresponding model–dataset combinations and include the long-context benchmark tables.
  • Without refresh, the sparse context selected at the last verification boundary becomes increasingly stale as draft length increases.The refresh variant updates the sparse context at every draft step to counteract this degradation.
  • At draft length 10, refresh raises acceptance from 60.3% to 76.6%, closing a 16.4 percentage-point gap on Qwen3-8B LB[16k-18k].At draft length 1, both variants are near 95%; the gap widens as drafting proceeds.
Loading 2609.17943v1…