| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """NVIDIA HiSparse resident, host, and hot-cache data plane.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import math |
| 8 | import mmap |
| 9 | import time |
| 10 | from collections.abc import Callable |
| 11 | from dataclasses import dataclass |
| 12 | from typing import TYPE_CHECKING, Any, TypeAlias |
| 13 | |
| 14 | import psutil |
| 15 | import torch |
| 16 | |
| 17 | from vllm.config import CUDAGraphMode, VllmConfig |
| 18 | from vllm.distributed import get_tp_group |
| 19 | from vllm.forward_context import get_forward_context, in_piecewise_cudagraph |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.platforms import current_platform |
| 22 | from vllm.utils.math_utils import cdiv, round_up |
| 23 | from vllm.utils.mem_utils import get_max_shared_memory_bytes |
| 24 | from vllm.utils.torch_utils import current_stream |
| 25 | from vllm.v1.attention.backend import max_decode_query_len |
| 26 | from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion |
| 27 | from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | HOST_REGISTER_CHUNK_BYTES = 256 * 2**30 |
| 32 | |
| 33 | if TYPE_CHECKING: |
| 34 | from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup |
| 35 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 36 | |
| 37 | # fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE. |
| 38 | FP8_DS_MLA_ROW_BYTES = 656 |
| 39 | HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor] |
| 40 | |
| 41 | |
| 42 | def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool: |
| 43 | return ( |
| 44 | attn_metadata is not None |
| 45 | and attn_metadata.num_decode_tokens == attn_metadata.num_actual_tokens |
| 46 | and attn_metadata.max_query_len == 1 |
| 47 | and attn_metadata.num_reqs == attn_metadata.num_actual_tokens |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | def _get_max_swap_rows(vllm_config: VllmConfig) -> int: |
| 52 | max_query_len = max_decode_query_len(vllm_config) |
| 53 | scheduler_config = vllm_config.scheduler_config |
| 54 | return min( |
| 55 | scheduler_config.max_num_batched_tokens, |
| 56 | scheduler_config.max_num_seqs * max_query_len, |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | @dataclass(frozen=True) |
| 61 | class PagedCacheView: |
| 62 | cache: torch.Tensor |
| 63 | attention_cache: torch.Tensor |
| 64 | block_size: int |
| 65 | attention_block_stride: int |
| 66 | |
| 67 | @classmethod |
| 68 | def bind( |
| 69 | cls, |
| 70 | raw_tensor: torch.Tensor, |
| 71 | *, |
| 72 | dtype: torch.dtype, |
| 73 | row_width: int, |
| 74 | byte_offset: int, |
| 75 | block_stride: int, |
| 76 | num_blocks: int, |
| 77 | block_size: int, |
| 78 | ) -> PagedCacheView: |
| 79 | itemsize = dtype.itemsize |
| 80 | assert byte_offset % itemsize == 0 and block_stride % itemsize == 0 |
| 81 | cache = torch.as_strided( |
| 82 | raw_tensor.view(dtype), |
| 83 | size=(num_blocks, block_size, row_width), |
| 84 | stride=(block_stride // itemsize, row_width, 1), |
| 85 | storage_offset=byte_offset // itemsize, |
| 86 | ) |
| 87 | row_bytes = row_width * itemsize |
| 88 | if block_stride % (block_size * row_bytes): |
| 89 | return cls(cache, cache, block_size, block_size) |
| 90 | attention_cache = ( |
| 91 | raw_tensor[byte_offset:].view(dtype).view(-1, block_size, row_width) |
| 92 | ) |
| 93 | return cls(cache, attention_cache, block_size, block_stride // row_bytes) |
| 94 | |
| 95 | |
| 96 | @dataclass(frozen=True) |
| 97 | class ResolvedHiSparseConfig: |
| 98 | top_k: int |
| 99 | device_buffer_size: int |
| 100 | # Most host rows one request's decode rows can select in a step: every |
| 101 | # verification row's top-k. |
| 102 | max_union_rows: int |
| 103 | eager_host_mirror: bool = True |
| 104 | |
| 105 | @classmethod |
| 106 | def from_vllm_config( |
| 107 | cls, |
| 108 | vllm_config: VllmConfig, |
| 109 | model_top_k: int, |
| 110 | block_size: int | None = None, |
| 111 | ) -> ResolvedHiSparseConfig | None: |
| 112 | config = vllm_config.attention_config.hisparse_config |
| 113 | if config is None: |
| 114 | return None |
| 115 | |
| 116 | max_query_len = max_decode_query_len(vllm_config) |
| 117 | min_device_buffer_size = max_query_len * model_top_k |
| 118 | # Retain the worst-case union of every speculative step's top-k and |
| 119 | # leave one additional top-k of LRU slack. |
| 120 | configured_size = config.device_buffer_size |
| 121 | if configured_size is None: |
| 122 | device_buffer_size = (max_query_len + 1) * model_top_k |
| 123 | else: |
| 124 | device_buffer_size = configured_size |
| 125 | |
| 126 | if device_buffer_size < min_device_buffer_size: |
| 127 | raise ValueError( |
| 128 | "HiSparse device_buffer_size must cover every decode query's " |
| 129 | "index_topk rows. Got " |
| 130 | f"device_buffer_size={device_buffer_size}, " |
| 131 | f"max_decode_query_len={max_query_len}, " |
| 132 | f"index_topk={model_top_k}; expected at least " |
| 133 | f"{min_device_buffer_size}." |
| 134 | ) |
| 135 | max_device_buffer_size = torch.iinfo(torch.int16).max + 1 |
| 136 | if device_buffer_size > max_device_buffer_size: |
| 137 | raise ValueError( |
| 138 | "HiSparse device_buffer_size exceeds the int16 slot-index " |
| 139 | f"limit: got {device_buffer_size}, maximum is " |
| 140 | f"{max_device_buffer_size}." |
| 141 | ) |
| 142 | if configured_size is not None and block_size is not None: |
| 143 | padding = -device_buffer_size % block_size |
| 144 | if padding: |
| 145 | logger.warning( |
| 146 | "HiSparse device_buffer_size=%d is not aligned to the " |
| 147 | "%d-token kernel block size and allocates %d unused " |
| 148 | "padding rows per hot group. Use %d to use all allocated " |
| 149 | "rows.", |
| 150 | device_buffer_size, |
| 151 | block_size, |
| 152 | padding, |
| 153 | device_buffer_size + padding, |
| 154 | ) |
| 155 | return cls( |
| 156 | top_k=model_top_k, |
| 157 | device_buffer_size=device_buffer_size, |
| 158 | max_union_rows=min_device_buffer_size, |
| 159 | eager_host_mirror=config.eager_host_mirror, |
| 160 | ) |
| 161 | |
| 162 | |
| 163 | def _check_residency_shared_memory( |
| 164 | device: torch.device, config: ResolvedHiSparseConfig |
| 165 | ) -> None: |
| 166 | """Reject a config whose residency resolver exceeds per-block shared memory.""" |
| 167 | # Mirrors hisparse_resolve_residency's layout: union hash keys and values |
| 168 | # (plus an empty slot), per-chunk offsets, five counters, done bits, and |
| 169 | # the int16 compacted LRU. |
| 170 | hot_size = config.device_buffer_size |
| 171 | num_chunks = cdiv(hot_size, 32) |
| 172 | required = 4 * (2 * (config.max_union_rows + 1) + 3 * num_chunks + 7) + 2 * hot_size |
| 173 | device_index = ( |
| 174 | device.index |
| 175 | if device.index is not None |
| 176 | else torch.accelerator.current_device_index() |
| 177 | ) |
| 178 | available = get_max_shared_memory_bytes(device_index) |
| 179 | if required > available: |
| 180 | raise ValueError( |
| 181 | "HiSparse residency resolution needs " |
| 182 | f"{required} B of shared memory per block for " |
| 183 | f"device_buffer_size={config.device_buffer_size} and " |
| 184 | f"{config.max_union_rows} rows per request " |
| 185 | "(decode query length x index_topk), but the device allows " |
| 186 | f"{available} B. Lower device_buffer_size or num_speculative_tokens." |
| 187 | ) |
| 188 | |
| 189 | |
| 190 | def check_hisparse_host_memory(pool_bytes: int) -> None: |
| 191 | """Reject a physical host-pool allocation that cannot fit in RAM.""" |
| 192 | mem = psutil.virtual_memory() |
| 193 | if pool_bytes > mem.available * 0.95: |
| 194 | raise ValueError( |
| 195 | f"HiSparse pinned host pool needs ~{pool_bytes / 2**30:.0f} GiB " |
| 196 | f"but only {mem.available / 2**30:.0f} GiB of RAM is available. " |
| 197 | "Lower the HiSparseConnector host_pool_gib or leave headroom " |
| 198 | "for co-tenants." |
| 199 | ) |
| 200 | |
| 201 | |
| 202 | def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]: |
| 203 | """Allocate and deterministically register an exact-size host KV region.""" |
| 204 | page = 4096 |
| 205 | padded_size = round_up(size, page) |
| 206 | backing = torch.empty(padded_size + page, dtype=torch.int8, device="cpu") |
| 207 | aligned_offset = (-backing.data_ptr()) % page |
| 208 | registered = backing[aligned_offset : aligned_offset + padded_size] |
| 209 | pin_tensor(registered) |
| 210 | return registered[:size], registered |
| 211 | |
| 212 | |
| 213 | class HiSparseHostPool: |
| 214 | """Keep the single host backing and its exact CUDA registration range.""" |
| 215 | |
| 216 | def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None: |
| 217 | self.vllm_config = vllm_config |
| 218 | self.kv_cache_config = kv_cache_config |
| 219 | self.backing: torch.Tensor | None = None |
| 220 | self.registered: torch.Tensor | None = None |
| 221 | self.shared_region: SharedOffloadRegion | None = None |
| 222 | |
| 223 | def allocate(self, size: int) -> torch.Tensor: |
| 224 | assert self.backing is None, "HiSparse host pool is already allocated." |
| 225 | config = self.kv_cache_config |
| 226 | if config.hisparse_shared_host_pool: |
| 227 | tensor_sizes = [ |
| 228 | size |
| 229 | for _, size in sorted( |
| 230 | (tensor.offset + i * tensor.layer_stride, tensor.layer_stride) |
| 231 | for tensor in config.kv_cache_tensors |
| 232 | if tensor.host_resident |
| 233 | for i in range(len(tensor.layers)) |
| 234 | ) |
| 235 | ] |
| 236 | assert config.hisparse_host_num_blocks is not None |
| 237 | assert config.hisparse_host_block_stride is not None |
| 238 | _, _, self.shared_region = allocate_hisparse_host_pools( |
| 239 | self.vllm_config, |
| 240 | tensor_sizes, |
| 241 | config.hisparse_host_num_blocks, |
| 242 | config.hisparse_host_block_stride, |
| 243 | use_shared_host_pool=True, |
| 244 | ) |
| 245 | assert self.shared_region is not None |
| 246 | self.registered = self.shared_region.base_tensor |
| 247 | self.backing = self.registered[:size] |
| 248 | else: |
| 249 | check_hisparse_host_memory(size) |
| 250 | self.backing, self.registered = allocate_pinned_host_pool(size) |
| 251 | return self.backing |
| 252 | |
| 253 | |
| 254 | def use_shared_hisparse_host_pool(vllm_config: VllmConfig) -> bool: |
| 255 | """Whether replicated MLA host KV can share one local mmap.""" |
| 256 | parallel = vllm_config.parallel_config |
| 257 | return ( |
| 258 | current_platform.is_cuda_alike() |
| 259 | and parallel.tensor_parallel_size > 1 |
| 260 | and parallel.pipeline_parallel_size == 1 |
| 261 | and parallel.prefill_context_parallel_size == 1 |
| 262 | and parallel.decode_context_parallel_size == 1 |
| 263 | and parallel.world_size == parallel.tensor_parallel_size |
| 264 | and parallel.distributed_executor_backend == "mp" |
| 265 | and parallel.nnodes_within_dp == 1 |
| 266 | ) |
| 267 | |
| 268 | |
| 269 | def get_hisparse_host_block_stride( |
| 270 | page_size_bytes: int, *, use_shared_host_pool: bool |
| 271 | ) -> int: |
| 272 | if use_shared_host_pool: |
| 273 | return round_up(page_size_bytes, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT) |
| 274 | return page_size_bytes |
| 275 | |
| 276 | |
| 277 | def _hisparse_registration_ranges( |
| 278 | tensor_sizes: list[int], |
| 279 | num_blocks: int, |
| 280 | host_block_stride: int, |
| 281 | max_chunk_bytes: int = HOST_REGISTER_CHUNK_BYTES, |
| 282 | ) -> tuple[tuple[int, int], ...]: |
| 283 | """Split a shared pool without bisecting any tensor's host block. |
| 284 | |
| 285 | CUDA batch copies reject descriptors spanning adjacent registrations. |
| 286 | """ |
| 287 | total_size = num_blocks * host_block_stride |
| 288 | if total_size % mmap.PAGESIZE: |
| 289 | raise ValueError("HiSparse host pool must be OS-page aligned.") |
| 290 | if max_chunk_bytes < mmap.PAGESIZE: |
| 291 | raise ValueError("HiSparse registration chunks must span at least one page.") |
| 292 | |
| 293 | tensors: list[tuple[int, int, int]] = [] |
| 294 | tensor_start = 0 |
| 295 | for tensor_size in tensor_sizes: |
| 296 | if tensor_size % num_blocks: |
| 297 | raise ValueError( |
| 298 | f"HiSparse host tensor size {tensor_size} is not divisible by " |
| 299 | f"{num_blocks} blocks." |
| 300 | ) |
| 301 | tensor_end = tensor_start + tensor_size |
| 302 | tensors.append((tensor_start, tensor_end, tensor_size // num_blocks)) |
| 303 | tensor_start = tensor_end |
| 304 | if tensor_start > total_size: |
| 305 | raise ValueError("HiSparse host tensors exceed the shared host pool.") |
| 306 | |
| 307 | def last_safe_boundary(lower: int, upper: int) -> int | None: |
| 308 | best = None |
| 309 | for start, end, block_bytes in tensors: |
| 310 | first_block = max(0, (lower - start) // block_bytes + 1) |
| 311 | last_block = min(num_blocks, (min(upper, end) - start) // block_bytes) |
| 312 | if first_block > last_block: |
| 313 | continue |
| 314 | divisor = math.gcd(block_bytes, mmap.PAGESIZE) |
| 315 | if start % divisor: |
| 316 | continue |
| 317 | period = mmap.PAGESIZE // divisor |
| 318 | if period == 1: |
| 319 | residue = 0 |
| 320 | else: |
| 321 | # Find block boundaries that are also OS-page boundaries. |
| 322 | residue = ( |
| 323 | (-start // divisor) * pow(block_bytes // divisor, -1, period) |
| 324 | ) % period |
| 325 | block = last_block - (last_block - residue) % period |
| 326 | if block >= first_block: |
| 327 | candidate = start + block * block_bytes |
| 328 | best = candidate if best is None else max(best, candidate) |
| 329 | |
| 330 | if upper >= tensor_start: |
| 331 | candidate = upper // mmap.PAGESIZE * mmap.PAGESIZE |
| 332 | if candidate >= tensor_start and candidate > lower: |
| 333 | best = candidate if best is None else max(best, candidate) |
| 334 | return best |
| 335 | |
| 336 | ranges: list[tuple[int, int]] = [] |
| 337 | start = 0 |
| 338 | while total_size - start > max_chunk_bytes: |
| 339 | end = last_safe_boundary(start, start + max_chunk_bytes) |
| 340 | if end is None: |
| 341 | raise ValueError( |
| 342 | "A HiSparse host block is too large for the registration chunk limit." |
| 343 | ) |
| 344 | ranges.append((start, end)) |
| 345 | start = end |
| 346 | ranges.append((start, total_size)) |
| 347 | return tuple(ranges) |
| 348 | |
| 349 | |
| 350 | def allocate_hisparse_host_pools( |
| 351 | vllm_config: VllmConfig, |
| 352 | tensor_sizes: list[int], |
| 353 | num_blocks: int, |
| 354 | host_block_stride: int, |
| 355 | *, |
| 356 | use_shared_host_pool: bool, |
| 357 | ) -> tuple[list[torch.Tensor], list[torch.Tensor], SharedOffloadRegion | None]: |
| 358 | """Allocate private host tensors or one mmap shared by local TP ranks.""" |
| 359 | if not use_shared_host_pool: |
| 360 | check_hisparse_host_memory(sum(tensor_sizes)) |
| 361 | private_pools = [allocate_pinned_host_pool(size) for size in tensor_sizes] |
| 362 | return ( |
| 363 | [pool for pool, _ in private_pools], |
| 364 | [registered for _, registered in private_pools], |
| 365 | None, |
| 366 | ) |
| 367 | |
| 368 | region = SharedOffloadRegion( |
| 369 | engine_id=( |
| 370 | f"hisparse_{vllm_config.instance_id}_" |
| 371 | f"dp{vllm_config.parallel_config.data_parallel_index}" |
| 372 | ), |
| 373 | num_chunks=1, |
| 374 | rank=0, |
| 375 | kv_bytes_per_chunk=num_blocks * host_block_stride, |
| 376 | cpu_page_size=sum(tensor_sizes), |
| 377 | barrier=get_tp_group().barrier, |
| 378 | creator_memory_check=check_hisparse_host_memory, |
| 379 | populate_only_on_creator=True, |
| 380 | ) |
| 381 | try: |
| 382 | for start, end in _hisparse_registration_ranges( |
| 383 | tensor_sizes, num_blocks, host_block_stride |
| 384 | ): |
| 385 | tensor = region.base_tensor[start:end] |
| 386 | pin_tensor(tensor) |
| 387 | region.pinned_addresses.append(tensor.data_ptr()) |
| 388 | region.is_pinned = True |
| 389 | pools = [ |
| 390 | region.create_next_canonical_view(size).view(-1) for size in tensor_sizes |
| 391 | ] |
| 392 | except Exception: |
| 393 | region.cleanup() |
| 394 | raise |
| 395 | return pools, [], region |
| 396 | |
| 397 | |
| 398 | def release_pinned_state( |
| 399 | runtimes: list[HiSparseRuntime], |
| 400 | pinned_host_pools: list[torch.Tensor], |
| 401 | shared_host_region: SharedOffloadRegion | None = None, |
| 402 | ) -> None: |
| 403 | """Synchronize and release registered host KV pools.""" |
| 404 | if pinned_host_pools or shared_host_region is not None: |
| 405 | try: |
| 406 | torch.accelerator.synchronize() |
| 407 | except RuntimeError as e: |
| 408 | logger.warning( |
| 409 | "HiSparse: CUDA context unusable at teardown (%s); leaving " |
| 410 | "%d host-pool tensors pinned for kernel exit reclaim.", |
| 411 | e, |
| 412 | len(pinned_host_pools) + int(shared_host_region is not None), |
| 413 | ) |
| 414 | return |
| 415 | |
| 416 | cudart = torch.cuda.cudart() |
| 417 | release_start = time.perf_counter() |
| 418 | freed_bytes = 0 |
| 419 | while pinned_host_pools: |
| 420 | tensor = pinned_host_pools[-1] |
| 421 | err = cudart.cudaHostUnregister(tensor.data_ptr()) |
| 422 | if err.value != 0: |
| 423 | logger.warning( |
| 424 | "HiSparse: cudaHostUnregister failed (code=%d); leaving " |
| 425 | "%d host-pool tensors pinned for kernel exit reclaim.", |
| 426 | err.value, |
| 427 | len(pinned_host_pools), |
| 428 | ) |
| 429 | cudart.cudaGetLastError() |
| 430 | break |
| 431 | freed_bytes += tensor.nbytes |
| 432 | pinned_host_pools.pop() |
| 433 | if freed_bytes: |
| 434 | logger.info( |
| 435 | "HiSparse: unpinned %.1f GiB of host pool in %.1fs.", |
| 436 | freed_bytes / 2**30, |
| 437 | time.perf_counter() - release_start, |
| 438 | ) |
| 439 | |
| 440 | for runtime in runtimes: |
| 441 | del runtime._host_cache |
| 442 | del runtime.registered_host_pool |
| 443 | del runtime.hot_backing |
| 444 | if shared_host_region is not None: |
| 445 | shared_host_region.cleanup() |
| 446 | |
| 447 | |
| 448 | def hisparse_prefill_staging_remap( |
| 449 | block_table: torch.Tensor, block_size: int |
| 450 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 451 | """Renumber blocks without a data-dependent CUDA output allocation.""" |
| 452 | block_ids = torch.cat( |
| 453 | ( |
| 454 | torch.zeros(1, dtype=torch.int32, device=block_table.device), |
| 455 | block_table.flatten().clamp(min=0).to(torch.int32), |
| 456 | ) |
| 457 | ) |
| 458 | sorted_ids, permutation = torch.sort(block_ids) |
| 459 | is_new = torch.cat( |
| 460 | ( |
| 461 | torch.ones(1, dtype=torch.bool, device=block_table.device), |
| 462 | sorted_ids[1:] != sorted_ids[:-1], |
| 463 | ) |
| 464 | ) |
| 465 | compact_sorted = torch.cumsum(is_new, dim=0, dtype=torch.int32) - 1 |
| 466 | inverse = torch.empty_like(compact_sorted) |
| 467 | inverse.scatter_(0, permutation, compact_sorted) |
| 468 | |
| 469 | unique_ids = torch.full_like(sorted_ids, -1) |
| 470 | unique_ids.scatter_(0, compact_sorted, sorted_ids) |
| 471 | valid = ( |
| 472 | torch.arange(sorted_ids.numel(), device=block_table.device) |
| 473 | <= compact_sorted[-1] |
| 474 | ) |
| 475 | unique_ids = torch.where(valid, unique_ids, -1) |
| 476 | new_bt = inverse[1:].view_as(block_table) |
| 477 | offsets = torch.arange(block_size, dtype=torch.int32, device=block_table.device) |
| 478 | row_ids = ( |
| 479 | torch.where( |
| 480 | valid[:, None], |
| 481 | unique_ids[:, None] * block_size + offsets, |
| 482 | -1, |
| 483 | ) |
| 484 | ).view(1, -1) |
| 485 | return new_bt, row_ids |
| 486 | |
| 487 | |
| 488 | @dataclass |
| 489 | class HiSparsePrefillStagingPlan: |
| 490 | block_table: torch.Tensor |
| 491 | row_ids: torch.Tensor |
| 492 | dst_rows: torch.Tensor |
| 493 | miss_mask: torch.Tensor |
| 494 | block_size: int |
| 495 | # Host rows with a GPU-resident copy (adopted shadow pages): the flat |
| 496 | # resident-cache row to read instead of DMAing from host, -1 for misses. |
| 497 | gpu_row_ids: torch.Tensor | None = None |
| 498 | gpu_source_key: tuple[int, int] | None = None |
| 499 | |
| 500 | def ensure_gpu_sources( |
| 501 | self, |
| 502 | resident_block_table: torch.Tensor, |
| 503 | resident_block_size: int, |
| 504 | ) -> None: |
| 505 | """Resolve which staged rows can be served from the resident cache. |
| 506 | |
| 507 | Computed once per plan (the resident block table is shared by every |
| 508 | layer in the group); non-null resident pages become miss_mask=0 rows |
| 509 | gathered device-to-device by ``gather_prefill_cache``. |
| 510 | """ |
| 511 | source_key = (resident_block_table.data_ptr(), resident_block_size) |
| 512 | if self.gpu_source_key == source_key: |
| 513 | return |
| 514 | block_size = self.block_size |
| 515 | if resident_block_size <= 0 or block_size % resident_block_size != 0: |
| 516 | return |
| 517 | device = self.row_ids.device |
| 518 | num_unique = self.row_ids.shape[1] // block_size |
| 519 | host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size |
| 520 | new_bt = self.block_table.to(torch.int64) |
| 521 | num_rows, num_cols = new_bt.shape |
| 522 | if num_rows == 0 or resident_block_table.shape[0] < num_rows: |
| 523 | return |
| 524 | # One representative (row, col) per unique host block: any request |
| 525 | # referencing the block holds an equivalent (refcounted) resident view. |
| 526 | flat_pos = torch.arange(num_rows * num_cols, device=device) |
| 527 | rep = torch.full( |
| 528 | (num_unique,), num_rows * num_cols, device=device, dtype=torch.int64 |
| 529 | ) |
| 530 | rep.scatter_reduce_( |
| 531 | 0, new_bt.reshape(-1), flat_pos, reduce="amin", include_self=True |
| 532 | ) |
| 533 | rep_row = (rep // num_cols).clamp(max=resident_block_table.shape[0] - 1) |
| 534 | rep_col = rep % num_cols |
| 535 | pages_per_block = block_size // resident_block_size |
| 536 | res_pos = rep_col[:, None] * pages_per_block + torch.arange( |
| 537 | pages_per_block, device=device |
| 538 | ) |
| 539 | res_cols = resident_block_table.shape[1] |
| 540 | res_blocks = torch.where( |
| 541 | res_pos < res_cols, |
| 542 | torch.gather( |
| 543 | resident_block_table.to(torch.int64)[rep_row], |
| 544 | 1, |
| 545 | res_pos.clamp(max=max(res_cols - 1, 0)), |
| 546 | ), |
| 547 | torch.zeros_like(res_pos), |
| 548 | ) |
| 549 | offsets = torch.arange(block_size, device=device) |
| 550 | per_off_block = res_blocks[:, offsets // resident_block_size] |
| 551 | gpu_rows = ( |
| 552 | per_off_block * resident_block_size |
| 553 | + (offsets % resident_block_size)[None, :] |
| 554 | ) |
| 555 | # Block id 0 is the null block in both id spaces. |
| 556 | valid_rows = self.row_ids.view(num_unique, block_size) >= 0 |
| 557 | hit = (per_off_block > 0) & (host_ids[:, None] > 0) & valid_rows |
| 558 | self.gpu_row_ids = torch.where(hit, gpu_rows, -1).reshape(1, -1).to(torch.int32) |
| 559 | self.miss_mask = ((self.gpu_row_ids < 0) & valid_rows.view(1, -1)).int() |
| 560 | self.gpu_source_key = source_key |
| 561 | |
| 562 | |
| 563 | def build_hisparse_prefill_staging_plan( |
| 564 | block_table: torch.Tensor, |
| 565 | seq_lens: torch.Tensor, |
| 566 | block_size: int, |
| 567 | staging_block_capacity: int, |
| 568 | ) -> HiSparsePrefillStagingPlan: |
| 569 | """Build an asynchronous layer-independent host-cache staging remap.""" |
| 570 | device = block_table.device |
| 571 | used = (seq_lens.to(torch.int64) + block_size - 1) // block_size |
| 572 | packed_offsets = torch.arange(staging_block_capacity, device=device) |
| 573 | row_ends = torch.cumsum(used, dim=0) |
| 574 | rows = torch.searchsorted(row_ends, packed_offsets, right=True) |
| 575 | valid = rows < block_table.shape[0] |
| 576 | rows = rows.clamp(max=block_table.shape[0] - 1) |
| 577 | row_starts = torch.cat((torch.zeros_like(row_ends[:1]), row_ends[:-1])) |
| 578 | columns = packed_offsets - row_starts[rows] |
| 579 | valid &= columns < block_table.shape[1] |
| 580 | flat_positions = rows * block_table.shape[1] + columns.clamp( |
| 581 | max=block_table.shape[1] - 1 |
| 582 | ) |
| 583 | referenced = torch.where(valid, block_table.flatten()[flat_positions], 0) |
| 584 | referenced_bt, row_ids = hisparse_prefill_staging_remap(referenced, block_size) |
| 585 | scratch_index = block_table.numel() |
| 586 | scatter_positions = torch.where(valid, flat_positions, scratch_index) |
| 587 | remapped = torch.where(valid, referenced_bt.flatten(), 0) |
| 588 | new_bt_storage = torch.zeros(scratch_index + 1, dtype=torch.int32, device=device) |
| 589 | new_bt_storage.scatter_(0, scatter_positions, remapped) |
| 590 | new_bt = new_bt_storage[:-1].view_as(block_table) |
| 591 | valid_rows = row_ids >= 0 |
| 592 | dst_rows = torch.arange(row_ids.shape[1], dtype=torch.int32, device=device).view( |
| 593 | 1, -1 |
| 594 | ) |
| 595 | return HiSparsePrefillStagingPlan( |
| 596 | block_table=new_bt, |
| 597 | row_ids=row_ids, |
| 598 | dst_rows=dst_rows, |
| 599 | miss_mask=valid_rows.to(torch.int32), |
| 600 | block_size=block_size, |
| 601 | ) |
| 602 | |
| 603 | |
| 604 | def _has_hisparse_ops() -> bool: |
| 605 | if not hasattr(torch.ops, "_C_cache_ops"): |
| 606 | return False |
| 607 | return ( |
| 608 | hasattr(torch.ops._C_cache_ops, "hisparse_resolve_residency") |
| 609 | and hasattr(torch.ops._C_cache_ops, "hisparse_invalidate_written_slots") |
| 610 | and hasattr(torch.ops._C_cache_ops, "hisparse_gather_plan") |
| 611 | and hasattr(torch.ops._C_cache_ops, "hisparse_gather_compact") |
| 612 | ) |
| 613 | |
| 614 | |
| 615 | class _SharedTopKState: |
| 616 | """CUDA-graph-safe residency result shared by index-sharing layers.""" |
| 617 | |
| 618 | def __init__(self, device: torch.device, max_rows: int, top_k: int) -> None: |
| 619 | self.device_topk_rows = torch.full( |
| 620 | (max_rows, top_k), -1, dtype=torch.int32, device=device |
| 621 | ) |
| 622 | self.physical_topk_indices = torch.empty_like(self.device_topk_rows) |
| 623 | self.swap_host_physical_rows = torch.empty( |
| 624 | (max_rows, top_k), dtype=torch.int32, device=device |
| 625 | ) |
| 626 | self.swap_device_physical_rows = torch.empty_like(self.swap_host_physical_rows) |
| 627 | self.swap_counts = torch.empty(max_rows, dtype=torch.int32, device=device) |
| 628 | self.valid_topk_counts = torch.empty(max_rows, dtype=torch.int32, device=device) |
| 629 | |
| 630 | |
| 631 | def _create_shared_topk_state( |
| 632 | device: torch.device, max_rows: int, top_k: int |
| 633 | ) -> _SharedTopKState: |
| 634 | return _SharedTopKState(device, max_rows, top_k) |
| 635 | |
| 636 | |
| 637 | def _create_copy_stream(device: torch.device) -> torch.Stream: |
| 638 | return torch.Stream(device=device) |
| 639 | |
| 640 | |
| 641 | class HiSparseIndexGroup: |
| 642 | """GPU index-resolution state shared by one leader and its followers.""" |
| 643 | |
| 644 | leader: HiSparseRuntime |
| 645 | |
| 646 | def __init__( |
| 647 | self, |
| 648 | device: torch.device, |
| 649 | max_num_reqs: int, |
| 650 | region_stride: int, |
| 651 | max_swap_rows: int, |
| 652 | top_k: int, |
| 653 | copy_stream: torch.Stream | None = None, |
| 654 | logical_topk_ready: torch.Event | None = None, |
| 655 | ) -> None: |
| 656 | self.device_global_indices = torch.full( |
| 657 | (max_num_reqs, region_stride), |
| 658 | -1, |
| 659 | dtype=torch.int32, |
| 660 | device=device, |
| 661 | ) |
| 662 | self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device) |
| 663 | self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous() |
| 664 | self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k) |
| 665 | # Request ids the resolver reads, built on the copy stream: a |
| 666 | # compute-stream write would race with its logical_topk_ready wait. |
| 667 | self.row_indices = torch.arange(max_swap_rows, dtype=torch.int32, device=device) |
| 668 | self.request_ids = torch.empty_like(self.row_indices) |
| 669 | self.copy_stream = ( |
| 670 | copy_stream if copy_stream is not None else _create_copy_stream(device) |
| 671 | ) |
| 672 | self.logical_topk_ready = logical_topk_ready |
| 673 | self.followers: list[HiSparseRuntime] = [] |
| 674 | self.swap_stats = torch.zeros(2, dtype=torch.uint64, device=device) |
| 675 | self.swap_stats_host = torch.empty( |
| 676 | 2, dtype=torch.uint64, device="cpu", pin_memory=True |
| 677 | ) |
| 678 | self.stats_row_bytes = 0 |
| 679 | |
| 680 | |
| 681 | class HiSparseRuntime: |
| 682 | """Per-cache host/hot data plane and GPU replacement state.""" |
| 683 | |
| 684 | hot: PagedCacheView |
| 685 | hot_backing: torch.Tensor |
| 686 | hot_block_table: torch.Tensor |
| 687 | _host_cache: torch.Tensor |
| 688 | registered_host_pool: torch.Tensor |
| 689 | |
| 690 | def __init__( |
| 691 | self, |
| 692 | config: ResolvedHiSparseConfig, |
| 693 | max_num_reqs: int, |
| 694 | row_width: int, |
| 695 | kv_dtype: torch.dtype, |
| 696 | device: torch.device | str, |
| 697 | max_swap_rows: int | None = None, |
| 698 | index_group: HiSparseIndexGroup | None = None, |
| 699 | copy_stream: torch.Stream | None = None, |
| 700 | logical_topk_ready: torch.Event | None = None, |
| 701 | ) -> None: |
| 702 | if not _has_hisparse_ops(): |
| 703 | raise RuntimeError( |
| 704 | "HiSparse requires its compiled _C_cache_ops CUDA kernels " |
| 705 | "(host-resident decode has no Python fallback). Rebuild vLLM " |
| 706 | "from source so " |
| 707 | "csrc/libtorch_stable/hisparse_kernels.cu is included." |
| 708 | ) |
| 709 | self.max_num_reqs = max_num_reqs |
| 710 | self.row_width = row_width |
| 711 | self.kv_dtype = kv_dtype |
| 712 | self.device = torch.device(device) |
| 713 | # Logical slots per request. Physical rows come from its ephemeral HMA |
| 714 | # block table and need not be contiguous in the shared slab. |
| 715 | self.region_stride = config.device_buffer_size |
| 716 | self.max_union_rows = config.max_union_rows |
| 717 | if self.device.type == "cuda": |
| 718 | _check_residency_shared_memory(self.device, config) |
| 719 | |
| 720 | row_bytes = row_width * kv_dtype.itemsize |
| 721 | if row_bytes % 16 != 0: |
| 722 | raise ValueError( |
| 723 | f"HiSparse requires 16-byte aligned KV rows, got {row_bytes}B." |
| 724 | ) |
| 725 | |
| 726 | self._layer_ready_event: torch.Event | None = None |
| 727 | self._swap_staged = False |
| 728 | self._swap_step = 0 |
| 729 | self.is_group_leader = index_group is None |
| 730 | if index_group is None: |
| 731 | index_group = HiSparseIndexGroup( |
| 732 | self.device, |
| 733 | max_num_reqs, |
| 734 | self.region_stride, |
| 735 | max_swap_rows or max_num_reqs, |
| 736 | config.top_k, |
| 737 | copy_stream, |
| 738 | logical_topk_ready, |
| 739 | ) |
| 740 | index_group.leader = self |
| 741 | else: |
| 742 | index_group.followers.append(self) |
| 743 | self.index_group = index_group |
| 744 | index_group.stats_row_bytes += row_bytes |
| 745 | |
| 746 | self.eager_host_mirror = config.eager_host_mirror |
| 747 | self.resident_source_index = -1 |
| 748 | self.request_state_indices: torch.Tensor | None = None |
| 749 | self.shared_host_region: SharedOffloadRegion | None = None |
| 750 | |
| 751 | @property |
| 752 | def host_cache(self) -> torch.Tensor: |
| 753 | return self._host_cache |
| 754 | |
| 755 | def bind_hot_cache( |
| 756 | self, |
| 757 | raw_tensor: torch.Tensor, |
| 758 | *, |
| 759 | byte_offset: int, |
| 760 | block_stride: int, |
| 761 | num_blocks: int, |
| 762 | block_size: int, |
| 763 | block_table: torch.Tensor, |
| 764 | ) -> None: |
| 765 | """Bind this runtime's strided view into the shared GPU HMA slab.""" |
| 766 | self.hot = PagedCacheView.bind( |
| 767 | raw_tensor, |
| 768 | dtype=self.kv_dtype, |
| 769 | row_width=self.row_width, |
| 770 | byte_offset=byte_offset, |
| 771 | block_stride=block_stride, |
| 772 | num_blocks=num_blocks, |
| 773 | block_size=block_size, |
| 774 | ) |
| 775 | self.hot_backing = raw_tensor |
| 776 | self.hot_block_table = block_table |
| 777 | |
| 778 | def bind_source_cache( |
| 779 | self, |
| 780 | kv_cache: torch.Tensor, |
| 781 | *, |
| 782 | registered_host_pool: torch.Tensor | None = None, |
| 783 | ) -> None: |
| 784 | if kv_cache.dtype != self.kv_dtype or kv_cache.shape[-1] != self.row_width: |
| 785 | raise ValueError( |
| 786 | "HiSparse runtime bound to a KV cache with mismatched " |
| 787 | f"layout: expected ({self.row_width}, {self.kv_dtype}), got " |
| 788 | f"({kv_cache.shape[-1]}, {kv_cache.dtype})." |
| 789 | ) |
| 790 | # Host-resident pool: the cache itself is the only full-size store; |
| 791 | # every allocated slot is written before the indexer can select it. |
| 792 | if kv_cache.device.type != "cpu": |
| 793 | raise ValueError( |
| 794 | "HiSparse requires a host-resident KV pool; got a KV cache " |
| 795 | f"on {kv_cache.device}." |
| 796 | ) |
| 797 | # The pool is pinned via cudaHostRegister (exact-size, deterministic |
| 798 | # unpin at shutdown); torch's is_pinned() only recognizes its own |
| 799 | # caching-host-allocator memory, so also accept ranges the model |
| 800 | # allocator explicitly registered. |
| 801 | if not (kv_cache.is_pinned() or registered_host_pool is not None): |
| 802 | raise ValueError("HiSparse host-resident KV pool must be pinned memory.") |
| 803 | |
| 804 | self._host_cache = kv_cache.view(-1, kv_cache.shape[-1]) |
| 805 | self.registered_host_pool = ( |
| 806 | registered_host_pool if registered_host_pool is not None else kv_cache |
| 807 | ) |
| 808 | |
| 809 | def gather_prefill_cache( |
| 810 | self, |
| 811 | kv_cache: torch.Tensor, |
| 812 | plan: HiSparsePrefillStagingPlan, |
| 813 | resident_cache: torch.Tensor | None = None, |
| 814 | ) -> torch.Tensor: |
| 815 | """Gather this runtime's host cache using a shared staging plan.""" |
| 816 | if kv_cache.shape[1] != plan.block_size: |
| 817 | raise ValueError( |
| 818 | f"HiSparse staging block size {plan.block_size} does not match " |
| 819 | f"the KV cache block size {kv_cache.shape[1]}." |
| 820 | ) |
| 821 | row_width = kv_cache.shape[-1] |
| 822 | |
| 823 | staged = torch.empty( |
| 824 | ( |
| 825 | plan.row_ids.shape[1] // plan.block_size, |
| 826 | plan.block_size, |
| 827 | row_width, |
| 828 | ), |
| 829 | dtype=kv_cache.dtype, |
| 830 | device=plan.block_table.device, |
| 831 | ) |
| 832 | if plan.gpu_row_ids is not None and resident_cache is not None: |
| 833 | gpu_rows = plan.gpu_row_ids[0].to(torch.long) |
| 834 | src = gpu_rows.clamp_min(0) |
| 835 | resident_block_size = resident_cache.shape[1] |
| 836 | rows = resident_cache[src // resident_block_size, src % resident_block_size] |
| 837 | if rows.dtype != staged.dtype: |
| 838 | rows = rows.contiguous().view(staged.dtype) |
| 839 | staged.view(-1, row_width).copy_(rows) |
| 840 | torch.ops._C_cache_ops.hisparse_gather_plan( |
| 841 | kv_cache.view(-1, row_width), |
| 842 | staged, |
| 843 | plan.row_ids, |
| 844 | plan.dst_rows, |
| 845 | plan.miss_mask, |
| 846 | None, |
| 847 | None, |
| 848 | 0, |
| 849 | ) |
| 850 | return staged |
| 851 | |
| 852 | def reset_hot_state(self) -> None: |
| 853 | """Drop all hot-buffer bookkeeping (hits become misses).""" |
| 854 | group = self.index_group |
| 855 | compute_stream = current_stream() |
| 856 | compute_stream.wait_stream(group.copy_stream) |
| 857 | group.device_global_indices.fill_(-1) |
| 858 | group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots)) |
| 859 | group.swap_stats.zero_() |
| 860 | group.copy_stream.wait_stream(compute_stream) |
| 861 | |
| 862 | def invalidate_sorted_slots( |
| 863 | self, |
| 864 | sorted_slots: torch.Tensor, |
| 865 | state_indices: torch.Tensor, |
| 866 | ) -> None: |
| 867 | """Drop hot copies using shared, preprocessed invalidation inputs.""" |
| 868 | device_global_indices = self.index_group.device_global_indices |
| 869 | active_indices = device_global_indices.index_select(0, state_indices) |
| 870 | positions = torch.searchsorted(sorted_slots, active_indices) |
| 871 | positions.clamp_(max=sorted_slots.numel() - 1) |
| 872 | active_indices.masked_fill_( |
| 873 | sorted_slots[positions] == active_indices, |
| 874 | -1, |
| 875 | ) |
| 876 | device_global_indices.index_copy_(0, state_indices, active_indices) |
| 877 | |
| 878 | def invalidate_written_slots( |
| 879 | self, |
| 880 | written_slots: torch.Tensor, |
| 881 | req_id_per_token: torch.Tensor, |
| 882 | ) -> None: |
| 883 | num_tokens = min(written_slots.numel(), req_id_per_token.numel()) |
| 884 | if num_tokens == 0: |
| 885 | return |
| 886 | assert self.request_state_indices is not None |
| 887 | torch.ops._C_cache_ops.hisparse_invalidate_written_slots( |
| 888 | self.index_group.device_global_indices, |
| 889 | self.request_state_indices, |
| 890 | req_id_per_token[:num_tokens], |
| 891 | written_slots[:num_tokens], |
| 892 | ) |
| 893 | |
| 894 | def begin_forward(self) -> None: |
| 895 | self._swap_step = 0 |
| 896 | |
| 897 | def _step_rows(self, num_tokens: int) -> slice: |
| 898 | start = self._swap_step * num_tokens |
| 899 | stop = start + num_tokens |
| 900 | if stop > self.index_group.shared_topk.physical_topk_indices.shape[0]: |
| 901 | raise ValueError( |
| 902 | "HiSparse swap rows exceed the configured speculative decode " |
| 903 | f"capacity: stop={stop}, capacity=" |
| 904 | f"{self.index_group.shared_topk.physical_topk_indices.shape[0]}." |
| 905 | ) |
| 906 | self._swap_step += 1 |
| 907 | return slice(start, stop) |
| 908 | |
| 909 | def _resolve_residency( |
| 910 | self, |
| 911 | *, |
| 912 | resident: HiSparseCacheHandle | None = None, |
| 913 | req_id_per_token: torch.Tensor, |
| 914 | block_table: torch.Tensor, |
| 915 | topk_indices: torch.Tensor, |
| 916 | block_size: int, |
| 917 | return_valid_counts: bool = False, |
| 918 | shared_rows: slice, |
| 919 | attention_indices_out: torch.Tensor | None = None, |
| 920 | valid_counts_out: torch.Tensor | None = None, |
| 921 | ) -> None: |
| 922 | """Resolve logical top-k positions and compact rows requiring swaps.""" |
| 923 | num_tokens = topk_indices.shape[0] |
| 924 | hot = self.hot |
| 925 | host_cache = self.host_cache |
| 926 | assert hot.block_size == block_size |
| 927 | request_state_indices = self.request_state_indices |
| 928 | assert request_state_indices is not None |
| 929 | group = self.index_group |
| 930 | shared = group.shared_topk |
| 931 | |
| 932 | device_topk_rows = shared.device_topk_rows[shared_rows] |
| 933 | valid_topk_counts = None |
| 934 | if return_valid_counts: |
| 935 | valid_topk_counts = ( |
| 936 | shared.valid_topk_counts[shared_rows] |
| 937 | if valid_counts_out is None |
| 938 | else valid_counts_out |
| 939 | ) |
| 940 | physical_topk_indices = ( |
| 941 | shared.physical_topk_indices[shared_rows] |
| 942 | if attention_indices_out is None |
| 943 | else attention_indices_out |
| 944 | ) |
| 945 | swap_host_physical_rows = shared.swap_host_physical_rows[shared_rows] |
| 946 | swap_device_physical_rows = shared.swap_device_physical_rows[shared_rows] |
| 947 | swap_counts = shared.swap_counts[shared_rows] |
| 948 | |
| 949 | # Padded rows are skipped by the kernel (request_state_indices) and must |
| 950 | # come out as -1 so the attention kernel masks them. |
| 951 | torch.ops._C_cache_ops.hisparse_resolve_residency( |
| 952 | host_cache, |
| 953 | hot.cache, |
| 954 | self.hot_block_table, |
| 955 | topk_indices, |
| 956 | device_topk_rows, |
| 957 | group.device_global_indices, |
| 958 | group.lru_slots, |
| 959 | request_state_indices, |
| 960 | self.region_stride, |
| 961 | self.max_union_rows, |
| 962 | None, |
| 963 | group.swap_stats, |
| 964 | physical_topk_indices, |
| 965 | hot.attention_block_stride, |
| 966 | req_id_per_token[:num_tokens].contiguous(), |
| 967 | block_table, |
| 968 | block_size, |
| 969 | None, |
| 970 | valid_topk_counts, |
| 971 | swap_host_physical_rows, |
| 972 | swap_device_physical_rows, |
| 973 | swap_counts, |
| 974 | resident.block_table if resident is not None else None, |
| 975 | resident.view.block_size |
| 976 | if resident is not None and resident.view is not None |
| 977 | else 0, |
| 978 | 0, |
| 979 | ) |
| 980 | |
| 981 | def _swap_rows(self, shared_rows: slice) -> None: |
| 982 | hot = self.hot |
| 983 | shared = self.index_group.shared_topk |
| 984 | torch.ops._C_cache_ops.hisparse_gather_compact( |
| 985 | self.host_cache, |
| 986 | hot.cache, |
| 987 | shared.swap_host_physical_rows[shared_rows], |
| 988 | shared.swap_device_physical_rows[shared_rows], |
| 989 | shared.swap_counts[shared_rows], |
| 990 | ) |
| 991 | |
| 992 | def _resolve_and_stage_group( |
| 993 | self, |
| 994 | *, |
| 995 | resident: HiSparseCacheHandle, |
| 996 | req_id_per_token: torch.Tensor, |
| 997 | block_table: torch.Tensor, |
| 998 | logical_topk_indices: torch.Tensor, |
| 999 | block_size: int, |
| 1000 | return_valid_counts: bool, |
| 1001 | shared_rows: slice, |
| 1002 | attention_indices_out: torch.Tensor | None, |
| 1003 | valid_counts_out: torch.Tensor | None, |
| 1004 | num_valid_rows: torch.Tensor, |
| 1005 | ) -> None: |
| 1006 | group = self.index_group |
| 1007 | compute_stream = current_stream() |
| 1008 | if group.logical_topk_ready is not None and not in_piecewise_cudagraph(): |
| 1009 | group.copy_stream.wait_event(group.logical_topk_ready) |
| 1010 | else: |
| 1011 | group.copy_stream.wait_stream(compute_stream) |
| 1012 | with group.copy_stream: |
| 1013 | # CUDA-graph padding rows past the batch's tokens map to request 0; |
| 1014 | # mark them -1 so residency resolution skips them. |
| 1015 | num_tokens = logical_topk_indices.shape[0] |
| 1016 | request_ids = group.request_ids[:num_tokens] |
| 1017 | request_ids.copy_(req_id_per_token) |
| 1018 | request_ids.masked_fill_( |
| 1019 | group.row_indices[:num_tokens] >= num_valid_rows, -1 |
| 1020 | ) |
| 1021 | self._resolve_residency( |
| 1022 | resident=resident, |
| 1023 | req_id_per_token=request_ids, |
| 1024 | block_table=block_table, |
| 1025 | topk_indices=logical_topk_indices, |
| 1026 | block_size=block_size, |
| 1027 | return_valid_counts=return_valid_counts, |
| 1028 | shared_rows=shared_rows, |
| 1029 | attention_indices_out=attention_indices_out, |
| 1030 | valid_counts_out=valid_counts_out, |
| 1031 | ) |
| 1032 | runtimes = [self] |
| 1033 | if resident.decode_batch: |
| 1034 | runtimes.extend(group.followers) |
| 1035 | for runtime in runtimes: |
| 1036 | runtime._swap_rows(shared_rows) |
| 1037 | if runtime._layer_ready_event is None: |
| 1038 | runtime._layer_ready_event = torch.Event() |
| 1039 | runtime._layer_ready_event.record(group.copy_stream) |
| 1040 | runtime._swap_staged = True |
| 1041 | |
| 1042 | def _stage_this_layer(self, shared_rows: slice) -> None: |
| 1043 | group = self.index_group |
| 1044 | compute_stream = current_stream() |
| 1045 | group.copy_stream.wait_stream(compute_stream) |
| 1046 | with group.copy_stream: |
| 1047 | self._swap_rows(shared_rows) |
| 1048 | if self._layer_ready_event is None: |
| 1049 | self._layer_ready_event = torch.Event() |
| 1050 | self._layer_ready_event.record(group.copy_stream) |
| 1051 | self._swap_staged = True |
| 1052 | |
| 1053 | def wait_for_staged_swap(self) -> None: |
| 1054 | if not self._swap_staged or self._layer_ready_event is None: |
| 1055 | raise RuntimeError("HiSparse layer swap was not prefetched.") |
| 1056 | current_stream().wait_event(self._layer_ready_event) |
| 1057 | self._swap_staged = False |
| 1058 | |
| 1059 | def swap_in( |
| 1060 | self, |
| 1061 | *, |
| 1062 | resident: HiSparseCacheHandle, |
| 1063 | req_id_per_token: torch.Tensor, |
| 1064 | block_table: torch.Tensor, |
| 1065 | logical_topk_indices: torch.Tensor, |
| 1066 | block_size: int, |
| 1067 | num_valid_rows: torch.Tensor, |
| 1068 | return_valid_counts: bool = False, |
| 1069 | attention_indices_out: torch.Tensor | None = None, |
| 1070 | valid_counts_out: torch.Tensor | None = None, |
| 1071 | ) -> HiSparseTopKResult: |
| 1072 | """Resolve once per group and ensure this layer's rows are available.""" |
| 1073 | num_tokens = logical_topk_indices.shape[0] |
| 1074 | shared_rows = self._step_rows(num_tokens) |
| 1075 | group = self.index_group |
| 1076 | hot = self.hot |
| 1077 | assert hot.block_size == block_size |
| 1078 | assert hot.attention_block_stride == group.leader.hot.attention_block_stride |
| 1079 | |
| 1080 | if self.is_group_leader: |
| 1081 | self._resolve_and_stage_group( |
| 1082 | resident=resident, |
| 1083 | req_id_per_token=req_id_per_token, |
| 1084 | block_table=block_table, |
| 1085 | logical_topk_indices=logical_topk_indices, |
| 1086 | block_size=block_size, |
| 1087 | return_valid_counts=return_valid_counts, |
| 1088 | shared_rows=shared_rows, |
| 1089 | attention_indices_out=attention_indices_out, |
| 1090 | valid_counts_out=valid_counts_out, |
| 1091 | num_valid_rows=num_valid_rows, |
| 1092 | ) |
| 1093 | |
| 1094 | if not self._swap_staged: |
| 1095 | self._stage_this_layer(shared_rows) |
| 1096 | self.wait_for_staged_swap() |
| 1097 | |
| 1098 | physical_topk_indices = group.shared_topk.physical_topk_indices[shared_rows] |
| 1099 | if return_valid_counts: |
| 1100 | return ( |
| 1101 | physical_topk_indices, |
| 1102 | group.shared_topk.valid_topk_counts[shared_rows], |
| 1103 | ) |
| 1104 | return physical_topk_indices |
| 1105 | |
| 1106 | |
| 1107 | class HiSparseCacheHandle: |
| 1108 | """Attention-facing handle for resident KV and sparse offload state.""" |
| 1109 | |
| 1110 | def __init__(self, runtime: HiSparseRuntime) -> None: |
| 1111 | self.view: PagedCacheView | None = None |
| 1112 | self.block_table: torch.Tensor | None = None |
| 1113 | self.source_block_table: torch.Tensor | None = None |
| 1114 | self.slot_mapping: torch.Tensor | None = None |
| 1115 | self.runtime = runtime |
| 1116 | self.dummy_batch = False |
| 1117 | self.decode_batch = False |
| 1118 | self.all_context_pages_resident = True |
| 1119 | self.num_actual_tokens = 0 |
| 1120 | self.num_decode_tokens = 0 |
| 1121 | self.req_id_per_token: torch.Tensor | None = None |
| 1122 | self.host_mirror_required = False |
| 1123 | self.mirror_from_resident = False |
| 1124 | self.mirror_slot_mapping: torch.Tensor | None = None |
| 1125 | self.mirror_staging_cache: torch.Tensor | None = None |
| 1126 | self.mirror_staging_slots: torch.Tensor | None = None |
| 1127 | self.submit_layer_mirror: Callable[[], None] | None = None |
| 1128 | # Speculator layers write their rows after the target forward. |
| 1129 | self.draft_layer = False |
| 1130 | self.index_group_caches: list[HiSparseCacheHandle] = [self] |
| 1131 | |
| 1132 | def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: |
| 1133 | assert self.runtime.is_group_leader |
| 1134 | for cache in self.index_group_caches: |
| 1135 | cache._prepare_for_batch(attn_metadata) |
| 1136 | |
| 1137 | def _prepare_for_batch(self, attn_metadata: Any | None) -> None: |
| 1138 | self.dummy_batch = attn_metadata is None |
| 1139 | self.runtime.begin_forward() |
| 1140 | self.num_actual_tokens = ( |
| 1141 | attn_metadata.num_actual_tokens if attn_metadata is not None else 0 |
| 1142 | ) |
| 1143 | self.num_decode_tokens = ( |
| 1144 | attn_metadata.num_decode_tokens if attn_metadata is not None else 0 |
| 1145 | ) |
| 1146 | self.req_id_per_token = ( |
| 1147 | attn_metadata.req_id_per_token if attn_metadata is not None else None |
| 1148 | ) |
| 1149 | self.decode_batch = is_hisparse_decode_batch(attn_metadata) |
| 1150 | self.host_mirror_required = attn_metadata is not None and ( |
| 1151 | not self.decode_batch or self.runtime.eager_host_mirror |
| 1152 | ) |
| 1153 | |
| 1154 | def write_target( |
| 1155 | self, num_input_rows: int, num_slot_rows: int |
| 1156 | ) -> tuple[torch.Tensor, torch.Tensor, int]: |
| 1157 | assert self.view is not None and self.slot_mapping is not None |
| 1158 | num_rows = min(num_input_rows, num_slot_rows, self.num_actual_tokens) |
| 1159 | if self.decode_batch: |
| 1160 | num_rows = min(num_rows, self.runtime.max_num_reqs) |
| 1161 | return self.view.cache, self.slot_mapping[:num_rows], num_rows |
| 1162 | |
| 1163 | def mirror_write_target( |
| 1164 | self, num_rows: int |
| 1165 | ) -> tuple[torch.Tensor, torch.Tensor] | None: |
| 1166 | if ( |
| 1167 | not self.host_mirror_required |
| 1168 | or self.decode_batch |
| 1169 | or self.mirror_from_resident |
| 1170 | ): |
| 1171 | return None |
| 1172 | cache = self.mirror_staging_cache |
| 1173 | slots = self.mirror_staging_slots |
| 1174 | if cache is None or slots is None: |
| 1175 | raise RuntimeError("HiSparse prefill mirror staging is not bound.") |
| 1176 | if num_rows > slots.numel(): |
| 1177 | raise RuntimeError( |
| 1178 | "HiSparse prefill mirror exceeds staging capacity: " |
| 1179 | f"{num_rows} > {slots.numel()}." |
| 1180 | ) |
| 1181 | return cache, slots[:num_rows] |
| 1182 | |
| 1183 | def finish_kv_update(self) -> None: |
| 1184 | if self.dummy_batch: |
| 1185 | return |
| 1186 | if ( |
| 1187 | self.submit_layer_mirror is not None |
| 1188 | and not self.decode_batch |
| 1189 | and get_forward_context().cudagraph_runtime_mode != CUDAGraphMode.FULL |
| 1190 | ): |
| 1191 | self.submit_layer_mirror() |
| 1192 | |
| 1193 | def bind_cache( |
| 1194 | self, |
| 1195 | raw_tensor: torch.Tensor, |
| 1196 | *, |
| 1197 | byte_offset: int, |
| 1198 | block_stride: int, |
| 1199 | num_blocks: int, |
| 1200 | block_size: int, |
| 1201 | block_table: torch.Tensor, |
| 1202 | slot_mapping: torch.Tensor, |
| 1203 | ) -> None: |
| 1204 | self.view = PagedCacheView.bind( |
| 1205 | raw_tensor, |
| 1206 | dtype=self.runtime.kv_dtype, |
| 1207 | row_width=self.runtime.row_width, |
| 1208 | byte_offset=byte_offset, |
| 1209 | block_stride=block_stride, |
| 1210 | num_blocks=num_blocks, |
| 1211 | block_size=block_size, |
| 1212 | ) |
| 1213 | self.block_table = block_table |
| 1214 | self.slot_mapping = slot_mapping |
| 1215 | |
| 1216 | def swap_in( |
| 1217 | self, |
| 1218 | req_id_per_token: torch.Tensor, |
| 1219 | block_table: torch.Tensor, |
| 1220 | logical_topk_indices: torch.Tensor, |
| 1221 | *, |
| 1222 | block_size: int, |
| 1223 | num_valid_rows: torch.Tensor, |
| 1224 | return_valid_counts: bool = False, |
| 1225 | attention_indices_out: torch.Tensor | None = None, |
| 1226 | valid_counts_out: torch.Tensor | None = None, |
| 1227 | ) -> HiSparseTopKResult: |
| 1228 | return self.runtime.swap_in( |
| 1229 | resident=self, |
| 1230 | req_id_per_token=req_id_per_token[: logical_topk_indices.shape[0]], |
| 1231 | block_table=block_table, |
| 1232 | logical_topk_indices=logical_topk_indices, |
| 1233 | block_size=block_size, |
| 1234 | return_valid_counts=return_valid_counts, |
| 1235 | attention_indices_out=attention_indices_out, |
| 1236 | valid_counts_out=valid_counts_out, |
| 1237 | num_valid_rows=num_valid_rows, |
| 1238 | ) |
| 1239 | |
| 1240 | |
| 1241 | def initialize_hisparse_runtime_buffers( |
| 1242 | cache_handles: list[HiSparseCacheHandle], |
| 1243 | *, |
| 1244 | max_num_reqs: int, |
| 1245 | max_num_batched_tokens: int, |
| 1246 | ) -> None: |
| 1247 | """Allocate shared request state and per-layer prefill staging buffers.""" |
| 1248 | assert cache_handles |
| 1249 | resident = cache_handles[0].view |
| 1250 | assert resident is not None |
| 1251 | device = resident.cache.device |
| 1252 | request_state_indices = torch.full( |
| 1253 | (max_num_reqs,), -1, dtype=torch.int32, device=device |
| 1254 | ) |
| 1255 | staging_blocks = ( |
| 1256 | max_num_batched_tokens + resident.block_size - 1 |
| 1257 | ) // resident.block_size |
| 1258 | mirror_staging_caches = torch.empty( |
| 1259 | ( |
| 1260 | len(cache_handles), |
| 1261 | staging_blocks, |
| 1262 | resident.block_size, |
| 1263 | resident.cache.shape[-1], |
| 1264 | ), |
| 1265 | dtype=resident.cache.dtype, |
| 1266 | device=device, |
| 1267 | ) |
| 1268 | mirror_staging_slots = torch.arange( |
| 1269 | max_num_batched_tokens, dtype=torch.int64, device=device |
| 1270 | ) |
| 1271 | for layer_index, cache_handle in enumerate(cache_handles): |
| 1272 | cache_handle.runtime.request_state_indices = request_state_indices |
| 1273 | cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index] |
| 1274 | cache_handle.mirror_staging_slots = mirror_staging_slots |
| 1275 | |
| 1276 | |
| 1277 | def create_hisparse_cache_handle( |
| 1278 | vllm_config: VllmConfig, |
| 1279 | model_top_k: int, |
| 1280 | *, |
| 1281 | is_index_group_leader: bool, |
| 1282 | row_width: int, |
| 1283 | kv_dtype: torch.dtype, |
| 1284 | index_group: HiSparseMLAIndexGroup | None = None, |
| 1285 | device: torch.device | str | None = None, |
| 1286 | ) -> HiSparseCacheHandle | None: |
| 1287 | config = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k) |
| 1288 | if config is None: |
| 1289 | return None |
| 1290 | |
| 1291 | max_num_reqs = vllm_config.scheduler_config.max_num_seqs |
| 1292 | max_swap_rows = _get_max_swap_rows(vllm_config) |
| 1293 | if device is None: |
| 1294 | device = torch.device( |
| 1295 | current_platform.device_type, torch.accelerator.current_device_index() |
| 1296 | ) |
| 1297 | |
| 1298 | hisparse_group = ( |
| 1299 | None |
| 1300 | if is_index_group_leader or index_group is None |
| 1301 | else index_group.hisparse_group |
| 1302 | ) |
| 1303 | runtime = HiSparseRuntime( |
| 1304 | config=config, |
| 1305 | max_num_reqs=max_num_reqs, |
| 1306 | max_swap_rows=max_swap_rows, |
| 1307 | row_width=row_width, |
| 1308 | kv_dtype=kv_dtype, |
| 1309 | device=device, |
| 1310 | index_group=hisparse_group, |
| 1311 | copy_stream=index_group.side_stream if index_group is not None else None, |
| 1312 | logical_topk_ready=( |
| 1313 | index_group.logical_topk_ready |
| 1314 | if index_group is not None and index_group.has_indexer |
| 1315 | else None |
| 1316 | ), |
| 1317 | ) |
| 1318 | if is_index_group_leader and index_group is not None: |
| 1319 | index_group.hisparse_group = runtime.index_group |
| 1320 | logger.info_once( |
| 1321 | "Enabled experimental HiSparse HMA hot cache: top_k=%d, " |
| 1322 | "device_buffer_size=%d (%d LRU rows), max_num_seqs=%d.", |
| 1323 | config.top_k, |
| 1324 | config.device_buffer_size, |
| 1325 | config.device_buffer_size, |
| 1326 | max_num_reqs, |
| 1327 | ) |
| 1328 | handle = HiSparseCacheHandle(runtime) |
| 1329 | speculative_config = vllm_config.speculative_config |
| 1330 | handle.mirror_from_resident = bool( |
| 1331 | vllm_config.scheduler_config.async_scheduling |
| 1332 | and speculative_config is not None |
| 1333 | and speculative_config.uses_draft_kv_cache() |
| 1334 | ) |
| 1335 | return handle |