minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/encode_receiver.py
+2
View File
@@ -2104,6 +2104,8 @@ class MMReceiverBase(ABC):
if self.scheduler.metrics_reporter.enable_metrics
else None
),
extra_key=recv_req.extra_key,
cache_salt=recv_req.cache_salt,
http_worker_ipc=recv_req.http_worker_ipc,
dllm_config=self.scheduler.dllm_config,
)
python/sglang/srt/disaggregation/kv_events.py
+20
View File
@@ -86,6 +86,12 @@ class StorageMedium(str, enum.Enum):
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:
"""
OffloadedState represents the state of a KV cache block offloaded to the hicache.
@@ -112,6 +118,16 @@ class BlockStored(KVCacheEvent):
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):
block_hashes: list[int]
medium: Optional[str] = None
@@ -122,6 +138,10 @@ class AllBlocksCleared(KVCacheEvent):
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]]
python/sglang/srt/entrypoints/engine.py
+6
View File
@@ -394,6 +394,8 @@ class Engine(EngineScoreMixin, EngineBase):
session_params: Optional[Dict] = None,
priority: Optional[int] = None,
session_id: Optional[str] = None,
*,
cache_salt: Optional[Union[List[str], str]] = None,
) -> Union[Dict, Iterator[Dict]]:
"""
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,
video_data=video_data,
mm_hashes=mm_hashes,
cache_salt=cache_salt,
return_logprob=return_logprob,
logprob_start_len=logprob_start_len,
top_logprobs_num=top_logprobs_num,
@@ -500,6 +503,8 @@ class Engine(EngineScoreMixin, EngineBase):
session_params: Optional[Dict] = None,
priority: Optional[int] = None,
session_id: Optional[str] = None,
*,
cache_salt: Optional[Union[List[str], str]] = None,
) -> Union[Dict, AsyncIterator[Dict]]:
"""
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,
video_data=video_data,
mm_hashes=mm_hashes,
cache_salt=cache_salt,
return_logprob=return_logprob,
logprob_start_len=logprob_start_len,
top_logprobs_num=top_logprobs_num,
python/sglang/srt/entrypoints/openai/protocol.py
+3 -3
View File
@@ -386,7 +386,7 @@ class CompletionRequest(BaseModel):
# For request id
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
# Cache salt for request caching
cache_salt: Optional[Union[List[str], str]] = None
@@ -848,7 +848,7 @@ class ChatCompletionRequest(BaseModel):
# For request id
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
# Cache salt for request caching
cache_salt: Optional[Union[List[str], str]] = None
@@ -1536,7 +1536,7 @@ class ResponsesRequest(BaseModel):
priority: int = Field(default=0, description="Request priority")
extra_key: Optional[str] = Field(
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(
default=None, description="Cache salt for request caching"
python/sglang/srt/entrypoints/openai/serving_base.py
-13
View File
@@ -148,19 +148,6 @@ class OpenAIServingBase(ABC):
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
def _convert_to_internal_request(
self,
python/sglang/srt/entrypoints/openai/serving_chat.py
+2 -1
View File
@@ -1018,7 +1018,8 @@ class OpenAIServingChat(OpenAIServingBase):
routed_experts_start_len=request.routed_experts_start_len,
rid=request.rid,
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,
priority=request.priority,
routing_key=self.extract_routing_key(raw_request),
python/sglang/srt/entrypoints/openai/serving_completions.py
+2 -1
View File
@@ -128,7 +128,8 @@ class OpenAIServingCompletion(OpenAIServingBase):
return_prompt_token_ids=request.return_token_ids,
rid=request.rid,
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,
routing_key=self.extract_routing_key(raw_request),
custom_labels=custom_labels,
python/sglang/srt/entrypoints/openai/serving_responses.py
+3 -1
View File
@@ -454,7 +454,8 @@ class OpenAIServingResponses(OpenAIServingChat):
stream=request.stream,
rid=request.request_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=request.background and not request.stream,
require_reasoning=require_reasoning,
@@ -2563,6 +2564,7 @@ class OpenAIServingResponses(OpenAIServingChat):
rid=request_id,
session_id=adapted_request.session_id,
extra_key=adapted_request.extra_key,
cache_salt=adapted_request.cache_salt,
return_logprob=adapted_request.return_logprob,
logprob_start_len=adapted_request.logprob_start_len,
top_logprobs_num=adapted_request.top_logprobs_num,
python/sglang/srt/managers/io_struct.py
+42 -3
View File
@@ -287,7 +287,7 @@ class GenerateReqInput:
# Priority for the request
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
# 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].
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):
"""Generate a new request ID and return it."""
if isinstance(self.rid, list):
@@ -488,6 +491,14 @@ class GenerateReqInput:
self.token_ids_logprob = None
if self.return_sampling_mask is None:
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):
"""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_custom_logit_processor(num)
self._normalize_extra_key(num)
self._normalize_cache_salt(num)
self._normalize_bootstrap_params(num)
def _expand_inputs(self, num):
@@ -703,16 +715,39 @@ class GenerateReqInput:
if self.extra_key is None:
return
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):
if len(self.extra_key) != self.batch_size:
raise ValueError(
"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
else:
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):
"""Normalize bootstrap parameters for batch processing."""
# Normalize bootstrap_host
@@ -837,6 +872,7 @@ class GenerateReqInput:
max_thinking_tokens=self.max_thinking_tokens,
priority=self.priority,
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,
custom_labels=self.custom_labels,
return_bytes=self.return_bytes,
@@ -925,7 +961,7 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
# Priority for the request
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
# 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]]
time_stats: Optional[PickleWrapper] = None
# Cache namespace used to isolate otherwise-identical prefixes.
cache_salt: Optional[str] = None
def wrap_pickle_fields(self):
self.mm_inputs = wrap_as_pickle(self.mm_inputs)
self.mm_data_mooncake = wrap_as_pickle(self.mm_data_mooncake)
python/sglang/srt/managers/schedule_batch.py
+5 -1
View File
@@ -855,6 +855,7 @@ class Req(ReqDllmMixin):
return_pooled_hidden_states: bool = False,
multi_item_delimiter_indices: Optional[List[int]] = None,
session_id: Optional[str] = None,
cache_salt: Optional[str] = None,
):
# Input and output info
self.rid = rid
@@ -924,13 +925,14 @@ class Req(ReqDllmMixin):
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:
extra_key = (
extra_key or ""
) + lora_id # lora_id is concatenated to the extra key
self.extra_key = extra_key
self.cache_salt = cache_salt or None
self.lora_id = lora_id
self.routing_key = routing_key
@@ -1359,6 +1361,7 @@ class Req(ReqDllmMixin):
token_ids=token_ids_to_match,
extra_key=self.extra_key,
limit=key_limit,
cache_salt=self.cache_salt,
),
req=self,
cow_mamba=cow_mamba,
@@ -1776,6 +1779,7 @@ class Req(ReqDllmMixin):
"bootstrap_room": self.bootstrap_room,
"priority": self.priority,
"extra_key": self.extra_key,
"cache_salt": self.cache_salt,
"routing_key": self.routing_key,
"disagg_prefill_dp_rank": self.disagg_prefill_dp_rank,
}
python/sglang/srt/managers/schedule_policy.py
+17 -3
View File
@@ -155,7 +155,12 @@ def match_prefix_for_req(
match_result = tree_cache.match_prefix(
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,
req=req if include_req else None,
)
@@ -319,6 +324,7 @@ class SchedulePolicy:
for r in waiting_queue:
prefix_ids = r.origin_input_ids + r.output_ids
extra_key = r.extra_key
cache_salt = r.cache_salt
match_result = match_prefix_for_req(
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:
match_result = self.waiting_queue_radix_tree.match_prefix(
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():
@@ -350,7 +360,11 @@ class SchedulePolicy:
# Insert with a dummy key
self.waiting_queue_radix_tree.insert(
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),
)
)
python/sglang/srt/managers/scheduler.py
+1
View File
@@ -2417,6 +2417,7 @@ class Scheduler(
),
routing_key=recv_req.routing_key,
extra_key=recv_req.extra_key,
cache_salt=recv_req.cache_salt,
http_worker_ipc=recv_req.http_worker_ipc,
dllm_config=self.dllm_config,
time_stats=recv_req.time_stats,
python/sglang/srt/managers/tokenizer_manager.py
+1
View File
@@ -1399,6 +1399,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
disagg_prefill_dp_rank=obj.disagg_prefill_dp_rank,
priority=obj.priority,
extra_key=obj.extra_key,
cache_salt=obj.cache_salt,
routing_key=obj.routing_key,
token_type_ids=token_type_ids,
need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs,
python/sglang/srt/mem_cache/events.py
+37 -16
View File
@@ -23,9 +23,12 @@ from sglang.srt.disaggregation.kv_events import (
AllBlocksCleared,
BlockRemoved,
BlockStored,
BlockStoredMetadata,
BlockStoredWithMetadata,
StorageMedium,
)
from sglang.srt.mem_cache.utils import (
compute_node_event_hash_values,
compute_node_hash_values,
hash_str_to_int64,
)
@@ -43,15 +46,22 @@ class KVCacheEventMixin:
# Compute hash_value lazily if not already set
if node.hash_value is None:
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
parent_block_hash = None
if node.parent is not None and node.parent != self.root_node:
if (
node.parent.hash_value is not None
and len(node.parent.hash_value) > 0
):
parent_block_hash = hash_str_to_int64(node.parent.hash_value[-1])
if node.key.cache_salt is not None:
parent_hash_values = node.parent.event_hash_value
assert parent_hash_values is not None
else:
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
logical_len = len(node.key)
@@ -67,18 +77,24 @@ class KVCacheEventMixin:
else:
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(
BlockStored(
block_hashes=[block_hash],
parent_block_hash=parent_block_hash,
token_ids=page_tokens,
block_size=len(page_tokens),
lora_id=None,
medium=medium,
event_args = {
"block_hashes": [block_hash],
"parent_block_hash": parent_block_hash,
"token_ids": page_tokens,
"block_size": len(page_tokens),
"lora_id": None,
"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
page_index += 1
@@ -94,6 +110,11 @@ class KVCacheEventMixin:
# Compute hash_value lazily if not already set (must match what was stored)
if node.hash_value is None:
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 = []
logical_len = len(node.key)
@@ -103,7 +124,7 @@ class KVCacheEventMixin:
if end <= start:
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
if block_hashes:
python/sglang/srt/mem_cache/hiradix_cache.py
+11 -1
View File
@@ -963,7 +963,12 @@ class HiRadixCache(RadixCache):
token_ids = []
for n in chain:
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):
hash_value = []
@@ -1474,6 +1479,7 @@ class HiRadixCache(RadixCache):
new_input_tokens,
extra_key=last_host_node.key.extra_key,
is_bigram=self.is_eagle,
cache_salt=last_host_node.key.cache_salt,
).page_aligned(self.page_size)
if len(prefetch_key) < self.prefetch_threshold:
return 0
@@ -1773,6 +1779,7 @@ class HiRadixCache(RadixCache):
new_input_tokens,
extra_key=last_host_node.key.extra_key,
is_bigram=self.is_eagle,
cache_salt=last_host_node.key.cache_salt,
)
# align the number of fetching tokens to the 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(
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.key = child.key[split_len:]
new_node.parent.children[key.child_key(self.page_size)] = new_node
python/sglang/srt/mem_cache/mamba_radix_cache.py
+22 -3
View File
@@ -104,6 +104,8 @@ class TreeNode:
self.host_value = None
# store hash values of each pages
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:
# 1. prev has greater last_access_time
@@ -629,7 +631,11 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
result = self.insert(
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,
mamba_value=mamba_value,
prev_prefix_len=req.cache_protected_len,
@@ -739,7 +745,11 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
result = self.insert(
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,
mamba_value=mamba_value_donated,
prev_prefix_len=req.cache_protected_len,
@@ -752,7 +762,13 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
# The prefix indices could be updated, reuse it
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 = (
match_result.device_indices,
@@ -1214,6 +1230,9 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
new_node.hash_value, child.hash_value = split_node_hash_value(
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
# parent first so that parent is after child in the lru list
python/sglang/srt/mem_cache/pure_swa_radix_cache.py
+4 -1
View File
@@ -89,7 +89,10 @@ class PureSWARadixCache(RadixCache):
]
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)
keys_len = len(radix_key)
python/sglang/srt/mem_cache/radix_cache.py
+40 -7
View File
@@ -59,7 +59,7 @@ if TYPE_CHECKING:
class RadixKey:
"""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__(
self,
@@ -67,11 +67,17 @@ class RadixKey:
extra_key: Optional[str] = None,
is_bigram: bool = False,
limit: Optional[int] = None,
cache_salt: Optional[str] = None,
):
# token ids sequence (raw ints in both modes)
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
# 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)
self.is_bigram = is_bigram
# 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);
# empty slice -> empty raw tokens (not a dangling boundary token).
raw = self.token_ids[start : stop + 1] if stop > start else array("q")
return RadixKey(raw, self.extra_key, is_bigram=True)
return RadixKey(self.token_ids[start:stop], self.extra_key)
return RadixKey(
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:
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:
if page_size == 1:
@@ -157,6 +172,11 @@ class RadixKey:
f"RadixKey operations require matching extra_key, but got "
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:
"""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))
else:
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)
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
# store hash values of each pages
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
self.priority = priority
@@ -455,7 +479,10 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
]
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)
key_len = len(radix_key)
values = kv_indices[:key_len].to(dtype=torch.int64, copy=True)
@@ -496,7 +523,10 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
]
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)
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(
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
python/sglang/srt/mem_cache/radix_cache_cpp.py
+10
View File
@@ -33,6 +33,13 @@ logger = logging.getLogger(__name__)
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__(
self,
params: CacheInitParams,
@@ -100,6 +107,7 @@ class RadixCacheCpp(BasePrefixCache):
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
key = params.key
self._reject_cache_salt(key.cache_salt)
device_indices_vec, host_indices_length, node_gpu, node_cpu = (
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
):
"""Cache request when it finishes."""
self._reject_cache_salt(req.cache_salt)
assert req.req_pool_idx is not None
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
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):
"""Cache request when it is unfinished."""
self._reject_cache_salt(req.cache_salt)
assert req.req_pool_idx is not None
token_ids = req.get_fill_ids()
prefill_len = len(token_ids) # prefill only (maybe chunked)
python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py
+13 -2
View File
@@ -216,7 +216,12 @@ class FlexKVRadixCache(RadixCache):
else:
token_ids_snap = token_ids
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,
)
return MatchResult(
@@ -410,7 +415,13 @@ class FlexKVRadixCache(RadixCache):
# Anchor on the new last_device_node so FlexKV's lock matches
# the node we'll later unlock when the store completes.
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
if new_last_node is None:
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
+13 -2
View File
@@ -229,7 +229,12 @@ class LMCRadixCache(RadixCache):
if token_ids is key.token_ids:
token_ids = token_ids[:]
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()),
)
return MatchResult(
@@ -463,7 +468,13 @@ class LMCRadixCache(RadixCache):
# Use super() to avoid a redundant LOOKUP — we only need new_last_node from radix.
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
assert new_last_node is not None
python/sglang/srt/mem_cache/swa_radix_cache.py
+20 -2
View File
@@ -80,6 +80,8 @@ class TreeNode:
self.host_value = None
# store hash values of each page
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:
# 1. prev has greater last_access_time
@@ -473,7 +475,10 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
]
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_len = len(radix_key)
values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True)
@@ -523,7 +528,10 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
]
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)
values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True)
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.extra_key,
is_bigram=node.key.is_bigram,
cache_salt=node.key.cache_salt,
)
node.value = torch.cat([node.value, child.value])
node.children = child.children
@@ -1033,6 +1042,12 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
node.hash_value = list(node.hash_value) + list(child.hash_value)
else:
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)
if not child.swa_tombstone:
@@ -1104,6 +1119,9 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
new_node.hash_value, child.hash_value = split_node_hash_value(
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
# parent first so that parent is after child in the lru list
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+12 -3
View File
@@ -117,6 +117,8 @@ class UnifiedTreeNode:
self.last_access_time = get_and_increase_time_counter()
self.creation_time = get_and_increase_time_counter()
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.priority = priority
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(
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:
component.redistribute_on_node_split(new_parent=new_node, child=child)
@@ -1867,10 +1872,14 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
comp_xfers[comp.component_type] = t
return kv_xfer, comp_xfers
def prefetch_anchor_info(self, node_id: NodeId) -> Optional[str]:
"""The anchor node's key extra_key."""
def prefetch_anchor_info(
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)
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(
self, node: UnifiedTreeNode, write_back: bool = False
python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py
+4 -2
View File
@@ -421,8 +421,10 @@ class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC):
...
@abstractmethod
def prefetch_anchor_info(self, node_id: NodeId) -> Optional[str]:
"""The anchor node's key extra_key."""
def prefetch_anchor_info(
self, node_id: NodeId
) -> tuple[Optional[str], Optional[str]]:
"""The anchor node's key extra_key and cache_salt."""
...
@abstractmethod
python/sglang/srt/mem_cache/unified_radix_cache.py
+7 -2
View File
@@ -697,7 +697,10 @@ class UnifiedRadixCache(BasePrefixCache):
kv_indices = kv_indices[:effective_cache_len]
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_len = len(radix_key)
values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True)
@@ -790,6 +793,7 @@ class UnifiedRadixCache(BasePrefixCache):
token_ids[:effective_cache_len],
req.extra_key,
is_bigram=self.tree_core.is_eagle,
cache_salt=req.cache_salt,
).page_aligned(self.page_size)
page_aligned_len = len(radix_key)
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:
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(
new_input_tokens,
extra_key=extra_key,
is_bigram=self.tree_core.is_eagle,
cache_salt=cache_salt,
).page_aligned(self.page_size)
prefetch_length = len(prefetch_key)
if (
python/sglang/srt/mem_cache/utils.py
+49
View File
@@ -13,6 +13,7 @@
# ==============================================================================
"""Common utilities."""
import hashlib
from typing import Any, Callable, List, Optional, Tuple
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
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(
child_hash_value: Optional[List[str]], split_len: int, page_size: int
) -> tuple[Optional[List[str]], Optional[List[str]]]:
python/sglang/srt/session/session_controller.py
+1
View File
@@ -312,6 +312,7 @@ class Session:
priority=req.priority,
routing_key=req.routing_key,
extra_key=req.extra_key,
cache_salt=req.cache_salt,
http_worker_ipc=req.http_worker_ipc,
time_stats=req.time_stats,
)
test/registered/unit/disaggregation/test_encode_receiver.py
+61
View File
@@ -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()
test/registered/unit/disaggregation/test_kv_events.py
+46
View File
@@ -9,7 +9,14 @@ the router can subscribe per replica (the `dp_size` it reads from
import unittest
import msgspec
from sglang.srt.disaggregation.kv_events import (
BlockStored,
BlockStoredMetadata,
BlockStoredWithMetadata,
KVEventBatch,
StorageMedium,
ZmqEventPublisher,
select_kv_publisher_dp_rank,
)
@@ -94,5 +101,44 @@ class TestSelectKvPublisherDpRank(CustomTestCase):
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__":
unittest.main()
test/registered/unit/entrypoints/openai/test_serving_chat.py
+4
View File
@@ -316,6 +316,8 @@ class ServingChatTestCase(unittest.TestCase):
input_ids=[101, 102, 103],
stop=["STOP"],
return_prompt_token_ids=True,
cache_salt="tenant-a",
extra_key="classification",
)
with patch(
@@ -329,6 +331,8 @@ class ServingChatTestCase(unittest.TestCase):
self.assertEqual(adapted.input_ids, [101, 102, 103])
self.assertTrue(adapted.return_prompt_token_ids)
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()
def test_kimi_k3_usage_excludes_assistant_generation_stub(self):
test/registered/unit/entrypoints/openai/test_serving_completions.py
+23
View File
@@ -65,6 +65,29 @@ class ServingCompletionTestCase(unittest.TestCase):
internal, _ = self.sc._convert_to_internal_request(req)
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 ----------
def test_echo_with_list_of_strings_streaming(self):
req = CompletionRequest(
test/registered/unit/managers/test_io_struct.py
+47
View File
@@ -519,6 +519,53 @@ class TestGenerateReqInputNormalization(CustomTestCase):
req.normalize_batch_and_arguments()
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):
"""Test normalization of logprob-related parameters."""
# Test single example
test/registered/unit/managers/test_priority_scheduling_disaggregation.py
+2
View File
@@ -244,6 +244,7 @@ class TestDecodePreallocQueueRebootstrapPayload(unittest.TestCase):
bootstrap_room=7,
priority=10,
extra_key=None,
cache_salt=None,
routing_key=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.assertEqual(payload["sampling_params"]["max_new_tokens"], 1)
self.assertEqual(payload["bootstrap_room"], 7)
self.assertIsNone(payload["cache_salt"])
# 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
# the decode-side override, so neither belongs in the payload.
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+1
View File
@@ -47,6 +47,7 @@ def _make_req(
req.logprob_start_len = -1
req.positional_embed_overrides = None
req.extra_key = None
req.cache_salt = None
req.mamba_pool_idx = None
req.sampling_params = SimpleNamespace(max_new_tokens=128, ignore_eos=False)
return req
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+1
View File
@@ -78,6 +78,7 @@ class MockReq:
self.cache_protected_len = cache_protected_len
self.last_node = last_node
self.extra_key = None
self.cache_salt = None
self.prefix_indices = torch.empty(0, dtype=torch.int64)
self.priority = 0
self.kv_committed_len = len(fill_ids)
test/registered/unit/mem_cache/test_mem_cache_utils.py
+65 -3
View File
@@ -5,6 +5,7 @@ import sys
import types
import unittest
from array import array
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from sglang.srt.mem_cache.evict_policy import (
@@ -17,6 +18,7 @@ from sglang.srt.mem_cache.evict_policy import (
SLRUStrategy,
)
from sglang.srt.mem_cache.utils import (
compute_node_event_hash_values,
compute_node_hash_values,
get_eviction_strategy,
get_hash_str,
@@ -54,9 +56,10 @@ def _legacy_page_hashes(key, page_size, prior_hash=None):
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.is_bigram = is_bigram
self.cache_salt = cache_salt
def __len__(self):
if self.is_bigram:
@@ -68,8 +71,12 @@ class _HashKey:
start = index.start or 0
stop = index.stop if index.stop is not None else len(self)
if self.is_bigram:
return _HashKey(self.token_ids[start : stop + 1], is_bigram=True)
return _HashKey(self.token_ids[start:stop])
return _HashKey(
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:
return (self.token_ids[index], self.token_ids[index + 1])
return self.token_ids[index]
@@ -298,6 +305,7 @@ class TestComputeNodeHashValues(unittest.TestCase):
node = MagicMock()
node.key = key
node.parent = parent
node.event_hash_value = None
if parent is not None:
parent.hash_value = parent_hash_values
return node
@@ -318,6 +326,60 @@ class TestComputeNodeHashValues(unittest.TestCase):
_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):
parent = MagicMock()
parent.key = _HashKey(array("q", range(1, 17)))
test/registered/unit/mem_cache/test_radix_cache_cpp_unit.py
+38
View File
@@ -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()
test/registered/unit/mem_cache/test_radix_cache_unit.py
+112
View File
@@ -100,6 +100,16 @@ class TestRadixKey(unittest.TestCase):
self.assertEqual(list(key[2:2].token_ids), []) # Empty 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):
"""Test __getitem__ with invalid indices."""
key = RadixKey(array("q", [1, 2, 3]))
@@ -424,6 +434,7 @@ class TestRadixCache(unittest.TestCase):
req_pool_idx=0,
cache_protected_len=0,
extra_key=None,
cache_salt=None,
priority=0,
last_node=cache.root_node,
)
@@ -579,6 +590,107 @@ class TestRadixCache(unittest.TestCase):
# Non-existent extra_key should not match
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):
"""Test lock reference counting operations."""
cache = RadixCache.create_simulated()
test/registered/unit/mem_cache/test_radix_force_miss.py
+1
View File
@@ -31,6 +31,7 @@ class _StubReq:
self.origin_input_ids = array("q", token_ids)
self.output_ids = array("q")
self.extra_key = None
self.cache_salt = None
self.prefix_indices = None
self.last_node = None
self.last_host_node = None
test/registered/unit/mem_cache/test_session_token_share_unit.py
+1
View File
@@ -46,6 +46,7 @@ def _recv(rid, input_ids, max_new_tokens=8):
priority=None,
routing_key=None,
extra_key=None,
cache_salt=None,
http_worker_ipc=None,
time_stats=None,
)
test/registered/unit/mem_cache/test_streaming_session_unit.py
+1
View File
@@ -66,6 +66,7 @@ class _FakeReq:
self.origin_input_ids = list(range(committed))
self.output_ids = []
self.extra_key = None
self.cache_salt = None
self.last_node = None
self.cache_protected_len = 0
self.swa_uuid_for_lock = None
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+1
View File
@@ -110,6 +110,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
swa_evicted_seqlen=0,
),
extra_key=None,
cache_salt=None,
last_node=tree.root_node,
swa_uuid_for_lock=None,
swa_prefix_lock_released=False,
test/registered/unit/mem_cache/test_swa_unittest.py
+2
View File
@@ -607,6 +607,7 @@ class TestSWA(unittest.TestCase):
(req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
)
req.extra_key = None
req.cache_salt = None
req.last_node = tree.root_node
req.swa_uuid_for_lock = None
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.extra_key = None
req2.cache_salt = None
req2.last_node = tree.root_node
req2.swa_uuid_for_lock = None
req2.kv.swa_evicted_seqlen = 0