feat: add cache salt support to KV cache events (#30827)

Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
This commit is contained in:
jthomson04
2026-08-12 16:14:04 -07:00
committed by GitHub
parent 198b7e9240
commit 385903b0ac
43 changed files with 753 additions and 70 deletions
@@ -2104,6 +2104,8 @@ class MMReceiverBase(ABC):
if self.scheduler.metrics_reporter.enable_metrics if self.scheduler.metrics_reporter.enable_metrics
else None else None
), ),
extra_key=recv_req.extra_key,
cache_salt=recv_req.cache_salt,
http_worker_ipc=recv_req.http_worker_ipc, http_worker_ipc=recv_req.http_worker_ipc,
dllm_config=self.scheduler.dllm_config, dllm_config=self.scheduler.dllm_config,
) )
@@ -86,6 +86,12 @@ class StorageMedium(str, enum.Enum):
EXTERNAL = "EXTERNAL" # L4: shared / remote pool (e.g. Mooncake) EXTERNAL = "EXTERNAL" # L4: shared / remote pool (e.g. Mooncake)
class BlockStoredMetadata(msgspec.Struct, omit_defaults=True, gc=False):
"""Typed request metadata attached to a stored KV block."""
cache_salt: str
class OffloadedState: class OffloadedState:
""" """
OffloadedState represents the state of a KV cache block offloaded to the hicache. OffloadedState represents the state of a KV cache block offloaded to the hicache.
@@ -112,6 +118,16 @@ class BlockStored(KVCacheEvent):
medium: Optional[str] = None medium: Optional[str] = None
class BlockStoredWithMetadata(BlockStored, tag="BlockStored", kw_only=True):
"""BlockStored wire extension used only when typed metadata is present.
A separate struct keeps unsalted events at their legacy array length; an
optional field on BlockStored would still serialize a trailing null.
"""
metadata: BlockStoredMetadata
class BlockRemoved(KVCacheEvent): class BlockRemoved(KVCacheEvent):
block_hashes: list[int] block_hashes: list[int]
medium: Optional[str] = None medium: Optional[str] = None
@@ -122,6 +138,10 @@ class AllBlocksCleared(KVCacheEvent):
class KVEventBatch(EventBatch): class KVEventBatch(EventBatch):
# BlockStoredWithMetadata deliberately stays out of this tagged union.
# Existing typed consumers decode its shared "BlockStored" tag as the base
# type and ignore the trailing metadata; adding both types would give
# msgspec duplicate tags and make the union invalid.
events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]] events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]
+6
View File
@@ -394,6 +394,8 @@ class Engine(EngineScoreMixin, EngineBase):
session_params: Optional[Dict] = None, session_params: Optional[Dict] = None,
priority: Optional[int] = None, priority: Optional[int] = None,
session_id: Optional[str] = None, session_id: Optional[str] = None,
*,
cache_salt: Optional[Union[List[str], str]] = None,
) -> Union[Dict, Iterator[Dict]]: ) -> Union[Dict, Iterator[Dict]]:
""" """
The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`. The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`.
@@ -411,6 +413,7 @@ class Engine(EngineScoreMixin, EngineBase):
audio_data=audio_data, audio_data=audio_data,
video_data=video_data, video_data=video_data,
mm_hashes=mm_hashes, mm_hashes=mm_hashes,
cache_salt=cache_salt,
return_logprob=return_logprob, return_logprob=return_logprob,
logprob_start_len=logprob_start_len, logprob_start_len=logprob_start_len,
top_logprobs_num=top_logprobs_num, top_logprobs_num=top_logprobs_num,
@@ -500,6 +503,8 @@ class Engine(EngineScoreMixin, EngineBase):
session_params: Optional[Dict] = None, session_params: Optional[Dict] = None,
priority: Optional[int] = None, priority: Optional[int] = None,
session_id: Optional[str] = None, session_id: Optional[str] = None,
*,
cache_salt: Optional[Union[List[str], str]] = None,
) -> Union[Dict, AsyncIterator[Dict]]: ) -> Union[Dict, AsyncIterator[Dict]]:
""" """
The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`. The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`.
@@ -517,6 +522,7 @@ class Engine(EngineScoreMixin, EngineBase):
audio_data=audio_data, audio_data=audio_data,
video_data=video_data, video_data=video_data,
mm_hashes=mm_hashes, mm_hashes=mm_hashes,
cache_salt=cache_salt,
return_logprob=return_logprob, return_logprob=return_logprob,
logprob_start_len=logprob_start_len, logprob_start_len=logprob_start_len,
top_logprobs_num=top_logprobs_num, top_logprobs_num=top_logprobs_num,
@@ -386,7 +386,7 @@ class CompletionRequest(BaseModel):
# For request id # For request id
rid: Optional[Union[List[str], str]] = None rid: Optional[Union[List[str], str]] = None
# Extra key for classifying the request (e.g. cache_salt) # Extra key for caller-defined request classification
extra_key: Optional[Union[List[str], str]] = None extra_key: Optional[Union[List[str], str]] = None
# Cache salt for request caching # Cache salt for request caching
cache_salt: Optional[Union[List[str], str]] = None cache_salt: Optional[Union[List[str], str]] = None
@@ -848,7 +848,7 @@ class ChatCompletionRequest(BaseModel):
# For request id # For request id
rid: Optional[Union[List[str], str]] = None rid: Optional[Union[List[str], str]] = None
# Extra key for classifying the request (e.g. cache_salt) # Extra key for caller-defined request classification
extra_key: Optional[Union[List[str], str]] = None extra_key: Optional[Union[List[str], str]] = None
# Cache salt for request caching # Cache salt for request caching
cache_salt: Optional[Union[List[str], str]] = None cache_salt: Optional[Union[List[str], str]] = None
@@ -1536,7 +1536,7 @@ class ResponsesRequest(BaseModel):
priority: int = Field(default=0, description="Request priority") priority: int = Field(default=0, description="Request priority")
extra_key: Optional[str] = Field( extra_key: Optional[str] = Field(
default=None, default=None,
description="Extra key for classifying the request (e.g. cache_salt)", description="Extra key for caller-defined request classification",
) )
cache_salt: Optional[str] = Field( cache_salt: Optional[str] = Field(
default=None, description="Cache salt for request caching" default=None, description="Cache salt for request caching"
@@ -148,19 +148,6 @@ class OpenAIServingBase(ABC):
return f"{self._request_id_prefix()}{uuid.uuid4().hex}" return f"{self._request_id_prefix()}{uuid.uuid4().hex}"
def _compute_extra_key(self, request: OpenAIServingRequest) -> Optional[str]:
"""Compute the final extra_key by concatenating cache_salt and extra_key if both are provided."""
parts = []
for key in ["cache_salt", "extra_key"]:
value = getattr(request, key, None)
if value:
if not isinstance(value, str):
raise TypeError(
f"Value of {key} must be a string, but got {type(value).__name__}"
)
parts.append(value)
return "".join(parts) if parts else None
@abstractmethod @abstractmethod
def _convert_to_internal_request( def _convert_to_internal_request(
self, self,
@@ -1018,7 +1018,8 @@ class OpenAIServingChat(OpenAIServingBase):
routed_experts_start_len=request.routed_experts_start_len, routed_experts_start_len=request.routed_experts_start_len,
rid=request.rid, rid=request.rid,
session_id=request.session_id, session_id=request.session_id,
extra_key=self._compute_extra_key(request), extra_key=request.extra_key,
cache_salt=request.cache_salt,
require_reasoning=processed_messages.require_reasoning, require_reasoning=processed_messages.require_reasoning,
priority=request.priority, priority=request.priority,
routing_key=self.extract_routing_key(raw_request), routing_key=self.extract_routing_key(raw_request),
@@ -128,7 +128,8 @@ class OpenAIServingCompletion(OpenAIServingBase):
return_prompt_token_ids=request.return_token_ids, return_prompt_token_ids=request.return_token_ids,
rid=request.rid, rid=request.rid,
session_id=request.session_id, session_id=request.session_id,
extra_key=self._compute_extra_key(request), extra_key=request.extra_key,
cache_salt=request.cache_salt,
priority=request.priority, priority=request.priority,
routing_key=self.extract_routing_key(raw_request), routing_key=self.extract_routing_key(raw_request),
custom_labels=custom_labels, custom_labels=custom_labels,
@@ -454,7 +454,8 @@ class OpenAIServingResponses(OpenAIServingChat):
stream=request.stream, stream=request.stream,
rid=request.request_id, rid=request.request_id,
session_id=request.session_id, session_id=request.session_id,
extra_key=self._compute_extra_key(request), extra_key=request.extra_key,
cache_salt=request.cache_salt,
# background+stream streams on this connection, so don't detach. # background+stream streams on this connection, so don't detach.
background=request.background and not request.stream, background=request.background and not request.stream,
require_reasoning=require_reasoning, require_reasoning=require_reasoning,
@@ -2563,6 +2564,7 @@ class OpenAIServingResponses(OpenAIServingChat):
rid=request_id, rid=request_id,
session_id=adapted_request.session_id, session_id=adapted_request.session_id,
extra_key=adapted_request.extra_key, extra_key=adapted_request.extra_key,
cache_salt=adapted_request.cache_salt,
return_logprob=adapted_request.return_logprob, return_logprob=adapted_request.return_logprob,
logprob_start_len=adapted_request.logprob_start_len, logprob_start_len=adapted_request.logprob_start_len,
top_logprobs_num=adapted_request.top_logprobs_num, top_logprobs_num=adapted_request.top_logprobs_num,
+42 -3
View File
@@ -287,7 +287,7 @@ class GenerateReqInput:
# Priority for the request # Priority for the request
priority: Optional[int] = None priority: Optional[int] = None
# Extra cache key for classifying the request (e.g. cache_salt) # Extra cache key for caller-defined request classification.
extra_key: Optional[Union[List[str], str]] = None extra_key: Optional[Union[List[str], str]] = None
# Whether to disallow logging for this request (e.g. due to ZDR) # Whether to disallow logging for this request (e.g. due to ZDR)
@@ -327,6 +327,9 @@ class GenerateReqInput:
# Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. # Batch-level: List[List[int]] (one per request). After __getitem__: List[int].
multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None
# Cache namespace used to isolate otherwise-identical prefixes.
cache_salt: Optional[Union[List[str], str]] = None
def regenerate_rid(self): def regenerate_rid(self):
"""Generate a new request ID and return it.""" """Generate a new request ID and return it."""
if isinstance(self.rid, list): if isinstance(self.rid, list):
@@ -488,6 +491,14 @@ class GenerateReqInput:
self.token_ids_logprob = None self.token_ids_logprob = None
if self.return_sampling_mask is None: if self.return_sampling_mask is None:
self.return_sampling_mask = False self.return_sampling_mask = False
for field_name in ("extra_key", "cache_salt"):
value = getattr(self, field_name)
if value is not None and not isinstance(value, str):
raise ValueError(
f"{field_name} should be a string for a single request."
)
if value == "":
setattr(self, field_name, None)
def _normalize_batch_inputs(self): def _normalize_batch_inputs(self):
"""Normalize inputs for a batch of examples, including parallel sampling expansion.""" """Normalize inputs for a batch of examples, including parallel sampling expansion."""
@@ -510,6 +521,7 @@ class GenerateReqInput:
self._normalize_return_hidden_states(num) self._normalize_return_hidden_states(num)
self._normalize_custom_logit_processor(num) self._normalize_custom_logit_processor(num)
self._normalize_extra_key(num) self._normalize_extra_key(num)
self._normalize_cache_salt(num)
self._normalize_bootstrap_params(num) self._normalize_bootstrap_params(num)
def _expand_inputs(self, num): def _expand_inputs(self, num):
@@ -703,16 +715,39 @@ class GenerateReqInput:
if self.extra_key is None: if self.extra_key is None:
return return
if isinstance(self.extra_key, str): if isinstance(self.extra_key, str):
self.extra_key = [self.extra_key] * num value = self.extra_key or None
self.extra_key = [value] * num
elif isinstance(self.extra_key, list): elif isinstance(self.extra_key, list):
if len(self.extra_key) != self.batch_size: if len(self.extra_key) != self.batch_size:
raise ValueError( raise ValueError(
"The length of extra_key should be equal to the batch size." "The length of extra_key should be equal to the batch size."
) )
if any(not isinstance(value, str) for value in self.extra_key):
raise ValueError("Every extra_key should be a string.")
self.extra_key = [value or None for value in self.extra_key]
self.extra_key = self.extra_key * self.parallel_sample_num self.extra_key = self.extra_key * self.parallel_sample_num
else: else:
raise ValueError("extra_key should be a list or a string.") raise ValueError("extra_key should be a list or a string.")
def _normalize_cache_salt(self, num):
"""Normalize cache_salt for batch processing."""
if self.cache_salt is None:
return
if isinstance(self.cache_salt, str):
value = self.cache_salt or None
self.cache_salt = [value] * num
elif isinstance(self.cache_salt, list):
if len(self.cache_salt) != self.batch_size:
raise ValueError(
"The length of cache_salt should be equal to the batch size."
)
if any(not isinstance(value, str) for value in self.cache_salt):
raise ValueError("Every cache_salt should be a string.")
self.cache_salt = [value or None for value in self.cache_salt]
self.cache_salt = self.cache_salt * self.parallel_sample_num
else:
raise ValueError("cache_salt should be a list or a string.")
def _normalize_bootstrap_params(self, num): def _normalize_bootstrap_params(self, num):
"""Normalize bootstrap parameters for batch processing.""" """Normalize bootstrap parameters for batch processing."""
# Normalize bootstrap_host # Normalize bootstrap_host
@@ -837,6 +872,7 @@ class GenerateReqInput:
max_thinking_tokens=self.max_thinking_tokens, max_thinking_tokens=self.max_thinking_tokens,
priority=self.priority, priority=self.priority,
extra_key=self.extra_key[i] if self.extra_key is not None else None, extra_key=self.extra_key[i] if self.extra_key is not None else None,
cache_salt=(self.cache_salt[i] if self.cache_salt is not None else None),
no_logs=self.no_logs, no_logs=self.no_logs,
custom_labels=self.custom_labels, custom_labels=self.custom_labels,
return_bytes=self.return_bytes, return_bytes=self.return_bytes,
@@ -925,7 +961,7 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
# Priority for the request # Priority for the request
priority: Optional[int] = None priority: Optional[int] = None
# Extra cache key for classifying the request (e.g. cache_salt) # Extra cache key for caller-defined request classification.
extra_key: Optional[str] = None extra_key: Optional[str] = None
# Whether to disallow logging for this request (e.g. due to ZDR) # Whether to disallow logging for this request (e.g. due to ZDR)
@@ -954,6 +990,9 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
# Pickled Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] # Pickled Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]]
time_stats: Optional[PickleWrapper] = None time_stats: Optional[PickleWrapper] = None
# Cache namespace used to isolate otherwise-identical prefixes.
cache_salt: Optional[str] = None
def wrap_pickle_fields(self): def wrap_pickle_fields(self):
self.mm_inputs = wrap_as_pickle(self.mm_inputs) self.mm_inputs = wrap_as_pickle(self.mm_inputs)
self.mm_data_mooncake = wrap_as_pickle(self.mm_data_mooncake) self.mm_data_mooncake = wrap_as_pickle(self.mm_data_mooncake)
+5 -1
View File
@@ -855,6 +855,7 @@ class Req(ReqDllmMixin):
return_pooled_hidden_states: bool = False, return_pooled_hidden_states: bool = False,
multi_item_delimiter_indices: Optional[List[int]] = None, multi_item_delimiter_indices: Optional[List[int]] = None,
session_id: Optional[str] = None, session_id: Optional[str] = None,
cache_salt: Optional[str] = None,
): ):
# Input and output info # Input and output info
self.rid = rid self.rid = rid
@@ -924,13 +925,14 @@ class Req(ReqDllmMixin):
return_hidden_states return_hidden_states
) )
# extra key for classifying the request (e.g. cache_salt) # Extra key for caller-defined request classification.
if lora_id is not None: if lora_id is not None:
extra_key = ( extra_key = (
extra_key or "" extra_key or ""
) + lora_id # lora_id is concatenated to the extra key ) + lora_id # lora_id is concatenated to the extra key
self.extra_key = extra_key self.extra_key = extra_key
self.cache_salt = cache_salt or None
self.lora_id = lora_id self.lora_id = lora_id
self.routing_key = routing_key self.routing_key = routing_key
@@ -1359,6 +1361,7 @@ class Req(ReqDllmMixin):
token_ids=token_ids_to_match, token_ids=token_ids_to_match,
extra_key=self.extra_key, extra_key=self.extra_key,
limit=key_limit, limit=key_limit,
cache_salt=self.cache_salt,
), ),
req=self, req=self,
cow_mamba=cow_mamba, cow_mamba=cow_mamba,
@@ -1776,6 +1779,7 @@ class Req(ReqDllmMixin):
"bootstrap_room": self.bootstrap_room, "bootstrap_room": self.bootstrap_room,
"priority": self.priority, "priority": self.priority,
"extra_key": self.extra_key, "extra_key": self.extra_key,
"cache_salt": self.cache_salt,
"routing_key": self.routing_key, "routing_key": self.routing_key,
"disagg_prefill_dp_rank": self.disagg_prefill_dp_rank, "disagg_prefill_dp_rank": self.disagg_prefill_dp_rank,
} }
+17 -3
View File
@@ -155,7 +155,12 @@ def match_prefix_for_req(
match_result = tree_cache.match_prefix( match_result = tree_cache.match_prefix(
MatchPrefixParams( MatchPrefixParams(
key=RadixKey(token_ids=token_ids, extra_key=req.extra_key, limit=key_limit), key=RadixKey(
token_ids=token_ids,
extra_key=req.extra_key,
limit=key_limit,
cache_salt=req.cache_salt,
),
cow_mamba=cow_mamba, cow_mamba=cow_mamba,
req=req if include_req else None, req=req if include_req else None,
) )
@@ -319,6 +324,7 @@ class SchedulePolicy:
for r in waiting_queue: for r in waiting_queue:
prefix_ids = r.origin_input_ids + r.output_ids prefix_ids = r.origin_input_ids + r.output_ids
extra_key = r.extra_key extra_key = r.extra_key
cache_salt = r.cache_salt
match_result = match_prefix_for_req( match_result = match_prefix_for_req(
self.tree_cache, r, prefix_ids, include_req=True self.tree_cache, r, prefix_ids, include_req=True
) )
@@ -333,7 +339,11 @@ class SchedulePolicy:
if len(r.prefix_indices) <= IN_BATCH_PREFIX_CACHING_CHECK_THRESHOLD: if len(r.prefix_indices) <= IN_BATCH_PREFIX_CACHING_CHECK_THRESHOLD:
match_result = self.waiting_queue_radix_tree.match_prefix( match_result = self.waiting_queue_radix_tree.match_prefix(
MatchPrefixParams( MatchPrefixParams(
key=RadixKey(token_ids=prefix_ids, extra_key=extra_key) key=RadixKey(
token_ids=prefix_ids,
extra_key=extra_key,
cache_salt=cache_salt,
)
) )
) )
if envs.SGLANG_RADIX_FORCE_MISS.get(): if envs.SGLANG_RADIX_FORCE_MISS.get():
@@ -350,7 +360,11 @@ class SchedulePolicy:
# Insert with a dummy key # Insert with a dummy key
self.waiting_queue_radix_tree.insert( self.waiting_queue_radix_tree.insert(
InsertParams( InsertParams(
key=RadixKey(token_ids=prefix_ids, extra_key=extra_key), key=RadixKey(
token_ids=prefix_ids,
extra_key=extra_key,
cache_salt=cache_salt,
),
value=torch.empty(len(prefix_ids), dtype=torch.bool), value=torch.empty(len(prefix_ids), dtype=torch.bool),
) )
) )
+1
View File
@@ -2417,6 +2417,7 @@ class Scheduler(
), ),
routing_key=recv_req.routing_key, routing_key=recv_req.routing_key,
extra_key=recv_req.extra_key, extra_key=recv_req.extra_key,
cache_salt=recv_req.cache_salt,
http_worker_ipc=recv_req.http_worker_ipc, http_worker_ipc=recv_req.http_worker_ipc,
dllm_config=self.dllm_config, dllm_config=self.dllm_config,
time_stats=recv_req.time_stats, time_stats=recv_req.time_stats,
@@ -1399,6 +1399,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
disagg_prefill_dp_rank=obj.disagg_prefill_dp_rank, disagg_prefill_dp_rank=obj.disagg_prefill_dp_rank,
priority=obj.priority, priority=obj.priority,
extra_key=obj.extra_key, extra_key=obj.extra_key,
cache_salt=obj.cache_salt,
routing_key=obj.routing_key, routing_key=obj.routing_key,
token_type_ids=token_type_ids, token_type_ids=token_type_ids,
need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs, need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs,
+37 -16
View File
@@ -23,9 +23,12 @@ from sglang.srt.disaggregation.kv_events import (
AllBlocksCleared, AllBlocksCleared,
BlockRemoved, BlockRemoved,
BlockStored, BlockStored,
BlockStoredMetadata,
BlockStoredWithMetadata,
StorageMedium, StorageMedium,
) )
from sglang.srt.mem_cache.utils import ( from sglang.srt.mem_cache.utils import (
compute_node_event_hash_values,
compute_node_hash_values, compute_node_hash_values,
hash_str_to_int64, hash_str_to_int64,
) )
@@ -43,15 +46,22 @@ class KVCacheEventMixin:
# Compute hash_value lazily if not already set # Compute hash_value lazily if not already set
if node.hash_value is None: if node.hash_value is None:
node.hash_value = compute_node_hash_values(node, self.page_size) node.hash_value = compute_node_hash_values(node, self.page_size)
event_hash_values = (
compute_node_event_hash_values(node, self.page_size)
if node.key.cache_salt is not None
else node.hash_value
)
# Get parent's last hash value for first page # Get parent's last hash value for first page
parent_block_hash = None parent_block_hash = None
if node.parent is not None and node.parent != self.root_node: if node.parent is not None and node.parent != self.root_node:
if ( if node.key.cache_salt is not None:
node.parent.hash_value is not None parent_hash_values = node.parent.event_hash_value
and len(node.parent.hash_value) > 0 assert parent_hash_values is not None
): else:
parent_block_hash = hash_str_to_int64(node.parent.hash_value[-1]) parent_hash_values = node.parent.hash_value
if parent_hash_values:
parent_block_hash = hash_str_to_int64(parent_hash_values[-1])
page_index = 0 page_index = 0
logical_len = len(node.key) logical_len = len(node.key)
@@ -67,18 +77,24 @@ class KVCacheEventMixin:
else: else:
page_tokens = list(raw[start:end]) page_tokens = list(raw[start:end])
block_hash = hash_str_to_int64(node.hash_value[page_index]) block_hash = hash_str_to_int64(event_hash_values[page_index])
self.kv_event_queue.append( event_args = {
BlockStored( "block_hashes": [block_hash],
block_hashes=[block_hash], "parent_block_hash": parent_block_hash,
parent_block_hash=parent_block_hash, "token_ids": page_tokens,
token_ids=page_tokens, "block_size": len(page_tokens),
block_size=len(page_tokens), "lora_id": None,
lora_id=None, "medium": medium,
medium=medium, }
if node.key.cache_salt is None:
event = BlockStored(**event_args)
else:
event = BlockStoredWithMetadata(
**event_args,
metadata=BlockStoredMetadata(cache_salt=node.key.cache_salt),
) )
) self.kv_event_queue.append(event)
parent_block_hash = block_hash parent_block_hash = block_hash
page_index += 1 page_index += 1
@@ -94,6 +110,11 @@ class KVCacheEventMixin:
# Compute hash_value lazily if not already set (must match what was stored) # Compute hash_value lazily if not already set (must match what was stored)
if node.hash_value is None: if node.hash_value is None:
node.hash_value = compute_node_hash_values(node, self.page_size) node.hash_value = compute_node_hash_values(node, self.page_size)
event_hash_values = (
compute_node_event_hash_values(node, self.page_size)
if node.key.cache_salt is not None
else node.hash_value
)
block_hashes = [] block_hashes = []
logical_len = len(node.key) logical_len = len(node.key)
@@ -103,7 +124,7 @@ class KVCacheEventMixin:
if end <= start: if end <= start:
continue continue
block_hashes.append(hash_str_to_int64(node.hash_value[page_index])) block_hashes.append(hash_str_to_int64(event_hash_values[page_index]))
page_index += 1 page_index += 1
if block_hashes: if block_hashes:
+11 -1
View File
@@ -963,7 +963,12 @@ class HiRadixCache(RadixCache):
token_ids = [] token_ids = []
for n in chain: for n in chain:
token_ids.extend(n.key.token_ids) token_ids.extend(n.key.token_ids)
key = RadixKey(token_ids, top.key.extra_key, top.key.is_bigram) key = RadixKey(
token_ids,
top.key.extra_key,
top.key.is_bigram,
cache_salt=top.key.cache_salt,
)
if all(n.hash_value is not None for n in chain): if all(n.hash_value is not None for n in chain):
hash_value = [] hash_value = []
@@ -1474,6 +1479,7 @@ class HiRadixCache(RadixCache):
new_input_tokens, new_input_tokens,
extra_key=last_host_node.key.extra_key, extra_key=last_host_node.key.extra_key,
is_bigram=self.is_eagle, is_bigram=self.is_eagle,
cache_salt=last_host_node.key.cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
if len(prefetch_key) < self.prefetch_threshold: if len(prefetch_key) < self.prefetch_threshold:
return 0 return 0
@@ -1773,6 +1779,7 @@ class HiRadixCache(RadixCache):
new_input_tokens, new_input_tokens,
extra_key=last_host_node.key.extra_key, extra_key=last_host_node.key.extra_key,
is_bigram=self.is_eagle, is_bigram=self.is_eagle,
cache_salt=last_host_node.key.cache_salt,
) )
# align the number of fetching tokens to the page size # align the number of fetching tokens to the page size
prefetch_key = prefetch_key.page_aligned(self.page_size) prefetch_key = prefetch_key.page_aligned(self.page_size)
@@ -1893,6 +1900,9 @@ class HiRadixCache(RadixCache):
new_node.hash_value, child.hash_value = split_node_hash_value( new_node.hash_value, child.hash_value = split_node_hash_value(
child.hash_value, split_len, self.page_size child.hash_value, split_len, self.page_size
) )
new_node.event_hash_value, child.event_hash_value = split_node_hash_value(
child.event_hash_value, split_len, self.page_size
)
child.parent = new_node child.parent = new_node
child.key = child.key[split_len:] child.key = child.key[split_len:]
new_node.parent.children[key.child_key(self.page_size)] = new_node new_node.parent.children[key.child_key(self.page_size)] = new_node
@@ -104,6 +104,8 @@ class TreeNode:
self.host_value = None self.host_value = None
# store hash values of each pages # store hash values of each pages
self.hash_value: Optional[List[str]] = None self.hash_value: Optional[List[str]] = None
# Namespace-aware hashes used only for external KV events.
self.event_hash_value: Optional[List[str]] = None
# for lru list, invariant: # for lru list, invariant:
# 1. prev has greater last_access_time # 1. prev has greater last_access_time
@@ -629,7 +631,11 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
result = self.insert( result = self.insert(
InsertParams( InsertParams(
key=RadixKey(token_ids[:page_aligned_len], req.extra_key), key=RadixKey(
token_ids[:page_aligned_len],
req.extra_key,
cache_salt=req.cache_salt,
),
value=page_aligned_kv_indices, value=page_aligned_kv_indices,
mamba_value=mamba_value, mamba_value=mamba_value,
prev_prefix_len=req.cache_protected_len, prev_prefix_len=req.cache_protected_len,
@@ -739,7 +745,11 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
result = self.insert( result = self.insert(
InsertParams( InsertParams(
key=RadixKey(page_aligned_token_ids, req.extra_key), key=RadixKey(
page_aligned_token_ids,
req.extra_key,
cache_salt=req.cache_salt,
),
value=page_aligned_kv_indices, value=page_aligned_kv_indices,
mamba_value=mamba_value_donated, mamba_value=mamba_value_donated,
prev_prefix_len=req.cache_protected_len, prev_prefix_len=req.cache_protected_len,
@@ -752,7 +762,13 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
# The prefix indices could be updated, reuse it # The prefix indices could be updated, reuse it
match_result = self.match_prefix( match_result = self.match_prefix(
MatchPrefixParams(key=RadixKey(page_aligned_token_ids, req.extra_key)) MatchPrefixParams(
key=RadixKey(
page_aligned_token_ids,
req.extra_key,
cache_salt=req.cache_salt,
)
)
) )
new_indices, new_last_node = ( new_indices, new_last_node = (
match_result.device_indices, match_result.device_indices,
@@ -1214,6 +1230,9 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
new_node.hash_value, child.hash_value = split_node_hash_value( new_node.hash_value, child.hash_value = split_node_hash_value(
child.hash_value, split_len, self.page_size child.hash_value, split_len, self.page_size
) )
new_node.event_hash_value, child.event_hash_value = split_node_hash_value(
child.event_hash_value, split_len, self.page_size
)
# insert the new node and child into the full lru list, insert # insert the new node and child into the full lru list, insert
# parent first so that parent is after child in the lru list # parent first so that parent is after child in the lru list
@@ -89,7 +89,10 @@ class PureSWARadixCache(RadixCache):
] ]
radix_key = RadixKey( radix_key = RadixKey(
token_ids, req.extra_key, is_bigram=self.is_eagle token_ids,
req.extra_key,
is_bigram=self.is_eagle,
cache_salt=req.cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
keys_len = len(radix_key) keys_len = len(radix_key)
+40 -7
View File
@@ -59,7 +59,7 @@ if TYPE_CHECKING:
class RadixKey: class RadixKey:
"""is_bigram=True: token_ids holds raw tokens (N+1 for N bigrams); slices share one boundary token.""" """is_bigram=True: token_ids holds raw tokens (N+1 for N bigrams); slices share one boundary token."""
__slots__ = ("token_ids", "extra_key", "is_bigram", "limit") __slots__ = ("token_ids", "extra_key", "cache_salt", "is_bigram", "limit")
def __init__( def __init__(
self, self,
@@ -67,11 +67,17 @@ class RadixKey:
extra_key: Optional[str] = None, extra_key: Optional[str] = None,
is_bigram: bool = False, is_bigram: bool = False,
limit: Optional[int] = None, limit: Optional[int] = None,
cache_salt: Optional[str] = None,
): ):
# token ids sequence (raw ints in both modes) # token ids sequence (raw ints in both modes)
self.token_ids = token_ids self.token_ids = token_ids
# extra key (e.g. lora_id, cache_salt) # Extra key for caller-defined cache classification.
self.extra_key = extra_key self.extra_key = extra_key
# Cache salt is kept distinct so it cannot collide with extra_key.
# It namespaces the in-process radix tree and external KV events;
# external L3/remote storage keys remain token-only and are outside
# this contract.
self.cache_salt = cache_salt or None
# bigram view over token_ids: length = max(0, len(token_ids) - 1) # bigram view over token_ids: length = max(0, len(token_ids) - 1)
self.is_bigram = is_bigram self.is_bigram = is_bigram
# Optional cap on raw tokens: behave as if token_ids were sliced to # Optional cap on raw tokens: behave as if token_ids were sliced to
@@ -125,12 +131,21 @@ class RadixKey:
# bigrams [start, stop) span raw tokens [start, stop + 1); # bigrams [start, stop) span raw tokens [start, stop + 1);
# empty slice -> empty raw tokens (not a dangling boundary token). # empty slice -> empty raw tokens (not a dangling boundary token).
raw = self.token_ids[start : stop + 1] if stop > start else array("q") raw = self.token_ids[start : stop + 1] if stop > start else array("q")
return RadixKey(raw, self.extra_key, is_bigram=True) return RadixKey(
return RadixKey(self.token_ids[start:stop], self.extra_key) raw,
self.extra_key,
is_bigram=True,
cache_salt=self.cache_salt,
)
return RadixKey(
self.token_ids[start:stop],
self.extra_key,
cache_salt=self.cache_salt,
)
def __repr__(self) -> str: def __repr__(self) -> str:
preview = self.token_ids[:10] preview = self.token_ids[:10]
return f"RadixKey(extra_key={self.extra_key!r}, token_ids={preview}{'...' if len(self.token_ids) > 10 else ''}, is_bigram={self.is_bigram})" return f"RadixKey(extra_key={self.extra_key!r}, cache_salt={self.cache_salt!r}, token_ids={preview}{'...' if len(self.token_ids) > 10 else ''}, is_bigram={self.is_bigram})"
def page_aligned(self, page_size: int) -> RadixKey: def page_aligned(self, page_size: int) -> RadixKey:
if page_size == 1: if page_size == 1:
@@ -157,6 +172,11 @@ class RadixKey:
f"RadixKey operations require matching extra_key, but got " f"RadixKey operations require matching extra_key, but got "
f"{self.extra_key=} != {other.extra_key=}" f"{self.extra_key=} != {other.extra_key=}"
) )
if self.cache_salt != other.cache_salt:
raise ValueError(
f"RadixKey operations require matching cache_salt, but got "
f"{self.cache_salt=} != {other.cache_salt=}"
)
def match(self, other: RadixKey, page_size: int = 1) -> int: def match(self, other: RadixKey, page_size: int = 1) -> int:
"""Logical-unit prefix length shared with ``other``. Result is rounded down to ``page_size``.""" """Logical-unit prefix length shared with ``other``. Result is rounded down to ``page_size``."""
@@ -204,6 +224,8 @@ class RadixKey:
plain = tuple((t[j], t[j + 1]) for j in range(page_size)) plain = tuple((t[j], t[j + 1]) for j in range(page_size))
else: else:
plain = t[0] if page_size == 1 else tuple(t[:page_size]) plain = t[0] if page_size == 1 else tuple(t[:page_size])
if self.cache_salt is not None:
return ((self.extra_key, self.cache_salt), plain)
return plain if self.extra_key is None else (self.extra_key, plain) return plain if self.extra_key is None else (self.extra_key, plain)
def hash_page(self, start: int, end: int, prior_hash: Optional[str] = None) -> str: def hash_page(self, start: int, end: int, prior_hash: Optional[str] = None) -> str:
@@ -235,6 +257,8 @@ class TreeNode:
self.write_through_pending_id: Optional[int] = None self.write_through_pending_id: Optional[int] = None
# store hash values of each pages # store hash values of each pages
self.hash_value: Optional[List[str]] = None self.hash_value: Optional[List[str]] = None
# Namespace-aware hashes used only for external KV events.
self.event_hash_value: Optional[List[str]] = None
# priority for priority-aware eviction # priority for priority-aware eviction
self.priority = priority self.priority = priority
@@ -455,7 +479,10 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
] ]
radix_key = RadixKey( radix_key = RadixKey(
token_ids, req.extra_key, is_bigram=self.is_eagle token_ids,
req.extra_key,
is_bigram=self.is_eagle,
cache_salt=req.cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
key_len = len(radix_key) key_len = len(radix_key)
values = kv_indices[:key_len].to(dtype=torch.int64, copy=True) values = kv_indices[:key_len].to(dtype=torch.int64, copy=True)
@@ -496,7 +523,10 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
] ]
radix_key = RadixKey( radix_key = RadixKey(
token_ids, req.extra_key, is_bigram=self.is_eagle token_ids,
req.extra_key,
is_bigram=self.is_eagle,
cache_salt=req.cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True) values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True)
@@ -690,6 +720,9 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
new_node.hash_value, child.hash_value = split_node_hash_value( new_node.hash_value, child.hash_value = split_node_hash_value(
child.hash_value, split_len, self.page_size child.hash_value, split_len, self.page_size
) )
new_node.event_hash_value, child.event_hash_value = split_node_hash_value(
child.event_hash_value, split_len, self.page_size
)
return new_node return new_node
@@ -33,6 +33,13 @@ logger = logging.getLogger(__name__)
class RadixCacheCpp(BasePrefixCache): class RadixCacheCpp(BasePrefixCache):
@staticmethod
def _reject_cache_salt(cache_salt: Optional[str]) -> None:
if cache_salt is not None:
raise ValueError(
"cache_salt is not supported by the experimental C++ radix tree"
)
def __init__( def __init__(
self, self,
params: CacheInitParams, params: CacheInitParams,
@@ -100,6 +107,7 @@ class RadixCacheCpp(BasePrefixCache):
def match_prefix(self, params: MatchPrefixParams) -> MatchResult: def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
key = params.key key = params.key
self._reject_cache_salt(key.cache_salt)
device_indices_vec, host_indices_length, node_gpu, node_cpu = ( device_indices_vec, host_indices_length, node_gpu, node_cpu = (
self.tree.match_prefix(key.raw_token_ids()) self.tree.match_prefix(key.raw_token_ids())
) )
@@ -173,6 +181,7 @@ class RadixCacheCpp(BasePrefixCache):
self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
): ):
"""Cache request when it finishes.""" """Cache request when it finishes."""
self._reject_cache_salt(req.cache_salt)
assert req.req_pool_idx is not None assert req.req_pool_idx is not None
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
@@ -210,6 +219,7 @@ class RadixCacheCpp(BasePrefixCache):
def cache_unfinished_req(self, req: Req, chunked=False): def cache_unfinished_req(self, req: Req, chunked=False):
"""Cache request when it is unfinished.""" """Cache request when it is unfinished."""
self._reject_cache_salt(req.cache_salt)
assert req.req_pool_idx is not None assert req.req_pool_idx is not None
token_ids = req.get_fill_ids() token_ids = req.get_fill_ids()
prefill_len = len(token_ids) # prefill only (maybe chunked) prefill_len = len(token_ids) # prefill only (maybe chunked)
@@ -216,7 +216,12 @@ class FlexKVRadixCache(RadixCache):
else: else:
token_ids_snap = token_ids token_ids_snap = token_ids
self._load_markers[req.rid] = _LoadBackMarker( self._load_markers[req.rid] = _LoadBackMarker(
key=RadixKey(token_ids_snap, key.extra_key, key.is_bigram), key=RadixKey(
token_ids_snap,
key.extra_key,
key.is_bigram,
cache_salt=key.cache_salt,
),
value_numel=device_len, value_numel=device_len,
) )
return MatchResult( return MatchResult(
@@ -410,7 +415,13 @@ class FlexKVRadixCache(RadixCache):
# Anchor on the new last_device_node so FlexKV's lock matches # Anchor on the new last_device_node so FlexKV's lock matches
# the node we'll later unlock when the store completes. # the node we'll later unlock when the store completes.
match_result = super().match_prefix( match_result = super().match_prefix(
MatchPrefixParams(key=RadixKey(token_ids, req.extra_key)) MatchPrefixParams(
key=RadixKey(
token_ids,
req.extra_key,
cache_salt=req.cache_salt,
)
)
) )
new_last_node = match_result.last_device_node new_last_node = match_result.last_device_node
if new_last_node is None: if new_last_node is None:
@@ -229,7 +229,12 @@ class LMCRadixCache(RadixCache):
if token_ids is key.token_ids: if token_ids is key.token_ids:
token_ids = token_ids[:] token_ids = token_ids[:]
self._mp_load_back_markers[req.rid] = _LMCacheLoadBackMarker( self._mp_load_back_markers[req.rid] = _LMCacheLoadBackMarker(
key=RadixKey(token_ids, key.extra_key, key.is_bigram), key=RadixKey(
token_ids,
key.extra_key,
key.is_bigram,
cache_salt=key.cache_salt,
),
value_numel=int(value.numel()), value_numel=int(value.numel()),
) )
return MatchResult( return MatchResult(
@@ -463,7 +468,13 @@ class LMCRadixCache(RadixCache):
# Use super() to avoid a redundant LOOKUP — we only need new_last_node from radix. # Use super() to avoid a redundant LOOKUP — we only need new_last_node from radix.
match_result = super().match_prefix( match_result = super().match_prefix(
MatchPrefixParams(key=RadixKey(token_ids, req.extra_key)) MatchPrefixParams(
key=RadixKey(
token_ids,
req.extra_key,
cache_salt=req.cache_salt,
)
)
) )
new_last_node = match_result.last_device_node new_last_node = match_result.last_device_node
assert new_last_node is not None assert new_last_node is not None
+20 -2
View File
@@ -80,6 +80,8 @@ class TreeNode:
self.host_value = None self.host_value = None
# store hash values of each page # store hash values of each page
self.hash_value: Optional[List[str]] = None self.hash_value: Optional[List[str]] = None
# Namespace-aware hashes used only for external KV events.
self.event_hash_value: Optional[List[str]] = None
# for lru list, invariant: # for lru list, invariant:
# 1. prev has greater last_access_time # 1. prev has greater last_access_time
@@ -473,7 +475,10 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
] ]
radix_key = RadixKey( radix_key = RadixKey(
token_ids, req.extra_key, is_bigram=self.is_eagle token_ids,
req.extra_key,
is_bigram=self.is_eagle,
cache_salt=req.cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
page_aligned_len = len(radix_key) page_aligned_len = len(radix_key)
values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True)
@@ -523,7 +528,10 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
] ]
radix_key = RadixKey( radix_key = RadixKey(
token_ids, req.extra_key, is_bigram=self.is_eagle token_ids,
req.extra_key,
is_bigram=self.is_eagle,
cache_salt=req.cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True) values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True)
old_prefix_len = req.cache_protected_len old_prefix_len = req.cache_protected_len
@@ -1020,6 +1028,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
node.key.token_ids + child.key.token_ids, node.key.token_ids + child.key.token_ids,
node.key.extra_key, node.key.extra_key,
is_bigram=node.key.is_bigram, is_bigram=node.key.is_bigram,
cache_salt=node.key.cache_salt,
) )
node.value = torch.cat([node.value, child.value]) node.value = torch.cat([node.value, child.value])
node.children = child.children node.children = child.children
@@ -1033,6 +1042,12 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
node.hash_value = list(node.hash_value) + list(child.hash_value) node.hash_value = list(node.hash_value) + list(child.hash_value)
else: else:
node.hash_value = None node.hash_value = None
if node.event_hash_value is not None and child.event_hash_value is not None:
node.event_hash_value = list(node.event_hash_value) + list(
child.event_hash_value
)
else:
node.event_hash_value = None
self.full_lru_list.remove_node(child) self.full_lru_list.remove_node(child)
if not child.swa_tombstone: if not child.swa_tombstone:
@@ -1104,6 +1119,9 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
new_node.hash_value, child.hash_value = split_node_hash_value( new_node.hash_value, child.hash_value = split_node_hash_value(
child.hash_value, split_len, self.page_size child.hash_value, split_len, self.page_size
) )
new_node.event_hash_value, child.event_hash_value = split_node_hash_value(
child.event_hash_value, split_len, self.page_size
)
# insert the new node and child into the lru lists, insert # insert the new node and child into the lru lists, insert
# parent first so that parent is after child in the lru list # parent first so that parent is after child in the lru list
@@ -117,6 +117,8 @@ class UnifiedTreeNode:
self.last_access_time = get_and_increase_time_counter() self.last_access_time = get_and_increase_time_counter()
self.creation_time = get_and_increase_time_counter() self.creation_time = get_and_increase_time_counter()
self.hash_value = None self.hash_value = None
# Namespace-aware hashes used only for external KV events.
self.event_hash_value: Optional[list[str]] = None
self.hit_count = 0 self.hit_count = 0
self.priority = priority self.priority = priority
self.lru_prev: list[UnifiedTreeNode | None] = [None] * ( self.lru_prev: list[UnifiedTreeNode | None] = [None] * (
@@ -1065,6 +1067,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
new_node.hash_value, child.hash_value = split_node_hash_value( new_node.hash_value, child.hash_value = split_node_hash_value(
child.hash_value, split_len, self.page_size child.hash_value, split_len, self.page_size
) )
new_node.event_hash_value, child.event_hash_value = split_node_hash_value(
child.event_hash_value, split_len, self.page_size
)
for component in self.components: for component in self.components:
component.redistribute_on_node_split(new_parent=new_node, child=child) component.redistribute_on_node_split(new_parent=new_node, child=child)
@@ -1867,10 +1872,14 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
comp_xfers[comp.component_type] = t comp_xfers[comp.component_type] = t
return kv_xfer, comp_xfers return kv_xfer, comp_xfers
def prefetch_anchor_info(self, node_id: NodeId) -> Optional[str]: def prefetch_anchor_info(
"""The anchor node's key extra_key.""" self, node_id: NodeId
) -> tuple[Optional[str], Optional[str]]:
"""The anchor node's key extra_key and cache_salt."""
node = self.node_by_id(node_id) node = self.node_by_id(node_id)
return node.key.extra_key if node.key else None if node.key is None:
return None, None
return node.key.extra_key, node.key.cache_salt
def _build_backup_kv_action( def _build_backup_kv_action(
self, node: UnifiedTreeNode, write_back: bool = False self, node: UnifiedTreeNode, write_back: bool = False
@@ -421,8 +421,10 @@ class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC):
... ...
@abstractmethod @abstractmethod
def prefetch_anchor_info(self, node_id: NodeId) -> Optional[str]: def prefetch_anchor_info(
"""The anchor node's key extra_key.""" self, node_id: NodeId
) -> tuple[Optional[str], Optional[str]]:
"""The anchor node's key extra_key and cache_salt."""
... ...
@abstractmethod @abstractmethod
@@ -697,7 +697,10 @@ class UnifiedRadixCache(BasePrefixCache):
kv_indices = kv_indices[:effective_cache_len] kv_indices = kv_indices[:effective_cache_len]
radix_key = RadixKey( radix_key = RadixKey(
token_ids, req.extra_key, is_bigram=self.tree_core.is_eagle token_ids,
req.extra_key,
is_bigram=self.tree_core.is_eagle,
cache_salt=req.cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
page_aligned_len = len(radix_key) page_aligned_len = len(radix_key)
values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True)
@@ -790,6 +793,7 @@ class UnifiedRadixCache(BasePrefixCache):
token_ids[:effective_cache_len], token_ids[:effective_cache_len],
req.extra_key, req.extra_key,
is_bigram=self.tree_core.is_eagle, is_bigram=self.tree_core.is_eagle,
cache_salt=req.cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
page_aligned_len = len(radix_key) page_aligned_len = len(radix_key)
values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True)
@@ -1219,11 +1223,12 @@ class UnifiedRadixCache(BasePrefixCache):
if not self.enable_storage or self.cache_controller is None: if not self.enable_storage or self.cache_controller is None:
return return
extra_key = self.tree_core.prefetch_anchor_info(last_host_node_id) extra_key, cache_salt = self.tree_core.prefetch_anchor_info(last_host_node_id)
prefetch_key = RadixKey( prefetch_key = RadixKey(
new_input_tokens, new_input_tokens,
extra_key=extra_key, extra_key=extra_key,
is_bigram=self.tree_core.is_eagle, is_bigram=self.tree_core.is_eagle,
cache_salt=cache_salt,
).page_aligned(self.page_size) ).page_aligned(self.page_size)
prefetch_length = len(prefetch_key) prefetch_length = len(prefetch_key)
if ( if (
+49
View File
@@ -13,6 +13,7 @@
# ============================================================================== # ==============================================================================
"""Common utilities.""" """Common utilities."""
import hashlib
from typing import Any, Callable, List, Optional, Tuple from typing import Any, Callable, List, Optional, Tuple
from sglang.kernels.ops.kvcache.mla_buffer import ( from sglang.kernels.ops.kvcache.mla_buffer import (
@@ -135,6 +136,54 @@ def compute_node_hash_values(node: Any, page_size: int) -> List[str]:
return hash_values return hash_values
def compute_node_event_hash_values(node: Any, page_size: int) -> List[str]:
"""Compute and memoize namespace-aware external KV-event hashes."""
cache_salt = node.key.cache_salt
if cache_salt is None:
return compute_node_hash_values(node, page_size)
if node.event_hash_value is not None:
return node.event_hash_value
missing_nodes = []
current = node
while (
current is not None
and current.key is not None
and len(current.key) > 0
and current.event_hash_value is None
):
if current.key.cache_salt != cache_salt:
raise ValueError("Radix path contains mismatched cache_salt values")
missing_nodes.append(current)
current = current.parent
if (
current is not None
and current.key is not None
and len(current.key) > 0
and current.key.cache_salt != cache_salt
):
raise ValueError("Radix path contains mismatched cache_salt values")
if current is not None and current.event_hash_value:
parent_hash = current.event_hash_value[-1]
else:
parent_hash = hashlib.sha256(
b"sglang-cache-salt-v1\0" + cache_salt.encode("utf-8")
).hexdigest()
for missing_node in reversed(missing_nodes):
hash_values = get_hash_str(missing_node.key, parent_hash, page_size=page_size)
assert isinstance(hash_values, list)
missing_node.event_hash_value = hash_values
if hash_values:
parent_hash = hash_values[-1]
assert node.event_hash_value is not None
return node.event_hash_value
def split_node_hash_value( def split_node_hash_value(
child_hash_value: Optional[List[str]], split_len: int, page_size: int child_hash_value: Optional[List[str]], split_len: int, page_size: int
) -> tuple[Optional[List[str]], Optional[List[str]]]: ) -> tuple[Optional[List[str]], Optional[List[str]]]:
@@ -312,6 +312,7 @@ class Session:
priority=req.priority, priority=req.priority,
routing_key=req.routing_key, routing_key=req.routing_key,
extra_key=req.extra_key, extra_key=req.extra_key,
cache_salt=req.cache_salt,
http_worker_ipc=req.http_worker_ipc, http_worker_ipc=req.http_worker_ipc,
time_stats=req.time_stats, time_stats=req.time_stats,
) )
@@ -0,0 +1,61 @@
"""Unit tests for request construction in the encode-disaggregation path."""
import unittest
from array import array
from types import SimpleNamespace
from sglang.srt.disaggregation.encode_receiver import MMReceiverBase
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestEncodeReceiverRequestConstruction(CustomTestCase):
def test_extra_key_and_cache_salt_are_forwarded(self):
scheduler = SimpleNamespace(
model_config=SimpleNamespace(hf_eos_token_id={2}, vocab_size=128),
disaggregation_mode=DisaggregationMode.NULL,
metrics_reporter=SimpleNamespace(enable_metrics=False),
metrics_collector=None,
dllm_config=None,
tokenizer=object(),
)
receiver = SimpleNamespace(scheduler=scheduler)
recv_req = SimpleNamespace(
rid="request-1",
input_text="hello",
input_ids=array("q", [1, 2]),
sampling_params=SamplingParams(max_new_tokens=1),
return_logprob=False,
top_logprobs_num=0,
token_ids_logprob=None,
stream=False,
lora_id=None,
input_embeds=None,
custom_logit_processor=None,
require_reasoning=False,
return_hidden_states=False,
return_routed_experts=False,
routed_experts_start_len=0,
bootstrap_host=None,
bootstrap_port=None,
bootstrap_room=None,
routed_dp_rank=None,
disagg_prefill_dp_rank=None,
priority=None,
extra_key="classification",
cache_salt="tenant-a",
http_worker_ipc=None,
)
req = MMReceiverBase.create_req(receiver, recv_req)
self.assertEqual(req.extra_key, "classification")
self.assertEqual(req.cache_salt, "tenant-a")
if __name__ == "__main__":
unittest.main()
@@ -9,7 +9,14 @@ the router can subscribe per replica (the `dp_size` it reads from
import unittest import unittest
import msgspec
from sglang.srt.disaggregation.kv_events import ( from sglang.srt.disaggregation.kv_events import (
BlockStored,
BlockStoredMetadata,
BlockStoredWithMetadata,
KVEventBatch,
StorageMedium,
ZmqEventPublisher, ZmqEventPublisher,
select_kv_publisher_dp_rank, select_kv_publisher_dp_rank,
) )
@@ -94,5 +101,44 @@ class TestSelectKvPublisherDpRank(CustomTestCase):
self.assertEqual(len(ranks), dp_size) self.assertEqual(len(ranks), dp_size)
class TestBlockStoredWireFormat(CustomTestCase):
def _event(self, metadata=None):
event_type = BlockStored if metadata is None else BlockStoredWithMetadata
kwargs = dict(
block_hashes=[123],
parent_block_hash=None,
token_ids=[1, 2],
block_size=2,
lora_id=None,
medium=StorageMedium.GPU,
)
if metadata is not None:
kwargs["metadata"] = metadata
return event_type(**kwargs)
def test_unsalted_event_keeps_legacy_array_shape(self):
decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(self._event()))
self.assertEqual(len(decoded), 7)
def test_salted_event_appends_typed_metadata(self):
event = self._event(BlockStoredMetadata(cache_salt="tenant-a"))
encoded = msgspec.msgpack.encode(event)
decoded = msgspec.msgpack.decode(encoded)
round_tripped = msgspec.msgpack.decode(encoded, type=BlockStoredWithMetadata)
self.assertEqual(len(decoded), 8)
self.assertEqual(decoded[7], {"cache_salt": "tenant-a"})
self.assertEqual(round_tripped.metadata.cache_salt, "tenant-a")
def test_salted_event_remains_compatible_with_typed_batch_consumers(self):
batch = KVEventBatch(
ts=1.0,
events=[self._event(BlockStoredMetadata(cache_salt="tenant-a"))],
)
round_tripped = msgspec.msgpack.decode(
msgspec.msgpack.encode(batch), type=KVEventBatch
)
self.assertEqual(round_tripped.events[0].block_hashes, [123])
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()
@@ -316,6 +316,8 @@ class ServingChatTestCase(unittest.TestCase):
input_ids=[101, 102, 103], input_ids=[101, 102, 103],
stop=["STOP"], stop=["STOP"],
return_prompt_token_ids=True, return_prompt_token_ids=True,
cache_salt="tenant-a",
extra_key="classification",
) )
with patch( with patch(
@@ -329,6 +331,8 @@ class ServingChatTestCase(unittest.TestCase):
self.assertEqual(adapted.input_ids, [101, 102, 103]) self.assertEqual(adapted.input_ids, [101, 102, 103])
self.assertTrue(adapted.return_prompt_token_ids) self.assertTrue(adapted.return_prompt_token_ids)
self.assertEqual(adapted.sampling_params["stop"], ["STOP"]) self.assertEqual(adapted.sampling_params["stop"], ["STOP"])
self.assertEqual(adapted.cache_salt, "tenant-a")
self.assertEqual(adapted.extra_key, "classification")
conv_mock.assert_not_called() conv_mock.assert_not_called()
def test_kimi_k3_usage_excludes_assistant_generation_stub(self): def test_kimi_k3_usage_excludes_assistant_generation_stub(self):
@@ -65,6 +65,29 @@ class ServingCompletionTestCase(unittest.TestCase):
internal, _ = self.sc._convert_to_internal_request(req) internal, _ = self.sc._convert_to_internal_request(req)
self.assertEqual(internal.input_ids, [1, 2, 3, 4]) self.assertEqual(internal.input_ids, [1, 2, 3, 4])
def test_cache_salt_and_extra_key_remain_distinct(self):
req = CompletionRequest(
model="x",
prompt=[1, 2, 3, 4],
max_tokens=1,
cache_salt="tenant-a",
extra_key="classification",
)
internal, _ = self.sc._convert_to_internal_request(req)
self.assertEqual(internal.cache_salt, "tenant-a")
self.assertEqual(internal.extra_key, "classification")
def test_single_request_rejects_batched_cache_salt(self):
req = CompletionRequest(
model="x",
prompt=[1, 2, 3, 4],
max_tokens=1,
cache_salt=["tenant-a"],
)
internal, _ = self.sc._convert_to_internal_request(req)
with self.assertRaisesRegex(ValueError, "single request"):
internal.normalize_batch_and_arguments()
# ---------- echo-handling ---------- # ---------- echo-handling ----------
def test_echo_with_list_of_strings_streaming(self): def test_echo_with_list_of_strings_streaming(self):
req = CompletionRequest( req = CompletionRequest(
@@ -519,6 +519,53 @@ class TestGenerateReqInputNormalization(CustomTestCase):
req.normalize_batch_and_arguments() req.normalize_batch_and_arguments()
self.assertEqual(req.extra_key, "solo") self.assertEqual(req.extra_key, "solo")
def test_cache_salt_normalization(self):
req = GenerateReqInput(
text=["Hello", "World"],
cache_salt=["tenant-A", ""],
sampling_params=[{}, {}],
)
req.normalize_batch_and_arguments()
self.assertEqual(req.cache_salt, ["tenant-A", None])
self.assertEqual(req[0].cache_salt, "tenant-A")
self.assertIsNone(req[1].cache_salt)
req = GenerateReqInput(
text=["Hello", "World"],
cache_salt="shared",
sampling_params={"n": 2},
)
req.normalize_batch_and_arguments()
self.assertEqual(req.cache_salt, ["shared", "shared"] * 2)
req = GenerateReqInput(text="Hello", cache_salt="")
req.normalize_batch_and_arguments()
self.assertIsNone(req.cache_salt)
req = GenerateReqInput(
text=["Hello", "World"],
cache_salt=["only-one"],
sampling_params=[{}, {}],
)
with self.assertRaisesRegex(ValueError, "batch size"):
req.normalize_batch_and_arguments()
def test_cache_key_normalization_rejects_invalid_types(self):
for field_name in ("extra_key", "cache_salt"):
with self.subTest(field_name=field_name, mode="single"):
req = GenerateReqInput(text="Hello", **{field_name: ["value"]})
with self.assertRaisesRegex(ValueError, "single request"):
req.normalize_batch_and_arguments()
with self.subTest(field_name=field_name, mode="batch"):
req = GenerateReqInput(
text=["Hello", "World"],
sampling_params=[{}, {}],
**{field_name: ["value", 1]},
)
with self.assertRaisesRegex(ValueError, "should be a string"):
req.normalize_batch_and_arguments()
def test_logprob_parameters_normalization(self): def test_logprob_parameters_normalization(self):
"""Test normalization of logprob-related parameters.""" """Test normalization of logprob-related parameters."""
# Test single example # Test single example
@@ -244,6 +244,7 @@ class TestDecodePreallocQueueRebootstrapPayload(unittest.TestCase):
bootstrap_room=7, bootstrap_room=7,
priority=10, priority=10,
extra_key=None, extra_key=None,
cache_salt=None,
routing_key=None, routing_key=None,
disagg_prefill_dp_rank=None, disagg_prefill_dp_rank=None,
) )
@@ -260,6 +261,7 @@ class TestDecodePreallocQueueRebootstrapPayload(unittest.TestCase):
self.assertTrue(all(type(x) is int for x in payload["input_ids"])) self.assertTrue(all(type(x) is int for x in payload["input_ids"]))
self.assertEqual(payload["sampling_params"]["max_new_tokens"], 1) self.assertEqual(payload["sampling_params"]["max_new_tokens"], 1)
self.assertEqual(payload["bootstrap_room"], 7) self.assertEqual(payload["bootstrap_room"], 7)
self.assertIsNone(payload["cache_salt"])
# The prefill /generate URL is derived from bootstrap info on the decode # The prefill /generate URL is derived from bootstrap info on the decode
# side, not sent in the payload; and the boundary token is replayed via # side, not sent in the payload; and the boundary token is replayed via
# the decode-side override, so neither belongs in the payload. # the decode-side override, so neither belongs in the payload.
@@ -47,6 +47,7 @@ def _make_req(
req.logprob_start_len = -1 req.logprob_start_len = -1
req.positional_embed_overrides = None req.positional_embed_overrides = None
req.extra_key = None req.extra_key = None
req.cache_salt = None
req.mamba_pool_idx = None req.mamba_pool_idx = None
req.sampling_params = SimpleNamespace(max_new_tokens=128, ignore_eos=False) req.sampling_params = SimpleNamespace(max_new_tokens=128, ignore_eos=False)
return req return req
@@ -78,6 +78,7 @@ class MockReq:
self.cache_protected_len = cache_protected_len self.cache_protected_len = cache_protected_len
self.last_node = last_node self.last_node = last_node
self.extra_key = None self.extra_key = None
self.cache_salt = None
self.prefix_indices = torch.empty(0, dtype=torch.int64) self.prefix_indices = torch.empty(0, dtype=torch.int64)
self.priority = 0 self.priority = 0
self.kv_committed_len = len(fill_ids) self.kv_committed_len = len(fill_ids)
@@ -5,6 +5,7 @@ import sys
import types import types
import unittest import unittest
from array import array from array import array
from types import SimpleNamespace
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
from sglang.srt.mem_cache.evict_policy import ( from sglang.srt.mem_cache.evict_policy import (
@@ -17,6 +18,7 @@ from sglang.srt.mem_cache.evict_policy import (
SLRUStrategy, SLRUStrategy,
) )
from sglang.srt.mem_cache.utils import ( from sglang.srt.mem_cache.utils import (
compute_node_event_hash_values,
compute_node_hash_values, compute_node_hash_values,
get_eviction_strategy, get_eviction_strategy,
get_hash_str, get_hash_str,
@@ -54,9 +56,10 @@ def _legacy_page_hashes(key, page_size, prior_hash=None):
class _HashKey: class _HashKey:
def __init__(self, token_ids, is_bigram=False): def __init__(self, token_ids, is_bigram=False, cache_salt=None):
self.token_ids = token_ids self.token_ids = token_ids
self.is_bigram = is_bigram self.is_bigram = is_bigram
self.cache_salt = cache_salt
def __len__(self): def __len__(self):
if self.is_bigram: if self.is_bigram:
@@ -68,8 +71,12 @@ class _HashKey:
start = index.start or 0 start = index.start or 0
stop = index.stop if index.stop is not None else len(self) stop = index.stop if index.stop is not None else len(self)
if self.is_bigram: if self.is_bigram:
return _HashKey(self.token_ids[start : stop + 1], is_bigram=True) return _HashKey(
return _HashKey(self.token_ids[start:stop]) self.token_ids[start : stop + 1],
is_bigram=True,
cache_salt=self.cache_salt,
)
return _HashKey(self.token_ids[start:stop], cache_salt=self.cache_salt)
if self.is_bigram: if self.is_bigram:
return (self.token_ids[index], self.token_ids[index + 1]) return (self.token_ids[index], self.token_ids[index + 1])
return self.token_ids[index] return self.token_ids[index]
@@ -298,6 +305,7 @@ class TestComputeNodeHashValues(unittest.TestCase):
node = MagicMock() node = MagicMock()
node.key = key node.key = key
node.parent = parent node.parent = parent
node.event_hash_value = None
if parent is not None: if parent is not None:
parent.hash_value = parent_hash_values parent.hash_value = parent_hash_values
return node return node
@@ -318,6 +326,60 @@ class TestComputeNodeHashValues(unittest.TestCase):
_legacy_page_hashes(key, page_size=page_size), _legacy_page_hashes(key, page_size=page_size),
) )
def test_cache_salt_seeds_root_hash_chain(self):
key = _HashKey(array("q", range(1, 17)), cache_salt="tenant-a")
seed = hashlib.sha256(b"sglang-cache-salt-v1\0tenant-a").hexdigest()
self.assertEqual(
compute_node_event_hash_values(self._make_node(key), page_size=8),
_legacy_page_hashes(key, page_size=8, prior_hash=seed),
)
self.assertEqual(
compute_node_hash_values(self._make_node(key), page_size=8),
_legacy_page_hashes(key, page_size=8),
)
other = _HashKey(array("q", range(1, 17)), cache_salt="tenant-b")
self.assertNotEqual(
compute_node_event_hash_values(self._make_node(key), page_size=8),
compute_node_event_hash_values(self._make_node(other), page_size=8),
)
def test_cache_salt_event_hashes_are_memoized(self):
node = self._make_node(
_HashKey(array("q", range(1, 17)), cache_salt="tenant-a")
)
with patch(
"sglang.srt.mem_cache.utils.get_hash_str", wraps=get_hash_str
) as mock_get_hash_str:
first = compute_node_event_hash_values(node, page_size=8)
second = compute_node_event_hash_values(node, page_size=8)
self.assertIs(first, second)
mock_get_hash_str.assert_called_once()
def test_cache_salt_event_hash_walk_is_iterative(self):
root = SimpleNamespace(
key=_HashKey(array("q")),
parent=None,
hash_value=[],
event_hash_value=None,
)
node = root
path = []
for token_id in range(1, 1102):
node = SimpleNamespace(
key=_HashKey(array("q", [token_id]), cache_salt="tenant-a"),
parent=node,
hash_value=None,
event_hash_value=None,
)
path.append(node)
result = compute_node_event_hash_values(node, page_size=1)
self.assertEqual(len(result), 1)
self.assertTrue(all(item.event_hash_value is not None for item in path))
def test_parent_hash_is_used_only_when_parent_has_nonempty_key_and_hash(self): def test_parent_hash_is_used_only_when_parent_has_nonempty_key_and_hash(self):
parent = MagicMock() parent = MagicMock()
parent.key = _HashKey(array("q", range(1, 17))) parent.key = _HashKey(array("q", range(1, 17)))
@@ -0,0 +1,38 @@
"""Unit tests for fail-closed C++ radix-cache request validation."""
import importlib
import sys
import types
import unittest
from unittest.mock import patch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestRadixCacheCppCacheSalt(CustomTestCase):
def test_cache_salt_is_rejected_without_loading_cpp_extension(self):
extension_name = "sglang.srt.mem_cache.cpp_radix_tree.radix_tree"
module_name = "sglang.srt.mem_cache.radix_cache_cpp"
fake_extension = types.ModuleType(extension_name)
fake_extension.IOHandle = object
fake_extension.RadixTreeCpp = object
fake_extension.TreeNodeCpp = object
original_module = sys.modules.pop(module_name, None)
try:
with patch.dict(sys.modules, {extension_name: fake_extension}):
module = importlib.import_module(module_name)
module.RadixCacheCpp._reject_cache_salt(None)
with self.assertRaisesRegex(ValueError, "experimental C\\+\\+"):
module.RadixCacheCpp._reject_cache_salt("tenant-a")
finally:
sys.modules.pop(module_name, None)
if original_module is not None:
sys.modules[module_name] = original_module
if __name__ == "__main__":
unittest.main()
@@ -100,6 +100,16 @@ class TestRadixKey(unittest.TestCase):
self.assertEqual(list(key[2:2].token_ids), []) # Empty slice self.assertEqual(list(key[2:2].token_ids), []) # Empty slice
self.assertEqual(list(key[:].token_ids), [1, 2, 3, 4, 5]) # Full slice self.assertEqual(list(key[:].token_ids), [1, 2, 3, 4, 5]) # Full slice
def test_cache_salt_is_preserved_by_slicing(self):
key = RadixKey(
array("q", [1, 2, 3, 4]),
extra_key="classification",
cache_salt="tenant-a",
)
sliced = key[1:3]
self.assertEqual(sliced.extra_key, "classification")
self.assertEqual(sliced.cache_salt, "tenant-a")
def test_getitem_invalid_index(self): def test_getitem_invalid_index(self):
"""Test __getitem__ with invalid indices.""" """Test __getitem__ with invalid indices."""
key = RadixKey(array("q", [1, 2, 3])) key = RadixKey(array("q", [1, 2, 3]))
@@ -424,6 +434,7 @@ class TestRadixCache(unittest.TestCase):
req_pool_idx=0, req_pool_idx=0,
cache_protected_len=0, cache_protected_len=0,
extra_key=None, extra_key=None,
cache_salt=None,
priority=0, priority=0,
last_node=cache.root_node, last_node=cache.root_node,
) )
@@ -579,6 +590,107 @@ class TestRadixCache(unittest.TestCase):
# Non-existent extra_key should not match # Non-existent extra_key should not match
self.assertEqual(len(result4.device_indices), 0) self.assertEqual(len(result4.device_indices), 0)
def test_cache_salt_isolation_is_independent_of_extra_key(self):
cache = RadixCache.create_simulated()
tokens = array("q", [1, 2, 3])
cache.insert(
InsertParams(
key=RadixKey(tokens, extra_key="bc", cache_salt="a"),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
cache.insert(
InsertParams(
key=RadixKey(tokens, extra_key="c", cache_salt="ab"),
value=torch.tensor([40, 50, 60], dtype=torch.int64),
)
)
first = cache.match_prefix(
MatchPrefixParams(key=RadixKey(tokens, extra_key="bc", cache_salt="a"))
)
second = cache.match_prefix(
MatchPrefixParams(key=RadixKey(tokens, extra_key="c", cache_salt="ab"))
)
torch.testing.assert_close(
first.device_indices, torch.tensor([10, 20, 30], dtype=torch.int64)
)
torch.testing.assert_close(
second.device_indices, torch.tensor([40, 50, 60], dtype=torch.int64)
)
def test_cache_salt_is_included_in_store_and_remove_events(self):
mock_allocator = unittest.mock.Mock()
mock_allocator.device = torch.device("cpu")
cache = RadixCache.create_simulated(
mock_allocator=mock_allocator,
page_size=2,
enable_kv_cache_events=True,
)
tokens = array("q", [1, 2, 3, 4])
cache.insert(
InsertParams(
key=RadixKey(tokens, cache_salt="tenant-a"),
value=torch.tensor([10, 20, 30, 40], dtype=torch.int64),
)
)
cache.evict(EvictParams(num_tokens=len(tokens)))
events = cache.take_events()
stored = [event for event in events if isinstance(event, BlockStored)]
removed = [event for event in events if isinstance(event, BlockRemoved)]
self.assertEqual(len(stored), 2)
self.assertTrue(
all(event.metadata.cache_salt == "tenant-a" for event in stored)
)
self.assertEqual(stored[1].parent_block_hash, stored[0].block_hashes[0])
self.assertEqual(
removed[0].block_hashes,
[event.block_hashes[0] for event in stored],
)
unsalted = RadixCache.create_simulated(page_size=2, enable_kv_cache_events=True)
unsalted.insert(InsertParams(key=RadixKey(tokens), value=None))
unsalted_hashes = [
event.block_hashes[0]
for event in unsalted.take_events()
if isinstance(event, BlockStored)
]
self.assertNotEqual(
unsalted_hashes, [event.block_hashes[0] for event in stored]
)
def test_cache_salt_event_hashes_are_preserved_across_node_split(self):
cache = RadixCache.create_simulated(page_size=2, enable_kv_cache_events=True)
original = RadixKey(array("q", [1, 2, 3, 4]), cache_salt="tenant-a")
cache.insert(
InsertParams(
key=original,
value=torch.tensor([10, 20, 30, 40], dtype=torch.int64),
)
)
original_node = cache.match_prefix(
MatchPrefixParams(key=original)
).last_device_node
original_hashes = list(original_node.event_hash_value)
cache.insert(
InsertParams(
key=RadixKey(array("q", [1, 2, 9, 10]), cache_salt="tenant-a"),
value=torch.tensor([10, 20, 90, 100], dtype=torch.int64),
)
)
split_child = cache.match_prefix(
MatchPrefixParams(key=original)
).last_device_node
split_parent = split_child.parent
self.assertEqual(
split_parent.event_hash_value + split_child.event_hash_value,
original_hashes,
)
def test_lock_ref_operations(self): def test_lock_ref_operations(self):
"""Test lock reference counting operations.""" """Test lock reference counting operations."""
cache = RadixCache.create_simulated() cache = RadixCache.create_simulated()
@@ -31,6 +31,7 @@ class _StubReq:
self.origin_input_ids = array("q", token_ids) self.origin_input_ids = array("q", token_ids)
self.output_ids = array("q") self.output_ids = array("q")
self.extra_key = None self.extra_key = None
self.cache_salt = None
self.prefix_indices = None self.prefix_indices = None
self.last_node = None self.last_node = None
self.last_host_node = None self.last_host_node = None
@@ -46,6 +46,7 @@ def _recv(rid, input_ids, max_new_tokens=8):
priority=None, priority=None,
routing_key=None, routing_key=None,
extra_key=None, extra_key=None,
cache_salt=None,
http_worker_ipc=None, http_worker_ipc=None,
time_stats=None, time_stats=None,
) )
@@ -66,6 +66,7 @@ class _FakeReq:
self.origin_input_ids = list(range(committed)) self.origin_input_ids = list(range(committed))
self.output_ids = [] self.output_ids = []
self.extra_key = None self.extra_key = None
self.cache_salt = None
self.last_node = None self.last_node = None
self.cache_protected_len = 0 self.cache_protected_len = 0
self.swa_uuid_for_lock = None self.swa_uuid_for_lock = None
@@ -110,6 +110,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
swa_evicted_seqlen=0, swa_evicted_seqlen=0,
), ),
extra_key=None, extra_key=None,
cache_salt=None,
last_node=tree.root_node, last_node=tree.root_node,
swa_uuid_for_lock=None, swa_uuid_for_lock=None,
swa_prefix_lock_released=False, swa_prefix_lock_released=False,
@@ -607,6 +607,7 @@ class TestSWA(unittest.TestCase):
(req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices (req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
) )
req.extra_key = None req.extra_key = None
req.cache_salt = None
req.last_node = tree.root_node req.last_node = tree.root_node
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.kv.swa_evicted_seqlen = 0 req.kv.swa_evicted_seqlen = 0
@@ -644,6 +645,7 @@ class TestSWA(unittest.TestCase):
(req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2 (req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2
) )
req2.extra_key = None req2.extra_key = None
req2.cache_salt = None
req2.last_node = tree.root_node req2.last_node = tree.root_node
req2.swa_uuid_for_lock = None req2.swa_uuid_for_lock = None
req2.kv.swa_evicted_seqlen = 0 req2.kv.swa_evicted_seqlen = 0