[Perf] unified-memory: Reduce allocation and SWA retirement overhead

python/sglang/kernels/ops/memory/virtual_slot.py

@@ -28,7 +28,7 @@
 # `.item()`, no tensor branching); runs on the scheduler thread.
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["v_pages_ptr"])
 def alloc_bind_inplace_kernel(
     v_pages_ptr,  # in: [N] int64 — virtual page ids
     v2p_ptr,  # in/out: int64 — virtual_to_physical table
@@ -126,7 +126,7 @@ def free_unbind_inplace_kernel(
     tl.store(p2v_ptr + p, -1, mask=mask)
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["v_pages_ptr", "p_pages_ptr"])
 def bind_inplace_kernel(
     v_pages_ptr,  # in: [N] int64 — virtual page ids
     p_pages_ptr,  # in: [N] int64 — physical page ids to bind them to

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

@@ -24,6 +24,7 @@
 import torch
 from torch.profiler import record_function
 
+from sglang.srt.environ import InvariantCheckLevel
 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 (
@@ -45,7 +46,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.invariants import expect, resolve_level
 from sglang.srt.utils.rank_consensus_checker import rank_consensus
 
 logger = logging.getLogger(__name__)
@@ -426,10 +427,10 @@ def _extend_in_virtual_space(
         if not self.ensure_capacity(need_tokens, swa_tokens):
             return None
 
-        # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps
-        # its view after the slice is consumed.
+        # free_virtual_ids is rebound (slice/cat), never mutated in place,
+        # so this prefix view stays valid for SWA after FULL consumes it.
         fa = self.full_attn_allocator
-        new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone()
+        new_virtual_pages = fa.free_virtual_ids[:num_new_pages]
 
         out_indices = fa.alloc_extend(
             prefix_lens,
@@ -630,11 +631,14 @@ def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
 
     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",
-        )
+        # Signal-only invariant: avoid launching its lookup and comparison
+        # when detection is disabled. This invariant has no recovery action.
+        if resolve_level() >= InvariantCheckLevel.WARN:
+            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:
@@ -655,11 +659,15 @@ def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None
         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",
-        )
+        # Signal-only invariant: avoid launching its lookup and comparison
+        # when detection is disabled. This invariant has no recovery action.
+        if resolve_level() >= InvariantCheckLevel.WARN:
+            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:
@@ -1685,10 +1693,32 @@ def _compute_available_size(self) -> int:
         h_s = sa._hole_pages()
         r_f = fa.num_pages - fa.min_page_index - fa._allocated_pages()
         r_s = sa.num_pages - sa.min_page_index - sa._allocated_pages()
+        # Allocator state is stable throughout this query. Snapshot the band
+        # geometry once instead of walking the chain at every binary-search step.
+        full_gap = fa._current_gap_bytes()
+        full_low = fa._byte_low_frontier()
+        full_page_bytes = fa.entry_bytes_per_page
+        swa_chain_low = sa._chain_high_frontier_below_bytes()
+        if sa._is_frontier_transparent():
+            swa_high_band_low = swa_chain_low
+            swa_low_room = 0
+        else:
+            swa_high_band_low = sa._byte_high_frontier()
+            swa_low_room = sa.pages_in_band(
+                low_byte=swa_chain_low, high_byte=sa._byte_low_frontier()
+            )
+
         lo_n, hi_n = 0, min(h_f + r_f, h_s + r_s)
         while lo_n < hi_n:
             mid = (lo_n + hi_n + 1) // 2
-            if self._fits_page_demand(mid, mid):
+            full_bytes = max(0, mid - h_f) * full_page_bytes
+            swa_room = max(
+                swa_low_room,
+                sa.pages_in_band(
+                    low_byte=swa_high_band_low, high_byte=full_low - full_bytes
+                ),
+            )
+            if full_bytes <= full_gap and max(0, mid - h_s) <= swa_room:
                 lo_n = mid
             else:
                 hi_n = mid - 1

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

@@ -1187,10 +1187,10 @@ def alloc_extend(
                 if not _relieve_for_alloc(self, need_tokens):
                     return None
             bs = len(prefix_lens)
-            # Snapshot the virtual pages the kernel will consume, to bind them
-            # to physical pages afterward.
+            # The free list is only rebound, not mutated in place. Retain its
+            # consumed slice for binding without launching a redundant copy.
             if num_new_pages > 0:
-                new_virtual_pages = self.free_virtual_ids[:num_new_pages].clone()
+                new_virtual_pages = self.free_virtual_ids[:num_new_pages]
             else:
                 new_virtual_pages = None
 
@@ -2424,10 +2424,13 @@ def free(
                 free_v_pages = free_v_pages_raw
             else:
                 free_v_pages = torch.unique(free_v_pages_raw // self.page_size)
-            freed_p_pages = self.virtual_to_physical[free_v_pages]
-            # `index_fill_`, never `t[idx] = -1`: see the END free path.
-            self.virtual_to_physical.index_fill_(0, free_v_pages, -1)
-            self.physical_to_virtual.index_fill_(0, freed_p_pages, -1)
+            # The same fused lookup + tombstones as the END free path avoids
+            # three separate launches on each sliding-window retirement.
+            freed_p_pages = free_unbind_inplace(
+                free_v_pages.contiguous(),
+                self.virtual_to_physical,
+                self.physical_to_virtual,
+            )
             if self.is_id_owner:
                 self.free_virtual_ids = torch.cat([self.free_virtual_ids, free_v_pages])
             self._free_phys_pages = torch.cat([self._free_phys_pages, freed_p_pages])

test/registered/kernels/ops/memory/test_virtual_slot_binding.py

@@ -0,0 +1,53 @@
+"""Free-list prefix views must not compile new bind kernels after odd consumes."""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.memory.virtual_slot import (
+    alloc_bind_inplace,
+    alloc_bind_inplace_kernel,
+    bind_inplace,
+    bind_inplace_kernel,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestVirtualSlotBinding(CustomTestCase):
+    def test_prefix_alignment_does_not_specialize_bind_kernels(self):
+        """Taking one int64 page from a free list changes its pointer's 16-byte
+        alignment. Copy-elided prefill must reuse the warmed bind kernel."""
+        virtual = torch.arange(1, 35, device="cuda", dtype=torch.int64)
+        physical = torch.arange(64, 98, device="cuda", dtype=torch.int64)
+        v2p = torch.empty(128, device="cuda", dtype=torch.int64)
+        p2v = torch.empty_like(v2p)
+        for kernel, allocate in (
+            (alloc_bind_inplace_kernel, True),
+            (bind_inplace_kernel, False),
+        ):
+            with self.subTest(kernel=kernel.__name__):
+                cache = kernel.device_caches[torch.cuda.current_device()][0]
+                saved = dict(cache)
+                cache.clear()
+                try:
+                    for voff, poff in ((0, 0), (1, 0), (0, 1), (1, 1)):
+                        v2p.fill_(-1)
+                        p2v.fill_(-1)
+                        v = virtual[voff : voff + 32]
+                        p = physical[poff : poff + 32]
+                        if allocate:
+                            p = alloc_bind_inplace(v, v2p, p2v, start_phys=64)
+                        else:
+                            bind_inplace(v, p, v2p, p2v)
+                        torch.testing.assert_close(v2p[v], p)
+                        torch.testing.assert_close(p2v[p], v)
+                        self.assertEqual(len(cache), 1, (voff, poff))
+                finally:
+                    cache.update(saved)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -18,9 +18,11 @@
 suite runs on CPU apart from the cases that skip themselves without CUDA.
 """
 
-from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=11, suite="base-a-test-cpu")
+register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small")
 
 import contextlib
 import random
@@ -2724,7 +2726,7 @@ def test_drainable_credit_reads_walked_neighbor(self):
         self.assertEqual(fa._peer_drainable_hole_bytes(), 0)
 
 
-class TestFloatMultiEndedAllocator(unittest.TestCase):
+class TestFloatMultiEndedAllocator(CustomTestCase):
     """Holes-first float middle: midpoint placement, in-place hole recycling,
     larger-gap extension, boundary absorption with park-on-empty transparency,
     and the on-demand movers `make_room` and `compact_holes`."""
@@ -2735,7 +2737,7 @@ def setUp(self):
         self.addCleanup(reset_context)
         publish(ServerArgs(model_path="dummy"), role="tokenizer")
 
-    def _build_tri(self, n_state=8, n_float=32, n_full=32):
+    def _build_tri(self, n_state=8, n_float=32, n_full=32, device=_DEV):
         state = _make_mamba_spec("state", "up", layer_num=2)
         fl = _make_mha_spec("swa", "float", layer_num=2)
         full = _make_mha_spec("full", "down", layer_num=2)
@@ -2747,30 +2749,30 @@ def _build_tri(self, n_state=8, n_float=32, n_full=32):
         pool = UnifiedKVPool(
             total_bytes=total,
             sub_pool_specs=[full, fl, state],
-            device=_DEV,
+            device=device,
             enable_memory_saver=False,
         )
         sa = MultiEndedAllocator(
             kvcache=_FakeKVCache(pool.max_slots("state")),
             unified_buffer=pool,
             sub_pool_name="state",
-            device=_DEV,
+            device=device,
             is_id_owner=True,
         )
         fkv = _FakeKVCache(pool.max_slots("swa"))
         fla = FloatMultiEndedAllocator(
             kvcache=fkv,
             unified_buffer=pool,
             sub_pool_name="swa",
-            device=_DEV,
+            device=device,
             is_id_owner=True,
         )
         dkv = _FakeKVCache(pool.max_slots("full"))
         da = MultiEndedAllocator(
             kvcache=dkv,
             unified_buffer=pool,
             sub_pool_name="full",
-            device=_DEV,
+            device=device,
             is_id_owner=True,
         )
         # Chain wiring state <-> float <-> full.
@@ -2805,6 +2807,30 @@ def _check_float_state(self, fla, kv):
             self.assertEqual(int(fla.virtual_to_physical[v].item()), p)
             self.assertEqual(int(kv.buf[p].item()), v, f"data lost at {p}")
 
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA tombstone kernel")
+    def test_free_strided_page_ids_preserves_live_bindings(self):
+        """Retiring alternate pages must not accidentally retire their neighbours."""
+        _, _, allocator, _, _ = self._build_tri(device="cuda")
+        pages = allocator.alloc(8)
+        self.assertIsNotNone(pages)
+        freed = pages[::2]
+        self.assertFalse(freed.is_contiguous())
+        physical = allocator.virtual_to_physical[freed].clone()
+        expected_v2p = allocator.virtual_to_physical.clone()
+        expected_p2v = allocator.physical_to_virtual.clone()
+        expected_v2p[freed] = -1
+        expected_p2v[physical] = -1
+
+        allocator.free(freed)
+        torch.testing.assert_close(allocator.virtual_to_physical, expected_v2p)
+        torch.testing.assert_close(allocator.physical_to_virtual, expected_p2v)
+        torch.testing.assert_close(
+            allocator._free_phys_pages.sort().values, physical.sort().values
+        )
+        self.assertEqual(allocator._live_pages(), 4)
+        allocator.free(pages[1::2])
+        self.assertTrue(allocator._is_frontier_transparent())
+
     def test_midpoint_initial_placement(self):
         _, _, fla, _, _ = self._build_tri()
         self.assertTrue(fla._is_frontier_transparent())

test/registered/unit/mem_cache/test_unified_capacity_memo.py

@@ -28,8 +28,9 @@
 import unittest
 
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
-register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+register_cpu_ci(est_time=15, suite="base-a-test-cpu")
 
 
 def _build(lazy: bool):
@@ -183,7 +184,7 @@ def test_float_only_span_move_invalidates_every_memo(self):
         self.assertLessEqual(fla.available_size(), float_cached)
 
 
-class TestTriCapacityMemoCoherence(unittest.TestCase):
+class TestTriCapacityMemoCoherence(CustomTestCase):
     """Tri-composite twins of the 2-pool cases: the joint view walks THREE
     bands (mamba end, swa float, full end), so a mutation on ANY of them must
     invalidate the composite memo -- including the two only the tri has, a
@@ -214,6 +215,50 @@ def _assert_memos_fresh(self, allocator):
             )
         self.assertEqual(allocator.verify_byte_accounting(), [])
 
+    def test_joint_capacity_boundary_after_fragmentation(self):
+        """Equal FULL/SWA demand must fit exactly at the advertised boundary.
+
+        State allocations move a differently aligned peer, and SWA retirement
+        creates reusable holes. The cached joint capacity must still fund an
+        actual reservation, while the next page must exceed it.
+        """
+        from test_unified_tri_pool import TestUnifiedTriPool
+
+        fixture = TestUnifiedTriPool("test_chain_wiring_and_float_swa")
+        for page_size in (1, 4, 128):
+            for lazy in (False, True):
+                with self.subTest(page_size=page_size, lazy=lazy):
+                    _, allocator, _, _ = fixture._build(
+                        n_full=256 * page_size,
+                        n_swa=128 * page_size,
+                        n_state=32,
+                        page_size=page_size,
+                        lazy_compaction=lazy,
+                    )
+
+                    def check_boundary():
+                        capacity = allocator.available_size()
+                        self.assertTrue(allocator.can_reserve(capacity, capacity))
+                        self.assertFalse(
+                            allocator.can_reserve(
+                                capacity + page_size, capacity + page_size
+                            )
+                        )
+
+                    check_boundary()  # empty float is frontier-transparent
+                    states = allocator.mamba_allocator.alloc(3)
+                    self.assertIsNotNone(states)
+                    check_boundary()
+                    tokens = allocator.alloc(8 * page_size)
+                    self.assertIsNotNone(tokens)
+                    check_boundary()  # float now splits the available band
+                    allocator.free_swa(tokens[2 * page_size : 5 * page_size])
+                    check_boundary()  # interior holes are reusable
+                    allocator.free(tokens)
+                    check_boundary()  # float parks again
+                    allocator.mamba_allocator.free(states)
+                    check_boundary()
+
     def test_memos_track_tri_mutation_kinds(self):
         for lazy in (False, True):
             with self.subTest(lazy_compaction=lazy):