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>

python/sglang/srt/disaggregation/common/conn.py

@@ -2,11 +2,14 @@ from __future__ import annotations
 
 import asyncio
 import dataclasses
+import hashlib
+import importlib.util
 import logging
 import threading
 import time
 from collections import defaultdict
 from functools import cache
+from pathlib import Path
 from typing import Dict, List, Optional, Set, Tuple, Union
 
 import numpy as np
@@ -44,6 +47,75 @@ from sglang.srt.utils.network import (
 logger = logging.getLogger(__name__)
 
 
+_RUNTIME_FINGERPRINT_MODULES = (
+    "sglang.srt.layers.attention.nsa_backend",
+    "sglang.srt.layers.attention.nsa.nsa_indexer",
+    "sglang.srt.layers.attention.nsa.utils",
+    "sglang.srt.mem_cache.memory_pool",
+    "sglang.srt.disaggregation.mooncake.conn",
+)
+
+
+@cache
+def _get_runtime_python_root() -> Optional[Path]:
+    spec = importlib.util.find_spec("sglang")
+    origin = getattr(spec, "origin", None) if spec is not None else None
+    if origin is None:
+        return None
+    path = Path(origin).resolve()
+    # .../<root>/python/sglang/__init__.py
+    for parent in path.parents:
+        if parent.name == "python":
+            return parent
+    return path.parent.parent
+
+
+@cache
+def get_runtime_source_root() -> Optional[str]:
+    python_root = _get_runtime_python_root()
+    if python_root is None:
+        return None
+    return str(python_root.parent)
+
+
+@cache
+def get_runtime_source_fingerprint() -> Optional[str]:
+    """Return a stable fingerprint for code that defines the P/D KV contract.
+
+    Prefill and decode may run from different mounted source trees in long-lived
+    containers.  CP shared-KV + FP8 + page-tail handling is not wire-compatible
+    across our recent development snapshots, so decode must fail fast when the
+    source contract differs instead of silently consuming malformed KV.
+
+    Do not import the target modules while fingerprinting: some package
+    ``__init__`` files import the connection modules themselves.  Build paths
+    from the already-imported ``sglang`` package root instead.
+    """
+
+    python_root = _get_runtime_python_root()
+    if python_root is None:
+        return None
+
+    digest = hashlib.sha256()
+    for module_name in _RUNTIME_FINGERPRINT_MODULES:
+        origin = python_root.joinpath(*module_name.split(".")).with_suffix(".py")
+        digest.update(module_name.encode("utf-8"))
+        digest.update(b"\0")
+        try:
+            digest.update(origin.read_bytes())
+        except OSError as exc:
+            logger.warning(
+                "Failed to read source for disaggregation runtime fingerprint: "
+                "module=%s origin=%s error=%s",
+                module_name,
+                origin,
+                exc,
+            )
+            return None
+        digest.update(b"\0")
+    return digest.hexdigest()
+
+
 @dataclasses.dataclass
 class PrefillServerInfo:
     # Topology fields (fetched from bootstrap server)
@@ -54,6 +126,8 @@ class PrefillServerInfo:
     page_size: Optional[int]
     kv_cache_dtype: Optional[str]
     follow_bootstrap_room: bool
+    runtime_source_fingerprint: Optional[str] = None
+    runtime_source_root: Optional[str] = None
 
     # Pre-computed rank mapping (set by try_ensure_parallel_info on decode side)
     target_tp_rank: Optional[int] = None
@@ -73,6 +147,16 @@ class PrefillServerInfo:
             str(self.kv_cache_dtype) if self.kv_cache_dtype is not None else None
         )
         self.follow_bootstrap_room = bool(self.follow_bootstrap_room)
+        self.runtime_source_fingerprint = (
+            str(self.runtime_source_fingerprint)
+            if self.runtime_source_fingerprint is not None
+            else None
+        )
+        self.runtime_source_root = (
+            str(self.runtime_source_root)
+            if self.runtime_source_root is not None
+            else None
+        )
 
 
 @dataclasses.dataclass
@@ -236,6 +320,22 @@ class CommonKVManager(BaseKVManager):
                 f"Both servers must use the same --kv-cache-dtype value."
             )
 
