[Test] Prune redundant unified-memory allocator and pool tests (#38093)

2026-09-04 20:47:01 -07:00
parent 92a4d8b5ee
commit d50e9a9756
21 changed files with 920 additions and 2487 deletions
@@ -1,28 +1,11 @@
 """Kimi-Linear (MLA full attention + KDA linear attention) served from the
 unified memory pool.
-Under `--enable-unified-memory` the MLA full side is exposed as per-layer
+Under `--enable-unified-memory` the MLA full side is exposed as per-layer views
-views and every loc the kernels see is a translated virtual id, so the whole
+and every loc the kernels see is a translated virtual id, so the whole
-read/write path differs from the static pool. The unit tests pin that pool in
+read/write path differs from the static pool. `test_prefix_cache_branching`
-isolation; this is the end-to-end guard. `test_prefix_cache_branching` carries
+carries most of the weight: a radix hit replays virtual locs whose physical
-most of the weight: a radix hit replays virtual locs whose physical pages may
+pages may have moved under compaction.
 have moved under compaction.
 No `--attention-backend` is pinned on purpose -- the test runs whatever the host
 resolves to (`fa3` on this suite's H100 runner, also the H200 default). Both
 defects found in review on #32972 were reachable only under a resolved default,
 which a pinned test hides by construction.
 Reference GSM8K, all with `--enable-unified-memory`:
  - 2x H200 TP2, resolved default (fa3): 0.917 @400, vs 0.915 static (1 sigma
    ~= 0.015). This file as written scores 0.920 @200.
  - 2x H200 TP2, `--attention-backend triton`: 0.900 @400.
  - 1x B300 TP1: 0.915 flashinfer, 0.900 trtllm_mla, @200.
 Nightly-only: a second full 48B server launch is too much for per-PR CI on top
 of the existing Kimi-Linear e2e coverage.
    python -m pytest test/registered/models_e2e/test_kimi_linear_unified_memory.py -v
 """
 import unittest
@@ -32,18 +15,23 @@ from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
 from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
 from sglang.test.server_fixtures.default_fixture import DefaultServerBase
-register_cuda_ci(est_time=1200, stage="nightly", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=800, stage="nightly", runner_config="4-gpu-h100")
 KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
-class TestKimiLinearUnifiedMemory(
+class TestKimiLinearUnifiedMemoryFlashMLA(
    GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
 ):
    """flashmla at its ps=64 snap: the block-table route
    (KVIndexTranslator.fill_read_table into flashmla's padded tables) plus the
    ps=64 sub-pool sizing (64-token sink floor, per-layer-view tail pad).
    Hopper-only, like the rest of this nightly suite."""
    model = KIMI_LINEAR_MODEL
    cache_chunk_size = 64
    # Same bar as the static-pool Kimi-Linear e2e test: unified memory must not
-    # cost accuracy (measured 0.917 vs 0.915 static, see the module docstring).
+    # cost accuracy.
    gsm8k_score_threshold = 0.88
    other_args = [
        "--trust-remote-code",
@@ -52,17 +40,6 @@ class TestKimiLinearUnifiedMemory(
        "--chunked-prefill-size",
        "2048",
        "--enable-unified-memory",
    ]
 class TestKimiLinearUnifiedMemoryFlashMLA(TestKimiLinearUnifiedMemory):
    """flashmla at its ps=64 snap: the canonical block-table route
    (KVIndexTranslator.fill_read_table into flashmla's padded tables) plus the
    ps=64 sub-pool sizing (64-token sink floor, per-layer-view tail pad) end to
    end.
    Hopper-only, like the rest of this nightly suite."""
    other_args = TestKimiLinearUnifiedMemory.other_args + [
        "--attention-backend",
        "flashmla",
        "--page-size",
@@ -73,13 +50,9 @@ class TestKimiLinearUnifiedMemoryFlashMLA(TestKimiLinearUnifiedMemory):
 class TestKimiLinearUnifiedMemoryDCP(
    GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
 ):
-    """Unified memory + decode context parallelism.
+    """Unified memory + decode context parallelism on flashinfer: on a radix hit
-
+    each rank must recover the same physical page from widened virtual locs
-    `test_prefix_cache_branching` is the sharp one here: a radix hit replays
+    while keeping a different row inside it."""
    widened virtual locs whose pages may have moved under compaction, and each
    rank must recover the same physical page from them while keeping a
    different row inside it.
    """
    model = KIMI_LINEAR_MODEL
    cache_chunk_size = 64
@@ -2,11 +2,8 @@
 gpt-oss-20b is uniform-row hybrid-SWA, so its MHA and SWA sub-pools are
 per-layer views and the fa3 cell reads them through the translator's read
-tables. The resolved-default cell pins the no-pin path, since a pinned
+tables. flashinfer is absent on purpose: gpt-oss uses attention sinks, which
-backend hides default-resolution breakage by construction. flashinfer is
+it does not support.
 absent on purpose: gpt-oss uses attention sinks, which it does not support.
 Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
 """
 import unittest
@@ -17,7 +14,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.server_fixtures.default_fixture import DefaultServerBase
 from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
-register_cuda_ci(est_time=1500, stage="extra-a", runner_config="1-gpu-large")
+register_cuda_ci(est_time=1000, stage="extra-a", runner_config="1-gpu-large")
 _UNIFIED_COMMON_ARGS = [
    "--enable-unified-memory",
@@ -68,12 +65,5 @@ class TestUnifiedGptOssFa3(TestUnifiedGptOssTriton):
    other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "fa3"]
 class TestUnifiedGptOssResolvedDefault(TestUnifiedGptOssTriton):
    """No backend pin: whatever the host resolves must be in the allow-list,
    or the server fails to boot under its own defaults."""
    other_args = _UNIFIED_COMMON_ARGS
 if __name__ == "__main__":
    unittest.main()
@@ -4,11 +4,7 @@ Qwen3.5-4B is a gated-delta-net / linear-attention hybrid, which exercises the
 path most prone to subtle bugs: the Mamba conv/SSM state stays a strided
 envelope view (its kernels are stride-aware by design) while the
 full-attention KV is per-layer views, which the fa3 / flashinfer cells read
-through the translator's read tables. The resolved-default cell pins the
+through the translator's read tables.
 no-pin path, since a pinned backend hides default-resolution breakage by
 construction.
 Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
 """
 import unittest
@@ -19,7 +15,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.server_fixtures.default_fixture import DefaultServerBase
 from sglang.test.test_utils import DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
-register_cuda_ci(est_time=1600, stage="extra-a", runner_config="1-gpu-large")
+register_cuda_ci(est_time=1200, stage="extra-a", runner_config="1-gpu-large")
 _UNIFIED_COMMON_ARGS = [
    "--trust-remote-code",
@@ -40,9 +36,8 @@ class TestUnifiedQwenHybridTriton(DefaultServerBase):
    model = DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
-    # Measured ~0.86 in this harness on both the static pools and the envelope
+    # Measured ~0.86 on both the static pools and the envelope layout; 0.80
-    # layout; 0.80 leaves noise margin and still catches a corrupted prefill
+    # leaves noise margin and still catches a corrupted prefill state (~0.61).
    # state, which reads ~0.61.
    gsm8k_threshold = 0.80
    num_gsm8k_questions = 200
    num_shots = 5
@@ -84,12 +79,5 @@ class TestUnifiedQwenHybridFlashinfer(TestUnifiedQwenHybridTriton):
    other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "flashinfer"]
 class TestUnifiedQwenHybridResolvedDefault(TestUnifiedQwenHybridTriton):
    """No backend pin: whatever the host resolves must be in the allow-list,
    or the server fails to boot under its own defaults."""
    other_args = _UNIFIED_COMMON_ARGS
 if __name__ == "__main__":
    unittest.main()
@@ -1,20 +1,16 @@
 """Nothing under layers/attention may translate KV ids for itself.
-Ownership is exactly two places: `KVIndexTranslator` for READS (indices are
+Ownership is exactly two places: `KVIndexTranslator` for READS and the
-born kernel-facing, backends consume its tables) and the ForwardBatch rebind
+ForwardBatch rebind (`rebind_write_loc`) for WRITES. Virtual and physical ids
-(`rebind_write_loc`) for WRITES. Virtual and physical ids share a value range,
+share a value range, so a backend that forgets a translate -- or does one
-so a backend that forgets a translate -- or does one twice -- reads the wrong
+twice -- reads the wrong rows and nothing crashes.
 rows and nothing crashes. This scan makes both unrepresentable.
-Out of scope, deliberately: the allocator-internal implementations
+Deliberately out of scope: the allocator-internal implementations
-(`allocator/unified_*` / `unified_memory_pool`), which ARE the mechanism the
+(`allocator/unified_*`, `unified_memory_pool`), which ARE the mechanism the
 translator calls; the PD transfer plane's `translate_kv_indices_for_transfer`,
 which stages for RDMA outside the forward path; and the STATIC SWA pool's
 legacy full->swa slot map, a different mapping kind with no virtual/physical
-ambiguity -- its call sites are count-pinned below so new ones are added
+ambiguity.
 consciously.
    python3 -m pytest test/registered/unit/layers/attention/test_kv_translate_ownership.py -v
 """
 import ast
@@ -60,30 +56,21 @@ def _iter_sources():
 class TestUnifiedTranslateBanned(CustomTestCase):
    def test_no_unified_translate_calls(self):
-        """No backend calls the unified translate surfaces. A hit here means
+        """No backend calls the unified translate surfaces, and none probes an
-        a backend re-grew its own id-space transition -- the design whose two
+        allocator for translate capability through getattr."""
        failure modes (forgotten translate, duplicated translate) this scan
        exists to prevent. Route reads through KVIndexTranslator views and
        writes through the ForwardBatch rebind instead."""
        banned = re.compile(r"\.translate_kv_loc(_kernel_id)?\(")
        hits = [
            f"{rel}: {m.group(0)}"
            for rel, src in _iter_sources()
            for m in banned.finditer(src)
        ]
        self.assertEqual(hits, [])
    def test_no_translate_capability_probing(self):
        """No backend probes an allocator for translate capability -- the
        getattr-hook pattern is how per-backend translation grew the first
        time."""
        probing = re.compile(r"""getattr\([^)]*['"]translate_kv_loc""")
-        hits = [rel for rel, src in _iter_sources() if probing.search(src)]
+        calls, probes = [], []
-        self.assertEqual(hits, [])
+        for rel, src in _iter_sources():
            calls += [f"{rel}: {m.group(0)}" for m in banned.finditer(src)]
            if probing.search(src):
                probes.append(rel)
        self.assertEqual(calls, [])
        self.assertEqual(probes, [])
    def test_hooks_module_deleted_and_unimported(self):
-        """The per-backend hooks module (the previous owner of backend-side
+        """The per-backend v2p hooks module stays deleted, and nothing imports
-        v2p knowledge) stays deleted, and nothing imports it."""
+        it."""
        self.assertFalse(
            os.path.exists(os.path.join(_ATTN_DIR, "unified_mem_hooks.py"))
        )
@@ -96,10 +83,9 @@ class TestUnifiedTranslateBanned(CustomTestCase):
 def _derive_wrapper_names():
-    """Wrapper classes, read off the source so a NEW one shows up the day it is
+    """Wrapper classes discovered from source, so a new one shows up the day it
-    written: an AttentionBackend subclass whose own __init__ takes another
+    is written: an AttentionBackend subclass whose own __init__ takes another
-    backend. Per class, not per file -- a file-wide scan passes as soon as any
+    backend. Per class, not per file."""
    one class in it forwards."""
    backend = re.compile(r"Att(?:ention|n)Backend")
    names = set()
    for _rel, src in _iter_sources():
@@ -155,8 +141,8 @@ class _Runner:
 def _build_wrappers(translator):
    """One live instance per wrapper. Only the inner that MUST supply the
-    translator carries it; every other inner carries None, so a wrapper that
+    translator carries it, so a wrapper copying off the linear / sparse / DSA
-    copies from the linear / sparse / DSA side ends up with None and fails."""
+    side ends up with None and fails."""
    carrier = _Inner(translator)
    # HybridAttnBackend reads the spec bag in __init__; the bag is unpublished
    # outside a launched server.
@@ -176,10 +162,9 @@ def _build_wrappers(translator):
 class TestWrapperBackendsForwardTranslator(CustomTestCase):
-    """BUG REGRESSION. `AttentionBackend.kv_index_translator` defaults to None,
+    """Bug regression: `AttentionBackend.kv_index_translator` defaults to None,
-    so a wrapper that does not re-expose its inner's copy reads as "needs no
+    so a wrapper that does not re-expose its inner's copy makes producers skip
-    translation" and producers that fetch it off `get_attn_backend()` skip the
+    the virtual->kernel-facing translation instead of failing."""
    virtual->kernel-facing translation instead of failing."""
    def test_every_wrapper_is_constructed_here(self):
        self.assertEqual(
@@ -15,21 +15,9 @@
 `HybridLinearKVPool`).
 All write-location info travels in the attention metadata (`KVWriteLoc`); the
-pools hold none and never translate — the write loc reaching `set_kv_buffer` is
+pools hold none and never translate, so the loc reaching `set_kv_buffer` is
-always PHYSICAL. Two routing contracts are pinned here:
+always PHYSICAL. Pure dispatch: the inner sub-pools are recording stubs, so no
-
+GPU and no real buffers are needed.
 1. Full-attention. The full-physical loc is carried in `KVWriteLoc.full_loc`
   (from `ForwardMetadata.out_cache_loc_full_physical`) and written directly.
   `UnifiedSWAKVPool` asserts it's present (the unified memory pool always precomputes
   it); `HybridLinearKVPool` falls back to `loc` for a static (non-shared) pool,
   where `loc` is itself already physical.
 2. SWA. The swa-physical loc rides the backend `swa_out_cache_loc` slot
   (`KVWriteLoc.swa_loc`) and is written directly.
 Pure dispatch tests: the inner sub-pools are recording stubs, so no GPU / real
 buffers are needed (CPU CI).
    python -m pytest test/registered/unit/mem_cache/test_full_loc_fast_path.py -v
 """
 import types
@@ -61,9 +49,8 @@ class _RecordingPool:
 class TestUnifiedSWARouting(unittest.TestCase):
    """`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write `full_loc`
    when present (triton's capture-stable buffer), else the rebound generic
-    `loc` -- the same id space once the loc is rebound; SWA layers write the swa-physical
+    `loc` -- the same id space; SWA layers write the swa-physical `swa_loc`,
-    `swa_loc`, which has no fallback (a different id space). The pool never
+    which has no fallback (a different id space)."""
    translates."""
    def _make_bare_pool(self):
        from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
@@ -93,18 +80,15 @@ class TestUnifiedSWARouting(unittest.TestCase):
        self.assertEqual(len(pool.full_kv_pool.calls), 1)
        forwarded, kwargs = pool.full_kv_pool.calls[0]
        # Forward the full-physical tensor from the write metadata, NOT the
-        # virtual loc. No `already_physical` — the pool only ever gets physical.
+        # virtual loc; no `already_physical`, the pool only ever gets physical.
        self.assertIs(forwarded, full_phys)
        self.assertIsNot(forwarded, virtual_loc)
        self.assertNotIn("already_physical", kwargs)
    def test_full_layer_falls_back_to_generic_loc(self):
-        """Bug regression: fa3 x unified-SWA crashed at gpt-oss
+        """Bug regression: a 2-arg `KVWriteLoc(loc, swa)` with no explicit
-        cuda-graph capture because every backend except triton bundles the
+        `full_loc` must fall back to the rebound `loc` -- which IS the full-side
-        2-arg KVWriteLoc(loc, swa) and the full-layer door demanded an explicit
+        kernel-facing id -- instead of failing the full-layer door."""
        full_loc. Once the loc is rebound the generic `loc` IS the full-side kernel-facing id
        (rebind_write_loc runs at ForwardBatch construction),
        so the door must fall back to it -- the pool still never translates."""
        pool = self._make_bare_pool()
        rebound_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
        swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64)
@@ -163,10 +147,9 @@ class TestUnifiedSWATombstoneClamp(unittest.TestCase):
    """`UnifiedSWAKVPool.translate_loc_from_full_to_swa` must clamp tombstoned
    ids to the reserved padding sink (0).
-    A token whose swa page was freed carries -1 in `virtual_to_physical`. Before
+    A token whose swa page was freed carries -1 in `virtual_to_physical`, and a
-    the clamp, that produced a negative id, which a captured graph stores at a
+    negative id makes a captured graph store at a negative offset from the
-    negative offset from the buffer base. The composite allocator's method of
+    buffer base.
    the same name already clamped; this path did not.
    """
    def _make_bare_pool(self, page_size, v2p, multiplier=1):
@@ -203,83 +186,66 @@ class TestUnifiedSWATombstoneClamp(unittest.TestCase):
 class TestHybridLinearFullLocRouting(unittest.TestCase):
-    """`HybridLinearKVPool.set_kv_buffer` (non-MLA) writes the full-physical
+    """`HybridLinearKVPool.set_kv_buffer` writes the full-physical `full_loc`
-    `full_loc` from the write metadata when present (unified memory pool), else the
+    when present (unified memory pool), else the already-physical `loc` (static
-    already-physical `loc` (static pool). No translate, no `already_physical`."""
+    pool) -- the same rule on the MHA and MLA branches."""
-    def _make_bare_pool(self):
+    def _make_bare_pool(self, use_mla):
        from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
        pool = object.__new__(HybridLinearKVPool)
        pool.full_kv_pool = _RecordingPool()
-        pool.use_mla = False
+        pool.use_mla = use_mla
        pool.full_attention_layer_id_mapping = {0: 0}
        return pool
    def test_writes_full_loc_from_write_loc(self):
-        pool = self._make_bare_pool()
+        for use_mla in (False, True):
-        virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
+            for has_full_loc in (True, False):
-        full_phys = torch.tensor([2, 3, 4], dtype=torch.int64)
+                with self.subTest(use_mla=use_mla, has_full_loc=has_full_loc):
                    pool = self._make_bare_pool(use_mla)
                    loc = torch.tensor([7, 8, 9], dtype=torch.int64)
                    full_phys = (
                        torch.tensor([2, 3, 4], dtype=torch.int64)
                        if has_full_loc
                        else None
                    )
-        layer = types.SimpleNamespace(layer_id=0)
+                    layer = types.SimpleNamespace(layer_id=0)
-        pool.set_kv_buffer(
+                    pool.set_kv_buffer(
-            layer,
+                        layer,
-            _loc_info(virtual_loc, full_phys=full_phys),
+                        _loc_info(loc, full_phys=full_phys),
-            torch.zeros(3, 4, 8),
+                        torch.zeros(3, 4, 8),
-            torch.zeros(3, 4, 8),
+                        None if use_mla else torch.zeros(3, 4, 8),
-        )
+                    )
-        self.assertEqual(len(pool.full_kv_pool.calls), 1)
+                    self.assertEqual(len(pool.full_kv_pool.calls), 1)
-        forwarded, kwargs = pool.full_kv_pool.calls[0]
+                    forwarded, kwargs = pool.full_kv_pool.calls[0]
-        self.assertIs(forwarded, full_phys)
+                    if has_full_loc:
-        self.assertIsNot(forwarded, virtual_loc)
+                        self.assertIs(forwarded, full_phys)
-        self.assertNotIn("already_physical", kwargs)
+                        self.assertIsNot(forwarded, loc)
-
+                    else:
-    def test_falls_back_to_loc_when_absent(self):
+                        # Static (non-shared) pool: no full_loc bundled; `loc`
-        # Static (non-shared) pool: no full_loc bundled; `loc` is already
+                        # is already physical, so write it directly.
-        # physical, so write it directly.
+                        self.assertIs(forwarded, loc)
-        pool = self._make_bare_pool()
+                    self.assertNotIn("already_physical", kwargs)
        phys_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
        layer = types.SimpleNamespace(layer_id=0)
        pool.set_kv_buffer(
            layer,
            _loc_info(phys_loc),
            torch.zeros(3, 4, 8),
            torch.zeros(3, 4, 8),
        )
        self.assertEqual(len(pool.full_kv_pool.calls), 1)
        forwarded, kwargs = pool.full_kv_pool.calls[0]
        self.assertIs(forwarded, phys_loc)
        self.assertNotIn("already_physical", kwargs)
 class _RecordingMLAPool(_RecordingPool):
-    """Also records the model-level MLA entry points."""
+    """Also records the model-level MLA write entry point."""
    def __init__(self):
        super().__init__()
        self.mla_set_calls = []
        self.mla_get_calls = []
    def set_mla_kv_buffer(self, layer, loc, cache_k_nope, cache_k_rope):
        self.mla_set_calls.append(loc)
    def get_mla_kv_buffer(self, layer, loc, dst_dtype=None):
        self.mla_get_calls.append(loc)
        return None, None
 class TestHybridLinearMLARouting(unittest.TestCase):
-    """MLA-side routing contracts of `HybridLinearKVPool`:
+    """MLA-side door contract of `HybridLinearKVPool`: `set_mla_kv_buffer`
-
+    forwards `loc` untouched -- writes are kernel-facing since the ForwardBatch
-    - `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the
+    rebind."""
      pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it
      carries the kernel-facing loc), else the raw `loc` (static pool, already physical).
    - `set_mla_kv_buffer` / `get_mla_kv_buffer` forward `loc` untouched:
      writes are kernel-facing since the ForwardBatch rebind, and read
      indices are translated at their production sites."""
    def _make_bare_pool(self):
        from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
@@ -290,46 +256,10 @@ class TestHybridLinearMLARouting(unittest.TestCase):
        pool.full_attention_layer_id_mapping = {0: 0}
        return pool
    def test_mla_writes_full_loc_from_write_loc(self):
        pool = self._make_bare_pool()
        virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
        kernel_phys = torch.tensor([21, 24, 27], dtype=torch.int64)
        layer = types.SimpleNamespace(layer_id=0)
        pool.set_kv_buffer(
            layer,
            _loc_info(virtual_loc, full_phys=kernel_phys),
            torch.zeros(3, 1, 8),
            None,
        )
        self.assertEqual(len(pool.full_kv_pool.calls), 1)
        forwarded, _ = pool.full_kv_pool.calls[0]
        self.assertIs(forwarded, kernel_phys)
        self.assertIsNot(forwarded, virtual_loc)
    def test_mla_falls_back_to_loc_when_absent(self):
        pool = self._make_bare_pool()
        phys_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
        layer = types.SimpleNamespace(layer_id=0)
        pool.set_kv_buffer(
            layer,
            _loc_info(phys_loc),
            torch.zeros(3, 1, 8),
            None,
        )
        self.assertEqual(len(pool.full_kv_pool.calls), 1)
        forwarded, _ = pool.full_kv_pool.calls[0]
        self.assertIs(forwarded, phys_loc)
    def test_set_mla_kv_buffer_door_never_translates(self):
-        """Physical-loc contract: the write door forwards `loc` UNTOUCHED.
+        """The translate happens exactly once, at ForwardBatch construction
-        The translate happens exactly once at ForwardBatch construction
+        (`rebind_write_loc`); a door that translated again would
-        (rebind_write_loc, kernel-facing-first); a door that translated
+        double-translate every unified MLA write."""
        again would double-translate every unified MLA write. Deleting the
        forward (or re-adding a door translate) turns this red."""
        pool = self._make_bare_pool()
        loc = torch.tensor([107, 108, 109], dtype=torch.int64)
        layer = types.SimpleNamespace(layer_id=0)
@@ -339,21 +269,6 @@ class TestHybridLinearMLARouting(unittest.TestCase):
        self.assertEqual(len(pool.full_kv_pool.mla_set_calls), 1)
        self.assertIs(pool.full_kv_pool.mla_set_calls[0], loc)
    def test_get_mla_kv_buffer_door_never_translates(self):
        """Kernel-facing contract, read side: `loc` is a read-index tensor
        already translated at its production site
        (fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
        door forwards it UNTOUCHED -- a re-added door translate would
        double-translate every unified MLA prefix read."""
        pool = self._make_bare_pool()
        loc = torch.tensor([104, 105], dtype=torch.int64)
        layer = types.SimpleNamespace(layer_id=0)
        pool.get_mla_kv_buffer(layer, loc)
        self.assertEqual(len(pool.full_kv_pool.mla_get_calls), 1)
        self.assertIs(pool.full_kv_pool.mla_get_calls[0], loc)
 class TestMlaWriteDoorsUnderDcp(unittest.TestCase):
    """Which MLA write door is DCP-aware, and which refuses.
@@ -362,10 +277,8 @@ class TestMlaWriteDoorsUnderDcp(unittest.TestCase):
    the DCP write. `set_kv_buffer` (the combined latent+rope row) cannot: the
    two backends that could reach it disagree on the loc space -- flashinfer's
    `k_rope is None` branch passes a WIDENED loc, the Triton backend one it
-    already collapsed -- so there is no single correct translation. It used to
+    already collapsed -- so there is no single correct translation and refusing
-    select `loc % dcp_size == dcp_rank` and then write WITHOUT dividing, i.e.
+    is the contract."""
    widened ids straight into a rank-local buffer. Refusing is the contract;
    a re-added masked-but-undivided write is what this guards."""
    def _bare_mla_pool(self):
        from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
@@ -13,29 +13,8 @@
 # ==============================================================================
 """KVIndexTranslator -- the read-path id translator.
-Covers, CPU-only (the builder's pure-torch reference path; GPU parity of the
+CPU-only: these exercise the builder's pure-torch reference path, not the
-Triton kernel is a later CUDA CI pin):
+Triton kernel.
  - strict passthrough: a non-unified source returns the SAME req_to_token /
    req_pool_indices objects -- zero tensor ops, no copies (the property that
    makes backend re-pointing byte-identical for every non-unified server);
  - static SWA pools keep their legacy full->swa mapping on the view;
  - the read table matches the hand formula
        entry[b, c] = clamp(v2p[req_to_token[req[b], c*ps] // ps] * mult, 0)
    over the REAL SWA composite's tables (full AND swa, ps in {1, 4},
    multiplier in {1, 2L}), with the swa table built from VIRTUAL ids;
  - sink routing: dead lanes (seq_len 0), -1 req_to_token entries, and
    tombstoned v2p pages all read entry 0;
  - the capture contract: buffers are zero-filled and idempotent; a refresh
    updates ONLY the live prefix (stale tails and rows beyond bs keep prior
    contents); the returned table is the WHOLE buffer (pointer-stable);
  - the eager-view memo: a single source-resident slot keyed by batch
    identity (same batch shares one build; the next batch replaces it; a
    dead batch never matches);
  - the two-phase write contract: the rebind touches only the full side, and
    the sliding-window write loc derives POINTWISE from the kernel-facing values
    (pads, slices, and fresh copies included), for both pool families.
    python -m pytest test/registered/unit/mem_cache/test_kv_index_translator.py -v
 """
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -137,9 +116,8 @@ def _reference_table(req_to_token, req_pool_indices, seq_lens, v2p, mult, ps, wi
 class TestPassthrough(unittest.TestCase):
    def test_non_unified_returns_same_objects(self):
-        """The strict-passthrough property: no copy, no branch, the exact
+        """Strict passthrough: no copy, no branch. Any tensor op on the
-        tensors backends read today. A regression here (any tensor op on the
+        non-unified path breaks byte-identity for every static-pool server."""
        non-unified path) breaks byte-identity for every static-pool server."""
        req_to_token = torch.arange(64, dtype=torch.int64).reshape(4, 16)
        src = KVIndexTranslator(
            req_to_token=req_to_token,
@@ -166,7 +144,7 @@ class TestPassthrough(unittest.TestCase):
 def _alloc_and_fill(allocator, ps, lens):
    """Allocate per-request virtual runs and write them into a fake
-    req_to_token; returns (req_to_token, req_pool_indices, seq_lens)."""
+    req_to_token."""
    width = 16 * ps
    req_to_token = torch.full((len(lens), width), -1, dtype=torch.int64)
    for r, n in enumerate(lens):
@@ -183,11 +161,10 @@ def _alloc_and_fill(allocator, ps, lens):
 class TestReadTableBuild(unittest.TestCase):
    def test_read_table_matches_reference_across_multipliers(self):
-        """The load-bearing formula pin: full AND swa read tables equal
+        """Both read tables must equal the independent per-element derivation,
-        the independent per-element derivation, across page sizes and both
+        across page sizes and both multiplier regimes (MLA=1, MHA=2L); the swa
-        multiplier regimes (MLA=1, MHA=2L). The swa table agreeing with
+        table agreeing over VIRTUAL ids proves it is never chained through
-        a formula over VIRTUAL ids is also the never-chained-through-
+        full-physical."""
        full-physical proof."""
        for ps in (1, 4):
            for collapse in (True, False):
                allocator = _build_composite(ps, collapse=collapse)
@@ -235,10 +212,8 @@ class TestReadTableBuild(unittest.TestCase):
                )
    def test_packed_stream_equals_the_rectangle_it_replaces(self):
-        """The packed builder and the rectangle must agree element for element:
+        """The two builders must agree element for element:
-        packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps. Consumers that
+        packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps."""
        plan over the stream and consumers that read the page table have to see
        the same KV, so a change to either builder alone turns this red."""
        for ps in (1, 4):
            allocator = _build_composite(ps)
            req_to_token, rows, seq_lens = _alloc_and_fill(
@@ -269,9 +244,9 @@ class TestReadTableBuild(unittest.TestCase):
                    )
    def test_sink_routing(self):
-        """Dead lanes (seq_len 0), -1 slots inside the live prefix, and
+        """Dead lanes, -1 slots inside the live prefix, and tombstoned v2p
-        tombstoned v2p pages must ALL read entry 0 -- one wild entry is a
+        pages must ALL read entry 0; one wild entry is a captured-graph OOB
-        captured-graph OOB read at replay."""
+        read at replay."""
        ps = 4
        allocator = _build_composite(ps)
        req_to_token, rows, seq_lens = _alloc_and_fill(
@@ -300,12 +275,9 @@ class TestBuildInto(unittest.TestCase):
    (their rows ARE the read table's rows)."""
    def test_prefix_filled_tail_sentinel_preserved_width_capped(self):
-        """Three contracts in one batch: entries equal the read-table formula,
+        """The -1 tail sentinel belongs to the backend, and a table padded
-        lanes past each row's live pages keep the backend's -1 sentinel
+        WIDER than the req_to_token page span (trtllm's LCM alignment) must be
-        (prefix-only -- a tail write scatters the trtllm sentinel contract),
+        capped rather than trip the builder's width assert."""
        and a table padded WIDER than the req_to_token page span (trtllm's
        LCM alignment) is capped instead of tripping the builder's width
        assert."""
        ps = 4
        allocator = _build_composite(ps)
        full_mult = allocator.kernel_page_multiplier
@@ -339,8 +311,8 @@ class TestBuildInto(unittest.TestCase):
            )
    def test_passthrough_source_refuses(self):
-        """Callers dispatch on `enabled`; a passthrough source has no v2p to
+        """Callers dispatch on `reads_are_translated`; a passthrough source has
-        build from and must fail loud, not fill garbage."""
+        no v2p to build from and must fail loud, not fill garbage."""
        src = KVIndexTranslator(
            req_to_token=torch.zeros((2, 4), dtype=torch.int64),
            token_to_kv_pool_allocator=SimpleNamespace(),
@@ -360,21 +332,18 @@ class TestPoolOwnership(unittest.TestCase):
    """A runner only gets the kernel-facing id space when the pool IT reads and
    writes is the one the allocator's ids address.
-    Guarded shape: a runner handed a SHARED allocator (one slot index space,
+    Guarded shape: a runner handed a SHARED allocator while owning a SEPARATE
-    one req_to_token) while owning a SEPARATE KV buffer sized to the
+    KV buffer sized to the allocator's SLOT count. Probing the allocator alone
-    allocator's SLOT count. Probing the allocator alone reports "unified" for
+    reports "unified" for that runner, so its indices would be mapped into the
-    that runner, so its indices would be mapped into the composite's
+    composite's kernel-facing space (ids up to num_pages * multiplier) and used
-    kernel-facing space (kernel-facing ids up to num_pages * multiplier) and then used
+    to address a buffer with only num_slots rows.
    to address a buffer with only num_slots rows -- out of bounds on both the
    read gather and the KV store.
    """
    def test_real_factory_bundle_satisfies_the_ownership_identity(self):
-        """The guard rests on `allocator.get_kvcache() is token_to_kv_pool`
+        """The guard rests on `allocator.get_kvcache() is token_to_kv_pool`, so
-        holding for a REAL target bundle. If a factory ever returned a pool
+        a factory returning a pool the allocator does not hold would silently
-        the allocator does not hold, the guard would silently disable the
+        disable the unified path for EVERY model. Pinned against the real
-        unified path for EVERY model -- so pin it against the real factory
+        factory, not this file's own construction."""
        rather than against this file's own construction."""
        from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
        bundle = init_unified_swa_pools(
@@ -424,9 +393,9 @@ class TestPoolOwnership(unittest.TestCase):
        self.assertFalse(src.is_translating)
    def test_disabled_source_is_strict_passthrough(self):
-        """Consequence of the guard: such a runner must see RAW virtual ids on
+        """Such a runner must see RAW virtual ids: they index its own pool
-        the read table -- they index its own pool directly. A translate here is
+        directly, and a translate here is the out-of-bounds bug the ownership
-        the out-of-bounds bug the ownership identity exists to prevent."""
+        identity exists to prevent."""
        alloc = _build_composite(ps=1)
        req_to_token = torch.arange(16, dtype=torch.int32, device=_DEV).view(2, 8)
        src = KVIndexTranslator(
@@ -483,10 +452,9 @@ class TestCaptureContract(unittest.TestCase):
        self.assertTrue(bool((cap[1:] == 7).all()), "rows beyond bs were touched")
    def test_row_ids_not_reallocated_across_builds(self):
-        """`row_ids` is a constant arange sized once from the request pool, so
+        """`row_ids` is one arange sized from the request pool and sliced per
-        builds at different batch sizes hand back slices of ONE buffer. A
+        build; a per-build `torch.arange` would be correct but costs an
-        per-build `torch.arange` would be correct but would spend an allocation
+        allocation and a launch on every replay prep."""
        and a launch on every replay prep."""
        allocator = _build_composite(1)
        req_to_token, rows, seq_lens = _alloc_and_fill(allocator, 1, lens=[4, 2, 3])
        src = _make_source(allocator, req_to_token, 1)
@@ -531,10 +499,9 @@ class _FakeForwardBatch:
 class TestViewMemo(unittest.TestCase):
-    """The eager view is memoized ON THE SOURCE in a single slot keyed by
+    """The eager view is memoized ON THE SOURCE in a single slot keyed by batch
-    batch identity -- per-batch state stays out of the ForwardBatch (it does
+    identity, so per-batch state stays out of the ForwardBatch while one
-    not scale with the number of id spaces), and one metadata build's many
+    metadata build's many consumers still share one table build."""
    consumers still share one table build."""
    def _fb(self, allocator, ps, lens):
        req_to_token, rows, seq_lens = _alloc_and_fill(allocator, ps, lens=lens)
@@ -590,12 +557,10 @@ class TestViewMemo(unittest.TestCase):
 class TestWriteLoc(unittest.TestCase):
-    """The two-phase write contract: phase 1 (`rebind_write_loc`) rebinds the
+    """The two-phase write contract: `rebind_write_loc` rebinds the full side
-    full side once at ForwardBatch construction; phase 2 derives the
+    once at ForwardBatch construction, and the sliding-window write loc derives
-    sliding-window write loc on demand, POINTWISE from the full-side
+    POINTWISE from the full-side values -- pads, slices, and fresh copies
-    values. Value-based derivation is the property under test: pads, slices,
+    included -- with no handover and no stored per-forward state."""
    and fresh copies of the loc must all derive correctly with no handover
    and no stored per-forward state."""
    def _built(self, ps=1, n=4):
        allocator = _build_composite(ps)
@@ -622,10 +587,8 @@ class TestWriteLoc(unittest.TestCase):
            self.assertTrue(torch.equal(virt, keep))
    def test_swa_write_loc_round_trips_from_full_side(self):
-        """The derived property behind phase 2: for any virtual run t,
+        """Derived property: `field(full(t)) == swa(t)` for any virtual run t,
-        deriving from the kernel-facing full-side values must equal the direct
+        across page sizes and multipliers."""
        virtual->swa translate — `field(full(t)) == swa(t)` across page sizes
        and multipliers."""
        for ps in (1, 4, 64):
            src, _, rows, seq_lens, _, want_full, want_swa = self._built(
                ps=ps, n=3 * ps
@@ -634,30 +597,15 @@ class TestWriteLoc(unittest.TestCase):
            self.assertTrue(torch.equal(got, want_swa))
    def test_pad_lanes_derive_to_sink(self):
-        """The DP pad appends zeros; kernel-facing 0 is the reserved padding slot in
+        """The DP pad appends zeros, and kernel-facing 0 is the reserved
-        every id space, so pad lanes must derive to swa slot 0 with no
+        padding slot in every id space, so pad lanes derive to swa slot 0 with
-        `num_live` bookkeeping."""
+        no `num_live` bookkeeping."""
        src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=3)
        padded = torch.cat([want_full, want_full.new_zeros(2)])
        got = self._field(src, rows, seq_lens, padded)
        self.assertTrue(torch.equal(got[:3], want_swa))
        self.assertTrue(bool((got[3:] == 0).all()), "pad lanes must land on slot 0")
    def test_slice_and_copy_derive_pointwise_without_handover(self):
        """REGRESSION (design): the retired identity-resolver refused any
        tensor it had not been handed -- a TBO child's re-padded slice or a
        registry's fresh copy raised. Value-based derivation must accept
        both, pointwise, with no adopt/handover call."""
        src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=4)
        padded = torch.cat([want_full, want_full.new_zeros(2)])
        # TBO-child shape: a slice crossing the pad boundary.
        got = self._field(src, rows, seq_lens, padded[2:6])
        self.assertTrue(torch.equal(got[:2], want_swa[2:4]))
        self.assertTrue(bool((got[2:] == 0).all()))
        # Registry shape: a fresh equal-value copy.
        got2 = self._field(src, rows, seq_lens, want_full.clone())
        self.assertTrue(torch.equal(got2, want_swa))
    def test_tombstoned_swa_page_clamps_to_sink(self):
        src, allocator, rows, seq_lens, virt, want_full, _ = self._built(ps=1, n=2)
        allocator.swa_v2p_page_table[int(virt[0])] = -1
@@ -11,22 +11,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Unit tests for the page-major envelope byte layout.
+"""``move_kv_cache_native`` -- the stock per-layer 3-D move static-pool
 compaction rides on -- must stay byte-exact. CPU-only.
-The subject here is the ENVELOPE — the byte layout the unified pool stores its
+The page-major envelope layout and the per-layer views over it are covered by
-KV in — pinned through ``MHASubPoolSpec``'s offset math. The 3-D per-layer views
+``test_unified_mha_views.py``, which pins the view addressing against the
-the pool exposes over the same bytes are covered by
+envelope formula byte for byte.
 ``test_unified_mha_views.py``, which also pins the view addressing
 against the envelope formula byte for byte.
 Verifies that:
 1. ``MHASubPoolSpec.layer_k_offset_in_page`` / ``layer_v_offset_in_page`` math
   matches the layout intent at ``page_size == 1`` and ``> 1``.
 2. ``move_kv_cache_native`` (the stock per-layer 3-D move) stays byte-exact.
 CPU-only — no GPU / Triton needed.
    python -m pytest test/registered/unit/mem_cache/test_layout_compat.py -v
 """
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -38,69 +28,12 @@ import unittest
 import torch
 from sglang.srt.mem_cache.memory_pool import move_kv_cache_native
 from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec
 _DEV = "cpu"
 def _make_mha_spec(name, grow, layer_num=2, head_num=2, head_dim=4):
    return MHASubPoolSpec(
        name=name,
        layer_num=layer_num,
        head_num=head_num,
        head_dim=head_dim,
        store_dtype=torch.float16,
        grow_direction=grow,
    )
 class TestMHASpecLayerOffsets(unittest.TestCase):
    """Verify ``layer_k_offset_in_page`` / ``layer_v_offset_in_page`` math."""
    def test_offsets_at_page_size_1_match_envelope(self):
        spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
        # At ps=1, layer-major within a 1-token page IS envelope-per-token.
        # Layer L's K offset = L * (k_row + v_row); V offset = +k_row.
        k_row = spec.k_row_bytes()
        v_row = spec.v_row_bytes()
        for L in range(spec.layer_num):
            self.assertEqual(
                spec.layer_k_offset_in_page(L, page_size=1),
                L * (k_row + v_row),
            )
            self.assertEqual(
                spec.layer_v_offset_in_page(L, page_size=1),
                L * (k_row + v_row) + k_row,
            )
    def test_offsets_at_page_size_gt_1(self):
        spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
        ps = 8
        k_row = spec.k_row_bytes()
        v_row = spec.v_row_bytes()
        # Layer L's K block within the page starts at L * ps * (k_row+v_row).
        # V block starts at +ps * k_row.
        for L in range(spec.layer_num):
            self.assertEqual(
                spec.layer_k_offset_in_page(L, page_size=ps),
                L * ps * (k_row + v_row),
            )
            self.assertEqual(
                spec.layer_v_offset_in_page(L, page_size=ps),
                L * ps * (k_row + v_row) + ps * k_row,
            )
    def test_page_bytes(self):
        spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
        # page_bytes = page_size * entry_bytes (preserved invariant)
        for ps in [1, 8, 64, 256]:
            self.assertEqual(spec.page_bytes(ps), ps * spec.entry_bytes())
 class TestMoveKVCacheNative(unittest.TestCase):
    def test_move_kv_cache_3d_path_unchanged(self):
        """The stock per-layer 3-D move must relocate exactly the named token
-        rows, byte-identically — compaction on static pools rides on it."""
+        rows, byte-identically; compaction on static pools rides on it."""
        k = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)]
        v = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)]
        for L in range(2):
@@ -11,31 +11,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Locked-full SWA tombstone-recovery under the unified pool (action handler).
+"""Locked-full SWA tombstone recovery under the unified pool (action handler).
-`RecoverSWAWithLockedFull` recovers a tombstoned SWA node whose full value is
+`RecoverSWAWithLockedFull` hands a tombstoned SWA node the INCOMING request's
-LOCKED: the node cannot adopt the incoming request's ids wholesale, so the
+swa pages and frees only their FULL pages. The static recipe re-points the
-static-pool recipe hands the node the INCOMING ids' swa pages, frees only their
+locked ids through `full_to_swa_index_mapping`; the unified composite has no
-FULL pages, and re-points the locked ids through `full_to_swa_index_mapping`.
+such tensor -- the swa sub-pool's v2p IS the mapping -- so it expresses the
 same move as a page-ownership rebind and then frees the incoming ids through
 the composite, whose `swa_v2p_pages > 0` filter skips the tombstoned swa side.
-The unified composite has no mapping tensor — the swa sub-pool's v2p IS the
+The recovery must succeed rather than decline: the TreeCore insert walk counts
-mapping — and its `set_full_to_swa_mapping` is an explicit no-op stub. The
+the node in `prefix_len` regardless of component consumption, while the SWA
-pre-fix handler therefore raised AttributeError on `full_to_swa_index_mapping`
+match validator rejects a `value is None` node, so a declined recovery reports
-(and, had that line been removed, would have silently skipped the rebind while
+a prefix `match_prefix` cannot honor and trips
 line 1 freed swa pages the kept ids still referenced). The fix expresses the
 same move as a page-ownership REBIND: bind the node's virtual pages to the
 incoming pages' physical pages, tombstone the incoming ones, then free the
 incoming ids through the composite — whose `swa_v2p_pages > 0` filter skips the
 tombstoned swa side, releasing ONLY the full side.
 Why the recovery must succeed rather than decline (the v1 lesson, still true on
 this branch): the TreeCore insert walk counts the node in `prefix_len`
 regardless of component consumption, while the SWA match validator rejects a
 `value is None` node — a declined recovery makes `insert` report a prefix the
 follow-up `match_prefix` cannot honor, tripping
 `new_prefix_len <= len(new_indices)` in `cache_unfinished_req`.
    python -m pytest test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py -v
 """
 import unittest
@@ -121,8 +110,8 @@ class _Probe(SWAComponent):
 class _StaticAllocRecorder:
    """Stands in for the STATIC SWATokenToKVPoolAllocator: has the mapping
-    tensor and a real set_full_to_swa_mapping. The handler must keep routing
+    tensor and a real `set_full_to_swa_mapping`, so the handler must keep
-    static pools through the original recipe."""
+    routing static pools through the original recipe."""
    def __init__(self, n=16):
        self.full_to_swa_index_mapping = torch.arange(n, dtype=torch.int64)
@@ -133,9 +122,8 @@ class _StaticAllocRecorder:
        self.full_attn_allocator = self
    def set_full_to_swa_mapping(self, full, swa):
-        # Honour the write like the real static allocator: the handler routes
+        # The handler routes every mapping write THROUGH the API, so the fake
-        # every mapping write THROUGH the API (never by indexing the tensor),
+        # must apply it for the mapping asserts to observe anything.
        # so the fake must apply it for the mapping asserts to observe it.
        self.mapping_calls.append((full, swa))
        self.full_to_swa_index_mapping[full.to(torch.int64)] = swa.to(torch.int64)
@@ -171,9 +159,9 @@ class _RecoverTestBase(unittest.TestCase):
 class TestPagePairing(_RecoverTestBase):
    def test_pairs_positionally_not_by_sorted_id(self):
-        """Allocation hands out virtual ids in no particular order; deduping
+        """Allocation hands out virtual ids in no particular order, so a
-        with `torch.unique` (which sorts) would bind the node's page k to an
+        `torch.unique` dedup (which sorts) would bind the node's page k to an
-        unrelated incoming page — silent wrong-KV."""
+        unrelated incoming page."""
        probe, _ = self._probe()
        kept = torch.tensor([9, 7, 5], dtype=torch.int64)  # descending
        incoming = torch.tensor([2, 4, 6], dtype=torch.int64)  # ascending
@@ -225,15 +213,17 @@ class TestOwnershipTransfer(_RecoverTestBase):
 class TestRecoverActionHandler(_RecoverTestBase):
    def test_recovery_sets_a_live_device_value_and_frees_only_the_full_side(self):
-        """End-to-end through apply_component_action — the pre-fix handler
+        """Bug regression: recovery must give the node a LIVE swa value and
-        raises AttributeError (`full_to_swa_index_mapping`) on this exact
+        return only the incoming ids' FULL side, allocating and freeing no swa
-        call. Post-fix: the node gets a LIVE swa value, the HANDLER neither
+        page of its own."""
        allocates nor frees any swa page (ownership only moves), and the
        incoming ids' FULL side returns to the pool."""
        probe, allocator = self._probe()
        swa = allocator.swa_attn_allocator
        kept, incoming = self._two_ranges(allocator)
        allocator.free_swa(kept)  # what eviction does when it tombstones
        self.assertTrue(
            bool((allocator.translate_loc_from_full_to_swa(kept) == 0).all()),
            "precondition: a tombstoned range translates to the sink",
        )
        # Snapshot AFTER the setup traffic: the invariant under test is that
        # the recovery handler itself moves ownership without moving capacity.
        swa_live = swa.allocated_count()
@@ -260,21 +250,6 @@ class TestRecoverActionHandler(_RecoverTestBase):
            full_avail + len(incoming),
            "the incoming ids' FULL side must come back",
        )
    def test_recovered_ids_translate_to_live_pages_not_the_sink(self):
        """The tombstoned range translates to the clamped sink before the
        recovery and to real pages after — recovering from the node's OWN
        already-freed ids (instead of the donated ones) reintroduces the sink."""
        probe, allocator = self._probe()
        kept, incoming = self._two_ranges(allocator)
        allocator.free_swa(kept)
        self.assertTrue(
            bool((allocator.translate_loc_from_full_to_swa(kept) == 0).all()),
            "precondition: a tombstoned range translates to the sink",
        )
        probe.apply_component_action(
            RecoverSWAWithLockedFull(node_id=1, kept_full=kept, incoming_full=incoming)
        )
        self.assertTrue(
            bool((allocator.translate_loc_from_full_to_swa(kept) > 0).all()),
            "after recovery the node's ids must address live swa pages",
@@ -283,8 +258,8 @@ class TestRecoverActionHandler(_RecoverTestBase):
 class TestStaticPoolPathUnchanged(unittest.TestCase):
    def test_static_allocator_keeps_the_mapping_recipe(self):
-        """A static SWA allocator (has the mapping tensor) must keep the
+        """A static SWA allocator must keep the mapping recipe; the unified
-        original recipe — the unified branch must not hijack it."""
+        branch must not hijack it."""
        static = _StaticAllocRecorder()
        probe = _Probe.__new__(_Probe)
        probe.cache = _Cache(static)
@@ -296,11 +271,8 @@ class TestStaticPoolPathUnchanged(unittest.TestCase):
            RecoverSWAWithLockedFull(node_id=3, kept_full=kept, incoming_full=incoming)
        )
-        # Both mapping writes go through the allocator API -- the kept remap
+        # Both mapping writes go through the allocator API, never by indexing
-        # via set_full_to_swa_mapping, the incoming tombstone via
+        # `full_to_swa_index_mapping` (absent on the unified composite).
        # clear_full_to_swa_mapping -- never by indexing
        # `full_to_swa_index_mapping` (the tensor is absent on the unified
        # composite by design).
        self.assertEqual(len(static.mapping_calls), 1, "static recipe must run")
        self.assertEqual(len(static.clear_calls), 1, "incoming must be tombstoned")
        self.assertTrue(
@@ -11,55 +11,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""`SWAKVPool.get_v_head_dim()` — the pool method a mambaish + hybrid-SWA
+"""`SWAKVPool.get_v_head_dim()` -- the pool method a mambaish + hybrid-SWA model
-model reaches on boot.
+reaches on boot.
-`TritonAttnBackend.__init__` picks its `v_head_dim` from one of three
+`TritonAttnBackend.__init__` asks the POOL for `v_head_dim` when the model is
-branches, and the middle one asks the POOL:
+mambaish and its full/SWA value head dims MATCH (Inkling-class), so an
-
+SWA-shaped pool without the method kills backend construction with
-    if sliding_window_size is not None and swa_v_head_dim != v_head_dim:
+AttributeError.
        ... from model_config ...                     # asymmetric hybrid SWA
    elif mambaish_config(model_config) is not None:
        v_head_dim = token_to_kv_pool.get_v_head_dim()   # <-- this one
    else:
        ... from get_value_buffer(start_layer) ...
 A model that is BOTH mambaish AND hybrid-SWA with MATCHING full/SWA value
 head dims (Inkling-class) skips the first branch and lands in the second —
 where its pool is an SWA-shaped pool, which had no `get_v_head_dim`. The
 server died at backend construction with
    AttributeError: 'SWAKVPool' object has no attribute 'get_v_head_dim'
 on the STATIC pool and, identically, on `UnifiedSWAKVPool`. Neither the
 mamba-hybrid pools (`HybridLinearKVPool` has the method) nor pure hybrid-SWA
 models (not mambaish, so the branch is never taken) can reach it, which is
 why it went unnoticed.
    python -m pytest test/registered/unit/mem_cache/test_swa_pool_v_head_dim.py -v
 """
 import inspect
 import unittest
 import torch
 from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
 from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=15, suite="base-a-test-cpu")
 _DEV = "cpu"
 _FULL_V_HEAD_DIM = 8
-_SWA_V_HEAD_DIM = 8  # MATCHING — this is what routes Inkling into the branch
+_SWA_V_HEAD_DIM = 8  # MATCHING: this is what routes Inkling into the branch
 def _swa_pool():
-    """A static SWAKVPool with the Inkling-class layer split: full and SWA
+    """Inkling-class layer split: full and SWA layers interleaved, layer 0 NOT a
-    layers interleaved, layer 0 NOT a full-attention layer (which is exactly
+    full-attention layer, which is why the backend asks the pool at all."""
    why the backend asks the pool instead of indexing layer 0)."""
    return SWAKVPool(
        size=32,
        size_swa=16,
@@ -76,49 +53,9 @@ def _swa_pool():
 class TestSWAPoolVHeadDim(unittest.TestCase):
    def test_static_pool_reports_the_full_side_value_head_dim(self):
        """Red before the fix with AttributeError; the value must be the FULL
        side's, since that is the geometry the caller means."""
        pool = _swa_pool()
        self.assertEqual(pool.get_v_head_dim(), _FULL_V_HEAD_DIM)
    def test_answer_matches_the_full_pool_buffer_not_layer_zero(self):
        """Layer 0 is an SWA layer here, so a naive `get_value_buffer(0)`
        would read the SWA side. Pin that the method routes through the FULL
        sub-pool at its own start_layer — the property that makes it correct
        under pipeline parallelism too."""
        pool = _swa_pool()
        want = pool.full_kv_pool.get_value_buffer(pool.full_kv_pool.start_layer).shape[
            -1
        ]
        self.assertEqual(pool.get_v_head_dim(), want)
        # And layer 0 really is the SWA side in this fixture.
        _, is_swa = pool.layers_mapping[0]
        self.assertTrue(is_swa, "fixture must keep layer 0 on the SWA side")
    def test_unified_swa_pool_inherits_it(self):
        """`UnifiedSWAKVPool` subclasses `SWAKVPool`, so the unified tri-pool
        path (mambaish + hybrid SWA in one buffer) is covered by the same
        method — no second implementation to drift."""
        self.assertTrue(issubclass(UnifiedSWAKVPool, SWAKVPool))
        self.assertIs(
            UnifiedSWAKVPool.get_v_head_dim,
            SWAKVPool.get_v_head_dim,
            "the unified pool must inherit the method, not shadow it",
        )
    def test_signature_matches_the_hybrid_linear_precedent(self):
        """The backend calls this method on whichever pool it holds, so every
        pool reachable from the mambaish branch must expose the SAME
        zero-argument shape. `HybridLinearKVPool` is the precedent this one
        mirrors; a future pool added to that branch has to match too."""
        for cls in (SWAKVPool, HybridLinearKVPool):
            sig = inspect.signature(cls.get_v_head_dim)
            self.assertEqual(
                [p for p in sig.parameters if p != "self"],
                [],
                f"{cls.__name__}.get_v_head_dim must take no arguments",
            )
 if __name__ == "__main__":
    unittest.main()
@@ -11,28 +11,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Byte-conservation verifier for the unified 2-pool composites.
+"""Byte-conservation verifier (`verify_byte_accounting`) for the unified
 2-pool composites.
-`verify_byte_accounting` is the idle-time tripwire the token-identity leak
+The unified pool's correctness rests on BYTE bookkeeping -- watermark spans,
-check cannot provide: the unified pool's correctness rests on BYTE bookkeeping
+holes, pending compaction, frontier ordering inside one shared buffer -- and
-(watermark spans, holes, pending compaction, frontier ordering inside one
+a drifted counter admits requests into memory that is not actually free.
-shared buffer), and a drifted counter admits requests into memory that is not
+Nothing about that is visible to the token-identity leak check, so this
-actually free — silent corruption territory, not a crash.
+verifier is the only idle-time tripwire for it.
-Derived properties pinned here:
+The conservation identity: on a lazy end pool the watermark span equals
-
+live + holes + pending pages at EVERY point of a lifecycle, not just at rest.
  * Conservation: on a lazy end pool the watermark span must equal
    live + holes + pending pages at EVERY point of a healthy lifecycle
    (alloc, partial free, group free, flush) — not just at rest.
  * The check is not vacuous: drifting any single term (live count, watermark,
    a leaked hole) reports loudly, naming the sub-pool.
  * Chain order: one member's low frontier clearing the other's high frontier
    is what "two pools share one buffer without overlap" MEANS; the pair check
    must hold regardless of which member grows up.
  * The strict escalation env defaults OFF: promoting the diagnostic to a
    RuntimeError is a validation posture, not the production one.
    python -m pytest test/registered/unit/mem_cache/test_unified_byte_accounting.py -v
 """
 import unittest
@@ -62,18 +51,6 @@ def _paged_pair(lazy: bool):
 class TestHealthyLifecycleReportsClean(unittest.TestCase):
    def test_swa_composite_clean_at_every_step(self):
        inst, allocator, kvcache = _swa_composite()
        self.assertEqual(allocator.verify_byte_accounting(), [])
        v = inst._alloc(allocator, kvcache, 8)
        self.assertEqual(allocator.verify_byte_accounting(), [])
        allocator.free_swa(v[:4])  # tombstone half the swa side
        self.assertEqual(allocator.verify_byte_accounting(), [])
        inst._free(allocator, kvcache, v)
        self.assertEqual(allocator.verify_byte_accounting(), [])
        allocator.clear()
        self.assertEqual(allocator.verify_byte_accounting(), [])
    def test_lazy_end_pool_clean_through_free_and_flush(self):
        full, _swa = _paged_pair(lazy=True)
        self.assertEqual(full._byte_accounting_violations(), [])
@@ -86,10 +63,10 @@ class TestHealthyLifecycleReportsClean(unittest.TestCase):
 class TestDriftReportsLoudly(unittest.TestCase):
-    """Each mutation below models a distinct bookkeeping bug; the verifier
+    """Each mutation models a distinct bookkeeping bug; the verifier must name
-    must name the drifted sub-pool. Without these, a regression in any single
+    the drifted sub-pool. Without these, a regression in any single counter
-    counter passes every other test (the pool still 'works' — it just lies
+    passes every other test -- the pool still 'works', it just lies about
-    about capacity)."""
+    capacity."""
    def _lazy_full(self):
        full, _swa = _paged_pair(lazy=True)
@@ -98,20 +75,27 @@ class TestDriftReportsLoudly(unittest.TestCase):
        self.assertEqual(full._byte_accounting_violations(), [])
        return full
-    def test_drifted_live_count(self):
+    def test_any_drifted_term_reports(self):
-        full = self._lazy_full()
+        def drift_live_count(full):
-        full.live_page_count += 1
+            full.live_page_count += 1
        self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
-    def test_leaked_hole(self):
+        def leak_hole(full):
-        full = self._lazy_full()
+            full._free_phys_pages = full._free_phys_pages[:-1]  # hole vanished
        full._free_phys_pages = full._free_phys_pages[:-1]  # hole vanished
        self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
-    def test_drifted_watermark(self):
+        def drift_watermark(full):
-        full = self._lazy_full()
+            full.watermark_physical += 1
-        full.watermark_physical += 1
+
-        self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
+        for term, mutate in (
            ("live_count", drift_live_count),
            ("leaked_hole", leak_hole),
            ("watermark", drift_watermark),
        ):
            with self.subTest(term=term):
                full = self._lazy_full()
                mutate(full)
                self.assertTrue(
                    any("span" in s for s in full._byte_accounting_violations())
                )
    def test_composite_report_names_the_sub_pool(self):
        """Frontier-bounds drift (checked in BOTH lazy and eager modes): push
@@ -129,11 +113,10 @@ class TestDriftReportsLoudly(unittest.TestCase):
 class TestChainFrontierOrder(unittest.TestCase):
    def test_overlapping_frontiers_report(self):
-        """Both bands hold pages, then the up member's watermark is pushed past
+        """Both bands hold pages, then the up member's watermark is pushed
-        the down member's LIVE low frontier: the two bands now claim the same
+        past the down member's LIVE low frontier. Both sides must be populated
-        bytes of one buffer. (An empty down band cannot overlap — its low
+        for this to be real corruption: an empty down band cannot overlap,
-        frontier IS the buffer top — so both sides must be populated for the
+        its low frontier IS the buffer top."""
        scenario to be a real corruption.)"""
        full, swa = _paged_pair(lazy=False)
        chain = mea._end_pair_chain(full, swa)
        up, down = chain
@@ -144,14 +127,6 @@ class TestChainFrontierOrder(unittest.TestCase):
        out = mea._chain_byte_accounting_violations(chain)
        self.assertTrue(any("overlap" in s for s in out), out)
    def test_pair_order_is_direction_agnostic(self):
        """The factories and the unit fixtures orient the pair differently;
        the check must order by grow direction, not by argument position."""
        full, swa = _paged_pair(lazy=False)
        a = mea._end_pair_chain(full, swa)
        b = mea._end_pair_chain(swa, full)
        self.assertEqual([x.sub_pool_name for x in a], [x.sub_pool_name for x in b])
 if __name__ == "__main__":
    unittest.main()
@@ -13,26 +13,17 @@
 # ==============================================================================
 """Byte-budget buffer sizing for the unified 2-pool factories.
-Derived properties pinned here:
+With ``unified_total_bytes`` set the buffer is that many bytes exactly (the
 mamba pair adds the state pool's bytes on top -- the budget is captured AFTER
 the state carve-out); without it, sizing falls back to the token-count re-sum.
 Sizing from the ratio-derived token counts instead would re-introduce the
 configurator's rounding, which floors by the cell size and then page-aligns
 EACH side, losing up to about a page of tokens per side.
-  * Budget honored EXACTLY: with ``unified_total_bytes`` set, the swa pair's
+The bs=1 feasibility floor raises at BOOT, before any pool construction: a
-    buffer is that many bytes (the mamba pair adds the state pool's bytes on
+budget that cannot fit ONE worst-case request (full KV at max context, plus
-    top — the budget is captured AFTER the state carve-out). Sizing from the
+one SWA window / the state slots a running request locks) is a retract
-    ratio-derived token counts instead re-introduces the configurator's
+LIVELOCK at runtime, not a perf bug.
    rounding: the swa split floors the budget by the cell size and then
    page-aligns EACH side's token count, so the re-sum reconstructs less
    than the profiled budget by up to about one page of tokens per side.
  * Fallback: without the budget, sizing is the historical token-count re-sum,
    bit-for-bit.
  * bs=1 feasibility floor: a budget that cannot fit ONE worst-case request
    (full KV at max context, plus one SWA window / the state slots a single
    running request locks) raises at BOOT, before any pool construction —
    under-sizing is a retract LIVELOCK at runtime, not a perf bug.
  * The 4096-byte alignment exists because the factories ``.view()`` the whole
    uint8 buffer as the KV dtype; an unaligned budget must be floored, never
    rounded up (rounding up overcommits profiled memory).
    python -m pytest test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py -v
 """
 import unittest
@@ -92,37 +83,32 @@ def _entry_bytes():
 class TestBudgetSizing(unittest.TestCase):
    def test_swa_factory_honors_the_budget_exactly(self):
        """The buffer IS the budget, including values the token-count re-sum
        cannot represent (not a whole-token multiple per side)."""
        e = _entry_bytes()
-        budget = 96 * e + 512  # deliberately NOT a token-count multiple
+        for budget in (
-        bundle = _swa_factory(unified_total_bytes=budget)
+            96 * e + 512,  # deliberately NOT a token-count multiple
-        self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
+            (64 + 32) * e + (e - 2),  # almost one more entry
        ):
            with self.subTest(budget=budget):
                bundle = _swa_factory(unified_total_bytes=budget)
                self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
    def test_fallback_is_the_token_count_resum(self):
        e = _entry_bytes()
        bundle = _swa_factory()
        self.assertEqual(bundle.unified_memory_pool.total_bytes, (64 + 32) * e)
    def test_budget_beats_resum_on_rounding(self):
        """The property that motivates the whole phase: the re-sum cannot
        represent a budget that is not a whole-token multiple per side, so it
        strands bytes the buffer could have held."""
        e = _entry_bytes()
        budget = (64 + 32) * e + (e - 2)  # almost one more entry
        bundle = _swa_factory(unified_total_bytes=budget)
        self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
        self.assertGreater(budget, (64 + 32) * e)
 class TestReservedFloorIsOneSourceOfTruth(unittest.TestCase):
    """The bs=1 floor charges the slot-0 sink, and MUST charge exactly what
    `UnifiedKVPool` actually reserves.
-    Regression (GPU eval_434/436, Falcon-H1 boot): the floor hand-copied the
+    Regression (Falcon-H1 boot): the floor hand-copied the formula as
-    formula as `page_size * max(entry_bytes)`, applying the page multiplier to
+    `page_size * max(entry_bytes)`, applying the page multiplier to the MAMBA
-    the MAMBA spec. The pool deliberately excludes mamba (it is page_size=1),
+    spec, which the pool deliberately excludes because it is page_size=1. At
-    so with page_size=256 and a ~139 MB state entry the floor over-charged the
+    page_size=256 that over-charged the sink 256x and a healthy config failed
-    sink by 256x — ~33 GiB of phantom requirement — and a healthy config
+    to boot with real headroom to spare.
    failed to boot with 25 GiB of real headroom.
    """
    def _specs(self, page_size):
@@ -134,7 +120,7 @@ class TestReservedFloorIsOneSourceOfTruth(unittest.TestCase):
            store_dtype=torch.float16,
            grow_direction="down",
        )
-        # A state entry vastly larger than a KV token entry — the real ratio
+        # A state entry vastly larger than a KV token entry: the real ratio
        # (~139 MB vs ~45 KB) is what made the over-charge fatal.
        mamba = MambaSubPoolSpec(
            name="mamba",
@@ -187,51 +173,22 @@ class TestBs1FeasibilityFloor(unittest.TestCase):
        self.assertIn("bs=1 floor", str(ctx.exception))
        self.assertIn("swa_window_kv", str(ctx.exception))
-    def test_context_longer_than_the_pool_is_not_rejected(self):
+    def test_feasible_floor_inputs_are_not_rejected(self):
-        """REGRESSION: the floor must NOT charge the full-attention token side.
+        """Floor inputs that must NOT raise: charging the full-attention token
-        `TpModelWorker.get_worker_info` clamps max_req_len to the pool, so a
+        side, or a window the context never clamps, failed these at boot."""
        context far larger than the buffer is refused at admission, not a
        livelock -- and it is an ordinary way to serve a long-context model on
        one GPU. Charging it here made such configs fail at boot."""
        e = _entry_bytes()
-        bundle = _swa_factory(
+        for case, model_context_len, sliding_window_size in (
-            unified_total_bytes=200 * e,
+            ("context_longer_than_the_pool", 1_000_000, 16),
-            model_context_len=1_000_000,  # far beyond what the buffer holds
+            ("feasible_config", 64, 16),
-            sliding_window_size=16,
+            ("window_larger_than_context", 64, 10_000),
-        )
+        ):
-        self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e)
+            with self.subTest(case=case):
-
+                bundle = _swa_factory(
-    def test_feasible_config_boots_with_floor_inputs_present(self):
+                    unified_total_bytes=200 * e,
-        e = _entry_bytes()
+                    model_context_len=model_context_len,
-        bundle = _swa_factory(
+                    sliding_window_size=sliding_window_size,
-            unified_total_bytes=200 * e,
+                )
-            model_context_len=64,
+                self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e)
            sliding_window_size=16,
        )
        self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e)
    def test_window_term_is_clamped_to_context(self):
        """A window larger than the context must charge at most the context —
        otherwise short-context models over-raise."""
        e = _entry_bytes()
        bundle = _swa_factory(
            unified_total_bytes=200 * e,
            model_context_len=64,
            sliding_window_size=10_000,  # window >> context
        )
        self.assertIsNotNone(bundle)
    def test_floor_message_itemizes_terms(self):
        with self.assertRaises(RuntimeError) as ctx:
            _check_bs1_feasibility_floor(
                total_bytes=10,
                floor_terms=[("a", 8), ("b", 8)],
                factory="test",
            )
        msg = str(ctx.exception)
        self.assertIn("a=8", msg)
        self.assertIn("b=8", msg)
        self.assertIn("16", msg)
    def test_exact_floor_passes(self):
        """Boundary: total == floor must NOT raise (>= is the contract)."""
@@ -11,29 +11,17 @@
 # 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).
+"""Epoch-memoized capacity views on the allocator chain.
-The capacity views (`available_size` / `schedulable_available_size` per band,
+The per-band and composite capacity views are pure functions of a handful of
-plus the composite joint view) are pure functions of a handful of
+CPU-resident fields that schedulers read O(queue) times between mutations;
-CPU-resident fields across the chain; schedulers read them O(queue) times
+`_CapacityField` descriptors bump `_capacity_epoch` on every rebind, so the
-between mutations. `_CapacityField` descriptors bump `_capacity_epoch` on
+memos invalidate by construction.
 every rebind, so the memos invalidate by construction.
-The failure mode being guarded: a memo serving a STALE value after a mutation
+The guarded failure mode is a memo serving a STALE value after a mutation the
-the epoch machinery missed — either a new mutation site writing a field the
+epoch machinery missed -- a new mutation site writing an uncovered field, or
-descriptors don't cover, or an in-place write that bypasses `__set__`. Stale
+an in-place write that bypasses `__set__`. Readers cannot detect either one;
-capacity is silent over-/under-admission, not a crash. Hence:
+stale capacity is silent over-/under-admission, not a crash.
  * 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
@@ -114,10 +102,8 @@ class TestCapacityMemoCoherence(unittest.TestCase):
                self._assert_memos_fresh(allocator)
    def test_random_op_sequence_value_identity(self):
-        """Property: after ANY mutation sequence, memoized views equal fresh
+        """Property: after ANY mutation sequence the memoized views equal
-        recomputes. Seeded (deterministic) — the sequences cover interleavings
+        fresh recomputes. Seeded, so the interleavings are deterministic."""
        (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):
@@ -162,29 +148,13 @@ class TestCapacityMemoCoherence(unittest.TestCase):
            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 --
+        """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
+        the span fields are its ONLY capacity state, so unless they are
-        `_CapacityField` descriptors, the float's own memo AND both
+        `_CapacityField` descriptors the float's own memo AND both neighbours'
-        neighbours' (the span flips transparency, walling off their gaps)
+        (the span flips transparency, walling off their gaps) keep serving
-        keep serving pre-move values.
+        pre-move values. Exercised on a hand-wired end+float+end chain so no
-
+        end-pool descriptor write can mask a missing span bump."""
        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(
@@ -212,21 +182,12 @@ class TestCapacityMemoCoherence(unittest.TestCase):
        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
+    invalidate the composite memo -- including the two only the tri has, a
-    has: a mamba-end state draw and a float span move behind the composite."""
+    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
@@ -281,19 +242,6 @@ class TestTriCapacityMemoCoherence(unittest.TestCase):
                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()
@@ -11,28 +11,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""The unified free path must not move anything across the PCIe bus.
+"""The unified allocator free path must not sync the host.
-Two independent host syncs lived in `MultiEndedAllocator`'s free path:
+Two syncs are guarded. A tombstone written as ``t[idx] = -1`` makes torch
 materialise the scalar as a CPU tensor and copy it H2D, which BLOCKS the host
 until the stream drains. And `torch.unique`, recovering distinct PAGE ids from
 freed TOKEN ids, has a data-dependent output shape and so must D2H the count;
 `free_segment` instead takes stride slices off the caller's `start_pos`, a
 page's tokens sitting consecutively in the kv row.
-  1. Tombstone scatters written as ``t[idx] = -1``. The scalar RHS makes torch
+Mirrors `test_paged_free_segment.py`, which pins the same properties for
-     materialise ``-1`` as a CPU tensor and copy it H2D, and a pageable H2D
+`PagedTokenToKVPoolAllocator`.
     copy BLOCKS the host until the stream drains. Invisible on decode-shaped
     work; ~16 ms per free behind an 8192-token prefill.
  2. `torch.unique` recovering distinct PAGE ids from freed TOKEN ids. Its
     output shape is data-dependent, so it must D2H the count
     (``_unique2 -> item -> _local_scalar_dense -> cudaStreamSynchronize``).
     `PagedTokenToKVPoolAllocator` already solved this with `free_segment`:
     a page's tokens sit consecutively in the kv row, so given `start_pos` the
     page representatives are stride slices. The unified allocators simply
     never implemented it and so were permanently on the syncing path.
 These tests mirror `test_paged_free_segment.py` -- the same sweep against the
 `torch.unique` reference, the same free-group deferral -- because the unified
 allocators now mirror that allocator's design rather than a parallel one.
    python -m pytest test/registered/unit/mem_cache/test_unified_free_no_host_sync.py -v
 """
 import ast
@@ -68,10 +57,9 @@ def _paged_allocator(lazy: bool):
 _TABLES = {"virtual_to_physical", "physical_to_virtual"}
-# Methods that MUST tombstone through index_fill_. Explicit, because "this
+# Methods that MUST tombstone through index_fill_; hand-listed because "writes
-# method writes a tombstone" is a design fact per method, not something a scan
+# a tombstone" is a per-method design fact a scan cannot infer. Completeness is
-# can infer -- but `test_every_allocator_free_path_is_listed` below fails if a
+# guarded by `test_every_allocator_free_path_is_listed` below.
 # new allocator arrives with its own free path and is not added here.
 _TOMBSTONE_METHODS = [
    (mea.MultiEndedAllocator, "_free_lazy"),
    (mea.MultiEndedAllocator, "free"),
@@ -82,9 +70,8 @@ _TOMBSTONE_METHODS = [
 ]
-# Ways to write `-1` into a table without the value crossing the bus:
+# Sanctioned no-sync tombstone forms: `index_fill_` takes the scalar as an
-# `index_fill_` takes the scalar as an argument torch keeps off the host, and
+# argument torch keeps off the host; the fused launcher stores it in-kernel.
 # the fused launchers store it from inside the kernel.
 _NO_SYNC_TOMBSTONE_FORMS = ("index_fill_", "free_unbind_inplace")
@@ -109,9 +96,8 @@ def _allocators_in_module():
 def _table_touching_methods():
    """Every own method of every allocator whose source names a page table.
-    DISCOVERY, not a list: a hardcoded list stops guarding the moment a new
+    Discovery rather than a hand list, so a new allocator class with its own
-    allocator class arrives with its own free path -- which is what happened
+    free path is covered the day it lands.
    when FloatMultiEndedAllocator was added and inherited no coverage.
    """
    out = []
    for cls in _allocators_in_module():
@@ -130,15 +116,13 @@ def _table_touching_methods():
 def _scalar_index_assignments(fn):
    """`self.<table>[<tensor idx>] = <scalar>` occurrences in fn's source.
-    Slice assignments (``t[a:b] = -1``) are excluded: a slice is a view, so the
+    Slices (``t[a:b] = -1``) and tensor-valued scatters are excluded: only a
-    fill needs no index tensor. Tensor-valued scatters are excluded too -- only
+    scalar RHS behind a tensor index materialises a CPU value tensor.
    the scalar RHS materialises a CPU value tensor.
    """
    def _is_scalar_literal(node):
-        # NOTE: `-1` parses as UnaryOp(USub, Constant(1)), NOT Constant. Testing
+        # `-1` parses as UnaryOp(USub, Constant(1)), not Constant; matching
-        # only for Constant silently skips every negative literal -- i.e. every
+        # only Constant silently skips every negative literal.
        # tombstone this scan exists to find.
        if isinstance(node, ast.Constant):
            return True
        return isinstance(node, ast.UnaryOp) and isinstance(node.operand, ast.Constant)
@@ -156,9 +140,8 @@ def _scalar_index_assignments(fn):
                continue
            if isinstance(tgt.slice, ast.Slice):
                continue
-            # A CONSTANT integer index (`t[0] = 0`, `t[-1] = -1`) is a
+            # A constant index (`t[0] = 0`) is a single-element sentinel write,
-            # single-element sentinel write, not the tensor-index tombstone this
+            # not a tensor-index tombstone; only `clear()` does it.
            # guard exists to find, and the only such writes are in `clear()`.
            if _is_scalar_literal(tgt.slice):
                continue
            bad.append(ast.unparse(node))
@@ -167,6 +150,13 @@ def _scalar_index_assignments(fn):
 class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
    def test_no_scalar_index_assignment(self):
        # Self-check first: a scan whose AST matching has drifted reports
        # clean on every method below.
        def _offender(self):
            self.virtual_to_physical[free_v_pages] = -1  # noqa: F821
        self.assertEqual(len(_scalar_index_assignments(_offender)), 1)
        discovered = _table_touching_methods()
        self.assertGreaterEqual(
            len(discovered), len(_TOMBSTONE_METHODS), "discovery scan went blind"
@@ -185,24 +175,10 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
                    ),
                )
    def test_the_scan_detects_the_scalar_form_it_guards(self):
        """Self-check. The scan is only as good as its AST matching, and it
        silently missed every tombstone until `-1` was recognised as
        UnaryOp(USub, Constant) rather than Constant. Pin that."""
        def _offender(self):
            self.virtual_to_physical[free_v_pages] = -1  # noqa: F821
        self.assertEqual(len(_scalar_index_assignments(_offender)), 1)
    def test_every_allocator_free_path_is_listed(self):
-        """The positive list must name every allocator that owns a free path.
+        """Bug regression: an allocator that owns a free path but is missing
-
+        from `_TOMBSTONE_METHODS` must fail loudly here rather than drop out of
-        REGRESSION: the list used to hold three MultiEndedAllocator methods, so
+        tombstone coverage."""
        adding FloatMultiEndedAllocator with its own `free` silently dropped that
        free path out of coverage -- and it shipped a scalar tombstone. Fail here
        instead, loudly, the next time an allocator arrives.
        """
        listed = {(cls.__name__, name) for cls, name in _TOMBSTONE_METHODS}
        for cls in _allocators_in_module():
            for name, fn in vars(cls).items():
@@ -212,9 +188,8 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
                    src = inspect.getsource(fn)
                except OSError:
                    continue
-                # WRITES a page table -- either correctly (index_fill_) or in the
+                # Only a method that WRITES a page table needs a tombstone;
-                # banned scalar form the scan below catches. A method that only
+                # one that merely READS has nothing to guard.
                # READS a table has nothing to tombstone.
                if not (
                    any(f"{t}.index_fill_" in src for t in _TABLES)
                    or _scalar_index_assignments(fn)
@@ -231,10 +206,8 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
                )
    def test_free_paths_actually_write_a_tombstone(self):
-        """Positive form, so deleting the scatter entirely cannot pass. The
+        """Positive form, so deleting the scatter entirely cannot pass; a new
-        mechanism is not the point -- keeping the tombstone value off the host
+        no-sync mechanism is added to `_NO_SYNC_TOMBSTONE_FORMS` deliberately."""
        is -- so this lists the sanctioned ways to do that and a new one is
        added here deliberately."""
        for cls, name in _TOMBSTONE_METHODS:
            with self.subTest(method=f"{cls.__name__}.{name}"):
                src = inspect.getsource(getattr(cls, name))
@@ -244,22 +217,6 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
                    f"{_NO_SYNC_TOMBSTONE_FORMS}",
                )
    def test_index_fill_matches_scalar_assign_semantics(self):
        """Behaviour-preserving, including the edge cases the free path hands
        it: empty index, duplicate pages, full table."""
        for idx in (
            torch.tensor([], dtype=torch.int64),
            torch.tensor([1, 3, 5], dtype=torch.int64),
            torch.tensor([2, 2, 3], dtype=torch.int64),  # duplicates
            torch.arange(6, dtype=torch.int64),
        ):
            with self.subTest(n=int(idx.numel())):
                a = torch.arange(6, dtype=torch.int64)
                b = a.clone()
                a[idx] = -1
                b.index_fill_(0, idx, -1)
                self.assertTrue(torch.equal(a, b))
 # --------------------------------------------------------------------------
 # 2. free_segment: stride page extraction instead of torch.unique
@@ -270,7 +227,6 @@ class TestFreeSegment(unittest.TestCase):
    """Mirrors `test_paged_free_segment.TestFreeSegment`."""
    def test_matches_unique_over_tail_alignments(self):
        """Page-aligned starts against every tail alignment."""
        for num_tokens in (1, PAGE_SIZE, PAGE_SIZE + 1, 3 * PAGE_SIZE - 1):
            for start in range(0, num_tokens, PAGE_SIZE):
                for end in (start + 1, num_tokens):
@@ -288,8 +244,6 @@ class TestFreeSegment(unittest.TestCase):
                        self.assertEqual(freed.numel(), expected.numel())
    def test_never_calls_unique(self):
        """The decisive check -- make `torch.unique` explode. A textual guard
        can be fooled; this cannot."""
        for start in (0, PAGE_SIZE, 2 * PAGE_SIZE):
            alloc = _paged_allocator(lazy=True)
            row = alloc.alloc(3 * PAGE_SIZE)
@@ -299,12 +253,6 @@ class TestFreeSegment(unittest.TestCase):
                ):
                    alloc.free_segment(row[start : start + PAGE_SIZE], start_pos=start)
    def test_unaligned_start_is_rejected(self):
        alloc = _paged_allocator(lazy=True)
        row = alloc.alloc(3 * PAGE_SIZE)
        with self.assertRaises(AssertionError):
            alloc.free_segment(row[1 : PAGE_SIZE + 1], start_pos=1)
    def test_empty_segment_is_noop(self):
        alloc = _paged_allocator(lazy=True)
        before = alloc._free_phys_pages.numel()
@@ -312,8 +260,8 @@ class TestFreeSegment(unittest.TestCase):
        self.assertEqual(alloc._free_phys_pages.numel(), before)
    def test_page_size_one_takes_the_plain_path(self):
-        """token == page: nothing to dedup, so `free_segment` must not invent
+        """token == page: a stride slice would drop tokens, so take the plain
-        a stride slice that would drop tokens."""
+        path."""
        alloc = _paged_allocator(lazy=True)
        alloc.page_size = 1
        v = alloc.alloc(PAGE_SIZE)
@@ -325,12 +273,9 @@ class TestFreeSegment(unittest.TestCase):
 class TestFreeGroupKeepsPositions(unittest.TestCase):
    """Mirrors `test_paged_free_segment.test_group_defers_until_group_end`.
-    Bug regression: buffering RAW tokens and `torch.cat`-ing them at
+    Bug regression: a free group must buffer page REPRESENTATIVES, not raw
-    `free_group_end` destroys each segment's shape, so the merged tensor has no
+    tokens -- concatenating raw tokens loses the page structure and sends
-    recoverable page structure and falls back to `torch.unique`. Measured as 71
+    `free_group_end` back to `torch.unique`.
    of 77 `_free_lazy` calls still syncing on gpt-oss and Qwen3.5 ps=256
    (eval_429), all attributed to `free_group_end` via the decode path. The fix
    buffers page REPRESENTATIVES, so the merge concatenates page ids.
    """
    def test_group_defers_until_group_end(self):
@@ -346,8 +291,6 @@ class TestFreeGroupKeepsPositions(unittest.TestCase):
        self.assertEqual(alloc._free_phys_pages.numel(), before + 2)
    def test_group_end_does_not_sync(self):
        """The property the whole fix exists for: a grouped segment free must
        complete with `torch.unique` disabled."""
        alloc = _paged_allocator(lazy=True)
        row = alloc.alloc(3 * PAGE_SIZE)
        alloc.free_group_begin()
@@ -360,8 +303,8 @@ class TestFreeGroupKeepsPositions(unittest.TestCase):
        self.assertGreater(alloc._free_phys_pages.numel(), 0)
    def test_positionless_group_still_uses_the_unique_path(self):
-        """Plain `free()` inside a group has no position to keep, so it must
+        """Plain `free()` carries no position to keep, so it must still take
-        still go through the (syncing) dedup -- correctness over speed."""
+        the syncing dedup -- correctness over speed."""
        alloc = _paged_allocator(lazy=True)
        row = alloc.alloc(2 * PAGE_SIZE)
        alloc.free_group_begin()
@@ -374,12 +317,9 @@ class TestFreeGroupKeepsPositions(unittest.TestCase):
 class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
-    """Completeness guard. The base `free_segment` DISCARDS `start_pos` and
+    """Completeness guard: the base `free_segment` DISCARDS `start_pos` and
    calls plain `free`, so an allocator that inherits it sends every segment
-    free into the syncing dedup -- silently, with no error and no wrong answer,
+    free into the syncing dedup -- silently, with no error and no wrong answer.
    just a stalled scheduler thread. That is exactly what happened: the SWA
    composite was overridden and the Mamba composite was not, and 77 of 77
    `_free_lazy` calls on Qwen3.5 ps=256 still synced (eval_428).
    """
    def test_all_overridden(self):
@@ -400,8 +340,8 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
                )
    def test_composites_buffer_reps_not_tokens_in_a_group(self):
-        """The group buffer must exist on every allocator that can receive a
+        """Every allocator that can receive a segment free needs the group
-        segment free, or `free_segment` raises inside a group."""
+        buffer, or `free_segment` raises inside a group."""
        for cls in (
            mea.MultiEndedAllocator,
            unified_mamba.UnifiedMambaTokenToKVPoolAllocator,
@@ -432,13 +372,10 @@ class TestUnifiedSwaFullSideGroup(unittest.TestCase):
 class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
-    """The per-decode-step SWA window ratchet frees a CONTIGUOUS row slice
+    """The per-decode-step SWA window ratchet frees a CONTIGUOUS row slice with
-    with host-int, page-aligned bounds — the same shape `free_segment` was
+    host-int, page-aligned bounds, so `free_swa(..., start_pos=)` must reach the
-    built for. `free_swa(..., start_pos=)` must therefore reach the swa side
+    swa side with caller-derived page ids: no `torch.unique` and no stale-slot
-    with caller-derived page ids: no `torch.unique` (data-dependent shape =
+    `.item()` on the per-step path.
    host sync) and no stale-slot `.item()` on the per-step path.
    Poisoning the ops is the decisive form (a textual guard can be fooled).
    """
    PS = 4
@@ -480,8 +417,7 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
        )
    def test_ratchet_shape_free_swa_never_syncs(self):
-        """Aligned bounds (the ratchet guarantees them at ps>1): no unique,
+        """Aligned bounds, which the ratchet guarantees at ps > 1."""
        no item -- on the lazy production config."""
        alloc = self._swa_composite(lazy=True)
        v = alloc.alloc(8 * self.PS)
        self.assertIsNotNone(v)
@@ -497,8 +433,8 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
            alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS)
    def test_full_only_segment_free_never_syncs(self):
-        """The request-finish dead half (swa already tombstoned) frees the full
+        """Request-finish shape: the swa side is already tombstoned, so the
-        side by page reps: no unique from free_full's token dedup."""
+        full side must free by page reps rather than `free_full`'s dedup."""
        alloc = self._swa_composite(lazy=True)
        v = alloc.alloc(8 * self.PS)
        alloc.free_swa(v, start_pos=0)
@@ -523,8 +459,8 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
            alloc.free_swa(v[1 : 5 * self.PS], start_pos=1)
    def test_start_pos_path_matches_the_fallback_end_state(self):
-        """Derived property: the stride-rep path and the dedup fallback must
+        """Derived property: the stride-rep path and the dedup fallback leave
-        leave IDENTICAL allocator state (v2p tombstones, capacity)."""
+        identical v2p tombstones and capacity."""
        for lazy in (True, False):
            with self.subTest(lazy=lazy):
                a1 = self._swa_composite(lazy=lazy)
@@ -559,10 +495,9 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
    torch.cuda.is_available(), "the fused tombstone is a Triton kernel"
 )
 class TestFusedTombstoneWritesBothTables(unittest.TestCase):
-    """The source scan above accepts `free_unbind_inplace` as a no-sync
+    """The source scan accepts `free_unbind_inplace` as a no-sync mechanism; on
-    mechanism; this is what makes that acceptance mean something. On CPU the
+    CPU the launcher takes its pure-torch reference path, so nothing else in
-    launcher takes its pure-torch reference path, so nothing else in the suite
+    the suite ever runs the kernel that does the tombstoning.
    ever runs the kernel that does the tombstoning.
    """
    def test_matches_the_reference_over_randomized_bindings(self):
@@ -602,19 +537,6 @@ class TestFusedTombstoneWritesBothTables(unittest.TestCase):
                f"a live binding was disturbed, trial {trial}",
            )
    def test_cuda_agrees_with_the_cpu_reference(self):
        from sglang.kernels.ops.memory.virtual_slot import free_unbind_inplace
        v = torch.tensor([3, 0, 5], dtype=torch.int64)
        cpu_v2p = torch.tensor([1, 2, 3, 4, 5, 0], dtype=torch.int64)
        cpu_p2v = torch.tensor([5, 0, 1, 2, 3, 4], dtype=torch.int64)
        cu_v2p, cu_p2v = cpu_v2p.cuda(), cpu_p2v.cuda()
        cpu_out = free_unbind_inplace(v, cpu_v2p, cpu_p2v)
        cu_out = free_unbind_inplace(v.cuda(), cu_v2p, cu_p2v)
        self.assertTrue(torch.equal(cpu_out, cu_out.cpu()))
        self.assertTrue(torch.equal(cpu_v2p, cu_v2p.cpu()))
        self.assertTrue(torch.equal(cpu_p2v, cu_p2v.cpu()))
    def test_empty_free_is_a_noop(self):
        from sglang.kernels.ops.memory.virtual_slot import free_unbind_inplace
@@ -11,49 +11,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Round-trip correctness of ``UnifiedKVPool._build_mamba_views`` — the
+"""Round-trip correctness of ``UnifiedKVPool._build_mamba_views`` -- the
 envelope-strided conv/temporal (SSM) state views that back ``UnifiedMambaPool``.
