Skip to content

English · Español

02 — Multi-Latent Attention (MLA)

🇪🇸 La KV cache crece linealmente con el contexto y, en modelos grandes, llena la HBM. MLA (DeepSeek-V2) la comprime: en lugar de cachear K y V completos, cachea un latente de baja dimensión, y reconstruye K y V sobre la marcha. Reducción típica: 4–8×. Coste: cómputo extra por token (la reconstrucción).

La KV cache es el mayor consumidor individual de memoria en tiempo de inferencia para cualquier modelo de tamaño razonable. Para un modelo de 32 capas, 32 heads y 128 head-dim en fp16:

\[\text{KV bytes/token} = 2 \cdot L \cdot H \cdot d_{\text{head}} \cdot 2 = 2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 512\,\text{KB}/\text{token}\]

Un contexto de 100k tokens: 50 GB. La KV cache, no los pesos del modelo, es lo que hace cara la inferencia en contexto largo.

La Multi-Latent Attention (MLA) de DeepSeek-V2 comprime la representación K-V a un latente de baja dimensión, cachea solo el latente y reconstruye K y V sobre la marcha cuando hace falta. Reducción típica de cache: 4–8×. Coste de cómputo: una pequeña proyección durante el cálculo de attention.


El mecanismo

Multi-head attention estándar (recordatorio)

Dado el input \(h_t \in \mathbb{R}^{d_{\text{model}}}\) en la posición \(t\):

\[q_t = W_Q h_t, \quad k_t = W_K h_t, \quad v_t = W_V h_t\]

Cada uno está en \(\mathbb{R}^{d_{\text{model}}}\), luego reshape a \(H\) heads de \(d_{\text{head}}\) cada uno (con \(H \cdot d_{\text{head}} = d_{\text{model}}\)).

La KV cache almacena \(k_t\) y \(v_t\) para todo \(t\) en la secuencia. Coste de memoria por token: \(2 \cdot d_{\text{model}}\) valores por capa.

La compresión de MLA

MLA introduce un "latente comprimido" \(c_t \in \mathbb{R}^{d_c}\) con \(d_c \ll d_{\text{model}}\). El latente es una proyección aprendida del input:

\[c_t = W_{DKV} h_t \in \mathbb{R}^{d_c}\]

(donde \(W_{DKV}\) es "Down-projection for KV", en la nomenclatura de DeepSeek).

En tiempo de attention, \(k_t\) y \(v_t\) se reconstruyen desde \(c_t\):

\[k_t = W_{UK} c_t, \quad v_t = W_{UV} c_t\]

(Up-projections.) Ahora: la cache almacena solo \(c_t\), no \(k_t, v_t\).

Coste de memoria por token: \(d_c\) valores por capa. Con \(d_c = 512\) (la elección de DeepSeek-V2 para \(d_{\text{model}} = 5120\)): reducción de 10×.

¿Y qué pasa con Q?

Q recibe un tratamiento similar en el MLA de DeepSeek, pero Q no se cachea (Q se recomputa en cada forward, no se reutiliza entre posiciones de token). Así que la compresión de Q es puramente un juego de ahorro de parámetros, no de ahorro de cache. Las matemáticas están en el paper de DeepSeek-V2; para la Fase 36 nos centramos en la compresión de K/V, que es la historia de la cache.

Truco de codificación posicional

Una sutileza: la reconstrucción de K/V estilo MLA estándar es incompatible con RoPE (rotary positional embeddings, Fase 16) porque RoPE rota K y Q de forma dependiente de la posición — pero el \(c_t\) cacheado no tiene noción de posición. La solución de DeepSeek: separar K en una "parte rotatoria desacoplada" \(k_t^{\text{rope}} \in \mathbb{R}^{d_h^R}\) y una "parte sin rotación" \(k_t^{\text{c}} \in \mathbb{R}^{d_h^C}\). La parte rotatoria se cachea tal cual (pequeña, \(d_h^R \approx 64\)); la parte sin rotación se recupera del latente.

Para los propósitos de la Fase 36: con saber que el truco existe es suficiente. La matemática de la separación rotatoria está en el paper y en el lab de lectura.


Las matemáticas, de punta a punta

Coste de memoria

Configuración KV cache por token KV cache para 100k tokens, 32 capas
MHA estándar, \(d_{\text{model}}=5120\) \(2 \cdot 5120 = 10240\) valores 64 GB en fp16
MLA, \(d_c = 512\), \(d_h^R = 64\) \(512 + 64 = 576\) valores 3.6 GB en fp16
Ratio 17.7× reducción de 17.7×

Por eso DeepSeek puede servir su modelo de 236B parámetros con contexto largo en hardware que sería imposible para una arquitectura estándar.

Coste de cómputo

