Multi-Head Latent Attention (MLA) e DeepSeekMoE: Compressão de KV Cache e Roteamento Esparso em Inferência

O consumo de memória do Key-Value (KV) Cache é o principal gargalo na inferência de LLMs. Em janelas longas, a memória para armazenar chaves e valores supera os próprios pesos do modelo.

A arquitetura Multi-Head Latent Attention (MLA) redefine essa dinâmica. Ela reduz o footprint de memória por token em 98,2% em relação ao Multi-Head Attention (MHA) tradicional.

Diferente do Grouped-Query Attention (GQA), o MLA não degrada a capacidade representacional. Ele preserva o espaço multidimensional de atenção com projeções de baixo posto e absorção matricial.

A arquitetura estreou no DeepSeek-V2 e amadureceu no DeepSeek-V3. Ela permite ao DeepSeek-R1 operar 671B de parâmetros. O modelo ativa apenas 37B por token em motores como vLLM e SGLang.

O gargalo de memória do KV Cache na inferência autoregressiva

A geração autoregressiva exige carregar todo o histórico de chaves e valores da memória HBM a cada novo token. Em janelas de 64k a 128k tokens, o KV Cache se torna o elemento mais custoso da GPU.

Em regime de batch unitário, a GPU opera limitada pela largura de banda de memória (memory-bound). O acelerador gasta a maior parte do ciclo transferindo tensores entre a HBM e a SRAM.

A fórmula exata para o tamanho do KV Cache por token em precisão FP16 (2 bytes por parâmetro) é:

S_KV = 2 × 2 × n_kv × d_h × L bytes por token

Na equação, n_kv é o número de cabeças de chave/valor. O termo d_h é a dimensão por cabeça e L é a quantidade de camadas.

Considere um modelo MHA com 128 cabeças (n_kv = 128), dimensão d_h = 128 e 61 camadas (L = 61). O consumo por camada é:

S_layer = 2 × 2 × 128 × 128 = 65.536 bytes = 64 KB por token

Para as 61 camadas do modelo, o custo consolidado atinge:

S_total = 61 × 64 KB = 3.904 KB por token ≈ 3,81 MB por token

Uma sequência de 64.000 tokens demanda cerca de 244 GB de VRAM apenas para a tabela de atenção. Em servidores com múltiplas requisições paralelas, esse volume esgota rapidamente a capacidade do hardware.

Por que o GQA degrada a capacidade representacional?

O Grouped-Query Attention (GQA) reduz a memória ao forçar múltiplas queries a compartilhar o mesmo conjunto de chaves e valores. No entanto, esse compartilhamento rígido prejudica a precisão do modelo em contextos longos.

Em testes como Needle In A Haystack e raciocínio multi-hop, aspectos semânticos distintos acabam forçados a convergir para as mesmas representações agrupadas.

Modelos como LLaMA 3 utilizam GQA com 8 grupos de chave/valor para mitigar a alocação de memória. Contudo, essa compressão impõe limites à capacidade do modelo de reter relações cruzadas finas.

O MLA supera essa restrição. Ele entrega uma compressão superior à do GQA mantendo a riqueza expressiva de um MHA de 128 cabeças completas.

A mecânica matemática do Multi-Head Latent Attention

O MLA comprime chaves e valores em um espaço latente de baixo posto (low-rank joint compression). O vetor de ativação h_t é comprimido em um vetor latente c_t^{KV} antes de ser gravado na memória.

A projeção descendente ocorre conforme a expressão:

c_t^{KV} = W^{DKV} × h_t

A matriz W^{DKV} projeta a ativação para uma dimensão compacta d_c = 512. Isso substitui as 16.384 dimensões exigidas por 128 cabeças convencionais.

As matrizes de conteúdo para cada cabeça i são calculadas por projeções ascendentes:

k_{t,i}^C = W^{UK}_i × c_t^{KV}, \quad v_{t,i}^C = W^{UV}_i × c_t^{KV}

Para viabilizar a fusão matricial na inferência, o MLA adota o Decoupled RoPE. Uma chave posicional dedicada k_t^R de 64 dimensões recebe o Rotary Embedding de maneira isolada:

k_t^R = RoPE(W^{KR} × h_t)

A query também é decomposta em parte de conteúdo e parte posicional:

q_{t,i}^C = W^{UQ}_i × c_t^Q, \quad q_{t,i}^R = RoPE(W^{QR}_i × c_t^Q)

Apenas o vetor latente c_t^{KV} (512 floats) e a chave posicional k_t^R (64 floats) são salvos no cache. O total por camada é de apenas 576 elementos.

Em precisão FP16, isso equivale a 1.152 bytes por token por camada. No modelo de 61 camadas, o consumo totaliza 68,63 KB por token.

Absorção matricial de pesos durante a geração de tokens

Durante a decodificação autoregressiva, o runtime não descompacta as chaves na memória HBM. O cálculo de atenção funde a matriz de projeção ascendente W^{UK}_i diretamente na query do token atual.

A pontuação de atenção combina as parcelas de conteúdo e posição:

Score_content = (q_{t,i}^C)^T × (W^{UK}_i × c_j^{KV})

Score_rope = (q_{t,i}^R)^T × k_j^R

Score_total = (Score_content + Score_rope) ÷ √(d_h + d_h^R)

Pela associatividade matricial, reorganiza-se o termo de conteúdo:

(q_{t,i}^C)^T × (W^{UK}_i × c_j^{KV}) = ((W^{UK}_i)^T × q_{t,i}^C)^T × c_j^{KV}

Define-se a query absorvida no espaço latente:

