TriForce: Accelerazione Senza Perdite della Generazione di Sequenze Lunghe con Decodifica Speculativa Gerarchica

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

April 18, 2024

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

cs.AI

Abstract

Con l'ampio utilizzo di modelli linguistici di grandi dimensioni (LLM) nella generazione di contenuti lunghi di recente, è emersa una crescente domanda di supporto efficiente per l'inferenza su sequenze lunghe. Tuttavia, la cache chiave-valore (KV), che viene memorizzata per evitare il ricalcolo, è diventata un collo di bottiglia critico, crescendo linearmente in dimensioni con la lunghezza della sequenza. A causa della natura auto-regressiva degli LLM, l'intera cache KV viene caricata per ogni token generato, risultando in una bassa utilizzazione dei core computazionali e un'elevata latenza. Sebbene siano stati proposti vari metodi di compressione per la cache KV per alleviare questo problema, essi soffrono di un degrado nella qualità della generazione. Introduciamo TriForce, un sistema di decodifica speculativa gerarchica scalabile per la generazione di sequenze lunghe. Questo approccio sfrutta i pesi originali del modello e una cache KV sparsa dinamica tramite recupero come modello di bozza, che funge da livello intermedio nella gerarchia ed è ulteriormente speculato da un modello più piccolo per ridurre la sua latenza di bozza. TriForce non solo facilita impressionanti accelerazioni per Llama2-7B-128K, raggiungendo fino a 2.31 volte su una GPU A100, ma dimostra anche scalabilità nella gestione di contesti ancora più lunghi. Per l'impostazione di offloading su due GPU RTX 4090, TriForce raggiunge 0.108s/token—solo la metà della lentezza rispetto alla baseline auto-regressiva su una A100, che ottiene 7.78 volte sul nostro sistema di offloading ottimizzato. Inoltre, TriForce performa 4.86 volte meglio di DeepSpeed-Zero-Inference su una singola GPU RTX 4090. La robustezza di TriForce è evidenziata dalla sua performance costantemente eccezionale a varie temperature. Il codice è disponibile su https://github.com/Infini-AI-Lab/TriForce.

English

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 to 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.31times 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/tokenx2014only half as slow as the auto-regressive baseline on an A100, which attains 7.78times on our optimized offloading system. Additionally, TriForce performs 4.86times 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.