[Perf] unified-memory: Free SWA segments without a host sync

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

@@ -29,6 +29,13 @@
 _SWA_PEER_MAPPED = Invariant("swa.peer_mapped", Bucket.FATAL_UNCONTAINABLE, IsTrue())
 # free_full leaves the mapping alone, so a live entry would strand its SWA peer.
 _SWA_PEER_RELEASED = Invariant("swa.peer_released", Bucket.GUARD, IsTrue())
+# The unified allocator's segment frees trust the caller's bound/released split.
+UNIFIED_SWA_SEGMENT_BOUND = Invariant(
+    "unified_swa.segment_bound", Bucket.FATAL_UNCONTAINABLE, IsTrue()
+)
+UNIFIED_SWA_SEGMENT_RELEASED = Invariant(
+    "unified_swa.segment_released", Bucket.FATAL_UNCONTAINABLE, IsTrue()
+)
 
 
 class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):

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

@@ -26,7 +26,11 @@
 
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.allocator.base import MambaFullCacheDonor
-from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.swa import (
+    UNIFIED_SWA_SEGMENT_BOUND,
+    UNIFIED_SWA_SEGMENT_RELEASED,
+    SWATokenToKVPoolAllocator,
+)
 from sglang.srt.mem_cache.allocator.unified_sub_pool import (
     FloatMultiEndedAllocator,
     MultiEndedAllocator,
@@ -41,6 +45,7 @@
 from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
 from sglang.srt.runtime_context import get_disagg, get_memory
 from sglang.srt.utils.common import get_num_new_pages
+from sglang.srt.utils.invariants import expect
 from sglang.srt.utils.rank_consensus_checker import rank_consensus
 
 logger = logging.getLogger(__name__)
@@ -172,6 +177,7 @@ def __init__(
 
         self.free_group = None
         self.free_page_reps_group: Optional[List[torch.Tensor]] = None
+        self.full_page_reps_group: List[torch.Tensor] = []
         self.full_free_group: List[torch.Tensor] = []
         # Empty (not None) for the leak checker.
         self.free_pages = torch.empty(0, dtype=torch.int64, device=device)
@@ -612,27 +618,25 @@ def free_swa(self, free_index: torch.Tensor) -> None:
         self.swa_attn_allocator.clear_inverse_history()
 
     def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
-        """free_swa() for a kv-row segment: `start_pos` promises a contiguous
-        ascending range, so page reps come from stride arithmetic, not `torch.unique`."""
+        """free_swa() for a kv-row segment whose swa pages are all bound: page reps
+        come from `start_pos`, so it is fixed-shape at every page size."""
         if free_index is None or free_index.numel() == 0:
             return
-        if self.page_size == 1:
-            self.free_swa(free_index)
-            return
-        ps = self.page_size
         reps = self.swa_attn_allocator._page_reps(
             free_index.detach().to(torch.int64), start_pos
         )
-        # Keep only pages still bound on swa; freeing a tombstoned one would
-        # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink.
-        rep_pages = reps // ps
-        swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages]
-        live_reps = reps[swa_v2p_pages > 0]
-        if live_reps.numel() == 0:
-            return
-        self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps)
+        self._free_swa_page_reps(reps)
         self.swa_attn_allocator.clear_inverse_history()
 
+    def _free_swa_page_reps(self, reps: torch.Tensor) -> None:
+        v_pages = reps // self.page_size
+        expect(
+            UNIFIED_SWA_SEGMENT_BOUND,
+            self.swa_attn_allocator.virtual_to_physical[v_pages] > 0,
+            msg="caller wants free_full_segment",
+        )
+        self.swa_attn_allocator.free(reps, _pages=v_pages)
+
     def free_full(self, free_index: torch.Tensor) -> None:
         """Release the full-physical page and the virtual id, leaving the swa
         side alone -- the caller already tombstoned it (`swa.v2p_page == -1`)."""
@@ -645,15 +649,21 @@ def free_full(self, free_index: torch.Tensor) -> None:
         self.full_attn_allocator.clear_inverse_history()
 
     def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        """Release a kv-row segment's full side; its swa side is already tombstoned."""
         if free_index is None or free_index.numel() == 0:
             return
