Move the Query, Not the Cache: Characterizing Cross-Instance Latent Attention Redistribution Across GPU Fabrics

Resumo Técnico: Mova a Consulta, Não o Cache

Declaração do Problema
À medida que os Grandes Modelos de Linguagem (LLMs) escalam, as cargas de trabalho de serviço dependem cada vez mais de mecanismos de Atenção de Cabeça Múltipla Latente (MLA) e atenção esparsa. Nessas arquiteturas, o cache de Chave-Valor (KV) é comprimido em vetores latentes estreitos ou selecionado via indexadores esparsos, criando uma "assimetria de bytes": uma única linha de consulta (ex: $\approx 1$ KB) é ordens de magnitude menor do que o bloco de cache KV ao qual ela atende (ex: megabytes).

Quando o corpus de interesse (um repositório de código canônico, documentos jurídicos ou uma base de conhecimento empresarial compartilhada) excede a Memória de Alta Largura de Banda (HBM) de uma única GPU, ele deve ser particionado entre múltiplas instâncias. Isso necessita de atenção entre instâncias, onde uma consulta de decodificação em uma GPU deve acessar blocos de KV residentes em outra. O reflexo predominante da indústria é mover o cache (buscar): puxar os blocos de KV remotos para o requerente, inseri-los no cache local e realizar a atenção localmente. No entanto, essa abordagem incorre em um custo significativo de "adaptação de posição" (re-rotacionar codificações posicionais em cache para o novo deslocamento) e uma grande transferência de dados. O artigo postula que, sob MLA e seleção esparsa, a aritmética pode inverter: pode ser mais barato rotear a consulta para o detentor, computar a atenção parcial lá e mesclar o pequeno resultado de volta.

Metodologia
Os autores caracterizam esse compromisso em um cluster real de H100 multi-nó conectado via NDR-200 InfiniBand com IBGDA (RDMA iniciado pelo dispositivo) habilitado. Sua metodologia envolve:

1. Decomposição de Custo: Eles modelam a redistribuição entre instâncias como uma função de custo aditiva: $T_{redist} = T_{probe} + T_{transfer} + T_{compute} + T_{return} + T_{merge}$.
2. Comparação de Primitivas: Eles comparam três primitivas:
   • Fetch (Busca): Puxar o bloco de KV inteiro, adaptar posições (inserção/splice) e atender localmente.
   • Route (Roteamento): Enviar a consulta para o detentor, computar a atenção parcial e mesclar o resultado via softmax online.
   • Local: Recomputar o bloco do zero (re-prefill).
3. Caracterização de Hardware: Eles medem a latência e a largura de banda do IBGDA para payloads com formato de atenção (lotes de consultas na escala de kilobytes) e comparam com caminhos de proxy de CPU. Eles também analisam o impacto da topologia (mesma folha vs. travessia de spine), tipo de fabric (NVLink, PCIe, InfiniBand) e concorrência.
4. Verificação de Correção: Eles verificam que a atenção parcial roteada, quando mesclada via softmax online, reproduz os resultados de atenção de instância única dentro das margens de ruído de bf16, mesmo para cenários de seleção esparsa.

Principais Contribuições

• Modelo de Custo Sensível à Topologia: Os autores derivam um modelo de custo de forma fechada para atenção entre instâncias que se ajusta às medições reais de IBGDA em H100 com $\approx 7\%$ de Erro Médio Absoluto Percentual (MAPE) para despacho em lote ($M_q \ge 512$). O modelo identifica que o custo é dominado por uma latência de probe fixa ($\approx 16 \mu s$) e uma largura de banda efetiva independente do payload ($\approx 25$ GB/s para despacho de bloco único), em vez da largura de banda de pico do link.
• Caracterização do Regime de RDMA Iniciado pelo Dispositivo (IBGDA): Diferente do trabalho anterior sobre despacho de MoE (onde caminhos de proxy de CPU às vezes vencem para mensagens minúsculas), este artigo estabelece que, para payloads de atenção ($\approx 2$ KB), o RDMA iniciado pelo dispositivo é superior, reduzindo a latência por busca em 40% em comparação com caminhos mediados pelo host.
• Predicado de Seleção de Forma Fechada: Eles fornecem uma regra de decisão para escalonadores escolherem entre route, fetch ou local com base no formato da requisição ($M_q$), tamanho do bloco ($c_t$) e características do fabric.
  • Route vence no decode (pequeno $M_q$) porque evita o custo de adaptação de posição de $\approx 3$ ms exigido pelo fetch e evita o custo de computação do local.
  • Fetch só se torna viável quando se amortiza o custo de adaptação (splice) ao longo de muitos passos locais subsequentes.
