Construye tu propio motor de inferencia LLM de alto rendimiento en C++ y CUDA

¿Alguna vez te has preguntado qué sucede bajo el capó cuando envías un prompt a un modelo de lenguaje grande? La magia no está solo en los pesos del modelo, sino en el motor de inferencia que orquesta el cómputo. Este artículo te guía en la construcción de tu propio motor de inferencia LLM de alto rendimiento en C++ y CUDA, inspirado en la arquitectura de vLLM pero diseñado como una herramienta de aprendizaje.

¿Por qué construir tu propio motor de inferencia?

La mayoría de los desarrolladores interactúan con LLMs a través de APIs de alto nivel como Hugging Face Transformers o OpenAI. Pero si quieres entender las características reales de rendimiento—latencia, rendimiento, uso de memoria—u optimizar para tu hardware específico, necesitas ir más profundo. Construir un motor de inferencia desde cero te enseña:

Cómo los LLMs realmente calculan predicciones
El papel de los kernels CUDA en la aceleración de operaciones matriciales
Estrategias de gestión de memoria para cargas de trabajo en GPU
Técnicas avanzadas como PagedAttention y batching continuo

La arquitectura de tiny-vllm

El proyecto, tiny-vllm, implementa un servidor de inferencia completo para Llama 3.2 1B Instruct. Cubre:

Cargar un LLM real desde el formato Safetensors
Paso forward completo (prefill + decode)
Todo el cómputo con kernels CUDA personalizados
Gestión de caché KV
Batching estático y continuo
Técnicas de softmax en línea y similares a FlashAttention
PagedAttention

Primeros pasos: Prerrequisitos técnicos

Para ejecutar esto en tu máquina, necesitarás:

Linux (probado en 6.19.8 x86_64)
CUDA Toolkit (13.1)
Compilador C++17 (GCC 15.2.1)
Una GPU NVIDIA (probado en RTX 5090)
El archivo del modelo: model.safetensors de Llama 3.2 1B Instruct

La única dependencia externa es nlohmann/json (cabecera única).

Entendiendo el formato Safetensors

Antes de escribir cualquier código, necesitas entender cómo se almacenan los pesos del modelo. Un archivo Safetensors tiene tres secciones:

Tamaño del encabezado (8 bytes): Un entero sin signo de 64 bits que indica la longitud del encabezado
Encabezado (JSON): Contiene metadatos sobre cada tensor—nombre, dtype, forma y desplazamientos
Datos del tensor: Los valores reales de los pesos

Así es como se lee el tamaño del encabezado:

std::ifstream safetensors_file("model.safetensors", std::ios_base::binary);
uint64_t header_size;
safetensors_file.read(reinterpret_cast<char *>(&header_size), 8);

El JSON del encabezado te dice dónde comienza y termina cada tensor dentro de la sección de datos. Esto te permite mapear tensores directamente a la memoria de la GPU sin copiar datos innecesarios.

La arquitectura del modelo: Llama 3.2 1B

De la tarjeta del modelo de Hugging Face, podemos inspeccionar la arquitectura:

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 2048)
    (layers): ModuleList(
      (0-15): 16 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
          (k_proj): Linear(in_features=2048, out_features=512, bias=False)
          (v_proj): Linear(in_features=2048, out_features=512, bias=False)
          (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((2048,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=2048, out_features=128256, bias=False)
)

El paso forward sigue esta secuencia:

Tokenizar el texto de entrada
Recuperar embeddings para cada token
Para cada una de las 16 capas del transformador:
- RMSNorm
- Atención de consulta agrupada (con RoPE)
- Conexión residual
- RMSNorm
- Red feed-forward (con SiLU)
- Conexión residual
RMSNorm final
Proyección lineal al tamaño del vocabulario
Argmax para seleccionar el siguiente token

Tu primer kernel CUDA: Recolección de Embeddings

Implementemos la búsqueda de embeddings. Para cada token de entrada, necesitamos recuperar un vector de 2048 elementos de la tabla de embeddings.

__global__ void embeddingGatherKernel(int *gpu_input_tokens, __nv_bfloat16 *input_embeddings, __nv_bfloat16 *embed_tokens)
{
    int workIndex = threadIdx.x + blockIdx.x * 2048;
    input_embeddings[workIndex] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x];
    input_embeddings[workIndex + 1024] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x + 1024];
}

Nota el truco: dado que la mayoría de las GPUs soportan un máximo de 1024 hilos por bloque, pero cada embedding tiene 2048 elementos, procesamos dos elementos por hilo. Esto evita lanzar más bloques de los necesarios.

RMSNorm con Reducción Paralela

RMSNorm normaliza cada vector de embedding. La fórmula es:

$$\text{RMS}(a) = \sqrt{\frac{1}{n}\sum_{i=1}^{n}a^2_i}$$

Para calcular esto eficientemente en GPU, usamos una reducción paralela (reducción en árbol) en memoria compartida:

__shared__ float rms_vector[1024];
int workIndex = threadIdx.x + blockIdx.x * 2048;
rms_vector[threadIdx.x] = (float)input[workIndex] * (float)input[workIndex] 
                        + (float)input[workIndex + 1024] * (float)input[workIndex + 1024];
__syncthreads();

// Reducción en árbol
for (int i = 1; i < 1024; i = i * 2)
{
    if (threadIdx.x % (i * 2) == 0)
    {
        rms_vector[threadIdx.x] = rms_vector[threadIdx.x] + rms_vector[threadIdx.x + i];
    }
    __syncthreads();
}

if (threadIdx.x == 0)
{
    rms_vector[0] = sqrt(rms_vector[0] / 2048.0 + 1.0e-5);
}
__syncthreads();

output[workIndex] = (__nv_bfloat16)(((float)input[workIndex] / rms_vector[0]) * (float)norm_weights[threadIdx.x]);
output[workIndex + 1024] = (__nv_bfloat16)(((float)input[workIndex + 1024] / rms_vector[0]) * (float)norm_weights[threadIdx.x + 1024]);

El truco de transposición de cuBLAS

cuBLAS espera matrices en formato column-major, pero los modelos suelen almacenarse en row-major. En lugar de convertir los datos, podemos usar propiedades de transposición de matrices:

$$C = A \times B^T \implies C^T = B \times A^T$$

Dado que cuBLAS lee datos row-major como si fueran column-major (efectivamente transponiéndolos), configuramos las banderas en consecuencia:

cublasGemmEx(cublas_handle, 
             CUBLAS_OP_T, CUBLAS_OP_N, 
             KV_DIM, num_active_slots, EMBEDDING_LENGTH, 
             &k_proj_alpha, weights.w_k[layer], CUDA_R_16BF, EMBEDDING_LENGTH, 
             rms_norms, CUDA_R_16BF, EMBEDDING_LENGTH, 
             &k_proj_beta, k_proj_batched_buffer, CUDA_R_16BF, KV_DIM, 
             CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);

Por qué existe la caché KV

Durante la inferencia, solo las proyecciones K y V de cada token necesitan persistir. La proyección Q, los puntajes de atención y las salidas feed-forward se descartan después de cada paso. Al almacenar en caché K y V, evitamos recalcularlos para cada nuevo token—una gran mejora de rendimiento.

Prefill vs. Decode

Prefill: Procesa todos los tokens de entrada para generar el primer token de salida. Esto está limitado por cómputo porque estás procesando muchos tokens a la vez.
Decode: Genera tokens subsiguientes uno a la vez. Esto está limitado por memoria porque estás recuperando valores K y V almacenados en caché.

Atención de Consulta Agrupada (GQA)

En GQA, múltiples cabezas de consulta comparten la misma cabeza de clave y valor. Para Llama 3.2 1B, 4 cabezas de consulta comparten 1 cabeza de clave/valor. Esto reduce el uso de memoria y mejora el rendimiento sin una pérdida significativa de calidad.

for (int i = 0; i < NUM_Q_HEADS; ++i)
{
    int k_head_idx = i / GQA_Q_TO_K_RATIO; // 4 cabezas Q usan la misma 1 cabeza K
    __nv_bfloat16 *q_head = q_proj + i * HEAD_DIM;
    __nv_bfloat16 *k_head = k_proj[layer] + k_head_idx * HEAD_DIM;
    // ... calcular puntajes de atención
}

Batching Continuo

El batching estático procesa N solicitudes juntas pero espera a la más lenta. El batching continuo resuelve esto usando ranuras: cuando una solicitud termina, su ranura se reutiliza inmediatamente para una nueva solicitud de la cola. Esto maximiza la utilización de la GPU.

¿Qué sigue?

El repositorio está en desarrollo activo. Las próximas secciones incluyen:

Derivación e implementación de softmax en línea
PagedAttention para gestión eficiente de memoria
Caché KV paginada

Conclusión

Construir un motor de inferencia LLM desde cero es una de las mejores maneras de entender profundamente cómo funcionan los sistemas modernos de IA. Aprenderás sobre programación de GPU, gestión de memoria, álgebra lineal y diseño de sistemas—todo mientras construyes algo que realmente ejecuta un modelo real.

Si eres profesor, siéntete libre de usar esto como recurso didáctico. Si eres estudiante, haz un fork del repositorio, hazlo funcionar en tu máquina y contribuye con tus mejoras.

¡Feliz codificación!

Fuente

jmaczan/tiny-vllm: Construye tu propio motor de inferencia LLM de alto rendimiento en C++ y CUDA - una versión más pequeña de vLLM