Tabla de bloques: mapeo de tokens lógicos a bloques KV físicos

源码版本v0.25.1

Responsabilidades

La idea central de PagedAttention es cortar la caché KV (KV cache) en bloques de tamaño fijo; en qué bloque y en qué slot cae cada token se resuelve consultando la "tabla de bloques (block table)". vLLM v1 implementa la tabla de bloques como la clase BlockTable: un tensor int32 de forma (max_num_reqs, max_num_blocks_per_req), donde cada fila corresponde a una petición y el j-ésimo elemento de la fila es el block_id físico correspondiente al j-ésimo bloque lógico de esa petición. MultiGroupBlockTable empaqueta las tablas de varios grupos de caché KV para que los modelos híbridos (por ejemplo full attention + sliding window) mantengan cada uno la suya. BlockTable también guarda el tensor slot_mapping, que mapea cada token de query directamente al id de slot físico dentro de la caché KV; el kernel de attention lo usa para leer y escribir KV de forma directa.

El ciclo de vida de la tabla de bloques va ligado a la petición: en InputBatch.add_request, block_table.add_row(request.block_ids, req_index) escribe en esa fila la lista de block_id devuelta por el manager (gpu_input_batch.py:336-379). Tras cada paso de planificación, los block_id recién asignados se añaden al final de la fila con block_table.append_row(new_block_ids, req_index) (gpu_model_runner.py:1425-1428). commit_block_table copia el arreglo numpy a la GPU (block_table.py:166-167), y compute_slot_mapping usa un kernel de Triton para calcular el id de slot de cada token (block_table.py:141-164). Durante el forward, el attention backend lee directamente este tensor desde common_attn_metadata.block_table_tensor (backend.py:420-421).

Motivación de diseño

¿Por qué mantener otra copia de la tabla de bloques en el worker en vez de usar directamente req_to_blocks del manager?

Doble buffer numpy + GPU: block_table usa CpuGpuBuffer para mantener tres copias a la vez (numpy / CPU / GPU); se escribe en CPU (numpy) y se lee en GPU (tensor int32). Así commit_block_table(num_reqs) solo necesita una copia H2D (block_table.py:70-77). En modo de captura de CUDA graph el tamaño es fijo y se reutiliza el tensor GPU tal cual.
kernel block_size != manager block_size: algunos kernels de attention (como TRTLLM) usan block_size 16, mientras que el KV cache manager usa 32; map_to_kernel_blocks expande un bloque del manager en blocks_per_kv_block=2 bloques del kernel (block_table.py:47-68) (block_table.py:173-201).
slot_mapping como kernel: compute_slot_mapping invoca el kernel de Triton _compute_slot_mapping_kernel, donde cada program procesa una petición, saca block_index = pos // block_size a partir de positions, luego lee block_table[req_idx, block_index] para obtener block_number, y finalmente slot_id = block_number * block_size + offset (block_table.py:325-380). Es más rápido que una implementación numpy y también corre dentro de un CUDA graph.
Interleaving DCP / PCP: las tres constantes TOTAL_CP_WORLD_SIZE, TOTAL_CP_RANK y CP_KV_CACHE_INTERLEAVE_SIZE hacen que, bajo decode context parallelism, el kernel solo escriba valores reales en los slots que pertenecen a este rank y rellene el resto con PAD_ID (block_table.py:368-379).
Padding para CUDA graph: cuando num_reqs < num_reqs_padded, _get_block_table rellena las filas de padding con NULL_BLOCK_ID (gpu_model_runner.py:2279-2295); el kernel de attention, al ver NULL_BLOCK_ID, se salta esas filas.
Cache de builder compartida entre grupos: los builders con supports_update_block_table=True no se reconstruyen en el segundo grupo con la misma spec, solo llaman update_block_table (gpu_model_runner.py:2477-2493), ahorrándose la construcción repetida de metadata.
Mantenimiento por filas: append_row / clear_row / move_row / swap_row tratan la tabla como un "arreglo dinámico por filas de petición", actualizándose in-place en escenarios como condense (reordenar peticiones) o preempt (swap-out de peticiones), sin reconstrucción total (block_table.py:102-139).

Archivos clave

