[diffusion] feat: disaggregated diffusion (#21701)

docs/diffusion/disaggregation.md

@@ -0,0 +1,237 @@
+# Disaggregated Diffusion Pipeline
+
+Split a monolithic text-to-video/image pipeline into independent **Encoder**, **Denoiser**, and **Decoder** roles, each running on its own GPU(s). A central **DiffusionServer** routes requests through the pipeline.
+
+## Quick Start
+
+Disaggregation is controlled by a single flag: `--disagg-role`. Each component is launched independently, just like LLM PD disaggregation.
+
+| `--disagg-role` | What it runs |
+|----------------|--------------|
+| `monolithic` | (Default) Standard single-server mode |
+| `encoder` | All stages with the default `RoleType.ENCODER` affinity: `InputValidationStage`, `TextEncodingStage` (plus `ImageEncodingStage` / `ImageVAEEncodingStage` for image-conditioned pipelines), `LatentPreparationStage`, `TimestepPreparationStage`, and any model-specific "before denoising" stage (e.g. `QwenImageLayeredBeforeDenoisingStage`, `GlmImageBeforeDenoisingStage`). |
+| `denoiser` | `DenoisingStage` (and its subclasses: `CausalDMDDenoisingStage`, `DmdDenoisingStage`, `LTX2AVDenoisingStage`, `LTX2RefinementStage`, `Hunyuan3DShapeDenoisingStage`, ...) — the DiT forward loop plus the scheduler stepping it drives. |
+| `decoder` | `DecodingStage` (VAE decode) and its subclasses (`LTX2AVDecodingStage`, `HeliosDecodingStage`, ...). |
+| `server` | DiffusionServer head node + HTTP server (no GPU) |
+
+> Each stage declares its role via the `role_affinity` property on `PipelineStage` (default `ENCODER`). When `--disagg-role` is not `monolithic`, the pipeline only instantiates stages whose affinity matches, so the above table is the source of truth for what actually runs in each process.
+
+### Single-Machine Example (Verified)
+
+The following commands have been tested end-to-end on an 8×H200 machine with
+`Wan-AI/Wan2.1-T2V-1.3B-Diffusers`. Each role runs on a separate GPU via
+`--base-gpu-id`; the `server` head node requires no GPU.
+
+```bash
+# Terminal 1: Encoder (GPU 0)
+sglang serve --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --disagg-role encoder \
+    --disagg-server-addr tcp://127.0.0.1:19655 \
+    --scheduler-port 19000 \
+    --num-gpus 1 --base-gpu-id 0
+
+# Terminal 2: Denoiser (GPU 1)
+sglang serve --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --disagg-role denoiser \
+    --disagg-server-addr tcp://127.0.0.1:19655 \
+    --scheduler-port 19001 \
+    --num-gpus 1 --base-gpu-id 1
+
+# Terminal 3: Decoder (GPU 2)
+sglang serve --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --disagg-role decoder \
+    --disagg-server-addr tcp://127.0.0.1:19655 \
+    --scheduler-port 19002 \
+    --num-gpus 1 --base-gpu-id 2
+
+# Terminal 4: DiffusionServer head (no GPU, receives HTTP requests)
+sglang serve --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --disagg-role server \
+    --encoder-urls  "tcp://127.0.0.1:19000" \
+    --denoiser-urls "tcp://127.0.0.1:19001" \
+    --decoder-urls  "tcp://127.0.0.1:19002" \
+    --host 0.0.0.0 --port 22000 \
+    --scheduler-port 19655
+
+# Send request (video generation)
+curl http://127.0.0.1:22000/v1/videos \
+    -H "Content-Type: application/json" \
+    -d '{"model": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "prompt": "A curious raccoon exploring a garden, cinematic", "size": "832x480"}'
+```
+
+> **Tested result (8×H200):**
+> Encoder 2.3 s (TextEncoding) → Denoiser 312.8 s (50 steps, layerwise offload) → Decoder 7.1 s (VAE decode).
+> Total ~322 s for 81-frame 1024×1024 video.
+
+> **Tip:** `--base-gpu-id` controls which physical GPU the role uses.
+> Encoder and Decoder can share a GPU (e.g. both `--base-gpu-id 0`) to save resources,
+> but make sure the combined GPU memory is sufficient.
+
+### Multi-Machine Example
+
+The exact same CLI pattern — just replace `127.0.0.1` with actual IPs and add
+RDMA flags for direct transfer:
+
+```bash
+# Machine A (10.0.0.1): Encoder
+sglang serve --model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \
+    --disagg-role encoder \
+    --disagg-server-addr tcp://10.0.0.4:19655 \
+    --scheduler-port 19000 \
+    --num-gpus 1 \
+    --disagg-p2p-hostname 10.0.0.1 --disagg-ib-device mlx5_0
+
+# Machine B (10.0.0.2): Denoiser (4 GPUs with SP)
+sglang serve --model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \
+    --disagg-role denoiser \
+    --disagg-server-addr tcp://10.0.0.4:19655 \
+    --scheduler-port 19001 \
+    --num-gpus 4 --denoiser-sp 4 --denoiser-ulysses 2 --denoiser-ring 2 \
+    --disagg-p2p-hostname 10.0.0.2 --disagg-ib-device mlx5_0
+
+# Machine C (10.0.0.3): Decoder
+sglang serve --model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \
+    --disagg-role decoder \
+    --disagg-server-addr tcp://10.0.0.4:19655 \
+    --scheduler-port 19002 \
+    --num-gpus 1 \
+    --disagg-p2p-hostname 10.0.0.3 --disagg-ib-device mlx5_0
+
+# Machine D (10.0.0.4): DiffusionServer head
+sglang serve --model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \
+    --disagg-role server \
+    --encoder-urls  "tcp://10.0.0.1:19000" \
+    --denoiser-urls "tcp://10.0.0.2:19001" \
+    --decoder-urls  "tcp://10.0.0.3:19002" \
+    --host 0.0.0.0 --port 30000 \
+    --scheduler-port 19655 \
+    --disagg-dispatch-policy max_free_slots
+```
+
+> ZMQ handles startup order gracefully — instances and head can start in any order.
+
+## Multiple Instances per Role
+
+Use semicolons in `--*-urls` to register multiple instances:
+
+```bash
+# 2 encoders + 2 denoisers (4-GPU SP each) + 1 decoder
+sglang serve --model-path ... --disagg-role server \
+    --encoder-urls  "tcp://10.0.0.1:35000;tcp://10.0.0.2:35000" \
+    --denoiser-urls "tcp://10.0.0.3:35000;tcp://10.0.0.4:35000" \
+    --decoder-urls  "tcp://10.0.0.5:35000"
+```
+
+## Port Convention
+
+Result endpoints are derived deterministically from the head node's `--scheduler-port` (default: 5555):
+
+| Socket | Port |
+|--------|------|
+| DS frontend (ROUTER) | `scheduler_port` |
+| Encoder result (PULL) | `scheduler_port + 1` |
+| Denoiser result (PULL) | `scheduler_port + 2` |
+| Decoder result (PULL) | `scheduler_port + 3` |
+
+Role instances derive their result endpoint automatically from `--disagg-server-addr`. No manual endpoint configuration needed.
+
+## Transfer Mechanism
+
+Tensor data between roles (encoder→denoiser, denoiser→decoder) is transferred via a P2P transfer engine. The DiffusionServer only routes lightweight control messages (alloc/push/ready); actual tensor data flows directly between instances.
+
+**mooncake-transfer-engine** is required for disaggregated diffusion. It provides RDMA for direct GPU-to-GPU data movement.
+
+```bash
+pip install mooncake-transfer-engine
+```
+
+### Transfer Flow
+
+1. **Sender** (encoder/denoiser) stages tensors: async copy to transfer buffer (GPU or CPU pinned, depending on GPUDirect support), overlapped with metadata JSON serialization.
+2. **Sender** sends `transfer_staged` control message to DiffusionServer (metadata only, no tensor data).
+3. **DiffusionServer** sends `transfer_alloc` to receiver → receiver allocates buffer slot → replies `transfer_allocated`.
+4. **DiffusionServer** sends `transfer_push` to receiver with sender's address info.
+5. **Receiver** pulls data via transfer engine (Mooncake RDMA or mock), sends `transfer_ready`.
+6. **Receiver** loads tensors async on a dedicated transfer stream, overlapped with the previous request's compute.
+
+Decoder results (final output) flow back through DiffusionServer as raw ZMQ frames to the HTTP client.
+
+### RDMA Flags
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--disagg-p2p-hostname` | `127.0.0.1` | RDMA-reachable hostname/IP of this instance |
+| `--disagg-ib-device` | `None` | InfiniBand device (e.g., `mlx5_0`, `mlx5_roce0`) |
+| `--disagg-transfer-pool-size` | 256 MiB | Pinned memory pool per instance |
+
+Set `--disagg-p2p-hostname` to the actual IP on each machine. For multi-machine, `--disagg-ib-device` specifies the RDMA NIC.
+
+## Per-Role Parallelism
+
+| Flag | Description |
+|------|-------------|
+| `--encoder-tp` | Encoder tensor parallelism |
+| `--denoiser-tp` / `--denoiser-sp` / `--denoiser-ulysses` / `--denoiser-ring` | Denoiser parallelism |
+| `--decoder-tp` | Decoder tensor parallelism |
+
+If not specified, parallelism is auto-derived from `--num-gpus`.
+
+## Other Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--disagg-timeout` | `600` | Timeout (seconds) for pending requests |
+| `--disagg-dispatch-policy` | `round_robin` | `round_robin` or `max_free_slots` |
+
+## Python API
+
+For programmatic single-machine deployment, `launch_pool_disagg_server()` is available:
+
+```python
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.launch_server import launch_pool_disagg_server
+
+server_args = ServerArgs.from_kwargs(
+    model_path="Wan-AI/Wan2.1-T2V-14B-Diffusers",
+    denoiser_sp=4, denoiser_ulysses=2, denoiser_ring=2,
+    disagg_ib_device="mlx5_0",
+)
+
+launch_pool_disagg_server(
+    server_args,
+    encoder_gpus=[[0]],
+    denoiser_gpus=[[1, 2, 3, 4], [5, 6, 7, 8]],
+    decoder_gpus=[[0]],
+)
+```
+
+## Architecture
+
+```
+Client ─── HTTP (port 30000) ──► FastAPI Server
+                                      │
+                                      ▼
+                              DiffusionServer (ROUTER, scheduler_port)
+                              ┌───────┼───────┐
+                   PUSH work  │       │       │  PUSH work
+                              ▼       │       ▼
+                    Encoder[0..N]     │    Decoder[0..K]
+                              │       │       ▲
+                   P2P tensor │       │       │ P2P tensor
+                   transfer   ▼       │       │ transfer
+                          Denoiser[0..M] ─────┘
+                                      │
+                    PULL results ◄────┘  (decoder → DS → client)
+```
+
+### Request State Machine
+
+```
+PENDING → ENCODER_WAITING → ENCODER_RUNNING → ENCODER_DONE
+                                                    │
+                        DENOISING_WAITING → DENOISING_RUNNING → DENOISING_DONE
+                                                                       │
+                                    DECODER_WAITING → DECODER_RUNNING → DONE
+```
+
+Any state can transition to `FAILED` or `TIMED_OUT`.

python/sglang/multimodal_gen/runtime/disaggregation/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Disaggregation support for diffusion pipelines."""

python/sglang/multimodal_gen/runtime/disaggregation/disagg_args.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Disaggregated diffusion CLI arguments and helper methods.
+
+All disagg-related dataclass fields, argparse registration, and endpoint
+derivation logic live here.  ``ServerArgs`` inherits from
+``DisaggArgsMixin`` so the fields appear on the top-level config object.
+"""
+
+from __future__ import annotations
+
+import argparse
+from typing import TYPE_CHECKING
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+
+if TYPE_CHECKING:
+    pass
+
+# ── Port offsets for disagg result endpoints (deterministic convention) ──
+DISAGG_RESULT_PORT_OFFSETS: dict[RoleType, int] = {
+    RoleType.ENCODER: 1,
+    RoleType.DENOISER: 2,
+    RoleType.DECODER: 3,
+}
+
+
+class DisaggArgsMixin:
+    """Methods for disaggregated diffusion, mixed into ``ServerArgs``.
+
+    The dataclass **fields** remain in ``ServerArgs`` (to avoid MRO
+    ordering issues with ``@dataclass`` inheritance).  This mixin only
+    provides the methods that operate on those fields.
+    """
+
+    def get_role_parallelism(self, role_type: RoleType) -> dict[str, int | None]:
+        """Return per-role parallelism overrides for the given role.
+
+        Returns a dict with keys tp_size, sp_degree, ulysses_degree,
+        ring_degree.  Values are ``None`` when not explicitly set
+        (auto-derive from ``num_gpus``).
+        """
+        _none: dict[str, int | None] = {
+            "tp_size": None,
+            "sp_degree": None,
+            "ulysses_degree": None,
+            "ring_degree": None,
+        }
+        if role_type == RoleType.ENCODER:
+            return {**_none, "tp_size": self.encoder_tp}
+        elif role_type == RoleType.DENOISER:
+            return {
+                "tp_size": self.denoiser_tp,
+                "sp_degree": self.denoiser_sp,
+                "ulysses_degree": self.denoiser_ulysses,
+                "ring_degree": self.denoiser_ring,
+            }
+        elif role_type == RoleType.DECODER:
+            return {**_none, "tp_size": self.decoder_tp}
+        return _none
+
+    def derive_pool_result_endpoint(self) -> str:
+        """Derive the result PUSH endpoint from ``disagg_server_addr`` + role.
+
+        Convention: DS binds result PULL on ``scheduler_port + {1,2,3}``
+        for encoder / denoiser / decoder.
+        """
+        if self.disagg_server_addr is None:
+            raise ValueError("disagg_server_addr is required for per-role launch")
+        addr = self.disagg_server_addr
+        if addr.startswith("tcp://"):
+            addr = addr[len("tcp://") :]
+        host, port_str = addr.rsplit(":", 1)
+        base_port = int(port_str)
+        offset = DISAGG_RESULT_PORT_OFFSETS[self.disagg_role]
+        return f"tcp://{host}:{base_port + offset}"
+
+    def derive_pool_work_endpoint(self) -> str:
+        """Derive the work PULL bind endpoint for a standalone role instance."""
+        return f"tcp://0.0.0.0:{self.scheduler_port}"
+
+
+# ── CLI registration ─────────────────────────────────────────────────
+
+
+def add_disagg_cli_args(parser: argparse.ArgumentParser) -> None:
+    """Register all disaggregated-diffusion CLI arguments as a group."""
+
+    g = parser.add_argument_group(
+        "Disaggregated diffusion",
+        "Split the pipeline into independent Encoder / Denoiser / Decoder "
+        "roles, each on its own GPU(s).  A DiffusionServer head node routes "
+        "requests.  See docs/disaggregation.md for details.",
+    )
+
+    # Core
+    g.add_argument(
+        "--base-gpu-id",
+        type=int,
+        default=0,
+        help="Starting GPU ID for this instance.  Used with --disagg-role "
+        "to place role instances on specific GPUs without CUDA_VISIBLE_DEVICES.",
+    )
+    g.add_argument(
+        "--disagg-role",
+        type=str,
+        default=RoleType.MONOLITHIC.value,
+        choices=RoleType.choices(),
+        help="Role for disaggregated pipeline.  "
+        "'monolithic' (default): single server.  "
+        "'encoder' / 'denoiser' / 'decoder': role instance.  "
+        "'server': DiffusionServer head node (no GPU).  "
+        "Role instances require --disagg-server-addr.  "
+        "Server requires --encoder-urls, --denoiser-urls, --decoder-urls.",
+    )
+    g.add_argument(
+        "--disagg-server-addr",
+        type=str,
+        default=None,
+        help="DiffusionServer head node address (tcp://HOST:PORT).  "
+        "Required for role instances.",
+    )
+    g.add_argument(
+        "--disagg-timeout",
+        type=int,
+        default=600,
+        help="Timeout in seconds for pending disagg requests (default: 600).",
+    )
+    g.add_argument(
+        "--disagg-dispatch-policy",
+        type=str,
+        default="round_robin",
+        choices=["round_robin", "max_free_slots"],
+        help="Dispatch policy: 'round_robin' or 'max_free_slots' (default: round_robin).",
+    )
+
+    # Server head: remote instance URLs
+    g.add_argument(
+        "--encoder-urls",
+        type=str,
+        default=None,
+        help="Encoder work endpoints (semicolon-separated).  "
+        "Example: 'tcp://10.0.0.1:35000;tcp://10.0.0.2:35000'.",
+    )
+    g.add_argument(
+        "--denoiser-urls",
+        type=str,
+        default=None,
+        help="Denoiser work endpoints (semicolon-separated).",
+    )
+    g.add_argument(
+        "--decoder-urls",
+        type=str,
+        default=None,
+        help="Decoder work endpoints (semicolon-separated).",
+    )
+
+    # Per-role parallelism
+    g.add_argument("--encoder-tp", type=int, default=None, help="Encoder TP degree.")
+    g.add_argument("--denoiser-tp", type=int, default=None, help="Denoiser TP degree.")
+    g.add_argument("--denoiser-sp", type=int, default=None, help="Denoiser SP degree.")
+    g.add_argument(
+        "--denoiser-ulysses", type=int, default=None, help="Denoiser Ulysses degree."
+    )
+    g.add_argument(
+        "--denoiser-ring", type=int, default=None, help="Denoiser Ring degree."
+    )
+    g.add_argument("--decoder-tp", type=int, default=None, help="Decoder TP degree.")
+
+    # P2P transfer engine
+    g.add_argument(
+        "--disagg-transfer-pool-size",
+        type=int,
+        default=256 * 1024 * 1024,
+        help="P2P transfer buffer pool size in bytes (default: 256 MiB).",
+    )
+    g.add_argument(
+        "--disagg-p2p-hostname",
+        type=str,
+        default="127.0.0.1",
+        help="RDMA-reachable hostname/IP of this instance (default: 127.0.0.1).",
+    )
+    g.add_argument(
+        "--disagg-ib-device",
+        type=str,
+        default=None,
+        help="InfiniBand device for RDMA transfers (e.g., mlx5_0).",
+    )
+
+
+def convert_disagg_role_string(kwargs: dict) -> None:
+    """Convert ``disagg_role`` from string to ``RoleType`` enum in-place."""
+    if "disagg_role" in kwargs and isinstance(kwargs["disagg_role"], str):
+        kwargs["disagg_role"] = RoleType.from_string(kwargs["disagg_role"])

python/sglang/multimodal_gen/runtime/disaggregation/dispatch_policy.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dispatch policies for multi-instance disaggregated diffusion pipelines."""
+
+import abc
+import logging
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+class DispatchPolicy(abc.ABC):
+    def __init__(self, num_instances: int):
+        if num_instances < 1:
+            raise ValueError(f"num_instances must be >= 1, got {num_instances}")
+        self._num_instances = num_instances
+
+    @property
+    def num_instances(self) -> int:
+        return self._num_instances
+
+    @abc.abstractmethod
+    def select(self, active_counts: list[int] | None = None) -> int: ...
+
+    def select_with_capacity(self, free_slots: list[int]) -> int | None:
+        """Select an instance that has free capacity, or None if all full."""
+        if not any(s > 0 for s in free_slots):
+            return None
+        return self.select(active_counts=None)
+
+    def record_completion(self, instance_id: int) -> None:
+        pass
+
+
+class RoundRobin(DispatchPolicy):
+    def __init__(self, num_instances: int):
+        super().__init__(num_instances)
+        self._lock = threading.Lock()
+        self._next = 0
+
+    def select(self, active_counts: list[int] | None = None) -> int:
+        with self._lock:
+            chosen = self._next
+            self._next = (self._next + 1) % self._num_instances
+        return chosen
+
+    def select_with_capacity(self, free_slots: list[int]) -> int | None:
+        with self._lock:
+            for _ in range(self._num_instances):
+                idx = self._next
+                self._next = (self._next + 1) % self._num_instances
+                if free_slots[idx] > 0:
+                    return idx
+            return None
+
+
+class MaxFreeSlotsFirst(DispatchPolicy):
+    """Dispatch to the instance with the most free slots."""
+
+    def __init__(self, num_instances: int, max_slots_per_instance: int = 1):
+        super().__init__(num_instances)
+        self._max_slots = max_slots_per_instance
+        self._lock = threading.Lock()
+        self._tiebreak = 0
+
+    def select(self, active_counts: list[int] | None = None) -> int:
+        with self._lock:
+            if active_counts is None or len(active_counts) != self._num_instances:
+                chosen = self._tiebreak % self._num_instances
+                self._tiebreak += 1
+                return chosen
+
+            best_id = 0
+            best_free = self._max_slots - active_counts[0]
+            for i in range(1, self._num_instances):
+                free = self._max_slots - active_counts[i]
+                if free > best_free:
+                    best_free = free
+                    best_id = i
+                elif free == best_free:
+                    if i == (self._tiebreak % self._num_instances):
+                        best_id = i
+
+            self._tiebreak += 1
+
+            if best_free <= 0:
+                logger.warning(
+                    "All %d instances are at capacity (%d slots each), "
+                    "dispatching to instance %d anyway",
+                    self._num_instances,
+                    self._max_slots,
+                    best_id,
+                )
+
+            return best_id
+
+    def select_with_capacity(self, free_slots: list[int]) -> int | None:
+        with self._lock:
+            best_id = -1
+            best_free = 0
+            for i in range(self._num_instances):
+                if free_slots[i] > best_free:
+                    best_free = free_slots[i]
+                    best_id = i
+                elif free_slots[i] == best_free and best_free > 0:
+                    if i == (self._tiebreak % self._num_instances):
+                        best_id = i
+
+            self._tiebreak += 1
+
+            if best_id < 0:
+                return None
+            return best_id
+
+
+class PoolDispatcher:
+    """Wraps three independent dispatch policies for encoder/denoiser/decoder pools."""
+
+    def __init__(
+        self,
+        num_encoders: int,
+        num_denoisers: int,
+        num_decoders: int,
+        policy_name: str = "round_robin",
+        **kwargs,
+    ):
+        self.encoder_policy = create_dispatch_policy(
+            policy_name, num_encoders, **kwargs
+        )
+        self.denoiser_policy = create_dispatch_policy(
+            policy_name, num_denoisers, **kwargs
+        )
+        self.decoder_policy = create_dispatch_policy(
+            policy_name, num_decoders, **kwargs
+        )
+
+    def select_encoder(self, active_counts: list[int] | None = None) -> int:
+        return self.encoder_policy.select(active_counts)
+
+    def select_denoiser(self, active_counts: list[int] | None = None) -> int:
+        return self.denoiser_policy.select(active_counts)
+
+    def select_decoder(self, active_counts: list[int] | None = None) -> int:
+        return self.decoder_policy.select(active_counts)
+
+    def select_encoder_with_capacity(self, free_slots: list[int]) -> int | None:
+        return self.encoder_policy.select_with_capacity(free_slots)
+
+    def select_denoiser_with_capacity(self, free_slots: list[int]) -> int | None:
+        return self.denoiser_policy.select_with_capacity(free_slots)
+
+    def select_decoder_with_capacity(self, free_slots: list[int]) -> int | None:
+        return self.decoder_policy.select_with_capacity(free_slots)
+
+
+def create_dispatch_policy(name: str, num_instances: int, **kwargs) -> DispatchPolicy:
+    policies = {
+        "round_robin": RoundRobin,
+        "max_free_slots": MaxFreeSlotsFirst,
+    }
+    cls = policies.get(name)
+    if cls is None:
+        raise ValueError(
+            f"Unknown dispatch policy '{name}'. Available: {list(policies.keys())}"
+        )
+    return cls(num_instances=num_instances, **kwargs)

python/sglang/multimodal_gen/runtime/disaggregation/metrics.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Observability metrics for disaggregated diffusion pipelines."""
+
+import threading
+import time
+from dataclasses import dataclass
+
+
+@dataclass
+class _RequestTiming:
+    start_time: float
+    stage_start: float = 0.0
+
+
+@dataclass
+class RoleStats:
+    role: str
+    requests_completed: int = 0
+    requests_failed: int = 0
+    requests_in_flight: int = 0
+    requests_timed_out: int = 0
+    queue_depth: int = 0
+    last_latency_s: float = 0.0
+    avg_latency_s: float = 0.0
+    max_latency_s: float = 0.0
+    throughput_rps: float = 0.0
+    uptime_s: float = 0.0
+
+    def to_dict(self) -> dict:
+        return {
+            "role": self.role,
+            "requests_completed": self.requests_completed,
+            "requests_failed": self.requests_failed,
+            "requests_in_flight": self.requests_in_flight,
+            "requests_timed_out": self.requests_timed_out,
+            "queue_depth": self.queue_depth,
+            "last_latency_s": round(self.last_latency_s, 4),
+            "avg_latency_s": round(self.avg_latency_s, 4),
+            "max_latency_s": round(self.max_latency_s, 4),
+            "throughput_rps": round(self.throughput_rps, 4),
+            "uptime_s": round(self.uptime_s, 1),
+        }
+
+
+class DisaggMetrics:
+    """Thread-safe metrics collector for a single disagg role."""
+
+    def __init__(self, role: str):
+        self._role = role
+        self._lock = threading.Lock()
+        self._start_time = time.monotonic()
+
+        self._completed = 0
+        self._failed = 0
+        self._timed_out = 0
+
+        self._in_flight: dict[str, _RequestTiming] = {}
+
+        self._last_latency = 0.0
+        self._max_latency = 0.0
+        self._total_latency = 0.0
+
+        self._completion_times: list[float] = []
+        self._throughput_window_s = 60.0
+
+        self._queue_depth = 0
+
+    @property
+    def role(self) -> str:
+        return self._role
+
+    def record_request_start(self, request_id: str) -> None:
+        with self._lock:
+            self._in_flight[request_id] = _RequestTiming(start_time=time.monotonic())
+
+    def record_request_complete(self, request_id: str) -> None:
+        now = time.monotonic()
+        with self._lock:
+            timing = self._in_flight.pop(request_id, None)
+            if timing is not None:
+                latency = now - timing.start_time
+                self._last_latency = latency
+                self._max_latency = max(self._max_latency, latency)
+                self._total_latency += latency
+
+            self._completed += 1
+            self._completion_times.append(now)
+            self._prune_completion_times(now)
+
+    def record_request_failed(self, request_id: str) -> None:
+        with self._lock:
+            self._in_flight.pop(request_id, None)
+            self._failed += 1
+
+    def record_request_timeout(self, request_id: str) -> None:
+        with self._lock:
+            self._in_flight.pop(request_id, None)
+            self._timed_out += 1
+
+    def update_queue_depth(self, depth: int) -> None:
+        with self._lock:
+            self._queue_depth = depth
+
+    def snapshot(self) -> RoleStats:
+        now = time.monotonic()
+        with self._lock:
+            self._prune_completion_times(now)
+            total = self._completed + self._failed
+            avg_latency = self._total_latency / total if total > 0 else 0.0
+            rps = (
+                len(self._completion_times) / self._throughput_window_s
+                if self._completion_times
+                else 0.0
+            )
+
+            return RoleStats(
+                role=self._role,
+                requests_completed=self._completed,
+                requests_failed=self._failed,
+                requests_in_flight=len(self._in_flight),
+                requests_timed_out=self._timed_out,
+                queue_depth=self._queue_depth,
+                last_latency_s=self._last_latency,
+                avg_latency_s=avg_latency,
+                max_latency_s=self._max_latency,
+                throughput_rps=rps,
+                uptime_s=now - self._start_time,
+            )
+
+    def _prune_completion_times(self, now: float) -> None:
+        cutoff = now - self._throughput_window_s
+        while self._completion_times and self._completion_times[0] < cutoff:
+            self._completion_times.pop(0)

python/sglang/multimodal_gen/runtime/disaggregation/request_state.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Request state machine for disaggregated diffusion pipelines."""
+
+import enum
+import logging
+import threading
+import time
+from dataclasses import dataclass, field
+
+logger = logging.getLogger(__name__)
+
+
+class RequestState(enum.Enum):
+    """Lifecycle states for a disagg pipeline request.
+
+    *_WAITING: request queued, awaiting a free buffer slot.
+    *_RUNNING: request dispatched to a specific instance.
+    """
+
+    PENDING = "pending"
+    ENCODER_WAITING = "encoder_waiting"
+    ENCODER_RUNNING = "encoder_running"
+    ENCODER_DONE = "encoder_done"
+    DENOISING_WAITING = "denoising_waiting"
+    DENOISING_RUNNING = "denoising_running"
+    DENOISING_DONE = "denoising_done"
+    DECODER_WAITING = "decoder_waiting"
+    DECODER_RUNNING = "decoder_running"
+    DONE = "done"
+    FAILED = "failed"
+    TIMED_OUT = "timed_out"
+
+
+_TERMINAL_STATES = {RequestState.DONE, RequestState.FAILED, RequestState.TIMED_OUT}
+_ACTIVE_STATES = set(RequestState) - _TERMINAL_STATES
+
+# Normal (non-failure) transitions.  FAILED and TIMED_OUT are handled
+# separately in transition() — any active state can reach them.
+_VALID_TRANSITIONS: dict[RequestState, set[RequestState]] = {
+    RequestState.PENDING: {RequestState.ENCODER_WAITING, RequestState.ENCODER_RUNNING},
+    RequestState.ENCODER_WAITING: {RequestState.ENCODER_RUNNING},
+    RequestState.ENCODER_RUNNING: {RequestState.ENCODER_DONE},
+    RequestState.ENCODER_DONE: {
+        RequestState.DENOISING_WAITING,
+        RequestState.DENOISING_RUNNING,
+    },
+    RequestState.DENOISING_WAITING: {RequestState.DENOISING_RUNNING},
+    RequestState.DENOISING_RUNNING: {RequestState.DENOISING_DONE},
+    RequestState.DENOISING_DONE: {
+        RequestState.DECODER_WAITING,
+        RequestState.DECODER_RUNNING,
+    },
+    RequestState.DECODER_WAITING: {RequestState.DECODER_RUNNING},
+    RequestState.DECODER_RUNNING: {RequestState.DONE},
+}
+
+
+@dataclass
+class RequestRecord:
+    request_id: str
+    state: RequestState = RequestState.PENDING
+    submit_time: float = field(default_factory=time.monotonic)
+    last_transition_time: float = field(default_factory=time.monotonic)
+    encoder_instance: int | None = None
+    denoiser_instance: int | None = None
+    decoder_instance: int | None = None
+    error: str | None = None
+
+    def elapsed_s(self) -> float:
+        return time.monotonic() - self.submit_time
+
+    def is_terminal(self) -> bool:
+        return self.state in _TERMINAL_STATES
+
+
+class RequestTracker:
+    """Thread-safe tracker for request state machines."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._requests: dict[str, RequestRecord] = {}
+
+    def submit(self, request_id: str) -> RequestRecord:
+        with self._lock:
+            if request_id in self._requests:
+                raise ValueError(f"Duplicate request_id: {request_id}")
+            record = RequestRecord(request_id=request_id)
+            self._requests[request_id] = record
+            return record
+
+    def transition(
+        self,
+        request_id: str,
+        new_state: RequestState,
+        *,
+        error: str | None = None,
+        encoder_instance: int | None = None,
+        denoiser_instance: int | None = None,
+        decoder_instance: int | None = None,
+    ) -> RequestRecord:
+        with self._lock:
+            record = self._requests.get(request_id)
+            if record is None:
+                raise ValueError(f"Unknown request_id: {request_id}")
+
+            old_state = record.state
+
+            if new_state in _TERMINAL_STATES and new_state != RequestState.DONE:
+                # FAILED / TIMED_OUT: allowed from any active state
+                if old_state not in _ACTIVE_STATES:
+                    raise ValueError(
+                        f"Cannot transition {request_id} from terminal state "
+                        f"{old_state.value} to {new_state.value}"
+                    )
+            elif new_state not in _VALID_TRANSITIONS.get(old_state, set()):
+                raise ValueError(
+                    f"Invalid transition for {request_id}: "
+                    f"{old_state.value} -> {new_state.value}"
+                )
+
+            record.state = new_state
+            record.last_transition_time = time.monotonic()
+            if error is not None:
+                record.error = error
+            if encoder_instance is not None:
+                record.encoder_instance = encoder_instance
+            if denoiser_instance is not None:
+                record.denoiser_instance = denoiser_instance
+            if decoder_instance is not None:
+                record.decoder_instance = decoder_instance
+
+            logger.debug(
+                "Request %s: %s -> %s", request_id, old_state.value, new_state.value
+            )
+            return record
+
+    def get(self, request_id: str) -> RequestRecord | None:
+        with self._lock:
+            return self._requests.get(request_id)
+
+    def remove(self, request_id: str) -> RequestRecord | None:
+        with self._lock:
+            return self._requests.pop(request_id, None)
+
+    def find_timed_out(self, timeout_s: float) -> list[str]:
+        now = time.monotonic()
+        with self._lock:
+            return [
+                r.request_id
+                for r in self._requests.values()
+                if r.state in _ACTIVE_STATES and (now - r.submit_time) > timeout_s
+            ]
+
+    def snapshot(self) -> dict:
+        with self._lock:
+            state_counts = {}
+            for r in self._requests.values():
+                state_counts[r.state.value] = state_counts.get(r.state.value, 0) + 1
+            return {
+                "total": len(self._requests),
+                "active": sum(
+                    1 for r in self._requests.values() if not r.is_terminal()
+                ),
+                "by_state": state_counts,
+            }

python/sglang/multimodal_gen/runtime/disaggregation/roles.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Role definitions for diffusion pipeline disaggregation."""
+
+from enum import Enum
+
+_ROLE_ALIASES = {"denoising": "denoiser"}
+
+
+class RoleType(str, Enum):
+    MONOLITHIC = "monolithic"
+    ENCODER = "encoder"
+    DENOISER = "denoiser"
+    DECODER = "decoder"
+    SERVER = "server"  # Head node (no GPU, routes requests)
+
+    @classmethod
+    def from_string(cls, value: str) -> "RoleType":
+        v = _ROLE_ALIASES.get(value.lower(), value.lower())
+        try:
+            return cls(v)
+        except ValueError:
+            raise ValueError(
+                f"Invalid role: {value}. Must be one of: {', '.join([r.value for r in cls])}"
+            ) from None
+
+    @classmethod
+    def choices(cls) -> list[str]:
+        return [role.value for role in cls]
+
+
+def get_module_role(module_name: str) -> "RoleType | None":
+    """Classify a module name to its primary role. Returns None for shared modules."""
+    encoder_prefixes = (
+        "text_encoder",
+        "tokenizer",
+        "image_encoder",
+        "image_processor",
+        "processor",
+        "connectors",
+    )
+    if any(
+        module_name == p or module_name.startswith(p + "_") for p in encoder_prefixes
+    ):
+        return RoleType.ENCODER
+
+    denoising_prefixes = ("transformer",)
+    if any(
+        module_name == p or module_name.startswith(p + "_") for p in denoising_prefixes
+    ):
+        return RoleType.DENOISER
+
+    decoder_prefixes = ("vae", "audio_vae", "video_vae", "vocoder")
+    if any(
+        module_name == p or module_name.startswith(p + "_") for p in decoder_prefixes
+    ):
+        return RoleType.DECODER
+
+    return None
+
+
+def filter_modules_for_role(module_names: list[str], role: "RoleType") -> list[str]:
+    """Filter module names to only those needed by the given role."""
+    if role in (RoleType.MONOLITHIC, RoleType.SERVER):
+        return module_names
+
+    filtered = []
+    for name in module_names:
+        module_role = get_module_role(name)
+
+        if module_role is None:
+            filtered.append(name)
+        elif module_role == role:
+            filtered.append(name)
+        elif role == RoleType.ENCODER and module_role == RoleType.DECODER:
+            # Encoder also needs VAE for ImageVAEEncoding stages
+            filtered.append(name)
+
+    return filtered

python/sglang/multimodal_gen/runtime/disaggregation/transport/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Transport layer for disaggregated diffusion pipelines."""

python/sglang/multimodal_gen/runtime/disaggregation/transport/allocator.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Buddy-system memory allocator for TransferTensorBuffer."""
+
+from __future__ import annotations
+
+import logging
+import threading
+from dataclasses import dataclass
+from functools import lru_cache
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Block:
+    offset: int  # byte offset from pool start
+    size: int
+    allocated: bool = False
+    request_id: str | None = None
+
+
+class BuddyAllocator:
+    """Power-of-2 buddy-system allocator for pinned memory."""
+
+    def __init__(self, pool_size: int, min_block_size: int = 1 << 20):
+        if min_block_size <= 0 or (min_block_size & (min_block_size - 1)) != 0:
+            raise ValueError(
+                f"min_block_size must be a power of 2, got {min_block_size}"
+            )
+
+        self._min_block_size = min_block_size
+        self._pool_size = self._next_power_of_2(max(pool_size, min_block_size))
+        self._lock = threading.Lock()
+
+        # Free lists indexed by order: order 0 = min_block_size, order 1 = 2*min_block_size, ...
+        self._max_order = self._size_to_order(self._pool_size)
+        self._free_lists: list[list[int]] = [[] for _ in range(self._max_order + 1)]
+
+        self._blocks: dict[int, Block] = {}
+
+        root = Block(offset=0, size=self._pool_size)
+        self._blocks[0] = root
+        self._free_lists[self._max_order].append(0)
+
+        self._allocated_bytes = 0
+        self._num_allocations = 0
+
+    @property
+    def pool_size(self) -> int:
+        return self._pool_size
+
+    def allocate(self, size: int, request_id: str | None = None) -> int | None:
+        """Allocate a block of at least `size` bytes. Returns offset or None."""
+        if size <= 0:
+            raise ValueError(f"Allocation size must be positive, got {size}")
+
+        alloc_size = max(self._next_power_of_2(size), self._min_block_size)
+        target_order = self._size_to_order(alloc_size)
+
+        if target_order > self._max_order:
+            logger.warning(
+                "Requested size %d exceeds pool size %d", size, self._pool_size
+            )
+            return None
+
+        with self._lock:
+            return self._allocate_locked(target_order, request_id)
+
+    def free(self, offset: int) -> bool:
+        """Free the block at the given offset and coalesce with buddy if possible."""
+        with self._lock:
+            return self._free_locked(offset)
+
+    def get_block_info(self, offset: int) -> Block | None:
+        with self._lock:
+            return self._blocks.get(offset)
+
+    def get_stats(self) -> dict:
+        with self._lock:
+            free_blocks_by_order = {}
+            for order, offsets in enumerate(self._free_lists):
+                if offsets:
+                    block_size = self._min_block_size << order
+                    free_blocks_by_order[block_size] = len(offsets)
+
+            return {
+                "pool_size": self._pool_size,
+                "min_block_size": self._min_block_size,
+                "allocated_bytes": self._allocated_bytes,
+                "free_bytes": self._pool_size - self._allocated_bytes,
+                "num_allocations": self._num_allocations,
+                "num_blocks": len(self._blocks),
+                "free_blocks_by_size": free_blocks_by_order,
+            }
+
+    def count_free_slots(self, slot_size: int) -> int:
+        """Count how many allocations of the given size can fit."""
+        if slot_size <= 0:
+            return 0
+        alloc_size = max(self._next_power_of_2(slot_size), self._min_block_size)
+
+        with self._lock:
+            count = 0
+            for order in range(self._size_to_order(alloc_size), self._max_order + 1):
+                for _ in self._free_lists[order]:
+                    block_size = self._min_block_size << order
+                    count += block_size // alloc_size
+            return count
+
+    # --- Internal (caller must hold self._lock) ---
+
+    def _allocate_locked(self, target_order: int, request_id: str | None) -> int | None:
+        found_order = -1
+        for order in range(target_order, self._max_order + 1):
+            if self._free_lists[order]:
+                found_order = order
+                break
+
+        if found_order < 0:
+            return None
+
+        offset = self._free_lists[found_order].pop(0)
+        block = self._blocks[offset]
+
+        # Split down to target_order
+        while found_order > target_order:
+            found_order -= 1
+            buddy_size = self._min_block_size << found_order
+            buddy_offset = offset + buddy_size
+
+            buddy = Block(offset=buddy_offset, size=buddy_size)
+            self._blocks[buddy_offset] = buddy
+            self._free_lists[found_order].append(buddy_offset)
+
+            block.size = buddy_size
+
+        block.allocated = True
+        block.request_id = request_id
+        self._allocated_bytes += block.size
+        self._num_allocations += 1
+
+        return offset
+
+    def _free_locked(self, offset: int) -> bool:
+        block = self._blocks.get(offset)
+        if block is None or not block.allocated:
+            return False
+
+        block.allocated = False
+        block.request_id = None
+        self._allocated_bytes -= block.size
+        self._num_allocations -= 1
+
+        self._coalesce(block)
+        return True
+
+    def _coalesce(self, block: Block) -> None:
+        """Recursively merge with buddy if both are free."""
+        while block.size < self._pool_size:
+            buddy_offset = block.offset ^ block.size
+            buddy = self._blocks.get(buddy_offset)
+
+            if buddy is None or buddy.allocated or buddy.size != block.size:
+                break
+
+            order = self._size_to_order(buddy.size)
+            self._free_lists[order].remove(buddy_offset)
+
+            if buddy_offset < block.offset:
+                del self._blocks[block.offset]
+                buddy.size *= 2
+                block = buddy
+            else:
+                del self._blocks[buddy_offset]
+                block.size *= 2
+
+        order = self._size_to_order(block.size)
+        self._free_lists[order].append(block.offset)
+
+    def _size_to_order(self, size: int) -> int:
+        order = 0
+        s = self._min_block_size
+        while s < size:
+            s <<= 1
+            order += 1
+        return order
+
+    @staticmethod
+    @lru_cache(maxsize=256)
+    def _next_power_of_2(n: int) -> int:
+        if n <= 0:
+            return 1
+        n -= 1
+        n |= n >> 1
+        n |= n >> 2
+        n |= n >> 4
+        n |= n >> 8
+        n |= n >> 16
+        n |= n >> 32
+        return n + 1

python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py

@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: Apache-2.0
+"""TransferTensorBuffer: memory staging area for disaggregated tensor transfer."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.transport.allocator import (
+    BuddyAllocator,
+)
+from sglang.multimodal_gen.runtime.disaggregation.transport.codec import (
+    str_to_dtype,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SlotHandle:
+    request_id: str
+    offset: int  # byte offset in the pool
+    size: int  # allocated size in bytes
+    tensor_views: dict[str, torch.Tensor | list[torch.Tensor]] = field(
+        default_factory=dict
+    )
+
+
+class TransferTensorBuffer:
+    """Memory pool for staging tensor payloads between roles.
+
+    Wraps a contiguous block of memory (CPU pinned or GPU) with a BuddyAllocator.
+    """
+
+    def __init__(
+        self,
+        pool_size: int,
+        min_block_size: int = 1 << 20,
+        role_name: str = "unknown",
+        device: str = "cpu",
+    ):
+        self._role_name = role_name
+        self._device = device
+        self._allocator = BuddyAllocator(pool_size, min_block_size)
+        actual_size = self._allocator.pool_size
+
+        if device == "cpu":
+            self._pool = torch.empty(actual_size, dtype=torch.uint8, pin_memory=True)
+        else:
+            self._pool = torch.empty(actual_size, dtype=torch.uint8, device=device)
+        self._pool_ptr = self._pool.data_ptr()
+
+        pool_location = "pinned CPU" if device == "cpu" else f"GPU ({device})"
+        logger.info(
+            "TransferTensorBuffer[%s]: allocated %d MiB %s memory "
+            "(min_block=%d KiB)",
+            role_name,
+            actual_size >> 20,
+            pool_location,
+            min_block_size >> 10,
+        )
+
+    @property
+    def pool_size(self) -> int:
+        return self._allocator.pool_size
+
+    @property
+    def device(self) -> str:
+        return self._device
+
+    @property
+    def pool_data_ptr(self) -> int:
+        return self._pool_ptr
+
+    def allocate(self, size: int, request_id: str) -> SlotHandle | None:
+        """Allocate a slot. Returns None if pool is full."""
+        offset = self._allocator.allocate(size, request_id=request_id)
+        if offset is None:
+            logger.warning(
+                "TransferTensorBuffer[%s]: allocation failed for %s (%d bytes). "
+                "Pool stats: %s",
+                self._role_name,
+                request_id,
+                size,
+                self._allocator.get_stats(),
+            )
+            return None
+
+        block = self._allocator.get_block_info(offset)
+        return SlotHandle(
+            request_id=request_id,
+            offset=offset,
+            size=block.size if block else size,
+        )
+
+    def free(self, handle: SlotHandle) -> bool:
+        return self._allocator.free(handle.offset)
+
+    def write_tensor(
+        self,
+        handle: SlotHandle,
+        name: str,
+        tensor: torch.Tensor,
+        byte_offset: int = 0,
+        stream: torch.cuda.Stream | None = None,
+    ) -> int:
+        """Copy a tensor into the pool slot. Returns bytes written."""
+        src_tensor = tensor.contiguous()
+        nbytes = src_tensor.numel() * src_tensor.element_size()
+
+        if byte_offset + nbytes > handle.size:
+            raise ValueError(
+                f"Write exceeds slot: offset={byte_offset}, nbytes={nbytes}, "
+                f"slot_size={handle.size}"
+            )
+
+        dst = self._pool[
+            handle.offset + byte_offset : handle.offset + byte_offset + nbytes
+        ]
+        src_bytes = src_tensor.view(torch.uint8).reshape(-1)
+
+        if stream is not None:
+            with torch.cuda.stream(stream):
+                dst.copy_(src_bytes, non_blocking=True)
+        else:
+            dst.copy_(src_bytes, non_blocking=True)
+
+        return nbytes
+
+    def read_tensor(
+        self,
+        handle: SlotHandle,
+        shape: list[int],
+        dtype: torch.dtype,
+        byte_offset: int = 0,
+        device: torch.device | str = "cpu",
+        stream: torch.cuda.Stream | None = None,
+    ) -> torch.Tensor:
+        """Read a tensor from the pool slot. Returns a clone on target device."""
+        nbytes = 1
+        for s in shape:
+            nbytes *= s
+        nbytes *= torch.tensor([], dtype=dtype).element_size()
+
+        raw = self._pool[
+            handle.offset + byte_offset : handle.offset + byte_offset + nbytes
+        ]
+        src = raw.view(dtype).reshape(shape)
+
+        pool_dev = str(self._pool.device)
+        target_dev = str(device)
+
+        same_device = pool_dev == target_dev
+
+        if same_device:
+            # Clone to decouple tensor lifetime from pool slot
+            if stream is not None:
+                with torch.cuda.stream(stream):
+                    return src.clone()
+            return src.clone()
+
+        if stream is not None:
+            with torch.cuda.stream(stream):
+                return src.to(device, non_blocking=True)
+        return src.to(device, non_blocking=True)
+
+    def write_tensors_from_gpu(
+        self,
+        handle: SlotHandle,
+        tensors: dict[str, torch.Tensor | list[torch.Tensor] | None],
+        stream: torch.cuda.Stream | None = None,
+    ) -> dict[str, list[dict]]:
+        """Batch-write GPU tensors into a slot. Returns a manifest for later reads."""
+        manifest: dict[str, list[dict]] = {}
+        byte_offset = 0
+
+        # Ensure copy stream sees all prior compute kernels
+        if stream is not None:
+            stream.wait_stream(torch.cuda.current_stream())
+
+        for name, value in tensors.items():
+            if value is None:
+                continue
+
+            entries = []
+            if isinstance(value, torch.Tensor):
+                nbytes = self.write_tensor(handle, name, value, byte_offset, stream)
+                entries.append(
+                    {
+                        "offset": byte_offset,
+                        "shape": list(value.shape),
+                        "dtype": str(value.dtype).replace("torch.", ""),
+                    }
+                )
+                byte_offset += nbytes
+                byte_offset = (byte_offset + 511) & ~511  # align to 512B
+
+            elif isinstance(value, list):
+                for i, t in enumerate(value):
+                    if t is None:
+                        continue
+                    nbytes = self.write_tensor(
+                        handle, f"{name}[{i}]", t, byte_offset, stream
+                    )
+                    entries.append(
+                        {
+                            "offset": byte_offset,
+                            "shape": list(t.shape),
+                            "dtype": str(t.dtype).replace("torch.", ""),
+                            "list_index": i,
+                        }
+                    )
+                    byte_offset += nbytes
+                    byte_offset = (byte_offset + 511) & ~511
+
+            if entries:
+                manifest[name] = entries
+
+        return manifest
+
+    def read_tensors_from_manifest(
+        self,
+        handle: SlotHandle,
+        manifest: dict[str, list[dict]],
+        device: torch.device | str = "cpu",
+        stream: torch.cuda.Stream | None = None,
+    ) -> dict[str, torch.Tensor | list[torch.Tensor]]:
+        """Batch-read tensors from a slot using a manifest."""
+        result: dict[str, torch.Tensor | list[torch.Tensor]] = {}
+
+        for name, entries in manifest.items():
+            if not entries:
+                continue
+            has_list_index = any("list_index" in e for e in entries)
+
+            if has_list_index:
+                max_idx = max(e.get("list_index", 0) for e in entries) + 1
+                tensors = [None] * max_idx
+                for entry in entries:
+                    t = self.read_tensor(
+                        handle,
+                        entry["shape"],
+                        str_to_dtype(entry["dtype"]),
+                        entry["offset"],
+                        device,
+                        stream,
+                    )
+                    tensors[entry["list_index"]] = t
+                result[name] = tensors
+            else:
+                entry = entries[0]
+                result[name] = self.read_tensor(
+                    handle,
+                    entry["shape"],
+                    str_to_dtype(entry["dtype"]),
+                    entry["offset"],
+                    device,
+                    stream,
+                )
+
+        return result
+
+    def free_slots_count(self, typical_request_size: int) -> int:
+        """Estimate how many requests of typical size can still be buffered."""
+        return self._allocator.count_free_slots(typical_request_size)
+
+    def get_stats(self) -> dict:
+        alloc_stats = self._allocator.get_stats()
+        alloc_stats["role"] = self._role_name
+        return alloc_stats

python/sglang/multimodal_gen/runtime/disaggregation/transport/codec.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Zero-copy tensor codec for ZMQ multipart messages.
+
+Frame 0: JSON metadata (tensor descriptors + scalar fields)
+Frame 1-N: Raw tensor data buffers (one per tensor)
+"""
+
+import ctypes
+import json
+import logging
+from dataclasses import dataclass
+
+import torch
+import zmq
+
+logger = logging.getLogger(__name__)
+
+_DTYPE_TO_STR = {
+    torch.float16: "float16",
+    torch.float32: "float32",
+    torch.float64: "float64",
+    torch.bfloat16: "bfloat16",
+    torch.int8: "int8",
+    torch.int16: "int16",
+    torch.int32: "int32",
+    torch.int64: "int64",
+    torch.uint8: "uint8",
+    torch.bool: "bool",
+}
+_STR_TO_DTYPE = {v: k for k, v in _DTYPE_TO_STR.items()}
+
+
+def dtype_to_str(dtype: torch.dtype) -> str:
+    s = _DTYPE_TO_STR.get(dtype)
+    if s is None:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+    return s
+
+
+def str_to_dtype(s: str) -> torch.dtype:
+    d = _STR_TO_DTYPE.get(s)
+    if d is None:
+        raise ValueError(f"Unknown dtype string: {s}")
+    return d
+
+
+class TensorWrapper:
+    """Expose a CPU-contiguous tensor's data buffer for zero-copy ZMQ send."""
+
+    def __init__(self, tensor: torch.Tensor):
+        if tensor.is_cuda:
+            tensor = tensor.cpu()
+        if not tensor.is_contiguous():
+            tensor = tensor.contiguous()
+        self.tensor = tensor
+        data_ptr = tensor.data_ptr()
+        total_bytes = tensor.numel() * tensor.element_size()
+        self._c_buf = (ctypes.c_char * total_bytes).from_address(data_ptr)
+        self._view = memoryview(self._c_buf)
+
+
+@dataclass
+class TensorDescriptor:
+    field_name: str
+    shape: list[int]
+    dtype: str
+    list_index: int = -1  # -1 means not part of a list
+
+    def to_dict(self) -> dict:
+        return {
+            "field_name": self.field_name,
+            "shape": self.shape,
+            "dtype": self.dtype,
+            "list_index": self.list_index,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "TensorDescriptor":
+        return cls(
+            field_name=d["field_name"],
+            shape=d["shape"],
+            dtype=d["dtype"],
+            list_index=d.get("list_index", -1),
+        )
+
+
+def pack_tensors(
+    tensor_fields: dict[str, torch.Tensor | list[torch.Tensor] | None],
+    scalar_fields: dict | None = None,
+) -> tuple[bytes, list[TensorWrapper]]:
+    """Pack tensor fields into metadata + buffer list for send_multipart."""
+    descriptors = []
+    buffers = []
+
+    for field_name, value in tensor_fields.items():
+        if value is None:
+            continue
+
+        if isinstance(value, torch.Tensor):
+            wrapper = TensorWrapper(value)
+            descriptors.append(
+                TensorDescriptor(
+                    field_name=field_name,
+                    shape=list(value.shape),
+                    dtype=dtype_to_str(value.dtype),
+                )
+            )
+            buffers.append(wrapper)
+
+        elif isinstance(value, list):
+            for i, t in enumerate(value):
+                if t is None:
+                    continue
+                if not isinstance(t, torch.Tensor):
+                    raise TypeError(
+                        f"Expected Tensor in list for field '{field_name}', "
+                        f"got {type(t)}"
+                    )
+                wrapper = TensorWrapper(t)
+                descriptors.append(
+                    TensorDescriptor(
+                        field_name=field_name,
+                        shape=list(t.shape),
+                        dtype=dtype_to_str(t.dtype),
+                        list_index=i,
+                    )
+                )
+                buffers.append(wrapper)
+
+    metadata = {
+        "tensor_descriptors": [d.to_dict() for d in descriptors],
+        "scalar_fields": scalar_fields or {},
+    }
+    metadata_bytes = json.dumps(metadata, separators=(",", ":")).encode("utf-8")
+    return metadata_bytes, buffers
+
+
+def send_tensors(
+    socket: zmq.Socket,
+    tensor_fields: dict[str, torch.Tensor | list[torch.Tensor] | None],
+    scalar_fields: dict | None = None,
+    flags: int = 0,
+) -> None:
+    """Send tensors over ZMQ using multipart with zero-copy."""
+    metadata_bytes, buffers = pack_tensors(tensor_fields, scalar_fields)
+    parts: list = [metadata_bytes]
+    parts.extend(w._view if isinstance(w, TensorWrapper) else w for w in buffers)
+    socket.send_multipart(parts, flags=flags, copy=True)
+
+
+def unpack_tensors(
+    parts: list,
+    device: str | torch.device = "cpu",
+) -> tuple[dict[str, torch.Tensor | list[torch.Tensor]], dict]:
+    """Unpack multipart message frames into tensor fields and scalar fields."""
+    metadata_frame = parts[0]
+    metadata_bytes = (
+        bytes(metadata_frame.buffer)
+        if hasattr(metadata_frame, "buffer")
+        else bytes(metadata_frame)
+    )
+    metadata = json.loads(metadata_bytes)
+
+    descriptors = [
+        TensorDescriptor.from_dict(d) for d in metadata["tensor_descriptors"]
+    ]
+    scalar_fields = metadata.get("scalar_fields", {})
+
+    if len(parts) - 1 != len(descriptors):
+        raise ValueError(
+            f"Expected {len(descriptors)} tensor frames, got {len(parts) - 1}"
+        )
+
+    tensor_fields: dict[str, torch.Tensor | list[torch.Tensor]] = {}
+    list_sizes: dict[str, int] = {}
+    for desc in descriptors:
+        if desc.list_index >= 0:
+            current_max = list_sizes.get(desc.field_name, 0)
+            list_sizes[desc.field_name] = max(current_max, desc.list_index + 1)
+
+    for field_name, size in list_sizes.items():
+        tensor_fields[field_name] = [None] * size
+
+    for i, desc in enumerate(descriptors):
+        frame = parts[i + 1]
+        buf = frame.buffer if hasattr(frame, "buffer") else bytes(frame)
+        dtype = str_to_dtype(desc.dtype)
+        # clone() to own the memory (decouple from ZMQ buffer lifetime)
+        tensor = torch.frombuffer(buf, dtype=dtype).reshape(desc.shape).clone()
+        if device != "cpu" and device != torch.device("cpu"):
+            tensor = tensor.to(device)
+
+        if desc.list_index >= 0:
+            tensor_fields[desc.field_name][desc.list_index] = tensor
+        else:
+            tensor_fields[desc.field_name] = tensor
+
+    return tensor_fields, scalar_fields

python/sglang/multimodal_gen/runtime/disaggregation/transport/engine.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Transfer engine abstraction for tensor transfer between role instances."""
+
+import logging
+from abc import ABC, abstractmethod
+
+logger = logging.getLogger(__name__)
+
+_MOONCAKE_AVAILABLE = None
+
+
+def _check_mooncake() -> bool:
+    global _MOONCAKE_AVAILABLE
+    if _MOONCAKE_AVAILABLE is None:
+        try:
+            from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (  # noqa: F401
+                MooncakeTransferEngine as _MTE,
+            )
+
+            _MOONCAKE_AVAILABLE = True
+        except ImportError:
+            _MOONCAKE_AVAILABLE = False
+    return _MOONCAKE_AVAILABLE
+
+
+class BaseTransferEngine(ABC):
+    """Abstract transfer engine for data movement between roles."""
+
+    @property
+    def supports_gpu_direct(self) -> bool:
+        return False
+
+    @property
+    @abstractmethod
+    def session_id(self) -> str: ...
+
+    @abstractmethod
+    def register_buffer(self, ptr: int, length: int) -> None: ...
+
+    @abstractmethod
+    def deregister_buffer(self, ptr: int) -> None: ...
+
+    @abstractmethod
+    def transfer_sync(
+        self, dst_session_id: str, src_addr: int, dst_addr: int, length: int
+    ) -> int:
+        """Returns 0 on success, negative on failure."""
+
+    @abstractmethod
+    def batch_transfer_sync(
+        self,
+        dst_session_id: str,
+        src_addrs: list[int],
+        dst_addrs: list[int],
+        lengths: list[int],
+    ) -> int: ...
+
+
+class MooncakeDiffusionEngine(BaseTransferEngine):
+    """Production engine backed by MooncakeTransferEngine (RDMA)."""
+
+    @property
+    def supports_gpu_direct(self) -> bool:
+        return True
+
+    def __init__(
+        self,
+        hostname: str,
+        gpu_id: int = 0,
+        ib_device: str | None = None,
+    ):
+        from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+            MooncakeTransferEngine,
+        )
+
+        self._engine = MooncakeTransferEngine(
+            hostname=hostname,
+            gpu_id=gpu_id,
+            ib_device=ib_device,
+        )
+        logger.info(
+            "MooncakeDiffusionEngine initialized: session_id=%s",
+            self._engine.session_id,
+        )
+
+    @property
+    def session_id(self) -> str:
+        return self._engine.session_id
+
+    def register_buffer(self, ptr: int, length: int) -> None:
+        self._engine.register(ptr, length)
+
+    def deregister_buffer(self, ptr: int) -> None:
+        self._engine.deregister(ptr)
+
+    def transfer_sync(
+        self, dst_session_id: str, src_addr: int, dst_addr: int, length: int
+    ) -> int:
+        return self._engine.transfer_sync(dst_session_id, src_addr, dst_addr, length)
+
+    def batch_transfer_sync(
+        self,
+        dst_session_id: str,
+        src_addrs: list[int],
+        dst_addrs: list[int],
+        lengths: list[int],
+    ) -> int:
+        return self._engine.batch_transfer_sync(
+            dst_session_id, src_addrs, dst_addrs, lengths
+        )
+
+
+def create_transfer_engine(
+    hostname: str = "127.0.0.1",
+    gpu_id: int = 0,
+    ib_device: str | None = None,
+) -> BaseTransferEngine:
+    """Factory: returns MooncakeDiffusionEngine if mooncake is available."""
+    if not _check_mooncake():
+        raise RuntimeError(
+            "Mooncake transfer engine is required for disaggregated diffusion "
+            "but is not installed. Please install mooncake first."
+        )
+    return MooncakeDiffusionEngine(
+        hostname=hostname, gpu_id=gpu_id, ib_device=ib_device
+    )

python/sglang/multimodal_gen/runtime/disaggregation/transport/manager.py

@@ -0,0 +1,387 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-instance transfer manager for disaggregated diffusion roles."""
+
+import logging
+import threading
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.transport.buffer import (
+    SlotHandle,
+    TransferTensorBuffer,
+)
+from sglang.multimodal_gen.runtime.disaggregation.transport.engine import (
+    BaseTransferEngine,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StagedTransfer:
+    request_id: str
+    slot: SlotHandle
+    manifest: dict
+    scalar_fields: dict = field(default_factory=dict)
+
+
+@dataclass
+class PendingReceive:
+    request_id: str
+    slot: SlotHandle
+
+
+class DiffusionTransferManager:
+    """Manages tensor transfers for a single role instance.
+
+    Owns a TransferTensorBuffer (memory pool) and a BaseTransferEngine (RDMA or mock).
+    """
+
+    def __init__(
+        self,
+        engine: BaseTransferEngine,
+        buffer: TransferTensorBuffer,
+    ):
+        self._engine = engine
+        self._buffer = buffer
+        self._lock = threading.Lock()
+
+        self._engine.register_buffer(self._buffer.pool_data_ptr, self._buffer.pool_size)
+
+        self._staged: dict[str, StagedTransfer] = {}
+        self._pending_receives: dict[str, PendingReceive] = {}
+
+        logger.info(
+            "DiffusionTransferManager initialized: session=%s, pool=%d bytes",
+            self._engine.session_id,
+            self._buffer.pool_size,
+        )
+
+    @property
+    def session_id(self) -> str:
+        return self._engine.session_id
+
+    @property
+    def pool_data_ptr(self) -> int:
+        return self._buffer.pool_data_ptr
+
+    @property
+    def pool_size(self) -> int:
+        return self._buffer.pool_size
+
+    def stage_tensors(
+        self,
+        request_id: str,
+        tensor_fields: dict[str, torch.Tensor | list[torch.Tensor] | None],
+        scalar_fields: dict | None = None,
+        stream: torch.cuda.Stream | None = None,
+    ) -> StagedTransfer | None:
+        """Stage GPU tensors into the local TransferBuffer. Returns None on allocation failure."""
+        total_size = 0
+        for name, t in tensor_fields.items():
+            if t is None:
+                continue
+            if isinstance(t, list):
+                for ti in t:
+                    total_size += ti.nelement() * ti.element_size()
+            else:
+                total_size += t.nelement() * t.element_size()
+
+        if total_size == 0:
+            staged = StagedTransfer(
+                request_id=request_id,
+                slot=None,
+                manifest={},
+                scalar_fields=scalar_fields or {},
+            )
+            with self._lock:
+                self._staged[request_id] = staged
+            return staged
+
+        slot = self._buffer.allocate(total_size, request_id)
+        if slot is None:
+            logger.warning(
+                "TransferManager: failed to allocate %d bytes for %s",
+                total_size,
+                request_id,
+            )
+            return None
+
+        manifest = self._buffer.write_tensors_from_gpu(slot, tensor_fields, stream)
+
+        if stream is not None:
+            stream.synchronize()
+        elif torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        staged = StagedTransfer(
+            request_id=request_id,
+            slot=slot,
+            manifest=manifest,
+            scalar_fields=scalar_fields or {},
+        )
+        with self._lock:
+            self._staged[request_id] = staged
+
+        logger.debug(
+            "TransferManager: staged %s (%d bytes, offset=%d)",
+            request_id,
+            total_size,
+            slot.offset,
+        )
+        return staged
+
+    def stage_tensors_async(
+        self,
+        request_id: str,
+        tensor_fields: dict[str, torch.Tensor | list[torch.Tensor] | None],
+        scalar_fields: dict | None = None,
+        stream: torch.cuda.Stream | None = None,
+    ) -> tuple[StagedTransfer | None, torch.cuda.Event | None]:
+        """Stage GPU tensors, returning a CUDA event instead of blocking.
+
+        Caller MUST wait on the event before reading buffer data.
+        """
+        total_size = 0
+        for name, t in tensor_fields.items():
+            if t is None:
+                continue
+            if isinstance(t, list):
+                for ti in t:
+                    total_size += ti.nelement() * ti.element_size()
+            else:
+                total_size += t.nelement() * t.element_size()
+
+        if total_size == 0:
+            staged = StagedTransfer(
+                request_id=request_id,
+                slot=None,
+                manifest={},
+                scalar_fields=scalar_fields or {},
+            )
+            with self._lock:
+                self._staged[request_id] = staged
+            return staged, None
+
+        slot = self._buffer.allocate(total_size, request_id)
+        if slot is None:
+            logger.warning(
+                "TransferManager: failed to allocate %d bytes for %s",
+                total_size,
+                request_id,
+            )
+            return None, None
+
+        manifest = self._buffer.write_tensors_from_gpu(slot, tensor_fields, stream)
+
+        d2h_event = None
+        if stream is not None:
+            d2h_event = torch.cuda.Event()
+            d2h_event.record(stream)
+        elif torch.cuda.is_available():
+            d2h_event = torch.cuda.Event()
+            d2h_event.record(torch.cuda.current_stream())
+
+        staged = StagedTransfer(
+            request_id=request_id,
+            slot=slot,
+            manifest=manifest,
+            scalar_fields=scalar_fields or {},
+        )
+        with self._lock:
+            self._staged[request_id] = staged
+
+        logger.debug(
+            "TransferManager: staged_async %s (%d bytes, offset=%d)",
+            request_id,
+            total_size,
+            slot.offset,
+        )
+        return staged, d2h_event
+
+    def load_tensors_async(
+        self,
+        request_id: str,
+        manifest: dict,
+        device: torch.device | str = "cuda",
+        stream: torch.cuda.Stream | None = None,
+    ) -> tuple[dict[str, torch.Tensor | list[torch.Tensor]], torch.cuda.Event | None]:
+        """Load tensors from receive slot to GPU, returning a CUDA event.
+
+        Caller MUST wait on the event before using the returned tensors.
+        """
+        with self._lock:
+            pending = self._pending_receives.get(request_id)
+
+        if pending is None:
+            raise ValueError(
+                f"TransferManager: no pending receive slot for {request_id}"
+            )
+
+        tensors = self._buffer.read_tensors_from_manifest(
+            pending.slot, manifest, device=device, stream=stream
+        )
+
+        load_event = None
+        if stream is not None:
+            load_event = torch.cuda.Event()
+            load_event.record(stream)
+        elif torch.cuda.is_available():
+            load_event = torch.cuda.Event()
+            load_event.record(torch.cuda.current_stream())
+
+        logger.debug(
+            "TransferManager: loaded_async %d tensor fields for %s to %s",
+            len(tensors),
+            request_id,
+            device,
+        )
+        return tensors, load_event
+
+    def push_to_peer(
+        self,
+        request_id: str,
+        dest_session_id: str,
+        dest_addr: int,
+        transfer_size: int,
+    ) -> bool:
+        """Push staged data to a remote peer's buffer via RDMA. Returns True on success."""
+        with self._lock:
+            staged = self._staged.get(request_id)
+
+        if staged is None:
+            logger.error("TransferManager: no staged transfer for %s", request_id)
+            return False
+
+        if staged.slot is None:
+            return True
+
+        src_addr = self._buffer.pool_data_ptr + staged.slot.offset
+        ret = self._engine.transfer_sync(
+            dest_session_id, src_addr, dest_addr, transfer_size
+        )
+
+        if ret == 0:
+            logger.debug(
+                "TransferManager: pushed %s (%d bytes) to %s",
+                request_id,
+                transfer_size,
+                dest_session_id,
+            )
+        else:
+            logger.error(
+                "TransferManager: RDMA push failed for %s (ret=%d)",
+                request_id,
+                ret,
+            )
+
+        return ret == 0
+
+    def free_staged(self, request_id: str) -> None:
+        with self._lock:
+            staged = self._staged.pop(request_id, None)
+
+        if staged and staged.slot is not None:
+            self._buffer.free(staged.slot)
+            logger.debug("TransferManager: freed staged slot for %s", request_id)
+
+    def allocate_receive_slot(
+        self, request_id: str, size: int
+    ) -> PendingReceive | None:
+        """Allocate a local buffer slot to receive incoming data."""
+        slot = self._buffer.allocate(size, request_id)
+        if slot is None:
+            logger.warning(
+                "TransferManager: failed to allocate receive slot (%d bytes) for %s",
+                size,
+                request_id,
+            )
+            return None
+
+        pending = PendingReceive(request_id=request_id, slot=slot)
+        with self._lock:
+            self._pending_receives[request_id] = pending
+
+        logger.debug(
+            "TransferManager: allocated receive slot for %s (offset=%d, size=%d)",
+            request_id,
+            slot.offset,
+            slot.size,
+        )
+        return pending
+
+    def load_tensors(
+        self,
+        request_id: str,
+        manifest: dict,
+        device: torch.device | str = "cuda",
+        stream: torch.cuda.Stream | None = None,
+    ) -> dict[str, torch.Tensor | list[torch.Tensor]]:
+        """Load tensors from a receive slot into GPU memory."""
+        with self._lock:
+            pending = self._pending_receives.get(request_id)
+
+        if pending is None:
+            raise ValueError(
+                f"TransferManager: no pending receive slot for {request_id}"
+            )
+
+        tensors = self._buffer.read_tensors_from_manifest(
+            pending.slot, manifest, device=device, stream=stream
+        )
+
+        if stream is not None:
+            stream.synchronize()
+        elif torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        logger.debug(
+            "TransferManager: loaded %d tensor fields for %s to %s",
+            len(tensors),
+            request_id,
+            device,
+        )
+        return tensors
+
+    def register_prealloc_as_receive(
+        self, request_id: str, slot: "SlotHandle"
+    ) -> "PendingReceive":
+        """Register a pre-allocated slot as a pending receive (fast path)."""
+        pending = PendingReceive(request_id=request_id, slot=slot)
+        with self._lock:
+            self._pending_receives[request_id] = pending
+        return pending
+
+    def free_receive_slot(self, request_id: str) -> None:
+        with self._lock:
+            pending = self._pending_receives.pop(request_id, None)
+
+        if pending:
+            self._buffer.free(pending.slot)
+            logger.debug("TransferManager: freed receive slot for %s", request_id)
+
+    def get_receive_slot_addr(self, request_id: str) -> int | None:
+        with self._lock:
+            pending = self._pending_receives.get(request_id)
+        if pending is None:
+            return None
+        return self._buffer.pool_data_ptr + pending.slot.offset
+
+    def get_receive_slot_offset(self, request_id: str) -> int | None:
+        with self._lock:
+            pending = self._pending_receives.get(request_id)
+        if pending is None:
+            return None
+        return pending.slot.offset
+
+    def get_staged_info(self, request_id: str) -> StagedTransfer | None:
+        with self._lock:
+            return self._staged.get(request_id)
+
+    def free_slots_count(self, typical_size: int = 64 * 1024 * 1024) -> int:
+        return self._buffer.free_slots_count(typical_size)
+
+    def cleanup(self) -> None:
+        self._engine.deregister_buffer(self._buffer.pool_data_ptr)
+        logger.info("DiffusionTransferManager cleaned up")

python/sglang/multimodal_gen/runtime/disaggregation/transport/protocol.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Transfer protocol messages for disaggregated diffusion.
+
+All messages are sent as ZMQ multipart with a b"__transfer__" discriminator
+in frame[0] and JSON payload in frame[1].
+"""
+
+import json
+import logging
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+TRANSFER_MAGIC = b"__transfer__"
+
+
+class TransferMsgType:
+    # Instance → DiffusionServer
+    STAGED = "transfer_staged"
+    ALLOCATED = "transfer_allocated"
+    PUSHED = "transfer_pushed"
+    DONE = "transfer_done"
+
+    # DiffusionServer → Instance
+    ALLOC = "transfer_alloc"
+    PUSH = "transfer_push"
+    READY = "transfer_ready"
+
+    # Registration
+    REGISTER = "transfer_register"
+    REGISTER_ACK = "transfer_register_ack"
+
+
+@dataclass
+class TransferStagedMsg:
+    msg_type: str = TransferMsgType.STAGED
+    request_id: str = ""
+    data_size: int = 0
+    manifest: dict = None
+    session_id: str = ""
+    pool_ptr: int = 0
+    slot_offset: int = 0
+
+    def __post_init__(self):
+        if self.manifest is None:
+            self.manifest = {}
+
+
+@dataclass
+class TransferAllocMsg:
+    msg_type: str = TransferMsgType.ALLOC
+    request_id: str = ""
+    data_size: int = 0
+    source_role: str = ""
+
+
+@dataclass
+class TransferAllocatedMsg:
+    msg_type: str = TransferMsgType.ALLOCATED
+    request_id: str = ""
+    session_id: str = ""
+    pool_ptr: int = 0
+    slot_offset: int = 0
+    slot_size: int = 0
+
+
+@dataclass
+class TransferPushMsg:
+    msg_type: str = TransferMsgType.PUSH
+    request_id: str = ""
+    dest_session_id: str = ""
+    dest_addr: int = 0
+    transfer_size: int = 0
+
+
+@dataclass
+class TransferPushedMsg:
+    msg_type: str = TransferMsgType.PUSHED
+    request_id: str = ""
+
+
+@dataclass
+class TransferReadyMsg:
+    msg_type: str = TransferMsgType.READY
+    request_id: str = ""
+    manifest: dict = None
+    slot_offset: int = 0
+    scalar_fields: dict = None
+
+    def __post_init__(self):
+        if self.manifest is None:
+            self.manifest = {}
+        if self.scalar_fields is None:
+            self.scalar_fields = {}
+
+
+@dataclass
+class TransferDoneMsg:
+    msg_type: str = TransferMsgType.DONE
+    request_id: str = ""
+    error: str | None = None
+
+
+@dataclass
+class TransferRegisterMsg:
+    msg_type: str = TransferMsgType.REGISTER
+    role: str = ""
+    session_id: str = ""
+    pool_ptr: int = 0
+    pool_size: int = 0
+    # The instance's own work endpoint (e.g. tcp://host:port). Used by the
+    # DiffusionServer to key peer info by URL index (i.e. the same index used
+    # to build the PUSH work-socket list), so the control plane and the RDMA
+    # data plane cannot drift when instances register in a different order
+    # than --*-urls.
+    work_endpoint: str = ""
+    # Pre-allocated receive slots: [{"offset": int, "size": int, "slot_id": int, "addr": int}]
+    preallocated_slots: list = field(default_factory=list)
+
+
+def encode_transfer_msg(msg: Any) -> list[bytes]:
+    """Encode as [TRANSFER_MAGIC, json_payload_bytes]."""
+    if hasattr(msg, "__dataclass_fields__"):
+        d = asdict(msg)
+    elif isinstance(msg, dict):
+        d = msg
+    else:
+        raise TypeError(f"Cannot encode transfer message: {type(msg)}")
+
+    return [TRANSFER_MAGIC, json.dumps(d, separators=(",", ":")).encode("utf-8")]
+
+
+def decode_transfer_msg(frames: list[bytes]) -> dict:
+    if len(frames) < 2 or frames[0] != TRANSFER_MAGIC:
+        raise ValueError(f"Not a transfer message: frame[0]={frames[0]!r}")
+    return json.loads(frames[1])
+
+
+def is_transfer_message(frames: list) -> bool:
+    return len(frames) >= 2 and (
+        frames[0] == TRANSFER_MAGIC
+        or (isinstance(frames[0], memoryview) and bytes(frames[0]) == TRANSFER_MAGIC)
+        or (hasattr(frames[0], "bytes") and frames[0].bytes == TRANSFER_MAGIC)
+    )

python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py

@@ -435,9 +435,9 @@ class GroupCoordinator:
         # Bypass the function if we are using only 1 GPU.
         if self.world_size == 1:
             return obj
-        if self.shm_broadcaster is not None:
+        if self.mq_broadcaster is not None:
             assert src == 0, "Shared memory broadcaster only supports src=0"
-            return self.shm_broadcaster.broadcast_object(obj)
+            return self.mq_broadcaster.broadcast_object(obj)
         if self.rank_in_group == src:
             torch.distributed.broadcast_object_list(
                 [obj], src=self.ranks[src], group=self.cpu_group

python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py

@@ -8,7 +8,9 @@ from typing import cast
 
 from sglang.multimodal_gen.apps.webui import run_sgl_diffusion_webui
 from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand
-from sglang.multimodal_gen.runtime.launch_server import launch_server
+from sglang.multimodal_gen.runtime.launch_server import (
+    dispatch_launch,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import FlexibleArgumentParser
@@ -31,7 +33,8 @@ def add_multimodal_gen_serve_args(parser: argparse.ArgumentParser):
 def execute_serve_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None):
     """The entry point for the serve command."""
     server_args = ServerArgs.from_cli_args(args, unknown_args)
-    launch_server(server_args)
+
+    dispatch_launch(server_args)
 
     if server_args.webui:
         run_sgl_diffusion_webui(server_args)

python/sglang/multimodal_gen/runtime/entrypoints/http_server.py

@@ -162,6 +162,32 @@ async def health_generate():
     return {"status": "ok"}
 
 
+@health_router.get("/stats")
+async def stats_endpoint(request: Request):
+    """Get runtime statistics including disagg pipeline metrics.
+
+    Returns queue depth, request counts, latency, throughput, etc.
+    Sends a GetDisaggStatsReq to the scheduler via ZMQ and returns the result.
+    """
+    from sglang.multimodal_gen.runtime.entrypoints.utils import GetDisaggStatsReq
+
+    server_args: ServerArgs = request.app.state.server_args
+    response: dict = {
+        "status": "ok",
+        "model_path": server_args.model_path,
+    }
+
+    # Query the scheduler for disagg metrics
+    try:
+        stats_response = await async_scheduler_client.forward(GetDisaggStatsReq())
+        if hasattr(stats_response, "output") and stats_response.output is not None:
+            response["disagg"] = stats_response.output
+    except Exception as e:
+        response["disagg"] = {"error": str(e)}
+
+    return response
+
+
 def make_serializable(obj):
     """Recursively converts Tensors to None for JSON serialization."""
     if isinstance(obj, torch.Tensor):

python/sglang/multimodal_gen/runtime/entrypoints/utils.py

@@ -69,6 +69,13 @@ class ShutdownReq:
     pass
 
 
+@dataclass
+class GetDisaggStatsReq:
+    """Request to get disagg pipeline metrics from the scheduler."""
+
+    pass
+
+
 def format_lora_message(
     lora_nickname: Union[str, List[str]],
     target: Union[str, List[str]],

python/sglang/multimodal_gen/runtime/launch_server.py

@@ -1,5 +1,6 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
+import dataclasses
 import multiprocessing as mp
 import os
 import signal
@@ -9,6 +10,10 @@ import threading
 import psutil
 import uvicorn
 
+from sglang.multimodal_gen.runtime.disaggregation.orchestrator import (
+    DiffusionServer,
+)
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.entrypoints.http_server import create_app
 from sglang.multimodal_gen.runtime.managers.gpu_worker import run_scheduler_process
 from sglang.multimodal_gen.runtime.server_args import (
@@ -16,9 +21,28 @@ from sglang.multimodal_gen.runtime.server_args import (
     prepare_server_args,
     set_global_server_args,
 )
+from sglang.multimodal_gen.runtime.utils.common import is_port_available
 from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger
 
 
+def _find_available_port(
+    start: int = 10000, avoid: set[int] | None = None, max_attempts: int = 100
+) -> int:
+    """Find an available port starting from *start*, skipping ports in *avoid*."""
+    if avoid is None:
+        avoid = set()
+    port = max(1024, min(start, 65535))
+    for _ in range(max_attempts):
+        if port not in avoid and is_port_available(port):
+            return port
+        port += 1
+        if port > 65535:
+            port = 1024
+    raise RuntimeError(
+        f"No available port found after {max_attempts} attempts (start={start})"
+    )
+
+
 def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
     """Kill the process and all its child processes."""
     # Remove sigchld handler to avoid spammy logs.
@@ -185,6 +209,237 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True):
     return processes
 
 
+def launch_pool_disagg_server(
+    server_args: ServerArgs,
+    encoder_gpus: list[list[int]],
+    denoiser_gpus: list[list[int]],
+    decoder_gpus: list[list[int]],
+    launch_http_server: bool = True,
+):
+    """Launch a pool-based disaggregated server with N:M:K independent role instances.
+
+    DiffusionServer orchestrates the full pipeline, dispatching at every
+    role transition (Encoder → Denoiser → Decoder).
+
+    Args:
+        server_args: Base server configuration
+        encoder_gpus: List of GPU ID lists, one per encoder instance.
+            e.g., [[0], [2]] for 2 encoder instances on GPUs 0 and 2.
+        denoiser_gpus: List of GPU ID lists, one per denoiser instance.
+            e.g., [[1], [3]] for 2 denoiser instances.
+        decoder_gpus: List of GPU ID lists, one per decoder instance.
+            e.g., [[0], [2]] for 2 decoder instances (can share with encoder).
+        launch_http_server: Whether to launch the HTTP server.
+
+    Example:
+        launch_pool_disagg_server(server_args,
+            encoder_gpus=[[0], [2]],
+            denoiser_gpus=[[1], [3]],
+            decoder_gpus=[[0], [2]],
+        )
+    """
+    configure_logger(server_args)
+
+    num_encoders = len(encoder_gpus)
+    num_denoisers = len(denoiser_gpus)
+    num_decoders = len(decoder_gpus)
+    logger.info(
+        "Starting pool disagg server: %d encoder(s), %d denoiser(s), %d decoder(s)...",
+        num_encoders,
+        num_denoisers,
+        num_decoders,
+    )
+
+    host = server_args.host or "127.0.0.1"
+
+    def find_port(start):
+        return _find_available_port(start)
+
+    # Allocate endpoints
+    port_cursor = server_args.scheduler_port + 3000
+
+    # Per-instance work endpoints (instance binds PULL, DS connects PUSH)
+    encoder_work_endpoints = []
+    for i in range(num_encoders):
+        p = find_port(port_cursor)
+        encoder_work_endpoints.append(f"tcp://{host}:{p}")
+        port_cursor = p + 1
+
+    denoiser_work_endpoints = []
+    for i in range(num_denoisers):
+        p = find_port(port_cursor)
+        denoiser_work_endpoints.append(f"tcp://{host}:{p}")
+        port_cursor = p + 1
+
+    decoder_work_endpoints = []
+    for i in range(num_decoders):
+        p = find_port(port_cursor)
+        decoder_work_endpoints.append(f"tcp://{host}:{p}")
+        port_cursor = p + 1
+
+    # Per-role-type result endpoints (DS binds PULL, instances connect PUSH)
+    # Use deterministic convention: scheduler_port + {1,2,3}
+    base_port = server_args.scheduler_port
+    encoder_result_ep = f"tcp://{host}:{base_port + 1}"
+    denoiser_result_ep = f"tcp://{host}:{base_port + 2}"
+    decoder_result_ep = f"tcp://{host}:{base_port + 3}"
+
+    logger.info(
+        "Pool endpoints allocated: %d work + 3 result endpoints",
+        num_encoders + num_denoisers + num_decoders,
+    )
+
+    # Launch all role instances
+    all_processes = []
+
+    role_configs = [
+        (RoleType.ENCODER, encoder_gpus, encoder_work_endpoints, encoder_result_ep),
+        (
+            RoleType.DENOISER,
+            denoiser_gpus,
+            denoiser_work_endpoints,
+            denoiser_result_ep,
+        ),
+        (RoleType.DECODER, decoder_gpus, decoder_work_endpoints, decoder_result_ep),
+    ]
+
+    for role_type, gpu_lists, work_eps, result_ep in role_configs:
+        for inst_idx, gpu_ids in enumerate(gpu_lists):
+            num_role_gpus = len(gpu_ids)
+
+            # Per-role parallelism: use explicit overrides if set, else None (auto-derive)
+            role_par = server_args.get_role_parallelism(role_type)
+
+            role_overrides = {
+                "disagg_role": role_type,
+                "disagg_mode": True,
+                "pool_work_endpoint": work_eps[inst_idx],
+                "pool_result_endpoint": result_ep,
+                "num_gpus": num_role_gpus,
+                "warmup": role_type == RoleType.ENCODER,
+                "scheduler_port": find_port(port_cursor),
+                "master_port": find_port(port_cursor + 100),
+                # Per-role parallelism (None = auto-derive from num_gpus)
+                "tp_size": role_par["tp_size"],
+                "sp_degree": role_par["sp_degree"],
+                "ulysses_degree": role_par["ulysses_degree"],
+                "ring_degree": role_par["ring_degree"],
+            }
+            port_cursor = role_overrides["master_port"] + 100
+
+            base_dict = {
+                f.name: getattr(server_args, f.name)
+                for f in dataclasses.fields(server_args)
+            }
+            base_dict.update(role_overrides)
+            base_dict.pop("pipeline_config", None)
+            role_args = ServerArgs.from_kwargs(**base_dict)
+
+            pool_ctx = mp.get_context("spawn")
+            inst_readers = []
+
+            # Spawn all ranks first — NCCL init blocks until all ranks connect
+            for rank_idx in range(num_role_gpus):
+                reader, writer = pool_ctx.Pipe(duplex=False)
+                gpu_id = gpu_ids[rank_idx]
+
+                process = pool_ctx.Process(
+                    target=_run_disagg_role_process,
+                    args=(gpu_id, rank_idx, rank_idx, role_args, writer, [], []),
+                    name=f"sglang-pool-{role_type.value}-{inst_idx}-r{rank_idx}",
+                    daemon=True,
+                )
+                process.start()
+                all_processes.append(process)
+                inst_readers.append(reader)
+
+            # Wait for all ranks to be ready (after all are spawned)
+            for rank_idx, reader in enumerate(inst_readers):
+                try:
+                    data = reader.recv()
+                except EOFError:
+                    logger.error(
+                        "Pool %s[%d] rank %d is dead.",
+                        role_type.value,
+                        inst_idx,
+                        rank_idx,
+                    )
+                    raise
+                if data.get("status") != "ready":
+                    raise RuntimeError(
+                        f"Pool {role_type.value}[{inst_idx}] rank {rank_idx} "
+                        "failed to initialize."
+                    )
+                reader.close()
+
+            logger.info(
+                "Pool %s[%d] ready on GPU(s) %s (work=%s)",
+                role_type.value.upper(),
+                inst_idx,
+                gpu_ids,
+                work_eps[inst_idx],
+            )
+
+    logger.info("All pool role instances ready")
+
+    # Start DiffusionServer
+    frontend_endpoint = f"tcp://{host}:{server_args.scheduler_port}"
+
+    diffusion_server = DiffusionServer(
+        frontend_endpoint=frontend_endpoint,
+        encoder_work_endpoints=encoder_work_endpoints,
+        denoiser_work_endpoints=denoiser_work_endpoints,
+        decoder_work_endpoints=decoder_work_endpoints,
+        encoder_result_endpoint=encoder_result_ep,
+        denoiser_result_endpoint=denoiser_result_ep,
+        decoder_result_endpoint=decoder_result_ep,
+        dispatch_policy_name=server_args.disagg_dispatch_policy,
+        timeout_s=float(server_args.disagg_timeout),
+    )
+    diffusion_server.start()
+
+    if not diffusion_server.wait_ready(timeout=30.0):
+        raise RuntimeError("DiffusionServer failed to bind sockets within 30 seconds")
+
+    if launch_http_server:
+        logger.info(
+            "Starting FastAPI server (connected to DiffusionServer at port %d).",
+            server_args.scheduler_port,
+        )
+        launch_http_server_only(server_args)
+
+    return all_processes
+
+
+def _run_disagg_role_process(
+    gpu_id: int,
+    _local_rank: int,
+    rank: int,
+    server_args: ServerArgs,
+    pipe_writer: mp.connection.Connection,
+    task_pipes: list,
+    result_pipes: list,
+):
+    """Entry point for a disagg role process.
+
+    Uses the physical GPU index (gpu_id) as local_rank so that
+    torch.cuda.set_device(local_rank) selects the correct GPU.
+    This avoids relying on CUDA_VISIBLE_DEVICES remapping, which
+    may not work if CUDA was pre-initialized in the parent process.
+    """
+    run_scheduler_process(
+        local_rank=gpu_id,
+        rank=rank,
+        master_port=server_args.master_port,
+        server_args=server_args,
+        pipe_writer=pipe_writer,
+        task_pipe_r=None,
+        result_pipe_w=None,
+        task_pipes_to_slaves=task_pipes,
+        result_pipes_from_slaves=result_pipes,
+    )
+
+
 def launch_http_server_only(server_args):
     # set for endpoints to access global_server_args
     set_global_server_args(server_args)
@@ -199,10 +454,224 @@ def launch_http_server_only(server_args):
     )
 
 
+def parse_url_string(url_str: str) -> list[str]:
+    """Parse a semicolon-separated URL string into a list.
+
+    Example: "tcp://10.0.0.1:35000;tcp://10.0.0.2:35000" -> ["tcp://...", "tcp://..."]
+    """
+    return [u.strip() for u in url_str.split(";") if u.strip()]
+
+
+def launch_disagg_server(server_args: ServerArgs):
+    """Launch DiffusionServer head node + HTTP server (--disagg-role server).
+
+    No GPU workers are spawned. Connects to remote role instances
+    specified by --encoder-urls, --denoiser-urls, --decoder-urls.
+
+    Result endpoints use deterministic convention:
+        encoder result: scheduler_port + 1
+        denoiser result: scheduler_port + 2
+        decoder result: scheduler_port + 3
+    """
+    configure_logger(server_args)
+
+    for name, val in [
+        ("--encoder-urls", server_args.encoder_urls),
+        ("--denoiser-urls", server_args.denoiser_urls),
+        ("--decoder-urls", server_args.decoder_urls),
+    ]:
+        if val is None:
+            raise ValueError(f"{name} is required for --disagg-role server")
+
+    host = server_args.host or "127.0.0.1"
+    base_port = server_args.scheduler_port
+
+    encoder_work_endpoints = parse_url_string(server_args.encoder_urls)
+    denoiser_work_endpoints = parse_url_string(server_args.denoiser_urls)
+    decoder_work_endpoints = parse_url_string(server_args.decoder_urls)
+
+    encoder_result_ep = f"tcp://{host}:{base_port + 1}"
+    denoiser_result_ep = f"tcp://{host}:{base_port + 2}"
+    decoder_result_ep = f"tcp://{host}:{base_port + 3}"
+
+    frontend_endpoint = f"tcp://{host}:{base_port}"
+
+    logger.info(
+        "Starting DiffusionServer: %d encoder(s), %d denoiser(s), %d decoder(s)",
+        len(encoder_work_endpoints),
+        len(denoiser_work_endpoints),
+        len(decoder_work_endpoints),
+    )
+    logger.info("  Frontend: %s", frontend_endpoint)
+    logger.info("  Encoder work endpoints: %s", encoder_work_endpoints)
+    logger.info("  Denoiser work endpoints: %s", denoiser_work_endpoints)
+    logger.info("  Decoder work endpoints: %s", decoder_work_endpoints)
+    logger.info(
+        "  Result endpoints: encoder=%s, denoiser=%s, decoder=%s",
+        encoder_result_ep,
+        denoiser_result_ep,
+        decoder_result_ep,
+    )
+
+    diffusion_server = DiffusionServer(
+        frontend_endpoint=frontend_endpoint,
+        encoder_work_endpoints=encoder_work_endpoints,
+        denoiser_work_endpoints=denoiser_work_endpoints,
+        decoder_work_endpoints=decoder_work_endpoints,
+        encoder_result_endpoint=encoder_result_ep,
+        denoiser_result_endpoint=denoiser_result_ep,
+        decoder_result_endpoint=decoder_result_ep,
+        dispatch_policy_name=server_args.disagg_dispatch_policy,
+        timeout_s=float(server_args.disagg_timeout),
+    )
+    diffusion_server.start()
+
+    if not diffusion_server.wait_ready(timeout=30.0):
+        raise RuntimeError("DiffusionServer failed to bind sockets within 30 seconds")
+
+    logger.info(
+        "Starting HTTP server (connected to DiffusionServer at port %d).",
+        base_port,
+    )
+    launch_http_server_only(server_args)
+
+
+def launch_disagg_role(server_args: ServerArgs):
+    """Launch a standalone disaggregated role instance (--disagg-role encoder/denoising/decoder).
+
+    The instance:
+    1. Binds its work PULL socket on tcp://0.0.0.0:{scheduler_port}
+    2. Connects its result PUSH socket to the DiffusionServer head node
+       (derived from --disagg-server-addr + role offset)
+    3. Spawns GPU worker processes for the assigned role.
+    """
+    configure_logger(server_args)
+
+    role_type = server_args.disagg_role
+    if server_args.disagg_server_addr is None:
+        raise ValueError(
+            "--disagg-server-addr is required for --disagg-role " f"{role_type.value}"
+        )
+
+    # Derive endpoints
+    work_endpoint = server_args.derive_pool_work_endpoint()
+    result_endpoint = server_args.derive_pool_result_endpoint()
+
+    logger.info(
+        "Starting disagg role: %s, num_gpus=%d",
+        role_type.value,
+        server_args.num_gpus,
+    )
+    logger.info("  Work endpoint (bind): %s", work_endpoint)
+    logger.info("  Result endpoint (connect): %s", result_endpoint)
+    logger.info(
+        "  P2P: hostname=%s, ib_device=%s, pool_size=%d",
+        server_args.disagg_p2p_hostname,
+        server_args.disagg_ib_device,
+        server_args.disagg_transfer_pool_size,
+    )
+
+    # Build role-specific ServerArgs
+    # Use a different port for the scheduler's internal ROUTER socket to avoid
+    # conflicting with the pool work PULL socket (both bind on scheduler_port).
+    internal_scheduler_port = _find_available_port(
+        start=server_args.scheduler_port + 100, avoid={server_args.scheduler_port}
+    )
+
+    role_par = server_args.get_role_parallelism(role_type)
+    role_overrides = {
+        "disagg_role": role_type,
+        "disagg_mode": True,
+        "pool_work_endpoint": work_endpoint,
+        "pool_result_endpoint": result_endpoint,
+        "warmup": role_type == RoleType.ENCODER,
+        "scheduler_port": internal_scheduler_port,
+        # Per-role parallelism (None = auto-derive from num_gpus)
+        "tp_size": role_par["tp_size"],
+        "sp_degree": role_par["sp_degree"],
+        "ulysses_degree": role_par["ulysses_degree"],
+        "ring_degree": role_par["ring_degree"],
+    }
+
+    base_dict = {
+        f.name: getattr(server_args, f.name) for f in dataclasses.fields(server_args)
+    }
+    base_dict.update(role_overrides)
+    base_dict.pop("pipeline_config", None)
+    role_args = ServerArgs.from_kwargs(**base_dict)
+
+    # Spawn GPU worker processes
+    # NOTE: All ranks must be spawned before waiting for ready signals,
+    # because NCCL init_process_group blocks until all ranks connect.
+    num_gpus = server_args.num_gpus
+    base_gpu_id = server_args.base_gpu_id
+    pool_ctx = mp.get_context("spawn")
+    processes = []
+    readers = []
+
+    for rank_idx in range(num_gpus):
+        reader, writer = pool_ctx.Pipe(duplex=False)
+        gpu_id = base_gpu_id + rank_idx
+
+        process = pool_ctx.Process(
+            target=_run_disagg_role_process,
+            args=(gpu_id, rank_idx, rank_idx, role_args, writer, [], []),
+            name=f"sglang-{role_type.value}-r{rank_idx}",
+            daemon=True,
+        )
+        process.start()
+        processes.append(process)
+        readers.append(reader)
+
+    # Wait for all ranks to be ready (after all are spawned)
+    for rank_idx, reader in enumerate(readers):
+        try:
+            data = reader.recv()
+        except EOFError:
+            logger.error(
+                "Role %s rank %d is dead.",
+                role_type.value,
+                rank_idx,
+            )
+            raise
+        if data.get("status") != "ready":
+            raise RuntimeError(
+                f"Role {role_type.value} rank {rank_idx} failed to initialize."
+            )
+        reader.close()
+
+    logger.info(
+        "Role %s ready (%d GPU(s), work=%s)",
+        role_type.value.upper(),
+        num_gpus,
+        work_endpoint,
+    )
+
+    # Block until interrupted
+    try:
+        for p in processes:
+            p.join()
+    except KeyboardInterrupt:
+        logger.info("Role %s shutting down.", role_type.value)
+
+
+def dispatch_launch(server_args: ServerArgs):
+    """Route to the correct launch function based on --disagg-role."""
+    role = server_args.disagg_role
+    if role == RoleType.MONOLITHIC:
+        launch_server(server_args)
+    elif role == RoleType.SERVER:
+        launch_disagg_server(server_args)
+    elif role in (RoleType.ENCODER, RoleType.DENOISER, RoleType.DECODER):
+        launch_disagg_role(server_args)
+    else:
+        raise ValueError(f"Unknown disagg_role: {role}")
+
+
 if __name__ == "__main__":
     server_args = prepare_server_args(sys.argv[1:])
 
     try:
-        launch_server(server_args)
+        dispatch_launch(server_args)
     finally:
         kill_process_tree(os.getpid(), include_parent=False)

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -208,9 +208,16 @@ class GPUWorker:
             f"Related offload server args to disable: {suggested_args_str}"
         )
 
-    def execute_forward(self, batch: List[Req]) -> OutputBatch:
+    def execute_forward(
+        self, batch: List[Req], return_req: bool = False
+    ) -> OutputBatch | Req:
         """
         Execute a forward pass.
+
+        Args:
+            batch: List of requests to process.
+            return_req: If True, return the raw Req instead of OutputBatch.
+                Used by disaggregated pipelines to access intermediate tensors.
         """
         assert self.pipeline is not None
         req = batch[0]
@@ -229,6 +236,11 @@ class GPUWorker:
             req.log(server_args=self.server_args)
             result = self.pipeline.forward(req, self.server_args)
 
+            # For disagg roles, return raw Req to let the caller handle
+            # the role-to-role tensor transfer before OutputBatch conversion.
+            if return_req and isinstance(result, Req):
+                return result
+
             if isinstance(result, Req):
                 output_batch = OutputBatch(
                     output=result.output,
@@ -261,7 +273,8 @@ class GPUWorker:
                 self.do_mem_analysis(output_batch)
 
             duration_ms = (time.monotonic() - start_time) * 1000
-            output_batch.metrics.total_duration_ms = duration_ms
+            if output_batch.metrics is not None:
+                output_batch.metrics.total_duration_ms = duration_ms
 
             # Save output to file and return file path only if requested. Avoid the serialization
             # and deserialization overhead between scheduler_client and gpu_worker.
@@ -526,6 +539,7 @@ def run_scheduler_process(
             port_args=port_args,
             task_pipes_to_slaves=task_pipes_to_slaves,
             result_pipes_from_slaves=result_pipes_from_slaves,
+            local_rank=local_rank,
         )
         logger.info(f"Worker {rank}: Scheduler loop started.")
         pipe_writer.send(

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -10,6 +10,10 @@ from typing import Any, List
 
 import zmq
 
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.disaggregation.scheduler_mixin import (
+    SchedulerDisaggMixin,
+)
 from sglang.multimodal_gen.runtime.distributed import get_world_group
 from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
     _parse_size,
@@ -20,6 +24,7 @@ from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import (
     UpdateWeightFromDiskReqInput,
 )
 from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    GetDisaggStatsReq,
     ListLorasReq,
     MergeLoraWeightsReq,
     SetLoraReq,
@@ -43,7 +48,7 @@ logger = init_logger(__name__)
 MINIMUM_PICTURE_BASE64_FOR_WARMUP = "data:image/jpg;base64,iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAACXBIWXMAAA7EAAAOxAGVKw4bAAAAbUlEQVRYhe3VsQ2AMAxE0Y/lIgNQULD/OqyCMgCihCKSG4yRuKuiNH6JLsoEbMACOGBcua9HOR7Y6w6swBwMy0qLTpkeI77qdEBpBFAHBBDAGH8WrwJKI4AAegUCfAKgEgpQDvh3CR3oQCuav58qlAw73kKCSgAAAABJRU5ErkJggg=="
 
 
-class Scheduler:
+class Scheduler(SchedulerDisaggMixin):
     """
     Runs the main event loop for the rank 0 worker.
     It listens for external requests via ZMQ and coordinates with other workers.
@@ -57,10 +62,17 @@ class Scheduler:
         port_args: PortArgs,
         task_pipes_to_slaves: list = None,
         result_pipes_from_slaves: list = None,
+        local_rank: int | None = None,
     ):
         self.server_args = server_args
         self.port_args = port_args
 
+        # local_rank is the physical GPU index for torch.cuda.set_device.
+        # In non-disagg mode, it equals gpu_id. In disagg mode, it may differ
+        # (e.g., denoiser rank 0 on physical GPU 1).
+        if local_rank is None:
+            local_rank = gpu_id
+
         set_global_server_args(server_args=server_args)
 
         # Inter-process Communication
@@ -76,7 +88,7 @@ class Scheduler:
             self.receiver = None
 
         worker = GPUWorker(
-            local_rank=gpu_id,
+            local_rank=local_rank,
             master_port=port_args.master_port,
             rank=gpu_id,
             server_args=server_args,
@@ -95,6 +107,7 @@ class Scheduler:
             List[Req]: self._handle_generation,
             ListLorasReq: self._handle_list_loras,
             ShutdownReq: self._handle_shutdown,
+            GetDisaggStatsReq: self._handle_get_disagg_stats,
             UpdateWeightFromDiskReqInput: self._handle_update_weights_from_disk,
             GetWeightsChecksumReqInput: self._handle_get_weights_checksum,
         }
@@ -114,6 +127,21 @@ class Scheduler:
         self._max_consecutive_errors = 3
         self._consecutive_error_count = 0
 
+        self._init_disagg_state(server_args, local_rank)
+
+    def get_disagg_metrics(self) -> dict | None:
+        """Return disagg role metrics snapshot, or None if not in disagg mode."""
+        if self._disagg_metrics is None:
+            return None
+        return self._disagg_metrics.snapshot().to_dict()
+
+    def _handle_get_disagg_stats(self, _reqs: List[Any]) -> OutputBatch:
+        """Handle stats request — return disagg metrics via OutputBatch.output."""
+        stats = self.get_disagg_metrics()
+        return OutputBatch(
+            output=stats or {"role": "monolithic", "message": "not in disagg mode"}
+        )
+
     def _handle_set_lora(self, reqs: List[Any]) -> OutputBatch:
         # TODO: return set status
         # TODO: return with SetLoRAResponse or something more appropriate
@@ -166,6 +194,7 @@ class Scheduler:
                 )
             else:
                 logger.info("Processing warmup req...")
+
         return self.worker.execute_forward(reqs)
 
     def return_result(
@@ -362,12 +391,20 @@ class Scheduler:
         The main event loop that listens for ZMQ requests.
         Handles abortion
         """
+        # Pool mode: all roles use the pool event loop
+        if self._disagg_role != RoleType.MONOLITHIC:
+            self._disagg_event_loop()
+            return
 
         logger.debug(
             f"Rank 0 scheduler listening on tcp://*:{self.server_args.scheduler_port}"
         )
 
         while self._running:
+            # Update queue depth for metrics
+            if self._disagg_metrics:
+                self._disagg_metrics.update_queue_depth(len(self.waiting_queue))
+
             # 1: receive requests
             try:
                 new_reqs = self.recv_reqs()
@@ -403,6 +440,10 @@ class Scheduler:
 
             try:
                 processed_req = reqs[0]
+                is_warmup = (
+                    processed_req.is_warmup if isinstance(processed_req, Req) else False
+                )
+
                 handler = self.request_handlers.get(type(processed_req))
                 if handler:
                     output_batch = handler(reqs)
@@ -415,16 +456,10 @@ class Scheduler:
                     f"Error executing request in scheduler event loop: {e}",
                     exc_info=True,
                 )
-                # Determine appropriate error response format
+                output_batch = OutputBatch(error=str(e))
-                output_batch = (
-                    OutputBatch(error=str(e))
-                    if reqs and isinstance(reqs[0], Req)
-                    else OutputBatch(error=str(e))
-                )
 
             # 3. return results
             try:
-                # log warmup info
                 is_warmup = (
                     processed_req.is_warmup if isinstance(processed_req, Req) else False
                 )
@@ -457,6 +492,7 @@ class Scheduler:
 
         if self.receiver is not None:
             self.receiver.close()
+        self._cleanup_disagg()
         self.context.destroy(linger=0)
 
     def _broadcast_task(self, payload: dict[str, Any]) -> None:

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -14,6 +14,10 @@ from typing import Any, Callable, Literal, cast
 import torch
 from tqdm import tqdm
 
+from sglang.multimodal_gen.runtime.disaggregation.roles import (
+    RoleType,
+    filter_modules_for_role,
+)
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     PipelineComponentLoader,
 )
@@ -82,6 +86,7 @@ class ComposedPipelineBase(ABC):
         use. The pipeline should be stateless and not hold any batch state.
         """
         self.server_args = server_args
+        self._disagg_role = server_args.disagg_role
 
         self.model_path: str = model_path
         self._stages: list[PipelineStage] = []
@@ -94,6 +99,20 @@ class ComposedPipelineBase(ABC):
         if self._required_config_modules is None:
             raise NotImplementedError("Subclass must set _required_config_modules")
 
+        # Filter modules based on disaggregation role
+        if self._disagg_role != RoleType.MONOLITHIC:
+            original_modules = list(self._required_config_modules)
+            self._required_config_modules = filter_modules_for_role(
+                self._required_config_modules, self._disagg_role
+            )
+            skipped = set(original_modules) - set(self._required_config_modules)
+            if skipped:
+                logger.info(
+                    "Disagg role=%s: skipping modules %s",
+                    self._disagg_role.value,
+                    sorted(skipped),
+                )
+
         # [module_name, gpu memory usage]
         self.memory_usages: dict[str, float] = {}
         # Load modules directly in initialization
@@ -169,6 +188,70 @@ class ComposedPipelineBase(ABC):
         """
         return
 
+    # --- Config-name → pipeline_config attribute mapping ---
+    _CONFIG_ATTR_MAP: dict[str, str] = {
+        "vae": "vae_config",
+        "video_vae": "vae_config",
+        "audio_vae": "audio_vae_config",
+    }
+
+    def _init_skipped_component_configs(
+        self,
+        full_model_index: dict[str, Any],
+        server_args: ServerArgs,
+    ) -> None:
+        """Read HF JSON configs for skipped components and run
+        update_model_arch + post_init so pipeline_config is fully
+        initialized without loading weights.
+        """
+        from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+            get_diffusers_component_config,
+        )
+
+        required = set(self.required_config_modules)
+        for module_name in full_model_index:
+            if module_name in required:
+                continue  # will be loaded normally
+            cfg_attr = self._CONFIG_ATTR_MAP.get(module_name)
+            if cfg_attr is None:
+                continue  # not a config we need to patch
+
+            pipeline_cfg = getattr(server_args.pipeline_config, cfg_attr, None)
+            if pipeline_cfg is None:
+                continue
+
+            try:
+                component_path = self._resolve_component_path(
+                    server_args, module_name, module_name
+                )
+                hf_config = get_diffusers_component_config(
+                    component_path=component_path
+                )
+                hf_config.pop("_class_name", None)
+                hf_config.pop("_diffusers_version", None)
+                pipeline_cfg.update_model_arch(hf_config)
+                if hasattr(pipeline_cfg, "post_init"):
+                    pipeline_cfg.post_init()
+                logger.info(
+                    "Disagg role=%s: initialized %s config from HF JSON "
+                    "(spatial_compression_ratio=%s)",
+                    self._disagg_role.value,
+                    module_name,
+                    getattr(
+                        getattr(pipeline_cfg, "arch_config", None),
+                        "spatial_compression_ratio",
+                        "N/A",
+                    ),
+                )
+            except Exception as e:
+                logger.warning(
+                    "Disagg role=%s: failed to read HF config for skipped "
+                    "component %s: %s",
+                    self._disagg_role.value,
+                    module_name,
+                    e,
+                )
+
     def _resolve_component_path(
         self, server_args: ServerArgs, module_name: str, load_module_name: str
     ) -> str:
@@ -214,7 +297,24 @@ class ComposedPipelineBase(ABC):
                     "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..."
                 )
                 if "transformer_2" not in self.required_config_modules:
-                    self.required_config_modules.append("transformer_2")
+                    # Re-apply disagg role filter: only add transformer_2 if the
+                    # role actually needs denoising modules.
+                    from sglang.multimodal_gen.runtime.disaggregation.roles import (
+                        get_module_role,
+                    )
+
+                    module_role = get_module_role("transformer_2")
+                    if (
+                        self._disagg_role == RoleType.MONOLITHIC
+                        or module_role is None
+                        or module_role == self._disagg_role
+                    ):
+                        self.required_config_modules.append("transformer_2")
+                    else:
+                        logger.info(
+                            "Disagg role=%s: skipping dynamically added module transformer_2",
+                            self._disagg_role.value,
+                        )
             else:
                 logger.info(
                     "Boundary ratio found in model_index.json without transformers; "
@@ -237,6 +337,11 @@ class ComposedPipelineBase(ABC):
             len(model_index) > 1
         ), "model_index.json must contain at least one pipeline module"
 
+        # In disagg mode, read HF config for skipped components (e.g., VAE)
+        # so that update_model_arch + post_init can derive pipeline_config.
+        if self._disagg_role != RoleType.MONOLITHIC:
+            self._init_skipped_component_configs(model_index, server_args)
+
         model_index = {
             required_module: model_index[required_module]
             for required_module in self.required_config_modules
@@ -341,6 +446,19 @@ class ComposedPipelineBase(ABC):
 
         assert self.modules is not None, "No modules are registered"
 
+        # Filter stages based on disaggregation role
+        if self._disagg_role != RoleType.MONOLITHIC:
+            if stage.role_affinity != self._disagg_role:
+                if stage_name is None:
+                    stage_name = self._infer_stage_name(stage)
+                logger.info(
+                    "Disagg role=%s: skipping stage %s (affinity=%s)",
+                    self._disagg_role.value,
+                    stage_name,
+                    stage.role_affinity.value,
+                )
+                return self
+
         if stage_name is None:
             stage_name = self._infer_stage_name(stage)
         if stage_name in self._stage_name_mapping:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py

@@ -13,6 +13,7 @@ from enum import Enum, auto
 
 import torch
 
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
@@ -103,6 +104,11 @@ class PipelineStage(ABC):
         """
         pass
 
+    # Default role affinity: ENCODER. Override in subclasses for DENOISING/DECODER.
+    @property
+    def role_affinity(self) -> RoleType:
+        return RoleType.ENCODER
+
     # execute on all ranks by default
     @property
     def parallelism_type(self) -> StageParallelismType:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -56,6 +56,12 @@ class DecodingStage(PipelineStage):
     output format (e.g., pixel values).
     """
 
+    @property
+    def role_affinity(self):
+        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+
+        return RoleType.DECODER
+
     def __init__(self, vae, pipeline=None, component_name: str = "vae") -> None:
         super().__init__()
         self.vae: ParallelTiledVAE = vae

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -35,6 +35,7 @@ from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
     refresh_context_on_dual_transformer,
     refresh_context_on_transformer,
 )
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.distributed import (
     cfg_model_parallel_all_reduce,
     get_local_torch_device,
@@ -152,6 +153,10 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
     the initial noise into the final output.
     """
 
+    @property
+    def role_affinity(self):
+        return RoleType.DENOISER
+
     def __init__(
         self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None
     ) -> None:

python/sglang/multimodal_gen/runtime/server_args.py

@@ -23,6 +23,12 @@ from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.configs.models.encoders import T5Config
 from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
 from sglang.multimodal_gen.configs.quantization.nunchaku import NunchakuSVDQuantArgs
+from sglang.multimodal_gen.runtime.disaggregation.disagg_args import (
+    DisaggArgsMixin,
+    add_disagg_cli_args,
+    convert_disagg_role_string,
+)
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
     NunchakuConfig,
 )
@@ -94,7 +100,7 @@ class Backend(str, Enum):
 
 
 @dataclasses.dataclass
-class ServerArgs:
+class ServerArgs(DisaggArgsMixin):
     # Model and path configuration (for convenience)
     model_path: str
 
@@ -230,10 +236,35 @@ class ServerArgs:
     # MoE parameters used by Wan2.2
     boundary_ratio: float | None = None
 
+    # Disaggregation — fields defined here, methods in DisaggArgsMixin,
+    # CLI registration in disagg_args.add_disagg_cli_args().
+    base_gpu_id: int = 0
+    disagg_role: RoleType = RoleType.MONOLITHIC
+    disagg_timeout: int = 600
+    disagg_dispatch_policy: str = "round_robin"
+    disagg_mode: bool = False
+    disagg_server_addr: str | None = None
+    encoder_urls: str | None = None
+    denoiser_urls: str | None = None
+    decoder_urls: str | None = None
+    encoder_tp: int | None = None
+    denoiser_tp: int | None = None
+    denoiser_sp: int | None = None
+    denoiser_ulysses: int | None = None
+    denoiser_ring: int | None = None
+    decoder_tp: int | None = None
+    disagg_transfer_pool_size: int = 256 * 1024 * 1024
+    disagg_p2p_hostname: str = "127.0.0.1"
+    disagg_ib_device: str | None = None
+    pool_work_endpoint: str | None = None
+    pool_result_endpoint: str | None = None
+
     # Logging
     log_level: str = "info"
     uvicorn_access_log_exclude_prefixes: list[str] = field(default_factory=list)
 
+    # get_role_parallelism, derive_pool_*_endpoint — from DisaggArgsMixin
+
     @property
     def broker_port(self) -> int:
         return self.port + 1
@@ -411,12 +442,22 @@ class ServerArgs:
             )
 
     def _adjust_network_ports(self):
