MQA, GQA e MLA

01 · Conceito

Conceito

Abra a config do Qwen3.8-27B e uma assimetria salta aos olhos: 24 query heads, mas apenas 4 key/value heads. Por que uma equipe de design daria a uma camada seis vezes mais perguntas do que respostas? Porque, durante o serving, as queries são computadas do zero a cada passo e custam apenas pesos — mas cada key e cada value precisam ser armazenados para todo token passado. Os KV heads são os caros, e esta lição é sobre a família de designs que economizam neles: MQA, GQA e MLA.

Monte a contabilidade. Com $H_q$ query heads e $H_{kv}$ KV heads, o tamanho do cache por token por camada é proporcional a $H_{kv}$ (vezes a head dimension, vezes dois para K e V, vezes os bytes por elemento). A multi-head attention clássica define $H_{kv} = H_q$: cada query head possui uma visão privada de key/value.

Aqui está o desvio a encenar deliberadamente, porque você vai vê-lo por aí. Alguém lê “24 heads, head dimension 256, bf16” e calcula o cache do Qwen como $24 \times 256 \times 2 \times 2\ \text{B} = 24\ \text{KiB}$ por token por camada. Esse é o número do MHA — ele ignora a grouped-query attention e superestima exatamente pelo fator que os projetistas eliminaram por engenharia. O número real, derivado uma vez na lição 7.2 e apenas referenciado aqui, usa os 4 KV heads: 4 KiB por token por camada com cache, 64 KiB por token através das 16 camadas de full attention, 16 GiB por sequência no contexto nativo de 262.144 tokens. Rode a variante hipotética de MHA pelo mesmo pipeline e a comparação é gritante:

um encolhimento de 6× — por token isso é 384 KiB contra 64 KiB, e em contexto completo hipotéticos 96 GiB por sequência contra os 16 GiB reais. A variante MHA não conseguiria segurar nem uma única sequência de comprimento total ao lado dos 54 GB de pesos em bf16 num acelerador de 80 GB; o design com GQA consegue. O mesmo fator governa a banda do decode, já que o cache precisa ser lido a cada passo por toda camada com cache.

A multi-query attention (MQA) leva o compartilhamento ao limite: $H_{kv} = 1$. Para a geometria do Qwen isso significaria $1 \times 256 \times 2 \times 2\ \text{B} = 1\ \text{KiB}$ por token por camada — 16 KiB por token, cerca de 4 GiB em contexto completo, um encolhimento de 24× em relação ao MHA. O preço é representacional: todos os 24 query heads precisam ler de um único subespaço compartilhado de key/value. Os query heads ainda podem fazer perguntas diferentes, mas todas consultam o mesmo índice e a mesma carga, e se a qualidade sobrevive depende do modelo, da escala e da receita de treinamento — não dá para ler isso no diagrama da arquitetura.

A grouped-query attention (GQA) é a interpolação, e a escolha do Qwen3.8-27B assenta bem nela: 24 query heads em 4 grupos de 6, cada grupo compartilhando um K/V head. Os grupos preservam alguma diversidade de subespaços de key/value enquanto dividem cache e banda por seis. O artigo de GQA (Ainslie et al.) mostrou que o design também pode ser convertido: agrupe os K/V heads de um checkpoint MHA por grupo, e então continue treinando brevemente. O mapeamento é estrutural — os kernels precisam saber quais query heads leem qual K/V head, e os checkpoints codificam os formatos das projeções (lembre da lição 4.2 que as projeções de K e V do Qwen produzem 1024 dimensões, 4 heads × 256, contra 12288 para Q — 6144 de queries interleaved por head com 6144 de gate).

A multi-head latent attention (MLA), introduzida com o DeepSeek-V2, abandona o eixo $H_{kv}$ por completo. Em vez de guardar keys e values por head, ela guarda em cache um vetor latente aprendido de baixa dimensão por token e reconstrói os componentes da attention por projeções, com tratamento desacoplado da parte posicional de RoPE. O tamanho do seu cache é definido pela dimensão latente, não por uma contagem de heads, e sua álgebra difere — o ponto de qualidade contra tamanho que ela alcança não está na linha MHA-GQA-MQA de forma alguma.

A implantação é o filtro final. Um cache menor só compensa quando kernels, quantização e layouts paralelos implementam a variante exata de forma eficiente; cair de volta em caminhos genéricos de attention pode apagar o ganho. Compare candidatos por bytes de cache por token, banda do decode, vazão sob concorrência realista e qualidade em contextos longos — um design de cache pequeno frequentemente vence no nível de sistema por admitir mais sequências concorrentes, um ganho que benchmarks de requisição única nunca enxergam.

O espectro, com o nosso espécime posicionado nele: o MHA dá a cada query head uma visão privada de K/V (24 KiB por token por camada, na geometria do Qwen). O GQA compartilha dentro de grupos — os 24-sobre-4 do Qwen3.8-27B compram um encolhimento de 6× e são metade da sua história de memória, ao lado do layout híbrido que limita o cache a 16 camadas em primeiro lugar. O MQA compartilha uma visão globalmente (1 KiB, 24×). O MLA recodifica o problema num cache latente. Cada um troca estado repetido por uma restrição estrutural mais forte — e os projetistas do Qwen, notavelmente, pegaram duas dessas economias de uma vez.