BlockTable.__init__:18-100 — construcción con doble buffer; cubre el modo de bloque híbrido cuando kernel_block_size != block_size.
BlockTable.append_row/add_row:102-122 — append al final de la fila / reescritura completa de la lista de block_id.
BlockTable.clear/move/swap_row:124-139 — mantenimiento a nivel de fila; soporta reordenamientos condense y preempt.
BlockTable.compute_slot_mapping:141-164 — invoca el kernel de Triton que convierte positions en el tensor slot_mapping.
BlockTable.commit_block_table/clear:166-171 — copia H2D / puesta a cero.
BlockTable.map_to_kernel_blocks:173-201 — expande la lista de block_id del manager a la lista de block_id del kernel.
BlockTable.get_device_tensor:203-213 — devuelve el slice del tensor GPU para el attention backend.
MultiGroupBlockTable:223-322 — envoltorio multi-group que difunde add_row / append_row / compute_slot_mapping etc. al BlockTable de cada grupo.
_compute_slot_mapping_kernel:325-380 — kernel de Triton: un program por petición, calcula slot_id = block_table[req, pos // block_size] * block_size + offset; bajo DCP filtra por TOTAL_CP_RANK.
InputBatch.block_table init:172-180 — punto de construcción de MultiGroupBlockTable.
InputBatch.add_request:336-379 — al entrar una petición nueva se llama block_table.add_row(request.block_ids, req_index).
_get_block_table:2279-2298 — prepara cada paso block_table_tensor, añadiendo padding con NULL_BLOCK_ID.
append_row on new blocks:1425-1428 — los block_id recién asignados se appendan al final de la fila existente.
commit_block_table:1933-1933 — tras planificar, vuelca el numpy a la GPU vía H2D.
CommonAttentionMetadata:2378-2396 — block_table_tensor se cuelga de la common metadata.
update_block_table:2477-2493 — grupos con la misma spec reutilizan la metadata ya construida.
block_table_tensor field:420-421 — campo block_table dentro de CommonAttentionMetadata.
FlashAttentionMetadata.block_table:234-234 — tensor block_table retenido por el backend Flash.

Flujo de datos

En cada paso del bucle de planificación, gpu_model_runner._prepare_inputs traduce el SchedulerOutput al estado del input batch: para cada petición existente llama block_table.append_row(new_block_ids, req_index) para añadir los bloques nuevos al final de la fila, y para los slots vacíos tras un condense usa move_row / swap_row para reordenar (gpu_input_batch.py:629-758). Luego commit_block_table(num_reqs) copia el arreglo numpy a la GPU vía H2D, y compute_slot_mapping corre el kernel de Triton para calcular 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) El último program se encarga específicamente de escribir PAD_ID en los slots de padding, evitando que datos residuales contaminen el attention al reutilizar un CUDA graph. Después _get_block_table(kv_cache_gid) recupera el tensor GPU de este grupo (gpu_model_runner.py:2279-2295) y lo inyecta en CommonAttentionMetadata.block_table_tensor (gpu_model_runner.py:2378-2396). Los backends Flash attention / Triton / rocM leen cada uno attn_metadata.block_table (flash_attn.py:234-234) (triton_attn.py:208-234) y, junto con slot_mapping, completan la lectura / escritura de la caché KV. Los block_id provienen del KVCacheManager tras ser asignados por el Coordinator y BlockPool.

Límites y fallos

kernel_block_size no divide a block_size: si block_size % kernel_block_size != 0, se lanza directamente raise ValueError (block_table.py:58-62).
append de lista vacía: append_row([]) hace return inmediato sin actualizar num_blocks_per_row (block_table.py:107-108).
clear_row no escribe 0 antes de poner a cero num_blocks_per_row: clear_row primero pone a cero block_table.np[row, :num_blocks] y luego reinicia el contador (block_table.py:124-128).
Padding para CUDA graph: las filas num_reqs:num_reqs_padded se rellenan con NULL_BLOCK_ID (gpu_model_runner.py:2292-2294); el kernel, al ver NULL_BLOCK_ID, se las salta.
EncoderOnlyAttentionSpec usa tensor de ceros: la caché KV encoder-only no tiene block_table, así que _get_block_table devuelve un tensor de ceros de forma (num_reqs_padded, 1) como marcador de posición (gpu_model_runner.py:2282-2287).
PCP / DCP sin inicializar: get_pcp_group() / get_dcp_group() pueden no estar inicializadas en tests; se cubre con try/except que las degrada a world_size=1 (block_table.py:86-99).
Mamba reutiliza block_table_tensor: los backends Mamba / GDN / Linear attention usan mamba_get_block_table_tensor para emplear block_table como state_indices; necesitan hacer gather en modo caché (utils.py:925-965).
max_num_blocks alineado a TRTLLM: MultiGroupBlockTable.__init__ alinea max_num_blocks a múltiplos de 128 // block_size, requisito de algunos backends de attention (TRTLLM) (block_table.py:260-265).

Resumen

BlockTable organiza la lista de block_id asignada por KVCacheManager / Coordinator y BlockPool en un tensor GPU, calcula el slot físico de cada token vía el kernel de Triton slot_mapping, y lo inyecta en CommonAttentionMetadata para que lo usen todos los backends de attention. No tiene acoplamiento con la carga de pesos del modelo (ver DefaultModelLoader); solo se acopla con la capa de gestión de la caché KV.

Véase la documentación oficial: Documentación de vLLM · README.

Tabla de bloques: mapeo de tokens lógicos a bloques KV físicos ​

Responsabilidades ​

Motivación de diseño ​

Archivos clave ​

Flujo de datos ​

Límites y fallos ​

Resumen ​