From 21065cdfdf75de31b6dada165aad12e08834451e Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 29 May 2026 22:04:01 +0800 Subject: [PATCH] Keep tiny CP cache-hit suffixes off async prefetch 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 --- ..._prefill_cp_page_aligned_cache_contract.md | 53 +++++++- .../attention/nsa/cp_shared_kv_prefetch.py | 4 +- .../attention/nsa/cp_shared_kv_runtime.py | 14 +- .../mem_cache/test_cp_shared_kv_runtime.py | 123 ++++++++++-------- 4 files changed, 135 insertions(+), 59 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 95496c41a..7fe3ab187 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -1490,8 +1490,10 @@ Correction: - If partial-current reuse is intentionally rejected, for example the current conservative draft/EAGLE cache-hit guard, that rejection must not disable current-only/full-current reuse for other request shapes. -- Default the async gate to one page when `page_size` is known; allow an env - override to lower or raise the threshold for experiments. +- C21 initially defaulted the async gate to one page when `page_size` is known; + C22 supersedes that default with one page per CP lane after the 65-token + repeated-request hang showed that one page was still too small. The env + override remains available for experiments. - This must not force tiny cache-hit suffixes back to compact full materialization. @@ -1509,3 +1511,50 @@ Verified: - Remote container `/sgl-workspace/sglang-tai` on `g0034`: - `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`: `73 passed` - `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`: `90 passed` + +## C22: One-page async prefetch gate still admits repeated-request hangs + +Finding: + +- Latest remote log `/mnt/beegfs/cjy/sglang_cp_hicache_20260529_134616.log` + shows the first 40k request completed and backed up CP HiCache node 4. +- The repeated request hit the cached prefix and created both MLA and index + prefetchers on every CP rank: + `prefix_lens=[40320]`, `extend_lens=[65]`, `prefix_pages=630`, + `min_async_extend_tokens=64`, `has_mla=True has_index=True`. +- No `index_start_request` or `forward_probe` appears after those `create_result` + lines; the next lines are detokenizer heartbeat timeout and shutdown cleanup. +- Therefore C21's one-page default threshold is insufficient. It blocks + sub-page suffixes but still allows a 65-token suffix, i.e. just over one page, + to enter async next-layer prefetch. + +Correction: + +- The intended production default is CP-lane scale, not one page: with + `cp_size=8,page_size=64`, async MLA/index prefetch should require at least + `8 * 64 = 512` current extend tokens unless explicitly overridden. +- This matches the owner-lane mental model: below one page per CP rank, the + prefetch collective and page-slot compose overhead is not amortized and is a + correctness risk while padded-tail ETE coverage is still incomplete. +- This gate only controls async prefetcher object creation. Target + partial-current reuse and full current reuse must continue through the + synchronous page-slot compose/current-splice path when no prefetcher exists. + +Required test update: + +- Change the min-extend helper contract from "one page" to "one page per CP + lane" when both `cp_size` and `page_size` are known. +- Add a regression case for `extend_len=65,page_size=64,cp_size=8` proving both + MLA and index prefetchers return `None` before touching pool getters/remap + builders. + +Verified C22: + +- Remote container `/sgl-workspace/sglang-tai` on `g0034`: + - `python -m py_compile` for `cp_shared_kv_runtime.py`, + `cp_shared_kv_prefetch.py`, and `test_cp_shared_kv_runtime.py` passed. + - Targeted C22 tests passed: `2 passed, 2 subtests passed`. + - Full `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py` passed: + `73 passed, 2 subtests passed`. +- Local pytest remains blocked by missing optional deps such as `orjson`; remote + container is the authoritative verification environment for this path. diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index f07517db6..977a6d78b 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -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( diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index df680775c..28aeebb46 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 0e1d04a81..c22955623 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -1404,7 +1404,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): 32, ) - def test_mla_prefetch_min_async_extend_tokens_defaults_to_one_page_and_can_override( + def test_mla_prefetch_min_async_extend_tokens_defaults_to_one_page_per_cp_lane_and_can_override( self, ): from sglang.srt.environ import envs @@ -1412,7 +1412,15 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS.clear() self.assertEqual( - runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens(page_size=64), + runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens( + cp_size=8, page_size=64 + ), + 512, + ) + self.assertEqual( + runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens( + cp_size=None, page_size=64 + ), 64, ) self.assertEqual( @@ -1422,13 +1430,17 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS.override(0): self.assertEqual( - runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens(page_size=64), + runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens( + cp_size=8, page_size=64 + ), 0, ) with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS.override(128): self.assertEqual( - runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens(page_size=64), + runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens( + cp_size=8, page_size=64 + ), 128, ) @@ -1439,16 +1451,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): def is_context_parallel_extend(self): return True - forward_batch = SimpleNamespace( - uses_cp_shared_kv=True, - hisparse_coordinator=None, - forward_mode=Mode(), - batch_size=1, - token_to_kv_pool=SimpleNamespace(page_size=64, start_layer=0), - cp_shared_kv_layout=SimpleNamespace(cp_size=8, cp_rank=0), - extend_prefix_lens_cpu=[16320], - extend_seq_lens_cpu=[16], - ) metadata = SimpleNamespace( real_page_table=torch.arange(256, dtype=torch.int64), page_table_1=torch.zeros((1, 16336), dtype=torch.int32), @@ -1461,45 +1463,62 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): dense_num_pages=257, ) - with patch.object( - prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True - ), patch.object( - prefetch, "cp_shared_kv_debug_enabled", return_value=False - ), patch.object( - prefetch.torch.cuda, "is_available", return_value=True - ), patch.object( - prefetch, "_is_cuda_stream_capturing", return_value=False - ), patch.object( - prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True - ), patch.object( - prefetch, "get_attention_cp_group", return_value=SimpleNamespace(pynccl_comm=object()) - ), patch.object( - prefetch.torch.cuda, "Stream", return_value=stream - ), patch.object( - prefetch, "_prefetch_pool_get_key_buffer", return_value=kv_cache - ) as mla_getter, patch.object( - prefetch, "get_or_build_shared_token_kv_slot_remap", return_value=remap - ) as token_remap, patch.object( - prefetch, - "_prefetch_pool_get_index_buffer", - side_effect=AssertionError("index getter should not be reached"), - ) as index_getter: - mla_result = prefetch.CpSharedKVMlaPrefetcher.maybe_create( - forward_batch=forward_batch, - metadata=metadata, - topk_transform_is_paged=True, - ) - index_result = prefetch.CpSharedKVIndexPrefetcher.maybe_create( - forward_batch=forward_batch, - metadata=metadata, - topk_transform_is_paged=True, - ) + for extend_len in (16, 65): + with self.subTest(extend_len=extend_len): + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + hisparse_coordinator=None, + forward_mode=Mode(), + batch_size=1, + token_to_kv_pool=SimpleNamespace(page_size=64, start_layer=0), + cp_shared_kv_layout=SimpleNamespace(cp_size=8, cp_rank=0), + extend_prefix_lens_cpu=[16320], + extend_seq_lens_cpu=[extend_len], + ) - self.assertIsNone(mla_result) - self.assertIsNone(index_result) - mla_getter.assert_not_called() - token_remap.assert_not_called() - index_getter.assert_not_called() + with patch.object( + prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True + ), patch.object( + prefetch, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + prefetch.torch.cuda, "is_available", return_value=True + ), patch.object( + prefetch, "_is_cuda_stream_capturing", return_value=False + ), patch.object( + prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True + ), patch.object( + prefetch, + "get_attention_cp_group", + return_value=SimpleNamespace(pynccl_comm=object()), + ), patch.object( + prefetch.torch.cuda, "Stream", return_value=stream + ), patch.object( + prefetch, "_prefetch_pool_get_key_buffer", return_value=kv_cache + ) as mla_getter, patch.object( + prefetch, + "get_or_build_shared_token_kv_slot_remap", + return_value=remap, + ) as token_remap, patch.object( + prefetch, + "_prefetch_pool_get_index_buffer", + side_effect=AssertionError("index getter should not be reached"), + ) as index_getter: + mla_result = prefetch.CpSharedKVMlaPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=True, + ) + index_result = prefetch.CpSharedKVIndexPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=True, + ) + + self.assertIsNone(mla_result) + self.assertIsNone(index_result) + mla_getter.assert_not_called() + token_remap.assert_not_called() + index_getter.assert_not_called() def test_fused_mla_store_uses_tai_kernel_when_enabled(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime