[PD] Support --enable-unified-memory with PD disaggregation (kimi-linear MLA hybrid-Mamba) (#33362)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/disaggregation/decode.py

@@ -327,6 +327,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
         self.transfer_queue = transfer_queue
         self.tree_cache = tree_cache
         self.gloo_group = gloo_group
+        # Destinations visible to prefill but not yet on the transfer queue.
+        self._num_published_destinations = 0
         self.tp_rank = tp_rank
         self.tp_size = tp_size
         self.dp_size = dp_size
@@ -1151,14 +1153,21 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 kv_indices = self.req_to_token_pool.req_to_token[
                     decode_req.req.req_pool_idx
                 ][total_prefix_len:origin_input_len]
+                kv_indices = (
+                    self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer(
+                        kv_indices
+                    )
+                )
 
             seq_len = origin_input_len
 
             def _mamba_payload():
                 return [
-                    self.req_to_token_pool.req_index_to_mamba_index_mapping[
+                    self.req_to_token_pool.translate_mamba_indices(
-                        decode_req.req.req_pool_idx
+                        self.req_to_token_pool.req_index_to_mamba_index_mapping[
-                    ]
+                            decode_req.req.req_pool_idx
+                        ]
+                    )
                     .cpu()
                     .numpy()
                 ]
@@ -1306,6 +1315,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                     decode_req.kv_receiver,
                     decode_req.req.build_rebootstrap_payload(),
                 )
+            self._num_published_destinations += 1
             preallocated_reqs.append(decode_req)
             indices_to_remove.add(i)
             decode_req.req.time_stats.set_decode_transfer_queue_entry_time()
@@ -1316,6 +1326,18 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
 
         return preallocated_reqs, failed_reqs
 
+    @property
+    def has_published_destinations(self) -> bool:
+        """Whether any destination address is visible to prefill but not yet
+        protected by the transfer queue."""
+        return self._num_published_destinations > 0
+
+    def note_destinations_queued(self, count: int) -> None:
+        """Hand `count` published destinations over to the transfer queue."""
+        self._num_published_destinations = max(
+            0, self._num_published_destinations - count
+        )
+
     @property
     def num_tokens_pre_allocated(self):
         return sum(
@@ -1798,6 +1820,10 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin):
 
     def extend(self, decode_reqs: List[DecodeRequest]) -> None:
         self.queue.extend(decode_reqs)
+        # This queue now covers them.
+        prealloc_queue = self.scheduler.disagg_decode_prealloc_queue
+        if prealloc_queue is not None:
+            prealloc_queue.note_destinations_queued(len(decode_reqs))
 
     def _commit_transfer_to_req(self, decode_req: DecodeRequest):
         idx = decode_req.metadata_buffer_index

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -8,7 +8,7 @@ import struct
 import threading
 import time
 from collections import defaultdict
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Set, Tuple, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -59,7 +59,7 @@ from sglang.srt.observability.trace import (
     TraceReqContext,
     trace_set_thread_info,
 )
-from sglang.srt.runtime_context import get_parallel, get_schedule
+from sglang.srt.runtime_context import get_memory, get_parallel, get_schedule
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils.network import NetworkAddress
 
@@ -282,35 +282,41 @@ class MooncakeKVManager(CommonKVManager):
     def init_engine(self):
         self.engine = get_mooncake_transfer_engine()
 
-    def register_buffer_to_engine(self):
+    def _registerable_regions(self) -> List[Tuple[int, int]]:
-        # Batch register KV data buffers
+        """(ptr, len) regions to (de)register, exact duplicates removed.
-        if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens:
-            self.engine.batch_register(
-                self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
-            )
 
-        # Batch register auxiliary data buffers
+        Deduped because the unified memory pool reports one raw buffer as both
-        if self.kv_args.aux_data_ptrs and self.kv_args.aux_data_lens:
+        its KV and its mamba state component, and double registration fails in
-            self.engine.batch_register(
+        the engine.
-                self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
+        """
-            )
+        regions: List[Tuple[int, int]] = []
+        seen: Set[Tuple[int, int]] = set()
 
+        def add(ptrs: List[int], lens: List[int]) -> None:
+            for ptr, length in zip(ptrs or [], lens or []):
+                if (ptr, length) not in seen:
+                    seen.add((ptr, length))
+                    regions.append((ptr, length))
+
+        add(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens)
+        add(self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens)
         for ptrs, lens in zip(
             self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
         ):
-            if ptrs and lens:
+            add(ptrs, lens)
-                self.engine.batch_register(ptrs, lens)
+        return regions
+
+    def register_buffer_to_engine(self):
+        regions = self._registerable_regions()
+        if regions:
+            ptrs, lens = zip(*regions)
+            self.engine.batch_register(list(ptrs), list(lens))
 
     def deregister_buffer_to_engine(self):
-        if self.kv_args.kv_data_ptrs:
+        regions = self._registerable_regions()
-            self.engine.batch_deregister(self.kv_args.kv_data_ptrs)
+        if regions:
-
+            ptrs, _ = zip(*regions)
-        if self.kv_args.aux_data_ptrs:
+            self.engine.batch_deregister(list(ptrs))
-            self.engine.batch_deregister(self.kv_args.aux_data_ptrs)
-
-        for ptrs in self.kv_args.state_data_ptrs or []:
-            if ptrs:
-                self.engine.batch_deregister(ptrs)
 
         if hasattr(self, "connection_pool"):
             with self.connection_lock:
@@ -775,6 +781,52 @@ class MooncakeKVManager(CommonKVManager):
             # compared to using multiple threads
             return process_layers(layers_params)
 
+    def _validate_envelope_kv_layout(
+        self,
+        dst_kv_ptrs: list[int],
+        dst_kv_item_len: Optional[int],
+        dst_attn_tp_size: Optional[int] = None,
+    ) -> None:
+        """Reject a peer whose KV registration shape differs from ours.
+
+        The unified memory pool registers ONE whole-envelope region and
+        addresses the destination as ``dst_ptr + page_id * item_len`` using OUR
+        ``item_len``, so a peer on a different page size / spec, or without
+        unified memory, would take envelope-sized blocks at the wrong offsets.
+        Must run before the first RDMA write.
+
+        Scoped to unified memory by config, not by region count: a non-unified
+        PP stage owning a single full-attention layer also registers one region,
+        and `_send_kvcache_generic` pairs that with the peer by layer id.
+        """
+        if not get_memory().enable_unified_memory:
+            return
+        if dst_attn_tp_size is not None and self.attn_tp_size != dst_attn_tp_size:
+            # The unified mamba state ships as one whole-slot envelope with no
+            # per-tensor dims, so `_send_mamba_state_slice` cannot reslice it and
+            # silently falls back to an unsliced copy. Reject here, before any KV
+            # is written, rather than in `maybe_send_extra` afterwards.
+            raise RuntimeError(
+                "--enable-unified-memory does not support different prefill / "
+                f"decode attention TP sizes (prefill={self.attn_tp_size}, "
+                f"decode={dst_attn_tp_size}): the whole-envelope state cannot "
+                "be TP-resliced."
+            )
+        src_item_lens = self.kv_args.kv_item_lens
+        if (
+            len(src_item_lens) != 1
+            or len(dst_kv_ptrs) != 1
+            or dst_kv_item_len is None
+            or src_item_lens[0] != dst_kv_item_len
+        ):
+            raise RuntimeError(
+                "PD KV layout mismatch on the whole-envelope path: prefill has "
+                f"{len(src_item_lens)} KV region(s) with item_lens="
+                f"{src_item_lens}, decode has {len(dst_kv_ptrs)} with item_len="
+                f"{dst_kv_item_len}. With --enable-unified-memory both sides "
+                "must enable it and use the same page size and model spec."
+            )
+
     def send_kvcache(
         self,
         mooncake_session_id: str,
@@ -784,7 +836,12 @@ class MooncakeKVManager(CommonKVManager):
         executor: concurrent.futures.ThreadPoolExecutor,
         dst_layer_ids: Optional[List[int]] = None,
         dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None,
+        dst_kv_item_len: Optional[int] = None,
+        dst_attn_tp_size: Optional[int] = None,
     ):
+        self._validate_envelope_kv_layout(
+            dst_kv_ptrs, dst_kv_item_len, dst_attn_tp_size
+        )
         dst_device_kv_ptrs = None
         if dst_device_kv_indices is not None:
             compression_ratios = self.kv_args.mla_compression_ratios
@@ -1243,6 +1300,17 @@ class MooncakeKVManager(CommonKVManager):
             )
 
             if st == StateType.MAMBA:
