Prefix Caching : réutilisation des blocs KV entre requêtes

源码版本v0.25.1

Responsabilités

Beaucoup de requêtes partagent un même préfixe : system prompt, chat template, définitions de fonctions, exemples few-shot. Les K/V de ces tokens ont déjà été calculés ; si la requête suivante peut les réutiliser, on économise tout un segment de prefill. C'est exactement le rôle du cache de préfixe (prefix caching) : on empreinte chaque bloc KV par un « hash de contenu », on consulte ce hash à l'arrivée d'une nouvelle requête, et en cas de hit on réutilise les blocs existants sans recalcul.

L'activation se fait via CacheConfig.enable_prefix_caching, True par défaut(cache.py:93). EngineCore.__init__ construit request_block_hasher quand enable_prefix_caching est vrai ou qu'un KV connector est présent(core.py:210-219) ; il s'agit en pratique d'une fermeture retournée par get_request_block_hasher(kv_cache_utils.py:673), qui découpe la séquence de tokens de la requête par tranches de hash_block_size et applique un chaînage de hachage. Le hash de chaque bloc alimente la fonction de hachage avec « hash du bloc parent + token ids du bloc courant + extra_keys (LoRA / multimodal) »(kv_cache_utils.py:577-604) ; l'empreinte d'un bloc est donc celle du préfixe complet « du début jusqu'à ce bloc ».

KVCacheManager délègue la recherche de hit à KVCacheCoordinator.find_longest_cache_hit(kv_cache_coordinator.py:364), qui redescend vers le manager monotype. En cas de hit, allocate_slots réutilise ces blocs et n'alloue de nouveaux blocs que pour la partie manquante. Scheduler.schedule appelle get_computed_blocks(kv_cache_manager.py:206) pour chaque requête planifiée, et réduit d'autant le budget de tokens à calculer quand il y a hit.

Motivation de conception

• Réutilisation au niveau bloc, pas au niveau token : le cache KV est découpé en blocs PagedAttention, le hash se calcule donc par bloc — granularité plus grossière mais recherche rapide.
• Hachage chaîné pour préserver la sémantique de préfixe : hash_block_tokens(parent_hash, tokens, extra_keys)(kv_cache_utils.py:577) fait que l'empreinte de chaque bloc encode tout le préfixe précédent ; impossible d'avoir un faux hit « milieu différent, fin identique ».
• Plusieurs algorithmes de hash : prefix_caching_hash_algo supporte sha256 / sha256_cbor / xxhash / xxhash_cbor(cache.py:95-110) ; xxhash est plus rapide mais sans garantie cryptographique, à utiliser avec prudence en multi-tenant.
• Coordinator dédié au chemin désactivé : KVCacheCoordinatorNoPrefixCache(kv_cache_coordinator.py:377) fait retourner ((), 0) à find_longest_cache_hit, pour éviter toute computation de hash quand enable_caching=False.
• extra_keys multimodal / LoRA : generate_block_hash_extra_keys(kv_cache_utils.py:714) prend les entrées multimodales et l'ID LoRA comme clés de hachage supplémentaires, pour éviter que des images différentes n'écrasent le même slot.
• Forçage à off pour les couches d'attention non causales : EngineCore._initialize_kv_caches met enable_prefix_caching = False quand un spec non_causal=True est détecté(core.py:255-269) ; l'attention bidirectionnelle des Prefix LM casserait la sémantique de hit de préfixe.
• Possibilité de skipper la lecture du cache : Request.get_skip_reading_prefix_cache(request.py:266-277) permet aux requêtes prompt logprobs / pooling all de ne pas consulter le cache, puisqu'elles doivent tout recalculer.
• Watermark anti-frétillement : KVCacheManager.watermark_blocks(kv_cache_manager.py:164-167) conserve une petite réserve de blocs libres pour limiter les préemptions à haute fréquence.

Fichiers clés

• enable_prefix_caching:93 — True par défaut, active ou non le cache de préfixe au démarrage.
• prefix_caching_hash_algo:95 — au choix parmi sha256 / sha256_cbor / xxhash / xxhash_cbor.
• EngineCore 构造 request_block_hasher:210-219 — sélectionne la fonction de hash + init_none_hash + get_request_block_hasher.
• 非因果注意力强制关闭:255-269 — enable_prefix_caching = False pour Prefix LM / attention bidirectionnelle.
• Request.update_block_hashes:237-240 — appelle _block_hasher(self) pour calculer le hash des blocs nouvellement remplis.
• get_skip_reading_prefix_cache:266-277 — prompt logprobs / pooling skippe explicitement les hits.
• hash_block_tokens:577 — hachage chaîné : hash_function((parent_hash, token_ids_tuple, extra_keys)).
• get_request_block_hasher:673 — retourne une fermeture qui ne hashe que les blocs complets, en chaînant le hash parent.
• get_computed_blocks:206 — point d'entrée : à l'arrivée d'une requête, recherche le plus long hit, retourne (KVCacheBlocks, num_computed_tokens).
• allocate_slots:248 — réutilise la partie touchée, prend de nouveaux blocs depuis BlockPool pour le reste.
• KVCacheCoordinatorNoPrefixCache:377 — coordinator stub pour le chemin désactivé.
• UnitaryKVCacheCoordinator.find_longest_cache_hit:480 — recherche de hit pour les modèles à groupe KV unique.
• Scheduler 调用 get_computed_blocks:724-726 — à la planification, récupère le nombre de hits pour réduire num_new_tokens.

Flux de données

À l'arrivée d'une requête, le Scheduler appelle d'abord get_computed_blocks, qui en interne passe par le coordinator et find_longest_cache_hit, lequel interroge pour chaque hash de block_hashes le BlockPool pour savoir si un bloc de même hash existe déjà :

# 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

Le hash de chaque bloc est calculé dans update_block_hashes de la requête ; request_block_hasher est une fermeture qui ne hashe que les blocs complets, en chaînant le hash parent :

# 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

Une fois le nombre de hits récupéré, le Scheduler diminue num_new_tokens d'autant, et allocate_slots ne prend dans le block pool que les blocs manquants. Les blocs en hit voient leur compteur de références incrémenté de 1, mais leurs tokens ne sont pas recalculés.

Limites et échecs

• Désactivation forcée pour les couches non causales : EngineCore._initialize_kv_caches met enable_prefix_caching = False quand non_causal=True est détecté(core.py:265-269) ; sinon l'attention bidirectionnelle lirait des KV erronés lors d'un hit de préfixe.
• Prompt logprobs / pooling skippe les hits : quand Request.get_skip_reading_prefix_cache retourne True, get_computed_blocks retourne immédiatement un résultat vide(kv_cache_manager.py:222-223) car ces scénarios doivent recalculer le KV du prompt.
• Hit complet → recalcul du dernier token : max_cache_hit_length = request.num_tokens - 1(kv_cache_manager.py:231) ; sinon aucun logit à échantillonner.
• xxhash non résistant cryptographiquement : prefix_caching_hash_algo = "xxhash" est plus rapide mais avec un risque de collision plus élevé en théorie, et peut divulguer des informations de préfixe en multi-tenant(cache.py:101-108).
• Stub dédié au chemin désactivé : KVCacheCoordinatorNoPrefixCache.find_longest_cache_hit retourne ((), 0)(kv_cache_coordinator.py:416-424), pour ne pas gaspiller de calcul de hash quand enable_caching=False.
• LRU dans BlockPool : les blocs réutilisés en hit voient leur compteur +1, décrémenté seulement en fin de requête ; les blocs sans référence depuis longtemps sont recyclés par LRU pour les nouvelles requêtes(kv_cache_manager.py:508).

Résumé

Le cache de préfixe est l'un des leviers clés du throughput de vLLM : il chaîne les blocs KV par « hash parent + token ids + extra_keys » et réutilise entre requêtes les calculs de préfixe identiques. À l'initialisation, EngineCore construit le hasher selon enable_prefix_caching, puis Scheduler.schedule appelle get_computed_blocks à chaque cycle pour mesurer les hits. Les blocs réutilisés sont gérés dans /kv-cache/kv-cache-manager ; les détails du BlockPool dans /kv-cache/coordinator-blockpool ; l'assemblage par le cœur du moteur dans /engine/engine-core.

Voir la documentation officielle : Documentation vLLM · README