Pooling Engram Conditional Memory in Large Language Models using CXL

Este artigo propõe o uso de pools de memória CXL para armazenar a memória condicional Engram em Grandes Modelos de Linguagem, integrando-a ao SGLang para oferecer uma solução de armazenamento escalável e econômica com desempenho de inferência próximo ao da memória DRAM.

O essencial

Imagine que você está tentando ensinar um robô superinteligente (um Modelo de Linguagem Grande, ou LLM) a conversar como um humano. Para isso, o robô precisa de duas coisas principais: cérebro (para pensar e raciocinar) e memória (para lembrar de fatos, palavras e regras).

Até agora, os robôs tentavam guardar tudo o que precisavam saber dentro da própria "cabeça" (na memória rápida do computador, chamada DRAM). O problema? À medida que o robô fica mais inteligente, ele precisa de uma biblioteca de conhecimento gigantesca. Guardar tudo na "cabeça" é caro, ocupa muito espaço e deixa o robô lento, como tentar correr com uma mochila cheia de pedras.

Aqui entra a ideia genial deste artigo: o "Engram" e o "CXL".

1. O Problema: A Mochila Cheia de Pedras

Os modelos de linguagem atuais usam uma técnica chamada "Mixture-of-Experts" (MoE). Pense nisso como ter vários especialistas em uma sala. Quando você faz uma pergunta, o robô acorda o especialista certo para responder.

O novo conceito, Engram, é como adicionar uma "biblioteca de consulta rápida" a esses especialistas. Em vez de o robô ter que pensar muito para lembrar uma palavra comum (como "cachorro" ou "sol"), ele apenas olha na biblioteca.

• O problema: Essa biblioteca é enorme (centenas de gigabytes). Se tentarmos colocá-la toda na memória rápida do computador (DRAM), o computador fica caro demais e lento.
• O desafio: A biblioteca é usada de um jeito estranho. O robô não lê a biblioteca inteira de uma vez. Ele dá "piscadinhas" rápidas e espalhadas: pega um pedacinho aqui, outro ali, outro acolá. São milhares de pequenos pedaços de informação.

2. A Solução: O "CXL" (O Elevador de Carga Inteligente)

Antes, para guardar essa biblioteca longe do computador (em um servidor separado), usávamos redes comuns (como RDMA). Imagine que tentar pegar um livro dessa biblioteca usando a rede antiga fosse como pedir um livro na biblioteca e ter que esperar o carteiro entregar, abrir o envelope, ler o endereço e só então pegar o livro. É lento para coisas pequenas e rápidas.

O artigo propõe usar uma tecnologia chamada CXL (Compute Express Link).

• A Analogia: Imagine que o CXL é como um elevador de carga ultra-rápido e direto que conecta a sala do robô (o computador) diretamente à biblioteca (o servidor de memória), sem passar pelo carteiro ou pelo porteiro.
• Com o CXL, o robô pode "puxar" esses pequenos pedaços de memória espalhados quase tão rápido quanto se estivessem dentro da própria cabeça dele.

3. Como Funciona na Prática?

Os pesquisadores criaram um sistema onde:

1. A Biblioteca Compartilhada: Em vez de cada robô ter sua própria biblioteca gigante (o que custaria uma fortuna), vários robôs em um data center compartilham uma única biblioteca central gigante guardada em memórias CXL.
2. O Truque de Velocidade: Eles criaram um "porteiro" especial (software otimizado) que sabe exatamente como pegar esses pedacinhos de memória espalhados sem perder tempo.
3. O Resultado: O robô consegue acessar essa biblioteca externa quase na mesma velocidade que acessaria a memória interna.

4. Por que isso é um "Milagre" de Economia?

Pense em um prédio de escritórios:

• Sem CXL: Cada escritório precisa ter seu próprio arquivo morto gigante no subsolo. Se você tem 100 escritórios, você gasta 100 vezes o dinheiro em arquivos e prateleiras.
• Com CXL: Você constrói um único arquivo morto gigante no subsolo, acessível por um elevador super-rápido. Todos os 100 escritórios usam o mesmo arquivo.
  • Economia: Você economiza uma quantidade enorme de dinheiro em hardware (memória).
  • Velocidade: O elevador (CXL) é tão rápido que ninguém sente que o arquivo está longe.

Resumo da Ópera

Este artigo diz: "Não precisamos mais colocar todo o conhecimento do robô dentro da cabeça dele, o que é caro e pesado. Podemos colocar essa memória em um lugar compartilhado e acessível via uma 'estrada de alta velocidade' (CXL). Assim, os robôs ficam mais baratos de construir, mais baratos de manter, e continuam pensando tão rápido quanto antes."

É como se a humanidade descobrisse que, em vez de cada pessoa ter que decorar a Enciclopédia Britânica inteira, poderíamos ter uma Enciclopédia Centralizada onde todos podem consultar qualquer página em milissegundos, sem precisar carregar o livro inteiro na mochila.

Resumo técnico

Aqui está um resumo técnico detalhado do artigo "Pooling Engram Conditional Memory in Large Language Models using CXL", apresentado em português:

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) modernos, especialmente aqueles baseados em arquiteturas Mixture-of-Experts (MoE), enfrentam desafios significativos na gestão de conhecimento estático.

• Ineficiência na Busca de Conhecimento: Modelos tradicionais simulam a recuperação de conhecimento vocabular através de computação intensiva, desperdiçando recursos.
• Memória Engram: A arquitetura Engram (proposta pela DeepSeek) introduz "Memória Condicional" para desacoplar o armazenamento de conhecimento (tabelas de N-Gram) da computação dinâmica. Embora isso melhore a escalabilidade, as tabelas de embeddings do Engram podem exigir centenas de GBs de memória.
• Padrão de Acesso Específico: O acesso ao Engram é caracterizado por ser esparsamente distribuído, minimamente granular (pequenos pacotes de dados) e apenas de leitura.
• Limitações das Soluções Atuais:
  • DRAM Local: Armazenar tabelas massivas em DRAM local é proibitivamente caro e limita a escalabilidade.
  • Pool de Memória via RDMA: Soluções existentes que usam RDMA (como o Mooncake) sofrem com alta latência e baixa eficiência ao lidar com pacotes pequenos e discretos (devido à sobrecarga da pilha de rede e semântica baseada em mensagens), o que é crítico para o Engram, que exige baixíssima latência para não atrasar a inferência.

2. Metodologia

Os autores propõem o uso de Compute Express Link (CXL) para criar um pool de memória compartilhado dedicado ao armazenamento das tabelas de embeddings do Engram.

• Arquitetura do Sistema:
  • Implementação de um pool de memória CXL centralizado conectado a múltiplos nós de computação via um switch CXL (ex: chip XConn XC50256).
  • O pool permite que CPUs e GPUs de diferentes servidores acessem a mesma memória física com semântica de load/store nativa, bypassando a pilha de rede tradicional.
• Otimização de Acesso (Routines):
  • Mapeamento DAX: A memória CXL é mapeada como um dispositivo de acesso direto (DAX) no sistema operacional, permitindo que a CPU acesse a memória remota como se fosse local.
  • CPU para CXL: Uso de estratégias de leitura paralela multithread (OpenMP) para saturar a largura de banda.
  • GPU para CXL: Implementação de kernels CUDA personalizados para transferências diretas Peer-to-Peer (P2P). Isso evita a sobrecarga de milhares de chamadas cudaMemcpy discretas, fundindo as operações em um único kernel de grade larga para maximizar a eficiência da banda PCIe.
• Integração com SGLang: O sistema foi integrado ao framework de inferência SGLang. O mecanismo de prefetch é assíncrono: as embeddings são buscadas do pool CXL e transferidas para a memória da GPU antes que a camada do Transformer que as utiliza precise delas, sobrepondo a latência de memória com a computação das camadas anteriores.

3. Principais Contribuições

1. Primeiro Sistema CXL para Engram: Apresentação do primeiro sistema capaz de descarregar parâmetros do Engram para um pool de memória baseado em CXL.
2. Análise Comparativa RDMA vs. CXL: Demonstração de que o CXL é superior ao RDMA para o padrão de acesso do Engram. Enquanto o RDMA sofre com latência em pacotes pequenos (devido à semântica de rede), o CXL oferece acesso de granularidade de linha de cache com latência próxima à da DRAM local.
3. Implementação e Otimização: Desenvolvimento de rotinas de acesso otimizadas (CPU e GPU) e integração bem-sucedida em um framework de produção (SGLang).
4. Análise de Custo e Escalabilidade: Prova de que a solução é economicamente viável para clusters grandes, reduzindo drasticamente o custo de capital (CapEx) em comparação com a replicação de DRAM em cada nó.

4. Resultados Experimentais

Os testes foram realizados em um cluster com processadores Intel Xeon e GPUs NVIDIA L20, utilizando um switch CXL e memória Samsung de 256 GB.

• Latência:
  • A latência de leitura do CXL para CPU é comparável à da DRAM local.
  • A leitura CXL para GPU (via kernel personalizado) apresenta uma latência ligeiramente maior, mas ainda dentro de uma margem aceitável para não causar stall na inferência.
  • O CXL supera o RDMA em ordens de magnitude para o padrão de acesso esparsa do Engram.
• Throughput de Inferência (End-to-End):
  • Ao integrar o Engram no SGLang, o sistema com pool CXL atingiu desempenho quase idêntico ao da DRAM local.
  • Exemplo (Qwen3-4B): O baseline foi 6183.9 tokens/s. Com Engram em DRAM: 5683.7 tokens/s. Com Engram em CXL: 5614.4 tokens/s (uma perda marginal de ~1.2% em relação à DRAM).
• Escalabilidade: O sistema manteve o desempenho ao aumentar o paralelismo de dados (DP) e o número de nós, indicando robustez em cenários multi-usuário.
• Análise de Custo:
  • Para modelos pequenos (ex: 100B parâmetros em 2 nós), o CXL pode ser mais caro devido ao custo fixo do switch e adaptadores.
  • No entanto, à medida que o número de nós e o tamanho do modelo aumentam (ex: 400B em 16 nós), o CXL oferece economias massivas (até $166.040 em comparação com a solução local), pois elimina a necessidade de replicar centenas de GBs de DRAM em cada servidor.

5. Significado e Impacto

Este trabalho é fundamental para o futuro da infraestrutura de LLMs por várias razões:

• Viabilidade Econômica: Permite que LLMs de próxima geração utilizem tabelas de conhecimento massivas (Engram) sem exigir que cada servidor tenha memória DRAM suficiente para armazená-las localmente, reduzindo drasticamente o custo de implantação.
• Desempenho sem Compromisso: Demonstra que é possível externalizar memória para camadas inferiores (CXL) sem sacrificar a latência crítica da inferência, superando as limitações de soluções baseadas em rede (RDMA) para cargas de trabalho esparsas.
• Nova Arquitetura de Memória: Estabelece um precedente para o uso de CXL não apenas para KV Cache, mas para componentes de memória condicional e embeddings estáticos, pavimentando o caminho para LLMs mais eficientes e escaláveis.

Em resumo, o artigo valida que o CXL é a tecnologia ideal para gerenciar a memória condicional do Engram, oferecendo o equilíbrio perfeito entre custo, escalabilidade e baixa latência necessária para a próxima geração de modelos de linguagem.