[Mem] Add KV-page double-free checks to the invariant checker (#27731)

python/sglang/srt/environ.py

@@ -279,6 +279,8 @@ class Envs:
     SGLANG_TEST_SKIP_CACHE_HIT_ASSERT = EnvBool(False)
     SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY = EnvInt(0)
     SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE = EnvBool(True)
+    # Physical KV-page checks: committed<=allocated + no page alias.
+    SGLANG_CHECK_KV_PAGE_INVARIANTS = EnvBool(False)
 
     # Load snapshot backend
     SGLANG_LOAD_SNAPSHOT_USE_ZMQ = EnvBool(False)

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -12,6 +12,8 @@ from typing import (
     Tuple,
 )
 
+import torch
+
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.environ import envs
 from sglang.srt.managers.scheduler_components.pool_stats_observer import (
@@ -263,6 +265,105 @@ class SchedulerInvariantChecker:
         assert not full_leak, f"Full Pool Mem Leak Detected! {full_msg}"
         assert not swa_leak, f"SWA Pool Mem Leak Detected! {swa_msg}"
 
+        if envs.SGLANG_CHECK_KV_PAGE_INVARIANTS.get():
+            self._check_kv_page_invariants()
+
+    def _check_kv_page_invariants(self):
+        """committed<=allocated for every req/slot, and no double free:
+          A. no owner references a page that is in the free pool (use-after-free).
+          B. the free pool has no duplicate pages (two owners freed the same page).
+        All heavy work runs on GPU to avoid per-token device->host sync."""
+        rtt = self.req_to_token_pool.req_to_token
+        row_width = rtt.shape[1]
+
+        def _add_owner(req_or_slot, label, rpi, committed, allocated):
+            assert 0 <= committed <= allocated <= row_width
+            owners.append((label, rpi, allocated))
+
+        owners: list[tuple[str, Optional[int], int]] = []
+        batch = self.get_last_batch()
+        if batch is not None:
+            for req in batch.reqs:
+                _add_owner(
+                    req,
+                    f"req {req.rid}",
+                    req.req_pool_idx,
+                    req.kv_committed_len,
+                    req.kv_allocated_len,
+                )
+        sess = getattr(self.tree_cache, "slots", None)
+        if sess:
+            for sid, slot in sess.items():
+                if getattr(slot, "is_holding_kv", False):
+                    _add_owner(
+                        slot,
+                        f"slot {sid[:8]}",
+                        slot.req_pool_idx,
+                        slot.kv_committed_len,
+                        slot.kv_allocated_len,
+                    )
+
+        active = [
+            (label, rpi, al) for label, rpi, al in owners if rpi is not None and al > 0
+        ]
+        if not active:
+            return
+
+        idx = torch.as_tensor([rpi for _, rpi, _ in active], device=rtt.device)
+        allocs = torch.as_tensor([al for _, _, al in active], device=rtt.device)
+        mask = torch.arange(row_width, device=rtt.device)[None, :] < allocs[:, None]
+        owner_pages = rtt[idx][mask] // self.page_size
+
+        # Sub-allocators to check: a flat allocator is its own single sub; a
+        # hybrid-SWA wrapper exposes full_attn_allocator + swa_attn_allocator.
+        alloc = self.token_to_kv_pool_allocator
+        sub_allocs = (
+            [alloc]
+            if getattr(alloc, "free_pages", None) is not None
+            else [
+                sub
+                for n in ("full_attn_allocator", "swa_attn_allocator")
+                if (sub := getattr(alloc, n, None)) is not None
+                and getattr(sub, "free_pages", None) is not None
+            ]
+        )
+        if not sub_allocs:
+            return
+
+        def _free_pages(a):
+            free = a.free_pages
+            release = getattr(a, "release_pages", None)
+            return (
+                torch.cat((free, release))
+                if release is not None and len(release) > 0
+                else free
+            )
+
+        # Check B: every sub-pool's free set has no duplicate pages.
+        for i, sub in enumerate(sub_allocs):
+            free = _free_pages(sub)
+            uniq = torch.unique(free)
+            if uniq.numel() != free.numel():
+                raise_error_or_warn(
+                    self,
+                    envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
+                    "count_memory_leak_warnings",
+                    f"KV double free: sub-pool {i} has {free.numel() - uniq.numel()} duplicate pages.",
+                )
+
+        # Check A: owner pages (full-pool indices) must not be in the full free
+        # set (sub_allocs[0] is the full pool, even on hybrid-SWA).
+        full_unique = torch.unique(_free_pages(sub_allocs[0]))
+        stale = owner_pages[torch.isin(owner_pages, full_unique)]
+        if stale.numel() > 0:
+            raise_error_or_warn(
+                self,
+                envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
+                "count_memory_leak_warnings",
+                f"KV page use-after-free: {stale.numel()} owner page refs are in "
+                f"the free pool, sample pages={torch.unique(stale)[:8].tolist()}.",
+            )
+
     def _check_req_pool(self):
         if self.disaggregation_mode == DisaggregationMode.DECODE:
             req_total_size = (

python/sglang/test/server_fixtures/streaming_session_fixture.py

@@ -419,6 +419,7 @@ class StreamingSessionServerBase(CustomTestCase):
             stack.enter_context(
                 envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1)
             )
+            stack.enter_context(envs.SGLANG_CHECK_KV_PAGE_INVARIANTS.override(True))
             for name, val in cls.env_overrides:
                 stack.enter_context(getattr(envs, name).override(val))
             cls.process = popen_launch_server(

test/registered/unit/managers/test_kv_page_invariants.py

@@ -0,0 +1,107 @@
+"""Unit tests for SGLANG_CHECK_KV_PAGE_INVARIANTS: watermark + double-free checks."""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+_PAGE_SIZE = 256
+
+
+def _make_checker(page_size=_PAGE_SIZE, row_width=4096, num_reqs=8, free_pages=None):
+    rtt = torch.zeros((num_reqs, row_width), dtype=torch.int32)
+    rtp = SimpleNamespace(req_to_token=rtt)
+    if free_pages is None:
+        free_pages = torch.arange(num_reqs * 4, dtype=torch.int64)
+    alloc = SimpleNamespace(
+        page_size=page_size,
+        free_pages=free_pages,
+        release_pages=torch.empty(0, dtype=torch.int64),
+    )
+    tc = SimpleNamespace(slots={})
+    _ps, _rtp, _alloc, _tc = page_size, rtp, alloc, tc
+
+    class _FakeChecker:
+        page_size = _ps
+        req_to_token_pool = _rtp
+        token_to_kv_pool_allocator = _alloc
+        tree_cache = _tc
+        get_last_batch = lambda self: None
+        count_memory_leak_warnings = 0
+
+        from sglang.srt.managers.scheduler_components.invariant_checker import (
+            SchedulerInvariantChecker as _RIC,
+        )
+
+        _check_kv_page_invariants = _RIC._check_kv_page_invariants
+
+    return _FakeChecker(), rtt, tc, alloc
+
+
+class _FakeReq:
+    def __init__(self, rid, rpi, committed, allocated):
+        self.rid = rid
+        self.req_pool_idx = rpi
+        self.kv_committed_len = committed
+        self.kv_allocated_len = allocated
+
+
+class _FakeSlot:
+    def __init__(self, rpi, committed, allocated):
+        self.req_pool_idx = rpi
+        self.kv_committed_len = committed
+        self.kv_allocated_len = allocated
+        self.is_holding_kv = True
+
+
+class TestKVPageInvariants(CustomTestCase):
+    def test_clean_layout_no_warning(self):
+        chk, rtt, tc, alloc = _make_checker(
+            free_pages=torch.arange(100, 200, dtype=torch.int64)
+        )
+        rtt[0, :256] = torch.arange(_PAGE_SIZE)  # req 0 owns page 0
+        rtt[1, :256] = torch.arange(_PAGE_SIZE, 2 * _PAGE_SIZE)  # req 1 owns page 1
+        chk.get_last_batch = lambda: SimpleNamespace(
+            reqs=[_FakeReq("a", 0, 256, 256), _FakeReq("b", 1, 200, 256)]
+        )
+        chk._check_kv_page_invariants()
+        self.assertEqual(chk.count_memory_leak_warnings, 0)
+
+    def test_committed_gt_allocated_raises(self):
+        chk, rtt, tc, alloc = _make_checker()
+        chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 145, 144)])
+        with self.assertRaises(AssertionError):
+            chk._check_kv_page_invariants()
+
+    def test_slot_committed_gt_allocated_raises(self):
+        chk, rtt, tc, alloc = _make_checker()
+        chk.get_last_batch = lambda: None
+        tc.slots = {"s1": _FakeSlot(0, 145, 144)}
+        with self.assertRaises(AssertionError):
+            chk._check_kv_page_invariants()
+
+    def test_owner_references_free_page_raises(self):
+        # req 0 owns page 5, but page 5 is in the free pool -> use-after-free.
+        chk, rtt, tc, alloc = _make_checker(free_pages=torch.tensor([5, 6, 7]))
+        rtt[0, :3] = torch.tensor(
+            [5 * _PAGE_SIZE, 5 * _PAGE_SIZE + 1, 5 * _PAGE_SIZE + 2]
+        )
+        chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 3, 3)])
+        with self.assertRaises(ValueError):
+            chk._check_kv_page_invariants()
+
+    def test_free_pool_duplicate_raises(self):
+        chk, rtt, tc, alloc = _make_checker(free_pages=torch.tensor([3, 3, 4]))
+        rtt[0, :1] = torch.tensor([10 * _PAGE_SIZE])  # owner page 10, not in free
+        chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 1, 1)])
+        with self.assertRaises(ValueError):
+            chk._check_kv_page_invariants()
+
+
+if __name__ == "__main__":
+    unittest.main()