[Perf] Unified memory: close the DCP decode gap on Blackwell (#37926)

test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py

@@ -432,6 +432,89 @@ class TestFusedFp8WriteGate(CustomTestCase):
                 )
 
 
+@unittest.skipUnless(torch.cuda.is_available(), "the fused translate is Triton")
+class TestFusedWriteLocTranslateCuda(CustomTestCase):
+    """The fused write-loc translate must agree with its own CPU branch.
+
+    The two implementations must not drift: Triton truncates division toward
+    zero where torch floors it, so a negative loc and a tombstoned v2p row are
+    where a divergence would appear -- and the CPU branch is all the CPU suites
+    ever exercise.
+    """
+
+    def test_cuda_matches_the_cpu_branch(self):
+        from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
+
+        for page_size in (1, 64):
+            span = page_size * 4
+            loc = torch.tensor(
+                [-1, 0, 1, page_size, span, span + 1, 2 * span + 3, 5 * page_size],
+                dtype=torch.int64,
+            )
+            # Size the table past the highest page any loc can name, then
+            # scramble it and tombstone one row (-1), so neither a dropped
+            # clamp nor a skipped gather can coincide with the right answer.
+            num_pages = int(loc.max()) // page_size + 2
+            v2p = torch.tensor(
+                [(5 * i + 2) % num_pages for i in range(num_pages)] + [-1],
+                dtype=torch.int64,
+            )
+            v2p[1] = -1
+            for dcp_size, dcp_rank in ((1, 0), (2, 1), (4, 2)):
+                kw = dict(
+                    page_size=page_size,
+                    stride=page_size * 3,
+                    dcp_size=dcp_size,
+                    dcp_rank=dcp_rank,
+                )
+                cpu = write_loc_to_kernel_ids(loc=loc, v2p=v2p, **kw)
+                gpu = write_loc_to_kernel_ids(loc=loc.cuda(), v2p=v2p.cuda(), **kw)
+                self.assertEqual(
+                    cpu.tolist(),
+                    gpu.cpu().tolist(),
+                    f"ps={page_size} dcp_size={dcp_size} rank={dcp_rank}",
+                )
+
+    def test_wide_out_clears_the_stale_tail(self):
+        """`out_width` past the batch must zero the tail in the same launch.
+
+        This is what lets a backend hand in its whole capture-stable buffer:
+        a shorter replay leaves stale kernel-facing ids past the batch, and the
+        captured write kernel consumes the full buffer, so an uncleared tail
+        scatters pad rows into live KV pages.
+        """
+        from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
+
+        v2p = torch.tensor([2, 5, 1, 3, 4], dtype=torch.int64, device="cuda")
+        loc = torch.tensor([0, 64, 128], dtype=torch.int64, device="cuda")
+        width = 8
+        # Poison the whole buffer so an unwritten or uncleared cell is visible.
+        buf = torch.full((width,), -999, dtype=torch.int64, device="cuda")
+        write_loc_to_kernel_ids(
+            loc=loc, v2p=v2p, page_size=64, stride=64 * 2, out=buf, out_width=width
+        )
+        self.assertEqual(buf[:3].tolist(), [2 * 128, 5 * 128, 1 * 128])
+        self.assertEqual(buf[3:].tolist(), [0] * (width - 3))
+
+        # And it must agree with the narrow call on the live prefix.
+        narrow = write_loc_to_kernel_ids(loc=loc, v2p=v2p, page_size=64, stride=64 * 2)
+        self.assertEqual(narrow.tolist(), buf[:3].tolist())
+
+    def test_out_is_written_in_place(self):
+        # The captured decode path hands in a capture-stable buffer; rebinding
+        # instead of filling it would leave the graph on a stale pointer.
+        from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
+
+        v2p = torch.tensor([2, 5, -1, 3], dtype=torch.int64, device="cuda")
+        loc = torch.tensor([0, 64, 128, 192], dtype=torch.int64, device="cuda")
+        dst = torch.full_like(loc, -7)
+        ret = write_loc_to_kernel_ids(
+            loc=loc, v2p=v2p, page_size=64, stride=64 * 2, out=dst
+        )
+        self.assertIs(ret, dst)
+        self.assertEqual(dst.tolist(), [2 * 128, 5 * 128, 0, 3 * 128])
+
+
 class TestDcpDecodeLayout(CustomTestCase):
     """Rank-local length math the decode page table above is built from."""
 

test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py

@@ -49,11 +49,15 @@ def _reference_chain(
     out_buf: torch.Tensor,
     valid_bs: int,
     total_bs: int,
+    v2p: torch.Tensor = None,
 ) -> None:
     """Replicates the _replay_metadata reference ops, in order, in place."""
     req_pool_indices[valid_bs:total_bs] = 0
     mamba_indices = mapping[req_pool_indices[:total_bs]]