-
+Guards against garbled greedy decode from a stride/offset/alignment bug in the
-This isolates the unified-memory-pool Mamba STATE layout from the full model. It guards
+state views. CPU-only.
 against a class of correctness defect where Falcon-H1 greedy decode is garbled
 under the unified memory pool, isolated to the Mamba conv/temporal state path: a
 stride/offset/alignment bug in the view construction (analogous to the fixed
 `_extract_kv_strides` MHA bug).
 Within one slot's envelope the bytes are
-``[conv[0]·L0 | conv[0]·L1 | ... | conv[1]·L0 | ... | temporal·L0 | ...]`` and
+``[conv[0] L0 | conv[0] L1 | ... | conv[1] L0 | ... | temporal L0 | ...]``,
-across slots the layout is envelope (slot stride == entry_bytes). Each returned
+the slot stride is ``entry_bytes``, and each returned view is
-view is ``(num_layers, max_slots, *inner_shape)``. The conv dtype (bf16, 2 B)
+``(num_layers, max_slots, *inner_shape)``. The conv dtype (bf16, 2 B) and
-and temporal dtype (fp32, 4 B) DIFFER, so the temporal view's byte offset must
+temporal dtype (fp32, 4 B) may DIFFER, so the temporal view's byte offset must
-be a multiple of the temporal itemsize — an alignment hazard that
+be a multiple of the temporal itemsize -- an alignment hazard that
-``_build_mamba_views`` now asserts.
+``_build_mamba_views`` asserts.
 These tests prove the views:
  - round-trip every (tensor, layer, slot) element with the Falcon-like
    bf16-conv / fp32-temporal dtype mix (catches stride/offset/alignment bugs);
  - do NOT alias each other (conv[i] vs conv[j] vs temporal) or across
    layers/slots (catches envelope-overlap);
  - match a contiguous ``(num_layers, max_slots, *inner)`` reference exactly
    (the shape `MambaPool.State.conv[i]` / `.temporal` expose);
  - reject a deliberately mis-aligned spec via the alignment assert.
 The round-trip class is skipped on CPU — those views back GPU kernels and we
 mirror the GPU path. ``TestKDAFlashInferEnvelopeStateContract`` is pure stride
 arithmetic and runs everywhere.
    python -m pytest test/registered/unit/mem_cache/test_shared_mamba_views.py -v
 """
 import unittest
 import torch
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cpu_ci
-_HAS_CUDA = torch.cuda.is_available()
+register_cpu_ci(est_time=30, suite="base-a-test-cpu")
 _DEV = "cuda" if _HAS_CUDA else "cpu"
 register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
 register_amd_ci(est_time=30, stage="stage-b", runner_config="1-gpu-small-amd")
@@ -65,11 +43,10 @@ def _make_pool(
    temporal_state_shape,
    temporal_dtype,
    want_slots=8,
-    device=_DEV,
+    device="cpu",
 ):
-    """Build a minimal 2-sub-pool ``UnifiedKVPool`` (a small MHA grow-up peer
+    """Minimal 2-sub-pool ``UnifiedKVPool`` sized to hold >= ``want_slots``
-    + the Mamba grow-down pool under test) sized to hold >= ``want_slots`` Mamba
+    Mamba slots; returns ``(pool, mamba_spec)``."""
    slots, and return ``(pool, mamba_spec)``."""
    from sglang.srt.mem_cache.unified_memory_pool import (
        MambaSubPoolSpec,
        MHASubPoolSpec,
@@ -97,9 +74,8 @@ def _make_pool(
    entry_mamba = mamba_spec.entry_bytes()
    entry_full = full_spec.entry_bytes()
    entry_max = max(entry_mamba, entry_full)
-    # Need max_slots("mamba") = total // entry_mamba >= want_slots, and total
+    # Headroom so BOTH pools clear their min_slot_index, then round up to a
-    # large enough that BOTH pools clear their min_slot_index. Add generous
+    # multiple of 8 so bf16/fp32 ``.view()`` stays legal.
    # headroom, then round up to a multiple of 8 (covers bf16/fp32 .view()).
    total_bytes = want_slots * entry_mamba + 8 * entry_max
    total_bytes = ((total_bytes + 7) // 8) * 8
    pool = UnifiedKVPool(
@@ -111,11 +87,8 @@ def _make_pool(
    return pool, mamba_spec
@unittest.skipUnless(_HAS_CUDA, "shared Mamba views back GPU kernels")
 class TestUnifiedMambaViews(unittest.TestCase):
-    # Falcon-H1-like dims: even conv_dim, bf16 conv, fp32 temporal, several
+    # Falcon-H1-like dims: even conv_dim, bf16 conv, fp32 temporal.
    # layers. (Mamba2 conv state is (conv_dim, kernel-1); temporal/SSM state is
    # (nheads, head_dim, ssm_state_size).)
    FALCON_KW = dict(
        mamba_layer_num=5,  # odd, to stress the temporal-offset alignment
        conv_state_shapes=[(48, 3)],  # conv_dim=48, kernel-1=3
@@ -125,11 +98,9 @@ class TestUnifiedMambaViews(unittest.TestCase):
    )
    def _fill_and_roundtrip(self, pool, mamba_spec):
-        """Write a distinct random tensor to each conv view + the temporal view
+        """Write a distinct random tensor to every view, then read all of them
-        (in their own dtypes), then read all back and assert exact equality.
+        back: writing all first makes any envelope overlap corrupt an earlier
-        Writing ALL views first and reading ALL after means any envelope overlap
+        write."""
        (conv[i]/conv[j]/temporal aliasing) corrupts an earlier write → mismatch.
        """
        conv_views, temporal_view = pool.mamba_views_for("mamba")
        torch.manual_seed(0)
        refs = []
@@ -155,34 +126,48 @@ class TestUnifiedMambaViews(unittest.TestCase):
            f"stride={temporal_view.stride()}",
        )
-    def test_roundtrip_falcon_like(self):
+    OTHER_GEOMETRIES = {
        pool, spec = _make_pool(**self.FALCON_KW)
        self._fill_and_roundtrip(pool, spec)
    def test_roundtrip_single_layer_single_slot_edges(self):
        # 1 layer, multiple conv tensors, same-dtype conv/temporal.
-        pool, spec = _make_pool(
+        "single_layer_single_slot_edges": dict(
            mamba_layer_num=1,
            conv_state_shapes=[(16, 3), (8, 3)],
            conv_dtype=torch.float32,
            temporal_state_shape=(4, 8, 16),
            temporal_dtype=torch.float32,
            want_slots=4,
-        )
+        ),
-        self._fill_and_roundtrip(pool, spec)
+        # Two conv tensors + bf16/fp32 mix: exercises the per-conv-tensor offset
    def test_roundtrip_multi_conv_tensors(self):
        # Two conv tensors + bf16/fp32 mix — exercises the per-conv-tensor offset
        # accumulation in _build_mamba_views.
-        pool, spec = _make_pool(
+        "multi_conv_tensors": dict(
            mamba_layer_num=3,
            conv_state_shapes=[(32, 3), (16, 3)],
            conv_dtype=torch.bfloat16,
            temporal_state_shape=(8, 8, 16),
            temporal_dtype=torch.float32,
            want_slots=6,
-        )
+        ),
-        self._fill_and_roundtrip(pool, spec)
+        # ALIGNED spec: conv region = 2 layers * 2*3*2 B = 24 B, a multiple of
        # the fp32 temporal itemsize, so it must build and round-trip.
        "alignment_ok": dict(
            mamba_layer_num=2,
            conv_state_shapes=[(2, 3)],
            conv_dtype=torch.bfloat16,
            temporal_state_shape=(2, 4, 4),
            temporal_dtype=torch.float32,
            want_slots=4,
        ),
    }
    def test_roundtrip_geometries(self):
        """Round-trip every (tensor, layer, slot) element across the geometry
        shapes the view builder has to handle."""
        for name, kwargs in (
            ("falcon_like", self.FALCON_KW),
            *self.OTHER_GEOMETRIES.items(),
        ):
            with self.subTest(geometry=name):
                pool, spec = _make_pool(**kwargs)
                self._fill_and_roundtrip(pool, spec)
    def test_no_cross_region_overlap(self):
        """Zero buffer; write a sentinel to ONE view; every OTHER view must read
@@ -207,62 +192,10 @@ class TestUnifiedMambaViews(unittest.TestCase):
                    f"(envelope regions overlap)",
                )
    def test_per_layer_per_slot_addressing(self):
        """Distinct value per (layer, slot) on the temporal view; verify exact
        addressing (no layer/slot aliasing). Uses small integers exactly
        representable in the view dtype.
        NB: ``temporal_view`` is a non-contiguous strided view, so we must NOT
        ``.reshape()`` it (that would COPY, breaking the alias) — we
        broadcast-assign into the view in place and read back via basic
        indexing (which keeps the view)."""
        pool, spec = _make_pool(**self.FALCON_KW)
        _, temporal_view = pool.mamba_views_for("mamba")
        N, S = temporal_view.shape[0], temporal_view.shape[1]
        inner_ndim = temporal_view.dim() - 2
        # value = layer*S + slot  (< N*S, small → exact in fp32)
        base = (
            torch.arange(N, device=temporal_view.device)[:, None] * S
            + torch.arange(S, device=temporal_view.device)[None, :]
        ).to(temporal_view.dtype)
        # Broadcast (N, S) over the inner dims, in place into the strided view.
        temporal_view[:] = base.view(N, S, *([1] * inner_ndim))
        # Read back the first inner element of every (layer, slot) via basic
        # indexing (stays a view).
        readback = temporal_view[(slice(None), slice(None)) + (0,) * inner_ndim]
        self.assertTrue(
            torch.equal(readback, base),
            "temporal (layer, slot) addressing wrong — layer/slot stride bug",
        )
    def test_matches_contiguous_reference(self):
        """The shared view must be a faithful relabeling of a contiguous
        ``(num_layers, max_slots, *inner)`` tensor: identical data written by the
        same logical index reads back identically."""
        pool, spec = _make_pool(**self.FALCON_KW)
        conv_views, temporal_view = pool.mamba_views_for("mamba")
        for v in conv_views + [temporal_view]:
            ref = torch.randn(v.shape, device=v.device).to(v.dtype)
            contig = ref.clone().contiguous()
            v.copy_(ref)
            self.assertEqual(tuple(v.shape), tuple(contig.shape))
            self.assertTrue(
                torch.equal(v.contiguous(), contig),
                "shared view not equivalent to its contiguous counterpart",
            )
    def test_alignment_guard_fires_on_misaligned_spec(self):
