Stabilize CP HiCache page-tail ownership under EAGLE reuse

CP shared KV and HiCache now keep page-aligned physical ownership while preserving valid-token radix semantics. Repeated tiny EAGLE exact hits free duplicate tail pages instead of leaking one allocator page, owner-lane load-back uses page-vector admission/eviction, and single-DP idle schedulers avoid entering an unnecessary MLP-sync collective.

The commit also records the current page-aligned cache contract and adds gated decode-side EAGLE accept diagnostics so future accept-length collapses can be tied to draft KV/state transfer evidence instead of more prefill cache speculation.

Constraint: CP HiCache allocator ownership is page-granular while radix matching remains valid-token based.

Constraint: New diagnostics must be gated and must not alter normal EAGLE, transfer, or cache behavior.

Rejected: Padding short requests to cp_size or 2*cp_size pages | wastes KV capacity and still hides valid-tail lifecycle bugs.

Rejected: Adding more unconditional collectives to prove CP consistency | hot-path collectives previously caused severe performance risk.

Confidence: medium

Scope-risk: broad

Directive: Do not reintroduce silent fallback for CP shared KV/HiCache paths; warning-level fallback or fail-fast is intentional.

Tested: git diff --check

Tested: local py_compile for all modified Python files

Tested: remote g0034 container py_compile for modified Python/test files

Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/managers/test_scheduler_dp_attn_mixin.py => 114 passed, 5 warnings, 2 subtests passed

Not-tested: full ETE traffic rerun after this commit

Not-tested: CUDA/TAI kernel benchmark coverage for all production shapes
This commit is contained in:
laoyao0822
2026-05-30 01:20:01 +08:00
parent 21065cdfdf
commit b56a4f2e6b
16 changed files with 964 additions and 17 deletions
@@ -233,7 +233,32 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
split_list, extend_prefix_len=54464, extend_len=256, page_size=64
)
def test_can_cp_split_keeps_cp_for_short_radix_hit_suffix(self):
def test_can_cp_split_skips_cp_when_radix_hit_suffix_has_too_few_pages(self):
class Mode:
def is_context_parallel_extend(self):
return True
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_seq_lens_cpu=[65],
extend_prefix_lens_cpu=[54464],
token_to_kv_pool=SimpleNamespace(page_size=64),
forward_mode=Mode(),
)
with (
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
),
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
return_value=True,
),
):
self.assertFalse(can_cp_split(128, 8, True, forward_batch))
def test_can_cp_split_skips_cp_when_page_units_do_not_cover_all_lanes(self):
class Mode:
def is_context_parallel_extend(self):
return True
@@ -256,7 +281,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
return_value=True,
),
):
self.assertTrue(can_cp_split(256, 8, True, forward_batch))
self.assertFalse(can_cp_split(256, 8, True, forward_batch))
def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self):
class Mode:
@@ -847,7 +872,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
"materialize/current merge when page-slot prefetch compose is unavailable.",
)
self.assertIn(
"[CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch]",
"[CP_SHARED_KV_FAIL_FAST][mla_partial_current_sync]",
source,
)