Keep draft cache-hit KV on current-suffix compose

Cache-hit EAGLE/NextN draft was falling back to full materialization while the target path used page-aligned prefix materialize plus fresh current-suffix splice. That creates a target/draft KV source asymmetry exactly on the high-cache-hit path where decode accept length collapsed. The draft model still does not get next-layer async prefetch; only the same-layer current suffix compose contract is made role-agnostic.

Constraint: EAGLE/NextN has no real next decoder layer, so draft async prefetch remains disabled.

Rejected: Restore draft full-materialize fallback | recreates the observed target/draft cache-hit asymmetry and hides stale-current-suffix bugs.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce draft cache-hit current-reuse fallback without proving the draft persistent pool has fully fresh suffix rows before attention.

Tested: Local py_compile for cp_shared_kv_runtime.py and test_cp_shared_kv_runtime.py.

Tested: Remote g0034 container py_compile for changed runtime/test files.

Tested: Remote g0034 pytest test_cp_shared_kv_runtime.py: 73 passed, 5 warnings, 2 subtests passed.

Tested: Remote g0034 related suite test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py: 128 passed, 5 warnings, 2 subtests passed.

Not-tested: ETE cache-hit accept-length validation after restarting prefill/decode with this synced code.
This commit is contained in:
laoyao0822
2026-05-30 03:06:17 +08:00
parent 5f343c65ca
commit b7364d23f9
3 changed files with 216 additions and 33 deletions
@@ -569,7 +569,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_cache_hit_suffix(self):
def test_should_reuse_current_extend_kv_allows_draft_cache_hit_suffix(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
@@ -593,15 +593,8 @@ 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.assertTrue(
any(
"draft_partial_current_reuse" in message
and "[CP_SHARED_KV_FALLBACK][current_reuse]" in message
for message in logs.output
)
)
self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch))
self.assertEqual(runtime._CURRENT_REUSE_FALLBACK_LOG_COUNTS, {})
forward_batch.spec_info = TargetSpecInfo()
forward_batch.cp_shared_kv_mla_prefetcher = object()