Dentro de un LLM: Inferencia en Python puro para Qwen3-0.6B en MLX
Construcción paso a paso de la inferencia de LLM en unas 200 líneas de Python con MLX: tokenizador, RMSNorm, RoPE, atención a consultas agrupadas, SwiGLU, caché KV y muestreo.
Fuentes: How Modern LLMs Actually Generate Text: A Pure Python Walkthrough
Escribe “La capital de Francia es” y un modelo lo completa con “París” en aproximadamente 900 milisegundos. ¿Qué sucede en ese breve intervalo de tiempo? ¿Qué secuencia de transformaciones convierte cinco palabras en una predicción?
La mejor forma de averiguarlo es construir todo el pipeline de inferencia por uno mismo en lugar de utilizar una biblioteca que oculta los detalles. La guía a continuación lo hace con pesos reales de Qwen3-0.6B, cargados en Python mediante MLX, el framework de arrays de Apple. El script completo consta de unas 200 líneas y funciona en cualquier Mac con Apple Silicon. Las instrucciones de configuración se encuentran al final.
Un vistazo al pipeline
Qwen3, GPT-4, Llama y sus variantes siguen todos un flujo común:
Text → Tokenize → Embed → [Transformer Block × 28] → LM Head → Logits → Sample → Next Token
↑ |
└-------------------------- KV Cache --------------------------┘
El texto se introduce, se convierte en números, pasa por 28 bloques de transformador y sale como una distribución de probabilidades sobre 151,936 tokens siguientes posibles. Se elige un token y se agrega, para luego reiniciar el bucle.
Ese bucle tiene dos fases:
- Prefill. Todos los tokens de entrada se procesan de una sola vez. El modelo lee la instrucción y construye su estado interno.
- Decode. Los tokens se generan uno por uno. Un nuevo token solo necesita tener en cuenta a los anteriores, cuyas claves y valores se encuentran en la caché KV, por lo que nada se vuelve a procesar.
Prefill se ejecuta en paralelo mientras que decode se ejecuta de forma secuencial. La caché KV es lo que hace que decode sea rápido. Sin ella, generar el token N significaría volver a calcular la atención con respecto a los N-1 tokens anteriores.
Ahora, analicemos cada etapa por turnos.
Paso 1: La tokenización convierte el texto en enteros
Un modelo nunca ve palabras, solo enteros.
from tokenizers import Tokenizer
tok = Tokenizer.from_file("Qwen3-0.6B/tokenizer.json")
ids = tok.encode("The capital of France is").ids
# [785, 6722, 315, 9625, 374]
Cada número entero es un índice hacia un vocabulario de 151,936 tokens. El tokenizador utiliza codificación de pares de bytes, por lo que los límites de los tokens no siempre coinciden con los límites de las palabras. “capital” es un token y “Francia” es otro, siendo el espacio inicial parte del token. Escribir BPE desde cero no es lo interesante aquí; el tokenizador es una tabla de búsqueda, y el script simplemente utiliza la ya existente.
Paso 2: Los embeddings convierten los números enteros en vectores
El primer cálculo real asigna a cada ID de token un vector de 1,024 dimensiones.
x = W["model.embed_tokens.weight"][mx.array(token_ids)] # (seq, 1024)
No es más que una búsqueda en una tabla: la matriz de incrustación tiene la forma (151936, 1024), con una fila por token y una columna por dimensión oculta, y el código hace referencia a ella. Piense en la matriz como el diccionario del modelo: cada token cuenta con una huella digital formada por 1,024 números aprendidos que codifican su significado. Las huellas digitales de “París” y “Francia” apuntan en direcciones similares, mientras que las de “París” y “computadora” apuntan en direcciones muy diferentes.
Paso 3: RMSNorm
Antes de aplicar la atención, los estados ocultos se normalizan. Los modelos modernos utilizan RMSNorm en lugar del antiguo LayerNorm.
En términos sencillos, la fórmula divide cada vector por su valor de raíz cuadrática media, que es la raíz cuadrada de la media de sus valores al cuadrado más una constante muy pequeña, y luego lo multiplica por un peso aprendido para cada característica.
def rms_norm(x, w):
return x * mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + EPS) * w
Cada vector se escala según su propia magnitud y luego se ajusta con pesos aprendidos por característica. No hay sustracción de media ni término de sesgo, lo que lo hace más simple que LayerNorm sin sacrificar calidad. EPS (1e-6) evita la división por cero, y el peso w, aprendido durante el entrenamiento, permite al modelo reescalar cada una de las 1,024 dimensiones de forma independiente.
Paso 4: RoPE integra la posición en los vectores
Aquí es donde muchos artículos sobre “LLM desde cero” pierden a sus lectores, pero en realidad es más simple de lo que parece.
El problema. La atención es invariante a las permutaciones. Si se barajan los tokens de entrada, el resultado será el mismo barajado, por lo que el modelo no puede determinar cuál token vino primero.
La solución antigua. Añadir un número de posición al embedding de cada token. Eso funciona, pero no es generalizable, y el modelo tiene dificultades para razonar sobre distancias relativas.
La solución moderna. Aplicar un ángulo a los vectores de consulta y clave que depende de la posición del token. Eso es el Rotary Position Embedding, o RoPE. Cada par de dimensiones se trata como un punto en un plano y se gira según un ángulo determinado por la posición del token.
def apply_rope(q, k, positions):
half = D_HEAD // 2 # 64
freqs = 1.0 / (THETA ** (mx.arange(half, dtype=mx.float32) / half))
angles = mx.array(positions, dtype=mx.float32)[:, None] * freqs[None, :]
cos = mx.cos(angles)[None, :, :]
sin = mx.sin(angles)[None, :, :]
def rotate(x):
x1 = x[..., :half]
x2 = x[..., half:]
return mx.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1)
return rotate(q), rotate(k)
Imagínese cada par de dimensiones como un plano 2D. Los tokens en posiciones diferentes giran por cantidades distintas, por lo que cuando la atención calcula el producto escalar entre una consulta y una clave, la respuesta depende de cuán separados estén y no solo de lo que contengan. Aquí hay tres tokens en las posiciones 0, 1 y 2, cada uno representado como una flecha 2D rotada un poco más que la anterior:
Position 0 (angle = 0°) Position 1 (angle = 24°) Position 2 (angle = 48°)
q q q
| / /
| / /
| / /
| / /
+------ x +------ x +------ x
Same content, 0° rotation Rotated 24° counterclockwise Rotated 48° counterclockwise
"I'm at the start" "I'm one step in" "I'm two steps in"
En el dibujo original, la tercera flecha no está rotada más que la segunda, por lo que debe considerarse el diagrama como una imagen aproximada y no como una representación exacta.
Attention compara una flecha de consulta con una flecha de clave mediante un producto escalar. Las flechas con el mismo ángulo generan un producto alto, mientras que las flechas con ángulos muy diferentes generan uno bajo. Así, el producto refleja la distancia entre dos tokens, sin que se almacene ningún número de posición explícito en ningún lugar. RoPE no agrega la posición como una señal separada. La rotación incorpora la posición en cada vector.
Qwen3 establece THETA = 1,000,000, una frecuencia base alta que permite al modelo manejar contextos muy largos, hasta 262K tokens sin extensiones.
Paso 5: Attention con cabezas de consulta agrupadas
Este es el elemento central. Si comprende la atención, comprende el transformer.
De MHA a MQA a GQA
Multi-Head Attention (MHA): 16 Q heads, 16 K heads, 16 V heads — expensive
Multi-Query Attention (MQA): 16 Q heads, 1 K head, 1 V head — too aggressive
Grouped Query Attention: 16 Q heads, 8 K heads, 8 V heads — sweet spot
Qwen3-0.6B utiliza GQA, con 16 cabezas de consulta y solo 8 cabezas clave-valor, por lo que cada cabeza KV sirve a dos cabezas de consulta. Eso reduce a la mitad la caché KV con poca pérdida en calidad.
Toda la función de atención
def attention(x, layer_idx, cache, positions):
p = f"model.layers.{layer_idx}.self_attn"
# Project to Q, K, V
q = x @ W[f"{p}.q_proj.weight"].T # (seq, 2048) = 16 heads × 128
k = x @ W[f"{p}.k_proj.weight"].T # (seq, 1024) = 8 heads × 128
v = x @ W[f"{p}.v_proj.weight"].T # (seq, 1024) = 8 heads × 128
# Split into heads
q = q.reshape(-1, N_HEADS, D_HEAD).transpose(1, 0, 2) # (16, seq, 128)
k = k.reshape(-1, N_KV, D_HEAD).transpose(1, 0, 2) # (8, seq, 128)
v = v.reshape(-1, N_KV, D_HEAD).transpose(1, 0, 2) # (8, seq, 128)
# QK-Norm (Qwen3-specific: normalize each head before RoPE)
q = rms_norm(q, W[f"{p}.q_norm.weight"])
k = rms_norm(k, W[f"{p}.k_norm.weight"])
# Apply RoPE
q, k = apply_rope(q, k, positions)
# KV Cache: append new keys/values
if cache[layer_idx] is not None:
ck, cv = cache[layer_idx]
k = mx.concatenate([ck, k], axis=1)
v = mx.concatenate([cv, v], axis=1)
cache[layer_idx] = (k, v)
# GQA: expand 8 KV heads → 16 (each KV head shared by 2 Q heads)
k = mx.repeat(k, N_HEADS // N_KV, axis=0) # (16, total, 128)
v = mx.repeat(v, N_HEADS // N_KV, axis=0)
# Scaled dot-product attention
scale = D_HEAD ** -0.5
scores = (q @ k.transpose(0, 2, 1)) * scale # (16, seq, total)
# Causal mask (only during prefill)
if seq_len > 1:
mask = mx.triu(mx.full((seq_len, total), -1e9, dtype=scores.dtype), k=1)
scores = scores + mask
# Softmax + weighted sum of V
attn = mx.softmax(scores, axis=-1)
out = attn @ v
# Merge heads + output projection
out = out.transpose(1, 0, 2).reshape(-1, N_HEADS * D_HEAD)
return out @ W[f"{p}.o_proj.weight"].T
Algunas partes de esto merecen una explicación.
QK-Norm. Qwen3 aplica RMSNorm a los vectores Q y K de cada cabeza antes de RoPE, una técnica más reciente que estabiliza el entrenamiento. Sus pesos q_norm y k_norm están formados por 128 números, que representan una escala por dimensión dentro de cada cabeza.
El caché KV. Esta es la optimización más importante en la inferencia. Durante el prefill, se calculan K y V para cada token de entrada y se almacenan. Durante la decodificación, solo el nuevo token necesita que se calculen K y V; estos se añaden al caché y la atención opera sobre todo el historial. Omitir el caché obligaría a volver a calcular K y V para cada token anterior, lo que convertiría un trabajo lineal en uno cuadrático.
La máscara causal. Durante el prefill, el token en la posición 3 no debe ver el token en la posición 4, que se encuentra en su futuro. El código agrega -1e9 al triángulo superior de la matriz de puntuaciones para garantizarlo. En la fase de decodificación, cuando la longitud de la secuencia es 1, no se necesita máscara, ya que el único nuevo token puede observar todo lo que está antes de él.
La fórmula. La atención calcula la función softmax de (Q multiplicado por K transpuesto, dividido entre la raíz cuadrada de la dimensión del head), y luego multiplica el resultado por V. El producto escalar mide qué tan bien cada consulta coincide con cada clave. Al dividir entre la raíz cuadrada del tamaño del head, se evita que las puntuaciones crezcan tanto que la función softmax se sature. La función softmax convierte las puntuaciones en probabilidades, y al multiplicar por V se obtiene un promedio ponderado de los vectores de valor.
Paso 6: La capa feed-forward SwiGLU
Después de la atención, cada token pasa por una red feed-forward. Los modelos modernos utilizan SwiGLU en lugar de la versión anterior basada en ReLU. En otras palabras: se calculan dos proyecciones a partir de la entrada, una se pasa por SiLU y se usa como compuerta, las dos se multiplican elemento por elemento, y una proyección final devuelve el resultado a un nivel inferior.
def ffn(x, layer_idx):
p = f"model.layers.{layer_idx}.mlp"
gate = x @ W[f"{p}.gate_proj.weight"].T # (seq, 3072)
up = x @ W[f"{p}.up_proj.weight"].T # (seq, 3072)
silu = gate * mx.sigmoid(gate) # SiLU(x) = x · σ(x)
return (silu * up) @ W[f"{p}.down_proj.weight"].T # (seq, 1024)
La red feed-forward es donde el modelo almacena gran parte de su conocimiento, lo cual es una de las razones por las que los modelos densos grandes parecen tan capaces. Expande el vector de 1,024 dimensiones a 3,072 dimensiones, aplica la compuerta SiLU y lo proyecta de vuelta a 1,024. La compuerta aprende qué información dejar pasar y cuál bloquear, algo similar a un portero para las neuronas. Al expandir de 1024 a 3072 y volver, el modelo cuenta con espacio para almacenar y transformar información, y la presencia de estas compuertas es la razón por la que las redes feed-forward modernas superan a las basadas en ReLU.
Paso 7: El bloque transformer
Un bloque combina todo lo anterior:
def transformer_block(x, layer_idx, cache, positions):
p = f"model.layers.{layer_idx}"
# Pre-norm attention + residual
h = rms_norm(x, W[f"{p}.input_layernorm.weight"])
x = x + attention(h, layer_idx, cache, positions)
# Pre-norm FFN + residual
h = rms_norm(x, W[f"{p}.post_attention_layernorm.weight"])
x = x + ffn(h, layer_idx)
return x
La secuencia se normaliza, se añade el residual, se vuelve a normalizar, se realiza una proyección hacia adelante y se añade nuevamente el residual.
Esas conexiones residuales (x = x + ...) son vitales. Permiten que los gradientes fluyan directamente a través de la red durante el entrenamiento y permiten que la información omita bloques individuales en el momento de la inferencia. Un modelo de 28 capas sin ellas estaría casi imposible de entrenar.
Qwen3-0.6B combina 28 de estos bloques, cada uno perfeccionando un poco más la representación. Para la capa 28, el estado oculto ya posee suficiente comprensión de la entrada como para predecir el siguiente token.
Paso 8: Cabeza LM y muestreo
Después del último bloque hay una última normalización, seguida de una proyección al tamaño del vocabulario:
x = rms_norm(x, W["model.norm.weight"])
logits = x @ W["lm_head.weight"].T # (seq, 151936)
El resultado son 151,936 puntuaciones, llamadas logits, una para cada token posible. Una puntuación más alta indica un token más probable. Para “La capital de Francia es”, las cinco probabilidades más altas después de aplicar softmax son las siguientes:
| Token | Probability |
|-----------|-------------|
| Paris | 0.6582 |
| located | 0.0274 |
| the | 0.0233 |
| ... | 0.0157 |
| ____ | 0.0136 |
“París” gana con el 65,8%, una predicción segura. A continuación, el muestreo convierte las probabilidades en un único ID de token:
def sample(logits, temperature=0.0, top_p=0.95):
if temperature == 0:
return int(mx.argmax(logits)) # Greedy: always pick the best
logits = logits / temperature
probs = mx.softmax(logits)
# Top-p: keep smallest set of tokens whose cumulative prob ≥ top_p
order = mx.argsort(-probs)
sorted_p = probs[order]
cumsum = mx.cumsum(sorted_p)
cutoff = int(mx.sum(cumsum < top_p)) + 1
sorted_p = sorted_p[:cutoff]
sorted_p = sorted_p / mx.sum(sorted_p)
order = order[:cutoff]
idx = mx.random.categorical(mx.log(sorted_p[None, :]))
return int(order[idx])
Temperatura controla la aleatoriedad. En 0, el token con mayor probabilidad siempre gana (estrategia codiciosa). En 1.0, los tokens se seleccionan en proporción a sus probabilidades, y en 0.7 la distribución se vuelve ligeramente más concentrada en los tokens con mayor probabilidad.
Top-p, o muestreo por núcleo, mantiene el grupo más pequeño de tokens cuya probabilidad acumulada alcanza p, por ejemplo 0.95. Esto reduce la cantidad de tokens poco probables mientras se conserva cierta diversidad.
Paso 9: El bucle de generación
El bucle completo se llena una vez de antemano y luego realiza la decodificación:
def generate(prompt, max_tokens=100, temperature=0.0, top_p=0.95):
ids = tok.encode(prompt).ids
# Prefill: process all input tokens at once
cache = [None] * N_LAYERS
logits = forward(ids, cache, 0)
next_id = sample(logits[-1], temperature, top_p)
ids.append(next_id)
# Decode: generate one token at a time
for _ in range(max_tokens - 1):
if next_id == EOS_ID:
break
pos = len(ids) - 1
logits = forward([next_id], cache, pos)
next_id = sample(logits[-1], temperature, top_p)
ids.append(next_id)
return tok.decode(ids)
Ese es todo el ciclo de inferencia. Los motores de servicio como vLLM, SGLang, TensorRT-LLM y mlx-lm son, en esencia, optimizaciones de este núcleo, y el núcleo en sí ocupa unas 198 líneas.
Lo que omite el script
El script abarca todo el pipeline de inferencia de un transformador estándar. Los modelos más recientes añaden elementos adicionales:
- Atención híbrida. Se dice que Qwen3.5 reemplaza el 75% de sus capas de atención por Gated DeltaNet, una variante de atención lineal que utiliza un estado recurrente de tamaño fijo en lugar de la matriz de atención completa. La atención completa es como la memoria fotográfica: perfecta pero costosa. DeltaNet es como notas comprimidas: eficiente pero con algo de pérdida de calidad. Una versión híbrida permite manejar contextos largos de forma eficiente además de ofrecer una recuperación casi perfecta periódicamente.
- Predicción de múltiples tokens. Algunos modelos predicen varios tokens por paso forward, sacrificando un poco de calidad a cambio de una gran mejora en velocidad.
- Cuantización. La inferencia en producción utiliza pesos de 4 bits o 8 bits para ahorrar memoria y tiempo. Este script mantiene los pesos originales en formato bfloat16.
- Batching. Los servidores manejan muchas solicitudes al mismo tiempo y comparten los cálculos entre ellas. Este script procesa una solicitud a la vez.
Ninguno de estos cambios afecta el algoritmo principal. Todos operan sobre la misma base.
¿Por qué molestarse?
Entender cómo funciona la inferencia te convierte en un profesional más competente. Una advertencia de “cache KV lleno” proveniente de un servidor deja de ser enigmática. El argumento relacionado con la memoria en GQA se vuelve fácil de explicar. Y cuando un modelo divaga y se repite, eso te indica que debes utilizar parámetros como temperatura y top-p.
Un LLM moderno no es magia. Consiste en 28 capas de normalización, atención y proyección forward, repetidas, con la posición integrada mediante rotación y el conocimiento almacenado en capas lineales con control de acceso. Las matemáticas son de nivel universitario. La ingeniería es donde reside la ingeniosidad, pero las matemáticas son la base. Solo se necesitan unas 200 líneas de Python, pesos reales y una salida real.
Échalo a andar tú mismo
Todo lo anterior se puede ejecutar en cualquier Mac con Apple Silicon (M1 o posterior).
1. Descarga el modelo
Obtén Qwen3-0.6B desde Hugging Face (el repositorio Qwen/Qwen3-0.6B). Solo se necesitan cuatro archivos, así que omite el resto:
Qwen3-0.6B/
├── config.json
├── model.safetensors (1.5 GB — the actual weights)
├── tokenizer.json
└── tokenizer_config.json
Coloque la carpeta del modelo junto a llm_inference.py, ya que el script resuelve ./Qwen3-0.6B/ desde el directorio desde el cual se ejecuta.
2. Crear un entorno virtual
# Create and activate a fresh venv
python3 -m venv llm_venv
source llm_venv/bin/activate
# Install the minimum packages - just two
pip install mlx tokenizers
No hay PyTorch, ni transformers, ni accelerate, ni el paquete safetensors. MLX se encarga de los cálculos con matrices y de la carga de pesos, mientras que la biblioteca tokenizers realiza la codificación y decodificación BPE.
3. Ejecútalo
python llm_inference.py
Esperar salida similar a esta:
Model: Qwen3-0.6B (28 layers, 1024 dim, 16 Q heads, 8 KV heads)
Prompt: The capital of France is
------------------------------------------------------------
Output: The capital of France is Paris. The capital of Italy is
Rome. The capital of Spain is Madrid...
Time: 0.89s
Si ve ModuleNotFoundError: No module named 'mlx', verifique que el entorno virtual esté activo (source llm_venv/bin/activate) y que la máquina tenga procesador Apple Silicon. MLX no funciona en Macs con procesador Intel, Linux ni Windows.
Todo esto funciona en una laptop, y la respuesta aparece en menos de un segundo.