PaperlayerView exact revision on arXiv

Source-linked AI summary

TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding

Hanshi Sun, Zhuoming Chen, Xinyu Yang, Yuandong Tian, Beidi Chen

arXiv:2404.11912v3cs.CLcs.LG

TL;DR

Long-context generation is bottlenecked by repeatedly loading growing KV caches and model weights, while eviction-based compression can lose information. TRIFORCE uses retrieval-based sparse KV drafting with hierarchical speculation, reporting speedups up to 2.31× on an A100 and 7.78× with offloading on two RTX 4090 GPUs. Its supported scope includes robust acceptance across temperatures and efficient long-context serving, but eviction-based alternatives remain limited for comprehensive context understanding.

  • Problem

    Long-context inference repeatedly loads growing KV caches and model weights, while existing eviction methods can irreversibly discard information and harm contextual coherence.

  • Method

    TRIFORCE uses original model weights with retrieval-selected partial KV cache as a draft and hierarchically speculates it with a lightweight model using StreamingLLM cache.

  • Results

    TRIFORCE achieves up to 2.31× speedup on an A100 and 7.78× on two RTX 4090 GPUs with offloading, reaching 0.108 s/token.

  • Takeaways & Limitations

    TRIFORCE provides scalable, lossless long-context serving while addressing both KV-cache and model-weight bottlenecks.

  • Takeaways & Limitations

    Eviction-based alternatives remain poorly suited to tasks requiring comprehensive long-context understanding because they discard tokens outside their retention policies.

Abstract

from arXiv · show

With large language models (LLMs) widely deployed in long content generation recently, there has emerged an increasing demand for efficient long-sequence inference support. However, key-value (KV) cache, which is stored to avoid re-computation, has emerged as a critical bottleneck by growing linearly in size with the sequence length. Due to the auto-regressive nature of LLMs, the entire KV cache will be loaded for every generated token, resulting in low utilization of computational cores and high latency. While various compression methods for KV cache have been proposed to alleviate this issue, they suffer from degradation in generation quality. We introduce TriForce, a hierarchical speculative decoding system that is scalable for long sequence generation. This approach leverages the original model weights and dynamic sparse KV cache via retrieval as a draft model, which serves as an intermediate layer in the hierarchy and is further speculated by a smaller model to reduce its drafting latency. TriForce not only facilitates impressive speedups for Llama2-7B-128K, achieving up to 2.31$\times$ on an A100 GPU but also showcases scalability in handling even longer contexts. For the offloading setting on two RTX 4090 GPUs, TriForce achieves 0.108s/token$\unicode{x2014}$only half as slow as the auto-regressive baseline on an A100, which attains 7.78$\times$ on our optimized offloading system. Additionally, TriForce performs 4.86$\times$ than DeepSpeed-Zero-Inference on a single RTX 4090 GPU. TriForce's robustness is highlighted by its consistently outstanding performance across various temperatures. The code is available at https://github.com/Infini-AI-Lab/TriForce.

1 Introduction

TRIFORCE addresses long-context inference bottlenecks from KV-cache growth and model-weight loading with retrieval-based drafting and hierarchical speculation. It reports substantial speedups while preserving lossless speculative decoding and robustness across temperatures.

  • Motivation: KV cache becomes a dominant long-context bottleneck because autoregressive generation repeatedly loads it together with model parameters for each token.Eviction-based compression can reduce memory but may lose information and degrade contextual coherence.
  • Motivation: Long-context speculative decoding must maintain high acceptance rates and low-cost drafting without training a context-matched draft model.The paper identifies these requirements as central technical challenges.
  • Method: TRIFORCE combines original model weights with a small dynamic KV-cache draft and further speculates that draft using a lightweight model with StreamingLLM cache.The hierarchy targets KV-cache and model-weight bottlenecks sequentially.
  • Method: Retrieval-based drafting selects relevant KV-cache chunks from the maintained full cache, avoiding the irreversible information loss of eviction-based methods.The approach scores chunks using attention to their average keys and gathers a fixed-budget partial cache.
  • Empirical results: 7.78× speedup on two RTX 4090 GPUs with offloading reaches 0.108 s/token, while single-A100 inference achieves up to 2.31× speedup.TRIFORCE also serves Llama2-13B with 128K contexts at 0.226 s/token and is 4.86× faster than DeepSpeed-Zero-Inference on one RTX 4090 GPU.
  • Empirical results: TRIFORCE remains robust across temperatures, maintaining an acceptance rate above 0.9 at temperature 1.0, and achieves a theoretical 13.1× upper bound.It also reports 1.9× speedup for batch size six with 19K-token contexts per sample.

