Resumo Técnico: HYDRA – Um Framework de DSE para Chiplets Heterogêneos no Atendimento de Cargas de Trabalho Híbridas de LLM Dinâmicas

1. Declaração do Problema

A rápida evolução dos Grandes Modelos de Linguagem (LLMs) introduziu desafios significativos para a aceleração de hardware, particularmente com a emergência de arquiteturas híbridas Transformer–Mamba (ex: Jamba, Nemotron-H). Esses modelos combinam o processamento de contexto longo escalável de Modelos de Espaço de Estados (SSMs) com a forte modelagem de dependência dos Transformers, criando padrões de computação e comunicação heterogêneos.

Os sistemas de hardware existentes enfrentam duas limitações primárias ao atender essas cargas de trabalho híbridas:

1. Heterogeneidade Arquitetural: Blocos SSM dependem cada vez mais de operações elemento a elemento e estados recorrentes, enquanto blocos Transformer são dominados por multiplicações de matrizes e tráfego de memória do cache KV. Além disso, as fases de prefill (preenchimento) e decode (decodificação) exibem demandas de recursos fundamentalmente diferentes. Por exemplo, o prefill é intensivo em computação para ambos, mas durante a decodificação, os blocos Transformer sofrem um overhead de largura de banda de memória que aumenta linearmente devido ao acesso ao cache KV, enquanto os blocos Mamba mantêm uma demanda constante.
2. Complexidade do Espaço de Design: A integração de diversas unidades de processamento em um único chip monolítico está se tornando proibitiva em termos de custo. Embora as arquiteturas baseadas em chiplets ofereçam uma solução escalável ao modularizar computação, memória e rede, o espaço de design para esses sistemas é proibitivamente grande. Ele abrange configurações macroestruturais estáticas (composição de chiplets, posicionamento, largura de banda entre chiplets) e políticas de execução dinâmica (batching, escalonamento). A exploração exaustiva desse espaço é inviável, muitas vezes exigindo semanas de tempo de simulação.
3. Imprevisibilidade em Tempo de Execução: O atendimento de LLM multi-inquilino envolve chegadas de requisições assíncronas e comportamentos dependentes da entrada (comprimentos variáveis de prefill/decode). Isso cria demandas de recursos altamente dinâmicas que modelos de desempenho estáticos e mapeamento de tarefas offline não conseguem capturar com precisão.

2. Metodologia: O Framework HYDRA

O HYDRA é um framework abrangente de Exploração de Espaço de Design (DSE) projetado para otimizar conjuntamente a arquitetura e as políticas de tempo de execução para LLMs híbridos em sistemas de chiplets 2.5D heterogêneos. Ele integra quatro dimensões principais:

A. Modelagem de Sistema e Restrições

O HYDRA recebe especificações de carga de trabalho (modelos, conjuntos de dados, rastros de requisição) e restrições de sistema (biblioteca de chiplets, topologia de Rede-no-Interposer (NoI), orçamento de área de die) como entradas.

• Biblioteca de Chiplets: Inclui aceleradores especializados para as fases de prefill/decode tanto para blocos de Atenção (Transformer) quanto para blocos Mamba (SSM), além de chiplets de memória HBM.
• Modelagem de NoI: Leva em conta a topologia (ex: mesh), largura de banda die-para-die (D2D) e latência, reconhecendo que uma maior largura de banda consome área de die que poderia ser usada para computação ou SRAM.

B. Estratégias de Otimização

O HYDRA emprega um processo de otimização de três estágios:

1. Posicionamento de Chiplets Consciente de Comunicação:
   • Passo 1 (Posicionamento de HBM): Os chiplets HBM são posicionados no anel externo do interposer. Eles são particionados em grupos para blocos Mamba e de Atenção para equilibrar a capacidade de memória e minimizar a interferência.
   • Passo 2 (Posicionamento de Computação): Os chiplets de computação são posicionados para minimizar a distância de Manhattan até seus parceiros de memória dominantes (grupos HBM), reduzindo os custos de comunicação entre chiplets e o tráfego de NoI.
2. Batching de Requisições Dinâmico:
   • Substitui o batching estático por uma política de admissão contínua. As requisições são admitidas em um lote ativo assim que os recursos (computação, memória, orçamento de tokens) estão disponíveis, sem esperar pela sincronização do lote.
   • Utiliza blocos de memória virtualizados (semelhante ao vLLM) para desacoplar a alocação de memória da admissão, permitindo o crescimento dinâmico do cache KV ou dos estados recorrentes.
   • Implementa preempção ao nível de requisição para evitar condições de falta de memória (out-of-memory), priorizando requisições mais próximas da conclusão.
3. Escalonamento de Tarefas Elástico:
   • Gera um mapeamento de pipeline espacial inicial baseado na afinidade tarefa-chiplet (otimização de prefill vs. decode).
   • Um escalonador elástico de tempo de execução monitora a utilização dos chiplets e reequilibra dinamicamente a carga. Se um chiplet estiver subutilizado devido a paradas de pipeline transientes ou chegadas explosivas, as tarefas são migradas para ele, desde que o overhead de comunicação permaneça limitado.

