Skip to content

KVCacheCoordinator + BlockPool : coordination et pool de blocs pour modèles hybrides

源码版本v0.25.1

Responsabilités

KVCacheCoordinator est, dans vLLM v1, la couche de coordination au-dessus des multiples types de cache KV (full attention / sliding window / état Mamba / MLA, etc.). Chaque groupe de cache KV dispose d'un SingleTypeKVCacheManager (par ex. FullAttentionManager / SlidingWindowManager / RSWAManager) ; le coordinator répartit « allouer N tokens de blocs », « trouver le plus long préfixe en cache », « libérer une requête », « cacher un bloc rempli » en appelant le bon single-type manager selon le groupe. BlockPool est le pool de blocs physiques partagé : tous les groupes partagent le même ensemble de KVCacheBlock — seule la clé de hash porte un group_id pour les distinguer, donc un même bloc peut être caché simultanément par plusieurs groupes (par ex. full attention et sliding window pour le même token).

À la construction, KVCacheCoordinator.__init__ instancie un BlockPool, puis get_manager_for_kv_cache_spec crée un single-type manager par groupe(kv_cache_coordinator.py:91-120). La stratégie est choisie par la fabrique get_kv_cache_coordinator : sans prefix caching → KVCacheCoordinatorNoPrefixCache, un seul groupe → UnitaryKVCacheCoordinator, plusieurs groupes → HybridKVCacheCoordinator(kv_cache_coordinator.py:782-822). BlockPool maintient lui-même une FreeKVCacheBlockQueue (liste doublement chaînée) pour l'éviction LRU, et une table de hachage cached_block_hash_to_block pour l'index du prefix caching(block_pool.py:144-197).

Motivation de conception

Pourquoi scinder le coordinator en NoPrefix / Unitary / Hybrid ?

  • Repli à vide : KVCacheCoordinatorNoPrefixCache supporte un nombre de groupes quelconque (y compris 0), n'implémente aucune interface liée au prefix caching — find_longest_cache_hit renvoie directement (blocks vides, 0), get_num_common_prefix_blocks renvoie toujours 0(kv_cache_coordinator.py:413-424). Quand le prefix caching est désactivé, on n'emprunte aucun chemin de cache.
  • Chemin rapide mono-groupe : UnitaryKVCacheCoordinator suppose hash_block_size == block_size et un seul groupe(kv_cache_coordinator.py:471-476) ; find_longest_cache_hit appelle directement le single manager sans fusion(kv_cache_coordinator.py:480-496) pour éviter l'overhead de fusion spec-group du mode hybrid.
  • Agrégation par spec en Hybrid : HybridKVCacheCoordinator.verify_and_split_kv_cache_groups fusionne les groupes de même spec en un SpecGroup et fait la recherche de cache hit conjointement(kv_cache_coordinator.py:560-588), en plaçant la full attention en tête pour fournir aux groupes suivants une borne supérieure serrée(kv_cache_coordinator.py:590-595).
  • BlockPool partagé : tous les groupes partagent num_gpu_blocks blocs physiques ; les hashes des différents groupes sont préfixés via make_block_hash_with_group_id(block_hash, group_id), get_cached_block interroge par liste de group_id(block_pool.py:199-224) ; un même hash dans BlockHashToBlockMap peut correspondre aux blocs de plusieurs groupes.
  • LRU + priorité aux blocs de queue : FreeKVCacheBlockQueue est une liste doublement chaînée pour une suppression intermédiaire en O(1) ; à la libération, on réinsère en ordre inverse afin que le bloc tail soit évicté en premier(kv_cache_utils.py:179-197) ; free_blocks évite d'abord les blocs sans hash avant ceux qui en ont un(block_pool.py:622-635).
  • Sentinelle null_block : block_id=0 est le null_block, utilisé par remove_skipped_blocks pour remplacer les blocs hors sliding window ou preempted ; son ref_cnt n'est pas maintenu, il faut sauter à la libération(block_pool.py:188-192).

Fichiers clés

Flux de données

BlockPool.get_new_blocks est la véritable entrée d'allocation physique : il pop N blocs en tête de la free queue, evict un à un les anciens hashes, et met ref_cnt à 1 :

python
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
    if num_blocks > self.get_num_free_blocks():
        raise ValueError(f"Cannot get {num_blocks} free blocks from the pool")

    ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)

    if self.enable_caching:
        for block in ret:
            self._maybe_evict_cached_block(block)
            assert block.ref_cnt == 0
            block.ref_cnt += 1
            if self.metrics_collector:
                self.metrics_collector.on_block_allocated(block)
    else:
        for block in ret:
            assert block.ref_cnt == 0
            block.ref_cnt += 1
            if self.metrics_collector:
                self.metrics_collector.on_block_allocated(block)
    return ret

(block_pool.py:542-572) Appelé par SingleTypeKVCacheManager.allocate_new_blocks(single_type_kv_cache_manager.py:279-311), qui append les nouveaux blocs à req_to_blocks[request_id]. La libération passe par BlockPool.free_blocks : on décrémente ref_cnt, et pour les blocs dont ref_cnt tombe à zéro, on les réinsère en queue de free queue en distinguant « avec hash / sans hash »(block_pool.py:614-635). Quand le prompt d'une requête remplit un bloc, cache_full_blocks hache le contenu et l'écrit dans cached_block_hash_to_block ; à la requête suivante avec le même préfixe, get_cached_block hit directement, et touch incrémente ref_cnt sans recalculer(block_pool.py:199-224)(block_pool.py:597-612).

Limites et échecs

  • Dépassement du nombre de blocs libres : get_new_blocks ne pré-vérifie pas, lève directement ValueError(block_pool.py:553-554) ; l'appelant (KVCacheManager.allocate_slots) doit d'abord calculer via get_num_blocks_to_allocate(kv_cache_coordinator.py:130-185).
  • Hybrid ne supporte pas DCP / PCP : assert dcp_world_size == 1 and pcp_world_size == 1(kv_cache_coordinator.py:556-557).
  • Unitary impose hash_block_size == block_size : assert not enable_caching or hash_block_size == self.block_size(kv_cache_coordinator.py:471-473).
  • Hybrid nécessite au moins deux attention groups : assert len(self.attention_groups) > 1(kv_cache_coordinator.py:586-588).
  • Un bloc ayant déjà un hash ne peut être reset : set_block_hash utilise un assert pour empêcher les doublons(kv_cache_utils.py:153-157).
  • reset_prefix_cache refuse de nettoyer : si d'autres blocs que null block sont encore occupés, on log en warning et on renvoie False(block_pool.py:665-672).
  • evict_blocks reçoit des block_id externes : les ids de blocs remontés par le KV connector hors plage déclenchent un assert(block_pool.py:647-652).
  • Validation de retention_interval : _validate_prefix_cache_retention_interval exige un entier positif multiple de scheduler_block_size(kv_cache_coordinator.py:31-58).

Résumé

Coordinator + BlockPool constituent la couche physique sous KVCacheManager : le coordinator répartit vers les single-type managers selon la stratégie, BlockPool gère le pool de blocs physiques et l'index du prefix cache. Le résultat de l'allocation de blocs devient une liste de block_id, que la block table emballe en tensors consommables directement par les kernels d'attention.

Voir la documentation officielle : Documentation vLLM · README