1148 lines
47 KiB
Python
1148 lines
47 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Scheduler-side manager for SimpleCPUOffloadConnector."""
|
|
|
|
import contextlib
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass, field, replace
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.distributed.kv_events import (
|
|
MEDIUM_CPU,
|
|
MEDIUM_STORAGE,
|
|
BlockRemoved,
|
|
BlockStored,
|
|
KVCacheEvent,
|
|
)
|
|
from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.math_utils import cdiv
|
|
from vllm.v1.core.block_pool import BlockPool
|
|
from vllm.v1.core.kv_cache_coordinator import (
|
|
KVCacheCoordinator,
|
|
get_kv_cache_coordinator,
|
|
)
|
|
from vllm.v1.core.kv_cache_utils import (
|
|
BlockHashWithGroupId,
|
|
ExternalBlockHash,
|
|
generate_block_hash_extra_keys,
|
|
get_block_hash,
|
|
get_group_id,
|
|
make_block_hash_with_group_id,
|
|
maybe_convert_block_hash,
|
|
resolve_block_hashes,
|
|
)
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
from vllm.v1.kv_cache_interface import (
|
|
FullAttentionSpec,
|
|
MambaSpec,
|
|
SlidingWindowSpec,
|
|
get_kv_cache_spec_kind,
|
|
get_kv_cache_spec_sliding_window,
|
|
)
|
|
from vllm.v1.outputs import KVConnectorOutput
|
|
from vllm.v1.simple_kv_offload.metadata import (
|
|
SimpleCPUOffloadMetadata,
|
|
SimpleCPUOffloadWorkerMetadata,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
|
|
from vllm.v1.core.kv_cache_utils import KVCacheBlock
|
|
from vllm.v1.kv_cache_interface import KVCacheConfig
|
|
from vllm.v1.request import Request
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class BlockStoreMeta:
|
|
"""Per-block metadata snapshot for BlockStored event emission.
|
|
|
|
Captured at store-prep time (when the Request is available) and carried
|
|
through TransferMeta to async completion, where the Request is gone.
|
|
"""
|
|
|
|
token_ids: list[int]
|
|
parent_block_hash: ExternalBlockHash | None
|
|
lora_id: int | None
|
|
lora_name: str | None
|
|
extra_keys: tuple[Any, ...] | None
|
|
block_size: int | None = None
|
|
|
|
|
|
@dataclass
|
|
class TransferMeta:
|
|
gpu_block_ids: list[int]
|
|
cpu_block_ids: list[int]
|
|
block_meta: list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None = None
|
|
|
|
|
|
@dataclass
|
|
class LoadRequestState:
|
|
request: "Request"
|
|
transfer_meta: TransferMeta
|
|
load_event: int | None = None
|
|
finished: bool = False
|
|
|
|
|
|
# NOTE: This per-request state is only used in eager mode.
|
|
@dataclass
|
|
class StoreRequestState:
|
|
request: "Request"
|
|
# Accumulated block IDs from scheduler_output via yield_req_data.
|
|
block_ids: tuple[list[int], ...]
|
|
# Per-group cursors tracking how many blocks have been stored/skipped.
|
|
num_stored_blocks: list[int]
|
|
store_events: set[int] = field(default_factory=set)
|
|
finished: bool = False
|
|
|
|
|
|
class SimpleCPUOffloadScheduler:
|
|
"""Scheduler-side manager for CPU offloading."""
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
kv_cache_config: "KVCacheConfig | None",
|
|
cpu_capacity_bytes: int,
|
|
scheduler_block_size: int,
|
|
hash_block_size: int,
|
|
lazy_offload: bool = False,
|
|
disk_capacity_bytes: int = 0,
|
|
):
|
|
self.vllm_config = vllm_config
|
|
self.kv_cache_config = kv_cache_config
|
|
# When disk mode is active, the offload pool size is disk-based.
|
|
offload_capacity = (
|
|
disk_capacity_bytes if disk_capacity_bytes > 0 else cpu_capacity_bytes
|
|
)
|
|
self.enable_kv_cache_events = (
|
|
vllm_config.kv_events_config is not None
|
|
and vllm_config.kv_events_config.enable_kv_cache_events
|
|
)
|
|
dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
|
|
self.cp_world_size = dcp_world_size
|
|
self.block_size = scheduler_block_size
|
|
self.hash_block_size = hash_block_size
|
|
assert self.block_size % self.hash_block_size == 0
|
|
# Derive a CPU KVCacheConfig from the GPU config and build a coordinator
|
|
assert kv_cache_config is not None
|
|
self.cpu_kv_cache_config = self._derive_cpu_config(
|
|
kv_cache_config, offload_capacity
|
|
)
|
|
self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
|
|
self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU
|
|
# Find the full attention kv group for prefix cache matching.
|
|
self.fa_gidx = -1
|
|
for g_idx, g in enumerate(self.cpu_kv_cache_config.kv_cache_groups):
|
|
if isinstance(g.kv_cache_spec, FullAttentionSpec):
|
|
self.fa_gidx = g_idx
|
|
break
|
|
assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups)
|
|
# FA group's own block_size; divides scheduler_block_size (the LCM)
|
|
# but is NOT assumed to equal it.
|
|
self.fa_block_size: int = (
|
|
self.cpu_kv_cache_config.kv_cache_groups[
|
|
self.fa_gidx
|
|
].kv_cache_spec.block_size
|
|
* self.cp_world_size
|
|
)
|
|
assert self.block_size % self.fa_block_size == 0
|
|
|
|
logger.info(
|
|
"SimpleCPUOffloadScheduler: Allocating %d offload blocks "
|
|
"(%.2f GB, mode=%s, backend=%s)",
|
|
self.num_cpu_blocks,
|
|
offload_capacity / (1024**3),
|
|
"lazy" if lazy_offload else "eager",
|
|
"disk" if disk_capacity_bytes > 0 else "cpu",
|
|
)
|
|
|
|
spec_config = vllm_config.speculative_config
|
|
use_eagle_block_drop = (
|
|
spec_config is not None and spec_config.use_eagle_block_drop()
|
|
)
|
|
self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
|
|
kv_cache_config=self.cpu_kv_cache_config,
|
|
max_model_len=vllm_config.model_config.max_model_len,
|
|
max_in_flight_tokens=vllm_config.max_in_flight_tokens,
|
|
use_eagle=use_eagle_block_drop,
|
|
enable_caching=True,
|
|
enable_kv_cache_events=self.enable_kv_cache_events,
|
|
dcp_world_size=dcp_world_size,
|
|
pcp_world_size=1,
|
|
scheduler_block_size=self.block_size,
|
|
hash_block_size=self.hash_block_size,
|
|
)
|
|
self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool
|
|
# GPU block pool reference - bound after scheduler builds kv_cache_manager
|
|
self._gpu_block_pool: BlockPool | None = None
|
|
|
|
# Load metadata
|
|
self._reqs_to_load: dict[str, LoadRequestState] = {}
|
|
# Inverse map: load_event_idx -> req_ids. Keyed by load_event_idx because
|
|
# the worker reports completions by event index, not request id.
|
|
self._load_event_to_reqs: dict[int, list[str]] = {}
|
|
|
|
# Pending (cpu_hit_blocks, hit_length) tuples from find_longest_cache_hit,
|
|
# kept pinned via touch() while awaiting update_state_after_alloc().
|
|
self._pending_cpu_hits: dict[
|
|
str, tuple[tuple[list[KVCacheBlock], ...], int]
|
|
] = {}
|
|
|
|
# Store metadata
|
|
self._lazy_mode = lazy_offload
|
|
# Lazy mode: use a cursor to track the last scanned block in the GPU free queue.
|
|
self._cursor: KVCacheBlock | None = None
|
|
if self._lazy_mode:
|
|
self._target_free = self._estimate_lazy_target_blocks(
|
|
kv_cache_config,
|
|
vllm_config.scheduler_config.max_num_batched_tokens,
|
|
self.cp_world_size,
|
|
)
|
|
else:
|
|
self._target_free = 0
|
|
self._store_event_to_blocks: dict[int, TransferMeta] = {}
|
|
self._abandoned_store_event_to_blocks: dict[int, TransferMeta] = {}
|
|
# Eager mode only
|
|
self._reqs_to_store: dict[str, StoreRequestState] = {}
|
|
self._store_event_to_reqs: dict[int, list[str]] = {}
|
|
self._in_flight_store_gpu_blocks: set[int] = set()
|
|
self._pending_finished_stores: list[TransferMeta] = []
|
|
self._abandoned_reqs_to_load: dict[str, LoadRequestState] = {}
|
|
|
|
# Event counters
|
|
self._load_event_counter: int = 0
|
|
self._store_event_counter: int = 0
|
|
|
|
# For TP/PP: track partial store completions across steps.
|
|
# Events must be reported by all world_size workers before considered complete.
|
|
self._expected_worker_count = vllm_config.parallel_config.world_size
|
|
self._store_event_pending_counts: dict[int, int] = {}
|
|
|
|
@staticmethod
|
|
def _derive_cpu_config(
|
|
gpu_config: "KVCacheConfig", cpu_capacity_bytes: int
|
|
) -> "KVCacheConfig":
|
|
"""Derive a CPU KVCacheConfig from the GPU config.
|
|
Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio."""
|
|
# Import here to avoid potential circular imports
|
|
from vllm.v1.kv_cache_interface import KVCacheTensor
|
|
|
|
assert len(gpu_config.kv_cache_tensors) > 0
|
|
|
|
# Every KVCacheTensor describes placement within the same backing allocation,
|
|
# so its size is the total GPU KV cache size.
|
|
gpu_total_bytes = gpu_config.kv_cache_tensors[0].size
|
|
num_gpu_blocks = gpu_config.num_blocks
|
|
num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes)
|
|
# Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally.
|
|
cpu_tensors = [
|
|
KVCacheTensor(
|
|
size=t.size // num_gpu_blocks * num_cpu_blocks,
|
|
layers=list(t.layers),
|
|
layer_stride=t.layer_stride,
|
|
block_stride=t.block_stride,
|
|
offset=t.offset,
|
|
)
|
|
for t in gpu_config.kv_cache_tensors
|
|
]
|
|
|
|
return replace(
|
|
gpu_config,
|
|
num_blocks=num_cpu_blocks,
|
|
kv_cache_tensors=cpu_tensors,
|
|
)
|
|
|
|
@staticmethod
|
|
def _estimate_lazy_target_blocks(
|
|
kv_cache_config: "KVCacheConfig",
|
|
max_num_batched_tokens: int,
|
|
cp_world_size: int = 1,
|
|
) -> int:
|
|
"""GPU blocks to keep available (free/offloaded) per step in lazy mode."""
|
|
WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks
|
|
target = 0
|
|
for g in kv_cache_config.kv_cache_groups:
|
|
spec = g.kv_cache_spec
|
|
block_size = spec.block_size * cp_world_size
|
|
if isinstance(spec, MambaSpec):
|
|
target += 2
|
|
elif isinstance(spec, SlidingWindowSpec):
|
|
target += cdiv(spec.sliding_window, block_size) + 1
|
|
else:
|
|
target += cdiv(max_num_batched_tokens, block_size)
|
|
return int(target * (1 + WATERMARK_RATIO))
|
|
|
|
def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:
|
|
"""Bind GPU block pool so that we can touch blocks during stores.
|
|
Called by Scheduler after kv_cache_manager is ready."""
|
|
self._gpu_block_pool = gpu_block_pool
|
|
|
|
def get_num_new_matched_tokens(
|
|
self, request: "Request", num_computed_tokens: int
|
|
) -> tuple[int | None, bool]:
|
|
"""Return (num_new_tokens, is_async) from consecutive CPU cache hits."""
|
|
|
|
# Pins found CPU blocks so they survive LRU eviction until
|
|
# update_state_after_alloc() consumes them. Any pin from an earlier
|
|
# call on the same request (e.g. retry after a failed allocate_slots)
|
|
# is dropped first.
|
|
if stale := self._pending_cpu_hits.pop(request.request_id, None):
|
|
self._free_pending_cpu_hit(stale)
|
|
|
|
num_skipped_hashes = num_computed_tokens // self.hash_block_size
|
|
remaining_hashes = request.block_hashes[num_skipped_hashes:]
|
|
|
|
if not remaining_hashes:
|
|
return 0, False
|
|
# Must recompute at least the last token, matching the logic in
|
|
# kv_cache_manager.get_computed_blocks().
|
|
max_hit_len = request.num_tokens - 1 - num_computed_tokens
|
|
if max_hit_len <= 0:
|
|
return 0, False
|
|
cpu_hit_blocks, hit_length, _ = self.cpu_coordinator.find_longest_cache_hit(
|
|
remaining_hashes, max_hit_len
|
|
)
|
|
|
|
if hit_length > 0:
|
|
pin_blocks = [
|
|
blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null
|
|
]
|
|
self.cpu_block_pool.touch(pin_blocks)
|
|
self._pending_cpu_hits[request.request_id] = (
|
|
cpu_hit_blocks,
|
|
hit_length,
|
|
)
|
|
return hit_length, True
|
|
return 0, False
|
|
|
|
# TODO(yifan): this API now only matches the suffix part of the prefix cache. A more
|
|
# general API should scan blocks in both GPU and CPU block pool in a single pass.
|
|
def update_state_after_alloc(
|
|
self,
|
|
request: "Request",
|
|
blocks: "KVCacheBlocks",
|
|
num_external_tokens: int,
|
|
) -> None:
|
|
req_id = request.request_id
|
|
block_ids_by_group = blocks.get_block_ids()
|
|
num_groups = len(block_ids_by_group)
|
|
|
|
# Store tracking (eager mode only). Register the request;
|
|
# block IDs are accumulated from scheduler_output in
|
|
# _prepare_eager_store_specs via yield_req_data.
|
|
if not self._lazy_mode and req_id not in self._reqs_to_store:
|
|
self._reqs_to_store[req_id] = StoreRequestState(
|
|
request=request,
|
|
block_ids=tuple([] for _ in range(num_groups)),
|
|
num_stored_blocks=[0] * num_groups,
|
|
)
|
|
|
|
# Pop the CPU hit cached by get_num_new_matched_tokens(). The
|
|
# found blocks were pinned there to survive LRU eviction in the window
|
|
# between get_num_new_matched_tokens() and this matching call.
|
|
pending = self._pending_cpu_hits.pop(req_id, None)
|
|
|
|
if num_external_tokens == 0:
|
|
if pending is not None:
|
|
logger.warning(
|
|
"SimpleCPUOffloadScheduler: update_state_after_alloc "
|
|
"called for req_id=%s with no external tokens but "
|
|
"get_num_new_matched_tokens() unexpectedly recorded "
|
|
"a pending CPU hit; releasing the stale pin.",
|
|
req_id,
|
|
)
|
|
self._free_pending_cpu_hit(pending)
|
|
return
|
|
|
|
if pending is None:
|
|
logger.warning(
|
|
"SimpleCPUOffloadScheduler: update_state_after_alloc called "
|
|
"for req_id=%s with num_external_tokens=%d but no pending "
|
|
"CPU hit from get_num_new_matched_tokens(); skipping load.",
|
|
req_id,
|
|
num_external_tokens,
|
|
)
|
|
return
|
|
|
|
cpu_hit_blocks_full, _ = pending
|
|
|
|
# ``num_external_tokens`` is LCM-aligned (checked per-group below),
|
|
# so this counts whole scheduler-aligned chunks of incoming tokens.
|
|
num_blocks_to_load = num_external_tokens // self.block_size
|
|
assert num_blocks_to_load > 0
|
|
num_cached_fa_blocks = sum(
|
|
blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx]
|
|
)
|
|
num_computed_tokens = num_cached_fa_blocks * self.fa_block_size
|
|
|
|
# Build transfer pairs across all groups.
|
|
total_computed_tokens = num_computed_tokens + num_external_tokens
|
|
kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups
|
|
|
|
# The scheduler may have accepted fewer blocks than
|
|
# get_num_new_matched_tokens() reported.
|
|
# (e.g. due to token budget in test_partial_gpu_prefix_plus_cpu_load).
|
|
# Take only the leading N blocks per group matching num_external_tokens;
|
|
# the rest will be released along with the temp pin below.
|
|
cpu_hit_blocks: list[list[KVCacheBlock]] = []
|
|
for g in range(num_groups):
|
|
g_block_size = (
|
|
kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size
|
|
)
|
|
assert num_external_tokens % g_block_size == 0, (
|
|
f"num_external_tokens={num_external_tokens} not aligned to "
|
|
f"group {g} block_size={g_block_size}"
|
|
)
|
|
n_take_g = num_external_tokens // g_block_size
|
|
cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])
|
|
|
|
gpu_block_ids: list[int] = []
|
|
cpu_block_ids: list[int] = []
|
|
cpu_blocks_to_touch: list[KVCacheBlock] = []
|
|
|
|
for g in range(num_groups):
|
|
cpu_blocks_g = cpu_hit_blocks[g]
|
|
n_ext_g = len(cpu_blocks_g)
|
|
if n_ext_g == 0:
|
|
continue
|
|
|
|
# Number of blocks in the computed range for this group.
|
|
g_block_size = (
|
|
kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size
|
|
)
|
|
n_computed_g = cdiv(total_computed_tokens, g_block_size)
|
|
|
|
# Back-trace: ext blocks sit at the tail of the computed range.
|
|
gpu_ext_start = n_computed_g - n_ext_g
|
|
group_gpu_ids = block_ids_by_group[g]
|
|
|
|
for i, cpu_blk in enumerate(cpu_blocks_g):
|
|
# Skip null blocks (e.g. sliding window or mamba padding).
|
|
if cpu_blk.is_null:
|
|
continue
|
|
gpu_block_ids.append(group_gpu_ids[gpu_ext_start + i])
|
|
cpu_block_ids.append(cpu_blk.block_id)
|
|
cpu_blocks_to_touch.append(cpu_blk)
|
|
|
|
# Touch CPU blocks to prevent eviction during async load.
|
|
self.cpu_block_pool.touch(cpu_blocks_to_touch)
|
|
# Release the temporary pin held since get_num_new_matched_tokens().
|
|
self._free_pending_cpu_hit(pending)
|
|
|
|
# Touch GPU blocks to prevent freeing during async load
|
|
assert self._gpu_block_pool is not None
|
|
self._gpu_block_pool.touch(
|
|
[self._gpu_block_pool.blocks[bid] for bid in gpu_block_ids]
|
|
)
|
|
|
|
assert self._reqs_to_load.get(req_id) is None
|
|
self._reqs_to_load[req_id] = LoadRequestState(
|
|
request=request, transfer_meta=TransferMeta(gpu_block_ids, cpu_block_ids)
|
|
)
|
|
|
|
def build_connector_meta(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
) -> SimpleCPUOffloadMetadata:
|
|
# --- Stores ---
|
|
store_event = -1
|
|
store_gpu, store_cpu, store_req_ids, store_meta = self.prepare_store_specs(
|
|
scheduler_output
|
|
)
|
|
for transfer in self._pending_finished_stores:
|
|
store_gpu.extend(transfer.gpu_block_ids)
|
|
store_cpu.extend(transfer.cpu_block_ids)
|
|
if store_meta is not None:
|
|
assert transfer.block_meta is not None
|
|
store_meta.extend(transfer.block_meta)
|
|
self._pending_finished_stores.clear()
|
|
|
|
if store_gpu:
|
|
store_event = self._store_event_counter
|
|
self._store_event_counter += 1
|
|
self._store_event_to_blocks[store_event] = TransferMeta(
|
|
store_gpu, store_cpu, store_meta
|
|
)
|
|
if store_req_ids: # For eager mode only, track req->blocks mapping
|
|
self._store_event_to_reqs[store_event] = store_req_ids
|
|
for req_id in store_req_ids:
|
|
store_state = self._reqs_to_store.get(req_id)
|
|
if store_state is not None:
|
|
store_state.store_events.add(store_event)
|
|
|
|
# --- Loads ---
|
|
load_event = -1
|
|
load_gpu: list[int] = []
|
|
load_cpu: list[int] = []
|
|
load_req_ids: list[str] = []
|
|
for req_id, load_state in self._reqs_to_load.items():
|
|
if load_state.load_event is not None:
|
|
continue
|
|
assert load_state.transfer_meta is not None
|
|
load_gpu.extend(load_state.transfer_meta.gpu_block_ids)
|
|
load_cpu.extend(load_state.transfer_meta.cpu_block_ids)
|
|
load_req_ids.append(req_id)
|
|
if load_req_ids:
|
|
load_event = self._load_event_counter
|
|
self._load_event_counter += 1
|
|
for req_id in load_req_ids:
|
|
self._reqs_to_load[req_id].load_event = load_event
|
|
self._load_event_to_reqs[load_event] = load_req_ids
|
|
|
|
result = SimpleCPUOffloadMetadata(
|
|
load_event=load_event,
|
|
load_gpu_blocks=load_gpu,
|
|
load_cpu_blocks=load_cpu,
|
|
load_event_to_reqs={
|
|
event_idx: list(req_ids)
|
|
for event_idx, req_ids in self._load_event_to_reqs.items()
|
|
},
|
|
store_event=store_event,
|
|
store_gpu_blocks=store_gpu,
|
|
store_cpu_blocks=store_cpu,
|
|
need_flush=bool(scheduler_output.preempted_req_ids),
|
|
)
|
|
return result
|
|
|
|
def prepare_store_specs(
|
|
self, scheduler_output: SchedulerOutput
|
|
) -> tuple[
|
|
list[int],
|
|
list[int],
|
|
list[str],
|
|
list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None,
|
|
]:
|
|
"""Prepare store specs for the store event."""
|
|
if self._lazy_mode:
|
|
return self._prepare_lazy_store_specs()
|
|
else:
|
|
return self._prepare_eager_store_specs(scheduler_output)
|
|
|
|
def _prepare_lazy_store_specs(
|
|
self,
|
|
) -> tuple[
|
|
list[int],
|
|
list[int],
|
|
list[str],
|
|
list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None,
|
|
]:
|
|
"""Single-pass cursor walk: offload cached GPU blocks near eviction.
|
|
|
|
Walks the GPU free queue from the cursor, counting blocks that are
|
|
free-or-offloaded (safe for the allocator to evict). Stops when
|
|
target_free blocks are covered or CPU capacity is reached.
|
|
"""
|
|
gpu_pool = self._gpu_block_pool
|
|
if gpu_pool is None or self._target_free <= 0:
|
|
return [], [], [], None
|
|
|
|
free_queue = gpu_pool.free_block_queue
|
|
cpu_pool = self.cpu_block_pool
|
|
num_cpu_free = cpu_pool.get_num_free_blocks()
|
|
|
|
# Validate cursor: stale if block was removed from free queue.
|
|
if self._cursor is not None and self._cursor.ref_cnt > 0:
|
|
self._cursor = None
|
|
|
|
gpu_ids: list[int] = []
|
|
last_visited = self._cursor
|
|
|
|
for covered, node in enumerate(free_queue.iter_blocks_after(self._cursor)):
|
|
if covered >= self._target_free or len(gpu_ids) >= num_cpu_free:
|
|
break
|
|
|
|
last_visited = node
|
|
bhash = node.block_hash
|
|
|
|
if (
|
|
bhash is not None
|
|
and not node.is_null
|
|
and cpu_pool.cached_block_hash_to_block.get_one_block(bhash) is None
|
|
):
|
|
gpu_ids.append(node.block_id)
|
|
|
|
self._cursor = last_visited
|
|
|
|
# Batch-allocate CPU blocks.
|
|
if gpu_ids:
|
|
cpu_blocks = cpu_pool.get_new_blocks(len(gpu_ids))
|
|
cpu_ids = [blk.block_id for blk in cpu_blocks]
|
|
# Touch GPU blocks to prevent eviction during async copy.
|
|
gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids])
|
|
else:
|
|
cpu_ids = []
|
|
|
|
return gpu_ids, cpu_ids, [], None
|
|
|
|
def _prepare_eager_store_specs(
|
|
self, scheduler_output: SchedulerOutput
|
|
) -> tuple[
|
|
list[int],
|
|
list[int],
|
|
list[str],
|
|
list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None,
|
|
]:
|
|
"""Identify newly computed blocks to offload from scheduler requests.
|
|
|
|
Only considers blocks whose KV data has been **confirmed computed** by
|
|
the GPU. Blocks from the current step are stored on a later step, or by
|
|
the finish-time flush if the request completes first.
|
|
|
|
Returns:
|
|
(gpu_block_ids, cpu_block_ids, req_ids, block_meta) for the store
|
|
event. ``block_meta`` is None when kv cache events are disabled.
|
|
"""
|
|
|
|
merged_gpu_block_ids: list[int] = []
|
|
merged_cpu_block_ids: list[int] = []
|
|
req_ids: list[str] = []
|
|
merged_block_meta: list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None = (
|
|
[] if self.enable_kv_cache_events else None
|
|
)
|
|
|
|
gpu_block_pool = self._gpu_block_pool
|
|
if gpu_block_pool is None:
|
|
return [], [], [], merged_block_meta
|
|
cpu_block_pool = self.cpu_block_pool
|
|
num_groups = len(self.cpu_kv_cache_config.kv_cache_groups)
|
|
# Dedup against blocks already scheduled.
|
|
in_flight = self._in_flight_store_gpu_blocks
|
|
|
|
for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
|
|
state = self._reqs_to_store.get(req_id)
|
|
if state is None or state.finished:
|
|
continue
|
|
|
|
if preempted:
|
|
state.block_ids = tuple([] for _ in range(num_groups))
|
|
state.num_stored_blocks = [0] * num_groups
|
|
if new_block_id_groups:
|
|
for g in range(min(num_groups, len(new_block_id_groups))):
|
|
if new_block_id_groups[g] is not None:
|
|
# Accumulate new block IDs.
|
|
state.block_ids[g].extend(new_block_id_groups[g])
|
|
|
|
num_new_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0)
|
|
if num_new_tokens == 0:
|
|
continue
|
|
|
|
block_ids_by_group = state.block_ids
|
|
if not block_ids_by_group:
|
|
continue
|
|
|
|
gpu_block_ids, advanced_per_group, block_meta = (
|
|
self._select_eager_blocks_to_store(state, block_ids_by_group)
|
|
)
|
|
|
|
# --- Phase 2: Batch allocate CPU blocks ---
|
|
n_to_alloc = len(gpu_block_ids)
|
|
if n_to_alloc > 0:
|
|
cpu_blocks_alloc = cpu_block_pool.get_new_blocks(n_to_alloc)
|
|
cpu_block_ids = [blk.block_id for blk in cpu_blocks_alloc]
|
|
else:
|
|
cpu_block_ids = []
|
|
|
|
if cpu_block_ids:
|
|
req_ids.append(req_id)
|
|
merged_gpu_block_ids.extend(gpu_block_ids)
|
|
merged_cpu_block_ids.extend(cpu_block_ids)
|
|
in_flight.update(gpu_block_ids)
|
|
if merged_block_meta is not None:
|
|
assert block_meta is not None
|
|
merged_block_meta.extend(block_meta)
|
|
|
|
# Touch GPU blocks to prevent freeing during async copy
|
|
gpu_block_pool.touch(
|
|
[gpu_block_pool.blocks[bid] for bid in gpu_block_ids]
|
|
)
|
|
|
|
logger.debug(
|
|
"Request %s: Scheduling store of %d blocks to CPU (%d groups)",
|
|
req_id,
|
|
len(cpu_block_ids),
|
|
num_groups,
|
|
)
|
|
|
|
# Advance per-group cursors (includes cached hits + newly stored)
|
|
for g in range(num_groups):
|
|
state.num_stored_blocks[g] += advanced_per_group[g]
|
|
|
|
return merged_gpu_block_ids, merged_cpu_block_ids, req_ids, merged_block_meta
|
|
|
|
def _select_eager_blocks_to_store(
|
|
self,
|
|
state: StoreRequestState,
|
|
block_ids_by_group: tuple[list[int], ...],
|
|
) -> tuple[
|
|
list[int],
|
|
list[int],
|
|
list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None,
|
|
]:
|
|
"""Return confirmed eager blocks, cursor advances, and event metadata."""
|
|
assert self._gpu_block_pool is not None
|
|
gpu_block_ids: list[int] = []
|
|
advanced_per_group = [0] * len(self.cpu_kv_cache_config.kv_cache_groups)
|
|
block_meta: list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None = (
|
|
[] if self.enable_kv_cache_events else None
|
|
)
|
|
request = state.request
|
|
confirmed_tokens = request.num_computed_tokens - request.num_output_placeholders
|
|
aligned_tokens = confirmed_tokens // self.block_size * self.block_size
|
|
num_free = self.cpu_block_pool.get_num_free_blocks()
|
|
|
|
for g, group_gpu_ids in enumerate(block_ids_by_group):
|
|
if len(gpu_block_ids) >= num_free:
|
|
break
|
|
# FIXME (yifan): handle CPU cache eviction, where
|
|
# num_stored_blocks can be stale and omit evicted blocks in
|
|
# the middle of the request.
|
|
group_size = (
|
|
self.cpu_kv_cache_config.kv_cache_groups[g].kv_cache_spec.block_size
|
|
* self.cp_world_size
|
|
)
|
|
ready = min(len(group_gpu_ids), aligned_tokens // group_size)
|
|
resolved_hashes = (
|
|
resolve_block_hashes(
|
|
request.block_hashes, self.hash_block_size, group_size
|
|
)
|
|
if block_meta is not None
|
|
else None
|
|
)
|
|
curr_mm_idx = 0
|
|
secondary_mm_idx = 0
|
|
start = state.num_stored_blocks[g]
|
|
for i, gpu_block_id in enumerate(group_gpu_ids[start:ready], start=start):
|
|
gpu_block = self._gpu_block_pool.blocks[gpu_block_id]
|
|
if gpu_block.is_null or gpu_block.block_hash is None:
|
|
advanced_per_group[g] += 1
|
|
continue
|
|
if (
|
|
gpu_block_id in self._in_flight_store_gpu_blocks
|
|
or self.cpu_block_pool.cached_block_hash_to_block.get_one_block(
|
|
gpu_block.block_hash
|
|
)
|
|
is not None
|
|
):
|
|
advanced_per_group[g] += 1
|
|
continue
|
|
if len(gpu_block_ids) >= num_free:
|
|
break
|
|
gpu_block_ids.append(gpu_block_id)
|
|
advanced_per_group[g] += 1
|
|
if block_meta is not None:
|
|
token_start = i * group_size
|
|
token_end = token_start + group_size
|
|
assert resolved_hashes is not None
|
|
parent_hash = (
|
|
None
|
|
if i == 0
|
|
else maybe_convert_block_hash(resolved_hashes[i - 1])
|
|
)
|
|
lora_req = request.lora_request
|
|
extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
|
|
request, token_start, token_end, curr_mm_idx
|
|
)
|
|
meta_by_hash = {
|
|
gpu_block.block_hash: BlockStoreMeta(
|
|
token_ids=list(
|
|
request.all_token_ids[token_start:token_end]
|
|
),
|
|
parent_block_hash=parent_hash,
|
|
lora_id=lora_req.adapter_id if lora_req else None,
|
|
lora_name=lora_req.name if lora_req else None,
|
|
extra_keys=extra_keys,
|
|
)
|
|
}
|
|
first_hash_idx = token_start // self.hash_block_size
|
|
last_hash_idx = token_end // self.hash_block_size
|
|
for hash_idx in range(first_hash_idx, last_hash_idx):
|
|
block_hash = make_block_hash_with_group_id(
|
|
request.block_hashes[hash_idx], g
|
|
)
|
|
if block_hash == gpu_block.block_hash:
|
|
continue
|
|
hash_start = hash_idx * self.hash_block_size
|
|
hash_end = hash_start + self.hash_block_size
|
|
secondary_extra_keys, secondary_mm_idx = (
|
|
generate_block_hash_extra_keys(
|
|
request,
|
|
hash_start,
|
|
hash_end,
|
|
secondary_mm_idx,
|
|
)
|
|
)
|
|
meta_by_hash[block_hash] = BlockStoreMeta(
|
|
token_ids=list(request.all_token_ids[hash_start:hash_end]),
|
|
parent_block_hash=(
|
|
None
|
|
if hash_idx == 0
|
|
else maybe_convert_block_hash(
|
|
request.block_hashes[hash_idx - 1]
|
|
)
|
|
),
|
|
lora_id=lora_req.adapter_id if lora_req else None,
|
|
lora_name=lora_req.name if lora_req else None,
|
|
extra_keys=secondary_extra_keys,
|
|
block_size=self.hash_block_size,
|
|
)
|
|
block_meta.append(meta_by_hash)
|
|
|
|
return gpu_block_ids, advanced_per_group, block_meta
|
|
|
|
def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
|
|
"""Handle async transfer completions from worker.
|
|
|
|
Load completions arrive via finished_recving (real req_ids).
|
|
Store completions arrive via kv_connector_worker_meta as
|
|
per-event worker counts. We accumulate across steps and process
|
|
a store event only when all workers have reported completion.
|
|
"""
|
|
# --- Load completions ---
|
|
for req_id in list(connector_output.finished_recving or []):
|
|
self._cleanup_load_request(req_id)
|
|
|
|
# --- Store completions ---
|
|
meta = connector_output.kv_connector_worker_meta
|
|
if not isinstance(meta, SimpleCPUOffloadWorkerMetadata):
|
|
return
|
|
for event_idx, count in meta.completed_store_events.items():
|
|
total = self._store_event_pending_counts.get(event_idx, 0) + count
|
|
if total >= self._expected_worker_count:
|
|
self._store_event_pending_counts.pop(event_idx, None)
|
|
self._process_store_event(event_idx)
|
|
else:
|
|
self._store_event_pending_counts[event_idx] = total
|
|
|
|
def _process_store_event(self, event_idx: int) -> None:
|
|
"""Process a fully-completed store event."""
|
|
transfer = self._store_event_to_blocks.pop(event_idx, None)
|
|
if transfer is None:
|
|
transfer = self._abandoned_store_event_to_blocks.pop(event_idx, None)
|
|
if transfer is None:
|
|
return # guard stale events from before a reset() call
|
|
self._release_transfer_refs(transfer)
|
|
return
|
|
|
|
if not self._lazy_mode:
|
|
self._in_flight_store_gpu_blocks.difference_update(transfer.gpu_block_ids)
|
|
|
|
self._process_store_completion(
|
|
transfer.gpu_block_ids,
|
|
transfer.cpu_block_ids,
|
|
transfer.block_meta,
|
|
)
|
|
logger.debug(
|
|
"Store event %d completed: cached %d blocks to CPU",
|
|
event_idx,
|
|
len(transfer.cpu_block_ids),
|
|
)
|
|
|
|
# Eager only: update per-req state
|
|
if not self._lazy_mode:
|
|
for req_id in self._store_event_to_reqs.pop(event_idx, []):
|
|
state = self._reqs_to_store.get(req_id)
|
|
if state is None:
|
|
continue
|
|
state.store_events.discard(event_idx)
|
|
if state.finished and not state.store_events:
|
|
self._cleanup_store_request(req_id)
|
|
|
|
def _process_store_completion(
|
|
self,
|
|
gpu_block_ids: list[int],
|
|
cpu_block_ids: list[int],
|
|
block_meta: list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None = None,
|
|
) -> None:
|
|
"""Register copied blocks in the CPU prefix cache and release refs."""
|
|
assert len(cpu_block_ids) == len(gpu_block_ids)
|
|
assert block_meta is None or len(block_meta) == len(gpu_block_ids)
|
|
|
|
cpu_blocks = [self.cpu_block_pool.blocks[bid] for bid in cpu_block_ids]
|
|
|
|
assert self._gpu_block_pool is not None
|
|
for i, (gpu_block_id, cpu_block) in enumerate(zip(gpu_block_ids, cpu_blocks)):
|
|
gpu_block = self._gpu_block_pool.blocks[gpu_block_id]
|
|
primary_hash = gpu_block.block_hash
|
|
assert primary_hash is not None
|
|
self.cpu_block_pool._insert_block_hash(
|
|
primary_hash,
|
|
cpu_block,
|
|
num_tokens=gpu_block.block_hash_num_tokens,
|
|
)
|
|
secondary_hashes = self._gpu_block_pool.cached_block_hashes_by_block.get(
|
|
gpu_block_id, ()
|
|
)
|
|
for block_hash in secondary_hashes:
|
|
self.cpu_block_pool._insert_block_hash(
|
|
block_hash, cpu_block, num_tokens=None
|
|
)
|
|
if self.enable_kv_cache_events:
|
|
meta_by_hash = block_meta[i] if block_meta is not None else {}
|
|
primary_group_idx = get_group_id(primary_hash)
|
|
primary_spec = self.cpu_kv_cache_config.kv_cache_groups[
|
|
primary_group_idx
|
|
].kv_cache_spec
|
|
primary_block_size = primary_spec.block_size * (
|
|
self.cp_world_size if not isinstance(primary_spec, MambaSpec) else 1
|
|
)
|
|
events_to_emit: list[
|
|
tuple[BlockHashWithGroupId, BlockStoreMeta | None, int]
|
|
] = [
|
|
(
|
|
primary_hash,
|
|
meta_by_hash.get(primary_hash),
|
|
primary_block_size,
|
|
)
|
|
]
|
|
for secondary_hash in secondary_hashes:
|
|
secondary_meta = meta_by_hash.get(secondary_hash)
|
|
events_to_emit.append(
|
|
(
|
|
secondary_hash,
|
|
secondary_meta,
|
|
secondary_meta.block_size
|
|
if secondary_meta is not None
|
|
and secondary_meta.block_size is not None
|
|
else 0,
|
|
)
|
|
)
|
|
|
|
for block_hash, meta, event_block_size in events_to_emit:
|
|
group_idx = get_group_id(block_hash)
|
|
spec = self.cpu_kv_cache_config.kv_cache_groups[
|
|
group_idx
|
|
].kv_cache_spec
|
|
if meta is not None and len(meta.token_ids) != event_block_size:
|
|
# token_ids were sliced with a different g_block_size
|
|
# (e.g. Mamba+DCP where capture uses spec.block_size *
|
|
# cp_world_size but the event emits spec.block_size * 1).
|
|
# Emit empty metadata until #49962 fixes the sizing.
|
|
token_ids = []
|
|
parent_block_hash = None
|
|
extra_keys = None
|
|
else:
|
|
token_ids = meta.token_ids if meta else []
|
|
parent_block_hash = meta.parent_block_hash if meta else None
|
|
extra_keys = meta.extra_keys if meta else None
|
|
self.cpu_block_pool.kv_event_queue.append(
|
|
BlockStored(
|
|
block_hashes=[
|
|
maybe_convert_block_hash(get_block_hash(block_hash))
|
|
],
|
|
parent_block_hash=parent_block_hash,
|
|
token_ids=token_ids,
|
|
block_size=event_block_size,
|
|
lora_id=meta.lora_id if meta else None,
|
|
medium=self.kv_event_medium,
|
|
lora_name=meta.lora_name if meta else None,
|
|
extra_keys=(
|
|
[extra_keys] if extra_keys is not None else None
|
|
),
|
|
group_idx=group_idx,
|
|
kv_cache_spec_kind=get_kv_cache_spec_kind(spec).value,
|
|
kv_cache_spec_sliding_window=(
|
|
get_kv_cache_spec_sliding_window(spec)
|
|
),
|
|
locality="LOCAL",
|
|
)
|
|
)
|
|
|
|
# Free CPU and GPU blocks' ref counts to turn them into prefix cache
|
|
self.cpu_block_pool.free_blocks(cpu_blocks)
|
|
self._gpu_block_pool.free_blocks(
|
|
self._gpu_block_pool.blocks[bid] for bid in gpu_block_ids
|
|
)
|
|
|
|
def _release_transfer_refs(self, transfer: TransferMeta) -> None:
|
|
"""Release transfer refs without making copied data cacheable."""
|
|
cpu_blocks = [self.cpu_block_pool.blocks[bid] for bid in transfer.cpu_block_ids]
|
|
self.cpu_block_pool.free_blocks(cpu_blocks)
|
|
assert self._gpu_block_pool is not None
|
|
self._gpu_block_pool.free_blocks(
|
|
self._gpu_block_pool.blocks[bid] for bid in transfer.gpu_block_ids
|
|
)
|
|
|
|
def has_pending_stores(self) -> bool:
|
|
"""Return True if a store transfer is queued or in flight."""
|
|
return bool(
|
|
self._pending_finished_stores
|
|
or self._store_event_to_blocks
|
|
or self._abandoned_store_event_to_blocks
|
|
)
|
|
|
|
def request_finished(
|
|
self,
|
|
request: "Request",
|
|
block_ids: list[int],
|
|
) -> tuple[bool, dict[str, Any] | None]:
|
|
"""Always returns (False, None). GPU blocks are protected by ref_cnt,
|
|
so the scheduler can free blocks immediately."""
|
|
req_id = request.request_id
|
|
|
|
# Release any temp CPU hit pin from get_num_new_matched_tokens()
|
|
# if request is canceled or preempted before update_state_after_alloc()
|
|
pending = self._pending_cpu_hits.pop(req_id, None)
|
|
if pending is not None:
|
|
self._free_pending_cpu_hit(pending)
|
|
|
|
# Handle load: defer cleanup if load is in-flight
|
|
load_state = self._reqs_to_load.get(req_id)
|
|
if load_state is not None:
|
|
if load_state.load_event is not None:
|
|
load_state.finished = True # Defer: load in-flight
|
|
else:
|
|
self._cleanup_load_request(req_id)
|
|
|
|
# Handle store (eager mode only): defer cleanup if stores in-flight
|
|
if not self._lazy_mode:
|
|
store_state = self._reqs_to_store.get(req_id)
|
|
if store_state is not None:
|
|
if store_state.store_events:
|
|
store_state.finished = True # Defer: stores in-flight
|
|
else:
|
|
self._cleanup_store_request(req_id)
|
|
|
|
return False, None
|
|
|
|
def request_finished_all_groups(
|
|
self,
|
|
request: "Request",
|
|
block_ids: tuple[list[int], ...],
|
|
) -> tuple[bool, dict[str, Any] | None]:
|
|
self._queue_finished_eager_store(request, block_ids)
|
|
return self.request_finished(request, block_ids=[])
|
|
|
|
def _queue_finished_eager_store(
|
|
self, request: "Request", block_ids: tuple[list[int], ...]
|
|
) -> None:
|
|
"""Queue confirmed eager blocks omitted after a request enters decode."""
|
|
state = self._reqs_to_store.get(request.request_id)
|
|
gpu_pool = self._gpu_block_pool
|
|
if state is None or gpu_pool is None:
|
|
return
|
|
gpu_ids, _, block_meta = self._select_eager_blocks_to_store(state, block_ids)
|
|
if not gpu_ids:
|
|
return
|
|
cpu_blocks = self.cpu_block_pool.get_new_blocks(len(gpu_ids))
|
|
self._pending_finished_stores.append(
|
|
TransferMeta(gpu_ids, [b.block_id for b in cpu_blocks], block_meta)
|
|
)
|
|
self._in_flight_store_gpu_blocks.update(gpu_ids)
|
|
gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids])
|
|
|
|
def _free_pending_cpu_hit(self, pending: tuple) -> None:
|
|
"""Release the temporary CPU block pin taken in get_num_new_matched_tokens()."""
|
|
cpu_hit_blocks, _ = pending
|
|
blocks_to_free = [
|
|
blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null
|
|
]
|
|
if blocks_to_free:
|
|
self.cpu_block_pool.free_blocks(blocks_to_free)
|
|
|
|
def _cleanup_load_request(self, req_id: str) -> None:
|
|
"""Release all load resources for a request.
|
|
|
|
Shared between request_finished() and update_connector_output() paths.
|
|
Removes the request from _reqs_to_load, cleans up event mappings,
|
|
and frees CPU/GPU touch refs.
|
|
"""
|
|
state = self._reqs_to_load.pop(req_id, None)
|
|
if state is None:
|
|
state = self._abandoned_reqs_to_load.pop(req_id, None)
|
|
if state is None:
|
|
return
|
|
# Remove from load event mapping (only this req, not whole event)
|
|
if state.load_event is not None:
|
|
reqs = self._load_event_to_reqs.get(state.load_event)
|
|
if reqs is not None:
|
|
with contextlib.suppress(ValueError):
|
|
reqs.remove(req_id)
|
|
if not reqs:
|
|
self._load_event_to_reqs.pop(state.load_event, None)
|
|
|
|
if state.transfer_meta is not None:
|
|
# Free CPU touch refs
|
|
self.cpu_block_pool.free_blocks(
|
|
self.cpu_block_pool.blocks[bid]
|
|
for bid in state.transfer_meta.cpu_block_ids
|
|
)
|
|
# Free GPU touch refs
|
|
assert self._gpu_block_pool is not None
|
|
self._gpu_block_pool.free_blocks(
|
|
self._gpu_block_pool.blocks[bid]
|
|
for bid in state.transfer_meta.gpu_block_ids
|
|
)
|
|
|
|
def _cleanup_store_request(self, req_id: str) -> None:
|
|
"""Release store metadata for a request.
|
|
|
|
Metadata-only cleanup but no block freeing. Job completion handles
|
|
block caching and GPU ref freeing via _process_store_completion().
|
|
"""
|
|
state = self._reqs_to_store.pop(req_id, None)
|
|
if state is None:
|
|
return
|
|
for event_idx in list(state.store_events):
|
|
if (reqs := self._store_event_to_reqs.get(event_idx)) is not None:
|
|
with contextlib.suppress(ValueError):
|
|
reqs.remove(req_id)
|
|
if not reqs:
|
|
self._store_event_to_reqs.pop(event_idx, None)
|
|
state.store_events.clear()
|
|
|
|
def take_events(self) -> Iterable[KVCacheEvent]:
|
|
events = self.cpu_block_pool.take_events()
|
|
for event in events:
|
|
if isinstance(event, BlockRemoved):
|
|
event.medium = self.kv_event_medium
|
|
event.locality = "LOCAL"
|
|
return events
|
|
|
|
def reset(self) -> bool:
|
|
"""Abandon pending transfers and reset the CPU cache when safe.
|
|
|
|
Worker-side DMA may still be using blocks after reset is requested.
|
|
Keep those block refs pinned until the existing completion path reports
|
|
the transfer finished, then release refs without caching abandoned
|
|
store results.
|
|
"""
|
|
|
|
self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks)
|
|
for transfer in self._pending_finished_stores:
|
|
self._release_transfer_refs(transfer)
|
|
self._pending_finished_stores.clear()
|
|
|
|
self._store_event_to_blocks.clear()
|
|
self._in_flight_store_gpu_blocks.clear()
|
|
|
|
# Loads that have not been sent to the worker cannot have running DMA.
|
|
# In-flight loads stay pinned and are cleaned up on completion.
|
|
for req_id in list(self._reqs_to_load):
|
|
state = self._reqs_to_load.pop(req_id)
|
|
if state.load_event is None:
|
|
self._reqs_to_load[req_id] = state
|
|
self._cleanup_load_request(req_id)
|
|
else:
|
|
self._abandoned_reqs_to_load[req_id] = state
|
|
|
|
self._reqs_to_store.clear()
|
|
self._store_event_to_reqs.clear()
|
|
self._store_event_pending_counts = {
|
|
event_idx: count
|
|
for event_idx, count in self._store_event_pending_counts.items()
|
|
if event_idx in self._abandoned_store_event_to_blocks
|
|
}
|
|
self._cursor = None
|
|
# NOTE: _load_event_counter / _store_event_counter are not
|
|
# reset as they are monotonic and must stay ahead of the workers
|
|
# high-water marks to avoid event index collisions
|
|
|
|
if self._abandoned_store_event_to_blocks or self._abandoned_reqs_to_load:
|
|
return False
|
|
|
|
return self.cpu_block_pool.reset_prefix_cache()
|