LLM 2026 · Lo que los cursos no cubren · 1

El bloque transformer de 2026

El Capítulo 1 de HF explica la atención del paper de 2017. Ningún modelo que uses hoy está armado así: cambió dónde se normaliza, con qué, cómo entra la posición, cuántas cabezas guardan K y V, y qué hay dentro de la FFN. Esta página recorre esos cinco cambios con un solo ejemplo, Gemma 4 31B, para que leas un config.json o un paper de arquitectura sin perderte.

Antes de leerHF cap. 1 §3, atención · HF cap. 1 §5, KV cache · fast.ai P2 U2 §4, normalización

🎯 Al terminar vas a poder

Dibujar el bloque de un decoder moderno (pre-norm, RMSNorm, GQA con RoPE, FFN con compuerta) y decir qué reemplaza a cada pieza del bloque de 2017 y por qué.
Explicar RoPE como una rotación por pares de dimensiones, y por qué eso hace que la atención dependa de la distancia relativa entre tokens.
Leer rope_theta, sliding_window, num_key_value_heads y head_dim en un config y calcular cuánta VRAM ocupa el KV cache para un contexto dado.
Explicar cómo Gemma 3 y 4 sostienen 128K-256K de contexto con capas locales y globales intercaladas, y qué se sacrifica a cambio.

0. Qué cambió desde 2017 y por qué nadie lo cuenta junto

El transformer original (Vaswani et al., 2017) era un encoder-decoder con LayerNorm después de cada sub-bloque, embeddings de posición sinusoidales sumados a la entrada, atención multi-cabeza donde cada cabeza tenía sus propias Q, K y V, y una FFN de dos capas con ReLU en el medio. Cada uno de esos cuatro detalles fue reemplazado entre 2019 y 2023, por razones distintas, en papers distintos, y casi ningún curso los junta: HF explica la atención, fast.ai explica la normalización, PyTorch te muestra nn.TransformerEncoderLayer con los defaults de 2017.

El resultado es que el bloque que efectivamente corre en Llama, Qwen, DeepSeek, Mistral y Gemma es el mismo hasta en el orden de las operaciones, y difiere del de 2017 en todo lo que sigue. Vale la pena verlo lado a lado antes de entrar en cada pieza:

El mismo bloque, con cada pieza reemplazada por una razón distinta: pre-norm por estabilidad en profundidad, RMSNorm por costo, RoPE por posición relativa, GQA por memoria de KV cache, la compuerta por calidad a igual cómputo. Las cinco secciones que siguen son las cinco cajas ámbar.

La lógica de las cinco secciones que siguen es la misma: cada cambio resolvió un problema concreto de entrenamiento o de inferencia, y el problema es lo que hay que recordar, no el nombre.

1. Pre-norm y RMSNorm: normalizar antes, y más barato

En 2017 la normalización iba después de la suma residual (post-norm): x = LayerNorm(x + Atención(x)). Eso funciona para 6 capas con warmup cuidadoso y se vuelve inestable al apilar decenas: el gradiente que llega a las primeras capas pasa por cada LayerNorm del camino y se reescala en cada una. Xiong et al. (2020) mostraron que mover la normalización adentro de la rama (pre-norm: x = x + Atención(LayerNorm(x))) deja un camino residual limpio de la entrada a la salida, exactamente el argumento de las skip connections que viste en fast.ai P2 U2. Con pre-norm el warmup deja de ser crítico y se pueden entrenar 80 capas sin trucos.

El segundo cambio es qué se normaliza. LayerNorm resta la media y divide por la desviación estándar de cada vector de activaciones; Zhang y Sennrich (2019) probaron que la resta de la media aporta poco, y que dividir solo por la raíz de la media de los cuadrados (root mean square) da la misma estabilidad con menos operaciones y sin necesidad de un pase para calcular la media:

def rmsnorm(x, peso, eps=1e-6):
    # x: (..., d_model) · una escala aprendida por dimensión, sin bias, sin resta de media
    rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps)
    return x / rms * peso

\mathrm{RMSNorm}(x)=\frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2+\epsilon}}\odot g

