From 385903b0acd69455cb688b5cb5e3afcc0fd91598 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Wed, 12 Aug 2026 16:14:04 -0700 Subject: [PATCH] feat: add cache salt support to KV cache events (#30827) Signed-off-by: jthomson04 --- .../srt/disaggregation/encode_receiver.py | 2 + python/sglang/srt/disaggregation/kv_events.py | 20 ++++ python/sglang/srt/entrypoints/engine.py | 6 + .../sglang/srt/entrypoints/openai/protocol.py | 6 +- .../srt/entrypoints/openai/serving_base.py | 13 -- .../srt/entrypoints/openai/serving_chat.py | 3 +- .../entrypoints/openai/serving_completions.py | 3 +- .../entrypoints/openai/serving_responses.py | 4 +- python/sglang/srt/managers/io_struct.py | 45 ++++++- python/sglang/srt/managers/schedule_batch.py | 6 +- python/sglang/srt/managers/schedule_policy.py | 20 +++- python/sglang/srt/managers/scheduler.py | 1 + .../sglang/srt/managers/tokenizer_manager.py | 1 + python/sglang/srt/mem_cache/events.py | 53 ++++++--- python/sglang/srt/mem_cache/hiradix_cache.py | 12 +- .../sglang/srt/mem_cache/mamba_radix_cache.py | 25 +++- .../srt/mem_cache/pure_swa_radix_cache.py | 5 +- python/sglang/srt/mem_cache/radix_cache.py | 47 ++++++-- .../sglang/srt/mem_cache/radix_cache_cpp.py | 10 ++ .../storage/flexkv/flexkv_radix_cache.py | 15 ++- .../storage/lmcache/lmc_radix_cache.py | 15 ++- .../sglang/srt/mem_cache/swa_radix_cache.py | 22 +++- .../unified_cache/unified_tree_core.py | 15 ++- .../unified_tree_core_interface.py | 6 +- .../srt/mem_cache/unified_radix_cache.py | 9 +- python/sglang/srt/mem_cache/utils.py | 49 ++++++++ .../sglang/srt/session/session_controller.py | 1 + .../disaggregation/test_encode_receiver.py | 61 ++++++++++ .../unit/disaggregation/test_kv_events.py | 46 +++++++ .../entrypoints/openai/test_serving_chat.py | 4 + .../openai/test_serving_completions.py | 23 ++++ .../unit/managers/test_io_struct.py | 47 ++++++++ ...test_priority_scheduling_disaggregation.py | 2 + .../test_scheduler_chunked_req_gate.py | 1 + .../mem_cache/test_decode_radix_lock_ref.py | 1 + .../unit/mem_cache/test_mem_cache_utils.py | 68 ++++++++++- .../mem_cache/test_radix_cache_cpp_unit.py | 38 ++++++ .../unit/mem_cache/test_radix_cache_unit.py | 112 ++++++++++++++++++ .../unit/mem_cache/test_radix_force_miss.py | 1 + .../test_session_token_share_unit.py | 1 + .../mem_cache/test_streaming_session_unit.py | 1 + .../mem_cache/test_swa_eviction_boundary.py | 1 + .../unit/mem_cache/test_swa_unittest.py | 2 + 43 files changed, 753 insertions(+), 70 deletions(-) create mode 100644 test/registered/unit/disaggregation/test_encode_receiver.py create mode 100644 test/registered/unit/mem_cache/test_radix_cache_cpp_unit.py diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index 95ebf8f4d..16ab5fae0 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -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, ) diff --git a/python/sglang/srt/disaggregation/kv_events.py b/python/sglang/srt/disaggregation/kv_events.py index fa0bdfc0a..5113d4759 100644 --- a/python/sglang/srt/disaggregation/kv_events.py +++ b/python/sglang/srt/disaggregation/kv_events.py @@ -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]] diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 4e4ba29cf..0c4fc805a 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -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, diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index ecb458fc6..8c8fd7888 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -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" diff --git a/python/sglang/srt/entrypoints/openai/serving_base.py b/python/sglang/srt/entrypoints/openai/serving_base.py index 164ccdd5a..aa50f8a71 100644 --- a/python/sglang/srt/entrypoints/openai/serving_base.py +++ b/python/sglang/srt/entrypoints/openai/serving_base.py @@ -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, diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 6919dadf6..dd0e3ab55 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -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), diff --git a/python/sglang/srt/entrypoints/openai/serving_completions.py b/python/sglang/srt/entrypoints/openai/serving_completions.py index 3867a1c8f..d5587765b 100644 --- a/python/sglang/srt/entrypoints/openai/serving_completions.py +++ b/python/sglang/srt/entrypoints/openai/serving_completions.py @@ -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, diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index 81f9c3e77..41ae749a3 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -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, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 9cc0a6272..7d9c94a5d 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 3b9152f4c..3f7465f4f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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, } diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 797a8dfcf..6780d914a 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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), ) ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a91433dd3..c7d8f3cdd 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index d0b893bf5..cd0c33fab 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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, diff --git a/python/sglang/srt/mem_cache/events.py b/python/sglang/srt/mem_cache/events.py index b1eded35e..15da2d40b 100644 --- a/python/sglang/srt/mem_cache/events.py +++ b/python/sglang/srt/mem_cache/events.py @@ -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: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 22daf99c1..1b9111608 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index bd9538a93..a8559db3e 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py index b3d1366df..938e8b5a6 100644 --- a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py @@ -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) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 2916f8c64..f7eebd75f 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index fb082c1f8..46d4c6774 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -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) diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py index 601cc5a66..aeaa90f17 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py @@ -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: diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 88caa246b..f8ce455d0 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 31cedca76..4268130e9 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 94805e2eb..0646734f3 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 8e7e20e20..367c111d7 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index f657897bc..375c17832 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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 ( diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 9cb4b1623..3ce48d3e3 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/python/sglang/srt/mem_cache/utils.py @@ -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]]]: diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index 789b9c3d6..95b6a5271 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -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, ) diff --git a/test/registered/unit/disaggregation/test_encode_receiver.py b/test/registered/unit/disaggregation/test_encode_receiver.py new file mode 100644 index 000000000..f1d654a96 --- /dev/null +++ b/test/registered/unit/disaggregation/test_encode_receiver.py @@ -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() diff --git a/test/registered/unit/disaggregation/test_kv_events.py b/test/registered/unit/disaggregation/test_kv_events.py index d935eb8cf..49db8967d 100644 --- a/test/registered/unit/disaggregation/test_kv_events.py +++ b/test/registered/unit/disaggregation/test_kv_events.py @@ -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() diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 63330b441..0875c0a30 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -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): diff --git a/test/registered/unit/entrypoints/openai/test_serving_completions.py b/test/registered/unit/entrypoints/openai/test_serving_completions.py index 0a7097d8e..3f0b6736f 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_completions.py +++ b/test/registered/unit/entrypoints/openai/test_serving_completions.py @@ -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( diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index 755cbf5ae..9c9896204 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -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 diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index cabe21066..1bd0a16af 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -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. diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 14ed036ec..35676f187 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index e54e69bc5..f7ef02a13 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -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) diff --git a/test/registered/unit/mem_cache/test_mem_cache_utils.py b/test/registered/unit/mem_cache/test_mem_cache_utils.py index b2ff08dcf..395a7965f 100644 --- a/test/registered/unit/mem_cache/test_mem_cache_utils.py +++ b/test/registered/unit/mem_cache/test_mem_cache_utils.py @@ -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))) diff --git a/test/registered/unit/mem_cache/test_radix_cache_cpp_unit.py b/test/registered/unit/mem_cache/test_radix_cache_cpp_unit.py new file mode 100644 index 000000000..48d204d58 --- /dev/null +++ b/test/registered/unit/mem_cache/test_radix_cache_cpp_unit.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 6e3782e87..d9a2808cf 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_radix_force_miss.py b/test/registered/unit/mem_cache/test_radix_force_miss.py index 82c77401f..4e1337428 100644 --- a/test/registered/unit/mem_cache/test_radix_force_miss.py +++ b/test/registered/unit/mem_cache/test_radix_force_miss.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_session_token_share_unit.py b/test/registered/unit/mem_cache/test_session_token_share_unit.py index 0413b5abe..b83c84489 100644 --- a/test/registered/unit/mem_cache/test_session_token_share_unit.py +++ b/test/registered/unit/mem_cache/test_session_token_share_unit.py @@ -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, ) diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index 46cd27433..4ed28eda5 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py index cd8114237..a884f0e44 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -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, diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 52822984d..0b451af5c 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -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