| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from dataclasses import dataclass |
| 6 | from enum import Enum |
| 7 | from typing import Any |
| 8 | |
| 9 | import numpy as np |
| 10 | import torch |
| 11 | |
| 12 | from vllm.distributed import get_dcp_group, get_pcp_group |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.warmup.jit_warmup import kernel_launcher |
| 15 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 16 | LaunchSpec, |
| 17 | TritonWarmupTensor, |
| 18 | VllmTritonJitKernel, |
| 19 | triton_scalar_specialization_rep, |
| 20 | ) |
| 21 | from vllm.triton_utils import tl, triton |
| 22 | from vllm.utils.math_utils import cdiv |
| 23 | from vllm.v1.attention.backends.utils import PAD_SLOT_ID |
| 24 | from vllm.v1.utils import CpuGpuBuffer |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | |
| 29 | def get_block_table_width( |
| 30 | max_num_blocks: int, |
| 31 | block_size: int, |
| 32 | kernel_block_size: int | None = None, |
| 33 | *, |
| 34 | token_alignment: int | None = 128, |
| 35 | ) -> int: |
| 36 | """Return the width after optional alignment and virtual block splitting.""" |
| 37 | if kernel_block_size is None: |
| 38 | kernel_block_size = block_size |
| 39 | if block_size % kernel_block_size != 0: |
| 40 | raise ValueError( |
| 41 | f"kernel_block_size {kernel_block_size} must divide " |
| 42 | f"block_size {block_size} evenly" |
| 43 | ) |
| 44 | if token_alignment is not None: |
| 45 | if token_alignment <= 0: |
| 46 | raise ValueError("token_alignment must be positive") |
| 47 | block_alignment = token_alignment // math.gcd(token_alignment, block_size) |
| 48 | max_num_blocks = cdiv(max_num_blocks, block_alignment) * block_alignment |
| 49 | return max_num_blocks * block_size // kernel_block_size |
| 50 | |
| 51 | |
| 52 | class SlotMappingMode(Enum): |
| 53 | TOKEN_TO_KV_SLOT = "token_to_kv_slot" |
| 54 | NONE = "none" |
| 55 | |
| 56 | |
| 57 | class BlockTable: |
| 58 | def __init__( |
| 59 | self, |
| 60 | block_size: int, |
| 61 | max_num_reqs: int, |
| 62 | max_num_blocks_per_req: int, |
| 63 | max_num_batched_tokens: int, |
| 64 | pin_memory: bool, |
| 65 | device: torch.device, |
| 66 | kernel_block_size: int, |
| 67 | cp_kv_cache_interleave_size: int, |
| 68 | slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT, |
| 69 | ): |
| 70 | """Args: |
| 71 | block_size: Block size used for KV cache memory allocation |
| 72 | max_num_reqs: Maximum number of concurrent requests supported. |
| 73 | max_num_blocks_per_req: Maximum number of blocks per request. |
| 74 | max_num_batched_tokens: Maximum number of tokens in a batch. |
| 75 | pin_memory: Whether to pin memory for faster GPU transfers. |
| 76 | device: Target device for the block table. |
| 77 | kernel_block_size: The block_size of underlying attention kernel. |
| 78 | Will be the same as `block_size` if `block_size` is supported |
| 79 | by the attention kernel. |
| 80 | slot_mapping_mode: How this cache group maps scheduled tokens to |
| 81 | cache slots. Mamba-like state caches do not use token slot |
| 82 | mappings and should use SlotMappingMode.NONE. |
| 83 | |
| 84 | """ |
| 85 | self.max_num_reqs = max_num_reqs |
| 86 | self.max_num_batched_tokens = max_num_batched_tokens |
| 87 | self.pin_memory = pin_memory |
| 88 | self.device = device |
| 89 | self.kv_cache_block_size = block_size |
| 90 | |
| 91 | if kernel_block_size == block_size: |
| 92 | # Standard case: allocation and computation use same block size |
| 93 | # No block splitting needed, direct mapping |
| 94 | self.block_size = block_size |
| 95 | self.blocks_per_kv_block = 1 |
| 96 | self.use_hybrid_blocks = False |
| 97 | else: |
| 98 | # Hybrid case: allocation block size differs from kernel block size |
| 99 | # Memory blocks are subdivided to match kernel requirements |
| 100 | # Example: 32-token memory blocks with 16-token kernel blocks |
| 101 | # → Each memory block corresponds to 2 kernel blocks |
| 102 | if block_size % kernel_block_size != 0: |
| 103 | raise ValueError( |
| 104 | f"kernel_block_size {kernel_block_size} must divide " |
| 105 | f"kv_manager_block_size size {block_size} evenly" |
| 106 | ) |
| 107 | |
| 108 | self.block_size = kernel_block_size |
| 109 | self.blocks_per_kv_block = block_size // kernel_block_size |
| 110 | self.use_hybrid_blocks = True |
| 111 | |
| 112 | self.max_num_blocks_per_req = max_num_blocks_per_req * self.blocks_per_kv_block |
| 113 | |
| 114 | self.block_table = self._make_buffer( |
| 115 | self.max_num_reqs, self.max_num_blocks_per_req, dtype=torch.int32 |
| 116 | ) |
| 117 | self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32) |
| 118 | |
| 119 | self.slot_mapping = self._make_buffer( |
| 120 | self.max_num_batched_tokens, dtype=torch.int64 |
| 121 | ) |
| 122 | |
| 123 | if self.use_hybrid_blocks: |
| 124 | self._kernel_block_arange = np.arange(0, self.blocks_per_kv_block).reshape( |
| 125 | 1, -1 |
| 126 | ) |
| 127 | else: |
| 128 | self._kernel_block_arange = None |
| 129 | |
| 130 | try: |
| 131 | self.pcp_world_size = get_pcp_group().world_size |
| 132 | self.pcp_rank = get_pcp_group().rank_in_group |
| 133 | except AssertionError: |
| 134 | # PCP might not be initialized in testing |
| 135 | self.pcp_world_size = 1 |
| 136 | self.pcp_rank = 0 |
| 137 | try: |
| 138 | self.dcp_world_size = get_dcp_group().world_size |
| 139 | self.dcp_rank = get_dcp_group().rank_in_group |
| 140 | except AssertionError: |
| 141 | # DCP might not be initialized in testing |
| 142 | self.dcp_world_size = 1 |
| 143 | self.dcp_rank = 0 |
| 144 | self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size |
| 145 | self.slot_mapping_mode = slot_mapping_mode |
| 146 | if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT: |
| 147 | _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup( |
| 148 | kv_cache_block_size=self.kv_cache_block_size, |
| 149 | blocks_per_kv_block=self.blocks_per_kv_block, |
| 150 | total_cp_world_size=self.dcp_world_size, |
| 151 | total_cp_rank=self.dcp_rank, |
| 152 | cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, |
| 153 | block_table_stride=self.block_table.gpu.stride(0), |
| 154 | block_size=self.block_size, |
| 155 | ) |
| 156 | |
| 157 | def append_row( |
| 158 | self, |
| 159 | block_ids: list[int], |
| 160 | row_idx: int, |
| 161 | ) -> None: |
| 162 | if not block_ids: |
| 163 | return |
| 164 | |
| 165 | if self.use_hybrid_blocks: |
| 166 | block_ids = self.map_to_kernel_blocks( |
| 167 | np.array(block_ids), self.blocks_per_kv_block, self._kernel_block_arange |
| 168 | ) |
| 169 | |
| 170 | num_blocks = len(block_ids) |
| 171 | start = self.num_blocks_per_row[row_idx] |
| 172 | self.num_blocks_per_row[row_idx] += num_blocks |
| 173 | self.block_table.np[row_idx, start : start + num_blocks] = block_ids |
| 174 | |
| 175 | def add_row(self, block_ids: list[int], row_idx: int) -> None: |
| 176 | self.num_blocks_per_row[row_idx] = 0 |
| 177 | self.append_row(block_ids, row_idx) |
| 178 | |
| 179 | def clear_row(self, row_idx: int) -> None: |
| 180 | num_blocks = self.num_blocks_per_row[row_idx] |
| 181 | if num_blocks > 0: |
| 182 | self.block_table.np[row_idx, :num_blocks] = 0 |
| 183 | self.num_blocks_per_row[row_idx] = 0 |
| 184 | |
| 185 | def move_row(self, src: int, tgt: int) -> None: |
| 186 | num_blocks = self.num_blocks_per_row[src] |
| 187 | block_table_np = self.block_table.np |
| 188 | block_table_np[tgt, :num_blocks] = block_table_np[src, :num_blocks] |
| 189 | self.num_blocks_per_row[tgt] = num_blocks |
| 190 | # Clear the vacated source row: dummy-run batches dereference stale |
| 191 | # rows as mamba state slots and write state in place there, possibly |
| 192 | # after the blocks have been freed and reallocated. |
| 193 | block_table_np[src, :num_blocks] = 0 |
| 194 | self.num_blocks_per_row[src] = 0 |
| 195 | |
| 196 | def swap_row(self, src: int, tgt: int) -> None: |
| 197 | src_tgt, tgt_src = [src, tgt], [tgt, src] |
| 198 | self.num_blocks_per_row[src_tgt] = self.num_blocks_per_row[tgt_src] |
| 199 | self.block_table.np[src_tgt] = self.block_table.np[tgt_src] |
| 200 | |
| 201 | def compute_slot_mapping( |
| 202 | self, |
| 203 | num_reqs: int, |
| 204 | query_start_loc: torch.Tensor, |
| 205 | positions: torch.Tensor, |
| 206 | ) -> None: |
| 207 | num_tokens = positions.shape[0] |
| 208 | if self.slot_mapping_mode == SlotMappingMode.NONE: |
| 209 | # Mamba/GDN groups consume the block table as recurrent state |
| 210 | # indices and do not use per-token slot mappings. |
| 211 | return |
| 212 | assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT |
| 213 | |
| 214 | _COMPUTE_SLOT_MAPPING_KERNEL( |
| 215 | num_reqs, |
| 216 | num_tokens, |
| 217 | self.max_num_batched_tokens, |
| 218 | query_start_loc, |
| 219 | positions, |
| 220 | self.block_table.gpu, |
| 221 | self.block_table.gpu.stride(0), |
| 222 | self.block_size, |
| 223 | self.slot_mapping.gpu, |
| 224 | self.kv_cache_block_size, |
| 225 | self.blocks_per_kv_block, |
| 226 | self.dcp_world_size, |
| 227 | self.dcp_rank, |
| 228 | self.cp_kv_cache_interleave_size, |
| 229 | ) |
| 230 | |
| 231 | def commit_block_table(self, num_reqs: int) -> None: |
| 232 | self.block_table.copy_to_gpu(num_reqs) |
| 233 | |
| 234 | def clear(self) -> None: |
| 235 | self.block_table.gpu.fill_(0) |
| 236 | self.block_table.cpu.fill_(0) |
| 237 | |
| 238 | @staticmethod |
| 239 | def map_to_kernel_blocks( |
| 240 | kv_manager_block_ids: np.ndarray, |
| 241 | blocks_per_kv_block: int, |
| 242 | kernel_block_arange: np.ndarray, |
| 243 | ) -> np.ndarray: |
| 244 | """Convert kv_manager_block_id IDs to kernel block IDs. |
| 245 | |
| 246 | Example: |
| 247 | # kv_manager_block_ids: 32 tokens, |
| 248 | # Kernel block size: 16 tokens |
| 249 | # blocks_per_kv_block = 2 |
| 250 | >>> kv_manager_block_ids = np.array([0, 1, 2]) |
| 251 | >>> Result: [0, 1, 2, 3, 4, 5] |
| 252 | |
| 253 | # Each kv_manager_block_id maps to 2 kernel block id: |
| 254 | # kv_manager_block_id 0 → kernel block id [0, 1] |
| 255 | # kv_manager_block_id 1 → kernel block id [2, 3] |
| 256 | # kv_manager_block_id 2 → kernel block id [4, 5] |
| 257 | |
| 258 | """ |
| 259 | if blocks_per_kv_block == 1: |
| 260 | return kv_manager_block_ids |
| 261 | |
| 262 | kernel_block_ids = ( |
| 263 | kv_manager_block_ids.reshape(-1, 1) * blocks_per_kv_block |
| 264 | + kernel_block_arange |
| 265 | ) |
| 266 | |
| 267 | return kernel_block_ids.reshape(-1) |
| 268 | |
| 269 | def get_device_tensor(self, num_reqs: int) -> torch.Tensor: |
| 270 | """Returns the device tensor of the block table.""" |
| 271 | return self.block_table.gpu[:num_reqs] |
| 272 | |
| 273 | def get_cpu_tensor(self) -> torch.Tensor: |
| 274 | """Returns the CPU tensor of the block table.""" |
| 275 | return self.block_table.cpu |
| 276 | |
| 277 | def get_numpy_array(self) -> np.ndarray: |
| 278 | """Returns the numpy array of the block table.""" |
| 279 | return self.block_table.np |
| 280 | |
| 281 | def _make_buffer( |
| 282 | self, *size: int | torch.SymInt, dtype: torch.dtype |
| 283 | ) -> CpuGpuBuffer: |
| 284 | return CpuGpuBuffer( |
| 285 | *size, dtype=dtype, device=self.device, pin_memory=self.pin_memory |
| 286 | ) |
| 287 | |
| 288 | |
| 289 | class MultiGroupBlockTable: |
| 290 | """The BlockTables for each KV cache group.""" |
| 291 | |
| 292 | def __init__( |
| 293 | self, |
| 294 | max_num_reqs: int, |
| 295 | max_num_batched_tokens: int, |
| 296 | pin_memory: bool, |
| 297 | device: torch.device, |
| 298 | block_sizes: list[int], |
| 299 | kernel_block_sizes: list[int], |
| 300 | max_num_blocks: list[int], |
| 301 | cp_kv_cache_interleave_size: int = 1, |
| 302 | slot_mapping_modes: list[SlotMappingMode] | None = None, |
| 303 | ) -> None: |
| 304 | if len(kernel_block_sizes) != len(block_sizes): |
| 305 | raise ValueError( |
| 306 | f"kernel_block_sizes length ({len(kernel_block_sizes)}) " |
| 307 | f"must match block_sizes length ({len(block_sizes)})" |
| 308 | ) |
| 309 | if slot_mapping_modes is None: |
| 310 | slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] * len(block_sizes) |
| 311 | if len(slot_mapping_modes) != len(block_sizes): |
| 312 | raise ValueError( |
| 313 | f"slot_mapping_modes length ({len(slot_mapping_modes)}) " |
| 314 | f"must match block_sizes length ({len(block_sizes)})" |
| 315 | ) |
| 316 | |
| 317 | if len(max_num_blocks) != len(block_sizes): |
| 318 | raise ValueError( |
| 319 | f"max_num_blocks length ({len(max_num_blocks)}) " |
| 320 | f"must match block_sizes length ({len(block_sizes)})" |
| 321 | ) |
| 322 | |
| 323 | max_num_blocks = [ |
| 324 | ( |
| 325 | get_block_table_width(n, block_size, token_alignment=None) |
| 326 | if slot_mapping_mode == SlotMappingMode.NONE |
| 327 | else get_block_table_width(n, block_size) |
| 328 | ) |
| 329 | for n, block_size, slot_mapping_mode in zip( |
| 330 | max_num_blocks, block_sizes, slot_mapping_modes |
| 331 | ) |
| 332 | ] |
| 333 | |
| 334 | self.block_tables = [ |
| 335 | BlockTable( |
| 336 | block_size, |
| 337 | max_num_reqs, |
| 338 | max_num_blocks_per_req, |
| 339 | max_num_batched_tokens, |
| 340 | pin_memory, |
| 341 | device, |
| 342 | kernel_block_size, |
| 343 | cp_kv_cache_interleave_size, |
| 344 | slot_mapping_mode=slot_mapping_mode, |
| 345 | ) |
| 346 | for ( |
| 347 | block_size, |
| 348 | kernel_block_size, |
| 349 | max_num_blocks_per_req, |
| 350 | slot_mapping_mode, |
| 351 | ) in zip( |
| 352 | block_sizes, kernel_block_sizes, max_num_blocks, slot_mapping_modes |
| 353 | ) |
| 354 | ] |
| 355 | |
| 356 | def append_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None: |
| 357 | for i, block_table in enumerate(self.block_tables): |
| 358 | block_table.append_row(block_ids[i], row_idx) |
| 359 | |
| 360 | def add_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None: |
| 361 | for i, block_table in enumerate(self.block_tables): |
| 362 | block_table.add_row(block_ids[i], row_idx) |
| 363 | |
| 364 | def clear_row(self, row_idx: int) -> None: |
| 365 | for block_table in self.block_tables: |
| 366 | block_table.clear_row(row_idx) |
| 367 | |
| 368 | def move_row(self, src: int, tgt: int) -> None: |
| 369 | for block_table in self.block_tables: |
| 370 | block_table.move_row(src, tgt) |
| 371 | |
| 372 | def swap_row(self, src: int, tgt: int) -> None: |
| 373 | for block_table in self.block_tables: |
| 374 | block_table.swap_row(src, tgt) |
| 375 | |
| 376 | def compute_slot_mapping( |
| 377 | self, |
| 378 | num_reqs: int, |
| 379 | query_start_loc: torch.Tensor, |
| 380 | positions: torch.Tensor, |
| 381 | ) -> None: |
| 382 | for block_table in self.block_tables: |
| 383 | block_table.compute_slot_mapping(num_reqs, query_start_loc, positions) |
| 384 | |
| 385 | def commit_block_table(self, num_reqs: int) -> None: |
| 386 | for block_table in self.block_tables: |
| 387 | block_table.commit_block_table(num_reqs) |
| 388 | |
| 389 | def clear(self) -> None: |
| 390 | for block_table in self.block_tables: |
| 391 | block_table.clear() |
| 392 | |
| 393 | def __getitem__(self, idx: int) -> "BlockTable": |
| 394 | """Returns the BlockTable for the i-th KV cache group.""" |
| 395 | return self.block_tables[idx] |
| 396 | |
| 397 | |
| 398 | class ComputeSlotMappingKernel( |
| 399 | VllmTritonJitKernel["ComputeSlotMappingKernel.CompileKey"] |
| 400 | ): |
| 401 | triton_block_size = 1024 |
| 402 | |
| 403 | @dataclass(frozen=True) |
| 404 | class CompileKey: |
| 405 | kv_cache_block_size: int |
| 406 | blocks_per_kv_block: int |
| 407 | total_cp_world_size: int |
| 408 | total_cp_rank: int |
| 409 | cp_kv_cache_interleave_size: int |
| 410 | block_table_stride: int |
| 411 | block_size: int |
| 412 | |
| 413 | @staticmethod |
| 414 | @triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) |
| 415 | def kernel( |
| 416 | num_tokens, |
| 417 | max_num_tokens, |
| 418 | query_start_loc_ptr, # [num_reqs + 1], int32 |
| 419 | positions_ptr, # [num_tokens], int64 |
| 420 | block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat) |
| 421 | block_table_stride, # max_num_blocks_per_req |
| 422 | block_size, |
| 423 | slot_mapping_ptr, # [max_num_tokens], int64 |
| 424 | KV_CACHE_BLOCK_SIZE: tl.constexpr, |
| 425 | BLOCKS_PER_KV_BLOCK: tl.constexpr, |
| 426 | TOTAL_CP_WORLD_SIZE: tl.constexpr, |
| 427 | TOTAL_CP_RANK: tl.constexpr, |
| 428 | CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr, |
| 429 | PAD_ID: tl.constexpr, |
| 430 | BLOCK_SIZE: tl.constexpr, |
| 431 | ): |
| 432 | req_idx = tl.program_id(0) |
| 433 | |
| 434 | if req_idx == tl.num_programs(0) - 1: |
| 435 | # Pad remaining slots for CUDA graph compatibility. |
| 436 | for i in range(num_tokens, max_num_tokens, BLOCK_SIZE): |
| 437 | offsets = i + tl.arange(0, BLOCK_SIZE) |
| 438 | tl.store( |
| 439 | slot_mapping_ptr + offsets, |
| 440 | PAD_ID, |
| 441 | mask=offsets < max_num_tokens, |
| 442 | ) |
| 443 | return |
| 444 | |
| 445 | start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64) |
| 446 | end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64) |
| 447 | |
| 448 | virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE |
| 449 | row_offset = req_idx * block_table_stride |
| 450 | for i in range(start_idx, end_idx, BLOCK_SIZE): |
| 451 | offsets = i + tl.arange(0, BLOCK_SIZE) |
| 452 | mask = offsets < end_idx |
| 453 | pos = tl.load(positions_ptr + offsets, mask=mask, other=0) |
| 454 | virtual_block_indices = pos // virtual_block_size |
| 455 | virtual_block_offsets = pos - virtual_block_indices * virtual_block_size |
| 456 | is_local = ( |
| 457 | virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE |
| 458 | ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK |
| 459 | local_block_offsets = ( |
| 460 | virtual_block_offsets |
| 461 | // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE) |
| 462 | ) * CP_KV_CACHE_INTERLEAVE_SIZE + ( |
| 463 | virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE |
| 464 | ) |
| 465 | |
| 466 | block_indices = ( |
| 467 | virtual_block_indices * BLOCKS_PER_KV_BLOCK |
| 468 | + local_block_offsets // block_size |
| 469 | ) |
| 470 | block_numbers = tl.load( |
| 471 | block_table_ptr + row_offset + block_indices, |
| 472 | mask=mask & is_local, |
| 473 | other=0, |
| 474 | ).to(tl.int64) |
| 475 | slot_offsets = local_block_offsets % block_size |
| 476 | slot_ids = block_numbers * block_size + slot_offsets |
| 477 | slot_ids = tl.where(is_local, slot_ids, PAD_ID) |
| 478 | tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask) |
| 479 | |
| 480 | def dispatch( # type: ignore[override] |
| 481 | self, |
| 482 | *, |
| 483 | block_table_stride: int, |
| 484 | block_size: int, |
| 485 | **compile_key_fields: int, |
| 486 | ) -> CompileKey: |
| 487 | return self.CompileKey( |
| 488 | **compile_key_fields, |
| 489 | block_table_stride=triton_scalar_specialization_rep(block_table_stride), |
| 490 | block_size=triton_scalar_specialization_rep(block_size), |
| 491 | ) |
| 492 | |
| 493 | def get_warmup_keys(self, **dispatch_kwargs: int) -> list[CompileKey]: |
| 494 | return self._trace_dispatch(self.dispatch)(**dispatch_kwargs) |
| 495 | |
| 496 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 497 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 498 | int64_ptr = TritonWarmupTensor(torch.int64) |
| 499 | return dict( |
| 500 | num_reqs=1, |
| 501 | num_tokens=2, # arbitrary, in do_not_specialize |
| 502 | max_num_tokens=2, # arbitrary, in do_not_specialize |
| 503 | query_start_loc=int32_ptr, |
| 504 | positions=int64_ptr, |
| 505 | block_table=TritonWarmupTensor( |
| 506 | torch.int32, |
| 507 | shape=(1, compile_key.block_table_stride), |
| 508 | ), |
| 509 | block_table_stride=compile_key.block_table_stride, |
| 510 | block_size=compile_key.block_size, |
| 511 | slot_mapping=int64_ptr, |
| 512 | kv_cache_block_size=compile_key.kv_cache_block_size, |
| 513 | blocks_per_kv_block=compile_key.blocks_per_kv_block, |
| 514 | total_cp_world_size=compile_key.total_cp_world_size, |
| 515 | total_cp_rank=compile_key.total_cp_rank, |
| 516 | cp_kv_cache_interleave_size=compile_key.cp_kv_cache_interleave_size, |
| 517 | ) |
| 518 | |
| 519 | @kernel_launcher |
| 520 | def __call__( |
| 521 | self, |
| 522 | num_reqs: int, |
| 523 | num_tokens: int, |
| 524 | max_num_tokens: int, |
| 525 | query_start_loc: torch.Tensor, |
| 526 | positions: torch.Tensor, |
| 527 | block_table: torch.Tensor, |
| 528 | block_table_stride: int, |
| 529 | block_size: int, |
| 530 | slot_mapping: torch.Tensor, |
| 531 | kv_cache_block_size: int, |
| 532 | blocks_per_kv_block: int, |
| 533 | total_cp_world_size: int, |
| 534 | total_cp_rank: int, |
| 535 | cp_kv_cache_interleave_size: int, |
| 536 | ) -> LaunchSpec: |
| 537 | return (num_reqs + 1,), dict( |
| 538 | KV_CACHE_BLOCK_SIZE=kv_cache_block_size, |
| 539 | BLOCKS_PER_KV_BLOCK=blocks_per_kv_block, |
| 540 | TOTAL_CP_WORLD_SIZE=total_cp_world_size, |
| 541 | TOTAL_CP_RANK=total_cp_rank, |
| 542 | CP_KV_CACHE_INTERLEAVE_SIZE=cp_kv_cache_interleave_size, |
| 543 | PAD_ID=PAD_SLOT_ID, |
| 544 | BLOCK_SIZE=self.triton_block_size, |
| 545 | ) |
| 546 | |
| 547 | |
| 548 | _COMPUTE_SLOT_MAPPING_KERNEL = ComputeSlotMappingKernel() |