q_{t,i}^{absorbed} = (W^{UK}_i)^T × q_{t,i}^C

# Execução do Kernel MLA com absorção matricial no vLLM
import torch

def mla_decode_step(q_content, q_rope, kv_cache_latent, kv_cache_rope, W_uk, W_uv, scale):
    # 1. Absorve os pesos da chave na query do token atual na SRAM
    q_absorbed = torch.einsum("bhd,hdl->bhl", q_content, W_uk)
    
    # 2. Produto escalar diretamente contra o cache latente comprimido da HBM
    scores_content = torch.einsum("bhl,bsl->bhs", q_absorbed, kv_cache_latent)
    scores_rope = torch.einsum("bhd,bsd->bhs", q_rope, kv_cache_rope)
    
    attn_weights = torch.softmax((scores_content + scores_rope) * scale, dim=-1)
    
    # 3. Agregação no espaço latente antes da projeção final de valores
    latent_out = torch.einsum("bhs,bsl->bhl", attn_weights, kv_cache_latent)
    output = torch.einsum("bhl,hdl->bhd", latent_out, W_uv)
    
    return output

O vetor q_{t,i}^{absorbed} é processado nos registradores locais da GPU para o token ativo. O loop de atenção lê unicamente o vetor comprimido c_j^{KV} da HBM.

Da mesma forma, a projeção de valores W^{UV}_i é aplicada apenas uma vez na saída condensada. Essa etapa elimina a necessidade de carregar matrizes intermediárias completas.

DeepSeekMoE: granularidade fina e balanceamento dinâmico sem perda auxiliar

O DeepSeekMoE expande a especialização particionando a camada Feed-Forward em experts granulares. No DeepSeek-V3, a arquitetura adota 256 routed experts e 1 shared expert fixo, ativando 8 routed experts por token.

O shared expert processa todos os tokens incondicionalmente. Ele absorve padrões gerais da linguagem, permitindo que os routed experts se concentrem em raciocínio especializado.

Essa divisão estrutural reduz a sobreposição funcional entre experts. Cada sub-rede roteada desenvolve competências refinadas em lógica formal, programação ou tradução técnica.

[Token Input h_t]
       │
       ├───► [Shared Expert Fixo (Conhecimento Geral)] ──────────┐
       │                                                         │
       └───► [Roteador Top-8 com Bias Adaptativo b_i]            │
                   │                                             │
                   ├──► [Routed Expert 014 (Lógica)] ────────────┼──► [Soma Ponderada] ──► Output
                   ├──► [Routed Expert 089 (Código)] ────────────┤
                   └──► [Routed Expert 201 (Causal)] ────────────┘

Para manter o cluster equilibrado sem degradar a capacidade do modelo, o DeepSeekMoE utiliza o Auxiliary-Loss-Free Load Balancing:

• Roteamento com Bias: O roteador calcula as afinidades com um viés aditivo: s_{i,t} = Softmax(TopK(h_t^T × e_i + b_i)).
• Ajuste Dinâmico: Um controlador proporcional ajusta b_i a cada passo com base na carga observada dos experts.
• Gradiente Puro: O termo b_i atua no despacho e não injeta gradientes nos pesos, preservando o objetivo de treino original.

Esse mecanismo atinge mais de 99% de equilíbrio de carga computacional entre os nós do cluster. Ele evita tanto o estrangulamento de nós quanto o colapso de rotas.

Multi-Token Prediction e pipeline DualPipe em escala

O DeepSeek-V3 incorpora o método de Multi-Token Prediction (MTP). A arquitetura treina módulos encadeados para prever dois ou mais tokens futuros simultaneamente a cada passo.

Durante o pré-treinamento, o MTP reforça o sinal de supervisão densa. Ele obriga as camadas profundas a planejar dependências contextuais com maior antecedência sintática.

Em motores como vLLM e SGLang, o MTP opera como decodificação especulativa nativa. A engine também se integra ao TensorRT-LLM para execução em clusters. O sistema emite múltiplos tokens candidatos em um único ciclo de inferência.

Para viabilizar treinamento e inferência em precisão mista FP8 com eficiência máxima, o framework introduz o pipeline DualPipe. Ele sobrepõe computação e comunicação entre forward e backward passes.

A tabela a seguir consolida as especificações operacionais de inferência em clusters modernos:

A combinação de MLA, balanceamento sem perda auxiliar e quantização nativa em FP8 consolida um padrão eficiente para sistemas inteligentes de larga escala.

Perguntas Frequentes sobre MLA e DeepSeekMoE

Qual é a diferença fundamental entre GQA e MLA?

O GQA compartilha chaves e valores entre queries, reduzindo a capacidade do modelo em contextos longos. O MLA projeta chaves e valores para um vetor latente compacto, mantendo a expressividade do MHA com menor consumo de memória.

O que é Decoupled RoPE e por que ele é necessário no MLA?

O Decoupled RoPE isola os dados posicionais do conteúdo semântico. Sem esse desacoplamento, a rotação posicional impediria a fusão matricial na query, forçando a descompactação do cache na memória HBM.

Como o Auxiliary-Loss-Free Load Balancing evita o colapso de experts?

Ele utiliza um bias aditivo dinâmico para distribuir tokens entre aceleradores. Isso equilibra a carga computacional sem distorcer os gradientes dos pesos nem prejudicar o aprendizado da rede.

É possível rodar modelos com arquitetura MLA localmente no vLLM?

Sim. O vLLM (a partir da versão 0.6) e o SGLang possuem suporte nativo a kernels fundidos de MLA, permitindo inferência otimizada de modelos DeepSeek com memória paginada.