[diffusion] refactor: simplify disaggregation transport hygiene (#33844)

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

@@ -44,6 +44,7 @@ from sglang.multimodal_gen.runtime.disaggregation.transport.protocol import (
     TransferMsgType,
     TransferPushedMsg,
     TransferRegisterMsg,
+    TransferStagedMsg,
     decode_transfer_msg,
     encode_transfer_msg,
     is_transfer_message,
@@ -106,15 +107,7 @@ _EXCLUDE_FIELDS = frozenset(
     }
 )
 
-# Sampling-params fields that should never be transferred across roles:
+# SamplingParams fields that are reconstructed locally or not JSON-safe.
-#   - data_type / supported_resolutions: enums / non-JSON classvars reconstructed on the receiver
-#   - teacache_params: model-specific object, not JSON-safe
-#   - output_* / save_output / return_*: output-side concerns owned by the decoder role
-#
-# Everything else on SamplingParams is forwarded automatically via a field-walk
-# below; this keeps new request-level features (e.g. Qwen-Image's
-# true_cfg_scale, guidance_rescale, cfg_normalization, ...) from silently
-# getting dropped just because nobody remembered to add them to a whitelist.
 _SAMPLING_PARAMS_EXCLUDE_FIELDS = frozenset(
     {
         "data_type",
@@ -123,11 +116,6 @@ _SAMPLING_PARAMS_EXCLUDE_FIELDS = frozenset(
     }
 )
 
-_BASE_SP_DEFAULTS: dict[str, Any] = {}
-for _f in dataclasses.fields(SamplingParams):
-    if _f.default is not dataclasses.MISSING:
-        _BASE_SP_DEFAULTS[_f.name] = _f.default
-
 
 def _is_tensor_like(value) -> bool:
     if isinstance(value, torch.Tensor):
@@ -290,8 +278,7 @@ def extract_transfer_fields(req) -> tuple[dict, dict]:
             value = getattr(sp, name, None)
             if value is None:
                 continue
-            base_default = _BASE_SP_DEFAULTS.get(name, dataclasses.MISSING)
+            if _is_default(value, f):
-            if base_default is not dataclasses.MISSING and value == base_default:
                 continue
             try:
                 scalar_fields[name] = _to_json_serializable(value)
@@ -669,12 +656,11 @@ class SchedulerDisaggMixin:
         Called from the recv prefetch thread. Loads on _transfer_stream
         and builds the Req, so the main thread can start compute immediately.
 
-        Returns (req, load_event, request_id, role_name, prealloc_slot_id).
+        Returns (req, load_event, request_id, prealloc_slot_id).
         """
         request_id = msg["request_id"]
         manifest = msg.get("manifest", {})
         scalar_fields = msg.get("scalar_fields", {})
-        role_name = self._disagg_role.value.upper()
 
         if self._disagg_metrics:
             self._disagg_metrics.record_request_start(request_id)
@@ -710,7 +696,7 @@ class SchedulerDisaggMixin:
         # running denoising loop on the main thread. Deferred to main thread
         # in _disagg_prefetch_event_loop, right before compute.
 
-        return (req, load_event, request_id, role_name, prealloc_slot_id, scalar_fields)
+        return req, load_event, request_id, prealloc_slot_id
 
     # ------------------------------------------------------------------
     # Broadcast
@@ -865,11 +851,7 @@ class SchedulerDisaggMixin:
           - queue timeout: broadcast "skip"
           - shutdown: broadcast None
         """
-        is_multi_rank = (
+        is_multi_rank = self._is_multi_rank()
-            self.server_args.sp_degree != 1
-            or self.server_args.tp_size > 1
-            or self.server_args.enable_cfg_parallel
-        )
 
         while self._running:
             try:
@@ -882,9 +864,7 @@ class SchedulerDisaggMixin:
 
                 if msg_type == "transfer_compute":
                     # Load already done by recv thread
-                    req, load_event, request_id, rn, prealloc_slot_id, scalar_fields = (
+                    req, load_event, request_id, prealloc_slot_id = data
-                        data
-                    )
                     # Wait for load to complete on compute stream
                     if load_event is not None:
                         torch.get_device_module().current_stream().wait_event(
@@ -913,9 +893,9 @@ class SchedulerDisaggMixin:
                         _init_disagg_request_scheduler(self, req)
                     # Run compute
                     if self._disagg_role == RoleType.DENOISER:
-                        self._disagg_denoiser_compute(req, request_id, rn)
+                        self._disagg_denoiser_compute(req, request_id)
                     elif self._disagg_role == RoleType.DECODER:
-                        self._disagg_decoder_compute(req, request_id, rn)
+                        self._disagg_decoder_compute(req, request_id)
 
                 elif msg_type == "transfer_control":
                     # alloc, push messages — handle on main thread (rank 0 only)
@@ -1249,8 +1229,6 @@ class SchedulerDisaggMixin:
         request_id = msg["request_id"]
         manifest = msg.get("manifest", {})
         scalar_fields = msg.get("scalar_fields", {})
-        role_name = self._disagg_role.value.upper()
-
         if self._disagg_metrics:
             self._disagg_metrics.record_request_start(request_id)
 
@@ -1299,9 +1277,9 @@ class SchedulerDisaggMixin:
 
         # 7. Run compute
         if self._disagg_role == RoleType.DENOISER:
-            self._disagg_denoiser_compute(req, request_id, role_name)
+            self._disagg_denoiser_compute(req, request_id)
         elif self._disagg_role == RoleType.DECODER:
-            self._disagg_decoder_compute(req, request_id, role_name)
+            self._disagg_decoder_compute(req, request_id)
 
     # ------------------------------------------------------------------
     # Compute
@@ -1407,9 +1385,7 @@ class SchedulerDisaggMixin:
         with trace_slice(ctx, DiffStage.SCHEDULER_DISPATCH, thread_finish_flag=True):
             yield
 
-    def _disagg_denoiser_compute(
+    def _disagg_denoiser_compute(self: Scheduler, req: Req, request_id: str) -> None:
-        self: Scheduler, req: Req, request_id: str, role_name: str
-    ) -> None:
         """Run denoiser compute in transfer mode, then stage output for decoder.
 
         Note: Scheduler timestep init is done in _handle_transfer_ready
@@ -1480,9 +1456,7 @@ class SchedulerDisaggMixin:
             duration_s,
         )
 
-    def _disagg_decoder_compute(
+    def _disagg_decoder_compute(self: Scheduler, req: Req, request_id: str) -> None:
-        self: Scheduler, req: Req, request_id: str, role_name: str
-    ) -> None:
         """Run decoder compute in transfer mode, send result to DS.
 
         Decoder result is sent as raw ZMQ multipart frames (same format as
@@ -1622,22 +1596,20 @@ class SchedulerDisaggMixin:
                 self._disagg_metrics.record_request_failed(request_id)
             return
 
-        # 2. Build transfer metadata dict while staging runs (CPU work, overlapped)
+        # 2. Build transfer metadata while staging runs (CPU work, overlapped)
-        staged_data = {
+        staged_msg = TransferStagedMsg(
-            "msg_type": "transfer_staged",
+            request_id=request_id,
-            "request_id": request_id,
+            data_size=staged.slot.size if staged.slot else 0,
-            "data_size": staged.slot.size if staged.slot else 0,
+            manifest=staged.manifest,
-            "manifest": staged.manifest,
+            session_id=self._transfer_manager.session_id,
-            "session_id": self._transfer_manager.session_id,
+            pool_ptr=self._transfer_manager.pool_data_ptr,
-            "pool_ptr": self._transfer_manager.pool_data_ptr,
+            slot_offset=staged.slot.offset if staged.slot else 0,
-            "slot_offset": staged.slot.offset if staged.slot else 0,
+            scalar_fields=staged.scalar_fields,
-            "scalar_fields": staged.scalar_fields,
+        )
-        }
-        msg_bytes = json.dumps(staged_data, separators=(",", ":")).encode("utf-8")
 
         # 3. Wait for staging to complete before sending (buffer must be ready)
         if stage_event is not None:
             stage_event.synchronize()
 
         # 4. Send transfer staged message
-        self._pool_result_push.send_multipart([TRANSFER_MAGIC, msg_bytes])
+        self._pool_result_push.send_multipart(encode_transfer_msg(staged_msg))

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

@@ -93,20 +93,6 @@ class BuddyAllocator:
                 "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:

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

@@ -101,7 +101,6 @@ class TransferTensorBuffer:
     def write_tensor(
         self,
         handle: SlotHandle,
-        name: str,
         tensor: torch.Tensor,
         byte_offset: int = 0,
         stream: torch.Stream | None = None,
@@ -186,7 +185,7 @@ class TransferTensorBuffer:
 
             entries = []
             if isinstance(value, torch.Tensor):
-                nbytes = self.write_tensor(handle, name, value, byte_offset, stream)
+                nbytes = self.write_tensor(handle, value, byte_offset, stream)
                 entries.append(
                     {
                         "offset": byte_offset,
@@ -201,9 +200,7 @@ class TransferTensorBuffer:
                 for i, t in enumerate(value):
                     if t is None:
                         continue
-                    nbytes = self.write_tensor(
+                    nbytes = self.write_tensor(handle, t, byte_offset, stream)
-                        handle, f"{name}[{i}]", t, byte_offset, stream
-                    )
                     entries.append(
                         {
                             "offset": byte_offset,
@@ -261,12 +258,3 @@ class TransferTensorBuffer:
                 )
 
         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/manager.py

@@ -22,7 +22,7 @@ logger = logging.getLogger(__name__)
 @dataclass
 class StagedTransfer:
     request_id: str
-    slot: SlotHandle
+    slot: SlotHandle | None
     manifest: dict
     scalar_fields: dict = field(default_factory=dict)
 
@@ -71,68 +71,6 @@ class DiffusionTransferManager:
     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,
@@ -315,39 +253,6 @@ class DiffusionTransferManager:
         )
         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":
@@ -365,27 +270,6 @@ class DiffusionTransferManager:
             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

@@ -6,12 +6,9 @@ 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__"
 
 
@@ -37,14 +34,11 @@ class TransferStagedMsg:
     msg_type: str = TransferMsgType.STAGED
     request_id: str = ""
     data_size: int = 0
-    manifest: dict = None
+    manifest: dict = field(default_factory=dict)
     session_id: str = ""
     pool_ptr: int = 0
     slot_offset: int = 0
-
+    scalar_fields: dict = field(default_factory=dict)
-    def __post_init__(self):
-        if self.manifest is None:
-            self.manifest = {}
 
 
 @dataclass