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/base/conn.py

@@ -31,6 +31,7 @@ class KVArgs:
     state_data_ptrs: List[int]
     state_data_lens: List[int]
     state_item_lens: List[int]
+    state_layer_ids: List[int]
     state_type: str  # "none", "mamba", "swa"
     # for mamba state different tp slice transfer
     state_dim_per_tensor: List[int]  # dimension to slice for each state tensor

python/sglang/srt/disaggregation/decode.py

@@ -136,6 +136,13 @@ def _state_buf_infos(pool):
     return state_type, state_data_ptrs, state_data_lens, state_item_lens
 
 
+def _state_layer_ids(pool):
+    get_state_layer_ids = getattr(pool, "get_state_layer_ids", None)
+    if get_state_layer_ids is None:
+        return []
+    return list(get_state_layer_ids())
+
+
 def _kv_locs_to_page_indices_cpu(
     kv_locs: torch.Tensor,
     page_size: int,
@@ -429,6 +436,7 @@ class DecodePreallocQueue:
             kv_args.state_data_ptrs = state_data_ptrs
             kv_args.state_data_lens = state_data_lens
             kv_args.state_item_lens = state_item_lens
+            kv_args.state_layer_ids = _state_layer_ids(self.token_to_kv_pool)
 
             if isinstance(self.token_to_kv_pool, SWAKVPool):
                 kv_args.state_type = "swa"
@@ -447,6 +455,7 @@ class DecodePreallocQueue:
             kv_args.state_data_ptrs = []
             kv_args.state_data_lens = []
             kv_args.state_item_lens = []
+            kv_args.state_layer_ids = []
             kv_args.state_type = "none"
 
         draft_state_type = "none"

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,
                     ]
                 )
 

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

@@ -84,6 +84,7 @@ class KVArgsRegisterInfo:
     decode_tp_size: int
     decode_tp_rank: int
     dst_kv_item_len: int
+    dst_state_layer_ids: list[int]
 
     @classmethod
     def from_zmq(cls, msg: List[bytes]):
@@ -106,6 +107,11 @@ class KVArgsRegisterInfo:
             decode_tp_size=int(msg[9].decode("ascii")),
             decode_tp_rank=int(msg[10].decode("ascii")),
             dst_kv_item_len=int(msg[11].decode("ascii")),
+            dst_state_layer_ids=(
+                list(struct.unpack(f"{len(msg[12]) // 4}i", msg[12]))
+                if len(msg) > 12 and msg[12] != b""
+                else []
+            ),
         )
 
 
@@ -680,6 +686,7 @@ class NixlKVManager(CommonKVManager):
         dst_gpu_id: int,
         notif: str,
         decode_tp_size: int,
+        dst_state_layer_ids: Optional[List[int]] = None,
     ):
         """Send state or extra pool data with type-specific handling."""
         state_type = getattr(self.kv_args, "state_type", "none")
@@ -702,6 +709,26 @@ class NixlKVManager(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(dst_state_layer_ids or [])
+            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} peer={peer_name} notif={notif}"
+                    )
+                if len(self.kv_args.state_data_ptrs) != len(dst_state_data_ptrs):
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][state_buffer_count] "
+                        f"src={len(self.kv_args.state_data_ptrs)} "
+                        f"dst={len(dst_state_data_ptrs)} "
+                        f"src_layers={src_state_layer_ids} "
+                        f"dst_layers={dst_state_layer_ids} "
+                        f"state_type={state_type} peer={peer_name} notif={notif}"
+                    )
             if len(prefill_state_indices) != len(dst_state_indices):
                 raise RuntimeError(
                     f"State index length mismatch: prefill={len(prefill_state_indices)}, "
@@ -791,6 +818,7 @@ class NixlKVManager(CommonKVManager):
                         dst_info.gpu_id,
                         f"{req.room}_state_{self.kv_args.pp_rank}",
                         decode_tp_size,
+                        dst_info.dst_state_layer_ids,
                     )
                     if state_xfer_handle is not None:
                         handles.append(state_xfer_handle)
@@ -1069,6 +1097,10 @@ class NixlKVReceiver(CommonKVReceiver):
             packed_state_data_ptrs = b"".join(
                 struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.state_data_ptrs
             )
+            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", [])
+            )
 
             with lock:
                 sock.send_multipart(
@@ -1086,6 +1118,7 @@ class NixlKVReceiver(CommonKVReceiver):
                         str(self.kv_mgr.attn_tp_size).encode("ascii"),
                         str(self.kv_mgr.kv_args.engine_rank).encode("ascii"),
                         str(self.kv_mgr.kv_args.kv_item_lens[0]).encode("ascii"),
+                        packed_state_layer_ids,
                     ]
                 )
 

python/sglang/srt/disaggregation/prefill.py

@@ -191,6 +191,13 @@ def _state_buf_infos(pool):
     return state_type, state_data_ptrs, state_data_lens, state_item_lens
 
 
+def _state_layer_ids(pool):
+    get_state_layer_ids = getattr(pool, "get_state_layer_ids", None)
+    if get_state_layer_ids is None:
+        return []
+    return list(get_state_layer_ids())
+
+
 def _kv_locs_to_page_indices_cpu(
     kv_locs: torch.Tensor,
     page_size: int,
@@ -376,6 +383,7 @@ class PrefillBootstrapQueue:
             kv_args.state_data_ptrs = state_data_ptrs
             kv_args.state_data_lens = state_data_lens
             kv_args.state_item_lens = state_item_lens
+            kv_args.state_layer_ids = _state_layer_ids(self.token_to_kv_pool)
 
             if isinstance(self.token_to_kv_pool, SWAKVPool):
                 kv_args.state_type = "swa"
@@ -394,6 +402,7 @@ class PrefillBootstrapQueue:
             kv_args.state_data_ptrs = []
             kv_args.state_data_lens = []
             kv_args.state_item_lens = []
+            kv_args.state_layer_ids = []
             kv_args.state_type = "none"
 
         draft_state_type = "none"

python/sglang/srt/disaggregation/utils.py

@@ -259,6 +259,8 @@ def append_cp_draft_state_buffers(
     kv_args.state_data_ptrs += draft_state_data_ptrs
     kv_args.state_data_lens += draft_state_data_lens
     kv_args.state_item_lens += draft_state_item_lens
+    if hasattr(kv_args, "state_layer_ids"):
+        kv_args.state_layer_ids += [-(i + 1) for i in range(draft_state_bufs)]
     kv_args.draft_state_buffer_count = draft_state_bufs
     return True
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -1913,6 +1913,11 @@ class NSATokenToKVPool(MLATokenToKVPool):
             if self.custom_mem_pool
             else nullcontext()
         ):
+            index_buffer_layer_num = (
+                len(self.index_active_layer_ids)
+                if self.index_compact_layers
+                else layer_num
+            )
             self.index_k_with_scale_buffer = [
                 torch.zeros(
                     # Layout:
@@ -1931,7 +1936,7 @@ class NSATokenToKVPool(MLATokenToKVPool):
                     dtype=self.index_k_with_scale_buffer_dtype,
                     device=device,
                 )
-                for _ in range(layer_num)
+                for _ in range(index_buffer_layer_num)
             ]
         self._finalize_allocation_log(size)
 
@@ -2054,17 +2059,24 @@ class NSATokenToKVPool(MLATokenToKVPool):
         )
 
     def get_state_buf_infos(self):
+        slots = [
+            self.get_index_layer_slot(layer_id)
+            for layer_id in self.index_active_layer_ids
+        ]
         data_ptrs = [
-            self.index_k_with_scale_buffer[i].data_ptr() for i in range(self.layer_num)
+            self.index_k_with_scale_buffer[slot].data_ptr() for slot in slots
         ]
         data_lens = [
-            self.index_k_with_scale_buffer[i].nbytes for i in range(self.layer_num)
+            self.index_k_with_scale_buffer[slot].nbytes for slot in slots
         ]
         item_lens = [
-            self.index_k_with_scale_buffer[i][0].nbytes for i in range(self.layer_num)
+            self.index_k_with_scale_buffer[slot][0].nbytes for slot in slots
         ]
         return data_ptrs, data_lens, item_lens
 
+    def get_state_layer_ids(self):
+        return list(self.index_active_layer_ids)
+
     def get_kv_size_bytes(self):
         kv_size_bytes = super().get_kv_size_bytes()
         for index_k_cache in self.index_k_with_scale_buffer:

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -3,6 +3,7 @@ import bisect
 import heapq
 import logging
 import threading
+import weakref
 from collections import defaultdict
 from functools import lru_cache, wraps
 from typing import Optional
@@ -196,12 +197,38 @@ def alloc_with_host_register(
     """
     buffer = allocator.allocate(dims, dtype=dtype, device=device)
     if pin_memory:
-        torch.cuda.cudart().cudaHostRegister(
-            buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0
+        ptr = buffer.data_ptr()
+        _check_torch_cudart(
+            torch.cuda.cudart().cudaHostRegister(
+                ptr, buffer.numel() * buffer.element_size(), 0
+            ),
+            "cudaHostRegister",
         )
+        weakref.finalize(buffer, _cuda_host_unregister, ptr)
     return buffer
 
 
+def _check_torch_cudart(err, op_name: str) -> None:
+    if int(err) != 0:
+        try:
+            err_str = torch.cuda.cudart().cudaGetErrorString(err)
+        except Exception:
+            err_str = repr(err)
+        raise RuntimeError(f"{op_name} failed: {err_str}")
+
+
+def _cuda_host_unregister(ptr: int) -> None:
+    try:
+        _check_torch_cudart(
+            torch.cuda.cudart().cudaHostUnregister(ptr), "cudaHostUnregister"
+        )
+    except Exception:
+        # Best-effort cleanup for host tensors.  The owning process may already
+        # be tearing down CUDA state at exit; production host pools are
+        # long-lived, while unit tests rely on timely unregister when pools die.
+        pass
+
+
 def alloc_with_pin_memory(
     dims,
     dtype: torch.dtype,
@@ -1834,6 +1861,18 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
         self.index_head_dim = device_pool.index_head_dim
         self.indexer_quant_block_size = device_pool.quant_block_size
         self.indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype
+        self.index_active_layer_ids = tuple(
+            int(layer_id)
+            for layer_id in getattr(
+                device_pool,
+                "index_active_layer_ids",
+                range(device_pool.start_layer, device_pool.start_layer + device_pool.layer_num),
+            )
+        )
+        self.index_active_layer_num = len(self.index_active_layer_ids)
+        self.index_logical_to_slot = {
+            layer_id: slot for slot, layer_id in enumerate(self.index_active_layer_ids)
+        }
         self.indexer_size_per_token = (
             self.index_head_dim
             + self.index_head_dim // self.indexer_quant_block_size * 4
@@ -1853,7 +1892,9 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
         self.indexer_page_stride_size = (
             self.indexer_size_per_token * self.page_size * self.indexer_dtype.itemsize
         )
-        self.indexer_layout_dim = self.indexer_page_stride_size * self.layer_num
+        self.indexer_layout_dim = (
+            self.indexer_page_stride_size * self.index_active_layer_num
+        )
         self.indexer_page_num = (self.size + self.page_size + 1) // self.page_size
         self._init_indexer_buffers()
         logger.info(
@@ -1864,7 +1905,9 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
         base = super().get_size_per_token()
         return (
             base
-            + self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize
+            + self.indexer_size_per_token
+            * self.index_active_layer_num
+            * self.indexer_dtype.itemsize
         )
 
     def _init_indexer_buffers(self):
@@ -1883,10 +1926,11 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
                     pin_memory=self.pin_memory,
                     allocator=self.allocator,
                 )
-                for _ in range(self.layer_num)
+                for _ in range(self.index_active_layer_num)
             ]
             self.index_k_data_refs = [
-                self.index_k_with_scale_buffer[i] for i in range(self.layer_num)
+                self.index_k_with_scale_buffer[i]
+                for i in range(self.index_active_layer_num)
             ]
             self.index_k_data_ptrs = torch.tensor(
                 [x.data_ptr() for x in self.index_k_data_refs],
@@ -1897,7 +1941,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
             self.index_k_with_scale_buffer = alloc_func(
                 (
                     self.indexer_page_num,
-                    self.layer_num,
+                    self.index_active_layer_num,
                     1,
                     self.indexer_page_stride_size,
                 ),
@@ -1909,7 +1953,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
         elif self.layout == "layer_page_first":
             self.index_k_with_scale_buffer = alloc_func(
                 (
-                    self.layer_num,
+                    self.index_active_layer_num,
                     self.indexer_page_num,
                     1,
                     self.indexer_page_stride_size,
@@ -1920,7 +1964,8 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
                 allocator=self.allocator,
             )
             self.index_k_data_refs = [
-                self.index_k_with_scale_buffer[i] for i in range(self.layer_num)
+                self.index_k_with_scale_buffer[i]
+                for i in range(self.index_active_layer_num)
             ]
             self.index_k_data_ptrs = torch.tensor(
                 [x.data_ptr() for x in self.index_k_data_refs],
@@ -1958,7 +2003,21 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
         return int(layer_id - getattr(device_pool, "start_layer", 0))
 
     def _host_index_layer_slot(self, layer_id: int) -> int:
-        return int(layer_id - getattr(self, "start_layer", 0))
+        mapping = getattr(self, "index_logical_to_slot", None)
+        if mapping is None:
+            # Unit-test stubs and legacy in-memory host pools created before the
+            # compact-index-layer metadata existed still use dense logical slots.
+            # Real NSATokenToKVPoolHost instances always install the mapping in
+            # __init__, so production compact paths retain the fail-fast below.
+            return int(layer_id - getattr(self, "start_layer", 0))
+        try:
+            return int(mapping[int(layer_id)])
+        except KeyError as exc:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][host_index_cache_layer] "
+                f"inactive host index layer requested: layer_id={layer_id} "
+                f"active_layer_ids={list(self.index_active_layer_ids)}"
+            ) from exc
 
     def _active_index_layer_ids_for_transfer(self, device_pool):
         active_layer_ids = getattr(device_pool, "index_active_layer_ids", None)

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -107,7 +107,17 @@ class ModelRunnerKVCacheMixin:
                 element_size = torch._utils._element_size(
                     NSATokenToKVPool.index_k_with_scale_buffer_dtype
                 )
-                cell_size += indexer_size_per_token * num_layers * element_size
+                index_layer_plan = build_nsa_index_layer_plan(
+                    self.model_config.hf_config,
+                    self.start_layer,
+                    self.end_layer,
+                    is_nextn=self.is_draft_worker,
+                )
+                cell_size += (
+                    indexer_size_per_token
+                    * len(index_layer_plan.active_layer_ids)
+                    * element_size
+                )
         else:
             if self.model_config.is_hybrid_swa:
                 full_layers_num = len(self.model_config.full_attention_layer_ids)
@@ -515,6 +525,7 @@ class ModelRunnerKVCacheMixin:
                 end_layer=self.end_layer,
                 index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config),
                 index_active_layer_ids=index_layer_plan.active_layer_ids,
+                compact_index_layers=True,
             )
             if self.enable_hisparse:
                 from sglang.srt.mem_cache.sparsity import parse_hisparse_config