+                if (not src_dim_per_tensor or not dst_dim_per_tensor) and list(
+                    src_item_lens
+                ) != list(dst_item_lens):
+                    raise RuntimeError(
+                        "Mamba state layouts differ between prefill and decode "
+                        f"(src item_lens={src_item_lens}, dst item_lens="
+                        f"{dst_item_lens}) and no per-tensor dim metadata is "
+                        "available to reslice. With --enable-unified-memory, "
+                        "prefill and decode must both enable it and use equal "
+                        "attention TP sizes."
+                    )
                 if (
                     target_rank_registration_info is not None
                     and self.attn_tp_size
@@ -1693,6 +1761,8 @@ class MooncakeKVManager(CommonKVManager):
                                 executor,
                                 dst_layer_ids=target_rank_registration_info.dst_kv_layer_ids,
                                 dst_device_kv_indices=chunked_dst_device_kv_indice,
+                                dst_kv_item_len=target_rank_registration_info.dst_kv_item_len,
+                                dst_attn_tp_size=target_rank_registration_info.dst_attn_tp_size,
                             )
                         elif (
                             self.enable_staging

python/sglang/srt/disaggregation/prefill.py

@@ -790,6 +790,7 @@ class SchedulerDisaggregationPrefillMixin:
                 # Optimistic bootstrap can fail while this overlapped chunk is
                 # already running. Drop aborted chunks instead of sending KV.
                 if is_aborted(req):
+                    self.clear_pending_chunk_send(req)
                     advance_logprob_pt(i, req)
                     req.time_stats.set_last_chunked_prefill_finish_time()
                     continue
@@ -980,7 +981,17 @@ class SchedulerDisaggregationPrefillMixin:
 
         return transferred_rids
 
+    def clear_pending_chunk_send(self: Scheduler, req: Req) -> None:
+        """Drop `req` from the sent-but-unconcluded chunk set.
+
+        Every path that retires a request without a `last_chunk=True` send must
+        call this: a stale entry holds the unified-memory compaction gate closed
+        for the process lifetime.
+        """
+        self.disagg_prefill_pending_chunk_rids.discard(req.rid)
+
     def handle_bootstrap_failure(self: Scheduler, req: Req) -> None:
+        self.clear_pending_chunk_send(req)
         error_message = (
             f"Prefill bootstrap failed for request rank={self.ps.tp_rank} "
             f"{req.rid=} {req.bootstrap_room=}"
@@ -1178,9 +1189,11 @@ class SchedulerDisaggregationPrefillMixin:
 
             def _mamba_payload():
                 return [
-                    self.req_to_token_pool.req_index_to_mamba_index_mapping[
+                    self.req_to_token_pool.translate_mamba_indices(
-                        req.req_pool_idx
+                        self.req_to_token_pool.req_index_to_mamba_index_mapping[
-                    ]
+                            req.req_pool_idx
+                        ]
+                    )
                     .cpu()
                     .numpy()
                 ]
@@ -1287,6 +1300,13 @@ class SchedulerDisaggregationPrefillMixin:
             kv_indices = self.req_to_token_pool.req_to_token[
                 req.req_pool_idx, seg_start:seg_end
             ]
+            # Unified memory: req_to_token holds VIRTUAL ids; the transfer needs
+            # physical ones. Per segment, since each is its own gather.
+            kv_indices = (
+                self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer(
+                    kv_indices
+                )
+            )
             page_indices = kv_to_page_indices(kv_indices, page_size)
             segment_is_last = last_chunk and is_final_segment
             if not req.disagg_kv_sender.should_send_kv_chunk(
@@ -1299,6 +1319,12 @@ class SchedulerDisaggregationPrefillMixin:
                 num_kv_tokens=seg_end - seg_start,
             )
         req.start_send_idx = end_idx
+        # A last chunk needs no entry: every `last_chunk=True` call site has
+        # already put the request on `disagg_prefill_inflight_queue`.
+        if last_chunk:
+            self.disagg_prefill_pending_chunk_rids.discard(req.rid)
+        else:
+            self.disagg_prefill_pending_chunk_rids.add(req.rid)
 
     def optimistic_release_and_requeue(self: Scheduler, req: Req) -> None:
         """Release KV cache and requeue an optimistic prefill request."""
@@ -1308,6 +1334,7 @@ class SchedulerDisaggregationPrefillMixin:
         req.reset_for_retract()
         req.output_ids = array("q")
         req.start_send_idx = 0
+        self.clear_pending_chunk_send(req)  # re-sends from scratch
         req.tmp_end_idx = -1
         req.disagg_decode_prefix_len = 0
         req.early_send_prefix_end = None

python/sglang/srt/disaggregation/utils.py

@@ -111,6 +111,50 @@ class DisaggregationMode(Enum):
         return "unified"
 
 
+def unified_memory_disagg_move_gate(scheduler):
+    """Compaction move gate for a PD node running the unified memory pool.
+
+    Returns a predicate that is True only when no transfer can be in flight, so
+    compaction never relocates a page the RDMA engine is reading or writing.
+    Safe to read this state from here: every mover runs on the scheduler thread.
+
+    A page is exposed from the moment its address reaches the peer until the
+    transfer concludes, and for part of that lifetime the request is in NEITHER
+    end's queue -- so queue emptiness alone is not enough:
+
+    - PREFILL: scheduling the final chunk clears `chunked_req` while earlier
+      chunks may still be draining, and the request only reaches the inflight
+      queue later, in the result path.
+    - DECODE: `pop_preallocated` publishes one request's destinations and keeps
+      allocating for the next, whose allocation can urgently flush the peer
+      sub-allocator; the batch reaches the transfer queue only after the loop.
+    """
+    if scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
+
+        def prefill_gate() -> bool:
+            return not (
+                scheduler.disagg_prefill_inflight_queue
+                or scheduler.disagg_prefill_pending_chunk_rids
+            )
+
+        return prefill_gate
+
+    if scheduler.disaggregation_mode == DisaggregationMode.DECODE:
+
+        def decode_gate() -> bool:
+            return not (
+                scheduler.disagg_decode_transfer_queue.queue
+                or scheduler.disagg_decode_prealloc_queue.has_published_destinations
+            )
+
+        return decode_gate
+
+    raise ValueError(
+        "unified_memory_disagg_move_gate: scheduler is not a PD node "
+        f"(mode={scheduler.disaggregation_mode})"
+    )
+
+
 #########################
 # Synchronization
 #########################

python/sglang/srt/managers/scheduler.py

@@ -25,7 +25,7 @@ from collections import deque
 from contextlib import contextmanager, nullcontext
 from functools import partial
 from http import HTTPStatus
-from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple, Union
 
 from sglang.srt.runtime_context import (
     get_device,
@@ -90,6 +90,7 @@ from sglang.srt.disaggregation.utils import (
     TransferBackend,
     get_dsa_seed_metadata_dim,
     prepare_abort,
+    unified_memory_disagg_move_gate,
 )
 from sglang.srt.distributed import get_pp_group, get_world_group
 from sglang.srt.distributed.parallel_state import get_tp_group
@@ -1396,9 +1397,19 @@ class Scheduler(
             )
             # The prefill requests that are in the middle of kv sending
             self.disagg_prefill_inflight_queue: List[Req] = []
+            # Requests with a sent chunk that are not yet on the inflight queue.
+            self.disagg_prefill_pending_chunk_rids: Set[str] = set()
 
             self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
 
+        if (
+            self.enable_unified_memory
+            and self.disaggregation_mode != DisaggregationMode.NULL
+        ):
+            self.token_to_kv_pool_allocator.set_disagg_move_gate(
+                unified_memory_disagg_move_gate(self)
+            )
+
         # Init mm receiver for EPD disaggregation mode
         if get_disagg().language_only and get_disagg().encoder_transfer_backend in [
             "zmq_to_scheduler",
@@ -2928,6 +2939,7 @@ class Scheduler(
         req.time_stats.trace_ctx.abort(abort_info={"reason": "Aborted"})
         req.to_finish = None
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            self.clear_pending_chunk_send(req)
             req.disagg_kv_sender.abort()
             maybe_release_metadata_buffer(
                 req, self.req_to_metadata_buffer_idx_allocator

python/sglang/srt/mem_cache/allocator/base.py

@@ -82,6 +82,16 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
                 (0,), dtype=self.release_pages.dtype, device=self.device
             )
 
+    def translate_kv_indices_for_transfer(
+        self, kv_indices: torch.Tensor
+    ) -> torch.Tensor:
+        """Token ids as the PD-disaggregation transfer engine addresses them.
+
+        Identity here: a static pool's token ids index its registered buffers
+        directly. Virtual-id pools must override.
+        """
+        return kv_indices
+
     def get_cpu_copy(self, indices, mamba_indices=None):
         # FIXME: reuse the get_cpu_copy after paged allocator is implemented
         raise NotImplementedError()

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -353,17 +353,29 @@ class KVCacheConfigurator:
         # Unified-pool fast path: build req_to_token + token_to_kv pool + allocator
         # from one byte buffer, then return. Gated to the target worker
         # (req_to_token_pool is None); supports hybrid Mamba and hybrid SWA (not DSV4).
-        if (
+        if get_memory().enable_unified_memory and req_to_token_pool is None:
-            get_memory().enable_unified_memory
+            pd_enabled = get_disagg().disaggregation_mode != "null"
-            and get_disagg().disaggregation_mode == "null"
-            and req_to_token_pool is None
-        ):
             if self.mambaish_config is not None:
+                if pd_enabled and not self.use_mla_backend:
+                    raise ValueError(
+                        "--enable-unified-memory with PD disaggregation "
+                        "currently supports only MLA hybrid-Mamba models "
+                        "(e.g. kimi-linear); this model uses the MHA full-"
+                        "attention pool. Drop --enable-unified-memory or run "
+                        "without PD disaggregation."
+                    )
                 bundle = self._init_unified_mamba_pools(
                     max_num_reqs=sizes.max_running_requests,
                     max_total_num_tokens=sizes.max_total_num_tokens,
                 )
             elif self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config):
+                if pd_enabled:
+                    raise ValueError(
+                        "--enable-unified-memory with PD disaggregation does "
+                        "not support hybrid-SWA models yet (no whole-envelope "
+                        "transfer scheme for the SWA sub-pool). Drop "
+                        "--enable-unified-memory or run without PD."
+                    )
                 bundle = self._init_unified_swa_pools(
                     max_num_reqs=sizes.max_running_requests,
                     full_max_total_num_tokens=sizes.full_max_total_num_tokens,
@@ -564,6 +576,11 @@ class KVCacheConfigurator:
             speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
             disable_overlap_schedule=get_schedule().disable_overlap_schedule,
             need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"),
+            decode_pre_alloc_size=(
+                get_disagg().disaggregation_decode_extra_slots
+                if get_disagg().disaggregation_mode == "decode"
+                else 0
+            ),
             mamba_full_memory_ratio=get_schedule().mamba_full_memory_ratio,
             # Overlap mode: the allocator's `free` drops a wait_stream(forward_stream)
             # barrier so eager compaction serializes after the in-flight forward's

python/sglang/srt/mem_cache/multi_ended_allocator.py

@@ -25,7 +25,7 @@ from __future__ import annotations
 import inspect
 import logging
 import os
-from typing import Dict, List, Optional, Set, Tuple
+from typing import Callable, Dict, List, Optional, Set, Tuple
 
 import torch
 from torch.profiler import record_function
@@ -218,6 +218,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             _STATS_INSTANCES.add(self)
             _install_signal_handlers_once()
         self.live_page_count = 0
+        # While this returns False, `_flush` must not relocate any page.
+        self.disagg_move_gate: Optional[Callable[[], bool]] = None
         self._latest_forward_done_event: Optional[torch.cuda.Event] = None
         # Most-recent forward's (done_event, out_cache_loc_virtual) for `_flush`'s
         # write-race check. Single slot: at most ONE forward in flight per call site.
@@ -393,6 +395,12 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         peer = self._peer
         if peer is None or not peer.lazy_compaction:
             return 0
+        if peer.disagg_move_gate is not None and not peer.disagg_move_gate():
+            # The peer cannot compact while a PD transfer is in flight, so these
+            # holes are not realizable. Crediting them would let the scheduler
+            # admit work that `_flush_peer_for_alloc` then cannot satisfy, and
+            # the caller treats a failed alloc as a memory-estimation bug.
+            return 0
         return len(peer._free_phys_pages) * peer.entry_bytes_per_page
 
     def schedulable_available_size(self) -> int:
@@ -1058,6 +1066,10 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             self._compact_pending_impl(freed_physical_pages)
 
     def _compact_pending_impl(self, freed_physical_pages: torch.Tensor) -> None:
+        assert self.disagg_move_gate is None, (
+            f"_compact_pending({self.sub_pool_name!r}): eager compaction ran with "
+            "a PD-disaggregation move gate installed; PD requires lazy_compaction."
+        )
         freed_set = set(int(x) for x in freed_physical_pages.tolist())
         if not freed_set:
             return
@@ -1440,6 +1452,9 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         """
         if not self.lazy_compaction:
             return 0
+        if self.disagg_move_gate is not None and not self.disagg_move_gate():
+            # Holes stay in the free list; the next flush picks them up.
+            return 0
         self._stats_n_flush_calls += 1
         with record_function("MultiEndedAlloc._flush"):
             self._drain_pending_reuse(urgent=urgent)
@@ -1936,6 +1951,26 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         to the physical translate when `kernel_page_multiplier == 1` (MHA)."""
         return self.full_attn_allocator.translate_kv_loc_dense(loc, out=out)
 
+    def translate_kv_indices_for_transfer(
+        self, kv_indices: torch.Tensor
+    ) -> torch.Tensor:
+        """Virtual TOKEN ids -> PHYSICAL token ids for the PD transfer engine.
+
+        PHYSICAL, not dense: the transfer registers page ENVELOPES (see
+        `UnifiedMLATokenToKVPool.get_contiguous_buf_infos`).
+        """
+        return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
+
+    def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None:
+        """Install the PD-disaggregation move gate on both sub-allocators."""
+        assert self.lazy_compaction, (
+            "PD disaggregation with the unified memory pool requires lazy "
+            "compaction (eager free-path compaction moves pages under "
+            "in-flight transfers)."
+        )
+        self.full_attn_allocator.disagg_move_gate = gate
+        self.mamba_allocator.disagg_move_gate = gate
+
     def is_slot_allocated(self, slot: int) -> bool:
         return self.full_attn_allocator.is_slot_allocated(slot)
 

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -42,6 +42,7 @@ from sglang.srt.mem_cache.layout.page_major import (
     build_page_major_mha_views,
 )
 from sglang.srt.mem_cache.memory_pool import (
+    HybridLinearKVPool,
     HybridReqToTokenPool,
     MambaPool,
     MHATokenToKVPool,
@@ -663,6 +664,20 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
     def get_kv_size_bytes(self):
         return 0  # UnifiedKVPool logs the total; per-sub-pool would double-count
 
+    def get_contiguous_buf_infos(self):
+        """PD-transfer registration: ONE entry, the raw buffer, addressed as
+        ``raw_ptr + physical_page_id * page_envelope_bytes``.
+
+        The transfer item is the whole page envelope (all layers of one page)
+        rather than a per-layer region, because the per-layer dense views
+        overlap and index in dense ids. Both sides must therefore build the
+        pool with identical specs.
+        """
+        # The address formula omits the anchor; a nonzero one would mis-address.
+        assert self._unified_buffer.anchor_bytes(self._sub_pool_name) == 0
+        raw = self._unified_buffer._raw
+        return [raw.data_ptr()], [raw.numel()], [self._page_bytes]
+
     def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
         """Relocate whole page envelopes.
 
@@ -812,6 +827,31 @@ class UnifiedMambaPool(MambaPool):
         # Physical-slot copy used by the allocator's `_compact_pending`.
         MambaPool.copy_from(self, src_index, dst_index)
 
+    # -- PD state transfer (StateType.MAMBA) --
+    # The transfer item is the whole per-slot envelope, addressed as
+    # `raw_ptr + physical_slot * entry_bytes`. An envelope cannot be TP-resliced
+    # or PP-subset, so the per-tensor metadata below stays empty and both sides
+    # must build identical mamba specs (equal attn TP, pp=1).
+
+    def get_contiguous_buf_infos(self):
+        # The address formula omits the anchor; a nonzero one would mis-address.
+        assert self._unified_buffer.anchor_bytes(self._sub_pool_name) == 0
+        spec = self._unified_buffer.mamba_spec(self._sub_pool_name)
+        raw = self._unified_buffer._raw
+        return [raw.data_ptr()], [raw.numel()], [spec.entry_bytes()]
+
+    def get_state_dim_per_tensor(self):
+        return []
+
+    def get_state_layer_ids(self):
+        return []
+
+    def get_state_slice_outer_counts(self):
+        return []
+
+    def get_state_conv_shard_groups(self):
+        return []
+
 
 class UnifiedMambaSlotAllocator:
     """Mamba slot allocator (PHYSICAL view) for the unified memory pool.
@@ -929,6 +969,7 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
         speculative_num_draft_tokens: Optional[int] = None,
         enable_overlap_schedule: bool = True,
         start_layer: Optional[int] = None,
+        pre_alloc_size: int = 0,
     ):
         self._unified_buffer = unified_buffer
         self._mamba_sub_pool_name = mamba_sub_pool_name
@@ -936,7 +977,10 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
             unified_buffer.max_slots(mamba_sub_pool_name) - 1
         )  # reserve slot 0
         super().__init__(
-            size=size,
+            # `DecodeReqToTokenPool` semantics: rows cover the preallocated
+            # requests too, while `self.size` (rebound below) stays the
+            # running-request cap the scheduler and leak invariant expect.
+            size=size + pre_alloc_size,
             mamba_size=self._shared_mamba_size,
             mamba_spec_state_size=mamba_spec_state_size,
             max_context_len=max_context_len,
@@ -949,6 +993,8 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
             enable_overlap_schedule=enable_overlap_schedule,
             start_layer=start_layer,
         )
+        self.size = size
+        self.pre_alloc_size = pre_alloc_size
 
     def _init_mamba_pool(
         self,
@@ -1013,6 +1059,17 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
         return self.mamba_allocator.translate(virtual_ids).to(torch.int32)
 
 
+class UnifiedHybridLinearKVPool(HybridLinearKVPool):
+    """`HybridLinearKVPool` over unified sub-pools (full = Unified{MLA,MHA},
+    mamba = UnifiedMambaPool)."""
+
+    def get_kv_layer_ids(self):
+        # Empty: the KV component is one whole-envelope entry, so there are no
+        # per-layer entries to pair by layer id (the sender falls back to
+        # positional pairing).
+        return []
+
+
 # ---------------------------------------------------------------------------
 # Factory
 # ---------------------------------------------------------------------------
@@ -1054,9 +1111,9 @@ def init_unified_mamba_pools(
     mamba_full_memory_ratio: Optional[float] = None,  # informational only
     forward_stream: Optional[torch.cuda.Stream] = None,
     lazy_compaction: bool = False,
+    decode_pre_alloc_size: int = 0,
 ) -> UnifiedPoolBundle:
     """Build the Mamba-hybrid unified-memory-pool stack."""
-    from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
     from sglang.srt.mem_cache.multi_ended_allocator import (
         UnifiedMambaTokenToKVPoolAllocator,
     )
@@ -1132,6 +1189,7 @@ def init_unified_mamba_pools(
         speculative_num_draft_tokens=speculative_num_draft_tokens,
         enable_overlap_schedule=not disable_overlap_schedule,
         start_layer=start_layer,
+        pre_alloc_size=decode_pre_alloc_size,
     )
     if use_mla_backend:
         # start_layer stays 0: HybridLinearKVPool patches layer ids to the dense
@@ -1153,7 +1211,7 @@ def init_unified_mamba_pools(
     full_attn_layer_ids_for_pool = (
         [0] if is_draft_worker else list(full_attention_layer_ids)
     )
-    token_to_kv_pool = HybridLinearKVPool(
+    token_to_kv_pool = UnifiedHybridLinearKVPool(
         page_size=page_size,
         size=max_total_num_tokens,
         dtype=kv_cache_dtype,

python/sglang/srt/server_args.py

@@ -8007,9 +8007,29 @@ class ServerArgs:
     def _handle_unified_memory_pool(self):
         if not self.enable_unified_memory:
             return
-        assert self.disaggregation_mode == "null", (
+        if self.disaggregation_mode != "null":
-            "--enable-unified-memory is not yet compatible with PD " "disaggregation."
+            # Constraints of the whole-envelope transfer; see
-        )
+            # UnifiedMLATokenToKVPool.get_contiguous_buf_infos.
+            assert self.disaggregation_transfer_backend == "mooncake", (
+                "--enable-unified-memory with PD disaggregation supports only "
+                "the mooncake transfer backend; got "
+                f"{self.disaggregation_transfer_backend!r}."
+            )
+            assert self.pp_size == 1, (
+                "--enable-unified-memory with PD disaggregation does not support "
+                "pipeline parallelism (whole-envelope transfer has no per-layer "
+                "entries to subset)."
+            )
+            assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), (
+                "--enable-unified-memory with PD disaggregation requires lazy "
+                "compaction; unset SGLANG_DISABLE_LAZY_COMPACTION."
+            )
+            assert not self.enable_hisparse, (
+                "--enable-unified-memory with PD disaggregation is not compatible "
+                "with --enable-hisparse: the decode-side HiSparse prealloc path "
+                "ships host/C4 rows straight from the allocator, bypassing the "
+                "virtual->physical translation the unified pool needs."
+            )
         assert self.speculative_algorithm in (None, "DSPARK"), (
             "--enable-unified-memory only supports --speculative-algorithm "
             "DSPARK (chain draft); other speculative algorithms are not yet "

test/registered/disaggregation/test_disaggregation_unified_memory.py

@@ -0,0 +1,152 @@
+"""PD disaggregation with --enable-unified-memory (MLA hybrid-Mamba).
+
+Guards the unified-memory PD transfer scheme end to end: whole page-envelope
+KV registration (`UnifiedMLATokenToKVPool.get_contiguous_buf_infos`), whole
+slot-envelope KDA/mamba state transfer, virtual->physical index translation at
+the prefill send / decode prealloc sites, and the compaction move gate. A
+regression in any of them shifts the decode-side KV/state bytes and breaks
+logprob parity with the non-PD unified-memory reference.
+
+`--attention-backend` is deliberately NOT pinned, matching
+`models_e2e/test_kimi_linear_unified_memory.py`, which documents that pinning
+hides defects reachable only under the resolved default. The transferred bytes
+are backend-independent, so the default (fa3 on this suite's H100 runner) covers
+this file's subject either way. The linear-attn/Mamba backends stay pinned to
+triton -- the page-major layout requires them.
+
+`--enable-deterministic-inference` is deliberately NOT set. It would only guard
+against batch-shape-dependent kernel variation, and the reference and P+D paths
+run the same shapes: measured, two fresh servers on separate GPUs produce
+bit-identical logits without it. Setting it would narrow the test to the
+batch-invariant op set and a non-default sampling backend -- a less
+representative config -- and couple a PD-transfer test to the deterministic code
+path, so a defect there would fail this file for an unrelated reason.
+"""
+
+import time
+import unittest
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+    assert_process_healthy,
+)
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-h100")
+
+KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random"
+SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"}
+SERVER_ARGS = [
+    "--skip-tokenizer-init",
+    "--random-seed",
+    "1",
+    "--enable-unified-memory",
+    "--linear-attn-backend",
+    "triton",
+    "--mamba-backend",
+    "triton",
+    "--max-mamba-cache-size",
+    "32",
+    "--max-total-tokens",
+    "4096",
+    "--cuda-graph-backend-decode",
+    "disabled",
+    "--cuda-graph-backend-prefill",
+    "disabled",
+]
+
+
+class TestUnifiedMemoryDisaggregation(PDDisaggregationServerBase):
+    """1 prefill + 1 decode, both with --enable-unified-memory, vs a non-PD
+    unified-memory reference server."""
+
+    prefill_tp_size = 1
+    decode_tp_size = 1
+    decode_base_gpu_id = 1
+    extra_prefill_args = SERVER_ARGS
+    extra_decode_args = SERVER_ARGS
+    extra_prefill_env = SERVER_ENV
+    extra_decode_env = SERVER_ENV
+    baseline_args = SERVER_ARGS
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.model = KIMI_LINEAR_MODEL
+
+    @staticmethod
+    def generate(base_url):
+        response = requests.post(
+            base_url + "/generate",
+            json={
+                "input_ids": [1] + [100 + i % 1000 for i in range(256)],
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 4,
+                    "ignore_eos": True,
+                },
+                "return_logprob": True,
+                "top_logprobs_num": 5,
+            },
+            timeout=120,
+        )
+        response.raise_for_status()
+        return response.json()["meta_info"]
+
+    def test_logprob_parity(self):
+        baseline = popen_launch_server(
+            self.model,
+            self.lb_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=["--trust-remote-code"] + self.baseline_args,
+            env=SERVER_ENV,
+        )
+        try:
+            reference = self.generate(self.lb_url)
+        finally:
+            kill_process_tree(baseline.pid, wait_timeout=60)
+        time.sleep(5)
+
+        self.launch_all()
+        disaggregated = self.generate(self.lb_url)
+
+        reference_logprobs = reference["output_token_logprobs"]
+        disaggregated_logprobs = disaggregated["output_token_logprobs"]
+        self.assertEqual(
+            [item[1] for item in reference_logprobs],
+            [item[1] for item in disaggregated_logprobs],
+        )
+        self.assertEqual(len(reference_logprobs), 4)
+        for reference_item, disaggregated_item in zip(
+            reference_logprobs, disaggregated_logprobs
+        ):
+            self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05)
+
+        assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url)
+        assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url)
+        assert_process_healthy(self, "decode", self.process_decode, self.decode_url)
+
+
+class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregation):
+    """Multi-chunk prefill (257-token prompt, 64-token chunks): each chunk's KV
+    pages are translated to physical ids and shipped while later chunks still
+    run, exercising the chunked send path and the prefill-side move gate
+    (`chunked_req.start_send_idx > 0`). The reference server uses the same
+    chunk size so any parity break isolates to the PD transfer.
+    """
+
+    _chunked_args = SERVER_ARGS + ["--chunked-prefill-size", "64"]
+    extra_prefill_args = _chunked_args
+    extra_decode_args = _chunked_args
+    baseline_args = _chunked_args
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_unified_memory_move_gate.py

@@ -0,0 +1,179 @@
+"""Regression tests for the unified-memory PD compaction move gate.
+
+The gate decides when lazy compaction may relocate physical pages. A page is
+exposed to the peer from the moment its address is published until the transfer
+concludes, and for part of that lifetime the request sits in NEITHER end's
+queue. Both cases below are exactly those windows: an earlier version of the
+predicates looked only at `disagg_prefill_inflight_queue` /
+`disagg_decode_transfer_queue` (plus `scheduler.chunked_req`) and returned True
+here, letting compaction move pages under in-flight RDMA -- silent KV
+corruption with no crash.
+"""
+
+import unittest
+from typing import List, Optional, Set
+
+from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    unified_memory_disagg_move_gate,
+)
+from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=30, suite="base-a-test-cpu")
+
+
+class _FakeTransferQueue:
+    def __init__(self):
+        self.queue: List[object] = []
+
+
+class _FakePreallocQueue:
+    """Mirrors the real queue's published-destination bookkeeping."""
+
+    def __init__(self):
+        self._num_published_destinations = 0
+
+    @property
+    def has_published_destinations(self) -> bool:
+        return self._num_published_destinations > 0
+
+    def note_destinations_published(self) -> None:
+        self._num_published_destinations += 1
+
+    def note_destinations_queued(self, count: int) -> None:
+        self._num_published_destinations = max(
+            0, self._num_published_destinations - count
+        )
+
+
+class _FakeScheduler:
+    def __init__(self, mode: DisaggregationMode):
+        self.disaggregation_mode = mode
+        self.chunked_req: Optional[object] = None
+        self.disagg_prefill_inflight_queue: List[object] = []
+        self.disagg_prefill_pending_chunk_rids: Set[str] = set()
+        self.disagg_decode_transfer_queue = _FakeTransferQueue()
+        self.disagg_decode_prealloc_queue = _FakePreallocQueue()
+
+
+class TestDecodeMoveGate(CustomTestCase):
+    def test_closed_while_destination_published_but_not_queued(self):
+        """`pop_preallocated` publishes request A's destination addresses via
+        `send_metadata`, then keeps allocating for request B in the same loop;
+        the batch only reaches the transfer queue after the loop returns. B's
+        allocation can urgently flush the peer sub-allocator, so the gate must
+        stay closed across that window even though the transfer queue is empty.
+        """
+        scheduler = _FakeScheduler(DisaggregationMode.DECODE)
+        gate = unified_memory_disagg_move_gate(scheduler)
+        self.assertTrue(gate(), "idle decode node should allow compaction")
+
+        # A's destination is now visible to prefill; transfer queue still empty.
+        scheduler.disagg_decode_prealloc_queue.note_destinations_published()
+        self.assertFalse(scheduler.disagg_decode_transfer_queue.queue)
+        self.assertFalse(gate())
+
+        # Handing the batch to the transfer queue transfers responsibility.
+        scheduler.disagg_decode_transfer_queue.queue.append(object())
+        scheduler.disagg_decode_prealloc_queue.note_destinations_queued(1)
+        self.assertFalse(gate(), "transfer queue still holds it")
+
+        scheduler.disagg_decode_transfer_queue.queue.clear()
+        self.assertTrue(gate())
+
+
+class TestPrefillMoveGate(CustomTestCase):
+    def test_closed_after_final_chunk_clears_chunked_req(self):
+        """Scheduling the final chunk clears `scheduler.chunked_req`, but the
+        request only reaches `disagg_prefill_inflight_queue` later in the result
+        path. Earlier middle chunks may still be draining in that window, so the
+        gate must not key off `chunked_req` alone.
+        """
+        scheduler = _FakeScheduler(DisaggregationMode.PREFILL)
+        gate = unified_memory_disagg_move_gate(scheduler)
+        self.assertTrue(gate(), "idle prefill node should allow compaction")
+
+        # A middle chunk went out for rid "r0".
+        scheduler.chunked_req = object()
+        scheduler.disagg_prefill_pending_chunk_rids.add("r0")
+        self.assertFalse(gate())
+
+        # Final chunk scheduled: chunked_req cleared, not yet inflight-queued.
+        scheduler.chunked_req = None
+        self.assertFalse(scheduler.disagg_prefill_inflight_queue)
+        self.assertFalse(gate())
+
+        # Last chunk sent: the request is on the inflight queue, which covers it.
+        scheduler.disagg_prefill_inflight_queue.append(object())
+        scheduler.disagg_prefill_pending_chunk_rids.discard("r0")
+        self.assertFalse(gate())
+
+        scheduler.disagg_prefill_inflight_queue.clear()
+        self.assertTrue(gate())
+
+    def test_reopens_when_middle_sent_request_is_retired_without_final_chunk(self):
+        """A request aborted after a middle chunk never reaches a `last_chunk`
+        send, so its rid is only dropped by the abort/release cleanup. Without
+        that discard the gate stays closed for the process lifetime and lazy
+        compaction never packs the free list again -- a liveness leak that ends
+        in allocation failure despite reclaimable space.
+        """
+        scheduler = _FakeScheduler(DisaggregationMode.PREFILL)
+        gate = unified_memory_disagg_move_gate(scheduler)
+
+        scheduler.chunked_req = object()
+        scheduler.disagg_prefill_pending_chunk_rids.add("r0")
+        self.assertFalse(gate())
+
+        # Aborted mid-chunking: chunked_req dropped, no final send, never queued.
+        scheduler.chunked_req = None
+        scheduler.disagg_prefill_pending_chunk_rids.discard("r0")
+        self.assertTrue(gate(), "abort cleanup must let compaction resume")
+
+
+class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
+    """`schedulable_available_size` credits holes a peer urgent-flush would
+    release. While the move gate is closed that flush relocates nothing, so
+    crediting them lets the scheduler admit work `_flush_peer_for_alloc` cannot
+    satisfy; the alloc then returns None and the decode prealloc path treats
+    that as a memory-estimation bug and aborts the scheduler.
+    """
+
+    class _Peer:
+        def __init__(self, gate):
+            self.lazy_compaction = True
+            self._free_phys_pages = [0, 1, 2, 3]  # only len() is read
+            self.entry_bytes_per_page = 512
+            self.disagg_move_gate = gate
+
+    class _Owner:
+        def __init__(self, peer):
+            self._peer = peer
+
+    def _credit(self, gate):
+        peer = self._Peer(gate)
+        owner = self._Owner(peer)
+        return MultiEndedAllocator._peer_drainable_hole_bytes(owner)
+
+    def test_credit_follows_the_gate(self):
+        # No PD gate installed (non-disagg): holes are realizable as before.
+        self.assertEqual(self._credit(gate=None), 4 * 512)
+        # Gate open: peer can compact, so the credit stands.
+        self.assertEqual(self._credit(gate=lambda: True), 4 * 512)
+        # Gate closed: an urgent flush would move nothing, so credit nothing.
+        self.assertEqual(self._credit(gate=lambda: False), 0)
+
+
+class TestMoveGateRejectsNonPdNode(CustomTestCase):
+    def test_null_mode_is_rejected(self):
+        """The gate is only meaningful on a PD node; a NULL-mode scheduler is a
+        wiring bug and must not silently produce an always-open predicate."""
+        scheduler = _FakeScheduler(DisaggregationMode.NULL)
+        with self.assertRaises(ValueError):
+            unified_memory_disagg_move_gate(scheduler)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/managers/test_priority_scheduling_disaggregation.py

@@ -109,6 +109,9 @@ class TestDecodePreallocQueuePriority(unittest.TestCase):
         queue.pending_reqs = []
         queue.retracted_queue = []
         queue.num_reserved_decode_tokens = 0
+        # `pop_preallocated` credits this counter; `__new__` skips the __init__
+        # that seeds it.
+        queue._num_published_destinations = 0
         queue._resolve_pending_reqs = MagicMock()
         queue._update_handshake_waiters = MagicMock()
         queue._allocatable_tokens = MagicMock(return_value=1000)

test/registered/unit/mem_cache/test_hisparse_allocator.py

@@ -9,6 +9,7 @@ import torch
 from sglang.srt.mem_cache.allocator.hisparse import (
     DeepSeekV4HiSparseTokenToKVPoolAllocator,
 )
+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
 
@@ -214,14 +215,18 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
         manager._send_kvcache_generic = MagicMock(return_value=0)
         executor = MagicMock()
 
-        manager.send_kvcache(
+        # send_kvcache reads the memory bag (the unified-memory envelope-layout
-            "session",
+        # check), so the context has to be published. This is the non-unified
-            np.array([1], dtype=np.int32),
+        # path -- pin that explicitly rather than leaning on the default.
-            [10000, 20000, 30000],
+        with get_context().override_server_args(enable_unified_memory=False):
-            np.array([7], dtype=np.int32),
+            manager.send_kvcache(
-            executor,
+                "session",
-            dst_device_kv_indices=np.array([21], dtype=np.int32),
+                np.array([1], dtype=np.int32),
-        )
+                [10000, 20000, 30000],
+                np.array([7], dtype=np.int32),
+                executor,
+                dst_device_kv_indices=np.array([21], dtype=np.int32),
+            )
 
         kwargs = manager._send_kvcache_generic.call_args.kwargs
         self.assertEqual(kwargs["dst_device_data_ptrs"], {20000, 30000})

