TriForce: Aceleração sem Perdas na Geração de Sequências Longas com Decodificação Especulativa Hierárquica

Resumo

Com a ampla implantação de modelos de linguagem de grande escala (LLMs) na geração de conteúdo longo recentemente, surgiu uma demanda crescente por suporte eficiente para inferência de sequências longas. No entanto, o cache de chave-valor (KV), que é armazenado para evitar recomputação, tornou-se um gargalo crítico ao crescer linearmente em tamanho com o comprimento da sequência. Devido à natureza autorregressiva dos LLMs, todo o cache KV será carregado para cada token gerado, resultando em baixa utilização dos núcleos computacionais e alta latência. Embora vários métodos de compressão para o cache KV tenham sido propostos para aliviar esse problema, eles sofrem com degradação na qualidade da geração. Apresentamos o TriForce, um sistema hierárquico de decodificação especulativa que é escalável para geração de sequências longas. Essa abordagem aproveita os pesos originais do modelo e o cache KV esparso dinâmico via recuperação como um modelo de rascunho, que serve como uma camada intermediária na hierarquia e é ainda especulado por um modelo menor para reduzir sua latência de rascunho. O TriForce não apenas facilita acelerações impressionantes para o Llama2-7B-128K, alcançando até 2,31 vezes em uma GPU A100, mas também demonstra escalabilidade ao lidar com contextos ainda mais longos. Para a configuração de offloading em duas GPUs RTX 4090, o TriForce alcança 0,108s/token—apenas metade da lentidão da linha de base autorregressiva em uma A100, que atinge 7,78 vezes em nosso sistema de offloading otimizado. Além disso, o TriForce performa 4,86 vezes melhor que o DeepSpeed-Zero-Inference em uma única GPU RTX 4090. A robustez do TriForce é destacada por seu desempenho consistentemente excepcional em várias temperaturas. O código está disponível em 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: Aceleração sem Perdas na Geração de Sequências Longas com Decodificação Especulativa Hierárquica

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

Resumo

Support