Keep CP HiCache reuse page-safe without eviction log churn

Page-aligned CP shared KV can pad out_cache_loc beyond valid current rows, so current reuse now gates MLA composition on the valid extend rows and permits draft partial-current reuse once the TAI sparse-page capability check passes. The TAI current-slot path self-tests sparse pages before use and falls back to the torch reference when the installed kernel is stale.

Eviction success and no-op diagnostics were also moved from INFO to DEBUG so owner-lane and host-admission churn does not flood production logs; true write failures remain WARNING.

Constraint: CP shared KV uses page-aligned physical reservations where valid suffix rows can be shorter than padded out_cache_loc.

Constraint: Production failure/fallback logs must remain visible, but hot successful eviction paths should not emit INFO per victim/rank.

Rejected: Keep draft partial-current reuse disabled | would preserve avoidable full materialization on draft cache-hit suffixes.

Rejected: Trust the TAI current-slot kernel unconditionally | stale kernels can corrupt sparse current-page composition.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce INFO logging in eviction hot paths without rate limiting and runtime evidence.

Tested: local py_compile for touched Python files

Tested: local git diff --check

Tested: remote container py_compile for touched Python files

Tested: remote PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py::TestHiCacheEvictLoggingLevels::test_evict_hot_path_success_logs_are_debug_only test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 82 passed, 5 warnings, 2 subtests passed

Not-tested: full ETE traffic after this commit; draft partial-current accept length still needs user-driven runtime validation
This commit is contained in:
laoyao0822
2026-05-31 03:31:06 +08:00
parent 3c14b1f127
commit 0fc95b6439
7 changed files with 704 additions and 119 deletions

View File