test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py

@@ -0,0 +1,126 @@
+"""Derived-property tests for the PD whole-envelope transfer addressing.
+
+PD disaggregation transfers the unified memory pool as whole envelopes with
+``addr = raw_ptr + physical_index * item_len`` (see
+``UnifiedMLATokenToKVPool.get_contiguous_buf_infos`` /
+``UnifiedMambaPool.get_contiguous_buf_infos`` and mooncake's
+``_send_kvcache_generic`` / ``_send_mamba_state``). That contract only holds if
+the page-major view builders keep (a) one page's data for ALL layers inside one
+contiguous ``page_envelope_bytes`` block, and (b) one mamba slot's conv+temporal
+state for all layers inside one contiguous ``entry_bytes`` block. A
+"looks equivalent" reordering of the view layout (e.g. layer-major across
+pages) would silently corrupt every PD transfer while all kernels keep working,
+because kernels read through the strided views, not through raw offsets.
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.layout.page_major import (
+    build_dense_mla_views,
+    build_page_major_mamba_views,
+    mamba_entry_bytes,
+    mla_entry_bytes,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=60, suite="base-a-test-cpu")
+
+
+class TestMLAEnvelopeTransferAddressing(CustomTestCase):
+    def test_page_envelope_matches_dense_views(self):
+        """Every (page, layer, slot) row written through the dense MLA views
+        must land at raw_ptr + page * page_envelope_bytes + layer-block offset,
+        i.e. inside the page's transfer envelope."""
+        layer_num, page_size, kv_dim, num_pages = 3, 4, 8, 6
+        store_dtype = torch.bfloat16
+        row_bytes = kv_dim * store_dtype.itemsize
+        page_bytes = page_size * layer_num * row_bytes
+        self.assertEqual(
+            page_bytes,
+            page_size
+            * mla_entry_bytes(
+                layer_num=layer_num,
+                kv_cache_dim=kv_dim,
+                itemsize=store_dtype.itemsize,
+            ),
+        )
+        # +1 page envelope of tail pad, as UnifiedKVPool allocates for MLA.
+        raw = torch.zeros((num_pages + 1) * page_bytes, dtype=torch.uint8)
+        views = build_dense_mla_views(
+            raw,
+            layer_num=layer_num,
+            kv_cache_dim=kv_dim,
+            store_dtype=store_dtype,
+            page_size=page_size,
+            num_pages=num_pages,
+            anchor_bytes=0,
+        )
+        torch.manual_seed(0)
+        for page in range(num_pages):
+            for layer in range(layer_num):
+                for off in range(page_size):
+                    dense_id = page * layer_num * page_size + off
+                    val = torch.randn(kv_dim, dtype=store_dtype)
+                    views[layer][dense_id, 0] = val
+                    start = (
+                        page * page_bytes
+                        + layer * page_size * row_bytes
+                        + off * row_bytes
+                    )
+                    got = raw[start : start + row_bytes].view(store_dtype)
+                    self.assertTrue(torch.equal(got, val), (page, layer, off))
+
+
+class TestMambaEnvelopeTransferAddressing(CustomTestCase):
+    def test_slot_envelope_is_self_contained(self):
+        """A slot's conv+temporal state for all layers must live exactly in
+        raw[slot * entry_bytes : (slot+1) * entry_bytes]: no byte outside the
+        envelope may change, and the payload byte count must fill it."""
+        layer_num, max_slots = 2, 5
+        conv_shapes = ((3, 4), (2, 6))
+        temporal_shape = (2, 3, 4)
+        conv_dtype = torch.bfloat16
+        temporal_dtype = torch.float32
+        entry = mamba_entry_bytes(
+            layer_num=layer_num,
+            conv_state_shapes=conv_shapes,
+            conv_dtype=conv_dtype,
+            temporal_state_shape=temporal_shape,
+            temporal_dtype=temporal_dtype,
+        )
+        raw = torch.zeros(max_slots * entry, dtype=torch.uint8)
+        conv_views, temporal_view = build_page_major_mamba_views(
+            raw,
+            layer_num=layer_num,
+            conv_state_shapes=conv_shapes,
+            conv_dtype=conv_dtype,
+            temporal_state_shape=temporal_shape,
+            temporal_dtype=temporal_dtype,
+            max_slots=max_slots,
+            anchor_bytes=0,
+        )
+        torch.manual_seed(0)
+        for slot in range(max_slots):
+            raw.zero_()
+            n_payload = 0
+            for i, conv_view in enumerate(conv_views):
+                val = torch.randn((layer_num,) + conv_shapes[i], dtype=conv_dtype)
+                conv_view[:, slot] = val
+                n_payload += val.numel() * val.element_size()
+            val = torch.randn((layer_num,) + temporal_shape, dtype=temporal_dtype)
+            temporal_view[:, slot] = val
+            n_payload += val.numel() * val.element_size()
+
+            outside = torch.cat([raw[: slot * entry], raw[(slot + 1) * entry :]])
+            self.assertTrue(
+                bool(outside.eq(0).all()),
+                f"slot {slot} state bled outside its transfer envelope",
+            )
+            self.assertEqual(n_payload, entry)
+
+
+if __name__ == "__main__":
+    unittest.main()