| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Iterable, Mapping, Sequence |
| 5 | from dataclasses import dataclass, field |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | from vllm.utils.math_utils import cdiv |
| 9 | from vllm.v1.core.block_pool import BlockPool |
| 10 | from vllm.v1.core.kv_cache_utils import ( |
| 11 | BlockHashWithGroupId, |
| 12 | KVCacheBlock, |
| 13 | KVCacheBlockCopy, |
| 14 | ) |
| 15 | from vllm.v1.core.single_type_kv_cache_manager import ( |
| 16 | HiSparseHotManager, |
| 17 | HiSparseResidentManager, |
| 18 | HiSparseSourceManager, |
| 19 | SingleTypeKVCacheManager, |
| 20 | ) |
| 21 | from vllm.v1.hisparse.types import ( |
| 22 | SparseKVOffloadCommand, |
| 23 | SparseKVPageTransfer, |
| 24 | SparseKVRowMirror, |
| 25 | ) |
| 26 | from vllm.v1.kv_cache_interface import ( |
| 27 | HiSparseResidentSpec, |
| 28 | KVCacheConfig, |
| 29 | ) |
| 30 | from vllm.v1.request import Request |
| 31 | |
| 32 | if TYPE_CHECKING: |
| 33 | from vllm.v1.core.kv_cache_manager import KVCacheManager |
| 34 | |
| 35 | # Sealed pages this many positions behind the block-table tail stay pinned so |
| 36 | # a page written by an in-flight step is never handed out under it. |
| 37 | _ACTIVE_TAIL_PAGES = 2 |
| 38 | |
| 39 | |
| 40 | @dataclass |
| 41 | class _PendingPublication: |
| 42 | request: Request |
| 43 | num_computed_tokens: int |
| 44 | num_pages: int |
| 45 | retention_interval: int | None |
| 46 | replay_boundaries: Sequence[int] |
| 47 | detached_blocks: list[KVCacheBlock] | None = None |
| 48 | num_cached_blocks: int = 0 |
| 49 | |
| 50 | |
| 51 | @dataclass |
| 52 | class _HiSparseRequestState: |
| 53 | """Residency bookkeeping for one request. |
| 54 | |
| 55 | A resident page moves through three states: dirty (only the GPU copy |
| 56 | exists), clean (``valid_pages``: the host copy is complete) and unpinned |
| 57 | (``unpinned_pages``: clean, and its allocation reference was released so |
| 58 | the pool may reuse it). ``pinned_clean`` tracks clean pages still holding |
| 59 | their reference, which only happens before the request can read from host. |
| 60 | """ |
| 61 | |
| 62 | valid_pages: set[int] = field(default_factory=set) |
| 63 | ready_prefix_pages: int = 0 |
| 64 | pending_pages: dict[int, int] = field(default_factory=dict) |
| 65 | publication: _PendingPublication | None = None |
| 66 | copies_recorded_blocks: int = 0 |
| 67 | pinned_clean: set[int] = field(default_factory=set) |
| 68 | unpinned_pages: set[int] = field(default_factory=set) |
| 69 | # Prefix pages whose GPU copies are adopted after the admitting allocation. |
| 70 | pages_to_adopt: int = 0 |
| 71 | |
| 72 | |
| 73 | @dataclass |
| 74 | class _PendingSpill: |
| 75 | transfer_id: int |
| 76 | page: tuple[str, int] |
| 77 | request_state: _HiSparseRequestState |
| 78 | host_block: KVCacheBlock |
| 79 | resident_blocks: tuple[KVCacheBlock, ...] |
| 80 | restore: bool = False |
| 81 | expected_worker_completions: int = 0 |
| 82 | worker_completions: int = 0 |
| 83 | enqueue_applied: bool = False |
| 84 | |
| 85 | |
| 86 | class HiSparseCoordinator: |
| 87 | """Own HiSparse host allocation, host publication, spills, and GPU residency. |
| 88 | |
| 89 | GPU-resident pages are a write-back cache of the host tier. Once a page's |
| 90 | host copy is durable and its owner can read from host (it has, or has |
| 91 | asked for, a hot region), the page's allocation reference is released with |
| 92 | ``BlockPool.unpin_blocks``: the block stays in the block table and keeps |
| 93 | being read, but the pool counts it as free and may hand it out, at which |
| 94 | point the owner's page is nulled and its block table republished. |
| 95 | """ |
| 96 | |
| 97 | def __init__( |
| 98 | self, |
| 99 | kv_cache_config: KVCacheConfig, |
| 100 | managers: tuple[SingleTypeKVCacheManager, ...], |
| 101 | max_model_len: int, |
| 102 | *, |
| 103 | num_reprefillable_tokens: int, |
| 104 | ) -> None: |
| 105 | self.managers = managers |
| 106 | self.max_model_len = max_model_len |
| 107 | self.num_reprefillable_tokens = num_reprefillable_tokens |
| 108 | groups = kv_cache_config.kv_cache_groups |
| 109 | |
| 110 | resident_managers: list[HiSparseResidentManager] = [] |
| 111 | hot_managers: list[HiSparseHotManager] = [] |
| 112 | resident_group_ids: list[int] = [] |
| 113 | for group_id, (group, manager) in enumerate(zip(groups, managers)): |
| 114 | if isinstance(group.kv_cache_spec, HiSparseResidentSpec): |
| 115 | if not isinstance(manager, HiSparseResidentManager): |
| 116 | raise TypeError( |
| 117 | "HiSparse resident specs require resident cache managers." |
| 118 | ) |
| 119 | resident_managers.append(manager) |
| 120 | resident_group_ids.append(group_id) |
| 121 | assert not group.host_resident |
| 122 | if isinstance(manager, HiSparseHotManager): |
| 123 | hot_managers.append(manager) |
| 124 | assert not group.host_resident |
| 125 | if isinstance(manager, (HiSparseHotManager, HiSparseResidentManager)): |
| 126 | manager.coordinator = self |
| 127 | self.resident_managers = tuple(resident_managers) |
| 128 | self.hot_managers = tuple(hot_managers) |
| 129 | self.host_manager: HiSparseSourceManager | None = None |
| 130 | self.host_group_id: int | None = None |
| 131 | self.max_spill_pages = 0 |
| 132 | for group_id, manager in enumerate(managers): |
| 133 | if isinstance(manager, HiSparseSourceManager): |
| 134 | if self.host_manager is not None: |
| 135 | raise ValueError("Only one HiSparse host group is supported.") |
| 136 | num_host_blocks = kv_cache_config.hisparse_host_num_blocks |
| 137 | if num_host_blocks is None: |
| 138 | raise ValueError("HiSparse host group needs host capacity.") |
| 139 | self.host_group_id = group_id |
| 140 | self.host_manager = manager |
| 141 | manager.coordinator = self |
| 142 | manager.bind_host_pool(num_host_blocks) |
| 143 | self.has_host_cache = self.host_manager is not None |
| 144 | self.gpu_pool: BlockPool | None = None |
| 145 | self.transition_watermark = 0 |
| 146 | if self.resident_managers: |
| 147 | assert self.host_manager is not None |
| 148 | assert self.host_group_id is not None |
| 149 | if self.host_group_id > min(resident_group_ids): |
| 150 | raise ValueError("HiSparse host group must precede resident groups.") |
| 151 | resident_block_sizes = { |
| 152 | manager.block_size for manager in self.resident_managers |
| 153 | } |
| 154 | if len(resident_block_sizes) != 1: |
| 155 | raise ValueError( |
| 156 | "HiSparse resident cache groups must use one block size." |
| 157 | ) |
| 158 | resident_block_size = resident_block_sizes.pop() |
| 159 | if self.host_manager.block_size != resident_block_size: |
| 160 | raise ValueError("HiSparse host and resident block sizes must match.") |
| 161 | self.max_spill_pages = max_model_len // resident_block_size |
| 162 | self.gpu_pool = self.resident_managers[0].block_pool |
| 163 | # Requests start reading from host once the shared pool runs this |
| 164 | # low, so admissions find evictable pages rather than pinned ones. |
| 165 | hot_cost = sum(manager.blocks_per_request for manager in hot_managers) |
| 166 | self.transition_watermark = max(hot_cost, kv_cache_config.num_blocks // 10) |
| 167 | |
| 168 | # Host copies handed to the worker, by destination block id, until it |
| 169 | # reports having run them. |
| 170 | self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {} |
| 171 | # request -> prefix length an external load is still filling in. |
| 172 | self._pending_imports: dict[str, int] = {} |
| 173 | self.block_table_updates: set[str] = set() |
| 174 | self.spills_to_send: list[SparseKVPageTransfer] = [] |
| 175 | self.pending_spills: dict[int, _PendingSpill] = {} |
| 176 | self.request_states: dict[str, _HiSparseRequestState] = {} |
| 177 | self.next_spill_id = 0 |
| 178 | # Published host block hash -> GPU copies of that page, readable until |
| 179 | # the pool reuses them. Lets a host prefix hit come back GPU-resident. |
| 180 | self.copies: dict[BlockHashWithGroupId, tuple[KVCacheBlock, ...]] = {} |
| 181 | self._copy_by_block: dict[int, BlockHashWithGroupId] = {} |
| 182 | # Unpinned GPU block id -> (request, page) pairs still reading it. |
| 183 | self._owners: dict[int, set[tuple[str, int]]] = {} |
| 184 | |
| 185 | def _get_request_state(self, request_id: str) -> _HiSparseRequestState: |
| 186 | state = self.request_states.get(request_id) |
| 187 | if state is None: |
| 188 | state = _HiSparseRequestState() |
| 189 | self.request_states[request_id] = state |
| 190 | return state |
| 191 | |
| 192 | def commit_computed_blocks(self, request_id: str, num_host_pages: int) -> None: |
| 193 | """Account for a prefix hit on host pages the request now references.""" |
| 194 | if not self.resident_managers or self.host_manager is None: |
| 195 | return |
| 196 | host_blocks = self.host_manager.req_to_blocks.get(request_id, ()) |
| 197 | num_host_pages = min(num_host_pages, len(host_blocks)) |
| 198 | if num_host_pages == 0: |
| 199 | return |
| 200 | state = self._get_request_state(request_id) |
| 201 | state.valid_pages.update(range(num_host_pages)) |
| 202 | state.ready_prefix_pages = max(num_host_pages, state.ready_prefix_pages) |
| 203 | # Adopting pins copies out of the free pool, so it must wait until the |
| 204 | # allocation admission counted on has run (see update_residency). |
| 205 | state.pages_to_adopt = max(state.pages_to_adopt, num_host_pages) |
| 206 | state.copies_recorded_blocks = max(state.copies_recorded_blocks, num_host_pages) |
| 207 | |
| 208 | def take_host_block_copies(self) -> tuple[KVCacheBlockCopy, ...]: |
| 209 | """Drain host copy-on-write work, retaining both endpoints. |
| 210 | |
| 211 | The worker runs the copies in the step that carries them, so the |
| 212 | endpoints are released when that step's output comes back. |
| 213 | """ |
| 214 | if self.host_manager is None: |
| 215 | return () |
| 216 | pairs = self.host_manager.take_host_cow_copies() |
| 217 | for source_block, cow_block in pairs: |
| 218 | self._retained_copies[cow_block.block_id] = (source_block, cow_block) |
| 219 | return tuple( |
| 220 | KVCacheBlockCopy( |
| 221 | src_block_id=source_block.block_id, |
| 222 | dst_block_id=cow_block.block_id, |
| 223 | ) |
| 224 | for source_block, cow_block in pairs |
| 225 | ) |
| 226 | |
| 227 | def release_completed_host_copies(self, dst_block_ids: Iterable[int]) -> None: |
| 228 | """Release the endpoints of host copies the worker reports having run.""" |
| 229 | blocks: list[KVCacheBlock] = [] |
| 230 | for dst_block_id in dst_block_ids: |
| 231 | pair = self._retained_copies.pop(dst_block_id, None) |
| 232 | if pair is not None: |
| 233 | blocks.extend(pair) |
| 234 | if not blocks: |
| 235 | return |
| 236 | assert self.host_manager is not None |
| 237 | self.host_manager.block_pool.free_blocks(reversed(blocks)) |
| 238 | |
| 239 | def get_host_block_pool(self) -> BlockPool | None: |
| 240 | manager = self.host_manager |
| 241 | return manager.block_pool if manager is not None else None |
| 242 | |
| 243 | # ------------------------------------------------------------------ |
| 244 | # GPU copies of published host pages |
| 245 | # ------------------------------------------------------------------ |
| 246 | |
| 247 | def _adopt_copies( |
| 248 | self, |
| 249 | request_id: str, |
| 250 | state: _HiSparseRequestState, |
| 251 | host_blocks: Sequence[KVCacheBlock], |
| 252 | ) -> None: |
| 253 | """Point null prefix pages at readable GPU copies of their contents. |
| 254 | |
| 255 | Runs after the hit length is fully reconciled, so copy hits and misses |
| 256 | can be scattered per page without affecting the prefix hit. |
| 257 | """ |
| 258 | if not self.copies: |
| 259 | return |
| 260 | for host_idx, host_block in enumerate(host_blocks): |
| 261 | if host_block.is_null or host_block.block_hash is None: |
| 262 | continue |
| 263 | blocks = self.copies.get(host_block.block_hash) |
| 264 | if blocks is None: |
| 265 | continue |
| 266 | for manager, block in zip(self.resident_managers, blocks): |
| 267 | if manager.adopt_resident_page(request_id, host_idx, block): |
| 268 | manager.block_pool.touch([block]) |
| 269 | state.pinned_clean.add(host_idx) |
| 270 | self.block_table_updates.add(request_id) |
| 271 | |
| 272 | def _record_copies(self, request_id: str, num_computed_tokens: int) -> None: |
| 273 | """Index the GPU copies of just-published host blocks for later hits.""" |
| 274 | if not self.resident_managers: |
| 275 | return |
| 276 | assert self.host_manager is not None |
| 277 | host_blocks = self.host_manager.req_to_blocks.get(request_id) |
| 278 | if not host_blocks: |
| 279 | return |
| 280 | state = self._get_request_state(request_id) |
| 281 | num_host_blocks = min( |
| 282 | num_computed_tokens // self.host_manager.block_size, len(host_blocks) |
| 283 | ) |
| 284 | for host_idx in range(state.copies_recorded_blocks, num_host_blocks): |
| 285 | if host_idx in state.pending_pages: |
| 286 | continue |
| 287 | self._record_copy(request_id, host_idx, host_blocks[host_idx]) |
| 288 | state.copies_recorded_blocks = max( |
| 289 | state.copies_recorded_blocks, num_host_blocks |
| 290 | ) |
| 291 | |
| 292 | def _record_copy( |
| 293 | self, request_id: str, page_idx: int, host_block: KVCacheBlock |
| 294 | ) -> None: |
| 295 | if host_block.is_null or host_block.block_hash is None: |
| 296 | return |
| 297 | blocks: list[KVCacheBlock] = [] |
| 298 | for manager in self.resident_managers: |
| 299 | block = manager.get_resident_page(request_id, page_idx) |
| 300 | if block is None: |
| 301 | return |
| 302 | blocks.append(block) |
| 303 | self._drop_copy(host_block.block_hash) |
| 304 | self.copies[host_block.block_hash] = tuple(blocks) |
| 305 | for block in blocks: |
| 306 | self._copy_by_block[block.block_id] = host_block.block_hash |
| 307 | |
| 308 | def _drop_copy(self, host_hash: BlockHashWithGroupId) -> None: |
| 309 | blocks = self.copies.pop(host_hash, None) |
| 310 | if blocks is None: |
| 311 | return |
| 312 | for block in blocks: |
| 313 | self._copy_by_block.pop(block.block_id, None) |
| 314 | |
| 315 | # ------------------------------------------------------------------ |
| 316 | # Residency: pinning, unpinning and losing pages |
| 317 | # ------------------------------------------------------------------ |
| 318 | |
| 319 | def _can_read_from_host(self, request_id: str) -> bool: |
| 320 | """Whether resident pages may be dropped under the request.""" |
| 321 | return bool(self.hot_managers) and all( |
| 322 | manager.has_hot(request_id) for manager in self.hot_managers |
| 323 | ) |
| 324 | |
| 325 | def _fills_admission_window(self, request_id: str) -> bool: |
| 326 | """Whether the request holds the resident pages admission reserved. |
| 327 | |
| 328 | Admission caps a request's resident pages at the in-flight window, |
| 329 | assuming older pages move to host; keeping them pinned past it lets a |
| 330 | chunked prefill outgrow the pool it was admitted into. |
| 331 | """ |
| 332 | return any( |
| 333 | len(manager.req_to_blocks.get(request_id, ())) |
| 334 | >= manager.max_admission_blocks_per_request |
| 335 | for manager in self.resident_managers |
| 336 | ) |
| 337 | |
| 338 | def _resident_page_blocks( |
| 339 | self, request_id: str, page_idx: int |
| 340 | ) -> list[KVCacheBlock] | None: |
| 341 | blocks: list[KVCacheBlock] = [] |
| 342 | for manager in self.resident_managers: |
| 343 | req_blocks = manager.req_to_blocks.get(request_id) |
| 344 | if ( |
| 345 | req_blocks is None |
| 346 | or page_idx >= len(req_blocks) - _ACTIVE_TAIL_PAGES |
| 347 | or req_blocks[page_idx].is_null |
| 348 | ): |
| 349 | return None |
| 350 | blocks.append(req_blocks[page_idx]) |
| 351 | return blocks |
| 352 | |
| 353 | def _unpin_page( |
| 354 | self, request_id: str, state: _HiSparseRequestState, page_idx: int |
| 355 | ) -> bool: |
| 356 | if page_idx in state.pending_pages or page_idx not in state.valid_pages: |
| 357 | return False |
| 358 | blocks = self._resident_page_blocks(request_id, page_idx) |
| 359 | if blocks is None: |
| 360 | return False |
| 361 | for manager, block in zip(self.resident_managers, blocks): |
| 362 | self._owners.setdefault(block.block_id, set()).add((request_id, page_idx)) |
| 363 | manager.block_pool.unpin_blocks([block], self._on_block_reused) |
| 364 | state.pinned_clean.discard(page_idx) |
| 365 | state.unpinned_pages.add(page_idx) |
| 366 | return True |
| 367 | |
| 368 | def _unpin_clean_pages(self, request_id: str, state: _HiSparseRequestState) -> None: |
| 369 | # Tail first, so a prefix loses its tail pages before its head. |
| 370 | for page_idx in sorted(state.pinned_clean, reverse=True): |
| 371 | self._unpin_page(request_id, state, page_idx) |
| 372 | |
| 373 | def _page_became_clean( |
| 374 | self, request_id: str, state: _HiSparseRequestState, page_idx: int |
| 375 | ) -> None: |
| 376 | state.pinned_clean.add(page_idx) |
| 377 | if self._can_read_from_host(request_id): |
| 378 | self._unpin_page(request_id, state, page_idx) |
| 379 | |
| 380 | def _on_block_reused(self, block: KVCacheBlock) -> None: |
| 381 | host_hash = self._copy_by_block.get(block.block_id) |
| 382 | if host_hash is not None: |
| 383 | self._drop_copy(host_hash) |
| 384 | for request_id, page_idx in self._owners.pop(block.block_id, ()): |
| 385 | self._lose_page(request_id, page_idx) |
| 386 | |
| 387 | def _lose_page(self, request_id: str, page_idx: int) -> None: |
| 388 | state = self.request_states.get(request_id) |
| 389 | if state is None: |
| 390 | return |
| 391 | for manager in self.resident_managers: |
| 392 | blocks = manager.req_to_blocks.get(request_id) |
| 393 | if blocks is None or page_idx >= len(blocks) or blocks[page_idx].is_null: |
| 394 | continue |
| 395 | block = blocks[page_idx] |
| 396 | blocks[page_idx] = manager._null_block |
| 397 | owners = self._owners.get(block.block_id) |
| 398 | if owners is not None: |
| 399 | owners.discard((request_id, page_idx)) |
| 400 | state.unpinned_pages.discard(page_idx) |
| 401 | self.block_table_updates.add(request_id) |
| 402 | for hot_manager in self.hot_managers: |
| 403 | hot_manager.require_hot(request_id) |
| 404 | |
| 405 | def update_residency(self, request_id: str) -> None: |
| 406 | """Per-step residency policy for a scheduled request. |
| 407 | |
| 408 | A request that can read from host releases every clean sealed page to |
| 409 | the pool. One that cannot keeps its pages pinned until the shared pool |
| 410 | runs low, or until it holds the resident pages admission reserved for |
| 411 | it, then asks for a hot region. Pages remain pinned until that region is |
| 412 | allocated on a subsequent scheduling pass. |
| 413 | """ |
| 414 | if not self.resident_managers: |
| 415 | return |
| 416 | state = self._get_request_state(request_id) |
| 417 | if state.pages_to_adopt: |
| 418 | assert self.host_manager is not None |
| 419 | host_blocks = self.host_manager.req_to_blocks.get(request_id, ()) |
| 420 | self._adopt_copies(request_id, state, host_blocks[: state.pages_to_adopt]) |
| 421 | state.pages_to_adopt = 0 |
| 422 | if not self._can_read_from_host(request_id): |
| 423 | assert self.gpu_pool is not None |
| 424 | if ( |
| 425 | self.gpu_pool.get_num_free_blocks() >= self.transition_watermark |
| 426 | and not self._fills_admission_window(request_id) |
| 427 | ): |
| 428 | return |
| 429 | for manager in self.hot_managers: |
| 430 | manager.require_hot(request_id) |
| 431 | return |
| 432 | self._unpin_clean_pages(request_id, state) |
| 433 | |
| 434 | # ------------------------------------------------------------------ |
| 435 | # Host publication and spills |
| 436 | # ------------------------------------------------------------------ |
| 437 | |
| 438 | def advance_scheduled(self, requests: Iterable[tuple[str, int]]) -> None: |
| 439 | """Run the per-step residency work for each scheduled request. |
| 440 | |
| 441 | ``requests`` pairs a request id with the number of tokens computed |
| 442 | once the step completes. As in ``KVCacheCoordinator.cache_blocks``, the |
| 443 | last ``num_reprefillable_tokens`` are not final yet and are not written |
| 444 | back. |
| 445 | """ |
| 446 | if not self.resident_managers: |
| 447 | return |
| 448 | for request_id, num_tokens in requests: |
| 449 | self.plan_prefix_materialization( |
| 450 | request_id, max(0, num_tokens - self.num_reprefillable_tokens) |
| 451 | ) |
| 452 | self.update_residency(request_id) |
| 453 | |
| 454 | def plan_prefix_materialization( |
| 455 | self, request_id: str, num_computed_tokens: int |
| 456 | ) -> None: |
| 457 | """Eagerly write back every sealed page that has no host copy yet.""" |
| 458 | if not self.resident_managers: |
| 459 | return |
| 460 | assert self.host_manager is not None |
| 461 | host_block_size = self.host_manager.block_size |
| 462 | num_pages = num_computed_tokens // host_block_size |
| 463 | importing_pages = cdiv( |
| 464 | self._pending_imports.get(request_id, 0), host_block_size |
| 465 | ) |
| 466 | state = self._get_request_state(request_id) |
| 467 | budget = max(self.max_spill_pages - len(self.spills_to_send), 0) |
| 468 | first_page = max(state.ready_prefix_pages, importing_pages) |
| 469 | for page_idx in range(first_page, num_pages): |
| 470 | if budget == 0: |
| 471 | break |
| 472 | if page_idx in state.pending_pages: |
| 473 | continue |
| 474 | if page_idx in state.valid_pages: |
| 475 | continue |
| 476 | if not all( |
| 477 | manager.get_resident_page(request_id, page_idx) is not None |
| 478 | for manager in self.resident_managers |
| 479 | ): |
| 480 | continue |
| 481 | if self._plan_page_transfer(request_id, page_idx, after_forward=True): |
| 482 | budget -= 1 |
| 483 | |
| 484 | def publish_when_ready( |
| 485 | self, |
| 486 | request: Request, |
| 487 | num_computed_tokens: int, |
| 488 | retention_interval: int | None, |
| 489 | *, |
| 490 | replay_boundaries: Sequence[int], |
| 491 | ) -> None: |
| 492 | """Publish host-source hashes of the pages that are already durable.""" |
| 493 | manager = self.host_manager |
| 494 | if manager is None: |
| 495 | return |
| 496 | request_id = request.request_id |
| 497 | self._get_request_state(request_id).publication = _PendingPublication( |
| 498 | request=request, |
| 499 | num_computed_tokens=num_computed_tokens, |
| 500 | num_pages=num_computed_tokens // manager.block_size, |
| 501 | retention_interval=retention_interval, |
| 502 | replay_boundaries=replay_boundaries, |
| 503 | ) |
| 504 | self._publish_host_blocks_if_ready(request_id) |
| 505 | |
| 506 | def _publish_host_blocks_if_ready(self, request_id: str) -> None: |
| 507 | """Publish the durable prefix of the pending publication. |
| 508 | |
| 509 | Host writes trail a long prefill by about a chunk, so waiting for the |
| 510 | whole computed prefix would publish nothing until the prefill ends, and |
| 511 | a prefill preempted before then would recompute from the start. |
| 512 | """ |
| 513 | state = self.request_states.get(request_id) |
| 514 | if state is None: |
| 515 | return |
| 516 | publication = state.publication |
| 517 | if publication is None: |
| 518 | return |
| 519 | assert self.host_manager is not None |
| 520 | num_tokens = min( |
| 521 | publication.num_computed_tokens, |
| 522 | state.ready_prefix_pages * self.host_manager.block_size, |
| 523 | ) |
| 524 | self.host_manager.publish_blocks( |
| 525 | publication.request, |
| 526 | num_tokens, |
| 527 | retention_interval=publication.retention_interval, |
| 528 | replay_boundaries=publication.replay_boundaries, |
| 529 | ) |
| 530 | self._record_copies(request_id, num_tokens) |
| 531 | if state.ready_prefix_pages >= publication.num_pages: |
| 532 | state.publication = None |
| 533 | |
| 534 | def record_pending_host_import(self, request_id: str, num_tokens: int) -> None: |
| 535 | """Note a prefix an external load is populating in host pages.""" |
| 536 | self._pending_imports[request_id] = num_tokens |
| 537 | |
| 538 | def finish_host_import(self, request_id: str, *, failed: bool) -> None: |
| 539 | num_tokens = self._pending_imports.pop(request_id, None) |
| 540 | if num_tokens is not None and not failed: |
| 541 | self.complete_host_import(request_id, num_tokens) |
| 542 | |
| 543 | def complete_host_import(self, request_id: str, num_computed_tokens: int) -> None: |
| 544 | """Publish externally populated host pages after connector completion.""" |
| 545 | if not self.resident_managers: |
| 546 | return |
| 547 | block_size = self.resident_managers[0].block_size |
| 548 | num_pages = num_computed_tokens // block_size |
| 549 | state = self._get_request_state(request_id) |
| 550 | state.valid_pages = set(range(num_pages)) |
| 551 | state.ready_prefix_pages = num_pages |
| 552 | if num_computed_tokens: |
| 553 | self._plan_page_transfer( |
| 554 | request_id, |
| 555 | (num_computed_tokens - 1) // block_size, |
| 556 | after_forward=False, |
| 557 | restore=True, |
| 558 | ) |
| 559 | self._publish_host_blocks_if_ready(request_id) |
| 560 | |
| 561 | def _plan_page_transfer( |
| 562 | self, |
| 563 | request_id: str, |
| 564 | page_idx: int, |
| 565 | *, |
| 566 | after_forward: bool, |
| 567 | restore: bool = False, |
| 568 | ) -> bool: |
| 569 | assert self.host_manager is not None |
| 570 | state = self._get_request_state(request_id) |
| 571 | if page_idx in state.pending_pages: |
| 572 | return False |
| 573 | host_blocks = self.host_manager.req_to_blocks.get(request_id) |
| 574 | if host_blocks is None or page_idx >= len(host_blocks): |
| 575 | return False |
| 576 | host_block = host_blocks[page_idx] |
| 577 | if host_block.is_null: |
| 578 | return False |
| 579 | |
| 580 | blocks: list[KVCacheBlock] = [] |
| 581 | for manager in self.resident_managers: |
| 582 | block = manager.get_resident_page(request_id, page_idx) |
| 583 | if block is None: |
| 584 | return False |
| 585 | blocks.append(block) |
| 586 | |
| 587 | self.host_manager.block_pool.touch([host_block]) |
| 588 | for manager, block in zip(self.resident_managers, blocks): |
| 589 | manager.block_pool.touch([block]) |
| 590 | spill_id = self.next_spill_id |
| 591 | self.next_spill_id += 1 |
| 592 | plan = SparseKVPageTransfer( |
| 593 | transfer_id=spill_id, |
| 594 | host_block_id=host_block.block_id, |
| 595 | resident_block_ids=tuple(block.block_id for block in blocks), |
| 596 | after_forward=after_forward, |
| 597 | restore=restore, |
| 598 | ) |
| 599 | self.pending_spills[spill_id] = _PendingSpill( |
| 600 | transfer_id=spill_id, |
| 601 | page=(request_id, page_idx), |
| 602 | request_state=state, |
| 603 | host_block=host_block, |
| 604 | resident_blocks=tuple(blocks), |
| 605 | restore=restore, |
| 606 | ) |
| 607 | state.pending_pages[page_idx] = spill_id |
| 608 | self.spills_to_send.append(plan) |
| 609 | return True |
| 610 | |
| 611 | # ------------------------------------------------------------------ |
| 612 | # Worker-facing views |
| 613 | # ------------------------------------------------------------------ |
| 614 | |
| 615 | def build_row_mirrors( |
| 616 | self, |
| 617 | requests: Iterable[tuple[str, int, int]], |
| 618 | ) -> tuple[SparseKVRowMirror, ...]: |
| 619 | """Map scheduled rows to scheduler-owned resident and host blocks.""" |
| 620 | if not self.resident_managers or self.host_manager is None: |
| 621 | return () |
| 622 | block_size = self.resident_managers[0].block_size |
| 623 | mirrors: list[SparseKVRowMirror] = [] |
| 624 | for request_id, num_computed_tokens, num_scheduled_tokens in requests: |
| 625 | host_blocks = self.host_manager.req_to_blocks.get(request_id) |
| 626 | resident_blocks = [ |
| 627 | manager.req_to_blocks.get(request_id) |
| 628 | for manager in self.resident_managers |
| 629 | ] |
| 630 | if host_blocks is None or any(blocks is None for blocks in resident_blocks): |
| 631 | continue |
| 632 | token_position = num_computed_tokens |
| 633 | end_position = token_position + num_scheduled_tokens |
| 634 | while token_position < end_position: |
| 635 | page_idx, row_offset = divmod(token_position, block_size) |
| 636 | num_rows = min(block_size - row_offset, end_position - token_position) |
| 637 | # A page without a GPU copy has no mirror source, and a page |
| 638 | # without a host block has no destination. Skip just that page: |
| 639 | # its neighbours in the window are still mirrorable, and a |
| 640 | # window only ever moves forward, so dropping them here would |
| 641 | # leave their host rows permanently stale. |
| 642 | source_starts = [] |
| 643 | for blocks in resident_blocks: |
| 644 | assert blocks is not None |
| 645 | if page_idx >= len(blocks) or blocks[page_idx].is_null: |
| 646 | break |
| 647 | source_starts.append( |
| 648 | blocks[page_idx].block_id * block_size + row_offset |
| 649 | ) |
| 650 | if len(source_starts) != len(resident_blocks): |
| 651 | token_position += num_rows |
| 652 | continue |
| 653 | host_block_idx = page_idx |
| 654 | if host_block_idx >= len(host_blocks): |
| 655 | token_position += num_rows |
| 656 | continue |
| 657 | host_block = host_blocks[host_block_idx] |
| 658 | if host_block.is_null: |
| 659 | token_position += num_rows |
| 660 | continue |
| 661 | mirrors.append( |
| 662 | SparseKVRowMirror( |
| 663 | source_starts=tuple(source_starts), |
| 664 | destination_start=( |
| 665 | host_block.block_id * block_size + row_offset |
| 666 | ), |
| 667 | num_rows=num_rows, |
| 668 | ) |
| 669 | ) |
| 670 | token_position += num_rows |
| 671 | return tuple(mirrors) |
| 672 | |
| 673 | def all_context_pages_resident( |
| 674 | self, |
| 675 | requests: Iterable[tuple[str, int, int]], |
| 676 | ) -> bool: |
| 677 | """Return whether every scheduled request can read only resident KV.""" |
| 678 | if not self.resident_managers: |
| 679 | return False |
| 680 | block_size = self.resident_managers[0].block_size |
| 681 | for request_id, num_computed_tokens, num_scheduled_tokens in requests: |
| 682 | num_pages = cdiv(num_computed_tokens + num_scheduled_tokens, block_size) |
| 683 | for page_idx in range(num_pages): |
| 684 | if any( |
| 685 | manager.get_resident_page(request_id, page_idx) is None |
| 686 | for manager in self.resident_managers |
| 687 | ): |
| 688 | return False |
| 689 | return True |
| 690 | |
| 691 | def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]: |
| 692 | updates = { |
| 693 | request_id: tuple( |
| 694 | [block.block_id for block in manager.req_to_blocks.get(request_id, [])] |
| 695 | for manager in self.managers |
| 696 | ) |
| 697 | for request_id in self.block_table_updates |
| 698 | } |
| 699 | self.block_table_updates.clear() |
| 700 | return updates |
| 701 | |
| 702 | def build_offload_command(self) -> SparseKVOffloadCommand | None: |
| 703 | """Package residency decisions for the worker connector.""" |
| 704 | if not self.resident_managers: |
| 705 | return None |
| 706 | command = SparseKVOffloadCommand(page_transfers=self.spills_to_send) |
| 707 | self.spills_to_send = [] |
| 708 | return command |
| 709 | |
| 710 | def has_pending_reclamation(self) -> bool: |
| 711 | """Whether an in-flight spill will free GPU or host blocks on completion.""" |
| 712 | return any( |
| 713 | pending.host_block.ref_cnt == 1 |
| 714 | or ( |
| 715 | pending.page[0] in self.request_states |
| 716 | and self._can_read_from_host(pending.page[0]) |
| 717 | ) |
| 718 | for pending in self.pending_spills.values() |
| 719 | ) |
| 720 | |
| 721 | def has_pending_work(self) -> bool: |
| 722 | return bool(self.spills_to_send or self.pending_spills or self._retained_copies) |
| 723 | |
| 724 | def update_spills( |
| 725 | self, |
| 726 | enqueued_counts: Mapping[int, int], |
| 727 | completed_counts: Mapping[int, int], |
| 728 | ) -> None: |
| 729 | for spill_id, count in enqueued_counts.items(): |
| 730 | pending = self.pending_spills.get(spill_id) |
| 731 | if pending is not None and pending.expected_worker_completions == 0: |
| 732 | pending.expected_worker_completions = count |
| 733 | for spill_id, count in completed_counts.items(): |
| 734 | pending = self.pending_spills.get(spill_id) |
| 735 | if pending is not None: |
| 736 | pending.worker_completions += count |
| 737 | |
| 738 | self._apply_enqueued_spills() |
| 739 | self._complete_host_writes() |
| 740 | |
| 741 | def _apply_enqueued_spills(self) -> None: |
| 742 | """Drop the GPU pins once every worker has completed the transfer.""" |
| 743 | for pending in self.pending_spills.values(): |
| 744 | if ( |
| 745 | pending.enqueue_applied |
| 746 | or not pending.expected_worker_completions |
| 747 | or pending.worker_completions < pending.expected_worker_completions |
| 748 | ): |
| 749 | continue |
| 750 | for manager, block in zip(self.resident_managers, pending.resident_blocks): |
| 751 | manager.block_pool.free_blocks([block]) |
| 752 | pending.enqueue_applied = True |
| 753 | |
| 754 | def _complete_host_writes(self) -> None: |
| 755 | completed = [ |
| 756 | pending |
| 757 | for pending in self.pending_spills.values() |
| 758 | if pending.enqueue_applied |
| 759 | and pending.expected_worker_completions |
| 760 | and pending.worker_completions >= pending.expected_worker_completions |
| 761 | ] |
| 762 | completed_request_ids: set[str] = set() |
| 763 | for pending in completed: |
| 764 | self.pending_spills.pop(pending.transfer_id, None) |
| 765 | request_id, page_idx = pending.page |
| 766 | pending.request_state.pending_pages.pop(page_idx, None) |
| 767 | if self.request_states.get(request_id) is pending.request_state: |
| 768 | if not pending.restore: |
| 769 | pending.request_state.valid_pages.add(page_idx) |
| 770 | if page_idx in pending.request_state.valid_pages: |
| 771 | self._page_became_clean(request_id, pending.request_state, page_idx) |
| 772 | if pending.restore: |
| 773 | self._record_copy(request_id, page_idx, pending.host_block) |
| 774 | completed_request_ids.add(request_id) |
| 775 | elif pending.request_state.publication is not None: |
| 776 | if not pending.restore: |
| 777 | pending.request_state.valid_pages.add(page_idx) |
| 778 | self._publish_detached_blocks(pending.request_state) |
| 779 | assert self.host_manager is not None |
| 780 | self.host_manager.block_pool.free_blocks([pending.host_block]) |
| 781 | for request_id in completed_request_ids: |
| 782 | state = self.request_states[request_id] |
| 783 | while state.ready_prefix_pages in state.valid_pages: |
| 784 | state.ready_prefix_pages += 1 |
| 785 | self._publish_host_blocks_if_ready(request_id) |
| 786 | |
| 787 | def _publish_detached_blocks(self, state: _HiSparseRequestState) -> None: |
| 788 | publication = state.publication |
| 789 | if publication is None or state.pending_pages: |
| 790 | return |
| 791 | blocks = publication.detached_blocks |
| 792 | if blocks is None: |
| 793 | return |
| 794 | assert self.host_manager is not None |
| 795 | while state.ready_prefix_pages in state.valid_pages: |
| 796 | state.ready_prefix_pages += 1 |
| 797 | # The host group uses full-attention retention. Only sealed pages with |
| 798 | # completed copies can outlive the request, including on cancellation. |
| 799 | self.host_manager.block_pool.cache_full_blocks( |
| 800 | request=publication.request, |
| 801 | blocks=blocks, |
| 802 | num_cached_blocks=publication.num_cached_blocks, |
| 803 | num_full_blocks=min(state.ready_prefix_pages, publication.num_pages), |
| 804 | block_size=self.host_manager.block_size, |
| 805 | kv_cache_group_id=self.host_manager.kv_cache_group_id, |
| 806 | ) |
| 807 | self.host_manager.block_pool.free_blocks(reversed(blocks)) |
| 808 | state.publication = None |
| 809 | |
| 810 | def free(self, request_id: str) -> None: |
| 811 | """Detach the request; its clean pages stay readable copies in the pool.""" |
| 812 | self._pending_imports.pop(request_id, None) |
| 813 | state = self.request_states.pop(request_id, None) |
| 814 | if state is None: |
| 815 | return |
| 816 | publication = state.publication |
| 817 | if ( |
| 818 | publication is not None |
| 819 | and publication.request.is_finished() |
| 820 | and state.pending_pages |
| 821 | ): |
| 822 | assert self.host_manager is not None |
| 823 | publication.detached_blocks = self.host_manager.req_to_blocks[request_id][ |
| 824 | : publication.num_pages |
| 825 | ] |
| 826 | publication.num_cached_blocks = self.host_manager.num_cached_block.get( |
| 827 | request_id, 0 |
| 828 | ) |
| 829 | # Completed pages need leases too: another allocation may evict |
| 830 | # them before the remaining copies finish and the prefix publishes. |
| 831 | self.host_manager.block_pool.touch(publication.detached_blocks) |
| 832 | else: |
| 833 | # Preemption can cancel a scheduled forward after its spills were |
| 834 | # planned, so its copies do not prove that the KV was computed. |
| 835 | state.publication = None |
| 836 | for manager in self.resident_managers: |
| 837 | blocks = manager.req_to_blocks.get(request_id) |
| 838 | if not blocks: |
| 839 | continue |
| 840 | # Tail first, so a prefix loses its tail pages before its head. |
| 841 | for page_idx in reversed(range(len(blocks))): |
| 842 | block = blocks[page_idx] |
| 843 | if block.is_null: |
| 844 | continue |
| 845 | if page_idx in state.unpinned_pages: |
| 846 | owners = self._owners.get(block.block_id) |
| 847 | if owners is not None: |
| 848 | owners.discard((request_id, page_idx)) |
| 849 | elif ( |
| 850 | page_idx in state.valid_pages |
| 851 | and page_idx not in state.pending_pages |
| 852 | ): |
| 853 | manager.block_pool.unpin_blocks([block], self._on_block_reused) |
| 854 | else: |
| 855 | continue |
| 856 | blocks[page_idx] = manager._null_block |
| 857 | |
| 858 | |
| 859 | def get_hisparse_coordinator( |
| 860 | kv_cache_manager: "KVCacheManager", |
| 861 | ) -> HiSparseCoordinator: |
| 862 | """Return the coordinator shared by a KV cache manager's HiSparse groups. |
| 863 | |
| 864 | Built on first call and memoised on the HiSparse managers it binds itself |
| 865 | to. Must run before the first request is admitted. |
| 866 | """ |
| 867 | managers = tuple(kv_cache_manager.coordinator.single_type_managers) |
| 868 | for manager in managers: |
| 869 | coordinator = getattr(manager, "coordinator", None) |
| 870 | if coordinator is not None: |
| 871 | assert isinstance(coordinator, HiSparseCoordinator) |
| 872 | return coordinator |
| 873 | coordinator = HiSparseCoordinator( |
| 874 | kv_cache_manager.kv_cache_config, |
| 875 | managers, |
| 876 | kv_cache_manager.max_model_len, |
| 877 | num_reprefillable_tokens=kv_cache_manager.coordinator.num_reprefillable_tokens, |
| 878 | ) |
| 879 | if coordinator.host_manager is None: |
| 880 | raise ValueError("No HiSparse cache group is configured.") |
| 881 | return coordinator |