Block table : mapping des tokens logiques vers les blocs KV physiques

源码版本v0.25.1

Responsabilités

L'idée centrale de PagedAttention est de découper le cache KV en blocs de taille fixe ; pour chaque token, le bloc et le slot auxquels il appartient sont retrouvés via la « block table ». vLLM v1 l'implémente comme la classe BlockTable : un tensor int32 de forme (max_num_reqs, max_num_blocks_per_req), chaque ligne correspondant à une requête, le j-ième élément de la ligne étant le block_id physique du j-ième bloc logique de la requête. MultiGroupBlockTable regroupe les block tables de plusieurs groupes de cache KV, pour qu'un modèle hybride (par exemple full attention + sliding window) en détienne chacune sa part. BlockTable porte également un tensor slot_mapping — qui mappe directement chaque token de query à un slot physique dans le cache KV, permettant au kernel d'attention de lire/écrire le KV directement.

Le cycle de vie de la block table est lié à la requête : à l'appel de InputBatch.add_request, block_table.add_row(request.block_ids, req_index) écrit la liste de block_ids renvoyée par le manager dans cette ligne(gpu_input_batch.py:336-379). Après chaque step d'ordonnancement, les nouveaux block_ids alloués sont ajoutés en bout de ligne via block_table.append_row(new_block_ids, req_index)(gpu_model_runner.py:1425-1428). commit_block_table copie le tableau numpy vers le GPU(block_table.py:166-167) ; compute_slot_mapping utilise un kernel Triton pour calculer le slot id de chaque token(block_table.py:141-164). Au forward, l'attention backend lit directement ce tensor depuis common_attn_metadata.block_table_tensor(backend.py:420-421).

Motivation de conception

Pourquoi maintenir une block table côté worker, plutôt que d'utiliser directement req_to_blocks du manager ?

Double buffer numpy + GPU : block_table utilise CpuGpuBuffer pour tenir simultanément trois copies numpy / CPU / GPU, écrit sur CPU (numpy), lu sur GPU (tensor int32) ; ainsi commit_block_table(num_reqs) ne fait qu'une seule copie H2D(block_table.py:70-77). En mode capture CUDA graph, la taille est fixe, on réutilise directement le tensor GPU.
kernel block_size ≠ manager block_size : certains kernels d'attention (ex. TRTLLM) ont un block_size de 16, alors que le manager de cache KV utilise 32 ; map_to_kernel_blocks déplie un bloc manager en blocks_per_kv_block=2 blocs kernel(block_table.py:47-68)(block_table.py:173-201).
slot_mapping kernelisé : compute_slot_mapping appelle le kernel Triton _compute_slot_mapping_kernel, chaque program traite une requête, calcule block_index = pos // block_size depuis positions, lit block_table[req_idx, block_index] pour obtenir le block_number, puis slot_id = block_number * block_size + offset(block_table.py:325-380). Plus rapide qu'une implémentation numpy, et exécutable dans un CUDA graph.
Entrelacement DCP / PCP : trois constexpr TOTAL_CP_WORLD_SIZE / TOTAL_CP_RANK / CP_KV_CACHE_INTERLEAVE_SIZE font que, en decode context parallelism, le kernel n'écrit la vraie valeur que sur les slots appartenant à ce rank, et PAD_ID ailleurs(block_table.py:368-379).
Padding CUDA graph : quand num_reqs < num_reqs_padded, _get_block_table remplit les lignes de padding avec NULL_BLOCK_ID(gpu_model_runner.py:2279-2295) ; le kernel d'attention saute ces lignes en voyant NULL_BLOCK_ID.
Cache de builder partagé entre groupes : les builders avec supports_update_block_table=True ne reconstruisent pas lors du deuxième passage sur un groupe de même spec, ils se contentent de update_block_table(gpu_model_runner.py:2477-2493), évitant une reconstruction répétée des métadonnées.
Maintenance au niveau ligne : append_row / clear_row / move_row / swap_row traitent la block table comme un « tableau dynamique par ligne de requête », mis à jour en place lors des condense (réarrangement de requêtes), preempt (éviction de requêtes), etc., sans reconstruction globale(block_table.py:102-139).

Fichiers clés

