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

Fix CP HiCache load_cp owner-pattern mismatch (cache-hit corruption)

Browse Source 
In CP=8 + NSA-shared-KV + HiCache disagg-prefill, cache-hit prefill produced
incoherent decode output. Cold prefill on CP was correct; pure CP without
HiCache was correct. The bug lived at the HiCache load_cp / device-alloc
interface.

Root cause: cache_controller.load_cp called the plain
mem_pool_device_allocator.alloc(logical_len), which returns logical pages
with no CP owner-pattern preservation. Cold prefill instead uses
alloc_extend_compute_owner with a zigzag owner pattern from
build_in_seq_page_compute_owners. The saved CpHiCacheNodeMetadata.owned_positions
records WHICH POSITIONS in the write-time alloc were owned by this rank. At
load time, those same positions are applied to a new alloc whose per-position
owner pattern is arbitrary -- each rank loads its host bytes into physical
slots whose corresponding logical page is owned by a DIFFERENT rank.
Attention's materialize_shared_token_kv_buffer reads from the owner's
physical slot, which was never loaded. Result: garbage.

Fix:
- CpHiCacheNodeMetadata gains two required fields: page_owners (int8 per
  logical page, identical on all CP ranks) and page_size. __post_init__
  validates; split() bisects page_owners by page index with a page-alignment
  check.
- _write_cp derives page_owners from device_indices (page-first slot of each
  page -> logical page id -> layout.owner_for_logical_pages) and stores in
  both metadata-construction sites (zero-owned and normal).
- New CPSharedPagedTokenToKVPoolAllocator.alloc_pages_with_owners() reuses
  _select_compute_owner_pages (with its tai-kernel fast path) and returns
  page-contiguous token locs whose per-page owner sequence equals the input.
- load_cp now concats page_owners across nodes_to_load and calls
  alloc_pages_with_owners. On None (lane exhausted) the caller hits the
  retry-with-eviction path; further failure returns None and degrades to
  cache miss. No silent fallback to plain alloc -- that recreated the bug.
- load_back retry path now calls _evict_for_compute_owner_lanes (module-top
  import) instead of plain evict(); this targets the deficit lane and gives
  the next alloc attempt a chance to satisfy it.
- envs import moved to module top in cache_controller.py per code-review
  feedback. Removed an over-defensive owned_check.all().item() in load_cp
  that would have re-introduced the host-sync anti-pattern 97a9f850c
  removed -- the invariant is already guaranteed by alloc_pages_with_owners.

Tests: 40 existing CpHiCacheNodeMetadata constructions migrated to pass the
new required fields. 9 new metadata tests (validators + split page-alignment).
10 new allocator tests in test_alloc_pages_with_owners.py covering input-order
preservation, lane exhaustion, release_pages fallback, debug-mode invariant.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-05-25 01:11:07 +08:00
parent f2834b3403
commit 1d630def95
6 changed files with 531 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
python/sglang/srt/managers/cache_controller.py
View File

@@ -32,6 +32,7 @@ from sglang.srt.distributed import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
)
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
get_attention_dp_rank,
get_attention_tp_rank,
@@ -836,6 +837,8 @@ class HiCacheController:
host_indices: torch.Tensor,
device_indices: torch.Tensor,
) -> None:
if not envs.SGLANG_DEBUG_HICACHE_VALIDATE.get():
return
validate_page_aligned_token_indices(
device_indices, self.page_size, "physical_device_indices"
)
@@ -853,6 +856,23 @@ class HiCacheController:
owned_mask = layout.owned_by_this_rank(device_indices)
owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu()
logical_len = len(device_indices)
# 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_size = self.page_size
if logical_len % page_size != 0:
raise ValueError(
f"_write_cp expects page-aligned device_indices, got "
f"logical_len={logical_len} page_size={page_size}"
)
page_first_locs = 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:
self._append_completed_write_ack(node_id)
logger.info(
@@ -865,6 +885,8 @@ class HiCacheController:
logical_len=logical_len,
owned_positions=owned_positions,
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=page_owners,
page_size=page_size,
draft_host_indices=(
torch.empty((0,), dtype=torch.int64)
if self.has_draft_hicache
@@ -935,6 +957,8 @@ class HiCacheController:
logical_len=logical_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
),
@@ -975,11 +999,42 @@ class HiCacheController:
return device_indices
def load_cp(self, nodes_to_load, node_id: int = -1) -> Optional[torch.Tensor]:
logical_len = sum(node.host_len for node in nodes_to_load)
device_indices = self.mem_pool_device_allocator.alloc(logical_len)
# 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())
device_indices = self.mem_pool_device_allocator.alloc_pages_with_owners(
page_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:
return None
logical_len_expected = sum(node.host_len for node in nodes_to_load)
if device_indices.numel() != logical_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 {logical_len_expected}"
)
host_chunks = []
draft_host_chunks = []
physical_chunks = []
@@ -991,6 +1046,18 @@ class HiCacheController:
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)
)

