diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py b/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py index 52f6878ba..27c9300e7 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py +++ b/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: -# - 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. +# SamplingParams fields that are reconstructed locally or not JSON-safe. _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 base_default is not dataclasses.MISSING and value == base_default: + if _is_default(value, f): 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 = ( - self.server_args.sp_degree != 1 - or self.server_args.tp_size > 1 - or self.server_args.enable_cfg_parallel - ) + is_multi_rank = self._is_multi_rank() 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 = ( - data - ) + req, load_event, request_id, prealloc_slot_id = 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( - self: Scheduler, req: Req, request_id: str, role_name: str - ) -> None: + def _disagg_denoiser_compute(self: Scheduler, req: Req, request_id: 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( - self: Scheduler, req: Req, request_id: str, role_name: str - ) -> None: + def _disagg_decoder_compute(self: Scheduler, req: Req, request_id: 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) - staged_data = { - "msg_type": "transfer_staged", - "request_id": request_id, - "data_size": staged.slot.size if staged.slot else 0, - "manifest": staged.manifest, - "session_id": self._transfer_manager.session_id, - "pool_ptr": self._transfer_manager.pool_data_ptr, - "slot_offset": staged.slot.offset if staged.slot else 0, - "scalar_fields": staged.scalar_fields, - } - msg_bytes = json.dumps(staged_data, separators=(",", ":")).encode("utf-8") + # 2. Build transfer metadata while staging runs (CPU work, overlapped) + staged_msg = TransferStagedMsg( + request_id=request_id, + data_size=staged.slot.size if staged.slot else 0, + manifest=staged.manifest, + session_id=self._transfer_manager.session_id, + pool_ptr=self._transfer_manager.pool_data_ptr, + slot_offset=staged.slot.offset if staged.slot else 0, + scalar_fields=staged.scalar_fields, + ) # 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)) diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/transport/allocator.py b/python/sglang/multimodal_gen/runtime/disaggregation/transport/allocator.py index 49a5399a2..85299af28 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/transport/allocator.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py b/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py index 49c3faeb5..5afba92e2 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py +++ b/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( - handle, f"{name}[{i}]", t, byte_offset, stream - ) + nbytes = self.write_tensor(handle, 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 diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/transport/manager.py b/python/sglang/multimodal_gen/runtime/disaggregation/transport/manager.py index b9241710b..57c3a2f03 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/transport/manager.py +++ b/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") diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/transport/protocol.py b/python/sglang/multimodal_gen/runtime/disaggregation/transport/protocol.py index 347bf0be5..edd84f5a4 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/transport/protocol.py +++ b/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 - - def __post_init__(self): - if self.manifest is None: - self.manifest = {} + scalar_fields: dict = field(default_factory=dict) @dataclass