From e86fdf3a3c43e311b7599ecfed7f30a7a2271a94 Mon Sep 17 00:00:00 2001 From: Ethan ZHU Date: Mon, 25 May 2026 11:15:10 +0800 Subject: [PATCH] [Bug Fix][HiCache] TreeNode.get_prefix_hash_values @lru_cache can return mutated list (#26177) Co-authored-by: Zhangheng --- .../sglang/srt/mem_cache/mamba_radix_cache.py | 2 -- python/sglang/srt/mem_cache/radix_cache.py | 2 -- .../unit/mem_cache/test_radix_cache_unit.py | 33 +++++++++++++++++++ 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index c92b6fc7b..be209c6a9 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -22,7 +22,6 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache. import heapq from array import array from collections import defaultdict -from functools import lru_cache from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -148,7 +147,6 @@ class TreeNode: return None return self.hash_value[-1] - @lru_cache(maxsize=1) def get_prefix_hash_values(self, node: "TreeNode") -> List[str]: if node is None or node.hash_value is None: return [] diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 883494068..6e35d1a31 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -28,7 +28,6 @@ import sys import time from array import array from collections import defaultdict -from functools import lru_cache from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union import torch @@ -258,7 +257,6 @@ class TreeNode: return None return self.hash_value[-1] - @lru_cache(maxsize=1) def get_prefix_hash_values(self, node: TreeNode) -> List[str]: if node is None or node.hash_value is None: return [] 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 9df04c12e..263670fb7 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -39,6 +39,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( InsertParams, MatchPrefixParams, ) +from sglang.srt.mem_cache.mamba_radix_cache import TreeNode as MambaTreeNode from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode # Test constants @@ -226,6 +227,38 @@ class TestTreeNode(unittest.TestCase): node.hash_value = ["hash1", "hash2", "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): """Test less than comparison based on last_access_time.""" node1 = TreeNode()