Skip to content

KVCacheManager: KV-Cache-Lebenszyklus pro Anfrage

源码版本v0.25.1

Verantwortung

KVCacheManager ist die KV-Cache (KV-Cache)-Schnittstelle, die der Scheduler (Scheduler) sieht. Es kapselt Aktionen wie "dieser Anfrage N Blocks zuweisen", "die Anfrage ist fertig, Blocks zurückgeben" und "prüfen, wie viel Prefix-Cache (Prefix-Caching) diese Prompt trifft" in den drei High-Level-APIs allocate_slots / free / get_computed_blocks. Intern delegiert es die Koordination unterschiedlicher KV-Cache-Typen (Full Attention / Sliding Window / Mamba-Zustand etc.) vollständig an KVCacheCoordinator und das physische Block-Pool-Management an BlockPool. Der Scheduler erhält lediglich ein unveränderliches KVCacheBlocks-Ergebnis(kv_cache_manager.py:29-50), ohne wissen zu müssen, ob darunter ein einzelner Typ oder ein hybrides Modell liegt.

Beim Aufbau setzt KVCacheManager.__init__ über get_kv_cache_coordinator den Coordinator auf und hängt block_pool / num_kv_cache_groups / kv_cache_config an sich selbst an(kv_cache_manager.py:147-183). Der Wasserzeichen-Puffer (watermark_blocks = int(watermark * num_blocks)) ist die Reserve für bereits verplante Anfragen und wirkt sich nur auf Anfragen im Status WAITING / PREEMPTED aus, damit eine neu eintreffende Anfrage nicht sofort bereits halb abgearbeitete Anfragen verdrängt(kv_cache_manager.py:164-167)(kv_cache_manager.py:367-374). Der Manager hält selbst kaum Per-Request-Zustand vor; die gesamten req_to_blocks liegen in SingleTypeKVCacheManager.

Entwurfsmotivation

Warum wirkt der Manager so dünn und die eigentliche Logik sitzt vollständig im Coordinator?

  • Aufgabenschichtung: KVCacheManager ist für die "anfrageorientierte" Semantik zuständig (allocate_slots / free / get_computed_blocks), KVCacheCoordinator für die "an KV-Cache-Gruppen orientierte" Koordination (wie bei hybriden Modellen mehrere Gruppen gemeinsam zugewiesen werden), und BlockPool für das "physische Block-orientierte" Pooling. Jede Schicht kümmert sich um ihre eigenen Belange, der Manager fungiert als Klebstoff.
  • KVCacheBlocks als stabile Schnittstelle: blocks[i][j] bezeichnet den j-ten Block der i-ten KV-Cache-Gruppe; get_block_ids() wandelt ihn in ein tuple[list[int], ...] für das Attention-Backend um(kv_cache_manager.py:73-88). __add__ erlaubt es, die Ergebnisse zweier allocate_slots-Aufrufe zusammenzufügen(kv_cache_manager.py:52-59); new_empty liefert den Platzhalter für "kein Block".
  • Verzögerte Cache-Einreichung: Der delay_cache_blocks=True-Zweig von allocate_slots überspringt bei asynchronem P/D-Transfer cache_blocks(kv_cache_manager.py:448-451). Erst wenn der entfernte KV tatsächlich eingetroffen ist, wird der Hash markiert – ein zu frühes Cachen würde zu falschen Treffern führen.
  • Wasserzeichen nur für wartende Anfragen: has_scheduled_reqs and request.status in (WAITING, PREEMPTED) löst das Hinzufügen des Wasserzeichens aus; für bereits laufende Anfallen gilt die inkrementelle Zuweisung ohne Reserve(kv_cache_manager.py:367-374).
  • Admissionsprüfung für die gesamte Sequenz: Bei full_sequence_must_fit=True wird zuerst geschätzt, wie viele Blocks die gesamte Prompt benötigt; bei Überschreitung wird direkt None zurückgegeben, damit Chunked-Prefill eine Anfrage nicht erst aufnimmt und dann stecken bleibt(kv_cache_manager.py:376-391).
  • reserved_blocks für asynchrone KV-Connectoren: Bereits in-flight befindliche Prefill-Sequenzen hängen von bestimmten Blocks ab; ein neu eintreffender asynchroner Connector-Load darf diese nicht überbieten. Bei required_blocks > available_blocks wird die Zuweisung abgelehnt(kv_cache_manager.py:420-426).

Schlüsseldateien

Datenfluss

In jedem Scheduling-Schritt ruft der Scheduler für jede Anfrage zuerst get_computed_blocks auf, um den Prefix-Treffer zu prüfen, und dann allocate_slots, um die Blocks für die neuen Token zuzuweisen. Nachfolgend die wenigen Zeilen, die allocate_slots nach Wasserzeichen- und Admissionsprüfung tatsächlich ausführt:

python
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) Zuerst werden die getroffenen Prefix-Blocks absorbiert (ref_cnt wird erhöht), dann lässt der Coordinator für die neuen Token physische Blocks zuteilen. Sofern das Caching nicht verzögert wird, wird anschließend cache_blocks aufgerufen, um die gerade gefüllten Blocks per Hash in den Prefix-Cache aufzunehmen (siehe Coordinator und BlockPool). Das zurückgegebene KVCacheBlocks wird über get_block_ids in eine block_table umgewandelt und in der Blocktabelle für den Attention-Kernel abgelegt.

Grenzen und Fehler

  • num_new_tokens == 0 und keine externen Token: löst ValueError aus. Der Aufrufer muss sicherstellen, dass entweder neue Token zu berechnen oder bereits berechnete externe Token zu speichern sind(kv_cache_manager.py:346-350).
  • get_computed_blocks fehlt bei vollem Treffer ein Token: max_cache_hit_length = request.num_tokens - 1, da bei vollem Treffer der letzte Token noch einmal berechnet werden muss, um die Logits zu erhalten. Die Trefferlänge ist also auf prompt_length - 1 begrenzt(kv_cache_manager.py:225-231).
  • Anfragen, die den Prefix-Cache überspringen: Bei request.skip_reading_prefix_cache=True oder enable_caching=False gibt get_computed_blocks direkt Leer zurück(kv_cache_manager.py:222-223), z. B. für Anfragen, die Prompt-Logprobs benötigen.
  • reset_prefix_cache abgelehnt: Wenn num_used_blocks != 1 (nur Null-Block), sind noch Blocks nicht freigegeben, der Reset schlägt fehl und liefert False(kv_cache_manager.py:516-530).
  • pop_blocks_for_free gibt nichts an den Pool zurück: Der Aufrufer muss selbst in umgekehrter Reihenfolge block_pool.free_blocks aufrufen, sonst lecken Blocks(kv_cache_manager.py:495-506).
  • take_events annotiert Gruppen-Metadaten: Ein BlockStored-Ereignis enthält nur group_idx; der Manager ergänzt über kv_cache_event_metadata den Kind und das Sliding-Window(kv_cache_manager.py:566-589).

Zusammenfassung

KVCacheManager ist die KV-Cache-Fassade, die der Scheduler sieht; die eigentliche Arbeit verrichten KVCacheCoordinator + BlockPool (siehe Coordinator und BlockPool). Die zugeteilten Block-IDs landen schließlich in der Blocktabelle, wo der Attention-Kernel sie zur Adressierung des physischen KV-Cache nutzt. Die Modellgewichte selbst werden in der DefaultModelLoader-Linie geladen und sind mit dieser Schicht nicht gekoppelt.

Siehe offizielle Dokumentation: vLLM 文档 · README.