vllm.v1.worker.block_table

Classes:

• BlockTable –
• MultiGroupBlockTable –

  The BlockTables for each KV cache group.

Functions:

• get_block_table_width –

  Return the width after optional alignment.

BlockTable

Methods:

• __init__ –

  Args:

• get_cpu_tensor –

  Returns the CPU tensor of the block table.

• get_device_tensor –

  Returns the device tensor of the block table.

• get_numpy_array –

  Returns the numpy array of the block table.

Source code in vllm/v1/worker/block_table.py

class BlockTable:
    def __init__(
        self,
        block_size: int,
        max_num_reqs: int,
        max_num_blocks_per_req: int,
        max_num_batched_tokens: int,
        pin_memory: bool,
        device: torch.device,
        cp_kv_cache_interleave_size: int,
        slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT,
    ):
        """Args:
        block_size: Block size used for KV cache memory allocation
        max_num_reqs: Maximum number of concurrent requests supported.
        max_num_blocks_per_req: Maximum number of blocks per request.
        max_num_batched_tokens: Maximum number of tokens in a batch.
        pin_memory: Whether to pin memory for faster GPU transfers.
        device: Target device for the block table.
        slot_mapping_mode: How this cache group maps scheduled tokens to
            cache slots. Mamba-like state caches do not use token slot
            mappings and should use SlotMappingMode.NONE.

        """
        self.max_num_reqs = max_num_reqs
        self.max_num_batched_tokens = max_num_batched_tokens
        self.pin_memory = pin_memory
        self.device = device
        self.block_size = block_size
        self.max_num_blocks_per_req = max_num_blocks_per_req

        self.block_table = self._make_buffer(
            self.max_num_reqs, self.max_num_blocks_per_req, dtype=torch.int32
        )
        self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)

        self.slot_mapping = self._make_buffer(
            self.max_num_batched_tokens, dtype=torch.int64
        )

        try:
            self.pcp_world_size = get_pcp_group().world_size
            self.pcp_rank = get_pcp_group().rank_in_group
        except AssertionError:
            # PCP might not be initialized in testing
            self.pcp_world_size = 1
            self.pcp_rank = 0
        try:
            self.dcp_world_size = get_dcp_group().world_size
            self.dcp_rank = get_dcp_group().rank_in_group
        except AssertionError:
            # DCP might not be initialized in testing
            self.dcp_world_size = 1
            self.dcp_rank = 0
        self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
        self.slot_mapping_mode = slot_mapping_mode
        if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT:
            _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup(
                kv_cache_block_size=self.block_size,
                total_cp_world_size=self.dcp_world_size,
                total_cp_rank=self.dcp_rank,
                cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
                block_table_stride=self.block_table.gpu.stride(0),
            )

    def append_row(
        self,
        block_ids: list[int],
        row_idx: int,
    ) -> None:
        if not block_ids:
            return

        num_blocks = len(block_ids)
        start = self.num_blocks_per_row[row_idx]
        self.num_blocks_per_row[row_idx] += num_blocks
        self.block_table.np[row_idx, start : start + num_blocks] = block_ids

    def add_row(self, block_ids: list[int], row_idx: int) -> None:
        self.num_blocks_per_row[row_idx] = 0
        self.append_row(block_ids, row_idx)

    def clear_row(self, row_idx: int) -> None:
        num_blocks = self.num_blocks_per_row[row_idx]
        if num_blocks > 0:
            self.block_table.np[row_idx, :num_blocks] = 0
        self.num_blocks_per_row[row_idx] = 0

    def move_row(self, src: int, tgt: int) -> None:
        num_blocks = self.num_blocks_per_row[src]
        block_table_np = self.block_table.np
        block_table_np[tgt, :num_blocks] = block_table_np[src, :num_blocks]
        self.num_blocks_per_row[tgt] = num_blocks
        # Clear the vacated source row: dummy-run batches dereference stale
        # rows as mamba state slots and write state in place there, possibly
        # after the blocks have been freed and reallocated.
        block_table_np[src, :num_blocks] = 0
        self.num_blocks_per_row[src] = 0

    def swap_row(self, src: int, tgt: int) -> None:
        src_tgt, tgt_src = [src, tgt], [tgt, src]
        self.num_blocks_per_row[src_tgt] = self.num_blocks_per_row[tgt_src]
        self.block_table.np[src_tgt] = self.block_table.np[tgt_src]

    def compute_slot_mapping(
        self,
        num_reqs: int,
        query_start_loc: torch.Tensor,
        positions: torch.Tensor,
    ) -> None:
        num_tokens = positions.shape[0]
        if self.slot_mapping_mode == SlotMappingMode.NONE:
            # Mamba/GDN groups consume the block table as recurrent state
            # indices and do not use per-token slot mappings.
            return
        assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT

        _COMPUTE_SLOT_MAPPING_KERNEL(
            num_reqs,
            num_tokens,
            self.max_num_batched_tokens,
            query_start_loc,
            positions,
            self.block_table.gpu,
            self.block_table.gpu.stride(0),
            self.slot_mapping.gpu,
            self.block_size,
            self.dcp_world_size,
            self.dcp_rank,
            self.cp_kv_cache_interleave_size,
        )

    def commit_block_table(self, num_reqs: int) -> None:
        self.block_table.copy_to_gpu(num_reqs)

    def clear(self) -> None:
        self.block_table.gpu.fill_(0)
        self.block_table.cpu.fill_(0)

    def get_device_tensor(self, num_reqs: int) -> torch.Tensor:
        """Returns the device tensor of the block table."""
        return self.block_table.gpu[:num_reqs]

    def get_cpu_tensor(self) -> torch.Tensor:
        """Returns the CPU tensor of the block table."""
        return self.block_table.cpu

    def get_numpy_array(self) -> np.ndarray:
        """Returns the numpy array of the block table."""
        return self.block_table.np

    def _make_buffer(
        self, *size: int | torch.SymInt, dtype: torch.dtype
    ) -> CpuGpuBuffer:
        return CpuGpuBuffer(
            *size, dtype=dtype, device=self.device, pin_memory=self.pin_memory
        )

