KVCacheManager: ciclo de vida de la caché KV a nivel de petición

源码版本v0.25.1

Responsabilidades

KVCacheManager es la interfaz de la caché KV (KV cache) que ve el planificador (scheduler). Encapsula acciones como "asignar N bloques a esta petición", "cuando esta petición termine devuelve sus bloques" o "mira cuánto prefijo cacheado hace match en este prompt" en tres APIs de alto nivel: allocate_slots / free / get_computed_blocks. Internamente delega toda la coordinación entre distintos tipos de caché KV (full attention / sliding window / estado Mamba, etc.) a KVCacheCoordinator, y deja el pool de bloques físicos en manos de BlockPool. Al planificador le basta con obtener el resultado inmutable KVCacheBlocks (kv_cache_manager.py:29-50), sin saber si debajo hay un solo tipo o un modelo híbrido.

En el constructor, KVCacheManager.__init__ monta el coordinator vía get_kv_cache_coordinator y se cuelga block_pool / num_kv_cache_groups / kv_cache_config (kv_cache_manager.py:147-183). El watermark (watermark_blocks = int(watermark * num_blocks)) es un margen reservado para peticiones ya planificadas; solo aplica a estados WAITING / PREEMPTED, evitando que una petición nueva recién llegada expulse a otra que ya llevaba la mitad del trabajo (kv_cache_manager.py:164-167) (kv_cache_manager.py:367-374). El manager casi no guarda estado por petición: todo el req_to_blocks vive dentro de SingleTypeKVCacheManager.

Motivación de diseño

¿Por qué el manager parece tan fino, con toda la lógica real en el coordinator?

• Separación de responsabilidades: KVCacheManager cubre la semántica "orientada a petición" (allocate_slots / free / get_computed_blocks); KVCacheCoordinator cubre la coordinación "orientada a grupo de caché KV" (en modelos híbridos, cómo asignan juntos varios grupos); BlockPool cubre el pool "orientado a bloque físico". Tres capas, cada una con su incumbencia; el manager hace de pegamento.
• KVCacheBlocks es la interfaz estable: blocks[i][j] representa el j-ésimo bloque del i-ésimo grupo de caché KV; get_block_ids() lo convierte a tuple[list[int], ...] para que lo consuma el attention backend (kv_cache_manager.py:73-88). __add__ permite concatenar los resultados de dos allocate_slots (<SrcLink path="vllm/v1/core/kv_cache_manager.py" lines="52-59" label="kv_cache_manager.py"/>); new_empty devuelve el marcador de "sin bloques".
• Commit diferido de la caché: la rama delay_cache_blocks=True de allocate_slots se salta cache_blocks durante la transferencia asíncrona en P/D (kv_cache_manager.py:448-451); el hash se marca solo cuando el KV remoto realmente llega, para evitar cacheos precipitados que producirían hits erróneos.
• El watermark solo aplica a peticiones en espera: has_scheduled_reqs and request.status in (WAITING, PREEMPTED) añade el watermark únicamente en esos casos; las asignaciones incrementales de peticiones ya en marcha no necesitan reserva (kv_cache_manager.py:367-374).
• Chequeo de admisión de secuencia completa: cuando full_sequence_must_fit=True, primero se estima cuántos bloques necesita todo el prompt; si se excede, se devuelve None directamente, evitando que el chunked prefill meta una petición que luego se queda colgada (kv_cache_manager.py:376-391).
• reserved_blocks deja sitio para conectores KV asíncronos: los prefills ya in-flight dependen de ciertos bloques; los connectors asíncronos que entran nuevos no pueden robarlos, así que si required_blocks > available_blocks se rechaza la asignación (kv_cache_manager.py:420-426).

Archivos clave

