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,
)