Support unified memory page-envelope transfers in PD (#39477)

Co-authored-by: yhzhuang <yhzhuang@fb.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
Co-authored-by: Yonghao Zhuang <yhzhuang@users.noreply.github.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

test/registered/unit/mem_cache/test_decode_radix_lock_ref.py

@@ -47,6 +47,9 @@ from sglang.srt.mem_cache.base_prefix_cache import (
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
 from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
 from sglang.srt.utils.common import Range
+from sglang.test.separate_buffer_allocator_double import (
+    bind_separate_buffer_capacity,
+)
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -128,6 +131,7 @@ class TestDecodeLockRefScenarios(CustomTestCase):
             server_args=SimpleNamespace(),
         )
         queue.token_to_kv_pool_allocator = MagicMock(page_size=64)
+        bind_separate_buffer_capacity(queue.token_to_kv_pool_allocator)
 
         tail_len = queue._swa_tail_len(895)
 
@@ -146,6 +150,7 @@ class TestDecodeLockRefScenarios(CustomTestCase):
         queue._need_space_for_single_req = MagicMock(return_value=0)
         queue._active_req_count = MagicMock(return_value=1)
         queue.token_to_kv_pool_allocator = MagicMock()
+        bind_separate_buffer_capacity(queue.token_to_kv_pool_allocator)
         queue.token_to_kv_pool_allocator.size_swa = 256
         queue.token_to_kv_pool_allocator.swa_available_size.return_value = 0
         queue.tree_cache = MagicMock()
@@ -162,6 +167,7 @@ class TestDecodeLockRefScenarios(CustomTestCase):
     def test_reclaim_swa_tail_capacity_page_rounds(self):
         queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
         queue.token_to_kv_pool_allocator = MagicMock(page_size=64)
+        bind_separate_buffer_capacity(queue.token_to_kv_pool_allocator)
         queue.token_to_kv_pool_allocator.swa_available_size.side_effect = [64, 192]
         queue.tree_cache = MagicMock()
 
@@ -175,6 +181,7 @@ class TestDecodeLockRefScenarios(CustomTestCase):
     def test_reclaim_swa_tail_capacity_fails_before_allocation(self):
         queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
         queue.token_to_kv_pool_allocator = MagicMock(page_size=64)
+        bind_separate_buffer_capacity(queue.token_to_kv_pool_allocator)
         queue.token_to_kv_pool_allocator.swa_available_size.side_effect = [64, 128]
         queue.tree_cache = MagicMock()
 
@@ -471,6 +478,7 @@ class TestDecodeLockRefScenarios(CustomTestCase):
         queue._update_handshake_waiters = MagicMock()
         queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
         queue._swa_tail_len = MagicMock(return_value=8)
+        queue._prealloc_required_tokens = MagicMock(return_value=(8, 8))
         queue._swa_aware_allocatable_token_budgets = MagicMock(return_value=(8, 8))
         queue._swa_tail_allocatable_token_budget = MagicMock(return_value=8)
         queue._match_prefix_and_lock = MagicMock(
@@ -497,6 +505,7 @@ class TestDecodeLockRefScenarios(CustomTestCase):
         queue.req_to_metadata_buffer_idx_allocator.available_size.return_value = 1
         queue.token_to_kv_pool = MagicMock()
         queue.token_to_kv_pool_allocator = MagicMock()
+        bind_separate_buffer_capacity(queue.token_to_kv_pool_allocator)
         queue.token_to_kv_pool_allocator.page_size = 4
 
         running_batch = MagicMock()
@@ -539,7 +548,7 @@ class TestDecodeLockRefScenarios(CustomTestCase):
             skip_swa=True,
         )
         self.assertFalse(req.swa_prefix_lock_released)
-        queue._swa_tail_len.assert_called_once_with(8)
+        queue._swa_tail_len.assert_called_with(8)
         queue._allocatable_token_budgets.assert_called_once()
 
     def test_hicache_restore_commit_hands_over_lock_with_receipt(self):

test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py

@@ -18,7 +18,10 @@ from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.separate_buffer_allocator_double import (
+    separate_buffer_allocator_double,
+)
+from sglang.test.test_utils import CustomTestCase, enter_override
 
 register_cpu_ci(est_time=12, suite="base-a-test-cpu")
 
@@ -102,7 +105,10 @@ def _make_prealloc_queue(
     """Build a minimal DecodePreallocQueue for _check_if_req_exceed_kv_capacity."""
     queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
     queue.max_total_num_tokens = max_total_num_tokens
-    queue.token_to_kv_pool_allocator = SimpleNamespace(size_swa=10**9)
+    queue.num_reserved_decode_tokens = 0
+    queue.token_to_kv_pool_allocator = separate_buffer_allocator_double(
+        page_size=1, size_swa=10**9
+    )
     # Disable the SWA-tail branch; this test only exercises the pool-length gate.
     queue._uses_swa_tail_prealloc = MagicMock(return_value=False)
 
@@ -128,6 +134,10 @@ def _make_req(rid: str, prompt_len: int):
 
 
 class TestCheckIfReqExceedKvCapacity(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        enter_override(self, get_context().override_server_args())
+
     def test_hisparse_admits_beyond_device_pool_up_to_host_backed_size(self):
         """Core regression: request longer than device-only
         `max_total_num_tokens` but within HiSparse host-backed

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -40,6 +40,7 @@ from sglang.srt.mem_cache.allocator.unified_sub_pool import (
     MultiEndedAllocator,
 )
 from sglang.srt.mem_cache.base_prefix_cache import EvictParams
+from sglang.srt.mem_cache.common import kv_to_page_indices
 from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens
 from sglang.srt.mem_cache.unified_cache.components import ComponentType
 from sglang.srt.mem_cache.unified_memory_pool import (
@@ -2522,6 +2523,29 @@ class TestSWACompositeKernelIdSurface(unittest.TestCase):
         expected = v2p_swa[v // self.PS] * (self.PS * mult) + v % self.PS
         self.assertTrue(torch.equal(a.translate_loc_from_full_to_swa(v), expected))
 
+    def test_swa_transfer_page_is_physical_not_kernel_scaled(self):
+        mult = 2 * self.SWA_L
+        a = self._build()
+        v = a.alloc(3 * self.PS)
+        self.assertIsNotNone(v)
+
+        physical_pages = a.swa_attn_allocator.virtual_to_physical[
+            v[:: self.PS] // self.PS
+        ]
+        physical_tokens = a.swa_attn_allocator.translate_kv_loc(v)
+        transfer_tokens = a.translate_swa_indices_for_transfer(v)
+        self.assertTrue(torch.equal(transfer_tokens, physical_tokens))
+        self.assertEqual(
+            kv_to_page_indices(transfer_tokens, self.PS).tolist(),
+            physical_pages.tolist(),
+        )
+
+        kernel_tokens = a.translate_loc_from_full_to_swa(v)
+        self.assertEqual(
+            kv_to_page_indices(kernel_tokens, self.PS).tolist(),
+            (physical_pages * mult).tolist(),
+        )
+
     def test_swa_kernel_tombstone_still_lands_on_sink(self):
         """The scaled stride must not break the tombstone clamp: a tombstoned
         page's ids (v2p == -1 -> -stride + offset, negative for every in-page

test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py

@@ -1,10 +1,13 @@
 import unittest
 from types import SimpleNamespace
+from unittest import mock
 
 import torch
+from test_unified_byte_budget_sizing import _swa_factory
 
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
@@ -57,5 +60,38 @@ class TestSWACpuCopyFilter(unittest.TestCase):
         )
 
 
+class TestUnifiedSWATransfers(CustomTestCase):
+    def test_cpu_copy_round_trip_with_request_index_and_swa_tombstone(self):
+        bundle = _swa_factory(page_size=4)
+        allocator = bundle.token_to_kv_pool_allocator
+        pool = bundle.token_to_kv_pool
+        indices = allocator.alloc(12)
+        allocator.free_swa(indices[:4])
+        raw = bundle.unified_memory_pool._raw
+        raw.copy_(torch.arange(raw.numel()).remainder(251).to(torch.uint8))
+        full_pages = allocator.translate_kv_indices_for_transfer(indices)[::4] // 4
+        swa_pages = allocator.translate_swa_indices_for_transfer(indices[4:])[::4] // 4
+        full_buffer, swa_buffer = (
+            raw[: side.num_pages * side.entry_bytes_per_page].view(
+                side.num_pages, side.entry_bytes_per_page
+            )
+            for side in (allocator.full_attn_allocator, allocator.swa_attn_allocator)
+        )
+        expected_full = full_buffer[full_pages].clone()
+        expected_swa = swa_buffer[swa_pages].clone()
+
+        # The buffers are CPU tensors; no device synchronization is needed.
+        with mock.patch(
+            "sglang.srt.mem_cache.memory_pool.current_platform.synchronize"
+        ):
+            saved = pool.get_cpu_copy(indices, req_pool_index=0)
+            raw.zero_()
+            pool.load_cpu_copy(saved, indices, req_pool_index=0)
+
+        self.assertTrue(torch.equal(full_buffer[full_pages], expected_full))
+        self.assertTrue(torch.equal(swa_buffer[swa_pages], expected_swa))
+        self.assertEqual(saved["swa_mask"].tolist(), [False] * 4 + [True] * 8)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_mha_views.py

@@ -344,17 +344,9 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
             "envelope move did not relocate exactly the named pages",
         )
 
-    def test_transfer_entry_points_fail_loud(self):
-        """The entry points that assume per-layer buffers indexed by TOKEN id
-        would silently mis-index against the row space (or hit a missing-attr
-        AttributeError), so each must raise. `get_contiguous_buf_infos` is NOT
-        among them: PD addresses this pool as whole page envelopes, pinned by
-        `test_pd_registration_is_one_whole_envelope` below."""
+    def test_prefix_valid_entry_point_fails_loud(self):
+        """Prefix-valid writes still assume token-major buffer indexing."""
         _, pool = _make_pool_and_kv(1)
-        with self.assertRaises(NotImplementedError):
-            pool.get_cpu_copy(torch.tensor([1]))
-        with self.assertRaises(NotImplementedError):
-            pool.load_cpu_copy(None, torch.tensor([1]))
         with self.assertRaises(NotImplementedError):
             pool.set_kv_buffer_prefix_valid()
 

test/registered/unit/mem_cache/test_unified_mla_views.py

@@ -26,9 +26,13 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
 import unittest
+from unittest import mock
 
 import torch
 
+from sglang.srt.mem_cache.allocator.unified_mamba import (
+    UnifiedMambaTokenToKVPoolAllocator,
+)
 from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator
 from sglang.srt.mem_cache.layout.page_major import build_mla_views
 from sglang.srt.mem_cache.unified_memory_pool import (
@@ -37,6 +41,7 @@ from sglang.srt.mem_cache.unified_memory_pool import (
     UnifiedKVPool,
     UnifiedMLATokenToKVPool,
 )
+from sglang.srt.runtime_context import get_parallel
 
 _DEV = "cpu"
 
@@ -219,6 +224,48 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase):
         k[7] = 2.5
         self.assertTrue(torch.all(v[7] == 2.5))
 
+    def test_cpu_copy_round_trips_through_physical_ids(self):
+        """REGRESSION: the host copy for decode retraction is addressed by
+        PHYSICAL token ids, but this pool's `kv_buffer` views are indexed by
+        kernel-facing ids. Without the rewrite the parent read a different row
+        and the restore silently returned other tokens' KV."""
+        for ps in (1, 4):
+            with self.subTest(page_size=ps):
+                pool, kv_pool = self._make(ps=ps)
+                phys = torch.tensor([0, 1, ps, ps + 1], dtype=torch.int64)
+                self.assertTrue(
+                    torch.equal(
+                        kv_pool._physical_to_kernel_indices(phys),
+                        torch.tensor(
+                            [_kernel_id(int(t), ps, _L) for t in phys],
+                            dtype=torch.int64,
+                        ),
+                    )
+                )
+                for layer in range(_L):
+                    kv_pool.get_key_buffer(layer)[
+                        kv_pool._physical_to_kernel_indices(phys)
+                    ] = float(layer + 1)
+
+                with (
+                    get_parallel().override(dcp_enabled=False),
+                    mock.patch(
+                        "sglang.srt.mem_cache.memory_pool.current_platform.synchronize"
+                    ),
+                ):
+                    saved = kv_pool.get_cpu_copy(phys)
+                    pool._raw.zero_()
+                    kv_pool.load_cpu_copy(saved, phys)
+
+                for layer in range(_L):
+                    restored = kv_pool.get_key_buffer(layer)[
+                        kv_pool._physical_to_kernel_indices(phys)
+                    ]
+                    self.assertTrue(
+                        torch.all(restored == float(layer + 1)),
+                        f"layer {layer} did not round-trip at page_size {ps}",
+                    )
+
     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."""
@@ -334,5 +381,58 @@ class TestTranslateKvLocForKernel(unittest.TestCase):
             self.assertTrue(torch.all(x == no_out))
 
 
+class _RecordingHybridPool:
+    """Stands in for `UnifiedHybridLinearKVPool`, recording the ids it is handed."""
+
+    def __init__(self, full_kv_pool, mamba_pool):
+        self.full_kv_pool = full_kv_pool
+        self.mamba_pool = mamba_pool
+        self.seen = None
+
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
+        self.seen = indices.clone()
+        return {"full": None}
+
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
+        self.seen = indices.clone()
+
+
+class TestMambaAllocatorCpuCopyIsPhysical(unittest.TestCase):
+    """REGRESSION: decode retraction calls the allocator's `get_cpu_copy` with
+    `req_to_token` rows, which hold VIRTUAL ids. This composite inherited the
+    raising base, and a plain delegate would have been just as wrong -- the
+    unified pools read those ids as PHYSICAL."""
+
+    def _build(self, ps=1):
+        pool, _, _ = _make_unified(page_size=ps)
+        kvcache = _RecordingHybridPool(
+            _FakeKVCache(pool.max_slots("full")),
+            _FakeKVCache(pool.max_slots("mamba")),
+        )
+        with get_parallel().override(dcp_enabled=False, attn_dcp_size=1):
+            allocator = UnifiedMambaTokenToKVPoolAllocator(
+                unified_buffer=pool, kvcache=kvcache, device=_DEV, page_size=ps
+            )
+        return allocator, kvcache
+
+    def test_pool_is_handed_physical_token_ids(self):
+        alloc, kvcache = self._build()
+        virtual = alloc.alloc(4)
+        self.assertIsNotNone(virtual)
+        virtual = virtual.to(torch.int64)
+        physical = alloc.full_attn_allocator.translate_kv_loc(virtual)
+        # Not identity here, so a delegate that passed the virtual ids straight
+        # through would read and restore other tokens' rows.
+        self.assertFalse(torch.equal(physical, virtual))
+
+        alloc.get_cpu_copy(virtual, req_pool_index=0)
+        self.assertTrue(torch.equal(kvcache.seen, physical))
+
+        alloc.load_cpu_copy({"full": None}, virtual, req_pool_index=0)
+        self.assertTrue(torch.equal(kvcache.seen, physical))
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_tri_pool.py

@@ -202,6 +202,58 @@ class TestUnifiedTriPool(unittest.TestCase):
         self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0)
         self.assertEqual(allocator.available_size(), before)
 
+    def test_pd_short_tail_fits_beyond_joint_capacity(self):
+        for lazy in (False, True):
+            for tail_len in (0, 5):
+                with self.subTest(lazy=lazy, tail_len=tail_len):
+                    _, allocator, _, _ = self._build(page_size=4, lazy_compaction=lazy)
+                    full = allocator.full_attn_allocator
+                    length = allocator.available_size() + 4
+                    self.assertFalse(allocator.can_reserve(length, length))
+                    self.assertTrue(allocator.can_reserve(length, tail_len))
+                    prefix = torch.tensor([0], dtype=torch.int64)
+                    seq = torch.tensor([length], dtype=torch.int64)
+                    with patch.object(
+                        full,
+                        "alloc_extend",
+                        side_effect=lambda *a, **kw: full.alloc(length),
+                    ):
+                        virtual = allocator.alloc_extend_swa_tail(
+                            prefix,
+                            prefix,
+                            seq,
+                            seq,
+                            torch.tensor([-1]),
+                            length,
+                            tail_len,
+                        )
+                    self.assertIsNotNone(virtual)
+                    self.assertEqual(full.allocated_count(), length)
+                    self.assertEqual(
+                        allocator.swa_attn_allocator.allocated_count(),
+                        -(-tail_len // 4) * 4,
+                    )
+                    self.assertEqual(allocator.verify_byte_accounting(), [])
+                    allocator.free(virtual)
+                    self.assertEqual(full.allocated_count(), 0)
+                    self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0)
+
+    def test_pd_tail_rejects_full_capacity_shortfall(self):
+        _, allocator, _, _ = self._build(page_size=4)
+        full = allocator.full_attn_allocator
+        length = full.available_size() + 4
+        prefix = torch.tensor([0], dtype=torch.int64)
+        seq = torch.tensor([length], dtype=torch.int64)
+        with patch.object(full, "alloc_extend") as extend:
+            self.assertIsNone(
+                allocator.alloc_extend_swa_tail(
+                    prefix, prefix, seq, seq, torch.tensor([-1]), length, 0
+                )
+            )
+        extend.assert_not_called()
+        self.assertEqual(full.allocated_count(), 0)
+        self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0)
+
     def test_empty_float_is_transparent_to_the_ends(self):
         _, allocator, _, _ = self._build()
         fa = allocator.full_attn_allocator
@@ -1408,5 +1460,64 @@ class TestFloatHoleCreditIsPerSide(unittest.TestCase):
         self.assertEqual(flt._byte_accounting_violations(), [])
 
 
+class TestPreallocIsPricedOnTheSharedGrid(unittest.TestCase):
+    """REGRESSION: PD admission compared FULL and SWA against per-side token
+    budgets, but each side's `available_size` credits the peer's drainable
+    holes, so a pair that each side can host alone can be jointly infeasible.
+    Such a pair was admitted and then refused inside `alloc_extend_swa_tail`."""
+
+    def _build(self, **kw):
+        return TestUnifiedTriPool._build(self, **kw)
+
+    def test_a_pair_each_side_can_host_alone_is_still_refused(self):
+        # page_size 1 leaves no slack between the per-side and joint views;
+        # the double-count only has room to show on a paged grid.
+        _, allocator, _, _ = self._build(page_size=4)
+        full_demand = allocator.full_available_size()
+        swa_demand = allocator.swa_available_size()
+        self.assertGreater(min(full_demand, swa_demand), 0)
+        # Each side alone reports room for its own half ...
+        self.assertLessEqual(full_demand, allocator.full_available_size())
+        self.assertLessEqual(swa_demand, allocator.swa_available_size())
+        # ... yet the two draw on the same bytes, so the grid refuses the pair.
+        self.assertFalse(
+            allocator._fits_page_demand(
+                -(-full_demand // allocator.page_size),
+                -(-swa_demand // allocator.page_size),
+            )
+        )
+        self.assertFalse(
+            allocator.prealloc_fits(
+                MagicMock(),
+                full_demand,
+                swa_demand,
+                full_budget_tokens=full_demand,
+                swa_budget_tokens=swa_demand,
+            )
+        )
+
+    def test_the_scheduler_budget_still_binds(self):
+        _, allocator, _, _ = self._build()
+        page_size = allocator.page_size
+        self.assertTrue(
+            allocator.prealloc_fits(
+                MagicMock(),
+                page_size,
+                page_size,
+                full_budget_tokens=page_size,
+                swa_budget_tokens=page_size,
+            )
+        )
+        self.assertFalse(
+            allocator.prealloc_fits(
+                MagicMock(),
+                page_size,
+                page_size,
+                full_budget_tokens=page_size - 1,
+                swa_budget_tokens=page_size,
+            )
+        )
+
+
 if __name__ == "__main__":
     unittest.main()