2 Background

The background contrasts speculative decoding and KV-cache efficiency methods for long-sequence inference. Existing eviction strategies reduce cache size but discard context and therefore may not support comprehensive long-context understanding, while quantization is orthogonal to TRIFORCE.

  • Speculative decoding: Speculative decoding accelerates autoregressive inference by letting a lightweight draft model predict tokens that the target model verifies in parallel while preserving the output distribution.Its motivation follows from the cost of repeatedly loading model weights and KV cache.
  • KV-cache eviction: StreamingLLM retains attention sinks and recent KV states to stabilize attention for effectively unbounded input streams without fine-tuning.Its cache policy prioritizes sink tokens together with recent context.
  • KV-cache eviction: H2O dynamically removes less critical KV pairs using cumulative attention scores to maintain a fixed cache size.The method uses a greedy policy based on local statistics.
  • KV-cache eviction: Eviction methods do not increase the model’s context window and may fail on tasks requiring comprehensive long-context understanding because they discard tokens outside their policies.They retain recent tokens plus attention sinks or heavy hitters rather than the full context.
  • KV-cache quantization: KV-cache quantization reduces activation bit width and memory consumption, but it is orthogonal to TRIFORCE’s approach.

3 Observation

The paper’s observations identify attention sparsity and contextual locality as foundations for efficient long-context speculation. Sparse partial caches can preserve most attention information, while cache selections can be reused across adjacent decoding steps.

  • Leveraging Attention Sparsity: Over 96% of attention score can be recovered with only 4K tokens from a 120K context across almost all layers.This establishes substantial attention sparsity in Llama2-7B-128K.
  • Leveraging Attention Sparsity: A 1K-token KV-cache budget theoretically reaches 97.6% acceptance with Top-K selection, while H2O and StreamingLLM exceed 90.5% in practical settings.The theoretical result is an upper bound, and the practical methods maintain the full cache for the initial two layers only in the illustration.
  • Exploiting Contextual Locality: Attention information needed by adjacent generated tokens tends to be similar, creating contextual locality that gradually diminishes as the sequence evolves.The locality observation is measured by reusing top-4K indices selected from the last prefilled token.
  • Exploiting Contextual Locality: A single draft-cache construction can support multiple decoding steps, amortizing construction latency and improving drafting efficiency.New cache entries can replace less significant ones, with rebuilding scheduled periodically or adaptively.

4 TRIFORCE

TRIFORCE combines retrieval-based KV-cache drafting with hierarchical speculation to address the memory bottlenecks of long-context inference while preserving the target model’s output distribution.

  • 4 TRIFORCE: TRIFORCE uses retrieval-based drafting and hierarchical speculation to address KV-cache and model-weight bottlenecks in long-context decoding.The system uses a retrieval cache for the target model and a lightweight model with StreamingLLM cache for initial speculation.
  • 4 TRIFORCE: StreamingLLM and H2O can lose acceptance on needle retrieval because evicted KV pairs cannot be restored for long-term contextual dependencies.Table 1 evaluates acceptance rates across tasks using a 120K context and 4K budget, with Top-K as the theoretical upper bound.
  • 4 TRIFORCE: Retrieval-based drafting scores chunked KV-cache keys against the query and selects a fixed-budget set of relevant chunks.The approach maintains the full cache, computes attention against average keys within chunks, and retrieves the highest-scoring chunks.
  • 4 TRIFORCE: Standard speculative decoding preserves the original output distribution during the final verification phase.The correctness argument identifies the final phase as identical to standard speculative decoding.
  • 4 TRIFORCE: The algorithm alternates lightweight-model drafting against the retrieval cache with target-model self-verification using the full cache.Each iteration performs two speculation stages before updating the retrieval and StreamingLLM caches.

