sglang/test/registered/unit/mem_cache/test_unified_capacity_memo.py

# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Epoch-memoized capacity views on the allocator chain (2-pool subset).

The capacity views (`available_size` / `schedulable_available_size` per band,
plus the composite joint view) are pure functions of a handful of
CPU-resident fields across the chain; schedulers read them O(queue) times
between mutations. `_CapacityField` descriptors bump `_capacity_epoch` on
every rebind, so the memos invalidate by construction.

The failure mode being guarded: a memo serving a STALE value after a mutation
the epoch machinery missed — either a new mutation site writing a field the
descriptors don't cover, or an in-place write that bypasses `__set__`. Stale
capacity is silent over-/under-admission, not a crash. Hence:

  * every mutation kind is followed by memo == fresh-recompute assertions;
  * a randomized op-sequence property test (seeded) catches interactions no
    hand-written sequence covers;
  * a deliberate descriptor-bypassing write must be caught by the IDLE check
    (`verify_byte_accounting`) — readers cannot detect it, the battery must.

The tri-pool cases (float span fields, three-band joint view) join at the
tri phase; this file pins the 2-pool chain form they build on.

    python -m pytest test/registered/unit/mem_cache/test_unified_capacity_memo.py -v
"""

import random
import unittest

from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=20, suite="base-a-test-cpu")


def _build(lazy: bool):
    # Function-scope import: the fixture is a TestCase subclass, and a
    # module-scope binding would make pytest collect its tests AGAIN here.
    from test_multi_ended_allocator import (
        TestUnifiedSWATokenToKVPoolAllocator as _SwaFixture,
    )

    inst = _SwaFixture([m for m in dir(_SwaFixture) if m.startswith("test_")][0])
    pool, allocator, kvcache = inst._build()
    allocator.full_attn_allocator.lazy_compaction = lazy
    allocator.swa_attn_allocator.lazy_compaction = lazy
    allocator.lazy_compaction = lazy
    return inst, allocator, kvcache


class TestCapacityMemoCoherence(unittest.TestCase):
    def _assert_memos_fresh(self, allocator):
        """Every memoized capacity view must equal a fresh recompute."""
        self.assertEqual(
            allocator.available_size(), allocator._compute_available_size()
        )
        for band in (
            allocator.full_attn_allocator,
            allocator.swa_attn_allocator,
        ):
            self.assertEqual(
                band.available_size(),
                band._available_tokens(),
                msg=f"stale available_size memo on {band.sub_pool_name!r}",
            )
            self.assertEqual(
                band.schedulable_available_size(),
                band._available_tokens(
                    extra_gap_bytes=band._peer_drainable_hole_bytes()
                ),
                msg=f"stale schedulable memo on {band.sub_pool_name!r}",
            )
        self.assertEqual(allocator.verify_byte_accounting(), [])

    def test_memos_track_every_mutation_kind(self):
        for lazy in (False, True):
            with self.subTest(lazy_compaction=lazy):
                inst, allocator, kvcache = _build(lazy)
                self._assert_memos_fresh(allocator)

                v1 = inst._alloc(allocator, kvcache, 8)  # both-side bind
                self.assertIsNotNone(v1)
                self._assert_memos_fresh(allocator)

                allocator.free_swa(v1[2:6])  # swa-side tombstones
                self._assert_memos_fresh(allocator)

                inst._free(allocator, kvcache, v1)  # both-side free
                self._assert_memos_fresh(allocator)

                v2 = inst._alloc(allocator, kvcache, 4)
                self.assertIsNotNone(v2)
                allocator.free_group_begin()  # grouped free path
                allocator.free(v2)
                allocator.free_group_end()
                self._assert_memos_fresh(allocator)

                if lazy:
                    allocator.full_attn_allocator._flush(urgent=True)
                    self._assert_memos_fresh(allocator)

                allocator.clear()
                self._assert_memos_fresh(allocator)

    def test_random_op_sequence_value_identity(self):
        """Property: after ANY mutation sequence, memoized views equal fresh
        recomputes. Seeded (deterministic) — the sequences cover interleavings
        (alloc / partial swa free / grouped free / flush / clear) that no
        hand-written case enumerates."""
        rng = random.Random(0xC0FFEE)
        for lazy in (False, True):
            with self.subTest(lazy_compaction=lazy):
                inst, allocator, kvcache = _build(lazy)
                live = []
                for step in range(60):
                    op = rng.choice(("alloc", "free", "free_swa", "flush", "clear"))
                    if op == "alloc":
                        v = inst._alloc(allocator, kvcache, rng.choice((1, 2, 4)))
                        if v is not None:
                            live.append(v)
                    elif op == "free" and live:
                        inst._free(allocator, kvcache, live.pop())
                    elif op == "free_swa" and live:
                        v = live[-1]
                        if v.numel() > 1:
                            allocator.free_swa(v[: v.numel() // 2])
                    elif op == "flush":
                        allocator.full_attn_allocator._flush(urgent=True)
                        allocator.swa_attn_allocator._flush(urgent=True)
                    elif op == "clear":
                        allocator.clear()
                        live.clear()
                    self._assert_memos_fresh(allocator)

    def test_bypassing_write_is_caught_by_the_idle_check(self):
        inst, allocator, kvcache = _build(lazy=False)
        v = inst._alloc(allocator, kvcache, 8)
        self.assertIsNotNone(v)
        fa = allocator.full_attn_allocator
        # Prime every memo at the current epoch.
        allocator.available_size()
        fa.available_size()
        fa.schedulable_available_size()
        # Mutate capacity state BYPASSING the _CapacityField descriptor -- the
        # epoch does not move, so the memos go stale undetectably for readers...
        fa.__dict__["watermark_physical"] = fa.watermark_physical + 2
        # ...but the idle-time coherence check must flag it.
        violations = allocator.verify_byte_accounting()
        self.assertTrue(
            any("stale" in msg for msg in violations),
            msg=f"bypassing write not caught: {violations}",
        )

    def test_joint_memo_invalidates_on_swa_only_mutation(self):
        """The joint view depends on the swa end's frontier through the chain
        walk; an swa-side-only mutation (tombstoning) must invalidate the
        composite memo even though the full side never moved."""
        inst, allocator, kvcache = _build(lazy=True)
        v = inst._alloc(allocator, kvcache, 8)
        self.assertIsNotNone(v)
        before = allocator.available_size()
        allocator.free_swa(v[:4])  # swa band only
        after = allocator.available_size()
        self.assertEqual(after, allocator._compute_available_size())
        self.assertGreaterEqual(after, before)  # holes only ever add room

    def test_float_only_span_move_invalidates_every_memo(self):
        """A hole-free float alloc rebinds NO free-list and has no watermark --
        the span fields are its ONLY capacity state. If they are not
        `_CapacityField` descriptors, the float's own memo AND both
        neighbours' (the span flips transparency, walling off their gaps)
        keep serving pre-move values.

        Driven on a hand-wired end+float+end chain (the composite arrives
        with the tri phase); the float is exercised alone so no end-pool
        descriptor write can mask a missing span bump."""
        from test_multi_ended_allocator import TestFloatMultiEndedAllocator

        inst = TestFloatMultiEndedAllocator(
            [m for m in dir(TestFloatMultiEndedAllocator) if m.startswith("test_")][0]
        )
        _pool, sa, fla, da, _kv = inst._build_tri()
        self.assertEqual(fla._hole_pages(), 0)  # hole-free extension path
        self.assertTrue(fla._is_frontier_transparent())

        # Prime every memo while the float is empty/transparent.
        float_cached = fla.available_size()
        low_end_cached = sa.available_size()
        high_end_cached = da.available_size()

        v = fla.alloc(4)  # float-only mutation: span move, no end-pool write
        self.assertIsNotNone(v)
        self.assertFalse(fla._is_frontier_transparent())  # span now opaque

        self.assertEqual(fla.available_size(), fla._available_tokens())
        self.assertEqual(sa.available_size(), sa._available_tokens())
        self.assertEqual(da.available_size(), da._available_tokens())
        # The opaque midpoint span must actually reduce what the neighbours
        # see, i.e. the memos above were not merely re-serving primed values.
        self.assertLess(sa.available_size(), low_end_cached)
        self.assertLess(da.available_size(), high_end_cached)
        self.assertLessEqual(fla.available_size(), float_cached)

    def test_bind_rewiring_bumps_the_epoch(self):
        """Rewiring changes what the chain walks see; a memo primed before a
        re-bind must not survive it."""
        inst, allocator, kvcache = _build(lazy=False)
        fa = allocator.full_attn_allocator
        e0 = fa._chain_capacity_epoch()
        fa.bind_peer(allocator.swa_attn_allocator)  # re-bind (same peer)
        self.assertGreater(fa._chain_capacity_epoch(), e0)


class TestTriCapacityMemoCoherence(unittest.TestCase):
    """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 mutations only the tri
    has: a mamba-end state draw and a float span move behind the composite."""

    def _build_tri(self, lazy=False):
        from test_unified_tri_pool import TestUnifiedTriPool

        inst = TestUnifiedTriPool(
            [m for m in dir(TestUnifiedTriPool) if m.startswith("test_")][0]
        )
        pool, allocator, kvcache, mamba_kv = inst._build(lazy_compaction=lazy)
        return inst, allocator

    def _assert_memos_fresh(self, allocator):
        self.assertEqual(
            allocator.available_size(), allocator._compute_available_size()
        )
        for band in (
            allocator.full_attn_allocator,
            allocator.swa_attn_allocator,
            allocator.mamba_allocator,
        ):
            self.assertEqual(
                band.available_size(),
                band._available_tokens(),
                msg=f"stale available_size memo on {band.sub_pool_name!r}",
            )
        self.assertEqual(allocator.verify_byte_accounting(), [])

    def test_memos_track_tri_mutation_kinds(self):
        for lazy in (False, True):
            with self.subTest(lazy_compaction=lazy):
                inst, allocator = self._build_tri(lazy)
                ma = allocator.mamba_allocator
                self._assert_memos_fresh(allocator)

                v1 = allocator.alloc(8)  # composite alloc (full + swa bind)
                self.assertIsNotNone(v1)
                self._assert_memos_fresh(allocator)

                s1 = ma.alloc(2)  # mamba end alloc
                self.assertIsNotNone(s1)
                self._assert_memos_fresh(allocator)

                allocator.free_swa(v1[2:6])  # interior float holes
                self._assert_memos_fresh(allocator)

                allocator.free(v1)  # both-side free
                self._assert_memos_fresh(allocator)

                ma.free(s1)  # mamba free
                self._assert_memos_fresh(allocator)

                allocator.clear()
                ma.clear()
                self._assert_memos_fresh(allocator)

    def test_joint_memo_invalidates_on_mamba_only_mutation(self):
        """The joint view depends on the mamba end's frontier through the
        chain walk; a mamba-only mutation must invalidate the composite memo
        even though neither KV side moved."""
        inst, allocator = self._build_tri()
        ma = allocator.mamba_allocator
        before = allocator.available_size()
        slots = ma.alloc(4)
        self.assertIsNotNone(slots)
        after = allocator.available_size()
        self.assertEqual(after, allocator._compute_available_size())
        self.assertLessEqual(after, before)


if __name__ == "__main__":
    unittest.main()