leaveletandClaude Opus 4.8
820790b97f
The CP shared-KV cache-hit corruption bug is fixed and verified end-to-end, so drop the per-layer/per-request trace probes added during the investigation. Each probe was a Python call plus an UNCACHED os.getenv (EnvField.get() re-reads os.environ on every call) in the hot path even when disabled -- ~1000 per forward from the 11 per-layer deepseek probes alone -- plus a latent risk that an accidentally-set SGLANG_NSA_DUMP_DIR would dump tensors in-forward and tank throughput. Delete cp_hicache_trace.py and every reference to it: - deepseek_v2: 11 per-layer fwd_hash probes + the final dump-flush block - nsa_backend: the per-compose NSA tensor-dump block - cp_shared_kv_runtime: 5 imports + the if _cptrace_enabled(...) compose blocks - cache_controller / hiradix_cache / allocator / memory_pool_host: the cptrace/khash/knz/rng round-trip + lifecycle hashes - environ: the SGLANG_CP_HICACHE_KV_TRACE and SGLANG_NSA_DUMP_DIR flags Kept: SGLANG_DEBUG_CP_SHARED_KV / SGLANG_CP_TRANSFER_LOG transfer logging and the x-request-id passthrough (independent of the trace module). Verified: imports clean, 235 mem_cache + reasoning unit tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2261 lines
90 KiB
Python
2261 lines
90 KiB
Python
from __future__ import annotations
|
|
|
|
"""
|
|
Copyright 2023-2025 SGLang Team
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from queue import Empty, Full, Queue
|
|
from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Set
|
|
|
|
import torch
|
|
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
|
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
|
|
|
|
from sglang.srt.distributed import (
|
|
get_tensor_model_parallel_rank,
|
|
get_tensor_model_parallel_world_size,
|
|
)
|
|
from sglang.srt.layers.dp_attention import (
|
|
get_attention_dp_rank,
|
|
get_attention_tp_rank,
|
|
get_attention_tp_size,
|
|
is_dp_attention_enabled,
|
|
)
|
|
from sglang.srt.mem_cache.cp_shared_kv_layout import (
|
|
CpSharedKVLayout,
|
|
pad_token_locs_to_page_boundary,
|
|
)
|
|
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool
|
|
from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices
|
|
from sglang.srt.utils import get_device_module
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
device_module = get_device_module()
|
|
_CP_SHARED_KV_BS_GT1_CACHE_TIMING_COUNTS: Dict[str, int] = {}
|
|
|
|
|
|
def _cp_shared_kv_bs_gt1_cache_timing(
|
|
key: str,
|
|
start_time: float,
|
|
message: str,
|
|
*args,
|
|
) -> None:
|
|
if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
|
|
return
|
|
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
|
slow_ms = float(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS.get())
|
|
if slow_ms > 0 and elapsed_ms < slow_ms:
|
|
return
|
|
limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT.get())
|
|
count = _CP_SHARED_KV_BS_GT1_CACHE_TIMING_COUNTS.get(key, 0)
|
|
if limit > 0 and count >= limit:
|
|
return
|
|
_CP_SHARED_KV_BS_GT1_CACHE_TIMING_COUNTS[key] = count + 1
|
|
logger.info(
|
|
"[CP_SHARED_KV_BS_GT1_TIMING] event=%s elapsed_ms=%.3f " + message,
|
|
key,
|
|
elapsed_ms,
|
|
*args,
|
|
)
|
|
|
|
|
|
def _cp_shared_kv_bs_gt1_cache_timing_start() -> Optional[float]:
|
|
if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
|
|
return None
|
|
return time.perf_counter()
|
|
|
|
|
|
def _cp_shared_kv_bs_gt1_cache_timing_log(
|
|
key: str,
|
|
start_time: Optional[float],
|
|
message: str,
|
|
*args,
|
|
) -> None:
|
|
if start_time is None:
|
|
return
|
|
_cp_shared_kv_bs_gt1_cache_timing(key, start_time, message, *args)
|
|
|
|
|
|
class LayerLoadingEvent:
|
|
def __init__(self, num_layers: int):
|
|
self._num_layers = num_layers
|
|
self.load_events = [device_module.Event() for _ in range(num_layers)]
|
|
self.start_event = device_module.Event() # start event on controller stream
|
|
|
|
def complete(self, layer_index: int):
|
|
assert 0 <= layer_index < self._num_layers
|
|
self.load_events[layer_index].record()
|
|
|
|
def wait(self, layer_index: int):
|
|
device_module.current_stream().wait_event(self.load_events[layer_index])
|
|
|
|
def wait_on_stream(self, layer_index: int, stream):
|
|
assert 0 <= layer_index < self._num_layers
|
|
stream.wait_event(self.load_events[layer_index])
|
|
|
|
@property
|
|
def finish_event(self):
|
|
return self.load_events[-1]
|
|
|
|
|
|
class LayerDoneCounter:
|
|
def __init__(self, num_layers: int):
|
|
self.num_layers = num_layers
|
|
# extra producer and consumer counters for overlap mode
|
|
self.num_counters = 5
|
|
self.events = [LayerLoadingEvent(num_layers) for _ in range(self.num_counters)]
|
|
self.producer_index = -1
|
|
self.consumer_index = -1
|
|
self.consumer_indices: List[int] = []
|
|
|
|
def update_producer(self):
|
|
self.producer_index = (self.producer_index + 1) % self.num_counters
|
|
assert self.events[
|
|
self.producer_index
|
|
].finish_event.query(), (
|
|
"Producer finish event should be ready before being reused."
|
|
)
|
|
return self.producer_index
|
|
|
|
def set_consumer(self, index):
|
|
if isinstance(index, (list, tuple, set)):
|
|
self.consumer_indices = [int(i) for i in index if int(i) >= 0]
|
|
self.consumer_index = (
|
|
self.consumer_indices[0] if self.consumer_indices else -1
|
|
)
|
|
return
|
|
self.consumer_index = int(index)
|
|
self.consumer_indices = [self.consumer_index] if self.consumer_index >= 0 else []
|
|
|
|
def wait_until(self, threshold: int):
|
|
if not self.consumer_indices:
|
|
return
|
|
for consumer_index in self.consumer_indices:
|
|
self.events[consumer_index].wait(threshold)
|
|
|
|
def wait_until_on_stream(self, threshold: int, stream) -> None:
|
|
if not self.consumer_indices:
|
|
return
|
|
for consumer_index in self.consumer_indices:
|
|
self.events[consumer_index].wait_on_stream(threshold, stream)
|
|
|
|
def reset(self):
|
|
self.producer_index = -1
|
|
self.consumer_index = -1
|
|
self.consumer_indices = []
|
|
|
|
|
|
class CacheOperation:
|
|
|
|
counter = 0
|
|
|
|
def __init__(
|
|
self,
|
|
host_indices: torch.Tensor,
|
|
device_indices: torch.Tensor,
|
|
node_id: int,
|
|
priority: Optional[int] = None,
|
|
):
|
|
self.host_indices = host_indices
|
|
self.device_indices = device_indices
|
|
self.node_ids = [node_id]
|
|
self.data = None
|
|
|
|
self.id = CacheOperation.counter
|
|
CacheOperation.counter += 1
|
|
# default priority is the order of creation
|
|
self.priority = priority if priority is not None else self.id
|
|
|
|
@staticmethod
|
|
def merge_ops(ops: List[CacheOperation]) -> CacheOperation:
|
|
assert len(ops) > 0
|
|
if len(ops) == 1:
|
|
return ops[0]
|
|
|
|
host_indices = torch.cat([op.host_indices for op in ops])
|
|
device_indices = torch.cat([op.device_indices for op in ops])
|
|
node_ids = []
|
|
priority = min(op.priority for op in ops)
|
|
for op in ops:
|
|
node_ids.extend(op.node_ids)
|
|
merged_op = CacheOperation(host_indices, device_indices, -1, priority)
|
|
merged_op.node_ids = node_ids
|
|
return merged_op
|
|
|
|
def __lt__(self, other: CacheOperation):
|
|
return self.priority < other.priority
|
|
|
|
|
|
class HiCacheAck(NamedTuple):
|
|
start_event: device_module.Event
|
|
finish_event: device_module.Event
|
|
node_ids: List[int]
|
|
|
|
|
|
class TransferBuffer:
|
|
"""
|
|
Overlapping buffer preparation and transfer operations to improve throughput.
|
|
"""
|
|
|
|
def __init__(
|
|
self, stop_event, buffer_count: int = 3, max_buffer_size: int = 1024
|
|
) -> None:
|
|
self.stop_event = stop_event
|
|
self.buffers = Queue(maxsize=buffer_count)
|
|
# todo: adjust the buffer size based on throughput profile of the system
|
|
self.max_buffer_size = max_buffer_size
|
|
|
|
def full(self) -> bool:
|
|
return self.buffers.full()
|
|
|
|
def empty(self) -> bool:
|
|
return self.buffers.empty()
|
|
|
|
def put(self, item, block=True, timeout=1) -> None:
|
|
while not self.stop_event.is_set():
|
|
try:
|
|
self.buffers.put(item, block=block, timeout=timeout)
|
|
break
|
|
except Full:
|
|
if not block:
|
|
break
|
|
continue
|
|
except Exception as e:
|
|
logger.error(e)
|
|
|
|
def get(self, block=True, timeout=1) -> Optional[CacheOperation]:
|
|
try:
|
|
return self.buffers.get(block=block, timeout=timeout)
|
|
except Empty:
|
|
return None
|
|
except Exception as e:
|
|
logger.error(e)
|
|
|
|
def clear(self):
|
|
self.buffers.queue.clear()
|
|
|
|
|
|
class StorageOperation:
|
|
counter = 0
|
|
|
|
def __init__(
|
|
self,
|
|
host_indices: torch.Tensor,
|
|
token_ids: List[int],
|
|
last_hash: Optional[str] = None,
|
|
hash_value: Optional[List[str]] = None,
|
|
prefix_keys: Optional[List[str]] = None,
|
|
):
|
|
self.host_indices = host_indices
|
|
self.token_ids = token_ids
|
|
self.last_hash = last_hash
|
|
self.completed_tokens = 0
|
|
self.hash_value = hash_value if hash_value is not None else []
|
|
self.prefix_keys = prefix_keys
|
|
|
|
self.id = StorageOperation.counter
|
|
StorageOperation.counter += 1
|
|
|
|
def __lt__(self, other: "StorageOperation"):
|
|
return self.id < other.id
|
|
|
|
|
|
class PrefetchOperation(StorageOperation):
|
|
def __init__(
|
|
self,
|
|
request_id: str,
|
|
host_indices: torch.Tensor,
|
|
token_ids: List[int],
|
|
last_hash: Optional[str] = None,
|
|
prefix_keys: Optional[List[str]] = None,
|
|
):
|
|
self.request_id = request_id
|
|
|
|
self._lock = threading.Lock()
|
|
self._terminated_flag = False
|
|
self.start_time = time.monotonic()
|
|
|
|
super().__init__(host_indices, token_ids, last_hash, prefix_keys=prefix_keys)
|
|
|
|
def increment(self, num_tokens: int):
|
|
with self._lock:
|
|
if self._terminated_flag:
|
|
return False
|
|
self.completed_tokens += num_tokens
|
|
return True
|
|
|
|
def mark_terminate(self):
|
|
with self._lock:
|
|
self._terminated_flag = True
|
|
|
|
def is_terminated(self) -> bool:
|
|
return self._terminated_flag
|
|
|
|
|
|
@dataclass
|
|
class HiCacheWriteResult:
|
|
metadata: object
|
|
required_host_slots: int = 0
|
|
|
|
|
|
@dataclass
|
|
class HiCacheWriteFailure:
|
|
required_host_slots: int
|
|
metadata: object = None
|
|
# Why the reservation failed. "host_capacity" failures are retryable via
|
|
# deterministic host eviction; "undrained_ack" means the node's previous
|
|
# backup ack has not been drained yet — callers must skip this backup
|
|
# round instead of entering the capacity retry/eviction path.
|
|
reason: str = "host_capacity"
|
|
|
|
|
|
@dataclass
|
|
class HiCacheWriteReservation:
|
|
metadata: object
|
|
host_indices: torch.Tensor
|
|
physical_device_indices: torch.Tensor
|
|
node_id: int = -1
|
|
priority: Optional[int] = None
|
|
draft_host_indices: Optional[torch.Tensor] = None
|
|
required_host_slots: int = 0
|
|
|
|
|
|
@dataclass
|
|
class HiCacheLayerWriteState:
|
|
reservation: HiCacheWriteReservation
|
|
total_layers: int
|
|
start_event: object
|
|
finish_event: object
|
|
layer_events: List[object] = field(default_factory=list)
|
|
completed_target_layers: Set[int] = field(default_factory=set)
|
|
completed_draft_layers: Set[int] = field(default_factory=set)
|
|
host_indices: Optional[torch.Tensor] = None
|
|
physical_device_indices: Optional[torch.Tensor] = None
|
|
draft_host_indices: Optional[torch.Tensor] = None
|
|
ack_appended: bool = False
|
|
|
|
|
|
class HiCacheController:
|
|
|
|
def __init__(
|
|
self,
|
|
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
|
|
mem_pool_host: HostKVCache,
|
|
page_size: int,
|
|
tp_group: torch.distributed.ProcessGroup,
|
|
load_cache_event: threading.Event,
|
|
write_policy: str = "write_through_selective",
|
|
io_backend: str = "",
|
|
storage_backend: Optional[str] = None,
|
|
prefetch_threshold: int = 256,
|
|
model_name: Optional[str] = None,
|
|
storage_backend_extra_config: Optional[dict] = None,
|
|
pp_rank: int = 0,
|
|
pp_size: int = 1,
|
|
enable_storage_metrics: bool = False,
|
|
cp_shared_kv_layout: Optional[CpSharedKVLayout] = None,
|
|
draft_mem_pool_host: Optional["HostKVCache"] = None,
|
|
draft_mem_pool_device=None,
|
|
):
|
|
self.tp_group = tp_group
|
|
self.mem_pool_device_allocator = token_to_kv_pool_allocator
|
|
mem_pool_device = token_to_kv_pool_allocator.get_kvcache()
|
|
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
|
|
|
|
if isinstance(mem_pool_device, HybridLinearKVPool):
|
|
mem_pool_device = mem_pool_device.full_kv_pool
|
|
self.mem_pool_device = mem_pool_device
|
|
self.cp_shared_kv_layout = cp_shared_kv_layout
|
|
|
|
self.has_draft = False
|
|
self.mem_pool_device_draft = None
|
|
self.mem_pool_host_draft = None
|
|
self.uses_cp_hicache = cp_shared_kv_layout is not None
|
|
if self.uses_cp_hicache and not isinstance(
|
|
self.mem_pool_device, NSATokenToKVPool
|
|
):
|
|
raise ValueError(
|
|
"CP shared KV HiCache host integration requires NSATokenToKVPool."
|
|
)
|
|
self.mem_pool_host = mem_pool_host
|
|
self.write_policy = write_policy
|
|
self.page_size = page_size
|
|
self.io_backend = io_backend
|
|
self.enable_storage = False
|
|
self.storage_backend = None
|
|
self.storage_backend_type = None
|
|
self.pp_rank = pp_rank
|
|
self.pp_size = pp_size
|
|
self.enable_storage_metrics = enable_storage_metrics
|
|
|
|
# Default storage page IO functions (may be overridden by attach).
|
|
self.page_get_func = self._generic_page_get
|
|
self.page_set_func = self._generic_page_set
|
|
|
|
# Dedicated stop event for storage background threads (prefetch/backup).
|
|
# NOTE: Do NOT reuse `self.stop_event` here since it also guards core HiCache
|
|
# transfer buffers (CPU<->GPU). We want to allow runtime attach/detach of
|
|
# storage without stopping the whole controller.
|
|
self.storage_stop_event = threading.Event()
|
|
|
|
self.device = self.mem_pool_device.device
|
|
self.layer_num = self.mem_pool_device.layer_num
|
|
self.layer_done_counter = LayerDoneCounter(self.layer_num)
|
|
self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter)
|
|
if hasattr(self.mem_pool_device, "register_layer_backup_notifier"):
|
|
self.mem_pool_device.register_layer_backup_notifier(
|
|
lambda layer_id: self.on_layer_end(layer_id, source="target")
|
|
)
|
|
self.draft_mem_pool_host = None
|
|
self.draft_mem_pool_device = None
|
|
|
|
if write_policy not in [
|
|
"write_through",
|
|
"write_through_selective",
|
|
"write_back",
|
|
"write_behind",
|
|
]:
|
|
raise ValueError(f"Invalid write policy: {write_policy}")
|
|
|
|
# self.write_queue = PriorityQueue[CacheOperation]()
|
|
self.load_queue: List[CacheOperation] = []
|
|
self.write_queue: List[CacheOperation] = []
|
|
self.draft_load_queue: List[CacheOperation] = []
|
|
self.draft_write_queue: List[CacheOperation] = []
|
|
self.ack_load_queue: List[HiCacheAck] = []
|
|
self.ack_write_queue: List[HiCacheAck] = []
|
|
self.pending_layer_writes: Dict[int, HiCacheLayerWriteState] = {}
|
|
# Highest node_id ever appended to ack_write_queue. TreeNode ids are
|
|
# strictly monotonic, so a reservation whose node_id exceeds this can
|
|
# never collide with a queued ack — the common (prepare-path) case
|
|
# short-circuits to one integer compare instead of a queue scan.
|
|
self._max_write_ack_node_id: int = -1
|
|
|
|
if draft_mem_pool_host is not None or draft_mem_pool_device is not None:
|
|
self.attach_draft_pool(draft_mem_pool_device, draft_mem_pool_host)
|
|
|
|
self.stop_event = threading.Event()
|
|
self.write_buffer = TransferBuffer(self.stop_event)
|
|
self.load_buffer = TransferBuffer(
|
|
self.stop_event, buffer_count=10, max_buffer_size=100
|
|
)
|
|
|
|
self.write_stream = device_module.Stream()
|
|
self.load_stream = device_module.Stream()
|
|
|
|
# If a storage backend is provided at startup, treat it as an implicit attach,
|
|
# so init/runtime share the same lifecycle semantics and code paths.
|
|
if storage_backend is not None:
|
|
try:
|
|
self.attach_storage_backend(
|
|
storage_backend=storage_backend,
|
|
prefetch_threshold=prefetch_threshold,
|
|
model_name=model_name,
|
|
storage_backend_extra_config=storage_backend_extra_config,
|
|
)
|
|
except ValueError as e:
|
|
# Preserve the historical error shape on init for unknown backends.
|
|
raise ValueError(f"Failed to create storage backend: {e}") from e
|
|
|
|
@property
|
|
def has_draft_hicache(self) -> bool:
|
|
return (
|
|
self.draft_mem_pool_host is not None
|
|
and self.draft_mem_pool_device is not None
|
|
)
|
|
|
|
def attach_draft_pool(self, draft_mem_pool_device, draft_mem_pool_host) -> None:
|
|
if draft_mem_pool_device is None and draft_mem_pool_host is None:
|
|
return
|
|
if draft_mem_pool_device is None or draft_mem_pool_host is None:
|
|
raise ValueError(
|
|
"draft_mem_pool_device and draft_mem_pool_host must be provided together"
|
|
)
|
|
self.draft_mem_pool_device = draft_mem_pool_device
|
|
self.draft_mem_pool_host = draft_mem_pool_host
|
|
if hasattr(draft_mem_pool_device, "register_layer_transfer_counter"):
|
|
draft_mem_pool_device.register_layer_transfer_counter(
|
|
self.layer_done_counter
|
|
)
|
|
if hasattr(draft_mem_pool_device, "register_layer_backup_notifier"):
|
|
draft_mem_pool_device.register_layer_backup_notifier(
|
|
lambda layer_id: self.on_layer_end(layer_id, source="draft")
|
|
)
|
|
|
|
def on_layer_end(self, layer_id: int, source: str = "target") -> None:
|
|
if not self.pending_layer_writes:
|
|
return
|
|
if source not in ("target", "draft"):
|
|
raise ValueError(f"Unknown CP HiCache layer backup source: {source}")
|
|
self._submit_write_cp_layer_states(
|
|
list(self.pending_layer_writes.values()),
|
|
layer_id,
|
|
submit_target=(source == "target"),
|
|
submit_draft=(source == "draft"),
|
|
)
|
|
|
|
def _start_storage_threads(self):
|
|
"""Start storage prefetch/backup threads and their queues.
|
|
|
|
This is used by runtime attach, and also by reset when storage is enabled.
|
|
"""
|
|
assert self.enable_storage
|
|
assert not self.storage_stop_event.is_set()
|
|
|
|
self.prefetch_thread = threading.Thread(
|
|
target=self.prefetch_thread_func, daemon=True
|
|
)
|
|
self.backup_thread = threading.Thread(
|
|
target=self.backup_thread_func, daemon=True
|
|
)
|
|
self.prefetch_queue = Queue()
|
|
self.backup_queue = Queue()
|
|
|
|
self.prefetch_revoke_queue = Queue()
|
|
self.ack_backup_queue = Queue()
|
|
self.host_mem_release_queue = Queue()
|
|
|
|
self.prefetch_thread.start()
|
|
self.backup_thread.start()
|
|
|
|
def _stop_storage_threads(self):
|
|
"""Stop storage prefetch/backup threads and drain internal queues.
|
|
|
|
Caller should ensure no in-flight requests.
|
|
"""
|
|
# Always request stop. This is safe even when storage is already disabled,
|
|
# and makes detach truly idempotent (previous partial detach may have left
|
|
# threads alive).
|
|
# NOTE: do NOT clear stop_event unless threads have fully stopped; otherwise
|
|
# a still-alive thread may resume and touch released state.
|
|
self.storage_stop_event.set()
|
|
|
|
# Best-effort wakeups so threads exit promptly even if blocked on queues.
|
|
try:
|
|
if hasattr(self, "prefetch_queue"):
|
|
self.prefetch_queue.put_nowait(None)
|
|
if hasattr(self, "backup_queue"):
|
|
self.backup_queue.put_nowait(None)
|
|
if hasattr(self, "prefetch_buffer"):
|
|
self.prefetch_buffer.put_nowait(None)
|
|
except Exception:
|
|
pass
|
|
|
|
# Best-effort joins (threads are daemon, but join keeps state clean).
|
|
threads = []
|
|
if hasattr(self, "prefetch_thread"):
|
|
threads.append(self.prefetch_thread)
|
|
if hasattr(self, "backup_thread"):
|
|
threads.append(self.backup_thread)
|
|
if hasattr(self, "prefetch_io_aux_thread"):
|
|
threads.append(self.prefetch_io_aux_thread)
|
|
|
|
for t in threads:
|
|
try:
|
|
t.join(timeout=10)
|
|
except Exception:
|
|
pass
|
|
|
|
alive = [t for t in threads if getattr(t, "is_alive", lambda: False)()]
|
|
if alive:
|
|
logger.error(
|
|
"Failed to stop HiCache storage threads cleanly: %s",
|
|
[getattr(t, "name", repr(t)) for t in alive],
|
|
)
|
|
raise RuntimeError("Failed to stop HiCache storage threads cleanly.")
|
|
|
|
def attach_storage_backend(
|
|
self,
|
|
storage_backend: str,
|
|
prefetch_threshold: int = 256,
|
|
model_name: Optional[str] = None,
|
|
storage_backend_extra_config: Optional[dict] = None,
|
|
):
|
|
"""Attach (enable) storage backend at runtime.
|
|
|
|
Requirement: no in-flight requests. This call is expected to run on the scheduler
|
|
thread (control path), not concurrently with prefetch/backup.
|
|
"""
|
|
if self.uses_cp_hicache:
|
|
raise RuntimeError(
|
|
"CP shared KV HiCache does not support attaching a storage backend at runtime."
|
|
)
|
|
|
|
if self.enable_storage:
|
|
raise RuntimeError("Storage backend already attached.")
|
|
|
|
# Defensive: a previous partial detach may have flipped `enable_storage` but
|
|
# left background threads alive. Attaching on top of them is unsafe.
|
|
try:
|
|
self._stop_storage_threads()
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
"Cannot attach storage backend: previous detach did not stop storage threads cleanly."
|
|
) from e
|
|
|
|
# Rollback-safe init: if creation fails, keep controller state consistent
|
|
# for future attach attempts.
|
|
self.storage_backend_type = storage_backend
|
|
from sglang.srt.mem_cache.hicache_storage import get_hash_str
|
|
|
|
self.get_hash_str = get_hash_str
|
|
self.storage_config = self._generate_storage_config(
|
|
model_name, storage_backend_extra_config
|
|
)
|
|
# for MLA models, only one rank needs to backup the KV cache
|
|
self.backup_skip = (
|
|
self.storage_config.is_mla_model
|
|
# todo: load balancing
|
|
and self.storage_config.tp_rank != 0
|
|
)
|
|
|
|
# Use storage backend factory for dynamic backend creation
|
|
from sglang.srt.mem_cache.storage import StorageBackendFactory
|
|
|
|
try:
|
|
self.storage_backend = StorageBackendFactory.create_backend(
|
|
storage_backend, self.storage_config, self.mem_pool_host
|
|
)
|
|
self.storage_backend.register_mem_pool_host(self.mem_pool_host)
|
|
|
|
self.enable_storage = True
|
|
# todo: threshold policy for prefetching
|
|
self.prefetch_threshold = max(prefetch_threshold, self.page_size)
|
|
self.prefetch_capacity_limit = max(
|
|
0, int(0.8 * (self.mem_pool_host.size - self.mem_pool_device.size))
|
|
)
|
|
# granularity of batch storage IO operations, in number of pages
|
|
self.storage_batch_size = 128
|
|
# tracking the number of tokens locked in prefetching, updated by the main scheduler thread
|
|
self.prefetch_tokens_occupied = 0
|
|
|
|
# create a new communication group for synchronizing storage operations across TP workers
|
|
self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
|
|
if self.tp_world_size > 1:
|
|
from sglang.srt.distributed.parallel_state import (
|
|
create_custom_parallel_group,
|
|
)
|
|
|
|
group_ranks = torch.distributed.get_process_group_ranks(self.tp_group)
|
|
self.prefetch_tp_group = create_custom_parallel_group(
|
|
group_ranks=group_ranks, backend="gloo"
|
|
)
|
|
|
|
# Select the get and set functions
|
|
self.page_get_func = self._generic_page_get
|
|
self.page_set_func = self._generic_page_set
|
|
|
|
if (self.storage_backend_type in ["hf3fs", "mooncake", "eic", "nixl"]) or (
|
|
self.storage_backend_type == "dynamic"
|
|
and bool(self.storage_config.extra_config.get("interface_v1", 0))
|
|
):
|
|
self.page_get_func = self._page_get_zero_copy
|
|
self.page_set_func = self._page_set_zero_copy
|
|
|
|
# Ensure stop_event is clear before starting threads.
|
|
self.storage_stop_event.clear()
|
|
self._start_storage_threads()
|
|
except Exception:
|
|
# Best-effort cleanup for partial init.
|
|
try:
|
|
self._stop_storage_threads()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
if hasattr(self, "prefetch_tp_group"):
|
|
try:
|
|
torch.distributed.destroy_process_group(self.prefetch_tp_group)
|
|
except Exception:
|
|
pass
|
|
self.prefetch_tp_group = None
|
|
except Exception:
|
|
pass
|
|
try:
|
|
if (
|
|
hasattr(self, "storage_backend")
|
|
and self.storage_backend is not None
|
|
):
|
|
if hasattr(self.storage_backend, "close"):
|
|
self.storage_backend.close()
|
|
except Exception:
|
|
pass
|
|
self.storage_backend = None
|
|
self.storage_backend_type = None
|
|
self.enable_storage = False
|
|
self.page_get_func = self._generic_page_get
|
|
self.page_set_func = self._generic_page_set
|
|
raise
|
|
|
|
def detach_storage_backend(self):
|
|
"""Detach (disable) storage backend at runtime.
|
|
|
|
Requirement: no in-flight requests. This will stop storage threads and release
|
|
the backend instance (best-effort close).
|
|
"""
|
|
# Idempotent cleanup: even if `enable_storage` is already False,
|
|
# we may still have leftover resources (threads/backend/process group) from a
|
|
# previous partial detach. We attempt cleanup whenever possible.
|
|
try:
|
|
self._stop_storage_threads()
|
|
except Exception as e:
|
|
# Do not proceed tearing down backend/process group if threads are not
|
|
# fully stopped; otherwise still-alive threads may touch released state.
|
|
# Caller can retry detach.
|
|
logger.exception("Stop storage threads failed: %s", e)
|
|
# IMPORTANT: Do not silently succeed. Upper layers rely on exceptions here
|
|
# to avoid flipping `enable_storage` flags while threads are still alive.
|
|
raise RuntimeError("Stop storage threads failed; detach aborted.") from e
|
|
|
|
# Best-effort destroy process group created for storage ops.
|
|
try:
|
|
if (
|
|
hasattr(self, "prefetch_tp_group")
|
|
and self.prefetch_tp_group is not None
|
|
):
|
|
try:
|
|
torch.distributed.destroy_process_group(self.prefetch_tp_group)
|
|
except Exception:
|
|
pass
|
|
self.prefetch_tp_group = None
|
|
except Exception:
|
|
pass
|
|
|
|
# Best-effort close (some backends rely on GC/destructor).
|
|
try:
|
|
if (
|
|
hasattr(self, "storage_backend")
|
|
and self.storage_backend is not None
|
|
and hasattr(self.storage_backend, "close")
|
|
):
|
|
self.storage_backend.close()
|
|
except Exception:
|
|
logger.exception("Failed to close storage backend cleanly.")
|
|
|
|
self.storage_backend = None
|
|
self.storage_backend_type = None
|
|
self.enable_storage = False
|
|
self.page_get_func = self._generic_page_get
|
|
self.page_set_func = self._generic_page_set
|
|
# Now it's safe to clear the stop event for future re-attach.
|
|
self.storage_stop_event.clear()
|
|
|
|
def _generate_storage_config(
|
|
self,
|
|
model_name: Optional[str] = None,
|
|
storage_backend_extra_config: Optional[dict] = None,
|
|
):
|
|
if storage_backend_extra_config is None:
|
|
storage_backend_extra_config = {}
|
|
|
|
if is_dp_attention_enabled():
|
|
self.tp_rank = get_attention_tp_rank()
|
|
self.tp_size = get_attention_tp_size()
|
|
self.dp_rank = get_attention_dp_rank()
|
|
else:
|
|
self.tp_rank = get_tensor_model_parallel_rank()
|
|
self.tp_size = get_tensor_model_parallel_world_size()
|
|
self.dp_rank = 0
|
|
|
|
# Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool.
|
|
is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool)
|
|
# Least Common Multiple among heterogeneous tp size
|
|
tp_lcm_size = storage_backend_extra_config.pop("tp_lcm_size", None)
|
|
should_split_heads = False
|
|
|
|
if tp_lcm_size:
|
|
assert (
|
|
tp_lcm_size % self.tp_size == 0
|
|
), "tp_lcm_size must be divisible by tp_size."
|
|
should_split_heads = (
|
|
not is_mla_backend
|
|
and self.mem_pool_host.layout == "page_head"
|
|
and tp_lcm_size > self.tp_size
|
|
)
|
|
|
|
return HiCacheStorageConfig(
|
|
tp_rank=self.tp_rank,
|
|
tp_size=self.tp_size,
|
|
pp_rank=self.pp_rank,
|
|
pp_size=self.pp_size,
|
|
is_mla_model=is_mla_backend,
|
|
enable_storage_metrics=self.enable_storage_metrics,
|
|
is_page_first_layout=self.mem_pool_host.layout == "page_first",
|
|
model_name=model_name,
|
|
tp_lcm_size=tp_lcm_size,
|
|
should_split_heads=should_split_heads,
|
|
extra_config=storage_backend_extra_config,
|
|
)
|
|
|
|
def reset(self):
|
|
self.stop_event.set()
|
|
self.storage_stop_event.set()
|
|
|
|
self.write_queue.clear()
|
|
self.load_queue.clear()
|
|
self.draft_write_queue.clear()
|
|
self.draft_load_queue.clear()
|
|
self.write_buffer.clear()
|
|
self.load_buffer.clear()
|
|
self.ack_write_queue.clear()
|
|
self.ack_load_queue.clear()
|
|
self.pending_layer_writes.clear()
|
|
# HiRadixCache.reset() restarts TreeNode.counter from 0; with the queue
|
|
# cleared, resetting the max keeps the duplicate-reservation fast path
|
|
# active for the re-minted ids instead of falling back to queue scans.
|
|
self._max_write_ack_node_id = -1
|
|
if self.enable_storage:
|
|
self.prefetch_thread.join()
|
|
self.backup_thread.join()
|
|
self.prefetch_queue.queue.clear()
|
|
self.backup_queue.queue.clear()
|
|
self.prefetch_revoke_queue.queue.clear()
|
|
self.ack_backup_queue.queue.clear()
|
|
|
|
self.stop_event.clear()
|
|
self.storage_stop_event.clear()
|
|
|
|
if self.enable_storage:
|
|
self.prefetch_thread = threading.Thread(
|
|
target=self.prefetch_thread_func, daemon=True
|
|
)
|
|
self.backup_thread = threading.Thread(
|
|
target=self.backup_thread_func, daemon=True
|
|
)
|
|
self.prefetch_thread.start()
|
|
self.backup_thread.start()
|
|
|
|
def clear_draft_host_pool(self) -> None:
|
|
if self.draft_mem_pool_host is not None:
|
|
self.draft_mem_pool_host.clear()
|
|
|
|
def write(
|
|
self,
|
|
device_indices: torch.Tensor,
|
|
priority: Optional[int] = None,
|
|
node_id: int = -1,
|
|
) -> Optional[torch.Tensor | HiCacheWriteResult | HiCacheWriteFailure]:
|
|
"""
|
|
Back up KV caches from device memory to host memory.
|
|
"""
|
|
logger.debug(
|
|
"[CacheCtrl-write] write: node_id=%d len=%d cp=%s",
|
|
node_id,
|
|
len(device_indices),
|
|
self.uses_cp_hicache,
|
|
)
|
|
if self.uses_cp_hicache:
|
|
return self._write_cp(device_indices, priority, node_id)
|
|
|
|
host_indices = self.mem_pool_host.alloc(len(device_indices))
|
|
if host_indices is None:
|
|
logger.debug(
|
|
"[CacheCtrl-write] write non-CP FAILED (host full): node_id=%d len=%d",
|
|
node_id,
|
|
len(device_indices),
|
|
)
|
|
return None
|
|
self.write_queue.append(
|
|
CacheOperation(host_indices, device_indices, node_id, priority)
|
|
)
|
|
if self.has_draft and not self.uses_cp_hicache:
|
|
self.draft_write_queue.append(
|
|
CacheOperation(host_indices, device_indices, node_id, priority)
|
|
)
|
|
self.start_writing()
|
|
logger.debug(
|
|
"[CacheCtrl-write] write non-CP submitted: node_id=%d len=%d",
|
|
node_id,
|
|
len(device_indices),
|
|
)
|
|
return host_indices
|
|
|
|
def start_writing(self) -> None:
|
|
if len(self.write_queue) == 0 and len(self.draft_write_queue) == 0:
|
|
return
|
|
|
|
op = CacheOperation.merge_ops(self.write_queue) if self.write_queue else None
|
|
draft_op = (
|
|
CacheOperation.merge_ops(self.draft_write_queue)
|
|
if self.draft_write_queue
|
|
else None
|
|
)
|
|
node_ids = op.node_ids if op is not None else draft_op.node_ids
|
|
if op is not None:
|
|
host_indices, device_indices = self.move_indices(op, self.mem_pool_host)
|
|
else:
|
|
host_indices = device_indices = None
|
|
if draft_op is not None:
|
|
draft_host_indices, draft_device_indices = self.move_indices(
|
|
draft_op, self.draft_mem_pool_host
|
|
)
|
|
else:
|
|
draft_host_indices = draft_device_indices = None
|
|
self.write_queue.clear()
|
|
self.draft_write_queue.clear()
|
|
|
|
start_event = device_module.Event()
|
|
finish_event = device_module.Event()
|
|
|
|
start_event.record()
|
|
with device_module.stream(self.write_stream):
|
|
start_event.wait(self.write_stream)
|
|
if op is not None:
|
|
self.mem_pool_host.backup_from_device_all_layer(
|
|
self.mem_pool_device, host_indices, device_indices, self.io_backend
|
|
)
|
|
if draft_op is not None:
|
|
self.draft_mem_pool_host.backup_from_device_all_layer(
|
|
self.draft_mem_pool_device,
|
|
draft_host_indices,
|
|
draft_device_indices,
|
|
self.io_backend,
|
|
)
|
|
finish_event.record()
|
|
# NOTE: We must save the host indices and device indices here,
|
|
# this is because we need to guarantee that these tensors are
|
|
# still alive when the write stream is executing.
|
|
if host_indices is not None and host_indices.is_cuda:
|
|
host_indices.record_stream(self.write_stream)
|
|
if device_indices is not None and device_indices.is_cuda:
|
|
device_indices.record_stream(self.write_stream)
|
|
if draft_host_indices is not None and draft_host_indices.is_cuda:
|
|
draft_host_indices.record_stream(self.write_stream)
|
|
if draft_device_indices is not None and draft_device_indices.is_cuda:
|
|
draft_device_indices.record_stream(self.write_stream)
|
|
|
|
self._append_write_ack(HiCacheAck(start_event, finish_event, node_ids))
|
|
|
|
def _append_write_ack(self, ack: HiCacheAck) -> None:
|
|
"""Single funnel for ack_write_queue appends.
|
|
|
|
Keeps _max_write_ack_node_id (the node_has_undrained_write_ack fast
|
|
path) covering every enqueued node id; appending around this helper
|
|
would re-open false negatives in the duplicate-reservation gate.
|
|
"""
|
|
self.ack_write_queue.append(ack)
|
|
for node_id in ack.node_ids:
|
|
if node_id > self._max_write_ack_node_id:
|
|
self._max_write_ack_node_id = node_id
|
|
|
|
def _append_completed_write_ack(self, node_id: int) -> None:
|
|
event = device_module.Event()
|
|
event.record()
|
|
self._append_write_ack(HiCacheAck(event, event, [node_id]))
|
|
|
|
def _append_completed_load_ack(self, node_id: int) -> None:
|
|
event = device_module.Event()
|
|
event.record()
|
|
self.ack_load_queue.append(HiCacheAck(event, event, [node_id]))
|
|
|
|
def _validate_cp_hicache_page_indices(
|
|
self,
|
|
host_indices: torch.Tensor,
|
|
device_indices: torch.Tensor,
|
|
) -> None:
|
|
"""Validate page-shaped CP HiCache indices without CUDA host sync.
|
|
|
|
The production invariant is construction-based: HostKVCache alloc/free
|
|
preserves page-shaped host spans, and CpSharedKVLayout preserves
|
|
per-page contiguity when mapping logical locs to physical locs. The
|
|
generic validator uses Tensor truth values and would synchronize CUDA
|
|
tensors, so keep the hot path sync-free and validate CPU/fake-test
|
|
tensors only.
|
|
"""
|
|
|
|
if not host_indices.is_cuda:
|
|
validate_page_aligned_token_indices(
|
|
host_indices, self.page_size, "host_indices"
|
|
)
|
|
if not device_indices.is_cuda:
|
|
validate_page_aligned_token_indices(
|
|
device_indices, self.page_size, "physical_device_indices"
|
|
)
|
|
|
|
def node_has_undrained_write_ack(self, node_id: int) -> bool:
|
|
"""Whether a final write ack for this node is still in ack_write_queue.
|
|
|
|
The radix side drains acks via writing_check; until then the node's
|
|
previous backup owns radix/host state. Creating a second layer-write
|
|
state for the same node while its ack is undrained would put two acks
|
|
for one radix registration into the queue (one inc_lock_ref, two
|
|
completions) — the invariant is at most one in-flight ack per node.
|
|
|
|
Cost: node ids are strictly monotonic, so fresh ids (the per-request
|
|
prepare path) exit on the integer compare below; the queue scan runs
|
|
only for re-reservations of old node ids (rare write_backup fallback).
|
|
"""
|
|
if node_id > getattr(self, "_max_write_ack_node_id", -1):
|
|
return False
|
|
ack_write_queue = getattr(self, "ack_write_queue", None)
|
|
if not ack_write_queue:
|
|
return False
|
|
return any(node_id in ack.node_ids for ack in ack_write_queue)
|
|
|
|
def cancel_layer_write_state(self, node_id: int) -> bool:
|
|
"""Drop a pending per-layer write state during backup rollback.
|
|
|
|
Without this, an exception between reservation and final ack leaves a
|
|
half-submitted state in pending_layer_writes: the next forward's layer
|
|
hooks keep driving it, eventually writing D2H into host slots the
|
|
rollback already evicted and appending an ack for a node the radix
|
|
side no longer has registered. Synchronize the write stream so any
|
|
already-issued per-layer copies finish before the caller frees the
|
|
host slots.
|
|
"""
|
|
state = self.pending_layer_writes.pop(node_id, None)
|
|
if state is None:
|
|
return False
|
|
self.write_stream.synchronize()
|
|
return True
|
|
|
|
def reserve_write_cp(
|
|
self,
|
|
device_indices: torch.Tensor,
|
|
priority: Optional[int] = None,
|
|
node_id: int = -1,
|
|
) -> HiCacheWriteReservation | HiCacheWriteFailure:
|
|
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
|
|
|
|
if node_id >= 0 and self.node_has_undrained_write_ack(node_id):
|
|
# Refusing here (instead of crashing later in writing_check) keeps
|
|
# the failure on the existing HiCacheWriteFailure path: the caller
|
|
# skips this backup round and may retry after the ack drains.
|
|
logger.warning(
|
|
"[CacheCtrl-write] reserve_write_cp refused: node_id=%d still "
|
|
"has an undrained write ack; skipping duplicate backup",
|
|
node_id,
|
|
)
|
|
return HiCacheWriteFailure(required_host_slots=0, reason="undrained_ack")
|
|
|
|
page_size = self.page_size
|
|
layout = self.cp_shared_kv_layout
|
|
logical_len = len(device_indices)
|
|
padded_device_indices = pad_token_locs_to_page_boundary(
|
|
device_indices,
|
|
page_size,
|
|
name="CP HiCache write device_indices",
|
|
)
|
|
padded_len = int(padded_device_indices.numel())
|
|
owned_mask = layout.owned_by_this_rank(padded_device_indices)
|
|
owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu()
|
|
# Capture the global owner pattern (one int8 per logical PAGE; identical
|
|
# on all CP ranks since it's a pure function of the logical page ids).
|
|
# load_cp will replay this pattern via alloc_pages_with_owners so the
|
|
# saved owned_positions correctly index the new allocation.
|
|
page_first_locs = padded_device_indices[::page_size]
|
|
logical_pages = torch.div(page_first_locs, page_size, rounding_mode="floor")
|
|
page_owners = layout.owner_for_logical_pages(logical_pages).to(
|
|
dtype=torch.int8, device="cpu"
|
|
)
|
|
if owned_positions.numel() == 0:
|
|
logger.debug(
|
|
"[CacheCtrl-write] reserve_write_cp zero-owned rank: node_id=%d logical_len=%d",
|
|
node_id,
|
|
logical_len,
|
|
)
|
|
empty = torch.empty((0,), dtype=torch.int64)
|
|
return HiCacheWriteReservation(
|
|
metadata=CpHiCacheNodeMetadata(
|
|
logical_len=logical_len,
|
|
padded_len=padded_len,
|
|
owned_positions=owned_positions,
|
|
host_indices=empty,
|
|
page_owners=page_owners,
|
|
page_size=page_size,
|
|
draft_host_indices=(empty.clone() if self.has_draft_hicache else None),
|
|
),
|
|
host_indices=empty,
|
|
physical_device_indices=empty,
|
|
node_id=node_id,
|
|
priority=priority,
|
|
draft_host_indices=(empty.clone() if self.has_draft_hicache else None),
|
|
)
|
|
|
|
owned_logical_indices = padded_device_indices[owned_mask]
|
|
physical_device_indices = layout.logical_locs_to_physical(owned_logical_indices)
|
|
host_alloc = getattr(
|
|
self.mem_pool_host,
|
|
"alloc_contiguous_preferred",
|
|
self.mem_pool_host.alloc,
|
|
)
|
|
host_indices = host_alloc(len(physical_device_indices))
|
|
if host_indices is None:
|
|
logger.debug(
|
|
"[CacheCtrl-write] reserve_write_cp FAILED (host full): node_id=%d logical_len=%d owned=%d",
|
|
node_id,
|
|
logical_len,
|
|
owned_positions.numel(),
|
|
)
|
|
return HiCacheWriteFailure(required_host_slots=len(physical_device_indices))
|
|
|
|
draft_host_indices = None
|
|
if self.has_draft_hicache:
|
|
draft_host_alloc = getattr(
|
|
self.draft_mem_pool_host,
|
|
"alloc_contiguous_preferred",
|
|
self.draft_mem_pool_host.alloc,
|
|
)
|
|
draft_host_indices = draft_host_alloc(len(physical_device_indices))
|
|
if draft_host_indices is None:
|
|
self.mem_pool_host.free(host_indices)
|
|
logger.debug(
|
|
"[CacheCtrl-write] reserve_write_cp FAILED (draft host full): node_id=%d logical_len=%d owned=%d",
|
|
node_id,
|
|
logical_len,
|
|
owned_positions.numel(),
|
|
)
|
|
return HiCacheWriteFailure(
|
|
required_host_slots=len(physical_device_indices)
|
|
)
|
|
|
|
try:
|
|
self._validate_cp_hicache_page_indices(
|
|
host_indices, physical_device_indices
|
|
)
|
|
if draft_host_indices is not None:
|
|
self._validate_cp_hicache_page_indices(
|
|
draft_host_indices, physical_device_indices
|
|
)
|
|
except Exception:
|
|
self.mem_pool_host.free(host_indices)
|
|
if draft_host_indices is not None:
|
|
self.draft_mem_pool_host.free(draft_host_indices)
|
|
raise
|
|
|
|
return HiCacheWriteReservation(
|
|
metadata=CpHiCacheNodeMetadata(
|
|
logical_len=logical_len,
|
|
padded_len=padded_len,
|
|
owned_positions=owned_positions,
|
|
host_indices=host_indices.cpu(),
|
|
page_owners=page_owners,
|
|
page_size=page_size,
|
|
draft_host_indices=(
|
|
draft_host_indices.cpu() if draft_host_indices is not None else None
|
|
),
|
|
),
|
|
host_indices=host_indices,
|
|
physical_device_indices=physical_device_indices,
|
|
node_id=node_id,
|
|
priority=priority,
|
|
draft_host_indices=draft_host_indices,
|
|
)
|
|
|
|
def submit_write_cp_all_layer(self, reservation: HiCacheWriteReservation) -> None:
|
|
logger.warning(
|
|
"[CP_HICACHE_FALLBACK][submit_write_cp_all_layer] "
|
|
"CP HiCache all-layer backup fallback. "
|
|
"node_id=%d logical_len=%d owned=%d physical=%d draft=%s",
|
|
reservation.node_id,
|
|
reservation.metadata.logical_len,
|
|
reservation.metadata.owned_positions.numel(),
|
|
len(reservation.physical_device_indices),
|
|
reservation.draft_host_indices is not None,
|
|
)
|
|
if len(reservation.physical_device_indices) == 0:
|
|
self._append_completed_write_ack(reservation.node_id)
|
|
logger.debug(
|
|
"[CacheCtrl-write] submit_write_cp_all_layer zero-owned ack: node_id=%d logical_len=%d",
|
|
reservation.node_id,
|
|
reservation.metadata.logical_len,
|
|
)
|
|
return
|
|
|
|
self.write_queue.append(
|
|
CacheOperation(
|
|
reservation.host_indices,
|
|
reservation.physical_device_indices,
|
|
reservation.node_id,
|
|
reservation.priority,
|
|
)
|
|
)
|
|
if reservation.draft_host_indices is not None:
|
|
self.draft_write_queue.append(
|
|
CacheOperation(
|
|
reservation.draft_host_indices,
|
|
reservation.physical_device_indices,
|
|
reservation.node_id,
|
|
reservation.priority,
|
|
)
|
|
)
|
|
self.start_writing()
|
|
logger.debug(
|
|
"[CacheCtrl-write] submit_write_cp_all_layer submitted: node_id=%d logical_len=%d owned=%d physical=%d draft=%s",
|
|
reservation.node_id,
|
|
reservation.metadata.logical_len,
|
|
reservation.metadata.owned_positions.numel(),
|
|
len(reservation.physical_device_indices),
|
|
reservation.draft_host_indices is not None,
|
|
)
|
|
|
|
def _get_or_create_layer_write_state(
|
|
self, reservation: HiCacheWriteReservation
|
|
) -> HiCacheLayerWriteState:
|
|
state = self.pending_layer_writes.get(reservation.node_id)
|
|
if state is not None:
|
|
if state.reservation is not reservation:
|
|
raise RuntimeError(
|
|
f"Conflicting CP HiCache per-layer backup reservation for "
|
|
f"node_id={reservation.node_id}"
|
|
)
|
|
return state
|
|
|
|
draft_layer_num = (
|
|
self.draft_mem_pool_device.layer_num
|
|
if reservation.draft_host_indices is not None
|
|
else 0
|
|
)
|
|
total_layers = max(self.layer_num, draft_layer_num)
|
|
if total_layers <= 0:
|
|
raise RuntimeError(
|
|
f"Invalid CP HiCache per-layer backup layer count: {total_layers}"
|
|
)
|
|
|
|
state = HiCacheLayerWriteState(
|
|
reservation=reservation,
|
|
total_layers=total_layers,
|
|
start_event=device_module.Event(),
|
|
finish_event=device_module.Event(),
|
|
)
|
|
if len(reservation.physical_device_indices) > 0:
|
|
op = CacheOperation(
|
|
reservation.host_indices,
|
|
reservation.physical_device_indices,
|
|
reservation.node_id,
|
|
reservation.priority,
|
|
)
|
|
state.host_indices, state.physical_device_indices = self.move_indices(
|
|
op, self.mem_pool_host
|
|
)
|
|
if reservation.draft_host_indices is not None:
|
|
draft_op = CacheOperation(
|
|
reservation.draft_host_indices,
|
|
reservation.physical_device_indices,
|
|
reservation.node_id,
|
|
reservation.priority,
|
|
)
|
|
state.draft_host_indices, _ = self.move_indices(
|
|
draft_op, self.draft_mem_pool_host
|
|
)
|
|
self.pending_layer_writes[reservation.node_id] = state
|
|
state.start_event.record()
|
|
return state
|
|
|
|
def _layer_write_state_done(self, state: HiCacheLayerWriteState) -> bool:
|
|
target_done = len(state.completed_target_layers) >= self.layer_num
|
|
draft_done = (
|
|
state.draft_host_indices is None
|
|
or len(state.completed_draft_layers)
|
|
>= self.draft_mem_pool_device.layer_num
|
|
)
|
|
return target_done and draft_done
|
|
|
|
@staticmethod
|
|
def _record_tensor_on_stream(tensor: Optional[torch.Tensor], stream) -> None:
|
|
if tensor is not None and tensor.is_cuda:
|
|
tensor.record_stream(stream)
|
|
|
|
@staticmethod
|
|
def _concat_layer_write_tensors(tensors: List[torch.Tensor]) -> torch.Tensor:
|
|
if len(tensors) == 1:
|
|
return tensors[0]
|
|
return torch.cat(tensors)
|
|
|
|
def _append_layer_write_ack(self, state: HiCacheLayerWriteState) -> None:
|
|
if state.ack_appended:
|
|
return
|
|
state.ack_appended = True
|
|
reservation = state.reservation
|
|
self.pending_layer_writes.pop(reservation.node_id, None)
|
|
self._append_write_ack(
|
|
HiCacheAck(state.start_event, state.finish_event, [reservation.node_id])
|
|
)
|
|
logger.debug(
|
|
"[CacheCtrl-write] submit_write_cp_layer final ack: node_id=%d logical_len=%d layers=%d draft=%s",
|
|
reservation.node_id,
|
|
reservation.metadata.logical_len,
|
|
state.total_layers,
|
|
reservation.draft_host_indices is not None,
|
|
)
|
|
|
|
def _submit_write_cp_layer_states(
|
|
self,
|
|
states: List[HiCacheLayerWriteState],
|
|
layer_id: int,
|
|
*,
|
|
submit_target: bool = True,
|
|
submit_draft: bool = True,
|
|
) -> None:
|
|
"""Submit one layer for a group of reserved CP host backups.
|
|
|
|
Radix/HiCache metadata remains per node. Only the transfer descriptor is
|
|
grouped so bs>1 batches issue one target D2H and one draft D2H per layer
|
|
instead of one D2H per request per layer.
|
|
"""
|
|
|
|
if not states:
|
|
return
|
|
|
|
unique_states: List[HiCacheLayerWriteState] = []
|
|
seen_state_ids = set()
|
|
for state in states:
|
|
state_id = id(state)
|
|
if state_id in seen_state_ids:
|
|
continue
|
|
seen_state_ids.add(state_id)
|
|
unique_states.append(state)
|
|
|
|
for state in unique_states:
|
|
if layer_id < 0 or layer_id >= state.total_layers:
|
|
raise ValueError(
|
|
f"layer_id={layer_id} is outside CP HiCache backup layer range "
|
|
f"[0, {state.total_layers}) for node_id={state.reservation.node_id}"
|
|
)
|
|
|
|
target_states = [
|
|
state
|
|
for state in unique_states
|
|
if submit_target
|
|
and layer_id < self.layer_num
|
|
and layer_id not in state.completed_target_layers
|
|
]
|
|
draft_states = [
|
|
state
|
|
for state in unique_states
|
|
if submit_draft
|
|
and state.draft_host_indices is not None
|
|
and layer_id < self.draft_mem_pool_device.layer_num
|
|
and layer_id not in state.completed_draft_layers
|
|
]
|
|
if not target_states and not draft_states:
|
|
return
|
|
|
|
active_states: List[HiCacheLayerWriteState] = []
|
|
seen_active_ids = set()
|
|
for state in target_states + draft_states:
|
|
state_id = id(state)
|
|
if state_id in seen_active_ids:
|
|
continue
|
|
seen_active_ids.add(state_id)
|
|
active_states.append(state)
|
|
|
|
layer_event = device_module.Event()
|
|
layer_event.record()
|
|
for state in active_states:
|
|
state.layer_events.append(layer_event)
|
|
|
|
target_transfer_states = [
|
|
state
|
|
for state in target_states
|
|
if state.physical_device_indices is not None
|
|
and len(state.physical_device_indices) > 0
|
|
]
|
|
draft_transfer_states = [
|
|
state
|
|
for state in draft_states
|
|
if state.physical_device_indices is not None
|
|
and len(state.physical_device_indices) > 0
|
|
]
|
|
grouped_tensors: List[torch.Tensor] = []
|
|
final_states: List[HiCacheLayerWriteState] = []
|
|
|
|
with device_module.stream(self.write_stream):
|
|
for state in active_states:
|
|
state.start_event.wait(self.write_stream)
|
|
layer_event.wait(self.write_stream)
|
|
|
|
if target_transfer_states:
|
|
target_host_indices = self._concat_layer_write_tensors(
|
|
[state.host_indices for state in target_transfer_states]
|
|
)
|
|
target_device_indices = self._concat_layer_write_tensors(
|
|
[
|
|
state.physical_device_indices
|
|
for state in target_transfer_states
|
|
]
|
|
)
|
|
self.mem_pool_host.backup_from_device_per_layer(
|
|
self.mem_pool_device,
|
|
target_host_indices,
|
|
target_device_indices,
|
|
layer_id,
|
|
self.io_backend,
|
|
)
|
|
grouped_tensors.extend([target_host_indices, target_device_indices])
|
|
if draft_transfer_states:
|
|
draft_host_indices = self._concat_layer_write_tensors(
|
|
[state.draft_host_indices for state in draft_transfer_states]
|
|
)
|
|
draft_device_indices = self._concat_layer_write_tensors(
|
|
[
|
|
state.physical_device_indices
|
|
for state in draft_transfer_states
|
|
]
|
|
)
|
|
self.draft_mem_pool_host.backup_from_device_per_layer(
|
|
self.draft_mem_pool_device,
|
|
draft_host_indices,
|
|
draft_device_indices,
|
|
layer_id,
|
|
self.io_backend,
|
|
)
|
|
grouped_tensors.extend([draft_host_indices, draft_device_indices])
|
|
for tensor in grouped_tensors:
|
|
self._record_tensor_on_stream(tensor, self.write_stream)
|
|
|
|
for state in target_states:
|
|
state.completed_target_layers.add(layer_id)
|
|
for state in draft_states:
|
|
state.completed_draft_layers.add(layer_id)
|
|
|
|
final_states = [
|
|
state
|
|
for state in active_states
|
|
if self._layer_write_state_done(state) and not state.ack_appended
|
|
]
|
|
for state in final_states:
|
|
state.finish_event.record()
|
|
self._record_tensor_on_stream(state.host_indices, self.write_stream)
|
|
self._record_tensor_on_stream(
|
|
state.physical_device_indices, self.write_stream
|
|
)
|
|
self._record_tensor_on_stream(
|
|
state.draft_host_indices, self.write_stream
|
|
)
|
|
|
|
# Value fingerprint of the EXACT device KV being backed up, taken on the
|
|
# DEFAULT stream (outside the write_stream block) so it reflects the
|
|
# correct post-store value -- NOT the racy write_stream copy's view. If a
|
|
# per-layer store/copy ordering race exists, the host copy will diverge
|
|
# from THIS hash, and the reload hash (post-H2D) will not match it.
|
|
for state in final_states:
|
|
self._append_layer_write_ack(state)
|
|
|
|
def submit_write_cp_layer(
|
|
self,
|
|
reservation: HiCacheWriteReservation,
|
|
layer_id: int,
|
|
*,
|
|
submit_target: bool = True,
|
|
submit_draft: bool = True,
|
|
) -> None:
|
|
"""Submit one layer of a reserved CP host backup.
|
|
|
|
Target and draft D2H are still one logical operation: this method only
|
|
queues a final write ack when all target/draft layers have been submitted.
|
|
"""
|
|
|
|
state = self._get_or_create_layer_write_state(reservation)
|
|
self._submit_write_cp_layer_states(
|
|
[state],
|
|
layer_id,
|
|
submit_target=submit_target,
|
|
submit_draft=submit_draft,
|
|
)
|
|
|
|
def submit_write_cp_per_layer(
|
|
self,
|
|
reservation: HiCacheWriteReservation,
|
|
*,
|
|
catch_up_all_layers: bool = True,
|
|
) -> None:
|
|
"""Register a CP write reservation for per-layer host backup.
|
|
|
|
Current radix insertion calls write_backup after the request KV has already
|
|
been materialized. For that path, catch up by submitting all layers
|
|
immediately through the per-layer API. Future early reservations can use
|
|
catch_up_all_layers=False and rely on explicit model layer-end hooks.
|
|
"""
|
|
|
|
state = self._get_or_create_layer_write_state(reservation)
|
|
if not catch_up_all_layers:
|
|
logger.debug(
|
|
"[CacheCtrl-write] submit_write_cp_per_layer registered: node_id=%d logical_len=%d layers=%d draft=%s",
|
|
reservation.node_id,
|
|
reservation.metadata.logical_len,
|
|
state.total_layers,
|
|
reservation.draft_host_indices is not None,
|
|
)
|
|
return
|
|
|
|
logger.warning(
|
|
"[CP_HICACHE_FALLBACK][catch_up_all_layers] "
|
|
"CP HiCache write is submitting all layers after KV materialization; "
|
|
"forward-overlap per-layer backup was not active. "
|
|
"node_id=%d logical_len=%d layers=%d draft=%s",
|
|
reservation.node_id,
|
|
reservation.metadata.logical_len,
|
|
state.total_layers,
|
|
reservation.draft_host_indices is not None,
|
|
)
|
|
total_layers = state.total_layers
|
|
for layer_id in range(total_layers):
|
|
self.submit_write_cp_layer(reservation, layer_id)
|
|
|
|
def _write_cp(
|
|
self,
|
|
device_indices: torch.Tensor,
|
|
priority: Optional[int] = None,
|
|
node_id: int = -1,
|
|
) -> HiCacheWriteResult | HiCacheWriteFailure:
|
|
reservation = self.reserve_write_cp(device_indices, priority, node_id)
|
|
if isinstance(reservation, HiCacheWriteFailure):
|
|
return reservation
|
|
self.submit_write_cp_per_layer(reservation)
|
|
return HiCacheWriteResult(metadata=reservation.metadata)
|
|
|
|
def set_draft_kv_pool(self, draft_device_pool, draft_host_pool) -> None:
|
|
"""Register draft KV pools so L2 ops piggyback draft transfers."""
|
|
self.has_draft = True
|
|
self.mem_pool_device_draft = draft_device_pool
|
|
self.mem_pool_host_draft = draft_host_pool
|
|
self.attach_draft_pool(draft_device_pool, draft_host_pool)
|
|
logger.info(
|
|
"HiCache draft KV registered: %s (host %d slots)",
|
|
type(draft_device_pool).__name__,
|
|
draft_host_pool.size,
|
|
)
|
|
|
|
def load(
|
|
self,
|
|
host_indices: torch.Tensor,
|
|
priority: Optional[int] = None,
|
|
node_id: int = -1,
|
|
) -> Optional[torch.Tensor]:
|
|
"""
|
|
Load KV caches from host memory to device memory.
|
|
"""
|
|
device_indices = self.mem_pool_device_allocator.alloc(len(host_indices))
|
|
if device_indices is None:
|
|
return None
|
|
self.load_queue.append(
|
|
CacheOperation(host_indices, device_indices, node_id, priority)
|
|
)
|
|
if self.has_draft and not self.uses_cp_hicache:
|
|
self.draft_load_queue.append(
|
|
CacheOperation(host_indices, device_indices, node_id, priority)
|
|
)
|
|
return device_indices
|
|
|
|
def load_cp(self, nodes_to_load, node_id: int = -1) -> Optional[torch.Tensor]:
|
|
start_time = time.perf_counter()
|
|
stage_start_time = start_time
|
|
stage_durations_ms: List[tuple[str, float]] = []
|
|
|
|
def record_stage(stage: str) -> None:
|
|
nonlocal stage_start_time
|
|
now = time.perf_counter()
|
|
stage_durations_ms.append((stage, (now - stage_start_time) * 1000.0))
|
|
stage_start_time = now
|
|
|
|
# Reproduce the original (write-time) CP owner pattern. Each node
|
|
# carries `page_owners` (one int8 per logical page, identical on all
|
|
# CP ranks) so we can ask the allocator for a fresh device range
|
|
# whose per-page owner sequence matches what was backed up. Without
|
|
# this, the saved `owned_positions` (positions WITHIN the original
|
|
# alloc that this rank owned) would index a new alloc with arbitrary
|
|
# owner pattern → each rank loads its host bytes into physical slots
|
|
# whose corresponding logical page is owned by some other rank →
|
|
# attention reads garbage at forward time.
|
|
page_owners: List[int] = []
|
|
for node in nodes_to_load:
|
|
meta = node.cp_hicache
|
|
if meta is None:
|
|
raise RuntimeError(
|
|
f"load_cp called with node {getattr(node, 'id', '?')} "
|
|
"that has no cp_hicache metadata"
|
|
)
|
|
# page_owners is CPU int8; .tolist() returns list[int] directly.
|
|
page_owners.extend(meta.page_owners.tolist())
|
|
record_stage("collect_page_owners")
|
|
|
|
device_indices = self.mem_pool_device_allocator.alloc_pages_with_owners(
|
|
page_owners
|
|
)
|
|
record_stage("alloc_pages_with_owners")
|
|
# Fail closed: returning None lets the caller drop to cache miss
|
|
# (cold prefill). Never proceed with a non-matching owner pattern.
|
|
if device_indices is None:
|
|
_cp_shared_kv_bs_gt1_cache_timing(
|
|
"hicache_load_cp_plan",
|
|
start_time,
|
|
"node_id=%d result=alloc_none pages=%d stages_ms=%s",
|
|
node_id,
|
|
len(page_owners),
|
|
stage_durations_ms,
|
|
)
|
|
return None
|
|
|
|
padded_len_expected = sum(
|
|
int(getattr(node.cp_hicache, "padded_len", node.host_len))
|
|
for node in nodes_to_load
|
|
)
|
|
if device_indices.numel() != padded_len_expected:
|
|
self.mem_pool_device_allocator.free(device_indices)
|
|
raise RuntimeError(
|
|
"alloc_pages_with_owners returned unexpected length: "
|
|
f"got {device_indices.numel()}, expected {padded_len_expected}"
|
|
)
|
|
|
|
host_chunks = []
|
|
draft_host_chunks = []
|
|
physical_chunks = []
|
|
visible_chunks = []
|
|
offset = 0
|
|
for node in nodes_to_load:
|
|
meta = node.cp_hicache
|
|
valid_len = int(getattr(meta, "valid_len", node.host_len))
|
|
padded_len = int(getattr(meta, "padded_len", node.host_len))
|
|
if valid_len != int(node.host_len):
|
|
self.mem_pool_device_allocator.free(device_indices)
|
|
raise RuntimeError(
|
|
"CP HiCache load metadata valid length mismatch: "
|
|
f"node_id={getattr(node, 'id', '?')} "
|
|
f"host_len={node.host_len} valid_len={valid_len}"
|
|
)
|
|
if padded_len < valid_len:
|
|
self.mem_pool_device_allocator.free(device_indices)
|
|
raise RuntimeError(
|
|
"CP HiCache load metadata padded length is shorter than "
|
|
"valid length: "
|
|
f"node_id={getattr(node, 'id', '?')} "
|
|
f"valid_len={valid_len} padded_len={padded_len}"
|
|
)
|
|
node_device_indices = device_indices[offset : offset + padded_len]
|
|
offset += padded_len
|
|
visible_chunks.append(node_device_indices[:valid_len])
|
|
if self.has_draft_hicache:
|
|
draft_host_indices = getattr(
|
|
meta, "draft_host_indices", None
|
|
)
|
|
if draft_host_indices is None:
|
|
self.mem_pool_device_allocator.free(device_indices)
|
|
raise RuntimeError(
|
|
"CP HiCache draft KV restore requested but node metadata "
|
|
"does not contain draft_host_indices"
|
|
)
|
|
else:
|
|
draft_host_indices = None
|
|
|
|
owned_positions = meta.owned_positions.to(device_indices.device)
|
|
if owned_positions.numel() == 0:
|
|
continue
|
|
selected_logical_locs = node_device_indices[owned_positions]
|
|
# Note: NOT re-validating owner_by_this_rank here. The invariant
|
|
# — every position in `owned_positions` lands on a page owned by
|
|
# this rank — is guaranteed by construction:
|
|
# (a) at write time, `owned_positions = owned_mask.nonzero(...)`
|
|
# where `owned_mask = owned_by_this_rank(device_indices)`;
|
|
# (b) at load time, `alloc_pages_with_owners(page_owners)` returns
|
|
# pages whose owner sequence matches `page_owners` by
|
|
# construction (and is debug-asserted inside the allocator).
|
|
# Both sides use the same `page_owners` (write-derived, identical
|
|
# on all CP ranks). A redundant `.all().item()` check here would
|
|
# force a CUDA host-sync — the exact anti-pattern commit 97a9f850c
|
|
# removed from the hot path.
|
|
physical_chunks.append(
|
|
self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs)
|
|
)
|
|
host_chunks.append(meta.host_indices)
|
|
if draft_host_indices is not None:
|
|
draft_host_chunks.append(draft_host_indices)
|
|
record_stage("build_chunks")
|
|
|
|
visible_device_indices = (
|
|
torch.cat(visible_chunks)
|
|
if visible_chunks
|
|
else torch.empty(
|
|
(0,), dtype=device_indices.dtype, device=device_indices.device
|
|
)
|
|
)
|
|
|
|
if not host_chunks:
|
|
# Keep CP load ACK rows identical across ranks. A zero-owned rank
|
|
# still queues a zero-length op with the logical node id so a later
|
|
# batched start_loading() merges the same node_ids as owning ranks.
|
|
self.load_queue.append(
|
|
CacheOperation(
|
|
torch.empty((0,), dtype=torch.int64),
|
|
torch.empty(
|
|
(0,), dtype=device_indices.dtype, device=device_indices.device
|
|
),
|
|
node_id,
|
|
)
|
|
)
|
|
_cp_shared_kv_bs_gt1_cache_timing(
|
|
"hicache_load_cp_plan",
|
|
start_time,
|
|
"node_id=%d result=zero_owned pages=%d visible_indices=%d "
|
|
"draft=%s stages_ms=%s",
|
|
node_id,
|
|
len(page_owners),
|
|
int(visible_device_indices.numel()),
|
|
self.has_draft_hicache,
|
|
[(stage, round(ms, 3)) for stage, ms in stage_durations_ms],
|
|
)
|
|
return visible_device_indices
|
|
|
|
host_indices = torch.cat(host_chunks)
|
|
physical_device_indices = torch.cat(physical_chunks)
|
|
draft_host_indices = None
|
|
if self.has_draft_hicache:
|
|
draft_host_indices = torch.cat(draft_host_chunks)
|
|
record_stage("concat_indices")
|
|
|
|
try:
|
|
self._validate_cp_hicache_page_indices(
|
|
host_indices, physical_device_indices
|
|
)
|
|
if draft_host_indices is not None:
|
|
self._validate_cp_hicache_page_indices(
|
|
draft_host_indices, physical_device_indices
|
|
)
|
|
record_stage("validate_indices")
|
|
except Exception:
|
|
self.mem_pool_device_allocator.free(device_indices)
|
|
raise
|
|
|
|
self.load_queue.append(
|
|
CacheOperation(
|
|
host_indices,
|
|
physical_device_indices,
|
|
node_id,
|
|
)
|
|
)
|
|
if draft_host_indices is not None:
|
|
self.draft_load_queue.append(
|
|
CacheOperation(
|
|
draft_host_indices,
|
|
physical_device_indices,
|
|
node_id,
|
|
)
|
|
)
|
|
|
|
total_duration = time.perf_counter() - start_time
|
|
if total_duration >= 1.0:
|
|
logger.warning(
|
|
"[HiCache-load] slow load_cp planning: node_id=%d duration_ms=%.3f "
|
|
"pages=%d host_indices=%d physical_indices=%d draft=%s stages_ms=%s",
|
|
node_id,
|
|
total_duration * 1000.0,
|
|
len(page_owners),
|
|
int(host_indices.numel()),
|
|
int(physical_device_indices.numel()),
|
|
draft_host_indices is not None,
|
|
[(stage, round(ms, 3)) for stage, ms in stage_durations_ms],
|
|
)
|
|
_cp_shared_kv_bs_gt1_cache_timing(
|
|
"hicache_load_cp_plan",
|
|
start_time,
|
|
"node_id=%d result=queued pages=%d host_indices=%d "
|
|
"physical_indices=%d visible_indices=%d draft=%s stages_ms=%s",
|
|
node_id,
|
|
len(page_owners),
|
|
int(host_indices.numel()),
|
|
int(physical_device_indices.numel()),
|
|
int(visible_device_indices.numel()),
|
|
draft_host_indices is not None,
|
|
[(stage, round(ms, 3)) for stage, ms in stage_durations_ms],
|
|
)
|
|
return visible_device_indices
|
|
|
|
def move_indices(self, op: CacheOperation, mem_pool_host=None):
|
|
mem_pool_host = mem_pool_host or self.mem_pool_host
|
|
host_indices, device_indices = op.host_indices, op.device_indices
|
|
# move indices to GPU if using kernels, to host if using direct indexing
|
|
if self.io_backend == "kernel":
|
|
if not host_indices.is_cuda:
|
|
host_indices = host_indices.to(self.device, non_blocking=True)
|
|
return host_indices, device_indices
|
|
elif self.io_backend == "direct":
|
|
if mem_pool_host.layout == "layer_first":
|
|
device_indices = device_indices.cpu()
|
|
host_indices, idx = host_indices.sort()
|
|
return host_indices, device_indices.index_select(0, idx)
|
|
elif mem_pool_host.layout in ("page_first_direct", "layer_page_first"):
|
|
return host_indices, device_indices.cpu()
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported layout {mem_pool_host.layout!r} for io backend 'direct'"
|
|
)
|
|
elif self.io_backend == "kernel_ascend":
|
|
return host_indices, device_indices.cpu()
|
|
else:
|
|
raise ValueError(f"Unsupported io backend")
|
|
|
|
def start_loading(self) -> int:
|
|
if len(self.load_queue) == 0 and len(self.draft_load_queue) == 0:
|
|
return -1
|
|
|
|
def begin_load_op(mem_pool_host, host_indices, device_indices) -> bool:
|
|
begin = getattr(mem_pool_host, "begin_load_to_device_op", None)
|
|
if begin is None:
|
|
return False
|
|
begin(host_indices, device_indices, self.io_backend)
|
|
return True
|
|
|
|
def end_load_op(mem_pool_host) -> None:
|
|
end = getattr(mem_pool_host, "end_load_to_device_op", None)
|
|
if end is not None:
|
|
end()
|
|
|
|
producer_id = self.layer_done_counter.update_producer()
|
|
op = CacheOperation.merge_ops(self.load_queue) if self.load_queue else None
|
|
draft_op = (
|
|
CacheOperation.merge_ops(self.draft_load_queue)
|
|
if self.draft_load_queue
|
|
else None
|
|
)
|
|
node_ids = op.node_ids if op is not None else draft_op.node_ids
|
|
if op is not None:
|
|
host_indices, device_indices = self.move_indices(op, self.mem_pool_host)
|
|
else:
|
|
host_indices = device_indices = None
|
|
if draft_op is not None:
|
|
draft_host_indices, draft_device_indices = self.move_indices(
|
|
draft_op, self.draft_mem_pool_host
|
|
)
|
|
else:
|
|
draft_host_indices = draft_device_indices = None
|
|
self.load_queue.clear()
|
|
self.draft_load_queue.clear()
|
|
producer_event = self.layer_done_counter.events[producer_id]
|
|
producer_event.start_event.record()
|
|
|
|
target_load_op_started = False
|
|
draft_load_op_started = False
|
|
try:
|
|
with device_module.stream(self.load_stream):
|
|
producer_event.start_event.wait(self.load_stream)
|
|
prepare_start = _cp_shared_kv_bs_gt1_cache_timing_start()
|
|
if op is not None:
|
|
target_load_op_started = begin_load_op(
|
|
self.mem_pool_host, host_indices, device_indices
|
|
)
|
|
if draft_op is not None:
|
|
draft_load_op_started = begin_load_op(
|
|
self.draft_mem_pool_host,
|
|
draft_host_indices,
|
|
draft_device_indices,
|
|
)
|
|
_cp_shared_kv_bs_gt1_cache_timing_log(
|
|
"prepare_load_descriptor",
|
|
prepare_start,
|
|
"target_tokens=%s draft_tokens=%s target_started=%s draft_started=%s",
|
|
int(host_indices.numel()) if host_indices is not None else 0,
|
|
int(draft_host_indices.numel())
|
|
if draft_host_indices is not None
|
|
else 0,
|
|
target_load_op_started,
|
|
draft_load_op_started,
|
|
)
|
|
draft_layer_num = (
|
|
self.draft_mem_pool_device.layer_num if draft_op is not None else 0
|
|
)
|
|
layer_loop_start = _cp_shared_kv_bs_gt1_cache_timing_start()
|
|
for i in range(max(self.layer_num, draft_layer_num)):
|
|
layer_submit_start = _cp_shared_kv_bs_gt1_cache_timing_start()
|
|
submitted_target = False
|
|
submitted_draft = False
|
|
if draft_op is not None and i < draft_layer_num:
|
|
if len(draft_host_indices) > 0:
|
|
self.draft_mem_pool_host.load_to_device_per_layer(
|
|
self.draft_mem_pool_device,
|
|
draft_host_indices,
|
|
draft_device_indices,
|
|
i,
|
|
self.io_backend,
|
|
)
|
|
submitted_draft = True
|
|
if op is not None and i < self.layer_num:
|
|
if len(host_indices) > 0:
|
|
self.mem_pool_host.load_to_device_per_layer(
|
|
self.mem_pool_device,
|
|
host_indices,
|
|
device_indices,
|
|
i,
|
|
self.io_backend,
|
|
)
|
|
submitted_target = True
|
|
# Value fingerprint of what actually landed on device
|
|
# after H2D (on load_stream, in-order after the load).
|
|
# Compare to backup_kv_hash(node,layer): inequality =
|
|
# round-trip delivered wrong KV. node_ids is the merged
|
|
# op's nodes (single node in the c=1 rehit pass).
|
|
producer_event.complete(i)
|
|
elif op is None and i < self.layer_num:
|
|
producer_event.complete(i)
|
|
_cp_shared_kv_bs_gt1_cache_timing_log(
|
|
"submit_h2d_layer_per_call_slow",
|
|
layer_submit_start,
|
|
"layer_id=%s target=%s draft=%s target_tokens=%s draft_tokens=%s",
|
|
i,
|
|
submitted_target,
|
|
submitted_draft,
|
|
int(host_indices.numel()) if host_indices is not None else 0,
|
|
int(draft_host_indices.numel())
|
|
if draft_host_indices is not None
|
|
else 0,
|
|
)
|
|
_cp_shared_kv_bs_gt1_cache_timing_log(
|
|
"submit_h2d_layer_loop",
|
|
layer_loop_start,
|
|
"layers=%s target_tokens=%s draft_tokens=%s",
|
|
max(self.layer_num, draft_layer_num),
|
|
int(host_indices.numel()) if host_indices is not None else 0,
|
|
int(draft_host_indices.numel())
|
|
if draft_host_indices is not None
|
|
else 0,
|
|
)
|
|
# NOTE: We must save the host indices and device indices here,
|
|
# this is because we need to guarantee that these tensors are
|
|
# still alive when the load stream is executing.
|
|
if host_indices is not None and host_indices.is_cuda:
|
|
host_indices.record_stream(self.load_stream)
|
|
if device_indices is not None and device_indices.is_cuda:
|
|
device_indices.record_stream(self.load_stream)
|
|
if draft_host_indices is not None and draft_host_indices.is_cuda:
|
|
draft_host_indices.record_stream(self.load_stream)
|
|
if draft_device_indices is not None and draft_device_indices.is_cuda:
|
|
draft_device_indices.record_stream(self.load_stream)
|
|
finally:
|
|
end_start = _cp_shared_kv_bs_gt1_cache_timing_start()
|
|
if draft_load_op_started:
|
|
end_load_op(self.draft_mem_pool_host)
|
|
if target_load_op_started:
|
|
end_load_op(self.mem_pool_host)
|
|
_cp_shared_kv_bs_gt1_cache_timing_log(
|
|
"end_load_descriptor",
|
|
end_start,
|
|
"target_started=%s draft_started=%s",
|
|
target_load_op_started,
|
|
draft_load_op_started,
|
|
)
|
|
|
|
self.ack_load_queue.append(
|
|
HiCacheAck(
|
|
start_event=producer_event.start_event,
|
|
finish_event=producer_event.finish_event,
|
|
node_ids=node_ids,
|
|
)
|
|
)
|
|
return producer_id
|
|
|
|
def evict_device(self, device_indices: torch.Tensor) -> int:
|
|
self.mem_pool_device_allocator.free(device_indices)
|
|
return len(device_indices)
|
|
|
|
def evict_host(self, host_indices: torch.Tensor, backup_only: bool = True) -> int:
|
|
if not backup_only:
|
|
raise ValueError("Other eviction policies are not supported yet.")
|
|
|
|
self.mem_pool_host.free(host_indices)
|
|
return len(host_indices)
|
|
|
|
def evict_cp_host(self, metadata, backup_only: bool = True) -> int:
|
|
if not backup_only:
|
|
raise ValueError("Other eviction policies are not supported yet.")
|
|
|
|
draft_host_indices = getattr(metadata, "draft_host_indices", None)
|
|
if self.has_draft_hicache and draft_host_indices is None:
|
|
raise RuntimeError(
|
|
"CP HiCache draft KV host evict requested but node metadata "
|
|
"does not contain draft_host_indices"
|
|
)
|
|
|
|
freed = self.evict_host(metadata.host_indices, backup_only=backup_only)
|
|
if self.has_draft_hicache:
|
|
self.draft_mem_pool_host.free(draft_host_indices)
|
|
return freed
|
|
|
|
def prefetch(
|
|
self,
|
|
request_id: str,
|
|
host_indices: torch.Tensor,
|
|
new_input_tokens: List[int],
|
|
last_hash: Optional[str] = None,
|
|
prefix_keys: Optional[List[str]] = None,
|
|
) -> PrefetchOperation:
|
|
"""
|
|
Prefetch KV caches from storage backend to host memory.
|
|
"""
|
|
operation = PrefetchOperation(
|
|
request_id, host_indices, new_input_tokens, last_hash, prefix_keys
|
|
)
|
|
self.prefetch_queue.put(operation)
|
|
return operation
|
|
|
|
def terminate_prefetch(self, operation):
|
|
operation.mark_terminate()
|
|
return operation.completed_tokens, operation.hash_value
|
|
|
|
def append_host_mem_release(self, host_indices: torch.Tensor):
|
|
if host_indices.numel() == 0:
|
|
return
|
|
pages = host_indices.split(self.mem_pool_host.page_size)
|
|
for page in pages:
|
|
self.host_mem_release_queue.put(page)
|
|
|
|
def _page_get_zero_copy(self, operation, hash_values, host_indices, extra_info=None):
|
|
results = self.storage_backend.batch_get_v1(
|
|
hash_values, host_indices, extra_info
|
|
)
|
|
inc = 0
|
|
for i in range(len(hash_values)):
|
|
if not results[i]:
|
|
logger.warning(
|
|
f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}."
|
|
)
|
|
break
|
|
inc += self.page_size
|
|
operation.increment(inc)
|
|
|
|
# todo: deprecate
|
|
def _generic_page_get(self, operation, hash_values, host_indices, extra_info=None):
|
|
dummy_page_dst = [
|
|
self.mem_pool_host.get_dummy_flat_data_page() for _ in hash_values
|
|
]
|
|
page_data = self.storage_backend.batch_get(hash_values, dummy_page_dst)
|
|
if page_data is None:
|
|
return
|
|
for i in range(len(hash_values)):
|
|
if page_data[i] is None:
|
|
logger.warning(
|
|
f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}."
|
|
)
|
|
break
|
|
# Must set the data before increasing the completed tokens.
|
|
# Otherwise this page may be read before being set.
|
|
self.mem_pool_host.set_from_flat_data_page(
|
|
host_indices[i * self.page_size],
|
|
page_data[i],
|
|
)
|
|
if not operation.increment(self.page_size):
|
|
break # Operation terminated by controller
|
|
|
|
def _page_transfer(self, operation):
|
|
# Transfer batch by batch
|
|
prefix_keys = operation.prefix_keys
|
|
for i in range(0, len(operation.hash_value), self.storage_batch_size):
|
|
batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
|
|
batch_host_indices = operation.host_indices[
|
|
i * self.page_size : (i + len(batch_hashes)) * self.page_size
|
|
]
|
|
prev_completed_tokens = operation.completed_tokens
|
|
# Get one batch token, and update the completed_tokens if succeed
|
|
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo
|
|
|
|
extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
|
|
self.page_get_func(operation, batch_hashes, batch_host_indices, extra_info)
|
|
# Check termination
|
|
if (
|
|
operation.completed_tokens
|
|
!= prev_completed_tokens + len(batch_hashes) * self.page_size
|
|
):
|
|
operation.mark_terminate()
|
|
break # Some operations fail or operation terminated by controller
|
|
|
|
if prefix_keys and len(prefix_keys) > 0:
|
|
prefix_keys += batch_hashes
|
|
|
|
def prefetch_io_aux_func(self):
|
|
"""
|
|
Auxiliary function conducting IO operations for prefetching.
|
|
"""
|
|
while not self.storage_stop_event.is_set():
|
|
try:
|
|
operation = self.prefetch_buffer.get(block=True, timeout=1)
|
|
if operation is None:
|
|
continue
|
|
self._page_transfer(operation)
|
|
# operation terminated by controller, release pre-allocated memory
|
|
self.append_host_mem_release(
|
|
operation.host_indices[operation.completed_tokens :]
|
|
)
|
|
except Empty:
|
|
continue
|
|
|
|
def prefetch_rate_limited(self) -> bool:
|
|
"""
|
|
Rate limit the prefetching operations to avoid overwhelming the storage backend.
|
|
"""
|
|
# cancel prefetch if too much memory is occupied
|
|
if self.prefetch_tokens_occupied >= self.prefetch_capacity_limit:
|
|
return True
|
|
# todo: more sophisticated rate limiting based on storage backend performance
|
|
return False
|
|
|
|
def _storage_hit_query(self, operation) -> tuple[list[str], int]:
|
|
last_hash = operation.last_hash
|
|
tokens_to_fetch = operation.token_ids
|
|
prefix_keys = operation.prefix_keys.copy() if operation.prefix_keys else None
|
|
|
|
storage_query_count = 0
|
|
hash_value = []
|
|
|
|
for start in range(
|
|
0, len(tokens_to_fetch), self.page_size * self.storage_batch_size
|
|
):
|
|
end = min(
|
|
start + self.page_size * self.storage_batch_size, len(tokens_to_fetch)
|
|
)
|
|
batch_tokens = tokens_to_fetch[start:end]
|
|
batch_hashes = []
|
|
for i in range(0, len(batch_tokens), self.page_size):
|
|
last_hash = self.get_hash_str(
|
|
batch_tokens[i : i + self.page_size], last_hash
|
|
)
|
|
batch_hashes.append(last_hash)
|
|
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo
|
|
|
|
extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
|
|
hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info)
|
|
hash_value.extend(batch_hashes[:hit_page_num])
|
|
storage_query_count += hit_page_num * self.page_size
|
|
if hit_page_num < len(batch_hashes):
|
|
break
|
|
if prefix_keys and len(prefix_keys) > 0:
|
|
prefix_keys += batch_hashes
|
|
|
|
return hash_value, storage_query_count
|
|
|
|
def prefetch_thread_func(self):
|
|
"""
|
|
Manage prefetching operations from storage backend to host memory.
|
|
"""
|
|
self.prefetch_buffer = Queue()
|
|
self.prefetch_io_aux_thread = threading.Thread(
|
|
target=self.prefetch_io_aux_func, daemon=True
|
|
)
|
|
self.prefetch_io_aux_thread.start()
|
|
while (not self.storage_stop_event.is_set()) or not self.prefetch_queue.empty():
|
|
try:
|
|
operation = self.prefetch_queue.get(block=True, timeout=1)
|
|
if operation is None:
|
|
continue
|
|
hash_value, storage_hit_count = self._storage_hit_query(operation)
|
|
if self.tp_world_size > 1:
|
|
storage_hit_count_tensor = torch.tensor(
|
|
storage_hit_count, dtype=torch.int
|
|
)
|
|
torch.distributed.all_reduce(
|
|
storage_hit_count_tensor,
|
|
op=torch.distributed.ReduceOp.MIN,
|
|
group=self.prefetch_tp_group,
|
|
)
|
|
storage_hit_count = storage_hit_count_tensor.item()
|
|
|
|
if storage_hit_count < self.prefetch_threshold:
|
|
# not to prefetch if not enough benefits
|
|
self.prefetch_revoke_queue.put(operation.request_id)
|
|
self.append_host_mem_release(operation.host_indices)
|
|
logger.info(
|
|
f"Revoking prefetch for request {operation.request_id} due to insufficient hits ({storage_hit_count})."
|
|
)
|
|
else:
|
|
operation.hash_value = hash_value[
|
|
: (storage_hit_count // self.page_size)
|
|
]
|
|
# free the pre-allocated memory for pages that are not hit
|
|
self.append_host_mem_release(
|
|
operation.host_indices[storage_hit_count:]
|
|
)
|
|
operation.host_indices = operation.host_indices[:storage_hit_count]
|
|
logger.info(
|
|
f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}."
|
|
)
|
|
self.prefetch_buffer.put(operation)
|
|
|
|
except Empty:
|
|
continue
|
|
|
|
def write_storage(
|
|
self,
|
|
host_indices: torch.Tensor,
|
|
token_ids: List[int],
|
|
hash_value: Optional[List[str]] = None,
|
|
prefix_keys: Optional[List[str]] = None,
|
|
) -> int:
|
|
"""
|
|
Write KV caches from host memory to storage backend.
|
|
"""
|
|
operation = StorageOperation(
|
|
host_indices, token_ids, hash_value=hash_value, prefix_keys=prefix_keys
|
|
)
|
|
self.backup_queue.put(operation)
|
|
return operation.id
|
|
|
|
# todo: deprecate
|
|
def _generic_page_set(self, hash_values, host_indices, extra_info=None) -> bool:
|
|
data = [
|
|
self.mem_pool_host.get_data_page(host_indices[i * self.page_size])
|
|
for i in range(len(hash_values))
|
|
]
|
|
return self.storage_backend.batch_set(hash_values, data)
|
|
|
|
def _page_set_zero_copy(self, hash_values, host_indices, extra_info=None) -> bool:
|
|
return all(
|
|
self.storage_backend.batch_set_v1(hash_values, host_indices, extra_info)
|
|
)
|
|
|
|
# Backup batch by batch
|
|
def _page_backup(self, operation):
|
|
# Backup batch by batch
|
|
prefix_keys = operation.prefix_keys
|
|
for i in range(0, len(operation.hash_value), self.storage_batch_size):
|
|
batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
|
|
batch_host_indices = operation.host_indices[
|
|
i * self.page_size : (i + len(batch_hashes)) * self.page_size
|
|
]
|
|
# Set one batch token, and record if success.
|
|
# todo: allow partial success
|
|
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo
|
|
|
|
extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
|
|
success = self.page_set_func(batch_hashes, batch_host_indices, extra_info)
|
|
if not success:
|
|
logger.warning(
|
|
f"Write page to storage: {len(batch_hashes)} pages failed."
|
|
)
|
|
break
|
|
|
|
if prefix_keys and len(prefix_keys) > 0:
|
|
prefix_keys += batch_hashes
|
|
operation.completed_tokens += self.page_size * len(batch_hashes)
|
|
|
|
def backup_thread_func(self):
|
|
"""
|
|
Manage backup operations from host memory to storage backend.
|
|
"""
|
|
while not self.storage_stop_event.is_set():
|
|
try:
|
|
operation = self.backup_queue.get(block=True, timeout=1)
|
|
if operation is None:
|
|
continue
|
|
|
|
if not self.backup_skip:
|
|
self._page_backup(operation)
|
|
self.ack_backup_queue.put(operation)
|
|
|
|
except Empty:
|
|
continue
|