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