KV-cache restore: el costo que nadie mide hasta que escala

Cuando un runtime de inferencia ofrece guardar y restaurar el KV-cache, la promesa implícita es que reanudar contexto sale más barato que recalcularlo. Esa promesa se sostiene en el papel hasta que alguien la corre contra una carga real. Esta semana, un fix en llama.cpp mostró qué tan lejos puede estar la implementación “correcta” de la implementación que efectivamente escala, y de paso destapó un segundo problema que la primera solución casi introduce sin querer.

Una copia por celda, hasta que deja de alcanzar

El PR #27991 (build b10724) describe el caso: al restaurar un estado guardado en celdas de destino no contiguas, por ejemplo una snapshot de prompt cache que cae sobre un ring fragmentado, state_read_data hacía una copia pequeña por cada celda del KV-cache. Para un restore de más de 40 mil tokens eso eran cerca de 1.363.616 copias de unos pocos KiB cada una, y tomaba entre 25 y 63 segundos en el backend CUDA.

La solución no cambia qué se restaura, cambia cómo se agrupa: como la snapshot guarda las filas de celdas en orden, una corrida máxima de índices de destino consecutivos mapea a un bloque contiguo y se puede restaurar con una sola copia. Precalculan esas corridas una vez y las reusan en los tres loops de scatter (K, V y V transpuesta). El resultado, medido en un setup de producción real: de 1.363.616 copias en 25-63 segundos a 224 copias en 221-424 milisegundos, para el mismo restore de 42.603 celdas en 4 corridas.

El chunking no es un detalle: casi rompe el propio fix

Lo interesante no es solo la ganancia, es lo que casi se rompe al lograrla. El lector on-device copiaba el estado guardado con una copia 1:1 por índice de tensor cuando ambos lados (guardado y restaurado) registraban el mismo número de tensores, protegido por un assert de tamaño por tensor. Pero que el conteo de tensores coincida no significa que el chunking coincida: el restore ahora agrupa por corrida contigua de destino, mientras que el guardado usa lecturas por rango, así que ambos lados pueden terminar con el mismo número de tensores partiendo los mismos datos de formas distintas. Eso hacía abortar el assert.

El commit de seguimiento compara los tamaños por tensor y solo toma el camino 1:1 cuando el chunking realmente coincide; si no, cae a una copia por cursor de bytes, válida porque ambos lados enumeran la misma data lógica en el mismo orden. Es el patrón clásico de cualquier sistema de snapshot/restore: optimizar un lado sin verificar que el otro lado asuma la misma partición de los datos es la forma más común de introducir un bug de integridad al mismo tiempo que arreglas uno de performance.

La misma semana, tres frentes contra el costo de arrancar

El restore del KV-cache no fue el único frente. El resumen de la versión v0.4.0 de llama.cpp agrupa, en el mismo ciclo, lectura de tensores bajo demanda (“lazy tensor reading”, PR #27794, con el flag renombrado de --tensor-read-lazy a --lazy-mode y el atajo -lzm en el PR #27969), un tope de RAM para el cuantizador (#27795), streaming de cuantización por slabs de filas en vez de tenerlo todo en memoria (#27830), y una corrección explícita para evitar picos de RAM durante la carga (#27483). Cuatro fixes distintos, un mismo blanco: el camino de arranque y restauración de estado, que normalmente se escribe simple y correcto primero, y solo se optimiza cuando alguien lo mete en un loop de producción con miles de restores por hora.

Qué revisar en tu propio stack

Si usas prompt caching o resume de sesión sobre cualquier motor de inferencia, pregunta cómo escala el restore: si es una llamada de copia por token o por celda, vas a pagarlo en cuanto el contexto crezca, aunque hoy no lo notes.

Si diseñas tu propio formato de snapshot, no asumas que el lado que guarda y el que restaura van a partir los datos de la misma forma. Un cambio de chunking sin ruta de fallback no corrompe datos en silencio, pero sí puede abortar el proceso entero: prefiere comparar tamaños explícitamente y caer a una copia por cursor cuando no coincidan.

Al evaluar el consumo de memoria de un runtime, mide el pico durante la carga del modelo, no solo el estado estable en inferencia. Lazy loading y streaming de cuantización son las mitigaciones concretas que adoptó llama.cpp esta semana, no una recomendación genérica.

Si actualizas llama.cpp, ajusta tus scripts: --tensor-read-lazy pasó a llamarse --lazy-mode, con -lzm como atajo.

Fuentes: https://github.com/ggml-org/llama.cpp/releases/tag/b10724, https://github.com/ggml-org/llama.cpp/releases/tag/b10700, https://github.com/ggml-org/llama.cpp/releases/tag/v0.4.0