Preserve draft NSA state during CP disaggregated transfer

CP shared KV already registered the draft model main KV buffer with the
prefill/decode Mooncake managers, but NSA draft state buffers were not part
of the state registration set. HiCache/cache-hit traffic could then transfer
pages from the draft pool without transferring the matching draft
index/scale state, which is a plausible cause of the EAGLE/MTP accept-length
collapse after cache hits.

This appends compatible draft NSA state buffers to the existing state
transfer registration on both prefill and decode, and extends transfer-side
diagnostics so source/destination state-buffer counts are visible. The
mismatch guard degrades to the common prefix of registered state buffers
instead of crashing if a rolling deployment exposes asymmetric registration.

Constraint: Scope is intentionally limited to target_state_type=nsa and draft_state_type=nsa.
Rejected: Treat draft main KV transfer as sufficient | NSA attention also needs draft index/scale state for transferred pages.
Rejected: Add Mamba/SWA draft-state semantics now | those state layouts need separate correctness analysis.
Confidence: medium
Scope-risk: moderate
Directive: Do not remove the draft_state_buffer_start/count fields without checking Mooncake source/destination registration symmetry.
Tested: PYTHONDONTWRITEBYTECODE=1 python3 -m py_compile python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/decode.py python/sglang/srt/disaggregation/mooncake/conn.py
Tested: git diff --check
Tested: Remote prefill log showed registered_state_bufs=79 and maybe_send_extra_state src_state_bufs=79 dst_state_bufs=79 with no state-buffer mismatch.
Not-tested: Full accept-length recovery; latest remote run hit an unrelated prefill KV allocator idle-check leak after transfer registration succeeded.

python/sglang/srt/disaggregation/decode.py

@@ -113,6 +113,26 @@ def _pool_summary(pool) -> str:
     return " ".join(parts)
 
 
+def _pool_state_type(pool) -> str:
+    if pool is None or not hasattr(pool, "get_state_buf_infos"):
+        return "none"
+    if isinstance(pool, SWAKVPool):
+        return "swa"
+    if isinstance(pool, HybridLinearKVPool):
+        return "mamba"
+    if isinstance(pool, NSATokenToKVPool):
+        return "nsa"
+    return "unknown"
+
+
+def _state_buf_infos(pool):
+    state_type = _pool_state_type(pool)
+    if state_type == "none":
+        return state_type, [], [], []
+    state_data_ptrs, state_data_lens, state_item_lens = pool.get_state_buf_infos()
+    return state_type, state_data_ptrs, state_data_lens, state_item_lens
+
+
 def _kv_locs_to_page_indices_cpu(
     kv_locs: torch.Tensor,
     page_size: int,
@@ -426,6 +446,68 @@ class DecodePreallocQueue:
             kv_args.state_item_lens = []
             kv_args.state_type = "none"
 
+        draft_state_type = "none"
+        draft_state_data_ptrs = []
+        draft_state_data_lens = []
+        draft_state_item_lens = []
+        if self.draft_token_to_kv_pool is not None:
+            (
+                draft_state_type,
+                draft_state_data_ptrs,
+                draft_state_data_lens,
+                draft_state_item_lens,
+            ) = _state_buf_infos(self.draft_token_to_kv_pool)
+
+        kv_args.draft_state_type = draft_state_type
+        kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs)
+        kv_args.draft_state_buffer_count = 0
+        if draft_state_data_ptrs:
+            if kv_args.state_type == "nsa" and draft_state_type == "nsa":
+                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
+                kv_args.draft_state_buffer_count = len(draft_state_data_ptrs)
+            else:
+                _cp_draft_shared_kv_debug(
+                    "decode_draft_state_skipped cp_rank=%s target_state_type=%s "
+                    "draft_state_type=%s draft_state_bufs=%s "
+                    "reason=unsupported_state_type",
+                    self.tp_rank,
+                    kv_args.state_type,
+                    draft_state_type,
+                    len(draft_state_data_ptrs),
+                )
+
+        if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
+            _cp_draft_shared_kv_debug(
+                "decode_state_manager cp_rank=%s target_state_type=%s "
+                "draft_state_type=%s draft_state_bufs=%s draft_state_lens=%s "
+                "draft_state_item_lens=%s draft_state_start=%s registered_state_bufs=%s "
+                "registered_state_lens=%s registered_state_item_lens=%s",
+                self.tp_rank,
+                kv_args.state_type,
+                kv_args.draft_state_type,
+                kv_args.draft_state_buffer_count,
+                _seq_summary(draft_state_data_lens),
+                _seq_summary(draft_state_item_lens),
+                kv_args.draft_state_buffer_start,
+                len(kv_args.state_data_ptrs),
+                _seq_summary(kv_args.state_data_lens),
+                _seq_summary(kv_args.state_item_lens),
+            )
+            if kv_args.draft_state_buffer_count > 0:
+                _cp_draft_shared_kv_debug(
+                    "decode_draft_state_registered cp_rank=%s "
+                    "draft_state_type=%s draft_state_bufs=%s draft_state_start=%s "
+                    "registered_state_type=%s registered_state_bufs=%s",
+                    self.tp_rank,
+                    kv_args.draft_state_type,
+                    kv_args.draft_state_buffer_count,
+                    kv_args.draft_state_buffer_start,
+                    kv_args.state_type,
+                    len(kv_args.state_data_ptrs),
+                )
+
         kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
         kv_args.gpu_id = self.scheduler.gpu_id
         kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)

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

@@ -259,7 +259,9 @@ class MooncakeKVManager(CommonKVManager):
         if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens:
             _cp_draft_shared_kv_debug(
                 "register_buffers mode=%s cp_rank=%s total_kv_bufs=%s "
-                "draft_start=%s draft_count=%s kv_lens=%s kv_item_lens=%s",
+                "draft_start=%s draft_count=%s kv_lens=%s kv_item_lens=%s "
+                "state_type=%s state_bufs=%s state_lens=%s state_item_lens=%s "
+                "draft_state_type=%s draft_state_bufs=%s",
                 self.disaggregation_mode,
                 self.attn_cp_rank,
                 len(self.kv_args.kv_data_ptrs),
@@ -267,6 +269,12 @@ class MooncakeKVManager(CommonKVManager):
                 getattr(self.kv_args, "draft_kv_buffer_count", None),
                 _np_summary(self.kv_args.kv_data_lens),
                 _np_summary(self.kv_args.kv_item_lens),
+                getattr(self.kv_args, "state_type", None),
+                len(getattr(self.kv_args, "state_data_ptrs", []) or []),
+                _np_summary(getattr(self.kv_args, "state_data_lens", [])),
+                _np_summary(getattr(self.kv_args, "state_item_lens", [])),
+                getattr(self.kv_args, "draft_state_type", None),
+                getattr(self.kv_args, "draft_state_buffer_count", None),
             )
             self.engine.batch_register(
                 self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
@@ -645,6 +653,33 @@ class MooncakeKVManager(CommonKVManager):
     ):
         """Send state or extra pool data with type-specific handling."""
         state_type = getattr(self.kv_args, "state_type", "none")
+        _cp_draft_shared_kv_debug(
+            "maybe_send_extra_state cp_rank=%s room=%s session=%s state_type=%s "
+            "src_state_bufs=%s dst_state_bufs=%s src_state_lens=%s src_state_item_lens=%s "
+            "prefill_state_indices=%s dst_state_indices=%s draft_state_type=%s "
+            "draft_state_bufs=%s target_registration_state_bufs=%s",
+            self.attn_cp_rank,
+            req.room,
+            req.mooncake_session_id,
+            state_type,
+            len(getattr(self.kv_args, "state_data_ptrs", []) or []),
+            len(dst_state_data_ptrs or []),
+            _np_summary(getattr(self.kv_args, "state_data_lens", [])),
+            _np_summary(getattr(self.kv_args, "state_item_lens", [])),
+            _np_summary(prefill_state_indices),
+            _np_summary(
+                dst_state_indices
+                if dst_state_indices is not None
+                else getattr(req, "dst_state_indices", [])
+            ),
+            getattr(self.kv_args, "draft_state_type", None),
+            getattr(self.kv_args, "draft_state_buffer_count", None),
+            (
+                len(target_rank_registration_info.dst_state_data_ptrs)
+                if target_rank_registration_info is not None
+                else None
+            ),
+        )
 
         if state_type == "mamba":
             # Check if we need slice transfer for different TP sizes
@@ -689,13 +724,34 @@ class MooncakeKVManager(CommonKVManager):
                 prefill_state_indices = prefill_state_indices[
                     : len(effective_dst_state_indices)
                 ]
+            src_state_data_ptrs = self.kv_args.state_data_ptrs
+            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):
+                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 "
+                    "state_type=%s draft_state_type=%s draft_state_bufs=%s room=%s "
+                    "session=%s; transferring first %s buffers only",
+                    len(src_state_data_ptrs),
+                    len(dst_state_ptrs),
+                    state_type,
+                    getattr(self.kv_args, "draft_state_type", None),
+                    getattr(self.kv_args, "draft_state_buffer_count", None),
+                    req.room,
+                    req.mooncake_session_id,
+                    transfer_buf_count,
+                )
+                src_state_data_ptrs = src_state_data_ptrs[:transfer_buf_count]
+                dst_state_ptrs = dst_state_ptrs[:transfer_buf_count]
+                state_item_lens = state_item_lens[:transfer_buf_count]
             # Reuse _send_kvcache_generic interface to send extra pool data
             prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32)
             return self._send_kvcache_generic(
                 mooncake_session_id=req.mooncake_session_id,
-                src_data_ptrs=self.kv_args.state_data_ptrs,
-                dst_data_ptrs=dst_state_data_ptrs,
-                item_lens=self.kv_args.state_item_lens,
+                src_data_ptrs=src_state_data_ptrs,
+                dst_data_ptrs=dst_state_ptrs,
+                item_lens=state_item_lens,
                 prefill_data_indices=prefill_state_indices,
                 dst_data_indices=effective_dst_state_indices,
                 executor=executor,
@@ -1454,6 +1510,24 @@ class MooncakeKVReceiver(CommonKVReceiver):
             dst_tp_rank = str(tp_rank).encode("ascii")
             dst_attn_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii")
             dst_kv_item_len = str(kv_item_len).encode("ascii")
+            _cp_draft_shared_kv_debug(
+                "decode_register_kv_args cp_rank=%s room=%s session=%s "
+                "kv_bufs=%s aux_bufs=%s state_type=%s state_bufs=%s "
+                "state_item_lens=%s draft_start=%s draft_count=%s "
+                "draft_state_type=%s draft_state_bufs=%s",
+                self.kv_mgr.attn_cp_rank,
+                self.bootstrap_room,
+                self.session_id,
+                len(self.kv_mgr.kv_args.kv_data_ptrs),
+                len(self.kv_mgr.kv_args.aux_data_ptrs),
+                getattr(self.kv_mgr.kv_args, "state_type", None),
+                len(self.kv_mgr.kv_args.state_data_ptrs),
+                _np_summary(self.kv_mgr.kv_args.state_item_lens),
+                getattr(self.kv_mgr.kv_args, "draft_kv_buffer_start", None),
+                getattr(self.kv_mgr.kv_args, "draft_kv_buffer_count", None),
+                getattr(self.kv_mgr.kv_args, "draft_state_type", None),
+                getattr(self.kv_mgr.kv_args, "draft_state_buffer_count", None),
+            )
 
             sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
             with lock:

python/sglang/srt/disaggregation/prefill.py

@@ -98,6 +98,26 @@ def _pool_summary(pool) -> str:
     return " ".join(parts)
 
 
+def _pool_state_type(pool) -> str:
+    if pool is None or not hasattr(pool, "get_state_buf_infos"):
+        return "none"
+    if isinstance(pool, SWAKVPool):
+        return "swa"
+    if isinstance(pool, HybridLinearKVPool):
+        return "mamba"
+    if isinstance(pool, NSATokenToKVPool):
+        return "nsa"
+    return "unknown"
+
+
+def _state_buf_infos(pool):
+    state_type = _pool_state_type(pool)
+    if state_type == "none":
+        return state_type, [], [], []
+    state_data_ptrs, state_data_lens, state_item_lens = pool.get_state_buf_infos()
+    return state_type, state_data_ptrs, state_data_lens, state_item_lens
+
+
 def _kv_locs_to_page_indices_cpu(
     kv_locs: torch.Tensor,
     page_size: int,
@@ -267,6 +287,68 @@ class PrefillBootstrapQueue:
             kv_args.state_item_lens = []
             kv_args.state_type = "none"
 
+        draft_state_type = "none"
+        draft_state_data_ptrs = []
+        draft_state_data_lens = []
+        draft_state_item_lens = []
+        if self.draft_token_to_kv_pool is not None:
+            (
+                draft_state_type,
+                draft_state_data_ptrs,
+                draft_state_data_lens,
+                draft_state_item_lens,
+            ) = _state_buf_infos(self.draft_token_to_kv_pool)
+
+        kv_args.draft_state_type = draft_state_type
+        kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs)
+        kv_args.draft_state_buffer_count = 0
+        if draft_state_data_ptrs:
+            if kv_args.state_type == "nsa" and draft_state_type == "nsa":
+                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
+                kv_args.draft_state_buffer_count = len(draft_state_data_ptrs)
+            else:
+                _cp_draft_shared_kv_debug(
+                    "prefill_draft_state_skipped cp_rank=%s target_state_type=%s "
+                    "draft_state_type=%s draft_state_bufs=%s "
+                    "reason=unsupported_state_type",
+                    self.tp_rank,
+                    kv_args.state_type,
+                    draft_state_type,
+                    len(draft_state_data_ptrs),
+                )
+
+        if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
+            _cp_draft_shared_kv_debug(
+                "prefill_state_manager cp_rank=%s target_state_type=%s "
+                "draft_state_type=%s draft_state_bufs=%s draft_state_lens=%s "
+                "draft_state_item_lens=%s draft_state_start=%s registered_state_bufs=%s "
+                "registered_state_lens=%s registered_state_item_lens=%s",
+                self.tp_rank,
+                kv_args.state_type,
+                kv_args.draft_state_type,
+                kv_args.draft_state_buffer_count,
+                _seq_summary(draft_state_data_lens),
+                _seq_summary(draft_state_item_lens),
+                kv_args.draft_state_buffer_start,
+                len(kv_args.state_data_ptrs),
+                _seq_summary(kv_args.state_data_lens),
+                _seq_summary(kv_args.state_item_lens),
+            )
+            if kv_args.draft_state_buffer_count > 0:
+                _cp_draft_shared_kv_debug(
+                    "prefill_draft_state_registered cp_rank=%s "
+                    "draft_state_type=%s draft_state_bufs=%s draft_state_start=%s "
+                    "registered_state_type=%s registered_state_bufs=%s",
+                    self.tp_rank,
+                    kv_args.draft_state_type,
+                    kv_args.draft_state_buffer_count,
+                    kv_args.draft_state_buffer_start,
+                    kv_args.state_type,
+                    len(kv_args.state_data_ptrs),
+                )
+
         kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
         kv_manager = kv_manager_class(
             kv_args,