5 Empirical Evaluation

TRIFORCE accelerates long-context generation across on-chip, offloading, batching, and longer-context settings, with performance supported by a 4K cache budget and hierarchical design.

  • 5 Empirical Evaluation: TRIFORCE scales to longer contexts, with theoretical speedup reaching 13.1× and remaining acceptance above 0.9 at temperature 1.0.The analysis attributes scalability to high acceptance rates and a growing latency gap between draft and target models.
  • 5 Empirical Evaluation: Using only 3% of the KV cache, TRIFORCE achieves the best reported long-sequence speedup against REST and Skipping Layers in the cited comparison.TRIFORCE retrieves dynamically from the KV cache, whereas REST uses an external predefined datastore; Skipping Layers uses 68% of the KV cache.
  • 5 Empirical Evaluation: A 4K KV-cache budget is optimal because acceptance rises toward 1.0 before larger budgets add drafting latency without comparable benefit.Smaller chunks can overfit individual tokens, whereas larger chunks dilute scores and reduce selection flexibility.
  • 5 Empirical Evaluation: Tree-based speculation can potentially improve end-to-end speedup by using additional speculation budgets.The study estimates this compatibility for Sequoia and Independent Sequences on Llama2-7B-128K.

6 Conclusion

TRIFORCE is a hierarchical speculative decoding system for efficient long-context LLM serving, combining attention sparsity and contextual locality to mitigate KV-cache and model-weight bottlenecks.

  • 6 Conclusion: TRIFORCE achieves up to 2.31× speedup on an A100 and 7.78× on two RTX 4090s with offloading, reaching 0.108 s/token.These results support its potential for efficient long-sequence generation.

A System Implementation

The implementation combines CUDA graphs and FlashAttention with a predetermined cache budget to reduce speculative-decoding overhead.

  • A System Implementation: A predetermined cache budget enables PyTorch CUDA graphs to minimize kernel-launch overhead during speculative decoding.FlashAttention is used to accelerate attention operations.

B.1 TRIFORCE’s Scalability for Longer Inputs

TRIFORCE is evaluated for longer inputs using KV-cache offloading on a single L40 GPU with LWM-Text models and PG-19.

  • B.1 TRIFORCE’s Scalability for Longer Inputs: TRIFORCE’s longer-input scalability is evaluated with KV-cache offloading on a single L40 GPU.The experiments use LWM-Text models and the PG-19 dataset.
  • B.1 TRIFORCE’s Scalability for Longer Inputs: The longer-input experiment targets contexts that exceed the practical on-GPU KV-cache capacity.Offloading the KV cache to CPU memory is used for these longer inputs.
  • B.1 TRIFORCE’s Scalability for Longer Inputs: The supplied table is identified as documenting TRIFORCE’s scalability for longer inputs.

B.2 TRIFORCE’s Scalability for Longer Outputs

TRIFORCE remains relatively stable as output sequences lengthen, with retrieval-cache reconstruction and reuse supporting scalability across varying output lengths.

  • B.2 TRIFORCE’s Scalability for Longer Outputs: TRIFORCE’s performance remains relatively stable for output sequences extended up to 2K tokens.The experiment uses Llama2-7B-128K on an A100 GPU with PG-19.
  • B.2 TRIFORCE’s Scalability for Longer Outputs: Fixed-interval reconstruction keeps the retrieval cache aligned with the evolving context during longer generation.
  • B.2 TRIFORCE’s Scalability for Longer Outputs: Reusing the retrieved cache across multiple decoding steps amortizes retrieval overhead.The results are presented as evidence of scalability and robustness across varying output lengths.

B.3 Ablation of γ1, γ2

The γ1, γ2 ablation shows that the first speculation phase has low acceptance because the lightweight and retrieval-cache models differ, with γ1=2 and γ2=6 optimal.

  • B.3 Ablation of γ1, γ2: Table 8 reports the joint ablation of γ1 and γ2.
Loading 2404.11912v3…