English · Español
05 — Matemática de memoria del KV cache, trabajada de punta a punta en el mini-GPT de la Fase 17¶
🇪🇸 Aplicamos la fórmula
bytes = 2 · L · H · d_h · S · B · sal mini-GPT exacto que construiste en la Fase 17. Vemos que el caché ocupa kilobytes, no gigabytes — y por qué incluso a esta escala las cuentas valen la pena: te dan el hábito que necesitas en producción.
Esta página extiende theory/02-memory-cost.md con un ejemplo numérico concreto trabajado usando el mini-GPT de la Fase 17 exactamente como está especificado. La aritmética es lo bastante pequeña como para hacerse a mano; la metodología escala a GPT-3.
Recap de la spec del mini-GPT de la Fase 17¶
De docs/phase-17-mini-gpt/README.md:
- \(d_\text{model} = 64\)
- \(n_\text{heads} = H = 4\)
- \(d_h = d_\text{model} / n_\text{heads} = 16\)
- \(n_\text{layers} = L = 2\)
- \(d_\text{ff} = 256\)
- Vocab (tras BPE sobre el corpus §A13): \(V \approx 512\)
- Longitud máxima de secuencia (según Fase 16): \(S_\text{max} = 128\)
Para inferencia usamos precisión completa fp32 (\(s = 4\) bytes/elemento) salvo indicación.
La fórmula, instanciada¶
Sustituyendo las partes fijas:
Para fp32 (\(s = 4\)):
Para fp16 / bf16 (\(s = 2\)):
Incremento por token, por secuencia (con \(B = 1\), fp32): 1 KiB por token generado.
Casos concretos¶
Caso A: decode de una sola secuencia, contexto máximo completo¶
\(B = 1\), \(S = S_\text{max} = 128\), fp32:
Trivial. Cabe en el caché L1 32 veces en el i5-8250U (que tiene 32 KiB de L1d por core, así que el caché se desborda a L2 — también bien con 256 KiB).
Caso B: decode con batch pequeño, contexto completo¶
\(B = 8\), \(S = 128\), fp32:
Aún trivial. Cabe en el L3 del i5-8250U (6 MiB).
Caso C: decode con batch máximo, fp16¶
Si empujáramos a \(B = 64\) (que excede lo que el corpus §A13 admite con sentido, pero útil para el ejercicio de escalado), \(S = 128\), fp16:
Aún en L3.
Caso D: comparación con los pesos del modelo¶
Conteo aproximado de pesos del mini-GPT: \(V \cdot d_\text{model} + L \cdot (4 \cdot d_\text{model}^2 + 2 \cdot d_\text{model} \cdot d_\text{ff} + 2 \cdot d_\text{model}^2_\text{FFN})\) — aproximado a ~50k pesos, más ~50k para embeddings y LM head. Llámalo ~100k pesos × 4 bytes = 400 KiB.
Así que en \(B = 1, S = 128\) fp32, el KV cache (128 KiB) es ~32% del tamaño de los pesos. El KV cache es no trivial como fracción, pero el tamaño absoluto es pequeño.
Esta es la lección §A13: el cociente caché vs pesos a nuestra escala es similar al de GPT-3 con contexto largo. La aritmética escala.
Comparación con GPT-2-small¶
Sanity-check de la metodología sobre un número familiar:
GPT-2 small: \(L = 12, H = 12, d_h = 64, S_\text{max} = 1024\), fp16, \(B = 1\):
Esto coincide con la cifra citada de 36 MiB en la literatura para el caché de inferencia de GPT-2-small. La metodología funciona.
Coste marginal por token: §A13 mini-GPT vs GPT-3¶
| Modelo | Coste de caché por token (fp16) |
|---|---|
| §A13 mini-GPT | \(2 \cdot 2 \cdot 4 \cdot 16 \cdot 2 = 512\) bytes / token |
| GPT-2 small | \(2 \cdot 12 \cdot 12 \cdot 64 \cdot 2 = 36{,}864\) bytes / token = 36 KiB / token |
| GPT-3 175B | \(2 \cdot 96 \cdot 96 \cdot 128 \cdot 2 = 4{,}718{,}592\) bytes / token = 4.5 MiB / token |
El coste por token crece 4 órdenes de magnitud a lo largo de estos modelos. A escala GPT-3, generar 100 tokens usa 450 MiB de caché por secuencia. A escala §A13, la misma generación usa 50 KiB. Misma fórmula, misma estructura de exponentes, escala absoluta vastamente distinta.
Implicaciones para prefill vs decode¶
En theory/01-prefill-vs-decode.md vimos que el decode está limitado por memoria porque el modelo carga sus pesos enteros (\(\sim\) 400 KiB) para producir un token nuevo, mientras que los FLOPs necesarios son diminutos (unos pocos miles). Los bytes-por-token son \(\sim 400\) KiB independientemente de \(S\), mientras que el caché añade \(\sim 1\) KiB por token generado. A \(S\) larga, el caché mueve la huella de memoria efectiva del modelo hacia arriba pero no cambia significativamente el coste de transferencia byte-por-token — los pesos dominan.
Por eso a escala §A13, el KV cache es una palanca pedagógica, no una palanca de rendimiento. Lo implementamos correctamente para que la metodología esté en tus dedos cuando te encuentres con modelos a escala GPT-3 en la Fase 23+. El mini-GPT en sí no necesita el caché para velocidad.
La corrupción silenciosa "off-by-one en el índice de posición"¶
El bug más común del KV cache es olvidar que cuando se genera el token \(t+1\), el caché contiene K, V para las posiciones \(0\) a \(t\), así que el nuevo slot K, V va en el índice \(t\) — no \(t+1\). Con un off-by-one, las K, V nuevas sobrescriben el slot de la posición \(t-1\), y la atención del siguiente paso consulta un historial corrupto. El resultado: la loss de decode parece bien para secuencias cortas pero degrada en las largas, porque la corrupción se compone.
El ejercicio de break de la Fase 22 (break/00-...) demuestra esto exactamente.
Tabla mnemónica para memorización¶
Para el §A13 mini-GPT en fp32, memoriza estos números:
- Tamaño K por capa por token: \(H \cdot d_h \cdot s = 4 \cdot 16 \cdot 4 = 256\) bytes.
- Igual para V. Total por capa por token: 512 bytes.
- × 2 capas: 1 KiB por token por secuencia. Memoriza esto — el mnemónico "1 KiB / token".
- Para B = 8 y S = 128: \(1\text{ KiB} \cdot 128 \cdot 8 = 1\) MiB.
Estos tres números bastan para estimar cualquier configuración de caché de cabeza.
Citación¶
Pope, R. et al. (2023). Efficiently Scaling Transformer Inference. https://arxiv.org/abs/2211.05102 — las secciones 2-3 derivan la fórmula de memoria del KV cache y analizan los compromisos prefill / decode a escala. Su fórmula coincide con la nuestra (solo descartamos su factor de group-attention ya que el mini-GPT vainilla es multi-head completo).
Recap en un párrafo¶
El KV cache del mini-GPT de la Fase 17 cuesta \(1024 \cdot S \cdot B\) bytes en fp32, o \(1\) KiB por celda (posición-de-secuencia, batch). Para \(B = 1, S = 128\) son 128 KiB — alrededor del 32% del tamaño de los pesos. La metodología generaliza: aplicar la misma fórmula a GPT-3-175B da 4.5 MiB por token generado por secuencia. A escala §A13, el caché es correcto, no rápido — su propósito es darte la metodología y los bugs (off-by-one, errores de dtype, intercambios de dim-batch) antes de que te cuesten reloj de pared a escala de producción.
Referencias cruzadas: theory/02-memory-cost.md (la derivación de la fórmula), theory/03-decode-as-memory-bound.md (por qué el caché ayuda menos de lo que pensarías en el mini-GPT), Fase 17 README.md (la spec del modelo que se cachea).