minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix(hicache): limit load-back pending to write-back (#34519)

Browse Source 
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
ziang663
2026-08-17 14:23:54 +08:00
committed by GitHub
co-authored by Zhangheng
parent eafbe2cb6f
commit 43226af812
2 changed files with 103 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+26 -18
View File
@@ -1958,19 +1958,20 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
rebuild is deferred to the orchestration layer."""
node = self.node_by_id(node_id)
cache_actions: list[CacheAction | ComponentAction] = []
# Pin every node whose host slots the in-flight DMA reads (including
# aux-only nodes) against reclaim until the ack.
for xfers in ([kv_xfer], *comp_xfers.values()):
for xfer in xfers:
for nid in xfer.nodes_to_load or ():
pinned = self.node_by_id(nid)
# One live load-back per node; only the same anchor may
# re-pin (a node can sit in Full and aux transfer lists).
assert pinned.load_back_pending_id in (None, node_id), (
f"node {nid} pinned by load-back "
f"{pinned.load_back_pending_id}, new anchor {node_id}"
)
pinned.load_back_pending_id = node_id
if self.is_write_back:
# Write-back may reclaim a duplicate host copy while H->D DMA is
# still reading it, so pin every source node until the ack.
for xfers in ([kv_xfer], *comp_xfers.values()):
for xfer in xfers:
for nid in xfer.nodes_to_load or ():
pinned = self.node_by_id(nid)
# One live load-back per node; only the same anchor may
# re-pin (a node can sit in Full and aux transfer lists).
assert pinned.load_back_pending_id in (None, node_id), (
f"node {nid} pinned by load-back "
f"{pinned.load_back_pending_id}, new anchor {node_id}"
)
pinned.load_back_pending_id = node_id
kv_xfer.device_indices = device_indices
self.components_by_type[BASE_COMPONENT_TYPE].commit_hicache_transfer(
node,
@@ -1991,14 +1992,21 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
return cache_actions
def finish_load_back(self, anchor_node_id: NodeId) -> None:
"""Clear the in-flight H->D marks along the anchor's root path at ack
time; split fragments stay on the path, so the walk covers them."""
"""Finalize H->D load-back state along the anchor's root path.
Write-back clears source-node pins at ack time. Write-through does not
use those pins, but still refreshes duplicate tracking after the device
copies become visible. Split fragments stay on the path, so the walk
covers them.
"""
node = self.node_by_id(anchor_node_id)
while node is not None and node is not self.root_node:
if node.load_back_pending_id == anchor_node_id:
if self.is_write_back:
if node.load_back_pending_id != anchor_node_id:
node = node.parent
continue
node.load_back_pending_id = None
# The loaded copies become tracked duplicates only now.
self._update_duplicate_tracking(node)
self._update_duplicate_tracking(node)
node = node.parent
def mark_write_through_pending(self, node_id: NodeId) -> None: