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

CP HiCache: kill the O(victims*candidates) GPU-sync storm in the owner-lane eviction planner

Browse Source 
A CP4 prefill server hung 79s in _plan_cp_load_back_owner_lane_evictions (the [HiCache-load]
slow-scan) and was killed by the detokenizer health-check. Root cause (verified from the b300
hang log + code): the L1 free-room watermark (--hicache-l1-free-room-ratio 0.25 over a 31507-page
lane) hands the planner a ~7877-page single-owner deficit; the planner then rescans all ~2000
evictable candidates once per victim (~195 iters), and for EACH candidate every iteration it
recomputes _cp_load_back_node_owner_page_counts -- which under the pooled shared-L2 path (no
page_owners on CpSharedL2NodeMetadata) takes the device-tensor fallback and does cp_size per-owner
.item() device->host syncs. That is ~195 * 2000 * 4 ~= 1.5M CUDA syncs on the synchronous load-back
admission path, blocking the scheduler for ~79s.

Fix (byte-identical victim selection, just fast):
- Memoize the owner-count histogram per planning call ({node.id: counts}); the counts are invariant
  while node.value is fixed (the plan does not mutate values), so node.id is a safe key for the plan's
  duration. Threaded explicitly to both call sites (planner loop + ancestor-unlock helper). Turns
  O(victims*candidates) recomputes into O(distinct nodes).
- Replace the cp_size per-owner sum().item() loop with one bincount().tolist() device->host sync.

Net: ~79s -> ~1s; the eviction plan (victims, planned_freed) is unchanged. bincount == the per-owner
loop proven over 8000 random vectors incl. the (-1)%cp==cp-1 zero-loc edge. New memo regression test;
existing count-fn + planner tests pass (the one pre-existing unrelated EAGLE-tail failure is unchanged).
(The 7877-page watermark magnitude is a separate, config-side issue: --hicache-l1-free-room-ratio 0.25
reserves ~25% of a ~2M-token lane -- ~30x more than a 64K chunk needs; lower it.)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-24 01:13:59 +00:00
parent ff1804f5cf
commit 4e9b4d05c0
2 changed files with 87 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -2496,28 +2496,52 @@ class HiRadixCache(RadixCache):
return self._cp_device_node_is_load_back_victim_after_plan(node, set())
def _cp_load_back_node_owner_page_counts(
self, node: TreeNode, cp_size: int
self, node: TreeNode, cp_size: int, *, memo: Optional[dict] = None
) -> Tuple[int, ...]:
"""Per-owner device page-count histogram for a node.
The owner-lane eviction planner calls this O(victims * candidates) times over an
UNCHANGING node set (the plan does not mutate node.value), so the per-node counts
are invariant for the whole planning call. Pass a planner-scoped ``memo``
({node.id: counts}) to compute each node's counts exactly once -- node.id is a safe
key for the duration of one plan. Without ``memo`` the behaviour is unchanged.
The histogram itself uses a single ``bincount().tolist()`` device->host sync instead
of ``cp_size`` per-owner ``.item()`` syncs -- under the pooled shared-L2 path (no
``page_owners`` on the metadata) this fallback is taken for every candidate, so the
per-owner sync loop (re-run O(victims * candidates) times) was the dominant cost.
"""
if memo is not None:
cached = memo.get(node.id)
if cached is not None:
return cached
metadata = getattr(node, "cp_hicache", None)
if metadata is not None and not self._is_cp_shared_l2_metadata(metadata):
return metadata.owner_page_counts(cp_size)
value = getattr(node, "value", None)
if value is None or int(value.numel()) == 0:
return tuple(0 for _ in range(cp_size))
padded_value = pad_token_locs_to_page_boundary(
value,
self.page_size,
name=(
"CP HiCache load-back eviction device values "
f"node_id={getattr(node, 'id', '?')}"
),
)
first_locs = padded_value[:: self.page_size]
logical_pages = torch.div(first_locs, self.page_size, rounding_mode="floor")
owners = torch.remainder(logical_pages - 1, cp_size)
return tuple(int((owners == owner).sum().item()) for owner in range(cp_size))
counts = metadata.owner_page_counts(cp_size)
else:
value = getattr(node, "value", None)
if value is None or int(value.numel()) == 0:
counts = tuple(0 for _ in range(cp_size))
else:
padded_value = pad_token_locs_to_page_boundary(
value,
self.page_size,
name=(
"CP HiCache load-back eviction device values "
f"node_id={getattr(node, 'id', '?')}"
),
)
first_locs = padded_value[:: self.page_size]
logical_pages = torch.div(
first_locs, self.page_size, rounding_mode="floor"
)
owners = torch.remainder(logical_pages - 1, cp_size)
# one device->host sync (bincount over [0, cp_size)) == the per-owner
# sum().item() loop, byte-identical counts, cp_size-fewer syncs.
counts = tuple(torch.bincount(owners, minlength=cp_size).tolist())
if memo is not None:
memo[node.id] = counts
return counts
def _cp_load_back_ancestor_unlock_contribution(
self,
@@ -2525,6 +2549,8 @@ class HiRadixCache(RadixCache):
deficits: List[int],
planned_evicted_nodes: set,
cp_size: int,
*,
memo: Optional[dict] = None,
) -> int:
ancestor = getattr(node, "parent", None)
while ancestor is not None and ancestor != getattr(self, "root_node", None):
@@ -2535,7 +2561,9 @@ class HiRadixCache(RadixCache):
continue
if not self._cp_device_node_is_load_back_victim_base(ancestor):
return 0
counts = self._cp_load_back_node_owner_page_counts(ancestor, cp_size)
counts = self._cp_load_back_node_owner_page_counts(
ancestor, cp_size, memo=memo
)
contribution = sum(
min(int(count), int(deficit))
for count, deficit in zip(counts, deficits)
@@ -2558,6 +2586,10 @@ class HiRadixCache(RadixCache):
candidate_nodes = set(getattr(self, "evictable_leaves", set()))
initial_candidate_count = len(candidate_nodes)
iteration = 0
# Per-plan memo {node.id: owner_counts}: the planner rescans the same node set
# once per victim, but each node's owner counts are invariant for the whole plan,
# so compute the (device-sync) histogram once instead of O(victims * candidates).
owner_counts_memo: dict = {}
while any(v > 0 for v in deficits):
iteration += 1
@@ -2589,7 +2621,9 @@ class HiRadixCache(RadixCache):
node, planned_evicted_nodes
):
continue
counts = self._cp_load_back_node_owner_page_counts(node, cp_size)
counts = self._cp_load_back_node_owner_page_counts(
node, cp_size, memo=owner_counts_memo
)
contribution = sum(
min(int(count), int(deficit))
for count, deficit in zip(counts, deficits)
@@ -2598,7 +2632,11 @@ class HiRadixCache(RadixCache):
if contribution <= 0:
unlock_contribution = (
self._cp_load_back_ancestor_unlock_contribution(
node, deficits, planned_evicted_nodes, cp_size
node,
deficits,
planned_evicted_nodes,
cp_size,
memo=owner_counts_memo,
)
)
if contribution <= 0 and unlock_contribution <= 0: