TriForce: Aceleración sin pérdidas en la generación de secuencias largas mediante decodificación especulativa jerárquica

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

April 18, 2024

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

cs.AI

Resumen

Con el despliegue generalizado de modelos de lenguaje de gran escala (LLMs) en la generación de contenido extenso recientemente, ha surgido una demanda creciente de soporte eficiente para inferencia de secuencias largas. Sin embargo, la caché de clave-valor (KV), que se almacena para evitar la re-computación, ha surgido como un cuello de botella crítico al crecer linealmente en tamaño con la longitud de la secuencia. Debido a la naturaleza auto-regresiva de los LLMs, toda la caché KV se cargará para cada token generado, lo que resulta en una baja utilización de los núcleos computacionales y una alta latencia. Aunque se han propuesto varios métodos de compresión para la caché KV para aliviar este problema, estos sufren de una degradación en la calidad de la generación. Introducimos TriForce, un sistema jerárquico de decodificación especulativa que es escalable para la generación de secuencias largas. Este enfoque aprovecha los pesos originales del modelo y una caché KV dispersa dinámica mediante recuperación como un modelo de borrador, que sirve como una capa intermedia en la jerarquía y es especulada aún más por un modelo más pequeño para reducir su latencia de borrador. TriForce no solo facilita impresionantes aceleraciones para Llama2-7B-128K, logrando hasta 2.31 veces en una GPU A100, sino que también muestra escalabilidad al manejar contextos aún más largos. Para la configuración de descarga en dos GPUs RTX 4090, TriForce logra 0.108s/token—solo la mitad de lento que la línea base auto-regresiva en una A100, que alcanza 7.78 veces en nuestro sistema de descarga optimizado. Además, TriForce realiza 4.86 veces mejor que DeepSpeed-Zero-Inference en una sola GPU RTX 4090. La robustez de TriForce se destaca por su rendimiento consistentemente sobresaliente en varias temperaturas. El código está disponible en 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.