Compact skipped NSA index-cache state safely

Index skip reduces the number of target layers that own NSA index state,
but PD transfer and HiCache still assumed dense full-layer state buffers.
This change carries explicit state layer IDs through prefill/decode
registration, compacts device and host index buffers to active layers,
and maps logical layer IDs to compact slots on transfer paths.

The PD side fails fast when prefill/decode disagree on NSA state layer
identity instead of silently truncating or copying mismatched buffers.
Host direct tests now use the same CPU-index descriptor contract required
by the TAI cudaMemcpyBatchAsync path, and host registered memory is
unregistered on tensor finalization to avoid stale cudaHostRegister state
across CUDA tests.

Constraint: CP shared-KV with index_topk skip must keep target/draft state identity explicit before compacting buffers
Constraint: Direct HiCache TAI transfer rejects CUDA indices to avoid hidden D2H copies on the control path
Rejected: Keep full-layer L1/L2 index buffers | wastes the memory/bandwidth that index skip is meant to save
Rejected: Infer state buffer order by count only | can silently corrupt cache when active layer sets differ
Confidence: high
Scope-risk: moderate
Directive: Do not compact or reorder NSA state buffers without carrying logical layer IDs through PD registration and validating both sides
Tested: Remote container py_compile for touched runtime files
Tested: Remote container pytest: test_nsa_pool_host_unit.py, test_model_runner_kv_cache_mixin.py, test_cp_shared_kv_transfer_mapping.py, test_pd_state_layer_ids.py, test_cp_per_layer_transfer.py, test_cp_shared_kv_runtime.py -> 200 passed, 2 subtests passed
Not-tested: Full ETE GSM8K/replay after compacted P3-P6 changes
Co-authored-by: OmX <omx@oh-my-codex.dev>

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

@@ -175,6 +175,7 @@ class KVArgsRegisterInfo:
     # for mamba state different tp slice transfer
     dst_state_item_lens: list[int]
     dst_state_dim_per_tensor: list[int]
+    dst_state_layer_ids: list[int]
 
     @classmethod
     def from_zmq(cls, msg: List[bytes]):
@@ -199,6 +200,11 @@ class KVArgsRegisterInfo:
                 if len(msg) > 11 and len(msg[11]) > 0
                 else []
             ),
+            dst_state_layer_ids=(
+                list(struct.unpack(f"{len(msg[12])//4}i", msg[12]))
+                if len(msg) > 12 and len(msg[12]) > 0
+                else []
+            ),
         )
 
 
@@ -957,6 +963,22 @@ class MooncakeKVManager(CommonKVManager):
                 raise RuntimeError(
                     f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet."
                 )
+            src_state_layer_ids = list(
+                getattr(self.kv_args, "state_layer_ids", []) or []
+            )
+            dst_state_layer_ids = (
+                list(getattr(target_rank_registration_info, "dst_state_layer_ids", []) or [])
+                if target_rank_registration_info is not None
+                else []
+            )
+            if src_state_layer_ids or dst_state_layer_ids:
+                if src_state_layer_ids != dst_state_layer_ids:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][state_layer_ids] "
+                        f"prefill={src_state_layer_ids} decode={dst_state_layer_ids} "
+                        f"state_type={state_type} room={req.room} "
+                        f"session={req.mooncake_session_id}"
+                    )
             effective_dst_state_indices = (
                 np.asarray(dst_state_indices, dtype=np.int32)
                 if dst_state_indices is not None
@@ -984,6 +1006,14 @@ class MooncakeKVManager(CommonKVManager):
             dst_state_ptrs = dst_state_data_ptrs
             state_item_lens = self.kv_args.state_item_lens
             if len(src_state_data_ptrs) != len(dst_state_ptrs):
+                if src_state_layer_ids or dst_state_layer_ids:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][state_buffer_count] "
+                        f"src={len(src_state_data_ptrs)} dst={len(dst_state_ptrs)} "
+                        f"src_layers={src_state_layer_ids} "
+                        f"dst_layers={dst_state_layer_ids} state_type={state_type} "
+                        f"room={req.room} session={req.mooncake_session_id}"
+                    )
                 transfer_buf_count = min(len(src_state_data_ptrs), len(dst_state_ptrs))
                 logger.warning(
                     "State buffer count mismatch during PD transfer: src=%s dst=%s "
@@ -1839,6 +1869,10 @@ class MooncakeKVReceiver(CommonKVReceiver):
             packed_state_dim_per_tensor = b"".join(
                 struct.pack("I", dim) for dim in state_dim_per_tensor
             )
+            packed_state_layer_ids = b"".join(
+                struct.pack("i", int(layer_id))
+                for layer_id in getattr(self.kv_mgr.kv_args, "state_layer_ids", [])
+            )
             # Note(shangming): No need to add pp rank here since decode pp size should be equal to prefill pp size or 1
             tp_rank = self.kv_mgr.kv_args.engine_rank
             kv_item_len = self.kv_mgr.kv_args.kv_item_lens[0]
@@ -1880,6 +1914,7 @@ class MooncakeKVReceiver(CommonKVReceiver):
                         dst_kv_item_len,
                         packed_state_item_lens,
                         packed_state_dim_per_tensor,
+                        packed_state_layer_ids,
                     ]
                 )