Fix SWA ownership across grouped frees (#36381)

python/sglang/srt/mem_cache/allocator/swa.py

@@ -354,10 +354,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         if free_index.numel() == 0:
             return
 
-        if not self.is_not_in_free_group:
-            self.swa_free_group.append(self._copy_for_free_group(free_index))
-            return
-
         if self.page_size == 1:
             mapping_indices = free_index
         else:
@@ -365,9 +361,16 @@ 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.clear_full_to_swa_mapping(mapping_indices)
 
+        if not self.is_not_in_free_group:
+            # Resolve ownership now. A cache action later in this group may
+            # install a new mapping for the same full index.
+            self.swa_free_group.append(swa_indices)
+            return
+
+        self.swa_attn_allocator.free(swa_indices)
+
     def free_group_begin(self):
         super().free_group_begin()
         self.swa_free_group = []
@@ -377,7 +380,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         if self.swa_free_group:
             swa_free_group = self.swa_free_group
             self.swa_free_group = []
-            self.free_swa(torch.cat(swa_free_group))
+            self.swa_attn_allocator.free(torch.cat(swa_free_group))
 
     def _expand_to_full_pages(self, indices: torch.Tensor) -> torch.Tensor:
         pages = torch.unique(indices // self.page_size)