Prevent stale CP shared-KV contracts from corrupting prefill

CP shared-KV now uses CP-local current rows consistently across MLA/index current reuse, passes fp8 current-index K through the tai-kernel uint8 ABI, and clears the transient EAGLE CP-local hidden marker after draft capture. The disaggregation bootstrap also fingerprints the runtime source contract so prefill/decode mismatches fail fast instead of silently exchanging incompatible KV metadata.

Constraint: CP shared-KV batch paths flatten current K/V rows in CP-rank-local valid order, not global request order.

Constraint: tai-kernel current-index prepare validates current_index_k as uint8 bytes for fp8 payloads.

Rejected: Keep using global extend offsets for bs>1 current-index reuse | corrupts request-local bases once current_index_kv is CP-local.

Rejected: Infer CP-local EAGLE hidden semantics from tensor shape | static padding and bs>1 can make shape-based inference unsafe.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce forward_batch.out_cache_loc slicing in CP shared-KV current reuse without verifying CP-local owner-lane layout.

Tested: Remote container py_compile for touched runtime/test files.

Tested: Remote PYTHONPATH=python pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py (198 passed, 2 subtests passed).

Not-tested: Full remote ETE traffic after this commit; accept length and garbage-output recovery still require a fresh prefill/decode run.

Co-authored-by: OmX <omx@oh-my-codex.dev>

test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py

@@ -729,7 +729,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
 
         self.assertIn("valid_current_rows=int(current_kv_rows_for_reuse)", body_source)
         self.assertIn("pack_current_mla_kv_for_reuse", body_source)
-        self.assertIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source)
+        self.assertIn("get_cp_shared_kv_local_out_cache_loc", body_source)
+        self.assertNotIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source)
 
     def test_flashmla_kv_current_only_reuse_keeps_page_slot_layout(self):
         from pathlib import Path
@@ -755,12 +756,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             Path(__file__).resolve().parents[4]
             / "python/sglang/srt/layers/attention/nsa_backend.py"
         ).read_text()
+        method_start = source.index("            current_locs_for_reuse = None")
         branch_start = source.index("eagle_draft_mla_branch = \"partial_current_sync\"")
         branch_end = source.index("forward_partial_current_sync_compose", branch_start)
         branch_source = source[branch_start:branch_end]
+        local_locs_source = source[method_start:branch_end]
 
         self.assertIn("build_batch_prefix_slot_span", source)
         self.assertIn("prefix_slot_span=", branch_source)
+        self.assertIn("get_cp_shared_kv_local_out_cache_loc", local_locs_source)
+        self.assertNotIn("current_locs = forward_batch.out_cache_loc", branch_source)
         self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source)
 
     def test_nsa_backend_topk_transform_uses_effective_forward_impl(self):
@@ -1504,12 +1509,18 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             Path(__file__).resolve().parents[4]
             / "python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
         ).read_text()
-        branch_start = source.index("        if current_index_kv is not None:")
+        method_start = source.index("    def _maybe_materialize_shared_index_buffer")
+        branch_start = source.index(
+            "        if current_index_kv is not None:", method_start
+        )
         branch_end = source.index("            return materialized, dense_pages", branch_start)
         branch_source = source[branch_start:branch_end]
+        branch_compact = "".join(branch_source.split())
 
         self.assertIn("build_batch_prefix_slot_span", source)
         self.assertIn("prefix_slot_span=", branch_source)
+        self.assertIn("get_cp_shared_kv_local_out_cache_loc", branch_source)
+        self.assertNotIn("current_locs=forward_batch.out_cache_loc", branch_compact)
         self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source)
 
     def test_ipc_page_descriptor_builder_maps_slots_to_owner_physical_pages(self):
@@ -2141,7 +2152,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         prepare_source = source[start:end]
 
         self.assertIn("valid_current_rows", prepare_source)
-        self.assertIn("key[:valid_current_rows]", prepare_source)
+        self.assertIn("get_cp_shared_kv_local_out_cache_loc", prepare_source)
+        self.assertIn("cp_split_and_rebuild_data", prepare_source)
+        self.assertNotIn("forward_batch.out_cache_loc[:valid_current_rows]", prepare_source)
         self.assertNotIn(
             "key.shape[0] == forward_batch.out_cache_loc.numel()",
             prepare_source,
@@ -2852,6 +2865,91 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         load_kernel.assert_not_called()
         self.assertTrue(any("debug_enabled" in message for message in logs.output))
 
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
+    def test_fp8_fused_mla_store_matches_sglang_fallback_for_cp_owned_tail_pages(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.layers.attention.nsa.quant_k_cache import (
+            quantize_k_cache_separate,
+        )
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+        from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton
+
+        try:
+            runtime._load_tai_fused_mla_store_kernel()
+        except Exception as exc:
+            self.skipTest(f"TAI fused MLA store kernel unavailable: {exc}")
+
+        page_size = 64
+        cp_size = 8
+        rank = 7
+        n_tokens = page_size * 3 + 17
+        layout = CpSharedKVLayout(
+            page_size=page_size,
+            cp_size=cp_size,
+            cp_rank=rank,
+        )
+        local_token_ids = torch.arange(n_tokens, device="cuda", dtype=torch.long)
+        local_pages = torch.div(local_token_ids, page_size, rounding_mode="floor")
+        logical_pages = 1 + rank + cp_size * local_pages
+        logical_locs = (
+            logical_pages * page_size + torch.remainder(local_token_ids, page_size)
+        ).to(torch.int64)
+        physical_locs = layout.logical_locs_to_physical(logical_locs)
+        capacity_tokens = int(physical_locs.max().item()) + page_size + 1
+
+        torch.manual_seed(20260604)
+        k_nope = (
+            torch.randn((n_tokens, 1, 512), device="cuda", dtype=torch.bfloat16)
+            * 2.0
+        ) + 0.25
+        latent_cache = torch.randn(
+            (n_tokens, 576), device="cuda", dtype=torch.bfloat16
+        )
+        k_rope = latent_cache[:, 512:].unsqueeze(1)
+        self.assertFalse(k_rope.is_contiguous())
+
+        expected = torch.zeros(
+            (capacity_tokens, 1, 656), dtype=torch.uint8, device="cuda"
+        )
+        nope_part, rope_part = quantize_k_cache_separate(k_nope, k_rope)
+        set_mla_kv_buffer_triton(expected, physical_locs, nope_part, rope_part)
+
+        class FakePool:
+            def __init__(self):
+                self.nsa_kv_cache_store_fp8 = True
+                self.page_size = page_size
+                self.start_layer = 0
+                self.kv_buffer = [
+                    torch.zeros(
+                        (capacity_tokens, 1, 656),
+                        dtype=torch.uint8,
+                        device="cuda",
+                    )
+                ]
+
+        pool = FakePool()
+        layer = SimpleNamespace(layer_id=0)
+        with patch.object(
+            runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=True
+        ), patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=False):
+            used = runtime.try_tai_fused_mla_store(
+                token_to_kv_pool=pool,
+                layer=layer,
+                layout=layout,
+                logical_locs=logical_locs,
+                k_nope=k_nope,
+                k_rope=k_rope,
+            )
+
+        torch.cuda.synchronize()
+        self.assertTrue(used)
+        torch.testing.assert_close(
+            pool.kv_buffer[0],
+            expected,
+            atol=0,
+            rtol=0,
+        )
+
     def test_token_range_materialize_uses_tai_kernel_when_enabled(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
         from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout