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

Keep tiny CP cache-hit suffixes off async prefetch

Browse Source 
The repeated-request hang reproduced on g0034 after the cached-prefix
request created both MLA and index async prefetchers with
prefix_lens=[40320] and extend_lens=[65]. The existing one-page default
only blocked sub-page suffixes, so a barely-over-one-page suffix still
entered the next-layer collective path before any later forward progress
was logged.

Raise the default async prefetch extend gate to one page per CP lane while
keeping the env override. This only gates async prefetcher object
creation; target partial-current reuse still uses the synchronous
page-slot compose/current-splice path when no prefetcher exists.

Constraint: cp_size=8,page_size=64 repeated prompt had extend_len=65 and hung immediately after has_mla=True has_index=True create_result logs
Rejected: Disable partial-current reuse for short extends | that would lose the cache-hit benefit and regress current/full reuse
Rejected: Disable all async prefetch by default | broader performance impact than the observed tiny-suffix failure
Confidence: medium
Scope-risk: moderate
Directive: Do not lower the default below one page per CP lane without ETE proof that repeated cache-hit tiny suffixes no longer hang
Tested: Remote g0034 container py_compile for touched runtime/prefetch/test files; targeted C22 tests passed 2 tests plus 2 subtests; full test_cp_shared_kv_runtime.py passed 73 tests plus 2 subtests
Not-tested: Full multi-node ETE repeated-request run after this threshold change
This commit is contained in:
laoyao0822
2026-05-29 22:04:01 +08:00
parent 40cf691c78
commit 21065cdfdf
4 changed files with 135 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
View File

@@ -437,7 +437,7 @@ class CpSharedKVMlaPrefetcher:
return None
extend_len = int(extend_seq_lens_cpu[0])
min_async_extend_tokens = cp_shared_kv_mla_prefetch_min_async_extend_tokens(
page_size=page_size
cp_size=layout.cp_size, page_size=page_size
)
if extend_len < min_async_extend_tokens:
_prefetch_log(
@@ -1226,7 +1226,7 @@ class CpSharedKVIndexPrefetcher:
return None
extend_len = int(extend_seq_lens_cpu[0])
min_extend_tokens = cp_shared_kv_mla_prefetch_min_async_extend_tokens(
page_size=page_size
cp_size=layout.cp_size, page_size=page_size
)
if extend_len < min_extend_tokens:
_prefetch_log(

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

@@ -97,7 +97,7 @@ def cp_shared_kv_mla_prefetch_min_prefix_pages(
def cp_shared_kv_mla_prefetch_min_async_extend_tokens(
*, page_size: int | None = None
*, cp_size: int | None = None, page_size: int | None = None
) -> int:
"""Minimum current-extend tokens required for async next-layer prefetch.
@@ -106,12 +106,20 @@ def cp_shared_kv_mla_prefetch_min_async_extend_tokens(
backend can still synchronously materialize prefix pages and splice current
KV rows for target partial-current reuse.
A negative env value uses the page size as the dynamic default. Set the env
to 0 to allow async prefetch even for sub-page extends during experiments.
A negative env value uses one page per CP lane as the dynamic default when
both ``cp_size`` and ``page_size`` are known. Set the env to 0 to allow
async prefetch even for tiny extends during experiments.
"""
configured = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS.get()
if configured < 0:
if (
cp_size is not None
and int(cp_size) > 0
and page_size is not None
and int(page_size) > 0
):
return int(cp_size) * int(page_size)
if page_size is not None and int(page_size) > 0:
return int(page_size)
return 0