-        """A spec whose conv region (bf16) is an odd multiple of 2 B makes the
+        """A per-slot entry that is not a multiple of the temporal itemsize would
-        per-slot entry (= conv_region + N*temporal_row = 2 B + 4 B = 6 B) NOT a
+        mis-offset the view and must trip the alignment assert; either guard can
-        multiple of the temporal itemsize (fp32, 4 B). The temporal/SSM-state
+        fire first, so match on the shared "misaligned" wording."""
        view's storage_offset is computed by integer-dividing a byte offset by
        the temporal itemsize, so this would silently mis-offset the view.
        ``_build_mamba_views`` must reject it with a loud alignment assert.
        NOTE: the ``entry_bytes % itemsize`` guard is what fires here, and it
        subsumes the conv-region offset check (see the comment in
        ``_build_mamba_views``). We assert on the shared "misaligned" wording
        rather than on which specific guard trips."""
        with self.assertRaises(AssertionError) as cm:
            _make_pool(
                mamba_layer_num=1,  # entry = 2 B conv + 4 B temporal = 6 B, not %4
@@ -274,30 +207,12 @@ class TestUnifiedMambaViews(unittest.TestCase):
            )
        self.assertIn("misalign", str(cm.exception).lower())
    def test_alignment_ok_for_aligned_spec(self):
        """An aligned spec (conv region a multiple of the temporal itemsize)
        must build and round-trip cleanly."""
        # conv region = N * conv_dim*(k-1) * 2 ; with conv_dim=2 -> per-layer 2*3*2=12,
        # times N=2 = 24, divisible by 4. Aligned.
        pool, spec = _make_pool(
            mamba_layer_num=2,
            conv_state_shapes=[(2, 3)],
            conv_dtype=torch.bfloat16,
            temporal_state_shape=(2, 4, 4),
            temporal_dtype=torch.float32,
            want_slots=4,
        )
        self._fill_and_roundtrip(pool, spec)
 def _k3_kda_mamba_geometry(heads_per_rank: int) -> dict:
