Este podría ser el momento que AMD lleva mucho tiempo esperando.
Recientemente, Wafer AI desplegó el modelo Kimi K3 en AMD MI355X. El resultado es que un modelo que originalmente requería 16 NVIDIA B200, distribuidos en dos servidores, ahora se puede desplegar en un solo servidor AMD equipado con 8 MI355X.
Lo más crucial es que no solo lograron empaquetar el modelo.
En pruebas con 1024 tokens de entrada y 400 tokens de salida, el MI355X alcanzó un rendimiento total de 952 tokens/s, con una velocidad de generación por usuario de 118 tokens/s.
Calculado por nodo individual, su rendimiento es aproximadamente 3.8 veces mayor que el promedio por nodo de la solución de 16 B200 en dos nodos, y también supera a los B200 y B300 en relación rendimiento-precio.
Y lo más sorprendente fue que, esta vez, ROCm no dio problemas.
Los modelos son tan grandes que la memoria empieza a ser más crucial que la potencia de cálculo
Kimi K3 tiene 2.8 billones de parámetros, y solo los pesos del modelo requieren más de 1.5 TB de memoria de video (VRAM), sin contar la KV Cache necesaria para el contexto de millones de tokens.
Un servidor con 8 B200, cada uno con 192 GB de VRAM, tiene una capacidad total de aproximadamente 1.5 TB. Esto significa que es difícil albergar siquiera los pesos del modelo, y mucho menos dejar espacio para la KV Cache. Por lo tanto, B200 necesitaba usar dos servidores con 16 GPUs.
B300 tiene 288 GB de VRAM por tarjeta, lo que permite albergar el modelo en un solo nodo. Casualidad o no, AMD MI355X también tiene 288 GB de VRAM. 8 MI355X suman aproximadamente 2.3 TB, suficiente para un solo servidor.
No se trata solo de usar una máquina menos. Cuando el modelo se ejecuta en múltiples nodos, cada token generado puede requerir sincronización de datos a través de la red. Incluso con una red RoCE v2 de aproximadamente 195 Gb/s, la comunicación entre nodos ralentiza la decodificación.
MI355X, con su mayor VRAM, mantuvo todo el modelo dentro de un solo nodo.
En los resultados finales, el rendimiento total máximo de 8 MI355X alcanzó 952 tokens/s, con una velocidad de generación por flujo de 118 tokens/s.
En comparación, el despliegue de dos nodos con 16 B200 tuvo un rendimiento total de 498 tokens/s, que, calculado por nodo individual, es de aproximadamente 249 tokens/s.
Es decir, el rendimiento por nodo del MI355X es aproximadamente 3.8 veces mayor que el rendimiento promedio por nodo del despliegue de dos nodos con B200. En cuanto a velocidad de generación por usuario, los 118 tokens/s del MI355X también superan los 90 tokens/s del B200.
B300 sigue siendo la solución con mayor rendimiento absoluto. Un nodo con 8 B300 alcanzó un rendimiento total de 1568 tokens/s y una velocidad de generación por flujo de 172 tokens/s, lo que supone aproximadamente 1.65 veces el rendimiento total del MI355X.
Pero el precio cambia la conclusión. Wafer realizó cálculos basándose en un precio hipotético de 2.5 dólares por hora por tarjeta MI355X, 4.25 dólares por B200 y 6 dólares por B300.
Bajo estos supuestos, MI355X proporciona aproximadamente 48 tokens/s de rendimiento máximo por dólar; B200 ofrece aproximadamente 7 tokens/s; y B300, aproximadamente 33 tokens/s.
B300 es más rápido, pero MI355X tiene una mayor eficiencia por costo unitario. Para los centros de datos que necesitan ejecutar modelos abiertos a gran escala, esto podría ser más importante que simplemente competir por el título de mayor rendimiento.
Lo más inesperado: ROCm funcionó prácticamente sin modificar
Históricamente, el mayor problema de las GPUs de AMD para centros de datos no ha sido el hardware, sino el software.
Un mismo modelo que se ejecuta directamente en CUDA puede requerir modificar marcos de trabajo, añadir operadores o incluso reescribir kernels de bajo nivel para funcionar en ROCm.
Pero el caso de Kimi K3 fue diferente.
AMD ofreció un soporte casi simultáneo a su lanzamiento. Wafer indicó que el modelo pudo ejecutarse básicamente sin cambios en MI355X, y el trabajo posterior se centró en algunos problemas menores de compatibilidad y optimización de rendimiento.
Uno de los problemas apareció en la decodificación especulativa. Kimi K3 no proporciona los parámetros del modelo borrador necesarios para MTP o EAGLE, por lo que Wafer utilizó un modelo borrador de difusión por bloques externo.
Esta solución funcionaba directamente en CUDA, pero en el entorno ROCm, la primera solicitud real hizo que el planificador fallara. La razón era que la rama de ROCm carecía de la definición de una función llamada `top_k_renorm_prob`.
Lo que hace esta función no es complejo: selecciona los k valores más altos de una distribución de probabilidad, establece a cero las demás probabilidades y luego renormaliza las probabilidades restantes.
Wafer finalmente suplió esta lógica usando una función estándar de PyTorch, sin necesidad de escribir un kernel de GPU personalizado ni rediseñar el sistema de decodificación especulativa.
Después de la corrección, la decodificación especulativa mejoró el rendimiento por flujo en aproximadamente 2.2 veces, el rendimiento por flujo único en concurrencia media en aproximadamente 1.7 veces, y el rendimiento total máximo aumentó aproximadamente un 18%.
Lo más importante es que el sistema pudo alcanzar su rendimiento máximo con un mayor nivel de concurrencia.
La primera palabra llegaba demasiado lenta; la solución fueron solo cuatro ceros
Por supuesto, el rendimiento no lo es todo en un servicio de inferencia. Para el usuario real, otra métrica que afecta directamente la experiencia es el TTFT (Time To First Token), el tiempo de espera entre enviar una solicitud y ver el primer token generado.
En este aspecto, el rendimiento inicial del MI355X no fue bueno. Frente a una tarea de prellenado en frío de aproximadamente 172,000 tokens, MI355X necesitaba unos 51 segundos, mientras que B300 solo requería unos 23 segundos.
En modelos que admiten un contexto de millones de tokens, las tareas de prellenado pueden ser enormes. Si el usuario tuviera que esperar decenas de segundos o más cada vez que procesa un contexto largo, incluso una alta velocidad de decodificación no compensaría la mala experiencia.
Wafer finalmente descubrió que la diferencia de rendimiento provenía casi en su totalidad de un kernel de atención. Kimi K3, en una configuración de paralelismo de tensores de 8 vías, asigna 12 cabezas de atención por GPU. Sin embargo, el kernel de prellenado MLA de AMD AITER, que es más rápido, solo admite formas específicas como múltiplos de 4, 8 o 16.
Las 12 cabezas no coincidían, por lo que el sistema recurría a la implementación genérica más lenta de Triton.
La solución fue sencilla: rellenar las 12 cabezas de atención con ceros hasta alcanzar 16, llamar al kernel de alta velocidad existente, y después de calcular, recuperar solo las 12 cabezas necesarias. No hubo modificación de la arquitectura del modelo ni escritura de nuevo código ensamblador, solo se añadieron cuatro ceros.
Después de la optimización, la velocidad estable de prellenado del kernel MLA de AITER alcanzó aproximadamente 13,000 tokens/s, mientras que la ruta alternativa original de Triton rondaba los 4000-7000 tokens/s. El tiempo de prellenado en frío se redujo entre dos y tres veces.
Esta optimización no cambia el rendimiento final de decodificación, pero reduce significativamente el tiempo que el usuario espera para ver aparecer la primera palabra.
Esto también demuestra que la aparente gran brecha de software entre AMD y NVIDIA a veces no se debe a una falta de capacidad subyacente, sino simplemente a que los kernels de alta velocidad existentes aún no cubren todas las nuevas formas que adoptan los modelos.
La fortaleza de CUDA persiste, pero ya apareció una grieta
Una sola prueba no puede probar que AMD haya alcanzado por completo a NVIDIA.
B200, por falta de VRAM, se vio forzado a ejecutar en múltiples nodos; B300 sigue teniendo el mayor rendimiento absoluto; y el ecosistema de herramientas, soporte de marcos y desarrolladores de ROCm todavía no iguala al de CUDA.
Pero los modelos abiertos están entrando rápidamente en la era de los billones de parámetros. Cuando un modelo es tan grande que no cabe en un solo servidor, la capacidad de VRAM deja de ser solo un número en una hoja de especificaciones y pasa a afectar directamente el costo de comunicación, la complejidad del despliegue y el rendimiento final.
La estrategia de AMD de dotar a cada tarjeta con más HBM se está convirtiendo en una ventaja sistémica tangible.
Si AMD puede seguir mejorando la estabilidad de ROCm, ampliar el soporte de formas de sus kernels de alta velocidad y ofrecer una adaptación más rápida para los nuevos modelos desde su lanzamiento, entonces los centros de datos tendrán que considerar seriamente estas GPUs. Precio más bajo, más VRAM, rendimiento suficiente, y el software ya no requiere meses de ajustes.
¿Qué opinas al respecto?
Enlaces de referencia:
https://x.com/wafer_ai/status/2083628389903315406
https://x.com/ChiragAsarpota/status/2083864019870634151
Este artículo proviene del WeChat Official Account "机器之心" (ID:almosthuman2014), autor: 关注LLM的.