Cómo se lee. Cada vector de activaciones se divide por la raíz de la media de sus cuadrados (su RMS) y se reescala dimensión a dimensión por un vector aprendido g. Comparado con LayerNorm, falta la resta de la media y el bias: es lo que la función rmsnorm de arriba calcula.

Llama, Qwen, DeepSeek, Mistral y Gemma usan RMSNorm en pre-norm. Gemma además agrega una segunda RMSNorm a la salida de cada rama (pre y post a la vez) y, desde Gemma 2, normaliza Q y K antes de la atención (QK-norm) para que los logits de atención no exploten con contexto largo. Ese detalle importa cuando cuantizás: cada norma extra es un tensor pequeño en float que hay que dejar sin tocar, como viste en tu pipeline NVFP4.

📐 Derivación analítica del flujo de gradientes: Pre-LN vs. Post-LN y eficiencia de RMSNorm

Para comprender por qué Post-LN colapsa en redes profundas, consideremos la ecuación de recurrencia en Post-LN:

\mathbf{x}_{l+1} = \text{LN}(\mathbf{x}_l + \mathcal{F}_l(\mathbf{x}_l))

Al aplicar la regla de la cadena para evaluar la derivada de la última capa $\mathbf{x}_L$ respecto a la entrada $\mathbf{x}_0$ , el gradiente atraviesa un producto continuo de jacobianos de normalización:

\frac{\partial \mathbf{x}_L}{\partial \mathbf{x}_0} = \prod_{l=0}^{L-1} \frac{\partial \text{LN}}{\partial (\mathbf{x}_l + \mathcal{F}_l(\mathbf{x}_l))} \left( \mathbf{I} + \frac{\partial \mathcal{F}_l(\mathbf{x}_l)}{\partial \mathbf{x}_l} \right)

Como la derivada de LayerNorm normaliza por la varianza de la suma en cada paso, la magnitud del gradiente se degrada exponencialmente en las capas iniciales ( $l \to 0$ ), exigiendo tasas de aprendizaje extremadamente bajas con warm-up prolongado.

En contraste, la formulación de Pre-LN traslada la normalización al interior de la rama:

\mathbf{x}_{l+1} = \mathbf{x}_l + \mathcal{F}_l(\text{RMSNorm}(\mathbf{x}_l))

Expandiendo telescópicamente la suma residual directa:

\mathbf{x}_L = \mathbf{x}_0 + \sum_{l=0}^{L-1} \mathcal{F}_l(\text{RMSNorm}(\mathbf{x}_l)) \implies \frac{\partial \mathcal{L}}{\partial \mathbf{x}_0} = \frac{\partial \mathcal{L}}{\partial \mathbf{x}_L} \left( \mathbf{I} + \sum_{l=0}^{L-1} \frac{\partial \mathcal{F}_l(\text{RMSNorm}(\mathbf{x}_l))}{\partial \mathbf{x}_0} \right)

La presencia de la matriz identidad $\mathbf{I}$ asegura un camino de gradiente directo e ininterrumpido desde la función de pérdida hasta la primera capa.

Eficiencia aritmética de RMSNorm: LayerNorm computa media $\mu$ y varianza $\sigma^2$ , lo que requiere dos pasadas de reducción sobre VRAM (o algoritmo de Welford en registros) y retener $\mu$ para el backward. RMSNorm omite el centrado de media:

\text{RMS}(\mathbf{x}) = \sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2 + \epsilon}

Esto reduce las operaciones de reducción por token de dos a una, eliminando escrituras intermedias y reduciendo la latencia del kernel GPU entre un 7% y un 10% en arquitecturas memory-bound.

🟢 Lo que hay que retener

Pre-norm es una decisión de optimización (gradientes limpios en redes profundas); RMSNorm es una decisión de costo (mismo efecto, menos cómputo y menos memoria de activaciones). Cuando un paper dice "we use pre-normalization with RMSNorm" está diciendo "el bloque estándar desde Llama 1", nada más.

2. SwiGLU y GeGLU: la FFN con compuerta

La FFN de 2017 era W2 · ReLU(W1 · x): proyectar a 4× el ancho, cortar los negativos, volver. Shazeer (2020) probó variantes donde la capa de subida se desdobla en dos proyecciones paralelas y una multiplica a la otra después de una activación suave, como una compuerta (gated linear unit, GLU):

