MagicDec: Quebrando o Tradeoff entre Latência e Throughput para Geração de Contexto Longo com Decodificação Especulativa

MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding

August 20, 2024

Autores: Jian Chen, Vashisth Tiwari, Ranajoy Sadhukhan, Zhuoming Chen, Jinyuan Shi, Ian En-Hsu Yen, Beidi Chen

cs.AI

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) tornaram-se mais comuns em aplicações de longo contexto, como chatbots interativos, análise de documentos e fluxos de agentes, mas é desafiador atender a solicitações de longo contexto com baixa latência e alta taxa de transferência. A Decodificação Especulativa (SD) é uma técnica amplamente utilizada para reduzir a latência sem comprometer o desempenho, mas a sabedoria convencional sugere que sua eficácia é limitada a tamanhos de lote pequenos. No MagicDec, demonstramos que, surpreendentemente, a SD pode alcançar aceleração mesmo para um regime de inferência de alta taxa de transferência para sequências moderadas a longas. Mais interessante ainda, uma estratégia inteligente de rascunho pode obter uma melhor aceleração com o aumento do tamanho do lote com base em nossa análise rigorosa. O MagicDec identifica primeiro as mudanças de gargalo com o aumento do tamanho do lote e da sequência, e utiliza essas percepções para implantar a decodificação especulativa de forma mais eficaz para inferência de alta taxa de transferência. Em seguida, ele aproveita modelos de rascunho com cache KV esparsa para lidar com o gargalo KV que aumenta tanto com o comprimento da sequência quanto com o tamanho do lote.

English

Large Language Models (LLMs) have become more prevalent in long-context applications such as interactive chatbots, document analysis, and agent workflows, but it is challenging to serve long-context requests with low latency and high throughput. Speculative decoding (SD) is a widely used technique to reduce latency without sacrificing performance but the conventional wisdom suggests that its efficacy is limited to small batch sizes. In MagicDec, we show that surprisingly SD can achieve speedup even for a high throughput inference regime for moderate to long sequences. More interestingly, an intelligent drafting strategy can achieve better speedup with increasing batch size based on our rigorous analysis. MagicDec first identifies the bottleneck shifts with increasing batch size and sequence length, and uses these insights to deploy speculative decoding more effectively for high throughput inference. Then, it leverages draft models with sparse KV cache to address the KV bottleneck that scales with both sequence length and batch size.