From 896acc8860ff60dd0471328bf598692bc38ee242 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 21 Aug 2026 01:23:26 -0700 Subject: [PATCH] [Fix] Clear full-to-SWA mapping with `index_fill_` to avoid a blocking H2D copy (#35773) --- python/sglang/srt/mem_cache/allocator/swa.py | 13 +++++-- .../srt/mem_cache/multi_ended_allocator.py | 4 ++ .../sglang/srt/mem_cache/swa_radix_cache.py | 2 +- .../unified_cache/components/swa_component.py | 2 +- .../unit/mem_cache/test_swa_unittest.py | 37 +++++++++++++++++++ .../test_unified_radix_cache_unittest.py | 4 +- 6 files changed, 53 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index b429b40d0..2f70b6216 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -280,9 +280,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): alloc_full_indices[-swa_tail_len:], alloc_swa_indices ) if swa_tail_len < extend_num_tokens: - self.full_to_swa_index_mapping[ - alloc_full_indices[:-swa_tail_len].to(torch.int64) - ] = 0 + self.clear_full_to_swa_mapping(alloc_full_indices[:-swa_tail_len]) return alloc_full_indices def alloc_decode( @@ -345,6 +343,13 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): swa_indices = swa_indices.to(self.full_to_swa_index_mapping.dtype) self.full_to_swa_index_mapping[full_indices] = swa_indices + def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: + if full_indices.numel() == 0: + return + # index_fill_ passes the 0 as a kernel argument; mapping[idx] = 0 copies a + # host-resident scalar and blocks until the stream drains. + self.full_to_swa_index_mapping.index_fill_(0, full_indices.to(torch.int64), 0) + def free_swa(self, free_index: torch.Tensor): if free_index.numel() == 0: return @@ -361,7 +366,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): swa_indices = self.full_to_swa_index_mapping[mapping_indices] swa_indices = swa_indices[swa_indices > 0] self.swa_attn_allocator.free(swa_indices) - self.full_to_swa_index_mapping[mapping_indices] = 0 + self.clear_full_to_swa_mapping(mapping_indices) def free_group_begin(self): super().free_group_begin() diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 98a7bb190..93b2a286b 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -2498,6 +2498,10 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): """ return + def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: + # Paired with set_full_to_swa_mapping: shared mode has no mapping tensor. + return + # -- free-group -- def free_group_begin(self) -> None: diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 4ecbd68b2..910551eee 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -1301,7 +1301,7 @@ class SWARadixCache(BasePrefixCache): allocator = self.token_to_kv_pool_allocator swa_value = allocator.translate_loc_from_full_to_swa(incoming_full) allocator.set_full_to_swa_mapping(node.value, swa_value) - allocator.full_to_swa_index_mapping[incoming_full.to(torch.int64)] = 0 + allocator.clear_full_to_swa_mapping(incoming_full) allocator.full_attn_allocator.free(incoming_full) node.swa_tombstone = False diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index ac49edef7..74bb8de90 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -1148,7 +1148,7 @@ class SWAComponent(TreeComponent): # freeing only the incoming full, then store the swa on the node. swa_value = self._translate_full_to_swa(action.incoming_full) alloc.set_full_to_swa_mapping(action.kept_full, swa_value) - alloc.full_to_swa_index_mapping[action.incoming_full.to(torch.int64)] = 0 + alloc.clear_full_to_swa_mapping(action.incoming_full) alloc.full_attn_allocator.free(action.incoming_full) self.tree_core.set_component_device_value( action.node_id, self.component_type, swa_value diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index ac1ccd04e..d75ee3ae9 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -226,6 +226,43 @@ class TestSWA(unittest.TestCase): allocator.free_swa(full_indices[1:2]) self.assertEqual(allocator.swa_available_size(), 16) + @unittest.skipUnless(torch.cuda.is_available(), "sync detection needs CUDA") + def test_clearing_the_mapping_does_not_synchronize(self): + """Clearing the full-to-SWA mapping must not block the stream; writing a + host-resident scalar into it does. + """ + _, allocator, _ = _build_swa_tree(is_eagle=False) + full_indices = _swa_alloc(allocator, 4) + mapping = allocator.full_to_swa_index_mapping + + # Warm up outside the window: a first-time cudaMalloc can synchronize on + # its own, which the detector would report as this call's fault. + allocator.clear_full_to_swa_mapping(full_indices) + + def sync_error(fn): + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + fn() + except RuntimeError as exc: + return exc + finally: + torch.cuda.set_sync_debug_mode("default") + torch.cuda.synchronize() + return None + + # Gate on the pre-fix form: a detector blind to this sync class would pass + # the assert below no matter how the mapping is cleared. + pre_fix_error = sync_error( + lambda: mapping.__setitem__(full_indices.to(torch.int64), 0) + ) + if pre_fix_error is None: + self.skipTest("sync debug mode does not flag a blocking H2D copy here") + + self.assertIsNone( + sync_error(lambda: allocator.clear_full_to_swa_mapping(full_indices)) + ) + def test_free_swa_group_owns_deferred_indices(self): _, allocator, _ = _build_swa_tree( is_eagle=False, diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index b05c1d110..b692b01b0 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -6147,9 +6147,7 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase): alloc.translate_loc_from_full_to_swa.assert_called_once_with(incoming_full) alloc.set_full_to_swa_mapping.assert_called_once_with(kept_full, swa_value) # the incoming full's stale mapping is cleared, then its slot freed (full-only) - key, val = alloc.full_to_swa_index_mapping.__setitem__.call_args.args - self.assertTrue(torch.equal(key, incoming_full)) - self.assertEqual(val, 0) + alloc.clear_full_to_swa_mapping.assert_called_once_with(incoming_full) alloc.full_attn_allocator.free.assert_called_once_with(incoming_full) alloc.free.assert_not_called() cache.tree_core.set_component_device_value.assert_called_once_with(