Crear desde cero un motor de inferencia para LLM usando C++ y CUDA

• Cómo construir un motor de inferencia para LLM usando C++ y CUDA sin librerías
• Esto permite entender toda la pila de inferencia de un LLM y comprobar en la práctica cómo distintas optimizaciones afectan la velocidad de inferencia
• Objetivo: implementar un modelo que pueda hacer inferencia rápida en lote único en un servidor con una sola CPU + GPU y lograr una velocidad de procesamiento de tokens superior a llama.cpp

1. Resumen de la arquitectura e inferencia de un LLM

• La mayoría de los LLM principales siguen la misma arquitectura, usando una secuencia continua de bloques Transformer.
• La carga del modelo consiste en definir una clase de bloque Transformer personalizable, organizarla en una secuencia e inicializarla con pesos en formato safetensors.
• La inferencia se realiza principalmente en lote único, y la "fase de decode" ocupa la mayor parte de la ejecución.

1.1 Resumen de la inferencia

• La inferencia se divide en una fase de prefill, donde los tokens del prompt se pasan al modelo para llenar la caché KV, y una fase de decode, donde el modelo se ejecuta repetidamente para generar tokens
  • Fase de prefill: procesa los tokens del prompt e inicializa la caché KV
  • Fase de decode: genera un token a la vez
• Caché KV: almacena pares clave/valor previos para calcular rápidamente la atención con el contexto pasado
• El forward pass del modelo usa una tabla de embeddings para mapear IDs de tokens a vectores de embedding y transforma el estado a través de una secuencia de bloques Transformer

1.2 Cuellos de botella y benchmarks

• Cuello de botella: en el hardware moderno, el ancho de banda de memoria es el factor limitante
  • Para generar cada token durante la inferencia del modelo, hay que leer el modelo completo, por lo que el ancho de banda de memoria pesa más que el cómputo
• La cuantización del modelo es efectiva para mejorar la velocidad de inferencia
• El rendimiento teórico máximo de tokens varía según el hardware, y el rendimiento real puede verificarse con varios motores de inferencia
• Límite teórico de velocidad:
  • AMD EPYC 7702P: máximo 13.6 tok/s (base FP16)
  • RTX 4090: máximo 67.1 tok/s (base FP16)
• Benchmarks:
  • llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
  • calm: GPU 66 tok/s

2. Inferencia basada en CPU

• La implementación inicial en CPU es de un solo hilo y solo soporta pesos FP32
• Se puede empezar a paralelizar el código con multithreading y mejorar el rendimiento usando SIMD

2.1 Multithreading

• Se usa OpenMP para paralelizar la multiplicación matriz-vector (matmul) y la atención multi-head, mejorando el rendimiento
• Resultado de la optimización: mejora de velocidad de 0.6 tok/s → 4.4 tok/s

2.2 Cuantización de pesos y optimización SIMD

• Cuantización: cuantizar pesos FP32 a FP16 reduce a la mitad el uso de memoria y mejora el rendimiento
• SIMD: optimización con AVX2 para procesar 8 valores FP32 al mismo tiempo
• Resultado: 8.4 tok/s

3. Inferencia basada en GPU

• Cuantizando el modelo a FP16 se puede cargar en una RTX 4090 y comenzar a implementar la inferencia en GPU
• Con CUDA, las funciones en C++ (kernels) pueden ejecutarse en paralelo en la GPU

3.1 Port sencillo a CUDA

• Se puede implementar un backend para GPU convirtiendo las operaciones de CPU 1 a 1 en kernels CUDA
• Los kernels CUDA se ejecutan de forma asíncrona, pero dentro del mismo stream se ejecutan secuencialmente
• Problema: por la ineficiencia de los hilos no se aprovechan bien los recursos de la GPU → lento, con 2.9 tok/s

3.2 Mejor multiplicación de matrices (matmul)

