[unified-memory] PD disaggregation for every unified pool shape (#37506)

python/sglang/srt/arg_groups/fields/memory.py

@@ -82,8 +82,10 @@ class Memory(msgspec.Struct):
         "Replace the statically-partitioned hybrid-model pools (full-attn KV + "
         "SWA/Mamba state) with one byte buffer split dynamically between "
         "sub-pools. Requires the Triton attention / linear-attn / Mamba "
-        "backends; not yet compatible with PD disaggregation or speculative "
-        "decoding.",
+        "backends. PD disaggregation is supported over mooncake at equal "
+        "attention TP with pp=1; not yet compatible with hierarchical / "
+        "host-tiered KV cache, prefill cuda-graph capture, or speculative "
+        "decoding other than DSPARK.",
     ] = False
     enable_session_radix_cache: A[
         bool,

python/sglang/srt/disaggregation/decode.py

@@ -1416,7 +1416,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                     decode_req.req.kv.req_pool_idx, window_start:seq_len
                 ]
                 window_kv_indices_swa = (
-                    self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
+                    self.token_to_kv_pool_allocator.translate_swa_indices_for_transfer(
                         window_kv_indices_full
                     )
                 )

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

@@ -957,6 +957,13 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
             prefill_data_indices=prefill_kv_indices,
             dst_data_indices=dst_kv_indices,
             executor=executor,
+            # The unified pool registers ONE region holding every layer's K and
+            # V inside each page envelope. The MHA branch would half-split that
+            # single region into K and V halves and compute num_kv_layers = 0,
+            # transferring nothing at all; the flat branch addresses the region
+            # as-is. MLA-unified already reaches the flat branch via
+            # is_mla_backend, so this only adds the MHA-unified peer.
+            force_flat=get_memory().enable_unified_memory,
             src_layer_ids=self.kv_args.kv_layer_ids,
             dst_layer_ids=dst_layer_ids,
             dst_device_data_indices=dst_device_kv_indices,
@@ -1598,8 +1605,16 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                         dst_data_indices=np.array(dst_indices_local, dtype=np.int32),
                         executor=executor,
                         state_type=st,
-                        force_flat=st
-                        in (StateType.QSA_PENDING, StateType.QSA_COMPRESSED),
+                        # Two independent reasons to keep the flat layout.
+                        # QSA's per-layer list must not be half-split into K/V;
+                        # neither must a unified sub-pool's single region, which
+                        # holds every layer's K and V per slot envelope -- the
+                        # MHA branch would compute zero layers and ship nothing
+                        # (same reason as in `send_kvcache`).
+                        force_flat=(
+                            st in (StateType.QSA_PENDING, StateType.QSA_COMPRESSED)
+                            or get_memory().enable_unified_memory
+                        ),
                         src_layer_ids=src_state_layer_ids,
                         dst_layer_ids=dst_state_layer_ids,
                     )

python/sglang/srt/disaggregation/prefill.py

@@ -1328,7 +1328,7 @@ class SchedulerDisaggregationPrefillMixin:
                     req.kv.req_pool_idx, window_start:seq_len
                 ]
                 window_kv_indices_swa = (
-                    self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
+                    self.token_to_kv_pool_allocator.translate_swa_indices_for_transfer(
                         window_kv_indices_full
                     )
                 )

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

@@ -346,6 +346,17 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
         return self.logical_attn_allocator.translate_loc_from_full_to_swa(kv_indices)
 
+    def translate_swa_indices_for_transfer(
+        self, kv_indices: torch.Tensor
+    ) -> torch.Tensor:
+        # Delegated like the read-path translate above: this composite is not a
+        # SWA allocator itself, so it inherits neither the default nor an
+        # override, and the PD payload path calls this on whatever allocator
+        # the scheduler holds.
+        return self.logical_attn_allocator.translate_swa_indices_for_transfer(
+            kv_indices
+        )
+
     def full_available_size(self):
         return min(
             self.logical_attn_allocator.full_available_size(),

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

@@ -203,6 +203,19 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         assert self._kvcache.full_to_swa_index_mapping is not None
         return self._kvcache.translate_loc_from_full_to_swa(kv_indices)
 
+    def translate_swa_indices_for_transfer(
+        self, kv_indices: torch.Tensor
+    ) -> torch.Tensor:
+        """Sliding-window token ids as the PD transfer engine addresses them.
+
+        The sibling of `translate_kv_indices_for_transfer` for the SWA state
+        component. On a static pool the sliding-window buffers are indexed by
+        the same ids the kernels use, so the read-path translate IS the answer.
+        A virtual-id pool must override: the transfer addresses raw bytes and
+        needs PHYSICAL ids, not kernel-facing ones.
+        """
+        return self.translate_loc_from_full_to_swa(kv_indices)
+
     def alloc(self, need_size: int):
         assert self.page_size == 1
         if need_size > self.full_attn_allocator.available_size():

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

@@ -17,7 +17,7 @@ sub-pools of one `UnifiedKVPool`, and the tri-pool variant that adds mamba state
 from __future__ import annotations
 
 import logging
-from typing import List, Optional, Sequence
+from typing import Callable, List, Optional, Sequence, Tuple
 
 import torch
 from torch.profiler import record_function
@@ -34,6 +34,7 @@ from sglang.srt.mem_cache.allocator.unified_sub_pool import (
     _flush_deferred_free_group,
     _full_tokens_before_mamba_recheck,
     _relieve_for_alloc,
+    install_move_gate,
 )
 from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
 from sglang.srt.utils.common import get_num_new_pages
@@ -320,6 +321,46 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         """Page-level physical->virtual table of the full sub-pool."""
         return self.full_attn_allocator.physical_to_virtual
 
+    def translate_kv_indices_for_transfer(
+        self, kv_indices: torch.Tensor
+    ) -> torch.Tensor:
+        """Virtual TOKEN ids -> full-sub-pool PHYSICAL token ids for the PD
+        transfer engine.
+
+        PHYSICAL, not kernel-facing: the transfer registers page ENVELOPES (see
+        `UnifiedMHATokenToKVPool.get_contiguous_buf_infos`). Without this
+        override the base identity would put VIRTUAL ids on the wire, which
+        address real bytes and so corrupt silently rather than fail.
+        """
+        return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
+
+    def translate_swa_indices_for_transfer(
+        self, kv_indices: torch.Tensor
+    ) -> torch.Tensor:
+        """Virtual TOKEN ids -> swa-sub-pool PHYSICAL token ids.
+
+        The SWA counterpart of the above. `translate_loc_from_full_to_swa`
+        cannot serve here: it returns KERNEL-FACING ids (the physical page
+        scaled by the sub-pool's per-page block count), which index the
+        per-layer views, whereas the SWA state component is registered as whole
+        page envelopes and addressed by physical page.
+        """
+        return self.swa_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
+
+    def _move_gate_targets(self):
+        """Every member a compaction gate must cover. A subclass that adds an
+        end overrides THIS, and every gate widens with it."""
+        return (self.full_attn_allocator, self.swa_attn_allocator)
+
+    def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None:
+        install_move_gate(
+            self._move_gate_targets(),
+            slot="disagg_move_gate",
+            gate=gate,
+            feature="PD disaggregation",
+            lazy_compaction=self.lazy_compaction,
+        )
+
     def translate_kv_loc_for_kernel(
         self,
         loc: torch.Tensor,
@@ -374,6 +415,52 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
             self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages)
             return v_tokens
 
+    def _extend_in_virtual_space(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+    ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+        """Run the full side's paged extend and report which virtual PAGES it
+        newly took. Returns (virtual TOKEN ids, new virtual PAGE ids), or None
+        when the joint capacity check cannot fund the allocation.
+
+        Both extend entries share this; they differ only in which of those pages
+        the sliding-window side then binds.
+        """
+        num_new_pages = get_num_new_pages(
+            seq_lens=seq_lens_cpu,
+            page_size=self.page_size,
+            prefix_lens=prefix_lens_cpu,
+        )
+        need_tokens = num_new_pages * self.page_size
+        if need_tokens > self.available_size():
+            if not _relieve_for_alloc(self, need_tokens):
+                return None
+
+        # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps
+        # its view after the slice is consumed.
+        fa = self.full_attn_allocator
+        new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone()
+
+        out_indices = fa.alloc_extend(
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc,
+            extend_num_tokens,
+            num_new_pages=num_new_pages,
+        )
+        assert out_indices is not None, (
+            "UnifiedSWA: full.alloc_extend returned None after joint pre-check "
+            "passed — internal-state inconsistency"
+        )
+        return out_indices, new_virtual_pages
+
     def alloc_extend(
         self,
         prefix_lens: torch.Tensor,
@@ -386,37 +473,75 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         """Paged extend; returns virtual TOKEN ids. The same virtual page maps to
         full- and swa-physical, so swa binds exactly what the full kernel consumed."""
         with record_function("UnifiedSWAAlloc.alloc_extend"):
-            num_new_pages = get_num_new_pages(
-                seq_lens=seq_lens_cpu,
-                page_size=self.page_size,
-                prefix_lens=prefix_lens_cpu,
-            )
-            need_tokens = num_new_pages * self.page_size
-            if need_tokens > self.available_size():
-                if not _relieve_for_alloc(self, need_tokens):
-                    return None
-
-            # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps
-            # its view after the slice is consumed.
-            fa = self.full_attn_allocator
-            new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone()
-
-            out_indices = fa.alloc_extend(
+            extended = self._extend_in_virtual_space(
                 prefix_lens,
                 prefix_lens_cpu,
                 seq_lens,
                 seq_lens_cpu,
                 last_loc,
                 extend_num_tokens,
-                num_new_pages=num_new_pages,
-            )
-            assert out_indices is not None, (
-                "UnifiedSWA.alloc_extend: full.alloc_extend returned None "
-                "after joint pre-check passed — internal-state inconsistency"
             )
+            if extended is None:
+                return None
+            out_indices, new_virtual_pages = extended
             self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages)
             return out_indices  # virtual TOKEN ids
 
+    def alloc_extend_swa_tail(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+        swa_tail_len: int,
+    ) -> Optional[torch.Tensor]:
+        """Decode-node prealloc: full KV for the whole sequence, sliding-window
+        KV for the live window tail only.
+
+        The static composite allocates the two sides independently and records
+        a full->swa index mapping. That is not representable here: the two
+        sides SHARE one virtual id space (a virtual page names a full-physical
+        page and, if bound, a swa-physical one), which is why
+        `set_full_to_swa_mapping` is a no-op on this allocator and
+        `translate_loc_from_full_to_swa` derives the swa id from the virtual id
+        instead of a table. Running the static body would call `alloc_extend`
+        on the swa sub-allocator, which asserts it is not the id owner.
+
+        The tail is expressed by binding swa for the TAIL's virtual pages only.
+        A new page left unbound has no swa-physical page, which reads as the
+        sink and is skipped by `free`'s `swa_v2p_page > 0` mask -- exactly the
+        out-of-window state the ratchet produces via `free_swa`.
+
+        Admission is priced at the FULL side's page count, as plain
+        `alloc_extend` is: pessimistic when the tail is short, but it reuses
+        the composite's audited joint capacity path, and the bytes actually
+        held still follow the tail.
+        """
+        assert len(prefix_lens_cpu) == 1
+        assert 0 <= swa_tail_len <= extend_num_tokens
+        with record_function("UnifiedSWAAlloc.alloc_extend_swa_tail"):
+            extended = self._extend_in_virtual_space(
+                prefix_lens,
+                prefix_lens_cpu,
+                seq_lens,
+                seq_lens_cpu,
+                last_loc,
+                extend_num_tokens,
+            )
+            if extended is None:
+                return None
+            out_indices, new_virtual_pages = extended
+            if swa_tail_len > 0 and new_virtual_pages.numel() > 0:
+                tail_pages = torch.unique(out_indices[-swa_tail_len:] // self.page_size)
+                # Only NEW pages need binding; a tail page carried in from the
+                # prefix is already bound on the swa side.
+                to_bind = new_virtual_pages[torch.isin(new_virtual_pages, tail_pages)]
+                if to_bind.numel() > 0:
+                    self.swa_attn_allocator.alloc_with_virtual(to_bind)
+            return out_indices  # virtual TOKEN ids
+
     def alloc_decode(
         self,
         seq_lens: torch.Tensor,
@@ -769,16 +894,9 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator):
         binary search; the order matches the alloc path (full takes the high band).
         """
         fa, sa = self.full_attn_allocator, self.swa_attn_allocator
-        e_f, e_s = fa.entry_bytes_per_page, sa.entry_bytes_per_page
+        e_f = fa.entry_bytes_per_page
         # full is grow-down: its chain gap IS the high band.
         b_high = fa._current_gap_bytes()
-        if sa._is_frontier_transparent():
-            b_low = 0
-        else:
-            b_low = max(
-                0,
-                sa._byte_low_frontier() - sa._chain_high_frontier_below_bytes(),
-            )
         h_f = len(fa._free_phys_pages) if fa.lazy_compaction else 0
         h_s = sa._hole_pages()
         r_f = fa.num_pages - fa.min_page_index - fa._allocated_pages()
@@ -819,6 +937,16 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator):
                 hi_n = mid - 1
         return lo_n * self.page_size
 
+    def _move_gate_targets(self):
+        """All three members. The mamba end compacts independently and its slot
+        envelopes move as `StateType.MAMBA`, so leaving it out of a gate would
+        let a conv/SSM slot relocate under an in-flight transfer."""
+        return (
+            self.full_attn_allocator,
+            self.swa_attn_allocator,
+            self.mamba_allocator,
+        )
+
     def _flush_targets(self):
         """All three members, float FIRST: its zero-copy boundary absorption must
         land before the deficit math prices a relocation it already covered."""

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

@@ -31,6 +31,7 @@ from sglang.srt.mem_cache.allocator.unified_sub_pool import (
     _flush_deferred_free_group,
     _full_tokens_before_mamba_recheck,
     _relieve_for_alloc,
+    install_move_gate,
 )
 from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
 from sglang.srt.runtime_context import get_parallel
@@ -306,15 +307,20 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         )
         return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
 
+    def _move_gate_targets(self):
+        """Every member a compaction gate must cover. The mamba end is gated
+        even where its state is not itself transferred: the gate is about the
+        MOVER, and the two ends compact as peers."""
+        return (self.full_attn_allocator, self.mamba_allocator)
+
     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)."
+        install_move_gate(
+            self._move_gate_targets(),
+            slot="disagg_move_gate",
+            gate=gate,
+            feature="PD disaggregation",
+            lazy_compaction=self.lazy_compaction,
         )
-        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/allocator/unified_sub_pool.py

@@ -239,6 +239,31 @@ def _full_tokens_before_mamba_recheck(
     return -(-minimum_missing_bytes * dcp_size // full_allocator.entry_bytes)
 
 
+def install_move_gate(
+    targets,
+    *,
+    slot: str,
+    gate: Callable[[], bool],
+    feature: str,
+    lazy_compaction: bool,
+) -> None:
+    """Point every member of a composite at one compaction gate.
+
+    A gate that reaches only some members is not a weaker gate, it is no gate:
+    the ungated end relocates its own pages under the same in-flight transfer.
+    So the member list is stated once per composite (`_move_gate_targets`) and
+    every gate installs over it, rather than each setter naming the members it
+    happens to remember.
+    """
+    assert lazy_compaction, (
+        f"{feature} with the unified memory pool requires lazy compaction "
+        "(eager free-path compaction moves pages under in-flight transfers)."
+    )
+    assert slot in ("disagg_move_gate", "host_transfer_move_gate"), slot
+    for target in targets:
+        setattr(target, slot, gate)
+
+
 class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
     """Allocator for one sub-pool over a `UnifiedKVPool`."""
 

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -155,11 +155,21 @@ def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str:
         )
         # Host-pool retraction transfers full and sliding-window components
         # only, so a model with recurrent state stays on cpu_tensor.
-        supports_host_pool = not uses_ssm_state(
-            tp_worker.model_runner.model_config
-        ) and (
-            isinstance(kv_cache, MHATokenToKVPool)
-            or (isinstance(kv_cache, SWAKVPool) and full_tokens_per_layer > 0)
+        #
+        # The unified pool is excluded for the same reason hierarchical cache is
+        # (see `handle_unified_memory_pool`): the host-transfer path indexes the
+        # device buffers with the ids it is handed, and under the unified pool
+        # those are VIRTUAL. It also cannot be sized from `kv_cache.size`, which
+        # is a KERNEL-FACING row count (`num_pages * 2 * layer_num * page_size`)
+        # rather than a token capacity -- gpt-oss-20b reports 85M "tokens" and
+        # asks for 418 GB of host memory per component.
+        supports_host_pool = (
+            not uses_ssm_state(tp_worker.model_runner.model_config)
+            and not memory.enable_unified_memory
+            and (
+                isinstance(kv_cache, MHATokenToKVPool)
+                or (isinstance(kv_cache, SWAKVPool) and full_tokens_per_layer > 0)
+            )
         )
         schedule = get_schedule()
         priority_preemption = (

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -446,24 +446,12 @@ class KVCacheConfigurator:
         # 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 get_memory().enable_unified_memory and req_to_token_pool is None:
-            pd_enabled = get_disagg().disaggregation_mode != "null"
             is_dsv4 = is_deepseek_v4(self.model_config.hf_config)
             # Order matters: an Inkling-class model is BOTH mambaish and
             # hybrid-SWA, and the mamba pair would store every SWA layer's KV at
             # FULL lifetime -- its branch reads the HF config's
             # full_attention_layer_ids, which for Inkling is ALL layers.
             if self.mambaish_config is not None and self.is_hybrid_swa and not is_dsv4:
-                if pd_enabled:
-                    # Same limitation as the 2-pool SWA branch below: the
-                    # tri-pool carries an SWA sub-pool, and there is no
-                    # whole-envelope transfer scheme for it.
-                    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); this model "
-                        "routes to the mamba+SWA tri-pool, which has one. Drop "
-                        "--enable-unified-memory or run without PD."
-                    )
                 bundle = self._init_unified_mamba_swa_pools(
                     max_num_reqs=sizes.max_running_requests,
                     full_max_total_num_tokens=sizes.full_max_total_num_tokens,
@@ -471,27 +459,12 @@ class KVCacheConfigurator:
                     unified_total_bytes=sizes.unified_total_bytes,
                 )
             elif 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,
                     unified_total_bytes=sizes.unified_total_bytes,
                 )
             elif self.is_hybrid_swa and not is_dsv4:
-                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,
@@ -835,6 +808,13 @@ class KVCacheConfigurator:
             unified_total_bytes=(None if self.is_draft_worker else unified_total_bytes),
             # bs=1 feasibility floor input (context len is already passed).
             sliding_window_size=self.model_config.sliding_window_size,
+            # Decode nodes hand out request rows to PREALLOCATED transfers on
+            # top of the running set; the 2-pool mamba factory takes the same.
+            decode_pre_alloc_size=(
+                get_disagg().disaggregation_decode_extra_slots
+                if get_disagg().disaggregation_mode == "decode"
+                else 0
+            ),
         )
 
     def _init_unified_swa_pools(
@@ -863,12 +843,28 @@ class KVCacheConfigurator:
         extra_max_context_len = 4
         if get_spec().speculative_num_draft_tokens is not None:
             extra_max_context_len += get_spec().speculative_num_draft_tokens
-        req_to_token_pool = ReqToTokenPool(
-            size=max_num_reqs,
-            max_context_len=self.model_config.context_len + extra_max_context_len,
-            device=self.device,
-            enable_memory_saver=get_exec().features.enable_memory_saver,
-        )
+        if get_disagg().disaggregation_mode == "decode":
+            # A decode node hands out request rows to PREALLOCATED transfers on
+            # top of its running set, so it needs the extra-slot pool (and the
+            # `pre_alloc_size` the scheduler's invariant checker reads). Mirrors
+            # `_build_req_to_token_pool`'s decode branch; the mamba composite
+            # already takes `decode_pre_alloc_size` the same way.
+            from sglang.srt.disaggregation.decode import DecodeReqToTokenPool
+
+            req_to_token_pool = DecodeReqToTokenPool(
+                size=max_num_reqs,
+                max_context_len=self.model_config.context_len + extra_max_context_len,
+                device=self.device,
+                enable_memory_saver=get_exec().features.enable_memory_saver,
+                pre_alloc_size=get_disagg().disaggregation_decode_extra_slots,
+            )
+        else:
+            req_to_token_pool = ReqToTokenPool(
+                size=max_num_reqs,
+                max_context_len=self.model_config.context_len + extra_max_context_len,
+                device=self.device,
+                enable_memory_saver=get_exec().features.enable_memory_saver,
+            )
 
         head_num = self.model_config.get_num_kv_heads(
             get_parallel().attn_tp_size, get_parallel().attn_dcp_size

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -624,10 +624,19 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool):
             env[tgt_pages] = env[src_pages]
 
     def get_contiguous_buf_infos(self):
-        raise NotImplementedError(
-            "unified layout has no per-layer contiguous regions; "
-            "KV transfer / disaggregation is unsupported."
-        )
+        """PD-transfer registration: ONE entry, the raw buffer, addressed as
+        ``raw_ptr + physical_page_id * page_envelope_bytes``.
+
+        Same whole-envelope contract as `UnifiedMLATokenToKVPool`: the transfer
+        item is one page across ALL layers and both K and V, because the
+        per-layer views overlap inside the envelope and index in kernel-facing
+        ids. A peer must therefore build an identical spec -- enforced on the
+        wire by `_validate_envelope_kv_layout`.
+        """
+        # 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 get_cpu_copy(self, indices, mamba_indices=None):
         raise NotImplementedError(
@@ -1882,6 +1891,7 @@ def init_unified_mamba_swa_pools(
     lazy_compaction: bool = False,
     unified_total_bytes: Optional[int] = None,
     sliding_window_size: Optional[int] = None,
+    decode_pre_alloc_size: int = 0,
 ) -> UnifiedPoolBundle:
     """Build the TRI-pool unified-memory-pool stack for models with full KV +
     SWA KV + mamba/conv state (Inkling-class: `mambaish_config` AND
@@ -2007,6 +2017,7 @@ def init_unified_mamba_swa_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,
     )
     allocator = UnifiedMambaSWATokenToKVPoolAllocator(
         unified_buffer=shared_pool,