KVCacheManager : cycle de vie du cache KV au niveau requête

源码版本v0.25.1

Responsabilités

KVCacheManager est l'interface du cache KV (KV cache) vue par l'ordonnanceur (scheduler). Il encapsule « allouer N blocs à cette requête », « la requête est terminée, rendre les blocs », « combien de préfixe est en cache pour ce prompt » en trois API de haut niveau : allocate_slots / free / get_computed_blocks. En interne, la coordination entre les différents types de cache KV (full attention / sliding window / état Mamba, etc.) est déléguée à KVCacheCoordinator, et la gestion du pool de blocs physiques à BlockPool. L'ordonnanceur ne récupère qu'un KVCacheBlocks immuable(kv_cache_manager.py:29-50), sans avoir à savoir s'il s'agit d'un modèle mono-type ou hybride.

À la construction, KVCacheManager.__init__ instancie le coordinator via get_kv_cache_coordinator, et attache block_pool / num_kv_cache_groups / kv_cache_config à lui-même(kv_cache_manager.py:147-183). Le watermark (watermark_blocks = int(watermark * num_blocks)) est une marge réservée aux requêtes déjà schedulées ; il ne s'applique qu'aux statuts WAITING / PREEMPTED, pour éviter qu'une nouvelle requête n'évince des requêtes déjà à mi-chemin(kv_cache_manager.py:164-167)(kv_cache_manager.py:367-374). Le manager ne détient quasiment aucun état per-request ; tous les req_to_blocks vivent dans SingleTypeKVCacheManager.

Motivation de conception

Pourquoi le manager paraît-il si fin, toute la logique étant dans le coordinator ?

• Séparation des responsabilités : KVCacheManager gère la sémantique « orientée requête » (allocate_slots / free / get_computed_blocks) ; KVCacheCoordinator gère la coordination « orientée groupe de cache KV » (en hybride, comment allouer conjointement entre plusieurs groupes) ; BlockPool gère le pool « orienté bloc physique ». Trois couches, chacune son rôle, le manager fait la colle.
• KVCacheBlocks est une interface stable : blocks[i][j] désigne le j-ième bloc du i-ème groupe de cache KV ; get_block_ids() le convertit en tuple[list[int], ...] pour l'attention backend(kv_cache_manager.py:73-88). __add__ permet de concaténer deux résultats d'allocate_slots(kv_cache_manager.py:52-59), et new_empty sert de placeholder « aucun bloc ».
• Commit différé du cache : la branche delay_cache_blocks=True d'allocate_slots saute cache_blocks lors du transfert asynchrone P/D(kv_cache_manager.py:448-451) ; on attend que le KV distant soit réellement arrivé avant de hacher, pour éviter un cache anticipé qui produirait des hits erronés.
• watermark uniquement pour les requêtes en attente : on n'ajoute le watermark que si has_scheduled_reqs and request.status in (WAITING, PREEMPTED) ; les requêtes déjà actives en allocation incrémentale n'ont pas besoin de réserve(kv_cache_manager.py:367-374).
• Admission sur séquence complète : quand full_sequence_must_fit=True, on estime d'abord le nombre de blocs nécessaires pour tout le prompt, et on renvoie None si ça dépasse — pour éviter qu'un chunked prefill ne laisse entrer une requête qui se bloque ensuite(kv_cache_manager.py:376-391).
• reserved_blocks pour les KV connectors asynchrones : les séquences de prefill déjà in-flight dépendent de certains blocs ; un connector asynchrone entrant ne peut pas les rafler — si required_blocks > available_blocks, l'allocation est refusée(kv_cache_manager.py:420-426).

Fichiers clés