-    """Kimi K3 KDA per-rank state geometry: 69 KDA layers, K = V = 128,
+    """Kimi K3 KDA per-rank state geometry: 69 KDA layers, K = V = 128, conv
-    conv width 4 (=> 3 cached tokens), conv row ``(kernel-1, q+k+v dim)``
+    width 4 (=> 3 cached tokens), conv row ``(kernel-1, q+k+v dim)`` per
-    in the KimiLinear layout (``KimiLinearStateShape.create`` with
+    ``KimiLinearStateShape.create``, temporal/SSM state ``(HV, V, K)``, and
-    num_k_heads == num_heads, head_k_dim == head_dim — see
+    ``heads_per_rank`` = 96 total KDA heads / attn_tp."""
    ``models/kimi_linear.py``), temporal/SSM state ``(HV, V, K)``.
    ``heads_per_rank`` = 96 total KDA heads / attn_tp (12 at the TP8
    deployment shape, cf. ``kernels/ops/attention/kda_fused_decode.py``)."""
    h = heads_per_rank
    return dict(
        layer_num=69,
@@ -311,46 +226,35 @@ def _k3_kda_mamba_geometry(heads_per_rank: int) -> dict:
 class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase):
-    """Derived property: the envelope-strided KDA temporal view (unified memory
+    """Derived property: the envelope-strided KDA temporal view must satisfy the
-    / page-major layout) must satisfy the state contract of FlashInfer
+    state contract of FlashInfer ``recurrent_kda``, because the KDA flashinfer
-    ``recurrent_kda``, because the KDA flashinfer decode wrapper
+    decode wrapper (``linear/kernels/kda_flashinfer.py``) passes the per-layer
-    (``linear/kernels/kda_flashinfer.py``) passes the committed per-layer pool
+    pool view straight into the kernel, with no gather/scatter copy.
    view straight into the kernel (in-place state update on the cu_seqlens path
    — no gather/scatter copy around the call).
    The kernel compiles its state argument as a CuTe fake tensor of shape
    ``[N, HV, V, K]`` with stride ``(sym_int64(divisibility=16), V*K, K, 1)``
-    and ``assumed_align=32`` (flashinfer ``kda_kernels/recurrent_kda.py``), so
+    and ``assumed_align=32``, so a per-layer pool view is readable only when:
    a per-layer pool view is only readable by the kernel when:
      * its inner strides are exactly compact ``(V*K, K, 1)``;
