Support unified memory page-envelope transfers in PD

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

@@ -82,10 +82,8 @@ 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. 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.",
+        "backends. Supported PD-disaggregation and speculative-decoding "
+        "configurations are validated at startup.",
     ] = False
     enable_session_radix_cache: A[
         bool,

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -200,18 +200,26 @@ def handle_unified_memory_pool(server_args: Any) -> None:
     if not cfg.enable_unified_memory:
         return
     if cfg.disaggregation_mode != "null":
-        # Constraints of the whole-envelope transfer; see
-        # UnifiedMLATokenToKVPool.get_contiguous_buf_infos.
-        assert cfg.disaggregation_transfer_backend == "mooncake", (
-            "--enable-unified-memory with PD disaggregation supports only "
-            "the mooncake transfer backend; got "
+        # Constraints of the whole-envelope transfer; see the unified MHA and
+        # MLA pool get_contiguous_buf_infos implementations.
+        supported_backends = server_args._unified_memory_pd_transfer_backends()
+        assert cfg.disaggregation_transfer_backend in supported_backends, (
+            "--enable-unified-memory with PD disaggregation supports only these "
+            f"transfer backends: {', '.join(sorted(supported_backends))}; got "
             f"{cfg.disaggregation_transfer_backend!r}."
         )
         assert cfg.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 (
+            cfg.disaggregation_transfer_backend == "mooncake"
+            and cfg.speculative_algorithm is not None
+        ), (
+            "--enable-unified-memory with PD disaggregation does not support "
+            "speculative decoding with the Mooncake transfer backend."
+        )
         assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), (
             "--enable-unified-memory with PD disaggregation requires lazy "
             "compaction; unset SGLANG_DISABLE_LAZY_COMPACTION."
@@ -222,6 +230,15 @@ def handle_unified_memory_pool(server_args: Any) -> None:
             "ships host/C4 rows straight from the allocator, bypassing the "
             "virtual->physical translation the unified pool needs."
         )
+        assert cfg.disaggregation_decode_retraction_backup != "host_pool", (
+            "--enable-unified-memory with PD disaggregation does not support "
+            "--disaggregation-decode-retraction-backup=host_pool; use "
+            "cpu_tensor (the automatic default for unified pools)."
+        )
+        assert not cfg.disaggregation_decode_enable_offload_kvcache, (
+            "--enable-unified-memory with PD disaggregation does not yet support "
+            "--disaggregation-decode-enable-offload-kvcache."
+        )
     assert cfg.speculative_algorithm in (None, "DSPARK"), (
         "--enable-unified-memory only supports --speculative-algorithm "
         "DSPARK (chain draft); other speculative algorithms are not yet "

python/sglang/srt/arg_groups/overrides.py

@@ -1808,6 +1808,9 @@ def post_capture_kv_sizing_planned(server_args: Any) -> bool:
     mla_enabled = use_mla_backend(server_args)
     if not envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get():
         return False
+    # Unified arenas are fully backed before capture and cannot resize afterward.
+    if cfg.enable_unified_memory:
+        return False
     if cfg.device != "cuda":
         return False
     if cfg.dcp_size != 1: