chore: import upstream snapshot with attribution

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,28 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Compatibility shim for disaggregated diffusion argument helpers."""
+
+from __future__ import annotations
+
+import argparse
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.server_args.disagg import DisaggServerArgsMixin
+
+# Keep the historical disagg_args import path working.
+DISAGG_RESULT_PORT_OFFSETS = DisaggServerArgsMixin.DISAGG_RESULT_PORT_OFFSETS
+DisaggArgsMixin = DisaggServerArgsMixin
+
+
+def add_disagg_cli_args(parser: argparse.ArgumentParser) -> None:
+    """Register disaggregated-diffusion CLI args through ServerArgs."""
+
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+    ServerArgs.add_disagg_cli_args(parser)
+
+
+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,99 @@
+# 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] + sorted(_ROLE_ALIASES)
+
+
+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",
+        "vision_language_encoder",
+    )
+    if any(
+        module_name == p or module_name.startswith(p + "_") for p in encoder_prefixes
+    ):
+        return RoleType.ENCODER
+
+    if module_name in {"hy3dshape_conditioner", "hy3dshape_image_processor"}:
+        return RoleType.ENCODER
+
+    denoising_prefixes = (
+        "transformer",
+        "unconditional_transformer",
+        "video_dit",
+        "audio_dit",
+        "dual_tower_bridge",
+    )
+    if any(
+        module_name == p or module_name.startswith(p + "_") for p in denoising_prefixes
+    ):
+        return RoleType.DENOISER
+
+    if module_name == "hy3dshape_model":
+        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
+
+    if module_name == "hy3dshape_vae":
+        return RoleType.DECODER
+
+    return None
+
+
+def filter_modules_for_role(
+    module_names: list[str],
+    role: "RoleType",
+    *,
+    extra_allowed_modules: set[str] | None = None,
+) -> list[str]:
+    """Filter module names to only those needed by the given role."""
+    if role in (RoleType.MONOLITHIC, RoleType.SERVER):
+        return module_names
+
+    extra_allowed_modules = extra_allowed_modules or set()
+    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 name in extra_allowed_modules:
+            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.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.get_device_module().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.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.get_device_module().stream(stream):
+                    return src.clone()
+            return src.clone()
+
+        if stream is not None:
+            with torch.get_device_module().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.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.get_device_module().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.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 or tensor.is_npu:
+            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,391 @@
+# 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,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+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.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.get_device_module().is_available():
+            torch.get_device_module().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.Stream | None = None,
+    ) -> tuple[StagedTransfer | None, torch.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.get_device_module().Event()
+            d2h_event.record(stream)
+        elif torch.get_device_module().is_available():
+            d2h_event = torch.get_device_module().Event()
+            d2h_event.record(torch.get_device_module().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 = current_platform.device_type,
+        stream: torch.Stream | None = None,
+    ) -> tuple[
+        dict[str, torch.Tensor | list[torch.Tensor]],
+        torch.get_device_module().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.get_device_module().Event()
+            load_event.record(stream)
+        elif torch.get_device_module().is_available():
+            load_event = torch.get_device_module().Event()
+            load_event.record(torch.get_device_module().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 = current_platform.device_type,
+        stream: torch.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.get_device_module().is_available():
+            torch.get_device_module().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)
+    )