-    # static pool: _translate_mamba_indices is the identity
+    if v2p is not None:
+        # Unified pool: virtual->physical slot translate, and it runs BEFORE
+        # the padding sentinel -- captured kernels read physical ids.
+        mamba_indices = v2p[mamba_indices].to(torch.int32)
     mamba_indices[valid_bs:] = -1
     out_buf[: len(mamba_indices)].copy_(mamba_indices)
 
@@ -139,6 +143,79 @@ class TestFusedReplayStateIndices(CustomTestCase):
                 f"{name} guard tail clobbered ({case}): {buf[total_bs:].tolist()}",
             )
 
+    def _run_v2p_case(self, total_bs: int, num_padding: int, seed: int) -> None:
+        """Same equivalence, with the unified pool's virtual slot ids.
+
+        The mapping yields VIRTUAL slots there and the kernel folds the v2p
+        gather in, so the two must still agree element for element.
+        """
+        device = torch.device("cuda")
+        gen = torch.Generator(device="cpu").manual_seed(seed + 1000)
+        valid_bs = total_bs - num_padding
+
+        req_pool = torch.randint(
+            0, _REQ_POOL_SIZE, (total_bs + _GUARD,), generator=gen, dtype=torch.int64
+        )
+        req_pool[total_bs:] = _GUARD_SENTINEL
+        mapping = torch.randint(
+            0, _MAMBA_POOL_SIZE, (_REQ_POOL_SIZE,), generator=gen, dtype=torch.int32
+        )
+        # Scrambled table with a tombstone, so a skipped gather cannot pass.
+        v2p = torch.randperm(_MAMBA_POOL_SIZE + 1, generator=gen).to(torch.int64)
+        v2p[7] = -1
+        out = torch.full((total_bs + _GUARD,), _OUT_POISON, dtype=torch.int32)
+
+        req_ref, req_fused = req_pool.clone().to(device), req_pool.clone().to(device)
+        out_ref, out_fused = out.clone().to(device), out.clone().to(device)
+        mapping_d, v2p_d = mapping.to(device), v2p.to(device)
+
+        _reference_chain(
+            req_pool_indices=req_ref,
+            mapping=mapping_d,
+            out_buf=out_ref,
+            valid_bs=valid_bs,
+            total_bs=total_bs,
+            v2p=v2p_d,
+        )
+        returned = fused_replay_state_indices(
+            req_pool_indices=req_fused,
+            mamba_index_mapping=mapping_d,
+            out_state_indices=out_fused,
+            valid_bs=valid_bs,
+            total_bs=total_bs,
+            v2p=v2p_d,
+        )
+        case = f"v2p {total_bs=} {num_padding=} {seed=}"
+        self.assertTrue(
+            torch.equal(out_ref[:total_bs], out_fused[:total_bs]),
+            f"state indices mismatch ({case}):\n"
+            f"  ref   {out_ref[:total_bs].tolist()}\n"
+            f"  fused {out_fused[:total_bs].tolist()}",
+        )
+        self.assertTrue(torch.equal(returned, out_fused[:total_bs]), case)
+        self.assertTrue(
+            torch.equal(req_pool_ref_tail := req_ref[total_bs:], req_fused[total_bs:]),
+            f"guard tail diverged ({case}): {req_pool_ref_tail.tolist()}",
+        )
+        self.assertTrue(
+            (out_fused[total_bs:] == _OUT_POISON).all(),
+            f"out guard tail clobbered ({case})",
+        )
+
+    def test_v2p_matrix(self):
+        for total_bs in (1, 7, 32, 33):
+            paddings = sorted(
+                {0, 1, total_bs // 2, total_bs - 1} & set(range(total_bs))
+            )
+            for num_padding in paddings:
+                for seed in (0, 1):
+                    with self.subTest(
+                        total_bs=total_bs, num_padding=num_padding, seed=seed
+                    ):
+                        self._run_v2p_case(
+                            total_bs=total_bs, num_padding=num_padding, seed=seed
+                        )
+
     def test_matrix(self):
         # Non-power-of-two sizes (7, 33) exercise the BS_UPPER in_range mask;
         # num_padding sweeps none / one / half / all-but-one padded rows.

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -3100,5 +3100,82 @@ class TestDcpWidening(unittest.TestCase):
                     self.assertLess(base_cost * full_entry, mamba_bytes * dcp_size)
 
 
+class TestFusedWriteLocTranslate(unittest.TestCase):
+    """`write_loc_to_kernel_ids` must equal the arithmetic it stands for.
+
+    Nothing downstream can tell the two apart except by value, so the reference
+    here is the definition rather than a recorded expectation. Triton truncates
+    division toward zero where torch floors it, so the negative-loc and
+    tombstoned-page cases are the ones that matter.
+    """
+
+    def _reference(self, loc, v2p, page_size, stride, dcp_size, dcp_rank):
+        out = []
+        for raw in loc.tolist():
+            if raw < 0 or (dcp_size > 1 and raw % dcp_size != dcp_rank):
+                out.append(0)
+                continue
+            collapsed = raw // dcp_size
+            page = collapsed // page_size
+            offset = collapsed % page_size if page_size > 1 else 0
+            out.append(max(int(v2p[page]) * stride + offset, 0))
+        return out
+
+    def _check(self, *, page_size, multiplier, dcp_size, dcp_rank, device):
+        from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
+
+        span = page_size * dcp_size
+        locs = [0, 1, span - 1, span, 2 * span + 3, 5 * span + dcp_rank, -1]
+        locs += [3 * span + dcp_rank]  # lands on the tombstoned page
+        loc = torch.tensor(locs, dtype=torch.int64, device=device)
+        # Table sized past the highest page any loc can name, scrambled, with
+        # one tombstone (-1) so a missing clamp shows up.
+        num_pages = max(locs) // (page_size * dcp_size) + 2
+        v2p = torch.tensor(
+            [(5 * i + 2) % num_pages for i in range(num_pages)] + [-1],
+            dtype=torch.int64,
+            device=device,
+        )
+        v2p[min(3, num_pages - 1)] = -1
+        stride = page_size * multiplier
+
+        got = write_loc_to_kernel_ids(
+            loc=loc,
+            v2p=v2p,
+            page_size=page_size,
+            stride=stride,
+            dcp_size=dcp_size,
+            dcp_rank=dcp_rank,
+        )
+        want = self._reference(
+            loc.cpu(), v2p.cpu(), page_size, stride, dcp_size, dcp_rank
+        )
+        self.assertEqual(got.tolist(), want, f"ps={page_size} dcp={dcp_size}")
+        # `out=` must write in place and agree (the cuda-graph-stable path).
+        dst = torch.full_like(loc, -7)
+        ret = write_loc_to_kernel_ids(
+            loc=loc,
+            v2p=v2p,
+            page_size=page_size,
+            stride=stride,
+            dcp_size=dcp_size,
+            dcp_rank=dcp_rank,
+            out=dst,
+        )
+        self.assertIs(ret, dst)
+        self.assertEqual(dst.tolist(), want)
+
+    def test_matches_reference_on_cpu(self):
+        for page_size in (1, 64):
+            for dcp_size, dcp_rank in ((1, 0), (2, 1), (4, 2)):
+                self._check(
+                    page_size=page_size,
+                    multiplier=7,
+                    dcp_size=dcp_size,
+                    dcp_rank=dcp_rank,
+                    device="cpu",
+                )
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py

@@ -55,6 +55,7 @@ def _run_handler(*, prefill_backend, explicit):
         "speculative_eagle_topk": None,
         "enable_hierarchical_cache": False,
         "enable_lmcache": False,
+        "enable_two_batch_overlap": False,
         "dcp_size": 1,
         "cuda_graph_config": cg,
         "cuda_graph_backend_prefill": prefill_backend if explicit else None,

test/registered/unit/server_args/test_unified_tbo_gate.py

@@ -0,0 +1,69 @@
+# 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` refuses `--enable-two-batch-overlap`.
+
+BUG REGRESSION. The combination launches and captures fine, then dies in the
+forward path on the first captured decode replay. Nothing else rejects the
+pair, so without this gate a running server crashes mid-serving.
+
+    python -m pytest test/registered/unit/server_args/test_unified_tbo_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(*, unified, tbo):
+    """Run just `handle_unified_memory_pool` over a minimal stand-in."""
+    sa = ServerArgs.__new__(ServerArgs)
+    for name, value in {
+        "enable_unified_memory": unified,
+        "enable_two_batch_overlap": tbo,
+        "disaggregation_mode": "null",
+        "speculative_algorithm": None,
+        "speculative_eagle_topk": None,
+        "enable_hierarchical_cache": False,
+        "enable_lmcache": False,
+        "dcp_size": 1,
+        "cuda_graph_config": SimpleNamespace(
+            prefill=SimpleNamespace(backend=Backend.DISABLED),
+            decode=SimpleNamespace(backend=Backend.FULL),
+        ),
+        "cuda_graph_backend_prefill": Backend.DISABLED,
+    }.items():
+        object.__setattr__(sa, name, value)
+    handle_unified_memory_pool(sa)
+
+
+class TestUnifiedTboGate(unittest.TestCase):
+    def test_tbo_with_unified_memory_is_refused(self):
+        with self.assertRaises(AssertionError) as ctx:
+            _run_handler(unified=True, tbo=True)
+        self.assertIn("two-batch-overlap", str(ctx.exception))
+
+    def test_gate_fires_only_on_the_pair(self):
+        """An inverted condition here would reject every unified launch."""
+        _run_handler(unified=True, tbo=False)
+        _run_handler(unified=False, tbo=True)
+
+
+if __name__ == "__main__":
+    unittest.main()