Blocktabelle: Abbildung von logischen Token auf physische KV-Blocks

源码版本v0.25.1

Verantwortung

Das Kernkonzept von PagedAttention ist, den KV-Cache in Blöcke fester Größe zu schneiden; für jedes Token wird über eine Blocktabelle (block table) nachgeschlagen, in welchem Block und an welchem Slot es liegt. vLLM v1 implementiert die Blocktabelle als Klasse BlockTable: einen int32-Tensor der Form (max_num_reqs, max_num_blocks_per_req). Jede Zeile entspricht einer Anfrage, das j-te Element der Zeile ist die physische block_id des j-ten logischen Blocks dieser Anfrage. MultiGroupBlockTable fasst die Blocktabellen mehrerer KV-Cache-Gruppen zusammen, sodass hybride Modelle (z. B. Full Attention + Sliding Window) jeweils eine eigene Tabelle führen. BlockTable hält zusätzlich den slot_mapping-Tensor, der jedes Query-Token direkt auf die physische Slot-ID im KV-Cache abbildet; der Attention-Kernel nutzt ihn, um KV direkt zu lesen und zu schreiben.

Der Lebenszyklus der Blocktabelle ist an die Anfrage gebunden: Bei InputBatch.add_request schreibt block_table.add_row(request.block_ids, req_index) die vom Manager gelieferte block_id-Liste in die entsprechende Zeile(gpu_input_batch.py:336-379). Nach jedem Scheduling-Schritt werden neu zugewiesene block_ids per block_table.append_row(new_block_ids, req_index) an das Zeilenende angehängt(gpu_model_runner.py:1425-1428). commit_block_table kopiert das numpy-Array auf die GPU(block_table.py:166-167), und compute_slot_mapping berechnet mit einem Triton-Kernel die Slot-ID jedes Tokens(block_table.py:141-164). Im Forward liest das Attention-Backend direkt aus common_attn_metadata.block_table_tensor(backend.py:420-421) diesen Tensor.

Entwurfsmotivation

Warum wird die Blocktabelle im Worker noch einmal gepflegt, anstatt req_to_blocks des Managers direkt zu verwenden?

• numpy + GPU doppelt gepuffert: block_table nutzt CpuGpuBuffer und hält numpy / CPU / GPU in dreifacher Ausfertigung; geschrieben wird auf der CPU (numpy), gelesen auf der GPU (int32-Tensor). So reduziert sich commit_block_table(num_reqs) auf einen einzigen H2D-Copy(block_table.py:70-77). Im CUDA-Graph-Modus mit fester Größe wird der GPU-Tensor direkt wiederverwendet.
• kernel_block_size ≠ manager block_size: Manche Attention-Kernel (wie TRTLLM) verwenden block_size=16, während der KV-Cache-Manager block_size=32 verwendet. map_to_kernel_blocks entfaltet einen Manager-Block in blocks_per_kv_block=2 Kernel-Blocks(block_table.py:47-68)(block_table.py:173-201).
• slot_mapping als Kernel: compute_slot_mapping ruft den Triton-Kernel _compute_slot_mapping_kernel auf. Jedes Programm verarbeitet eine Anfrage, berechnet aus positions den block_index = pos // block_size, schlägt in block_table[req_idx, block_index] die block_number nach und berechnet schließlich slot_id = block_number * block_size + offset(block_table.py:325-380). Das ist schneller als eine numpy-Implementierung und lässt sich innerhalb eines CUDA-Graphen ausführen.
• DCP / PCP-Interleaving: Die drei Constexprs TOTAL_CP_WORLD_SIZE / TOTAL_CP_RANK / CP_KV_CACHE_INTERLEAVE_SIZE sorgen dafür, dass der Kernel im Decode Context Parallelism nur für Slots, die zum eigenen Rank gehören, echte Werte schreibt; alle anderen erhalten PAD_ID(block_table.py:368-379).
• CUDA-Graph-Padding: Bei num_reqs < num_reqs_padded füllt _get_block_table die Padding-Zeilen mit NULL_BLOCK_ID(gpu_model_runner.py:2279-2295); der Attention-Kernel überspringt Zeilen mit NULL_BLOCK_ID.
• Mehrere Gruppen teilen sich den Builder-Cache: Ein Builder mit supports_update_block_table=True baut bei der zweiten Gruppe mit identischer Spec nicht neu auf, sondern ruft nur update_block_table auf(gpu_model_runner.py:2477-2493), was den wiederholten Metadata-Aufbau einspart.
• Zeilenweises Maintenance: append_row / clear_row / move_row / swap_row behandeln die Blocktabelle als "dynamisches Array pro Anfragezeile" und erlauben in Szenarien wie Condense (Anfrage-Umsortierung) oder Preempt (Anfrage-Verdrängung) direkte Aktualisierungen ohne vollständigen Neuaufbau(block_table.py:102-139).