# FFN clásica (2017):        FFN con compuerta (Llama, Qwen, DeepSeek: SwiGLU; Gemma: GeGLU)
h = relu(W1 @ x)              h = silu(W_gate @ x) * (W_up @ x)     # silu(z) = z·sigmoid(z); GeGLU usa gelu
y = W2 @ h                    y = W_down @ h

Tres matrices en vez de dos, así que para mantener el mismo número de parámetros el ancho intermedio baja de 4·d a unos 8/3·d (Llama 3 8B: d = 4096, intermedio 14336, que es 3,5·d redondeado a múltiplo de 256). A igualdad de parámetros y cómputo, la versión con compuerta da mejor perplejidad; el paper de Shazeer no ofrece una explicación teórica y lo dice con ironía en el título, pero el resultado se replicó en cada familia desde entonces. La intuición útil: la compuerta permite que una parte de la FFN decida cuánto de la otra pasa, token por token, en vez de que un ReLU fijo decida por signo.

📐 Formulación rigurosa de SwiGLU y presupuesto dimensional

La activación SwiGLU (Shazeer, 2020) se define formalmente como:

\text{SwiGLU}(\mathbf{x}) = \left( \text{Swish}_1(\mathbf{x}\mathbf{W}_{\text{gate}}) \right) \odot (\mathbf{x}\mathbf{W}_{\text{up}}) \mathbf{W}_{\text{down}}

donde $\text{Swish}_1(\mathbf{z}) = \mathbf{z} \odot \sigma(\mathbf{z}) = \frac{\mathbf{z}}{1 + e^{-\mathbf{z}}}$ actúa como compuerta suave y no monótona.

Deducción de la dimensión intermedia: En un transformer estándar con FFN clásica de 2 matrices, el número de parámetros es:

P_{\text{FFN}} = 2 \times d \times d_{\text{inter}} = 2 \times d \times 4d = 8d^2

En SwiGLU, al tener 3 matrices ( $\mathbf{W}_{\text{gate}}, \mathbf{W}_{\text{up}} \in \mathbb{R}^{d \times d_{\text{inter}}}$ y $\mathbf{W}_{\text{down}} \in \mathbb{R}^{d_{\text{inter}} \times d}$ ), el conteo paramétrico es:

P_{\text{SwiGLU}} = 3 \times d \times d_{\text{inter}}

Igualando ambos presupuestos para preservar idénticos FLOPs y huella de memoria:

3 d \cdot d_{\text{inter}} = 8 d^2 \implies d_{\text{inter}} = \frac{8}{3} d \approx 2.67 d

En arquitecturas reales de LLMs (como Llama 3), los autores aumentaron la capacidad a $d_{\text{inter}} \approx 3.5d$ redondeado a múltiplos de 256 (e.g., 14.336 para $d=4096$ ) para garantizar alineamiento de memoria contigua en los Tensor Cores de la GPU.

En un state_dict de Llama esto se ve como mlp.gate_proj, mlp.up_proj y mlp.down_proj. Esos tres tensores son, en un modelo denso, alrededor de dos tercios de los parámetros totales por bloque; por eso la página siguiente los reemplaza por expertos, y por eso son lo primero que se cuantiza.

3. RoPE: la posición como rotación

La atención en sí no sabe en qué orden llegan los tokens: softmax(QKᵀ) da el mismo resultado si permutás las filas. En 2017 la posición entraba sumando un vector fijo (senos y cosenos de distintas frecuencias) al embedding de cada token, o con un embedding aprendido por posición como en BERT y GPT-2. Los dos tienen el mismo defecto: la posición se mezcla con el contenido en la entrada y después el modelo tiene que aprender a desenredarla, y el segundo además no puede representar una posición que no vio en entrenamiento.

Rotary Position Embedding (Su et al., 2021) cambia dónde y cómo entra la posición: no toca el embedding, sino Q y K justo antes del producto punto, y no suma sino que rota. Se agrupan las dimensiones de cada cabeza en pares (d₀,d₁), (d₂,d₃), …; a cada par le toca una frecuencia θᵢ, y el vector del token en la posición m se rota en ese plano un ángulo m·θᵢ:

RoPE no suma un vector de posición: rota cada par de dimensiones de Q y K un ángulo m·θᵢ. Como el producto punto de dos vectores rotados depende solo de la diferencia de ángulos, la atención ve (m − n): distancia, no posición. Cada par tiene su frecuencia; los lentos son los que codifican distancias largas y los que hay que estirar para contexto largo.

La propiedad que lo hace funcionar es de geometría elemental: el producto punto entre dos vectores rotados depende solo de la diferencia de ángulos. Si q está en la posición m y k en la posición n, el puntaje de atención después de RoPE depende de q, de k y de (m − n), no de m y n por separado. La posición absoluta desaparece del cálculo y queda la relativa, que es lo que un modelo de lenguaje necesita: "el token de tres lugares atrás", no "el token 1.047".

Las frecuencias siguen la misma escala geométrica que los senos de 2017: θᵢ = base^−2i/d, con base 10.000. Los primeros pares giran rápido (una vuelta cada pocos tokens) y codifican distancias cortas; los últimos giran tan despacio que en 4.000 tokens no completan una vuelta y codifican distancias largas. Ese detalle es el que abre la sección siguiente.

\theta_i=\text{base}^{-2i/d},\quad f_m(x_{2i},x_{2i+1})=\big(x_{2i}\cos m\theta_i-x_{2i+1}\sin m\theta_i,\ x_{2i}\sin m\theta_i+x_{2i+1}\cos m\theta_i\big),\quad \langle f_m(q),f_n(k)\rangle=g(q,k,m-n)

Cómo se lee. El par i de dimensiones se rota un ángulo m·θ_i, con m la posición del token; cada par tiene su frecuencia θ_i. La tercera igualdad es la propiedad que importa: el producto punto entre una Query rotada en m y una Key rotada en n solo depende de m − n. Es la función rope de arriba, vista como matemática.

📐 Matriz ortogonal de rotación de RoPE y demostración de invariancia relativa

Sea un vector de activación $\mathbf{x} \in \mathbb{R}^d$ . RoPE descompone el espacio en $d/2$ subespacios bidimensionales ortogonales e introduce la matriz de rotación diagonal por bloques $\mathbf{R}_{\Theta, m}^{d}$ :

\mathbf{R}_{\Theta, m}^{d} = \text{diag}\big(\mathbf{R}_{m\theta_0}, \mathbf{R}_{m\theta_1}, \dots, \mathbf{R}_{m\theta_{d/2 - 1}}\big)

donde cada bloque de rotación 2D para la posición $m$ con frecuencia $\theta_i = \text{base}^{-2i/d}$ es:

\mathbf{R}_{m\theta_i} = \begin{pmatrix} \cos(m\theta_i) & -\sin(m\theta_i) \\ \sin(m\theta_i) & \cos(m\theta_i) \end{pmatrix}

Dado que cada bloque pertenece al grupo ortogonal especial $\text{SO}(2)$ , la matriz es ortogonal: $(\mathbf{R}_m)^\top = (\mathbf{R}_m)^{-1} = \mathbf{R}_{-m}$ .

Demostración de preservación de posición relativa: Evaluemos el producto punto de atención entre una Query en posición $m$ y una Key en posición $n$ :

\langle \mathbf{R}_m \mathbf{q}, \mathbf{R}_n \mathbf{k} \rangle = (\mathbf{R}_m \mathbf{q})^\top (\mathbf{R}_n \mathbf{k}) = \mathbf{q}^\top \mathbf{R}_m^\top \mathbf{R}_n \mathbf{k} = \mathbf{q}^\top \mathbf{R}_{-m} \mathbf{R}_n \mathbf{k} = \mathbf{q}^\top \mathbf{R}_{n - m} \mathbf{k}

Dado que $\mathbf{R}_{n-m} = \mathbf{R}_{-(m-n)}$ , el resultado es invariante bajo traslación absoluta y depende estrictamente del desplazamiento relativo $m - n$ :

\langle \mathbf{R}_m \mathbf{q}, \mathbf{R}_n \mathbf{k} \rangle = g(\mathbf{q}, \mathbf{k}, m - n)

def rope(x, pos, base=10_000.0):
    # x: (..., seq, head_dim) · pos: (seq,) · rota los pares (x[2i], x[2i+1]) un ángulo pos·theta_i
    d = x.shape[-1]
    theta = base ** (-torch.arange(0, d, 2).float() / d)          # (d/2,)
    ang = pos[:, None].float() * theta[None, :]                        # (seq, d/2)
    cos, sin = ang.cos(), ang.sin()
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], -1).flatten(-2)

Los senos y cosenos no dependen de los datos, así que se precalculan una vez y se guardan como buffer del módulo, no como parámetro: son las tablas cos_cached / sin_cached que aparecieron en la página de nn.Module como ejemplo de register_buffer.

4. Contexto largo: base de RoPE, ventanas y capas globales

Un modelo entrenado con secuencias de 4K tokens nunca vio los ángulos que RoPE produce en la posición 30.000 para los pares lentos, y al extrapolar se degrada rápido. Hay dos familias de solución, y los modelos actuales usan las dos:

Estirar RoPE

Interpolación de posiciones (Chen et al., 2023): en vez de dejar que el ángulo crezca más allá de lo entrenado, se comprimen las posiciones (m → m/s) para que 32K posiciones ocupen el rango angular de 4K. Con un fine-tuning corto el modelo se adapta.
Subir la base: con base 500.000 (Llama 3) o 1.000.000 (Gemma 3 y 4 en sus capas globales) los pares lentos giran todavía más despacio, así que a 128K tokens siguen dentro de ángulos que el modelo puede aprender. Gemma 3 lo combina con un entrenamiento a 32K y una extensión al final a 128K reescalando RoPE por un factor 8.
YaRN (Peng et al., 2023): interpola solo las frecuencias bajas, deja las altas intactas (las distancias cortas no necesitan cambiar) y corrige la temperatura del softmax. Es lo que usan Qwen y DeepSeek para llegar a 128K; en el config aparece como rope_scaling con type: "yarn" y un factor.

No atender a todo

Aunque RoPE aguante, la atención completa cuesta O(n²) por capa y el KV cache crece una columna por token en cada capa. La solución de Gemma desde la versión 2, y ya de Mistral en 2023, es que la mayoría de las capas atiendan solo a una ventana deslizante de los últimos w tokens, y que unas pocas capas globales atiendan a todo:

Modelo	Patrón local : global	Ventana	Base de RoPE (local / global)	Contexto
Gemma 2 (2024)	1 : 1	4096	10k / 10k	8K
Gemma 3 (2025)	5 : 1	1024	10k / 1M	128K
Gemma 4 31B (2026)	5 : 1, la última capa siempre global	1024	10k / 1M, y en las globales RoPE parcial: solo el 25 % de las dimensiones rota (`partial_rotary_factor: 0.25` en el config)	256K
Gemma 4 E2B	4 : 1	512	10k / 1M	128K

Las ablaciones del reporte de Gemma 3 son el argumento: pasar de 1:1 a 5:1 y de ventana 4096 a 1024 casi no mueve la perplejidad, y baja el sobrecosto de memoria del KV cache a 32K de contexto del 60 % a menos del 15 % del tamaño del modelo. La intuición es que la mayor parte de lo que un token necesita está cerca; las capas globales integran lo lejano sobre representaciones que las locales ya refinaron. El precio es fijo: si una tarea necesita cruzar el documento entero en cada capa, el modelo no puede pedir más capas globales de las que tiene.

El RoPE parcial de las capas globales de Gemma 4 (p = 0,25) merece una línea: rotar solo un cuarto de las dimensiones deja tres cuartos sin posición, y esas dimensiones hacen matching de contenido puro, insensible a la distancia, que es lo que una capa global quiere a 200K tokens. Como las dimensiones sin rotar no necesitan tablas de senos, el reporte lo cuenta como parte del ahorro de KV cache global (37,5 % junto con compartir K y V, que es la sección siguiente).

5. GQA y la cuenta del KV cache

