| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Iterable |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.triton_utils import tl, triton |
| 8 | from vllm.v1.attention.backends.utils import PAD_SLOT_ID |
| 9 | from vllm.v1.worker.gpu.buffer_utils import ( |
| 10 | FusedStagedWriter, |
| 11 | StagedWriteTensor, |
| 12 | UvaBackedTensor, |
| 13 | _load_ptr, |
| 14 | ) |
| 15 | |
| 16 | |
| 17 | class BlockTables: |
| 18 | # Set by the elastic EP warmup so KV writes land in the null block. |
| 19 | redirect_writes_to_null_block = False |
| 20 | |
| 21 | def __init__( |
| 22 | self, |
| 23 | block_sizes: list[int], |
| 24 | max_num_reqs: int, |
| 25 | max_num_batched_tokens: int, |
| 26 | max_num_blocks_per_group: list[int], |
| 27 | device: torch.device, |
| 28 | kernel_block_sizes: list[int], |
| 29 | cp_size: int = 1, |
| 30 | cp_rank: int = 0, |
| 31 | cp_interleave: int = 1, |
| 32 | slot_mapping_enabled: list[bool] | None = None, |
| 33 | dcp_sharded: list[bool] | None = None, |
| 34 | ): |
| 35 | self.block_sizes = block_sizes |
| 36 | self.kernel_block_sizes = kernel_block_sizes |
| 37 | self.max_num_reqs = max_num_reqs |
| 38 | self.max_num_batched_tokens = max_num_batched_tokens |
| 39 | self.device = device |
| 40 | |
| 41 | self.cp_size = cp_size |
| 42 | self.cp_rank = cp_rank |
| 43 | self.cp_interleave = cp_interleave |
| 44 | |
| 45 | self.num_kv_cache_groups = len(self.block_sizes) |
| 46 | assert len(max_num_blocks_per_group) == self.num_kv_cache_groups |
| 47 | if slot_mapping_enabled is None: |
| 48 | slot_mapping_enabled = [True] * self.num_kv_cache_groups |
| 49 | assert len(slot_mapping_enabled) == self.num_kv_cache_groups |
| 50 | self._slot_mapping_enabled = slot_mapping_enabled |
| 51 | if dcp_sharded is None: |
| 52 | dcp_sharded = [True] * self.num_kv_cache_groups |
| 53 | assert len(dcp_sharded) == self.num_kv_cache_groups |
| 54 | self.dcp_sharded = torch.tensor(dcp_sharded, dtype=torch.bool, device=device) |
| 55 | |
| 56 | self.blocks_per_kv_block = [ |
| 57 | bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) |
| 58 | ] |
| 59 | |
| 60 | # num_kv_cache_groups x [max_num_reqs, max_num_blocks] |
| 61 | self.block_tables: list[StagedWriteTensor] = [] |
| 62 | for i in range(self.num_kv_cache_groups): |
| 63 | max_num_blocks = max_num_blocks_per_group[i] * self.blocks_per_kv_block[i] |
| 64 | block_table = StagedWriteTensor( |
| 65 | (self.max_num_reqs, max_num_blocks), dtype=torch.int32, device=device |
| 66 | ) |
| 67 | self.block_tables.append(block_table) |
| 68 | |
| 69 | self.num_blocks = UvaBackedTensor( |
| 70 | (self.num_kv_cache_groups, self.max_num_reqs), dtype=torch.int32 |
| 71 | ) |
| 72 | self.fused_writer: FusedStagedWriter | None = None |
| 73 | if self.num_kv_cache_groups > 1: |
| 74 | # Only the multi-group path uses the fused writer. |
| 75 | self.fused_writer = FusedStagedWriter( |
| 76 | self.device, self.num_kv_cache_groups * self.max_num_reqs |
| 77 | ) |
| 78 | |
| 79 | # Block tables used for model's forward pass. |
| 80 | # num_kv_cache_groups x [max_num_reqs, max_num_blocks] |
| 81 | self.input_block_tables: list[torch.Tensor] = [ |
| 82 | torch.zeros_like(b.gpu) for b in self.block_tables |
| 83 | ] |
| 84 | |
| 85 | self.slot_mappings = torch.zeros( |
| 86 | self.num_kv_cache_groups, |
| 87 | self.max_num_batched_tokens, |
| 88 | dtype=torch.int64, |
| 89 | device=self.device, |
| 90 | ) |
| 91 | |
| 92 | self.init_block_table_layout_tensors() |
| 93 | |
| 94 | def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor: |
| 95 | # NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses. |
| 96 | return torch.tensor( |
| 97 | [t.data_ptr() for t in x], dtype=torch.uint64, device=self.device |
| 98 | ) |
| 99 | |
| 100 | def init_block_table_layout_tensors(self) -> None: |
| 101 | self.block_table_ptrs = self._make_ptr_tensor( |
| 102 | [b.gpu for b in self.block_tables] |
| 103 | ) |
| 104 | self.block_table_strides = torch.tensor( |
| 105 | [b.gpu.stride(0) for b in self.block_tables], |
| 106 | dtype=torch.int64, |
| 107 | device=self.device, |
| 108 | ) |
| 109 | self.block_sizes_tensor = torch.tensor( |
| 110 | self.block_sizes, dtype=torch.int32, device=self.device |
| 111 | ) |
| 112 | self.kernel_block_sizes_tensor = torch.tensor( |
| 113 | self.kernel_block_sizes, dtype=torch.int32, device=self.device |
| 114 | ) |
| 115 | self.slot_mapping_enabled = torch.tensor( |
| 116 | self._slot_mapping_enabled, dtype=torch.bool, device=self.device |
| 117 | ) |
| 118 | self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables) |
| 119 | |
| 120 | def append_block_ids( |
| 121 | self, |
| 122 | req_index: int, |
| 123 | new_block_ids: tuple[list[int], ...], |
| 124 | overwrite: bool, |
| 125 | ) -> None: |
| 126 | for i in range(self.num_kv_cache_groups): |
| 127 | start = self.num_blocks.np[i, req_index] if not overwrite else 0 |
| 128 | block_ids = new_block_ids[i] |
| 129 | bpk = self.blocks_per_kv_block[i] |
| 130 | if bpk > 1: |
| 131 | block_ids = [b * bpk + k for b in block_ids for k in range(bpk)] |
| 132 | if self.redirect_writes_to_null_block: |
| 133 | block_ids = [0] * len(block_ids) |
| 134 | end = start + len(block_ids) |
| 135 | row_capacity = self.block_tables[i].gpu.shape[1] |
| 136 | if end > row_capacity: |
| 137 | raise RuntimeError( |
| 138 | f"Block table write for request {req_index}, group {i} exceeds " |
| 139 | f"row capacity ({end} > {row_capacity})" |
| 140 | ) |
| 141 | self.block_tables[i].stage_write(req_index, start, block_ids) |
| 142 | self.num_blocks.np[i, req_index] = end |
| 143 | |
| 144 | def apply_staged_writes(self) -> None: |
| 145 | if self.num_kv_cache_groups == 0: |
| 146 | return |
| 147 | if self.num_kv_cache_groups == 1: |
| 148 | # Single group: write directly, skipping the per-write group lookup. |
| 149 | self.block_tables[0].apply_write() |
| 150 | elif self.num_kv_cache_groups > 1: |
| 151 | # Multiple groups: apply all block tables with one fused kernel. |
| 152 | assert self.fused_writer is not None |
| 153 | self.fused_writer.apply( |
| 154 | self.block_tables, self.block_table_ptrs, self.block_table_strides |
| 155 | ) |
| 156 | self.num_blocks.copy_to_uva() |
| 157 | |
| 158 | def gather_block_tables( |
| 159 | self, |
| 160 | idx_mapping: torch.Tensor, |
| 161 | num_reqs_padded: int, |
| 162 | out: tuple[torch.Tensor, ...] | None = None, |
| 163 | out_ptrs: torch.Tensor | None = None, |
| 164 | ) -> tuple[torch.Tensor, ...]: |
| 165 | if self.num_kv_cache_groups == 0: |
| 166 | return () |
| 167 | if out is None: |
| 168 | out = tuple(self.input_block_tables) |
| 169 | out_ptrs = self.input_block_table_ptrs |
| 170 | else: |
| 171 | assert out_ptrs is not None |
| 172 | assert len(out) == self.num_kv_cache_groups |
| 173 | num_reqs = idx_mapping.shape[0] |
| 174 | # Launch kernel with num_reqs_padded to fuse zeroing of padded rows. |
| 175 | _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs_padded)]( |
| 176 | idx_mapping, |
| 177 | self.block_table_ptrs, |
| 178 | out_ptrs, |
| 179 | self.block_table_strides, |
| 180 | self.num_blocks.gpu, |
| 181 | self.num_blocks.gpu.stride(0), |
| 182 | num_reqs, |
| 183 | BLOCK_SIZE=1024, # type: ignore |
| 184 | ) |
| 185 | return tuple(bt[:num_reqs_padded] for bt in out) |
| 186 | |
| 187 | def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]: |
| 188 | # NOTE(woosuk): The output may be used for CUDA graph capture. |
| 189 | # Therefore, this method must return the persistent tensor |
| 190 | # with the same memory address as that used during the model's forward pass, |
| 191 | # rather than allocating a new tensor. |
| 192 | # |
| 193 | # Zero the rows so dummy runs write mamba state to the reserved null |
| 194 | # block rather than through the previous real step's (stale) block |
| 195 | # ids, which may point at blocks since freed and reallocated. |
| 196 | return tuple( |
| 197 | block_table[:num_reqs].zero_() for block_table in self.input_block_tables |
| 198 | ) |
| 199 | |
| 200 | def compute_slot_mappings( |
| 201 | self, |
| 202 | idx_mapping: torch.Tensor, |
| 203 | query_start_loc: torch.Tensor, |
| 204 | positions: torch.Tensor, |
| 205 | num_tokens_padded: int, |
| 206 | out: torch.Tensor | None = None, |
| 207 | ) -> torch.Tensor: |
| 208 | if self.num_kv_cache_groups == 0: |
| 209 | return (self.slot_mappings if out is None else out)[:, :num_tokens_padded] |
| 210 | num_reqs = idx_mapping.shape[0] |
| 211 | num_groups = self.num_kv_cache_groups |
| 212 | slot_mappings = self.slot_mappings if out is None else out |
| 213 | _compute_slot_mappings_kernel[(num_groups, num_reqs + 1)]( |
| 214 | slot_mappings.shape[1], |
| 215 | idx_mapping, |
| 216 | query_start_loc, |
| 217 | positions, |
| 218 | self.block_table_ptrs, |
| 219 | self.block_table_strides, |
| 220 | self.block_sizes_tensor, |
| 221 | self.kernel_block_sizes_tensor, |
| 222 | self.slot_mapping_enabled, |
| 223 | self.dcp_sharded, |
| 224 | slot_mappings, |
| 225 | slot_mappings.stride(0), |
| 226 | self.cp_rank, |
| 227 | CP_SIZE=self.cp_size, |
| 228 | CP_INTERLEAVE=self.cp_interleave, |
| 229 | PAD_ID=PAD_SLOT_ID, |
| 230 | TRITON_BLOCK_SIZE=1024, # type: ignore |
| 231 | ) |
| 232 | return slot_mappings[:, :num_tokens_padded] |
| 233 | |
| 234 | def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor: |
| 235 | # Fill the entire slot_mappings tensor, not just the first `num_tokens` entries. |
| 236 | # This is because the padding logic is complex and kernels may access beyond |
| 237 | # the requested range. |
| 238 | self.slot_mappings.fill_(PAD_SLOT_ID) |
| 239 | # NOTE(woosuk): The output may be used for CUDA graph capture. |
| 240 | # Therefore, this method must return the persistent tensor |
| 241 | # with the same memory address as that used during the model's forward pass, |
| 242 | # rather than allocating a new tensor. |
| 243 | return self.slot_mappings[:, :num_tokens] |
| 244 | |
| 245 | |
| 246 | @triton.jit(do_not_specialize=["num_reqs"]) |
| 247 | def _gather_block_tables_kernel( |
| 248 | batch_idx_to_req_idx, # [batch_size] |
| 249 | src_block_table_ptrs, # [num_kv_cache_groups] |
| 250 | dst_block_table_ptrs, # [num_kv_cache_groups] |
| 251 | block_table_strides, # [num_kv_cache_groups] |
| 252 | num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs] |
| 253 | num_blocks_stride, |
| 254 | num_reqs, # actual number of requests (for padding) |
| 255 | BLOCK_SIZE: tl.constexpr, |
| 256 | ): |
| 257 | # kv cache group id |
| 258 | group_id = tl.program_id(0) |
| 259 | batch_idx = tl.program_id(1) |
| 260 | |
| 261 | stride = tl.load(block_table_strides + group_id) |
| 262 | max_num_blocks = stride # stride equals max_num_blocks for this group. |
| 263 | dst_block_table_ptr = _load_ptr(dst_block_table_ptrs + group_id, tl.int32) |
| 264 | dst_row_ptr = dst_block_table_ptr + batch_idx * stride |
| 265 | |
| 266 | if batch_idx >= num_reqs: |
| 267 | # Zero out padded rows. |
| 268 | for i in tl.range(0, max_num_blocks, BLOCK_SIZE): |
| 269 | offset = i + tl.arange(0, BLOCK_SIZE) |
| 270 | tl.store(dst_row_ptr + offset, 0, mask=offset < max_num_blocks) |
| 271 | return |
| 272 | |
| 273 | req_idx = tl.load(batch_idx_to_req_idx + batch_idx) |
| 274 | group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride |
| 275 | num_blocks = tl.load(group_num_blocks_ptr + req_idx) |
| 276 | |
| 277 | src_block_table_ptr = _load_ptr(src_block_table_ptrs + group_id, tl.int32) |
| 278 | src_row_ptr = src_block_table_ptr + req_idx * stride |
| 279 | |
| 280 | for i in tl.range(0, num_blocks, BLOCK_SIZE): |
| 281 | offset = i + tl.arange(0, BLOCK_SIZE) |
| 282 | block_ids = tl.load(src_row_ptr + offset, mask=offset < num_blocks) |
| 283 | tl.store(dst_row_ptr + offset, block_ids, mask=offset < num_blocks) |
| 284 | |
| 285 | |
| 286 | @triton.jit |
| 287 | def _compute_slot_mappings_kernel( |
| 288 | max_num_tokens, |
| 289 | idx_mapping, # [num_reqs] |
| 290 | query_start_loc, # [num_reqs + 1] |
| 291 | pos, # [num_tokens] |
| 292 | block_table_ptrs, # [num_kv_cache_groups] |
| 293 | block_table_strides, # [num_kv_cache_groups] |
| 294 | block_sizes, # [num_kv_cache_groups] |
| 295 | kernel_block_sizes, # [num_kv_cache_groups] |
| 296 | slot_mapping_enabled, # [num_kv_cache_groups] |
| 297 | dcp_sharded, # [num_kv_cache_groups] |
| 298 | slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] |
| 299 | slot_mappings_stride, |
| 300 | cp_rank, |
| 301 | CP_SIZE: tl.constexpr, |
| 302 | CP_INTERLEAVE: tl.constexpr, |
| 303 | PAD_ID: tl.constexpr, |
| 304 | TRITON_BLOCK_SIZE: tl.constexpr, |
| 305 | ): |
| 306 | # kv cache group id |
| 307 | group_id = tl.program_id(0) |
| 308 | batch_idx = tl.program_id(1) |
| 309 | slot_mapping_ptr = slot_mappings_ptr + group_id * slot_mappings_stride |
| 310 | |
| 311 | if batch_idx == tl.num_programs(1) - 1: |
| 312 | # Pad remaining slots to -1. This is needed for CUDA graphs. |
| 313 | # Start from actual token count (not padded) to cover the gap |
| 314 | # between actual tokens and padded tokens that can contain stale |
| 315 | # valid slot IDs from previous chunks during chunked prefill. |
| 316 | actual_num_tokens = tl.load(query_start_loc + batch_idx) |
| 317 | for i in range(actual_num_tokens, max_num_tokens, TRITON_BLOCK_SIZE): |
| 318 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 319 | tl.store(slot_mapping_ptr + offset, PAD_ID, mask=offset < max_num_tokens) |
| 320 | return |
| 321 | |
| 322 | block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32) |
| 323 | block_table_stride = tl.load(block_table_strides + group_id) |
| 324 | kv_block_size = tl.load(block_sizes + group_id) |
| 325 | kernel_block_size = tl.load(kernel_block_sizes + group_id) |
| 326 | mapping_enabled = tl.load(slot_mapping_enabled + group_id) |
| 327 | if CP_SIZE != 1: |
| 328 | sharded = tl.load(dcp_sharded + group_id) |
| 329 | |
| 330 | req_state_idx = tl.load(idx_mapping + batch_idx) |
| 331 | # idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns |
| 332 | # no blocks: never read its block-table row and emit PAD for its tokens. |
| 333 | is_real_req = req_state_idx >= 0 |
| 334 | start_idx = tl.load(query_start_loc + batch_idx) |
| 335 | end_idx = tl.load(query_start_loc + batch_idx + 1) |
| 336 | for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): |
| 337 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 338 | positions = tl.load(pos + offset, mask=offset < end_idx, other=0) |
| 339 | |
| 340 | if CP_SIZE == 1: |
| 341 | # Common case: Context parallelism is not used. |
| 342 | local_positions = positions |
| 343 | is_local = True |
| 344 | else: |
| 345 | # Context parallelism is used. |
| 346 | virtual_block_size = kv_block_size * CP_SIZE |
| 347 | virtual_block_indices = positions // virtual_block_size |
| 348 | virtual_block_offsets = positions % virtual_block_size |
| 349 | is_local = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank |
| 350 | rounds = virtual_block_offsets // (CP_INTERLEAVE * CP_SIZE) |
| 351 | remainder = virtual_block_offsets % CP_INTERLEAVE |
| 352 | local_offsets = rounds * CP_INTERLEAVE + remainder |
| 353 | local_positions = virtual_block_indices * kv_block_size + local_offsets |
| 354 | local_positions = tl.where(sharded, local_positions, positions) |
| 355 | is_local = ~sharded | is_local |
| 356 | |
| 357 | block_indices = tl.where( |
| 358 | mapping_enabled, local_positions // kernel_block_size, 0 |
| 359 | ) |
| 360 | block_offsets = local_positions % kernel_block_size |
| 361 | block_numbers = tl.load( |
| 362 | block_table_ptr + req_state_idx * block_table_stride + block_indices, |
| 363 | mask=is_local & is_real_req, |
| 364 | other=0, |
| 365 | ) |
| 366 | slot_ids = block_numbers * kernel_block_size + block_offsets |
| 367 | if CP_SIZE != 1: |
| 368 | slot_ids = tl.where(is_local, slot_ids, PAD_ID) |
| 369 | |
| 370 | slot_ids = tl.where(mapping_enabled & is_real_req, slot_ids, PAD_ID) |
| 371 | tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) |