Cómputo de attention estándar por token (decode): \(O(d_{\text{model}}^2 + d_{\text{model}} \cdot T)\) donde \(T\) es la longitud pasada.

MLA añade: - Down-projection: \(d_{\text{model}} \cdot d_c\) por token (poco, ya que \(d_c \ll d_{\text{model}}\)). - Up-projection de K y V en tiempo de attention: \(2 \cdot d_c \cdot d_{\text{model}}\) por token atendido. Como la attention atiende a todos los tokens pasados, son \(2 \cdot d_c \cdot d_{\text{model}} \cdot T\).

Comparado con el \(d_{\text{model}} \cdot T\) estándar (el coste del producto escalar QK): MLA es aproximadamente \(2 d_c\) veces más cómputo por token atendido. Con \(d_c = 512\), eso son ~1024×... espera, eso es malo.

La realidad es que el MLA de DeepSeek reorganiza los matmuls para fusionar la up-projection en las matemáticas de la attention, eliminando el coste aparente. La identidad matricial:

\[\text{attn}(q_t, k_\tau, v_\tau) = q_t^\top (W_{UK} c_\tau) \cdot (W_{UV} c_\tau)\]

puede reorganizarse (con algunos productos de matrices precomputados) para computar la attention directamente desde \(c_\tau\) a un coste de \(O(d_c \cdot T)\) en lugar de \(O(d_{\text{model}} \cdot T)\). Reducción neta de cómputo, no aumento. El truco está en la implementación.

El lab de lectura (lab/01-mla-math-exercise.md) recorre la reorganización de matmuls.


Lo que MLA no es

  • No es lo mismo que Multi-Query Attention (MQA) o Grouped-Query Attention (GQA). Esas reducen la cache compartiendo K y V entre heads. Son complementarias; puedes tener GQA + MLA (DeepSeek-V2 lo tiene).
  • No es una aproximación. La K y V reconstruidas son exactamente \(W_{UK} c_t\) y \(W_{UV} c_t\). La información que se pierde es la que la proyección de bajo rango descarte — pero el modelo aprendió esa compresión conjuntamente con el resto, así que no es lossy en ningún sentido de calidad de salida (DeepSeek reporta paridad).
  • No es gratis. Dos matrices de pesos extra por capa (\(W_{UK}, W_{UV}\)). El overhead de parámetros es modesto, ~5% del tamaño del modelo. Vale la pena por la reducción de cache de 10×.

¿Ayudaría MLA al tutor de gramática?

Aplicar el test:

  1. Cuello de botella de MLA: "memoria de KV cache en contexto largo, en un modelo grande".
  2. ¿El tutor de gramática lo tiene? Cuentas de memoria: \(L=4\) capas, \(H=4\) heads, \(d_{\text{head}}=16\), contexto máximo 32 tokens en fp32: KV cache = \(2 \cdot 4 \cdot 4 \cdot 16 \cdot 32 \cdot 4 = 32\) KB en total. No hay problema de KV cache que resolver. Una reducción de 17× sobre 32 KB son 1.9 KB. Vaya cosa.
  3. ¿Qué costaría MLA? Dos matrices nuevas por capa (~8 KB más de parámetros). Nuevas matemáticas de attention. Nuevo code-path que mantener.
  4. Veredicto: Nunca.

MLA es el ejemplo de libro de "herramienta correcta, escala equivocada". DeepSeek-V2 tiene el cuello de botella (236B parámetros, contexto de 128k, serving multiusuario en clusters de H100). El tutor de gramática de Lynx-Cortex no tiene nada de eso. Reconocerlo en 30 segundos, en lugar de construir MLA "para ser modernos", es la habilidad que premia la fase.

Cuándo ayudaría MLA (el contrafáctico)

Si el vocabulario del tutor de gramática creciese a 600k formas (5 idiomas × paráfrasis × etc.) y quisieras servir a millones de usuarios concurrentes con contextos de 1k tokens: ahora tienes un problema de KV cache. Entonces MLA se gana su sitio. Hasta entonces, es una solución bellamente diseñada a un problema que no tienes.


Lo que esta fase NO cubre

  • La separación rotatoria desacoplada de MLA con profundidad matemática. El mecanismo se explica; la derivación completa está en la lectura del lab.
  • Comparar MLA con otras técnicas de compresión de K/V (p. ej., KV cache cuantizada, ver Fase 26 / Q4_K_M). La Fase 36 las cubre solo conceptualmente.
  • Implementar MLA en PyTorch. Solo lectura. El matmul de attention reorganizado no es trivial; la implementación oficial de DeepSeek tiene ~200 líneas y es más fácil de leer que de reescribir.
  • Interacción MLA + speculative decoding. El theory/04-speculative-and-reasoning.md de la Fase 36 lo menciona; no se deriva.

Siguiente: theory/03-state-space-models.md.