• La multiplicación de matrices ocupa una gran parte del tiempo de ejecución en CPU y puede optimizarse con OpenMP
• En GPU, se puede aumentar el aprovechamiento de hilos haciendo que cada bloque procese 1 fila
• Métodos de optimización:
  1. Un bloque procesa una fila y los hilos dentro del bloque colaboran en el cálculo
  2. Aplicar warp reduction
• Resultado: mejora de velocidad a 51.7 tok/s

3.3 Fusión de kernels y optimizaciones adicionales

• Se puede mejorar el rendimiento fusionando kernels
  • Fusión de kernels: combinar operaciones consecutivas en un solo kernel para minimizar accesos a memoria y tiempo de cómputo
• Con la optimización del patrón de acceso a memoria y la reutilización de espacio, se lograron 56.1 tok/s

3.4 Optimización de attention y manejo de contexto largo

• Problema: con contexto largo, el kernel de attention se vuelve un cuello de botella de rendimiento
• Soluciones:
  1. Optimización del acceso a memoria: rediseñarlo para leer bloques contiguos de memoria
  2. Usar memoria compartida en lugar de atomicAdd para resolver el problema de valores decimales perdidos
• Resultado de la optimización:
  • Contexto corto: 63.8 tok/s (más rápido que los 61.0 tok/s de llama.cpp)
  • Contexto largo: 58.8 tok/s

3.5 Cuantización de la caché KV y problemas de optimización del compilador

• Cuantizar la caché KV a FP16 provoca una caída de rendimiento (por falta de optimización del compilador)
• Solución: desenrollado manual de loops y prefetching de memoria
• Resultado: aprox. 2x más rápido que FP32 y se mantiene el rendimiento en contexto largo de 58.8 tok/s

4. Direcciones de mejora futura

• Optimización del prefill del prompt: procesar varios tokens al mismo tiempo para reducir el tiempo hasta el primer token
• Fusión de kernels de attention: aplicar técnicas de optimización como FlashAttention
• Cuantización más agresiva: aplicar FP8, INT8, INT4 y cuantización de activaciones/caché
• Optimización de kernels: introducir técnicas avanzadas para maximizar el ancho de banda de memoria y la eficiencia de cómputo
• Uso de librerías: aprovechar librerías como cuDNN y cuBLAS para reducir el tiempo de optimización

Resumen de resultados:

• Se alcanzó una velocidad de 63.8 tok/s mediante varias optimizaciones en CPU y GPU
• Se registró un rendimiento similar o superior al de llama.cpp y calm
• Se implementó un motor de inferencia para LLM de alto rendimiento usando solo C++ y CUDA, sin librerías

[GN⁺]

Comentarios de Hacker News

• El autor está contento de que su entrada de blog haya llamado la atención y quiere escuchar comentarios
• Un lector elogia el artículo y pregunta cuánto tiempo tomó escribirlo
  • Como alguien que trabaja en GPGPU, le gustaría escribir algo similar, pero duda por la incertidumbre sobre el tiempo requerido
• Otro lector cree que el código no aprovecha tensor cores ni las instrucciones wgmma
  • Explica que este tipo de programación es difícil porque hay que manejar varias tareas al mismo tiempo
  • Menciona que, debido a las limitaciones de ancho de banda, quizá no haga falta cómputo adicional
  • Considera que el código del blog probablemente funcionaría bien al portarlo a otros aceleradores
  • Le preocupa que usar wgmma pueda reducir la portabilidad entre generaciones de Nvidia
• Otro lector está buscando material similar en Python y quiere compartirlo con su equipo
  • Busca material conciso, estilo tutorial y conceptualmente completo, más que centrado en el rendimiento
• Un usuario quiere comparar su versión de Mistral y el rendimiento en tokens/segundo
  • Se recomienda revisar la sección de cuantización del README
• Hay una opinión de que __shfl_down ya no se recomienda hoy en día por problemas de sincronización de warp