+        local_source_fingerprint = get_runtime_source_fingerprint()
+        if (
+            info.runtime_source_fingerprint is not None
+            and local_source_fingerprint is not None
+            and info.runtime_source_fingerprint != local_source_fingerprint
+        ):
+            raise RuntimeError(
+                "Runtime source fingerprint mismatch between prefill and decode. "
+                "This can corrupt disaggregated KV transfer contracts. "
+                f"prefill_fingerprint={info.runtime_source_fingerprint}, "
+                f"decode_fingerprint={local_source_fingerprint}, "
+                f"prefill_source_root={info.runtime_source_root}, "
+                f"decode_source_root={get_runtime_source_root()}. "
+                "Restart both prefill and decode from the same source tree."
+            )
+
         self._resolve_rank_mapping(info)
         self.prefill_info_table[bootstrap_addr] = info
         logger.debug(f"Prefill parallel info for [{bootstrap_addr}]: {info}")
@@ -348,6 +448,8 @@ class CommonKVManager(BaseKVManager):
             "page_size": self.kv_args.page_size,
             "kv_cache_dtype": self.server_args.kv_cache_dtype,
             "load_balance_method": self.server_args.load_balance_method,
+            "runtime_source_fingerprint": get_runtime_source_fingerprint(),
+            "runtime_source_root": get_runtime_source_root(),
         }
 
         try:
@@ -668,6 +770,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
         self.page_size = None
         self.kv_cache_dtype: Optional[str] = None
         self.follow_bootstrap_room: Optional[bool] = None
+        self.runtime_source_fingerprint: Optional[str] = get_runtime_source_fingerprint()
+        self.runtime_source_root: Optional[str] = get_runtime_source_root()
         self.prefill_port_table: Dict[
             int, Dict[int, Dict[int, Dict[int, PrefillRankInfo]]]
         ] = {}
@@ -734,6 +838,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
         rank_port = int(data["rank_port"])
         page_size = int(data["page_size"])
         kv_cache_dtype = data["kv_cache_dtype"]
+        runtime_source_fingerprint = data.get("runtime_source_fingerprint")
+        runtime_source_root = data.get("runtime_source_root")
 
         if self.attn_tp_size is None:
             self.attn_tp_size = attn_tp_size
@@ -753,6 +859,15 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
         if self.kv_cache_dtype is None and kv_cache_dtype is not None:
             self.kv_cache_dtype = kv_cache_dtype
 
+        if (
+            self.runtime_source_fingerprint is None
+            and runtime_source_fingerprint is not None
+        ):
+            self.runtime_source_fingerprint = str(runtime_source_fingerprint)
+
+        if self.runtime_source_root is None and runtime_source_root is not None:
+            self.runtime_source_root = str(runtime_source_root)
+
         if self.follow_bootstrap_room is None:
             load_balance_method = data.get(
                 "load_balance_method", "follow_bootstrap_room"
@@ -822,6 +937,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
                     if self.follow_bootstrap_room is not None
                     else True
                 ),
+                runtime_source_fingerprint=self.runtime_source_fingerprint,
+                runtime_source_root=self.runtime_source_root,
             )
             return web.json_response(dataclasses.asdict(info), status=200)
 

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -68,6 +68,7 @@ from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.nsa.utils import (
     cp_all_gather_rerange_output,
+    cp_split_and_rebuild_data,
     get_cp_shared_kv_batch_plan,
     get_cp_shared_kv_local_out_cache_loc,
     get_cp_shared_kv_local_physical_out_cache_loc,
@@ -132,6 +133,65 @@ class BatchTopKQueryLengths:
     compute_token_count: int
 
 
+def _current_index_k_for_tai(current_index_k: torch.Tensor) -> torch.Tensor:
+    """Return the current-index K tensor in the uint8 ABI expected by tai-kernel."""
+
+    if current_index_k.dtype == torch.uint8:
+        return current_index_k
+    fp8_dtypes = {
+        dtype
+        for dtype in (
+            getattr(torch, "float8_e4m3fn", None),
+            getattr(torch, "float8_e5m2", None),
+            getattr(torch, "float8_e4m3fnuz", None),
+            getattr(torch, "float8_e5m2fnuz", None),
+        )
+        if dtype is not None
+    }
+    if current_index_k.dtype in fp8_dtypes:
+        return current_index_k.view(torch.uint8)
+    return current_index_k
+
+
+def _build_current_index_request_bases(forward_batch: ForwardBatch) -> List[int]:
+    """Build request bases for flattened current-index K/S rows.
+
+    Without CP shared-KV batch metadata the legacy current-index tensor is a
+    global flattened `[req0 current, req1 current, ...]` buffer, so global
+    extend offsets are correct.  With CP shared-KV batch metadata the tensor is
+    CP-rank-local valid rows, so request bases must come from the local valid
+    owner-lane plan instead of global extend lengths.
+    """
+
+    extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
+    if extend_seq_lens_cpu is None:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+            "reason=current_index_missing_extend_lens"
+        )
+
+    batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+    if batch_plan is not None:
+        offsets = getattr(batch_plan, "request_valid_rank_local_offsets", None)
+        if offsets is None:
+            offsets = getattr(batch_plan, "request_rank_local_offsets", None)
+        if offsets is None or len(offsets) != len(extend_seq_lens_cpu):
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                "reason=current_index_missing_local_bases "
+                f"batch_size={len(extend_seq_lens_cpu)} "
+                f"offsets={offsets}"
+            )
+        return [int(x) for x in offsets]
+
+    current_req_offsets = []
+    current_cursor = 0
+    for extend_len in extend_seq_lens_cpu:
+        current_req_offsets.append(current_cursor)
+        current_cursor += int(extend_len)
+    return current_req_offsets
+
+
 def _select_batch_topk_query_lengths(
     *,
     cp_metadata,
@@ -503,29 +563,33 @@ class Indexer(MultiPlatformOp):
                         f"current_locs_shape={tuple(current_locs.shape)}"
                     )
             else:
