Prefix Caching: reutilización de KV blocks entre peticiones

源码版本v0.25.1

Responsabilidades

Muchas peticiones comparten un mismo prefijo: system prompt, chat template, definiciones de funciones, ejemplos few-shot. Los tokens K/V de ese tramo ya se calcularon antes; si la siguiente petición puede reutilizarlos, se ahorra un prefill entero. La caché de prefijo (prefix caching) consiste precisamente en eso: para cada KV block se calcula una huella basada en «hash de contenido», y cuando llega una nueva petición se consulta primero ese hash; en caso de acierto se reutiliza la referencia al block existente, sin recalcular.

El interruptor que la activa es CacheConfig.enable_prefix_caching, por defecto True(cache.py:93).EngineCore.__init__, cuando enable_prefix_caching está activo o existe un KV connector, construye request_block_hasher(core.py:210-219); se trata en realidad de una clausura devuelta por get_request_block_hasher(kv_cache_utils.py:673) que parte la secuencia de tokens de la petición en bloques de tamaño hash_block_size y aplica un hash encadenado. El hash de cada block combina «hash del block padre + token ids del block actual + extra_keys(LoRA / multimodal)» y se lo pasa a la función de hash(kv_cache_utils.py:577-604), de modo que la huella de un block equivale a la huella de «todo el prefijo desde el inicio hasta ese block».

KVCacheManager delega la búsqueda de aciertos a KVCacheCoordinator.find_longest_cache_hit(kv_cache_coordinator.py:364), que a su vez lo baja al manager de un solo tipo. Tras un acierto, allocate_slots reutiliza esos blocks y solo reparte blocks nuevos para la parte no acertada. Scheduler.schedule, al planificar cada petición, invoca primero get_computed_blocks(kv_cache_manager.py:206); si hay aciertos, reduce el presupuesto de tokens a calcular.

Motivación de diseño

• Reutilización a nivel de block, no de token: el KV cache se parte en blocks según PagedAttention, así que el block es la unidad mínima natural, y el hash también se calcula por block: grano más grueso y consulta más rápida.
• Hash encadenado para preservar la semántica de prefijo: hash_block_tokens(parent_hash, tokens, extra_keys)(kv_cache_utils.py:577) hace que la huella de cada block incluya implícitamente la del prefijo completo, evitando falsos aciertos del tipo «parte central distinta pero extremo común».
• Varios algoritmos de hash: prefix_caching_hash_algo admite sha256 / sha256_cbor / xxhash / xxhash_cbor(cache.py:95-110); xxhash es más rápido pero no es criptográficamente seguro, así que en escenarios multiusuario debe usarse con cuidado.
• Coordinador específico para la ruta desactivada: KVCacheCoordinatorNoPrefixCache(kv_cache_coordinator.py:377) hace que find_longest_cache_hit devuelva directamente ((), 0), evitando cómputo de hash cuando enable_caching=False.
• extra_keys para multimodal / LoRA: generate_block_hash_extra_keys(kv_cache_utils.py:714) añade entradas multimodales y LoRA ID como claves extra, para que imágenes distintas no caigan en el mismo slot de block.
• Capas de atención no causal se fuerzan a off: EngineCore._initialize_kv_caches, cuando detecta un spec con non_causal=True, reescribe enable_prefix_caching a False(core.py:255-269), porque la atención bidireccional de los Prefix LM rompería la semántica de acierto de prefijo.
• Saltar la lectura de prefix cache: Request.get_skip_reading_prefix_cache(request.py:266-277) hace que las peticiones de prompt logprobs o pooling all no lean el prefijo, porque necesitan recalcularlo completo.
• watermark anti-thrashing: KVCacheManager.watermark_blocks(kv_cache_manager.py:164-167) reserva un puñado de blocks libres para evitar preempt a alta frecuencia.

Archivos clave

• enable_prefix_caching:93 — por defecto True, activa la caché de prefijo al arrancar.
• prefix_caching_hash_algo:95 — uno entre sha256 / sha256_cbor / xxhash / xxhash_cbor.
• EngineCore 构造 request_block_hasher:210-219 — elige la función de hash + init_none_hash + get_request_block_hasher.
• 非因果注意力强制关闭:255-269 — para Prefix LM / atención bidireccional, enable_prefix_caching = False.
• Request.update_block_hashes:237-240 — invoca _block_hasher(self) para hashear los nuevos blocks ya completos.
• get_skip_reading_prefix_cache:266-277 — prompt logprobs / pooling saltan el acierto explícitamente.
• hash_block_tokens:577 — hash encadenado: hash_function((parent_hash, token_ids_tuple, extra_keys)).
• get_request_block_hasher:673 — devuelve una clausura que solo hashea blocks completos y encadena el hash del padre.
• get_computed_blocks:206 — entrada: llega petición → se busca el acierto más largo → se devuelve (KVCacheBlocks, num_computed_tokens).
• allocate_slots:248 — reutiliza la parte acertada y pide blocks nuevos al BlockPool para el resto.
• KVCacheCoordinatorNoPrefixCache:377 — stub del coordinator para la ruta desactivada.
• UnitaryKVCacheCoordinator.find_longest_cache_hit:480 — búsqueda de aciertos para modelos con un único grupo de KV.
• Scheduler 调用 get_computed_blocks:724-726 — al planificar, obtiene el número de aciertos y reduce num_new_tokens.

Flujo de datos

