class KVCRSecondaryTierManager(SecondaryTierManager):
"""Secondary tier wrapper around the KVCR KV P2P API."""
# Serves framework pin requests from serve_external_requests(), but only
# once per step for now: opting in would also poll the KVCR library from
# the tiering manager's polling thread, which it is not yet verified to
# tolerate.
tolerates_off_thread_polling: ClassVar[bool] = False
@classmethod
@override
def build_metric_definitions(
cls, extra_config: dict[str, Any]
) -> dict[str, OffloadingMetricMetadata]:
if not extra_config.get("enable_telemetry", False):
return {}
return _kvcr_metric_definitions()
def __init__(
self,
offloading_spec: "OffloadingSpec",
primary_kv_view: memoryview,
tier_type: str,
router_capabilities: Iterable[str] | None = None,
control_host: str = "0.0.0.0",
control_ports: list[int] | None = None,
control_advertise_host: str | None = None,
eager_ctrl_connect: bool = True,
opportunistic_query: bool = True,
enable_telemetry: bool = False,
operation_timeout_ms: int = 1000,
abandon_timeout_ms: int = 5000,
metadata_retry_interval_ms: int = 100,
secondary_g2_slots: int = 0,
kvcr_service_socket_path: str | None = None,
compatibility_digest: str | None = None,
policy: str | None = None,
g3: dict[str, Any] | None = None,
local_dram_backend: str = "UCX",
remote_fw_dram_backend: str = "UCX",
backpressure_detector: "BackpressureDetector | None" = None,
) -> None:
super().__init__(
offloading_spec, primary_kv_view, tier_type, backpressure_detector
)
selected_policy = _resolve_policy(policy)
if (kvcr_service_socket_path is None) != (compatibility_digest is None):
raise ValueError(
"kvcr_service_socket_path and compatibility_digest must be "
"configured together"
)
missing = _REQUIRED_ROUTER_CAPABILITIES - set(router_capabilities or ())
if missing:
raise ValueError(
f"KVCR secondary tier requires router capabilities: {sorted(missing)}"
)
events_enabled = offloading_spec.kv_events_config.enable_kv_cache_events
self_describing_events = (
offloading_spec.kv_events_config.self_describing_kv_events
)
if (
events_enabled
and (
secondary_g2_slots > 0
or kvcr_service_socket_path is not None
or g3 is not None
)
and not self_describing_events
):
raise ValueError(
"KVCR local G2/G3 inventory requires self_describing_kv_events"
)
self._key_adapter: KeyAdapter = _VllmKeyAdapter()
advertise_host = control_advertise_host or socket.gethostname()
dp_local_rank = offloading_spec.config.parallel.data_parallel_rank_local
if dp_local_rank is None:
if control_ports is None or len(control_ports) != 1:
raise ValueError(
"control_ports must contain exactly one port when "
"data_parallel_rank_local is unset"
)
dp_local_rank = 0
if control_ports is None or not 0 <= dp_local_rank < len(control_ports):
raise ValueError(
"control_ports must contain a port for "
f"local data-parallel rank {dp_local_rank}"
)
if any(not 1 <= port <= 65535 for port in control_ports):
raise ValueError("control_ports must be between 1 and 65535")
control = ZmqPeerControlChannel(
control_host,
control_ports[dp_local_rank],
advertise_host,
)
# Give colocated workers distinct transport listen ports.
with socket.socket() as _s:
_s.bind(("", 0))
_nixl_listen_port = _s.getsockname()[1]
if secondary_g2_slots < 0:
raise ValueError("secondary_g2_slots must be non-negative")
if (
secondary_g2_slots or kvcr_service_socket_path is not None
) and not events_enabled:
logger.warning(
"KVCR local DRAM is enabled, but KV cache events are disabled; "
"local DRAM inventory will not be published"
)
local_mapping: mmap.mmap | None = None
local_dram: LocalDramOptions | None = None
if secondary_g2_slots:
if kvcr_service_socket_path is not None:
logger.warning(
"secondary_g2_slots is ignored when "
"kvcr_service_socket_path is configured"
)
else:
local_mapping = mmap.mmap(
-1, secondary_g2_slots * self.block_size_bytes
)
local_dram = LocalDramOptions(
pools=[
(
"",
ctypes.addressof(ctypes.c_char.from_buffer(local_mapping)),
len(local_mapping),
)
],
backend=local_dram_backend,
)
guard_config = (
KVCRGuardConfig(
kvcr_service_socket_path=kvcr_service_socket_path,
guard_index=dp_local_rank,
compatibility_digest=compatibility_digest,
)
if kvcr_service_socket_path is not None and compatibility_digest is not None
else None
)
g3_config = None
if g3 is not None:
g3_values = dict(g3)
g3_values["paths"] = tuple(Path(path) for path in g3_values["paths"])
g3_config = G3Options(**g3_values)
nixl_agent_name = f"KVCR-{uuid.uuid4()}"
self._framework_pin_adapter = _FrameworkPinAdapter(self)
self._inventory_events: list[OffloadingEvent] = []
try:
self._kvcr = KVCR(
KVCRConfig(
pool_layouts=[("", self.block_size_bytes)],
enable_telemetry=enable_telemetry,
operation_timeout_ms=operation_timeout_ms,
abandon_timeout_ms=abandon_timeout_ms,
nixl_agent_name=nixl_agent_name,
nixl_listen_port=_nixl_listen_port,
),
KVCRBindings(
request_pin=self._framework_pin_adapter.request_pin,
poll_pin_results=self._framework_pin_adapter.poll_pin_results,
release_pin=self._framework_pin_adapter.release_pin,
cancel_pin_request=(self._framework_pin_adapter.cancel_pin_request),
framework_control=control,
key_adapter=self._key_adapter,
inventory_sink=(self._record_inventory if events_enabled else None),
stats_factory=(
OffloadingConnectorStats if enable_telemetry else None
),
policy=selected_policy,
),
KVCRBackendConfigs(
framework_regions=[
RegionDescriptor(
addr=ctypes.addressof(
ctypes.c_char.from_buffer(primary_kv_view)
),
size=self.block_size_bytes,
count=len(primary_kv_view),
)
],
local_dram=local_dram,
g3=g3_config,
remote_fw_dram=RemoteFWDramOptions(
eager_ctrl_connect=eager_ctrl_connect,
opportunistic_query=opportunistic_query,
metadata_retry_interval_ms=metadata_retry_interval_ms,
backend=remote_fw_dram_backend,
),
),
guard_config,
)
except BaseException as exc:
if not isinstance(exc, Exception) or isinstance(exc, KVCRStartupError):
# Native work may still touch these.
_RETAINED_NATIVE_BUFFERS.append(
ctypes.c_char.from_buffer(primary_kv_view)
)
if local_mapping is not None:
_RETAINED_NATIVE_BUFFERS.append(
ctypes.c_char.from_buffer(local_mapping)
)
raise
control.close()
if local_mapping is not None:
local_mapping.close()
raise
self._nixl_agent_name = self._kvcr.config.nixl_agent_name
self._local_dram_mmap = local_mapping
self._finished_jobs: list[JobResult] = []
self._jobs_by_op: dict[OpHandle, _JobState] = {}
@override
def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
block_key = self._key_adapter.encode(key)
status, _ = self._kvcr.query((block_key,), req_context.req_id)[0]
if status is QueryStatus.FETCHING:
return LookupResult.RETRY
if status in (QueryStatus.HIT, QueryStatus.FETCHABLE):
return LookupResult.HIT
return LookupResult.MISS
@override
def submit_load(self, job_metadata: TransferJob) -> None:
blocks = {
self._key_adapter.encode(key): [self._make_descriptor(chunk_id)]
for key, chunk_id in zip(
job_metadata.keys, job_metadata.chunk_ids.tolist(), strict=True
)
}
if not blocks:
self._finished_jobs.append(
JobResult(job_id=job_metadata.job_id, success=True)
)
return
op_handle = self._kvcr.deliver(
blocks, request_id=job_metadata.req_context.req_id
)
self._jobs_by_op[op_handle] = (
job_metadata,
len(blocks),
True,
set(),
)
@override
def submit_store(self, job_metadata: TransferJob) -> None:
blocks = {
self._key_adapter.encode(key): [self._make_descriptor(chunk_id)]
for key, chunk_id in zip(
job_metadata.keys, job_metadata.chunk_ids.tolist(), strict=True
)
}
if not blocks:
self._finished_jobs.append(
JobResult(job_id=job_metadata.job_id, success=True)
)
return
op_handle = self._kvcr.deposit(blocks)
self._jobs_by_op[op_handle] = (
job_metadata,
len(blocks),
True,
None,
)
@override
def get_finished_jobs(self) -> Iterable[JobResult]:
results = self._finished_jobs
self._finished_jobs = []
results.extend(self._poll_finished_jobs())
return results
@override
def get_stats(self) -> OffloadingConnectorStats | None:
stats = cast(OffloadingConnectorStats | None, self._kvcr.get_stats())
if stats is not None:
# Add the framework namespace only at the vLLM boundary.
for metrics in stats.metric_sections():
for name in tuple(metrics):
metrics[_vllm_metric_name(name)] = metrics.pop(name)
return stats
@override
def serve_external_requests(self, parent: ParentManager) -> None:
self._framework_pin_adapter.process_pending(parent)
def _poll_finished_jobs(self) -> list[JobResult]:
results: list[JobResult] = []
completed_prefixes: list[tuple[ReqContext, int]] = []
for op_handle, entries in self._kvcr.poll_completed():
job_state = self._jobs_by_op.get(op_handle)
if job_state is None:
continue
job, remaining, success, successful_keys = job_state
remaining -= len(entries)
success = success and all(entry.success for entry in entries.values())
if successful_keys is not None and (remaining > 0 or not success):
successful_keys.update(
OffloadKey(bytes(key))
for key, entry in entries.items()
if entry.success
)
if remaining > 0:
self._jobs_by_op[op_handle] = (
job,
remaining,
success,
successful_keys,
)
continue
self._jobs_by_op.pop(op_handle, None)
positions = job.req_context._offload_key_positions
end_token = max((positions.get(key, -1) for key in job.keys), default=-1)
# Preserve ordering across requests, which may share prefix keys.
if completed_prefixes and completed_prefixes[-1][0] is job.req_context:
end_token = max(end_token, completed_prefixes.pop()[1])
completed_prefixes.append((job.req_context, end_token))
results.append(
JobResult(
job_id=job.job_id,
success=success,
successful_keys=successful_keys if not success else None,
)
)
for req_context, end_token in completed_prefixes:
# Include earlier resident keys, even when only the tail transferred.
positions = req_context._offload_key_positions
self._align_sequence(
[key for key, position in positions.items() if position <= end_token],
req_context,
)
results.extend(self._framework_pin_adapter.take_pin_job_results())
return results
@override
def has_pending_work(self) -> bool:
# Keep scheduler polling: control messages may arrive without a load job.
return True
@override
def take_events(self) -> Iterable[OffloadingEvent]:
events = self._inventory_events
self._inventory_events = []
return events
@override
def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
if (hint := req_context.kv_hints) is not None:
self._kvcr.submit_hint(
request_id=req_context.req_id,
hints=msgspec.to_builtins(hint),
)
return RequestOffloadingContext()
@override
def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
self._align_sequence(keys, req_context, use_current_time=True)
def _align_sequence(
self,
keys: Collection[OffloadKey],
req_context: ReqContext,
use_current_time: bool = False,
) -> None:
if not keys:
return
ordered_keys = order_request_keys((list(keys),), req_context)
self._kvcr.align_sequence(
[self._key_adapter.encode(key) for key in ordered_keys],
use_current_time=use_current_time,
)
@override
def on_request_finished(self, req_context: ReqContext) -> None:
# The scheduler finalizes request recency instead of issuing touches.
computed_tokens = req_context.num_computed_tokens
in_flight = req_context.num_in_flight_tokens
prompt_tokens = req_context.num_prompt_tokens
num_tokens = req_context.num_tokens
processed_tokens = None
if (
computed_tokens is not None
and in_flight is not None
and prompt_tokens is not None
and num_tokens is not None
):
# Exclude in-flight work and the uncommitted final sampled token.
processed_tokens = min(
max(0, computed_tokens - in_flight),
max(prompt_tokens, num_tokens - 1),
)
keys = [
key
for key, position in req_context._offload_key_positions.items()
if processed_tokens is None or position <= processed_tokens
]
self.touch(keys, req_context)
self._kvcr.discard_hint(req_context.req_id)
@override
def drain_jobs(self) -> None:
while self._jobs_by_op or self._framework_pin_adapter.has_active_pins():
drained = self._poll_finished_jobs()
if drained:
self._finished_jobs.extend(drained)
else:
time.sleep(0.001)
@override
def shutdown(self) -> None:
self.drain_jobs()
# TODO: Keep registered buffers alive if KVCR.close() fails.
self._kvcr.close()
if self._local_dram_mmap is not None:
self._local_dram_mmap.close()
self._local_dram_mmap = None
def _record_inventory(self, event: InventoryEvent) -> None:
if event.tier is CacheTier.LOCAL_G2:
medium = Medium.CPU
elif event.tier is CacheTier.G3:
medium = Medium.STORAGE
else:
raise ValueError(f"unsupported KVCR inventory tier: {event.tier}")
self._inventory_events.append(
OffloadingEvent(
keys=[OffloadKey(bytes(key)) for key in event.keys],
medium=medium,
removed=event.removed,
ownership="kvcr",
removal_expected=True,
)
)
def _make_descriptor(self, chunk_id: int) -> MemoryRef:
return MemoryRef(
end_point_name=self._nixl_agent_name,
element_index=chunk_id,
)