-                current_locs = forward_batch.out_cache_loc
-                valid_current_rows = current_extend_kv_rows_for_reuse(
-                    forward_batch,
-                    current_index_kv[0],
-                    current_index_kv[1],
-                )
-                if valid_current_rows is None:
+                current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+                if current_locs is None:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
+                        "CP shared KV index partial-current compose requires local "
+                        "valid out_cache_loc. "
+                        f"cp_rank={layout.cp_rank} layer_id={layer_id} "
+                        f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
+                        f"current_k_shape={tuple(current_index_kv[0].shape)} "
+                        f"current_scale_shape={tuple(current_index_kv[1].shape)}"
+                    )
+                valid_current_rows = int(current_locs.numel())
+                if (
+                    int(current_index_kv[0].shape[0]) != valid_current_rows
+                    or int(current_index_kv[1].shape[0]) != valid_current_rows
+                ):
                     raise RuntimeError(
                         "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
                         "CP shared KV index partial-current compose received "
-                        "current_index_kv that does not satisfy current reuse "
-                        "metadata. "
+                        "current_index_kv rows that do not match local valid "
+                        "out_cache_loc. "
                         f"cp_rank={layout.cp_rank} layer_id={layer_id} "
                         f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
                         f"current_k_shape={tuple(current_index_kv[0].shape)} "
                         f"current_scale_shape={tuple(current_index_kv[1].shape)} "
                         f"out_cache_loc_shape={tuple(current_locs.shape)}"
                     )
-                current_locs = current_locs[:valid_current_rows]
-                current_index_kv = (
-                    current_index_kv[0][:valid_current_rows],
-                    current_index_kv[1][:valid_current_rows],
-                )
             prefix_slot_span = None
             if len(prefix_lens_cpu) == 1:
                 prefix_pages = int(prefix_lens_cpu[0]) // page_size
@@ -1333,13 +1397,11 @@ class Indexer(MultiPlatformOp):
         if cp_index is not None:
             current_req_offsets: Optional[List[int]] = None
             if current_index_kv is not None:
-                current_req_offsets = []
-                current_cursor = 0
-                for extend_len in forward_batch.extend_seq_lens_cpu:
-                    current_req_offsets.append(current_cursor)
-                    current_cursor += int(extend_len)
+                current_req_offsets = _build_current_index_request_bases(
+                    forward_batch
+                )
 
-            segment_records: List[Tuple[int, int, int, int, int, int, int]] = []
+            segment_records: List[Tuple[int, int, int, int, int, int, int, int]] = []
             batch_idx_list = []
             kv_lens_list = []
             q_starts_list = []
@@ -1374,6 +1436,7 @@ class Indexer(MultiPlatformOp):
                         kv_len_i,
                         k_cursor,
                         q_cursor,
+                        int(pre_chunk_offset),
                     )
                 )
                 batch_idx_list.append(batch_idx)
