SqueezeAttention: Gerenciamento 2D do KV-Cache na Inferência de LLM via Orçamento Ótimo por Camada

Otimizar o cache Key-Value (KV) do Large Language Model (LLM) tem sido considerado crítico para economizar o custo da inferência. A maioria dos algoritmos existentes de compressão de KV-cache tentou esparsificar a sequência de tokens aproveitando a importância diferente dos tokens. Neste trabalho, descobrimos que, ao identificar a importância das camadas de atenção, podemos otimizar o KV-cache conjuntamente a partir de duas dimensões. Baseado em nossas observações sobre a importância por camada na inferência, propomos o SqueezeAttention para otimizar precisamente a alocação do orçamento do KV-cache entre as camadas em tempo real e, em seguida, incorporar três algoritmos representativos de esparsificação de tokens para comprimir o KV-cache para cada camada com seu próprio orçamento. Ao otimizar o KV-cache tanto nas dimensões da sequência quanto da camada, o SqueezeAttention alcança cerca de 30% a 70% de redução de memória e até 2,2 vezes de melhoria na taxa de processamento em uma ampla gama de LLMs e benchmarks. O código está disponível em https://github.com/hetailang/SqueezeAttention.