52
python/sglang/srt/mem_cache/allocator.py
View File

@@ -763,3 +763,55 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
assert torch.equal(selected_owners, expected_owners)
return out_indices
def alloc_pages_with_owners(
self,
page_compute_owners: List[int],
) -> Optional[torch.Tensor]:
"""Allocate ``len(page_compute_owners)`` whole logical pages where
page ``i``'s owner equals ``page_compute_owners[i]``.
Returns a flat int64 tensor of logical token locs of length
``len(owners) * page_size``, page-contiguous in input order, or
``None`` when an owner lane is exhausted (caller handles
eviction-retry / cache-miss).
Used by CP HiCache ``load_cp`` to reproduce the write-time owner
pattern recorded in ``CpHiCacheNodeMetadata.page_owners``. Without
this, plain ``alloc()`` would return pages with an arbitrary owner
pattern, the saved per-position owner mask would index the wrong
slots, and each rank would load its host bytes into physical slots
whose corresponding logical page is owned by another rank → attention
reads garbage.
Distinguished from ``alloc_extend_compute_owner``: that variant
expects extend semantics (prefix_lens / seq_lens / last_loc) and
invokes the ``alloc_extend_naive`` Triton kernel to splice the
partial last page of a prefix with new pages. HiCache reload has
no prefix and no partial page — host pages are always page-aligned
— so a pure page-fresh allocation is enough.
"""
if not page_compute_owners:
return torch.empty((0,), dtype=torch.int64, device=self.device)
selected = self._select_compute_owner_pages(page_compute_owners)
if selected is None:
return None
selected_pages, selected_free_mask, selected_release_mask = selected
page_size = self.page_size
base = selected_pages.to(torch.int64).unsqueeze(1) * page_size
offsets = torch.arange(
page_size, dtype=torch.int64, device=self.device
).unsqueeze(0)
out_indices = (base + offsets).reshape(-1)
self.free_pages = self.free_pages[~selected_free_mask]
self.release_pages = self.release_pages[~selected_release_mask]
if self.debug_mode:
assert torch.unique(out_indices).numel() == out_indices.numel()
check_owners = torch.remainder(selected_pages - 1, self.cp_size)
expected = torch.tensor(
page_compute_owners,
dtype=check_owners.dtype,
device=check_owners.device,
)
assert torch.equal(check_owners, expected)
return out_indices

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

