From 737a1183d69ae2b46cb7e04156610d3f6d18c881 Mon Sep 17 00:00:00 2001 From: Yi Zhang <1109276519@qq.com> Date: Sat, 17 Jan 2026 16:47:37 +0800 Subject: [PATCH] [BUGFIX] fix radix cache memory consumption to avoid OOM (#17191) Co-authored-by: Liangsheng Yin --- python/sglang/srt/mem_cache/hiradix_cache.py | 8 ++--- .../sglang/srt/mem_cache/mamba_radix_cache.py | 4 +-- python/sglang/srt/mem_cache/radix_cache.py | 6 ++-- .../sglang/srt/mem_cache/swa_radix_cache.py | 4 +-- .../attention/test_radix_cache_unit.py | 34 +++++++++++++++++++ 5 files changed, 45 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index cb3c8b24b..b3ea8bdff 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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] - child.value = child.value[split_len:] + new_node.value = child.value[:split_len].clone() + child.value = child.value[split_len:].clone() if child.backuped: - new_node.host_value = child.host_value[:split_len] - child.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:].clone() new_node.hash_value, child.hash_value = split_node_hash_value( child.hash_value, split_len, self.page_size diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 598dd348f..acd2fc1cb 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index eb53c8beb..24dc1a741 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index a5d736145..9f2a75fae 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/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 diff --git a/test/registered/attention/test_radix_cache_unit.py b/test/registered/attention/test_radix_cache_unit.py index 9408fd513..4a165cd90 100644 --- a/test/registered/attention/test_radix_cache_unit.py +++ b/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()