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

Keep CP HiCache reuse page-safe without eviction log churn

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+55
View File
@@ -1,3 +1,5 @@
import inspect
import re
import sys
import types
import unittest
@@ -97,6 +99,8 @@ from sglang.srt.managers.cache_controller import (
HiCacheWriteFailure,
HiCacheWriteReservation,
)
import sglang.srt.mem_cache.common as mem_cache_common
import sglang.srt.mem_cache.hiradix_cache as hiradix_cache
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
InsertParams,
@@ -117,6 +121,57 @@ from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=2, suite="stage-a-test-cpu")
class TestHiCacheEvictLoggingLevels(CustomTestCase):
def assert_debug_marker(self, source: str, marker: str):
self.assertIn(marker, source)
self.assertRegex(
source,
r"logger\.debug\(\s*\n\s*\"" + re.escape(marker),
msg=f"{marker} should be debug-only on the success/no-op hot path",
)
self.assertNotRegex(
source,
r"logger\.info\(\s*\n\s*\"" + re.escape(marker),
msg=f"{marker} must not stay at INFO on the success/no-op hot path",
)
def test_evict_hot_path_success_logs_are_debug_only(self):
common_source = inspect.getsource(mem_cache_common)
hiradix_source = inspect.getsource(hiradix_cache)
for marker in (
"[MemCache-evict] evict_from_tree_cache:",
"[MemCache-evict] _evict_for_compute_owner_lanes: evictable_size",
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d:",
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d result:",
):
self.assert_debug_marker(common_source, marker)
for marker in (
"[HiCache-load] owner-lane device eviction before CP load-back: ",
"[HiCache-evict] owner-lane evict found no contributing victims: ",
"[HiCache-evict] owner-lane evict victim no longer evictable: ",
"[HiCache-evict] owner-lane evict END: requested_tokens=%d ",
"[HiCache-evict] deterministic CP host eviction before write: ",
"[HiCache-write] write_backup CP retry after deterministic host eviction: ",
"[HiCache-evict] evict START:",
"[HiCache-evict] evict END:",
"[HiCache-evict] _evict_backuped:",
"[HiCache-evict] _evict_regular:",
"[HiCache-evict] _evict_host_for_physical_slots:",
):
self.assert_debug_marker(hiradix_source, marker)
self.assertRegex(
hiradix_source,
r"logger\.warning\(\s*\n\s*\"\[HiCache-write\] write_backup CP FAILED after deterministic retry:",
)
self.assertNotRegex(
hiradix_source,
r"logger\.info\(\s*\n\s*\"\[HiCache-write\] write_backup CP FAILED after deterministic retry:",
)
class TestCpHiCacheImports(CustomTestCase):
def test_cp_hicache_public_imports_without_sgl_kernel(self):
import subprocess
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+173 -44
View File
@@ -54,9 +54,7 @@ for _name in ("flash_attn_varlen_func", "flash_attn_with_kvcache"):
if not hasattr(flash_attn_stub, _name):
setattr(flash_attn_stub, _name, lambda *args, **kwargs: None)
sgl_kernel_stub = sys.modules.setdefault(
"sgl_kernel", types.ModuleType("sgl_kernel")
)
sgl_kernel_stub = sys.modules.setdefault("sgl_kernel", types.ModuleType("sgl_kernel"))
if not hasattr(sgl_kernel_stub, "__path__"):
sgl_kernel_stub.__path__ = []
if not hasattr(sgl_kernel_stub, "flash_attn"):
@@ -361,9 +359,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with patch.dict(
sys.modules,
{
"sglang.srt.layers.attention.nsa.index_buf_accessor": index_accessor_stub
},
{"sglang.srt.layers.attention.nsa.index_buf_accessor": index_accessor_stub},
):
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
@@ -457,7 +453,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertIs(out, buffer)
dist_all_reduce.assert_called_once()
self.assertIs(dist_all_reduce.call_args.args[0], buffer)
self.assertIs(dist_all_reduce.call_args.kwargs["group"], dummy_group.device_group)
self.assertIs(
dist_all_reduce.call_args.kwargs["group"], dummy_group.device_group
)
def test_build_dense_page_remap_preserves_sentinels(self):
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
@@ -492,9 +490,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
current_locs = torch.tensor([100, 64, 256, 128], dtype=torch.int64)
query_locs = torch.tensor(
[[128, -1, 64], [512, 100, 256]], dtype=torch.int32
)
query_locs = torch.tensor([[128, -1, 64], [512, 100, 256]], dtype=torch.int32)
is_current, compact_rows = build_current_loc_remap(query_locs, current_locs)
@@ -569,7 +565,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
forward_batch.out_cache_loc = torch.arange(64, dtype=torch.int64)
self.assertFalse(can_reuse_current_extend_kv(forward_batch))
def test_should_reuse_current_extend_kv_disables_draft_partial_cache_hit_suffix(
def test_should_reuse_current_extend_kv_enables_draft_partial_cache_hit_suffix(
self,
):
from sglang.srt.environ import envs
@@ -595,13 +591,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True):
with self.assertLogs(runtime.logger.name, level="WARNING") as logs:
self.assertFalse(runtime.should_reuse_current_extend_kv(forward_batch))
self.assertIn(
"[CP_SHARED_KV_FALLBACK][current_reuse]",
logs.output[0],
)
self.assertIn("draft_partial_current_reuse_disabled", logs.output[0])
self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch))
forward_batch.spec_info = TargetSpecInfo()
forward_batch.cp_shared_kv_mla_prefetcher = object()
@@ -619,6 +609,73 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
forward_batch.cp_shared_kv_mla_prefetcher = None
self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch))
def test_current_extend_kv_rows_for_reuse_accepts_padded_out_cache_loc(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
class DraftSpecInfo:
def is_draft_input(self):
return True
forward_batch = SimpleNamespace(
forward_mode=_FakeExtendForwardMode(),
batch_size=1,
extend_prefix_lens_cpu=[40384],
extend_seq_lens_cpu=[65],
seq_lens_cpu=torch.tensor([40384 + 65], dtype=torch.int32),
out_cache_loc=torch.arange(128, dtype=torch.int64),
spec_info=DraftSpecInfo(),
)
k = torch.empty((65, 2, 4), dtype=torch.float32)
k_rope = torch.empty((65, 2, 1), dtype=torch.float32)
with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True):
self.assertEqual(
runtime.current_extend_kv_rows_for_reuse(forward_batch, k, k_rope),
65,
)
self.assertIsNone(
runtime.current_extend_kv_rows_for_reuse(
forward_batch,
k[:64],
k_rope,
)
)
def test_mla_current_reuse_gate_accepts_padded_out_cache_loc(self):
from pathlib import Path
source = (
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa_backend.py"
).read_text()
start = source.index(" current_kv_rows_for_reuse =")
end = source.index(
" if cp_shared_kv_mla_prefetch_log_enabled()", start
)
gate_source = source[start:end]
self.assertIn("current_extend_kv_rows_for_reuse", gate_source)
self.assertNotIn(
"k.shape[0] == forward_batch.out_cache_loc.numel()",
gate_source,
)
self.assertNotIn(
"k_rope.shape[0] == forward_batch.out_cache_loc.numel()",
gate_source,
)
body_start = source.index(" if can_reuse_current_kv:", end)
body_end = source.index(
" logical_page_table_1 = page_table_1", body_start
)
body_source = "".join(source[body_start:body_end].split())
self.assertIn("valid_current_rows=int(current_kv_rows_for_reuse)", body_source)
self.assertIn("k[:valid_current_rows]", body_source)
self.assertIn("k_rope[:valid_current_rows]", body_source)
self.assertIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source)
def test_runtime_fallback_helpers_use_standard_warning_marker(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
@@ -738,6 +795,53 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertEqual(current_mask.tolist(), [[False, True, True, False, False]])
self.assertEqual(mixed_locs.tolist(), [[4, 12, 13, -1, -1]])
def test_tai_current_slot_fill_is_skipped_when_sparse_page_self_test_fails(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
class BadKernels:
@staticmethod
def fill_current_token_kv_page_slots_and_remap_locs(*args, **kwargs):
raise AssertionError("stale TAI fill kernel should not be called")
with envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.override(True), patch.object(
runtime,
"_tai_current_slot_fill_supports_sparse_pages",
return_value=False,
), patch.object(
runtime,
"_load_tai_materialize_kernels",
return_value=BadKernels,
):
result = runtime._try_tai_fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=torch.zeros((16, 1), dtype=torch.float32),
materialized_dense_locs=torch.tensor([[4, 5, 8, 9]], dtype=torch.int64),
current_kv_cache=torch.ones((2, 1), dtype=torch.float32),
logical_locs=torch.tensor([[20, 21, 40, 41]], dtype=torch.int64),
current_locs=torch.tensor([20, 21], dtype=torch.int64),
page_inverse=torch.tensor(
[0, -1, -1, -1, -1, 1, -1, -1, -1, -1, 2],
dtype=torch.long,
),
page_size=4,
mask_non_current_in_current_pages=True,
)
self.assertIsNone(result)
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
def test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel(
self,
):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
runtime._tai_current_slot_fill_sparse_pages_self_test.cache_clear()
self.assertTrue(
runtime._tai_current_slot_fill_supports_sparse_pages(
torch.device("cuda", torch.cuda.current_device())
)
)
def test_materialize_prefix_and_reuse_current_kv_page_slots_without_prefetcher(
self,
):
@@ -913,9 +1017,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with patch.object(
prefetch.torch.cuda, "current_stream", return_value=current_stream
), patch.object(
prefetcher, "launch_pending_reduce"
) as launch_pending_reduce:
), patch.object(prefetcher, "launch_pending_reduce") as launch_pending_reduce:
prefetcher.wait_attention_window()
launch_pending_reduce.assert_not_called()
@@ -1402,8 +1504,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(
any(
"token slot remap cache not reused (missing_cached_value)"
in message
"token slot remap cache not reused (missing_cached_value)" in message
for message in logs.output
)
)
@@ -1549,13 +1650,17 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
self.assertIn("prefix_not_page_aligned", logger.call_args.args[1])
def test_mla_prefetch_min_prefix_pages_uses_cached_token_default_and_can_override(self):
def test_mla_prefetch_min_prefix_pages_uses_cached_token_default_and_can_override(
self,
):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.clear()
default_tokens = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_TOKENS.get()
self.assertEqual(runtime._MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS, default_tokens)
self.assertEqual(
runtime._MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS, default_tokens
)
expected_pages = (default_tokens + 63) // 64
self.assertEqual(
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(8, page_size=64),
@@ -1844,9 +1949,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with patch.object(
runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=False
), patch.object(
runtime, "_load_tai_fused_mla_store_kernel"
) as load_kernel:
), patch.object(runtime, "_load_tai_fused_mla_store_kernel") as load_kernel:
used = runtime.try_tai_fused_mla_store(
token_to_kv_pool=FakePool(),
layer=SimpleNamespace(layer_id=0),
@@ -1962,7 +2065,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.inverse_calls = []
self.remap_calls = []
def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity):
def build_slot_page_inverse(
self, slot_logical_pages, logical_page_capacity
):
self.inverse_calls.append((slot_logical_pages, logical_page_capacity))
return torch.tensor([0, 1, 2, -1], dtype=torch.long)
@@ -2056,7 +2161,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with patch.object(
runtime, "cp_shared_kv_debug_enabled", return_value=True
), patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
), patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
):
with self.assertRaisesRegex(
RuntimeError,
"CP shared KV materialize got logical token locs outside the physical pool",
@@ -2068,7 +2175,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
page_size=4,
)
def test_materialize_token_kv_skips_physical_pool_validation_when_debug_disabled(self):
def test_materialize_token_kv_skips_physical_pool_validation_when_debug_disabled(
self,
):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -2081,7 +2190,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with patch.object(
runtime, "cp_shared_kv_debug_enabled", return_value=False
), patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
), patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
):
_, dense_locs = runtime.materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=logical_locs,
@@ -2100,7 +2211,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
logical_locs = torch.tensor([4, -1, 8], dtype=torch.int64)
with patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=True):
with patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
with patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
):
_, dense_locs = runtime.materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=logical_locs,
@@ -2172,7 +2285,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
logical_locs = torch.tensor([8, 20, -1], dtype=torch.int64)
remap_source_locs = torch.tensor([4, 8, 12, 16, 20], dtype=torch.int64)
with patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
with patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
):
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=logical_locs,
@@ -2223,7 +2338,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
kv_cache = torch.arange(0, 32, dtype=torch.float32).view(32, 1, 1)
remap_source_locs = torch.tensor([4, 8, 12, 16, 20], dtype=torch.int64)
with patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
with patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
):
dense_kv_a, dense_locs_a = runtime.materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=torch.tensor([4, 20], dtype=torch.int64),
@@ -2254,7 +2371,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with patch.object(
runtime, "cp_shared_kv_debug_enabled", return_value=True
), patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
), patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
):
with self.assertRaisesRegex(
RuntimeError,
"CP shared KV materialize got logical pages outside the physical page buffer",
@@ -2401,7 +2520,9 @@ class TestCpSharedKVLazyDebugLogging(unittest.TestCase):
self.assertEqual(k_to_write.shape[0], 2)
self.assertEqual(k_rope_to_write.shape[0], 2)
def test_index_write_filter_does_not_build_debug_summaries_when_debug_disabled(self):
def test_index_write_filter_does_not_build_debug_summaries_when_debug_disabled(
self,
):
from sglang.srt.layers.attention.nsa import nsa_indexer
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -2422,10 +2543,12 @@ class TestCpSharedKVLazyDebugLogging(unittest.TestCase):
"sglang.srt.environ.envs.SGLANG_DEBUG_CP_SHARED_KV.get",
return_value=False,
):
physical_locs, key_to_write = nsa_indexer.Indexer._filter_shared_index_write(
None,
forward_batch,
key,
physical_locs, key_to_write = (
nsa_indexer.Indexer._filter_shared_index_write(
None,
forward_batch,
key,
)
)
self.assertEqual(physical_locs.tolist(), [4, 8])
@@ -2493,8 +2616,12 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
self.remap_calls = []
self.token_calls = []
def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity):
self.page_inverse_calls.append((slot_logical_pages, logical_page_capacity))
def build_slot_page_inverse(
self, slot_logical_pages, logical_page_capacity
):
self.page_inverse_calls.append(
(slot_logical_pages, logical_page_capacity)
)
return torch.tensor([0, 1, 2, -1, 3], dtype=torch.long)
def remap_logical_locs_to_slot_dense_locs(
@@ -2567,7 +2694,9 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
class FakeTaiKernels:
def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity):
def build_slot_page_inverse(
self, slot_logical_pages, logical_page_capacity
):
return torch.tensor([0, 1, 2, -1, 3], dtype=torch.long)
def remap_logical_locs_to_slot_dense_locs(