Prevent incompatible CP shared-KV transfer mapping
Mooncake is the only disaggregation transfer backend in this branch with CP shared-KV owner filtering plus logical-page-position destination selection. NIXL still slices destination pages by the original chunk slice, so allowing CP shared-KV prefill on NIXL can silently pair filtered prefill pages with the wrong decode pages. This keeps the supported path narrow while preserving the page-aligned transfer contract: non-page-aligned valid tails transfer their physical tail page, but do not get padded to CP-size pages. Constraint: CP shared-KV transfer remaps prefill logical pages to per-rank physical pages while decode metadata remains request-position based. Rejected: Let NIXL continue through the generic slice path | it lacks logical-page-position selection and can silently corrupt CP shared-KV transfers. Confidence: high Scope-risk: narrow Directive: Do not enable CP shared-KV on another PD transfer backend until its sender filters owner pages and selects decode pages by logical request-page position. Tested: Local py_compile for server_args and touched tests. Tested: Remote g0034 pytest test_cp_shared_kv_transfer_mapping.py test_req_to_token_pool.py TestHiCacheArgs: 22 passed, 8 subtests passed. Not-tested: End-to-end PD transfer with a live non-page-aligned prompt. Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
@@ -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 = [
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user