En la atención multi-cabeza de 2017 cada cabeza tiene su propia K y su propia V, y las dos hay que guardarlas para cada token pasado durante el decode, como viste en HF cap. 1 §5. Multi-query attention (Shazeer, 2019) fue al otro extremo: todas las cabezas de Q comparten una sola K y una sola V, y el cache se divide por el número de cabezas. Pierde algo de calidad. Grouped-query attention (Ainslie et al., 2023) es el punto medio que usa todo el mundo: las cabezas de Q se reparten en grupos y cada grupo comparte una K y una V. Llama 3 8B: 32 cabezas de Q, 8 de KV; Llama 3 70B: 64 y 8.

La fórmula del cache es la misma que en tu script de NVFP4 · Gemma 4, ahora con el detalle de que las capas locales solo guardan la ventana:

bytes_por_token_y_capa = 2 · num_kv_heads · head_dim · bytes_por_valor      # 2 = K y V (1 si K = V)
cache_capa_global      = bytes_por_token_y_capa · contexto
cache_capa_local       = bytes_por_token_y_capa · min(contexto, ventana)

Con la geometría del config.json de google/gemma-4-31B-it (60 capas, 50 locales con 16 cabezas KV de 256 y ventana 1024, 10 globales con 4 cabezas de 512 donde V = K; el 12B tiene 8 y 1, que es la geometría que aparece en tus fixes de vLLM), la cuenta a 256K de contexto en bf16 da esto:

La misma cantidad de capas y de cabezas KV, tres cuentas distintas. La ventana deslizante hace que 50 de las 60 capas dejen de crecer a partir de 1024 tokens, y K = V con 4 cabezas de 512 deja las 10 globales en 4 KB por token. De 258 GB a 11.6 GB: es la razón por la que 256K de contexto entran en una GPU, y por la que head_dim y num_kv_heads pasaron a ser propiedades de cada capa.

🔧 En tu proyecto

Esta es la razón de fondo de la "geometría heterogénea" que rompió el backend de Transformers en vLLM: head_dim y num_kv_heads ya no son constantes del modelo sino propiedades de cada capa, porque el diseño del KV cache se hace capa por capa. Y es también por qué las capas globales con head_dim = 512 no entran en FlashAttention-2, que soporta hasta 256, y vLLM cae al kernel de Triton para ese modelo.

Las cinco piezas se cierran en un solo bloque, y así se lee un config.json moderno: rms_norm_eps (§1), hidden_act: "silu" con intermediate_size (§2), rope_theta (§3), rope_scaling y sliding_window o layer_types (§4), num_attention_heads y num_key_value_heads (§5). Todo lo demás es tamaño. Y un cambio de 2026 que viene de esta misma geometría: transformers 5.16 declara num_key_value_heads y head_dim como atributos por capa; leerlos como valor global en un modelo heterogéneo lanza AmbiguousGlobalPerLayerAttributeError, y la forma correcta es config.per_layer_config[i] o config.per_layer_config["full_attention"]. Es el mismo supuesto roto que arreglaste a mano en vLLM, ahora formalizado en la librería.

6. Errores comunes

Error	Por qué pasa	Qué hacer
Extender el contexto solo con `max_position_embeddings` más grande	Con RoPE ese campo no crea nada: las posiciones nuevas producen ángulos que el modelo no vio y la calidad cae después del largo de entrenamiento.	Usar `rope_scaling` (YaRN o interpolación) con el factor que corresponde, y si es posible un fine-tuning corto a la longitud nueva (§4).
Calcular el KV cache con `num_attention_heads`	Con GQA las cabezas de Q no guardan nada: el cache lo dimensionan las cabezas KV, que son 4 a 8 veces menos.	`num_key_value_heads × head_dim`, y por capa si el modelo tiene capas locales y globales (§5).
Asumir `head_dim = hidden_size / num_heads`	Desde Gemma 2 el `head_dim` es un campo propio y puede ser mayor que ese cociente; en Gemma 4 además difiere entre capas.	Leer `head_dim` del config y, si hay `layer_types`, la geometría de cada tipo (§5).
Cuantizar o fusionar las normas junto con las matrices grandes	RMSNorm y QK-norm son vectores chicos en float cuya escala fija el rango de la capa siguiente; tocarlos desestabiliza todo el bloque.	Dejar normas, embeddings y `lm_head` fuera de la cuantización de 4 bits, como hace la receta NVFP4 (§1).
Leer "sliding window 1024" como "el modelo solo ve 1024 tokens"	Es por capa: la información viaja más lejos a través de las capas y las globales ven todo.	Mirar el patrón local:global y dónde están las globales; la última capa de Gemma 4 siempre es global (§4).
Aplicar RoPE a V, o después del producto punto	RoPE rota Q y K antes del producto para que el puntaje dependa de la posición relativa; V lleva contenido y no debe rotarse.	Rotar solo Q y K, justo antes de `QKᵀ`, y con las mismas tablas para los dos (§3).

7. Resumen

El bloque moderno es pre-norm con RMSNorm, atención GQA con RoPE, y FFN con compuerta (SwiGLU o GeGLU). Cada pieza reemplazó a una de 2017 por una razón concreta: estabilidad en profundidad, costo, posición relativa, memoria de cache, calidad por parámetro.
RoPE rota pares de dimensiones de Q y K un ángulo proporcional a la posición; el producto punto de dos rotaciones depende solo de la diferencia de ángulos, así que la atención ve la distancia relativa.
Para contexto largo se estira RoPE (base más alta, interpolación, YaRN) y se evita atender a todo: capas locales con ventana y pocas capas globales, en Gemma 3 y 4 con patrón 5:1 y ventana 1024.
El KV cache lo dimensionan las cabezas KV y el head_dim de cada capa; con GQA y ventanas, Gemma 4 31B guarda unos 12 GB a 256K de contexto donde un diseño con todas las capas globales y las mismas cabezas necesitaría unos 260 GB.
Esa geometría por capa es lo que rompió supuestos de "un valor global" en el código que arreglaste en vLLM.

8. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

1. ¿Por qué pre-norm permite entrenar redes más profundas que post-norm?

Porque con pre-norm la suma residual queda fuera de la normalización: hay un camino de identidad desde la entrada hasta la salida por el que el gradiente pasa sin reescalarse en cada capa. Con post-norm cada LayerNorm está en el camino principal y el gradiente hacia las primeras capas se atenúa o amplifica en cada una; por eso necesitaba warmup fino y no escalaba en profundidad. (§1)

2. ¿Qué propiedad matemática hace que RoPE codifique posición relativa y no absoluta?

Que el producto punto de dos vectores rotados en el mismo plano depende solo de la diferencia de sus ángulos. Si q se rotó m·θ y k se rotó n·θ, q·k después de rotar es función de (m − n)·θ. La posición absoluta se cancela en el producto y queda la distancia. (§3)

3. Llama 3 8B tiene 32 cabezas de atención, 8 cabezas KV, head_dim 128 y 32 capas. ¿Cuántos bytes de KV cache por token en bf16?

Por capa: 2 × 8 × 128 × 2 bytes = 4.096 bytes. Por 32 capas: 131.072 bytes, 128 KB por token. Con 8K de contexto son 1 GB por secuencia; con multi-head (32 cabezas KV) sería 4 GB. (§5)

4. ¿Por qué Gemma 3 usa base 1M para RoPE en las capas globales y deja 10k en las locales?

Las locales solo ven 1024 tokens: base 10k ya cubre esas distancias con buena resolución. Las globales tienen que distinguir posiciones hasta 128K; con base 1M los pares lentos giran más despacio y a 128K siguen en ángulos que el modelo puede aprender, en vez de dar vueltas completas que confunden posiciones lejanas con cercanas. (§4)

5. En Gemma 4 las capas globales rotan solo el 25 % de las dimensiones. ¿Qué gana el modelo con las que no rota?

Dimensiones que hacen matching por contenido puro, sin ninguna penalización por distancia: dos tokens semánticamente iguales a 200K de distancia se parecen tanto como si estuvieran juntos. Para una capa cuya función es integrar lo lejano es lo que se quiere; la posición fina la aportan las capas locales. (§4)

6. Por qué la FFN con compuerta tiene tres matrices y aun así no tiene más parámetros que la de 2017.

Porque el ancho intermedio se reduce de 4·d a alrededor de 8/3·d para compensar la tercera matriz: 3 × (8/3)·d² = 8·d², igual que 2 × 4·d². A igual costo, la versión con compuerta rinde mejor. (§2)

