Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator

O artigo apresenta o Helios, um acelerador de hardware-software baseado em 3D-DRAM e emparelhamento híbrido que supera as limitações de designs de processamento próximo à memória existentes ao introduzir um fluxo de execução de atenção distribuído e uma alocação de cache KV espacialmente consciente, resultando em significativos ganhos de desempenho e eficiência energética para o atendimento dinâmico de modelos de linguagem grandes (LLMs).

O essencial

Imagine que você está tentando organizar uma festa gigante onde os convidados (os dados da Inteligência Artificial) chegam de forma imprevisível: às vezes vêm em grupos pequenos, às vezes em multidões, e cada um fica por um tempo diferente.

O papel que você leu descreve um novo sistema chamado Helios, criado para ser o "gerente de festas" perfeito para os modelos de linguagem grandes (como o ChatGPT), que estão mudando a forma como usamos computadores.

Aqui está a explicação simples, usando analogias do dia a dia:

1. O Problema: A Festa Caótica

Atualmente, os computadores que rodam essas IAs (como as GPUs) são como cozinhas de restaurante muito rápidas, mas com um problema: elas são ótimas para cozinhar pratos complexos (cálculos), mas lentas quando precisam buscar ingredientes que estão longe na despensa (memória).

Além disso, os pedidos chegam de forma bagunçada:

• Variação de tamanho: Alguns pedidos são curtos (uma pergunta simples), outros são longos (um livro inteiro).
• Variação de chegada: Às vezes chega 1 pedido, às vezes chegam 100 de uma vez.

Os sistemas antigos tentam resolver isso alocando um "quartinho" fixo para cada pedido na memória. O problema é que, se o pedido for pequeno, o quartinho fica metade vazio (desperdício). Se o pedido for grande, ele não cabe e precisa esperar. Isso cria um gargalo e deixa a festa lenta.

2. A Solução: O Helios (O Gerente Inteligente)

Os pesquisadores criaram o Helios, que é como uma nova cozinha com uma despensa integrada e um gerente superinteligente.

A. A Tecnologia de "Colagem" (Hybrid Bonding)

Imagine que a memória (onde os dados ficam) e o processador (onde os cálculos acontecem) são dois andares de um prédio. Nos computadores antigos, eles estão em prédios separados e você precisa pegar um elevador lento para ir de um ao outro.

O Helios usa uma tecnologia chamada Hybrid Bonding. É como se eles colassem o processador e a memória um em cima do outro, criando "elevadores" super-rápidos e curtos entre os andares. Isso significa que os dados viajam quase instantaneamente, sem perder tempo.

B. O Gerenciamento de Memória Dinâmica (O "Bloco de Notas" Flexível)

Aqui está a mágica principal.

• Sistemas Antigos: Davam um "quarto inteiro" para cada pedido, não importa o tamanho. Se o pedido fosse pequeno, o resto do quarto ficava vazio. Se fosse grande, o pedido tinha que esperar um quarto inteiro ficar livre.
• Helios: Usa uma abordagem de blocos. Imagine que a memória é um grande tapete de peças de Lego. O Helios não reserva um quarto inteiro; ele reserva apenas as peças de Lego necessárias para o tamanho do pedido.
  • Se o pedido é pequeno, ele usa poucas peças.
  • Se o pedido cresce, ele pega mais peças de onde estiverem livres.
  • Se o pedido termina, ele devolve as peças imediatamente para serem usadas por outros.

Isso elimina o desperdício e permite que o sistema atenda muito mais pessoas ao mesmo tempo.

C. A Dança dos Dados (Atenção Distribuída)

Quando a IA "lê" um texto, ela precisa conectar todas as palavras entre si. Nos sistemas antigos, essa tarefa era feita de forma rígida, onde cada máquina fazia uma parte fixa. Se uma máquina terminasse rápido, ela ficava parada esperando a mais lenta.

O Helios divide essa tarefa como uma dança coordenada.

• Ele quebra o texto em pequenos pedaços (blocos).
• Distribui esses pedaços entre todos os "trabalhadores" (processadores) disponíveis.
• Se um trabalhador termina rápido, ele pega outro pedaço imediatamente.
• Eles se comunicam de forma super-rápida (como se estivessem sussurrando entre si) para juntar as peças do quebra-cabeça.

3. Os Resultados: Por que isso importa?

O Helios não é apenas uma melhoria pequena; é uma revolução na eficiência:

• Velocidade: É 3,25 vezes mais rápido que os melhores sistemas atuais. É como trocar de ir de bicicleta para ir de avião.
• Economia de Energia: É 3,36 vezes mais eficiente energeticamente. Isso significa que você pode fazer o mesmo trabalho gastando muito menos eletricidade (e gerando menos calor).
• Estabilidade: Mesmo quando a "festa" fica lotada (muitos usuários ao mesmo tempo), o Helios mantém o ritmo, enquanto os sistemas antigos começam a engasgar e ficar lentos.

Resumo Final

Pense no Helios como a evolução de um sistema de transporte público.

• Antes: Ônibus grandes que saem apenas quando estão cheios, deixando muita gente parada na parada, e que param em todas as estações, mesmo que ninguém desça.
• Helios: Um sistema de "uber" dinâmico. O veículo se ajusta ao tamanho do passageiro, pega o caminho mais curto, e se um passageiro desce, o espaço é imediatamente usado por outro.

Essa inovação permite que as IAs do futuro sejam mais rápidas, mais baratas de rodar e capazes de atender milhões de pessoas simultaneamente sem travar.

Resumo técnico

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) são amplamente utilizados em serviços de geração online, onde múltiplas instâncias de modelos lidam com cargas de trabalho altamente dinâmicas. Essas cargas são caracterizadas por:

• Flutuações na taxa de chegada de solicitações: Variações de até uma ordem de magnitude entre picos e vales.
• Comprimentos de solicitação variáveis: Janelas de contexto que podem variar de algumas centenas a dezenas de milhares de tokens (ex: 10K tokens).
• Natureza mista da computação: O serviço envolve operadores intensivos em computação (fase de prefill) e operadores intensivos em memória (fase de decoding, especialmente as camadas de atenção).

Limitações das Soluções Existentes (NMP e GPUs):
As abordagens atuais de Near-Memory Processing (NMP) e GPUs enfrentam gargalos críticos ao lidar com essa dinâmica:

1. Gerenciamento de KV Cache de Granularidade Grossa: Designs NMP existentes alocam o cache completo de um head de atenção para um módulo de memória fixo (ex: um cubo HBM). Isso causa:
   • Desequilíbrio de Carga: Se o número de heads de atenção for menor que o número de módulos NMP, recursos ficam ociosos.
   • Fragmentação de Memória: Solicitações de tamanhos variados não preenchem eficientemente os módulos fixos, desperdiçando capacidade e limitando o número de solicitações simultâneas que podem ser atendidas.
2. Fluxo de Execução de Atenção Inflexível: A maioria dos designs NMP executa a atenção como duas GEMMs (multiplicações de matriz) independentes e estáticas. Eles exigem que todo o cache KV de um head resida no mesmo módulo durante o cálculo, impedindo a partição fina e dinâmica necessária para gerenciar caches de tamanhos variáveis.

---

2. Metodologia: Helios

O artigo propõe o Helios, um acelerador de serviço de LLM baseado em Hybrid Bonding (ligação híbrida 3D) que realiza um co-design completo de arquitetura e sistema para superar as limitações acima.

A. Arquitetura de Hardware (HB-Device)

• Tecnologia de Ligação Híbrida (Hybrid Bonding): Utiliza a fusão direta Cu-Cu para empilhar até 4 camadas de DRAM sobre uma die de lógica. Isso oferece largura de banda extremamente alta (110.000 I/Os/mm²) e baixa latência/energia (0,66 pJ/bit).
• Array de Processadores (PEs): A die lógica contém uma matriz de PEs (ex: 4x4) conectados por uma rede em malha (NoC). Cada PE tem acesso direto à sua memória local de DRAM, mas pode comunicar-se com outros PEs.
• Unidades Especializadas: Cada PE possui unidades vetoriais, matrizes (para GEMM), uma unidade de Softmax Online e uma unidade de redução, permitindo a execução eficiente de atenção em blocos.

B. Execução de Operadores (Software/Algoritmo)

• Atenção em Blocos Distribuídos (Distributed Tiled Attention): Em vez de calcular a atenção sobre a sequência inteira em um único módulo, o Helios divide o cache KV em blocos finos (ex: 64 tokens).
  • A computação de atenção é realizada iterativamente sobre esses blocos distribuídos entre todos os PEs.
  • Utiliza Softmax Online para permitir a acumulação de resultados parciais sem precisar carregar toda a sequência na memória.
• Otimização de Fluxo: O design sobrepõe operações vetoriais (como softmax) com operações matriciais (GEMM) e utiliza double buffering para esconder a latência de transferência de dados.

