[Bug Fix][HiCache] TreeNode.get_prefix_hash_values @lru_cache can return mutated list (#26177)

Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
Ethan ZHU
2026-05-25 11:15:10 +08:00
committed by GitHub
co-authored by Zhangheng
parent 821d5f4a5b
commit e86fdf3a3c
3 changed files with 33 additions and 4 deletions
@@ -22,7 +22,6 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache.
import heapq import heapq
from array import array from array import array
from collections import defaultdict from collections import defaultdict
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple from typing import TYPE_CHECKING, List, Optional, Tuple
import torch import torch
@@ -148,7 +147,6 @@ class TreeNode:
return None return None
return self.hash_value[-1] return self.hash_value[-1]
@lru_cache(maxsize=1)
def get_prefix_hash_values(self, node: "TreeNode") -> List[str]: def get_prefix_hash_values(self, node: "TreeNode") -> List[str]:
if node is None or node.hash_value is None: if node is None or node.hash_value is None:
return [] return []
@@ -28,7 +28,6 @@ import sys
import time import time
from array import array from array import array
from collections import defaultdict from collections import defaultdict
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
import torch import torch
@@ -258,7 +257,6 @@ class TreeNode:
return None return None
return self.hash_value[-1] return self.hash_value[-1]
@lru_cache(maxsize=1)
def get_prefix_hash_values(self, node: TreeNode) -> List[str]: def get_prefix_hash_values(self, node: TreeNode) -> List[str]:
if node is None or node.hash_value is None: if node is None or node.hash_value is None:
return [] return []
@@ -39,6 +39,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
InsertParams, InsertParams,
MatchPrefixParams, MatchPrefixParams,
) )
from sglang.srt.mem_cache.mamba_radix_cache import TreeNode as MambaTreeNode
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
# Test constants # Test constants
@@ -226,6 +227,38 @@ class TestTreeNode(unittest.TestCase):
node.hash_value = ["hash1", "hash2", "hash3"] node.hash_value = ["hash1", "hash2", "hash3"]
self.assertEqual(node.get_last_hash_value(), "hash3") self.assertEqual(node.get_last_hash_value(), "hash3")
def test_get_prefix_hash_values_not_shared_across_calls(self):
"""Regression guard for cached mutable prefix hash lists."""
for node_cls in (TreeNode, MambaTreeNode):
with self.subTest(node_cls=node_cls.__module__):
root = node_cls()
n1 = node_cls()
n1.parent = root
n1.hash_value = ["h1"]
n2 = node_cls()
n2.parent = n1
n2.hash_value = ["h2"]
n3 = node_cls()
n3.parent = n2
n3.hash_value = ["h3"]
first = n3.get_prefix_hash_values(n2)
self.assertEqual(first, ["h1", "h2"])
# Downstream storage code extends prefix_keys in place while
# processing pages. A cached list must not be observable by a
# later call.
first += ["h3"]
second = n3.get_prefix_hash_values(n2)
self.assertEqual(second, ["h1", "h2"])
self.assertIsNot(second, first)
n4 = node_cls()
n4.parent = n3
n4.hash_value = ["h4"]
self.assertEqual(n4.get_prefix_hash_values(n3), ["h1", "h2", "h3"])
def test_lt_comparison(self): def test_lt_comparison(self):
"""Test less than comparison based on last_access_time.""" """Test less than comparison based on last_access_time."""
node1 = TreeNode() node1 = TreeNode()