| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import time |
| 4 | from collections.abc import Iterable, Sequence |
| 5 | from dataclasses import dataclass, field, replace |
| 6 | from itertools import chain, islice |
| 7 | from typing import Any, NamedTuple |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.distributed.kv_events import KVCacheEvent |
| 11 | from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data |
| 12 | from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata |
| 13 | from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( |
| 14 | OffloadingConnectorMetadata, |
| 15 | OffloadingWorkerMetadata, |
| 16 | ReqId, |
| 17 | TransferJob, |
| 18 | ) |
| 19 | from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import ( |
| 20 | OffloadingEventGroupSpec, |
| 21 | OffloadingEventsTracker, |
| 22 | get_offloading_event_group_spec, |
| 23 | ) |
| 24 | from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( |
| 25 | OffloadingConnectorStats, |
| 26 | _ConnectorMetricName, |
| 27 | _TransferMetricName, |
| 28 | ) |
| 29 | from vllm.logger import init_logger |
| 30 | from vllm.utils.math_utils import cdiv, round_down |
| 31 | from vllm.v1.core.kv_cache_manager import KVCacheBlocks |
| 32 | from vllm.v1.core.sched.output import SchedulerOutput |
| 33 | from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager |
| 34 | from vllm.v1.kv_cache_interface import ( |
| 35 | ChunkedLocalAttentionSpec, |
| 36 | FullAttentionSpec, |
| 37 | KVCacheConfig, |
| 38 | KVCacheSpec, |
| 39 | MambaSpec, |
| 40 | SlidingWindowSpec, |
| 41 | ) |
| 42 | from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry |
| 43 | from vllm.v1.kv_offload.base import ( |
| 44 | GPULoadStoreSpec, |
| 45 | Locality, |
| 46 | LookupResult, |
| 47 | Medium, |
| 48 | OffloadingManager, |
| 49 | OffloadingSpec, |
| 50 | OffloadKey, |
| 51 | OffloadPolicy, |
| 52 | ReqContext, |
| 53 | RequestOffloadingContext, |
| 54 | ScheduleEndContext, |
| 55 | TierFilter, |
| 56 | TierMatcher, |
| 57 | make_offload_key, |
| 58 | ) |
| 59 | from vllm.v1.outputs import KVConnectorOutput |
| 60 | from vllm.v1.request import Request, RequestStatus |
| 61 | |
| 62 | logger = init_logger(__name__) |
| 63 | |
| 64 | KV_LOAD_TIERS_KEY = "kv_load_tiers" |
| 65 | MATCHER_MEDIUM_KEY = "medium" |
| 66 | MATCHER_LOCALITY_KEY = "locality" |
| 67 | |
| 68 | |
| 69 | @dataclass(slots=True) |
| 70 | class TransferJobStatus: |
| 71 | """Tracks scheduler-side state for a single transfer job.""" |
| 72 | |
| 73 | req_id: ReqId |
| 74 | # Number of workers still pending. Starts at num_workers, |
| 75 | # decremented as each worker reports completion. Job is done at 0. |
| 76 | pending_count: int |
| 77 | # Offload keys this job covers; passed to manager.complete_*(). |
| 78 | keys: set[OffloadKey] |
| 79 | is_store: bool |
| 80 | # Store source blocks fenced after the request finishes. |
| 81 | deferred_fence_block_ids: list[int] | None = None |
| 82 | # Store source blocks fenced when the transfer is created. |
| 83 | fenced_block_ids: list[int] | None = None |
| 84 | |
| 85 | |
| 86 | class GroupOffloadConfig(NamedTuple): |
| 87 | group_idx: int |
| 88 | tokens_per_block: int |
| 89 | tokens_per_chunk: int |
| 90 | hashes_per_chunk: int |
| 91 | # KV cache spec metadata propagated onto emitted BlockStored events so |
| 92 | # KV-aware consumers can classify and filter the group. |
| 93 | kv_event_group_spec: OffloadingEventGroupSpec |
| 94 | # None below means full attention |
| 95 | sliding_window_size_in_chunks: int | None |
| 96 | kv_cache_spec: KVCacheSpec |
| 97 | # Cached manager class for this group's KV cache spec, resolved at init time |
| 98 | manager_cls: type[SingleTypeKVCacheManager] |
| 99 | # Partial-tail data for this group comes from the scheduler's CoW hand-off |
| 100 | # rather than the request block table. |
| 101 | requires_cow_source: bool = False |
| 102 | # True for EAGLE/MTP draft-model attention groups. The trailing chunk |
| 103 | # of these groups is volatile and lacks a stable hash, so it must |
| 104 | # be excluded from store and load scheduling. |
| 105 | is_eagle_group: bool = False |
| 106 | |
| 107 | def load_window_size_in_chunks(self, num_tokens: int) -> int | None: |
| 108 | window = self.sliding_window_size_in_chunks |
| 109 | if isinstance(self.kv_cache_spec, SlidingWindowSpec): |
| 110 | assert window is not None |
| 111 | # Lookup rounds the right edge up, but allocation retains the |
| 112 | # window ending at the actual hit boundary, possibly one chunk left. |
| 113 | right_padding = -num_tokens % self.tokens_per_chunk |
| 114 | window = max( |
| 115 | window, |
| 116 | cdiv( |
| 117 | self.kv_cache_spec.sliding_window |
| 118 | - 1 |
| 119 | + self.kv_cache_spec.extra_retained_tokens |
| 120 | + right_padding, |
| 121 | self.tokens_per_chunk, |
| 122 | ), |
| 123 | ) |
| 124 | return window |
| 125 | |
| 126 | |
| 127 | def get_sliding_window_size_in_chunks( |
| 128 | kv_cache_spec: KVCacheSpec, tokens_per_chunk: int |
| 129 | ) -> int | None: |
| 130 | if isinstance(kv_cache_spec, SlidingWindowSpec): |
| 131 | assert kv_cache_spec.sliding_window > 0 |
| 132 | return max( |
| 133 | cdiv(kv_cache_spec.sliding_window, tokens_per_chunk), |
| 134 | cdiv( |
| 135 | kv_cache_spec.sliding_window - 1 + kv_cache_spec.extra_retained_tokens, |
| 136 | tokens_per_chunk, |
| 137 | ), |
| 138 | ) |
| 139 | |
| 140 | if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec): |
| 141 | # Attention never reaches back past one chunk |
| 142 | assert kv_cache_spec.attention_chunk_size > 0 |
| 143 | return cdiv(kv_cache_spec.attention_chunk_size, tokens_per_chunk) |
| 144 | |
| 145 | if isinstance(kv_cache_spec, MambaSpec): |
| 146 | # Mamba depends on a single state |
| 147 | return 1 |
| 148 | |
| 149 | assert isinstance(kv_cache_spec, FullAttentionSpec) |
| 150 | return None |
| 151 | |
| 152 | |
| 153 | def resolve_mamba_align_size( |
| 154 | spec: "OffloadingSpec", kv_cache_config: KVCacheConfig |
| 155 | ) -> int | None: |
| 156 | """Scan all KV cache groups in *spec* and return the single mamba alignment |
| 157 | size, or None if no group requires mamba alignment. |
| 158 | |
| 159 | For MambaSpec groups in "align" cache mode the hit window must be |
| 160 | rounded down to a multiple of the offloaded chunk size. Asserts that all |
| 161 | such groups agree on the same value. |
| 162 | """ |
| 163 | mamba_align_size: int | None = None |
| 164 | for group in spec.config.groups: |
| 165 | tokens_per_block = group.tokens_per_block |
| 166 | kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec |
| 167 | if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align": |
| 168 | tokens_per_chunk = tokens_per_block * spec.blocks_per_chunk |
| 169 | assert mamba_align_size is None or mamba_align_size == tokens_per_chunk |
| 170 | mamba_align_size = tokens_per_chunk |
| 171 | return mamba_align_size |
| 172 | |
| 173 | |
| 174 | class SchedulerOffloadConfig(NamedTuple): |
| 175 | kv_group_configs: tuple[GroupOffloadConfig, ...] |
| 176 | blocks_per_chunk: int |
| 177 | tokens_per_hash: int |
| 178 | num_workers: int |
| 179 | offload_prompt_only: bool |
| 180 | supports_partial_tail: bool |
| 181 | alignment_tokens: int | None = None |
| 182 | retention_interval: int | None = None |
| 183 | dcp_world_size: int = 1 |
| 184 | |
| 185 | @classmethod |
| 186 | def from_spec( |
| 187 | cls, |
| 188 | spec: OffloadingSpec, |
| 189 | vllm_config: VllmConfig, |
| 190 | kv_cache_config: KVCacheConfig, |
| 191 | ) -> "SchedulerOffloadConfig": |
| 192 | # Determine the alignment token count from the full-attention group(s). |
| 193 | # This is the tokens_per_chunk of the full-attention group; load |
| 194 | # hits are always aligned to this boundary, so SWA blocks earlier in |
| 195 | # each segment can never serve a load hit. Relevant for hybrid |
| 196 | # architectures like DeepSeek V4 (MLA + SWA groups). |
| 197 | full_attn_tokens_per_chunk: set[int] = set() |
| 198 | for group in spec.config.groups: |
| 199 | tokens_per_block = group.tokens_per_block |
| 200 | kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec |
| 201 | sw = get_sliding_window_size_in_chunks( |
| 202 | kv_spec, tokens_per_block * spec.blocks_per_chunk |
| 203 | ) |
| 204 | if sw is None: |
| 205 | full_attn_tokens_per_chunk.add(tokens_per_block * spec.blocks_per_chunk) |
| 206 | |
| 207 | # Only apply the optimization if there's a single consistent |
| 208 | # full-attention alignment size. For pure SWA/Mamba models (no |
| 209 | # full-attention groups), fall back to the largest SWA/Mamba |
| 210 | # tokens_per_chunk so retention_interval has a boundary |
| 211 | # granularity to work with. |
| 212 | alignment_tokens: int | None = None |
| 213 | if len(full_attn_tokens_per_chunk) == 1: |
| 214 | alignment_tokens = full_attn_tokens_per_chunk.pop() |
| 215 | elif not full_attn_tokens_per_chunk: |
| 216 | all_tokens_per_chunk = [ |
| 217 | tpb * spec.blocks_per_chunk for tpb in spec.tokens_per_block |
| 218 | ] |
| 219 | if all_tokens_per_chunk: |
| 220 | alignment_tokens = max(all_tokens_per_chunk) |
| 221 | |
| 222 | retention_interval = vllm_config.cache_config.prefix_cache_retention_interval |
| 223 | |
| 224 | eagle_groups = { |
| 225 | group.group_id |
| 226 | for group in spec.config.groups |
| 227 | if kv_cache_config.kv_cache_groups[group.group_id].is_eagle_group |
| 228 | } |
| 229 | |
| 230 | use_eagle_block_drop = ( |
| 231 | vllm_config.speculative_config is not None |
| 232 | and vllm_config.speculative_config.use_eagle_block_drop() |
| 233 | ) |
| 234 | if use_eagle_block_drop and not eagle_groups: |
| 235 | # No group is annotated as holding drafter layers. This happens |
| 236 | # for shared-group MTP models (e.g. Qwen3.5-style), whose drafter |
| 237 | # is a regular decoder layer merged into the target's |
| 238 | # full-attention group, and for separate-draft models that miss |
| 239 | # annotation. Treating every group as a drafter group here makes |
| 240 | # the volatile-tail pop apply to ALL groups, which can zero the |
| 241 | # whole request's offload hit whenever any group's servable |
| 242 | # window is a single chunk (issue #52735). Drafter KV served one |
| 243 | # chunk stale can only lower speculative acceptance rates — the |
| 244 | # target model verifies every draft — so fail toward serving. |
| 245 | logger.info_once( |
| 246 | "KV offloading: speculative decoding is enabled but no " |
| 247 | "KV-cache group is annotated as a drafter group; treating " |
| 248 | "all groups as non-draft for offloading." |
| 249 | ) |
| 250 | |
| 251 | if eagle_groups: |
| 252 | logger.info( |
| 253 | "KV offloading: EAGLE/MTP draft attention groups %s " |
| 254 | "detected. The trailing chunk of these groups will be " |
| 255 | "excluded from offloading due to volatility.", |
| 256 | sorted(eagle_groups), |
| 257 | ) |
| 258 | |
| 259 | kv_group_configs_list: list[GroupOffloadConfig] = [] |
| 260 | for group in spec.config.groups: |
| 261 | group_id = group.group_id |
| 262 | tokens_per_block = group.tokens_per_block |
| 263 | kv_cache_group = kv_cache_config.kv_cache_groups[group_id] |
| 264 | kv_spec = kv_cache_group.kv_cache_spec |
| 265 | sw = get_sliding_window_size_in_chunks( |
| 266 | kv_spec, tokens_per_block * spec.blocks_per_chunk |
| 267 | ) |
| 268 | manager_cls = KVCacheSpecRegistry.get_manager_class(kv_spec) |
| 269 | assert manager_cls is not None, ( |
| 270 | f"No manager found for KV cache spec {type(kv_spec).__name__}" |
| 271 | ) |
| 272 | kv_group_configs_list.append( |
| 273 | GroupOffloadConfig( |
| 274 | group_idx=group_id, |
| 275 | tokens_per_block=tokens_per_block, |
| 276 | tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk, |
| 277 | hashes_per_chunk=( |
| 278 | (tokens_per_block * spec.blocks_per_chunk) |
| 279 | // spec.tokens_per_hash |
| 280 | ), |
| 281 | sliding_window_size_in_chunks=sw, |
| 282 | kv_cache_spec=kv_spec, |
| 283 | manager_cls=manager_cls, |
| 284 | kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group), |
| 285 | is_eagle_group=group_id in eagle_groups, |
| 286 | requires_cow_source=( |
| 287 | isinstance(kv_spec, MambaSpec) |
| 288 | and kv_spec.mamba_cache_mode == "align" |
| 289 | ), |
| 290 | ) |
| 291 | ) |
| 292 | kv_group_configs = tuple(kv_group_configs_list) |
| 293 | group_block_sizes = {config.tokens_per_block for config in kv_group_configs} |
| 294 | has_partial_recurrent_group = any( |
| 295 | config.requires_cow_source |
| 296 | and config.tokens_per_block > spec.tokens_per_hash |
| 297 | for config in kv_group_configs |
| 298 | ) |
| 299 | # Partial tails currently require one physical block per offload chunk |
| 300 | # and uniform, non-windowed groups so one boundary identifies every |
| 301 | # group's source. EAGLE and DCP need additional hand-off semantics. |
| 302 | supports_partial_tail = ( |
| 303 | spec.blocks_per_chunk == 1 |
| 304 | and len(group_block_sizes) == 1 |
| 305 | and has_partial_recurrent_group |
| 306 | and all( |
| 307 | config.sliding_window_size_in_chunks is None |
| 308 | or config.requires_cow_source |
| 309 | for config in kv_group_configs |
| 310 | ) |
| 311 | and not any(config.is_eagle_group for config in kv_group_configs) |
| 312 | and vllm_config.parallel_config.decode_context_parallel_size == 1 |
| 313 | ) |
| 314 | |
| 315 | if retention_interval is not None: |
| 316 | if retention_interval < 0: |
| 317 | raise ValueError( |
| 318 | f"prefix_cache_retention_interval " |
| 319 | f"({retention_interval}) must be non-negative." |
| 320 | ) |
| 321 | for config in kv_group_configs: |
| 322 | if ( |
| 323 | config.sliding_window_size_in_chunks is not None |
| 324 | and retention_interval % config.tokens_per_chunk != 0 |
| 325 | ): |
| 326 | raise ValueError( |
| 327 | f"prefix_cache_retention_interval " |
| 328 | f"({retention_interval}) must be a multiple of " |
| 329 | f"tokens_per_chunk ({config.tokens_per_chunk})." |
| 330 | ) |
| 331 | |
| 332 | return cls( |
| 333 | num_workers=vllm_config.parallel_config.world_size, |
| 334 | kv_group_configs=kv_group_configs, |
| 335 | blocks_per_chunk=spec.blocks_per_chunk, |
| 336 | tokens_per_hash=spec.tokens_per_hash, |
| 337 | offload_prompt_only=spec.offload_prompt_only, |
| 338 | supports_partial_tail=supports_partial_tail, |
| 339 | alignment_tokens=alignment_tokens, |
| 340 | retention_interval=retention_interval, |
| 341 | dcp_world_size=vllm_config.parallel_config.decode_context_parallel_size, |
| 342 | ) |
| 343 | |
| 344 | |
| 345 | @dataclass |
| 346 | class RequestGroupState: |
| 347 | offload_keys: list[OffloadKey] = field(default_factory=list) |
| 348 | block_ids: list[int] = field(default_factory=list) |
| 349 | # Index of the next chunk to offload. |
| 350 | next_stored_chunk_idx: int = 0 |
| 351 | # Number of offloaded chunks hit (including GPU prefix cache) |
| 352 | # when the request first started |
| 353 | num_hit_chunks: int = 0 |
| 354 | |
| 355 | |
| 356 | @dataclass(slots=True) |
| 357 | class RequestOffloadState: |
| 358 | config: SchedulerOffloadConfig |
| 359 | req: Request |
| 360 | req_context: ReqContext |
| 361 | offloading_context: RequestOffloadingContext |
| 362 | group_states: tuple[RequestGroupState, ...] = field(init=False) |
| 363 | # upper bound on tokens to load for this request; None means no cap |
| 364 | max_load_tokens: int | None = None |
| 365 | # upper bound on tokens to offload for this request; None means no cap |
| 366 | max_offload_tokens: int | None = None |
| 367 | # number of hits in the GPU cache |
| 368 | num_locally_computed_tokens: int = 0 |
| 369 | # In-flight job IDs. Per the connector's invariant, at any given time |
| 370 | # this contains either a single load job, or one or more store jobs. |
| 371 | transfer_jobs: set[int] = field(default_factory=set) |
| 372 | # time.monotonic() of this request's first deferred offload lookup; |
| 373 | # None once consumed (observed) or while no lookup is pending. |
| 374 | deferred_lookup_start_time: float | None = None |
| 375 | # Fine-grained token boundary selected beyond the last complete offload |
| 376 | # chunk. It is consumed when the corresponding load is scheduled. |
| 377 | partial_tail_boundary: int | None = None |
| 378 | # True once on_request_finished has been signaled to the manager. |
| 379 | finished_signaled: bool = False |
| 380 | |
| 381 | def __post_init__(self) -> None: |
| 382 | self.group_states = tuple( |
| 383 | RequestGroupState() for _ in self.config.kv_group_configs |
| 384 | ) |
| 385 | params = self.req.kv_transfer_params |
| 386 | |
| 387 | # NOTE: This field is experimental and subject to change in the future. |
| 388 | raw = params.get("max_load_tokens") if params else None |
| 389 | if type(raw) is int and raw >= 0: |
| 390 | self.max_load_tokens = raw |
| 391 | logger.debug( |
| 392 | "Request %s: max_load_tokens set to %d", |
| 393 | self.req.request_id, |
| 394 | raw, |
| 395 | ) |
| 396 | elif raw is not None: |
| 397 | logger.warning( |
| 398 | "max_load_tokens must be a non-negative int, got %r; ignoring", raw |
| 399 | ) |
| 400 | |
| 401 | # NOTE: This field is experimental and subject to change in the future. |
| 402 | raw = params.get("max_offload_tokens") if params else None |
| 403 | if type(raw) is int and raw >= 0: |
| 404 | self.max_offload_tokens = raw |
| 405 | logger.debug( |
| 406 | "Request %s: max_offload_tokens set to %d", |
| 407 | self.req.request_id, |
| 408 | raw, |
| 409 | ) |
| 410 | elif raw is not None: |
| 411 | logger.warning( |
| 412 | "max_offload_tokens must be a non-negative int, got %r; ignoring", raw |
| 413 | ) |
| 414 | |
| 415 | def update_offload_keys(self) -> None: |
| 416 | for group_config, group_state in zip( |
| 417 | self.config.kv_group_configs, self.group_states |
| 418 | ): |
| 419 | for req_block_hash in islice( |
| 420 | self.req.block_hashes, |
| 421 | group_config.hashes_per_chunk * len(group_state.offload_keys) |
| 422 | + group_config.hashes_per_chunk |
| 423 | - 1, |
| 424 | None, |
| 425 | group_config.hashes_per_chunk, |
| 426 | ): |
| 427 | key = make_offload_key(req_block_hash, group_config.group_idx) |
| 428 | group_state.offload_keys.append(key) |
| 429 | self.req_context.set_offload_key_position( |
| 430 | key, |
| 431 | len(group_state.offload_keys) * group_config.tokens_per_chunk, |
| 432 | ) |
| 433 | |
| 434 | def update_block_id_groups( |
| 435 | self, new_block_id_groups: tuple[list[int], ...] | None |
| 436 | ) -> None: |
| 437 | if new_block_id_groups is None: |
| 438 | return |
| 439 | |
| 440 | for group_config, group_state in zip( |
| 441 | self.config.kv_group_configs, self.group_states |
| 442 | ): |
| 443 | group_state.block_ids.extend(new_block_id_groups[group_config.group_idx]) |
| 444 | |
| 445 | def storable_chunks( |
| 446 | self, |
| 447 | group_config: "GroupOffloadConfig", |
| 448 | group_state: RequestGroupState, |
| 449 | num_offloadable_tokens: int, |
| 450 | ) -> int: |
| 451 | """Number of allocated and keyed leading chunks eligible for store. |
| 452 | |
| 453 | For eagle/MTP groups the volatile trailing chunk of the offloadable |
| 454 | range is excluded while decoding: the draft-layer KV of the last |
| 455 | accepted position may be rewritten after spec-token rejection. During |
| 456 | prefill the trailing chunk is stable (the draft input for a chunk's |
| 457 | last position is the next prompt token), so it is stored immediately. |
| 458 | Once the request has finished, no further spec-token rejection can |
| 459 | rewrite the tail, so the exclusion is lifted and the final chunk |
| 460 | becomes storable (issue #52735). The exclusion must be applied |
| 461 | consistently everywhere ``next_stored_chunk_idx`` is derived: |
| 462 | otherwise the trailing chunk of each step is skipped on collection but |
| 463 | jumped over by ``next_stored_chunk_idx``, so it is never re-considered |
| 464 | and a permanent hole breaks prefix-reuse lookup. ``is_finished`` is |
| 465 | monotonic, so the finish-time calls all see the lifted exclusion. |
| 466 | """ |
| 467 | num_chunks = num_offloadable_tokens // group_config.tokens_per_chunk |
| 468 | is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens |
| 469 | # Finished requests have no pending speculation. |
| 470 | if group_config.is_eagle_group and is_decoding and not self.req.is_finished(): |
| 471 | num_chunks = max(0, num_chunks - 1) |
| 472 | num_allocated_chunks = ( |
| 473 | len(group_state.block_ids) // self.config.blocks_per_chunk |
| 474 | ) |
| 475 | num_keyed_chunks = len(group_state.offload_keys) |
| 476 | return min(num_chunks, num_allocated_chunks, num_keyed_chunks) |
| 477 | |
| 478 | def advance_stored_idx(self, num_offloadable_tokens: int) -> None: |
| 479 | # max(): at the prefill->decode transition of a chunk-aligned prompt, |
| 480 | # storable_chunks drops by one (the eagle exclusion kicks in), and the |
| 481 | # index must not move backwards past already-stored chunks. |
| 482 | for group_config, group_state in zip( |
| 483 | self.config.kv_group_configs, self.group_states |
| 484 | ): |
| 485 | group_state.next_stored_chunk_idx = max( |
| 486 | group_state.next_stored_chunk_idx, |
| 487 | self.storable_chunks(group_config, group_state, num_offloadable_tokens), |
| 488 | ) |
| 489 | |
| 490 | def update_num_hit_chunks(self, num_cached_tokens: int) -> None: |
| 491 | for group_config, group_state in zip( |
| 492 | self.config.kv_group_configs, self.group_states |
| 493 | ): |
| 494 | group_state.num_hit_chunks = ( |
| 495 | num_cached_tokens // group_config.tokens_per_chunk |
| 496 | ) |
| 497 | |
| 498 | |
| 499 | def _parse_tier_filter(raw: Any) -> TierFilter: |
| 500 | """Parse raw kv_transfer_params tier matchers into a TierFilter.""" |
| 501 | if not isinstance(raw, list): |
| 502 | logger.warning( |
| 503 | "_parse_tier_filter: expected list, got %s; ignoring", |
| 504 | type(raw).__name__, |
| 505 | ) |
| 506 | return TierFilter.ALL |
| 507 | matchers: list[TierMatcher] = [] |
| 508 | for entry in raw: |
| 509 | if not isinstance(entry, dict): |
| 510 | logger.warning("_parse_tier_filter: entry is not a dict; skipping") |
| 511 | continue |
| 512 | medium: Medium | None = None |
| 513 | locality: Locality | None = None |
| 514 | raw_medium = entry.get(MATCHER_MEDIUM_KEY) |
| 515 | if raw_medium is not None: |
| 516 | try: |
| 517 | medium = Medium(raw_medium.upper()) |
| 518 | except (ValueError, AttributeError): |
| 519 | logger.warning( |
| 520 | "_parse_tier_filter: unknown medium %r; skipping entry", |
| 521 | raw_medium, |
| 522 | ) |
| 523 | continue |
| 524 | raw_locality = entry.get(MATCHER_LOCALITY_KEY) |
| 525 | if raw_locality is not None: |
| 526 | try: |
| 527 | locality = Locality(raw_locality.upper()) |
| 528 | except (ValueError, AttributeError): |
| 529 | logger.warning( |
| 530 | "_parse_tier_filter: unknown locality %r; skipping entry", |
| 531 | raw_locality, |
| 532 | ) |
| 533 | continue |
| 534 | matchers.append(TierMatcher(medium=medium, locality=locality)) |
| 535 | if not matchers: |
| 536 | if not raw: # input was [] — user explicitly wants nothing |
| 537 | return TierFilter(matchers=()) |
| 538 | # all entries were invalid — fall back to ALL |
| 539 | return TierFilter.ALL |
| 540 | return TierFilter(matchers=tuple(matchers)) |
| 541 | |
| 542 | |
| 543 | def _create_req_context(req: Request) -> ReqContext: |
| 544 | params = req.kv_transfer_params |
| 545 | load_filter = TierFilter.ALL |
| 546 | if params: |
| 547 | raw = params.get(KV_LOAD_TIERS_KEY) |
| 548 | if raw is not None: |
| 549 | load_filter = _parse_tier_filter(raw) |
| 550 | return ReqContext( |
| 551 | req_id=req.request_id, |
| 552 | kv_transfer_params=params, |
| 553 | kv_hints=req.kv_hints, |
| 554 | load_tier_filter=load_filter, |
| 555 | ) |
| 556 | |
| 557 | |
| 558 | class OffloadingConnectorScheduler: |
| 559 | """Implementation of Scheduler side methods""" |
| 560 | |
| 561 | def __init__( |
| 562 | self, |
| 563 | spec: OffloadingSpec, |
| 564 | vllm_config: VllmConfig, |
| 565 | kv_cache_config: KVCacheConfig, |
| 566 | ): |
| 567 | self.config = SchedulerOffloadConfig.from_spec( |
| 568 | spec, vllm_config, kv_cache_config |
| 569 | ) |
| 570 | self.manager: OffloadingManager = spec.get_manager() |
| 571 | self._connector_stats = OffloadingConnectorStats() |
| 572 | |
| 573 | full_attention_groups: list[int] = [] |
| 574 | sliding_window_groups: list[int] = [] |
| 575 | for config_idx, group_config in enumerate(self.config.kv_group_configs): |
| 576 | if group_config.sliding_window_size_in_chunks is None: |
| 577 | full_attention_groups.append(config_idx) |
| 578 | else: |
| 579 | sliding_window_groups.append(config_idx) |
| 580 | |
| 581 | # sort sliding window groups by window size in decreasing order |
| 582 | def _sliding_window_sort_key(i: int) -> int: |
| 583 | val = self.config.kv_group_configs[i].sliding_window_size_in_chunks |
| 584 | assert val is not None |
| 585 | return val |
| 586 | |
| 587 | sliding_window_groups.sort(key=_sliding_window_sort_key, reverse=True) |
| 588 | |
| 589 | # used by _lookup |
| 590 | self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups) |
| 591 | self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups |
| 592 | self._mamba_align_size: int | None = resolve_mamba_align_size( |
| 593 | spec, kv_cache_config |
| 594 | ) |
| 595 | self._partial_tail_block_size = ( |
| 596 | self.config.kv_group_configs[0].tokens_per_block |
| 597 | if self.config.supports_partial_tail |
| 598 | else 0 |
| 599 | ) |
| 600 | self._cow_source_groups = frozenset( |
| 601 | config.group_idx |
| 602 | for config in self.config.kv_group_configs |
| 603 | if config.requires_cow_source |
| 604 | ) |
| 605 | |
| 606 | self._req_status: dict[ReqId, RequestOffloadState] = {} |
| 607 | self._current_batch_load_jobs: dict[int, TransferJob] = {} |
| 608 | self._current_batch_jobs_to_flush: set[int] = set() |
| 609 | # GPU block IDs allocated in the current engine step |
| 610 | self._current_batch_allocated_block_ids: set[int] = set() |
| 611 | # if GPU prefix caching is enabled, |
| 612 | # Track loaded chunks to avoid redundant loads. |
| 613 | self._chunks_being_loaded: set[OffloadKey] | None = ( |
| 614 | set() if vllm_config.cache_config.enable_prefix_caching else None |
| 615 | ) |
| 616 | |
| 617 | # Job ID counter shared by loads and stores. |
| 618 | self._job_counter: int = 0 |
| 619 | # Threshold value for stale jobs. All job ids >= _stale_job_threshold are |
| 620 | # active jobs. |
| 621 | self._stale_job_threshold: int = 0 |
| 622 | self._jobs: dict[int, TransferJobStatus] = {} |
| 623 | |
| 624 | # block_id -> pending store job_ids. Used to track jobs that needs |
| 625 | # flushing in case a block is re-allocated by the KV cache manager. |
| 626 | # Populated only for finished requests (running-request blocks are |
| 627 | # protected by their ref_cnt) and for sliding window blocks (which can |
| 628 | # be freed before a request finishes). |
| 629 | self._block_id_to_pending_jobs: dict[int, set[int]] = {} |
| 630 | |
| 631 | self._events_tracker = OffloadingEventsTracker(spec.kv_events_config) |
| 632 | |
| 633 | def _maybe_observe_lookup_async_delay( |
| 634 | self, req_status: RequestOffloadState |
| 635 | ) -> None: |
| 636 | start_time = req_status.deferred_lookup_start_time |
| 637 | if start_time is None: |
| 638 | return |
| 639 | req_status.deferred_lookup_start_time = None |
| 640 | self._connector_stats.observe_histogram( |
| 641 | _ConnectorMetricName.LOOKUP_ASYNC_DELAY, |
| 642 | time.monotonic() - start_time, |
| 643 | ) |
| 644 | |
| 645 | def _generate_job_id(self) -> int: |
| 646 | job_id = self._job_counter |
| 647 | self._job_counter += 1 |
| 648 | return job_id |
| 649 | |
| 650 | def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None: |
| 651 | for bid in block_ids or (): |
| 652 | pending = self._block_id_to_pending_jobs[bid] |
| 653 | pending.remove(job_id) |
| 654 | if not pending: |
| 655 | del self._block_id_to_pending_jobs[bid] |
| 656 | |
| 657 | def _calc_num_offloadable_tokens( |
| 658 | self, req_status: RequestOffloadState, num_computed_tokens: int |
| 659 | ) -> int: |
| 660 | num = min(num_computed_tokens, req_status.req.num_tokens) |
| 661 | max_offload_tokens = req_status.max_offload_tokens |
| 662 | if max_offload_tokens is not None: |
| 663 | num = min(num, max_offload_tokens) |
| 664 | if self.config.offload_prompt_only: |
| 665 | num = min(num, req_status.req.num_prompt_tokens) |
| 666 | return num |
| 667 | |
| 668 | def _maximal_prefix_lookup( |
| 669 | self, |
| 670 | keys: Iterable[OffloadKey], |
| 671 | req_context: ReqContext, |
| 672 | req: Request, |
| 673 | group_config: GroupOffloadConfig, |
| 674 | start_chunk_idx: int, |
| 675 | ) -> int | None: |
| 676 | """Return the number of consecutive offloaded chunks from the start, |
| 677 | or None if the backend deferred a lookup.""" |
| 678 | hit_count = 0 |
| 679 | defer_lookup = False |
| 680 | for local_idx, key in enumerate(keys): |
| 681 | result = self.manager.lookup(key, req_context) |
| 682 | match result: |
| 683 | case LookupResult.HIT: |
| 684 | self._events_tracker.record_lookup( |
| 685 | req, |
| 686 | group_config, |
| 687 | start_chunk_idx + local_idx, |
| 688 | key, |
| 689 | ) |
| 690 | hit_count += 1 |
| 691 | case LookupResult.HIT_PENDING: |
| 692 | defer_lookup = True |
| 693 | hit_count += 1 |
| 694 | case LookupResult.RETRY: |
| 695 | # Don't break: keep scanning to let manager kick off |
| 696 | # async lookups (until a miss is detected). |
| 697 | defer_lookup = True |
| 698 | case LookupResult.MISS: |
| 699 | break |
| 700 | return hit_count if not defer_lookup else None |
| 701 | |
| 702 | def _sliding_window_lookup( |
| 703 | self, |
| 704 | keys: Sequence[OffloadKey], |
| 705 | sliding_window_size: int, |
| 706 | req_context: ReqContext, |
| 707 | initial_window_size: int | None = None, |
| 708 | ) -> int | None: |
| 709 | """Return the end index (in `keys`) of the last run of |
| 710 | `sliding_window_size` consecutive hits, scanning from the end. |
| 711 | The first run may need a larger window for a partial rightmost chunk. |
| 712 | Returns 0 on miss, None if the backend deferred a lookup.""" |
| 713 | defer_lookup = False |
| 714 | pending_in_window = False |
| 715 | consecutive_hits = 0 |
| 716 | required_window = initial_window_size or sliding_window_size |
| 717 | for idx in range(len(keys) - 1, -1, -1): |
| 718 | match self.manager.lookup(keys[idx], req_context): |
| 719 | case LookupResult.HIT: |
| 720 | consecutive_hits += 1 |
| 721 | case LookupResult.HIT_PENDING: |
| 722 | # Block is in cache, just not readable yet — counts |
| 723 | # as hit for the consecutive streak. Don't break: |
| 724 | # keep scanning to let manager kick off async lookups. |
| 725 | pending_in_window = True |
| 726 | consecutive_hits += 1 |
| 727 | case LookupResult.RETRY: |
| 728 | # Block location uncertain — does not count as hit. |
| 729 | # Don't break: keep scanning to let manager kick off |
| 730 | # async lookups. |
| 731 | defer_lookup = True |
| 732 | consecutive_hits = 0 |
| 733 | pending_in_window = False |
| 734 | required_window = sliding_window_size |
| 735 | case LookupResult.MISS: |
| 736 | consecutive_hits = 0 |
| 737 | # This gap rules out the incomplete window to its right. |
| 738 | pending_in_window = False |
| 739 | required_window = sliding_window_size |
| 740 | if consecutive_hits == required_window: |
| 741 | return ( |
| 742 | None if defer_lookup or pending_in_window else idx + required_window |
| 743 | ) |
| 744 | return None if defer_lookup or pending_in_window else consecutive_hits |
| 745 | |
| 746 | def _lookup_complete_chunks( |
| 747 | self, |
| 748 | req_status: RequestOffloadState, |
| 749 | max_num_new_tokens: int | None = None, |
| 750 | ) -> int | None: |
| 751 | """Find how many tokens beyond num_locally_computed_tokens can be loaded. |
| 752 | |
| 753 | Iterates full-attention groups first (prefix lookup), then sliding-window |
| 754 | groups (suffix lookup). Each group may tighten max_hit_size_tokens, which |
| 755 | can invalidate an earlier group's result, so the loop re-runs when that |
| 756 | happens until num_hit_tokens converges. |
| 757 | """ |
| 758 | num_computed_tokens = req_status.num_locally_computed_tokens |
| 759 | max_hit_size_tokens: int = req_status.req.num_tokens |
| 760 | if max_num_new_tokens is not None: |
| 761 | max_hit_size_tokens = min( |
| 762 | max_hit_size_tokens, num_computed_tokens + max_num_new_tokens |
| 763 | ) |
| 764 | if req_status.max_load_tokens is not None: |
| 765 | max_hit_size_tokens = min( |
| 766 | max_hit_size_tokens, |
| 767 | num_computed_tokens + req_status.max_load_tokens, |
| 768 | ) |
| 769 | if self._sliding_window_groups: |
| 770 | # the last prompt token has to be recomputed to get the logprobs |
| 771 | # for sliding window attention, we must reduce by 1 to make sure |
| 772 | # we still have a hit after reduction |
| 773 | max_hit_size_tokens -= 1 |
| 774 | if self._mamba_align_size is not None: |
| 775 | # Constrain hit-window to the mamba block size. |
| 776 | max_hit_size_tokens = round_down( |
| 777 | max_hit_size_tokens, self._mamba_align_size |
| 778 | ) |
| 779 | |
| 780 | num_hit_tokens: int = 0 |
| 781 | defer_lookup = False |
| 782 | lookup_groups = self._lookup_groups |
| 783 | |
| 784 | # Tracks which eagle groups have already popped their volatile trailing chunk |
| 785 | # in the current convergence iteration. Reset when a non-eagle group |
| 786 | # tightens the hit boundary, requiring a fresh pop. |
| 787 | eagle_verified: set[int] = set() |
| 788 | while lookup_groups: |
| 789 | looked_up_sliding_window: bool = False |
| 790 | groups_iter = iter(lookup_groups) |
| 791 | lookup_groups = () |
| 792 | for group_idx in groups_iter: |
| 793 | group_config: GroupOffloadConfig = self.config.kv_group_configs[ |
| 794 | group_idx |
| 795 | ] |
| 796 | group_state: RequestGroupState = req_status.group_states[group_idx] |
| 797 | tokens_per_chunk = group_config.tokens_per_chunk |
| 798 | offload_keys = group_state.offload_keys |
| 799 | |
| 800 | assert ( |
| 801 | len(offload_keys) >= req_status.req.num_tokens // tokens_per_chunk |
| 802 | ) |
| 803 | |
| 804 | is_eagle_unverified = ( |
| 805 | group_config.is_eagle_group and group_idx not in eagle_verified |
| 806 | ) |
| 807 | |
| 808 | # Constrain to a chunk-aligned boundary for this group. |
| 809 | max_hit_size_tokens = min( |
| 810 | max_hit_size_tokens, len(offload_keys) * tokens_per_chunk |
| 811 | ) |
| 812 | if req_status.max_load_tokens is not None: |
| 813 | max_hit_size_tokens = round_down( |
| 814 | max_hit_size_tokens, tokens_per_chunk |
| 815 | ) |
| 816 | if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk: |
| 817 | # We can only load less than a chunk, so skip. |
| 818 | return 0 |
| 819 | |
| 820 | sliding_window_size_in_chunks = ( |
| 821 | group_config.sliding_window_size_in_chunks |
| 822 | ) |
| 823 | |
| 824 | # For eagle groups, query one extra chunk that will be popped. |
| 825 | # Widening applies to every group type: without it, the pop |
| 826 | # below shrinks max_hit_size_tokens past what was queried, |
| 827 | # which can push the confirmed boundary under a coarser |
| 828 | # sibling group's chunk granularity and zero the whole |
| 829 | # request's hit (issue #52735). |
| 830 | query_max = max_hit_size_tokens |
| 831 | if is_eagle_unverified: |
| 832 | query_max = min( |
| 833 | max_hit_size_tokens + tokens_per_chunk, |
| 834 | len(offload_keys) * tokens_per_chunk, |
| 835 | ) |
| 836 | |
| 837 | num_chunks = min(cdiv(query_max, tokens_per_chunk), len(offload_keys)) |
| 838 | start_chunk_idx = num_computed_tokens // tokens_per_chunk |
| 839 | offload_keys = offload_keys[start_chunk_idx:num_chunks] |
| 840 | |
| 841 | # end index (in the sliced offload_keys) up to which we |
| 842 | # have backend-confirmed hits |
| 843 | num_hit_chunks: int | None |
| 844 | if sliding_window_size_in_chunks is None: |
| 845 | num_hit_chunks = self._maximal_prefix_lookup( |
| 846 | offload_keys, |
| 847 | req_status.req_context, |
| 848 | req_status.req, |
| 849 | group_config, |
| 850 | start_chunk_idx, |
| 851 | ) |
| 852 | else: |
| 853 | required_window = sliding_window_size_in_chunks |
| 854 | if is_eagle_unverified: |
| 855 | required_window += 1 |
| 856 | candidate_end = min( |
| 857 | max_hit_size_tokens, |
| 858 | (num_chunks - int(is_eagle_unverified)) * tokens_per_chunk, |
| 859 | ) |
| 860 | initial_window = group_config.load_window_size_in_chunks( |
| 861 | candidate_end |
| 862 | ) |
| 863 | assert initial_window is not None |
| 864 | num_hit_chunks = self._sliding_window_lookup( |
| 865 | offload_keys, |
| 866 | required_window, |
| 867 | req_status.req_context, |
| 868 | initial_window + int(is_eagle_unverified), |
| 869 | ) |
| 870 | if num_hit_chunks == 0: |
| 871 | return 0 |
| 872 | |
| 873 | if num_hit_chunks is None: |
| 874 | defer_lookup = True |
| 875 | else: |
| 876 | if is_eagle_unverified: |
| 877 | num_hit_chunks -= 1 |
| 878 | eagle_verified.add(group_idx) |
| 879 | |
| 880 | max_hit_size_tokens = min( |
| 881 | max_hit_size_tokens, |
| 882 | tokens_per_chunk * (start_chunk_idx + num_hit_chunks), |
| 883 | ) |
| 884 | |
| 885 | new_num_hit_tokens = max_hit_size_tokens - num_computed_tokens |
| 886 | if new_num_hit_tokens < tokens_per_chunk: |
| 887 | # We can only load less than a chunk, so skip. |
| 888 | return 0 |
| 889 | |
| 890 | if new_num_hit_tokens < num_hit_tokens: |
| 891 | if not group_config.is_eagle_group: |
| 892 | eagle_verified.clear() |
| 893 | if defer_lookup: |
| 894 | # make another iteration on all groups to check |
| 895 | # if we still need to defer lookup |
| 896 | defer_lookup = False |
| 897 | lookup_groups = self._lookup_groups |
| 898 | elif looked_up_sliding_window and not lookup_groups: |
| 899 | # we need another iteration to confirm previously looked up |
| 900 | # sliding window works with the new_num_hit_tokens |
| 901 | lookup_groups = self._sliding_window_groups |
| 902 | |
| 903 | looked_up_sliding_window |= sliding_window_size_in_chunks is not None |
| 904 | num_hit_tokens = new_num_hit_tokens |
| 905 | |
| 906 | if defer_lookup: |
| 907 | logger.debug( |
| 908 | "Offloading manager delayed request %s as backend requested", |
| 909 | req_status.req.request_id, |
| 910 | ) |
| 911 | return None |
| 912 | |
| 913 | # Possibly delay the request if any hit chunk is already being loaded. |
| 914 | if self._chunks_being_loaded: |
| 915 | for group_config, group_state in zip( |
| 916 | self.config.kv_group_configs, req_status.group_states |
| 917 | ): |
| 918 | tokens_per_chunk = group_config.tokens_per_chunk |
| 919 | sliding_window_size_in_chunks = group_config.load_window_size_in_chunks( |
| 920 | num_computed_tokens + num_hit_tokens |
| 921 | ) |
| 922 | offload_keys = group_state.offload_keys |
| 923 | num_chunks = cdiv( |
| 924 | num_computed_tokens + num_hit_tokens, tokens_per_chunk |
| 925 | ) |
| 926 | start_chunk_idx = num_computed_tokens // tokens_per_chunk |
| 927 | offload_keys = offload_keys[start_chunk_idx:num_chunks] |
| 928 | if sliding_window_size_in_chunks is not None: |
| 929 | offload_keys = offload_keys[-sliding_window_size_in_chunks:] |
| 930 | if any(key in self._chunks_being_loaded for key in offload_keys): |
| 931 | # Hit chunks are being loaded, so delay the request. |
| 932 | logger.debug( |
| 933 | "Delaying request %s since some of its" |
| 934 | " chunks are already being loaded", |
| 935 | req_status.req.request_id, |
| 936 | ) |
| 937 | return None |
| 938 | |
| 939 | logger.debug( |
| 940 | "Request %s hit %s offloaded tokens after %s GPU hit tokens", |
| 941 | req_status.req.request_id, |
| 942 | num_hit_tokens, |
| 943 | num_computed_tokens, |
| 944 | ) |
| 945 | |
| 946 | return num_hit_tokens |
| 947 | |
| 948 | def _make_boundary_key( |
| 949 | self, |
| 950 | request: Request, |
| 951 | group_idx: int, |
| 952 | boundary_tokens: int, |
| 953 | req_context: ReqContext, |
| 954 | ) -> OffloadKey: |
| 955 | hash_idx = boundary_tokens // self.config.tokens_per_hash - 1 |
| 956 | key = make_offload_key(request.block_hashes[hash_idx], group_idx) |
| 957 | req_context.set_offload_key_position(key, boundary_tokens) |
| 958 | return key |
| 959 | |
| 960 | def _lookup( |
| 961 | self, |
| 962 | req_status: RequestOffloadState, |
| 963 | max_num_new_tokens: int | None = None, |
| 964 | ) -> int | None: |
| 965 | complete_hit = self._lookup_complete_chunks(req_status, max_num_new_tokens) |
| 966 | req_status.partial_tail_boundary = None |
| 967 | if complete_hit is None or not self.config.supports_partial_tail: |
| 968 | return complete_hit |
| 969 | |
| 970 | local_tokens = req_status.num_locally_computed_tokens |
| 971 | complete_boundary = local_tokens + complete_hit |
| 972 | tokens_per_hash = self.config.tokens_per_hash |
| 973 | block_end = complete_boundary + self._partial_tail_block_size |
| 974 | max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1) |
| 975 | if max_num_new_tokens is not None: |
| 976 | max_boundary = min(max_boundary, local_tokens + max_num_new_tokens) |
| 977 | if req_status.max_load_tokens is not None: |
| 978 | max_boundary = min( |
| 979 | max_boundary, |
| 980 | local_tokens + req_status.max_load_tokens, |
| 981 | ) |
| 982 | max_boundary = round_down(max_boundary, tokens_per_hash) |
| 983 | if max_boundary <= complete_boundary: |
| 984 | return complete_hit |
| 985 | |
| 986 | pending = False |
| 987 | for boundary in range(max_boundary, complete_boundary, -tokens_per_hash): |
| 988 | boundary_pending = False |
| 989 | boundary_missed = False |
| 990 | boundary_keys = [] |
| 991 | for group_config in self.config.kv_group_configs: |
| 992 | key = self._make_boundary_key( |
| 993 | req_status.req, |
| 994 | group_config.group_idx, |
| 995 | boundary, |
| 996 | req_status.req_context, |
| 997 | ) |
| 998 | boundary_keys.append(key) |
| 999 | result = self.manager.lookup(key, req_status.req_context) |
| 1000 | if result is LookupResult.MISS: |
| 1001 | boundary_missed = True |
| 1002 | break |
| 1003 | if result in (LookupResult.HIT_PENDING, LookupResult.RETRY): |
| 1004 | boundary_pending = True |
| 1005 | |
| 1006 | pending |= boundary_pending |
| 1007 | if not boundary_missed and not boundary_pending: |
| 1008 | for group_config, key in zip( |
| 1009 | self.config.kv_group_configs, boundary_keys |
| 1010 | ): |
| 1011 | self._events_tracker.record_partial_lookup( |
| 1012 | req_status.req, group_config, boundary, key |
| 1013 | ) |
| 1014 | req_status.partial_tail_boundary = boundary |
| 1015 | return boundary - local_tokens |
| 1016 | |
| 1017 | if pending and complete_hit == 0: |
| 1018 | return None |
| 1019 | return complete_hit |
| 1020 | |
| 1021 | def on_new_request(self, request: Request) -> None: |
| 1022 | """Called when a new request is added to the scheduler.""" |
| 1023 | req_context = _create_req_context(request) |
| 1024 | offloading_context = self.manager.on_new_request(req_context) |
| 1025 | req_status = RequestOffloadState( |
| 1026 | config=self.config, |
| 1027 | req=request, |
| 1028 | req_context=req_context, |
| 1029 | offloading_context=offloading_context, |
| 1030 | ) |
| 1031 | self._req_status[request.request_id] = req_status |
| 1032 | |
| 1033 | def get_num_new_matched_tokens( |
| 1034 | self, |
| 1035 | request: Request, |
| 1036 | num_computed_tokens: int, |
| 1037 | max_num_new_tokens: int | None = None, |
| 1038 | ) -> tuple[int | None, bool]: |
| 1039 | """Get number of new tokens that can be loaded beyond the |
| 1040 | num_computed_tokens. |
| 1041 | |
| 1042 | Args: |
| 1043 | request (Request): the request object. |
| 1044 | num_computed_tokens (int): the number of locally |
| 1045 | computed tokens for this request |
| 1046 | max_num_new_tokens (int | None): cap on the number of tokens that |
| 1047 | may be loaded beyond `num_computed_tokens`, if any. |
| 1048 | |
| 1049 | Returns: |
| 1050 | A tuple with the following elements: |
| 1051 | - The number of tokens that can be loaded beyond what is |
| 1052 | already computed. |
| 1053 | If None, it means that the connector needs more time to |
| 1054 | determine the number of matched tokens, and the scheduler |
| 1055 | should query for this request again later. |
| 1056 | - `True` if tokens will be loaded asynchronously |
| 1057 | (between scheduler steps). |
| 1058 | |
| 1059 | """ |
| 1060 | req_status = self._req_status[request.request_id] |
| 1061 | for group_state in req_status.group_states: |
| 1062 | group_state.block_ids.clear() |
| 1063 | |
| 1064 | if req_status.transfer_jobs: |
| 1065 | logger.debug( |
| 1066 | "Delaying request %s since it still has in-flight transfers", |
| 1067 | request.request_id, |
| 1068 | ) |
| 1069 | return None, False |
| 1070 | |
| 1071 | req_status.update_offload_keys() |
| 1072 | req_status.num_locally_computed_tokens = num_computed_tokens |
| 1073 | |
| 1074 | num_hit_tokens: int | None |
| 1075 | if request.skip_reading_prefix_cache: |
| 1076 | num_hit_tokens = 0 |
| 1077 | else: |
| 1078 | lookup_start = time.monotonic() |
| 1079 | num_hit_tokens = self._lookup(req_status, max_num_new_tokens) |
| 1080 | self._connector_stats.observe_histogram( |
| 1081 | _ConnectorMetricName.LOOKUP_SYNC_DELAY, |
| 1082 | time.monotonic() - lookup_start, |
| 1083 | ) |
| 1084 | if num_hit_tokens is None: |
| 1085 | if req_status.deferred_lookup_start_time is None: |
| 1086 | req_status.deferred_lookup_start_time = lookup_start |
| 1087 | else: |
| 1088 | self._maybe_observe_lookup_async_delay(req_status) |
| 1089 | req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0)) |
| 1090 | |
| 1091 | return num_hit_tokens, bool(num_hit_tokens) |
| 1092 | |
| 1093 | def update_state_after_alloc( |
| 1094 | self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int |
| 1095 | ): |
| 1096 | if num_external_tokens == 0: |
| 1097 | return |
| 1098 | |
| 1099 | req_status = self._req_status[request.request_id] |
| 1100 | |
| 1101 | num_locally_computed_tokens = req_status.num_locally_computed_tokens |
| 1102 | num_cached_tokens = num_locally_computed_tokens + num_external_tokens |
| 1103 | partial_tail_boundary = req_status.partial_tail_boundary |
| 1104 | if partial_tail_boundary is not None: |
| 1105 | assert partial_tail_boundary == num_cached_tokens |
| 1106 | |
| 1107 | keys_to_load: list[OffloadKey] = [] |
| 1108 | dst_block_ids: list[int] = [] |
| 1109 | # per group |
| 1110 | group_sizes: list[int] = [] |
| 1111 | block_indices: list[int] = [] |
| 1112 | for group_config, group_state in zip( |
| 1113 | self.config.kv_group_configs, |
| 1114 | req_status.group_states, |
| 1115 | ): |
| 1116 | group_blocks = blocks.blocks[group_config.group_idx] |
| 1117 | self._current_batch_allocated_block_ids.update( |
| 1118 | block.block_id for block in group_blocks if block.block_id != 0 |
| 1119 | ) |
| 1120 | |
| 1121 | tokens_per_block = group_config.tokens_per_block |
| 1122 | tokens_per_chunk = group_config.tokens_per_chunk |
| 1123 | offload_keys = group_state.offload_keys |
| 1124 | num_gpu_blocks = cdiv(num_cached_tokens, tokens_per_block) |
| 1125 | |
| 1126 | assert len(group_blocks) >= num_gpu_blocks |
| 1127 | # ``load_start_gpu_block_idx``: the index in ``group_blocks`` where the |
| 1128 | # load region begins -- a slice bound, not a count of computed blocks. |
| 1129 | # Scan from the computed boundary, not 0: sparse groups (Mamba, |
| 1130 | # SWA) legitimately hold non-null unhashed blocks below it. |
| 1131 | # Skip nulls (sentinel / out-of-retention padding). |
| 1132 | first_fresh_gpu_block_idx = cdiv( |
| 1133 | num_locally_computed_tokens, tokens_per_block |
| 1134 | ) |
| 1135 | load_start_gpu_block_idx = num_gpu_blocks |
| 1136 | for i in range(first_fresh_gpu_block_idx, num_gpu_blocks): |
| 1137 | block = group_blocks[i] |
| 1138 | if not block.is_null and block.block_hash is None: |
| 1139 | load_start_gpu_block_idx = i |
| 1140 | break |
| 1141 | |
| 1142 | assert num_locally_computed_tokens % tokens_per_block == 0 |
| 1143 | num_pending_gpu_blocks = num_gpu_blocks - load_start_gpu_block_idx |
| 1144 | |
| 1145 | if group_config.sliding_window_size_in_chunks is not None: |
| 1146 | assert ( |
| 1147 | num_pending_gpu_blocks |
| 1148 | <= group_config.sliding_window_size_in_chunks |
| 1149 | * self.config.blocks_per_chunk |
| 1150 | + 1 |
| 1151 | ) |
| 1152 | |
| 1153 | num_chunks = cdiv(num_cached_tokens, tokens_per_chunk) |
| 1154 | if num_pending_gpu_blocks: |
| 1155 | start_chunk_idx = ( |
| 1156 | load_start_gpu_block_idx // self.config.blocks_per_chunk |
| 1157 | ) |
| 1158 | end_chunk_idx = num_chunks - (partial_tail_boundary is not None) |
| 1159 | assert len(offload_keys) >= end_chunk_idx |
| 1160 | keys_to_load.extend(offload_keys[start_chunk_idx:end_chunk_idx]) |
| 1161 | if partial_tail_boundary is not None: |
| 1162 | keys_to_load.append( |
| 1163 | self._make_boundary_key( |
| 1164 | request, |
| 1165 | group_config.group_idx, |
| 1166 | partial_tail_boundary, |
| 1167 | req_status.req_context, |
| 1168 | ) |
| 1169 | ) |
| 1170 | |
| 1171 | dst_block_ids.extend( |
| 1172 | block.block_id |
| 1173 | for block in group_blocks[load_start_gpu_block_idx:num_gpu_blocks] |
| 1174 | ) |
| 1175 | group_sizes.append(num_pending_gpu_blocks) |
| 1176 | block_indices.append(load_start_gpu_block_idx) |
| 1177 | |
| 1178 | # Skip prefix-hit chunks for block-level policy; for |
| 1179 | # request-level, next_stored_chunk_idx stays at 0 so all |
| 1180 | # chunks (including hits) are offloaded. |
| 1181 | if req_status.offloading_context.policy == OffloadPolicy.CHUNK_LEVEL: |
| 1182 | group_state.next_stored_chunk_idx = num_chunks |
| 1183 | |
| 1184 | src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context) |
| 1185 | dst_spec = GPULoadStoreSpec( |
| 1186 | dst_block_ids, group_sizes=group_sizes, block_indices=block_indices |
| 1187 | ) |
| 1188 | |
| 1189 | load_job_id = self._generate_job_id() |
| 1190 | self._current_batch_load_jobs[load_job_id] = TransferJob( |
| 1191 | req_id=request.request_id, |
| 1192 | src_spec=src_spec, |
| 1193 | dst_spec=dst_spec, |
| 1194 | ) |
| 1195 | # a load can only be issued when no other jobs are pending. |
| 1196 | assert not req_status.transfer_jobs |
| 1197 | req_status.transfer_jobs.add(load_job_id) |
| 1198 | self._jobs[load_job_id] = TransferJobStatus( |
| 1199 | req_id=request.request_id, |
| 1200 | pending_count=self.config.num_workers, |
| 1201 | keys=set(keys_to_load), |
| 1202 | is_store=False, |
| 1203 | ) |
| 1204 | |
| 1205 | if self._chunks_being_loaded is not None: |
| 1206 | self._chunks_being_loaded.update(keys_to_load) |
| 1207 | req_status.partial_tail_boundary = None |
| 1208 | |
| 1209 | def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: |
| 1210 | """Update request states from the Scheduler's output.""" |
| 1211 | # new_block_ids_end[req_id][i] = end of pre-existing block_ids for |
| 1212 | # the i-th sliding window group (before this step's extend). |
| 1213 | # Used to detect sliding window blocks that got re-allocated. |
| 1214 | new_block_ids_end: dict[str, tuple[int, ...]] = {} |
| 1215 | |
| 1216 | for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output): |
| 1217 | req_status = self._req_status[req_id] |
| 1218 | req_status.update_offload_keys() |
| 1219 | |
| 1220 | if preempted: |
| 1221 | for group_state in req_status.group_states: |
| 1222 | group_state.block_ids.clear() |
| 1223 | |
| 1224 | if new_block_id_groups: |
| 1225 | if self._sliding_window_groups: |
| 1226 | new_block_ids_end[req_id] = tuple( |
| 1227 | len(req_status.group_states[grp_idx].block_ids) |
| 1228 | for grp_idx in self._sliding_window_groups |
| 1229 | ) |
| 1230 | req_status.update_block_id_groups(new_block_id_groups) |
| 1231 | for group_config in self.config.kv_group_configs: |
| 1232 | new_blocks = new_block_id_groups[group_config.group_idx] |
| 1233 | for bid in new_blocks: |
| 1234 | if bid != 0: |
| 1235 | self._current_batch_allocated_block_ids.add(bid) |
| 1236 | |
| 1237 | for copy in scheduler_output.kv_cache_block_copies or (): |
| 1238 | self._current_batch_allocated_block_ids.add(copy.dst_block_id) |
| 1239 | |
| 1240 | # Zero out stale block_ids in sliding window groups' pending-store |
| 1241 | # positions. Only sliding window groups can have stale entries (blocks |
| 1242 | # freed by remove_skipped_blocks then reallocated). Only positions in |
| 1243 | # [next_stored_chunk_idx * bsf, end) need checking where end is the |
| 1244 | # pre-extend length: earlier positions were already offloaded, later |
| 1245 | # ones are fresh allocations from this step. |
| 1246 | if self._sliding_window_groups and self._current_batch_allocated_block_ids: |
| 1247 | blocks_per_chunk = self.config.blocks_per_chunk |
| 1248 | for req_id, req_status in self._req_status.items(): |
| 1249 | ends = new_block_ids_end.get(req_id) |
| 1250 | for i, grp_idx in enumerate(self._sliding_window_groups): |
| 1251 | group_state = req_status.group_states[grp_idx] |
| 1252 | start = group_state.next_stored_chunk_idx * blocks_per_chunk |
| 1253 | end = ends[i] if ends is not None else len(group_state.block_ids) |
| 1254 | for j in range(start, end): |
| 1255 | if ( |
| 1256 | group_state.block_ids[j] |
| 1257 | in self._current_batch_allocated_block_ids |
| 1258 | ): |
| 1259 | group_state.block_ids[j] = 0 |
| 1260 | |
| 1261 | def _build_aligned_boundary_store_jobs( |
| 1262 | self, handoffs: dict[str, list[tuple[int, int, int]]] |
| 1263 | ) -> dict[int, TransferJob]: |
| 1264 | store_jobs: dict[int, TransferJob] = {} |
| 1265 | num_groups = len(self.config.kv_group_configs) |
| 1266 | config_idx_by_group = { |
| 1267 | config.group_idx: idx |
| 1268 | for idx, config in enumerate(self.config.kv_group_configs) |
| 1269 | } |
| 1270 | for req_id, entries in handoffs.items(): |
| 1271 | req_status = self._req_status.get(req_id) |
| 1272 | if req_status is None: |
| 1273 | continue |
| 1274 | req = req_status.req |
| 1275 | max_boundary = self._calc_num_offloadable_tokens(req_status, req.num_tokens) |
| 1276 | for group_idx, block_id, boundary in entries: |
| 1277 | config_idx = config_idx_by_group.get(group_idx) |
| 1278 | if config_idx is None: |
| 1279 | continue |
| 1280 | group_config = self.config.kv_group_configs[config_idx] |
| 1281 | if ( |
| 1282 | block_id == 0 |
| 1283 | or boundary > max_boundary |
| 1284 | or boundary % group_config.tokens_per_chunk != 0 |
| 1285 | ): |
| 1286 | continue |
| 1287 | |
| 1288 | key = self._make_boundary_key( |
| 1289 | req, group_idx, boundary, req_status.req_context |
| 1290 | ) |
| 1291 | store_output = self.manager.prepare_store([key], req_status.req_context) |
| 1292 | if store_output is None: |
| 1293 | self._connector_stats.increase_counter( |
| 1294 | _ConnectorMetricName.ALLOCATION_FAILURE |
| 1295 | ) |
| 1296 | continue |
| 1297 | if not store_output.keys_to_store: |
| 1298 | continue |
| 1299 | |
| 1300 | job_id = self._generate_job_id() |
| 1301 | req_status.transfer_jobs.add(job_id) |
| 1302 | self._block_id_to_pending_jobs.setdefault(block_id, set()).add(job_id) |
| 1303 | self._jobs[job_id] = TransferJobStatus( |
| 1304 | req_id=req_id, |
| 1305 | pending_count=self.config.num_workers, |
| 1306 | keys={key}, |
| 1307 | is_store=True, |
| 1308 | fenced_block_ids=[block_id], |
| 1309 | ) |
| 1310 | group_sizes = [0] * num_groups |
| 1311 | group_sizes[config_idx] = 1 |
| 1312 | block_indices = [0] * num_groups |
| 1313 | block_indices[config_idx] = ( |
| 1314 | boundary // group_config.tokens_per_block - 1 |
| 1315 | ) |
| 1316 | store_jobs[job_id] = TransferJob( |
| 1317 | req_id=req_id, |
| 1318 | src_spec=GPULoadStoreSpec( |
| 1319 | [block_id], |
| 1320 | group_sizes=group_sizes, |
| 1321 | block_indices=block_indices, |
| 1322 | ), |
| 1323 | dst_spec=store_output.store_spec, |
| 1324 | ) |
| 1325 | self._events_tracker.record_store( |
| 1326 | req, |
| 1327 | group_config, |
| 1328 | boundary // group_config.tokens_per_chunk - 1, |
| 1329 | key, |
| 1330 | ) |
| 1331 | return store_jobs |
| 1332 | |
| 1333 | def _build_partial_tail_store_jobs( |
| 1334 | self, scheduler_output: SchedulerOutput |
| 1335 | ) -> dict[int, TransferJob]: |
| 1336 | block_state = scheduler_output.kv_connector_block_state |
| 1337 | handoffs = block_state.boundary_state_offloads if block_state else None |
| 1338 | if not handoffs: |
| 1339 | return {} |
| 1340 | |
| 1341 | store_jobs = self._build_aligned_boundary_store_jobs(handoffs) |
| 1342 | if not self.config.supports_partial_tail: |
| 1343 | return store_jobs |
| 1344 | |
| 1345 | for req_id, entries in handoffs.items(): |
| 1346 | entries = [ |
| 1347 | entry |
| 1348 | for entry in entries |
| 1349 | if entry[2] % self._partial_tail_block_size != 0 |
| 1350 | ] |
| 1351 | if not entries: |
| 1352 | continue |
| 1353 | req_status = self._req_status.get(req_id) |
| 1354 | assert req_status is not None |
| 1355 | boundaries = {boundary for _, _, boundary in entries} |
| 1356 | assert len(boundaries) == 1 |
| 1357 | boundary = boundaries.pop() |
| 1358 | req = req_status.req |
| 1359 | group_states = { |
| 1360 | group.group_idx: state |
| 1361 | for group, state in zip( |
| 1362 | self.config.kv_group_configs, req_status.group_states |
| 1363 | ) |
| 1364 | } |
| 1365 | max_boundary = min( |
| 1366 | req.num_prompt_tokens, |
| 1367 | req_status.max_offload_tokens or req.num_prompt_tokens, |
| 1368 | ) |
| 1369 | assert boundary > 0 |
| 1370 | assert boundary % self.config.tokens_per_hash == 0 |
| 1371 | assert boundary <= max_boundary |
| 1372 | |
| 1373 | cow_blocks = {group_idx: block_id for group_idx, block_id, _ in entries} |
| 1374 | assert self._cow_source_groups.issubset(cow_blocks) |
| 1375 | |
| 1376 | assert boundary % self._partial_tail_block_size != 0 |
| 1377 | block_idx = boundary // self._partial_tail_block_size |
| 1378 | if any( |
| 1379 | group.group_idx not in self._cow_source_groups |
| 1380 | and block_idx >= len(group_states[group.group_idx].block_ids) |
| 1381 | for group in self.config.kv_group_configs |
| 1382 | ): |
| 1383 | continue |
| 1384 | keys = [ |
| 1385 | self._make_boundary_key( |
| 1386 | req, group.group_idx, boundary, req_status.req_context |
| 1387 | ) |
| 1388 | for group in self.config.kv_group_configs |
| 1389 | ] |
| 1390 | block_ids = [ |
| 1391 | cow_blocks[group.group_idx] |
| 1392 | if group.group_idx in self._cow_source_groups |
| 1393 | else group_states[group.group_idx].block_ids[block_idx] |
| 1394 | for group in self.config.kv_group_configs |
| 1395 | ] |
| 1396 | assert all(block_id != 0 for block_id in block_ids) |
| 1397 | |
| 1398 | store_output = self.manager.prepare_store(keys, req_status.req_context) |
| 1399 | if store_output is None: |
| 1400 | self._connector_stats.increase_counter( |
| 1401 | _ConnectorMetricName.ALLOCATION_FAILURE |
| 1402 | ) |
| 1403 | continue |
| 1404 | if not store_output.keys_to_store: |
| 1405 | continue |
| 1406 | |
| 1407 | for group_config, key in zip(self.config.kv_group_configs, keys): |
| 1408 | if key in store_output.keys_to_store: |
| 1409 | self._events_tracker.record_partial_store( |
| 1410 | req, group_config, boundary, key |
| 1411 | ) |
| 1412 | |
| 1413 | group_by_key = {key: idx for idx, key in enumerate(keys)} |
| 1414 | accepted_groups = [group_by_key[key] for key in store_output.keys_to_store] |
| 1415 | group_sizes = [0] * len(self.config.kv_group_configs) |
| 1416 | block_indices = [0] * len(self.config.kv_group_configs) |
| 1417 | for group_idx in accepted_groups: |
| 1418 | group_sizes[group_idx] = 1 |
| 1419 | block_indices[group_idx] = block_idx |
| 1420 | source_blocks = [block_ids[group_idx] for group_idx in accepted_groups] |
| 1421 | |
| 1422 | job_id = self._generate_job_id() |
| 1423 | req_status.transfer_jobs.add(job_id) |
| 1424 | for block_id in source_blocks: |
| 1425 | self._block_id_to_pending_jobs.setdefault(block_id, set()).add(job_id) |
| 1426 | self._jobs[job_id] = TransferJobStatus( |
| 1427 | req_id=req_id, |
| 1428 | pending_count=self.config.num_workers, |
| 1429 | keys=set(store_output.keys_to_store), |
| 1430 | is_store=True, |
| 1431 | fenced_block_ids=source_blocks, |
| 1432 | ) |
| 1433 | store_jobs[job_id] = TransferJob( |
| 1434 | req_id=req_id, |
| 1435 | src_spec=GPULoadStoreSpec( |
| 1436 | source_blocks, |
| 1437 | group_sizes=group_sizes, |
| 1438 | block_indices=block_indices, |
| 1439 | ), |
| 1440 | dst_spec=store_output.store_spec, |
| 1441 | ) |
| 1442 | |
| 1443 | return store_jobs |
| 1444 | |
| 1445 | def _reachable_store_block_mask( |
| 1446 | self, |
| 1447 | group_config: GroupOffloadConfig, |
| 1448 | start_chunk_idx: int, |
| 1449 | end_chunk_idx: int, |
| 1450 | final_segment_end_chunk_idx: int | None, |
| 1451 | reachable_boundaries: tuple[int, ...], |
| 1452 | ) -> list[bool] | None: |
| 1453 | """Build the block mask for a range of candidate offload chunks.""" |
| 1454 | blocks_per_chunk = self.config.blocks_per_chunk |
| 1455 | kv_cache_spec = group_config.kv_cache_spec |
| 1456 | if isinstance(kv_cache_spec, SlidingWindowSpec) and ( |
| 1457 | kv_cache_spec.extra_retained_tokens |
| 1458 | ): |
| 1459 | # Offload restores all allocated history, including MTP re-prefill |
| 1460 | # tokens. Widen only the store mask, not the model's attention window. |
| 1461 | kv_cache_spec = replace( |
| 1462 | kv_cache_spec, |
| 1463 | sliding_window=( |
| 1464 | kv_cache_spec.sliding_window + kv_cache_spec.extra_retained_tokens |
| 1465 | ), |
| 1466 | extra_retained_tokens=0, |
| 1467 | ) |
| 1468 | return group_config.manager_cls.reachable_block_mask( |
| 1469 | start_block=start_chunk_idx * blocks_per_chunk, |
| 1470 | end_block=end_chunk_idx * blocks_per_chunk, |
| 1471 | alignment_tokens=self.config.alignment_tokens, |
| 1472 | kv_cache_spec=kv_cache_spec, |
| 1473 | use_eagle=group_config.is_eagle_group, |
| 1474 | retention_interval=self.config.retention_interval, |
| 1475 | reachable_boundaries=reachable_boundaries, |
| 1476 | dcp_world_size=self.config.dcp_world_size, |
| 1477 | final_segment_end_block=( |
| 1478 | final_segment_end_chunk_idx * blocks_per_chunk |
| 1479 | if final_segment_end_chunk_idx is not None |
| 1480 | else None |
| 1481 | ), |
| 1482 | ) |
| 1483 | |
| 1484 | def _final_swa_alignment_blocks( |
| 1485 | self, group_config: GroupOffloadConfig |
| 1486 | ) -> int | None: |
| 1487 | """Return representable final-tail alignment for a non-EAGLE SWA group.""" |
| 1488 | alignment_tokens = self.config.alignment_tokens |
| 1489 | if ( |
| 1490 | alignment_tokens is None |
| 1491 | or group_config.is_eagle_group |
| 1492 | or not isinstance(group_config.kv_cache_spec, SlidingWindowSpec) |
| 1493 | or alignment_tokens % group_config.tokens_per_block != 0 |
| 1494 | ): |
| 1495 | return None |
| 1496 | return alignment_tokens // group_config.tokens_per_block |
| 1497 | |
| 1498 | def _build_store_jobs( |
| 1499 | self, |
| 1500 | scheduler_output: SchedulerOutput, |
| 1501 | ) -> dict[int, TransferJob]: |
| 1502 | blocks_per_chunk = self.config.blocks_per_chunk |
| 1503 | store_jobs: dict[int, TransferJob] = {} |
| 1504 | for req_id in chain( |
| 1505 | scheduler_output.num_scheduled_tokens, |
| 1506 | scheduler_output.finished_req_ids or (), |
| 1507 | ): |
| 1508 | req_status = self._req_status.get(req_id) |
| 1509 | if req_status is None: |
| 1510 | continue |
| 1511 | req = req_status.req |
| 1512 | |
| 1513 | if req.status is RequestStatus.FINISHED_ABORTED: |
| 1514 | num_tokens_after_batch = req.num_computed_tokens |
| 1515 | elif req.is_finished(): |
| 1516 | # Clamp to the GPU prefix cache's commit point. The final sampled |
| 1517 | # token's slot is never committed (under spec decode it holds a |
| 1518 | # rejected draft's KV), so a block ending there must not be stored. |
| 1519 | num_tokens_after_batch = max(req.num_prompt_tokens, req.num_tokens - 1) |
| 1520 | else: |
| 1521 | num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] |
| 1522 | num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens |
| 1523 | |
| 1524 | num_offloadable_tokens = self._calc_num_offloadable_tokens( |
| 1525 | req_status, num_tokens_after_batch |
| 1526 | ) |
| 1527 | prompt_offloadable_tokens = self._calc_num_offloadable_tokens( |
| 1528 | req_status, req.num_prompt_tokens |
| 1529 | ) |
| 1530 | |
| 1531 | # Filter out chunks skipped due to sliding window attention / SSM |
| 1532 | # or unreachable by the load path's alignment constraints. |
| 1533 | new_offload_keys: list[OffloadKey] = [] |
| 1534 | group_store_ranges: list[tuple[int, int]] = [] |
| 1535 | |
| 1536 | reachable_boundaries: tuple[int, ...] = () |
| 1537 | if self.config.retention_interval is not None: |
| 1538 | reachable_boundaries = (req.num_prompt_tokens - 1,) |
| 1539 | if req.shared_prefix_boundary: |
| 1540 | reachable_boundaries += (req.shared_prefix_boundary,) |
| 1541 | |
| 1542 | for group_config, group_state in zip( |
| 1543 | self.config.kv_group_configs, req_status.group_states |
| 1544 | ): |
| 1545 | num_chunks = req_status.storable_chunks( |
| 1546 | group_config, group_state, num_offloadable_tokens |
| 1547 | ) |
| 1548 | |
| 1549 | start_chunk_idx = group_state.next_stored_chunk_idx |
| 1550 | prompt_horizon_chunks = ( |
| 1551 | prompt_offloadable_tokens // group_config.tokens_per_chunk |
| 1552 | ) |
| 1553 | final_swa_alignment_blocks = self._final_swa_alignment_blocks( |
| 1554 | group_config |
| 1555 | ) |
| 1556 | reconsider_final_swa_tail = ( |
| 1557 | req.is_finished() |
| 1558 | and num_chunks != prompt_horizon_chunks |
| 1559 | and self.config.retention_interval is None |
| 1560 | and final_swa_alignment_blocks is not None |
| 1561 | ) |
| 1562 | if ( |
| 1563 | self.config.retention_interval is not None |
| 1564 | or group_config.is_eagle_group |
| 1565 | ): |
| 1566 | store_horizon_chunks = None |
| 1567 | elif req.is_finished(): |
| 1568 | store_horizon_chunks = num_chunks |
| 1569 | elif num_chunks <= prompt_horizon_chunks: |
| 1570 | store_horizon_chunks = prompt_horizon_chunks |
| 1571 | else: |
| 1572 | # An active decode frontier is not a final request |
| 1573 | # boundary. Only fixed alignment tails are reachable. |
| 1574 | store_horizon_chunks = None |
| 1575 | if reconsider_final_swa_tail: |
| 1576 | assert final_swa_alignment_blocks is not None |
| 1577 | horizon_blocks = num_chunks * blocks_per_chunk |
| 1578 | partial_segment_start_block = ( |
| 1579 | horizon_blocks - horizon_blocks % final_swa_alignment_blocks |
| 1580 | ) |
| 1581 | partial_segment_start_chunk = ( |
| 1582 | partial_segment_start_block // blocks_per_chunk |
| 1583 | ) |
| 1584 | start_chunk_idx = min(start_chunk_idx, partial_segment_start_chunk) |
| 1585 | group_store_ranges.append((start_chunk_idx, num_chunks)) |
| 1586 | |
| 1587 | if group_config.requires_cow_source: |
| 1588 | continue |
| 1589 | |
| 1590 | if num_chunks <= start_chunk_idx: |
| 1591 | continue |
| 1592 | offload_keys = group_state.offload_keys[start_chunk_idx:num_chunks] |
| 1593 | # For each chunk, take the last corresponding GPU block. For |
| 1594 | # blocks_per_chunk=3 and GPU block IDs 1 5 6 7 2 4 9 3 8, |
| 1595 | # this selects GPU blocks 6 4 8. |
| 1596 | # A block_id of 0 means either a sliding window / SSM skip |
| 1597 | # or a stale entry that was zeroed out — skip it either way. |
| 1598 | offload_block_ids = group_state.block_ids[ |
| 1599 | start_chunk_idx * blocks_per_chunk |
| 1600 | + blocks_per_chunk |
| 1601 | - 1 : num_chunks * blocks_per_chunk : blocks_per_chunk |
| 1602 | ] |
| 1603 | assert len(offload_keys) == len(offload_block_ids) |
| 1604 | |
| 1605 | # Use reachable_block_mask to filter unreachable chunks |
| 1606 | # (SWA/Mamba sparsity + retention interval). |
| 1607 | # reachable_block_mask operates in KV-block coordinates, |
| 1608 | # so convert chunk indices to block indices. |
| 1609 | block_mask = self._reachable_store_block_mask( |
| 1610 | group_config=group_config, |
| 1611 | start_chunk_idx=start_chunk_idx, |
| 1612 | end_chunk_idx=num_chunks, |
| 1613 | final_segment_end_chunk_idx=store_horizon_chunks, |
| 1614 | reachable_boundaries=reachable_boundaries, |
| 1615 | ) |
| 1616 | |
| 1617 | prompt_horizon_block_mask: list[bool] | None = None |
| 1618 | if reconsider_final_swa_tail: |
| 1619 | prompt_horizon_block_mask = self._reachable_store_block_mask( |
| 1620 | group_config=group_config, |
| 1621 | start_chunk_idx=start_chunk_idx, |
| 1622 | end_chunk_idx=num_chunks, |
| 1623 | final_segment_end_chunk_idx=prompt_horizon_chunks, |
| 1624 | reachable_boundaries=reachable_boundaries, |
| 1625 | ) |
| 1626 | |
| 1627 | for key_idx, (offload_key, block_id) in enumerate( |
| 1628 | zip(offload_keys, offload_block_ids) |
| 1629 | ): |
| 1630 | if block_id == 0: |
| 1631 | continue |
| 1632 | abs_chunk_idx = start_chunk_idx + key_idx |
| 1633 | if ( |
| 1634 | reconsider_final_swa_tail |
| 1635 | and abs_chunk_idx < group_state.next_stored_chunk_idx |
| 1636 | and ( |
| 1637 | prompt_horizon_block_mask is None |
| 1638 | or any( |
| 1639 | prompt_horizon_block_mask[ |
| 1640 | key_idx * blocks_per_chunk + block_idx |
| 1641 | ] |
| 1642 | for block_idx in range(blocks_per_chunk) |
| 1643 | ) |
| 1644 | ) |
| 1645 | ): |
| 1646 | continue |
| 1647 | # A chunk is reachable if any of its constituent |
| 1648 | # blocks is reachable. |
| 1649 | if block_mask is not None and not any( |
| 1650 | block_mask[key_idx * blocks_per_chunk + b] |
| 1651 | for b in range(blocks_per_chunk) |
| 1652 | ): |
| 1653 | continue |
| 1654 | new_offload_keys.append(offload_key) |
| 1655 | |
| 1656 | if not new_offload_keys: |
| 1657 | req_status.advance_stored_idx(num_offloadable_tokens) |
| 1658 | continue |
| 1659 | |
| 1660 | store_output = self.manager.prepare_store( |
| 1661 | new_offload_keys, req_status.req_context |
| 1662 | ) |
| 1663 | if store_output is None: |
| 1664 | self._connector_stats.increase_counter( |
| 1665 | _ConnectorMetricName.ALLOCATION_FAILURE |
| 1666 | ) |
| 1667 | logger.warning("Request %s: cannot store chunks", req_id) |
| 1668 | continue |
| 1669 | |
| 1670 | if not store_output.keys_to_store: |
| 1671 | req_status.advance_stored_idx(num_offloadable_tokens) |
| 1672 | continue |
| 1673 | |
| 1674 | keys_to_store = set(store_output.keys_to_store) |
| 1675 | |
| 1676 | group_sizes: list[int] = [] |
| 1677 | block_indices: list[int] = [] |
| 1678 | src_block_ids: list[int] = [] |
| 1679 | fenced_block_ids: list[int] = [] |
| 1680 | deferred_fence_block_ids: list[int] = [] |
| 1681 | for group_config, group_state, store_range in zip( |
| 1682 | self.config.kv_group_configs, |
| 1683 | req_status.group_states, |
| 1684 | group_store_ranges, |
| 1685 | ): |
| 1686 | is_sliding_window = ( |
| 1687 | group_config.sliding_window_size_in_chunks is not None |
| 1688 | ) |
| 1689 | start_chunk_idx, num_chunks = store_range |
| 1690 | block_ids = group_state.block_ids |
| 1691 | num_group_blocks = 0 |
| 1692 | start_gpu_block_idx: int | None = None |
| 1693 | for idx, offload_key in enumerate( |
| 1694 | group_state.offload_keys[start_chunk_idx:num_chunks] |
| 1695 | ): |
| 1696 | if offload_key not in keys_to_store: |
| 1697 | continue |
| 1698 | |
| 1699 | chunk_idx = start_chunk_idx + idx |
| 1700 | |
| 1701 | self._events_tracker.record_store( |
| 1702 | req, group_config, chunk_idx, offload_key |
| 1703 | ) |
| 1704 | |
| 1705 | gpu_block_idx = chunk_idx * blocks_per_chunk |
| 1706 | for i in range(blocks_per_chunk): |
| 1707 | block_id = block_ids[gpu_block_idx + i] |
| 1708 | if block_id == 0: |
| 1709 | continue |
| 1710 | if start_gpu_block_idx is None: |
| 1711 | start_gpu_block_idx = gpu_block_idx + i |
| 1712 | src_block_ids.append(block_id) |
| 1713 | num_group_blocks += 1 |
| 1714 | if is_sliding_window: |
| 1715 | fenced_block_ids.append(block_id) |
| 1716 | else: |
| 1717 | deferred_fence_block_ids.append(block_id) |
| 1718 | |
| 1719 | group_sizes.append(num_group_blocks) |
| 1720 | block_indices.append(start_gpu_block_idx or 0) |
| 1721 | group_state.next_stored_chunk_idx = max( |
| 1722 | group_state.next_stored_chunk_idx, num_chunks |
| 1723 | ) |
| 1724 | |
| 1725 | src_spec = GPULoadStoreSpec( |
| 1726 | src_block_ids, group_sizes=group_sizes, block_indices=block_indices |
| 1727 | ) |
| 1728 | dst_spec = store_output.store_spec |
| 1729 | |
| 1730 | job_id = self._generate_job_id() |
| 1731 | # a store can only be issued when no load is pending. |
| 1732 | if req_status.transfer_jobs: |
| 1733 | any_jid = next(iter(req_status.transfer_jobs)) |
| 1734 | assert self._jobs[any_jid].is_store |
| 1735 | req_status.transfer_jobs.add(job_id) |
| 1736 | |
| 1737 | # Watch sliding window blocks as they may get evicted |
| 1738 | # before the request finishes |
| 1739 | for bid in fenced_block_ids: |
| 1740 | self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id) |
| 1741 | |
| 1742 | # the non-sliding window blocks will be watched only |
| 1743 | # when the request finishes |
| 1744 | self._jobs[job_id] = TransferJobStatus( |
| 1745 | req_id=req_id, |
| 1746 | pending_count=self.config.num_workers, |
| 1747 | keys=set(keys_to_store), |
| 1748 | is_store=True, |
| 1749 | deferred_fence_block_ids=deferred_fence_block_ids, |
| 1750 | fenced_block_ids=fenced_block_ids or None, |
| 1751 | ) |
| 1752 | |
| 1753 | store_jobs[job_id] = TransferJob( |
| 1754 | req_id=req_id, src_spec=src_spec, dst_spec=dst_spec |
| 1755 | ) |
| 1756 | |
| 1757 | logger.debug( |
| 1758 | "Request %s offloading %s chunks upto %d tokens (job %d)", |
| 1759 | req_id, |
| 1760 | len(keys_to_store), |
| 1761 | num_offloadable_tokens, |
| 1762 | job_id, |
| 1763 | ) |
| 1764 | |
| 1765 | if req.is_finished(): |
| 1766 | # Register non-sliding-window blocks for flush detection. |
| 1767 | for bid in deferred_fence_block_ids: |
| 1768 | self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id) |
| 1769 | if bid in self._current_batch_allocated_block_ids: |
| 1770 | self._current_batch_jobs_to_flush.add(job_id) |
| 1771 | |
| 1772 | return store_jobs |
| 1773 | |
| 1774 | def build_connector_meta( |
| 1775 | self, scheduler_output: SchedulerOutput |
| 1776 | ) -> KVConnectorMetadata: |
| 1777 | self._update_req_states(scheduler_output) |
| 1778 | schedule_end_context = ScheduleEndContext( |
| 1779 | new_req_ids=[req.req_id for req in scheduler_output.scheduled_new_reqs], |
| 1780 | preempted_req_ids=scheduler_output.preempted_req_ids or (), |
| 1781 | ) |
| 1782 | self.manager.on_schedule_end(schedule_end_context) |
| 1783 | |
| 1784 | # Flush jobs for preempted requests. |
| 1785 | for req_id in scheduler_output.preempted_req_ids or (): |
| 1786 | req_status = self._req_status.get(req_id) |
| 1787 | if req_status is None or not req_status.transfer_jobs: |
| 1788 | continue |
| 1789 | any_jid = next(iter(req_status.transfer_jobs)) |
| 1790 | assert self._jobs[any_jid].is_store |
| 1791 | self._current_batch_jobs_to_flush.update(req_status.transfer_jobs) |
| 1792 | |
| 1793 | # Flush jobs that contain re-allocated blocks. |
| 1794 | if ( |
| 1795 | self._block_id_to_pending_jobs |
| 1796 | and not self._block_id_to_pending_jobs.keys().isdisjoint( |
| 1797 | self._current_batch_allocated_block_ids |
| 1798 | ) |
| 1799 | ): |
| 1800 | self._current_batch_jobs_to_flush.update( |
| 1801 | jid |
| 1802 | for bid in self._current_batch_allocated_block_ids |
| 1803 | if bid in self._block_id_to_pending_jobs |
| 1804 | for jid in self._block_id_to_pending_jobs[bid] |
| 1805 | ) |
| 1806 | |
| 1807 | partial_store_jobs = self._build_partial_tail_store_jobs(scheduler_output) |
| 1808 | normal_store_jobs = self._build_store_jobs(scheduler_output) |
| 1809 | meta = OffloadingConnectorMetadata( |
| 1810 | load_jobs=self._current_batch_load_jobs, |
| 1811 | store_jobs=partial_store_jobs | normal_store_jobs, |
| 1812 | jobs_to_flush=self._current_batch_jobs_to_flush, |
| 1813 | ) |
| 1814 | |
| 1815 | # All prepare_store calls for finished requests have been issued. |
| 1816 | # Signal on_request_finished and clean up state where possible. |
| 1817 | for req_id in scheduler_output.finished_req_ids or (): |
| 1818 | req_status = self._req_status.get(req_id) |
| 1819 | if req_status is None: |
| 1820 | continue |
| 1821 | req_status.finished_signaled = True |
| 1822 | self.manager.on_request_finished(req_status.req_context) |
| 1823 | if not req_status.transfer_jobs: |
| 1824 | del self._req_status[req_id] |
| 1825 | self._current_batch_load_jobs = {} |
| 1826 | self._current_batch_jobs_to_flush = set() |
| 1827 | self._current_batch_allocated_block_ids = set() |
| 1828 | return meta |
| 1829 | |
| 1830 | def has_pending_push_work(self) -> bool: |
| 1831 | """Whether the engine must keep stepping. |
| 1832 | |
| 1833 | While True, build_connector_meta() and update_connector_output() |
| 1834 | continue to be called even when no requests are scheduled. |
| 1835 | """ |
| 1836 | return bool(self._jobs) or self.manager.has_pending_work() |
| 1837 | |
| 1838 | def update_connector_output(self, connector_output: KVConnectorOutput): |
| 1839 | """Update KVConnector state from worker-side connectors output. |
| 1840 | |
| 1841 | Args: |
| 1842 | connector_output (KVConnectorOutput): the worker-side |
| 1843 | connectors output. |
| 1844 | |
| 1845 | """ |
| 1846 | meta = connector_output.kv_connector_worker_meta |
| 1847 | if not isinstance(meta, OffloadingWorkerMetadata): |
| 1848 | assert meta is None |
| 1849 | meta = OffloadingWorkerMetadata() |
| 1850 | if not meta.transfer_stats.is_empty(): |
| 1851 | transfer_stats = OffloadingConnectorStats() |
| 1852 | if not meta.transfer_stats.load.is_empty(): |
| 1853 | transfer_stats.increase_counter( |
| 1854 | _TransferMetricName.LOAD_BYTES, |
| 1855 | meta.transfer_stats.load.bytes, |
| 1856 | ) |
| 1857 | transfer_stats.increase_counter( |
| 1858 | _TransferMetricName.LOAD_TIME, |
| 1859 | meta.transfer_stats.load.time, |
| 1860 | ) |
| 1861 | for size in meta.transfer_stats.load.sizes: |
| 1862 | transfer_stats.observe_histogram( |
| 1863 | _TransferMetricName.LOAD_SIZE, size |
| 1864 | ) |
| 1865 | if not meta.transfer_stats.store.is_empty(): |
| 1866 | transfer_stats.increase_counter( |
| 1867 | _TransferMetricName.STORE_BYTES, |
| 1868 | meta.transfer_stats.store.bytes, |
| 1869 | ) |
| 1870 | transfer_stats.increase_counter( |
| 1871 | _TransferMetricName.STORE_TIME, |
| 1872 | meta.transfer_stats.store.time, |
| 1873 | ) |
| 1874 | for size in meta.transfer_stats.store.sizes: |
| 1875 | transfer_stats.observe_histogram( |
| 1876 | _TransferMetricName.STORE_SIZE, size |
| 1877 | ) |
| 1878 | self._connector_stats.aggregate(transfer_stats) |
| 1879 | |
| 1880 | for job_id, count in meta.completed_jobs.items(): |
| 1881 | assert count > 0 |
| 1882 | if job_id < self._stale_job_threshold: |
| 1883 | logger.debug( |
| 1884 | "Skipping stale completed job %d (pre-reset counter: %d)", |
| 1885 | job_id, |
| 1886 | self._stale_job_threshold, |
| 1887 | ) |
| 1888 | continue |
| 1889 | job_status = self._jobs[job_id] |
| 1890 | job_status.pending_count -= count |
| 1891 | if job_status.pending_count > 0: |
| 1892 | continue |
| 1893 | assert job_status.pending_count == 0 |
| 1894 | |
| 1895 | req_status = self._req_status[job_status.req_id] |
| 1896 | if job_status.is_store: |
| 1897 | self.manager.complete_store(job_status.keys, req_status.req_context) |
| 1898 | else: |
| 1899 | self.manager.complete_load(job_status.keys, req_status.req_context) |
| 1900 | if self._chunks_being_loaded: |
| 1901 | self._chunks_being_loaded.difference_update(job_status.keys) |
| 1902 | if self._block_id_to_pending_jobs: |
| 1903 | # Sliding window blocks are tracked from store creation |
| 1904 | # and must be cleaned up unconditionally. |
| 1905 | self._remove_pending_job(job_id, job_status.fenced_block_ids) |
| 1906 | # Non-sliding-window blocks are only tracked after |
| 1907 | # request_finished, so only clean up for finished requests. |
| 1908 | if req_status.req.is_finished(): |
| 1909 | self._remove_pending_job( |
| 1910 | job_id, job_status.deferred_fence_block_ids |
| 1911 | ) |
| 1912 | |
| 1913 | del self._jobs[job_id] |
| 1914 | req_status.transfer_jobs.remove(job_id) |
| 1915 | if req_status.finished_signaled and not req_status.transfer_jobs: |
| 1916 | del self._req_status[job_status.req_id] |
| 1917 | |
| 1918 | def get_stats(self) -> OffloadingConnectorStats | None: |
| 1919 | stats: OffloadingConnectorStats | None = None |
| 1920 | if not self._connector_stats.is_empty(): |
| 1921 | stats = self._connector_stats |
| 1922 | self._connector_stats = OffloadingConnectorStats() |
| 1923 | |
| 1924 | manager_stats = self.manager.get_stats() |
| 1925 | if manager_stats is not None: |
| 1926 | if stats is None: |
| 1927 | stats = manager_stats |
| 1928 | else: |
| 1929 | stats.aggregate(manager_stats) |
| 1930 | |
| 1931 | return stats |
| 1932 | |
| 1933 | def request_finished( |
| 1934 | self, |
| 1935 | request: Request, |
| 1936 | ) -> tuple[bool, dict[str, Any] | None]: |
| 1937 | """Called when a request has finished, before its blocks are freed. |
| 1938 | |
| 1939 | Returns: |
| 1940 | True if the request is being saved/sent asynchronously and blocks |
| 1941 | should not be freed until the request_id is returned from |
| 1942 | get_finished(). |
| 1943 | Optional KVTransferParams to be included in the request outputs |
| 1944 | returned by the engine. |
| 1945 | |
| 1946 | """ |
| 1947 | req_status = self._req_status.get(request.request_id) |
| 1948 | |
| 1949 | if req_status is None: |
| 1950 | # Untracked request (offloading never started): no in-flight jobs, |
| 1951 | # nothing was deferred, so finalize immediately. |
| 1952 | req_context = _create_req_context(request) |
| 1953 | self.manager.on_new_request(req_context) |
| 1954 | self.manager.on_request_finished(req_context) |
| 1955 | return False, None |
| 1956 | |
| 1957 | self._maybe_observe_lookup_async_delay(req_status) |
| 1958 | |
| 1959 | # Update offload keys with final block hash so _build_store_jobs can |
| 1960 | # create store jobs for the last block(s) on the next schedule step. |
| 1961 | req_status.update_offload_keys() |
| 1962 | |
| 1963 | # Keep req_status alive: _build_store_jobs will process finished_req_ids |
| 1964 | # on the next step and handle cleanup after creating store jobs. |
| 1965 | # Register deferred fences so future block reuse triggers a flush via |
| 1966 | # _block_id_to_pending_jobs. |
| 1967 | for job_id in req_status.transfer_jobs: |
| 1968 | job_status = self._jobs[job_id] |
| 1969 | for bid in job_status.deferred_fence_block_ids or (): |
| 1970 | self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id) |
| 1971 | |
| 1972 | return False, None |
| 1973 | |
| 1974 | def take_events(self) -> Iterable[KVCacheEvent]: |
| 1975 | """Drain pending KV cache events. |
| 1976 | |
| 1977 | Complete metadata is available only when self-describing KV events |
| 1978 | are enabled, and only for full-attention groups. Other shapes retain |
| 1979 | the previous placeholder payload so consumers can ignore them. |
| 1980 | |
| 1981 | Yields: |
| 1982 | ``BlockStored`` or ``BlockRemoved`` events corresponding to |
| 1983 | the underlying :class:`OffloadingEvent` stream. |
| 1984 | |
| 1985 | """ |
| 1986 | yield from self._events_tracker.take_events(self.manager.take_events()) |
| 1987 | |
| 1988 | def reset_cache(self) -> None: |
| 1989 | """Reset the offloading manager cache, evicting all stored chunks.""" |
| 1990 | # reset_cache cannot be called in the middle of a schedule step |
| 1991 | assert not self._current_batch_load_jobs |
| 1992 | assert not self._current_batch_jobs_to_flush |
| 1993 | assert not self._current_batch_allocated_block_ids |
| 1994 | |
| 1995 | # Flush all in-flight jobs |
| 1996 | self._current_batch_jobs_to_flush.update(self._jobs.keys()) |
| 1997 | |
| 1998 | for req_id, status in list(self._req_status.items()): |
| 1999 | if status.req.is_finished(): |
| 2000 | if not status.finished_signaled: |
| 2001 | self.manager.on_request_finished(status.req_context) |
| 2002 | del self._req_status[req_id] |
| 2003 | |
| 2004 | # Reset offloading manager cache |
| 2005 | self.manager.reset_cache() |
| 2006 | |
| 2007 | # Reset store progress so active requests re-offload from chunk 0. |
| 2008 | for status in self._req_status.values(): |
| 2009 | for group_state in status.group_states: |
| 2010 | group_state.next_stored_chunk_idx = 0 |
| 2011 | status.transfer_jobs.clear() |
| 2012 | status.partial_tail_boundary = None |
| 2013 | |
| 2014 | # Discard jobs and save job_counter to be able to discard worker responses |
| 2015 | self._stale_job_threshold = self._job_counter |
| 2016 | self._jobs.clear() |
| 2017 | self._block_id_to_pending_jobs.clear() |
| 2018 | |
| 2019 | # The manager pool is empty; pending event payloads and announced |
| 2020 | # reference counts are stale. |
| 2021 | self._events_tracker.reset() |
| 2022 | |
| 2023 | # Note: _current_batch_jobs_to_flush is intentionally NOT cleared. |
| 2024 | # The load flush IDs collected above must be delivered to workers. |
| 2025 | if self._chunks_being_loaded is not None: |
| 2026 | self._chunks_being_loaded.clear() |
| 2027 | |
| 2028 | def shutdown(self) -> None: |
| 2029 | self.manager.shutdown() |