Por que a arquitetura de atenção agora define a velocidade da inferência

E aí, dev! Workloads agênticos e de longo contexto já não são exceção—são o padrão. Conforme o contexto cresce, a atenção consome uma fatia cada vez maior do tempo de inferência. A forma como você projeta a atenção (não só implementa) impacta diretamente a performance do modelo em GPUs. Essa é a premissa do co-design de modelos de IA: moldar a arquitetura do modelo em torno de como a GPU realmente executa.

Neste deep dive, analisamos como três parâmetros—group size (G), dimensão de cabeça (Hsz) e comprimento de sequência (ISL/KVSL)—afetam a atenção densa. Tudo baseado em fórmulas analíticas (aritmética de formas GEMM) e dados medidos de kernels de prefill e decode com FP8. Para uma discussão mais ampla sobre design de LLMs amigável ao hardware, confira nossa análise anterior sobre co-design de modelos de IA.

Antes de mergulhar, vamos esclarecer as duas fases:

  • Prefill: Processa o prompt completo em paralelo → GEMM-M grande (ISL × G), compute-bound.
  • Decode: Gera um token por vez → GEMM-M pequeno (G), memory-bound pelas leituras do KV cache na HBM.

Decodificação especulativa pode deslocar o decode para compute-bound aumentando o GEMM-M efetivo. Cada parâmetro abaixo é analisado separadamente para prefill e decode porque os gargalos diferem.

Diagram illustrating attention head grouping (GQA) and its impact on GPU memory bandwidth during LLM inference Dev Environment Setup

As quatro diretrizes de co-design (com código e matemática)

Diretriz 1: Escolha o Group Size (G) para eficiência no decode—e aumente

Group size (G) é o número de query heads que compartilham uma KV head. MHA tem G=1, GQA tem G=4,8,16…, MQA tem G=QH.

Prefill é dominado pelo ISL, não pelo G. Intensidade aritmética ≈ 2 × ISL / (1 + 1/G). Com ISL=32K, subir G de 8 para 16 melhora a IA em <6%. O runtime varia <1% entre G=1 e 64.

Decode tem intensidade aritmética ≈ 2 × G (quando KVSL ≫ G). Dobrar G dobra a IA do decode e praticamente reduz o runtime pela metade (até overheads dominarem).

# Cálculo rápido: intensidade aritmética do decode vs group size
# Assume FP8 (1 byte por elemento)
def decode_arithmetic_intensity(G, KVSL):
    # FLOPs = 4 * DB * QH * KVSL * Hsz (constante em G)
    # Bytes = 2 * DB * KH * Hsz * (G + KVSL)
    # AI = 2 * G * KVSL / (G + KVSL)
    return 2 * G * KVSL / (G + KVSL)

for G in [1, 8, 16, 32]:
    ai = decode_arithmetic_intensity(G, 32768)
    print(f"G={G:2d} -> AI ≈ {ai:.1f} FLOPs/byte")
# Saída: G=1 -> 2.0, G=8 -> 16.0, G=16 -> 32.0, G=32 -> 64.0

Moral da história: G mais alto melhora a velocidade do decode e a utilização da GPU. Decodificação especulativa é outra alavanca.

Diretriz 2: Use dimensão de cabeça (Hsz) = 128 ou 256

Hsz não muda a intensidade aritmética (dobra FLOPs e bytes). Mas o runtime cresce com Hsz porque:

  • Matmul cresce em passos alinhados; Hsz=64 ainda paga por um tile de 128. Hsz≥512 atinge o limite do TMEM.
  • Memória se move em unidades de 128 bytes; Hsz deve ser múltiplo de 128.
  • Softmax é independente de Hsz (opera na matriz de scores).
# Ops de softmax escalam com comprimento de sequência, não com dimensão de cabeça
def softmax_ops_prefill(PB, QH, ISL):
    return PB * QH * ISL * ISL  # independente de Hsz

def softmax_ops_decode(DB, QH, KVSL):
    return DB * QH * KVSL      # independente de Hsz

Moral da história: Hsz=128 ou 256 é o ponto ideal para alinhamento com GPU e orçamento de TMEM.

Diretriz 3: Reduza o estado KV efetivo quando possível

Prefill escala quadraticamente com ISL (cada token atende a todos), enquanto decode escala linearmente com KVSL (cada passo lê o KV cache completo). O custo é assimétrico.

