TriForce: Verliesvrije versnelling van lange sequentiegeneratie met hiërarchisch speculatief decoderen

Samenvatting

Met de wijdverspreide inzet van grote taalmodellen (LLMs) voor het genereren van lange inhoud is er recentelijk een toenemende vraag ontstaan naar efficiënte ondersteuning voor inferentie van lange sequenties. Echter, de key-value (KV) cache, die wordt opgeslagen om herberekening te voorkomen, is een kritieke bottleneck geworden doordat deze lineair in omvang toeneemt met de sequentielengte. Vanwege het auto-regressieve karakter van LLMs wordt de volledige KV cache geladen voor elk gegenereerd token, wat resulteert in een lage benutting van rekenkernen en hoge latentie. Hoewel verschillende compressiemethoden voor KV cache zijn voorgesteld om dit probleem te verlichten, leiden deze tot een afname in de generatiekwaliteit. Wij introduceren TriForce, een hiërarchisch speculatief decodering systeem dat schaalbaar is voor het genereren van lange sequenties. Deze aanpak maakt gebruik van de originele modelgewichten en een dynamische sparse KV cache via retrieval als een draft model, dat fungeert als een tussenlaag in de hiërarchie en verder wordt gespeculeerd door een kleiner model om de drafting latentie te verminderen. TriForce zorgt niet alleen voor indrukwekkende snelheidsverbeteringen voor Llama2-7B-128K, met een versnelling tot 2,31 keer op een A100 GPU, maar toont ook schaalbaarheid in het omgaan met nog langere contexten. Voor de offloading-instelling op twee RTX 4090 GPU's behaalt TriForce 0,108s/token—slechts half zo langzaam als de auto-regressieve baseline op een A100, die 7,78 keer haalt op ons geoptimaliseerde offloading systeem. Daarnaast presteert TriForce 4,86 keer beter dan DeepSpeed-Zero-Inference op een enkele RTX 4090 GPU. De robuustheid van TriForce wordt benadrukt door zijn consistent uitstekende prestaties bij verschillende temperaturen. De code is beschikbaar op 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.

TriForce: Verliesvrije versnelling van lange sequentiegeneratie met hiërarchisch speculatief decoderen

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

Samenvatting

Support