• Análise de Dispersão e Seleção: Para atenção esparsa (ex: DeepSeek-V3.2, GLM-5.1), o artigo demonstra que o fetch torna-se um gather disperso através de múltiplos detentores, incorrendo em uma penalidade que cresce com o número de detentores. O route mantém um custo constante porque a consulta é enviada uma vez, e o detentor atende aos blocos dispersos no local sem re-rotação.
• Cotovelos de Capacidade do Detentor (Holder Capacity Elbows): Eles identificam gargalos específicos para o lado do detentor:
  • Cotovelo de Cópia (Copy Elbow): Um pool de $K=8$ streams CUDA é ideal para preparar dados para requisições concorrentes.
  • Cotovelo de Computação (Compute Elbow): Um detentor pode servir $\approx 8$ requerentes roteados simultâneos com aumento negligível de latência antes da saturação de computação ocorrer.

Resultos

• Dominância de Latência: Em tamanhos típicos de lote de decode ($M_q = 256$), um ciclo de ida e volta (round trip) roteado custa 31–48 $\mu s$ através de vários fabrics (NVLink, PCIe, IBGDA). Isso é 60–100$\times$ mais rápido que o custo de adaptação de $\approx 3$ ms de buscar um bloco.
• Eficiência de Bytes no Fio (Wire Byte Efficiency): O roteamento move $\ge 76\%$ menos bytes no fio do que o fetch durante o decode em lotes.
• Independência de Fabric: No decode, o despacho é limitado pelo despacho (dispatch-bound) (limitado pela capacidade de um único bloco de thread emitir transferências, $\approx 18\text{--}25$ GB/s) em vez de limitado pela largura de banda (bandwidth-bound). Consequentemente, um link NVLink 4.0 de 900 GB/s performa de forma semelhante a um link InfiniBand de 25 GB/s entre nós para consultas de bloco único; a largura de banda de pico só importa para grandes tamanhos de lote.
• Resiliência à Congestão: A primitiva de roteamento é robusta à congestão de rede. A latência permanece estável até que o link esteja totalmente saturado ($K=3$ fluxos concorrentes e, mesmo assim, permanece uma ordem de magnitude mais rápida que o fetch.
• Advertência de Overhead do Host: Embora o roteamento vença em bytes no fio e latência de transporte, os autores observam que, em seu protótipo atual, o overhead do host (cliente Python, lógica de despacho) domina o tempo de primeiro token (time-to-first-token). Em tamanhos de lote pequenos, um simples fetch de "bytes-de-volta" pode atualmente vencer o tempo total devido ao menor overhead do host, embora três reduções de implementação específicas (resposta colapsada, amortização de computação, loteamento) sejam identificadas para fechar essa lacuna.

Significância e Alegações
O artigo alega que o reflexo de "mover o cache" é suboptimal para o regime emergente de atenção comprimida e esparsa. Ao inverter a aritmética, rotear a consulta torna-se a primitiva dominante para atenção entre instâncias, particularmente para cargas de trabalho de agentes onde muitos sub-agentes consultam um corpus compartilhado e imutável.

Os autores enfatizam que sua contribuição não é um novo algoritmo de roteamento (que já existe em trabalhos como o DistAttention), mas uma caracterização e um modelo de custo que valida o caso econômico para o roteamento sob condições específicas de hardware e arquitetura (MLA, IBGDA, seleção esparsa). Eles argumentam que essa assimetria de bytes é a propriedade chave que motiva o uso de RDMA iniciado pelo dispositivo para atenção. O artigo conclui que, embora o roteamento seja teoricamente superior, realizar seus benefícios totais de latência de ponta a ponta em uma pilha de serviço de produção requer abordar os overheads do lado do host, uma tarefa deixada para esforços futuros de engenharia. O modelo de custo e o predicado fornecidos servem como artefatos reutilizáveis para projetistas de sistemas tomarem decisões de escalonamento sensíveis à topologia sem a necessidade de re-perfilamento a cada nova configuração.