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
This commit is contained in:
laoyao0822
2026-05-27 05:46:34 +08:00
parent 71c4f66968
commit 367dff06f3
9 changed files with 448 additions and 42 deletions

View File

@@ -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(

View File

@@ -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(

View File

@@ -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."""

View File

@@ -1261,9 +1261,15 @@ class HiCacheController:
if not backup_only:
raise ValueError("Other eviction policies are not supported yet.")
freed = self.evict_host(metadata.host_indices, backup_only=backup_only)
draft_host_indices = getattr(metadata, "draft_host_indices", None)
if self.has_draft_hicache and draft_host_indices is not None:
if self.has_draft_hicache and draft_host_indices is None:
raise RuntimeError(
"CP HiCache draft KV host evict requested but node metadata "
"does not contain draft_host_indices"
)
freed = self.evict_host(metadata.host_indices, backup_only=backup_only)
if self.has_draft_hicache:
self.draft_mem_pool_host.free(draft_host_indices)
return freed

View File

@@ -944,11 +944,20 @@ class HiRadixCache(RadixCache):
)
return True
def _node_host_write_pending(self, node: TreeNode) -> bool:
return getattr(node, "id", None) in getattr(self, "ongoing_write_through", {})
def _node_host_write_ready(self, node: TreeNode) -> bool:
return not self._node_host_write_pending(node)
def _node_backuped(self, node: TreeNode) -> bool:
if self._uses_cp_hicache:
if node.host_len <= 0:
return False
return self._node_has_required_draft_hicache(node)
return (
self._node_has_required_draft_hicache(node)
and self._node_host_write_ready(node)
)
return node.host_value is not None
def _node_host_len(self, node: TreeNode) -> int:
@@ -1077,6 +1086,9 @@ class HiRadixCache(RadixCache):
return
node.hit_count += 1
if self._uses_cp_hicache and self._node_host_write_pending(node):
return
if not self._node_backuped(node):
if node.hit_count >= self.write_through_threshold:
logger.info(