Skip to content

English · Español

03 — PagedAttention y sliding window

🇪🇸 Flash optimiza el kernel. Paged optimiza el almacenamiento del KV cache entre peticiones. Sliding-window cambia la complejidad de O(N²) a O(N·W). Tres optimizaciones independientes, todas argumentos de roofline desde perspectivas distintas.

PagedAttention: el KV cache como memoria virtual

El problema que PagedAttention resuelve

Durante la generación autoregresiva, cada paso añade una nueva fila K, V por capa por cabeza de attention. Para batch=1 esto está bien — KV crece linealmente. Para un servidor que batchea N peticiones de longitudes variadas:

  • Pre-asignar max_seq_len × n_layers × n_heads × d por slot desperdicia la mayor parte de la memoria (la mayoría de peticiones son cortas).
  • Asignar dinámicamente y redimensionar causa fragmentación: incluso si la memoria libre total es suficiente, no existe un bloque contiguo del tamaño requerido.

Números concretos del paper de vLLM: un modelo de 13B parámetros con max_seq_len=2048 en una A100 (80GB) puede servir ~10 peticiones concurrentes con pre-asignación naive. Con PagedAttention, el mismo hardware sirve ~30 peticiones concurrentes. 3× throughput sólo desde gestión de memoria, sin cambios en el kernel.

El mecanismo

vLLM trata el KV cache como memoria virtual:

  1. La memoria física se divide en páginas de tamaño fijo, cada una almacenando block_size tokens de K (o V) para una capa una cabeza. Default block_size = 16.
  2. Cada secuencia tiene una tabla de páginas que mapea sus posiciones lógicas a páginas físicas.
  3. Los kernels de attention siguen las indirecciones de la tabla de páginas: en vez de K[t, :] (plano), calculan K[page_table[t / block_size]][t % block_size, :].

Propiedades que esto da:

  • Sin fragmentación externa. Todas las páginas son del mismo tamaño; el asignador nunca tiene que encontrar un bloque contiguo.
  • Copy-on-write para cachear prefijos. Dos secuencias que comparten un prefijo (p. ej., el mismo system prompt) comparten páginas físicas; al divergir, sólo las páginas que divergen se copian.
  • Append barato. Añadir un token asigna una página (si la página actual está llena) o escribe en una página existente. Sin re-asignación del tensor KV entero.

El coste

Indirección de tabla de páginas por acceso de attention. En GPUs, este es un patrón de acceso a memoria no-coalesced — más lento que un array plano. El kernel de vLLM mitiga:

  1. Cargando las páginas relevantes en shared memory una vez por bloque de query.
  2. Usando block_size = 16 (suficientemente grande para que los lookups de tabla de páginas amorticen sobre múltiples accesos a memoria).

Overhead real: ~5% por llamada de attention, eclipsado por la ganancia de throughput desde mejor utilización.

Composición con Flash

FlashAttention calcula la salida de una cabeza tileando K, V. PagedAttention almacena K, V en páginas. Composición: el paso del kernel Flash "cargar un tile B_c × d de K" se convierte en "buscar los índices de página, cargar las filas de las páginas".

La implementación real de vLLM tiene un kernel a medida (paged_attention_v2) que fusiona ambas ideas. No la re-implementamos en la Fase 27 — la leemos.

Sliding window attention

La afirmación de complejidad

La attention estándar es O(N²) en tiempo y (para implementaciones naive) O(N²) en memoria. Para contextos largos (N=32k, 100k, 1M), esto es brutal incluso en hardware moderno.

Sliding window attention restringe cada token a atender sólo a los últimos W tokens (típicamente W=4096 para Mistral 7B). La complejidad cae a O(N·W) — lineal en N, que es la única forma en que la inferencia de contexto largo es tratable.

Lo que cambia mecánicamente

En el cómputo de attention, la máscara ya no es "causal" (triángulo inferior) sino "causal dentro de una banda" (triángulo inferior recortado a anchura W). Concretamente: la posición i sólo puede atender a las posiciones j ∈ [max(0, i-W+1), i].

En el kernel Flash: el bucle interno sobre los tiles K, V sólo necesita procesar los tiles dentro de la ventana. El bucle externo sigue siendo sobre los tiles de Q. Para Q_tile i en filas i*B_r a (i+1)*B_r - 1, los tiles K, V válidos abarcan columnas max(0, i*B_r - W) a (i+1)*B_r - 1. Salta el resto.

Composición con Flash

Trivialmente componible. El enmascarado se calcula a nivel de tile (un tile totalmente dentro de la ventana → sin máscara por elemento; un tile que cruza la frontera → enmascara sólo esa frontera). En el kernel: un parámetro extra constexpr SLIDING_WINDOW: int y una condicional sobre qué tiles entrar.

El lab 02 (kernel Triton) tiene sliding window como objetivo extra — no está en el DoD, pero se menciona como una extensión natural.

Cacheo de prefijos

Un caso específico donde PagedAttention brilla: muchas peticiones comparten el mismo system prompt o ejemplos few-shot. El KV para esos tokens se calcula una vez, se almacena en páginas, y se reutiliza entre peticiones vía la tabla de páginas.

