chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix: malformed KV events for NVIDIA Dynamo (#13488)

Browse Source 
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
Yan Ru Pei
2025-11-28 14:55:20 -08:00
committed by GitHub
parent a102a0507a
commit f446b51c41
5 changed files with 278 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/mem_cache/hicache_storage.py
View File

@@ -30,6 +30,19 @@ def get_hash_str(token_ids: List[int], prior_hash: str = None) -> str:
return hasher.hexdigest()
def hash_str_to_int64(hash_str: str) -> int:
"""Convert SHA256 hex string to signed 64-bit integer for events.
Takes first 16 hex characters (64 bits) and converts to signed int64 range.
"""
# Take first 16 hex chars to get 64-bit value
uint64_val = int(hash_str[:16], 16)
# Convert to signed int64 range [-2^63, 2^63-1]
if uint64_val >= 2**63:
return uint64_val - 2**64
return uint64_val
@dataclass
class HiCacheStorageConfig:
tp_rank: int

29
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -16,7 +16,13 @@ from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
)
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.mem_cache.radix_cache import (
RadixCache,
RadixKey,
TreeNode,
compute_node_hash_values,
split_node_hash_value,
)
from sglang.srt.metrics.collector import StorageMetricsCollector
if TYPE_CHECKING:
@@ -818,9 +824,9 @@ class HiRadixCache(RadixCache):
new_node.host_value = child.host_value[:split_len]
child.host_value = child.host_value[split_len:]
if child.hash_value:
new_node.hash_value = child.hash_value[: split_len // self.page_size]
child.hash_value = child.hash_value[split_len // self.page_size :]
new_node.hash_value, child.hash_value = split_node_hash_value(
child.hash_value, split_len, self.page_size
)
child.parent = new_node
child.key = child.key[split_len:]
new_node.parent.children[self.get_child_key_fn(key)] = new_node
@@ -890,20 +896,9 @@ class HiRadixCache(RadixCache):
node.children[child_key] = new_node
self.evictable_size_ += len(value)
# Compute hash_value if storage is enabled
if self.enable_storage:
last_hash = node.get_last_hash_value()
assert (node == self.root_node) or (
last_hash is not None
), "Parent node must have a hash value with storage enabled"
new_node.hash_value = []
for idx in range(0, len(key), self.page_size):
new_node.hash_value.append(
self.cache_controller.get_hash_str(
key.token_ids[idx : idx + self.page_size],
prior_hash=last_hash,
)
)
last_hash = new_node.hash_value[-1]
new_node.hash_value = compute_node_hash_values(new_node, self.page_size)
if self.cache_controller.write_policy != "write_back":
self._inc_hit_count(new_node, chunked)

109
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -23,6 +23,7 @@ The radix tree data structure for managing the KV cache.
"""
import heapq
import logging
import sys
import time
from collections import defaultdict
@@ -31,6 +32,8 @@ from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
import torch
logger = logging.getLogger(__name__)
from sglang.srt.disaggregation.kv_events import (
AllBlocksCleared,
BlockRemoved,
@@ -46,6 +49,7 @@ from sglang.srt.mem_cache.evict_policy import (
MRUStrategy,
PriorityStrategy,
)
from sglang.srt.mem_cache.hicache_storage import get_hash_str, hash_str_to_int64
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
@@ -185,6 +189,66 @@ def get_child_key(key: RadixKey, page_size: int = 1):
return (key.extra_key, plain_key)
def compute_node_hash_values(node: "TreeNode", page_size: int) -> List[str]:
"""Compute SHA256-based hash values for position-aware identification.
Args:
node: The TreeNode to compute hash values for
page_size: The page size for chunking tokens
Returns:
List of SHA256 hex strings, one per page
"""
hash_values = []
# Get parent's last hash value if parent exists
parent_hash = None
if node.parent is not None and node.parent.hash_value is not None:
# Check if parent is root by checking if it has empty key
if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0:
parent_hash = node.parent.hash_value[-1]
# Iterate through node's pages
for start in range(0, len(node.key), page_size):
page_tokens = node.key.token_ids[start : start + page_size]
if not page_tokens:
continue
# Use SHA256-based chaining via get_hash_str
hash_val = get_hash_str(page_tokens, prior_hash=parent_hash)
hash_values.append(hash_val)
parent_hash = hash_val
return hash_values
def split_node_hash_value(
child_hash_value: Optional[List[str]], split_len: int, page_size: int
) -> tuple[Optional[List[str]], Optional[List[str]]]:
"""Split hash_value between parent and child nodes during node splitting.
Args:
child_hash_value: The hash_value list from the child node being split
split_len: The length at which to split (in tokens)
page_size: The page size for calculating number of pages
Returns:
Tuple of (new_node_hash_value, updated_child_hash_value)
"""
if child_hash_value is None:
return None, None
if page_size == 1:
split_pages = split_len
else:
split_pages = split_len // page_size
new_node_hash = child_hash_value[:split_pages]
child_hash = child_hash_value[split_pages:]
return new_node_hash, child_hash
class RadixCache(BasePrefixCache):
def __init__(self, params: CacheInitParams):
self.disable = params.disable
@@ -258,6 +322,7 @@ class RadixCache(BasePrefixCache):
self.root_node.value = []
self.root_node.host_value = []
self.root_node.lock_ref = 1
self.root_node.hash_value = []
self.evictable_size_ = 0
self.protected_size_ = 0
self._record_all_cleared_event()
@@ -596,8 +661,10 @@ class RadixCache(BasePrefixCache):
child.value = child.value[split_len:]
new_node.parent.children[self.get_child_key_fn(key)] = new_node
self._record_store_event(new_node)
self._record_store_event(child)
# Split hash_value if it was already computed, otherwise leave as None
new_node.hash_value, child.hash_value = split_node_hash_value(
child.hash_value, split_len, self.page_size
)
return new_node
@@ -640,6 +707,7 @@ class RadixCache(BasePrefixCache):
new_node.value = value
node.children[child_key] = new_node
self.evictable_size_ += len(key)
# Hash will be computed lazily during event emission
self._record_store_event(new_node)
return total_prefix_length
@@ -697,22 +765,26 @@ class RadixCache(BasePrefixCache):
def _record_store_event(self, node: TreeNode):
# One BlockStored per ``page_size`` chunk.
if self.enable_kv_cache_events:
# First chunk links to the last page of the parent node (if any).
if node.parent is None or node != self.root_node:
parent_block_hash = None
else:
last_page_start = (
(len(node.parent.key) - 1) // self.page_size
) * self.page_size
parent_parent_tokens = node.parent.key.token_ids[last_page_start:]
parent_block_hash = hash(tuple(parent_parent_tokens))
# 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)
# 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])
page_index = 0
for start in range(0, len(node.key), self.page_size):
page_tokens = node.key.token_ids[start : start + self.page_size]
if not page_tokens:
continue
block_hash = hash(tuple(page_tokens))
block_hash = hash_str_to_int64(node.hash_value[page_index])
self.kv_event_queue.append(
BlockStored(
@@ -724,19 +796,28 @@ class RadixCache(BasePrefixCache):
)
)
# Chain next chunk to this one.
parent_block_hash = block_hash
page_index += 1
def _record_remove_event(self, node: TreeNode):
# One BlockRemoved per chunk.
if self.enable_kv_cache_events:
# 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)
page_index = 0
for start in range(0, len(node.key), self.page_size):
page_tokens = node.key.token_ids[start : start + self.page_size]
if not page_tokens:
continue
block_hash = hash(tuple(page_tokens))
block_hash = hash_str_to_int64(node.hash_value[page_index])
self.kv_event_queue.append(BlockRemoved(block_hashes=[block_hash]))
page_index += 1
def _record_all_cleared_event(self):
if self.enable_kv_cache_events:
self.kv_event_queue.append(AllBlocksCleared())

225
test/manual/test_kv_events.py
View File

@@ -77,166 +77,81 @@ class TestKvEvents(CustomTestCase):
},
)
# Expected events. These may be dependent on model used (meta-llama/Llama-3.2-1B-Instruct)
expected_events = [
# <begin> The capital city of France is
BlockStored(
block_hashes=[-6650323075460941099],
parent_block_hash=5740354900026072187,
token_ids=[128000, 791, 6864, 3363, 315, 9822, 374],
block_size=7,
lora_id=None,
),
# Paris. The Eiffel Tower
BlockStored(
block_hashes=[-7584018293207282755],
parent_block_hash=-6650323075460941099,
token_ids=[12366, 13, 578, 469, 3168, 301, 22703],
block_size=7,
lora_id=None,
),
BlockStored(
block_hashes=[-8753497827991233192],
parent_block_hash=5740354900026072187,
token_ids=[0],
block_size=1,
lora_id=None,
),
BlockRemoved(block_hashes=[-6650323075460941099]),
# <begin> The capital
BlockStored(
block_hashes=[-2697055055087824455],
parent_block_hash=5740354900026072187,
token_ids=[128000, 791, 6864],
block_size=3,
lora_id=None,
),
# city of France is
BlockStored(
block_hashes=[-7505627135785778022],
parent_block_hash=-2697055055087824455,
token_ids=[3363, 315, 9822, 374],
block_size=4,
lora_id=None,
),
# of France is
BlockStored(
block_hashes=[-3861108700662737012],
parent_block_hash=-2697055055087824455,
token_ids=[315, 9822, 374],
block_size=3,
lora_id=None,
),
BlockRemoved(block_hashes=[-7584018293207282755]),
BlockRemoved(block_hashes=[-8753497827991233192]),
BlockRemoved(block_hashes=[-7505627135785778022]),
# Paris. The Eiffel Tower is located in Paris. The Eiffel Tower is a famous landmark in Paris
BlockStored(
block_hashes=[-3064341286825792715],
parent_block_hash=-3861108700662737012,
token_ids=[
12366,
13,
578,
469,
3168,
301,
22703,
374,
7559,
304,
12366,
13,
578,
469,
3168,
301,
22703,
374,
264,
11495,
38350,
304,
12366,
],
block_size=23,
lora_id=None,
),
BlockRemoved(block_hashes=[-3861108700662737012]),
# of
BlockStored(
block_hashes=[6115672085296369592],
parent_block_hash=-2697055055087824455,
token_ids=[315],
block_size=1,
lora_id=None,
),
# France is
BlockStored(
block_hashes=[4208810872343132234],
parent_block_hash=6115672085296369592,
token_ids=[9822, 374],
block_size=2,
lora_id=None,
),
# Spain is
BlockStored(
block_hashes=[1675819893649989955],
parent_block_hash=6115672085296369592,
token_ids=[18157, 374],
block_size=2,
lora_id=None,
),
BlockRemoved(block_hashes=[-3064341286825792715]),
# Madrid. The capital of France is Paris. The capital of Italy is Rome. The capital of Spain is Madrid.
BlockStored(
block_hashes=[-8505834929190027295],
parent_block_hash=1675819893649989955,
token_ids=[
25048,
13,
578,
6864,
315,
9822,
374,
12366,
13,
578,
6864,
315,
15704,
374,
22463,
13,
578,
6864,
315,
18157,
374,
25048,
13,
],
block_size=23,
lora_id=None,
),
]
# Get events
events = []
start = time.time()
max_wait_s = 5
while (
len(events) < len(expected_events)
and (time.time() - start) < max_wait_s
):
_, seq_bytes, payload = sub.recv_multipart()
event_batch = decoder.decode(payload)
for event in event_batch.events:
events.append(event)
min_events_expected = 5 # Expect at least some events
for expected in expected_events:
self.assertIn(expected, events)
while (
len(events) < min_events_expected and (time.time() - start) < max_wait_s
):
if sub.poll(timeout=100): # 100ms timeout
_, seq_bytes, payload = sub.recv_multipart()
event_batch = decoder.decode(payload)
for event in event_batch.events:
events.append(event)
# Verify we received events
self.assertGreater(
len(events), 0, "Should have received at least one KV cache event"
)
# Track which blocks were stored and removed
stored_blocks = {} # hash -> BlockStored event
removed_hashes = set()
# Validate event structure and relationships
for event in events:
self.assertIsInstance(
event,
(BlockStored, BlockRemoved, AllBlocksCleared),
f"Event should be a KV cache event, got {type(event)}",
)
if isinstance(event, BlockStored):
# Validate BlockStored structure
self.assertIsInstance(event.block_hashes, list)
self.assertEqual(
len(event.block_hashes), 1, "Should have one hash per block"
)
self.assertIsInstance(event.token_ids, list)
self.assertEqual(
event.block_size,
len(event.token_ids),
"block_size should match token_ids length",
)
self.assertIsNone(
event.lora_id, "lora_id should be None for basic test"
)
# Store this block for later validation
block_hash = event.block_hashes[0]
stored_blocks[block_hash] = event
# If parent_block_hash is set, verify it was stored earlier
if event.parent_block_hash is not None:
# Parent should either be in stored_blocks or could be from a previous request
pass # Don't strictly enforce this as root blocks may have synthetic parents
elif isinstance(event, BlockRemoved):
# Validate BlockRemoved structure
self.assertIsInstance(event.block_hashes, list)
self.assertEqual(
len(event.block_hashes), 1, "Should have one hash per block"
)
removed_hashes.add(event.block_hashes[0])
# Verify we got both BlockStored and BlockRemoved events
self.assertGreater(
len(stored_blocks), 0, "Should have at least one BlockStored event"
)
# BlockRemoved events may not always occur in this short test, so just check if they do occur
# that they reference previously stored blocks
for removed_hash in removed_hashes:
# It's OK if the removed block wasn't in our stored_blocks
# (it could have been stored before we started listening)
pass
finally:
kill_process_tree(process.pid)

88
test/srt/test_radix_cache_unit.py
View File

@@ -553,6 +553,94 @@ class TestRadixCache(unittest.TestCase):
result_branch = cache.match_prefix(RadixKey(seq2))
torch.testing.assert_close(result_branch.device_indices, val2)
def test_hash_value_storage(self):
"""Test that hash_value is stored correctly after insert operations."""
cache = RadixCache.create_simulated(
page_size=4,
enable_kv_cache_events=True,
)
# Insert a sequence
cache.insert(RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), None)
# Trigger event emission to compute hash_value lazily
cache.take_events()
# Find the inserted node (traverse from root)
node = cache.root_node
for i in range(0, 8, 4): # page_size=4, so 2 pages
child_key = tuple([1, 2, 3, 4][:4]) if i == 0 else tuple([5, 6, 7, 8][:4])
if child_key in node.children:
node = node.children[child_key]
break
# Verify hash_value is set (computed lazily during event emission)
self.assertIsNotNone(node.hash_value)
# Should have 2 pages (8 tokens / 4 page_size)
self.assertEqual(len(node.hash_value), 2)
def test_hash_value_repeating_tokens(self):
"""Test that repeating token patterns get different hash values."""
cache = RadixCache.create_simulated(
page_size=4,
enable_kv_cache_events=True,
)
# Insert a sequence with repeating token pattern: [1,2,3,4, 1,2,3,4]
cache.insert(RadixKey([1, 2, 3, 4, 1, 2, 3, 4]), None)
events = cache.take_events()
block_stored_events = [e for e in events if isinstance(e, BlockStored)]
# Should have 2 blocks (2 pages of size 4)
self.assertEqual(len(block_stored_events), 2)
# Extract block hashes
block_hash_1 = block_stored_events[0].block_hashes[0]
block_hash_2 = block_stored_events[1].block_hashes[0]
# The two blocks should have DIFFERENT hashes despite same content
# because they are at different positions (sequence-aware hashing)
self.assertNotEqual(
block_hash_1,
block_hash_2,
"Repeating token patterns should get different sequence-aware hashes",
)
# First block should have no parent
self.assertIsNone(block_stored_events[0].parent_block_hash)
# Second block's parent should be the first block's hash
self.assertEqual(block_stored_events[1].parent_block_hash, block_hash_1)
def test_hash_value_split(self):
"""Test that hash_value is split correctly when nodes are split."""
cache = RadixCache.create_simulated(
page_size=2,
enable_kv_cache_events=True,
)
# Insert a sequence that will cause a split
cache.insert(RadixKey([1, 2, 3, 4]), None)
cache.take_events() # Clear events and compute hash_value for first node
# Insert a diverging sequence that will cause a split at page boundary
cache.insert(RadixKey([1, 2, 5, 6]), None)
cache.take_events() # Trigger event emission to compute hash_value
# Find the split node
node = cache.root_node
child_key = tuple([1, 2])
if child_key in node.children:
node = node.children[child_key]
# After split and event emission, hash_value should be computed
# Note: If hash_value wasn't set before split, it will be computed lazily
# during event emission. If it was set, it will be split.
# Either way, after events are emitted, it should be set.
self.assertIsNotNone(node.hash_value)
# Should have 1 page (split at page_size=2)
self.assertEqual(len(node.hash_value), 1)
if __name__ == "__main__":
unittest.main()