English · Español

Lab 01 — Implementar el KV Cache¶

Objetivo: escribir src/minicache/cache.py a partir del BLUEPRINT.md. Hacer que tests/test_minicache.py pase.

Tiempo estimado: 4–8 horas en 2–3 sesiones.

Prerrequisito: lab/00-derive-cache-size.md commiteado. src/minicache/BLUEPRINT.md leído; cualquier pregunta abierta sobre él resuelta con /phase-checkpoint antes de empezar.

Lo que produces¶

src/minicache/cache.py — implementación según el API del BLUEPRINT.md.
tests/test_minicache.py — Claude hace scaffold de los tests fallidos; tú los haces pasar.
src/minimodel/attention.py actualizado — attention(...) acepta un cache: KVCache | None opcional.
Todos los tests verdes; mypy --strict src/minicache limpio; ruff check src/minicache limpio.

TODOs¶

Bloque A — leer el blueprint¶

Abre src/minicache/BLUEPRINT.md. Relee cada § (Purpose, API, Alternatives, Complexity, Test plan, Anti-goals, Open questions).
Si alguna pregunta abierta no está resuelta, detente aquí y ejecuta /phase-checkpoint. No programes contra preguntas abiertas.

Bloque B — escribir los tests fallidos (TDD)¶

El scaffold de los tests está en tests/test_minicache.py (Claude commitea esto con cuerpos vacíos; tú rellenas los cuerpos primero, antes de cualquier código del caché). Cada test empieza como un docstring que describe la propiedad que comprueba; tú conviertes el docstring en asserts.

Lista de tests (de BLUEPRINT.md §Test plan):

test_allocate_shapes — KVCache.allocate(layers=4, heads=4, head_dim=32, max_seq=128, batch=2, dtype=np.float32) devuelve un objeto cuyos buffers K y V por capa tienen forma (2, 4, 128, 32).
test_initial_cursor_zero — cache.current_length() == 0 inmediatamente tras allocate.
test_append_advances_cursor — añadir 5 tokens deja el cursor en 5.
test_append_one_token_per_layer — añadir escribe el slice correcto; las lecturas devuelven los mismos bytes.
test_read_returns_only_filled_rows — cache.read(layer=0) devuelve forma (B, H, cursor, d_h), no el tensor pre-asignado entero.
test_capacity_exceeded_raises — añadir más allá de max_seq lanza un CacheFullError personalizado.
test_dtype_preserved — allocate fp16, añadir tokens fp16, la lectura devuelve fp16.
test_reset_empties_cursor — cache.reset() pone el cursor a cero; el buffer subyacente no se requiere que esté a cero.
test_independent_layers — escribir en la capa 1 no perturba la capa 0.
test_independent_batch_entries — escribir para el índice de batch 0 no perturba el índice de batch 1.
test_memory_footprint_matches_formula — cache.bytes_allocated() es exactamente igual a \(2 \cdot L \cdot H \cdot d_h \cdot S_\text{max} \cdot B \cdot s\).

Cada test debería tener 5–15 líneas. Si el tuyo es más largo, la implementación está peleando con el test.

Bloque C — implementar `KVCache`¶

API según BLUEPRINT.md:

class KVCache:
    @classmethod
    def allocate(cls, *, layers: int, heads: int, head_dim: int,
                 max_seq: int, batch: int, dtype: np.dtype) -> "KVCache": ...
    def append(self, layer: int, k_new: np.ndarray, v_new: np.ndarray) -> None: ...
    def read(self, layer: int) -> tuple[np.ndarray, np.ndarray]: ...
    def current_length(self) -> int: ...
    def reset(self) -> None: ...
    def bytes_allocated(self) -> int: ...

Restricciones:

Todos los métodos públicos anotados con tipos. mypy --strict limpio.
Sin dependencias externas más allá de numpy + stdlib.
append es O(1) por token (sin concatenación). Este es el único propósito del caché; si te encuentras alcanzando np.concatenate, relee theory/02.
El cursor avanza una vez por llamada a append. Añadir K y V juntos avanza el cursor en 1, no en 2. (Off-by-one fácil — los appends de K y V son conceptualmente el mismo "paso".)
Todas las capas comparten el mismo cursor (todas procesan el mismo token en el mismo paso).

Bloque D — conectar a la atención¶

Actualizar src/minimodel/attention.py:

def attention(q: np.ndarray, k: np.ndarray, v: np.ndarray, *,
              mask: np.ndarray | None = None,
              cache: KVCache | None = None,
              layer_idx: int | None = None) -> np.ndarray:
    """If cache is None: original Phase-15 path (training).
       If cache is not None: append (k, v) to cache, read full cached K, V, do attention.
       layer_idx must be supplied if cache is supplied."""

Restricciones:

El camino de training (cache=None) está sin cambios en lo numérico. Los tests de la Fase 15 deben seguir pasando byte a byte idénticamente.
El camino de decode (cache is not None) debe tomar q de forma (B, H, 1, d_h) (longitud de secuencia 1) y añadir (k, v) de la misma forma.
No hace falta máscara nueva en el camino de decode (ver theory/01-prefill-vs-decode.md §Pseudo-pseudocódigo).

Bloque E — manifest¶

Commitea un manifest.json en experiments/22-cache-impl/:

{
  "experiment": "22-cache-impl",
  "date": "YYYY-MM-DD",
  "seed": 42,
  "versions": {"python": "3.11.x", "numpy": "X.Y.Z"},
  "tests": {"total": 11, "passed": null, "skipped": 0},
  "lines_added": null,
  "mypy_strict_clean": null,
  "ruff_clean": null
}

Rellena los nulls después de que pasen los tests.

Restricciones¶

Solo NumPy. Aún no PyTorch (la Fase 24 lo introduce).
Sin np.concatenate en append. Escritura O(1) en buffer pre-asignado; este es el contrato entero de corrección/perf.
Sin pasada hacia atrás. El caché es solo para inferencia. El training nunca lo usa (la Fase 17 entrena sin caché).
Sin threading. El caché es de un solo stream.

Condiciones de parada¶

Hecho cuando:

Los 11 tests en tests/test_minicache.py pasan.
mypy --strict src/minicache limpio.
ruff check src/minicache limpio.
Los tests de atención de la Fase 15 siguen pasando (es decir, no rompiste el camino de training).
manifest.json commiteado.
src/minicache/README.md refleja el API final (mantenido sincronizado con BLUEPRINT.md según A5).

Escollos (leer antes de depurar)¶

append(layer, k, v) avanza el cursor dos veces si lo incrementas por llamada. El cursor avanza una vez por token, no una vez por par (layer, k_or_v). Lleva la cuenta con cuidado — ver BLUEPRINT.md §Pitfalls.
Olvidar layer_idx en la llamada de atención. Sin él, el caché no puede saber qué K, V de qué capa leer.
Mutar slices devueltos. cache.read(layer) devuelve una vista del buffer pre-asignado. Mutarla corrompe el caché. Documenta el contrato: read devuelve una vista; no escribas.
Dependencia de orden de tests. Si test_capacity_exceeded_raises corre tras otros, el caché podría estar ya en cursor=0 por las fixtures; asegúrate de que cada test crea un caché nuevo.

Cuándo consultar `solutions/`¶

Después de que se cumplan todas las condiciones de parada. La referencia en solutions/01-implement-cache-ref.md (escrita al abrir la fase) recorre la decisión de layout y el invariante de gestión del cursor.

Siguiente lab: lab/02-correctness-test.md.