-      * its slot stride — the per-slot envelope pitch, NOT ``HV*V*K`` — is a
+      * its slot stride -- the per-slot envelope pitch, NOT ``HV*V*K`` -- is a
        multiple of 16 elements (32 bytes at bf16);
      * its base byte offset is 32-byte aligned (for every layer).
    Any envelope-layout change that breaks one of these (per-slot padding that
    is not a 32 B multiple, a conv-shape change misaligning the temporal
    region, a transposed/padded temporal inner layout) would silently
    mis-address every KDA state read/write on SM100 flashinfer decode; this
    test turns such a diff red without a GPU.
    """
-    # 32 B: recurrent_kda's assumed_align AND its slot-stride divisibility
+    # recurrent_kda's assumed_align AND its slot-stride divisibility (16
-    # (16 elements * 2 B bf16). External-source literal from flashinfer
+    # elements * 2 B bf16), from flashinfer kda_kernels/recurrent_kda.py.
    # kda_kernels/recurrent_kda.py (S_batch = cute.sym_int64(divisibility=16),
    # make_fake_tensor(..., assumed_align=32)).
    _KERNEL_ALIGN_BYTES = 32
    @staticmethod
-    def _build_tp8_views():
+    def _build_kda_views(heads_per_rank: int = 12):
-        """Real TP8 K3 KDA envelope views on CPU (2 slots suffice — the
+        """Real K3 KDA envelope views on CPU for one attn-TP shard; 2 slots
-        per-slot geometry is slot-count independent)."""
+        suffice, the per-slot geometry is slot-count independent."""
        from sglang.srt.mem_cache.layout.page_major import (
            build_page_major_mamba_views,
            mamba_entry_bytes,
        )
-        geom = _k3_kda_mamba_geometry(12)  # 96 heads / TP8
+        geom = _k3_kda_mamba_geometry(heads_per_rank)
        entry_bytes = mamba_entry_bytes(**geom)
        max_slots = 2
        raw = torch.empty(max_slots * entry_bytes, dtype=torch.uint8, device="cpu")
@@ -360,76 +264,44 @@ class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase):
        return geom, entry_bytes, temporal_view
    def test_k3_tp8_envelope_view_matches_recurrent_kda_contract(self):
-        """Check every per-layer temporal view against the kernel contract."""
+        """Check every per-layer temporal view against the kernel contract, for
-        geom, entry_bytes, temporal_view = self._build_tp8_views()
+        every plausible attn-TP shard of K3's 96 KDA heads."""
        itemsize = temporal_view.element_size()
        _, v, k = geom["temporal_state_shape"]
        for layer in (0, geom["layer_num"] - 1):
            view = temporal_view[layer]  # [slots, HV, V, K], what decode() gets
            self.assertEqual(
                view.stride()[1:],
                (v * k, k, 1),
                "temporal inner strides must stay compact (V*K, K, 1): "
                "recurrent_kda compiles them as constants",
            )
            self.assertEqual(
                view.stride(0),
                entry_bytes // itemsize,
                "slot stride must be the envelope pitch (entry_bytes)",
            )
            self.assertEqual(
                view.stride(0) % (self._KERNEL_ALIGN_BYTES // itemsize),
                0,
                "slot stride must satisfy recurrent_kda's "
                "sym_int64(divisibility=16) — 16 elements = 32 B at bf16",
            )
            self.assertEqual(
                (view.storage_offset() * itemsize) % self._KERNEL_ALIGN_BYTES,
                0,
                f"layer {layer} temporal view base is not 32 B aligned "
                "(recurrent_kda assumed_align=32)",
            )
    def test_k3_entry_and_temporal_offset_32B_multiples_across_tp(self):
        """The two byte quantities that feed the contract above — the per-slot
        envelope pitch and the temporal region's offset inside the envelope
        (= all-layers conv region, temporal comes last) — must be 32 B
        multiples for every plausible attn-TP shard of K3's 96 KDA heads."""
        import math
        from sglang.srt.mem_cache.layout.page_major import mamba_entry_bytes
        for heads_per_rank in (96, 48, 24, 12):  # attn_tp 1 / 2 / 4 / 8
-            geom = _k3_kda_mamba_geometry(heads_per_rank)
+            with self.subTest(heads_per_rank=heads_per_rank):
-            entry_bytes = mamba_entry_bytes(**geom)
+                geom, entry_bytes, temporal_view = self._build_kda_views(heads_per_rank)
-            conv_region_bytes = (
+                itemsize = temporal_view.element_size()
-                geom["layer_num"]
+                _, v, k = geom["temporal_state_shape"]
-                * math.prod(geom["conv_state_shapes"][0])
+                for layer in (0, geom["layer_num"] - 1):
-                * geom["conv_dtype"].itemsize
+                    # [slots, HV, V, K], what decode() gets
-            )
+                    view = temporal_view[layer]
-            self.assertEqual(
+                    self.assertEqual(
-                entry_bytes % self._KERNEL_ALIGN_BYTES,
+                        view.stride()[1:],
-                0,
+                        (v * k, k, 1),
-                f"tp shard h={heads_per_rank}: envelope pitch {entry_bytes} B "
+                        "temporal inner strides must stay compact (V*K, K, 1): "
-                "breaks recurrent_kda's slot-stride divisibility",
+                        "recurrent_kda compiles them as constants",
-            )
+                    )
-            self.assertEqual(
+                    self.assertEqual(
-                conv_region_bytes % self._KERNEL_ALIGN_BYTES,
+                        view.stride(0),
-                0,
+                        entry_bytes // itemsize,
-                f"tp shard h={heads_per_rank}: temporal region offset "
+                        "slot stride must be the envelope pitch (entry_bytes)",
-                f"{conv_region_bytes} B breaks assumed_align=32",
+                    )
-            )
+                    self.assertEqual(
                        view.stride(0) % (self._KERNEL_ALIGN_BYTES // itemsize),
                        0,
                        "slot stride must satisfy recurrent_kda's "
                        "sym_int64(divisibility=16) — 16 elements = 32 B at bf16",
                    )
                    self.assertEqual(
                        (view.storage_offset() * itemsize) % self._KERNEL_ALIGN_BYTES,
                        0,
                        f"layer {layer} temporal view base is not 32 B aligned "
                        "(recurrent_kda assumed_align=32)",
                    )
    def test_wrapper_state_contract_check_matches_layout(self):
-        """The KDA flashinfer decode wrapper enforces this same contract at
+        """``FlashInferKDAKernel._check_state_stride_contract`` enforces the same
-        runtime (``FlashInferKDAKernel._check_state_stride_contract``, called
+        contract at runtime, and a regression there would only surface on SM100
-        once per pool view before handing the pool to ``recurrent_kda``). A
+        hardware, so pin its accept/reject behavior on CPU."""
        regression in that check would only surface on SM100 hardware, so pin
        its accept/reject behavior here: it must ACCEPT exactly what the
        layouts produce — the envelope-strided per-layer view and a plain
        contiguous pool — and REJECT views the kernel would silently
        mis-address (wrong inner strides; a slot stride off the divisibility)."""
        import types
        from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import (
@@ -442,7 +314,7 @@ class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase):
            # Fresh stub per call: the real kernel caches approvals by id().
            check(types.SimpleNamespace(_state_contract_ok=set()), view)
-        _, _, temporal_view = self._build_tp8_views()
+        _, _, temporal_view = self._build_kda_views()
        envelope = temporal_view[0]  # what forward_decode hands to the kernel
        run(envelope)  # must not raise
@@ -11,26 +11,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""MHA K/V views for the unified memory pool (uniform-row hybrid models).
+"""MHA K/V views for the unified memory pool (uniform-row hybrid models), CPU-only.
-Covers, CPU-only (pure torch — no GPU / Triton kernels):
+Addressing law under test:
  - `build_mha_views` refuses an asymmetric-KV spec: its addressing
    assumes one uniform row width, so it is the boundary that checks;
  - `build_mha_views` addressing: view_l[kernel_id(t)] must land exactly at
    the page-major envelope byte offset the STRIDED builder assigns to the same
    (page, slot, layer, K|V) cell — the two builders are views over one truth;
  - K and V of one token share ONE kernel-facing id (per-layer origin shift does the
    disambiguation), with no aliasing across the 2*L overlapping views;
  - the missing-tail-pad and asymmetric-dims cases fail loud at construction.
 Addressing law under test (the derived property everything else builds on):
    kernel_id(t) = (t // ps) * (ps * 2L) + t % ps
-    K of layer l at block 2l, V at block 2l+1, blocks are ps rows of
+    K of layer l at block 2l, V at block 2l+1; blocks are ps rows of
-    head_num*head_dim elements — offsets identical to
+    head_num*head_dim elements, at offsets identical to
    MHASubPoolSpec.layer_k/v_offset_in_page when rows are uniform.
    python -m pytest test/registered/unit/mem_cache/test_unified_mha_views.py -v
 """
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -43,10 +31,7 @@ from types import SimpleNamespace
 import torch
 from sglang.srt.environ import envs
-from sglang.srt.mem_cache.layout.page_major import (
+from sglang.srt.mem_cache.layout.page_major import build_mha_views
    build_mha_views,
    mha_entry_bytes,
 )
 from sglang.srt.mem_cache.unified_memory_pool import (
    MHASubPoolSpec,
    UnifiedKVPool,
@@ -55,12 +40,10 @@ from sglang.srt.mem_cache.unified_memory_pool import (
 _DEV = "cpu"
 # `set_kv_buffer` dispatches on the PLATFORM (memory_pool._is_cuda, resolved at
-# import), not on the tensors it is handed, so cases driving it must build on
+# import), not on the tensors it is handed, so cases driving it build there.
 # the platform's device. The rest of this file is byte arithmetic, so CPU.
 _STORE_DEV = "cuda" if torch.cuda.is_available() else "cpu"
-# Small-but-nontrivial MHA geometry: L=2 layers, H=2 heads, D=4, so every byte
+# Geometry kept tiny so every byte offset is hand-checkable.
 # offset is hand-checkable. blocks = 2L = 4 per page.
 _L = 2
 _H = 2
 _D = 4
@@ -108,13 +91,8 @@ def _build_views(raw, ps, num_pages, head_dim=_D, v_head_dim=_D, layer_num=_L):
 def _reference_strided_views(raw, *, page_size, num_pages, anchor_bytes=0):
-    """Independent 4-D strided description of the page-major envelope.
+    """Independent 4-D strided description of the same page-major envelope,
-
+    addressed by ``(page, slot)`` -- the oracle for the view builder."""
    This is the retired production strided builder, kept here as the oracle:
    per-layer ``(num_pages, page_size, head_num, head_dim)`` views addressed by
    ``(page, slot)``, so the builder's addressing can be cross-checked
    against a second, independently-derived description of the same bytes.
    """
    k_row_bytes = _ROW * _ITEM
    v_row_bytes = _ROW * _ITEM
    page_bytes = page_size * _L * (k_row_bytes + v_row_bytes)
@@ -147,11 +125,8 @@ def _reference_strided_views(raw, *, page_size, num_pages, anchor_bytes=0):
 class TestMHASpecSurface(unittest.TestCase):
    def test_asymmetric_rows_refused_by_the_view_builder(self):
-        """The row-block array exists only for uniform rows, so the builder
+        """ServerArgs screens asymmetric-KV models out of
-        whose addressing depends on it is the one that refuses (the MiMoV2
+        --enable-unified-memory; this guards a caller reaching the builder."""
        shape, scaled down). ServerArgs screens such models out of
        --enable-unified-memory long before we get here; this is the check for
        a caller that reaches the builder directly."""
        spec = _mha_spec()
        raw = torch.zeros(1 << 16, dtype=torch.uint8)
        with self.assertRaises(AssertionError):
@@ -166,28 +141,6 @@ class TestMHASpecSurface(unittest.TestCase):
                num_pages=4,
            )
    def test_spec_offsets_equal_block_origins(self):
        """The spec's byte math and the view builder's origins are two
        independent derivations of the envelope; under uniform rows they must
        agree: layer_k_offset(l) == (2l)*ps*row, layer_v_offset(l) == (2l+1)*ps*row."""
        spec = _mha_spec()
        for ps in (1, 4):
            row = spec.k_row_bytes()
            for l in range(_L):
                self.assertEqual(spec.layer_k_offset_in_page(l, ps), (2 * l) * ps * row)
                self.assertEqual(
                    spec.layer_v_offset_in_page(l, ps), (2 * l + 1) * ps * row
                )
    def test_entry_bytes_matches_layout_helper(self):
        spec = _mha_spec()
        self.assertEqual(
            spec.entry_bytes(),
            mha_entry_bytes(
                layer_num=_L, head_num=_H, head_dim=_D, v_head_dim=_D, itemsize=_ITEM
            ),
        )
 class TestMHAViews(unittest.TestCase):
    def test_view_shapes_are_stock_mha(self):
@@ -202,11 +155,8 @@ class TestMHAViews(unittest.TestCase):
            self.assertEqual(v.stride(), (_ROW, _D, 1))
    def test_addressing_matches_strided_reference(self):
-        """Cross-readback: bytes written through the reference STRIDED views at
+        """Cross-readback both ways: a write through the strided oracle at
-        (page, slot) must be read back through the views at kernel_id(t),
+        (page, slot) must read back through the view at kernel_id(t)."""
        for both K and V of every layer — and vice versa. This pins that the
        view builder and the independent strided description agree on the
        same physical envelope."""
        for ps in (1, 4):
            num_pages = 5
            raw = _make_raw(ps, num_pages)
@@ -240,37 +190,6 @@ class TestMHAViews(unittest.TestCase):
                    torch.all(sv[l][p, s] == float(p * 100 + l * 10 + s + 4))
                )
    def test_byte_addresses_match_envelope_formula(self):
        """The per-layer view's byte address for token ``t``, layer ``L`` must equal
        the hand-computed envelope formula: page origin + layer-block origin +
        slot offset. Independent of any view builder — this is the raw layout
        contract every envelope consumer (moves, sizing, transfer math) relies
        on."""
        k_row = _ROW * _ITEM
        v_row = _ROW * _ITEM
        for ps in (1, 4):
            num_pages = 5
            page_bytes = ps * _L * (k_row + v_row)
            dk, dv = _build_views(_make_raw(ps, num_pages), ps, num_pages)
            for t in (0, 1, ps, 3 * ps + (ps - 1), 4 * ps):
                d = _kernel_id(t, ps)
                for L in range(_L):
                    expected_k = (
                        (t // ps) * page_bytes
                        + L * ps * (k_row + v_row)
                        + (t % ps) * k_row
                    )
                    expected_v = (
                        (t // ps) * page_bytes
                        + L * ps * (k_row + v_row)
                        + ps * k_row
                        + (t % ps) * v_row
                    )
                    got_k = (dk[L].storage_offset() + d * dk[L].stride(0)) * _ITEM
                    got_v = (dv[L].storage_offset() + d * dv[L].stride(0)) * _ITEM
                    self.assertEqual(got_k, expected_k, f"K t={t} L={L} ps={ps}")
                    self.assertEqual(got_v, expected_v, f"V t={t} L={L} ps={ps}")
    def test_k_and_v_share_one_kernel_id_without_aliasing(self):
        """One kernel-facing id, 2L distinct cells (K and V of every layer): writes
        through all 2L views at the SAME id must not clobber each other."""
@@ -291,12 +210,6 @@ class TestMHAViews(unittest.TestCase):
        with self.assertRaises(AssertionError):
            _build_views(raw, ps, num_pages)
    def test_asymmetric_dims_rejected(self):
        ps, num_pages = 2, 4
        raw = _make_raw(ps, num_pages)
        with self.assertRaises(AssertionError):
            _build_views(raw, ps, num_pages, head_dim=6, v_head_dim=4)
 # ---- pool level ----
@@ -341,26 +254,6 @@ class TestUnifiedKVPoolViews(unittest.TestCase):
                self.assertEqual(v[0].dim(), 3, f"{name} V at ps={ps}")
                self.assertTrue(k[0].is_contiguous())
    def test_tail_pad_is_derived_from_the_specs(self):
        """The per-layer views hang past the last page envelope, so the pool
        over-allocates one envelope of the widest sub-pool. Derived here, not
        passed in, so no construction site can under-allocate it."""
        for ps in (1, 4):
            kv = _make_pool(ps)
            full, swa = _mha_spec(), _swa_spec()
            self.assertEqual(
                kv.view_tail_pad_bytes,
                ps * max(full.entry_bytes(), swa.entry_bytes()),
                f"tail pad at ps={ps}",
            )
            self.assertEqual(
                kv._raw.numel(),
                full.entry_bytes() * _N_FULL
                + swa.entry_bytes() * _N_SWA
                + kv.view_tail_pad_bytes,
                "the pad extends the allocation only",
            )
 def _layer(l):
    return SimpleNamespace(layer_id=l)
@@ -386,12 +279,8 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
            self.assertEqual(pool_under_test.size, n_rows - ps)
    def test_stock_write_lands_on_envelope_truth(self):
-        """Byte-identity: the pool's stock inherited `set_kv_buffer` at kernel-facing
+        """Byte-identity: the inherited `set_kv_buffer` at kernel-facing locs must
-        locs must produce exactly the bytes that direct writes through STRIDED
+        produce the same bytes as writes through the strided oracle."""
        views over the same envelope produce at the same (page, slot, layer)
        cells. The strided views are built here purely as the independent
        description of the envelope — pins the whole write path (loc -> view ->
        raw bytes) end to end."""
        for ps in (1, 4):
            kv, pool = _make_pool_and_kv(ps, device=_STORE_DEV)
            # An independent strided view of the SAME sub-pool region.
@@ -431,10 +320,8 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
                    )
    def test_move_kv_cache_relocates_whole_envelopes(self):
-        """Compaction hands PHYSICAL token runs, not kernel-facing ids. The override
+        """Compaction hands PHYSICAL token runs, not kernel-facing ids; the
-        must relocate exactly the page envelopes those runs name — red if it is
+        override must relocate exactly the page envelopes those runs name."""
        lost, since the inherited per-layer move would apply physical ids to
        the row space."""
        ps = 4
        kv, pool = _make_pool_and_kv(ps)
        live = kv._raw.numel() - kv.view_tail_pad_bytes
@@ -459,8 +346,7 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
    def test_transfer_entry_points_fail_loud(self):
        """PD / CPU-copy entry points assume per-layer buffers indexed by TOKEN
-        id; against the row space they would silently mis-index (or hit a
+        id and would silently mis-index the row space, so each must raise."""
        missing-attr AttributeError). Every one of them must raise."""
        _, pool = _make_pool_and_kv(1)
        with self.assertRaises(NotImplementedError):
            pool.get_contiguous_buf_infos()
@@ -472,10 +358,8 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
            pool.set_kv_buffer_prefix_valid()
    def test_hnd_env_cannot_hijack_layout(self):
-        """SGLANG_USE_HND_KVCACHE=1 used to flip the inherited env-driven
+        """SGLANG_USE_HND_KVCACHE must not flip this pool's layout: HND indexes
-        layout selector, putting the pool in a mode whose code paths do not
+        4-D while the per-layer views are 3-D, so the pinned label has to win."""
        match its buffers (HND indexes 4-D; the per-layer views are 3-D). The
        pinned label must win."""
        with envs.SGLANG_USE_HND_KVCACHE.override(True):
            _, pool = _make_pool_and_kv(1)
            self.assertFalse(pool.use_hnd)
@@ -483,17 +367,15 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
 class TestFactoryViews(unittest.TestCase):
-    """The real SWA factory builds the sub-pools and wires the matching
+    """Over the real SWA factory: matching kernel-facing multipliers in the
-    kernel-facing multipliers into the composite allocator. End-to-end over
+    composite allocator, and a rebind that emits both write locs."""
    that factory, the rebind must emit BOTH kernel-facing write locs."""
    # _swa_factory geometry: L_full = L_swa = 2, uniform 8/8 dims, ps = 1.
    FULL_MULT = 4  # 2 * L_full
    SWA_MULT = 4  # 2 * L_swa
    def _bundle(self):
-        # Self-contained tiny SWA-factory bundle (L_full = L_swa = 2, uniform
+        # Kept tiny so the per-layer views build on CPU.
        # 8/8 dims, ps = 1) — small enough that per-layer views build on CPU.
        from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
        return init_unified_swa_pools(
@@ -528,11 +410,8 @@ class TestFactoryViews(unittest.TestCase):
        self.assertGreater(pool.view_tail_pad_bytes, 0)
    def test_rebind_emits_kernel_facing_full_and_build_derives_swa(self):
-        """End-to-end over the real factory: rebind_write_loc rebinds
+        """rebind_write_loc rebinds out_cache_loc to FULL-kernel-facing ids, and
-        out_cache_loc to FULL-kernel-facing ids (phase 1), and the per-batch build
+        the SWA write loc is derived pointwise from those kernel-facing values."""
        derives the SWA write loc pointwise from those kernel-facing values
        (phase 2) — both checked against the formulas over the VIRTUAL
        ids."""
        from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
        b = self._bundle()
@@ -14,19 +14,11 @@
 """GPU parity of the per-layer-view `UnifiedMLATokenToKVPool` against the stock
 `MLATokenToKVPool` on real K3 MLA geometry (L=24, D=512+64).
-The unified pool receives kernel-facing locs (kernel_id(t) = (t//ps)*(ps*L) + t%ps); the
+The unified pool receives kernel-facing locs (kernel_id(t) = (t//ps)*(ps*L) +
-reference pool receives the raw token ids. Every (layer, token) cell must hold
+t%ps) where the reference pool receives raw token ids; every (layer, token)
-identical bytes afterwards. Covers:
+cell must hold identical bytes afterwards. The TMA JIT fast path (n_loc >= 768)
-
+flattens the buffer via `.view(shape[0], -1)`, which is legal only because the
-  - `set_mla_kv_buffer` under BOTH kernel paths — the Triton fallback
+per-layer views are contiguous.
    (n_loc < 768) and the TMA JIT fast path (n_loc >= 768, which flattens the
    buffer via `.view(shape[0], -1)`, only legal because per-layer views are
    contiguous);
  - `set_kv_buffer` (combined pre-concatenated write, the Triton-backend path);
  - `get_mla_kv_buffer` roundtrip;
  - page_size 1 and 64.
    python -m pytest test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py -v
 """
 import types
@@ -134,17 +126,12 @@ class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
        torch.cuda.synchronize()
        self._assert_parity(unified, ref, locs, ps)
-    def test_set_mla_kv_buffer_triton_fallback_ps1(self):
+    def test_set_mla_kv_buffer_matches_stock_pool(self):
-        self._run_set_mla(ps=1, n_loc=256)  # < 768 -> Triton fallback kernel
+        """Both kernel paths at both page sizes: n_loc < 768 takes the Triton
-
+        fallback, n_loc >= 768 the TMA JIT fast path."""
-    def test_set_mla_kv_buffer_tma_jit_ps1(self):
+        for ps, n_loc in ((1, 256), (1, 1024), (64, 256), (64, 1024)):
-        self._run_set_mla(ps=1, n_loc=1024)  # >= 768 -> TMA JIT fast path
+            with self.subTest(page_size=ps, n_loc=n_loc):
-
+                self._run_set_mla(ps=ps, n_loc=n_loc)
    def test_set_mla_kv_buffer_triton_fallback_ps64(self):
        self._run_set_mla(ps=64, n_loc=256)
    def test_set_mla_kv_buffer_tma_jit_ps64(self):
        self._run_set_mla(ps=64, n_loc=1024)
    def test_set_kv_buffer_combined_write(self):
        for ps in (1, 64):
@@ -175,29 +162,6 @@ class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
            torch.testing.assert_close(got_nope, nope, rtol=0, atol=0)
            torch.testing.assert_close(got_rope, rope, rtol=0, atol=0)
    def test_move_kv_cache_page_envelope_gpu(self):
        for ps in (1, 64):
            unified, ref, max_tokens = _make_pools(ps)
            num_pages = max_tokens // ps
            n_loc = ps  # one full page of tokens
            src_page, dst_page = num_pages - 2, 2
            src_t = torch.arange(ps, device=_DEV, dtype=torch.int64) + src_page * ps
            dst_t = torch.arange(ps, device=_DEV, dtype=torch.int64) + dst_page * ps
            torch.manual_seed(17)
            for l in range(_L):
                layer = types.SimpleNamespace(layer_id=l)
                k = torch.randn(n_loc, 1, _D, dtype=_DTYPE, device=_DEV)
                unified.set_kv_buffer(layer, _kernel_id(src_t, ps), k, None)
            before = [
                unified.get_key_buffer(l)[_kernel_id(src_t, ps)].clone()
                for l in range(_L)
            ]
            unified.move_kv_cache(dst_t, src_t)
            torch.cuda.synchronize()
            for l in range(_L):
                got = unified.get_key_buffer(l)[_kernel_id(dst_t, ps)]
                torch.testing.assert_close(got, before[l], rtol=0, atol=0)
 if __name__ == "__main__":
    unittest.main()
@@ -11,26 +11,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3).
+"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3), CPU-only.
-Covers, CPU-only (pure torch — no GPU / Triton kernels):
+Addressing law under test: the (page, layer, slot) cell sits at envelope byte
-  - `MLASubPoolSpec` byte math;
+offset `p*(L*ps*D) + l*(ps*D) + s*D`, reached through the kernel-facing id
-  - `build_mla_views` addressing: view_l[kernel_id(t)] must land exactly at
+`(t // ps) * (ps * L) + t % ps`.
    the page-major envelope byte offset `p*(L*ps*D) + l*(ps*D) + s*D`, the
    overlapping per-layer views must not alias at equal kernel-facing ids, and the
    missing-tail-pad case must fail loud;
  - `UnifiedKVPool` MLA plumbing: `view_tail_pad_bytes` extends the allocation
    only, and the reserved sink floor covers the whole page-0 envelope;
  - `UnifiedMLATokenToKVPool`: buffer wiring, V-as-prefix-slice, and the
    page-envelope `move_kv_cache` (REAL physical token ids, page-major runs);
  - `MultiEndedAllocator.translate_kv_loc_for_kernel`: kernel id = v2p-page * (ps*L) +
    offset, tombstone clamp to the sink, `out=` contract, multiplier-1
    fallback, and correctness across eager compaction.
-GPU parity of the actual read/write kernels (set_mla_kv_buffer TMA path etc.)
+GPU parity of the read/write kernels (set_mla_kv_buffer TMA path etc.) lives in
-lives in the server-level tests, not here.
+`test_unified_mla_gpu_parity.py`.
    python -m pytest test/registered/unit/mem_cache/test_unified_mla_views.py -v
 """
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -42,10 +30,7 @@ import unittest
 import torch
 from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator
-from sglang.srt.mem_cache.layout.page_major import (
+from sglang.srt.mem_cache.layout.page_major import build_mla_views
    build_mla_views,
    mla_entry_bytes,
 )
 from sglang.srt.mem_cache.unified_memory_pool import (
    MambaSubPoolSpec,
    MLASubPoolSpec,
@@ -55,8 +40,8 @@ from sglang.srt.mem_cache.unified_memory_pool import (
 _DEV = "cpu"
-# Small-but-nontrivial MLA geometry: L=3 layers, D=8 (=6+2), so every byte
+# Geometry kept tiny so every byte offset is hand-checkable; real K3 is
-# offset is hand-checkable. Real K3 is L=24, D=576 (=512+64).
+# L=24, D=576 (=512+64).
 _L = 3
 _LORA = 6
 _ROPE = 2
@@ -107,16 +92,6 @@ def _kernel_id(t, ps, layer_num):
 class TestMLASubPoolSpec(unittest.TestCase):
    def test_entry_bytes_and_dim(self):
        spec = _mla_spec()
        self.assertEqual(spec.kv_cache_dim, _D)
        self.assertEqual(spec.entry_bytes(), _L * _D * _ITEM)
        self.assertEqual(
            spec.entry_bytes(),
            mla_entry_bytes(layer_num=_L, kv_cache_dim=_D, itemsize=_ITEM),
        )
        self.assertEqual(spec.get_dtype(), _DTYPE)
    def test_rejects_nonpositive_dims(self):
        with self.assertRaises(AssertionError):
            MLASubPoolSpec(
@@ -153,7 +128,7 @@ class TestMLAViews(unittest.TestCase):
            n_rows = num_pages * _L * ps
            for v in views:
                self.assertEqual(tuple(v.shape), (n_rows, 1, _D))
-                # contiguous in the (row, dim) sense — .view(-1, ps, D) legality
+                # contiguous in the (row, dim) sense: .view(-1, ps, D) legality
                self.assertEqual(v.stride(0), _D)
                self.assertEqual(v.stride(2), 1)
            flat = raw.view(_DTYPE)
@@ -203,14 +178,6 @@ class TestMLAViews(unittest.TestCase):
 class TestUnifiedKVPoolMLA(unittest.TestCase):
    def test_max_slots_ignore_tail_pad(self):
        pool, full, mamba = _make_unified(page_size=4)
        total = full.entry_bytes() * 64 + mamba.entry_bytes() * 8
        self.assertEqual(pool.max_slots("full"), total // full.entry_bytes())
        self.assertEqual(pool.max_slots("mamba"), total // mamba.entry_bytes())
        # allocation actually carries the pad
        self.assertEqual(pool._raw.numel(), total + 4 * full.entry_bytes())
    def test_reserved_floor_covers_page0_envelope(self):
        ps = 4
        pool, full, mamba = _make_unified(page_size=ps)
@@ -253,15 +220,24 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase):
        self.assertTrue(torch.all(v[7] == 2.5))
    def test_move_kv_cache_moves_page_envelopes(self):
        """Whole page envelopes relocate, in raw bytes and (at ps=4) as read
        back through the per-layer views at the destination kernel ids."""
        for ps in (1, 4):
            pool, kv_pool = self._make(ps=ps)
            num_pages = pool.max_slots("full") // ps
            page_bytes = ps * _L * _D * _ITEM
            env = pool._raw[: num_pages * page_bytes].view(num_pages, page_bytes)
-            src_pages = torch.tensor([num_pages - 2, num_pages - 4])
+            src_pages = torch.tensor([num_pages - 2, num_pages - 4, num_pages - 3])
-            dst_pages = torch.tensor([2, 3])
+            dst_pages = torch.tensor([2, 3, 5])
            env[src_pages[0]] = 7
            env[src_pages[1]] = 9
            if ps == 4:
                # write through the views at src, expect it at dst after the move
                for l in range(_L):
                    for s in range(ps):
                        kv_pool.kv_buffer[l][
                            _kernel_id(int(src_pages[2]) * ps + s, ps, _L)
                        ] = float(l * ps + s + 1)
            # page-major token runs, exactly how compaction expands pages
            offsets = torch.arange(ps, dtype=torch.int64)
            src_t = (src_pages[:, None] * ps + offsets).reshape(-1)
@@ -269,29 +245,15 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase):
            kv_pool.move_kv_cache(dst_t, src_t)
            self.assertTrue(torch.all(env[dst_pages[0]] == 7), f"ps={ps}")
            self.assertTrue(torch.all(env[dst_pages[1]] == 9), f"ps={ps}")
-
+            if ps == 4:
-    def test_move_then_readback(self):
+                for l in range(_L):
-        ps = 4
+                    for s in range(ps):
-        pool, kv_pool = self._make(ps=ps)
+                        got = kv_pool.kv_buffer[l][
-        num_pages = pool.max_slots("full") // ps
+                            _kernel_id(int(dst_pages[2]) * ps + s, ps, _L)
-        src_page, dst_page = num_pages - 3, 5
+                        ]
-        # write through the views at src, expect it at dst after the move
+                        self.assertTrue(
-        for l in range(_L):
+                            torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})"
-            for s in range(ps):
+                        )
                kv_pool.kv_buffer[l][_kernel_id(src_page * ps + s, ps, _L)] = float(
                    l * ps + s + 1
                )
        offsets = torch.arange(ps, dtype=torch.int64)
        kv_pool.move_kv_cache(
            (torch.tensor([dst_page])[:, None] * ps + offsets).reshape(-1),
            (torch.tensor([src_page])[:, None] * ps + offsets).reshape(-1),
        )
        for l in range(_L):
            for s in range(ps):
                got = kv_pool.kv_buffer[l][_kernel_id(dst_page * ps + s, ps, _L)]
                self.assertTrue(
                    torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})"
                )
 class _FakeKVCache:
@@ -325,23 +287,29 @@ class TestTranslateKvLocForKernel(unittest.TestCase):
        mamba_alloc.bind_peer(full_alloc)
        return full_alloc
-    def test_kernel_id_matches_formula_ps1(self):
+    def test_kernel_id_matches_formula(self):
-        alloc = self._build(ps=1)
+        """kernel id = (phys // ps) * (ps * multiplier) + phys % ps, across
-        v = alloc.alloc(8)
+        page sizes, the multiplier-1 physical fallback, and eager compaction."""
-        self.assertIsNotNone(v)
+        for ps, multiplier in ((1, _L), (4, _L), (1, 1)):
-        phys = alloc.translate_kv_loc(v)
+            with self.subTest(page_size=ps, multiplier=multiplier):
-        kernel = alloc.translate_kv_loc_for_kernel(v)
+                alloc = self._build(ps=ps, multiplier=multiplier)
-        self.assertTrue(torch.all(kernel == phys * _L))
+                a = alloc.alloc(4 * ps)
                b = alloc.alloc(4 * ps)
                c = alloc.alloc(4 * ps)
                self.assertIsNotNone(c)
-    def test_kernel_id_matches_formula_paged(self):
+                def check(virt):
-        ps = 4
+                    phys = alloc.translate_kv_loc(virt)
-        alloc = self._build(ps=ps)
+                    expected = (phys // ps) * (ps * multiplier) + phys % ps
-        v = alloc.alloc(3 * ps)
+                    self.assertTrue(
-        self.assertIsNotNone(v)
+                        torch.all(alloc.translate_kv_loc_for_kernel(virt) == expected)
-        phys = alloc.translate_kv_loc(v)
+                    )
-        kernel = alloc.translate_kv_loc_for_kernel(v)
+
-        expected = (phys // ps) * (ps * _L) + phys % ps
+                for virt in (a, b, c):
-        self.assertTrue(torch.all(kernel == expected))
+                    check(virt)
                alloc.free(b)  # eager compaction relocates survivors
                for virt in (a, c):
                    check(virt)
    def test_tombstone_clamps_to_sink(self):
        alloc = self._build(ps=1)
@@ -365,26 +333,6 @@ class TestTranslateKvLocForKernel(unittest.TestCase):
            alloc.translate_kv_loc_for_kernel(x, out=x)
            self.assertTrue(torch.all(x == no_out))
    def test_multiplier_one_falls_back_to_physical(self):
        alloc = self._build(ps=1, multiplier=1)
        v = alloc.alloc(4)
        self.assertIsNotNone(v)
        self.assertTrue(
            torch.all(alloc.translate_kv_loc_for_kernel(v) == alloc.translate_kv_loc(v))
        )
    def test_kernel_id_follows_compaction(self):
        alloc = self._build(ps=1)
        a = alloc.alloc(4)
        b = alloc.alloc(4)
        c = alloc.alloc(4)
        self.assertIsNotNone(c)
        alloc.free(b)  # eager compaction relocates survivors
        phys_a = alloc.translate_kv_loc(a)
        phys_c = alloc.translate_kv_loc(c)
        self.assertTrue(torch.all(alloc.translate_kv_loc_for_kernel(a) == phys_a * _L))
        self.assertTrue(torch.all(alloc.translate_kv_loc_for_kernel(c) == phys_c * _L))
 if __name__ == "__main__":
    unittest.main()
@@ -13,25 +13,13 @@
 # ==============================================================================
 """N-sub-pool construction sweep for ``UnifiedKVPool``.
-The pool accepts N sub-pool specs: exactly one grow-up END, exactly one
+The pool accepts N sub-pool specs -- exactly one grow-up END, exactly one
-grow-down END, and >= 0 "float" MIDDLE pools between their frontiers. These
+grow-down END, and >= 0 "float" MIDDLE pools between their frontiers -- and
-tests pin the constructor contract the N-pool chain machinery builds on:
+sorts them into the canonical chain order
 ``[up end, floats (input order), down end]``. Every sub-pool view spans the
 whole buffer at anchor 0; keeping the bands disjoint is the allocators' job.
-  - canonical chain order ``[up end, floats (input order), down end]`` —
+Pure CPU geometry -- no allocator, no GPU.
    input list order is irrelevant (2-pool configs stay byte-identical);
  - by-name geometry (``max_slots = total_bytes // entry_bytes``,
    ``min_slot_index`` past the shared reserved floor) independent of N;
  - the reserved slot-0 sink covers EVERY sub-pool's page-0 dummy-write
    envelope, floats included (mamba stays page_size=1);
  - validation: unique names, exactly one up + one down, >= 2 specs, and
    per-spec ``_allowed_grow_directions`` narrowing;
  - float sub-pool views build and round-trip like end-pool views (all views
    span the whole buffer at anchor 0; keeping the bands disjoint is the
    allocators' job).
 Pure CPU geometry — no allocator, no GPU.
    python -m pytest test/registered/unit/mem_cache/test_unified_npool_sweep.py -v
 """
 import unittest
@@ -46,10 +34,8 @@ from sglang.srt.mem_cache.unified_memory_pool import (
 )
 from sglang.test.ci.ci_register import register_cpu_ci
-# Plain unittest.TestCase, importing only ci_register -- the deliberate
+# Hermetic convention of this directory's pool tests: plain unittest.TestCase,
-# hermetic convention of the pool-geometry tests in this directory (see
+# only ci_register imported (no heavy sglang.test.test_utils chain).
 # test_multi_ended_allocator.py): no heavy sglang.test.test_utils import
 # chain, so the suite runs in a lean torch-only environment.
 register_cpu_ci(est_time=30, suite="base-a-test-cpu")
 _DEV = "cpu"
@@ -111,34 +97,46 @@ def _chain_names(pool: UnifiedKVPool):
 class TestNPoolCanonicalOrder(unittest.TestCase):
-    def test_two_pool_input_order_irrelevant(self):
+    def test_chain_order_is_canonical(self):
-        for specs in (
+        """Ends canonical (up first, down last) whatever the input order;
-            [_mha("full", "down"), _mamba("mamba", "up")],
+        floats keep INPUT order between them, on every cache-spec kind."""
-            [_mamba("mamba", "up"), _mha("full", "down")],
+        for specs, expect in (
            ([_mha("full", "down"), _mamba("mamba", "up")], ["mamba", "full"]),
            ([_mamba("mamba", "up"), _mha("full", "down")], ["mamba", "full"]),
            (
                [_mha("full", "down"), _mha("swa", "float"), _mamba("conv", "up")],
                ["conv", "swa", "full"],
            ),
            (
                [_mha("swa", "float"), _mamba("conv", "up"), _mha("full", "down")],
                ["conv", "swa", "full"],
            ),
            (
                [_mamba("conv", "up"), _mha("full", "down"), _mha("swa", "float")],
                ["conv", "swa", "full"],
            ),
            (
                [
                    _mha("full", "down"),
                    _mha("f1", "float"),
                    _mamba("state", "up"),
                    _mha("f0", "float", layer_num=1),
                ],
                ["state", "f1", "f0", "full"],
            ),
            (
                [
                    _mamba("state", "up"),
                    _mha("f_mha", "float"),
                    _mla("f_mla", "float", layer_num=1),
                    _mamba("f_mamba", "float", layer_num=1),
                    _mha("full", "down"),
                ],
                ["state", "f_mha", "f_mla", "f_mamba", "full"],
            ),
        ):
-            pool = _make_pool(specs)
+            with self.subTest(inputs=[s.name for s in specs]):
-            self.assertEqual(_chain_names(pool), ["mamba", "full"])
+                self.assertEqual(_chain_names(_make_pool(specs)), expect)
    def test_three_pool_float_in_the_middle(self):
        for specs in (
            [_mha("full", "down"), _mha("swa", "float"), _mamba("conv", "up")],
            [_mha("swa", "float"), _mamba("conv", "up"), _mha("full", "down")],
            [_mamba("conv", "up"), _mha("full", "down"), _mha("swa", "float")],
        ):
            pool = _make_pool(specs)
            self.assertEqual(_chain_names(pool), ["conv", "swa", "full"])
    def test_four_pool_float_input_order_preserved(self):
        pool = _make_pool(
            [
                _mha("full", "down"),
                _mha("f1", "float"),
                _mamba("state", "up"),
                _mha("f0", "float", layer_num=1),
            ]
        )
        # Ends canonical; floats keep INPUT order between them.
        self.assertEqual(_chain_names(pool), ["state", "f1", "f0", "full"])
    def test_by_name_geometry_independent_of_n(self):
        two = _make_pool([_mha("full", "down"), _mamba("mamba", "up")])
@@ -151,9 +149,6 @@ class TestNPoolCanonicalOrder(unittest.TestCase):
                two.total_bytes // two.spec(name).entry_bytes(),
            )
            self.assertEqual(two.max_slots(name), three.max_slots(name))
        for pool in (two, three):
            for s in pool.sub_pool_specs:
                self.assertEqual(pool.anchor_bytes(s.name), 0)
 class TestNPoolValidation(unittest.TestCase):
@@ -165,37 +160,24 @@ class TestNPoolValidation(unittest.TestCase):
        with self.assertRaisesRegex(AssertionError, ">= 2 sub-pools"):
            _make_pool([_mha("full", "down")])
-    def test_two_ups_rejected(self):
+    def test_end_direction_counts_rejected(self):
-        with self.assertRaisesRegex(AssertionError, "exactly one grow-up"):
+        """Exactly one grow-up END and one grow-down END, no more, no less."""
-            _make_pool([_mha("a", "up"), _mamba("b", "up")])
+        for case, specs in (
-
+            ("two_ups", [_mha("a", "up"), _mamba("b", "up")]),
-    def test_missing_down_end_rejected(self):
+            ("missing_down", [_mha("a", "up"), _mha("b", "float")]),
-        with self.assertRaisesRegex(AssertionError, "exactly one grow-up"):
+            (
-            _make_pool([_mha("a", "up"), _mha("b", "float")])
+                "missing_up",
-
+                [_mha("a", "down"), _mha("b", "float"), _mha("c", "float")],
-    def test_missing_up_end_rejected(self):
+            ),
-        with self.assertRaisesRegex(AssertionError, "exactly one grow-up"):
+        ):
-            _make_pool([_mha("a", "down"), _mha("b", "float"), _mha("c", "float")])
+            with self.subTest(case=case):
                with self.assertRaisesRegex(AssertionError, "exactly one grow-up"):
                    _make_pool(specs)
    def test_bogus_direction_rejected_at_spec_level(self):
        with self.assertRaisesRegex(AssertionError, "grow_direction"):
            _mha("a", "sideways")
    def test_float_accepted_on_all_cache_spec_kinds(self):
        # Every cache-class spec kind may float (the chain decides placement).
        pool = _make_pool(
            [
                _mamba("state", "up"),
                _mha("f_mha", "float"),
                _mla("f_mla", "float", layer_num=1),
                _mamba("f_mamba", "float", layer_num=1),
                _mha("full", "down"),
            ]
        )
        self.assertEqual(
            _chain_names(pool), ["state", "f_mha", "f_mla", "f_mamba", "full"]
        )
 class TestReservedFloorWithFloats(unittest.TestCase):
    def test_float_page_envelope_extends_the_sink(self):
@@ -217,8 +199,7 @@ class TestReservedFloorWithFloats(unittest.TestCase):
    def test_too_small_buffer_fails_loud(self):
        # 2048 B with page_size=16 and 128 B/entry MHA specs: the page-0 sink
-        # (16*128 = 2048 B) consumes the whole buffer -> min_slot_index ==
+        # (16*128 = 2048 B) consumes the whole buffer, leaving no slot.
        # max_slots for the MHA pools -> no allocatable slot -> loud error.
        with self.assertRaisesRegex(RuntimeError, "no room"):
            _make_pool(
                [_mamba("state", "up"), _mha("swa", "float"), _mha("full", "down")],
@@ -258,14 +239,6 @@ class TestFloatViews(unittest.TestCase):
        k_views[1][row] = pattern
        torch.testing.assert_close(k_views[1][row], pattern)
    def test_float_mamba_views_zero_visible(self):
        pool = _make_pool(
            [_mamba("state", "up"), _mamba("fstate", "float"), _mha("full", "down")]
        )
        conv_views, temporal = pool.mamba_views_for("fstate")
        self.assertTrue(all(v.eq(0).all() for v in conv_views))
        self.assertTrue(temporal.eq(0).all())
 if __name__ == "__main__":
    unittest.main()
@@ -15,23 +15,7 @@
 SWA KV + mamba/conv state in ONE unified byte buffer, chain
 ``[mamba (up END) | swa (FLOAT) | full (down END)]``.
 Pinned contracts (each guards a distinct failure mode):
  - chain wiring + the swa side being a `FloatMultiEndedAllocator`;
  - the JOINT `available_size()` feasibility contract: `alloc(N)` for
    N == available_size() must succeed (full extends into the high band
    first, then the float extends a single side — the predicate models
    exactly that order; over-promising here is the fail-loud
    `alloc_with_virtual` assert, i.e. a crash in production);
  - `free_swa` tombstones become float HOLES recycled IN PLACE by later
    allocs (steady-state SWA churn == zero copies), while the full side
    keeps the token;
  - the per-request state surface (`UnifiedMambaSlotAllocator` over the
    mamba END) and its independence from the token surface;
  - urgent flushes drain the two ENDS but never touch the float's holes.
 Pure CPU; fakes stand in for the KV pools (data markers verify moves).
    python -m pytest test/registered/unit/mem_cache/test_unified_tri_pool.py -v
 """
 import inspect
@@ -39,11 +23,13 @@ import unittest
 import torch
 import sglang.srt.mem_cache.allocator.unified_sub_pool as mea
 from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
    UnifiedMambaSWATokenToKVPoolAllocator,
 )
-from sglang.srt.mem_cache.allocator.unified_sub_pool import FloatMultiEndedAllocator
+from sglang.srt.mem_cache.allocator.unified_sub_pool import (
    FloatMultiEndedAllocator,
    MultiEndedAllocator,
 )
 from sglang.srt.mem_cache.unified_memory_pool import (
    MambaSubPoolSpec,
    MHASubPoolSpec,
@@ -195,22 +181,6 @@ class TestUnifiedTriPool(unittest.TestCase):
    # -- the joint availability contract --
    def test_available_size_alloc_contract(self):
        for lazy in (False, True):
            _, allocator, kvcache, _ = self._build(lazy_compaction=lazy)
            avail = allocator.available_size()
            self.assertGreater(avail, 0)
            v = allocator.alloc(avail)
            self.assertIsNotNone(
                v, f"alloc(available_size()={avail}) must succeed (lazy={lazy})"
            )
            self.assertEqual(int(v.numel()), avail)
            # Both sides bound for every allocated virtual id.
            fa = allocator.full_attn_allocator
            sa = allocator.swa_attn_allocator
            self.assertTrue(bool((fa.virtual_to_physical[v] >= 0).all()))
            self.assertTrue(bool((sa.virtual_to_physical[v] >= 0).all()))
    def test_available_shrinks_as_state_slots_grow(self):
        _, allocator, _, _ = self._build()
        before = allocator.available_size()
@@ -224,9 +194,8 @@ class TestUnifiedTriPool(unittest.TestCase):
    # -- steady-state SWA churn: tombstones -> holes -> in-place reuse --
    def _swa_interior_block(self, allocator, blocks):
-        """The block whose SWA-physical pages touch neither float boundary --
+        """The block whose SWA-physical pages touch neither float boundary:
-        `free_swa` on it must create interior holes (a boundary block would be
+        `free_swa` on it holes; a boundary block would be absorbed instead."""
        absorbed instead; both are zero-copy, different mechanisms)."""
        sa = allocator.swa_attn_allocator
        for v in blocks:
            pages = set(int(x) for x in sa.virtual_to_physical[v].tolist())
@@ -262,10 +231,9 @@ class TestUnifiedTriPool(unittest.TestCase):
        self.assertEqual(len(sa._inverse_history), 0)  # zero copies
    def test_free_swa_boundary_block_absorbed_zero_copy(self):
-        # The OTHER zero-copy mechanism: a boundary block's tombstones shrink
+        # A boundary block's tombstones shrink the span instead of holing it.
-        # the span, handing bytes back to the neighbours. The shrink is
+        # The shrink needs the hole set on the host, so it is deferred out of
-        # DEFERRED out of the per-step free (it needs the hole set on the
+        # the per-step free and lands in the opportunistic flush.
        # host); the per-step opportunistic flush is where it lands.
        _, allocator, kvcache, _ = self._build()
        blocks = [allocator.alloc(4) for _ in range(2)]
        for v in blocks:
@@ -329,16 +297,7 @@ class TestUnifiedTriPool(unittest.TestCase):
            (pool.max_slots("mamba") - 1) - 1,
        )
-    # -- cost + flush semantics --
+    # -- flush semantics --
    def test_mamba_slot_full_token_cost_formula(self):
        _, allocator, _, _ = self._build()
        e_tok = (
            allocator.full_attn_allocator.entry_bytes
            + allocator.swa_attn_allocator.entry_bytes
        )
        m = allocator.mamba_allocator.entry_bytes_per_page
        self.assertEqual(allocator.mamba_slot_full_token_cost(), -(-m // e_tok))
    def test_urgent_flush_preserves_float_holes(self):
        _, allocator, kvcache, _ = self._build(lazy_compaction=True)
@@ -359,18 +318,11 @@ class TestTriPagedFreeGroup(unittest.TestCase):
    """The tri composite at PAGE SIZE > 1, driven through the production free
    path: free_group_begin -> free_segment -> free_group_end.
-    Regression (GPU eval_440, Inkling ps=128 boot crash): every other tri test
+    Regression: every other tri test runs at page_size=1, where the
-    runs at page_size=1, where the page-REPRESENTATIVE machinery
+    page-representative machinery is dead code. At ps>1 the composite releases
-    (`free_page_reps_group` / `_release_page_reps`, added when the free path
+    reps via `swa_attn_allocator.free(..., _pages=...)`, which the float must
-    was made host-sync-free) is entirely dead code — `free_segment` frees
+    accept or the first decode batch dies; honouring `_pages` is also what
-    directly. At ps>1 the composite releases reps by calling
+    keeps the free path off the data-dependent `torch.unique` host sync.
    `swa_attn_allocator.free(..., _pages=...)`, and the float allocator was
    ported from a base that predates that keyword, so the first real decode
    batch died with `TypeError: unexpected keyword argument '_pages'`.
    `_pages` is not cosmetic: honouring it is what keeps the free path free of
    the data-dependent `torch.unique` host sync, so this also pins that the
    float takes the caller's page ids rather than re-deriving them.
    """
    def _build_paged(self, page_size=4, n_full=64, n_swa=32, n_state=8):
@@ -418,22 +370,7 @@ class TestTriPagedFreeGroup(unittest.TestCase):
        # Capacity fully recovered: the float parked, both ends rewound.
        self.assertTrue(allocator.swa_attn_allocator._is_frontier_transparent())
    def test_float_free_honours_caller_supplied_pages(self):
        """`_pages` must be USED, not merely accepted -- re-deriving it is the
        host sync the paged free path exists to avoid."""
        pool, allocator = self._build_paged()
        v = allocator.alloc(8)
        self.assertIsNotNone(v)
        sa = allocator.swa_attn_allocator
        ps = allocator.page_size
        pages = (v[::ps] // ps).clone()
        live_before = sa._live_pages()
        sa.free(v[::ps] * 0 + v[::ps], _pages=pages)
        self.assertEqual(sa._live_pages(), live_before - pages.numel())
    def test_ungrouped_segment_free_also_reaches_the_float(self):
        """`free_segment` outside a free group releases reps immediately --
        the same float call, one frame shallower."""
        pool, allocator = self._build_paged()
        v = allocator.alloc(8)
        self.assertIsNotNone(v)
@@ -442,13 +379,10 @@ class TestTriPagedFreeGroup(unittest.TestCase):
 class TestTriFreeSwaNoHostSync(unittest.TestCase):
-    """The tri's swa side is the FLOAT, and the float can never run the lazy
+    """The tri's swa side is the FLOAT, which can never run the lazy event
-    event pipeline — so unless the per-step frees carry caller-derived page
+    pipeline: unless the per-step frees carry caller-derived page ids, the tri
-    ids, the tri silently reintroduces the host syncs the sync-free free
+    reintroduces the host syncs the sync-free free path removed. Fixtures run
-    path removed. Poison the ops to pin the property.
+    at page_size > 1; ps==1 short-circuits the page machinery and hides this.
    (Fixtures at page_size > 1 on purpose: ps==1 short-circuits the whole
    page machinery and hides exactly this class of bug.)
    """
    PS = 4
@@ -476,24 +410,9 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase):
            alloc.free_swa(v[: 4 * self.PS], start_pos=0)
        self.assertEqual(alloc.verify_byte_accounting(), [])
    def test_float_free_has_no_stale_slot_item_sync(self):
        """The float's free must not `.item()`-assert per free (the lazy-path
        contract: callers must not double-free; the idle span == p2v-bound +
        holes conservation catches violations without a per-free sync)."""
        alloc = self._tri()
        v = alloc.alloc(4 * self.PS)
        sa = alloc.swa_attn_allocator
        from unittest import mock
        with mock.patch.object(
            torch.Tensor, "item", side_effect=AssertionError("item = host sync")
        ):
            sa.free(v[:: self.PS], _pages=v[:: self.PS] // self.PS)
    def test_fallback_free_swa_still_correct_for_radix_shapes(self):
        """Radix eviction hands arbitrary node values (no start_pos): the
-        dedup fallback must keep working and end in the same state as the
+        dedup fallback must end in the same state as the stride path."""
        stride path."""
        a1, a2 = self._tri(), self._tri()
        v1, v2 = a1.alloc(6 * self.PS), a2.alloc(6 * self.PS)
        self.assertTrue(torch.equal(v1, v2))
@@ -512,12 +431,9 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase):
 class TestGeneralizedRebalance(unittest.TestCase):
    """The float must yield to WHICHEVER end is short, with the direction
-    computed from the layout — not only to the token path's hard-coded side.
+    derived from the caller's growth direction rather than hard-coded to the
-
+    token path's side. `make_room` was always side-agnostic; these pin the
-    The mechanism (`make_room`) was always side-agnostic; these pin the
+    POLICY in `_ask_float_for_room`."""
    POLICY: any end pool's own-alloc shortfall reaches
    `_ask_float_for_room`, which derives the side from the caller's
    growth direction."""
    PS = 4
@@ -527,34 +443,10 @@ class TestGeneralizedRebalance(unittest.TestCase):
        )
        return inst._build_paged(page_size=self.PS)[1]
    def test_state_end_shortfall_slides_the_float_low(self):
        """The previously-missing direction: mamba (grow-up END) starved while
        free bytes idle ABOVE the float. The remedy must slide the float up
        (open its LOW side) and let the state alloc succeed."""
        alloc = self._tri()
        v = alloc.alloc(4 * self.PS)  # places the float mid-region
        self.assertIsNotNone(v)
        ma = alloc.mamba_allocator
        sa = alloc.swa_attn_allocator
        self.assertFalse(sa._is_frontier_transparent())
        # Fill the LOW band exactly: as many state slots as fit below the
        # float's low frontier.
        e_m = ma.entry_bytes_per_page
        fit = (sa._byte_low_frontier() - ma._byte_high_frontier()) // e_m
        self.assertGreater(fit, 0)
        got = ma.alloc(int(fit) * ma.page_size)
        self.assertIsNotNone(got)
        low_before = sa.low_wm_page
        # One more slot does NOT fit below the float -- only a rebalance helps.
        more = ma.alloc(ma.page_size)
        self.assertIsNotNone(more, "state alloc must succeed via float rebalance")
        self.assertGreater(sa.low_wm_page, low_before)  # float slid UP
        self.assertEqual(alloc.verify_byte_accounting(), [])
    def test_direction_is_derived_from_growth_on_both_ends(self):
        """Raw end+float+end chain, BOTH orientations in one fixture: the
-        up-growing end opens the float's LOW side; the down-growing end opens
+        up-growing end opens the float's LOW side, the down-growing end its
-        its HIGH side. No layout assumption survives."""
+        HIGH side."""
        from test_multi_ended_allocator import TestFloatMultiEndedAllocator
        inst = TestFloatMultiEndedAllocator(
@@ -583,8 +475,6 @@ class TestGeneralizedRebalance(unittest.TestCase):
        self.assertLess(fla.high_wm_page, high_before)  # opened HIGH side
    def test_two_pool_chain_rebalance_is_a_noop(self):
        """No float in the chain => the remedy must change nothing (the
        2-pool composites keep their exact pre-existing behavior)."""
        from test_multi_ended_allocator import (
            TestPagedMultiEndedAllocator as _PagedFixture,
        )
@@ -617,9 +507,8 @@ class TestGeneralizedRebalance(unittest.TestCase):
 class TestComputedShortSide(unittest.TestCase):
    """`_ask_float_for_room` must open the side that MEASURES short -- never
-    "the side facing full". These pin the per-side computation, including
+    "the side facing full" -- including the shape with coupled ends on BOTH
-    the coupled-ends-on-both-sides shape a DSV4-style composite
+    sides of the float.
    (C128 | swa-float | C4) will need.
    """
    PS = 4
@@ -638,9 +527,9 @@ class TestComputedShortSide(unittest.TestCase):
    def test_float_share_short_opens_the_state_side(self):
        """RED-LINE: full's demand fits its band, the float's own share fits
-        NEITHER band, and the state side has the larger surplus — the policy
+        NEITHER band, and the state side has the larger surplus -- so the
-        must open the STATE side (the float slides toward full during a
+        policy must open the STATE side, i.e. slide the float toward full
-        TOKEN alloc), which the old "side facing full" policy could never do.
+        during a TOKEN alloc.
        """
        from unittest import mock
@@ -654,10 +543,9 @@ class TestComputedShortSide(unittest.TestCase):
        # Position: slide the float LOW (setup uses the mechanism directly),
        # so the low band is small and the geometry below is expressible.
        b_low0, b_high0 = self._sides(alloc)
-        # Two positioning moves: pack the float low (leapfrog over-opens by
+        # Pack the float low (leapfrog over-opens by design), then reopen LOW
-        # design), then open the LOW side back to ~2 full-pages -- small
+        # to ~2 full-pages: F must outgrow it, yet the need_n window below
-        # enough that F outgrows it, wide enough that the integer need_n
+        # must stay non-empty.
        # window below is non-empty.
        sa.make_room(side="high", min_bytes=b_low0 + b_high0 - 2 * e_f)
        sa.make_room(side="low", min_bytes=2 * e_f)
@@ -686,10 +574,8 @@ class TestComputedShortSide(unittest.TestCase):
        self.assertEqual(calls[0]["side"], "low")  # the STATE side
    def test_full_side_short_target_matches_the_closed_form(self):
-        """Equivalence: when the full side is the short one (today's only
+        """When the full side is the short one, the ask must equal the
-        reachable end-shortage), the ask must equal the documented formula
+        closed form demand + max(0, F - far_surplus) + slack."""
        demand + max(0, F - far_surplus) + slack — i.e. the historical
        behavior is the special case, preserved."""
        from unittest import mock
        alloc = self._tri()
@@ -717,9 +603,9 @@ class TestComputedShortSide(unittest.TestCase):
        self.assertEqual(calls[0]["min_bytes"], want)
    def test_two_coupled_ends_lands_demand_on_both_sides(self):
-        """DSV4 shape (C128 | float | C4): a coupled set with ends on BOTH
+        """A coupled set with ends on BOTH sides of the float: one-side-short
-        sides. One-side-short must open that side; BOTH-sides-short must not
+        must open that side; BOTH-sides-short must not move at all
-        move at all (relocation is zero-sum between the bands)."""
+        (relocation is zero-sum between the bands)."""
        from unittest import mock
        alloc = self._tri()
@@ -730,8 +616,8 @@ class TestComputedShortSide(unittest.TestCase):
            alloc.full_attn_allocator,
            alloc.mamba_allocator,
        )
-        # Synthetic coupling: the state end joins the demand vector, exactly
+        # Synthetic coupling: the state end joins the demand vector, the
-        # the override a DSV4-style composite would ship.
+        # override a composite with ends on both sides would ship.
        need = lambda self, t: {
            fa: -(-t // self.page_size),
            sa: -(-t // self.page_size),
@@ -764,11 +650,14 @@ class TestComputedShortSide(unittest.TestCase):
                self.assertEqual(calls[0]["side"], "high")
    def test_nothing_short_means_no_relocation(self):
-        """Everything fits -> the policy must not move a single page."""
+        """Everything fits -> no page moves. The tri's token vector carries
        {mamba: 0}: a zero entry must not move the float either."""
        from unittest import mock
        alloc = self._tri()
        alloc.alloc(4 * self.PS)
        demand = alloc._alloc_demand(2 * self.PS)
        self.assertEqual(demand[alloc.mamba_allocator], 0)
        sa = alloc.swa_attn_allocator
        with mock.patch.object(
            sa, "make_room", side_effect=AssertionError("needless move")
@@ -779,12 +668,9 @@ class TestComputedShortSide(unittest.TestCase):
 class TestFloatPolicyTotalTarget(unittest.TestCase):
    """`make_room`'s min_bytes is a TARGET for the whole band, not a delta.
-    Regression: the band-level policy passed `deficit + one page` — with a
+    Regression: an ask shaped as `deficit + one page` lands BELOW the current
-    PARTIALLY free band that is below the current gap, so `make_room`
+    gap when the band is only PARTIALLY free, so `make_room` no-ops and the
-    no-oped and the allocation failed even though the float had room to
+    allocation fails though the float had room to slide.
    slide. (Its own test missed this because it filled the band exactly,
    making deficit ≈ the whole need.) The demand-vector policy computes the
    total target, so a partial gap under-asks never.
    """
    PS = 4
@@ -804,34 +690,17 @@ class TestFloatPolicyTotalTarget(unittest.TestCase):
        gap_slots = int((sa._byte_low_frontier() - ma._byte_high_frontier()) // e_m)
        self.assertGreater(gap_slots, 2)
        low_before = sa.low_wm_page
-        # Need = partial-gap + 3: the old delta-ask was BELOW the current
+        # partial-gap + 3: a delta-shaped ask lands below the current gap.
        # gap, so nothing moved and this returned None.
        got = ma.alloc((gap_slots + 3) * ma.page_size)
        self.assertIsNotNone(got, "partial-gap shortfall must relocate, not fail")
        self.assertGreater(sa.low_wm_page, low_before)
        self.assertEqual(alloc.verify_byte_accounting(), [])
    def test_zero_demand_bands_are_inert(self):
        """The tri's token vector carries {mamba: 0}: a zero entry must
        neither move the float for mamba's sake nor trip the index guard."""
        from unittest import mock
        alloc = self._tri()
        alloc.alloc(4 * self.PS)
        demand = alloc._alloc_demand(2 * self.PS)
        self.assertEqual(demand[alloc.mamba_allocator], 0)
        sa = alloc.swa_attn_allocator
        with mock.patch.object(
            sa, "make_room", side_effect=AssertionError("needless move")
        ):
            alloc._ask_float_for_room(1)  # nothing short -> no relocation
 class TestTriDeferredAbsorption(unittest.TestCase):
    """Boundary absorption is deferred out of the per-step free and paid once
-    at a quiescent point — the base allocator's model (its lazy free does "no
+    at a quiescent point. These pin WHERE it is paid, and that skipping it
-    boundary absorb" and `_flush` pays a single D2H). These pin WHERE it is
+    stays merely conservative."""
    now paid, and that skipping it stays merely conservative."""
    PS = 4
@@ -855,9 +724,8 @@ class TestTriDeferredAbsorption(unittest.TestCase):
        self.assertEqual(alloc.verify_byte_accounting(), [])
    def test_shortfall_ladder_absorbs_before_the_deficit_math(self):
-        """The zero-copy rung must run FIRST: a stale-wide span would inflate
+        """The zero-copy rung must run FIRST: a stale-wide span inflates the
-        the rebalance deficit and buy a `make_room` relocation the shrink
+        rebalance deficit and buys a relocation the shrink already covers."""
        already covers."""
        alloc = self._tri()
        v = alloc.alloc(8 * self.PS)
        sa = alloc.swa_attn_allocator
@@ -884,9 +752,9 @@ class TestTriDeferredAbsorption(unittest.TestCase):
    def test_clean_flush_skips_the_d2h_entirely(self):
        """Only `free` can put a hole ON a boundary (alloc DRAINS holes into
-        live pages; extension adds live pages), so with nothing freed since
+        live pages, extension adds live pages), so with nothing freed since
-        the last absorb the walk provably finds nothing — and must not pay
+        the last absorb the walk provably finds nothing and must not pay the
-        the D2H. Steady churn with only interior holes then costs no sync."""
+        D2H."""
        from unittest import mock
        alloc = self._tri()
@@ -902,9 +770,8 @@ class TestTriDeferredAbsorption(unittest.TestCase):
            self.assertEqual(sa._flush(urgent=False), 0)
    def test_alloc_between_frees_cannot_hide_a_boundary_hole(self):
-        """Soundness of the skip: an alloc drains holes and can change the
+        """An alloc drains holes and can restore a previously-seen hole
-        hole COUNT back to a previously-seen value, so the flag must be armed
+        COUNT, so the flag must be armed by `free`, not read off `numel()`."""
        by `free`, not inferred from `numel()`."""
        alloc = self._tri()
        v = alloc.alloc(8 * self.PS)
        sa = alloc.swa_attn_allocator
@@ -1019,9 +886,9 @@ class TestTriFactorySizing(unittest.TestCase):
 class TestTriPoolHardening(unittest.TestCase):
-    """C1.7 pressure lanes: the planned-rebalance remedy in the alloc path
+    """Pressure lanes: the planned-rebalance remedy in the alloc path (a
-    (a mis-positioned float must not fail an alloc that fits in total bytes),
+    mis-positioned float must not fail an alloc that fits in total bytes),
-    retract-loop convergence through check_decode_capacity, and bounded copy
+    retract-loop convergence through `check_decode_capacity`, and bounded copy
    traffic under alternating end pressure.
    """
@@ -1029,9 +896,9 @@ class TestTriPoolHardening(unittest.TestCase):
        return TestUnifiedTriPool._build(self, **kw)
    def test_alloc_rebalances_a_blocking_float(self):
-        # Fill much of the high band so the float (midpoint-placed) walls off
+        # Fill much of the high band so the midpoint-placed float walls the
-        # the low band's free bytes from `full`; the next alloc must succeed
+        # low band's free bytes off from `full`; the next alloc must SLIDE it
-        # by SLIDING the float, not fail while total bytes suffice.
+        # rather than fail while total bytes suffice.
        _, allocator, kvcache, _ = self._build(n_full=32, n_swa=24, n_state=8)
        sa = allocator.swa_attn_allocator
        v0 = allocator.alloc(4)  # places the float at the region midpoint
@@ -1043,9 +910,8 @@ class TestTriPoolHardening(unittest.TestCase):
        b_high_pages = fa._current_gap_bytes() // fa.entry_bytes_per_page
        grab = fa.alloc(max(0, (b_high_pages - 2)))
        self.assertIsNotNone(grab)
-        # The honest gate under-reports (no slide credit) -- asking BEYOND it
+        # The gate under-reports (no slide credit), so asking BEYOND it fires
-        # is what fires the rebalance remedy; the ask still fits total free
+        # the remedy; the ask still fits the free bytes the LOW band holds.
        # bytes because the LOW band holds them behind the float.
        avail = allocator.available_size()
        need = avail + 4
        live_before = sa._live_pages()
@@ -1063,9 +929,8 @@ class TestTriPoolHardening(unittest.TestCase):
        self.assertEqual(allocator.verify_byte_accounting(), [])
    def test_check_decode_capacity_retract_convergence(self):
-        # Simulated retract loop: requests' token blocks freed one at a time
+        # Retract loop: token blocks freed one at a time until the next step
-        # until the next-step allocation fits; must converge before bs=1 and
+        # fits; must converge before bs=1 and never report capacity early.
        # never report capacity while the gate is short.
        _, allocator, _, _ = self._build(n_full=32, n_swa=24, n_state=8)
        reqs = []
        while True:
@@ -1085,9 +950,8 @@ class TestTriPoolHardening(unittest.TestCase):
        self.assertEqual(allocator.verify_byte_accounting(), [])
    def test_alternating_pressure_copy_traffic_bounded(self):
-        # Alternating full-grow / swa-churn cycles: total float moves stay
+        # Alternating full-grow / swa-churn: hole recycling and absorption do
-        # bounded (hole recycling + absorption do the steady-state work; the
+        # the steady-state work, so total float moves stay bounded.
        # rebalance fires only on real positional deficits).
        _, allocator, kvcache, _ = self._build(n_full=48, n_swa=32, n_state=8)
        sa = allocator.swa_attn_allocator
        fa = allocator.full_attn_allocator
@@ -1124,13 +988,11 @@ class TestJointCapacityIsHonoured(unittest.TestCase):
    """`alloc(available_size())` must never fail.
    REGRESSION: the joint predicate priced the swa float's extension in RAW
-    BYTES while `take_physical_pages` can only use whole pages on the float's
+    BYTES, but `take_physical_pages` yields only whole pages on the float's
-    OWN grid -- `_region_bounds_pages` rounds the band's low edge UP. The
+    OWN grid, whose low edge is rounded up to a multiple of the NEIGHBOUR's
-    bounding frontier is a multiple of the NEIGHBOUR's entry size, which is
+    entry size -- so the budget credited a page the grid could not yield.
-    unrelated to the float's, so the byte budget credited a page the grid could
+    Swept over geometries so a symmetric mistake on the FULL side surfaces
-    not yield and the very first alloc tripped `alloc_with_virtual`'s backstop
+    here too.
    assert. Swept over geometries rather than pinned to one, so a symmetric
    mistake on the FULL side would surface here too.
    """
    def _build(self, *, page_size, n_full, n_swa, n_state, lazy, specs):
@@ -1162,8 +1024,7 @@ class TestJointCapacityIsHonoured(unittest.TestCase):
    def test_fresh_boot_alloc_of_available_size_succeeds(self):
        # Geometries chosen so the mamba end's frontier (a multiple of the
-        # STATE entry size) lands off the swa float's page grid -- the
+        # STATE entry size) lands off the swa float's page grid.
        # misalignment the byte budget used to ignore.
        for page_size in (1, 2, 4):
            for fl, sl, ml in ((4, 3, 1), (4, 2, 2), (6, 3, 1), (3, 5, 2)):
                for n_full, n_swa, n_state in ((24, 16, 4), (32, 16, 8), (20, 12, 6)):
@@ -1200,6 +1061,15 @@ class TestJointCapacityIsHonoured(unittest.TestCase):
                                f"alloc(available_size()={n}) returned None",
                            )
                            self.assertEqual(out.numel(), n)
                            # Both sides bound for every allocated virtual id.
                            fa = alloc.full_attn_allocator
                            sa = alloc.swa_attn_allocator
                            self.assertTrue(
                                bool((fa.virtual_to_physical[out] >= 0).all())
                            )
                            self.assertTrue(
                                bool((sa.virtual_to_physical[out] >= 0).all())
                            )
    def test_available_size_never_exceeds_the_float_page_grid(self):
        """Direct form: the joint answer, converted to float pages, must fit
@@ -1235,13 +1105,11 @@ class TestJointCapacityIsHonoured(unittest.TestCase):
 class TestFloatRelocationIsOrderedAgainstTheForward(unittest.TestCase):
    """Float relocation must settle the in-flight forward BEFORE its first copy.
-    REGRESSION: `make_room` / `compact_holes` issued `move_kv_cache` and rebound
+    REGRESSION: `make_room` / `compact_holes` issued `move_kv_cache` and
-    `virtual_to_physical` with no ordering against the running forward, so the
+    rebound `virtual_to_physical` with no ordering against the running
-    copy could carry pre-write bytes and the rebind then pointed every later
+    forward, so a copy could carry pre-write bytes and every later reader saw
-    reader at a destination that never received those writes -- silently wrong
+    a destination that never received those writes -- silently wrong KV, no
-    KV, no crash. The END pools guard exactly this hazard in
+    crash. The END pools guard the same hazard via `_settle_inflight_forward`.
    `_flush(urgent=True)` via `_settle_inflight_forward`; the float had no
    `forward_stream` / `wait_event` / settle call anywhere in its body.
    """
    def _tri(self, lazy=True):
@@ -1318,11 +1186,9 @@ class TestFloatRelocationIsOrderedAgainstTheForward(unittest.TestCase):
        self.assertEqual(order[0], "settle", f"first action was not a settle: {order}")
    def test_the_settle_is_a_stream_wait_not_a_host_sync(self):
-        """Pin the mechanism: `_settle_inflight_forward` must stream-wait, so
+        """The settle before a float move must be a stream wait: a host sync
-        the fix costs no host sync on the shortfall path."""
+        there would land on the alloc-shortfall path every time the float moves."""
-        src = inspect.getsource(
+        src = inspect.getsource(MultiEndedAllocator._settle_inflight_forward)
            mea.MultiEndedAllocator._settle_inflight_forward  # noqa: SLF001
        )
        self.assertIn("wait_event", src)
        self.assertNotIn(".item()", src)
        self.assertNotIn("synchronize()", src)
@@ -1331,14 +1197,12 @@ class TestFloatRelocationIsOrderedAgainstTheForward(unittest.TestCase):
 class TestFloatHoleCreditIsPerSide(unittest.TestCase):
    """A float's schedulable credit must follow the side the holes are on.
-    REGRESSION: the base `_peer_drainable_hole_bytes` asks
+    REGRESSION: the base `_peer_drainable_hole_bytes` picks its neighbour from
-    `_growth_side_neighbor()`, which reads `grow_direction`. A float's is
+    `grow_direction`, which for a float is "float" -- it fell through to
-    "float", so the base fell through to `low_peer` -- it never saw the HIGH
+    `low_peer`, never saw the HIGH neighbour, and its single scalar was added
-    neighbour, and the single scalar it returned was then added to
+    to `max(gap_low, gap_high)`, landing a LOW neighbour's holes on the HIGH
-    `max(gap_low, gap_high)`, landing a LOW neighbour's holes on the HIGH gap.
+    gap. Over-reporting `schedulable_available_size` admits work the shortfall
-    Over-reporting `schedulable_available_size` makes the scheduler admit work
+    ladder cannot satisfy.
    the shortfall ladder cannot satisfy, which the caller treats as a
    memory-estimation bug.
    """
    def _float(self):
@@ -1365,15 +1229,6 @@ class TestFloatHoleCreditIsPerSide(unittest.TestCase):
        )
        return alloc, alloc.swa_attn_allocator
    def test_credit_sees_both_neighbours(self):
        _alloc, flt = self._float()
        self.assertIsInstance(flt, FloatMultiEndedAllocator)
        low = flt._side_drainable_hole_bytes("low")
        high = flt._side_drainable_hole_bytes("high")
        self.assertEqual(flt._peer_drainable_hole_bytes(), max(low, high))
        # The base would have answered with the LOW side alone.
        self.assertGreaterEqual(flt._peer_drainable_hole_bytes(), high)
    def test_schedulable_never_exceeds_the_sum_of_the_two_sides(self):
        """Upper bound that the undirected scalar could violate: no side may be
        credited with the other side's holes on top of its own gap."""
@@ -13,30 +13,19 @@
 # ==============================================================================
 """`--enable-page-major-kv-layout` full-attention backend allowlist.
-Two-way gate (see `_handle_page_major_kv_layout`), because the unified pool
+`handle_page_major_kv_layout` gates two ways, because the per-layer views the
-exposes per-layer views and nothing else:
+unified pool exposes are all the allowlisted backends can read: an MLA arm
-  * unified-memory MLA models (`build_mla_views`) allow the whole wired
+(`build_mla_views`), an MHA/SWA arm (`build_mha_views`), and no page-major arm
-    paged MLA family -- `fa3`, `flashinfer`'s MLA backend, `trtllm_mla` with
+at all without the unified pool. `fa3` is the resolved default on pre-Blackwell
-    its `cutedsl_mla` / `tokenspeed_mla` subclasses, and `flashmla` (ps=64
+hosts, so its absence from an arm makes `--enable-unified-memory` fail to boot
-    snap);
+under its own default configuration.
  * unified-memory MHA/SWA models (`build_mha_views`) allow `fa3` /
    `fa4` / `flashinfer` / `trtllm_mha` alongside Triton;
  * plain `--enable-page-major-kv-layout` without the unified pool keeps the
    envelope-strided 4-D views only the stride-aware Triton kernels read.
-The same handler also screens the pool itself: the MHA/SWA per-layer views need
+The same handler screens the pool itself: the MHA/SWA per-layer views need
 uniform K/V rows, so an asymmetric-K/V model (MiMoV2: head_dim 192 !=
-v_head_dim 128) cannot run `--enable-unified-memory` at all and is rejected on
+v_head_dim 128) is rejected on EVERY backend, Triton included. MLA models are
-EVERY backend, Triton included. MLA models are exempt -- their sub-pool keeps
+exempt -- their sub-pool keeps one latent row per layer, and real MLA configs
-one latent row per layer, and several MLA configs (Kimi-Linear: head_dim 72,
+(Kimi-Linear: head_dim 72, v_head_dim 128) report asymmetric dims while running
-v_head_dim 128) report asymmetric dims while running the unified pool today.
+the unified pool today.
 Pinned here so no arm silently widens to an unwired backend (`cutlass_mla`,
 `aiter`) and no arm silently narrows: `fa3` is the resolved default on
 pre-Blackwell hosts, so its absence from an arm makes `--enable-unified-memory`
 fail to boot under its own default configuration.
    python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
 """
 import unittest
@@ -59,11 +48,8 @@ def _accepts(
    linear_prefill: str | None = None,
    has_asymmetric_kv: bool = False,
 ) -> bool:
-    """Run just `_handle_page_major_kv_layout` against a minimal stand-in.
+    """Run just `handle_page_major_kv_layout` against a minimal stand-in, since
-
+    ServerArgs' real constructor pulls in a model config."""
    ServerArgs' real constructor pulls in a model config; this exercises the
    single handler under test with the fields it reads.
    """
    sa = ServerArgs.__new__(ServerArgs)
    for name, value in {
        "enable_unified_memory": unified,
@@ -145,21 +131,9 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
                f"{backend} is an MLA kernel and must stay out of the MHA arm",
            )
    def test_plain_page_major_arm_is_gated_at_boot(self):
        """The strided views were removed: --enable-page-major-kv-layout
        without --enable-unified-memory must be rejected up front for EVERY
        backend, Triton included, until the per-layer-view reimplementation."""
        for backend in ("triton",) + self.PER_LAYER_VIEW_MLA_BACKENDS:
            for use_mla in (True, False):
                self.assertFalse(
                    _accepts(backend, use_mla=use_mla, unified=False),
                    f"{backend} must be rejected on the static page-major arm",
                )
    def test_asymmetric_kv_mha_model_cannot_use_unified_memory(self):
-        """head_dim != v_head_dim (MiMoV2): no uniform rows, so no per-layer views
+        """The rejection is the POOL's, not a backend's, so it must fire on
-        and no unified pool. The rejection is the POOL's, not a backend's, so
+        every backend -- Triton included."""
        it must fire on every backend -- Triton included."""
        for backend in ("triton",) + self.PER_LAYER_VIEW_MHA_BACKENDS:
            self.assertFalse(
                _accepts(backend, use_mla=False, has_asymmetric_kv=True),
@@ -168,10 +142,8 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
            )
    def test_asymmetric_dims_do_not_screen_out_mla(self):
-        """MLA stores one latent row per layer, so its K/V head dims never have
+        """Screening on `has_asymmetric_kv` alone would lock every real MLA
-        to agree -- and real MLA configs report them as unequal (Kimi-Linear:
+        config out of the unified pool."""
        head_dim 72, v_head_dim 128). Screening on `has_asymmetric_kv` alone
        would lock every one of them out of the unified pool."""
        for backend in ("triton",) + self.PER_LAYER_VIEW_MLA_BACKENDS:
            self.assertTrue(
                _accepts(backend, use_mla=True, has_asymmetric_kv=True),
@@ -180,9 +152,8 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
            )
    def test_page_major_rejected_without_unified_memory(self):
-        """--enable-page-major-kv-layout without the unified pool is rejected
+        """There is no static page-major arm today, so the flag alone is
-        outright, Triton included: the static page-major arm went away with
+        rejected outright -- Triton included."""
        the strided views and awaits its per-layer-view reimplementation."""
        for backend in ("triton",) + tuple(
            set(self.PER_LAYER_VIEW_MLA_BACKENDS + self.PER_LAYER_VIEW_MHA_BACKENDS)
        ):