Cuando llega una petición, Scheduler invoca primero get_computed_blocks; por dentro, este paso pasa por find_longest_cache_hit en el coordinator, recorriendo la lista de block_hashes y consultando uno a uno al BlockPool si ya existe un block con ese mismo hash:

# vllm/v1/core/kv_cache_manager.py L206-L246
def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]:
    """Get the computed (cached) blocks for the request.
    Note that the computed blocks must be full.
    """
    # We skip finding the prefix cache hit when prefix caching is
    # disabled or the request is marked as skipping kv cache read
    # (which happens when the request requires prompt logprobs
    # or calls a pooling model with all pooling).
    if not self.enable_caching or request.skip_reading_prefix_cache:
        return self.empty_kv_cache_blocks, 0

    # NOTE: When all tokens hit the cache, we must recompute the last token
    # to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1.
    # This can trigger recomputation of an entire block, rather than just
    # the single last token, because allocate_slots() requires
    # num_computed_tokens to be block-size aligned. Removing this limitation
    # could slightly improve performance in the future.
    max_cache_hit_length = request.num_tokens - 1
    computed_blocks, num_new_computed_tokens = (
        self.coordinator.find_longest_cache_hit(
            request.block_hashes, max_cache_hit_length
        )
    )

    if self.log_stats:
        assert self.prefix_cache_stats is not None
        self.prefix_cache_stats.record(
            num_tokens=request.num_tokens,
            num_hits=num_new_computed_tokens,
            preempted=request.num_preemptions > 0,
        )

    return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens

El hash de cada block se calcula dentro de la propia petición en update_block_hashes; request_block_hasher es una clausura que solo hashea blocks completos y va encadenando el hash del padre:

# vllm/v1/core/kv_cache_utils.py L687-L728
def request_block_hasher(request: Request) -> list[BlockHash]:
    start_token_idx = len(request.block_hashes) * hash_block_size
    num_tokens = request.num_tokens

    if start_token_idx + hash_block_size > num_tokens:
        # Early stop when there no new full blocks created.
        return []

    curr_mm_idx = 0
    if start_token_idx > 0:
        # Set curr_mm_idx = -1 to indicate the last mm input.
        # Note that since we reach to this branch only when the block is
        # completed with generated tokens, we only need to consider the
        # last mm input.
        curr_mm_idx = -1

    prev_block_hash_value = (
        request.block_hashes[-1] if request.block_hashes else None
    )
    new_block_hashes: list[BlockHash] = []
    while True:
        end_token_idx = start_token_idx + hash_block_size
        if end_token_idx > num_tokens:
            # We only hash full blocks
            break

        # MM and LoRA requests need extra keys for block-hash computation.
        extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
            request, start_token_idx, end_token_idx, curr_mm_idx
        )

        # Compute the hash of the current block
        block_tokens = request.all_token_ids[start_token_idx:end_token_idx]
        block_hash = hash_block_tokens(
            caching_hash_fn, prev_block_hash_value, block_tokens, extra_keys
        )

        new_block_hashes.append(block_hash)
        start_token_idx += hash_block_size
        prev_block_hash_value = block_hash

    return new_block_hashes

Una vez que Scheduler tiene el número de aciertos, resta esa parte de num_new_tokens; allocate_slots solo pide al block pool los blocks no acertados. Los blocks acertados incrementan su contador de referencias en +1, pero no se recalculan sus tokens.

Límites y fallos

• Capas de atención no causal fuerzan prefix caching off: EngineCore._initialize_kv_caches, al detectar non_causal=True, pone enable_prefix_caching = False(core.py:265-269); de lo contrario, los aciertos de prefijo en atención bidireccional leerían KV incorrectos.
• Prompt logprobs / pooling saltan el acierto: si Request.get_skip_reading_prefix_cache devuelve True, get_computed_blocks retorna vacío(kv_cache_manager.py:222-223), porque estos escenarios deben recalcular el KV del prompt.
• Aciero total también recalcula el último token: max_cache_hit_length = request.num_tokens - 1(kv_cache_manager.py:231); si no, no habría logits para muestrear.
• xxhash no es criptográficamente seguro: prefix_caching_hash_algo = "xxhash" es más rápido, pero tiene mayor probacidad teórica de colisión y podría filtrar información de prefijos en escenarios multiusuario(cache.py:101-108).
• Ruta desactivada con stub propio: KVCacheCoordinatorNoPrefixCache.find_longest_cache_hit devuelve ((), 0)(kv_cache_coordinator.py:416-424), para que con enable_caching=False no se malgaste cómputo de hash.
• LRU en BlockPool: los blocks reutilizados por acierto suman +1 al contador de referencias y se decrementan al terminar la petición; los que nadie referencia a largo plazo son reclamados por LRU para nuevas peticiones(kv_cache_manager.py:508).

Resumen

La caché de prefijo es uno de los mecanismos clave de throughput en vLLM: aplica un hash encadenado «hash del padre + token ids + extra_keys» sobre cada KV block para reutilizar entre peticiones el cálculo de prefijos comunes. EngineCore, al inicializar, decide según enable_prefix_caching si construye el hasher; Scheduler.schedule llama get_computed_blocks en cada iteración para buscar aciertos. Los blocks reutilizados se gestionan en /kv-cache/kv-cache-manager;los detalles del BlockPool, en /kv-cache/coordinator-blockpool;la integración en el núcleo del motor, en /engine/engine-core.

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