-        if self.page_size == 1:
-            # token == page: free_full already frees by exact ids, no dedup.
-            self.free_full(free_index)
-            return
-        # The swa v2p is the mapping, so a tombstoned swa page drops out of the
-        # two-sided segment path by itself; full-only is the same call.
-        self.free_segment(free_index, start_pos=start_pos)
+        reps = self.full_attn_allocator._page_reps(
+            free_index.detach().to(torch.int64), start_pos
+        )
+        expect(
+            UNIFIED_SWA_SEGMENT_RELEASED,
+            self.swa_attn_allocator.virtual_to_physical[reps // self.page_size] <= 0,
+            msg="caller wants free_segment",
+        )
+        if self.free_page_reps_group is None:
+            self._release_page_reps((), (reps,))
+        else:
+            self.full_page_reps_group.append(reps)
 
     def set_full_to_swa_mapping(
         self, full_indices: torch.Tensor, swa_indices: torch.Tensor
@@ -676,44 +686,44 @@ def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None:
     def free_group_begin(self) -> None:
         BaseTokenToKVPoolAllocator.free_group_begin(self)
         self.free_page_reps_group = []
+        self.full_page_reps_group = []
         self.full_free_group = []
 
     def free_group_end(self) -> None:
         pending, self.free_page_reps_group = self.free_page_reps_group, None
+        full_pending, self.full_page_reps_group = self.full_page_reps_group, []
         full_free_group, self.full_free_group = self.full_free_group, []
         BaseTokenToKVPoolAllocator.free_group_end(self)
         if full_free_group:
             self.full_attn_allocator.free(torch.cat(full_free_group))
             self.full_attn_allocator.clear_inverse_history()
-        if pending:
-            self._release_page_reps(pending)
+        if pending or full_pending:
+            self._release_page_reps(pending, full_pending)
 
     def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
-        """Fixed-shape counterpart of `free()`; see `MultiEndedAllocator._page_reps`.
-        Both sides share one page-rep derivation instead of dedup'ing twice."""
+        """Fixed-shape counterpart of `free()` for a segment whose swa pages are all
+        bound; see `MultiEndedAllocator._page_reps`."""
         if free_index is None or free_index.numel() == 0:
             return
-        if self.page_size == 1:
-            self.free(free_index)
-            return
         reps = self.full_attn_allocator._page_reps(
             free_index.detach().to(torch.int64), start_pos
         )
         if self.free_page_reps_group is None:
-            self._release_page_reps((reps,))
+            self._release_page_reps((reps,), ())
         else:
             self.free_page_reps_group.append(reps)
 
-    def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None:
-        reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces))
-        v_pages = reps // self.page_size
-        # Same tombstone filter as `free`, but at PAGE granularity (page_size
-        # times smaller): `> 0` strict -- -1 = tombstoned, 0 = padding sink.
-        swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages]
-        live_pages = v_pages[swa_v2p_pages > 0]
-        if live_pages.numel() > 0:
-            self.swa_attn_allocator.free(live_pages * self.page_size, _pages=live_pages)
-        self.full_attn_allocator.free(reps, _pages=v_pages)
+    def _release_page_reps(
+        self, both: Sequence[torch.Tensor], full_only: Sequence[torch.Tensor]
+    ) -> None:
+        """Free page reps: `both` on the swa and full sides, `full_only` on full."""
+        if both:
+            self._free_swa_page_reps(
+                both[0] if len(both) == 1 else torch.cat(tuple(both))
+            )
+        pieces = (*both, *full_only)
+        reps = pieces[0] if len(pieces) == 1 else torch.cat(pieces)
+        self.full_attn_allocator.free(reps, _pages=reps // self.page_size)
         self.full_attn_allocator.clear_inverse_history()
         self.swa_attn_allocator.clear_inverse_history()
 
@@ -738,6 +748,7 @@ def clear(self) -> None:
         self.swa_attn_allocator.clear()
         self.free_group = None
         self.free_page_reps_group = None
+        self.full_page_reps_group = []
         self.full_free_group = []
 
     # -- Lazy compaction hooks --
@@ -1217,7 +1228,12 @@ def evict_to_free_tokens(
         required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens
         _flush_deferred_free_group(
             self,
-            (self.free_group, self.free_page_reps_group, self.full_free_group),
+            (
+                self.free_group,
+                self.free_page_reps_group,
+                self.full_free_group,
+                self.full_page_reps_group,
+            ),
         )
         reclaim_plan = self.reclaim_plan(
             num_tokens,
@@ -1724,7 +1740,12 @@ def flush_deferred_full_frees(self) -> None:
         """Expose grouped composite frees while preserving the group scope."""
         _flush_deferred_free_group(
             self,
-            (self.free_group, self.free_page_reps_group, self.full_free_group),
+            (
+                self.free_group,
+                self.free_page_reps_group,
+                self.full_page_reps_group,
+                self.full_free_group,
+            ),
         )
 
     def full_tokens_before_mamba_recheck(self, target_size: int) -> int:

test/registered/unit/mem_cache/test_unified_free_no_host_sync.py

@@ -25,17 +25,22 @@
 """
 
 import ast
+import importlib.util
 import inspect
 import textwrap
 import unittest
 from unittest import mock
 
 import torch
 from test_multi_ended_allocator import TestPagedMultiEndedAllocator as _PagedFixture
+from torch.utils._python_dispatch import TorchDispatchMode
 
+from sglang.srt.environ import InvariantCheckLevel, envs
 from sglang.srt.mem_cache.allocator import unified_hybrid_swa, unified_mamba
 from sglang.srt.mem_cache.allocator import unified_sub_pool as mea
 from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
+from sglang.srt.runtime_context import publish, reset_context
+from sglang.srt.server_args import ServerArgs
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
@@ -51,6 +56,39 @@ def _paged_allocator(lazy: bool):
     return full
 
 
+_aten = torch.ops.aten
+
+# Output shape depends on tensor values: on CUDA each one copies a count to the
+# host and blocks it until the stream drains.
+_SHAPE_SYNC_OPS = {
+    _aten.nonzero.default,
+    _aten.masked_select.default,
+    _aten._local_scalar_dense.default,
+    _aten._unique2.default,
+    _aten.unique_dim.default,
+    _aten.unique_consecutive.default,
+}
+_INDEX_OPS = {_aten.index.Tensor, _aten.index_put.default, _aten.index_put_.default}
+
+
+class _ShapeSyncRecorder(TorchDispatchMode):
+    """Records every host-syncing op a call dispatches. A boolean-mask index is
+    one: its `nonzero` runs inside the index kernel, below this mode."""
+
+    def __init__(self):
+        super().__init__()
+        self.hits = []
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+        if func in _SHAPE_SYNC_OPS or (
+            func in _INDEX_OPS
+            and any(i is not None and i.dtype == torch.bool for i in args[1])
+        ):
+            self.hits.append(str(func))
+        return func(*args, **kwargs)
+
+
 # --------------------------------------------------------------------------
 # 1. tombstone scatters
 # --------------------------------------------------------------------------
@@ -331,6 +369,33 @@ class TestUnifiedSwaFullSideGroup(unittest.TestCase):
     """`free_full` inside a free group must defer and land at `free_group_end`;
     the composite carries that pile itself, not through the SWA parent's hooks."""
 
+    def setUp(self):
+        reset_context()
+        self.addCleanup(reset_context)
+        publish(ServerArgs(model_path="dummy"), role="tokenizer")
+
+    def test_allocation_reclaim_flushes_full_only_segment_frees(self):
+        """FULL-only queued pages already satisfy this allocation. Ignoring
+        their queue makes two-pool reclaim report insufficient capacity."""
+        from test_unified_swa_shared_virtual_ids import _build
+
+        alloc = _build(64, 32, 2, 4)
+        indices = alloc.alloc(8)
+        alloc.free_swa(indices)
+        before = alloc.full_available_size()
+        alloc.free_group_begin()
+        self.addCleanup(alloc.free_group_end)
+        alloc.free_full_segment(indices, start_pos=0)
+        self.assertEqual(alloc.full_available_size(), before)
+        cache = mock.Mock()
+        cache.supports_prefix_sharing.return_value = True
+        cache.full_evictable_size.return_value = 0
+        cache.swa_evictable_size.return_value = 0
+        cache.evict_for_alloc.side_effect = AssertionError("no tree eviction needed")
+
+        self.assertTrue(alloc.evict_to_free_tokens(cache, before + 4, swa_num_tokens=0))
+        self.assertEqual(alloc.full_available_size(), before + 8)
+
     def test_full_only_frees_defer_until_group_end(self):
         from test_unified_swa_shared_virtual_ids import _build
 
@@ -356,7 +421,7 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
 
     PS = 4
 
-    def _swa_composite(self, lazy=True):
+    def _swa_composite(self, lazy=True, ps=None):
         from test_multi_ended_allocator import _FakeKVCache, _make_mha_spec
 
         from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
@@ -369,7 +434,7 @@ def _swa_composite(self, lazy=True):
             sub_pool_specs=[full, swa],
             device="cpu",
             enable_memory_saver=False,
-            page_size=self.PS,
+            page_size=ps or self.PS,
         )
 
         class _KV:
@@ -386,7 +451,7 @@ def attach_allocators(self, **kwargs):
             device="cpu",
             full_max_total_num_tokens=64,
             swa_max_total_num_tokens=64,
-            page_size=self.PS,
+            page_size=ps or self.PS,
             need_sort=False,
             forward_stream=None,
             lazy_compaction=lazy,
@@ -458,13 +523,58 @@ def test_start_pos_path_matches_the_fallback_end_state(self):
                     a2.swa_attn_allocator.schedulable_available_size(),
                 )
 
-    def test_double_ratchet_is_filtered_not_crashed(self):
-        """Freeing an already-tombstoned range again must no-op through the
-        liveness filter (radix eviction and the ratchet can overlap)."""
+    def test_segment_frees_dispatch_no_data_dependent_shape(self):
+        """Bug regression: the ratchet, a request's release and a tree eviction
+        must not filter by liveness on the device; a boolean-mask index blocks
+        the scheduler on the in-flight forward at every prefill."""
+        recorder = _ShapeSyncRecorder()
+        with recorder:
+            torch.arange(4)[torch.arange(4) > 1]
+        self.assertEqual(len(recorder.hits), 1, "the recorder went blind")
+        for ps in (1, self.PS):
+            with self.subTest(page_size=ps):
+                alloc = self._swa_composite(lazy=True, ps=ps)
+                full_before = alloc.full_available_size()
+                swa_before = alloc.swa_available_size()
+                v = alloc.alloc(8 * ps)
+                recorder = _ShapeSyncRecorder()
+                with recorder:
+                    alloc.free_swa_segment(v[: 2 * ps], start_pos=0)
+                    alloc.free_group_begin()
+                    alloc.free_full_segment(v[: 2 * ps], start_pos=0)
+                    alloc.free_segment(v[2 * ps : 5 * ps], start_pos=2 * ps)
+                    alloc.free_group_end()
+                    alloc.free_segment(v[5 * ps :], start_pos=5 * ps)
+                self.assertEqual(recorder.hits, [])
+                self.assertEqual(alloc.full_available_size(), full_before)
+                self.assertEqual(alloc.swa_available_size(), swa_before)
+
+    def test_a_double_ratchet_is_reported(self):
+        """A segment free trusts its caller that the swa pages are bound (the
+        ratchet's cursor only moves up); a second free must be loud under
+        strict checks, not corrupt the swa hole list."""
         alloc = self._swa_composite(lazy=True)
         v = alloc.alloc(4 * self.PS)
         alloc.free_swa_segment(v, start_pos=0)
-        alloc.free_swa_segment(v, start_pos=0)  # all tombstoned -> filtered to empty
+        with envs.SGLANG_INVARIANT_CHECK.override(int(InvariantCheckLevel.STRICT)):
+            with self.assertRaises(RuntimeError):
+                alloc.free_swa_segment(v, start_pos=0)
+
+    def test_a_full_only_free_of_a_bound_swa_page_is_reported(self):
+        """`free_full_segment` releases the full side only; a swa page still
+        bound under it would leak, so the caller wanted `free_segment`."""
+        alloc = self._swa_composite(lazy=True)
+        v = alloc.alloc(4 * self.PS)
+        with envs.SGLANG_INVARIANT_CHECK.override(int(InvariantCheckLevel.STRICT)):
+            with self.assertRaises(RuntimeError):
+                alloc.free_full_segment(v, start_pos=0)
+
+    def test_importing_the_module_again_registers_nothing_twice(self):
+        """A fixture that patches `sys.modules` drops the modules first imported
+        under the patch, so the next test imports them again; the invariant
+        registry refuses a name it already holds."""
+        spec = importlib.util.find_spec(unified_hybrid_swa.__name__)
+        spec.loader.exec_module(importlib.util.module_from_spec(spec))
 
 
 @unittest.skipUnless(