• KVCacheBlocks dataclass:29-50 — resultado inmutable visto por el planificador; blocks[i][j] = bloque j del grupo i.
• KVCacheBlocks.__add__:52-59 — concatena dos resultados de asignación; muy usado en asignación incremental.
• KVCacheBlocks.get_block_ids:73-88 — convierte a tuple[list[int], ...] para el attention backend.
• KVCacheManager.__init__:114-183 — monta el coordinator, calcula watermark_blocks, construye la caché empty_kv_cache_blocks para evitar GC.
• get_computed_blocks:206-246 — llama a coordinator.find_longest_cache_hit para buscar coincidencia de prefijo; devuelve (KVCacheBlocks, num_new_computed_tokens).
• allocate_slots:248-343 — tres fases: liberar bloques saltados → procesar tokens del prefijo → asignar bloques a los tokens nuevos; si falla, devuelve None.
• full_sequence_must_fit:376-391 — chequeo de admisión: bloques del prompt completo + watermark deben ser ≤ bloques libres.
• required_blocks gate:410-426 — available = free - reserved; si required > available, rechaza la asignación.
• allocate_new + cache_blocks:428-464 — delega al coordinator la asignación de bloques nuevos + llama cache_blocks para confirmar el hash.
• free:466-474 — invoca coordinator.free(request_id); libera en orden inverso para que el tail block se evicte primero.
• pop_blocks_for_free:495-506 — retira la contabilidad por petición sin devolver bloques al BlockPool; se usa en transferencias P/D.
• reset_prefix_cache:516-530 — invalida toda la caché tras un cambio de pesos en RLHF.

Flujo de datos

En cada paso y para cada petición, el planificador primero llama get_computed_blocks para ver el hit de prefijo, y luego allocate_slots para apartar los bloques de los tokens nuevos. Esto es lo que hace allocate_slots tras el chequeo de watermark + admisión:

if (
    new_computed_block_list is not self.empty_kv_cache_blocks.blocks
    or num_external_computed_tokens > 0
):
    # Append the new computed blocks to the request blocks until now to
    # avoid the case where the new blocks cannot be allocated.
    self.coordinator.allocate_new_computed_blocks(
        request_id=request.request_id,
        new_computed_blocks=new_computed_block_list,
        num_local_computed_tokens=num_local_computed_tokens,
        num_external_computed_tokens=num_external_computed_tokens,
    )

new_blocks = self.coordinator.allocate_new_blocks(
    request.request_id,
    num_tokens_need_slot,
    num_tokens_main_model,
    num_encoder_tokens,
)

# P/D: delay caching blocks if we have to recv from
# remote. Update state for locally cached blocks.
if not self.enable_caching or delay_cache_blocks:
    return self.create_kv_cache_blocks(new_blocks)

(kv_cache_manager.py:428-451) Primero absorbe los bloques del hit de prefijo (subiendo ref_cnt), luego pide al coordinator que asigne bloques físicos para los tokens nuevos, y finalmente, salvo que haya que diferir el cacheo, llama cache_blocks para marcar el hash de los bloques recién llenados en la caché de prefijo (ver Coordinator y BlockPool). El KVCacheBlocks devuelto se convierte en block_table vía get_block_ids y se inyecta en la tabla de bloques para que la consuma el kernel de attention.

Límites y fallos

• num_new_tokens == 0 y sin tokens externos: lanza ValueError; el llamador debe garantizar que o bien hay tokens nuevos que calcular, o bien hay tokens ya calculados externos que almacenar (kv_cache_manager.py:346-350).
• get_computed_blocks con hit completo a falta de un token: max_cache_hit_length = request.num_tokens - 1, porque en hit completo hay que recalcular el último token para sacar logits; por eso la longitud del hit se acota a prompt_length - 1 (kv_cache_manager.py:225-231).
• Peticiones que se saltan la prefix cache: con request.skip_reading_prefix_cache=True o enable_caching=False, get_computed_blocks devuelve vacío directamente (kv_cache_manager.py:222-223), por ejemplo en peticiones que necesitan prompt logprobs.
• reset_prefix_cache se niega: si num_used_blocks != 1 (solo null block contado) es que aún hay bloques sin liberar; el reset falla y devuelve False (kv_cache_manager.py:516-530).
• pop_blocks_for_free no devuelve al pool: el llamador debe llamar block_pool.free_blocks en orden inverso por su cuenta; si no, los bloques se filtran (kv_cache_manager.py:495-506).
• take_events anota metadatos de grupo: el evento BlockStored solo trae group_idx; el manager completa kind y sliding_window con kv_cache_event_metadata (kv_cache_manager.py:566-589).

Resumen

KVCacheManager es la fachada de la caché KV que ve el planificador; quienes realmente trabajan son KVCacheCoordinator + BlockPool, desglosados en Coordinator y BlockPool. Los block id asignados acaban en la tabla de bloques, donde el kernel de attention los usa para localizar la caché KV física; los pesos del modelo se cargan por la línea de DefaultModelLoader, sin acoplamiento con esta capa.

Véase la documentación oficial: Documentación de vLLM · README.