class BlockTables:
# Set by the elastic EP warmup so KV writes land in the null block.
redirect_writes_to_null_block = False
def __init__(
self,
block_sizes: list[int],
max_num_reqs: int,
max_num_batched_tokens: int,
max_num_blocks_per_group: list[int],
device: torch.device,
cp_size: int = 1,
cp_rank: int = 0,
cp_interleave: int = 1,
slot_mapping_enabled: list[bool] | None = None,
dcp_sharded: list[bool] | None = None,
sliding_window_reach: list[int] | None = None,
):
self.block_sizes = block_sizes
self.max_num_reqs = max_num_reqs
self.max_num_batched_tokens = max_num_batched_tokens
self.device = device
self.cp_size = cp_size
self.cp_rank = cp_rank
self.cp_interleave = cp_interleave
self.num_kv_cache_groups = len(self.block_sizes)
assert len(max_num_blocks_per_group) == self.num_kv_cache_groups
if slot_mapping_enabled is None:
slot_mapping_enabled = [True] * self.num_kv_cache_groups
assert len(slot_mapping_enabled) == self.num_kv_cache_groups
self._slot_mapping_enabled = slot_mapping_enabled
if dcp_sharded is None:
dcp_sharded = [True] * self.num_kv_cache_groups
assert len(dcp_sharded) == self.num_kv_cache_groups
self.dcp_sharded = torch.tensor(dcp_sharded, dtype=torch.bool, device=device)
# Per group, how many tokens below a request's first scheduled position
# attention can still read (SlidingWindowSpec.num_retained_tokens);
# 0 disables it for the group. Under context parallelism a
# block-table column does not map to a contiguous token range, so the
# masking in gather_block_tables is disabled.
if sliding_window_reach is None or cp_size > 1:
sliding_window_reach = [0] * self.num_kv_cache_groups
assert len(sliding_window_reach) == self.num_kv_cache_groups
self._sliding_window_reach = sliding_window_reach
self.has_sliding_window_groups = any(r > 0 for r in sliding_window_reach)
# num_kv_cache_groups x [max_num_reqs, max_num_blocks]
self.block_tables: list[StagedWriteTensor] = []
for max_num_blocks in max_num_blocks_per_group:
block_table = StagedWriteTensor(
(self.max_num_reqs, max_num_blocks), dtype=torch.int32, device=device
)
self.block_tables.append(block_table)
self.num_blocks = UvaBackedTensor(
(self.num_kv_cache_groups, self.max_num_reqs), dtype=torch.int32
)
self.fused_writer: FusedStagedWriter | None = None
if self.num_kv_cache_groups > 1:
# Only the multi-group path uses the fused writer.
self.fused_writer = FusedStagedWriter(
self.device, self.num_kv_cache_groups * self.max_num_reqs
)
# Block tables used for model's forward pass.
# num_kv_cache_groups x [max_num_reqs, max_num_blocks]
self.input_block_tables: list[torch.Tensor] = [
torch.zeros_like(b.gpu) for b in self.block_tables
]
self._num_dirty_rows = 0
self.slot_mappings = torch.zeros(
self.num_kv_cache_groups,
self.max_num_batched_tokens,
dtype=torch.int64,
device=self.device,
)
self.init_block_table_layout_tensors()
def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
# NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.
return torch.tensor(
[t.data_ptr() for t in x], dtype=torch.uint64, device=self.device
)
def init_block_table_layout_tensors(self) -> None:
self.block_table_ptrs = self._make_ptr_tensor(
[b.gpu for b in self.block_tables]
)
self.block_table_strides = torch.tensor(
[b.gpu.stride(0) for b in self.block_tables],
dtype=torch.int64,
device=self.device,
)
self.block_sizes_tensor = torch.tensor(
self.block_sizes, dtype=torch.int32, device=self.device
)
self.slot_mapping_enabled = torch.tensor(
self._slot_mapping_enabled, dtype=torch.bool, device=self.device
)
self.sliding_window_reach = torch.tensor(
self._sliding_window_reach, dtype=torch.int32, device=self.device
)
self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
def append_block_ids(
self,
req_index: int,
new_block_ids: tuple[list[int], ...],
overwrite: bool,
) -> None:
for i in range(self.num_kv_cache_groups):
start = self.num_blocks.np[i, req_index] if not overwrite else 0
block_ids = new_block_ids[i]
if self.redirect_writes_to_null_block:
block_ids = [0] * len(block_ids)
end = start + len(block_ids)
row_capacity = self.block_tables[i].gpu.shape[1]
if end > row_capacity:
raise RuntimeError(
f"Block table write for request {req_index}, group {i} exceeds "
f"row capacity ({end} > {row_capacity})"
)
self.block_tables[i].stage_write(req_index, start, block_ids)
self.num_blocks.np[i, req_index] = end
def apply_staged_writes(self) -> None:
if self.num_kv_cache_groups == 0:
return
if self.num_kv_cache_groups == 1:
# Single group: write directly, skipping the per-write group lookup.
self.block_tables[0].apply_write()
elif self.num_kv_cache_groups > 1:
# Multiple groups: apply all block tables with one fused kernel.
assert self.fused_writer is not None
self.fused_writer.apply(
self.block_tables, self.block_table_ptrs, self.block_table_strides
)
self.num_blocks.copy_to_uva()
def gather_block_tables(
self,
idx_mapping: torch.Tensor,
num_reqs_padded: int,
out: tuple[torch.Tensor, ...] | None = None,
out_ptrs: torch.Tensor | None = None,
num_computed_tokens: torch.Tensor | None = None,
) -> tuple[torch.Tensor, ...]:
"""Gather the batch's block-table rows into the forward-pass tables.
When num_computed_tokens ([max_num_reqs], indexed like the block tables)
is given, sliding-window groups get every entry wholly below the window
of the request's first scheduled token replaced by the null block. Those
blocks may already have been freed by the sliding-window manager and
reused, while attention kernels can still load them as part of a KV
tile that straddles the window start: masked-out NaN in such a tile
propagates into the output.
"""
if self.num_kv_cache_groups == 0:
return ()
mask_sliding_window = (
self.has_sliding_window_groups and num_computed_tokens is not None
)
if not mask_sliding_window:
# Unused by the kernel; any int32 tensor satisfies the signature.
num_computed_tokens = self.sliding_window_reach
num_reqs = idx_mapping.shape[0]
num_rows = num_reqs_padded
if out is None:
out = tuple(self.input_block_tables)
out_ptrs = self.input_block_table_ptrs
num_rows = max(num_rows, self._num_dirty_rows)
self._num_dirty_rows = num_reqs
else:
assert out_ptrs is not None
assert len(out) == self.num_kv_cache_groups
# Also clear rows left by a larger batch in the shared buffer.
_gather_block_tables_kernel[(self.num_kv_cache_groups, num_rows)](
idx_mapping,
self.block_table_ptrs,
out_ptrs,
self.block_table_strides,
self.num_blocks.gpu,
self.num_blocks.gpu.stride(0),
num_reqs,
num_computed_tokens,
self.sliding_window_reach,
self.block_sizes_tensor,
MASK_SLIDING_WINDOW=mask_sliding_window,
BLOCK_SIZE=1024, # type: ignore
)
return tuple(bt[:num_reqs_padded] for bt in out)
def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
# NOTE(woosuk): The output may be used for CUDA graph capture.
# Therefore, this method must return the persistent tensor
# with the same memory address as that used during the model's forward pass,
# rather than allocating a new tensor.
#
# Zero the rows so dummy runs write mamba state to the reserved null
# block rather than through the previous real step's (stale) block
# ids, which may point at blocks since freed and reallocated.
num_rows = max(num_reqs, self._num_dirty_rows)
for block_table in self.input_block_tables:
block_table[:num_rows].zero_()
# Profiling may populate the returned rows with valid block IDs.
self._num_dirty_rows = num_reqs
return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
def compute_slot_mappings(
self,
idx_mapping: torch.Tensor,
query_start_loc: torch.Tensor,
positions: torch.Tensor,
num_tokens_padded: int,
out: torch.Tensor | None = None,
) -> torch.Tensor:
if self.num_kv_cache_groups == 0:
return (self.slot_mappings if out is None else out)[:, :num_tokens_padded]
num_reqs = idx_mapping.shape[0]
num_groups = self.num_kv_cache_groups
slot_mappings = self.slot_mappings if out is None else out
_compute_slot_mappings_kernel[(num_groups, num_reqs + 1)](
slot_mappings.shape[1],
idx_mapping,
query_start_loc,
positions,
self.block_table_ptrs,
self.block_table_strides,
self.block_sizes_tensor,
self.slot_mapping_enabled,
self.dcp_sharded,
slot_mappings,
slot_mappings.stride(0),
self.cp_rank,
CP_SIZE=self.cp_size,
CP_INTERLEAVE=self.cp_interleave,
PAD_ID=PAD_SLOT_ID,
TRITON_BLOCK_SIZE=1024, # type: ignore
)
return slot_mappings[:, :num_tokens_padded]
def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
# Fill the entire slot_mappings tensor, not just the first `num_tokens` entries.
# This is because the padding logic is complex and kernels may access beyond
# the requested range.
self.slot_mappings.fill_(PAD_SLOT_ID)
# NOTE(woosuk): The output may be used for CUDA graph capture.
# Therefore, this method must return the persistent tensor
# with the same memory address as that used during the model's forward pass,
# rather than allocating a new tensor.
return self.slot_mappings[:, :num_tokens]