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

Keep CP shared-KV cache hits page-aligned without short fallback

Browse Source 
CP shared KV and HiCache need a stable contract where physical cache coverage is page-aligned, while scheduler/radix-visible hit length remains the valid token length. This records the contract, adds page-aligned extent metadata, keeps owner assignment on actual tail pages instead of short-prefix fallback, and updates partial current reuse tests around tail-page masking.

Constraint: CP owner lanes operate on page units while scheduler and radix hit accounting must remain token-valid.

Rejected: Pad short suffixes to cp_size or 2*cp_size pages | wastes KV capacity and can turn a small tail into a much larger physical span.

Rejected: Silent direct-write or prefetch fallback | production fallback must be warning-visible for diagnosis.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce replicated short-radix fallback without checking docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md.

Tested: local py_compile for touched runtime, utility, owner, and unit-test files.

Tested: remote g0034 container three-file suite: 122 passed, 5 warnings.

Not-tested: full local pytest, blocked by missing runtime dependencies such as orjson.

Not-tested: CUDA E2E runtime for this commit.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-29 05:04:50 +08:00
co-authored by OmX
parent 1e1807f2bc
commit c551623ca8
11 changed files with 1369 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
+19 -6
View File
@@ -150,7 +150,7 @@ class TestCPSharedPagedAllocator(CustomTestCase):
[0, 0, 1, 1, 2, 2, 3, 3, 3, 3, 2, 2, 1, 1, 0, 0],
)
def test_compute_owner_page_assignment_falls_back_for_short_extend(self):
def test_compute_owner_page_assignment_keeps_page_aligned_short_extend(self):
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_in_seq_page_compute_owners,
get_in_seq_page_compute_owner_unavailable_reason,
@@ -163,17 +163,30 @@ class TestCPSharedPagedAllocator(CustomTestCase):
cp_size=4,
)
self.assertIsNone(owners)
self.assertEqual(
self.assertEqual(owners, [0, 1, 2, 3, 3, 2, 1])
self.assertIsNone(
get_in_seq_page_compute_owner_unavailable_reason(
extend_len=64 * 7,
extend_prefix_len=0,
page_size=64,
cp_size=4,
),
"too_short_for_page_aligned",
)
)
def test_compute_owner_page_assignment_uses_tail_page_not_cp_padding(self):
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_in_seq_page_compute_owners,
)
owners = build_in_seq_page_compute_owners(
extend_len=100,
extend_prefix_len=0,
page_size=64,
cp_size=8,
)
self.assertEqual(owners, [0, 1])
def test_compute_owner_page_assignment_allows_radix_hit_suffix_with_one_page_per_rank(
self,
):
@@ -199,7 +212,7 @@ class TestCPSharedPagedAllocator(CustomTestCase):
)
)
def test_compute_owner_page_assignment_allows_short_radix_hit_suffix_with_replicated_compute(
def test_compute_owner_page_assignment_keeps_short_radix_hit_suffix_page_aligned(
self,
):
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+169 -26
View File
@@ -48,7 +48,7 @@ class _FakeExtendForwardMode:
class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
def test_mla_prefetch_materializes_and_reduces_on_prefetch_stream(
def test_mla_prefetch_materializes_on_current_stream_and_reduces_on_prefetch_stream(
self,
):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
@@ -137,7 +137,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")])
self.assertEqual(
calls,
[("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")],
[("materialize", "current"), ("reduce", "prefetch", "prefetch")],
)
self.assertEqual(prefetch_stream.waited, ["current"])
@@ -145,11 +145,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertEqual(
calls,
[("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")],
[("materialize", "current"), ("reduce", "prefetch", "prefetch")],
)
self.assertEqual(prefetch_stream.waited, ["current"])
def test_index_prefetch_materializes_and_reduces_on_prefetch_stream(
def test_index_prefetch_materializes_on_current_stream_and_reduces_on_prefetch_stream(
self,
):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
@@ -239,7 +239,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")])
self.assertEqual(
calls,
[("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")],
[("materialize", "current"), ("reduce", "prefetch", "prefetch")],
)
self.assertEqual(prefetch_stream.waited, ["current"])
@@ -247,7 +247,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertEqual(
calls,
[("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")],
[("materialize", "current"), ("reduce", "prefetch", "prefetch")],
)
self.assertEqual(prefetch_stream.waited, ["current"])
@@ -539,6 +539,41 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
self.assertEqual(mixed_locs.tolist(), [[4, 8, -1], [9, 7, -1]])
def test_fill_current_kv_page_slots_masks_suffix_page_slack(self):
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
fill_current_kv_page_slots_and_remap_locs,
)
page_size = 4
dense_kv = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1)
original_dense_kv = dense_kv.clone()
current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1)
logical_locs = torch.tensor([[4, 20, 21, 22, 23]], dtype=torch.int32)
# The prefetched dense buffer has physical rows for the whole suffix page,
# but only locs 20 and 21 are real current tokens. Rows for 22 and 23 are
# page slack and must not stay visible.
materialized_locs = torch.tensor([[4, 12, 13, 14, 15]], dtype=torch.int32)
current_locs = torch.tensor([20, 21], dtype=torch.int64)
page_inverse = torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64)
mixed_kv, mixed_locs, current_mask = fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=dense_kv,
materialized_dense_locs=materialized_locs,
current_kv_cache=current_kv,
logical_locs=logical_locs,
current_locs=current_locs,
page_inverse=page_inverse,
page_size=page_size,
mask_non_current_in_current_pages=True,
)
expected_kv = original_dense_kv.clone()
expected_kv[12:14] = current_kv
self.assertEqual(int(mixed_kv.shape[0]), 16)
self.assertTrue(torch.equal(mixed_kv, expected_kv))
self.assertEqual(current_mask.tolist(), [[False, True, True, False, False]])
self.assertEqual(mixed_locs.tolist(), [[4, 12, 13, -1, -1]])
def test_mla_prefetch_consume_prefix_with_current_skips_suffix_materialize(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -553,6 +588,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
current_stream = FakeCurrentStream()
fake_event = object()
dense_kv = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1)
original_dense_kv = dense_kv.clone()
current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1)
page_inverse = torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64)
prefetcher = prefetch.CpSharedKVMlaPrefetcher(
@@ -583,7 +619,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
mixed_kv, mixed_locs = prefetcher.consume_prefix_with_current(
layer_id=1,
kv_cache=torch.zeros((64, 1, 1), dtype=torch.float32),
logical_locs=torch.tensor([[4, 20], [21, 7]], dtype=torch.int32),
logical_locs=torch.tensor(
[[4, 20], [21, 7], [22, 23]], dtype=torch.int32
),
current_kv_cache=current_kv,
current_locs=torch.tensor([20, 21], dtype=torch.int64),
)
@@ -591,11 +629,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertEqual(current_stream.events, [fake_event])
self.assertEqual(prefetcher.handles, {})
self.assertIsNone(prefetcher.pending_attention_handle)
self.assertTrue(torch.equal(mixed_kv[:16], dense_kv))
self.assertTrue(torch.equal(mixed_kv[16:], current_kv))
self.assertEqual(mixed_locs.tolist(), [[4, 16], [17, 7]])
expected_kv = original_dense_kv.clone()
expected_kv[12:14] = current_kv
self.assertTrue(torch.equal(mixed_kv, expected_kv))
self.assertEqual(mixed_locs.tolist(), [[4, 12], [13, 7], [-1, -1]])
def test_mla_prefetch_attention_window_waits_on_pending_event(self):
def test_mla_prefetch_attention_window_defers_pending_event_wait(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -631,11 +670,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
):
prefetcher.wait_attention_window()
self.assertEqual(current_stream.events, [fake_event])
self.assertIsNone(prefetcher.pending_attention_handle)
self.assertEqual(current_stream.events, [])
self.assertIs(prefetcher.pending_attention_handle, handle)
self.assertIs(prefetcher.handles[1], handle)
def test_mla_prefetch_attention_window_launches_pending_reduce_before_wait(self):
def test_mla_prefetch_attention_window_does_not_launch_pending_reduce(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -647,7 +686,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.events.append(event)
current_stream = FakeCurrentStream()
fake_event = object()
prefetcher = prefetch.CpSharedKVMlaPrefetcher(
layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
page_size=4,
@@ -666,21 +704,19 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
prefetcher.handles[1] = handle
prefetcher.pending_attention_handle = handle
def finish_reduce():
handle.event = fake_event
with patch.object(
prefetch.torch.cuda, "current_stream", return_value=current_stream
), patch.object(
prefetcher, "launch_pending_reduce", side_effect=finish_reduce
prefetcher, "launch_pending_reduce"
) as launch_pending_reduce:
prefetcher.wait_attention_window()
launch_pending_reduce.assert_called_once_with()
self.assertEqual(current_stream.events, [fake_event])
self.assertIsNone(prefetcher.pending_attention_handle)
launch_pending_reduce.assert_not_called()
self.assertEqual(current_stream.events, [])
self.assertIs(prefetcher.pending_attention_handle, handle)
self.assertIsNone(handle.event)
def test_index_prefetch_attention_window_waits_on_pending_event(self):
def test_index_prefetch_attention_window_defers_pending_event_wait(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -715,10 +751,69 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
):
prefetcher.wait_attention_window()
self.assertEqual(current_stream.events, [fake_event])
self.assertIsNone(prefetcher.pending_attention_handle)
self.assertEqual(current_stream.events, [])
self.assertIs(prefetcher.pending_attention_handle, handle)
self.assertIs(prefetcher.handles[1], handle)
def test_index_prefetch_tail_page_keeps_valid_length_unpadded(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
class FakeCurrentStream:
def __init__(self):
self.events = []
def wait_event(self, event):
self.events.append(event)
page_size = 4
valid_tokens = 10
padded_pages = (valid_tokens + page_size - 1) // page_size
cp_size = 8
page_buffer = torch.arange(0, 8 * 3, dtype=torch.uint8).view(8, 3)
dense_page_buffer = torch.zeros((padded_pages + 1, 3), dtype=torch.uint8)
dense_page_buffer[1] = page_buffer[1]
fake_event = object()
prefetcher = prefetch.CpSharedKVIndexPrefetcher(
layout=CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=0),
prefix_pages=1,
slot_logical_pages=torch.tensor([1, 2, 3], dtype=torch.int64),
page_inverse=torch.tensor([-1, 1, 2, 3], dtype=torch.int64),
dense_num_pages=padded_pages,
stream=object(),
)
handle = prefetch.CpSharedKVIndexPrefetchHandle(
layer_id=1,
dense_page_buffer=dense_page_buffer,
prefix_rows=slice(1, 2),
event=fake_event,
)
prefetcher.handles[1] = handle
prefetcher.pending_attention_handle = handle
with patch.object(
prefetch.torch.cuda, "current_stream", return_value=FakeCurrentStream()
), patch.object(
prefetch, "_all_reduce_materialized_buffer_range", _identity_all_reduce
):
dense_pages_buffer, dense_pages = prefetcher.consume(
layer_id=1,
page_buffer=page_buffer,
logical_pages=torch.tensor([[1, 2, 3]], dtype=torch.int32),
)
self.assertEqual(valid_tokens, 10)
self.assertEqual(padded_pages, 3)
self.assertEqual(dense_pages.tolist(), [[1, 2, 3]])
self.assertEqual(dense_pages.shape[1], padded_pages)
self.assertLess(dense_pages.shape[1], cp_size)
self.assertEqual(list(dense_pages_buffer.shape), [padded_pages + 1, 3])
self.assertTrue(torch.equal(dense_pages_buffer[1], page_buffer[1]))
# Index kernels receive valid sequence lengths separately; the dense
# table includes the tail page, but no fake 12-token length is exposed.
self.assertLess(valid_tokens, padded_pages * page_size)
self.assertEqual(valid_tokens, 10)
def test_materialize_local_token_kv_pages(self):
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_dense_page_remap,
@@ -1023,6 +1118,54 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True):
self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled())
def test_mla_prefetch_prefix_misalignment_logs_warning_fallback(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
class Mode:
def is_context_parallel_extend(self):
return True
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
hisparse_coordinator=None,
forward_mode=Mode(),
batch_size=1,
token_to_kv_pool=SimpleNamespace(page_size=64),
cp_shared_kv_layout=SimpleNamespace(cp_size=1, cp_rank=0),
extend_prefix_lens_cpu=[65],
)
metadata = SimpleNamespace(
real_page_table=torch.tensor([1, 2], dtype=torch.int64),
page_table_1=torch.tensor([[64, 65]], dtype=torch.int32),
)
with patch.object(
prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True
), patch.object(
prefetch, "cp_shared_kv_debug_enabled", return_value=False
), patch.object(
prefetch.torch.cuda, "is_available", return_value=True
), patch.object(
prefetch, "_is_cuda_stream_capturing", return_value=False
), patch.object(
prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True
), patch.object(
prefetch.logger, "warning"
) as logger:
result = prefetch.CpSharedKVMlaPrefetcher.maybe_create(
forward_batch=forward_batch,
metadata=metadata,
topk_transform_is_paged=True,
)
self.assertIsNone(result)
logger.assert_called_once()
self.assertIn(
"[CP_SHARED_KV_FALLBACK][mla_prefetch]",
logger.call_args.args[0],
)
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):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
@@ -1053,7 +1196,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(-2):
self.assertEqual(
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(4, page_size=64),
16,
max(4, expected_pages),
)
with patch.object(runtime, "_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS", 2048):