@@ -28,6 +28,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams,
MatchResult,
)
from sglang.srt.mem_cache.common import _evict_for_compute_owner_lanes
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.memory_pool import (
MHATokenToKVPool,
@@ -57,21 +58,50 @@ class CpHiCacheNodeMetadata:
logical_len: int
owned_positions: torch.Tensor
host_indices: torch.Tensor
# NEW: one int8 per logical PAGE; identical on all CP ranks (function of
# the original logical page ids only). Captures the owner pattern of the
# write-time allocation so load_cp can reproduce it via owner-aware alloc.
# Without this, the saved owned_positions index a fresh alloc whose
# per-position owner pattern is arbitrary → load writes to wrong physical
# slots → attention reads garbage. See _write_cp for the derivation and
# alloc_pages_with_owners for the consumer.
page_owners: torch.Tensor
# NEW: needed at split() time to convert split_len into a page index
# without plumbing it from the caller.
page_size: int
draft_host_indices: Optional[torch.Tensor] = None
def __post_init__(self):
if self.logical_len < 0:
raise ValueError(f"logical_len must be non-negative, got {self.logical_len}")
if self.page_size <= 0:
raise ValueError(f"page_size must be positive, got {self.page_size}")
if self.logical_len % self.page_size != 0:
raise ValueError(
f"logical_len ({self.logical_len}) must be a multiple of "
f"page_size ({self.page_size})"
)
self.owned_positions = self.owned_positions.to(
device="cpu", dtype=torch.int64
).detach().clone()
self.host_indices = self.host_indices.to(
device="cpu", dtype=torch.int64
).detach().clone()
self.page_owners = self.page_owners.to(
device="cpu", dtype=torch.int8
).detach().clone()
if self.draft_host_indices is not None:
self.draft_host_indices = self.draft_host_indices.to(
device="cpu", dtype=torch.int64
).detach().clone()
expected_num_pages = self.logical_len // self.page_size
if self.page_owners.numel() != expected_num_pages:
raise ValueError(
f"page_owners length ({self.page_owners.numel()}) must equal "
f"logical_len/page_size ({expected_num_pages})"
)
if expected_num_pages > 0 and bool((self.page_owners < 0).any()):
raise ValueError("page_owners entries must be non-negative")
if self.owned_positions.numel() != self.host_indices.numel():
raise ValueError(
"owned_positions and host_indices must have same length, got "
@@ -102,13 +132,21 @@ class CpHiCacheNodeMetadata:
raise ValueError(
f"split_len must be in [0, {self.logical_len}], got {split_len}"
)
if split_len % self.page_size != 0:
raise ValueError(
f"split_len ({split_len}) must be a multiple of page_size "
f"({self.page_size})"
)
parent_mask = self.owned_positions < split_len
child_mask = ~parent_mask
split_pages = split_len // self.page_size
return (
CpHiCacheNodeMetadata(
logical_len=split_len,
owned_positions=self.owned_positions[parent_mask],
host_indices=self.host_indices[parent_mask],
page_owners=self.page_owners[:split_pages],
page_size=self.page_size,
draft_host_indices=(
self.draft_host_indices[parent_mask]
if self.draft_host_indices is not None
@@ -119,6 +157,8 @@ class CpHiCacheNodeMetadata:
logical_len=self.logical_len - split_len,
owned_positions=self.owned_positions[child_mask] - split_len,
host_indices=self.host_indices[child_mask],
page_owners=self.page_owners[split_pages:],
page_size=self.page_size,
draft_host_indices=(
self.draft_host_indices[child_mask]
if self.draft_host_indices is not None
@@ -1414,11 +1454,25 @@ class HiRadixCache(RadixCache):
)
if device_indices is None:
logger.info(
"[HiCache-load] load_back CP retry with eviction: node_id=%d tokens_needed=%d",
"[HiCache-load] load_back CP retry with lane-aware eviction: "
"node_id=%d tokens_needed=%d",
last_hit_node.id,
host_hit_len,
)
self.evict(EvictParams(num_tokens=host_hit_len))
# Lane-aware eviction: alloc_pages_with_owners failed because
# at least one owner lane is short of free pages. Targeted
# eviction frees pages from THOSE specific lanes, leaving
# other lanes untouched. Mirrors cold prefill's retry path
# in common.py:alloc_paged_token_slots_extend.
_retry_page_owners: List[int] = []
for _node in nodes_to_load:
# page_owners is CPU int8; tolist() returns list[int].
_retry_page_owners.extend(_node.cp_hicache.page_owners.tolist())
_evict_for_compute_owner_lanes(
tree_cache=self,
allocator=self.token_to_kv_pool_allocator,
page_compute_owners=_retry_page_owners,
)
device_indices = self.cache_controller.load_cp(
nodes_to_load, node_id=last_hit_node.id
)