C. Estimativa de Desempenho Rápida (Baseada em Markov)

Para superar o custo proibitivo da simulação completa orientada a eventos para cada configuração candidata, o HYDRA introduz um estimador de desempenho baseado em Markov:

• Abstração: Modela o sistema como uma Cadeia de Markov de Tempo Contínuo (CTMC), onde os estados representam o número de chiplets alocados para as fases de prefill vs. decode.
• Modelo de Fila Fluida: Em vez de rastrear eventos individuais, utiliza um modelo fluido para aproximar o trabalho agregado e as transições de estado ao longo do tempo.
• Saída: Prevê o throughput (vazão) do sistema e o Tempo-até-o-Primeiro-Token (TTFT) calculando as probabilidades de estado estacionário e as capacidades de serviço. Este estimador serve como um mecanismo de poda rápida para identificar regiões promissoras do espaço de design, que são então validadas por simulação detalhada.

3. Principais Contribuições

1. Framework HYDRA: Um framework de DSE especificamente para cargas de trabalho de LLM híbrido em sistemas de chiplets 2.5D heterogêneos, disponível como uma ferramenta de código aberto.
2. Estratégia de Co-Design: Uma abordagem unificada combinando posicionamento consciente de comunicação, batching dinâmico e escalonamento elástico para melhorar a utilização de computação e memória em ambientes multi-inquilino.
3. Estimador Baseado em Markov: Uma técnica de modelagem de desempenho inovadora que captura a dinâmica de tempo de execução (escalonamento elástico, batching dinâmico) com alta precisão (similaridade de cosseno média de 0,9 com simulação completa) enquanto reduz o tempo de exploração de dias para minutos.
4. Avaliação Abrangente: Experimentos extensos em modelos puramente Transformer (LLaMA3), puramente SSM (Mamba2) e híbridos (Nemotron-H) através de diversos conjuntos de dados (ArXiv, Chat, LongWriter, etc.).

4. Resultados Experimentais

Os autores avaliaram o HYDRA contra baselines de última geração e configurações estáticas em 12 pares de modelo-conjunto de dados.

• Ganhos de Desempenho: Comparado aos baselines, o HYDRA alcança um melhoria média de 1,55× no throughput e um TTFT 43,7% menor. Em casos ótimos, os ganhos de throughput chegaram a até 2,3×.
• Eficiência do Estimador: O estimador baseado em Markov reduziu o tempo de DSE de dias/semanas para minutos (ex: 8 dias e 8 horas para simulação exaustiva vs. 15 minutos para o HYDRA). Ele recuperou 93% do throughput máximo encontrado pela simulação exaustiva, em média.
• Impacto do Posicionamento: O posicionamento consciente de comunicação melhorou o throughput em 3%–100% e a razão throughput-para-TTFT em comparação com o posicionamento round-robin ou aleatório, principalmente ao reduzir a utilização e a latência do link NoI.
• Adaptabilidade de Tempo de Execução: A combinação de escalonamento elástico e batching dinâmico foi crítica. O escalonamento elástico melhorou o throughput em 3%–150% sobre o mapeamento estático, enquanto o batching dinâmico contribuiu com um ganho adicional de 15% no throughput e 23% de redução no TTFT, em média.
• Especialização vs. Generalidade: O estudo descobriu que, embora configurações especializadas em um único modelo gerem o melhor desempenho para aquele modelo específico, o "DSE consciente de mistura" (otimização para múltiplos modelos) melhora significamente o desempenho de pior caso em modelos não vistos (ex: melhorando o desempenho de transferência em até 2,84×) com perda mínima para o ótimo do modelo alvo.

5. Significância e Alegações

O artigo argumenta que co-projetar arquitetura e políticas de tempo de execução é crítico para o atendimento eficiente de LLMs em larga escala em sistemas de chiplets heterogêneos.

• Além de Modelos Estáticos: O trabalho demonstra que designs arquiteturais estáticos e escalonamento offline são insuficientes para a natureza dinâmica e multi-inquilino do atendimento moderno de LLM. A adaptabilidade em tempo de execução (escalonamento elástico e batching dinâmico) é essencial para explorar configurações macro-arquiteturais otimizadas.
• Localidade de Comunicação: Os resultados destacam que a localidade de comunicação (minimizar a distância entre chiplets) é frequentemente um gargalo mais dominante do que o provisionamento bruto de computação em sistemas baseados em chiplets.
• Exploração Escalável: O HYDRA fornece um caminho prático para explorar o vasto espaço de design de sistemas de chiplets heterogêneos, que anteriormente era intratável devido aos custos de simulação. Ao modelar com precisão o acoplamento entre batching, escalonamento e contenção de recursos, o framework permite a identificação de designs Pareto-ótimos que equilibram throughput e latência.

Os autores concluem que o HYDRA oferece um framework escalável para explorar os sistemas de LLM multi-inquilino de próxima geração, observando que trabalhos futuros podem estender o framework para incluir modelos de potência, térmica e confiabilidade.