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/layers/test_nsa_cp_utils.py

@@ -1322,6 +1322,28 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2)))
 
+    def test_cp_split_and_rebuild_data_preserves_non_token_dimensions(self):
+        import torch
+
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=2,
+                request_extend_lens=[4, 9],
+                request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
+                request_zigzag_indices=[[0, 3], [0, 3]],
+            )
+        )
+        tensor = torch.arange(13 * 2 * 3, dtype=torch.float32).view(13, 2, 3)
+
+        with patch(
+            "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+            return_value=False,
+        ):
+            local = cp_split_and_rebuild_data(forward_batch, tensor)
+
+        self.assertEqual(tuple(local.shape), (8, 2, 3))
+        self.assertTrue(torch.equal(local[0], tensor[0]))
+
     def test_cp_split_and_rebuild_data_uses_compute_padding_rows(self):
         import torch
 
@@ -3309,6 +3331,122 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10])
         self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14])
 
+    def test_indexer_ragged_cp_index_current_batch_uses_cp_local_bases_and_uint8_k(
+        self,
+    ):
+        import contextlib
+        import torch
+        from types import SimpleNamespace
+        from unittest.mock import patch
+
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        indexer.index_topk = 2
+        indexer._with_real_sm_count = lambda: contextlib.nullcontext()
+
+        prepare_calls = []
+
+        def fake_prepare(**kwargs):
+            prepare_calls.append(kwargs)
+            total_kv_len = int(kwargs["total_kv_len"])
+            return (
+                torch.zeros((total_kv_len, 1), dtype=torch.uint8),
+                torch.zeros((total_kv_len,), dtype=torch.float32),
+                torch.tensor([0, 0, 3, 3, 3], dtype=torch.int32),
+                torch.tensor([2, 3, 2, 3, 4], dtype=torch.int32),
+            )
+
+        class Metadata:
+            def get_page_table_64(self):
+                raise AssertionError("current cp_index path must not materialize index pages")
+
+            def topk_transform(self, logits, topk, **kwargs):
+                return (
+                    torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
+                    .view(-1, 1)
+                    .repeat(1, topk)
+                )
+
+        forward_batch = SimpleNamespace(
+            token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1),
+            seq_lens_cpu=torch.tensor([5, 7], dtype=torch.int64),
+            extend_seq_lens_cpu=[5, 7],
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=2,
+                batch_plan=SimpleNamespace(
+                    request_rank_local_offsets=[0, 2],
+                    request_valid_rank_local_offsets=[0, 2],
+                ),
+            ),
+        )
+        current_index_k = (
+            torch.arange(6, dtype=torch.float32)
+            .to(torch.float8_e4m3fn)
+            .view(6, 1)
+        )
+        current_index_scale = torch.arange(6, dtype=torch.float32).view(6, 1)
+        current_index_kv = (current_index_k, current_index_scale)
+
+        with patch.object(
+            nsa_indexer,
+            "try_tai_prepare_cp_mqa_current_index_batch",
+            side_effect=fake_prepare,
+            create=True,
+        ), patch.object(
+            nsa_indexer,
+            "deep_gemm",
+            SimpleNamespace(
+                fp8_mqa_logits=lambda q_fp8, kv_fp8, weights, ks, ke, clean_logits=False: torch.zeros(
+                    (int(q_fp8.shape[0]), 8), dtype=torch.float32
+                )
+            ),
+        ):
+            result = Indexer._get_topk_ragged_with_cp(
+                indexer,
+                forward_batch,
+                layer_id=7,
+                q_fp8=torch.empty((5, 1), dtype=torch.float32),
+                weights=torch.empty((5, 1, 1), dtype=torch.float32),
+                metadata=Metadata(),
+                kv_len=0,
+                actual_seq_q=5,
+                cp_index=[(0, 1, 3), (1, 1, 4)],
+                current_index_kv=current_index_kv,
+            )
+
+        self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
+        self.assertEqual(len(prepare_calls), 1)
+        call = prepare_calls[0]
+        self.assertEqual(call["current_index_k"].dtype, torch.uint8)
+        self.assertEqual(call["current_bases"].tolist(), [0, 2])
+        self.assertEqual(call["kv_lens"].tolist(), [3, 4])
+        self.assertEqual(call["q_lens"].tolist(), [2, 3])
+
+    def test_eagle_capture_for_decode_clears_cp_local_hidden_marker(self):
+        import torch
+
+        from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+        from sglang.srt.speculative.eagle_worker import EAGLEWorker
+
+        worker = object.__new__(EAGLEWorker)
+        worker.topk = 1
+        draft_input = EagleDraftInput(
+            hidden_states=torch.full((4, 2), -1.0),
+            cp_local_hidden_states=True,
+        )
+        draft_output_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        logits_output = LogitsProcessorOutput(
+            next_token_logits=torch.tensor([[0.1, 0.9], [0.7, 0.3]]),
+            hidden_states=draft_output_hidden,
+        )
+
+        worker.capture_for_decode(logits_output, draft_input)
+
+        self.assertIs(draft_input.hidden_states, draft_output_hidden)
+        self.assertFalse(draft_input.cp_local_hidden_states)
+
     def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self):
         import torch