Keep CP HiCache draft KV invisible until joint readiness

CP HiCache now treats draft KV as a strict target-owned payload through pending write visibility, host eviction, and state-buffer registration. Host metadata created before async D2H ack is no longer request-visible, so match_prefix cannot select an in-flight host node. Draft host eviction now fails before target cleanup when draft metadata is missing, and prefill/decode share one helper for draft NSA state buffers so shared-KV mode cannot silently skip mismatched draft state.

Constraint: CP shared KV + HiCache + EAGLE/MTP must not expose target-only host hits or skipped draft state as valid cache hits

Rejected: Rely on event-loop ordering and lock_ref to hide in-flight writes | match_prefix does not consult lock_ref and can observe host_len/cp_hicache directly

Rejected: Keep draft state mismatch as debug-only skip | it can poison speculative acceptance while looking like a successful cache hit

Confidence: high

Scope-risk: moderate

Directive: Do not reintroduce silent draft/target fallback in CP shared-KV HiCache paths; malformed strong-sync metadata should fail fast

Tested: python -m py_compile targeted modified files

Tested: remote g0034 container pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py -q (91 passed)

Not-tested: Full CP shared KV + HiCache + EAGLE/MTP ETE server run after this commit

python/sglang/srt/disaggregation/decode.py

@@ -41,6 +41,7 @@ from sglang.srt.disaggregation.utils import (
     MetadataBuffers,
     ReqToMetadataIdxAllocator,
     TransferBackend,
+    append_cp_draft_state_buffers,
     get_kv_class,
     is_mla_backend,
     poll_and_all_reduce,
@@ -458,25 +459,25 @@ class DecodePreallocQueue:
                 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),
-                )
+        draft_state_registered = append_cp_draft_state_buffers(
+            kv_args,
+            draft_state_type,
+            draft_state_data_ptrs,
+            draft_state_data_lens,
+            draft_state_item_lens,
+            role="decode",
+            cp_rank=self.tp_rank,
+        )
+        if draft_state_data_ptrs and not draft_state_registered:
+            _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(

python/sglang/srt/disaggregation/prefill.py

@@ -36,6 +36,7 @@ from sglang.srt.disaggregation.utils import (
     MetadataBuffers,
     ReqToMetadataIdxAllocator,
     TransferBackend,
+    append_cp_draft_state_buffers,
     get_kv_class,
     is_mla_backend,
     kv_to_page_num,
@@ -299,25 +300,25 @@ class PrefillBootstrapQueue:
                 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),
-                )
+        draft_state_registered = append_cp_draft_state_buffers(
+            kv_args,
+            draft_state_type,
+            draft_state_data_ptrs,
+            draft_state_data_lens,
+            draft_state_item_lens,
+            role="prefill",
+            cp_rank=self.tp_rank,
+        )
+        if draft_state_data_ptrs and not draft_state_registered:
+            _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(

python/sglang/srt/disaggregation/utils.py

@@ -5,7 +5,7 @@ import random
 from collections import deque
 from contextlib import nullcontext
 from enum import Enum
-from typing import TYPE_CHECKING, Literal, Optional, Tuple, Type, overload
+from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, Type, overload
 
 import numpy as np
 import torch
@@ -85,6 +85,50 @@ def poll_and_all_reduce_attn_cp_tp_group(
 #########################
 
 
+def append_cp_draft_state_buffers(
+    kv_args: Any,
+    draft_state_type: str,
+    draft_state_data_ptrs,
+    draft_state_data_lens,
+    draft_state_item_lens,
+    *,
+    role: str,
+    cp_rank: int,
+) -> bool:
+    """Append draft NSA state buffers as payload attached to target state buffers.
+
+    Under CP draft shared KV, draft state is not allowed to silently diverge
+    from target state. If draft exposes state buffers, target and draft must both
+    be NSA state so the transfer still has one target-driven state stream.
+    """
+
+    target_state_type = getattr(kv_args, "state_type", "none")
+    draft_state_bufs = len(draft_state_data_ptrs)
+    can_append = target_state_type == "nsa" and draft_state_type == "nsa"
+
+    if draft_state_bufs and not can_append:
+        if envs.SGLANG_CP_DRAFT_SHARED_KV.get():
+            raise RuntimeError(
+                "CP draft shared KV state mismatch during "
+                f"{role}: cp_rank={cp_rank} target_state_type={target_state_type} "
+                f"draft_state_type={draft_state_type} "
+                f"draft_state_bufs={draft_state_bufs}"
+            )
+
+    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 not draft_state_bufs or not can_append:
+        return False
+
+    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 = draft_state_bufs
+    return True
+
+
 class ReqToMetadataIdxAllocator:
     """A memory pool that maps a request to its first output token location."""