BlockTable.__init__:18-100 — construction en double buffer, gère le mode hybrid block quand kernel_block_size != block_size.
BlockTable.append_row/add_row:102-122 — ajout en bout de ligne / réécriture complète de la liste de block_ids.
BlockTable.clear/move/swap_row:124-139 — maintenance au niveau ligne, supporte condense et preempt.
BlockTable.compute_slot_mapping:141-164 — appelle un kernel Triton pour convertir positions en tensor slot_mapping.
BlockTable.commit_block_table/clear:166-171 — copie H2D / mise à zéro.
BlockTable.map_to_kernel_blocks:173-201 — déplie la liste de block_ids manager en liste de block_ids kernel.
BlockTable.get_device_tensor:203-213 — renvoie la slice GPU à l'attention backend.
MultiGroupBlockTable:223-322 — emballage multi-groupes, broadcaste add_row / append_row / compute_slot_mapping vers chaque BlockTable de groupe.
_compute_slot_mapping_kernel:325-380 — kernel Triton : un program par requête, calcule slot_id = block_table[req, pos // block_size] * block_size + offset ; en DCP, filtre selon TOTAL_CP_RANK.
InputBatch.block_table init:172-180 — point de construction de MultiGroupBlockTable.
InputBatch.add_request:336-379 — à l'arrivée d'une nouvelle requête, block_table.add_row(request.block_ids, req_index).
_get_block_table:2279-2298 — prépare le block_table_tensor à chaque step, ajoute le padding NULL_BLOCK_ID.
append_row on new blocks:1425-1428 — les nouveaux block_ids sont ajoutés en bout de ligne existante.
commit_block_table:1933-1933 — après l'ordonnancement, copie le numpy en H2D vers le GPU.
CommonAttentionMetadata:2378-2396 — block_table_tensor attaché aux common metadata.
update_block_table:2477-2493 — groupes de même spec réutilisent les métadonnées déjà construites.
block_table_tensor field:420-421 — champ block_table dans CommonAttentionMetadata.
FlashAttentionMetadata.block_table:234-234 — tensor block_table détenu par le backend Flash.

Flux de données

À chaque boucle d'ordonnancement, gpu_model_runner._prepare_inputs traduit le SchedulerOutput en état d'input batch : pour chaque requête existante, on appelle block_table.append_row(new_block_ids, req_index) pour ajouter les nouveaux blocs en bout de ligne ; pour les slots libérés par le condense, on réarrange avec move_row / swap_row(gpu_input_batch.py:629-758). Ensuite, commit_block_table(num_reqs) copie le tableau numpy en H2D vers le GPU, et compute_slot_mapping lance le kernel Triton pour calculer slot_mapping :

python

@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) Le dernier program est dédié à l'écriture de PAD_ID sur les slots de padding, pour éviter que des données résiduelles ne polluent l'attention lors de la réutilisation du CUDA graph. Ensuite, _get_block_table(kv_cache_gid) récupère le tensor GPU du groupe(gpu_model_runner.py:2279-2295) et l'injecte dans CommonAttentionMetadata.block_table_tensor(gpu_model_runner.py:2378-2396). Les backends Flash attention / Triton / rocM lisent chacun attn_metadata.block_table(flash_attn.py:234-234)(triton_attn.py:208-234) et, avec slot_mapping, accomplissent la lecture / écriture du cache KV. Les block_ids proviennent du KVCacheManager, alloués via le Coordinator et le BlockPool.

Limites et échecs

kernel_block_size ne divise pas block_size : si block_size % kernel_block_size != 0, on lève directement ValueError(block_table.py:58-62).
append d'une liste vide : append_row([]) return immédiatement, ne met pas à jour num_blocks_per_row(block_table.py:107-108).
clear_row ne met pas à zéro num_blocks_per_row avant écriture : clear_row met d'abord block_table.np[row, :num_blocks] à zéro, puis remet le compteur à zéro(block_table.py:124-128).
Padding CUDA graph : les lignes num_reqs:num_reqs_padded sont remplies avec NULL_BLOCK_ID(gpu_model_runner.py:2292-2294) ; le kernel saute en voyant NULL_BLOCK_ID.
EncoderOnlyAttentionSpec renvoie un tensor zéro : le cache KV encoder-only n'a pas de block_table ; _get_block_table fournit un placeholder (num_reqs_padded, 1) rempli de zéros(gpu_model_runner.py:2282-2287).
PCP / DCP non initialisés : get_pcp_group() / get_dcp_group() peuvent ne pas être initialisés en test ; un try/except dégrade en world_size=1(block_table.py:86-99).
Mamba réutilise block_table_tensor : les backends Mamba / GDN / Linear attention utilisent mamba_get_block_table_tensor pour employer block_table comme state_indices, ce qui nécessite un gather en cache mode(utils.py:925-965).
Alignement max_num_blocks pour TRTLLM : MultiGroupBlockTable.__init__ aligne max_num_blocks sur un multiple de 128 // block_size, exigence de certains backends d'attention (TRTLLM)(block_table.py:260-265).

Résumé

BlockTable organise la liste de block_id allouée par le KVCacheManager / Coordinator et BlockPool en tensor GPU, calcule le slot physique de chaque token via un kernel Triton slot_mapping, et l'injecte dans CommonAttentionMetadata pour tous les backends d'attention. Elle n'a aucun couplage avec le chargement des poids du modèle (voir DefaultModelLoader), uniquement avec la couche de gestion du cache KV.

Voir la documentation officielle : Documentation vLLM · README

Block table : mapping des tokens logiques vers les blocs KV physiques ​

Responsabilités ​

Motivation de conception ​

Fichiers clés ​

Flux de données ​

Limites et échecs ​

Résumé ​