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

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

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()