[BUGFIX] fix radix cache memory consumption to avoid OOM (#17191)

Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -839,11 +839,11 @@ class HiRadixCache(RadixCache):
         if child.evicted:
             new_node.value = None
         else:
-            new_node.value = child.value[:split_len]
+            new_node.value = child.value[:split_len].clone()
-            child.value = child.value[split_len:]
+            child.value = child.value[split_len:].clone()
         if child.backuped:
-            new_node.host_value = child.host_value[:split_len]
+            new_node.host_value = child.host_value[:split_len].clone()
-            child.host_value = child.host_value[split_len:]
+            child.host_value = child.host_value[split_len:].clone()
 
         new_node.hash_value, child.hash_value = split_node_hash_value(
             child.hash_value, split_len, self.page_size

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -982,7 +982,7 @@ class MambaRadixCache(BasePrefixCache):
         new_node.full_lock_ref = child.full_lock_ref
         new_node.mamba_lock_ref = 0
         new_node.key = child.key[:split_len]
-        new_node.value = child.value[:split_len]
+        new_node.value = child.value[:split_len].clone()
 
         # child time should be later than parent's time for mamba tombstone
         child.last_access_time = get_last_access_time()
@@ -992,7 +992,7 @@ class MambaRadixCache(BasePrefixCache):
             self.mamba_lru_list.remove_node(child)
         child.parent = new_node
         child.key = child.key[split_len:]
-        child.value = child.value[split_len:]
+        child.value = child.value[split_len:].clone()
         new_node.parent.children[self.get_child_key_fn(key)] = new_node
 
         # insert the new node and child into the lru lists, insert

python/sglang/srt/mem_cache/radix_cache.py

@@ -654,10 +654,10 @@ class RadixCache(BasePrefixCache):
         new_node.parent = child.parent
         new_node.lock_ref = child.lock_ref
         new_node.key = child.key[:split_len]
-        new_node.value = child.value[:split_len]
+        new_node.value = child.value[:split_len].clone()
         child.parent = new_node
         child.key = child.key[split_len:]
-        child.value = child.value[split_len:]
+        child.value = child.value[split_len:].clone()
         new_node.parent.children[self.get_child_key_fn(key)] = new_node
 
         # Split hash_value if it was already computed, otherwise leave as None
@@ -703,7 +703,7 @@ class RadixCache(BasePrefixCache):
             new_node = TreeNode(priority=priority)
             new_node.parent = node
             new_node.key = key
-            new_node.value = value
+            new_node.value = value.clone()
             node.children[child_key] = new_node
             self.evictable_size_ += len(key)
             # Hash will be computed lazily during event emission

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -878,7 +878,7 @@ class SWARadixCache(BasePrefixCache):
         new_node.full_lock_ref = child.full_lock_ref
         new_node.swa_lock_ref = child.swa_lock_ref
         new_node.key = child.key[:split_len]
-        new_node.value = child.value[:split_len]
+        new_node.value = child.value[:split_len].clone()
         # parent inherits the swa_uuid from child for swa lock ref
         new_node.swa_uuid = child.swa_uuid
         child.swa_uuid = None
@@ -891,7 +891,7 @@ class SWARadixCache(BasePrefixCache):
             self.swa_lru_list.remove_node(child)
         child.parent = new_node
         child.key = child.key[split_len:]
-        child.value = child.value[split_len:]
+        child.value = child.value[split_len:].clone()
         new_node.parent.children[self.get_child_key_fn(key)] = new_node
 
         # insert the new node and child into the lru lists, insert

test/registered/attention/test_radix_cache_unit.py

@@ -23,6 +23,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 register_cuda_ci(est_time=5, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=5, suite="stage-b-test-small-1-gpu-amd")
 
+import random
 import time
 import unittest
 import unittest.mock
@@ -647,6 +648,39 @@ class TestRadixCache(unittest.TestCase):
             # Should have 1 page (split at page_size=2)
             self.assertEqual(len(node.hash_value), 1)
 
+    def test_memory_allocated(self):
+        keys, values = [], []
+
+        num_seqs = 10000
+        vocab_size = 1000
+        base_prefix_len = 10000
+        suffix_len = 100
+
+        torch_allocated_before = torch.cuda.memory_allocated()
+
+        # build dataset with common prefix
+        common_prefix = [random.randint(1, vocab_size) for _ in range(base_prefix_len)]
+        for _ in range(num_seqs):
+            suffix = [random.randint(1, vocab_size) for _ in range(suffix_len)]
+            seq = common_prefix + suffix
+            keys.append(seq)
+            values.append(torch.zeros(len(seq), device="cuda", dtype=torch.int32))
+
+        cache: RadixCache = RadixCache.create_simulated()
+
+        for key, value in zip(keys, values):
+            cache.insert(RadixKey(key), value)
+
+        del values
+
+        torch_allocated = torch.cuda.memory_allocated() - torch_allocated_before
+        cache_size_bytes = cache.total_size() * 4
+        print(f"\nCache size (MB): {cache_size_bytes / (1024 * 1024)}")
+        print(f"Torch allocated (MB): {torch_allocated / (1024 * 1024)}")
+
+        # The cache size should be within reasonable bounds of the actual allocated memory.
+        self.assertLess(torch_allocated, cache_size_bytes * 2)
+
 
 if __name__ == "__main__":
     unittest.main()