__init__(block_size, max_num_reqs, max_num_blocks_per_req, max_num_batched_tokens, pin_memory, device, cp_kv_cache_interleave_size, slot_mapping_mode=SlotMappingMode.TOKEN_TO_KV_SLOT)

Args: block_size: Block size used for KV cache memory allocation max_num_reqs: Maximum number of concurrent requests supported. max_num_blocks_per_req: Maximum number of blocks per request. max_num_batched_tokens: Maximum number of tokens in a batch. pin_memory: Whether to pin memory for faster GPU transfers. device: Target device for the block table. slot_mapping_mode: How this cache group maps scheduled tokens to cache slots. Mamba-like state caches do not use token slot mappings and should use SlotMappingMode.NONE.

Source code in vllm/v1/worker/block_table.py

def __init__(
    self,
    block_size: int,
    max_num_reqs: int,
    max_num_blocks_per_req: int,
    max_num_batched_tokens: int,
    pin_memory: bool,
    device: torch.device,
    cp_kv_cache_interleave_size: int,
    slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT,
):
    """Args:
    block_size: Block size used for KV cache memory allocation
    max_num_reqs: Maximum number of concurrent requests supported.
    max_num_blocks_per_req: Maximum number of blocks per request.
    max_num_batched_tokens: Maximum number of tokens in a batch.
    pin_memory: Whether to pin memory for faster GPU transfers.
    device: Target device for the block table.
    slot_mapping_mode: How this cache group maps scheduled tokens to
        cache slots. Mamba-like state caches do not use token slot
        mappings and should use SlotMappingMode.NONE.

    """
    self.max_num_reqs = max_num_reqs
    self.max_num_batched_tokens = max_num_batched_tokens
    self.pin_memory = pin_memory
    self.device = device
    self.block_size = block_size
    self.max_num_blocks_per_req = max_num_blocks_per_req

    self.block_table = self._make_buffer(
        self.max_num_reqs, self.max_num_blocks_per_req, dtype=torch.int32
    )
    self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)

    self.slot_mapping = self._make_buffer(
        self.max_num_batched_tokens, dtype=torch.int64
    )

    try:
        self.pcp_world_size = get_pcp_group().world_size
        self.pcp_rank = get_pcp_group().rank_in_group
    except AssertionError:
        # PCP might not be initialized in testing
        self.pcp_world_size = 1
        self.pcp_rank = 0
    try:
        self.dcp_world_size = get_dcp_group().world_size
        self.dcp_rank = get_dcp_group().rank_in_group
    except AssertionError:
        # DCP might not be initialized in testing
        self.dcp_world_size = 1
        self.dcp_rank = 0
    self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
    self.slot_mapping_mode = slot_mapping_mode
    if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT:
        _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup(
            kv_cache_block_size=self.block_size,
            total_cp_world_size=self.dcp_world_size,
            total_cp_rank=self.dcp_rank,
            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
            block_table_stride=self.block_table.gpu.stride(0),
        )

get_cpu_tensor()

Returns the CPU tensor of the block table.

Source code in vllm/v1/worker/block_table.py

def get_cpu_tensor(self) -> torch.Tensor:
    """Returns the CPU tensor of the block table."""
    return self.block_table.cpu

get_device_tensor(num_reqs)

Returns the device tensor of the block table.

Source code in vllm/v1/worker/block_table.py

def get_device_tensor(self, num_reqs: int) -> torch.Tensor:
    """Returns the device tensor of the block table."""
    return self.block_table.gpu[:num_reqs]

get_numpy_array()

Returns the numpy array of the block table.

Source code in vllm/v1/worker/block_table.py

def get_numpy_array(self) -> np.ndarray:
    """Returns the numpy array of the block table."""
    return self.block_table.np

MultiGroupBlockTable

The BlockTables for each KV cache group.

Methods:

• __getitem__ –

  Returns the BlockTable for the i-th KV cache group.

Source code in vllm/v1/worker/block_table.py

class MultiGroupBlockTable:
    """The BlockTables for each KV cache group."""

    def __init__(
        self,
        max_num_reqs: int,
        max_num_batched_tokens: int,
        pin_memory: bool,
        device: torch.device,
        block_sizes: list[int],
        max_num_blocks: list[int],
        cp_kv_cache_interleave_size: int = 1,
        slot_mapping_modes: list[SlotMappingMode] | None = None,
    ) -> None:
        if slot_mapping_modes is None:
            slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] * len(block_sizes)
        if len(slot_mapping_modes) != len(block_sizes):
            raise ValueError(
                f"slot_mapping_modes length ({len(slot_mapping_modes)}) "
                f"must match block_sizes length ({len(block_sizes)})"
            )

        if len(max_num_blocks) != len(block_sizes):
            raise ValueError(
                f"max_num_blocks length ({len(max_num_blocks)}) "
                f"must match block_sizes length ({len(block_sizes)})"
            )

        max_num_blocks = [
            (
                get_block_table_width(n, block_size, token_alignment=None)
                if slot_mapping_mode == SlotMappingMode.NONE
                else get_block_table_width(n, block_size)
            )
            for n, block_size, slot_mapping_mode in zip(
                max_num_blocks, block_sizes, slot_mapping_modes
            )
        ]

        self.block_tables = [
            BlockTable(
                block_size,
                max_num_reqs,
                max_num_blocks_per_req,
                max_num_batched_tokens,
                pin_memory,
                device,
                cp_kv_cache_interleave_size,
                slot_mapping_mode=slot_mapping_mode,
            )
            for block_size, max_num_blocks_per_req, slot_mapping_mode in zip(
                block_sizes, max_num_blocks, slot_mapping_modes
            )
        ]

    def append_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None:
        for i, block_table in enumerate(self.block_tables):
            block_table.append_row(block_ids[i], row_idx)

    def add_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None:
        for i, block_table in enumerate(self.block_tables):
            block_table.add_row(block_ids[i], row_idx)

    def clear_row(self, row_idx: int) -> None:
        for block_table in self.block_tables:
            block_table.clear_row(row_idx)

    def move_row(self, src: int, tgt: int) -> None:
        for block_table in self.block_tables:
            block_table.move_row(src, tgt)

    def swap_row(self, src: int, tgt: int) -> None:
        for block_table in self.block_tables:
            block_table.swap_row(src, tgt)

    def compute_slot_mapping(
        self,
        num_reqs: int,
        query_start_loc: torch.Tensor,
        positions: torch.Tensor,
    ) -> None:
        for block_table in self.block_tables:
            block_table.compute_slot_mapping(num_reqs, query_start_loc, positions)

    def commit_block_table(self, num_reqs: int) -> None:
        for block_table in self.block_tables:
            block_table.commit_block_table(num_reqs)

    def clear(self) -> None:
        for block_table in self.block_tables:
            block_table.clear()

    def __getitem__(self, idx: int) -> "BlockTable":
        """Returns the BlockTable for the i-th KV cache group."""
        return self.block_tables[idx]

__getitem__(idx)

Returns the BlockTable for the i-th KV cache group.

Source code in vllm/v1/worker/block_table.py

def __getitem__(self, idx: int) -> "BlockTable":
    """Returns the BlockTable for the i-th KV cache group."""
    return self.block_tables[idx]

get_block_table_width(max_num_blocks, block_size, *, token_alignment=128)

Return the width after optional alignment.

Source code in vllm/v1/worker/block_table.py

def get_block_table_width(
    max_num_blocks: int,
    block_size: int,
    *,
    token_alignment: int | None = 128,
) -> int:
    """Return the width after optional alignment."""
    if token_alignment is not None:
        if token_alignment <= 0:
            raise ValueError("token_alignment must be positive")
        block_alignment = token_alignment // math.gcd(token_alignment, block_size)
        max_num_blocks = cdiv(max_num_blocks, block_alignment) * block_alignment
    return max_num_blocks