C. Gerenciamento de Sistema e Alocação Espacial

• Alocação de Cache KV Consciente de Espaço (Spatially-Aware KV Cache Allocation):
  • O sistema não aloca blocos aleatoriamente. Um algoritmo (Algoritmo 2) distribui os blocos de cache entre os PEs considerando dois fatores:
    1. Balanceamento de Carga Computacional: Distribui tokens para equilibrar o trabalho entre todos os PEs.
    2. Minimização de Tráfego de Rede: Posiciona os "últimos blocos" (que precisam ser acessados frequentemente para projeção de K/V) em PEs centrais ou simétricos para reduzir o número de saltos na rede NoC durante a comunicação.
• Gerenciamento Dinâmico: O cache é alocado sob demanda conforme os tokens são gerados, eliminando a necessidade de reservar espaço para o tamanho máximo de contexto (janela de contexto), o que reduz drasticamente o desperdício de memória.

---

3. Contribuições Principais

1. Análise de Limitações: Identificação profunda de como a abstração de memória não uniforme dos designs NMP atuais conflita com os algoritmos de gerenciamento de cache dinâmico baseados em GPU.
2. Arquitetura Helios: Proposta de um acelerador baseado em Hybrid Bonding que suporta nativamente a atenção em blocos distribuídos e o gerenciamento de cache KV dinâmico.
3. Co-design de Sistema: Desenvolvimento de mecanismos de alocação espacial que equilibram a carga computacional e minimizam a sobrecarga de comunicação entre PEs.
4. Eficiência de Execução: Criação de fluxos de execução intra-PE e primitivas de comunicação inter-PE otimizadas para variáveis de atenção (MHA, GQA, MQA, MLA) e modelos MoE.

---

4. Resultados Experimentais

O Helios foi avaliado em comparação com GPUs (A100), designs NMP existentes (AttAcc, Duplex) e outras abordagens de ligação híbrida, utilizando modelos como LLaMA3 70B, Mixtral, Qwen e DeepSeek.

• Aceleração de Desempenho:
  • O Helios alcançou um speedup geométrico de 3,25x em comparação com os melhores designs de GPU/NMP existentes.
  • Em comparação com um design HB sem atenção em blocos (HB-non-Tiled), houve um speedup de 1,67x.
• Eficiência Energética:
  • Melhoria de 3,36x na eficiência energética (geomean).
• Desempenho de Serviço (Serving):
  • Redução significativa na latência entre tokens (Time-Between-Tokens - TBT):
    • P50: Até 72% de melhoria (redução de latência).
    • P99: Até 76% de melhoria.
• Capacidade de Serviço:
  • O Helios suporta taxas de chegada de solicitações muito mais altas antes de saturar, devido à eliminação da fragmentação de memória.
  • A taxa de fragmentação de memória em designs NMP tradicionais foi de 4,67x a 23,10x maior do que no Helios.
• Balanceamento de Carga:
  • O desequilíbrio de carga entre PEs foi mantido abaixo de 13,3% no Helios, enquanto designs concorrentes atingiram até 100% de desequilíbrio em cargas baixas.

---

5. Significado e Impacto

O trabalho do Helios é significativo porque:

1. Preenche a Lacuna Arquitetural: Resolve o conflito fundamental entre a natureza dinâmica do serviço de LLM (tamanhos de solicitação variáveis) e a arquitetura de memória não uniforme dos aceleradores NMP.
2. Viabiliza NMP para Serviços Reais: Demonstra que o processamento próximo à memória (NMP) pode ser eficaz não apenas para inferência estática, mas para ambientes de produção com alta dinâmica, superando as limitações de designs anteriores que eram rígidos demais.
3. Escalabilidade e Eficiência: Ao utilizar Hybrid Bonding e co-design de hardware/software, o Helios oferece uma rota escalável para acelerar modelos de linguagem cada vez maiores e mais complexos (incluindo MoE e MLA) com menor consumo de energia e menor latência.
4. Direção Futura: Estabelece um novo paradigma para aceleradores de IA, onde o gerenciamento de memória e a execução de operadores são tratados de forma integrada e espacialmente consciente, em vez de serem tratamentos separados.

Em resumo, o Helios representa um avanço crucial na otimização de hardware para LLMs, transformando o NMP de uma solução de nicho para uma arquitetura viável e superior para o serviço massivo de modelos de linguagem na nuvem.