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 5b3c6dd18..1faf086ec 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 @@ -732,6 +732,15 @@ Current state: token-length oriented. - There is no current checklist proving that decode never infers prompt length from transferred padded page count. +- Mooncake has CP shared-KV page filtering in + `MooncakeKVSender.send()` / `filter_kv_pages_for_cp_shared_kv()`, where + prefill logical pages are filtered by owner lane and remapped to this rank's + physical page ids while destination pages are selected by request-page + position. +- NIXL does not have the equivalent CP shared-KV logical-page-position mapping. + Its sender still slices `req.dst_kv_indices[index_slice]` directly. Under CP + shared KV this can silently pair owner-filtered prefill physical pages with the + wrong decode destination pages. Correction: @@ -739,6 +748,9 @@ Correction: - Request/bootstrap metadata must carry valid token length separately. - Decode-visible prompt length and max-context checks must use valid length. - Draft transfer follows target page pattern. +- Until NIXL gets the same logical-position mapping as Mooncake, CP shared-KV PD + transfer must fail fast for non-Mooncake transfer backends instead of silently + taking an incompatible path. Tests: @@ -746,6 +758,21 @@ Tests: - transferred pages include the padded tail page; - decode sees the original valid length; - draft and target transfer page counts match when draft KV is enabled. +- Unit-level guard: + - CP shared-KV page filtering keeps the padded tail page for a valid-tail + prompt without padding to CP-size pages; + - `ServerArgs` rejects CP shared-KV prefill with NIXL transfer until that + backend implements logical-position selection. + +Implemented C10 slice: + +- Added unit coverage that a 100-token / 64-page prompt covers two transfer + pages, and owner rank 1 transfers only the tail page selected by logical + request-page position. This locks the "pad only to page boundary, not + cp-size pages" transfer behavior. +- Added a startup validation guard: CP shared-KV prefill disaggregation is + currently Mooncake-only. NIXL now fails fast because it still lacks + Mooncake's logical-page-position mapping for owner-filtered transfer pages. ### C11. Fallback logging must be audited diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9acb65d1d..5d580017b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -917,6 +917,16 @@ class ServerArgs: "hicache_storage_backend in the host-only CP HiCache stage. " "Disable hicache_storage_backend or disable CP shared KV." ) + assert not ( + self.enable_nsa_prefill_cp_shared_kv + and self.disaggregation_mode == "prefill" + and self.disaggregation_transfer_backend != "mooncake" + ), ( + "enable_nsa_prefill_cp_shared_kv with prefill disaggregation " + "currently requires disaggregation_transfer_backend='mooncake'. " + "Other backends do not implement CP shared-KV logical-page-position " + "transfer mapping yet." + ) def _handle_cp_hicache_layout_validation(self): if not ( diff --git a/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py b/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py index bd7e78f53..eceb72a84 100644 --- a/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py +++ b/test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py @@ -27,6 +27,26 @@ class TestCPSharedKVTransferMapping(unittest.TestCase): self.assertEqual(src_pages.tolist(), [2, 3]) self.assertEqual(positions.tolist(), [102, 106]) + def test_filter_valid_tail_keeps_tail_page_without_cp_size_padding(self): + layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=1) + # A 100-token prompt physically covers exactly two pages. The second + # page is a valid tail page and must be transferred by its owner rank, + # but the transfer path must not pad the request to cp_size pages. + logical_pages = np.array([1, 2], dtype=np.int32) + src_pages, positions = filter_kv_pages_for_cp_shared_kv( + layout=layout, + logical_pages=logical_pages, + chunk_page_start=0, + ) + selected_decode_pages = select_pages_by_request_positions( + np.array([41, 42], dtype=np.int32), + positions, + ) + + self.assertEqual(src_pages.tolist(), [1]) + self.assertEqual(positions.tolist(), [1]) + self.assertEqual(selected_decode_pages.tolist(), [42]) + def test_filter_kv_pages_for_cp_shared_kv_rejects_negative_pages(self): layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 7937c0589..8a9d26112 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -415,6 +415,29 @@ class TestHiCacheArgs(CustomTestCase): hicache_storage_backend="mooncake", ) + def test_cp_shared_kv_prefill_rejects_nixl_transfer_backend(self): + with ( + self.assertRaisesRegex( + AssertionError, + "enable_nsa_prefill_cp_shared_kv.*disaggregation_transfer_backend.*mooncake", + ), + patch.object( + ServerArgs, + "get_model_config", + side_effect=AssertionError("get_model_config should not be called"), + ), + ): + ServerArgs( + model_path="dummy", + enable_nsa_prefill_context_parallel=True, + enable_nsa_prefill_cp_shared_kv=True, + nsa_prefill_cp_mode="in-seq-split", + disaggregation_mode="prefill", + disaggregation_transfer_backend="nixl", + page_size=64, + enable_hisparse=False, + ) + def test_hicache_io_backend_and_mem_layout_compatibility(self): cases = [ {