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

Keep draft cache-hit KV on current-suffix compose

Browse Source 
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:04:34 +08:00
parent 5f343c65ca
commit b7364d23f9
3 changed files with 216 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -978,14 +978,13 @@ def can_reuse_current_extend_kv(forward_batch) -> bool:
def should_reuse_current_extend_kv(forward_batch) -> bool:
"""Return whether MLA should splice current extend KV into materialized KV.
Partial current reuse appends the freshly computed suffix KV to the
materialized prefix buffer. That path is safe for the target model, but the
EAGLE/NextN draft layer has a different lifetime contract: it mirrors target
KV state and immediately enters DeepEP MoE after attention. A cache-hit
draft suffix observed in production can leave all ranks stuck in that MoE
collective after the partial-reuse attention path returns. Keep draft
cache-hit suffixes on the older full-materialize path until draft gets an
explicit same-layer reuse contract.
The contract is model-role agnostic: target and prefill-time EAGLE/NextN
draft both write only the CP-owned suffix rows into the persistent pool, so
cache-hit attention should materialize the page-aligned prefix and splice the
freshly computed current suffix explicitly. Draft async prefetch remains
disabled elsewhere because NextN has no next layer to prefetch, but disabling
that optimization must not force draft cache-hit suffixes back to full
materialization from the pool.
"""
if not cp_shared_kv_current_reuse_enabled():
@@ -993,21 +992,7 @@ def should_reuse_current_extend_kv(forward_batch) -> bool:
current_only = is_current_only_extend_batch(forward_batch)
partial_current = can_reuse_current_extend_kv(forward_batch)
if not (current_only or partial_current):
return False
if cp_shared_kv_is_draft_input(forward_batch) and not current_only:
_log_current_reuse_fallback(
"draft_partial_current_reuse_disabled",
"CP shared KV current-reuse fallback (draft_partial_current_reuse): "
"cache-hit EAGLE/NextN draft uses full materialize instead of "
"partial current reuse. prefix_lens=%s extend_lens=%s",
getattr(forward_batch, "extend_prefix_lens_cpu", None),
getattr(forward_batch, "extend_seq_lens_cpu", None),
)
return False
return True
return current_only or partial_current
def current_loc_remap_fast_path_args(