[Fix] Clear full-to-SWA mapping with index_fill_ to avoid a blocking H2D copy (#35773)

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

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:

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

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