refactor(unified-memory): translate the KV write location once, at ForwardBatch construction (#35245)

Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

test/registered/attention/test_create_kvindices.py

@@ -75,6 +75,60 @@ class TestCreateKvIndices(CustomTestCase):
         for batch in BATCH:
             self._run_test(batch, MAX_BATCH, MAX_CONTEXT_LEN)
 
+    def _run_page_table_test(self, batch, ps, with_window_start):
+        """ENTRY_PAGE_SIZE > 1: the source is a PAGE-granular table (the unified
+        pool's read table); the kernel must reconstruct token ids by the affine
+        rule token = entry * ps + pos % ps -- including the kv_start_idx
+        (sliding-window) offset path, whose pos is an absolute token position."""
+        max_batch, max_pages = 64, 128
+        page_table = torch.randint(
+            0, 1 << 20, (max_batch, max_pages), dtype=torch.int32
+        )
+        req_pool_indices = torch.tensor(
+            np.random.choice(range(max_batch), size=batch, replace=False),
+            dtype=torch.int32,
+        )
+        lens = torch.tensor(
+            np.random.randint(1, max_pages * ps, size=batch), dtype=torch.int32
+        )
+        if with_window_start:
+            start = torch.clamp(
+                lens - torch.randint(1, ps * 3, (batch,), dtype=torch.int32), min=0
+            )
+            gather_lens = lens - start
+        else:
+            start, gather_lens = None, lens
+        kv_indptr = torch.zeros((batch + 1,), dtype=torch.int32)
+        kv_indptr[1:] = torch.cumsum(gather_lens, dim=0)
+
+        # ref: absolute positions [start, start+len) through the affine rule
+        refs = []
+        for i in range(batch):
+            s = int(start[i]) if start is not None else 0
+            pos = torch.arange(s, s + int(gather_lens[i]), dtype=torch.int64)
+            entry = page_table[int(req_pool_indices[i])][pos // ps].to(torch.int64)
+            refs.append(entry * ps + pos % ps)
+        ref = torch.cat(refs).contiguous()
+
+        out = torch.empty(int(kv_indptr[-1]), dtype=torch.int64)
+        create_flashinfer_kv_indices_triton[(batch,)](
+            page_table,
+            req_pool_indices,
+            gather_lens,
+            kv_indptr,
+            start,
+            out,
+            page_table.size(1),
+            ENTRY_PAGE_SIZE=ps,
+        )
+        self.assertTrue(torch.equal(ref, out))
+
+    def test_page_table_source_reconstruction(self):
+        for batch in (1, 37):
+            for ps in (4, 64, 256):
+                self._run_page_table_test(batch, ps, with_window_start=False)
+                self._run_page_table_test(batch, ps, with_window_start=True)
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_full_loc_fast_path.py

@@ -19,7 +19,7 @@ pools hold none and never translate — the write loc reaching `set_kv_buffer` i
 always PHYSICAL. Two routing contracts are pinned here:
 
 1. Full-attention. The full-physical loc is carried in `KVWriteLoc.full_loc`
-   (from `ForwardBatch.out_cache_loc_full_physical`) and written directly.
+   (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.
@@ -264,8 +264,10 @@ class TestHybridLinearMLARouting(unittest.TestCase):
     - `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the
       pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it
       carries the DENSE loc), else the raw `loc` (static pool, already physical).
-    - `set_mla_kv_buffer` / `get_mla_kv_buffer` receive VIRTUAL locs and apply
-      `_full_translate` exactly once (identity for a static pool)."""
+    - `set_mla_kv_buffer` forwards `loc` untouched (kernel-facing since the
+      ForwardBatch rebind); `get_mla_kv_buffer` applies `_full_translate`
+      exactly once (its indices are req_to_token-produced, virtual under the
+      unified pool)."""
 
     def _make_bare_pool(self, translate=None):
         from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
@@ -274,7 +276,7 @@ class TestHybridLinearMLARouting(unittest.TestCase):
         pool.full_kv_pool = _RecordingMLAPool()
         pool.use_mla = True
         pool.full_attention_layer_id_mapping = {0: 0}
-        pool._full_translate = translate if translate is not None else (lambda x: x)
+        pool._full_translate = translate if translate is not None else (lambda ids: ids)
         return pool
 
     def test_mla_writes_full_loc_from_write_loc(self):
@@ -311,28 +313,26 @@ class TestHybridLinearMLARouting(unittest.TestCase):
         forwarded, _ = pool.full_kv_pool.calls[0]
         self.assertIs(forwarded, phys_loc)
 
-    def test_set_mla_kv_buffer_translates_exactly_once(self):
-        calls = []
-
-        def translate(ids):
-            calls.append(ids)
-            return ids + 100
-
-        pool = self._make_bare_pool(translate=translate)
-        virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
+    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
+        (rebind_write_loc, kernel-facing-first); a door that translated
+        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)
 
-        pool.set_mla_kv_buffer(
-            layer, virtual_loc, torch.zeros(3, 1, 6), torch.zeros(3, 1, 2)
-        )
+        pool.set_mla_kv_buffer(layer, loc, torch.zeros(3, 1, 6), torch.zeros(3, 1, 2))
 
-        self.assertEqual(len(calls), 1)
         self.assertEqual(len(pool.full_kv_pool.mla_set_calls), 1)
-        self.assertTrue(
-            torch.all(pool.full_kv_pool.mla_set_calls[0] == virtual_loc + 100)
-        )
+        self.assertIs(pool.full_kv_pool.mla_set_calls[0], loc)
 
     def test_get_mla_kv_buffer_translates_exactly_once(self):
+        """READ door: `loc` is produced from req_to_token (VIRTUAL under the
+        unified pool), so the get side still translates here — exactly once.
+        The WRITE door (case above) never translates: the split is the write
+        flip's contract."""
         calls = []
 
         def translate(ids):

test/registered/unit/mem_cache/test_kv_index_translator.py

@@ -0,0 +1,684 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""KVIndexTranslator -- the read-path id translator.
+
+Covers, CPU-only (the builder's pure-torch reference path; GPU parity of the
+Triton kernel is a later CUDA CI pin):
+  - 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
+
+register_cpu_ci(est_time=8, suite="base-a-test-cpu")
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+from test_multi_ended_allocator import _FakeUnifiedSWAKVPool
+
+from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
+from sglang.srt.mem_cache.multi_ended_allocator import (
+    UnifiedSWATokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool
+
+_DEV = "cpu"
+_FULL_L = 2
+_SWA_L = 3
+
+
+def _build_composite(ps, collapse=False, n_full_pages=16, n_swa_pages=8):
+    full_spec = MHASubPoolSpec(
+        name="full",
+        layer_num=_FULL_L,
+        head_num=2,
+        head_dim=4,
+        store_dtype=torch.float16,
+        grow_direction="up",
+    )
+    swa_spec = MHASubPoolSpec(
+        name="swa",
+        layer_num=_SWA_L,
+        head_num=2,
+        head_dim=4,
+        store_dtype=torch.float16,
+        grow_direction="down",
+    )
+    n_full, n_swa = n_full_pages * ps, n_swa_pages * ps
+    total = n_full * full_spec.entry_bytes() + n_swa * swa_spec.entry_bytes()
+    pool = UnifiedKVPool(
+        total_bytes=total,
+        sub_pool_specs=[full_spec, swa_spec],
+        device=_DEV,
+        enable_memory_saver=False,
+        page_size=ps,
+    )
+    kvcache = _FakeUnifiedSWAKVPool(pool)
+    allocator = UnifiedSWATokenToKVPoolAllocator(
+        unified_buffer=pool,
+        kvcache=kvcache,
+        device=_DEV,
+        full_max_total_num_tokens=n_full,
+        swa_max_total_num_tokens=n_swa,
+        page_size=ps,
+        need_sort=False,
+        forward_stream=None,
+    )
+    if collapse:
+        # The multiplier-1 arm, where kernel-facing ids ARE the physical ones.
+        # No unified sub-pool reports 1 today, so pin the regime here.
+        allocator.full_attn_allocator.kernel_page_multiplier = 1
+        allocator.swa_attn_allocator.kernel_page_multiplier = 1
+    # The fake IS the runner's token_to_kv_pool, and the real UnifiedSWAKVPool
+    # carries the pool-level full->swa translate, so the fake must too.
+    kvcache.translate_loc_from_full_to_swa = allocator.translate_loc_from_full_to_swa
+    return allocator
+
+
+def _make_source(allocator, req_to_token, ps):
+    """The owning runner's source: its token_to_kv_pool IS the allocator's own
+    kvcache. (A runner can share the allocator while owning a different pool --
+    see TestPoolOwnership.)"""
+    return KVIndexTranslator(
+        req_to_token=req_to_token,
+        token_to_kv_pool_allocator=allocator,
+        token_to_kv_pool=allocator.get_kvcache(),
+        page_size=ps,
+        device=_DEV,
+    )
+
+
+def _reference_table(req_to_token, req_pool_indices, seq_lens, v2p, mult, ps, width):
+    """Independent python derivation of the read-table formula."""
+    bs = req_pool_indices.numel()
+    out = torch.zeros((bs, width), dtype=torch.int32)
+    for b in range(bs):
+        req = int(req_pool_indices[b])
+        n_pages = -(-int(seq_lens[b]) // ps)
+        for c in range(min(n_pages, width)):
+            tok = int(req_to_token[req, c * ps])
+            page = 0 if tok < 0 else tok // ps
+            out[b, c] = max(int(v2p[page]) * mult, 0)
+    return out
+
+
+class TestPassthrough(unittest.TestCase):
+    def test_non_unified_returns_same_objects(self):
+        """The strict-passthrough property: no copy, no branch, the exact
+        tensors backends read today. A regression here (any tensor op on the
+        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,
+            token_to_kv_pool_allocator=SimpleNamespace(),  # not a composite
+            token_to_kv_pool=SimpleNamespace(),  # not an SWAKVPool
+            page_size=1,
+            device=_DEV,
+        )
+        self.assertFalse(src.is_translating)
+        rows = torch.tensor([2, 0])
+        view = src.build_index_table(
+            req_pool_indices=rows, seq_lens=torch.tensor([5, 3])
+        )
+        self.assertIs(view.ids, req_to_token)
+        self.assertIs(view.row_ids, rows)
+        self.assertEqual(view.row_stride, req_to_token.stride(0))
+        self.assertEqual(view.entry_page_size, 1)
+        self.assertFalse(view.is_translated)
+        self.assertIsNone(view.sliding_window_ids)
+        # And the translate surface is the identity, not a wrapped copy.
+        t = torch.tensor([1, 2, 3])
+        self.assertIs(src.translate_full_attn_ids(t), t)
+
+
+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)."""
+    width = 16 * ps
+    req_to_token = torch.full((len(lens), width), -1, dtype=torch.int64)
+    for r, n in enumerate(lens):
+        n_alloc = -(-n // ps) * ps  # page-aligned virtual run
+        v = allocator.alloc(n_alloc)
+        assert v is not None
+        req_to_token[r, :n] = v[:n]
+    return (
+        req_to_token,
+        torch.arange(len(lens), dtype=torch.int64),
+        torch.tensor(lens, dtype=torch.int64),
+    )
+
+
+class TestReadTableBuild(unittest.TestCase):
+
+    def test_read_table_matches_reference_dense_and_strided(self):
+        """The load-bearing formula pin: full AND swa read tables equal
+        the independent per-element derivation, across page sizes and both
+        multiplier regimes (strided=1, dense=2L). The swa table agreeing with
+        a formula over VIRTUAL ids is also the never-chained-through-
+        full-physical proof."""
+        for ps in (1, 4):
+            for collapse in (True, False):
+                allocator = _build_composite(ps, collapse=collapse)
+                full_mult = allocator.kernel_page_multiplier
+                swa_mult = allocator.swa_kernel_page_multiplier
+                req_to_token, rows, seq_lens = _alloc_and_fill(
+                    allocator, ps, lens=[5 * ps, 2 * ps, 3 * ps - 1]
+                )
+                src = _make_source(allocator, req_to_token, ps)
+                self.assertTrue(src.is_translating)
+                width = 6
+                view = src.build_index_table(
+                    req_pool_indices=rows, seq_lens=seq_lens, max_pages=width
+                )
+                self.assertTrue(view.is_translated)
+                self.assertEqual(view.entry_page_size, ps)
+                self.assertTrue(
+                    torch.equal(view.row_ids, torch.arange(3, dtype=torch.int64))
+                )
+                want_full = _reference_table(
+                    req_to_token,
+                    rows,
+                    seq_lens,
+                    allocator.full_v2p_page_table,
+                    full_mult,
+                    ps,
+                    width,
+                )
+                want_swa = _reference_table(
+                    req_to_token,
+                    rows,
+                    seq_lens,
+                    allocator.swa_v2p_page_table,
+                    swa_mult,
+                    ps,
+                    width,
+                )
+                self.assertTrue(
+                    torch.equal(view.ids, want_full),
+                    f"full read table off-formula (ps={ps}, mult={full_mult})",
+                )
+                self.assertTrue(
+                    torch.equal(view.sliding_window_ids, want_swa),
+                    f"swa read table off-formula (ps={ps}, mult={swa_mult})",
+                )
+
+    def test_sink_routing(self):
+        """Dead lanes (seq_len 0), -1 slots inside the live prefix, and
+        tombstoned v2p pages must ALL read entry 0 -- one wild entry is a
+        captured-graph OOB read at replay."""
+        ps = 4
+        allocator = _build_composite(ps)
+        req_to_token, rows, seq_lens = _alloc_and_fill(
+            allocator, ps, lens=[3 * ps, 2 * ps, ps]
+        )
+        seq_lens[1] = 0  # dead lane
+        req_to_token[0, ps] = -1  # unwritten slot inside the live prefix
+        # Tombstone row 2's first page on BOTH sides.
+        tomb_page = int(req_to_token[2, 0]) // ps
+        allocator.full_v2p_page_table[tomb_page] = -1
+        allocator.swa_v2p_page_table[tomb_page] = -1
+        src = _make_source(allocator, req_to_token, ps)
+        view = src.build_index_table(
+            req_pool_indices=rows, seq_lens=seq_lens, max_pages=4
+        )
+        for table in (view.ids, view.sliding_window_ids):
+            self.assertTrue(bool((table >= 0).all()))
+            self.assertTrue(bool((table[1] == 0).all()), "dead lane not sunk")
+            self.assertEqual(int(table[0, 1]), 0, "-1 slot not sunk")
+            self.assertEqual(int(table[2, 0]), 0, "tombstone not sunk")
+
+
+class TestBuildInto(unittest.TestCase):
+    """fill_read_table fills a backend-owned padded block table's live prefix with
+    FULL-side read-table entries -- the trtllm_mla / flashmla consumption route
+    (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,
+        lanes past each row's live pages keep the backend's -1 sentinel
+        (prefix-only -- a tail write scatters the trtllm sentinel contract),
+        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
+        lens = [5, 2 * ps + 1, 1]
+        req_to_token, rows, seq_lens = _alloc_and_fill(allocator, ps, lens=lens)
+        src = _make_source(allocator, req_to_token, ps)
+        self.assertTrue(src.is_translating)
+
+        width_pages = req_to_token.shape[1] // ps + 3  # wider than the span
+        out = torch.full((len(lens), width_pages), -1, dtype=torch.int32)
+        src.fill_read_table(out=out, req_pool_indices=rows, seq_lens=seq_lens)
+
+        want = _reference_table(
+            req_to_token,
+            rows,
+            seq_lens,
+            allocator.full_v2p_page_table,
+            full_mult,
+            ps,
+            width_pages,
+        )
+        for b, n in enumerate(lens):
+            n_pages = -(-n // ps)
+            self.assertTrue(
+                torch.equal(out[b, :n_pages], want[b, :n_pages]),
+                f"row {b} live prefix off-formula",
+            )
+            self.assertTrue(
+                bool((out[b, n_pages:] == -1).all()),
+                f"row {b} tail sentinel clobbered",
+            )
+
+    def test_passthrough_source_refuses(self):
+        """Callers dispatch on `enabled`; a passthrough source has 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(),
+            token_to_kv_pool=SimpleNamespace(),
+            page_size=1,
+            device=_DEV,
+        )
+        with self.assertRaises(AssertionError):
+            src.fill_read_table(
+                out=torch.zeros((1, 4), dtype=torch.int32),
+                req_pool_indices=torch.tensor([0]),
+                seq_lens=torch.tensor([1]),
+            )
+
+
+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,
+    one req_to_token) while owning a SEPARATE KV buffer sized to the
+    allocator's SLOT count. Probing the allocator alone reports "unified" for
+    that runner, so its indices would be mapped into the composite's
+    kernel-facing space (kernel-facing ids up to num_pages * multiplier) and then used
+    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`
+        holding for a REAL target bundle. If a factory ever returned a pool
+        the allocator does not hold, the guard would silently disable the
+        unified path for EVERY model -- so pin it against the real factory
+        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(
+            device="cpu",
+            kv_cache_dtype=torch.float16,
+            head_num=2,
+            head_dim=8,
+            v_head_dim=8,
+            swa_head_num=2,
+            swa_head_dim=8,
+            swa_v_head_dim=8,
+            page_size=1,
+            start_layer=0,
+            end_layer=4,
+            swa_attention_layer_ids=[1, 3],
+            full_attention_layer_ids=[0, 2],
+            full_max_total_num_tokens=64,
+            swa_max_total_num_tokens=32,
+            enable_memory_saver=False,
+            need_sort=False,
+        )
+        self.assertIs(
+            bundle.token_to_kv_pool_allocator.get_kvcache(),
+            bundle.token_to_kv_pool,
+        )
+        src = KVIndexTranslator(
+            req_to_token=torch.zeros((2, 8), dtype=torch.int32, device=_DEV),
+            token_to_kv_pool_allocator=bundle.token_to_kv_pool_allocator,
+            token_to_kv_pool=bundle.token_to_kv_pool,
+            page_size=1,
+            device=_DEV,
+        )
+        self.assertTrue(src.is_translating)
+
+    def test_runner_with_its_own_pool_is_disabled(self):
+        """Same allocator, different pool: must stay disabled."""
+        alloc = _build_composite(ps=1)
+        req_to_token = torch.zeros((2, 8), dtype=torch.int32, device=_DEV)
+        own_pool = SimpleNamespace()  # a separate buffer, not the composite's
+        src = KVIndexTranslator(
+            req_to_token=req_to_token,
+            token_to_kv_pool_allocator=alloc,
+            token_to_kv_pool=own_pool,
+            page_size=1,
+            device=_DEV,
+        )
+        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
+        the read table -- they index its own pool directly. A translate here is
+        the out-of-bounds bug the ownership 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(
+            req_to_token=req_to_token,
+            token_to_kv_pool_allocator=alloc,
+            token_to_kv_pool=SimpleNamespace(),
+            page_size=1,
+            device=_DEV,
+        )
+        rows = torch.tensor([1, 0], dtype=torch.int32, device=_DEV)
+        view = src.build_index_table(
+            req_pool_indices=rows,
+            seq_lens=torch.tensor([3, 2], dtype=torch.int32, device=_DEV),
+        )
+        # Read table: the EXACT objects a static-pool backend reads today.
+        self.assertIs(view.ids, req_to_token)
+        self.assertIs(view.row_ids, rows)
+        self.assertFalse(view.is_translated)
+        # And the token-level surface is the identity, same guard.
+        t = torch.tensor([5, 6], dtype=torch.int64, device=_DEV)
+        self.assertIs(src.translate_full_attn_ids(t), t)
+
+
+class TestCaptureContract(unittest.TestCase):
+    def test_caller_owned_table_is_returned_whole_and_filled_prefix_only(self):
+        ps = 4
+        allocator = _build_composite(ps)
+        req_to_token = torch.full((4, 16 * ps), -1, dtype=torch.int64)
+        v = allocator.alloc(2 * ps)
+        req_to_token[1, : 2 * ps] = v
+        src = _make_source(allocator, req_to_token, ps)
+
+        tables = src.make_capture_tables(max_bs=4, max_context_len=8 * ps)
+        cap, cap_swa = tables.full, tables.sliding_window
+        self.assertTrue(bool((cap == 0).all()), "read tables must start zeroed")
+        self.assertIsNotNone(cap_swa, "the SWA composite has a second id space")
+
+        # Poison everything, then refresh a 1-row batch: ONLY its live prefix
+        # may change -- stale tails and other rows are the fa3 contract.
+        cap.fill_(7)
+        cap_swa.fill_(7)
+        view = src.build_index_table(
+            req_pool_indices=torch.tensor([1]),
+            seq_lens=torch.tensor([2 * ps]),
+            into=tables,
+        )
+        self.assertIs(view.ids, cap, "the caller's table comes back WHOLE")
+        want = allocator.full_v2p_page_table[req_to_token[1, ::ps][:2] // ps] * (
+            2 * _FULL_L
+        )
+        self.assertTrue(torch.equal(cap[0, :2], want.to(torch.int32)))
+        self.assertTrue(bool((cap[0, 2:] == 7).all()), "stale tail was cleared")
+        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
+        builds at different batch sizes hand back slices of ONE buffer. A
+        per-build `torch.arange` would be correct but would spend an allocation
+        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)
+        self.assertTrue(src.is_translating)
+        first = src.build_index_table(
+            req_pool_indices=rows[:2], seq_lens=seq_lens[:2], max_pages=4
+        )
+        second = src.build_index_table(
+            req_pool_indices=rows, seq_lens=seq_lens, max_pages=4
+        )
+        self.assertEqual(first.row_ids.data_ptr(), second.row_ids.data_ptr())
+        self.assertTrue(torch.equal(first.row_ids, torch.arange(2, device=_DEV)))
+        self.assertTrue(torch.equal(second.row_ids, torch.arange(3, device=_DEV)))
+        # Sized to bound any batch the request pool can hold.
+        self.assertGreaterEqual(src._rows.numel(), req_to_token.shape[0])
+
+
+class _FakeForwardBatch:
+    """Weakref-able stand-in (SimpleNamespace is not) carrying the fields
+    `index_table_for_batch` and `rebind_write_loc` read. `seq_lens_sum`
+    defaults to the real sum: it is the signal that the CPU mirror is live,
+    and a real ForwardBatch always carries it (None only when gpu_only)."""
+
+    def __init__(
+        self,
+        *,
+        req_pool_indices=None,
+        seq_lens=None,
+        seq_lens_cpu=None,
+        out_cache_loc=None,
+        seq_lens_sum=-1,
+    ):
+        self.req_pool_indices = req_pool_indices
+        self.seq_lens = seq_lens
+        self.seq_lens_cpu = seq_lens_cpu
+        self.out_cache_loc = out_cache_loc
+        self.seq_lens_sum = (
+            (None if seq_lens is None else int(seq_lens.sum()))
+            if seq_lens_sum == -1
+            else seq_lens_sum
+        )
+
+
+class TestViewMemo(unittest.TestCase):
+    """The eager view is memoized ON THE SOURCE in a single slot keyed by
+    batch identity -- per-batch state stays out of the ForwardBatch (it does
+    not scale with the number of id spaces), and one metadata build's many
+    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)
+        fb = _FakeForwardBatch(
+            req_pool_indices=rows,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens,
+        )
+        return fb, req_to_token
+
+    def test_same_batch_returns_the_memoized_view(self):
+        ps = 1
+        allocator = _build_composite(ps)
+        fb, req_to_token = self._fb(allocator, ps, lens=[3, 2])
+        src = _make_source(allocator, req_to_token, ps)
+        v1 = src.index_table_for_batch(fb)
+        v2 = src.index_table_for_batch(fb)
+        self.assertIs(v1, v2)
+
+    def test_next_batch_replaces_the_single_slot(self):
+        ps = 1
+        allocator = _build_composite(ps)
+        fb1, req_to_token = self._fb(allocator, ps, lens=[3, 2])
+        src = _make_source(allocator, req_to_token, ps)
+        v1 = src.index_table_for_batch(fb1)
+        fb2 = _FakeForwardBatch(
+            req_pool_indices=fb1.req_pool_indices,
+            seq_lens=fb1.seq_lens,
+            seq_lens_cpu=fb1.seq_lens_cpu,
+        )
+        v2 = src.index_table_for_batch(fb2)
+        self.assertIsNot(v1, v2)
+        # Single slot: fb1 no longer matches and rebuilds.
+        v1_again = src.index_table_for_batch(fb1)
+        self.assertIsNot(v1_again, v1)
+
+    def test_dead_batch_never_matches(self):
+        """A garbage-collected batch's slot must not serve a later batch: the
+        weakref key goes dead and the build runs fresh."""
+        import gc
+
+        ps = 1
+        allocator = _build_composite(ps)
+        fb1, req_to_token = self._fb(allocator, ps, lens=[3, 2])
+        src = _make_source(allocator, req_to_token, ps)
+        v1 = src.index_table_for_batch(fb1)
+        del fb1
+        gc.collect()
+        fb2, _ = self._fb(allocator, ps, lens=[2])
+        v2 = src.index_table_for_batch(fb2)
+        self.assertIsNot(v2, v1)
+        self.assertEqual(v2.ids.shape[0], 1)
+
+
+class TestWriteLoc(unittest.TestCase):
+    """The two-phase write contract: phase 1 (`rebind_write_loc`) rebinds the
+    full side once at ForwardBatch construction; phase 2 derives the
+    sliding-window write loc on demand, POINTWISE from the full-side
+    values. Value-based derivation is the property under test: pads, slices,
+    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)
+        req_to_token, rows, seq_lens = _alloc_and_fill(allocator, ps, lens=[max(n, 1)])
+        src = _make_source(allocator, req_to_token, ps)
+        virt = allocator.alloc(-(-n // ps) * ps)[:n]
+        want_full = allocator.translate_kv_loc_for_kernel(virt)
+        want_swa = allocator.translate_loc_from_full_to_swa(virt)
+        return src, allocator, rows, seq_lens, virt, want_full, want_swa
+
+    def _field(self, src, rows, seq_lens, kernel_loc):
+        return src.sliding_window_write_loc_for(kernel_loc)
+
+    def test_rebind_translates_full_side_only(self):
+        for ps in (1, 4):
+            src, _, _, _, virt, want_full, _ = self._built(ps=ps, n=3 * ps)
+            keep = virt.clone()
+            fb = _FakeForwardBatch(out_cache_loc=virt)
+            src.rebind_write_loc(fb)
+            # Full side: rebound to a FRESH kernel-facing tensor; the
+            # ScheduleBatch's aliased virtual tensor is untouched.
+            self.assertIsNot(fb.out_cache_loc, virt)
+            self.assertTrue(torch.equal(fb.out_cache_loc, want_full))
+            self.assertTrue(torch.equal(virt, keep))
+
+    def test_swa_write_loc_round_trips_from_dense(self):
+        """The derived property behind phase 2: for any virtual run t,
+        deriving from the dense full-side values must equal the direct
+        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
+            )
+            got = self._field(src, rows, seq_lens, want_full)
+            self.assertTrue(torch.equal(got, want_swa))
+
+    def test_pad_lanes_derive_to_sink(self):
+        """The DP pad appends zeros; dense 0 is the reserved padding slot in
+        every id space, so pad lanes must derive to swa slot 0 with 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
+        got = self._field(src, rows, seq_lens, want_full[:1])
+        self.assertEqual(int(got[0]), 0)
+
+    def test_static_swa_pool_derives_via_pool_translate(self):
+        """Static SWA pools: the field is the pool's own legacy full->swa
+        translate, computed at the same build; the rebind stays a no-op."""
+        pool = SWAKVPool.__new__(SWAKVPool)
+        pool.full_to_swa_index_mapping = torch.arange(10, dtype=torch.int64)
+        pool.translate_loc_from_full_to_swa = lambda t: t + 100
+        src = KVIndexTranslator(
+            req_to_token=torch.zeros((2, 4), dtype=torch.int64),
+            token_to_kv_pool_allocator=SimpleNamespace(),
+            token_to_kv_pool=pool,
+            page_size=1,
+            device=_DEV,
+        )
+        loc = torch.tensor([5, 6], dtype=torch.int64)
+        fb = _FakeForwardBatch(out_cache_loc=loc)
+        src.rebind_write_loc(fb)
+        self.assertIs(fb.out_cache_loc, loc, "disabled rebind must be a no-op")
+        self.assertTrue(torch.equal(src.sliding_window_write_loc_for(loc), loc + 100))
+
+    def test_no_loc_or_no_swa_side_yields_none(self):
+        # Unified swa composite, but there is no write loc this forward.
+        src, _, rows, seq_lens, _, _, _ = self._built(n=2)
+        self.assertIsNone(src.sliding_window_write_loc_for(None))
+        # Passthrough on a non-SWA pool: a loc is given, but there is no swa
+        # id space to derive into.
+        plain = KVIndexTranslator(
+            req_to_token=torch.zeros((2, 4), dtype=torch.int64),
+            token_to_kv_pool_allocator=SimpleNamespace(),
+            token_to_kv_pool=SimpleNamespace(),
+            page_size=1,
+            device=_DEV,
+        )
+        self.assertIsNone(
+            plain.sliding_window_write_loc_for(torch.tensor([3], dtype=torch.int64))
+        )
+
+    def test_rebind_retires_the_view_memo(self):
+        ps = 1
+        allocator = _build_composite(ps)
+        req_to_token, rows, seq_lens = _alloc_and_fill(allocator, ps, lens=[3, 2])
+        src = _make_source(allocator, req_to_token, ps)
+        fb = _FakeForwardBatch(
+            req_pool_indices=rows, seq_lens=seq_lens, seq_lens_cpu=seq_lens
+        )
+        v1 = src.index_table_for_batch(fb)
+        src.rebind_write_loc(_FakeForwardBatch(out_cache_loc=None))
+        v2 = src.index_table_for_batch(fb)
+        self.assertIsNot(v2, v1, "rebind starts the next forward: stale views die")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -539,6 +539,31 @@ class TestMultiEndedAllocator(unittest.TestCase):
         )
         self.assertEqual(int(buf[1].item()), 0)
 
+    def test_slot_zero_sink_invariant_survives_churn(self):
+        """PINNED INVARIANT: virtual 0 <-> physical 0 (the padding sink), so
+        `translate_kv_loc(zeros) == zeros` -- after init AND after alloc/free/
+        compaction churn. The cuda-graph capture path RELIES on this: the
+        physical-loc contract replaced capture-time translate with a plain
+        copy of the zero-filled static buffer, which is only equivalent while
+        v2p[0] == 0. If an allocator change breaks this, captured stores would
+        write pad lanes to a live slot."""
+        _, full_alloc, _, full_kv, _ = self._build_pair()
+        zeros = torch.zeros(4, dtype=torch.int64)
+
+        self.assertEqual(int(full_alloc.virtual_to_physical[0].item()), 0)
+        self.assertTrue(torch.equal(full_alloc.translate_kv_loc(zeros), zeros))
+
+        # Churn: allocate, free interior (forces compaction moves), re-allocate.
+        a = self._alloc(full_alloc, full_kv, 6)
+        b = self._alloc(full_alloc, full_kv, 6)
+        self._free(full_alloc, full_kv, a)
+        c = self._alloc(full_alloc, full_kv, 4)
+        self._free(full_alloc, full_kv, b)
+        self._free(full_alloc, full_kv, c)
+
+        self.assertEqual(int(full_alloc.virtual_to_physical[0].item()), 0)
+        self.assertTrue(torch.equal(full_alloc.translate_kv_loc(zeros), zeros))
+
 
 # ---------------------------------------------------------------------------
 # Shared SWA composite — unit tests

test/registered/unit/mem_cache/test_unified_mha_views.py

@@ -483,8 +483,9 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
 
 
 class TestFactoryDenseViews(unittest.TestCase):
-    """The real SWA factory builds both sub-pools and wires the matching
-    kernel-facing multipliers into the composite allocator."""
+    """The real SWA factory builds dense sub-pools and wires the matching
+    kernel-facing multipliers into the composite allocator. End-to-end over
+    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
@@ -526,6 +527,42 @@ class TestFactoryDenseViews(unittest.TestCase):
         self.assertEqual(b.token_to_kv_pool.swa_kv_pool.k_buffer[0].dim(), 3)
         self.assertGreater(pool.view_tail_pad_bytes, 0)
 
+    def test_rebind_emits_dense_full_and_build_derives_swa(self):
+        """End-to-end over the real factory: rebind_write_loc rebinds
+        out_cache_loc to FULL-kernel-facing ids (phase 1), and the per-batch build
+        derives the SWA-DENSE 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()
+        alloc = b.token_to_kv_pool_allocator
+        v = alloc.alloc(4)
+        self.assertIsNotNone(v)
+        expected_full = alloc.full_v2p_page_table[v] * self.FULL_MULT  # ps=1
+        expected_swa = alloc.swa_v2p_page_table[v] * self.SWA_MULT
+
+        class _FB:
+            pass
+
+        fb = _FB()
+        fb.out_cache_loc = v.clone()
+        source = KVIndexTranslator(
+            req_to_token=torch.zeros((2, 8), dtype=torch.int64),
+            token_to_kv_pool_allocator=alloc,
+            token_to_kv_pool=b.token_to_kv_pool,
+            page_size=1,
+            device="cpu",
+        )
+        self.assertTrue(source.is_translating)
+        source.rebind_write_loc(fb)
+        self.assertTrue(torch.equal(fb.out_cache_loc, expected_full))
+        self.assertTrue(
+            torch.equal(
+                source.sliding_window_write_loc_for(fb.out_cache_loc), expected_swa
+            )
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py

@@ -0,0 +1,181 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""ForwardBatch construction wires the unified write-loc rebind.
+
+`init_new` must call `kv_index_translator.rebind_write_loc`: a construction
+path that skips it ships VIRTUAL write ids to the kernels, a silent
+wrong-slot store. Also runs the REAL `_pad_inputs_to_size` against a live
+translator, since pad lanes are zeros and zeros must derive to the slot-0
+sink. Sliding-window semantics are pinned in test_kv_index_translator.py.
+
+    python -m pytest test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py -v
+"""
+
+import ast
+import inspect
+import textwrap
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+_DEV = "cpu"
+
+
+def _make_fb(out_cache_loc, **kw):
+    """Minimal ForwardBatch with only the required core fields."""
+    n = 0 if out_cache_loc is None else out_cache_loc.shape[0]
+    defaults = dict(
+        forward_mode=ForwardMode.DECODE,
+        batch_size=max(n, 1),
+        input_ids=torch.zeros(max(n, 1), dtype=torch.int64),
+        req_pool_indices=torch.zeros(max(n, 1), dtype=torch.int64),
+        seq_lens=torch.ones(max(n, 1), dtype=torch.int64),
+        out_cache_loc=out_cache_loc,
+        seq_lens_sum=max(n, 1),
+    )
+    defaults.update(kw)
+    return ForwardBatch(**defaults)
+
+
+def _armed_source(v2p, swa_map):
+    """A KVIndexTranslator hand-armed with fake translates: this file pins the
+    ForwardBatch-side wiring, not the composite's formulas (those are pinned
+    in test_kv_index_translator.py over the real allocator)."""
+    src = KVIndexTranslator(
+        req_to_token=torch.zeros((1, 4), dtype=torch.int64),
+        token_to_kv_pool_allocator=SimpleNamespace(),
+        token_to_kv_pool=SimpleNamespace(),
+        page_size=1,
+        device=_DEV,
+    )
+    src.is_translating = True
+    src._translate_full = lambda t, out=None: v2p[t.to(torch.int64)]
+    # Phase 2 derives from DENSE values through p2v + the swa v2p; arm the
+    # inverse of the fake v2p (ps=1, both multipliers 1: dense == physical,
+    # and the expected swa loc for virtual t is swa_map[t]).
+    p2v = torch.zeros(int(v2p.max()) + 1, dtype=torch.int64)
+    p2v[v2p] = torch.arange(v2p.numel(), dtype=torch.int64)
+    src._full_p2v_table = p2v
+    src._swa_v2p_table = swa_map
+    src._full_page_multiplier = 1
+    src._swa_page_multiplier = 1
+    return src
+
+
+def _call_names(func) -> list:
+    """Dotted call targets appearing in `func`'s body, e.g.
+    'model_runner.kv_index_translator.rebind_write_loc'."""
+    tree = ast.parse(textwrap.dedent(inspect.getsource(func)))
+    names = []
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Call):
+            parts = []
+            cur = node.func
+            while isinstance(cur, ast.Attribute):
+                parts.append(cur.attr)
+                cur = cur.value
+            if isinstance(cur, ast.Name):
+                parts.append(cur.id)
+            names.append(".".join(reversed(parts)))
+    return names
+
+
+class TestForwardBatchWiring(CustomTestCase):
+    """Critical-path bookkeeping: the construction-time call sites."""
+
+    def test_init_new_calls_the_rebind(self):
+        self.assertIn(
+            "model_runner.kv_index_translator.rebind_write_loc",
+            _call_names(ForwardBatch.init_new.__func__),
+            "init_new must rebind the write loc through the source; a batch "
+            "built without it ships virtual ids to the kernels",
+        )
+
+
+class TestPadComposesWithDerivation(CustomTestCase):
+    def _fake_runner_for_pad(self, src):
+        return SimpleNamespace(
+            attn_backend=SimpleNamespace(get_cuda_graph_seq_len_fill_value=lambda: 0),
+            kv_index_translator=src,
+        )
+
+    def test_pad_lanes_derive_to_sink_and_slices_stay_pointwise(self):
+        """The REAL `_pad_inputs_to_size` composes with phase 2: pad lanes are
+        zeros, zeros derive to the slot-0 sink, and any slice of the padded
+        tensor (the TBO-child shape) derives pointwise -- no handover call
+        exists for the pad to make."""
+        n, padded = 3, 6
+        v2p = torch.arange(64, dtype=torch.int64) * 3
+        swa_map = torch.arange(64, dtype=torch.int64) * 5
+        src = _armed_source(v2p, swa_map)
+        virt = torch.tensor([11, 12, 13], dtype=torch.int64)
+        fb = _make_fb(virt.clone())
+        fb.positions = torch.arange(n, dtype=torch.int64)
+        fb.lora_ids = [None] * fb.batch_size
+        src.rebind_write_loc(fb)
+        self.assertTrue(torch.equal(fb.out_cache_loc, v2p[virt]))
+
+        fb._pad_inputs_to_size(self._fake_runner_for_pad(src), padded, fb.batch_size)
+
+        self.assertEqual(fb.out_cache_loc.shape[0], padded)
+        # Padded tail lanes go to slot 0 -- the reserved dummy-write sink.
+        self.assertTrue(bool((fb.out_cache_loc[n:] == 0).all()))
+        loc = src._swa_write_loc_unified(fb.out_cache_loc)
+        self.assertTrue(torch.equal(loc[:n], swa_map[virt]))
+        self.assertTrue(bool((loc[n:] == 0).all()))
+        self.assertEqual(loc.dtype, torch.int64)
+        # The TBO-child shape: a slice of the PADDED tensor derives pointwise.
+        sub = src._swa_write_loc_unified(fb.out_cache_loc[1:5])
+        self.assertTrue(torch.equal(sub, loc[1:5]))
+
+    def test_the_probe_separates_kernel_facing_from_virtual_ids(self):
+        """A skipped rebind is the failure mode this contract has no other
+        guard against: virtual ids stay inside the OOB probe's bounds (they are
+        `blocks_per_page` times SMALLER than a kernel-facing id), so the store lands on
+        the wrong slots and only the output is wrong. The kernel-facing probe
+        is what separates them -- the in-page offset of a kernel-facing id is always
+        below page_size, and a virtual id's is not unless it happens to fall in
+        the first block."""
+        for page_size, blocks in ((1, 8), (4, 6)):
+            with self.subTest(page_size=page_size, blocks=blocks):
+                stride = page_size * blocks
+                virt = torch.arange(1, 2 * stride, dtype=torch.int64)
+                dense = (virt // page_size) * stride + virt % page_size
+                in_space = dense % stride < page_size
+                self.assertTrue(bool(in_space.all()), "kernel-facing ids must pass")
+                # Virtual ids pass only in the first block; that is why the
+                # probe needs a batch, not one id, to be conclusive.
+                caught = ~(virt % stride < page_size)
+                self.assertTrue(bool(caught.any()), "virtual ids must be caught")
+
+    def test_empty_loc_rebinds_to_empty(self):
+        src = _armed_source(
+            torch.arange(8, dtype=torch.int64), torch.arange(8, dtype=torch.int64)
+        )
+        fb = _make_fb(torch.empty(0, dtype=torch.int64))
+        src.rebind_write_loc(fb)
+        self.assertEqual(fb.out_cache_loc.numel(), 0)
+        self.assertEqual(src._swa_write_loc_unified(fb.out_cache_loc).numel(), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py

@@ -0,0 +1,92 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""`--enable-unified-memory` disables PREFILL cuda-graph capture.
+
+BUG REGRESSION. Only decode capture is wired: the prefill graph runner builds
+its ForwardBatch directly, so it never runs the unified pool's write-loc
+rebind (rebind_write_loc) and the captured batch holds VIRTUAL ids -- the
+captured store would silently write wrong slots.
+
+The old gate only rejected `TC_PIECEWISE`, but the generic prefill default is
+`BREAKABLE` -- so the DEFAULT unified invocation was broken; it only ever
+worked when `--disable-piecewise-cuda-graph` (a deprecated alias for
+`--cuda-graph-backend-prefill=disabled`) happened to be passed.
+
+Pinned: the default is auto-disabled with a warning (unified boots out of the
+box), an EXPLICIT prefill backend still raises (never silently override a
+user's stated intent), and decode capture is untouched either way.
+
+    python -m pytest test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py -v
+"""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.arg_groups.kv_cache_hook import handle_unified_memory_pool
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _run_handler(*, prefill_backend, explicit):
+    """Run just `handle_unified_memory_pool` over a minimal stand-in."""
+    sa = ServerArgs.__new__(ServerArgs)
+    cg = SimpleNamespace(
+        prefill=SimpleNamespace(backend=prefill_backend),
+        decode=SimpleNamespace(backend=Backend.FULL),
+    )
+    for name, value in {
+        "enable_unified_memory": True,
+        "disaggregation_mode": "null",
+        "speculative_algorithm": None,
+        "speculative_eagle_topk": None,
+        "enable_hierarchical_cache": False,
+        "enable_lmcache": False,
+        "dcp_size": 1,
+        "cuda_graph_config": cg,
+        "cuda_graph_backend_prefill": prefill_backend if explicit else None,
+    }.items():
+        object.__setattr__(sa, name, value)
+    handle_unified_memory_pool(sa)
+    return cg
+
+
+class TestUnifiedPrefillCudaGraphGate(unittest.TestCase):
+    def test_default_prefill_capture_is_auto_disabled(self):
+        """The generic default (BREAKABLE) must be turned off, not crash the
+        server 30 seconds later inside graph capture."""
+        for backend in (Backend.BREAKABLE, Backend.FULL, Backend.TC_PIECEWISE):
+            cg = _run_handler(prefill_backend=backend, explicit=False)
+            self.assertEqual(cg.prefill.backend, Backend.DISABLED)
+            # Decode capture is the wired path and must survive untouched.
+            self.assertEqual(cg.decode.backend, Backend.FULL)
+
+    def test_explicit_prefill_backend_is_refused(self):
+        """A user who explicitly asked for prefill graphs gets a clear error,
+        not a silent override of their stated intent."""
+        for backend in (Backend.BREAKABLE, Backend.FULL, Backend.TC_PIECEWISE):
+            with self.assertRaises(ValueError) as ctx:
+                _run_handler(prefill_backend=backend, explicit=True)
+            self.assertIn("prefill capture is not wired", str(ctx.exception))
+
+    def test_already_disabled_is_a_no_op(self):
+        cg = _run_handler(prefill_backend=Backend.DISABLED, explicit=True)
+        self.assertEqual(cg.prefill.backend, Backend.DISABLED)
+        self.assertEqual(cg.decode.backend, Backend.FULL)
+
+
+if __name__ == "__main__":
+    unittest.main()