Mecanismo: al enviar la petición, hashea los tokens del prefijo. Busca en un cache global (LRU). Si hay hit, la tabla de páginas de la petición empieza con punteros a las páginas cacheadas; sólo los tokens del sufijo disparan nuevo cómputo de KV.

Para un servidor con muchas consultas cortas que comparten un system prompt largo, el cacheo de prefijos puede dar 5–10× de ganancia de throughput sobre el 3× de PagedAttention. Lo mencionamos pero no lo implementamos.

Prefill troceado

La fase de prefill (calcular KV para el prompt inicial) es la fase más compute-intensiva de la inferencia. El prefill troceado procesa el prefill en trozos de C tokens a la vez, intercalado con pasos de decoding de otras peticiones en el batch.

La ganancia: mientras una petición está en la fase lenta de prefill, la GPU no está ociosa para los pasos rápidos de decode de otras peticiones. El throughput sube significativamente para batches mixtos.

El coste: scheduling más complejo; trocear demasiado fino introduce overhead. vLLM ajusta C automáticamente.

Ejercicio de lectura: block_manager.py de vLLM

El lab 03 le pide a Borja leer block_manager.py (el archivo que asigna y libera páginas KV) y anotarlo. Preguntas clave a contestar en las anotaciones:

  1. ¿Qué estructuras de datos rastrean páginas libres vs asignadas?
  2. ¿Cómo se representan las tablas de páginas? ¿Lista por secuencia de IDs de bloque físicos?
  3. ¿Qué pasa cuando una secuencia necesita crecer pero no quedan páginas libres? (Pista: eviction.)
  4. ¿Cómo se dispara el copy-on-write? (Pista: cuando dos secuencias que comparten una página divergen.)
  5. ¿Cuál es el ciclo de vida de una página a lo largo de la vida de una petición?

La lectura se evalúa por la calidad de la anotación, no por re-implementar el archivo.

Interpretación del roofline

Estas tres ideas — PagedAttention, sliding window, cacheo de prefijos — no cambian el punto del roofline de un único kernel. Cambian cuántos kernels de attention de peticiones pueden ejecutarse en el mismo hardware a la vez (PagedAttention), qué tan grande es N efectivamente por kernel (sliding window), y con qué frecuencia un kernel se ejecuta en absoluto (cacheo de prefijos).

En otras palabras, son argumentos de roofline a nivel de sistema, no a nivel de kernel:

  • PagedAttention: más peticiones por GPU → mayor throughput agregado.
  • Sliding window: cada kernel procesa menos trabajo → menor latencia por kernel.
  • Cacheo de prefijos: menos llamadas a kernel → menos FLOPs totales.

Cada uno es un "eje" distinto de optimización. Flash + Paged + Sliding + Prefix Cache componen; vLLM hace los cuatro.

Problemas de práctica

Soluciones al abrir la fase en solutions/03-paged-sliding-ref.md.

  1. Un servidor ejecuta LLaMA-13B con N=2048 longitud máxima. Tamaño de página = 16 tokens. KV por token (fp16, 40 capas, 40 cabezas, d=128) = 2 × 40 × 40 × 128 × 2 = 819 KiB. ¿Tamaño de página en bytes? ¿Cuántas páginas caben en 70 GiB?
  2. Sliding window W=4096 en N=32k. La complejidad de attention por query cae de N a W. ¿Por qué factor cae el compute total de attention para toda la secuencia? (Pista: N × N a N × W.)
  3. Cacheo de prefijos: 100 peticiones comparten todas un system prompt de 500 tokens. Sin cacheo de prefijos, KV se recomputa 100 veces. Con cacheo de prefijos, una vez. ¿Cuál es el ahorro en compute de prefill de attention?
  4. El lookup de tabla de páginas de PagedAttention añade un acceso extra a memoria por elemento de attention. Estima el overhead de latencia en ciclos asumiendo que la tabla de páginas cabe en cache L1. Argumenta si importa.

Recap de un párrafo

PagedAttention trata el KV cache como memoria virtual: páginas de tamaño fijo, tablas de páginas por secuencia, copy-on-write para prefijos compartidos. Esto elimina la fragmentación entre peticiones, subiendo el throughput del servidor ~3× sin cambiar ningún kernel. Sliding window attention restringe el campo receptivo a W tokens recientes, bajando la complejidad por paso de O(N) a O(W) — la única forma en que contextos de millones de tokens son tratables. El cacheo de prefijos reutiliza KV entre peticiones que comparten un prompt. Los tres componen con FlashAttention a nivel de kernel; juntos son lo que hace vLLM. Para la Fase 27 leemos el asignador de PagedAttention en block_manager.py de vLLM en lugar de re-implementarlo — el contenido algorítmico es pequeño pero la superficie de ingeniería es grande. El siguiente archivo de teoría mira GQA, MQA, y MLA — trucos ortogonales de reducción de KV.

Siguiente: theory/04-gqa-mqa-mla.md.