El problema fundamental que aborda este trabajo es la creciente demanda de inferencia de modelos de lenguaje grandes (LLM) y la escasez de hardware de alto rendimiento, específicamente GPUs NVIDIA, que ha llevado a un aumento significativo de los costos. Históricamente, NVIDIA ha mantenido una ventaja dominante no solo en el hardware, sino crucialmente en su ecosistema de software CUDA, que ofrece herramientas de desarrollo maduras, bibliotecas optimizadas y soporte 'day-0' para nuevos modelos. Esto crea una barrera de entrada para alternativas de hardware como las GPUs de AMD, a pesar de que sus especificaciones de silicio pueden ser competitivas.

La tesis central es que, mediante una ingeniería de software intensiva y la optimización de la pila ROCm de AMD, es posible cerrar la brecha de rendimiento y lograr una eficiencia costo-rendimiento superior para la inferencia de LLM. Esto implica seleccionar cuidadosamente frameworks de inferencia, aplicar técnicas de cuantificación y realizar ajustes a nivel de kernel para sortear las limitaciones iniciales del soporte de software. El contexto actual de la computación distribuida y la IA a gran escala exige explorar todas las vías para reducir el costo operativo y aumentar la disponibilidad de recursos computacionales, haciendo que la optimización de plataformas alternativas sea una necesidad estratégica.

Arquitectura del Sistema

La arquitectura de inferencia se basa en una GPU AMD MI355X, utilizando la pila de software ROCm. El modelo GLM5.2 se cuantifica a MXFP4 usando AMD Quark, una decisión clave para reducir el footprint de memoria y aumentar el throughput sin pérdida significativa de precisión. El framework de inferencia seleccionado es sglang, elegido por su capacidad de integrar la cuantificación MXFP4 y su menor fricción para el soporte nativo en comparación con vLLM o ATOM.

Para optimizar el throughput, se habilitó la decodificación especulativa (speculative decode) en sglang. Esto requirió dos correcciones a nivel de código: primero, un ajuste en la lógica de cuantificación de sglang para manejar correctamente el 'shared expert' del MTP head, que se mantenía en bf16 mientras el resto del modelo estaba en MXFP4, evitando un crash por 'shape mismatch'. Segundo, se añadió un 'ifdef USE_ROCM' para proteger una inclusión de 'cuda_runtime.h' en el kernel de metadatos multi-paso, permitiendo el funcionamiento de la decodificación especulativa profunda. Finalmente, para mejorar el rendimiento en cargas de trabajo 'prefill-bound' (dominadas por la fase de procesamiento de la entrada), se ajustó la configuración de paralelización de TP8 a TP4xDP2 y se optimizó la selección del kernel MoE (Mixture of Experts) para las formas específicas de GLM5.2 en fp4, superando un fallback heurístico lento en aiter.

Flujo de Optimización de Inferencia de LLM en AMD ROCm

  1. 1 Selección de Modelo y Cuantificación Modelo GLM5.2 cuantificado a MXFP4 con AMD Quark para eficiencia.
  2. 2 Selección de Framework de Inferencia Elección de sglang por soporte MXFP4 y menor fricción con ROCm.
  3. 3 Habilitar Decodificación Especulativa Ajustes en sglang para corregir 'shape mismatch' en MTP head y añadir 'ifdef'...
  4. 4 Optimización de Configuración Ajustes como '--kv-cache-dtype fp8_e4m3' y '--enable-aiter-allreduce-fusion'.
  5. 5 Optimización de Carga de Trabajo Pref... Cambio de TP8 a TP4xDP2 para mejor rendimiento en prefill.
  6. 6 Ajuste de Kernel MoE Sintonización manual del kernel MoE para formas específicas de GLM5.2 en fp4.
CapaTecnologíaJustificación
compute AMD MI355X GPU Hardware principal para la inferencia de LLM, ofreciendo una alternativa de menor costo a NVIDIA Blackwell. vs NVIDIA Blackwell (B300/B200)
compute ROCm Pila de software de AMD para computación GPU, equivalente a CUDA de NVIDIA, que requiere optimización manual para modelos de vanguardia. vs CUDA
data-processing MXFP4 Quantization (AMD Quark) Técnica de cuantificación para reducir el tamaño del modelo GLM5.2 y mejorar el rendimiento de inferencia sin pérdida de precisión. vs FP8 quantization (z-ai)
data-processing sglang Framework de inferencia de LLM seleccionado por su compatibilidad con MXFP4 y menor fricción para soporte nativo en ROCm. vs vLLM, ATOM --kv-cache-dtype fp8_e4m3, --enable-aiter-allreduce-fusion

Trade-offs

Ganancias
  • ▲▲ Costo por GPU
  • ▲ Rendimiento por dólar
  • ▲ Throughput agregado
Costes
  • ▲ Tiempo de ingeniería para optimización
  • ▲ Soporte 'day-0' de software
  • △ Rendimiento absoluto (vs. Blackwell) 
#ifdef USE_ROCM
#include <rocm_runtime.h> // o equivalente para ROCm
#else
#include <cuda_runtime.h>
#endif
Corrección para el kernel de metadatos multi-paso que incluía 'cuda_runtime.h' sin guardas para ROCm.

Fundamentos Teóricos

El concepto de cuantificación, como la utilizada aquí con MXFP4, se basa en principios fundamentales de la teoría de la información y la compresión de datos, buscando representar valores con menos bits para reducir el consumo de memoria y ancho de banda, a menudo con un compromiso mínimo en la precisión. Trabajos seminales en cuantificación para redes neuronales, como los de Courbariaux et al. (2015) con 'Binarized Neural Networks', sentaron las bases para la investigación en la reducción de la precisión de los pesos y activaciones.

La decodificación especulativa (speculative decoding) es una técnica que se inspira en la ejecución especulativa en arquitecturas de CPU, donde se predice el siguiente estado o instrucción para ejecutarla antes de que se confirme su necesidad, mejorando la latencia. En el contexto de LLMs, esto se traduce en generar tokens con un modelo 'draft' más pequeño y rápido, y luego verificar estas predicciones con el modelo 'target' más grande y preciso, como se describe en trabajos como 'Accelerating Large Language Model Decoding with Speculative Sampling' (Leviathan et al., 2023). La optimización de kernels específicos para operaciones como Mixture of Experts (MoE) se relaciona con la investigación en compiladores y optimización de rendimiento a nivel de hardware, donde se buscan transformaciones de código para explotar la arquitectura subyacente de la GPU, como la coalescencia de memoria, el uso eficiente de registros y la paralelización masiva, principios explorados en la literatura sobre computación paralela y GPGPU.