• KVCacheBlocks dataclass:29-50 — résultat immuable côté ordonnanceur, blocks[i][j] = j-ème bloc du i-ème groupe.
• KVCacheBlocks.__add__:52-59 — concaténation de deux allocations, utilisée pour l'allocation incrémentale.
• KVCacheBlocks.get_block_ids:73-88 — convertit en tuple[list[int], ...] pour l'attention backend.
• KVCacheManager.__init__:114-183 — installe le coordinator, calcule watermark_blocks, construit le cache empty_kv_cache_blocks pour éviter l'overhead GC.
• get_computed_blocks:206-246 — appelle coordinator.find_longest_cache_hit, renvoie (KVCacheBlocks, num_new_computed_tokens).
• allocate_slots:248-343 — en trois temps : libération des blocs skippés → traitement des tokens préfixe → allocation des blocs pour les nouveaux tokens ; renvoie None en cas d'échec.
• full_sequence_must_fit:376-391 — contrôle d'admission : nombre de blocs du prompt complet + watermark doit être ≤ nombre de blocs libres.
• required_blocks gate:410-426 — available = free - reserved, required > available refuse l'allocation.
• allocate_new + cache_blocks:428-464 — délègue au coordinator l'allocation de nouveaux blocs + appelle cache_blocks pour soumettre le hash.
• free:466-474 — appelle coordinator.free(request_id) ; libération en ordre inverse pour que les blocs tail soient évictés en premier.
• pop_blocks_for_free:495-506 — récupère la comptabilité per-request sans rendre au BlockPool, utilisé pour les transferts P/D.
• reset_prefix_cache:516-530 — invalide tout le pool de cache après un changement de poids en RLHF.

Flux de données

À chaque step, l'ordonnanceur appelle d'abord get_computed_blocks pour voir le hit de préfixe, puis allocate_slots pour allouer les blocs des nouveaux tokens. Voici les quelques lignes d'allocate_slots après le contrôle watermark + admission :

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) On absorbe d'abord les blocs hit par préfixe (en incrémentant ref_cnt), puis le coordinator alloue les blocs physiques pour les nouveaux tokens, et sauf si le cache doit être différé, on appelle cache_blocks pour hacher les blocs fraîchement remplis et les pousser dans le prefix cache (voir Coordinator et BlockPool). Le KVCacheBlocks renvoyé est converti en block_table via get_block_ids, puis injecté dans la block table pour les kernels d'attention.

Limites et échecs

• num_new_tokens == 0 sans token externe : lève ValueError ; l'appelant doit garantir soit qu'il y a des nouveaux tokens à calculer, soit qu'il y a des tokens externes déjà calculés à stocker(kv_cache_manager.py:346-350).
• get_computed_blocks un token de moins en cas de hit total : max_cache_hit_length = request.num_tokens - 1, car en cas de hit total il faut recalculer le dernier token pour obtenir les logits, donc la longueur hit est plafonnée à prompt_length - 1(kv_cache_manager.py:225-231).
• Requêtes qui sautent le prefix cache : si request.skip_reading_prefix_cache=True ou enable_caching=False, get_computed_blocks renvoie vide(kv_cache_manager.py:222-223) — par exemple les requêtes qui nécessitent des prompt logprobs.
• reset_prefix_cache refuse : si num_used_blocks != 1 (uniquement null block), c'est qu'il reste des blocs non libérés ; le reset échoue et renvoie False(kv_cache_manager.py:516-530).
• pop_blocks_for_free ne rend pas au pool : l'appelant doit lui-même appeler block_pool.free_blocks en ordre inverse, sinon les blocs fuitent(kv_cache_manager.py:495-506).
• take_events annote les métadonnées de groupe : l'événement BlockStored ne porte que group_idx ; le manager ajoute kind et sliding_window via kv_cache_event_metadata(kv_cache_manager.py:566-589).

Résumé

KVCacheManager est la façade du cache KV vue par l'ordonnanceur ; les véritables acteurs, KVCacheCoordinator + BlockPool, sont décrits dans Coordinator et BlockPool. Les block ids alloués atterrissent dans la block table, où les kernels d'attention s'en servent pour localiser le cache KV physique ; les poids du modèle, eux, sont chargés via DefaultModelLoader, sans couplage avec cette couche.

Voir la documentation officielle : Documentation vLLM · README