Schlüsseldateien

• BlockTable.__init__:18-100 — doppelt gepufferte Konstruktion, behandelt den Hybrid-Block-Modus mit kernel_block_size != block_size.
• BlockTable.append_row/add_row:102-122 — block_id-Liste am Zeilenende anfügen / gesamte Zeile neu schreiben.
• BlockTable.clear/move/swap_row:124-139 — zeilenweises Maintenance, stützt Condense- und Preempt-Umsortierung.
• BlockTable.compute_slot_mapping:141-164 — ruft den Triton-Kernel auf, der positions in den slot_mapping-Tensor übersetzt.
• BlockTable.commit_block_table/clear:166-171 — H2D-Copy / Nullen.
• BlockTable.map_to_kernel_blocks:173-201 — entfaltet die Liste der manager-block_ids in die Liste der kernel-block_ids.
• BlockTable.get_device_tensor:203-213 — liefert dem Attention-Backend einen Slice des GPU-Tensors.
• MultiGroupBlockTable:223-322 — Wrapper für mehrere Gruppen, der add_row / append_row / compute_slot_mapping usw. an die BlockTable jeder Gruppe weiterreicht.
• _compute_slot_mapping_kernel:325-380 — Triton-Kernel: Programm pro Anfrage, berechnet slot_id = block_table[req, pos // block_size] * block_size + offset; bei DCP wird nach TOTAL_CP_RANK gefiltert.
• InputBatch.block_table init:172-180 — Konstruktionspunkt von MultiGroupBlockTable.
• InputBatch.add_request:336-379 — bei neuer Anfrage block_table.add_row(request.block_ids, req_index).
• _get_block_table:2279-2298 — bereitet jeden Schritt den block_table_tensor vor und füllt NULL_BLOCK_ID-Padding auf.
• append_row on new blocks:1425-1428 — neu zugewiesene block_ids werden am Ende der bestehenden Zeile angehängt.
• commit_block_table:1933-1933 — nach dem Scheduling wird numpy per H2D auf die GPU kopiert.
• CommonAttentionMetadata:2378-2396 — block_table_tensor wird an die Common-Metadata angehängt.
• update_block_table:2477-2493 — Gruppe mit identischer Spec gibt bereits aufgebaute Metadata wieder.
• block_table_tensor field:420-421 — das block_table-Feld in CommonAttentionMetadata.
• FlashAttentionMetadata.block_table:234-234 — der vom Flash-Backend gehaltene block_table-Tensor.

Datenfluss

In jedem Scheduling-Schleifendurchlauf übersetzt gpu_model_runner._prepare_inputs den SchedulerOutput in den Input-Batch-Zustand: Für jede bestehende Anfrage wird block_table.append_row(new_block_ids, req_index) aufgerufen, um neue Blocks an das Zeilenende anzufügen; für die nach Condense leeren Slots wird mit move_row / swap_row aufgeräumt(gpu_input_batch.py:629-758). Anschließend kopiert commit_block_table(num_reqs) das numpy-Array per H2D auf die GPU, und compute_slot_mapping führt den Triton-Kernel aus, um das slot_mapping zu berechnen:

@triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"])
def _compute_slot_mapping_kernel(
    num_tokens,
    max_num_tokens,
    query_start_loc_ptr,  # [num_reqs + 1], int32
    positions_ptr,  # [num_tokens], int64
    block_table_ptr,  # [max_num_reqs, max_num_blocks_per_req], int32 (flat)
    block_table_stride,  # max_num_blocks_per_req
    block_size,
    slot_mapping_ptr,  # [max_num_tokens], int64
    TOTAL_CP_WORLD_SIZE: tl.constexpr,
    TOTAL_CP_RANK: tl.constexpr,
    CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
    PAD_ID: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    req_idx = tl.program_id(0)

    if req_idx == tl.num_programs(0) - 1:
        # Pad remaining slots for CUDA graph compatibility.
        for i in range(num_tokens, max_num_tokens, BLOCK_SIZE):
            offsets = i + tl.arange(0, BLOCK_SIZE)
            tl.store(
                slot_mapping_ptr + offsets,
                PAD_ID,
                mask=offsets < max_num_tokens,
            )
        return

(block_table.py:325-352) Das letzte Programm ist ausschließlich dafür zuständig, die Padding-Slots mit PAD_ID zu befüllen, damit beim Wiederverwenden des CUDA-Graphen keine Reste die Attention verunreinigen. Anschließend holt _get_block_table(kv_cache_gid) den GPU-Tensor der jeweiligen Gruppe(gpu_model_runner.py:2279-2295) und hängt ihn in CommonAttentionMetadata.block_table_tensor ein(gpu_model_runner.py:2378-2396). Die Backends Flash-Attention / Triton / rocM lesen jeweils attn_metadata.block_table(flash_attn.py:234-234)(triton_attn.py:208-234) und schließen zusammen mit slot_mapping das Lesen / Schreiben des KV-Cache ab. Die block_id stammt vom KVCacheManager und wird über Coordinator und BlockPool zugewiesen.

Grenzen und Fehler

• kernel_block_size teilt block_size nicht: block_size % kernel_block_size != 0 wirft direkt raise ValueError(block_table.py:58-62).
• append mit leerer Liste: append_row([]) kehrt sofort zurück, ohne num_blocks_per_row zu aktualisieren(block_table.py:107-108).
• clear_row setzt num_blocks_per_row erst nach dem Nullen: clear_row nullt zuerst block_table.np[row, :num_blocks] und setzt erst dann den Zähler auf Null(block_table.py:124-128).
• CUDA-Graph-Padding: Die Zeilen num_reqs:num_reqs_padded werden mit NULL_BLOCK_ID gefüllt(gpu_model_runner.py:2292-2294); der Kernel überspringt bei NULL_BLOCK_ID.
• EncoderOnlyAttentionSpec liefert Null-Tensor: Encoder-only-KV-Cache hat keine block_table; _get_block_table liefert einen (num_reqs_padded, 1)-Null-Tensor als Platzhalter(gpu_model_runner.py:2282-2287).
• PCP / DCP nicht initialisiert: get_pcp_group() / get_dcp_group() sind in Tests u. U. nicht initialisiert; über try/except wird mit world_size=1 gefallbackt(block_table.py:86-99).
• Mamba verwendet block_table_tensor weiter: Die Backends Mamba / GDN / Linear-Attention verwenden mamba_get_block_table_tensor, um block_table als state_indices zu nutzen; im Cache-Modus ist ein Gather erforderlich(utils.py:925-965).
• max_num_blocks an TRTLLM ausgerichtet: MultiGroupBlockTable.__init__ richtet max_num_blocks an einem Vielfachen von 128 // block_size aus, wie es manche Attention-Backends (TRTLLM) verlangen(block_table.py:260-265).

Zusammenfassung

BlockTable organisiert die vom KVCacheManager / Coordinator und BlockPool zugewiesene block_id-Liste als GPU-Tensor und berechnet über den slot_mapping-Triton-Kernel den physischen Slot jedes Tokens. Diese werden in CommonAttentionMetadata abgelegt und stehen allen Attention-Backends zur Verfügung. Es ist nicht mit dem Laden der Modellgewichte gekoppelt (siehe DefaultModelLoader), sondern nur mit der KV-Cache-Verwaltungsschicht.

Siehe offizielle Dokumentation: vLLM 文档 · README.