class CPUOffloadingSpec(OffloadingSpec):
BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
@classmethod
def build_metric_definitions(
cls, extra_config: dict[str, Any]
) -> dict[str, OffloadingMetricMetadata]:
definitions: dict[str, OffloadingMetricMetadata] = {
CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC: OffloadingGaugeMetadata(
documentation=(
"Fraction of CPU KV-cache space currently pinned by active "
"transfers (0.0 = idle, 1.0 = saturated). Sustained high "
"values indicate transfers (stores or promotions) may be "
"dropped due to insufficient capacity."
),
),
CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC: OffloadingGaugeMetadata(
documentation=(
"Fraction of CPU KV-cache space currently pinned by "
"in-flight stores that have not yet "
"completed (0.0 = idle, 1.0 = saturated)."
),
),
CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC: OffloadingGaugeMetadata(
documentation=(
"Fraction of CPU KV-cache space currently pinned by "
"in-flight loads that have not yet "
"completed (0.0 = idle, 1.0 = saturated)."
),
),
CPUOffloadingMetrics.CPU_ALLOCATION_SIZE: OffloadingHistogramMetadata(
documentation=(
"Histogram of the number of CPU blocks requested by each "
"KV offload prepare_store call."
),
buckets=(1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144),
),
}
store_threshold = int(extra_config.get("store_threshold", 0))
if store_threshold >= 2:
definitions[CPUOffloadingMetrics.STORES_SKIPPED] = (
OffloadingCounterMetadata(
documentation=(
"Number of KV offload stores skipped because the reuse "
"threshold was not reached."
),
)
)
return definitions
def __init__(self, config: OffloadingConfig):
super().__init__(config)
cpu_bytes_to_use = self.extra_config.get("cpu_bytes_to_use")
if not cpu_bytes_to_use:
raise Exception(
"cpu_bytes_to_use must be specified in kv_connector_extra_config"
)
world_size = config.parallel.world_size
self.num_blocks = 0
self.kv_bytes_per_chunk = 0
self.cpu_page_size_per_worker = 0
self.replicated_layout = config.replicated_layout and self._uses_shared_region()
if config.worker_kv_bytes_per_block > 0 and world_size > 0:
num_copies = 1 if self.replicated_layout else world_size
kv_bytes_per_block = config.worker_kv_bytes_per_block * num_copies
kv_bytes_per_chunk = kv_bytes_per_block * self.blocks_per_chunk
# calculate cpu_page_size_per_worker
self.cpu_page_size_per_worker = kv_bytes_per_chunk // num_copies
# calculate num_blocks
aligned_kv_bytes_per_chunk = round_up(
kv_bytes_per_chunk, self.BLOCK_SIZE_ALIGNMENT
)
self.num_blocks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk
# Expose aligned_kv_bytes_per_chunk as
# kv_bytes_per_chunk. Note that this might contain
# some padding. i.e. each offloaded block is of the form,
# |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** |
# or |--- B0 (single copy) ---| *** maybe-pad *** |
self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk
# scheduler-side
self._manager: OffloadingManager | None = None
# worker-side
self._worker: CPUOffloadingWorker | None = None
self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru")
self.cache_policy_module_path: str | None = self.extra_config.get(
"cache_policy_module_path"
)
@override
def get_manager(self) -> OffloadingManager:
if not self._manager:
# store_threshold: how many times a block must appear in lookup()
# before it is eligible for CPU offloading. Values < 2 disable
# filtering (a threshold of 1 equals no filter; 0 is the default).
store_threshold = int(self.extra_config.get("store_threshold", 0))
# Maximum entries in the internal tracker's LRU table.
max_tracker_size = int(self.extra_config.get("max_tracker_size", 64_000))
self._manager = CPUOffloadingManager(
num_blocks=self.num_blocks,
cache_policy=self.eviction_policy,
cache_policy_module_path=self.cache_policy_module_path,
enable_events=self.kv_events_config.enable_kv_cache_events,
store_threshold=store_threshold,
max_tracker_size=max_tracker_size,
)
return self._manager
def _uses_shared_region(self) -> bool:
"""Whether the worker CPU buffer is the shared mmap region (vs a private
per-rank tensor); replicated-layout dedup is gated on this being True."""
return current_platform.is_cuda_alike()
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
mmap_region: SharedOffloadRegion | None = None
# num_blocks == 0 would size the region to zero bytes, which cannot be
# mmap'd; fall back to the tensor path (empty tensors) as before.
if self._uses_shared_region() and self.num_blocks > 0:
# Replicated layout puts all ranks on slot 0 (single MLA copy);
# otherwise each rank takes its own slot by physical device index.
if self.replicated_layout:
rank = 0
else:
world_size = self.config.parallel.world_size
rank = torch.accelerator.current_device_index() % world_size
mmap_region = SharedOffloadRegion(
engine_id=self.config.engine_id,
num_blocks=self.num_blocks,
rank=rank,
kv_bytes_per_block=self.kv_bytes_per_chunk,
cpu_page_size=self.cpu_page_size_per_worker,
)
try:
return CPUOffloadingWorker(
kv_caches=kv_caches,
blocks_per_chunk=self.blocks_per_chunk,
num_cpu_blocks=self.num_blocks,
mmap_region=mmap_region,
)
except Exception:
if mmap_region is not None:
mmap_region.cleanup()
raise
@override
def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:
if not self._worker:
if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
raise Exception(
"CPU Offloading is currently only supported on CUDA-alike "
"and XPU GPUs"
)
self._worker = self.create_worker(kv_caches)
assert self._worker is not None
return self._worker