@@ -1455,7 +1518,7 @@ class Indexer(MultiPlatformOp):
                     int(current_index_kv[0].reshape(current_index_kv[0].shape[0], -1).shape[1]),
                 )
                 tai_current_prepared = try_tai_prepare_cp_mqa_current_index_batch(
-                    current_index_k=current_index_kv[0],
+                    current_index_k=_current_index_k_for_tai(current_index_kv[0]),
                     current_index_scale=current_index_kv[1],
                     current_bases=torch.tensor(
                         current_bases_list, dtype=torch.int32, device=descriptor_device
@@ -1516,6 +1579,7 @@ class Indexer(MultiPlatformOp):
                 _kv_len_i,
                 segment_k_base,
                 _segment_q_base,
+                pre_chunk_offset,
             ) in segment_records:
                 if current_index_kv is None:
                     k_fp8 = index_buf_accessor.GetK.execute(
@@ -2437,19 +2501,22 @@ class Indexer(MultiPlatformOp):
                     int(current_locs.numel()) if current_locs is not None else None
                 )
             else:
-                valid_current_rows = current_extend_kv_rows_for_reuse(
+                global_valid_current_rows = current_extend_kv_rows_for_reuse(
                     forward_batch, key
                 )
-                current_key = (
-                    key[:valid_current_rows]
-                    if valid_current_rows is not None
-                    else None
-                )
-                current_locs = (
-                    forward_batch.out_cache_loc[:valid_current_rows]
-                    if valid_current_rows is not None
-                    else None
+                current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+                valid_current_rows = (
+                    int(current_locs.numel()) if current_locs is not None else None
                 )
+                current_key = None
+                if global_valid_current_rows is not None and valid_current_rows is not None:
+                    if int(local_key.shape[0]) == valid_current_rows:
+                        current_key = local_key
+                    else:
+                        current_key = cp_split_and_rebuild_data(
+                            forward_batch,
+                            key[: int(global_valid_current_rows)].contiguous(),
+                        )
             if (
                 valid_current_rows is not None
                 and current_key is not None

python/sglang/srt/layers/attention/nsa/utils.py

@@ -1558,7 +1558,7 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
     )
     result = torch.cat(
         [input_list[i] for i in metadata.zigzag_index], dim=0
-    ).view(-1, input_.shape[-1])
+    ).view(-1, *input_.shape[1:])
     return result
 
 

python/sglang/srt/layers/attention/nsa_backend.py

@@ -50,6 +50,7 @@ from sglang.srt.layers.attention.nsa.transform_index import (
 )
 from sglang.srt.layers.attention.nsa.utils import (
     can_nsa_prefill_cp_round_robin_split,
+    cp_split_and_rebuild_data,
     compute_nsa_seqlens,
     get_cp_shared_kv_batch_plan,
     get_cp_shared_kv_local_out_cache_loc,
@@ -1883,11 +1884,41 @@ class NativeSparseAttnBackend(
                 assert current_kv_rows_for_reuse is not None
                 valid_current_rows = int(current_kv_rows_for_reuse)
                 if not compute_padding_current:
-                    current_k_nope = k[:valid_current_rows]
-                    current_k_rope = k_rope[:valid_current_rows]
-                    current_locs_for_reuse = forward_batch.out_cache_loc[
-                        :valid_current_rows
-                    ]
+                    current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc(
+                        forward_batch
+                    )
+                    if current_locs_for_reuse is None:
+                        raise RuntimeError(
+                            "[CP_SHARED_KV_FAIL_FAST][mla_current_reuse_locs] "
+                            "CP shared KV MLA current reuse requires local "
+                            "out_cache_loc in the prefill CP path."
+                        )
+                    local_current_rows = int(current_locs_for_reuse.numel())
+                    if int(k.shape[0]) == local_current_rows and int(
+                        k_rope.shape[0]
+                    ) == local_current_rows:
+                        current_k_nope = k
+                        current_k_rope = k_rope
+                    else:
+                        current_k_nope = cp_split_and_rebuild_data(
+                            forward_batch, k[:valid_current_rows].contiguous()
+                        )
+                        current_k_rope = cp_split_and_rebuild_data(
+                            forward_batch, k_rope[:valid_current_rows].contiguous()
+                        )
+                    if (
+                        int(current_k_nope.shape[0]) != local_current_rows
+                        or int(current_k_rope.shape[0]) != local_current_rows
+                    ):
+                        raise RuntimeError(
+                            "[CP_SHARED_KV_FAIL_FAST][mla_current_reuse_shape] "
+                            "CP shared KV MLA current rows do not match local "
+                            "out_cache_loc after CP split. "
+                            f"k_rows={int(current_k_nope.shape[0])} "
+                            f"k_rope_rows={int(current_k_rope.shape[0])} "
+                            f"local_locs={local_current_rows} "
+                            f"valid_current_rows={valid_current_rows}"
+                        )
                 assert current_k_nope is not None
                 assert current_k_rope is not None
                 assert current_locs_for_reuse is not None

python/sglang/srt/speculative/eagle_worker.py

@@ -1072,6 +1072,7 @@ class EAGLEWorker(TpModelWorker):
         probs = torch.softmax(logits_output.next_token_logits, dim=-1)
         draft_input.topk_p, draft_input.topk_index = fast_topk(probs, self.topk, dim=-1)
         draft_input.hidden_states = logits_output.hidden_states
+        draft_input.cp_local_hidden_states = False
 
     def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
         monkey_patch_torch_reductions()