TriForce: Verlustfreie Beschleunigung der langen Sequenzgenerierung mit hierarchischem spekulativem Decodieren

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

April 18, 2024

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

cs.AI

Zusammenfassung

Mit der weit verbreiteten Verwendung großer Sprachmodelle (LLMs) zur Generierung langer Inhalte ist in letzter Zeit eine zunehmende Nachfrage nach effizienter Unterstützung für Inferenz von langen Sequenzen entstanden. Allerdings hat sich der Schlüssel-Wert (KV)-Cache, der gespeichert wird, um erneute Berechnungen zu vermeiden, als kritischer Engpass erwiesen, da er mit der Sequenzlänge linear anwächst. Aufgrund der autoregressiven Natur von LLMs wird der gesamte KV-Cache für jedes generierte Token geladen, was zu einer geringen Auslastung der Rechenkerne und hohen Latenzzeiten führt. Obwohl verschiedene Kompressionsmethoden für den KV-Cache vorgeschlagen wurden, um dieses Problem zu mildern, leiden sie unter einer Verschlechterung der Generierungsqualität. Wir stellen TriForce vor, ein hierarchisches spekulatives Dekodierungssystem, das für die Generierung langer Sequenzen skalierbar ist. Dieser Ansatz nutzt die ursprünglichen Modellgewichte und den dynamischen spärlichen KV-Cache über Abruf als Entwurfsmodell, das als Zwischenschicht in der Hierarchie dient und weiter von einem kleineren Modell spekuliert wird, um die Entwurfslatenz zu reduzieren. TriForce ermöglicht nicht nur beeindruckende Beschleunigungen für Llama2-7B-128K, die auf einer A100 GPU bis zu 2,31-mal erreichen, sondern zeigt auch Skalierbarkeit bei der Behandlung noch längerer Kontexte. Für die Auslagerungseinstellung auf zwei RTX 4090 GPUs erreicht TriForce 0,108s/Token - nur halb so langsam wie die autoregressive Basislinie auf einer A100, die auf unserem optimierten Auslagerungssystem 7,78-mal erreicht. Darüber hinaus übertrifft TriForce DeepSpeed-Zero-Inference auf einer einzelnen RTX 4090 GPU um das 4,86-fache. Die Robustheit von TriForce wird durch seine durchweg herausragende Leistung bei verschiedenen Temperaturen hervorgehoben. Der Code ist verfügbar unter 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.