Estratégias: compressão de KV cache, atenção esparsa/sliding-window, arquiteturas híbridas (ex.: só algumas camadas carregam estado KV global).

Diretriz 4: Deixe o KH determinar a estratégia de paralelismo

Tensor parallelism (TP) divide as attention heads entre GPUs. Limite prático: TP ≤ KH para que cada GPU tenha pelo menos um grupo completo (uma KV head + suas G query heads). Quando TP > KH, o estado KV é duplicado entre ranks, desperdiçando memória e banda.

Para modelos com poucas KV heads (MQA com 1, GQA com 2), o TP esgota rápido. Use Attention Data Parallelism (ADP) ou KV Parallelism (KVP) para atenção, mais Expert Parallelism (EP) para o FFN MoE. O TensorRT-LLM implementa isso como Wide EP (ADP+EP) e Helix Parallelism (KVP+EP).

# Checagem de paralelismo: garanta que TP não exceda o número de KV heads
def validate_tp(KH, TP):
    if TP > KH:
        raise ValueError(f"TP={TP} excede KH={KH}. Estado KV será duplicado.")
    return True

# Exemplo: GQA com 2 KV heads
validate_tp(KH=2, TP=2)  # OK
# validate_tp(KH=2, TP=4)  # Levanta erro

GPU server rack with high-bandwidth memory handling long-context KV cache for AI model decoding Programming Illustration

Limitações e cuidados críticos

  • Prefill e decode não são simétricos. Otimizações que ajudam o decode (G alto) quase não movem o prefill. Sempre profile as duas fases separadamente.
  • Decodificação especulativa muda o jogo. Ela eleva o GEMM-M efetivo para (1+D)×G, onde D é o número de draft tokens. Quando grande o suficiente para preencher tiles de compute, o decode tende a compute-bound—reduzindo o benefício relativo de G alto.
  • Overheads dominam em sequências curtas. Com KVSL=32K, o runtime do decode achata além de G=16 devido a setup fixo, pós-processamento e redução do flash-decoding. Contextos maiores (128K) amortizam melhor esses custos.
  • TP ≤ KH é limite rígido. Violar isso duplica o estado KV e anula os ganhos de escala. Modelos com poucas KV heads precisam adotar ADP/KVP + EP.
  • Hsz=64 é ineficiente. Ainda paga por um tile de 128—você recebe metade do compute pelo mesmo custo. Evite a menos que esteja restrito.

Próximos passos de aprendizado

  • Experimente diferentes group sizes no seu modelo: meça o throughput de decode com G=1, 4, 8, 16 na sua GPU alvo.
  • Profile prefill e decode separadamente com NVIDIA Nsight Systems para identificar se você está compute-bound ou memory-bound.
  • Explore Wide EP e Helix Parallelism do TensorRT-LLM para escalar modelos com poucas KV heads.
  • Leia o post complementar sobre atenção esparsa (em breve) para entender como a esparsidade muda o cálculo de intensidade aritmética.
  • Para uma perspectiva relacionada sobre integração de IA empresarial, confira nossa análise sobre como agentes de IA estão redefinindo a integração empresarial.

Performance graph showing prefill vs decode runtime scaling with sequence length on NVIDIA GPUs Algorithm Concept Visual

Conclusão: co-design não é opcional

As quatro diretrizes acima formam um checklist prático para desenvolvedores de modelos que miram GPUs NVIDIA:

  1. Aumente G para eficiência no decode (prefill é insensível).
  2. Use Hsz=128 ou 256 para alinhar com tiles da GPU e transferências de 128 bytes.
  3. Reduza o estado KV efetivo via compressão, esparsidade ou arquiteturas híbridas.
  4. Case o paralelismo com o KH—mantenha TP ≤ KH e use Wide EP/Helix para modelos com poucas KV heads.

Com workloads agênticos se tornando padrão, a arquitetura de atenção não é mais um detalhe—é a principal alavanca de performance de inferência. Co-designar a arquitetura do modelo com a execução na GPU em mente é como você entrega inferência rápida e interativa em escala.

Fique seguro: Ao integrar agentes de IA na sua stack, esteja ciente de novos riscos de supply chain como ataques de injeção via AGENTS.md. Segurança também precisa de co-design.

Este conteúdo foi elaborado com o auxílio de ferramentas de IA, com base em fontes confiáveis, e revisado pela nossa equipe editorial antes da publicação. Não substitui o aconselhamento de um profissional especializado.