@@ -9,7 +9,9 @@ from typing import Any
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import log_cp_draft_shared_kv_debug # noqa: F401
from sglang.srt.layers.attention.nsa.utils import (
log_cp_draft_shared_kv_debug,
) # noqa: F401
from sglang.srt.layers.dp_attention import get_attention_cp_group
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -85,11 +87,7 @@ def cp_shared_kv_mla_prefetch_min_prefix_pages(
if page_size is not None and int(page_size) > 0:
min_pages = max(
min_pages,
(
_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS
+ int(page_size)
- 1
)
(_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS + int(page_size) - 1)
// int(page_size),
)
return min_pages
@@ -302,6 +300,110 @@ def _load_tai_materialize_kernels():
return None
@lru_cache(maxsize=16)
def _tai_current_slot_fill_sparse_pages_self_test(
device_type: str,
device_index: int | None,
) -> bool:
"""Return whether TAI current-slot fill handles sparse current pages.
Older tai-kernel builds filled current rows correctly but masked every page
between first and last current page as current-page slack. That silently
hid unrelated prefix pages for cache-hit suffixes and was enough to corrupt
CP shared-KV reuse. Validate the installed kernel once before allowing it
on the hot path.
"""
if device_type != "cuda":
return True
if not torch.cuda.is_available():
return False
kernels = _load_tai_materialize_kernels()
fill_kernel = (
getattr(kernels, "fill_current_token_kv_page_slots_and_remap_locs", None)
if kernels is not None
else None
)
if fill_kernel is None:
return False
device = torch.device(device_type, device_index)
try:
page_size = 4
dense_kv = torch.zeros((16, 1), device=device, dtype=torch.float32)
materialized_locs = torch.tensor(
[[4, 5, 8, 9, 12, 13, 14, 15]],
device=device,
dtype=torch.int64,
)
current_kv = torch.tensor(
[[10.0], [11.0], [12.0], [13.0]],
device=device,
dtype=torch.float32,
)
logical_locs = torch.tensor(
[[20, 21, 40, 41, 100, 101, 102, 103]],
device=device,
dtype=torch.int64,
)
current_locs = torch.tensor(
[20, 21, 100, 101],
device=device,
dtype=torch.int64,
)
page_inverse = torch.full((32,), -1, device=device, dtype=torch.long)
page_inverse[0] = 0
page_inverse[5] = 1
page_inverse[10] = 2
page_inverse[25] = 3
mixed_kv, mixed_locs, current_mask = fill_kernel(
dense_kv,
materialized_locs,
current_kv,
logical_locs,
current_locs,
page_inverse,
page_size=page_size,
mask_non_current_in_current_pages=True,
)
expected_locs = torch.tensor(
[[4, 5, 8, 9, 12, 13, -1, -1]],
device=device,
dtype=torch.int64,
)
expected_mask = torch.tensor(
[[True, True, False, False, True, True, False, False]],
device=device,
dtype=torch.bool,
)
expected_kv = torch.zeros_like(dense_kv)
expected_kv[4:6] = current_kv[0:2]
expected_kv[12:14] = current_kv[2:4]
return bool(
torch.equal(mixed_locs, expected_locs)
and torch.equal(current_mask, expected_mask)
and torch.equal(mixed_kv, expected_kv)
)
except Exception as exc:
_log_tai_materialize_fallback(
"fill_current_sparse_page_self_test_failed",
"CP shared KV tai current-slot fill sparse-page self-test failed; "
"falling back to torch reference. error=%s",
exc,
limit=1,
)
return False
def _tai_current_slot_fill_supports_sparse_pages(device: torch.device) -> bool:
return _tai_current_slot_fill_sparse_pages_self_test(
device.type,
device.index,
)
def _tai_materialize_runtime_enabled() -> bool:
# Keep the debug path on the existing torch implementation. The debug path
# intentionally preserves tensor summaries and value assertions used for
@@ -719,6 +821,22 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs(
)
return None
if dense_kv_cache.is_cuda and not _tai_current_slot_fill_supports_sparse_pages(
dense_kv_cache.device
):
_log_tai_materialize_fallback(
"fill_current_sparse_page_unsupported",
"CP shared KV tai current-slot fill kernel failed the sparse-page "
"capability check; falling back to torch reference to avoid hiding "
"prefix pages in partial-current reuse. page_size=%s current_rows=%s "
"query_locs=%s",
page_size,
int(current_kv_cache.shape[0]),
int(logical_locs.numel()),
limit=1,
)
return None
try:
return fill_kernel(
_contiguous_for_tai(dense_kv_cache),
@@ -866,8 +984,10 @@ def fill_current_index_page_slots(
)
current_pages = torch.div(current_locs, page_size, rounding_mode="floor")
valid_pages = (current_locs >= 0) & (current_pages >= 0) & (
current_pages < int(page_inverse.numel())
valid_pages = (
(current_locs >= 0)
& (current_pages >= 0)
& (current_pages < int(page_inverse.numel()))
)
safe_pages = torch.clamp(
current_pages,
@@ -1075,12 +1195,12 @@ def can_reuse_current_extend_kv(forward_batch) -> bool:
def should_reuse_current_extend_kv(forward_batch) -> bool:
"""Return whether MLA should splice current extend KV into materialized KV.
Current-only reuse is safe for both target and draft because there is no
cached prefix to compose. Partial current reuse is currently a target-model
contract only. EAGLE/NextN draft cache-hit suffixes keep using the older
full-materialize path until the draft splice path has value-level ETE proof;
the 2026-05-30 accept-length regression correlated with enabling that draft
partial splice path.
Current-only and partial-current reuse are both cache-layout operations: the
cached page-aligned prefix is materialized from the shared KV cache and the
current valid suffix is spliced from ``out_cache_loc``. Draft/NextN uses the
same CP shared-KV layout contract as the target model. The stale TAI
sparse-current-page corruption that made this unsafe is blocked at the
current-slot fill capability check before any TAI result is used.
"""
if not cp_shared_kv_current_reuse_enabled():
@@ -1091,25 +1211,40 @@ def should_reuse_current_extend_kv(forward_batch) -> bool:
return True
partial_current = can_reuse_current_extend_kv(forward_batch)
if partial_current and cp_shared_kv_is_draft_input(forward_batch):
prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
extend_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
_log_current_reuse_fallback(
"draft_partial_current_reuse_disabled",
"cache-hit EAGLE/NextN draft uses full materialize instead of "
"partial current reuse. prefix_lens=%s extend_lens=%s",
[int(x) for x in prefix_lens_cpu]
if prefix_lens_cpu is not None
else None,
[int(x) for x in extend_lens_cpu]
if extend_lens_cpu is not None
else None,
)
return False
return current_only or partial_current
def current_extend_kv_rows_for_reuse(
forward_batch,
*current_kv_tensors: torch.Tensor | None,
) -> int | None:
"""Return valid current rows if current KV tensors can be spliced.
``out_cache_loc`` may be page padded while MLA/index projections only produce
valid-token rows. Partial-current reuse should therefore validate tensor
rows against ``extend_seq_lens_cpu`` rather than the padded loc tensor length.
"""
if not should_reuse_current_extend_kv(forward_batch):
return None
extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
if extend_seq_lens_cpu is None or len(extend_seq_lens_cpu) != 1:
return None
valid_current_rows = int(extend_seq_lens_cpu[0])
if valid_current_rows <= 0:
return None
out_cache_loc = getattr(forward_batch, "out_cache_loc", None)
if out_cache_loc is None or int(out_cache_loc.numel()) < valid_current_rows:
return None
for tensor in current_kv_tensors:
if tensor is None or int(tensor.shape[0]) < valid_current_rows:
return None
return valid_current_rows
def current_loc_remap_fast_path_args(
forward_batch,
) -> tuple[int | None, int | None]:
@@ -1366,7 +1501,9 @@ def _debug_assert_no_negative_tensor_values(
)
def build_dense_page_remap(logical_pages: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
def build_dense_page_remap(
logical_pages: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build a dense per-call page remap for shared-KV runtime materialization."""
dense_pages = logical_pages.clone()
positive_mask = logical_pages > 0
@@ -1389,17 +1526,25 @@ def remap_logical_pages_to_dense_pages(
insert_positions = torch.searchsorted(unique_logical_pages, positive_pages)
if cp_shared_kv_debug_enabled() and insert_positions.numel() > 0:
if unique_logical_pages.numel() == 0:
raise ValueError("unique_logical_pages is empty but logical_pages contains data")
raise ValueError(
"unique_logical_pages is empty but logical_pages contains data"
)
if torch.any(insert_positions >= unique_logical_pages.numel()):
raise ValueError("logical_pages contains entries outside unique_logical_pages")
raise ValueError(
"logical_pages contains entries outside unique_logical_pages"
)
if not torch.equal(unique_logical_pages[insert_positions], positive_pages):
raise ValueError("logical_pages contains entries outside unique_logical_pages")
raise ValueError(
"logical_pages contains entries outside unique_logical_pages"
)
dense_pages[positive_mask] = insert_positions.to(dense_pages.dtype) + 1
return dense_pages
def build_slot_page_remap(logical_pages: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
def build_slot_page_remap(
logical_pages: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build a fixed-shape page remap without dynamic CUDA output ops.
The compact remap path uses boolean compaction + unique/searchsorted. Those
@@ -1618,9 +1763,7 @@ def build_shared_paged_buffer_slot_remap(
physical_page_capacity=page_buffer.shape[0],
)
slot_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * (
layout.cp_size
) + 1
logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * (layout.cp_size) + 1
page_inverse = build_slot_page_inverse_optimized(
slot_logical_pages,
logical_page_capacity=logical_page_capacity,
@@ -1803,9 +1946,7 @@ def build_current_loc_remap(
safe_query_locs = torch.where(
valid_query, query_flat_long, torch.zeros_like(query_flat_long)
)
query_pages = torch.div(
safe_query_locs, page_size, rounding_mode="floor"
)
query_pages = torch.div(safe_query_locs, page_size, rounding_mode="floor")
query_offsets = torch.remainder(safe_query_locs, page_size)
query_pages_in_range = query_pages < logical_page_capacity
safe_query_pages = torch.clamp(query_pages, max=logical_page_capacity - 1)
@@ -1822,7 +1963,9 @@ def build_current_loc_remap(
row_values.to(compact_row_ids.dtype),
torch.full_like(compact_row_ids.reshape(-1), -1),
)
return matched.reshape(query_locs.shape), compact_flat.reshape(query_locs.shape)
return matched.reshape(query_locs.shape), compact_flat.reshape(
query_locs.shape
)
finally:
if cp_shared_kv_sort_nvtx_enabled():
torch.cuda.nvtx.range_pop()
@@ -1975,7 +2118,9 @@ def materialize_local_token_kv_pages(
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
torch.long
)
dense_page_ids = torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
dense_page_ids = (
torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
)
page_offsets = torch.arange(page_size, device=kv_cache.device, dtype=torch.long)
src_tokens = (owned_physical_pages[:, None] * page_size + page_offsets).reshape(-1)
dst_tokens = (dense_page_ids[:, None] * page_size + page_offsets).reshape(-1)
@@ -2236,9 +2381,9 @@ def token_page_copy_debug_checksum(
if not torch.any(owned_mask):
return "owned_pages=0"
owned_logical_pages = unique_logical_pages[owned_mask].to(torch.int64)
owned_physical_pages = layout.logical_pages_to_physical(
owned_logical_pages
).to(torch.long)
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
torch.long
)
dense_page_ids = (
torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
)
@@ -2270,7 +2415,9 @@ def materialize_local_paged_buffer(
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
torch.long
)
dense_page_ids = torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
dense_page_ids = (
torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
)
dense_page_buffer[dense_page_ids] = page_buffer[owned_physical_pages]
return dense_page_buffer
@@ -2369,9 +2516,9 @@ def paged_copy_debug_checksum(
if not torch.any(owned_mask):
return "owned_pages=0"
owned_logical_pages = unique_logical_pages[owned_mask].to(torch.int64)
owned_physical_pages = layout.logical_pages_to_physical(
owned_logical_pages
).to(torch.long)
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
torch.long
)
dense_page_ids = (
torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
)

View File

@@ -23,13 +23,13 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_mla_prefetch_should_log_layer,
cp_shared_kv_is_draft_input,
cp_shared_kv_should_prefetch_next_layer,
current_extend_kv_rows_for_reuse,
current_loc_remap_fast_path_args,
filter_owned_logical_locs,
get_or_build_shared_token_kv_slot_remap,
is_current_only_extend_batch,
materialize_prefix_and_reuse_current_kv_page_slots,
materialize_shared_token_kv_buffer,
should_reuse_current_extend_kv,
tensor_debug_checksum,
tensor_debug_summary,
)
@@ -1747,13 +1747,12 @@ class NativeSparseAttnBackend(
mla_prefetcher = getattr(
forward_batch, "cp_shared_kv_mla_prefetcher", None
)
can_reuse_current_kv = (
should_reuse_current_extend_kv(forward_batch)
and k is not None
and k_rope is not None
and k.shape[0] == forward_batch.out_cache_loc.numel()
and k_rope.shape[0] == forward_batch.out_cache_loc.numel()
current_kv_rows_for_reuse = current_extend_kv_rows_for_reuse(
forward_batch,
k,
k_rope,
)
can_reuse_current_kv = current_kv_rows_for_reuse is not None
if cp_shared_kv_mla_prefetch_log_enabled():
if cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id):
prefix_lens_cpu = getattr(
@@ -1788,8 +1787,15 @@ class NativeSparseAttnBackend(
else None,
)
if can_reuse_current_kv:
current_kv_cache = _cat([k, k_rope], dim=-1)
current_locs_for_reuse = forward_batch.out_cache_loc
assert k is not None and k_rope is not None
assert current_kv_rows_for_reuse is not None
valid_current_rows = int(current_kv_rows_for_reuse)
current_kv_cache = _cat(
[k[:valid_current_rows], k_rope[:valid_current_rows]], dim=-1
)
current_locs_for_reuse = forward_batch.out_cache_loc[
:valid_current_rows
]
logical_page_table_1 = page_table_1
current_remap_page_size, current_remap_logical_page_capacity = (
current_loc_remap_fast_path_args(forward_batch)

View File

@@ -306,7 +306,7 @@ def evict_from_tree_cache(
# Standard allocator
available = allocator.available_size()
if available < num_tokens:
logger.info(
logger.debug(
"[MemCache-evict] evict_from_tree_cache: available=%d < num_tokens=%d deficit=%d, triggering eviction",
available,
num_tokens,
@@ -343,7 +343,7 @@ def _evict_for_compute_owner_lanes(
if isinstance(evictable_size, tuple):
evictable_size = evictable_size[0]
if evictable_size <= 0:
logger.info(
logger.debug(
"[MemCache-evict] _evict_for_compute_owner_lanes: evictable_size=%d <= 0, giving up",
evictable_size,
)
@@ -355,7 +355,7 @@ def _evict_for_compute_owner_lanes(
# load-back pressure.
evict_tokens = max(allocator.page_size, deficit_pages * allocator.page_size)
before_available = allocator.available_size()
logger.info(
logger.debug(
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d: deficit_pages=%d evict_tokens=%d before_available=%d evictable_size=%d",
attempt,
deficit_pages,
@@ -371,7 +371,7 @@ def _evict_for_compute_owner_lanes(
)
after_available = allocator.available_size()
evicted_tokens = getattr(evict_result, "num_tokens_evicted", 0)
logger.info(
logger.debug(
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d result: evicted=%d after_available=%d",
attempt,
evicted_tokens,

View File

@@ -1220,7 +1220,7 @@ class HiRadixCache(RadixCache):
num_evicted += self._evict_backuped(victim)
refreshed = self._refresh_cp_load_back_plan(plan)
logger.info(
logger.debug(
"[HiCache-load] owner-lane device eviction before CP load-back: "
"node_id=%d victims=%s num_evicted=%d required_by_owner=%s "
"before_available_by_owner=%s before_deficit_by_owner=%s "
@@ -1254,7 +1254,7 @@ class HiRadixCache(RadixCache):
)
eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan)
if len(eviction_plan.victims) == 0:
logger.info(
logger.debug(
"[HiCache-evict] owner-lane evict found no contributing victims: "
"num_tokens=%d deficits=%s evictable_size=%d available_size=%d",
params.num_tokens,
@@ -1276,7 +1276,7 @@ class HiRadixCache(RadixCache):
self._clear_pin(victim)
if not self._cp_device_leaf_is_load_back_victim(victim):
logger.info(
logger.debug(
"[HiCache-evict] owner-lane evict victim no longer evictable: "
"victim_id=%s deficits=%s",
getattr(victim, "id", None),
@@ -1319,7 +1319,7 @@ class HiRadixCache(RadixCache):
if self._node_backuped(victim):
num_evicted += self._evict_backuped(victim)
logger.info(
logger.debug(
"[HiCache-evict] owner-lane evict END: requested_tokens=%d "
"deficits=%s victims=%s planned_freed_by_owner=%s "
"remaining_deficit_by_owner=%s num_evicted=%d "
@@ -1925,7 +1925,7 @@ class HiRadixCache(RadixCache):
victim.host_value = None
self._remove_host_leaf(victim)
logger.info(
logger.debug(
"[HiCache-evict] deterministic CP host eviction before write: "
"node_id=%d phase=%s victims=%s local_freed=%d planned_freed=%s",
node_id,
@@ -1978,7 +1978,7 @@ class HiRadixCache(RadixCache):
)
)
logger.info(
logger.debug(
"[HiCache-write] write_backup CP retry after deterministic host eviction: "
"node_id=%d deficit_by_owner=%s",
node_id,
@@ -1991,7 +1991,7 @@ class HiRadixCache(RadixCache):
if not isinstance(result, HiCacheWriteFailure):
return result
logger.info(
logger.warning(
"[HiCache-write] write_backup CP FAILED after deterministic retry: "
"node_id=%d len=%d needed_slots=%d",
node_id,
@@ -2800,7 +2800,7 @@ class HiRadixCache(RadixCache):
]
heapq.heapify(eviction_heap)
logger.info(
logger.debug(
"[HiCache-evict] evict START: num_tokens=%d heap_size=%d evictable_size=%d available_size=%d",
num_tokens,
len(eviction_heap),
@@ -2864,7 +2864,7 @@ class HiRadixCache(RadixCache):
self._evict_backuped(node)
self.update_eviction_metrics(num_evicted, start_time)
logger.info(
logger.debug(
"[HiCache-evict] evict END: num_tokens=%d num_evicted=%d num_locked_skipped=%d evictable_size_after=%d available_size_after=%d",
num_tokens,
num_evicted,
@@ -2880,7 +2880,7 @@ class HiRadixCache(RadixCache):
freed_len = self.cache_controller.evict_device(node.value)
assert freed_len > 0
self.evictable_size_ -= device_resident_len
logger.info(
logger.debug(
"[HiCache-evict] _evict_backuped: node_id=%d num_evicted=%d physical_tokens=%d lock_ref=%d backed=%s",
node.id,
freed_len,
@@ -2899,7 +2899,7 @@ class HiRadixCache(RadixCache):
def _evict_regular(self, node: TreeNode):
# evict a node not initiated write to host -- emit BlockRemoved
num_evicted = self._node_device_resident_len(node)
logger.info(
logger.debug(
"[HiCache-evict] _evict_regular: node_id=%d num_evicted=%d",
node.id,
num_evicted,
@@ -2943,7 +2943,7 @@ class HiRadixCache(RadixCache):
return 0
leaves = list(self.evictable_host_leaves)
logger.info(
logger.debug(
"[HiCache-evict] _evict_host_for_physical_slots: required_slots=%d sync=%s leaves=%d",
required_host_slots,
synchronize_across_ranks,