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:
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user