chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Overlap CP HiCache backup without exposing partial host state

Browse Source 
CP shared KV with HiCache and EAGLE needs host backup to overlap forward while keeping radix visibility synchronous. The change reserves host slots before forward, drives target and draft backup from explicit layer-end hooks, and commits host visibility only after the final target/draft ack. It also probes the final insertion prefix before early reservation so repeated EAGLE prompts do not prepare duplicate suffix backups that later rollback as insert_miss.

Constraint: CP ranks use independent shared-KV pools, so target/draft host state must remain atomically visible at the radix boundary.

Constraint: Fused MLA and NSA store paths can bypass store-side notifier hooks, so layer end is the safer backup progress boundary.

Rejected: Store-side backup notifier as the primary trigger | fused store and zero-local paths made notifier coverage fragile.

Rejected: Reserve from cache_protected_len alone | EAGLE bigram/page alignment can make final insertion find a longer existing prefix and force duplicate rollback work.

Confidence: medium

Scope-risk: moderate

Directive: Do not add per-layer CP collectives here; keep radix state synchronous and data transfer asynchronous/local-event driven.

Tested: local git diff --check

Tested: local py_compile for touched CP HiCache/cache-controller/deepseek/test files

Tested: remote pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q (115 passed, 5 warnings)

Not-tested: full GLM5 ETE server rerun after this commit
This commit is contained in:
laoyao0822
2026-05-27 09:50:47 +08:00
parent 03529319a1
commit f355fdd39e
11 changed files with 747 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md
+62 -12
View File
@@ -226,6 +226,18 @@ all_reduce per node victim in a tight host-eviction loop
all_reduce every scheduler tick when no completion prefix advanced
```
Current correctness note: CP host reservation now synchronizes
`required_host_slots` with `all_reduce(MAX)` before the host-eviction retry.
This forces every rank into the same reserve/evict/retry branch and avoids
collective mismatches when one rank is host-full and another rank reserves
successfully. It is intentionally a coarse slow-path collective, not a
per-layer collective, but it can become a performance cost when host pressure
is frequent because every reservation failure pays at least one rank-wide MAX
sync and retry failures pay a second one. Treat this as a correctness guard to
be amortized later with batched reservation epochs, deterministic host
watermarks, or less frequent proactive host eviction; do not move it into the
per-layer data path.
## Per-Layer Backup Data Plane
The existing host pool exposes `load_to_device_per_layer()` but not a matching
@@ -415,6 +427,13 @@ Refactor host full handling so reserve failure returns required local physical
slots, host eviction frees only valid host-only victims, and reservation retries
once. Add logs/counters for backup skipped due to host pressure.
The current first implementation uses rank-wide
`MAX(required_host_slots)` synchronization before eviction/retry so all CP ranks
observe the same capacity-pressure branch. This fixes correctness under
rank-skewed host pressure, but it is a known performance risk if host-full
becomes common. Follow-up optimization should reduce how often this slow path
is reached rather than adding more collectives around it.
### Phase 7: Split safety
Implement defer/requeue-on-pending-split for nodes with pending backup. Add
@@ -436,10 +455,14 @@ paths:
## Open Questions / Decisions
1. **Layer completion hook.** The chosen semantic hook is layer-end after that
layer's KV store is ordered. The exact code location still needs
implementation-time confirmation. A next-layer-start hook is a fallback only
if it preserves the same ordering and adds an explicit final-layer trigger.
1. **Layer completion hook.** The chosen semantic hook is explicit model
layer-end after that layer's target/draft KV and NSA indexer writes have
been enqueued. Store-side notifiers are intentionally not the ownership
boundary because MLA, NSA indexer, TAI fused store, Triton store, and
zero-local paths can bypass one another. The first implementation hooks
`DeepseekV2DecoderLayer.forward()` and the TBO `op_comm_postprocess_layer()`
path; TBO children only notify from subbatch 1 so a layer is not backed up
after only the first child has written its KV.
2. **Pending split behavior.** The chosen first pass is defer/requeue the
request that would split a node with pending backup. Do not split the
in-flight backup op, including for future `bs > 1`.
@@ -466,9 +489,8 @@ Implemented in the first P4-P8 pass:
- pending backup split attempts are deferred instead of mutating the radix tree;
- request admission marks deferred prefix matches and keeps those requests in
the waiting path for a later retry;
- KV pools expose a no-op-safe layer-store notifier. NSA defers notification
until `set_index_k_scale_buffer()` so the indexer payload is ordered after the
KV payload.
- KV pools expose a no-op-safe explicit layer-end notifier. Store methods and
fused store kernels do not drive backup progress.
Important limitation: the current `sgl_kernel.kvcacheio` surface has H2D
per-layer page-first kernels and D2H all-layer page-first kernels, but not a
@@ -503,14 +525,16 @@ This covers the GLM5/NSA production shape:
MHA page-first remains fail-fast until it has a dedicated kernel.
The current radix insertion path creates host reservations after request KV has
already been materialized. Therefore `submit_write_cp_per_layer()` performs a
safe catch-up submission of all layers through the per-layer API for current
call sites. The layer-store notifier path is also wired:
The first production forward-overlap path reserves host slots after
`prepare_for_extend()` and before the model forward, then calls
`submit_write_cp_per_layer(..., catch_up_all_layers=False)`. Radix insertion
only attaches that prepared reservation; it does not replay an all-layer
catch-up copy. The model layer-end hook is the forward progress driver:
```text
submit_write_cp_per_layer(..., catch_up_all_layers=False)
-> on_layer_kv_stored(layer_id, source="target" | "draft")
-> DeepseekV2 layer-end calls token_to_kv_pool.notify_layer_end_for_backup(layer_id)
-> HiCacheController.on_layer_end(layer_id, source="target" | "draft")
-> submit_write_cp_layer(reservation, layer_id, submit_target=..., submit_draft=...)
-> final ack only after all target/draft layers are submitted
```
@@ -521,6 +545,32 @@ payloads remain invisible until the final write ack commits
separately so an early target hook cannot copy draft KV before draft has stored
the same layer.
The previous store-notifier attempt was rejected as the primary correctness
boundary. It was too easy for fused paths to skip it: NSA fused indexer store
bypassed `set_index_k_scale_buffer()`, MLA can use `SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=1`,
and zero-local paths do not necessarily touch the same setter. Missing one
store notifier left the prepared backup in `pending_layer_writes` forever: the
first request could finish prefill without a final write ack, and a second
identical request could hit the pending node and wedge the scheduler. The
layer-end hook makes backup progress independent of the chosen store kernel;
store-side hooks should remain absent or debug-only.
One more correctness/performance edge was observed with repeated EAGLE requests.
Scheduler prefix matching and final radix insertion can have different effective
lengths because matching is based on the schedulable prefix while insertion uses
the final page-aligned key. With EAGLE bigram keys this can leave
`req.cache_protected_len` shorter than the prefix that insertion will find in the
tree. If early backup reserves `[cache_protected_len, insertion_key_len)` blindly,
the system copies a duplicate page, then final insertion discovers that no new
node was created and rolls the prepared backup back with `reason=insert_miss`.
This is safe but wastes host reservation, per-layer callbacks, and write-stream
work on repeated prompts.
The early backup path now probes the existing radix prefix for the final
insertion key without splitting the tree, and only reserves the suffix beyond
`max(cache_protected_len, existing_insertion_prefix_len)`. This preserves the
cold-request behavior while avoiding duplicate prepared backups on exact repeats.
## Summary
Per-layer backup should be implemented as a data-plane refinement, not a radix
python/sglang/srt/managers/cache_controller.py
+4 -4
View File
@@ -361,7 +361,7 @@ class HiCacheController:
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_kv_stored(layer_id, source="target")
lambda layer_id: self.on_layer_end(layer_id, source="target")
)
self.draft_mem_pool_host = None
self.draft_mem_pool_device = None
@@ -431,10 +431,10 @@ class HiCacheController:
)
if hasattr(draft_mem_pool_device, "register_layer_backup_notifier"):
draft_mem_pool_device.register_layer_backup_notifier(
lambda layer_id: self.on_layer_kv_stored(layer_id, source="draft")
lambda layer_id: self.on_layer_end(layer_id, source="draft")
)
def on_layer_kv_stored(self, layer_id: int, source: str = "target") -> None:
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"):
@@ -1236,7 +1236,7 @@ class HiCacheController:
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 on_layer_kv_stored().
catch_up_all_layers=False and rely on explicit model layer-end hooks.
"""
state = self._get_or_create_layer_write_state(reservation)
python/sglang/srt/managers/schedule_batch.py
+1
View File
@@ -644,6 +644,7 @@ class Req(ReqDllmMixin):
self.last_host_backup_node: Any = None
self.host_hit_length = 0
self.prefix_match_deferred_by_pending_backup = False
self.cp_hicache_prepared_backup = None
# Tokens loaded from storage backend (L3) during prefetch for this request
self.storage_hit_length = 0
# The node to lock until for swa radix tree lock ref
python/sglang/srt/managers/scheduler.py
+5
View File
@@ -2457,6 +2457,11 @@ class Scheduler(
)
new_batch.prepare_for_extend()
if self.enable_hierarchical_cache and hasattr(
self.tree_cache, "prepare_write_backup_for_req"
):
for req in can_run_list:
self.tree_cache.prepare_write_backup_for_req(req)
# Record prefill stats for logging after forward
new_batch.prefill_stats = PrefillStats.from_adder(
python/sglang/srt/mem_cache/base_prefix_cache.py
+1
View File
@@ -60,6 +60,7 @@ class InsertParams:
# General
chunked: bool = False
priority: int = 0
cp_hicache_prepared_backup: Optional[object] = None
@dataclasses.dataclass
python/sglang/srt/mem_cache/hiradix_cache.py
+286 -23
View File
@@ -17,6 +17,7 @@ from sglang.srt.environ import envs
from sglang.srt.managers.cache_controller import (
HiCacheController,
HiCacheWriteFailure,
HiCacheWriteReservation,
PrefetchOperation,
)
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
@@ -44,6 +45,7 @@ from sglang.srt.mem_cache.radix_cache import (
RadixKey,
TreeNode,
compute_node_hash_values,
page_align_keys,
split_node_hash_value,
)
from sglang.srt.mem_cache.utils import convert_to_bigram_key
@@ -246,6 +248,15 @@ class PendingHiCacheBackup:
locked: bool = True
@dataclass
class PreparedCpHiCacheBackup:
node_id: int
reservation: HiCacheWriteReservation
metadata: CpHiCacheNodeMetadata
logical_len: int
attached: bool = False
class HiCachePendingBackupSplit(Exception):
def __init__(self, node: TreeNode):
self.node = node
@@ -1010,6 +1021,260 @@ class HiRadixCache(RadixCache):
self.dec_node_lock_ref(node)
return node
def _sync_cp_write_required_host_slots(self, result) -> int:
"""Return the max host slots any CP rank needs before write backup.
CP ranks must make the reserve/evict/retry decision collectively. A
local successful reservation is not enough: another CP rank may be full
and enter host eviction, whose implementation contains collective
all_reduce calls. If successful ranks skip that branch they can instead
enter unrelated collectives (for example disagg polling), causing a
process-group mismatch.
"""
required_host_slots = (
int(result.required_host_slots)
if isinstance(result, HiCacheWriteFailure)
else 0
)
if (
getattr(self, "tp_group", None) is None
or getattr(self, "tp_world_size", 1) <= 1
):
return required_host_slots
required = torch.tensor(required_host_slots, dtype=torch.int64, device="cpu")
torch.distributed.all_reduce(
required, op=torch.distributed.ReduceOp.MAX, group=self.tp_group
)
return int(required.item())
def _release_cp_write_reservation_for_retry(self, result, node_id: int) -> None:
if isinstance(result, HiCacheWriteFailure):
return
logger.info(
"[HiCache-write] release local CP reservation before collective retry: node_id=%d owned_positions=%d",
node_id,
result.metadata.owned_positions.numel(),
)
self.cache_controller.evict_cp_host(result.metadata)
def _reserve_write_cp_indices_collectively(
self, device_indices: torch.Tensor, node_id: int
):
result = self.cache_controller.reserve_write_cp(
device_indices=device_indices,
node_id=node_id,
)
required_host_slots = self._sync_cp_write_required_host_slots(result)
if required_host_slots <= 0:
return result
self._release_cp_write_reservation_for_retry(result, node_id)
self._evict_host_for_physical_slots(
required_host_slots,
synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1,
)
logger.info(
"[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d",
node_id,
required_host_slots,
)
result = self.cache_controller.reserve_write_cp(
device_indices=device_indices,
node_id=node_id,
)
required_host_slots = self._sync_cp_write_required_host_slots(result)
if required_host_slots <= 0:
return result
self._release_cp_write_reservation_for_retry(result, node_id)
logger.info(
"[HiCache-write] write_backup CP FAILED after collective retry: node_id=%d len=%d needed_slots=%d",
node_id,
len(device_indices) if device_indices is not None else 0,
required_host_slots,
)
return HiCacheWriteFailure(required_host_slots=required_host_slots)
def _reserve_write_cp_collectively(self, node: TreeNode):
return self._reserve_write_cp_indices_collectively(node.value, node.id)
def _attach_prepared_cp_backup(
self, node: TreeNode, prepared: PreparedCpHiCacheBackup
) -> None:
if prepared.attached:
raise RuntimeError(
f"CP HiCache prepared backup node_id={prepared.node_id} was attached twice"
)
if len(node.value) != prepared.logical_len:
raise RuntimeError(
f"Prepared CP HiCache backup length mismatch for node_id={prepared.node_id}: "
f"node_len={len(node.value)} prepared_len={prepared.logical_len}"
)
node.id = prepared.node_id
self.ongoing_write_through[node.id] = node
self.inc_node_lock_ref(node)
self.pending_host_backups[node.id] = PendingHiCacheBackup(
node=node,
metadata=prepared.metadata,
logical_len=prepared.logical_len,
submitted=True,
locked=True,
)
prepared.attached = True
logger.info(
"[HiCache-write] attached prepared CP backup: node_id=%d logical_len=%d owned_positions=%d pending_backups=%d",
node.id,
prepared.logical_len,
prepared.metadata.owned_positions.numel(),
len(self.pending_host_backups),
)
def _rollback_prepared_cp_backup(
self, prepared: Optional[PreparedCpHiCacheBackup], reason: str
) -> None:
if prepared is None or prepared.attached:
return
if prepared.node_id in self.cache_controller.pending_layer_writes:
raise RuntimeError(
f"Cannot rollback prepared CP HiCache backup node_id={prepared.node_id} "
f"while per-layer writes are still pending; reason={reason}"
)
retained_acks = []
removed_ack = False
for ack in self.cache_controller.ack_write_queue:
if prepared.node_id not in ack.node_ids:
retained_acks.append(ack)
continue
ack.finish_event.synchronize()
remaining_node_ids = [
node_id for node_id in ack.node_ids if node_id != prepared.node_id
]
if remaining_node_ids:
retained_acks.append(ack._replace(node_ids=remaining_node_ids))
removed_ack = True
if removed_ack:
self.cache_controller.ack_write_queue = retained_acks
logger.info(
"[HiCache-write] rollback unattached prepared CP backup: node_id=%d logical_len=%d reason=%s",
prepared.node_id,
prepared.logical_len,
reason,
)
self.cache_controller.evict_cp_host(prepared.metadata)
def _probe_existing_radix_prefix_len_no_split(self, key: RadixKey) -> int:
"""Return the already-present radix prefix length without mutating the tree.
`prepare_write_backup_for_req` runs before forward and must reserve host
slots only for KV that will become a *new* radix node at insertion time.
The request's `cache_protected_len` is based on scheduler prefix-match
semantics, which can be shorter than the prefix that a final insertion
will find (for example EAGLE/bigram plus page alignment). Probing the
full insertion key here avoids backing up duplicate pages that insertion
will later reject and roll back.
"""
if self.disable or len(key) == 0:
return 0
key, _ = self.maybe_bigram_convert(key)
if self.page_size != 1:
page_aligned_len = len(key) // self.page_size * self.page_size
key = key[:page_aligned_len]
if len(key) == 0:
return 0
node = getattr(self, "root_node", None)
if node is None:
return 0
child_key = self.get_child_key_fn(key)
total_prefix_len = 0
while len(key) > 0 and child_key in node.children:
child = node.children[child_key]
prefix_len = self.key_match_fn(child.key, key)
if prefix_len <= 0:
break
total_prefix_len += prefix_len
if prefix_len < len(child.key):
break
node = child
key = key[prefix_len:]
if len(key):
child_key = self.get_child_key_fn(key)
return total_prefix_len
def prepare_write_backup_for_req(self, req) -> None:
if (
self.disable
or not self._uses_cp_hicache
or self.cache_controller.write_policy == "write_back"
or getattr(req, "is_chunked", 0) > 0
or getattr(req, "cp_hicache_prepared_backup", None) is not None
):
return
token_ids = req.fill_ids
keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
keys = page_align_keys(keys, self.page_size)
existing_prefix_len = self._probe_existing_radix_prefix_len_no_split(
RadixKey(
keys,
getattr(req, "extra_key", None),
is_bigram=self.is_eagle,
)
)
start = min(max(req.cache_protected_len, existing_prefix_len), len(keys))
if len(keys) <= start:
return
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, start : len(keys)
].to(dtype=torch.int64, copy=True)
if len(kv_indices) == 0:
return
node_id = TreeNode.counter
TreeNode.counter += 1
result = self._reserve_write_cp_indices_collectively(kv_indices, node_id)
if isinstance(result, HiCacheWriteFailure):
logger.info(
"[HiCache-write] prepare CP backup skipped after host reservation failure: node_id=%d rid=%s len=%d",
node_id,
getattr(req, "rid", "<unknown>"),
len(kv_indices),
)
return
try:
self.cache_controller.submit_write_cp_per_layer(
result, catch_up_all_layers=False
)
except Exception:
self.cache_controller.evict_cp_host(result.metadata)
raise
req.cp_hicache_prepared_backup = PreparedCpHiCacheBackup(
node_id=node_id,
reservation=result,
metadata=result.metadata,
logical_len=len(kv_indices),
)
logger.info(
"[HiCache-write] prepared CP per-layer backup before forward: node_id=%d rid=%s logical_len=%d owned_positions=%d",
node_id,
getattr(req, "rid", "<unknown>"),
len(kv_indices),
result.metadata.owned_positions.numel(),
)
def _node_host_len(self, node: TreeNode) -> int:
if self._uses_cp_hicache:
return node.host_len
@@ -1029,25 +1294,7 @@ class HiRadixCache(RadixCache):
write_back,
)
if self._uses_cp_hicache:
result = self.cache_controller.reserve_write_cp(
device_indices=node.value,
node_id=node.id,
)
if isinstance(result, HiCacheWriteFailure):
required_host_slots = result.required_host_slots
self._evict_host_for_physical_slots(
required_host_slots,
synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1,
)
logger.info(
"[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d",
node.id,
required_host_slots,
)
result = self.cache_controller.reserve_write_cp(
device_indices=node.value,
node_id=node.id,
)
result = self._reserve_write_cp_collectively(node)
if isinstance(result, HiCacheWriteFailure):
logger.info(
"[HiCache-write] write_backup CP FAILED (host full): node_id=%d len=%d",
@@ -1191,6 +1438,7 @@ class HiRadixCache(RadixCache):
if not finish_event.query():
break
finish_count += 1
local_finish_count = finish_count
queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu")
if self.tp_world_size > 1:
# synchronize TP workers to make the same update to radix cache
@@ -1201,11 +1449,11 @@ class HiRadixCache(RadixCache):
)
finish_count = int(queue_size.item())
logger.info(
logger.debug(
"[HiCache-write] writing_check: ongoing=%d ack_queue=%d local_finished=%d sync_finished=%d tp_size=%d",
len(self.ongoing_write_through),
ack_queue_len,
finish_count if self.tp_world_size <= 1 else queue_size.item(),
local_finish_count,
finish_count,
getattr(self, "tp_world_size", 1),
)
@@ -2264,6 +2512,7 @@ class HiRadixCache(RadixCache):
value = params.value
chunked = params.chunked
priority = params.priority
prepared_cp_backup = params.cp_hicache_prepared_backup
if priority is None:
priority = 0
@@ -2325,7 +2574,18 @@ class HiRadixCache(RadixCache):
child_key = self.get_child_key_fn(key)
if len(key):
new_node = TreeNode(priority=priority)
use_prepared_cp_backup = (
prepared_cp_backup is not None
and getattr(prepared_cp_backup, "logical_len", -1) == len(value)
)
new_node = TreeNode(
id=(
prepared_cp_backup.node_id
if use_prepared_cp_backup
else None
),
priority=priority,
)
new_node.parent = node
new_node.key = key
new_node.value = value.clone()
@@ -2342,7 +2602,10 @@ class HiRadixCache(RadixCache):
self._record_store_event(new_node)
if self.cache_controller.write_policy != "write_back":
self._inc_hit_count(new_node, chunked)
if use_prepared_cp_backup:
self._attach_prepared_cp_backup(new_node, prepared_cp_backup)
else:
self._inc_hit_count(new_node, chunked)
return InsertResult(prefix_len=total_prefix_length)
def release_aborted_request(self, rid: str):
python/sglang/srt/mem_cache/memory_pool.py
+3 -10
View File
@@ -695,7 +695,6 @@ class KVCache(abc.ABC):
# default state for optional layer-wise transfer control
self.layer_transfer_counter = None
self.layer_backup_notifiers = []
self.layer_backup_notify_after_indexer = False
# for disagg with nvlink
self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
@@ -750,11 +749,10 @@ class KVCache(abc.ABC):
def register_layer_backup_notifier(self, notifier):
self.layer_backup_notifiers.append(notifier)
def _notify_layer_kv_stored(self, layer_id: int, source: str = "kv"):
if self.layer_backup_notify_after_indexer and source != "indexer":
return
def notify_layer_end_for_backup(self, layer_id: int):
local_layer_id = layer_id - self.start_layer
for notifier in self.layer_backup_notifiers:
notifier(layer_id)
notifier(local_layer_id)
def get_cpu_copy(self, indices):
raise NotImplementedError()
@@ -1047,7 +1045,6 @@ class MHATokenToKVPool(KVCache):
alt_stream=self.alt_stream,
same_kv_dim=self.same_kv_dim,
)
self._notify_layer_kv_stored(layer_id - self.start_layer)
def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
if envs.SGLANG_NATIVE_MOVE_KV_CACHE.get():
@@ -1595,7 +1592,6 @@ class MLATokenToKVPool(KVCache):
)
else:
self.kv_buffer[layer_id - self.start_layer][loc] = cache_k
self._notify_layer_kv_stored(layer_id - self.start_layer)
def set_mla_kv_buffer(
self,
@@ -1637,7 +1633,6 @@ class MLATokenToKVPool(KVCache):
cache_k_nope,
cache_k_rope,
)
self._notify_layer_kv_stored(layer_id - self.start_layer)
def get_mla_kv_buffer(
self,
@@ -1859,7 +1854,6 @@ class NSATokenToKVPool(MLATokenToKVPool):
use_nsa=True,
override_kv_cache_dim=override_dim,
)
self.layer_backup_notify_after_indexer = True
# self.index_k_dtype = torch.float8_e4m3fn
# self.index_k_scale_dtype = torch.float32
self.index_head_dim = index_head_dim
@@ -1966,7 +1960,6 @@ class NSATokenToKVPool(MLATokenToKVPool):
index_buf_accessor.SetKAndS.execute(
pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale
)
self._notify_layer_kv_stored(layer_id - self.start_layer, source="indexer")
def get_state_buf_infos(self):
data_ptrs = [
python/sglang/srt/mem_cache/radix_cache.py
+29 -1
View File
@@ -489,22 +489,40 @@ class RadixCache(BasePrefixCache):
keys = page_align_keys(keys, self.page_size)
values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)
# Radix Cache takes one ref in memory pool
if is_insert:
priority = getattr(req, "priority", 0) or 0
result = self.insert(
InsertParams(key=radix_key, value=values, priority=priority)
InsertParams(
key=radix_key,
value=values,
priority=priority,
cp_hicache_prepared_backup=prepared_cp_backup,
)
)
if (
prepared_cp_backup is not None
and not getattr(prepared_cp_backup, "attached", False)
and hasattr(self, "_rollback_prepared_cp_backup")
):
self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss")
new_prefix_len = result.prefix_len
# Free the duplicates that were already in the tree
self.token_to_kv_pool_allocator.free(
kv_indices[req.cache_protected_len : new_prefix_len]
)
else:
if prepared_cp_backup is not None and hasattr(
self, "_rollback_prepared_cp_backup"
):
self._rollback_prepared_cp_backup(prepared_cp_backup, "no_insert")
self.token_to_kv_pool_allocator.free(
kv_indices[req.cache_protected_len : len(keys)]
)
if prepared_cp_backup is not None:
req.cp_hicache_prepared_backup = None
# free the unaligned tail
self.token_to_kv_pool_allocator.free(kv_indices[len(keys) :])
@@ -527,6 +545,7 @@ class RadixCache(BasePrefixCache):
keys = page_align_keys(keys, self.page_size)
values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)
# Radix Cache takes one ref in memory pool
result = self.insert(
@@ -535,8 +554,17 @@ class RadixCache(BasePrefixCache):
value=values,
chunked=chunked,
priority=getattr(req, "priority", 0) or 0,
cp_hicache_prepared_backup=prepared_cp_backup,
)
)
if (
prepared_cp_backup is not None
and not getattr(prepared_cp_backup, "attached", False)
and hasattr(self, "_rollback_prepared_cp_backup")
):
self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss")
if prepared_cp_backup is not None:
req.cp_hicache_prepared_backup = None
new_prefix_len = result.prefix_len
self.token_to_kv_pool_allocator.free(
python/sglang/srt/models/deepseek_v2.py
+20
View File
@@ -1620,6 +1620,22 @@ class DeepseekV2DecoderLayer(nn.Module):
and layer_id % self.config.moe_layer_freq == 0
)
def _notify_cp_hicache_layer_end(
self,
forward_batch: ForwardBatch,
tbo_subbatch_index: Optional[int] = None,
) -> None:
if (
getattr(forward_batch, "tbo_parent_token_range", None) is not None
and tbo_subbatch_index != 1
):
return
notifier = getattr(
forward_batch.token_to_kv_pool, "notify_layer_end_for_backup", None
)
if notifier is not None:
notifier(self.layer_id)
def forward(
self,
positions: torch.Tensor,
@@ -1707,6 +1723,7 @@ class DeepseekV2DecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
self._notify_cp_hicache_layer_end(forward_batch)
return hidden_states, residual
def op_comm_prepare_attn(
@@ -1761,6 +1778,9 @@ class DeepseekV2DecoderLayer(nn.Module):
state.pop("residual_after_comm_pre_mlp"),
state.forward_batch,
)
self._notify_cp_hicache_layer_end(
state.forward_batch, tbo_subbatch_index=state.tbo_subbatch_index
)
output = dict(
positions=state.positions,
test/registered/unit/managers/test_hicache_controller_cp.py
+8 -4
View File
@@ -175,6 +175,10 @@ class FakeDevicePool:
def register_layer_backup_notifier(self, notifier):
self.layer_backup_notifiers.append(notifier)
def notify_layer_end_for_backup(self, layer_id):
for notifier in self.layer_backup_notifiers:
notifier(layer_id)
class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
def test_mla_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self):
@@ -677,12 +681,12 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
self.assertEqual(host_pool.layer_backups, [])
self.assertEqual(controller.ack_write_queue, [])
controller.on_layer_kv_stored(0)
allocator.device_pool.notify_layer_end_for_backup(0)
self.assertEqual(host_pool.layer_backups[0][2], 0)
self.assertEqual(controller.ack_write_queue, [])
controller.on_layer_kv_stored(1)
allocator.device_pool.notify_layer_end_for_backup(1)
self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1])
self.assertEqual(len(controller.ack_write_queue), 1)
@@ -708,13 +712,13 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
reservation = controller.reserve_write_cp(logical_locs, node_id=85)
controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False)
controller.on_layer_kv_stored(0, source="target")
allocator.device_pool.notify_layer_end_for_backup(0)
self.assertEqual([x[2] for x in host_pool.layer_backups], [0])
self.assertEqual(draft_host_pool.layer_backups, [])
self.assertEqual(controller.ack_write_queue, [])
controller.on_layer_kv_stored(0, source="draft")
draft_device_pool.notify_layer_end_for_backup(0)
self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0])
self.assertEqual(len(controller.ack_write_queue), 1)
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+328 -3
View File
@@ -91,14 +91,20 @@ for _schema in (
raise
from sglang.srt.managers.cache_controller import (
HiCacheAck,
HiCacheWriteFailure,
HiCacheWriteReservation,
)
from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.hiradix_cache import (
CpHiCacheNodeMetadata,
HiRadixCache,
PendingHiCacheBackup,
PreparedCpHiCacheBackup,
_compute_shared_hicache_token_capacities,
)
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
@@ -544,18 +550,25 @@ class FakeReserveWriteController:
self.results = list(results)
self.reservations = []
self.submitted = []
self.submit_kwargs = []
self.evicted_host_indices = []
def reserve_write_cp(self, device_indices, priority=None, node_id=-1):
self.reservations.append((device_indices.clone(), node_id))
result = self.results.pop(0)
return result(device_indices) if callable(result) else result
if not callable(result):
return result
try:
return result(device_indices, node_id=node_id)
except TypeError:
return result(device_indices)
def submit_write_cp_all_layer(self, reservation):
self.submitted.append(reservation)
def submit_write_cp_per_layer(self, reservation):
def submit_write_cp_per_layer(self, reservation, **kwargs):
self.submitted.append(reservation)
self.submit_kwargs.append(kwargs)
def evict_cp_host(self, metadata):
self.evicted_host_indices.append(metadata.host_indices.clone())
@@ -883,6 +896,221 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
self.assertIn(node.id, cache.pending_host_backups)
self.assertEqual(len(cache.cache_controller.submitted), 1)
def test_write_backup_rolls_back_local_success_when_peer_needs_host_eviction(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.tp_world_size = 2
cache.tp_group = object()
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices: make_write_reservation(
device_indices, node_id=130, host_start=90
),
lambda device_indices: make_write_reservation(
device_indices, node_id=130, host_start=100
),
]
)
evictions = []
cache._evict_host_for_physical_slots = lambda required, synchronize_across_ranks=False: (
evictions.append((required, synchronize_across_ranks)) or required
)
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
cache.inc_node_lock_ref = lambda node: None
reduce_values = iter([4, 0])
def fake_all_reduce(tensor, op=None, group=None):
tensor.fill_(next(reduce_values))
node = TreeNode()
node.id = 130
node.value = torch.arange(16, dtype=torch.int64)
with patch("torch.distributed.all_reduce", side_effect=fake_all_reduce):
backed_len = cache.write_backup(node)
self.assertEqual(backed_len, 16)
# The first local success must be released before all ranks enter the
# collective host eviction/retry branch. Otherwise CP ranks can diverge:
# successful ranks submit backup while failing ranks enter eviction
# all_reduce, which matches the observed Gloo 4-vs-1 mismatch.
self.assertEqual(
cache.cache_controller.evicted_host_indices[0].tolist(),
list(range(90, 106)),
)
self.assertEqual(evictions, [(4, True)])
self.assertEqual(len(cache.cache_controller.submitted), 1)
self.assertEqual(
cache.cache_controller.submitted[0].metadata.host_indices.tolist(),
list(range(100, 116)),
)
def test_insert_attaches_prepared_cp_backup_without_catchup_copy(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.root_node = TreeNode()
cache.root_node.key = RadixKey([])
cache.evictable_size_ = 0
cache.get_child_key_fn = lambda key: key.token_ids[0]
cache.key_match_fn = lambda lhs, rhs: 0
cache.maybe_bigram_convert = lambda key, value: (key, value)
cache.is_eagle = False
cache.enable_storage = False
cache.enable_kv_cache_events = False
cache._update_leaf_status = lambda node: None
cache._update_host_leaf_status = lambda node: None
cache._record_store_event = lambda node: None
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
locked_nodes = []
cache.inc_node_lock_ref = lambda node: locked_nodes.append(node)
cache.cache_controller = types.SimpleNamespace(write_policy="write_through")
value = torch.arange(16, dtype=torch.int64)
reservation = make_write_reservation(value, node_id=131, host_start=110)
prepared = PreparedCpHiCacheBackup(
node_id=131,
reservation=reservation,
metadata=reservation.metadata,
logical_len=16,
)
result = cache.insert(
InsertParams(
key=RadixKey(list(range(16))),
value=value,
cp_hicache_prepared_backup=prepared,
)
)
node = cache.root_node.children[0]
self.assertEqual(result.prefix_len, 0)
self.assertEqual(node.id, 131)
self.assertTrue(prepared.attached)
self.assertIs(cache.ongoing_write_through[131], node)
self.assertIs(cache.pending_host_backups[131].node, node)
self.assertIs(cache.pending_host_backups[131].metadata, reservation.metadata)
self.assertEqual(cache.pending_host_backups[131].logical_len, 16)
self.assertEqual(locked_nodes, [node])
def test_prepare_write_backup_for_req_registers_before_forward_without_catchup(
self,
):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 1
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16)
)
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices, node_id: make_write_reservation(
device_indices, node_id=node_id, host_start=130
)
]
)
req = types.SimpleNamespace(
rid="rid-prepare",
fill_ids=list(range(16)),
cache_protected_len=0,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backup_for_req(req)
prepared = req.cp_hicache_prepared_backup
self.assertIsNotNone(prepared)
self.assertEqual(prepared.logical_len, 16)
self.assertEqual(len(cache.cache_controller.submitted), 1)
self.assertIs(cache.cache_controller.submitted[0], prepared.reservation)
self.assertEqual(
cache.cache_controller.submit_kwargs,
[{"catch_up_all_layers": False}],
)
def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 1
cache.root_node = TreeNode()
cache.root_node.key = RadixKey([])
cache.root_node.children = {}
cache.get_child_key_fn = lambda key: key.token_ids[0]
def key_match(lhs, rhs):
matched = 0
for left, right in zip(lhs.token_ids, rhs.token_ids):
if left != right:
break
matched += 1
return matched
cache.key_match_fn = key_match
cache.maybe_bigram_convert = lambda key, value=None: (key, value)
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16)
)
cache.cache_controller = FakeReserveWriteController([])
existing = TreeNode()
existing.id = 201
existing.parent = cache.root_node
existing.key = RadixKey(list(range(16)))
existing.value = torch.arange(16, dtype=torch.int64)
cache.root_node.children[0] = existing
req = types.SimpleNamespace(
rid="rid-existing-prefix",
fill_ids=list(range(16)),
extra_key=None,
cache_protected_len=8,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backup_for_req(req)
self.assertIsNone(req.cp_hicache_prepared_backup)
self.assertEqual(cache.cache_controller.reservations, [])
self.assertEqual(cache.cache_controller.submitted, [])
def test_rollback_unattached_prepared_cp_backup_removes_orphan_ack(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
value = torch.arange(4, dtype=torch.int64)
reservation = make_write_reservation(value, node_id=132, host_start=120)
prepared = PreparedCpHiCacheBackup(
node_id=132,
reservation=reservation,
metadata=reservation.metadata,
logical_len=4,
)
evicted = []
class ReadyEvent:
def synchronize(self):
pass
cache.cache_controller = types.SimpleNamespace(
pending_layer_writes={},
ack_write_queue=[HiCacheAck(ReadyEvent(), ReadyEvent(), [132])],
evict_cp_host=lambda metadata: evicted.append(metadata),
)
cache._rollback_prepared_cp_backup(prepared, "test")
self.assertEqual(cache.cache_controller.ack_write_queue, [])
self.assertEqual(evicted, [reservation.metadata])
def test_write_backup_cp_retry_failure_leaves_node_device_only(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
@@ -1641,5 +1869,102 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertEqual(cache.ongoing_load_back, {})
class TestCPHiCacheLayerBackupNotifications(CustomTestCase):
def _call_store_index_fast_path(self, out_loc):
from sglang.srt.layers.attention.nsa import nsa_indexer
notifications = []
fused_calls = []
class FakePool:
page_size = 64
start_layer = 0
def get_index_k_with_scale_buffer(self, layer_id):
return torch.empty((1,), dtype=torch.uint8)
def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"):
notifications.append((layer_id, source))
forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool())
indexer = types.SimpleNamespace()
def fake_fused_store(key, buf, loc, page_size):
fused_calls.append((key, buf, loc.clone(), page_size))
with (
patch.object(nsa_indexer, "_is_cuda", True),
patch.object(nsa_indexer, "_is_fp8_fnuz", False),
patch.object(nsa_indexer, "can_use_nsa_fused_store", return_value=True),
patch.object(
nsa_indexer, "fused_store_index_k_cache", side_effect=fake_fused_store
),
):
nsa_indexer.Indexer._store_index_k_cache(
indexer,
forward_batch,
layer_id=5,
key=torch.empty((max(out_loc.numel(), 1), 128), dtype=torch.float32),
out_loc_override=out_loc,
)
return notifications, fused_calls
def test_nsa_indexer_fused_store_does_not_notify_cp_hicache_layer_backup(self):
notifications, fused_calls = self._call_store_index_fast_path(
torch.tensor([1, 2, 3], dtype=torch.int64)
)
self.assertEqual(notifications, [])
self.assertEqual(len(fused_calls), 1)
self.assertEqual(fused_calls[0][2].tolist(), [1, 2, 3])
def test_nsa_indexer_empty_store_does_not_notify_cp_hicache_layer_backup(self):
notifications, fused_calls = self._call_store_index_fast_path(
torch.empty((0,), dtype=torch.int64)
)
self.assertEqual(notifications, [])
self.assertEqual(fused_calls, [])
def test_cp_shared_zero_local_index_store_does_not_notify_layer_backup(self):
from sglang.srt.layers.attention.nsa import nsa_indexer
notifications = []
class FakePool:
page_size = 64
def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"):
notifications.append((layer_id, source))
forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool())
indexer = types.SimpleNamespace(nsa_enable_prefill_cp=True)
with (
patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True),
patch.object(
nsa_indexer,
"get_cp_shared_kv_local_out_cache_loc",
return_value=torch.empty((0,), dtype=torch.int64),
),
patch.object(
nsa_indexer,
"get_cp_shared_kv_local_physical_out_cache_loc",
side_effect=AssertionError("must not require physical locs"),
),
):
handled = nsa_indexer.Indexer._store_cp_shared_local_index_k_cache(
indexer,
forward_batch=forward_batch,
layer_id=6,
local_key=torch.empty((0, 128), dtype=torch.float32),
act_quant=None,
)
self.assertTrue(handled)
self.assertEqual(notifications, [])
if __name__ == "__main__":
unittest.main()