Como o Google otimizou o Qwen 3.5 de 397B no TPU7x (Ironwood)

O time de performance do Google publicou um relatório técnico detalhando como colocou o Qwen 3.5-397B, um modelo Mixture-of-Experts (MoE), para rodar de forma eficiente na plataforma Ironwood (TPU v7x). O documento, hospedado no Google Developers Blog, é menos sobre "mais um modelo grande" e mais sobre uma metodologia: em vez de otimizar cada família de modelo do zero, a equipe decompõe a arquitetura em blocos independentes e reutilizáveis (Batched RPA, Grouped GEMMs, unpermutation em SparseCore) com modelos de custo cientes do hardware. Quando chega uma arquitetura nova, esses módulos pré-otimizados são portados com pouco atrito.

O resultado prático: entre abril e junho de 2026, a performance de inferência melhorou ~3,1x em workloads decode-heavy e ~4,7x em prefill-heavy (tier de 512 requisições concorrentes). E, importante para quem constrói produto, as otimizações foram integradas nativamente ao vLLM e ao SGLang.

Por que esse modelo é difícil de servir

O Qwen 3.5 tem 397 bilhões de parâmetros totais, mas ativa apenas 17 bilhões por token (taxa de ativação de 4,3%). Ou seja: capacidade de um modelo 400B com footprint de execução mais próximo de um 20B. O problema é a estrutura híbrida:

• 60 camadas organizadas em blocos 3:1, onde 75% usam Gated DeltaNet (GDN), uma atenção linear com estado recorrente de tamanho constante (escala O(S) em memória, não O(S²)), e 25% usam GQA tradicional com apenas 2 KV heads.
• A FFN é dividida em 512 experts pequenos, com o router selecionando top-10 por token, mais um expert compartilhado sempre executado.

Essas escolhas quebram o sharding uniforme. É aí que o relatório fica interessante.

O problema do sharding e a solução

Com Tensor Parallelism de 8 (TP=8), tentar dividir 2 KV heads entre 8 devices dá 0,25 head por device, algo fisicamente impossível. Replicar as heads em todos os cores duplica o KV cache em cada device, o que come a HBM e derruba a concorrência real de 512 planejadas para cerca de 200.

A solução foi um esquema híbrido: Attention Data Parallelism (DP=8) replicando os pesos de GQA e GDN em todos os cores, combinado com Expert Parallelism (EP=8) distribuindo os 512 experts (64 por device). Isso preserva o KV cache local e evita duplicar os 400 GB de parâmetros.

Para a transição entre atenção (DP) e MoE (EP), o time evitou o All-to-All (imprevisível sob carga variável) e escolheu All-Gather → MoE local → Reduce-Scatter, priorizando latência determinística, que é o que importa em produção. Dois detalhes de engenharia valem destaque:

• Otimização 3-para-2 no All-Gather: como os índices de experts (int) e os pesos topk (float) têm shape idêntico [1024,10], eles são empacotados via bitcast num único blob int32, reduzindo três coletivas para duas.
• Reduce-Scatter hierárquico escrito em Pallas/Mosaic: primeiro reduz entre chiplets do mesmo chip (memória compartilhada 6x mais rápida que ICI chip-a-chip), depois entre chips físicos, com micro-batches para esconder DMA atrás de compute e evitar OOM na VMEM.

Roofline: onde cada fase trava

O relatório separa bem os dois regimes, algo útil para qualquer um que dimensiona inferência:

• Prefill (8K in / 1K out) é compute-bound: limitado pelos MXUs do TensorCore (4.614 TFLOPS FP8 por chip). O gargalo é distribuição irregular de tokens entre experts, que cria stragglers. Roofline estimado: ~4.500 tokens/s/chip com de-rate.
• Decode (1K in / 8K out) é memory-bound: para gerar um token, é preciso streamar os 400 GB da HBM. Intensidade aritmética perto de 1 FLOP/Byte. Latência de 16,36 ms por passo, com limite realista de ~850 tokens/s/chip.

Kernels sob medida em Pallas

Escrevendo kernels custom em Pallas (a linguagem de kernel do JAX), o time contornou o lowering padrão do XLA para controlar layout de VMEM, registradores e escalonamento de memória. Um exemplo concreto: no Ragged Page Attention, aumentar o KV page size de 16 para 256 tokens reduziu a latência do passo de decode sob concorrência 512 de 428µs para 283µs, um ganho de 33,8% no nível de kernel. Blocos pequenos economizam fragmentação, mas geram overhead de indexação que trava o VPU no decode.

O que isso significa para quem constrói no Brasil

Pouca gente por aqui vai plugar um cluster Ironwood amanhã, mas a lição transcende o hardware específico. Primeiro, o valor de modelos de custo por primeiros princípios em vez de tentativa e erro: saber se você está compute-bound ou memory-bound muda toda a estratégia de otimização. Segundo, a integração nativa em vLLM e SGLang significa que boa parte desses ganhos chega a stacks que os times BR já usam. E terceiro, o padrão de decompor o modelo em blocos portáveis é replicável mesmo em escala menor: separar atenção, roteamento e coletivas em módulos com contratos claros paga dividendos quando o próximo modelo chegar.

Fonte: Google Developers Blog