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:
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user