9. Para practicar

✏️ Ejercicio 1 — verificar la propiedad relativa de RoPE numéricamente

Con la función rope de §3, generá un q y un k al azar de dimensión 64, rotá q en la posición 10 y k en la 3, y compará el producto punto con el de q en la posición 110 y k en la 103. Después repetí con posiciones (10, 3) y (10, 7).

import torch
torch.manual_seed(0)
q, k = torch.randn(64), torch.randn(64)
def puntaje(m, n):
    qm = rope(q[None], torch.tensor([m]))[0]
    kn = rope(k[None], torch.tensor([n]))[0]
    return (qm @ kn).item()
print(puntaje(10, 3), puntaje(110, 103))   # misma distancia 7
print(puntaje(10, 3), puntaje(10, 7))     # distancias 7 y 3

Pista / qué esperar

Los dos primeros coinciden hasta el error de float32 (la diferencia es 7 en los dos casos); los dos segundos difieren. Si no coinciden, lo más probable es que estés rotando pares que no son adyacentes o usando tablas distintas para q y k.

✏️ Ejercicio 2 — leer la geometría de un config real y calcular el cache

Escribí una función que reciba un id del Hub y un largo de contexto y devuelva los bytes de KV cache en bf16, leyendo solo el config: layer_types, sliding_window, attention_k_eq_v, num_kv_shared_layers, y las cabezas KV y el head_dim de cada capa con config.per_layer_config[i]. Corré Gemma 4 E2B, 26B A4B y 31B contra un modelo homogéneo (SmolLM2-360M o Qwen2.5-0.5B) a 8K, 32K, 128K y 256K. La función completa está en el ejercicio 3 de HF cap. 1.

Pista / qué esperar

Salida real con transformers 5.16, a 8K / 32K / 128K / 256K: 31B 1,09 / 2,03 / 5,78 / 10,78 GiB; 26B A4B 0,27 / 0,51 / 1,45 / 2,70; E2B 0,05 / 0,19 / 0,76 / 1,51; SmolLM2-360M 0,31 / 1,25 / 5,00 / 10,00. El homogéneo crece lineal; en Gemma la curva se aplana porque las locales se clavan en 1024 tokens y solo las globales siguen. Si tu 31B a 256K da mucho más de 11 GiB, no estás aplicando la ventana; si da lo mismo que a 8K, la aplicaste también a las globales. Y si leés c.num_key_value_heads a secas en el 31B, transformers 5.16 lanza AmbiguousGlobalPerLayerAttributeError: esos atributos son por capa desde esa versión.

10. Fuentes y para seguir

Xiong et al. — On Layer Normalization in the Transformer Architecture (2020, arXiv:2002.04745) — el análisis de por qué post-norm necesita warmup y pre-norm no.
Zhang & Sennrich — Root Mean Square Layer Normalization (2019, arXiv:1910.07467) — RMSNorm: misma estabilidad sin centrar.
Shazeer — GLU Variants Improve Transformer (2020, arXiv:2002.05202) — SwiGLU y GeGLU comparadas a igual cómputo.
Su et al. — RoFormer: Enhanced Transformer with Rotary Position Embedding (2021, arXiv:2104.09864) — RoPE y la derivación de la propiedad relativa.
Chen et al. — Extending Context Window via Position Interpolation (2023, arXiv:2306.15595) y Peng et al. — YaRN (2023, arXiv:2309.00071) — las dos maneras de estirar RoPE.
Ainslie et al. — GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023, arXiv:2305.13245) — grouped-query attention y cómo convertir un checkpoint multi-head.
Gemma Team — Gemma 3 Technical Report (2025, arXiv:2503.19786) — las ablaciones de local:global, tamaño de ventana y memoria de KV cache citadas en §4.
Gemma Team — Gemma 4 Technical Report (2026, arXiv:2607.02770) y la model card — K = V en capas globales, RoPE parcial p = 0,25, bases 1M/10k, y la tabla de tamaños y ventanas usada en esta página.
Llama Team — The Llama 3 Herd of Models (2024, arXiv:2407.21783) — la arquitectura de referencia densa: GQA, SwiGLU, RoPE con base 500k.