English · Español

Lab 01 — Ejercicio matemático de MLA (lápiz y papel, sin código de modelo)¶

Objetivo: derivar desde cero la reducción de KV cache de Multi-Latent Attention. Dibujar la curva cache-vs-contexto a escala de tutor de gramática vs. escala de DeepSeek-V2. Confirmar: MLA es la herramienta correcta a la escala correcta, herramienta equivocada a la nuestra.

Tiempo estimado: 2 horas.

Prerrequisito: theory/02-mla.md leído; Fase 15 (attention) y Fase 22 (KV cache) entendidas. Lápiz + papel recomendado para la derivación; matplotlib para los gráficos.

Lo que produces¶

Un directorio experiments/36-mla-math-exercise/ que contenga:

derivation.md — tu derivación a mano del tamaño de KV cache de MLA, la reorganización de attention y el overhead de parámetros.
kv-cache-curves.py — script que dibuja el tamaño de KV cache vs. longitud de contexto para: (a) MHA estándar a escala de tutor de gramática, (b) MHA estándar a escala de DeepSeek-V2, © MLA a escala de tutor de gramática, (d) MLA a escala de DeepSeek-V2.
kv-cache-curves.png — el gráfico.
findings.md — informe corto: ¿a qué escala empieza MLA a merecer su sitio?

TODOs¶

Bloque A — derivar el tamaño de KV cache de MLA¶

Partiendo de la fórmula estándar de KV cache de MHA:

\[M_{\text{MHA,token}} = 2 \cdot L \cdot H \cdot d_{\text{head}} \cdot \text{bytes\_per\_value}\]

Muestra en derivation.md:

Paso 1. El MHA estándar almacena tanto \(K\) como \(V\) por capa, por head. Coste por token = \(2 L H d_{\text{head}}\) valores.
Paso 2. MLA introduce la down-projection \(W_{DKV} \in \mathbb{R}^{d_c \times d_{\text{model}}}\). Por token, solo se cachea \(c_t \in \mathbb{R}^{d_c}\). Coste por token = \(L \cdot d_c\) valores.
Paso 3. Tener en cuenta el componente K rotatorio desacoplado \(k_t^{\text{rope}} \in \mathbb{R}^{d_h^R}\) que sí se cachea por separado. El coste por token pasa a \(L \cdot (d_c + d_h^R)\).
Paso 4. Calcular el ratio \(M_{\text{MHA}} / M_{\text{MLA}}\) como función de \(H, d_{\text{head}}, d_c, d_h^R\). Mostrar la simplificación algebraica.

Ejemplo trabajado para rellenar:

Escala	\(L\)	\(H\)	\(d_{\text{head}}\)	\(d_c\)	\(d_h^R\)	\(M_{\text{MHA}}/\text{tok}\)	\(M_{\text{MLA}}/\text{tok}\)	Ratio
Tutor de gramática	4	4	16	(probar 16, 32, 64)	(probar 8, 16)	...	...	...
DeepSeek-V2 (paper)	60	128	128	512	64	...	...	...

Rellena los números. Comenta lo que ves.

Bloque B — derivar la reorganización de attention¶

El forward ingenuo de MLA (descrito en theory/02-mla.md) reconstruye \(k_\tau = W_{UK} c_\tau\) en tiempo de attention, lo que costaría \(d_c \cdot d_{\text{model}}\) FLOPs por cada token atendido. Eso es caro.

El truco de matmul: definiendo \(\hat{q}_t = W_{UK}^\top q_t\) solo en el token actual, puedes reescribir el producto escalar de attention como:

\[q_t^\top k_\tau = q_t^\top W_{UK} c_\tau = (W_{UK}^\top q_t)^\top c_\tau = \hat{q}_t^\top c_\tau\]

Así que precomputas \(\hat{q}_t\) una vez por token (coste \(O(d_c \cdot d_{\text{model}})\) solo en el token actual) y luego la attention se hace directamente sobre \(c_\tau\) (coste \(O(d_c \cdot T)\) en lugar de \(O(d_{\text{model}} \cdot T)\)). Neto: menos FLOPs por token en decode, no más.

En derivation.md:

Escribir la identidad de matmul de arriba.
Calcular los FLOPs de attention por token, MLA ingenuo (descomprimido) vs. la forma reorganizada.
Confirmar que la forma reorganizada es estrictamente mejor que MHA estándar en FLOPs por token en decode, asumiendo \(d_c < d_{\text{model}}\).

Bloque C — calcular el overhead de parámetros¶

Comparar conteos de parámetros:

MHA estándar: \(W_Q, W_K, W_V, W_O\) por capa. Total: \(4 L d_{\text{model}}^2\).
MLA: \(W_Q\) (o su variante de bajo rango), \(W_{DKV}, W_{UK}, W_{UV}, W_O\), más la separación rotatoria. Total: \(L \cdot (d_{\text{model}} \cdot (d_c + 2 d_c + d_h^R) + d_{\text{model}}^2)\) aproximadamente.

Mostrar:

Ratio de overhead de parámetros MLA/MHA como función de \(d_c/d_{\text{model}}\).
A escala DeepSeek (\(d_c/d_{\text{model}} = 512/5120 = 0.1\)): overhead de parámetros = X%.
A escala tutor de gramática (\(d_c/d_{\text{model}} = 16/64 = 0.25\)): overhead de parámetros = X%.

La escala del tutor de gramática tiene mayor overhead relativo de parámetros porque \(d_c\) no puede encogerse por debajo de "útil" — hay un suelo.

Bloque D — dibujar las curvas¶

kv-cache-curves.py produce una figura única de 2 paneles:

Panel 1: tamaño de KV cache (MB) vs. longitud de contexto (tokens), a la escala del tutor de gramática. Dos líneas: MHA, MLA. Anota el punto de contexto máximo (32 tokens para el tutor).

Panel 2: mismo gráfico a escala DeepSeek-V2. Dos líneas: MHA, MLA. Anota en 4k, 32k, 128k tokens.

Entrega kv-cache-curves.png.

Qué buscar:

Panel 1: ambas líneas son ambas diminutas (kilobytes). La reducción de MLA es real pero irrelevante.
Panel 2: el gap es dramático. MHA pasa de 50 GB en 128k; MLA se mantiene por debajo de 5 GB.

Bloque E — el informe de findings¶

findings.md (~300 palabras):

¿Cuál es la longitud de contexto de cruce (en MB de KV cache) a la que MLA empieza a importar?
En el contexto máximo del tutor de gramática (32 tokens), ¿cuál es el tamaño absoluto de KV cache con y sin MLA?
¿Qué overhead de parámetros añade MLA a escala de tutor de gramática? ¿Vale la pena?
Veredicto: ¿deberíamos añadir MLA al tutor de gramática? ¿Por qué sí o no?
¿Cuándo ayudaría MLA, hipotéticamente? (El contrafáctico del tutor de gramática de 5 idiomas de theory/02.)

Restricciones¶

Sin código PyTorch. Este es un ejercicio de matemáticas + gráficos. NumPy puro + matplotlib.
Sin entrenamiento de modelo. No intentes "verificar" MLA entrenando uno — no es el objetivo del lab.
Muestra tu álgebra. Si derivation.md salta pasos, no estás aprendiendo. Escribe cada sustitución explícitamente.

Condiciones de parada¶

Has terminado cuando:

experiments/36-mla-math-exercise/{derivation.md, kv-cache-curves.py, kv-cache-curves.png, findings.md} existen todos.
derivation.md recorre los cuatro bloques (tamaño de KV cache, reorganización de attention, overhead de parámetros, ratios).
El gráfico tiene ambos paneles y ambas escalas de longitud de contexto.
findings.md responde a las cinco preguntas.
Puedes recitar, de memoria, "MLA reduce la KV cache por [ratio]× a [parámetros] / [latent dim] / [layers]". Ese es el número mágico que sale del álgebra.

Pista de último recurso¶

Si te atascas en la reorganización de matmul (Bloque B): relee la sección 2.1 del paper de DeepSeek-V2. La reorganización es la idea clave; el resto del paper es implementación. Después de pasar 30 minutos luchando con ello, permítete leer la exposición del paper.

Si tus ratios de KV cache no coinciden con el esperado ~10–18× a escala DeepSeek: revisa dos veces si tuviste en cuenta \(d_h^R\) (el componente K rotatorio cacheado por separado). Olvidar ese término da un ratio excesivamente optimista.

Cuándo consultar `solutions/`¶

Tras commitear findings.md. La solución vive en solutions/01-mla-math-ref.md — escrita al abrir la fase. La referencia incluye los números trabajados a ambas escalas, exactamente la tabla del Bloque A rellena.

Siguiente lab: lab/02-mamba-walkthrough.md.