Kimi K3 da Moonshot roda mais barato em GPUs AMD MI355X

Kimi K3 na MI355X: 952 tok/s por nó

A Wafer publicou um benchmark do modelo Kimi K3 (2,8 trilhões de parâmetros) rodando em GPUs AMD MI355X. O modelo exige mais de 1,5 TB de VRAM só para os pesos, sem contar o KV cache de 1M tokens - não cabe num nó B200 (8 GPUs de 192 GB). A MI355X, com 288 GB de VRAM por GPU, resolve isso num único nó de 8 GPUs.

Resultados de desempenho

Num benchmark de 1.024 tokens de input e 400 de output:

• MI355X (TP8): 952 tok/s agregado por nó, 118 tok/s single stream.
• B200 (TP16, 2 nós): 498 tok/s agregado (249 tok/s por nó), 90 tok/s single stream.
• B300 (TP8+DCP8): 1.568 tok/s agregado, 172 tok/s single stream.

O B300 ainda vence em throughput bruto (~1,65×), mas custa ~2,4× mais por GPU. Em performance por dólar, a MI355X entrega 48 tok/s por $/GPU-hr contra 33 do B300 e 7 do B200.

O que foi preciso fazer

O Kimi K3 roda out-of-the-box na MI355X com suporte day-0 da AMD, mas dois ajustes foram necessários:

1. Speculative decode: o modelo não traz draft tensors embutidos. A solução foi usar o draft externo Kimi-K3-DSpark (RadixArk). No ROCm, um NameError derrubava o scheduler porque a função top_k_renorm_prob não estava definida no build ROCm do sglang. A correção foi uma função PyTorch simples (sort + masked_fill + divide) - sem kernel customizado. Ganho: ~2,2× single-stream, +18% no agregado de pico.

2. Prefill: o kernel rápido AITER MLA não carregava por incompatibilidade de shape (12 heads por rank vs. múltiplos de 4/8/16 esperados). A solução: zero-pad de 12→16 heads. O prefill de 172k tokens caiu de ~51s para ~17-25s (2-3× mais rápido).

Conclusão da Wafer

A MI355X se posiciona como a opção com melhor custo-benefício para modelos de fronteira que exigem muita VRAM. Os problemas de software no ecossistema AMD estão diminuindo - neste caso, sem necessidade de kernels customizados. A Wafer sugere que o "moat" da CUDA está cada vez mais frágil.