+        # Disagg role instances (encoder/denoiser/decoder) don't serve HTTP,
+        # so skip settling the HTTP port to avoid unnecessary port collisions.
+        needs_http = self.disagg_role in (
+            RoleType.MONOLITHIC,
+            RoleType.SERVER,
+        )
+
         if self.strict_ports:
-            self._require_port(self.port, "HTTP")
+            if needs_http:
+                self._require_port(self.port, "HTTP")
             self._require_port(self.scheduler_port, "Scheduler")
-            self._require_port(self.master_port, "Master")
+            if self.master_port is not None:
+                self._require_port(self.master_port, "Master")
         else:
-            self.port = self.settle_port(self.port)
+            if needs_http:
+                self.port = self.settle_port(self.port)
             initial_scheduler_port = self.scheduler_port + (
                 random.randint(0, 100) if self.scheduler_port == 5555 else 0
             )
@@ -611,6 +652,9 @@ class ServerArgs:
         # configure logger before use
         configure_logger(server_args=self)
 
+        # Convert string disagg_role to enum (from CLI/config)
+        convert_disagg_role_string(self.__dict__)
+
         # 1. adjust parameters
         self._adjust_parameters()
 
@@ -699,6 +743,7 @@ class ServerArgs:
             default=ServerArgs.num_gpus,
             help="The number of GPUs to use.",
         )
+
         parser.add_argument(
             "--tp-size",
             type=int,
@@ -758,6 +803,9 @@ class ServerArgs:
             "Increase this value if you encounter 'Connection closed by peer' errors after the service is idle. ",
         )
 
+        # Disaggregated diffusion args (defined in disagg_args.py)
+        add_disagg_cli_args(parser)
+
         # Prompt text file for batch processing
         parser.add_argument(
             "--prompt-file-path",
@@ -1139,6 +1187,9 @@ class ServerArgs:
         if "backend" in kwargs and isinstance(kwargs["backend"], str):
             kwargs["backend"] = Backend.from_string(kwargs["backend"])
 
+        # Convert disagg_role string to enum if necessary
+        convert_disagg_role_string(kwargs)
+
         kwargs["pipeline_config"] = PipelineConfig.from_kwargs(kwargs)
         return cls(**kwargs)
 

python/sglang/multimodal_gen/runtime/utils/common.py

@@ -130,6 +130,7 @@ def get_zmq_socket(
     endpoint: str,
     bind: bool,
     max_bind_retries: int = 10,
+    same_port: bool = False,
 ) -> tuple[zmq.Socket, str]:
     """
     Create and configure a ZMQ socket.
@@ -140,10 +141,13 @@ def get_zmq_socket(
         endpoint: Endpoint string (e.g., "tcp://localhost:5555")
         bind: Whether to bind (True) or connect (False)
         max_bind_retries: Maximum number of retries if bind fails due to address already in use
+        same_port: If True, retry on the same port instead of incrementing.
+            Useful when the port must be fixed (e.g., disagg sockets where
+            DiffusionServer connects to a pre-determined port).
 
     Returns:
         A tuple of (socket, actual_endpoint). The actual_endpoint may differ from the
-        requested endpoint if bind retry was needed.
+        requested endpoint if bind retry was needed (and same_port is False).
     """
     mem = psutil.virtual_memory()
     total_mem = mem.total / 1024**3
@@ -182,13 +186,15 @@ def get_zmq_socket(
         port_match = re.search(r":(\d+)$", endpoint)
 
         if port_match and max_bind_retries > 1:
+            import time as _time
+
             original_port = int(port_match.group(1))
             last_exception = None
 
             for attempt in range(max_bind_retries):
                 try:
                     current_endpoint = endpoint
-                    if attempt > 0:
+                    if attempt > 0 and not same_port:
                         # Try next port (increment by 42 to match settle_port logic)
                         current_port = original_port + attempt * 42
                         current_endpoint = re.sub(
@@ -198,6 +204,11 @@ def get_zmq_socket(
                             f"ZMQ bind failed for port {original_port + (attempt - 1) * 42}, "
                             f"retrying with port {current_port} (attempt {attempt + 1}/{max_bind_retries})"
                         )
+                    elif attempt > 0:
+                        logger.info(
+                            f"ZMQ bind attempt {attempt + 1}/{max_bind_retries} "
+                            f"on same port {original_port}..."
+                        )
 
                     socket.bind(current_endpoint)
 
@@ -212,7 +223,21 @@ def get_zmq_socket(
                 except zmq.ZMQError as e:
                     last_exception = e
                     if e.errno == zmq.EADDRINUSE and attempt < max_bind_retries - 1:
-                        # Address already in use, try next port
+                        # Address already in use, retry
+                        # Longer sleep for same_port (waiting for TIME_WAIT release)
+                        _time.sleep(1.0 if same_port else 0.5)
+                        # Re-create socket since ZMQ socket state may be invalid after failed bind
+                        socket.close()
+                        socket = context.socket(socket_type)
+                        if endpoint.find("[") != -1:
+                            socket.setsockopt(zmq.IPV6, 1)
+                        if socket_type == zmq.PUSH:
+                            set_send_opt()
+                        elif socket_type == zmq.PULL:
+                            set_recv_opt()
+                        elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP, zmq.ROUTER]:
+                            set_send_opt()
+                            set_recv_opt()
                         continue
                     elif attempt == max_bind_retries - 1:
                         # Last attempt failed

python/sglang/multimodal_gen/test/run_suite.py

@@ -88,7 +88,9 @@ STANDALONE_FILES = {
         "../cli/test_generate_t2i_perf.py",
         "test_update_weights_from_disk.py",
     ],
-    "2-gpu": [],
+    "2-gpu": [
+        "test_disagg_server.py",
+    ],
 }
 
 # New standalone files may omit an estimate once to learn the real CI runtime.
@@ -99,7 +101,11 @@ STANDALONE_FILE_EST_TIMES = {
         "../cli/test_generate_t2i_perf.py": 240.0,
         "test_update_weights_from_disk.py": 480.0,
     },
-    "2-gpu": {},
+    "2-gpu": {
+        # Two disagg clusters × (~3 min startup + ~1 min generate) ≈ 8 min.
+        # Raise if CI reports a higher measured time.
+        "test_disagg_server.py": 600.0,
+    },
 }
 
 # Backward-compatible suite view for scripts that still operate on file lists.

python/sglang/multimodal_gen/test/server/test_disagg_server.py

@@ -0,0 +1,388 @@
+"""End-to-end tests for disaggregated diffusion.
+
+Launches encoder / denoiser / decoder role instances plus a DiffusionServer
+head, sends a generation request through the HTTP front-end, and verifies
+that a non-empty output comes back.
+
+Two configurations are covered:
+
+1. :class:`TestDisaggZImage1Rank` — 1 rank per role (baseline disagg path).
+2. :class:`TestDisaggZImage2RankDenoiser` — denoiser with
+   ``--denoiser-sp 2`` across 2 GPUs. Exercises the multi-rank receive path
+   where only rank 0 owns the RDMA TransferManager and must broadcast
+   prompt/image tensors to non-rank-0 ranks before
+   ``execute_forward`` — without that broadcast the denoising stage fails
+   ``verify_input`` on an empty ``prompt_embeds``.
+
+Run directly:
+
+    pytest -v python/sglang/multimodal_gen/test/server/test_disagg_server.py
+    pytest -v ... -k ZImage1Rank              # one class
+    pytest -v ... -k test_generates_image     # one test
+"""
+
+from __future__ import annotations
+
+import base64
+import os
+import signal
+import subprocess
+import time
+import unittest
+from pathlib import Path
+
+import requests
+import torch
+
+from sglang.multimodal_gen.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    find_free_port,
+    wait_for_server_health,
+)
+from sglang.test.test_utils import CustomTestCase
+
+HOST = "127.0.0.1"
+_LOG_DIR = Path(os.environ.get("SGLANG_TEST_LOG_DIR", "/tmp"))
+
+# Env knob: bump if a cold HF download is needed on a fresh CI runner.
+_STARTUP_TIMEOUT_S = float(os.environ.get("SGLANG_DISAGG_STARTUP_TIMEOUT", "600"))
+
+
+# ---------------------------------------------------------------------------
+# Process management
+# ---------------------------------------------------------------------------
+
+
+def _kill_tree(pid: int) -> None:
+    try:
+        os.killpg(os.getpgid(pid), signal.SIGKILL)
+    except (ProcessLookupError, PermissionError):
+        pass
+
+
+def _wait_for_log(path: Path, message: str, timeout: float) -> bool:
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        if path.exists():
+            try:
+                if message in path.read_text(errors="ignore"):
+                    return True
+            except OSError:
+                pass
+        time.sleep(2)
+    return False
+
+
+def _tail_log(path: Path, n: int = 50) -> str:
+    if not path.exists():
+        return f"<no log at {path}>"
+    try:
+        lines = path.read_text(errors="ignore").splitlines()
+    except OSError as e:
+        return f"<log read failed: {e}>"
+    return "\n".join(lines[-n:])
+
+
+# ---------------------------------------------------------------------------
+# Disagg cluster helper
+# ---------------------------------------------------------------------------
+
+
+class DisaggCluster:
+    """Launch encoder / denoiser / decoder / server as separate processes.
+
+    ``gpu_layout`` is a mapping role → list of physical GPU ids. The length
+    of each list determines ``--num-gpus`` for that role, and the first id is
+    passed as ``--base-gpu-id``. For a multi-rank role the GPUs must be
+    contiguous starting from ``base-gpu-id`` (sglang derives local_rank from
+    ``base-gpu-id + rank``).
+    """
+
+    def __init__(
+        self,
+        model: str,
+        name: str,
+        gpu_layout: dict[str, list[int]],
+        extra_role_args: dict[str, list[str]] | None = None,
+        startup_timeout: float = _STARTUP_TIMEOUT_S,
+    ) -> None:
+        self.model = model
+        self.name = name
+        self.gpu_layout = gpu_layout
+        self.extra_role_args = extra_role_args or {}
+        self.startup_timeout = startup_timeout
+        self._procs: list[subprocess.Popen] = []
+        self._fhs: list = []
+        self._logs: dict[str, Path] = {}
+        self._alloc_ports()
+
+    def _alloc_ports(self) -> None:
+        self.base_port = find_free_port(HOST)
+        self.api_port = find_free_port(HOST)
+        self._role_ports = {
+            "encoder": find_free_port(HOST),
+            "denoiser": find_free_port(HOST),
+            "decoder": find_free_port(HOST),
+        }
+
+    # -- context manager -----------------------------------------------------
+
+    def __enter__(self) -> "DisaggCluster":
+        for attempt in range(3):
+            try:
+                self._launch_roles()
+                self._launch_server_head()
+                self._warmup()
+                return self
+            except Exception as e:
+                print(
+                    f"[disagg-test] Cluster {self.name} attempt {attempt + 1} "
+                    f"failed: {e}",
+                    flush=True,
+                )
+                self.stop()
+                self._alloc_ports()
+                if attempt == 2:
+                    raise
+        return self  # unreachable
+
+    def __exit__(self, *exc) -> None:
+        self.stop()
+
+    # -- internals -----------------------------------------------------------
+
+    def _start_proc(self, cmd: list[str], log_path: Path) -> subprocess.Popen:
+        fh = open(log_path, "w")
+        proc = subprocess.Popen(
+            cmd,
+            stdout=fh,
+            stderr=subprocess.STDOUT,
+            preexec_fn=os.setsid,
+            env=os.environ.copy(),
+        )
+        self._procs.append(proc)
+        self._fhs.append(fh)
+        return proc
+
+    def _launch_roles(self) -> None:
+        for role in ("encoder", "denoiser", "decoder"):
+            port = self._role_ports[role]
+            gpus = self.gpu_layout[role]
+            log = _LOG_DIR / f"disagg_{self.name}_{role}.log"
+            self._logs[role] = log
+
+            cmd = [
+                "sglang",
+                "serve",
+                "--model-path",
+                self.model,
+                "--disagg-role",
+                role,
+                "--disagg-server-addr",
+                f"tcp://{HOST}:{self.base_port}",
+                "--scheduler-port",
+                str(port),
+                "--num-gpus",
+                str(len(gpus)),
+                "--base-gpu-id",
+                str(gpus[0]),
+                "--log-level",
+                "info",
+                *self.extra_role_args.get(role, []),
+            ]
+            self._start_proc(cmd, log)
+
+            ready_msg = f"Role {role.upper()} ready"
+            if not _wait_for_log(log, ready_msg, self.startup_timeout):
+                raise RuntimeError(
+                    f"{role} failed to start for {self.name}. Log tail:\n"
+                    f"{_tail_log(log)}"
+                )
+
+    def _launch_server_head(self) -> None:
+        log = _LOG_DIR / f"disagg_{self.name}_server.log"
+        self._logs["server"] = log
+        cmd = [
+            "sglang",
+            "serve",
+            "--model-path",
+            self.model,
+            "--disagg-role",
+            "server",
+            "--encoder-urls",
+            f"tcp://{HOST}:{self._role_ports['encoder']}",
+            "--denoiser-urls",
+            f"tcp://{HOST}:{self._role_ports['denoiser']}",
+            "--decoder-urls",
+            f"tcp://{HOST}:{self._role_ports['decoder']}",
+            "--scheduler-port",
+            str(self.base_port),
+            "--port",
+            str(self.api_port),
+            "--host",
+            HOST,
+            "--disagg-timeout",
+            "120",
+            "--log-level",
+            "info",
+        ]
+        self._start_proc(cmd, log)
+        try:
+            wait_for_server_health(
+                f"http://{HOST}:{self.api_port}",
+                path="/v1/models",
+                timeout=self.startup_timeout,
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"server head failed to become healthy for {self.name}: {e}\n"
+                f"Server log tail:\n{_tail_log(log)}"
+            ) from e
+
+    def _warmup(self) -> None:
+        """Send a warmup request to establish RDMA connections."""
+        try:
+            _generate_image(self.api_port, self.model)
+        except Exception as e:
+            raise RuntimeError(
+                f"Warmup request failed for {self.name}: {e}\n"
+                f"Server log tail:\n{_tail_log(self._logs.get('server', Path('/dev/null')))}"
+            ) from e
+
+    def stop(self) -> None:
+        for proc in self._procs:
+            _kill_tree(proc.pid)
+        for fh in self._fhs:
+            try:
+                fh.close()
+            except OSError:
+                pass
+        # Give OS a moment to release ports before the next test.
+        time.sleep(3)
+        self._procs.clear()
+        self._fhs.clear()
+
+
+# ---------------------------------------------------------------------------
+# Request helpers
+# ---------------------------------------------------------------------------
+
+
+def _generate_image(api_port: int, model: str) -> bytes:
+    # Use raw requests (openai SDK pulls in a lot and complicates CI deps).
+    resp = requests.post(
+        f"http://{HOST}:{api_port}/v1/images/generations",
+        json={
+            "model": model,
+            "prompt": "A sunset over mountains",
+            "n": 1,
+            "size": "1024x1024",
+            "response_format": "b64_json",
+        },
+        timeout=600,
+    )
+    if resp.status_code != 200:
+        print(
+            f"[disagg-test] Server returned {resp.status_code}: {resp.text[:2000]}",
+            flush=True,
+        )
+    resp.raise_for_status()
+    data = resp.json()
+    return base64.b64decode(data["data"][0]["b64_json"])
+
+
+# ---------------------------------------------------------------------------
+# Test classes
+# ---------------------------------------------------------------------------
+
+
+def _require_gpus(n: int) -> None:
+    available = torch.cuda.device_count() if torch.cuda.is_available() else 0
+    if available < n:
+        raise unittest.SkipTest(f"need {n} GPUs, have {available}")
+
+
+class _DisaggTestBase(CustomTestCase):
+    """Shared setup: launch cluster once per class, tear down at the end."""
+
+    model: str = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+    required_gpus: int = 2
+    cluster_name: str = ""
+    gpu_layout: dict[str, list[int]] = {}
+    extra_role_args: dict[str, list[str]] = {}
+
+    cluster: DisaggCluster | None = None
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        _require_gpus(cls.required_gpus)
+        cls.cluster = DisaggCluster(
+            model=cls.model,
+            name=cls.cluster_name,
+            gpu_layout=cls.gpu_layout,
+            extra_role_args=cls.extra_role_args,
+        )
+        cls.cluster.__enter__()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        if cls.cluster is not None:
+            # Dump log tails for debugging CI failures
+            for role_name, log_path in cls.cluster._logs.items():
+                print(
+                    f"\n=== [{cls.cluster_name}] {role_name} log tail ===",
+                    flush=True,
+                )
+                print(_tail_log(log_path, n=80), flush=True)
+            cls.cluster.stop()
+            cls.cluster = None
+        super().tearDownClass()
+
+
+class TestDisaggZImage1Rank(_DisaggTestBase):
+    """Baseline: 1 rank per role, 2 physical GPUs."""
+
+    cluster_name = "zimage_1rank"
+    required_gpus = 2
+    gpu_layout = {
+        "encoder": [0],
+        "denoiser": [1],
+        "decoder": [0],
+    }
+
+    def test_generates_image(self) -> None:
+        assert self.cluster is not None
+        img = _generate_image(self.cluster.api_port, self.model)
+        # A real PNG is well above 1 KB; catches empty / error responses.
+        self.assertGreater(len(img), 1_000, f"image too small: {len(img)} bytes")
+
+
+class TestDisaggZImage2RankDenoiser(_DisaggTestBase):
+    """Multi-rank denoiser (``--denoiser-sp 2``) on 2 GPUs.
+
+    Regression guard for the bug where non-rank-0 denoiser ranks entered
+    ``execute_forward`` with an empty Req because ``ParallelExecutor``'s
+    REPLICATED stage does not broadcast the batch. With the fix, rank 0
+    broadcasts both scalar and tensor fields over NCCL before compute.
+    """
+
+    cluster_name = "zimage_sp2"
+    required_gpus = 2
+    gpu_layout = {
+        "encoder": [0],
+        "denoiser": [0, 1],
+        "decoder": [0],
+    }
+    extra_role_args = {
+        "denoiser": ["--denoiser-sp", "2"],
+    }
+
+    def test_generates_image_with_sp2_denoiser(self) -> None:
+        assert self.cluster is not None
+        img = _generate_image(self.cluster.api_port, self.model)
+        self.assertGreater(len(img), 1_000, f"image too small: {len(img)} bytes")
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sglang/multimodal_gen/test/unit/test_server_args.py

@@ -81,6 +81,115 @@ class TestModelIdResolution(unittest.TestCase):
             _get_config_info("/data/no-such-model", model_id="NonExistentModelXYZ")
 
 
+class TestPerRoleParallelism(unittest.TestCase):
+    """Test per-role parallelism args and get_role_parallelism helper."""
+
+    def _from_dict(self, kwargs):
+        with patch.object(
+            PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
+        ):
+            return ServerArgs.from_dict(kwargs)
+
+    def test_defaults_are_none(self):
+        args = self._from_dict({"model_path": "/fake"})
+        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+
+        for role in [RoleType.ENCODER, RoleType.DENOISER, RoleType.DECODER]:
+            par = args.get_role_parallelism(role)
+            self.assertIsNone(par["tp_size"])
+            self.assertIsNone(par["sp_degree"])
+            self.assertIsNone(par["ulysses_degree"])
+            self.assertIsNone(par["ring_degree"])
+
+    def test_encoder_overrides(self):
+        args = self._from_dict({"model_path": "/fake", "encoder_tp": 2})
+        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+
+        par = args.get_role_parallelism(RoleType.ENCODER)
+        self.assertEqual(par["tp_size"], 2)
+        self.assertIsNone(par["sp_degree"])
+        self.assertIsNone(par["ulysses_degree"])
+        self.assertIsNone(par["ring_degree"])
+
+    def test_denoiser_overrides(self):
+        args = self._from_dict(
+            {
+                "model_path": "/fake",
+                "denoiser_tp": 1,
+                "denoiser_sp": 8,
+                "denoiser_ulysses": 4,
+                "denoiser_ring": 2,
+            }
+        )
+        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+
+        par = args.get_role_parallelism(RoleType.DENOISER)
+        self.assertEqual(par["tp_size"], 1)
+        self.assertEqual(par["sp_degree"], 8)
+        self.assertEqual(par["ulysses_degree"], 4)
+        self.assertEqual(par["ring_degree"], 2)
+
+    def test_decoder_overrides(self):
+        args = self._from_dict({"model_path": "/fake", "decoder_tp": 2})
+        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+
+        par = args.get_role_parallelism(RoleType.DECODER)
+        self.assertEqual(par["tp_size"], 2)
+        self.assertIsNone(par["sp_degree"])
+        self.assertIsNone(par["ulysses_degree"])
+        self.assertIsNone(par["ring_degree"])
+
+    def test_monolithic_returns_all_none(self):
+        args = self._from_dict({"model_path": "/fake", "encoder_tp": 2})
+        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+
+        par = args.get_role_parallelism(RoleType.MONOLITHIC)
+        self.assertIsNone(par["tp_size"])
+        self.assertIsNone(par["sp_degree"])
+
+    def test_mixed_roles_independent(self):
+        """Per-role args don't interfere with each other."""
+        args = self._from_dict(
+            {
+                "model_path": "/fake",
+                "encoder_tp": 1,
+                "denoiser_tp": 2,
+                "decoder_tp": 4,
+            }
+        )
+        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+
+        self.assertEqual(args.get_role_parallelism(RoleType.ENCODER)["tp_size"], 1)
+        self.assertEqual(args.get_role_parallelism(RoleType.DENOISER)["tp_size"], 2)
+        self.assertEqual(args.get_role_parallelism(RoleType.DECODER)["tp_size"], 4)
+
+    def test_cli_args_parsed(self):
+        """Per-role parallelism args are parsed from CLI."""
+        parser = FlexibleArgumentParser()
+        ServerArgs.add_cli_args(parser)
+        argv = [
+            "--model-path",
+            "/fake",
+            "--denoiser-tp",
+            "2",
+            "--denoiser-sp",
+            "4",
+            "--denoiser-ulysses",
+            "2",
+            "--denoiser-ring",
+            "2",
+            "--encoder-tp",
+            "1",
+        ]
+        args, unknown = parser.parse_known_args(argv)
+        self.assertEqual(args.denoiser_tp, 2)
+        self.assertEqual(args.denoiser_sp, 4)
+        self.assertEqual(args.denoiser_ulysses, 2)
+        self.assertEqual(args.denoiser_ring, 2)
+        self.assertEqual(args.encoder_tp, 1)
+        self.assertIsNone(args.decoder_tp)
+
+
 class TestPipelineResolutionCliOverride(unittest.TestCase):
     def setUp(self):
         _get_config_info.cache_clear()
@@ -102,25 +211,5 @@ class TestPipelineResolutionCliOverride(unittest.TestCase):
         self.assertEqual(server_args.pipeline_config.resolution, 768)
 
 
-class TestComponentPathParsing(unittest.TestCase):
-    def test_extract_component_paths_accepts_config_expanded_keys(self):
-        component_paths, remaining = ServerArgs._extract_component_paths(
-            [
-                "--component-paths.spatial-upsampler",
-                "/tmp/latent_upsampler",
-                "--component_paths.distilled-lora=/tmp/distilled.safetensors",
-            ]
-        )
-
-        self.assertEqual(
-            component_paths,
-            {
-                "spatial_upsampler": "/tmp/latent_upsampler",
-                "distilled_lora": "/tmp/distilled.safetensors",
-            },
-        )
-        self.assertEqual(remaining, [])
-
-
 if __name__ == "__main__":
     unittest.main()