Resumo Técnico: PTStore (Prefix Tensor Store)

Declaração do Problema

As cargas de trabalho de inferência de Grandes Modelos de Linguagem (LLMs) tornaram-se o domínio principal no consumo de energia e demanda de recursos em centros de dados de computação de alto desempenho (HPC), superando o treinamento. A inferência de LLM consiste em duas fases: prefill (processamento do prompt de entrada em paralelo) e decode (geração de tokens sequencialmente). Para evitar a computação redundante de mecanismos de atenção, os sistemas utilizam um cache de Chave-Valor (KV) para armazenar resultados intermediários.

Embora os runtimes de última geração (ex: vLLM) otimizem o cache de KV dentro de uma única GPU ou nó, eles enfrentam limitações significativas em escala:

1. Falta de Reuso entre Nós: Sistemas existentes frequentemente falham em agregar memória entre nós de computação distribuídos. Se uma requisição em um nó compartilha um prefixo com uma requisição em outro nó, o segundo nó tipicamente recomputa o prefixo em vez de reutilizar os tensores em cache.
2. Gargalos de Metadados e Latência: Abordagens que tentam o cache distribuído (ex: LMCache, EvoStore) frequentemente sofrem com altos overheads de I/O devido ao acesso à memória remota ou sincronização complexa de metadados (ex: escalar o Radix-Attention além de um único nó).
3. Restrições de Memória: A memória individual de uma GPU é insuficiente para janelas de contexto grandes, e o descarregamento (offloading) para a memória do host ou SSDs introduz latência que anula os benefícios do cache.

O desafio central é permitir o reuso escalável e de baixa latência de prefixos de cache de KV através de um grande número de GPUs distribuídas sobre muitos nós de computação sem incorrer em overheads proibitivos de I/O ou metadados.

Metodologia: Arquitetura PTStore

O PTStore (Prefix Tensor Store) é um armazenamento de tensores distribuído e replicado, projetado para abordar essas limitações ao distribuir e replicar prefixos populares de cache de KV. O sistema emprega um modelo cliente-servidor onde cada nó de computação executa um servidor que agrega a memória local do host e SSDs para servir tanto clientes de GPU locais quanto remotos.

Princípios de Design Chave

1. Armazenamento Incremental de Tensores (Estrutura do tipo Trie):

   • Em vez de armazenar blocos de KV completos, o PTStore armazena diferenças incrementais (tensores) entre um novo objeto e o maior prefixo comum (LCP - Longest Common Prefix) de objetos armazenados anteriormente.
   • Isso permite que os prefixos cresçam sem redundância ao longo do tempo em direções divergentes, de forma semelhante a uma trie, mas implementada via granularidade de nível de tensor.
   • Metadados Consolidados: Para evitar o custoso travessia de trie distribuída, o PTStore utiliza uma estrutura de metadados plana. O metadado de cada objeto contém uma lista de IDs de tensores únicos. Uma operação de carga itera por esses IDs para verificar a existência local na cache de replicação; se ausente, ela busca esses itens remotamente do servidor "dono".

2. Cache Hierárquico Distribuído com Replicação:

   • Cache Proprietário (Owned Cache): Armazena os tensores incrementais pelos quais um servidor específico é responsável.
   • Cache de Replicação (Replication Cache): Armazena cópias de prefixos "quentes" (populares) localmente no servidor para melhorar a localidade de acesso.
   • Gestão de Trade-off: O sistema gerencia um limiar configurável entre os caches proprietários e de replicação. Ele prioriza o descarte de tensores replicados (que podem ser buscados novamente) sobre a expulsão de tensores proprietários (que exigem o descarregamento para o armazenamento mais lento) para equilibrar a velocidade de recuperação com a capacidade de armazenamento.

3. Expulsão Consciente do Padrão de Acesso:

   • O PTStore utiliza uma política de expulsão baseada em frequência (adaptada do GDSF) em vez de LRU (Least Recently Used), pois as estruturas de prefixo significam que tensores iniciais são acessados com mais frequência.
   • Ele considera o trade-off entre tamanho vs. frequência, garantindo que tensores pequenos e frequentes não desloquem tensores maiores e caros de buscar.

4. Consolidação Consciente de RDMA:

   • Para minimizar a dispersão, os incrementos anexados a um LCP são consolidados em uma única região contígua no servidor dono.
   • Operações de carga usam RDMA em lote (bulk RDMA) para buscar segmentos dispersos em paralelo via uma única RPC, evitando o overhead de copiar dados para uma região contígua antes da transferência.

Principais Contribuições

1. Princípios de Design: Um conjunto de princípios de alto nível para um repositório distribuído que integra armazenamento incremental de tensores, metadados consolidados e replicação de prefixo.
2. Protótipo PTStore: Um protótipo de pesquisa implementando esses princípios, apresentando uma API de baixo nível em C++ e uma interface Python para integração fluida com runtimes de LLM como o vLLM.
3. Validação de Desempenho: Experimentos extensos demonstrando reduções significativas no overhead de I/O e no tempo de execução de ponta a ponta em comparação com baselines de última geração.

Resultados Experimentais

Os autores avaliaram o PTStore no ambiente de teste ALCF Polaris HPC (560 nós, GPUs A100) usando duas cargas de trabalho de QA extrativa: WikiQA (contexto longo) e SQUAD (alto volume de perguntas). O LLM utilizado foi o Mistral-7B-instruct-V2.

Baselines

• vLLM Vanilla: vLLM padrão sem compartilhamento de prefixo entre requisições.
• vLLM Prefix: vLLM com compartilhamento de prefixo local (dentro de um nó).
• EvoStore: Um armazenamento de tensores distribuído que utiliza armazenamento incremental e RDMA, mas carece de replicação de prefixo local.
• PTStore: O sistema proposto com consciência distribuída e replicação local.

Constatações

• Escalabilidade Fraca (8–32 GPUs): O PTStore superou significativamente o EvoStore e o vLLM Prefix. Enquanto o EvoStore sofreu com altos overheads de I/O de RDMA ao buscar prefixos remotos, a replicação local do PTStore mitigou isso, resultando em uma "vantagem descolada" no Tempo para o Primeiro Token (TTFT).
• Escalabilidade de Comprimento de Sequência (1k–8k tokens):
  • Para sequências curtas (1k), o cache local do vLLM foi competitivo.
  • À medida que o comprimento da sequência aumentava, a vantagem do PTStore crescia. Em 8k tokens, o PTStore foi quase 2x mais rápido que o cache de prefixo do vLLM e 20% mais rápido que o EvoStore.
  • A lacuna de desempenho aumentou com contextos mais longos porque o custo de recomputação ou I/O remoto superou os benefícios do cache apenas local.
• Ganhos de Eficiência: Em conjuntos de dados de QA de passagens longas, o PTStore executou inferências 5 a 6 vezes mais eficientemente do que os baselines que não agregam memória entre nós e exigem a regeneração dos caches de KV.

Significância e Alegações

O artigo afirma que o PTStore aborda uma lacuna crítica no serviço escalável de inferência de LLM: a incapacidade dos sistemas atuais de reutilizar eficientemente os prefixos de cache de KV entre nós distribuídos. Ao combinar armazenamento incremental para minimizar a redundância, metadados consolidados para consultas rápidas e uma estratégia de replicação para otimizar a localidade, o PTStore permite:

• Expansão de ordens de magnitude no tamanho efetivo do cache de KV ao agregar a memória por todo o cluster.
• Redução significativa no TTFT, particularmente para cargas de trabalho de contexto longo, onde a recomputação é custosa.
• Escalabilidade que evita os gargalos de comunicação e os problemas de sincronização de metadados que assolam abordagens distribuídas anteriores.

Os autores posicionam o PTStore como um passo fundamental em direção à inferência de IA escalável, observando que trabalhos futuros focarão em equilíbrio dinâmico de memória, políticas de expulsão baseadas em ML e benchmarks mais amplos contra sistemas como LMCache e Mooncake em traces reais de conversação e completagem de código.