diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index ef5c46119..cc5851a2a 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 4238ae281..a7b47b703 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/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 = ( diff --git a/python/sglang/test/server_fixtures/streaming_session_fixture.py b/python/sglang/test/server_fixtures/streaming_session_fixture.py index d26415b97..6c95548ed 100644 --- a/python/sglang/test/server_fixtures/streaming_session_fixture.py +++ b/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( diff --git a/test/registered/unit/managers/test_kv_page_invariants.py b/test/registered/unit/managers/test_kv_page_invariants.py new file mode 100644 index 000000000..939acf1e3 --- /dev/null +++ b/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()